diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index 1bfc167d9..87b66dc14 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -1,16 +1,24 @@
 version: 2
 # Version updates for GitHub Actions only; CVEs are handled by Dependabot
-# security updates (enabled repo-wide).
+# security updates (enabled repo-wide), which are unaffected by the schedule
+# and ignore rules below.
 updates:
   - package-ecosystem: "github-actions"
     directory: "/"
     schedule:
-      interval: "weekly"
+      interval: "monthly"
     open-pull-requests-limit: 5
     groups:
       github-actions:
         patterns:
           - "*"
+    # Only open PRs for major version bumps; skip the churn of minor/patch
+    # version updates (security fixes still come through regardless).
+    ignore:
+      - dependency-name: "*"
+        update-types:
+          - "version-update:semver-minor"
+          - "version-update:semver-patch"
     labels:
       - "dependencies"
       - "github-actions"
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index f3916a33d..02b991ae2 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -2,15 +2,30 @@ name: Primus-CI-TAS
 
 on:
   workflow_dispatch:
+    inputs:
+      full_tests:
+        description: "Run the full test matrix (every model E2E, including @pytest.mark.weekly siblings, slow unit tests and the MaxText models)"
+        type: boolean
+        default: false
   push:
     branches:
       - main
     tags:
       - "v*"
   pull_request:
+    # Re-list the implicit defaults (opened/synchronize/reopened) and add
+    # ready_for_review so flipping a Draft PR to Ready triggers its first run.
+    types: [opened, synchronize, reopened, ready_for_review]
+  schedule:
+    # Saturday 18:00 UTC (Sunday 02:00 UTC+8) — the weekend full run. Kept far
+    # from Primus-Benchmark's daily 16:00 UTC cron; it also uses other runners.
+    - cron: "0 18 * * 6"
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.event.merge_group.head_ref || github.ref }}
+  # Scheduled runs get their own group: they share `github.ref` (refs/heads/main)
+  # with pushes to main, and cancel-in-progress would otherwise let any merge
+  # kill a multi-hour weekend run.
+  group: ${{ github.workflow }}-${{ github.event_name == 'schedule' && 'weekly' || github.head_ref || github.event.merge_group.head_ref || github.ref }}
   cancel-in-progress: true
 
 # Default to read-only. No job uses GITHUB_TOKEN for writes (Docker Hub push
@@ -20,12 +35,17 @@ permissions:
   contents: read
 
 env:
-  PRIMUS_TURBO_COMMIT: a04a233cbfb468dbe21600cbf9db70953428b25c # feat: force use nt layout gemm in bwd (#386)
-  PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # v0.14.0.post1
-  ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355)
+  # Test tier, the single switch every test step derives its pytest args from.
+  # "0" (PR/push) keeps one model E2E per architecture and feature path; "1"
+  # (weekend cron, or a manual run with full_tests=true) adds the
+  # @pytest.mark.weekly siblings and the slow unit tests.
+  PRIMUS_CI_FULL: ${{ (github.event_name == 'schedule' || github.event.inputs.full_tests == 'true') && '1' || '0' }}
+  PRIMUS_TURBO_COMMIT: 2f622de6a29ab711925e8fe6bcc763be76ac2699 # fix: loss NaN issue (Primus-Turbo bump); fix: fixed flydsl version (#446)
+  PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py
+  #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355)
   UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710)
-  TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0
-  BASE_IMAGE: docker.io/rocm/primus:v26.3
+  TRITON_COMMIT: 09500db9f0fe66fd176d1f080e2017b37e7e995d # pinned triton-lang/triton commit built from source in Dockerfile
+  BASE_IMAGE: docker.io/rocm/primus:v26.4
   MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2
 
 jobs:
@@ -52,6 +72,24 @@ jobs:
         run: pre-commit run --all-files --show-diff-on-failure
       - name: Check version/commit pin consistency
         run: python tools/ci/check_version_consistency.py
+      # The three submodules the drift check reads a schema from; `submodules:
+      # recursive` on the checkout would also clone Megatron-Bridge & friends
+      # for nothing. A fetch failure is deliberately left to the check step,
+      # which names the backend it could not check instead of dropping a bare
+      # git error; drift stays warn-only, but a backend it could not check
+      # fails the job, since a check that did not run must not read as a pass.
+      - name: Fetch backend schema sources
+        continue-on-error: true
+        run: >-
+          git submodule update --init --depth 1
+          third_party/torchtitan third_party/maxtext third_party/Megatron-LM
+      - name: Check backend config schema drift
+        # `pipefail` is not on by default, and without it `tee` would report
+        # success for a check that failed -- the same silent pass this step
+        # exists to stop.
+        run: |
+          set -o pipefail
+          python tools/ci/check_config_schema.py --warn-only | tee -a "$GITHUB_STEP_SUMMARY"
 
   # Flag PRs that introduce known-vulnerable dependencies. Warn-only during
   # burn-in; drop warn-only to turn it into a hard gate once the team is ready.
@@ -69,7 +107,11 @@ jobs:
 
   build-docker:
     needs: [code-lint]
+    # Skip on Draft PRs (keep code-lint ungated); the event_name guard preserves
+    # push/tag and workflow_dispatch runs where github.event.pull_request is absent.
+    if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
     runs-on: build-docker
+    # runs-on: build-docker-crusoe
     strategy:
       matrix:
         python-version: ["3.12"]
@@ -92,6 +134,11 @@ jobs:
             echo "IMAGE_TAG=pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
             echo "IMAGE_TAG=latest" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "schedule" ]]; then
+            # Own tag so the weekend run cannot race a concurrent main push over
+            # ":latest" (both would build the same commit, but not necessarily
+            # at the same time).
+            echo "IMAGE_TAG=weekly" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "release" ]]; then
             TAG_NAME="${{ github.ref }}"
             TAG="${TAG_NAME#refs/tags/}"
@@ -104,20 +151,33 @@ jobs:
           # echo "> Login to ROCm Docker Hub"
           # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
+          # This build-docker runner lives under /apps, which the Docker daemon
+          # cannot see (only its bind-mount /mnt/apps_proxy is visible to docker).
+          # Use the /mnt/apps_proxy view of GITHUB_WORKSPACE for every docker
+          # build context / -f path so BuildKit can read them.
+          CWS="${GITHUB_WORKSPACE}"
+          if [[ "${GITHUB_WORKSPACE}" == /apps/* ]]; then
+            CWS="/mnt/apps_proxy${GITHUB_WORKSPACE#/apps}"
+          fi
+
           echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}"
           echo "> Build dependencies"
           start_time=$(date +%s)
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
+          # Optional build-args, off by default (kept as an array so a commented-out
+          # element here can't break the docker build \-continuation below).
+          EXTRA_BUILD_ARGS=()
+          # EXTRA_BUILD_ARGS+=(--build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT})
+          docker build -f $CWS/.github/workflows/docker/Dockerfile \
                --network=host \
                -t tasimage/primus:${{env.IMAGE_TAG}} \
                --build-arg BASE_IMAGE=${BASE_IMAGE} \
                --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
                --build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \
-               --build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \
                --build-arg PRIMUS_TURBO_FRAMEWORK=PYTORCH \
                --build-arg UCCL_COMMIT=${UCCL_COMMIT} \
                --build-arg TRITON_COMMIT=${TRITON_COMMIT} \
-               $GITHUB_WORKSPACE/.github/workflows/docker
+               "${EXTRA_BUILD_ARGS[@]}" \
+               $CWS/.github/workflows/docker
           end_time=$(date +%s)
           elapsed=$((end_time - start_time))
           echo "⏱️ [build primus docker] Total elapsed time: ${elapsed} seconds"
@@ -127,18 +187,19 @@ jobs:
           docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}
           # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
-          # Primus v26.3 already includes AINIC under /workspace. Re-enable this
+          # Primus v26.4 already includes AINIC under /workspace. Re-enable this
           # Dockerfile.ainic build only when we need to refresh the tasimage -ainic image.
           echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-ainic"
           start_time=$(date +%s)
-          mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
-          cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-77.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile.ainic \
+          mkdir -p $CWS/.github/workflows/docker/ainic
+          cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-77.tar.gz $CWS/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
+          # cp /shared_nfs/tas/primus_github_runner/build-docker/ainic/ainic_bundle_1.117.5-a-77.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
+          docker build -f $CWS/.github/workflows/docker/Dockerfile.ainic \
                --network=host \
                -t tasimage/primus:${{env.IMAGE_TAG}}-ainic \
                --build-arg BASE_IMAGE=docker.io/tasimage/primus:${{env.IMAGE_TAG}} \
                --build-arg AINIC_BUNDLE_PATH=ainic \
-               $GITHUB_WORKSPACE/.github/workflows/docker
+               $CWS/.github/workflows/docker
           end_time=$(date +%s)
           elapsed=$((end_time - start_time))
           echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds"
@@ -148,140 +209,100 @@ jobs:
           docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic
           # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
-          echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-v25.09-ainic"
-          start_time=$(date +%s)
-          mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
-          cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile_v25.09_ainic \
-               --network=host \
-               -t tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic \
-               --build-arg AINIC_BUNDLE_PATH=ainic \
-               --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
-               $GITHUB_WORKSPACE/.github/workflows/docker
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "⏱️ [build primus docker-v25.09-ainic] Total elapsed time: ${elapsed} seconds"
-
-          docker tag tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic
-          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
-          docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic
-          # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
+          # echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax"
+          # start_time=$(date +%s)
+          # docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
+          #      --network=host \
+          #      -t tasimage/primus:${{env.IMAGE_TAG}}-jax \
+          #      --build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
+          #      --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
+          #      --build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \
+          #      --build-arg PRIMUS_TURBO_FRAMEWORK=JAX \
+          #      --build-arg UCCL_COMMIT=${UCCL_COMMIT} \
+          #      --build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \
+          #      --build-arg TRITON_COMMIT=${TRITON_COMMIT} .
+          # end_time=$(date +%s)
+          # elapsed=$((end_time - start_time))
+          # echo "⏱️ [build primus docker-jax] Total elapsed time: ${elapsed} seconds"
 
-          echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax"
-          start_time=$(date +%s)
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
-               --network=host \
-               -t tasimage/primus:${{env.IMAGE_TAG}}-jax \
-               --build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
-               --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \
-               --build-arg PRIMUS_TURBO_AITER_COMMIT=${PRIMUS_TURBO_AITER_COMMIT} \
-               --build-arg PRIMUS_TURBO_FRAMEWORK=JAX \
-               --build-arg UCCL_COMMIT=${UCCL_COMMIT} \
-               --build-arg ROCSHMEM_COMMIT=${ROCSHMEM_COMMIT} \
-               --build-arg TRITON_COMMIT=${TRITON_COMMIT} .
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "⏱️ [build primus docker-jax] Total elapsed time: ${elapsed} seconds"
+          # echo "> Docker tag image for Docker Hub"
+          # docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
+          # docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+          # docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
+          # # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
-          echo "> Docker tag image for Docker Hub"
-          docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
-          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
-          docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax
-          # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
+          # echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax-ainic"
+          # start_time=$(date +%s)
+          # mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
+          # cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
+          # # cp /shared_nfs/tas/primus_github_runner/build-docker/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
+          # docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile_jax.ainic \
+          #      --network=host \
+          #      -t tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic \
+          #      --build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
+          #      --build-arg AINIC_BUNDLE_PATH=ainic \
+          #      $GITHUB_WORKSPACE/.github/workflows/docker
+          # end_time=$(date +%s)
+          # elapsed=$((end_time - start_time))
+          # echo "⏱️ [build primus docker-jax-ainic] Total elapsed time: ${elapsed} seconds"
 
-          echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax-ainic"
-          start_time=$(date +%s)
-          mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic
-          cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; }
-          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile_jax.ainic \
-               --network=host \
-               -t tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic \
-               --build-arg BASE_IMAGE=${MAXTEXT_BASE_IMAGE} \
-               --build-arg AINIC_BUNDLE_PATH=ainic \
-               $GITHUB_WORKSPACE/.github/workflows/docker
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "⏱️ [build primus docker-jax-ainic] Total elapsed time: ${elapsed} seconds"
-
-          docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
-          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
-          docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
-          # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
+          # docker tag tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
+          # docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+          # docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-jax-ainic
+          # # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
           # echo "> Docker cleanup local images"
           # docker rmi tasimage/primus:${{env.IMAGE_TAG}}
-          # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic
-          # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-v25.10-ainic
           # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-jax
           echo "> build-docker success"
 
   run-unittest-torch:
+    permissions:
+      contents: read
+      actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing
     env:
+      # PRIMUS_WORKDIR: /shared_nfs/tas/primus_github_runner/primus-lm-cicd/actions-runner
       PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_ci/actions-runner-torch
       # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch
       PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32: 1
-    needs: [code-lint]
+      # All test steps run inside a sibling container launched from the image that
+      # build-docker built+pushed (docker.io/tasimage/primus:${IMAGE_TAG}). That
+      # image already bundles turbo / aiter / triton / rocSHMEM / uccl / origami,
+      # so their runtime installs are removed here.
+      UT_CONTAINER: primus_ut_torch
+    needs: [build-docker]
+    if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
     # runs-on: [primus-lm-cicd-torch-j8knc]
     runs-on: [primus-lm-cicd-v26.3-tas8n-a16-40]
+    # runs-on: [primus-lm-cicd-v26.3-crusoe-m2m-321]
     steps:
-      - run: echo "🎉 Begin Primus-Turbo Checkout."
-      - name: Clean stale Primus-Turbo checkout
-        run: rm -rf "${GITHUB_WORKSPACE}/Primus-Turbo"
-      - name: Set commit hash to env
-        run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV
-      - name: Checkout Repo Primus-Turbo
-        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-        with:
-          repository: AMD-AGI/Primus-Turbo
-          submodules: "recursive"
-          path: Primus-Turbo
-          ref: ${{ env.PRIMUS_TURBO_COMMIT }}
-      - run: echo "Begin AITER + Primus-Turbo Install."
-      - name: Install AITER
-        run: |
-          : > "$RUNNER_TEMP/runtime.tsv"  # reset the CI runtime log for this job
-          echo "✅ [Uninstall old aiter] started at: $(date)"
-          pip3 uninstall aiter amd-aiter -y || true
-          rm -rf /tmp/aiter || true
-          cd /tmp
-          git clone https://github.com/ROCm/aiter.git
-          cd aiter
-          git checkout -f ${PRIMUS_TURBO_AITER_COMMIT}
-          git submodule sync
-          git submodule update --init --recursive
-          start_time=$(date +%s)
-          echo "✅ [Build aiter] started at: $(date)"
-          PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 .
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "✅ [Build aiter] ended at: $(date)"
-          echo "⏱️ [Build aiter] Total elapsed time: ${elapsed} seconds"
-          echo -e "Build aiter\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
-      - name: Install Primus-Turbo
-        run: |
-          rm -rf /tmp/Primus-Turbo || true
-          mv Primus-Turbo /tmp/
-          echo "Primus-Turbo dir: /tmp/Primus-Turbo"
-          git config --global --add safe.directory /tmp/Primus-Turbo || true
-          cd /tmp/Primus-Turbo || true
-          start_time=$(date +%s)
-          echo "✅ [Pip install requirements] started at: $(date)"
-          mkdir -p ${PRIMUS_WORKDIR}/primus-cache
-          MAX_JOBS=128 pip install --cache-dir=${PRIMUS_WORKDIR}/primus-cache --no-build-isolation --no-clean -r requirements.txt
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "✅ [Pip install requirements] ended at: $(date)"
-          echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds"
-          echo -e "primus-turbo: pip install requirements\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
-          start_time=$(date +%s)
-          echo "✅ [build primus-turbo] started at: $(date)"
-          pip3 install --no-build-isolation -e . -v
-          end_time=$(date +%s)
-          elapsed=$((end_time - start_time))
-          echo "✅ [build primus-turbo] ended at: $(date)"
-          echo "⏱️ [build primus-turbo] Total elapsed time: ${elapsed} seconds"
-          echo -e "primus-turbo: build/install\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv"
-      - run: echo "🎉 Begin Primus Unit Test."
+      - run: echo "🎉 Begin Primus Unit Test (containerized)."
+      - name: Pre-checkout cleanup of root-owned leftovers
+        # Safety net for when a PREVIOUS run never reached its Clean step (runner
+        # crash/reboot/OOM, or cancellation past the grace period): it can leave
+        # root-owned files (E2E data/, output/, pp_simulation_result/,
+        # .pytest_cache, ...) in the workspace, and actions/checkout then fails
+        # with EACCES while deleting them ("Deleting the contents of ..."). Fix
+        # ownership here BEFORE checkout so the current run self-heals regardless
+        # of how the previous one ended. Runs as root via docker because the
+        # runner user cannot chown root-owned files. Best-effort (never fails).
+        run: |
+          WS="${GITHUB_WORKSPACE}"
+          CWS="$WS"; [[ "$WS" == /apps/* ]] && CWS="/mnt/apps_proxy${WS#/apps}"
+          if [ ! -d "$CWS" ]; then echo "Workspace $CWS not present yet; nothing to clean."; exit 0; fi
+          owner="$(stat -c '%u:%g' "$CWS")"
+          echo "Workspace=$CWS owner=$owner root-owned-before=$(find "$CWS" -user 0 2>/dev/null | wc -l)"
+          FIX="find '$CWS' -user 0 -exec chown $owner {} + 2>/dev/null; true"
+          if [ "$(docker inspect -f '{{.State.Running}}' "${UT_CONTAINER}" 2>/dev/null)" = "true" ]; then
+            echo "Fixing ownership via running container ${UT_CONTAINER}."
+            docker exec "${UT_CONTAINER}" bash -lc "$FIX" || true
+          else
+            IMG="$(docker images --format '{{.Repository}}:{{.Tag}}' | grep -m1 -E 'tasimage/primus|rocm/primus' || true)"
+            [ -n "$IMG" ] || IMG="${BASE_IMAGE}"
+            echo "Fixing ownership via throwaway container from image: $IMG"
+            docker run --rm -v /mnt/apps_proxy:/mnt/apps_proxy "$IMG" bash -lc "$FIX" || true
+          fi
+          echo "root-owned-after=$(find "$CWS" -user 0 2>/dev/null | wc -l)"
       - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           submodules: recursive
@@ -293,24 +314,73 @@ jobs:
           echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE"
           echo "Runner Temp Dir: $RUNNER_TEMP"
           echo "Runner Tool Cache: $RUNNER_TOOL_CACHE"
+      - name: Resolve container workspace path
+        run: |
+          # This runner lives under /apps, so GITHUB_WORKSPACE is an /apps path.
+          # On this host /apps cannot be bind-mounted into a container (the mount
+          # shows up empty); only its identical bind-mount /mnt/apps_proxy can.
+          # Export CWS = the /mnt/apps_proxy view of GITHUB_WORKSPACE and use it
+          # as the container-side workspace (docker -w and -e GITHUB_WORKSPACE).
+          # Host-side steps keep using GITHUB_WORKSPACE (/apps) since both paths
+          # point at the same files.
+          if [[ "${GITHUB_WORKSPACE}" == /apps/* ]]; then
+            echo "CWS=/mnt/apps_proxy${GITHUB_WORKSPACE#/apps}" >> "$GITHUB_ENV"
+          else
+            echo "CWS=${GITHUB_WORKSPACE}" >> "$GITHUB_ENV"
+          fi
+      - name: Resolve image tag
+        run: |
+          # build-docker sets IMAGE_TAG in its own job env only; re-derive it here
+          # with the identical logic so we can target the image it pushed.
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "IMAGE_TAG=pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
+            echo "IMAGE_TAG=latest" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "schedule" ]]; then
+            echo "IMAGE_TAG=weekly" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "release" ]]; then
+            TAG_NAME="${{ github.ref }}"
+            TAG="${TAG_NAME#refs/tags/}"
+            echo "IMAGE_TAG=$TAG" >> $GITHUB_ENV
+          else
+            echo "IMAGE_TAG=others" >> $GITHUB_ENV
+          fi
+      - name: Start (or reuse) unit-test container
+        run: |
+          : > "$RUNNER_TEMP/runtime.tsv"  # reset the CI runtime log for this job
+          IMG="docker.io/tasimage/primus:${{ env.IMAGE_TAG }}"
+          echo "Target image: ${IMG}"
+          # Pull first, then compare image id (not tag name): IMAGE_TAG is
+          # constant ("latest") on main, so a tag-only check can't detect a
+          # same-tag repush and would keep reusing a stale container.
+          docker pull "${IMG}"
+          target_id="$(docker image inspect -f '{{.Id}}' "${IMG}" 2>/dev/null || true)"
+          cur_id="$(docker inspect -f '{{.Image}}' "${UT_CONTAINER}" 2>/dev/null || true)"
+          running="$(docker inspect -f '{{.State.Running}}' "${UT_CONTAINER}" 2>/dev/null || true)"
+          if [ "${running}" = "true" ] && [ -n "${target_id}" ] && [ "${cur_id}" = "${target_id}" ]; then
+            echo "Reusing container ${UT_CONTAINER} (already on ${IMG}, id=${target_id})."
+          else
+            echo "Recreating container ${UT_CONTAINER} (running=${running}, cur_id=${cur_id}, target_id=${target_id})."
+            docker rm -f "${UT_CONTAINER}" >/dev/null 2>&1 || true
+            docker run -d --name "${UT_CONTAINER}" \
+              --network host --ipc host --privileged \
+              --group-add video \
+              --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
+              --shm-size 128G \
+              -e PYTHONDONTWRITEBYTECODE=1 \
+              --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+              -v /mnt/apps_proxy:/mnt/apps_proxy \
+              -w "${CWS}" \
+              "${IMG}" sleep infinity
+          fi
+          echo "Container ready:"
+          docker exec "${UT_CONTAINER}" bash -lc 'python3 --version; python3 -c "import triton; print('\''triton'\'', triton.__version__)" 2>/dev/null || true'
       - name: Install Primus
         run: |
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
           pip install -r requirements.txt
-      - name: Install fixed origami
-        run: |
-          # rocm/primus:v26.3 bundles origami 0.1.0, whose rank_configs() raises
-          # ValueError: vector::reserve during MoE grouped-gemm kernel selection and
-          # crashes training (turbo's _safe_rank_configs only catches RuntimeError).
-          # Install the fixed origami (rocm-libraries@223648a) over the bundled 0.1.0.
-          # TODO: drop this once a base image ships origami with the fix.
-          rm -rf /tmp/rocm-libraries
-          git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries
-          cd /tmp/rocm-libraries
-          git sparse-checkout init --cone
-          git sparse-checkout set shared/origami
-          git checkout 223648a26928ebed7f3dd0ccdc044c09f1dccf9b
-          pip uninstall -y origami || true
-          pip install ./shared/origami/python
+          IN
       - name: Set UT_LOG_PATH
         run: |
           ts="$(date +%Y%m%d-%H%M%S)"
@@ -319,6 +389,8 @@ jobs:
             echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/pr-${{ github.event.pull_request.number }}-${ts}-${commit_id}" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
             echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/main-${ts}-${commit_id}" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "schedule" ]]; then
+            echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/weekly-${ts}-${commit_id}" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "release" ]]; then
             TAG_NAME="${{ github.ref }}"
             TAG="${TAG_NAME#refs/tags/}"
@@ -328,70 +400,84 @@ jobs:
           fi
       - name: Run CLI Shell Tests
         run: |
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
           echo "Running Primus CLI shell tests..."
           bash ./tests/runner/run_all_tests.sh
+          IN
       - name: Run Primus Core Tests
         run: |
-          echo "Running Primus Core tests..."
           # Note: The tests `test_fp8_te_linear` and `test_te_linear` are temporarily skipped due to intermittent failures.
+          # Note: `test_limit_layers_moe_with_dense_keeps_two` is temporarily disabled pending a fix to
+          #       `_limit_layers_for_projection` dense-layer handling (num_layers collapses to 1 instead of 2).
+          # Note: `test_single_pass_te_vs_local` is temporarily disabled pending a fix to the
+          #       diffusion TE-vs-local spec attention parity check.
           # Note HSA_NO_SCRATCH_RECLAIM=1 must be set to avoid RCCL perf hit (TAS-8N Node), rocm ver:70125424
-          export HSA_NO_SCRATCH_RECLAIM=1
+          docker exec -i \
+            -e GITHUB_WORKSPACE="${CWS}" -e HSA_NO_SCRATCH_RECLAIM=1 -e PRIMUS_CI_FULL \
+            -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
+          echo "Running Primus Core tests..."
           mkdir -p "${GITHUB_WORKSPACE}/test-reports"
-          # Component-aware selection on PRs; full suite on push/release/dispatch.
-          # Fail-safe: any failure to compute the diff falls back to the full suite.
-          TARGETS="./tests/unit_tests/"
-          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
-            base_sha="${{ github.event.pull_request.base.sha }}"
-            git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true
-            changed="$(git diff --name-only "${base_sha}" HEAD 2>/dev/null || true)"
-            if [[ -n "${changed}" ]]; then
-              sel="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)"
-              [[ -n "${sel}" ]] && TARGETS="${sel}"
-            fi
-          fi
-          echo "Selected unit-test targets: ${TARGETS}"
-          # shellcheck disable=SC2086 # intentional word-splitting of multiple paths
-          pytest --maxfail=1 -s ${TARGETS} \
+          # Weekend full run also takes the @pytest.mark.slow release-tier shape
+          # gates (they self-skip off MI355X).
+          TIER_ARGS=()
+          [[ "${PRIMUS_CI_FULL:-0}" == "1" ]] && TIER_ARGS=(--run-slow)
+          pytest --maxfail=1 -s ./tests/unit_tests/ "${TIER_ARGS[@]}" \
             --cov=primus --cov-report=term-missing:skip-covered \
             --junitxml="${GITHUB_WORKSPACE}/test-reports/core-unit.xml" \
             --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \
             --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_te_linear \
             --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_forward_backward \
-            --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_capacity_forward_backward
+            --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_capacity_forward_backward \
+            --deselect=tests/unit_tests/core/projection/test_performance_projection.py::test_limit_layers_moe_with_dense_keeps_two \
+            --deselect=tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py::TestTEvsLocalSpecAttention::test_single_pass_te_vs_local
+          IN
       - name: Snapshot unit-test coverage
         if: always()
         continue-on-error: true
         run: |
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set +e
           # Keep the unit data for the final unit-vs-E2E comparison table
           # (.coverage.unit is fed to `coverage combine` after the E2E steps).
           python -m coverage json -o coverage_unit.json 2>/dev/null || true
-          cp .coverage "$GITHUB_WORKSPACE/.coverage.unit" 2>/dev/null || true
+          cp .coverage "${GITHUB_WORKSPACE}/.coverage.unit" 2>/dev/null || true
+          IN
       - name: Setup E2E training coverage
         if: always()
         continue-on-error: true
         run: |
           # Record coverage of the training subprocesses (primus-cli -> torchrun
-          # -> python). A .pth makes every new interpreter call
-          # coverage.process_startup(); COVERAGE_PROCESS_START points it at the rc
-          # and COVERAGE_FILE keeps E2E data separate from the unit .coverage.
+          # -> python). A .pth (written inside the container's site-packages) makes
+          # every new interpreter call coverage.process_startup(); COVERAGE_PROCESS_START
+          # points it at the rc and COVERAGE_FILE keeps E2E data separate from unit.
           # parallel=true -> one data file per rank; sigterm=true captures ranks
-          # killed on teardown. These env vars are inherited by the E2E steps.
+          # killed on teardown. These env vars are passed to the E2E steps via -e.
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
           SP=$(python -c "import site; print(site.getsitepackages()[0])")
           echo "import coverage; coverage.process_startup()" > "$SP/primus_e2e_coverage.pth"
-          printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e"
-          echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV"
-          echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV"
+          printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "${GITHUB_WORKSPACE}/.coveragerc_e2e"
+          IN
+          echo "COVERAGE_PROCESS_START=${CWS}/.coveragerc_e2e" >> "$GITHUB_ENV"
+          echo "COVERAGE_FILE=${CWS}/.coverage_e2e" >> "$GITHUB_ENV"
       - name: Decide E2E scope (torch)
         run: |
           # Fail-safe: default to running all E2E (push/release/dispatch, or any
           # diff failure). On PRs, narrow to the suites the changed files affect.
+          # Computed inside the container (git + select_tests.py), result read back
+          # on the runner to export RUN_*_E2E for the (host-side) skip guards.
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
           M=1; T=1
+          git config --global --add safe.directory "${GITHUB_WORKSPACE}" 2>/dev/null || true
           if [[ "${{ github.event_name }}" == "pull_request" ]]; then
             base="${{ github.event.pull_request.base.sha }}"
             git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true
             changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)"
             if [[ -n "${changed}" ]]; then
-              e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)"
+              e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)"
               echo "Selected torch E2E scope: ${e2e:-<none>}"
               if [[ "${e2e}" != "all" ]]; then
                 echo "${e2e}" | grep -qw megatron   || M=0
@@ -399,33 +485,79 @@ jobs:
               fi
             fi
           fi
-          echo "RUN_MEGATRON_E2E=${M}" >> "$GITHUB_ENV"
-          echo "RUN_TORCHTITAN_E2E=${T}" >> "$GITHUB_ENV"
+          printf '%s %s\n' "$M" "$T" > "${GITHUB_WORKSPACE}/.e2e_scope"
+          IN
+          read RM RT < "${GITHUB_WORKSPACE}/.e2e_scope"
+          echo "RUN_MEGATRON_E2E=${RM:-1}" >> "$GITHUB_ENV"
+          echo "RUN_TORCHTITAN_E2E=${RT:-1}" >> "$GITHUB_ENV"
+          if [[ "${PRIMUS_CI_FULL:-0}" == "1" ]]; then
+            echo "**Test tier:** full — every model E2E plus the slow unit tests." >> "$GITHUB_STEP_SUMMARY"
+          else
+            echo "**Test tier:** slim — one model E2E per architecture (\`-m 'not weekly'\`); the siblings run in the weekend scheduled build." >> "$GITHUB_STEP_SUMMARY"
+          fi
       - name: Run Primus Model Tests -- Megatron-LM
         env:
           HF_TOKEN: ${{secrets.HF_TOKEN}}
         run: |
           if [[ "${RUN_MEGATRON_E2E:-1}" != "1" ]]; then echo "Skipping Megatron-LM E2E: no relevant changes."; exit 0; fi
-          echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}"
-          rm -rf "${{ env.UT_LOG_PATH }}"
-          mkdir -p "${{ env.UT_LOG_PATH }}"
-          # MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
-          MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data HSA_NO_SCRATCH_RECLAIM=1 \
-          GPU_ARCHS=gfx942 \
-          pytest  --maxfail=1 -s ./tests/trainer/test_megatron_trainer.py \
+          docker exec -i \
+            -e GITHUB_WORKSPACE="${CWS}" -e UT_LOG_PATH -e HF_TOKEN \
+            -e COVERAGE_PROCESS_START -e COVERAGE_FILE -e PRIMUS_CI_FULL \
+            -e MASTER_PORT=10009 -e DATA_PATH=/mnt/apps_proxy/tas/0_public/data \
+            -e HSA_NO_SCRATCH_RECLAIM=1 \
+            -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
+          # Detect the real GPU arch (gfx950 on MI355X, gfx942 on MI300X, ...) so the
+          # aiter/turbo JIT kernels are built for the actual hardware. Hardcoding the
+          # wrong arch (e.g. gfx942 on an MI355X node) makes those kernels SIGSEGV.
+          export GPU_ARCHS="$(python3 -c 'import torch; print(torch.cuda.get_device_properties(0).gcnArchName.split(":")[0])' 2>/dev/null || echo gfx942)"
+          echo "Detected GPU_ARCHS=${GPU_ARCHS}"
+          echo "Set UT_LOG_PATH: ${UT_LOG_PATH}"
+          rm -rf "${UT_LOG_PATH}"
+          mkdir -p "${UT_LOG_PATH}"
+          # Temporarily deselect the Mamba / SSM tests: on gfx950 the pinned Triton
+          # crashes (SIGABRT) while compiling the mamba_ssm SSD backward kernel with an
+          # LLVM assertion (Sequence.h: `Begin <= End`) in make_amdgcn. Re-enable once
+          # the Triton/LLVM toolchain is fixed for this kernel.
+          TIER_ARGS=()
+          [[ "${PRIMUS_CI_FULL:-0}" == "1" ]] || TIER_ARGS=(-m "not weekly")
+          pytest  --maxfail=1 -s ./tests/trainer/test_megatron_trainer.py "${TIER_ARGS[@]}" \
+            --deselect tests/trainer/test_megatron_trainer.py::TestMegatronTrainer::test_mamba_130M_bridge_pretrain \
+            --deselect tests/trainer/test_megatron_trainer.py::TestMegatronTrainer::test_mamba_370M \
+            --deselect tests/trainer/test_megatron_trainer.py::TestMegatronTrainer::test_zebra_llama_1B_hybrid \
             --junitxml="${GITHUB_WORKSPACE}/test-reports/megatron-e2e.xml"
+          IN
       - name: Run Primus Model Tests -- TorchTitan
         env:
           HF_TOKEN: ${{secrets.HF_TOKEN}}
         run: |
           if [[ "${RUN_TORCHTITAN_E2E:-1}" != "1" ]]; then echo "Skipping TorchTitan E2E: no relevant changes."; exit 0; fi
-          echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}"
-          rm -rf "${{ env.UT_LOG_PATH }}"
-          mkdir -p "${{ env.UT_LOG_PATH }}"
-          # MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
-          MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data HSA_NO_SCRATCH_RECLAIM=1 \
-          pytest  --maxfail=1 -s ./tests/trainer/test_torchtitan_trainer.py \
+          docker exec -i \
+            -e GITHUB_WORKSPACE="${CWS}" -e UT_LOG_PATH -e HF_TOKEN \
+            -e COVERAGE_PROCESS_START -e COVERAGE_FILE -e PRIMUS_CI_FULL \
+            -e MASTER_PORT=10009 -e DATA_PATH=/mnt/apps_proxy/tas/0_public/data \
+            -e HSA_NO_SCRATCH_RECLAIM=1 \
+            -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set -e
+          echo "Set UT_LOG_PATH: ${UT_LOG_PATH}"
+          rm -rf "${UT_LOG_PATH}"
+          mkdir -p "${UT_LOG_PATH}"
+          # Temporarily deselect the deepseek_v3_671b config: under activation
+          # checkpointing the MoE dispatch is non-deterministic on recompute, so the
+          # recomputed dispatch tensors get a different token count than the forward
+          # pass (e.g. 635136 vs 635144) and torch.utils.checkpoint aborts with a
+          # CheckpointError. Re-enable once the MoE dispatch is recompute-stable.
+          # Temporarily deselect the deepseek_v3_16b_fp8 config: the aiter flash
+          # attention v3 backward kernel raises "RuntimeError: invalid argument for
+          # fmha_v3_bwd" during loss.backward() on the fp8 path (all ranks).
+          # Re-enable once the primus_turbo/aiter fp8 attention backward is fixed.
+          TIER_ARGS=()
+          [[ "${PRIMUS_CI_FULL:-0}" == "1" ]] || TIER_ARGS=(-m "not weekly")
+          pytest  --maxfail=1 -s ./tests/trainer/test_torchtitan_trainer.py "${TIER_ARGS[@]}" \
+            --deselect tests/trainer/test_torchtitan_trainer.py::TestTorchTitanTrainer::test_deepseek_v3_671b \
+            --deselect tests/trainer/test_torchtitan_trainer.py::TestTorchTitanTrainer::test_deepseek_v3_16b_fp8 \
             --junitxml="${GITHUB_WORKSPACE}/test-reports/torchtitan-e2e.xml"
+          IN
       - name: Run Primus CLI tool smoke (benchmark / preflight, E2E coverage)
         if: always()
         continue-on-error: true
@@ -433,8 +565,12 @@ jobs:
           # Best-effort CLI smoke so primus/tools (GEMM / attention / RCCL
           # benches + preflight probes) is exercised under the inherited E2E
           # coverage injection. Never fails CI: coverage is kept up to any exit.
-          export HSA_NO_SCRATCH_RECLAIM=1
-          OUT="${{ env.UT_LOG_PATH }}/cli_smoke"; mkdir -p "$OUT"
+          docker exec -i \
+            -e GITHUB_WORKSPACE="${CWS}" -e UT_LOG_PATH -e HSA_NO_SCRATCH_RECLAIM=1 \
+            -e COVERAGE_PROCESS_START -e COVERAGE_FILE \
+            -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set +e
+          OUT="${UT_LOG_PATH}/cli_smoke"; mkdir -p "$OUT"
           cli() { python -m primus.cli.main "$@" || true; }
           # single-GPU benches + preflight info
           cli benchmark gemm --M 512 --N 512 --K 512 --duration 2 --output-file "$OUT/gemm.md"
@@ -446,6 +582,7 @@ jobs:
           NGPU=$(python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1)
           [ "$NGPU" -ge 2 ] && torchrun --nproc_per_node="$NGPU" -m primus.cli.main benchmark rccl --op all_reduce all_gather --min-bytes 1K --max-bytes 256K --num-sizes 2 --iters 3 --warmup 1 --output-file "$OUT/rccl.md" || true
           [ "$NGPU" -ge 8 ] && torchrun --nproc_per_node=8 -m primus.cli.main benchmark strided-allgather --sizes-mb 8 --iters 3 --warmup 1 || true
+          IN
       - name: Upload test reports
         if: always()
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
@@ -457,23 +594,30 @@ jobs:
       - name: Write test summary
         if: always()
         run: |
-          python tools/ci/junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true
+          docker exec -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" \
+            bash -lc 'python tools/ci/junit_summary.py --title torch test-reports/*.xml' >> "$GITHUB_STEP_SUMMARY" || true
       - name: Write runtime summary
         if: always()
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          python tools/ci/runtime_summary.py --title torch "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true
+          docker exec -e GITHUB_WORKSPACE="${CWS}" -e GITHUB_TOKEN -e RUNNER_TEMP -w "${CWS}" "${UT_CONTAINER}" \
+            bash -lc 'python tools/ci/runtime_summary.py --title torch' >> "$GITHUB_STEP_SUMMARY" || true
       - name: Build torch coverage json (unit + E2E)
         if: always()
         continue-on-error: true
         run: |
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
+          set +e
           # Don't instrument coverage's own helper processes here.
           unset COVERAGE_PROCESS_START
           # Merge per-rank E2E data, then line-merge unit + E2E into one dataset.
           # The JSON is consumed by the coverage-summary job (no per-job summary).
           python -m coverage combine 2>/dev/null || true
-          COVERAGE_FILE="$GITHUB_WORKSPACE/.coverage_all" python -m coverage combine --keep \
-            "$GITHUB_WORKSPACE/.coverage.unit" "$GITHUB_WORKSPACE/.coverage_e2e" 2>/dev/null || true
-          COVERAGE_FILE="$GITHUB_WORKSPACE/.coverage_all" python -m coverage json -o coverage_combined.json 2>/dev/null || true
+          COVERAGE_FILE="${GITHUB_WORKSPACE}/.coverage_all" python -m coverage combine --keep \
+            "${GITHUB_WORKSPACE}/.coverage.unit" "${GITHUB_WORKSPACE}/.coverage_e2e" 2>/dev/null || true
+          COVERAGE_FILE="${GITHUB_WORKSPACE}/.coverage_all" python -m coverage json -o coverage_combined.json 2>/dev/null || true
+          IN
       - name: Upload torch coverage json
         if: always()
         continue-on-error: true
@@ -487,18 +631,43 @@ jobs:
       - name: Clean
         if: always()
         run: |
-          rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo
-          rm -rf ${PRIMUS_WORKDIR}/Primus
-          # Remove E2E coverage artifacts and, importantly, the subprocess
-          # injection .pth so it never leaks into later jobs on this persistent runner.
-          rm -f "${GITHUB_WORKSPACE}/.coverage_e2e"* "${GITHUB_WORKSPACE}/.coverage.unit" "${GITHUB_WORKSPACE}/.coverage_all" "${GITHUB_WORKSPACE}/.coveragerc_e2e" coverage_unit.json coverage_combined.json || true
-          SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth" || true
+          # The container steps run as ROOT and write root-owned files into the
+          # mounted GITHUB_WORKSPACE (__pycache__, .pytest_cache/, .hypothesis/,
+          # logs/, test-reports/, ut_out/, .coverage*, plus E2E-generated dirs like
+          # data/, output/, pp_simulation_result/). The next run's actions/checkout
+          # runs as the runner user and cannot delete those, so it fails at
+          # "Deleting the contents of ..." (e.g. EACCES on data/huggingface).
+          # Remove the well-known artifacts here AS ROOT (via the container), then
+          # chown whatever root-owned files remain back to the runner user so the
+          # next checkout can clean them. Artifacts have already been uploaded by the
+          # preceding upload steps. The container is left running for reuse.
+          docker exec -i -e GITHUB_WORKSPACE="${CWS}" "${UT_CONTAINER}" bash -s <<'IN' || true
+          set +e
+          cd "${GITHUB_WORKSPACE}" || exit 0
+          rm -rf logs test-reports ut_out .pytest_cache .hypothesis \
+                 .coverage .coverage.* .coverage_e2e* .coverage.unit .coverage_all \
+                 .coveragerc_e2e .e2e_scope coverage_unit.json coverage_combined.json
+          # Bytecode caches written during the run (root-owned; block checkout clean).
+          find . -type d -name __pycache__ -prune -exec rm -rf {} + 2>/dev/null
+          # Coverage subprocess-injection .pth in site-packages.
+          SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth"
+          # Safety net: hand any remaining root-owned files (E2E-generated data/,
+          # output/, pp_simulation_result/, .git internals, ...) back to the runner
+          # user so the next actions/checkout can delete them instead of hitting
+          # EACCES. Owner is taken from GITHUB_WORKSPACE itself (the runner user).
+          owner="$(stat -c '%u:%g' "${GITHUB_WORKSPACE}" 2>/dev/null)"
+          [ -n "$owner" ] && find "${GITHUB_WORKSPACE}" -user 0 -exec chown "$owner" {} + 2>/dev/null
+          IN
 
   run-unittest-jax:
+    permissions:
+      contents: read
+      actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing
     env:
       # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/jax
       PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_docker_jax_ci/actions-runner
     needs: [code-lint]
+    if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
     runs-on: [primus-jax-tas-runner] # docker container primus_jax_github_runner on tas a16-31
     steps:
       - run: echo "🎉 Begin Primus-Turbo Checkout."
@@ -565,6 +734,8 @@ jobs:
             echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/pr-${{ github.event.pull_request.number }}-${ts}-${commit_id}" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
             echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/main-${ts}-${commit_id}" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "schedule" ]]; then
+            echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/weekly-${ts}-${commit_id}" >> $GITHUB_ENV
           elif [[ "${{ github.event_name }}" == "release" ]]; then
             TAG_NAME="${{ github.ref }}"
             TAG="${TAG_NAME#refs/tags/}"
@@ -599,7 +770,7 @@ jobs:
             git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true
             changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)"
             if [[ -n "${changed}" ]]; then
-              e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)"
+              e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)"
               echo "Selected jax E2E scope: ${e2e:-<none>}"
               if [[ "${e2e}" != "all" ]]; then
                 echo "${e2e}" | grep -qw maxtext || X=0
@@ -619,6 +790,8 @@ jobs:
           # run_unit_tests.py shells out to pytest; PYTEST_ADDOPTS injects the
           # JUnit report without having to touch the wrapper script.
           export PYTEST_ADDOPTS="--junitxml=${GITHUB_WORKSPACE}/test-reports/maxtext-e2e.xml"
+          # JAX_SKIP_UT=1 in both tiers: the MaxText models it hides stay hidden
+          # even in the weekend full run, like the --deselect'ed cases.
           # MASTER_PORT=10009 DATA_PATH=/wekafs/primus-data \
           MASTER_PORT=10009 DATA_PATH=/mnt/apps_proxy/tas/0_public/data \
           JAX_SKIP_UT=1 python ./tests/run_unit_tests.py --jax
@@ -636,8 +809,10 @@ jobs:
           python tools/ci/junit_summary.py --title jax test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true
       - name: Write runtime summary
         if: always()
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          python tools/ci/runtime_summary.py --title jax "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true
+          python tools/ci/runtime_summary.py --title jax >> "$GITHUB_STEP_SUMMARY" || true
       - name: Build jax coverage json (MaxText E2E)
         if: always()
         continue-on-error: true
diff --git a/.github/workflows/deploy-backend-gap-dashboard.yml b/.github/workflows/deploy-backend-gap-dashboard.yml
index 2c61ea52d..958c9d655 100644
--- a/.github/workflows/deploy-backend-gap-dashboard.yml
+++ b/.github/workflows/deploy-backend-gap-dashboard.yml
@@ -1,15 +1,28 @@
 name: Deploy Backend Gap Dashboard
 
+# Report content lives on the orphan `dashboard-data` branch; tooling + site
+# shell live on `main`. This job checks out both and overlays the data branch's
+# `docs/` before building, so report runs never touch `main`.
+#
+# Pre-merge preview: run this workflow via "Run workflow" on a dev branch; it
+# builds and deploys the full site (dashboard + that branch's sections) so new
+# content can be validated before merging. Note: one Pages site per repo, so a
+# preview temporarily replaces production until the next normal run.
 on:
   workflow_dispatch:
+  # Companion `notify-dashboard-deploy.yml` on the dashboard-data branch fires
+  # this after a data push, so report updates publish without a timed schedule.
+  repository_dispatch:
+    types: [dashboard-data-updated]
+  # Tooling/shell/workflow changes and mounted section sources (declared in
+  # pages-sections.json) on main. Add a new section's source path here too.
   push:
     branches:
       - main
     paths:
-      - "docs/backend-gap/**"
-      - "docs/weekly_reports/**"
-      - "docs/monthly_reports/**"
       - "tools/backend_gap_report/**"
+      - "tools/pip_index/**"
+      - "examples/deepseek-v4/projection/site/**"
       - ".github/workflows/deploy-backend-gap-dashboard.yml"
 
 permissions:
@@ -21,12 +34,30 @@ concurrency:
   group: backend-gap-dashboard-pages
   cancel-in-progress: true
 
+env:
+  DASHBOARD_DATA_BRANCH: dashboard-data
+
 jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - name: Check out repository
+      - name: Check out repository (tooling + site shell from main)
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+
+      - name: Check out dashboard data branch
         uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ env.DASHBOARD_DATA_BRANCH }}
+          path: _dashboard_data
+
+      - name: Overlay report data onto the checkout
+        run: |
+          if [ -d _dashboard_data/docs ]; then
+            cp -r _dashboard_data/docs/. docs/
+          else
+            echo "::error::dashboard-data branch has no docs/ directory"
+            exit 1
+          fi
 
       - name: Set up Python
         uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
diff --git a/.github/workflows/docker-release/Dockerfile.jax-v26.5 b/.github/workflows/docker-release/Dockerfile.jax-v26.5
new file mode 100644
index 000000000..5d6e8a334
--- /dev/null
+++ b/.github/workflows/docker-release/Dockerfile.jax-v26.5
@@ -0,0 +1,281 @@
+ARG DEFAULT_BASE=ubuntu:24.04
+ARG BASE_IMAGE=${DEFAULT_BASE}
+FROM  ${BASE_IMAGE:-$DEFAULT_BASE} AS jax_base
+
+WORKDIR /workspace/
+ENV MAX_JOBS=128
+ENV PYTORCH_ROCM_ARCH="gfx942;gfx950"
+ENV ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+
+# Flag to fix profiler hang issue.
+ENV ROCPROFILER_QUEUE_INTERPOSITION=0
+ENV DEBUG_HIP_DYNAMIC_QUEUES=0
+
+# Install Ubuntu dependencies
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt update \
+	&& apt install -y \
+	gfortran \
+	git \
+	git-lfs \
+	ninja-build \
+	g++ \
+	pkg-config \
+	xxd \
+	patchelf \
+	automake \
+	libtool \
+	python3-venv \
+	python3-dev \
+	python3-pip \
+	python-is-python3 \
+	libegl1-mesa-dev \
+	wget \
+    sudo \
+    flex \
+    liblzma-dev \
+    ccache \
+    libdw1 \
+    libdrm-dev \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Setup python env
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+RUN pip install --upgrade pip \
+    && pip uninstall -y wheel \
+    && pip install \
+        cmake==3.31.6 \
+        ninja==1.11.1.3 \
+        wheel==0.45.1 \
+        packaging==25.0 \
+        setuptools==69.5.1 \
+    && rm -rf /root/.cache
+
+ENV THEROCK_TARBALL=https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz
+
+RUN mkdir -p /opt/rocm \
+    && cd /opt/rocm \
+    && wget ${THEROCK_TARBALL} \
+    && tar -xvf *.tar.gz -C /opt/rocm \
+    && rm *.tar.gz
+
+# Setup ROCm stack related envs
+ENV PATH="/opt/rocm/lib:/opt/venv/bin:/opt/rocm/bin:$PATH"
+ENV LD_LIBRARY_PATH="/opt/rocm/lib:/opt/rocm/lib/rocm_sysdeps/lib"
+ENV ROCM_PATH=/opt/rocm
+ENV HIP_PLATFORM=amd
+ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
+
+# Install LLVM 18
+RUN echo 'deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main' | sudo tee /etc/apt/sources.list.d/llvm.list
+RUN wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
+RUN sudo apt-get update
+RUN sudo apt-get install -y clang-18 lld-18 llvm-18-dev llvm-18-tools
+
+# amdsmi
+RUN pip install amdsmi==7.0.2 \
+    && rm -rf /root/.cache
+
+WORKDIR /workspace
+
+# Install maxtext
+ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git
+ARG MAXTEXT_BRANCH=release/v26.5
+
+RUN git clone ${MAXTEXT_REPO} \
+    && cd maxtext \
+    && git checkout ${MAXTEXT_BRANCH} \
+    && chmod +x src/dependencies/scripts/setup.sh \
+    && ./src/dependencies/scripts/setup.sh \
+    && cd .. \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/* \
+    && rm -rf /root/.cache
+
+# Fix (root cause): the prebuilt PyPI TensorFlow wheel bundles an LLVM whose symbols
+# collide with ROCm 7.14's libLLVM inside Grain "spawn" workers -> SIGSEGV on
+# `import tensorflow` after `import jax`. Rebuild the CPU-only TF wheel from ROCm's
+# fork (correct symbol visibility) and install it in place of the stock wheel. A CPU
+# build also has no bundled NCCL, so it preserves the previous XLA-RCCL fix.
+# NOTE: bumps TF to 2.21.0 and adds a long (~30-60 min) bazel build to the image.
+ARG TF_REPO=https://github.com/ROCm/tensorflow-upstream.git
+ARG TF_BRANCH=upstream-v2.21.0
+ARG BAZELISK_VERSION=v1.25.0
+# Bazelisk auto-selects the bazel version pinned by TF's .bazelversion.
+RUN apt-get update && apt-get install -y unzip zip \
+    && wget -O /usr/local/bin/bazel \
+        https://github.com/bazelbuild/bazelisk/releases/download/${BAZELISK_VERSION}/bazelisk-linux-amd64 \
+    && chmod +x /usr/local/bin/bazel \
+    && apt clean && rm -rf /var/lib/apt/lists/*
+RUN git clone --depth 1 --branch ${TF_BRANCH} ${TF_REPO} tensorflow-upstream \
+    && cd tensorflow-upstream \
+    && bazel build //tensorflow/tools/pip_package:wheel \
+        --repo_env=WHEEL_NAME=tensorflow_cpu \
+        --repo_env=HERMETIC_PYTHON_VERSION=3.12 \
+    && pip uninstall -y tensorflow tensorflow-cpu tensorflow_cpu \
+    && pip install --no-deps \
+        bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow_cpu-2.21.0-cp312-cp312-linux_x86_64.whl \
+    && cd .. \
+    && rm -rf tensorflow-upstream /root/.cache/bazel /root/.cache
+
+# JAX
+# Note: JAX and related libraries need to be installed before TE
+ARG JAX_VERSION=0.10.0
+ENV JAX_VERSION=${JAX_VERSION}
+# See: https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-pjrt/
+ENV JAX_PJRT_VERSION=0.10.0+rocm7.14.0
+# See: https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-plugin/
+ENV JAX_PLUGIN_VERSION=0.10.0+rocm7.14.0
+
+RUN pip install jax==${JAX_VERSION} jaxlib==${JAX_VERSION} scipy==1.16 \
+    && pip install \
+        --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
+        --pre jax_rocm7_pjrt==${JAX_PJRT_VERSION} \
+        --pre jax_rocm7_plugin==${JAX_PLUGIN_VERSION}  \
+    && rm -rf /root/.cache
+
+# TransformerEngine
+# See: https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-jax/
+ARG TE_VERSION=2.15.0.dev0+rocm7.15.0a20260707.72d01a0
+
+RUN pip install \
+        pybind11==3.0.4 \
+        importlib-metadata==8.7.1 \
+        pydantic==2.13.4 \
+        flax==0.12.2 \
+    && pip install \
+        --index-url https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/ \
+        --pre \
+        --no-build-isolation \
+        transformer_engine_rocm_jax==${TE_VERSION} \
+    && rm -rf /root/.cache
+
+# Configure required envs
+ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
+ENV NVTE_USE_HIPBLASLT=1
+ENV GPU_MAX_HW_QUEUES=2
+ENV HIP_FORCE_DEV_KERNARG=1
+ENV HSA_FORCE_FINE_GRAIN_PCIE=1
+ENV NVTE_FUSED_ATTN=1
+ENV NCCL_DEBUG=VERSION
+ENV NVTE_CK_USES_BWD_V3=1
+ENV NVTE_CK_USES_FWD_V3=1
+ENV NVTE_CK_IS_V3_ATOMIC_FP32=1
+ENV NVTE_CK_HOW_V3_BF16_CVT=2
+ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.9
+ENV XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_enable_command_buffer=''"
+
+WORKDIR /workspace
+
+# Primus
+ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git
+ARG PRIMUS_BRANCH=main
+
+RUN git clone --recurse-submodules ${PRIMUS_REPO} \
+    && cd Primus \
+    && git checkout ${PRIMUS_BRANCH} \
+    && git submodule update --init third_party/maxtext/ \
+    && cd .. \
+    && pip uninstall -y dataclasses dataclasses_json \
+    && rm -rf /root/.cache
+
+####### AINIC related installations #######
+RUN apt update \
+    && apt install -y \
+    libibverbs-dev \
+    jq \
+    dpkg-dev \
+    kmod \
+    xz-utils \
+    ibverbs-utils \
+    infiniband-diags \
+    rdma-core \
+    ethtool \
+    libevent-dev \
+    libhwloc-dev \
+    libmunge-dev \
+    software-properties-common \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install AMD AINIC library
+ARG AINIC_BUNDLE_VERSION="1.117.5-a-77"
+
+RUN add-apt-repository -y "deb https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_BUNDLE_VERSION} noble main" \
+    && apt update --allow-insecure-repositories \
+    && apt install -y --allow-unauthenticated libionic-dev \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Allow insecure repositories by default
+RUN echo 'Acquire::AllowInsecureRepositories "true";' >> /etc/apt/apt.conf.d/99allow-insecure-repositories
+
+# Install UCX
+ARG UCX_VERSION="1.18.0"
+
+RUN wget https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz \
+    && mkdir -p ucx-${UCX_VERSION} \
+    && tar -zxf ucx-${UCX_VERSION}.tar.gz -C ucx-${UCX_VERSION} --strip-components=1 \
+    && cd ucx-${UCX_VERSION} \
+    && mkdir build \
+    && cd build \
+    && ../configure --prefix=/workspace/ucx-${UCX_VERSION}/install --with-rocm=${ROCM_PATH} \
+    && make -j 16 \
+    && make install \
+    && cd ../.. \
+    && rm ucx-${UCX_VERSION}.tar.gz
+
+ENV UCX_INSTALL_DIR=/workspace/ucx-${UCX_VERSION}/install
+
+# Install OpenMPI
+ARG MPI_VERSION="4.1.6"
+
+RUN wget https://download.open-mpi.org/release/open-mpi/v$(echo "${MPI_VERSION}" | cut -d. -f1-2)/openmpi-${MPI_VERSION}.tar.gz \
+    && mkdir -p ompi-${MPI_VERSION} \
+    && tar -zxf openmpi-${MPI_VERSION}.tar.gz -C ompi-${MPI_VERSION} --strip-components=1 \
+    && cd ompi-${MPI_VERSION} \
+    && mkdir build \
+    && cd build \
+    && ../configure --prefix=${WORKDIR}/ompi-${MPI_VERSION}/install --with-ucx=${UCX_INSTALL_DIR} --disable-oshmem --disable-mpi-fortran \
+    && make -j 16 \
+    && make install \
+    && cd ../.. \
+    && rm openmpi-${MPI_VERSION}.tar.gz \
+    && rm -rf ompi-${MPI_VERSION}
+
+# Build RCCL with https://github.com/ROCm/rocm-systems/pull/8484
+# This commit `9e5e408` is on `develop` branch, commited on 2026-07-13.
+RUN rm -rf rocm-systems \
+    && git clone https://github.com/ROCm/rocm-systems.git \
+    && cd rocm-systems \
+    && git checkout 9e5e4084a4b8e1e86551b0eb054725c62354a926 \
+    && cd projects/rccl \
+    && ./install.sh -l \
+        --prefix build/ \
+        --amdgpu_targets="gfx942;gfx950" \
+    && cp -r build/release/librccl* /opt/rocm/lib/ \
+    && cd /workspace \
+    && rm -rf rocm-systems
+
+####### End of AINIC related installations #######
+
+# Secret scan will complain those files to be secrets. Explicitly remove them.
+RUN rm -f /workspace/Primus/.git/packed-refs /workspace/Primus/.git/modules/third_party/**/packed-refs
+
+# Clean cache
+RUN rm -rf /root/.cache
+
+# Training docker manifest
+ARG GIT_COMMIT_TAG=DEV
+ARG DOCKERFILE_PATH=Dockerfile
+
+# Copy patch files for training docker versioning
+RUN mkdir -p /workspace/.manifest \
+    && env > /workspace/.manifest/env.txt \
+    && pip list > /workspace/.manifest/requirements.txt \
+    && dpkg -l > /workspace/.manifest/dpkg-list.txt \
+    && echo "${GIT_COMMIT_TAG}" > /workspace/.manifest/training_docker_version
+
+COPY ${DOCKERFILE_PATH} /workspace/.manifest/Dockerfile
diff --git a/.github/workflows/docker-release/Dockerfile.primus-v26.5 b/.github/workflows/docker-release/Dockerfile.primus-v26.5
new file mode 100644
index 000000000..02529d57a
--- /dev/null
+++ b/.github/workflows/docker-release/Dockerfile.primus-v26.5
@@ -0,0 +1,502 @@
+# Primus Training Dockerfile for ROCm on TheRock
+ARG DEFAULT_BASE=ubuntu:24.04
+ARG BASE_IMAGE=${DEFAULT_BASE}
+FROM  ${BASE_IMAGE:-$DEFAULT_BASE} AS pytorch_base
+
+WORKDIR /workspace/
+ENV MAX_JOBS=128
+ENV PYTORCH_ROCM_ARCH="gfx942;gfx950"
+ENV ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+
+# Install Ubuntu dependencies
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt update \
+    && apt install -y \
+    gfortran \
+    git \
+    git-lfs \
+    ninja-build \
+    g++ \
+    pkg-config \
+    xxd \
+    patchelf \
+    automake \
+    libtool \
+    python3-venv \
+    python3-dev \
+    python3-pip \
+    python-is-python3 \
+    libegl1-mesa-dev \
+    wget \
+    sudo \
+    flex \
+    liblzma-dev \
+    ccache \
+    libdw1 \
+    libdrm-dev \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Setup python env
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+RUN pip install --upgrade pip \
+    && pip install \
+        pybind11 \
+        typeguard \
+        wheel==0.45.1 \
+        cmake==3.31.6 \
+        ninja==1.11.1.3 \
+        packaging==25.0 \
+        setuptools==75.1.0 \
+    && rm -rf /root/.cache
+
+RUN apt update \
+    && apt install -y ninja-build autoconf libtool flex rdma-core ffmpeg \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Workaround for the HSA_STATUS_ERROR_OUT_OF_RESOURCES issue
+ENV HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+ENV HSA_NO_SCRATCH_RECLAIM=1
+
+# Pytorch + rocm
+# See: https://rocm.nightlies.amd.com/whl-multi-arch/torch/
+ARG PYTORCH_VERSION=2.12.0+rocm7.15.0a20260720
+
+RUN pip install \
+    cxxfilt==0.3.0 \
+    tqdm==4.67.3 \
+    pyyaml==6.0.3 \
+    pytest==9.0.3 \
+    matplotlib==3.10.9 \
+    pandas==2.3.3 \
+    py-cpuinfo==9.0.0 \
+    build==1.5.0 \
+    && rm -rf /root/.cache
+
+RUN python -m pip uninstall -y torch && \
+    python -m pip install \
+        --index-url https://rocm.nightlies.amd.com/whl-multi-arch \
+        --pre \
+        torch==${PYTORCH_VERSION} \
+        amd-torch-device-gfx942==${PYTORCH_VERSION} \
+        amd-torch-device-gfx950==${PYTORCH_VERSION} \
+        rocm-sdk-devel \
+        rocm-sdk-device-gfx942 \
+        rocm-sdk-device-gfx950 \
+        torchaudio \
+        torchvision==0.27 \
+        amd-torchvision-device-gfx942==0.27 \
+        amd-torchvision-device-gfx950==0.27 \
+        apex \
+    && rm -rf /root/.cache
+
+RUN rocm-sdk init
+ENV ROCM_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel
+# Note: Primus uses `ROCM_HOME` instead of `ROCM_PATH`. Set both to avoid potential issues.
+ENV ROCM_HOME=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel
+ENV HIP_PLATFORM=amd
+ENV HIP_PATH=$ROCM_PATH
+ENV HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
+ENV HIP_INCLUDE_PATH=$ROCM_PATH/include
+ENV HIP_LIB_PATH=$ROCM_PATH/lib
+ENV HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
+ENV PATH="$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
+ENV LD_LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/host-math/lib:$ROCM_PATH/lib/rocm_sysdeps/lib"
+ENV LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64"
+ENV CPATH=$HIP_INCLUDE_PATH
+ENV PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+
+# FA
+ARG FA_REPO=https://github.com/ROCm/flash-attention.git
+ARG FA_BRANCH=6387433156558135a998d5568a9d74c1778666d8
+ENV GPU_ARCHS="${PYTORCH_ROCM_ARCH}"
+
+RUN git clone --recursive ${FA_REPO} \
+    && cd flash-attention \
+    && git checkout ${FA_BRANCH} \
+    && python setup.py install \
+    && cd .. \
+    && rm -rf flash-attention \
+    && rm -rf /root/.cache
+
+# TransformerEngine
+# See: https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-torch/
+ARG TE_VERSION=2.15.0.dev0+rocm7.15.0a20260716.a07e607
+
+# Those envs are required for getting the best performance
+ENV NVTE_USE_CAST_TRANSPOSE_TRITON=1
+ENV NVTE_CK_IS_V3_ATOMIC_FP32=0
+ENV NVTE_CK_USES_BWD_V3=1
+ENV NVTE_CK_USES_FWD_V3=1
+ENV CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2
+ENV NVTE_CK_HOW_V3_BF16_CVT=2
+
+RUN pip install \
+        pybind11==3.0.4 \
+        importlib-metadata==8.7.1 \
+        onnxscript==0.7.0 \
+        pydantic==2.13.4 \
+        nvdlfw_inspect==0.2.2 \
+    && rm -rf /root/.cache
+RUN pip install \
+        --index-url https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/ \
+        --pre \
+        --no-build-isolation \
+        transformer_engine_rocm_torch==${TE_VERSION} \
+    && sed -i 's|    mv -n "$_TMP_SO" "$OUTPUT"$|    mv -n "$_TMP_SO" "$OUTPUT" 2>/dev/null \|\| true|' \
+        /opt/venv/lib/python3.12/site-packages/transformer_engine/lib/ck_jit/ck_jit_compile.sh \
+    && rm -rf /root/.cache
+
+# Flux
+ARG FLUX_REPO=https://github.com/ROCm/AMDiffusionBenchmark.git
+ARG FLUX_BRANCH=26.4_RC_Patch_torchcodec
+
+RUN git clone ${FLUX_REPO} \
+    && cd AMDiffusionBenchmark \
+    && git checkout ${FLUX_BRANCH} \
+    && pip install -r requirements.txt \
+    && cd .. \
+    && rm -rf /root/.cache
+
+# DLRM
+ARG DLRM_REPO=https://github.com/AMD-AGI/DLRMBenchmark.git
+ARG DLRM_BRANCH=main
+
+RUN git clone ${DLRM_REPO} \
+    && cd DLRMBenchmark \
+    && git checkout ${DLRM_BRANCH} \
+    && cd .. \
+    && rm -rf /root/.cache
+
+# torchtune
+ARG TORCHTUNE_REPO=https://github.com/pytorch/torchtune.git
+ARG TORCHTUNE_BRANCH=b4c98ac2a37f0397d64c22579aed415ce7264db6
+
+RUN git clone ${TORCHTUNE_REPO} \
+    && if [ "${TORCHTUNE_BRANCH}" != "" ]; then cd torchtune && git checkout ${TORCHTUNE_BRANCH} && cd ..; fi \
+    && cd torchtune \
+    && git checkout ${TORCHTUNE_BRANCH} \
+    && find torchtune/modules/moe/utils.py -type f -print0 | xargs -0 sed -i 's/use_grouped_mm = True/use_grouped_mm = False/g' \
+    && pip install . \
+    && cd .. \
+    && rm -rf torchtune \
+    && rm -rf /root/.cache
+
+# Note: we need to change `pad_inner_dim` default explicitly to address fp8 in torchtune issue
+# torchao
+ARG TORCHAO_REPO=https://github.com/pytorch/ao.git
+ARG TORCHAO_BRANCH=e9c7bead90b840b280f97374308255957108ce47
+
+RUN git clone ${TORCHAO_REPO} \
+    && cd ao && git checkout ${TORCHAO_BRANCH} \
+    && find torchao/float8/config.py -type f -print0 | xargs -0 sed -i 's/pad_inner_dim: bool = False/pad_inner_dim: bool = True/g' \
+    && find torchao/csrc/rocm/swizzle/swizzle.cpp -type f -print0 | xargs -0 sed -i 's/if defined(HIPBLASLT_VEC_EXT)/if false/g' \
+    && pip install --no-build-isolation . \
+    && cd .. \
+    && rm -rf ao \
+    && rm -rf /root/.cache
+
+RUN pip install \
+    datasets==3.6.0 \
+    av==16.0.1 \
+    transformers==4.55.0 \
+    optree==0.18.0 \
+    sympy \
+    accelerate==1.9.0 \
+    trl==0.21.0 \
+    tensorboard==2.20.0 \
+    peft \
+    scipy \
+    einops \
+    flask-restful \
+    nltk \
+    pytest \
+    pytest-cov \
+    pytest_mock \
+    pytest-csv \
+    pytest-random-order \
+    sentencepiece \
+    wrapt \
+    zarr==2.18.7 \
+    numcodecs==0.12.1 \
+    xarray \
+    wandb \
+    tensorstore==0.1.45 \
+    pytest_mock \
+    pybind11 \
+    tiktoken \
+    pynvml \
+    huggingface_hub[cli] \
+    && python3 -m nltk.downloader punkt_tab \
+    && rm -rf /root/.cache
+
+# Groupped GEMM
+ARG GROUPED_GEMM_REPO=https://github.com/caaatch22/grouped_gemm.git
+ARG GROUPED_GEMM_BRANCH=rocm
+
+RUN git clone ${GROUPED_GEMM_REPO} \
+    && cd grouped_gemm \
+    && git checkout ${GROUPED_GEMM_BRANCH} \
+    && git submodule update --init --recursive \
+    && pip install --no-build-isolation . \
+    && cd .. \
+    && rm -rf grouped_gemm \
+    && rm -rf /root/.cache
+
+# Install Causal-Conv1d and its dependencies
+ARG CAUSAL_CONV1D_REPO=https://github.com/Dao-AILab/causal-conv1d
+ARG CAUSAL_CONV1D_BRANCH=e940ead2fd962c56854455017541384909ca669f
+
+ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE
+ENV MAMBA_FORCE_BUILD=TRUE
+ENV HIP_ARCHITECTURES=gfx942,gfx950
+
+RUN git clone ${CAUSAL_CONV1D_REPO} causal-conv1d \
+    && cd causal-conv1d \
+    && git checkout ${CAUSAL_CONV1D_BRANCH} \
+    && git show --oneline -s \
+    && pip install --no-build-isolation . \
+    && cd .. \
+    && rm -fr causal-conv1d \
+    && rm -rf /root/.cache
+
+# Mamba
+ARG MAMBA_REPO=https://github.com/AndreasKaratzas/mamba.git
+ARG MAMBA_BRANCH=enable-primus-hybrid-models
+
+RUN pip install apache-tvm-ffi==0.1.11 \
+    && git clone --branch ${MAMBA_BRANCH} ${MAMBA_REPO} \
+    && cd mamba \
+    && python setup.py install \
+    && cd .. \
+    && rm -rf /root/.cache
+
+# Primus
+ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git
+# Latest commit on `release/v26.5` branch. Commited on 2026-07-22.
+ARG PRIMUS_BRANCH=b511d1b66b0068715308ea9bfe8ba147ea1a3860
+# Note: those envs are required to resolve ther issue around the primus update
+# post v26.2
+ENV NVTE_FLASH_ATTN=0
+ENV NVTE_FUSED_ATTN=1
+
+RUN git clone --recurse-submodules ${PRIMUS_REPO} \
+    && cd Primus \
+    && git checkout ${PRIMUS_BRANCH} \
+    && git submodule update --init --recursive \
+    && pip install -r requirements.txt \
+    && cd .. \
+    && rm -rf /root/.cache
+
+# Aiter
+ARG AITER_REPO=https://github.com/ROCm/aiter.git
+ARG AITER_COMMIT=0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e
+
+RUN pip uninstall aiter amd-aiter -y \
+    && rm -rf aiter \
+    && git clone --recursive ${AITER_REPO} \
+    && cd aiter \
+    && git checkout ${AITER_COMMIT} \
+    && git submodule update --init --recursive \
+    && PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 . \
+    && cd .. \
+    && rm -rf aiter
+
+# Primus-Turbo
+ARG TURBO_REPO=https://github.com/AMD-AGI/Primus-Turbo.git
+# Latest commit on `main` branch. Commited on 2026-07-20.
+ARG TURBO_COMMIT=edc8d2ccb0be4888e80ee7c6e765fd3956026a32
+ENV MAX_JOBS=128
+ENV HCC_AMDGPU_TARGET="gfx942,gfx950"
+
+RUN git clone ${TURBO_REPO} --recursive \
+    && cd Primus-Turbo \
+    && git checkout ${TURBO_COMMIT} \
+    && git submodule update --init --recursive \
+    && pip3 install -r requirements.txt \
+    && pip3 install --no-build-isolation . -v \
+    && cd .. \
+    && rm -rf /root/.cache
+
+RUN pip install \
+    boto3==1.35.42 \
+    botocore==1.35.99 \
+    && rm -rf /root/.cache
+
+WORKDIR /workspace/
+
+# Install torchrec
+RUN pip install --no-deps torchrec \
+    && pip install \
+        tensordict \
+        iopath \
+        torchmetrics==1.0.3 \
+        git+https://github.com/mlperf/logging.git \
+        --extra-index-url https://rocm.nightlies.amd.com/whl-multi-arch \
+    && rm -rf /root/.cache
+
+# Install FBGEMM
+ARG FBGEMM_REPO=https://github.com/pytorch/FBGEMM.git
+# Latest known working commit on the `main` branch. Committed on 2026-06-17.
+ARG FBGEMM_COMMIT=80bd3c077dc41b55cd16ed4dcad15cf7c1c1d76a
+ENV BUILD_ROCM_VERSION='7.14'
+
+RUN apt update \
+    && apt install -y libtbb-dev \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone ${FBGEMM_REPO} \
+    && cd FBGEMM \
+    && git checkout ${FBGEMM_COMMIT} \
+    && cd fbgemm_gpu \
+    && git clean -dfx \
+    && git submodule sync \
+    && git submodule update --init --recursive \
+    && pip install -r requirements.txt \
+    && pip install setuptools==75.1.0 \
+    && python setup.py install \
+        --build-variant=rocm \
+        --build-target=default \
+        -DAMDGPU_TARGETS=$PYTORCH_ROCM_ARCH \
+        -DHIP_ROOT_DIR=$ROCM_PATH \
+        -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
+        -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA" \
+    && cd ../.. \
+    && rm -rf /root/.cache
+
+# Install numactl
+RUN apt update \
+    && apt install -y numactl \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# This is improtant to unset and reset the arc so aiter can
+# pick up the correct architecture
+ENV GPU_ARCHS=native
+
+RUN apt update \
+    && apt install -y \
+    libz3-dev \
+    pciutils \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Uninstall tilelang to avoid dependency issues
+# tilelang is installed as part of mamba-ssm package: https://github.com/state-spaces/mamba/blob/0048fbf2e7b2f214dcbe703ea3dec2b9647595e1/pyproject.toml#L20
+RUN pip uninstall -y tilelang
+
+####### AINIC related installations #######
+RUN apt update \
+    && apt install -y \
+    libibverbs-dev \
+    jq \
+    dpkg-dev \
+    kmod \
+    xz-utils \
+    ibverbs-utils \
+    infiniband-diags \
+    rdma-core \
+    ethtool \
+    libevent-dev \
+    libhwloc-dev \
+    libmunge-dev \
+    software-properties-common \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install AMD AINIC library
+ARG AINIC_BUNDLE_VERSION="1.117.5-a-77"
+
+RUN add-apt-repository -y "deb https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_BUNDLE_VERSION} noble main" \
+    && apt update --allow-insecure-repositories \
+    && apt install -y --allow-unauthenticated libionic-dev \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Allow insecure repositories by default
+RUN echo 'Acquire::AllowInsecureRepositories "true";' >> /etc/apt/apt.conf.d/99allow-insecure-repositories
+
+# Install UCX
+ARG UCX_VERSION="1.18.0"
+
+RUN wget https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz \
+    && mkdir -p ucx-${UCX_VERSION} \
+    && tar -zxf ucx-${UCX_VERSION}.tar.gz -C ucx-${UCX_VERSION} --strip-components=1 \
+    && cd ucx-${UCX_VERSION} \
+    && mkdir build \
+    && cd build \
+    && ../configure --prefix=/workspace/ucx-${UCX_VERSION}/install --with-rocm=${ROCM_PATH} \
+    && make -j 16 \
+    && make install \
+    && cd ../.. \
+    && rm ucx-${UCX_VERSION}.tar.gz
+
+ENV UCX_INSTALL_DIR=/workspace/ucx-${UCX_VERSION}/install
+
+# Install OpenMPI
+ARG MPI_VERSION="4.1.6"
+
+RUN wget https://download.open-mpi.org/release/open-mpi/v$(echo "${MPI_VERSION}" | cut -d. -f1-2)/openmpi-${MPI_VERSION}.tar.gz \
+    && mkdir -p ompi-${MPI_VERSION} \
+    && tar -zxf openmpi-${MPI_VERSION}.tar.gz -C ompi-${MPI_VERSION} --strip-components=1 \
+    && cd ompi-${MPI_VERSION} \
+    && mkdir build \
+    && cd build \
+    && ../configure --prefix=/opt/openmpi --with-ucx=${UCX_INSTALL_DIR} --disable-oshmem --disable-mpi-fortran \
+    && make -j 16 \
+    && make install \
+    && cd ../.. \
+    && rm openmpi-${MPI_VERSION}.tar.gz \
+    && rm -rf ompi-${MPI_VERSION}
+
+ENV PATH="/opt/openmpi/bin:${PATH}"
+ENV LD_LIBRARY_PATH="/opt/openmpi/lib:${LD_LIBRARY_PATH}"
+
+####### End of AINIC related installations #######
+
+# MLPerf related changes
+ARG TRAINING_RESULTS_REPO=https://github.com/mlcommons/training_results_v6.0.git
+# Latest commit on `main` branch. Commited on 2026-06-15.
+ARG TRAINING_RESULTS_COMMIT=eabf23a07b2a0c60a289ff871dc3a46fff0d0421
+
+RUN git clone ${TRAINING_RESULTS_REPO} --recursive \
+    && cd training_results_v6.0 \
+    && git checkout ${TRAINING_RESULTS_COMMIT} \
+    && pip install AMD/benchmarks/llama31_8b/implementations/MI355X_EPYC_9575F_primus/primus_mllog-0.1.21-py3-none-any.whl \
+    && cd .. \
+    && rm -rf training_results_v6.0 \
+    && rm -rf /root/.cache
+
+# Some secrets could leak from these files. Explicitly clean them out.
+RUN rm -f /workspace/Megatron-LM/.gitlab/stages/00.pre.yml \
+ /workspace/Primus/third_party/Megatron-LM/.gitlab/stages/00.pre.yml
+
+# Secret scan will complain those files to be secrets. Explicitly remove them.
+RUN rm -f /workspace/Primus/.git/packed-refs /workspace/Primus/.git/modules/third_party/**/packed-refs
+
+# Clean cache
+RUN rm -rf /root/.cache
+
+# Clean build files
+RUN rm -rf /*.whl \
+    && rm -rf /root/.triton \
+    && rm -rf /workspace/aiter/aiter/jit/build \
+    && rm -rf /workspace/TransformerEngine
+
+# Training docker manifest
+ARG GIT_COMMIT_TAG=DEV
+ARG DOCKERFILE_PATH=Dockerfile
+
+# Copy patch files for training docker versioning
+RUN mkdir -p /workspace/.manifest \
+    && env > /workspace/.manifest/env.txt \
+    && pip list > /workspace/.manifest/requirements.txt \
+    && dpkg -l > /workspace/.manifest/dpkg-list.txt \
+    && echo "${GIT_COMMIT_TAG}" > /workspace/.manifest/training_docker_version
+
+COPY ${DOCKERFILE_PATH} /workspace/.manifest/Dockerfile
+
+CMD ["/usr/bin/bash"]
diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile
index 84f7c5f39..89c25d508 100644
--- a/.github/workflows/docker/Dockerfile
+++ b/.github/workflows/docker/Dockerfile
@@ -1,10 +1,10 @@
-ARG BASE_IMAGE=docker.io/rocm/primus:v26.3
+ARG BASE_IMAGE=docker.io/rocm/primus:v26.4
 FROM ${BASE_IMAGE}
 
 ARG PRIMUS_TURBO_COMMIT
 ARG PRIMUS_TURBO_AITER_COMMIT
 ARG PRIMUS_TURBO_FRAMEWORK
-ARG ROCSHMEM_COMMIT
+# ARG ROCSHMEM_COMMIT # only needed if the rocSHMEM rebuild below is restored
 ARG UCCL_COMMIT
 ARG TRITON_COMMIT
 # Non-interactive APT
@@ -21,34 +21,25 @@ RUN rm -rf /var/lib/apt/lists/*
 # ---------------------------------------------------------------------------
 # Enviroment variables
 # ---------------------------------------------------------------------------
-ENV ROCSHMEM_HOME=/opt/rocshmem
-ENV UCX_HOME=/opt/ucx
-# ENV MPI_HOME=/opt/ompi
-# Use the system OpenMPI prefix from the v26.3 base image.
-ENV MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi
-ENV ROCM_HOME=/opt/rocm
 ENV PRIMUS_TURBO_FRAMEWORK=${PRIMUS_TURBO_FRAMEWORK}
 
-# ENV PATH="/opt/ompi/bin:/opt/ompi/sbin:${PATH}"
-# ENV LD_LIBRARY_PATH="/opt/ompi/lib:${LD_LIBRARY_PATH}"
 ENV GPU_ARCHS="gfx942;gfx950"
 ENV PYTORCH_ROCM_ARCH="gfx942;gfx950"
 ENV HCC_AMDGPU_TARGET="gfx942,gfx950"
-# ---------------------------------------------------------------------------
-# Install rocSHMEM
-# ---------------------------------------------------------------------------
-RUN mkdir -p /opt && cd /opt && \
-    git clone https://github.com/ROCm/rocSHMEM.git && \
-    cd rocSHMEM && \
-    git checkout ${ROCSHMEM_COMMIT} && \
-    mkdir build && \
-    cd build && \
-    MPI_ROOT=${MPI_HOME} UCX_ROOT=${UCX_HOME} INSTALL_PREFIX=${ROCSHMEM_HOME} ../scripts/build_configs/gda \
-                                                                                            -DGDA_IONIC=ON \
-                                                                                            -DGDA_MLX5=ON  \
-                                                                                            -DGDA_BNXT=ON  \
-                                                                                            -DUSE_IPC=ON
-RUN rm -rf /opt/rocSHMEM
+
+# rocSHMEM: used by Primus ODC via Primus-Turbo's odc_rocshmem_* extensions.
+# This base image already ships a prebuilt install (commit 17ff985c, version
+# 3.2.0) at $ROCSHMEM_HOME=/opt/rocshmem, so building it again is redundant.
+#
+# If a future base image drops it, restore by uncommenting below (and the
+# `ARG ROCSHMEM_COMMIT` above). Note: cmake reads the ROCm version from
+# -DROCM_PATH=, not $ROCM_PATH; omitting it falls back to the nonexistent
+# /opt/rocm and fails.
+# RUN mkdir -p /opt && cd /opt && git clone https://github.com/ROCm/rocSHMEM.git && \
+#     cd rocSHMEM && git checkout ${ROCSHMEM_COMMIT} && mkdir build && cd build && \
+#     MPI_ROOT=${MPI_HOME} UCX_ROOT=${UCX_HOME} INSTALL_PREFIX=${ROCSHMEM_HOME} \
+#     ../scripts/build_configs/gda -DROCM_PATH=${ROCM_PATH} \
+#       -DGDA_IONIC=ON -DGDA_MLX5=ON -DGDA_BNXT=ON -DUSE_IPC=ON
 
 # ---------------------------------------------------------------------------
 # Install Primus-Turbo
@@ -73,18 +64,41 @@ RUN cd /opt && \
 
 RUN rm -rf /opt/Primus-Turbo
 
+# Primus-Turbo links librocshmem.a with -l: (not --whole-archive) under
+# -fgpu-rdc/--hip-link, so rocSHMEM's team.cpp.o (which defines the plain,
+# non-namespaced rocshmem::ROCSHMEM_TEAM_WORLD symbol) never gets pulled into
+# libprimus_turbo_kernels.so; since it's a shared object the missing symbol
+# only surfaces as an ImportError at dlopen/import time, not at build time.
+# Work around it by re-exporting that symbol from a tiny shim .so, placed on
+# libprimus_turbo_kernels.so's existing rpath and wired in via patchelf.
+RUN PT_LIB=$(python3 -c "import primus_turbo, os; print(os.path.join(os.path.dirname(primus_turbo.__file__), 'lib', 'libprimus_turbo_kernels.so'))") && \
+    mkdir -p /tmp/rocshmem_team_shim && cd /tmp/rocshmem_team_shim && \
+    ar x ${ROCSHMEM_HOME}/lib/librocshmem.a team.cpp.o && \
+    g++ -shared -fPIC -o ${ROCSHMEM_HOME}/lib/librocshmem_team_shim.so team.cpp.o \
+      -L${ROCSHMEM_HOME}/lib -l:librocshmem.a -L${MPI_HOME}/lib -l:libmpi.so \
+      -Wl,-rpath,${ROCSHMEM_HOME}/lib -Wl,-rpath,${MPI_HOME}/lib && \
+    patchelf --add-needed librocshmem_team_shim.so "${PT_LIB}" && \
+    cd / && rm -rf /tmp/rocshmem_team_shim
+
 # ---------------------------------------------------------------------------
 # Install Triton
 # ---------------------------------------------------------------------------
 RUN cd /opt && \
-    git clone -b release/3.7.x https://github.com/triton-lang/triton.git && \
+    git clone https://github.com/triton-lang/triton.git && \
     cd triton && \
     git checkout ${TRITON_COMMIT} && \
-    pip3 install ninja cmake && \
-    pip3 install --no-build-isolation -v .
+    pip3 install -r python/requirements.txt && \
+    MAX_JOBS=96 pip3 install --no-build-isolation --force-reinstall --no-deps .
 
 RUN rm -rf /opt/triton
 
+# torch's dist-info still pins the base image's stock triton build; without this,
+# any later "pip install" (e.g. CI unit tests) sees torch's triton requirement as
+# unsatisfiable and silently upgrades torch from PyPI, dragging in a CUDA stack.
+RUN TORCH_META=$(python3 -c "import importlib.metadata as m; print(m.distribution('torch')._path / 'METADATA')") && \
+    TRITON_VER=$(python3 -c "import importlib.metadata as m; print(m.version('triton'))") && \
+    sed -i -E "s/^Requires-Dist: triton==.*/Requires-Dist: triton==${TRITON_VER}/" "${TORCH_META}"
+
 # ---------------------------------------------------------------------------
 # Install UCCL-EP (skip for JAX framework)
 # ---------------------------------------------------------------------------
@@ -100,12 +114,14 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
     rm -rf /opt/uccl; \
     fi
 
-# ---------------------------------------------------------------------------
-# Install fixed origami (rocm-libraries@223648a) over the base image's bundled
-# 0.1.0. The bundled origami's rank_configs() raises
-# `ValueError: vector::reserve` during MoE grouped-gemm kernel selection and
-# crashes training (turbo's _safe_rank_configs only catches RuntimeError).
-# Skipped for JAX. TODO: drop once a base image ships origami with the fix.
+# Install origami (rocm-libraries@223648a). The v26.4 base image ships no
+# origami at all. primus_turbo's MoE grouped-gemm kernel selection tolerates
+# its absence (falls back to a heuristic, no measurable perf loss on gfx942),
+# but Primus's performance-projection tool (primus projection performance)
+# hard-depends on origami as its only supported GEMM simulation backend and
+# has no fallback (see primus/core/projection/simulation_backends/factory.py)
+# — dropping this install broke TestProjectionSimulate::* outright. Skipped
+# for JAX, where the projection tool isn't used.
 RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
     rm -rf /tmp/rocm-libraries && \
     git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries && \
diff --git a/.github/workflows/docker/Dockerfile.ainic b/.github/workflows/docker/Dockerfile.ainic
index ca103df2d..5308ae2b9 100644
--- a/.github/workflows/docker/Dockerfile.ainic
+++ b/.github/workflows/docker/Dockerfile.ainic
@@ -7,12 +7,16 @@ ARG AINIC_BUNDLE_PATH
 # Non-interactive APT
 ENV DEBIAN_FRONTEND=noninteractive
 
+# `cmd | tee log` below reports only tee's exit code by default, silently
+# swallowing real failures; pipefail makes the RUN fail with `cmd` instead.
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 
 # ---------------------------------------------------------------------------
 # Install build dependencies
 # ---------------------------------------------------------------------------
+# initramfs-tools: required by the AINIC install.sh below.
 RUN apt-get update && \
-    apt-get install jq dpkg-dev kmod xz-utils \
+    apt-get install jq dpkg-dev kmod xz-utils initramfs-tools \
     libfmt-dev libboost-all-dev \
     libibverbs-dev ibverbs-utils infiniband-diags -y
 
@@ -20,8 +24,18 @@ RUN apt-get update && \
 # Enviroment variables
 # ---------------------------------------------------------------------------
 ENV WORKDIR=/workspace
-ENV ROCM_PATH=/opt/rocm
-ENV MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi
+# ROCM_PATH/MPI_HOME already come from the base image; mirror MPI_HOME here
+# instead of hardcoding the old v26.3 OpenMPI path.
+ENV MPI_PATH=${MPI_HOME}
+
+# v26.4 moved ROCm to a pip-installed rocm-sdk-devel, so /opt/rocm no longer
+# exists, and rocm-sdk-devel/bin/amdclang++ can't find its sibling clang++
+# (now under llvm/bin/). rccl/amd-anp still hardcode both, so symlink them in
+# rather than patching those repos.
+RUN ln -sfn "${ROCM_PATH}" /opt/rocm && \
+    for f in "${ROCM_PATH}"/llvm/bin/clang*; do \
+        ln -sf "$f" "${ROCM_PATH}/bin/$(basename "$f")"; \
+    done
 
 # =============================== Build AINIC Driver ===============================
 # WARNING: Please ensure the following environment variables are correctly set:
diff --git a/.github/workflows/docker/Dockerfile_v25.09_ainic b/.github/workflows/docker/Dockerfile_v25.09_ainic
deleted file mode 100644
index d8c5934a7..000000000
--- a/.github/workflows/docker/Dockerfile_v25.09_ainic
+++ /dev/null
@@ -1,106 +0,0 @@
-# Base image
-FROM docker.io/rocm/megatron-lm:v25.9_gfx950
-
-# Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
-ARG PRIMUS_TURBO_COMMIT
-ARG AINIC_BUNDLE_PATH
-
-# Install basic dependencies
-RUN apt-get update
-
-# Clone and install the Primus-Turbo
-WORKDIR /opt
-RUN mkdir -p /opt && cd /opt && \
-    git clone https://github.com/AMD-AGI/Primus-Turbo.git && \
-    cd Primus-Turbo && \
-    git checkout ${PRIMUS_TURBO_COMMIT} && \
-    git submodule update --init --recursive && \
-    pip3 install -r requirements.txt && \
-    GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
-
-RUN apt-get install --reinstall binutils -y && apt-get install numactl -y
-
-WORKDIR /opt
-ENV WORKDIR=/opt
-ENV ROCM_PATH=/opt/rocm
-
-RUN apt-get update && \
-    apt-get install jq dpkg-dev kmod xz-utils \
-    libfmt-dev libboost-all-dev \
-    libibverbs-dev ibverbs-utils infiniband-diags -y
-
-# =============================== Build AINIC Driver ===============================
-# WARNING: Please ensure the following environment variables are correctly set:
-# WARNING:     1. PATH: /usr/sbin must be included.
-# WARNING:     2. LD_LIBRARY_PATH: /usr/lib must be included.
-# WARNING: If these paths are missing, tools and libraries may not function correctly.
-# INFO: Installation completed successfully
-
-COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-56.tar.gz ${WORKDIR}
-RUN cd ${WORKDIR} && \
-    echo "Building ainic bundle... current directory: ${WORKDIR}" && \
-    tar zxf ainic_bundle_1.117.5-a-56.tar.gz && \
-    cd ainic_bundle_1.117.5-a-56 && \
-    tar zxf host_sw_pkg.tar.gz && \
-    cd host_sw_pkg && \
-    ./install.sh --domain=user -y 2>&1 | tee log_install.txt && \
-    cd ${WORKDIR} && \
-    apt-get install -y ./amd/ainic/deb-repo/libionic*.deb
-
-# =============================== Test AINIC Driver ===============================
-# ibv_devices
-# rdma link
-# ethtool -i enp9s0
-# ibv_devinfo -vv | grep GID
-
-# =============================== Build UCX ===============================
-RUN cd ${WORKDIR} && wget https://github.com/openucx/ucx/releases/download/v1.18.0/ucx-1.18.0.tar.gz && \
-    mkdir -p ucx-1.18.0 && \
-    tar -zxf ucx-1.18.0.tar.gz -C ucx-1.18.0 --strip-components=1 && \
-    cd ucx-1.18.0 && mkdir build && cd build && \
-    ../configure --prefix=${WORKDIR}/ucx-1.18.0/install --with-rocm=${ROCM_PATH} 2>&1 | tee log_ucx_configure.txt && \
-    make -j 16 2>&1 | tee log_ucx_build.txt && \
-    make install && \
-    cd ${WORKDIR}
-
-ENV UCX_INSTALL_DIR=${WORKDIR}/ucx-1.18.0/install
-
-# =============================== Build MPI ===============================
-RUN cd ${WORKDIR} && \
-    wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.6.tar.gz && \
-    mkdir -p ompi-4.1.6 && \
-    tar -zxf openmpi-4.1.6.tar.gz -C ompi-4.1.6 --strip-components=1 && \
-    cd ompi-4.1.6 && mkdir build && cd build && \
-    ../configure --prefix=${WORKDIR}/ompi-4.1.6/install --with-ucx=${UCX_INSTALL_DIR} \
-        --disable-oshmem --disable-mpi-fortran 2>&1 | tee log_mpi_configure.txt && \
-    make -j 16 2>&1 | tee log_mpi_build.txt && \
-    make install && \
-    cd ${WORKDIR}
-
-ENV MPI_PATH=${WORKDIR}/ompi-4.1.6/install
-
-# =============================== Build RCCL ===============================
-RUN cd ${WORKDIR} && \
-    git clone https://github.com/ROCm/rccl.git && \
-    cd rccl && git checkout drop/2025-08 && \
-    ./install.sh -l --prefix build/ --disable-mscclpp \
-        --disable-msccl-kernel --amdgpu_targets="gfx950" 2>&1 | tee log_rccl_install.txt && \
-    cd ${WORKDIR}
-
-ENV RCCL_HOME=${WORKDIR}/rccl
-
-# =============================== Build AMD ANP ===============================
-
-RUN cd ${WORKDIR} && git clone https://github.com/rocm/amd-anp.git && \
-cd amd-anp && git checkout tags/v1.1.0-5 && \
-sed -i '5a CFLAGS += --offload-arch=gfx950' ./Makefile && head -10 ./Makefile && \
-make -j 16 RCCL_BUILD=${RCCL_HOME}/build/release \
-           MPI_INCLUDE=${MPI_PATH}/include/ \
-           MPI_LIB_PATH=${MPI_PATH}/lib/ \
-           ROCM_PATH=${ROCM_PATH} 2>&1 | tee log_amd_anp_build.txt
-
-# Set the default working directory
-WORKDIR /opt
-
-# check the installed Primus-Turbo package
-RUN python3 -m pip show primus-turbo || true
diff --git a/.github/workflows/release-build-wheel.yml b/.github/workflows/release-build-wheel.yml
index 71ae789a2..9727d7c1f 100644
--- a/.github/workflows/release-build-wheel.yml
+++ b/.github/workflows/release-build-wheel.yml
@@ -149,6 +149,12 @@ jobs:
           format: spdx-json
           artifact-name: primus-sbom.spdx.json
           output-file: dist/primus-sbom.spdx.json
+          # The "Attach wheel to release" step below is the single place that
+          # uploads release assets (with clobber semantics it controls). Let
+          # this action only produce the SBOM as a workflow artifact, so it
+          # doesn't race that step by also auto-uploading to the release on
+          # `release` events -- that always fails with "asset already exists".
+          upload-release-assets: false
 
       - name: Attest build provenance
         uses: actions/attest-build-provenance@0f67c3f4856b2e3261c31976d6725780e5e4c373 # v4.1.1
diff --git a/.gitignore b/.gitignore
index 97c0687f4..48d575a0a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,10 +6,34 @@ dist/
 *.egg-info
 *~
 logs
+.pip_cache
 .vscode
 local/
 .gitmodules
 output
 experiment
-data
+/data/*
+!primus/backends/diffusion/data/
+!primus/backends/diffusion/data/**
+primus/backends/diffusion/data/**/__pycache__/
+primus/backends/diffusion/data/**/*.pyc
 pp_simulation_result
+
+# Allow the projection site's breakdown JSON (consumed by the static site /
+# GitHub Pages) despite the generic `data` ignore above.
+!examples/deepseek-v4/projection/site/data/
+!examples/deepseek-v4/projection/site/data/*.json
+
+*.log
+*.nohup
+*.zip
+
+# Local run artifacts dropped at the repo root. These capture node FQDNs and
+# absolute paths from whichever cluster the run happened on, so keep them out
+# of the history.
+/log
+/log.*
+/nohup.*
+/core.*.gpu
+.triton_cache_shared/
+.cursor/
diff --git a/.gitmodules b/.gitmodules
index 1d234e782..226dbc1ce 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -8,7 +8,7 @@
 [submodule "third_party/maxtext"]
 	path = third_party/maxtext
 	url = https://github.com/ROCm/maxtext.git
-	branch = release/v26.3
+	branch = release/v26.5
 [submodule "third_party/Emerging-Optimizers"]
 	path = third_party/Emerging-Optimizers
 	url = https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git
@@ -16,6 +16,14 @@
 [submodule "third_party/Megatron-Bridge"]
 	path = third_party/Megatron-Bridge
 	url = https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
+[submodule "third_party/mamba"]
+	path = third_party/mamba
+	url = https://github.com/AndreasKaratzas/mamba.git
+	branch = enable-primus-hybrid-models
 [submodule "third_party/HummingbirdXT"]
 	path = third_party/HummingbirdXT
 	url = https://github.com/AMD-AGI/HummingbirdXT.git
+[submodule "third_party/Automodel"]
+	path = third_party/Automodel
+	url = https://github.com/NVIDIA-NeMo/Automodel.git
+	branch = main
diff --git a/LICENSE b/LICENSE
index 297acbff2..96a4fb715 100644
--- a/LICENSE
+++ b/LICENSE
@@ -41,6 +41,10 @@ Primus uses or references the following third-party projects:
    - License: MIT License (Sea AI Lab, NVIDIA)
    - Repository: https://github.com/sail-sg/zero-bubble-pipeline-parallelism
 
+5. odc (On-Demand Communication)
+   - License: MIT License (Sea AI Lab)
+   - Repository: https://github.com/sail-sg/odc
+
 User must comply with the respective licenses of these third-party projects when using or distributing Primus.
 
 --------------------------------------------------------------------------------
diff --git a/README.md b/README.md
index 481149035..698034e0f 100644
--- a/README.md
+++ b/README.md
@@ -12,30 +12,48 @@
 
 ## ✨ Key Features
 
-- **🔄 Multi-Backend Support**: Seamlessly switch between Megatron-LM, TorchTitan, and other training frameworks
-- **🚀 Unified CLI**: One command interface for local development, containers, and Slurm clusters ([Docs](./docs/README.md))
-- **⚡ ROCm Optimized**: Deep integration with AMD ROCm stack and optimized kernels from Primus-Turbo
-- **📦 Production Ready**: Battle-tested on large-scale training with hundreds of GPUs
-- **🔌 Extensible Architecture**: Plugin-based design for easy integration of custom models and workflows
-- **🛡️ Enterprise Features**: Built-in fault tolerance, checkpoint management, and monitoring
+- **🔄 Multi-Backend Support**: Seamlessly switch between Megatron-LM, TorchTitan, JAX MaxText, Megatron-Bridge, and the diffusion backend under one configuration system
+- **🎯 Full Training Lifecycle**: Pretraining, SFT / LoRA post-training ([Docs](./docs/02-user-guide/posttraining.md)), and image / video diffusion training ([Docs](./docs/04-technical-guides/diffusion-models/README.md))
+- **🚀 Unified CLI**: One command interface for local development, containers, and Slurm clusters ([Docs](./docs/02-user-guide/cli-reference.md))
+- **⚡ ROCm Optimized**: Deep integration with AMD ROCm stack, Primus-Turbo kernels, DeepEP, fused [MegaMoE](./docs/04-technical-guides/mega-moe.md), and FP8 / MXFP8 / MXFP4 recipes
+- **📐 Plan Before You Train**: [Projection](./docs/02-user-guide/projection.md) estimates parallelism and memory fit before you allocate a cluster, and the [tuning agent](./docs/02-user-guide/tuning-agent.md) searches configs automatically
+- **📦 Production Ready**: Battle-tested on large-scale training with hundreds of GPUs; shipped as ROCm Docker images and a pip wheel
+- **🔌 Extensible Architecture**: Patch-based backend integration for custom models and workflows ([Docs](./docs/06-developer-guide/extending-backends.md))
+- **🛡️ Enterprise Features**: Built-in fault tolerance, checkpoint management, and monitoring with MLflow / TraceLens
 
 ---
 
 ## ✅ Supported Models (high level)
 
-- **Megatron-LM**: LLaMA2 / LLaMA3 / LLaMA4 families, DeepSeek-V2/V3, Mixtral-style MoE, and other GPT-style models
-- **TorchTitan**: LLaMA3 / LLaMA4, DeepSeek-V3, and related decoder-only architectures
-- **MaxText (JAX)**: LLaMA3.x and other MaxText-supported transformer models (subset; see MaxText docs for details)
+- **Megatron-LM**: LLaMA2 / LLaMA3.x / LLaMA4 families, DeepSeek-V2 / V3 / V4, Qwen2.5 and Qwen3 (dense and MoE), Mixtral, Grok, GPT-OSS 20B/120B, GLM, Kimi K2, MiniMax, LFM2, plus hybrid and linear-attention stacks (Mamba, Zebra-LLaMA with GDN / KDA)
+- **TorchTitan**: LLaMA3.x / LLaMA4, DeepSeek-V3 (16B to 671B), and Qwen3 0.6B to 32B
+- **MaxText (JAX)**: LLaMA2 / LLaMA3.x, DeepSeek-V2 16B, Mixtral-8x7B, Grok1, and Qwen3 14B / 30B-A3B (subset; see MaxText docs for details)
+- **Megatron-Bridge**: SFT and LoRA post-training for Qwen3 8B/32B, LLaMA3.1 70B, Zebra-LLaMA, and Mamba
+- **Diffusion**: Flux.1 (schnell / dev) text-to-image and Wan 2.1 / 2.2 text- and image-to-video
 
-For the full and up-to-date model matrix, see [Supported Models](./docs/backends/overview.md#supported-models).
+For the full and up-to-date model matrix, see [Supported Models](./docs/06-developer-guide/model-support-matrix.md).
 
 ---
 
 ## 🆕 What's New
 
+- **[2026/07/29]** ⚡ **MegaMoE** - FlyDSL-based fused MoE layer that folds expert all-to-all into the grouped GEMMs, plus FP4 grouped GEMM support ([MegaMoE guide](./docs/04-technical-guides/mega-moe.md))
+- **[2026/07/29]** Hybrid linear-attention models: Gated Delta Net (GDN) and Kimi Delta Attention (KDA) on Megatron-LM ([Hybrid models](./docs/04-technical-guides/hybrid-models/README.md))
+- **[2026/07/22]** Backend upgrades: TorchTitan v0.2.2 (PyTorch 2.12) with GPT-OSS, and MaxText v26.5
+- **[2026/07/17]** 🚀 **DeepSeek-V4 training support** - model definition, fused attention/MoE kernels, Muon optimizer, FP8/FP4 recipes, and a projection toolkit ([examples](./examples/deepseek-v4))
+- **[2026/07/16]** 🎨 **Diffusion backend** - Flux.1 image and Wan video training with FP8/MXFP4, FSDP2, and Energon data pipelines ([Diffusion docs](./docs/04-technical-guides/diffusion-models/README.md))
+- **[2026/07/14]** MLPerf Training 6.0 examples on MI355X: Llama2-70B LoRA, Llama3.1-8B, and GPT-OSS-20B ([examples](./examples/mlperf))
+- **[2026/06/15]** [Tuning agent](./docs/02-user-guide/tuning-agent.md) with memory-aware benchmarking for automatic config search
+- **[2026/06/08]** Primus is published as a pip wheel with a bundled `primus-cli` ([install](#install-as-a-python-package-pip))
+- **[2026/01/22]** Post-training via Megatron-Bridge - SFT and LoRA workflows ([Post-training](./docs/02-user-guide/posttraining.md))
+- **[2026/01/06]** MXFP4 low-precision training in the Megatron-LM backend, with MXFP8 recipes following in June
 - **[2025/12/17]** MoE Training Best Practices on AMD GPUs - [MoE Package Blog](https://rocm.blogs.amd.com/software-tools-optimization/primus-moe-package/README.html)
 - **[2025/11/14]** 🎉 **Primus CLI 1.0 Released** - Unified command-line interface with comprehensive documentation
 - **[2025/08/22]** Primus introduction [blog](https://rocm.blogs.amd.com/software-tools-optimization/primus/README.html)
+
+<details>
+<summary>Earlier updates</summary>
+
 - **[2025/06/18]** Added TorchTitan backend support
 - **[2025/05/16]** Added benchmark suite for performance evaluation
 - **[2025/04/18]** Added [Preflight](./primus/tools/preflight/README.md) cluster sanity checker
@@ -43,6 +61,8 @@ For the full and up-to-date model matrix, see [Supported Models](./docs/backends
 - **[2025/04/09]** Extended support for LLaMA2, LLaMA3, DeepSeek-V2/V3 models
 - **[2025/03/04]** Released Megatron trainer module
 
+</details>
+
 ---
 
 ## 🚀 Setup & Deployment
@@ -91,9 +111,9 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party
 
     ```bash
     # For Megatron-LM and TorchTitan backends
-    docker pull rocm/primus:v26.3
+    docker pull rocm/primus:v26.4
     # For MaxText backend
-    docker pull rocm/jax-training:v26.3
+    docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
     ```
 
 2. **Clone the repository**
@@ -102,7 +122,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party
     git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
     cd Primus
     # checkout the branch for the specific release
-    git checkout release/v26.3
+    git checkout release/v26.4
     git submodule update --init --recursive
     ```
 
@@ -113,12 +133,12 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party
     # NOTE: If your config downloads weights/tokenizer from Hugging Face Hub,
     #       you typically need to pass HF_TOKEN into the container.
     # Run in the Primus repository root directory
-    ./primus-cli container --image rocm/primus:v26.3 \
+    ./primus-cli container --image rocm/primus:v26.4 \
       --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
       -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
     ```
 
-For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS-CLI-GUIDE.md).
+For more detailed usage instructions, see the [CLI User Guide](./docs/02-user-guide/cli-reference.md).
 
 #### Option 2: wheel installation of Primus and run training in container
 
@@ -131,18 +151,18 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS
     python -m venv primus-env
     source primus-env/bin/activate
     # Install Primus
-    pip install "primus==26.3.1" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/
+    pip install "primus==26.4.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/
 
     ```
 
     >**Note**: this will only install the Primus CLI in your virtual environment under the `site-packages` directory, without other dependencies. The third party submodules will be downloaded on the first run. The complete dependencies and training software stack is provided in the AMD published training Docker images. You can use `primus-cli` to launch the training in container from any directory.
 
-    >**Note**: If you don't want to use docker container to run training, and want to install the complete dependencies and training software stack on your host machine, please refer to the instruction: [Install training environment on your host machine](docs/install-on-host.md). The automated installation script is under development and will be released soon.
+    >**Note**: If you don't want to use docker container to run training, and want to install the complete dependencies and training software stack on your host machine, please refer to the instruction: [Install training environment on your host machine](docs/01-getting-started/installation.md#bare-metal-host-setup). The automated installation script is under development and will be released soon.
 
 2. **Run training in container using pip-installed Primus**
 
     ```bash
-    primus-cli container --image rocm/primus:v26.3 \
+    primus-cli container --image rocm/primus:v26.4 \
     --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
     --volume /path/to/your/data:/data  -- --log_file /data/run.log \
     -- train pretrain --config /data/your/config.yaml
@@ -162,10 +182,10 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS
 
 Comprehensive documentation is available in the [`docs/`](./docs/) directory:
 
-- **[Quick Start Guide](./docs/quickstart.md)** - Get started in 5 minutes
-- **[Primus CLI User Guide](./docs/cli/PRIMUS-CLI-GUIDE.md)** - Complete CLI reference and usage
-- **[CLI Architecture](./docs/cli/CLI-ARCHITECTURE.md)** - Technical design and architecture
-- **[Backend Patch Notes](./docs/backends/overview.md)** - Primus-specific backend arguments
+- **[Quick Start Guide](./docs/01-getting-started/quickstart.md)** - Get started in 5 minutes
+- **[Primus CLI User Guide](./docs/02-user-guide/cli-reference.md)** - Complete CLI reference and usage
+- **[CLI Architecture](./docs/06-developer-guide/cli-architecture.md)** - Technical design and architecture
+- **[Backend Patch Notes](./docs/06-developer-guide/backend-patch-notes.md)** - Primus-specific backend arguments
 - **[Full Documentation Index](./docs/README.md)** - Browse all available documentation
 
 ---
diff --git a/benchmark/kernel/rccl/run_benchmark_b200.sh b/benchmark/kernel/rccl/run_benchmark_b200.sh
index f3a97d36b..996d2a549 100644
--- a/benchmark/kernel/rccl/run_benchmark_b200.sh
+++ b/benchmark/kernel/rccl/run_benchmark_b200.sh
@@ -27,8 +27,8 @@ set -euo pipefail
 
 GPUS_PER_NODE=8
 # Pre-imported sqsh avoids enroot aufs whiteout errors on ext4 /data filesystem.
-# The sqsh is pre-imported on: GPUA806, GPUA817, GPUA818, GPUA863, GPUA7DD, GPUA81E
-# If running on other nodes, root must run: enroot import --output /scratch/enroot/data/nvcr+nvidia+nemo+25.11.01.sqsh docker://nvcr.io#nvidia/nemo:25.11.01
+# It must exist on every node in the allocation; where it does not, root must run:
+#   enroot import --output /scratch/enroot/data/nvcr+nvidia+nemo+25.11.01.sqsh docker://nvcr.io#nvidia/nemo:25.11.01
 DOCKER_IMAGE="/scratch/enroot/data/nvcr+nvidia+nemo+25.11.01.sqsh"
 MASTER_PORT=29500
 SCRIPT_DIR="${SLURM_SUBMIT_DIR}"
diff --git a/benchmark/kernel/rccl/run_slurm.sh b/benchmark/kernel/rccl/run_slurm.sh
index 2a6710147..2e6b52078 100644
--- a/benchmark/kernel/rccl/run_slurm.sh
+++ b/benchmark/kernel/rccl/run_slurm.sh
@@ -10,7 +10,7 @@
 # Usage:
 #   DOCKER_IMAGE=<image> sbatch run_slurm.sh
 #   DOCKER_IMAGE=<image> NNODES=2 PARTITION=my-gpu sbatch run_slurm.sh
-#   DOCKER_IMAGE=rocm/primus:v26.3 NNODES=2 sbatch -N2 -w smci355-ccs-aus-n04-[25,29] -p Compute-DCPT ./run_slurm.sh
+#   DOCKER_IMAGE=rocm/primus:v26.4 NNODES=2 sbatch -N2 -w node[01-02] -p my-gpu ./run_slurm.sh
 #
 # Environment variables (all optional except DOCKER_IMAGE):
 #   DOCKER_IMAGE        Docker image to use (required)
@@ -89,7 +89,7 @@ NODE_RANK="${SLURM_NODEID}"
 
 # ---- Build short NODE_TAG like: n05-29_n05-33 ----
 short_node() {
-  # input: smci355-ccs-aus-n05-29  -> output: n05-29
+  # keep the last two dash-separated fields, e.g. <prefix>-n05-29 -> n05-29
   local h="$1"
   local n_part id_part
   n_part="$(echo "$h" | awk -F"-" "{print \$(NF-1)}")"  # n05
@@ -179,7 +179,7 @@ docker run --rm \
         MEGATRON_PATH=\"\${PRIMUS_ROOT_PATH}/third_party/Megatron-LM\"
         export PYTHONPATH=\"\${MEGATRON_PATH}:\${PYTHONPATH:-}\"
 
-        # Final CSV names you want (no smci prefix, include both nodes)
+        # Final CSV names (short node tag, no cluster prefix, includes every node)
         FINAL_ALLREDUCE=\"${OUTPUT_DIR}/allreduce_\${NODE_TAG}.csv\"
         FINAL_ALLGATHER=\"${OUTPUT_DIR}/allgather_\${NODE_TAG}.csv\"
         FINAL_REDUCESCATTER=\"${OUTPUT_DIR}/reducescatter_\${NODE_TAG}.csv\"
diff --git a/benchmark/kernel/rccl/submit_b200.sh b/benchmark/kernel/rccl/submit_b200.sh
deleted file mode 100644
index db2b67097..000000000
--- a/benchmark/kernel/rccl/submit_b200.sh
+++ /dev/null
@@ -1,2 +0,0 @@
-#!/bin/bash
-sbatch -w  GPUA7DD,GPUA817 run_benchmark_gb200.sh
diff --git a/benchmark/kernel/rccl/submit_gb200.sh b/benchmark/kernel/rccl/submit_gb200.sh
deleted file mode 100644
index 52e7bf92a..000000000
--- a/benchmark/kernel/rccl/submit_gb200.sh
+++ /dev/null
@@ -1,11 +0,0 @@
-#!/bin/bash
-# Slurm -w accepts hostlist patterns; quote so shellcheck does not treat brackets as globs.
-sbatch -w 'slurm-compute-node-[0-1]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[2-3]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[4-5]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[6-7]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[8-9]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[10-11]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[12-13]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[14-15]' run_bench_gb200.sh
-sbatch -w 'slurm-compute-node-[16-17]' run_bench_gb200.sh
diff --git a/benchmark/kernel/rccl/submit_pairs.sh b/benchmark/kernel/rccl/submit_pairs.sh
deleted file mode 100755
index ba64c93a1..000000000
--- a/benchmark/kernel/rccl/submit_pairs.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-DOCKER_IMAGE="rocm/primus:v26.3"
-#DOCKER_IMAGE="docker.gpuperf:5000/aai_2026_training/rocm/primus_megatron:v25.11_gpt_oss_sink"
-#DOCKER_IMAGE="docker.gpuperf:5000/gpuperf/primus:v26.1_sinkfa"
-NNODES=2
-PARTITION="Compute-DCPT"
-SCRIPT="./run_slurm.sh"
-
-NODELISTS=(
-  "smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-25"
-  "smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21"
-  "smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29"
-  "smci355-ccs-aus-n03-25,smci355-ccs-aus-n03-33"
-  "smci355-ccs-aus-n04-21,smci355-ccs-aus-n04-25"
-  "smci355-ccs-aus-n04-29,smci355-ccs-aus-n04-33"
-  "smci355-ccs-aus-n05-21,smci355-ccs-aus-n05-29"
-  "smci355-ccs-aus-n05-33,smci355-ccs-aus-n06-25"
-  "smci355-ccs-aus-n06-33,smci355-ccs-aus-n10-29"
-)
-
-echo "Submitting ${#NODELISTS[@]} jobs..."
-export DOCKER_IMAGE NNODES
-for nodelist in "${NODELISTS[@]}"; do
-  echo ">> $nodelist"
-  sbatch -N "${NNODES}" -w "$nodelist" -p "$PARTITION" "$SCRIPT"
-done
-
-echo "Done."
diff --git a/benchmark/megatron/model/benchmark_report.py b/benchmark/megatron/model/benchmark_report.py
index c918541a1..64cb7d7d0 100644
--- a/benchmark/megatron/model/benchmark_report.py
+++ b/benchmark/megatron/model/benchmark_report.py
@@ -13,11 +13,22 @@
 
 ANSI_ESCAPE_PATTERN = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
 ARG_PATTERN = re.compile(r"\]:\s+([a-zA-Z0-9_]+)\s+\.{3,}\s+(.*)$")
+# Field order matches the current Megatron training-log line, where the memory
+# segment is appended last (after tokens) rather than sitting between elapsed and
+# throughput as in the old format. Groups: elapsed inst/avg, tflops inst/avg,
+# tokens inst/avg, memory.
 ITERATION_PATTERN = re.compile(
     r"iteration\s+\d+/.*?elapsed time per iteration \(ms\): ([\d.]+)/([\d.]+).*?"
-    r"mem usages: ([\d.]+).*?"
-    r"throughput per GPU \(TFLOP/s/GPU\): ([\d.]+)/([\d.]+).*?"
-    r"tokens per GPU \(tokens/s/GPU\): ([\d.]+)/([\d.]+)",
+    # Compute (TFLOP/s/GPU) accepts both the current "compute per GPU (...): X (avg Y)"
+    # label and the legacy "throughput per GPU (...): X/Y" label. The space before
+    # "(avg" is optional (some log sinks drop it).
+    r"(?:compute|throughput) per GPU \(TFLOP/s/GPU\): ([\d.]+)(?:/|\s*\(avg\s*)([\d.]+)\)?.*?"
+    # Tokens accepts both the current "tokens/s/GPU inst/harmonic mean: X/Y" label
+    # and the legacy "tokens per GPU (tokens/s/GPU): X/Y" label.
+    r"(?:tokens per GPU \(tokens/s/GPU\)|tokens/s/GPU inst/harmonic mean): ([\d.]+)/([\d.]+).*?"
+    # Memory accepts both the current "hip mem usage/free/total/usage_ratio: X GB/..."
+    # segment and the legacy "mem usages: X" segment. Only the used value is captured.
+    r"(?:hip mem usage/free/total/usage_ratio:\s*|mem usages:\s*)([\d.]+)(?:\s*G(?:i)?B)?",
     re.DOTALL,
 )
 
@@ -49,9 +60,9 @@ def parse_last_metrics_from_log(file_path: str) -> Dict[str, float]:
 
     last = matches[-1]
     step_time_s = (float(last[0]) + float(last[1])) / 2000
-    mem_usage = float(last[2])
-    tflops = max(float(last[3]), float(last[4]))
-    tokens_per_gpu = (float(last[5]) + float(last[6])) / 2
+    tflops = max(float(last[2]), float(last[3]))
+    tokens_per_gpu = (float(last[4]) + float(last[5])) / 2
+    mem_usage = float(last[6])
 
     return {
         "TFLOP/s/GPU": round(tflops, 2),
diff --git a/docs/.gitignore b/docs/.gitignore
new file mode 100644
index 000000000..c6cf22b0f
--- /dev/null
+++ b/docs/.gitignore
@@ -0,0 +1,2 @@
+_build
+sphinx/_toc.yml
diff --git a/docs/.readthedocs.yaml b/docs/.readthedocs.yaml
new file mode 100644
index 000000000..75a0bb1c4
--- /dev/null
+++ b/docs/.readthedocs.yaml
@@ -0,0 +1,16 @@
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+build:
+  os: ubuntu-24.04
+  tools:
+    python: "3.12"
+
+sphinx:
+  configuration: docs/conf.py
+
+python:
+  install:
+    - requirements: docs/sphinx/requirements.txt
diff --git a/docs/01-getting-started/README.md b/docs/01-getting-started/README.md
new file mode 100644
index 000000000..559b9b52f
--- /dev/null
+++ b/docs/01-getting-started/README.md
@@ -0,0 +1,12 @@
+# Getting started
+
+Start here if you are new to Primus.
+
+- [Project overview](overview.md): what Primus does, who it is for, key capabilities
+- [Installation guide](installation.md): prerequisites, Docker/bare-metal/Slurm setup
+- [Quickstart](quickstart.md): first training run in 5 minutes
+- [Glossary](glossary.md): terms, acronyms, and domain concepts
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/01-getting-started/bare-metal-installation-jax.md b/docs/01-getting-started/bare-metal-installation-jax.md
new file mode 100644
index 000000000..c6cf91529
--- /dev/null
+++ b/docs/01-getting-started/bare-metal-installation-jax.md
@@ -0,0 +1,901 @@
+# Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)
+
+This guide explains how to build the **Primus JAX / MaxText training software
+stack directly on a host machine**, without using the AMD published JAX training
+Docker image. It is intended for users who, for policy or operational reasons,
+cannot run containers and need to reproduce the same environment on bare metal.
+
+It is derived from the official JAX training `Dockerfile` and
+installs the same components and versions. Wherever possible, everything is
+installed **inside a Python virtual environment and without `sudo`**. The only
+steps that require root are a small set of OS-level system libraries (installed
+with `apt`), and a couple of optional networking packages used for multi-node
+training.
+
+> **Looking for the PyTorch/Megatron/TorchTitan stack instead?** See
+> [bare-metal-installation.md](bare-metal-installation.md). This document is the
+> JAX **MaxText** counterpart. It is leaner than the PyTorch stack — no Flash
+> Attention / aiter / Primus-Turbo / FBGEMM / rocSHMEM builds — but v26.5 still
+> compiles **TensorFlow (CPU) and RCCL from source** (and TransformerEngine from
+> source on hosts with glibc < 2.38), so it is not build-free.
+
+> **Python 3.12+ required.** MaxText requires Python ≥ 3.12 (the reference image
+> is built on Ubuntu 24.04). Unlike the PyTorch recipe, Python 3.10 is **not**
+> sufficient here.
+
+> **⚠️ Host OS: Ubuntu 24.04 / glibc ≥ 2.38 strongly recommended.** The prebuilt
+> `transformer_engine_rocm_jax` wheel (and other JAX training wheels) are built
+> against the Dockerfile's `ubuntu:24.04` base — they need **`glibc ≥ 2.38`** and
+> **`libstdc++` with `GLIBCXX_3.4.32`** (GCC 13/14). On an older host such as
+> **Ubuntu 22.04 (glibc 2.35)** the TransformerEngine shared library fails to
+> load with `version 'GLIBC_2.38' not found` (and this failure is *silently*
+> swallowed by the launcher — training just exits right after JAX initializes
+> the GPUs). `libstdc++` can be side-loaded via `LD_LIBRARY_PATH`, but **glibc
+> cannot**. For a **manual** install on Ubuntu 22.04 you must either (a) run on a
+> 24.04 host, or (b) build TransformerEngine **from source** against the host
+> toolchain — see
+> [Section 3.7](#37-install-transformerengine-jax-from-the-prebuilt-rocm-wheel).
+> Check your host with `ldd --version`.
+>
+> **The automated `setup.sh` does (b) for you:** its `te` stage detects the host
+> glibc and, when it is < 2.38, transparently builds TransformerEngine from
+> source instead of installing the prebuilt wheel — so `bash setup.sh` works on
+> both Ubuntu 22.04 and 24.04. Only the prebuilt TE wheel needs glibc ≥ 2.38; the
+> ROCm JAX/PJRT wheels load fine on glibc 2.35.
+
+---
+
+## Quick path: automated install scripts
+
+If you just want the environment built for you, use the helper scripts in
+[tools/installation-jax/](https://github.com/AMD-AGI/Primus/tree/main/tools/installation-jax).
+They automate everything in Section 3 (Python venv) of this guide — venv
+creation, the ROCm release tarball, MaxText and its dependencies, TensorFlow
+(built from source), JAX + ROCm PJRT/plugin, TransformerEngine, RCCL (built from
+source), and Primus itself — and provide a single `env.sh` to activate the
+environment before each job. Read the rest of this document if you want to
+understand or customize what they do, or if you need the multi-node networking
+stack (Section 4), which the scripts do not build.
+
+There are two files:
+
+- **env.sh** — defines the install location and exports every environment
+  variable the build and runtime need (ROCm paths, `NVTE_*` flags, `XLA_FLAGS`,
+  `MAXTEXT_PATH`, cache locations, etc.). Source it both during the build and
+  every time you use the environment.
+- **setup.sh** — runs the install in re-runnable **stages**. It sources `env.sh`
+  automatically.
+
+### Choose where it installs (important)
+
+Everything lives under `PRIMUS_JAX_BASE` (venv, kept checkouts), with transient
+build sources on `SRC_DIR` (defaults to local `/tmp` for fast I/O).
+**`PRIMUS_JAX_BASE` is required — there is no default**, so set it to a directory
+you can write to that has tens of GB free (`env.sh` errors out if it is unset):
+
+```bash
+export PRIMUS_JAX_BASE=/path/to/big/disk/primus-jax-env   # venv + checkouts (persistent)
+export SRC_DIR=/tmp/primus-jax-build                      # transient sources (optional override)
+```
+
+The scripts **auto-detect your GPU architecture** (`env.sh` reads `rocminfo`, or
+falls back to the kernel KFD sysfs `gfx_target_version` when no ROCm is
+installed yet), and install the matching device wheels — `gfx942` (MI300X/MI325X),
+`gfx950` (MI350X/MI355X), or both. To force a target, export it before running:
+
+```bash
+export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+```
+
+> `PYTORCH_ROCM_ARCH` is the variable name the ROCm SDK and TransformerEngine
+> read to select gfx targets — it applies to the JAX build too, despite the name.
+
+**If your default `python3` is older than 3.12** (e.g. Ubuntu 22.04 ships 3.10),
+you do **not** need `sudo` or a PPA. The scripts use [`uv`](https://docs.astral.sh/uv/)
+to provide Python 3.12:
+
+- If `uv` is already installed, `env.sh` auto-detects a uv-managed
+  `>= 3.12` interpreter, and `setup.sh` runs `uv python install 3.12`
+  automatically when none is present yet. Just run `bash setup.sh`.
+- If `uv` is not installed, install it once (no root) and re-run:
+
+  ```bash
+  python3 -m pip install --user uv      # or: curl -LsSf https://astral.sh/uv/install.sh | sh
+  bash setup.sh                          # setup.sh will fetch Python 3.12 via uv
+  ```
+
+To force a specific interpreter instead, export it before running:
+
+```bash
+export PRIMUS_PYTHON=/path/to/python3.12
+```
+
+### Build the environment
+
+```bash
+cd tools/installation-jax
+
+bash setup.sh            # run all default stages, in order
+bash setup.sh --list     # list available stages
+bash setup.sh te         # re-run a single stage (e.g. reinstall TransformerEngine)
+bash setup.sh venv rocm jax  # run a subset of stages
+```
+
+Stages are idempotent and re-runnable, so if a step fails you can fix the cause
+and re-run just that stage. On failure the script stops immediately and prints
+which stage failed.
+
+Default stages (v26.5):
+
+```
+venv → rocm → maxtext → tf_source → jax → te → primus → jaxreqs → rccl → manifest
+```
+
+This mirrors the v26.5 image, which **builds TensorFlow (2.21 CPU) and RCCL from
+source** (`tf_source`, `rccl`). Those are heavy: the TF bazel build alone is
+~30–60 min. Lighter/alternative stages:
+
+- `tf_cpu_fix` — pip `tensorflow-cpu` instead of the `tf_source` bazel build.
+- `te_source` — force the from-source TransformerEngine build regardless of
+  glibc. You normally don't need to pass this: the default `te` stage
+  **auto-falls-back** to a from-source build on glibc < 2.38 hosts (e.g. Ubuntu
+  22.04), where the prebuilt wheel won't load (see the host-OS note above and
+  Section 3.7). Heavy (~30–60 min).
+
+```bash
+# Ubuntu 22.04 (glibc < 2.38): `bash setup.sh` already builds TE from source
+# automatically. If you also want to skip the heavy tf_source bazel build, swap
+# in the lighter tf_cpu_fix:
+bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest
+```
+
+### Use the environment for a training job
+
+```bash
+# Use the SAME PRIMUS_JAX_BASE you built with
+export PRIMUS_JAX_BASE=/path/to/big/disk/primus-jax-env
+source tools/installation-jax/env.sh    # activates the venv + sets ROCm/NVTE/XLA vars
+
+python -c "import jax; print('devices:', jax.devices())"
+
+# Primus is checked out under $WORKSPACE_DIR
+cd "$WORKSPACE_DIR/Primus"
+./primus-cli direct -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+```
+
+> Pick the config directory that matches your GPU: `examples/maxtext/configs/MI300X/`
+> for **gfx942** (MI300X/MI325X) and `examples/maxtext/configs/MI355X/` for
+> **gfx950** (MI350X/MI355X), e.g.
+> `--config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml`.
+
+### What the scripts do NOT do
+
+- **System (`apt`) packages** (Section 2): skipped — they need root. A C++
+  compiler (`g++`/`make`), `git`, and the build basics must already be present,
+  along with the small extra set MaxText expects (`numactl`, `curl`, `lsb-release`, …).
+- **Multi-node networking** (Section 4: UCX, OpenMPI, AINIC): not built.
+  Single-node training works without them; follow Section 4 manually if you need
+  distributed-over-RDMA.
+- **gcsfuse**: only needed to mount Google Cloud Storage buckets for data. Not
+  required for synthetic-data or local-data runs.
+
+The scripts also adapt a few host-specific details beyond the Dockerfile (the
+Python part of MaxText's `setup.sh` is run directly to skip its `apt`/interactive
+steps, and the ROCm/MaxText/Primus checkouts are collapsed onto a single
+`MAXTEXT_PATH`). See
+[tools/installation-jax/README.md](https://github.com/AMD-AGI/Primus/blob/main/tools/installation-jax/README.md)
+for the full rationale.
+
+---
+
+## 0. The key idea: ROCm comes from a tarball, not from a system install
+
+This build **does not require a system-wide ROCm installation**. In v26.5 ROCm
+is delivered as a **release tarball** (AMD "TheRock" multi-arch dist) that is
+extracted into a user-writable directory (`$ROCM_DIR`, default
+`$PRIMUS_JAX_BASE/rocm`) — no `/opt/rocm`, no root:
+
+- The tarball (`repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz`)
+  provides the full ROCm toolchain (HIP, hipBLASLt, compilers, headers,
+  libraries). `ROCM_PATH` points at the extraction dir.
+- `jax` + `jaxlib` (upstream) plus the ROCm `jax_rocm7_pjrt` and
+  `jax_rocm7_plugin` wheels (from `repo.amd.com/rocm/whl-multi-arch/`) provide
+  GPU-accelerated JAX built against that ROCm.
+
+> The pinned **release tarball** is stable and avoids any pip version-skew that
+> could break hipBLASLt GEMMs. RCCL is then rebuilt from source (see Section 3.11)
+> and dropped into this ROCm tree.
+
+This means almost the entire stack can be installed **without root** into a
+venv. The only host-level requirements from the system administrator are:
+
+- The **AMD GPU kernel driver (amdgpu / ROCm KMD)** must already be installed and
+  loaded (`/dev/kfd` and `/dev/dri` must exist, and the user must be in the
+  `video` and `render` groups).
+- A small set of **build/runtime system libraries** (see Section 2).
+
+---
+
+## 1. Required software stack for JAX MaxText training
+
+The complete environment is composed of the following layers:
+
+
+| Layer                   | Component                                                          | Source                    | Needs root?               |
+| ----------------------- | ----------------------------------------------------------------- | ------------------------- | ------------------------- |
+| Kernel / hardware       | AMD GPU driver (amdgpu KMD), GPU device access                    | OS / admin                | Yes (one-time, by admin)  |
+| OS libraries            | Build toolchain + runtime libs (`g++`, `git`, `numactl`, RDMA, …) | `apt`                     | Yes (one-time)            |
+| ROCm user-space         | TheRock ROCm dist (HIP, hipBLASLt, compilers, libs)              | **release tarball** ($ROCM_DIR) | No (user dir)       |
+| Deep learning framework | JAX (`jax`, `jaxlib`) + ROCm `jax_rocm7_pjrt` / `jax_rocm7_plugin`| pip (upstream + repo.amd.com) | No (venv)             |
+| Accelerated kernels     | TransformerEngine (JAX) — prebuilt ROCm wheel (or from source)   | pip (staging index) / build | No (venv)               |
+| Training framework      | MaxText (ROCm fork) + its Python deps                             | git + pip (`uv`)          | No (venv)                 |
+| Collectives fix         | `tensorflow-cpu` 2.21 (built from source; no bundled NCCL/LLVM)   | build from source (bazel) | No (venv)                 |
+| Collectives lib         | RCCL (rebuilt from source into the ROCm tree)                    | build from source         | No (user dir)             |
+| Multi-node comms        | UCX, OpenMPI, AMD AINIC (libionic)                                | build from source / apt   | Mostly no (AINIC needs root) |
+| Primus                  | Primus + `third_party/maxtext` submodule                         | git + pip                 | No (venv)                 |
+
+
+### 1.1 Version requirements (pins and host prerequisites)
+
+These are the exact versions the v26.5 reference `Dockerfile` pins. The install
+scripts use the same pins; change one and you may have to change the others.
+
+| Component                         | Pinned version / source                                                 | Notes |
+| --------------------------------- | ----------------------------------------------------------------------- | ----- |
+| ROCm (TheRock dist tarball)       | `therock-dist-linux-multiarch-7.14.0.tar.gz`                            | Multi-arch (gfx942 + gfx950). Extracted into `$ROCM_DIR`. |
+| JAX / jaxlib                      | `0.10.0`                                                                | Upstream PyPI. |
+| ROCm PJRT / plugin                | `jax_rocm7_pjrt` / `jax_rocm7_plugin` `0.10.0+rocm7.14.0`               | From `repo.amd.com/rocm/whl-multi-arch/`. |
+| TransformerEngine (JAX)           | `transformer_engine_rocm_jax 2.15.0.dev0+rocm7.15.0a20260707.72d01a0`  | Prebuilt wheel **needs glibc ≥ 2.38**; else build from source (`te_source`). |
+| TensorFlow (CPU, from source)     | ROCm `tensorflow-upstream` branch `upstream-v2.21.0`                    | Built with bazelisk `v1.25.0`. Needs host `clang-18`/`lld-18`. |
+| RCCL (from source)                | `rocm-systems` @ `9e5e4084a4b8e1e86551b0eb054725c62354a926`            | Installed into `$ROCM_PATH/lib`. Needs host `clang-18`/`lld-18`. |
+| MaxText (ROCm fork)               | `release/v26.5`                                                         | 2-value `initialize()`/`run()` API; Primus `main` supports it (fix #912). |
+| Primus                            | `main`                                                                  | Includes the MaxText `initialize()`/`run()` compatibility shim. |
+| scipy                             | `1.16`                                                                  | |
+| amdsmi                            | `7.0.2`                                                                 | pip, after the ROCm tarball. |
+| Build front-end                   | `cmake 3.31.6`, `ninja 1.11.1.3`, `wheel 0.45.1`, `packaging 25.0`, `setuptools 69.5.1` | Plus `uv` (used by MaxText's dep install). |
+| TE deps                           | `pybind11 3.0.4`, `importlib-metadata 8.7.1`, `pydantic 2.13.4`, `flax 0.12.2` | |
+
+**Host prerequisites** (independent of the pins above):
+
+| Requirement            | Minimum / recommended                          | Why |
+| ---------------------- | ---------------------------------------------- | --- |
+| **Python**             | **≥ 3.12** — **3.12 recommended/pinned**       | MaxText requires ≥ 3.12; the prebuilt TE/JAX wheels are `cp312`, so on a 3.13 venv you must build TE from source. `uv` provides 3.12 with no root. |
+| **glibc**              | **≥ 2.38** for the prebuilt TE wheel           | Ubuntu 24.04 = glibc 2.38+. On **Ubuntu 22.04 (glibc 2.35)** the prebuilt TE wheel won't load — the `te` stage auto-falls-back to a from-source build. Check with `ldd --version`. `glibc` cannot be side-loaded via `LD_LIBRARY_PATH`. |
+| **libstdc++ (GCC)**    | `GLIBCXX_3.4.32` (GCC 13/14) for prebuilt TE   | Can be side-loaded via `LD_LIBRARY_PATH` if glibc itself is new enough. |
+| **C/C++ toolchain**    | `g++` with C++17; `clang-18`/`lld-18` for the TF/RCCL source builds | Source builds (`tf_source`, `rccl`, `te_source`) need LLVM 18. |
+| **GPU arch**           | AMD Instinct **gfx942** (MI300/MI325) or **gfx950** (MI350/MI355) | The pinned ROCm tarball + JAX wheels target these. Other archs (e.g. gfx90a/MI250) need matching wheels not pinned here. |
+| **GPU driver (KMD)**   | amdgpu / ROCm kernel driver loaded             | `/dev/kfd` + `/dev/dri` present; user in `video`/`render` groups. |
+
+> **Validated:** the scripts' default flow (with `te_source` substituted for the
+> prebuilt `te`) has been run end-to-end for single-node MaxText pretraining on
+> **gfx942 / Ubuntu 22.04 (glibc 2.35) / Python 3.12**. The prebuilt-`te` flow is
+> the path for Ubuntu 24.04 hosts.
+
+For a **distributed (multi-node) JAX MaxText job** specifically, beyond JAX and
+ROCm you additionally need:
+
+- **RCCL** (AMD's collective library) — rebuilt from source into the ROCm tree
+  (see Section 3.11). This is what MaxText's collectives run over.
+- **AMD AINIC / RDMA stack** (`libibverbs`, `rdma-core`, `libionic`) — for
+  high-performance networking on AMD Pensando NICs.
+- Correct GPU/NIC device permissions and (often) hugepages / `ulimit -l unlimited`.
+- **UCX + OpenMPI** (Section 4) — **optional for MaxText**; carried over from the
+  reference image for MPI-launched / other JAX workloads. MaxText itself does not
+  use them (see the note in Section 4).
+
+> Unlike the PyTorch stack, the JAX MaxText image does **not** build rocSHMEM, and
+> MaxText does not launch via `mpirun`. JAX forms its process group through the
+> **JAX distributed coordinator** (`JAX_COORDINATOR_IP`, which Primus sets from
+> `MASTER_ADDR`) and runs collectives over **RCCL**.
+
+---
+
+## 2. System packages (require `sudo` / administrator, one-time)
+
+These are OS-level libraries needed to *build* the rest of the stack and to run
+MaxText / RDMA networking. They must be installed by someone with root, but this
+is a **one-time** action; everything afterward is done unprivileged in a venv.
+
+> If you genuinely cannot get root at all, these packages must already be present
+> on the host. The remainder of the guide (Sections 3+) then runs entirely
+> without root.
+
+### 2.1 Build toolchain and core libraries
+
+```bash
+sudo apt update
+sudo apt install -y \
+    gfortran git git-lfs ninja-build g++ pkg-config xxd patchelf \
+    automake libtool flex ccache \
+    python3-venv python3-dev python3-pip python-is-python3 \
+    libegl1-mesa-dev liblzma-dev libdw1 libdrm-dev \
+    wget unzip zip
+```
+
+> **Source builds (v26.5): LLVM 18 toolchain.** The `tf_source` (TensorFlow 2.21
+> bazel) and `rccl` source builds want a host `clang-18`/`lld-18`. The reference
+> image adds the LLVM apt repo and installs them:
+>
+> ```bash
+> echo 'deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main' | sudo tee /etc/apt/sources.list.d/llvm.list
+> wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
+> sudo apt update && sudo apt install -y clang-18 lld-18 llvm-18-dev llvm-18-tools
+> ```
+>
+> (Use `jammy` for Ubuntu 22.04, `noble` for 24.04.) You can skip this if you use
+> the lighter `tf_cpu_fix` stage instead of `tf_source`.
+
+> **Python 3.12+**: MaxText needs Python ≥ 3.12. Ubuntu 24.04 ships 3.12 by
+> default. On Ubuntu 22.04 (which ships 3.10) `apt install python3.12` fails
+> because jammy has no such package — do **not** rely on it. Two options:
+>
+> - **Recommended, no sudo — `uv`** (this is what the automated scripts use):
+>   ```bash
+>   python3 -m pip install --user uv    # or: curl -LsSf https://astral.sh/uv/install.sh | sh
+>   uv python install 3.12              # downloads a standalone CPython 3.12 (no root)
+>   export PRIMUS_PYTHON="$(uv python find '>=3.12')"
+>   ```
+>   The manual venv step below then uses `"$PRIMUS_PYTHON" -m venv ...`.
+> - **Alternative — deadsnakes PPA** (needs sudo, and the PPA must be reachable
+>   from your network):
+>   ```bash
+>   sudo add-apt-repository ppa:deadsnakes/ppa
+>   sudo apt update
+>   sudo apt install -y python3.12 python3.12-venv python3.12-dev
+>   export PRIMUS_PYTHON=python3.12
+>   ```
+
+### 2.2 Extra packages MaxText's setup expects
+
+MaxText's own `setup.sh` installs these via `apt`; on bare metal, pre-install
+them once:
+
+```bash
+sudo apt install -y \
+    numactl lsb-release gnupg curl net-tools iproute2 procps lsof ethtool
+```
+
+Optional — only if you read training data from Google Cloud Storage:
+
+```bash
+# gcsfuse (mount GCS buckets)
+export GCSFUSE_REPO=gcsfuse-$(lsb_release -c -s)
+echo "deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main" | \
+    sudo tee /etc/apt/sources.list.d/gcsfuse.list
+curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
+sudo apt update && sudo apt install -y gcsfuse
+```
+
+### 2.3 RDMA / networking libraries (needed for multi-node training)
+
+```bash
+sudo apt install -y \
+    rdma-core libibverbs-dev ibverbs-utils infiniband-diags \
+    ethtool kmod dpkg-dev jq xz-utils \
+    libevent-dev libhwloc-dev libmunge-dev \
+    software-properties-common
+```
+
+### 2.4 AMD AINIC library (optional, for AMD Pensando NICs)
+
+This pulls a vendor `.deb` from the AMD radeon repository. Skip it if you are
+not using AMD AINIC networking.
+
+```bash
+# Pin to the version used by the reference image
+AINIC_BUNDLE_VERSION="1.117.5-a-77"
+
+sudo add-apt-repository -y \
+  "deb https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_BUNDLE_VERSION} noble main"
+sudo apt update --allow-insecure-repositories
+sudo apt install -y --allow-unauthenticated libionic-dev
+```
+
+---
+
+## 3. Build the Python environment (no `sudo` from here on)
+
+Everything below runs as a regular user inside a virtual environment.
+
+### 3.1 Create and activate the virtual environment
+
+```bash
+# Pick a stable location, e.g. ~/primus-jax-env. MaxText needs Python >= 3.12.
+# On Ubuntu 22.04, get a 3.12 interpreter via uv first (see the Python 3.12+
+# note in Section 2.1): export PRIMUS_PYTHON="$(uv python find '>=3.12')"
+"${PRIMUS_PYTHON:-python3.12}" -m venv ~/primus-jax-env
+source ~/primus-jax-env/bin/activate
+
+# Build/runtime knobs (match the Dockerfile)
+export MAX_JOBS=128                              # lower this if you have fewer CPU cores / less RAM
+export PYTORCH_ROCM_ARCH="gfx942;gfx950"         # MI300/MI325 = gfx942, MI350/MI355 = gfx950
+export ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+```
+
+### 3.2 Bootstrap build tooling
+
+```bash
+pip install --upgrade pip
+pip uninstall -y wheel
+pip install \
+    cmake==3.31.6 \
+    ninja==1.11.1.3 \
+    wheel==0.45.1 \
+    packaging==25.0 \
+    setuptools==69.5.1 \
+    uv
+```
+
+### 3.3 Workaround environment variables
+
+```bash
+# Avoids HSA_STATUS_ERROR_OUT_OF_RESOURCES on some configurations
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+export HSA_NO_SCRATCH_RECLAIM=1
+
+# Fix the ROCm profiler hang issue
+export ROCPROFILER_QUEUE_INTERPOSITION=0
+export DEBUG_HIP_DYNAMIC_QUEUES=0
+```
+
+### 3.4 Install ROCm from the TheRock release tarball
+
+This step replaces a system ROCm install. v26.5 uses a **release tarball**
+(not pip wheels) extracted into a user-writable dir — no `/opt/rocm`, no root.
+
+```bash
+# Extract into a user-writable location (the automated env.sh uses
+# $PRIMUS_JAX_BASE/rocm). ROCM_PATH will point here.
+export ROCM_DIR=~/primus-jax-env/rocm
+mkdir -p "$ROCM_DIR"
+wget -O /tmp/therock-dist.tar.gz \
+    https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz
+tar -xzf /tmp/therock-dist.tar.gz -C "$ROCM_DIR"
+
+# amdsmi (installed via pip in the reference image)
+pip install amdsmi==7.0.2
+```
+
+> The tarball is multi-arch (gfx942 + gfx950), so there is no per-arch package to
+> pick, and there is **no pip version-skew to worry about** — a pinned tarball
+> cannot drift out of sync and break hipBLASLt GEMMs.
+
+### 3.5 Export ROCm paths
+
+Point the rest of the build/runtime at the extracted ROCm. **These must be set
+every time you use the environment** — Section 5 shows how to make them
+persistent (the automated `env.sh` does all of this for you).
+
+```bash
+export ROCM_PATH=$ROCM_DIR
+export ROCM_HOME=$ROCM_PATH
+export HIP_PLATFORM=amd
+export HIP_PATH=$ROCM_PATH
+export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
+export HIP_INCLUDE_PATH=$ROCM_PATH/include
+export HIP_LIB_PATH=$ROCM_PATH/lib
+export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
+# The reference Dockerfile puts $ROCM_PATH/lib on PATH too; mirror that.
+export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
+export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib"
+export LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64"
+export CPATH=$HIP_INCLUDE_PATH
+export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+```
+
+Quick check before continuing:
+
+```bash
+hipcc --version
+```
+
+### 3.6 Install MaxText, then JAX + the ROCm PJRT/plugin
+
+> **Order matters (v26.5).** The subsections below are numbered for reference,
+> but the correct install order is: **MaxText deps (§3.8) → TensorFlow from
+> source (§3.9) → ROCm JAX/PJRT/plugin (§3.6, right here) → TransformerEngine
+> (§3.7) → Primus (§3.10) → RCCL from source (§3.11)**. MaxText's `setup.sh`
+> pulls in a stock `jax`/`tensorflow`, so the ROCm JAX must be installed *after*
+> MaxText (to override it) and *before* TE (or a later step overwrites
+> `jaxlib`). The automated `setup.sh` enforces this ordering for you.
+
+MaxText install is in Section 3.8 (deps) below; the JAX packages are:
+
+```bash
+JAX_VERSION=0.10.0
+JAX_PJRT_VERSION=0.10.0+rocm7.14.0       # https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-pjrt/
+JAX_PLUGIN_VERSION=0.10.0+rocm7.14.0     # https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-plugin/
+
+pip install jax==${JAX_VERSION} jaxlib==${JAX_VERSION} scipy==1.16
+pip install \
+    --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
+    --pre jax_rocm7_pjrt==${JAX_PJRT_VERSION} \
+    --pre jax_rocm7_plugin==${JAX_PLUGIN_VERSION}
+```
+
+### 3.7 Install TransformerEngine (JAX) from the prebuilt ROCm wheel
+
+TransformerEngine is installed as a prebuilt ROCm JAX wheel. Check
+[the staging index](https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-jax/)
+for the current pin.
+
+```bash
+TE_VERSION=2.15.0.dev0+rocm7.15.0a20260707.72d01a0
+
+pip install \
+    pybind11==3.0.4 \
+    importlib-metadata==8.7.1 \
+    pydantic==2.13.4 \
+    flax==0.12.2
+
+pip install \
+    --index-url https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/ \
+    --pre \
+    --no-build-isolation \
+    transformer_engine_rocm_jax==${TE_VERSION}
+```
+
+> **The prebuilt wheel needs `glibc ≥ 2.38` (Ubuntu 24.04).** Verify it actually
+> loads before continuing:
+>
+> ```bash
+> python -c "import transformer_engine.jax; print('TE JAX OK')"
+> ```
+>
+> If you see `OSError: ... version 'GLIBC_2.38' not found` (typical on Ubuntu
+> 22.04, glibc 2.35), the wheel is incompatible with your host. **Build
+> TransformerEngine from source instead** so it links against your host's glibc.
+> This is exactly what the automated `te_source` stage does
+> (`bash setup.sh ... te_source ...`); the manual equivalent is:
+>
+> ```bash
+> pip uninstall -y transformer_engine transformer_engine_rocm_jax
+> export USE_ROCM=1 NVTE_FRAMEWORK=jax NVTE_USE_ROCM=1
+> export NVTE_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" CMAKE_BUILD_PARALLEL_LEVEL=${MAX_JOBS}
+> git clone --recursive https://github.com/ROCm/TransformerEngine.git
+> cd TransformerEngine
+> git checkout 635d7c085c39a6d9bfe4881c7d3efab7a46d7129   # last known-good ROCm JAX TE source commit
+> git submodule update --init --recursive
+> python3 setup.py bdist_wheel && pip install dist/*.whl
+> cd ..
+> ```
+>
+> If you only hit a `GLIBCXX_3.4.32` (libstdc++) error but glibc is new enough,
+> you can instead side-load a newer `libstdc++.so.6` (e.g. extracted from a
+> newer distro's `libstdc++6` package) via `LD_LIBRARY_PATH` — no rebuild needed.
+
+### 3.8 Install MaxText and its dependencies
+
+Clone the ROCm MaxText fork and install its Python dependencies. The reference
+image runs MaxText's `src/dependencies/scripts/setup.sh`; on bare metal we run
+the **Python portion** of that script directly (the `apt`/`gcsfuse` steps are the
+one-time root action from Section 2, and the venv already exists).
+
+> **MaxText `release/v26.5` and the Primus API.** MaxText v26.5 uses a
+> **2-value** `initialize()`/`run()` API (`config, recorder`). Primus `main`
+> handles it: `MaxTextPretrainTrainer` forwards `initialize()`'s tuple verbatim
+> to `run()` (fix #912), so v26.5 trains out of the box. Override `MAXTEXT_BRANCH`
+> only if you deliberately need to pin a different MaxText release.
+
+```bash
+cd ~/primus-jax-env   # or your $WORKSPACE_DIR
+git clone https://github.com/ROCm/maxtext.git
+cd maxtext
+git checkout release/v26.5   # matches the v26.5 image; Primus main supports its 2-value API
+
+# MaxText installs its deps with uv. The default (tpu) requirements set contains
+# the framework-agnostic Python deps WITHOUT any CUDA packages, which is what the
+# ROCm image uses.
+pip install -U setuptools wheel uv
+python -m uv pip install --resolution=lowest \
+    -r src/dependencies/requirements/generated_requirements/tpu-requirements.txt
+python -m src.dependencies.scripts.install_pre_train_extra_deps
+python -m uv pip install --no-deps -e .
+cd ..
+```
+
+> This pulls in the full `tensorflow` package (via `tensorflow-text`); the next
+> step swaps it for the CPU build. It may also nudge `jax`/`jaxlib`/`scipy`
+> within their allowed ranges — the ROCm PJRT/plugin installed in 3.6 remain in
+> place.
+
+### 3.9 Build TensorFlow (CPU) from source
+
+v26.5 rebuilds TensorFlow 2.21 (CPU) from ROCm's fork. The stock PyPI TF wheel
+bundles an LLVM whose symbols collide with ROCm's `libLLVM` in Grain "spawn"
+workers → **SIGSEGV on `import tensorflow` after `import jax`**. A CPU build has
+correct symbol visibility and no bundled NCCL (so it also preserves the XLA→RCCL
+collective fix). This is a **heavy bazel build (~30–60 min)** and needs a host
+`clang`/`lld` (LLVM 18) plus `unzip`/`zip` (Section 2).
+
+```bash
+# bazelisk (to a user-writable location) auto-picks the bazel version TF pins.
+wget -O ~/primus-jax-env/bin/bazel \
+    https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64
+chmod +x ~/primus-jax-env/bin/bazel
+export PATH="$HOME/primus-jax-env/bin:$PATH"
+
+git clone --depth 1 --branch upstream-v2.21.0 https://github.com/ROCm/tensorflow-upstream.git
+cd tensorflow-upstream
+bazel --output_user_root=/tmp/primus-jax-build/bazel build //tensorflow/tools/pip_package:wheel \
+    --repo_env=WHEEL_NAME=tensorflow_cpu \
+    --repo_env=HERMETIC_PYTHON_VERSION=3.12
+pip uninstall -y tensorflow tensorflow-cpu tensorflow_cpu
+pip install --no-deps bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow_cpu-2.21.0-cp312-cp312-linux_x86_64.whl
+cd ..
+```
+
+> **Lighter alternative:** if you don't want the bazel build, `pip install
+> --no-deps tensorflow-cpu==$(pip show tensorflow | awk '/^Version:/{print $2}')`
+> installs a CPU wheel that avoids the bundled-NCCL clash. It may still hit the
+> LLVM-symbol SIGSEGV on some ROCm 7.14 configs — the from-source build is the
+> robust fix. The automated recipe exposes this as the `tf_cpu_fix` stage.
+
+### 3.10 Install Primus
+
+```bash
+cd ~/primus-jax-env   # or your $WORKSPACE_DIR
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+git checkout main
+git submodule update --init third_party/maxtext/
+
+# The JAX path does not install Primus' torch-oriented requirements.txt.
+# Remove stale dataclasses backports that conflict on modern Python:
+pip uninstall -y dataclasses dataclasses_json
+
+# Primus' JAX runtime deps (also installed by the MaxText pre-train hook at
+# launch time):
+pip install -r requirements-jax.txt
+```
+
+If you already have a local Primus checkout (e.g. this repository), you can skip
+the clone and just run the `git submodule update`, `pip uninstall`, and
+`pip install -r requirements-jax.txt` steps from its root.
+
+> **Which MaxText does Primus run?** At launch, Primus resolves the MaxText
+> backend from the `MAXTEXT_PATH` environment variable, falling back to its own
+> `third_party/maxtext` submodule. Set `MAXTEXT_PATH` to the checkout you
+> installed dependencies into (Section 3.8) so the code and the installed deps
+> match — the automated `env.sh` does this for you.
+
+### 3.11 Build RCCL from source (into the ROCm tree)
+
+v26.5 rebuilds RCCL from `rocm-systems` and drops the libraries into the ROCm
+tree so JAX/XLA collectives use it. Requires the ROCm toolchain (`hipcc`) from
+Section 3.4.
+
+```bash
+git clone https://github.com/ROCm/rocm-systems.git
+cd rocm-systems
+git checkout 9e5e4084a4b8e1e86551b0eb054725c62354a926
+cd projects/rccl
+./install.sh -l --prefix build/ --amdgpu_targets="${PYTORCH_ROCM_ARCH}"
+cp -r build/release/librccl* "$ROCM_PATH/lib/"
+cd ../../..
+```
+
+---
+
+## 4. Multi-node communication stack (UCX, OpenMPI)
+
+These are only needed for **multi-node distributed** training. They build from
+source and install into user-writable prefixes (no root needed, except the AINIC
+`.deb` already handled in Section 2.4).
+
+> **Does JAX MaxText actually need UCX/OpenMPI?** For MaxText itself, **no** — JAX
+> forms its process group through the **JAX distributed coordinator**
+> (`JAX_COORDINATOR_IP`/`JAX_COORDINATOR_PORT`, which Primus sets from
+> `MASTER_ADDR`/`MASTER_PORT`) and runs collectives over **RCCL**; there is no
+> `mpirun` launch and no rocSHMEM. UCX/OpenMPI are carried over from the shared
+> reference image (used by MPI-launched / other JAX-based workloads) and are
+> installed here only for parity. **You can skip Section 4 entirely for
+> single- and multi-node MaxText pretraining.**
+
+> **Multi-node: make all local GPUs visible to each process.** On each node,
+> every rank/process must see all local GPUs, otherwise JAX enumerates only a
+> single device per node. Export `CUDA_VISIBLE_DEVICES` covering every local GPU
+> (the ROCm PJRT plugin honors the CUDA-named variable) before launching:
+>
+> ```bash
+> export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((GPUS_PER_NODE - 1)))   # e.g. 0,1,2,3,4,5,6,7
+> ```
+>
+> Add it to your activation script (Section 5) or your job launcher. It is
+> intentionally **not** hard-coded in `env.sh`, since the right value depends on
+> how ranks are pinned to GPUs on your host/scheduler.
+
+### 4.1 UCX
+
+```bash
+cd ~/primus-jax-env
+UCX_VERSION="1.18.0"
+wget https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz
+mkdir -p ucx-${UCX_VERSION}
+tar -zxf ucx-${UCX_VERSION}.tar.gz -C ucx-${UCX_VERSION} --strip-components=1
+cd ucx-${UCX_VERSION}
+mkdir build && cd build
+../configure --prefix=$HOME/primus-jax-env/ucx-${UCX_VERSION}/install --with-rocm=${ROCM_PATH}
+make -j 16 && make install
+cd ../..
+
+export UCX_INSTALL_DIR=$HOME/primus-jax-env/ucx-${UCX_VERSION}/install
+```
+
+### 4.2 OpenMPI
+
+```bash
+MPI_VERSION="4.1.6"
+wget https://download.open-mpi.org/release/open-mpi/v$(echo "${MPI_VERSION}" | cut -d. -f1-2)/openmpi-${MPI_VERSION}.tar.gz
+mkdir -p ompi-${MPI_VERSION}
+tar -zxf openmpi-${MPI_VERSION}.tar.gz -C ompi-${MPI_VERSION} --strip-components=1
+cd ompi-${MPI_VERSION}
+mkdir build && cd build
+# Install to a user-writable prefix instead of /opt to avoid sudo
+../configure --prefix=$HOME/primus-jax-env/openmpi --with-ucx=${UCX_INSTALL_DIR} \
+    --disable-oshmem --disable-mpi-fortran
+make -j 16 && make install
+cd ../..
+
+export PATH="$HOME/primus-jax-env/openmpi/bin:${PATH}"
+export LD_LIBRARY_PATH="$HOME/primus-jax-env/openmpi/lib:${LD_LIBRARY_PATH}"
+```
+
+> The Dockerfile installs OpenMPI under `/workspace`. The `$HOME/...` prefixes
+> above keep it unprivileged.
+
+---
+
+## 5. Make the environment reproducible (activation script)
+
+Many of the variables above (especially the ROCm paths and the `NVTE_*` / `XLA_*`
+runtime flags) must be present in **every** shell that runs training. Append them
+to the venv's activation script so they're set whenever you
+`source ~/primus-jax-env/bin/activate`:
+
+```bash
+cat >> ~/primus-jax-env/bin/activate <<'EOF'
+
+# ---- Primus JAX host environment ----
+export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+export ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+export HSA_NO_SCRATCH_RECLAIM=1
+export ROCPROFILER_QUEUE_INTERPOSITION=0
+export DEBUG_HIP_DYNAMIC_QUEUES=0
+
+# v26.5: ROCm lives in the extracted tarball dir (Section 3.4), NOT a pip wheel.
+export ROCM_PATH=$HOME/primus-jax-env/rocm
+export ROCM_HOME=$ROCM_PATH
+export HIP_PLATFORM=amd
+export HIP_PATH=$ROCM_PATH
+export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
+export HIP_INCLUDE_PATH=$ROCM_PATH/include
+export HIP_LIB_PATH=$ROCM_PATH/lib
+export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
+export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$HOME/primus-jax-env/openmpi/bin:$PATH"
+export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$HOME/primus-jax-env/openmpi/lib"
+export LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib64"
+export CPATH=$HIP_INCLUDE_PATH
+export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+
+# Point Primus at the MaxText checkout whose deps we installed
+export MAXTEXT_PATH=$HOME/primus-jax-env/maxtext
+
+# TransformerEngine (ROCm) runtime flags for JAX
+export NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH"
+export NVTE_USE_ROCM=1
+export NVTE_USE_HIPBLASLT=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
+export NVTE_FUSED_ATTN=1
+export NVTE_CK_USES_BWD_V3=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+export NVTE_CK_HOW_V3_BF16_CVT=2
+
+# AMD GPU runtime knobs
+export GPU_MAX_HW_QUEUES=2
+export HIP_FORCE_DEV_KERNARG=1
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_DEBUG=VERSION
+# Bare-metal only: force RCCL to use its built-in ROCm IB/RoCE transport. On a
+# host, /usr/local/lib/librccl-net.so is on the default loader path and is
+# ABI-incompatible with the from-source RCCL (undefined symbol
+# ncclNetPlugin_v11/_v10 -> falls back to v9 -> segfault at clique init). The
+# v26.5 image ships no librccl-net.so, so this matches it.
+export NCCL_NET_PLUGIN=none
+
+# XLA / JAX runtime settings (v26.5 uses .9)
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.9
+export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_enable_command_buffer=''"
+# ---- end Primus JAX host environment ----
+EOF
+```
+
+> The automated `tools/installation-jax/env.sh` sets all of the above (and
+> auto-detects the GPU arch); prefer sourcing it over hand-editing `activate`.
+
+---
+
+## 6. Verify the installation
+
+```bash
+source ~/primus-jax-env/bin/activate   # or: source tools/installation-jax/env.sh
+
+# GPUs visible to ROCm?
+rocm-smi || ls -l /dev/kfd /dev/dri
+
+# JAX sees the GPUs?
+python -c "import jax; print('jax', jax.__version__); \
+print('backend:', jax.default_backend()); \
+print('devices:', jax.devices())"
+
+# Key libraries import cleanly? (import transformer_engine.jax — this actually
+# loads TE's shared lib, which is what fails on glibc < 2.38 with the prebuilt wheel)
+python -c "import jax, jaxlib, flax, transformer_engine.jax; print('JAX/flax/TE OK')"
+
+# Run a Primus MaxText training directly (no container). Use the config dir that
+# matches your GPU: MI300X/ for gfx942 (MI300X/MI325X), MI355X/ for gfx950 (MI350X/MI355X).
+cd ~/primus-jax-env/Primus   # or your Primus checkout
+./primus-cli direct -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+  # gfx950: --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml
+```
+
+`jax.default_backend()` should report `gpu` (ROCm), and `jax.devices()` should
+list your AMD GPUs. Use `primus-cli direct` (not `container`) since you are
+running on bare metal with everything installed in your environment.
+
+---
+
+## 7. Other important considerations
+
+- **Python version**: MaxText requires Python ≥ 3.12. If your venv is older, the
+  build will fail; get a 3.12 interpreter with `uv` (no sudo — see Section 2.1)
+  and recreate it. The automated `setup.sh` does this for you (it provisions
+  Python 3.12 via `uv` and recreates a too-old venv automatically).
+- **GPU device access without root**: the user running training must be able to
+  read/write `/dev/kfd` and `/dev/dri/*` — usually via membership in the `video`
+  and `render` groups (`sudo usermod -aG video,render $USER`, then re-login).
+- **Hugging Face access**: for gated models/tokenizers, export your token
+  (`export HF_TOKEN=hf_xxx` and/or `huggingface-cli login`).
+- **Install order is load-bearing (v26.5)**: MaxText → TensorFlow (from source)
+  → ROCm JAX/PJRT/plugin → TransformerEngine → RCCL (from source). MaxText's
+  `setup.sh` pulls in a stock `jax`/`tensorflow`, so the ROCm JAX must be
+  installed *after* MaxText (to override it) and *before* TE (or `jaxlib` gets
+  clobbered). The automated `setup.sh` stage order enforces this.
+- **RDMA / multi-node limits**: high-performance networking typically requires
+  `ulimit -l unlimited` and possibly hugepages, configured in
+  `/etc/security/limits.conf` (admin help). Verify NICs with `ibv_devinfo` /
+  `ibstat`. JAX uses the distributed coordinator (`JAX_COORDINATOR_IP` /
+  `JAX_COORDINATOR_PORT`), which Primus sets from `MASTER_ADDR` / `MASTER_PORT`.
+- **Version drift**: the ROCm release tarball, the JAX/PJRT/plugin versions, the
+  TransformerEngine wheel, the TensorFlow/RCCL source revisions, and the MaxText
+  branch are all pinned to one release (see the table in Section 1.1). If you
+  change one, you may need to update the others. The Docker image is the
+  authoritative, tested combination — match its `Dockerfile` ARGs when in doubt.
+- **Automated scripts**: the manual steps in Section 3 are automated by
+  [tools/installation-jax/](https://github.com/AMD-AGI/Primus/tree/main/tools/installation-jax)
+  (see *Quick path* above). The multi-node networking stack in Section 4 is still
+  manual.
+
+---
+
+## 8. Quick reference: minimal vs. full install
+
+If you only need **single-node MaxText pretraining**, you can skip the
+multi-node components:
+
+
+| Component                                              | Needed for                          |
+| ------------------------------------------------------ | ----------------------------------- |
+| ROCm (tarball), JAX + PJRT/plugin, TransformerEngine (JAX) | Core MaxText training (install these) |
+| MaxText + its deps, TensorFlow (from source), RCCL (from source) | Core MaxText training (install these) |
+| Primus + `third_party/maxtext`                         | Running MaxText via Primus          |
+| gcsfuse                                                | Reading data from GCS buckets only  |
+| UCX, OpenMPI, AINIC                                     | Multi-node distributed training     |
+
+
+Install the core rows first, validate with Section 6, then add the optional
+components as your workload requires.
diff --git a/docs/01-getting-started/bare-metal-installation.md b/docs/01-getting-started/bare-metal-installation.md
new file mode 100644
index 000000000..725160c28
--- /dev/null
+++ b/docs/01-getting-started/bare-metal-installation.md
@@ -0,0 +1,620 @@
+# Bare-metal installation: build the Primus training stack from source (no Docker)
+
+This guide explains how to build the **full Primus training software stack directly on a host machine**, without using the AMD published training Docker image. It is intended for users who, for policy or operational reasons, cannot run containers and need to reproduce the same environment on bare metal.
+
+It is derived from the official training Dockerfile — currently
+[`Dockerfile.primus-v26.5`](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/docker-release/Dockerfile.primus-v26.5) — and installs the same components and versions. Wherever possible, everything is installed **inside a Python virtual environment and without `sudo`**. The only steps that require root are a small set of OS-level system libraries (installed with `apt`), and a couple of optional networking packages used for multi-node training.
+
+> **Important**: This is a long, build-heavy process. A full from-source build (Flash Attention, aiter, Primus-Turbo, and on older hosts TransformerEngine) can take **several hours** and needs a machine with many CPU cores, plenty of RAM, and tens of GB of free disk. The Docker image remains the recommended and best-supported path. Use this guide only when containers are not an option.
+
+> **Python 3.12 is required.** The pinned torch nightly publishes a **cp312 Linux wheel and nothing else**, so a Python 3.10 environment (the default on Ubuntu 22.04) cannot install it. The automated scripts provision a private CPython 3.12 for you without root; if you build manually you must supply one yourself. See [Section 4.1](#41-create-the-virtual-environment-python-312).
+
+> **Using the JAX MaxText backend instead?** This guide builds the PyTorch stack (Megatron-LM / TorchTitan). For the JAX / MaxText backend, follow the leaner [JAX bare-metal installation guide](bare-metal-installation-jax.md) instead.
+
+---
+
+## 0. The key idea: ROCm comes from pip, not from a system install
+
+The most important thing to understand is that this stack **does not require a system-wide ROCm installation**. ROCm is delivered as Python wheels (AMD "TheRock" multi-arch nightly wheels):
+
+- `rocm-sdk-devel`, `rocm-sdk-device-gfx942`, `rocm-sdk-device-gfx950` provide the ROCm toolchain (HIP, hipBLASLt, compilers, headers, libraries) **inside the virtual environment**.
+- `torch`, `torchvision`, `torchaudio` and the `amd-*-device-gfx*` packages provide a ROCm-enabled PyTorch built against those wheels.
+
+This means almost the entire stack can be installed **without root** into a venv. The only host-level requirements from the system administrator are:
+
+- The **AMD GPU kernel driver (amdgpu / ROCm KMD)** must already be installed and loaded on the host (`/dev/kfd` and `/dev/dri` must exist, and the user must have permission to access them — typically by being in the `video` and `render` groups).
+- A small set of **build/runtime system libraries** (see [Section 3](#3-system-packages-need-root-one-time)).
+
+You do **not** need to install the full ROCm user-space stack system-wide.
+
+---
+
+## 1. Recommended path: the automated scripts
+
+The helper scripts in [tools/installation/](https://github.com/AMD-AGI/Primus/tree/main/tools/installation) build the entire single-node environment for you: the Python 3.12 interpreter, the venv, the ROCm/PyTorch wheels, every source-built kernel library, and Primus itself. They are the maintained path and are kept in sync with the reference Dockerfile — prefer them over the manual reference in [Section 4](#4-manual-build-reference).
+
+There are two files:
+
+- **`env.sh`** — defines the install location and exports every environment variable the build and runtime need (ROCm paths, `NVTE_*` flags, cache locations). Source it both during the build and every time you use the environment.
+- **`setup.sh`** — runs the install in re-runnable **stages**. It sources `env.sh` automatically.
+
+### 1.1 Before you start
+
+- System packages from [Section 3](#3-system-packages-need-root-one-time) must already be present (a C++ compiler, `git`, `make` and the build basics). These need root, so the scripts do not install them.
+- The GPU driver must be loaded and your user must be able to access `/dev/kfd` and `/dev/dri/*`.
+- Python 3.12 is handled for you: if no suitable interpreter is found, `setup.sh` fetches a standalone CPython 3.12 with [`uv`](https://docs.astral.sh/uv/). No root, no `apt`.
+- **`uv` itself is optional to pre-install.** If it is missing, `setup.sh` downloads it into `$PRIMUS_BASE/bin` from `astral.sh`. Install it yourself if that download is blocked, or if you prefer not to pipe a script into a shell — see [Installing `uv`](#installing-uv).
+
+### 1.2 Choose where it installs (required)
+
+`PRIMUS_BASE` **has no default and must be exported.** The right location is site-specific, so the scripts refuse to guess and stop with an error if it is unset. Point it at a directory you can write to with tens of GB free; it holds the venv, the provisioned interpreter, and the kept checkouts.
+
+```bash
+export PRIMUS_BASE="$HOME/envs/primus-env"   # venv + interpreter + checkouts (persistent)
+export SRC_DIR=/tmp/primus-build             # transient build sources (optional override)
+```
+
+The scripts detect your GPU architecture automatically and build only for it — `gfx942` (MI300X/MI325X) or `gfx950` (MI350X/MI355X) — which keeps the build as short as possible. Override it with `export PYTORCH_ROCM_ARCH="gfx942;gfx950"` only if you need a different target, such as one environment shared across both.
+
+### 1.3 Build the environment
+
+```bash
+cd tools/installation
+
+bash setup.sh            # run all default stages, in order
+bash setup.sh --list     # list available stages (works without PRIMUS_BASE)
+bash setup.sh te         # re-run a single stage (e.g. reinstall TransformerEngine)
+bash setup.sh venv torch # run a subset of stages
+```
+
+Expect hours rather than minutes. Running it detached avoids losing the build to a dropped connection:
+
+```bash
+nohup bash setup.sh > ~/primus-setup.log 2>&1 &
+tail -f ~/primus-setup.log
+```
+
+Stages are idempotent, so if a build fails you can fix the cause and re-run just that stage — exporting the same `PRIMUS_BASE` again, since that is how it finds the venv. On failure the script stops immediately and prints which stage failed.
+
+Default stages (single-node training path):
+
+```
+venv → torch → flash_attn → te → torchtune → torchao → pydeps
+     → grouped_gemm → causal_conv1d → mamba → primus → aiter → turbo
+     → boto → cleanup → manifest
+```
+
+Optional: `torchrec` (DLRM / recommendation stack).
+
+Order matters if you cherry-pick: `te` needs `torch` and `flash_attn` to have run first, because the TransformerEngine package index serves only TE packages and cannot supply anything else.
+
+### 1.4 Use the environment for a training job
+
+```bash
+# Use the SAME PRIMUS_BASE you built with
+export PRIMUS_BASE="$HOME/envs/primus-env"
+source tools/installation/env.sh    # activates the venv + sets all ROCm/NVTE vars
+
+python -c "import torch; print('gpu:', torch.cuda.is_available())"
+
+# Primus is checked out under $WORKSPACE_DIR
+cd "$WORKSPACE_DIR/Primus"
+./primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+Use `primus-cli direct` (not `container`), since you are running on bare metal with everything installed in your environment.
+
+### 1.5 What the scripts do NOT do
+
+- **System (`apt`) packages** ([Section 3](#3-system-packages-need-root-one-time)): skipped — they need root.
+- **Multi-node networking** ([Section 5](#5-multi-node-communication-stack-ucx-and-openmpi)): UCX, OpenMPI and AMD AINIC are not built. Single-node training works without them.
+- **FBGEMM / Flux / DLRM benchmarks**: not built (`torchrec` is an optional stage; FBGEMM additionally needs apt `libtbb-dev`).
+- **MLPerf `primus_mllog`**: v26.5 installs it from `training_results_v6.0`; it is not part of the core training path and is not installed.
+
+### 1.6 Where the scripts deliberately differ from the Dockerfile
+
+A handful of places diverge on purpose, because copying the Dockerfile exactly produces an environment that does not work outside the container — several of its floating (unpinned) dependencies have since drifted to versions that break. The short list:
+
+- **TransformerEngine** is installed from the wheel index only where glibc is new enough, otherwise built from the equivalent source commit ([Section 4.6](#46-transformerengine)).
+- **Companion wheels are pinned** (`torchaudio`, `torchvision`, `apex`) to the same nightly as `torch`, because the Dockerfile's floating versions no longer resolve.
+- **`nvidia-cutlass-dsl` is pinned to 4.5.3** and **`flydsl` to 0.2.4**, both to versions that the Dockerfile itself resolved to when it was built. Newer releases silently break `import mamba_ssm` and aiter's CK/HIP kernels respectively.
+- **`NVTE_CK_IS_V3_ATOMIC_FP32=1` on gfx942** instead of the Dockerfile's `0`: without fp32 atomics the CK v3 backward attention kernel produces `Inf` gradients on MI300X/MI325X at the first training step. This matches the MI300X/MI325X block in [training recipes](../02-user-guide/training-recipes.md); gfx950 keeps the Dockerfile value.
+- **pip is constrained** so no later install can replace the ROCm `torch`/`triton` with upstream CUDA builds.
+
+Each one is documented with the exact failure it avoids in
+[tools/installation/README.md](https://github.com/AMD-AGI/Primus/blob/main/tools/installation/README.md). Read that before "correcting" any of them back.
+
+---
+
+## 2. Required software stack for distributed LLM training
+
+The complete environment is composed of the following layers:
+
+| Layer                   | Component                                                                                   | Source                  | Needs root?                      |
+| ----------------------- | ------------------------------------------------------------------------------------------- | ----------------------- | -------------------------------- |
+| Kernel / hardware       | AMD GPU driver (amdgpu KMD), GPU device access                                              | OS / admin              | Yes (one-time, by admin)         |
+| OS libraries            | Build toolchain + runtime libs (`g++`, `git`, RDMA, hwloc, etc.)                            | `apt`                   | Yes (one-time)                   |
+| Python                  | CPython 3.12 (required by the pinned wheels)                                                | `uv` / standalone build | No                               |
+| ROCm user-space         | `rocm-sdk-devel` + device packages                                                          | pip (TheRock wheels)    | No (venv)                        |
+| Deep learning framework | PyTorch ROCm (`torch`, `torchvision`, `torchaudio`, `apex`)                                 | pip (TheRock wheels)    | No (venv)                        |
+| Accelerated kernels     | Flash Attention, TransformerEngine, aiter, grouped_gemm, Primus-Turbo, causal-conv1d, mamba | pip / build from source | No (venv)                        |
+| Training frameworks     | torchtune, torchao, torchrec, FBGEMM                                                        | build from source / pip | No (venv)                        |
+| Multi-node comms        | UCX, OpenMPI, AMD AINIC (libionic)                                                          | build from source / apt | Mostly no (AINIC libs need root) |
+| Primus                  | Primus + submodules (Megatron-LM, TorchTitan, etc.)                                         | git + pip               | No (venv)                        |
+| Python deps             | datasets, transformers, accelerate, trl, wandb, etc.                                        | pip                     | No (venv)                        |
+
+For a **distributed (multi-node) training job** specifically, beyond PyTorch and ROCm you additionally need:
+
+- **RCCL** (AMD's collective library) — provided by the ROCm SDK wheels.
+- **UCX + OpenMPI** — point-to-point transport and the MPI launcher.
+- **AMD AINIC / RDMA stack** (`libibverbs`, `rdma-core`, `libionic`) — for high-performance networking on AMD Pensando NICs.
+- Correct GPU/NIC device permissions and (often) hugepages / `ulimit -l unlimited` for RDMA.
+
+---
+
+## 3. System packages (need root, one-time)
+
+These are OS-level libraries needed to *build* the rest of the stack and to run RDMA networking. They must be installed by someone with root, but this is a **one-time** action; everything afterward is done unprivileged in a venv. On a shared/managed host, ask your administrator to install them once.
+
+> If you genuinely cannot get root at all, these packages must already be present on the host. There is no supported way to install system `.deb` packages without root. The remainder of the guide then runs entirely without root.
+
+### 3.1 Build toolchain and core libraries
+
+```bash
+sudo apt update
+sudo apt install -y \
+    gfortran git git-lfs ninja-build g++ pkg-config xxd patchelf \
+    automake libtool autoconf flex ccache \
+    python3-venv python3-dev python3-pip python-is-python3 \
+    libegl1-mesa-dev liblzma-dev libdw1 libdrm-dev libz3-dev \
+    wget xz-utils ffmpeg numactl pciutils
+```
+
+`libtbb-dev` is additionally required if you intend to build FBGEMM.
+
+### 3.2 RDMA / networking libraries (needed for multi-node training)
+
+```bash
+sudo apt install -y \
+    rdma-core libibverbs-dev ibverbs-utils infiniband-diags \
+    ethtool kmod dpkg-dev jq \
+    libevent-dev libhwloc-dev libmunge-dev \
+    software-properties-common
+```
+
+### 3.3 AMD AINIC library (optional, for AMD Pensando NICs)
+
+This pulls a vendor `.deb` from the AMD radeon repository. Skip it if you are not using AMD AINIC networking.
+
+```bash
+# Pin to the version used by the reference image
+AINIC_BUNDLE_VERSION="1.117.5-a-77"
+
+sudo add-apt-repository -y \
+  "deb https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_BUNDLE_VERSION} noble main"
+sudo apt update --allow-insecure-repositories
+sudo apt install -y --allow-unauthenticated libionic-dev
+```
+
+---
+
+## 4. Manual build reference
+
+This section is **reference material**, not a maintained walkthrough: it records the exact versions and the non-obvious build steps so you can audit or adapt the process. For an actual install, use the scripts in [Section 1](#1-recommended-path-the-automated-scripts) — they encode everything below plus the workarounds listed in [Section 1.6](#16-where-the-scripts-deliberately-differ-from-the-dockerfile).
+
+The authoritative source of versions is always the reference Dockerfile.
+
+### 4.1 Create the virtual environment (Python 3.12)
+
+The pinned torch nightly ships a cp312 Linux wheel only, so the interpreter must be 3.12. Ubuntu 22.04 has 3.10, and installing 3.12 with `apt` needs root — the no-root option is a standalone build, which is what `uv` provides.
+
+#### Installing `uv`
+
+Either method works and neither needs root. `setup.sh` performs the second one automatically when `uv` is missing, so you only need to do this by hand for a manual build, or if the `astral.sh` download is blocked on your network.
+
+```bash
+# A. via pip, using whatever Python you already have.
+#    --user puts it in ~/.local/bin, one of the locations setup.sh searches.
+python3 -m pip install --user uv
+
+# B. via the standalone installer (what setup.sh uses)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+Make sure the result is on your `PATH` (`uv --version` should work). If you cannot install `uv` at all, supply your own 3.12 interpreter instead and skip it entirely — the scripts accept `PRIMUS_PYTHON=/path/to/python3.12`, and for a manual build just point `PY312` at it below.
+
+#### Create the venv
+
+```bash
+uv python install 3.12
+PY312="$(uv python find 3.12)"
+
+"$PY312" -m venv ~/primus-env
+source ~/primus-env/bin/activate
+python --version          # must report 3.12.x
+
+# Build/runtime knobs (match the Dockerfile)
+export MAX_JOBS=128                              # lower if you have fewer cores / less RAM
+export PYTORCH_ROCM_ARCH="gfx942;gfx950"         # MI300/MI325 = gfx942, MI350/MI355 = gfx950
+export ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0        # avoids HSA_STATUS_ERROR_OUT_OF_RESOURCES
+export HSA_NO_SCRATCH_RECLAIM=1
+```
+
+> Set `PYTORCH_ROCM_ARCH` to only the architecture(s) you actually have to speed up source builds (e.g. `"gfx942"` for MI300X-only).
+
+### 4.2 Bootstrap build tooling
+
+```bash
+pip install --upgrade pip
+pip install \
+    pybind11 typeguard \
+    wheel==0.45.1 cmake==3.31.6 ninja==1.11.1.3 \
+    packaging==25.0 setuptools==75.1.0
+```
+
+### 4.3 Install ROCm + PyTorch from the TheRock multi-arch wheels
+
+This replaces a system ROCm install. Install the base deps first — `apex` declares `cxxfilt`, `pytest` and `ninja` requirements that the nightly index does not serve, so they must be present before the torch resolve.
+
+```bash
+pip install \
+    cxxfilt==0.3.0 tqdm==4.67.3 pyyaml==6.0.3 pytest==9.0.3 \
+    matplotlib==3.10.9 pandas==2.3.3 py-cpuinfo==9.0.0 build==1.5.0
+
+# One coherent nightly. The Dockerfile pins only `torch` and lets the companions
+# float, which no longer resolves; pin them all to the same date.
+NIGHTLY="rocm7.15.0a20260720"
+
+python -m pip uninstall -y torch
+python -m pip install \
+    --index-url https://rocm.nightlies.amd.com/whl-multi-arch --pre \
+    torch==2.12.0+${NIGHTLY} \
+    amd-torch-device-gfx942==2.12.0+${NIGHTLY} \
+    amd-torch-device-gfx950==2.12.0+${NIGHTLY} \
+    rocm-sdk-devel==7.15.0a20260720 \
+    rocm-sdk-device-gfx942==7.15.0a20260720 \
+    rocm-sdk-device-gfx950==7.15.0a20260720 \
+    torchaudio==2.11.0+${NIGHTLY} \
+    torchvision==0.27.0+${NIGHTLY} \
+    amd-torchvision-device-gfx942==0.27.0+${NIGHTLY} \
+    amd-torchvision-device-gfx950==0.27.0+${NIGHTLY} \
+    apex==1.12.0+${NIGHTLY}
+```
+
+> Install only the `*-gfx942` **or** `*-gfx950` device packages matching your hardware for a smaller install.
+>
+> Nightly indexes are pruned. If this date has disappeared, pick a newer one that publishes a *complete* cp312 set — `torch`, `amd-torch-device-*`, `rocm-sdk-*`, `torchaudio`, `torchvision`, `amd-torchvision-device-*` and `apex` must all come from the same date.
+
+### 4.4 Initialize the ROCm SDK and export ROCm paths
+
+`rocm-sdk init` materializes the ROCm toolchain inside the venv. The variables below point the rest of the build at that in-venv ROCm and **must be set every time you use the environment** — see [Section 6](#6-make-the-environment-reproducible).
+
+```bash
+rocm-sdk init
+
+export ROCM_PATH=$(python -c 'import _rocm_sdk_devel, os; print(os.path.dirname(_rocm_sdk_devel.__file__))')
+export ROCM_HOME=$ROCM_PATH          # Primus uses ROCM_HOME; set both
+export HIP_PLATFORM=amd
+export HIP_PATH=$ROCM_PATH
+export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
+export HIP_INCLUDE_PATH=$ROCM_PATH/include
+export HIP_LIB_PATH=$ROCM_PATH/lib
+export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
+export PATH="$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
+export LD_LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/host-math/lib:$ROCM_PATH/lib/rocm_sysdeps/lib"
+export LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64"
+export CPATH=$HIP_INCLUDE_PATH
+export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+```
+
+> The Dockerfile hardcodes `ROCM_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel`. The `python -c ...` form derives it from the venv instead, so it works wherever you put the environment.
+
+Quick check before continuing:
+
+```bash
+hipcc --version
+python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"
+```
+
+### 4.5 Source-built components and their pins
+
+Build these against your in-venv ROCm and selected `PYTORCH_ROCM_ARCH`, in a scratch directory:
+
+```bash
+export GPU_ARCHS="${PYTORCH_ROCM_ARCH}"   # cross-compile during builds; see Section 6 for runtime
+mkdir -p ~/primus-build && cd ~/primus-build
+```
+
+| Component | Repository | Pin | Install command | Notes |
+|---|---|---|---|---|
+| Flash Attention | `ROCm/flash-attention` | `6387433156558135a998d5568a9d74c1778666d8` | `python setup.py install` | clone `--recursive` |
+| grouped_gemm | `caaatch22/grouped_gemm` | branch `rocm` | `pip install --no-build-isolation .` | MoE models |
+| causal-conv1d | `Dao-AILab/causal-conv1d` | `e940ead2fd962c56854455017541384909ca669f` | `pip install --no-build-isolation .` | needs `CAUSAL_CONV1D_FORCE_BUILD=TRUE`, `HIP_ARCHITECTURES=gfx942,gfx950` |
+| mamba | `AndreasKaratzas/mamba` | branch `enable-primus-hybrid-models` | `pip install --no-build-isolation .` | see note below |
+| aiter | `ROCm/aiter` | `0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e` | `PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 .` | clone `--recursive`; `pip uninstall aiter amd-aiter` first |
+| Primus-Turbo | `AMD-AGI/Primus-Turbo` | `edc8d2ccb0be4888e80ee7c6e765fd3956026a32` | `pip install -r requirements.txt` then `pip install --no-build-isolation . -v` | needs `HCC_AMDGPU_TARGET="gfx942,gfx950"`; see note below |
+| torchtune | `pytorch/torchtune` | `b4c98ac2a37f0397d64c22579aed415ce7264db6` | `pip install .` | patch first: `sed -i 's/use_grouped_mm = True/use_grouped_mm = False/g' torchtune/modules/moe/utils.py` |
+| torchao | `pytorch/ao` | `e9c7bead90b840b280f97374308255957108ce47` | `pip install --no-build-isolation .` | two patches: `pad_inner_dim` → `True` in `torchao/float8/config.py`, and `if defined(HIPBLASLT_VEC_EXT)` → `if false` in `torchao/csrc/rocm/swizzle/swizzle.cpp` |
+
+**mamba.** Install with `pip`, **not** `python setup.py install`: the legacy `easy_install` path ignores pip-installed packages and re-fetches the latest of every unpinned dependency as `.egg`s, which clobbers the pins. Pin `apache-tvm-ffi==0.1.11` beforehand, and `nvidia-cutlass-dsl==4.5.3` — newer CUTLASS DSL releases removed a symbol that `mamba_ssm`'s `quack-kernels` dependency imports, breaking `import mamba_ssm`. `mamba_ssm` pulls in `tilelang`, which v26.5 uninstalls again once everything is built.
+
+**Primus-Turbo.** Its `setup.py` hard-pins upstream `triton==3.7.0`, which conflicts with the ROCm `triton` that `torch` requires; installing it replaces the ROCm build and then segfaults on import, because ROCm's HIP runtime already loads its own LLVM. Install with `--no-deps` and supply `scipy` and `flydsl==0.2.4` yourself. Primus-Turbo also probes for rocSHMEM and mistakes the pip ROCm SDK for one; set `ROCSHMEM_HOME` to a non-existent path to make the probe fail cleanly and disable internode DeepEP.
+
+### 4.6 TransformerEngine
+
+v26.5 changed this: TE is no longer built from source but installed from the ROCm staging index.
+
+```bash
+# Runtime tuning flags (see Section 6 for the full set)
+export NVTE_USE_CAST_TRANSPOSE_TRITON=1
+
+# TE's own dependencies must be present first: the staging index serves only the
+# transformer-engine packages, so pip cannot fetch anything else while it is selected.
+pip install pybind11==3.0.4 importlib-metadata==8.7.1 onnxscript==0.7.0 \
+            pydantic==2.13.4 nvdlfw_inspect==0.2.2 einops onnx
+
+pip install \
+    --index-url https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/ \
+    --pre --no-build-isolation \
+    transformer_engine_rocm_torch==2.15.0.dev0+rocm7.15.0a20260716.a07e607
+```
+
+> **Those wheels need glibc ≥ 2.38.** They are built on Ubuntu 24.04, and `libtransformer_engine.so` requires `GLIBC_2.38` plus `GLIBCXX_3.4.32`. Ubuntu 22.04 has glibc 2.35, and glibc cannot be side-loaded via `LD_LIBRARY_PATH`, so on 22.04 the wheels install fine but fail at import with `version 'GLIBC_2.38' not found`.
+>
+> On such hosts, build the equivalent source commit instead — the same commit the version label refers to:
+>
+> ```bash
+> git clone --recursive https://github.com/ROCm/TransformerEngine.git
+> cd TransformerEngine
+> git checkout a07e607f14a5330807ffdafeeb6224f2d7dffacc
+> git submodule update --init --recursive
+> pip install psutil
+> MAX_JOBS=${MAX_JOBS} NVTE_FRAMEWORK=pytorch NVTE_USE_ROCM=1 \
+>   NVTE_USE_HIPBLASLT=1 NVTE_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \
+>   pip install --no-build-isolation .
+> ```
+>
+> `setup.sh` picks between the two automatically from the host's glibc; override with `PRIMUS_TE_MODE=wheel|source`.
+
+Either way, apply the same fix v26.5 applies inside the image: concurrent CK JIT compiles race to publish the same `.so`, and without this the loser aborts the run.
+
+```bash
+CK_JIT=$(python -c 'import os, sysconfig; print(os.path.join(sysconfig.get_paths()["purelib"], "transformer_engine/lib/ck_jit/ck_jit_compile.sh"))')
+sed -i 's|    mv -n "$_TMP_SO" "$OUTPUT"$|    mv -n "$_TMP_SO" "$OUTPUT" 2>/dev/null \|\| true|' "$CK_JIT"
+```
+
+### 4.7 Training-framework Python dependencies
+
+```bash
+pip install \
+    datasets==3.6.0 av==16.0.1 transformers==4.55.0 optree==0.18.0 sympy \
+    accelerate==1.9.0 trl==0.21.0 tensorboard==2.20.0 peft scipy einops \
+    flask-restful nltk pytest pytest-cov pytest_mock pytest-csv \
+    pytest-random-order sentencepiece wrapt \
+    zarr==2.18.7 numcodecs==0.12.1 xarray wandb tensorstore==0.1.45 \
+    pybind11 tiktoken pynvml "huggingface_hub[cli]"
+
+python3 -m nltk.downloader punkt_tab
+
+# AWS SDK (used by some data pipelines)
+pip install boto3==1.35.42 botocore==1.35.99
+```
+
+### 4.8 Install Primus and its submodules
+
+```bash
+cd ~/primus-build
+# Required to resolve a post-v26.2 attention backend issue
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=1
+
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+git checkout b511d1b66b0068715308ea9bfe8ba147ea1a3860   # release/v26.5
+git submodule update --init --recursive
+pip install -r requirements.txt
+
+# Megatron's dataset indexing needs a compiled pybind11 extension, or training
+# fails with "MockGPTDataset failed to build as a mock data generator".
+# Pass LIBEXT explicitly: the Makefile reads it from `python3-config`, which a
+# venv does not provide, so it otherwise emits a filename 3.12 will never import.
+EXT=$(python -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))')
+make -C third_party/Megatron-LM/megatron/core/datasets LIBEXT="$EXT"
+```
+
+If you already have a local Primus checkout, run `pip install -r requirements.txt` from its root and skip the clone — but still build `helpers_cpp` in *that* checkout, since it is the one you will launch training from.
+
+### 4.9 Optional: torchrec + FBGEMM (DLRM / recommendation workloads)
+
+```bash
+pip install --no-deps torchrec
+pip install tensordict iopath torchmetrics==1.0.3 \
+    git+https://github.com/mlperf/logging.git \
+    --extra-index-url https://rocm.nightlies.amd.com/whl-multi-arch
+
+# FBGEMM (GPU) — needs apt libtbb-dev
+export BUILD_ROCM_VERSION='7.14'
+
+git clone https://github.com/pytorch/FBGEMM.git
+cd FBGEMM
+git checkout 80bd3c077dc41b55cd16ed4dcad15cf7c1c1d76a
+cd fbgemm_gpu
+git clean -dfx && git submodule sync && git submodule update --init --recursive
+pip install -r requirements.txt
+pip install setuptools==75.1.0
+python setup.py install \
+    --build-variant=rocm \
+    --build-target=default \
+    -DAMDGPU_TARGETS=$PYTORCH_ROCM_ARCH \
+    -DHIP_ROOT_DIR=$ROCM_PATH \
+    -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
+    -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"
+cd ../..
+```
+
+> The Flux (`AMDiffusionBenchmark`) and DLRM (`DLRMBenchmark`) repos in the Dockerfile are benchmark workloads, not core training dependencies. Clone them only if you need those specific benchmarks.
+
+---
+
+## 5. Multi-node communication stack (UCX and OpenMPI)
+
+These are only needed for **multi-node distributed** training. They build from source into user-writable prefixes (no root needed, except the AINIC `.deb` already handled in [Section 3.3](#33-amd-ainic-library-optional-for-amd-pensando-nics)).
+
+### 5.1 UCX
+
+```bash
+cd ~/primus-build
+UCX_VERSION="1.18.0"
+wget https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz
+mkdir -p ucx-${UCX_VERSION}
+tar -zxf ucx-${UCX_VERSION}.tar.gz -C ucx-${UCX_VERSION} --strip-components=1
+cd ucx-${UCX_VERSION}
+mkdir build && cd build
+../configure --prefix=$HOME/primus-build/ucx-${UCX_VERSION}/install --with-rocm=${ROCM_PATH}
+make -j 16 && make install
+cd ../..
+
+export UCX_INSTALL_DIR=$HOME/primus-build/ucx-${UCX_VERSION}/install
+```
+
+### 5.2 OpenMPI
+
+```bash
+MPI_VERSION="4.1.6"
+wget https://download.open-mpi.org/release/open-mpi/v$(echo "${MPI_VERSION}" | cut -d. -f1-2)/openmpi-${MPI_VERSION}.tar.gz
+mkdir -p ompi-${MPI_VERSION}
+tar -zxf openmpi-${MPI_VERSION}.tar.gz -C ompi-${MPI_VERSION} --strip-components=1
+cd ompi-${MPI_VERSION}
+mkdir build && cd build
+# Install to a user-writable prefix instead of /opt to avoid sudo
+../configure --prefix=$HOME/primus-build/openmpi --with-ucx=${UCX_INSTALL_DIR} \
+    --disable-oshmem --disable-mpi-fortran
+make -j 16 && make install
+cd ../..
+
+export PATH="$HOME/primus-build/openmpi/bin:${PATH}"
+export LD_LIBRARY_PATH="$HOME/primus-build/openmpi/lib:${LD_LIBRARY_PATH}"
+```
+
+> The Dockerfile installs OpenMPI to `/opt/openmpi` (needs root). The `$HOME/primus-build/openmpi` prefix above keeps it unprivileged. Use `/opt/openmpi` only if you have root and want to match the image exactly.
+
+---
+
+## 6. Make the environment reproducible
+
+The ROCm paths and `NVTE_*` flags must be present in **every** shell that runs training.
+
+If you used the scripts, this is already handled — `source tools/installation/env.sh` sets everything (and refuses to run if `PRIMUS_BASE` is unset). For a manual install, append the equivalent to the venv's activation script:
+
+```bash
+cat >> ~/primus-env/bin/activate <<'EOF'
+
+# ---- Primus host environment ----
+export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+export ROCM_AMDGPU_TARGETS="gfx942,gfx950"
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+export HSA_NO_SCRATCH_RECLAIM=1
+
+export ROCM_PATH=$(python -c 'import _rocm_sdk_devel, os; print(os.path.dirname(_rocm_sdk_devel.__file__))')
+export ROCM_HOME=$ROCM_PATH
+export HIP_PLATFORM=amd
+export HIP_PATH=$ROCM_PATH
+export HIP_CLANG_PATH=$ROCM_PATH/llvm/bin
+export HIP_INCLUDE_PATH=$ROCM_PATH/include
+export HIP_LIB_PATH=$ROCM_PATH/lib
+export HIP_DEVICE_LIB_PATH=$ROCM_PATH/lib/llvm/amdgcn/bitcode
+export PATH="$ROCM_PATH/bin:$HIP_CLANG_PATH:$HOME/primus-build/openmpi/bin:$PATH"
+export LD_LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/host-math/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$HOME/primus-build/openmpi/lib"
+export LIBRARY_PATH="$HIP_LIB_PATH:$ROCM_PATH/lib:$ROCM_PATH/lib64"
+export CPATH=$HIP_INCLUDE_PATH
+export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+
+# TransformerEngine: attention backend selection used by Primus
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=1
+
+# TransformerEngine: CK performance knobs
+export NVTE_USE_CAST_TRANSPOSE_TRITON=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_USES_BWD_V3=1
+export CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2
+export NVTE_CK_HOW_V3_BF16_CVT=2
+# gfx942 (MI300X/MI325X) needs fp32 atomics for the CK v3 backward kernel: with
+# the Dockerfile's 0 it produces Inf gradients at the first step. Use 0 on gfx950.
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+
+# Runtime: let aiter detect the local GPU architecture.
+# NOTE: this is the runtime value. If you rebuild any source kernel later,
+#       re-export GPU_ARCHS="$PYTORCH_ROCM_ARCH" first, then set it back to native.
+export GPU_ARCHS=native
+
+# Multi-node comms (only if built)
+export UCX_HOME=$HOME/primus-build/ucx-1.18.0/install
+export MPI_HOME=$HOME/primus-build/openmpi
+# ---- end Primus host environment ----
+EOF
+```
+
+---
+
+## 7. Verify the installation
+
+```bash
+# Scripts: export the same PRIMUS_BASE and source env.sh
+# Manual:  source ~/primus-env/bin/activate
+
+# GPUs visible to ROCm?
+rocm-smi || ls -l /dev/kfd /dev/dri
+
+# PyTorch sees the GPUs?
+python -c "import torch; print('torch', torch.__version__); \
+print('gpu available:', torch.cuda.is_available()); \
+print('device count:', torch.cuda.device_count()); \
+print('device 0:', torch.cuda.get_device_name(0))"
+
+# Key kernel libs import cleanly?
+python -c "import transformer_engine, flash_attn, aiter, primus_turbo, mamba_ssm; print('kernels OK')"
+
+# Run a Primus benchmark / training directly (no container)
+cd "$WORKSPACE_DIR/Primus"   # or your Primus checkout
+./primus-cli direct -- benchmark gemm -M 4096 -N 4096 -K 4096
+./primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml \
+  --train_iters 10
+```
+
+A healthy run reports a decreasing `lm loss`, `number of nan iterations: 0`, and a steady `throughput per GPU (TFLOP/s/GPU)` after the first couple of iterations.
+
+Use `primus-cli direct` (not `container`) since you are running on bare metal with everything installed in your environment.
+
+---
+
+## 8. Other important considerations
+
+- **GPU device access without root**: the user running training must be able to read/write `/dev/kfd` and `/dev/dri/*`. This usually means membership in the `video` and `render` groups (`sudo usermod -aG video,render $USER`, then re-login). This is a one-time admin action.
+- **Hugging Face access**: if your config downloads weights or tokenizers from the Hub, export your token: `export HF_TOKEN=hf_xxx` (and/or `huggingface-cli login`). The token is needed for gated models like Llama.
+- **RDMA / multi-node limits**: high-performance networking typically requires locked-memory limits raised (`ulimit -l unlimited`) and possibly hugepages. These are configured in `/etc/security/limits.conf` and need admin help. Verify NICs with `ibv_devinfo` and `ibstat`.
+- **Disk and time**: source builds of aiter, Primus-Turbo, FBGEMM and (on older glibc) TransformerEngine are large and slow. Reserve plenty of disk and expect a multi-hour first build. Lower `MAX_JOBS` if the build runs out of memory.
+- **`ccache`**: installed in [Section 3.1](#31-build-toolchain-and-core-libraries); it dramatically speeds up rebuilds, with no extra configuration needed for a basic speedup.
+- **Architecture pinning**: building for only your actual GPU arch (e.g. `gfx942` for MI300X/MI325X, `gfx950` for MI350X/MI355X) significantly reduces build time and binary size versus building both. The scripts already do this automatically; it only needs stating explicitly for a manual build.
+- **Version drift is the main hazard.** The nightly wheels and source commits are a tested combination; several of the Dockerfile's *unpinned* transitive dependencies have since released versions that break the build or silently disable kernels. The scripts pin those explicitly — see [Section 1.6](#16-where-the-scripts-deliberately-differ-from-the-dockerfile). Treat the reference `Dockerfile` as the source of truth for versions, and the scripts as the source of truth for the workarounds.
+
+---
+
+## 9. Quick reference: minimal vs. full install
+
+If you only need **single-node Megatron/TorchTitan LLM pretraining**, you can skip several optional components:
+
+| Component                                               | Needed for                           |
+| ------------------------------------------------------- | ------------------------------------ |
+| Flash Attention, TransformerEngine, aiter, Primus-Turbo | Core LLM training (install these)    |
+| grouped_gemm                                            | MoE models                           |
+| causal-conv1d, mamba                                    | Hybrid / Mamba-family models         |
+| torchtune, torchao                                      | Post-training (SFT/LoRA), fp8        |
+| torchrec, FBGEMM, DLRM                                  | Recommendation (DLRM) workloads only |
+| Flux / AMDiffusionBenchmark                             | Diffusion benchmark only             |
+| UCX, OpenMPI, AINIC                                     | Multi-node distributed training      |
+
+Install the core rows first, validate with [Section 7](#7-verify-the-installation), then add the optional components as your workload requires.
diff --git a/docs/01-getting-started/glossary.md b/docs/01-getting-started/glossary.md
new file mode 100644
index 000000000..4ede3de7b
--- /dev/null
+++ b/docs/01-getting-started/glossary.md
@@ -0,0 +1,250 @@
+# Glossary
+
+Alphabetical reference for terms used in Primus documentation and configuration. Cross-links point to other production docs where applicable.
+
+---
+
+### AINIC
+
+**AMD AI NIC**—AMD’s AI-optimized network interface for multi-node GPU communication (for example, the **AMD Pensando™ Pollara 400 AI NIC**).
+
+---
+
+### Backend
+
+A training framework integrated into Primus (for example **Megatron-LM**, **TorchTitan**, **MaxText**, **Megatron Bridge**, **HummingbirdXT**).
+
+---
+
+### BackendAdapter
+
+Abstract class in Primus that connects a backend: discovery of setup paths, config conversion, and trainer loading.
+
+---
+
+### BackendRegistry
+
+Registry mapping backend names to adapter classes, often with **lazy import** to avoid loading unused frameworks.
+
+---
+
+### BaseTrainer
+
+Abstract trainer defining the lifecycle: **setup** → **init** → **train** → **cleanup**.
+
+---
+
+### BF16 / FP16 / FP8 / FP4
+
+Floating-point precisions: **Brain Float 16**, **IEEE half**, **8-bit float**, and **4-bit float** training or inference formats (exact support depends on backend and hardware).
+
+---
+
+### CP (context parallelism)
+
+Parallelism that **splits the sequence dimension** across devices for long-context training.
+
+---
+
+### DP (data parallelism)
+
+Replicates the model across GPUs; each rank processes **different data batches**.
+
+---
+
+### DeepEP
+
+**Deep Expert Parallelism**—Primus-Turbo’s acceleration path for **MoE token dispatch** and related expert-parallel work.
+
+---
+
+### EP (expert parallelism)
+
+Distributes **Mixture-of-Experts** expert networks across devices.
+
+---
+
+### Experiment config
+
+Top-level **YAML** describing `work_group`, **modules**, and **overrides** for a training run.
+
+---
+
+### FSDP
+
+**Fully Sharded Data Parallel**—shards parameters, gradients, and optimizer states across devices (PyTorch FSDP and similar concepts per backend).
+
+---
+
+### GBS (global batch size)
+
+Total **batch size across all data-parallel ranks** for one optimizer step (may combine micro-batching and gradient accumulation).
+
+---
+
+### Gradient accumulation
+
+Accumulates gradients over **multiple micro-batches** before an optimizer update.
+
+---
+
+### HipBLASLt
+
+AMD’s high-performance **BLAS** library with **autotuning** for GEMM and related kernels.
+
+---
+
+### Hook
+
+Shell or Python scripts under `runner/helpers/hooks/` executed at defined **lifecycle** points.
+
+---
+
+### LoRA
+
+**Low-Rank Adaptation**—parameter-efficient fine-tuning that trains small adapter matrices.
+
+---
+
+### MBS (micro batch size)
+
+Batch size **per GPU** (per rank) for **one forward/backward pass** within a gradient-accumulation window.
+
+---
+
+### MLA (multi-latent attention)
+
+Compressed **KV-cache** attention architecture used in models such as DeepSeek.
+
+---
+
+### MoE (mixture of experts)
+
+Architecture with **multiple expert** sub-networks and a **router** that assigns tokens to experts.
+
+---
+
+### Model config
+
+YAML **preset** describing architecture (hidden size, layers, attention heads, and so on).
+
+---
+
+### Module config
+
+YAML **preset** for training behavior: learning rate, batch sizes, optimizer, schedules.
+
+---
+
+### NCCL / RCCL
+
+**NVIDIA Collective Communications Library** / **ROCm** equivalent—libraries for **GPU collective** operations in distributed training.
+
+---
+
+### PP (pipeline parallelism)
+
+Splits **model layers** into **stages** on different devices.
+
+---
+
+### Patch
+
+Runtime **monkey-patch** registered in **PatchRegistry** and applied at a named training phase.
+
+---
+
+### PatchRegistry
+
+Registry of **phase-aware** patches (for example `build_args`, `setup`, `before_train`, `after_train`).
+
+---
+
+### Platform config
+
+YAML describing **cluster environment** mappings (for example `platform_azure.yaml`): env vars, paths, and scheduler hints.
+
+---
+
+### Preflight
+
+Cluster **diagnostic** tooling that checks host, GPU, network, and baseline performance before long jobs. See `primus/tools/preflight/` in the repository.
+
+---
+
+### Preset
+
+Reusable YAML fragment under `primus/configs/` (**module**, **model**, or **platform**).
+
+---
+
+### PrimusRuntime
+
+Core **orchestrator**: loads configuration, resolves the backend, applies patches, and drives the **trainer lifecycle**.
+
+---
+
+### Primus-SaFE
+
+**Stability and Fault-tolerance Engine**—external ecosystem component for **cluster management** and resilience. This repository references it in auxiliary tooling but does not include a production integration guide.
+
+---
+
+### Primus-Turbo
+
+High-performance **operator** library (for example FlashAttention-style kernels, GEMM, collectives, grouped GEMM).
+
+---
+
+### Projection
+
+Tools that **estimate memory** and **training performance** without requiring a full production cluster.
+
+---
+
+### ROCm
+
+**Radeon Open Compute**—AMD’s GPU computing platform (drivers, compilers, libraries).
+
+---
+
+### SFT (supervised fine-tuning)
+
+Supervised fine-tuning that typically **updates all** (or a defined subset of) model parameters, as opposed to adapter-only methods.
+
+---
+
+### SP (sequence parallelism)
+
+Parallelism that extends tensor-parallel regions to **non-TP** parts of the model to **reduce activation memory**.
+
+---
+
+### TP (tensor parallelism)
+
+Splits **layer weights** across GPUs within a node (or defined process group).
+
+---
+
+### Transformer engine (TE)
+
+Library stack for **FP8** and related training optimizations (availability depends on backend and build).
+
+---
+
+### VPP (virtual pipeline parallelism)
+
+**Interleaved** pipeline parallelism with **multiple virtual stages** per device to improve utilization.
+
+---
+
+### Zero-bubble
+
+Pipeline scheduling that **reduces or eliminates pipeline bubbles** (idle time between micro-batches).
+
+---
+
+## Related documentation
+
+- [Overview](./overview.md)
+- [Configuration system](../02-user-guide/configuration-system.md)
diff --git a/docs/01-getting-started/installation.md b/docs/01-getting-started/installation.md
new file mode 100644
index 000000000..ca054197b
--- /dev/null
+++ b/docs/01-getting-started/installation.md
@@ -0,0 +1,266 @@
+# Installation and setup
+
+This guide covers supported platforms, prerequisites, and how to set up the training environment: **container (recommended)** and **bare metal**, plus **multi-node distributed training** (Slurm recommended).
+
+---
+
+## Supported platforms
+
+
+| Requirement | Notes                                                                                                       |
+| ----------- | ----------------------------------------------------------------------------------------------------------- |
+| **OS**      | Linux (ROCm-supported distributions per AMD documentation).                                                 |
+| **ROCm**    | **≥ 7.0** recommended.                                                                                      |
+| **GPUs**    | AMD Instinct™ **MI300X**, **MI325X**, **MI355X** (or other ROCm-supported Instinct SKUs your site supports). |
+
+
+---
+
+## Prerequisites
+
+
+| Prerequisite                                                  | Purpose                                       |
+| ------------------------------------------------------------- | --------------------------------------------- |
+| **AMD Instinct GPUs**                                         | Training and benchmarks execute on GPU.       |
+| **ROCm drivers and user-space stack**                         | Required for HIP, RCCL, and ML frameworks.    |
+| **Docker ≥ 24.0** (or Podman with compatible GPU passthrough) | Container mode and reproducible environments. |
+| **git**                                                       | Clone the repository and submodules.          |
+
+
+### Quick environment checks
+
+```bash
+rocm-smi
+docker --version
+```
+
+`rocm-smi` should list your GPUs; `docker --version` should report **24.0** or newer.
+
+---
+
+## Container setup (recommended)
+
+AMD publishes training Docker images monthly, providing a consistent, ready-to-run environment optimized for AMD GPUs. It is recommended to use the AMD-published training Docker images together with this Primus-LM repository to run your training jobs. The images support pre-training and post-training workflows with multiple backends including Megatron-LM, TorchTitan, and JAX MaxText, alongside ROCm-optimized components.
+
+Check the AMD-published training Docker images here:
+
+- For Megatron-LM and TorchTitan backends: [https://hub.docker.com/r/rocm/primus/tags](https://hub.docker.com/r/rocm/primus/tags)
+- For MaxText backend: [https://hub.docker.com/r/rocm/jax-training/tags](https://hub.docker.com/r/rocm/jax-training/tags)
+
+### 1. Pull the image
+
+```bash
+# For Megatron-LM and TorchTitan backends
+docker pull rocm/primus:v26.4
+# For MaxText backend
+docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+```
+
+### 2. Clone the repository
+
+Submodules are required for third-party backends and tools:
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+# checkout the branch for the specific release
+git checkout release/v26.4
+git submodule update --init --recursive
+```
+
+### 3. Run a verification benchmark
+
+From the repository root:
+
+```bash
+./primus-cli container --image rocm/primus:v26.4 -- \
+  benchmark gemm --M 4096 --N 4096 --K 4096
+```
+
+A successful run validates the GPU stack and Primus CLI wiring without launching a full training job.
+
+---
+
+## Bare-metal (host) setup
+
+> **The [container setup](#container-setup-recommended) above is strongly recommended.** The AMD-published training Docker image is the tested, reproducible, and best-supported path. Build the full stack on a bare-metal host only when containers are not an option (for example, due to policy or operational constraints).
+
+### What to expect
+
+Reproducing the training environment on the host means building the **same stack the Docker image ships**, mostly from source. This is a **long, build-heavy process** that requires:
+
+- Several **source-built kernel libraries** (Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, causal-conv1d, Mamba) compiled against ROCm.
+- A **machine with many CPU cores, ample RAM, and tens of GB of free disk**.
+- A long time (expect a **multi-hour first build**).
+
+### What a complete host environment needs
+
+| Layer                   | What it provides                                                                 | How it is installed              |
+| ----------------------- | -------------------------------------------------------------------------------- | -------------------------------- |
+| Kernel/hardware         | AMD GPU driver (amdgpu KMD) and device access (`/dev/kfd`, `/dev/dri`)            | OS/administrator (root, one-time) |
+| OS libraries            | Build toolchain and runtime libraries (`g++`, `git`, RDMA, hwloc, etc.)          | `apt` (root, one-time)           |
+| ROCm user-space         | `rocm-sdk-devel` + device wheels—**no system-wide ROCm install required**       | `pip` (TheRock wheels, in `venv`) |
+| Deep learning framework | ROCm-enabled PyTorch (`torch`, `torchvision`, `torchaudio`, `apex`)              | `pip` (TheRock wheels, in `venv`) |
+| Accelerated kernels     | Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, Mamba     | build from source (in `venv`)    |
+| Multi-node communications | UCX, OpenMPI, rocSHMEM, AMD AINIC—only for distributed (multi-node) training  | build from source or `apt`       |
+| Primus + Python dependencies | Primus, submodules, and training libraries (datasets, transformers, wandb, etc.) | `git` + `pip` (in `venv`)   |
+
+### General approach
+
+1. **System packages (root, one-time):** install the build toolchain and, for multi-node, the RDMA/networking libraries via `apt`. The GPU kernel driver must already be loaded.
+2. **Python virtual environment (no root):** create a `venv`, then install ROCm and PyTorch from AMD's TheRock multi-arch wheels—this replaces a system ROCm install and keeps everything unprivileged.
+3. **Build the accelerated kernels from source** against the ROCm in `venv` and your GPU architecture (`gfx942` for MI300X/MI325X, `gfx950` for MI350X/MI355X).
+4. **Install Primus and its Python dependencies**, then persist the required environment variables (ROCm paths, `NVTE_*` flags) in your `venv` activation script.
+5. **(Optional) Build the multi-node communication stack** (UCX, OpenMPI, rocSHMEM) only if you require RDMA-based distributed training.
+
+### Detailed instructions
+
+Follow the full, step-by-step guide here, which includes the exact pinned versions, environment variables, and automated install scripts:
+
+- **[Bare-metal installation (PyTorch: Megatron-LM / TorchTitan): build the Primus training stack from source (no Docker)](./bare-metal-installation.md)**
+- **[Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)](./bare-metal-installation-jax.md)**
+
+### Verify
+
+After the build, validate the environment and run a benchmark directly on the host (no container):
+
+```bash
+./primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096
+```
+
+---
+
+## Multi-node distributed training (Slurm recommended)
+
+For training jobs that span **multiple nodes**, we recommend using **[Slurm](https://slurm.schedmd.com/)**. Slurm is a cluster workload manager and job scheduler: it allocates nodes and GPUs, places your job on them, launches one task per node, and injects the topology information (node list, node count, per-node rank) that distributed PyTorch needs. `primus-cli` has a built-in **`slurm` mode** that wraps `srun`/`sbatch` and wires this topology into the training launcher for you.
+
+### Cluster baseline
+
+Before launching distributed jobs, ensure every participating node has:
+
+- The **same software stack**—use the **same container image** on all nodes (recommended), or an identical bare-metal install (see sections above).
+- A **shared filesystem** for code, datasets, checkpoints, and logs (e.g. NFS/Lustre), mounted at the same path on every node.
+- **Working inter-node networking**: for best performance, use a high-speed RDMA fabric (InfiniBand, RoCE, or AMD AINIC) and ensure RCCL can select the right interface.
+
+### Setting up Slurm
+
+Setting up Slurm itself is a cluster-administration task and is outside Primus's scope. Follow the official documentation:
+
+- [Slurm Quick Start (users)](https://slurm.schedmd.com/quickstart.html)
+- [Slurm Quick Start Administrator Guide (install & configure)](https://slurm.schedmd.com/quickstart_admin.html)
+
+Once `sinfo` and `srun` work on your login node, Primus can submit jobs to it. If you don't administer the cluster, your site administrator typically provides the partition, account, and reservation names you need.
+
+### Launching with `primus-cli slurm`
+
+The Slurm wrapper uses a single `--` separator:
+
+- **Before the `--`**: put the launcher (`srun` or `sbatch`, default `srun`) and Slurm flags (`-N`, `-p`, `--nodelist`, `--account`, `--qos`, `--reservation`, …).
+- **After the `--`**: put the Primus command to run (`train` / `benchmark` / …). It runs inside the container image (see [Selecting the container image](#selecting-the-container-image) below).
+
+```bash
+cd /path/to/Primus
+
+# Pretrain on 2 nodes via srun
+./primus-cli slurm srun -N 2 \
+  -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+Add the global `--dry-run` flag before the launcher to print the exact command without executing it:
+
+```bash
+./primus-cli --dry-run slurm srun -N 2 \
+  -- train pretrain --config <your-config>.yaml
+```
+
+> See `runner/README.md` and the [CLI reference](../02-user-guide/cli-reference.md) for the full set of launcher flags and scenarios.
+
+#### Selecting the container image
+
+The image used for the job is resolved according to the following priority order (highest first):
+
+1. **`DOCKER_IMAGE` environment variable**—overrides everything else. This is the simplest way to switch images and it propagates to all nodes:
+
+```bash
+export DOCKER_IMAGE=rocm/primus:v26.4
+./primus-cli slurm srun -N 2 \
+  -- train pretrain --config <your-config>.yaml
+```
+
+2. **`--image` CLI flag**—place it immediately after the `--`, before the Primus command (ignored if `DOCKER_IMAGE` is set):
+
+```bash
+./primus-cli slurm srun -N 2 \
+  -- --image rocm/primus:v26.4 train pretrain --config <your-config>.yaml
+```
+
+3. **Config file default**—`container.options.image` in `runner/.primus.yaml` (or your `~/.primus.yaml`), which is set to `rocm/primus:v26.4` by default.
+
+### Distributed environment variables
+
+Primus launches training with `torchrun`, which needs to know the cluster topology. These are the key variables:
+
+
+| Variable        | Role                                                      | Default     |
+| --------------- | -------------------------------------------------------- | ----------- |
+| `MASTER_ADDR`   | Hostname/IP of rank 0; all ranks rendezvous here.        | `localhost` |
+| `MASTER_PORT`   | Port on the master used for rendezvous.                  | `1234`      |
+| `NNODES`        | Number of nodes in the job.                              | `1`         |
+| `NODE_RANK`     | Index of this node (0-based, unique per node).           | `0`         |
+| `GPUS_PER_NODE` | GPUs (processes) to launch per node.                     | `8`         |
+
+The total number of training processes (world size) is `NNODES × GPUS_PER_NODE`.
+
+**Under Slurm, the values of these variables are derived automatically.** `primus-cli slurm` reads Slurm's own variables and sets the Primus ones for you: `NNODES` from `SLURM_NNODES`, `NODE_RANK` from `SLURM_NODEID`, and `MASTER_ADDR` from the first host in `SLURM_NODELIST` (with `MASTER_PORT` defaulting to `1234`). You normally only need to set `GPUS_PER_NODE` if your nodes don't have 8 GPUs. You can still override `MASTER_PORT` (e.g. to avoid a port clash) via `--env`.
+
+### Without Slurm: Kubernetes or `parallel-ssh`
+
+Slurm is recommended but not required. The mechanism underneath is simple: **set the same distributed environment variables on every node, point them all at the same `MASTER_ADDR`, give each node a unique `NODE_RANK`, and run the same `primus-cli direct` command on each node.** Any tool that can run a command across nodes works—for example Kubernetes (e.g. a `PyTorchJob` / indexed Job) or `parallel-ssh`/`pdsh`.
+
+For a 2-node job you would run, on the master (rank 0):
+
+```bash
+export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR=<node0-host> MASTER_PORT=1234
+./primus-cli direct -- train pretrain --config <your-config>.yaml
+```
+
+and on the worker (rank 1), the same command with `NODE_RANK=1` and the same `MASTER_ADDR`:
+
+```bash
+export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=1 MASTER_ADDR=<node0-host> MASTER_PORT=1234
+./primus-cli direct -- train pretrain --config <your-config>.yaml
+```
+
+With Kubernetes, inject these as container environment variables (deriving `NODE_RANK` from the pod's index); with `parallel-ssh`, pass them per host. The training command itself is identical on every node.
+
+### Other important considerations
+
+- **Cluster validation.** Run the built-in preflight check across your nodes before a long job: `./primus-cli slurm srun -N <N> -- preflight`.
+- **Networking/RCCL.** On RDMA fabrics, make sure the correct interface is selected (Primus should auto-detect this; if it doesn't, set `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA`). Use `NCCL_DEBUG=INFO` (passed via `--env`) to diagnose hangs at startup.
+- **RDMA limits.** High-performance networking usually needs locked-memory limits raised (`ulimit -l unlimited`) and sometimes hugepages—configured by your admin.
+- **`MASTER_PORT` must be free** on the master node and reachable from all workers; firewalls between nodes will cause rendezvous timeouts.
+- **Hugging Face access.** If your configuration downloads gated models or tokenizers, export `HF_TOKEN` (and ensure it is propagated to all nodes and into the containers).
+
+---
+
+## Post-installation verification checklist (for all setup approaches)
+
+
+| Step                 | Check                                                                                                   |
+| -------------------- | ------------------------------------------------------------------------------------------------------- |
+| **ROCm**             | `rocm-smi` shows expected GPUs and no driver errors.                                                    |
+| **Container engine** | `docker run --rm ... rocm/primus:v26.4` (or your site’s GPU test) succeeds.                             |
+| **GEMM benchmark**   | `./primus-cli` **container** or **direct** benchmark completes (see sections above).                    |
+| **Preflight**        | Run preflight diagnostics: `./primus-cli direct -- preflight` (single node) or `./primus-cli slurm srun -N <N> -- preflight` (cluster).         |
+
+
+If your training pulls models or tokenizers from Hugging Face Hub, configure tokens (for example `HF_TOKEN`) in the environment or container flags as required by your configuration.
+
+---
+
+## Related documentation
+
+- [Overview](./overview.md)
+- [Quickstart](./quickstart.md)
+- [CLI reference](../02-user-guide/cli-reference.md)
diff --git a/docs/01-getting-started/overview.md b/docs/01-getting-started/overview.md
new file mode 100644
index 000000000..3bbabe351
--- /dev/null
+++ b/docs/01-getting-started/overview.md
@@ -0,0 +1,137 @@
+# Primus overview
+
+## Executive summary
+
+**Primus** is a YAML-driven training framework for large-scale foundation model work on AMD GPUs. It targets **machine learning engineers**, **researchers**, and **platform/operations teams** who need reproducible, multi-backend training pipelines on AMD Instinct™ hardware. This project, also referred to as **Primus-LM**, serves as the training component within the broader [Primus ecosystem](#primus-ecosystem).
+
+**Repository:** [https://github.com/AMD-AGI/Primus](https://github.com/AMD-AGI/Primus)
+
+---
+
+## What Primus provides
+
+| Area | Description |
+|------|-------------|
+| **Multi-backend training** | One workflow surface over Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge, and HummingbirdXT. |
+| **Unified CLI** | `primus-cli` with **direct** (bare metal), **container** (Docker/Podman), and **slurm** (cluster) execution modes. |
+| **YAML-driven configuration** | Experiment, model, module, and platform presets composed from reusable fragments under `primus/configs/`. |
+| **Benchmark suite** | Built-in benchmarks (for example, GEMM) for quick hardware and stack validation. |
+| **Preflight diagnostics** | Cluster-oriented checks for host, GPU, and network health before long jobs. |
+| **Performance projection** | Tools to estimate memory use and throughput without occupying a full cluster. |
+
+Workflows span **pretraining** and **post-training** (including SFT and LoRA). Some Megatron configuration files expose RL-related parameters, but reinforcement-learning workflows are outside the scope of this documentation set: they are not part of the tested, supported paths described here, and this documentation does not cover how to run them.
+
+---
+
+## Primus ecosystem
+
+The training component (Primus-LM) serves as the layer between the stability and platform services above it and the low-level operator libraries below it:
+
+```
+                    +------------------+
+                    |   Primus-SaFE    |
+                    | (stability /     |
+                    |  cluster mgmt)   |
+                    +--------+---------+
+                             |
+                    +--------v---------+
+                    |    Primus-LM     |
+                    |  (this repo:     |
+                    |   training)      |
+                    +--------+---------+
+                             |
+                    +--------v---------+
+                    |  Primus-Turbo    |
+                    | (operators /     |
+                    |  kernels)        |
+                    +------------------+
+```
+
+- **Primus-SaFE**: Stability and fault-tolerance oriented cluster management referenced by auxiliary tooling (not documented here).
+- **Primus-LM**: Training orchestration, backends, CLI, and configurations (maintained in the [AMD-AGI/Primus](https://github.com/AMD-AGI/Primus/) repository).
+- **Primus-Turbo**: High-performance operators (for example, FlashAttention-style kernels, GEMM, and collectives).
+
+---
+
+## Supported backends
+
+| Backend | Typical use |
+|---------|-------------|
+| **Megatron-LM** | Broadest model coverage; default for many GPT-style and MoE recipes. |
+| **TorchTitan** | PyTorch-native large-model training paths. |
+| **JAX MaxText** | JAX/Flax training stacks aligned with MaxText. |
+| **Megatron Bridge** | Post-training and bridge workflows on top of Megatron-related stacks. |
+| **HummingbirdXT** | Additional integrated training path when enabled by your deployment. |
+
+Backend choice is specified in the configuration YAML and resolved through Primus’s adapter layer (see [glossary](./glossary.md)).
+
+---
+
+## Supported hardware and stack
+
+| Item | Requirement |
+|------|----------------|
+| **GPUs** | AMD Instinct™ **MI300X**, **MI325X**, **MI355X** |
+| **Platform** | **ROCm** (version **≥ 7.0** recommended) |
+| **Container image (reference)** | `docker.io/rocm/primus:v26.4` |
+
+Exact kernel and driver packages should match AMD’s documentation for your GPU SKU and ROCm release.
+
+---
+
+## Key dependencies
+
+Primus depends on the following categories of software (the list is non-exhaustive):
+
+| Category | Examples |
+|----------|----------|
+| **Framework** | PyTorch (Megatron-LM, TorchTitan paths); JAX/Flax (MaxText path). |
+| **AMD stack** | ROCm, RCCL, HipBLASLt (GEMM), GPU drivers. |
+| **Execution** | Docker or Podman for container mode; Slurm for cluster mode. |
+| **Observability & tooling** | **loguru** (logging), **Weights & Biases** (`wandb`) and other optional trackers (see `requirements.txt`). |
+
+Install specifics are covered in [Installation and setup](./installation.md).
+
+---
+
+## Runtime model
+
+At a high level, a run follows this pipeline:
+
+1. **YAML configuration** defines work group, modules, model preset, and overrides.
+2. **`primus-cli`** selects execution mode (direct, container, or slurm) and forwards to the runner.
+3. **Backend adapter** maps the resolved configuration to the target framework (Megatron-LM, TorchTitan, and so on).
+4. **Distributed launch** typically uses **`torchrun`** (or the backend’s equivalent) to start workers across GPUs and nodes.
+
+For CLI shape and options, see [Quickstart](./quickstart.md) and [CLI reference](../02-user-guide/cli-reference.md).
+
+---
+
+## Repository layout (top level)
+
+The following table shows the top-level layout of the Primus project repository, [AMD-AGI/Primus](https://github.com/AMD-AGI/Primus/):
+
+| Path | Role |
+|------|------|
+| `primus/` | Core library: configurations, runtime, trainers, backend adapters, tools (including preflight). |
+| `runner/` | CLI implementation, helpers, hooks, and launch glue. |
+| `examples/` | End-to-end example YAML and recipes per backend and GPU SKU. |
+| `docs/` | Project documentation (this documentation set). |
+| `tests/` | Automated tests. |
+| `tools/` | Auxiliary scripts and utilities. |
+| `benchmark/` | Benchmark drivers and related assets. |
+| `third_party/` | Vendored or submodule dependencies. |
+
+---
+
+## Next steps
+
+- [Installation and setup](./installation.md): ROCm, Docker, pip, and Slurm setup.
+- [Quickstart](./quickstart.md): run a minimal training job in minutes.
+- [Glossary](./glossary.md): terms used across Primus documentation.
+
+---
+
+## Licensing
+
+Primus is distributed under the terms described in the project's `LICENSE` file and `README`. If you encounter differing license references between the `README` and the repository root `LICENSE` file, treat licensing as project-specific: confirm the intended terms with the maintainers and your own compliance process before redistributing.
diff --git a/docs/01-getting-started/quickstart.md b/docs/01-getting-started/quickstart.md
new file mode 100644
index 000000000..bb9a4789a
--- /dev/null
+++ b/docs/01-getting-started/quickstart.md
@@ -0,0 +1,171 @@
+# Quickstart (about five minutes)
+
+This guide runs a **small Megatron-LM pretraining example** with **mock data** so you can validate the stack without preparing a full dataset. The same sample YAML works across **direct**, **container**, and **Slurm** modes.
+
+> **Recommended: run this example inside AMD-published training Docker images.** AMD publishes ready-to-run ROCm training images (`rocm/primus` for Megatron-LM and TorchTitan, `rocm/jax-training` for MaxText) with all dependencies and the complete training software stack already installed and validated. Using them means you don't have to build or tune the environment yourself, and—most important for multi-node jobs—**every node runs an identical, tested environment**. That consistency helps avoid version-skew and configuration issues that often occur with per-host installations. Host-based installation is supported, but is recommended only for advanced users (see [Installation and setup](./installation.md)).
+
+See [Installation and setup](./installation.md) for prerequisites and environment setup.
+
+---
+
+## Prerequisites
+
+- AMD ROCm drivers (version ≥ 7.0 recommended)
+- Docker (version ≥ 24.0) with ROCm support
+- ROCm-compatible AMD GPUs (e.g., Instinct MI300 series)
+- Proper permissions for Docker and GPU device access
+
+---
+
+## Option 1: Clone the repository and run training in a container (recommended)
+
+### Step 1: Pull the container image
+
+Check the AMD published training Docker images:
+
+- Megatron-LM and TorchTitan backends: <https://hub.docker.com/r/rocm/primus/tags>
+- MaxText backend: <https://hub.docker.com/r/rocm/jax-training/tags>
+
+```bash
+# For Megatron-LM and TorchTitan backends
+docker pull rocm/primus:v26.4
+# For MaxText backend
+docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+```
+
+### Step 2: Clone the repository
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+# checkout the branch for the specific release
+git checkout release/v26.4
+git submodule update --init --recursive
+```
+
+### Step 3: Run training inside container
+
+Run the training from the repository root. If your configuration downloads weights or tokenizers from Hugging Face Hub, pass `HF_TOKEN` into the container:
+
+```bash
+./primus-cli container --image rocm/primus:v26.4 \
+  --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
+  -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+---
+
+## Option 2: Install Primus from a wheel and run training in a container
+
+### Step 1: Install Primus as a Python package
+
+Install Primus in a virtual environment:
+
+```bash
+python -m venv primus-env
+source primus-env/bin/activate
+pip install "primus==26.4.0" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/
+```
+
+> **Note:** This installs only the Primus CLI into your virtual environment (under `site-packages`), without other dependencies. Third-party submodules are downloaded on the first run of the container, and the complete training software stack is provided in the AMD-published Docker images. You can launch `primus-cli` from any directory.
+
+### Step 2: Run training in a container using the pip-installed Primus
+
+```bash
+primus-cli container --image rocm/primus:v26.4 \
+  --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \
+  --volume /path/to/your/data:/data -- --log_file /data/run.log \
+  -- train pretrain --config /data/your/config.yaml
+```
+
+> **Note:** `--volume` mounts a local data directory into the container. `--log_file` writes the training log there; if omitted, logs go to the Primus install directory (`site-packages/primus/logs` by default).
+
+---
+
+## Expected output
+
+You should see the backend initialize distributed processes, load the training configuration, and emit **iteration-level logs** (with loss, throughput, step index, etc.). Exact fields depend on the backend and logging configuration; a typical pattern resembles:
+
+```
+... [INFO] starting training ...
+... iteration      1 | loss: 10.xxx | ...
+... iteration      2 | loss:  9.xxx | ...
+```
+
+Let the job run briefly to confirm stability; stop with `Ctrl+C` when satisfied.
+
+---
+
+## Same configuration, three execution modes
+
+Use one configuration file: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml`.
+
+| Mode | Example command |
+|------|------------------|
+| **Container** | `./primus-cli container --image rocm/primus:v26.4 -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` |
+| **Direct** | `./primus-cli direct -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` |
+| **Slurm** | `./primus-cli slurm srun -N <nodes> ... -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` |
+
+Replace `<nodes>` and Slurm resource flags with values appropriate for your cluster. **Container** and **Slurm** runs execute inside a Docker image (Slurm dispatches through the container launcher on each node); **Direct** runs execute inside a Docker container (you start the container yourself and run the command) or directly on the host if the training environment is already set up.
+
+> **Multi-node networking:** Primus auto-detects RDMA settings (`NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, …) on each node. If auto-detection selects the wrong NIC, or your fabric needs specific values (RoCE `NCCL_IB_GID_INDEX`, AMD AINIC, etc.), override them via `--env` or your config. See [Multi-node networking](../04-technical-guides/multi-node-networking.md).
+
+### Selecting the container image
+
+For container and Slurm runs, Primus resolves which Docker image to use in the following order (**highest priority first**):
+
+1. **`DOCKER_IMAGE` environment variable**—if set, it overrides every other source (including `--image`):
+
+```bash
+export DOCKER_IMAGE=rocm/primus:v26.4
+./primus-cli container -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+2. **`--image` command-line argument**—the usual per-run override, passed as a container mode argument (before `--`):
+
+```bash
+./primus-cli container --image rocm/primus:v26.4 -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+3. **`image` field in a config file** (`container.options.image`)—used when neither of the above is set:
+
+```yaml
+container:
+  options:
+    image: "rocm/primus:v26.4"
+```
+
+Primus loads a **single** config file—the first that exists among `--config <file>`, then `~/.primus.yaml`, then the shipped `runner/.primus.yaml` (these files are **not** merged). Because `runner/.primus.yaml` ships with a default image, a bare `./primus-cli container -- ...` works out of the box.
+
+>**Check the logs to make sure actual image being used is the one you wanted.**
+
+
+---
+
+## Command structure
+
+`primus-cli` parses **global options**, a **mode** (`direct`, `container`, `slurm`, …), optional **mode-specific arguments**, then a **`--` separator** followed by the **subcommand and its arguments** (for example `train` or `benchmark`).
+
+```
+primus-cli [global-options] <mode> [mode-args] -- <command> [command-args...]
+```
+
+Example:
+
+```text
+primus-cli  container  --image rocm/primus:v26.4  --  train pretrain --config path/to/experiment.yaml
+            ^^^^^^^^^^  ^^^^^^^^^^^^^^^^^^^^^^^^^     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+            mode        mode args                     command + args
+```
+
+---
+
+## Next steps
+
+| Topic | Document |
+|-------|----------|
+| Full CLI flags and subcommands | [CLI reference](../02-user-guide/cli-reference.md) |
+| YAML presets, overrides, and composition | [Configuration system](../02-user-guide/configuration-system.md) |
+| Pretraining workflows and backend notes | [Pretraining workflows](../02-user-guide/pretraining.md) |
+| Terminology | [Glossary](./glossary.md) |
diff --git a/docs/02-user-guide/README.md b/docs/02-user-guide/README.md
new file mode 100644
index 000000000..3c34cb574
--- /dev/null
+++ b/docs/02-user-guide/README.md
@@ -0,0 +1,20 @@
+# User guide
+
+Core workflows and day-to-day usage.
+
+- [Primus tools](primus-tools.md): start here—an at-a-glance catalog of all Primus tools and ecosystem projects with how-to starting points
+- [CLI reference](cli-reference.md): `primus-cli` modes, flags, and subcommands
+- [Configuration system](configuration-system.md): YAML configuration model, presets, overrides, inheritance
+- [Pretraining](pretraining.md): pretraining **concepts**: backends, YAML structure, parallelism, configuration inventory
+- [Backend training recipes](training-recipes.md): pretraining **commands**: copy-paste, GPU-arch-specific run commands
+- [Post-training](posttraining.md): SFT and LoRA fine-tuning via Megatron Bridge
+- [Node-smoke test instruction](node-smoke-test-instruction.md): screen a cluster fast and exclude bad nodes before launching a real training job
+- [Preflight](preflight.md): cluster diagnostics and environment validation
+- [Run preflight without a container](preflight-without-container.md): run cluster-diagnostic tool directly on the host
+- [Benchmarking](benchmarking.md): GEMM, RCCL, and dense-GEMM benchmark suites
+- [Projection](projection.md): memory and performance projection tools
+- [Tuning agent](tuning-agent.md): LLM-driven search for an optimal training configuration (uses projection as an oracle)
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/02-user-guide/benchmarking.md b/docs/02-user-guide/benchmarking.md
new file mode 100644
index 000000000..722865593
--- /dev/null
+++ b/docs/02-user-guide/benchmarking.md
@@ -0,0 +1,204 @@
+# Benchmark suite
+
+Primus ships microbenchmarks for GPU compute and distributed communication. They are exposed as the `benchmark` subcommand of the Primus CLI. Use them to validate a node or cluster before long training jobs.
+
+**Implementation:** `primus/cli/subcommands/benchmark.py` (initializes distributed execution, runs the selected suite, then finalizes).
+
+Related documentation: [Preflight diagnostics](./preflight.md) (broader cluster checks), [Memory and performance projection](./projection.md) (training-scale estimates), [Installation](../01-getting-started/installation.md) (environment setup).
+
+---
+
+## Overview and command syntax
+
+```bash
+primus-cli [global-options] <mode> [mode-args] -- benchmark <suite> [suite-specific-args]
+```
+
+- **`<mode>`** is typically `direct`, `container`, or `slurm` so that `WORLD_SIZE`, `RANK`, `MASTER_ADDR`, and related variables are set consistently.
+- **`benchmark`** runs inside the Primus Python CLI; the runner wires up the process environment the same way as training.
+
+The CLI also registers an `attention` suite; the subsections below cover **`gemm`**, **`gemm-dense`**, **`gemm-deepseek`**, **`strided-allgather`**, and **`rccl`**.
+
+---
+
+## Quick start
+
+Single-node GEMM:
+
+```bash
+primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096 --dtype bf16 --duration 10
+```
+
+Multi-node RCCL on Slurm:
+
+```bash
+primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M
+```
+
+---
+
+## Suite reference
+
+### `gemm`
+
+Single-shape general matrix multiply (GEMM) microbenchmark.
+
+| Argument | Description |
+|----------|-------------|
+| `--M`, `--N`, `--K` | Matrix dimensions (defaults: 4096 / 4096 / 4096). |
+| `--trans_a` | Transpose the A matrix. |
+| `--trans_b` | Transpose the B matrix. |
+| `--dtype` | `bf16`, `fp16`, `fp32`, or `fp8` (`fp8` requires torchao). Default: `bf16`. |
+| `--duration` | Run duration in seconds (default: 10). |
+| `--output-file` | Destination for results (`.md`, `.csv`, `.tsv`, `.jsonl`, `.jsonl.gz`). Default: `./gemm_report.md`. Use `-` or omit for Markdown on stdout. |
+
+**Example**
+
+```bash
+primus-cli direct -- benchmark gemm --M 8192 --N 8192 --K 8192 --dtype bf16 --duration 10 --output-file ./gemm_report.md
+```
+
+---
+
+### `gemm-dense`
+
+Dense GEMM workload using Llama-like shape parameters (model-derived GEMMs).
+
+| Argument | Description |
+|----------|-------------|
+| `--model` | Optional label (for example `Llama3.1_8B`). |
+| `--seqlen` | Sequence length (default: 2048). |
+| `--hidden-size` | Hidden size (default: 4096). |
+| `--intermediate-size` | FFN intermediate size (default: 11008). |
+| `--num-attention-heads` | Attention heads (default: 32). |
+| `--num-key-value-heads` | KV heads (default: 32). |
+| `--head-dim` | Per-head dimension (default: 128). |
+| `--vocab-size` | Vocabulary size (default: 32000). |
+| `--dtype` | `bf16`, `fp16`, `fp32`, or `fp8` (`fp8` requires torchao). Default: `bf16`. |
+| `--mbs` | Microbatch size (default: 1). |
+| `--duration` | Seconds per shape (default: 3). |
+| `--output-file` | Report path (default: `./gemm-dense_report.md`). |
+
+**Example**
+
+```bash
+primus-cli direct -- benchmark gemm-dense --model Llama3.1_8B --seqlen 4096 --dtype bf16
+```
+
+---
+
+### `gemm-deepseek`
+
+Dense GEMM workload using DeepSeek-style shapes (MoE / MLA-related dimensions).
+
+| Argument | Description |
+|----------|-------------|
+| `--model` | Label (for example `Deepseek_V2`, `Deepseek_V3`). |
+| `--seqlen` | Sequence length (default: 4096). |
+| `--hidden-size` | Hidden size (default: 4096). |
+| `--intermediate-size` | Dense FFN intermediate (default: 12288). |
+| `--kv-lora-rank` | KV LoRA rank (default: 512). |
+| `--moe-intermediate-size` | MoE expert intermediate (default: 1536). |
+| `--num-attention-heads` | Attention heads (default: 64). |
+| `--num-experts-per-tok` | Experts per token (default: 6). |
+| `--n-routed-experts` | Number of routed experts (default: 128). |
+| `--n-shared-experts` | Shared experts (default: 2). |
+| `--q-lora-rank` | Optional Q LoRA rank. |
+| `--qk-nope-head-dim`, `--qk-rope-head-dim`, `--v-head-dim` | Head dimensions for MLA-style attention (defaults: 128 / 64 / 128). |
+| `--vocab-size` | Vocabulary size (default: 128256). |
+| `--dtype` | `bf16` or `fp16` (default: `bf16`). |
+| `--mbs` | Microbatch size (default: 1). |
+| `--duration` | Seconds per shape (default: 3). |
+| `--output-file` | Report path (default: `./gemm-deepseek_report.md`). |
+| `--append` | Append to an existing report instead of overwriting. |
+
+**Example**
+
+```bash
+primus-cli direct -- benchmark gemm-deepseek --model Deepseek_V3 --dtype bf16 --append
+```
+
+---
+
+### `strided-allgather`
+
+Strided all-gather microbenchmark (useful for multi-rank communication patterns).
+
+| Argument | Description |
+|----------|-------------|
+| `--sizes-mb` | Comma-separated message sizes in MB per rank (default: `64,128,256`). |
+| `--stride` | Rank stride for group formation (default: 8). |
+| `--parallel` | Run multiple groups’ all-gathers in parallel. |
+| `--iters` | Timed iterations per size (default: 50). |
+| `--warmup` | Warmup iterations per size (default: 10). |
+| `--dtype` | `fp16`, `bf16`, or `fp32` (default: `bf16`). |
+| `--backend` | `nccl`, `gloo`, or `mpi` (default: `nccl`). |
+
+**Example**
+
+```bash
+primus-cli slurm srun -N 2 -- benchmark strided-allgather --sizes-mb 64,128 --stride 8 --iters 50
+```
+
+---
+
+### `rccl`
+
+RCCL collective benchmark: sweeps message sizes and reports bandwidth and latency statistics.
+
+| Argument | Description |
+|----------|-------------|
+| `--op` | One or more of: `all_reduce`, `broadcast`, `reduce_scatter`, `all_gather`, `alltoall` (default: `all_reduce`). |
+| `--sizes` | Explicit size list (for example `1K,2K,4K,8K,1M`). Overrides generated sweep. |
+| `--min-bytes` | Minimum message size for generated sweep (default: `1K`). |
+| `--max-bytes` | Maximum message size (default: `128M`). |
+| `--num-sizes` | Number of points in generated sweep (default: 12). |
+| `--scale` | `log2` or `linear` for generated sweeps (default: `log2`). |
+| `--dtype` | `bf16`, `fp16`, or `fp32` (default: `bf16`). |
+| `--warmup` | Warmup iterations (default: 20). |
+| `--iters` | Timed iterations (default: 100). |
+| `--repeat` | Repeat each `(op, size)` for stability (default: 1). |
+| `--aggregate-repeat` | Emit an extra summary row aggregating repeat runs. |
+| `--check` | Enable lightweight correctness checks. |
+| `--output-file` | Report path (`.md`, `.csv`, `.tsv`, `.jsonl`, `.jsonl.gz`; default: `./rccl_report.md`). |
+| `--append` | Append instead of overwrite. |
+| `--per-rank` | Per-rank summary lines. |
+| `--per-rank-file` | Path for per-rank stats (if empty, derived from `--output-file` with `_rank` suffix). |
+| `--per-iter-trace` | Emit per-iteration trace (can be large). |
+| `--trace-file` | Trace output path (if empty, derived from `--output-file`). |
+| `--trace-limit` | Max iterations to record per `(op, size)`; `0` means all. |
+| `--trace-ops` | Comma-separated ops to include in trace (empty = all). |
+| `--trace-sizes` | Comma-separated sizes to include in trace (empty = all). |
+| `--cluster` | Label for the report preamble. Defaults to `$PRIMUS_CLUSTER`, falling back to a built-in placeholder (`amd-aig-poolside`) when it is unset—set `PRIMUS_CLUSTER` or pass `--cluster` to record your own cluster name. |
+
+**Example**
+
+```bash
+primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M --dtype bf16
+```
+
+---
+
+## Understanding results
+
+- **GEMM suites** emit throughput-oriented metrics suitable for comparing dtypes, shapes, and durations across runs. Keep `duration` long enough to smooth variance on shared clusters.
+- **`rccl`** reports collective latency and bandwidth across a size sweep; use it to verify inter-node behavior and to compare against expected NIC bandwidth.
+- **Markdown / CSV / TSV / JSONL** output formats support post-processing in notebooks or CI; gzip JSONL is supported for large traces.
+
+---
+
+## Tips
+
+1. **Distributed initialization:** If jobs hang or report uninitialized distributed state, launch through `primus-cli` (`direct` / `container` / `slurm`) rather than calling Python entrypoints manually without the right environment.
+2. **Paths:** Prefer **absolute** paths for `--output-file` when using containers or Slurm so the working directory matches your expectations.
+3. **Multi-node:** Use your scheduler integration (`primus-cli slurm …`) so rank and address assignment matches your cluster.
+4. **Full cluster validation:** Combine targeted `benchmark` runs with [Preflight](./preflight.md) for host, GPU, network, and integrated perf checks.
+
+---
+
+## Related documentation
+
+- [Preflight diagnostics](./preflight.md)
+- [Memory and performance projection](./projection.md)
+- [Post-training workflows](./posttraining.md)
+- [Installation](../01-getting-started/installation.md)
diff --git a/docs/02-user-guide/cli-reference.md b/docs/02-user-guide/cli-reference.md
new file mode 100644
index 000000000..e7928ba8f
--- /dev/null
+++ b/docs/02-user-guide/cli-reference.md
@@ -0,0 +1,233 @@
+# CLI reference
+
+This section describes the unified Primus launcher (`runner/primus-cli`) and how it invokes the Python CLI (`primus/cli/main.py`). For deeper background, see [CLI architecture](../06-developer-guide/cli-architecture.md).
+
+---
+
+## Command structure
+
+```text
+primus-cli [global-options] <mode> [mode-args] -- [command]
+```
+
+- **Global options** go before the mode name and they affect configuration loading and logging for the whole run.
+- **Mode** is one of `direct`, `container`, or `slurm`.
+- **`--` (required)** separates launcher options from the Primus Python CLI. Everything after the first `--` is passed to `primus/cli/main.py` (or another script if you override it in direct mode).
+
+From the repository root, invoke the launcher as `./runner/primus-cli` (or install or link it as `primus-cli` on your `PATH`).
+
+---
+
+## Global options
+
+These flags are parsed in `runner/primus-cli` before the mode name is read and passed on to `runner/primus-cli-<mode>.sh`.
+
+| Option | Description |
+| --- | --- |
+| `--config FILE` | Load a YAML file for launcher defaults (see [Configuration precedence](#configuration-precedence-launcher-yaml)). |
+| `--debug` | Verbose logging; sets `PRIMUS_LOG_LEVEL=DEBUG`. |
+| `--dry-run` | Print the command that would run and exit without executing the mode script. |
+| `--version` | Print the CLI version and exit. |
+| `-h`, `--help` | Show top-level usage and exit. |
+
+Mode-specific help:
+
+```bash
+./runner/primus-cli direct --help
+./runner/primus-cli container --help
+./runner/primus-cli slurm --help
+```
+
+Primus Python CLI help (after `--`):
+
+```bash
+./runner/primus-cli direct -- --help
+./runner/primus-cli direct -- train --help
+./runner/primus-cli direct -- benchmark --help
+```
+
+---
+
+## Direct mode
+
+Run training, benchmarks, or diagnostics on the current host (or inside an environment you already prepared). GPU-specific tuning is applied via `runner/helpers/envs/<GPU_MODEL>.sh` when present.
+
+### Syntax
+
+```bash
+primus-cli direct [options] -- <command>
+```
+
+### Options
+
+| Option | Description |
+| --- | --- |
+| `--config FILE` | Launcher YAML (same resolution as [global `--config`](#configuration-precedence-launcher-yaml)). |
+| `--debug` | Debug logging for the direct launcher. |
+| `--dry-run` | Show the resolved command that would be launched without running training. |
+| `--single` | Run with `python3` instead of `torchrun` (single process). |
+| `--script PATH` | Python entry script (default: `primus/cli/main.py`). |
+| `--env KEY=VALUE` | Set an environment variable before launch (repeatable). A path without `=` is treated as an env file (`--env_file`), loaded later in the launch sequence. |
+| `--patch script.sh` | Run a shell snippet before the main script (repeatable). |
+| `--log_file PATH` | Redirect logs to a file. |
+| `--numa` | Force NUMA binding on. |
+| `--no-numa` | Force NUMA binding off. |
+
+### Distributed environment variables
+
+For multi-node or multi-process runs, set these via `export` or `--env`:
+
+| Variable | Role | Typical default |
+| --- | --- | --- |
+| `NNODES` | Number of nodes | `1` |
+| `NODE_RANK` | Rank of this node | `0` |
+| `GPUS_PER_NODE` | GPUs per node | `8` (see `runner/.primus.yaml` `direct.gpus_per_node`) |
+| `MASTER_ADDR` | Hostname or IP of rank 0 | `localhost` |
+| `MASTER_PORT` | TCP port for the process group | `1234` |
+
+---
+
+## Container mode
+
+Run the same Python CLI inside Docker or Podman with ROCm-oriented defaults from `runner/.primus.yaml`.
+
+### Syntax
+
+```bash
+primus-cli container [options] -- <command>
+```
+
+### Common options
+
+| Option | Description |
+| --- | --- |
+| `--image NAME` | Image tag (default from config: `rocm/primus:v26.4`). |
+| `--volume HOST[:CONTAINER]` | Bind mount (repeatable). |
+| `--env KEY=VALUE` | Pass into the **inner** `primus-cli direct` as `--env` (repeatable). |
+| `--device PATH` | Extra device nodes (repeatable; defaults include GPU/RDMA devices). |
+| `--name`, `--user`, `--network`, `--ipc` | Standard container runtime options. |
+| `--clean` | Remove all containers before launch. |
+| `--cpus N` | CPU limit. |
+| `--memory SIZE` | Memory limit (e.g. `128G`). |
+| `--shm-size SIZE` | Shared memory size. |
+| `--gpus N` | GPU limit (when using a runtime that supports this flag). |
+
+### Auto-mounted devices
+
+When using `runner/.primus.yaml`, the default container section includes:
+
+- `/dev/kfd`—ROCm kernel fusion driver
+- `/dev/dri`—GPU render nodes
+- `/dev/infiniband`—InfiniBand character devices (when present)
+
+### Environment forwarding
+
+`container.options.env` in `runner/.primus.yaml` lists **names** that are forwarded into the container as inner `--env` arguments when the variable is set in the host environment (for example `MASTER_ADDR`, `HF_TOKEN`, `NCCL_SOCKET_IFNAME`). The container script also auto-forwards host variables whose names start with `PRIMUS_`, `NCCL_`, `RCCL_`, `GLOO_`, `IONIC_`, or `HIPBLASLT_` when not already listed.
+
+---
+
+## Slurm mode
+
+Launch distributed jobs with `srun` or `sbatch`. The Slurm launcher builds `srun` or `sbatch` flags, merges them with `slurm.*` entries from the loaded YAML, then runs `runner/primus-cli-slurm-entry.sh` on allocated nodes.
+
+### Syntax
+
+```text
+primus-cli slurm [--config FILE] [--debug] [--dry-run] [srun|sbatch] [SLURM_FLAGS...] -- <command>
+```
+
+| Part | Meaning |
+| --- | --- |
+| First `--` | Separates Slurm launcher flags from the Primus Python CLI command (for example `train pretrain ...`). |
+| Default launcher | If you omit `srun` and `sbatch`, **`srun` is used** (`LAUNCH_CMD` in `runner/primus-cli-slurm.sh`). |
+
+### Examples
+
+```bash
+# Interactive multi-node training
+./runner/primus-cli slurm srun -N 4 -p gpu -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+
+# Batch job
+./runner/primus-cli slurm sbatch -N 8 -t 8:00:00 -o train.log -- train pretrain --config exp.yaml
+```
+
+On each node, `primus-cli-slurm-entry.sh` sets `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `MASTER_ADDR`, and `MASTER_PORT` from Slurm and invokes `primus-cli-container.sh` with matching `--env` injections (see `runner/primus-cli-slurm-entry.sh`). Container options such as `--image` should come from `runner/.primus.yaml` or the launcher config file rather than appearing as an inner `container` command after the Slurm separator.
+
+---
+
+## Python subcommands (after `--`)
+
+These run under `primus/cli/main.py` unless you change `--script` in direct mode.
+
+| Subcommand | Purpose |
+| --- | --- |
+| `train pretrain --config <yaml>` | Pretraining (Megatron-LM, TorchTitan, MaxText, Megatron Bridge, etc., per configuration YAML). |
+| `train posttrain --config <yaml>` | Post-training (SFT or LoRA-style workflows; same top-level flags as pretrain in the parser). |
+| `benchmark <suite> [args]` | Performance microbenchmarks (see table below). |
+| `preflight [--host] [--gpu] [--network] [--perf-test]` | Cluster and node diagnostics. |
+| `projection memory --config <yaml>` | Memory estimation from a merged config. |
+| `projection performance --config <yaml>` | Performance projection from a merged config. |
+| `projection both --config <yaml>` | Single benchmark → both performance and memory projections (cluster sizing). |
+
+### Benchmark suites
+
+Implemented in `primus/cli/subcommands/benchmark.py`:
+
+| Suite | Notes |
+| --- | --- |
+| `gemm` | General GEMM microbenchmark. |
+| `gemm-dense` | Dense GEMM variant. |
+| `gemm-deepseek` | DeepSeek-style dense GEMM. |
+| `strided-allgather` | Communication microbenchmark. |
+| `rccl` | RCCL collective microbenchmark. |
+
+The same file also registers an `attention` suite for attention microbenchmarks.
+
+---
+
+## Configuration precedence (launcher YAML)
+
+Resolution is implemented in `runner/lib/config.sh` (functions `resolve_config_file` and `load_config_auto`):
+
+1. **`--config FILE`** on the command line (if given).
+2. **`~/.primus.yaml`** if it exists.
+3. **`runner/.primus.yaml`** (system default).
+
+Within a chosen file, nested keys follow normal YAML structure. Slurm and container scripts merge CLI flags with their sections so that **explicit CLI arguments override file values** where applicable.
+
+**Note:** This precedence applies to the **shell launcher** YAML. Training YAML merge order for configurations is documented in [Configuration system](configuration-system.md).
+
+---
+
+## Common examples
+
+| Goal | Example |
+| --- | --- |
+| Direct pretrain | `./runner/primus-cli direct -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` |
+| Direct GEMM | `./runner/primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096` |
+| Container pretrain | `./runner/primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml` |
+| Slurm training | `./runner/primus-cli slurm srun -N 4 -- train pretrain --config exp.yaml` |
+| Preflight (fast) | `./runner/primus-cli slurm srun -N 4 -- preflight --host --gpu --network` |
+| Inspect launch command | `./runner/primus-cli --dry-run direct -- train pretrain --config exp.yaml` |
+| Dry-run Slurm | `./runner/primus-cli --dry-run slurm srun -N 2 -- train pretrain --config exp.yaml` |
+
+---
+
+## Exit codes
+
+From `runner/primus-cli`:
+
+| Code | Meaning |
+| --- | --- |
+| 0 | Success |
+| 1 | Library or dependency failure |
+| 2 | Invalid arguments or configuration |
+| 3 | Runtime execution failure |
+
+---
+
+## Related documentation
+
+- [Getting started: Quickstart](../01-getting-started/quickstart.md): installation and first steps
+- [Configuration system](configuration-system.md): YAML configuration model, presets, overrides, inheritance
+- [Pretraining](pretraining.md): pretraining workflows and backend notes
diff --git a/docs/02-user-guide/configuration-system.md b/docs/02-user-guide/configuration-system.md
new file mode 100644
index 000000000..90b8cb6fe
--- /dev/null
+++ b/docs/02-user-guide/configuration-system.md
@@ -0,0 +1,170 @@
+# Configuration system
+
+Primus experiments are described in YAML. The loader resolves **environment variables**, **`extends:` inheritance**, and **module/model/platform presets** before training starts. This document focuses on the Python configuration pipeline (`primus/core/config/` and `primus/core/launcher/parser.py`).
+
+**Related documentation**
+
+| Topic | Location |
+| --- | --- |
+| CLI launcher and `--config` | [CLI Reference](cli-reference.md) |
+| Backend parameter references | [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md), [MaxText parameters](../03-configuration-reference/maxtext-parameters.md) |
+
+---
+
+## Overview: Three-layer YAML
+
+A typical experiment ties together:
+
+1. **Experiment YAML**—your run: identity, workspace, and a `modules` section naming framework presets plus overrides.
+2. **Module preset**—training defaults for a backend (optimizer, schedule, parallelism hooks) under `primus/configs/modules/<framework>/`.
+3. **Model preset**—architecture and tokenizer metadata under `primus/configs/models/<framework>/`.
+
+All of these are **deep-merged** (see `primus/core/config/yaml_loader.py` and `primus/core/config/merge_utils.py`). A **platform preset** (`primus/configs/platforms/`) maps distributed environment variable names and logging defaults; if omitted, the parser injects `platform_azure.yaml` (see `PrimusParser.parse_platform` in `primus/core/launcher/parser.py`).
+
+---
+
+## Experiment config structure
+
+```yaml
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:my_experiment}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron          # backend name (megatron, torchtitan, maxtext, megatron_bridge, …)
+    config: pre_trainer.yaml     # module preset file under primus/configs/modules/<framework>/
+    model: llama3_8B.yaml        # model preset file under primus/configs/models/<framework>/
+    overrides:                   # training overrides (deep-merged last)
+      train_iters: 50
+      micro_batch_size: 4
+```
+
+Required top-level keys are validated in `PrimusParser.parse_meta_info`: `work_group`, `user_name`, `exp_name`, `workspace`.
+
+---
+
+## Module presets
+
+- **Location:** `primus/configs/modules/<framework>/` (for example `primus/configs/modules/megatron/pre_trainer.yaml`).
+- **Purpose:** Default training behavior (iterations, batching, optimizer, logging, parallelism-related flags for that backend).
+- **Inheritance:** Use `extends:` to compose files in the same directory (or relative paths). Multiple entries merge in order; the **current file wins** on conflicts (`_apply_extends` in `yaml_loader.py`).
+
+Example chain (excerpt): `pre_trainer.yaml` extends `trainer_base.yaml`, which extends `../module_base.yaml` and other shared fragments (`primus/configs/modules/megatron/trainer_base.yaml`).
+
+---
+
+## Model presets
+
+- **Location:** `primus/configs/models/<framework>/` (for example `primus/configs/models/megatron/llama3_8B.yaml`).
+- **Purpose:** Architecture dimensions, tokenizer identifiers, and other model metadata.
+- **Inheritance:** Same `extends:` mechanism as modules (for example `llama3_8B.yaml` → `llama3_base.yaml` → …).
+
+---
+
+## Platform presets
+
+- **Location:** `primus/configs/platforms/` (for example `primus/configs/platforms/platform_azure.yaml`).
+- **Purpose:** Names of environment variables used for distributed launch (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, …) and defaults such as `master_sink_level` and `workspace`.
+- **Default:** If the experiment omits `platform`, the parser sets `config: platform_azure.yaml` (`primus/core/launcher/parser.py`).
+
+---
+
+## Environment variable substitution
+
+`primus/core/config/yaml_loader.py` expands:
+
+| Pattern | Behavior |
+| --- | --- |
+| `${VAR}` | **Required.** Raises if `VAR` is unset. |
+| `${VAR:default}` | Uses `default` when `VAR` is unset. |
+
+After substitution, purely numeric strings may be converted to `int` or `float`.
+
+---
+
+## `extends:` inheritance
+
+For each YAML file:
+
+1. Each path in `extends:` is resolved **relative to the directory of the current file**.
+2. Presets are loaded recursively (each may have its own `extends:`).
+3. Merge order: earlier presets in the list are merged first; **later presets override earlier ones**; the **current file overrides all** (`_apply_extends`).
+
+`PresetLoader.load` (`primus/core/config/preset_loader.py`) resolves `primus/configs/<config_type>/<framework>/<name>.yaml` and runs the same `parse_yaml` pipeline.
+
+---
+
+## CLI overrides (training)
+
+After the main arguments are parsed, unknown tokens are interpreted as **key=value overrides** and deep-merged into the active training module namespace (`module_cfg.params`). The core runtime applies them in `PrimusRuntime._apply_overrides` (`primus/core/runtime/train_runtime.py`) using `parse_cli_overrides` (`primus/core/utils/arg_utils.py`) followed by `deep_merge`. Both `key=value` and `--key value` forms are accepted. Unknown keys are merged in (not rejected) and forwarded to the backend; the stricter key-existence check in `parse_args` / `_check_keys_exist` (`primus/core/launcher/parser.py`) belongs to a legacy path that the `train` subcommand does not exercise.
+
+Example (conceptual):
+
+```bash
+./runner/primus-cli direct -- train pretrain --config exp.yaml \
+  --train_iters 100 --micro_batch_size 2
+```
+
+---
+
+## Merge priority (training config)
+
+The effective ordering of training parameters is:
+
+1. **CLI overrides** (key=value after the main `train` arguments)—highest. Applied last by the runtime (`PrimusRuntime._apply_overrides`), after preset and experiment merging.
+2. **`modules.pre_trainer.overrides`** in the experiment YAML (applied in `PrimusParser.parse_trainer_module`).
+3. **Module preset with model preset additions**: the module preset is loaded first, then the model preset is merged in with `allow_override=False`, so duplicate top-level module keys are preserved while non-duplicate model keys are added (`merge_namespace` in `parse_trainer_module`).
+4. **Preset chains** via `extends:` inside those files—base layers first, specialized layers later, file body last.
+
+A concise mental model:
+
+**CLI overrides > experiment `overrides` > module preset with model preset additions > each preset's own `extends:` chain > shared bases such as `module_base.yaml`.**
+
+---
+
+## Config resolution walkthrough: `llama2_7B-BF16-pretrain.yaml`
+
+Example experiment: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml`.
+
+1. **Load experiment**—`parse_yaml` reads the file; `${PRIMUS_*:…}` placeholders resolve.
+2. **Meta**—`work_group`, `user_name`, `exp_name`, `workspace` are checked.
+3. **Platform**—If not present, default `platform_azure.yaml` loads from `primus/configs/platforms/`.
+4. **Module preset**—`PresetLoader.load("pre_trainer.yaml", "megatron", "modules")` loads `primus/configs/modules/megatron/pre_trainer.yaml` and applies its `extends:` chain (for example `trainer_base.yaml` → …).
+5. **Model preset**—`PresetLoader.load("llama2_7B.yaml", "megatron", "models")` loads `primus/configs/models/megatron/llama2_7B.yaml` (which extends `llama2_base.yaml` → `llama_base.yaml` → …).
+6. **Merge**—Module and model namespaces are merged for `pre_trainer`.
+7. **Experiment overrides**—Keys under `modules.pre_trainer.overrides` in the example file (for example `mock_data: true`, parallelism, LR) are applied on top.
+8. **CLI overrides**—Any key=value pairs from the command line are merged last.
+
+---
+
+## How to write a new config for a new model
+
+1. **Add or reuse a model preset** under `primus/configs/models/<framework>/`, using `extends:` from the closest existing architecture (for example copy `llama3_8B.yaml` and adjust hidden size, layers, tokenizer).
+2. **Point an experiment YAML at it**—set `modules.pre_trainer.framework`, `config: <module_preset>.yaml`, and `model: <your_model>.yaml`.
+3. **Set overrides** in the experiment file for run-specific values (batch sizes, paths, `mock_data`, parallelism). Prefer small experiment files that reference presets instead of duplicating hundreds of keys.
+4. **Validate** with `--dry-run` and by tracing the referenced experiment, module, and model presets (see below).
+5. **Optional:** add an example under `examples/<backend>/configs/<platform>/` for others to copy.
+
+---
+
+## Debugging config issues
+
+| Technique | What it does |
+| --- | --- |
+| `--export_config PATH` | Parsed by the training config parser, but the default core `PrimusRuntime` path does not currently write the resolved YAML. Treat this as legacy/future functionality unless your deployment implements it. |
+| `./runner/primus-cli --dry-run …` | Shows the launcher command without executing (shell layer). |
+| `--debug` | Enables verbose logging for launcher and Python (`PRIMUS_LOG_LEVEL=DEBUG`). |
+| Inspect presets | Open the resolved `extends:` chain under `primus/configs/modules/` and `primus/configs/models/` for the framework you use. |
+
+If `${VAR}` substitution fails, set the variable or switch to `${VAR:default}` in the YAML.
+
+---
+
+## Cross-references
+
+- Default launcher YAML: `runner/.primus.yaml`
+- YAML loader (env + extends): `primus/core/config/yaml_loader.py`
+- Parser and merge: `primus/core/launcher/parser.py`
+- Preset paths: `primus/core/config/preset_loader.py`
diff --git a/docs/02-user-guide/node-smoke-test-instruction.md b/docs/02-user-guide/node-smoke-test-instruction.md
new file mode 100644
index 000000000..41f163be1
--- /dev/null
+++ b/docs/02-user-guide/node-smoke-test-instruction.md
@@ -0,0 +1,452 @@
+# Node-smoke test instruction
+
+A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. It produces a **single PASS/FAIL verdict per node** plus SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`.
+
+Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node surfaces as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers.
+
+- **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves the distributed env, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` (bare `srun` + `direct`) is equivalent and handy for ad-hoc runs.
+- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)).
+
+---
+
+## 1. What it does
+
+Node-smoke answers one question fast: **"which nodes are healthy enough to run anything?"** Because training jobs allocate whole nodes, a single degraded GPU (or NIC, or wedged driver) takes an entire node out of rotation. Node-smoke checks each node independently and emits a per-node verdict plus a ready-to-use exclude list, so you can prune broken nodes before committing a large job to a global rendezvous.
+
+It deliberately does **not** measure cross-node bandwidth — that's what [`preflight`](./preflight.md) is for.
+
+---
+
+## 2. How it works
+
+- **Per-node and independent** — every node runs the checks on its own. No `MASTER_ADDR`, no `MASTER_PORT`, no global `torch.distributed` rendezvous, so a stuck node cannot wedge its peers.
+- **Per-GPU isolation** — each GPU's checks run in their own Python subprocess with a hard timeout (`--per-gpu-timeout-sec`, default 15 s). A stuck `torch.cuda.set_device()` (which `signal.alarm` cannot interrupt because it sits inside a driver syscall) is `SIGKILL`'d from the parent without affecting the rest of the node's checks.
+- **Local-only RCCL** — the optional Tier 2 all-reduce uses `torch.multiprocessing.spawn` over `tcp://127.0.0.1`. No cross-node communication.
+- **Rank-0 aggregation** — `NODE_RANK==0` polls for the expected number of per-node JSONs (with a timeout), computes cluster-wide drift, writes the Markdown report + pass/fail lists, and returns non-zero if any node FAILs or never reports.
+
+---
+
+## 3. Prerequisites
+
+| Prerequisite                        | How                                                                                                                                                                                                                                                                                                                                                                                              |
+| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| Python venv on a shared filesystem  | Same venv used by `primus-cli direct -- preflight` (see [`preflight-without-container.md`](./preflight-without-container.md) §2).                                                                                                                                                                                                                                                                 |
+| `VENV_ACTIVATE` exported            | `export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate` (optional inside the container path).                                                                                                                                                                                                                                                                                                 |
+| Inside an existing SLURM allocation | One task per node. Recommended: `runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 -- direct -- node_smoke ...`. Equivalent bare form: `srun ... --ntasks-per-node=1 runner/primus-cli direct -- node_smoke ...`. Either way the `direct -- node_smoke` path auto-selects `--single`, so each task spawns one Python process and per-GPU subprocesses are launched internally. |
+
+No `MASTER_ADDR`, no `MASTER_PORT`, no global rendezvous required.
+
+---
+
+## 4. Quick start
+
+**Git clone the Primus repository to a shared filesystem that all nodes can read.**
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+```
+
+**Note: remember to set up the Python virtual environment and NCCL / fabric environment variables as described in [§3 Prerequisites](#3-prerequisites).**
+
+> ⚠ **Set the NCCL / RCCL environment first** if you plan to run with `--tier2-perf` (the local 8-GPU RCCL all-reduce). Even though the smoke test never opens a cross-node rendezvous, the Tier 2 RCCL step calls `dist.init_process_group(backend="nccl", ...)`, and RCCL **enumerates every transport at init** (XGMI / PCIe P2P + IB + sockets). A misconfigured `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME` / `NCCL_IB_GID_INDEX` can stall init or make the all-reduce silently fall back to a slow path. The launcher's `base_env.sh` auto-detects these, **but auto-detect sometimes picks the wrong values inside a container** (devices masked by the network namespace, frontend NICs picked up instead of fabric NICs, etc.), so check them and set them explicitly if auto-detection is wrong.
+>
+> Minimum-viable checklist before running with `--tier2-perf`:
+>
+> ```bash
+> # Pin the RDMA / RoCE training NICs the container can actually see.
+> # On a bare-metal host the auto-detect in base_env.sh usually picks
+> # the right set; inside a container or on a multi-role node, list
+> # them explicitly. Use the same set you would pass to a training job.
+> export NCCL_IB_HCA="rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7"
+>
+> # Pick the RoCE v2 GID index for your fabric:
+> #   - Mellanox / Broadcom: typically 3 (base_env.sh default).
+> #   - Pensando Pollara (AINIC): 1.
+> export NCCL_IB_GID_INDEX=3
+>
+> # The bootstrap socket interface. Auto-detect prefers the first
+> # non-loopback interface from `hostname -I`; override when that
+> # picks a frontend NIC instead of the data-plane interface.
+> export NCCL_SOCKET_IFNAME=eno0
+> export GLOO_SOCKET_IFNAME=eno0
+> ```
+>
+> See [`preflight-without-container.md` §4 Cluster-specific NCCL configuration](./preflight-without-container.md#4-cluster-specific-nccl-configuration) for the canonical Broadcom / Pensando Pollara values (the same `NCCL_*` set is used by both tools). If you skip `--tier2-perf`, the RCCL step is not executed and none of the above applies — Tier 1 (host limits, RDMA roll-call, leaked-process detection, etc.) does not depend on RCCL.
+>
+> Quick verification: `runner/primus-cli direct --dry-run -- node_smoke --tier2-perf` prints the resolved `NCCL_*` block under "Environment Variables" so you can confirm the values before launching for real.
+
+Recommended — through the `primus-cli slurm srun` wrapper (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK`, applies `slurm.*` config defaults):
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Basic Tier 1 check (~5 s/GPU, ~30 s total)
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke
+
+# Tier 1 + Tier 2 perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL)
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf
+
+# Then re-run training, excluding any node the smoke test failed:
+srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job
+```
+
+Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke
+
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf
+```
+
+Single-node sanity check (no SLURM):
+
+```bash
+runner/primus-cli direct -- node_smoke
+```
+
+> **Both forms produce the same workload.** The wrapper form is recommended because it resolves the distributed env once on the launching node and propagates it via `--env`, and applies any `slurm.`* config defaults (partition / time / etc.). The `direct` keyword between the two `--`s is **mandatory** to take the no-container path — without it the wrapper routes through the container path. See [`preflight-without-container.md` § Wrapper vs. bare-srun](./preflight-without-container.md#wrapper-vs-bare-srun) for the precedence table.
+
+---
+
+## 5. What's checked
+
+A `level='fail'` finding in any check FAILs the node. Everything else is reported as info / warn.
+
+### Tier 1 — always runs (~5 s/GPU)
+
+**Per-GPU liveness** (each GPU in its own subprocess with a hard timeout):
+
+- `torch.cuda.set_device(i)` — proves the device is bindable (a stale / wedged GPU often fails here).
+- 256 MB allocation.
+- Tiny 2048² bf16 GEMM with an `isfinite()` check on the result.
+
+**Host / GPU / network inventory** (no rendezvous):
+
+- **dmesg recent-error scan** — greps the last `--dmesg-minutes` (default 15) of `dmesg` for known patterns (`xid`, `gpu reset`, `hung_task`, `mce:`, `amdgpu.*error`, ...). Matches are surfaced in the report.
+- **A. Software-stack fingerprint** — kernel / OS / Python, ROCm version, amdgpu kernel-module version, PyTorch / `torch.version.hip` / RCCL versions, and per-IB-device firmware + HCA model. Used for cluster drift detection.
+- **B. NIC / RDMA roll-call** — per-port state read from `/sys/class/infiniband` (works inside containers; no `ibv_devinfo` / `ibstat` dependency). Many clusters expose more RDMA ports than the training job uses, so the hard-fail rules only run against the *training-NIC* subset, selected by this precedence:
+  1. `--rdma-nic-allowlist` (`NCCL_IB_HCA` syntax: comma-separated `device[:port]`, `^...` denylist, `=dev` exact-match).
+  2. `NCCL_IB_HCA` env (same syntax) — so the smoke test and the training launch agree by construction.
+  3. Heuristic: auto-exclude any port whose `phys_state` is `Disabled` or `Sleep` (admin-disabled).
+  4. Fallback: every IB port must be ACTIVE / LinkUp.
+
+  **Hard-fail rules** (on the included set only): port not ACTIVE / not LinkUp; active port with zero RoCE v2 GIDs (RoCE) or zero valid GIDs (IB); included-NIC count ≠ `--expected-rdma-nics N` (when set). If *every* discovered port gets excluded, the node still fails — a node with zero training NICs cannot participate in inter-node training. Excluded ports stay visible in the report for diagnostics but don't contribute to the FAIL signal.
+- **C. Host limits / system tunables** — `RLIMIT_MEMLOCK` below `--ulimit-l-min-gb` (default 32 GiB) → "RDMA pin will fail under load"; `/dev/shm` below `--shm-min-gb` (default 8 GiB) → "NCCL shared-mem may fail". NUMA node count, CPU count, and `cpu0` scaling governor are collected for drift detection.
+- **Foreign / leaked process detection** — foreign PIDs holding a GPU FAIL the node by default (the most common cause of "training fails to launch on a healthy-looking node"). Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`. See the [container note in §6.5](#65-allow--extend-the-foreign-process-whitelist) — running inside a container almost always needs `--allow-foreign-procs`.
+- **rocm-smi self-latency** — a `rocm-smi --version` call slower than `--rocm-smi-timeout-sec` (default 5 s) is a node FAIL; a wedging amdgpu driver typically hangs `rocm-smi` for 30–60 s before the GPU itself stops responding.
+
+### Tier 2 — optional perf sanity (`--tier2-perf`)
+
+Per-GPU steady-state metrics, with iteration counts aligned to the preflight `--quick` preset so smoke and preflight numbers are directly comparable. It's a single switch — you cannot enable just one half.
+
+- **GEMM TFLOPS** — 8192³ bf16 `torch.matmul`; FAIL below `--gemm-tflops-min` (default 600).
+- **HBM GB/s** — 512 MB device-to-device `torch.Tensor.copy_` (counts read + write); FAIL below `--hbm-gbs-min` (default 2000; a healthy MI300X is ≈ 4500–5000).
+- **Local 8-GPU RCCL all-reduce GB/s** — algorithmic bandwidth `2·S·(P-1)/P / t / 1e9` at `--rccl-size-mb` (default 64 MB); FAIL below `--rccl-gbs-min` (default 100). Local only, no cross-node traffic.
+
+---
+
+## 6. Examples (by configuration knob)
+
+> **Convention used below.** The examples are written with bare `srun` for brevity. Anywhere you see `srun <flags> runner/primus-cli direct -- node_smoke ...`, the equivalent wrapper form is `runner/primus-cli slurm srun <flags> -- direct -- node_smoke ...`. Pick whichever matches your habits; both target the same launcher.
+
+### 6.1 Hard-fail on partial NIC enumeration
+
+Catches "7 of 8 RDMA NICs visible" — a common cause of crashes after RoCE init. The count is compared against the *training-NIC* set (after the selector chain), so frontend / storage RoCE NICs do not inflate it.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf --expected-rdma-nics 8
+```
+
+### 6.2 Pin the training-NIC set explicitly
+
+When auto-detection picks the wrong ports, name the training NICs directly (otherwise `NCCL_IB_HCA` env is used; otherwise admin-disabled ports are auto-excluded):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf \
+    --rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1,rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1'
+```
+
+### 6.3 Tighten Tier 2 perf thresholds
+
+Reject GPUs that come in below your acceptance bar. Defaults: GEMM 600 TFLOPS, HBM 2000 GB/s, local RCCL 100 GB/s.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf \
+        --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180
+```
+
+### 6.4 Tighten host limits
+
+Fail nodes whose `RLIMIT_MEMLOCK` or `/dev/shm` is too small for production training.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke \
+        --ulimit-l-min-gb 64 --shm-min-gb 16
+```
+
+### 6.5 Allow / extend the foreign-process whitelist
+
+By default, leaked / foreign processes holding a GPU FAIL the node. Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`.
+
+```bash
+# Add a site-specific monitoring agent to the whitelist
+srun ... runner/primus-cli direct -- node_smoke \
+    --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,my-monitor
+
+# Don't fail at all on foreign processes (still reported in the markdown)
+srun ... runner/primus-cli direct -- node_smoke --allow-foreign-procs
+```
+
+> ⚠ **Running node_smoke inside a container almost always trips this check.** `amd-smi process --json` reports `name="N/A"` for kernel/system PIDs like `gpuagent` whose `/proc/<host_pid>/comm` it cannot read, and the fallback name resolution inside `node_smoke` also fails because the container's `/proc` typically does not expose host PIDs (private PID namespace without `--pid=host`, or a `hidepid=2` mount). The unresolved name doesn't match the allowlist, so the check fires and the node FAILs — even though the only "foreign" processes are well-known system daemons holding zero HBM.
+>
+> **In the container path, pass `--allow-foreign-procs`:**
+>
+> ```bash
+> srun ... runner/primus-cli direct -- node_smoke --tier2-perf --allow-foreign-procs
+> ```
+>
+> The processes are still listed in `smoke_report.md` under "Busy GPUs / leaked processes", so a real leak is still visible; only the FAIL verdict is downgraded.
+>
+> **Narrower alternative** — add the literal sentinel `N/A` to the allowlist so the check still catches leaks with resolvable names (e.g. a leftover `python` rank):
+>
+> ```bash
+> srun ... runner/primus-cli direct -- node_smoke --tier2-perf \
+>     --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,N/A
+> ```
+>
+> Name resolution runs first, so whenever a real name *can* be resolved it overrides `N/A` and the normal allowlist applies — the `N/A` entry only matches PIDs whose name genuinely could not be recovered.
+>
+> **Root-cause fix** (preferred long-term): grant the container access to host PIDs so names resolve and the report shows `gpuagent` etc. instead of `N/A`. Typical fixes: launch with `--pid=host` (Docker / Podman); mount `/proc` without `hidepid=2`; or loosen `ptrace_scope` / grant `CAP_SYS_PTRACE`.
+
+### 6.6 Require specific tools
+
+Make missing CLI tools a hard FAIL (default: warn-only).
+
+```bash
+srun ... runner/primus-cli direct -- node_smoke --require-tools amd-smi,rocm-smi,lsof
+```
+
+### 6.7 Skip dmesg scan (containers with no privileges)
+
+```bash
+srun ... runner/primus-cli direct -- node_smoke --skip-dmesg
+```
+
+### 6.8 Custom dump path
+
+Keep one report per smoke run instead of overwriting the default location.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf \
+        --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+### 6.9 Re-aggregate from existing per-node JSONs (no re-run)
+
+The primus-cli wrapper always runs both phases (per-node run + rank-0 aggregate). To *only* re-render the report from JSONs collected earlier, use the standalone `aggregate` subcommand — it reads the existing `<dump>/smoke/*.json` without re-running the per-node step:
+
+```bash
+python -m primus.tools.preflight.node_smoke aggregate \
+    --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5
+```
+
+### 6.10 Silent mode (for CI)
+
+Suppresses wrapper stdout, but the **final report path is still printed** and stderr / exit code are preserved.
+
+```bash
+srun ... runner/primus-cli direct --silent -- node_smoke --tier2-perf
+```
+
+### 6.11 Combined "production-ready screen"
+
+A representative one-shot for a production cluster screen:
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct --silent -- node_smoke --tier2-perf \
+        --expected-rdma-nics 8 \
+        --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 \
+        --ulimit-l-min-gb 64 --shm-min-gb 16 \
+        --require-tools amd-smi,rocm-smi,lsof \
+        --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+---
+
+## 7. Outputs
+
+All written under `--dump-path` (default `output/preflight/`).
+
+| File                      | Purpose                                                                                                              |
+| ------------------------- | -------------------------------------------------------------------------------------------------------------------- |
+| `smoke/<short-host>.json` | Per-node verdict + every collected metric. One file per node.                                                        |
+| `smoke_report.md`         | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail).                     |
+| `passing_nodes.txt`       | Newline-separated short hostnames. Pipe into `srun --nodelist=`.                                                     |
+| `failing_nodes.txt`       | Newline-separated short hostnames. Pipe into `srun --exclude=`.                                                      |
+| `expected_nodes.txt`      | Auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`. Lets the report name nodes that never reported. |
+
+Read the cluster verdict at a glance:
+
+```bash
+head -10 output/preflight/smoke_report.md
+```
+
+Feed bad nodes into a re-run:
+
+```bash
+srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job
+```
+
+---
+
+## 8. Understanding the report
+
+`smoke_report.md` renders in a stable order. Each section short-circuits to a placeholder (e.g. `*All nodes match.*`, `*No NIC issues.*`) on a healthy cluster, so a clean report stays short. In order:
+
+1. **Status table** — one row per node: `node_rank`, hostname, PASS/FAIL, duration, top fail reason.
+2. **Stack drift across cluster** — per fingerprint key, outliers vs the cluster majority (catches "1 of N nodes on a different RCCL build").
+3. **NIC firmware drift across cluster** — per-IB-device firmware drift.
+4. **NIC / RDMA roll-call issues** — every offending node + port (included set only).
+5. **NIC port-count summary** — cluster-majority *training-NIC* count and any node that disagrees (catches partial-NIC degradation even without `--expected-rdma-nics`).
+6. **NIC excluded ports (informational)** — ports the selector chain dropped, grouped by source. Does not contribute to FAIL.
+7. **Host limits issues** — per-node hard-limit violations.
+8. **GPU visibility issues** — nodes where torch couldn't see the GPUs, or amd-smi sees more GPUs than torch (stale ROCm / wedged driver).
+9. **GPU low-level outliers (PCIe link / HBM)** — per-GPU outliers vs the cluster majority on PCIe width/speed and HBM total.
+10. **XGMI link issues** — any non-XGMI GPU pair (intra-node collectives silently fall back to PCIe).
+11. **Cluster clock + time daemons** — wall-clock spread plus per-node time-daemon health.
+12. **Tooling self-latency (`rocm-smi --version`)** — slow / timed-out tool calls (precursor to a wedged driver).
+13. **Tooling availability** — inventory of `amd-smi` / `rocm-smi` / `lsof` per node, plus which Tier 1 checks have no working tool on each node.
+14. **Busy GPUs / leaked processes** — foreign PIDs holding GPUs at smoke start.
+15. **GPU pre-touch HBM usage outliers** — GPUs with non-trivial HBM in use *before* smoke touched the device.
+16. **GPU compute-activity outliers** — GPUs above `--gpu-activity-warn-pct` at smoke start (warn-only).
+17. **Tier 2 perf summary** (only when at least one node ran Tier 2) — per-node GEMM TFLOPS / HBM GB/s as `min / median / max`, plus local RCCL GB/s.
+18. **Failing nodes — full reasons** (only when there are failing nodes) — every fail reason, expanded per node.
+
+---
+
+## 9. Configuration reference
+
+### 9.1 Common knobs (cheat sheet)
+
+| Flag                     | Default                                          | When you'd change it                                                                  |
+| ------------------------ | ------------------------------------------------ | ------------------------------------------------------------------------------------- |
+| `--tier2-perf`           | off                                              | Always on for production screens — adds GEMM TFLOPS, HBM GB/s, local RCCL all-reduce. |
+| `--gemm-tflops-min N`    | 600                                              | Site-specific acceptance bar.                                                         |
+| `--hbm-gbs-min N`        | 2000                                             | Site-specific acceptance bar (MI300X healthy ≈ 4500–5000).                            |
+| `--rccl-gbs-min N`       | 100                                              | Site-specific acceptance bar.                                                         |
+| `--expected-rdma-nics N` | unset                                            | Hard-fail on partial NIC enumeration.                                                 |
+| `--ulimit-l-min-gb GB`   | 32                                               | Raise for production training profiles.                                               |
+| `--shm-min-gb GB`        | 8                                                | Raise for large-batch / many-rank profiles.                                           |
+| `--allow-foreign-procs`  | off                                              | Co-tenant clusters, shared GPUs, or the container path (see [§6.5](#65-allow--extend-the-foreign-process-whitelist)). |
+| `--allowed-procs LIST`   | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Add site-specific monitoring agents.                                                  |
+| `--require-tools LIST`   | `""`                                             | Fail-fast if a CLI tool is missing in PATH.                                           |
+| `--skip-dmesg`           | off                                              | Inside unprivileged containers.                                                       |
+| `--dump-path DIR`        | `output/preflight`                               | Archive each run separately.                                                          |
+| `--silent` (launcher)    | off                                              | CI / scripted runs.                                                                   |
+
+### 9.2 Full `node_smoke` (per-node) flags
+
+Authoritative source: `python -m primus.tools.preflight.node_smoke run --help`.
+
+| Flag | Default | Purpose |
+|---|---|---|
+| `--dump-path` | `output/preflight` | Output directory. |
+| `--expected-gpus N` | auto | Override GPU count (auto-detected from `LOCAL_WORLD_SIZE` / `GPUS_PER_NODE` / `torch.cuda.device_count()`). |
+| `--per-gpu-timeout-sec` | 15 | Hard timeout per per-GPU subprocess. |
+| `--tier2-perf` | off | Enable Tier 2 perf sanity (per-GPU GEMM TFLOPS + HBM GB/s + node-local RCCL all-reduce). Single switch. |
+| `--gemm-tflops-min` | 600 | Tier 2 GEMM threshold. |
+| `--hbm-gbs-min` | 2000 | Tier 2 HBM threshold. |
+| `--rccl-size-mb` | 64 | Local RCCL message size. |
+| `--rccl-gbs-min` | 100 | Local RCCL bandwidth threshold. |
+| `--rccl-timeout-sec` | 120 | Hard timeout for the RCCL phase. |
+| `--skip-dmesg` | off | Skip dmesg scan (e.g. inside containers). |
+| `--dmesg-minutes` | 15 | dmesg `--since` window. |
+| `--expected-rdma-nics N` | auto (report-only) | When set, a mismatch between the included (training-NIC) count and N becomes a node FAIL. |
+| `--rdma-nic-allowlist LIST` | unset | Explicit training-NIC selector in `NCCL_IB_HCA` syntax (`device[:port],...`, `^...` denylist, `=dev` exact-match). Wins over `NCCL_IB_HCA` env. When neither is set, ports whose `phys_state` is `Disabled` / `Sleep` are auto-excluded. |
+| `--ulimit-l-min-gb GB` | 32 | `RLIMIT_MEMLOCK` threshold (0 disables). |
+| `--shm-min-gb GB` | 8 | `/dev/shm` size threshold (0 disables). |
+| `--rocm-smi-timeout-sec SEC` | 5.0 | Hard timeout for the `rocm-smi --version` self-latency canary; hitting it is a node FAIL. |
+| `--hbm-busy-threshold-gib GiB` | 2.0 | FAIL if any GPU has ≥ this much HBM in use before smoke touches the device. Boundary inclusive. |
+| `--allow-foreign-procs` | off | Do NOT FAIL on foreign processes holding a GPU. They are still reported. |
+| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Process names OK to find holding the GPU. Set to `""` to disable the whitelist. |
+| `--gpu-activity-warn-pct PCT` | 20.0 | Warn (does NOT fail) if any GPU's `gfx_activity_pct` exceeds this at smoke start. |
+| `--require-tools LIST` | `""` (warn-only) | CLI tools that MUST be in PATH (`amd-smi`, `rocm-smi`, `lsof`); anything missing becomes a hard node FAIL. |
+| `--no-clean-dump-path` | off | Do NOT auto-wipe stale per-node JSONs / aggregator outputs from `--dump-path` on rank 0 at startup. |
+
+### 9.3 Standalone `aggregate` flags
+
+The primus-cli wrapper runs the aggregator automatically on rank 0 and fills `--expected-nodes` / `--expected-nodelist-file` from SLURM. These matter only when you invoke `python -m primus.tools.preflight.node_smoke aggregate` yourself (see [§6.9](#69-re-aggregate-from-existing-per-node-jsons-no-re-run)).
+
+| Flag | Default | Purpose |
+|---|---|---|
+| `--dump-path` | `output/preflight` | Same as `run`. |
+| `--expected-nodes N` | none | If fewer JSONs land within `--wait-timeout-sec`, missing nodes are added as FAIL placeholders. |
+| `--wait-timeout-sec` | 60 | Polling timeout. |
+| `--rocm-smi-warn-sec SEC` | 1.0 | Flag (warn-only) any node where `rocm-smi --version` took longer than this. |
+| `--clock-skew-warn-sec SEC` | 30.0 | Warn when wall-clock spread across nodes exceeds this many seconds (includes srun launch jitter). |
+| `--hbm-busy-threshold-gib GiB` | 2.0 | Mirrors the `run` default; labels the "GPU pre-touch HBM usage outliers" section. |
+| `--gpu-activity-warn-pct PCT` | 20.0 | Mirrors the `run` default; labels the "GPU compute-activity outliers" section. |
+| `--expected-nodelist-file FILE` | none | One short hostname per line. Missing nodes get their real short hostname in the report and `failing_nodes.txt`. The wrapper auto-populates this from `scontrol show hostnames "$SLURM_JOB_NODELIST"` under SLURM. |
+
+### 9.4 Launcher-level knobs (`primus-cli direct`)
+
+Consumed by `primus-cli-direct.sh` **before** the `--` separator (not forwarded to the `node_smoke` Python tool):
+
+| Flag | Purpose |
+|---|---|
+| `--silent` | Redirect launcher + tool stdout to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN` on stderr) and the log file are preserved. Exit code propagated. |
+| `--debug` | Verbose launcher logging. |
+| `--dry-run` | Print the resolved configuration and command without executing. |
+| `--env KEY=VALUE` | Inject an env var into the Python process. |
+
+> **Run vs. aggregate.** The primus-cli `node_smoke` subcommand always runs the per-node checks on every rank, then aggregates on rank 0 — which is what you want ~100% of the time, so there is no `--aggregate-only` wrapper flag. For the rare single-phase cases, call the standalone CLI directly: `python -m primus.tools.preflight.node_smoke run ...` (per-node only, no report) or `... aggregate ...` (report only, from existing JSONs).
+
+---
+
+## 10. Comparison with the full `preflight`
+
+| Aspect | `node_smoke` | full `preflight` |
+|---|---|---|
+| Rendezvous | None — every node independent | Global `torch.distributed` |
+| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests |
+| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) |
+| GEMM | Hard threshold per GPU | Reports per-GPU numbers, no auto-fail |
+| HBM bandwidth | Yes (D2D `copy_`) | Not measured |
+| Inter-node all-reduce / all-to-all | Not tested (intentionally) | Yes |
+| Drift detection | Yes (versions, NIC firmware, port count) | No |
+| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only |
+| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF |
+
+Use `node_smoke` to **screen** a cluster fast and exclude bad nodes. Use the full [`preflight`](./preflight.md) when you want **deep cross-node measurements** (inter-node bandwidth matrix, ring-P2P, etc.). The recommended sequence is node-smoke first, then `preflight --quick` on the surviving nodes.
+
+---
+
+## 11. Troubleshooting
+
+| Symptom                                                                       | Likely cause / fix                                                                                                                                                                            |
+| ----------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...`          | Fix the path (`export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate`), or `unset VENV_ACTIVATE` to fall back to system / container Python.                                                |
+| Every node FAILs with `gpu_processes: ... name='N/A'`                         | Container `/proc` can't resolve host PID names. See [§6.5](#65-allow--extend-the-foreign-process-whitelist): pass `--allow-foreign-procs`, or grant host-PID visibility.                       |
+| Some nodes never produce a JSON                                               | The aggregator names them in `failing_nodes.txt` via `expected_nodes.txt`. If `scontrol` was unavailable, they appear as `<missing-N>`.                                                       |
+| Tier 2 perf numbers below threshold on a known-good node                      | Almost always insufficient CPU cores on `srun` — pass `-c <cores-per-node>` so RCCL proxy threads have CPU.                                                                                   |
+| Re-run on a smaller nodelist still shows the previously removed nodes as PASS | Default behavior cleans stale JSONs on rank 0. If you passed `--no-clean-dump-path`, either remove it or `rm -rf output/preflight` between runs.                                              |
+
+---
+
+## 12. See also
+
+- [`preflight.md`](./preflight.md) — the heavier `preflight` tool with a global rendezvous and inter-node bandwidth tests.
+- [`preflight-without-container.md`](./preflight-without-container.md) — running `preflight` directly on the host (no container), including the shared venv + NCCL setup.
+- [`primus/cli/subcommands/node_smoke.py`](../../primus/cli/subcommands/node_smoke.py) — the primus-cli subcommand wiring (two-phase dispatch: per-rank run + rank-0 aggregate).
+- [`primus/tools/preflight/node_smoke/cli.py`](../../primus/tools/preflight/node_smoke/cli.py) — canonical flag definitions and per-node / aggregate phase bodies.
diff --git a/docs/02-user-guide/posttraining.md b/docs/02-user-guide/posttraining.md
new file mode 100644
index 000000000..dedd39235
--- /dev/null
+++ b/docs/02-user-guide/posttraining.md
@@ -0,0 +1,216 @@
+# Post-training workflows
+
+Post-training (supervised fine-tuning) adapts a pre-trained foundation model to new tasks or domains. In Primus, post-training runs through the **Megatron Bridge** backend using the `train posttrain` subcommand. Example YAML configurations live under `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus).
+
+For YAML field details, see [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md). For related tooling, see [Benchmark suite](./benchmarking.md), [Preflight diagnostics](./preflight.md), [Memory and performance projection](./projection.md).
+
+---
+
+## Overview: SFT vs LoRA
+
+The following table shows how SFT (Supervised Fine-Tuning) and LoRA (Low-Rank Adaptation) differ in different aspects.
+
+| Aspect | SFT (full fine-tuning) | LoRA (parameter-efficient) |
+|--------|------------------------|----------------------------|
+| **PEFT setting** | `peft: "none"` | `peft: lora` |
+| **What is trained** | All model parameters | Low-rank adapters only |
+| **Memory** | Higher | Lower |
+| **Throughput** | Typically slower per step | Often faster iteration |
+| **Learning rate** | Lower: roughly `5e-6` to `1e-5` | Higher: roughly `1e-4` to `5e-4` |
+| **Typical use** | Maximum adaptation when memory allows | Limited GPU memory, many task-specific adapters, rapid experimentation |
+
+---
+
+## Quick start commands
+
+General form:
+
+```bash
+./primus-cli <mode> -- train posttrain --config <path-to-yaml>
+```
+
+From a clone of the Primus repository, the same entrypoint is often invoked as `./runner/primus-cli`.
+
+**Prerequisites:** AMD ROCm (recommended ≥ 7.0) and Docker with ROCm support (optional but typical) installed on systems with AMD Instinct™ GPUs (for example MI300X, MI355X). Run this for a quick check on these prerequisites: `rocm-smi && docker --version`. See [Installation and setup](../01-getting-started/installation.md) for the full prerequisites and container setup.
+
+### Direct mode (bare metal or inside a Docker container)
+
+```bash
+# SFT — example: Qwen3 32B on MI355X
+./runner/primus-cli direct -- train posttrain \
+  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml
+
+# LoRA — same model family
+./runner/primus-cli direct -- train posttrain \
+  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml
+```
+
+### Container mode
+
+```bash
+./runner/primus-cli container --image rocm/primus:v26.4 -- \
+  train posttrain \
+  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml
+```
+
+---
+
+## Configuration reference
+
+These keys are commonly set under `modules.post_trainer.overrides` in your configuration YAML (see the examples in the repository under `examples/megatron_bridge/configs/`).
+
+| Area | Parameters | Notes |
+|------|------------|--------|
+| **Method** | `peft` | `"none"` for SFT; `lora` for LoRA. |
+| **Learning rate** | `finetune_lr`, `min_lr`, `lr_warmup_iters`, `lr_decay_iters` | LoRA usually needs a higher `finetune_lr` than SFT. |
+| **Precision** | `precision_config` | Typical: `bf16_mixed`. Alternatives include `fp16_mixed` and `fp32` depending on backend support. |
+| **Parallelism** | `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `context_parallel_size`, `sequence_parallel` | Increase TP/PP when the model does not fit on fewer GPUs. |
+| **Recompute (memory)** | `recompute_granularity`, `recompute_method`, `recompute_num_layers` | Use to trade compute for activation memory (for example `recompute_granularity: full` with uniform recompute). |
+| **Batching / length** | `train_iters`, `global_batch_size`, `micro_batch_size`, `seq_length` | `micro_batch_size` is per-GPU; tune with sequence length and memory. |
+
+Snippet of an SFT configuration (illustrative only):
+
+```yaml
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+    model: qwen3_32b.yaml
+    overrides:
+      peft: "none"
+      finetune_lr: 5.0e-6
+      precision_config: bf16_mixed
+      tensor_model_parallel_size: 1
+      global_batch_size: 8
+      micro_batch_size: 1
+      seq_length: 8192
+```
+
+Snippet of a LoRA configuration (illustrative only):
+
+```yaml
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+    model: qwen3_32b.yaml
+    overrides:
+      peft: lora
+      finetune_lr: 1.0e-4
+      precision_config: bf16_mixed
+      recompute_granularity: full
+      recompute_method: uniform
+      recompute_num_layers: 1
+```
+
+---
+
+> The reference configurations below are organized by GPU architecture. **MI325X** uses the same configurations as **MI300X** (both are `gfx942`), and **MI350X** uses the same configurations as **MI355X** (both are `gfx950`).
+
+## MI300X configurations
+
+Paths are relative to `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus).
+
+| Model | Method | Config path | TP | GBS | MBS | Seq len |
+|-------|--------|-------------|----|-----|-----|---------|
+| Qwen3 32B | SFT | `MI300X/qwen3_32b_sft_posttrain.yaml` | 2 | 8 | 2 | 8192 |
+| Qwen3 32B | LoRA | `MI300X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 2 | 8192 |
+
+**Legend:** TP = tensor parallel size; GBS = global batch size; MBS = micro batch size per GPU; Seq len = `seq_length`.
+
+Sample command for running the post-training:
+
+```bash
+./runner/primus-cli direct -- train posttrain \
+  --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml
+```
+
+---
+
+## MI355X configurations
+
+Paths are relative to `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus).
+
+| Model | Method | Config path | TP | GBS | MBS | Seq len |
+|-------|--------|-------------|----|-----|-----|---------|
+| Qwen3 32B | SFT | `MI355X/qwen3_32b_sft_posttrain.yaml` | 1 | 8 | 1 | 8192 |
+| Qwen3 32B | LoRA | `MI355X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 4 | 8192 |
+
+Sample command for running the post-training:
+
+```bash
+./runner/primus-cli direct -- train posttrain \
+  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml
+```
+
+---
+
+## Best practices
+
+### Use SFT or LoRA?
+
+- **SFT is preferred** when you need the strongest possible task fit, have enough GPU memory, and can afford longer runs.
+- **LoRA is preferred** when memory is tight, you want fast iteration, or you plan to maintain multiple adapters for different tasks.
+
+### Learning rates
+
+- **SFT:** start in the `5e-6`–`1e-5` range; adjust with validation loss.
+- **LoRA:** often `1e-4`–`5e-4`; still use warmup (`lr_warmup_iters`) for stability.
+
+### Batch sizes
+
+- **SFT**: starting with `global_batch_size: 8` is a reasonable default for development; scale up when stable (for example to 64, 128, or higher) if memory and throughput allow.
+- **LoRA**: larger global batches are often feasible (for example 32 in the reference configs); align `micro_batch_size` with sequence length and available HBM.
+- Very long sequences (for example 8192) may require smaller micro-batches or more parallelism.
+
+### Parallelism
+
+- **SFT:** large models may need higher `tensor_model_parallel_size` (for example TP 8 for very large models). The bundled 32B examples use TP 2 on MI300X and TP 1 on MI355X for SFT.
+- **LoRA:** adapters reduce memory pressure; lower TP is often sufficient for a given model size.
+
+---
+
+## Troubleshooting
+
+### Out of memory (OOM)
+
+**SFT**
+
+1. Increase `tensor_model_parallel_size` (and/or pipeline parallelism for very large models).
+2. Reduce `micro_batch_size` or `seq_length`.
+3. Enable activation recomputation (`recompute_granularity`, `recompute_method`, `recompute_num_layers`).
+
+**LoRA**
+
+1. Confirm `peft: lora` is set.
+2. Reduce `micro_batch_size` if OOM persists.
+3. Apply the same recompute settings as for SFT.
+
+### Training instability (loss spikes, NaNs)
+
+1. Decrease `finetune_lr`.
+2. Increase `lr_warmup_iters`.
+3. Keep mixed precision stable (`precision_config: bf16_mixed` where supported).
+4. Monitor gradients and clipping settings if exposed by your trainer config.
+
+### Slow training
+
+1. Increase effective batch size where memory allows (`global_batch_size` / `micro_batch_size` tuning).
+2. Revisit TP/PP/CP for your cluster topology.
+3. Run [benchmarks](./benchmarking.md) or [preflight](./preflight.md) to isolate network or GPU issues.
+
+### Configuration errors
+
+1. Verify YAML paths and indentation.
+2. Set `PRIMUS_WORKSPACE` and other environment variables expected by your team’s templates.
+3. Confirm checkpoint and data paths by reviewing the experiment YAML and any presets it references. The current core training runtime parses `--export_config`, but resolved-config export is not implemented on the default `PrimusRuntime` path.
+
+---
+
+## Related documentation
+
+- [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md)
+- [Native SFT / LoRA quick start](../04-technical-guides/native-sft-lora.md)
+- [Benchmark suite](./benchmarking.md)
+- [Preflight diagnostics](./preflight.md)
+- [Memory and performance projection](./projection.md)
diff --git a/docs/02-user-guide/preflight-without-container.md b/docs/02-user-guide/preflight-without-container.md
new file mode 100644
index 000000000..9dbb23192
--- /dev/null
+++ b/docs/02-user-guide/preflight-without-container.md
@@ -0,0 +1,792 @@
+# Run preflight without a container
+
+> ⚠ **Run the [node-smoke test](./node-smoke-test-instruction.md) first.** `preflight` opens a global `torch.distributed` rendezvous, so a single sick node (wedged driver, leaked rank holding HBM, partial NIC enumeration, time-sync drift, etc.) can stall the whole job for up to `--dist-timeout-sec` seconds — long before any cross-node bandwidth number is produced. The node-smoke test catches those exact failure modes *without* a rendezvous in ~30–60 s and emits a SLURM-ready `failing_nodes.txt` you can pipe straight into `srun --exclude=`. Treat node-smoke as a hard prerequisite; only run `preflight` on the nodes node-smoke marked PASS. See [§0 "Which test should I run?"](#0-which-test-should-i-run) for the side-by-side comparison and the recommended 3-step workflow.
+
+This guide explains how to run Primus's [`preflight`](./preflight.md) cluster-diagnostic tool **directly on the host** (no Docker / Podman), via the standard Primus launcher.
+
+**Git clone the Primus repository to a shared filesystem that all nodes can read.**
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+cd Primus
+```
+
+**Recommended (through the primus-cli SLURM wrapper):**
+
+For some clusters, you may need to explicitly request CPU and GPU resources with `srun -N <nodes> -c <cpus-per-node> --gpus-per-node=<gpus-per-node>`.
+
+```
+runner/primus-cli slurm srun -N <NNODES> --ntasks-per-node=1 -- direct -- preflight [PREFLIGHT_ARGS...]
+```
+
+**Equivalent (bare srun, useful when composing with custom srun flags):**
+
+```
+srun -N <NNODES> --ntasks-per-node=1 runner/primus-cli direct -- preflight [PREFLIGHT_ARGS...]
+```
+
+Both forms produce the **same workload** on the same ranks. The wrapper form is recommended because it auto-resolves `MASTER_ADDR` / `MASTER_PORT` / `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` once on the launching node and passes them to every rank via `--env`, applies any `slurm.`* config defaults (partition / time / etc.) from your YAML, and is the same pattern used for `train` / `benchmark` / `node_smoke`. See [§ Wrapper vs. bare-srun](#wrapper-vs-bare-srun) below for the exact precedence / caveats.
+
+`primus-cli direct` activates an optional Python virtualenv (`VENV_ACTIVATE`), auto-derives the distributed environment variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`) from `SLURM_*` when running inside a SLURM allocation, and then launches the `preflight` Python subcommand via `torchrun` (one worker per GPU). It is the recommended entry point when:
+
+- You're running on a SLURM cluster but cannot (or don't want to) use the container-based path.
+- Your nodes share a Python virtual environment on a network-mounted filesystem.
+- You want a single-node sanity check with no extra configuration.
+
+---
+
+## 0. Which test should I run?
+
+Primus ships **two** complementary cluster screens. Pick the right one — and ideally run them in this order.
+
+
+| Aspect            | `node-smoke` (start here)                                                          | `preflight` (this doc)                                              |
+| ----------------- | ---------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
+| Purpose           | "Which nodes are healthy enough to run anything?"                                  | "What is the actual cross-node performance on the surviving nodes?" |
+| Rendezvous        | None — every node independent                                                      | Global `torch.distributed` rendezvous                               |
+| Wall clock        | ~30–60 s for 6 nodes (Tier 1+2)                                                    | A few minutes; scales with N for inter-node tests                   |
+| Granularity       | Per-node PASS/FAIL                                                                 | Per-rank perf measurements                                          |
+| Safety            | A stuck node cannot wedge its peers                                                | A single hung NIC can stall the whole rendezvous                    |
+| Output            | Per-node JSON + cluster md + SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` | Markdown + PDF perf report                                          |
+| Entry point       | `primus-cli direct -- node_smoke`                                                  | `primus-cli direct -- preflight` (this doc)                         |
+| Quick-start guide | [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md)               | This doc, §3+                                                       |
+
+
+### Recommended workflow
+
+> **Before running any of the commands below, complete the one-time setup:**
+>
+> 1. **Python virtualenv** on a shared filesystem — see [§2 Set up the Python virtual environment](#2-set-up-the-python-virtual-environment), then point the launcher at it via `export VENV_ACTIVATE=...` (details in [§2 → Tell the launcher where the venv is](#tell-the-launcher-where-the-venv-is)).
+> 2. **NCCL / fabric environment variables** — usually the defaults in `base_env.sh` are fine, but multi-NIC nodes may need `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` / `NCCL_SOCKET_IFNAME` overrides. See [§4 Cluster-specific NCCL configuration](#4-cluster-specific-nccl-configuration) for known-good values per fabric (Broadcom, Pensando Pollara/AINIC).
+
+Through the `primus-cli slurm srun -- direct --` wrapper (recommended):
+
+```bash
+# 1) Prune broken nodes with node-smoke (fast, no rendezvous).
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf
+
+# 2) Re-allocate excluding the bad nodes, and run preflight --quick
+#    for a fast cross-node sanity check.
+runner/primus-cli slurm srun -N <good-nnodes> -c 128 --gpus-per-node=8 \
+    --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    -- direct -- preflight --quick
+
+# 3) Optional: full preflight on the same set if --quick numbers
+#    look off, or if you want the full bandwidth matrix.
+runner/primus-cli slurm srun -N <good-nnodes> -c 128 --gpus-per-node=8 \
+    --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    -- direct -- preflight
+```
+
+Equivalent with bare `srun` (works identically; useful when scripting around custom srun flags that don't compose with the wrapper):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf
+
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct -- preflight --quick
+
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct -- preflight
+```
+
+Why this ordering matters:
+
+- A single broken node can stall a `torch.distributed.init_process_group()` for `--dist-timeout-sec` seconds (default 120), so feeding a known-good list to preflight is much faster.
+- `node-smoke` catches things preflight cannot — leaked / foreign processes, wedged drivers, partial NIC enumeration, time-sync drift, RDMA roll-call issues — that produce *misleading* preflight failures.
+- `preflight --quick` adds the cross-node bandwidth signal that `node-smoke` deliberately does not measure.
+
+---
+
+## 1. Prerequisites
+
+- A working AMD ROCm installation on every node.
+- Network reachability between nodes (Ethernet for bootstrap, RDMA / InfiniBand recommended for perf tests).
+- A Python ≥ 3.10 virtual environment **on a shared filesystem** that all nodes can read (the same path is sourced on every node).
+- The Primus repository checked out somewhere readable from every node.
+
+---
+
+## 2. Set up the Python virtual environment
+
+The environment must live on a path visible from every node (e.g. NFS-mounted home, Lustre, or any shared filesystem). All nodes will `source` the same activation script.
+
+You can use any tool you like; `uv` is the fastest. Either of the following works.
+
+### What you actually need to install
+
+The `preflight` and `node-smoke` tools deliberately use **only a small subset** of Primus's full dependency tree. You do **not** need to install the entire `requirements.txt` — that pulls in trainer / dataset / experiment-tracking packages that neither tool ever imports.
+
+
+| Package              | Required for                                                                          | Skip when                                                                                  |
+| -------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ |
+| `torch` (ROCm build) | Both tools — perf measurements (`torch.matmul`, `torch.distributed`, `torch.cuda.`*). | Never (mandatory).                                                                         |
+| `markdown2`          | `preflight` PDF report only (Markdown → HTML).                                        | You always pass `--disable-pdf`, or you only run `node-smoke` (which never produces PDFs). |
+| `weasyprint`         | `preflight` PDF report only (HTML → PDF).                                             | Same as above.                                                                             |
+| `matplotlib`         | `preflight --plot` only (per-test bandwidth bar charts).                              | You don't pass `--plot`.                                                                   |
+
+
+Everything else in the preflight / node-smoke code path is Python stdlib (`os`, `subprocess`, `socket`, `argparse`, `dataclasses`, `json`, `time`, ...) — no extra installs needed.
+
+### Option A — `uv` (recommended), minimal install
+
+```bash
+mkdir -p ~/envs/preflight
+cd ~/envs/preflight
+
+uv venv --python 3.12
+source .venv/bin/activate
+
+# 1) ROCm-built PyTorch (pin to your ROCm version; rocm7.1 shown here)
+uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir
+
+# 2) Optional: only if you want preflight PDF reports (omit to use --disable-pdf)
+uv pip install markdown2 weasyprint
+
+# 3) Optional: only if you want preflight --plot bar charts
+uv pip install matplotlib
+```
+
+### Option B — `python -m venv`, minimal install
+
+```bash
+mkdir -p ~/envs/preflight
+python3.12 -m venv ~/envs/preflight/.venv
+source ~/envs/preflight/.venv/bin/activate
+
+pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir
+pip install markdown2 weasyprint   # optional, for preflight PDFs
+pip install matplotlib             # optional, for preflight --plot
+```
+
+### Option C — full Primus runtime (only if you also want the rest of Primus)
+
+```bash
+cd /path/to/Primus
+uv pip install -r requirements.txt   # or: pip install -r requirements.txt
+```
+
+This installs every Primus runtime dependency (trainer, dataset loaders, experiment trackers, ...). Use only if you're going to run more than just preflight / node-smoke from this environment.
+
+### Per-tool minimum install matrix
+
+If you want the absolute smallest footprint, install only what your intended invocations need:
+
+
+| Invocation                                       | `torch`  | `markdown2`                       | `weasyprint`                      | `matplotlib` |
+| ------------------------------------------------ | -------- | --------------------------------- | --------------------------------- | ------------ |
+| `node-smoke` (any flags)                         | required | —                                 | —                                 | —            |
+| `preflight --host --gpu --network --disable-pdf` | required | —                                 | —                                 | —            |
+| `preflight --host --gpu --network` (with PDF)    | required | required                          | required                          | —            |
+| `preflight --quick --disable-pdf`                | required | —                                 | —                                 | —            |
+| `preflight --quick` (with PDF)                   | required | required                          | required                          | —            |
+| `preflight ... --plot`                           | required | required (unless `--disable-pdf`) | required (unless `--disable-pdf`) | required     |
+
+
+### Tell the launcher where the venv is
+
+`primus-cli direct` reads the `VENV_ACTIVATE` environment variable. When set, it sources the path before launching the Python process; when unset, it is a no-op (the container path, which uses the container's bundled Python, leaves this unset):
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+```
+
+`VENV_ACTIVATE` is the only optional environment variable specific to the direct flow. Everything else has a sensible default; distributed-env variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, ...) are auto-derived from SLURM when not pre-exported.
+
+---
+
+## 3. Run preflight
+
+### Single node (no SLURM)
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Info report only (fast)
+runner/primus-cli direct -- preflight --host --gpu --network
+
+# Info + perf report
+runner/primus-cli direct -- preflight
+
+# Perf report only
+runner/primus-cli direct -- preflight --perf-test
+```
+
+When SLURM is not detected the script defaults to `NNODES=1`, `NODE_RANK=0`, `MASTER_ADDR=localhost`. Any of those can be overridden by exporting them before calling the script.
+
+### Multi-node without SLURM (parallel SSH)
+
+When no scheduler is available (bare-metal, cloud VMs, lab nodes), launch
+`primus-cli direct` on each node yourself via SSH. The script works
+identically — you just pre-export the distributed variables that SLURM
+would normally provide.
+
+#### Requirements
+
+- All nodes share the same filesystem (or at least the same Primus checkout + venv path).
+- Nodes can reach each other on a **data-plane** network interface (not the management NIC).
+- SSH key-based access to each node from the launching host.
+
+#### Required environment variables
+
+
+| Variable             | Description                                                            |
+| -------------------- | ---------------------------------------------------------------------- |
+| `NNODES`             | Total number of nodes                                                  |
+| `NODE_RANK`          | This node's rank (`0` through `NNODES-1`)                              |
+| `MASTER_ADDR`        | IP of rank-0 node **on the data-plane interface**                      |
+| `MASTER_PORT`        | Rendezvous port (default `1234`; increment between concurrent runs)    |
+| `GPUS_PER_NODE`      | GPUs per node (default `8`)                                            |
+| `NCCL_SOCKET_IFNAME` | Data-plane NIC name (e.g. `enp159s0np0`) — **critical for multi-node** |
+| `GLOO_SOCKET_IFNAME` | Same as `NCCL_SOCKET_IFNAME`                                           |
+| `VENV_ACTIVATE`      | Path to virtualenv `activate` script                                   |
+
+
+> **Warning**: `NCCL_SOCKET_IFNAME` auto-detection often picks a management interface
+> (e.g. `enp28s0np0`, `eno8303`) instead of the high-bandwidth data NIC. For multi-node
+> runs this causes `init_process_group` to hang or NCCL to fail silently. Always set it
+> explicitly.
+
+#### Identifying the correct data-plane interface
+
+```bash
+# On any node, find the interface whose IP matches the MASTER_ADDR subnet:
+ip -4 addr show | grep "10.245.134"
+# → enp159s0np0  inet 10.245.134.129/24
+
+# Or check which interface routes to the master:
+ip route get 10.245.134.129 | awk '{print $5; exit}'
+```
+
+### Multi-node via SLURM
+
+`primus-cli direct` auto-detects a SLURM allocation (via `SLURM_JOB_ID`) and derives all distributed variables from `SLURM_*` automatically. **Pre-exported values always win**, so the same launcher script also works inside the `primus-cli slurm srun ... -- direct -- ...` chain (where `slurm-entry` has already set these via `--env`):
+
+| Variable        | Resolved as                                                          |
+| --------------- | -------------------------------------------------------------------- |
+| `NNODES`        | `NNODES` → `SLURM_NNODES` → `SLURM_JOB_NUM_NODES` → `1`              |
+| `NODE_RANK`     | `NODE_RANK` → `SLURM_NODEID` → `SLURM_PROCID` → `0`                  |
+| `MASTER_ADDR`   | `MASTER_ADDR` (if not empty / not `localhost`) → first hostname from `scontrol show hostnames "$SLURM_NODELIST"` |
+| `MASTER_PORT`   | `MASTER_PORT` → `1234`                                               |
+| `GPUS_PER_NODE` | `GPUS_PER_NODE` → `8`                                                |
+
+Run it as a single task per node (the script invokes `torchrun` internally, which spawns one worker per GPU):
+
+> **Verify NCCL / network env first.** The script sets sensible `NCCL_`* defaults via `base_env.sh`, but auto-detection can pick the wrong device on multi-NIC nodes. Always confirm `NCCL_IB_HCA`, `NCCL_IB_GID_INDEX`, `NCCL_SOCKET_IFNAME`, and `GLOO_SOCKET_IFNAME` (set to the same value as `NCCL_SOCKET_IFNAME`) are correct for your fabric, and `export` overrides before running. See [§4](#4-cluster-specific-nccl-configuration) for cluster-specific values.
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# export NCCL_IB_HCA=rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
+# export NCCL_IB_GID_INDEX=3
+# export NCCL_SOCKET_IFNAME=eno0
+# export GLOO_SOCKET_IFNAME=eno0
+
+# Recommended: through the primus-cli SLURM wrapper.
+runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 \
+    --nodelist <nodes> --ntasks-per-node=1 \
+    -- direct -- preflight --perf-test
+
+# Or, equivalently, with bare srun:
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight --perf-test
+```
+
+
+
+#### Wrapper vs. bare-srun
+
+Both forms target the **same** `primus-cli-direct.sh` launcher and produce identical workloads. The difference is only in how the SLURM context is constructed:
+
+
+| Aspect                                   | `primus-cli slurm srun -- direct --` (recommended)                                                                                                             | Bare `srun ... primus-cli direct --`                                                                                                |
+| ---------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
+| `MASTER_ADDR` resolution                 | Resolved **once** on the launching node via `scontrol show hostnames "$SLURM_NODELIST" | head -n1`, then propagated to every rank via `--env MASTER_ADDR=...`. | Each rank re-derives it inside `primus-cli-direct.sh` STEP 4.7 from `SLURM_`* (same result, more `scontrol` calls).                 |
+| `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` | Set explicitly by `slurm-entry.sh` via `--env`.                                                                                                                | Derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_PROCID` inside `direct.sh`.                                                   |
+| `slurm.*` config defaults                | Honored (partition, time, ntasks-per-node, etc. from the active YAML).                                                                                         | Not consulted — you pass every flag explicitly to `srun`.                                                                           |
+| Default wall-time                        | `-t 4:00:00` is auto-added if you don't pass `--time`.                                                                                                         | None — `srun` uses the cluster default (may reject the job).                                                                        |
+| `direct` keyword                         | **Required**: `primus-cli slurm srun ... -- direct -- <cmd>`. Without `direct`, the wrapper routes through the **container** path.                             | N/A — there's only one path.                                                                                                        |
+| `--ntasks-per-node=1`                    | **Not auto-added**. Pass it on the CLI (before the first `--`) or set it in the `slurm.`* config.                                                              | **Not auto-added**. Pass it as an `srun` flag.                                                                                      |
+| Best for                                 | Production / repeatable runs. Same pattern as `train` / `benchmark` / `node_smoke`.                                                                            | Ad-hoc runs where you want to compose with arbitrary `srun` flags (`--nodelist=$(...)`, `--exclude=...` from a runtime file, etc.). |
+
+
+For the rest of this doc the examples use bare `srun` for brevity, but every example also works with the wrapper form by substituting `srun <flags> runner/primus-cli direct --` → `runner/primus-cli slurm srun <flags> -- direct --`.
+
+### Key `srun` flags
+
+
+| Flag                  | Why it's necessary                                                                                                                                                                                      |
+| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `-c 128`              | Allocate all CPU cores per task. Without this, SLURM may default to 1 core, which starves the RCCL network proxy threads and can cause >30× slowdown on perf tests. Set this to your node's core count. |
+| `--gpus-per-node=8`   | Grants GPU device access (`/dev/kfd`, `/dev/dri`). Required for non-container execution.                                                                                                                |
+| `--ntasks-per-node=1` | One launcher invocation per node; `primus-cli direct` then spawns 8 workers per node via `torchrun`.                                                                                                    |
+| `-t 00:45:00`         | Wall-clock limit. Full perf tests on 8N usually finish well under 10 min.                                                                                                                               |
+
+
+> Tip — check core count: `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`
+
+---
+
+## 4. Cluster-specific NCCL configuration
+
+`primus-cli direct` sources `runner/helpers/envs/base_env.sh`, which sets sensible defaults for `NCCL_`* and auto-detects `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME`. Pre-exported values from your shell take precedence, so the standard pattern is:
+
+```bash
+export VAR=value
+runner/primus-cli direct -- preflight ...
+```
+
+### Broadcom NICs (no AINIC)
+
+Most clusters fall here. The defaults from `base_env.sh` are usually fine, but the two values most commonly worth overriding are:
+
+```bash
+export NCCL_CROSS_NIC=1     # default in base_env.sh is 0
+export NCCL_PXN_DISABLE=0   # default in base_env.sh is 1
+
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight --perf-test
+```
+
+### Pensando Pollara (AINIC) RDMA
+
+```bash
+export USING_AINIC=1
+export NCCL_IB_GID_INDEX=1   # AINIC uses index 1 (default in base_env.sh is 3)
+export NCCL_PXN_DISABLE=0
+
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight
+```
+
+> `primus-cli direct` *does* accept `--env KEY=VALUE` on its own command line (placed before `--`), in addition to the conventional `export`/`srun --export=` approaches.
+
+---
+
+## 5. Launcher flags vs. preflight flags
+
+Anything you place **after** the `--` separator is forwarded verbatim to the `preflight` Python tool. The launcher (`primus-cli-direct.sh`) consumes a small set of flags **before** `--`. The one most users care about is `--silent`.
+
+### Launcher-only flags (before `--`)
+
+
+| Flag              | Effect                                                                                                                                                                                                                                                                                                                                                                      |
+| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `--silent`        | Back-pocket knob: redirect the launcher's and the Python tool's `stdout` to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN`, written to `stderr`) are preserved so real failures still surface; the log file under `logs/` captures everything. Exit code is propagated unchanged. **Not recommended** for normal use — you lose live progress; prefer the log file. |
+| `--debug`         | Verbose launcher logging (`PRIMUS_LOG_LEVEL=DEBUG`). Forwarded to the Python tool as `--debug` too.                                                                                                                                                                                                                                                                         |
+| `--dry-run`       | Print the resolved configuration and final `torchrun` / `python3` command without executing.                                                                                                                                                                                                                                                                                |
+| `--single`        | Force `python3` instead of `torchrun`. `node_smoke` auto-selects this; for `preflight` you usually want the default (`torchrun`).                                                                                                                                                                                                                                           |
+| `--env KEY=VALUE` | Inject an env var into the Python process (in addition to anything `export`-ed in the shell).                                                                                                                                                                                                                                                                               |
+| `--log_file PATH` | Redirect the captured tee log to a specific path (default: `logs/log_<timestamp>.txt`).                                                                                                                                                                                                                                                                                     |
+
+
+See `runner/primus-cli direct --help` for the full set.
+
+### Forwarded `preflight` flags (after `--`, most common)
+
+See [Preflight](./preflight.md) for the full list. The most common are:
+
+- Mode selection: `--host`, `--gpu`, `--network`, `--perf-test`, `--tests`, `--quick`
+- Test tuning: `--comm-sizes-mb`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--intra-group-sizes`, `--inter-group-sizes`, `--ring-p2p-sizes-mb`
+- Reporting: `--dump-path`, `--report-file-name`, `--disable-pdf`, `--plot`
+- Reliability: `--comm-cleanup-delay-sec`, `--dist-timeout-sec`
+
+If you do not pass `--report-file-name`, `preflight` auto-generates a unique one of the form:
+
+```
+preflight-${NNODES}N-YYYYMMDD-HHMMSS
+```
+
+This guarantees that each run lands in its own files and prevents stale leftovers from earlier runs from being mistaken for fresh output. The auto-name logic now lives in the Python tool itself, so every call site (host `srun ... primus-cli direct`, `primus-cli slurm ... -- direct`, `primus-cli slurm ... -- container`) gets the same fresh name.
+
+### Examples
+
+The examples below all assume one of the two equivalent shell-prefix conventions. Pick whichever matches your habits — every example block in this section works with either definition:
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Recommended: through the primus-cli SLURM wrapper. Auto-resolves
+# MASTER_ADDR/NNODES/NODE_RANK once on the launching node and propagates
+# them via --env; honors slurm.* config defaults.
+SRUN="runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist <nodes> --"
+# Then in every example below, replace `$SRUN runner/primus-cli direct --`
+# with just `$SRUN direct --`. (The wrapper expects the entry-mode keyword
+# `direct` as the first token after the inner `--`.)
+
+# Equivalent: bare srun. NNODES/NODE_RANK/MASTER_ADDR get derived inside
+# primus-cli-direct.sh's STEP 4.7 directly from SLURM_*; same net effect.
+SRUN="srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist <nodes>"
+```
+
+The examples in this section use the **bare-srun** form below for brevity (since `$SRUN runner/primus-cli direct -- preflight` reads naturally as one command line). To use the wrapper form instead, substitute `$SRUN runner/primus-cli direct --` → `$SRUN direct --` after exporting `SRUN` to the wrapper variant.
+
+#### A. Mode selection
+
+```bash
+# Default: info report + every perf test
+$SRUN runner/primus-cli direct -- preflight
+
+# Info-only (fast, no torch.distributed rendezvous)
+$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf
+
+# Perf-only, every test
+$SRUN runner/primus-cli direct -- preflight --perf-test
+
+# Fast pre-launch sanity preset (gemm + intra-AR + inter-AR @ 64,1024 MB,
+# full intra-node group, full N-node inter group, low warmup/iter)
+$SRUN runner/primus-cli direct -- preflight --quick
+```
+
+> **Note**: Mixing perf-mode flags (`--perf-test` / `--tests` / `--quick`) with info selectors (`--host` / `--gpu` / `--network`) makes preflight drop the info selectors with a `WARN`. Run two invocations if you want both reports.
+
+#### B. Test selection (`--tests`)
+
+```bash
+# Only GEMM
+$SRUN runner/primus-cli direct -- preflight --tests gemm
+
+# Only the inter-node bandwidth tests
+$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall
+
+# Only the inter-node ring P2P
+$SRUN runner/primus-cli direct -- preflight --tests inter-ring-p2p
+
+# Combine: GEMM + inter-AR with overridden sizes/groups
+$SRUN runner/primus-cli direct -- preflight \
+    --tests gemm,inter-allreduce \
+    --comm-sizes-mb 64,1024 \
+    --inter-group-sizes all
+```
+
+Valid `--tests` tokens: `gemm`, `intra-allreduce`, `intra-alltoall`, `inter-allreduce`, `inter-alltoall`, `inter-p2p`, `inter-ring-p2p`, `all`. Unknown tokens fail fast (before NCCL init).
+
+#### C. Message sizes
+
+```bash
+# One CSV applied to both intra and inter
+$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128
+
+# Different sizes for intra vs inter (override wins over --comm-sizes-mb)
+$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128 --intra-comm-sizes-mb 4,32
+
+# Inter-only override (also covers inter-p2p when enabled)
+$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-p2p \
+    --comm-sizes-mb 8,128 --inter-comm-sizes-mb 16,512
+```
+
+#### D. Group sizes
+
+```bash
+# Custom intra-node group sizes (each must divide LOCAL_WORLD_SIZE)
+$SRUN runner/primus-cli direct -- preflight \
+    --tests intra-allreduce \
+    --intra-group-sizes 4,8
+
+# Custom inter-node groups: 2-node pairs and the full N-node group
+$SRUN runner/primus-cli direct -- preflight \
+    --tests inter-allreduce \
+    --inter-group-sizes 2,all
+```
+
+> Note: for `inter-alltoall` only, every requested per-group node count is internally capped at **16** (real-world MoE training rarely dispatches across more nodes; see [`preflight.md` §5.2](./preflight.md#52-group-sizes) for the rationale). The other inter-node tests use the requested sizes unchanged. So on a 128-node cluster, `--tests inter-alltoall --inter-group-sizes all` actually runs at 16-node sub-groups, while `--tests inter-allreduce --inter-group-sizes all` runs at 128 nodes as written.
+
+#### E. Ring P2P sizes
+
+```bash
+$SRUN runner/primus-cli direct -- preflight \
+    --tests inter-ring-p2p \
+    --ring-p2p-sizes-mb 5,20,80
+```
+
+#### F. Plotting
+
+```bash
+# Generate per-test bandwidth bar charts under <dump-path>/<test>/*.png
+$SRUN runner/primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce --plot
+```
+
+#### G. Reliability knobs
+
+```bash
+# Bump the per-phase cleanup delay. Default 2.0 is sufficient at every
+# cluster size for the comm shapes preflight exercises (inter-alltoall
+# is internally capped at 16 nodes; see preflight.md §5.2). Only bump
+# this on very flaky networks or unusual kernel TIME_WAIT settings.
+$SRUN runner/primus-cli direct -- preflight --quick --comm-cleanup-delay-sec 5
+
+# Fail fast if torch.distributed rendezvous can't complete in 30s
+$SRUN runner/primus-cli direct -- preflight --perf-test --dist-timeout-sec 30
+```
+
+> Operating clusters at ≥ 128 nodes? See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the recommended OS-level tuning (widening `ip_local_port_range`) and per-test invocation patterns. With the §5.2 inter-alltoall cap in place, a default invocation is safe at every cluster size; the §7.3 OS tuning remains best-practice for any RDMA host.
+
+#### H. Reporting & output layout
+
+```bash
+# Quick info-only check on 4 nodes, no PDF
+$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf \
+    --report-file-name info-4N
+
+# Perf test only, silenced (CI-friendly), explicit name. Note that --silent
+# is consumed by primus-cli-direct.sh and must appear BEFORE the `--`
+# separator; everything after `--` is forwarded to the preflight Python tool.
+$SRUN runner/primus-cli direct --silent -- preflight --perf-test \
+    --report-file-name nightly-4N-perf
+
+# Archive each run under its own directory
+$SRUN runner/primus-cli direct -- preflight --quick \
+    --dump-path /shared/preflight-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+#### I. Backward-compat aliases
+
+These still work and are equivalent to flags above. Use them only when retrofitting older scripts.
+
+```bash
+# Same as --host --gpu --network
+$SRUN runner/primus-cli direct -- preflight --check-host --check-gpu --check-network
+
+# Same as --inter-group-sizes all AND drops inter-p2p
+$SRUN runner/primus-cli direct -- preflight --perf-test --no-split-nodes-subgroup
+```
+
+#### J. Combined "production-ready" pre-launch screen
+
+```bash
+# 1) Smoke first to prune broken nodes (note: --silent goes BEFORE `--`)
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct --silent -- node_smoke --tier2-perf
+
+# 2) Quick perf sanity on the survivors
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct --silent -- preflight --quick \
+        --comm-cleanup-delay-sec 5 --dist-timeout-sec 60 \
+        --report-file-name screen-$(date +%Y%m%d-%H%M%S)
+```
+
+---
+
+## 6. Outputs
+
+Reports are written to `--dump-path` (default: `output/preflight/`), with the basename from `--report-file-name` and a `_perf` suffix for performance reports:
+
+
+| File              | Produced by                                     | Notes                     |
+| ----------------- | ----------------------------------------------- | ------------------------- |
+| `<name>.md`       | `--host --gpu --network` (or default selection) | Info report               |
+| `<name>.pdf`      | same, unless `--disable-pdf`                    | Info report PDF           |
+| `<name>_perf.md`  | `--perf-test`                                   | Perf report (GEMM + comm) |
+| `<name>_perf.pdf` | same, unless `--disable-pdf`                    | Perf report PDF           |
+
+
+Only **rank 0** writes the report. After preflight completes, the Python tool prints the absolute path of every report file it produced to stdout. Under `--silent` these prints go to `/dev/null` along with everything else (one of the trade-offs of using `--silent`); without `--silent` the announcement is visible live. Sample output:
+
+```
+[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925.md
+[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925_perf.md
+```
+
+---
+
+## 7. Environment variable reference
+
+Variables read by `primus-cli direct` itself:
+
+
+| Variable        | Required | Default                                                                      | Purpose                                                                                                                |
+| --------------- | -------- | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------- |
+| `VENV_ACTIVATE` | no       | —                                                                            | Path to the venv `bin/activate` script. Unset = no-op (use system / container Python). Set + missing file = fail-fast. |
+| `NNODES`        | no       | `1` (or auto-derived from `SLURM_NNODES` / `SLURM_JOB_NUM_NODES`)            | Number of nodes. Pre-exported always wins.                                                                             |
+| `NODE_RANK`     | no       | `0` (or auto-derived from `SLURM_NODEID` / `SLURM_PROCID`)                   | This node's rank. Pre-exported always wins.                                                                            |
+| `GPUS_PER_NODE` | no       | `8`                                                                          | GPUs per node                                                                                                          |
+| `MASTER_ADDR`   | no       | `localhost` (or first host from `scontrol show hostnames "$SLURM_NODELIST"`) | Rendezvous host. Pre-exported always wins.                                                                             |
+| `MASTER_PORT`   | no       | `1234`                                                                       | Rendezvous port                                                                                                        |
+
+
+Variables consumed downstream by `primus-cli direct` / `base_env.sh` (set them via `export`):
+
+
+| Variable             | Default in `base_env.sh` | When to override                                  |
+| -------------------- | ------------------------ | ------------------------------------------------- |
+| `NCCL_SOCKET_IFNAME` | auto-detected            | Force a specific Ethernet interface for bootstrap |
+| `NCCL_IB_HCA`        | auto-detected            | Force specific RDMA HCAs                          |
+| `NCCL_IB_GID_INDEX`  | `3`                      | `1` on AINIC clusters                             |
+| `NCCL_CROSS_NIC`     | `0`                      | `1` for multi-rail IB fabrics                     |
+| `NCCL_PXN_DISABLE`   | `1`                      | `0` to enable PXN multi-hop NIC sharing           |
+| `USING_AINIC`        | unset                    | `1` on Pensando Pollara clusters                  |
+| `NCCL_DEBUG`         | unset                    | `INFO` for verbose NCCL logging                   |
+
+
+---
+
+## 8. Troubleshooting
+
+### `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...`
+
+`VENV_ACTIVATE` was set in the environment but the path it points at doesn't exist on this node. This is a fail-fast guard to prevent a silent fallback to system Python (which usually has the wrong `torch` / no ROCm). Either fix the path:
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+```
+
+… or unset it to fall back to the container / system Python:
+
+```bash
+unset VENV_ACTIVATE
+```
+
+If the path looks right but the file still appears missing, confirm the venv lives on a filesystem visible from the node SLURM scheduled you onto.
+
+### `[Primus:Preflight] FAIL: No GPUs detected`
+
+The Python process inside the venv can't find ROCm. Diagnose with:
+
+```bash
+srun --nodes=1 --nodelist=<node> bash -c '
+echo "=== PATH ==="; echo $PATH
+echo "=== LD_LIBRARY_PATH ==="; echo $LD_LIBRARY_PATH
+echo "=== rocm-smi ==="; rocm-smi --showid 2>&1
+echo "=== Python torch check ==="
+source ~/envs/preflight/.venv/bin/activate
+python3 -c "import torch; print(\"hip:\", torch.version.hip); print(\"available:\", torch.cuda.is_available()); print(\"count:\", torch.cuda.device_count())"
+'
+```
+
+If `LD_LIBRARY_PATH` is empty, set it explicitly:
+
+```bash
+export LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-}
+```
+
+### Report announcement points at stale files
+
+This shouldn't happen with the current Python tool — the auto-generated unique report name (`preflight-${NNODES}N-<timestamp>`) ensures every run gets a fresh path. If you explicitly pass `--report-file-name X`, you're responsible for choosing a name that doesn't collide with prior runs.
+
+### Slow perf tests (~30× expected)
+
+Almost always a symptom of insufficient CPU cores. Pass `-c <cores-per-node>` to `srun` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`.
+
+### Using `conda` instead of venv
+
+`primus-cli direct` does `source "$VENV_ACTIVATE"`, which works for venv/uv but not directly for conda. Two options:
+
+1. Create a venv inside the conda env and point `VENV_ACTIVATE` at that venv's activate script.
+2. Write a small shim activate script (e.g. `~/envs/conda-shim.sh`) that activates conda and the desired env, then point `VENV_ACTIVATE` at it:
+  ```bash
+   # ~/envs/conda-shim.sh
+   source "$HOME/miniconda3/etc/profile.d/conda.sh"
+   conda activate <env_name>
+  ```
+
+### "Address already in use" during perf tests
+
+This means a node's kernel ephemeral-port pool was momentarily exhausted while preflight was building many communicators in a short window, so an outgoing `bind()` could not find a free port. It is a preflight-specific artifact of repeated communicator setup/teardown — a real training job builds its communicators once and reuses them — not a training failure mode. See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for details.
+
+Preflight has two complementary defenses:
+
+1. The **inter-node alltoall sub-group is internally capped at 16 nodes** (see [`preflight.md` §5.2](./preflight.md#52-group-sizes)) — the only test that, at large scale, opens enough simultaneous connections to approach the per-node ephemeral-port pool. The cap eliminates this failure mode by construction.
+2. A **global barrier + `--comm-cleanup-delay-sec` sleep** (default 2 s) is inserted after every comm destroy, primarily for cross-rank synchronization across the destroy → setup transition.
+
+If you still see `Address already in use` (e.g. on a network with an unusually narrow ephemeral-port range), the directly relevant **OS-level tuning** is widening that range — best-practice for any RDMA host:
+
+```bash
+# Widen the ephemeral port range from ~28k to ~64k. This is the
+# OS knob that directly addresses the binding constraint.
+sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535"
+```
+
+As a fallback, raise the per-phase delay:
+
+```bash
+# Bump the per-phase delay (default 2 s) on a particularly stressed
+# network. Rarely needed in practice with the §5.2 alltoall cap.
+runner/primus-cli direct -- preflight --comm-cleanup-delay-sec 5
+```
+
+See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for persistence and recommended large-cluster invocation patterns (split tests into separate runs, etc.).
+
+If the error occurs at `init_process_group` (before tests even start), it typically means a previous job left the rendezvous port (`MASTER_PORT`, default `1234`) in `TIME_WAIT`. Either wait ~60 s or use a different port:
+
+```bash
+export MASTER_PORT=1235
+```
+
+### Capturing full output
+
+The launcher already writes a complete log to `logs/log_<timestamp>.txt` (configurable via `--log_file PATH`), even under `--silent`. If you also want a copy at the call site, redirect there:
+
+```bash
+srun ... runner/primus-cli direct -- preflight --perf-test \
+    2>&1 | tee preflight-$(date +%Y%m%d-%H%M%S).log
+```
+
+---
+
+## 9. Automated node bisection (finding the bad node in an NCCL hang)
+
+When a cluster-wide preflight run hangs or fails, use
+[`tools/preflight_bisect/bisect.py`](https://github.com/AMD-AGI/Primus/blob/main/tools/preflight_bisect/bisect.py) to
+run `preflight --perf-test` on smaller Slurm node subsets until suspect nodes
+are isolated.
+
+### Prerequisites
+
+1. Working non-container preflight setup from the sections above, with
+  `VENV_ACTIVATE` exported from a shared filesystem path.
+2. Run from the SLURM login/head node, where both `scontrol` and `srun` are
+  available.
+3. Run from inside a Slurm allocation, or provide a Slurm nodelist explicitly.
+
+### Example from inside an allocation
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+mkdir -p output
+
+python tools/preflight_bisect/bisect.py \
+    --nodelist "$SLURM_NODELIST" \
+    --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \
+    --trial-timeout-sec 600 \
+    --slurm-time 00:15:00 \
+    --preflight-env USING_AINIC=1 \
+    --preflight-env NCCL_IB_GID_INDEX=1 \
+    --preflight-env NCCL_CROSS_NIC=1 \
+    --preflight-env NCCL_PXN_DISABLE=0 \
+    2>&1 | tee output/bisect-latest.log
+```
+
+Adjust the `--preflight-env` lines to match your cluster. Per-trial logs and a
+final `summary.txt` are written under `--output-dir`.
+
+> Note: Set `--trial-timeout-sec` high enough for a healthy subset to finish.
+> Too small a timeout can turn slow-but-good trials into false failures, causing
+> the bisection to explore extra paths.
+>
+> Note: `--preflight-env KEY=VALUE` values are concatenated into a single
+> `srun --export=ALL,...` argument, so values must not contain commas or
+> whitespace. Keep comma-containing values as normal exported environment
+> variables.
+
+---
+
+## 10. See also
+
+- [Preflight](./preflight.md)—full reference for the `preflight` subcommand and its flags
+- [CLI User Guide](./cli-reference.md)—container-based and `primus-cli slurm` workflows
+- [`runner/primus-cli-direct.sh`](https://github.com/AMD-AGI/Primus/blob/main/runner/primus-cli-direct.sh)—the direct launcher itself (`primus-cli direct` dispatches here)
+- [`primus/tools/preflight/`](https://github.com/AMD-AGI/Primus/tree/main/primus/tools/preflight)—preflight implementation
+- [`tools/preflight_bisect/bisect.py`](https://github.com/AMD-AGI/Primus/blob/main/tools/preflight_bisect/bisect.py)—bisect wrapper for narrowing down failing nodes in multi-node preflight runs
diff --git a/docs/02-user-guide/preflight.md b/docs/02-user-guide/preflight.md
new file mode 100644
index 000000000..eb1b325b4
--- /dev/null
+++ b/docs/02-user-guide/preflight.md
@@ -0,0 +1,417 @@
+# Preflight
+
+`preflight` is Primus' cluster diagnostic tool. It produces:
+
+- A **fast info report** (host / GPU / network configuration), and
+- A configurable suite of **performance tests** (GEMM TFLOPS, intra-node and inter-node communication bandwidth, P2P, ring P2P).
+
+Use it to spot misconfiguration, hardware degradation, or perf outliers **before** committing a large distributed training run to a global rendezvous.
+
+- **User-facing entry**: `primus-cli ... -- preflight [args]`
+- **No-container launcher**: `runner/primus-cli direct -- preflight ...` — see [`preflight-without-container.md`](./preflight-without-container.md).
+- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py` → `primus/tools/preflight/preflight_perf_test.py`.
+
+> Looking for a faster, distributed-rendezvous-free per-node screen? See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md). The recommended workflow is **smoke first, preflight second** — see [§10 Comparison with node-smoke](#10-comparison-with-node-smoke).
+
+---
+
+## 1. Two run modes (and how preflight picks one)
+
+Preflight has two report types, controlled by a single precedence rule:
+
+| Mode | Triggered by | What it does |
+|---|---|---|
+| **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. Emits a per-node report **without requiring a rendezvous**; multi-node aggregation then uses a **timeout-bounded** rendezvous (`--dist-timeout-sec`), so it never hangs indefinitely on network misconfig. |
+| **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. |
+| **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. |
+
+### Mode precedence
+
+1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.**
+   If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations.
+2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.**
+   Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them.
+3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests.
+
+The default order ensures you always get a report even if `torch.distributed` initialization later hangs.
+
+---
+
+## 2. Quick start
+
+### Info report only (fast)
+
+```bash
+primus-cli direct -- preflight --host --gpu --network
+```
+
+### Full preflight (info + every perf test)
+
+```bash
+primus-cli direct -- preflight
+```
+
+### Perf tests only
+
+```bash
+primus-cli direct -- preflight --perf-test
+```
+
+### Fast pre-launch sanity check
+
+```bash
+primus-cli direct -- preflight --quick
+```
+
+Equivalent on SLURM via `primus-cli slurm`:
+
+```bash
+primus-cli slurm srun -N 4 -- preflight --quick
+```
+
+Without a container, see [`preflight-without-container.md`](./preflight-without-container.md) for the equivalent `runner/primus-cli direct -- preflight ...` invocations.
+
+---
+
+## 3. Test selection (`--tests`)
+
+`--tests` takes a comma-separated list of canonical tokens (or `all`). Implies `--perf-test`.
+
+| Token | What it runs |
+|---|---|
+| `gemm` | Single-GPU square GEMM TFLOPS sweep. |
+| `intra-allreduce` | Intra-node `all_reduce` bandwidth at every selected `--intra-group-sizes` x `--intra-comm-sizes-mb`. |
+| `intra-alltoall` | Intra-node `all_to_all` bandwidth, same configuration matrix. |
+| `inter-allreduce` | Inter-node `all_reduce` bandwidth at every selected `--inter-group-sizes` x `--inter-comm-sizes-mb`. |
+| `inter-alltoall` | Inter-node `all_to_all` bandwidth, same configuration matrix. |
+| `inter-p2p` | Inter-node point-to-point send/recv between fixed **adjacent 2-node pairs** (does not use `--inter-group-sizes`). Sized by `--inter-comm-sizes-mb`, falling back to `--comm-sizes-mb`. |
+| `inter-ring-p2p` | Inter-node ring-pattern P2P, sized by `--ring-p2p-sizes-mb`. |
+| `all` | Every token above. Default when `--tests` is omitted. |
+
+Examples:
+
+```bash
+# GEMM only
+primus-cli direct -- preflight --tests gemm
+
+# Just the inter-node bandwidth tests
+primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall
+
+# Combine with size overrides
+primus-cli direct -- preflight \
+    --tests gemm,inter-allreduce \
+    --comm-sizes-mb 64,1024 \
+    --inter-group-sizes all
+```
+
+Unknown tokens fail fast (before any rendezvous):
+
+```text
+[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'.
+Valid tokens: gemm, intra-allreduce, intra-alltoall, inter-allreduce,
+inter-alltoall, inter-p2p, inter-ring-p2p, all
+```
+
+---
+
+## 4. Quick preset (`--quick`)
+
+`--quick` is the recommended **pre-launch sanity** preset. Implies `--perf-test`. It substitutes:
+
+| Knob | `--quick` value |
+|---|---|
+| `--tests` | `gemm,intra-allreduce,inter-allreduce` |
+| `--comm-sizes-mb` | `64,1024` |
+| `--intra-group-sizes` | `LOCAL_WORLD_SIZE` (full intra-node group only) |
+| `--inter-group-sizes` | `all` (full N-node group only) |
+| `warmup` | `5` |
+| `iteration` | `20` |
+
+**User-supplied flags override the preset.** For example:
+
+```bash
+# Quick preset, but with a custom size set
+primus-cli direct -- preflight --quick --comm-sizes-mb 32,256
+```
+
+A full perf run with default knobs takes minutes; `--quick` typically finishes in <60s on healthy hardware.
+
+---
+
+## 5. Tuning the perf tests
+
+All perf tuning knobs default to `None` so preflight can tell whether you set them. When unset, the documented defaults below apply.
+
+### 5.1 Message sizes (collective + P2P)
+
+| Flag | Default | Applies to |
+|---|---|---|
+| `--comm-sizes-mb CSV` | `2,4,8,16,32,64,128,256,512,1024` | Default for both intra- and inter-node `allreduce` / `alltoall` and `inter-p2p` when no specific override is given. |
+| `--intra-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **intra-node** `allreduce` / `alltoall`. |
+| `--inter-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **inter-node** `allreduce` / `alltoall` / `inter-p2p`. |
+
+```bash
+# Smaller, focused sweep
+primus-cli direct -- preflight --comm-sizes-mb 8,128
+
+# Different sizes for intra vs inter
+primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128 \
+    --intra-comm-sizes-mb 4,32
+```
+
+### 5.2 Group sizes
+
+| Flag | Default | Notes |
+|---|---|---|
+| `--intra-group-sizes CSV` | `2,4,8` | Each value must divide `LOCAL_WORLD_SIZE`. |
+| `--inter-group-sizes CSV` | `2,4,all` | `all` means the full N-node group. Other values are subgroup sizes. **Only `inter-allreduce` and `inter-alltoall` consult this flag** — for `inter-alltoall`, every requested per-group node count is internally capped at **16** before deduping (see "Inter-node alltoall is capped at 16 nodes" below), while `inter-allreduce` uses the requested sizes unchanged. `inter-p2p` and `inter-ring-p2p` ignore this flag (fixed adjacent 2-node pairs and a full-cluster ring, respectively). |
+
+```bash
+# All-GPU intra + full N-node inter only
+primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce \
+    --intra-group-sizes 8 \
+    --inter-group-sizes all
+```
+
+Validation is gated by which tests are actually selected. For example, `--tests gemm --intra-group-sizes 3` does **not** abort on a host with `LOCAL_WORLD_SIZE=8`; the intra-group constraint is only checked when an intra test is enabled.
+
+#### Inter-node alltoall is capped at 16 nodes
+
+Regardless of the cluster size or what `--inter-group-sizes` requests, the `inter-alltoall` test always runs on per-group node counts of at most **16**. Concretely, every requested value `G` is replaced with `min(G, 16)`, and the resulting list is deduped. Examples:
+
+| Cluster | `--inter-group-sizes` | Requested (resolved) | `inter-alltoall` actually runs |
+|---|---|---|---|
+| 8 N | `all` | `[8]` | `[8]` (no change) |
+| 64 N | `2,4,all` | `[2, 4, 64]` | `[2, 4, 16]` |
+| 128 N | `2,4,16,32,all` | `[2, 4, 16, 32, 128]` | `[2, 4, 16]` |
+| 128 N | `64` | `[64]` | `[16]` |
+
+When the cap actually changes the list, preflight emits a single one-line WARN to stdout so the row labels in the report (e.g. `alltoall-16nodes` instead of `alltoall-128nodes`) are not surprising.
+
+Why the cap, and why 16:
+
+- **It matches real-world usage.** Production MoE training rarely dispatches tokens across more than ~8 nodes (for example, DeepSeek-V3's largest published configuration uses `EP=64` over 8 nodes with per-token dispatch capped at 4 nodes). A 16-node ceiling covers every published configuration with comfortable headroom.
+- **It keeps the test from exhausting per-node network resources.** A large `inter-alltoall` sub-group opens a near-full mesh of connections per rank during communicator setup. Capping it at 16 keeps that well within a node's ephemeral-port budget and avoids spurious `Address already in use` failures at scale (see [§7](#7-running-on-very-large-clusters--64-nodes)).
+- **Other inter-node tests are unaffected.** The cap applies only to `inter-alltoall`: `inter-allreduce` uses `--inter-group-sizes` unchanged, while `inter-p2p` and `inter-ring-p2p` don't consult it at all. All three also open far fewer simultaneous connections than alltoall.
+- **It is intentionally not configurable.** This is a known-safe ceiling for the communication shapes preflight characterizes, not a tuning knob.
+
+### 5.3 Ring P2P sizes
+
+| Flag | Default | Applies to |
+|---|---|---|
+| `--ring-p2p-sizes-mb CSV` | `10,20,40,80,160` | `inter-ring-p2p` only. |
+
+```bash
+primus-cli direct -- preflight \
+    --tests inter-ring-p2p \
+    --ring-p2p-sizes-mb 5,20,80
+```
+
+### 5.4 Plotting
+
+| Flag | Effect |
+|---|---|
+| `--plot` | After each perf test, write per-size bandwidth bar charts under `<dump-path>/<test>/` and reference them in the markdown report. |
+
+---
+
+## 6. Reliability knobs
+
+Two knobs that are inert under happy-path conditions but matter at scale or on flaky networks.
+
+### 6.1 `--comm-cleanup-delay-sec FLOAT` (default `2.0`)
+
+Delay (seconds) inserted between destroying NCCL/RCCL process groups and creating new ones. It provides cross-rank synchronization across the destroy → setup transition, so a rank doesn't try to connect to a peer whose listener hasn't finished closing.
+
+- Default `2.0` is essentially free and worth keeping at every cluster size.
+- Set to `0` to disable the sleep entirely (barrier only).
+- Bump to e.g. `5` only on very flaky networks.
+
+```bash
+# Small/medium clusters: the default is fine. Override only if you
+# see port-reuse races on a very flaky network.
+primus-cli slurm srun -N 8 -- preflight --quick --comm-cleanup-delay-sec 5
+```
+
+See [§7](#7-running-on-very-large-clusters--64-nodes) for guidance on running at very large scale.
+
+### 6.2 `--dist-timeout-sec INT` (default `120`)
+
+Timeout (seconds) for `torch.distributed.init_process_group`. If init does not complete within this many seconds, preflight writes the info report (when applicable) plus a `Distributed Init` failure section to the markdown report, prints a clear error, and exits `2` — instead of hanging indefinitely.
+
+```bash
+# Fail fast if rendezvous does not work
+primus-cli direct -- preflight --perf-test --dist-timeout-sec 30
+```
+
+---
+
+## 7. Running on very large clusters (≥ 64 nodes)
+
+At very large scale there are a few practical considerations beyond what smaller runs encounter. A default `preflight` invocation still runs correctly at every scale we test (up to 128 nodes) without special flags — the points below are limitations to be aware of, plus recommendations that make large-cluster runs faster and easier to interpret.
+
+### 7.1 Limitations
+
+- **`inter-alltoall` is measured on at most 16 nodes per sub-group.** Regardless of cluster size or `--inter-group-sizes`, the alltoall test is capped at 16-node sub-groups (see [§5.2](#52-group-sizes)). This is intentional — it matches real-world MoE dispatch patterns and keeps the test from exhausting per-node network resources during communicator setup — but it does mean preflight will not report alltoall bandwidth for a larger topology. The cap applies only to `inter-alltoall`: `inter-allreduce` honors `--inter-group-sizes` unchanged, while `inter-p2p` and `inter-ring-p2p` don't use it at all.
+- **preflight briefly builds many communicators.** Unlike a real training job — which creates its communicators once at startup and reuses them — preflight repeatedly builds and tears down large communicators in a short window. On a cluster with an unusually narrow ephemeral-port range this can occasionally surface as `Address already in use` during setup. It is a preflight-specific artifact rather than a training failure mode; the one-line OS fix is in [§7.3](#73-optional-os-tuning).
+
+### 7.2 Recommended: split large runs by test family
+
+For clusters at or beyond ~128 nodes, run one test family per invocation instead of one large run. Each invocation stays short, and it becomes easy to see which specific communication shape is degraded if a number looks off.
+
+```bash
+# 1) GPU + intra-node fabric first (cheap, no inter-node OOB churn).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests gemm,intra-allreduce,intra-alltoall
+
+# 2) Inter-node DP-style collectives, all-nodes group only.
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-allreduce \
+    --inter-group-sizes all
+# Note: --inter-group-sizes all is honored here for inter-allreduce.
+# For inter-alltoall it would be capped at 16 (see §5.2).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-alltoall \
+    --inter-group-sizes all
+
+# 3) Inter-node PP-style ring P2P (the test that benefits most from
+#    isolation — it's the closest match to what real pipeline-parallel
+#    training actually exercises).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-ring-p2p
+
+# 4) Optional: pairwise inter-node P2P scan (useful for finding a
+#    single bad link, slower because it walks many pairs).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-p2p
+```
+
+Each invocation tears down its own `WORLD` on exit and touches only one `--tests` value, so you get a per-test wall clock, can re-run a single phase in isolation, and get a separate report per run via `--report-file-name`.
+
+### 7.3 Optional OS tuning
+
+`preflight` runs fine with default OS settings at every scale we test. If you do hit `Address already in use` on a cluster with a narrow ephemeral-port range, widen the range — this is good general practice for any RDMA host regardless of preflight:
+
+```bash
+# Widen the per-node ephemeral port range (default ~28k → ~64k ports).
+sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535"
+
+# Persist across reboots:
+echo 'net.ipv4.ip_local_port_range = 1024 65535' | sudo tee /etc/sysctl.d/99-large-cluster.conf
+sudo sysctl --system
+```
+
+---
+
+## 8. Reporting
+
+| Flag | Default | Effect |
+|---|---|---|
+| `--dump-path DIR` | `output/preflight` | Output directory for reports + plots. |
+| `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. |
+| `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. |
+
+Output files:
+
+| File | Produced when | Notes |
+|---|---|---|
+| `<name>.md` / `<name>.pdf` | Info-only mode, or default mode | Info report. |
+| `<name>_perf.md` / `<name>_perf.pdf` | Perf-only mode, or default mode | Perf report (GEMM + comm). |
+
+Only **rank 0** writes the report.
+
+### Perf report layout
+
+A `<name>_perf.md` produced by a default run contains, in order:
+
+1. (Optional) `> Note:` line listing dropped info selectors.
+2. `# Nodes` legend — `Node N → Hostname` table, used by every subsequent table to keep host columns compact.
+3. `=======IB Bandwidth roofline (GB/s)=======` — bandwidth of the first IB device on Node 0.
+4. Per enabled test, in this order: `gemm`, `intra-comm`, `inter-comm`, `inter-p2p`, `inter-ring-p2p`. Each section has a configuration line, a results table (Node / Rank / hostname / per-size GB/s), optional plots, and a per-rank wall-clock summary.
+5. `[Primus:Preflight] <test> done in <T>s` lines on stdout for at-a-glance progress on the launching shell.
+
+---
+
+## 9. Backward-compat aliases
+
+| Flag | Equivalent | Notes |
+|---|---|---|
+| `--check-host`, `--check-gpu`, `--check-network` | `--host`, `--gpu`, `--network` | Same behavior. Keep working for older scripts. |
+| `--no-split-nodes-subgroup` | `--inter-group-sizes all` **and** drops `inter-p2p` | Pre-`--tests`/`--inter-group-sizes` alias. Use the new flags in new scripts. |
+
+---
+
+## 10. Comparison with node-smoke
+
+| Aspect | `node-smoke` | `preflight` |
+|---|---|---|
+| Rendezvous | None — every node independent | Global `torch.distributed` |
+| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests |
+| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) |
+| Inter-node bandwidth matrix | Not tested (intentionally) | Yes (allreduce/alltoall/p2p/ring-p2p) |
+| Drift detection | Yes (versions, NIC firmware, port count) | No |
+| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only |
+| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF |
+
+**Recommended workflow**: run `node-smoke` first to exclude broken nodes, then run `preflight` on the surviving set to get cross-node bandwidth measurements. See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) §4 ("Quick start") for the integration commands.
+
+---
+
+## 11. Validation & error handling
+
+Preflight resolves the perf config **before** any distributed rendezvous. This means typos and bad sizes/group-sizes fail in seconds, not after a 120s NCCL init:
+
+```text
+[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'.
+[Primus:Preflight] ERROR: invalid perf config:
+    --intra-group-sizes: [3] do not divide LOCAL_WORLD_SIZE=8
+[Primus:Preflight] ERROR: invalid perf config: --comm-sizes-mb: values must be positive (got 0)
+```
+
+In info-only mode, perf-only tuning knobs trigger a single warning so you notice them but they don't abort:
+
+```text
+[Primus:Preflight] WARN: --comm-sizes-mb,--intra-group-sizes have no effect
+in info-only mode (no --perf-test/--tests/--quick).
+```
+
+In default mode where info selectors are dropped because perf intent was set, the preserved warning is also written into the perf report header:
+
+```text
+> Note: info selectors --host were dropped because perf mode
+> (--perf-test/--tests/--quick) takes precedence. Run them in a separate
+> invocation if you want both reports.
+```
+
+---
+
+## 12. Operational tips
+
+- **For multi-node runs, always use `primus-cli slurm` or `primus-cli direct` under `srun`** so distributed environment variables (`NNODES` / `NODE_RANK` / `MASTER_ADDR`) are set correctly.
+- **Make sure slurm requests GPU resources**. For some clusters, you may need to explicitly request GPU resources with `srun -N <nodes> --gpus-per-node=<gpus-per-node>`.
+- **Insufficient CPU cores cause >30x perf slowdowns** — pass `srun -c <cores-per-node>` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`.
+- **For a quick environment snapshot**, prefer `--host --gpu --network` — you always get a local per-node report even on a broken network, and any multi-node aggregation is timeout-bounded (`--dist-timeout-sec`), so the command never hangs.
+- **Between each communication test phase**, preflight performs a global barrier + `--comm-cleanup-delay-sec` sleep (default 2 s) for cross-rank sync across the destroy → setup transition. The default works at every cluster size we test up to 128 nodes. See [§7](#7-running-on-very-large-clusters--64-nodes) for large-cluster guidance.
+- **For pre-launch screening of a large cluster**, the recommended sequence is:
+  1. `node-smoke` to prune broken nodes (`failing_nodes.txt`).
+  2. `preflight --quick` on the surviving nodes for the perf sanity numbers.
+  3. `preflight` (full) on the same set if the `--quick` numbers raise a flag.
+
+---
+
+## 13. Running preflight without a container
+
+If you cannot (or prefer not to) use a container, see [`preflight-without-container.md`](./preflight-without-container.md) for the step-by-step `runner/primus-cli direct -- preflight ...` walkthrough — Python virtual-environment setup, SLURM invocation patterns, NCCL configuration for Broadcom and Pensando (AINIC) clusters, and many configurable-knob examples.
+
+---
+
+## 14. See also
+
+- [`preflight-without-container.md`](./preflight-without-container.md) — quick-start guide for `primus-cli direct -- preflight` (no container).
+- [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) — full guide for the per-node smoke test (screen + exclude bad nodes).
+- [`runner/primus-cli-direct.sh`](../../runner/primus-cli-direct.sh) — non-container launcher (`primus-cli direct` dispatches here).
+- [`primus/tools/preflight/`](../../primus/tools/preflight/) — implementation.
+- [`primus/tools/preflight/preflight_args.py`](../../primus/tools/preflight/preflight_args.py) — canonical CLI definition (single source of truth for flags + defaults).
diff --git a/docs/02-user-guide/pretraining.md b/docs/02-user-guide/pretraining.md
new file mode 100644
index 000000000..6734e98e5
--- /dev/null
+++ b/docs/02-user-guide/pretraining.md
@@ -0,0 +1,291 @@
+# Pretraining workflows
+
+Primus is a YAML-driven training stack for AMD GPUs. You select a **backend** (Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge), point `train pretrain` at a **configuration YAML**, and launch Primus with the unified CLI (`runner/primus-cli`) in **direct**, **container**, or **Slurm** mode. See [CLI reference](cli-reference.md) and [Configuration system](configuration-system.md).
+
+This section helps you understand concepts related to the Primus workflow: how backends work, YAML structure and inheritance, parallelism vocabulary, the full per-backend configuration inventory, and so on. If you already understand the concepts and just need the specific commands to run your training with Primus, see [Backend training recipes](training-recipes.md).
+
+---
+
+## Overview
+
+The following table describes the four backend types supported by Primus and their typical uses.
+
+| Backend | Framework | Typical use |
+| --- | --- | --- |
+| Megatron-LM | `framework: megatron` | Large-scale transformer pretraining with Megatron-style parallelism (TP/PP/EP). |
+| TorchTitan | `framework: torchtitan` | PyTorch-native scaled training (FSDP / tensor / pipeline / expert parallelism per config). |
+| MaxText (JAX) | `framework: maxtext` | JAX/MaxText single- and multi-node runs; parallelism via MaxText `ici_*` / `dcn_*` settings. |
+| Megatron Bridge | `framework: megatron_bridge` | Bridge-oriented workflows (configure like other backends; see parameter reference). |
+
+> Several setup steps apply to **all** backends (mock vs. real data, Hugging Face tokens, scaling to multiple nodes, and HipBLASLt autotuning). After you read the backend section that applies to you, see [Common patterns](#common-patterns) below.
+
+---
+
+## Megatron-LM pretraining
+
+### Quick start (container mode)
+
+From the root of the clone of the [Primus repository](https://github.com/AMD-AGI/Primus), with Docker or Podman available, the following command starts the training in container mode:
+
+```bash
+./runner/primus-cli container -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+This uses the default image from `runner/.primus.yaml` (`rocm/primus:v26.4` unless overridden). The project tree is mounted into the container automatically by `runner/primus-cli-container.sh`.
+
+### Example configurations under `examples/megatron/configs/MI300X/`
+
+The following files ship in the repository (sorted by name). Parallelism columns are taken from `tensor_model_parallel_size` / `pipeline_model_parallel_size` / `expert_model_parallel_size` in each file (literals or `${PRIMUS_TP:…}` defaults).
+
+| Config | TP | PP | EP |
+| --- | --- | --- | --- |
+| `deepseek_v2-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:4}` | `${PRIMUS_EP:8}` |
+| `deepseek_v2-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:4}` | `${PRIMUS_EP:8}` |
+| `deepseek_v2_lite-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `deepseek_v2_lite-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `deepseek_v3-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `deepseek_v3-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `gpt_oss_20B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `gpt_oss_20B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `grok1-BF16-pretrain.yaml` | `1` | `4` | `8` |
+| `grok1-FP8-pretrain.yaml` | `1` | `4` | `8` |
+| `grok2-BF16-pretrain.yaml` | `1` | `4` | `8` |
+| `grok2-FP8-pretrain.yaml` | `1` | `4` | `8` |
+| `llama2_13B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama2_13B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama2_70B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama2_70B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama2_7B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama2_7B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.1_405B-BF16-pretrain.yaml` | `8` | `8` | `1` |
+| `llama3.1_405B-FP8-pretrain.yaml` | `8` | `8` | `1` |
+| `llama3.1_70B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.1_70B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.1_8B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.1_8B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.2_1B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.2_1B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.2_3B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.2_3B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.3_70B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3.3_70B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3_70B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3_70B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3_8B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `llama3_8B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `llama4_17B128E-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `llama4_17B128E-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `llama4_17B16E-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `llama4_17B16E-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `mamba_370M-pretrain.yaml` | `1` | `1` | `1` |
+| `mixtral_8x22B_v0.1-BF16-pretrain.yaml` | `1` | `4` | `8` |
+| `mixtral_8x22B_v0.1-FP8-pretrain.yaml` | `1` | `4` | `8` |
+| `mixtral_8x7B_v0.1-BF16-pretrain.yaml` | `1` | `1` | `8` |
+| `mixtral_8x7B_v0.1-FP8-pretrain.yaml` | `1` | `1` | `8` |
+| `qwen2.5_14B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_14B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_32B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_32B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_3B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_3B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_72B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_72B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_7B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen2.5_7B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_14B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_14B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_235B_A22B-BF16-pretrain.yaml` | `1` | `1` | `8` |
+| `qwen3_235B_A22B-FP8-pretrain.yaml` | `1` | `1` | `8` |
+| `qwen3_30B_A3B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `qwen3_30B_A3B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `qwen3_32B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_32B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_4B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_4B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_5_35B_A3B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `qwen3_5_35B_A3B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` |
+| `qwen3_8B-BF16-pretrain.yaml` | `1` | `1` | `1` |
+| `qwen3_8B-FP8-pretrain.yaml` | `1` | `1` | `1` |
+| `zebra_llama_1B-pretrain.yaml` | `1` | `1` | `1` |
+| `zebra_llama_3B-pretrain.yaml` | `1` | `1` | `1` |
+| `zebra_llama_8B-pretrain.yaml` | `1` | `1` | `1` |
+
+### Sample YAML file (`llama2_7B-BF16-pretrain.yaml`) explained
+
+Path: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml`
+
+| Section | Role |
+| --- | --- |
+| `work_group`, `user_name`, `exp_name`, `workspace` | Run identity and output root (supports `${VAR:default}` substitution). |
+| `modules.pre_trainer.framework` | `megatron` selects Megatron-LM integration. |
+| `config: pre_trainer.yaml` | Module preset under `primus/configs/modules/megatron/`. |
+| `model: llama2_7B.yaml` | Model preset under `primus/configs/models/megatron/` (extends `llama2_base.yaml` → …). |
+| `overrides` | Run-specific training knobs: iterations, batching, LR, **parallelism** (`tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size`), data paths, checkpoints, Primus Turbo flags, etc. |
+
+The sample sets `mock_data: true` and `train_data_path: null` so you can validate the stack without real corpora.
+
+### Mock data versus real data
+
+- **Mock data:** Set `mock_data: true` and leave `train_data_path` / `valid_data_path` empty (as in `llama2_7B-BF16-pretrain.yaml`).
+- **Real data:** Set `mock_data: false` and populate Megatron-compatible data paths (and tokenizer assets) in `overrides`. Use paths visible inside your container mounts.
+
+### Multi-node training with Slurm
+
+```bash
+./runner/primus-cli slurm srun -N 4 -p <partition> -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+`runner/primus-cli-slurm-entry.sh` derives `MASTER_ADDR`, `NNODES`, and `NODE_RANK` from Slurm and forwards them into the container. Align `tensor_model_parallel_size`, `pipeline_model_parallel_size`, and `expert_model_parallel_size` with your cluster width and job size.
+
+---
+
+## TorchTitan pretraining
+
+### Quick start
+
+```bash
+./runner/primus-cli container -- train pretrain \
+  --config examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+### Example configurations under `examples/torchtitan/configs/MI300X/`
+
+| File |
+| --- |
+| `deepseek_v3_16b-BF16-pretrain.yaml` |
+| `deepseek_v3_16b-FP8-pretrain.yaml` |
+| `deepseek_v3_236b-BF16-pretrain.yaml` |
+| `deepseek_v3_236b-FP8-pretrain.yaml` |
+| `deepseek_v3_671b-pretrain.yaml` |
+| `llama3.1_405B-BF16-pretrain.yaml` |
+| `llama3.1_405B-FP8-pretrain.yaml` |
+| `llama3.1_70B-BF16-pretrain.yaml` |
+| `llama3.1_70B-FP8-pretrain.yaml` |
+| `llama3.1_8B-BF16-pretrain.yaml` |
+| `llama3.1_8B-FP8-pretrain.yaml` |
+| `llama4_17Bx128E-BF16-pretrain.yaml` |
+| `llama4_17Bx128E-FP8-pretrain.yaml` |
+| `llama4_17Bx16E-BF16-pretrain.yaml` |
+| `llama4_17Bx16E-FP8-pretrain.yaml` |
+| `qwen3_0.6B-pretrain.yaml` |
+| `qwen3_1.7B-pretrain.yaml` |
+| `qwen3_14B-pretrain.yaml` |
+| `qwen3_32B-pretrain.yaml` |
+| `qwen3_4B-pretrain.yaml` |
+| `qwen3_8B-pretrain.yaml` |
+
+### Sample YAML file (`llama3.1_8B-BF16-pretrain.yaml`) explained
+
+Path: `examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml`
+
+| Section | Role |
+| --- | --- |
+| `framework: torchtitan` | Selects the TorchTitan integration. |
+| `config: pre_trainer.yaml` | Module preset under `primus/configs/modules/torchtitan/`. |
+| `model: llama3.1_8B.yaml` | Model preset under `primus/configs/models/torchtitan/`. |
+| `overrides.training`, `lr_scheduler`, `activation_checkpoint`, `primus_turbo` | Run-specific batching, steps, checkpointing, and Turbo options. |
+
+Some configurations omit an explicit `parallelism:` block; in that case the default values come from the **module and model presets** (`primus/configs/modules/torchtitan/pre_trainer.yaml` and the chosen model YAML). Other examples (for example DeepSeek and Qwen) set `parallelism:` inline with `tensor_parallel_degree`, `pipeline_parallel_degree`, `expert_parallel_degree`, etc.
+
+---
+
+## MaxText (JAX) pretraining
+
+### Quick start
+
+```bash
+./runner/primus-cli container -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+```
+
+### JAX-specific requirements
+
+Install JAX/MaxText dependencies from the repository root:
+
+```bash
+pip install -r requirements-jax.txt
+```
+
+### Example configurations under `examples/maxtext/configs/MI300X/`
+
+| File | Key parallelism (`ici_*` intra-node, `dcn_*` inter-node) |
+| --- | --- |
+| `deepseek_v2_16B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `grok1-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama2_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama2_7B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3.3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `llama3_8B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `mixtral_8x7B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `qwen3_14B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+| `qwen3_30B_A3B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` |
+
+The `llama2_7B-pretrain.yaml` example also sets `dataset_type: "synthetic"` and `hf_access_token: ${HF_TOKEN:""}` for gated Hugging Face assets when you switch to real data.
+
+---
+
+## Common patterns
+
+### Testing with mock data
+
+Set `mock_data: true` (Megatron/TorchTitan) or synthetic dataset settings (MaxText) to validate the configurations and infrastructure without I/O-heavy datasets.
+
+### Real training data
+
+- Megatron: Configure `train_data_path` / `valid_data_path` and tokenizer assets in `overrides` once `mock_data` is false.
+- For **all backends**, ensure host paths are mounted in **container** mode (`--volume` or `container.options.volume` in YAML).
+- TorchTitan/MaxText: Follow backend-specific dataset fields in the `overrides` and presets.
+
+### Scaling from single-node to multi-node
+
+- Use **Slurm** mode for allocation; keep the **container** entry if you want the same image on every node.
+- Set environment variables consistently (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`); the Slurm entry script injects them when using `primus-cli slurm`.
+- Increase values in the parallelism fields (Megatron TP/PP/EP; TorchTitan `parallelism`; MaxText `ici_*` / `dcn_*`) to match topology.
+
+### Hugging Face token for gated models
+
+Export `HF_TOKEN` on the host before launching **container** mode; `runner/.primus.yaml` lists `HF_TOKEN` under `container.options.env` so it can be forwarded into the container. MaxText configurations may reference `${HF_TOKEN:""}` directly.
+
+### hipBLASLt autotuning (three stages)
+
+Controlled with `PRIMUS_HIPBLASLT_TUNING_STAGE` (see `examples/README.md`):
+
+| Stage | Purpose |
+| --- | --- |
+| 1 | Dump GEMM shapes seen during training (reduce `train_iters` for faster collection). |
+| 2 | Tune kernels from dumped shapes (offline tooling under `examples/offline_tune`). |
+| 3 | Train using tuned kernel artifacts from `./output/tune_hipblaslt/...`. |
+
+Example (from in-repo docs):
+
+```bash
+export PRIMUS_HIPBLASLT_TUNING=1        # master switch (required; tuning is skipped without it)
+export PRIMUS_HIPBLASLT_TUNING_STAGE=1
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+---
+
+## Supported models
+
+The tables above in the Megatron, TorchTitan, and MaxText sections are curated MI300X examples from the [Primus repository](https://github.com/AMD-AGI/Primus). Use `examples/<backend>/configs/` in the repository as the authoritative inventory, as new presets and hardware-specific examples may be added there before this document is updated to reflect their additions.
+
+| Backend | Example region | Parallelism vocabulary |
+| --- | --- | --- |
+| Megatron-LM | `examples/megatron/configs/MI300X/` | `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size` (and env-driven `${PRIMUS_TP:…}` variants). |
+| TorchTitan | `examples/torchtitan/configs/MI300X/` | `parallelism.*` (e.g. `tensor_parallel_degree`, `pipeline_parallel_degree`, `expert_parallel_degree`, FSDP shard settings). |
+| MaxText | `examples/maxtext/configs/MI300X/` | `ici_fsdp_parallelism`, `ici_data_parallelism`, `dcn_fsdp_parallelism`, `dcn_data_parallelism`. |
+
+For scripting patterns that predate `primus-cli`, the repository still documents `examples/run_local_pretrain.sh` and `examples/run_slurm_pretrain.sh` in `examples/README.md`; equivalent launches are shown above using `./runner/primus-cli`.
+
+---
+
+## Related documentation
+
+- [CLI reference](cli-reference.md): launcher usage
+- [Configuration system](configuration-system.md): YAML merge rules
+- Backend parameter references: [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md), [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)
diff --git a/docs/02-user-guide/primus-tools.md b/docs/02-user-guide/primus-tools.md
new file mode 100644
index 000000000..212a3fd73
--- /dev/null
+++ b/docs/02-user-guide/primus-tools.md
@@ -0,0 +1,29 @@
+# Primus tools
+
+A quick catalog of the tools that ship with Primus and the sibling projects
+around it—command-line tools, the tuning agent, ecosystem projects, and
+auxiliary utilities. Each row gives a short description and a how-to starting
+point; follow a tool's link for the full reference.
+
+| Tool | Type | What it does | How to use |
+|------|------|--------------|------------|
+| [`train`](./pretraining.md) | CLI | Launch pretraining or post-training on any backend from a YAML configuration. | `primus-cli <mode> -- train pretrain --config <yaml>` |
+| [`benchmark`](./benchmarking.md) | CLI | GEMM, RCCL, and attention microbenchmarks for hardware and stack validation. | `primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096` |
+| [`preflight`](./preflight.md) | CLI | Host, GPU, and network health checks before long jobs. | `primus-cli slurm srun -N 4 -- preflight --host --gpu --network` |
+| [`projection`](./projection.md) | CLI | Estimate per-GPU memory and throughput without occupying a full cluster. | `primus-cli direct -- projection both --config <yaml>` |
+| [Tuning agent](./tuning-agent.md) | Agent | LLM-driven search for a near-optimal training configuration, scored by projection. | `python -m primus.agents.tuning_agent --workload <yaml> --target-cluster <yaml>` |
+| [Primus-LM](../01-getting-started/quickstart.md) | Ecosystem | The training framework in this repository (multi-backend, unified CLI). | See the [Quickstart](../01-getting-started/quickstart.md) |
+| [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) | Ecosystem | High-performance ROCm operators (attention, GEMM, grouped GEMM, DeepEP, FP8/FP4). | Bundled in the `rocm/primus` image; enabled via configuration flags |
+| [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) | Ecosystem | Kubernetes-native stability, scheduling, and fault-tolerance platform. | Deployed separately on Kubernetes (Helm) |
+| [IRLens](../../tools/IRLens/README.md) | Auxiliary | Parse XLA HLO dumps into a communication-vs-compute execution skeleton. | See the [README](../../tools/IRLens/README.md) |
+| [model_stats](../../tools/model_stats/README.md) | Auxiliary | Chart model dimensions from the config registry. | See the [README](../../tools/model_stats/README.md) |
+| [Pipeline visualization](../../tools/visualization/pp_vis/README.md) | Auxiliary | Render pipeline-parallel schedules in a local web UI. | See the [README](../../tools/visualization/pp_vis/README.md) |
+| [Auto benchmark](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) | Auxiliary | Interactive Megatron/TorchTitan benchmark menu with metrics collection. | See the [README](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) |
+
+---
+
+## Related documentation
+
+- [CLI reference](./cli-reference.md)—full launcher grammar and subcommand options.
+- [Tooling](../06-developer-guide/tooling.md)—developer-guide index of the `tools/` utilities.
+- [Project overview](../01-getting-started/overview.md#primus-ecosystem)—how the Primus ecosystem layers fit together.
diff --git a/docs/02-user-guide/projection.md b/docs/02-user-guide/projection.md
new file mode 100644
index 000000000..af64941e5
--- /dev/null
+++ b/docs/02-user-guide/projection.md
@@ -0,0 +1,177 @@
+# Memory and performance projection
+
+Primus projection tools estimate **per-GPU memory** and **training throughput** for large-scale distributed jobs without requiring the full target cluster. Two modes are available: analytical **memory** projection and **performance** projection that combines profiling with simulation.
+
+**Implementation:** `primus/cli/subcommands/projection.py`
+
+| Mode | Command | Role |
+|------|---------|------|
+| **Memory** | `projection memory` | Estimates per-GPU memory (parameters, optimizer state, activations) using analytical formulas. |
+| **Performance** | `projection performance` | Benchmarks on a single node (or sub-node), then projects training time to multi-node configurations. |
+| **Both** | `projection both` | Runs a single benchmark and produces **both** the performance and (benchmark-anchored) memory projections from it. Recommended for cluster-sizing workflows. |
+
+**Core logic**
+
+- Memory: `primus/core/projection/memory_projection/`
+- Performance: `primus/core/projection/performance_projection/`
+
+Related: [Benchmark suite](./benchmarking.md), [Preflight diagnostics](./preflight.md), [Megatron parameters](../03-configuration-reference/megatron-parameters.md).
+
+---
+
+## Memory projection
+
+### Quick start
+
+```bash
+export NNODES=1
+export HSA_NO_SCRATCH_RECLAIM=1
+
+./runner/primus-cli direct --script primus/cli/main.py -- \
+  projection memory \
+  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml
+```
+
+Adjust `--config` to your experiment YAML. Memory estimation is analytical; the CLI still expects a normal Primus launch path (including distributed initialization where applicable).
+
+### What it estimates
+
+| Component | Meaning |
+|-----------|---------|
+| **Parameter memory** | Model weights assigned to this GPU (respecting parallelism). |
+| **Optimizer memory** | Optimizer state (for example Adam moments), accounting for sharding across data-parallel groups. |
+| **Activation memory** | Activations retained for the backward pass for a given microbatch and sequence length. |
+
+The tool walks a hierarchical profiler structure aligned with the model (embeddings, dense and MoE layers, output head, loss) and aggregates per-component contributions.
+
+### How to interpret results
+
+Console output includes per-component breakdowns and a summary such as parameter count, param+optimizer memory, activation memory for the configured batch size and sequence length, and a projected total. Use these to answer whether a configuration fits in HBM before you allocate large clusters.
+
+---
+
+## Performance projection
+
+### Quick start
+
+Minimum required nodes (derived from parallelism):
+
+```bash
+export NNODES=1
+export HSA_NO_SCRATCH_RECLAIM=1
+
+./runner/primus-cli direct --script primus/cli/main.py -- \
+  projection performance \
+  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml
+```
+
+### How it works
+
+1. **Profile** layer-level behavior on **one node** (or a subset of GPUs with automatic scaling rules).
+2. **Simulate** pipeline scheduling, data parallelism, and communication using analytical models.
+3. **Project** iteration time and tokens/s to a **target** node count when you specify one.
+
+### Projecting to a specific node count
+
+```bash
+./runner/primus-cli direct --script primus/cli/main.py -- \
+  projection performance \
+  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \
+  --target-nodes 4
+```
+
+If `--target-nodes` is omitted, the tool defaults to the **minimum** number of nodes implied by your parallelism configuration (TP, PP, EP, CP, GPUs per node).
+
+### Parallelism overrides (environment)
+
+You can override parallelism for what-if analysis:
+
+```bash
+export PRIMUS_TP=1
+export PRIMUS_PP=3
+export PRIMUS_EP=8
+
+./runner/primus-cli direct --script primus/cli/main.py -- \
+  projection performance \
+  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \
+  --target-nodes 6
+```
+
+---
+
+## Command reference
+
+### Syntax
+
+```bash
+primus-cli [global-options] <mode> [mode-args] -- projection {memory,performance,both} [options]
+```
+
+### Shared options (both modes)
+
+| Option | Description |
+|--------|-------------|
+| `--config` / `--exp` | Path to the Primus YAML configuration (**required**). |
+| `--data_path` | Data directory (default `./data` when included on the parser). |
+| `--backend_path` | Optional Megatron/TorchTitan import path appended to `PYTHONPATH`. |
+| `--export_config` | Accepted by the shared pretrain parser, but the default core runtime does not currently write a resolved YAML file. |
+
+### Performance-only options
+
+| Option | Description |
+|--------|-------------|
+| `--target-nodes` | Target number of nodes for scaling projection. Defaults to the minimum nodes required by TP/PP/EP/CP and GPUs per node. |
+| `--target-num-nodes` | Alias-style projection override for target node count. |
+| `--target-ep-size` | Override `expert_model_parallel_size` for the projection target. |
+| `--benchmark-gpus` | Use fewer than `GPUS_PER_NODE` GPUs for benchmarking; results are scaled analytically back to a full node. |
+| `--hardware-config` | YAML file with hardware parameters for communication modeling. |
+| `--profiling-mode` | `benchmark` (default, uses GPU), `simulate` (analytical / Origami GEMM + SDPA models, no GPU), or `both` (side-by-side). |
+| `--gemm-backend` | GEMM simulation backend when profiling is simulated (`origami`). |
+| `--gpu-arch` | Target architecture for simulation (for example `mi300x`, `gfx942`, `mi355x`, `gfx950`); can use `PRIMUS_GPU_ARCH`. |
+| `--gpu-clock-mhz` | Override GPU clock in MHz for simulation; can use `PRIMUS_GPU_CLOCK_MHZ`. |
+| `--pipeline-schedule-algorithm` | Pipeline simulation scheduler (`auto`, zero-bubble variants, or `all` for comparison). |
+| `--enable-zero-bubble` | Enable zero-bubble pipeline scheduling for projection. |
+| `--enable-deepep` | Enable DeepEP overlap modeling. |
+| `--sync-free-stage` | Override Sync-Free MoE stage (`0` off; stages `1`-`3` enable additional modeling assumptions). |
+| `--num-virtual-stages-per-pipeline-rank` | Override virtual pipeline stage count for projection. |
+| `--micro-batch-size`, `--global-batch-size` | Override batch sizes for projection without editing the YAML. |
+
+---
+
+## Assumptions and limitations
+
+### Assumptions (performance projection)
+
+1. **Data-parallel scaling**—Compute time scales with ideal weak-scaling assumptions versus data-parallel width.
+2. **Communication model**—Uses simplified bandwidth and latency models (defaults such as efficiency factors may apply).
+3. **Pipeline scheduling**—Bubble and overlap behavior is modeled with fixed splits; real frameworks may differ.
+4. **Gradients and MoE**—Gradient all-reduce overlap and MoE all-to-all behavior follow the implemented model (for example overlap flags, EP scaling).
+
+### Limitations
+
+1. **Single-node benchmark accuracy**—Reduced PP/EP on the benchmark GPU count may not capture every production behavior.
+2. **Contention**—Network contention between jobs is not modeled.
+3. **Memory vs speed**—Activation recomputation reduces memory but adds compute; performance projection may not fully reflect that trade-off unless modeled.
+4. **Heterogeneity**—Assumes homogeneous nodes; GPU frequency drift across nodes is not modeled.
+
+---
+
+## Tips
+
+1. Run **`projection memory`** first to confirm a configuration is feasible in HBM before spending time on performance projection.
+2. Always establish a **single-node** baseline before interpreting multi-node projections.
+3. **Data-parallel scaling** is bounded by batching: if you run out of microbatches (`global_batch_size` / `micro_batch_size`), adding nodes may not increase throughput.
+4. If the YAML **requires** multiple nodes (for example large PP), the performance path may automatically reduce parallelism for benchmarking and restore it analytically—read the console summary carefully.
+5. **No GPU available:** use `--profiling-mode simulate` for CPU-side analytical timing.
+6. **Validate models:** use `--profiling-mode both` to compare GPU benchmark timing with simulation on the same config.
+7. For **MoE** models, activation memory from MoE layers often dominates; memory projection highlights when recomputation is worth considering.
+
+---
+
+## Related documentation
+
+- [Benchmark suite](./benchmarking.md)
+- [Preflight diagnostics](./preflight.md)
+- [Post-training workflows](./posttraining.md)
+- [Tuning agent](./tuning-agent.md)
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md)
diff --git a/docs/02-user-guide/training-recipes.md b/docs/02-user-guide/training-recipes.md
new file mode 100644
index 000000000..dec02128a
--- /dev/null
+++ b/docs/02-user-guide/training-recipes.md
@@ -0,0 +1,260 @@
+# Backend training recipes
+
+Task-oriented, copy-paste commands for launching pretraining runs with each Primus backend on AMD Instinct™ GPUs.
+
+This section is for users who already know what they want to run and need the specific command for a given model, precision, and GPU. To understand concepts related to the Primus workflow (how backends work, YAML structure and inheritance, parallelism vocabulary, the full per-backend configuration inventory, etc.), see **[Pretraining](pretraining.md)**.
+
+> **Authoritative full matrices.** AMD publishes per-model reproduction pages with verified images, commits, and tuned batch sizes for every supported model at the following locations—treat them as the source of truth for achieving the expected performance; this page only gives the canonical *pattern* plus a representative example per backend and links to other reference materials.
+>
+> - [Training with Primus + Megatron-LM](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.html)
+> - [Training with Primus + PyTorch (TorchTitan)](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.html)
+> - [Training with Primus + JAX MaxText](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.html)
+
+---
+
+## How recipes are structured
+
+Every recipe follows the same four-step pattern:
+
+1. **Pull and launch the AMD Docker image** (for a reproducible environment).
+2. **Set the GPU-architecture environment** (performance environment variable settings differ by GPU).
+3. **Pick the configuration YAML** for your GPU architecture under `examples/<backend>/configs/<ARCH>/` in the [Primus repository](https://github.com/AMD-AGI/Primus).
+4. **Launch** with `runner/primus-cli` in `direct`, `container`, or `slurm` mode.
+
+### GPU-architecture config folders
+
+Configuration YAMLs are organized by GPU architecture. Always pick the folder that matches your hardware:
+
+
+| Backend                             | `MI300X` | `MI325X` | `MI355X` / `MI350X` |
+| ----------------------------------- | -------- | -------- | ------------------- |
+| `examples/megatron/configs/`        | yes      | yes      | yes                 |
+| `examples/torchtitan/configs/`      | yes      | yes      | yes                 |
+| `examples/maxtext/configs/`         | yes      | —        | yes                 |
+| `examples/megatron_bridge/configs/` | yes      | —        | yes                 |
+
+
+> MI350X uses the same configurations as MI355X because both are based on the gfx950 architecture. If a configuration for your model is not available in the architecture-specific folder, use the closest match from the same generation as a starting point.
+
+### GPU-architecture environment variables
+
+MI300X and MI325X benefit from the following performance settings; MI355X/MI350X do **not** need them:
+
+```bash
+# MI300X / MI325X only -- improves performance
+export HSA_NO_SCRATCH_RECLAIM=1
+export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+```
+
+### Choosing the Docker image
+
+For **container** and **Slurm** modes (direct mode runs in whatever environment you launched it from), the default image is `rocm/primus:v26.4` (`runner/.primus.yaml`). For reproducing published benchmarks, use the AMD-published tag for your release (the AMD pages under **Authoritative full matrices** above list the most current tag). JAX MaxText has its own separate image family of `rocm/jax-training:maxtext-...`.
+
+Image is picked in the priority order of `DOCKER_IMAGE` environment variable > `--image` CLI argument > config file. See [Selecting the container image](../01-getting-started/quickstart.md#selecting-the-container-image) for a full explanation, and [Configuration system](configuration-system.md) for configuration loading.
+
+---
+
+## Shared setup for all backends
+
+These apply across all backends. Set them up before running the recipes below.
+
+### Hugging Face token (for gated models or real data)
+
+```bash
+export HF_TOKEN=<your_hf_token>
+```
+
+`runner/.primus.yaml` forwards `HF_TOKEN` into the container automatically. MaxText configurations might also read `${HF_TOKEN:""}` directly.
+
+### Mock vs. real data
+
+- **Mock/synthetic data** (default for most examples): validates the stack without datasets. Megatron and TorchTitan set `mock_data: true`; MaxText sets `dataset_type: "synthetic"`.
+- **Real data:** set `mock_data: false` and point `train_data_path` (for Megatron) or the backend's dataset fields at paths visible *inside* your container mounts.
+
+### Multi-node networking checklist
+
+The `primus-cli` launcher sets sensible `NCCL_`* defaults, but auto-detection can pick the wrong device on multi-NIC nodes. Before multi-node, confirm and export if needed:
+
+```bash
+export NCCL_IB_HCA=<your_rdma_interfaces>      # from `ibv_devices`
+export NCCL_SOCKET_IFNAME=<your_net_interface> # from `ip a`
+export GLOO_SOCKET_IFNAME=<same_as_NCCL_SOCKET_IFNAME>
+export NCCL_IB_GID_INDEX=3                      # 3 for RoCE (1 for AMD AINIC)
+```
+
+For AMD AINIC clusters also set `USING_AINIC=1`, `NCCL_PXN_DISABLE=0`, `NCCL_IB_GID_INDEX=1`. See [Multi-Node Networking](../04-technical-guides/multi-node-networking.md) for the full reference.
+
+---
+
+## Megatron-LM
+
+**Image:** `rocm/primus`  |  **Configurations:** `examples/megatron/configs/<ARCH>/`  |  **Precisions:** BF16, FP8
+
+### 1. Launch the container
+
+```bash
+docker pull rocm/primus:v26.4
+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.4
+```
+
+Access the container later with `docker start primus_training_env && docker exec -it primus_training_env bash`.
+
+### 2. Run pretraining (direct mode, inside the container)
+
+Pretrain Llama 3.1 8B BF16 on **MI355X / MI350X**:
+
+```bash
+./runner/primus-cli direct \
+  --log_file /tmp/primus_llama3.1_8B.log \
+  -- train pretrain \
+  --config examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+Pretrain the same model on **MI300X / MI325X** (add the performance environment variables):
+
+```bash
+export HSA_NO_SCRATCH_RECLAIM=1
+export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+
+./runner/primus-cli direct \
+  --log_file /tmp/primus_llama3.1_8B.log \
+  -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+Switch model or precision by changing the config filename (e.g. `llama3.1_70B-FP8-pretrain.yaml`, `mixtral_8x7B_v0.1-BF16-pretrain.yaml`). The full configuration inventory is the repository's `examples/megatron/configs/<ARCH>/` directory. See the parallelism table in [Pretraining](pretraining.md#example-configurations-under-examplesmegatronconfigsmi300x).
+
+**Model-specific notes:**
+
+- **Zebra-Llama** (hybrid Mamba+MLA) pretrain presets ship at `examples/megatron/configs/<ARCH>/zebra_llama_{1B,3B,8B}-pretrain.yaml` and run via the standard core runtime; Megatron Bridge SFT variants live under `examples/megatron_bridge/configs/<ARCH>/`.
+- **MoE models** (DeepSeek-V2-Lite, Mixtral) might need extra grouped-GEMM or router flags; [Training with Primus + Megatron-LM](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.html) lists the exact flags per model.
+
+### 3. Multi-node (Slurm mode)
+
+```bash
+./runner/primus-cli slurm srun -N 8 -p <partition> -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml \
+  --micro_batch_size 4 --global_batch_size 1024
+```
+
+Scale batch size with node count and align `tensor_model_parallel_size`, `pipeline_model_parallel_size`, and `expert_model_parallel_size` to your topology. See the [multi-node networking checklist](#multi-node-networking-checklist) above.
+
+---
+
+## TorchTitan (PyTorch)
+
+**Image:** `rocm/primus`  |  **Configurations:** `examples/torchtitan/configs/<ARCH>/`  |  **Precisions:** BF16, FP8
+
+Use the same `rocm/primus` container as Megatron (step 1 above). TorchTitan parameters use a dotted namespace (e.g. `--training.local_batch_size`).
+
+### Run pretraining (direct mode)
+
+Pretrain Llama 3.1 8B BF16 on **MI355X / MI350X**:
+
+```bash
+./runner/primus-cli direct \
+  --log_file /tmp/primus_llama3.1_8B.log \
+  -- train pretrain \
+  --config examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml
+```
+
+On **MI300X / MI325X**, export the performance environment variables first (see above) and use the `MI300X` config path.
+
+### Multi-node (Slurm mode)
+
+```bash
+./runner/primus-cli slurm srun -N 4 -- train pretrain \
+  --config examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml \
+  --training.local_batch_size 6 \
+  --training.global_batch_size 192 \
+  --training.mock_data True
+```
+
+Available models include Llama 3.1 (8B/70B/405B), Llama 4, DeepSeek V3, and Qwen 3. See the `examples/torchtitan/configs/<ARCH>/` directory in the repository.
+
+---
+
+## JAX MaxText
+
+**Image:** `rocm/jax-training:maxtext-...` (separate family from the other backends)  |  **Configurations:** `examples/maxtext/configs/<ARCH>/`
+
+MaxText uses a different Docker image than Megatron and TorchTitan, and it is **not** the default image pointed to in `runner/.primus.yaml`. In container or Slurm mode, you must point Primus at your MaxText image explicitly.
+
+### 1. Launch the container
+
+```bash
+docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+docker run -it \
+    --device /dev/dri --device /dev/kfd \
+    --network host --ipc host \
+    --group-add video --cap-add SYS_PTRACE \
+    --security-opt seccomp=unconfined --privileged \
+    -v $HOME:$HOME -v $HOME/.ssh:/root/.ssh \
+    --shm-size 64G \
+    --name training_env \
+    rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+```
+
+If you run Primus directly on the host instead of inside the prebuilt Docker image, install the JAX dependencies first by `pip install -r requirements-jax.txt`.
+
+### 2. Run pretraining
+
+Direct mode (inside the container)—pretraining Llama 3 8B on **MI355X**:
+
+```bash
+./runner/primus-cli direct \
+  -- train pretrain \
+  --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml
+```
+
+Container mode—passing the MaxText image with `--image`:
+
+```bash
+./runner/primus-cli container --image rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 \
+  -- train pretrain \
+  --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml
+```
+
+Slurm mode—supplying the image (and any environment variables) via a config file:
+
+```bash
+./runner/primus-cli --config my_maxtext_config.yaml slurm srun -N 8 \
+  -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama3_8B-pretrain.yaml
+```
+
+MaxText parallelism is set with `ici_*` (intra-node) and `dcn_*` (inter-node) fields—see the [MaxText config table](pretraining.md#maxtext-jax-pretraining) and [MaxText parameters](../03-configuration-reference/maxtext-parameters.md).
+
+---
+
+## Megatron Bridge (post-training)
+
+Megatron Bridge configurations are under `examples/megatron_bridge/configs/<ARCH>/` in the repository and are primarily **SFT and LoRA post-training** recipes (e.g. `qwen3_32b_sft_posttrain.yaml`, `llama31_70b_lora_posttrain.yaml`). Launch with `train posttrain`:
+
+```bash
+./runner/primus-cli direct \
+  --log_file /tmp/primus_qwen3_32b_sft.log \
+  -- train posttrain \
+  --config examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml
+```
+
+See [Post-training](posttraining.md) for the full SFT/LoRA workflow.
+
+---
+
+## Related documentation
+
+- [Pretraining](pretraining.md): backend concepts, configuration walkthroughs, parallelism vocabulary, full configuration inventories.
+- [Post-training](posttraining.md): SFT and LoRA via Megatron Bridge.
+- [CLI reference](cli-reference.md): `direct` / `container` / `slurm` modes and flags.
+- [Configuration system](configuration-system.md): YAML inheritance, overrides, image/env precedence.
+- [Performance tuning](../04-technical-guides/performance-tuning.md): HipBLASLt autotuning, Primus-Turbo, FP8, MoE.
diff --git a/docs/tuning_agent.md b/docs/02-user-guide/tuning-agent.md
similarity index 89%
rename from docs/tuning_agent.md
rename to docs/02-user-guide/tuning-agent.md
index e346de1d2..08baa04d2 100644
--- a/docs/tuning_agent.md
+++ b/docs/02-user-guide/tuning-agent.md
@@ -1,18 +1,18 @@
-# Tuning Agent
+# Tuning agent
 
 The **Tuning Agent** is an LLM-driven search for a near-optimal Primus
-**training configuration** — the full parallelism strategy *plus* the coupled
+**training configuration**—the full parallelism strategy *plus* the coupled
 batching, pipeline-schedule, memory, MoE-communication, and precision knobs —
 on a target GPU cluster, **without running the workload at scale**. It drives
 the [Primus Projection](./projection.md) tool as an evaluation oracle. Projection
-provides two estimates — **memory** and **performance** — each of which runs
+provides two estimates—**memory** and **performance**—each of which runs
 **benchmark-anchored by default** (measuring what fits on a sub-node run and
 scaling the rest analytically) with a fully analytical **no-GPU `simulate`**
 fallback. The agent returns the configuration that maximizes `tokens/s/GPU`
 subject to a per-GPU memory safety margin. See [Knobs Searched](#knobs-searched)
 for the full set of levers it tunes.
 
-- **Package**: [`primus/agents/tuning_agent/`](../primus/agents/tuning_agent/)
+- **Package**: [`primus/agents/tuning_agent/`](../../primus/agents/tuning_agent/)
 - **Entry point**: `python -m primus.agents.tuning_agent`
 
 This document is both the **user/operator guide** (installation, configuration,
@@ -22,36 +22,36 @@ features) for the agent.
 
 ---
 
-## Table of Contents
+## Table of contents
 
-1. [Why a Tuning Agent](#why-a-tuning-agent)
-2. [How It Works](#how-it-works)
-3. [Knobs Searched](#knobs-searched)
+1. [Why a tuning agent](#why-a-tuning-agent)
+2. [How it works](#how-it-works)
+3. [Knobs searched](#knobs-searched)
 4. [Installation](#installation)
-5. [LLM Setup](#llm-setup)
+5. [LLM setup](#llm-setup)
 6. [Quickstart](#quickstart)
-7. [Execution Modes](#execution-modes)
-8. [CLI Reference](#cli-reference)
-9. [Target-Cluster YAML](#target-cluster-yaml)
-10. [The Search Loop](#the-search-loop)
-11. [Evaluator and Projection Modes](#evaluator-and-projection-modes)
-12. [Output Artefacts](#output-artefacts)
-13. [Worked Example](#worked-example)
+7. [Execution modes](#execution-modes)
+8. [CLI reference](#cli-reference)
+9. [Target-cluster YAML](#target-cluster-yaml)
+10. [The search loop](#the-search-loop)
+11. [Evaluator and projection modes](#evaluator-and-projection-modes)
+12. [Output artefacts](#output-artefacts)
+13. [Worked example](#worked-example)
 14. [Troubleshooting](#troubleshooting)
 15. [Limitations](#limitations)
-16. [Design Notes & Future Features](#design-notes--future-features)
+16. [Design notes and future features](#design-notes-and-future-features)
 
 ---
 
-## Why a Tuning Agent
+## Why a tuning agent
 
 Choosing a training configuration for a large training (or inference) workload
 is a combinatorial problem. The configuration is the joint choice of the
 parallelism dimensions:
 
-- **Data Parallel (DP)** — derived from world size and the other axes,
+- **Data Parallel (DP)**—derived from world size and the other axes,
 - **Tensor Parallel (TP)**,
-- **Expert Parallel (EP)** — for MoE models,
+- **Expert Parallel (EP)**—for MoE models,
 - **Context Parallel (CP)**,
 - **Pipeline Parallel (PP)**, with virtual pipeline (**VPP**) and the
   **pipeline schedule** (1F1B / interleaved / zero-bubble / ZBV-\* /
@@ -62,7 +62,7 @@ dimensions translate into in-flight work and memory pressure: global batch
 size (**GBS**), micro batch size (**MBS**), activation recomputation
 (`recompute_granularity`, `recompute_num_layers`), the overlap flags
 (`overlap_grad_reduce`, `overlap_param_gather`), and a set of higher-impact
-levers — FP8 precision, MoE DeepEP / sync-free communication, fused
+levers—FP8 precision, MoE DeepEP / sync-free communication, fused
 cross-entropy, and optimizer-state sharding (distributed optimizer / FSDP2).
 The full set is enumerated in [Knobs Searched](#knobs-searched).
 
@@ -77,7 +77,7 @@ legal configuration within a user-specified budget.
 
 ---
 
-## How It Works
+## How it works
 
 ```
 ┌────────────────────────────────────────────────────────────────────────┐
@@ -108,7 +108,7 @@ recompute for MBS, whether CP helps a given MoE shape).
 
 ---
 
-## Knobs Searched
+## Knobs searched
 
 The agent sweeps far more than the five parallelism dimensions. Its trial
 configuration (`TrialConfig` in `legality.py`) carries the full set of knobs
@@ -117,7 +117,7 @@ agent leaves it unset) or **overridden for a trial** and translated into the
 corresponding `projection` flags by the evaluator. Each is legality-checked in
 code before it ever reaches the projection tool.
 
-### Parallelism & batching
+### Parallelism and batching
 
 | Knob | Legal values | What it controls |
 |------|--------------|------------------|
@@ -143,11 +143,11 @@ code before it ever reaches the projection tool.
 |------|--------------|------------------|
 | `recompute_granularity` | `none` / `selective` / `full` | Activation recomputation strategy |
 | `recompute_num_layers` | int (≤ layers per VPP stage) | Layers recomputed per stage under `full` |
-| `cross_entropy_loss_fusion` | `true` / `false` / inherit | Fused cross-entropy — large-vocab memory + compute win |
+| `cross_entropy_loss_fusion` | `true` / `false` / inherit | Fused cross-entropy—large-vocab memory + compute win |
 | `use_distributed_optimizer` | `true` / `false` / inherit | ZeRO-1 optimizer-state sharding across DP |
 | `use_torch_fsdp2` | `true` / `false` / inherit | FSDP2 sharding (mutually exclusive with `use_distributed_optimizer`) |
 
-### MoE communication — *MoE only, high impact*
+### MoE communication—*MoE only, high impact*
 
 | Knob | Legal values | What it controls |
 |------|--------------|------------------|
@@ -155,11 +155,11 @@ code before it ever reaches the projection tool.
 | `sync_free_stage` | `0` / `1` / `2` / `3` | Sync-free MoE pipelining; stage ≥ 2 auto-enables DeepEP |
 | `target_ep_size` | positive int / inherit | EP override used for All-to-All modeling |
 
-### Precision — *high impact*
+### Precision—*high impact*
 
 | Knob | Legal values | What it controls |
 |------|--------------|------------------|
-| `fp8` | `none` / `hybrid` (also `e4m3`, `delayed`) | FP8 on linear layers — roughly 2× compute on GEMMs |
+| `fp8` | `none` / `hybrid` (also `e4m3`, `delayed`) | FP8 on linear layers—roughly 2× compute on GEMMs |
 
 ### Coupling rules enforced in code
 
@@ -206,10 +206,10 @@ Origami) are only needed by the evaluator paths you actually use:
 
 ---
 
-## LLM Setup
+## LLM setup
 
 The agent uses [DSPy](https://dspy.ai), which routes LLM calls through
-[LiteLLM](https://docs.litellm.ai/docs/providers) internally — **no separate
+[LiteLLM](https://docs.litellm.ai/docs/providers) internally—**no separate
 proxy process is required**. Set credentials for whichever provider you use:
 
 ```bash
@@ -266,7 +266,7 @@ python -m primus.agents.tuning_agent \
 
 ---
 
-## Execution Modes
+## Execution modes
 
 The agent has two orthogonal mode switches: **what the evaluator does**
 (`--mode`) and **whether the LLM stage runs** (`--seed-only` / `--agent-only`).
@@ -293,7 +293,7 @@ The agent has two orthogonal mode switches: **what the evaluator does**
 
 ---
 
-## CLI Reference
+## CLI reference
 
 ```bash
 python -m primus.agents.tuning_agent \
@@ -322,15 +322,15 @@ python -m primus.agents.tuning_agent \
 
 ---
 
-## Target-Cluster YAML
+## Target-cluster YAML
 
 A thin wrapper around the existing Primus `hardware_config` convention, so no
 new networking format has to be invented; topology, bandwidths, and latencies
 are consumed by the analytical communication model (see
-[`projection.md` → Communication Modeling](./projection.md#communication-modeling)).
+[`projection.md` → Assumptions (performance projection)](./projection.md#assumptions-performance-projection)).
 
 A complete example ships at
-[`examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml`](../examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml):
+[`examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml`](../../examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml):
 
 ```yaml
 target_cluster:
@@ -396,7 +396,7 @@ agent:
 
 ---
 
-## The Search Loop
+## The search loop
 
 1. **Resolve the workload.** Load the workload YAML, follow
    `modules.pre_trainer.model` into `primus/configs/models/megatron/<model>.yaml`,
@@ -437,7 +437,7 @@ agent:
 
 ---
 
-## Evaluator and Projection Modes
+## Evaluator and projection modes
 
 The evaluator wraps the Primus Projection CLI behind a uniform interface, so
 the agent does not need to know which mode produced a number:
@@ -463,21 +463,21 @@ pre-filter to reject infeasible configs before paying for a performance call,
 then `simulate` (or `benchmark` for promising candidates when a GPU is
 available). The tool belt exposed to the LLM mirrors this:
 
-- `evaluate_memory_only(config_json)` — cheap pre-filter
-- `evaluate_simulate(config_json)` — primary scoring path
-- `evaluate_with_benchmark(config_json)` — only if `has_gpu: true`
+- `evaluate_memory_only(config_json)`—cheap pre-filter
+- `evaluate_simulate(config_json)`—primary scoring path
+- `evaluate_with_benchmark(config_json)`—only if `has_gpu: true`
 - `get_history`, `get_best`, `get_legal_axes`, `get_architecture`,
   `get_cluster`, `get_budget_status`
 - `note_to_scratchpad`, `read_scratchpad`
-- `query_llm(prompt, system?)` — one-shot "LLM-inside-LLM" consultation
+- `query_llm(prompt, system?)`—one-shot "LLM-inside-LLM" consultation
 
-For the underlying projection math — memory components, the simulate vs.
+For the underlying projection math—memory components, the simulate vs.
 benchmark trade-off, and the benchmark-based memory projection the agent uses
-for OOM-accurate feasibility — see [`projection.md`](./projection.md).
+for OOM-accurate feasibility—see [`projection.md`](./projection.md).
 
 ---
 
-## Output Artefacts
+## Output artefacts
 
 Everything lands in `--out-dir`:
 
@@ -509,7 +509,7 @@ The run also prints the best configuration and ready-to-paste exports:
 
 ---
 
-## Worked Example
+## Worked example
 
 Search for the best Mixtral 8×22B configuration on a 4-node MI355X pod, with a
 single idle 8-GPU node available for benchmarking:
@@ -566,17 +566,17 @@ These are honest caveats, not future features:
    `memory_safety_margin` compensates conservatively; the benchmark-based
    memory path (see `projection.md`) closes most of this gap.
 3. **Search-space explosion** with all axes on. The agent mitigates with an
-   impact-ordered deterministic seed plan — high-leverage levers first
+   impact-ordered deterministic seed plan—high-leverage levers first
    (recompute, MoE DeepEP / sync-free, FP8, then schedule and VPP/MBS
    neighbors), with the broad TP×PP×EP×CP grid evaluated last and capped by
-   `--seed-budget` — so the LLM starts from an informed incumbent and spends
+   `--seed-budget`—so the LLM starts from an informed incumbent and spends
    its budget on polish.
 4. **Cluster-description lossiness**: averaged bandwidth/latency cannot capture
    contention or per-rail asymmetry.
 
 ---
 
-## Design Notes & Future Features
+## Design notes and future features
 
 > This section captures the design rationale and the paper-ready problem
 > statement, plus the list of features deliberately deferred from v1.
@@ -584,8 +584,8 @@ These are honest caveats, not future features:
 ### Problem statement (paper-ready)
 
 We address the problem of **automatically selecting a near-optimal
-training configuration — the parallelism strategy plus the coupled batching,
-schedule, memory, MoE-communication, and precision knobs — for a large-scale
+training configuration—the parallelism strategy plus the coupled batching,
+schedule, memory, MoE-communication, and precision knobs—for a large-scale
 distributed training workload on a target GPU cluster, without executing the
 workload at scale**. The
 configuration space is combinatorial: for each axis only a small set of values
@@ -608,7 +608,7 @@ We formulate the search as **LLM-as-policy over a hybrid analytical /
 benchmark-driven evaluator**. The evaluator is the Primus Projection tool,
 which provides three calls of increasing fidelity and cost: (i) a **memory
 projection** that runs either analytically (no GPU) or, by default,
-benchmark-anchored — measuring the real per-rank peak on a sub-node run and
+benchmark-anchored—measuring the real per-rank peak on a sub-node run and
 extrapolating it for an OOM-accurate estimate; (ii) a fully analytical
 **performance projection** built on the Origami GEMM model and an SDPA
 simulator; and (iii) a **hybrid benchmark** that measures per-layer compute on
@@ -622,7 +622,7 @@ to a configurable per-GPU memory safety margin.
 The contribution is a configuration-search methodology that exploits an LLM's
 ability to reason over architectural priors (topk dominance of MoE activations,
 the impact of MQA on attention activation, the inter-node/intra-node boundary
-for All-to-All) to direct an analytical oracle — replacing exhaustive sweeps on
+for All-to-All) to direct an analytical oracle—replacing exhaustive sweeps on
 real hardware with a small number of informed analytical evaluations. The
 system runs either entirely on a CPU-only host (simulate backend only) or in a
 mixed mode where a small number of real-hardware benchmark runs calibrate the
@@ -640,7 +640,7 @@ analytical predictions.
 4. Agent search over **all** parallelism and coupled axes (TP, PP, EP, CP, MBS,
    GBS, VPP, pipeline schedule, recompute, overlap flags) plus the higher-impact
    levers (FP8, MoE DeepEP / sync-free, fused cross-entropy,
-   distributed-optimizer / FSDP2) — restricted by per-architecture legality.
+   distributed-optimizer / FSDP2)—restricted by per-architecture legality.
    See [Knobs Searched](#knobs-searched) for the complete list.
 5. Two evaluator paths: a **no-GPU path** (`projection memory --memory-mode
    simulate` + `projection performance --profiling-mode simulate`), always
@@ -657,60 +657,60 @@ analytical predictions.
 These are recorded so they are not lost; they are deliberately left out to keep
 the agent small.
 
-- **F1. Multi-objective Pareto** — replace the scalar `tokens/s/GPU` objective
+- **F1. Multi-objective Pareto**—replace the scalar `tokens/s/GPU` objective
   with a Pareto frontier over (throughput, MFU, memory headroom, projected
   $/token, projected energy/token).
-- **F2. Online cluster-spec retrieval** — pull the cluster description from an
+- **F2. Online cluster-spec retrieval**—pull the cluster description from an
   internal registry or known-archs catalog (MI300X / MI325X / MI355X reference
   pods) and fall back to user overrides; optionally infer topology from a small
   DCGM / ROCm-SMI dump.
-- **F3. Persistent memory / configuration cache** — cache
+- **F3. Persistent memory / configuration cache**—cache
   `(model_signature, cluster_signature) → best_known_configs` across runs;
   invalidate when ROCm / hipBLASLt / framework versions change.
-- **F4. Agent-proposed scale-downs and microbenchmarks** — let the agent
+- **F4. Agent-proposed scale-downs and microbenchmarks**—let the agent
   *propose* reduced-model proxies and targeted microbenchmarks to reduce the
   uncertainty of its current top-k (reusing the `moe_proxy_single_node.yaml`
   pattern).
-- **F5. Telemetry plug-ins (rocprofiler / TraceLens / Magpie)** — after a
+- **F5. Telemetry plug-ins (rocprofiler / TraceLens / Magpie)**—after a
   benchmark run, optionally extract per-kernel time, GEMM efficiency, A2A bytes,
   NIC utilization to calibrate the analytical models and explain
   underperformance back to the agent. Exposed as a `SKILL.md`-described plug-in.
   References: [TraceLens](https://github.com/AMD-AGI/TraceLens-internal),
   [Magpie](https://github.com/AMD-AGI/Magpie).
-- **F6. Calibration learning** — under `--profiling-mode both`, record
+- **F6. Calibration learning**—under `--profiling-mode both`, record
   per-(model, arch, dim) residuals between simulate and benchmark, fit a small
   correction model, and report a confidence band on subsequent simulate runs.
-- **F7. Robustness / sensitivity report** — for the winning config, sweep ±1
+- **F7. Robustness / sensitivity report**—for the winning config, sweep ±1
   step on each axis and report whether the optimum is a sharp peak or a broad
   basin (cheap; only `simulate` calls).
-- **F8. Cross-axis priors as DSPy modules** — a library of tunable "rules of
+- **F8. Cross-axis priors as DSPy modules**—a library of tunable "rules of
   thumb" that DSPy's optimizer can refine over time using the trial logs.
-- **F9. Sub-agent "test-proposer"** — delegate targeted experiments (e.g. a
+- **F9. Sub-agent "test-proposer"**—delegate targeted experiments (e.g. a
   2-layer scale-down forward+backward microbenchmark, or a stand-alone A2A probe
   at the proposed EP × hidden_size × topk), profiling the *test* with explicit
   synchronisation rather than the sandbox. A `run_proposed_experiment(plan,
   code)` tool can be added to `tools.py` without restructuring the loop.
-- **F10. Sub-LLM expert router** — extend the existing `query_llm` tool into a
+- **F10. Sub-LLM expert router**—extend the existing `query_llm` tool into a
   *named expert* router (`query_llm(expert='moe', …)`).
 
 ### Known design holes
 
-1. **Simulator-vs-reality gap** — see Limitations above; in no-GPU mode the
+1. **Simulator-vs-reality gap**—see Limitations above; in no-GPU mode the
    agent reports a confidence caveat.
-2. **Memory-projection blind spots** — A2A buffers, allocator fragmentation,
+2. **Memory-projection blind spots**—A2A buffers, allocator fragmentation,
    and comm scratch are not modeled analytically; the benchmark-based memory
    projection closes most of this gap by anchoring on a measured peak.
-3. **Search-space explosion** — mitigated by an impact-ordered deterministic
+3. **Search-space explosion**—mitigated by an impact-ordered deterministic
    seed plan (high-leverage levers first, broad parallelism grid last) plus
    LLM-guided search from the seed incumbent.
-4. **Cluster-description lossiness** — averaged bandwidth/latency cannot capture
+4. **Cluster-description lossiness**—averaged bandwidth/latency cannot capture
    contention or per-rail asymmetry.
 
 ---
 
-## Related Documentation
+## Related documentation
 
-- [Projection](./projection.md) — memory + performance projection internals,
+- [Projection](./projection.md)—memory + performance projection internals,
   including the benchmark-based memory projection the agent relies on.
-- [Tuning Agent package README](../primus/agents/tuning_agent/README.md) —
+- [Tuning Agent package README](../../primus/agents/tuning_agent/README.md) —
   quickstart reference inside the source tree.
diff --git a/docs/03-configuration-reference/README.md b/docs/03-configuration-reference/README.md
new file mode 100644
index 000000000..487d6e526
--- /dev/null
+++ b/docs/03-configuration-reference/README.md
@@ -0,0 +1,13 @@
+# Configuration reference
+
+Parameter references for Primus presets, backend-facing keys, and commonly used environment variables.
+
+- [Megatron parameters](megatron-parameters.md): Megatron-LM backend YAML parameters and Primus overrides
+- [TorchTitan parameters](torchtitan-parameters.md): Primus TorchTitan preset keys and common JobConfig fields
+- [MaxText parameters](maxtext-parameters.md): Primus MaxText overlay defaults and common fields
+- [Megatron Bridge parameters](megatron-bridge-parameters.md): Megatron Bridge recipe, SFT, and pretraining fields surfaced through Primus
+- [Environment variables](environment-variables.md): practical reference for commonly encountered environment variables
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/03-configuration-reference/environment-variables.md b/docs/03-configuration-reference/environment-variables.md
new file mode 100644
index 000000000..61aca7f1b
--- /dev/null
+++ b/docs/03-configuration-reference/environment-variables.md
@@ -0,0 +1,239 @@
+# Environment variables reference
+
+This document catalogs the main environment variables you might encounter when running Primus on AMD GPUs: distributed launchers, Primus runners and CLI, YAML substitution, libraries (NCCL/RCCL, ROCm, PyTorch, JAX), and optional integrations (Hugging Face, WandB, MLflow). It is a practical reference, not a complete list of every variable accepted by upstream libraries.
+
+**Legend**
+
+- **Required**: Must be set for the stated workflow; otherwise the job fails or mis-ranks.
+- **Optional**: Has a safe default or is only needed for specific features.
+- **Set by**: Typical source (launcher, `runner/helpers/envs/*.sh`, user shell, container host).
+- **Used in**: Representative Primus paths; many variables are also read by NVIDIA NCCL, AMD RCCL, PyTorch, or JAX without Primus wrapping them.
+
+---
+
+## 1. PyTorch distributed
+
+Set by `torchrun`, Slurm launchers, or `runner/primus-cli-direct.sh` / `runner/primus-cli-slurm-entry.sh`. Consumed by PyTorch distributed, RCCL, and Primus helpers.
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `MASTER_ADDR` | `localhost` (direct / `base_env.sh`) | User, Slurm entry (`primus-cli-slurm-entry.sh`), or validation fallback (`runner/lib/validation.sh`) | `primus/pretrain.py`, `primus/core/base_module.py`, `primus/core/utils/env.py`, `primus/tools/preflight/network/network_probe.py`, PyTorch rendezvous | Rendezvous hostname or IP for process group initialization. **Required** for multi-node if not using Slurm auto-detection. |
+| `MASTER_PORT` | `1234` (direct), `29500` in some Python defaults | Config / CLI / user | Same as `MASTER_ADDR`; `validation.sh` enforces 1024–65535 | TCP port for the store backing `torch.distributed`. |
+| `RANK` | `0` if unset in helpers | `torchrun` | `primus/tools/utils.py`, `primus/tools/preflight/global_vars.py`, projection and profiler code | Global rank index. |
+| `WORLD_SIZE` | `1` | `torchrun` | Preflight, projection, `primus/core/base_module.py` | Total number of processes. |
+| `LOCAL_RANK` | `0` | `torchrun` | `primus/core/base_module.py`, GPU selection in benchmarks and trainers | GPU index on this node. |
+| `LOCAL_WORLD_SIZE` | `1` (Python) / `8` in benchmarks default | `torchrun` | `primus/tools/preflight/*.py`, `strided_allgather_bench.py` | Processes (GPUs) per node. |
+| `NODE_RANK` | `0` | `primus-cli-direct` / `primus-cli-slurm-entry.sh` | `primus/pretrain.py`, logging in `runner/lib/common.sh` | Zero-based node index in multi-node jobs. |
+| `NNODES` | `1` | Direct config (`runner/.primus.yaml`), `primus-cli-slurm-entry.sh` | `primus/pretrain.py`, `primus/core/projection/training_config.py` | Number of nodes in the job. |
+| `GPUS_PER_NODE` | `8` | `runner/.primus.yaml` direct section, `primus-cli-slurm-entry.sh`, `validation.sh` | `primus/core/projection/module_profilers/*.py`, training config helpers | GPUs per node for world-size math and binding. |
+
+---
+
+## 2. Primus core
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `PRIMUS_PATCHES` | `""` / `"all"` | User | `primus/core/patches/patch_runner.py` | `"all"` or empty enables all patches; `"none"` disables; comma list enables subset. |
+| `PRIMUS_LOG_LEVEL` | `INFO` | User; debug paths in `runner/primus-cli-*.sh` set `DEBUG` | `runner/lib/common.sh` | Log verbosity: `DEBUG`, `INFO`, `WARN`, `ERROR`. |
+| `PRIMUS_LOG_TIMESTAMP` | `1` | User | `runner/lib/common.sh` | `1` prefixes logs with timestamps; `0` disables. |
+| `PRIMUS_LOG_COLOR` | `1` (auto-off if not a TTY) | User; tests may set `0` | `runner/lib/common.sh` | ANSI colors in runner logs. |
+| `PRIMUS_DEBUG` | `0` | User | `runner/helpers/envs/primus-env.sh` | `1` enables `set -x` in the env loader for shell tracing. |
+| `PRIMUS_SKIP_VALIDATION` | `0` | User / tests | `runner/helpers/envs/primus-env.sh` | `1` skips `validate_distributed_params` (not recommended). |
+| `PRIMUS_EXPECT_IB` | (unset) | User | `primus/tools/preflight/network/network_standard.py` | When `1`, preflight treats InfiniBand as expected for validation. |
+| `PRIMUS_CLUSTER` | `amd-aig-poolside` (CLI default) | User | `primus/tools/benchmark/rccl_bench_args.py` | Cluster label for RCCL benchmark tooling. |
+| `PRIMUS_GPU_ARCH` | (auto / `"mi300x"` in simulators) | User / CLI | `primus/core/projection/simulation_backends/origami_backend.py`, `sdpa_simulator.py`, `projection.py` CLI | GPU architecture string for performance projection. |
+| `PRIMUS_GPU_CLOCK_MHZ` | (unset) | User | Same as `PRIMUS_GPU_ARCH` | Optional clock override for projection. |
+| `PRIMUS_GPU_DEVICE` | `0` | User | `origami_backend.py` | GPU index for hardware detection in projection. |
+| `PRIMUS_GEMM_BACKEND` | (unset) | User | `primus/core/projection/simulation_backends/factory.py` | Selects GEMM simulation backend by name. |
+| `PRIMUS_PREFLIGHT_MIN_FREE_MEM_GB` | `1` | User | `primus/tools/preflight/gpu/utils.py` | Minimum free GPU memory (GB) for preflight checks. |
+| `PRIMUS_PREFLIGHT_MIN_TFLOPS` | `10.0` | User | `primus/tools/preflight/gpu/utils.py` | Minimum TFLOPS threshold for preflight GEMM checks. |
+| `PRIMUS_TURBO_AUTO_TUNE` | (unset) | User / tests | `tests/trainer/test_megatron_trainer.py` (integration) | Enables Turbo auto-tuning in supported Turbo/Megatron test flows; not referenced in core `primus/` Python outside tests. **Optional**. |
+| `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND` | `TURBO` | User; hooks may set `DEEP_EP` | `primus/backends/megatron/patches/args/rocm_arg_validation.py`, `examples/run_pretrain.sh`, `runner/helpers/hooks/05_using_uep.sh` | MoE dispatch/combine backend selector. |
+
+---
+
+## 3. Primus YAML substitution
+
+Parsed by `primus/core/config/yaml_loader.py` for patterns `${VAR}` (required) and `${VAR:default}` (optional). Typical experiment YAMLs under `examples/` use these for sweep-friendly overrides.
+
+| Variable | Typical default in YAML | Where set | Where used | Description |
+|----------|-------------------------|-----------|------------|-------------|
+| `PRIMUS_TEAM` | `"amd"` | User | Resolved before module merge in experiment YAML | Work group / team segment in paths. |
+| `PRIMUS_USER` | `"root"` | User | Experiment YAML | User name segment. |
+| `PRIMUS_EXP_NAME` | per-example | User | Experiment YAML | Experiment folder name. |
+| `PRIMUS_WORKSPACE` | `"./output"` | User | Experiment YAML | Root workspace for artifacts. |
+| `PRIMUS_TP` | `1` | User | Megatron example YAMLs | `tensor_model_parallel_size` override. |
+| `PRIMUS_PP` | `1` | User | Megatron example YAMLs | `pipeline_model_parallel_size` override. |
+| `PRIMUS_EP` | `1` | User | Megatron example YAMLs | `expert_model_parallel_size` override. |
+| `PRIMUS_SEQ_LENGTH` | per-model | User | Megatron example YAMLs | Sequence length override. |
+| `PRIMUS_MAX_POSITION_EMBEDDINGS` | `4096` or `131072` | User | `examples/megatron/**/*.yaml`, tests | Position embedding cap override. |
+| `PRIMUS_GLOBAL_BATCH_SIZE` | per-model | User | Megatron example YAMLs | Global batch override. |
+| `PRIMUS_NUM_LAYERS` | per-model | User | Tests and MoE examples | Transformer layer count override. |
+| `PRIMUS_MOE_LAYER_FREQ` | MoE patterns | User | MoE examples / tests | MoE layer frequency pattern. |
+| `PRIMUS_TOKENIZED_DATA_PATH` | `null` | User | Megatron examples | Path to tokenized training data. |
+| `PRIMUS_MODEL` | per-stack | User | Megatron examples | Model preset stem (e.g. `llama3_8B`). |
+| `PRIMUS_VPP` | `null` | User | `tests/trainer/test_megatron_trainer.yaml` | Virtual pipeline stages override. |
+
+---
+
+## 4. NCCL / RCCL
+
+Primus seeds many of these in `runner/helpers/envs/base_env.sh`. RCCL honors NCCL-compatible variables on AMD GPUs. See [NCCL environment](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) and [RCCL environment](https://rocm.docs.amd.com/projects/rccl/en/develop/api-reference/env-variables.html).
+
+| Variable | Default (Primus base) | Where set | Where used | Description |
+|----------|------------------------|-----------|------------|-------------|
+| `NCCL_DEBUG` | unset | User / `base_env.sh` empty default | Preflight reports, RCCL runtime | Log verbosity: `NONE`, `WARN`, `INFO`, `TRACE`, etc. **Optional** unless debugging comms. |
+| `NCCL_SOCKET_IFNAME` | derived from `IP_INTERFACE` | `base_env.sh` | `primus/tools/preflight/network/*.py`, GPU topology helpers | Socket NIC for host networking. |
+| `GLOO_SOCKET_IFNAME` | same as NCCL if unset | `base_env.sh` | Preflight | Gloo TCP backend interface. |
+| `NCCL_IB_HCA` | auto via `get_nccl_ib_hca.sh` if empty | `base_env.sh`, container passthrough | Preflight, multi-node tuning | InfiniBand HCAs to use. |
+| `NCCL_IB_GID_INDEX` | `3` | `base_env.sh` | RCCL | GID index for IB/RoCE; many sites use `1` for RoCE v2 (override as needed). |
+| `NCCL_IB_TC` | (unset) | User | RCCL | InfiniBand traffic class. |
+| `NCCL_IB_FIFO_TC` | (unset) | User | RCCL | InfiniBand FIFO traffic class. |
+| `NCCL_IB_ROCE_VERSION_NUM` | (unset) | User | RCCL | RoCE version selection. |
+| `NCCL_PXN_DISABLE` | `1` | `base_env.sh` | RCCL | Disable PXN (PCIe cross-NIC); set `0` to enable. |
+| `NCCL_P2P_NET_CHUNKSIZE` | `524288` | `base_env.sh` | RCCL | P2P network chunk size tuning. |
+| `NCCL_PROTO` | (unset) | User | RCCL | Protocol selection (e.g. `Simple`, `LL`, `LL128`). |
+| `NCCL_CROSS_NIC` | `0` | `base_env.sh` | RCCL | Cross-NIC communication policy. |
+| `NCCL_IB_RETRY_CNT` | (unset) | User | RCCL | IB retry count. |
+| `NCCL_IB_TIMEOUT` | (unset) | User | RCCL | IB timeout. |
+| `NCCL_NET_GDR_LEVEL` | (unset) | User | Preflight summaries | GPUDirect RDMA level. |
+| `NCCL_IB_DISABLE` | `0` | User / env | Preflight | Disable IB; use sockets only. |
+| `NCCL_DMABUF_ENABLE` | (unset) | User | RCCL | DMA-BUF registration path. |
+| `NCCL_IGNORE_CPU_AFFINITY` | (unset) | User | RCCL | Ignore CPU affinity hints. |
+| `NCCL_IB_QPS_PER_CONNECTION` | (unset) | User | RCCL | IB QPs per connection. |
+| `NCCL_MAX_P2P_CHANNELS` | (unset) | User | RCCL | Cap P2P channels. |
+| `NCCL_GDR_FLUSH_DISABLE` | (unset) | User | RCCL | Disable GDR flush. |
+| `NCCL_IB_USE_INLINE` | (unset) | User | RCCL | Inline IB sends. |
+| `NCCL_NET_PLUGIN` | (unset) | User | RCCL | Alternate network plugin (e.g. `librccl-anp.so`). |
+| `RCCL_MSCCL_ENABLE` | `0` | `base_env.sh` | RCCL | Enable MSCCL algorithms. |
+| `RCCL_MSCCLPP_THRESHOLD` | `1GiB` default | `base_env.sh` | RCCL | MSCCL++ message-size threshold. |
+| `RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING` | `0` in hooks | `runner/helpers/hooks/03_enable_ainic.sh` | RCCL | Stricter GDR flush memory ordering; relevant for some NIC/GPU combos. |
+| `TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK` | `0` | `base_env.sh` | PyTorch + RCCL | Tensor allocator hook for NCCL registration. |
+| `TORCH_NCCL_HIGH_PRIORITY` | `1` | `base_env.sh` | PyTorch | High-priority NCCL streams. |
+
+---
+
+## 5. ROCm / HSA / HIP
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `HSA_ENABLE_SDMA` | `1` | `base_env.sh` | ROCm runtime | Enable SDMA engines for copies. |
+| `HSA_NO_SCRATCH_RECLAIM` | `1` | `base_env.sh`, container passthrough | ROCm runtime; documented for MoE stability | `1` keeps scratch allocated (often used for MoE stability). See [ROCR environment](https://rocm.docs.amd.com/projects/ROCR-Runtime/en/docs-7.1.1/environment_variables.html). |
+| `HIP_VISIBLE_DEVICES` | `0..GPUS_PER_NODE-1` | `base_env.sh` | ROCm device visibility | Restricts which GPU indices ROCm exposes. |
+| `ROCBLAS_DEFAULT_ATOMICS_MODE` | (unset) | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Read for deterministic / accuracy-sensitive GEMM behavior. |
+
+---
+
+## 6. CUDA / PyTorch
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `CUDA_DEVICE_MAX_CONNECTIONS` | `1` | `base_env.sh`; Megatron patches may adjust | `primus/backends/megatron/patches/env_patches.py`, Megatron patches | Limits concurrent CUDA connections; often `1` for TP/PP overlap. |
+| `TORCH_COMPILE_DISABLE` | `0` | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Disable `torch.compile` when `1`. |
+
+---
+
+## 7. Transformer engine
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `NVTE_ROCM_ENABLE_MXFP8` | `1` | `base_env.sh` | Transformer Engine on ROCm | Enable MXFP8 paths. |
+| `NVTE_CK_USES_BWD_V3` | `1` | `base_env.sh`, container passthrough | TE / CK | Use CK backward v3 kernels. |
+| `NVTE_CK_IS_V3_ATOMIC_FP32` | (unset; examples print `0`) | User / `examples/run_pretrain.sh`, container passthrough | TE / CK | Atomic FP32 mode for CK v3 backward. |
+| `PATCH_TE_FLASH_ATTN` | `0` | `base_env.sh`, container passthrough | `runner/helpers/hooks/01_patch_te_flash_attn_max_version.sh` | Trigger TE flash-attn patch hook when `1`. |
+
+---
+
+## 8. Caches and authentication
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `HF_HOME` | `${DATA_PATH}/huggingface` | `base_env.sh`, `primus/core/utils/env_setup.py`, `primus/pretrain.py` | Hugging Face libraries | Cache for models and datasets. |
+| `HF_TOKEN` | (unset) | User, container passthrough | Hugging Face Hub | Auth for gated models. **Required** for private/gated assets. |
+| `TORCH_HOME` | under workspace | `primus/core/utils/env_setup.py` | PyTorch Hub | Torch Hub cache root. |
+| `TRANSFORMERS_CACHE` | aligned with HF layout | `primus/core/utils/env_setup.py` | `transformers` | Model cache for Transformers. |
+| `WANDB_API_KEY` | (unset) | User, container passthrough | Weights & Biases client, Megatron trainer checks | API key for logging. **Required** for Weights & Biases when enabled. |
+| `WANDB_PROJECT` | (unset) | User / TorchTitan patch | `primus/backends/torchtitan/patches/wandb_patches.py` | Project name. |
+| `WANDB_RUN_NAME` | (unset) | User / patches | Same | Run display name. |
+| `WANDB_TEAM` | (unset) | User | TorchTitan metrics (entity) | WandB team/entity. |
+| `DATABRICKS_HOST` | (unset) | User | `mlflow` client (via `primus/backends/megatron/training/global_vars.py` MLflow setup) | Required for Databricks-hosted MLflow when MLflow logging is enabled. |
+| `DATABRICKS_TOKEN` | (unset) | User | Databricks APIs | Auth token paired with host. |
+| `MLFLOW_TRACKING_URI` | (unset) | User | `mlflow` (via Megatron integrations) | MLflow tracking server URI. **Optional** unless using MLflow. |
+| `MLFLOW_REGISTRY_URI` | (unset) | User | MLflow | Model registry endpoint. |
+| `NLTK_DATA` | (unset) | User | `runner/helpers/hooks/train/pretrain/megatron/preprocess_data.py`, Megatron-LM tools | Punkt and other tokenizer data for preprocessing. |
+| `TOKENIZED_DATA_PATH` | per-hook default | User | `runner/helpers/hooks/train/pretrain/megatron/prepare.py` | Pre-tokenized dataset location for Megatron data prep hooks. |
+
+---
+
+## 9. hipBLASLt tuning
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `PRIMUS_HIPBLASLT_TUNING` | `0` | User | `examples/run_pretrain.sh` | **Master switch** for the HipBLASLt tuning flow (`1` enables). Must be set before `PRIMUS_HIPBLASLT_TUNING_STAGE` takes effect, and is mutually exclusive with deterministic mode (`PRIMUS_DETERMINISTIC=1`). |
+| `PRIMUS_HIPBLASLT_TUNING_STAGE` | `0` | User | `examples/run_pretrain.sh` | Stages `0` off, `1` dump shapes, `2` offline tune, `3` apply tuned kernels. |
+| `HIPBLASLT_TUNING_OVERRIDE_FILE` | (unset) | User / tuning scripts | `examples/run_pretrain.sh` | Path to tuned-kernel override file for stage `3`. |
+| `TE_HIPBLASLT_TUNING_RUN_COUNT` | varies | User | `examples/run_pretrain.sh` | Number of benchmark runs per shape during TE hipBLASLt tuning. |
+| `TE_HIPBLASLT_TUNING_ALGO_COUNT` | varies | User | `examples/run_pretrain.sh` | Transformer Engine hipBLASLt search breadth. |
+| `TE_HIPBLASLT_TUNING_ALGO_FILE` | (unset) | User | TE + HipBLASLt | Algorithm file for TE tuning flows. |
+| `TE_HIPBLASLT_TUNING` | (unset) | User | `examples/run_pretrain.sh` | When set, interacts with deterministic mode and tuning stages (disable conflicting modes per script comments). |
+| `HIPBLASLT_LOG_LEVEL` | (unset) | User | HipBLASLt | Library log level. |
+| `HIPBLASLT_LOG_MASK` | (unset) | User | HipBLASLt | Bitmask for log categories. |
+
+---
+
+## 10. Build and rebuild
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `REBUILD_PRIMUS_TURBO` | `0` | User, container passthrough | `runner/helpers/hooks/00_rebuild_primus_turbo.sh` | `1` rebuilds Primus-Turbo on startup. |
+| `REBUILD_BNXT` | `0` | User, container passthrough | `runner/helpers/hooks/02_rebuild_bnxt.sh` | `1` rebuilds BNXT driver artifacts when packaged. |
+| `USING_AINIC` | (unset) | User | `runner/helpers/hooks/03_enable_ainic.sh` | `1` enables AINIC-oriented networking hooks. |
+| `MAX_JOBS` | (unset) | User / tooling | `tools/daily/safe_wrapper.py` | Parallel compile jobs for pip builds. |
+| `BACKEND_PATH` | (unset) | User | `primus/pretrain.py`, `primus/core/backend/backend_adapter.py` | Override checkout path for third-party backends (Megatron, TorchTitan, MaxText). |
+
+---
+
+## 11. Container passthrough
+
+`runner/.primus.yaml` lists names forwarded from the host into training containers (`container.options.env`). Primus does not assign values here; it only allowlists keys for `--env` forwarding.
+
+Forwarded keys:
+
+`MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `DOCKER_IMAGE`, `HF_TOKEN`, `WANDB_API_KEY`, `ENABLE_NUMA_BINDING`, `REBUILD_PRIMUS_TURBO`, `USING_AINIC`, `PATCH_TE_FLASH_ATTN`, `REBUILD_BNXT`, `HSA_NO_SCRATCH_RECLAIM`, `NVTE_CK_USES_BWD_V3`, `GPU_MAX_HW_QUEUES`, `HSA_KERNARG_POOL_SIZE`, `PRIMUS_TURBO_DEEPEP_TIMEOUT`, `NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, `GLOO_SOCKET_IFNAME`, `NCCL_IB_GID_INDEX`, `PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32`, `NVTE_CK_IS_V3_ATOMIC_FP32`, `PATH_TO_BNXT_TAR_PACKAGE`, `ANP_HOME_DIR`, `RCCL_HOME_DIR`, `MPI_HOME_DIR`, `DUMP_HLO`, `DUMP_HLO_DIR`, `PRIMUS_DETERMINISTIC`, `PRIMUS_HIPBLASLT_TUNING`, `PRIMUS_HIPBLASLT_TUNING_STAGE`, `TE_HIPBLASLT_TUNING_RUN_COUNT`, `TE_HIPBLASLT_TUNING_ALGO_COUNT`, `HIPBLASLT_LOG_MASK`, `HIPBLASLT_LOG_FILE`, `HIPBLASLT_LOG_LEVEL`, `HIPBLASLT_TUNING_OVERRIDE_FILE`
+
+---
+
+## 12. Slurm
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `SLURM_NNODES` / `SLURM_JOB_NUM_NODES` | job-dependent | Slurm | `primus-cli-slurm-entry.sh` (`NNODES` export), preflight probes | Node count for the allocation. |
+| `SLURM_NODEID` | job-dependent | Slurm | Mapped to `NODE_RANK` in `primus-cli-slurm-entry.sh` | Node index. |
+| `SLURM_PROCID` | job-dependent | Slurm | Fallback for `NODE_RANK` when `SLURM_NODEID` is unset | Process ID within the Slurm step (entry script). |
+| `SLURM_JOB_ID` | job-dependent | Slurm | `primus/tools/preflight/host/host_probe.py` | Job identifier string. |
+
+---
+
+## 13. Debug and pipeline
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `DUMP_PP_DIR` | `output/pp_data` | User | `primus/backends/megatron/megatron_pretrain_trainer.py`, `primus/backends/megatron/patches/pp_dump_data_patches.py` | Directory for pipeline-parallel debug dumps. |
+| `DEBUG_SIMULATOR` | `0` | User | `primus/core/projection/performance_projection/simulator.py` | `1` enables verbose projection simulator logging. |
+| `RECORD_OFFLOAD_MEMORY_INFO` | `0` | User | `primus/core/pipeline_parallel/handler/offload_handler.py` | Record offload memory stats when `1`. |
+| `RECORD_OFFLOAD_MEMORY_INFO_DIR` | `output` | User | `primus/core/pipeline_parallel/scheduler/scheduler.py` | Output directory for offload memory logs. |
+| `USE_PINNED_OFFLOAD` | `0` | User | `offload_handler.py` | Use pinned host memory for offload buffers when `1`. |
+
+---
+
+## 14. JAX / XLA (MaxText)
+
+Primus MaxText hooks print recommended values in `runner/helpers/hooks/train/pretrain/maxtext/prepare.py`; MaxText and JAX read them directly.
+
+| Variable | Default | Where set | Where used | Description |
+|----------|---------|-----------|------------|-------------|
+| `XLA_PYTHON_CLIENT_MEM_FRACTION` | e.g. `.97` in prepare hook | User / hook output | JAX / XLA allocator | Fraction of GPU memory pre-allocated for JAX. |
+| `DUMP_HLO_DIR` | `${PRIMUS_PATH}/output/xla_dump_hlo` (example) | User | XLA via `XLA_FLAGS` composition | Directory for HLO dumps when enabled. |
+| `DUMP_HLO` | `0` | User | Prepare hook → XLA flags | Gate HLO dumping (`1` enables in hook samples). |
+
+**Note:** MaxText also propagates many knobs through `XLA_FLAGS` and `LIBTPU_INIT_ARGS` upstream; see MaxText sources for the full list.
diff --git a/docs/03-configuration-reference/maxtext-parameters.md b/docs/03-configuration-reference/maxtext-parameters.md
new file mode 100644
index 000000000..17f03bdbd
--- /dev/null
+++ b/docs/03-configuration-reference/maxtext-parameters.md
@@ -0,0 +1,146 @@
+# MaxText backend configuration reference
+
+Primus routes experiment YAML into the [MaxText](https://maxtext.readthedocs.io/) stack (JAX / XLA). Configuration is a **flat map of keys** (no nested `training.`* trees like TorchTitan): Primus merges module and model presets, writes a temporary YAML, and MaxText’s `pyconfig.initialize` loads it on top of upstream defaults.
+
+The Primus overlay keeps `base_config: "base.yml"` so MaxText loads its own [`configs/base.yml`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/base.yml) at runtime. This page lists **Primus-defined defaults and commonly overridden Primus fields**. For the full upstream parameter set (hundreds of keys), see the [MaxText documentation](https://maxtext.readthedocs.io/) and upstream `base.yml`.
+
+---
+
+## How parameters flow
+
+1. YAML presets under `primus/configs/modules/maxtext/` and `primus/configs/models/maxtext/` are merged with CLI overrides.
+2. `MaxTextAdapter.convert_config` passes the merged namespace through `MaxTextConfigBuilder` (currently a thin pass-through).
+3. `export_params_to_yaml` writes a flat YAML file; MaxText ignores unknown Primus-private keys via pydantic filtering.
+4. Unknown keys from upstream still resolve through environment overrides inside MaxText (`pyconfig`), not shown here.
+
+---
+
+## 1. Base module parameters
+
+Shared with all Primus modules via `module_base.yaml` and trainer extensions.
+
+
+| Parameter           | Default (Primus)                                                  | Description                                                                                  |
+| ------------------- | ----------------------------------------------------------------- | -------------------------------------------------------------------------------------------- |
+| `trainable`         | `true` in `trainer_base.yaml` (overrides `module_base`’s `false`) | When `true`, the module participates in training orchestration.                              |
+| `sink_level`        | `null`                                                            | Structured logging sink level for the module (if the logging stack is configured to use it). |
+| `file_sink_level`   | `DEBUG`                                                           | File sink verbosity.                                                                         |
+| `stderr_sink_level` | `INFO`                                                            | Stderr sink verbosity.                                                                       |
+
+
+---
+
+## 2. Training
+
+From `pre_trainer.yaml` (extends `trainer_base.yaml`).
+
+
+| Parameter     | Default      | Description                                                   |
+| ------------- | ------------ | ------------------------------------------------------------- |
+| `base_config` | `"base.yml"` | Upstream MaxText base file loaded by `pyconfig._load_config`. |
+| `hardware`    | `"gpu"`      | Hardware target string consumed by MaxText.                   |
+| `steps`       | `1000`       | Global optimizer steps for the run.                           |
+| `log_period`  | `100`        | Steps between log emissions.                                  |
+
+
+---
+
+## 3. Data
+
+
+| Parameter        | Default        | Description                                                          |
+| ---------------- | -------------- | -------------------------------------------------------------------- |
+| `dataset_type`   | `"hf"`         | Dataset backend selector (Hugging Face in the default path).         |
+| `hf_path`        | `"allenai/c4"` | Hugging Face dataset repo or identifier.                             |
+| `hf_data_dir`    | `"en"`         | Subdirectory / config slice within the HF dataset.                   |
+| `hf_train_files` | `""`           | Optional explicit train file list (format per MaxText HF loader).    |
+| `packing`        | `true`         | Sequence packing for efficiency when supported by the data pipeline. |
+
+
+---
+
+## 4. Checkpointing
+
+These are Primus overlay defaults. MaxText also loads upstream `base.yml` at runtime through `base_config: "base.yml"`, where upstream checkpoint defaults might differ. When debugging effective behavior, distinguish the Primus YAML written by the adapter from the upstream MaxText defaults loaded afterward.
+
+
+| Parameter              | Default | Description                                                        |
+| ---------------------- | ------- | ------------------------------------------------------------------ |
+| `enable_checkpointing` | `false` | See Training section.                                              |
+| `async_checkpointing`  | `false` | When `enable_checkpointing` is true, use async checkpoint workers. |
+
+
+---
+
+## 5. Profiling
+
+
+| Parameter                         | Default    | Description                             |
+| --------------------------------- | ---------- | --------------------------------------- |
+| `profiler`                        | `"xplane"` | Profiler backend (e.g. XPlane for JAX). |
+| `skip_first_n_steps_for_profiler` | `3`        | Warmup steps excluded from capture.     |
+| `profiler_steps`                  | `1`        | Number of steps to profile once active. |
+
+
+---
+
+## 6. Memory and recomputation
+
+
+| Parameter                       | Default  | Description                                                                        |
+| ------------------------------- | -------- | ---------------------------------------------------------------------------------- |
+| `remat_policy`                  | `'full'` | Activation rematerialization policy (`none`, `minimal`, `full`, etc.—see MaxText). |
+| `optimizer_memory_host_offload` | `false`  | Offload optimizer state to host memory when supported.                             |
+| `scan_layers`                   | `true`   | Use scanned layer implementation where applicable.                                 |
+| `param_scan_axis`               | `1`      | Axis for parameter scanning / partitioning layout.                                 |
+
+
+---
+
+## 7. Precision and quantization
+
+
+| Parameter                 | Default           | Description                                                           |
+| ------------------------- | ----------------- | --------------------------------------------------------------------- |
+| `dtype`                   | `"bfloat16"`      | Default compute dtype for many ops.                                   |
+| `quantization`            | `""`              | Quantization mode string (empty = none; set per MaxText AQT recipes). |
+| `quantize_kvcache`        | `false`           | Quantize KV cache tensors.                                            |
+| `kv_quant_axis`           | `"heads_and_dkv"` | KV quantization axis naming for kernels.                              |
+| `kv_quant_dtype`          | `"int8"`          | Storage dtype for KV cache when quantization is on.                   |
+| `weight_dtype`            | `bfloat16`        | Weight storage / compute dtype for non-quantized paths.               |
+| `checkpoint_is_quantized` | `false`           | Set `true` when loading an AQT-quantized checkpoint.                  |
+| `logits_dot_in_fp32`      | `false`           | Compute logits matmul in `float32` for numerical stability.           |
+
+
+---
+
+## 8. Model
+
+From `model_base.yaml` and per-model files such as `llama3_8B.yaml`.
+
+
+| Parameter               | Default                                                             | Description                                                             |
+| ----------------------- | ------------------------------------------------------------------- | ----------------------------------------------------------------------- |
+| `model_name`            | `"default"` in `model_base`; e.g. `"llama3-8b"` in `llama3_8B.yaml` | Selects MaxText’s bundled model YAML when present.                      |
+| `override_model_config` | `true`                                                              | When `true`, CLI / kwargs override values from the loaded model config. |
+| `attention`             | `"cudnn_flash_te"`                                                  | Attention implementation (Primus default favors TE flash on AMD GPUs).  |
+| `use_iota_embed`        | `true`                                                              | Use iota-based embedding for performance on accelerator backends.       |
+| `tokenizer_path`        | e.g. `"meta-llama/Meta-Llama-3-8B"`                                 | Hugging Face tokenizer id or local path.                                |
+
+
+---
+
+## 9. Advanced
+
+
+| Parameter | Default | Description                                                           |
+| --------- | ------- | --------------------------------------------------------------------- |
+| `shardy`  | `false` | Enable Shardy-related integration in MaxText when building shardings. |
+
+
+---
+
+## Related reading
+
+- [MaxText documentation](https://maxtext.readthedocs.io/)—full parameter reference and recipes.
+- Primus implementation: `primus/backends/maxtext/argument_builder.py`, `maxtext_pretrain_trainer.py`, `maxtext_adapter.py`.
diff --git a/docs/03-configuration-reference/megatron-bridge-parameters.md b/docs/03-configuration-reference/megatron-bridge-parameters.md
new file mode 100644
index 000000000..156b82a5b
--- /dev/null
+++ b/docs/03-configuration-reference/megatron-bridge-parameters.md
@@ -0,0 +1,167 @@
+# Megatron Bridge backend configuration reference
+
+Megatron Bridge integrates [Megatron-Core](https://github.com/NVIDIA/Megatron-LM) training with Hugging Face–centric workflows. In Primus, the **`megatron_bridge`** framework is used for post-training with module preset `sft_trainer.yaml`, and the repository also ships a pretraining preset at `primus/configs/modules/megatron_bridge/pretrain_trainer.yaml`.
+
+## Recipe system
+
+Megatron Bridge resolves training defaults through a **recipe** and **flavor**:
+
+- `recipe` is a Python module path under `megatron.bridge.recipes` (e.g. `qwen.qwen3`).
+- `flavor` is the function name inside that module (e.g. `qwen3_8b_finetune_config`) that returns a `ConfigContainer`.
+
+At runtime, `load_recipe_config` in `primus/backends/megatron_bridge/config_utils.py`:
+
+1. Imports `megatron.bridge.recipes.<recipe>` and calls `<flavor>(**filtered_backend_args)` to build the baseline `ConfigContainer`.
+2. **Deep-merges** Primus `backend_args` (from YAML + CLI) into that dataclass via `_merge_dict_to_dataclass`, so user overrides sit on top of recipe defaults.
+
+You normally specify `recipe`, `flavor`, `hf_path`, and `dataset` in the model YAML; training hyperparameters and parallelism go in module overrides or experiment module overrides (`modules.post_trainer.overrides` for SFT/post-training, `modules.pre_trainer.overrides` for pretraining examples).
+
+---
+
+## 1. Base module parameters
+
+From `primus/configs/modules/megatron_bridge/sft_trainer.yaml` (extends `module_base.yaml`). Pretraining examples use `pretrain_trainer.yaml` instead.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `trainable` | `true` | Module participates in the training graph. |
+| `sink_level` | `null` | Inherited from `module_base.yaml`; structured logging sink level. |
+| `file_sink_level` | `DEBUG` | File sink verbosity. |
+| `stderr_sink_level` | `INFO` | Stderr sink verbosity. |
+
+---
+
+## 2. Training
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `stage` | `"sft"` | Backend stage selector. Primus dispatches post-training via `primus train posttrain` and loads the Megatron Bridge posttrain trainer when this module is used under `post_trainer`. |
+| `trainable` | `true` | See Base module parameters. |
+
+**CLI note:** The user-facing suite is **`posttrain`** (`primus train posttrain --config ...`). The YAML `stage` field selects the Megatron Bridge trainer implementation (`sft`), not the CLI suite name.
+
+For Bridge pretraining, use the normal pretraining suite (`primus train pretrain --config ...`) with experiments that reference `modules.pre_trainer.config: pretrain_trainer.yaml`.
+
+---
+
+## 3. Fine-tuning method (PEFT)
+
+Primus examples set these under `modules.post_trainer.overrides` (see `examples/megatron_bridge/configs/`).
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `peft` | `"none"`, `"lora"` | Parameter-efficient fine-tuning mode. |
+| `peft_dim` | `16` | LoRA rank (example: `llama31_70b_lora_posttrain.yaml`). |
+| `peft_alpha` | `32` | LoRA scaling alpha (same example). |
+| `packed_sequence` | `false` | Pack multiple short sequences per microbatch when supported. |
+
+Additional keys such as `pretrained_checkpoint`, `use_distributed_optimizer`, or `cross_entropy_loss_fusion` appear in larger examples and are merged into the recipe `ConfigContainer` when the dataclass exposes matching fields.
+
+---
+
+## 4. Parallelism
+
+Typical overrides from Megatron Bridge examples:
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `tensor_model_parallel_size` | `1`, `2`, `8` | Tensor parallelism degree. |
+| `pipeline_model_parallel_size` | `1` | Pipeline parallelism degree. |
+| `virtual_pipeline_model_parallel_size` | `null` | Virtual pipeline stages per rank when PP > 1. |
+| `context_parallel_size` | `1` | Context parallelism degree. |
+| `sequence_parallel` | `false` | Sequence parallelism within TP groups. |
+| `use_megatron_fsdp` | `false` | Optional Megatron FSDP path. |
+
+---
+
+## 5. Training hyperparameters
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `train_iters` | `200`, `1000` | Total training iterations. |
+| `global_batch_size` | `8`, `128` | Global batch across data-parallel groups. |
+| `micro_batch_size` | `1`, `2` | Per-GPU microbatch before gradient accumulation. |
+| `seq_length` | `2048`, `8192` | Training sequence length. |
+| `eval_interval` | `30` | Steps between evaluations. |
+| `save_interval` | `50` | Steps between checkpoint saves. |
+
+---
+
+## 6. Learning rate
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `finetune_lr` | `1.0e-4`, `5.0e-6` | Peak learning rate for fine-tuning. |
+| `min_lr` | `0.0` | Floor learning rate after decay. |
+| `lr_warmup_iters` | `50` | Linear warmup length in iterations. |
+| `lr_decay_iters` | `null` | Optional decay span; `null` defers to recipe defaults. |
+
+---
+
+## 7. Precision
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `precision_config` | `bf16_mixed`, `fp16_mixed`, `fp32` | Mixed-precision recipe for Megatron Bridge. |
+| `comm_overlap_config` | `null` | Optional communication/compute overlap policy object. |
+| `pipeline_dtype` | `null` | Dtype for pipeline stages when PP is enabled. |
+
+---
+
+## 8. Memory optimization
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `recompute_granularity` | `full` | Activation recomputation granularity. |
+| `recompute_method` | `uniform` | How recomputation is scheduled across layers. |
+| `recompute_num_layers` | `1` | Number of layers per recompute group (workload-dependent). |
+
+---
+
+## 9. Primus-Turbo
+
+From `sft_trainer.yaml` (defaults shown).
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `enable_primus_turbo` | `true` | Master flag for Primus-Turbo optimized kernels and paths. |
+| `use_turbo_attention` | `false` | Turbo attention implementation. |
+| `use_turbo_parallel_linear` | `false` | Turbo parallel linear layers. |
+| `use_turbo_grouped_gemm` | `false` | Turbo grouped GEMM flag for MoE paths (the preset ships this key). The former `use_turbo_grouped_mlp` alias has been removed. |
+| `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary loss handling. |
+| `enable_turbo_attention_float8` | `false` | FP8 path inside Turbo attention. |
+| `use_turbo_deepep` | `false` | DeepEP-style expert-parallel integration. |
+| `turbo_deepep_num_cu` | `32` | Compute-unit count hint for DeepEP. |
+| `turbo_deepep_use_comm_stream` | `false` | Use dedicated communication streams. |
+| `turbo_sync_free_moe_stage` | `0` | Sync-free MoE scheduling stage. |
+| `use_turbo_fused_act_with_probs` | `false` | Fuse activation with probability tensors where applicable. |
+| `use_turbo_rms_norm` | `false` | Turbo RMSNorm path. |
+
+**Environment:** `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND` (default `TURBO`) is read in `primus/backends/megatron/patches/args/rocm_arg_validation.py` to select MoE dispatch/combine behavior when Turbo MoE is active.
+
+---
+
+## 10. Model and dataset
+
+Model YAML files (`qwen3_8b.yaml`, `qwen3_32b.yaml`, `llama31_70b.yaml`) supply:
+
+| Parameter | Example | Description |
+|-----------|---------|-------------|
+| `recipe` | `qwen.qwen3`, `llama.llama3` | Recipe module under `megatron.bridge.recipes`. |
+| `flavor` | `qwen3_8b_finetune_config`, `llama31_70b_finetune_config` | Flavor function producing the baseline `ConfigContainer`. |
+| `hf_path` | `Qwen/Qwen3-8B`, `meta-llama/Meta-Llama-3.1-70B` | Hugging Face model id for weights/tokenizer flows. |
+| `dataset` | nested | Example: `dataset_name: "rajpurkar/squad"` for SQuAD-style fine-tuning. |
+
+**Logging (optional overrides in examples):** `wandb_project`, `wandb_entity`, `wandb_exp_name` might be set under `overrides` for experiment tracking when Weights & Biases is configured.
+
+---
+
+## Argument merge mechanics
+
+`MegatronBridgeArgBuilder` (`primus/backends/megatron_bridge/argument_builder.py`) performs a **deep merge** of CLI and YAML into a single dict/namespace before `load_recipe_config` runs. Nested dicts (for example dataset or optimizer sections) combine recursively; explicit `None` in the merged structure can clear fields depending on merge rules in `_merge_dict_to_dataclass`.
+
+---
+
+## Example layouts
+
+Under `examples/megatron_bridge/configs/`, per-GPU directories (for example `MI300X/`, `MI355X/`) contain full experiment YAMLs that set `work_group`, `user_name`, `exp_name`, `workspace`, and Megatron Bridge modules. Post-training examples use `modules.post_trainer` with `config: sft_trainer.yaml`; MI300X pretraining examples use `modules.pre_trainer` with `config: pretrain_trainer.yaml`. Both patterns set `framework: megatron_bridge`, `model: <preset>.yaml`, and an `overrides` block for parallelism, LR, precision, and related options.
diff --git a/docs/03-configuration-reference/megatron-parameters.md b/docs/03-configuration-reference/megatron-parameters.md
new file mode 100644
index 000000000..97e0bfc04
--- /dev/null
+++ b/docs/03-configuration-reference/megatron-parameters.md
@@ -0,0 +1,857 @@
+# Megatron backend configuration reference
+
+This page lists the flat configuration keys exposed by Primus when `framework: megatron`. Unless a section says otherwise, values are the defaults from `primus/configs/modules/megatron/trainer_base.yaml` and related model presets. The effective pretraining preset is `pre_trainer.yaml`, which extends `trainer_base.yaml` and overrides several high-impact training defaults.
+
+**Where parameters live.** Set overrides under `modules.pre_trainer.overrides:` in your experiment YAML. Model architecture keys usually come from `models.<role>.overrides:` (or your chosen model preset), but the same names map to Megatron’s argparse namespace either way.
+
+**Presets.**
+
+- Module presets: `primus/configs/modules/megatron/` (the main pretraining bundle is `pre_trainer.yaml`, which extends `trainer_base.yaml` and Primus Megatron add-ons).
+- Model presets: `primus/configs/models/megatron/` (for example `language_model.yaml`).
+
+**Mapping to Megatron-LM.** Keys are passed through **1:1** to Megatron’s training arguments (same names as `argparse` / `Namespace`). Primus builds that namespace with `MegatronArgBuilder`.
+
+**Upstream reference.** Full flag semantics and newer options are defined in Megatron-LM: [`megatron/training/arguments.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py).
+
+### Example (experiment YAML)
+
+```yaml
+framework: megatron
+
+modules:
+  pre_trainer:
+    overrides:
+      global_batch_size: 256
+      train_iters: 50000
+      tensor_model_parallel_size: 2
+
+models:
+  pre_train:
+    overrides:
+      hidden_size: 2048
+      num_layers: 32
+```
+
+---
+
+## 1. Base module parameters
+
+*Source: `primus/configs/modules/module_base.yaml` (merged into Megatron presets; `trainer_base.yaml` sets `trainable: true`).*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `trainable` | `true` | When `true`, this module participates in training workflows. (`module_base.yaml` alone defaults to `false`; Megatron `trainer_base.yaml` overrides to `true`.) |
+| `sink_level` | `null` | Log level for the structured sink (Primus module plumbing); `null` uses framework default. |
+| `file_sink_level` | `DEBUG` | Minimum level for file-backed logging. |
+| `stderr_sink_level` | `INFO` | Minimum level for stderr logging. |
+
+---
+
+## 2. Training and batching
+
+*Source: `primus/configs/modules/megatron/trainer_base.yaml`; effective `pre_trainer.yaml` overrides are noted where they differ.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `yaml_cfg` | `null` | Reserved; not supported as a Megatron override in this preset. |
+| `spec` | `null` | Optional trainer spec hook (unused in defaults). |
+| `micro_batch_size` | `2` | Samples per microbatch per data-parallel rank (per forward/backward step before gradient accumulation). |
+| `batch_size` | `null` | Deprecated; use `micro_batch_size` / `global_batch_size`. |
+| `global_batch_size` | `128` (`16` in `pre_trainer.yaml`) | Total batch size across the data-parallel world (before or after splitting, per Megatron semantics). |
+| `rampup_batch_size` | `null` | Optional batch-size ramp schedule string / config. |
+| `decrease_batch_size_if_needed` | `false` | Allow shrinking batch if memory is insufficient. |
+| `check_for_nan_in_loss_and_grad` | `true` | Abort on NaNs in loss or gradients. |
+| `check_for_spiky_loss` | `false` | Detect abnormal loss spikes. |
+| `check_for_large_grads` | `false` | Detect abnormally large gradients. |
+| `make_vocab_size_divisible_by` | `128` | Pads vocabulary size for efficient kernels / partitioning. |
+| `exit_signal_handler` | `false` | Install handlers for graceful shutdown signals. |
+| `exit_duration_in_mins` | `null` | Stop training after this many minutes. |
+| `exit_interval` | `null` | Exit after this many iterations (if set). |
+| `onnx_safe` | `null` | ONNX export compatibility tweaks. |
+| `bert_binary_head` | `true` | Use BERT binary classification head when applicable. |
+| `use_flash_attn` | `false` (`true` in `pre_trainer.yaml`) | Prefer FlashAttention kernels when available. |
+| `seed` | `1234` | RNG seed for reproducibility. |
+| `data_parallel_random_init` | `false` | Random init that varies across data-parallel ranks. |
+| `init_method_xavier_uniform` | `false` | Use Xavier uniform for some weights. |
+| `test_mode` | `false` | Lightweight test path (fewer steps / checks). |
+| `train_iters` | `null` (`1000` in `pre_trainer.yaml`) | Total training iterations (mutually exclusive with sample-based stopping in typical setups). |
+| `train_samples` | `null` | Total training samples (when using sample-based training). |
+| `eval_iters` | `32` (`0` in `pre_trainer.yaml`) | Validation iterations per eval. |
+| `eval_interval` | `2000` (`1000` in `pre_trainer.yaml`) | Run validation every this many iterations. |
+| `full_validation` | `false` | Run a full pass over validation data. |
+| `multiple_validation_sets` | `false` | Multiple validation datasets / passes. |
+| `skip_train` | `false` | Only run eval / test, no training updates. |
+| `train_sync_interval` | `null` | Periodic distributed sync barrier for debugging. |
+| `adlr_autoresume` | `false` | ADLR autoresume integration. |
+| `adlr_autoresume_interval` | `1000` | Autoresume checkpoint interval. |
+| `manual_gc` | `false` | Force Python GC on a schedule. |
+| `manual_gc_interval` | `1` | GC every N steps when `manual_gc` is enabled. |
+| `manual_gc_eval` | `false` | Run manual GC during evaluation. |
+| `mask_type` | `random` | Masking strategy for MLM / similar objectives. |
+| `mask_factor` | `1.0` | Masking strength multiplier. |
+| `iter_per_epoch` | `1250` | Iterations interpreted as one “epoch” for logging. |
+
+---
+
+## 3. Mixed precision
+
+*Source: `trainer_base.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `fp16` | `false` | Enable FP16 mixed precision training. |
+| `bf16` | `true` | Enable BF16 mixed precision training. |
+| `grad_reduce_in_bf16` | `false` | All-reduce gradients in BF16 (saves bandwidth). |
+| `calculate_per_token_loss` | `false` | Normalize loss per token instead of per sample. |
+| `loss_scale` | `null` | Static loss scale for FP16; `null` uses dynamic scaling. |
+| `initial_loss_scale` | `4294967296` | Initial dynamic loss scale. |
+| `min_loss_scale` | `1.0` | Floor for dynamic loss scale. |
+| `loss_scale_window` | `1000` | Window for dynamic loss scaling updates. |
+| `hysteresis` | `2` | Hysteresis steps for loss-scale decreases. |
+| `accumulate_allreduce_grads_in_fp32` | `false` | Accumulate and reduce gradients in FP32. |
+| `fp16_lm_cross_entropy` | `false` | Compute LM cross-entropy in FP16. |
+| `fp8` | `null` | FP8 recipe selection (`e4m3`, `hybrid`, etc.); `null` disables. |
+| `fp8_margin` | `0` | FP8 scaling margin. |
+| `fp8_recipe` | `delayed` | FP8 recipe variant (e.g. delayed scaling). |
+| `fp8_interval` | `1` | Deprecated FP8 interval (kept for compatibility). |
+| `fp8_amax_history_len` | `1024` | History length for FP8 amax statistics. |
+| `fp8_amax_compute_algo` | `"max"` | How to combine amax history (`max`, etc.). |
+| `fp8_wgrad` | `true` | Run weight gradients in FP8 where supported. |
+| `fp8_param_gather` | `false` | FP8 parameter gather for distributed optimizer paths. |
+| `te_rng_tracker` | `false` | Transformer Engine RNG tracker for FP8. |
+| `inference_rng_tracker` | `false` | Separate RNG tracker for inference FP8. |
+| `fp4` | `null` | FP4 mode; `null` disables. |
+| `fp4_recipe` | `nvfp4` | FP4 recipe name. |
+| `fp4_param` | `false` | Store parameters in FP4. |
+| `first_last_layers_bf16` | `false` | Keep first/last layers in BF16 for stability. |
+| `num_layers_at_start_in_bf16` | `1` | Count of early layers forced to BF16 when enabled. |
+| `num_layers_at_end_in_bf16` | `1` | Count of final layers forced to BF16 when enabled. |
+| `no_fp8_weight_transpose_cache` | `false` | *Primus:* disable FP8 weight transpose cache (see `primus_megatron_module.yaml`). |
+
+---
+
+## 4. Optimizer and learning rate
+
+*Source: `trainer_base.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `optimizer` | `adam` | Optimizer family (`adam`, `sgd`, etc.). |
+| `lr` | `2.5e-4` (`2.0e-05` in `pre_trainer.yaml`) | Peak learning rate. |
+| `lr_decay_style` | `cosine` | LR decay schedule (`cosine`, `linear`, `constant`, WSD, etc.). |
+| `lr_decay_iters` | `null` | Decay duration in iterations. |
+| `lr_decay_samples` | `null` | Decay duration in samples. |
+| `lr_warmup_fraction` | `null` | Warmup as a fraction of total train steps. |
+| `lr_warmup_iters` | `0` (`40` in `pre_trainer.yaml`) | Linear warmup steps. |
+| `lr_warmup_samples` | `0` | Warmup in samples. |
+| `lr_warmup_init` | `0.0` | LR at the start of warmup. |
+| `min_lr` | `2.5e-5` (`0.0` in `pre_trainer.yaml`) | Minimum LR after decay. |
+| `lr_wsd_decay_style` | `exponential` | Weight-decay schedule style for WSD when used. |
+| `lr_wsd_decay_samples` | `null` | WSD decay window in samples. |
+| `lr_wsd_decay_iters` | `null` | WSD decay window in iterations. |
+| `head_lr_mult` | `1.0` | LR multiplier for attention/head modules when supported. |
+| `weight_decay` | `0.01` (`0.0` in `pre_trainer.yaml`) | AdamW / L2-style weight decay. |
+| `start_weight_decay` | `null` | Starting weight decay for schedules. |
+| `end_weight_decay` | `null` | Ending weight decay for schedules. |
+| `weight_decay_incr_style` | `constant` | How weight decay changes between start/end. |
+| `clip_grad` | `1.0` | Global gradient norm clip. |
+| `adam_beta1` | `0.9` | Adam first moment decay. |
+| `adam_beta2` | `0.95` (`0.999` in `pre_trainer.yaml`) | Adam second moment decay. |
+| `adam_eps` | `1.0e-08` | Adam epsilon. |
+| `sgd_momentum` | `0.9` | SGD momentum when `optimizer` is SGD. |
+| `override_opt_param_scheduler` | `false` (`true` in `pre_trainer.yaml`) | Override optimizer parameter groups’ schedulers. |
+| `use_checkpoint_opt_param_scheduler` | `false` | Load optimizer scheduler state strictly from checkpoint. |
+| `warmup` | `null` | Alternate warmup specification (legacy / schedule hooks). |
+| `decoupled_lr` | `null` | Decoupled LR for certain param groups. |
+| `decoupled_min_lr` | `null` | Minimum for decoupled LR. |
+| `muon_extra_scale_factor` | `1.0` | Muon optimizer scaling. |
+| `muon_scale_mode` | `"spectral"` | Muon scaling mode. |
+| `muon_fp32_matmul_prec` | `"medium"` | Muon matmul precision hint. |
+| `muon_num_ns_steps` | `5` | Muon Newton–Schulz iterations. |
+| `muon_tp_mode` | `"blockwise"` | Muon tensor-parallel mode. |
+| `muon_use_nesterov` | `false` | Muon Nesterov momentum. |
+| `muon_split_qkv` | `true` | Split QKV for Muon. |
+| `muon_momentum` | `0.95` | Muon momentum. |
+| `muon_weight_decay` | `0.01` | Muon-specific decay. |
+| `muon_weight_decay_method` | `"decoupled"` | How Muon applies decay. |
+| `optimizer_cpu_offload` | `false` | Offload optimizer state to CPU. |
+| `optimizer_offload_fraction` | `1.0` | Fraction of optimizer state offloaded. |
+| `use_torch_optimizer_for_cpu_offload` | `false` | Use PyTorch optimizer for offload path. |
+| `overlap_cpu_optimizer_d2h_h2d` | `false` | Overlap CPU optimizer device transfers. |
+| `pin_cpu_grads` | `true` | Pin memory for CPU gradients. |
+| `pin_cpu_params` | `true` | Pin memory for CPU params in offload. |
+| `use_precision_aware_optimizer` | `false` | Use precision-aware optimizer (main grads/params in lower precision). |
+| `main_grads_dtype` | `fp32` | Dtype for main gradients (`fp32`, `bf16`). |
+| `main_params_dtype` | `fp32` | Dtype for master params. |
+| `exp_avg_dtype` | `fp32` | Optimizer first moment dtype (`fp32`, `fp16`, `fp8`). |
+| `exp_avg_sq_dtype` | `fp32` | Optimizer second moment dtype. |
+
+---
+
+## 5. Parallelism and distribution
+
+*Sources: `trainer_base.yaml` (distributed runtime) and `primus/configs/models/megatron/language_model.yaml` (model-parallel sizes and TP communication).*
+
+### 5.1 Data / distributed runtime (trainer)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `overlap_p2p_comm` | `true` | Overlap pipeline P2P with compute. |
+| `distributed_backend` | `nccl` | Process-group backend (`nccl`, `gloo`, …). |
+| `distributed_timeout_minutes` | `10` (`60` in `pre_trainer.yaml`) | Collective timeout. |
+| `defer_embedding_wgrad_compute` | `false` | Defer embedding weight gradients. |
+| `wgrad_deferral_limit` | `0` | Max deferred embedding wgrad steps. |
+| `align_grad_reduce` | `true` | Align gradient reductions for efficiency. |
+| `ddp_num_buckets` | `null` | Number of DDP buckets. |
+| `ddp_bucket_size` | `null` | DDP bucket size in elements. |
+| `ddp_pad_buckets_for_high_nccl_busbw` | `false` | Pad buckets for NCCL bus bandwidth. |
+| `ddp_average_in_collective` | `false` | Average inside collective vs outside. |
+| `overlap_grad_reduce` | `false` | Overlap gradient all-reduce with backward. |
+| `overlap_param_gather` | `false` | Overlap param all-gather (distributed optimizer). |
+| `overlap_param_gather_with_optimizer_step` | `false` | Overlap param gather with optimizer step. |
+| `align_param_gather` | `true` | Align param gather for distributed optimizer. |
+| `scatter_gather_tensors_in_pipeline` | `true` | Scatter/gather tensors across PP ranks. |
+| `use_ring_exchange_p2p` | `false` | Ring-exchange P2P for PP. |
+| `local_rank` | `null` | Local rank override (normally from launcher). |
+| `lazy_mpu_init` | `null` | Defer Megatron parallel state init. |
+| `account_for_embedding_in_pipeline_split` | `false` | Account for embedding in PP partition. |
+| `account_for_loss_in_pipeline_split` | `false` | Account for loss partition in PP. |
+| `empty_unused_memory_level` | `0` | Aggressiveness of `torch.cuda.empty_cache`. |
+| `standalone_embedding_stage` | `false` | Dedicated PP stage for embeddings. |
+| `use_distributed_optimizer` | `false` (`true` in `pre_trainer.yaml`) | Shard optimizer state across data parallel. |
+| `use_sharp` | `false` | Use SHARP for collectives when available. |
+| `sharp_enabled_group` | `null` | Which group SHARP applies to (`dp`, `dp_replica`). |
+| `use_custom_fsdp` | `false` | Custom FSDP integration path. |
+| `use_megatron_fsdp` | `false` | Megatron FSDP path. |
+| `init_model_with_meta_device` | `false` | Build model on `meta` device first. |
+| `data_parallel_sharding_strategy` | `no_shard` | FSDP / ZeRO style sharding (`no_shard`, `optim`, …). |
+| `gradient_reduce_div_fusion` | `true` | Fuse division into reduce-scatter. |
+| `suggested_communication_unit_size` | `400000000` | Suggested communication chunk size. |
+| `keep_fp8_transpose_cache_when_using_custom_fsdp` | `false` | Keep FP8 transpose cache with custom FSDP. |
+| `num_distributed_optimizer_instances` | `1` | Sharded optimizer instances per rank group. |
+| `use_torch_fsdp2` | `false` | Use PyTorch FSDP2 integration. |
+| `nccl_communicator_config_path` | `null` | JSON config for NCCL communicators. |
+| `use_tp_pp_dp_mapping` | `false` | Custom TP/PP/DP process mapping. |
+| `replication` | `false` | Data replication mode for certain schedules. |
+| `replication_jump` | `null` | Stride between replicated ranks. |
+| `replication_factor` | `null` | Replication factor. |
+| `deterministic_mode` | `false` | Prefer deterministic algorithms (slower). |
+| `check_weight_hash_across_dp_replicas_interval` | `null` | Periodically hash weights across DP replicas for debugging. |
+| `overlap_moe_expert_parallel_comm` | `false` | Overlap MoE expert-parallel communication. |
+| `decoder_pipeline_manual_split_list` | `null` | *Primus:* manual PP split points for decoder (list of ints). |
+| `patch_moe_overlap` | `false` | *Primus:* patch MoE compute/comm overlap. |
+
+### 5.2 Model parallelism (model preset)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `model_parallel_size` | `null` | Legacy combined MP size override. |
+| `tensor_model_parallel_size` | `1` | Tensor parallelism degree (intra-layer split). |
+| `encoder_tensor_model_parallel_size` | `0` | Encoder TP size when encoder/decoder differ. |
+| `pipeline_model_parallel_size` | `1` | Pipeline parallelism stages. |
+| `pipeline_model_parallel_layout` | `null` | Optional explicit PP layout string. |
+| `pipeline_model_parallel_comm_backend` | `null` | `nccl` or `ucc` for PP collectives. |
+| `encoder_pipeline_model_parallel_size` | `0` | Encoder PP stages (encoder–decoder models). |
+| `pipeline_model_parallel_split_rank` | `null` | Rank where encoder/decoder split. |
+| `decoder_first_pipeline_num_layers` | `null` | Layers on first decoder PP stage. |
+| `decoder_last_pipeline_num_layers` | `null` | Layers on last decoder PP stage. |
+| `virtual_pipeline_model_parallel_size` | `null` | Virtual PP (interleaved) depth. |
+| `num_layers_per_virtual_pipeline_stage` | `null` | Layers per virtual stage. |
+| `num_virtual_stages_per_pipeline_rank` | `null` | Virtual stages per physical PP rank. |
+| `microbatch_group_size_per_vp_stage` | `null` | Microbatch grouping for interleaved PP. |
+| `sequence_parallel` | `true` | Sequence parallelism when TP > 1. |
+| `context_parallel_size` | `1` | Context (sequence) parallelism degree. |
+| `cp_comm_type` | `p2p` | Context-parallel comm pattern (`p2p`, `a2a`, `allgather`, `a2a+p2p`). |
+| `hierarchical_context_parallel_sizes` | `null` | Hierarchical CP group sizes. |
+| `expert_model_parallel_size` | `1` | Expert parallelism for MoE. |
+| `expert_tensor_parallel_size` | `null` | Expert tensor-parallel degree. |
+| `high_priority_stream_groups` | `[]` | Named groups that get high-priority CUDA streams. |
+
+### 5.3 Tensor-parallel communication overlap (model)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `async_tensor_model_parallel_allreduce` | `true` | Async TP all-reduces for column-parallel layers. |
+| `tp_comm_overlap` | `false` | Enable TP communication overlap planner. |
+| `tp_comm_overlap_cfg` | `null` | Extra JSON / path for overlap configuration. |
+| `tp_comm_overlap_ag` | `true` | Overlap all-gather in TP backward. |
+| `tp_comm_overlap_rs` | `true` | Overlap reduce-scatter in TP backward. |
+| `tp_comm_overlap_rs_dgrad` | `false` | Overlap RS for data-grad path. |
+| `tp_comm_split_ag` | `true` | Split all-gather for overlap. |
+| `tp_comm_split_rs` | `true` | Split reduce-scatter for overlap. |
+| `tp_comm_bulk_wgrad` | `true` | Bulk weight-gradient path for TP comm. |
+| `tp_comm_bulk_dgrad` | `true` | Bulk data-gradient path for TP comm. |
+| `barrier_with_L1_time` | `true` | Barrier using L1 timing hooks for TP comm profiling. |
+| `tp_comm_bootstrap_backend` | `nccl` | Backend used to bootstrap TP communicators. |
+
+---
+
+## 6. Checkpointing
+
+*Source: `trainer_base.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `save` | `null` | Path prefix / pattern for checkpoints to write. |
+| `save_interval` | `20000` (`1000` in `pre_trainer.yaml`) | Save every N iterations. |
+| `save_retain_interval` | `null` | Retain checkpoints at this interval. |
+| `no_save_optim` | `null` | Skip optimizer state in checkpoints when truthy. |
+| `no_save_rng` | `null` | Skip RNG state in checkpoints when truthy. |
+| `load` | `null` | Checkpoint path to load. |
+| `load_main_params_from_ckpt` | `false` | Load only main parameters. |
+| `no_load_optim` | `null` | Skip loading optimizer state. |
+| `no_load_rng` | `null` | Skip loading RNG state. |
+| `finetune` | `false` (`true` in `pre_trainer.yaml`) | Finetune mode (do not require full optimizer match). |
+| `use_checkpoint_args` | `false` | When `true`, restore training args from checkpoint metadata. |
+| `use_mp_args_from_checkpoint_args` | `false` | Restore model-parallel args from checkpoint. |
+| `use_tokenizer_model_from_checkpoint_args` | `true` | Restore tokenizer path from checkpoint args. |
+| `exit_on_missing_checkpoint` | `true` | Fail if `load` is set but checkpoint is missing. |
+| `non_persistent_save_interval` | `null` | Ephemeral checkpoint interval. |
+| `non_persistent_ckpt_type` | `null` | `global`, `local`, `in_memory`, or `null`. |
+| `non_persistent_global_ckpt_dir` | `null` | Directory for non-persistent global checkpoints. |
+| `non_persistent_local_ckpt_dir` | `null` | Directory for non-persistent local checkpoints. |
+| `non_persistent_local_ckpt_algo` | `"fully_parallel"` | `fully_parallel` or `atomic`. |
+| `pretrained_checkpoint` | `null` | Load weights from a pretrained checkpoint path. |
+| `ckpt_step` | `null` | Specific step to load within a distributed checkpoint. |
+| `use_dist_ckpt_deprecated` | `false` | Use deprecated distributed checkpoint format. |
+| `use_persistent_ckpt_worker` | `false` | Background worker for checkpoint IO. |
+| `auto_detect_ckpt_format` | `false` | Infer checkpoint format automatically. |
+| `dist_ckpt_format_deprecated` | `null` | Legacy format hint. |
+| `ckpt_format` | `torch_dist` | `torch`, `torch_dist`, or `zarr`. |
+| `ckpt_convert_format` | `null` | Target format for one-shot conversion. |
+| `ckpt_convert_save` | `null` | Output path for conversion. |
+| `ckpt_convert_update_legacy_dist_opt_format` | `false` | Update legacy distributed-optimizer layout when converting. |
+| `ckpt_fully_parallel_save_deprecated` | `false` | Deprecated fully-parallel save toggle. |
+| `ckpt_fully_parallel_save` | `true` | Save shards in parallel across ranks. |
+| `async_save` | `null` | Async checkpoint save (`null` = framework default). |
+| `ckpt_fully_parallel_load` | `false` | Load shards in parallel. |
+| `ckpt_assume_constant_structure` | `false` | Assume identical layer structure across ranks. |
+| `dist_ckpt_strictness` | `assume_ok_unexpected` | How to handle unexpected keys in distributed ckpt. |
+| `dist_ckpt_save_pre_mcore_014` | `null` | Compatibility flag for older Megatron-Core checkpoints. |
+| `dist_ckpt_optim_fully_reshardable` | `null` | Optimizer state fully reshardable layout. |
+| `auto_continue_train` | `false` | *Primus:* resume from latest checkpoint in the save directory when enabled. |
+| `disable_last_saving` | `false` | *Primus:* skip writing the final checkpoint at shutdown. |
+
+---
+
+## 7. Data
+
+*Source: `trainer_base.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `data_path` | `null` | Single blended dataset path / list. |
+| `data_sharding` | `true` | Shard data across ranks. |
+| `split` | `"99,1,0"` (`null` in `pre_trainer.yaml`) | Train/valid/test split ratios as comma string. |
+| `train_data_path` | `null` | Training data blend. |
+| `valid_data_path` | `null` | Validation data blend. |
+| `test_data_path` | `null` | Test data blend. |
+| `data_args_path` | `null` | External JSON/YAML of dataset arguments. |
+| `per_split_data_args_path` | `null` | Per-split dataset args file. |
+| `data_cache_path` | `null` | On-disk cache for indexed datasets. |
+| `mock_data` | `false` | Use synthetic data (no real files). |
+| `merge_file` | `null` | Merge file for blended datasets. |
+| `seq_length` | `4096` (`1024` in `pre_trainer.yaml`) | Training sequence length. |
+| `encoder_seq_length` | `null` | Encoder sequence length (encoder–decoder). |
+| `decoder_seq_length` | `null` | Decoder sequence length. |
+| `retriever_seq_length` | `256` | Sequence length for retriever models. |
+| `sample_rate` | `1.0` | Sampling rate for dataset blending. |
+| `mask_prob` | `0.15` | MLM mask probability. |
+| `short_seq_prob` | `0.1` | Probability of shorter sequences in BERT-style data. |
+| `num_workers` | `8` | DataLoader worker processes per rank. |
+| `reset_position_ids` | `false` | Reset position IDs at document boundaries. |
+| `reset_attention_mask` | `false` | Reset attention mask at boundaries. |
+| `eod_mask_loss` | `false` | Mask loss at end-of-document tokens. |
+| `dataloader_type` | `null` (`cyclic` in `pre_trainer.yaml`) | Dataloader implementation (`single`, `cyclic`, `external`, …). |
+| `mmap_bin_files` | `true` | Memory-map `.bin` index files when supported. |
+| `create_attention_mask_in_dataloader` | `true` | Build attention masks in the dataloader. |
+| `num_dataset_builder_threads` | `1` | Threads to build dataset indices. |
+
+---
+
+## 8. Recomputation (activation checkpointing)
+
+*Sources: `trainer_base.yaml` and `primus_megatron_module.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `recompute_activations` | `false` | Enable activation recomputation globally. |
+| `recompute_granularity` | `null` | `full` or `selective` checkpointing. |
+| `recompute_method` | `null` | `uniform` or `block` selective recomputation. |
+| `recompute_num_layers` | `null` | Layers to recompute per block / schedule. |
+| `recompute_layer_ids` | `null` | *Primus:* explicit **global** layer indices to recompute (`0 … num_layers-1`). |
+| `distribute_saved_activations` | `false` | Distribute saved activations across TP/PP for memory balance. |
+| `checkpoint_activations` | `false` | Deprecated alias for activation checkpointing. |
+| `moe_layer_recompute` | `false` | Recompute MoE layer activations (model preset). |
+
+---
+
+## 9. Logging and profiling
+
+*Sources: `trainer_base.yaml` and `primus_megatron_module.yaml`.*
+
+### 9.1 Logging
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `log_avg_skip_iterations` | `2` | Skip first N iterations for throughput averaging. |
+| `log_avg_reset_interval` | `10` | Reset moving averages periodically. |
+| `log_params_norm` | `false` | Log L2 norms of parameters. |
+| `log_num_zeros_in_grad` | `false` | Log fraction of zero gradients. |
+| `log_throughput` | `false` (`true` in `pre_trainer.yaml`) | Log tokens/sec and timing. |
+| `log_progress` | `false` | Verbose progress logging. |
+| `timing_log_level` | `0` | Verbosity for timing logs. |
+| `timing_log_option` | `minmax` | Aggregate style for timing (`minmax`, `all`, …). |
+| `tensorboard_log_interval` | `1` | Steps between TensorBoard scalars. |
+| `tensorboard_queue_size` | `1000` | TensorBoard event queue size. |
+| `log_timers_to_tensorboard` | `false` (`true` in `pre_trainer.yaml`) | Write timer stats to TensorBoard. |
+| `log_batch_size_to_tensorboard` | `false` (`true` in `pre_trainer.yaml`) | Log batch size. |
+| `log_learning_rate_to_tensorboard` | `true` | Log LR. |
+| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. |
+| `log_memory_to_tensorboard` | `false` | Log memory usage. |
+| `log_world_size_to_tensorboard` | `false` | Log distributed world size. |
+| `log_loss_scale_to_tensorboard` | `true` | Log FP16/FP8 loss scale. |
+| `wandb_project` | `null` | Weights & Biases project name. |
+| `wandb_exp_name` | `null` | W&B run name. |
+| `wandb_save_dir` | `null` | W&B local directory. |
+| `wandb_entity` | `null` | W&B entity / team. |
+| `enable_one_logger` | `true` | Enable NVIDIA OneLogger integration. |
+| `one_logger_project` | `megatron-lm` | OneLogger project string. |
+| `one_logger_run_name` | `null` | OneLogger run name. |
+| `log_interval` | `100` (`1` in `pre_trainer.yaml`) | Console log interval in iterations. |
+| `tensorboard_dir` | `null` | TensorBoard output directory. |
+| `logging_level` | `null` | Python logging level override. |
+| `config_logger_dir` | `""` | Directory for dumped config logs. |
+| `one_logger_async` | `false` | Async OneLogger flushing. |
+| `app_tag_run_name` | `null` | Application tag for telemetry. |
+| `app_tag_run_version` | `0.0.0` | Application tag version. |
+| `disable_tensorboard` | `true` | *Primus:* disable TensorBoard integration in Primus-wrapped runs. |
+| `disable_wandb` | `true` | *Primus:* disable W&B. |
+| `disable_mlflow` | `true` | *Primus:* disable MLflow. |
+| `mlflow_run_name` | `null` | *Primus:* MLflow run name. |
+| `mlflow_experiment_name` | `null` | *Primus:* MLflow experiment name. |
+| `use_rocm_mem_info` | `false` | *Primus:* collect ROCm memory info via `rocm-smi` every step when `true`. |
+| `use_rocm_mem_info_iters` | `[1, 2]` | *Primus:* iterations at which to log memory if `use_rocm_mem_info` is `false`. |
+
+### 9.2 Profiling
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `profile` | `false` | Enable lightweight Nsight / CUDA profiling hooks. |
+| `use_pytorch_profiler` | `false` | Enable `torch.profiler` regions. |
+| `profile_ranks` | `[0]` | Ranks to profile. |
+| `profile_step_start` | `10` | First step to profile. |
+| `profile_step_end` | `12` | Last step to profile. |
+| `iterations_to_skip` | `null` | Skip listed iterations in profiling. |
+| `result_rejected_tracker_filename` | `null` | Log rejected samples to this file. |
+| `enable_gloo_process_groups` | `true` | Create auxiliary Gloo groups for CPU-side ops. |
+| `record_memory_history` | `false` | Record CUDA memory history (debug). |
+| `memory_snapshot_path` | `snapshot.pickle` | Path for memory snapshot dumps. |
+| `disable_profiler_activity_cpu` | `false` | *Primus:* omit CPU activities from profiler traces. |
+| `torch_profiler_record_shapes` | `true` | *Primus:* record tensor shapes in PyTorch profiler. |
+| `torch_profiler_with_stack` | `true` | *Primus:* capture Python stacks in profiler. |
+| `torch_profiler_use_gzip` | `false` | *Primus:* gzip profiler outputs. |
+
+---
+
+## 10. Model architecture
+
+*Sources: `primus/configs/models/megatron/language_model.yaml` and `primus/configs/models/megatron/primus_megatron_model.yaml`.*
+
+### 10.1 Core architecture
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `use_legacy_models` | `false` | Use legacy Megatron model code paths. |
+| `deprecated_use_mcore_models` | `false` | Deprecated flag for Megatron-Core models; prefer current `transformer_impl` + stack. |
+| `model_type` | `gpt` | `gpt` or `mamba` family. |
+| `num_layers` | `24` | Transformer layers (decoder or unified stack). |
+| `encoder_num_layers` | `null` | Encoder depth (encoder–decoder). |
+| `decoder_num_layers` | `null` | Decoder depth. |
+| `hidden_size` | `1024` | Hidden / model width. |
+| `num_attention_heads` | `16` | Attention heads. |
+| `attention_backend` | `auto` | Attention kernel backend selection. |
+| `group_query_attention` | `false` | Enable grouped-query attention (GQA). |
+| `qk_layernorm` | `false` | LayerNorm on Q/K projections. |
+| `qk_l2_norm` | `false` | L2-normalize Q/K vectors. |
+| `num_query_groups` | `null` | Number of query groups for GQA; `null` means MHA. |
+| `add_position_embedding` | `false` | Add absolute position embeddings (non-RoPE stacks). |
+| `position_embedding_type` | `learned_absolute` | Position embedding style. |
+| `max_position_embeddings` | `null` | Maximum sequence positions (context length cap). |
+| `original_max_position_embeddings` | `null` | Original pretrained length for interpolation / scaling. |
+| `untie_embeddings_and_output_weights` | `true` | Separate input embedding and LM head weights. |
+| `ffn_hidden_size` | `null` | FFN hidden size; `null` often defaults via `hidden_size` heuristics. |
+| `kv_channels` | `null` | Per-head KV channels override. |
+| `hidden_dropout` | `0.1` | Dropout on residual / hidden states. |
+| `attention_dropout` | `0.1` | Attention dropout. |
+| `fp32_residual_connection` | `false` | Accumulate residuals in FP32. |
+| `apply_residual_connection_post_layernorm` | `false` | Apply residual after (vs before) norm where supported. |
+| `add_bias_linear` | `false` | Biases in linear / column-parallel layers. |
+| `add_qkv_bias` | `false` | Biases in QKV projections. |
+| `swiglu` | `true` | SwiGLU activation in FFN. |
+| `quick_geglu` | `false` | Faster GeGLU path. |
+| `openai_gelu` | `false` | OpenAI GELU variant. |
+| `squared_relu` | `false` | Squared ReLU activation. |
+| `rotary_base` | `10000` | RoPE base frequency. |
+| `rotary_percent` | `1.0` | Fraction of head dim spanned by RoPE. |
+| `rotary_interleaved` | `false` | Interleaved RoPE layout. |
+| `rotary_seq_len_interpolation_factor` | `null` | Positional interpolation factor for long contexts. |
+| `use_rotary_position_embeddings` | `null` | Force RoPE on/off; `null` follows model type. |
+| `use_rope_scaling` | `false` | Enable LLaMA-style rope scaling. |
+| `rope_scaling_factor` | `8.0` | Scaling factor for extended contexts (LLaMA-3 style). |
+| `transformer_impl` | `transformer_engine` | Backend library (`transformer_engine`, `local`, …). |
+| `rope_type` | `null` | `rope` or `yarn` style extensions. |
+| `norm_epsilon` | `1.0e-05` | LayerNorm / RMSNorm epsilon. |
+| `normalization` | `"LayerNorm"` | Norm type (`LayerNorm`, `RMSNorm` with TE, …). |
+| `apply_layernorm_1p` | `false` | LayerNorm with +1 offset trick. |
+| `clone_scatter_output_in_embedding` | `true` | Clone embedding scatter for autograd safety. |
+| `perform_initialization` | `true` | Run weight initialization. |
+| `use_cpu_initialization` | `null` | Initialize on CPU then move to GPU. |
+| `use_te_activation_func` | `false` | Use Transformer Engine activation kernels. |
+| `gradient_accumulation_fusion` | `true` | Fuse gradient accumulation kernels. |
+| `delay_wgrad_compute` | `false` | Delay weight-gradient computation for scheduling. |
+
+### 10.2 Tokenizer and vocabulary
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `tokenizer_type` | `null` | Tokenizer class name (`GPT2BPETokenizer`, `HuggingFaceTokenizer`, …). |
+| `tokenizer_model` | `null` | Path to tokenizer model / vocabulary file. |
+| `vocab_size` | `null` | Vocabulary size (often inferred from tokenizer). |
+| `vocab_file` | `null` | Vocabulary file path for BPE/WP tokenizers. |
+| `vocab_extra_ids` | `0` | Extra reserved token slots. |
+| `tiktoken_pattern` | `null` | Regex pattern for tiktoken. |
+| `tiktoken_num_special_tokens` | `1000` | Special token count for tiktoken setup. |
+| `tiktoken_special_tokens` | `null` | Serialized special tokens for tiktoken. |
+| `legacy_tokenizer` | `false` | Legacy tokenizer behavior. |
+| `trust_remote_code` | `false` | `trust_remote_code` for Hugging Face tokenizers. |
+
+### 10.3 Initialization and attention numerics
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `init_method_std` | `0.02` | Standard deviation for weight init. |
+| `apply_query_key_layer_scaling` | `false` | Scale Q/K by layer index (deprecated GPT-3 trick). |
+| `attention_softmax_in_fp32` | `false` | Force softmax in FP32. |
+
+### 10.4 Kernel fusion flags
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `bias_gelu_fusion` | `true` | Fuse bias + GELU. |
+| `cross_entropy_loss_fusion` | `false` | Fused cross-entropy + softmax. |
+| `cross_entropy_fusion_impl` | `"native"` | `native` or `te` fused CE. |
+| `bias_swiglu_fusion` | `true` | Fuse bias + SwiGLU. |
+| `masked_softmax_fusion` | `true` | Fused masked softmax. |
+| `no_persist_layer_norm` | `false` | Non-persistent LayerNorm mode in TE. |
+| `bias_dropout_fusion` | `true` | Fuse bias + dropout. |
+| `apply_rope_fusion` | `true` | Fused RoPE kernels. |
+
+### 10.5 Multi-latent attention (MLA)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `multi_latent_attention` | `false` | Enable MLA blocks instead of standard MHA. |
+| `q_lora_rank` | `null` | Low-rank query projection rank. |
+| `kv_lora_rank` | `32` | Low-rank KV compression rank. |
+| `qk_head_dim` | `128` | Q/K head dimension for MLA. |
+| `qk_pos_emb_head_dim` | `64` | Positional head dimension for MLA. |
+| `v_head_dim` | `128` | Value head dimension for MLA. |
+| `rotary_scaling_factor` | `1.0` | RoPE scaling inside MLA (distinct from `rope_scaling_factor` above). |
+| `mscale` | `1.0` | Yarn / scaling m-factor. |
+| `mscale_all_dim` | `1.0` | Yarn scaling on all dims. |
+
+### 10.6 Mixture-of-experts (MoE)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `num_experts` | `null` | Experts per MoE layer; `null` means dense model. |
+| `moe_layer_freq` | `1` | Every Nth layer is MoE (1 = every layer). |
+| `moe_ffn_hidden_size` | `null` | Expert FFN hidden size. |
+| `moe_shared_expert_overlap` | `false` | Shared expert overlaps routing. |
+| `moe_shared_expert_intermediate_size` | `null` | Shared expert FFN size. |
+| `moe_grouped_gemm` | `false` | Grouped GEMM for experts. |
+| `moe_router_load_balancing_type` | `"aux_loss"` | Router balancing (`aux_loss`, `seq_aux_loss`, `sinkhorn`, `none`). |
+| `moe_router_dtype` | `null` | Router activation dtype (`fp32`, `fp64`). |
+| `moe_router_score_function` | `softmax` | `softmax` or `sigmoid` routing scores. |
+| `moe_router_topk` | `2` | Experts to select per token. |
+| `moe_router_pre_softmax` | `false` | Apply softmax before top-k. |
+| `moe_router_num_groups` | `null` | Group-limited routing: number of expert groups. |
+| `moe_router_group_topk` | `null` | Groups to pick before top-k inside groups. |
+| `moe_router_topk_scaling_factor` | `null` | Scaling for routing logits. |
+| `moe_router_enable_expert_bias` | `false` | Learnable per-expert bias. |
+| `moe_router_bias_update_rate` | `1.0e-03` | Update rate for expert bias. |
+| `moe_use_legacy_grouped_gemm` | `false` | Legacy grouped GEMM path. |
+| `moe_aux_loss_coeff` | `0.0` | Auxiliary load-balancing loss weight. |
+| `moe_z_loss_coeff` | `null` | Router z-loss coefficient. |
+| `moe_input_jitter_eps` | `null` | Input jitter for router stability. |
+| `moe_token_dispatcher_type` | `allgather` | Token dispatch algorithm (`allgather`, `alltoall`, `flex`, `alltoall_seq`). |
+| `moe_enable_deepep` | `false` | DeepEP-style expert parallelism. |
+| `moe_per_layer_logging` | `false` | Per-layer MoE statistics logging. |
+| `moe_expert_capacity_factor` | `null` | Capacity factor for token dropping / padding. |
+| `moe_pad_expert_input_to_capacity` | `false` | Pad expert batches to capacity. |
+| `moe_token_drop_policy` | `probs` | Token dropping policy when over capacity. |
+| `moe_extended_tp` | `false` | Extended tensor-parallel for experts. |
+| `moe_use_upcycling` | `false` | Expert upcycling initialization. |
+| `moe_permute_fusion` | `false` | Fuse token permutation for MoE. |
+| `disable_primus_topk_router` | `false` | *Primus:* disable Primus top-k router patch. |
+| `moe_router_force_load_balancing` | `false` | *Primus:* force load-balanced routing. |
+| `use_deprecated_20241209_moe_layer` | `false` | *Primus:* legacy MoE layer implementation. |
+| `moe_router_force_load_balancing_type` | `even` | *Primus:* Control the force load balancing type for the MoE router. Choices: even, uniform. |
+
+
+### 10.7 Logit softcapping (Primus / Grok-style)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `final_logit_softcapping` | `null` | Softcap value for final logits; `null` disables. |
+| `attn_logit_softcapping` | `null` | Softcap for attention logits. |
+| `router_logit_softcapping` | `null` | Softcap for MoE router logits. |
+
+---
+
+## 11. Primus extensions
+
+### 11.1 Build and compile
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `disable_compile_dependencies` | `true` | *Primus:* avoid compiling dependency stacks in the trainer wrapper. |
+
+### 11.2 Primus-Turbo (`primus_turbo.yaml`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `enable_primus_turbo` | `false` | Master switch for Primus-Turbo integrations. Many sub-features require this plus specific kernels. |
+| `use_turbo_attention` | `false` | Turbo attention implementation. |
+| `use_sink_attention` | `false` | GPT-OSS-style learned sink attention. |
+| `sink_sliding_window` | `0` | Sliding-window size for sink attention (GPT-OSS uses `128`). |
+| `sink_window_even_layers_only` | `true` | Apply the sliding window only to even layers (GPT-OSS pattern). |
+| `use_turbo_gemm` | `false` | Active Turbo GEMM flag for Dense paths. |
+| `use_turbo_parallel_linear` | *(removed)* | Removed—use `use_turbo_gemm`. Passing this key now raises an assertion error (`use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead`). |
+| `use_turbo_grouped_gemm` | `false` | Active Turbo grouped GEMM flag for MoE paths. |
+| `use_turbo_grouped_mlp` | *(removed)* | Removed—use `use_turbo_grouped_gemm`. Passing this key now raises an assertion error (`use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead`). |
+| `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary scores. |
+| `enable_turbo_attention_float8` | `false` | FP8 path inside Turbo attention (spacing in YAML is normalized to this key). |
+| `use_turbo_deepep` | `false` | Turbo DeepEP expert communication. |
+| `turbo_deepep_num_cu` | `32` | DeepEP compute units / channels. |
+| `turbo_deepep_use_comm_stream` | `false` | Use a dedicated communication stream for DeepEP. |
+| `turbo_sync_free_moe_stage` | `0` | Stage selector for sync-free MoE. |
+| `use_turbo_fused_act_with_probs` | `false` | Fuse activation + probability tensors to remove redundant work. |
+| `use_turbo_rms_norm` | `false` | Turbo RMSNorm kernels. |
+
+### 11.3 Zero-bubble pipeline (`zero_bubble.yaml`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `patch_zero_bubble` | `false` | Install Primus zero-bubble PP patches when `true`. |
+| `debug_scheduler_table` | `false` | Print PP scheduler tables (also in `primus_pipeline.yaml`; last merge wins—defaults match). |
+| `enable_zb_runtime` | `true` | Unified runtime for zero-bubble and related schedules. |
+| `pre_communication_optimization` | `false` | Issue a tiny comm before real comm to tune overlap. |
+| `zero_bubble_pipeline_timers_start_iter` | `100` | Start iter for auto-scheduler timers. |
+| `zero_bubble_pipeline_timers_end_iter` | `110` | End iter for auto-scheduler timers. |
+| `zero_bubble_max_pending_backward` | `auto` | Max pending backward ops (ZB1p vs ZB2p style); `auto` adapts. |
+| `zero_bubble_adaptive_memory_limit_percentile` | `85` | GPU memory percentile cap for adaptive ZB. |
+| `enable_optimizer_post_validation` | `false` | Post-optimizer validation step (needs FSDP path). |
+| `enable_exactly_numeric_match` | `true` | Require bitwise match in post validation when enabled. |
+| `enable_zero_bubble` | `true` | Enable zero-bubble schedule features in the ZB runtime. |
+| `zero_bubble_v_schedule` | `false` | Zero-bubble “V” schedule without extra memory vs some baselines. |
+| `zero_bubble_v_schedule_mem_setup` | `half` | Memory setup variant: `half`, `min`, or `zb`. |
+| `enable_1f1b_v` | `false` | 1F1B-V schedule variant. |
+| `allow_padding_num_layers` | `true` | Allow PP layer padding for divisibility. |
+| `profile_memory_iter` | `-1` | Iteration to profile memory (`-1` disables). |
+| `interleave_group_size` | `0` | Interleaved PP group size. |
+| `offload_chunk_num` | `0` | Activation offload chunk count. |
+| `offload_time` | `1.0` | Time budget for offload (scheduler hint). |
+| `auto_offload_time` | `true` | Auto-tune offload timing. |
+| `offload_overlap_sr` | `true` | Overlap save/resume in offload path. |
+| `num_seq_splits` | `1` | Splits along sequence dimension for ZB. |
+| `cpu_offload` | `false` | CPU offload of activations in ZB path. |
+
+### 11.4 Primus pipeline (`primus_pipeline.yaml`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `patch_primus_pipeline` | `false` | Enable Primus pipeline scheduling patches. |
+| `pp_algorithm` | `"1f1b-interleaved"` | Schedule name (`1f1b`, `1f1b-interleaved`, `zero-bubble`, `zero-bubble-heuristic`, `zbv-formatted`, `v-half`, `v-min`). |
+| `communication_method` | `"async_p2p"` | `async_p2p` or `batch_p2p` PP transfers. |
+| `offload` | `false` | Generic PP activation offload toggle in Primus pipeline. |
+| `offload_ops` | `""` | Comma-separated offload targets (`attn` today; other ops listed in-file are not supported yet). |
+| `pp_max_mem` | `null` | `zero-bubble-heuristic` only: max activation memory per stage (`null` = unlimited). |
+| `pp_cost_f` | `null` | `zero-bubble-heuristic` only: forward cost per stage (scalar or list; `null` = default 1000). |
+| `pp_cost_b` | `null` | `zero-bubble-heuristic` only: backward cost per stage (scalar or list; `null` = default 1000). |
+| `pp_cost_w` | `null` | `zero-bubble-heuristic` only: weight-grad cost per stage (scalar or list; `null` = default 1000). |
+
+`pp_warmup` and `dump_pp_data` are *Primus* helpers defined in `primus_megatron_module.yaml` (not `primus_pipeline.yaml`):
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `pp_warmup` | `false` | *Primus:* warm-up PP stages to reduce first-iteration latency. |
+| `dump_pp_data` | `false` | *Primus:* dump PP tensors for debugging. |
+
+---
+
+## 12. Reinforcement learning and GRPO-related settings
+
+*Source: `trainer_base.yaml`. Names follow Megatron’s `grpo_*` / `rl_*` prefixes (there is no `rl_grpo` single flag in these presets).*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `perform_rl_step` | `false` | Run RL / preference optimization steps (GRPO / LangRL integration). |
+| `rl_prompts_per_eval` | `32` | Prompts per RL evaluation pass. |
+| `grpo_prompts_per_step` | `32` | GRPO prompts sampled per training step. |
+| `grpo_group_size` | `2` | Samples per prompt group for GRPO. |
+| `grpo_iterations` | `2` | Inner GRPO iterations. |
+| `grpo_clamp_eps_lower` | `0.01` | PPO-style lower clip epsilon. |
+| `grpo_clamp_eps_upper` | `0.01` | Upper clip epsilon. |
+| `grpo_kl_beta` | `0.001` | KL penalty weight toward reference policy. |
+| `grpo_entropy_term_weight` | `0.0` | Entropy bonus weight. |
+| `grpo_filter_groups_with_same_reward` | `false` | Drop groups with identical rewards. |
+| `grpo_default_temperature` | `1.0` | Default softmax temperature for rollouts. |
+| `grpo_default_top_p` | `0` | Top-p sampling (`0` often means disabled / greedy—see Megatron RL docs). |
+| `langrl_inference_server_type` | `inplace_megatron` | LangRL inference backend. |
+| `langrl_inference_server_conversation_template` | `null` | Conversation template path / name. |
+| `langrl_env_config` | `null` | Environment / task YAML for LangRL. |
+| `rl_offload_optimizer_during_inference` | `false` | Offload optimizer to CPU during rollout inference. |
+| `rl_offload_kv_cache_during_training` | `false` | Offload KV cache while training forward runs. |
+| `rl_remove_kv_cache_during_training` | `false` | Drop KV cache between RL phases to save memory. |
+| `rl_reset_cuda_graphs` | `false` | Reset CUDA graphs when switching RL phases. |
+| `rl_partial_rollouts` | `false` | Partial sequence rollouts. |
+| `rl_inference_logprobs_is_correction` | `false` | Interpret inference logprobs as IS correction term. |
+| `rl_importance_sampling_truncation_coef` | `null` | Truncate importance ratios at this value. |
+| `rl_calculate_intra_group_similarity` | `false` | Log similarity within GRPO groups. |
+
+---
+
+## 13. Additional specialized parameters
+
+*Source: `trainer_base.yaml` (remaining domains).*
+
+### 13.1 Vision pretraining
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `vision_pretraining` | `false` | Enable vision backbone pretraining. |
+| `vision_pretraining_type` | `classify` | Objective (`classify`, etc.). |
+| `vision_backbone_type` | `vit` | Vision backbone family. |
+| `swin_backbone_type` | `tiny` | Swin variant size. |
+| `num_classes` | `1000` | Classification classes. |
+| `img_h` | `224` | Image height. |
+| `img_w` | `224` | Image width. |
+| `num_channels` | `3` | Input channels. |
+| `patch_dim` | `16` | ViT patch size. |
+| `classes_fraction` | `1.0` | Fraction of classes used. |
+| `data_per_class_fraction` | `1.0` | Fraction of data per class. |
+
+### 13.2 RETRO
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `retro_project_dir` | `null` | RETRO project directory with indices. |
+| `retro_add_retriever` | `false` | Add frozen retriever tower. |
+| `retro_cyclic_train_iters` | `null` | Cyclic iterator length. |
+| `retro_encoder_layers` | `2` | Retriever encoder layers. |
+| `retro_encoder_hidden_dropout` | `0.1` | Retriever dropout. |
+| `retro_encoder_attention_dropout` | `0.1` | Retriever attention dropout. |
+| `retro_num_neighbors` | `2` | Neighbors per query chunk. |
+| `retro_num_retrieved_chunks` | `2` | Chunks concatenated per neighbor set. |
+| `retro_attention_gate` | `1` | Gating between retrieval and LM. |
+| `retro_verify_neighbor_count` | `true` | Assert neighbor counts for debugging. |
+
+### 13.3 DINO self-supervised
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `dino_local_img_size` | `96` | Local crop size. |
+| `dino_local_crops_number` | `10` | Number of local crops. |
+| `dino_head_hidden_size` | `2048` | Projection head width. |
+| `dino_bottleneck_size` | `256` | Bottleneck dimension. |
+| `dino_freeze_last_layer` | `1` | Freeze last layer epochs. |
+| `dino_norm_last_layer` | `false` | Normalize last layer weights. |
+| `dino_warmup_teacher_temp` | `0.04` | Teacher temperature warmup start. |
+| `dino_teacher_temp` | `0.07` | Teacher temperature. |
+| `dino_warmup_teacher_temp_epochs` | `30` | Epochs to warm teacher temperature. |
+
+### 13.4 Biencoder / ICT / retriever utilities
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `ict_head_size` | `null` | ICT projection head width. |
+| `biencoder_projection_dim` | `0` | Biencoder shared projection dimension. |
+| `biencoder_shared_query_context_model` | `false` | Share query/context encoders. |
+| `ict_load` | `null` | ICT checkpoint path. |
+| `bert_load` | `null` | BERT encoder checkpoint for biencoder. |
+| `titles_data_path` | `null` | Titles file for ICT datasets. |
+| `query_in_block_prob` | `0.1` | Probability of in-block queries. |
+| `use_one_sent_docs` | `false` | Single-sentence pseudo documents. |
+| `evidence_data_path` | `null` | Evidence passages for open-domain QA. |
+| `retriever_report_topk_accuracies` | `[]` | k values for top-k accuracy logging. |
+| `retriever_score_scaling` | `false` | Scale retriever scores. |
+| `block_data_path` | `null` | Block JSON data for retrieval. |
+| `embedding_path` | `null` | Precomputed embeddings path. |
+| `indexer_batch_size` | `128` | Batch size when building ANN index. |
+| `indexer_log_interval` | `1000` | Indexer progress log interval. |
+
+### 13.5 Straggler detection
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `log_straggler` | `false` | Log straggler diagnostics. |
+| `disable_straggler_on_startup` | `false` | Skip straggler detection at startup. |
+| `straggler_ctrlr_port` | `65535` | Controller port for straggler service. |
+| `straggler_minmax_count` | `1` | Min/max samples for straggler stats. |
+
+### 13.6 Inference-oriented options
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `inference_batch_times_seqlen_threshold` | `-1` | Heuristic threshold tying batch and sequence length. |
+| `inference_dynamic_batching` | `false` | Dynamic batching for inference server. |
+| `inference_dynamic_batching_buffer_size_gb` | `40.0` | GPU buffer budget (GB). |
+| `inference_dynamic_batching_buffer_guaranteed_fraction` | `0.2` | Minimum reserved fraction of buffer. |
+| `inference_dynamic_batching_buffer_overflow_factor` | `null` | Overflow growth factor. |
+| `inference_dynamic_batching_max_requests_override` | `null` | Hard cap on concurrent requests. |
+| `inference_dynamic_batching_max_tokens_override` | `null` | Hard cap on tokens in flight. |
+| `max_tokens_to_oom` | `12000` | Token limit guard before OOM abort. |
+| `output_bert_embeddings` | `false` | Return BERT pooled embeddings. |
+| `bert_embedder_type` | `megatron` | `megatron` or `huggingface` embedder. |
+| `flash_decode` | `false` | Flash decode kernels for incremental generation. |
+| `enable_cuda_graph` | `false` | Capture CUDA graphs for inference. |
+| `cuda_graph_warmup_steps` | `3` | Warm-up steps before capturing graphs. |
+| `external_cuda_graph` | `false` | External graph provider hooks. |
+| `cuda_graph_scope` | `full` | Graph scope (`full` or `attn`). |
+| `inference_max_requests` | `8` | Max concurrent requests. |
+| `inference_max_seq_length` | `2560` | Max prefill + decode tokens per request. |
+
+### 13.7 Fault tolerance package and tooling
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `enable_ft_package` | `false` | NVIDIA fault-tolerance package hooks. |
+| `calc_ft_timeouts` | `false` | Auto-calculate FT timeouts. |
+| `run_workload_inspector_server` | `false` | Run workload inspector sidecar. |
+
+### 13.8 Heterogeneous layers and process resilience
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `heterogeneous_layers_config_path` | `null` | JSON describing variable layer widths/types per layer. |
+| `heterogeneous_layers_config_encoded_json` | `null` | Inline base64/JSON blob for heterogeneous layers. |
+| `inprocess_restart` | `false` | In-process restart for fault recovery experiments. |
+
+### 13.9 Experimental and rerun controls
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `enable_experimental` | `false` | Gate experimental Megatron features. |
+| `error_injection_rate` | `0` | Fraction of iterations with injected errors (testing). |
+| `error_injection_type` | `transient_error` | `correct_result`, `transient_error`, or `persistent_error`. |
+| `rerun_mode` | `disabled` | `disabled`, `validate_results`, or `report_stats` for rerun harness. |
+
+---
+
+### Related documentation
+
+- Megatron-LM argument definitions: [`megatron/training/arguments.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py)
+- Primus Megatron presets: `primus/configs/modules/megatron/`
+- Primus Megatron model presets: `primus/configs/models/megatron/`
diff --git a/docs/03-configuration-reference/torchtitan-parameters.md b/docs/03-configuration-reference/torchtitan-parameters.md
new file mode 100644
index 000000000..c838fa84a
--- /dev/null
+++ b/docs/03-configuration-reference/torchtitan-parameters.md
@@ -0,0 +1,370 @@
+# TorchTitan backend configuration reference
+
+This page lists Primus preset keys and common TorchTitan `JobConfig` fields used when `framework: torchtitan`. Defaults are taken from the TorchTitan module preset (`pre_trainer.yaml`), its `extends` chain (`module_base.yaml`, `quantize.yaml`), and the example model preset `llama3_8B.yaml`. It is not a complete upstream TorchTitan `JobConfig` reference.
+
+**Where parameters live.** Provide overrides under `modules.pre_trainer.overrides:` in your experiment YAML. TorchTitan’s `JobConfig` is hierarchical: use **dot notation** for flat overrides, or nest YAML objects under `overrides`—both are equivalent when merged.
+
+**Example (flat dot paths):**
+
+```yaml
+framework: torchtitan
+
+modules:
+  pre_trainer:
+    overrides:
+      training.steps: 20000
+      training.global_batch_size: 512
+      optimizer.lr: 0.00015
+      parallelism.tensor_parallel_degree: 2
+```
+
+**Example (nested YAML):**
+
+```yaml
+modules:
+  pre_trainer:
+    overrides:
+      training:
+        steps: 20000
+        global_batch_size: 512
+      optimizer:
+        lr: 0.00015
+      parallelism:
+        tensor_parallel_degree: 2
+```
+
+**Presets.**
+
+- Module presets: `primus/configs/modules/torchtitan/` (main entry: `pre_trainer.yaml`).
+- Model presets: `primus/configs/models/torchtitan/` (example: `llama3_8B.yaml`).
+
+**Mapping to TorchTitan.** Keys are translated into TorchTitan’s `JobConfig` via `TorchTitanJobConfigBuilder` (same nested structure as upstream TorchTitan).
+
+**Upstream reference.** TorchTitan repository and documentation: [https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan) (vendored as the `third_party/torchtitan` submodule).
+
+---
+
+## 1. Base module parameters
+
+*Source: `primus/configs/modules/module_base.yaml` (merged before TorchTitan-specific keys; `pre_trainer.yaml` does not override `trainable`).*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `trainable` | `false` | Whether this module is active in training orchestration. (TorchTitan preset inherits `false` from `module_base.yaml`.) |
+| `sink_level` | `null` | Structured logging sink level; `null` uses defaults. |
+| `file_sink_level` | `DEBUG` | Minimum level for file logging. |
+| `stderr_sink_level` | `INFO` | Minimum level for stderr logging. |
+
+---
+
+## 2. Training (`training.*`)
+
+*Source: `primus/configs/modules/torchtitan/pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `training.mock_data` | `true` | Primus preset extension: use synthetic data instead of reading `dataset_path`. |
+| `training.dataset` | `c4` | Dataset name key for TorchTitan dataset loaders. |
+| `training.dataset_path` | `null` | Filesystem or remote path to dataset assets. |
+| `training.enable_cpu_offload` | `false` | Offload optimizer or activations to CPU when supported. |
+| `training.gc_debug` | `false` | Extra garbage-collection diagnostics. |
+| `training.gc_freq` | `50` | Run Python GC every N steps when enabled. |
+| `training.global_batch_size` | `-1` | Global batch size across all ranks (`-1` often means “auto” / unset in TorchTitan). |
+| `training.local_batch_size` | `8` | Per-rank microbatch size before gradient accumulation. |
+| `training.max_norm` | `1.0` | Gradient clipping max norm (global). |
+| `training.mixed_precision_param` | `bfloat16` | Parameter dtype for mixed precision (`bfloat16`, `float16`, etc.). |
+| `training.mixed_precision_reduce` | `float32` | Dtype for reduction / gradient accumulation. |
+| `training.seq_len` | `2048` | Sequence length per sample. |
+| `training.steps` | `10000` | Total optimizer steps. |
+
+---
+
+## 3. Optimizer (`optimizer.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `optimizer.name` | `AdamW` | Optimizer class (`AdamW`, `Adam`, …). |
+| `optimizer.lr` | `0.0008` | Base learning rate. |
+| `optimizer.beta1` | `0.9` | First moment decay. |
+| `optimizer.beta2` | `0.95` | Second moment decay. |
+| `optimizer.eps` | `1.0e-08` | Numerical stability term. |
+| `optimizer.weight_decay` | `0.1` | Weight decay coefficient. |
+| `optimizer.implementation` | `fused` | Kernel implementation (`fused`, `foreach`, …). |
+| `optimizer.early_step_in_backward` | `false` | Experimental: step optimizer during backward when supported. |
+
+---
+
+## 4. Learning rate scheduler (`lr_scheduler.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `lr_scheduler.decay_ratio` | `null` | Fraction of training at the end used for decay; `null` uses framework default. |
+| `lr_scheduler.decay_type` | `linear` | LR decay curve (`linear`, `cosine`, etc.). |
+| `lr_scheduler.min_lr_factor` | `0.0` | LR floor as a fraction of base LR after decay. |
+| `lr_scheduler.warmup_steps` | `200` | Linear warmup steps before decay. |
+
+---
+
+## 5. Parallelism (`parallelism.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `parallelism.tensor_parallel_degree` | `1` | Tensor parallelism (intra-layer) degree. |
+| `parallelism.pipeline_parallel_degree` | `1` | Pipeline parallelism stages. |
+| `parallelism.pipeline_parallel_microbatch_size` | `1` | Microbatches per pipeline round. |
+| `parallelism.pipeline_parallel_schedule` | `1F1B` | Pipeline schedule name (`1F1B`, `GPipe`, …). |
+| `parallelism.pipeline_parallel_schedule_csv` | `''` | Optional CSV schedule definition. |
+| `parallelism.pipeline_parallel_layers_per_stage` | `null` | Layers per stage when auto-balanced. |
+| `parallelism.pipeline_parallel_first_stage_less_layers` | `1` | Fewer layers on first PP stage (for imbalance). |
+| `parallelism.pipeline_parallel_last_stage_less_layers` | `1` | Fewer layers on last PP stage. |
+| `parallelism.data_parallel_shard_degree` | `-1` | FSDP / shard degree (`-1` = auto). |
+| `parallelism.data_parallel_replicate_degree` | `1` | Replicated data-parallel groups. |
+| `parallelism.expert_parallel_degree` | `1` | Expert parallelism for MoE models. |
+| `parallelism.expert_tensor_parallel_degree` | `1` | Tensor parallelism inside experts. |
+| `parallelism.context_parallel_degree` | `1` | Context (sequence) parallelism degree. |
+| `parallelism.context_parallel_rotate_method` | `allgather` | Communication pattern for context parallel. |
+| `parallelism.disable_loss_parallel` | `false` | Disable loss parallel layout when TP is used. |
+| `parallelism.enable_async_tensor_parallel` | `false` | Overlap TP collectives with compute. |
+| `parallelism.fsdp_reshard_after_forward` | `default` | FSDP reshard policy (`default`, `always`, `never`). |
+| `parallelism.module_fqns_per_model_part` | `null` | Map of pipeline stage → module FQNs for multi-part models. |
+
+---
+
+## 6. Checkpoint (`checkpoint.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `checkpoint.enable` | `false` | Master switch for checkpointing. |
+| `checkpoint.folder` | `checkpoint` | Output directory for checkpoints. |
+| `checkpoint.interval` | `500` | Save every N steps. |
+| `checkpoint.initial_load_path` | `null` | Path to load from at startup. |
+| `checkpoint.initial_load_model_only` | `true` | Load weights only (skip optimizer/scheduler). |
+| `checkpoint.initial_load_in_hf` | `false` | Load initial weights from Hugging Face format. |
+| `checkpoint.last_save_model_only` | `true` | Final save stores weights only. |
+| `checkpoint.last_save_in_hf` | `false` | Export final weights in Hugging Face format. |
+| `checkpoint.export_dtype` | `float32` | Dtype for exported checkpoints. |
+| `checkpoint.async_mode` | `disabled` | Async checkpoint (`disabled`, `async`, …). |
+| `checkpoint.keep_latest_k` | `10` | Retain only the newest k checkpoints. |
+| `checkpoint.load_step` | `-1` | Step index to load (`-1` = latest). |
+| `checkpoint.exclude_from_loading` | `[]` | FQNs or keys to skip when loading. |
+| `checkpoint.enable_first_step_checkpoint` | `false` | Save checkpoint at step 0 for debugging. |
+| `checkpoint.create_seed_checkpoint` | `false` | Save a seed checkpoint before training starts. |
+
+---
+
+## 7. Activation checkpoint (`activation_checkpoint.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `activation_checkpoint.mode` | `none` | Activation checkpointing mode (`none`, `selective`, `full`). |
+| `activation_checkpoint.selective_ac_option` | `"2"` | Selective AC policy string (TorchTitan-specific). |
+| `activation_checkpoint.per_op_sac_force_recompute_mm_shapes_by_fqns` | `["moe.router.gate"]` | FQNs that always recompute matmuls in selective AC. |
+| `activation_checkpoint.early_stop` | `false` | Stop AC early in certain subgraphs. |
+
+---
+
+## 8. Metrics and profiling
+
+### 8.1 Metrics (`metrics.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `metrics.disable_color_printing` | `false` | Disable ANSI colors in logs. |
+| `metrics.enable_tensorboard` | `false` | Write TensorBoard scalars. |
+| `metrics.enable_wandb` | `false` | Log to Weights & Biases. |
+| `metrics.log_freq` | `10` | Steps between metric logs. |
+| `metrics.save_for_all_ranks` | `false` | Save metric files per rank (not just rank 0). |
+| `metrics.save_tb_folder` | `tb` | TensorBoard subdirectory / name. |
+
+### 8.2 Profiling (`profiling.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `profiling.enable_profiling` | `false` | Enable PyTorch profiler traces. |
+| `profiling.enable_memory_snapshot` | `false` | Capture CUDA memory snapshots. |
+| `profiling.profile_freq` | `10` | Steps between profiler activations. |
+| `profiling.save_traces_folder` | `profile_traces` | Directory for profiler traces. |
+| `profiling.save_memory_snapshot_folder` | `memory_snapshot` | Directory for memory snapshots. |
+
+---
+
+## 9. Quantization (`quantize.*`)
+
+*Source: `primus/configs/modules/torchtitan/quantize.yaml` (merged into the module preset).*
+
+### 9.1 Linear FP8 (`quantize.linear.float8.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `quantize.linear.float8.enable_fsdp_float8_all_gather` | `false` | FP8 all-gather for FSDP sharded params (recommended for tensorwise scaling). |
+| `quantize.linear.float8.precompute_float8_dynamic_scale_for_fsdp` | `false` | Precompute dynamic scales for FSDP FP8. |
+| `quantize.linear.float8.recipe_name` | `null` | Recipe (`tensorwise`, `rowwise`, `rowwise_with_gw_hp`); `null` disables. |
+| `quantize.linear.float8.filter_fqns` | `[]` | Module FQNs to skip for FP8 training. |
+| `quantize.linear.float8.emulate` | `false` | Emulate FP8 in FP32 (no FP8 HW); not compatible with `torch.compile`. |
+
+### 9.2 Linear MX (`quantize.linear.mx.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `quantize.linear.mx.mxfp8_dim1_cast_kernel_choice` | `"triton"` | Kernel backend for MXFP8 dim-1 cast (`triton`, `cuda`, `torch`). |
+| `quantize.linear.mx.recipe_name` | `"mxfp8_cublas"` | MX recipe name (see torchao `mx_formats`). |
+| `quantize.linear.mx.filter_fqns` | `["output"]` | FQNs to skip; output layer skipped by default. |
+
+### 9.3 Grouped GEMM FP8 (`quantize.grouped_mm.float8.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `quantize.grouped_mm.float8.fqns` | `[]` | MoE layer FQNs for FP8 grouped GEMM (prototype; may require torchao nightly). |
+
+### 9.4 Grouped GEMM MX (`quantize.grouped_mm.mx.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `quantize.grouped_mm.mx.recipe_name` | `"mxfp8"` | MX recipe for grouped GEMMs. |
+| `quantize.grouped_mm.mx.fqns` | `[]` | MoE module FQNs for MXFP8 grouped GEMM (prototype). |
+
+---
+
+## 10. Compile (`compile.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `compile.enable` | `true` | Enable `torch.compile` on selected subsystems. |
+| `compile.components` | `["model", "loss"]` | Which components to compile. |
+
+---
+
+## 11. Communication and fault tolerance
+
+### 11.1 Communicator (`comm.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `comm.init_timeout_seconds` | `300` | Timeout for initial process-group setup. |
+| `comm.train_timeout_seconds` | `100` | Timeout for training collectives. |
+| `comm.trace_buf_size` | `20000` | Flight recorder buffer size for NCCL traces. |
+| `comm.save_traces_folder` | `comm_traces` | Where to dump communication traces. |
+
+### 11.2 Fault tolerance (`fault_tolerance.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `fault_tolerance.enable` | `false` | Enable fault-tolerant training hooks. |
+| `fault_tolerance.process_group` | `gloo` | Backend for control-plane process group. |
+| `fault_tolerance.process_group_timeout_ms` | `10000` | Control-plane timeout. |
+| `fault_tolerance.replica_id` | `0` | Replica index in elastic setups. |
+| `fault_tolerance.group_size` | `0` | Group size (0 = unset / default). |
+| `fault_tolerance.min_replica_size` | `1` | Minimum replicas to continue. |
+| `fault_tolerance.semi_sync_method` | `null` | Optional semi-synchronous strategy name. |
+
+### 11.3 Memory estimation (`memory_estimation.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `memory_estimation.enable` | `false` | Run memory-estimation fake-mode passes. |
+| `memory_estimation.disable_fake_mode` | `false` | Disable fake tensor mode inside estimation. |
+
+### 11.4 Experimental (`experimental.*`)
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `experimental.custom_import` | `""` | Optional Python module import path for custom extensions. |
+| `experimental.custom_args_module` | `"primus.backends.torchtitan.primus_turbo_extensions.config_extension"` | Module providing extra `JobConfig` fields for Primus-Turbo. |
+
+---
+
+## 12. Primus-Turbo (`primus_turbo.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `primus_turbo.enable_primus_turbo` | `true` | Master switch for Primus-Turbo integrations in TorchTitan. |
+| `primus_turbo.enable_attention_float8` | `false` | FP8 attention path inside Turbo attention. |
+| `primus_turbo.use_turbo_attention` | `true` | Use Turbo attention kernels. |
+| `primus_turbo.use_classic_attention` | `false` | Fall back to classic attention implementation. |
+| `primus_turbo.use_turbo_async_tp` | `true` | Async tensor-parallel communication in Turbo. |
+| `primus_turbo.use_turbo_mx_linear` | `true` | MX linear layers via Turbo. |
+| `primus_turbo.use_turbo_float8_linear` | `true` | FP8 linear layers via Turbo. |
+| `primus_turbo.use_turbo_grouped_mm` | `false` | Turbo grouped GEMM for MoE (off by default). |
+| `primus_turbo.use_moe_fp8` | `true` | FP8 paths for MoE experts when applicable. |
+| `primus_turbo.enable_embedding_autocast` | `true` | Autocast policy around embeddings for Turbo. |
+
+---
+
+## 13. Model (`model.*` and `job.*`)
+
+### 13.1 Model preset (`models.*` / `model.*`)
+
+*Example defaults from `primus/configs/models/torchtitan/llama3_8B.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `model.name` | `"llama3"` | Model family key for TorchTitan recipes. |
+| `model.flavor` | `"8B"` | Size / variant within the family. |
+| `model.hf_assets_path` | `"meta-llama/Meta-Llama-3-8B"` | Hugging Face Hub repo or local path for weights/tokenizer. |
+| `model.converters` | `["primus_turbo"]` | Weight converter pipeline stages applied at load. |
+
+### 13.2 Job metadata (`job.*`)
+
+*Source: `llama3_8B.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `job.dump_folder` | `"./outputs"` | Root directory for logs, checkpoints, and exports. |
+| `job.description` | `"Llama 3 8B training"` | Human-readable label for run metadata. |
+
+---
+
+## 14. Validation (`validation.*`)
+
+*Source: `pre_trainer.yaml`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `validation.enable` | `false` | Run periodic validation loops. |
+| `validation.dataset` | `c4_validation` | Validation dataset key. |
+| `validation.dataset_path` | `null` | Filesystem path to validation data. |
+| `validation.local_batch_size` | `8` | Per-rank validation batch size. |
+| `validation.seq_len` | `2048` | Validation sequence length. |
+| `validation.freq` | `10` | Run validation every N training steps. |
+| `validation.steps` | `-1` | Max validation steps (`-1` = full pass / framework default). |
+
+---
+
+## 15. Debug (`debug.*`)
+
+*Source: `pre_trainer.yaml`. Upstream added this section in TorchTitan v0.2.2; the keys previously lived under `training.*`.*
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `debug.deterministic` | `false` | Prefer deterministic algorithms (often slower). |
+| `debug.deterministic_warn_only` | `false` | Warn instead of erroring when an op has no deterministic implementation. |
+| `debug.moe_force_load_balance` | `false` | Round-robin tokens across MoE experts so every expert gets the same amount; debugging only. |
+| `debug.seed` | `null` | RNG seed; `null` lets the framework choose. |
+
+---
+
+### Related documentation
+
+- TorchTitan repository and documentation: [https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan)
+- Primus TorchTitan presets: `primus/configs/modules/torchtitan/`
+- Primus TorchTitan model presets: `primus/configs/models/torchtitan/`
diff --git a/docs/04-technical-guides/README.md b/docs/04-technical-guides/README.md
new file mode 100644
index 000000000..2f45c99c3
--- /dev/null
+++ b/docs/04-technical-guides/README.md
@@ -0,0 +1,24 @@
+# Technical guides
+
+Deep technical topics for advanced users.
+
+- [Parallelism strategies](parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained
+- [Parallelism configuration](parallelism-configuration.md): per-backend parallelism setup and batch size relationships
+- [Collective operations](collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy
+- [Performance tuning](performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization
+- [MoE training deep-dive](moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models
+- [MegaMoE fused MoE layer](mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction
+- [Data preparation](data-preparation.md): tokenization, data formats, mock data
+- [Checkpoint management](checkpoint-management.md): formats, save/load, distributed checkpointing
+- [Multi-node networking](multi-node-networking.md): InfiniBand, RoCE, AINIC configuration
+- [Profiling and observability](profiling-and-observability.md): Torch profiler, TraceLens, memory snapshots, projection, pp_vis
+- [Logging and experiment tracking](logging-and-experiment-tracking.md): TensorBoard, WandB, MLflow setup per backend
+- [Fault tolerance and elastic training](fault-tolerance-and-elastic-training.md): graceful exit, auto-resume, in-process restart, torchft
+- [Determinism and reproducibility](determinism-and-reproducibility.md): deterministic mode, seeds, trade-offs
+- [Diffusion models](diffusion-models/README.md): Flux diffusion architecture, data pipeline, and FP8 / MXFP4 training
+- [Hybrid models](hybrid-models/README.md): Zebra-Llama hybrid recurrent-attention (Mamba/KDA/GDN + MLA) models, FLA-parity recipes, and checkpoint conversion
+- [Native SFT and LoRA](native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/04-technical-guides/checkpoint-management.md b/docs/04-technical-guides/checkpoint-management.md
new file mode 100644
index 000000000..9e6f29cbc
--- /dev/null
+++ b/docs/04-technical-guides/checkpoint-management.md
@@ -0,0 +1,203 @@
+# Checkpoint management
+
+Checkpoints capture **model state**, **optimizer state**, and **training progress** (iteration or step counters, schedulers, and related metadata). They are essential for **fault tolerance** (resume after failure), **experiment management** (reproducibility and comparison), and **hand-offs** between pretraining, fine-tuning, and conversion workflows.
+
+Primus is YAML-driven: checkpoint behavior is configured per backend. **Megatron-LM**, **TorchTitan**, and **MaxText** each expose their own checkpoint surfaces; this guide maps the knobs you set in Primus configs.
+
+**Primary sources in this repository**
+
+| Area | File |
+|------|------|
+| Megatron trainer defaults | `primus/configs/modules/megatron/trainer_base.yaml` |
+| Primus Megatron extensions | `primus/configs/modules/megatron/primus_megatron_module.yaml` |
+| TorchTitan defaults | `primus/configs/modules/torchtitan/pre_trainer.yaml` |
+| Megatron checkpoint benchmark | `benchmark/megatron/checkpoint/README.md` |
+
+---
+
+## 1. Overview
+
+- **What is saved:** Typically model parameters, optimizer state, RNG state, and iteration/step tracking—exact contents depend on flags such as `no_save_optim` / `no_save_rng` (Megatron) or `initial_load_model_only` (TorchTitan).
+- **Why it matters:** Long runs on AMD GPU clusters benefit from periodic saves to durable storage; resuming or branching experiments requires consistent paths and formats.
+- **Backend-specific systems:** Each training backend integrates its own checkpoint pipeline; Primus wires YAML into those backends without forcing a single universal format across Megatron, TorchTitan, and MaxText.
+
+---
+
+## 2. Megatron checkpoint configuration
+
+Megatron-related options live on the trainer configuration merged from `trainer_base.yaml` and `primus_megatron_module.yaml`. Defaults below are taken from `trainer_base.yaml` unless noted.
+
+### Core paths and cadence
+
+| Parameter | Default (`trainer_base.yaml`) | Description |
+|-----------|-------------------------------|-------------|
+| `save` | `null` | Directory where new checkpoints are written. |
+| `load` | `null` | Directory to load from when **resuming** training. |
+| `save_interval` | `20000` | Save every *N* iterations. |
+| `finetune` | `false` | When `true`, loads weights but **resets** the iteration counter (typical fine-tune entry). |
+
+### Format and detection
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `ckpt_format` | `torch_dist` | Checkpoint format: `torch` (legacy single-file style), `torch_dist` (distributed), or `zarr`. |
+| `auto_detect_ckpt_format` | `false` | When loading, infer format automatically. |
+| `pretrained_checkpoint` | `null` | Path to a **pretrained** checkpoint. |
+| `ckpt_step` | `null` | Load a specific step from the pretrained checkpoint when applicable. |
+
+### Optimizer and RNG inclusion
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `no_save_optim` | `null` | When set truthy, **omit** optimizer state from saves. |
+| `no_save_rng` | `null` | When set truthy, **omit** RNG state from saves. |
+| `no_load_optim` | `null` | When set truthy, **do not** restore optimizer from checkpoint. |
+| `no_load_rng` | `null` | When set truthy, **do not** restore RNG from checkpoint. |
+
+### Performance and distributed I/O
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `async_save` | `null` | Asynchronous checkpoint saving to reduce time blocking the training loop. |
+| `ckpt_fully_parallel_save` | `true` | Parallel save path for distributed checkpoints. |
+| `ckpt_fully_parallel_load` | `false` | Parallel load for distributed checkpoints. |
+| `ckpt_assume_constant_structure` | `false` | Optimization when model structure is fixed across saves/loads. |
+| `non_persistent_save_interval` | `null` | Save to **fast local** storage on a different cadence than persistent saves. |
+
+Related keys in `trainer_base.yaml` for non-persistent checkpoints include `non_persistent_ckpt_type`, `non_persistent_global_ckpt_dir`, `non_persistent_local_ckpt_dir`, and `non_persistent_local_ckpt_algo` (default `"fully_parallel"`).
+
+### Format conversion
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `ckpt_convert_format` | `null` | Target format for conversion (`torch`, `torch_dist`, or `zarr`). |
+| `ckpt_convert_save` | `null` | Output directory for converted checkpoints. |
+| `ckpt_convert_update_legacy_dist_opt_format` | `false` | Update legacy distributed optimizer layout when converting. |
+
+### Primus extensions
+
+Defined in `primus/configs/modules/megatron/primus_megatron_module.yaml` and implemented in `primus/backends/megatron/patches/checkpoint_patches.py` and `primus/backends/megatron/patches/args/checkpoint_path_patches.py`.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `auto_continue_train` | `false` | When `true`, **automatically resume** from the latest checkpoint under `save` (adjusts load/finetune and related flags). |
+| `disable_last_saving` | `false` | When `true`, **skip** the final checkpoint at shutdown (useful for benchmarking or when only periodic saves matter). |
+
+---
+
+## 3. TorchTitan checkpoint configuration
+
+TorchTitan checkpoint options are grouped under `checkpoint` in `primus/configs/modules/torchtitan/pre_trainer.yaml`.
+
+| Parameter | Default (`pre_trainer.yaml`) | Description |
+|-----------|------------------------------|-------------|
+| `checkpoint.enable` | `false` | Master switch for checkpointing. |
+| `checkpoint.folder` | `checkpoint` | Output directory (relative to run layout unless given as absolute). |
+| `checkpoint.interval` | `500` | Save every *N* **steps**. |
+| `checkpoint.initial_load_path` | `null` | Path for **initial** load (cold start or migration). |
+| `checkpoint.initial_load_model_only` | `true` | Load **weights only**, not optimizer state. |
+| `checkpoint.initial_load_in_hf` | `false` | Load initial weights from **Hugging Face** layout. |
+| `checkpoint.last_save_model_only` | `true` | On last save, write **model only**. |
+| `checkpoint.last_save_in_hf` | `false` | Write final checkpoint in **Hugging Face** format. |
+| `checkpoint.export_dtype` | `float32` | Dtype for exported checkpoints. |
+| `checkpoint.async_mode` | `disabled` | Asynchronous checkpoint mode. |
+| `checkpoint.keep_latest_k` | `10` | Retain only the **K** most recent checkpoints. |
+| `checkpoint.load_step` | `-1` | Load a specific step (`-1` typically means latest or default behavior per backend). |
+| `checkpoint.exclude_from_loading` | `[]` | Glob or pattern list to **exclude** from restore. |
+| `checkpoint.enable_first_step_checkpoint` | `false` | Optional checkpoint at step 0. |
+| `checkpoint.create_seed_checkpoint` | `false` | Create a seed checkpoint when enabled. |
+
+TorchTitan also defines `activation_checkpoint` (activation recomputation) separately from persistent training checkpoints—do not confuse the two sections in `pre_trainer.yaml`.
+
+---
+
+## 4. MaxText checkpoint configuration
+
+MaxText (JAX) uses configuration keys surfaced in Primus documentation and MaxText configs under `third_party/maxtext`. Primus overlay presets set the defaults shown below, while upstream MaxText `base.yml` is still loaded at runtime via `base_config: "base.yml"` and might define different upstream defaults. Typical training flags:
+
+| Parameter | Typical default | Description |
+|-----------|-----------------|-------------|
+| `enable_checkpointing` | `false` | Enable Orbax (or configured) checkpoint saves. |
+| `async_checkpointing` | `false` | When checkpointing is enabled, use **async** checkpoint workers. |
+
+See `docs/03-configuration-reference/maxtext-parameters.md` for the full MaxText parameter table and interaction with training runs.
+
+---
+
+## 5. Checkpoint formats (Megatron)
+
+| Format | Behavior | Notes |
+|--------|----------|--------|
+| `torch` | Classic PyTorch save/load; often **one file per rank** in distributed settings. | Simple but less flexible for topology changes. |
+| `torch_dist` | **Distributed** checkpoint format with **resharding** support (e.g., changing tensor/pipeline parallel degree between save and load). | **Recommended** for many production flows that might change parallelism. |
+| `zarr` | Zarr-backed checkpoint storage. | Useful when the stack and storage backend support it. |
+
+**Recommendation:** Prefer `torch_dist` for production when you need **flexibility across parallel layouts** and scalable I/O (see `ckpt_fully_parallel_save` / `ckpt_fully_parallel_load` in Megatron config).
+
+---
+
+## 6. Common workflows
+
+**Resume training**
+
+- Set `load` to the checkpoint directory produced by a previous run.
+- Keep `save` pointed at the directory for **new** checkpoints (often the same tree with a new run id, depending on your layout).
+- Ensure `finetune` is `false` when you want to **continue** iteration counts.
+
+**Fine-tune from a pretrained checkpoint**
+
+- Set `load` (and optionally `pretrained_checkpoint` / `ckpt_step` as appropriate).
+- Set `finetune: true` so iteration counters reset while weights load.
+
+**Auto-resume (Primus Megatron extension)**
+
+- Set `auto_continue_train: true` in the Megatron module config.
+- Primus searches for the latest checkpoint under `save` and aligns load/optimizer flags; see `primus/backends/megatron/patches/checkpoint_patches.py` for behavior details.
+
+**Convert checkpoint format**
+
+- Set `ckpt_convert_format` (for example `torch_dist`) and `ckpt_convert_save` to the output directory.
+
+**Import Hugging Face weights (TorchTitan)**
+
+- Set `checkpoint.initial_load_in_hf: true` and `checkpoint.initial_load_path` to the HF model directory.
+
+---
+
+## 7. Benchmarking checkpoints
+
+The Megatron checkpoint benchmark lives in `benchmark/megatron/checkpoint/`.
+
+**Entry points**
+
+- `benchmark/megatron/checkpoint/ckpt_launch.py`—main launcher (requires a Primus YAML config).
+- `benchmark/megatron/checkpoint/ckpt_report.py`—reporting utility (can be run separately).
+
+**Example** (from `benchmark/megatron/checkpoint/README.md`):
+
+```bash
+export DATA_PATH=/PATH/TO/DATA
+python3 benchmark/megatron/checkpoint/ckpt_launch.py \
+    --yaml-config-path examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-pretrain.yaml \
+    --nnodes 1
+```
+
+The tool reports save/load times, bandwidth, and configuration echoes (world size, `ckpt_format`, `async_save`, paths, and more). Truncate or clean leftover output directories between runs if permissions or stale outputs cause issues.
+
+---
+
+## 8. Best practices
+
+- Enable **`async_save`** (Megatron) for large models when supported, to limit training stalls during checkpoint windows.
+- Set **`save_interval`** from **economic** criteria: frequent enough to limit lost work, infrequent enough to avoid storage and throughput bottlenecks (Megatron default in `trainer_base.yaml` is `20000`—override per job).
+- Use **`non_persistent_save_interval`** with fast **local SSD** for frequent snapshots and a slower interval to **NFS** or object storage for durability.
+- **Validate** resume and fine-tune paths on short runs before multi-week jobs; confirm `finetune` and `auto_continue_train` behave as intended.
+- For TorchTitan, enable **`checkpoint.enable`** explicitly and set **`checkpoint.keep_latest_k`** to bound disk usage.
+
+---
+
+## Related documentation
+
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md)
+- [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)
+- [Benchmark suite](../02-user-guide/benchmarking.md)
diff --git a/docs/04-technical-guides/collective-operations.md b/docs/04-technical-guides/collective-operations.md
new file mode 100644
index 000000000..b048ecd98
--- /dev/null
+++ b/docs/04-technical-guides/collective-operations.md
@@ -0,0 +1,289 @@
+# NCCL/RCCL collective operations guide
+
+Distributed training spends a large fraction of wall time in **collective communication**: many GPUs must exchange gradients, parameters, or activations in coordinated patterns. On AMD GPUs, **RCCL** (ROCm Collective Communications Library) provides these operations with an API aligned to **NCCL** (NVIDIA Collective Communications Library), so most concepts and environment variables carry over between vendors.
+
+This guide explains core collectives, how they map to parallelism strategies in Primus (Megatron-LM, TorchTitan), and how to benchmark and troubleshoot communication.
+
+For Megatron knobs like `overlap_grad_reduce` and TorchTitan parallelism flags, see [Megatron parameters](../03-configuration-reference/megatron-parameters.md). For `NCCL_*` / `RCCL_*` environment variables, see [Environment variables](../03-configuration-reference/environment-variables.md).
+
+---
+
+## 1. Introduction
+
+### What are collective operations?
+
+A **collective** is a multi-party communication pattern where **every participant** (or a defined **process group**) follows the same operation: combine tensors, broadcast, scatter pieces, or exchange shards. Unlike a single **Send/Recv** pair, collectives are **synchronized** by construction and are implemented with optimized algorithms (ring, tree, etc.).
+
+### NCCL vs RCCL
+
+| | NCCL | RCCL |
+|---|------|------|
+| Vendor | NVIDIA CUDA | AMD ROCm |
+| Role | GPU collective communication | GPU collective communication |
+| Typical API surface | C/C++ and bindings used by PyTorch distributed | ROCm stack; PyTorch uses similar backends |
+
+Application code written for **PyTorch distributed** (e.g. `torch.distributed`) generally selects the backend provided by the stack (**nccl** on NVIDIA, **rccl** on AMD). **Operation names and semantics** (AllReduce, AllGather, …) align so that **framework-level** code and tuning guides are largely **portable**.
+
+### Process groups
+
+Not every rank talks to every other rank in every step. **Process groups** define **subsets of ranks** that participate in a collective (e.g. only **tensor-parallel** ranks, only **data-parallel** ranks). Correctness and performance depend on **matching ranks** to the same group for each layer or phase of training.
+
+---
+
+## 2. Core collective operations
+
+Below, \(n\) is the number of ranks in the process group, and \(S\) is the size of the logical tensor being reduced or moved (per-rank message size in ring formulations). **Complexity** expressions are **standard ring-style** approximations for **amount of data moved per rank** relative to \(S\); real implementations pick algorithms based on message size, topology, and environment.
+
+---
+
+### AllReduce
+
+**What it does:** Each rank contributes a tensor; the **element-wise reduction** (typically **sum**) is applied across ranks, and the **full result** is **replicated** on every rank.
+
+```
+Rank 0: [a0]     Rank 1: [a1]     Rank 2: [a2]
+        \           |           /
+         \          |          /
+          --> REDUCE(sum) <--
+                    |
+        All ranks: [a0+a1+a2]
+```
+
+**Complexity (ring, per-rank data moved):** about \(\frac{2(n-1)}{n} S\).
+
+**Where used:** **Data-parallel** gradient synchronization; **tensor-parallel** partial sums; any step that needs **identical** tensors on all ranks after a reduction.
+
+---
+
+### AllGather
+
+**What it does:** Each rank holds **one shard**; every rank receives the **concatenation** (or stacked layout) of **all shards**.
+
+```
+Rank 0: [x0]  Rank 1: [x1]  Rank 2: [x2]
+          \       |       /
+           \      |      /
+            --> ALL GATHER -->
+Each rank: [x0 | x1 | x2]
+```
+
+**Complexity (ring):** about \(\frac{n-1}{n} S\) **if** each rank contributes \(S/n\); more generally scales with gathering \(n-1\) other shards of comparable size.
+
+**Where used:** **FSDP / ZeRO-3** parameter gather before forward; **TP** weight or activation assembly depending on layout.
+
+---
+
+### ReduceScatter
+
+**What it does:** Conceptually **AllReduce** then **split**: each rank ends with **one shard** of the reduced result (each rank’s shard is the reduction over corresponding positions from all ranks’ inputs).
+
+```
+Inputs per rank: full-sized chunks (partial sums local)
+        |
+   Reduce + partition
+        |
+Rank i gets shard i of the fully reduced tensor
+```
+
+**Complexity (ring):** about \(\frac{n-1}{n} S\) for the common balanced case.
+
+**Where used:** **FSDP** gradient **sharding** after backward; **sequence parallelism** with TP (activation distribution); distributed optimizer flows that **scatter** reduced pieces.
+
+---
+
+### AllToAll
+
+**What it does:** Each rank sends a **distinct slice** to every other rank; every rank receives from every rank (matrix transpose of data ownership).
+
+```
+        From rank 0..n-1
+              |
+    +---------+---------+
+    | scatter per dest |
+    v         v         v
+Each rank receives its column/row of the logical matrix
+```
+
+**Where used:** **Expert parallelism** token **dispatch** and **combine** in MoE; some **sparsity** and **parallel embedding** layouts.
+
+---
+
+### Broadcast
+
+**What it does:** One **root** rank’s tensor is copied to all other ranks.
+
+```
+Root: [w] ----copy----> all other ranks: [w]
+```
+
+**Where used:** **Weight initialization**, **loading checkpoints** to a group, distributing hyperparameters or small metadata.
+
+---
+
+### Reduce
+
+**What it does:** Like AllReduce, but the **full result** appears on **one root** rank only.
+
+**Where used:** **Logging** or **metrics** where only rank 0 needs the scalar (e.g. reduced loss on one process).
+
+---
+
+### Send / Recv (point-to-point)
+
+**What it does:** **One** rank sends a buffer to **one** other rank (possibly bidirectional with two ops).
+
+```
+Stage i ----Send/Recv----> Stage i+1
+```
+
+**Where used:** **Pipeline parallelism** activations in forward, gradients in backward; **ring attention** steps in **context parallelism** (often implemented as a ring of Send/Recv with careful ordering).
+
+---
+
+## 3. Which collectives are used in each parallelism strategy
+
+| Parallelism | Forward Pass | Backward Pass | Optimizer Step |
+|---------------|--------------|----------------|----------------|
+| Data Parallel | — | AllReduce (gradients) | — |
+| FSDP/ZeRO-3 | AllGather (params) | ReduceScatter (grads) + AllGather (params) | — |
+| Tensor Parallel | AllReduce or AllGather+ReduceScatter | AllReduce or AllGather+ReduceScatter | — |
+| Sequence Parallel | AllGather (activations) | ReduceScatter (activations) | — |
+| Pipeline Parallel | Send/Recv (activations) | Send/Recv (gradients) | — |
+| Expert Parallel | AllToAll (token dispatch) | AllToAll (gradient dispatch) | — |
+| Context Parallel | Ring Send/Recv (KV chunks) | Ring Send/Recv | — |
+
+Exact fusion and overlap depend on the backend (Megatron vs TorchTitan) and flags such as async TP or overlapped gradient reduction.
+
+---
+
+## 4. Communication patterns in Megatron-LM
+
+Primus trains with **Megatron-LM** patches and configurations. Typical patterns:
+
+| Mode | Pattern |
+|------|---------|
+| **TP** | Column-parallel and row-parallel **linear** layers use **AllReduce** or **ReduceScatter/AllGather** sequences; with **sequence_parallel**, activations follow the Megatron **scatter/gather** pattern around TP regions. |
+| **PP** | **Point-to-point** Send/Recv (or backend equivalents) between **pipeline stages** for activations and backward tensors. |
+| **DP** | **AllReduce** for gradients when not using distributed optimizer; with **distributed optimizer**, **ReduceScatter**-style paths for shard-sized gradients. |
+| **EP** | **AllToAll** for MoE **routing** (dispatch/combine) when experts are parallelized. |
+
+### Overlap knobs
+
+Megatron integrates **communication/compute overlap** options such as:
+
+- `overlap_grad_reduce`—overlap gradient reduction with computation where supported.
+- `overlap_param_gather`—overlap parameter gathering (e.g. with distributed optimizer / FSDP-style paths) with computation.
+
+See [Megatron parameters](../03-configuration-reference/megatron-parameters.md) for defaults and compatibility with `use_distributed_optimizer`, `use_torch_fsdp2`, and checkpoint formats.
+
+---
+
+## 5. Communication patterns in TorchTitan
+
+TorchTitan (used as a backend in Primus) relies on **PyTorch** distributed primitives and **DTensor**-style layouts:
+
+| Area | Pattern |
+|------|---------|
+| **FSDP / sharding** | **AllGather** / **ReduceScatter** orchestrated by **FSDP2** (`fully_shard` and related APIs) when `data_parallel_shard_degree` and sharding are enabled. |
+| **TP** | Tensor parallelism is integrated with **DTensor** and model parallel helpers; schedules may use **collectives** inside module forward/backward. |
+| **PP** | **Pipeline schedules** (`parallelism.pipeline_parallel_schedule`, `parallelism.pipeline_parallel_degree`) determine stage boundaries and buffering; communication is managed by the pipeline implementation. |
+
+### Async tensor parallelism
+
+Set `parallelism.enable_async_tensor_parallel: true` (where supported) to **overlap** TP communication with computation in eligible layers.
+
+---
+
+## 6. RCCL-specific features and tuning
+
+The following appear in ROCm / AMD deployments and partner integrations; availability depends on your **driver**, **RCCL build**, and **network** stack.
+
+| Feature | Notes |
+|---------|--------|
+| **MSCCL** | Microsoft Collective Communication Library: **custom algorithms** and patterns; might be used when the stack is built and configured for them. |
+| **MSCCL++** | User-space collective paths aimed at **lower latency** for specific patterns and hardware. |
+| **ANP (AMD Network Plugin)** | Network backend integration (e.g. **AINIC**-oriented paths). Example: `NCCL_NET_PLUGIN` might point to `librccl-anp.so` or similar when installed (see Primus `examples/run_pretrain.sh` patterns). |
+
+### Environment variables
+
+Many deployments tune behavior with **NCCL-prefixed** variables (honored by RCCL for compatibility), for example:
+
+- `NCCL_PROTO`—protocol selection hints.
+- `NCCL_P2P_NET_CHUNKSIZE`—chunking for P2P/network paths.
+- `NCCL_IB_*`—InfiniBand / RDMA-related settings when applicable.
+- `NCCL_SOCKET_IFNAME`—**socket** interface selection for TCP fallback or hybrid setups.
+
+Document your cluster’s recommended values in [Environment variables](../03-configuration-reference/environment-variables.md).
+
+---
+
+## 7. Benchmarking collectives with Primus
+
+Primus includes an **RCCL microbenchmark** suite to measure **latency and bandwidth** for common collectives across message sizes.
+
+### Command
+
+Invoke through **`primus-cli`** (after `runner/` / container setup per your installation):
+
+```bash
+./primus-cli direct -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M
+```
+
+Useful flags (see `primus/tools/benchmark/rccl_bench_args.py`):
+
+| Flag | Purpose |
+|------|---------|
+| `--op` | One or more of: `all_reduce`, `broadcast`, `reduce_scatter`, `all_gather`, `alltoall` |
+| `--min-bytes`, `--max-bytes` | Sweep range (e.g. `1K`, `1M`, `128M`) |
+| `--num-sizes`, `--scale` | Generated sweep (`log2` or `linear`) |
+| `--dtype` | `bf16`, `fp16`, `fp32` |
+| `--output-file` | Write Markdown/CSV/JSONL report (default `./rccl_report.md`) |
+| `--check` | Lightweight correctness checks |
+
+Example with multiple ops:
+
+```bash
+./primus-cli direct -- benchmark rccl --op all_reduce all_gather reduce_scatter --min-bytes 1M --max-bytes 128M
+```
+
+### Reading results
+
+- **Bandwidth** (GB/s or similar): higher is better for large messages; compare against **peak NIC** or **GPU-GPU** limits for your topology.
+- **Latency** (µs): dominates for **small** messages; important for **frequent small** collectives (e.g. some TP patterns).
+
+Use results to spot **unexpected drops** (wrong NIC, congestion, fallback to TCP) before scaling full training.
+
+---
+
+## 8. Troubleshooting communication issues
+
+| Symptom | Checks |
+|---------|--------|
+| Hangs / timeouts | Enable **`NCCL_DEBUG=INFO`** (or `TRACE` for deep dives) and inspect which collective stalls. |
+| Wrong interface | Set **`NCCL_SOCKET_IFNAME`** to the intended **cluster** interface; verify with `ip link` / admin docs. |
+| IB / RDMA not used | Confirm **`NCCL_IB_*`**, HCA names, and permissions; run **preflight** (below). |
+| Slow AllReduce | Compare **`benchmark rccl`** to baseline; check **topology** (NVLink vs network), **contention**, **message sizes**. |
+
+### Preflight: Network validation
+
+Primus **preflight** can aggregate host/GPU/network info:
+
+```bash
+./primus-cli direct -- preflight --network
+```
+
+Combine with GPU checks as needed:
+
+```bash
+./primus-cli direct -- preflight --gpu --network
+```
+
+Use this to confirm **RCCL/NCCL-related environment** snapshots and **connectivity expectations** before long jobs.
+
+---
+
+## Related documentation
+
+- [Parallelism strategies](./parallelism-strategies.md)—how TP, PP, DP, FSDP, EP, and CP fit together.
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md)
+- [Environment variables](../03-configuration-reference/environment-variables.md)
diff --git a/docs/04-technical-guides/data-preparation.md b/docs/04-technical-guides/data-preparation.md
new file mode 100644
index 000000000..1dc6acbab
--- /dev/null
+++ b/docs/04-technical-guides/data-preparation.md
@@ -0,0 +1,159 @@
+# Data preparation guide
+
+Primus routes training through **Megatron-LM**, **TorchTitan**, and **MaxText**. Each backend expects its own data format and preprocessing pipeline. This guide summarizes how to prepare data, how to use **mock** data for smoke tests, and which environment variables commonly apply.
+
+Scripts referenced below live under the Primus repository root, for example:
+
+- `examples/megatron/preprocess_data.py`
+- `examples/megatron/prepare.py`
+- `examples/megatron/prepare_bookcorpus_megatron_dataset.py`
+- `examples/torchtitan/prepare.py`
+
+---
+
+## 1. Overview
+
+| Backend | Format | Typical entry |
+|---------|--------|----------------|
+| Megatron | Indexed `.bin` + `.idx` datasets | `data_path`, `train_data_path`, tokenizer args |
+| TorchTitan | Hugging Face datasets + local assets | `training.dataset`, `training.dataset_path`, `model.hf_assets_path` |
+| MaxText | TFDS / Hugging Face / Grain / synthetic | `dataset_type`, paths per pipeline |
+
+All backends support **synthetic or mock** data for configuration and scaling tests without large downloads.
+
+---
+
+## 2. Mock data (testing)
+
+### Megatron
+
+Set in the trainer module:
+
+```yaml
+mock_data: true
+```
+
+Default in `primus/configs/modules/megatron/trainer_base.yaml` is `false`. When `true`, training uses generated data matching configured dimensions so you can validate YAML, parallelism, and throughput without real corpora.
+
+### TorchTitan
+
+```yaml
+training:
+  mock_data: true
+```
+
+Default in `primus/configs/modules/torchtitan/pre_trainer.yaml` is `true` (useful for quick runs; set `false` and supply real datasets for production).
+
+### MaxText
+
+Use `dataset_type: synthetic` (or other synthetic paths in MaxText configs). See `third_party/maxtext/src/MaxText/configs/base.yml` and model YAMLs under `third_party/maxtext/src/MaxText/configs/`.
+
+---
+
+## 3. Megatron data pipeline
+
+### Inputs
+
+- Raw **JSON** or **JSONL** text (one JSON object per line for JSONL).
+- Optional **sentence splitting** via NLTK when `--split-sentences` is used (requires NLTK data; see [Environment variables](#6-environment-variables-for-data)).
+
+### Preprocessing: `examples/megatron/preprocess_data.py`
+
+The script tokenizes input and writes **Megatron indexed datasets** (`.bin` + `.idx`). It uses `build_tokenizer` from Primus’s Megatron tokenizer integration and accepts tokenizer flags from `_add_tokenizer_args`.
+
+**Important arguments** (from the script’s argparse):
+
+| Argument | Description |
+|----------|-------------|
+| `--input` | Path to input JSON (required). |
+| `--json-keys` | Keys to read (default `text`). |
+| `--output-prefix` | Output path **without** suffix; produces `{prefix}_{key}_{document|sentence}.bin` and `.idx`. |
+| `--workers` | Number of worker processes (required). |
+| `--partitions` | Split input for parallel preprocessing (default `1`). |
+| `--split-sentences` | Run NLTK sentence splitting before encode. |
+| `--append-eod` | Append end-of-document token. |
+
+**Example** (mirrors `examples/megatron/prepare.py` for BookCorpus-style flows):
+
+```bash
+python3 examples/megatron/preprocess_data.py \
+  --input /path/to/train.json \
+  --tokenizer-type HuggingFaceTokenizer \
+  --tokenizer-model /path/to/tokenizer \
+  --output-prefix /path/to/out/bookcorpus_train \
+  --workers "$(nproc)" \
+  --split-sentences \
+  --partitions 2
+```
+
+### Configuring training runs
+
+| Parameter | Notes |
+|-----------|--------|
+| `data_path` | Single path or **weighted blend**: `0.5 /path/a 0.5 /path/b` |
+| `train_data_path`, `valid_data_path`, `test_data_path` | Separate splits when used |
+| `split` | Train/valid/test ratio string, e.g. `"99,1,0"` (default in `trainer_base.yaml`) or `"98,2,0"` for train/valid/test |
+| `dataloader_type` | Megatron dataloader type; default in `trainer_base.yaml` is `null` (set explicitly in experiments as needed) |
+
+### BookCorpus example scripts
+
+- **`examples/megatron/prepare_bookcorpus_megatron_dataset.py`**—downloads BookCorpus to JSON via Hugging Face `datasets`, optional `--out-dir`.
+- **`examples/megatron/prepare.py`**—orchestrates download, train/valid split, and calls `preprocess_data.py` with tokenizer settings from Primus config; respects `TOKENIZED_TRAIN_DATA_PATH` / `TOKENIZED_EVAL_DATA_PATH` for output locations.
+
+### Tokenizers
+
+Tokenizer type and model path are set on the model preset (for example `tokenizer_type`, `tokenizer_model` in `primus/configs/models/megatron/language_model.yaml` comments list `Llama2Tokenizer`, `HuggingFaceTokenizer`, etc.).
+
+---
+
+## 4. TorchTitan data pipeline
+
+TorchTitan uses **Hugging Face datasets** style identifiers and local paths.
+
+| Key | Default (`pre_trainer.yaml`) | Description |
+|-----|------------------------------|-------------|
+| `training.dataset` | `c4` | Dataset identifier for TorchTitan loaders. |
+| `training.dataset_path` | `null` | Local directory for dataset assets when needed. |
+
+Tokenizer and model assets are resolved from **`model.hf_assets_path`** (or equivalent in your model preset). The preparation script **`examples/torchtitan/prepare.py`**:
+
+- Resolves the TorchTitan checkout path.
+- Runs `scripts/download_hf_assets.py` inside TorchTitan to fetch tokenizer assets for a given `repo_id`.
+- Uses `HF_TOKEN` when the model or dataset is gated.
+
+---
+
+## 5. MaxText data pipeline
+
+MaxText configuration is defined in upstream YAML (for example `third_party/maxtext/src/MaxText/configs/base.yml`).
+
+| Parameter | Meaning |
+|-----------|---------|
+| `dataset_type` | One of `synthetic`, `hf`, `grain`, `tfds` (per `base.yml` comments). |
+| `hf_path`, `hf_data_dir`, `hf_train_files` | Hugging Face pipeline inputs when `dataset_type: hf`. |
+| `per_device_batch_size` | Batch sizing on each device. |
+| `packing` | Sequence packing for efficiency (default `True` in `base.yml`). |
+
+See MaxText’s data input documentation for Grain and TFDS specifics.
+
+---
+
+## 6. Environment variables for data
+
+| Variable | Usage |
+|----------|--------|
+| `TOKENIZED_DATA_PATH` / `PRIMUS_TOKENIZED_DATA_PATH` | Tokenized dataset locations for Megatron hooks and examples (see `docs/03-configuration-reference/environment-variables.md`). |
+| `TOKENIZED_TRAIN_DATA_PATH`, `TOKENIZED_EVAL_DATA_PATH` | Override output paths in `examples/megatron/prepare.py`. |
+| `DATA_PATH` | General data root used in scripts and CI-style launches. |
+| `HF_TOKEN` | **Required** for gated Hugging Face models and some datasets (TorchTitan `prepare.py`, Kubernetes examples in `examples/README.md`). |
+| `HF_HOME` | Hugging Face cache directory (used in `examples/megatron/prepare.py`). |
+| `NLTK_DATA` | NLTK tokenizer data directory for sentence splitting in `preprocess_data.py` when `NLTK_DATA` is set. |
+
+---
+
+## Summary
+
+1. Use **mock or synthetic** data to validate configs and performance before investing in large preprocessing jobs.
+2. For **Megatron**, convert JSON/JSONL to `.bin`/`.idx` with `preprocess_data.py` and point `data_path` or split paths at the outputs.
+3. For **TorchTitan**, set `training.dataset` / `dataset_path` and run **`examples/torchtitan/prepare.py`** to fetch tokenizer assets.
+4. For **MaxText**, configure `dataset_type` and `per_device_batch_size` per upstream `base.yml` and model YAMLs.
diff --git a/docs/04-technical-guides/determinism-and-reproducibility.md b/docs/04-technical-guides/determinism-and-reproducibility.md
new file mode 100644
index 000000000..e81cb59a5
--- /dev/null
+++ b/docs/04-technical-guides/determinism-and-reproducibility.md
@@ -0,0 +1,120 @@
+# Determinism and reproducibility
+
+Reproducibility—getting bit-identical (or run-to-run stable) results—matters for debugging divergence, validating optimizations, and regression testing. This guide covers Primus's deterministic mode, the environment variables it sets, the per-backend seed/determinism knobs, and the performance trade-offs. Parameters and behavior are grounded in `examples/run_pretrain.sh`, `primus/configs/modules/megatron/trainer_base.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`.
+
+---
+
+## 1. What "deterministic" means here
+
+There are two distinct goals:
+
+- **Reproducible (seeded)**—same seed + same config + same hardware/software gives the *same trajectory*. Achieved with fixed seeds; cheap.
+- **Bitwise-deterministic**—kernels avoid non-deterministic reductions/atomics and tuning so results don't vary between runs. Requires deterministic algorithms and disabling autotuning; **slower**.
+
+Full determinism also generally requires the **same world size, parallelism layout, and library versions**. Changing TP/PP/DP, GPU count, or ROCm/Megatron versions can change numerics even with everything else fixed.
+
+---
+
+## 2. Primus deterministic mode (`PRIMUS_DETERMINISTIC`)
+
+Setting `PRIMUS_DETERMINISTIC=1` configures the GPU/communication stack for deterministic behavior. The CLI/runner path applies this through the hook `runner/helpers/hooks/05_deterministic.sh`; the `examples/run_pretrain.sh` script applies an equivalent inline block. The exported variables are:
+
+```bash
+# when PRIMUS_DETERMINISTIC=1 (runner/helpers/hooks/05_deterministic.sh)
+export NCCL_ALGO="Ring"                    # deterministic collective algorithm
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0  # Transformer Engine: forbid non-deterministic kernels
+export ROCBLAS_DEFAULT_ATOMICS_MODE=0      # rocBLAS: disable atomic (non-deterministic) reductions
+export TORCH_COMPILE_DISABLE=1             # avoid torch.compile/Triton race conditions
+export PRIMUS_TURBO_AUTO_TUNE=0            # disable Primus-Turbo autotuning (stable kernel choice)
+```
+
+> `PRIMUS_TURBO_AUTO_TUNE` also defaults to `0` in `runner/helpers/envs/base_env.sh`. The inline block in `examples/run_pretrain.sh` sets the first four variables and relies on that default for the fifth.
+
+Additionally, **HipBLASLt autotuning is disabled** in deterministic mode: tuning only runs when `PRIMUS_DETERMINISTIC != 1` *and* `PRIMUS_HIPBLASLT_TUNING=1` (`examples/run_pretrain.sh`). This prevents run-to-run kernel-selection differences. See [Performance tuning](./performance-tuning.md).
+
+`PRIMUS_DETERMINISTIC` is on the container passthrough allowlist (`runner/.primus.yaml`), so it reaches the training container. See [Environment variables](../03-configuration-reference/environment-variables.md).
+
+```bash
+export PRIMUS_DETERMINISTIC=1
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+> The MoE example scripts explicitly set `PRIMUS_DETERMINISTIC=0` because deterministic mode disables the performance kernels/tuning they rely on.
+
+---
+
+## 3. Seeds and deterministic algorithms (Megatron)
+
+In `primus/configs/modules/megatron/trainer_base.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `seed` | `1234` | Master RNG seed (Python/NumPy/Torch, data order, init). |
+| `deterministic_mode` | `false` | Force deterministic kernels/algorithms inside Megatron (slower; pairs with `PRIMUS_DETERMINISTIC`). |
+| `data_parallel_random_init` | `false` | When `false`, parameters are initialized identically and broadcast across DP ranks; keep `false` for reproducible init. |
+
+For a fully reproducible Megatron run: set a fixed `seed`, `deterministic_mode: true`, and launch with `PRIMUS_DETERMINISTIC=1`.
+
+> **Startup assertion.** When `deterministic_mode: true`, Primus validates (`primus/backends/megatron/patches/args/rocm_arg_validation.py`, `validate_args_on_rocm`) that these environment variables are set, and **fails fast** otherwise: `TORCH_COMPILE_DISABLE=1`, `ROCBLAS_DEFAULT_ATOMICS_MODE=0`, `PRIMUS_TURBO_AUTO_TUNE=0`, and `PRIMUS_DETERMINISTIC=1`. Launching with `PRIMUS_DETERMINISTIC=1` (above) sets all of them, so always pair `deterministic_mode: true` with `PRIMUS_DETERMINISTIC=1`.
+
+---
+
+## 4. Seeds and determinism (TorchTitan)
+
+Under `debug:` in `primus/configs/modules/torchtitan/pre_trainer.yaml` (TorchTitan v0.2.2 moved these keys out of `training:`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `seed` | `null` | RNG seed; set an integer for reproducible runs. |
+| `deterministic` | `false` | Enable deterministic algorithms (disables some optimized kernels; slower). |
+| `deterministic_warn_only` | `false` | With `deterministic: true`, warn instead of erroring when an op has no deterministic implementation. |
+
+Related: `checkpoint.create_seed_checkpoint` (`false`) creates a deterministic seed checkpoint that all ranks load, ensuring identical initialization across a distributed run.
+
+Note `compile.enable: true` is the TorchTitan default; for strict determinism prefer launching with `PRIMUS_DETERMINISTIC=1` (which sets `TORCH_COMPILE_DISABLE=1`) or disable compilation.
+
+---
+
+## 5. MaxText
+
+MaxText determinism is governed by the upstream MaxText seed/data options surfaced through the MaxText config (see [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)). The GPU-stack environment effects of `PRIMUS_DETERMINISTIC` (rocBLAS atomics, deterministic collectives) still apply at the launcher level.
+
+---
+
+## 6. Performance trade-offs
+
+Determinism is not free:
+
+| Setting | Cost |
+|---------|------|
+| `NCCL_ALGO=Ring` | Forgoes faster topology-aware collective algorithms. |
+| `ROCBLAS_DEFAULT_ATOMICS_MODE=0` | Disables atomic reductions—slower GEMMs. |
+| `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` | Restricts TE to deterministic (often slower) kernels. |
+| `TORCH_COMPILE_DISABLE=1` | No `torch.compile` fusion/codegen speedups. |
+| `PRIMUS_TURBO_AUTO_TUNE=0` | No Primus-Turbo kernel autotuning. |
+| HipBLASLt tuning disabled | No autotuned GEMM kernels. |
+| `deterministic_mode` / `deterministic` | Deterministic algorithm variants are generally slower. |
+
+**Use deterministic mode for debugging and validation, not production throughput runs.** Once a result is reproduced/diagnosed, disable it to recover performance.
+
+---
+
+## 7. Reproducibility checklist
+
+1. **Pin the environment**—same container image, ROCm version, and backend (Megatron/TorchTitan) commit.
+2. **Fix seeds**—Megatron `seed`; TorchTitan `debug.seed`.
+3. **Hold the layout constant**—same world size and TP/PP/DP/EP/CP degrees.
+4. **Enable determinism**—`PRIMUS_DETERMINISTIC=1` plus backend `deterministic_mode`/`deterministic`.
+5. **Disable autotuning**—automatic in deterministic mode (HipBLASLt tuning off).
+6. **Use mock or fixed data ordering**—ensure the data pipeline is seeded; see [Data preparation](./data-preparation.md).
+7. **Record everything**—log the full resolved config and env (see [Logging & experiment tracking](./logging-and-experiment-tracking.md)).
+
+---
+
+## Related documentation
+
+- [Performance tuning](./performance-tuning.md)—HipBLASLt tuning and its interaction with deterministic mode.
+- [Environment variables](../03-configuration-reference/environment-variables.md)—`PRIMUS_DETERMINISTIC` and related flags.
+- [Data preparation](./data-preparation.md)—deterministic data ordering.
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md).
diff --git a/docs/04-technical-guides/diffusion-models/README.md b/docs/04-technical-guides/diffusion-models/README.md
new file mode 100644
index 000000000..26daa6cb3
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/README.md
@@ -0,0 +1,371 @@
+# Diffusion models in Primus - developer and architecture guide
+
+**Purpose:** Developer-focused documentation for understanding Primus diffusion architecture, design decisions, and implementation details.
+
+**For training/usage instructions, see:** [examples/megatron/diffusion/README.md](../../../examples/megatron/diffusion/README.md)
+
+**For test documentation, see:** [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/)
+
+---
+
+## Architecture philosophy
+
+Primus diffusion models are built as **Megatron-Core native implementations**, designed for:
+- Production-scale distributed training
+- Seamless integration with Megatron parallelism strategies (TP, PP, DP, EP)
+- Advanced checkpoint management with heterogeneous layers
+- Clean separation of concerns (no framework dependencies like PyTorch Lightning)
+
+### Key design decisions
+
+**1. Megatron-Core Integration**
+- Models in `core/models/diffusion/` follow Megatron-Core patterns
+- Extends `TransformerConfig` for configurations (inherits all Megatron features)
+- Uses `TransformerBlock` with heterogeneous layer support
+- Compatible with Megatron's distributed checkpointing
+
+**2. Unified TransformerBlock Architecture**
+- Unlike HuggingFace's ModuleLists, uses Megatron's unified TransformerBlock
+- Simplifies checkpoint management
+- More efficient gradient synchronization
+- Note: pipeline parallelism is not supported for diffusion models (`pipeline_model_parallel_size` must be 1)
+
+**3. No Framework Dependencies**
+- Direct PyTorch implementation (no PyTorch Lightning)
+- Uses Megatron's distributed primitives directly
+- Simpler debugging and profiling
+- Better control over distributed training
+
+**4. Extensibility First**
+- Base classes designed for multiple diffusion models (Flux, DiT, MovieGen)
+- Clear shared vs model-specific separation
+- Hierarchical encoder registry for easy extension
+
+---
+
+## Supported models
+
+### Flux ✅ production ready
+Flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer) architecture.
+
+- **Architecture**: Dual-stream with joint and single transformer blocks
+- **Sizes**: 535M (testing) and 12B (production)
+- **Reference**: [Black Forest Labs FLUX.1](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+- **Status**: Fully implemented and tested (390 tests)
+
+### Future models ⏳ planned
+- **DiT**: Diffusion Transformer for image generation
+- **MovieGen**: Video diffusion models
+- **Custom Models**: Extensible framework for new architectures
+
+---
+
+## Project structure
+
+```
+primus/backends/megatron/
+├── core/models/
+│   ├── common/diffusion_module/    # DiffusionModule (base class with sharded state dict)
+│   │   └── diffusion_module.py
+│   └── diffusion/                  # Model implementations (Megatron-Core style)
+│       ├── common/                 # Shared components (MMDiT layers, attention)
+│       │   ├── config.py           # BaseDiffusionConfig (extends TransformerConfig)
+│       │   └── layers.py           # Shared layers (if any)
+│       └── flux/                   # Flux-specific code
+│           ├── config.py           # FluxConfig with factory methods (535M, 12B)
+│           ├── model.py            # Flux model (extends DiffusionModule)
+│           └── layer_spec.py       # Flux layer specifications
+│
+├── training/diffusion/             # Training utilities
+│   ├── noise_utils.py              # Noise application (flow matching, DDPM)
+│   ├── loss_computation.py         # Loss functions (flow matching, epsilon, v-prediction)
+│   ├── timestep_sampling.py        # Timestep sampling strategies
+│   └── schedulers/
+│       ├── base.py                 # BaseScheduler
+│       └── flow_matching.py        # FlowMatchEulerDiscreteScheduler
+│
+└── data/
+    ├── energon/                    # Shared Energon infrastructure
+    └── diffusion/                  # Diffusion-specific data
+        ├── encoders/               # Hierarchical encoder registry
+        │   ├── image/vae/          # VAE variants (SD VAE, custom VAEs)
+        │   ├── text/t5/            # T5 variants (XXL, etc.)
+        │   └── text/clip/          # CLIP variants (L, H, etc.)
+        ├── preprocessing/          # Data preprocessing utilities
+        │   ├── download.py         # Reusable download utils (retry, MD5, manifests)
+        │   ├── finalize.py         # Energon dataset finalization
+        │   ├── validate.py         # Dataset structure validation
+        │   └── pipelines/          # Dataset preparation pipelines
+        │       ├── base.py         # DatasetPipeline abstract base class
+        │       ├── raw.py          # Raw image pipeline
+        │       ├── encoded.py      # Pre-encoded pipeline
+        │       └── ingest.py       # StreamingIngestPipeline (MLPerf Arrow->WDS)
+        └── task_encoders/          # Energon TaskEncoders for diffusion
+
+primus/configs/models/megatron/diffusion/
+├── flux_535m.yaml                  # Flux 535M config
+├── flux_12b.yaml                   # Flux 12B config
+└── encoders.yaml                   # Encoder configuration
+
+tests/unit_tests/backends/megatron/diffusion/  # Comprehensive test suite (390 tests)
+├── models/                         # Model-level tests
+├── layers/                         # Layer-level tests
+├── unit/                          # Unit tests for utilities
+├── distributed/                    # Distributed training tests
+├── functional/                     # End-to-end functional tests
+└── checkpointing/                  # Checkpoint tests
+
+docs/04-technical-guides/diffusion-models/   # This directory
+├── README.md                       # This file (developer guide)
+├── architecture_overview.md        # Detailed architecture
+├── data_preprocessing.md           # Data pipeline guide (includes Flux-specific section)
+├── energon_integration.md          # Energon patterns
+├── flux_architecture.md            # Flux deep dive
+├── fp8_training.md                 # FP8 training guide (benchmarks, tuning, troubleshooting)
+├── api_reference.md                # API documentation
+├── adding_new_models.md            # Extension guide
+└── STRUCTURE.md                    # Directory tree and organization
+```
+
+---
+
+## Key technical features
+
+### 1. DiffusionModule base class
+
+All diffusion models inherit from `DiffusionModule`, which provides:
+- Megatron-Core integration (process groups, parallelism)
+- Sharded state dict support for distributed checkpointing
+- Gradient checkpointing
+- Mixed precision support
+- Device placement utilities
+
+**Location:** `primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py`
+
+### 2. BaseDiffusionConfig
+
+Configuration class extending `TransformerConfig`:
+- Inherits all Megatron-Core configuration (TP, PP, sequence_parallel, etc.)
+- Adds diffusion-specific parameters (channels, patch_size, etc.)
+- Factory methods for common presets
+
+**Location:** `primus/backends/megatron/core/models/diffusion/common/config.py`
+
+### 3. Hierarchical encoder registry
+
+Organized by modality → type → variant:
+```
+encoders/
+├── image/vae/
+│   ├── sd_vae.py              # Standard SD VAE
+│   └── (future: custom VAEs)
+├── text/t5/
+│   ├── t5_xxl.py              # T5-XXL encoder
+│   └── (future: T5 variants)
+└── text/clip/
+    ├── clip_l.py              # CLIP-L encoder
+    └── (future: CLIP-H, etc.)
+```
+
+Benefits:
+- Easy to add new encoder variants (5+ planned per modality)
+- Config-driven selection via `encoders.yaml`
+- Lazy loading (encoders loaded only when needed)
+- Shared base classes for common functionality
+
+### 4. Training utilities structure
+
+**Noise Application** (`noise_utils.py`):
+- `apply_flow_matching_noise()`: For flow matching models (Flux)
+- `apply_ddpm_noise()`: For DDPM-based models
+- Support for different noise schedules
+
+**Loss Computation** (`loss_computation.py`):
+- `compute_flow_matching_loss()`: For flow matching
+- `compute_epsilon_loss()`: For epsilon prediction (DDPM)
+- `compute_v_prediction_loss()`: For v-prediction
+- Unified interface for different loss types
+
+**Timestep Sampling** (`timestep_sampling.py`):
+- `LogitNormalSampler`: Logit-normal distribution
+- `UniformSampler`: Uniform distribution
+- `ModeSampler`: Mode-focused sampling
+- Base class for custom samplers
+
+### 5. Shared Energon infrastructure
+
+Located in `data/energon/` for reusability across models:
+- Shared data loading utilities
+- Common preprocessing functions
+- WebDataset integration
+- Model-specific TaskEncoders in `data/diffusion/task_encoders/`
+
+### 6. Precalculated data support
+
+**Performance**: 5-10x faster training than on-the-fly encoding
+
+**Supported encodings**:
+- `preencoded` -- Primus-encoded PyTorch `.pth` format (VAE latents + text embeddings)
+- `preencoded_numpy` -- MLPerf NumPy uint16 format (bfloat16 tensors as `.bytes` entries)
+
+**Workflow**:
+1. Precompute VAE latents and text embeddings offline
+2. Store in WebDataset/Energon format
+3. Load directly during training (no encoder overhead)
+
+**Benefits**:
+- Faster training iteration
+- Consistent encoder versions across runs
+- Lower GPU memory (no encoders loaded during training)
+- Better reproducibility
+
+### 7. MLPerf streaming ingest pipeline
+
+**Location:** `data/diffusion/preprocessing/pipelines/ingest.py`
+
+The `StreamingIngestPipeline` downloads Apache Arrow IPC files from MLCommons R2 storage and converts them directly into Energon WebDataset tar shards in a single streaming pass. This avoids storing the full ~6 TB raw Arrow dataset on disk.
+
+**Architecture**: Producer-consumer with concurrent download and sequential conversion:
+- **Producer thread**: Acquires a semaphore permit, submits downloads to a `ThreadPoolExecutor`, passes completed futures to a drain thread
+- **Drain thread**: Processes futures in submission order and feeds the prefetch queue
+- **Consumer (main thread)**: Converts Arrow data to tar shards, deletes temporary files, releases semaphore permits
+
+**Key properties**:
+- Bounded disk usage: `threading.Semaphore(prefetch_depth)` limits Arrow files on disk
+- Deterministic shard ordering preserved via in-order future draining
+- Retry with exponential backoff for HTTP 429/503 and MD5 mismatches (`download.py`)
+- Skip-and-log: individual failures are recorded in `failed_files.json`
+- Resume: re-running skips shards that already exist on disk
+
+**Related modules**:
+- `download.py`: `download_with_backoff()`, `fetch_manifest()`, `parse_md5_manifest()`
+- `pipelines/base.py`: `DatasetPipeline` ABC (shared by `raw.py`, `encoded.py`, `ingest.py`)
+- `finalize.py`: Energon dataset finalization (`.nv-meta/dataset.yaml` + `energon prepare`)
+- `validate.py`: Post-finalization structural validation
+
+---
+
+## Implementation status
+
+### Core infrastructure ✅
+- ✅ Directory structure with 25+ directories
+- ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler)
+- ✅ DiffusionModule with Megatron-Core integration
+- ✅ FluxConfig with factory methods (flux_535m, flux_12b)
+- ✅ FlowMatchEulerDiscreteScheduler
+- ✅ Configuration system (YAML files)
+- ✅ Testing framework (390 tests)
+- ✅ Comprehensive documentation
+
+### Flux model implementation ✅
+- ✅ Flux model architecture (dual-stream MMDiT)
+- ✅ MMDiT layers and attention (joint + single blocks)
+- ✅ Embeddings (3D RoPE, timestep, vector)
+- ✅ Hierarchical encoder registry
+- ✅ Data pipeline and TaskEncoders
+- ✅ Training utilities (noise, loss, sampling)
+- ✅ Checkpoint conversion (HF <-> Megatron)
+
+---
+
+## Documentation map
+
+### Core guides
+
+📖 **[Architecture Overview](architecture_overview.md)**
+High-level design, directory structure, and architectural decisions.
+
+📖 **[Directory Structure](STRUCTURE.md)**
+Complete directory tree and file organization.
+
+📖 **[Data Preprocessing Guide](data_preprocessing.md)**
+How to prepare datasets, precalculate latents, and use Energon.
+
+📖 **[Energon Integration](energon_integration.md)**
+Megatron-Energon patterns and TaskEncoder implementation.
+
+📖 **[Adding New Models](adding_new_models.md)**
+Step-by-step guide for implementing new diffusion models.
+
+### Advanced documentation
+
+📖 **[Flux Architecture Deep Dive](flux_architecture.md)**
+Mathematical formulation, detailed component descriptions, and performance optimizations.
+
+📖 **[API Reference](api_reference.md)**
+Complete API documentation with function signatures and usage examples.
+
+📖 **[FP8 Training Guide](fp8_training.md)**
+FP8 precision training on AMD MI300X: configuration, benchmarks, tuning recipes, and troubleshooting.
+
+### Related documentation
+
+📖 **[Training Guide](../../../examples/megatron/diffusion/README.md)**
+User-facing guide for training Flux models (quick start, configurations, troubleshooting).
+
+📖 **[Test Directory](../../../tests/unit_tests/backends/megatron/diffusion/)**
+Test suite for diffusion models.
+
+---
+
+## Testing architecture
+
+**Test Organization** (following Megatron-LM patterns):
+- One comprehensive file per model (`test_flux_model.py`)
+- Unit tests for utilities (`unit/test_utils.py`, etc.)
+- Distributed tests in separate directory (`distributed/`)
+- Functional tests for workflows (`functional/`)
+
+**Test Status**: ✅ 390 tests passing
+
+See [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/) for details.
+
+---
+
+## Hardware requirements
+
+### Flux 535M (testing)
+- **Training**: 1x MI300X 192GB (compatible with H100/A100)
+- **Inference**: 1x MI300X 192GB
+- **Batch Size**: 1-8 per GPU
+
+### Flux 12B (production)
+- **Training**: 8x MI300X 192GB (recommended) or 4x MI300X 192GB with TP=2
+- **Inference**: 1x MI300X 192GB
+- **Batch Size**: 1-2 per GPU for training, 1-4 for inference
+
+---
+
+## Contributing
+
+See the main guide: [Adding New Models](adding_new_models.md)
+
+**To contribute:**
+1. Follow the established directory structure
+2. Extend base classes (DiffusionModule, BaseDiffusionConfig)
+3. Add comprehensive tests in `tests/unit_tests/backends/megatron/diffusion/`
+4. Update documentation (architecture guide + API reference)
+5. Submit PR with clear description
+
+---
+
+## License
+
+- **Primus Code**: AMD Copyright 2025, Apache License 2.0
+- **Flux Encoders**:
+  - FLUX.1 [dev]: Non-commercial license
+  - FLUX.1 [schnell]: Apache 2.0 (commercial use allowed)
+  - Individual components (T5, CLIP, VAE): Check respective licenses
+
+---
+
+## Resources
+
+- **Megatron-Core**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) - Core framework
+- **Flux Model**: [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+- **Flow Matching**: Rectified flow and flow matching papers
+- **NeMo**: [nvidia/NeMo](https://github.com/NVIDIA/NeMo) - Alternative diffusion implementation
+
+---
+
+**Last Updated**: January 2026
diff --git a/docs/04-technical-guides/diffusion-models/STRUCTURE.md b/docs/04-technical-guides/diffusion-models/STRUCTURE.md
new file mode 100644
index 000000000..9bed81156
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/STRUCTURE.md
@@ -0,0 +1,251 @@
+# Flux diffusion infrastructure - directory structure
+
+**Created**: December 5, 2025
+**Status**: ✓ Implementation Complete
+
+## Overview
+
+This document describes the directory structure created for Flux diffusion model support in Primus, following Megatron-Core conventions with production-ready enhancements.
+
+---
+
+## Directory tree
+
+```
+Primus/
+├── primus/backends/megatron/
+│   ├── core/models/
+│   │   ├── common/diffusion_module/         # DiffusionModule base class
+│   │   │   └── diffusion_module.py
+│   │   └── diffusion/                      # Diffusion models (Megatron-Core convention)
+│   │       ├── common/                     # Shared building blocks (config, embeddings, normalization)
+│   │       │   ├── __init__.py
+│   │       │   ├── config.py               # ✓ BaseDiffusionConfig
+│   │       │   ├── embeddings.py           # ✓ TimeStepEmbedder, MLPEmbedder
+│   │       │   └── normalization.py        # ✓ AdaLN, AdaLNContinuous, RMSNorm
+│   │       ├── flux/                       # Flux-specific components
+│   │       │   ├── __init__.py
+│   │       │   ├── config.py               # ✓ FluxConfig (with factory methods)
+│   │       │   ├── model.py                # ✓ Flux model
+│   │       │   ├── layers.py               # ✓ EmbedND, embedders
+│   │       │   ├── layer_spec.py           # ✓ get_flux_layer_spec, get_flux_*_spec_for_backend, MMDiTLayer
+│   │       │   ├── attention.py            # ✓ JointSelfAttention, FluxSingleAttention
+│   │       │   ├── utils.py                # ✓ generate_image_position_ids
+│   │       │   └── checkpoint_converter.py # ✓ HF <-> Megatron conversion
+│   │       └── __init__.py
+│   │
+│   ├── training/diffusion/                 # Training utilities
+│   │   ├── schedulers/
+│   │   │   ├── __init__.py
+│   │   │   ├── base.py                     # ✓ BaseScheduler
+│   │   │   └── flow_matching.py            # ✓ FlowMatchEulerDiscreteScheduler
+│   │   ├── noise_utils.py                  # ✓ apply_flow_matching_noise, apply_ddpm_noise
+│   │   ├── loss_computation.py             # ✓ compute_flow_matching_loss, etc.
+│   │   ├── timestep_sampling.py             # ✓ LogitNormalSampler, UniformSampler
+│   │   └── __init__.py
+│   │
+│   └── data/
+│       ├── energon/                        # Shared Energon infrastructure
+│       │   └── __init__.py                 # ✓ Energon wrappers
+│       │
+│       └── diffusion/                      # Diffusion-specific data
+│           ├── encoders/                   # Hierarchical encoder registry
+│           │   ├── image/
+│           │   │   ├── vae/                # VAE variants
+│           │   │   │   └── __init__.py     # ✓ AutoencoderKL, VQVAE, etc.
+│           │   │   └── __init__.py
+│           │   ├── text/
+│           │   │   ├── t5/                 # T5 variants
+│           │   │   │   └── __init__.py     # ✓ T5-XXL, T5-Large, etc.
+│           │   │   ├── clip/               # CLIP variants
+│           │   │   │   └── __init__.py     # ✓ CLIP-L, CLIP-H, etc.
+│           │   │   └── __init__.py
+│           │   └── __init__.py             # ✓ EncoderRegistry
+│           │
+│           ├── preprocessing/
+│           │   ├── image/
+│           │   │   └── __init__.py         # ✓ Resizing, augmentation
+│           │   └── __init__.py
+│           │
+│           ├── task_encoders/              # Energon TaskEncoders
+│           │   ├── __init__.py
+│           │   └── image.py                # ✓ EncodedDiffusionTaskEncoder, RawDiffusionTaskEncoder
+│           │
+│           └── __init__.py
+│
+├── primus/backends/megatron/
+│   └── megatron_pretrain_trainer.py        # ✓ Shared Megatron pretrain trainer (drives diffusion pretraining)
+│
+├── primus/configs/models/megatron/
+│   └── diffusion/                          # YAML configs
+│       ├── __init__.py
+│       ├── flux_535m.yaml                  # ✓ Flux 535M config
+│       ├── flux_12b.yaml                   # ✓ Flux 12B config
+│       └── encoders.yaml                   # ✓ Encoder configs
+│
+├── examples/megatron/
+│   ├── diffusion/
+│   │   └── README.md                       # ✓ Training guide (consolidated)
+│   ├── configs/MI300X/diffusion/           # MI300X training configs
+│   │   ├── flux_535m_pretrain.yaml
+│   │   ├── flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml
+│   │   ├── flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml
+│   │   └── ...
+│   ├── configs/MI355X/diffusion/           # MI355X training configs (mirrors MI300X + MXFP4/MLPerf)
+│   │   ├── flux_12b_ddp_energon_schnell_resample_*.yaml
+│   │   ├── flux_12b_fsdp2_energon_schnell_resample_*.yaml
+│   │   └── ...
+│   └── prepare.py
+│
+├── examples/run_pretrain.sh                # Main training script
+│
+├── tests/
+│   ├── unit_tests/backends/megatron/diffusion/   # Unit test suite
+│   │   ├── test_flux_model.py
+│   │   ├── test_flux_config.py
+│   │   ├── test_flux_layers.py
+│   │   ├── test_flux_embeddings.py
+│   │   ├── test_flux_normalization.py
+│   │   ├── test_flux_utils.py
+│   │   ├── test_flux_checkpoint_converter.py
+│   │   ├── test_flux_checkpoint_utils.py
+│   │   ├── test_flux_layer_spec_backend_selection.py
+│   │   ├── test_flux_compile_checkpoint_keys.py
+│   │   ├── training/
+│   │   ├── data/
+│   │   └── distributed/
+│   └── integration_tests/backends/megatron/diffusion/
+│       ├── data/
+│       └── distributed/
+│
+└── docs/backends/megatron/
+    └── diffusion/                          # Documentation
+        ├── README.md                       # ✓ Overview
+        ├── STRUCTURE.md                    # ✓ This file
+        ├── architecture_overview.md        # ✓ Design details
+        ├── data_preprocessing.md           # ✓ Data guide (includes Flux-specific section)
+        ├── energon_integration.md          # ✓ Energon patterns
+        ├── flux_architecture.md            # ✓ Flux deep dive
+        ├── fp8_training.md                 # ✓ FP8 training guide
+        ├── api_reference.md                # ✓ API documentation
+        └── adding_new_models.md            # ✓ Extension guide
+```
+
+---
+
+## Completed components
+
+### ✓ Base classes
+
+1. **DiffusionModule** (`core/models/common/diffusion_module/diffusion_module.py`)
+   - Base class for all diffusion models (extends MegatronModule)
+   - Provides Megatron-Core integration
+   - Required methods: `forward()`
+   - Loss computation: Use standalone functions from `loss_computation.py`
+   - Utility methods: `get_num_params()`, `set_requires_grad()`
+
+2. **BaseDiffusionConfig** (`common/config.py`)
+   - Extends `megatron.core.transformer.transformer_config.TransformerConfig`
+   - Common parameters: `in_channels`, `out_channels`, `patch_size`
+   - Validation method for configuration integrity
+
+3. **FluxConfig** (`flux/config.py`)
+   - Flux-specific configuration
+   - Parameters: `num_joint_layers`, `num_single_layers`, `context_dim`, `vec_in_dim`
+   - Factory methods: `flux_535m()`, `flux_12b()`
+   - 3D RoPE configuration: `axes_dim`, `theta`
+
+3. **BaseScheduler** (`schedulers/base.py`)
+   - Abstract base for diffusion schedulers
+   - Required: `add_noise()`, `get_velocity_target()`, `sample_timesteps()`
+   - Optional: `scale_model_input()`, `get_snr()`, `get_alpha()`, `get_sigma()`
+
+4. **FlowMatchEulerDiscreteScheduler** (`schedulers/flow_matching.py`)
+   - Concrete implementation for Flux
+   - Linear interpolation: `x_t = (1-t)*noise + t*data`
+   - Velocity target: `v = data - noise`
+
+### ✓ Directory structure
+
+- **25 `__init__.py` files** with comprehensive docstrings
+- **Multiple implementation files** (models, configs, schedulers, data pipeline)
+- **Complete test suite** with fixtures and helpers
+
+---
+
+## Architectural decisions
+
+### 1. Models under `core/models/`
+- Follows Megatron-Core convention (`megatron/core/models/gpt/`, etc.)
+- Easier upstream tracking when Megatron-Core adds diffusion support
+
+### 2. Shared components in `common/`
+- Standard approach stores shared code in model-specific directories
+- Primus: `common/` for shared config, embeddings, and normalization
+- Flux-specific: model class, MMDiT/single-block layer specs, joint attention, and `EmbedND`
+
+### 3. Hierarchical encoder structure
+- `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/`
+- Registry pattern for config-driven selection
+- Easy to add new encoder variants (5+ planned per modality)
+
+### 4. Shared Energon infrastructure
+- `data/energon/` for cross-model utilities (VLM, diffusion, future)
+- `data/diffusion/task_encoders/` for diffusion-specific TaskEncoders
+- Traditional approach nests Energon under model-specific directories
+
+### 5. Synthetic (mock) data
+- Synthetic datasets live in `primus/backends/megatron/data/synthetic/mock_datasets.py`, wired through `primus/backends/megatron/data/synthetic_dataset_provider.py`, so training can run without real data
+- Unit tests exercise them under `tests/unit_tests/backends/megatron/diffusion/data/`
+
+### 6. No PyTorch lightning
+- Pure Megatron patterns (no PTL DataModules)
+- Better integration with Megatron training loop
+
+---
+
+## Import examples
+
+```python
+# Base classes
+from primus.backends.megatron.core.models.diffusion.common import (
+    BaseDiffusionConfig,
+)
+
+# Flux configuration
+from primus.backends.megatron.core.models.diffusion.flux import FluxConfig
+
+# Create configs
+config_535m = FluxConfig.flux_535m()
+config_12b = FluxConfig.flux_12b()
+
+# Schedulers
+from primus.backends.megatron.training.diffusion.schedulers import (
+    BaseScheduler,
+    FlowMatchEulerDiscreteScheduler,
+)
+
+# Create scheduler
+scheduler = FlowMatchEulerDiscreteScheduler()
+timesteps = scheduler.sample_timesteps(batch_size=8, device='cuda')
+```
+
+---
+
+## Validation status
+
+✓ All Python files syntactically correct
+✓ No linter errors detected
+✓ All imports properly structured
+✓ Comprehensive docstrings
+✓ Copyright headers applied (AMD 2025, Apache 2.0)
+
+---
+
+## Files summary
+
+All infrastructure files, model implementations, data pipeline components, tests, and documentation are complete and ready for production use.
+
+---
+
+**End of Structure Document**
diff --git a/docs/04-technical-guides/diffusion-models/adding_new_models.md b/docs/04-technical-guides/diffusion-models/adding_new_models.md
new file mode 100644
index 000000000..e975bae64
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/adding_new_models.md
@@ -0,0 +1,711 @@
+# Adding new diffusion models
+
+This guide explains how to add new diffusion models to Primus, following the established patterns and architecture.
+
+---
+
+## Table of contents
+
+1. [Overview](#overview)
+2. [Prerequisites](#prerequisites)
+3. [Step-by-step guide](#step-by-step-guide)
+4. [Example: Adding DiT](#example-adding-dit)
+5. [Testing your model](#testing-your-model)
+6. [Best practices](#best-practices)
+
+---
+
+## Overview
+
+Adding a new diffusion model involves:
+1. Creating model configuration
+2. Implementing model class
+3. Adding necessary layers
+4. Creating data pipeline components
+5. Writing tests
+6. Updating documentation
+
+**Time Estimate**: 5-10 days depending on model complexity
+
+---
+
+## Prerequisites
+
+Before adding a new model, ensure you have:
+- ✅ Understanding of the model architecture (paper, reference implementation)
+- ✅ Access to pretrained weights (if applicable)
+- ✅ Sample dataset for testing
+- ✅ Familiarity with Primus diffusion architecture
+- ✅ Development environment setup
+
+**Required Reading**:
+- [Architecture Overview](architecture_overview.md)
+- [Data Preprocessing Guide](data_preprocessing.md)
+- [Energon Integration](energon_integration.md)
+
+---
+
+## Step-by-step guide
+
+### Step 1: Create model directory
+
+Create a directory for your model under `core/models/diffusion/`:
+
+```bash
+mkdir -p primus/backends/megatron/core/models/diffusion/dit
+cd primus/backends/megatron/core/models/diffusion/dit
+```
+
+Create files:
+```bash
+touch __init__.py
+touch config.py
+touch model.py
+touch layers.py  # If model-specific layers needed
+```
+
+### Step 2: Implement configuration
+
+**File**: `config.py`
+
+```python
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""Configuration for DiT (Diffusion Transformer) model."""
+
+from dataclasses import dataclass
+from typing import Optional
+from ..common.config import BaseDiffusionConfig
+
+
+@dataclass
+class DiTConfig(BaseDiffusionConfig):
+    """
+    DiT-specific configuration.
+
+    DiT uses a standard transformer architecture for diffusion.
+    """
+
+    # Model identification
+    model_type: str = "dit"
+
+    # Architecture: Number of layers
+    num_layers: int = 28  # DiT-XL/2 default
+
+    # Architecture: Dimensions
+    hidden_size: int = 1152
+    num_attention_heads: int = 16
+
+    # Input dimensions
+    in_channels: int = 4  # VAE latent channels (standard SD VAE)
+
+    # Context dimensions
+    context_dim: int = 768  # CLIP text embedding dimension
+
+    # Patchification
+    patch_size: int = 2  # DiT uses 2x2 patches
+
+    # Class conditioning (for conditional generation)
+    num_classes: int = 1000  # ImageNet classes
+    class_dropout_prob: float = 0.1
+
+    # Adaptive LayerNorm (DiT-specific)
+    use_adaptive_layernorm: bool = True
+
+    def validate(self):
+        """Validate DiT-specific configuration."""
+        # Call parent validation
+        super().validate()
+
+        # DiT-specific validations
+        if self.num_layers <= 0:
+            raise ValueError(f"num_layers must be positive, got {self.num_layers}")
+
+        if self.patch_size <= 0:
+            raise ValueError(f"patch_size must be positive, got {self.patch_size}")
+
+        if self.num_classes < 0:
+            raise ValueError(f"num_classes must be non-negative, got {self.num_classes}")
+
+    @classmethod
+    def dit_xl_2(cls, **kwargs):
+        """
+        Create configuration for DiT-XL/2.
+
+        Args:
+            **kwargs: Override default parameters
+
+        Returns:
+            DiTConfig instance
+        """
+        defaults = {
+            'num_layers': 28,
+            'hidden_size': 1152,
+            'num_attention_heads': 16,
+            'patch_size': 2,
+        }
+        defaults.update(kwargs)
+        return cls(**defaults)
+
+    @classmethod
+    def dit_l_2(cls, **kwargs):
+        """Create configuration for DiT-L/2."""
+        defaults = {
+            'num_layers': 24,
+            'hidden_size': 1024,
+            'num_attention_heads': 16,
+            'patch_size': 2,
+        }
+        defaults.update(kwargs)
+        return cls(**defaults)
+```
+
+### Step 3: Implement model class
+
+**File**: `model.py`
+
+```python
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""DiT model implementation."""
+
+import torch
+import torch.nn as nn
+from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import DiffusionModule
+from megatron.core.process_groups_config import ProcessGroupCollection
+from ..common.layers import MMDiTLayer  # Reuse shared components if applicable
+from .config import DiTConfig
+
+
+class DiT(DiffusionModule):
+    """
+    DiT (Diffusion Transformer) model.
+
+    Reference: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023)
+
+    Note: Inherits from DiffusionModule for Megatron-Core integration
+    (process groups, distributed checkpointing, attention backend config)
+    """
+
+    def __init__(
+        self,
+        config: DiTConfig,
+        encoder_configs: Optional[Dict[str, Any]] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+    ):
+        """
+        Initialize DiT model.
+
+        Args:
+            config: DiT configuration
+            encoder_configs: Optional encoder configurations (VAE, T5, CLIP)
+            pg_collection: Process group collection for distributed training
+        """
+        super().__init__(config, pg_collection=pg_collection, encoder_configs=encoder_configs)
+
+        self.config = config
+
+        # Input projection (patchify)
+        self.input_proj = nn.Conv2d(
+            config.in_channels,
+            config.hidden_size,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+        )
+
+        # Positional embedding
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, (config.seq_length // config.patch_size) ** 2, config.hidden_size)
+        )
+
+        # Timestep embedding
+        self.time_embed = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(config.hidden_size * 4, config.hidden_size),
+        )
+
+        # Class embedding (for conditional generation)
+        if config.num_classes > 0:
+            self.class_embed = nn.Embedding(config.num_classes, config.hidden_size)
+
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            DiTBlock(config) for _ in range(config.num_layers)
+        ])
+
+        # Output projection
+        self.output_proj = nn.Sequential(
+            nn.LayerNorm(config.hidden_size),
+            nn.Linear(config.hidden_size, config.patch_size ** 2 * config.out_channels),
+        )
+
+        # Initialize weights
+        self._init_weights()
+
+    def forward(self, x, timesteps, context=None, class_labels=None, **kwargs):
+        """
+        Forward pass through DiT.
+
+        Args:
+            x: Noisy latents [B, C, H, W]
+            timesteps: Diffusion timesteps [B]
+            context: Text conditioning [B, S, D] (optional)
+            class_labels: Class labels [B] (optional)
+
+        Returns:
+            Model prediction [B, C, H, W]
+        """
+        B, C, H, W = x.shape
+
+        # Patchify input
+        x = self.input_proj(x)  # [B, hidden_size, H/p, W/p]
+        x = x.flatten(2).transpose(1, 2)  # [B, N, hidden_size]
+
+        # Add positional embedding
+        x = x + self.pos_embed
+
+        # Embed timesteps
+        t_emb = self.time_embed(self._timestep_embedding(timesteps))  # [B, hidden_size]
+
+        # Embed class labels (if provided)
+        if class_labels is not None and self.config.num_classes > 0:
+            c_emb = self.class_embed(class_labels)  # [B, hidden_size]
+            # Combine with timestep embedding
+            cond = t_emb + c_emb
+        else:
+            cond = t_emb
+
+        # Transformer blocks
+        for block in self.blocks:
+            x = block(x, cond, context)
+
+        # Output projection
+        x = self.output_proj(x)  # [B, N, p^2 * C]
+
+        # Unpatchify
+        x = self._unpatchify(x, H, W)  # [B, C, H, W]
+
+        return x
+
+            target: Ground truth target [B, C, H, W]
+
+        Returns:
+            Loss scalar
+        """
+        # Simple MSE loss (can be extended)
+        loss = nn.functional.mse_loss(model_output, target)
+        return loss
+
+    def _timestep_embedding(self, timesteps):
+        """Create sinusoidal timestep embeddings."""
+        # Standard sinusoidal embedding
+        half_dim = self.config.hidden_size // 2
+        emb = torch.exp(
+            -torch.arange(half_dim, device=timesteps.device) *
+            (torch.log(torch.tensor(10000.0)) / (half_dim - 1))
+        )
+        emb = timesteps[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        return emb
+
+    def _unpatchify(self, x, H, W):
+        """Convert patched tensor back to image."""
+        p = self.config.patch_size
+        h = H // p
+        w = W // p
+        x = x.reshape(x.shape[0], h, w, p, p, self.config.out_channels)
+        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
+        x = x.reshape(x.shape[0], self.config.out_channels, H, W)
+        return x
+
+    def _init_weights(self):
+        """Initialize model weights."""
+        # Standard initialization (customize as needed)
+        pass
+
+
+class DiTBlock(nn.Module):
+    """DiT transformer block with adaptive LayerNorm."""
+
+    def __init__(self, config):
+        super().__init__()
+        # Implementation details...
+        pass
+
+    def forward(self, x, cond, context=None):
+        # Block forward pass
+        pass
+```
+
+### Step 4: Add model-specific layers (if needed)
+
+If your model has unique layers not shared with other models, add them to `layers.py`.
+
+### Step 5: Export model
+
+**File**: `__init__.py`
+
+```python
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""DiT model implementation."""
+
+from .config import DiTConfig
+from .model import DiT
+
+__all__ = [
+    'DiTConfig',
+    'DiT',
+]
+```
+
+### Step 6: Create configuration files
+
+**File**: `primus/configs/models/megatron/diffusion/dit_xl_2.yaml`
+
+```yaml
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# DiT-XL/2 Configuration
+
+model_type: dit
+
+# Architecture: Layers
+num_layers: 28
+
+# Architecture: Dimensions
+hidden_size: 1152
+num_attention_heads: 16
+
+# Input/Output Channels
+in_channels: 4        # VAE latent channels
+out_channels: 4
+
+# Patchification
+patch_size: 2
+
+# Class Conditioning
+num_classes: 1000     # ImageNet classes
+class_dropout_prob: 0.1
+
+# Adaptive LayerNorm
+use_adaptive_layernorm: true
+
+# Precision Settings
+bf16: true
+fp16: false
+
+# Training
+seq_length: 4096
+micro_batch_size: 8
+global_batch_size: 256
+learning_rate: 1.0e-4
+```
+
+### Step 7: Write tests
+
+**File**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py`
+
+```python
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""Unit tests for DiT model."""
+
+import pytest
+import torch
+from primus.backends.megatron.core.models.diffusion.dit import DiT, DiTConfig
+
+
+class TestDiTConfig:
+    """Tests for DiT configuration."""
+
+    def test_dit_xl_2_factory(self):
+        """Test DiT-XL/2 configuration factory method."""
+        config = DiTConfig.dit_xl_2()
+
+        assert config.model_type == "dit"
+        assert config.num_layers == 28
+        assert config.hidden_size == 1152
+        assert config.num_attention_heads == 16
+        assert config.patch_size == 2
+
+    def test_dit_config_validation(self):
+        """Test configuration validation."""
+        config = DiTConfig.dit_xl_2()
+        config.validate()  # Should not raise
+
+        # Invalid configuration
+        with pytest.raises(ValueError):
+            config = DiTConfig(num_layers=-1)
+            config.validate()
+
+
+class TestDiTModel:
+    """Tests for DiT model."""
+
+    def test_dit_initialization(self):
+        """Test DiT model initialization."""
+        config = DiTConfig.dit_xl_2()
+        model = DiT(config)
+
+        assert model is not None
+        assert isinstance(model, DiffusionModule)
+
+    def test_dit_forward_shapes(self):
+        """Test forward pass produces correct output shapes."""
+        config = DiTConfig.dit_xl_2()
+        model = DiT(config)
+
+        batch_size = 2
+        latents = torch.randn(batch_size, 4, 32, 32)
+        timesteps = torch.rand(batch_size)
+        class_labels = torch.randint(0, 1000, (batch_size,))
+
+        output = model(latents, timesteps, class_labels=class_labels)
+
+        assert output.shape == latents.shape
+
+### Loss Computation
+
+Use standalone loss functions from `loss_computation.py` instead of implementing loss as a method:
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+
+# In your forward_step_func
+target = noise - clean_latents
+loss = compute_flow_matching_loss(prediction, clean_latents, noise)
+```
+
+Models do NOT implement loss as a method. Loss computation is:
+- Separate from model architecture
+- Reusable across models
+- Testable independently
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
+```
+
+### Step 8: Update documentation
+
+1. Add model to `README.md` supported models list
+2. Update `architecture_overview.md` with model-specific details
+3. Create model-specific training guide (e.g., `dit_training.md`)
+
+### Step 9: Add example scripts
+
+**File**: Use `examples/run_pretrain.sh` with appropriate config
+
+```python
+#!/usr/bin/env python3
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""Training script for DiT model."""
+
+import argparse
+import yaml
+
+from primus.backends.megatron.core.models.diffusion.dit import DiT, DiTConfig
+from primus.backends.megatron.training.diffusion.schedulers import DDPMScheduler
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+# ... other imports
+
+
+def main():
+    # Use examples/run_pretrain.sh with config from examples/megatron/configs/MI300X/diffusion/
+    # MegatronDataloaderWrapper wraps an existing iterable (from dataset provider):
+    #   dataloader = MegatronDataloaderWrapper(energon_loader_or_pytorch_loader)
+    # ...
+
+
+if __name__ == "__main__":
+    main()
+```
+
+---
+
+## Example: Adding DiT
+
+See the complete example in the step-by-step guide above.
+
+**Key Files Created**:
+1. `core/models/diffusion/dit/config.py` - DiTConfig
+2. `core/models/diffusion/dit/model.py` - DiT model
+3. `configs/models/megatron/diffusion/dit_xl_2.yaml` - Config file
+4. `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` - Tests
+5. `examples/run_pretrain.sh` - Use with config from `examples/megatron/configs/MI300X/diffusion/`
+
+---
+
+## Testing your model
+
+### Unit tests
+
+Run tests to verify implementation:
+
+```bash
+# Run all DiT tests
+pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py -v
+
+# Run specific test
+pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py::TestDiTModel::test_dit_forward_shapes -v
+```
+
+### Integration tests
+
+Test with actual data:
+
+```bash
+# Small dataset test
+./examples/run_pretrain.sh --config examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml
+```
+
+### Validation
+
+Compare with reference implementation:
+1. Load reference weights
+2. Run same inputs through both models
+3. Compare outputs (should match within tolerance)
+
+---
+
+## Best practices
+
+### 1. Code organization
+- ✅ Separate configuration from model code
+- ✅ Reuse shared components from `common/`
+- ✅ Keep model-specific code minimal
+- ✅ Follow existing naming conventions
+
+### 2. Configuration
+- ✅ Extend `BaseDiffusionConfig`
+- ✅ Add factory methods for common sizes
+- ✅ Implement validation
+- ✅ Document all parameters
+
+### 3. Model implementation
+- ✅ Extend `DiffusionModule` from `primus.backends.megatron.core.models.common.diffusion_module.diffusion_module`
+- ✅ Implement required method: `forward()`
+- ✅ Use standalone loss functions from `loss_computation.py`
+- ✅ Add comprehensive docstrings
+- ✅ Use type hints
+- ✅ Include `pg_collection` parameter in `__init__` for distributed training
+
+### 4. Testing
+- ✅ Test configuration validation
+- ✅ Test model initialization
+- ✅ Test forward pass shapes
+- ✅ Test loss computation
+- ✅ Test with mock data first
+
+### 5. Documentation
+- ✅ Update README with new model
+- ✅ Document architecture specifics
+- ✅ Provide usage examples
+- ✅ Reference original paper
+
+### 6. Performance
+- ✅ Profile memory usage
+- ✅ Optimize critical paths
+- ✅ Support mixed precision
+- ✅ Enable gradient checkpointing
+
+---
+
+## Common pitfalls
+
+### 1. Import errors
+❌ **Wrong**: Absolute imports
+```python
+from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import DiffusionModule
+```
+
+✅ **Right**: Relative imports
+```python
+from ...common.diffusion_module.diffusion_module import DiffusionModule
+```
+
+### 2. Configuration validation
+❌ **Wrong**: No validation
+```python
+class DiTConfig(BaseDiffusionConfig):
+    pass  # No validation
+```
+
+✅ **Right**: Validate parameters
+```python
+def validate(self):
+    super().validate()
+    if self.num_layers <= 0:
+        raise ValueError(f"num_layers must be positive")
+```
+
+### 3. Shape mismatches
+❌ **Wrong**: Assuming fixed shapes
+```python
+def forward(self, x):
+    # Assumes x is always [B, 4, 32, 32]
+    pass
+```
+
+✅ **Right**: Handle variable shapes
+```python
+def forward(self, x):
+    B, C, H, W = x.shape
+    # Handle any valid shape
+    pass
+```
+
+### 4. Testing
+❌ **Wrong**: No tests
+```python
+# Just implement and hope it works
+```
+
+✅ **Right**: Comprehensive tests
+```python
+def test_forward_shapes(self):
+    # Test various input shapes
+    pass
+```
+
+---
+
+## Checklist
+
+Before submitting your new model:
+
+- [ ] Configuration class implemented and validated
+- [ ] Model class extends `DiffusionModule`
+- [ ] Required method implemented: `forward()`
+- [ ] Loss computation uses standalone functions from `loss_computation.py`
+- [ ] Process group support added (`pg_collection` parameter)
+- [ ] YAML configuration files created
+- [ ] Unit tests written and passing
+- [ ] Integration tests run successfully
+- [ ] Documentation updated
+- [ ] Example training script provided
+- [ ] Code follows Primus style guide
+- [ ] No linter errors
+- [ ] PR description includes architecture details
+
+---
+
+## Getting help
+
+If you encounter issues:
+1. Review existing models (Flux) for patterns
+2. Check documentation in `docs/04-technical-guides/diffusion-models/`
+3. Run tests in debug mode: `pytest --pdb`
+4. Consult architecture overview for design principles
+
+---
+
+**Last Updated**: December 2025
diff --git a/docs/04-technical-guides/diffusion-models/api_reference.md b/docs/04-technical-guides/diffusion-models/api_reference.md
new file mode 100644
index 000000000..c08525eb2
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/api_reference.md
@@ -0,0 +1,1059 @@
+# Flux model API reference
+
+## Overview
+
+This document provides comprehensive API reference for the Flux diffusion model implementation in Primus. Flux is a flow-based diffusion model that uses MMDiT (Multimodal Diffusion Transformer) architecture for high-quality text-to-image generation.
+
+---
+
+## Base classes
+
+### DiffusionModule
+
+**Location**: `primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py`
+
+Base class for all diffusion models, providing Megatron-Core integration.
+
+```python
+from primus.backends.megatron.core.models.common.diffusion_module import DiffusionModule
+```
+
+**Key Features**:
+- Process group management (TP, PP, CP, DP)
+- Attention backend configuration
+- Distributed checkpointing support
+- Common loss computation utilities
+
+**Inherited Methods** (available to all diffusion models):
+- `get_num_params()` - Count trainable/total parameters
+- `set_requires_grad()` - Freeze/unfreeze model
+- `compute_diffusion_loss()` - Common loss helper (MSE, MAE, Huber)
+- `sharded_state_dict()` - Distributed checkpointing
+
+> Note: in-model encoder loading (`load_encoders()`, `get_encoder_by_type()`,
+> `get_encoders_by_type()`) is not implemented and raises `NotImplementedError`.
+> Encode VAE/T5/CLIP inputs via the offline diffusion preprocessing pipeline
+> (`primus.backends.megatron.data.diffusion.preprocessing`) instead.
+
+---
+
+## Model architecture
+
+### Flux class
+
+**Location**: `primus/backends/megatron/core/models/diffusion/flux/model.py`
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+
+# Create model
+config = FluxConfig.flux_535m()
+model = Flux(config)
+```
+
+#### Architecture diagram
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                          Flux Model                              │
+├─────────────────────────────────────────────────────────────────┤
+│                                                                   │
+│  Input Processing:                                               │
+│  ┌────────────┐        ┌────────────┐                           │
+│  │ Image      │──┐  ┌──│ Text       │                           │
+│  │ Latents    │  │  │  │ Embeddings │                           │
+│  │ [B,64,H,W] │  │  │  │ [B,S,4096] │                           │
+│  └────────────┘  │  │  └────────────┘                           │
+│                  ▼  ▼                                            │
+│            ┌──────────────┐                                      │
+│            │ Linear Embed │                                      │
+│            └──────┬───────┘                                      │
+│                   │ [B,seq,3072]                                 │
+│                   ▼                                              │
+│            ┌──────────────┐                                      │
+│            │ 3D RoPE      │                                      │
+│            │ Position Emb │                                      │
+│            └──────┬───────┘                                      │
+│                   │                                              │
+│  Conditioning:    │                                              │
+│  ┌────────────┐  │                                              │
+│  │ Timestep   │──┼──┐                                           │
+│  │ Embedding  │  │  │                                           │
+│  └────────────┘  │  │                                           │
+│  ┌────────────┐  │  │                                           │
+│  │ CLIP       │──┼──┤                                           │
+│  │ Pooled     │  │  │ vec_emb [B,3072]                         │
+│  └────────────┘  │  │                                           │
+│  ┌────────────┐  │  │                                           │
+│  │ Guidance   │──┼──┘                                           │
+│  │ (optional) │  │                                              │
+│  └────────────┘  │                                              │
+│                   │                                              │
+│  Double Blocks    │                                              │
+│  (Joint):         │                                              │
+│  ┌────────────────▼──────────┐                                  │
+│  │   MMDiTLayer x N_joint    │ N_joint = 1 (535M), 19 (12B)    │
+│  │  ┌──────────────────────┐ │                                  │
+│  │  │ Joint Self-Attention │ │ (image + text together)         │
+│  │  └──────────────────────┘ │                                  │
+│  │  ┌──────────────────────┐ │                                  │
+│  │  │ Image MLP            │ │                                  │
+│  │  └──────────────────────┘ │                                  │
+│  │  ┌──────────────────────┐ │                                  │
+│  │  │ Text MLP             │ │                                  │
+│  │  └──────────────────────┘ │                                  │
+│  └────────────┬──────────────┘                                  │
+│               │                                                  │
+│  Single Blocks│                                                  │
+│  (Combined):  │                                                  │
+│  ┌────────────▼──────────────┐                                  │
+│  │ FluxSingleTransformer     │ N_single = 1 (535M), 38 (12B)   │
+│  │       Block x N_single    │                                  │
+│  │  ┌──────────────────────┐ │                                  │
+│  │  │ Self-Attention       │ │ (image + text concatenated)     │
+│  │  └──────────────────────┘ │                                  │
+│  │  ┌──────────────────────┐ │                                  │
+│  │  │ MLP                  │ │                                  │
+│  │  └──────────────────────┘ │                                  │
+│  └────────────┬──────────────┘                                  │
+│               │ (extract image tokens)                           │
+│               ▼                                                  │
+│  Output Processing:                                              │
+│  ┌────────────────────────┐                                     │
+│  │ AdaLNContinuous        │                                     │
+│  │ (timestep conditioned) │                                     │
+│  └────────────┬───────────┘                                     │
+│               ▼                                                  │
+│  ┌────────────────────────┐                                     │
+│  │ Linear Projection      │                                     │
+│  └────────────┬───────────┘                                     │
+│               ▼                                                  │
+│         [B,64,H,W]                                               │
+│       Predicted Velocity                                         │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+#### Parameter counts
+
+| Variant | Joint Layers | Single Layers | Total Parameters | Use Case |
+|---------|-------------|---------------|------------------|----------|
+| Flux 535M | 1 | 1 | ~535 million | Testing, debugging, prototyping |
+| Flux 12B | 19 | 38 | ~12 billion | Production training |
+
+---
+
+## Configuration
+
+### FluxConfig
+
+**Location**: `primus/backends/megatron/core/models/diffusion/flux/config.py`
+
+Complete configuration class for Flux models, inheriting from `BaseDiffusionConfig`.
+
+#### Key parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `num_joint_layers` | int | 19 | Number of joint (MMDiT) transformer layers |
+| `num_single_layers` | int | 38 | Number of single transformer layers |
+| `hidden_size` | int | 3072 | Hidden dimension size |
+| `num_attention_heads` | int | 24 | Number of attention heads |
+| `in_channels` | int | 64 | Input channels (VAE latent dimension) |
+| `context_dim` | int | 4096 | Text context dimension (T5-XXL) |
+| `vec_in_dim` | int | 768 | Vector input dimension (CLIP pooled) |
+| `model_channels` | int | 256 | Channels for timestep embedding |
+| `guidance_embed` | bool | False | Enable guidance embedding for CFG |
+| `guidance_scale` | float | 3.5 | Guidance scale for classifier-free guidance |
+| `theta` | int | 10000 | Base for RoPE frequency computation |
+| `axes_dim` | tuple | (16, 56, 56) | Dimensions for 3D RoPE axes |
+| `patch_size` | int | 1 | Patch size for image tokens |
+| `add_qkv_bias` | bool | True | Add bias to QKV projections |
+| `rotary_interleaved` | bool | True | Interleave RoPE dimensions |
+| `layernorm_epsilon` | float | 1e-6 | Epsilon for layer normalization |
+| `hidden_dropout` | float | 0.0 | Hidden layer dropout rate |
+| `attention_dropout` | float | 0.0 | Attention dropout rate |
+
+#### Configuration examples
+
+**Flux 535M (Testing)**:
+```python
+config = FluxConfig.flux_535m()
+# Equivalent to:
+config = FluxConfig(
+    num_joint_layers=1,
+    num_single_layers=1,
+    hidden_size=3072,
+    num_attention_heads=24,
+)
+```
+
+**Flux 12B (Production)**:
+```python
+config = FluxConfig.flux_12b()
+# Equivalent to:
+config = FluxConfig(
+    num_joint_layers=19,
+    num_single_layers=38,
+    hidden_size=3072,
+    num_attention_heads=24,
+)
+```
+
+**Custom Configuration**:
+```python
+config = FluxConfig(
+    num_joint_layers=4,
+    num_single_layers=8,
+    hidden_size=2048,
+    num_attention_heads=16,
+    guidance_embed=True,
+    patch_size=2,
+)
+```
+
+---
+
+## Components API
+
+### Embeddings
+
+#### TimeStepEmbedder
+
+**Location**: `primus/backends/megatron/core/models/diffusion/common/embeddings.py`
+
+Converts scalar timesteps to high-dimensional embeddings using sinusoidal encoding.
+
+```python
+from primus.backends.megatron.core.models.diffusion.common.embeddings import TimeStepEmbedder
+
+embedder = TimeStepEmbedder(embedding_dim=256, hidden_dim=3072)
+timesteps = torch.tensor([0, 100, 500, 999])  # [B]
+t_emb = embedder(timesteps)  # [B, 3072]
+```
+
+**Input**: Timesteps [B] in range [0, 1000]
+**Output**: Embeddings [B, hidden_dim]
+
+#### MLPEmbedder
+
+Embeds vector conditioning (e.g., CLIP pooled embeddings) via 2-layer MLP.
+
+```python
+from primus.backends.megatron.core.models.diffusion.common.embeddings import MLPEmbedder
+
+embedder = MLPEmbedder(in_dim=768, hidden_dim=3072)
+clip_pooled = torch.randn(4, 768)  # [B, 768]
+embedded = embedder(clip_pooled)  # [B, 3072]
+```
+
+**Input**: Vectors [B, in_dim]
+**Output**: Embeddings [B, hidden_dim]
+
+---
+
+### Normalization
+
+#### AdaLN (adaptive layer normalization)
+
+**Location**: `primus/backends/megatron/core/models/diffusion/common/normalization.py`
+
+Applies layer normalization conditioned on timestep embeddings.
+
+```python
+from megatron.core.transformer.transformer_config import TransformerConfig
+from primus.backends.megatron.core.models.diffusion.common.normalization import AdaLN
+
+config = TransformerConfig(hidden_size=3072)
+adaln = AdaLN(config, n_adaln_chunks=6)
+
+timestep_emb = torch.randn(4, 3072)
+shift, scale, gate, shift_mlp, scale_mlp, gate_mlp = adaln(timestep_emb)
+# Each output: [B, 3072]
+```
+
+**Methods**:
+- `forward(timestep_emb)`: Generate modulation parameters
+- `modulate(x, shift, scale)`: Apply adaptive modulation
+- `scale_add(residual, x, gate)`: Gated residual addition
+
+#### AdaLNContinuous
+
+Continuous variant of AdaLN for Flux output normalization.
+
+```python
+from primus.backends.megatron.core.models.diffusion.common.normalization import AdaLNContinuous
+
+adaln = AdaLNContinuous(config, conditioning_embedding_dim=3072)
+x = torch.randn(4, 256, 3072)  # [B, seq, hidden]
+cond = torch.randn(4, 3072)  # [B, cond_dim]
+x_norm = adaln(x, cond)  # [B, seq, hidden]
+```
+
+#### RMSNorm
+
+Root Mean Square Layer Normalization (simpler, faster than LayerNorm).
+
+```python
+from primus.backends.megatron.core.models.diffusion.common.normalization import RMSNorm
+
+norm = RMSNorm(hidden_size=3072)
+x = torch.randn(4, 256, 3072)
+x_norm = norm(x)
+```
+
+---
+
+### Position embeddings
+
+#### EmbedND (3D RoPE)
+
+**Location**: `primus/backends/megatron/core/models/diffusion/flux/layers.py`
+
+Multi-dimensional Rotary Position Embedding for image patches.
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.layers import (
+    EmbedND,
+)
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+)
+
+# Initialize
+embed_nd = EmbedND(dim=3072, theta=10000, axes_dim=[16, 56, 56])
+
+# Generate position IDs for 56x56 image patches
+batch_size = 2
+height, width = 112, 112  # Unpacked dimensions (56*2, 56*2)
+img_ids = generate_image_position_ids(batch_size, height, width)
+# img_ids: [B, H*W/4, 3] where dimension 0 is always 0
+
+# Get RoPE frequencies
+rope_freqs = embed_nd(img_ids)  # [3, B, H*W, 3072]
+```
+
+**Axes**:
+- Axis 0: Channel groups (16 for 64 channels)
+- Axis 1: Height positions (56 for 1024px image)
+- Axis 2: Width positions (56 for 1024px image)
+
+---
+
+### Attention mechanisms
+
+#### JointSelfAttention
+
+**Location**: `primus/backends/megatron/core/models/diffusion/flux/attention.py`
+
+Joint attention over image and text tokens (MMDiT architecture).
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.attention import (
+    JointSelfAttention,
+    JointSelfAttentionSubmodules,
+)
+
+submodules = JointSelfAttentionSubmodules(...)
+joint_attn = JointSelfAttention(config, submodules, layer_number=0)
+
+# Forward
+img_tokens = torch.randn(3136, 2, 3072)  # [seq_img, B, hidden]
+txt_tokens = torch.randn(512, 2, 3072)   # [seq_txt, B, hidden]
+img_out, txt_out = joint_attn(
+    img_tokens,
+    attention_mask=None,
+    additional_hidden_states=txt_tokens,
+)
+```
+
+**Input**:
+- `hidden_states`: Image tokens [seq_img, B, hidden]
+- `additional_hidden_states`: Text tokens [seq_txt, B, hidden]
+
+**Output**: Tuple of (img_output, txt_output)
+
+#### FluxSingleAttention
+
+Single-stream self-attention for image tokens only.
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.attention import FluxSingleAttention
+
+single_attn = FluxSingleAttention(config, submodules, layer_number=0)
+
+img_tokens = torch.randn(3136, 2, 3072)
+output = single_attn(img_tokens, attention_mask=None)
+```
+
+---
+
+### Layer specifications
+
+#### MMDiTLayer
+
+**Location**: `primus/backends/megatron/core/models/diffusion/flux/layer_spec.py`
+
+Joint image-text transformer block.
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    MMDiTLayer,
+    get_flux_double_transformer_spec_for_backend,
+)
+
+# Using factory function (recommended)
+spec = get_flux_double_transformer_spec_for_backend(backend)
+mmdit_layer = MMDiTLayer(
+    config=config,
+    submodules=spec.submodules,
+    layer_number=0,
+)
+
+# Forward
+img_tokens = torch.randn(3136, 2, 3072)
+txt_tokens = torch.randn(512, 2, 3072)
+emb = torch.randn(2, 3072)
+img_out, txt_out = mmdit_layer(img_tokens, txt_tokens, emb=emb)
+```
+
+#### FluxSingleTransformerBlock
+
+Image-only transformer block.
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    FluxSingleTransformerBlock,
+    get_flux_single_transformer_spec_for_backend,
+)
+
+spec = get_flux_single_transformer_spec_for_backend(backend)
+single_block = FluxSingleTransformerBlock(
+    config=config,
+    submodules=spec.submodules,
+    layer_number=0,
+)
+
+# Forward
+img_tokens = torch.randn(3136, 2, 3072)
+emb = torch.randn(2, 3072)
+output, _ = single_block(img_tokens, emb=emb)
+```
+
+---
+
+## Training utilities
+
+### Noise application
+
+**Location**: `primus/backends/megatron/training/diffusion/noise_utils.py`
+
+Pure functions for applying noise according to different diffusion forward processes.
+
+#### apply_flow_matching_noise()
+
+```python
+from primus.backends.megatron.training.diffusion.noise_utils import apply_flow_matching_noise
+
+clean_latents = torch.randn(2, 16, 64, 64)
+noise = torch.randn(2, 16, 64, 64)
+sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1)
+
+noisy = apply_flow_matching_noise(clean_latents, noise, sigma)
+# Formula: noisy = (1 - sigma) * clean + sigma * noise
+```
+
+**Parameters**:
+- `clean_latents` (Tensor): Clean latents [any shape]
+- `noise` (Tensor): Sampled noise [same shape as clean_latents]
+- `sigma` (Tensor): Noise schedule values [broadcast compatible], range [0, 1]
+
+**Returns**: Noisy latents [same shape as clean_latents]
+
+**Reference**: [Flow Matching for Generative Modeling](https://arxiv.org/abs/2210.02747)
+
+#### apply_ddpm_noise()
+
+```python
+from primus.backends.megatron.training.diffusion.noise_utils import apply_ddpm_noise
+
+clean = torch.randn(2, 3, 256, 256)
+noise = torch.randn(2, 3, 256, 256)
+alpha_bar = torch.tensor([0.9, 0.5]).reshape(2, 1, 1, 1)
+
+noisy = apply_ddpm_noise(clean, noise, alpha_bar)
+# Formula: noisy = sqrt(alpha_bar) * clean + sqrt(1 - alpha_bar) * noise
+```
+
+**Parameters**:
+- `clean_latents` (Tensor): Clean latents
+- `noise` (Tensor): Sampled noise (same shape)
+- `alpha_bar` (Tensor): Cumulative product of alphas, range (0, 1]
+
+**Returns**: Noisy latents
+
+**Reference**: [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
+
+---
+
+### Loss computation
+
+**Location**: `primus/backends/megatron/training/diffusion/loss_computation.py`
+
+Reusable loss computation logic for different diffusion training objectives.
+
+#### compute_flow_matching_loss()
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+
+prediction = torch.randn(2, 16, 64, 64)  # Model output
+clean = torch.randn(2, 16, 64, 64)
+noise = torch.randn(2, 16, 64, 64)
+
+loss = compute_flow_matching_loss(prediction, clean, noise)
+# Formula: target = noise - clean
+#          loss = MSE(prediction, target)
+```
+
+**Parameters**:
+- `prediction` (Tensor): Model output (predicted velocity) [any shape]
+- `clean_latents` (Tensor): Original clean latents [same shape]
+- `noise` (Tensor): Sampled noise [same shape]
+
+**Returns**: Scalar loss value (mean squared error)
+
+**Used by**: Flux, SD3, video models
+
+#### compute_epsilon_loss()
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_epsilon_loss
+
+prediction = torch.randn(2, 3, 256, 256)
+noise = torch.randn(2, 3, 256, 256)
+
+loss = compute_epsilon_loss(prediction, noise)
+# Formula: loss = MSE(prediction, noise)
+```
+
+**Parameters**:
+- `prediction` (Tensor): Model output (predicted noise)
+- `noise` (Tensor): Sampled noise (ground truth)
+
+**Returns**: Scalar loss value
+
+**Used by**: DDPM and older models
+
+#### compute_v_prediction_loss()
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_v_prediction_loss
+
+prediction = torch.randn(2, 16, 64, 64)
+clean = torch.randn(2, 16, 64, 64)
+noise = torch.randn(2, 16, 64, 64)
+sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1)
+
+loss = compute_v_prediction_loss(prediction, clean, noise, sigma)
+# Formula: v = sigma * noise - (1 - sigma) * clean
+#          loss = MSE(prediction, v)
+```
+
+**Parameters**:
+- `prediction` (Tensor): Model output
+- `clean_latents` (Tensor): Clean latents
+- `noise` (Tensor): Sampled noise
+- `sigma` (Tensor): Noise schedule values [broadcast compatible]
+
+**Returns**: Scalar loss value
+
+**Reference**: [Progressive Distillation for Fast Sampling](https://arxiv.org/abs/2202.00512)
+
+---
+
+### Timestep sampling
+
+**Location**: `primus/backends/megatron/training/diffusion/timestep_sampling.py`
+
+Sampling strategies for training timesteps (hyperparameter optimization, separate from inference).
+
+#### LogitNormalSampler
+
+```python
+from primus.backends.megatron.training.diffusion.timestep_sampling import LogitNormalSampler
+
+sampler = LogitNormalSampler(mean=0.0, std=1.0)
+timesteps, sigmas = sampler.sample(
+    batch_size=32,
+    device='cuda',
+    scheduler=flow_scheduler
+)
+```
+
+**Description**: Logit-normal distribution sampling, emphasizes boundary timesteps (t≈0 and t≈1000).
+
+**Parameters**:
+- `mean` (float): Mean of normal distribution (default: 0.0)
+- `std` (float): Standard deviation (default: 1.0)
+
+**Returns**: Tuple of (timesteps [B], sigmas [B])
+
+**Reference**: [Stable Diffusion 3](https://arxiv.org/abs/2403.03206v1), Section 3.1
+
+**Used by**: Flux, SD3
+
+#### UniformSampler
+
+```python
+from primus.backends.megatron.training.diffusion.timestep_sampling import UniformSampler
+
+sampler = UniformSampler()
+timesteps, sigmas = sampler.sample(batch_size=32, device='cuda', scheduler=flow_scheduler)
+```
+
+**Description**: Uniform timestep sampling (baseline approach for comparison).
+
+**Returns**: Tuple of (timesteps [B], sigmas [B])
+
+#### ModeSampler
+
+```python
+from primus.backends.megatron.training.diffusion.timestep_sampling import ModeSampler
+
+sampler = ModeSampler(mode_scale=1.29)
+timesteps, sigmas = sampler.sample(batch_size=32, device='cuda', scheduler=flow_scheduler)
+```
+
+**Description**: Mode-based sampling from SD3 paper (alternative to logit-normal).
+
+**Parameters**:
+- `mode_scale` (float): Scaling factor (default: 1.29 from SD3 paper)
+
+**Returns**: Tuple of (timesteps [B], sigmas [B])
+
+#### create_timestep_sampler()
+
+```python
+from primus.backends.megatron.training.diffusion.timestep_sampling import create_timestep_sampler
+
+# Factory function for easy experimentation
+sampler = create_timestep_sampler("logit_normal", mean=0.0, std=1.0)
+sampler = create_timestep_sampler("uniform")
+sampler = create_timestep_sampler("mode", mode_scale=1.5)
+```
+
+**Parameters**:
+- `strategy` (str): Sampling strategy ("logit_normal", "uniform", "mode")
+- `**kwargs`: Additional arguments for the sampler
+
+**Returns**: TimestepSampler instance
+
+---
+
+### Training workflow example
+
+Complete training step using the utilities:
+
+```python
+import torch
+from torch.optim import AdamW
+from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig
+from primus.backends.megatron.training.diffusion.noise_utils import apply_flow_matching_noise
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+from primus.backends.megatron.training.diffusion.timestep_sampling import LogitNormalSampler
+from primus.backends.megatron.training.diffusion.schedulers.flow_match_euler import (
+    FlowMatchEulerDiscreteScheduler
+)
+
+# Setup
+config = FluxConfig.flux_535m()
+model = Flux(config).cuda()
+optimizer = AdamW(model.parameters(), lr=1e-4)
+
+# Initialize scheduler and timestep sampler
+scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
+timestep_sampler = LogitNormalSampler(mean=0.0, std=1.0)
+
+# Training loop
+for batch in dataloader:
+    clean_latents = batch['latents'].cuda()  # [B, 16, 64, 64]
+    txt_embeddings = batch['text'].cuda()     # [B, 512, 4096]
+    clip_pooled = batch['clip'].cuda()        # [B, 768]
+
+    batch_size = clean_latents.shape[0]
+
+    # 1. Sample timesteps
+    timesteps, sigmas = timestep_sampler.sample(
+        batch_size=batch_size,
+        device='cuda',
+        scheduler=scheduler
+    )
+
+    # 2. Sample noise
+    noise = torch.randn_like(clean_latents)
+
+    # 3. Apply noise
+    sigma_reshaped = sigmas.view(-1, 1, 1, 1)
+    noisy_latents = apply_flow_matching_noise(clean_latents, noise, sigma_reshaped)
+
+    # 4. Prepare position IDs
+    img_ids = generate_image_position_ids(batch_size, 128, 128).cuda()
+    txt_ids = torch.zeros(batch_size, 512, 3).cuda()
+
+    # 5. Forward pass
+    predicted_velocity = model(
+        img=noisy_latents,
+        txt=txt_embeddings,
+        y=clip_pooled,
+        timesteps=sigmas,
+        img_ids=img_ids,
+        txt_ids=txt_ids,
+    )
+
+    # 6. Compute loss
+    loss = compute_flow_matching_loss(predicted_velocity, clean_latents, noise)
+
+    # 7. Backward and optimize
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    print(f"Step {step}, Loss: {loss.item():.4f}")
+```
+
+---
+
+## Usage examples
+
+### Basic inference
+
+```python
+import torch
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.utils import generate_image_position_ids
+
+# 1. Setup model
+config = FluxConfig.flux_535m()
+model = Flux(config)
+model.eval()
+
+# 2. Prepare inputs
+batch_size = 1
+img_latents = torch.randn(batch_size, 64, 128, 128)  # VAE latents
+txt_embeddings = torch.randn(batch_size, 512, 4096)  # T5-XXL embeddings
+clip_pooled = torch.randn(batch_size, 768)  # CLIP-L pooled
+timesteps = torch.tensor([0.5])  # Diffusion timestep [0, 1]
+
+# 3. Generate position IDs
+img_ids = generate_image_position_ids(batch_size, 256, 256)
+txt_ids = torch.zeros(batch_size, 512, 3)
+
+# 4. Forward pass
+with torch.no_grad():
+    predicted_velocity = model(
+        img=img_latents,
+        txt=txt_embeddings,
+        y=clip_pooled,
+        timesteps=timesteps,
+        img_ids=img_ids,
+        txt_ids=txt_ids,
+    )
+
+print(f"Output shape: {predicted_velocity.shape}")  # [1, 64, 128, 128]
+```
+
+### Training step
+
+```python
+import torch
+from torch.optim import AdamW
+
+# Setup
+config = FluxConfig.flux_535m()
+model = Flux(config)
+model.train()
+
+optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
+
+# Prepare batch
+batch_size = 4
+img = torch.randn(batch_size, 64, 128, 128)
+txt = torch.randn(batch_size, 512, 4096)
+y = torch.randn(batch_size, 768)
+timesteps = torch.rand(batch_size)
+
+img_ids = generate_image_position_ids(batch_size, 256, 256)
+txt_ids = torch.zeros(batch_size, 512, 3)
+
+# Add noise for flow matching
+original = img.clone()
+noise = torch.randn_like(img)
+t = timesteps.view(-1, 1, 1, 1)
+noisy_img = original + noise * t
+
+# Velocity target for flow matching
+velocity_target = noise - original
+
+# Training step
+optimizer.zero_grad()
+
+# Forward pass
+output = model(noisy_img, txt, y, timesteps, img_ids, txt_ids)
+
+# Loss computation (using standalone function)
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+target = noise - clean_latents
+loss = compute_flow_matching_loss(output, clean_latents, noise)
+
+# Backward pass
+loss.backward()
+
+# Optimizer step
+optimizer.step()
+
+print(f"Loss: {loss.item():.4f}")
+```
+
+### With guidance (classifier-free guidance)
+
+```python
+# Enable guidance in config
+config = FluxConfig.flux_535m(guidance_embed=True)
+model = Flux(config)
+model.eval()
+
+# Prepare inputs with guidance
+guidance_scale = torch.tensor([3.5])  # Typical guidance scale
+
+with torch.no_grad():
+    output = model(
+        img=img_latents,
+        txt=txt_embeddings,
+        y=clip_pooled,
+        timesteps=timesteps,
+        img_ids=img_ids,
+        txt_ids=txt_ids,
+        guidance=guidance_scale,  # Add guidance
+    )
+```
+
+### Different resolutions
+
+```python
+# Flux can handle different resolutions
+resolutions = [
+    (64, 64),    # 512x512 pixels
+    (128, 128),  # 1024x1024 pixels
+    (192, 192),  # 1536x1536 pixels
+]
+
+for height, width in resolutions:
+    img = torch.randn(1, 64, height, width)
+    img_ids = generate_image_position_ids(1, height, width)
+
+    with torch.no_grad():
+        output = model(img, txt, y, timesteps, img_ids, txt_ids)
+
+    assert output.shape == img.shape
+```
+
+---
+
+## Methods reference
+
+### Flux.forward()
+
+```python
+def forward(
+    self,
+    img: Tensor,                    # [B, C, H, W] Image latents from VAE
+    txt: Tensor,                    # [B, S_txt, D_txt] T5-XXL embeddings
+    y: Tensor,                      # [B, D_pool] CLIP pooled embeddings
+    timesteps: Tensor,              # [B] Timesteps in [0, 1]
+    img_ids: Tensor,                # [B, H*W, 3] Image position IDs
+    txt_ids: Tensor,                # [B, S_txt, 3] Text position IDs
+    guidance: Optional[Tensor] = None,  # [B] Guidance scale
+    controlnet_double_block_samples: Optional[Tensor] = None,
+    controlnet_single_block_samples: Optional[Tensor] = None,
+) -> Tensor:                        # Returns: [B, C, H, W] Predicted velocity
+```
+
+## Loss computation
+
+Loss is computed using standalone functions from `loss_computation.py`:
+
+### compute_flow_matching_loss()
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+
+def compute_flow_matching_loss(
+    prediction: Tensor,                 # [any shape] Model prediction
+    clean_latents: Tensor,              # [same shape] Original clean latents
+    noise: Tensor,                      # [same shape] Sampled noise
+) -> Tensor:                            # Returns: Scalar loss
+```
+
+### Flux.get_num_params()
+
+```python
+def get_num_params(
+    self,
+    trainable_only: bool = True,
+) -> int:                           # Returns: Number of parameters
+```
+
+---
+
+## Testing
+
+### Running tests
+
+```bash
+# Run all Flux tests
+pytest tests/unit_tests/backends/megatron/diffusion/ -v
+
+# Run specific test files
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py -v
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py -v
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_config.py -v
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py -v
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_model.py -v
+pytest tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py -v
+
+# Run with coverage
+pytest tests/unit_tests/backends/megatron/diffusion/ --cov=primus.backends.megatron.core.models.diffusion --cov-report=html
+```
+
+### Test coverage
+
+- **Component tests**: 80+ tests covering all components
+- **Model tests**: 17+ tests for full model
+- **Integration tests**: 11+ tests for complete workflows
+
+---
+
+## Performance considerations
+
+### Memory usage
+
+| Configuration | Model Weights | Training (bf16) | Training (fp32) |
+|--------------|---------------|-----------------|-----------------|
+| Flux 535M | ~2 GB | ~6-8 GB | ~10-12 GB |
+| Flux 12B | ~24 GB | ~60-80 GB | ~100-120 GB |
+
+### Throughput (estimated)
+
+On MI300X (192GB):
+- **Flux 535M**: ~5-10 samples/sec (depends on resolution)
+- **Flux 12B**: ~0.5-1 samples/sec (requires multi-GPU)
+
+### Optimization tips
+
+1. **Use mixed precision**: `torch.autocast(device_type='cuda', dtype=torch.bfloat16)`
+2. **Enable CUDA graphs**: Set `enable_cuda_graph=True` in config
+3. **Use Transformer Engine**: Automatically used with factory functions
+4. **Gradient checkpointing**: Can be enabled for memory savings
+
+---
+
+## Common issues
+
+### Issue: Import errors
+
+```python
+# Old import style
+from some_other_library import Flux
+
+# Correct Primus import
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+```
+
+### Issue: Position ID shape mismatch
+
+```python
+# Position IDs must be [B, seq, 3] for 3D RoPE
+img_ids = generate_image_position_ids(batch_size, height, width)
+# Not: img_ids = torch.randn(batch_size, height * width, 2)  # Wrong!
+```
+
+### Issue: Timestep range
+
+```python
+# Flux expects timesteps in [0, 1]
+timesteps = torch.rand(batch_size)  # Correct: [0, 1]
+# Not: timesteps = torch.randint(0, 1000, (batch_size,))  # Wrong range!
+```
+
+---
+
+## API compatibility
+
+### Primus architecture features
+
+| Aspect | Primus Implementation |
+|--------|----------------------|
+| Import path | `primus.backends.megatron.core.models.diffusion.flux` |
+| Base class | `DiffusionModule` (extends MegatronModule) |
+| Config parent | `BaseDiffusionConfig` (extends TransformerConfig) |
+| Layer organization | Unified `TransformerBlock` with heterogeneous specs |
+| Checkpoint format | `transformer.layers.{0-56}` unified namespace |
+| Process groups | Via `pg_collection` parameter |
+
+### Key design choices
+
+**TransformerBlock Architecture**:
+- Primus uses Megatron-Core's `TransformerBlock` with heterogeneous layer specifications
+- Unified checkpoint format for simpler distributed training
+- Note: pipeline parallelism is not supported for diffusion models (`pipeline_model_parallel_size` must be 1)
+
+**Example Usage**:
+
+```python
+# Primus native approach
+from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig
+
+# Primus
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+config = FluxConfig()
+model = Flux(config)
+```
+
+For more advanced examples, see `examples/run_pretrain.sh`.
+
+---
+
+## References
+
+### Papers
+- **Flux**: "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis"
+- **MMDiT**: "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis"
+- **Flow Matching**: "Flow Matching for Generative Modeling"
+- **RoPE**: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+- **DiT**: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023)
+
+### Source code
+- **Primus Implementation**: `primus/backends/megatron/core/models/diffusion/flux/`
+- **Megatron-Core**: `megatron/core/transformer/`
+- **Official Flux**: Black Forest Labs (HuggingFace)
+
+---
+
+## Version information
+
+- **Primus Version**: Current
+- **Megatron-Core Version**: Latest
+- **Transformer Engine**: Optional (recommended for performance)
+
+---
+
+## Support
+
+For issues or questions:
+1. Check this API reference
+2. Read architecture guide: `docs/04-technical-guides/diffusion-models/flux_architecture.md`
+3. See examples in docstrings
+4. Check test files for usage patterns
diff --git a/docs/04-technical-guides/diffusion-models/architecture_overview.md b/docs/04-technical-guides/diffusion-models/architecture_overview.md
new file mode 100644
index 000000000..8ef25764f
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/architecture_overview.md
@@ -0,0 +1,499 @@
+# Architecture overview
+
+This document provides a detailed overview of the diffusion model architecture in Primus, including design decisions, directory structure, and implementation patterns.
+
+---
+
+## Table of contents
+
+1. [Design philosophy](#design-philosophy)
+2. [Directory structure](#directory-structure)
+3. [Architectural decisions](#architectural-decisions)
+4. [Component hierarchy](#component-hierarchy)
+5. [Data flow](#data-flow)
+6. [Comparison with alternative implementations](#comparison-with-alternative-implementations)
+
+---
+
+## Design philosophy
+
+### Core principles
+
+1. **Megatron-Core Native**: Built on Megatron-Core patterns and conventions
+2. **Extensibility**: Easy to add new models (DiT, MovieGen, custom)
+3. **Clarity**: Clear separation between shared and model-specific code
+4. **Reusability**: Shared components across multiple models
+5. **Performance**: Support for precalculated data and multi-GPU training
+
+### Architectural advantages
+
+| Aspect | Primus Design Choice |
+|--------|---------------------|
+| Model Location | `core/models/diffusion/` (Megatron-Core convention) |
+| Layer Organization | Unified `TransformerBlock` with heterogeneous specs |
+| Shared Code | Dedicated `common/` directory |
+| Encoders | Hierarchical `encoders/{type}/{variant}/` |
+| Energon | Shared `data/energon/` for all models |
+| Mock Data | Synthetic providers under `data/synthetic/` |
+| Framework | Pure Megatron (no PyTorch Lightning dependency) |
+
+---
+
+## Directory structure
+
+### High-level organization
+
+```
+primus/backends/megatron/
+├── core/models/diffusion/          # Core model implementations
+├── training/diffusion/             # Training utilities (schedulers, etc.)
+└── data/
+    ├── energon/                    # Shared Energon infrastructure
+    └── diffusion/                  # Diffusion-specific data
+```
+
+### Detailed breakdown
+
+#### 1. Core models (`core/models/diffusion/`)
+
+Following Megatron-Core convention (`megatron/core/models/gpt/`, `megatron/core/models/multimodal/`):
+
+```
+core/models/diffusion/
+├── common/                         # Shared building blocks
+│   ├── __init__.py
+│   ├── config.py                   # BaseDiffusionConfig
+│   ├── embeddings.py               # TimeStepEmbedder, Timesteps, MLPEmbedder
+│   └── normalization.py            # RMSNorm, AdaLN, AdaLNContinuous
+│
+└── flux/                           # Flux-specific components
+    ├── __init__.py
+    ├── config.py                   # FluxConfig
+    ├── model.py                    # Flux (main model class)
+    ├── layer_spec.py               # MMDiTLayer, FluxSingleTransformerBlock
+    ├── attention.py                # JointSelfAttention, FluxSingleAttention
+    ├── layers.py                   # EmbedND (RoPE), embedders
+    ├── checkpoint_converter.py     # HF <-> Megatron conversion
+    └── utils.py                    # image position ids, helpers
+```
+
+**Rationale**:
+- `common/`: shared building blocks (base config, timestep/positional embeddings, normalization) reusable across diffusion models
+- `flux/`: Flux-specific code (model class, MMDiT and single-block layer specs, joint attention, RoPE embedders)
+- Clear separation keeps model-specific code isolated and shared code reusable for future models
+
+#### 2. Training (`training/diffusion/`)
+
+```
+training/diffusion/
+├── __init__.py
+└── schedulers/
+    ├── __init__.py
+    ├── base.py                     # BaseScheduler
+    ├── flow_matching.py            # FlowMatchEulerDiscreteScheduler
+    ├── ddpm.py                     # [Future] DDPMScheduler
+    ├── edm.py                      # [Future] EDMScheduler
+    └── euler.py                    # [Future] EulerDiscreteScheduler
+```
+
+**Rationale**:
+- Schedulers define noise schedules and training targets
+- Separate from model code for clarity
+- Easy to add new schedulers (DDPM, EDM, etc.)
+
+#### 3. Data pipeline (`data/`)
+
+```
+data/
+├── energon/                        # Shared Energon utilities
+│   └── __init__.py
+├── dataloader.py                   # MegatronDataloaderWrapper (wraps any iterable)
+│
+└── diffusion/
+    ├── __init__.py
+    ├── encoders/                   # Hierarchical encoder registry
+    │   ├── __init__.py
+    │   ├── registry.py             # EncoderRegistry
+    │   ├── image/
+    │   │   └── vae/
+    │   │       ├── __init__.py
+    │   │       ├── autoencoder_kl.py   # AutoencoderKL
+    │   │       └── vqvae.py            # VQVAE
+    │   └── text/
+    │       ├── t5/
+    │       │   ├── __init__.py
+    │       │   ├── t5_xxl.py           # T5-XXL
+    │       │   └── t5_large.py         # T5-Large
+    │       └── clip/
+    │           ├── __init__.py
+    │           ├── clip_l.py           # CLIP-L
+    │           └── clip_h.py           # CLIP-H
+    │
+    ├── preprocessing/
+    │   ├── __init__.py
+    │   └── image/
+    │       ├── __init__.py
+    │       ├── transforms.py           # Resizing, normalization
+    │       └── augmentation.py         # Data augmentation
+    │
+    └── task_encoders/
+        ├── __init__.py
+        └── image.py  # EncodedDiffusionTaskEncoder
+```
+
+**Rationale**:
+- **Energon shared**: VLM and other models can reuse Energon utilities
+- **Hierarchical encoders**: Organized by modality and variant type
+- **Registry pattern**: Config-driven encoder selection
+- **Preprocessing separated**: Clear pipeline stages
+
+---
+
+## Architectural decisions
+
+### Decision 1: Models under `core/models/`
+
+**Choice**: `primus/backends/megatron/core/models/diffusion/`
+**Not**: `primus/backends/megatron/models/diffusion/`
+
+**Reasoning**:
+- Aligns with Megatron-Core structure
+- Easier upstream tracking
+- Clear that these are Megatron-Core compatible
+- Consistent with existing Primus structure (`core/models/gpt/`)
+
+### Decision 2: Separate `common/` directory
+
+**Choice**: Shared components in `common/`
+**Not**: Everything in `flux/` or flat structure
+
+**Reasoning**:
+- Standard approach puts shared code in model-specific directories
+- Primus: `common/` makes it explicit what's shared
+- Future DiT implementation trivial (reuse from `common/`)
+- Clear contract: if in `common/`, must work for all models
+
+**Shared Components**:
+- `JointSelfAttention`: Used by DiT and Flux joint layers
+- `FluxSingleAttention`: Used by DiT and Flux single layers
+- `MMDiTLayer`: Joint (multimodal) transformer block
+- `FluxSingleTransformerBlock`: Single-modality transformer block
+
+**Flux-Only Components**:
+- `EmbedND`: 3D RoPE position embedding (Flux-specific)
+- `Flux` model class
+
+### Decision 3: Hierarchical encoder structure
+
+**Choice**: `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/`
+**Not**: Flat `encoders/conditioner.py`
+
+**Reasoning**:
+- Traditional approach uses flat structure with all encoders in one file
+- User requirement: support 5+ variants per modality
+- Registry pattern enables config-driven selection
+- Easy to add new encoders without modifying existing files
+
+**Registry Pattern**:
+```python
+from primus.backends.megatron.data.diffusion.encoders import get_encoder
+
+# Config-driven selection
+vae = get_encoder('autoencoder_kl', config=vae_config)
+t5 = get_encoder('t5_xxl', config=t5_config)
+clip = get_encoder('clip_l', config=clip_config)
+```
+
+### Decision 4: Shared Energon infrastructure
+
+**Choice**: `data/energon/` for shared utilities
+**Not**: Nested under `data/diffusion/`
+
+**Reasoning**:
+- Energon is general-purpose (VLM, diffusion, future models)
+- Traditional approach nests under model-specific directories
+- Megatron-LM has Energon at example level (not in core)
+- Primus approach: shared infra + model-specific TaskEncoders
+
+**Pattern**:
+```python
+# Shared: data/dataloader.py
+# MegatronDataloaderWrapper wraps any iterable (Energon loader, PyTorch DataLoader, etc.)
+wrapper = MegatronDataloaderWrapper(dataloader)
+
+# Model-specific: data/diffusion/task_encoders/image.py
+class EncodedDiffusionTaskEncoder:
+    """Diffusion-specific encoding logic"""
+    pass
+```
+
+### Decision 5: Model provider at adapter level
+
+**Choice**: construct the model through provider functions in the Megatron trainer/adapter layer under `primus/backends/megatron/`, above `core/models/`.
+
+**Reasoning**:
+- Model providers are adapter/wrapper functions rather than a standalone module; they live alongside the trainers (for example `primus/backends/megatron/megatron_pretrain_trainer.py`)
+- They sit above the core models to add Primus-specific functionality
+- Example: wrap the model with a custom loss, precision handling, or checkpoint logic
+
+### Decision 6: No PyTorch lightning
+
+**Choice**: Pure Megatron patterns
+**Not**: PyTorch Lightning DataModules
+
+**Reasoning**:
+- Primus doesn't use PyTorch Lightning
+- Framework-specific implementations reduce flexibility
+- Better integration with Megatron training loop
+- Follows Megatron-LM's `MegatronDataloaderWrapper` wrapper pattern
+
+---
+
+## Component hierarchy
+
+### 1. Model hierarchy
+
+```
+nn.Module (PyTorch)
+└── MegatronModule
+    └── DiffusionModule (abstract)
+        └── Flux (concrete)
+        ├── Joint layers: MMDiTLayer × num_joint_layers
+        │   └── JointSelfAttention (shared)
+        ├── Single layers: FluxSingleTransformerBlock × num_single_layers
+        │   └── FluxSingleAttention (shared)
+        └── Embeddings:
+            ├── TimeStepEmbedder (shared)
+            ├── MLPEmbedder (shared)
+            └── EmbedND (Flux-specific, 3D RoPE)
+```
+
+### 2. Configuration hierarchy
+
+```
+TransformerConfig (Megatron-Core)
+└── BaseDiffusionConfig
+    └── FluxConfig
+        ├── flux_535m() factory
+        └── flux_12b() factory
+```
+
+### 3. Scheduler hierarchy
+
+```
+BaseScheduler (abstract)
+├── FlowMatchEulerDiscreteScheduler (Flux)
+├── DDPMScheduler (future)
+├── EDMScheduler (future)
+└── EulerDiscreteScheduler (future)
+```
+
+### 4. Encoder hierarchy
+
+```
+BaseEncoder (abstract)
+├── ImageEncoder
+│   └── VAE
+│       ├── AutoencoderKL (Flux)
+│       └── VQVAE (future)
+└── TextEncoder
+    ├── T5
+    │   ├── T5-XXL (Flux)
+    │   └── T5-Large (future)
+    └── CLIP
+        ├── CLIP-L (Flux)
+        └── CLIP-H (future)
+```
+
+---
+
+## Data flow
+
+### Training pipeline
+
+```
+Raw Data (images + captions)
+    ↓
+[Optional] Precalculation
+    ├─ VAE → latents [B, 64, H/8, W/8]
+    ├─ T5-XXL → embeddings [B, 512, 4096]
+    └─ CLIP-L → pooled [B, 768]
+    ↓
+WebDataset/Energon Format (.tar files)
+    ↓
+MegatronDataloaderWrapper (from data/dataloader.py)
+    ↓
+EncodedDiffusionTaskEncoder (from data/diffusion/task_encoders/)
+    ├─ Load precalculated data, OR
+    └─ Encode on-the-fly (slower)
+    ↓
+Training Batch:
+    ├─ latents: [B, 64, H, W]
+    ├─ t5_embeddings: [B, S, 4096]
+    └─ clip_pooled: [B, 768]
+    ↓
+FlowMatchEulerDiscreteScheduler
+    ├─ Sample timesteps: t ~ U(0, 1)
+    ├─ Sample noise: ε ~ N(0, I)
+    ├─ Add noise: x_t = (1-t)*ε + t*x_0
+    └─ Compute target: v = x_0 - ε
+    ↓
+Flux Model Forward Pass
+    ├─ Embed timesteps
+    ├─ Embed pooled text (CLIP)
+    ├─ Joint layers (process latents + T5 embeddings)
+    ├─ Single layers (process latents only)
+    └─ Output: v_pred [B, 64, H, W]
+    ↓
+Loss Computation: MSE(v_pred, v_target)
+    ↓
+Backward Pass & Optimizer Step
+```
+
+### Inference pipeline
+
+```
+Text Prompt
+    ↓
+Text Encoders
+    ├─ T5-XXL → embeddings [1, S, 4096]
+    └─ CLIP-L → pooled [1, 768]
+    ↓
+Initialize Noise: x_0 ~ N(0, I)
+    ↓
+Sampling Loop (t = 1.0 → 0.0)
+    ├─ Model forward: v_t = Flux(x_t, t, embeddings)
+    ├─ Update: x_{t-dt} = x_t + v_t * dt
+    └─ Repeat until t = 0
+    ↓
+Latents: x_0 [1, 64, H, W]
+    ↓
+VAE Decoder
+    ↓
+Generated Image [1, 3, H*8, W*8]
+```
+
+---
+
+## Comparison with alternative implementations
+
+### Primus architectural advantages
+
+#### 1. TransformerBlock architecture (Primus innovation)
+- **Primus**: Unified `TransformerBlock` with heterogeneous layer specs
+- **Others**: Separate `nn.ModuleList` containers for double/single blocks
+- **Benefit**: Better PP slicing, unified checkpointing, future-proof
+
+#### 2. Megatron-core native
+- **Primus**: Pure Megatron-Core, no framework dependencies
+- **Others**: Often integrated with PyTorch Lightning or other frameworks
+- **Benefit**: Tighter integration, simpler training loops
+
+#### 3. Checkpoint format
+- **Primus**: Unified `transformer.layers.{0-56}` structure
+- **Others**: Separate `double_blocks.{i}` and `single_blocks.{j}`
+- **Benefit**: Simpler distributed checkpointing
+
+#### 4. Encoder architecture
+- **Primus**: Registry-based, hierarchical organization
+- **Others**: Direct imports from monolithic files
+- **Benefit**: Easy extensibility for new encoder variants
+
+### File organization comparison
+
+| Component | Standard Location | Primus Location | Improvement |
+|-----------|------------------|-----------------|-------------|
+| Model | `models/diffusion/flux/` | `core/models/diffusion/flux/` | Megatron-Core convention |
+| Layers | Mixed locations | `common/` for shared, `flux/` for specific | Clear boundaries |
+| Encoders | Single file | `data/diffusion/encoders/{type}/{variant}/` | Hierarchical, extensible |
+| Tests | Mixed with code | `tests/unit_tests/backends/megatron/diffusion/` | Proper separation |
+
+---
+
+## Implementation status
+
+### Core infrastructure ✅
+- ✅ Directory structure
+- ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler)
+- ✅ DiffusionModule with Megatron-Core integration
+- ✅ FluxConfig with factory methods
+- ✅ FlowMatchEulerDiscreteScheduler implementation
+- ✅ Configuration files (YAML)
+- ✅ Testing framework (290+ tests)
+- ✅ Documentation structure
+
+### Flux model implementation ✅
+- ✅ Flux model architecture
+- ✅ MMDiT layers and attention
+- ✅ Embeddings (RoPE, timestep, vector)
+- ✅ Encoder registry and loaders
+- ✅ TaskEncoder for Energon
+- ✅ Data pipeline
+
+---
+
+## Extension points
+
+### Adding a new model (e.g., DiT)
+
+1. **Create model directory**: `core/models/diffusion/dit/`
+2. **Add config**: Extend `BaseDiffusionConfig`
+3. **Implement model**: Extend `DiffusionModule` (which extends MegatronModule)
+   - Inherit process group management
+   - Get distributed checkpointing support
+   - Access attention backend configuration
+4. **Reuse shared components**: Import from `common/`
+5. **Add tests**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py`
+6. **Update configs**: Add `dit_config.yaml`
+
+### Adding a new encoder variant
+
+1. **Create encoder file**: e.g., `data/diffusion/encoders/text/t5/t5_large.py`
+2. **Implement encoder class**: Extend `BaseEncoder`
+3. **Register**: Add to `ENCODER_REGISTRY`
+4. **Add config**: Update `encoders.yaml`
+5. **Add tests**: Test in `tests/unit_tests/backends/megatron/diffusion/data/encoders/`
+
+### Adding a new scheduler
+
+1. **Create scheduler file**: `training/diffusion/schedulers/ddpm.py`
+2. **Implement**: Extend `BaseScheduler`
+3. **Export**: Add to `__init__.py`
+4. **Add tests**: `tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py`
+5. **Document**: Update this file and README
+
+---
+
+## Performance considerations
+
+### Memory optimization
+- **Precalculated data**: 5-10x faster, lower memory
+- **Frozen encoders**: Only train diffusion model
+- **Gradient checkpointing**: Trade compute for memory
+- **Mixed precision**: bf16 on MI300X (compatible with H100/A100)
+
+### Multi-GPU scaling
+- **Tensor Parallelism**: Split model across GPUs
+- **Pipeline Parallelism**: Split layers across GPUs
+- **Data Parallelism**: Replicate model, split data
+- **Sequence Parallelism**: For very long sequences
+
+### Best practices
+1. Use precalculated mode for training
+2. Freeze encoders (standard practice)
+3. Use bf16 on modern hardware
+4. Start with TP=1, PP=1, scale as needed
+5. Profile before optimizing
+
+---
+
+## References
+
+- **Megatron-Core**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) transformer patterns
+- **Flux**: [black-forest-labs/FLUX.1](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+- **Flow Matching**: Rectified flow and flow matching papers
+- **NeMo**: [nvidia/NeMo](https://github.com/NVIDIA/NeMo) - Alternative diffusion implementation
+
+---
+
+**Last Updated**: December 2025
diff --git a/docs/04-technical-guides/diffusion-models/data_preprocessing.md b/docs/04-technical-guides/diffusion-models/data_preprocessing.md
new file mode 100644
index 000000000..a99a9754f
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/data_preprocessing.md
@@ -0,0 +1,629 @@
+# Data preprocessing guide
+
+This guide explains how to prepare datasets for Flux and other diffusion models in Primus, including pre-encoding of VAE latents and text embeddings into Energon WebDataset format.
+
+---
+
+## Table of contents
+
+1. [Overview](#overview)
+2. [Quick start](#quick-start)
+3. [Two pipelines](#two-pipelines)
+4. [Running preprocessing](#running-preprocessing)
+5. [Configuration](#configuration)
+6. [Authentication](#authentication)
+7. [Finalization](#finalization)
+8. [Output format](#output-format)
+9. [Validation](#validation)
+10. [Troubleshooting](#troubleshooting)
+
+---
+
+## Overview
+
+Diffusion models require three types of encodings:
+1. **VAE latents**: Images encoded to latent space
+2. **Text embeddings**: Captions encoded with T5-XXL (sequence)
+3. **Pooled embeddings**: Captions encoded with CLIP-L (pooled)
+
+### Why pre-encode?
+
+**Benefits**:
+- 5-10x faster training (no online encoding)
+- Lower GPU memory usage (encoders not loaded during training)
+- Deterministic inputs (same preprocessing for all runs)
+- Eliminates encoder differences as a variable
+
+**When to Use**:
+- Training (highly recommended)
+- Fine-tuning on fixed datasets
+- Benchmarking and reproducibility
+
+**When NOT to Use**:
+- Interactive data augmentation needed
+- Dataset too large to store pre-encoded
+- Rapid prototyping with changing data
+
+---
+
+## Quick start
+
+Preprocess the Pokemon dataset with a single command:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+This will:
+1. Download the `diffusers/pokemon-gpt4-captions` dataset from HuggingFace
+2. Encode all images with VAE and captions with T5/CLIP
+3. Write Energon WebDataset tar shards to `/workspace/Primus/data/quickstart_pokemon`
+4. Automatically finalize the dataset (create `dataset.yaml`, run `energon prepare`, validate)
+
+The default encoder model (`black-forest-labs/FLUX.1-dev`) is gated and requires a HuggingFace token. Get one at https://huggingface.co/settings/tokens and save it to a file.
+
+---
+
+## Two pipelines
+
+Primus provides two preprocessing pipelines via the `primus data` CLI:
+
+### `diffusion-encoded` (recommended for training)
+
+Pre-encodes images with VAE and text with T5/CLIP. Produces larger datasets but enables faster training since encoders are not needed at training time.
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+**Output**: WebDataset shards containing `latents.pth`, `prompt_embeds.pth`, `pooled_prompt_embeds.pth`, `caption.txt`
+
+**Use when**: Training on production datasets, maximum training speed is needed, storage is available.
+
+### `diffusion-raw`
+
+Stores raw images and captions without encoding. Smaller datasets but encoding happens on-the-fly during training (requires GPU + encoders loaded in memory).
+
+```bash
+primus-cli direct -- data diffusion-raw \
+  --source-type huggingface \
+  --hf-dataset diffusers/pokemon-gpt4-captions \
+  --output-dir /workspace/Primus/data/raw_pokemon \
+  --hf-token-file /path/to/.hf_token
+```
+
+**Output**: WebDataset shards containing `jpg` (or `png`/`webp`) and `txt` files.
+
+**Use when**: Storage is limited, experimenting with different encoders, rapid prototyping.
+
+**Note**: `diffusion-raw` does not support `--config` files. All options must be passed as CLI arguments.
+
+---
+
+## Running preprocessing
+
+### Using a config file (recommended)
+
+The `--config` flag is supported by `diffusion-encoded` only. The simplest approach uses a YAML config file:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+Available example configs in `primus/configs/data/megatron/diffusion/preprocessing/`:
+
+| Config | Source | Description |
+|--------|--------|-------------|
+| `quickstart_pokemon.yaml` | HuggingFace | Minimal config, 256px, fast |
+| `example_huggingface.yaml` | HuggingFace | Full example with all options |
+| `example_directory.yaml` | Local directory | Images + captions from disk |
+| `example_webdataset.yaml` | WebDataset | Existing tar archives |
+| `example_base.yaml` | N/A | Comprehensive reference with all fields |
+| `text_to_image_2m_10k.yaml` | HuggingFace | 10K subset of text-to-image-2M (1024px) |
+
+### Using CLI arguments directly
+
+All config values can be provided as CLI arguments:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --source-type huggingface \
+  --hf-dataset diffusers/pokemon-gpt4-captions \
+  --output-dir /workspace/Primus/data/encoded_pokemon \
+  --model-path black-forest-labs/FLUX.1-dev \
+  --batch-size 8 \
+  --precision bf16 \
+  --hf-token-file /path/to/.hf_token
+```
+
+### CLI overrides config values
+
+When using both `--config` and CLI arguments, CLI arguments take priority:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --output-dir /my/custom/path \
+  --batch-size 16 \
+  --max-samples 1000
+```
+
+Priority order (highest to lowest):
+1. Explicitly provided CLI arguments
+2. YAML config file values
+3. CLI default values
+
+### Multi-GPU processing
+
+Use `--nproc-per-node` for data-parallel preprocessing across multiple GPUs:
+
+```bash
+primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+Each GPU processes a subset of the data. Shards are named to avoid conflicts across ranks.
+
+---
+
+## Configuration
+
+### YAML config structure
+
+Preprocessing configs have four sections:
+
+```yaml
+source:
+  type: huggingface          # huggingface | directory | webdataset
+  hf_dataset: diffusers/pokemon-gpt4-captions
+  hf_split: train
+
+output:
+  output_dir: /workspace/Primus/data/encoded_pokemon
+  shard_size: 1000           # samples per tar shard
+  max_samples: null          # null = process all
+  compress: false
+
+model:
+  model_path: black-forest-labs/FLUX.1-dev   # HF repo or local path
+  precision: bf16            # bf16 | fp16 | fp32
+  batch_size: 8
+  # Optional per-encoder overrides:
+  vae_path: null
+  t5_path: null
+  clip_path: null
+
+image:
+  image_size: 1024
+  center_crop: false
+```
+
+### Source types
+
+**HuggingFace** (`type: huggingface`):
+```yaml
+source:
+  type: huggingface
+  hf_dataset: diffusers/pokemon-gpt4-captions
+  hf_split: train
+  hf_data_files: null        # optional: specific files within dataset
+```
+
+**Local Directory** (`type: directory`):
+```yaml
+source:
+  type: directory
+  input_dir: /data/my_images
+```
+
+Expected directory structure for `directory` source:
+```
+my_images/
+├── images/
+│   ├── 00001.jpg
+│   ├── 00002.png
+│   └── ...
+└── captions/
+    ├── 00001.txt
+    ├── 00002.txt
+    └── ...
+```
+
+**WebDataset** (`type: webdataset`):
+```yaml
+source:
+  type: webdataset
+  input_path: /data/existing_shards/*.tar
+```
+
+### Model configuration
+
+The `model_path` defaults to `black-forest-labs/FLUX.1-dev`, which downloads VAE, T5-XXL, and CLIP-L encoders from HuggingFace. This model is gated and requires authentication (see [Authentication](#authentication)).
+
+Individual encoder paths can be overridden:
+
+```yaml
+model:
+  model_path: black-forest-labs/FLUX.1-dev
+  vae_path: /local/models/vae        # use local VAE instead
+  t5_path: null                       # falls back to model_path
+  clip_path: null                     # falls back to model_path
+```
+
+---
+
+## Authentication
+
+The default encoder model (`FLUX.1-dev`) is gated on HuggingFace and requires authentication. Primus supports three authentication methods, checked in priority order:
+
+### 1. Token file (recommended)
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config your_config.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+The token file must have secure permissions (600 or 400). Create it with:
+
+```bash
+echo "hf_your_token_here" > /path/to/.hf_token
+chmod 600 /path/to/.hf_token
+```
+
+### 2. Environment variable
+
+```bash
+export HF_TOKEN=hf_your_token_here
+primus-cli direct -- data diffusion-encoded --config your_config.yaml
+```
+
+### 3. HuggingFace CLI login
+
+```bash
+huggingface-cli login
+primus-cli direct -- data diffusion-encoded --config your_config.yaml
+```
+
+If authentication fails, Primus provides a clear error message indicating which encoder failed and how to fix it.
+
+---
+
+## Finalization
+
+Finalization is **automatic by default**. After preprocessing completes, Primus automatically:
+
+1. **Creates `.nv-meta/dataset.yaml`** with `CrudeWebdataset` sample type and encoding subflavor
+2. **Runs `energon prepare`** to index the tar shards and create split assignments
+3. **Validates the dataset** using Primus's custom validation (metadata checks, sample count verification, energon API spot-check)
+
+### Skipping finalization
+
+To skip automatic finalization (e.g., for manual post-processing):
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config your_config.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --no-finalize
+```
+
+### Custom train/val/test splits
+
+By default, 100% of data goes to the training split. To create validation and test splits:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config your_config.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --train-split 0.8
+```
+
+This creates an 80% train / 10% val / 10% test split.
+
+---
+
+## Output format
+
+### Directory structure
+
+After preprocessing and finalization, the output directory contains:
+
+```
+encoded_pokemon/
+├── 000000.tar              # WebDataset shard
+├── 000001.tar
+├── 000000.tar.idx          # Energon index files
+├── 000001.tar.idx
+└── .nv-meta/               # Energon metadata
+    ├── dataset.yaml         # Dataset type configuration
+    ├── split.yaml           # Train/val/test split assignments
+    └── .info.json           # Shard counts and sample counts
+```
+
+### Pre-encoded shard contents (`diffusion-encoded`)
+
+Each tar shard contains samples with these keys:
+
+```
+000000.tar:
+├── 0000000000.latents.pth              # VAE latents tensor
+├── 0000000000.prompt_embeds.pth        # T5-XXL embeddings tensor
+├── 0000000000.pooled_prompt_embeds.pth # CLIP-L pooled embeddings tensor
+├── 0000000000.caption.txt              # Original caption text
+├── 0000000001.latents.pth
+├── 0000000001.prompt_embeds.pth
+├── ...
+```
+
+Tensor shapes:
+- `latents.pth`: `[64, H/8, W/8]` (e.g., `[64, 128, 128]` for 1024x1024 images)
+- `prompt_embeds.pth`: `[seq_len, 4096]` (T5-XXL hidden dim)
+- `pooled_prompt_embeds.pth`: `[768]` (CLIP-L pooled dim)
+
+### Raw shard contents (`diffusion-raw`)
+
+```
+000000.tar:
+├── 0000000000.jpg       # Preprocessed image
+├── 0000000000.txt       # Caption text
+├── 0000000001.jpg
+├── 0000000001.txt
+├── ...
+```
+
+### Dataset YAML
+
+The auto-generated `.nv-meta/dataset.yaml` uses `CrudeWebdataset` format:
+
+```yaml
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: preencoded    # or 'raw' for diffusion-raw
+```
+
+---
+
+## Validation
+
+### Automatic validation
+
+Validation runs automatically as part of finalization. It performs four checks:
+
+1. **Metadata check**: Verifies `.nv-meta/` files exist and are valid (`.info.json`, `split.yaml`, `dataset.yaml`)
+2. **Sample count check**: Spot-checks that tar shard entry counts match `.info.json`
+3. **Sample load check**: Loads one sample through Energon's Python API (same code path as training)
+4. **Summary report**: Prints dataset statistics (encoding, total samples, splits, data shapes, size)
+
+### Standalone validation
+
+To validate a dataset independently:
+
+```bash
+python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset
+
+# For raw datasets:
+python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset --encoding raw
+```
+
+### Programmatic validation
+
+```python
+from primus.backends.megatron.data.diffusion.preprocessing.validate import validate_energon_dataset
+
+ok = validate_energon_dataset('/path/to/dataset', encoding='preencoded')
+```
+
+### Known Energon CLI limitations
+
+The standard Energon CLI tools (`energon info`, `energon preview`, `energon lint`) do **not** work correctly with `CrudeWebdataset` format. Primus uses custom validation instead:
+
+- `energon info` raises `KeyError: 'sample_type'`
+- `energon preview` raises `TypeError` (expects dataclass, `CrudeSample` is a dict)
+- `energon lint` raises `AssertionError` (expects registered cookers)
+
+Use Primus's built-in validation or the standalone script above.
+
+---
+
+## Troubleshooting
+
+### HuggingFace authentication failure
+
+**Symptoms**: Error mentioning "token", "gated", "401", or "403" when downloading encoders.
+
+**Solutions**:
+1. Provide a token: `--hf-token-file /path/to/.hf_token`
+2. Set environment variable: `export HF_TOKEN=hf_xxx`
+3. Run `huggingface-cli login`
+4. Accept the model's license on https://huggingface.co/black-forest-labs/FLUX.1-dev
+
+### Out of memory during preprocessing
+
+**Symptoms**: CUDA out of memory error during encoding.
+
+**Solutions**:
+1. Reduce batch size: `--batch-size 4` or `--batch-size 1`
+2. Use smaller image size: `--image-size 512`
+3. Use fp16 precision: `--precision fp16`
+4. Use multi-GPU to distribute work: `--nproc-per-node=8`
+
+### Missing dependencies
+
+**Symptoms**: `ModuleNotFoundError` for `webdataset`, `megatron-energon`, `tqdm`, etc.
+
+**Solution**:
+```bash
+pip install -r requirements.txt
+```
+
+Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages.
+
+### Finalization fails
+
+**Symptoms**: Error during `energon prepare` or validation after preprocessing.
+
+**Solutions**:
+1. Check that tar shards exist in the output directory
+2. Ensure `megatron-energon` is installed (`pip install megatron-energon`)
+3. Re-run with `--no-finalize`, then manually inspect the output before finalizing
+
+### Slow preprocessing
+
+**Symptoms**: Low throughput (< 10 samples/sec).
+
+**Solutions**:
+1. Increase batch size (limited by VRAM): `--batch-size 16`
+2. Use multiple GPUs: `--nproc-per-node=8`
+3. Use bf16 precision (faster on supported hardware): `--precision bf16`
+4. For raw pipeline, reduce image quality: `--image-quality 85`
+
+---
+
+## Best practices
+
+1. **Always pre-encode for production training**: 5-10x speedup is worth the storage
+2. **Test on small dataset first**: Use `--max-samples 100` to verify the pipeline works
+3. **Use bf16 precision**: Good balance of speed, storage, and quality
+4. **Start with quickstart_pokemon.yaml**: Verify your setup before processing large datasets
+5. **Keep raw data**: Pre-encoding is a one-way transformation
+6. **Version your configs**: Track which preprocessing config produced each dataset
+
+---
+
+## Flux-specific data preparation (torchrun)
+
+This section covers preparing datasets for Flux training using `torchrun` directly
+inside a Docker container, as an alternative to the `primus-cli` workflow above.
+
+### Prerequisites
+
+Start the container and optionally install requirements:
+
+```bash
+bash tools/docker/start_container.sh
+docker exec dev_primus bash -c 'pip install -r /workspace/Primus/requirements.txt'
+```
+
+This mounts the repository to `/workspace/Primus` inside the container.
+Override the image with `DOCKER_IMAGE`:
+
+```bash
+DOCKER_IMAGE=docker.io/rocm/primus:v26.1 bash tools/docker/start_container.sh
+```
+
+### Input directory structure
+
+When using `--source-type directory`, organize your data as follows:
+
+```
+dataset/
+├── images/
+│   ├── 0000000.png
+│   ├── 0000001.png
+│   └── ...
+└── captions/
+    ├── 0000000.txt
+    ├── 0000001.txt
+    └── ...
+```
+
+The file stem links each image to its caption (e.g., `images/0000000.png` pairs
+with `captions/0000000.txt`). Images can be `.jpg`, `.jpeg`, `.png`, or `.webp`.
+Captions must be UTF-8 `.txt` files. Samples without a matching caption are skipped.
+
+Other supported source types:
+- **huggingface** -- Load directly from HuggingFace Hub. Requires `--hf-dataset`.
+- **webdataset** -- Read from existing WebDataset tar archives. Requires `--input-path`.
+
+### Image sizing
+
+**Fixed size (default):** Every image is resized to `--image-size` pixels square
+(default 1024). Use `--center-crop` to control center-cropping.
+
+**Variable size (`--variable-size`):** Preserves aspect ratio by scaling the
+longest side to `--max-size` (default 1024) and rounding dimensions to multiples
+of 16. Only use this when all images share the same dimensions -- mixed tensor
+sizes cause load imbalance across GPUs.
+
+### Key parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--source-type` | — | `directory`, `huggingface`, or `webdataset` |
+| `--input-dir` | — | Dataset root (for `directory` source) |
+| `--image-size` | 1024 | Square target size |
+| `--variable-size` | off | Preserve aspect ratio |
+| `--max-size` | 1024 | Maximum dimension in variable-size mode |
+| `--model-path` | `black-forest-labs/FLUX.1-dev` | Base model for encoder weights |
+| `--t5-max-length` | 512 | T5 token limit (use 256 for FLUX.1-schnell) |
+| `--batch-size` | 8 | Encoding batch size per GPU |
+| `--output-dir` | — | Destination for encoded WebDataset shards |
+| `--shard-size` | 1000 | Samples per tar shard |
+
+### Examples
+
+**From host (via `docker exec`):**
+
+```bash
+docker exec dev_primus bash -c '\
+  export HF_HOME=/workspace/Primus/checkpoints/flux; \
+PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \
+  torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \
+    data diffusion-encoded \
+    --source-type directory \
+    --input-dir /workspace/Primus/data/dataset \
+    --output-dir /workspace/Primus/data/dataset_encoded_256 \
+    --image-size 256  \
+    --t5-max-length 256'
+```
+
+**Inside the container:**
+
+```bash
+export HF_HOME=/workspace/Primus/checkpoints/flux
+PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \
+  torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \
+    data diffusion-encoded \
+    --source-type directory \
+    --input-dir /workspace/Primus/data/dataset \
+    --output-dir /workspace/Primus/data/dataset_encoded_256 \
+    --image-size 256 \
+    --t5-max-length 256
+```
+
+Both commands encode a local directory dataset at 256px on 8 GPUs, then create an Energon dataset by default.
+
+### Limitations
+
+**Uniform output size only.** When using `--variable-size`, images may produce
+tensors of different shapes. The current data pipeline does not support mixed
+tensor sizes because samples are distributed evenly across GPUs and unequal
+shapes cause load imbalance and eventual timeout. Use fixed `--image-size` if
+your images have various sizes.
+
+---
+
+## Next steps
+
+After preprocessing:
+1. **Training**: Use the preprocessed dataset path in your training config
+2. **Validation**: The dataset is ready for training immediately after finalization
+
+See:
+- [Energon Integration](energon_integration.md) for TaskEncoder and dataloader details
+- [Config Directory Guide](../../../primus/configs/data/megatron/diffusion/README.md) for config file reference
+- Example configs in `primus/configs/data/megatron/diffusion/preprocessing/`
+
+---
+
+**Last Updated**: June 2026
diff --git a/docs/04-technical-guides/diffusion-models/energon_integration.md b/docs/04-technical-guides/diffusion-models/energon_integration.md
new file mode 100644
index 000000000..97941a67c
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/energon_integration.md
@@ -0,0 +1,514 @@
+# Energon integration guide
+
+This guide explains how Megatron-Energon is integrated with Primus diffusion models, including the Cooker/TaskEncoder pattern, dataloader configuration, and dataset format.
+
+---
+
+## Table of contents
+
+1. [Overview](#overview)
+2. [Energon architecture](#energon-architecture)
+3. [Dataset format](#dataset-format)
+4. [TaskEncoder and cooker pattern](#taskencoder-and-cooker-pattern)
+5. [Dataloader setup](#dataloader-setup)
+6. [Dataset configuration](#dataset-configuration)
+7. [Implementation examples](#implementation-examples)
+8. [Best practices](#best-practices)
+
+---
+
+## Overview
+
+### What is Megatron-Energon?
+
+Megatron-Energon is NVIDIA's data loading framework for large-scale multimodal training. It provides:
+- **WebDataset integration**: Efficient streaming from .tar archives
+- **Task encoding**: Flexible data transformation pipeline via Cookers
+- **Multi-worker support**: Parallel data loading
+- **Deterministic iteration**: Reproducible training
+- **Checkpoint resumption**: Resume from any step
+
+### Why Energon for diffusion?
+
+- **Proven at scale**: Used by NVIDIA for LLM and multimodal training
+- **Flexible**: Supports various data formats and transformations
+- **Efficient**: Optimized for multi-GPU training
+- **Compatible**: Works with Megatron parallelism (TP, PP, DP)
+
+### Primus integration strategy
+
+```
+Shared Infrastructure (data/)
+  ├─ MegatronDataloaderWrapper
+  ├─ Dataset configuration parsers
+  └─ Common utilities
+      ↓
+Model-Specific TaskEncoders (data/diffusion/task_encoders/)
+  ├─ EncodedDiffusionTaskEncoder (preencoded data)
+  ├─ RawDiffusionTaskEncoder (raw images + text)
+  └─ Custom task encoders
+```
+
+**Key Principle**: Share infrastructure, separate domain logic.
+
+---
+
+## Energon architecture
+
+### Component stack
+
+```
+Training Loop
+    ↓
+MegatronDataloaderWrapper (Primus wrapper, cyclic iteration)
+    ↓
+Megatron-Energon Core
+    ├─ WebDataset Reader (reads .tar shards)
+    ├─ Cooker (transforms raw dict → typed Sample)
+    ├─ TaskEncoder.batch() (stacks samples → batch)
+    └─ Worker Pool
+    ↓
+.tar Shards (CrudeWebdataset format)
+```
+
+### Data flow
+
+```
+1. Load from .tar shard
+   → raw_dict = {'__key__': '0000000000',
+                  'latents.pth': bytes,
+                  'prompt_embeds.pth': bytes,
+                  'pooled_prompt_embeds.pth': bytes,
+                  'caption.txt': bytes}
+
+2. Cooker function (e.g., cook_preencoded_diffusion)
+   → Deserializes bytes to tensors
+   → Returns typed DiffusionSample dataclass
+   → DiffusionSample(latents=Tensor[64,128,128],
+                      prompt_embeds=Tensor[512,4096],
+                      pooled_prompt_embeds=Tensor[768],
+                      caption="a photo of...")
+
+3. TaskEncoder.batch()
+   → Stacks list of DiffusionSample into batch dict
+   → batch = {'latents': [B, 64, H, W],
+              'prompt_embeds': [B, seq_len, 4096],
+              'pooled_prompt_embeds': [B, 768]}
+
+4. Return to training loop via MegatronDataloaderWrapper
+   → Forward pass, loss, backprop
+```
+
+---
+
+## Dataset format
+
+### CrudeWebdataset
+
+Primus uses Energon's `CrudeWebdataset` format for preprocessed datasets. This is the simplest Energon format -- it stores raw key-value pairs in tar shards without requiring a strict schema.
+
+The dataset type is configured in `.nv-meta/dataset.yaml`:
+
+```yaml
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: preencoded    # or 'raw'
+```
+
+The `subflavors.encoding` field tells the TaskEncoder which Cooker to use:
+- `preencoded`: Uses `cook_preencoded_diffusion` -- loads pre-encoded tensors
+- `raw`: Uses `cook_raw_images` -- loads raw images and text
+
+### Pre-encoded shard contents
+
+Each tar shard contains samples with these keys:
+
+| Key | Type | Shape | Description |
+|-----|------|-------|-------------|
+| `latents.pth` | Tensor | `[64, H/8, W/8]` | VAE-encoded image latents |
+| `prompt_embeds.pth` | Tensor | `[seq_len, 4096]` | T5-XXL text embeddings |
+| `pooled_prompt_embeds.pth` | Tensor | `[768]` | CLIP-L pooled embeddings |
+| `caption.txt` | str | N/A | Original caption text |
+
+### Raw shard contents
+
+| Key | Type | Description |
+|-----|------|-------------|
+| `jpg` / `png` / `webp` | bytes | Preprocessed image |
+| `txt` | str | Caption text |
+
+### Known Energon CLI limitations
+
+The standard Energon CLI tools have issues with `CrudeWebdataset`:
+- `energon info` raises `KeyError: 'sample_type'`
+- `energon preview` raises `TypeError` (expects dataclass, `CrudeSample` is a dict)
+- `energon lint` raises `AssertionError` (expects registered cookers)
+
+Primus includes custom validation (`primus.backends.megatron.data.diffusion.preprocessing.validate`) as a replacement. See the [Data Preprocessing Guide](data_preprocessing.md#validation) for details.
+
+---
+
+## TaskEncoder and cooker pattern
+
+Primus uses Energon's Cooker pattern rather than the older `encode_sample()` approach. Cookers are `@stateless` functions that transform raw sample dicts into typed dataclass instances, dispatched based on `subflavors`.
+
+### DiffusionSample dataclass
+
+Location: `primus/backends/megatron/data/diffusion/task_encoders/image.py`
+
+```python
+# (imports like torch omitted for brevity)
+from dataclasses import dataclass
+from megatron.energon import Sample
+
+@dataclass
+class DiffusionSample(Sample):
+    """
+    Diffusion training sample with framework-standard field names.
+
+    Inherits from megatron.energon.Sample to ensure __key__, __restore_key__,
+    and __subflavors__ are properly tracked for deterministic training resumption.
+    """
+    latents: torch.Tensor              # [C, H, W] VAE latents
+    prompt_embeds: torch.Tensor        # [seq_len, hidden_dim] T5 embeddings
+    pooled_prompt_embeds: torch.Tensor # [hidden_dim] CLIP pooled
+    caption: str = ""
+```
+
+### Cooker functions
+
+Cookers are `@stateless` functions registered with a `Cooker` wrapper that specifies which `subflavors` they handle:
+
+```python
+# (imports like torch, io omitted for brevity)
+from megatron.energon import Cooker, basic_sample_keys, stateless
+
+@stateless
+def cook_preencoded_diffusion(sample: dict) -> DiffusionSample:
+    """Load precalculated VAE latents and text embeddings from disk."""
+
+    def load_tensor(data):
+        if isinstance(data, bytes):
+            return torch.load(io.BytesIO(data), map_location='cpu')
+        return data
+
+    latents = load_tensor(sample.get('latents.pth'))
+    prompt_embeds = load_tensor(sample.get('prompt_embeds.pth'))
+    pooled_prompt_embeds = load_tensor(sample.get('pooled_prompt_embeds.pth'))
+
+    caption = sample.get('caption.txt', b'')
+    if isinstance(caption, bytes):
+        caption = caption.decode('utf-8')
+
+    return DiffusionSample(
+        **basic_sample_keys(sample),
+        latents=latents,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        caption=caption,
+    )
+
+
+@stateless
+def cook_raw_images(sample: dict) -> Dict[str, Any]:
+    """Load raw images and text -- NO encoding, just data loading."""
+    return {
+        **basic_sample_keys(sample),
+        'images': sample.get('images'),
+        'txt': sample.get('txt'),
+    }
+```
+
+The cooker code above is simplified for clarity. See `primus/backends/megatron/data/diffusion/task_encoders/image.py` for the full implementation, which includes additional input type handling and validation.
+
+### EncodedDiffusionTaskEncoder
+
+The TaskEncoder registers Cookers and provides the `batch()` method:
+
+```python
+from megatron.energon import DefaultTaskEncoder, SampleDecoder, Cooker, WorkerConfig
+
+class EncodedDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]):
+    """TaskEncoder for PRE-ENCODED diffusion data."""
+
+    decoder = SampleDecoder(image_decode="pil")
+
+    cookers = [
+        Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}),
+    ]
+
+    def __init__(self, worker_config: Optional[WorkerConfig] = None):
+        super().__init__()
+        self.worker_config = worker_config
+
+    def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]:
+        return {
+            'latents': torch.stack([s.latents for s in samples]),
+            'prompt_embeds': torch.stack([s.prompt_embeds for s in samples]),
+            'pooled_prompt_embeds': torch.stack([s.pooled_prompt_embeds for s in samples]),
+        }
+```
+
+### RawDiffusionTaskEncoder
+
+For raw (on-the-fly encoding) datasets:
+
+```python
+class RawDiffusionTaskEncoder(DefaultTaskEncoder):
+    """TaskEncoder for RAW diffusion data (images and text)."""
+
+    decoder = SampleDecoder(image_decode="pil")
+
+    cookers = [
+        Cooker(cook_raw_images, has_subflavors={"encoding": "raw"}),
+    ]
+
+    def __init__(self, worker_config: Optional[WorkerConfig] = None):
+        super().__init__()
+        self.worker_config = worker_config
+
+    def batch(self, samples: List[Dict]) -> Dict[str, Any]:
+        return {
+            'images': [s['images'] for s in samples],
+            'txt': [s['txt'] for s in samples],
+        }
+```
+
+### How cooker dispatch works
+
+The Cooker framework matches samples to cooker functions based on `subflavors`:
+
+1. `dataset.yaml` specifies `subflavors: { encoding: preencoded }`
+2. Energon reads a sample from the tar shard
+3. The `has_subflavors` on each `Cooker` is checked against the sample's subflavors
+4. The matching cooker function is called to transform the raw dict into a typed sample
+5. `TaskEncoder.batch()` stacks multiple samples into a training batch
+
+This decouples data format (what's in the tar) from data loading logic (how to interpret it).
+
+---
+
+## Dataloader setup
+
+### MegatronDataloaderWrapper
+
+Location: `primus/backends/megatron/data/dataloader.py`
+
+The `MegatronDataloaderWrapper` is a generic wrapper that makes any iterable compatible with Megatron's training loop. It provides:
+- Cyclic iteration (never raises StopIteration)
+- Optional checkpoint support via duck typing (`save_state_rank()` / `restore_state_rank()`)
+- Works with PyTorch DataLoader, Energon loaders, and synthetic data
+
+```python
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+
+wrapper = MegatronDataloaderWrapper(pytorch_or_energon_loader)
+```
+
+Note: Originally named `EnergonDataloader`, renamed to `MegatronDataloaderWrapper` to reflect its generic nature (it has no Energon dependencies). The old name is available as a deprecated alias.
+
+### Usage example
+
+```python
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+
+# Create Energon loader (with TaskEncoder configured)
+energon_loader = get_loader(...)
+dataloader = MegatronDataloaderWrapper(energon_loader)
+
+# Use in training loop
+for batch in dataloader:
+    latents = batch['latents']                          # [B, 64, H, W]
+    prompt_embeds = batch['prompt_embeds']              # [B, seq_len, 4096]
+    pooled_prompt_embeds = batch['pooled_prompt_embeds'] # [B, 768]
+
+    output = model(latents, timesteps, prompt_embeds, pooled_prompt_embeds)
+    loss = criterion(output, target)
+    loss.backward()
+```
+
+---
+
+## Dataset configuration
+
+### Per-dataset dataset.yaml
+
+Each dataset directory has a `.nv-meta/dataset.yaml` that specifies the Energon dataset type:
+
+```yaml
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: preencoded
+```
+
+This file is auto-generated by Primus during finalization. See [Data Preprocessing Guide](data_preprocessing.md#finalization).
+
+### Metadataset (multi-dataset mixing)
+
+For combining multiple datasets with different weights:
+
+```yaml
+__module__: megatron.energon
+__class__: Metadataset
+
+splits:
+  train:
+    datasets:
+      - weight: 0.7
+        path: /data/laion_precalculated/
+      - weight: 0.3
+        path: /data/coco_precalculated/
+```
+
+Energon samples proportionally to weights (70% from LAION, 30% from COCO).
+
+---
+
+## Implementation examples
+
+### Example 1: Pre-encoded training
+
+```python
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+from primus.backends.megatron.data.diffusion.task_encoders import EncodedDiffusionTaskEncoder
+
+# TaskEncoder is configured by the dataset provider
+# The cooker automatically handles subflavors dispatch
+energon_loader = get_loader(...)
+dataloader = MegatronDataloaderWrapper(energon_loader)
+
+for batch in dataloader:
+    output = model(
+        batch['latents'],
+        batch['prompt_embeds'],
+        batch['pooled_prompt_embeds'],
+    )
+```
+
+### Example 2: Raw data training (on-the-fly encoding)
+
+```python
+from primus.backends.megatron.data.diffusion.task_encoders import RawDiffusionTaskEncoder
+
+# Raw TaskEncoder passes through images and text without encoding
+# Encoding happens in the model's forward_step
+energon_loader = get_loader(...)
+dataloader = MegatronDataloaderWrapper(energon_loader)
+
+for batch in dataloader:
+    images = batch['images']   # List of PIL Images
+    captions = batch['txt']    # List of caption strings
+    # Model handles VAE/T5/CLIP encoding in forward_step
+```
+
+### Example 3: Multi-GPU training
+
+```python
+import torch.distributed as dist
+
+dist.init_process_group(backend='nccl')
+
+energon_loader = get_loader(...)
+dataloader = MegatronDataloaderWrapper(energon_loader)
+
+for batch in dataloader:
+    output = model(batch['latents'], ...)
+```
+
+---
+
+## Best practices
+
+### 1. Pre-encoded mode
+- **Always use for production training**: 5-10x faster
+- **Validate first**: Test with small dataset using `quickstart_pokemon.yaml`
+- **Version datasets**: Track which preprocessing config produced each dataset
+
+### 2. Cooker design
+- **Use `@stateless`**: Cookers must be stateless and side-effect-free
+- **Use `basic_sample_keys()`**: Always forward `__key__`, `__restore_key__`, `__subflavors__`
+- **Keep it simple**: Cookers should only deserialize and restructure data, not transform it
+- **Use subflavors for dispatch**: Let the framework choose the right cooker
+
+### 3. TaskEncoder design
+- **Single responsibility**: One task encoder per data format family
+- **Minimal `batch()`**: Only stack tensors, avoid computation in the batch method
+- **Separate concerns**: Data loading in Cooker, encoding in model forward_step
+
+### 4. Dataloader configuration
+- **Num workers**: Match CPU cores (typically 4-8)
+- **Batch size**: Max out GPU memory
+- **Shuffle**: Always true for training
+- **Drop last**: True to avoid irregular batches
+
+### 5. Debugging
+- **Small dataset**: Test with 100 samples first (`--max-samples 100`)
+- **Single worker**: Set `num_workers=1` for debugging
+- **Validate**: Run `python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset`
+
+---
+
+## Customization patterns
+
+### Custom cooker function
+
+Add a new cooker for a different data format:
+
+```python
+from megatron.energon import Cooker, basic_sample_keys, stateless
+
+@stateless
+def cook_my_custom_format(sample: dict) -> DiffusionSample:
+    """Custom cooker for a different data layout."""
+    latents = torch.load(io.BytesIO(sample['vae_output.pth']), map_location='cpu')
+    prompt = torch.load(io.BytesIO(sample['text_embed.pth']), map_location='cpu')
+    pooled = torch.load(io.BytesIO(sample['clip_embed.pth']), map_location='cpu')
+
+    return DiffusionSample(
+        **basic_sample_keys(sample),
+        latents=latents,
+        prompt_embeds=prompt,
+        pooled_prompt_embeds=pooled,
+    )
+
+class CustomDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]):
+    decoder = SampleDecoder(image_decode="pil")
+    cookers = [
+        Cooker(cook_my_custom_format, has_subflavors={"encoding": "custom_v2"}),
+    ]
+
+    def batch(self, samples):
+        return {
+            'latents': torch.stack([s.latents for s in samples]),
+            'prompt_embeds': torch.stack([s.prompt_embeds for s in samples]),
+            'pooled_prompt_embeds': torch.stack([s.pooled_prompt_embeds for s in samples]),
+        }
+```
+
+### Multiple cookers in one TaskEncoder
+
+A single TaskEncoder can register multiple cookers for different subflavors:
+
+```python
+class MultiFormatTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]):
+    cookers = [
+        Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}),
+        Cooker(cook_my_custom_format, has_subflavors={"encoding": "custom_v2"}),
+    ]
+```
+
+---
+
+## References
+
+- **Megatron-Energon**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) multimodal examples
+- **WebDataset**: [webdataset/webdataset](https://github.com/webdataset/webdataset)
+- **NeMo Implementation**: `nemo/collections/diffusion/data/`
+- **Primus TaskEncoders**: `primus/backends/megatron/data/diffusion/task_encoders/image.py`
+- **Primus Dataloader**: `primus/backends/megatron/data/dataloader.py`
+
+---
+
+**Last Updated**: March 2026
diff --git a/docs/04-technical-guides/diffusion-models/flux_architecture.md b/docs/04-technical-guides/diffusion-models/flux_architecture.md
new file mode 100644
index 000000000..61dcb6bb6
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/flux_architecture.md
@@ -0,0 +1,926 @@
+# Flux architecture deep dive
+
+## Table of contents
+
+1. [Overview](#overview)
+2. [Architecture principles](#architecture-principles)
+3. [Model components](#model-components)
+4. [Data flow](#data-flow)
+5. [Mathematical formulation](#mathematical-formulation)
+6. [Implementation details](#implementation-details)
+7. [Megatron-core integration](#megatron-core-integration)
+8. [Performance optimizations](#performance-optimizations)
+
+---
+
+## Overview
+
+Flux is a **flow-based diffusion model** for high-quality text-to-image generation. It uses an innovative **MMDiT (Multimodal Diffusion Transformer)** architecture that jointly processes image and text tokens through shared transformer blocks.
+
+### Key innovations
+
+1. **Flow Matching**: Uses rectified flow instead of traditional diffusion
+2. **MMDiT Architecture**: Joint image-text attention in early layers
+3. **3D RoPE**: Multi-dimensional rotary position embeddings for spatial awareness
+4. **Two-Stage Processing**: Joint layers followed by image-only layers
+
+### Model variants
+
+| Variant | Joint Layers | Single Layers | Parameters | Use Case |
+|---------|--------------|---------------|------------|----------|
+| Flux 535M | 1 | 1 | ~535M | Development, testing |
+| Flux 12B | 19 | 38 | ~12B | Production deployment |
+
+---
+
+## Architecture principles
+
+### 1. Flow matching framework
+
+Unlike traditional diffusion (which adds Gaussian noise), Flux uses **rectified flow**:
+
+```
+Forward process:  z_t = (1-t) * z_0 + t * z_1
+where:
+  - z_0 = original image (latent)
+  - z_1 = random noise
+  - t ∈ [0, 1] is the flow timestep
+```
+
+The model predicts the **velocity field** v_θ:
+
+```
+v_θ(z_t, t, c) ≈ z_1 - z_0
+```
+
+**Advantages**:
+- Straight-line interpolation paths (more efficient than diffusion curves)
+- Faster sampling (fewer steps needed)
+- Better training stability
+
+### 2. MMDiT (multimodal diffusion transformer)
+
+Traditional DiT processes image tokens independently. MMDiT jointly processes image and text:
+
+```
+┌─────────────┐     ┌─────────────┐
+│ Image Tokens│     │ Text Tokens │
+└──────┬──────┘     └──────┬──────┘
+       │                   │
+       └───────┬───────────┘
+               │
+        ┌──────▼──────┐
+        │ Joint Attn  │ ← Cross-attend image ↔ text
+        └──────┬──────┘
+               │
+        ┌──────▼──────┐
+        │   Split     │
+        └──┬───────┬──┘
+           │       │
+   ┌───────▼──┐ ┌─▼────────┐
+   │ Image MLP│ │ Text MLP │ ← Separate processing
+   └──────────┘ └──────────┘
+```
+
+**Benefits**:
+- Better text-image alignment
+- Richer cross-modal interactions
+- Improved compositional understanding
+
+### 3. Two-stage processing
+
+Flux uses a unique two-stage architecture:
+
+**Stage 1: Joint Processing (Double Blocks)**
+- Both image and text tokens
+- Multi-modal attention
+- Rich semantic understanding
+
+**Stage 2: Image Refinement (Single Blocks)**
+- Image tokens only (text concatenated but not separated)
+- Focus on spatial coherence
+- Fine-grained detail generation
+
+---
+
+## Model components
+
+### 1. Input embeddings
+
+#### Image path
+
+```
+Image (RGB) → VAE Encoder → Latents [B, 64, H/8, W/8]
+                           ↓
+                    Patchify + Linear
+                           ↓
+              Image Tokens [B, H*W, 3072]
+```
+
+**VAE**: AutoencoderKL (8x downsampling)
+- Input: 1024×1024 RGB image
+- Output: 64×128×128 latent
+
+**Linear Projection**: Maps 64 channels → 3072 hidden dim
+
+#### Text path
+
+```
+Caption → T5-XXL Encoder → Embeddings [B, S, 4096]
+                          ↓
+                    Linear Projection
+                          ↓
+              Text Tokens [B, S, 3072]
+
+Caption → CLIP-L Encoder → Pooled [B, 768]
+                          ↓
+                    MLP Embedder
+                          ↓
+                Vector Embedding [B, 3072]
+```
+
+**T5-XXL**: Context-rich text embeddings (max 512 tokens)
+**CLIP-L**: Global style/semantic vector (768 dim)
+
+#### Conditioning
+
+```
+Timestep t ∈ [0, 1] → Sinusoidal Encoding → [B, 256]
+                                            ↓
+                                          MLP
+                                            ↓
+                              Timestep Embedding [B, 3072]
+
+(Optional) Guidance scale g → Linear → [B, 3072]
+```
+
+**Combined Conditioning Vector**:
+```
+vec = timestep_emb + clip_pooled_emb + [guidance_emb]
+```
+
+### 2. Position embeddings (3D RoPE)
+
+Flux uses **3D Rotary Position Embeddings** for spatial awareness:
+
+**Axes**:
+1. **Axis 0**: Channel groups (16 groups for 64 channels)
+2. **Axis 1**: Height positions (e.g., 128 for 1024px)
+3. **Axis 2**: Width positions (e.g., 128 for 1024px)
+
+**Implementation**:
+```python
+# Generate position IDs for each axis
+pos_ids = [
+    (h * W + w) // (16 * patch_size),  # Axis 0: channel group
+    h,                                  # Axis 1: height
+    w,                                  # Axis 2: width
+]
+
+# Compute frequencies for each axis
+theta_i = theta ^ (2i / dim_axis)
+freqs = pos_id / theta_i
+
+# Apply RoPE rotation
+cos_freq = cos(freqs)
+sin_freq = sin(freqs)
+```
+
+**Advantages**:
+- Encodes spatial structure (height × width)
+- Encodes channel relationships
+- Works for any resolution (generalization)
+
+### 3. MMDiT layer (double block)
+
+Each MMDiT layer performs:
+
+```
+Input: img [B, H*W, D], txt [B, S, D], vec [B, D]
+
+1. Pre-normalization (AdaLN with timestep conditioning)
+   img_norm = AdaLN(img, vec)
+   txt_norm = AdaLN(txt, vec)
+
+2. Joint Self-Attention
+   img_qkv = Linear(img_norm)  # [B, H*W, 3*D]
+   txt_qkv = Linear(txt_norm)  # [B, S, 3*D]
+
+   # Concatenate for joint attention
+   joint_qkv = concat([img_qkv, txt_qkv], dim=1)  # [B, H*W+S, 3*D]
+
+   # Apply attention
+   joint_out = Attention(joint_qkv)  # [B, H*W+S, D]
+
+   # Split back
+   img_attn, txt_attn = split(joint_out, [H*W, S])
+
+3. Gated Residual Addition
+   img = img + gate_img * img_attn
+   txt = txt + gate_txt * txt_attn
+
+4. Feed-Forward Networks (separate for img and txt)
+   img_mlp = AdaLN(img, vec) → Linear → GELU → Linear
+   txt_mlp = AdaLN(txt, vec) → Linear → GELU → Linear
+
+   img = img + gate_img_mlp * img_mlp
+   txt = txt + gate_txt_mlp * txt_mlp
+
+Output: img [B, H*W, D], txt [B, S, D]
+```
+
+**Key Features**:
+- **Shared attention space**: Image and text attend to each other
+- **Adaptive gating**: Timestep-conditioned residual connections
+- **Separate MLPs**: Modality-specific processing
+
+### 4. Flux single block
+
+After joint processing, image tokens go through single blocks:
+
+```
+Input: img [B, H*W, D], txt [B, S, D], vec [B, D]
+
+1. Concatenate (but don't split later)
+   combined = concat([img, txt], dim=1)  # [B, H*W+S, D]
+
+2. Pre-normalization (AdaLN)
+   combined_norm = AdaLN(combined, vec)
+
+3. Self-Attention
+   qkv = Linear(combined_norm)  # [B, H*W+S, 3*D]
+   attn_out = Attention(qkv)     # [B, H*W+S, D]
+
+4. Gated Residual
+   combined = combined + gate * attn_out
+
+5. Feed-Forward
+   mlp_norm = AdaLN(combined, vec)
+   mlp_out = Linear → GELU → Linear
+   combined = combined + gate_mlp * mlp_out
+
+6. Extract image tokens
+   img = combined[:, :H*W, :]  # Only use image part
+
+Output: img [B, H*W, D]  (text tokens discarded)
+```
+
+**Rationale**:
+- Text still influences attention (as keys/values)
+- Output focuses on image generation
+- More efficient than full joint processing
+
+### 5. Output processing
+
+```
+Image Tokens [B, H*W, D]
+        ↓
+    AdaLNContinuous(vec)  ← Final timestep conditioning
+        ↓
+    Linear Projection
+        ↓
+     [B, H*W, 64]
+        ↓
+    Reshape
+        ↓
+  [B, 64, H, W]  ← Predicted velocity field
+```
+
+---
+
+## Data flow
+
+### Complete forward pass
+
+```
+                            Input
+                              │
+          ┌───────────────────┼───────────────────┐
+          │                   │                   │
+        Image               Text               Timestep
+    [B,64,H,W]           [B,S,4096]             [B]
+          │                   │                   │
+          ▼                   ▼                   ▼
+     img_linear          txt_linear          time_emb
+          │                   │                   │
+          ├─────── + ─────────┴─────── vec ──────┤
+          │              (conditioning)           │
+          ▼                                       │
+    img [B,H*W,3072]                             │
+    txt [B,S,3072]                               │
+          │                                       │
+          │              ┌──────────┐            │
+          └──────────────► MMDiT    ├────◄───────┤
+          ┌──────────────◄ Layer 1  ├────────────┘
+          │              └──────────┘
+          │              ┌──────────┐
+          └──────────────► MMDiT    ├────◄───────┐
+          ┌──────────────◄ Layer 2  ├────────────┤
+          │              └──────────┘            │
+          ...                                    vec
+          │              ┌──────────┐            │
+          └──────────────► MMDiT    ├────◄───────┘
+          ┌──────────────◄ Layer N  ├────────────┐
+          │              └──────────┘            │
+          │                                       │
+    img [B,H*W,3072]                             │
+    txt [B,S,3072]                               │
+          │                                       │
+          │              ┌──────────┐            │
+          └──────────────► Single   ├────◄───────┤
+          ┌──────────────◄ Block 1  ├────────────┘
+          │              └──────────┘
+          │              ┌──────────┐
+          └──────────────► Single   ├────◄───────┐
+          ┌──────────────◄ Block 2  ├────────────┤
+          │              └──────────┘            │
+          ...                                    vec
+          │              ┌──────────┐            │
+          └──────────────► Single   ├────◄───────┘
+          ┌──────────────◄ Block M  ├────────────┐
+          │              └──────────┘            │
+          │                                       │
+    img [B,H*W,3072]                             │
+          │                                       │
+          ▼                                       │
+    AdaLNContinuous ◄───────────────────────────┘
+          │
+          ▼
+    Linear Projection
+          │
+          ▼
+    Reshape to [B,64,H,W]
+          │
+          ▼
+   Predicted Velocity
+```
+
+### Training data flow
+
+```
+Original Image
+      │
+      ▼
+   VAE Encode → z_0 [B,64,H,W]
+      │
+      ├─────────────┐
+      │             │
+      ▼             ▼
+   z_0        Sample Noise → z_1
+      │             │
+      │    Sample t ~ Uniform(0,1)
+      │             │
+      └──────┬──────┘
+             │
+        z_t = (1-t)*z_0 + t*z_1  ← Noisy latent
+             │
+             ▼
+      Flux Model(z_t, text, t)
+             │
+             ▼
+       v_pred [B,64,H,W]  ← Predicted velocity
+             │
+             ▼
+    Loss = MSE(v_pred, z_1 - z_0)  ← Flow matching loss
+```
+
+---
+
+## Mathematical formulation
+
+### Flow matching objective
+
+**Forward Process**:
+```
+z_t = (1 - t) * z_0 + t * z_1,  where t ~ Uniform(0, 1)
+```
+
+**Velocity Target**:
+```
+v* = dz_t/dt = z_1 - z_0
+```
+
+**Training Loss**:
+```
+L = E_{z_0, z_1, t, c} [ ||v_θ(z_t, t, c) - (z_1 - z_0)||² ]
+
+where:
+  z_0 = VAE(image)                    # Original latent
+  z_1 ~ N(0, I)                       # Random noise
+  c = (text_emb, clip_pooled)         # Conditioning
+  v_θ = Flux model                    # Predicted velocity
+```
+
+### Sampling (inference)
+
+**Euler Integration** (first-order ODE solver):
+```
+z_0 = z_1  # Start from noise
+for t in [1.0, 0.9, ..., 0.1, 0.0]:
+    v_t = Flux(z_t, t, c)
+    z_{t-Δt} = z_t - Δt * v_t
+```
+
+**Higher-Order Solvers** (optional):
+- Heun's method (2nd order)
+- DPM-Solver (adaptive)
+
+### Classifier-free guidance
+
+During inference, use guidance scale `w`:
+
+```
+v_guided = v_uncond + w * (v_cond - v_uncond)
+
+where:
+  v_cond = Flux(z_t, t, c_text)       # With text
+  v_uncond = Flux(z_t, t, c_empty)    # Without text (null prompt)
+  w = guidance scale (typically 3-5)
+```
+
+**Implementation**: Use guidance embedding in config:
+```python
+config = FluxConfig(guidance_embed=True, guidance_scale=3.5)
+```
+
+### 3D RoPE mathematics
+
+For position `(h, w)` in image:
+
+**Position IDs**:
+```
+pid = [floor((h*W + w) / 16), h, w]  # [channel_group, height, width]
+```
+
+**Frequencies**:
+```
+θ_i = θ_base ^ (2i / d_axis),  for i = 0, ..., d_axis/2
+
+freq_{axis,i} = pid[axis] / θ_i
+```
+
+**RoPE Rotation**:
+```
+q_rot = [q[:d/2] * cos(freq) - q[d/2:] * sin(freq),
+         q[:d/2] * sin(freq) + q[d/2:] * cos(freq)]
+
+k_rot = [k[:d/2] * cos(freq) - k[d/2:] * sin(freq),
+         k[:d/2] * sin(freq) + k[d/2:] * cos(freq)]
+```
+
+---
+
+## Implementation details
+
+### Memory layout
+
+Megatron-Core uses **sequence-first format**: `[seq, batch, hidden]`
+
+**Conversions**:
+```python
+# User format: [B, C, H, W]
+img_latents = torch.randn(B, 64, H, W)
+
+# Reshape to [B, H*W, 64]
+img_seq = rearrange(img_latents, 'b c h w -> b (h w) c')
+
+# Project to hidden_size
+img_tokens = linear(img_seq)  # [B, H*W, 3072]
+
+# Convert to Megatron format: [seq, batch, hidden]
+img_megatron = rearrange(img_tokens, 'b s d -> s b d')
+```
+
+### Adaptive layer normalization
+
+**Standard AdaLN**:
+```python
+class AdaLN:
+    def forward(self, timestep_emb):
+        modulation = MLP(SiLU(timestep_emb))
+        shift, scale, gate = split(modulation, 3)
+        return shift, scale, gate
+
+    @staticmethod
+    def modulate(x, shift, scale):
+        return LayerNorm(x) * (1 + scale) + shift
+```
+
+**Usage in Layer**:
+```python
+shift, scale, gate = adaln(timestep_emb)
+x_norm = AdaLN.modulate(x, shift, scale)
+x_attn = attention(x_norm)
+x = x + gate * x_attn  # Gated residual
+```
+
+### Attention implementation
+
+**Using Megatron SelfAttention**:
+```python
+from megatron.core.transformer.attention import SelfAttention
+
+attn = SelfAttention(
+    config=config,
+    submodules=submodules,
+    layer_number=layer_idx,
+    attn_mask_type=AttnMaskType.no_mask,  # Flux uses no mask
+)
+
+# Megatron expects [seq, batch, hidden]
+output = attn(hidden_states)
+```
+
+**Joint Attention Trick**:
+```python
+# Concatenate image and text
+combined = torch.cat([img, txt], dim=0)  # [seq_img+seq_txt, batch, hidden]
+
+# Single attention call processes both
+joint_output = attention(combined)
+
+# Split back
+img_out = joint_output[:seq_img]
+txt_out = joint_output[seq_img:]
+```
+
+---
+
+## Megatron-core integration
+
+### TransformerConfig
+
+Flux uses Megatron's `TransformerConfig`:
+
+```python
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+config = TransformerConfig(
+    num_layers=19 + 38,  # joint + single
+    hidden_size=3072,
+    num_attention_heads=24,
+    ffn_hidden_size=3072 * 4,  # Standard 4x expansion
+    layernorm_epsilon=1e-6,
+    hidden_dropout=0.0,
+    attention_dropout=0.0,
+    add_qkv_bias=True,
+    # ... other Megatron params
+)
+```
+
+### Layer specs
+
+Flux provides factory functions for layer specs:
+
+```python
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    get_flux_double_transformer_spec_for_backend,
+    get_flux_single_transformer_spec_for_backend,
+    get_flux_layer_spec,
+)
+
+# For MMDiT layers (pass backend from config)
+double_spec = get_flux_double_transformer_spec_for_backend(backend)
+
+# For single blocks
+single_spec = get_flux_single_transformer_spec_for_backend(backend)
+
+# Or use high-level API for full TransformerBlock
+layer_specs = get_flux_layer_spec(config, backend=backend)
+```
+
+### Distributed training support
+
+Flux inherits Megatron's parallelism:
+
+**Tensor Parallelism** (TP):
+```python
+config = TransformerConfig(
+    tensor_model_parallel_size=8,  # 8-way TP
+    sequence_parallel=True,        # Sequence parallelism
+)
+```
+
+**Pipeline Parallelism** (PP): not supported for diffusion models. The forward
+path runs embeddings/output head on every rank and does not relay activations
+between stages, so `pipeline_model_parallel_size` must be 1 (PP > 1 is rejected
+at config construction).
+
+**Data Parallelism**: Handled automatically by trainer
+
+---
+
+## Performance optimizations
+
+### 1. Transformer engine
+
+Flux uses NVIDIA Transformer Engine for FP8 training:
+
+```python
+from megatron.core.extensions.transformer_engine import (
+    TELayerNormColumnParallelLinear,
+    TERowParallelLinear,
+)
+
+# Automatically used in layer specs
+linear_qkv = TELayerNormColumnParallelLinear(...)
+linear_proj = TERowParallelLinear(...)
+```
+
+**Benefits**:
+- FP8 matmul (faster, less memory)
+- Fused operations (LayerNorm + Linear)
+- Automatic scaling for numerical stability
+
+### 2. Flash attention
+
+Enabled via Megatron:
+
+```python
+config = TransformerConfig(
+    attention_type='flash_attention',  # Use Flash Attention 2
+)
+```
+
+**Speedup**: 2-3x faster attention, 4x less memory
+
+### 3. Gradient checkpointing
+
+For large models:
+
+```python
+config = TransformerConfig(
+    recompute_granularity='selective',  # Checkpoint expensive ops
+    recompute_method='uniform',
+    recompute_num_layers=19,  # Checkpoint all joint layers
+)
+```
+
+**Memory Savings**: ~40% reduction, ~20% slower
+
+### 4. Fused operations
+
+```python
+config = TransformerConfig(
+    bias_activation_fusion=True,      # Fuse bias + activation
+    masked_softmax_fusion=True,       # Fuse mask + softmax
+    gradient_accumulation_fusion=True, # Fuse grad accumulation
+)
+```
+
+### 5. Mixed precision
+
+```python
+from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss
+
+# Using PyTorch autocast
+with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+    output = model(img, txt, y, timesteps, img_ids, txt_ids)
+    target = noise - clean_latents
+    loss = compute_flow_matching_loss(output, clean_latents, noise)
+
+# Backward pass (handled automatically)
+scaler = torch.cuda.amp.GradScaler()
+scaler.scale(loss).backward()
+scaler.step(optimizer)
+scaler.update()
+```
+
+---
+
+## Comparison with other models
+
+### Flux vs DiT
+
+| Aspect | DiT | Flux |
+|--------|-----|------|
+| Architecture | Image-only transformer | MMDiT (image + text joint) |
+| Conditioning | AdaLN (injected) | Joint attention |
+| Position Encoding | Learned 2D | 3D RoPE |
+| Diffusion Type | DDPM | Flow matching |
+
+### Flux vs stable diffusion
+
+| Aspect | Stable Diffusion (UNet) | Flux (Transformer) |
+|--------|-------------------------|---------------------|
+| Backbone | UNet with ResNet blocks | Full transformer |
+| Text Integration | Cross-attention | Joint self-attention |
+| Scalability | Limited (U-Net bottleneck) | Excellent (transformer scaling) |
+| Efficiency | Fast (fewer params) | Slower but higher quality |
+
+---
+
+## Design decisions
+
+### Why two-stage (joint + single)?
+
+**Joint Blocks**:
+- Deep semantic understanding
+- Text-image alignment
+- Compositional reasoning
+
+**Single Blocks**:
+- Spatial refinement
+- Detail generation
+- Efficient (no text processing)
+
+**Alternative**: All joint blocks → slower, marginal quality gain
+
+### Why flow matching?
+
+**Advantages over DDPM**:
+1. **Simpler training**: Straight-line interpolation (no schedule design)
+2. **Faster sampling**: Fewer steps (10-20 vs 50-100)
+3. **Better mode coverage**: Straighter paths → less error accumulation
+
+**Math**: Rectified flow is ODE-based (vs SDE for DDPM)
+
+### Why 3D RoPE?
+
+**Advantages**:
+1. **Spatial awareness**: Encodes (height, width) structure
+2. **Resolution flexibility**: Works for any image size
+3. **Channel grouping**: Models relationships between channels
+
+**Alternative**: Learned absolute embeddings → less flexible
+
+---
+
+## Primus implementation highlights
+
+### TransformerBlock architecture
+
+Primus's key architectural enhancement is the use of Megatron-Core's **TransformerBlock with heterogeneous layer specifications**:
+
+```mermaid
+graph TD
+    subgraph traditional[Traditional Approach]
+        A[Model] --> B[double_blocks: ModuleList]
+        A --> C[single_blocks: ModuleList]
+        B --> D[Manual iteration]
+        C --> D
+        D --> E[Manual PP splitting]
+    end
+
+    subgraph primus[Primus Approach]
+        F[Model] --> G[transformer: TransformerBlock]
+        G --> H[layer_specs: heterogeneous]
+        H --> J[Unified checkpoint format]
+    end
+```
+
+**Benefits**:
+1. **Unified Checkpointing**: Single `transformer.layers.{0-56}` namespace
+2. **Future-Proof**: Native support for new Megatron-Core features
+3. **Cleaner Code**: No manual iteration over separate block lists
+
+### Layer specification pattern
+
+```python
+# Primus approach
+layer_specs = get_flux_layer_spec(config, backend=backend)
+# Or manually:
+# layer_specs = [
+#     *[get_flux_double_transformer_spec_for_backend(backend) for _ in range(19)],
+#     *[get_flux_single_transformer_spec_for_backend(backend) for _ in range(38)],
+# ]
+
+transformer = TransformerBlock(
+    config=config,
+    spec=TransformerBlockSubmodules(layer_specs=layer_specs)
+)
+
+# Automatic PP slicing handled by TransformerBlock
+# No manual offset calculation needed
+```
+
+### Checkpoint format comparison
+
+**Traditional Format**:
+```
+double_blocks.0.attn.qkv.weight
+double_blocks.18.mlp.fc2.bias
+single_blocks.0.attn.qkv.weight
+single_blocks.37.mlp.fc2.bias
+```
+
+**Primus Format** (TransformerBlock):
+```
+transformer.layers.0.self_attention.linear_qkv.weight   # Joint layer 0
+transformer.layers.18.mlp.linear_fc2.bias               # Joint layer 18
+transformer.layers.19.self_attention.linear_qkv.weight  # Single layer 0
+transformer.layers.56.mlp.linear_fc2.bias               # Single layer 37
+```
+
+Benefits: Simpler distributed checkpointing, easier layer inspection, consistent with Megatron GPT models.
+
+---
+
+## Future enhancements
+
+### Planned features
+
+1. **ControlNet Support**:
+   - Spatial conditioning (pose, depth, edges)
+   - Residual connections from control encoder
+
+2. **Multi-Resolution Training**:
+   - Dynamic image sizes during training
+   - Aspect ratio bucketing
+
+3. **Efficient Sampling**:
+   - DPM-Solver integration
+   - Distillation for 1-step generation
+
+4. **LoRA Fine-Tuning**:
+   - Low-rank adaptation for custom styles
+   - Efficient personalization
+
+### Research directions
+
+- **Sparse Attention**: Reduce quadratic complexity for high-res
+- **Mixture of Experts**: Conditional computation for efficiency
+- **3D Extension**: Video generation with Flux architecture
+
+---
+
+## References
+
+### Papers
+
+1. **Flow Matching**: Lipman et al., "Flow Matching for Generative Modeling", 2022
+2. **Rectified Flow**: Liu et al., "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", 2022
+3. **DiT**: Peebles & Xie, "Scalable Diffusion Models with Transformers", 2023
+4. **RoPE**: Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding", 2021
+5. **Transformer Engine**: NVIDIA, "Transformer Engine: Accelerating Transformer Training", 2022
+
+### Code references
+
+- **Primus Flux**: `primus/backends/megatron/core/models/diffusion/flux/`
+- **Megatron-Core**: `megatron/core/transformer/`
+- **Transformer Engine**: `transformer_engine/pytorch/`
+- **Official Flux**: Black Forest Labs (HuggingFace)
+
+---
+
+## Appendix: Hyperparameters
+
+### Flux 535M training
+
+```yaml
+model:
+  num_joint_layers: 1
+  num_single_layers: 1
+  hidden_size: 3072
+  num_attention_heads: 24
+
+training:
+  batch_size: 256  # global
+  learning_rate: 1e-4
+  weight_decay: 0.01
+  lr_schedule: cosine
+  warmup_steps: 10000
+  total_steps: 500000
+
+optimization:
+  optimizer: AdamW
+  beta1: 0.9
+  beta2: 0.999
+  epsilon: 1e-8
+  gradient_clip: 1.0
+```
+
+### Flux 12B training
+
+```yaml
+model:
+  num_joint_layers: 19
+  num_single_layers: 38
+  hidden_size: 3072
+  num_attention_heads: 24
+
+training:
+  batch_size: 2048  # global, multi-node
+  learning_rate: 1e-4
+  weight_decay: 0.01
+  lr_schedule: cosine
+  warmup_steps: 10000
+  total_steps: 1000000
+
+optimization:
+  optimizer: AdamW
+  beta1: 0.9
+  beta2: 0.95
+  epsilon: 1e-8
+  gradient_clip: 1.0
+
+parallelism:
+  tensor_parallel: 8
+  pipeline_parallel: 4
+  data_parallel: 8
+  sequence_parallel: True
+```
+
+---
+
+*For API usage, see [api_reference.md](api_reference.md).*
diff --git a/docs/04-technical-guides/diffusion-models/fp8_training.md b/docs/04-technical-guides/diffusion-models/fp8_training.md
new file mode 100644
index 000000000..9067bf201
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/fp8_training.md
@@ -0,0 +1,514 @@
+# FP8 training for Flux models
+
+Complete guide for training Flux diffusion models with FP8 (8-bit floating point) precision on AMD MI300X GPUs using Transformer Engine's delayed scaling recipe.
+
+## Overview
+
+FP8 training provides significant memory and speed improvements while maintaining numerical stability through delayed scaling:
+
+- **~2x memory reduction** (activations and weights)
+- **1.5-2x training speedup** on AMD MI300X
+- **Maintains numerical stability** via delayed scaling
+- **Enables larger batch sizes** or higher resolutions
+
+## Table of contents
+
+- [Prerequisites](#prerequisites)
+- [Quick start](#quick-start)
+- [Configuration](#configuration)
+- [Performance benchmarks](#performance-benchmarks)
+- [Troubleshooting](#troubleshooting)
+- [Best practices](#best-practices)
+- [AMD MI300X specific](#amd-mi300x-specific)
+
+---
+
+## Prerequisites
+
+### Hardware requirements
+
+- **AMD MI300X GPUs** with ROCm 6.0+ support
+- **Minimum GPUs:**
+  - Flux 535M: 1x MI300X (testing)
+  - Flux 12B: 2x MI300X with TP=2 (can train with FP8)
+
+### Software requirements
+
+1. **ROCm 6.0+** with FP8 tensor core support
+2. **Transformer Engine 2.1.0+** with ROCm backend
+3. **PyTorch** with ROCm support
+4. **Megatron-LM** (included in Primus)
+
+### Verification
+
+Verify your environment has:
+- Transformer Engine 2.1.0+ with ROCm backend
+- FP8 support (run `python3 -c "import transformer_engine.pytorch as te; print(te.fp8.is_fp8_available())"`)
+
+---
+
+## Quick start
+
+### Test FP8 with Flux 535M (recommended first step)
+
+```bash
+# 1. Prepare test dataset (or use existing)
+# See primus/configs/data/megatron/diffusion/README.md
+
+# 2. Train Flux 535M with FP8
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \
+GPUS_PER_NODE=1 \
+bash examples/run_pretrain.sh
+```
+
+### Production training with Flux 12B
+
+```bash
+# After validating with 535M, scale to 12B (TransformerEngine FP8)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml \
+GPUS_PER_NODE=8 \
+NNODES=4 \
+bash examples/run_slurm_pretrain.sh
+
+# Or local-spec FP8 (no TransformerEngine dependency)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \
+GPUS_PER_NODE=8 \
+NNODES=4 \
+bash examples/run_slurm_pretrain.sh
+```
+
+---
+
+## Configuration
+
+### FP8 model configuration
+
+FP8 is configured at the model level. Two pre-configured files are available:
+
+- `primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml`
+- `primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml`
+
+**Key FP8 Parameters:**
+
+```yaml
+# Enable FP8
+fp8: "e4m3"  # E4M3 format (recommended)
+             # Alternative: "hybrid" (E4M3 activations + E5M2 gradients)
+
+# FP8 Recipe
+fp8_recipe: "delayed"  # Delayed scaling (most stable)
+                       # Alternatives: "tensorwise", "blockwise", "mxfp8"
+
+# Scaling Configuration
+fp8_margin: 0  # Margin for scaling factor (0 = no margin)
+fp8_amax_history_len: 1024  # History window for delayed scaling
+                            # Larger = more stable, smaller = adapts faster
+fp8_amax_compute_algo: "most_recent"  # or "max"
+
+# Gradient Precision
+fp8_wgrad: true  # Enable FP8 for weight gradients (recommended)
+
+# Attention Precision
+fp8_dot_product_attention: false  # Keep attention in higher precision
+fp8_multi_head_attention: false   # Keep MHA in higher precision
+```
+
+### Training configuration adjustments
+
+**Batch Sizes with FP8:**
+
+```yaml
+# Flux 12B - can increase batch size with FP8 memory savings
+micro_batch_size: 2  # vs 1 for BF16
+global_batch_size: 256  # same as BF16
+
+# Flux 535M - can increase significantly
+micro_batch_size: 4  # vs 2 for BF16
+global_batch_size: 32
+```
+
+**Optimizer Settings (Same as BF16):**
+
+```yaml
+optimizer: adamw
+lr: 1.0e-4
+min_lr: 1.0e-5
+weight_decay: 0.01
+clip_grad: 1.0  # Gradient clipping still important!
+```
+
+**Parallelism with FP8:**
+
+```yaml
+# Flux 12B - can potentially reduce TP with FP8
+tensor_model_parallel_size: 2  # or reduce to 1 with FP8
+pipeline_model_parallel_size: 1
+context_parallel_size: 1
+```
+
+---
+
+## Performance benchmarks
+
+### Memory usage
+
+| Model | Precision | Memory/GPU | Batch Size | Notes |
+|-------|-----------|------------|------------|-------|
+| Flux 535M | BF16 | ~7-10GB | 2 | Baseline |
+| Flux 535M | FP8 | ~3-5GB | 4 | ~50% reduction |
+| Flux 12B | BF16 | ~40-50GB | 1 | TP=2 required |
+| Flux 12B | FP8 | ~20-25GB | 2 | TP=2, ~50% reduction |
+
+### Training speed
+
+| Model | Precision | Steps/sec | Speedup | Hardware |
+|-------|-----------|-----------|---------|----------|
+| Flux 535M | BF16 | ~20-30 | 1.0x | 1x MI300X |
+| Flux 535M | FP8 | ~30-50 | 1.5-2x | 1x MI300X |
+| Flux 12B | BF16 | ~0.5-1.0 | 1.0x | 32x MI300X |
+| Flux 12B | FP8 | ~0.8-1.5 | 1.5-2x | 32x MI300X |
+
+### Expected results
+
+- **Memory:** ~50% reduction vs BF16
+- **Speed:** 1.5-2x faster training
+- **Quality:** Loss curves within 5% of BF16
+- **Convergence:** Similar or faster than BF16
+
+---
+
+## Troubleshooting
+
+### NaN or inf in losses
+
+**Problem:** Training becomes unstable with NaN/Inf values
+
+**Solutions:**
+
+1. **Increase scaling history:**
+   ```yaml
+   fp8_amax_history_len: 2048  # or 4096
+   ```
+
+2. **Disable FP8 for weight gradients:**
+   ```yaml
+   fp8_wgrad: false
+   ```
+
+3. **Use more conservative scaling:**
+   ```yaml
+   fp8_amax_compute_algo: "max"  # instead of "most_recent"
+   ```
+
+4. **Add scaling margin:**
+   ```yaml
+   fp8_margin: 1  # or 2
+   ```
+
+5. **Keep first/last layers in BF16:**
+   ```yaml
+   first_last_layers_bf16: true
+   num_layers_at_start_in_bf16: 2
+   num_layers_at_end_in_bf16: 2
+   ```
+
+### Out of memory even with FP8
+
+**Problem:** Still hitting OOM errors with FP8 enabled
+
+**Solutions:**
+
+1. **Reduce micro batch size:**
+   ```yaml
+   micro_batch_size: 1  # back to minimum
+   ```
+
+2. **Enable gradient checkpointing:**
+   ```yaml
+   recompute_granularity: "selective"  # or "full"
+   recompute_method: "block"
+   ```
+
+3. **Increase tensor parallelism:**
+   ```yaml
+   tensor_model_parallel_size: 4  # distribute more
+   ```
+
+4. **Reduce sequence length:**
+   ```yaml
+   seq_length: 2048  # if applicable
+   ```
+
+### FP8 not available
+
+**Problem:** Setup script shows "FP8 not available"
+
+**Checks:**
+
+1. **Verify GPU model:**
+   ```bash
+   rocm-smi --showproductname
+   # Should show MI300X
+   ```
+
+2. **Check ROCm version:**
+   ```bash
+   rocm-smi --showversion
+   # Should be 6.0+
+   ```
+
+3. **Verify Transformer Engine:**
+   ```bash
+   python3 -c "import transformer_engine; print(transformer_engine.__version__)"
+   # Should be 2.1.0+
+   ```
+
+4. **Test FP8 directly:**
+   ```python
+   import transformer_engine.pytorch as te
+   print(te.fp8.is_fp8_available())  # Should be True
+   ```
+
+### Slower than expected
+
+**Problem:** FP8 training is not faster than BF16
+
+**Checks:**
+
+1. **Verify FP8 is actually enabled:**
+   - Check logs for FP8 context messages
+   - Run with `NCCL_DEBUG=INFO` to see precision info
+
+2. **Check batch size:**
+   - Ensure you increased micro_batch_size with FP8
+   - Small batches may not show speedup
+
+3. **Verify tensor cores:**
+   - FP8 requires tensor core support
+   - Check ROCm driver configuration
+
+4. **Profile training:**
+   ```yaml
+   log_timers_to_tensorboard: true
+   ```
+   - Compare FP8 vs BF16 step times
+
+---
+
+## Best practices
+
+### Recommended workflow
+
+1. **Start with 535M:**
+   - Validate FP8 works correctly
+   - Test for 100-1000 steps
+   - Verify no NaN/Inf
+
+2. **Validate on small 12B run:**
+   - Train for 1000-5000 steps
+   - Compare loss with BF16 baseline
+   - Check memory and speed improvements
+
+3. **Production training:**
+   - Monitor closely for first 10K steps
+   - Watch for numerical issues
+   - Compare checkpoints with BF16
+
+### Training configuration
+
+**Conservative (stable):**
+```yaml
+fp8_recipe: "delayed"
+fp8_amax_history_len: 2048
+fp8_amax_compute_algo: "max"
+fp8_wgrad: false
+```
+
+**Balanced (recommended):**
+```yaml
+fp8_recipe: "delayed"
+fp8_amax_history_len: 1024
+fp8_amax_compute_algo: "most_recent"
+fp8_wgrad: true
+```
+
+**Aggressive (maximum performance):**
+```yaml
+fp8_recipe: "tensorwise"  # Requires TE 2.2.0+
+fp8_amax_history_len: 512
+fp8_amax_compute_algo: "most_recent"
+fp8_wgrad: true
+```
+
+### Monitoring
+
+**Key metrics to watch:**
+
+1. **Loss curves:**
+   - Should be smooth (no spikes)
+   - Should decrease normally
+   - Compare with BF16 baseline
+
+2. **Gradient norms:**
+   - Should be stable
+   - No sudden jumps to infinity
+
+3. **Memory usage:**
+   - Should be ~50% of BF16
+   - Check with `rocm-smi`
+
+4. **Training speed:**
+   - Should be 1.5-2x faster
+   - Measure steps/second
+
+### Checkpointing
+
+- **FP8 checkpoints are compatible with BF16**
+- Can switch between FP8/BF16 training
+- Optimizer state includes FP8 scaling factors
+- Checkpoints are same size as BF16
+
+---
+
+## Autotune (local spec FP8)
+
+> **Scope:** This section covers the **local-spec** FP8 path (`PrimusTurboFloat8LocalSpecProvider`, no TransformerEngine), e.g. `flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml`. The TransformerEngine prerequisites and checks elsewhere in this guide (`te.fp8.is_fp8_available()`, "FP8 Not Available") do **not** apply here -- this path quantizes via Primus Turbo directly.
+
+### Enable autotune (and do not pin the FP8 GEMM backend)
+
+The local-spec FP8 kernels benefit from the Primus-Turbo autotuner, which picks the best backend per GEMM shape. Enable it with `PRIMUS_TURBO_AUTO_TUNE=1`.
+
+`PRIMUS_TURBO_AUTO_TUNE=1` is necessary but **not sufficient**: an explicit `PRIMUS_TURBO_GEMM_BACKEND` short-circuits autotune (the FP8 kernel dispatcher returns the user-specified backend before the autotune step), so it must be unset (or scoped so it does not cover FP8) for autotune to engage.
+
+**Note:** some base images bake `PRIMUS_TURBO_GEMM_BACKEND` as an *empty string* rather than leaving it unset. An empty value is not treated as "unset" and can raise `KeyError ''` on the first FP8 GEMM. If you hit this, `unset PRIMUS_TURBO_GEMM_BACKEND` before launching.
+
+```bash
+unset PRIMUS_TURBO_GEMM_BACKEND          # or scope it so it does not cover FP8
+export PRIMUS_TURBO_AUTO_TUNE=1
+
+EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \
+  bash examples/run_pretrain.sh
+```
+
+### Contrast with MXFP4
+
+For MXFP4/FP4 + AITER with a tuned CSV, do the **opposite**: leave `PRIMUS_TURBO_AUTO_TUNE` unset, because autotune disables the AITER preshuffle fast path. See the [MXFP4 Training Guide](mxfp4_training.md) ("Preshuffle fast path"). Do not copy the MXFP4 env recipe for FP8.
+
+---
+
+## AMD MI300X specific
+
+### Environment variables
+
+```bash
+# Optional: set for better performance
+export HSA_FORCE_FINE_GRAIN_PCIE=1  # Better PCIe performance
+export NCCL_DEBUG=INFO  # For debugging
+export HSA_ENABLE_SDMA=0  # Disable SDMA for stability
+```
+
+### ROCm optimization
+
+1. **HipBLASLt tuning:**
+   ```bash
+   # Generate optimal GEMM kernels for your hardware
+   # See ROCm documentation for details
+   ```
+
+2. **NCCL configuration:**
+   ```bash
+   export NCCL_IB_DISABLE=0  # Enable InfiniBand if available
+   export NCCL_NET_GDR_LEVEL=3  # GPU Direct RDMA
+   ```
+
+3. **Memory management:**
+   ```bash
+   export HSA_OVERRIDE_GFX_VERSION=9.4.2  # For MI300X
+   ```
+
+### Known issues
+
+1. **Transformer Engine ROCm support:**
+   - Verify TE version supports ROCm FP8
+   - Some recipes may require specific TE versions
+
+2. **Numerical stability:**
+   - MI300X may require longer history (fp8_amax_history_len)
+   - Start conservative and tune
+
+3. **Multi-node training:**
+   - Ensure RCCL/NCCL properly configured
+   - Test single-node first
+
+---
+
+## Testing
+
+### Integration test
+
+```bash
+# Quick 100-step validation run
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \
+GPUS_PER_NODE=1 \
+bash examples/run_pretrain.sh
+```
+
+### Convergence test
+
+1. Train both BF16 and FP8 for 5000 steps
+2. Compare loss curves (should be within 5%)
+3. Generate images from checkpoints
+4. Compare quality visually
+
+---
+
+## Numerical verification status
+
+To set expectations for external users, the precision/convergence claims in this
+codebase fall into two tiers:
+
+- **Backed by in-repo tests** (CI-runnable on supported hardware): structural and
+  convention checks—attention TE-vs-local-spec equivalence, RNG/seed
+  determinism, chimera init, VAE resample reproducibility, fused delayed-scale
+  update, and MLPerf warmup FP8 state.
+- **Asserted, not yet backed by an in-repo test:** end-to-end *tensor parity*
+  against HuggingFace/Diffusers FLUX from a real checkpoint, and the exact MLPerf
+  v5.1 eval sample-count / validation-timestep semantics. These are validated by
+  internal reference runs but no committed test reproduces them.
+
+Tracked follow-ups (file as public-repo issues):
+
+1. A real-checkpoint forward-parity test (535M minimum) comparing Primus Flux
+   against HF/Diffusers within a documented tolerance.
+2. A robustness test for the MLPerf validation-timestep fallback path.
+
+Treat any "bit-exact / matches NeMo / matches MLPerf / within X%" statement in
+source comments as *asserted, unverified* until the parity test above lands.
+
+## References
+
+- [Transformer Engine Documentation](https://docs.nvidia.com/deeplearning/transformer-engine/)
+- [AMD ROCm Documentation](https://rocm.docs.amd.com/)
+- [Megatron-LM FP8 Guide](https://github.com/NVIDIA/Megatron-LM/blob/main/docs/llm/fp8.md)
+- [FP8 Formats Explained (E4M3 vs E5M2)](https://arxiv.org/abs/2209.05433)
+
+---
+
+## Support
+
+For issues or questions:
+
+1. Check this guide first
+2. Verify Transformer Engine and FP8 support (see Prerequisites)
+3. Review logs for error messages
+4. File issue with:
+   - Hardware specs (GPU model, ROCm version)
+   - Software versions (TE, PyTorch, Megatron-LM)
+   - Config files used
+   - Error logs
+
+---
+
+**Happy FP8 Training! 🚀**
+
+*Last updated: 2026-01-10*
diff --git a/docs/04-technical-guides/diffusion-models/mxfp4_training.md b/docs/04-technical-guides/diffusion-models/mxfp4_training.md
new file mode 100644
index 000000000..d45f64f52
--- /dev/null
+++ b/docs/04-technical-guides/diffusion-models/mxfp4_training.md
@@ -0,0 +1,210 @@
+# MXFP4 training for Flux models
+
+Guide for training Flux diffusion models in **MXFP4** (E2M1 mantissa + E8M0 block-of-32 scales) on AMD MI355X GPUs using Primus's local-spec MXFP4 implementation backed by Primus-Turbo and AITER.
+
+## Overview
+
+MXFP4 stores activations and weights in 4-bit microscale floating-point with one E8M0 exponent shared per block of 32 elements. The Primus integration:
+
+- Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency**—MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks.
+- Keeps **attention, optimizer state / main params, and inter-rank communication in BF16**. Only the MMA inputs of the column- and row-parallel linears are quantized.
+- Supports two backward modes via `mxfp4_backward_precision`: pure **MXFP4** (default) or **FP8** hybrid (E5M2 backward with tensorwise scaling on HipBLASLt).
+- Dispatches the FP4 GEMM through Primus-Turbo's pluggable backend layer, which can route to either AITER (recommended for MI355X) or HipBLASLt.
+
+## Table of contents
+
+- [Prerequisites](#prerequisites)
+- [Quick start](#quick-start)
+- [Configuration](#configuration)
+- [Primus-Turbo backend selection](#primus-turbo-backend-selection)
+- [Tuned GEMMs](#tuned-gemms)
+- [Troubleshooting](#troubleshooting)
+- [Verification status](#verification-status)
+
+---
+
+## Prerequisites
+
+### Hardware
+
+- **AMD Instinct MI355X** (gfx950) with FP4 tensor-core support. The MXFP4 linear-layer modules assert `check_mxfp4_support()` at construction and will refuse to initialize on unsupported devices ([`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)).
+- Single node (the local-spec layers require `tensor_model_parallel_size: 1`).
+
+### Software
+
+- ROCm-compatible install of `aiter` (provides `aiter.gemm_a4w4` and the tuned-config loader in `aiter/jit/core.py`).
+- Primus-Turbo with FP4 backend registered (`primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl`).
+- `enable_primus_turbo: true` and `use_turbo_attention: true` in the training config.
+
+---
+
+## Quick start
+
+The verified MXFP4 config is `examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`. Launch with the AITER backend and the pre-tuned GEMM CSV:
+
+```bash
+# Path to a checkout of the `tuned_gemm_configs` directory.
+# Set TUNED_GEMM_DIR to wherever you have the tuned configs available.
+export TUNED_GEMM_DIR=${TUNED_GEMM_DIR:-/path/to/tuned_gemm_configs}
+
+export EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml
+export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER
+export AITER_CONFIG_GEMM_A4W4=$TUNED_GEMM_DIR/mi355x/flux_12b.csv
+export AITER_LOG_TUNED_CONFIG=1   # recommended: confirms each shape hits the CSV
+
+bash examples/run_pretrain.sh
+```
+
+The pre-tuned CSV is distributed via an internal tuned-config source (`tuned_gemm_configs/mi355x/flux_12b.csv`). If you do not have access, omit `AITER_CONFIG_GEMM_A4W4` and AITER will fall back to its bundled `a4w4_blockscale_tuned_gemm.csv` (slower for Flux 12B shapes).
+
+---
+
+## Configuration
+
+The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`](../../../examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml):
+
+```yaml
+# MXFP4 precision
+fp4: "mxfp4"
+fp4_recipe: "mxfp4"              # default is "nvfp4" in trainer_base.yaml; must override
+mxfp4_backward_precision: "mxfp4"  # "mxfp4" (pure) or "fp8" (hybrid)
+
+# Local spec + Primus-Turbo
+transformer_impl: "local"
+enable_primus_turbo: true
+use_turbo_attention: true
+
+# Required by the MXFP4 linear-layer modules
+tensor_model_parallel_size: 1
+gradient_accumulation_fusion: false
+# sequence_parallel must remain false
+```
+
+### Knob semantics
+
+| Knob | Values | Notes |
+|------|--------|-------|
+| `fp4` | `"mxfp4"` | Top-level switch to enable FP4. |
+| `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. |
+| `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. |
+| `mxfp4_gradient_stochastic_rounding` | `true` / `false` | Optional. Enables SR on FP4 gradient quantization. |
+
+---
+
+## Primus-Turbo backend selection
+
+The FP4 GEMM call is routed by `GEMMFP4KernelDispatcher` in `Primus-Turbo/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. Backends are selected with the precision-scoped env var `PRIMUS_TURBO_GEMM_BACKEND` (declared in `Primus-Turbo/primus_turbo/common/constants.py`):
+
+```bash
+# Single backend for every precision:
+export PRIMUS_TURBO_GEMM_BACKEND=AITER
+
+# Precision-scoped (recommended): route FP4 GEMMs to AITER, leave others to defaults:
+export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER
+
+# Per-precision routing:
+export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER,FP8:HIPBLASLT
+```
+
+The dispatcher (`GlobalBackendManager` / `AutoKernelDispatcher.dispatch` in `Primus-Turbo/primus_turbo/pytorch/core/backend.py`) resolves the backend in this order: **explicit env > code-set > auto-tune > registered default > fallback**.
+
+### Preshuffle fast path
+
+When **all** of the following are true, MXFP4 GEMMs take the preshuffled fast path with no per-call shuffle overhead:
+
+- `PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER` (or `AITER`) is set.
+- `PRIMUS_TURBO_AUTO_TUNE` is unset or `0`.
+
+The `_enable_preshuffle()` helper in `primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py` returns `True` under these conditions (FP4 backend pinned to AITER and auto-tune off), and the call becomes `aiter.gemm_a4w4(..., bpreshuffle=True)`. This helper reproduces the upstream `enable_preshuffle()` that Primus-Turbo removed in PR #383 ("refactor preshuffle ..."), which moved per-call preshuffle control onto `Float4QuantConfig.use_preshuffle`; Primus keeps the runtime probe locally because `MXFP4LinearFunction` passes a plain `bool` into its custom ops.
+
+> **Do not combine `PRIMUS_TURBO_AUTO_TUNE=1` with a tuned CSV.** Auto-tune disables the preshuffle fast path, so each call pays the shuffle cost while AITER still picks the same kernel internally. For production runs, leave `PRIMUS_TURBO_AUTO_TUNE` unset.
+
+---
+
+## Tuned GEMMs
+
+AITER reads its tuned-GEMM CSV from the `AITER_CONFIG_GEMM_A4W4` env var (handled in `aiter/jit/core.py`; the default is the bundled `aiter/configs/a4w4_blockscale_tuned_gemm.csv`). Each row maps `(cu_num, M, N, K)` to a profiled kernel and split-K factor.
+
+For Flux 12B on MI355X, the pre-tuned CSV is provided by an internal tuned-config source (`tuned_gemm_configs/mi355x/flux_12b.csv`). See that directory's `README.md` for the tuning runbook, CSV schema, ASM-vs-CK kernel distinction, and re-tuning triggers.
+
+### Verifying the CSV is being used
+
+Set `AITER_LOG_TUNED_CONFIG=1`. AITER will log one line per **hit**:
+
+```
+shape is M:16384, N:9216, K:3072, found padded_M: 16384, N:9216, K:3072 is tuned on cu_num = 256 in /path/to/tuned_gemm_configs/mi355x/flux_12b.csv, kernel name is _ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E, splitK is 0!
+```
+
+**Miss** lines are printed unconditionally (no env var required) and look like:
+
+```
+shape is M:..., N:..., K:..., not found tuned config in /path/to/flux_12b.csv, will use default config!
+```
+
+Any miss line means the CSV needs re-tuning for that shape—follow the runbook in `tuned_gemm_configs/README.md`.
+
+### First-run JIT compile
+
+The two `a4w4_blockscale_*_intrawave_v3` CK kernels are JIT-compiled on first use (~2-5 min). Subsequent runs reuse the cached `.so` files. The ASM `f4gemm_bf16_per1x32Fp4_BpreShuffle_*` kernels are pre-compiled blobs shipped with AITER and incur no JIT cost.
+
+---
+
+## Troubleshooting
+
+### `not found tuned config in {file}, will use default config!`
+
+The (M, N, K) shape is missing from your CSV. AITER will fall back to its compiled-in default, which is typically slow. Re-tune for that shape per the internal `tuned_gemm_configs/README.md` runbook (capture the shape via the same log line, append to the untuned CSV, re-run the AITER tuner, commit the new CSV).
+
+### Slow first iteration (~minutes), normal afterwards
+
+Expected—the first call to a CK-based `a4w4_blockscale_*` kernel triggers JIT compilation. Cached `.so` files are reused on subsequent starts.
+
+### `User specified backend AITER cannot handle the given inputs`
+
+Raised by `AutoKernelDispatcher.dispatch` when `GEMMFP4AITERBackend.can_handle` rejects the input. Common causes:
+
+- `M` not a multiple of 16, or `N` not a multiple of 16 (constants `AITER_FP4GEMM_M_MULTIPLE` / `AITER_FP4GEMM_N_MULTIPLE` in `gemm_fp4_impl.py`).
+- Unsupported dtype combination (only `(float4_e2m1fn_x2, float4_e2m1fn_x2, fp16/bf16)` is supported).
+- Non-NT layout (`trans_a=False, trans_b=True, trans_c=False`).
+
+Workaround: switch to `PRIMUS_TURBO_GEMM_BACKEND=FP4:HIPBLASLT` for unsupported shapes, or pad/reshape inputs.
+
+### `MXFP4ColumnParallelLinear requires tensor_model_parallel_size=1`
+
+The MXFP4 linear-layer modules assert on `tensor_model_parallel_size == 1`, `gradient_accumulation_fusion == False`, and `sequence_parallel == False` ([`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). Adjust the config accordingly.
+
+### NaN losses
+
+Switch to the hybrid backward mode, which keeps the FP4 forward but does the gradient GEMM in FP8 (E5M2 tensorwise on HipBLASLt):
+
+```yaml
+mxfp4_backward_precision: "fp8"
+```
+
+If NaNs persist, also try `mxfp4_gradient_stochastic_rounding: true`.
+
+---
+
+## Verification status
+
+The public config has been smoke-tested end-to-end: 1000 iters on 8x MI355X (single node, micro-batch 64 / global 512, sequence length 512) completes in ~16-20 minutes with `PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER` and the tuned CSV. No errors across ranks; `pretrain() completed successfully`.
+
+Formal A/B benchmarks vs BF16 and FP8 (delayed and tensorwise) are pending and will be published once a representative suite is run; do not rely on the wall-clock numbers above as performance characterizations.
+
+---
+
+## Source code pointers
+
+- MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`).
+- MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py).
+- Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml).
+- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py).
+- FP4 backend selection (Primus-Turbo): `primus_turbo/common/constants.py`, `primus_turbo/pytorch/core/backend.py`, `primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`.
+- AITER tuned-config loader: `aiter/jit/core.py` (`AITER_CONFIG_GEMM_A4W4`).
+- AITER A4W4 dispatch + hit/miss logging: `aiter/ops/gemm_op_a4w4.py`.
+
+## Related documentation
+
+- [FP8 Training Guide](fp8_training.md)—companion guide for FP8.
+- [Diffusion Architecture / Developer Guide](README.md).
+- [Diffusion Examples README](../../../examples/megatron/diffusion/README.md).
diff --git a/docs/04-technical-guides/fault-tolerance-and-elastic-training.md b/docs/04-technical-guides/fault-tolerance-and-elastic-training.md
new file mode 100644
index 000000000..1c107f072
--- /dev/null
+++ b/docs/04-technical-guides/fault-tolerance-and-elastic-training.md
@@ -0,0 +1,123 @@
+# Fault tolerance and elastic training
+
+Large-scale jobs run for days across thousands of GPUs, where hardware faults, NIC flaps, and node loss are routine. This guide covers the mechanisms Primus exposes to survive and recover from failures: graceful exit + checkpoint-based resume, Megatron's fault-tolerance package and in-process restart, and TorchTitan's [torchft](https://github.com/pytorch/torchft)-based elastic training. Parameters are grounded in `primus/configs/modules/megatron/trainer_base.yaml`, `primus_megatron_module.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`.
+
+The foundation of all recovery is checkpointing—read [Checkpoint management](./checkpoint-management.md) first.
+
+---
+
+## 1. The recovery model
+
+There are three layers, from simplest to most advanced:
+
+1. **Checkpoint + restart**—periodically save state; on failure, relaunch the job and resume from the last checkpoint. Works on every backend; relies on the scheduler (Slurm `--requeue`, Kubernetes restart policy) to relaunch.
+2. **Graceful exit**—detect a signal or time/iteration budget, save a final checkpoint, and exit cleanly so the restart resumes with no lost work.
+3. **In-job fault tolerance / elastic**—detect a failed rank and restart in-process (Megatron) or continue with a reduced/replaced replica group (TorchTitan + torchft) without tearing down the whole job.
+
+---
+
+## 2. Graceful exit and auto-resume (Megatron)
+
+Controls in `trainer_base.yaml` let a run stop cleanly at a boundary so the next launch resumes seamlessly:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `exit_signal_handler` | `false` | Install a signal handler that saves a checkpoint and exits gracefully on SIGTERM (e.g. Slurm preemption). |
+| `exit_duration_in_mins` | `null` | Exit (after saving) once the job has run this many minutes—useful to fit scheduler time limits. |
+| `exit_interval` | `null` | Exit after this many iterations. |
+| `adlr_autoresume` | `false` | Enable ADLR auto-resume integration. |
+| `adlr_autoresume_interval` | `1000` | Iterations between auto-resume checks. |
+
+Primus-level continuation (`primus_megatron_module.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `auto_continue_train` | `false` | Automatically continue from the latest checkpoint in the experiment's save directory on relaunch. |
+| `disable_last_saving` | `false` | Disable the final end-of-run checkpoint save (leave `false` so resume points exist). |
+
+**Pattern:** enable `exit_signal_handler` + `exit_duration_in_mins` (or rely on preemption signals), set a reasonable checkpoint `save_interval`, and turn on `auto_continue_train` so requeued jobs pick up where they left off.
+
+---
+
+## 3. Megatron fault-tolerance package and in-process restart
+
+Megatron integrates an optional fault-tolerance package and in-process restart (`trainer_base.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `enable_ft_package` | `false` | Enable the Megatron fault-tolerance package (rank monitoring / heartbeat). |
+| `calc_ft_timeouts` | `false` | Auto-calculate fault-tolerance timeouts from observed step times. |
+| `run_workload_inspector_server` | `false` | Run the workload inspector server for health/diagnostics. |
+| `inprocess_restart` | `false` | Restart failed ranks **in process** to avoid a full job teardown. |
+
+In-process restart reduces recovery time by re-initializing the process group and reloading state without re-scheduling the whole allocation. Combine with frequent checkpoints so the restarted ranks have a recent resume point.
+
+---
+
+## 4. Numerical safety nets (Megatron)
+
+Detecting corruption early prevents wasted compute and divergence (`trainer_base.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `check_for_nan_in_loss_and_grad` | `true` | Abort/handle on NaN/Inf in loss or gradients. |
+| `check_for_spiky_loss` | `false` | Detect anomalous loss spikes. |
+| `check_for_large_grads` | `false` | Detect abnormally large gradients. |
+| `decrease_batch_size_if_needed` | `false` | Reduce batch size when needed instead of failing. |
+
+These don't recover from hardware faults but stop a corrupted run before it pollutes downstream checkpoints.
+
+---
+
+## 5. Elastic training with torchft (TorchTitan)
+
+TorchTitan supports semi-synchronous, replica-based fault tolerance via [torchft](https://github.com/pytorch/torchft). Configured under `fault_tolerance:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `enable` | `false` | Enable torchft fault tolerance. |
+| `process_group` | `gloo` | Process group backend for fault-tolerance coordination. |
+| `process_group_timeout_ms` | `10000` | Coordination timeout (ms). |
+| `replica_id` | `0` | This replica's ID. |
+| `group_size` | `0` | Number of replica groups (`0` = auto). |
+| `min_replica_size` | `1` | Minimum replicas required to keep training. |
+| `semi_sync_method` | `null` | Semi-synchronous algorithm (e.g. DiLoCo-style), `null` = standard. |
+
+With replica groups, the loss of one replica can be tolerated as long as `min_replica_size` is still satisfied—training continues while the failed replica recovers/rejoins, rather than crashing the whole job.
+
+**Install** the optional dependencies before enabling (`requirements-torchft.txt`):
+
+```bash
+pip install -r requirements-torchft.txt   # torchft-nightly + OpenTelemetry exporters
+```
+
+TorchTitan also exposes communication timeouts under `comm:` (`init_timeout_seconds: 300`, `train_timeout_seconds: 100`) that govern how long collectives wait before declaring a fault.
+
+---
+
+## 6. Scheduler integration
+
+In-job mechanisms still need the scheduler to relaunch on full-job failure:
+
+- **Slurm**—submit with `--requeue` so preempted/failed jobs are re-queued; pair with `exit_signal_handler` to checkpoint on SIGTERM. See [Deployment](../05-operations/deployment.md).
+- **Kubernetes**—use a restart policy / operator that recreates pods; mount checkpoint storage on a shared/persistent volume.
+- **Shared checkpoint storage**—all ranks must read the same checkpoint directory after relaunch (NFS, Lustre, or object storage). See [Checkpoint management](./checkpoint-management.md).
+
+---
+
+## 7. Recommended setup
+
+1. **Always checkpoint**—set a `save_interval` matched to your mean-time-between-failures; use async/distributed checkpointing to keep overhead low.
+2. **Exit cleanly**—`exit_signal_handler: true` (+ `exit_duration_in_mins` for time-boxed allocations).
+3. **Resume automatically**—`auto_continue_train: true` (Megatron) and `--requeue` (Slurm).
+4. **Reduce recovery time at scale**—`enable_ft_package` + `inprocess_restart` (Megatron) or torchft replica groups (TorchTitan).
+5. **Guard numerics**—keep `check_for_nan_in_loss_and_grad` on; consider spiky/large-grad checks for unstable configs.
+
+---
+
+## Related documentation
+
+- [Checkpoint management](./checkpoint-management.md)—save/load formats, async and distributed checkpointing.
+- [Deployment](../05-operations/deployment.md)—Slurm/Kubernetes restart and requeue.
+- [Multi-node networking](./multi-node-networking.md)—NIC faults and collective timeouts.
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md).
diff --git a/docs/04-technical-guides/hybrid-models/README.md b/docs/04-technical-guides/hybrid-models/README.md
new file mode 100644
index 000000000..0e1afae98
--- /dev/null
+++ b/docs/04-technical-guides/hybrid-models/README.md
@@ -0,0 +1,618 @@
+# Zebra-Llama: Hybrid Recurrent-Attention Models on AMD GPUs
+
+Zebra-Llama is a family of hybrid models that combine **recurrent layers** (Mamba SSM, KDA, or GDN) with **Multi-Latent Attention (MLA)** and **SwiGLU MLP** layers. These models are designed to achieve competitive quality with sub-quadratic inference cost.
+
+This guide covers the complete workflow: environment setup, data preparation, pretraining, checkpoint conversion, and evaluation.
+
+> **FLA-validated recipes**: For runnable, FLA-parity-validated walkthroughs of the pure-recurrent
+> variants, see [Pure GDN guide](gdn-guide.md) and [Pure KDA guide](kda-guide.md). For the exhaustive
+> list of code/config/runtime changes required for exact parity with the
+> [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention) reference
+> implementation, see [GDN ⇄ FLA parity](gdn-fla-parity.md) and [KDA ⇄ FLA parity](kda-fla-parity.md).
+
+---
+
+## Table of Contents
+
+- [Related Guides](#related-guides)
+
+- [Architecture Overview](#architecture-overview)
+- [Available Configurations](#available-configurations)
+- [Prerequisites](#prerequisites)
+- [Step 1: Environment Setup](#step-1-environment-setup)
+- [Step 2: Dataset Preparation](#step-2-dataset-preparation)
+- [Step 3: Pretraining](#step-3-pretraining)
+  - [Single-Node (Local / Docker)](#single-node-local--docker)
+  - [Multi-Node (Slurm)](#multi-node-slurm)
+  - [Mock Data (Smoke Test)](#mock-data-smoke-test)
+- [Step 4: Checkpoint Conversion to HuggingFace](#step-4-checkpoint-conversion-to-huggingface)
+- [Step 5: Evaluation with lm-eval-harness](#step-5-evaluation-with-lm-eval-harness)
+- [Configuration Reference](#configuration-reference)
+- [Troubleshooting](#troubleshooting)
+
+---
+
+## Related Guides
+
+| Guide | Description |
+|-------|-------------|
+| [Pure GDN guide](gdn-guide.md) | End-to-end 300M pure Gated DeltaNet (GDN) recipe, FLA-validated on 8× MI300X |
+| [Pure KDA guide](kda-guide.md) | End-to-end 300M pure Kimi Delta Attention (KDA) recipe, FLA-validated on 8× MI300X |
+| [GDN ⇄ FLA parity](gdn-fla-parity.md) | Every Primus/Megatron-LM change required for GDN to match the FLA reference implementation |
+| [KDA ⇄ FLA parity](kda-fla-parity.md) | Every Primus/Megatron-LM change required for KDA to match the FLA reference implementation |
+
+---
+
+## Architecture Overview
+
+Zebra-Llama interleaves three types of layers in a repeating pattern:
+
+```
+[Attention] [MLP] [Recurrent] [MLP] [Recurrent] [MLP] ... [Attention] [MLP] ...
+```
+
+- **Recurrent layers** — one of Mamba SSM, Kimi Delta Attention (KDA), or Gated Delta Net (GDN)
+- **Attention layers** — Multi-Latent Attention with YaRN rotary embeddings and LoRA-compressed KV
+- **MLP layers** — SwiGLU feed-forward with Transformer Engine fused norms
+
+The `hybrid_attention_ratio` parameter controls what fraction of recurrent+attention layer pairs use attention (default 0.25 = 1 attention layer per 3 recurrent layers). Setting it to `0.0` yields a pure recurrent model (all KDA or GDN), while `1.0` yields a pure MLA attention model.
+
+---
+
+## Available Configurations
+
+### Pretrain Configs (`examples/megatron/configs/MI300X/`)
+
+| Config | Model | Recurrent Type | Seq Length | Params | Tokenizer |
+|--------|-------|---------------|------------|--------|-----------|
+| `zebra_llama_1B-pretrain.yaml` | 1B (Mamba+MLA) | Mamba SSM | 2048 | ~1B | `meta-llama/Llama-3.2-1B` |
+| `zebra_llama_1B_kda-pretrain.yaml` | 1B (KDA+MLA) | Kimi Delta Attention | 8192 | ~1B | `meta-llama/Llama-3.2-1B` |
+| `zebra_llama_1B_kda_pure-pretrain.yaml` | 1B (pure KDA) | Kimi Delta Attention | 2048 | ~1.2B | `meta-llama/Llama-3.2-1B` |
+| `zebra_llama_1B_gdn-pretrain.yaml` | 1B (GDN only) | Gated Delta Net | 8192 | ~1B | `fla-hub/gla-1.3B-100B` |
+| `zebra_llama_1B_gdn_pure-pretrain.yaml` | 1B (pure GDN) | Gated Delta Net | 2048 | ~1.2B | `meta-llama/Llama-3.2-1B` |
+| `zebra_llama_3B-pretrain.yaml` | 3B (Mamba+MLA) | Mamba SSM | 8192 | ~3B | `meta-llama/Llama-3.2-3B` |
+| `zebra_llama_8B-pretrain.yaml` | 8B (Mamba+MLA) | Mamba SSM | 8192 | ~8B | `meta-llama/Llama-3.1-8B` |
+
+### Model Configs (`primus/configs/models/megatron/`)
+
+| Config | Layers | Hidden | FFN | Attention Ratio | Attention Type |
+|--------|--------|--------|-----|----------------|----------------|
+| `zebra_llama_1B.yaml` | 32 | 2048 | 8192 | 0.25 | MLA |
+| `zebra_llama_1B_kda_pure.yaml` | 32 | 2048 | 8192 | 0.0 (pure KDA) | None |
+| `zebra_llama_1B_gdn.yaml` | 32 | 2048 | 8192 | 0.0 (pure GDN) | None |
+| `zebra_llama_1B_gdn_pure.yaml` | 32 (16 GDN+16 MLP) | 2048 | 8192 | 0.0 (pure GDN) | None |
+| `zebra_llama_3B.yaml` | 56 | 3072 | 8192 | 0.25 | MLA |
+| `zebra_llama_8B.yaml` | 64 | 4096 | 14436 | 0.25 | MLA |
+
+> **Note on pure KDA**: The `zebra_llama_1B_kda_pure` config matches FLA's `kda_1B_pure.json`
+> architecture (16 KDA layers, `head_dim=32` for keys, `head_dim=64` for values, tied
+> embeddings, `norm_eps=1e-6`). It uses the FLA Triton kernel (`use_fla_triton_kda: true`)
+> for fused forward+backward during training.
+
+> **Note on pure GDN**: The `zebra_llama_1B_gdn_pure` config matches FLA's
+> `gated_deltanet_1B_pure.json` architecture (16 GDN + 16 MLP layers, `num_heads=8`,
+> `num_v_heads=16`, short convolution with kernel size 4, tied embeddings). This config
+> has been validated end-to-end against FLA on MI300X — the training loss curves match
+> within ~1% across 76K steps on FineWeb-Edu 10BT. See
+> [Step 4](#step-4-checkpoint-conversion-to-huggingface) for conversion to FLA's HuggingFace
+> format.
+
+---
+
+## Prerequisites
+
+- **Hardware**: AMD Instinct MI300X (or compatible ROCm GPUs)
+- **Software**: ROCm drivers >= 7.0, Docker >= 24.0
+- **HuggingFace Token**: Required for gated tokenizers (`HF_TOKEN`)
+- **Disk Space**: ~50 GB for FineWeb-Edu 10BT tokenized data
+
+---
+
+## Step 1: Environment Setup
+
+### 1.1 Pull the Docker Image
+
+```bash
+docker pull docker.io/rocm/primus:v25.10
+```
+
+### 1.2 Clone the Repository
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git
+cd Primus
+```
+
+### 1.3 Start a Development Container
+
+```bash
+# Quick start (mounts Primus into /workspace/Primus)
+bash tools/docker/start_container.sh
+```
+
+This creates a persistent container named `dev_primus_<user>`. You can customize it with environment variables:
+
+```bash
+DOCKER_IMAGE=docker.io/rocm/primus:v25.10 \
+DATA_PATH=/path/to/data \
+bash tools/docker/start_container.sh
+```
+
+Then exec into the container:
+
+```bash
+docker exec -it dev_primus_$(whoami) bash
+cd /workspace/Primus
+```
+
+### 1.4 Install Python Dependencies (inside container)
+
+```bash
+pip install -r requirements.txt
+```
+
+For GDN models (required for the Triton kernel and FLA model classes):
+
+```bash
+pip install flash-linear-attention
+```
+
+For evaluation, also install:
+
+```bash
+pip install lm-eval
+```
+
+---
+
+## Step 2: Dataset Preparation
+
+Zebra-Llama uses the [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) dataset, preprocessed into Megatron binary format.
+
+### 2.1 Set Up Environment
+
+```bash
+export HF_TOKEN="hf_your_token_here"
+export PYTHONPATH="$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"
+```
+
+### 2.2 Run Data Preparation
+
+```bash
+python examples/megatron/prepare_fineweb_edu.py \
+    --primus-path . \
+    --data-path ./data \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model meta-llama/Llama-3.2-1B \
+    --sample-size 10BT
+```
+
+This will:
+1. Download the FineWeb-Edu 10BT dataset from HuggingFace
+2. Tokenize it into Megatron binary format (`.bin` + `.idx` files)
+3. Output files to `./data/fineweb-edu-10BT/HuggingFaceTokenizer/`
+
+Available sample sizes: `10BT`, `100BT`, `350BT`
+
+The script uses all available CPU cores by default. To limit parallelism, add `--workers N`.
+
+> **Note**: The context length (sequence length) is not set during data prep. It is configured at training time via `seq_length` in your pretrain YAML.
+
+### 2.3 Using a Different Tokenizer
+
+For the GDN config which uses `fla-hub/gla-1.3B-100B`:
+
+```bash
+python examples/megatron/prepare_fineweb_edu.py \
+    --primus-path . \
+    --data-path ./data \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model fla-hub/gla-1.3B-100B \
+    --sample-size 10BT
+```
+
+### 2.4 FLA-Aligned Data for Pure GDN (Recommended)
+
+When training a pure GDN model for comparison against FLA, it is critical that both frameworks see **identical tokens in the same order**. The standard Megatron data pipeline produces different token ordering than FLA's `preprocess.py` (which shuffles with `seed=42` and concatenates into fixed-length chunks without EOD tokens). This difference alone can cause persistent training loss divergence.
+
+To ensure exact alignment:
+
+**Step 1: Preprocess with FLA** (if not already done):
+
+```bash
+cd /path/to/flash-linear-attention/legacy/training
+python preprocess.py \
+    --dataset HuggingFaceFW/fineweb-edu \
+    --name sample-10BT \
+    --tokenizer meta-llama/Llama-3.2-1B \
+    --seq_len 2048 --num_proc 64
+```
+
+This produces an Arrow dataset under `data/HuggingFaceFW/fineweb-edu/sample-10BT/train/`.
+
+**Step 2: Convert FLA's Arrow data to Megatron binary format**:
+
+```bash
+python convert_fla_to_megatron.py
+```
+
+> **Note**: Edit the `FLA_DATA` and `OUT_PREFIX` paths at the top of `convert_fla_to_megatron.py` to match your environment before running.
+
+This reads the Arrow shard files directly with PyArrow and produces Megatron-compatible `.bin` + `.idx` files. Each FLA 2048-token sequence becomes one Megatron "document". The script verifies token-level consistency after writing.
+
+**Step 3: Point the pretrain config at the converted data**:
+
+```yaml
+train_data_path: >
+  /path/to/data/fla_aligned/fla_fineweb_edu_10BT_text_sentence
+mock_data: false
+```
+
+### 2.5 Update Data Paths in Config
+
+After preparation (standard or FLA-aligned), update the `train_data_path` in your pretrain config YAML to point to the generated files:
+
+```yaml
+# Standard Megatron data prep (multiple shards)
+train_data_path: >
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_0_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_1_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_2_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_3_text_sentence
+mock_data: false
+```
+
+---
+
+## Step 3: Pretraining
+
+### Single-Node (Local / Docker)
+
+Launch training inside a Docker container on a single node:
+
+```bash
+# Zebra-Llama 1B with KDA (Kimi Delta Attention)
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml \
+DATA_PATH=./data \
+GPUS_PER_NODE=8 \
+HF_TOKEN=$HF_TOKEN \
+bash examples/run_local_pretrain.sh
+```
+
+Other model variants:
+
+```bash
+# Zebra-Llama 1B with Mamba SSM
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml \
+bash examples/run_local_pretrain.sh
+
+# Zebra-Llama 1B with pure KDA (no attention layers)
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_kda_pure-pretrain.yaml \
+bash examples/run_local_pretrain.sh
+
+# Zebra-Llama 1B with GDN (pure recurrent, no attention)
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_gdn-pretrain.yaml \
+bash examples/run_local_pretrain.sh
+
+# Zebra-Llama 1B pure GDN (FLA-validated, 4-GPU)
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure-pretrain.yaml \
+GPUS_PER_NODE=4 \
+bash examples/run_local_pretrain.sh
+
+# Zebra-Llama 3B
+EXP=examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml \
+bash examples/run_local_pretrain.sh
+
+# Zebra-Llama 8B
+EXP=examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml \
+bash examples/run_local_pretrain.sh
+```
+
+### Multi-Node (Slurm)
+
+For multi-node training on a Slurm cluster:
+
+```bash
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml \
+DATA_PATH=/shared/data \
+NNODES=2 \
+bash examples/run_slurm_pretrain.sh
+```
+
+Ensure the `global_batch_size` in your config is divisible by `micro_batch_size * GPUS_PER_NODE * NNODES`.
+
+### If Already Inside a Container
+
+If you are already inside a Docker container or on a bare-metal node with the environment set up:
+
+```bash
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml \
+bash examples/run_pretrain.sh
+```
+
+### Mock Data (Smoke Test)
+
+To quickly verify the model runs without real data, the 3B and 8B configs come with `mock_data: true` by default. For the 1B configs, you can override:
+
+```bash
+EXP=examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml \
+bash examples/run_local_pretrain.sh \
+  --mock_data true --train_iters 10
+```
+
+### Key Training Parameters
+
+| Parameter | Description | Typical Values |
+|-----------|-------------|---------------|
+| `train_iters` | Total training iterations | 38147 (1B KDA pure), 400000 (1B Mamba) |
+| `micro_batch_size` | Per-GPU batch size | 4 (KDA/GDN), 16 (Mamba 1B) |
+| `global_batch_size` | Total batch size across all GPUs | micro_batch_size * num_gpus |
+| `seq_length` | Sequence length | 2048, 4096, 8192 |
+| `lr` | Peak learning rate | 2.0e-4 |
+| `save_interval` | Checkpoint save frequency | 1000 |
+| `auto_continue_train` | Auto-resume from last checkpoint on crash | `true` / `false` |
+| `hybrid_attention_ratio` | Fraction of attention layers (0.0 = pure recurrent) | 0.0, 0.25 |
+
+---
+
+## Step 4: Checkpoint Conversion to HuggingFace
+
+Convert a Megatron checkpoint to HuggingFace format for inference and evaluation.
+
+### 4.1 Convert Checkpoint
+
+#### Pure GDN Models
+
+Pure GDN models use a dedicated converter that maps Primus's fused projections to FLA's native `GatedDeltaNetForCausalLM` format:
+
+```bash
+python tools/hybrid/convert_gdn_to_fla_hf.py \
+    --checkpoint-path output/amd/root/zebra_llama_1B_gdn_pure-pretrain/checkpoints/iter_0076294 \
+    --output-dir output/gdn_pure_1B_fla_hf \
+    --config /path/to/gated_deltanet_1B_pure.json
+```
+
+This handles:
+- Splitting the fused `in_proj` (3104 → q/k/v/gate/beta/alpha projections)
+- Splitting the fused `conv1d` (q/k/v convolutions)
+- Splitting the fused SwiGLU `fc1` (gate_proj + up_proj)
+- Mapping alternating GDN/MLP sublayers to combined FLA layers
+- Handling tied embeddings
+
+After conversion, verify with the sanity check:
+
+```bash
+python tools/hybrid/verify_gdn_conversion.py --model-path output/gdn_pure_1B_fla_hf
+```
+
+Expected output: Loss ~2-4, top prediction for "The capital of France is" should be "Paris".
+
+#### KDA / Hybrid Models
+
+The general converter auto-detects architecture from the checkpoint's saved arguments:
+
+```bash
+# KDA+MLA hybrid model
+python tools/hybrid/convert_zebra_llama_to_hf.py \
+    --checkpoint-path output/zebra_llama_1B_kda-pretrain/iter_0028000 \
+    --output-dir output/zebra_llama_1B_kda_hf_iter_0028000
+
+# Pure KDA model
+python tools/hybrid/convert_zebra_llama_to_hf.py \
+    --checkpoint-path output/zebra_llama_1B_kda_pure-pretrain/iter_0038000 \
+    --output-dir output/zebra_llama_1B_kda_pure_hf
+```
+
+The converter will:
+- Read the Megatron checkpoint and training arguments
+- Auto-detect architecture parameters (`hybrid_attention_ratio`, `kda_num_heads`, `q_lora_rank`, etc.)
+- Remap parameter names from Megatron conventions to HuggingFace conventions
+- Save `pytorch_model.bin`, `config.json`, and a model card `README.md` in the output directory
+- Copy `modeling_zebra_llama.py` into the output directory for `trust_remote_code` loading
+
+### 4.2 Verify Conversion
+
+The script prints a summary of missing, extra, and shape-mismatched keys. A successful conversion shows:
+
+```
+0 missing, 0 extra, 0 shape mismatches
+```
+
+### 4.3 Supported Architectures
+
+| Architecture | `hybrid_attention_ratio` | Layer pattern |
+|---|---|---|
+| Pure KDA | `0.0` | All KDA + MLP |
+| KDA + MLA hybrid | `0.0 < r < 1.0` | Mix of KDA and MLA + MLP |
+| Pure MLA | `1.0` | All MLA + MLP |
+| Pure GDN | `0.0` (with GDN spec) | All GDN + MLP |
+| Mamba + MLA hybrid | `0.0 < r < 1.0` (with Mamba spec) | Mix of Mamba and MLA + MLP |
+
+---
+
+## Step 5: Evaluation with lm-eval-harness
+
+### 5.1 Pure GDN Models (FLA format)
+
+Pure GDN models use a dedicated eval wrapper (`tools/hybrid/eval_gdn_lm_eval.py`) that pre-registers FLA's `GatedDeltaNetForCausalLM` with transformers' `AutoModel` and patches compatibility issues with transformers >= 4.55:
+
+```bash
+python tools/hybrid/eval_gdn_lm_eval.py \
+    --model hf \
+    --model_args pretrained=output/gdn_pure_1B_fla_hf,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+    --batch_size auto \
+    --output_path eval_results/gdn_pure_1B
+```
+
+> **Note**: Do not use `lm_eval --model hf` directly — it will fail because `AutoConfig` does not recognize `gated_deltanet` without FLA being imported first. The wrapper handles this. The `tokenizer=meta-llama/Llama-3.2-1B` argument is required since the converted model directory does not contain tokenizer files.
+
+### 5.2 KDA / Hybrid Models (Zebra-Llama format)
+
+KDA and hybrid models use the custom `ZebraLlamaForCausalLM` architecture, which requires a dedicated lm-eval wrapper:
+
+```bash
+python3 tools/hybrid/lm_harness_eval.py --model zebra_llama \
+    --model_args pretrained=output/zebra_llama_1B_kda_pure_hf,dtype=bfloat16 \
+    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+    --batch_size auto
+```
+
+### 5.3 Using the Eval Shell Script (KDA/Hybrid)
+
+```bash
+bash tools/hybrid/eval_zebra_llama_lm_eval.sh \
+    --checkpoint output/zebra_llama_1B_kda_pure_hf \
+    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+    --batch-size auto \
+    --dtype bfloat16 \
+    --output eval_results/zebra_llama_1B_kda_pure
+```
+
+> **Important**: The eval script internally invokes `python3 tools/hybrid/lm_harness_eval.py --model zebra_llama` (not `lm_eval --model hf`). This ensures the custom model architecture is properly registered.
+
+### 5.4 Available Benchmarks
+
+| Task | Description | Metric |
+|------|-------------|--------|
+| `arc_easy` | ARC Easy (science QA) | acc, acc_norm |
+| `arc_challenge` | ARC Challenge (harder science QA) | acc, acc_norm |
+| `hellaswag` | HellaSwag (commonsense NLI) | acc, acc_norm |
+| `mmlu` | MMLU (57 subject knowledge benchmark) | acc |
+| `openbookqa` | OpenBookQA | acc, acc_norm |
+| `piqa` | PIQA (physical intuition QA) | acc, acc_norm |
+| `race` | RACE (reading comprehension) | acc |
+| `winogrande` | Winogrande (coreference resolution) | acc |
+
+### 5.5 Memory Considerations
+
+The pure-PyTorch KDA chunked attention is memory-intensive. If you encounter OOM errors:
+
+- Use `--batch_size auto` to let lm-eval find the largest fitting batch size
+- Reduce `max_length` (e.g., `max_length=1024` in `--model_args`)
+- Reduce `--batch_size` to 1
+
+---
+
+## Configuration Reference
+
+### Hybrid Layer Specs
+
+The `spec` field in the pretrain config selects the layer arrangement:
+
+| Spec | Description |
+|------|-------------|
+| `hybrid_stack_spec` | Mamba SSM + MLA hybrid |
+| `kda_hybrid_stack_spec` | KDA + MLA hybrid (or pure KDA with `hybrid_attention_ratio: 0.0`) |
+| `gdn_hybrid_stack_spec` | GDN + MLA hybrid (or pure GDN with `hybrid_attention_ratio: 0.0`) |
+
+### Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `DOCKER_IMAGE` | Docker image for training | `docker.io/rocm/primus:v25.10` |
+| `EXP` | Path to experiment config YAML | `examples/megatron/exp_pretrain.yaml` |
+| `DATA_PATH` | Path to dataset directory | `./data` |
+| `HF_TOKEN` | HuggingFace API token | (required for gated models) |
+| `WANDB_API_KEY` | Weights & Biases API key | (optional) |
+| `GPUS_PER_NODE` | Number of GPUs per node | 8 |
+| `NNODES` | Number of nodes | 1 |
+| `MASTER_ADDR` | Master node address | `localhost` |
+| `MASTER_PORT` | Master node port | `1234` |
+
+---
+
+## Troubleshooting
+
+### OOM During Training
+
+- Reduce `micro_batch_size` or `seq_length`
+- Enable activation checkpointing: add `recompute_granularity: selective` to the config
+
+### OOM During Evaluation
+
+- Use `--batch_size 1` or `--batch_size auto`
+- Add `max_length=1024` to `--model_args`
+
+### `ModuleNotFoundError: No module named 'megatron'`
+
+Set the Python path before running data preparation:
+
+```bash
+export PYTHONPATH="$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"
+```
+
+### Checkpoint Conversion Shape Mismatches
+
+Ensure the `modeling_zebra_llama.py` model definition matches the architecture of your checkpoint (Mamba vs KDA vs GDN). The converter auto-detects architecture from checkpoint args, but the HF model code in `tools/hybrid/modeling_zebra_llama.py` must support the target architecture. Common causes of shape mismatches:
+
+- Mismatched `hybrid_attention_ratio` between config and checkpoint
+- Incorrect `kda_num_heads` or head dimension settings
+- Using a `modeling_zebra_llama.py` that doesn't support the checkpoint's attention type
+
+### `ValueError: model type 'zebra_llama' not recognized`
+
+This occurs when using `lm_eval --model hf` directly instead of the custom wrapper. Always use:
+
+```bash
+python3 tools/hybrid/lm_harness_eval.py --model zebra_llama ...
+```
+
+Or the eval shell script, which handles this automatically.
+
+### Truncation Warnings During Eval
+
+Messages like `Combined length of context and continuation exceeds model's maximum length` mean some eval samples are being truncated. This has minimal impact on most benchmarks but can affect long-context tasks like RACE. To avoid truncation, increase `max_length` in `--model_args`.
+
+### NCCL / RCCL Timeout During Training
+
+On MI300X, intermittent RCCL hangs can occur (typically during checkpoint saves). Mitigations:
+
+- Set `auto_continue_train: true` in the pretrain config to auto-resume from the last checkpoint
+- Increase the heartbeat timeout: `export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=7200`
+
+---
+
+## File Reference
+
+```
+Primus/
+├── examples/megatron/
+│   ├── configs/MI300X/
+│   │   ├── zebra_llama_1B-pretrain.yaml        # 1B Mamba+MLA
+│   │   ├── zebra_llama_1B_kda-pretrain.yaml     # 1B KDA+MLA hybrid
+│   │   ├── zebra_llama_1B_kda_pure-pretrain.yaml # 1B pure KDA
+│   │   ├── zebra_llama_1B_gdn-pretrain.yaml     # 1B GDN
+│   │   ├── zebra_llama_1B_gdn_pure-pretrain.yaml # 1B pure GDN (FLA-validated)
+│   │   ├── zebra_llama_3B-pretrain.yaml         # 3B Mamba+MLA
+│   │   └── zebra_llama_8B-pretrain.yaml         # 8B Mamba+MLA
+│   ├── prepare_fineweb_edu.py                   # Data preparation script
+│   ├── prepare_fineweb_edu.sh                   # Data prep shell wrapper
+│   └── preprocess_data.py                       # Megatron tokenizer
+├── primus/configs/models/megatron/
+│   ├── zebra_llama_1B.yaml                      # 1B model architecture
+│   ├── zebra_llama_1B_kda_pure.yaml             # 1B pure KDA architecture
+│   ├── zebra_llama_1B_gdn.yaml                  # 1B GDN architecture
+│   ├── zebra_llama_1B_gdn_pure.yaml             # 1B pure GDN (FLA-validated)
+│   ├── zebra_llama_3B.yaml                      # 3B model architecture
+│   └── zebra_llama_8B.yaml                      # 8B model architecture
+├── tools/
+│   ├── hybrid/
+│   │   ├── convert_zebra_llama_to_hf.py         # Megatron → HF converter (KDA/hybrid)
+│   │   ├── convert_gdn_to_fla_hf.py             # Megatron → FLA HF converter (pure GDN)
+│   │   ├── verify_gdn_conversion.py             # Post-conversion sanity check (pure GDN)
+│   │   ├── eval_gdn_lm_eval.py                  # lm-eval wrapper for GDN (registers FLA)
+│   │   ├── convert_zebra_llama_to_hf.sh         # Converter shell wrapper
+│   │   ├── modeling_zebra_llama.py              # HF model definition (KDA/hybrid)
+│   │   ├── lm_harness_eval.py                   # lm-eval wrapper
+│   │   ├── eval_zebra_llama_lm_eval.sh          # Eval shell wrapper
+│   │   ├── run_zebra_eval.sh                    # Quick eval script
+│   │   ├── chat_zebra_llama.py                  # Interactive chat
+│   │   └── convert_fla_to_megatron.py           # FLA Arrow → Megatron binary converter
+│   └── docker/start_container.sh                # Dev container launcher
+├── examples/
+│   ├── run_local_pretrain.sh                    # Single-node Docker launcher
+│   ├── run_slurm_pretrain.sh                    # Slurm launcher
+│   └── run_pretrain.sh                          # Core training entrypoint
+└── requirements.txt                             # Python dependencies
+```
diff --git a/docs/04-technical-guides/hybrid-models/gdn-fla-parity.md b/docs/04-technical-guides/hybrid-models/gdn-fla-parity.md
new file mode 100644
index 000000000..8bb66710d
--- /dev/null
+++ b/docs/04-technical-guides/hybrid-models/gdn-fla-parity.md
@@ -0,0 +1,313 @@
+# GDN ⇄ FLA Parity in Primus
+
+This document captures every change required in Primus and the vendored
+Megatron-LM submodule to make a 300M Gated DeltaNet (GDN) pretraining run
+match the [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention)
+reference implementation on **both** loss trajectory and step throughput
+on 8× MI300X.
+
+## Final result
+
+| Axis | FLA reference | Primus (this branch) | Δ |
+|------|---------------|----------------------|----|
+| Per-iteration time (avg over 4768 iters) | **1434.6 ms** | **1431.6 ms** | **−0.21% (Primus faster)** |
+| Throughput | 182,729 tok/s/GPU | **183,213 tok/s/GPU** | **+0.27%** |
+| TFLOP/s/GPU | (not logged) | 642 | — |
+| Total wall time (4768 iters) | 1h 54m 00s | **1h 53m 42s** | **−18s (Primus faster)** |
+| Loss @ iter 1 | 11.9654 | **11.9652** | **−0.00% (bit-perfect)** |
+| Loss @ iter 1000 | 4.0012 | 4.0497 | +1.21% |
+| Loss @ iter 2000 | 3.6067 | 3.6144 | +0.21% |
+| Loss late-training (iter 3700–4700 avg) | 3.3795 | 3.3829 | +0.10% |
+| First crossover (Primus < FLA) | — | iter 2100 | — |
+
+**Loss curves overlap from iter ~2000 onward**, with batch-to-batch
+oscillation of ±0.25%. The only persistent gap is in the LR-warmup region
+(iter 50–500), and that gap closes monotonically with no instability.
+Both forward and gradient at iter 1 are bit-identical to FLA.
+
+---
+
+## How to run
+
+Inside the `rocm/primus:v26.2` container with the repo mounted at
+`/home/<user>/Primus`:
+
+```bash
+# Launch training (8 GPUs by default). The Megatron-LM behavioral patches
+# below are applied automatically at startup via Primus's patch system --
+# no separate apply step needed.
+EXP=examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml \
+  bash examples/run_pretrain.sh 2>&1 | tee primus_gdn.log
+```
+
+Optional toggles (all default off unless noted).  Each is exposed at
+TWO equivalent surfaces — pick whichever is more convenient:
+
+- **YAML knob** (canonical, declarative — co-located with the rest of the
+  run config; see `primus/configs/models/megatron/mamba_base.yaml` for the
+  full set of `null` defaults, and the GDN/KDA `*-pretrain.yaml`
+  overrides for resolved values).
+- **Environment variable** (ad-hoc, for one-off A/B without editing a
+  YAML).  When both are set, the env var wins (backward compat).
+
+The mapping is plumbed by
+`primus/backends/megatron/patches/fla_runtime_patches.py` at
+`phase="build_args"` which copies any non-`null` YAML field into the
+corresponding env var before any FLA module is imported.
+
+| YAML knob | Env var | Default | Effect |
+|--|--|--|--|
+| `fused_ce_mode` | `PRIMUS_FUSED_CE` | `1` | `1` = FLA `FusedLinearCrossEntropyLoss` (chunked, no full logits tensor); `2` = FLA `FusedCrossEntropyLoss` (matches FLA exactly); `0` = native Megatron CE. |
+| `fused_ce_chunks` | `PRIMUS_FUSED_CE_CHUNKS` | `32` | Number of chunks the FLA CE splits the logits across.  Lower = faster but bigger peak allocation. |
+| `use_fla_fused_swiglu` | `PRIMUS_FLA_SWIGLU` | `1` | Replaces Megatron's naive SwiGLU with FLA's Triton-fused kernel (≈20 ms/step saved). |
+| `use_fla_fused_rmsnorm` | `PRIMUS_FLA_NORM` | `0` | Use FLA's `RMSNorm` in `WrappedTorchNorm`. |
+| `use_fla_fused_gated_norm` | `PRIMUS_FLA_NORM` | `0` | Use FLA's `FusedRMSNormGated` for GDN's gated output norm.  Also enables a fused pre-norm/MLP path inside `HybridStack` (saves one normalization launch per GDN block).  Same env var as `use_fla_fused_rmsnorm` — kept as a separate YAML alias for clarity. |
+| `use_fla_short_conv` | `PRIMUS_FLA_CONV` | `0` | Route the depthwise short conv1d through FLA's Triton `causal_conv1d` instead of Tri-Dao's CUDA package. |
+| `use_fla_data` + `fla_cache_dir` | `PRIMUS_FLA_DATA` + `PRIMUS_FLA_CACHE_DIR` | `0` / `""` | When `use_fla_data=true` and `fla_cache_dir=<HF dataset cache path>`, replace Megatron's `GPTDataset` with the `FLAOrderGPTDataset` shim that emits tokens in the exact same order as FLA's HuggingFace `DistributedSampler`. |
+| `fla_mla_attn` | `PRIMUS_FLA_MLA_ATTN` | unset | MLA `core_attention` calls `flash_attn_func` directly (skips TE's CK fallback). |
+| _(env-only)_ | `PRIMUS_TORCH_OPTIM` | `0` | Use `torch.optim.AdamW(fused=True)` instead of TE/Apex `FusedAdam` (for bit-level reproducibility experiments). |
+
+All env-var paths are inert when the variable is unset (cost: a few
+`os.environ.get()` lookups per iteration — microseconds vs seconds).
+
+---
+
+## What changed and why
+
+The work splits cleanly across four layers: model code, Megatron-LM
+submodule, YAML configs, and runtime knobs.
+
+### A. Primus model code
+
+| File | Change | Reason |
+|------|--------|--------|
+| `primus/backends/megatron/core/models/hybrid/gated_delta_net.py` | Pass `g=alpha`, `use_gate_in_kernel=True`, `A_log=…`, `dt_bias=…` directly to `chunk_gated_delta_rule`; add optional FLA Triton `causal_conv1d` path under `args.use_fla_short_conv`; add optional FLA `FusedRMSNormGated` path under `args.use_fla_fused_gated_norm`; remove `@jit_fuser` on `_apply_gated_norm` so the gated path can branch. | Match FLA's exact kernel call signature (it folds gate+softplus+log into the kernel) and let users opt into FLA's Triton kernels when bit-level parity is required. |
+| `primus/backends/megatron/core/models/hybrid/gated_delta_net_layer.py` | Forward `eps=self.config.layernorm_epsilon` to the pre-norm `build_module(...)` call; defer the `residual.to(fp32)` cast until after the optional pre-norm fusion path; expose `_fuse_prenorm_with_next` flag. | `WrappedTorchNorm`'s default `eps=1e-5` was silently overriding the YAML's `1e-6`, causing a ~1.1% per-layer divergence from FLA. The deferred fp32 cast lets the pre-norm/MLP fusion in `HybridStack` work correctly. |
+| `primus/backends/megatron/core/models/hybrid/hybrid_block.py` | If `config.fp32_residual_connection` is set, force `residual_in_fp32=True`; under `args.use_fla_fused_rmsnorm`, mark every GDN layer with `_fuse_prenorm_with_next=True` and rewrite the forward loop to fuse a GDN block's mixer-out with the next MLP block's pre-MLP layernorm in a single op. | The fp32-residual handling was previously silently dropped. The pre-norm fusion saves one normalization launch per GDN block when FLA-norm is enabled. (For TE-free builds use the `gdn_hybrid_stack_spec_no_te` spec from the YAML instead.) |
+| `primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py` | Add a new `gdn_hybrid_stack_spec_no_te` ModuleSpec that uses `WrappedTorchNorm` and plain `Column/RowParallelLinear` everywhere, with the same submodule wiring as `gdn_hybrid_stack_spec`. | YAML can now select TE-free layers via `spec: [..., gdn_hybrid_stack_spec_no_te]` for FLA loss-curve alignment without touching code. |
+| `primus/backends/megatron/patches/mamba_fla_data_patches.py` | Wraps `pretrain_mamba.train_valid_test_datasets_provider`, branching to `tools.hybrid.fla_order_dataset.FLAOrderGPTDataset` when `args.use_fla_data=True` + `args.fla_cache_dir=<path>`. | Lets us bypass Megatron's `GPTDataset` shuffler and drive Primus with the exact same token order FLA's `DistributedSampler` produces, isolating data-ordering effects from model effects during comparison. |
+
+### B. Megatron-LM behavioral patches (Primus patch system)
+
+Primus never forks the vendored `third_party/Megatron-LM` submodule.
+Instead these six patches are runtime monkey-patches registered with
+`@register_patch` and applied automatically at `phase="before_train"` by
+`primus/core/patches` -- see the
+[Backend Patch Explorer](../../.cursor/skills/backend-patch-explorer/SKILL.md)
+skill for how the engine works in general. Each patch's `condition=` gates
+it on the relevant config flag, so it's a no-op unless that flag is set.
+
+| Patch id | File | Change | Reason |
+|-------|------|--------|--------|
+| `megatron.mamba.fla_fused_ce` | `mamba_fused_ce_patches.py` | Add `_use_fused_cross_entropy` path to `MambaModel`. Mode 1 = `FusedLinearCrossEntropyLoss` (chunked, never materializes the full logits tensor). Mode 2 = `FusedCrossEntropyLoss` (matches FLA exactly, materializes bf16 logits). Selected by `args.fused_ce_mode` via `get_args()`. | Megatron always materializes a `(batch*seq, vocab)` fp32 logits tensor before CE — for 1024 batch × 2048 seq × 32k vocab this is 256 GB at fp32. FLA chunks it. Massive memory + speed win. |
+| `megatron.optimizer.torch_fused_adam` | `torch_fused_adam_patches.py` | Add `PRIMUS_TORCH_OPTIM=1` opt-in path that selects `torch.optim.AdamW(fused=True)` over TE/Apex `FusedAdam`. | TE's FusedAdam has slightly different epsilon-handling internally; toggling this lets us prove that Primus's AdamW is bit-identical to FLA's when both use torch's fused kernel. |
+| `megatron.mlp.fla_swiglu` | `mlp_fla_swiglu_patches.py` | Replace the naive `silu(x_glu) * x_linear` (2 separate kernel launches + intermediate tensor) with FLA's Triton-fused `swiglu(x_glu, x_linear)` (1 fwd + 1 bwd kernel) in `MLP`. Toggle: `args.use_fla_fused_swiglu` (default True) via `get_args()`. | Profiler shows ~3.8× fewer GPU cycles spent on the activation step. Saves ~20 ms/iter at our batch size. |
+| `megatron.torch_norm.fla_rmsnorm` | `torch_norm_fla_rmsnorm_patches.py` | When `args.use_fla_fused_rmsnorm=True`, return `fla.modules.RMSNorm` from `WrappedTorchNorm` instead of `torch.nn.RMSNorm`. Reads from `get_args()`. | FLA's RMSNorm is a fused Triton kernel that matches the reference run's normalization semantics bit-for-bit. |
+| `megatron.transformer.hybrid_output_init` | `gdn_config_patches.py` | For `is_hybrid_model`, set `output_layer_init_method = init_method_normal(self.init_method_std)` (uniform std, no depth scaling) after `TransformerConfig.__post_init__` runs. | Megatron's default `scaled_init_method_normal` divides std by `sqrt(2 * num_layers)` — that's correct for transformers but **wrong** for hybrid GDN models, where FLA uses a uniform `initializer_range`. Without this fix the output layer started ~24× smaller than FLA's, causing the iter-1 loss to be 11.971 instead of 11.965. |
+| `megatron.mamba.fla_order_dataset` | `mamba_fla_data_patches.py` | Add the FLA-order dataset shim (`args.use_fla_data` + `args.fla_cache_dir`) to `pretrain_mamba.train_valid_test_datasets_provider` — `pretrain_mamba.py` provides its own provider used for Mamba/GDN models. Reads from `get_args()`. | Lets Mamba/GDN training consume the exact same token order FLA's `DistributedSampler` produces. |
+
+### C. YAML configuration changes
+
+#### `primus/configs/models/megatron/{mamba_base,zebra_llama_*_gdn*}.yaml`
+
+Renamed `bases:` → `extends:` (4 files). The Primus YAML resolver was
+silently dropping inheritance from `bases:` lists, which meant model
+configs were missing the dropout/normalization defaults from
+`mamba_base.yaml` → `language_model.yaml`. Verified empirically by
+checking that `hidden_dropout` was leaking through as `0.1` despite
+`mamba_base.yaml` setting it to `0.0`.
+
+#### `examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml`
+
+The training-side config picked up these settings during the
+parity work:
+
+```yaml
+# Logging
+num_workers: 8                    # was 2; FLA uses 8 dataloader workers
+log_interval: 100
+check_for_nan_in_loss_and_grad: false
+
+# Per-rank serialization removal — Megatron defaults insert a
+# dist.barrier() before every L1 timer measurement (~5–10/iter).
+barrier_with_L1_time: false
+
+# Match FLA's seed for bit-perfect iter-1 comparison
+seed: 42
+
+# Norm — Megatron's default is 1e-5; FLA uses 1e-6
+layernorm_epsilon: 1.0e-6
+
+# Force dropout to 0 at the YAML level.
+# language_model.yaml sets these to 0.1 and that was leaking through
+# even when mamba_base.yaml inherited from it (`bases:` bug, see above).
+hidden_dropout: 0.0
+attention_dropout: 0.0
+
+# Training schedule matched to FLA (8 GPUs):
+#   FLA: per_device_train_batch_size=128, 8 GPUs → global=1024
+train_iters: 4768
+micro_batch_size: 128
+global_batch_size: 1024
+
+# Use the no-TE spec for layer alignment with FLA's native PyTorch layers
+spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+no_persist_layer_norm: true
+
+# Distributed-optimizer (ZeRO-1) costs allreduce bandwidth and saves
+# only ~3.6 GB/rank for a 300M model — disable to match FLA's plain DDP.
+use_distributed_optimizer: false
+overlap_grad_reduce: true
+overlap_param_gather: false       # requires distributed optimizer
+gradient_accumulation_fusion: false
+ddp_average_in_collective: true   # divide gradients in NCCL collective
+
+# Load FLA-initialized weights to compare apples-to-apples
+finetune: true
+auto_continue_train: false
+no_load_optim: true
+no_load_rng: true
+load: /home/<user>/Primus/output/fla_init_ckpt_300M
+```
+
+---
+
+## Reproducing the loss-curve match plot
+
+The full per-iteration log lives at `primus_gdn.log` once training
+finishes. Compare against FLA's log
+(`/home/<user>/flash-linear-attention/legacy/training/train_gdn_bs32.log`)
+using the parser in `tools/compare_losses.py` (or the inline parser
+documented in this file's history).
+
+Notable comparison points (FLA loss is divided by 8 to undo the
+DeepSpeed sum-across-ranks):
+
+| iter | FLA / 8 | Primus | Δ% | Notes |
+|-----:|--------:|-------:|---:|-------|
+| 1    | 11.9654 | 11.9652 | **−0.00%** | bit-perfect |
+| 100  | 7.471   | 9.601 | +28.5% | warmup gap (peak) |
+| 500  | 4.625   | 4.728 | +2.2% | warmup closing |
+| 1000 | 4.001   | 4.050 | +1.21% | LR-warmup done |
+| 2000 | 3.607   | 3.614 | +0.21% | converged |
+| 2100 | 3.600   | 3.592 | **−0.22%** | first Primus < FLA crossover |
+| 3000 | 3.448   | 3.460 | +0.35% | matched |
+| 4000 | 3.396   | 3.390 | −0.19% | Primus slightly lower |
+| 4500 | 3.373   | 3.373 | −0.01% | identical |
+| 4700 | 3.351   | 3.366 | +0.45% | identical |
+
+The only persistent gap (iter 50–500) is attributable to dataloader
+ordering — Megatron `GPTDataset` uses its own random shuffler while
+FLA uses HuggingFace's `DistributedSampler`. With `use_fla_data: true`
+the gap closes further but Primus has been verified to converge to
+within ±0.5% by iter 1000 even without it.
+
+---
+
+## Files in the repo for this work
+
+```
+primus/backends/megatron/patches/
+  mamba_fused_ce_patches.py         # FLA fused cross-entropy for MambaModel
+  torch_fused_adam_patches.py       # PRIMUS_TORCH_OPTIM opt-in
+  mlp_fla_swiglu_patches.py         # FLA Triton SwiGLU for MLP
+  torch_norm_fla_rmsnorm_patches.py # FLA RMSNorm for WrappedTorchNorm
+  gdn_config_patches.py             # linear-attention config fields + hybrid init
+  fla_runtime_patches.py            # resolves PRIMUS_FLA_* knobs onto args
+  mamba_fla_data_patches.py         # FLA-order dataset shim wiring
+tools/hybrid/fla_order_dataset.py        # FLA-order dataset shim
+tools/profile_training.py         # NSight Compute / rocprof launcher
+tools/run_profiled_training.sh    # one-shot profiling driver
+tools/hybrid/convert_fla_to_megatron.py  # FLA HF checkpoint → Megatron sharded ckpt
+tools/hybrid/convert_gdn_to_fla_hf.py    # Megatron sharded ckpt → FLA HF checkpoint
+tools/hybrid/verify_gdn_conversion.py    # validates round-trip checkpoint conversion
+tools/hybrid/eval_gdn_lm_eval.py         # lm-eval-harness wrapper for GDN models
+```
+
+The `tools/compare_*.py`, `tools/diff_*.py`, `tools/dump_*.py`,
+`tools/forensic_*.py`, `tools/inspect_*.py`, `tools/hybrid/convert_fla_gdn_init_to_megatron.py`,
+`tools/prove_*.py`, `tools/single_*.py` and `tools/check_*.py` scripts
+were used as one-off forensics during the parity hunt and are kept
+untracked under `tools/`. They reference the env-var-gated dump paths
+documented above.
+
+---
+
+## Hybrid (3 MLA + 9 GDN) parity delta
+
+Everything above applies as-is to the 75% Hybrid GDN+MLA configuration.
+On top of the pure-GDN parity stack, the hybrid run needs two more pieces
+to match FLA's `gated_deltanet_300M_hybrid.json` reference:
+
+### Spec-level fix — LoRA RMSNorm in MLA
+
+FLA's MLA wraps every LoRA projection in a `nn.Sequential` chain:
+
+```python
+self.q_proj = nn.Sequential(
+    nn.Linear(hidden_size, q_lora_rank, bias=False),
+    RMSNorm(q_lora_rank, dtype=torch.float32),
+    nn.Linear(q_lora_rank, num_heads * qk_head_dim, bias=False),
+)
+self.kv_proj = nn.Sequential(
+    nn.Linear(hidden_size, kv_lora_rank, bias=False),
+    RMSNorm(kv_lora_rank, dtype=torch.float32),
+    nn.Linear(kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), bias=False),
+)
+```
+
+Megatron's `MLASelfAttention` constructs the equivalent intermediate
+norm from its `q_layernorm` / `kv_layernorm` submodules:
+
+```python
+self.q_layernorm  = submodules.q_layernorm( hidden_size=config.q_lora_rank,  config=config, eps=config.layernorm_epsilon)
+self.kv_layernorm = submodules.kv_layernorm(hidden_size=config.kv_lora_rank, config=config, eps=config.layernorm_epsilon)
+# ... and applied between linear_*_down_proj and linear_*_up_proj.
+```
+
+Earlier hybrid specs declared both as `IdentityOp`, which silently
+skipped FLA's per-LoRA RMSNorm. Iter-1 still matched bit-perfect
+(both models start from the same init and the missing norm only kicks
+in once the LoRA weights drift from their init), but from iter 100
+onward Primus plateaued ~0.12 above FLA's loss curve.
+
+Fix in `primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py`:
+flip `q_layernorm` / `kv_layernorm` to `TENorm` (TE specs) or
+`WrappedTorchNorm` (no-TE specs) in all four MLA-bearing specs.
+Under `use_fla_fused_rmsnorm: true`, `WrappedTorchNorm` resolves to FLA's
+Triton `RMSNorm`, giving bit-exact FLA semantics.
+
+### Launcher-level fix — full FLA fusion stack
+
+The YAML overrides block is now the canonical surface (all consumers
+read `args.*` via `get_args()`):
+
+```yaml
+# YAML overrides (canonical)
+use_fla_fused_swiglu: true
+use_fla_fused_rmsnorm: true
+use_fla_fused_gated_norm: true
+use_fla_short_conv: true
+use_fla_data: true
+fla_cache_dir: /path/to/fla/cache
+fused_ce_mode: 1
+fused_ce_chunks: 32
+fla_mla_attn: "1"
+```
+
+Legacy env vars are still accepted as ad-hoc overrides (env wins over
+YAML) for backward compatibility:
+
+```bash
+export PRIMUS_FLA_MLA_ATTN=1   # MLA → flash_attn_func directly (TE 2.8.1 cap)
+export PRIMUS_FUSED_CE=1       # FLA chunked fused-LCE (mem + speed)
+export PRIMUS_FLA_SWIGLU=1     # Triton SwiGLU (~20 ms/iter)
+export PRIMUS_FLA_NORM=1       # FLA RMSNorm + FusedRMSNormGated + prenorm/MLP fusion
+export PRIMUS_FLA_CONV=1       # FLA Triton causal_conv1d
+export PRIMUS_FLA_DATA=1       # same token order as FLA's DistributedSampler
+```
+
+With these flags on, the same Megatron stack that ran pure-KDA at
+1.46 s/iter runs the hybrid at FLA-parity speed (∼1.47 s/iter) and
+loss curve (Δ ≤ 0.5% from iter 100 onward), no other changes
+required.
diff --git a/docs/04-technical-guides/hybrid-models/gdn-guide.md b/docs/04-technical-guides/hybrid-models/gdn-guide.md
new file mode 100644
index 000000000..663484565
--- /dev/null
+++ b/docs/04-technical-guides/hybrid-models/gdn-guide.md
@@ -0,0 +1,578 @@
+# Pure GDN 300M on Primus — End-to-End Guide (FLA-validated)
+
+This document is a runnable walkthrough for the **300M pure Gated DeltaNet (GDN)** pretraining recipe in Primus, validated on 8× AMD MI300X against the [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention) reference implementation. It covers every step from raw dataset → tokenization → training → checkpoint conversion → lm-eval benchmark.
+
+The same recipe scales up to the 1B pure-GDN config (`zebra_llama_1B_gdn_pure-pretrain.yaml`) — just swap the config file at training time and the FLA config JSON at conversion time.
+
+---
+
+## Final result
+
+After 4768 iterations (≈10B tokens) on FineWeb-Edu sample-10BT:
+
+
+| Axis                                            | FLA reference     | Primus (this branch)  | Δ                           |
+| ----------------------------------------------- | ----------------- | --------------------- | --------------------------- |
+| Per-iteration time (avg over 4768 iters)        | **1434.6 ms**     | **1431.6 ms**         | **−0.21 % (Primus faster)** |
+| Throughput                                      | 182,729 tok/s/GPU | **183,213 tok/s/GPU** | **+0.27 %**                 |
+| TFLOP/s/GPU                                     | —                 | 642                   | —                           |
+| Wall time (4768 iters, 8× MI300X, healthy node) | 1h 54m 00s        | **1h 53m 42s**        | **−18s**                    |
+| Loss @ iter 1                                   | 11.9654           | **11.9652**           | **−0.00 % (bit-perfect)**   |
+| Loss @ iter 4700 (final logged)                 | 3.3511            | **3.3590**            | **+0.24 %**                 |
+| First Primus-below-FLA crossover                | —                 | iter 2100             | —                           |
+
+
+Loss trajectories overlap from iter ~2000 onward; the only persistent gap is in the LR-warmup region (iter 50–500) and closes monotonically. See [gdn-fla-parity.md](gdn-fla-parity.md) for the deep-dive on every patch and env var.
+
+---
+
+## Table of contents
+
+- [Overview](#overview)
+- [Prerequisites](#prerequisites)
+- [Step 1: Environment](#step-1-environment)
+- [Step 2: Dataset preparation](#step-2-dataset-preparation)
+- [Step 3: Apply Megatron-LM patches](#step-3-apply-megatron-lm-patches)
+- [Step 4: (Optional) Initialize from FLA weights](#step-4-optional-initialize-from-fla-weights)
+- [Step 5: Train](#step-5-train)
+- [Step 6: Monitor and compare against FLA](#step-6-monitor-and-compare-against-fla)
+- [Step 7: Convert checkpoint to HuggingFace format](#step-7-convert-checkpoint-to-huggingface-format)
+- [Step 8: Verify conversion](#step-8-verify-conversion)
+- [Step 9: Run lm-eval-harness benchmarks](#step-9-run-lm-eval-harness-benchmarks)
+- [Configs and tools used](#configs-and-tools-used)
+- [Troubleshooting](#troubleshooting)
+
+---
+
+## Overview
+
+The 300M pure-GDN model has:
+
+- 12 Gated DeltaNet blocks + 12 MLP blocks → 24 Megatron "sublayers"
+- `hidden_size = 1024`, `ffn_hidden_size = 4096`
+- `num_heads = 4` (Q/K), `num_v_heads = 8` (V, grouped-value attention)
+- `head_dim = 64`, short-conv kernel size 4
+- Tied embeddings, no positional encoding (delta-rule recurrence), RMSNorm with `eps = 1e-6`
+- Tokenizer: `meta-llama/Llama-3.2-1B` (128k vocab)
+- Total parameters: **0.308B**
+
+Training schedule (matched to FLA's `gated_deltanet_300M_pure.json`):
+
+- 4768 iterations × 1024 global batch × 2048 seq len = **10.0 B tokens**
+- AdamW (β1=0.9, β2=0.95, wd=0.01), peak LR `2e-4`, cosine decay, 200-step warmup
+- bf16 mixed-precision, no dropout, gradient clip 1.0
+
+---
+
+## Prerequisites
+
+- **Hardware**: 8× AMD MI300X (or compatible ROCm GPU) on a single node
+- **Software**: ROCm ≥ 7.0, Docker ≥ 24.0
+- **Container image**: `rocm/primus:v26.2` (or `v25.10` with the same patches)
+- **HF token**: `HF_TOKEN` set for the gated `meta-llama/Llama-3.2-1B` tokenizer
+- **Disk**: ~20 GB for the FLA-aligned tokenized dataset + ~5 GB per saved checkpoint
+- **Optional**: a local clone of [flash-linear-attention](https://github.com/fla-org/flash-linear-attention) checked out at `legacy/training` — only needed if you want to reuse FLA's preprocessed Arrow files (recommended for bit-identical iter-1 comparison)
+
+---
+
+## Step 1: Environment
+
+### 1.1 Start the dev container
+
+```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 -v $(pwd):$(pwd) -w $(pwd) --shm-size 64G --name primus_hybrid_new \
+  rocm/primus:v26.2
+```
+
+This runs the `rocm/primus:v26.2` image with `/dev/dri`, `/dev/kfd`, IB devices, `--privileged`, your `$HOME` mounted in-place, and `--shm-size 64G`. The container is named `primus_hybrid_new`.
+
+To re-attach later:
+
+```bash
+docker exec -it primus_hybrid_new bash
+cd /home/<user>/Primus
+```
+
+### 1.2 Install Python dependencies inside the container
+
+```bash
+pip install -r requirements.txt
+pip install flash-linear-attention   # FLA model classes + Triton kernels
+pip install lm-eval                  # for benchmark evaluation
+```
+
+The `flash-linear-attention` package supplies the FLA `GatedDeltaNetForCausalLM` class (needed for HF conversion + lm-eval) and the Triton kernels that the optional `PRIMUS_FLA_*` toggles route into.
+
+---
+
+## Step 2: Dataset preparation
+
+You have two choices. **For exact loss-curve parity with the FLA reference run, use Option B.** For a quick first run that won't bit-match FLA in the warmup region but will converge to the same final loss, Option A is fine.
+
+### Option A — Standard Megatron data prep (faster setup)
+
+```bash
+export HF_TOKEN="hf_your_token_here"
+export PYTHONPATH="$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"
+
+python examples/megatron/prepare_fineweb_edu.py \
+    --primus-path . \
+    --data-path ./data \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model meta-llama/Llama-3.2-1B \
+    --sample-size 10BT
+```
+
+Output ends up at `./data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_{0..3}_text_sentence.{bin,idx}` (4 shards). Then point the YAML at it:
+
+```yaml
+train_data_path: >
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_0_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_1_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_2_text_sentence
+  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_3_text_sentence
+```
+
+### Option B — FLA-aligned data (recommended for parity)
+
+The Megatron `GPTDataset` shuffler produces a different token order than FLA's `DistributedSampler` (`seed=42`, fixed 2048-token chunks, no EOD tokens). To match exactly, reuse FLA's already-preprocessed Arrow shards and re-encode them into Megatron `.bin`/`.idx` format.
+
+**Step B.1 — Preprocess with FLA's script** (one-time, ~10 min on 64 cores):
+
+```bash
+cd /path/to/flash-linear-attention/legacy/training
+python preprocess.py \
+    --dataset HuggingFaceFW/fineweb-edu \
+    --name sample-10BT \
+    --tokenizer meta-llama/Llama-3.2-1B \
+    --seq_len 2048 --num_proc 64
+```
+
+This writes Arrow shard files to `legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train/data-*.arrow`.
+
+**Step B.2 — Convert the Arrow shards to Megatron binary** using the script at `[tools/hybrid/convert_fla_to_megatron.py](../../tools/hybrid/convert_fla_to_megatron.py)`:
+
+```bash
+cd /home/<user>/Primus
+# Edit FLA_DATA and OUT_PREFIX at the top of the script if your paths differ
+python tools/hybrid/convert_fla_to_megatron.py
+```
+
+The script reads each Arrow shard directly with PyArrow (zero HuggingFace `datasets` overhead), writes a single `.bin` containing flat int32 token IDs, and emits a Megatron `.idx` file where each 2048-token chunk is one document. It cross-checks the first 10 tokens of the output against the first sample of the first Arrow shard before finishing.
+
+Output: `data/fla_aligned/fla_fineweb_edu_10BT_text_sentence.{bin,idx}` (~19 GB binary).
+
+The default 300M YAML already points at this path:
+
+```yaml
+train_data_path: >
+  /home/<user>/Primus/data/fla_aligned/fla_fineweb_edu_10BT_text_sentence
+```
+
+(adjust the user prefix in `examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml` to match your home directory).
+
+---
+
+## Step 3: Megatron-LM patches (automatic — no action needed)
+
+GDN parity training needs six behavioral patches on top of the vendored
+`third_party/Megatron-LM` submodule. Unlike a typical vendored fork, Primus
+never modifies the third-party source: these patches are implemented as
+runtime monkey-patches using Primus's own patch system
+(`primus/core/patches`) and live under
+`[primus/backends/megatron/patches/](../../primus/backends/megatron/patches/)`.
+They register unconditionally but each carries a `condition=` gate on the
+relevant config flag, so they're a no-op unless you actually enable that
+flag — nothing to run by hand.
+
+| Patch id                                   | File                            | Touches                 | Purpose                                                                                                                                                          |
+| ------------------------------------------- | -------------------------------- | ----------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `megatron.mamba.fla_fused_ce`               | `mamba_fused_ce_patches.py`      | `MambaModel`            | Wires `FusedLinearCrossEntropyLoss` / `FusedCrossEntropyLoss` — never materializes the (batch×seq, vocab) logits tensor; gated by `fused_ce_mode` (default 1)   |
+| `megatron.optimizer.torch_fused_adam`       | `torch_fused_adam_patches.py`    | `optimizer/__init__.py` | Adds `PRIMUS_TORCH_OPTIM=1` opt-in for `torch.optim.AdamW(fused=True)` over TE/Apex FusedAdam                                                                    |
+| `megatron.mlp.fla_swiglu`                   | `mlp_fla_swiglu_patches.py`      | `MLP`                   | Routes SwiGLU through FLA's Triton-fused kernel (saves ~20 ms/iter); `use_fla_fused_swiglu` (default true)                                                       |
+| `megatron.torch_norm.fla_rmsnorm`           | `torch_norm_fla_rmsnorm_patches.py` | `WrappedTorchNorm`   | Routes RMSNorm through `fla.modules.RMSNorm` when `use_fla_fused_rmsnorm=true`                                                                                   |
+| `megatron.transformer.hybrid_output_init`   | `gdn_config_patches.py`          | `TransformerConfig`     | For hybrid models, uses uniform `init_method_normal` (no depth-scaled std) — required for bit-perfect iter-1 loss vs FLA                                         |
+| `megatron.mamba.fla_order_dataset`          | `mamba_fla_data_patches.py`      | `pretrain_mamba.py`     | FLA-order dataset shim (`use_fla_data=true` + `fla_cache_dir=<path>`)                                                                                            |
+
+The `fused_ce_mode` / `use_fla_fused_swiglu` / `use_fla_fused_rmsnorm` /
+`use_fla_data` / `fla_cache_dir` knobs above are resolved once (env var >
+YAML field > default) by `fla_runtime_patches.py` before the patches in this
+table run. See `[gdn-fla-parity.md](gdn-fla-parity.md)` for the full
+per-patch deep-dive.
+
+---
+
+## Step 4: (Optional) Initialize from FLA weights
+
+For bit-perfect iter-1 loss alignment, the validated run loads FLA's *initialized but untrained* checkpoint and then trains from there. The YAML's `load:` field points at this directory:
+
+```yaml
+load: /home/<user>/Primus/output/fla_init_ckpt_300M
+finetune: true            # load weights, ignore optimizer state and iteration count
+no_load_optim: true
+no_load_rng: true
+```
+
+The Primus repo includes `tools/hybrid/convert_fla_gdn_init_to_megatron.py` (the GDN counterpart of `tools/hybrid/convert_fla_kda_init_to_megatron.py`) that takes the FLA HuggingFace random-init checkpoint and writes a Megatron-shape `iter_0000000/mp_rank_00/model_optim_rng.pt`. Skip this step if you're happy with Primus's own random init — final loss is identical, only iter-1 drifts by `~5e-3`.
+
+---
+
+## Step 5: Train
+
+### 5.1 Inspect the config
+
+The training config lives at `[examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml](../../examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml)`. Key parameters (matched to FLA):
+
+```yaml
+train_iters: 4768                 # ≈ 10B tokens at global_batch=1024, seq=2048
+micro_batch_size: 128             # per-GPU
+global_batch_size: 1024           # 8 GPUs × 128 = 1024
+seq_length: 2048
+lr: 2.0e-4
+min_lr: 2.0e-5                    # min_lr_rate=0.1 → 2e-5
+lr_warmup_iters: 200
+lr_decay_iters: 4768
+lr_decay_style: cosine
+adam_beta1: 0.9
+adam_beta2: 0.95
+weight_decay: 0.01
+clip_grad: 1.0
+seed: 42
+layernorm_epsilon: 1.0e-6         # MUST be explicit — TransformerConfig default 1e-5 silently overrides the model YAML
+hidden_dropout: 0.0               # MUST be explicit — language_model.yaml default 0.1 leaks through
+attention_dropout: 0.0
+spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+use_distributed_optimizer: false  # 300M fits — ZeRO-1 adds allreduce overhead
+```
+
+The architecture-only YAML it extends from is `[primus/configs/models/megatron/zebra_llama_300M_gdn_pure.yaml](../../primus/configs/models/megatron/zebra_llama_300M_gdn_pure.yaml)`.
+
+### 5.2 Launch
+
+```bash
+# inside the container, in /home/<user>/Primus
+EXP=examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml \
+  bash examples/run_pretrain.sh 2>&1 | tee primus_gdn.log
+```
+
+This brings up `torchrun` with 8 ranks on the local node. Expected wall time on a healthy MI300X box: **~1h 54m** for the full 4768 iters.
+
+### 5.3 Recommended toggle profile (for FLA parity)
+
+The defaults are already good; for *bit-level* parity with FLA's
+optimizer/CE/SwiGLU kernels, set the following.  Either surface works
+(YAML wins for declarative runs; env vars win for ad-hoc overrides
+since the patch `fla_runtime_patches.py` does not overwrite an
+already-set env var).
+
+Preferred (canonical) — add to the experiment YAML's `overrides:` block:
+
+```yaml
+use_fla_fused_swiglu: true       # FLA Triton SwiGLU
+use_fla_fused_rmsnorm: true      # FLA fused RMSNorm
+use_fla_fused_gated_norm: true   # FLA FusedRMSNormGated path
+fused_ce_mode: 1                 # FLA FusedLinearCrossEntropyLoss (chunked, no full logits tensor)
+fused_ce_chunks: 32              # Chunk count for FLA fused CE
+# Only if you did Option B (FLA-aligned data) AND want bit-identical iter-1:
+use_fla_data: true
+fla_cache_dir: /home/<user>/Primus/data/huggingface
+```
+
+Legacy (still supported):
+
+```bash
+export PRIMUS_FUSED_CE=1          # FLA FusedLinearCrossEntropyLoss (chunked, no full logits tensor)
+export PRIMUS_FLA_SWIGLU=1        # FLA Triton SwiGLU
+export PRIMUS_FLA_NORM=1          # FLA fused RMSNorm + fused pre-norm/MLP path
+export PRIMUS_TORCH_OPTIM=1       # torch.optim.AdamW(fused=True), matches FLA exactly
+# Only if you did Option B (FLA-aligned data) AND want bit-identical iter-1:
+export PRIMUS_FLA_DATA=1
+export PRIMUS_FLA_CACHE_DIR=/home/<user>/Primus/data/huggingface
+```
+
+These add roughly +1.2 % per-iter overhead vs the all-defaults run, but they pin the loss curve to FLA's. On healthy hardware (tw006 in our cluster), the absolute wall is still ~18 s **below** FLA. On a slower node (tw029) it's ~75 s above. See [gdn-fla-parity.md](gdn-fla-parity.md) for the cost-of-each-flag breakdown.
+
+### 5.4 Output layout
+
+Checkpoints land under Primus's `work_group/user_name/exp_name` template:
+
+```
+output/amd/root/zebra_llama_300M_gdn_pure-pretrain/
+├── checkpoints/
+│   ├── iter_0001024/
+│   ├── iter_0002048/
+│   ├── iter_0003072/
+│   ├── iter_0004096/
+│   ├── iter_0004768/                  ← FINAL (4.1 GB)
+│   │   └── mp_rank_00/
+│   │       └── model_optim_rng.pt
+│   └── latest_checkpointed_iteration.txt  → "4768"
+└── logs/
+    └── pre_trainer/
+```
+
+`save_interval: 1024` in the YAML produces 4 mid-training checkpoints plus the final one.
+
+---
+
+## Step 6: Monitor and compare against FLA
+
+Megatron logs `iteration / elapsed_ms_inst / elapsed_ms_avg / TFLOP/s/GPU / tok/s/GPU / lm loss` every 100 steps. A representative tail looks like:
+
+```
+ iteration  4700/ 4768 | elapsed time per iteration (ms): 1460.8/1450.4 |
+   TFLOP/s/GPU: 633.7 | tokens per GPU (tokens/s/GPU): 180743.2 | lm loss: 3.3632
+```
+
+To diff against FLA's reference log (`train_gdn_bs32.log`, `train_runtime=6840.2s`):
+
+
+| iter | FLA / 8 | Primus  | Δ %         | Notes                        |
+| ---- | ------- | ------- | ----------- | ---------------------------- |
+| 1    | 11.9654 | 11.9652 | **−0.00 %** | bit-perfect                  |
+| 100  | 7.471   | 9.601   | +28.5 %     | warmup gap (peak)            |
+| 500  | 4.625   | 4.728   | +2.2 %      | warmup closing               |
+| 1000 | 4.001   | 4.050   | +1.21 %     | LR-warmup done               |
+| 2000 | 3.607   | 3.614   | +0.21 %     | converged                    |
+| 2100 | 3.600   | 3.592   | **−0.22 %** | first Primus < FLA crossover |
+| 3000 | 3.448   | 3.460   | +0.35 %     | matched                      |
+| 4000 | 3.396   | 3.390   | −0.19 %     | Primus slightly lower        |
+| 4500 | 3.373   | 3.373   | −0.01 %     | identical                    |
+| 4700 | 3.351   | 3.366   | +0.45 %     | identical                    |
+
+
+Final wall time on a healthy MI300X box: **6832 s vs FLA 6840 s** = Primus 8 s faster. On a slower node it's +75 s (~1.1 %). Both are within the run-to-run noise of FLA itself.
+
+---
+
+## Step 7: Convert checkpoint to HuggingFace format
+
+Use `[tools/hybrid/convert_gdn_to_fla_hf.py](../../tools/hybrid/convert_gdn_to_fla_hf.py)` to translate the Megatron checkpoint into FLA's native `GatedDeltaNetForCausalLM` HF format:
+
+```bash
+python tools/hybrid/convert_gdn_to_fla_hf.py \
+    --checkpoint-path output/amd/root/zebra_llama_300M_gdn_pure-pretrain/checkpoints/iter_0004768 \
+    --output-dir      output/gdn_pure_300M_fla_hf_final
+```
+
+The converter auto-detects 300M from the path and uses `gated_deltanet_300M_pure.json`. What it does:
+
+- Reads `mp_rank_00/model_optim_rng.pt` and pulls the `model` state dict
+- For each of the 12 FLA layers, pairs the alternating Megatron sublayers:
+  - GDN sublayer (even index) → FLA `model.layers.<i>.attn.*`
+  - MLP sublayer (odd index) → FLA `model.layers.<i>.mlp.*`
+- Splits Primus's **fused** projections into FLA's separate ones:
+  - `mixer.in_proj.weight` (rows = `2·key_dim + 2·value_dim + 2·num_v_heads`) → `q_proj / k_proj / v_proj / g_proj / b_proj / a_proj`
+  - `mixer.conv1d.weight` → `q_conv1d / k_conv1d / v_conv1d`
+  - `mlp.linear_fc1.weight` (rows = `2·intermediate_size`) → `gate_proj / up_proj`
+- Handles **both** layer-spec variants:
+  - TE spec (`gdn_hybrid_stack_spec`): norm fused into linear (`mixer.in_proj.layer_norm_weight`, `mlp.linear_fc1.layer_norm_weight`)
+  - No-TE spec (`gdn_hybrid_stack_spec_no_te`, **used by the validated run**): separate `WrappedTorchNorm` modules (`norm.weight`, `pre_mlp_layernorm.weight`)
+- Preserves `A_log`, `dt_bias`, per-head `out_norm`, `out_proj`, embeddings, tied `lm_head`, final norm
+
+Output:
+
+```
+output/gdn_pure_300M_fla_hf_final/
+├── config.json              # GatedDeltaNetConfig, architectures=["GatedDeltaNetForCausalLM"]
+├── model.safetensors        # ~870 MB
+└── tokenizer_config.json    # placeholder — point to meta-llama/Llama-3.2-1B at load time
+```
+
+For the 1B pure-GDN model, same command but use the 1B checkpoint path — the converter auto-selects `gated_deltanet_1B_pure.json`.
+
+---
+
+## Step 8: Verify conversion
+
+Run the sanity check at `[tools/hybrid/verify_gdn_conversion.py](../../tools/hybrid/verify_gdn_conversion.py)`:
+
+```bash
+python tools/hybrid/verify_gdn_conversion.py \
+    --model-path output/gdn_pure_300M_fla_hf_final
+```
+
+It loads the converted model in bf16 on GPU, runs three test prompts, and reports per-prompt loss, top-5 next-token IDs, and a 40-token greedy continuation. **Expected output** for a healthy 300M-on-10B model:
+
+
+| Prompt                                      | Loss | Top-1                                       | Verdict          |
+| ------------------------------------------- | ---- | ------------------------------------------- | ---------------- |
+| "The capital of France is"                  | ~3.7 | `Paris`                                     | knows the answer |
+| "Machine learning is a field of"            | ~2.8 | `artificial`                                | knows the domain |
+| "The largest planet in our solar system is" | ~2.3 | one of `[the, Jupiter, a, called, located]` | knows the topic  |
+
+
+Loss <6.0 = PASS. Greedy decoding will produce *grammatical but repetitive* English (e.g. *"Paris. The capital is Paris. The capital is Paris..."*) — this is the canonical small-undertrained-LM failure mode with no repetition penalty and is **not** a sign of conversion error.
+
+### Optional — logit parity vs the FLA reference checkpoint
+
+If you have the FLA reference HF checkpoint locally (e.g. trained by FLA itself, or downloaded from `fla-hub`), compare per-token logits:
+
+```bash
+python - <<'PY'
+import torch, fla
+from fla.models.gated_deltanet import GatedDeltaNetForCausalLM
+
+ids = torch.tensor([[1, 791, 6864, 315, 9822, 374]])  # "The capital of France is"
+
+hf  = GatedDeltaNetForCausalLM.from_pretrained("output/gdn_pure_300M_fla_hf_final",
+                                               torch_dtype=torch.bfloat16).cuda().eval()
+ref = GatedDeltaNetForCausalLM.from_pretrained("/path/to/fla/checkpoints/gdn_pure_300M_10BT",
+                                               torch_dtype=torch.bfloat16).cuda().eval()
+
+with torch.no_grad():
+    h = hf (ids.cuda()).logits[0, -1].float().cpu()
+    r = ref(ids.cuda()).logits[0, -1].float().cpu()
+
+print(f"Cosine sim:        {torch.nn.functional.cosine_similarity(h, r, dim=0).item():.4f}")
+print(f"Top-1 (converted): {h.argmax().item()}   Top-1 (FLA ref): {r.argmax().item()}")
+print(f"Top-5 (converted): {h.topk(5).indices.tolist()}")
+print(f"Top-5 (FLA ref):   {r.topk(5).indices.tolist()}")
+PY
+```
+
+**Expected:**
+
+
+| Metric            | Value       | Interpretation                                                     |
+| ----------------- | ----------- | ------------------------------------------------------------------ |
+| Cosine similarity | **≥ 0.95**  | conversion is correct                                              |
+| Top-5 set overlap | **≥ 3 / 5** | distributions agree                                                |
+| Top-1 exact match | optional    | a single-token disagreement is within 0.24 % loss divergence noise |
+
+
+If cosine < 0.5 → permutation bug. If 0.5–0.95 → likely missing-key issue (check that all 12 layers got `A_log`, `dt_bias`, `out_norm`).
+
+---
+
+## Step 9: Run lm-eval-harness benchmarks
+
+Use `[tools/hybrid/eval_gdn_lm_eval.py](../../tools/hybrid/eval_gdn_lm_eval.py)`, which imports `fla` first (so `AutoConfig` recognizes the `gated_deltanet` model type) and patches the FLA model `__init__` to accept the `dtype` kwarg that `transformers ≥ 4.55` passes internally.
+
+**Do not** invoke `lm_eval --model hf ...` directly — `AutoConfig.from_pretrained` will fail with `model type gated_deltanet not recognized`.
+
+### 9.1 Standard six-task suite (~15–30 min on one MI300X)
+
+```bash
+python tools/hybrid/eval_gdn_lm_eval.py \
+    --model hf \
+    --model_args pretrained=output/gdn_pure_300M_fla_hf_final,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+    --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
+    --batch_size auto \
+    --output_path output/gdn_pure_300M_eval_results_final
+```
+
+### 9.2 Full FLA-paper suite (adds MMLU + RACE, ~1–2 h)
+
+```bash
+python tools/hybrid/eval_gdn_lm_eval.py \
+    --model hf \
+    --model_args pretrained=output/gdn_pure_300M_fla_hf_final,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+    --batch_size auto \
+    --output_path output/gdn_pure_300M_eval_results_final
+```
+
+### 9.3 Diff against the FLA reference run
+
+If you also evaluated FLA's own checkpoint (`output/gdn_pure_300M_fla_eval_results/`), compare the JSONs:
+
+```bash
+python - <<'PY'
+import json, glob
+def load_latest(d): return json.load(open(sorted(glob.glob(f"{d}/**/results_*.json", recursive=True))[-1]))
+fla    = load_latest("output/gdn_pure_300M_fla_eval_results")
+primus = load_latest("output/gdn_pure_300M_eval_results_final")
+print(f"{'task':<18} {'FLA':>8} {'Primus':>8} {'Δ':>+8}")
+for task in sorted(set(fla['results']) & set(primus['results'])):
+    for k in ('acc,none', 'acc_norm,none'):
+        if k in fla['results'][task]:
+            f, p = fla['results'][task][k], primus['results'][task][k]
+            print(f"{task[:17]:<18} {f:>8.4f} {p:>8.4f} {p-f:>+8.4f}  ({k})")
+PY
+```
+
+**Expected:** each task within ±1.5 absolute accuracy points (consistent with the 0.24 % loss delta at the end of training).
+
+---
+
+## Configs and tools used
+
+```
+docs/04-technical-guides/hybrid-models/
+├── gdn-guide.md                                      ← this file
+└── gdn-fla-parity.md                                  ← deep-dive on every patch & env var
+primus/backends/megatron/patches/
+├── mamba_fused_ce_patches.py                          ← FLA fused cross-entropy for MambaModel
+├── torch_fused_adam_patches.py                        ← PRIMUS_TORCH_OPTIM opt-in
+├── mlp_fla_swiglu_patches.py                          ← FLA Triton SwiGLU for MLP
+├── torch_norm_fla_rmsnorm_patches.py                  ← FLA RMSNorm for WrappedTorchNorm
+├── gdn_config_patches.py                              ← linear-attention config fields + hybrid init
+├── fla_runtime_patches.py                             ← resolves PRIMUS_FLA_* knobs onto args
+└── mamba_fla_data_patches.py                          ← FLA-order dataset shim wiring
+examples/megatron/configs/MI300X/
+└── zebra_llama_300M_gdn_pure-pretrain.yaml            ← training config
+primus/configs/models/megatron/
+└── zebra_llama_300M_gdn_pure.yaml                     ← architecture-only config
+primus/backends/megatron/core/models/hybrid/
+├── gated_delta_net.py                                 ← FLA-aligned mixer (FLA Triton paths)
+├── gated_delta_net_layer.py                           ← eps propagation, pre-norm fusion
+├── hybrid_block.py                                    ← HybridStack, fp32-residual + fusion
+└── hybrid_mamba_mla_layer_specs.py                    ← gdn_hybrid_stack_spec_no_te
+tools/hybrid/
+├── patch_fla_triton_autotune_hang.sh                  ← MI300X FLA Triton autotune-hang workaround
+├── convert_fla_to_megatron.py                         ← FLA Arrow → Megatron .bin/.idx
+├── fla_order_dataset.py                               ← FLA-order dataset shim
+├── convert_gdn_to_fla_hf.py                           ← Megatron → FLA HF (handles TE + no-TE)
+├── verify_gdn_conversion.py                           ← loss + greedy generation sanity check
+└── eval_gdn_lm_eval.py                                ← lm-eval wrapper (registers FLA)
+```
+
+---
+
+## Troubleshooting
+
+### `KeyError: 'decoder.layers.0.mixer.in_proj.layer_norm_weight'` during conversion
+
+You trained with `gdn_hybrid_stack_spec_no_te` (separate `WrappedTorchNorm`) but are running an old version of `convert_gdn_to_fla_hf.py` that only knew the TE spec. Pull the latest converter — it now tries TE keys first and falls back to `norm.weight` / `pre_mlp_layernorm.weight`.
+
+### Loss is flat near 11.97 for many iterations
+
+LR warmup misconfigured. Confirm `lr_warmup_iters: 200` matches your `train_iters`, and verify the YAML override block resolved correctly by checking the logged config near the top of the training log.
+
+### Iter 1 loss ~12.1 instead of ~11.97
+
+The `layernorm_epsilon: 1.0e-6` override is being silently overwritten by the TransformerConfig default of `1e-5`. Confirm it's in the *training* YAML's `overrides:` block (not just the model YAML) — see `[zebra_llama_300M_gdn_pure-pretrain.yaml](../../examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml)` for the canonical placement.
+
+### Iter 1 loss not bit-matching FLA but converges fine
+
+You probably didn't enable `PRIMUS_FLA_DATA=1` — the Megatron `GPTDataset` shuffler is producing a different first batch than FLA's `DistributedSampler`. Either enable that env var (with `PRIMUS_FLA_CACHE_DIR` set) or accept the ~0.0002 loss delta at iter 1 (it disappears by iter ~2000).
+
+### Iter 1 takes ~60 seconds, subsequent iters are fast
+
+Cold MIOpen + Triton autotune caches. Normal on a freshly-rebooted node. The run-averaged ms/iter takes ~1500 iters to fully wash out this cold-start tax; the instantaneous ms/iter is at steady state by iter 200.
+
+### Eval fails with `model type gated_deltanet not recognized`
+
+You ran `lm_eval --model hf` directly instead of the wrapper. Use `python tools/hybrid/eval_gdn_lm_eval.py --model hf ...` — it imports `fla` first to register the model class.
+
+### Eval truncation warnings
+
+Some samples exceed the model's `max_position_embeddings = 2048`. Add `max_length=1024` to `--model_args` if it bothers you; it only meaningfully affects RACE.
+
+### Per-iter time +1–2 % above FLA on healthy hardware
+
+Expected with all four `PRIMUS_FLA_`* env vars set. The biggest single cost is `PRIMUS_TORCH_OPTIM=1` (torch fused AdamW vs Apex FusedAdam). Drop it if you don't need bit-level optimizer parity; you keep the loss-curve match and recover ~1 % perf.
+
+---
+
+## See also
+
+- `[docs/04-technical-guides/hybrid-models/README.md](README.md)` — full Zebra-Llama family overview (1B / 3B / 8B Mamba+MLA, KDA variants)
+- `[gdn-fla-parity.md](gdn-fla-parity.md)` — exhaustive list of code/config/runtime changes that made parity possible
+- FLA upstream: [https://github.com/fla-org/flash-linear-attention](https://github.com/fla-org/flash-linear-attention)
diff --git a/docs/04-technical-guides/hybrid-models/kda-fla-parity.md b/docs/04-technical-guides/hybrid-models/kda-fla-parity.md
new file mode 100644
index 000000000..d365215c1
--- /dev/null
+++ b/docs/04-technical-guides/hybrid-models/kda-fla-parity.md
@@ -0,0 +1,272 @@
+# KDA ⇄ FLA Parity in Primus
+
+This document captures every change required in Primus and the vendored
+Megatron-LM submodule to make a 300M Kimi Delta Attention (KDA) pretraining
+run match the [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention)
+reference implementation on **loss trajectory, step throughput, and
+downstream lm-eval accuracy** on 8× MI300X.
+
+This is the KDA-side companion to [`gdn-fla-parity.md`](gdn-fla-parity.md);
+because KDA shares Megatron-LM submodule patches with GDN, the architecture
+and tooling sections below focus on the KDA-specific deltas.
+
+## Final result
+
+| Axis | FLA reference | Primus (this branch) | Δ |
+|------|---------------|----------------------|----|
+| Per-iteration time (steady state, iter > 200) | **1493 ms** | **1466.8 ms** | **−1.8% (Primus faster)** |
+| Throughput (tok/s/GPU) | 175,617 | **178,810** | **+1.8%** |
+| TFLOP/s/GPU | — | 626.9 | — |
+| Total wall time (4768 iters) | 1h 58m 39s (7119.2 s) | **1h 56m 33s** (~6993 s) | **−126 s (Primus faster)** |
+| Loss @ iter 1 | 11.9673 | **11.9669** | **−0.00% (bit-perfect)** |
+| Loss @ iter 1000 | 4.0357 | 4.0720 | +0.90% |
+| Loss @ iter 2000 | 3.6009 | 3.6141 | +0.37% |
+| Loss late-training (iter 3700–4700 avg) | 3.3681 | 3.3846 | +0.49% |
+| First crossover (Primus < FLA) | — | iter 2600 (and 3600) | — |
+
+**Loss curves overlap from iter ~2000 onward**, with batch-to-batch
+oscillation of ±0.5%. The only persistent gap is in the LR-warmup region
+(iter 50–500), and that gap closes monotonically with no instability.
+Iter-1 forward at fp32 is bit-identical to FLA when the FLA-init checkpoint
+is loaded.
+
+### Downstream lm-eval parity
+
+After full training (4768 iters / ~10B tokens), both the Primus-trained
+KDA-300M and the FLA-trained KDA-300M were converted to HuggingFace
+`KDAForCausalLM` and evaluated with `lm-eval-harness` on the FLA-paper
+8-task suite. Every task is within ±1.4 absolute accuracy points, well
+inside the ±1.5 pp tolerance set by the 0.49% loss delta.
+
+The `Random` column is `100 / num_choices` for the task (25 % for
+4-choice tasks, 50 % for 2-choice tasks) — anything above it means the
+model has learned something. arc_easy / hellaswag / openbookqa / piqa
+clearly clear the bar; mmlu / race / arc_challenge sit at random for
+*both* training stacks (a 300 M model on 10 B tokens is below those
+benchmarks' lift-off threshold), which is exactly the regime the FLA
+paper reports.
+
+| Task                     | Metric     | Random | FLA    | Primus | Δ (Primus − FLA) |
+|--------------------------|------------|-------:|-------:|-------:|-----------------:|
+| arc_challenge            | acc_norm   |  25.00 | 25.17  | 25.00  | −0.17 pp         |
+| arc_easy                 | acc        |  25.00 | 48.78  | 47.94  | −0.84 pp         |
+| arc_easy                 | acc_norm   |  25.00 | 42.76  | 43.39  | +0.63 pp         |
+| hellaswag                | acc_norm   |  25.00 | 29.16  | 29.18  | +0.02 pp         |
+| openbookqa               | acc_norm   |  25.00 | 30.40  | 29.00  | −1.40 pp         |
+| piqa                     | acc_norm   |  50.00 | 60.99  | 60.34  | −0.65 pp         |
+| winogrande               | acc        |  50.00 | 51.85  | 52.72  | **+0.87 pp**     |
+| mmlu (aggregate)         | acc        |  25.00 | 22.88  | 23.12  | +0.24 pp         |
+| race                     | acc        |  25.00 | 25.07  | 25.45  | +0.38 pp         |
+| **mean absolute Δ**      |            |        |        |        | **0.58 pp**      |
+
+See [`kda-guide.md`](kda-guide.md) for
+the exact `lm_eval` invocation that produced both rows.
+
+---
+
+## How to run
+
+Inside the `rocm/primus:v26.2` container with the repo mounted at
+`/home/<user>/Primus`:
+
+```bash
+# 1. (one time) build the FLA-init KDA-300M checkpoint
+python tools/hybrid/convert_fla_kda_init_to_megatron.py
+#  → output/fla_init_kda_300M/iter_0000000/mp_rank_00/model_optim_rng.pt
+
+# 2. Launch training (8 GPUs by default). The Megatron-LM behavioral
+#    patches (same set as GDN) are applied automatically at startup via
+#    Primus's patch system -- no separate apply step needed.
+EXP=examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml \
+  bash examples/run_pretrain.sh 2>&1 | tee primus_kda.log
+```
+
+### Recommended toggle profile (YAML or env var)
+
+KDA uses the same toggle set as GDN.  Each knob is exposed at two
+equivalent surfaces — the YAML knob (canonical, declarative; co-located
+with the rest of the run config) and the legacy env var (ad-hoc, for
+one-off A/B without editing a YAML).  When both are set, the env var
+wins (backward compat); see
+`primus/backends/megatron/patches/fla_runtime_patches.py` for the
+precedence rules.  Defaults below match FLA's numerics on MI300X:
+
+| YAML knob | Env var | Default | Effect |
+|--|--|--|--|
+| `fused_ce_mode` | `PRIMUS_FUSED_CE` | `1` | `1` = FLA `FusedLinearCrossEntropyLoss` (chunked, no full logits tensor); `2` = FLA `FusedCrossEntropyLoss` (matches FLA exactly); `0` = native Megatron CE. |
+| `fused_ce_chunks` | `PRIMUS_FUSED_CE_CHUNKS` | `32` | Number of chunks the FLA CE splits the logits across.  Lower = faster but bigger peak allocation. |
+| `use_fla_fused_swiglu` | `PRIMUS_FLA_SWIGLU` | `1` | Replaces Megatron's naive SwiGLU with FLA's Triton-fused kernel (≈20 ms/step saved). |
+| `use_fla_fused_rmsnorm` | `PRIMUS_FLA_NORM` | `1` | Use FLA's `RMSNorm` Triton kernel via `WrappedTorchNorm`.  KDA's gated output norm is selected separately via `use_fla_fused_norm_gated` in the model YAML. |
+| `use_fla_short_conv` | `PRIMUS_FLA_CONV` | `1` | Route KDA's depthwise short conv1d through FLA's Triton `causal_conv1d` (saves ~35 ms/iter by accepting `[B, T, D]` directly — no `transpose+contiguous` round-trip). |
+| _(env-only)_ | `PRIMUS_TORCH_OPTIM` | `1` | Use `torch.optim.AdamW(fused=True)` instead of TE/Apex `FusedAdam` (matches FLA bit-for-bit). |
+| `use_fla_data` + `fla_cache_dir` | `PRIMUS_FLA_DATA` + `PRIMUS_FLA_CACHE_DIR` | `0` / `""` | When `use_fla_data=true` and `fla_cache_dir=<HF dataset cache path>`, replace Megatron's `GPTDataset` with the `FLAOrderGPTDataset` shim that emits tokens in the exact same order as FLA's HuggingFace `DistributedSampler`. |
+
+KDA's TE/no-TE selection is done by the `spec:` line in the YAML
+(`kda_hybrid_stack_spec_no_te` for no-TE, which is the default).
+
+---
+
+## What changed and why
+
+The work splits into three layers: KDA-specific model code, KDA-specific
+runtime config flags, and shared Megatron-LM patches (already documented
+in `gdn-fla-parity.md`).
+
+### A. Primus model code (KDA-specific)
+
+| File | Change | Reason |
+|------|--------|--------|
+| `primus/backends/megatron/core/models/hybrid/kimi_delta_attention.py` | Replace six separate `hidden_states → X` projections (q, k, v, beta, f_a, g_a) with a single fused `in_proj: ColumnParallelLinear` of width `2·qk_dim + v_dim + 2·head_v_dim + num_v_heads`. Split downstream into `[qkv | f_a | g_a | beta]`. The two low-rank-bottleneck expansion projections (`f_b`, `g_b`) stay separate because their input is the 64-dim bottleneck output, not `hidden_states`. | Matches GDN's fusion recipe. On ROCm each separate matmul pays ~3-5 ms of HIP dispatch + autograd overhead; the GDN parity work measured this same fusion at ~250 ms/iter saved. KDA was originally launching 6 matmuls/layer × 12 layers = 72 dispatches; now it's 12. |
+| same file | Add optional `FusedRMSNormGated` (RMSNorm + sigmoid-gate + multiply in ONE Triton kernel) for the per-head output gate, gated by `use_fla_fused_norm_gated` (default `True` when `use_fla_triton_kda=True`). | Avoids materializing the post-norm tensor and the fp32-upcast gate for backward — saves ~6.4 GiB activation memory per rank at micro_batch=128. Matches `fla/layers/kda.py` exactly. |
+| same file | Add optional in-kernel gate fusion path: when `use_fla_kda_in_kernel_gate=True`, call `chunk_kda(..., A_log=…, dt_bias=…, use_gate_in_kernel=True)` and let the kernel fuse `−exp(A_log) · softplus(g + dt_bias) + cumsum` internally (recomputed in backward). The pre-fusion `fused_kda_gate()` path is kept under `use_fla_kda_in_kernel_gate=False` for bit-identical comparison with FLA's old code. | Smallest activation footprint. The bf16 in-kernel accumulator drifts ~+0.2 lm-loss vs the explicit-gate path on ROCm at 12 layers depth; the FLA-init checkpoint cancels the drift, giving GDN-style parity. |
+| same file | Add optional FLA Triton `causal_conv1d` path under `args.use_fla_short_conv` (was `PRIMUS_FLA_CONV`). The FLA kernel accepts `[B, T, D]` directly (no `transpose+contiguous` round-trip). | Matches the conv backend FLA's `ShortConvolution` uses. Saves ~35 ms/iter (two avoided full-qkv buffer copies × ~17 ms each). |
+| same file | `g_b_proj.bias=True` and `dt_bias` initialised by FLA's log-uniform + inverse-softplus recipe (was `nn.init.ones_` → `dt ≈ 1.31`, ~20× larger than FLA's intended range). `beta = b_proj(h).float().sigmoid()` (fp32 sigmoid stops bf16 drift across 12 layers). Removed the `@torch.compiler.disable` decorator on `forward()`. | (a) `g_b_proj` bias matches `fla/layers/kda.py:189`. (b) `dt_bias` init matches `fla/layers/kda.py:180-184`; without it the gate's initial decay step is ~20× too large and the loss curve drifts visibly by iter 100. (c) fp32 sigmoid eliminates ~+0.2 lm-loss bf16 drift. (d) the compiler-disable was a leftover from debugging and cost ~25 ms/iter in dispatch overhead. |
+| same file | Materialize `q.contiguous() / k.contiguous() / v.contiguous()` after the `torch.split` on the fused in_proj output. | The `torch.split` along `dim=-1` returns non-contiguous views; passing them into `chunk_kda` as views makes the Triton kernel allocate a second internal contiguous copy while autograd still pins the original views. Net 2× activation memory for Q/K/V (~29 GiB extra at micro_batch=128). The explicit `.contiguous()` here gives autograd a single canonical buffer to save. Tested: 184 GiB → 155 GiB at iter 1. |
+| `primus/backends/megatron/core/models/hybrid/kimi_delta_attention_layer.py` | Add `KimiDeltaAttentionLayerSubmodules.norm` field (default `IdentityOp`). When set to `WrappedTorchNorm`, the layer applies an explicit pre-norm matching `fla/models/kda/modeling_kda.py:113` `hidden_states = self.attn_norm(...)`. `eps` is forwarded explicitly because `WrappedTorchNorm` defaults to `1e-5` while KDA configs (and FLA) use `1e-6`. | Required for the no-TE spec (which uses plain `ColumnParallelLinear` for `in_proj`) to apply the pre-norm separately. Without this fix the no-TE path skipped the pre-norm entirely, producing nonsense at iter 1. |
+| `primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py` | Add a new `kda_hybrid_stack_spec_no_te` ModuleSpec — plain `WrappedTorchNorm`, plain `ColumnParallelLinear`, plain `RowParallelLinear`, mixer `gate_norm=IdentityOp` (FLA has no re-norm for the gate path). | YAML can now select TE-free KDA layers via `spec: [..., kda_hybrid_stack_spec_no_te]` for FLA loss-curve alignment without touching code. Mirrors `gdn_hybrid_stack_spec_no_te`. |
+| `primus/backends/megatron/patches/gdn_config_patches.py` | Register `use_fla_kda_in_kernel_gate` (default `True`) and `use_fla_fused_norm_gated` (default `None` → auto when `use_fla_triton_kda=True`) as `TransformerConfig` fields. | Lets the YAML `overrides:` block toggle the two KDA-specific fusion paths without touching code. |
+
+### B. Megatron-LM behavioral patches (shared with GDN, Primus patch system)
+
+KDA reuses the **exact same six patches** that GDN uses; no KDA-specific
+megatron-LM patch is required. These are runtime monkey-patches registered
+with `@register_patch` under `primus/backends/megatron/patches/` and applied
+automatically at `phase="before_train"` -- see `gdn-fla-parity.md` section B
+for the patch-by-patch breakdown. Nothing needs to be run by hand.
+
+### C. YAML configuration changes
+
+#### `primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml` (new)
+
+300M architecture-only YAML matched to FLA's `kda_300M_pure.json`:
+
+```yaml
+extends: [mamba_base.yaml]
+
+num_layers: 24                  # 12 KDA + 12 MLP sublayers
+hidden_size: 1024
+ffn_hidden_size: 4096
+
+# Pure KDA — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+
+# KDA params (match FLA exactly)
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 32         # 8 heads × 32 = 256 qk_dim
+linear_value_head_dim: 64       # 8 heads × 64 = 512 v_dim (expand_v=2.0)
+linear_num_key_heads: 8
+linear_num_value_heads: 8
+
+# Tied embeddings, all linear bias=False, RMSNorm eps=1e-6
+untie_embeddings_and_output_weights: false
+add_bias_linear: false
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+position_embedding_type: none
+```
+
+#### `examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml` (new)
+
+The training-side config sets:
+
+```yaml
+# Training schedule matched to FLA (8 GPUs)
+train_iters: 4768                       # ≈10B tokens at 1024×2048 = 2.1M tok/iter
+micro_batch_size: 128
+global_batch_size: 1024
+
+# FLA optimizer / LR schedule
+lr: 2.0e-4
+min_lr: 2.0e-5                          # min_lr_rate=0.1
+lr_warmup_iters: 200
+lr_decay_iters: 4768
+lr_decay_style: cosine
+adam_beta1: 0.9; adam_beta2: 0.95
+weight_decay: 0.01; clip_grad: 1.0
+seed: 42
+
+# Norm — Megatron default is 1e-5; FLA uses 1e-6
+layernorm_epsilon: 1.0e-6
+hidden_dropout: 0.0; attention_dropout: 0.0
+
+# Pure KDA, no-TE spec (matches FLA KDABlock layout exactly)
+spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs',
+       'kda_hybrid_stack_spec_no_te']
+use_fla_triton_kda: true
+use_fla_kda_in_kernel_gate: true
+use_fla_fused_norm_gated: true
+
+# Plain DDP, matches FLA — distributed optimizer (ZeRO-1) costs allreduce
+# bandwidth and saves only ~3.6 GiB/rank for a 300M model
+use_distributed_optimizer: false
+overlap_grad_reduce: true
+ddp_average_in_collective: true
+
+# FLA-init checkpoint — bit-perfect iter-1 forward
+finetune: true; no_load_optim: true; no_load_rng: true
+load: /home/<user>/Primus/output/fla_init_kda_300M
+```
+
+---
+
+## Reproducing the loss-curve match plot
+
+The full per-iteration log lives at `primus_kda.log` once training
+finishes. Compare against FLA's `trainer_state.json` log_history
+(`/home/<user>/checkpoints/kda_pure_300M_10B/trainer_state.json`).
+
+Notable comparison points (FLA loss is divided by 8 to undo the
+DeepSpeed sum-across-ranks):
+
+| iter | FLA / 8 | Primus | Δ% | Notes |
+|-----:|--------:|-------:|---:|-------|
+| 1    | 11.9673 | 11.9669 | **−0.00%** | bit-perfect (forward fp32) |
+| 100  | 7.7171  | 9.6903  | +25.6%     | warmup gap (peak) |
+| 500  | 4.7349  | 4.8390  | +2.20%     | warmup closing |
+| 1000 | 4.0357  | 4.0720  | +0.90%     | LR-warmup done |
+| 2000 | 3.6009  | 3.6141  | +0.37%     | converged |
+| 2600 | 3.5056  | 3.5047  | **−0.03%** | first Primus < FLA crossover |
+| 3000 | 3.4356  | 3.4571  | +0.63%     | matched |
+| 3600 | 3.4107  | 3.4075  | **−0.09%** | Primus slightly lower |
+| 4000 | 3.3831  | 3.3861  | +0.09%     | identical |
+| 4500 | 3.3603  | 3.3694  | +0.27%     | identical |
+| 4700 | 3.3388  | 3.3624  | +0.71%     | identical |
+
+The persistent gap (iter 50–500) is attributable to dataloader ordering —
+Megatron `GPTDataset` uses its own random shuffler while FLA uses
+HuggingFace's `DistributedSampler`. With `use_fla_data: true` the gap
+closes further but Primus has been verified to converge to within ±1% by
+iter 1000 even without it.
+
+---
+
+## Files in the repo for this work
+
+```
+primus/backends/megatron/core/models/hybrid/
+  kimi_delta_attention.py                          # FLA-aligned mixer
+  kimi_delta_attention_layer.py                    # wrapper w/ pre-norm
+  hybrid_mamba_mla_layer_specs.py                  # kda_hybrid_stack_spec_no_te
+primus/backends/megatron/patches/                  # 6 patches (same as GDN), Primus patch system
+  gdn_config_patches.py                            # registers KDA fusion flags + hybrid init
+  mamba_fused_ce_patches.py                        # FLA fused cross-entropy for MambaModel
+  torch_fused_adam_patches.py                      # PRIMUS_TORCH_OPTIM opt-in
+  mlp_fla_swiglu_patches.py                        # FLA Triton SwiGLU for MLP
+  torch_norm_fla_rmsnorm_patches.py                # FLA RMSNorm for WrappedTorchNorm
+  fla_runtime_patches.py                           # resolves PRIMUS_FLA_* knobs onto args
+  mamba_fla_data_patches.py                        # FLA-order dataset shim wiring
+primus/configs/models/megatron/
+  zebra_llama_300M_kda_pure.yaml                   # architecture-only
+examples/megatron/configs/MI300X/
+  zebra_llama_300M_kda_pure-pretrain.yaml          # training config
+tools/hybrid/
+  convert_fla_to_megatron.py                       # FLA Arrow → Megatron .bin/.idx (shared)
+  fla_order_dataset.py                             # FLA-order dataset shim (shared)
+  convert_fla_kda_init_to_megatron.py              # FLA HF init → Megatron sharded ckpt
+  convert_kda_to_fla_hf.py                         # Megatron sharded ckpt → FLA HF
+  eval_kda_lm_eval.py                              # lm-eval wrapper (registers KDA)
+docs/04-technical-guides/hybrid-models/
+  kda-guide.md                                    # step-by-step recipe
+  kda-fla-parity.md                                # this file
+```
diff --git a/docs/04-technical-guides/hybrid-models/kda-guide.md b/docs/04-technical-guides/hybrid-models/kda-guide.md
new file mode 100644
index 000000000..54b2e6f59
--- /dev/null
+++ b/docs/04-technical-guides/hybrid-models/kda-guide.md
@@ -0,0 +1,646 @@
+# Pure KDA 300M on Primus — End-to-End Guide (FLA-validated)
+
+This document is a runnable walkthrough for the **300M pure Kimi Delta
+Attention (KDA)** pretraining recipe in Primus, validated on 8× AMD MI300X
+against the [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention)
+reference implementation. It covers every step from raw dataset →
+tokenization → training → checkpoint conversion → lm-eval benchmark.
+
+The same recipe scales up to the 1B pure-KDA config (`zebra_llama_1B_kda_pure-pretrain.yaml`).
+
+It mirrors [`gdn-guide.md`](gdn-guide.md) and reuses the same Megatron-LM
+patches, dataset shim, FLA-init flow, and lm-eval wrapper pattern.
+
+---
+
+## Final result
+
+After 4768 iterations (≈10B tokens) on FineWeb-Edu sample-10BT:
+
+| Axis                                            | FLA reference     | Primus (this branch)  | Δ                           |
+| ----------------------------------------------- | ----------------- | --------------------- | --------------------------- |
+| Per-iteration time (steady state, iter > 200)   | **1493 ms**       | **1466.8 ms**         | **−1.8 % (Primus faster)**  |
+| Throughput                                      | 175,617 tok/s/GPU | **178,810 tok/s/GPU** | **+1.8 %**                  |
+| TFLOP/s/GPU                                     | —                 | 626.9                 | —                           |
+| Wall time (4768 iters, 8× MI300X, healthy node) | 1h 58m 39s        | **1h 56m 33s**        | **−126 s**                  |
+| Loss @ iter 1                                   | 11.9673           | **11.9669**           | **−0.00 % (bit-perfect)**   |
+| Loss @ iter 4700 (final logged)                 | 3.3388            | **3.3624**            | **+0.71 %**                 |
+| First Primus-below-FLA crossover                | —                 | iter 2600             | —                           |
+
+Loss trajectories overlap from iter ~2000 onward; the only persistent gap
+is in the LR-warmup region (iter 50–500) and closes monotonically. See
+[`kda-fla-parity.md`](kda-fla-parity.md) for the deep-dive on every
+patch and env var.
+
+### lm-eval-harness (FLA-paper 8-task suite)
+
+Random chance is `100 / num_choices` — 25 % for the 4-choice tasks
+(arc, hellaswag, openbookqa, mmlu, race) and 50 % for the 2-choice tasks
+(piqa, winogrande). Any score above random shows the model has learned
+*something*; the FLA and Primus rows show how closely the two training
+stacks track each other on the same 10 B-token diet.
+
+| Task                     | Metric     | Random | FLA    | Primus | Δ (Primus − FLA) |
+|--------------------------|------------|-------:|-------:|-------:|-----------------:|
+| arc_challenge            | acc_norm   |  25.00 | 25.17  | 25.00  | −0.17 pp         |
+| arc_easy                 | acc        |  25.00 | 48.78  | 47.94  | −0.84 pp         |
+| arc_easy                 | acc_norm   |  25.00 | 42.76  | 43.39  | +0.63 pp         |
+| hellaswag                | acc_norm   |  25.00 | 29.16  | 29.18  | +0.02 pp         |
+| openbookqa               | acc_norm   |  25.00 | 30.40  | 29.00  | −1.40 pp         |
+| piqa                     | acc_norm   |  50.00 | 60.99  | 60.34  | −0.65 pp         |
+| winogrande               | acc        |  50.00 | 51.85  | 52.72  | **+0.87 pp**     |
+| mmlu (aggregate)         | acc        |  25.00 | 22.88  | 23.12  | +0.24 pp         |
+| race                     | acc        |  25.00 | 25.07  | 25.45  | +0.38 pp         |
+| **mean absolute Δ**      |            |        |        |        | **0.58 pp**      |
+
+Every task within ±1.4 pp — well inside the ±1.5 pp tolerance set by the
+0.49% mid-training loss delta. Both stacks comfortably beat random on
+arc_easy, hellaswag, openbookqa and piqa; mmlu/race/arc_challenge are at
+random-chance for *both* training stacks (expected for a 300 M model on
+only 10 B tokens — those benchmarks need 7 B+ parameters and/or
+trillion-token training to lift above 25 %).
+
+---
+
+## Table of contents
+
+- [Overview](#overview)
+- [Prerequisites](#prerequisites)
+- [Step 1: Environment](#step-1-environment)
+- [Step 2: Dataset preparation](#step-2-dataset-preparation)
+- [Step 3: Apply Megatron-LM patches](#step-3-apply-megatron-lm-patches)
+- [Step 4: (Optional) Initialize from FLA weights](#step-4-optional-initialize-from-fla-weights)
+- [Step 5: Train](#step-5-train)
+- [Step 6: Monitor and compare against FLA](#step-6-monitor-and-compare-against-fla)
+- [Step 7: Convert checkpoint to HuggingFace format](#step-7-convert-checkpoint-to-huggingface-format)
+- [Step 8: Verify conversion](#step-8-verify-conversion)
+- [Step 9: Run lm-eval-harness benchmarks](#step-9-run-lm-eval-harness-benchmarks)
+- [Configs and tools used](#configs-and-tools-used)
+- [Troubleshooting](#troubleshooting)
+
+---
+
+## Overview
+
+The 300M pure-KDA model has:
+
+- 12 Kimi-Delta-Attention blocks + 12 MLP blocks → 24 Megatron "sublayers"
+- `hidden_size = 1024`, `ffn_hidden_size = 4096`
+- `num_heads = num_v_heads = 8` (Q, K, V all share head count)
+- `head_k_dim = 32`, `head_v_dim = 64` (expand_v = 2.0)
+- Short-conv kernel size 4 (depthwise, on the concatenated QKV)
+- Per-head output gate (`g_a → g_b`) + per-head decay gate (`f_a → f_b`,
+  combined with learnable `A_log` and `dt_bias` via `softplus`)
+- Tied embeddings, no positional encoding (delta-rule recurrence),
+  RMSNorm with `eps = 1e-6`
+- Tokenizer: `meta-llama/Llama-3.2-1B` (128k vocab)
+- Total parameters: **0.302 B**
+
+Training schedule (matched to FLA's `kda_300M_pure.json`):
+
+- 4768 iterations × 1024 global batch × 2048 seq len = **10.0 B tokens**
+- AdamW (β1=0.9, β2=0.95, wd=0.01), peak LR `2e-4`, cosine decay, 200-step warmup
+- bf16 mixed-precision, no dropout, gradient clip 1.0
+
+---
+
+## Prerequisites
+
+- **Hardware**: 8× AMD MI300X (or compatible ROCm GPU) on a single node
+- **Software**: ROCm ≥ 7.0, Docker ≥ 24.0
+- **Container image**: `rocm/primus:v26.2` (or `v25.10` with the same patches)
+- **HF token**: `HF_TOKEN` set for the gated `meta-llama/Llama-3.2-1B` tokenizer
+- **Disk**: ~20 GB for the FLA-aligned tokenized dataset + ~5 GB per saved checkpoint
+- **flash-linear-attention** checked out at
+  `/home/<user>/flash-linear-attention` (or installed via
+  `pip install -e .`) — provides the FLA `KDAForCausalLM` class for
+  HF conversion + lm-eval, plus the Triton kernels that the `PRIMUS_FLA_*`
+  toggles route into.
+
+---
+
+## Step 1: Environment
+
+### 1.1 Start the dev container
+
+```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 -v $(pwd):$(pwd) -w $(pwd) --shm-size 64G --name primus_hybrid_new \
+  rocm/primus:v26.2
+```
+
+This runs the `rocm/primus:v26.2` image with `/dev/dri`, `/dev/kfd`, IB
+devices, `--privileged`, your `$HOME` mounted in-place, and `--shm-size 64G`.
+The container is named `primus_hybrid_new`.
+
+To re-attach later:
+
+```bash
+docker exec -it primus_hybrid_new bash
+cd /home/<user>/Primus
+```
+
+### 1.2 Install Python dependencies inside the container
+
+```bash
+pip install -r requirements.txt
+pip install -e /home/<user>/flash-linear-attention   # FLA model classes + Triton kernels
+pip install lm-eval                                  # for benchmark evaluation
+```
+
+The editable FLA install removes the need to set `PYTHONPATH` for every
+later command.
+
+---
+
+## Step 2: Dataset preparation
+
+Identical to the GDN recipe — see
+[`gdn-guide.md`](gdn-guide.md#step-2-dataset-preparation). KDA reuses the
+same FineWeb-Edu sample-10BT preprocessed Arrow shards and the same
+Llama-3.2-1B tokenizer.
+
+The default 300M YAML already points at the FLA-aligned binary:
+
+```yaml
+train_data_path: >
+  /home/<user>/Primus/data/fla_aligned/fla_fineweb_edu_10BT_text_sentence
+```
+
+(adjust the user prefix in
+`examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml`
+to match your home directory).
+
+---
+
+## Step 3: Megatron-LM patches (automatic — no action needed)
+
+KDA uses the **same six patches** as GDN — no KDA-specific Megatron patch
+is required. They're implemented as runtime monkey-patches using Primus's
+own patch system (`primus/core/patches`), living under
+`[primus/backends/megatron/patches/](../../primus/backends/megatron/patches/)`.
+Each patch registers unconditionally but is gated behind a `condition=` on
+the relevant config flag, so nothing needs to be run by hand.
+
+See [`gdn-guide.md`](gdn-guide.md#step-3-megatron-lm-patches-automatic--no-action-needed)
+§3 for the patch-by-patch breakdown.
+
+---
+
+## Step 4: (Optional) Initialize from FLA weights
+
+For bit-perfect iter-1 loss alignment, the validated run loads FLA's
+*initialized but untrained* KDA-300M checkpoint and then trains from
+there. The YAML's `load:` field points at this directory:
+
+```yaml
+load: /home/<user>/Primus/output/fla_init_kda_300M
+finetune: true            # load weights, ignore optimizer state and iteration count
+no_load_optim: true
+no_load_rng: true
+```
+
+Generate it once with:
+
+```bash
+python tools/hybrid/convert_fla_kda_init_to_megatron.py
+#   → output/fla_init_kda_300M/iter_0000000/mp_rank_00/model_optim_rng.pt
+```
+
+The script instantiates FLA's `KDAForCausalLM` with `seed=42`, harvests
+its randomly-initialized weights, concatenates the six FLA `hidden_states
+→ X` projections into Primus's single fused `in_proj`, and writes a
+Megatron-shape checkpoint. Skip this step if you're happy with Primus's
+own random init — final loss is identical, only iter-1 drifts by `~5e-3`.
+
+---
+
+## Step 5: Train
+
+### 5.1 Inspect the config
+
+The training config lives at
+[`examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml`](../../examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml).
+Key parameters (matched to FLA):
+
+```yaml
+train_iters: 4768                 # ≈ 10B tokens at global_batch=1024, seq=2048
+micro_batch_size: 128             # per-GPU
+global_batch_size: 1024           # 8 GPUs × 128 = 1024
+seq_length: 2048
+lr: 2.0e-4
+min_lr: 2.0e-5                    # min_lr_rate=0.1 → 2e-5
+lr_warmup_iters: 200
+lr_decay_iters: 4768
+lr_decay_style: cosine
+adam_beta1: 0.9
+adam_beta2: 0.95
+weight_decay: 0.01
+clip_grad: 1.0
+seed: 42
+layernorm_epsilon: 1.0e-6         # MUST be explicit — TransformerConfig default 1e-5 silently overrides the model YAML
+hidden_dropout: 0.0               # MUST be explicit — language_model.yaml default 0.1 leaks through
+attention_dropout: 0.0
+spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec_no_te']
+use_fla_triton_kda: true
+use_fla_kda_in_kernel_gate: true
+use_fla_fused_norm_gated: true
+use_distributed_optimizer: false  # 300M fits — ZeRO-1 adds allreduce overhead
+finetune: true
+load: /home/<user>/Primus/output/fla_init_kda_300M
+no_load_optim: true
+no_load_rng: true
+```
+
+The architecture-only YAML it extends from is
+[`primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml`](../../primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml).
+
+### 5.2 Launch
+
+```bash
+# inside the container, in /home/<user>/Primus
+EXP=examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml \
+  bash examples/run_pretrain.sh 2>&1 | tee primus_kda.log
+```
+
+Expected wall time on a healthy MI300X box: **~1h 56m** for the full 4768
+iters (about 2 min faster than FLA's HF-Trainer reference run).
+
+### 5.3 Recommended toggle profile (for FLA parity)
+
+Preferred (canonical) — add to the experiment YAML's `overrides:` block:
+
+```yaml
+# FLA runtime knobs (consumed by primus.backends.megatron.patches.fla_runtime_patches)
+use_fla_fused_swiglu: true       # FLA Triton SwiGLU
+use_fla_fused_rmsnorm: true      # FLA fused RMSNorm
+use_fla_fused_gated_norm: true   # FLA FusedRMSNormGated for KDA gated output norm
+use_fla_short_conv: true         # FLA Triton causal_conv1d (no transpose round-trip)
+fused_ce_mode: 1                 # FLA FusedLinearCrossEntropyLoss (chunked, no full logits tensor)
+fused_ce_chunks: 32              # Chunk count for FLA fused CE
+# Only if you want bit-identical iter-1 batch ordering:
+use_fla_data: true
+fla_cache_dir: /home/<user>/Primus/data/huggingface
+```
+
+Legacy (still supported, env-var wins over YAML when set):
+
+```bash
+export PRIMUS_FUSED_CE=1          # FLA FusedLinearCrossEntropyLoss (chunked, no full logits tensor)
+export PRIMUS_FLA_SWIGLU=1        # FLA Triton SwiGLU
+export PRIMUS_FLA_NORM=1          # FLA fused RMSNorm
+export PRIMUS_FLA_CONV=1          # FLA Triton causal_conv1d (no transpose round-trip)
+export PRIMUS_TORCH_OPTIM=1       # torch.optim.AdamW(fused=True), matches FLA exactly
+# Only if you want bit-identical iter-1 batch ordering:
+export PRIMUS_FLA_DATA=1
+export PRIMUS_FLA_CACHE_DIR=/home/<user>/Primus/data/huggingface
+```
+
+See [`kda-fla-parity.md`](kda-fla-parity.md) for the cost-of-each-flag
+breakdown.
+
+### 5.4 Output layout
+
+Checkpoints land under Primus's `work_group/user_name/exp_name` template:
+
+```
+output/amd/root/zebra_llama_300M_kda_pure-pretrain/
+├── checkpoints/
+│   ├── iter_0001024/
+│   ├── iter_0002048/
+│   ├── iter_0003072/
+│   ├── iter_0004096/
+│   ├── iter_0004768/                  ← FINAL (~4.5 GB)
+│   │   └── mp_rank_00/
+│   │       └── model_optim_rng.pt
+│   └── latest_checkpointed_iteration.txt  → "4768"
+└── logs/
+    └── pre_trainer/
+```
+
+`save_interval: 1024` in the YAML produces 4 mid-training checkpoints plus
+the final one.
+
+---
+
+## Step 6: Monitor and compare against FLA
+
+Megatron logs `iteration / elapsed_ms_inst / elapsed_ms_avg / TFLOP/s/GPU
+/ tok/s/GPU / lm loss` every 100 steps. A representative tail looks like:
+
+```
+iteration  4700/ 4768 | elapsed time per iteration (ms): 1467.8/1466.1 |
+  TFLOP/s/GPU: 626.1 | tokens per GPU (tokens/s/GPU): 178596.5 | lm loss: 3.362445E+00
+```
+
+To diff against FLA's reference log
+(`/home/<user>/checkpoints/kda_pure_300M_10B/trainer_state.json`), divide
+the FLA `loss` field by 8 (DeepSpeed reports sum-across-ranks):
+
+| iter | FLA / 8 | Primus  | Δ %         | Notes                            |
+| ---- | ------- | ------- | ----------- | -------------------------------- |
+| 1    | 11.9673 | 11.9669 | **−0.00 %** | bit-perfect                      |
+| 100  | 7.7171  | 9.6903  | +25.6 %     | warmup gap (peak)                |
+| 500  | 4.7349  | 4.8390  | +2.20 %     | warmup closing                   |
+| 1000 | 4.0357  | 4.0720  | +0.90 %     | LR-warmup done                   |
+| 2000 | 3.6009  | 3.6141  | +0.37 %     | converged                        |
+| 2600 | 3.5056  | 3.5047  | **−0.03 %** | first Primus < FLA crossover     |
+| 3000 | 3.4356  | 3.4571  | +0.63 %     | matched                          |
+| 3600 | 3.4107  | 3.4075  | **−0.09 %** | Primus slightly lower            |
+| 4000 | 3.3831  | 3.3861  | +0.09 %     | identical                        |
+| 4500 | 3.3603  | 3.3694  | +0.27 %     | identical                        |
+| 4700 | 3.3388  | 3.3624  | +0.71 %     | identical                        |
+
+Final wall time on a healthy MI300X box: **6993 s vs FLA 7119 s** =
+Primus 126 s faster.
+
+---
+
+## Step 7: Convert checkpoint to HuggingFace format
+
+Use [`tools/hybrid/convert_kda_to_fla_hf.py`](../../tools/hybrid/convert_kda_to_fla_hf.py)
+to translate the Megatron checkpoint into FLA's native
+`KDAForCausalLM` HF format:
+
+```bash
+python tools/hybrid/convert_kda_to_fla_hf.py \
+    --checkpoint-path output/amd/root/zebra_llama_300M_kda_pure-pretrain/checkpoints/iter_0004768 \
+    --output-dir      output/kda_pure_300M_fla_hf \
+    --config          /home/<user>/flash-linear-attention/legacy/training/configs/kda_300M_pure.json \
+    --tokenizer-src   /home/<user>/checkpoints/kda_pure_300M_10B
+```
+
+What it does:
+
+- Reads `mp_rank_00/model_optim_rng.pt` and pulls the `model` state dict
+- For each of the 12 FLA layers, pairs the alternating Megatron sublayers:
+  - KDA sublayer (even index) → FLA `model.layers.<i>.attn.*`
+  - MLP sublayer (odd index) → FLA `model.layers.<i>.mlp.*`
+- Splits Primus's **fused** projections into FLA's separate ones:
+  - `mixer.in_proj.weight` (rows = `2·qk_dim + v_dim + 2·head_v_dim +
+    num_v_heads`) → `q_proj / k_proj / v_proj / f_proj.0 / g_proj.0 / b_proj`
+  - `mlp.linear_fc1.weight` (rows = `2·intermediate_size`) →
+    `gate_proj / up_proj`
+- Preserves `A_log`, `dt_bias`, per-head `g_norm` (FLA's
+  `FusedRMSNormGated`), `o_proj`, `f_proj.1`, `g_proj.1`, embeddings,
+  tied `lm_head`, final norm
+- Copies tokenizer files from `--tokenizer-src` into the output dir
+
+Output:
+
+```
+output/kda_pure_300M_fla_hf/
+├── config.json              # KDAConfig, architectures=["KDAForCausalLM"]
+├── model.safetensors        # ~870 MB
+└── tokenizer{,_config}.json + special_tokens_map.json
+```
+
+---
+
+## Step 8: Verify conversion
+
+Quick smoke test in the container (with FLA importable):
+
+```bash
+PYTHONPATH=/home/<user>/flash-linear-attention \
+python - <<'PY'
+import torch
+import fla   # auto-registers "kda" with transformers.AutoConfig
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ckpt = "output/kda_pure_300M_fla_hf"
+tok = AutoTokenizer.from_pretrained(ckpt)
+model = AutoModelForCausalLM.from_pretrained(
+    ckpt, trust_remote_code=True, torch_dtype=torch.bfloat16
+).cuda().eval()
+
+for prompt in [
+    "The capital of France is",
+    "Once upon a time, there was a small",
+    "The first law of thermodynamics states that",
+]:
+    inp = tok(prompt, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+        out = model.generate(**inp, max_new_tokens=40, do_sample=False)
+    print("---"); print(tok.decode(out[0], skip_special_tokens=True))
+PY
+```
+
+**Expected output** for a healthy 300 M-on-10 B model: grammatical but
+repetitive English (canonical small-undertrained-LM failure mode under
+greedy decoding with no repetition penalty). Knowing "capital of France"
+→ "Paris" is the standard sanity-check pass.
+
+If `AutoConfig` raises `model type kda not recognized`, FLA was not
+imported before `AutoModelForCausalLM`. Either prepend
+`PYTHONPATH=/home/<user>/flash-linear-attention` or run
+`pip install -e /home/<user>/flash-linear-attention` so the
+auto-registration in `fla/models/kda/__init__.py` fires on import.
+
+---
+
+## Step 9: Run lm-eval-harness benchmarks
+
+Use [`tools/hybrid/eval_kda_lm_eval.py`](../../tools/hybrid/eval_kda_lm_eval.py), which
+imports `fla` first (so `AutoConfig` recognizes the `kda` model type) and
+patches `KDAForCausalLM.__init__` / `KDAModel.__init__` to accept the
+`dtype` kwarg that `transformers ≥ 4.55` passes internally.
+
+**Do not** invoke `lm_eval --model hf ...` directly — `AutoConfig.from_pretrained`
+will fail with `model type kda not recognized`.
+
+### 9.1 Evaluate the Primus checkpoint (~15–30 min on one MI300X)
+
+```bash
+mkdir -p output/kda_pure_300M_eval_results_primus
+
+PYTHONPATH=/home/<user>/flash-linear-attention \
+HIP_VISIBLE_DEVICES=0 \
+TOKENIZERS_PARALLELISM=false \
+python tools/hybrid/eval_kda_lm_eval.py \
+    --model hf \
+    --model_args pretrained=output/kda_pure_300M_fla_hf,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+    --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
+    --batch_size auto \
+    --output_path output/kda_pure_300M_eval_results_primus \
+    2>&1 | tee output/kda_pure_300M_eval_results_primus/lm_eval.log
+```
+
+### 9.2 Evaluate the FLA reference checkpoint (apples-to-apples)
+
+```bash
+mkdir -p output/kda_pure_300M_eval_results_fla
+
+PYTHONPATH=/home/<user>/flash-linear-attention \
+HIP_VISIBLE_DEVICES=1 \
+TOKENIZERS_PARALLELISM=false \
+python tools/hybrid/eval_kda_lm_eval.py \
+    --model hf \
+    --model_args pretrained=/home/<user>/checkpoints/kda_pure_300M_10B,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+    --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
+    --batch_size auto \
+    --output_path output/kda_pure_300M_eval_results_fla \
+    2>&1 | tee output/kda_pure_300M_eval_results_fla/lm_eval.log
+```
+
+### 9.3 Diff the two result JSONs
+
+```bash
+python - <<'PY'
+import json, glob
+def load_latest(d):
+    return json.load(open(sorted(glob.glob(f"{d}/**/results_*.json", recursive=True))[-1]))
+fla    = load_latest("output/kda_pure_300M_eval_results_fla")
+primus = load_latest("output/kda_pure_300M_eval_results_primus")
+print(f"{'task':<18} {'FLA':>8} {'Primus':>8} {'Δ':>+8}")
+for task in sorted(set(fla['results']) & set(primus['results'])):
+    for k in ('acc,none', 'acc_norm,none'):
+        if k in fla['results'][task] and k in primus['results'][task]:
+            f, p = fla['results'][task][k], primus['results'][task][k]
+            print(f"{task[:17]:<18} {f:>8.4f} {p:>8.4f} {p-f:>+8.4f}  ({k})")
+PY
+```
+
+**Measured result** (validated on `tw006`, this branch). The `Random`
+column is `100 / num_choices` for the lm-eval task — anything above it
+means the model learned something:
+
+| Task                     | Metric     | Random | FLA    | Primus | Δ (Primus − FLA) |
+|--------------------------|------------|-------:|-------:|-------:|-----------------:|
+| arc_challenge            | acc_norm   |  25.00 | 25.17  | 25.00  | −0.17 pp         |
+| arc_easy                 | acc        |  25.00 | 48.78  | 47.94  | −0.84 pp         |
+| arc_easy                 | acc_norm   |  25.00 | 42.76  | 43.39  | +0.63 pp         |
+| hellaswag                | acc_norm   |  25.00 | 29.16  | 29.18  | +0.02 pp         |
+| openbookqa               | acc_norm   |  25.00 | 30.40  | 29.00  | −1.40 pp         |
+| piqa                     | acc_norm   |  50.00 | 60.99  | 60.34  | −0.65 pp         |
+| winogrande               | acc        |  50.00 | 51.85  | 52.72  | **+0.87 pp**     |
+| mmlu (aggregate)         | acc        |  25.00 | 22.88  | 23.12  | +0.24 pp         |
+| race                     | acc        |  25.00 | 25.07  | 25.45  | +0.38 pp         |
+| **mean absolute Δ**      |            |        |        |        | **0.58 pp**      |
+
+Every task within ±1.4 pp — consistent with the 0.49% loss delta at the
+end of training. mmlu / race / arc_challenge are at random-chance for
+*both* stacks (300 M params + 10 B tokens is below the threshold those
+benchmarks need to lift above noise).
+
+---
+
+## Configs and tools used
+
+```
+docs/04-technical-guides/hybrid-models/
+├── kda-guide.md                                  ← this file
+└── kda-fla-parity.md                              ← deep-dive on every change
+examples/megatron/configs/MI300X/
+└── zebra_llama_300M_kda_pure-pretrain.yaml        ← training config
+primus/configs/models/megatron/
+└── zebra_llama_300M_kda_pure.yaml                 ← architecture-only config
+primus/backends/megatron/core/models/hybrid/
+├── kimi_delta_attention.py                        ← FLA-aligned mixer (fused in_proj, FLA Triton paths)
+├── kimi_delta_attention_layer.py                  ← eps propagation, optional pre-norm
+└── hybrid_mamba_mla_layer_specs.py                ← kda_hybrid_stack_spec_no_te
+primus/backends/megatron/patches/                  ← same 6 patches as GDN (Primus patch system, shared)
+├── gdn_config_patches.py                          ← registers use_fla_triton_kda + fusion flags + hybrid init
+├── mamba_fused_ce_patches.py                      ← FLA fused cross-entropy for MambaModel
+├── torch_fused_adam_patches.py                    ← PRIMUS_TORCH_OPTIM opt-in
+├── mlp_fla_swiglu_patches.py                      ← FLA Triton SwiGLU for MLP
+├── torch_norm_fla_rmsnorm_patches.py              ← FLA RMSNorm for WrappedTorchNorm
+├── fla_runtime_patches.py                         ← resolves PRIMUS_FLA_* knobs onto args
+└── mamba_fla_data_patches.py                      ← FLA-order dataset shim wiring
+tools/hybrid/
+├── patch_fla_triton_autotune_hang.sh              ← MI300X FLA Triton autotune-hang workaround
+├── convert_fla_to_megatron.py                     ← FLA Arrow → Megatron .bin/.idx (shared)
+├── fla_order_dataset.py                           ← FLA-order dataset shim (shared)
+├── convert_fla_kda_init_to_megatron.py            ← FLA HF init → Megatron sharded ckpt
+├── convert_kda_to_fla_hf.py                       ← Megatron sharded ckpt → FLA HF
+└── eval_kda_lm_eval.py                            ← lm-eval wrapper (registers KDA)
+```
+
+---
+
+## Troubleshooting
+
+### `KeyError: 'kda'` at `AutoModelForCausalLM.from_pretrained`
+
+You imported `transformers` before `fla` (or didn't import `fla` at all).
+`fla/models/kda/__init__.py` runs
+`AutoConfig.register(KDAConfig.model_type, KDAConfig, exist_ok=True)`
+on import. Either:
+
+- Prepend `PYTHONPATH=/home/<user>/flash-linear-attention` and `import fla`
+  in your script BEFORE the `transformers` import, OR
+- `pip install -e /home/<user>/flash-linear-attention` once and forget
+  about `PYTHONPATH`, OR
+- Use the wrapper: `python tools/hybrid/eval_kda_lm_eval.py ...`
+
+### Conversion: `KeyError: 'decoder.layers.0.mixer.in_proj.weight'`
+
+You trained with an older code branch that still had six separate
+projections. Either re-train with the current fused-in_proj branch or
+patch the converter to read the unfused `q_proj_weight`/`k_proj_weight`/…
+keys (see git history of `tools/hybrid/convert_kda_to_fla_hf.py`).
+
+### Iter 1 loss ~12.05 instead of ~11.97
+
+The `layernorm_epsilon: 1.0e-6` override is being silently overwritten by
+the `TransformerConfig` default of `1e-5`. Confirm it's in the *training*
+YAML's `overrides:` block (not just the model YAML).
+
+### Iter 1 loss not bit-matching FLA but converges fine
+
+You probably didn't load the FLA-init checkpoint (Step 4) or didn't set
+`PRIMUS_FLA_DATA=1`. Without either, the first batch differs (Megatron
+shuffler vs HF `DistributedSampler`) and the per-parameter `nn.init.normal_`
+draw order differs (Megatron traverses Primus's fused `in_proj`, FLA
+traverses 6 separate `nn.Linear` modules). The gap disappears by iter
+~2000 even without either fix.
+
+### Loss is +0.2–0.4 above FLA across the whole run (with FLA-init loaded)
+
+You probably have `use_fla_kda_in_kernel_gate: false` or
+`use_fla_fused_norm_gated: false`. Those toggles select the bit-identical-
+to-old-FLA `fused_kda_gate` + `_apply_gated_norm` paths, which run the
+gate compute in fp32 (slightly different rounding than the in-kernel bf16
+accumulator). Set both to `true` to match the current FLA reference.
+
+### Per-iter time ≫ 1500 ms
+
+Most likely you have `PRIMUS_FLA_CONV=0`. The Tri-Dao `causal_conv1d_fn`
+on ROCm requires `[B, D, T]` layout, so each iteration pays two
+`transpose+contiguous` copies of the (B, qk_dim·2 + v_dim, T) tensor —
+about 35 ms wasted per iter at micro_batch=128. Set `PRIMUS_FLA_CONV=1`
+to switch to FLA's Triton `causal_conv1d` (accepts `[B, T, D]` natively).
+
+### Out-of-memory at iter 1
+
+Two common culprits:
+
+1. `PYTORCH_ALLOC_CONF=expandable_segments:True` is unset — set it.
+2. `q.contiguous()/k.contiguous()/v.contiguous()` removed from KDA forward
+   — the Triton kernel will allocate its own copies while autograd still
+   pins the original views, doubling Q/K/V activation memory. Restore
+   the explicit contiguous calls (see `kimi_delta_attention.py` around
+   the `chunk_kda` call site).
+
+### Eval truncation warnings
+
+Some samples exceed the model's `max_position_embeddings = 2048`. Add
+`max_length=1024` to `--model_args` if it bothers you; it only
+meaningfully affects RACE.
+
+---
+
+## See also
+
+- [`docs/04-technical-guides/hybrid-models/README.md`](README.md) — full Zebra-Llama family
+  overview (1 B / 3 B / 8 B Mamba+MLA, KDA variants)
+- [`docs/04-technical-guides/hybrid-models/gdn-guide.md`](gdn-guide.md) — the GDN companion
+  recipe (shares Megatron patches and dataset shim with this one)
+- [`kda-fla-parity.md`](kda-fla-parity.md) — exhaustive list of
+  code/config/runtime changes that made KDA parity possible
+- FLA upstream: [https://github.com/fla-org/flash-linear-attention](https://github.com/fla-org/flash-linear-attention)
diff --git a/docs/04-technical-guides/logging-and-experiment-tracking.md b/docs/04-technical-guides/logging-and-experiment-tracking.md
new file mode 100644
index 000000000..fd88ca9e6
--- /dev/null
+++ b/docs/04-technical-guides/logging-and-experiment-tracking.md
@@ -0,0 +1,163 @@
+# Logging and experiment tracking
+
+This guide covers how Primus emits training metrics and logs, and how to wire up the supported experiment trackers—**TensorBoard**, **Weights & Biases (WandB)**, and **MLflow** (including Databricks-hosted MLflow)—across the Megatron and TorchTitan backends. Parameters are grounded in `primus/configs/modules/megatron/trainer_base.yaml`, `primus_megatron_module.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`.
+
+For an operations-oriented overview, see [Monitoring and logging](../05-operations/monitoring-logging.md). For required credentials/keys, see [Environment variables](../03-configuration-reference/environment-variables.md).
+
+---
+
+## 1. Tracker toggles at a glance (Megatron)
+
+All three trackers are **opt-in** and disabled by default (`primus/configs/modules/megatron/primus_megatron_module.yaml`):
+
+```yaml
+disable_tensorboard: true
+disable_wandb: true
+disable_mlflow: true
+```
+
+Set the relevant `disable_*` to `false` to enable a tracker. Primus performs validation checks at startup—e.g. it warns if WandB is enabled but `WANDB_API_KEY` is unset (`primus/backends/megatron/patches/args/wandb_config_patches.py`). MLflow logging is initialized in `primus/backends/megatron/training/global_vars.py`; Databricks-hosted MLflow additionally requires `DATABRICKS_HOST` (read by the `mlflow` client).
+
+---
+
+## 2. Console and step logging (Megatron)
+
+Core logging cadence and content (`trainer_base.yaml`, overridden by `pre_trainer.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `log_interval` | `100` (`pre_trainer.yaml` sets `1`) | Steps between log lines. |
+| `log_throughput` | `false` (`pre_trainer.yaml` sets `true`) | Log tokens/s and TFLOP/s throughput. |
+| `log_progress` | `false` | Log progress/ETA. |
+| `log_params_norm` | `false` | Log parameter L2 norm. |
+| `log_num_zeros_in_grad` | `false` | Log gradient sparsity. |
+| `log_avg_skip_iterations` | `2` | Warmup iterations excluded from averages. |
+| `log_avg_reset_interval` | `10` | Reset window for running averages. |
+| `timing_log_level` | `0` | Verbosity of timer breakdowns. |
+| `timing_log_option` | `minmax` | Timer aggregation across ranks. |
+| `logging_level` | `null` | Python logging level override. |
+
+---
+
+## 3. TensorBoard (Megatron)
+
+Enable with `disable_tensorboard: false` and set an output directory:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `tensorboard_dir` | `null` | Output directory for event files (**required** when enabled). |
+| `tensorboard_log_interval` | `1` | Steps between TensorBoard writes. |
+| `tensorboard_queue_size` | `1000` | Event queue size before flush. |
+| `log_learning_rate_to_tensorboard` | `true` | Log LR. |
+| `log_loss_scale_to_tensorboard` | `true` | Log loss scale (mixed precision). |
+| `log_timers_to_tensorboard` | `false` | Log per-stage timers. |
+| `log_batch_size_to_tensorboard` | `false` | Log batch size. |
+| `log_memory_to_tensorboard` | `false` | Log GPU memory. |
+| `log_world_size_to_tensorboard` | `false` | Log world size. |
+| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. |
+
+---
+
+## 4. Weights and biases (Megatron)
+
+Enable with `disable_wandb: false`. Configuration:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `wandb_project` | `null` | WandB project name. |
+| `wandb_exp_name` | `null` | Run/experiment name. |
+| `wandb_save_dir` | `null` | Local directory for WandB files. |
+| `wandb_entity` | `null` | Team/entity. |
+
+**Credentials** (see [Environment variables](../03-configuration-reference/environment-variables.md)):
+
+```bash
+export WANDB_API_KEY=...     # required when WandB is enabled
+export WANDB_PROJECT=...     # optional
+export WANDB_RUN_NAME=...    # optional
+export WANDB_TEAM=...        # optional (entity)
+```
+
+`WANDB_API_KEY` is on the container passthrough allowlist (`runner/.primus.yaml`), so it propagates into the training container.
+
+---
+
+## 5. MLflow (Megatron)
+
+Enable with `disable_mlflow: false`. Run identification and upload behavior:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `mlflow_run_name` | `null` | MLflow run name. |
+| `mlflow_experiment_name` | `null` | MLflow experiment name. |
+| `mlflow_upload_traces` | `false` | Upload profiler trace files. |
+| `mlflow_upload_logs` | `false` | Upload training log files. |
+| `mlflow_upload_performance_metrics` | `false` | Upload the comprehensive perf/memory/utilization metric set (implicitly enables throughput calc). |
+| `mlflow_upload_tracelens_report` | `false` | Generate + upload TraceLens reports (see [Profiling & observability](./profiling-and-observability.md)). |
+
+**Credentials and endpoints:**
+
+```bash
+export MLFLOW_TRACKING_URI=...     # tracking server URI
+export MLFLOW_REGISTRY_URI=...     # optional model registry
+# Databricks-hosted MLflow:
+export DATABRICKS_HOST=...         # checked at startup when MLflow is enabled
+export DATABRICKS_TOKEN=...
+```
+
+> The `mlflow_upload_*` flags are designed so MLflow stays opt-in: they only take effect when `disable_mlflow: false`.
+
+---
+
+## 6. One-logger (Megatron)
+
+NVIDIA One-Logger telemetry is enabled by default in `trainer_base.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `enable_one_logger` | `true` | Enable One-Logger collection. |
+| `one_logger_project` | `megatron-lm` | Project tag. |
+| `one_logger_run_name` | `null` | Run name. |
+| `one_logger_async` | `false` | Async upload. |
+| `app_tag_run_name` / `app_tag_run_version` | `null` / `0.0.0` | Application tags. |
+
+---
+
+## 7. Metrics and logging (TorchTitan)
+
+Configured under `metrics:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `enable_tensorboard` | `false` | Enable TensorBoard logging. |
+| `enable_wandb` | `false` | Enable WandB logging. |
+| `log_freq` | `10` | Steps between metric logs. |
+| `save_tb_folder` | `tb` | TensorBoard output subfolder. |
+| `save_for_all_ranks` | `false` | Write metrics from every rank (default: rank 0 only). |
+| `disable_color_printing` | `false` | Disable ANSI colors in console output. |
+
+TorchTitan reads WandB settings from the environment (`WANDB_PROJECT`, `WANDB_RUN_NAME`, `WANDB_TEAM`) via `primus/backends/torchtitan/patches/wandb_patches.py` and `third_party/torchtitan/torchtitan/components/metrics.py`.
+
+---
+
+## 8. Logging (MaxText)
+
+MaxText logging cadence is controlled by `log_period` (`primus/configs/modules/maxtext/pre_trainer.yaml`, default `100`). See [MaxText parameters](../03-configuration-reference/maxtext-parameters.md).
+
+---
+
+## 9. Recommended setup
+
+1. **Local-only:** enable TensorBoard (`disable_tensorboard: false`, set `tensorboard_dir`)—no credentials required.
+2. **Team tracking:** enable WandB (`disable_wandb: false`) + export `WANDB_API_KEY` and `wandb_project`/`wandb_entity`.
+3. **Enterprise / scaling studies:** enable MLflow (`disable_mlflow: false`) + `MLFLOW_TRACKING_URI` (or Databricks host/token), and turn on `mlflow_upload_performance_metrics` for throughput/memory/utilization dashboards.
+4. **Keep `WANDB_API_KEY` and tokens out of YAML**—pass them as environment variables (allowlisted for container passthrough). See [Security](../05-operations/security.md).
+
+---
+
+## Related documentation
+
+- [Monitoring and logging](../05-operations/monitoring-logging.md)—operational view of trackers.
+- [Profiling & observability](./profiling-and-observability.md)—traces, TraceLens, perf metrics.
+- [Environment variables](../03-configuration-reference/environment-variables.md)—credentials and passthrough.
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md).
diff --git a/docs/04-technical-guides/mega-moe.md b/docs/04-technical-guides/mega-moe.md
new file mode 100644
index 000000000..2b5ec8c5e
--- /dev/null
+++ b/docs/04-technical-guides/mega-moe.md
@@ -0,0 +1,238 @@
+# MegaMoE
+
+MegaMoE is a **FlyDSL**-based fused MoE layer that replaces Megatron's native `MoELayer`. It fuses
+the expert-parallel all-to-all communication into the grouped GEMMs via two fused kernels:
+
+- **dispatch grouped GEMM** (`dispatch_grouped_gemm`): fuses token dispatch (all-to-all) into the
+  L1 grouped GEMM.
+- **grouped GEMM combine** (`grouped_gemm_combine`): fuses the L2 grouped GEMM into combine
+  (all-to-all) + weighted reduce.
+
+Together with a **fused router** (score function + group-limited top-k + aux score) and the
+intermediate SwiGLU, the full expert path is `dispatch_grouped_gemm → SwiGLU → grouped_gemm_combine`;
+the load-balancing aux loss is computed internally and returned. Runtime target is **EP-only
+(TP=1) + bf16**.
+
+## Prerequisites
+
+- **Runtime**: ROCm ≥ 7.0, Python ≥ 3.10, PyTorch ≥ 2.6.0 (ROCm build); gfx950. The DeepEP
+  baseline additionally needs the optional **rocSHMEM**. Image `rocm/primus:v26.3` is recommended.
+- **Primus-Turbo with MegaMoE**: MegaMoE requires Primus-Turbo
+  (`https://github.com/AMD-AGI/Primus-Turbo.git`) at commit
+  **`9b5d3092efcbc087657b233d8e9ae662cee6ec6b` or newer** on `main`. The default image does not ship
+  this kernel, so Primus-Turbo must be rebuilt from source. See upstream
+  [main README](https://github.com/AMD-AGI/Primus-Turbo/blob/main/README.md) and
+  [MegaMoE doc](https://github.com/AMD-AGI/Primus-Turbo/blob/main/docs/README_Mega_MoE.md)
+  for build details. Install it via the rebuild hook:
+
+**Primus rebuild hook (build from source).** The system hook
+`runner/helpers/hooks/00_rebuild_primus_turbo.sh` clones + builds + installs the given ref before
+the training command; each node builds in a node-local dir so multi-node runs avoid shared-fs
+conflicts. Point it at a commit for reproducibility, or at `main` to track the latest code:
+
+```bash
+export REBUILD_PRIMUS_TURBO=1                       # trigger the hook
+export PRIMUS_TURBO_REF=9b5d3092efcbc087657b233d8e9ae662cee6ec6b   # min required commit (or main)
+export GPU_ARCHS="gfx950"                           # build only target arch (multiple: semicolon-separated)
+# Optional: custom build dir (default /tmp/primus_turbo_<hostname>)
+# export PRIMUS_TURBO_BUILD_DIR=/tmp/primus_turbo_build
+```
+
+## Design: two stages + weight modules (for DDP overlap)
+
+The expert path could be exposed as a *single* fused op taking both `w1` and `w2`
+(`fused_mega_moe(x, topk_idx, topk_weights, w1, w2, group)` still exists in Primus-Turbo). Primus
+instead drives it as **two stages**, each owning one weight, wrapped in **two tiny weight modules**
+(`primus/backends/megatron/core/extensions/mega_moe.py`):
+
+```
+MegaMoEExperts
+├── fc1_weight : MegaMoEWeightModule   # w1 [g, 2I, H] gate+up ; forward() -> w1
+└── fc2_weight : MegaMoEWeightModule   # w2 [g, H, I]  down    ; forward() -> w2
+
+
+FORWARD (in order)                          BACKWARD (in order)
+──────────────────────────────────────      ──────────────────────────────────────
+w1 = fc1_weight()                           stage2.backward  ->  dW2
+  hook: all-gather(w1), wait
+                                            hook on fc2_weight
+stage1: dispatch + GEMM1        <──┐          -> reduce-scatter(dW2)  ──┐
+                                   │ overlap                            │ overlap
+w2 = fc2_weight()                  │        stage1.backward  ->  dW1  ──┘
+  hook: all-gather(w2)  ───────────┘
+                                            hook on fc1_weight
+stage2: SwiGLU + GEMM2 + combine              -> reduce-scatter(dW1)
+  (waits for w2 here)                         (overlaps the next layer)
+```
+
+`MegaMoEWeightModule` holds a single `torch.nn.Parameter` and its `forward()` just returns it. It
+computes nothing — its only job is to be a **hook site** for the two collectives Megatron's
+distributed optimizer wants to overlap, both of which are driven at module / parameter granularity:
+
+- **Parameter all-gather** (`overlap_param_gather`) rides the *forward pre-hook*, which is per
+  module. With a single call site taking `(w1, w2)` both gathers must land before any compute; split,
+  `w2`'s gather is issued at `fc2_weight` and overlaps stage1.
+- **Gradient reduce-scatter** (`overlap_grad_reduce`) rides the *grad hook*, which fires when a
+  parameter's `.grad` appears. One fused autograd node emits `dW1` and `dW2` together at the end of
+  the layer backward; split, `dW2` lands early and its reduce-scatter hides under stage1's backward.
+
+The split is purely at the Python/autograd level — the kernels are unchanged.
+
+## Configuration
+
+Enable the fused MegaMoE layer in the training config:
+
+```yaml
+enable_primus_turbo: true
+use_turbo_mega_moe: true   # MegaMoE layer replacement (EP-only / TP=1 / bf16)
+```
+
+The patch is applied only when **all** of these hold: `enable_primus_turbo=True`,
+`use_turbo_mega_moe=True`, `tensor_model_parallel_size==1`, `params_dtype==bf16`, and an EP process
+group exists.
+
+The following model settings are **required** — MegaMoE asserts on anything else:
+
+```yaml
+tensor_model_parallel_size: 1           # EP-only, TP=1
+add_bias_linear: false                  # no bias in linear layers
+# gated SwiGLU + SiLU activation
+```
+
+Unsupported (each raises an error): sequence-level / global aux loss, z-loss, sinkhorn, and input
+jitter (only the standard `aux_loss` is supported); aux-loss-free expert bias
+(`enable_expert_bias=True` raises `NotImplementedError`).
+
+### Router force load balancing (`moe_router_force_load_balancing_type`)
+
+The benchmark config sets `moe_router_force_load_balancing: true`, which discards the real router
+decision and forces every expert to receive a similar number of tokens — this removes run-to-run
+expert-imbalance noise so throughput numbers are comparable. `moe_router_force_load_balancing_type`
+selects *how* the balancing is done (it only has an effect when force load balancing is on):
+
+| value | behavior |
+| --- | --- |
+| `even` (Primus default) | Deterministic round-robin `(token_idx * topk + k) % num_experts`. Per-expert token counts are exactly equal **and identical every step**, so the grouped-GEMM shapes (`M_total`, per-expert `M`) never change. |
+| `uniform` | Megatron-LM's original behavior: the logits are replaced with random values before routing, so token counts are balanced only *statistically* and fluctuate step to step, like real routing. |
+
+**Use `uniform` for benchmarking.** `even` produces perfectly constant, aligned per-expert shapes
+that the non-fused (DeepEP / grouped-GEMM) baseline benefits from disproportionately — no autotune
+or shape-recompile churn, no padding waste — while MegaMoE is designed to absorb ragged, varying
+token counts. Measuring under `even` therefore understates MegaMoE's gain; `uniform` keeps the
+balancing (for reproducibility) but preserves the step-to-step shape variation of real training.
+
+The examples below use the rebuild hook (`REBUILD_PRIMUS_TURBO=1
+PRIMUS_TURBO_REF=9b5d3092efcbc087657b233d8e9ae662cee6ec6b`) to build Primus-Turbo from source.
+
+### Example 1 — single-node EP8, 4 layers (`run_pretrain_cli.sh`)
+
+1 node × 8 GPUs, `TP=1 / PP=1 / EP=8`, DeepSeek-V3 BF16, `GBS = MBS*GPUS*GA = 2*8*64 = 1024`.
+Minimal fused-MegaMoE run from `Primus/`:
+
+```bash
+#!/bin/bash
+set -e
+
+# Model config
+export EXP=examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml
+# Build Primus-Turbo from source before training (hook)
+export REBUILD_PRIMUS_TURBO=1
+export PRIMUS_TURBO_REF=9b5d3092efcbc087657b233d8e9ae662cee6ec6b
+export GPU_ARCHS=gfx950
+
+# Parallelism (EP-only) + fused MegaMoE
+bash examples/run_pretrain_cli.sh \
+  --num_layers 4 \
+  --micro_batch_size 2 \
+  --global_batch_size 1024 \
+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 1 \
+  --expert_model_parallel_size 8 \
+  --moe_layer_freq 1 \
+  --moe_shared_expert_intermediate_size None \
+  --pipeline_model_parallel_layout null \
+  --recompute_granularity null \
+  --recompute_num_layers 0 \
+  --recompute_layer_ids null \
+  --moe_router_force_load_balancing_type uniform \
+  --enable_primus_turbo True \
+  --use_turbo_mega_moe True \
+  --mock_data True
+```
+
+
+### Example 2 — single-node EP8 with full CUDA graph
+
+Same geometry as Example 1, plus TE full-scope CUDA graph capture
+(`--external_cuda_graph True` maps to `cuda_graph_impl="transformer_engine"`; scope `full` is
+normalized to `[]` by Megatron, i.e. capture the whole layer). MegaMoE is fully sync-free — no
+device-to-host sync or CPU-side wait in the expert path — so it captures cleanly and is
+`torch.compile`-friendly.
+
+```bash
+#!/bin/bash
+set -e
+
+export EXP=examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml
+export REBUILD_PRIMUS_TURBO=1
+export PRIMUS_TURBO_REF=9b5d3092efcbc087657b233d8e9ae662cee6ec6b
+export GPU_ARCHS=gfx950
+
+bash examples/run_pretrain_cli.sh \
+  --num_layers 4 \
+  --micro_batch_size 2 \
+  --global_batch_size 1024 \
+  --train_iters 15 \
+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 1 \
+  --expert_model_parallel_size 8 \
+  --moe_layer_freq 1 \
+  --moe_shared_expert_intermediate_size None \
+  --pipeline_model_parallel_layout null \
+  --recompute_granularity null \
+  --recompute_num_layers 0 \
+  --recompute_layer_ids null \
+  --moe_router_force_load_balancing_type uniform \
+  --external_cuda_graph True \
+  --cuda_graph_scope full \
+  --cuda_graph_warmup_steps 3 \
+  --enable_primus_turbo True \
+  --use_turbo_mega_moe True \
+  --mock_data True
+```
+
+### Example 3 — 8-node EP8/PP8
+
+8 nodes × 8 GPUs = 64 GPU, `TP=1 / PP=8 / EP=8` → `DP = 64/(TP*PP) = 8`,
+`GBS = MBS*DP*GA = 2*8*64 = 1024`.
+
+```bash
+set -e
+
+# Cluster geometry (8 nodes x 8 GPUs = 64 GPUs)
+export NNODES=8
+
+export USING_AINIC=1
+
+# Toggle the fused MegaMoE layer + model config
+export EXP=examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml
+
+bash examples/run_slurm_pretrain_cli.sh \
+  --train_iters 15 \
+  --micro_batch_size 2 \
+  --global_batch_size 1024 \
+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 8 \
+  --pipeline_model_parallel_layout "Ett|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttL" \
+  --expert_model_parallel_size 8 \
+  --moe_shared_expert_intermediate_size None \
+  --mtp_num_layers 0 \
+  --recompute_granularity full \
+  --recompute_method uniform \
+  --recompute_num_layers 1 \
+  --recompute_layer_ids null \
+  --moe_router_force_load_balancing_type uniform \
+  --enable_primus_turbo True \
+  --use_turbo_mega_moe True \
+  --mock_data True
+```
diff --git a/docs/04-technical-guides/moe-training.md b/docs/04-technical-guides/moe-training.md
new file mode 100644
index 000000000..b31af0b2c
--- /dev/null
+++ b/docs/04-technical-guides/moe-training.md
@@ -0,0 +1,193 @@
+# MoE training deep-dive
+
+This guide covers Mixture-of-Experts (MoE) training in Primus on AMD Instinct GPUs: the bottlenecks unique to sparse models, the Primus/Primus-Turbo optimizations that address them, and a model-by-model tuning walkthrough. It is adapted from the AMD blog [MoE Training Best Practices on AMD GPU](https://rocm.blogs.amd.com/software-tools-optimization/primus-moe-package/README.html) (`examples/moe_package/README.md`) and grounded in the actual Primus configs and run scripts.
+
+All flags shown here are the **real CLI/YAML keys** used by `examples/moe_package/run_*_pretrain_mi355x.sh` and the Megatron module configs (`primus/configs/modules/megatron/`). The Primus-Turbo MoE optimizations in this guide (DeepEP, sync-free MoE, Turbo grouped GEMM) are **Megatron-backend** features. TorchTitan also supports MoE via expert parallelism (`expert_parallel_degree`, `expert_tensor_parallel_degree`), but its tuning is out of scope here.
+
+---
+
+## 1. Why MoE training is different
+
+MoE scales model capacity by routing each token through a small subset of "expert" sub-networks instead of activating the whole network. A gating/router picks the top-k experts per token, so a model can hold many billions of parameters while only a fraction are active per token.
+
+This sparsity creates performance challenges that dense models do not have:
+
+- **Grouped GEMM overhead**—each expert is a separate GEMM; naive multi-stream execution leaves scheduling gaps.
+- **All-to-all (A2A) communication**—token dispatch/combine across expert-parallel ranks can dominate runtime, especially with `EP >= 8` and multi-node.
+- **CPU sync & launch delays**—dynamic shapes (token counts per expert) force device-to-host syncs that stall the kernel launch queue.
+- **Too many small kernels**—fine-grained MoE ops stress the CPU launch path.
+- **Pipeline load imbalance**—uneven layer distribution across pipeline stages quietly degrades throughput.
+- **Memory pressure**—activations dominate memory at large scale, forcing recomputation.
+
+---
+
+## 2. Representative model configs
+
+Primus ships Megatron model presets for DeepSeek-style MoE models plus two ultra-large research configs:
+
+| Model | Total / Active params | Model config (`primus/configs/models/megatron/`) |
+|-------|-----------------------|--------------------------------------------------|
+| DeepSeek-V2-Lite | 16B / 2.4B | `deepseek_v2_lite.yaml` |
+| DeepSeek-V2 | 236B / 21B | `deepseek_v2.yaml` |
+| DeepSeek-V3 | 671B / 37B | `deepseek_v3.yaml` |
+| MoE-1T | 1T / 44B | `moe_1T.yaml` |
+| MoE-2T | 2T / 80B | `moe_2T.yaml` |
+
+Ready-to-run pretrain scripts live in `examples/moe_package/`, e.g.:
+
+- `examples/moe_package/run_deepseek_v2_lite_pretrain_mi355x.sh`
+- `examples/moe_package/run_deepseek_v2_pretrain_mi355x.sh`
+- `examples/moe_package/run_deepseek_v3_pretrain_mi355x.sh`
+
+Each example script is a convenience wrapper: it sets environment + parallelism and selects an experiment YAML under `examples/moe_package/configs/`, then launches training. You can run the same training directly with the unified CLI, passing the experiment YAML with `--config` and the MoE feature toggles (Section 4) as overrides:
+
+```bash
+# DeepSeek-V2-Lite baseline + DeepEP + sync-free + loss fusion + manual GC, via primus-cli
+export ENABLE_NUMA_BINDING=1 HSA_KERNARG_POOL_SIZE=12582912   # feature 6 (env, not CLI flags)
+./runner/primus-cli direct -- train pretrain \
+  --config examples/moe_package/configs/MI355X/deepseek_v2_lite-pretrain-baseline.yaml \
+  --enable_primus_turbo True \
+  --use_turbo_deepep True --turbo_deepep_num_cu 64 --moe_router_dtype fp32 \
+  --turbo_sync_free_moe_stage 1 \
+  --cross_entropy_fusion_impl te --cross_entropy_loss_fusion True \
+  --manual_gc True --manual_gc_interval 1
+```
+
+Use `./runner/primus-cli slurm srun -N <nodes> -- train pretrain --config ...` for multi-node. The feature tables below list the exact flags so you can compose your own command. (Optimizations that are environment variables—NUMA binding, `HSA_KERNARG_POOL_SIZE`, UCCL-EP—are exported before the command rather than passed as `--flags`.)
+
+---
+
+## 3. Profiling and analysis workflow
+
+Diagnose before optimizing. The recommended order:
+
+1. **Torch Profiler**—capture operator times, memory, and GPU utilization. Enable through the Megatron profiling flags (`--profile`, `--use_pytorch_profiler`, `--profile_step_start`, `--profile_step_end`, `--disable_profiler_activity_cpu`). Load the trace in [Perfetto](https://ui.perfetto.dev/) to inspect CPU/GPU overlap, launch delays, and idle gaps. See [Profiling & observability](./profiling-and-observability.md).
+2. **TraceLens**—AMD's automated trace analyzer for hierarchical breakdowns, roofline/efficiency, communication-vs-sync separation, and trace diffing. Wired into Primus via `generate_tracelens_report` / `mlflow_upload_tracelens_report` (see the profiling guide).
+3. **Memory projection**—model VRAM across params, gradients, activations, and optimizer state *before* launching, via `./primus-cli direct -- projection memory --config <exp>.yaml`. See [Projection](../02-user-guide/projection.md).
+4. **Pipeline visualization**—dump pipeline schedule data (`--dump_pp_data true`) and render stage utilization with `tools/visualization/pp_vis/vis.py` to find bubbles and stage imbalance.
+
+---
+
+## 4. Primus MoE optimizations
+
+The `examples/moe_package/run_*` scripts expose these as composable "MoE features." The table maps each feature to the **actual `--flags` (or environment variables)** you pass to `train pretrain`—the same toggles the example scripts set.
+
+| Feature | Flags (real keys) | What it does |
+|---------|-------------------|--------------|
+| Turbo attention | `--enable_primus_turbo True --use_turbo_attention True` | Optimized attention kernels (Primus-Turbo). |
+| Turbo grouped GEMM | `--enable_primus_turbo True --use_turbo_grouped_gemm True` | Fused CK grouped GEMM processes all experts in one launch instead of multi-stream. |
+| Loss fusion | `--cross_entropy_fusion_impl te --cross_entropy_loss_fusion True` | Fuses large-vocab loss into one kernel to cut memory + launch overhead. |
+| DeepEP acceleration | `--enable_primus_turbo True --use_turbo_deepep True --turbo_deepep_num_cu 64 --turbo_deepep_use_comm_stream False --moe_shared_expert_overlap False --moe_router_dtype fp32` | GPU-side index calc + sync-free dispatch to cut redundant cross-node A2A traffic. |
+| Sync-free MoE | `--enable_primus_turbo True --turbo_sync_free_moe_stage <N>` | Removes CPU D2H syncs across Router → Dispatcher → Permutation → GroupMLP. |
+| NUMA binding | `export ENABLE_NUMA_BINDING=1` | Pins each GPU process to its NUMA socket for better memory bandwidth/stability. |
+| HIP kernarg pool | `export HSA_KERNARG_POOL_SIZE=12582912` | Enlarges the kernel-argument pool (12 MB) to avoid launch stalls under many small kernels. |
+| Manual GC | `--manual_gc True --manual_gc_interval 1` | Periodic host GC to remove iteration-time jitter on long runs. |
+| UCCL-EP | `export USING_UEP=1` | Use the UCCL transport for DeepEP dispatch/combine (sets `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP` + UCCL network env). Requires the `uccl` and `deep_ep` packages. |
+
+> **Grouped GEMM keys.** Enable Turbo grouped GEMM for MoE with `use_turbo_grouped_gemm` (`--use_turbo_grouped_gemm True`). The older `use_turbo_grouped_mlp` alias has been **removed**—passing it now raises an assertion error (`use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead`).
+>
+> **Legacy path.** The legacy multi-stream grouped GEMM path is selected with `--moe_use_legacy_grouped_gemm True` (the scripts default `LEGACY_GG=True`). Turbo grouped GEMM is **incompatible** with the legacy path—set `--moe_use_legacy_grouped_gemm False` whenever `use_turbo_grouped_gemm` is enabled (Primus raises an error otherwise).
+
+### Sync-free MoE stages
+
+`turbo_sync_free_moe_stage` is a single knob with four levels (`0`–`3`, validated in `primus/backends/megatron/patches/args/rocm_arg_validation.py`). Each stage **auto-enables** a set of fusion flags:
+
+| Level | Auto-enabled flags | Behavior |
+|-------|--------------------|----------|
+| `0` (default) | — | Disabled—standard baseline. |
+| `1` | `moe_use_fused_router_with_aux_score`, `moe_permute_fusion` | Sync-free **Router** + **Permutation** fusion. |
+| `2` | stage 1 + `use_turbo_deepep`, `use_turbo_grouped_gemm` | Adds sync-free **DeepEP** dispatch and **Turbo grouped GEMM**. |
+| `3` | stage 2 + `use_turbo_fused_act_with_probs` | Full sync-free pipeline (adds fused activation). Per the blog this **uses significantly more GPU memory**—only enable with headroom. |
+
+Requirements (enforced at startup):
+
+- All stages require `--enable_primus_turbo True`.
+- Stages `2` and `3` require Turbo grouped GEMM and are therefore **incompatible with `--moe_use_legacy_grouped_gemm True`**.
+
+Practical guidance from the run scripts:
+
+- **MI355X:** `--turbo_sync_free_moe_stage 1` (compatible with the default legacy grouped GEMM path, since stage 1 does not enable Turbo grouped GEMM).
+- **MI300X / MI325X:** the example scripts suggest stage `2` (with `--moe_shared_expert_overlap False --moe_router_dtype fp32`). Because stage 2 auto-enables Turbo grouped GEMM, you must also set `--moe_use_legacy_grouped_gemm False`.
+
+### Scheduling and memory features
+
+- **1F1B A2A overlap**—interleaves micro-batch N's expert communication with micro-batch N-1's backward compute on top of interleaved-1F1B pipeline parallelism, hiding A2A behind compute while preserving the bubble rate and roughly the same peak memory.
+- **Arbitrary pipeline partition**—manual stage layout instead of automatic even splits, to balance per-stage memory/compute. Use the Megatron-core `--pipeline_model_parallel_layout` flag (as the DeepSeek-V3 script does) or the Primus `decoder_pipeline_manual_split_list` config key (`primus/configs/modules/megatron/primus_megatron_module.yaml`).
+- **Selective layer recompute**—recompute specific transformer layers with `--recompute_layer_ids 0,1,2,3` (keep `RECOMPUTE_LAYERS=0` so this is the only recompute control), or full block recompute via `--recompute_granularity full --recompute_method block --recompute_num_layers N`.
+- **MoE expert-parallel comm overlap**—`overlap_moe_expert_parallel_comm: true` (`trainer_base.yaml`).
+
+See [Performance tuning](./performance-tuning.md) for the full Primus-Turbo flag reference.
+
+---
+
+## 5. Model-specific tuning
+
+### DeepSeek-V2-Lite (16B / 2.4B, 27 layers)
+
+A compute/memory-efficient variant ideal for high-throughput pretraining. AMD Instinct's large HBM (192 GB on MI300X, 288 GB on MI355X) lets you push **micro-batch size (MBS)** high to maximize throughput.
+
+Recommended optimization stack (matches `run_deepseek_v2_lite_pretrain_mi355x.sh`, where `MoE_Features=(3 4 5 6 7 8)`):
+
+1. Manual GC for stable iteration time.
+2. Loss fusion for the large vocabulary.
+3. DeepEP for A2A.
+4. Sync-free mode (stage 1 on MI355X) to remove D2H syncs.
+5. NUMA binding for CPU affinity (`ENABLE_NUMA_BINDING=1` + `HSA_KERNARG_POOL_SIZE`).
+6. MBS scaling using the memory freed by the above (the blog reports peak memory dropping from ~99.8% to ~84.3% at MBS=12, enabling MBS=14).
+
+The script's default `MoE_Features=(3 4 5 6 7 8)` also enables feature `8` = UCCL-EP (see the feature table above). Default parallelism in the script: `TP=1 ETP=1 PP=1 EP=8 CP=1`, `MBS=14 GBS=896 SEQ=4096`.
+
+### DeepSeek-V2 (236B / 21B, 60 layers)
+
+Scale up with parallelism for max throughput across nodes. Recommended stack:
+
+1. Manual GC, 2) Loss fusion, 3) DeepEP, 4) NUMA binding, 5) Sync-free mode, plus **interleaved pipeline parallelism (VPP)** to cut the pipeline bubble ratio. Enabling VPP (`--num_virtual_stages_per_pipeline_rank > 1`) also improves sync-free mode effectiveness.
+
+Default parallelism in the script: `TP=1 ETP=1 PP=4 VPP=5 EP=8 CP=1` (interleaved PP), `SEQ=4096`.
+
+### 1T+ parameter models (MoE-1T / MoE-2T, 96 layers)
+
+Ultra-large training combines every advanced technique. Use **memory projection first**—at this scale **activations dominate memory**, not parameters/optimizer state.
+
+Findings from the blog's projections (768–1024 GPUs):
+
+- **Context parallelism (CP2)** roughly halves activation memory (~76 GB/GPU saved for 1T, ~131 GB/GPU for 2T)—the most effective single lever.
+- **Increasing EP** (8→16) barely reduces memory but adds A2A time.
+- **Increasing PP** (24→48) doesn't materially cut memory and raises pipeline bubbles + activation memory.
+
+Suggested configs:
+
+| Model | MI300X | MI355X |
+|-------|--------|--------|
+| MoE-1T | PP24 EP8 CP2 | PP24 EP8 (no checkpointing) |
+| MoE-2T | PP24 EP16 CP2 | PP24 EP8 CP2 (benefits from larger DP) |
+
+**Pipeline bubble at scale.** When the global batch is constrained, gradient-accumulation (GA) per iteration drops and the bubble ratio rises. Interleaved PP (VPP) mitigates this:
+
+$$
+\text{bubble ratio} = \frac{PP-1}{(PP-1) + GA \times VPP}
+$$
+
+For PP=16, GA=16: VPP=1 gives ~48% bubble; VPP=6 gives ~14%—a large efficiency win verified on a 64-node setup.
+
+**Inter-node dispatch.** Profiling on a 2T/1024-GPU run showed A2A consuming **25–30%** of step time; DeepEP delivered roughly **1.05×–7.66×** end-to-end speedup over plain A2A and kept EP scaling nearly flat.
+
+---
+
+## 6. Quick checklist
+
+1. **Profile first**—Torch Profiler + TraceLens; project memory before launching ultra-large runs.
+2. **Turn on Turbo**—`--enable_primus_turbo True`, then grouped GEMM, attention, DeepEP as needed (requires the external `primus_turbo` package).
+3. **Kill CPU syncs**—`--turbo_sync_free_moe_stage` (1 on MI355X, 2 on MI300X/MI325X).
+4. **Stabilize + bind**—`--manual_gc True`, `ENABLE_NUMA_BINDING=1`, `HSA_KERNARG_POOL_SIZE=12582912`.
+5. **Scale memory headroom into throughput**—raise MBS; use CP2 + selective recompute for 1T+; use VPP to cut pipeline bubbles.
+
+---
+
+## Related documentation
+
+- [Performance tuning](./performance-tuning.md)—Primus-Turbo flags, HipBLASLt, precision, recompute.
+- [Parallelism strategies](./parallelism-strategies.md) and [Parallelism configuration](./parallelism-configuration.md)—EP, PP, CP, VPP.
+- [Collective operations](./collective-operations.md)—A2A and DeepEP context.
+- [Profiling & observability](./profiling-and-observability.md) and [Projection](../02-user-guide/projection.md).
+- Source blog: `examples/moe_package/README.md`.
diff --git a/docs/04-technical-guides/multi-node-networking.md b/docs/04-technical-guides/multi-node-networking.md
new file mode 100644
index 000000000..ffc61d5ab
--- /dev/null
+++ b/docs/04-technical-guides/multi-node-networking.md
@@ -0,0 +1,203 @@
+# Multi-node networking guide
+
+Multi-node training depends on **high-bandwidth, low-latency** communication between GPUs. On AMD systems, **RCCL** (ROCm Collective Communications Library) provides GPU collectives with an API aligned to **NCCL**, so most **NCCL-prefixed** environment variables apply to RCCL as well.
+
+This guide summarizes how Primus configures networking, how **InfiniBand**, **RoCE**, and **AINIC (AMD AI NIC)** fit in, and how to validate and troubleshoot cluster connectivity.
+
+**Primary sources in this repository**
+
+| Topic | File |
+|-------|------|
+| Default NCCL/RCCL and socket setup | `runner/helpers/envs/base_env.sh` |
+| IB HCA detection | `runner/helpers/envs/get_nccl_ib_hca.sh` |
+| Socket / interface detection | `runner/helpers/envs/get_ip_interface.sh` |
+| AINIC hook (container/CLI integration) | `runner/helpers/hooks/03_enable_ainic.sh` |
+| AINIC CLI defaults | `runner/use_ainic.yaml` |
+| ANP / `NCCL_NET_PLUGIN` example | `examples/run_pretrain.sh` |
+
+---
+
+## 1. Overview
+
+- **Goal:** Keep gradient and parameter exchanges from becoming the bottleneck when scaling across nodes.
+- **Stack:** PyTorch distributed uses the ROCm **NCCL** backend name in many configs; the implementation is **RCCL** on AMD GPUs.
+- **Transports:** Common fabrics include **InfiniBand (IB)**, **RoCE** (RDMA over Converged Ethernet), and **AINIC** on supported AMD platforms. Primus scripts set or detect **HCAs**, **socket interfaces**, and optional **AINIC** tuning.
+
+---
+
+## 2. InfiniBand configuration
+
+These variables are standard in NCCL/RCCL deployments. Primus seeds several from `runner/helpers/envs/base_env.sh` when that script is sourced.
+
+| Variable | Role |
+|----------|------|
+| `NCCL_IB_HCA` | Selects **InfiniBand Host Channel Adapters** (device:port list). |
+| `NCCL_IB_GID_INDEX` | **GID index** for the active port (RoCE and IB differ; see vendor docs). |
+| `NCCL_IB_TC` | **Traffic class** for InfiniBand. |
+| `NCCL_IB_FIFO_TC` | Traffic class for FIFO traffic. |
+| `NCCL_IB_RETRY_CNT` | Retry count for IB operations (tune with vendor guidance). |
+| `NCCL_IB_TIMEOUT` | Timeout for IB operations. |
+| `NCCL_IB_QPS_PER_CONNECTION` | Queue pairs per connection. |
+| `NCCL_NET_GDR_LEVEL` | **GPUDirect RDMA** level for NIC/GPU transfers. |
+| `NCCL_DMABUF_ENABLE` | Use **DMA-BUF** path where supported. |
+
+### Auto-detection in Primus
+
+If `NCCL_IB_HCA` is **unset**, `base_env.sh` runs `runner/helpers/envs/get_nccl_ib_hca.sh`, which enumerates `/sys/class/infiniband/`, skips bonded/storage-style devices, and builds a comma-separated `device:port` list for `NCCL_IB_HCA`.
+
+Default in `base_env.sh`:
+
+```bash
+export NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX:-3}
+```
+
+AINIC-oriented configs often override `NCCL_IB_GID_INDEX` to `1` (see `runner/use_ainic.yaml` and `03_enable_ainic.sh`).
+
+---
+
+## 3. RoCE (RDMA over converged ethernet)
+
+RoCE reuses much of the **IB verb** stack; the same **`NCCL_IB_*`** knobs apply.
+
+| Variable | Typical use |
+|----------|-------------|
+| `NCCL_IB_ROCE_VERSION_NUM` | RoCE version (commonly **2** for RoCE v2). |
+
+GID selection (`NCCL_IB_GID_INDEX`) and traffic classes (`NCCL_IB_TC`, `NCCL_IB_FIFO_TC`) remain important on RoCE fabrics. Follow your network team’s mapping (often **GID index 1** for RoCE v2 vs **3** for some IB fabrics—your site might differ).
+
+---
+
+## 4. AINIC (AMD AI NIC)
+
+**AINIC** refers to AMD’s AI-optimized NIC path (for example, the **AMD Pensando™ Pollara 400 AI NIC**) used in some clusters. Enabling it is a combination of **environment**, **container image**, and **device pass-through**.
+
+### Enable AINIC
+
+- Set **`USING_AINIC=1`**. The hook `runner/helpers/hooks/03_enable_ainic.sh` runs when this is set and exports AINIC-related variables back to the caller (`env.VAR=VALUE` lines).
+- Use container images built for AINIC when required by your site. Examples in this repository use tags such as `docker.io/tasimage/primus:<version>-ainic` (see `examples/customer_package/` and `.github/workflows/ci.yaml`). Match the image to your ROCm and ANP bundle.
+
+### `runner/use_ainic.yaml`
+
+Primus CLI system defaults for AINIC-oriented runs include:
+
+- Container **`device`** mounts: `/dev/kfd`, `/dev/dri`, `/dev/infiniband` (required for GPU and IB access in the container).
+- Environment entries such as `USING_AINIC=1`, `NCCL_PXN_DISABLE=0`, and `NCCL_IB_GID_INDEX=1`.
+
+Adjust **`NCCL_IB_GID_INDEX`** and **`container.options.image`** to match your cluster; comments in `runner/use_ainic.yaml` call this out explicitly.
+
+### AINIC hook
+
+**`runner/helpers/hooks/03_enable_ainic.sh`** is the supported hook path: it sets ANP/RCCL/MPI home directories, IB QoS, RoCE version, P2P channel counts, GDR flush behavior, `LD_LIBRARY_PATH` (including `libibverbs` and RCCL/ANP/MPI build paths), and related flags. Default `NCCL_IB_FIFO_TC` in the hook is **192**; align this value with your fabric.
+
+### RCCL network plugin (ANP)
+
+For ANP-based networking, clusters often set **`NCCL_NET_PLUGIN`** to **`librccl-anp.so`** when that library is present under `ANP_HOME_DIR`, falling back to `librccl-net.so` otherwise—see the logic in `examples/run_pretrain.sh`. This complements the library paths from `03_enable_ainic.sh`.
+
+### Variables commonly set for AINIC
+
+From `03_enable_ainic.sh` (non-exhaustive):
+
+| Variable | Purpose |
+|----------|---------|
+| `ANP_HOME_DIR`, `RCCL_HOME_DIR`, `MPI_HOME_DIR` | Install roots for ANP, RCCL, and Open MPI. |
+| `NCCL_IB_TC`, `NCCL_IB_FIFO_TC` | Traffic classes for IB/RoCE. |
+| `NCCL_IB_GID_INDEX` | Often **1** for AINIC-oriented configs in Primus examples. |
+| `NCCL_IB_ROCE_VERSION_NUM` | RoCE v2. |
+| `RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING` | Stricter GDR flush ordering (set to `0` in these scripts). |
+| `LD_LIBRARY_PATH` | Prepends `libibverbs`, RCCL, ANP, and MPI library paths. |
+
+---
+
+## 5. Socket configuration
+
+CPU-side and fallback socket traffic uses interface selection:
+
+| Variable | Role |
+|----------|------|
+| `NCCL_SOCKET_IFNAME` | Interface name or pattern for NCCL socket transport (e.g. `eth0`, or `^docker0,lo` to **exclude** virtual interfaces). |
+| `GLOO_SOCKET_IFNAME` | Interface for **Gloo** process groups (CPU barriers and related). |
+
+**Primus behavior:** `base_env.sh` sets `IP_INTERFACE` via `runner/helpers/envs/get_ip_interface.sh` (fallback: first address from `hostname -I`). Both `NCCL_SOCKET_IFNAME` and `GLOO_SOCKET_IFNAME` default to **`IP_INTERFACE`** when unset.
+
+**Requirement:** All nodes must agree on a **reachable** address family and interface choice; mismatched bindings are a frequent source of hangs.
+
+---
+
+## 6. PCIe cross-NIC (PXN)
+
+| Variable | Default in `base_env.sh` | Meaning |
+|----------|--------------------------|---------|
+| `NCCL_PXN_DISABLE` | `1` | **PXN disabled** by default (saves GPU memory per comment in `base_env.sh`). |
+
+When **`NCCL_PXN_DISABLE=0`**, **PCIe cross-NIC** is enabled: GPUs might use NICs attached to **other** PCIe switches, which can improve **multi-rail** bandwidth at the cost of **higher GPU memory** use. `runner/use_ainic.yaml` sets `NCCL_PXN_DISABLE=0` for AINIC-oriented runs.
+
+---
+
+## 7. Network diagnostics
+
+### Preflight
+
+```bash
+primus-cli direct -- preflight --network
+```
+
+For multi-node (Slurm example):
+
+```bash
+primus-cli slurm srun -N 4 -- preflight --host --gpu --network
+```
+
+See `docs/02-user-guide/preflight.md` for flags, output locations (`output/preflight` by default), and interpretation.
+
+Set **`PRIMUS_EXPECT_IB=1`** when InfiniBand is **required** for validation; preflight uses this in `primus/tools/preflight/network/network_standard.py`.
+
+### RCCL benchmark
+
+```bash
+primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M
+```
+
+This exercises collective bandwidth and latency across a message-size sweep. See `docs/02-user-guide/benchmarking.md` and `primus/tools/benchmark/rccl_bench_args.py` for options (dtypes, operations, output files).
+
+### Verbose RCCL logs
+
+```bash
+export NCCL_DEBUG=INFO
+```
+
+Use for short, controlled runs; **TRACE** can be extremely verbose.
+
+---
+
+## 8. Multi-node setup checklist
+
+- **ROCm version** matches across all nodes (driver and container image).
+- **`NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`** (or auto-detected `IP_INTERFACE`) identify the **same logical network** on every node.
+- **InfiniBand or RoCE** is up (`ibstat`, `/dev/infiniband`, kernel modules such as `ib_core` / `mlx5_core` as appropriate).
+- **Firewall** allows ports required by your launcher and collective tests (`MASTER_ADDR` / `MASTER_PORT` reachable).
+- **`MASTER_ADDR`** resolves and is reachable from **all** nodes.
+- **`GPUS_PER_NODE`** matches physical GPUs per node.
+- **Containers** mount `/dev/kfd`, `/dev/dri`, and `/dev/infiniband` when using IB/RoCE/AINIC (see `runner/use_ainic.yaml`).
+
+---
+
+## 9. Troubleshooting network issues
+
+| Symptom | What to check |
+|---------|----------------|
+| **Timeout or hang** at init | `MASTER_ADDR` / `MASTER_PORT`, firewall, VPN, wrong `NCCL_SOCKET_IFNAME`, or inconsistent interface across nodes. |
+| **Slow** collectives | IB vs Ethernet path, `NCCL_NET_GDR_LEVEL`, fabric errors, or contention; compare **`benchmark rccl`** to baseline. |
+| **IB not detected** | `/dev/infiniband` missing, modules not loaded, or wrong container devices. |
+| **Wrong interface** | Restrict with `NCCL_SOCKET_IFNAME=^docker0,lo` (exclude loopback and Docker bridges). |
+| **GID / RoCE issues** | `NCCL_IB_GID_INDEX` vs site documentation; RoCE v2 settings (`NCCL_IB_ROCE_VERSION_NUM`). |
+
+For a consolidated list of `NCCL_*` / `RCCL_*` variables, see `docs/03-configuration-reference/environment-variables.md` and the upstream [RCCL environment variables](https://rocm.docs.amd.com/projects/rccl/en/develop/api-reference/env-variables.html) documentation.
+
+---
+
+## Related documentation
+
+- [Preflight diagnostics](../02-user-guide/preflight.md)
+- [Benchmark suite](../02-user-guide/benchmarking.md)
+- [NCCL/RCCL collective operations](./collective-operations.md)
+- [Environment variables](../03-configuration-reference/environment-variables.md)
diff --git a/docs/README_NATIVE_SFT_LORA_EN.md b/docs/04-technical-guides/native-sft-lora.md
similarity index 95%
rename from docs/README_NATIVE_SFT_LORA_EN.md
rename to docs/04-technical-guides/native-sft-lora.md
index 9ff918858..ec160614b 100644
--- a/docs/README_NATIVE_SFT_LORA_EN.md
+++ b/docs/04-technical-guides/native-sft-lora.md
@@ -1,4 +1,4 @@
-# Primus Native SFT LoRA — Quick Start
+# Primus native SFT LoRA—quick start
 
 > **Branch**: `feat/megatron/support-sft-native` (PR701)
 > **Backend**: Megatron-LM **native** (no Megatron-Bridge runtime dependency)
@@ -38,13 +38,13 @@ Entry point: `primus/backends/megatron/megatron_sft_trainer.py` (`MegatronSFTTra
 
 ---
 
-## 2. Runtime Environment
+## 2. Runtime environment
 
 ### 2.1 Docker container
 
 | Container | Image | Notes |
 |---|---|---|
-| **`sft_primus_0507_native`** | `rocm/primus:v26.2` | Recommended; verified |
+| **`sft_primus_0507_native`** | `rocm/primus:v26.3` | Recommended; verified |
 
 Container mounts (set once at container start):
 
@@ -63,13 +63,13 @@ export EXP_NAME="llama2_70b_native_$(date +%Y%m%d_%H%M%S)"
 ```
 
 Set automatically inside the container by `examples/run_pretrain.sh` (you don't need to touch these):
-- `TRITON_CACHE_DIR`, `MIOPEN_USER_DB_PATH`, `PRIMUS_CACHE_ROOT` — persistent JIT cache
-- `NCCL_*` / `RCCL_*` — communication tuning
-- `HSA_*` / `GPU_MAX_HW_QUEUES` — AMD GPU performance tuning
+- `TRITON_CACHE_DIR`, `MIOPEN_USER_DB_PATH`, `PRIMUS_CACHE_ROOT`—persistent JIT cache
+- `NCCL_*` / `RCCL_*`—communication tuning
+- `HSA_*` / `GPU_MAX_HW_QUEUES`—AMD GPU performance tuning
 
 ---
 
-## 3. Launch Commands (verified)
+## 3. Launch commands (verified)
 
 ### 3.1 BF16 / FP8 (existing yaml configs, ready to run)
 
@@ -130,7 +130,7 @@ The plumbing is already in place:
 
 Hard constraints:
 
-1. **TransformerEngine ≥ 2.7.0.dev0** required (the `rocm/primus:v26.2` image already satisfies this)
+1. **TransformerEngine ≥ 2.7.0.dev0** required (the `rocm/primus:v26.3` image already satisfies this)
 2. **FP4 and FP8 are mutually exclusive**: `args.fp4 and args.fp8` raises in Megatron (`arguments.py:885-887`)
 3. **`fp4_param` must be paired with `fp4`**: enabling `fp4_param` alone raises (`arguments.py:889-891`)
 
@@ -310,7 +310,7 @@ modules:
       # =====================================================================
       enable_primus_turbo: true
       use_turbo_attention: false
-      use_turbo_grouped_mlp: false
+      use_turbo_grouped_gemm: false
       use_turbo_rms_norm: false
 
       # ---------- Cross-entropy fusion ----
@@ -403,11 +403,11 @@ grep -E "throughput per GPU" "$RANK0" | head -5
 
 ### Q1: `--fp4-format requires Transformer Engine >= 2.7.0.dev0`
 
-Upgrade TE inside the container, or switch to image `rocm/primus:v26.2`+.
+Upgrade TE inside the container, or switch to image `rocm/primus:v26.3`+.
 
 ### Q2: `--fp4-format and --fp8-format cannot be used simultaneously`
 
-Leftover `fp8: hybrid` / `fp8: e4m3` in the yaml — must be removed.
+Leftover `fp8: hybrid` / `fp8: e4m3` in the yaml—must be removed.
 
 ### Q3: `--fp4-param-gather must be used together with --fp4-format`
 
@@ -437,9 +437,10 @@ done
 
 ---
 
-## 6. References / Further Reading
+## 6. References / further reading
 
-- **PR #701** — Full implementation of this native SFT stack:
+- **Post-training overview**: [Post-Training (SFT / LoRA / DPO)](../02-user-guide/posttraining.md)—how this native SFT LoRA path fits into the broader fine-tuning workflow.
+- **PR #701**—Full implementation of this native SFT stack:
   https://github.com/AMD-AGI/Primus/pull/701
 - **Megatron-LM FP4 design**:
   `third_party/Megatron-LM/megatron/core/fp4_utils.py` +
@@ -454,6 +455,6 @@ done
 
 ## 7. Maintainers
 
-- @wenxie-amd  — PR #701 main author
-- @Xiaoming-AMD — co-author (trainer + dataset core)
-- @botaohu001  — packing / mlperf-aligned recipe / diagnostic tools
+- @wenxie-amd—PR #701 main author
+- @Xiaoming-AMD—co-author (trainer + dataset core)
+- @botaohu001—packing / mlperf-aligned recipe / diagnostic tools
diff --git a/docs/04-technical-guides/parallelism-configuration.md b/docs/04-technical-guides/parallelism-configuration.md
new file mode 100644
index 000000000..0bd455360
--- /dev/null
+++ b/docs/04-technical-guides/parallelism-configuration.md
@@ -0,0 +1,255 @@
+# Parallelism configuration guide
+
+Primus is a YAML-driven training framework for AMD GPUs. Megatron-LM, TorchTitan, and MaxText each expose parallelism through different configuration namespaces. This guide explains how to set parallelism and batch-related parameters, how global batch size relates to micro batch size and data parallel width, and how to choose a parallel strategy for common model sizes.
+
+Default values cited below come from Primus module presets:
+
+- Megatron trainer: `primus/configs/modules/megatron/trainer_base.yaml`
+- Megatron model (tensor/pipeline/expert/context parallel): `primus/configs/models/megatron/language_model.yaml`
+- TorchTitan: `primus/configs/modules/torchtitan/pre_trainer.yaml`
+
+Experiment YAMLs in `examples/` often override these defaults for specific models and hardware.
+
+---
+
+## 1. Megatron parallelism configuration
+
+Model-parallel degrees live on the **model** config (for example under `model:` in your experiment YAML, merged from `language_model.yaml`). Training batch and overlap settings live on the **trainer** module (`trainer_base.yaml`).
+
+### Core parallel degrees
+
+| Parameter | Default (Primus `language_model.yaml`) | Description |
+|-----------|------------------------------------------|-------------|
+| `tensor_model_parallel_size` | `1` | Tensor parallelism (TP): shards attention and MLP across this many GPUs. |
+| `pipeline_model_parallel_size` | `1` | Pipeline parallelism (PP): number of pipeline stages. |
+| `expert_model_parallel_size` | `1` | Expert parallelism (EP) for MoE: shards experts across this many GPUs. |
+| `context_parallel_size` | `1` | Context parallelism (CP) for long sequences. |
+| `sequence_parallel` | `true` | Sequence parallelism (SP); typically used with TP greater than 1. |
+
+### Virtual pipeline (VPP) and pipeline communication
+
+| Parameter | Description |
+|-----------|-------------|
+| `virtual_pipeline_model_parallel_size` | Interleaved pipeline depth (null disables VPP). |
+| `num_layers_per_virtual_pipeline_stage` | Layers per virtual stage when using VPP. |
+| `overlap_p2p_comm` | Overlap pipeline P2P with compute (default `true` in `trainer_base.yaml`). |
+
+### Optimizer, FSDP, and overlap (trainer module)
+
+| Parameter | Default (`trainer_base.yaml`) | Description |
+|-----------|-------------------------------|-------------|
+| `use_distributed_optimizer` | `false` | ZeRO-1 style optimizer state sharding when enabled. |
+| `use_torch_fsdp2` | `false` | Full FSDP2 integration. |
+| `overlap_grad_reduce` | `false` | Overlap gradient all-reduce with backward. |
+| `overlap_param_gather` | `false` | Overlap parameter gathering with forward. |
+
+Set these to `true` in your experiment when you want communication/compute overlap; many production configs enable `use_distributed_optimizer` and overlap flags for large runs.
+
+### Data parallel size (implicit)
+
+For Megatron, data parallel size is not a single YAML key; it is implied by the world size and the product of parallel degrees:
+
+\[
+\text{DP} = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}}
+\]
+
+(Adjust if you also use context parallelism or other groupings; your job’s process layout must match the configured degrees.)
+
+### Batch parameters
+
+| Parameter | Default (`trainer_base.yaml`) | Description |
+|-----------|-------------------------------|-------------|
+| `micro_batch_size` | `2` | Micro batch size per data-parallel rank (MBS). |
+| `global_batch_size` | `128` | Target global batch size (GBS) across the data parallel group. |
+
+Megatron derives **gradient accumulation** from `global_batch_size`, `micro_batch_size`, and the effective data parallel size so that:
+
+\[
+\text{GBS} = \text{MBS} \times \text{DP} \times \text{gradient\_accumulation\_steps}
+\]
+
+Equivalently:
+
+\[
+\text{gradient\_accumulation\_steps} = \frac{\text{GBS}}{\text{MBS} \times \text{DP}}
+\]
+
+You normally set `global_batch_size` and `micro_batch_size` in YAML; Megatron computes the number of accumulation steps automatically.
+
+---
+
+## 2. TorchTitan parallelism configuration
+
+TorchTitan parallelism is grouped under the `parallelism:` key in the TorchTitan module (see `primus/configs/modules/torchtitan/pre_trainer.yaml`).
+
+### `parallelism.*` parameters
+
+| Key | Default | Description |
+|-----|---------|-------------|
+| `parallelism.tensor_parallel_degree` | `1` | Tensor parallelism degree. |
+| `parallelism.pipeline_parallel_degree` | `1` | Pipeline parallelism degree. |
+| `parallelism.data_parallel_shard_degree` | `-1` | FSDP shard degree; `-1` lets the framework choose. |
+| `parallelism.data_parallel_replicate_degree` | `1` | DDP-style replication degree. |
+| `parallelism.expert_parallel_degree` | `1` | Expert parallelism for MoE. |
+| `parallelism.context_parallel_degree` | `1` | Context parallelism. |
+| `parallelism.fsdp_reshard_after_forward` | `default` | FSDP reshard policy (`default` uses TorchTitan’s default behavior). |
+| `parallelism.enable_async_tensor_parallel` | `false` | Async tensor-parallel communication. |
+| `parallelism.pipeline_parallel_schedule` | `1F1B` | Pipeline schedule (for example `1F1B`). |
+| `parallelism.pipeline_parallel_microbatch_size` | `1` | Microbatch size for pipeline stages. |
+
+### Batch parameters under `training.*`
+
+| Key | Default | Description |
+|-----|---------|-------------|
+| `training.global_batch_size` | `-1` | Global batch size; `-1` typically means unset or derived. |
+| `training.local_batch_size` | `8` | Per-rank local (micro) batch size. |
+
+### Global batch relationship
+
+For TorchTitan, a useful relationship when using replicate and shard degrees explicitly is:
+
+\[
+\text{global\_batch\_size} \approx \text{local\_batch\_size} \times \text{data\_parallel\_replicate\_degree} \times \text{data\_parallel\_shard\_degree}
+\]
+
+Exact semantics follow TorchTitan’s distributed layout; set `training.global_batch_size` and parallelism degrees consistently with your launcher’s world size.
+
+---
+
+## 3. MaxText parallelism configuration
+
+MaxText (JAX) uses a **device mesh** with **ICI** (intra-node / “in-cluster interconnect”) and **DCN** (inter-node / “data center network”) axes for parallelism. Defaults and parameter names come from upstream MaxText, for example `third_party/maxtext/src/MaxText/configs/base.yml`, not from Primus presets alone.
+
+### Common parallelism keys (from `base.yml`)
+
+Examples include:
+
+- `ici_tensor_parallelism`—tensor parallelism within a node
+- `ici_fsdp_parallelism`—FSDP-style sharding on ICI (default `-1` for auto in many layouts)
+- `dcn_data_parallelism`—data parallelism across nodes (default `-1` for auto)
+- `dcn_fsdp_parallelism`—FSDP across DCN
+
+### Batch sizing
+
+- `per_device_batch_size`—primary knob for per-device batch (see `base.yml`).
+
+Consult MaxText’s mesh documentation and your chosen model YAML for valid combinations of ICI/DCN axes.
+
+---
+
+## 4. Batch size relationships
+
+### Megatron-style identity
+
+\[
+\text{GBS} = \text{MBS} \times \text{DP} \times \text{grad\_accum}
+\]
+
+\[
+\text{DP} = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}}
+\]
+
+(Subject to your exact parallel groups; CP and custom layouts can introduce additional groups.)
+
+### How GBS, MBS, and DP interact
+
+| Goal | What to change |
+|------|----------------|
+| Increase global batch without more per-GPU memory | Increase `gradient_accumulation_steps` (Megatron) or increase accumulation / GBS while keeping MBS fixed. |
+| Increase throughput per step | Increase `micro_batch_size` if memory allows; might require lowering accumulation to keep GBS fixed. |
+| Scale to more GPUs | Increase world size; often increase DP; keep GBS stable by adjusting accumulation. |
+
+### Memory and convergence
+
+| Factor | Effect |
+|--------|--------|
+| **MBS** | Strongly affects per-GPU activation memory; larger MBS often improves GPU utilization but can OOM. |
+| **GBS** | Affects effective noise in the gradient and optimal learning rate scaling; many recipes scale LR with GBS. |
+
+**Practical recommendation:** start with `micro_batch_size` of `1` or `2`, verify stability and memory. Increase `global_batch_size` (via accumulation or more DP ranks) gradually while monitoring loss and adjusting learning rate per your recipe.
+
+### Example numeric table (Megatron-style)
+
+Assume TP=1, PP=1, EP=1, so DP equals world size.
+
+| World size (DP) | MBS | Grad accum | GBS |
+|-----------------|-----|------------|-----|
+| 8 | 1 | 16 | 128 |
+| 8 | 2 | 8 | 128 |
+| 16 | 1 | 8 | 128 |
+| 16 | 2 | 4 | 128 |
+
+---
+
+## 5. Decision guide: Choosing parallelism
+
+| Situation | Suggested direction |
+|-----------|----------------------|
+| Model fits on **one GPU** | Use DP and/or FSDP only; TP=1, PP=1. |
+| Model fits on **one node** but not one GPU | **TP** within the node; **DP** across any remaining replicas. |
+| Model needs **multiple nodes** | **TP** within node where possible; **PP** across nodes for very large depth; **DP** for remaining width. |
+| **MoE** | Add **EP**; align expert count and routing with `expert_model_parallel_size` / `parallelism.expert_parallel_degree`. |
+| **Very long sequences** | Increase **CP** (`context_parallel_size` / `context_parallel_degree`) as supported by the backend. |
+
+### Example configurations (illustrative)
+
+These are representative topologies; always validate with your checkpoint format, memory profile, and hardware interconnect.
+
+| Profile | GPUs | TP | PP | EP | DP (illustrative) |
+|---------|------|----|----|----|---------------------|
+| ~7B | 8 | 1 | 1 | 1 | 8 |
+| ~70B | 64 | 8 | 2 | 1 | 4 |
+| Large MoE (~671B class) | many | 8 | 4 | 8 | remainder |
+
+---
+
+## 6. Common parallelism recipes (YAML snippets)
+
+### Megatron: 8-GPU data parallel only
+
+```yaml
+# model (or merged language_model section)
+tensor_model_parallel_size: 1
+pipeline_model_parallel_size: 1
+expert_model_parallel_size: 1
+context_parallel_size: 1
+sequence_parallel: false
+
+# trainer
+micro_batch_size: 2
+global_batch_size: 128
+```
+
+### Megatron: Tensor + pipeline + data parallel
+
+```yaml
+tensor_model_parallel_size: 8
+pipeline_model_parallel_size: 2
+expert_model_parallel_size: 1
+context_parallel_size: 1
+sequence_parallel: true
+
+micro_batch_size: 1
+global_batch_size: 512
+```
+
+### TorchTitan: TP + PP with explicit schedule
+
+```yaml
+parallelism:
+  tensor_parallel_degree: 4
+  pipeline_parallel_degree: 2
+  data_parallel_replicate_degree: 1
+  data_parallel_shard_degree: -1
+  expert_parallel_degree: 1
+  context_parallel_degree: 1
+  pipeline_parallel_schedule: 1F1B
+  pipeline_parallel_microbatch_size: 1
+  enable_async_tensor_parallel: false
+
+training:
+  global_batch_size: 256
+  local_batch_size: 4
+```
+
+For full worked examples, see `examples/megatron/configs/` and `examples/torchtitan/configs/` under your target hardware (for example `MI300X/`).
diff --git a/docs/04-technical-guides/parallelism-strategies.md b/docs/04-technical-guides/parallelism-strategies.md
new file mode 100644
index 000000000..45d61fe37
--- /dev/null
+++ b/docs/04-technical-guides/parallelism-strategies.md
@@ -0,0 +1,361 @@
+# Parallelism strategies for distributed training
+
+This guide explains the parallelism dimensions used when training large foundation models on AMD GPUs with Primus. It moves from basic data parallelism to advanced combinations of tensor, pipeline, context, and expert parallelism, including how Primus exposes these options through Megatron-LM and TorchTitan.
+
+For Megatron YAML flags and environment tuning, see [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [Environment variables](../03-configuration-reference/environment-variables.md).
+
+---
+
+## 1. Introduction
+
+### Why parallelism is needed
+
+Modern foundation models often exceed the memory of a single GPU: parameters, activations, optimizer states, and KV caches cannot all reside on one device at useful batch sizes. Even when a model *fits*, training throughput might be too low without scaling across many GPUs. Parallelism splits the problem along several independent **dimensions** so that:
+
+- **Memory** is shared across devices (sharding, pipeline stages, sequence splits).
+- **Compute** is scaled by processing more data in parallel or by overlapping communication with computation.
+
+### Overview of parallelism dimensions
+
+| Dimension | What is split | Primary goal |
+|-----------|----------------|--------------|
+| **Data parallelism (DP)** | Input batches | Throughput; same model on each GPU |
+| **FSDP / ZeRO** | Parameters, gradients, optimizer (by stage) | Memory; keep DP semantics |
+| **Tensor parallelism (TP)** | Individual weight matrices and matmuls | Memory per layer; needs fast links |
+| **Sequence parallelism (SP)** | Sequence in non-TP regions | Activation memory with TP |
+| **Pipeline parallelism (PP)** | Layer groups across stages | Memory; depth-wise split |
+| **Context parallelism (CP)** | Sequence for attention (e.g. ring) | Very long contexts |
+| **Expert parallelism (EP)** | MoE experts across devices | Memory and compute for MoE |
+
+These can be **combined**. The product of parallel degrees must match how processes are laid out on the cluster (see [Section 9](#9-combining-parallelism-strategies)).
+
+---
+
+## 2. Data parallelism (DP)
+
+In **classic data parallelism**, every GPU holds a **full copy** of the model. Each rank receives a **different mini-batch** of data. After the backward pass, **gradients are synchronized** so that all ranks apply the same update.
+
+```
+  Batch shard 0     Batch shard 1     Batch shard 2     Batch shard 3
+       |                  |                  |                  |
+       v                  v                  v                  v
+   +--------+        +--------+        +--------+        +--------+
+   | GPU 0  |        | GPU 1  |        | GPU 2  |        | GPU 3  |
+   | full   |        | full   |        | full   |        | full   |
+   | model  |        | model  |        | model  |        | model  |
+   +--------+        +--------+        +--------+        +--------+
+       |                  |                  |                  |
+       +------------------+------------------+------------------+
+                          |
+                    AllReduce(gradients)
+                          |
+                          v
+              Same weights on all ranks after optimizer step
+```
+
+**Properties**
+
+- Simple to reason about and widely supported.
+- Requires the **full model, activations for one micro-batch, and optimizer state** to fit in **one GPU’s memory** (unless combined with other strategies).
+
+**Effective batch size**
+
+For a single update that aggregates over data-parallel ranks and gradient accumulation:
+
+\[
+\text{effective\_batch\_size} = \text{micro\_batch\_size} \times \text{num\_GPUs}_{\text{DP}} \times \text{gradient\_accumulation\_steps}
+\]
+
+Here `num_GPUs_DP` is the **data-parallel group size** (not always the same as `world_size` when TP/PP/EP are also used).
+
+---
+
+## 3. Fully sharded data parallel (FSDP / ZeRO)
+
+**ZeRO** (Zero Redundancy Optimizer) reduces redundant storage by **sharding** optimizer states, gradients, and/or parameters across data-parallel ranks.
+
+| ZeRO stage | Sharded | Idea |
+|------------|---------|------|
+| **Stage 1** | Optimizer states | Each rank keeps only \(1/N\) of optimizer tensors |
+| **Stage 2** | + Gradients | Gradients are sharded; reduced where needed |
+| **Stage 3** | + Parameters | Each rank holds \(1/N\) of parameters; gather before use |
+
+**FSDP** (Fully Sharded Data Parallel) in PyTorch is the common **implementation** of sharded data parallel training; in the Megatron ecosystem, **ZeRO-3-style** behavior is often discussed alongside **FSDP** for full parameter sharding.
+
+**Typical execution pattern (conceptual)**
+
+1. **Forward:** **AllGather** (or equivalent) to materialize parameters needed for the current layer/batch on each rank.
+2. **Backward:** **ReduceScatter** (or equivalent) to write shard-sized gradient pieces back to ranks.
+
+**Memory intuition**
+
+If replicated training used \(M\) memory per rank for parameters+gradients+optimizer, **ideal** full sharding across \(N\) ranks approaches **\(M/N\)** for the sharded pieces (plus buffers and fragmentation). Moving from **full replication** to **\(1/N\)** sharding for those tensors saves roughly **\((N-1)/N\)** of that component—for **8 GPUs**, about **87.5%** of the replicated footprint for the sharded tensors.
+
+### In Primus
+
+| Backend | Configuration |
+|---------|----------------|
+| **Megatron-LM** | `use_distributed_optimizer: true` enables the **distributed optimizer** (ZeRO-1–style optimizer sharding in Megatron). For **PyTorch FSDP2**, set `use_torch_fsdp2: true` (see Megatron constraints: FSDP2 and distributed optimizer are not used together). |
+| **TorchTitan** | `data_parallel_shard_degree` controls how ranks participate in **FSDP-style** sharding (see TorchTitan job config; `-1` often means auto). |
+
+Exact interactions with checkpoint formats and DDP are documented in [Megatron parameters](../03-configuration-reference/megatron-parameters.md).
+
+---
+
+## 4. Tensor parallelism (TP)
+
+**Tensor parallelism** splits **individual layers** (usually linear / attention projections) across GPUs so **no single GPU stores the full weight matrix** for that layer.
+
+### Column-parallel vs row-parallel
+
+Consider a linear layer \(Y = X W\) with weight matrix \(W\). **Column-parallel** splits \(W\) **along the output dimension** (columns). **Row-parallel** splits \(W\) **along the input dimension** (rows) and **splits \(X\)** so each rank’s matmul dimensions match.
+
+**Column-parallel linear**—each rank holds **disjoint columns** of \(W\); each rank's output is a **disjoint column shard** (half the width for 2-way TP). To recover the full-width tensor the shards are **concatenated (All-Gather along the output dim)**—this is only done when the full tensor is actually needed:
+
+```
+                    SAME full X replicated on each TP rank
+                              |
+            +-----------------+-----------------+
+            |                                   |
+            v                                   v
+      Rank 0: X @ W[:,0:h/2]              Rank 1: X @ W[:,h/2:h]
+            |                                   |
+            v                                   v
+      partial Y_0                         partial Y_1
+      (narrow)                            (narrow)
+            |                                   |
+            +-----------------+-----------------+
+                              |
+             All-Gather (concatenate) on output dim
+            (only when the full tensor is needed; with
+             gather_output=False the output stays column-
+             sharded and feeds the next layer with no comm)
+                              |
+                              v
+              full-width Y (concatenation of shards)
+```
+
+**Row-parallel linear**—each rank holds **disjoint rows** of \(W\); **input \(X\)** is **split** along the **input feature** dimension so each rank computes part of the reduction:
+
+```
+      Rank 0: X_0 @ W[0:r/2,:] ----+
+                                   +-- AllReduce --> Y
+      Rank 1: X_1 @ W[r/2:r,:] ----+
+      (X split along features)     (partial sums add to full Y)
+```
+
+Typical **transformer block** pattern: **column-parallel** for the first projection—its column-sharded output is fed directly into the next layer **without communication**—then **row-parallel** for the second projection, which performs the single **AllReduce** that reconstructs the full output. Column-parallel itself only communicates when `gather_output=True`.
+
+**Communication**
+
+- Often **AllReduce** of partial outputs, or **ReduceScatter** + **AllGather** sequences depending on implementation and **sequence parallelism** (see next section).
+
+**When to use**
+
+- Best **within a node** (NVLink / high-bandwidth GPU–GPU paths). Multi-node TP is possible but latency-sensitive.
+
+### In Primus
+
+| Backend | Parameter |
+|---------|-----------|
+| Megatron-LM | `tensor_model_parallel_size` |
+| TorchTitan | `parallelism.tensor_parallel_degree` |
+
+---
+
+## 5. Sequence parallelism (SP)
+
+**Sequence parallelism** extends TP by splitting the **sequence dimension** in regions that are not covered by tensor-parallel matmuls—commonly **LayerNorm**, **dropout**, and sometimes **residual** paths—so **activation memory** scales better when **TP > 1**.
+
+**Interaction with TP**
+
+- After a **column-parallel** region, partial activations can be **ReduceScatter**d along the sequence.
+- Before a **row-parallel** region, activations might be **AllGather**d along the sequence.
+
+So SP trades **extra collectives** for **lower per-rank activation footprint** on long sequences.
+
+### In Primus
+
+| Backend | Parameter |
+|---------|-----------|
+| Megatron-LM | `sequence_parallel: true` (used with TP) |
+| TorchTitan | Sequence-parallel behavior is integrated with TP/parallelization pipelines in supported models |
+
+---
+
+## 6. Pipeline parallelism (PP)
+
+**Pipeline parallelism** assigns **disjoint subsets of layers** to **stages** on different devices. Activations (and gradients) move **between stages** with **point-to-point** communication.
+
+```
+  Microbatch 1:  Stage0 -> Stage1 -> Stage2 -> Stage3
+  Microbatch 2:       Stage0 -> Stage1 -> Stage2 -> Stage3
+  ...
+```
+
+### Pipeline bubbles
+
+If a stage waits for input while other stages compute, **idle time** appears (**pipeline bubble**). Schedulers reduce bubbles by overlapping forwards and backwards across microbatches.
+
+**Common schedules**
+
+| Schedule | Idea |
+|----------|------|
+| **1F1B** | One forward, one backward; classic **warmup / steady / cooldown** phases |
+| **1F1B interleaved (VPP)** | **Virtual pipeline** stages: multiple chunks per device to improve utilization |
+| **Zero-bubble (ZB)** | Reorders / splits backward so **forward and backward** hide each other better; might separate **input-gradient** vs **weight-gradient** phases |
+| **V-Schedule / V-Half / V-Min** | Variants reducing bubbles further (names vary by codebase) |
+| **DualPipe** | Bidirectional pipeline scheduling (e.g. DeepSeek-style) to overlap forward/backward paths |
+
+**Bubble rate**
+
+\[
+\text{bubble\_rate} = \frac{\text{idle time}}{\text{total time}}
+\]
+
+Lower is better; large **microbatch counts** and better schedules reduce bubble overhead.
+
+### In Primus (Megatron)
+
+| Parameter | Role |
+|-----------|------|
+| `pipeline_model_parallel_size` | Number of pipeline stages |
+| `patch_zero_bubble` | Enable Primus/Megatron **zero-bubble** pipeline patches |
+| `patch_primus_pipeline` | Use Primus pipeline implementation for schedule logic |
+| `pp_algorithm` | e.g. `1f1b`, `1f1b-interleaved`, `zero-bubble`, `zero-bubble-heuristic`, `zbv-formatted`, `v-half`, `v-min` |
+
+See `primus/configs/modules/megatron/primus_pipeline.yaml` and `zero_bubble.yaml` in the repo for defaults.
+
+### In Primus (TorchTitan)
+
+| Parameter | Role |
+|-----------|------|
+| `parallelism.pipeline_parallel_degree` | Pipeline depth |
+| `parallelism.pipeline_parallel_schedule` | e.g. `1F1B`, `Interleaved1F1B`, `GPipe`, zero-bubble variants where supported |
+
+---
+
+## 7. Context parallelism (CP)
+
+**Context parallelism** splits the **sequence length** across devices for long-context training. A common pattern is **ring attention**: each rank holds a **chunk** of queries/keys/values and participates in a **ring** of message passing so attention covers the full sequence without centralizing all activations on one GPU.
+
+**Use cases**
+
+- Long documents, 32K–128K+ tokens, where **per-layer activation memory** and **attention compute** must be distributed.
+
+### In Primus
+
+| Backend | Parameter |
+|---------|-----------|
+| Megatron-LM | `context_parallel_size` |
+| TorchTitan | `parallelism.context_parallel_degree` |
+
+---
+
+## 8. Expert parallelism (EP)
+
+**Mixture-of-Experts (MoE)** models route each token to a small subset of **experts**. **Expert parallelism** assigns **different experts** to **different GPUs** so expert weights are not duplicated on every device.
+
+**Communication**
+
+- **AllToAll** (or equivalent) is typical: **dispatch** tokens to expert ranks and **combine** expert outputs back.
+
+**Expert tensor parallelism (ETP)**
+
+- Experts can be further **tensor-parallel** within a subset of GPUs, analogous to TP for dense layers.
+
+### In Primus
+
+| Backend | Parameter |
+|---------|-----------|
+| Megatron-LM | `expert_model_parallel_size` |
+| TorchTitan | `parallelism.expert_parallel_degree`, `parallelism.expert_tensor_parallel_degree` |
+
+---
+
+## 9. Combining parallelism strategies
+
+### Common pattern
+
+- **TP within a node** (fast interconnect).
+- **PP across nodes or across groups** when layers do not fit on one device.
+- **DP / FSDP** for scaling batch size and sharding optimizer state or parameters.
+
+### GPU count (simplified)
+
+For dense models (ignoring CP and detailed MoE layout):
+
+\[
+\text{world\_size} \approx \text{TP} \times \text{PP} \times \text{DP}
+\]
+
+For MoE-heavy setups, you often see:
+
+\[
+\text{world\_size} \approx \text{TP} \times \text{PP} \times \text{EP} \times \text{DP}
+\]
+
+**Context parallelism** introduces another multiplicative factor in layouts where CP ranks are part of the global mesh (exact rank ordering is implementation-specific).
+
+### Memory vs communication
+
+- **More TP** → smaller matrices per GPU but **more frequent** collectives within layers.
+- **More PP** → less memory per stage but **pipeline bubbles** and **latency** between stages.
+- **More DP/FSDP** → better throughput scaling if communication is not saturated.
+
+### Example configurations (illustrative)
+
+| Scenario | TP | PP | DP / notes |
+|----------|----|----|------------|
+| ~7B on 8 GPUs | 1 | 1 | 8-way DP (or FSDP) |
+| ~70B on 64 GPUs (8 nodes × 8) | 8 | 2 | 4-way DP |
+| Large MoE (e.g. 671B-class) on 256 GPUs | 8 | 4 | EP 8 (example; real jobs vary widely) |
+
+Always validate against **memory profiling**, **checkpoint sharding**, and **network** on your cluster.
+
+---
+
+## 10. Batch size relationships
+
+Let:
+
+- \(B_{\text{micro}}\) = micro-batch size per forward/backward **per data-parallel rank** (per step inside accumulation),
+- \(D\) = **data parallel size** (ranks that share the same model split for DP),
+- \(G\) = **gradient accumulation** steps,
+- \(B_{\text{global}}\) = **global batch size** across all DP ranks for one optimizer update.
+
+Then:
+
+\[
+B_{\text{global}} = B_{\text{micro}} \times D \times G
+\]
+
+**Data parallel size** from world size (when using TP, PP, EP):
+
+\[
+D = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}}
+\]
+
+(If **context parallelism** is present, the denominator must include **CP** in the same way your trainer defines the mesh.)
+
+Solve for accumulation:
+
+\[
+G = \frac{B_{\text{global}}}{B_{\text{micro}} \times D}
+\]
+
+**Practical notes**
+
+- **Micro batch** drives **per-GPU activation memory** (often linearly in sequence length for attention).
+- **Global batch** affects **convergence** and learning dynamics; scaling laws often refer to global batch.
+- **Gradient accumulation** increases **time per optimizer step** but **reduces memory** by using smaller \(B_{\text{micro}}\).
+
+Megatron-specific names for batch arguments appear in [Megatron parameters](../03-configuration-reference/megatron-parameters.md).
+
+---
+
+## Related documentation
+
+- [NCCL/RCCL collective operations guide](./collective-operations.md)—which collectives each strategy uses.
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md)
+- [Environment variables](../03-configuration-reference/environment-variables.md)
diff --git a/docs/04-technical-guides/performance-tuning.md b/docs/04-technical-guides/performance-tuning.md
new file mode 100644
index 000000000..5505d9a81
--- /dev/null
+++ b/docs/04-technical-guides/performance-tuning.md
@@ -0,0 +1,242 @@
+# Performance tuning guide
+
+This guide covers AMD-focused performance work in Primus: HipBLASLt autotuning for GEMMs, **Primus-Turbo** optional kernels, mixed precision, activation recomputation, communication overlap, memory settings, and MoE-specific flags. It references Primus examples and Megatron module YAMLs.
+
+---
+
+## 1. HipBLASLt autotuning
+
+Transformer Engine and GEMM-heavy training benefit from HipBLASLt kernel selection. Primus integrates a **three-stage** workflow controlled by `PRIMUS_HIPBLASLT_TUNING_STAGE` (see `examples/README.md` and `examples/run_pretrain.sh`).
+
+> **Activate tuning first.** The stage variable is only honored when the master switch `PRIMUS_HIPBLASLT_TUNING=1` is set (and `PRIMUS_DETERMINISTIC` is not `1`). Without `PRIMUS_HIPBLASLT_TUNING=1`, both `run_pretrain.sh` and the CLI hook `runner/helpers/hooks/train/pretrain/prepare_experiment.sh` skip tuning entirely and force `TE_HIPBLASLT_TUNING_RUN_COUNT=0` / `TE_HIPBLASLT_TUNING_ALGO_COUNT=0`. Export `PRIMUS_HIPBLASLT_TUNING=1` alongside the stage in every command below.
+
+### Stage 0 (default)
+
+No tuning:
+
+```bash
+export PRIMUS_HIPBLASLT_TUNING_STAGE=0  # default
+```
+
+### Stage 1: Dump GEMM shapes
+
+Run a **short** training job so shapes are collected during real forward/backward passes. Reduce `train_iters` (or equivalent) for faster shape collection.
+
+```bash
+export PRIMUS_HIPBLASLT_TUNING=1
+export PRIMUS_HIPBLASLT_TUNING_STAGE=1
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+Output layout (from `examples/README.md`):
+
+- `./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_shape`
+
+### Stage 2: Offline tuning
+
+Runs offline tuning from dumped shapes (often 10–30 minutes depending on model and shapes):
+
+```bash
+export PRIMUS_HIPBLASLT_TUNING=1
+export PRIMUS_HIPBLASLT_TUNING_STAGE=2
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+Expected output:
+
+- `./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_tune/tune_hipblas_gemm_results.txt`
+
+### Stage 3: Train with tuned kernels
+
+Point the runtime at the tuned override file:
+
+```bash
+export PRIMUS_HIPBLASLT_TUNING=1
+export PRIMUS_HIPBLASLT_TUNING_STAGE=3
+export HIPBLASLT_TUNING_OVERRIDE_FILE=/path/to/tune_hipblas_gemm_results.txt
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+```
+
+### Related environment variables
+
+| Variable | Role |
+|----------|------|
+| `TE_HIPBLASLT_TUNING_ALGO_COUNT` | Breadth of algorithm search for TE HipBLASLt tuning (see `examples/run_pretrain.sh` defaults). |
+| `TE_HIPBLASLT_TUNING_RUN_COUNT` | Number of benchmark runs per shape during TE tuning. |
+| `TE_HIPBLASLT_TUNING_ALGO_FILE` | Optional algorithm file for TE tuning flows. |
+| `TE_HIPBLASLT_TUNING` | When set, interacts with deterministic mode; avoid conflicting settings with shape dump (see script comments in `examples/run_pretrain.sh`). |
+| `HIPBLASLT_TUNING_OVERRIDE_FILE` | Override file for stage 3 training. |
+
+### Standalone offline tool
+
+For manual HipBLASLt bench workflows, see `examples/offline_tune/offline_tune_gemm.py` and `examples/offline_tune/README.md` (hipblaslt-bench integration and `HIPBLASLT_TUNING_OVERRIDE_FILE` usage).
+
+---
+
+## 2. Primus-Turbo optimization
+
+**Primus-Turbo** is a separate package of optimized AMD GPU kernels used by Primus Megatron and TorchTitan integrations. It is controlled by the master flag `enable_primus_turbo` in Megatron configs (`primus/configs/modules/megatron/primus_turbo.yaml` extends into trainer/model as needed). **You must install the external `primus_turbo` package** for these paths to be available.
+
+### Master flag (Megatron)
+
+```yaml
+enable_primus_turbo: true
+```
+
+### Feature flags (Megatron)
+
+Defaults in `primus/configs/modules/megatron/primus_turbo.yaml` are mostly `false` until enabled.
+
+| Flag | Purpose |
+|------|---------|
+| `use_turbo_attention` | Optimized attention kernels. |
+| `use_turbo_parallel_linear` | Optimized tensor-parallel linear layers. |
+| `use_turbo_grouped_gemm` | Optimized grouped GEMM for MoE. |
+| `use_turbo_grouped_mlp` | Removed—use `use_turbo_grouped_gemm` (passing this key now raises an error). |
+| `use_turbo_rms_norm` | Optimized RMSNorm. |
+| `moe_use_fused_router_with_aux_score` | Fused MoE router (requires Primus-Turbo backend; see [Backend Patch Notes](../06-developer-guide/backend-patch-notes.md)). |
+| `use_turbo_deepep` | DeepEP token dispatcher; set with `enable_primus_turbo: true`. |
+| `turbo_deepep_num_cu` | Compute units for DeepEP (patch notes suggest practices such as 64 or 80 for EP8, 32 for EP16–64). |
+| `turbo_sync_free_moe_stage` | Sync-free MoE stages (`0`–`3`; `0` disables, stage `2` recommended for performance per patch notes). See [MoE training deep-dive](./moe-training.md). |
+| `use_turbo_fused_act_with_probs` | Fused activation with probabilities to reduce redundant work. |
+
+### Feature flags (TorchTitan)
+
+TorchTitan presets include `primus_turbo` in `primus/configs/modules/torchtitan/pre_trainer.yaml`
+
+Example keys:
+
+```yaml
+primus_turbo:
+  enable_primus_turbo: true
+  use_turbo_attention: true
+  use_turbo_async_tp: true
+  use_turbo_float8_linear: true
+  use_turbo_grouped_mm: false
+```
+
+### Documentation
+
+Extended Megatron arguments and Turbo-related behavior are summarized in [Backend Patch Notes](../06-developer-guide/backend-patch-notes.md).
+
+---
+
+## 3. Mixed precision training
+
+### Megatron (`trainer_base.yaml` patterns)
+
+| Setting | Description |
+|---------|-------------|
+| `bf16: true` | BFloat16 training (default `true` in `trainer_base.yaml`). |
+| `fp16: false` | FP16 training (optional). |
+| `fp8` | FP8 recipe control (`null` / recipes such as delayed scaling in upstream Megatron). |
+| `fp8_recipe`, `fp8_margin`, `fp8_interval` | FP8 scaling behavior. |
+| `fp4`, `fp4_recipe` | Experimental FP4 paths. |
+| `first_last_layers_bf16: true` | Keep first/last layers in BF16 for stability (`num_layers_at_start_in_bf16`, `num_layers_at_end_in_bf16` fine-tune). |
+
+### TorchTitan
+
+| Setting | Location |
+|---------|----------|
+| `training.mixed_precision_param: bfloat16` | `pre_trainer.yaml` default |
+| `training.mixed_precision_reduce: float32` | Reduce precision |
+| FP8 / quantization | `quantize.linear.float8.*` in `primus/configs/modules/torchtitan/quantize.yaml` |
+
+### Loss fusion (Megatron model)
+
+From `primus/configs/models/megatron/language_model.yaml`:
+
+- Default: `cross_entropy_loss_fusion: false` with `cross_entropy_fusion_impl: "native"`.
+- To use Transformer Engine fused cross entropy where supported, enable the fusion explicitly and set `cross_entropy_fusion_impl: "te"`.
+
+---
+
+## 4. Activation recomputation
+
+### Megatron
+
+| Parameter | Typical values | Notes |
+|-----------|----------------|-------|
+| `recompute_granularity` | `full`, `selective` | `full` recomputes more; max memory savings. |
+| `recompute_method` | `uniform`, `block` | How recomputation is distributed. |
+| `recompute_num_layers` | integer | Layers to recompute when using selective or uniform strategies. |
+| `recompute_layer_ids` | list or null | Primus extension: **global** layer indices from `0` to `num_layers - 1` (the patch resolves block-local indices to global ids via `layer_offset`). Use with `recompute_granularity: full` and supported recompute methods. |
+
+### TorchTitan
+
+| Parameter | Location |
+|-----------|----------|
+| `activation_checkpoint.mode` | `none` in default `pre_trainer.yaml` |
+| `activation_checkpoint.selective_ac_option` | Selective AC options |
+
+---
+
+## 5. Communication overlap
+
+### Megatron
+
+From `trainer_base.yaml` and model settings:
+
+| Setting | Purpose |
+|---------|---------|
+| `overlap_grad_reduce` | Overlap gradient reduction with backward. |
+| `overlap_param_gather` | Overlap parameter gather with forward. |
+| `overlap_p2p_comm` | Pipeline P2P overlap. |
+| `async_tensor_model_parallel_allreduce` | Async TP all-reduce (model config). |
+
+### TorchTitan
+
+| Setting | Purpose |
+|---------|---------|
+| `parallelism.enable_async_tensor_parallel: true` | Async tensor parallelism. |
+
+### Environment
+
+`CUDA_DEVICE_MAX_CONNECTIONS=1` is commonly required for **correct** overlap behavior in TP/PP stacks (see `docs/03-configuration-reference/environment-variables.md` and Megatron tests). Primus launch scripts or your cluster setup might set this.
+
+---
+
+## 6. Memory optimization
+
+### Megatron
+
+| Parameter | Purpose |
+|-----------|---------|
+| `optimizer_cpu_offload: true` | Offload optimizer state to CPU. |
+| `optimizer_offload_fraction` | Fraction to offload (`1.0` in `trainer_base.yaml`). |
+| `use_distributed_optimizer: true` | Shards optimizer state across DP ranks (when enabled). |
+| `empty_unused_memory_level` | Aggressive emptying of unused memory (`0` default). |
+| `global_batch_size` + `micro_batch_size` | Increase global batch via **gradient accumulation** without increasing per-step activation memory. |
+
+### TorchTitan
+
+| Parameter | Purpose |
+|-----------|---------|
+| `training.enable_cpu_offload` | CPU offload path in `pre_trainer.yaml` |
+
+---
+
+## 7. MoE-specific optimization
+
+### Megatron (model + turbo)
+
+| Setting | Purpose |
+|---------|---------|
+| `moe_permute_fusion: true` | Fuse permutation / unpermutation (`patch-notes.md`). |
+| `moe_use_fused_router_with_aux_score: true` | Fused router + aux loss (Primus-Turbo). |
+| `use_turbo_deepep: true` | DeepEP dispatcher (`enable_primus_turbo` must be true). |
+| `turbo_sync_free_moe_stage: 2` | Recommended stage for sync-free MoE (per patch notes). |
+| `overlap_moe_expert_parallel_comm: true` | Overlap expert parallel communication (`trainer_base.yaml`). |
+
+---
+
+## Quick checklist
+
+1. **GEMMs:** run HipBLASLt stages 1–3 or use `offline_tune_gemm.py` for custom workflows.
+2. **Kernels:** enable `primus_turbo` after installing `primus_turbo`; turn on attention/MoE flags as needed.
+3. **Precision:** BF16 by default; add FP8/FP4 only with recipe testing.
+4. **Memory:** recomputation + distributed optimizer + CPU offload + accumulation before buying more GPUs.
+5. **MoE:** fusion + DeepEP + sync-free stages + EP comm overlap when supported.
diff --git a/docs/04-technical-guides/profiling-and-observability.md b/docs/04-technical-guides/profiling-and-observability.md
new file mode 100644
index 000000000..42560cbbf
--- /dev/null
+++ b/docs/04-technical-guides/profiling-and-observability.md
@@ -0,0 +1,170 @@
+# Profiling and observability
+
+This guide covers how to capture and analyze performance data in Primus: the PyTorch/Kineto profiler, GPU memory snapshots, AMD's TraceLens trace analysis, ROCm memory sampling, memory/performance projection, and pipeline-schedule visualization. Parameters are grounded in `primus/configs/modules/megatron/`, `primus/configs/modules/torchtitan/pre_trainer.yaml`, and `primus/configs/modules/maxtext/pre_trainer.yaml`.
+
+---
+
+## 1. Torch profiler (Megatron)
+
+Megatron training integrates the PyTorch profiler. Defaults live in `primus/configs/modules/megatron/trainer_base.yaml` and `primus_megatron_module.yaml`.
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `profile` | `false` | Master switch for profiling. |
+| `use_pytorch_profiler` | `false` | Use the PyTorch (Kineto) profiler path. |
+| `profile_ranks` | `[0]` | Which global ranks to profile. |
+| `profile_step_start` | `10` | First step to capture. |
+| `profile_step_end` | `12` | Last step to capture. |
+| `disable_profiler_activity_cpu` | `false` | Drop CPU-side activity to shrink traces (GPU-only trace). |
+| `torch_profiler_record_shapes` | `true` | Record tensor shapes per op. |
+| `torch_profiler_with_stack` | `true` | Record Python/C++ stacks (larger traces). |
+| `torch_profiler_use_gzip` | `false` | Gzip the exported trace. |
+
+Enable via CLI overrides on `train pretrain`:
+
+```bash
+./runner/primus-cli direct -- train pretrain \
+  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml \
+  --profile True \
+  --use_pytorch_profiler True \
+  --profile_step_start 5 \
+  --profile_step_end 6 \
+  --disable_profiler_activity_cpu False
+```
+
+Keep the capture window **short** (a few steps after warmup)—traces grow quickly, especially with `with_stack` and CPU activity enabled. Load the resulting trace in [Perfetto](https://ui.perfetto.dev/) to inspect CPU/GPU overlap, kernel launch delays, and idle gaps.
+
+---
+
+## 2. GPU memory profiling (Megatron)
+
+Two complementary mechanisms:
+
+**Memory history snapshot** (`trainer_base.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `record_memory_history` | `false` | Record the CUDA/HIP allocator history for snapshot analysis. |
+| `memory_snapshot_path` | `snapshot.pickle` | Output path for the allocator snapshot. |
+
+Load the pickle with PyTorch's memory visualizer to find fragmentation and peak allocations.
+
+**ROCm memory sampling** (`primus_megatron_module.yaml`):
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `use_rocm_mem_info` | `false` | When `true`, collect ROCm memory info via `rocm-smi` **every** iteration. |
+| `use_rocm_mem_info_iters` | `[1, 2]` | When `use_rocm_mem_info=false`, only sample at these iterations. |
+
+Also relevant: `log_memory_to_tensorboard` (`trainer_base.yaml`) writes memory metrics to TensorBoard.
+
+---
+
+## 3. TraceLens automated trace analysis (Megatron)
+
+[TraceLens](https://github.com/AMD-AGI/TraceLens) turns raw profiler traces into hierarchical breakdowns (roofline/efficiency, compute-vs-memory bound kernels, communication-vs-sync separation, trace diffing). Primus can generate and optionally upload these reports. Configured in `primus/configs/modules/megatron/primus_megatron_module.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `generate_tracelens_report` | `false` | Generate TraceLens reports locally (auto-enabled when upload is on). |
+| `mlflow_upload_tracelens_report` | `false` | Upload reports to MLflow (auto-enables generation, profiling, tensorboard). |
+| `mlflow_tracelens_ranks` | `null` | Ranks to analyze (`null` = all; e.g. `[0, 8]` for one rank/node). |
+| `mlflow_tracelens_output_format` | `xlsx` | `xlsx` (fastest), `csv`, or `all`. |
+| `mlflow_tracelens_cleanup_after_upload` | `false` | Delete local reports after upload to save disk. |
+| `mlflow_tracelens_auto_install` | `true` | Auto-install TraceLens if missing (set `false` to disable). |
+
+Related profiler/log uploads: `mlflow_upload_traces` (upload raw trace files) and `mlflow_upload_logs` (upload training logs). See [Logging & experiment tracking](./logging-and-experiment-tracking.md) for MLflow setup.
+
+---
+
+## 4. Performance metrics to MLflow (Megatron)
+
+`mlflow_upload_performance_metrics: false` (`primus_megatron_module.yaml`) enables a comprehensive scaling-test metric set when turned on (implicitly enabling throughput calculation):
+
+- `perf/throughput_tflops_per_gpu`, `perf/tps_tokens_per_sec_per_gpu`, `perf/iteration_time_ms`
+- `perf/{rocm,hip}_current_mem_gb`, `perf/{rocm,hip}_mem_utilization_pct`
+- `perf/gpu_utilization_pct_rank{N}`, `perf/gpu_utilization_pct_avg`
+
+> GPU utilization collection uses an `all_gather` every `log_interval`, which synchronizes ranks—keep this in mind for throughput-sensitive runs.
+
+---
+
+## 5. Profiling (TorchTitan)
+
+Configured under `profiling:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `enable_profiling` | `false` | Enable the Torch profiler. |
+| `profile_freq` | `10` | Capture every N steps. |
+| `save_traces_folder` | `profile_traces` | Output folder for traces. |
+| `enable_memory_snapshot` | `false` | Capture GPU memory snapshots. |
+| `save_memory_snapshot_folder` | `memory_snapshot` | Output folder for snapshots. |
+
+Communication tracing is configured under `comm:` (`trace_buf_size`, `save_traces_folder: comm_traces`).
+
+---
+
+## 6. Profiling (MaxText)
+
+Configured in `primus/configs/modules/maxtext/pre_trainer.yaml`:
+
+| Parameter | Default | Purpose |
+|-----------|---------|---------|
+| `profiler` | `xplane` | Profiler backend (XPlane traces). |
+| `skip_first_n_steps_for_profiler` | `3` | Warmup steps to skip before capture. |
+| `profiler_steps` | `1` | Number of steps to capture. |
+
+---
+
+## 7. Memory and performance projection
+
+Project resource usage **before** launching, without consuming a full cluster. Exposed through the Primus CLI `projection` subcommand (`primus/cli/subcommands/projection.py`).
+
+```bash
+# Memory projection from an experiment config
+./primus-cli direct -- projection memory --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml
+
+# Performance projection (single-node benchmarking)
+./primus-cli direct -- projection performance --config <exp>.yaml
+
+# Performance projection scaled to N nodes (simulation)
+./primus-cli direct -- projection performance --config <exp>.yaml --target-nodes 4
+```
+
+Memory projection breaks VRAM down across parameters, gradients, activations, optimizer states, and mixed-precision overhead—invaluable for MoE and ultra-large models where activations dominate. See [Projection](../02-user-guide/projection.md) for the full reference.
+
+---
+
+## 8. Pipeline schedule visualization
+
+Diagnose pipeline bubbles and stage imbalance with the built-in tool `tools/visualization/pp_vis/`.
+
+1. Dump per-rank schedule data during training with `--dump_pp_data true` (Megatron flag `dump_pp_data`, `primus_megatron_module.yaml`). Output lands under `output/pp_data/` (`config.json`, `pp_rank_*.json`).
+2. Install the tool deps and run the viewer:
+
+```bash
+pip install -r tools/visualization/pp_vis/requirements.txt
+python tools/visualization/pp_vis/vis.py   # open http://127.0.0.1:8988
+```
+
+Configure `task_list` in `vis.py` to point at your dumped `log_path` and the iterations to render. The tool can also visualize the PP simulator output (see `tools/visualization/pp_vis/README.md`).
+
+---
+
+## 9. Recommended workflow
+
+1. **Project first**—run `projection memory` to confirm the config fits before booking GPUs.
+2. **Capture a short trace**—a few steps after warmup with `profile` + `use_pytorch_profiler`.
+3. **Inspect**—Perfetto for the timeline; TraceLens for automated hierarchical analysis and trace diffs.
+4. **Check memory**—`record_memory_history` / ROCm sampling for fragmentation and peaks.
+5. **Check pipelines**—`dump_pp_data` + `pp_vis` for bubbles and stage imbalance.
+
+---
+
+## Related documentation
+
+- [Logging & experiment tracking](./logging-and-experiment-tracking.md)—WandB / TensorBoard / MLflow setup.
+- [MoE training deep-dive](./moe-training.md)—applying this workflow to sparse models.
+- [Performance tuning](./performance-tuning.md)—what to change once you've found the bottleneck.
+- [Projection](../02-user-guide/projection.md) and [Monitoring and logging](../05-operations/monitoring-logging.md).
diff --git a/docs/05-operations/README.md b/docs/05-operations/README.md
new file mode 100644
index 000000000..6b9c1e016
--- /dev/null
+++ b/docs/05-operations/README.md
@@ -0,0 +1,12 @@
+# Operations
+
+Production deployment and operational guidance.
+
+- [Deployment](deployment.md): container, Slurm, and Kubernetes deployment
+- [Monitoring and logging](monitoring-logging.md): WandB, TensorBoard, MLflow, Primus logging
+- [Troubleshooting](troubleshooting.md): common failures, diagnostics, and fixes
+- [Security](security.md): secrets handling, container security, dependencies
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/05-operations/deployment.md b/docs/05-operations/deployment.md
new file mode 100644
index 000000000..28948b621
--- /dev/null
+++ b/docs/05-operations/deployment.md
@@ -0,0 +1,237 @@
+# Deployment guide
+
+This guide describes how to deploy Primus training across **container**, **direct (bare metal)**, and **Slurm** environments using the unified `primus-cli` launcher. For environment variable semantics, see [Environment variables](../03-configuration-reference/environment-variables.md). For YAML hierarchy and precedence, see [Configuration system](../02-user-guide/configuration-system.md).
+
+---
+
+## 1. Deployment overview
+
+Primus supports three deployment modes:
+
+| Mode | Description | Typical use |
+|------|-------------|-------------|
+| **Container** | Docker/Podman with ROCm-capable GPU devices and capabilities | Recommended default; reproducible images |
+| **Direct** | Runs on the current host (or inside an existing container) | Local debugging, single-node, clusters with ROCm on nodes |
+| **Slurm** | Wraps `srun`/`sbatch` and launches per-node entry scripts | Multi-node clusters with Slurm |
+
+**Container image:** `docker.io/rocm/primus:v26.4` (default in `runner/.primus.yaml`). For clusters using **AINIC**, use `runner/use_ainic.yaml` and tune the image and NCCL-related variables (for example `USING_AINIC`, `NCCL_IB_GID_INDEX`) to match your fabric.
+
+**Prerequisites (baseline):**
+
+- **AMD ROCm** >= 7.0 on the host (or in the image when using containers)
+- **Docker** or **Podman** >= 24.0 when using container mode
+- **AMD Instinct** GPUs and working ROCm stack (`rocm-smi` should report devices)
+
+---
+
+## 2. Container deployment
+
+### 2.1 Pull the image
+
+```bash
+docker pull docker.io/rocm/primus:v26.4
+```
+
+The default `container.options.image` in `runner/.primus.yaml` is `rocm/primus:v26.4` (equivalent to `docker.io/rocm/primus:v26.4` when the registry is omitted).
+
+### 2.2 Required device mounts
+
+System defaults (`runner/.primus.yaml`, `container.options.device`) pass each path as `--device` to the runtime:
+
+| Device | Purpose |
+|--------|---------|
+| `/dev/kfd` | Kernel Fusion Driver (ROCm core) |
+| `/dev/dri` | Direct Rendering Infrastructure (GPU access) |
+| `/dev/infiniband` | InfiniBand character devices (multi-node / RDMA) |
+
+### 2.3 Required capabilities
+
+Defaults (`container.options.cap-add`):
+
+| Capability | Purpose |
+|------------|---------|
+| `SYS_PTRACE` | Debugging and profiling tools |
+| `CAP_SYS_ADMIN` | Administrative operations required by some ROCm/GPU workflows |
+
+### 2.4 Container runtime options
+
+Defaults in `runner/.primus.yaml` include:
+
+| Option | Value |
+|--------|--------|
+| `ipc` | `host` |
+| `network` | `host` |
+| `privileged` | `true` |
+| `security-opt` | `seccomp=unconfined` |
+| `group-add` | `video` |
+
+`primus-cli-container.sh` always mounts the **Primus repository root** into the container at the same path (`-v $PRIMUS_PATH:$PRIMUS_PATH`). Mount additional paths for **datasets**, **model weights**, and **outputs** with `--volume` (or `container.options.volume` in YAML).
+
+### 2.5 Environment passthrough
+
+`container.options.env` lists names that are forwarded into the **inner** `primus-cli` invocation as `--env` when set on the host (see `runner/.primus.yaml`). Examples include:
+
+`MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `DOCKER_IMAGE`, `HF_TOKEN`, `WANDB_API_KEY`, `ENABLE_NUMA_BINDING`, `USING_AINIC`, and NCCL/GLOO socket and IB-related variables (`NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, `GLOO_SOCKET_IFNAME`, `NCCL_IB_GID_INDEX`, and others).
+
+Additionally, `primus-cli-container.sh` auto-forwards any environment variable whose name starts with `PRIMUS_`, `NCCL_`, `RCCL_`, `GLOO_`, `IONIC_`, or `HIPBLASLT_` when present on the host.
+
+### 2.6 Single-node example
+
+```bash
+./primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml
+```
+
+### 2.7 Multi-node container deployment
+
+Set `MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, and `GPUS_PER_NODE` on each node (Slurm or your orchestrator sets these; see `runner/primus-cli-slurm-entry.sh`). Example pattern when launching manually:
+
+```bash
+export MASTER_ADDR=<head-node-hostname>
+export MASTER_PORT=1234
+export NNODES=4
+export NODE_RANK=<0-based index for this node>
+export GPUS_PER_NODE=8
+
+./primus-cli container -- train pretrain --config /path/to/config.yaml
+```
+
+Use `--clean` before launch to remove existing containers (`primus-cli-container.sh`).
+
+---
+
+## 3. Slurm deployment
+
+### 3.1 `srun` (interactive or blocking)
+
+```bash
+./primus-cli slurm srun -N <nodes> -p <partition> -- train pretrain --config <yaml>
+```
+
+The Slurm entry script invokes the container launcher on each allocated node. Set the image through `runner/.primus.yaml`, a custom launcher config file, or site policy; the default is `rocm/primus:v26.4`.
+
+### 3.2 `sbatch` (batch jobs)
+
+```bash
+./primus-cli slurm sbatch -N <nodes> -p <partition> --time <HH:MM:SS> --job-name <name> -o <logfile> -- \
+  train pretrain --config <yaml>
+```
+
+Add `-e <errfile>` if you want separate stderr.
+
+### 3.3 Slurm-to-Primus environment mapping
+
+`runner/primus-cli-slurm-entry.sh` sets:
+
+| Variable | Source (typical) |
+|----------|-------------------|
+| `NNODES` | `SLURM_NNODES`, or `SLURM_JOB_NUM_NODES`, or existing `NNODES` |
+| `NODE_RANK` | `SLURM_NODEID`, or `SLURM_PROCID`, or existing `NODE_RANK` |
+| `GPUS_PER_NODE` | Default `8` if unset |
+| `MASTER_ADDR` | First host in `SLURM_NODELIST` if unset |
+| `MASTER_PORT` | Default `1234` if unset |
+
+The entry script then exports `MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, and `GPUS_PER_NODE` into the container launcher.
+
+### 3.4 Slurm YAML defaults (`runner/.primus.yaml`)
+
+| Key | Default |
+|-----|---------|
+| `slurm.nodes` | `1` |
+| `slurm.gpus_per_node` | `8` |
+| `slurm.time` | `"4:00:00"` |
+| `slurm.partition` | (commented; set per site) |
+
+CLI Slurm flags override YAML when both are specified (see `runner/primus-cli-slurm.sh`).
+
+### 3.5 Entry after the first `--`
+
+Production examples pass the Primus Python command after the Slurm `--` separator, for example:
+
+```bash
+./primus-cli slurm srun -N 4 -p gpu -- train pretrain --config exp.yaml
+```
+
+The shipped `primus-cli-slurm-entry.sh` invokes **`primus-cli-container.sh`** with distributed variables set from Slurm. Container options should come from launcher configuration instead of a literal `container` token in the inner command. For **bare-metal** nodes without Docker, run `primus-cli direct` under your allocation and ensure the same distributed variables and ROCm layout as in [Multi-node configuration](#5-multi-node-configuration).
+
+---
+
+## 4. Kubernetes deployment
+
+Kubernetes integration is **not** shipped as a Helm chart or operator in this repository. The repo includes **`examples/run_k8s_pretrain.sh`**, a client script that talks to a Kubernetes **API** to create and manage training workloads (image default `docker.io/rocm/primus:v26.4`).
+
+Use that script as a reference for your platform; adapt networking, storage, and scheduling to your cluster policies.
+
+---
+
+## 5. Multi-node configuration
+
+Required variables for distributed training:
+
+| Variable | Role |
+|----------|------|
+| `MASTER_ADDR` | Hostname or IP of rank-0 process |
+| `MASTER_PORT` | TCP port for the process group rendezvous |
+| `NNODES` | Number of nodes |
+| `NODE_RANK` | Zero-based index of this node |
+| `GPUS_PER_NODE` | GPUs per node used by `torchrun` |
+
+**Flow:** User or Slurm sets the environment → `primus-cli` and `primus-cli-direct.sh` load GPU and comm settings → **`torchrun`** launches `primus/cli/main.py` with the distributed topology.
+
+**Defaults from `runner/.primus.yaml` (`direct` section):**
+
+| Key | Default |
+|-----|---------|
+| `direct.master_port` | `1234` |
+| `direct.gpus_per_node` | `8` |
+| `direct.nnodes` | `1` |
+| `direct.master_addr` | `"localhost"` |
+
+---
+
+## 6. Startup and shutdown
+
+**Lifecycle (high level):**
+
+1. Parse CLI and load YAML (`--config` chain: see [Configuration system](../02-user-guide/configuration-system.md)).
+2. Load environment (GPU detection, hooks, patches in `primus-cli-direct.sh`).
+3. Launch training via **`torchrun`** into the Python CLI.
+
+**Verification:**
+
+- `--dry-run` prints the command that would run without executing (supported in container and Slurm scripts).
+- `--debug` sets `PRIMUS_LOG_LEVEL=DEBUG` for verbose launcher and shell logging.
+
+**Shutdown:**
+
+- Normal completion or **Ctrl+C** terminates the training process.
+- In container mode, **`--clean`** removes existing containers before launch (`primus-cli-container.sh`).
+
+**Timeouts (config):**
+
+| Backend | Parameter | Location |
+|---------|-----------|----------|
+| Megatron | `distributed_timeout_minutes` | `primus/configs/modules/megatron/trainer_base.yaml` (default `10`) |
+| TorchTitan | `comm.init_timeout_seconds` | `primus/configs/modules/torchtitan/pre_trainer.yaml` (default `300`) |
+
+---
+
+## 7. Production checklist
+
+| Item | Action |
+|------|--------|
+| ROCm drivers | Install and verify with `rocm-smi` |
+| Container image | Pulled and aligned with host ROCm expectations |
+| Network | Run `preflight --network` (see [Preflight](../02-user-guide/preflight.md)) |
+| Shared data | Paths visible and consistent on all nodes |
+| Hugging Face | Set `HF_TOKEN` if using gated models |
+| Checkpoints | Save directory on shared or replicated storage with sufficient space |
+| Monitoring | Configure Weights & Biases or TensorBoard (see [Monitoring and Logging](./monitoring-logging.md)) |
+| Resources | Slurm time limits, partitions, and GPU counts match your YAML and hardware |
+
+---
+
+## Related documentation
+
+- [CLI reference](../02-user-guide/cli-reference.md)
+- [Troubleshooting](./troubleshooting.md)
+- [Multi-node networking](../04-technical-guides/multi-node-networking.md)
diff --git a/docs/05-operations/monitoring-logging.md b/docs/05-operations/monitoring-logging.md
new file mode 100644
index 000000000..581ddfe44
--- /dev/null
+++ b/docs/05-operations/monitoring-logging.md
@@ -0,0 +1,244 @@
+# Monitoring and logging
+
+This page summarizes how Primus configures application logging, experiment tracking (Weights & Biases, TensorBoard, MLflow), training metrics, profilers, ROCm memory probes, and how to capture a reproducible configuration snapshot.
+
+---
+
+## 1. Primus logging system
+
+Primus uses **loguru** for structured logging. Initialization wires **file sinks** (per log level) and a **stderr** sink, binds experiment and distributed context (`team`, `user`, `exp`, `module_name`, `node_ip`, `rank`, `world_size`), and installs an **intercept handler** so legacy `logging` output from frameworks such as Megatron is forwarded to loguru with consistent formatting.
+
+**Rank-aware behavior**
+
+- Worker processes write under `{exp_root}/logs/{module_name}/rank-{rank}/` with separate rotated files for `debug`, `info`, `warning`, and `error` (subject to `file_sink_level`).
+- The launcher **master** process can use `logs/master/` when the master logger is configured with `is_head=True`.
+
+**Levels from module configuration** (`primus/configs/modules/module_base.yaml`)
+
+| Parameter | Default | Role |
+|-----------|---------|------|
+| `sink_level` | `null` | If set, overrides both file and stderr sink levels. |
+| `file_sink_level` | `DEBUG` | Minimum level for file sinks when `sink_level` is unset. |
+| `stderr_sink_level` | `INFO` | Minimum level for stderr when `sink_level` is unset. |
+
+`init_worker_logger` in `primus/core/runtime/logging.py` reads `sink_level`, `file_sink_level`, and `stderr_sink_level` from the merged module config. The Megatron trainer maps `stderr_sink_level` to Megatron’s numeric `logging_level` (the `logging_level` field in `trainer_base.yaml` is deprecated; this mapping supersedes it)..
+
+**Shell / runner environment** (see `docs/03-configuration-reference/environment-variables.md`)
+
+| Variable | Purpose |
+|----------|---------|
+| `PRIMUS_LOG_LEVEL` | Runner verbosity: `DEBUG`, `INFO`, `WARN`, `ERROR` (default `INFO`). |
+| `PRIMUS_LOG_TIMESTAMP` | `1` enables timestamps on runner logs; `0` disables. |
+| `PRIMUS_LOG_COLOR` | `1` enables ANSI colors when appropriate; often `0` in non-TTY contexts. |
+
+**CLI**
+
+- `primus-cli --debug` sets `PRIMUS_LOG_LEVEL=DEBUG` so launcher and shell logging are verbose (see `docs/02-user-guide/cli-reference.md`).
+
+---
+
+## 2. Weights & Biases
+
+### Megatron
+
+Configuration files: `primus/configs/modules/megatron/trainer_base.yaml` and `primus_megatron_module.yaml`.
+
+Defaults in `primus_megatron_module.yaml` disable Weights & Biases; trainer fields in `trainer_base.yaml` supply names and paths when enabled.
+
+| Parameter | Default (module / trainer) | Description |
+|-----------|----------------------------|-------------|
+| `disable_wandb` | `true` (`primus_megatron_module.yaml`) | Master switch; when `false`, Primus sets paths and default project/run names from experiment metadata. |
+| `wandb_project` | `null` | If unset when Weights & Biases is enabled, defaults to `{work_group}_{user_name}`. |
+| `wandb_exp_name` | `null` | If unset, defaults to `exp_name`. |
+| `wandb_entity` | `null` | Optional WandB entity/team. |
+| `wandb_save_dir` | `null` | Deprecated in favor of `{exp_root}`; artifacts use `{exp_root}/wandb`. |
+
+**Environment**
+
+- `WANDB_API_KEY` is **required** when Weights & Biases is enabled; Primus emits a warning if it is missing (`primus/backends/megatron/patches/args/wandb_config_patches.py`).
+
+### TorchTitan
+
+Configuration file: `primus/configs/modules/torchtitan/pre_trainer.yaml`.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `metrics.enable_wandb` | `false` | Enables WandB in the TorchTitan metrics stack. |
+
+When enabled, `primus/backends/torchtitan/patches/wandb_patches.py` can set `WANDB_PROJECT` and `WANDB_RUN_NAME` from Primus experiment metadata if unset. Use `WANDB_API_KEY` for authentication.
+
+---
+
+## 3. TensorBoard
+
+### Megatron
+
+**Module toggles**
+
+Configuration file: `primus_megatron_module.yaml`
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `disable_tensorboard` | `true` | When `false`, TensorBoard output is placed under `{exp_root}/tensorboard` (Primus overrides deprecated `tensorboard_dir` with this path). |
+
+**Trainer**
+
+Configuration file: `trainer_base.yaml`
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `tensorboard_log_interval` | `1` | Steps between TensorBoard writes. |
+| `tensorboard_queue_size` | `1000` | Event file queue size. |
+| `log_timers_to_tensorboard` | `false` | Log timer stats. |
+| `log_batch_size_to_tensorboard` | `false` | Log batch size. |
+| `log_learning_rate_to_tensorboard` | `true` | Log learning rate. |
+| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. |
+| `log_memory_to_tensorboard` | `false` | Log memory stats. |
+| `log_world_size_to_tensorboard` | `false` | Log world size. |
+| `log_loss_scale_to_tensorboard` | `true` | Log loss scale. |
+| `tensorboard_dir` | `null` | Deprecated; Primus sets the directory under `exp_root`. |
+
+**Note:** Enabling Megatron **profiling** (`profile: true`) forces `disable_tensorboard` off in `update_primus_config` so TensorBoard is available for profile-related views.
+
+### TorchTitan
+
+Configuration file: `pre_trainer.yaml`.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `metrics.enable_tensorboard` | `false` | Enables TensorBoard logging. |
+| `metrics.save_tb_folder` | `tb` | Subfolder name (typically under the job dump directory in TorchTitan layouts). |
+
+**Launch TensorBoard locally**
+
+```bash
+tensorboard --logdir <path-to-tensorboard-or-tb-folder>
+```
+
+Point `<path>` at the Megatron `tensorboard` directory under the experiment root, or at the TorchTitan metrics folder that contains the `save_tb_folder` subtree.
+
+---
+
+## 4. MLflow
+
+MLflow integration is **Megatron-only** in the paths described here.
+
+**Module**
+
+Configuration file: `primus_megatron_module.yaml`
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `disable_mlflow` | `true` | When `false`, MLflow run setup runs on the **last** global rank (`world_size - 1`). |
+| `mlflow_run_name` | `null` | If unset when enabled, defaults to `{work_group}_{user_name}`. |
+| `mlflow_experiment_name` | `null` | Passed to `mlflow.set_experiment` when set. |
+
+**Startup behavior**
+
+Configuration file: `primus/backends/megatron/training/global_vars.py`
+
+- Logs training `args` as parameters.
+- Logs filtered environment variables with an `env__` prefix.
+- Collects git metadata, sets MLflow source tags, and writes `system/git_metadata.json` as a run artifact.
+
+**Environment** (typical Databricks / hosted tracking)
+
+| Variable | Role |
+|----------|------|
+| `DATABRICKS_HOST` | Checked by the Megatron trainer when MLflow is enabled; a warning is printed if unset. |
+| `DATABRICKS_TOKEN` | Authentication for Databricks-hosted tracking (see environment reference). |
+| `MLFLOW_TRACKING_URI` | Tracking server URI; optional depending on deployment. |
+
+---
+
+## 5. Training metrics
+
+### Megatron
+
+Configuration file: `trainer_base.yaml`
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `log_interval` | `100` | Steps between standard training log lines. |
+| `log_throughput` | `false` | Log throughput metrics. |
+| `log_avg_skip_iterations` | `2` | Skip initial iterations when computing averages. |
+| `log_avg_reset_interval` | `10` | Interval for resetting running averages. |
+| `log_params_norm` | `false` | Log parameter norm. |
+| `log_num_zeros_in_grad` | `false` | Log count of zero gradients. |
+| `log_progress` | `false` | Progress-style logging. |
+| `timing_log_level` | `0` | Timing log verbosity. |
+| `timing_log_option` | `minmax` | Timing aggregation option. |
+
+### TorchTitan
+
+Configuration file: `pre_trainer.yaml`.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `metrics.log_freq` | `10` | Metric logging frequency (steps). |
+| `metrics.disable_color_printing` | `false` | Disable colored console metrics. |
+| `metrics.save_for_all_ranks` | `false` | Save metrics from every rank vs. reduced ranks. |
+
+---
+
+## 6. Profiling
+
+### Megatron
+
+Configuration files: `trainer_base.yaml` and `primus_megatron_module.yaml`
+
+| Parameter | Source | Default | Description |
+|-----------|--------|---------|-------------|
+| `profile` | `trainer_base.yaml` | `false` | Enables Megatron profiling path; also forces TensorBoard on when `true`. |
+| `use_pytorch_profiler` | `trainer_base.yaml` | `false` | Use PyTorch profiler integration. |
+| `profile_ranks` | `trainer_base.yaml` | `[0]` | Ranks to profile. |
+| `profile_step_start` | `trainer_base.yaml` | `10` | First step to profile. |
+| `profile_step_end` | `trainer_base.yaml` | `12` | Last step to profile. |
+| `record_memory_history` | `trainer_base.yaml` | `false` | Record memory history. |
+| `memory_snapshot_path` | `trainer_base.yaml` | `snapshot.pickle` | Memory snapshot file name. |
+| `disable_profiler_activity_cpu` | `primus_megatron_module.yaml` | `false` | Disable CPU activities in the profiler. |
+| `torch_profiler_record_shapes` | `primus_megatron_module.yaml` | `true` | Record tensor shapes. |
+| `torch_profiler_with_stack` | `primus_megatron_module.yaml` | `true` | Capture Python stacks. |
+| `torch_profiler_use_gzip` | `primus_megatron_module.yaml` | `false` | Gzip profiler traces. |
+
+### TorchTitan
+
+Configuration file: `pre_trainer.yaml`.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `profiling.enable_profiling` | `false` | Master profiling toggle. |
+| `profiling.profile_freq` | `10` | How often to capture traces. |
+| `profiling.enable_memory_snapshot` | `false` | Enable memory snapshots. |
+| `profiling.save_memory_snapshot_folder` | `memory_snapshot` | Output folder for snapshots. |
+| `profiling.save_traces_folder` | `profile_traces` | Folder for profiler traces. |
+
+---
+
+## 7. ROCm memory monitoring
+
+Configured in `primus/configs/modules/megatron/primus_megatron_module.yaml` and applied in the Megatron trainer when logging throughput.
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `use_rocm_mem_info` | `false` | When `true`, collect ROCm memory information via `rocm-smi` on **every** iteration that hits the throughput logging branch. |
+| `use_rocm_mem_info_iters` | `[1, 2]` | When `use_rocm_mem_info` is `false`, `rocm-smi` runs only on these iteration numbers (same branch). |
+
+Collection is evaluated where `log_throughput` drives the extended iteration log (see `primus/backends/megatron/patches/training_log/print_rank_last_patches.py`): enable `log_throughput` in `trainer_base.yaml` (or overrides) when you need ROCm memory lines in the training log.
+
+---
+
+## 8. Experiment snapshots
+
+**On disk (every run)**
+
+- **Experiment root**: `{workspace}/{work_group}/{user_name}/{exp_name}` is created at config load time (`PrimusConfig`).
+- **Per-rank logs**: `{exp_root}/logs/{module_name}/rank-{rank}/` with rotated level-specific files.
+- **Checkpoints**: Megatron uses `{exp_root}/checkpoints` (trainer sets `save` to this path).
+- **TensorBoard / WandB**: Under `exp_root` as described above when those features are enabled.
+
+**MLflow** (Megatron, when enabled): Parameters, environment snapshot, and git metadata artifact provide a structured record of the run configuration and repository state.
+
+**Resolved configuration**
+
+The launcher and parser accept `--export_config`, but the default core training path (`primus/cli/subcommands/train.py` into `PrimusRuntime`) does not currently write a resolved YAML file. Archive the submitted experiment YAML, any referenced presets, launcher config, and runtime logs with each run. Treat resolved-config export as a legacy or future capability unless your deployment has implemented it on the core runtime path.
diff --git a/docs/05-operations/security.md b/docs/05-operations/security.md
new file mode 100644
index 000000000..c51d362d0
--- /dev/null
+++ b/docs/05-operations/security.md
@@ -0,0 +1,154 @@
+# Security considerations
+
+This document describes security-relevant properties of Primus as a **YAML-driven training framework** for AMD GPUs (ROCm, RCCL, containers). It is intended for operators, platform engineers, and security reviewers. It does not replace organizational policies, threat models, or vendor hardening guides.
+
+**Related documentation:** [Environment variables](../03-configuration-reference/environment-variables.md), [Installation](../01-getting-started/installation.md), [CLI reference](../02-user-guide/cli-reference.md).
+
+---
+
+## 1. Overview
+
+| Aspect | Description |
+|--------|-------------|
+| Role | Primus orchestrates distributed **training** jobs; it is **not** a general user-facing network service. |
+| Authentication / authorization | **No built-in** authentication, authorization, or multi-tenant isolation in Primus itself. |
+| Responsibility | **Security posture is determined by** the scheduler, container runtime, network, storage, identity systems, and operational practices of the deployment environment. |
+
+Treat Primus like privileged infrastructure software: run it on appropriately isolated hosts and networks, and govern secrets and data the same way you would for large-scale ML training elsewhere.
+
+---
+
+## 2. Secrets management
+
+Secrets are commonly passed as **environment variables** consumed by Primus, launchers, or third-party libraries.
+
+| Variable (examples) | Typical use |
+|---------------------|-------------|
+| `HF_TOKEN` | Hugging Face token for **gated models** and authenticated downloads. |
+| `WANDB_API_KEY` | Weights & Biases API key for experiment logging. |
+| `DATABRICKS_HOST` / `DATABRICKS_TOKEN` | Databricks or MLflow-related credentials when those integrations are used. |
+
+**Practices**
+
+| Practice | Detail |
+|----------|--------|
+| Do not hardcode secrets | Avoid putting tokens or passwords directly in YAML, shell history, or committed scripts. |
+| Prefer indirection | Use **`${VAR}`** substitution in configs to reference environment-injected values rather than literals. |
+| Slurm | Use **`--export`** deliberately; prefer site-specific **secret injection** or **credential helpers** where available. |
+| Containers | Pass secrets with **`--env`** or via **`runner/.primus.yaml`** env forwarding—never bake them into images. |
+| Rotation | Rotate API keys and tokens on a schedule and after personnel or scope changes. |
+
+A broader catalog of variables appears in [Environment variables](../03-configuration-reference/environment-variables.md).
+
+---
+
+## 3. Container security
+
+Primus-oriented container runs often require **elevated access** so ROCm, profilers, and high-performance networking behave correctly.
+
+**Common high-privilege options**
+
+| Option | Typical purpose |
+|--------|-----------------|
+| `--privileged true` | Broad device access (often required for ROCm workflows on some setups). |
+| `--cap-add SYS_PTRACE` | Debugging and profiling tooling. |
+| `--cap-add CAP_SYS_ADMIN` | Administrative operations expected by parts of the ROCm/tooling stack. |
+| `--security-opt seccomp=unconfined` | Relaxes seccomp constraints for compatibility with drivers and tools. |
+| `--ipc host` | Shared memory semantics for large tensors and collectives. |
+| `--network host` | Host networking—frequently used for **multi-node RCCL** performance and simplicity. |
+
+**Device access (examples)**
+
+| Device | Role |
+|--------|------|
+| `/dev/kfd` | ROCm kernel interface. |
+| `/dev/dri` | GPU render nodes. |
+| `/dev/infiniband` | InfiniBand character devices when using IB. |
+
+**Risks**
+
+| Risk | Why it matters |
+|------|----------------|
+| Privileged containers | Substantial **host** access; container escape or compromise has high impact. |
+| Host networking | Exposes the container to the **host’s network namespace**; services may bind broadly. |
+| Shared IPC | Potential for **cross-process interference** or information leakage if workloads share hosts improperly. |
+
+**Mitigations**
+
+| Mitigation | Detail |
+|------------|--------|
+| Dedicated training nodes | Run training on **isolated** machines rather than mixed with user-facing services. |
+| Network controls | Apply **firewall rules** and **segmentation** so only required ports and peers are reachable. |
+| Trusted images | Pull from **trusted registries**, pin digests, and verify image provenance. |
+| Monitoring | Track **CPU, memory, GPU, and network** usage; alert on anomalous processes or egress. |
+
+---
+
+## 4. Third-party dependencies
+
+Primus integrates **third-party submodules** and Python packages; each carries its own license and maintenance cadence.
+
+**Representative submodules**
+
+| Component | Notes (non-exhaustive) |
+|-----------|-------------------------|
+| Megatron-LM | MIT License; NVIDIA upstream. |
+| TorchTitan | Apache 2.0; PyTorch / Meta ecosystem. |
+| MaxText | Apache 2.0; Google upstream. |
+| Megatron-Bridge | NVIDIA NeMo ecosystem. |
+| Emerging-Optimizers | NVIDIA NeMo ecosystem. |
+| HummingbirdXT | AMD AGI ecosystem. |
+
+**Python dependencies**
+
+Runtime tooling often includes packages such as **loguru**, **wandb**, **nltk**, **matplotlib**, **mlflow**, and others as declared in project requirements—verify the canonical list in the repository’s `requirements.txt` (or lockfile) for your revision.
+
+**Recommendations**
+
+| Recommendation | Rationale |
+|----------------|-----------|
+| Pin versions | Reproducible builds and controlled upgrade paths. |
+| Update submodules | Security and correctness fixes flow from upstream projects. |
+| Monitor advisories | Subscribe to upstream security notices for frameworks you enable. |
+
+---
+
+## 5. Network security
+
+| Property | Detail |
+|----------|--------|
+| RCCL / NCCL traffic | **Not encrypted** at the application layer; assumes a **trusted network path**. |
+| Coordination | **`MASTER_ADDR`** and **`MASTER_PORT`** should reside on a **private** or otherwise **trusted** segment. |
+| InfiniBand | Often on a **dedicated fabric**; still treat adjacent compromised hosts as in-scope for lateral movement. |
+| TLS / mTLS | Primus does **not** provide TLS or mTLS for inter-node training traffic by default. |
+
+For physical and logical networking topics, see [Multi-node networking](../04-technical-guides/multi-node-networking.md).
+
+---
+
+## 6. Data security
+
+| Asset | Consideration |
+|-------|-----------------|
+| Training data | May include **PII**, licensed corpora, or export-controlled material—classify and restrict accordingly. |
+| Checkpoints | Contain **full model state**; treat as sensitive intellectual property. |
+| Storage permissions | Use **least privilege** on shared filesystems and object stores. |
+| `HF_TOKEN` | Grants access to **gated** Hugging Face assets—protect like any other long-lived credential. |
+
+Checkpoint formats and operational practices are described in [Checkpoint management](../04-technical-guides/checkpoint-management.md).
+
+---
+
+## 7. What is not verified
+
+The following items reflect **typical gaps** in public-facing evidence for many research and infrastructure codebases; confirm against your organization’s audits and CI for your fork and deployment.
+
+| Topic | Status (evidence-based caveat) |
+|-------|--------------------------------|
+| Independent security audit | **No** comprehensive third-party audit of this codebase is asserted here. |
+| CI secrets scanning | **No** guarantee of automated secret detection in CI unless your pipeline adds it. |
+| Dependency vulnerability scanning | **No** guarantee of continuous SCA unless your pipeline adds it. |
+| Container images | Images may contain **unpatched** OS or Python packages—scan and rebuild on a schedule. |
+| RCCL / NCCL traffic | **Not** encrypted or mutually authenticated by default; rely on network trust boundaries. |
+
+Use this section as a checklist for **your** production controls: add scanning, signing, policy-as-code, and periodic reviews appropriate to your threat model.
diff --git a/docs/05-operations/troubleshooting.md b/docs/05-operations/troubleshooting.md
new file mode 100644
index 000000000..313a0a634
--- /dev/null
+++ b/docs/05-operations/troubleshooting.md
@@ -0,0 +1,193 @@
+# Troubleshooting guide
+
+This guide is the primary reference for diagnosing and resolving common failures when running Primus on AMD GPUs (ROCm, RCCL, Docker). It complements the [CLI reference](../02-user-guide/cli-reference.md), [Preflight](../02-user-guide/preflight.md), [Benchmarking](../02-user-guide/benchmarking.md), and [Configuration system](../02-user-guide/configuration-system.md) documentation.
+
+---
+
+## 1. Diagnostic tools
+
+Use these tools before scaling a job or when a failure is hard to localize.
+
+| Tool | Purpose |
+|------|---------|
+| `primus-cli --debug` | Enables verbose logging in `primus-cli` for command construction, delegation, and runtime details. |
+| `primus-cli --dry-run` | Prints the commands Primus would run without executing them—useful to verify wrappers, paths, and MPI/launcher wiring. |
+| `--export_config <path>` | Parsed by the training config parser, but not written by the current default `PrimusRuntime` path. Treat it as legacy/future functionality unless your deployment has implemented it. |
+| `NCCL_DEBUG=INFO` | Surfaces detailed RCCL/NCCL connection and collective logs (set in the job environment). |
+| `primus-cli direct -- preflight --host --gpu --network` | Fast host, GPU, and network validation. See [Preflight](../02-user-guide/preflight.md). |
+| `PRIMUS_PATCHES=none` | Disables Primus patches to the selected backend to isolate whether a failure is Primus-specific or upstream. |
+
+**Examples**
+
+```bash
+# Verbose CLI + dry run (no training executed)
+primus-cli --debug --dry-run direct -- train pretrain --config path/to/config.yaml
+
+# Preflight: environment validation only
+primus-cli direct -- preflight --host --gpu --network
+
+# RCCL/NCCL verbose logs in the training process environment
+export NCCL_DEBUG=INFO
+
+# Isolate backend vs. Primus patch layer
+export PRIMUS_PATCHES=none
+```
+
+For end-to-end checks including optional performance probes, see [Preflight](../02-user-guide/preflight.md) (`preflight --perf-test`).
+
+---
+
+## 2. Out of memory (OOM) errors
+
+**Symptoms:** HIP/CUDA OOM messages, worker processes killed by the OOM killer, or abrupt exits during forward/backward.
+
+**Primary levers and mitigations**
+
+| Approach | What to change |
+|----------|------------------|
+| Reduce per-device activation memory | Lower **`micro_batch_size`** (often the first knob). |
+| Shard weights/activations across devices | Increase **tensor parallelism** (`tensor_model_parallel_size` or `parallelism.tensor_parallel_degree`, depending on backend). See [Parallelism configuration](../04-technical-guides/parallelism-configuration.md). |
+| Trade compute for memory | **Activation recomputation:** `recompute_granularity: full`, `recompute_method: uniform`, `recompute_num_layers: <N>`. |
+| Shard optimizer state | **`use_distributed_optimizer: true`** (Megatron-style stacks). |
+| FSDP / sharded data parallel | Megatron: **`use_torch_fsdp2: true`**. TorchTitan: increase **`data_parallel_shard_degree`**. |
+| CPU offload | **`optimizer_cpu_offload: true`** where supported. |
+| Sequence length | Reduce **maximum sequence length** if the workload allows. |
+| MoE / scratch memory | Set **`HSA_NO_SCRATCH_RECLAIM=1`** to reduce scratch-memory conflicts on some MoE workloads. |
+| Plan before you run | **`primus-cli ... -- projection memory --config <yaml>`**—see [Projection](../02-user-guide/projection.md). |
+
+**Related references:** [Parallelism strategies](../04-technical-guides/parallelism-strategies.md), [Performance tuning](../04-technical-guides/performance-tuning.md), [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md).
+
+---
+
+## 3. Distributed communication failures
+
+**Symptoms:** Hangs at initialization, NCCL/RCCL timeouts, connection errors, or inconsistent ranks.
+
+| Cause | What to verify / fix |
+|-------|----------------------|
+| Wrong network interface | Set `NCCL_SOCKET_IFNAME` to the correct interface; exclude virtual interfaces, e.g. `^docker0,lo`. |
+| `MASTER_ADDR` unreachable | From every node, resolve DNS/IP consistently; verify firewalls and routing. |
+| Port already in use | Change **`MASTER_PORT`** to a free port on all nodes. |
+| InfiniBand not used or missing | Confirm `/dev/infiniband` exists where expected; run `ibstat`; ensure IB kernel modules are loaded. See [Multi-node networking](../04-technical-guides/multi-node-networking.md). |
+| Timeout too aggressive | Megatron: increase **`distributed_timeout_minutes`**. TorchTitan: increase **`comm.init_timeout_seconds`**. |
+| Mismatched world size | Align **`NNODES`**, **`GPUS_PER_NODE`**, and launcher settings across **all** nodes. |
+
+**Debugging**
+
+```bash
+export NCCL_DEBUG=INFO
+```
+
+**Validation**
+
+```bash
+primus-cli direct -- preflight --network
+primus-cli direct -- benchmark rccl --op all_reduce
+```
+
+Collective behavior and RCCL roles are summarized in [Collective operations](../04-technical-guides/collective-operations.md).
+
+---
+
+## 4. Container issues
+
+**Symptoms:** Permission denied on devices, GPUs not visible inside the container, immediate exit, or RCCL failures only under Docker.
+
+| Symptom | Typical cause | Mitigation |
+|---------|-----------------|------------|
+| GPU not visible | Devices not passed through | Ensure **`--device /dev/kfd`** and **`--device /dev/dri`** (Primus container mode typically sets these). |
+| Permission denied on GPU | Group membership / permissions | Add **`--group-add video`**; ensure the user has **video/render** access on the host. |
+| InfiniBand missing in container | Device not mounted | Add **`--device /dev/infiniband`** (and related uverbs devices as required by your site). |
+| Debugger/profiler failures | Missing capabilities | **`--cap-add SYS_PTRACE`** and **`--cap-add CAP_SYS_ADMIN`** are required for many ROCm tooling paths. |
+| Driver/library mismatch | Image vs. host ROCm | Match **container image ROCm** to **host ROCm driver** version. |
+| Data or code not found | Bind mounts | Use **`--volume /host/path:/container/path`** for datasets and workspace. |
+| Env vars missing in container | Forwarding | Check **`runner/.primus.yaml`** `container.options.env` for auto-forwarded variables; add extras with **`--env KEY=VALUE`**. |
+
+Installation and container-oriented setup are covered in [Installation](../01-getting-started/installation.md).
+
+---
+
+## 5. Configuration errors
+
+**Symptoms:** YAML parse failures, "unknown key" or type errors, silent wrong behavior after edits.
+
+| Issue | Resolution |
+|-------|------------|
+| Unset `${VAR}` | `${VAR}` with no default fails if `VAR` is unset. Use **`${VAR:default}`** or **export** the variable before launch. |
+| Broken `extends:` chain | Confirm every referenced file exists and paths resolve relative to the expected directory. |
+| Wrong parameter name | Cross-check backend docs: [Megatron](../03-configuration-reference/megatron-parameters.md), [TorchTitan](../03-configuration-reference/torchtitan-parameters.md), [MaxText](../03-configuration-reference/maxtext-parameters.md), [Megatron Bridge](../03-configuration-reference/megatron-bridge-parameters.md). |
+| Override not applied | Review merge order: **CLI > experiment overrides > module preset with model preset additions**; duplicate top-level module keys win over model preset keys. See [Configuration system](../02-user-guide/configuration-system.md). |
+| Effective config unknown | Use **`--dry-run`** to inspect the launch command and manually trace the experiment YAML, module preset, model preset, and overrides. Resolved-config export is not currently written by the default core runtime. |
+
+---
+
+## 6. Backend-specific issues
+
+### Megatron
+
+| Issue | Mitigation |
+|-------|------------|
+| Suspected Primus patch interaction | `export PRIMUS_PATCHES=none` and retry with vanilla Megatron behavior. |
+| Custom kernel compile failures | `disable_compile_dependencies: true` skips custom kernel compilation where applicable. |
+| Wrong third-party path | Set **`BACKEND_PATH`** to override third-party resolution. |
+
+### TorchTitan
+
+| Issue | Mitigation |
+|-------|------------|
+| Submodule drift | `git submodule update --recursive` so `third_party/torchtitan` matches the Primus revision you run. |
+| `torch.compile` instability | `export TORCH_COMPILE_DISABLE=1` or set **`compile.enable: false`** in config. |
+
+### MaxText (JAX)
+
+| Issue | Mitigation |
+|-------|------------|
+| JAX / jaxlib vs ROCm | Verify JAX and jaxlib builds match your ROCm stack. |
+| XLA memory pressure | Tune **`XLA_PYTHON_CLIENT_MEM_FRACTION`** to cap client-side allocator use. |
+
+---
+
+## 7. Performance issues
+
+**Symptoms:** Low tokens/sec, long iteration time, or poor scaling versus expectations.
+
+**Diagnosis**
+
+| Step | Command / action |
+|------|-------------------|
+| Compute sanity | `primus-cli direct -- benchmark gemm` |
+| Interconnect | `primus-cli direct -- benchmark rccl --op all_reduce` |
+| Broader probe | `primus-cli direct -- preflight --perf-test` (see [Preflight](../02-user-guide/preflight.md)) |
+
+**Common fixes**
+
+| Area | Action |
+|------|--------|
+| GEMM / kernels | Enable **hipBLASLt tuning** (multi-stage workflow—see [Performance tuning](../04-technical-guides/performance-tuning.md)). |
+| Primus stack | Enable **`enable_primus_turbo: true`** where supported. |
+| Communication overlap | **`overlap_grad_reduce: true`**, **`overlap_param_gather: true`** (when applicable to your backend). |
+| MoE / scratch | Confirm **`HSA_NO_SCRATCH_RECLAIM=1`** when recommended for your model class. |
+| FP8 | Enable FP8 when hardware and backend support it. |
+
+---
+
+## 8. Data issues
+
+| Symptom | Checks |
+|---------|--------|
+| Mock data works; real data fails | **`data_path`** format: Megatron typically expects **`.bin` / `.idx`** pairs. See [Data preparation](../04-technical-guides/data-preparation.md). |
+| Tokenizer errors | **`tokenizer_type`** and **`tokenizer_model`** must match (e.g., Hugging Face tokenizer ID for `HuggingFaceTokenizer`). |
+| Hugging Face download failures | Set **`HF_TOKEN`** for gated models; verify outbound network and cache directories. |
+
+---
+
+## 9. Known limitations
+
+| Area | Note |
+|------|------|
+| MaxText | Parameter completeness depends on upstream MaxText **`base.yml`**; some keys are inherited from upstream defaults. |
+| Megatron Bridge | Recipe parameters may be loaded dynamically; not every key appears in static reference tables. |
+| HummingbirdXT | Less mature than other backends; expect sharper edges in configs and tooling. |
+| Primus-Turbo | Requires a **separate installation** step; not always present by default. |
+
+For terminology, see the [Glossary](../01-getting-started/glossary.md). For checkpoint-related failures, see [Checkpoint management](../04-technical-guides/checkpoint-management.md).
diff --git a/docs/06-developer-guide/README.md b/docs/06-developer-guide/README.md
new file mode 100644
index 000000000..a8c21c5a6
--- /dev/null
+++ b/docs/06-developer-guide/README.md
@@ -0,0 +1,17 @@
+# Developer guide
+
+For contributors and maintainers.
+
+- [Architecture](architecture.md): system design, runtime, backends, patch system
+- [Contributing](contributing.md): development setup, code style, PR process
+- [Testing](testing.md): test types, running tests, CI pipeline
+- [Extending backends](extending-backends.md): adding new training backends
+- [Adding models](adding-models.md): adding model configurations per backend
+- [Model support matrix](model-support-matrix.md): supported models per backend and GPU
+- [CLI architecture](cli-architecture.md): CLI internals: subcommand discovery, dispatch, and launch wrappers
+- [Backend patch notes](backend-patch-notes.md): Primus-specific backend arguments and the files they patch
+- [Tooling](tooling.md): auxiliary analysis, benchmarking, visualization, and diagnostics tools
+
+---
+
+[← Documentation home](../README.md)
diff --git a/docs/06-developer-guide/adding-models.md b/docs/06-developer-guide/adding-models.md
new file mode 100644
index 000000000..d76010615
--- /dev/null
+++ b/docs/06-developer-guide/adding-models.md
@@ -0,0 +1,379 @@
+# Adding model configurations
+
+This guide explains how to add **model configuration YAML** for each Primus training backend. Model presets live under `primus/configs/models/<framework>/` and are referenced from **experiment** YAML under `examples/<framework>/configs/...`. Backend-specific parameter references:
+
+- [Megatron parameters](../03-configuration-reference/megatron-parameters.md)
+- [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md)
+- [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)
+- [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md)
+
+---
+
+## Overview: Three-layer configuration
+
+For each backend, Primus composes configuration in three layers:
+
+1. **Experiment config** (entry point): `examples/<backend>/configs/<GPU SKU>/<name>.yaml`—selects `framework`, `config` (module preset), `model` (model preset), and `overrides`.
+2. **Module config** (trainer defaults): `primus/configs/modules/<framework>/<module>.yaml`—training loop defaults, logging, optimizer blocks, and backend-specific knobs.
+3. **Model config** (architecture and assets): `primus/configs/models/<framework>/<model>.yaml`—architecture fields and tokenizer or Hugging Face paths, shaped differently per backend (see sections below).
+
+At runtime, `modules.<module>.model: <file>.yaml` resolves to `primus/configs/models/<framework>/<file>.yaml` and is merged into module parameters before the backend adapter converts them.
+
+---
+
+## Adding a Megatron model
+
+### How Megatron configs are wired
+
+1. **Experiment config** (entry point):
+
+   ```yaml
+   # examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+   modules:
+     pre_trainer:
+       framework: megatron
+       config: pre_trainer.yaml           # module-level trainer config
+
+       # model to run
+       model: llama3.1_8B.yaml            # model config name
+   ```
+
+2. **Module config** (trainer-level defaults): `primus/configs/modules/megatron/pre_trainer.yaml`—extends shared bases and sets Megatron training defaults.
+
+3. **Model config** (architecture + tokenizer):
+
+   ```yaml
+   # primus/configs/models/megatron/llama3.1_8B.yaml
+   extends:
+     - llama3_8B.yaml
+
+   tokenizer_type: HuggingFaceTokenizer
+   tokenizer_model: meta-llama/Llama-3.1-8B
+
+   max_position_embeddings: 131072
+   ```
+
+At runtime, `modules.pre_trainer.model: llama3.1_8B.yaml` resolves to `primus/configs/models/megatron/llama3.1_8B.yaml`. The `extends` chain pulls in parent files (for example `llama3_8B.yaml` → `llama3_base.yaml` → `llama_base.yaml`).
+
+### Files you typically add
+
+| Artifact | Purpose |
+| -------- | ------- |
+| **Model preset** (required) | New YAML under `primus/configs/models/megatron/`—architecture, tokenizer, and optional `extends`. |
+| **Experiment config** (required) | New or copied YAML under `examples/megatron/configs/MI300X/` or `MI355X/`—points `model:` at your preset and sets `overrides` (batch size, precision, parallelism, `mock_data`, and so on). |
+| **Module preset** (optional) | Only if you need trainer defaults that differ from `pre_trainer.yaml`—new file under `primus/configs/modules/megatron/` and reference it as `config:` in the experiment. |
+
+### Example: TinyLlama 1.1B from Hugging Face
+
+Assume Hugging Face repo `TinyLlama/TinyLlama-1.1B-Chat-v1.0` is not yet represented in Primus. You can add a local model preset (not necessarily committed upstream) as follows.
+
+**1. Decide the architecture**
+
+Because TinyLlama is not shipped as a Megatron preset in Primus, you can:
+
+- **Option A (recommended):** extend `language_model.yaml` and set all architecture fields explicitly.
+- **Option B:** extend the closest existing model (for example a LLaMA-style preset) and override differing fields.
+
+**2. Map Hugging Face `config.json` to Megatron keys**
+
+Read from the Hugging Face model (typically `config.json` or the model card):
+
+| Hugging Face / concept | Megatron YAML (typical keys) |
+| ---------------------- | ---------------------------- |
+| `hidden_size` | `hidden_size` |
+| `intermediate_size` | `ffn_hidden_size` |
+| `num_attention_heads` | `num_attention_heads` |
+| `num_hidden_layers` | `num_layers` |
+| `num_key_value_heads` | Use with `num_attention_heads` to set `num_query_groups` (often `num_attention_heads / num_key_value_heads`) |
+| `max_position_embeddings` | `max_position_embeddings` |
+
+**3. Create `tinyllama_1.1B.yaml`**
+
+Path: `primus/configs/models/megatron/tinyllama_1.1B.yaml`
+
+```yaml
+extends:
+  - language_model.yaml        # generic Megatron language model base
+
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+
+hidden_size: 2048
+ffn_hidden_size: 5632                # intermediate_size in HF config.json
+num_attention_heads: 32
+num_layers: 22                       # num_hidden_layers in HF config.json
+num_query_groups: 8                  # e.g. 32 / 4 if HF has 4 KV heads
+
+max_position_embeddings: 2048
+position_embedding_type: rope
+```
+
+**4. Point an experiment at the new model**
+
+Copy an existing experiment (for example `examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml`) and set `model:` to your preset. Use **mock data** first for a quick smoke test:
+
+```yaml
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:tinyllama_1.1B-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: tinyllama_1.1B.yaml
+    overrides:
+      save: null
+      disable_last_saving: true
+      stderr_sink_level: DEBUG
+
+      mock_data: true
+      train_iters: 50
+      micro_batch_size: 2
+      global_batch_size: 128
+
+      seq_length: 2048
+```
+
+**5. Run verification**
+
+```bash
+./primus-cli direct -- \
+  train pretrain \
+  --config examples/megatron/configs/MI300X/tinyllama_1.1B-pretrain.yaml
+```
+
+Confirm in logs that `framework` is `megatron`, the resolved model file is `tinyllama_1.1B.yaml`, and the tokenizer matches your preset.
+
+### Megatron checklist
+
+- [ ] Choose an appropriate base under `primus/configs/models/megatron/` (`language_model.yaml` or a close LLaMA-style model).
+- [ ] Set `tokenizer_type`, `tokenizer_model`, and architecture fields aligned with Hugging Face.
+- [ ] Add or update an experiment YAML under `examples/megatron/configs/...` with `model: <your_model>.yaml`.
+- [ ] Run `./primus-cli direct -- train pretrain --config ...` to validate resolution and a short run.
+
+---
+
+## Adding a TorchTitan model
+
+### How TorchTitan configs are wired in Primus
+
+1. **Experiment config**:
+
+   ```yaml
+   # examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
+   modules:
+     pre_trainer:
+       framework: torchtitan
+       config: pre_trainer.yaml
+
+       model: llama3.1_8B.yaml
+       overrides:
+         training:
+           local_batch_size: 4
+           seq_len: 8192
+           mock_data: false
+           steps: 50
+   ```
+
+2. **Module config**: `primus/configs/modules/torchtitan/pre_trainer.yaml`—training defaults, quantization fragments, and TorchTitan-oriented structure.
+
+3. **Model config**—`job` and `model` sections consumed by the TorchTitan launcher:
+
+   ```yaml
+   # primus/configs/models/torchtitan/llama3.1_8B.yaml
+   job:
+     dump_folder: "./outputs"
+     description: "Llama 3.1 8B training"
+
+   model:
+     name: "llama3"
+     flavor: "8B"
+     hf_assets_path: "meta-llama/Llama-3.1-8B"
+     converters:
+       - primus_turbo
+   ```
+
+At runtime, `modules.pre_trainer.model: llama3.1_8B.yaml` resolves to `primus/configs/models/torchtitan/llama3.1_8B.yaml`. The launcher uses `job` and `model` to wire the PyTorch model and training loop.
+
+### Mapping from Hugging Face to TorchTitan
+
+You need:
+
+- **Model family** (`model.name`): must match a family implemented in TorchTitan (for example `llama3`, `qwen3`, `deepseek_v3`).
+- **Flavor** (`model.flavor`): a size key defined in TorchTitan code (for example `8B`, `70B`, `1.7b`)—see `third_party/torchtitan/torchtitan/models/<family>/`.
+- **Hugging Face assets** (`model.hf_assets_path`): repository used to load weights and tokenizer.
+
+**Important limitations**
+
+- TorchTitan can only train models that are **implemented in the TorchTitan codebase**. The YAML under `primus/configs/models/torchtitan/` does **not** define new architectures; it selects and configures existing `*ModelArgs` entries.
+- If a family or flavor is missing in TorchTitan, you cannot enable it with YAML alone—extend TorchTitan first, then add a Primus preset.
+
+### Example pattern: Qwen3 8B preset
+
+Qwen3 8B already exists in this repository as a TorchTitan preset and example. Use it as a pattern when adding a different TorchTitan model or flavor that is implemented upstream but not yet represented in Primus.
+
+**Existing file:** `primus/configs/models/torchtitan/qwen3_8b.yaml`
+
+```yaml
+job:
+  dump_folder: "./outputs"
+  description: "Qwen 3 8B training"
+
+model:
+  name: "qwen3"
+  flavor: "8B"
+  hf_assets_path: "Qwen/Qwen3-8B"
+  converters:
+    - primus_turbo
+```
+
+**Field meanings:**
+
+- **`job.dump_folder`**—where TorchTitan writes logs and checkpoints for the job.
+- **`job.description`**—free-form description shown in logs and metadata.
+- **`model.name` / `model.flavor`**—the TorchTitan family and size key; both must exist in the TorchTitan code.
+- **`model.hf_assets_path`**—Hugging Face repository used to load weights and tokenizer.
+- **`model.converters`**—extra TorchTitan converters; `primus_turbo` is the default used in the Primus examples.
+
+The architecture itself lives in TorchTitan code, not in this YAML. For example, Qwen3 8B is declared in `third_party/torchtitan/torchtitan/models/qwen3/__init__.py`:
+
+```python
+"8B": Qwen3ModelArgs(
+    vocab_size=151936,
+    max_seq_len=4096,
+    head_dim=128,
+    dim=4096,
+    n_layers=36,
+    n_heads=32,
+    n_kv_heads=8,
+    qk_norm=True,
+    hidden_dim=12288,
+    rope_theta=1000000,
+),
+```
+
+The Primus preset only selects and configures such a definition.
+
+**Experiment snippet** (copy from an existing TorchTitan example and change `model:`):
+
+```yaml
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+    model: qwen3_8b.yaml
+    overrides:
+      training:
+        local_batch_size: 4
+        seq_len: 4096
+        mock_data: true
+        steps: 50
+```
+
+**Run:**
+
+```bash
+./primus-cli direct -- \
+  train pretrain \
+  --config examples/torchtitan/configs/MI300X/qwen3_8B-pretrain.yaml
+```
+
+For a new model, create a new preset and example path that matches the upstream TorchTitan family/flavor you are adding.
+
+### TorchTitan checklist
+
+- [ ] Define `job` (for example `dump_folder`, `description`) and `model` (`name`, `flavor`, `hf_assets_path`, `converters`).
+- [ ] Add an experiment under `examples/torchtitan/configs/...` referencing `model: <your_model>.yaml`.
+- [ ] Run `./primus-cli direct -- train pretrain --config ...` for a short job.
+
+---
+
+## Adding a MaxText model
+
+MaxText (JAX) model presets in Primus are intentionally thin: they set **`model_name`** and **`tokenizer_path`** (and extend `model_base.yaml`) so MaxText can load its own architecture tables when available.
+
+**Typical model preset**
+
+Path pattern: `primus/configs/models/maxtext/<name>.yaml`
+
+```yaml
+extends:
+  - model_base.yaml
+
+model_name: "llama3-8b"
+tokenizer_path: "meta-llama/Meta-Llama-3-8B"
+```
+
+Comments in `primus/configs/models/maxtext/model_base.yaml` explain that architecture parameters are resolved from MaxText’s `configs/models/<model_name>.yml` when present, or from Primus overrides as appropriate.
+
+**Experiment wiring**
+
+Experiments reference the preset the same way as other backends, for example:
+
+```yaml
+modules:
+  pre_trainer:
+    framework: maxtext
+    config: pre_trainer.yaml
+    model: llama3_8B.yaml
+```
+
+**Supported architectures**
+
+For the authoritative list of model names and architectures MaxText supports, see the [MaxText](https://github.com/AI-Hypercomputer/maxtext) repository and upstream documentation. Primus examples under `examples/maxtext/configs/MI300X/` and `MI355X/` illustrate which presets are exercised in this tree.
+
+---
+
+## Adding a Megatron Bridge model (post-training)
+
+Megatron Bridge model presets are small YAML files that select a **recipe**, **flavor**, and **Hugging Face path**, plus optional dataset blocks.
+
+**Example preset** (`primus/configs/models/megatron_bridge/qwen3_8b.yaml`):
+
+```yaml
+recipe: qwen.qwen3
+flavor: qwen3_8b_finetune_config
+hf_path: Qwen/Qwen3-8B
+
+dataset:
+  dataset_name: "rajpurkar/squad"
+```
+
+| Field | Role |
+| ----- | ---- |
+| `recipe` | Logical recipe module (for example `qwen.qwen3`, `llama.llama3`). |
+| `flavor` | Named configuration within the recipe (for example `qwen3_8b_finetune_config`). |
+| `hf_path` | Hugging Face model id for weights and tokenizer. |
+
+**Experiment** (pattern from `examples/megatron_bridge/configs/`):
+
+```yaml
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+    model: qwen3_8b.yaml
+    overrides:
+      precision_config: bf16_mixed
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+```
+
+See [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md) for the full override surface.
+
+---
+
+## Testing new models
+
+Use a staged approach so failures are easy to localize.
+
+| Stage | Goal | Typical settings |
+| ----- | ---- | ---------------- |
+| **Mock or synthetic data** | Validate config resolution, tokenizer, and a few steps without real datasets. | Megatron: `mock_data: true`. TorchTitan: `training.mock_data: true`. MaxText: `dataset_type: "synthetic"` in overrides where applicable. Keep `train_iters` / `steps` small. |
+| **Single-GPU** | Confirm numerics and memory before scaling. | Set tensor and pipeline parallelism to 1 in overrides; use one process / one device per your launcher docs. |
+| **Multi-GPU** | Match production parallelism. | Set `tensor_model_parallel_size`, `pipeline_model_parallel_size`, expert / context parallel sizes, or MaxText `ici_*` / `dcn_*` fields as required by the model size and hardware. |
+
+Cross-check [Parallelism configuration](../04-technical-guides/parallelism-configuration.md) and [Model support matrix](./model-support-matrix.md) when moving from single-device to multi-device runs.
diff --git a/docs/06-developer-guide/architecture.md b/docs/06-developer-guide/architecture.md
new file mode 100644
index 000000000..a82520b5e
--- /dev/null
+++ b/docs/06-developer-guide/architecture.md
@@ -0,0 +1,119 @@
+# Architecture overview
+
+This document describes how the Primus training framework is structured: CLI and configuration, the core runtime orchestrator, backend adapters, trainer lifecycle, and the patch system.
+
+## 1. System overview
+
+Primus is organized into three conceptual layers: runtime launch (how processes and GPUs are started), hooks and patches (environment and in-process adjustments), and task execution (CLI subcommands that drive training and utilities).
+
+```
+┌─────────────────────────────────────────────────────┐
+│              Runtime Layer (runner/)                  │
+│         direct | container | slurm                   │
+│    GPU detection, env setup, distributed launch      │
+├─────────────────────────────────────────────────────┤
+│           Hook / Patch System                        │
+│    runner/helpers/hooks/ | primus/core/patches/       │
+│    Pre/post processing, runtime monkey-patches       │
+├─────────────────────────────────────────────────────┤
+│           Task Execution Layer                       │
+│        primus/cli/subcommands/                       │
+│    train | benchmark | preflight | projection        │
+└─────────────────────────────────────────────────────┘
+```
+
+The repository also provides shell entrypoints under `runner/` (for example `primus-cli-direct.sh`, `primus-cli-container.sh`, `primus-cli-slurm.sh`) that prepare the environment and invoke the Python CLI.
+
+## 2. CLI and plugin system
+
+- **Entry point:** `primus/cli/main.py` is the unified CLI entry. It discovers subcommand modules under `primus/cli/subcommands/` with `pkgutil.walk_packages`, skipping modules whose leaf name starts with `_`.
+- **Registration contract:** Each subcommand module exposes `register_subcommand(subparsers)` and must return the configured parser. The parser must call `set_defaults(func=run)` so `main()` can dispatch to the handler.
+- **Parsing:** The CLI uses the standard library `argparse` only (no Click or Typer).
+- **Unknown arguments:** `main()` calls `parse_known_args()`. For selected subcommands (`train`, `projection`, `preflight`), trailing tokens are passed through to the handler as overrides; for other commands, unknown arguments are rejected.
+
+## 3. Configuration pipeline
+
+Configuration flows from experiment YAML to a resolved structure consumed by the runtime.
+
+1. **CLI** parses `--config` / `--exp` (required for train flows) pointing at an experiment YAML file.
+2. **`load_primus_config()`** (used by `PrimusRuntime`) delegates to **`PrimusParser.parse()`** in `primus/core/launcher/parser.py`. The parser loads the experiment file via **`yaml_utils.parse_yaml_to_namespace()`**, which uses **`primus/core/config/yaml_loader.parse_yaml()`** for `${VAR}` / `${VAR:default}` substitution and `extends:` inheritance with deep merge.
+3. **Per trainer module** (names containing `trainer`, for example `pre_trainer`):
+   - **`PresetLoader.load()`** loads the module preset from `primus/configs/modules/<framework>/<config>.yaml`.
+   - **`PresetLoader.load()`** loads the model preset from `primus/configs/models/<framework>/<model>.yaml`.
+   - Each preset is loaded through the same YAML pipeline (env substitution and `extends:` chains).
+4. **`parse_platform()`** merges platform settings from `primus/configs/platforms/` (defaulting to `platform_azure.yaml` when the experiment omits `platform`).
+5. **CLI overrides:** For `primus train`, `main()` passes `unknown_args` into the train handler. `PrimusRuntime` parses them with `parse_cli_overrides()` and **deep-merges** them into `module_config.params`.
+6. **Result:** A resolved configuration where each module exposes a **`params`** namespace (`SimpleNamespace`) for training parameters, produced by `_normalize_module_for_runtime()` in `primus/core/config/primus_config.py`.
+
+The object returned from `load_primus_config()` is a lightweight `SimpleNamespace` (not `PrimusConfig`), with `modules` as a **list** of module configs, each tagged with a `.name` field.
+
+## 4. Core runtime (PrimusRuntime)
+
+`primus/core/runtime/train_runtime.py` defines **`PrimusRuntime`**, the main orchestrator for the new core training path. Execution for a single module follows this flow:
+
+1. **`load_primus_config()`** loads and validates the experiment; **`get_module_config()`** selects the requested module (for example `pre_trainer` or `post_trainer`).
+2. **`_apply_overrides()`** merges CLI overrides into `module_config.params`.
+3. **`_initialize_environment()`** ensures the data directory exists and calls **`setup_training_env()`** (Hugging Face cache and related setup).
+4. **`_initialize_distributed_context()`** reads torchrun-style rank and master information via **`get_torchrun_env()`**.
+5. **`_initialize_logging()`** initializes worker logging.
+6. **`BackendRegistry.get_adapter(framework)`** resolves the **`BackendAdapter`** (lazy-importing `primus.backends.<name>` if needed).
+7. **`adapter.setup_backend_path()`** inserts the backend tree on `sys.path`. Resolution order: CLI `--backend_path`, then the `BACKEND_PATH` env var, then the default—`third_party/<dir>` under the repo root, followed by `$PRIMUS_THIRDPARTY_DIR` or `~/.cache/Primus/third_party` (the `primus-cli deps sync` location).
+8. **`adapter.prepare_backend()`** runs backend setup hooks (via **`BackendRegistry.run_setup()`** by default).
+9. **`adapter.convert_config(module_config.params)`** produces **`backend_args`** for the trainer.
+10. **`run_patches(phase="build_args", ...)`** runs registered patches; backend version detection runs when patches first need it (**`adapter.detect_backend_version()`** via **`_get_backend_version()`**).
+11. **`merge_namespace()`** merges `backend_args` into `module_config.params` (backend wins on conflicts); **`adapter.load_trainer_class(stage)`** resolves the trainer class (default stage `pretrain`).
+12. **`TrainerClass(backend_args=backend_args)`** constructs the trainer.
+13. **`run_patches(phase="setup")`** then **`trainer.setup()`**.
+14. **`trainer.init()`**.
+15. **`run_patches(phase="before_train")`** then **`trainer.train()`** then **`run_patches(phase="after_train")`** then **`trainer.cleanup()`**.
+16. On failure, **`_safe_cleanup()`** calls **`trainer.cleanup(on_error=True)`** when possible.
+
+## 5. Backend system
+
+- **`BackendAdapter`** (`primus/core/backend/backend_adapter.py`) is the abstract integration surface. Subclasses implement **`convert_config()`**, **`load_trainer_class()`**, and **`detect_backend_version()`**. Shared behavior includes **`setup_backend_path()`** and a default **`prepare_backend()`** that runs registered setup hooks.
+- **`BackendRegistry`** (`primus/core/backend/backend_registry.py`) maps backend names to adapter classes, supports **lazy import** of `primus.backends.<backend>`, and stores optional **setup hooks** per backend.
+- **Registered adapters** (via each backend package’s `__init__.py` calling **`BackendRegistry.register_adapter()`**): **`megatron`**, **`torchtitan`**, **`maxtext`**, **`megatron_bridge`**, **`hummingbirdxt`**.
+- Backend code lives under **`primus/backends/<name>/`**. Importing the package registers the adapter and any trainers or hooks that package defines.
+
+## 6. Trainer lifecycle
+
+- **`BaseTrainer`** (`primus/core/trainer/base_trainer.py`) defines the lifecycle. **`setup()`**, **`init()`**, and **`train()`** are **abstract** (subclasses must implement them); **`cleanup(on_error=False)`** is **optional**—it ships a default (no-op) implementation that subclasses may override. The constructor stores **`backend_args`** and reads distributed settings from **`get_torchrun_env()`**.
+- Concrete trainers (for example Megatron or TorchTitan pretrain classes) subclass **`BaseTrainer`** and implement the abstract methods.
+- **`PrimusRuntime`** drives **`setup` → `init` → `train` → `cleanup`**, with patch phases **`build_args`** (before the trainer is created), **`setup`**, and **`before_train`/`after_train`** (around `train`). No patch phase runs around `cleanup`—`after_train` fires before `cleanup` (see §4).
+
+## 7. Patch system
+
+- **`PatchRegistry`** (`primus/core/patches/patch_registry.py`) stores **`FunctionPatch`** objects keyed by backend and phase, with wildcard buckets (`None`) for patches that apply broadly.
+- The **`@register_patch`** decorator registers a patch with metadata (priority, optional version patterns, tags).
+- **`run_patches()`** (`primus/core/patches/patch_runner.py`) collects applicable patches, filters by **`PatchContext`**, sorts by **priority**, and runs handlers. It accepts an optional **`enabled_ids`** list; if omitted, behavior is controlled by **`PRIMUS_PATCHES`**:
+  - unset or **`all`**: all patches
+  - **`none`**: disable all
+  - comma-separated IDs: only those patches
+- **Phases** used by the core runtime include **`build_args`**, **`setup`**, **`before_train`**, and **`after_train`**.
+- Patch implementations are typically colocated with backends under **`primus/backends/<backend>/patches/`**.
+
+## 8. Legacy runtime
+
+The legacy pretrain path—previously selected with **`PRIMUS_TRAIN_RUNTIME=legacy`** and backed by the **`primus/modules/`** stack (**`BaseModule`**-style composition)—has been **removed**. `primus/modules/` no longer contains any source code, and **`primus/cli/subcommands/train.py`** no longer reads `PRIMUS_TRAIN_RUNTIME` or resolves a legacy-vs-core runtime.
+
+All training now runs exclusively through the **core runtime**: both `primus train pretrain` and `primus train posttrain` construct a **`PrimusRuntime`** (**`primus/core/runtime/train_runtime.py`**). **`primus/pretrain.py`** now only provides shared backend-path / environment helpers (for example **`setup_backend_path()`**) used by the training and projection entry points; it no longer defines a `launch_pretrain_from_cli()` legacy launcher.
+
+## 9. Key source files
+
+| Path | Role |
+|------|------|
+| `primus/cli/main.py` | CLI entry, subcommand discovery, dispatch |
+| `primus/cli/subcommands/train.py` | `train` subcommand; chooses core vs legacy pretrain; `posttrain` via `PrimusRuntime` |
+| `primus/core/launcher/parser.py` | **`PrimusParser`**: experiment, platform, and module preset loading |
+| `primus/core/config/preset_loader.py` | **`PresetLoader`**: load framework presets from `primus/configs/` |
+| `primus/core/config/yaml_loader.py` | YAML load with env substitution and `extends` |
+| `primus/core/config/primus_config.py` | **`load_primus_config()`**, **`get_module_config()`**, module normalization |
+| `primus/core/runtime/train_runtime.py` | **`PrimusRuntime`**, **`TrainContext`** |
+| `primus/core/backend/backend_adapter.py` | **`BackendAdapter`** ABC |
+| `primus/core/backend/backend_registry.py` | **`BackendRegistry`** |
+| `primus/core/trainer/base_trainer.py` | **`BaseTrainer`** ABC |
+| `primus/core/patches/patch_registry.py` | **`PatchRegistry`**, **`@register_patch`** |
+| `primus/core/patches/patch_runner.py` | **`run_patches()`**, **`PRIMUS_PATCHES`** parsing |
+| `runner/primus-cli-*.sh` | Shell wrappers for direct, container, and Slurm launch |
+
+For a deep dive on the CLI internals (subcommand discovery, dispatch, and the launch wrappers), see [CLI Architecture](cli-architecture.md). For day-to-day contribution workflows (style, tests, CI), see [Contributing Guide](contributing.md) and [Testing Guide](testing.md).
diff --git a/docs/06-developer-guide/backend-patch-notes.md b/docs/06-developer-guide/backend-patch-notes.md
new file mode 100644
index 000000000..da8a68ad1
--- /dev/null
+++ b/docs/06-developer-guide/backend-patch-notes.md
@@ -0,0 +1,143 @@
+# Backend patch notes
+
+Primus integrates several large-model backends (Megatron-LM, TorchTitan, JAX MaxText, …) and applies a lightweight patch layer to keep configuration flags consistent with the Primus CLI. This page captures those backend-specific switches so they live alongside the rest of the documentation (instead of the historical `primus/README_patch.md` file).
+
+## How to read these notes
+
+- Start with the **Base Module Parameters** table below—every backend module inherits these knobs.
+- Jump to the backend-specific section for details on extra CLI/config options and links to the patched source files.
+- When editing configs or CLI presets, cross-reference the [Primus CLI Reference](../02-user-guide/cli-reference.md) so command examples and backend parameters stay in sync.
+
+## Supported models
+
+This section lists, at a high level, the model families Primus currently targets on each backend. For more details and configuration examples, refer to the backend-specific patch notes below.
+
+### Megatron-LM
+
+- **LLaMA family**: LLaMA2, LLaMA3, LLaMA3.1, LLaMA3.3, LLaMA4 (various sizes from 7B up to 405B+)
+- **DeepSeek family**: DeepSeek-V2 (lite/base/full) and DeepSeek-V3
+- **MoE / Mixtral**: Mixtral-8x7B / 8x22B, large MoE configs (515B, 1T, 2T, 4T) and DeepSeek-style MoE
+- **Qwen family**: Qwen2.5 (7B/72B) and Qwen3 (8B/30B/235B variants)
+- **Other GPT-style models**: Grok1/2, GPT-OSS 20B and generic `language_model.yaml`
+
+### TorchTitan
+
+- **LLaMA family**: LLaMA3, LLaMA3.1, LLaMA3.3 (various sizes, including FP8 variants)
+- **DeepSeek family**: DeepSeek-V3 (16B and 671B, FP8 and BF16 configs)
+- **Qwen family**: Qwen3 small/medium models (0.6B, 1.7B, 32B)
+
+### JAX MaxText
+
+- **LLaMA family**: LLaMA2 (7B/70B), LLaMA3 (8B/70B), LLaMA3.3 (70B)
+- **DeepSeek family**: DeepSeek-V2 16B
+- **MoE / Mixtral**: Mixtral-8x7B
+- **Other models**: Grok1 and additional MaxText-supported transformers (see MaxText docs for the full list)
+
+## Base module parameters
+
+All modules inherit the options defined in [`primus/configs/modules/module_base.yaml`](https://github.com/AMD-AGI/Primus/blob/main/primus/configs/modules/module_base.yaml):
+
+| Argument Name       | Default Value | Description                                                                                |
+| ------------------- | ------------- | ------------------------------------------------------------------------------------------ |
+| `trainable`         | `false`       | Whether the module participates in training.                                              |
+| `sink_level`        | `null`        | Global sink level for logging. Overrides `file_sink_level` and `stderr_sink_level` if set. |
+| `file_sink_level`   | `DEBUG`       | Logging level for file sink (log files).                                                  |
+| `stderr_sink_level` | `INFO`        | Logging level for stderr/console output.                                                  |
+
+### Backend index
+
+- [Megatron-LM patch notes](#megatron-lm-patch-notes)
+- [TorchTitan patch notes](#torchtitan-patch-notes)
+- [JAX MaxText patch notes](#jax-maxtext-patch-notes)
+
+---
+
+## Megatron-LM patch notes
+
+Primus keeps a curated patch layer on top of upstream Megatron-LM so CLI presets and configs can expose additional controls. Use this section with the [Base Module Parameters](#base-module-parameters) above for shared module parameters, and the [Primus CLI Reference](../02-user-guide/cli-reference.md) for CLI/config usage patterns.
+
+> ℹ️ The **Version** column maps to Primus internal patch tags (v0.x.y) so you know when a flag landed.
+
+### 1. Module-level parameters
+
+These arguments are introduced in the Megatron module logic (e.g., training loop, logging, resume logic). They are defined via patching and can be configured to control training behavior and logging utilities.
+
+| New Argument                         | Default Value | Version | Description                                                                                    | Patched Files                                                                                                                                                                                                                                                                                                | Notes                                            |
+| ------------------------------------ | ------------- | ------- | ---------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------ |
+| `disable_tensorboard`                | `true`        | v0.1.0  | Whether to disable TensorBoard. Set to `false` if you want to enable profiling or torch trace. | NA                                                                                                                                                                                                                                                                                                           | Required for timeline and performance debugging. |
+| `disable_wandb`                      | `true`        | v0.1.0  | Whether to disable Weights & Biases logging.                                                   | NA                                                                                                                                                                                                                                                                                                           | Useful for internal benchmarking.                |
+| `disable_compile_dependencies`       | `true`        | v0.1.0  | Disables Megatron’s custom kernel compilation. Most ops are already covered by TE.             | NA                                                                                                                                                                                                                                                                                                           | Avoids redundant compilation steps.              |
+| `auto_continue_train`                | `false`       | v0.1.0  | Automatically resume training from the latest checkpoint if found in the `--save` path.        | NA                                                                                                                                                                                                                                                                                                           | Simplifies job restarts.                         |
+| `disable_last_saving`                | `false`       | v0.1.0  | Skip saving the final checkpoint at the last iteration.                                        | NA                                                                                                                                                                                                                                                                                                           | Useful for profiling or benchmarking runs.       |
+| `no_fp8_weight_transpose_cache`      | `false`       | v0.2.0  | Disable the FP8 weight transpose cache to save memory.                                         | `megatron.core.extensions.transformer_engine.TELinear`, `megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear`, `megatron.core.extensions.transformer_engine.TEDelayedScaling`                                                                                                        | May affect performance but reduce memory use.    |
+| `decoder_pipeline_manual_split_list` | `null`        | v0.2.0  | Enable manual pipeline split in (interleaved) 1F1B pipeline parallelism.                       | `megatron.core.transformer.transformer_block.get_num_layers_to_build`, `megatron.core.transformer.transformer_layer.get_transformer_layer_offset`                                                                                                                                                            | Deprecated. Use `pipeline_model_parallel_layout` instead.    |
+| `pp_warmup`                          | `false`       | v0.2.0  | Add fwd/bwd warmup to save iter1's time when pp degree is large.                             | NA                                                                                                                                                                                                                                                                                                           | Can save much time for pipeline debug.           |
+| `dump_pp_data`                       | `false`       | v0.2.0  | Enable dumping pp schedule data for visualization.                                             | `megatron.core.pipeline_parallel.schedules.forward_step`, `megatron.core.pipeline_parallel.schedules.backward_step`, `megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_with_interleaving`, `megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_without_interleaving` | Useful for pipeline schedule visualization.      |
+| `disable_profiler_activity_cpu`                | `false`       | v0.2.0  | Disable CPU activity in torch profiling.                                        |  NA                                                                                                                                                                                                                                                                                                         | If you only want to trace CUDA kernels and get a smaller trace JSON file, you can enable this option. However, if you plan to run with TraceLen, please do not enable it.  more torch profiler args: <br>`torch_profiler_record_shapes: true`, <br>`torch_profiler_with_stack: true`, <br>`torch_profiler_use_gzip: true`      |
+| `use_rocm_mem_info`                        | `false`       | v0.2.0  | Logging ROCm memory information in Megatron-LM Trainer                             | NA                                                                                                                                                                                                                                                                                                           | If `use_rocm_mem_info = True`, ROCm memory information will be collected with `rocm-smi` at every iteration.           |
+| `use_rocm_mem_info_iters`                        | `[1,2]`       | v0.2.0  | Logging ROCm memory information in Megatron-LM Trainer for some iterations                            | NA                                                                                                                                                                                                                                                                                                           | If `use_rocm_mem_info = False`, ROCm memory information will be collected at the iterations specified in `use_rocm_mem_info_iters`.           |
+| `patch_zero_bubble`                        | `false`       | v0.2.0  | Using Zero-Bubble pipeline parallism  | `megatron.core.optimizer.ChainedOptimizer`, `megatron.core.pipeline_parallel.get_forward_backward_func`, `megatron.core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication`, `megatron.core.parallel_stat.default_embedding_ranks`, `megatron.core.parallel_stat.is_pipeline_last_stage`, `megatron.core.parallel_stat.is_rank_in_embedding_group`, `megatron.core.distributed.finalize_model_grads`, `megatron.core.transformer.transformer_layer.get_transformer_layer_offset` | If `patch_zero_bubble = True`, Zero bubble pipeline parallism will be enable to use. See more detail at [ZeroBubble User Guide](../../primus/backends/megatron/core/pipeline_parallel/zerobubble/README.md)           |
+| `disable_mlflow`                        | `true`       | v0.3.0  | Track model development using MLflow                             | NA                                                                                                                                                                                                                                                                                                         |  Envs: <br> `export DATABRICKS_TOKEN=your_token`<br>`export DATABRICKS_HOST=your_host`<br>`export MLFLOW_TRACKING_URI=databricks`<br>`export MLFLOW_REGISTRY_URI=databricks-uc`<br>Arguments: <br>`mlflow_run_name: null`,<br>`mlflow_experiment_name: null`     |
+| `recompute_layer_ids`                    | `null`       | v0.4.0  | Specify the exact IDs of layers to recompute, enabling more flexible memory reduction                            | NA                                                                                                                                                                                                                                                                                                           | Using `recompute_layer_ids=[layer_id_0, layer_id_1,...]` together with `recompute_granularity=full`, where layer_id ranges from 0 to num_layers - 1.           |
+| `dataloader_mp_context`                  | `null`       | v0.5.0  | Set `DataLoader.multiprocessing_context` to avoid SIGSEGV caused by fork()-hostile native state (RDMA MRs, HIP runtime, IPC handles). | `torch.utils.data.DataLoader.__init__`                                                                                                                                                                                                                                                                       | Accepted values: `"forkserver"`, `"spawn"`, `"fork"`, or `null` (keep PyTorch default). Only takes effect when `num_workers > 0` and no explicit `multiprocessing_context` is passed. |
+
+---
+
+### 2. Model-definition parameters
+
+These arguments affect the internal architecture or layer implementations. They are patched into the model construction logic and used for tuning or debugging specific variants.
+
+| New Argument                        | Default Value | Version | Description                                                               | Patched Files                                                                                                                                                                                                                                                                                                                   | Notes                                       |
+| ----------------------------------- | ------------- | ------- | ------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------- |
+| `disable_primus_topk_router`   | `false`       | v0.1.0  | Disable PrimusTopkRouter and use TopkRouter implemented by megatron. | `megatron.core.transformer.moe.router.TopKRouter`                                                                                                                                                                                                                                                                               | Used to debug internal.         |
+| `moe_router_force_load_balancing`   | `false`       | v0.1.0  | Force token redistribution in MoE to achieve load balance across experts. | `megatron.core.transformer.moe.router.TopKRouter`                                                                                                                                                                                                                                                                               | Use to debug MoE imbalance issues.          |
+| `use_deprecated_20241209_moe_layer` | `false`       | v0.1.0  | Enable legacy MoE implementation for debugging/perf comparison.           | `megatron.core.transformer.moe.moe_layer.MoELayer`, `megatron.core.transformer.moe.moe_layer.MoESubmodules`, `megatron.core.transformer.moe.experts.GroupedMLP`, `megatron.core.transformer.moe.experts.SequentialMLP`, `megatron.core.transformer.moe.experts.TEGroupedMLP`, `megatron.core.transformer.moe.router.TopKRouter` | Deprecated, used for internal testing only. |
+| `moe_permute_fusion`   | `false`       | v0.1.0  | Permutation and unpermutation fusion. | `megatron.core.extensions.transformer_engine`, `megatron.core.transformer.moe.moe_utils`                                                                                                                                                                                                                                                                               | Fuse permutation and unpermutation in moe layer.         |
+| `moe_use_fused_router_with_aux_score`   | `false`       | v0.2.0  | Fused router topk and calculation of moe aux loss score. Need Primus turbo backend | `megatron.core.transformer.moe.router.TopKRouter`                                                                                                                                                                                                                                                                               | Used to reduce launch overhead of the small kernels in router.         |
+
+### 3. Primus-Turbo related options
+
+| New Argument                        | Default Value | Version | Description                                                               | Patched Files                                                                                                                                                                                                                                                                                                                   | Notes                                       |
+| ----------------------------------- | ------------- | ------- | ------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------- |
+| `use_turbo_gemm`   | `false`       | v0.8.0  | Use Primus-Turbo linear modules (`PrimusTurboLinear`, `PrimusTurboColumnParallelLinear`, `PrimusTurboRowParallelLinear`, `PrimusTurboLayerNormColumnParallelLinear`) in place of the TE linear modules. | `megatron.core.extensions.transformer_engine.TELinear`, `megatron.core.extensions.transformer_engine.TEColumnParallelLinear`, `megatron.core.extensions.transformer_engine.TERowParallelLinear`, `megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear` | Accelerates dense GEMMs. Supports FP8 recipes (`tensorwise`, `blockwise`, `mxfp8`) and FP4 recipe (`mxfp4`). Replaces the deprecated `use_turbo_parallel_linear`. **Please set `enable_primus_turbo=True` first.** |
+| `use_turbo_grouped_gemm`   | `false`       | v0.8.0  | Use Primus-Turbo grouped GEMM (`PrimusGroupedMLP` with `PrimusTurboColumnParallelGroupedLinear` / `PrimusTurboRowParallelGroupedLinear`) for MoE experts in place of `TEGroupedMLP`. | `megatron.core.transformer.moe.experts.TEGroupedMLP`, `megatron.core.extensions.transformer_engine.TEColumnParallelGroupedLinear`, `megatron.core.extensions.transformer_engine.TERowParallelGroupedLinear` | Accelerates MoE grouped GEMMs. Incompatible with `moe_use_legacy_grouped_gemm=True`. Required by Sync-Free MoE stage 2/3. Replaces the deprecated `use_turbo_grouped_mlp`. **Please set `enable_primus_turbo=True` first.** |
+| `use_turbo_permute_padding`   | `false`       | v0.8.0  | Pad tokens of every experts to 16 or 32 multiple to reduce d2h and h2d. | `megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher`, `megatron.core.transformer.moe.experts.TEGroupedMLP` | Only effective under FP8/FP4 with `use_turbo_deepep=True`. Pad multiple is 16/32 (FP8, depending on recipe) or 32 (FP4). **Please set `enable_primus_turbo=True` first.** |
+| `use_turbo_deepep`   | `false`       | v0.4.0  | Use Primus-turbo `DeepEPTokenDispatcher`. | `megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher`                                                                                                                                                                                                                                                              | Used Primus-Turbo DeepEP to accelerate MoE token dispatcher.  **You must both set`enable_primus_turbo=True` and `use_turbo_deepep=True` to enable this function.**         |
+| `turbo_deepep_num_cu`   | `32`       | v0.4.0  | Set the number of CUs to use for Primus-Turbo DeepEP. |   | 64 or 80 for ep8, 32 for ep16-64 is best practice.  |
+| `turbo_deepep_use_comm_stream`   | `false`       | v0.4.0  | Primus-Turbo DeepEP will use an internal stream to dispatch/combine when enabled, default used `current_stream` |   |  **Please both set`enable_primus_turbo=True` and `use_turbo_deepep=True` first**
+| `turbo_sync_free_moe_stage`   | `0`       | v0.4.0  | Primus Sync-Free MoE has 4 stages. See [RFC: Primus-Megatron SyncFree MoE](https://github.com/AMD-AGI/Primus/issues/203) for more details. |   |   stage 2 is recommended for better performance. **Please set`enable_primus_turbo=True` first**   |
+| `use_turbo_mega_moe`   | `false`       | v0.8.0  | Replace the whole MoE layer with the FlyDSL-based fused Primus-Turbo MegaMoE layer (fused router + `dispatch_grouped_gemm` → SwiGLU → `grouped_gemm_combine`). | `megatron.core.transformer.moe.moe_layer.MoELayer` | EP-only: requires `tensor_model_parallel_size=1`, `params_dtype=bf16`, and an EP process group; asserts on unsupported router options (sequence/global aux loss, z-loss, sinkhorn, input jitter, expert bias). Needs a Primus-Turbo build with MegaMoE. See [MegaMoE guide](../04-technical-guides/mega-moe.md). **Please set `enable_primus_turbo=True` first.** |
+
+---
+
+## TorchTitan patch notes
+
+TorchTitan integration uses the same Primus configuration surface (CLI flags + YAML) but exposes a few extra knobs via patches. Pair this with the [Base Module Parameters](#base-module-parameters) above for shared module parameters.
+
+| New Argument | Default Value | Version | Description | Patched Files | Notes |
+| ------------ | ------------- | ------- | ----------- | ------------- | ----- |
+| `primus_turbo.enable_embedding_autocast` | `true` | v0.4.0 | Automatically casts `nn.Embedding` outputs to the AMP dtype (e.g., bf16) during training so downstream layers stay in sync. | (Primus TorchTitan patch set) | Disable only if you manage casting manually. |
+
+---
+
+## JAX MaxText patch notes
+
+Primus integrates JAX MaxText as a backend for running LLaMA and related transformer models on AMD GPUs. At the moment, Primus does not apply any additional patch-layer arguments on top of MaxText—the MaxText configuration surface (YAML + CLI) is used as-is.
+
+Use this section together with:
+
+- The [Base Module Parameters](#base-module-parameters) and [Supported Models](#supported-models) above for a high-level model overview
+- The [Primus CLI Reference](../02-user-guide/cli-reference.md) for Primus CLI usage patterns
+- The official MaxText documentation for the full set of MaxText-specific arguments
+
+### MaxText-specific notes
+
+- Primus currently wires MaxText via `primus/configs/models/maxtext` and `primus/configs/modules/maxtext`.
+- Model families currently exercised in examples include:
+  - LLaMA2 7B/70B
+  - LLaMA3 8B/70B
+  - LLaMA3.3 70B
+  - DeepSeek-V2 16B
+  - Mixtral-8x7B
+  - Grok1
+- There are no extra Primus-only flags for MaxText yet; as we add MaxText-specific patches (e.g., ROCm optimizations, logging helpers), they will be documented in tables here in the same style as the Megatron-LM and TorchTitan patch notes.
diff --git a/docs/cli/CLI-ARCHITECTURE.md b/docs/06-developer-guide/cli-architecture.md
similarity index 93%
rename from docs/cli/CLI-ARCHITECTURE.md
rename to docs/06-developer-guide/cli-architecture.md
index 944988aa3..8dd86e78f 100644
--- a/docs/cli/CLI-ARCHITECTURE.md
+++ b/docs/06-developer-guide/cli-architecture.md
@@ -1,4 +1,4 @@
-# 🚀 From Chaos to Order: Building a Unified Entry Point for AMD GPU LLM Training
+# 🚀 From chaos to order: Building a unified entry point for AMD GPU LLM training
 
 > ⚠️ **NOTE**: This is a draft version and not the final release.
 >
@@ -8,7 +8,7 @@
 
 ---
 
-## 📖 The Beginning: Pain Points in Training Workflows
+## 📖 The beginning: Pain points in training workflows
 
 Imagine this scenario:
 
@@ -41,7 +41,7 @@ The traditional approach is to use a large number of Bash scripts to handle thes
 
 ---
 
-## 💡 Design Philosophy: One Command, Done
+## 💡 Design philosophy: One command, done
 
 Our core philosophy is simple: **One command, from environment configuration to training launch, fully automated.**
 
@@ -50,7 +50,7 @@ Our core philosophy is simple: **One command, from environment configuration to
 primus-cli direct -- train pretrain --config deepseek_v2.yaml
 ```
 
-### 🏗️ Three-Layer Architecture Design
+### 🏗️ Three-layer architecture design
 
 Primus CLI adopts a clear **three-layer structure + plugin system**:
 
@@ -74,7 +74,7 @@ Primus CLI adopts a clear **three-layer structure + plugin system**:
 └─────────────────────────────────────────────────────┘
 ```
 
-### 🎯 Four Design Goals
+### 🎯 Four design goals
 
 | Goal | Implementation | User Benefits |
 |------|---------------|---------------|
@@ -85,9 +85,9 @@ Primus CLI adopts a clear **three-layer structure + plugin system**:
 
 ---
 
-## 🔍 Deep Dive: Architecture Dissection
+## 🔍 Deep dive: Architecture dissection
 
-### ⚙️ Layer 1: Intelligent Runtime Abstraction
+### ⚙️ Layer 1: Intelligent runtime abstraction
 
 Different scenarios require different runtime environments, but users shouldn't have to worry about these details. Primus CLI provides three seamlessly switchable runtime modes:
 
@@ -114,7 +114,7 @@ primus-cli slurm srun -N 8 -- benchmark gemm -M 4096
 
 ---
 
-### 🔁 Layer 2: Hook & Patch System
+### 🔁 Layer 2: Hook and patch system
 
 Training is more than just running a Python script. You might need to:
 - 🗂️ Preprocess datasets before training
@@ -147,7 +147,7 @@ This is especially useful when you need to quickly apply temporary fixes or make
 
 ---
 
-### 🧩 Layer 3: Task Execution Layer
+### 🧩 Layer 3: Task execution layer
 
 This layer is responsible for executing specific business logic—training, testing, environment checks, and other actual tasks. Remember we said "zero-intrusion extension"? How is this achieved?
 
@@ -185,15 +185,15 @@ This plugin-based design allows Primus CLI to quickly respond to new requirement
 
 ---
 
-## 🌐 The Magic Behind: Intelligent Environment Detection
+## 🌐 The magic behind: Intelligent environment detection
 
 This is probably the most "black tech" part of Primus CLI.
 
-### Problem: Different GPUs Need Different Configurations
+### Problem: Different GPUs need different configurations
 
 AMD's GPU family is rich: MI300X, MI250X, MI210... Each GPU has its optimal ROCm configuration and environment variable settings. The traditional approach is to let users manually select configurations, but this is both error-prone and insufficiently automated.
 
-### Solution: Three-Step Auto-Configuration
+### Solution: Three-step auto-configuration
 
 **Step 1: Load Common Environment**
 
@@ -226,7 +226,7 @@ Now, `MI300X.sh` can contain all best practices for this GPU model:
 
 **Users don't need to worry about these details at all - everything is automatic.**
 
-### Real-World Example
+### Real-world example
 
 ```bash
 # On MI300X cluster
@@ -242,7 +242,7 @@ primus-cli direct -- train pretrain --config config.yaml
 
 ---
 
-## 🧪 Foundation of Scientific Experiments: Reproducibility
+## 🧪 Foundation of scientific experiments: Reproducibility
 
 In machine learning research, reproducibility is crucial. But reality is harsh:
 
@@ -250,7 +250,7 @@ In machine learning research, reproducibility is crucial. But reality is harsh:
 
 Does this sound familiar? Primus CLI completely solves this problem with an **automated snapshot mechanism**.
 
-### Auto-Record Everything
+### Auto-record everything
 
 Every time training starts, Primus CLI automatically saves the complete runtime context:
 
@@ -270,7 +270,7 @@ output/exp_2025_11_10_134522/
 └── metadata.json                # Runtime metadata
 ```
 
-### One-Click Reproduction
+### One-click reproduction
 
 Three months later, when you want to reproduce this experiment:
 
@@ -285,7 +285,7 @@ Primus CLI will automatically:
 3. Verify GPU and system environment
 4. Start training (if environment is compatible)
 
-### Real-World Value
+### Real-world value
 
 | Scenario | Traditional Approach | Using Primus CLI |
 |----------|---------------------|------------------|
@@ -297,11 +297,11 @@ Primus CLI will automatically:
 
 ---
 
-## 📊 Real-World Case: From Development to Production
+## 📊 Real-world case: From development to production
 
 Let's see how Primus CLI simplifies the entire workflow through a real scenario.
 
-### Scenario: Training DeepSeek-V2 Model
+### Scenario: Training DeepSeek-V2 model
 
 **Step 1: Local Development & Validation** 🖥️
 
@@ -346,7 +346,7 @@ primus-cli slurm sbatch \
   -- train pretrain --config configs/deepseek_v2_prod.yaml
 ```
 
-### Key Insight
+### Key insight
 
 Notice? **From development to production, the core command structure remains unchanged**:
 ```
@@ -357,7 +357,7 @@ Only the runtime environment (`direct` → `container` → `slurm`) changes - ev
 
 ---
 
-## 🎯 Core Advantages Summary
+## 🎯 Core advantages summary
 
 After the detailed introduction above, let's summarize the core value Primus CLI brings:
 
@@ -373,11 +373,11 @@ After the detailed introduction above, let's summarize the core value Primus CLI
 
 ---
 
-## 🛣️ Future Roadmap
+## 🛣️ Future roadmap
 
 Primus CLI continues to evolve, and our near-term plans include:
 
-### Short-Term Goals (2025)
+### Short-term goals (2025)
 - 🎯 **Python Hook API**: Support writing Hooks in Python scripts for more flexible extension capabilities
 - 🎯 **Intelligent Preflight**: Auto-check GPU health, network topology, InfiniBand connectivity before launch
 - 🎯 **Configuration Template System**: Built-in best practice config templates for common models
@@ -386,12 +386,12 @@ Primus CLI continues to evolve, and our near-term plans include:
 - 🎯 **Extended Framework Support**: Improve support for more training frameworks like TorchTitan, JAX/Flax
 - 🎯 **CI/CD Integration**: Provide standardized testing and validation workflows, support automated regression testing
 
-### Long-Term Vision
+### Long-term vision
 - 🌟 Become the **standard training entry point for the ROCm ecosystem**
 
 ---
 
-## 🎓 Summary: The Power of One Command
+## 🎓 Summary: The power of one command
 
 Back to the question at the beginning: How do we make large model training go from complex to simple?
 
@@ -415,9 +415,10 @@ primus-cli direct -- train pretrain --config deepseek_v2.yaml
 
 ---
 
-## 📚 Learn More
+## 📚 Learn more
 
-- 📖 **User Guide**: [PRIMUS-CLI-GUIDE.md](./PRIMUS-CLI-GUIDE.md)
+- 📖 **CLI Reference (user guide)**: [cli-reference.md](../02-user-guide/cli-reference.md)
+- 🏛 **System Architecture**: [architecture.md](./architecture.md)
 - 🔧 **Quick Start**: `primus-cli --help`
 - 💬 **Issue Reporting**: GitHub Issues
 - 🌐 **ROCm Ecosystem**: [rocm.github.io](https://rocm.github.io)
diff --git a/docs/06-developer-guide/contributing.md b/docs/06-developer-guide/contributing.md
new file mode 100644
index 000000000..9cc07f781
--- /dev/null
+++ b/docs/06-developer-guide/contributing.md
@@ -0,0 +1,142 @@
+# Contributing guide
+
+This guide summarizes how to set up a development environment, follow project conventions, run checks locally, and align with the CI pipeline. For test commands and layout, see [Testing Guide](testing.md). The repository root [CONTRIBUTING.md](../../CONTRIBUTING.md) summarizes key contribution guidelines, such as branch naming conventions, commit message style, and pull request requirements.
+
+## 1. Development setup
+
+1. **Clone the repository** (include submodules):
+
+   ```bash
+   git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
+   cd Primus
+   ```
+
+2. **Install Python dependencies:**
+
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. **Install pre-commit hooks** (recommended):
+
+   ```bash
+   pip install pre-commit
+   pre-commit install
+   ```
+
+4. **Optional—JAX / MaxText work:**
+
+   ```bash
+   pip install -r requirements-jax.txt
+   ```
+
+5. **Quick verification** (from the repository root, with Primus on your `PATH` or via the bundled launcher):
+
+   ```bash
+   ./primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096
+   ```
+
+## 2. Code style
+
+Configuration lives in `.pre-commit-config.yaml`. Hooks run automatically on `git commit` after `pre-commit install`.
+
+| Tool | Version | Purpose |
+|------|---------|---------|
+| **black** | 24.8.0 | Python formatter, line length 110 |
+| **isort** | 5.13.2 | Import sorting, profile `black` |
+| **autoflake** | 2.3.1 | Removes unused imports and variables (see hook args for star imports and `__init__`) |
+| **shellcheck** | 0.10.0.1 (shellcheck-py) | Shell script analysis |
+| **pre-commit-hooks** | v4.0.1 | `trailing-whitespace`, `end-of-file-fixer`, `check-yaml`, `check-added-large-files`, `check-merge-conflict` |
+
+Manual one-off runs (repository root):
+
+```bash
+black --line-length=110 .
+isort --profile black .
+autoflake --remove-all-unused-imports --remove-unused-variables --expand-star-imports --ignore-init-module-imports --recursive --in-place .
+```
+
+CI runs `pre-commit run --all-files --show-diff-on-failure`, so lint behavior follows `.pre-commit-config.yaml` rather than a separate hand-written list of formatter commands.
+
+## 3. Branch naming convention
+
+Format:
+
+```text
+<type>/<scope>/<short-description>
+```
+
+**Types:** `feat`, `fix`, `docs`, `refactor`, `test`, `chore`, `ci`
+
+**Scope (optional):** `engine`, `model`, `scheduler`, `docs`, `tests`, `config`, or another short area name.
+
+**Examples:**
+
+- `feat/model/implement-moe-routing`
+- `fix/engine/init-error`
+
+## 4. Commit convention
+
+Follow [Conventional Commits](https://www.conventionalcommits.org/):
+
+```text
+<type>(<scope>): <description>
+```
+
+**Examples:**
+
+- `feat(model): add MOE routing functionality`
+- `fix(engine): resolve initialization error`
+
+## 5. Testing requirements
+
+Before opening a pull request, run the following from the repository root:
+
+- **Shell integration tests:**
+
+  ```bash
+  bash ./tests/runner/run_all_tests.sh
+  ```
+
+- **Python unit tests:**
+
+  ```bash
+  pytest tests/unit_tests/ --maxfail=1 -s
+  ```
+
+- **Backend / trainer tests** (GPU, datasets, and sometimes Hugging Face tokens): run the relevant file under `tests/trainer/` when your change touches that backend. See [Testing Guide](testing.md).
+
+- **Pre-commit on all files:**
+
+  ```bash
+  pre-commit run --all-files
+  ```
+
+## 6. Pull request process
+
+1. Fork the repository (unless you have write access and use a feature branch).
+2. Create a branch that follows the naming convention above.
+3. Implement changes and commit using the commit message convention.
+4. Run tests and pre-commit locally.
+5. Push and open a pull request with a clear description.
+6. Reference related issues when applicable.
+7. Request reviewers.
+8. Address review feedback.
+9. Ensure CI passes (lint and unit tests on the paths your PR triggers).
+
+## 7. CI pipeline
+
+The workflow **`.github/workflows/ci.yaml`** defines how changes are validated.
+
+**Triggers:** `workflow_dispatch`, pushes to `main`, tags matching `v*`, and pull requests.
+
+**Jobs (high level):**
+
+- **`code-lint`:** Ubuntu, Python 3.12—installs `pre-commit` and runs `pre-commit run --all-files --show-diff-on-failure`, so the checks (black, isort, autoflake, shellcheck, and the `pre-commit-hooks` set) follow `.pre-commit-config.yaml` exactly.
+- **`build-docker`:** Builds and pushes Docker images (depends on `code-lint`).
+- **`run-unittest-torch`:** Self-hosted GPU runner—installs dependencies (including Primus-Turbo and AITER as defined in the workflow), runs `bash ./tests/runner/run_all_tests.sh`, `pytest` on `tests/unit_tests/` (with a few deselected tests), then Megatron and TorchTitan trainer tests with `DATA_PATH`, `MASTER_PORT`, `HSA_NO_SCRATCH_RECLAIM=1`, and `HF_TOKEN` where required.
+- **`run-unittest-jax`:** JAX runner—installs `requirements-jax.txt`, runs shell tests and `python ./tests/run_unit_tests.py --jax` with CI-specific environment variables.
+
+Lint checks mirror the pre-commit stack. Trainer jobs require GPU resources and shared secrets (for example `HF_TOKEN`) in the hosted environment.
+
+For a focused description of local vs CI test commands, see [Testing Guide](testing.md).
diff --git a/docs/06-developer-guide/extending-backends.md b/docs/06-developer-guide/extending-backends.md
new file mode 100644
index 000000000..cc42151fc
--- /dev/null
+++ b/docs/06-developer-guide/extending-backends.md
@@ -0,0 +1,419 @@
+# Extending backends
+
+This guide explains how to add a **new training backend** to Primus using the current runtime architecture. It complements the high-level picture in [Primus overview](../01-getting-started/overview.md): adapters sit under the unified CLI and configuration system ([Configuration system](../02-user-guide/configuration-system.md)), and each backend plugs in through the same lifecycle and hook points as Megatron-LM, TorchTitan, MaxText, and the other integrated stacks.
+
+The runtime is built around:
+
+- **`BackendAdapter`** – integrates a backend framework
+- **`BackendRegistry`** – discovers and instantiates adapters
+- **`BaseTrainer`** – defines the minimal training lifecycle that all backends follow
+- **`PrimusRuntime`** – orchestrates config loading, environment setup, patches, adapter, and trainer
+
+The examples below use a minimal **`dummy`** backend as a template. The dummy files are illustrative and are not checked into this repository; existing backends such as Megatron, TorchTitan, MaxText, Megatron Bridge, and HummingbirdXT show the production pattern.
+
+---
+
+## What happens when you run Primus?
+
+When you run:
+
+```bash
+primus train pretrain --config <exp.yaml>
+```
+
+the runtime (`PrimusRuntime`) does roughly:
+
+1. Load the experiment config—`load_primus_config()` returns a lightweight `SimpleNamespace` (not a `PrimusConfig`)—and select the `module_config`
+2. Apply CLI overrides to `module_config.params`
+3. Initialize environment (HF, logging, distributed environment, data directory)
+4. Resolve backend adapter via `BackendRegistry.get_adapter(framework)`
+5. Call `adapter.setup_backend_path(...)` to put the backend on `sys.path`
+6. Call `adapter.prepare_backend(module_config)` (usually runs backend setup hooks)
+7. Build backend arguments:
+
+   ```python
+   backend_args = adapter.convert_config(module_config.params)
+   # run "build_args" patches and merge back into module_config.params
+   ```
+
+8. Load and construct the trainer:
+
+   ```python
+   TrainerClass = adapter.load_trainer_class(stage=module_config.params.stage or "pretrain")
+   trainer = TrainerClass(backend_args=backend_args)
+   ```
+
+9. Execute the trainer lifecycle (with patches around it). Backend version detection is lazy during patch handling through `adapter.detect_backend_version()` rather than a separate pre-trainer step:
+
+    ```python
+    # PrimusRuntime (the "build_args" patches from step 7 already ran
+    # before the trainer was constructed):
+    run_patches(phase="setup",        backend_args=backend_args)
+    trainer.setup()
+
+    trainer.init()
+
+    run_patches(phase="before_train", backend_args=backend_args)
+    trainer.train()
+    run_patches(phase="after_train",  backend_args=backend_args)
+
+    trainer.cleanup()
+    ```
+
+So a complete backend must provide:
+
+- An **adapter** subclassing `BackendAdapter`
+- A **trainer** subclassing `BaseTrainer` and implementing `setup`, `init`, and `train` (and optionally overriding `cleanup`, which has a default no-op implementation)
+- A small `primus.backends.<name>.__init__` that calls `BackendRegistry.register_adapter(...)`
+
+---
+
+## Minimal backend layout
+
+Create a new backend folder under `primus/backends/`:
+
+```text
+primus/backends/dummy/
+├── __init__.py
+├── dummy_adapter.py
+└── dummy_pretrain_trainer.py
+```
+
+This mirrors the pattern used by existing backends (for example Megatron, TorchTitan).
+
+---
+
+## Implement the adapter (`BackendAdapter`)
+
+**File**: `primus/backends/dummy/dummy_adapter.py`
+
+```python
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+from primus.core.backend.backend_adapter import BackendAdapter
+from primus.core.backend.backend_registry import BackendRegistry
+from primus.core.utils.module_utils import log_rank_0
+
+
+class DummyAdapter(BackendAdapter):
+    """Minimal adapter for a 'dummy' backend."""
+
+    def __init__(self, framework: str = "dummy"):
+        super().__init__(framework)
+
+    def setup_backend_path(self, backend_path=None) -> str:
+        """
+        Dummy backend lives inside the Primus tree (no third_party submodule),
+        so we don't need to modify sys.path or resolve any external path here.
+
+        For real backends that live under third_party/<backend>, you can rely on
+        the default implementation in BackendAdapter instead.
+        """
+        log_rank_0("[Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend")
+        return ""
+
+    def convert_config(self, params: Any) -> Any:
+        """
+        Convert Primus module params → backend-specific args.
+
+        For a real backend you would build a structured args object. Here we
+        just wrap the incoming params in a SimpleNamespace.
+        """
+        if isinstance(params, dict):
+            backend_args = SimpleNamespace(**params)
+        else:
+            backend_args = params
+        log_rank_0("[Primus:DummyAdapter] Converted Primus params -> dummy backend_args")
+        return backend_args
+
+    def detect_backend_version(self) -> str:
+        """Return a version string used by patch filtering."""
+        return "dummy-0.1"
+
+    def load_trainer_class(self, stage: str = "pretrain"):
+        """Return the Trainer class for the specified training stage."""
+        from primus.backends.dummy.dummy_pretrain_trainer import DummyPretrainTrainer
+
+        log_rank_0("[Primus:DummyAdapter] Loaded trainer class: DummyPretrainTrainer")
+        return DummyPretrainTrainer
+```
+
+Key points:
+
+- Since the dummy backend is implemented directly under `primus.backends.dummy` (not in `third_party/`), it overrides `setup_backend_path()` as a **no-op** so that the default third-party path resolution is skipped.
+- `convert_config()` returns whatever your trainer expects as `backend_args`.
+- `load_trainer_class()` imports and returns `DummyPretrainTrainer` directly (similar to `MegatronAdapter`), without going through a registry lookup.
+
+---
+
+## Implement a runnable trainer (`BaseTrainer`)
+
+**File**: `primus/backends/dummy/dummy_pretrain_trainer.py`
+
+```python
+from typing import Any
+
+from primus.core.trainer.base_trainer import BaseTrainer
+from primus.core.utils.module_utils import log_rank_0
+
+
+class DummyPretrainTrainer(BaseTrainer):
+    """Minimal runnable trainer for the dummy backend."""
+
+    def __init__(self, backend_args: Any):
+        # BaseTrainer stores backend_args and reads torchrun env (rank, world_size, etc.)
+        super().__init__(backend_args=backend_args)
+        self._initialized = False
+
+    def setup(self):
+        # Optional pre-init setup (e.g., logging, sanity checks)
+        log_rank_0(f"[DummyPretrainTrainer] setup() on rank={self.rank}")
+
+    def init(self):
+        # Build your model / optimizer / dataloader here in a real backend.
+        log_rank_0("[DummyPretrainTrainer] init()")
+        self._initialized = True
+
+    def train(self):
+        if not self._initialized:
+            raise RuntimeError("DummyPretrainTrainer.init() must be called before train().")
+
+        log_rank_0("[DummyPretrainTrainer] train()")
+        # Example: access a custom param (e.g., 'hello') from backend_args.
+        hello_value = getattr(self.backend_args, "hello", "<missing>")
+        log_rank_0(f"[DummyPretrainTrainer] hello={hello_value}")
+        # Real training loop would go here.
+        log_rank_0("[DummyPretrainTrainer] training finished successfully.")
+
+    def cleanup(self, on_error: bool = False):
+        # Optional cleanup logic (close files, finalize logging, etc.)
+        status = "error" if on_error else "success"
+        log_rank_0(f"[DummyPretrainTrainer] cleanup(on_error={status})")
+```
+
+Why this matches the core architecture:
+
+- `BaseTrainer.__init__` reads distributed environment from `get_torchrun_env()`.
+- `PrimusRuntime` drives the lifecycle: `setup` → `init` → `train` → `cleanup` and runs patch phases around these steps.
+- Your trainer only needs to implement `setup`, `init`, `train`, and `cleanup` using `backend_args` and the resolved environment information.
+
+---
+
+## Register the adapter in `BackendRegistry`
+
+**File**: `primus/backends/dummy/__init__.py`
+
+```python
+from primus.backends.dummy.dummy_adapter import DummyAdapter
+from primus.core.backend.backend_registry import BackendRegistry
+
+
+# Register adapter (backend name → adapter class)
+BackendRegistry.register_adapter("dummy", DummyAdapter)
+```
+
+At runtime, when `framework: dummy` is requested:
+
+- `BackendRegistry.get_adapter("dummy")` lazily imports `primus.backends.dummy` (this file), which calls `register_adapter("dummy", DummyAdapter)`.
+- The adapter instance is created and used by `PrimusRuntime` to set up the backend path, run setup hooks, build `backend_args`, and load and construct the trainer.
+
+---
+
+## Minimal config example
+
+Create an experiment YAML (simplified; full template in the next section):
+
+```yaml
+modules:
+  pre_trainer:
+    framework: dummy
+    config: dummy_trainer.yaml
+    model: dummy_8B.yaml
+```
+
+Run:
+
+```bash
+./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml
+```
+
+Because this dummy backend is an in-tree template and `setup_backend_path()` is a no-op, you should see logs similar to:
+
+- `[Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend`
+- `[Primus:DummyAdapter] Converted Primus params -> dummy backend_args`
+- `[DummyPretrainTrainer] setup()`
+- `[DummyPretrainTrainer] init()`
+- `[DummyPretrainTrainer] train()`
+
+---
+
+## Example end-to-end YAML configs
+
+This template mirrors the Megatron pattern. Create these files only when you are actually adding a dummy backend for local development or tests:
+
+- The **top-level experiment config** lives under `examples/<backend>/configs/...`
+- The **module config** is resolved from `primus/configs/modules/{framework}/`
+- The **model config** is resolved from `primus/configs/models/{framework}/`
+
+### Top-level experiment config
+
+**File 1**: `examples/dummy/configs/dummy_8B-pretrain.yaml`
+
+```yaml
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:dummy_8B-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: dummy
+    config: dummy_trainer.yaml
+
+    # model to run
+    model: dummy_8B.yaml
+
+    overrides:
+      # log / debug
+      stderr_sink_level: DEBUG
+
+      # example training overrides (merged into module params)
+      train_iters: 100
+      global_batch_size: 32
+      micro_batch_size: 4
+      seq_length: 1024
+      hello: world
+```
+
+### Module-level trainer config
+
+**File 2**: `primus/configs/modules/dummy/dummy_trainer.yaml`
+
+```yaml
+extends:
+  - trainer_base.yaml        # optional, if you have a shared base; otherwise omit
+
+train_iters: 1000
+global_batch_size: 16
+micro_batch_size: 1
+seq_length: 512
+
+log_interval: 1
+save_interval: 100
+```
+
+This file defines **default training hyperparameters** for the `dummy` backend. Fields under `modules.pre_trainer.overrides` in the top-level config are deep-merged on top of these defaults.
+
+### Model-level config
+
+**File 3**: `primus/configs/models/dummy/dummy_8B.yaml`
+
+```yaml
+extends: []
+
+model_name: dummy_8B
+vocab_size: 32000
+hidden_size: 4096
+num_layers: 32
+num_attention_heads: 32
+```
+
+This file plays the same role as Megatron model configs under `primus/configs/models/megatron/`. It is loaded via:
+
+- `modules.pre_trainer.model: dummy_8B.yaml`
+- resolved as `primus/configs/models/{framework}/dummy_8B.yaml`
+
+### Running the example
+
+Run:
+
+```bash
+./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml
+```
+
+Primus will:
+
+- load `examples/dummy/configs/dummy_8B-pretrain.yaml`
+- resolve `modules.pre_trainer.config` → `primus/configs/modules/dummy/dummy_trainer.yaml`
+- resolve `modules.pre_trainer.model` → `primus/configs/models/dummy/dummy_8B.yaml`
+- build `module_config.params` from these sources plus `overrides`
+- call `DummyAdapter.convert_config(params)` to build `backend_args`
+- construct `DummyPretrainTrainer(backend_args=...)`
+- execute `setup` → `init` → `train` → `cleanup`.
+
+For adding model YAML for existing backends (Megatron, TorchTitan, and others), see [Adding model configurations](./adding-models.md).
+
+---
+
+## Checklist for a complete backend
+
+Use this as a quick checklist when adding a new backend:
+
+- [ ] Adapter subclass of `BackendAdapter` implements:
+      - `load_trainer_class(stage: str)`
+      - `convert_config(params)`
+      - `detect_backend_version()`
+      - (optionally) overrides `prepare_backend()` / `third_party_dir_name`
+- [ ] Trainer subclass of `BaseTrainer` implements:
+      - `setup()`, `init()`, `train()`, and optional `cleanup(on_error: bool)`
+- [ ] `BackendRegistry.register_adapter(backend, AdapterClass)` is called in `primus.backends.<backend>.__init__`
+- [ ] At least one unit test is added under `tests/unit_tests/backends/`
+
+Once these are in place, your backend is fully integrated into the Primus runtime and follows the same lifecycle and patch phases as the built-in backends.
+
+---
+
+## Advanced: Backend-specific setup with train hooks
+
+For more advanced scenarios (for example installing extra Python packages or configuring backend-specific environment variables at runtime), you can use **train hooks** under `runner/helpers/hooks`.
+
+- **Hook locations for training**:
+  - Global hooks (run for all commands): `runner/helpers/hooks/*.sh` and `runner/helpers/hooks/*.py`. These are discovered with `find ... -maxdepth 1 \( -name "*.sh" -o -name "*.py" \)` and executed in **lexicographical order** of their filenames (see `runner/helpers/execute_hooks.sh`).
+  - Command-specific hooks: `runner/helpers/hooks/train/pretrain/*.sh|*.py` (and `.../posttrain/...`), discovered and ordered the same way. For pretrain, this directory contains the dispatcher `prepare_experiment.sh`.
+  - Per-framework hooks: `runner/helpers/hooks/train/pretrain/<framework>/` and `runner/helpers/hooks/train/posttrain/<framework>/`, where `<framework>` is `megatron`, `torchtitan`, `dummy`, and so on. These are **not** run directly by `execute_hooks`; instead `prepare_experiment.sh` detects the framework from the experiment config, runs that framework folder's `*.sh` files in lexicographical order, and then invokes the framework's `prepare.py` dispatcher.
+
+When you run:
+
+```bash
+./primus-cli direct -- train pretrain --config <exp.yaml>
+```
+
+Primus will:
+
+- Call `execute_hooks train pretrain ...`, which:
+  - Runs global hooks under `runner/helpers/hooks/` (lexicographical order)
+  - Then runs command-specific hooks under `runner/helpers/hooks/train/pretrain/`, including `prepare_experiment.sh`
+  - `prepare_experiment.sh` resolves the framework from the config and runs the per-framework hooks under `runner/helpers/hooks/train/pretrain/<framework>/` (its `*.sh` files in lexicographical order, then `prepare.py`)
+
+Each hook script can **emit control lines on stdout** that Primus parses (the framework hooks' stdout is captured through `prepare_experiment.sh`):
+
+- **`env.*=value` → environment variables**
+
+  ```bash
+  # inside runner/helpers/hooks/train/pretrain/<framework>/<something>.sh
+  echo "env.MY_BACKEND_FLAG=1"        # becomes: export MY_BACKEND_FLAG=1
+  echo "env.PYTHONPATH=/opt/mylib:$PYTHONPATH"
+  ```
+
+  These are exported into the environment of the `primus-cli direct` process, so downstream backend code and trainers see them.
+
+- **`extra.*=value` → extra CLI arguments**
+
+  ```bash
+  # inside the same hook
+  echo "extra.backend_path=/opt/my-backend"   # becomes: --backend_path /opt/my-backend
+  echo "extra.train_data_path=/my/data"       # becomes: --train_data_path /my/data
+  ```
+
+  These `extra.*` pairs are appended to the Primus CLI invocation as `--<name> <value>` after hook execution.
+
+Typical pattern to install or configure packages for a backend:
+
+- Add a script under `runner/helpers/hooks/train/pretrain/<framework>/<NNN>-setup.sh` (use a numeric prefix such as `000-` or `010-` to control ordering).
+- In that script:
+  - Optionally run `python -m pip install ...` or other setup commands.
+  - Emit `env.*=...` lines to export any required environment variables.
+  - Emit `extra.*=...` lines if you need to pass additional CLI arguments (for example `backend_path`) into the Primus runtime for this run.
diff --git a/docs/06-developer-guide/model-support-matrix.md b/docs/06-developer-guide/model-support-matrix.md
new file mode 100644
index 000000000..f749f8343
--- /dev/null
+++ b/docs/06-developer-guide/model-support-matrix.md
@@ -0,0 +1,211 @@
+# Model support matrix
+
+This document summarizes which model families Primus targets per backend and lists representative checked-in model presets and example experiment YAML under the repository. It distinguishes **curated examples** from **theoretical** support (a preset or upstream stack may exist without a matching `examples/` entry). Use the filesystem under `primus/configs/models/` and `examples/*/configs/` as the authoritative live inventory.
+
+For how to add presets, see [Adding model configurations](./adding-models.md). Backend parameter references: [Megatron](../03-configuration-reference/megatron-parameters.md), [TorchTitan](../03-configuration-reference/torchtitan-parameters.md), [MaxText](../03-configuration-reference/maxtext-parameters.md), [Megatron Bridge](../03-configuration-reference/megatron-bridge-parameters.md).
+
+---
+
+## Overview: Supported model families (high level)
+
+The following aligns with the backend overview and the configs present in this tree.
+
+| Backend | Model families (documentation / stack scope) |
+| ------- | ---------------------------------------------- |
+| **Megatron-LM** | LLaMA2 / LLaMA3 / LLaMA3.1 / LLaMA3.3 / LLaMA4 (sizes from small to 405B+), DeepSeek-V2 (including lite), DeepSeek-V3, and DeepSeek-V4 (flash / pro), Mixtral MoE and large MoE recipe YAML, Qwen2.5 and Qwen3 (dense and MoE), Grok, GPT-OSS (20B / 120B), GLM, Kimi K2, LFM2, MiniMax, Zebra LLaMA (including GDN and KDA linear-attention variants), Mamba, and generic `language_model.yaml` bases. |
+| **TorchTitan** | LLaMA3 family (including 3.1), LLaMA4 examples, DeepSeek-V3 examples, and Qwen3 examples including 0.6B, 1.7B, 4B, 8B, 14B, and 32B variants where present. Additional presets exist under `primus/configs/models/torchtitan/` without being exhaustively listed here. |
+| **MaxText (JAX)** | LLaMA2 / LLaMA3 / LLaMA3.3, DeepSeek-V2 16B, Mixtral-8x7B, Grok1, Qwen3 14B / 30B-A3B (per presets and examples). Broader coverage may exist in upstream MaxText; see [MaxText](https://github.com/AI-Hypercomputer/maxtext). |
+| **Megatron Bridge** | Qwen3 pretraining and post-training examples, plus post-training examples for Zebra LLaMA and Mamba where present. LLaMA 3.1 70B Bridge examples appear under MI355X. |
+| **Diffusion** | Flux.1 (`schnell` / `dev`) text-to-image and Wan 2.1 / 2.2 text- and image-to-video presets under `primus/configs/models/diffusion/`, with examples under `examples/diffusion/configs/` and `examples/megatron/configs/*/diffusion/`. See [Diffusion models](../04-technical-guides/diffusion-models/README.md). |
+| **HummingbirdXT** | Registered backend with a post-training trainer and one checked-in example; user-facing support level still needs maintainer confirmation. |
+
+**Interpretation:** “Supported” in upstream code can exceed what this repository ships as YAML. Rows below reference representative files that exist under `primus/configs/models/` and `examples/`; they should not be treated as a complete generated inventory.
+
+---
+
+## Megatron model configs
+
+Model presets live in `primus/configs/models/megatron/`. Example experiments that reference those presets appear under `examples/megatron/configs/MI300X/`, `MI325X/`, and `MI355X/`.
+
+For **TorchTitan**, the MI300X, MI325X, and MI355X example directories carry the same model set (21 configs each). For **Megatron**, MI300X and MI325X are nearly identical **except** that MI325X omits `qwen3_5_35B_A3B` (BF16 and FP8)—so MI300X has 70 example configs while MI325X has 68—and **MI355X** is a superset (99 configs; it adds models such as `glm5`, `gpt_oss_120B`, `kimi_k2`, `lfm2_8B_A1B`, and `minimax_m2.5`). Each row's SKU list below reflects exactly which SKUs ship a curated example (see, for example, `qwen3_5_35B_A3B`, which is MI300X/MI355X only).
+
+| Model name (file) | Preset path | Role | Example experiment dirs | Precision in examples |
+| ----------------- | ----------- | ---- | ----------------------- | ---------------------- |
+| `deepseek_v2.yaml` | `primus/configs/models/megatron/deepseek_v2.yaml` | Dense model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `deepseek_v2_base.yaml` | `primus/configs/models/megatron/deepseek_v2_base.yaml` | Base fragment (`extends` only) | — | — |
+| `deepseek_v2_lite.yaml` | `primus/configs/models/megatron/deepseek_v2_lite.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `deepseek_v3.yaml` | `primus/configs/models/megatron/deepseek_v3.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `deepseek_v3_base.yaml` | `primus/configs/models/megatron/deepseek_v3_base.yaml` | Base fragment | — | — |
+| `glm4_7.yaml` | `primus/configs/models/megatron/glm4_7.yaml` | Model preset | No curated example in this repo | — |
+| `glm5.yaml` | `primus/configs/models/megatron/glm5.yaml` | Model preset | MI355X | BF16, FP8 |
+| `gpt_oss_20B.yaml` | `primus/configs/models/megatron/gpt_oss_20B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `gpt_oss_120B.yaml` | `primus/configs/models/megatron/gpt_oss_120B.yaml` | Model preset | MI355X | BF16, FP8 |
+| `grok1.yaml` | `primus/configs/models/megatron/grok1.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `grok2.yaml` | `primus/configs/models/megatron/grok2.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `grok_base.yaml` | `primus/configs/models/megatron/grok_base.yaml` | Base fragment | — | — |
+| `hybrid_model_base.yaml` | `primus/configs/models/megatron/hybrid_model_base.yaml` | Base fragment | — | — |
+| `kimi_k2.yaml` | `primus/configs/models/megatron/kimi_k2.yaml` | MoE model preset | MI355X | BF16, FP8 |
+| `language_model.yaml` | `primus/configs/models/megatron/language_model.yaml` | Generic Megatron LM defaults | Used via `extends` | — |
+| `lfm2_8B_A1B.yaml` | `primus/configs/models/megatron/lfm2_8B_A1B.yaml` | MoE model preset | MI355X | BF16, FP8 |
+| `lfm_base.yaml` | `primus/configs/models/megatron/lfm_base.yaml` | Base fragment | — | — |
+| `llama2_7B.yaml` | `primus/configs/models/megatron/llama2_7B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama2_13B.yaml` | `primus/configs/models/megatron/llama2_13B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama2_70B.yaml` | `primus/configs/models/megatron/llama2_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama2_base.yaml` | `primus/configs/models/megatron/llama2_base.yaml` | Base fragment | — | — |
+| `llama_base.yaml` | `primus/configs/models/megatron/llama_base.yaml` | Base fragment | — | — |
+| `llama3_8B.yaml` | `primus/configs/models/megatron/llama3_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3_70B.yaml` | `primus/configs/models/megatron/llama3_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3_base.yaml` | `primus/configs/models/megatron/llama3_base.yaml` | Base fragment | — | — |
+| `llama3.1_8B.yaml` | `primus/configs/models/megatron/llama3.1_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3.1_70B.yaml` | `primus/configs/models/megatron/llama3.1_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3.1_405B.yaml` | `primus/configs/models/megatron/llama3.1_405B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3.2_1B.yaml` | `primus/configs/models/megatron/llama3.2_1B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3.2_3B.yaml` | `primus/configs/models/megatron/llama3.2_3B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama3.3_70B.yaml` | `primus/configs/models/megatron/llama3.3_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama4_17B128E.yaml` | `primus/configs/models/megatron/llama4_17B128E.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama4_17B16E.yaml` | `primus/configs/models/megatron/llama4_17B16E.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `llama4_base.yaml` | `primus/configs/models/megatron/llama4_base.yaml` | Base fragment | — | — |
+| `mamba_370M.yaml` | `primus/configs/models/megatron/mamba_370M.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides |
+| `mamba_base.yaml` | `primus/configs/models/megatron/mamba_base.yaml` | Base fragment | — | — |
+| `minimax_m2.5.yaml` | `primus/configs/models/megatron/minimax_m2.5.yaml` | MoE model preset | MI355X | BF16, FP8 |
+| `mixtral_8x7B_v0.1.yaml` | `primus/configs/models/megatron/mixtral_8x7B_v0.1.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `mixtral_8x22B_v0.1.yaml` | `primus/configs/models/megatron/mixtral_8x22B_v0.1.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `mixtral_base.yaml` | `primus/configs/models/megatron/mixtral_base.yaml` | Base fragment | — | — |
+| `moe_515B.yaml` | `primus/configs/models/megatron/moe_515B.yaml` | Large MoE template | No curated example in this repo | — |
+| `moe_1T.yaml` | `primus/configs/models/megatron/moe_1T.yaml` | Large MoE template | No curated example in this repo | — |
+| `moe_2T.yaml` | `primus/configs/models/megatron/moe_2T.yaml` | Large MoE template | No curated example in this repo | — |
+| `moe_4T.yaml` | `primus/configs/models/megatron/moe_4T.yaml` | Large MoE template | No curated example in this repo | — |
+| `moe_proxy_single_node.yaml` | `primus/configs/models/megatron/moe_proxy_single_node.yaml` | MoE proxy / test template | No curated example in this repo | — |
+| `primus_megatron_model.yaml` | `primus/configs/models/megatron/primus_megatron_model.yaml` | Primus Megatron root defaults | Used via `extends` | — |
+| `qwen2.5_3B.yaml` | `primus/configs/models/megatron/qwen2.5_3B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen2.5_7B.yaml` | `primus/configs/models/megatron/qwen2.5_7B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen2.5_14B.yaml` | `primus/configs/models/megatron/qwen2.5_14B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen2.5_32B.yaml` | `primus/configs/models/megatron/qwen2.5_32B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen2.5_72B.yaml` | `primus/configs/models/megatron/qwen2.5_72B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen2.5_base.yaml` | `primus/configs/models/megatron/qwen2.5_base.yaml` | Base fragment | — | — |
+| `qwen3_4B.yaml` | `primus/configs/models/megatron/qwen3_4B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen3_8B.yaml` | `primus/configs/models/megatron/qwen3_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen3_14B.yaml` | `primus/configs/models/megatron/qwen3_14B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen3_32B.yaml` | `primus/configs/models/megatron/qwen3_32B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen3_30B_A3B.yaml` | `primus/configs/models/megatron/qwen3_30B_A3B.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `qwen3_5_35B_A3B.yaml` | `primus/configs/models/megatron/qwen3_5_35B_A3B.yaml` | MoE model preset | MI300X, MI355X | BF16, FP8 |
+| `qwen3_235B_A22B.yaml` | `primus/configs/models/megatron/qwen3_235B_A22B.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 |
+| `zebra_llama_1B.yaml` | `primus/configs/models/megatron/zebra_llama_1B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides |
+| `zebra_llama_3B.yaml` | `primus/configs/models/megatron/zebra_llama_3B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides |
+| `zebra_llama_8B.yaml` | `primus/configs/models/megatron/zebra_llama_8B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides |
+
+**Parallelism:** Tensor, pipeline, and expert parallel sizes are **not** fixed in model presets; they are set in experiment `overrides` (for example `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size`). MoE presets such as `qwen3_235B_A22B.yaml` typically require non-default expert parallelism in real runs—see the matching experiment YAML.
+
+---
+
+## TorchTitan model configs
+
+Presets: `primus/configs/models/torchtitan/`. Examples: `examples/torchtitan/configs/MI300X/`, `MI325X/`, and `MI355X/`.
+
+| Model name (file) | Preset path | Example experiment dirs | Precision in examples |
+| ----------------- | ----------- | ------------------------- | ---------------------- |
+| `deepseek_v3_16b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_16b.yaml` | MI300X, MI325X, MI355X | BF16 |
+| `deepseek_v3_16b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_16b-fp8.yaml` | MI300X, MI325X, MI355X | FP8 |
+| `deepseek_v3_236b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_236b.yaml` | MI300X, MI325X, MI355X | BF16 |
+| `deepseek_v3_236b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_236b-fp8.yaml` | MI300X, MI325X, MI355X | FP8 |
+| `deepseek_v3_671b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_671b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `deepseek_v3_671b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_671b-fp8.yaml` | Preset only; stock examples use `deepseek_v3_671b.yaml` | — |
+| `llama3_8B.yaml` | `primus/configs/models/torchtitan/llama3_8B.yaml` | No example in this repo | — |
+| `llama3_8B-fp8.yaml` | `primus/configs/models/torchtitan/llama3_8B-fp8.yaml` | No example in this repo | — |
+| `llama3_70B.yaml` | `primus/configs/models/torchtitan/llama3_70B.yaml` | No example in this repo | — |
+| `llama3_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3_70B-fp8.yaml` | No example in this repo | — |
+| `llama3.1_8B.yaml` | `primus/configs/models/torchtitan/llama3.1_8B.yaml` | MI300X, MI325X, MI355X | BF16 |
+| `llama3.1_8B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_8B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 |
+| `llama3.1_70B.yaml` | `primus/configs/models/torchtitan/llama3.1_70B.yaml` | MI300X, MI325X, MI355X | BF16 |
+| `llama3.1_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_70B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 |
+| `llama3.1_405B.yaml` | `primus/configs/models/torchtitan/llama3.1_405B.yaml` | MI300X, MI325X, MI355X | BF16 |
+| `llama3.1_405B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_405B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 |
+| `llama3.2_1B.yaml` | `primus/configs/models/torchtitan/llama3.2_1B.yaml` | No example in this repo | — |
+| `llama3.3_70B.yaml` | `primus/configs/models/torchtitan/llama3.3_70B.yaml` | No example in this repo | — |
+| `llama3.3_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.3_70B-fp8.yaml` | No example in this repo | — |
+| `llama4_17Bx128E.yaml` | `primus/configs/models/torchtitan/llama4_17Bx128E.yaml` | MoE; MI300X, MI325X, MI355X | BF16 |
+| `llama4_17Bx128E-fp8.yaml` | `primus/configs/models/torchtitan/llama4_17Bx128E-fp8.yaml` | MoE; MI300X, MI325X, MI355X | FP8 |
+| `llama4_17Bx16E.yaml` | `primus/configs/models/torchtitan/llama4_17Bx16E.yaml` | MoE; MI300X, MI325X, MI355X | BF16 |
+| `llama4_17Bx16E-fp8.yaml` | `primus/configs/models/torchtitan/llama4_17Bx16E-fp8.yaml` | MoE; MI300X, MI325X, MI355X | FP8 |
+| `qwen3_0.6b.yaml` | `primus/configs/models/torchtitan/qwen3_0.6b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `qwen3_1.7b.yaml` | `primus/configs/models/torchtitan/qwen3_1.7b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `qwen3_4b.yaml` | `primus/configs/models/torchtitan/qwen3_4b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `qwen3_8b.yaml` | `primus/configs/models/torchtitan/qwen3_8b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `qwen3_14b.yaml` | `primus/configs/models/torchtitan/qwen3_14b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+| `qwen3_32b.yaml` | `primus/configs/models/torchtitan/qwen3_32b.yaml` | MI300X, MI325X, MI355X | (see experiment) |
+
+**Parallelism:** Controlled by TorchTitan launch configuration and Primus module overrides (see TorchTitan patch notes and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md)); not embedded in the small `job` / `model` preset alone.
+
+---
+
+## MaxText model configs
+
+Presets: `primus/configs/models/maxtext/`. Examples: `examples/maxtext/configs/MI300X/` and `examples/maxtext/configs/MI355X/`.
+
+| Model name (file) | Preset path | Example experiment dirs |
+| ----------------- | ----------- | ------------------------ |
+| `deepseek_v2_16B.yaml` | `primus/configs/models/maxtext/deepseek_v2_16B.yaml` | MI300X, MI355X |
+| `grok1.yaml` | `primus/configs/models/maxtext/grok1.yaml` | MI300X |
+| `llama2_7B.yaml` | `primus/configs/models/maxtext/llama2_7B.yaml` | MI300X, MI355X |
+| `llama2_70B.yaml` | `primus/configs/models/maxtext/llama2_70B.yaml` | MI300X, MI355X |
+| `llama3_8B.yaml` | `primus/configs/models/maxtext/llama3_8B.yaml` | MI300X, MI355X |
+| `llama3_70B.yaml` | `primus/configs/models/maxtext/llama3_70B.yaml` | MI300X, MI355X |
+| `llama3.1_405B.yaml` | `primus/configs/models/maxtext/llama3.1_405B.yaml` | MI355X |
+| `llama3.3_70B.yaml` | `primus/configs/models/maxtext/llama3.3_70B.yaml` | MI300X, MI355X |
+| `mixtral_8x7B.yaml` | `primus/configs/models/maxtext/mixtral_8x7B.yaml` | MI300X, MI355X |
+| `qwen3_14B.yaml` | `primus/configs/models/maxtext/qwen3_14B.yaml` | MI300X, MI355X |
+| `qwen3_30B_A3B.yaml` | `primus/configs/models/maxtext/qwen3_30B_A3B.yaml` | MI300X, MI355X |
+| `model_base.yaml` | `primus/configs/models/maxtext/model_base.yaml` | Extended by other presets (not a standalone run) |
+
+**Parallelism:** JAX / MaxText sharding is configured in experiment overrides (for example `ici_fsdp_parallelism`, `ici_data_parallelism`, `dcn_*` in sample experiments). See [MaxText parameters](../03-configuration-reference/maxtext-parameters.md).
+
+---
+
+## Megatron Bridge model configs
+
+Presets: `primus/configs/models/megatron_bridge/`. Examples: `examples/megatron_bridge/configs/MI300X/` and `examples/megatron_bridge/configs/MI355X/`.
+
+| Model name (file) | Preset path | Recipe / flavor (from preset) | Example experiment dirs |
+| ----------------- | ----------- | ----------------------------- | ------------------------ |
+| `qwen3_8b.yaml` | `primus/configs/models/megatron_bridge/qwen3_8b.yaml` | `qwen.qwen3` / `qwen3_8b_finetune_config` | MI300X pretrain, MI355X posttrain |
+| `qwen3_32b.yaml` | `primus/configs/models/megatron_bridge/qwen3_32b.yaml` | `qwen.qwen3` / `qwen3_32b_finetune_config` | MI300X, MI355X |
+| `llama31_70b.yaml` | `primus/configs/models/megatron_bridge/llama31_70b.yaml` | `llama.llama3` / `llama31_70b_finetune_config` | MI355X |
+| `zebra_llama_1B.yaml`, `zebra_llama_3B.yaml`, `zebra_llama_8B.yaml` | `primus/configs/models/megatron_bridge/` | Zebra LLaMA presets | MI300X posttrain |
+| `mamba_370M.yaml` | `primus/configs/models/megatron_bridge/mamba_370M.yaml` | Mamba preset | MI300X posttrain |
+
+Example filenames include `*_pretrain.yaml`, `*_sft_posttrain.yaml`, and `*_lora_posttrain.yaml`; precision such as `bf16_mixed` is set in experiment `overrides`.
+
+---
+
+## Hardware compatibility (example directories)
+
+Curated example layouts under `examples/` use GPU SKU subdirectories. As of this document:
+
+| GPU SKU | `examples/megatron/configs/` | `examples/torchtitan/configs/` | `examples/maxtext/configs/` | `examples/megatron_bridge/configs/` |
+| ------- | ---------------------------- | ------------------------------ | ---------------------------- | ----------------------------------- |
+| **MI300X** | Yes | Yes | Yes | Yes |
+| **MI355X** | Yes | Yes | Yes | Yes |
+| **MI325X** | Yes | Yes | No | No |
+
+Megatron and TorchTitan ship MI325X example directories in addition to MI300X and MI355X examples. MaxText includes MI300X and MI355X examples, including MI355X-only entries such as `llama3.1_405B-pretrain.yaml`. Megatron Bridge MI300X examples include Qwen3 8B and 32B pretraining plus Qwen3 32B, Zebra LLaMA, and Mamba post-training examples; LLaMA 3.1 70B Bridge examples appear under MI355X.
+
+Absence of a SKU directory for a given backend does **not** imply the backend cannot run there; it means this tree does not currently provide a checked-in example path to copy from.
+
+---
+
+## Model architecture reference (Megatron presets)
+
+Values below come from `primus/configs/models/megatron/` presets (merged through `extends`). **Vocabulary size** is usually defined by the tokenizer / Hugging Face config, not duplicated in every YAML; **context** is `max_position_embeddings` where set in the chain. Use this table as a quick reference for common sizes—not an exhaustive spec of every parameter.
+
+| Model family | Example preset | Hidden size | Layers | Attention heads | KV heads (GQA) | Max position (context) |
+| ------------ | -------------- | ----------- | ------ | ----------------- | ---------------- | ------------------------ |
+| LLaMA 2 7B | `llama2_7B.yaml` | 4096 | 32 | 32 | 32 (no GQA) | From `llama2_base` / tokenizer |
+| LLaMA 3 8B | `llama3_8B.yaml` | 4096 | 32 | 32 | 8 | 8192 (`llama3_base`) |
+| LLaMA 3 70B | `llama3_70B.yaml` | 8192 | 80 | 64 | 8 | 8192 |
+| LLaMA 3.1 405B | `llama3.1_405B.yaml` | 16384 | 126 | 128 | 8 | 8192 |
+| Qwen3 8B | `qwen3_8B.yaml` | 4096 | 36 | 32 | 8 | 131072 (`qwen2.5_base` chain) |
+| Mixtral 8x7B | `mixtral_8x7B_v0.1.yaml` | 4096 | 32 | 32 | — | 4096 |
+| DeepSeek-V3 (MoE) | `deepseek_v3.yaml` | 7168 | 61 | 128 (MLA) | — | See preset / HF |
+| Mamba 370M | `mamba_370M.yaml` | (Mamba stack) | — | — | — | — |
+
+For MoE and hybrid architectures (LLaMA 4, Qwen3-MoE, large `moe_*.yaml` templates), refer to the full YAML and upstream model cards; headline dimensions alone do not capture expert layout or MLA.
diff --git a/docs/06-developer-guide/testing.md b/docs/06-developer-guide/testing.md
new file mode 100644
index 000000000..aa8215ca4
--- /dev/null
+++ b/docs/06-developer-guide/testing.md
@@ -0,0 +1,175 @@
+# Testing guide
+
+This guide describes where tests live, how to run them locally, and how they map to CI. For coding standards and PR workflow, see [Contributing Guide](contributing.md). The canonical CI definition is `.github/workflows/ci.yaml`.
+
+## 1. Test organization
+
+Layout (simplified from the repository root):
+
+```text
+Primus/
+├── runner/
+│   └── lib/
+│       └── common.sh              # Shared logging/helpers sourced by the shell test runner
+├── tests/
+│   ├── runner/                    # Shell integration tests
+│   │   ├── run_all_tests.sh       # Master shell test runner
+│   │   ├── lib/                   # test_common.sh, test_config.sh, test_validation.sh
+│   │   ├── helpers/               # Hook and env tests
+│   │   ├── test_primus_cli.sh
+│   │   ├── test_primus_cli_direct.sh
+│   │   ├── test_primus_cli_container.sh
+│   │   └── test_primus_cli_slurm.sh
+│   ├── unit_tests/                # Python unit tests (pytest)
+│   │   ├── agents/                # Tuning-agent tests
+│   │   ├── backends/
+│   │   ├── ci/
+│   │   ├── cli/
+│   │   ├── core/                  # config, backend, launcher, patches, projection, pipeline_parallel, runtime, trainer, utils
+│   │   ├── megatron/              # Megatron-specific unit tests
+│   │   ├── modules/
+│   │   └── tools/
+│   ├── trainer/                   # Integration tests (typically need GPU + data)
+│   │   ├── test_megatron_trainer.py
+│   │   ├── test_torchtitan_trainer.py
+│   │   └── test_maxtext_trainer.py
+│   ├── scripts/                   # CI unit/integration launch scripts and UT patches
+│   ├── utils.py                   # Shared test utilities
+│   └── run_unit_tests.py          # Optional orchestrator (walks tests/, see below)
+```
+
+`tests/runner/run_all_tests.sh` sources shared helpers from **`runner/lib/common.sh`** at the repository root (not under `tests/runner/`).
+
+## 2. Running tests
+
+**Shell integration tests** (CLI behavior, config loading, hooks, environment):
+
+```bash
+bash ./tests/runner/run_all_tests.sh
+```
+
+**Python unit tests:**
+
+```bash
+pytest tests/unit_tests/ --maxfail=1 -s
+```
+
+**Trainer integration tests** (GPU and data; might require Hugging Face access):
+
+```bash
+# Megatron
+DATA_PATH=<path> pytest tests/trainer/test_megatron_trainer.py -s
+
+# TorchTitan
+DATA_PATH=<path> pytest tests/trainer/test_torchtitan_trainer.py -s
+
+# MaxText (JAX) — often run via the orchestrator in CI
+python ./tests/run_unit_tests.py --jax
+```
+
+`tests/run_unit_tests.py` walks **`tests/`** and runs every `test_*.py` it finds, except for **`tests/trainer/test_maxtext_trainer.py`** in the default mode (that file is only selected when **`--jax`** is set). That means the default orchestrator run includes **`tests/unit_tests/`** and **`tests/trainer/`** (and any other matching tests), which is broader than `pytest tests/unit_tests/` alone.
+
+**Orchestrator (default—all discovered tests except MaxText trainer):**
+
+```bash
+python ./tests/run_unit_tests.py
+```
+
+**Orchestrator (JAX / MaxText trainer only):**
+
+```bash
+python ./tests/run_unit_tests.py --jax
+```
+
+## 3. Test types
+
+- **Shell tests:** Exercise runner scripts, CLI wiring, configuration loading, hook execution, and environment setup. Implemented as bash scripts under `tests/runner/` and orchestrated by `run_all_tests.sh`.
+- **Unit tests:** Cover configuration parsing, preset loading, CLI behavior, patch registration, adapters, and other library logic under `tests/unit_tests/`.
+- **Trainer tests:** End-to-end training against real backends; require AMD GPUs and appropriate data paths (and sometimes tokens). See `.github/workflows/ci.yaml` for CI values such as `DATA_PATH`, `MASTER_PORT`, and `HSA_NO_SCRATCH_RECLAIM`.
+
+### Test tiers (slim on PRs, full on weekends)
+
+Each trainer E2E case is a full training launch, so the suites hold one model per
+architecture and per feature path. A model that only scales the dims of an
+existing test is **deleted**, not kept in a slower tier — its recipe stays
+schema-checked by the example-config smoke test below. What remains is split by
+the **`@pytest.mark.weekly`** marker (registered in `tests/conftest.py`):
+
+| Tier | Trigger | Model E2E | Unit tests |
+| --- | --- | --- | --- |
+| slim | pull request, push to `main`/tags | `-m "not weekly"` | default (slow gates skipped) |
+| full | `schedule` cron (Sat 18:00 UTC), or `workflow_dispatch` with `full_tests=true` | no filter | `--run-slow` |
+
+The workflow derives every test step's arguments from a single `PRIMUS_CI_FULL`
+workflow-level env var, so there is one switch to flip.
+
+Two kinds of case belong in the weekly tier: extended coverage that is not worth
+a PR's wall clock — secondary architectures, extra precisions, variants of a
+feature whose primary path stays per-PR — and new E2E cases during burn-in,
+promoted to the per-PR tier by deleting the marker once a weekend run has shown
+them green. Each marked case carries a comment saying what still covers its path
+on PRs.
+
+Reproduce either tier locally:
+
+```bash
+pytest tests/trainer/test_megatron_trainer.py -m "not weekly" -s   # what PR CI runs
+pytest tests/trainer/test_megatron_trainer.py -s                   # full matrix
+```
+
+Two things are deliberately outside this mechanism and stay hidden in **both**
+tiers: the cases `--deselect`ed in `ci.yaml` because they are broken on the
+current toolchain, and the MaxText models hidden by `JAX_SKIP_UT=1`.
+
+Do not delete the last remaining test of an architecture or of a feature path.
+Two cases with identical `extra_args` are not necessarily isomorphic: the flag
+that makes one unique often lives in its recipe yaml, so compare those too.
+
+### Example-config smoke test
+
+`tests/unit_tests/configs/test_example_configs.py` loads **every** yaml under
+`examples/**/configs/` through the real config stack — env interpolation,
+`extends:` merge, module/model preset merge, experiment overrides — and asserts
+each resolves to a well-formed experiment whose declared `framework` matches its
+backend directory. It is CPU-only (`load_primus_config` does not import
+megatron/torchtitan/jax) and runs in seconds.
+
+This is what makes deleting an isomorphic E2E safe: the recipe stops being
+*trained*, but a renamed model preset, a broken `extends:` chain or a `${VAR}`
+without a default still fails CI, naming the exact yaml.
+
+## 4. Writing new tests
+
+- **Pytest:** Add files named `test_*.py` under `tests/unit_tests/`, following existing patterns and reusing fixtures from `conftest.py` where present.
+- **Shell:** Add scripts under `tests/runner/` or extend `tests/runner/run_all_tests.sh` to invoke new suites, consistent with existing `test_primus_cli*.sh` scripts.
+- **Backends:** Prefer `tests/unit_tests/backends/<backend>/` for adapter-focused tests.
+
+## 5. CI pipeline details
+
+From `.github/workflows/ci.yaml`:
+
+- **`code-lint`:** Python 3.12 on GitHub-hosted runners. Runs `pre-commit run --all-files --show-diff-on-failure`, so checks follow `.pre-commit-config.yaml`.
+- **`dependency-review`:** Runs `actions/dependency-review-action` on pull requests to flag dependency changes.
+- **`run-unittest-torch`:** Self-hosted GPU runner. Installs `requirements.txt`, runs `bash ./tests/runner/run_all_tests.sh`, then `pytest tests/unit_tests/` under coverage (`--cov=primus --cov-report=term-missing`) with specific tests `--deselect`ed (currently some `megatron/cco` TP-overlap and `megatron/transformer/moe` dispatcher cases—see the workflow file). Trainer steps set `MASTER_PORT`, `DATA_PATH`, `HSA_NO_SCRATCH_RECLAIM=1`, and `HF_TOKEN` for Megatron and TorchTitan trainer tests. A follow-up **coverage** step combines unit and E2E coverage.
+- **`run-unittest-jax`:** JAX runner. Installs `requirements-jax.txt`, runs the same shell test script, then `python ./tests/run_unit_tests.py --jax` with CI environment variables (for example `JAX_SKIP_UT=1` and `DATA_PATH` as defined in the workflow).
+
+The torch job picks its tier from `PRIMUS_CI_FULL` and prints the resolved tier in the job summary. See [Test tiers](#test-tiers-slim-on-prs-full-on-weekends).
+
+The **`build-docker`** job builds images after lint passes; unit test jobs depend on **`code-lint`**, not on **`build-docker`**.
+
+## 6. Pre-commit hooks
+
+Install once per clone:
+
+```bash
+pip install pre-commit
+pre-commit install
+```
+
+Run manually on the whole tree:
+
+```bash
+pre-commit run --all-files
+```
+
+Hooks include: `trailing-whitespace`, `end-of-file-fixer`, `check-yaml`, `check-added-large-files`, `check-merge-conflict`, `isort`, `autoflake`, `black`, and `shellcheck` (as configured in `.pre-commit-config.yaml`). These align with the **`code-lint`** job in CI; see [Contributing Guide](contributing.md) for manual equivalents.
diff --git a/docs/06-developer-guide/tooling.md b/docs/06-developer-guide/tooling.md
new file mode 100644
index 000000000..982755a41
--- /dev/null
+++ b/docs/06-developer-guide/tooling.md
@@ -0,0 +1,24 @@
+# Tooling
+
+Primus ships a set of auxiliary tools for analysis, benchmarking, visualization, installation, and diagnostics. They live under [`tools/`](../../tools/README.md) in the repository; each tool keeps its own README with detailed usage instructions. This page is a lightweight index so the tooling is discoverable from the documentation.
+
+## Available tools
+
+| Tool | Directory | What it does | Docs |
+|------|-----------|--------------|------|
+| **IRLens** | `tools/IRLens/` | Parses XLA HLO text dumps and prints an execution skeleton with control flow, separating communication vs compute ops. | [README](../../tools/IRLens/README.md) |
+| **model_stats** | `tools/model_stats/` | Generates charts from the model config registry under `primus/configs/models`. | [README](../../tools/model_stats/README.md) |
+| **Pipeline Visualization** | `tools/visualization/pp_vis/` | Visualizes pipeline-parallelism schedules from dumped data or PP-simulator JSON via a local web UI. | [README](../../tools/visualization/pp_vis/README.md) |
+| **Auto Benchmark** | `tools/auto_benchmark/` | Interactive benchmark menu for Megatron/TorchTitan on MI300X/MI355X with metrics collection. | [README](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) |
+| **Backend Gap Report / Engineering Dashboard** | `tools/backend_gap_report/` | Generation and publishing toolchain for the shared Primus engineering dashboard and backend-gap reports. | [README](../../tools/backend_gap_report/README.md) |
+| **Installation (venv)** | `tools/installation/` | Reproduces the Primus training Docker environment in a Python virtual environment (no Docker, no sudo). | [README](../../tools/installation/README.md) |
+| **Daily Report** | `tools/daily/` | Benchmark summary CSV generation used by CI workflows. | — |
+| **Docker Helpers** | `tools/docker/` | Container startup and proxy scripts. | — |
+| **Profile Trace** | `tools/profile_trace/` | Trace-file merging utility. | — |
+
+## Related documentation
+
+- [Primus tools](../02-user-guide/primus-tools.md)—the full catalog of Primus tools (CLI, tuning agent, ecosystem) with how-to starting points.
+- [Tools overview](../../tools/README.md)—the top-level index maintained alongside the code.
+- [Profiling and observability](../04-technical-guides/profiling-and-observability.md)—how these tools fit into performance analysis.
+- [Benchmarking](../02-user-guide/benchmarking.md)—running the benchmark suites the tools summarize.
diff --git a/docs/README.md b/docs/README.md
index 48c6c271e..dd4da3e9b 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,70 +1,148 @@
-# Primus Documentation
+# Primus documentation
 
-Welcome to the Primus documentation! This guide will help you get started with training large-scale foundation models on AMD GPUs.
+Documentation for **Primus**, a large-scale foundation model training framework for AMD GPUs.
 
-## 📚 Documentation Structure
+---
+
+## Choose your starting point
+
+| I am a... | Start here |
+|-----------|------------|
+| **New user** | [Getting started](./01-getting-started/overview.md) |
+| **User** running training jobs | [User guide](./02-user-guide/pretraining.md) |
+| **User** writing YAML configurations | [Configuration reference](./03-configuration-reference/megatron-parameters.md) |
+| **Engineer** tuning performance | [Technical guides](./04-technical-guides/performance-tuning.md) |
+| **Operator** deploying to production | [Operations](./05-operations/deployment.md) |
+| **Contributor** to the codebase | [Developer guide](./06-developer-guide/architecture.md) |
+
+---
+
+## Documentation structure
+
+### [Getting started](./01-getting-started/)
+
+Start here if you are new to Primus.
 
-### 🚀 Getting Started
+- [Project overview](./01-getting-started/overview.md): what Primus does, who it is for, key capabilities
+- [Installation guide](./01-getting-started/installation.md): prerequisites, Docker/bare-metal/Slurm setup
+- [Quickstart](./01-getting-started/quickstart.md): first training run in 5 minutes
+- [Glossary](./01-getting-started/glossary.md): terms, acronyms, and domain concepts
 
-Start here if you're new to Primus:
+### [User guide](./02-user-guide/)
 
-- **[Quick Start Guide](./quickstart.md)** - Get up and running in 5 minutes
-- **[CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)** - Complete command-line reference
-- **[CLI Architecture](./cli/CLI-ARCHITECTURE.md)** - Design philosophy and deep dive
+Core workflows and day-to-day usage.
 
-### 📖 User Guides
+- [CLI reference](./02-user-guide/cli-reference.md): `primus-cli` modes, flags, and subcommands
+- [Configuration system](./02-user-guide/configuration-system.md): YAML configuration model, presets, overrides, inheritance
+- [Pretraining](./02-user-guide/pretraining.md): pretraining **concepts**: backends, YAML structure, parallelism, configuration inventory
+- [Backend training recipes](./02-user-guide/training-recipes.md): pretraining **commands**: copy-paste, GPU-arch-specific run commands
+- [Post-training](./02-user-guide/posttraining.md): SFT and LoRA fine-tuning via Megatron Bridge
+- [Benchmarking](./02-user-guide/benchmarking.md): GEMM, RCCL, and dense-GEMM benchmark suites
+- [Preflight](./02-user-guide/preflight.md): cluster diagnostics and environment validation
+- [Projection](./02-user-guide/projection.md): memory and performance projection tools
+- [Tuning agent](./02-user-guide/tuning-agent.md): LLM-driven search for an optimal training configuration (uses projection as an oracle)
+- [Primus tools](./02-user-guide/primus-tools.md): catalog of all Primus tools and ecosystem projects with how-to starting points
 
-Guides for common workflows and features:
+### [Configuration reference](./03-configuration-reference/)
 
-- **[Configuration Guide](./configuration.md)** - YAML/TOML configuration, recommended patterns, and examples
-- **[Slurm & Container Usage](./slurm-container.md)** - Distributed training and containerization workflows
-- **[Experiment Management](./experiments.md)** - Organizing and tracking your training runs
+Parameter references for Primus presets, backend-facing keys, and commonly used environment variables.
 
-### 🔧 Technical References
+- [Megatron parameters](./03-configuration-reference/megatron-parameters.md): Megatron-LM backend YAML parameters and Primus overrides
+- [TorchTitan parameters](./03-configuration-reference/torchtitan-parameters.md): Primus TorchTitan preset keys and common JobConfig fields
+- [MaxText parameters](./03-configuration-reference/maxtext-parameters.md): Primus MaxText overlay defaults and common fields
+- [Megatron Bridge parameters](./03-configuration-reference/megatron-bridge-parameters.md): Megatron Bridge recipe, SFT, and pretraining fields surfaced through Primus
+- [Environment variables](./03-configuration-reference/environment-variables.md): practical reference for commonly encountered environment variables
 
-In-depth technical documentation:
+### [Technical guides](./04-technical-guides/)
 
-- **[Post-Training Guide](./posttraining.md)** - Fine-tuning with SFT and LoRA using Primus CLI
-- **[Native SFT & LoRA Quick Start](./README_NATIVE_SFT_LORA_EN.md)** - Megatron-native SFT/LoRA launch guide (BF16/FP8/FP4), no Megatron-Bridge runtime dependency
-- **[Performance Projection](./projection.md)** - Project training performance and memory to multi-node configurations
-- **[Tuning Agent](./tuning_agent.md)** - LLM-driven search for an optimal training config — parallelism plus batching, schedule, memory, MoE-comm, and precision knobs (drives the projection tool as an oracle)
-- **[Preflight](./preflight.md)** - Cluster diagnostics (host/GPU/network info + perf tests)
-- **[Benchmark Suite](./benchmark.md)** - GEMM, RCCL, end-to-end benchmarks and profiling
-- **[Supported Models](./backends/overview.md#supported-models)** - Supported LLM architectures and feature compatibility matrix
-- **[Advanced Features](./advanced.md)** - Mixed precision, parallelism strategies, optimization techniques
-- **[Backend Patch Notes](./backends/overview.md)** - Primus-specific arguments for Megatron, TorchTitan, etc.
-- **[Backend Extension Guide](./backends/extending-backends.md)** - How to add a new backend using the current adapter/trainer architecture
- - **[Megatron Model Extension Guide](./backends/adding-megatron-models.md)** - How to add a new Megatron model config
- - **[TorchTitan Model Extension Guide](./backends/adding-torchtitan-models.md)** - How to add a new TorchTitan model config
+Deep technical topics for advanced users.
 
-### 💡 Help & Support
+- [Parallelism strategies](./04-technical-guides/parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained
+- [Parallelism configuration](./04-technical-guides/parallelism-configuration.md): per-backend parallelism setup and batch size relationships
+- [Collective operations](./04-technical-guides/collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy
+- [Performance tuning](./04-technical-guides/performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization
+- [MoE training deep-dive](./04-technical-guides/moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models
+- [MegaMoE fused MoE layer](./04-technical-guides/mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction
+- [Data preparation](./04-technical-guides/data-preparation.md): tokenization, data formats, mock data
+- [Checkpoint management](./04-technical-guides/checkpoint-management.md): formats, save/load, distributed checkpointing
+- [Multi-node networking](./04-technical-guides/multi-node-networking.md): InfiniBand, RoCE, AINIC configuration
+- [Profiling and observability](./04-technical-guides/profiling-and-observability.md): Torch profiler, TraceLens, memory snapshots, projection, pp_vis
+- [Logging and experiment tracking](./04-technical-guides/logging-and-experiment-tracking.md): TensorBoard, WandB, MLflow setup per backend
+- [Fault tolerance and elastic training](./04-technical-guides/fault-tolerance-and-elastic-training.md): graceful exit, auto-resume, in-process restart, torchft
+- [Determinism and reproducibility](./04-technical-guides/determinism-and-reproducibility.md): deterministic mode, seeds, trade-offs
+- [Diffusion models](./04-technical-guides/diffusion-models/README.md): Flux diffusion architecture, data pipeline, and FP8 / MXFP4 training
+- [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency
 
-Get help and find answers:
+### [Operations](./05-operations/)
 
-- **[FAQ](./faq.md)** - Frequently asked questions and troubleshooting
-- **[Examples](../examples/README.md)** - Real-world training examples and templates
-- **[Preflight Tool](../primus/tools/preflight/README.md)** - Cluster sanity checker to verify environment readiness
+Production deployment and operational guidance.
 
-## 🎯 Quick Navigation by Use Case
+- [Deployment](./05-operations/deployment.md): container, Slurm, and Kubernetes deployment
+- [Monitoring and logging](./05-operations/monitoring-logging.md): WandB, TensorBoard, MLflow, Primus logging
+- [Troubleshooting](./05-operations/troubleshooting.md): common failures, diagnostics, and fixes
+- [Security](./05-operations/security.md): secrets handling, container security, dependencies
+
+### [Developer guide](./06-developer-guide/)
+
+For contributors and maintainers.
+
+- [Architecture](./06-developer-guide/architecture.md): system design, runtime, backends, patch system
+- [Contributing](./06-developer-guide/contributing.md): development setup, code style, PR process
+- [Testing](./06-developer-guide/testing.md): test types, running tests, CI pipeline
+- [Extending backends](./06-developer-guide/extending-backends.md): adding new training backends
+- [Adding models](./06-developer-guide/adding-models.md): adding model configurations per backend
+- [Model support matrix](./06-developer-guide/model-support-matrix.md): supported models per backend and GPU
+- [CLI architecture](./06-developer-guide/cli-architecture.md): CLI internals: subcommand discovery, dispatch, and launch wrappers
+- [Backend patch notes](./06-developer-guide/backend-patch-notes.md): Primus-specific backend arguments and the files they patch
+- [Tooling](./06-developer-guide/tooling.md): auxiliary analysis, benchmarking, visualization, and diagnostics tools
+
+---
+
+## Common use cases
 
 ### I want to...
 
-- **Train a model locally** → [Quick Start](./quickstart.md) + [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)
-- **Run distributed training on Slurm** → [Slurm & Container Usage](./slurm-container.md)
-- **Configure my training run** → [Configuration Guide](./configuration.md)
-- **Project performance to multi-node** → [Performance Projection](./projection.md)
-- **Auto-tune my training config (parallelism + knobs)** → [Tuning Agent](./tuning_agent.md)
-- **Benchmark performance** → [Benchmark Suite](./benchmark.md)
-- **Understand the CLI design** → [CLI Architecture](./cli/CLI-ARCHITECTURE.md)
-- **Troubleshoot issues** → [FAQ](./faq.md)
+| Goal | Document |
+|------|----------|
+| Understand what Primus is | [Overview](./01-getting-started/overview.md) |
+| Browse all Primus tools | [Primus tools](./02-user-guide/primus-tools.md) |
+| Install Primus | [Installation](./01-getting-started/installation.md) |
+| Run my first training | [Quickstart](./01-getting-started/quickstart.md) |
+| Get an exact run command for my model/GPU | [Backend training recipes](./02-user-guide/training-recipes.md) |
+| Write a training YAML configuration | [Configuration system](./02-user-guide/configuration-system.md) |
+| Look up a Megatron parameter | [Megatron parameters](./03-configuration-reference/megatron-parameters.md) |
+| Look up a TorchTitan parameter | [TorchTitan parameters](./03-configuration-reference/torchtitan-parameters.md) |
+| Look up an environment variable | [Environment variables](./03-configuration-reference/environment-variables.md) |
+| Understand parallelism strategies | [Parallelism strategies](./04-technical-guides/parallelism-strategies.md) |
+| Configure parallelism for my model | [Parallelism configuration](./04-technical-guides/parallelism-configuration.md) |
+| Tune training performance | [Performance tuning](./04-technical-guides/performance-tuning.md) |
+| Train a Mixture-of-Experts model | [MoE training deep-dive](./04-technical-guides/moe-training.md) |
+| Use the fused MegaMoE layer | [MegaMoE fused MoE layer](./04-technical-guides/mega-moe.md) |
+| Train a diffusion (Flux) model | [Diffusion models](./04-technical-guides/diffusion-models/README.md) |
+| Fine-tune with native SFT / LoRA | [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md) |
+| Auto-tune my training configuration | [Tuning agent](./02-user-guide/tuning-agent.md) |
+| Profile a training run | [Profiling and observability](./04-technical-guides/profiling-and-observability.md) |
+| Track experiments (WandB/MLflow/TensorBoard) | [Logging and experiment tracking](./04-technical-guides/logging-and-experiment-tracking.md) |
+| Survive node failures on long runs | [Fault tolerance and elastic training](./04-technical-guides/fault-tolerance-and-elastic-training.md) |
+| Reproduce results bit-for-bit | [Determinism and reproducibility](./04-technical-guides/determinism-and-reproducibility.md) |
+| Prepare training data | [Data preparation](./04-technical-guides/data-preparation.md) |
+| Deploy to a Slurm cluster | [Deployment](./05-operations/deployment.md) |
+| Debug a training failure | [Troubleshooting](./05-operations/troubleshooting.md) |
+| Contribute to Primus | [Contributing](./06-developer-guide/contributing.md) |
+| Understand the code architecture | [Architecture](./06-developer-guide/architecture.md) |
+| Add a new training backend | [Extending backends](./06-developer-guide/extending-backends.md) |
+
+---
 
-## 🔗 External Resources
+## External resources
 
-- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) - High-performance operators & modules
-- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) - Stability & platform layer
-- [AMD ROCm Documentation](https://rocm.docs.amd.com/)
-- [TorchTitan Documentation](https://github.com/pytorch/torchtitan)
+- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo): high-performance operators and kernels
+- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE): external stability/platform layer; this repository does not include a production integration guide
+- [AMD ROCm documentation](https://rocm.docs.amd.com/)
+- [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
+- [TorchTitan](https://github.com/pytorch/torchtitan)
+- [MaxText](https://github.com/AI-Hypercomputer/maxtext)
 
 ---
 
-**Need help?** Check the [FAQ](./faq.md) or open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues).
+**Need help?** Open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues).
diff --git a/docs/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json b/docs/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json
deleted file mode 100644
index 4c0f1313f..000000000
--- a/docs/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json
+++ /dev/null
@@ -1,87 +0,0 @@
-{
-  "id": "megatron-upstream-main-2026-04-30",
-  "title": "Primus Megatron vs upstream main",
-  "backend": {
-    "key": "megatron",
-    "label": "Megatron"
-  },
-  "generated_at": "2026-04-30",
-  "status": "verified",
-  "scope": "Current Megatron-LM bundled in Primus vs upstream NVIDIA/Megatron-LM origin/main.",
-  "local": {
-    "source_path": "third_party/Megatron-LM",
-    "version": "0.16.0rc0",
-    "commit": "d3528a21",
-    "commit_date": "2026-03-06"
-  },
-  "upstream": {
-    "repo": "https://github.com/NVIDIA/Megatron-LM",
-    "ref": "origin/main",
-    "version": "0.18.0",
-    "commit": "0d98cb83",
-    "commit_date": "2026-04-29"
-  },
-  "stats": {
-    "commit_gap": 403,
-    "diff_files": 1300,
-    "insertions": 615916,
-    "deletions": 326835
-  },
-  "integration": {
-    "backend_files": 136,
-    "tracked_files": 365,
-    "integration_model": "third_party/Megatron-LM + Primus outer trainer / adapter / patches"
-  },
-  "highlights": [
-    "Primus pins Megatron Core source version 0.16.0rc0 while upstream main declares 0.18.0.",
-    "Upstream main is 403 commits ahead with broad FSDP, MoE, Hybrid/Mamba, inference, RL, MiMo, and resharding changes.",
-    "Primus has a large Megatron integration layer with direct dependencies on upstream training, distributed, optimizer, MoE, and pipeline internals."
-  ],
-  "dashboard_summary": {
-    "headline": "Megatron is 403 commits behind upstream main, with broad API and capability drift across training, MoE, FSDP, inference, and Hybrid paths.",
-    "recommendation": "plan sync",
-    "why_it_matters": [
-      "Primus pins Megatron Core source version 0.16.0rc0 while upstream main declares 0.18.0.",
-      "The gap spans 1300 changed files, so this is a broad upstream movement.",
-      "Package dependency metadata is mostly unchanged; the main risk is API and capability drift.",
-      "Primus trainer, adapter, optimizer, MoE, and patch layers depend on upstream internals."
-    ],
-    "feature_deltas": [
-      "Megatron-FSDP fixes for mixed precision, MXFP8, uneven DTensor, frozen parameters, and async save",
-      "MoE shared expert overlap, FlexDispatcher support, router score function, and A2A combine backprop overlap",
-      "Hybrid/Mamba model path expansion and Mamba-to-Hybrid naming movement",
-      "Inference CUDA graphs, prefix caching, per-block MoE routing storage, and text generation controller updates",
-      "MiMo, RL, resharding, and broader test coverage updates"
-    ],
-    "dependency_deltas": [
-      "pyproject.toml project dependencies are unchanged in the compared range.",
-      "megatron/core/requirements.txt remains torch and packaging.",
-      "Source-declared Megatron Core version moved from 0.16.0rc0 to 0.18.0.",
-      "Selected workflow and documentation paths changed, but the primary drift is upstream API/capability movement."
-    ],
-    "integration_risks": [
-      "Primus MegatronTrainer imports and wraps megatron.training, checkpointing, initialization, distributed, optimizer, and dataset internals.",
-      "Primus patches Megatron argument parsing, training flow, scheduler, recompute, Muon optimizer, and runtime hooks.",
-      "Primus Turbo paths depend on transformer engine spec provider, MoE token dispatcher, transformer config, and tensor-parallel internals.",
-      "FSDP, MoE/router, inference, and Hybrid/Mamba upstream movement overlap with existing Primus patch surfaces."
-    ]
-  },
-  "artifacts": [
-    {
-      "label": "Detailed Report (PDF)",
-      "path": "./reports/megatron/upstream-main/report.pdf",
-      "format": "pdf",
-      "language": "en",
-      "kind": "detail",
-      "primary": true
-    },
-    {
-      "label": "One-Page Summary (PDF)",
-      "path": "./reports/megatron/upstream-main/summary.pdf",
-      "format": "pdf",
-      "language": "en",
-      "kind": "summary",
-      "primary": false
-    }
-  ]
-}
diff --git a/docs/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json b/docs/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json
deleted file mode 100644
index 83fd07bee..000000000
--- a/docs/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json
+++ /dev/null
@@ -1,88 +0,0 @@
-{
-  "id": "torchtitan-upstream-main-2026-04-21",
-  "title": "Primus TorchTitan vs upstream main",
-  "backend": {
-    "key": "torchtitan",
-    "label": "TorchTitan"
-  },
-  "generated_at": "2026-04-21",
-  "status": "verified",
-  "scope": "Current TorchTitan bundled in Primus vs upstream pytorch/torchtitan origin/main.",
-  "local": {
-    "source_path": "third_party/torchtitan",
-    "version": "0.1.0",
-    "commit": "5fb7cc2e",
-    "commit_date": "2025-10-15"
-  },
-  "upstream": {
-    "repo": "https://github.com/pytorch/torchtitan",
-    "ref": "origin/main",
-    "version": "0.2.2",
-    "commit": "fc54b897",
-    "commit_date": "2026-04-20"
-  },
-  "stats": {
-    "commit_gap": 493,
-    "diff_files": 447,
-    "insertions": 56071,
-    "deletions": 17716
-  },
-  "integration": {
-    "backend_files": 90,
-    "tracked_files": 147,
-    "integration_model": "third_party/torchtitan + Primus outer adapter / trainer / patches"
-  },
-  "highlights": [
-    "Primus bundles TorchTitan 0.1.0 while upstream main is at 0.2.2.",
-    "Upstream main tracks the latest nightly on CUDA cu130 and ROCm 7.1, while the current Primus baseline is centered on cu126.",
-    "Primus has a non-trivial outer integration layer with direct dependencies on upstream internal paths."
-  ],
-  "dashboard_summary": {
-    "headline": "TorchTitan is 493 commits behind upstream main, with dependency-channel and Primus integration impact.",
-    "recommendation": "urgent sync",
-    "why_it_matters": [
-      "Primus bundles TorchTitan 0.1.0 while upstream main is at 0.2.2.",
-      "Upstream main moved to nightly cu130 and ROCm 7.1 while the Primus baseline remains centered on cu126.",
-      "The gap includes 447 changed files, so this is not a narrow dependency bump.",
-      "Primus trainer, adapter, and patch layers directly depend on upstream internal paths."
-    ],
-    "feature_deltas": [
-      "GraphTrainer precompile and cudagraph paths",
-      "MoE token dispatcher and DeepEP / HybridEP distributed runtime updates",
-      "Fused QKV GQAttention and FlexAttention context-parallel work",
-      "FSDP2 fully_shard and broader distributed runtime updates",
-      "New model families and reorganized shared model abstractions"
-    ],
-    "dependency_deltas": [
-      "README nightly channel moved from nightly/cu126 to nightly/cu130.",
-      "ROCm workflow coverage now targets nightly/rocm7.1.",
-      "torchdata is explicitly installed as 0.12.0.dev20260327 in workflow.",
-      "datasets lower bound moved from >=2.21.0 to >=3.6.0 with CI constraints.",
-      "safetensors, einops, pillow, av, torchvision, numpy, and pyrefly were added or expanded."
-    ],
-    "integration_risks": [
-      "Primus trainer imports torchtitan.config.job_config.JobConfig and torchtitan.train.Trainer.",
-      "Turbo attention patches reference upstream llama3, llama4, and deepseek_v3 internals.",
-      "FP8 and MX patches depend on torchtitan.components.quantization.float8 and mx.",
-      "MoE grouped-mm patch depends on torchtitan.models.moe.moe."
-    ]
-  },
-  "artifacts": [
-    {
-      "label": "Detailed Report (PDF)",
-      "path": "./reports/torchtitan/upstream-main/report.pdf",
-      "format": "pdf",
-      "language": "en",
-      "kind": "detail",
-      "primary": true
-    },
-    {
-      "label": "One-Page Summary (PDF)",
-      "path": "./reports/torchtitan/upstream-main/summary.pdf",
-      "format": "pdf",
-      "language": "en",
-      "kind": "summary",
-      "primary": false
-    }
-  ]
-}
diff --git a/docs/backend-gap/reports/megatron/upstream-main/report.md b/docs/backend-gap/reports/megatron/upstream-main/report.md
deleted file mode 100644
index 94c0e6672..000000000
--- a/docs/backend-gap/reports/megatron/upstream-main/report.md
+++ /dev/null
@@ -1,122 +0,0 @@
-# Primus Megatron vs Upstream `main` Comparison Report
-
-> Date: 2026-04-30
-> Scope: Current Megatron-LM bundled in `Primus` vs upstream `NVIDIA/Megatron-LM` `origin/main`
-
-## High-Level Comparison
-
-| Item | Current Megatron-LM in Primus | Upstream `NVIDIA/Megatron-LM` `main` |
-| --- | --- | --- |
-| Source-declared Megatron Core version | `0.16.0rc0` from `megatron/core/package_info.py` | `0.18.0` from `origin/main:megatron/core/package_info.py` |
-| Pinned commit | `d3528a21` | `0d98cb83` |
-| Commit date | 2026-03-06 | 2026-04-29 |
-| Commit gap | Behind by `403` commits | - |
-| Git relation | `merge-base(HEAD, origin/main) = HEAD` | - |
-| Diff size | `1300 files changed, 615916 insertions, 326835 deletions` | - |
-| Integration model | `third_party/Megatron-LM` + Primus outer trainer / adapter / patches | Upstream mainline |
-| Integration footprint | `primus/backends/megatron/` has about `136` files; report-covered Primus Megatron directories have about `365` files | No Primus integration layer |
-| Private submodule commits | None | - |
-
-## Dependency and Package Metadata Comparison
-
-| Item | Current Megatron-LM in Primus | Upstream `main` |
-| --- | --- | --- |
-| `pyproject.toml` project dependencies | `torch>=2.6.0`, `numpy`, `packaging>=24.2` | No diff in the compared range |
-| `megatron/core/requirements.txt` | `torch`, `packaging` | No diff in the compared range |
-| Source-declared version source | `megatron/core/package_info.py`: `0.16.0rc0` | `0.18.0` with git SHA suffix logic |
-| Dev dependency groups | Present in `pyproject.toml` | No direct diff in the compared range |
-| Workflow / docs surface | Existing CI and docs | About `42` changed entries across selected workflow/docs paths |
-
-## Directory and Capability Differences
-
-### Megatron-FSDP and Distributed Runtime
-
-Upstream `main` has continued Megatron-FSDP and distributed-runtime work:
-
-- Unified and refactored Megatron-FSDP documentation.
-- Fixed FusedAdam `use_decoupled_grad` handling in Megatron-FSDP.
-- Added and fixed mixed-precision / MXFP8 / uneven DTensor / frozen parameter paths.
-- Added support around DCP and FSDP async save.
-- Refined parameter layout, all-gather / reduce-scatter overlap, and precision-aware optimizer behavior.
-
-### MoE, Router, and Expert Parallelism
-
-The upstream gap includes active MoE and expert-parallel changes:
-
-- Improved shared expert overlap and FlexDispatcher support.
-- Added a new router score function.
-- Added NVFP4 native weights for DDP.
-- Fixed non-quantized MoE dispatch padding.
-- Added overlap for A2A combine backprop with wgrad GEMM.
-- Added broader MoE inference and prefix-cache related updates.
-
-### Hybrid / Mamba and Inference
-
-Upstream `main` has expanded the Hybrid/Mamba and inference surface:
-
-- Added `megatron/core/models/hybrid/`.
-- Renamed Mamba model / stack concepts toward Hybrid naming.
-- Added YARN support and DeepSeek Sparse Attention paths for Hybrid/Mamba work.
-- Added CUDA graph support for MTP inference and prefix caching.
-- Added per-block MoE routing storage for prefix caching.
-- Reorganized order of operations in inference context and text generation controller.
-
-### MiMo, RL, and Resharding
-
-Upstream has broader non-core-training surface area:
-
-- Updated `examples/mimo/` and `megatron/core/models/mimo/`.
-- Added `ColocatedBridgeCommunicator` and `MimoOptimizer` related paths.
-- Updated RL agents, token throughput / packing metrics, and RL inference flows.
-- Added and evolved `megatron/core/resharding/` APIs and copy services.
-
-## Change Hotspots
-
-| Area | Representative changes |
-| --- | --- |
-| `megatron/core/distributed/fsdp/` | Megatron-FSDP fixes, docs, DTensor conversion, mixed precision, MXFP8, async save |
-| `megatron/core/transformer/moe/` | Shared expert overlap, router score function, MoE dispatch fixes |
-| `megatron/core/inference/` | CUDA graph / prefix caching / hybrid inference / text generation controller updates |
-| `megatron/core/models/hybrid/` | Added Hybrid model/block/layer allocation/spec files |
-| `megatron/core/models/mimo/` | MiMo communication, config, optimizer, and submodule updates |
-| `megatron/core/resharding/` | Added / evolved resharding planner, execution, transforms, and copy services |
-| `examples/rl/` and `megatron/rl/` | RL agents, inference API, packing metrics, and reward fixes |
-| `tests/` | Broad functional and unit-test churn across MoE, inference, FSDP, SSM, resharding, and training |
-
-## Primus Outer Integration Layer
-
-### Related Directories
-
-The Primus outer integration layer is mainly distributed across:
-
-- `primus/backends/megatron/`
-- `primus/modules/trainer/megatron/`
-- `primus/configs/models/megatron/`
-- `examples/megatron/`
-- `tests/unit_tests/backends/megatron/`
-
-### Directly Referenced Upstream Paths
-
-| Primus code location | Direct upstream dependency path |
-| --- | --- |
-| `primus/modules/trainer/megatron/trainer.py` | `megatron.core.distributed`, `megatron.core.optimizer`, `megatron.training.*`, dataset builders, checkpointing, initialization, parallel state |
-| `primus/modules/trainer/megatron/utils.py` | `megatron.core.parallel_state`, `megatron.training.global_vars`, pipeline-parallel schedules, transformer block/layer helpers |
-| `primus/backends/megatron/megatron_base_trainer.py` | `megatron.training.arguments`, `megatron.training.initialize` |
-| `primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py` | `megatron.core.extensions.transformer_engine`, `megatron.core.transformer.moe.experts`, `megatron.core.models.backends` |
-| `primus/backends/megatron/core/extensions/primus_turbo.py` | `megatron.core.tensor_parallel`, `megatron.core.transformer.moe.token_dispatcher`, `megatron.core.transformer.transformer_config`, `megatron.training.global_vars` |
-| `primus/backends/megatron/patches/turbo/te_spec_provider_patches.py` | `megatron.core.extensions`, `megatron.core.models.gpt.gpt_layer_specs`, `megatron.core.models.gpt.moe_module_specs`, `megatron.core.transformer.multi_token_prediction` |
-| `primus/backends/megatron/patches/muon_optimizer_patches.py` | `megatron.training.training.get_megatron_optimizer` namespace |
-
-## Evidence Sources
-
-- `third_party/Megatron-LM/megatron/core/package_info.py`
-- `third_party/Megatron-LM/pyproject.toml`
-- `third_party/Megatron-LM/megatron/core/requirements.txt`
-- `third_party/Megatron-LM/README.md`
-- `third_party/Megatron-LM/.github/workflows/*`
-- [NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
-- `primus/backends/megatron/*`
-- `primus/modules/trainer/megatron/*`
-- `primus/configs/models/megatron/*`
-- `examples/megatron/*`
-- `tests/unit_tests/backends/megatron/*`
diff --git a/docs/backend-gap/reports/megatron/upstream-main/summary.md b/docs/backend-gap/reports/megatron/upstream-main/summary.md
deleted file mode 100644
index 57819f857..000000000
--- a/docs/backend-gap/reports/megatron/upstream-main/summary.md
+++ /dev/null
@@ -1,69 +0,0 @@
-# Primus Megatron Upstream Gap One-Page Summary
-
-> Date: 2026-04-30
-
-## High-Level Comparison
-
-| Item | Current Megatron-LM in Primus | Upstream `NVIDIA/Megatron-LM` `main` | Impact |
-| --- | --- | --- | --- |
-| Source-declared Megatron Core version | `0.16.0rc0` from `megatron/core/package_info.py` | `0.18.0` | Spans multiple upstream version steps |
-| Pinned commit | `d3528a21` | `0d98cb83` | Current submodule is materially behind |
-| Commit gap | Behind by `403` commits | - | Not a small patch-level gap |
-| Diff size | `1300 files changed, 615916 insertions, 326835 deletions` | - | Very broad upstream churn |
-| Integration model | `third_party/Megatron-LM` + Primus trainer / adapter / patches | Upstream mainline | Upgrade is more than a submodule bump |
-| Integration footprint | `primus/backends/megatron/` has about `136` files; report-covered Primus Megatron directories have about `365` files | No Primus integration layer | Upgrade blast radius is large |
-| Dependency metadata | `pyproject.toml` and `megatron/core/requirements.txt` unchanged in compared range | Same dependency metadata | Main concern is API / capability drift, not dependency metadata drift |
-
-## Dependency / Package Facts
-
-| Item | Current Megatron-LM in Primus | Upstream `main` |
-| --- | --- | --- |
-| `pyproject.toml` project dependencies | `torch>=2.6.0`, `numpy`, `packaging>=24.2` | No direct diff |
-| `megatron/core/requirements.txt` | `torch`, `packaging` | No direct diff |
-| Source-declared version source | `megatron/core/package_info.py`: `0.16.0rc0` | `0.18.0` with git SHA suffix logic |
-
-## Representative Upstream Changes
-
-| Area | Representative changes |
-| --- | --- |
-| Megatron-FSDP | Documentation refactor; mixed precision, MXFP8, uneven DTensor, frozen parameter, async save, and optimizer fixes |
-| MoE / Router | Shared expert overlap, FlexDispatcher support, new router score function, non-quantized dispatch padding fix, A2A combine backprop overlap |
-| Hybrid / Mamba | Added `megatron/core/models/hybrid`; renamed Mamba concepts toward Hybrid; added YARN and DeepSeek Sparse Attention paths |
-| Inference | CUDA graph support for MTP inference, prefix caching, per-block MoE routing storage, text generation controller reordering |
-| MiMo / RL | MiMo communication and optimizer updates; RL agents, inference flows, token throughput and packing metrics |
-| Resharding | Updated planner, execution, transforms, and copy services |
-| Tests | Broad functional and unit-test changes across MoE, inference, FSDP, SSM, resharding, and training |
-
-## Primus Outer Integration Layer
-
-The Primus Megatron integration layer is broad and directly coupled to upstream internals:
-
-- `primus/backends/megatron/`
-- `primus/modules/trainer/megatron/`
-- `primus/configs/models/megatron/`
-- `examples/megatron/`
-- `tests/unit_tests/backends/megatron/`
-
-Direct upstream dependency paths include:
-
-- `megatron.training.*`
-- `megatron.core.distributed.*`
-- `megatron.core.optimizer.*`
-- `megatron.core.transformer.moe.*`
-- `megatron.core.extensions.transformer_engine`
-- `megatron.core.models.gpt.*`
-- `megatron.core.pipeline_parallel.*`
-
-## Recommendation
-
-Treat this as a planned sync, not a trivial dependency bump. The largest risks are API compatibility in trainer initialization, distributed/FSDP paths, MoE/router behavior, inference/Hybrid changes, and Primus patch compatibility.
-
-## Evidence Sources
-
-- `third_party/Megatron-LM/megatron/core/package_info.py`
-- `third_party/Megatron-LM/pyproject.toml`
-- `third_party/Megatron-LM/megatron/core/requirements.txt`
-- [NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
-- `primus/backends/megatron/*`
-- `primus/modules/trainer/megatron/*`
-- `tests/unit_tests/backends/megatron/*`
diff --git a/docs/backend-gap/reports/torchtitan/upstream-main/report.md b/docs/backend-gap/reports/torchtitan/upstream-main/report.md
deleted file mode 100644
index 8e9ea2df3..000000000
--- a/docs/backend-gap/reports/torchtitan/upstream-main/report.md
+++ /dev/null
@@ -1,147 +0,0 @@
-# Primus TorchTitan vs Upstream `main` Comparison Report
-
-> Date: 2026-04-21
-> Scope: Current TorchTitan bundled in `Primus` vs upstream `pytorch/torchtitan` `origin/main`
-
-## High-Level Comparison
-
-| Item | Current TorchTitan in Primus | Upstream `pytorch/torchtitan` `main` |
-| --- | --- | --- |
-| Submodule version | `0.1.0` | `0.2.2` |
-| Pinned commit | `5fb7cc2e` | `fc54b897` |
-| Commit date | 2025-10-15 | 2026-04-20 |
-| Commit gap | Behind by `493` commits | - |
-| Git relation | `merge-base(HEAD, origin/main) = HEAD` | - |
-| Diff size | `447 files changed, 56071 insertions, 17716 deletions` | - |
-| Integration model | `third_party/torchtitan` + Primus outer layer (`adapter / trainer / patches`) | Upstream mainline |
-| Integration footprint | `primus/backends/torchtitan/` has about `90` files; total across report-covered directories is about `147` files | No Primus integration layer |
-| Private submodule commits | None | - |
-| Extra private requirements | `requirements-torchtitan.txt` has comments only, with no effective dependency entries | - |
-
-## Torch / TorchAO / Dependency Comparison
-
-### Install Channels and Version Semantics
-
-| Item | Current TorchTitan in Primus | Upstream `main` |
-| --- | --- | --- |
-| README nightly channel | `nightly/cu126` | `nightly/cu130` |
-| Workflow install channel | `nightly/cu126` | `nightly/cu130`; ROCm uses `nightly/rocm7.1` |
-| Workflow `torch-version` parameter | No explicit fixed version | Empty string in `set-matrix.yaml` |
-| `v0.1.0` release anchor | `torch-2.8.0.dev20250617+cu126` / `torchao-0.12.0.dev20250617+cu126` | - |
-| `v0.2.2` release anchor | - | `torch-2.12.0.dev20260220+cu126` / `torchao-0.17.0.dev20260220+cu126` |
-
-### Python Dependency Differences
-
-| Dependency | Current TorchTitan in Primus | Upstream `main` |
-| --- | --- | --- |
-| `torchdata` | `>=0.8.0` | Explicitly installed as `0.12.0.dev20260327` in workflow |
-| `datasets` | `>=2.21.0` | `>=3.6.0`, constrained to `<4.8.0` in CI |
-| `tokenizers` | No fixed lower bound | `>=0.15.0` |
-| `safetensors` | Not present | Added |
-| `wandb` | Dev-only dependency | Moved into runtime dependencies |
-| `einops` | Not present | Added |
-| `pillow` | Not present | Added |
-| `av` | Not present | Added for VLM-related dependencies |
-| `torchvision` | Not present | Added in VLM / CPU tests |
-| `expecttest` | Pinned at `0.1.6` | `>=0.2.0` |
-| `pyrefly` | Not present | `==0.45.1` |
-| `numpy` | Not present | Added to dev / CI dependencies |
-| `tyro` | No fixed lower bound | Raised to `>=1.0.5` in CI dependencies |
-| `tomli` | Present in runtime dependencies | Removed from runtime dependencies |
-
-## Directory and Capability Differences
-
-### Model Directories
-
-The current TorchTitan in Primus still follows an earlier model layout. Upstream `main` has added new model families and shared abstractions:
-
-- Added `torchtitan/models/common/` with shared modules: `attention / decoder / embedding / feed_forward / linear / moe / param_init / rmsnorm / rope / token_dispatcher`
-- Added `torchtitan/models/gpt_oss/`
-- Added `torchtitan/models/qwen3_vl/`
-- Added `torchtitan/models/flux/`; `flux` was moved from `experiments/flux`
-- `llama3 / llama4 / qwen3 / deepseek_v3` were reorganized into a more unified shape around `config_registry.py`, `parallelize.py`, and `state_dict_adapter.py`
-
-### `experiments/` Directory
-
-The current TorchTitan in Primus keeps an earlier experiments layout. Upstream `main` added or moved the following:
-
-- Added `torchtitan/experiments/autoparallel/`
-- Added `torchtitan/experiments/graph_trainer/`
-- Added `torchtitan/experiments/rl/`
-- Added `torchtitan/experiments/transformers_modeling_backend/`
-- Added `torchtitan/experiments/ft/`; content from `components/ft` moved here
-- Main content from `torchtitan/experiments/flux/` moved into `torchtitan/models/flux/`
-- Multiple files in `torchtitan/experiments/simple_fsdp/` were removed or moved
-- Multiple files in `torchtitan/experiments/torchcomms/` were removed
-
-### `distributed/` and `components/`
-
-The current TorchTitan in Primus keeps older distributed/components layouts. Upstream `main` added or continuously evolved these paths:
-
-- Added `torchtitan/distributed/compile.py`
-- Added `torchtitan/distributed/context_parallel.py`
-- Added `torchtitan/distributed/deepep/`
-- Added `torchtitan/distributed/fsdp.py`
-- `torchtitan/distributed/tensor_parallel.py` continuously updated
-- `torchtitan/distributed/pipeline_parallel.py` continuously updated
-- `torchtitan/distributed/expert_parallel.py` continuously updated
-- Added `torchtitan/components/quantization/module_utils.py`
-- `torchtitan/components/quantization/float8.py` continuously updated
-- `torchtitan/components/quantization/mx.py` continuously updated
-- `torchtitan/components/metrics.py` continuously updated
-- `torchtitan/components/optimizer.py` continuously updated
-- `torchtitan/components/tokenizer.py` continuously updated
-
-## Change Hotspots
-
-| Area | Representative changes |
-| --- | --- |
-| `tests/unit_tests/` | Broader unit-test coverage |
-| `.github/workflows/` | Added `integration_test_4gpu_rl.yaml`, `integration_test_8gpu_autoparallel.yaml`, `integration_test_8gpu_graph_trainer.yaml`, `integration_test_8gpu_rl_h100.yaml`, `integration_test_8gpu_transformers_modeling_backend.yaml`, `set-matrix.yaml` |
-| `torchtitan/experiments/graph_trainer/` | Added `compile.py / cudagraph.py / precompile.py / passes.py / storage.py`; covers `llama3 / deepseek_v3 / qwen3`; includes `tests/test_profiler.py` |
-| `torchtitan/models/common/` | Added `attention / decoder / embedding / feed_forward / linear / moe / param_init / rmsnorm / rope / token_dispatcher` |
-| `docs/` | Added `docs/mxfp8.md`, `docs/bf16_optimizer_states.md`, and updated docs such as `release / debugging / metrics / datasets / checkpoint` |
-| `torchtitan/models/flux/` | `flux` moved from `experiments/flux` into `models/flux` |
-| `torchtitan/distributed/` | Added `compile.py`, `context_parallel.py`, `deepep/deepep.py`, `deepep/hybridep.py`, `fsdp.py` |
-| `torchtitan/experiments/rl/` | Added `actors/`, `models/vllm_wrapper.py`, `simple_grpo_sum_digits.py` |
-| `torchtitan/components/` | Added `module_utils`; continuous updates in `float8 / mx / metrics / optimizer / tokenizer / validate` |
-| `torchtitan/models/gpt_oss/` | Added `gpt_oss` |
-| `torchtitan/models/qwen3_vl/` | Added `qwen3_vl` |
-| `torchtitan/experiments/autoparallel/` | Added `llama3`, `deepseek_v3`, `local_map_deepseek_v3` |
-| `torchtitan/experiments/transformers_modeling_backend/` | Added transformers-based modeling backend experiments |
-
-## Primus Outer Integration Layer
-
-### Related Directories
-
-The Primus outer integration layer is mainly distributed across:
-
-- `primus/backends/torchtitan/`
-- `primus/modules/trainer/torchtitan/`
-- `primus/configs/modules/torchtitan/`
-- `examples/torchtitan/`
-- `runner/helpers/hooks/train/pretrain/torchtitan/`
-- `tests/unit_tests/backends/torchtitan/`
-- `tests/trainer/test_torchtitan_trainer.py`
-
-### Directly Referenced Upstream Paths
-
-| Primus code location | Direct upstream dependency path |
-| --- | --- |
-| `primus/modules/trainer/torchtitan/pre_trainer.py` | `torchtitan.config.job_config.JobConfig`, `torchtitan.train.Trainer` |
-| `primus/backends/torchtitan/patches/turbo/attention_patches.py` | `torchtitan.models.llama3.model.model`, `torchtitan.models.llama4.model.model`, `torchtitan.models.deepseek_v3.model.model` |
-| `primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py` | `torchtitan.components.quantization.float8` |
-| `primus/backends/torchtitan/patches/turbo/mx_linear_patches.py` | `torchtitan.components.quantization.mx` |
-| `primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py` | `torchtitan.models.moe.moe` |
-
-## Evidence Sources
-
-- `third_party/torchtitan/pyproject.toml`
-- `third_party/torchtitan/README.md`
-- `third_party/torchtitan/docs/release.md`
-- `third_party/torchtitan/.github/workflows/*`
-- [pytorch/torchtitan](https://github.com/pytorch/torchtitan)
-- [pytorch/torchtitan/docs/release.md](https://github.com/pytorch/torchtitan/blob/main/docs/release.md)
-- `primus/backends/torchtitan/*`
-- `primus/modules/trainer/torchtitan/*`
-- `docs/backends/torchtitan/patch-notes.md`
diff --git a/docs/backend-gap/reports/torchtitan/upstream-main/summary.md b/docs/backend-gap/reports/torchtitan/upstream-main/summary.md
deleted file mode 100644
index 69ae788cd..000000000
--- a/docs/backend-gap/reports/torchtitan/upstream-main/summary.md
+++ /dev/null
@@ -1,82 +0,0 @@
-# Primus TorchTitan Upstream Gap One-Page Summary
-
-> Date: 2026-04-21
-
-## High-Level Comparison
-
-| Item | Current TorchTitan in Primus | Upstream `pytorch/torchtitan` `main` | Impact |
-| --- | --- | --- | --- |
-| Submodule version | `0.1.0` | `0.2.2` | Spans multiple upstream release iterations |
-| Pinned commit | `5fb7cc2e` | `fc54b897` | Current submodule is significantly behind |
-| Commit gap | Behind by `493` commits | - | Not a small patch-level gap anymore |
-| Integration model | `third_party/torchtitan` + Primus outer layer (`adapter / trainer / patches`) | Upstream mainline | Upgrade is more than a submodule bump |
-| Integration footprint | `primus/backends/torchtitan/` has about `90` files; total across summary-covered directories is about `147` files | No Primus integration layer | Upgrade blast radius is large |
-| Private submodule commits | None | - | Git history remains traceable and clean |
-| Extra private requirements | `requirements-torchtitan.txt` has comments only, with no effective dependency entries | - | Primus mainly reuses upstream dependencies |
-| Torch / TorchAO version semantics | `nightly/cu126`; `v0.1.0` anchors `torch-2.8.0.dev20250617+cu126` / `torchao-0.12.0.dev20250617+cu126` | GitHub current `main` tracks the latest nightly at the time (`CUDA cu130`, `ROCm rocm7.1`) | The key comparison is `cu126` generation vs `cu130/rocm7.1` generation |
-| `torchdata` | `>=0.8.0` | `0.12.0.dev20260327` (workflow) | Data stack baseline increased |
-| `datasets` | `>=2.21.0` | `>=3.6.0, <4.8.0` | Compatibility and streaming behavior changed |
-| `tokenizers` | No fixed lower bound | `>=0.15.0` | Tokenizer baseline increased |
-| Runtime added deps | No runtime `safetensors / wandb / einops / pillow` | Added | Runtime environment is heavier |
-| Multimodal added deps | No `av / torchvision` | Added | Multimodal environment requirements increased |
-| New upstream capabilities | Existing training main path | `graph_trainer / autoparallel / rl / qwen3_vl / gpt_oss` | Upstream capability surface is much broader |
-
-## Torch / TorchAO / Dependencies
-
-| Item | Current TorchTitan in Primus | Upstream `main` |
-| --- | --- | --- |
-| README nightly channel | `nightly/cu126` | `nightly/cu130` |
-| Workflow install channel | `nightly/cu126` | `nightly/cu130`; ROCm uses `nightly/rocm7.1` |
-| Workflow `torch-version` parameter | No explicit fixed version | Empty string in `set-matrix.yaml` |
-| `v0.1.0` release anchor | `torch-2.8.0.dev20250617+cu126` / `torchao-0.12.0.dev20250617+cu126` | - |
-| `torchdata` | `>=0.8.0` | Explicitly installed as `0.12.0.dev20260327` in workflow |
-| `datasets` | `>=2.21.0` | `>=3.6.0, <4.8.0` |
-| `tokenizers` | No fixed lower bound | `>=0.15.0` |
-| `safetensors` | No | Added |
-| `wandb` | Dev-only dependency | Moved into runtime dependencies |
-| `einops` | No | Added |
-| `pillow` | No | Added |
-| `av` | No | Added |
-| `torchvision` | No | Added |
-
-## Representative Upstream Changes
-
-| Area | Representative changes |
-| --- | --- |
-| `models/` | Added `gpt_oss`, `qwen3_vl`; moved `flux` into `models/flux`; added shared `models/common` layer |
-| `distributed/` | Added `compile.py`, `context_parallel.py`, `deepep/`, `fsdp.py` |
-| `experiments/` | Added `autoparallel`, `graph_trainer`, `rl`, `transformers_modeling_backend`, `ft` |
-| `components/` | Added `module_utils`; continuous updates in `float8 / mx / metrics / optimizer / tokenizer / validate` |
-| `.github/workflows/` | Added `integration_test_4gpu_rl.yaml`, `integration_test_8gpu_autoparallel.yaml`, `integration_test_8gpu_graph_trainer.yaml`, `integration_test_8gpu_transformers_modeling_backend.yaml`, `set-matrix.yaml` |
-| `docs/` | Added `docs/mxfp8.md`, `docs/bf16_optimizer_states.md`, and updated docs such as `release / debugging / metrics / datasets / checkpoint` |
-
-## Primus Outer Integration Layer
-
-The Primus outer integration layer is mainly distributed across:
-
-- `primus/backends/torchtitan/`
-- `primus/modules/trainer/torchtitan/`
-- `primus/configs/modules/torchtitan/`
-- `examples/torchtitan/`
-- `runner/helpers/hooks/train/pretrain/torchtitan/`
-- `tests/unit_tests/backends/torchtitan/`
-- `tests/trainer/test_torchtitan_trainer.py`
-
-Directly referenced upstream paths include:
-
-- `torchtitan.config.job_config.JobConfig`
-- `torchtitan.train.Trainer`
-- `torchtitan.models.llama3.model.model`
-- `torchtitan.models.llama4.model.model`
-- `torchtitan.models.deepseek_v3.model.model`
-- `torchtitan.components.quantization.float8`
-- `torchtitan.components.quantization.mx`
-- `torchtitan.models.moe.moe`
-
-## Evidence Sources
-
-- `third_party/torchtitan/README.md`
-- `third_party/torchtitan/docs/release.md`
-- `third_party/torchtitan/.github/workflows/*`
-- [pytorch/torchtitan](https://github.com/pytorch/torchtitan)
-- [pytorch/torchtitan/docs/release.md](https://github.com/pytorch/torchtitan/blob/main/docs/release.md)
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 000000000..76cefc389
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,47 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+import re
+from pathlib import Path
+
+
+def _get_version():
+    init = Path(__file__).parent.parent / "primus" / "__init__.py"
+    match = re.search(r'^__version__\s*=\s*["\']([^"\']+)["\']', init.read_text(), re.MULTILINE)
+    if not match:
+        raise ValueError("Could not find __version__ in primus/__init__.py")
+    return match.group(1)
+
+
+# Project info
+version = _get_version()
+release = version
+project = f"AMD Primus {version}"
+author = "Advanced Micro Devices, Inc."
+copyright = "Copyright (c) %Y Advanced Micro Devices, Inc. All rights reserved."
+
+# Theme-related configs
+html_theme = "rocm_docs_theme"
+html_theme_options = {
+    "flavor": "ai-ecosystem",
+    "link_main_doc": True,
+    "repository_url": "https://github.com/AMD-AGI/Primus",
+    "use_repository_button": True,
+    "use_issues_button": True,
+    "use_download_button": True,
+}
+html_title = project
+
+# Sphinx extension-related configs
+extensions = ["rocm_docs"]
+external_toc_path = "./sphinx/_toc.yml"
+external_projects_current_project = "primus"
+
+# Publish the llms.txt index at the docs site root and let
+# rocm-docs-core generate llms-full.txt after each build (the llms.txt standard,
+# https://llmstxt.org/). See the rocm-docs-core guide:
+# https://rocm.docs.amd.com/projects/rocm-docs-core/en/latest/user_guide/llms.html
+rocm_docs_generate_llms = True
diff --git a/docs/license.md b/docs/license.md
new file mode 100644
index 000000000..1f8761f24
--- /dev/null
+++ b/docs/license.md
@@ -0,0 +1,4 @@
+# License
+
+```{include} ../LICENSE
+```
diff --git a/docs/preflight.md b/docs/preflight.md
deleted file mode 100644
index b4290e04e..000000000
--- a/docs/preflight.md
+++ /dev/null
@@ -1,78 +0,0 @@
-# Preflight
-
-`preflight` is Primus’ cluster diagnostic tool. It can generate a **fast info report** (host/GPU/network) and can also run **performance tests** (GEMM + intra/inter-node comm) to help spot misconfiguration or outliers before large distributed runs.
-
-- **User-facing entry**: `primus-cli … -- preflight [args]`
-- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py`
-
-## Quick start
-
-### Info report only (fast)
-
-```bash
-primus-cli direct -- preflight --host --gpu --network
-```
-
-### Full preflight (info + perf tests)
-
-```bash
-primus-cli direct -- preflight
-```
-
-### Perf tests only
-
-```bash
-primus-cli direct -- preflight --perf-test
-```
-
-## Common usage (Slurm)
-
-Info report only (fast):
-
-```bash
-primus-cli slurm srun -N 4 -- preflight --host --gpu --network
-```
-
-Full preflight (info + perf tests):
-
-```bash
-primus-cli slurm srun -N 4 -- preflight
-```
-
-Perf tests only:
-
-```bash
-primus-cli slurm srun -N 4 -- preflight --perf-test
-```
-
-## CLI flags
-
-Selection:
-- `--host`: host info (CPU, memory, PCIe)
-- `--gpu`: GPU info
-- `--network`: network info
-- `--perf-test`: run perf tests only (GEMM + comm). This is slower.
-
-Reporting:
-- `--dump-path`: output directory (default: `output/preflight`)
-- `--report-file-name`: base report name (default: `preflight_report`)
-- `--disable-pdf`: disable PDF generation
-
-Perf-test extras:
-- `--plot`: generate plots (only used with `--perf-test`)
-
-Backward compatibility:
-- `--check-host/--check-gpu/--check-network` are supported as aliases for `--host/--gpu/--network`.
-
-## Outputs
-
-By default, outputs are written under `output/preflight`.
-
-Typical report files:
-- `preflight_report.md` / `preflight_report.pdf`: **info report** (host/GPU/network)
-- `preflight_report_perf.md` / `preflight_report_perf.pdf`: **perf report** (GEMM + comm tests)
-
-## Notes
-
-- For multi-node runs, use `primus-cli slurm …` (or your preferred launcher) so distributed environment variables are set correctly.
-- If you only want a quick environment snapshot, prefer `--host --gpu --network`.
diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in
new file mode 100644
index 000000000..c033923bf
--- /dev/null
+++ b/docs/sphinx/_toc.yml.in
@@ -0,0 +1,145 @@
+# Variables of the form ${<variable>} are substituted, currently the following
+# list is supported:
+#   - ${branch} (or {branch}) the name of the current branch
+#   - ${url} (or {url}) github url of the current project
+#   - ${project:<project_name>} base url of the documentation of <project_name>
+#     based on intersphinx_mapping.
+# These comments will also be removed.
+defaults:
+  numbered: false
+root: 01-getting-started/overview.md
+subtrees:
+  - caption: Getting started
+    entries:
+      - file: 01-getting-started/quickstart.md
+        title: Quickstart
+      - file: 01-getting-started/installation.md
+        title: Installation and setup
+      - file: 01-getting-started/glossary.md
+        title: Glossary
+
+  - caption: User guide
+    entries:
+      - file: 02-user-guide/primus-tools.md
+        title: Primus tools
+      - file: 02-user-guide/cli-reference.md
+        title: CLI reference
+      - file: 02-user-guide/configuration-system.md
+        title: Configuration system
+      - file: 02-user-guide/pretraining.md
+        title: Pretraining workflows
+      - file: 02-user-guide/training-recipes.md
+        title: Backend training recipes
+      - file: 02-user-guide/posttraining.md
+        title: Post-training workflows
+      - file: 02-user-guide/node-smoke-test-instruction.md
+        title: Node-smoke test instruction
+      - file: 02-user-guide/preflight.md
+        title: Preflight diagnostics
+      - file: 02-user-guide/preflight-without-container.md
+        title: Run preflight without a container
+      - file: 02-user-guide/benchmarking.md
+        title: Benchmark suite
+      - file: 02-user-guide/projection.md
+        title: Memory and performance projection
+      - file: 02-user-guide/tuning-agent.md
+        title: Tuning agent
+
+  - caption: Configuration reference
+    entries:
+      - file: 03-configuration-reference/megatron-parameters.md
+        title: Megatron backend
+      - file: 03-configuration-reference/torchtitan-parameters.md
+        title: TorchTitan backend
+      - file: 03-configuration-reference/megatron-bridge-parameters.md
+        title: Megatron Bridge backend
+      - file: 03-configuration-reference/maxtext-parameters.md
+        title: MaxText backend
+      - file: 03-configuration-reference/environment-variables.md
+        title: Environment variables
+
+  - caption: Technical guides
+    entries:
+      - file: 04-technical-guides/parallelism-strategies.md
+        title: Parallelism strategies
+      - file: 04-technical-guides/parallelism-configuration.md
+        title: Parallelism configuration
+      - file: 04-technical-guides/multi-node-networking.md
+        title: Multi-node networking
+      - file: 04-technical-guides/collective-operations.md
+        title: Collective operations (NCCL/RCCL)
+      - file: 04-technical-guides/data-preparation.md
+        title: Data preparation
+      - file: 04-technical-guides/checkpoint-management.md
+        title: Checkpoint management
+      - file: 04-technical-guides/moe-training.md
+        title: MoE training
+      - file: 04-technical-guides/mega-moe.md
+        title: MegaMoE fused MoE layer
+      - file: 04-technical-guides/performance-tuning.md
+        title: Performance tuning
+      - file: 04-technical-guides/profiling-and-observability.md
+        title: Profiling and observability
+      - file: 04-technical-guides/logging-and-experiment-tracking.md
+        title: Logging and experiment tracking
+      - file: 04-technical-guides/fault-tolerance-and-elastic-training.md
+        title: Fault tolerance and elastic training
+      - file: 04-technical-guides/determinism-and-reproducibility.md
+        title: Determinism and reproducibility
+      - file: 04-technical-guides/native-sft-lora.md
+        title: Native SFT LoRA
+      - file: 04-technical-guides/diffusion-models/architecture_overview.md
+        title: Diffusion model architecture
+        subtrees:
+          - entries:
+              - file: 04-technical-guides/diffusion-models/flux_architecture.md
+                title: Flux architecture
+              - file: 04-technical-guides/diffusion-models/fp8_training.md
+                title: FP8 training
+              - file: 04-technical-guides/diffusion-models/mxfp4_training.md
+                title: MXFP4 training
+              - file: 04-technical-guides/diffusion-models/data_preprocessing.md
+                title: Data preprocessing
+              - file: 04-technical-guides/diffusion-models/energon_integration.md
+                title: Energon integration
+              - file: 04-technical-guides/diffusion-models/adding_new_models.md
+                title: Adding new models
+              - file: 04-technical-guides/diffusion-models/api_reference.md
+                title: API reference
+
+  - caption: Operations
+    entries:
+      - file: 05-operations/deployment.md
+        title: Deployment
+      - file: 05-operations/monitoring-logging.md
+        title: Monitoring and logging
+      - file: 05-operations/troubleshooting.md
+        title: Troubleshooting
+      - file: 05-operations/security.md
+        title: Security considerations
+
+  - caption: Developer guide
+    entries:
+      - file: 06-developer-guide/architecture.md
+        title: Architecture overview
+      - file: 06-developer-guide/cli-architecture.md
+        title: CLI architecture
+      - file: 06-developer-guide/contributing.md
+        title: Contributing
+      - file: 06-developer-guide/adding-models.md
+        title: Adding model configurations
+      - file: 06-developer-guide/extending-backends.md
+        title: Extending backends
+      - file: 06-developer-guide/model-support-matrix.md
+        title: Model support matrix
+      - file: 06-developer-guide/testing.md
+        title: Testing
+      - file: 06-developer-guide/tooling.md
+        title: Tooling
+      - file: 06-developer-guide/backend-patch-notes.md
+        title: Backend patch notes
+
+  - caption: About
+    entries:
+      - file: license.md
+        title: License
diff --git a/docs/sphinx/requirements.in b/docs/sphinx/requirements.in
new file mode 100644
index 000000000..2e037fb21
--- /dev/null
+++ b/docs/sphinx/requirements.in
@@ -0,0 +1 @@
+rocm-docs-core[llms]==1.38.0
diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt
new file mode 100644
index 000000000..f7587f5c5
--- /dev/null
+++ b/docs/sphinx/requirements.txt
@@ -0,0 +1,285 @@
+#
+# This file is autogenerated by pip-compile with Python 3.12
+# by the following command:
+#
+#    pip-compile docs/sphinx/requirements.in
+#
+accessible-pygments==0.0.5
+    # via pydata-sphinx-theme
+alabaster==1.0.0
+    # via sphinx
+asttokens==3.0.2
+    # via stack-data
+attrs==26.1.0
+    # via
+    #   jsonschema
+    #   jupyter-cache
+    #   referencing
+babel==2.18.0
+    # via
+    #   pydata-sphinx-theme
+    #   sphinx
+beautifulsoup4==4.15.0
+    # via pydata-sphinx-theme
+breathe==4.36.0
+    # via rocm-docs-core
+certifi==2026.6.17
+    # via requests
+cffi==2.1.0
+    # via
+    #   cryptography
+    #   pynacl
+charset-normalizer==3.4.9
+    # via requests
+click==8.4.2
+    # via
+    #   jupyter-cache
+    #   sphinx-external-toc
+comm==0.2.3
+    # via ipykernel
+cryptography==49.0.0
+    # via pyjwt
+debugpy==1.8.21
+    # via ipykernel
+decorator==5.3.1
+    # via ipython
+docutils==0.22.4
+    # via
+    #   myst-parser
+    #   pydata-sphinx-theme
+    #   sphinx
+    #   sphinx-markdown-builder
+executing==2.2.1
+    # via stack-data
+fastjsonschema==2.21.2
+    # via
+    #   nbformat
+    #   rocm-docs-core
+gitdb==4.0.12
+    # via gitpython
+gitpython==3.1.53
+    # via rocm-docs-core
+greenlet==3.5.3
+    # via sqlalchemy
+idna==3.18
+    # via requests
+imagesize==2.0.0
+    # via sphinx
+importlib-metadata==9.0.0
+    # via
+    #   jupyter-cache
+    #   myst-nb
+ipykernel==7.3.0
+    # via myst-nb
+ipython==9.15.0
+    # via
+    #   ipykernel
+    #   myst-nb
+ipython-pygments-lexers==1.1.1
+    # via ipython
+jedi==0.20.0
+    # via ipython
+jinja2==3.1.6
+    # via
+    #   myst-parser
+    #   sphinx
+jsonschema==4.26.0
+    # via nbformat
+jsonschema-specifications==2025.9.1
+    # via jsonschema
+jupyter-cache==1.0.1
+    # via myst-nb
+jupyter-client==8.9.1
+    # via
+    #   ipykernel
+    #   nbclient
+jupyter-core==5.9.1
+    # via
+    #   ipykernel
+    #   jupyter-client
+    #   nbclient
+    #   nbformat
+markdown-it-py==4.2.0
+    # via
+    #   mdit-py-plugins
+    #   myst-parser
+markupsafe==3.0.3
+    # via jinja2
+matplotlib-inline==0.2.2
+    # via
+    #   ipykernel
+    #   ipython
+mdit-py-plugins==0.6.1
+    # via myst-parser
+mdurl==0.1.2
+    # via markdown-it-py
+myst-nb==1.4.0
+    # via rocm-docs-core
+myst-parser==5.1.0
+    # via myst-nb
+nbclient==0.11.0
+    # via
+    #   jupyter-cache
+    #   myst-nb
+nbformat==5.10.4
+    # via
+    #   jupyter-cache
+    #   myst-nb
+    #   nbclient
+nest-asyncio2==1.7.2
+    # via ipykernel
+packaging==26.2
+    # via
+    #   ipykernel
+    #   pydata-sphinx-theme
+    #   sphinx
+parso==0.8.7
+    # via jedi
+pexpect==4.9.0
+    # via ipython
+platformdirs==4.11.0
+    # via jupyter-core
+prompt-toolkit==3.0.52
+    # via ipython
+psutil==7.2.2
+    # via
+    #   ipykernel
+    #   ipython
+ptyprocess==0.7.0
+    # via pexpect
+pure-eval==0.2.3
+    # via stack-data
+pycparser==3.0
+    # via cffi
+pydata-sphinx-theme==0.15.4
+    # via
+    #   rocm-docs-core
+    #   sphinx-book-theme
+pygithub==2.9.1
+    # via rocm-docs-core
+pygments==2.20.0
+    # via
+    #   accessible-pygments
+    #   ipython
+    #   ipython-pygments-lexers
+    #   pydata-sphinx-theme
+    #   sphinx
+pyjwt[crypto]==2.13.0
+    # via pygithub
+pynacl==1.6.2
+    # via pygithub
+python-dateutil==2.9.0.post0
+    # via jupyter-client
+pyyaml==6.0.3
+    # via
+    #   jupyter-cache
+    #   myst-nb
+    #   myst-parser
+    #   rocm-docs-core
+    #   sphinx-external-toc
+pyzmq==27.1.0
+    # via
+    #   ipykernel
+    #   jupyter-client
+referencing==0.37.0
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+requests==2.34.2
+    # via
+    #   pygithub
+    #   sphinx
+rocm-docs-core[llms]==1.38.0
+    # via -r docs/sphinx/requirements.in
+roman-numerals==4.1.0
+    # via sphinx
+rpds-py==2026.6.3
+    # via
+    #   jsonschema
+    #   referencing
+six==1.17.0
+    # via python-dateutil
+smmap==5.0.3
+    # via gitdb
+snowballstemmer==3.1.1
+    # via sphinx
+soupsieve==2.9.1
+    # via beautifulsoup4
+sphinx==9.1.0
+    # via
+    #   breathe
+    #   myst-nb
+    #   myst-parser
+    #   pydata-sphinx-theme
+    #   rocm-docs-core
+    #   sphinx-book-theme
+    #   sphinx-copybutton
+    #   sphinx-design
+    #   sphinx-external-toc
+    #   sphinx-markdown-builder
+    #   sphinx-multitoc-numbering
+    #   sphinx-notfound-page
+sphinx-book-theme==1.1.4
+    # via rocm-docs-core
+sphinx-copybutton==0.5.2
+    # via rocm-docs-core
+sphinx-design==0.7.0
+    # via rocm-docs-core
+sphinx-external-toc==1.1.0
+    # via rocm-docs-core
+sphinx-markdown-builder==0.6.10
+    # via rocm-docs-core
+sphinx-multitoc-numbering==0.1.3
+    # via sphinx-external-toc
+sphinx-notfound-page==1.1.0
+    # via rocm-docs-core
+sphinxcontrib-applehelp==2.0.0
+    # via sphinx
+sphinxcontrib-devhelp==2.0.0
+    # via sphinx
+sphinxcontrib-htmlhelp==2.1.0
+    # via sphinx
+sphinxcontrib-jsmath==1.0.1
+    # via sphinx
+sphinxcontrib-qthelp==2.0.0
+    # via sphinx
+sphinxcontrib-serializinghtml==2.0.0
+    # via sphinx
+sqlalchemy==2.0.51
+    # via jupyter-cache
+stack-data==0.6.3
+    # via ipython
+tabulate==0.10.0
+    # via
+    #   jupyter-cache
+    #   sphinx-markdown-builder
+tornado==6.5.7
+    # via
+    #   ipykernel
+    #   jupyter-client
+traitlets==5.15.1
+    # via
+    #   ipykernel
+    #   ipython
+    #   jupyter-client
+    #   jupyter-core
+    #   matplotlib-inline
+    #   nbclient
+    #   nbformat
+typing-extensions==4.16.0
+    # via
+    #   beautifulsoup4
+    #   jupyter-client
+    #   myst-nb
+    #   pydata-sphinx-theme
+    #   pygithub
+    #   referencing
+    #   sqlalchemy
+urllib3==2.7.0
+    # via
+    #   pygithub
+    #   requests
+wcwidth==0.8.2
+    # via prompt-toolkit
+zipp==4.1.0
+    # via importlib-metadata
diff --git a/docs/weekly_reports/2026-W17-primus-weekly.md b/docs/weekly_reports/2026-W17-primus-weekly.md
deleted file mode 100644
index 08a1db41e..000000000
--- a/docs/weekly_reports/2026-W17-primus-weekly.md
+++ /dev/null
@@ -1,147 +0,0 @@
-# Primus Weekly Engineering Report — 2026-W17
-
-## 1. Time Window
-
-- Start: Monday 2026-04-20 00:00:00 Asia/Shanghai (GMT+8)
-- End: Friday 2026-04-24 16:56 Asia/Shanghai (GMT+8) (report generation time)
-- Branch observed: `origin/main`
-
-## 2. Executive Summary
-
-- **9 PRs merged to `main`** in the weekly window (Mon 2026-04-20 00:00 GMT+8 → Fri 2026-04-24 16:56 GMT+8).
-- Category breakdown: **Bug Fix: 3**, **Performance Optimization: 2**, **CI/Infra: 2**, **Docs: 2**; Turbo/Dependency Version Update, Refactor, Other: 0.
-- **No Primus-Turbo version bump this week.** Current `PRIMUS_TURBO_COMMIT` pins in both `ci.yaml` (`333b68d7`) and `benchmark.yaml` (`a4488f6c`) are unchanged since March; month-to-date Turbo drift on `main` is zero.
-- **Megatron-LM upstream drift: `plan sync` recommended.** Pin is `d3528a21` (2026-03-06); upstream `main` HEAD is `a1165fab` with **365 commits ahead**, including MoE routing/dispatch improvements, MTP layers in token-per-expert logging, `MambaModel` → `HybridModel` rename, NVRx async compatibility, DDP parameter-layout refactor, TE `release_v2.14`, and several inference and checkpoint fixes.
-- **torchtitan upstream drift: `urgent sync` recommended.** Pin is `5fb7cc2e` (2025-10-15); upstream `main` HEAD is `d40df991` with **514 commits ahead**, spanning GraphTrainer precompile (SAC + FSDP bucketing, CooR for DeepSeek-V3, CUDA-graph annotation pass), MoE token-dispatcher rewrite, Fused QKV GQAttention, FlexAttention + CP, FSDP2 `fully_shard`, and VLM workflow dependency pinning.
-- **Primus-Turbo month-to-date drift: `monitor` (no action needed).** Both CI and benchmark pins are identical to their 2026-03-30 values.
-- This week was dominated by **Megatron backend hardening** against upstream API drift: #672 realigns the patched MoE overlap `TransformerLayerSchedulePlan` after the upstream `post_attn` removal; #671 makes the muon optimizer wrapper runtime-signature aware against `get_megatron_optimizer()` drift; #674 rewrites `recompute_layer_patches` to be byte-identical to upstream `TransformerBlock._checkpointed_forward` and seeds a SHA256 fingerprint guard.
-- Perf-category PRs: **#673** (`pp_warmup` optimization — parallel fwd+bwd warm-up on every PP rank, with a bit-identical loss-parity UT); **#684** introduces the opt-in `PRIMUS_EXIT_FAST` env to skip Python interpreter teardown after a successful Megatron trainer `cleanup()`, shaving ~22s of wall-time tail per process on MI355X DSV3 EP8 and ~2m off the MI300X Megatron-LM E2E UT suite.
-- Tooling/infra: **#687** lands the shared backend-gap dashboard publishing toolchain (`tools/backend_gap_report/`) plus the initial torchtitan baseline report; this same shared site is being extended this week to also surface the weekly-report series. **#678** switches the JAX CI runner to a TAS node. Docs: **#683** published the initial W17 run, and **#690** refreshed it mid-week.
-
-## 3. Weekly PR Update Table
-
-| PR | Merged Time (GMT+8) | Category | Key Update |
-| --- | --- | --- | --- |
-| [#690](https://github.com/AMD-AGI/Primus/pull/690) `[Primus weekly report] 2026-W17` (author: `cursor[bot]`) | 2026-04-24 14:16 | Docs | Mid-week refresh of the automated W17 weekly report: extends the window through Fri 2026-04-24 09:10 GMT+8, brings the total to 6 PRs at the time of that run, and refreshes upstream drift snapshots for Megatron-LM and torchtitan. Superseded by this report. |
-| [#684](https://github.com/AMD-AGI/Primus/pull/684) `feat(megatron): add PRIMUS_EXIT_FAST to exit training much faster` (author: `lhzhang333`) | 2026-04-24 11:17 | Performance Optimization | Adds an opt-in `PRIMUS_EXIT_FAST=1` env to `MegatronBaseTrainer.cleanup()`: after a successful cleanup and only when `on_error=False`, Primus calls `os._exit(0)` to skip normal Python interpreter/torchrun teardown. Measured ~22s saved on MI355X DSV3 4L EP8 post-train tail and ~2m saved on the MI300X Megatron-LM E2E UT suite; documented as experimental (not recommended for production graceful-shutdown paths). |
-| [#687](https://github.com/AMD-AGI/Primus/pull/687) `feat(backend-gap): add dashboard publishing toolchain and torchtitan baseline report` (author: `WangLingxun`) | 2026-04-24 10:37 | CI/Infra | Ships the shared backend-gap dashboard publishing toolchain under `tools/backend_gap_report/` (static site shell + `build_dashboard_index.py` + `build_site_bundle.py`), seeds project-level skill docs, and publishes the initial TorchTitan-vs-upstream-main report set (`docs/backend-gap/reports/torchtitan/upstream-main/{report,summary}.md` + `dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json`). This is the shared dashboard that this week's weekly-report flow extends to surface `Weekly Reports` as a first-class section. |
-| [#672](https://github.com/AMD-AGI/Primus/pull/672) `fix(megatron): drop stale post_attn usage from patched MoE overlap schedule` (author: `WangLingxun`) | 2026-04-23 15:13 | Bug Fix | Aligns Primus `TransformerModelChunkSchedulePlan` with the current upstream `TransformerLayerSchedulePlan` after Megatron removed the `post_attn` node, eliminating a runtime `AttributeError` under `patch_moe_overlap=True`; preserves `ep_overlap_early_attn_memory_release` ordering and adds a regression test for the patched MoE-overlap schedule-plan path. |
-| [#671](https://github.com/AMD-AGI/Primus/pull/671) `Fix/megatron muon optimizer signature` (author: `WangLingxun`) | 2026-04-23 15:07 | Bug Fix | Hardens the muon optimizer wrapper against `get_megatron_optimizer()` runtime signature drift; uses `inspect.signature` to support both keyword and mixed positional/keyword call patterns; resolves muon-specific args only when a muon optimizer is detected and passes non-muon optimizers through unchanged; adds regression tests for keyword vs positional `config_overrides`, muon vs non-muon, and correct parameter binding under different signatures. |
-| [#678](https://github.com/AMD-AGI/Primus/pull/678) `[CICD]switch jax runner to tas node` (author: `llying-001`) | 2026-04-23 10:32 | CI/Infra | Switches the JAX unit-test runner to a TAS node in `.github/workflows/ci.yaml`. |
-| [#683](https://github.com/AMD-AGI/Primus/pull/683) `[Primus weekly report] 2026-W17` (author: `cursor[bot]`) | 2026-04-23 07:11 | Docs | Adds the initial automated W17 weekly engineering report (`docs/weekly_reports/2026-W17-primus-weekly.md`) covering the Mon 2026-04-20 → Wed 2026-04-22 window. Superseded by #690 and by the current run. |
-| [#673](https://github.com/AMD-AGI/Primus/pull/673) `opt(megatron): optimize pp_warmup and add corresponding UT` (author: `lhzhang333`) | 2026-04-22 20:32 | Performance Optimization | Rewrites `run_pp_warmup` to fabricate per-rank synthetic activations/output grads using `get_tensor_shapes`, so every PP rank runs one warm-up fwd+bwd **in parallel** (bypassing p2p) and exercises all CUDA/TE/FP8/NCCL lazy init paths concurrently; adds rigorous state-isolation (RNG + grad buffers) plus an end-to-end UT validating bit-for-bit loss parity and iter-1 speedup on `PP_SIZE>1`. |
-| [#674](https://github.com/AMD-AGI/Primus/pull/674) `fix(megatron): adapt recompute_layer_patches to the upstream Megatron and add UT` (author: `lhzhang333`) | 2026-04-22 16:30 | Bug Fix | Rewrites inner `custom`/`checkpoint_handler` closures to be byte-identical to Megatron's latest `TransformerBlock._checkpointed_forward`; keeps the delegation fast-path when `recompute_layer_ids` is unset; seeds the upstream-source SHA256 fingerprint guard; adds UTs for pipeline-stage offset mapping and the FP8-no-grad skip rule; removes the stale `PrimusTransformerBlock` subclass. |
-
-## 4. Megatron-LM Drift Overview
-
-- Upstream: `https://github.com/NVIDIA/Megatron-LM.git` (`main`)
-- Pinned in Primus `main` (`third_party/Megatron-LM`): `d3528a21301db2d12e92912b3ec025dc8a2ed4d6` — *fix(moe): fix TE general_gemm API change (#3582)*, 2026-03-06
-- Upstream `main` HEAD: `a1165fabcad97eae3778f386839c233dfabf3f8b` — *Inference: Fix broken functional tests on gitlab (#4454)* (2026-04-24)
-- Commit gap: **upstream is 365 commits ahead of Primus pin**
-- Month-to-date movement on Primus side: pin advanced from `3bec9aa9` (2026-02-26) → `d3528a21` (2026-03-06) inside PR #654 merged 2026-04-10 (282-commit upstream catch-up). No further submodule SHA change in April.
-- Recommendation: **plan sync** — several releases' worth of MoE, precision, Mamba→Hybrid rename, and FSDP/DistOpt changes have accumulated; schedule a controlled bump rather than urgent. This week's Primus-side fixes (#671, #672, #674) already pre-harden Primus against concrete upstream API drift points that must be validated during the sync.
-
-### Notable upstream areas that have moved since the pin
-
-- **MoE routing / dispatch**: new router score function (#3673), shared-expert overlap improvements including FlexDispatcher support (#2207), fix for unnecessary permute padding in non-quantized dispatch (#4038); MTP layers now counted in token-per-expert logging (#4412).
-- **Mamba / Hybrid models**: `MambaModel`/`MambaStack` → `HybridModel`/`HybridStack` rename including the outside-of-`megatron/core` rename (#4159, follow-up to #4099); Mamba inference optimization (#4414); QK layernorm for DPA in `MambaModel` (#4067); port DeepSeek Sparse Attention to `MambaModel` (#3553); fine-grained activation offloading (#4173).
-- **Low precision / TE**: NVFP4 native DDP weights (#4005); Enable FP8 DPA for MXFP8 recipe (#4066); TransformerEngine bumped to `release_v2.14` (#4331).
-- **Checkpointing**: add `--async-ckpt-use-cpu-shm` (#4355); remove cross-rank sync during checkpoint load & deprecate `state_dict_loader.load_state_dict` (#2864); fix potential coredump on save (#1871); RL onload optimizer after logprobs (#4235).
-- **FSDP / Distributed Optimizer / DDP**: Megatron-FSDP 0.5.0; `expt_device_mesh` fix for MoE-only (#3831); MFSDP `decoupled_grad`/DistOpt mechanics (#4133); layerwise-optimizer fixes (#4272, #4138); DDP refactor extracting parameter-layout computation into an optimizer classmethod (#3812); MFSDP log mcore detection only after imports succeed (#4400).
-- **Inference / RL / misc**: NVRx async compatibility + defer resiliency import (#4420); misc inference fixes (#4397, #4454); RL token-throughput & packing metrics (#3877); nvtx_decorator checks `_nvtx_enabled` at call time (#4184); NullTokenizer for pretraining to reduce I/O (#4057); rampup batch-size scheduler replaced with custom step schedule (#4411, reverted earlier work #3779 via #4404).
-
-### Megatron-LM upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| MoE | Router new score function;<br>shared-expert overlap for FlexDispatcher;<br>MoE DPA / permute-padding fixes;<br>MTP layers in token-per-expert logging | NVIDIA/Megatron-LM #3673, #2207, #4038, #4412 | Could unlock additional MoE recipes (DeepSeek / Mixtral variants) already referenced in Primus `examples/megatron/configs/MI300X/deepseek_v*`. |
-| Mamba / Hybrid | `MambaModel` → `HybridModel` rename (outside-core rename landed);<br>Mamba inference opt;<br>DeepSeek Sparse Attention port;<br>fine-grained activation offloading | NVIDIA/Megatron-LM #4099, #4159, #4414, #3553, #4173 | **Breaking import rename** — Primus patch system (`primus/backends/megatron`) must audit any `MambaModel`/`MambaStack` references before the next bump. |
-| Low precision / TE | NVFP4 native DDP weights;<br>FP8 DPA for MXFP8;<br>TE bumped to `release_v2.14` | NVIDIA/Megatron-LM #4005, #4066, #4331 | Requires matching TE/AITER versions in Primus Dockerfile; likely coordinated with next Turbo+aiter bump. |
-| Checkpoint I/O | Async checkpoint CPU-SHM;<br>removed cross-rank sync load;<br>save coredump fix;<br>RL optimizer onload ordering | NVIDIA/Megatron-LM #4355, #2864, #1871, #4235 | Expected improvement for Primus pretrain at scale; validate with Primus async-ckpt configs. |
-| FSDP / DistOpt / DDP | Megatron-FSDP 0.5.0;<br>MFSDP `decoupled_grad`/DistOpt fixes;<br>layerwise-optimizer fixes;<br>DDP parameter-layout refactor | NVIDIA/Megatron-LM #3831, #4133, #4272, #4138, #3812, #4400 | Primus FSDP/DDP launch paths (`primus/modules/trainer/megatron/*`) should be re-benchmarked post-sync; this directly touches paths exercised by this week's #673 `pp_warmup` optimization. |
-| Inference / resiliency | NVRx async compatibility;<br>misc inference fixes;<br>nvtx decorator call-time gating | NVIDIA/Megatron-LM #4420, #4397, #4454, #4184 | Aligns with Primus inference-adjacent tooling; low-risk but worth validating Primus resiliency paths during the controlled bump. |
-| Schedule plan API | `post_attn` node already removed from `TransformerLayerSchedulePlan` (consumed in Primus by #672) | Primus/#672, upstream schedule-plan change | Primus-side fix has already shipped; treat as confirmed drift that future bumps must keep in sync. |
-
-## 5. torchtitan Drift Overview
-
-- Upstream: `https://github.com/pytorch/torchtitan.git` (`main`)
-- Pinned in Primus `main` (`third_party/torchtitan`): `5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021` — *Deepseek-V3 toml file minor fix (#1894)*, 2025-10-15
-- Upstream `main` HEAD: `d40df991ac535108e428b0746a08b74a3cf6afc7` — *[GraphTrainer] Skip bucketing pass in precompile and re-enable tests (#3079)* (2026-04-24)
-- Commit gap: **upstream is 514 commits ahead of Primus pin**
-- Month-to-date movement on Primus side: none (submodule SHA unchanged in April).
-- Recommendation: **urgent sync** — the pin is six months stale; upstream has undergone major refactors (GraphTrainer precompile, MoE token dispatcher, FlexAttention CP, FSDP2) that block adopting any new Primus torchtitan-backend features. This week, the shared backend-gap dashboard (#687) publishes the first tracked torchtitan-vs-upstream baseline report, which can be used to drive the sync plan.
-
-### Notable upstream areas that have moved since the pin
-
-- **GraphTrainer / precompile**: CooR precompile for DeepSeek V3 (#2916) and `aot_fx_trace` compile mode (#2975); precompile integration tests in CI (#3043); regional-inductor precompile (#2883); `enable_cudagraph` config flag (#3049); FSDP bucketing pass toggled (#3044), re-enabled with SAC improvements (#3060), and most recently skipped again in precompile with tests re-enabled (#3079); CUDA-graph kernel annotation pass (#2926); CPU offload pass for activation memory savings (#3064, reland); SAC pass refactor using `module_fqn` for layer boundaries (#3050).
-- **MoE**: token dispatcher introduced replacing token reorderer (#2842); EP setup moved from trainer to config registry (#2960); revert `torch.bmm` → scatter-add (#2775); remove unnecessary MoE padding (#2774).
-- **Attention / FlexAttention**: Fused QKV GQAttention (#2878); combine `q_norm`+`k_norm` into `qk_norm` (#2872); FlexAttention bitwise-deterministic tests (#2903, #2989); 2-tier compilation with FlexAttention (#2929); refactor inner attention module (#2761); CP + block_causal + FlexAttention position fix (#2780).
-- **FSDP2 & compile**: replace `amp` + `replicate` with `fully_shard` (#2900); lazy import of FSDP mesh helpers for older PyTorch (#2888); SimpleFSDP wrapper shared across same-type modules (#2754); migrate to `.compile()` API (#2688); full DTensor for Qwen3 and Llama4 at TP region (#2149).
-- **RL / trainer refactor**: RL trainer and generator refactors (#2985, #3001); rename inference example + consolidate vllm logical elements (#3045); drop `get_model_state_dict` in `push_model_state_dict` (#3066).
-- **Datasets / CI / ROCm**: shuffle `HuggingFaceTextDataset` on re-loop and replay on resume (#3023); VLM 8-GPU workflow pins `torchvision` alongside `torch` (#3047); tj-actions version bumps (#3048); MI350 label used for all ROCm workflows (#2740); JIT/AOT tests gated off upstream partitioner regression (#3061).
-
-### torchtitan upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| GraphTrainer / precompile | CooR precompile for DeepSeek V3;<br>precompile for `aot_fx_trace` + SAC/FSDP bucketing improvements;<br>regional-inductor precompile;<br>`enable_cudagraph` flag;<br>CUDA-graph kernel annotation pass;<br>CPU-offload activation pass | pytorch/torchtitan #2916, #2975, #3060, #3079, #2883, #3049, #2926, #3064 | Major perf/UX upgrade for torchtitan-backed training in Primus; currently unavailable behind stale pin. |
-| MoE | New token dispatcher replacing token reorderer;<br>EP setup moved to config registry | pytorch/torchtitan #2842, #2960 | API surface change for torchtitan MoE configs in `primus/backends/torchtitan/**`; patch notes will need an update after sync. |
-| Attention | Fused QKV GQAttention;<br>`qk_norm` consolidation;<br>FlexAttention bitwise-deterministic tests;<br>2-tier compilation with FlexAttention;<br>CP + block_causal + FlexAttention position fix | pytorch/torchtitan #2878, #2872, #2903, #2929, #2780 | Potential perf uplift for Primus torchtitan attention path; determinism tests useful for CI. |
-| FSDP2 / compile / TP | `fully_shard` replaces amp+replicate;<br>`.compile()` migration;<br>shared SimpleFSDP wrapper;<br>full DTensor for Qwen3 / Llama4 TP | pytorch/torchtitan #2900, #2688, #2754, #2149 | Breaking public-API adjustments; Primus torchtitan launcher and patches must be re-validated. |
-| RL trainer | RL trainer + generator refactors;<br>drop `get_model_state_dict` in `push_model_state_dict`;<br>inference-example rename + vllm consolidation | pytorch/torchtitan #2985, #3001, #3066, #3045 | Relevant to any Primus post-training / RL integrations on the torchtitan backend. |
-| Datasets / CI / ROCm | Shuffle HF text dataset on re-loop;<br>`torchvision` pin in VLM workflow;<br>MI350 label across workflows;<br>tj-actions version bumps | pytorch/torchtitan #3023, #3047, #2740, #3048 | Good hygiene reference for Primus torchtitan CI and MI-series Docker dependency pinning. |
-
-## 6. Primus-Turbo Monthly Drift Overview
-
-- Drift type: **in-repo**, not upstream — compares Turbo version/SHA referenced on Primus `main` now vs the latest commit at or before `month_start_ts = 2026-04-01 00:00 Asia/Shanghai` (`2026-03-31 16:00 UTC`).
-- Turbo is **not a submodule** in Primus. Canonical version source:
-  - `.github/workflows/ci.yaml` → `PRIMUS_TURBO_COMMIT` (also wired through `.github/workflows/docker/Dockerfile`)
-  - `.github/workflows/benchmark.yaml` → `PRIMUS_TURBO_COMMIT`
-- Reference commit at month start on `main`: `76651575` (*[WIP][Megatron-LM] feat: reduce extra qkv transpose in attn (#625)*, 2026-03-31 14:19 GMT+8). The underlying Turbo pins at that commit are the same as the current values.
-- Current state on `origin/main`:
-  - `ci.yaml` `PRIMUS_TURBO_COMMIT`: `333b68d7c81b722b21b4aad10cd250c45f15027c` — *fix sm_scale none bug (#263)*
-  - `ci.yaml` `PRIMUS_TURBO_AITER_COMMIT`: `e83f9903c07001a0ec29e85d223f6e6cdbe00859`
-  - `benchmark.yaml` `PRIMUS_TURBO_COMMIT`: `a4488f6cdb15cfff4383c61af7922bb50803f0ea` — *feat: update triton impl for mi300 & mi355 (#252)*
-- Month-start state on `main`: all three values identical to current.
-- **No Primus-Turbo drift in this comparison window.**
-- Recommendation: **monitor**. Note the *pre-existing* skew between the two YAML pins (CI pin `333b68d` is 5 commits ahead of benchmark pin `a4488f6c` in Primus-Turbo history) — tracked separately and unchanged this month.
-
-### Notable areas changed since month start
-
-- **No changes this window** — both `ci.yaml` and `benchmark.yaml` Turbo pins on `main` are byte-identical to their 2026-03-30 values.
-- **No Turbo drift in this comparison window.**
-
-### Primus-Turbo monthly drift table
-
-| Component | Current Version/SHA | Month-start Version/SHA | Delta Summary | Key Changes | Evidence |
-| --- | --- | --- | --- | --- | --- |
-| `PRIMUS_TURBO_COMMIT` (CI build) | `333b68d7c81b722b21b4aad10cd250c45f15027c` (*fix sm_scale none bug (#263)*) | `333b68d7c81b722b21b4aad10cd250c45f15027c` | No drift (0 commits) | No changes this window. | [`.github/workflows/ci.yaml` L17](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L17) |
-| `PRIMUS_TURBO_AITER_COMMIT` (CI build) | `e83f9903c07001a0ec29e85d223f6e6cdbe00859` | `e83f9903c07001a0ec29e85d223f6e6cdbe00859` | No drift | No changes this window. | [`.github/workflows/ci.yaml` L18](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L18) |
-| `PRIMUS_TURBO_COMMIT` (benchmark) | `a4488f6cdb15cfff4383c61af7922bb50803f0ea` (*feat: update triton impl for mi300 & mi355 (#252)*) | `a4488f6cdb15cfff4383c61af7922bb50803f0ea` | No drift | No changes this window. | [`.github/workflows/benchmark.yaml` L9](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/benchmark.yaml#L9) |
-
-## 7. Source Links
-
-- Primus main branch: https://github.com/AMD-AGI/Primus/tree/main
-- Primus weekly PR listing (window): https://github.com/AMD-AGI/Primus/pulls?q=is%3Apr+is%3Amerged+base%3Amain+merged%3A%3E%3D2026-04-19T16%3A00%3A00Z
-- PR #671: https://github.com/AMD-AGI/Primus/pull/671
-- PR #672: https://github.com/AMD-AGI/Primus/pull/672
-- PR #673: https://github.com/AMD-AGI/Primus/pull/673
-- PR #674: https://github.com/AMD-AGI/Primus/pull/674
-- PR #678: https://github.com/AMD-AGI/Primus/pull/678
-- PR #683: https://github.com/AMD-AGI/Primus/pull/683
-- PR #684: https://github.com/AMD-AGI/Primus/pull/684
-- PR #687: https://github.com/AMD-AGI/Primus/pull/687
-- PR #690: https://github.com/AMD-AGI/Primus/pull/690
-- Megatron-LM pin: https://github.com/NVIDIA/Megatron-LM/commit/d3528a21301db2d12e92912b3ec025dc8a2ed4d6
-- Megatron-LM upstream HEAD (at report time): https://github.com/NVIDIA/Megatron-LM/commit/a1165fabcad97eae3778f386839c233dfabf3f8b
-- Megatron-LM compare: https://github.com/NVIDIA/Megatron-LM/compare/d3528a21301db2d12e92912b3ec025dc8a2ed4d6...main
-- torchtitan pin: https://github.com/pytorch/torchtitan/commit/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021
-- torchtitan upstream HEAD (at report time): https://github.com/pytorch/torchtitan/commit/d40df991ac535108e428b0746a08b74a3cf6afc7
-- torchtitan compare: https://github.com/pytorch/torchtitan/compare/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021...main
-- Primus-Turbo CI pin: https://github.com/AMD-AGI/Primus-Turbo/commit/333b68d7c81b722b21b4aad10cd250c45f15027c
-- Primus-Turbo benchmark pin: https://github.com/AMD-AGI/Primus-Turbo/commit/a4488f6cdb15cfff4383c61af7922bb50803f0ea
-- Month-start reference commit on `main`: https://github.com/AMD-AGI/Primus/commit/76651575
-
----
-
-*Generated automatically by the Primus weekly report automation. Factual statements are derived from `git log origin/main`, the pinned submodule SHAs in `third_party/`, and the `PRIMUS_TURBO_COMMIT` values in `.github/workflows/{ci,benchmark}.yaml` as observed at 2026-04-24 16:56 GMT+8. Upstream-HEAD SHAs and commit counts are snapshots at report generation time.*
diff --git a/docs/weekly_reports/2026-W18-primus-weekly.md b/docs/weekly_reports/2026-W18-primus-weekly.md
deleted file mode 100644
index 5277e0874..000000000
--- a/docs/weekly_reports/2026-W18-primus-weekly.md
+++ /dev/null
@@ -1,138 +0,0 @@
-# Primus Weekly Engineering Report — 2026-W18
-
-## 1. Time Window
-
-- Start: Monday 2026-04-27 00:00:00 Asia/Shanghai (GMT+8)
-- End: Friday 2026-05-01 09:10 Asia/Shanghai (GMT+8) (report generation time)
-- Branch observed: `origin/main`
-
-## 2. Executive Summary
-
-- **2 PRs merged to `main`** in the weekly window (Mon 2026-04-27 00:00 GMT+8 → Fri 2026-05-01 09:10 GMT+8).
-- Category breakdown: **Bug Fix: 1**, **Other (feature): 1**; Performance Optimization, Turbo/Dependency Version Update, CI/Infra, Refactor, Docs: 0.
-- **No backend pin or Turbo pin changes this week.** `third_party/Megatron-LM` is still pinned at `d3528a21` and `third_party/torchtitan` at `5fb7cc2e`. `PRIMUS_TURBO_COMMIT` is still `333b68d7` in `ci.yaml` and `a4488f6c` in `benchmark.yaml`. The week's only `.github/workflows/ci.yaml` change (in #693) commented out the ROCm Docker Hub login lines and does not move any version pin.
-- **Megatron-LM upstream drift: `plan sync` (unchanged).** Pin is `d3528a21` (2026-03-06); upstream `main` HEAD is `3460bba1` (2026-05-01) — **413 commits ahead** (+48 since W17 snapshot `a1165fab`). Notable new upstream activity in this window: AllgatherV inference dispatcher (#4258), permute fusion in hybrid EP (#4089), DeepSeek/MoE prep (per-block MoE routing storage for prefix caching #4301, MTP CUDA graphs #4260), Megatron-FSDP doc unification (#4418), inference graph standardization (#4485) and inference RL graph fix (#4323), heterogeneous TP/DP MIMO `ColocatedBridgeCommunicator` (#4368), checkpoint integrity verification (#4305), `--global-batch-size` removal in step-batch-size release tests (#4545), and a build move of `mamba-ssm`/`causal-conv1d` to optional `[ssm]` extras (#4517).
-- **torchtitan upstream drift: `urgent sync` (unchanged).** Pin is `5fb7cc2e` (2025-10-15); upstream `main` HEAD is `70340f4e` (2026-04-30) — **566 commits ahead** (+52 since W17 snapshot `d40df991`). Major new upstream activity in this window: GraphTrainer unified activation-memory framework (#3118), full inductor compilation pass (#3141), async-TP / micro-pipeline TP graph pass (#3129), HybridEP integration with GraphTrainer (#3007, #3177), MoE-series consolidation (`All2AllTokenDispatcher` for EP=1 and EP>1 #3125, ETP deprecation #3167), `ChunkedCELoss` (#2937), Full DTensor config-based sharding for Llama3/Qwen3/Llama4/DSV3/GPT-OSS (#2963, #2969), `MeshDimName → MeshAxisName` rename (#3113), and a CI-side `replace-imports-with-any` pattern to keep CI green when optional packages are missing (#3180).
-- **Primus-Turbo month-to-date drift (April → May 1 GMT+8): `monitor` (no action needed).** Both CI and benchmark `PRIMUS_TURBO_COMMIT` pins remain byte-identical to their values at month start (2026-04-01 00:00 GMT+8) and to their previously reported W17 values.
-- This week was the **second pass on the upstream Primus-Turbo `TEGroupedMLP` integration on the Megatron backend** (#693): the `te_spec_provider` patch no longer disables PrimusTurbo when `use_turbo_grouped_mlp=True` + `moe_grouped_gemm=True` + TE>=1.9.0; `validate_args_on_rocm` now hard-rejects `use_turbo_grouped_mlp=True` combined with `moe_use_legacy_grouped_gemm=True`; the legacy GroupedMLP/`grouped_gemm_util` code is preserved under `primus/backends/megatron/core/transformer/moe/deprecated_2caa681a1/` for backward compatibility with older configs that still set `moe_use_legacy_grouped_gemm=True`; and `PrimusTurboDeepEPTokenDispatcher` no longer gates `deepep_use_cuda_num_tokens_per_expert` on the legacy GroupedGEMM flag. Several MI300X/MI355X model YAMLs (`deepseek_v2_lite`, `deepseek_v3`, `gpt_oss_20B`, `qwen3_30B_A3B`, `qwen3_235B_A22B`) had their stale `moe_use_legacy_grouped_gemm: true` lines dropped to align with the new validation.
-- **New model family: LFM/LFM2** (#651) — Megatron backend gains LiquidAI LFM2 support: `primus/configs/models/megatron/lfm_base.yaml` + `lfm2_8B_A1B.yaml`, a Primus implementation of the LFM2 short-convolution layer (`Lfm2ShortConv`) wired through a new `primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py`, and three example MI355X pretrain YAMLs (`lfm2_8B_A1B-{BF16,FP8,FP8-te-precision}-pretrain.yaml`). Also includes a small docker-build CI fix (commenting out `rocmshared` Docker Hub logins in the build/push workflow).
-- Tooling/infra/docs: no merged tooling, infra-only, refactor, or docs PRs in this window. The shared backend-gap dashboard from W17's #687 is left unmodified.
-
-## 3. Weekly PR Update Table
-
-| PR | Merged Time (GMT+8) | Category | Key Update |
-| --- | --- | --- | --- |
-| [#693](https://github.com/AMD-AGI/Primus/pull/693) `fix: keep legacy groupedgemm on megatron backend` (author: `zhenhuang12`) | 2026-04-28 15:27 | Bug Fix | Second-pass alignment of PrimusTurbo with the upstream `TEGroupedMLP` path on the Megatron backend. Removes the patch-level guard that disabled PrimusTurbo when `use_turbo_grouped_mlp + moe_grouped_gemm + TE>=1.9.0`, and adds a hard `validate_args_on_rocm` check that forbids `use_turbo_grouped_mlp=True` combined with `moe_use_legacy_grouped_gemm=True`. `PrimusTurboDeepEPTokenDispatcher` no longer requires `moe_use_legacy_grouped_gemm=True` to enable `deepep_use_cuda_num_tokens_per_expert`. The legacy GroupedMLP path is preserved under `primus/backends/megatron/core/transformer/moe/deprecated_2caa681a1/` for backward-compat with older configs that still set the legacy flag. Sync-Free MoE stage 2/3 now requires `use_turbo_grouped_mlp=True` (instead of the old legacy flag). Strips the now-stale `moe_use_legacy_grouped_gemm: true` lines from the MI300X/MI355X DeepSeek-V2/V3, GPT-OSS-20B, and Qwen3 30B/235B example YAMLs. |
-| [#651](https://github.com/AMD-AGI/Primus/pull/651) `LFM model support` (author: `wenxie-amd`) | 2026-04-28 11:44 | Other | Adds initial Megatron-backend support for LiquidAI's LFM2 model family. Introduces `primus/configs/models/megatron/lfm_base.yaml` + `lfm2_8B_A1B.yaml`, a Primus implementation of the LFM2 short-convolution "attention" layer (`Lfm2ShortConv`) plus an LFM-aware GPT layer-spec builder, and a new Megatron patch (`primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py`) that routes `get_gpt_decoder_layer_specs` to the Primus implementation when LFM-specific layer types are configured. Ships three MI355X example pretrain configs: `lfm2_8B_A1B-BF16-pretrain.yaml`, `lfm2_8B_A1B-FP8-pretrain.yaml`, `lfm2_8B_A1B-FP8-te-precision.yaml`. Also includes a docker-build CI fix in `.github/workflows/ci.yaml` (comments out four `docker login -u rocmshared ... ROCM_DOCKER_HUB_TOKEN` lines in the image build/push job). |
-
-## 4. Megatron-LM Drift Overview
-
-- Upstream: `https://github.com/NVIDIA/Megatron-LM.git` (`main`)
-- Pinned in Primus `main` (`third_party/Megatron-LM`): `d3528a21301db2d12e92912b3ec025dc8a2ed4d6` — *fix(moe): fix TE general_gemm API change (#3582)*, 2026-03-06
-- Upstream `main` HEAD: `3460bba1d6f945ec04f47fdb1dcee6da3259fcd8` — *Update copy-pr-bot.yaml [skip ci]* (2026-05-01)
-- Last upstream functional change in this window: `83e7466c` — *Fixes for modelopt examples and SFTTokenizer for transformers v5 (#4450)* (2026-04-30)
-- Commit gap: **upstream is 413 commits ahead of Primus pin** (+48 since the W17 snapshot `a1165fab`).
-- Month-to-date movement on Primus side: pin unchanged in April; last submodule SHA bump on `main` was `3bec9aa9` → `d3528a21` inside PR #654 (merged 2026-04-10).
-- Recommendation: **plan sync** (unchanged from W17). The accumulated upstream change set continues to grow (now MoE permute fusion in hybrid EP, AllgatherV inference dispatcher, MTP CUDA graphs, RL inference graph fixes, MFSDP doc unification) without changing Primus's existing sync risk profile.
-
-### Notable upstream areas that have moved since the pin
-
-- **MoE / EP**: permute fusion in hybrid EP (#4089) on top of W17's router, FlexDispatcher, and MTP-token-per-expert work; per-block MoE routing storage for prefix caching (#4301); previously reported router score function (#3673), shared-expert overlap incl. FlexDispatcher (#2207), permute-padding fix (#4038), MTP token-per-expert logging (#4412).
-- **Inference / CUDA graphs**: AllgatherV inference dispatcher and old-dispatcher simplification (#4258); CUDA graphs for MTP inference (#4260); avoid nsys profile crash with CUDA graphs (#4541); standardize misc graph interface (#4485); fix inference graph override in RL flow (#4323); local-CG bugfixes for latent MoE loss-curve gaps (#4433); embedding/output layer in `full_iteration_inference` graph for hybrid models (#4440).
-- **Mamba / Hybrid models**: avoid redundant HBM reloads in `causal_conv1d_update` shift loop (#4460); build move of `mamba-ssm`/`causal-conv1d` to optional `[ssm]` extra (#4517); on top of W17's `MambaModel`/`MambaStack` → `HybridModel`/`HybridStack` rename (#4099, #4159), Mamba inference opt (#4414), QK layernorm in `MambaModel` DPA (#4067), DeepSeek Sparse Attention port (#3553), fine-grained activation offloading (#4173), YARN support for hybrid_model (#4244).
-- **Megatron-FSDP / DistOpt / DDP**: documentation unified and refactored (#4418); fix `FusedAdam.use_decoupled_grad` mis-set for Megatron-FSDP (#4427); add `reduce_scatter_with_fp32_accumulation` knob (#4410); on top of W17's MFSDP 0.5.0, MFSDP `decoupled_grad`/DistOpt fixes (#4133), layerwise-optimizer fixes (#4272, #4138), DDP parameter-layout refactor (#3812), MFSDP log gating (#4400).
-- **Checkpoint / safety**: checkpoint integrity verification (#4305); SafeUnpickler class for safe pickle usage (#4319); SHA-256 prefix-cache hash replacing polynomial rolling hash (#4158); `weights_only=False` removal (#4434); `validate_access_integrity` knob on dist-ckpt load (#4422); fix checkpoint loading with rerun state machine (#4448); on top of W17's async-ckpt CPU-SHM (#4355) and cross-rank-sync removal (#2864).
-- **Heterogeneous training / RL / misc**: `ColocatedBridgeCommunicator` for heterogeneous TP/DP MIMO training (#4368); ModelOpt examples + SFTTokenizer fixes for transformers v5 (#4450); ModelOpt list-format `quant_cfg` fix (#4187); YAML quant recipe in PTQ + first/last layer modifier removal (#4503); `--global-batch-size` removal from step-batch-size schedule release tests (#4545); training-migration container/serialization classes (#4227, #4309); upstream skill-doc updates (#4502, #4542).
-
-### Megatron-LM upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| MoE / EP | Permute fusion in hybrid EP;<br>per-block MoE routing storage for prefix caching;<br>(carries) router score function, FlexDispatcher overlap, MTP token-per-expert logging | NVIDIA/Megatron-LM #4089, #4301, #3673, #2207, #4412 | Primus DeepSeek/Mixtral configs (`examples/megatron/configs/MI300X/deepseek_v*`, `qwen3_*`) may pick up additional EP perf knobs; aligns with this week's Primus-side groupedgemm cleanup in #693. |
-| Inference / CUDA graphs | New AllgatherV inference dispatcher;<br>CUDA graphs for MTP inference;<br>standardized misc graph interface;<br>RL inference-graph override fix;<br>local-CG bugfixes for latent MoE loss curves | NVIDIA/Megatron-LM #4258, #4260, #4485, #4323, #4433, #4541 | Inference/post-train paths for Primus DSV3/MoE configs benefit; expand validation when sync lands. |
-| Mamba / Hybrid | `causal_conv1d_update` HBM-reload reduction;<br>`mamba-ssm`/`causal-conv1d` move to optional `[ssm]` extra;<br>(carries) outside-core `MambaModel` → `HybridModel` rename | NVIDIA/Megatron-LM #4460, #4517, #4159, #4244 | The `[ssm]` extra split affects Primus Dockerfile install layering; Primus must audit `MambaModel`/`MambaStack` references in `primus/backends/megatron` before the next bump. |
-| FSDP / DistOpt / DDP | MFSDP doc unification;<br>FusedAdam `use_decoupled_grad` Megatron-FSDP fix;<br>`reduce_scatter_with_fp32_accumulation` knob;<br>(carries) MFSDP 0.5.0 + layout refactor | NVIDIA/Megatron-LM #4418, #4427, #4410, #3812, #4400 | Primus FSDP/DDP launch paths (`primus/modules/trainer/megatron/*`) should re-validate post-sync; this week's Primus-side validation (#693) makes the Turbo grouped-MLP path the new default for MoE, which interacts with these MFSDP changes. |
-| Checkpoint / safety | Checkpoint integrity verification;<br>SafeUnpickler;<br>SHA-256 prefix-cache hash;<br>`weights_only=False` removal;<br>`validate_access_integrity` knob | NVIDIA/Megatron-LM #4305, #4319, #4158, #4434, #4422, #4448 | Tighter checkpoint hardening; Primus pretrain at scale should benefit; coordinate when bumping the pin. |
-| Heterogeneous / RL / misc | `ColocatedBridgeCommunicator` (NMFW-17);<br>ModelOpt fixes for transformers v5;<br>YAML quant recipe in PTQ;<br>step-batch-size release-test fix | NVIDIA/Megatron-LM #4368, #4450, #4187, #4503, #4545 | New heterogeneous-training entry point may inform future Primus multi-pod configs; ModelOpt fixes are required when Primus moves to transformers v5. |
-| Schedule plan API (carry-over) | `post_attn` node already removed from `TransformerLayerSchedulePlan` (consumed in Primus by W17 #672) | Primus/#672, upstream schedule-plan change | Primus-side fix already shipped; treat as confirmed drift that future bumps must keep in sync. |
-
-## 5. torchtitan Drift Overview
-
-- Upstream: `https://github.com/pytorch/torchtitan.git` (`main`)
-- Pinned in Primus `main` (`third_party/torchtitan`): `5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021` — *Deepseek-V3 toml file minor fix (#1894)*, 2025-10-15
-- Upstream `main` HEAD: `70340f4ec31d9e1dbd448506cc4423934c3cd62f` — *[CI] Use replace-imports-with-any to avoid missing packages causing CI failures (#3180)* (2026-04-30)
-- Commit gap: **upstream is 566 commits ahead of Primus pin** (+52 since the W17 snapshot `d40df991`).
-- Month-to-date movement on Primus side: none (submodule SHA unchanged in April).
-- Recommendation: **urgent sync** (unchanged from W17). The pin is now ~6.5 months stale; another major refactor wave landed upstream this week (GraphTrainer activation-memory framework, full inductor compilation pass, async-TP graph pass, MoE token-dispatcher consolidation, ETP deprecation, ChunkedCELoss, Full DTensor config-based sharding) that further widens the gap from the existing baseline report.
-
-### Notable upstream areas that have moved since the pin
-
-- **GraphTrainer / precompile**: unified framework for activation memory management (#3118); full inductor compilation pass (#3141); joint graph bucketing + prefetching that composes with SAC (#3056); async-TP / micro-pipeline TP graph pass (#3129); start deprecating jit/aot compile modes in graph_trainer (#2788); `apply_graph_ac` removed (#3147); `log_activation_memory_policy` for per-tensor inspection (#3062); on top of W17's CooR precompile for DSV3 (#2916), `aot_fx_trace` (#2975), regional-inductor precompile (#2883), `enable_cudagraph` (#3049), CUDA-graph annotation pass (#2926), CPU-offload activation pass (#3064), and SAC pass refactor (#3050).
-- **HybridEP / MoE**: HybridEP enabled with GraphTrainer (#3007); HybridEP reads `comm_backend` from model config (#3177); `[MoE][3/n]` consolidate EP=1 and EP>1 to all use `All2AllTokenDispatcher` (#3125); `[MoE][4/n]` deprecate Expert Tensor Parallel (ETP) (#3167); on top of W17's token-dispatcher introduction (#2842) and EP-config-registry move (#2960).
-- **Module / DTensor / TP**: Full DTensor config-based sharding infrastructure with Llama3 adoption (#2963) and follow-up for Qwen3, Llama4, DSV3, GPT-OSS (#2969); `MeshDimName → MeshAxisName` rename (#3113); remove `LocalMapInnerAttention`, use static `LocalMapSpec` (#2986); on top of W17's `fully_shard` migration (#2900), `.compile()` migration (#2688), and SimpleFSDP wrapper sharing (#2754).
-- **Attention / loss**: `ChunkedCELoss` (#2937) plus disable chunked CE in graph trainer (#3115); revert "Remove MATH from ScaledDotProductAttention default backends" (#3135) after the original change (#3080); fused QKV support in Qwen3-VL state-dict adapter (#3102); on top of W17's Fused QKV GQAttention (#2878), `qk_norm` consolidation (#2872), FlexAttention bitwise-deterministic tests (#2903, #2989), 2-tier compilation with FlexAttention (#2929), CP+block_causal+FlexAttention position fix (#2780).
-- **RL / inference / experiments**: RL env-rollout-based controller refactor (#3073); patched-Qwen3 parallel plan removed and merged into core torchtitan Qwen3 parallel plan (#3070); reset_prefix_cache (#3095); deprecate VLM experiment (#3151); remove stale autoparallel/deepseek_v3 experiment (#2271); on top of W17's RL trainer/generator refactors (#2985, #3001) and inference-example/vllm consolidation (#3045).
-- **Datasets / CI / ROCm**: CI uses `replace-imports-with-any` to avoid missing-package failures (#3180); shuffle ChatDataset before splitting across nodes (#3131); reproducible training resume across epoch boundaries for map and streaming datasets (#3008); MoE loss comparison guard added to CI (#3081); ROCm experimental workflows toggled off again (#3140) after a brief revert (#3097); MI350 loss numbers updated (#3078); SAC test compatibility with PyTorch indexed storage (#3098); on top of W17's HF text-dataset reshuffle (#3023), VLM `torchvision` pin (#3047), MI350 label rollout (#2740), and tj-actions bumps (#3048).
-
-### torchtitan upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| GraphTrainer / precompile | Unified activation-memory framework;<br>full inductor compilation pass;<br>joint graph bucketing + prefetching that composes with SAC;<br>async-TP / micro-pipeline TP graph pass;<br>jit/aot compile-mode deprecation;<br>activation-memory policy logging | pytorch/torchtitan #3118, #3141, #3056, #3129, #2788, #3062 | Major continuing perf/UX upgrade; remains unavailable behind the stale pin. Will require coordinated patches in `primus/backends/torchtitan/**` when Primus syncs. |
-| HybridEP / MoE | HybridEP integration with GraphTrainer;<br>HybridEP `comm_backend` from model config;<br>`All2AllTokenDispatcher` consolidation across EP=1 and EP>1;<br>ETP deprecation | pytorch/torchtitan #3007, #3177, #3125, #3167 | Direct impact on Primus torchtitan MoE configs and any planned EP topology in `primus/modules/trainer/torchtitan/*`; ETP removal is a breaking config-level change. |
-| Module / DTensor / TP | Full DTensor config-based sharding for Llama3/Qwen3/Llama4/DSV3/GPT-OSS;<br>`MeshDimName → MeshAxisName` rename;<br>`LocalMapInnerAttention` removed in favor of static `LocalMapSpec` | pytorch/torchtitan #2963, #2969, #3113, #2986 | Public-API rename + refactor; Primus torchtitan launcher and patches must be re-validated. |
-| Attention / loss | `ChunkedCELoss`;<br>graph-trainer chunked-CE gating;<br>SDP-default-backend revert (re-include MATH);<br>fused QKV in Qwen3-VL adapter | pytorch/torchtitan #2937, #3115, #3135, #3102 | Loss-side knob useful for memory-bound training; SDP backend revert reduces breakage risk for Primus tests when sync lands. |
-| RL / experiments | RL env-rollout-based controller refactor;<br>upstream Qwen3 parallel plan absorbs patched RL plan;<br>VLM and stale autoparallel/dsv3 experiments deprecated/removed | pytorch/torchtitan #3073, #3070, #3151, #2271 | Relevant for Primus post-training/RL on torchtitan; may unblock removing internal RL patches after sync. |
-| Datasets / CI / ROCm | `replace-imports-with-any` CI pattern;<br>reproducible epoch-boundary resume for map/streaming datasets;<br>ChatDataset cross-node shuffle;<br>MoE loss comparison CI guard;<br>MI350 loss-number refresh | pytorch/torchtitan #3180, #3008, #3131, #3081, #3078 | Useful CI/MI350 hygiene reference for Primus torchtitan CI; the loss-comparison guard pattern can be mirrored in Primus torchtitan UTs. |
-
-## 6. Primus-Turbo Monthly Drift Overview
-
-- Drift type: **in-repo**, not upstream — compares Turbo version/SHA referenced on Primus `main` now vs the latest commit at or before `month_start_ts = 2026-04-01 00:00 Asia/Shanghai` (`2026-03-31 16:00 UTC`).
-- Turbo is **not a submodule** in Primus. Canonical version source:
-  - `.github/workflows/ci.yaml` → `PRIMUS_TURBO_COMMIT` (also wired through `.github/workflows/docker/Dockerfile`)
-  - `.github/workflows/benchmark.yaml` → `PRIMUS_TURBO_COMMIT`
-- Reference commit at month start on `main`: `76651575` (*[WIP][Megatron-LM] feat: reduce extra qkv transpose in attn (#625)*, 2026-03-31 14:19 GMT+8). The underlying Turbo pins at that commit are byte-identical to today's values.
-- Current state on `origin/main`:
-  - `ci.yaml` `PRIMUS_TURBO_COMMIT`: `333b68d7c81b722b21b4aad10cd250c45f15027c` — *fix sm_scale none bug (#263)*
-  - `ci.yaml` `PRIMUS_TURBO_AITER_COMMIT`: `e83f9903c07001a0ec29e85d223f6e6cdbe00859`
-  - `benchmark.yaml` `PRIMUS_TURBO_COMMIT`: `a4488f6cdb15cfff4383c61af7922bb50803f0ea` — *feat: update triton impl for mi300 & mi355 (#252)*
-- Month-start state on `main`: all three values identical to current.
-- **No Primus-Turbo drift in this comparison window.**
-- Recommendation: **monitor**. The pre-existing skew between the two YAML pins (CI pin `333b68d` is 5 commits ahead of benchmark pin `a4488f6c` in Primus-Turbo history) is unchanged this month. This week's Primus-side change in #693 modifies how Primus *uses* Turbo's grouped-MLP path on the Megatron backend but does not bump any Turbo pin.
-
-### Notable areas changed since month start
-
-- **No changes this window** — both `ci.yaml` and `benchmark.yaml` Turbo pins on `main` are byte-identical to their 2026-03-30 values.
-- **Primus-side Turbo integration moved**: PR #693 (this window) drops the legacy GroupedGEMM gating in `_is_primus_turbo_enabled` and `PrimusTurboDeepEPTokenDispatcher` so that Turbo's `TEGroupedMLP` path is now the default, but no `PRIMUS_TURBO_COMMIT` was bumped.
-- **AITER pin unchanged**: `PRIMUS_TURBO_AITER_COMMIT` is identical to the month-start value.
-- **Benchmark pin unchanged**: `benchmark.yaml` `PRIMUS_TURBO_COMMIT` is identical to the month-start value.
-- **No Turbo drift in this comparison window.**
-
-### Primus-Turbo monthly drift table
-
-| Component | Current Version/SHA | Month-start Version/SHA | Delta Summary | Key Changes | Evidence |
-| --- | --- | --- | --- | --- | --- |
-| `PRIMUS_TURBO_COMMIT` (CI build) | `333b68d7c81b722b21b4aad10cd250c45f15027c` (*fix sm_scale none bug (#263)*) | `333b68d7c81b722b21b4aad10cd250c45f15027c` | No drift (0 commits) | No changes this window. Primus-side use of Turbo grouped-MLP changed in #693 without bumping the pin. | [`.github/workflows/ci.yaml` L17](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L17) |
-| `PRIMUS_TURBO_AITER_COMMIT` (CI build) | `e83f9903c07001a0ec29e85d223f6e6cdbe00859` | `e83f9903c07001a0ec29e85d223f6e6cdbe00859` | No drift | No changes this window. | [`.github/workflows/ci.yaml` L18](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L18) |
-| `PRIMUS_TURBO_COMMIT` (benchmark) | `a4488f6cdb15cfff4383c61af7922bb50803f0ea` (*feat: update triton impl for mi300 & mi355 (#252)*) | `a4488f6cdb15cfff4383c61af7922bb50803f0ea` | No drift | No changes this window. | [`.github/workflows/benchmark.yaml` L9](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/benchmark.yaml#L9) |
-
-## 7. Source Links
-
-- Primus main branch: https://github.com/AMD-AGI/Primus/tree/main
-- Primus weekly PR listing (window): https://github.com/AMD-AGI/Primus/pulls?q=is%3Apr+is%3Amerged+base%3Amain+merged%3A%3E%3D2026-04-26T16%3A00%3A00Z
-- PR #651: https://github.com/AMD-AGI/Primus/pull/651
-- PR #693: https://github.com/AMD-AGI/Primus/pull/693
-- Megatron-LM pin: https://github.com/NVIDIA/Megatron-LM/commit/d3528a21301db2d12e92912b3ec025dc8a2ed4d6
-- Megatron-LM upstream HEAD (at report time): https://github.com/NVIDIA/Megatron-LM/commit/3460bba1d6f945ec04f47fdb1dcee6da3259fcd8
-- Megatron-LM compare: https://github.com/NVIDIA/Megatron-LM/compare/d3528a21301db2d12e92912b3ec025dc8a2ed4d6...main
-- torchtitan pin: https://github.com/pytorch/torchtitan/commit/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021
-- torchtitan upstream HEAD (at report time): https://github.com/pytorch/torchtitan/commit/70340f4ec31d9e1dbd448506cc4423934c3cd62f
-- torchtitan compare: https://github.com/pytorch/torchtitan/compare/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021...main
-- Primus-Turbo CI pin: https://github.com/AMD-AGI/Primus-Turbo/commit/333b68d7c81b722b21b4aad10cd250c45f15027c
-- Primus-Turbo benchmark pin: https://github.com/AMD-AGI/Primus-Turbo/commit/a4488f6cdb15cfff4383c61af7922bb50803f0ea
-- Month-start reference commit on `main`: https://github.com/AMD-AGI/Primus/commit/76651575
-- Last week's report (W17): https://github.com/AMD-AGI/Primus/blob/main/docs/weekly_reports/2026-W17-primus-weekly.md
-
----
-
-*Generated automatically by the Primus weekly report automation. Factual statements are derived from `git log origin/main`, the pinned submodule SHAs in `third_party/`, and the `PRIMUS_TURBO_COMMIT` values in `.github/workflows/{ci,benchmark}.yaml` as observed at 2026-05-01 09:10 GMT+8. Upstream-HEAD SHAs and commit counts are snapshots at report generation time.*
diff --git a/docs/weekly_reports/2026-W19-primus-weekly.md b/docs/weekly_reports/2026-W19-primus-weekly.md
deleted file mode 100644
index ce44f7996..000000000
--- a/docs/weekly_reports/2026-W19-primus-weekly.md
+++ /dev/null
@@ -1,151 +0,0 @@
-# Primus Weekly Engineering Report — 2026-W19
-
-## 1. Time Window
-
-- Start: Monday 2026-05-04 00:00:00 Asia/Shanghai (GMT+8)
-- End: Friday 2026-05-08 09:01 Asia/Shanghai (GMT+8) (report generation time)
-- Branch observed: `origin/main`
-
-## 2. Executive Summary
-
-- **3 PRs merged to `main`** in the weekly window (Mon 2026-05-04 00:00 GMT+8 → Fri 2026-05-08 09:01 GMT+8).
-- Category breakdown: **Bug Fix: 2**, **Turbo/Dependency Version Update: 1**; Performance Optimization, CI/Infra, Refactor, Docs, Other: 0.
-- **Backend dependency pin change this week (Primus-Turbo + Triton).** PR #694 bumped `PRIMUS_TURBO_COMMIT` from `333b68d7` (CI) / `a4488f6c` (benchmark) to a single shared SHA `ef5b58ea3de0a2956d57dba518be466b7a092442` ([Primus-Turbo #320](https://github.com/AMD-AGI/Primus-Turbo/pull/320), *feat(grouped_gemm): switch default backend to Triton for BF16 and FP8*, 2026-04-29), bumped `PRIMUS_TURBO_AITER_COMMIT` from `e83f9903` to `857f4d1577`, and introduced a new `TRITON_COMMIT=88b227e23f0445f3f695bad05bbf1a363b4f50e0` build-arg that compiles Triton from `release/3.7.x` source inside the docker image (wired through both the main and JAX docker builds). The CI runner now also exports `PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1`. The same PR adds a Primus-Turbo Attention adapter for torchtitan Qwen3 (new `primus/backends/torchtitan/models/qwen3/model/model.py`, plus `patch_turbo_attention` extending the Qwen3 path).
-- **Submodule SHAs unchanged.** `third_party/Megatron-LM` is still pinned at `d3528a21` and `third_party/torchtitan` at `5fb7cc2e` on `origin/main`.
-- **Megatron-LM upstream drift: `plan sync` (unchanged).** Pin is `d3528a21` (2026-03-06); upstream `main` HEAD is `7cdf652c` (2026-05-07) — **458 commits ahead** (+45 since the W18 snapshot `3460bba1`). Notable new upstream activity in this window: removal of the legacy `transformer`/`modules` and legacy GPT code (#4207, #4322), inference cache of input/position ID views (#4634), inference vLLM grouped-gemm kernel tuning + shared-expert overlap in latent MoE (#4603), checkpoint conversion between `GPT_model` and `Hybrid_model` (#4482), `nvidia_resiliency_ext` fault-injection support (#4370), InJob-restart-on-failure support (#4594), partial-cudagraphs+HybridEP DDP-hook fix (#4500), gradient corruption fix for layerwise param all-gather overlap (#4609), Flextron code (#4429), FlashInfer sampling (#2456), and a vLLM grouped-gemm MoE inference backend (#4566).
-- **torchtitan upstream drift: `urgent sync` (unchanged).** Pin is `5fb7cc2e` (2025-10-15); upstream `main` HEAD is `4ebb0895` (2026-05-07) — **610 commits ahead** (+44 since the W18 snapshot `70340f4e`). Major new upstream activity in this window: GraphTrainer memory-policy registry + extra-graph-passes hook (#3215), graph-pass timing/log of names + total time (#3261), HybridEP-with-GraphTrainer integration tests (#3184), regional Inductor RMSNorm fusion (#3132), bucketing pass enabled in precompile path (re-landed #3213/#3107), CPU-offloading prefetch pass (#3166), `[graph_trainer] FSDP AG RS overlap` (#3156), MXFP8 GroupedExperts swap fix (#3199), float8 quantization of GroupedExperts during conversion (#3233), `AllToAllTokenDispatcher` token-count `sp_size` padding (#3193), CP `AllGather` wrong-dimension fix (#3206), llama4 `fsdp_mesh_info` bug fix (#3231), Qwen3 MoE bitwise-deterministic tests + weight-tying gradient fix (#3174), structured-logging observability (#3176), CI parallel integration tests (#3144), and a 50s→15s vLLM compile-time win (#3145).
-- **Primus-Turbo monthly drift: pin advanced — recommendation is now `monitor (new pin landed)`.** Both CI and benchmark `PRIMUS_TURBO_COMMIT` are now `ef5b58ea` (2026-04-29), which is **32 commits ahead** of the month-start CI pin `333b68d7` (2026-03-27) and **37 commits ahead** of the month-start benchmark pin `a4488f6c` (2026-03-19). The headline upstream changes between month-start and the new pin are: switch grouped-gemm default backend to Triton for BF16 and FP8 (#320), `[Attention] Torch.compile custom wrappers` (#310), MoE EPBackend Protocol + EPBufferConfig refactor (#297), `attn add bhsd layout` (#304), Symmetric-Memory rewrite (#276), `feat(moe): add back deepep_use_comm_stream` (#314), and several MXFP8/FP8 quantization fixes (#306, #307, #308). The CI `PRIMUS_TURBO_AITER_COMMIT` was also bumped (`e83f9903` → `857f4d15`), and a new `TRITON_COMMIT=88b227e` build-arg compiles Triton from source.
-- **Megatron-side Primus interface debt closed (#675):** custom HF tokenizer types now route through Megatron's official `build_tokenizer` first with a fallback that injects `unique_identifiers`; data path args are normalized for `str`/`list`/`tuple` using upstream-compatible semantics; legacy `core_gpt_dataset_config_from_args` is realigned with upstream `pretrain_gpt` via `get_blend_and_blend_per_split`; missing upgraded fields (`object_storage_cache_path`, `per_dataset_sequences_path`, `dataloader_fast_cache_load`, etc.) are wired to prevent silent mock-fallback drift.
-- **Preflight reliability fix (#668):** the preflight IP-address probe replaces a substring `in` check with a precise regex match in `primus/tools/preflight/network/network_probe.py`, eliminating a class of false-positives where a target IP is a substring of another active IP.
-- **No backend-gap report regenerated this week.** No submodule SHA changed (`third_party/Megatron-LM`, `third_party/torchtitan`, `third_party/Megatron-Bridge`, `third_party/Emerging-Optimizers`, `third_party/HummingbirdXT`, `third_party/maxtext` are all unchanged). The Primus-Turbo pin bump is a tracked-config change but Primus-Turbo is not currently surfaced as a separate backend under `docs/backend-gap/`, so the existing torchtitan baseline report (`5fb7cc2e` vs upstream `main`) is unaffected.
-
-## 3. Weekly PR Update Table
-
-| PR | Merged Time (GMT+8) | Category | Key Update |
-| --- | --- | --- | --- |
-| [#694](https://github.com/AMD-AGI/Primus/pull/694) `Build Triton from source and bump Primus-Turbo` (author: `kyle-256`) | 2026-05-07 17:19 | Turbo/Dependency Version Update | Bumps `PRIMUS_TURBO_COMMIT` to a single shared SHA `ef5b58ea3de0a2956d57dba518be466b7a092442` in both `.github/workflows/ci.yaml` and `.github/workflows/benchmark.yaml` (CI was at `333b68d7`, benchmark was at `a4488f6c`). Also bumps `PRIMUS_TURBO_AITER_COMMIT` (`e83f9903` → `857f4d1577`). Introduces a new `TRITON_COMMIT=88b227e23f0445f3f695bad05bbf1a363b4f50e0` env, wired through both the main and JAX docker builds, plus Dockerfile steps that clone `triton-lang/triton@release/3.7.x`, check out the pinned commit, and `pip install --no-build-isolation -v .`. Adds `PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1` to the torch CI runner job env. Adds Primus-Turbo Attention support for torchtitan Qwen3 by introducing `primus/backends/torchtitan/models/qwen3/model/model.py` (subclasses upstream `Attention` and routes through `inner_attention` with FlexAttention/non-FlexAttention branches) and extending `patch_turbo_attention` to swap `torchtitan.models.qwen3.model.model.Attention` to the Primus version. Also drops the recursive `submodules: "recursive"` from the Primus-Turbo checkout step in favor of an explicit `git submodule sync --recursive && git submodule update --init --recursive` after cleaning `3rdparty/composable_kernel` to avoid stale submodule state. |
-| [#668](https://github.com/AMD-AGI/Primus/pull/668) `Fix/preflight: fix IP address matching` (author: `alexsu52`) | 2026-05-07 17:26 | Bug Fix | Replaces the substring `in` check used to validate a target IP address against the host's active IPs with a precise regex match (`re.search(r"\b<ip>\b", ...)`) in `primus/tools/preflight/network/network_probe.py`. Eliminates a class of false-positives where a target IP is a substring of another active IP on the same machine (e.g. `10.0.0.1` falsely matching `10.0.0.10`). Single-file, 2-insertion change. |
-| [#675](https://github.com/AMD-AGI/Primus/pull/675) `fix(megatron): align tokenizer and dataset path interfaces with upgraded Megatron` (author: `WangLingxun`) | 2026-05-06 14:00 | Bug Fix | Aligns Primus Megatron-backend behavior with upstream Megatron after the recent upgrade. Routes custom HF tokenizer types through Megatron's official `build_tokenizer` first; on failure, falls back to a Primus path and injects `unique_identifiers` when the constructed tokenizer is missing it (fixes `'HuggingFaceTokenizer' object has no attribute 'unique_identifiers'` during dataset construction). Normalizes `--data-path` style args for `str`, `list`, and `tuple` inputs using upstream-compatible semantics, eliminating `.split()` crashes on list input. Realigns the legacy `core_gpt_dataset_config_from_args` path with upstream `pretrain_gpt` via `get_blend_and_blend_per_split`. Wires missing upgraded fields (`object_storage_cache_path`, `per_dataset_sequences_path`, `dataloader_fast_cache_load`, etc.) into the trainer-side `GPTDatasetConfig` build to prevent real-data config drift and unintended mock-dataset fallback. Touches `primus/backends/megatron/patches/args/data_path_split_patches.py`, `primus/backends/megatron/training/tokenizer/tokenizer.py`, and `primus/modules/trainer/megatron/trainer.py`. |
-
-## 4. Megatron-LM Drift Overview
-
-- Upstream: `https://github.com/NVIDIA/Megatron-LM.git` (`main`)
-- Pinned in Primus `main` (`third_party/Megatron-LM`): `d3528a21301db2d12e92912b3ec025dc8a2ed4d6` — *fix(moe): fix TE general_gemm API change (#3582)*, 2026-03-06
-- Upstream `main` HEAD: `7cdf652c16a0283b42ee13b7700481d0cfe2f291` — *ci: Update Gitlab base image to 26.04 pytorch (#4688)* (2026-05-07)
-- Last upstream functional change in this window: `1df264c4` — *Inference: Cache input + position ID views (#4634)* (2026-05-07)
-- Commit gap: **upstream is 458 commits ahead of Primus pin** (+45 since the W18 snapshot `3460bba1`).
-- Submodule SHA on Primus side: unchanged in W19; last submodule SHA bump on `main` was `3bec9aa9` → `d3528a21` inside PR #654 (merged 2026-04-10).
-- Recommendation: **plan sync** (unchanged from W17/W18). Upstream continues to land large structural removals (legacy `transformer`/`modules`, legacy GPT code) and inference perf/correctness work, growing the eventual sync diff but not changing Primus's overall sync risk profile this week.
-
-### Notable upstream areas that have moved since the pin
-
-- **Legacy code removal**: legacy `transformer`/`modules` removed (#4207); legacy GPT code removed (#4322); these are large structural deletes that will dominate the eventual Primus sync diff.
-- **Inference**: cache input + position ID views (#4634); propagate errors for failed inference requests (#4679); fix crash with evicted requests + tpot (#4645); FlashInfer sampling (#2456); make `last_token_logits` graphable (#4552); shared-expert overlap with `allgatherv` in inference (#4570); MoE dispatcher buffer-size fix from actual tensor sizes (#4576); inference EP-sync re-enable + simplify (#4587); `Move inference context bookkeeping to CPU with ContextGPUView` (#4306); on top of W18's AllgatherV inference dispatcher (#4258), MTP CUDA graphs (#4260) and inference graph standardization (#4485).
-- **MoE / EP**: Tune vLLM grouped gemm + `moe_sum` + shared-expert overlap in latent MoEs (#4603); Add vLLM grouped gemm backend for MoE inference (#4566); fix partial cudagraphs + HybridEP DDP-hook trigger (#4500); fix EP sync regression (#4607); chunked MLP during training (#3656); on top of W18's permute fusion in hybrid EP (#4089) and per-block MoE routing storage for prefix caching (#4301).
-- **Resiliency / training stability**: `nvidia_resiliency_ext` fault-injection (#4370); InJob restart on failures (#4594); fix gradient corruption with layerwise param all-gather overlap (#4609); fix Hang in tests (#4575); KD teacher loading moved after `Float16Module` (#4394); guard vocab `reduce_scatter` on TP > 1 (#4565); `Remove invalid timeout argument for dist.barrier` (#4512).
-- **Hybrid / Mamba / SSM**: Checkpoint conversion between `GPT_model` and `Hybrid_model` (#4482); Handle SSM sharded tensor merge OOM with CPU fallback (#4442); fix `mtp_use_repeated_layer` behavior for GPT models (#3965); on top of W18's `causal_conv1d_update` HBM-reload reduction (#4460), `mamba-ssm`/`causal-conv1d` move to optional `[ssm]` extra (#4517).
-- **Tokenizer / data / experiment**: convert tokenizer args to config (#4406); `Finalize all builders in preprocess_data, not just the last key` (#4573); Named validation sets (#4578); Adding code for Flextron (#4429); refit buffer fix (#4580); skill-doc refactor (#4574).
-
-### Megatron-LM upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| Legacy code removal | Legacy `transformer`/`modules` removed;<br>legacy GPT code removed | NVIDIA/Megatron-LM #4207, #4322 | Large structural deletes that will dominate the eventual Primus sync diff; Primus patches under `primus/backends/megatron/**` that touch the legacy paths must be audited before the next pin bump. |
-| Inference / CUDA graphs | Cache input + position ID views;<br>propagate errors for failed inference requests;<br>FlashInfer sampling;<br>`Move inference context bookkeeping to CPU with ContextGPUView`;<br>(carries) AllgatherV inference dispatcher, MTP CUDA graphs | NVIDIA/Megatron-LM #4634, #4679, #2456, #4306, #4258, #4260 | Inference/post-train paths for Primus DSV3/MoE configs benefit; expand validation when sync lands. |
-| MoE / EP | Tune vLLM grouped gemm + `moe_sum` + shared-expert overlap in latent MoEs;<br>Add vLLM grouped gemm MoE-inference backend;<br>partial-cudagraphs + HybridEP DDP-hook fix;<br>EP-sync regression fix;<br>chunked MLP during training | NVIDIA/Megatron-LM #4603, #4566, #4500, #4607, #3656 | Aligns with this week's Primus-side Turbo bump (#694) which already moves grouped-gemm default to Triton in the Turbo path; once Primus syncs, Megatron's vLLM grouped-gemm MoE backend becomes available for Primus DSV3/MoE inference experiments. |
-| Resiliency / stability | `nvidia_resiliency_ext` fault injection;<br>InJob restart on failures;<br>gradient-corruption fix for layerwise param AG overlap;<br>Hang-in-tests fix | NVIDIA/Megatron-LM #4370, #4594, #4609, #4575 | Useful for Primus pretrain-at-scale on MI300X/MI355X clusters; coordinate with Primus preflight/launcher when bumping the pin. |
-| Hybrid / Mamba / SSM | Checkpoint conversion between `GPT_model` and `Hybrid_model`;<br>SSM sharded-tensor merge OOM CPU fallback;<br>(carries) `MambaModel`/`MambaStack` → `HybridModel`/`HybridStack` rename | NVIDIA/Megatron-LM #4482, #4442, #4159 | The `[ssm]` extra split + Hybrid rename remains a known sync risk for Primus Mamba/Hybrid configs; Primus must audit `MambaModel`/`MambaStack` references before the next bump. |
-| Tokenizer / data / experiment | Convert tokenizer args to config;<br>`Finalize all builders in preprocess_data`;<br>Named validation sets;<br>Flextron code | NVIDIA/Megatron-LM #4406, #4573, #4578, #4429 | This week's Primus tokenizer/data alignment (#675) deliberately matches upstream `pretrain_gpt`; the upstream tokenizer-args→config refactor is the next candidate Primus must adopt at sync time. |
-| Schedule plan API (carry-over) | `post_attn` node already removed from `TransformerLayerSchedulePlan` (consumed in Primus by W17 #672) | Primus/#672, upstream schedule-plan change | Primus-side fix already shipped; treat as confirmed drift that future bumps must keep in sync. |
-
-## 5. torchtitan Drift Overview
-
-- Upstream: `https://github.com/pytorch/torchtitan.git` (`main`)
-- Pinned in Primus `main` (`third_party/torchtitan`): `5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021` — *Deepseek-V3 toml file minor fix (#1894)*, 2025-10-15
-- Upstream `main` HEAD: `4ebb0895a18859ad0f0642901dba0cc888d82b97` — *Tests HybridEP integration with GraphTrainer (#3184)* (2026-05-07)
-- Commit gap: **upstream is 610 commits ahead of Primus pin** (+44 since the W18 snapshot `70340f4e`).
-- Submodule SHA on Primus side: unchanged in W19. PR #694 in this window adds a Primus-Turbo Qwen3 attention adapter on the **Primus outer adapter** layer (`primus/backends/torchtitan/models/qwen3/`) but does not bump the upstream pin.
-- Recommendation: **urgent sync** (unchanged from W17/W18). The pin is now ~6.5 months stale; another wave of GraphTrainer / precompile / float8 / MoE landings this week further widens the gap.
-
-### Notable upstream areas that have moved since the pin
-
-- **GraphTrainer / precompile**: memory-policy registry + extra graph-passes hook (#3215); log graph-pass names + total time (#3261); regional Inductor RMSNorm fusion (#3132); bucketing pass enabled in precompile path (re-landed via #3213, originally #3107); CPU-offloading prefetch pass (#3166); FSDP AG RS overlap (#3156); FlexAttention precompile bitwise-deterministic tests re-landed via #3214 (originally #3178); `Port remat_using_tags_for_fwd_loss_bwd_graph pass locally` (#3260); `Skip identity-slice rewrite when start/end/step are dynamic Nodes` (#3195); annotate generated FX with user source lines (#3194); annotate loss region with `module_fqn` (#3207); on top of W18's GraphTrainer activation-memory framework (#3118), full inductor compilation pass (#3141), async-TP graph pass (#3129).
-- **Compilation / vLLM**: improve compile time (~50s → ~15s for vLLM) (#3145); `Refactor pipeline parallel helpers for graph PP reuse` (#2724); `Re-land "Enable bucketing pass in precompile path"` (#3213).
-- **MoE / float8 / mxfp8**: `[float8] Quantize GroupedExperts params during conversion` (#3233); `[mxfp8] Fix MXFP8GroupedExpertsConverter to actually swap GroupedExperts params` (#3199); `[MoE] Pad token count to a multiple of sp_size in AllToAllTokenDispatcher` (#3193); HybridEP-with-GraphTrainer integration tests (#3184); on top of W18's `All2AllTokenDispatcher` consolidation (#3125) and ETP deprecation (#3167).
-- **Attention / Qwen3 / Llama4**: `[GraphTrainer] Add Qwen3 MoE bitwise deterministic tests and fix weight-tying gradient bug` (#3174); `[Qwen3VL] remove local_map for dtensor interpolation` (#2957); `[torchtitan][llama4_parallelize] fixing fsdp_mesh_info bug` (#3231); `refactor get_attention_masks to use positions instead of eos_id` (#3149); CP `AllGather on the wrong dimension` (#3206); `Wrong token count in validation when CP>1` (#3257).
-- **Observability / RL / CI**: structured logging + training instrumentation (#3176); `[CI] Run integration tests in parallel` (#3144); `ci: regenerate qwen3_moe_rocm_mi350x.txt baseline with actual MI350 losses` (#3196); `[GraphTrainer] Disable cpu_offload_all CI test for upstream cuDNN regression` (#3252); RL CI loss=0/logprob=NaN fix (#3232); `[rl] Add a tag to trigger RL CI runs` (#3230); RL `expandable_segments:True` for monarch RDMA (#3221); torchtitan-ubuntu docker upgraded to 22.04 + CTK 13.0 (#3183).
-
-### torchtitan upstream feature delta table
-
-| Area | New Upstream Capability | Evidence (PR/Commit) | Potential Impact to Primus |
-| --- | --- | --- | --- |
-| GraphTrainer / precompile | Memory policy registry + extra graph passes hook;<br>graph-pass timing + name logging;<br>regional Inductor RMSNorm fusion;<br>bucketing pass in precompile path (re-land);<br>CPU-offloading prefetch;<br>FSDP AG RS overlap;<br>FlexAttention precompile bitwise-deterministic tests (re-land) | pytorch/torchtitan #3215, #3261, #3132, #3213, #3166, #3156, #3214 | Major continuing GraphTrainer infrastructure; remains unavailable behind the stale pin. Will require coordinated patches in `primus/backends/torchtitan/**` when Primus syncs. |
-| Compilation / vLLM | 50s → 15s vLLM compile-time win;<br>pipeline-parallel helpers refactored for graph PP reuse | pytorch/torchtitan #3145, #2724 | Direct training-loop wall-clock win when Primus syncs the pin. |
-| MoE / float8 / mxfp8 | float8 quantize GroupedExperts params during conversion;<br>MXFP8GroupedExpertsConverter swap fix;<br>`AllToAllTokenDispatcher` `sp_size` token-count padding;<br>HybridEP-with-GraphTrainer integration tests | pytorch/torchtitan #3233, #3199, #3193, #3184 | Useful for Primus MoE configs (DSV3, GPT-OSS, Qwen3) on the torchtitan backend; aligns conceptually with this week's Primus-side Qwen3 attention adapter (#694) but the upstream consolidation of MoE dispatch is still gated by the urgent sync. |
-| Attention / Qwen3 / Llama4 | Qwen3 MoE bitwise-deterministic tests + weight-tying gradient bug fix;<br>`get_attention_masks` refactored to positions;<br>llama4 `fsdp_mesh_info` fix;<br>CP `AllGather` wrong-dim fix;<br>CP>1 validation token-count fix | pytorch/torchtitan #3174, #3149, #3231, #3206, #3257 | Touches the same code paths as this week's Primus-Turbo Qwen3 attention adapter (#694); after sync, Primus's Qwen3 patch will need to be re-validated against the upstream weight-tying fix. |
-| Observability / RL / CI | Structured logging + training instrumentation;<br>parallel CI integration tests;<br>MI350 baseline regenerated with actual losses;<br>RL CI loss/logprob fix | pytorch/torchtitan #3176, #3144, #3196, #3232 | Reference patterns for Primus torchtitan CI; the MI350 baseline regen pattern is directly applicable to Primus MI350-class CI runners. |
-
-## 6. Primus-Turbo Monthly Drift Overview
-
-- Drift type: **in-repo**, not upstream — compares Turbo version/SHA referenced on Primus `main` now vs the latest commit at or before `month_start_ts = 2026-05-01 00:00 Asia/Shanghai` (`2026-04-30 16:00 UTC`). For consistency with the W17/W18 reports, the **April→May** comparison anchor is the prior month-start (2026-04-01 00:00 GMT+8); both anchors give the same conclusion this week because the pin only moved on 2026-05-07.
-- Turbo is **not a submodule** in Primus. Canonical version source:
-  - `.github/workflows/ci.yaml` → `PRIMUS_TURBO_COMMIT`, `PRIMUS_TURBO_AITER_COMMIT`, `TRITON_COMMIT` (also wired through `.github/workflows/docker/Dockerfile`)
-  - `.github/workflows/benchmark.yaml` → `PRIMUS_TURBO_COMMIT`
-- Reference Primus commit at month start (April) on `main`: `766515755e9c29c11ed55e213dbc82f6581ca31e` (*[WIP][Megatron-LM] feat: reduce extra qkv transpose in attn (#625)*, 2026-03-31 14:19 GMT+8). The Turbo pins at that commit are byte-identical to the values reported in W17/W18.
-- Current state on `origin/main` (W19):
-  - `ci.yaml` `PRIMUS_TURBO_COMMIT`: `ef5b58ea3de0a2956d57dba518be466b7a092442` — *feat(grouped_gemm): switch default backend to Triton for BF16 and FP8 (#320)*, 2026-04-29
-  - `ci.yaml` `PRIMUS_TURBO_AITER_COMMIT`: `857f4d15775a29af153a2c68a2f8e8a8d696c986`
-  - `ci.yaml` `TRITON_COMMIT` (new env var, introduced this week): `88b227e23f0445f3f695bad05bbf1a363b4f50e0`
-  - `benchmark.yaml` `PRIMUS_TURBO_COMMIT`: `ef5b58ea3de0a2956d57dba518be466b7a092442`
-- Month-start (2026-04-01) state on `main`:
-  - `ci.yaml` `PRIMUS_TURBO_COMMIT`: `333b68d7c81b722b21b4aad10cd250c45f15027c` — *fix sm_scale none bug (#263)*, 2026-03-27
-  - `ci.yaml` `PRIMUS_TURBO_AITER_COMMIT`: `e83f9903c07001a0ec29e85d223f6e6cdbe00859`
-  - `benchmark.yaml` `PRIMUS_TURBO_COMMIT`: `a4488f6cdb15cfff4383c61af7922bb50803f0ea` — *feat: update triton impl for mi300 & mi355 (#252)*, 2026-03-19
-  - `ci.yaml` `TRITON_COMMIT`: not present (new env var introduced by #694 in this window).
-- **Primus-Turbo pin advanced this week.** CI pin moved 32 commits ahead of month start; benchmark pin moved 37 commits ahead of month start; AITER pin changed; a new explicit Triton source build was added.
-- Recommendation: **monitor (new pin landed)**. The two YAML pins are now in sync (CI and benchmark both `ef5b58ea`), removing the previously reported skew. Track the next Turbo upstream bump and validate the new Triton-backed grouped-gemm default path on Primus MoE configs.
-
-### Notable areas changed since month start
-
-- **Grouped GEMM default backend**: Primus-Turbo #320 switches the default grouped-gemm backend to Triton for BF16 and FP8; combined with the new `TRITON_COMMIT=88b227e` source build, this changes the kernel path used by Primus-Turbo's `TEGroupedMLP` integration on the Megatron backend (the Primus-side gating was already removed in W18 #693).
-- **Attention `torch.compile` wrappers**: Primus-Turbo #310 adds `torch.compile` custom wrappers for the attention path; relevant for both the Megatron backend and the new torchtitan Qwen3 attention adapter introduced in #694.
-- **MoE EPBackend Protocol + `EPBufferConfig`**: Primus-Turbo #297 refactors the dispatcher around an extensible `EPBackend` Protocol with an explicit `EPBufferConfig`; Primus's `PrimusTurboDeepEPTokenDispatcher` (`primus/backends/megatron/core/transformer/moe/`) consumes this surface and may need a follow-up after this bump.
-- **Attention layouts**: Primus-Turbo #275 (sbhd) and #304 (bhsd) extend attention layout coverage; aligns with this week's Primus Qwen3 attention adapter.
-- **Symmetric Memory rewrite**: Primus-Turbo #276 replaces `hip-python` and `AMDSymmetricMemory` with a new `SymmetricMemory`; relevant when running DeepEP/MoE paths on ROCm.
-- **MXFP8/FP8 quantization fixes**: Primus-Turbo #298 (MXScalingRecipe rename + grouped-gemm FP8 API), #306 (`reduce_row_kernel` partial-tile indexing), #307 (mxfp8 int32 overflow), #308 (mxfp8 gemm WAR race) tighten correctness on FP8 paths used by Primus MoE configs.
-- **DeepEP single-stream toggle**: Primus-Turbo #314 brings back `deepep_use_comm_stream` to force single-stream comm; useful escape hatch for DeepEP correctness debugging on ROCm.
-- **AITER pin moved**: `PRIMUS_TURBO_AITER_COMMIT` changed from `e83f9903` (month-start) to `857f4d15` (current); the Turbo build pulls AITER at this SHA inside the docker image.
-- **Triton built from source**: a new explicit `TRITON_COMMIT=88b227e` is wired through `ci.yaml`, `benchmark.yaml` does not pin Triton today; the Dockerfile clones `triton-lang/triton@release/3.7.x`, checks out `${TRITON_COMMIT}`, and `pip install --no-build-isolation -v .` to compile Triton from source.
-
-### Primus-Turbo monthly drift table
-
-| Component | Current Version/SHA | Month-start Version/SHA | Delta Summary | Key Changes | Evidence |
-| --- | --- | --- | --- | --- | --- |
-| `PRIMUS_TURBO_COMMIT` (CI build) | `ef5b58ea3de0a2956d57dba518be466b7a092442` (*feat(grouped_gemm): switch default backend to Triton for BF16 and FP8 (#320)*, 2026-04-29) | `333b68d7c81b722b21b4aad10cd250c45f15027c` (*fix sm_scale none bug (#263)*, 2026-03-27) | **+32 commits** | Grouped-GEMM default → Triton (#320);<br>`Torch.compile` attention wrappers (#310);<br>EPBackend Protocol + EPBufferConfig refactor (#297);<br>SymmetricMemory rewrite (#276);<br>MXFP8/FP8 quant fixes (#306, #307, #308);<br>`deepep_use_comm_stream` re-introduction (#314). | [`.github/workflows/ci.yaml` L17](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L17), [Primus-Turbo compare](https://github.com/AMD-AGI/Primus-Turbo/compare/333b68d7c81b722b21b4aad10cd250c45f15027c...ef5b58ea3de0a2956d57dba518be466b7a092442) |
-| `PRIMUS_TURBO_AITER_COMMIT` (CI build) | `857f4d15775a29af153a2c68a2f8e8a8d696c986` | `e83f9903c07001a0ec29e85d223f6e6cdbe00859` | **Pin advanced** | AITER pin advanced by #694; build pulls AITER at the new SHA inside the docker image. | [`.github/workflows/ci.yaml` L18](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L18) |
-| `PRIMUS_TURBO_COMMIT` (benchmark) | `ef5b58ea3de0a2956d57dba518be466b7a092442` (*feat(grouped_gemm): switch default backend to Triton for BF16 and FP8 (#320)*, 2026-04-29) | `a4488f6cdb15cfff4383c61af7922bb50803f0ea` (*feat: update triton impl for mi300 & mi355 (#252)*, 2026-03-19) | **+37 commits** | Same set of upstream changes as the CI pin plus the earlier interval (`a4488f6c`→`333b68d7`); benchmark and CI pins are now in sync. | [`.github/workflows/benchmark.yaml` L9](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/benchmark.yaml#L9), [Primus-Turbo compare](https://github.com/AMD-AGI/Primus-Turbo/compare/a4488f6cdb15cfff4383c61af7922bb50803f0ea...ef5b58ea3de0a2956d57dba518be466b7a092442) |
-| `TRITON_COMMIT` (CI build, new) | `88b227e23f0445f3f695bad05bbf1a363b4f50e0` (`triton-lang/triton@release/3.7.x`) | not present | **New env var** | Compiles Triton from source inside the docker image; wired through both the main and JAX docker builds. | [`.github/workflows/ci.yaml` L21](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/ci.yaml#L21), [`Dockerfile`](https://github.com/AMD-AGI/Primus/blob/main/.github/workflows/docker/Dockerfile) |
-
-## 7. Source Links
-
-- Primus main branch: https://github.com/AMD-AGI/Primus/tree/main
-- Primus weekly PR listing (window): https://github.com/AMD-AGI/Primus/pulls?q=is%3Apr+is%3Amerged+base%3Amain+merged%3A%3E%3D2026-05-03T16%3A00%3A00Z
-- PR #694 (Build Triton from source and bump Primus-Turbo): https://github.com/AMD-AGI/Primus/pull/694
-- PR #675 (fix(megatron): align tokenizer and dataset path interfaces with upgraded Megatron): https://github.com/AMD-AGI/Primus/pull/675
-- PR #668 (Fix/preflight: fix IP address matching): https://github.com/AMD-AGI/Primus/pull/668
-- Megatron-LM pin: https://github.com/NVIDIA/Megatron-LM/commit/d3528a21301db2d12e92912b3ec025dc8a2ed4d6
-- Megatron-LM upstream HEAD (at report time): https://github.com/NVIDIA/Megatron-LM/commit/7cdf652c16a0283b42ee13b7700481d0cfe2f291
-- Megatron-LM compare: https://github.com/NVIDIA/Megatron-LM/compare/d3528a21301db2d12e92912b3ec025dc8a2ed4d6...main
-- torchtitan pin: https://github.com/pytorch/torchtitan/commit/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021
-- torchtitan upstream HEAD (at report time): https://github.com/pytorch/torchtitan/commit/4ebb0895a18859ad0f0642901dba0cc888d82b97
-- torchtitan compare: https://github.com/pytorch/torchtitan/compare/5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021...main
-- Primus-Turbo (new shared) pin: https://github.com/AMD-AGI/Primus-Turbo/commit/ef5b58ea3de0a2956d57dba518be466b7a092442
-- Primus-Turbo CI pin (W18 → W19) compare: https://github.com/AMD-AGI/Primus-Turbo/compare/333b68d7c81b722b21b4aad10cd250c45f15027c...ef5b58ea3de0a2956d57dba518be466b7a092442
-- Primus-Turbo benchmark pin (W18 → W19) compare: https://github.com/AMD-AGI/Primus-Turbo/compare/a4488f6cdb15cfff4383c61af7922bb50803f0ea...ef5b58ea3de0a2956d57dba518be466b7a092442
-- Triton release branch (new pin): https://github.com/triton-lang/triton/commit/88b227e23f0445f3f695bad05bbf1a363b4f50e0
-- Month-start reference commit on `main` (April): https://github.com/AMD-AGI/Primus/commit/76651575
-- Last week's report (W18): https://github.com/AMD-AGI/Primus/blob/main/docs/weekly_reports/2026-W18-primus-weekly.md
-
----
-
-*Generated automatically by the Primus weekly report automation. Factual statements are derived from `git log origin/main`, the pinned submodule SHAs in `third_party/`, and the `PRIMUS_TURBO_COMMIT`/`PRIMUS_TURBO_AITER_COMMIT`/`TRITON_COMMIT` values in `.github/workflows/{ci,benchmark}.yaml` as observed at 2026-05-08 09:01 GMT+8. Upstream-HEAD SHAs and commit counts are snapshots at report generation time.*
diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W17.json b/docs/weekly_reports/dashboard-data/reports/2026-W17.json
deleted file mode 100644
index ad97b3c16..000000000
--- a/docs/weekly_reports/dashboard-data/reports/2026-W17.json
+++ /dev/null
@@ -1,69 +0,0 @@
-{
-  "report_id": "2026-W17",
-  "content_type": "weekly-report",
-  "title": "Primus Weekly Engineering Report — 2026-W17",
-  "report_path": "docs/weekly_reports/2026-W17-primus-weekly.md",
-  "report_github_url": "https://github.com/AMD-AGI/Primus/blob/main/docs/weekly_reports/2026-W17-primus-weekly.md",
-  "time_window": {
-    "timezone": "Asia/Shanghai",
-    "start": "2026-04-20T00:00:00+08:00",
-    "end": "2026-04-24T16:56:00+08:00"
-  },
-  "generated_at": "2026-04-24T16:56:00+08:00",
-  "merged_pr_count": 9,
-  "category_breakdown": {
-    "Bug Fix": 3,
-    "Performance Optimization": 2,
-    "Turbo/Dependency Version Update": 0,
-    "CI/Infra": 2,
-    "Refactor": 0,
-    "Docs": 2,
-    "Other": 0
-  },
-  "megatron_status": {
-    "pin_commit": "d3528a21301db2d12e92912b3ec025dc8a2ed4d6",
-    "pin_date": "2026-03-06",
-    "upstream_repo": "https://github.com/NVIDIA/Megatron-LM",
-    "upstream_ref": "main",
-    "upstream_head": "a1165fabcad97eae3778f386839c233dfabf3f8b",
-    "upstream_head_date": "2026-04-24",
-    "commit_gap": 365,
-    "recommendation": "plan sync"
-  },
-  "torchtitan_status": {
-    "pin_commit": "5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021",
-    "pin_date": "2025-10-15",
-    "upstream_repo": "https://github.com/pytorch/torchtitan",
-    "upstream_ref": "main",
-    "upstream_head": "d40df991ac535108e428b0746a08b74a3cf6afc7",
-    "upstream_head_date": "2026-04-24",
-    "commit_gap": 514,
-    "recommendation": "urgent sync"
-  },
-  "primus_turbo_status": {
-    "drift_type": "monthly_in_repo",
-    "month_start_ts": "2026-04-01T00:00:00+08:00",
-    "current_ci_commit": "333b68d7c81b722b21b4aad10cd250c45f15027c",
-    "month_start_ci_commit": "333b68d7c81b722b21b4aad10cd250c45f15027c",
-    "current_benchmark_commit": "a4488f6cdb15cfff4383c61af7922bb50803f0ea",
-    "month_start_benchmark_commit": "a4488f6cdb15cfff4383c61af7922bb50803f0ea",
-    "current_aiter_commit": "e83f9903c07001a0ec29e85d223f6e6cdbe00859",
-    "month_start_aiter_commit": "e83f9903c07001a0ec29e85d223f6e6cdbe00859",
-    "drift_commits": 0,
-    "recommendation": "monitor"
-  },
-  "recommendations": {
-    "megatron": "plan sync",
-    "torchtitan": "urgent sync",
-    "primus_turbo": "monitor"
-  },
-  "key_findings": [
-    "9 PRs merged to main in the 2026-W17 window (Bug Fix: 3, Performance Optimization: 2, CI/Infra: 2, Docs: 2).",
-    "Week dominated by Megatron backend hardening against upstream API drift: #672 realigns the patched MoE overlap schedule after upstream removed post_attn; #671 makes the muon optimizer wrapper signature-aware; #674 rewrites recompute_layer_patches byte-identical to upstream with a SHA256 fingerprint guard.",
-    "Performance: #673 parallelizes per-rank pp_warmup with a bit-identical loss-parity UT; #684 introduces opt-in PRIMUS_EXIT_FAST shaving ~22s off post-train teardown on MI355X DSV3 EP8 and ~2m off the MI300X Megatron-LM E2E UT suite.",
-    "Infra: #687 lands the shared backend-gap dashboard publishing toolchain under tools/backend_gap_report/ plus the initial torchtitan baseline report; this weekly run extends the same shared site to surface Weekly Reports as a first-class section.",
-    "Megatron-LM upstream drift: pin d3528a21 (2026-03-06) is 365 commits behind upstream HEAD a1165fab; MoE router, Mamba→Hybrid rename (outside core), TE v2.14, DDP parameter-layout refactor, NVRx async compat. Recommendation: plan sync.",
-    "torchtitan upstream drift: pin 5fb7cc2e (2025-10-15) is 514 commits behind upstream HEAD d40df991; GraphTrainer precompile, MoE token dispatcher, Fused QKV GQAttention, FlexAttention CP, FSDP2 fully_shard. Recommendation: urgent sync.",
-    "Primus-Turbo month-to-date drift: none. CI and benchmark PRIMUS_TURBO_COMMIT pins byte-identical to 2026-03-30 values. Recommendation: monitor."
-  ]
-}
diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W18.json b/docs/weekly_reports/dashboard-data/reports/2026-W18.json
deleted file mode 100644
index c995e49ca..000000000
--- a/docs/weekly_reports/dashboard-data/reports/2026-W18.json
+++ /dev/null
@@ -1,69 +0,0 @@
-{
-  "report_id": "2026-W18",
-  "content_type": "weekly-report",
-  "title": "Primus Weekly Engineering Report — 2026-W18",
-  "report_path": "docs/weekly_reports/2026-W18-primus-weekly.md",
-  "report_github_url": "https://github.com/AMD-AGI/Primus/blob/main/docs/weekly_reports/2026-W18-primus-weekly.md",
-  "time_window": {
-    "timezone": "Asia/Shanghai",
-    "start": "2026-04-27T00:00:00+08:00",
-    "end": "2026-05-01T09:10:00+08:00"
-  },
-  "generated_at": "2026-05-01T09:10:00+08:00",
-  "merged_pr_count": 2,
-  "category_breakdown": {
-    "Bug Fix": 1,
-    "Performance Optimization": 0,
-    "Turbo/Dependency Version Update": 0,
-    "CI/Infra": 0,
-    "Refactor": 0,
-    "Docs": 0,
-    "Other": 1
-  },
-  "megatron_status": {
-    "pin_commit": "d3528a21301db2d12e92912b3ec025dc8a2ed4d6",
-    "pin_date": "2026-03-06",
-    "upstream_repo": "https://github.com/NVIDIA/Megatron-LM",
-    "upstream_ref": "main",
-    "upstream_head": "3460bba1d6f945ec04f47fdb1dcee6da3259fcd8",
-    "upstream_head_date": "2026-05-01",
-    "commit_gap": 413,
-    "recommendation": "plan sync"
-  },
-  "torchtitan_status": {
-    "pin_commit": "5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021",
-    "pin_date": "2025-10-15",
-    "upstream_repo": "https://github.com/pytorch/torchtitan",
-    "upstream_ref": "main",
-    "upstream_head": "70340f4ec31d9e1dbd448506cc4423934c3cd62f",
-    "upstream_head_date": "2026-04-30",
-    "commit_gap": 566,
-    "recommendation": "urgent sync"
-  },
-  "primus_turbo_status": {
-    "drift_type": "monthly_in_repo",
-    "month_start_ts": "2026-04-01T00:00:00+08:00",
-    "current_ci_commit": "333b68d7c81b722b21b4aad10cd250c45f15027c",
-    "month_start_ci_commit": "333b68d7c81b722b21b4aad10cd250c45f15027c",
-    "current_benchmark_commit": "a4488f6cdb15cfff4383c61af7922bb50803f0ea",
-    "month_start_benchmark_commit": "a4488f6cdb15cfff4383c61af7922bb50803f0ea",
-    "current_aiter_commit": "e83f9903c07001a0ec29e85d223f6e6cdbe00859",
-    "month_start_aiter_commit": "e83f9903c07001a0ec29e85d223f6e6cdbe00859",
-    "drift_commits": 0,
-    "recommendation": "monitor"
-  },
-  "recommendations": {
-    "megatron": "plan sync",
-    "torchtitan": "urgent sync",
-    "primus_turbo": "monitor"
-  },
-  "key_findings": [
-    "2 PRs merged to main in the 2026-W18 window (Bug Fix: 1, Other/feature: 1). No backend pin or Turbo pin changed this week.",
-    "PR #693 finishes the Megatron-backend Turbo TEGroupedMLP migration: removes the patch-level guard that disabled PrimusTurbo for use_turbo_grouped_mlp + moe_grouped_gemm + TE>=1.9.0; adds a hard validate_args_on_rocm rejection of use_turbo_grouped_mlp=True combined with moe_use_legacy_grouped_gemm=True; preserves the legacy GroupedMLP under primus/backends/megatron/core/transformer/moe/deprecated_2caa681a1/; drops stale moe_use_legacy_grouped_gemm: true lines from MI300X/MI355X DeepSeek-V2/V3, GPT-OSS-20B, and Qwen3 30B/235B example YAMLs.",
-    "PR #651 adds initial LFM/LFM2 Megatron support: lfm_base.yaml + lfm2_8B_A1B.yaml configs, Lfm2ShortConv layer, gpt_decoder_layer_specs_patches.py routing get_gpt_decoder_layer_specs to the Primus implementation, and three MI355X lfm2_8B_A1B example pretrain configs (BF16, FP8, FP8-te-precision). Also commented out four rocmshared docker login lines in the build/push CI job.",
-    "Megatron-LM upstream drift: pin d3528a21 (2026-03-06) is now 413 commits behind upstream HEAD 3460bba1 (2026-05-01), +48 since W17. New since W17: AllgatherV inference dispatcher (#4258), permute fusion in hybrid EP (#4089), MTP CUDA graphs (#4260), MFSDP doc unification (#4418), heterogeneous TP/DP MIMO ColocatedBridgeCommunicator (#4368), checkpoint integrity verification (#4305), mamba-ssm/causal-conv1d move to optional [ssm] extra (#4517). Recommendation: plan sync (unchanged).",
-    "torchtitan upstream drift: pin 5fb7cc2e (2025-10-15) is now 566 commits behind upstream HEAD 70340f4e (2026-04-30), +52 since W17. New since W17: GraphTrainer unified activation-memory framework (#3118), full inductor compilation pass (#3141), async-TP graph pass (#3129), HybridEP integration with GraphTrainer (#3007), [MoE] All2AllTokenDispatcher consolidation across EP=1/EP>1 (#3125), ETP deprecation (#3167), Full DTensor config-based sharding for Llama3/Qwen3/Llama4/DSV3/GPT-OSS (#2963, #2969), MeshDimName→MeshAxisName rename (#3113), ChunkedCELoss (#2937). Recommendation: urgent sync (unchanged).",
-    "Primus-Turbo month-to-date drift (April → May 1 GMT+8): none. CI and benchmark PRIMUS_TURBO_COMMIT pins byte-identical to 2026-03-30 values. Recommendation: monitor.",
-    "Backend-gap reports were not refreshed this week: no submodule SHA change, no PRIMUS_TURBO_* pin change, and no other tracked backend-version source moved in the weekly window."
-  ]
-}
diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W19.json b/docs/weekly_reports/dashboard-data/reports/2026-W19.json
deleted file mode 100644
index 81bca2e62..000000000
--- a/docs/weekly_reports/dashboard-data/reports/2026-W19.json
+++ /dev/null
@@ -1,73 +0,0 @@
-{
-  "report_id": "2026-W19",
-  "content_type": "weekly-report",
-  "title": "Primus Weekly Engineering Report — 2026-W19",
-  "report_path": "docs/weekly_reports/2026-W19-primus-weekly.md",
-  "report_github_url": "https://github.com/AMD-AGI/Primus/blob/main/docs/weekly_reports/2026-W19-primus-weekly.md",
-  "time_window": {
-    "timezone": "Asia/Shanghai",
-    "start": "2026-05-04T00:00:00+08:00",
-    "end": "2026-05-08T09:01:00+08:00"
-  },
-  "generated_at": "2026-05-08T09:01:00+08:00",
-  "merged_pr_count": 3,
-  "category_breakdown": {
-    "Bug Fix": 2,
-    "Performance Optimization": 0,
-    "Turbo/Dependency Version Update": 1,
-    "CI/Infra": 0,
-    "Refactor": 0,
-    "Docs": 0,
-    "Other": 0
-  },
-  "megatron_status": {
-    "pin_commit": "d3528a21301db2d12e92912b3ec025dc8a2ed4d6",
-    "pin_date": "2026-03-06",
-    "upstream_repo": "https://github.com/NVIDIA/Megatron-LM",
-    "upstream_ref": "main",
-    "upstream_head": "7cdf652c16a0283b42ee13b7700481d0cfe2f291",
-    "upstream_head_date": "2026-05-07",
-    "commit_gap": 458,
-    "recommendation": "plan sync"
-  },
-  "torchtitan_status": {
-    "pin_commit": "5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021",
-    "pin_date": "2025-10-15",
-    "upstream_repo": "https://github.com/pytorch/torchtitan",
-    "upstream_ref": "main",
-    "upstream_head": "4ebb0895a18859ad0f0642901dba0cc888d82b97",
-    "upstream_head_date": "2026-05-07",
-    "commit_gap": 610,
-    "recommendation": "urgent sync"
-  },
-  "primus_turbo_status": {
-    "drift_type": "monthly_in_repo",
-    "month_start_ts": "2026-04-01T00:00:00+08:00",
-    "current_ci_commit": "ef5b58ea3de0a2956d57dba518be466b7a092442",
-    "month_start_ci_commit": "333b68d7c81b722b21b4aad10cd250c45f15027c",
-    "current_benchmark_commit": "ef5b58ea3de0a2956d57dba518be466b7a092442",
-    "month_start_benchmark_commit": "a4488f6cdb15cfff4383c61af7922bb50803f0ea",
-    "current_aiter_commit": "857f4d15775a29af153a2c68a2f8e8a8d696c986",
-    "month_start_aiter_commit": "e83f9903c07001a0ec29e85d223f6e6cdbe00859",
-    "current_triton_commit": "88b227e23f0445f3f695bad05bbf1a363b4f50e0",
-    "month_start_triton_commit": null,
-    "drift_commits": 32,
-    "benchmark_drift_commits": 37,
-    "recommendation": "monitor (new pin landed)"
-  },
-  "recommendations": {
-    "megatron": "plan sync",
-    "torchtitan": "urgent sync",
-    "primus_turbo": "monitor (new pin landed)"
-  },
-  "key_findings": [
-    "3 PRs merged to main in the 2026-W19 window (Bug Fix: 2, Turbo/Dependency Version Update: 1).",
-    "Backend dependency pin change: PR #694 bumps PRIMUS_TURBO_COMMIT in both ci.yaml (333b68d7 → ef5b58ea, +32) and benchmark.yaml (a4488f6c → ef5b58ea, +37) to a single shared SHA, bumps PRIMUS_TURBO_AITER_COMMIT (e83f9903 → 857f4d15), and introduces a new TRITON_COMMIT=88b227e23f0445f3f695bad05bbf1a363b4f50e0 build-arg that compiles Triton from triton-lang/triton@release/3.7.x source inside the docker image. PR #694 also adds a Primus-Turbo Qwen3 attention adapter under primus/backends/torchtitan/models/qwen3/ and exports PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 on the torch CI runner.",
-    "PR #675 closes Megatron-side interface debt after the recent Megatron upgrade: routes custom HF tokenizer types through Megatron's official build_tokenizer with a unique_identifiers fallback; normalizes data-path args for str/list/tuple to upstream-compatible semantics; realigns legacy core_gpt_dataset_config_from_args with upstream pretrain_gpt via get_blend_and_blend_per_split; wires missing upgraded fields (object_storage_cache_path, per_dataset_sequences_path, dataloader_fast_cache_load) to prevent silent mock-dataset fallback.",
-    "PR #668 replaces the substring `in` check in the preflight IP-address probe with a precise regex match in primus/tools/preflight/network/network_probe.py, eliminating false-positives where a target IP is a substring of another active IP.",
-    "Megatron-LM upstream drift: pin d3528a21 (2026-03-06) is now 458 commits behind upstream HEAD 7cdf652c (2026-05-07), +45 since W18. Highlights since W18: legacy transformer/modules removed (#4207), legacy GPT code removed (#4322), inference cache of input + position ID views (#4634), vLLM grouped-gemm MoE inference backend (#4566), checkpoint conversion between GPT_model and Hybrid_model (#4482), nvidia_resiliency_ext fault injection (#4370), gradient corruption fix for layerwise param AG overlap (#4609). Recommendation: plan sync (unchanged).",
-    "torchtitan upstream drift: pin 5fb7cc2e (2025-10-15) is now 610 commits behind upstream HEAD 4ebb0895 (2026-05-07), +44 since W18. Highlights since W18: GraphTrainer memory-policy registry + extra graph-passes hook (#3215), regional Inductor RMSNorm fusion (#3132), bucketing pass enabled in precompile path (re-land #3213), CPU-offloading prefetch pass (#3166), `[graph_trainer] FSDP AG RS overlap` (#3156), MXFP8GroupedExpertsConverter swap fix (#3199), float8 GroupedExperts quantization (#3233), AllToAllTokenDispatcher sp_size padding (#3193), 50s→15s vLLM compile-time win (#3145). Recommendation: urgent sync (unchanged).",
-    "Primus-Turbo monthly drift: pin advanced this week. CI pin moved 32 commits ahead of month start; benchmark pin moved 37 commits ahead of month start; the two pins are now in sync at ef5b58ea (Primus-Turbo #320, default grouped-gemm backend → Triton for BF16/FP8). AITER pin moved (e83f9903 → 857f4d15). New TRITON_COMMIT introduced. Recommendation: monitor (new pin landed).",
-    "Backend-gap reports were not regenerated this week: no submodule SHA changed (Megatron-LM, torchtitan, Megatron-Bridge, Emerging-Optimizers, HummingbirdXT, maxtext are all unchanged on origin/main); the Primus-Turbo pin bump is a tracked-config backend-version change but Primus-Turbo is not currently surfaced as a separate backend under docs/backend-gap/, so the existing torchtitan baseline report (5fb7cc2e vs upstream main) is unaffected."
-  ]
-}
diff --git a/docs_deprecated/README.md b/docs_deprecated/README.md
new file mode 100644
index 000000000..d2659d274
--- /dev/null
+++ b/docs_deprecated/README.md
@@ -0,0 +1,82 @@
+# Primus Documentation
+
+Welcome to the Primus documentation! This guide will help you get started with training large-scale foundation models on AMD GPUs.
+
+> **Comprehensive Documentation**: For the complete production documentation set, see [`docs/`](../docs/README.md). It includes configuration references, parallelism guides, environment variable documentation, and more.
+
+## Documentation Structure
+
+### Getting Started
+
+Start here if you're new to Primus:
+
+- **[Quick Start Guide](./quickstart.md)** - Get up and running in 5 minutes
+- **[CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)** - Complete command-line reference
+- **[CLI Architecture](../docs/06-developer-guide/cli-architecture.md)** - Design philosophy and deep dive
+
+### User Guides
+
+Guides for common workflows and features:
+
+- **[Configuration System](../docs/02-user-guide/configuration-system.md)** - YAML configuration, presets, overrides, and inheritance
+- **[Deployment Guide](../docs/05-operations/deployment.md)** - Container, Slurm, and Kubernetes deployment
+
+### Technical References
+
+In-depth technical documentation:
+
+- **[Post-Training Guide](./posttraining.md)** - Fine-tuning with SFT and LoRA using Primus CLI
+- **[Native SFT & LoRA Quick Start](../docs/04-technical-guides/native-sft-lora.md)** - Megatron-native SFT/LoRA launch guide (BF16/FP8/FP4), no Megatron-Bridge runtime dependency
+- **[Performance Projection](./projection.md)** - Project training performance and memory to multi-node configurations
+- **[Tuning Agent](../docs/02-user-guide/tuning-agent.md)** - LLM-driven search for an optimal training config — parallelism plus batching, schedule, memory, MoE-comm, and precision knobs (drives the projection tool as an oracle)
+- **[Preflight](./preflight.md)** - Cluster diagnostics (host/GPU/network info + perf tests)
+- **[Benchmark Suite](./benchmark.md)** - GEMM, RCCL, end-to-end benchmarks and profiling
+- **[Supported Models](./backends/overview.md#supported-models)** - Supported LLM architectures and feature compatibility matrix
+- **[Backend Patch Notes](./backends/overview.md)** - Primus-specific arguments for Megatron, TorchTitan, etc.
+- **[Backend Extension Guide](./backends/extending-backends.md)** - How to add a new backend using the current adapter/trainer architecture
+ - **[Megatron Model Extension Guide](./backends/adding-megatron-models.md)** - How to add a new Megatron model config
+ - **[TorchTitan Model Extension Guide](./backends/adding-torchtitan-models.md)** - How to add a new TorchTitan model config
+- **[Flux Diffusion Models](../docs/04-technical-guides/diffusion-models/README.md)** - Flux diffusion model architecture, training, and API reference
+- **[FP8 Training Guide](../docs/04-technical-guides/diffusion-models/fp8_training.md)** - FP8 precision training on AMD MI300X/MI355X: configuration, benchmarks, and tuning
+
+### Production Documentation
+
+For comprehensive coverage, see the [Production Documentation](../docs/README.md):
+
+- **[Configuration References](../docs/03-configuration-reference/megatron-parameters.md)** - Per-backend YAML parameter documentation
+- **[Environment Variables](../docs/03-configuration-reference/environment-variables.md)** - Complete environment variable reference
+- **[Parallelism Strategies](../docs/04-technical-guides/parallelism-strategies.md)** - Distributed training parallelism explained
+- **[Performance Tuning](../docs/04-technical-guides/performance-tuning.md)** - HipBLASLt, Primus-Turbo, FP8, MoE optimization
+- **[Troubleshooting](../docs/05-operations/troubleshooting.md)** - Common issues and solutions
+- **[Architecture](../docs/06-developer-guide/architecture.md)** - System design and code architecture
+
+### Help and Support
+
+- **[Troubleshooting Guide](../docs/05-operations/troubleshooting.md)** - Common issues and solutions
+- **[Examples](../examples/README.md)** - Real-world training examples and templates
+- **[Preflight Tool](../primus/tools/preflight/README.md)** - Cluster sanity checker to verify environment readiness
+
+## Quick Navigation by Use Case
+
+### I want to...
+
+- **Train a model locally** → [Quick Start](./quickstart.md) + [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)
+- **Run distributed training on Slurm** → [Deployment Guide](../docs/05-operations/deployment.md)
+- **Configure my training run** → [Configuration System](../docs/02-user-guide/configuration-system.md)
+- **Look up YAML parameters** → [Configuration References](../docs/03-configuration-reference/megatron-parameters.md)
+- **Project performance to multi-node** → [Performance Projection](./projection.md)
+- **Auto-tune my training config (parallelism + knobs)** → [Tuning Agent](../docs/02-user-guide/tuning-agent.md)
+- **Benchmark performance** → [Benchmark Suite](./benchmark.md)
+- **Understand the CLI design** → [CLI Architecture](../docs/06-developer-guide/cli-architecture.md)
+- **Troubleshoot issues** → [Troubleshooting](../docs/05-operations/troubleshooting.md)
+
+## External Resources
+
+- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) - High-performance operators and modules
+- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) - Stability and platform layer
+- [AMD ROCm Documentation](https://rocm.docs.amd.com/)
+- [TorchTitan Documentation](https://github.com/pytorch/torchtitan)
+
+---
+
+**Need help?** Check the [FAQ](./faq.md) or open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues).
diff --git a/docs/backends/adding-megatron-models.md b/docs_deprecated/backends/adding-megatron-models.md
similarity index 100%
rename from docs/backends/adding-megatron-models.md
rename to docs_deprecated/backends/adding-megatron-models.md
diff --git a/docs/backends/adding-torchtitan-models.md b/docs_deprecated/backends/adding-torchtitan-models.md
similarity index 100%
rename from docs/backends/adding-torchtitan-models.md
rename to docs_deprecated/backends/adding-torchtitan-models.md
diff --git a/docs/backends/extending-backends.md b/docs_deprecated/backends/extending-backends.md
similarity index 99%
rename from docs/backends/extending-backends.md
rename to docs_deprecated/backends/extending-backends.md
index 129e57a51..e0ffc0bc6 100644
--- a/docs/backends/extending-backends.md
+++ b/docs_deprecated/backends/extending-backends.md
@@ -95,7 +95,7 @@ from typing import Any, Dict
 
 from primus.core.backend.backend_adapter import BackendAdapter
 from primus.core.backend.backend_registry import BackendRegistry
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class DummyAdapter(BackendAdapter):
@@ -154,7 +154,7 @@ Key points:
 from typing import Any
 
 from primus.core.trainer.base_trainer import BaseTrainer
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class DummyPretrainTrainer(BaseTrainer):
diff --git a/docs/backends/maxtext/patch-notes.md b/docs_deprecated/backends/maxtext/patch-notes.md
similarity index 100%
rename from docs/backends/maxtext/patch-notes.md
rename to docs_deprecated/backends/maxtext/patch-notes.md
diff --git a/docs/backends/megatron/patch-notes.md b/docs_deprecated/backends/megatron/patch-notes.md
similarity index 100%
rename from docs/backends/megatron/patch-notes.md
rename to docs_deprecated/backends/megatron/patch-notes.md
diff --git a/docs/backends/overview.md b/docs_deprecated/backends/overview.md
similarity index 100%
rename from docs/backends/overview.md
rename to docs_deprecated/backends/overview.md
diff --git a/docs/backends/torchtitan/patch-notes.md b/docs_deprecated/backends/torchtitan/patch-notes.md
similarity index 100%
rename from docs/backends/torchtitan/patch-notes.md
rename to docs_deprecated/backends/torchtitan/patch-notes.md
diff --git a/docs/benchmark.md b/docs_deprecated/benchmark.md
similarity index 100%
rename from docs/benchmark.md
rename to docs_deprecated/benchmark.md
diff --git a/docs/cli/PRIMUS-CLI-GUIDE.md b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md
similarity index 94%
rename from docs/cli/PRIMUS-CLI-GUIDE.md
rename to docs_deprecated/cli/PRIMUS-CLI-GUIDE.md
index c48915e6f..82046e9ca 100644
--- a/docs/cli/PRIMUS-CLI-GUIDE.md
+++ b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md
@@ -69,8 +69,20 @@ Primus CLI supports three execution modes, each suitable for different scenarios
 
 # Environment check (info only)
 ./primus-cli direct -- preflight --host --gpu --network
+
+# Per-node smoke test (auto-selects `--single` since node_smoke runs one
+# process per node by design; rank 0 also aggregates the per-node JSONs):
+./primus-cli direct -- node_smoke --tier2-perf
+
+# Suppress launcher + tool stdout (--silent goes BEFORE `--`; errors and
+# the launcher log file are preserved; not recommended for normal use):
+./primus-cli direct --silent -- preflight --quick
 ```
 
+**Optional environment variables (direct mode)**:
+- `VENV_ACTIVATE` — Path to a Python virtualenv `bin/activate` script. If set, sourced before launching; if unset, no-op (the container path uses the container's bundled Python and never sets this).
+- `NNODES` / `NODE_RANK` / `MASTER_ADDR` / `MASTER_PORT` / `GPUS_PER_NODE` — Pre-export to override SLURM-derived values. Inside a SLURM allocation, they are auto-derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_NODELIST` when not pre-exported.
+
 **Suitable for**:
 - ✅ Local development and debugging
 - ✅ Single-node training
@@ -166,13 +178,23 @@ Primus CLI supports three execution modes, each suitable for different scenarios
 # Run distributed GEMM benchmark
 ./primus-cli slurm srun -N 2 -- benchmark gemm --M 16384 --N 16384 --K 16384
 
-# Multi-node environment check (info only)
-# this will generate a fast info report of the host, GPU, and network
+# Multi-node environment check (info only). Anything after `--` that isn't
+# the keyword `container` or `direct` is treated as a primus subcommand and
+# routed through the default container entry chain.
 ./primus-cli slurm srun -N 4 -- preflight --host --gpu --network
 
 # this will generate a full preflight report of the host, GPU, and network, as well as the performance tests
 ./primus-cli slurm srun -N 4 -- preflight --report-file-name preflight-report-4N
 
+# Explicit entry-mode keyword: route through primus-cli-direct.sh instead of
+# the container chain. Useful when nodes share a Python venv on a shared FS
+# (see docs/preflight-direct.md for the setup).
+./primus-cli slurm srun -N 4 -- direct -- preflight --quick
+
+# Per-node smoke via the direct entry (node_smoke auto-runs in single mode;
+# rank 0 aggregates after every rank finishes):
+./primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf
+
 # if you are using AINIC in your cluster, use the appropriate configuration file
 # for preflight test, set docker image to rocm/primus:v26.3 in the configuration file
 ./primus-cli --config runner/use_ainic.yaml slurm srun -N 2 -- preflight --report-file-name preflight-report-2N
@@ -761,7 +783,7 @@ Final result:
 | **Entry Script** | primus-cli-direct.sh | primus-cli-container.sh | primus-cli-slurm.sh |
 | **Environment Prep** | Load local GPU env | Start container + mount + devices | Allocate nodes + network config |
 | **Execution Location** | Current host | Inside container | Slurm-allocated nodes |
-| **Final Call** | Direct torchrun execution | Execute direct.sh in container | Each node executes slurm-entry.sh → direct.sh |
+| **Final Call** | Direct torchrun execution (single mode auto-selected for `node_smoke`) | Execute direct.sh in container | Each node executes slurm-entry.sh → container.sh or direct.sh (via `direct` keyword) |
 | **Distributed Support** | Single-node multi-GPU | Single-node multi-GPU | Multi-node multi-GPU |
 | **Use Case** | Dev debugging | Environment isolation | Production training |
 
@@ -982,7 +1004,7 @@ export PRIMUS_LOG_LEVEL=DEBUG
 ## Reference Resources
 
 ### Related Documentation
-- [CLI Architecture](./CLI-ARCHITECTURE.md) - Primus CLI architecture deep dive
+- [CLI Architecture](../../docs/06-developer-guide/cli-architecture.md) - Primus CLI architecture deep dive
 - [Main Documentation](../README.md) - Complete Primus documentation index
 - [.primus.yaml](../../runner/.primus.yaml) - Default configuration example
 
diff --git a/docs/cli/README.md b/docs_deprecated/cli/README.md
similarity index 73%
rename from docs/cli/README.md
rename to docs_deprecated/cli/README.md
index bdb5e618e..a7e322a46 100644
--- a/docs/cli/README.md
+++ b/docs_deprecated/cli/README.md
@@ -13,7 +13,7 @@ The Primus CLI provides a unified command-line interface for training, benchmark
   - Configuration files and options
   - Best practices and troubleshooting
 
-- **[Architecture Deep Dive](./CLI-ARCHITECTURE.md)**
+- **[Architecture Deep Dive](../../docs/06-developer-guide/cli-architecture.md)**
   - Design philosophy and principles
   - Three-layer architecture explained
   - Plugin system and extensibility
@@ -43,6 +43,22 @@ If you're running from the Primus repo root (after `git clone ... && cd Primus`)
 primus-cli direct -- benchmark gemm -M 4096 -N 4096 -K 4096
 ```
 
+### Data Preparation Commands
+
+```bash
+# Prepare a raw WebDataset (smaller, on-the-fly encoding during training)
+primus-cli direct -- data diffusion-raw \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml
+
+# Prepare a pre-encoded dataset from HuggingFace
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml
+
+# Ingest MLPerf Flux1 pre-encoded data (streaming download + conversion)
+primus-cli direct -- data diffusion-ingest \
+  --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
+```
+
 ## 🎯 Three Execution Modes
 
 | Mode | Use Case | Command Example |
@@ -54,7 +70,7 @@ primus-cli direct -- benchmark gemm -M 4096 -N 4096 -K 4096
 ## 📖 Learn More
 
 - For detailed usage instructions, see the [User Guide](./PRIMUS-CLI-GUIDE.md)
-- For architecture and design details, see [Architecture Deep Dive](./CLI-ARCHITECTURE.md)
+- For architecture and design details, see [Architecture Deep Dive](../../docs/06-developer-guide/cli-architecture.md)
 - For the main Primus documentation, see [Primus README](../../README.md)
 
 ## 🔗 Related Documentation
diff --git a/docs/install-on-host.md b/docs_deprecated/install-on-host.md
similarity index 100%
rename from docs/install-on-host.md
rename to docs_deprecated/install-on-host.md
diff --git a/docs_deprecated/node-smoke-test-instruction.md b/docs_deprecated/node-smoke-test-instruction.md
new file mode 100644
index 000000000..639a8dd5b
--- /dev/null
+++ b/docs_deprecated/node-smoke-test-instruction.md
@@ -0,0 +1,313 @@
+# Node-Smoke Test — Quick-Start Instructions
+
+A short get-started guide for the per-node preflight smoke test. For the full design / aggregator section reference / implementation history, see [node-smoke.md](./node-smoke.md).
+
+---
+
+## 1. What it does
+
+A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. It produces a **single PASS/FAIL verdict per node** plus SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`.
+
+Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node will surface as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers.
+
+---
+
+## 2. Prerequisites
+
+
+| Prerequisite                        | How                                                                                                                                                                                                                                                                                                                                                                                              |
+| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| Python venv on a shared filesystem  | Same venv used by `primus-cli direct -- preflight` (see `[preflight-direct.md](./preflight-direct.md)` §2).                                                                                                                                                                                                                                                                                      |
+| `VENV_ACTIVATE` exported            | `export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate` (optional inside the container path).                                                                                                                                                                                                                                                                                                 |
+| Inside an existing SLURM allocation | One task per node. Recommended: `runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 -- direct -- node_smoke ...`. Equivalent bare form: `srun ... --ntasks-per-node=1 runner/primus-cli direct -- node_smoke ...`. Either way the `direct -- node_smoke` path auto-selects `--single`, so each task spawns one Python process and per-GPU subprocesses are launched internally. |
+
+
+No `MASTER_ADDR`, no `MASTER_PORT`, no global rendezvous required.
+
+---
+
+## 3. Quick start
+
+**Git clone the Primus repository to a shared filesystem that all nodes can read.**
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git
+cd Primus
+git checkout dev/preflight-direct-test
+```
+
+**Note: remember to setup the Python virtual environment and NCCL / fabric environment variables as described in [§2 Prerequisites](#2-prerequisites).**
+
+> ⚠ **Set the NCCL / RCCL environment first** if you plan to run with `--tier2-perf` (the local 8-GPU RCCL all-reduce). Even though the smoke test never opens a cross-node rendezvous, the Tier 2 RCCL step calls `dist.init_process_group(backend="nccl", ...)`, and RCCL **enumerates every transport at init** (XGMI / PCIe P2P + IB + sockets). A misconfigured `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME` / `NCCL_IB_GID_INDEX` can stall init or make the all-reduce silently fall back to a slow path. The launcher's `base_env.sh` auto-detects these via `get_nccl_ib_hca.sh` + `get_ip_interface.sh`, **but auto-detect sometimes picks the wrong values inside a container** (devices masked by the network namespace, frontend NICs picked up instead of fabric NICs, etc.) so you usually want to check these settings and set them explicitly if auto-detection is wrong.
+>
+> Minimum-viable checklist before running with `--tier2-perf`:
+>
+> ```bash
+> # Pin the RDMA / RoCE training NICs the container can actually see.
+> # On a bare-metal host the auto-detect in base_env.sh usually picks
+> # the right set; inside a container or on a multi-role node, list
+> # them explicitly. Use the same set you would pass to a training job.
+> export NCCL_IB_HCA="rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7"
+>
+> # Pick the RoCE v2 GID index for your fabric:
+> #   - Mellanox / Broadcom: typically 3 (base_env.sh default).
+> #   - Pensando Pollara (AINIC): 1.
+> export NCCL_IB_GID_INDEX=3
+>
+> # The bootstrap socket interface. Auto-detect prefers the first
+> # non-loopback interface from `hostname -I`; override when that
+> # picks a frontend NIC instead of the data-plane interface.
+> export NCCL_SOCKET_IFNAME=eno0
+> export GLOO_SOCKET_IFNAME=eno0
+> ```
+>
+> See `[preflight-direct.md` § 4 Cluster-specific NCCL configuration](./preflight-direct.md#4-cluster-specific-nccl-configuration) for the canonical Broadcom / Pensando Pollara values (the same `NCCL_*` set is used by both tools). If you skip `--tier2-perf`, the RCCL step is not executed and none of the above applies — Tier 1 (host limits, RDMA roll-call, leaked-process detection, etc.) does not depend on RCCL.
+>
+> Quick verification: `runner/primus-cli direct --dry-run -- node_smoke --tier2-perf` prints the resolved `NCCL_*` block under "Environment Variables" so you can confirm the values before launching for real.
+
+Recommended — through the `primus-cli slurm srun` wrapper (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK`, applies `slurm.*` config defaults):
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Basic Tier 1 check (~5 s/GPU, ~30 s total)
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke
+
+# Tier 1 + Tier 2 perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL)
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf
+
+# Then re-run training, excluding any node the smoke test failed:
+srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job
+```
+
+Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke
+
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf
+```
+
+Single-node sanity check (no SLURM):
+
+```bash
+runner/primus-cli direct -- node_smoke
+```
+
+> **Both forms produce the same workload.** The wrapper form is recommended because it resolves the distributed env once on the launching node and propagates it via `--env`, and applies any `slurm.`* config defaults (partition / time / etc.). See `[preflight-direct.md` § Wrapper vs. bare-srun](./preflight-direct.md#wrapper-vs-bare-srun) for the precedence table.
+
+---
+
+## 4. More examples (by configuration knob)
+
+> **Convention used below.** The examples in this section are written with bare `srun` for brevity. Anywhere you see `srun <flags> runner/primus-cli direct -- node_smoke ...`, the equivalent wrapper form is `runner/primus-cli slurm srun <flags> -- direct -- node_smoke ...`. Pick whichever matches your habits; both target the same launcher.
+
+### 4.1 Hard-fail on partial NIC enumeration
+
+Catches "7 of 8 RDMA NICs visible" — common cause of crashes after RoCE init.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf --expected-rdma-nics 8
+```
+
+### 4.2 Tighten Tier 2 perf thresholds
+
+Reject GPUs that come in below your acceptance bar. Defaults: GEMM 600 TFLOPS, HBM 2000 GB/s, local RCCL 100 GB/s.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf \
+        --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180
+```
+
+### 4.3 Tighten host limits
+
+Fail nodes whose `RLIMIT_MEMLOCK` or `/dev/shm` is too small for production training.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke \
+        --ulimit-l-min-gb 64 --shm-min-gb 16
+```
+
+### 4.4 Custom dump path
+
+Keep one report per smoke run instead of overwriting the default location.
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf \
+        --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+### 4.5 Allow / extend the foreign-process whitelist
+
+By default, leaked / foreign processes holding a GPU FAIL the node (most common cause of "training fails to launch on a healthy-looking node"). Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`.
+
+```bash
+# Add a site-specific monitoring agent to the whitelist
+srun ... runner/primus-cli direct -- node_smoke \
+    --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,my-monitor
+
+# Don't fail at all on foreign processes (still reported in the markdown)
+srun ... runner/primus-cli direct -- node_smoke --allow-foreign-procs
+```
+
+#### Containers: `name='N/A'` false positives → use `--allow-foreign-procs`
+
+> ⚠ **Running node_smoke inside a container almost always trips this check.** `amd-smi process --json` reports `name="N/A"` for kernel/system PIDs like `gpuagent` whose `/proc/<host_pid>/comm` it cannot read, and the fallback `_resolve_proc_name(pid)` inside `node_smoke` then also fails because the container's `/proc` typically does not expose host PIDs (private PID namespace without `--pid=host`, or a `hidepid=2` mount). The unresolved name doesn't match the allowlist (`gpuagent,rocm-smi-daemon,...`), so the check fires and the node FAILs — even though the only "foreign" processes are well-known system daemons holding zero HBM.
+>
+> **In the container path, pass `--allow-foreign-procs`:**
+>
+> ```bash
+> srun ... runner/primus-cli direct -- node_smoke --tier2-perf --allow-foreign-procs
+> ```
+>
+> The processes are still listed in `smoke_report.md` under "Busy GPUs / leaked processes" so a real leak is still visible; only the FAIL verdict is downgraded.
+>
+> **Narrower alternative** if you want the check to still catch leaks with resolvable names (e.g. a leftover `python` rank), add the literal sentinel `N/A` to the allowlist:
+>
+> ```bash
+> srun ... runner/primus-cli direct -- node_smoke --tier2-perf \
+>     --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,N/A
+> ```
+>
+> The annotator runs `_resolve_proc_name` first, so whenever a real name *can* be resolved (on the host, or after fixing `/proc` visibility) it overrides "N/A" and the normal allowlist applies. The `N/A` entry only matches PIDs whose name genuinely could not be recovered — strictly narrower than `--allow-foreign-procs`.
+>
+> **Root-cause fix** (preferred long-term): grant the container access to host PIDs so `_resolve_proc_name` works and the report shows real names (`gpuagent`, etc.) instead of `N/A`. Typical fixes:
+>
+> - Launch with `--pid=host` (Docker / Podman) so host PIDs are directly addressable.
+> - Mount `/proc` without `hidepid=2`.
+> - Loosen `ptrace_scope` or grant `CAP_SYS_PTRACE`.
+>
+> Once any of those is in place, `_resolve_proc_name` finds the names, the default allowlist matches them, and you no longer need `--allow-foreign-procs`.
+
+### 4.6 Require specific tools
+
+Make missing CLI tools a hard FAIL (default: warn-only).
+
+```bash
+srun ... runner/primus-cli direct -- node_smoke --require-tools amd-smi,rocm-smi,lsof
+```
+
+### 4.7 Skip dmesg scan (containers with no privileges)
+
+```bash
+srun ... runner/primus-cli direct -- node_smoke --skip-dmesg
+```
+
+### 4.8 Re-aggregate from existing per-node JSONs (no re-run)
+
+Useful when you only want to refresh the markdown report, or when you've collected JSONs separately.
+
+```bash
+# From any node, no allocation needed if you're just reading local files.
+# The primus-cli wrapper always runs both phases, so use the standalone
+# aggregate subcommand for "aggregate only" -- it reads the existing
+# <dump>/smoke/*.json without re-running the per-node smoke step.
+python -m primus.tools.preflight.node_smoke aggregate \
+    --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5
+```
+
+### 4.9 Silent mode (for CI)
+
+Suppresses wrapper stdout, but the **final report path is still printed** and stderr / exit code are preserved.
+
+```bash
+srun ... runner/primus-cli direct --silent -- node_smoke --tier2-perf
+```
+
+### 4.10 Combined "production-ready screen"
+
+A representative one-shot for a production cluster screen:
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct --silent -- node_smoke --tier2-perf \
+        --expected-rdma-nics 8 \
+        --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 \
+        --ulimit-l-min-gb 64 --shm-min-gb 16 \
+        --require-tools amd-smi,rocm-smi,lsof \
+        --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+---
+
+## 5. Outputs
+
+All written under `--dump-path` (default `output/preflight/`).
+
+
+| File                      | Purpose                                                                                                              |
+| ------------------------- | -------------------------------------------------------------------------------------------------------------------- |
+| `smoke/<short-host>.json` | Per-node verdict + every collected metric. One file per node.                                                        |
+| `smoke_report.md`         | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail).                     |
+| `passing_nodes.txt`       | Newline-separated short hostnames. Pipe into `srun --nodelist=`.                                                     |
+| `failing_nodes.txt`       | Newline-separated short hostnames. Pipe into `srun --exclude=`.                                                      |
+| `expected_nodes.txt`      | Auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`. Lets the report name nodes that never reported. |
+
+
+Read the cluster verdict at a glance:
+
+```bash
+head -10 output/preflight/smoke_report.md
+```
+
+Feed bad nodes into a re-run:
+
+```bash
+srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job
+```
+
+---
+
+## 6. Common knobs (cheat sheet)
+
+
+| Flag                         | Default                                          | When you'd change it                                                                  |
+| ---------------------------- | ------------------------------------------------ | ------------------------------------------------------------------------------------- |
+| `--tier2-perf`               | off                                              | Always on for production screens — adds GEMM TFLOPS, HBM GB/s, local RCCL all-reduce. |
+| `--gemm-tflops-min N`        | 600                                              | Site-specific acceptance bar.                                                         |
+| `--hbm-gbs-min N`            | 2000                                             | Site-specific acceptance bar (MI300X healthy ≈ 4500–5000).                            |
+| `--rccl-gbs-min N`           | 100                                              | Site-specific acceptance bar.                                                         |
+| `--expected-rdma-nics N`     | unset                                            | Hard-fail on partial NIC enumeration.                                                 |
+| `--ulimit-l-min-gb GB`       | 32                                               | Raise for production training profiles.                                               |
+| `--shm-min-gb GB`            | 8                                                | Raise for large-batch / many-rank profiles.                                           |
+| `--allow-foreign-procs`      | off                                              | Co-tenant clusters or shared GPUs.                                                    |
+| `--allowed-procs LIST`       | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Add site-specific monitoring agents.                                                  |
+| `--require-tools LIST`       | `""`                                             | Fail-fast if a CLI tool is missing in PATH.                                           |
+| `--skip-dmesg`               | off                                              | Inside unprivileged containers.                                                       |
+| `--dump-path DIR`            | `output/preflight`                               | Archive each run separately.                                                          |
+| `--silent` (wrapper)         | off                                              | CI / scripted runs.                                                                   |
+| `--aggregate-only` (wrapper) | off                                              | Re-render report without re-running per-node checks.                                  |
+
+
+For the full flag list and the aggregator subcommand, see `python -m primus.tools.preflight.node_smoke run --help` and `... aggregate --help`, or `[node-smoke.md](./node-smoke.md)` §"Configuration knobs".
+
+---
+
+## 7. Troubleshooting
+
+
+| Symptom                                                                       | Likely cause / fix                                                                                                                                                                            |
+| ----------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...`          | Fix the path (`export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate`), or `unset VENV_ACTIVATE` to fall back to system / container Python.                                                |
+| Every node FAILs with `gpu_processes: ... name='N/A'`                         | Should no longer happen after the `/proc/<pid>/comm` fallback fix. If it does, check that `/proc/<pid>/comm` is readable on the node (`hidepid` mount?). Workaround: `--allow-foreign-procs`. |
+| Some nodes never produce a JSON                                               | Aggregator names them in `failing_nodes.txt` via `expected_nodes.txt`. If `scontrol` was unavailable, they'll appear as `<missing-N>`.                                                        |
+| Tier 2 perf numbers below threshold on a known-good node                      | Almost always insufficient CPU cores on `srun` — pass `-c <cores-per-node>` so RCCL proxy threads have CPU.                                                                                   |
+| Re-run on a smaller nodelist still shows the previously removed nodes as PASS | Default behavior cleans stale JSONs on rank 0. If you passed `--no-clean-dump-path`, either remove it or `rm -rf output/preflight` between runs.                                              |
+
+
+---
+
+## 8. See also
+
+- `[node-smoke.md](./node-smoke.md)` — full design, aggregator sections, configuration reference, implementation history.
+- `[preflight-direct.md](./preflight-direct.md)` — the heavier `preflight` tool with global rendezvous and inter-node bandwidth tests.
+- `[primus/cli/subcommands/node_smoke.py](../primus/cli/subcommands/node_smoke.py)` — the primus-cli subcommand wiring (two-phase dispatch: rank-N run + rank-0 aggregate).
+- `[primus/tools/preflight/node_smoke/cli.py](../primus/tools/preflight/node_smoke/cli.py)` — canonical flag definitions and per-node / aggregate phase bodies.
diff --git a/docs_deprecated/node-smoke.md b/docs_deprecated/node-smoke.md
new file mode 100644
index 000000000..98471ea80
--- /dev/null
+++ b/docs_deprecated/node-smoke.md
@@ -0,0 +1,429 @@
+# Node-Local Smoke Test
+
+> **Just want to run it?** See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) for the short quick-start guide. This document is the full reference (architecture, every report section, every flag, design history).
+
+A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. Designed to **quickly identify broken nodes before a large training job commits to a global rendezvous**. Because training jobs allocate whole nodes, a under-performing GPU (or NIC, or wedged driver) takes the entire node out of rotation -- so the smoke test produces a single PASS/FAIL verdict per node and SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`.
+
+- **Implementation**: `primus/tools/preflight/node_smoke/` (Python sub-package; entry point `python -m primus.tools.preflight.node_smoke`).
+- **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK` via `--env`, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` form (bare `srun` + direct) is equivalent and useful for ad-hoc runs.
+- **Companion**: see `docs/preflight.md` for the full preflight tool (with global rendezvous and richer perf tests).
+
+## Quick start
+
+Recommended — through the `primus-cli slurm srun` wrapper:
+
+```bash
+# Inside an existing SLURM allocation (the normal case):
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke
+
+# With perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL all-reduce):
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf
+
+# Hard-fail on partial NIC enumeration (e.g. 7 of 8 RDMA NICs).
+# The count is compared against the *training-NIC* set after the
+# selector chain runs, so frontend / storage RoCE NICs do not count.
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf --expected-rdma-nics 8
+
+# Explicitly pin the training-NIC selector (otherwise NCCL_IB_HCA env
+# is used; otherwise admin-disabled phys_state ports are auto-excluded):
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf \
+    --rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1,rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1'
+```
+
+Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke
+
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf
+```
+
+Single-node local check (no SLURM, both forms collapse to the same call):
+
+```bash
+runner/primus-cli direct -- node_smoke
+```
+
+> **Note on the `direct` keyword**: with the `primus-cli slurm srun` wrapper, the entry-mode keyword `direct` between the two `--`s is mandatory to take the direct (no-container) path. Without it the wrapper routes through the **container** path. See [`preflight-direct.md` § Wrapper vs. bare-srun](./preflight-direct.md#wrapper-vs-bare-srun) for the full precedence and caveats.
+
+When `VENV_ACTIVATE` is set, `primus-cli direct` sources it before launching `node_smoke` (same convention as `primus-cli direct -- preflight`). When unset (e.g. inside the container path), it is a no-op.
+
+## Outputs
+
+After a run, `<dump>/` (default `output/preflight/`) contains:
+
+| File | Purpose |
+|---|---|
+| `smoke/<short-host>.json` | Per-node verdict + every collected metric. One file per node. |
+| `smoke_report.md` | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail). |
+| `passing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --nodelist=`. |
+| `failing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --exclude=`. |
+| `expected_nodes.txt` | (auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`) Used by the aggregator to name nodes that never reported. |
+
+```bash
+# Re-run training, excluding the bad nodes from the previous smoke:
+srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job
+```
+
+## Architecture
+
+- **Per-node Python entry** (`node_smoke.py run`) — runs independently on every node. No `MASTER_ADDR`, no global `torch.distributed` rendezvous; a stuck node cannot wedge its peers.
+- **Per-GPU isolation** — each GPU's checks run in their own Python subprocess with a hard timeout. A stuck `torch.cuda.set_device()` (which can't be aborted by `signal.alarm` because it sits inside a non-interruptible driver syscall) is `SIGKILL`'d from the parent without affecting the rest of the node's checks.
+- **Local-only RCCL** — Tier 2 all-reduce uses `torch.multiprocessing.spawn` over `tcp://127.0.0.1`. No cross-node communication.
+- **Aggregator on `NODE_RANK==0`** — polls `<dump>/smoke/` for the expected number of JSONs (with a timeout), computes drift across the cluster, writes the markdown report and pass/fail txt files. Returns non-zero if any node FAILs or never reports.
+
+## Module layout — where each check lives
+
+The implementation is a Python sub-package under `primus/tools/preflight/node_smoke/`, split so each Tier 1 sub-section, the per-GPU subprocess body, the orchestrator, and the aggregator each live in their own file. The single public entry point is `main` (re-exported from `__init__`); `python -m primus.tools.preflight.node_smoke ...` resolves to `__main__.py` which calls it.
+
+```
+primus/tools/preflight/node_smoke/
+├── __init__.py          # re-export `main`
+├── __main__.py          # `python -m primus.tools.preflight.node_smoke`
+├── cli.py               # `_build_parser`, `_cmd_run`, `_cmd_aggregate`,
+│                        # `_cmd_per_gpu`, `main`
+├── types.py             # `GPUResult`, `NodeResult` dataclasses
+├── logging_utils.py     # `_ts`, `_log`, `_warn`, hostname normalisation
+├── shell_utils.py       # `_which`, `_read_text`, `_resolve_gpu_bdf`,
+│                        # `_systemctl_is_active`, `_parse_size_with_unit`,
+│                        # `_findings_to_dicts`
+├── per_gpu.py           # `_per_gpu_body` (Tier 1 + optional Tier 2 perf,
+│                        # GEMM/HBM bandwidth measurement)
+├── rccl_local.py        # node-local RCCL all-reduce (Tier 2)
+├── orchestrator.py      # `_spawn_per_gpu`, `_node_status_from`,
+│                        # `_clean_dump_path`
+├── collectors/          # one module per Tier 1 sub-section
+│   ├── dmesg.py             # recent dmesg error scan
+│   ├── fingerprint.py       # Tier 1 A — software-stack fingerprint
+│   ├── nics.py              # Tier 1 B — NIC / RDMA roll-call
+│   ├── host_limits.py       # Tier 1 C — ulimit / shm / NUMA / governor
+│   ├── gpu_low_level.py     # Tier 1 D-1 — amd-smi metric (ECC, throttle,
+│   │                        #               clocks, power)
+│   ├── xgmi.py              # Tier 1 D-2 — XGMI link matrix
+│   ├── clock.py             # Tier 1 E — wall time + time-daemon health
+│   ├── rocm_smi.py          # Tier 1 F + cross-tool fallbacks for D-1/2/G
+│   ├── gpu_processes.py     # Tier 1 G — foreign / leaked PID detection
+│   ├── tooling.py           # tooling availability inventory
+│   └── reused_info.py       # reused gpu/host/network info collectors
+└── aggregator/
+    ├── summarizers.py       # `_*_rows` / `_*_summary` data shapers
+    └── report.py            # `write_smoke_report` + one `_write_<section>`
+                             # helper per Markdown `##` section
+```
+
+Dependency graph (acyclic; arrows mean "imports"):
+
+```mermaid
+flowchart TD
+    cli[cli.py] --> orchestrator[orchestrator.py]
+    cli --> per_gpu[per_gpu.py]
+    cli --> rccl_local[rccl_local.py]
+    cli --> aggReport["aggregator/report.py"]
+    cli --> collectorsAll["collectors/*"]
+    cli --> logging[logging_utils.py]
+    cli --> types[types.py]
+
+    aggReport --> aggSum["aggregator/summarizers.py"]
+    aggReport --> logging
+    aggSum --> tooling["collectors/tooling.py"]
+
+    orchestrator --> types
+
+    per_gpu --> shell[shell_utils.py]
+
+    collectorsAll --> shell
+    collectorsAll --> rocmsmi["collectors/rocm_smi.py"]
+    rocmsmi --> shell
+    logging -.->|stdlib only| std[(socket / sys / time)]
+    shell -.->|stdlib only| std
+    types -.->|stdlib only| std
+```
+
+* `collectors/` are leaf modules (depend on `shell_utils` + sometimes `collectors/rocm_smi`); they never import `cli` / `orchestrator` / `per_gpu`.
+* `per_gpu` is the body of the `_per_gpu` subprocess and is the ONLY module loaded inside that subprocess via `python -m primus.tools.preflight.node_smoke _per_gpu N` — so its dependency surface is intentionally narrow (only `shell_utils`).
+* `aggregator/` only depends on its own `summarizers` plus `logging_utils` (and `collectors/tooling` for the static `_TRACKED_TOOLS` constant).
+
+## What's checked
+
+### Tier 1 — mandatory (~5 s / GPU, always runs)
+
+**Per-GPU subprocess (with hard timeout):**
+- `torch.cuda.set_device(i)` — proves the device is bindable (a stale GPU often fails here)
+- 256 MB allocation
+- Tiny GEMM (2048² bf16) with `isfinite()` check on the result
+
+**Reused from existing preflight collectors** (no rendezvous needed):
+- `collect_gpu_info` — `level='fail'` Findings cause node FAIL
+- `collect_host_info` — same
+- `collect_network_info(expect_distributed=False)` — same
+
+**dmesg recent-error scan** — greps the last `--dmesg-minutes` (default 15) of `dmesg` for known patterns (`xid`, `gpu reset`, `hung_task`, `mce:`, `amdgpu.*error`, ...). Matches are surfaced in the report.
+
+**A. Software-stack fingerprint** (`tier1.fingerprint`):
+- Kernel, OS, Python
+- ROCm version (`/opt/rocm/.info/version`)
+- amdgpu kernel-module version (`/sys/module/amdgpu/version`)
+- PyTorch version, `torch.version.hip`, RCCL version (via `torch.cuda.nccl.version()`), librccl path
+- Per-IB-device firmware (`/sys/class/infiniband/<dev>/fw_ver`) and HCA model
+
+**B. NIC / RDMA roll-call** (`tier1.nics`):
+- Per port (read entirely from `/sys/class/infiniband` — no `ibv_devinfo`/`ibstat` dependency, works inside containers): `state`, `phys_state`, `rate`, netdev + MTU, total non-zero GIDs, RoCE v2 GID count
+- **Training-NIC selector** — many clusters expose more RDMA-capable ports than the training job uses (frontend / management / storage NICs). The hard-fail rules only run against the *included* subset. Precedence:
+  1. `--rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,...'` (full `NCCL_IB_HCA` syntax: comma-separated `device[:port]`, `^...` for denylist, `=dev` for exact-match, no `:port` to match any port on the device).
+  2. `NCCL_IB_HCA` env (same syntax) — mirrors what NCCL/RCCL itself will use, so the smoke test and the training launch agree by construction.
+  3. Heuristic: auto-exclude any port whose `phys_state` is `Disabled` or `Sleep` (admin-disabled at firmware/driver level — no SFP, BIOS port-disable, netdev admin-down). Real failure modes on a port that *is* meant to be used produce a different `phys_state` (`Polling`, `LinkErrorRecovery`, or `LinkUp` with `state!=ACTIVE`), so this heuristic does not mask cable / driver problems.
+  4. Fallback: every IB port must be ACTIVE / LinkUp.
+- Excluded ports stay visible in `tier1.nics.ports` and are summarised in `tier1.nics.excluded_ports` + `info_issues` for diagnostics. They do NOT contribute to the node FAIL signal.
+- **Hard fail rules** (only on the included set): port not `ACTIVE` / not `LinkUp`, active port with zero RoCE v2 GIDs (RoCE) or zero valid GIDs (IB), included-NIC count ≠ `--expected-rdma-nics N` (when set).
+- **Empty-set guard**: if every discovered port gets excluded, the node still hard-fails — a node with zero training NICs cannot participate in inter-node training.
+
+**C. Host limits / system tunables** (`tier1.host_limits`):
+- `RLIMIT_MEMLOCK`, `RLIMIT_NOFILE`, `RLIMIT_NPROC`
+- `/dev/shm` size + free
+- NUMA node count, CPU count, `cpu0` scaling governor
+- **Hard fail rules**: `RLIMIT_MEMLOCK` finite and below `--ulimit-l-min-gb` (default 32 GiB) → "RDMA pin will fail under load"; `/dev/shm` size below `--shm-min-gb` (default 8 GiB) → "NCCL shared-mem may fail"
+
+### Tier 2 — optional perf sanity (`--tier2-perf`)
+
+Per-GPU steady-state metrics, with iteration counts aligned to the preflight `--quick` preset (`warmup=5, iters=20` for GEMM/RCCL; `warmup=10, iters=20` for HBM) so smoke and preflight numbers are directly comparable.
+
+- **GEMM TFLOPS** — 8192³ bf16 `torch.matmul`, threshold `--gemm-tflops-min` (default 600).
+- **HBM GB/s** — 512 MB device-to-device `torch.Tensor.copy_` (counts read + write), threshold `--hbm-gbs-min` (default 2000).
+- **Local 8-GPU RCCL all-reduce GB/s** — algorithmic bandwidth `2·S·(P-1)/P / t / 1e9` at 64 MB, threshold `--rccl-gbs-min` (default 100).
+
+## Aggregator report sections
+
+Every section short-circuits to a placeholder (`*All nodes match.*` / `*No NIC issues.*` / `*No host-limit issues.*`) on a healthy cluster, so the report stays short. Each section header is part of the operator-facing contract — order and wording are stable across releases (some Slack bots / CI scripts grep for them).
+
+In order:
+
+1. **Status table** — one row per node with `node_rank`, hostname, PASS/FAIL, duration, top fail reason.
+2. **Stack drift across cluster** — for every scalar fingerprint key, outliers vs the cluster majority.
+3. **NIC firmware drift across cluster** — per-IB-device firmware drift.
+4. **NIC / RDMA roll-call issues** — every offending node + port (included set only).
+5. **NIC port-count summary** — cluster-majority *training-NIC* count and any node that disagrees (catches partial-NIC degradation without `--expected-rdma-nics`). The count is taken from the included set, so nodes that legitimately have extra frontend / storage RoCE NICs don't show up as anomalies.
+6. **NIC excluded ports (informational)** — ports the selector chain dropped from the training-NIC set, grouped by source (`--rdma-nic-allowlist` / `NCCL_IB_HCA` / heuristic). Informational only; does not contribute to FAIL.
+7. **Host limits issues** — per-node hard-limit violations.
+8. **GPU visibility issues** — nodes where torch couldn't see the GPUs or amd-smi sees more GPUs than torch (stale ROCm / wedged amdgpu driver). Independent of every other collector.
+9. **GPU low-level outliers (PCIe link / HBM)** — per-GPU outliers vs the cluster majority on PCIe width/speed and HBM total.
+10. **XGMI link issues** — any non-XGMI GPU pair (intra-node collectives silently fall back to PCIe).
+11. **Cluster clock + time daemons** — wall-clock spread plus per-node time-daemon health.
+12. **Tooling self-latency (`rocm-smi --version`)** — slow / timed-out tool calls (precursor to a wedged amdgpu driver).
+13. **Tooling availability** — always-on inventory of `amd-smi` / `rocm-smi` / `lsof` per node, plus which Tier 1 checks have NO working tool on each node.
+14. **Busy GPUs / leaked processes** — foreign PIDs holding GPUs at smoke start (most common cause of training failing to launch on an otherwise-healthy node).
+15. **GPU pre-touch HBM usage outliers** — GPUs with non-trivial HBM in use BEFORE smoke touched the device.
+16. **GPU compute-activity outliers** — GPUs with `gfx_activity_pct >= --gpu-activity-warn-pct` at smoke start (warn-only).
+17. **Tier 2 perf summary** (conditional, only when at least one node ran Tier 2) — per-node GEMM TFLOPS / HBM GB/s as `min / median / max`, plus local RCCL GB/s.
+18. **Failing nodes — full reasons** (conditional, only when there are failing nodes) — every fail reason, expanded per node.
+
+Each section that does pure data shaping is wrapped in its own `try / except`, so a future schema bug in one section can't truncate the rest of the report. The two intentional EXCEPTIONS are **Tier 2 perf summary** and **Failing nodes — full reasons** — both deliberately propagate exceptions so a regression in either bubbles up rather than silently rendering a half-empty section.
+
+## Configuration knobs
+
+The authoritative source of flags + defaults is `python -m primus.tools.preflight.node_smoke run --help` (and `... aggregate --help`). The tables below mirror the parser as of the package-split refactor.
+
+### `run` subcommand
+
+| Flag | Default | Purpose |
+|---|---|---|
+| `--dump-path` | `output/preflight` | Output directory. |
+| `--expected-gpus N` | auto | Override GPU count (auto-detected from `LOCAL_WORLD_SIZE` / `GPUS_PER_NODE` / `torch.cuda.device_count()`). |
+| `--per-gpu-timeout-sec` | 15 | Hard timeout per per-GPU subprocess. |
+| `--tier2-perf` | off | Enable Tier 2 perf sanity (per-GPU GEMM TFLOPS + HBM GB/s + node-local RCCL all-reduce). Single switch — you cannot enable just one half. |
+| `--gemm-tflops-min` | 600 | Tier 2 GEMM threshold. |
+| `--hbm-gbs-min` | 2000 | Tier 2 HBM threshold. |
+| `--rccl-size-mb` | 64 | Local RCCL message size. |
+| `--rccl-gbs-min` | 100 | Local RCCL bandwidth threshold. |
+| `--rccl-timeout-sec` | 120 | Hard timeout for the RCCL phase. |
+| `--skip-dmesg` | off | Skip dmesg scan (e.g. inside containers). |
+| `--dmesg-minutes` | 15 | dmesg `--since` window. |
+| `--expected-rdma-nics N` | auto-report-only | When set, a mismatch between the **included (training-NIC) count** and N becomes a node FAIL. Compares against the post-selector count, not the raw number of devices under `/sys/class/infiniband`. |
+| `--rdma-nic-allowlist LIST` | unset | Explicit training-NIC selector in `NCCL_IB_HCA` syntax (`device[:port],...`, `^...` denylist, `=dev` exact-match). Wins over `NCCL_IB_HCA` env. When neither this flag nor the env is set, the collector auto-excludes ports whose `phys_state` is `Disabled` or `Sleep`. |
+| `--ulimit-l-min-gb GB` | 32 | RLIMIT_MEMLOCK threshold (0 disables). |
+| `--shm-min-gb GB` | 8 | `/dev/shm` size threshold (0 disables). |
+| `--rocm-smi-timeout-sec SEC` | 5.0 | Hard timeout for the `rocm-smi --version` self-latency canary; hitting it is a node FAIL (driver likely wedging). |
+| `--hbm-busy-threshold-gib GiB` | 2.0 | FAIL the node if any GPU has at least this many GiB of HBM in use BEFORE smoke touches the device (i.e. someone else is holding it). Boundary is inclusive. |
+| `--allow-foreign-procs` | off | Do NOT FAIL the node when foreign processes are found holding a GPU. They will still be reported. |
+| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Comma-separated process names that are OK to find holding the GPU. Set to `""` to disable the whitelist. |
+| `--gpu-activity-warn-pct PCT` | 20.0 | Warn (does NOT fail) if amd-smi reports any GPU's `gfx_activity_pct` above this when smoke starts. |
+| `--require-tools LIST` | `""` (warn-only) | Comma-separated CLI tool names that MUST be in PATH (`amd-smi`, `rocm-smi`, `lsof`); anything missing becomes a hard node FAIL. |
+| `--no-clean-dump-path` | off | Do NOT auto-wipe stale per-node JSONs / aggregator outputs from `--dump-path` on rank 0 at startup. Default behavior is to clean so re-runs on a different (smaller) nodelist don't inherit ghost PASS verdicts from removed nodes. |
+
+### `aggregate` subcommand
+
+| Flag | Default | Purpose |
+|---|---|---|
+| `--dump-path` | `output/preflight` | Same as `run`. |
+| `--expected-nodes N` | none | If fewer JSONs land within `--wait-timeout-sec`, missing nodes are added as FAIL placeholders. |
+| `--wait-timeout-sec` | 60 | Polling timeout. |
+| `--rocm-smi-warn-sec SEC` | 1.0 | Flag (warn-only) any node where `rocm-smi --version` took longer than this. |
+| `--clock-skew-warn-sec SEC` | 30.0 | Warn when wall-clock spread across nodes exceeds this many seconds. Includes srun launch jitter so the default is loose. |
+| `--hbm-busy-threshold-gib GiB` | 2.0 | Mirrors the `run`-side default; used to label the **GPU pre-touch HBM usage outliers** section. |
+| `--gpu-activity-warn-pct PCT` | 20.0 | Mirrors the `run`-side default; used to label the **GPU compute-activity outliers** section. |
+| `--expected-nodelist-file FILE` | none | One short hostname per line. Missing nodes get their **real short hostname** in the report and `failing_nodes.txt` (instead of `<missing-N>` placeholders). The primus-cli wrapper auto-populates this from `scontrol show hostnames "$SLURM_JOB_NODELIST"` under SLURM. |
+
+### Launcher-level knobs (`primus-cli direct`)
+
+These are consumed by `primus-cli-direct.sh` **before** the `--` separator (not forwarded to the `node_smoke` Python tool):
+
+| Flag       | Purpose |
+|---|---|
+| `--silent` | Back-pocket knob: redirect launcher + tool stdout to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN` on stderr) and the log file are preserved. Exit code propagated. |
+| `--debug`  | Verbose launcher logging. |
+| `--dry-run`| Show the resolved command without executing. |
+| `--env KEY=VALUE` | Inject an env var into the Python process. |
+
+### Rare advanced control (run-only / aggregate-only / no-aggregate)
+
+The primus-cli `node_smoke` subcommand always runs `_cmd_run` on every rank followed by rank-0 `_cmd_aggregate`. That is what users want ~100% of the time. For the rare cases where you need just one phase (e.g. re-aggregate yesterday's JSONs without re-running per-node, or smoke a single node without producing a cluster report), reach for the standalone CLI directly:
+
+```bash
+# Per-node only, no aggregator (useful when scheduling phases separately):
+python -m primus.tools.preflight.node_smoke run --tier2-perf
+
+# Aggregate only (read existing <dump>/smoke/*.json, produce cluster report):
+python -m primus.tools.preflight.node_smoke aggregate \
+    --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5
+```
+
+## Comparison with the full `preflight`
+
+| Aspect | `node_smoke` | full `preflight` |
+|---|---|---|
+| Rendezvous | None — every node independent | Global `torch.distributed` |
+| Wall clock | ~50–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests |
+| GEMM threshold | Hard threshold per GPU | Reports per-GPU numbers, no auto-fail |
+| HBM bandwidth | Yes (D2D `copy_`) | Not measured |
+| Inter-node all-reduce/all-to-all | Not tested (intentionally) | Yes |
+| Drift detection | Yes (versions, NIC firmware, port count) | No |
+| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only |
+| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF |
+
+Use `node_smoke` to **screen** a cluster fast and exclude bad nodes. Use the full `preflight` when you want **deep cross-node measurements** (inter-node bandwidth matrix, ring-P2P, etc.).
+
+---
+
+## Implementation history
+
+Captured here so future contributors understand *why* the design looks the way it does.
+
+### 1. Configurable preflight (predecessor work)
+
+Before `node_smoke` existed, the goal was simply to make the full `preflight` perf phase configurable: which tests to run, which message sizes, which subgroup sizes. Outcome (committed before `node_smoke`):
+
+- `--tests gemm,intra-allreduce,inter-allreduce,...` to select tests
+- `--comm-sizes-mb 2,8,64,1024`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--ring-p2p-sizes-mb`
+- `--intra-group-sizes 2,4,8`, `--inter-group-sizes 2,4,all`
+- `--quick` preset (small warmup/iters, single message size)
+- New flag-precedence rules: `--tests` / `--quick` imply `--perf-test`; mixing perf and info selectors warns and drops info; tuning knobs without perf intent are inert with a quieter warn
+- Report improvements: Node/Rank columns compressed into ranges (e.g. `0-7`), a Node→Hostname legend at the top, "Leader hostname" column showing only the first host of each group
+
+This work is in `primus/tools/preflight/preflight_perf_test.py` and the comm modules. It set up the global accessors (`set_warmup` / `set_iteration` / `get_*`) that `node_smoke` later mirrored to keep iteration counts comparable.
+
+### 2. Why a separate node-local smoke test
+
+A user pointed out that for large jobs, what really matters is "which node has a problem", not which GPU. They proposed a per-node distributed-environment with simple checks (`set_device`, quick bandwidth/speed tests) that returns a success/fail flag per node, *before* the real training job opens its global rendezvous. This was the motivation for `node_smoke.py` — much faster, no dependency on a healthy cluster, and a stuck node can't take down its peers.
+
+### 3. Tiering decision
+
+Two tiers, picked interactively:
+
+- **Tier 1**: mandatory, fast (~5 s/GPU) — `set_device`, alloc, tiny GEMM, plus reused info collectors. The bar is "the GPU enumerated and runs ops".
+- **Tier 2**: optional perf sanity (`--tier2 / --tier2-rccl`) — GEMM TFLOPS, HBM bandwidth, local 8-GPU RCCL. The bar is "the GPU is at expected steady-state performance".
+
+HBM bandwidth was clarified to mean device-to-device `torch.Tensor.copy_` of a 512 MB buffer, counting read + write, which gives ~70–80 % of the MI300X HBM3 roofline (~5300 GB/s) on healthy hardware.
+
+### 4. First scaled run + measurement-quality bug
+
+The first 6-node run produced:
+
+- GEMM 8192³ bf16: smoke median **724 TFLOPS**, full preflight median **765 TFLOPS** (~6 % gap)
+- Local AR 64 MB / 8 GPU: smoke median **199 GB/s**, full preflight median **230 GB/s** (~12 % gap)
+
+Formula audit confirmed `2·S·(P-1)/P / t` and `2·N³/t` are identical to `intra_node_comm.py` and `square_gemm.py`. The systematic offset traced to **iteration counts being too low**. The original RCCL loop was `1 warmup + 1 timed iter` — basically a kernel-launch latency test, not a bandwidth test. Fixed by aligning to the preflight `--quick` preset:
+
+- GEMM: `warmup 3→5, iters 10→20`
+- HBM: `warmup 5→10, iters 10→20`
+- RCCL: `warmup 1→5, iters 1→20`
+
+Per-node runtime cost rounded to <1 s additional. Aggregator gained a "Tier 2 perf summary" section so per-node GEMM/HBM/RCCL outliers are visible without grepping JSONs.
+
+### 5. A + B + C (drift, NIC roll-call, host limits)
+
+Discussion of what the smoke test was *missing* led to six categories. A, B, C landed:
+
+- **A. Stack drift detection** — per-node `fingerprint` (kernel, ROCm, amdgpu, RCCL, torch, NIC firmware, HCA model) + aggregator-side cluster-majority-vs-outlier comparison. Catches "1 of N nodes on a different RCCL build" — a frequent cause of "job dies at minute 3".
+- **B. NIC / RDMA roll-call** — sysfs-only inventory of `/sys/class/infiniband` (no `ibv_devinfo` dependency), per-port hard-fail rules for state ≠ ACTIVE, missing RoCE v2 GIDs, count mismatch (when `--expected-rdma-nics` is set).
+- **C. Host limits** — `RLIMIT_MEMLOCK` (32 GiB default threshold), `/dev/shm` (8 GiB default threshold), plus collected-only NUMA / governor / kernel for drift detection.
+
+Verified with a synthetic two-node drift test (one real + one edited copy with mismatched RCCL, mismatched amdgpu, mismatched `rdma3` firmware, `rdma2:1` DOWN, and `memlock=64 MiB`): every section lit up correctly, `fail_reasons` were prefixed with `nic:` / `host_limits:` for traceability, exit code propagated.
+
+### 6. Aggregator crash on heterogeneous fingerprints
+
+An 18-node run on a different cluster crashed the aggregator with `TypeError: unhashable type: 'dict'`. Root cause: `_stack_drift_rows()` added a key to its scalar-comparison set whenever **any** node reported it as `None` (or a scalar), then iterated **all** nodes' values for that key into `Counter(...)`. On the failing cluster `nic_fw` was `None` on one node and a dict on others — the dict wasn't hashable. Fix:
+
+1. Only collect a key when at least one node reports it as a real scalar (drop the "None counts as scalar" path).
+2. Defense-in-depth `isinstance(v, (str, int, float))` check inside the per-host loop.
+3. Each report section wrapped in its own `try / except` so one section's bug can't truncate the rest of the report.
+4. New "NIC port-count summary" section that always renders and lists nodes whose port count differs from the cluster majority (so partial-NIC degradation like 7-of-8 is visible without `--expected-rdma-nics`).
+
+### 7. Package split (refactor of the 4.5k-line monolith)
+
+`node_smoke.py` had grown to ~4500 lines with all collectors, the orchestrator, the per-GPU subprocess body, and the ~700-line aggregator markdown writer in a single file. The refactor turned it into a Python sub-package (`primus/tools/preflight/node_smoke/`) with one module per Tier 1 sub-section (`collectors/`), the per-GPU subprocess body, the orchestrator, and the aggregator's data shapers (`aggregator/summarizers.py`) and Markdown writer (`aggregator/report.py`, with one `_write_<section>` helper per `##` heading). The single public entry point — `main` — is re-exported from `__init__.py`, so the existing `python -m primus.tools.preflight.node_smoke ...` invocation (used by the primus-cli wrapper and by `_spawn_per_gpu` for per-GPU subprocesses) keeps working unchanged. Behavior parity was checked by diffing the per-node JSON and `smoke_report.md` against a baseline (with a small allowlist for run-variant fields like PIDs, hardware cycle counters, and `available_gb`/`free_gb`/`cached_gb`); CLI help text, JSON schema, report section order, and exit-code semantics for `run` / `_per_gpu` / `aggregate` are byte-identical to pre-refactor.
+
+### 8. Short hostnames + naming nodes that never reported
+
+`failing_nodes.txt` held FQDNs (`socket.gethostname()` returned the FQDN on the failing cluster) — not pipeable into `srun --exclude=`. Nodes that never produced a JSON only showed up as `<missing-N>` placeholders, so operators couldn't act on them.
+
+Fix:
+
+1. Normalize `host = socket.gethostname().split(".", 1)[0]` in `_cmd_run` for both the JSON filename and the `host` field; logs use the short name too.
+2. Aggregator defensively short-normalizes every loaded JSON, so legacy FQDN files produce SLURM-ready txt outputs without re-running the smoke step.
+3. New `aggregate --expected-nodelist-file FILE` flag — missing nodes appended with their real short hostname (and a self-describing `expected hostname '<host>' from --expected-nodelist-file` reason), written to `failing_nodes.txt` directly.
+4. Wrapper resolves `SLURM_JOB_NODELIST` via `scontrol show hostnames` into `<dump>/expected_nodes.txt` and forwards it to the aggregator. Best-effort: silent fallback to count-only behaviour when `scontrol` is unavailable.
+
+This also makes "the node that SLURM marked as `task X: unknown`" visible in the report under its real hostname.
+
+---
+
+## Future work
+
+These were proposed but not yet built. In rough priority order:
+
+### D. GPU low-level health (beyond "alloc + small GEMM works")
+
+Reveals hardware that *enumerates* but is degraded. Most map to one `rocm-smi` query or one sysfs read in the existing per-GPU subprocess.
+
+- GPU count == expected and `lspci -d 1002:` agrees
+- PCIe link width/speed per GPU (`/sys/bus/pci/devices/<bdf>/current_link_{speed,width}`) — catches "GPU at Gen3 x8 because the slot needs reseating"
+- XGMI link matrix between every GPU pair (reuse `primus/tools/preflight/gpu/gpu_topology.py`)
+- HBM size per GPU matches expected
+- ECC counters: uncorrectable as hard fail, correctable as info-only with a cluster-median baseline
+- GPU clock state: flag any GPU stuck at idle GFX clock (stale-state symptom)
+- Throttle reasons from `rocm-smi --showperflevel` (`power_throttle` / `thermal_throttle`)
+- Power cap drift across the cluster
+
+Aggregator gets a "GPU-level drift" section that pinpoints `host:gpu` outliers, not just node-level.
+
+### E. Time / cluster sync
+
+- Wall-clock skew vs `node_rank=0`: each node writes its `time.time()` into its JSON; aggregator computes `max - min` and warns at > 1 s, fails at > 5 s
+- Time-daemon health (`systemctl is-active chronyd / ntpd / systemd-timesyncd`)
+
+### F. Storage / runtime liveness (site-specific)
+
+- Shared-FS latency probe: 1 KB write + `stat` to a unique path, aggregator flags nodes far above the cluster median (Lustre/NFS hiccups)
+- Shared-FS quota / free space
+- DNS resolution sanity for peer hostnames
+- `rocm-smi --version` self-latency (5 s timeout) — catches drivers that have started to wedge but haven't crashed yet (we've seen 30–60 s `rocm-smi` calls precede a full GPU hang by minutes)
+- Container / image hash drift — if the launcher exports `CONTAINER_IMAGE_TAG`, fold it into the existing fingerprint
+
+### Bigger architectural item (lower priority)
+
+- If `NODE_RANK==0` itself fails to start, no aggregator runs anywhere. Possible mitigations: separate aggregator step submitted after the smoke step, or polling watchdog on the submit host. Out of scope for now — handled in practice by always passing `--time=` to `srun` so SLURM force-terminates a stuck job and you can re-run the aggregator alone with `--aggregate-only` + `--expected-nodelist-file`.
diff --git a/docs/posttraining.md b/docs_deprecated/posttraining.md
similarity index 100%
rename from docs/posttraining.md
rename to docs_deprecated/posttraining.md
diff --git a/docs_deprecated/preflight-direct.md b/docs_deprecated/preflight-direct.md
new file mode 100644
index 000000000..416ce4457
--- /dev/null
+++ b/docs_deprecated/preflight-direct.md
@@ -0,0 +1,794 @@
+# Run Preflight Without a Container
+
+> ⚠ **Run the [node-smoke test](./node-smoke-test-instruction.md) first.** `preflight` opens a global `torch.distributed` rendezvous, so a single sick node (wedged driver, leaked rank holding HBM, partial NIC enumeration, time-sync drift, etc.) can stall the whole job for up to `--dist-timeout-sec` seconds — long before any cross-node bandwidth number is produced. The node-smoke test catches those exact failure modes *without* a rendezvous in ~30–60 s and emits a SLURM-ready `failing_nodes.txt` you can pipe straight into `srun --exclude=`. Treat node-smoke as a hard prerequisite; only run `preflight` on the nodes node-smoke marked PASS. See [§0 "Which test should I run?"](#0-which-test-should-i-run) for the side-by-side comparison and the recommended 3-step workflow.
+
+This guide explains how to run Primus's `[preflight](./preflight.md)` cluster-diagnostic tool **directly on the host** (no Docker / Podman), via the standard Primus launcher.
+
+**Git clone the Primus repository to a shared filesystem that all nodes can read.**
+
+```bash
+git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git
+cd Primus
+git checkout dev/preflight-direct-test
+```
+
+**Recommended (through the primus-cli SLURM wrapper):**
+
+```
+runner/primus-cli slurm srun -N <NNODES> --ntasks-per-node=1 -- direct -- preflight [PREFLIGHT_ARGS...]
+```
+
+**Equivalent (bare srun, useful when composing with custom srun flags):**
+
+```
+srun -N <NNODES> --ntasks-per-node=1 runner/primus-cli direct -- preflight [PREFLIGHT_ARGS...]
+```
+
+Both forms produce the **same workload** on the same ranks. The wrapper form is recommended because it auto-resolves `MASTER_ADDR` / `MASTER_PORT` / `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` once on the launching node and passes them to every rank via `--env`, applies any `slurm.`* config defaults (partition / time / etc.) from your YAML, and is the same pattern used for `train` / `benchmark` / `node_smoke`. See [§ Wrapper vs. bare-srun](#wrapper-vs-bare-srun) below for the exact precedence / caveats.
+
+`primus-cli direct` activates an optional Python virtualenv (`VENV_ACTIVATE`), auto-derives the distributed environment variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`) from `SLURM_*` when running inside a SLURM allocation, and then launches the `preflight` Python subcommand via `torchrun` (one worker per GPU). It is the recommended entry point when:
+
+- You're running on a SLURM cluster but cannot (or don't want to) use the container-based path.
+- Your nodes share a Python virtual environment on a network-mounted filesystem.
+- You want a single-node sanity check with no extra configuration.
+
+---
+
+## 0. Which test should I run?
+
+Primus ships **two** complementary cluster screens. Pick the right one — and ideally run them in this order.
+
+
+| Aspect            | `node-smoke` (start here)                                                          | `preflight` (this doc)                                              |
+| ----------------- | ---------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
+| Purpose           | "Which nodes are healthy enough to run anything?"                                  | "What is the actual cross-node performance on the surviving nodes?" |
+| Rendezvous        | None — every node independent                                                      | Global `torch.distributed` rendezvous                               |
+| Wall clock        | ~30–60 s for 6 nodes (Tier 1+2)                                                    | A few minutes; scales with N for inter-node tests                   |
+| Granularity       | Per-node PASS/FAIL                                                                 | Per-rank perf measurements                                          |
+| Safety            | A stuck node cannot wedge its peers                                                | A single hung NIC can stall the whole rendezvous                    |
+| Output            | Per-node JSON + cluster md + SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` | Markdown + PDF perf report                                          |
+| Entry point       | `primus-cli direct -- node_smoke`                                                  | `primus-cli direct -- preflight` (this doc)                         |
+| Quick-start guide | `[node-smoke-test-instruction.md](./node-smoke-test-instruction.md)`               | This doc, §3+                                                       |
+
+
+### Recommended workflow
+
+> **Before running any of the commands below, complete the one-time setup:**
+>
+> 1. **Python virtualenv** on a shared filesystem — see [§2 Set up the Python virtual environment](#2-set-up-the-python-virtual-environment), then point the launcher at it via `export VENV_ACTIVATE=...` (details in [§2 → Tell the launcher where the venv is](#tell-the-launcher-where-the-venv-is)).
+> 2. **NCCL / fabric environment variables** — usually the defaults in `base_env.sh` are fine, but multi-NIC nodes may need `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` / `NCCL_SOCKET_IFNAME` overrides. See [§4 Cluster-specific NCCL configuration](#4-cluster-specific-nccl-configuration) for known-good values per fabric (Broadcom, Pensando Pollara/AINIC).
+
+Through the `primus-cli slurm srun -- direct --` wrapper (recommended):
+
+```bash
+# 1) Prune broken nodes with node-smoke (fast, no rendezvous).
+runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    -- direct -- node_smoke --tier2-perf
+
+# 2) Re-allocate excluding the bad nodes, and run preflight --quick
+#    for a fast cross-node sanity check.
+runner/primus-cli slurm srun -N <good-nnodes> -c 128 --gpus-per-node=8 \
+    --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    -- direct -- preflight --quick
+
+# 3) Optional: full preflight on the same set if --quick numbers
+#    look off, or if you want the full bandwidth matrix.
+runner/primus-cli slurm srun -N <good-nnodes> -c 128 --gpus-per-node=8 \
+    --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    -- direct -- preflight
+```
+
+Equivalent with bare `srun` (works identically; useful when scripting around custom srun flags that don't compose with the wrapper):
+
+```bash
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct -- node_smoke --tier2-perf
+
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct -- preflight --quick
+
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct -- preflight
+```
+
+Why this ordering matters:
+
+- A single broken node can stall a `torch.distributed.init_process_group()` for `--dist-timeout-sec` seconds (default 120), so feeding a known-good list to preflight is much faster.
+- `node-smoke` catches things preflight cannot — leaked / foreign processes, wedged drivers, partial NIC enumeration, time-sync drift, RDMA roll-call issues — that produce *misleading* preflight failures.
+- `preflight --quick` adds the cross-node bandwidth signal that `node-smoke` deliberately does not measure.
+
+---
+
+## 1. Prerequisites
+
+- A working AMD ROCm installation on every node.
+- Network reachability between nodes (Ethernet for bootstrap, RDMA / InfiniBand recommended for perf tests).
+- A Python ≥ 3.10 virtual environment **on a shared filesystem** that all nodes can read (the same path is sourced on every node).
+- The Primus repository checked out somewhere readable from every node.
+
+---
+
+## 2. Set up the Python virtual environment
+
+The environment must live on a path visible from every node (e.g. NFS-mounted home, Lustre, or any shared filesystem). All nodes will `source` the same activation script.
+
+You can use any tool you like; `uv` is the fastest. Either of the following works.
+
+### What you actually need to install
+
+The `preflight` and `node-smoke` tools deliberately use **only a small subset** of Primus's full dependency tree. You do **not** need to install the entire `requirements.txt` — that pulls in trainer / dataset / experiment-tracking packages that neither tool ever imports.
+
+
+| Package              | Required for                                                                          | Skip when                                                                                  |
+| -------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ |
+| `torch` (ROCm build) | Both tools — perf measurements (`torch.matmul`, `torch.distributed`, `torch.cuda.`*). | Never (mandatory).                                                                         |
+| `markdown2`          | `preflight` PDF report only (Markdown → HTML).                                        | You always pass `--disable-pdf`, or you only run `node-smoke` (which never produces PDFs). |
+| `weasyprint`         | `preflight` PDF report only (HTML → PDF).                                             | Same as above.                                                                             |
+| `matplotlib`         | `preflight --plot` only (per-test bandwidth bar charts).                              | You don't pass `--plot`.                                                                   |
+
+
+Everything else in the preflight / node-smoke code path is Python stdlib (`os`, `subprocess`, `socket`, `argparse`, `dataclasses`, `json`, `time`, ...) — no extra installs needed.
+
+### Option A — `uv` (recommended), minimal install
+
+```bash
+mkdir -p ~/envs/preflight
+cd ~/envs/preflight
+
+uv venv --python 3.12
+source .venv/bin/activate
+
+# 1) ROCm-built PyTorch (pin to your ROCm version; rocm7.1 shown here)
+uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir
+
+# 2) Optional: only if you want preflight PDF reports (omit to use --disable-pdf)
+uv pip install markdown2 weasyprint
+
+# 3) Optional: only if you want preflight --plot bar charts
+uv pip install matplotlib
+```
+
+### Option B — `python -m venv`, minimal install
+
+```bash
+mkdir -p ~/envs/preflight
+python3.12 -m venv ~/envs/preflight/.venv
+source ~/envs/preflight/.venv/bin/activate
+
+pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir
+pip install markdown2 weasyprint   # optional, for preflight PDFs
+pip install matplotlib             # optional, for preflight --plot
+```
+
+### Option C — full Primus runtime (only if you also want the rest of Primus)
+
+```bash
+cd /path/to/Primus
+uv pip install -r requirements.txt   # or: pip install -r requirements.txt
+```
+
+This installs every Primus runtime dependency (trainer, dataset loaders, experiment trackers, ...). Use only if you're going to run more than just preflight / node-smoke from this environment.
+
+### Per-tool minimum install matrix
+
+If you want the absolute smallest footprint, install only what your intended invocations need:
+
+
+| Invocation                                       | `torch`  | `markdown2`                       | `weasyprint`                      | `matplotlib` |
+| ------------------------------------------------ | -------- | --------------------------------- | --------------------------------- | ------------ |
+| `node-smoke` (any flags)                         | required | —                                 | —                                 | —            |
+| `preflight --host --gpu --network --disable-pdf` | required | —                                 | —                                 | —            |
+| `preflight --host --gpu --network` (with PDF)    | required | required                          | required                          | —            |
+| `preflight --quick --disable-pdf`                | required | —                                 | —                                 | —            |
+| `preflight --quick` (with PDF)                   | required | required                          | required                          | —            |
+| `preflight ... --plot`                           | required | required (unless `--disable-pdf`) | required (unless `--disable-pdf`) | required     |
+
+
+### Tell the launcher where the venv is
+
+`primus-cli direct` reads the `**VENV_ACTIVATE**` environment variable. When set, it sources the path before launching the Python process; when unset, it is a no-op (the container path, which uses the container's bundled Python, leaves this unset):
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+```
+
+`VENV_ACTIVATE` is the only optional environment variable specific to the direct flow. Everything else has a sensible default; distributed-env variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, ...) are auto-derived from SLURM when not pre-exported.
+
+---
+
+## 3. Run preflight
+
+### Single node (no SLURM)
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Info report only (fast)
+runner/primus-cli direct -- preflight --host --gpu --network
+
+# Info + perf report
+runner/primus-cli direct -- preflight
+
+# Perf report only
+runner/primus-cli direct -- preflight --perf-test
+```
+
+When SLURM is not detected the script defaults to `NNODES=1`, `NODE_RANK=0`, `MASTER_ADDR=localhost`. Any of those can be overridden by exporting them before calling the script.
+
+### Multi-node without SLURM (parallel SSH)
+
+When no scheduler is available (bare-metal, cloud VMs, lab nodes), launch
+`primus-cli direct` on each node yourself via SSH. The script works
+identically — you just pre-export the distributed variables that SLURM
+would normally provide.
+
+#### Requirements
+
+- All nodes share the same filesystem (or at least the same Primus checkout + venv path).
+- Nodes can reach each other on a **data-plane** network interface (not the management NIC).
+- SSH key-based access to each node from the launching host.
+
+#### Required environment variables
+
+
+| Variable             | Description                                                            |
+| -------------------- | ---------------------------------------------------------------------- |
+| `NNODES`             | Total number of nodes                                                  |
+| `NODE_RANK`          | This node's rank (`0` through `NNODES-1`)                              |
+| `MASTER_ADDR`        | IP of rank-0 node **on the data-plane interface**                      |
+| `MASTER_PORT`        | Rendezvous port (default `1234`; increment between concurrent runs)    |
+| `GPUS_PER_NODE`      | GPUs per node (default `8`)                                            |
+| `NCCL_SOCKET_IFNAME` | Data-plane NIC name (e.g. `enp159s0np0`) — **critical for multi-node** |
+| `GLOO_SOCKET_IFNAME` | Same as `NCCL_SOCKET_IFNAME`                                           |
+| `VENV_ACTIVATE`      | Path to virtualenv `activate` script                                   |
+
+
+> **Warning**: `NCCL_SOCKET_IFNAME` auto-detection often picks a management interface
+> (e.g. `enp28s0np0`, `eno8303`) instead of the high-bandwidth data NIC. For multi-node
+> runs this causes `init_process_group` to hang or NCCL to fail silently. Always set it
+> explicitly.
+
+#### Identifying the correct data-plane interface
+
+```bash
+# On any node, find the interface whose IP matches the MASTER_ADDR subnet:
+ip -4 addr show | grep "10.245.134"
+# → enp159s0np0  inet 10.245.134.129/24
+
+# Or check which interface routes to the master:
+ip route get 10.245.134.129 | awk '{print $5; exit}'
+### Multi-node via SLURM
+
+`primus-cli direct` auto-detects a SLURM allocation (via `SLURM_JOB_ID`) and derives all distributed variables from `SLURM_*` automatically. **Pre-exported values always win**, so the same launcher script also works inside the `primus-cli slurm srun ... -- direct -- ...` chain (where `slurm-entry` has already set these via `--env`):
+
+| Variable        | Resolved as                                                          |
+| --------------- | -------------------------------------------------------------------- |
+| `NNODES`        | `NNODES` → `SLURM_NNODES` → `SLURM_JOB_NUM_NODES` → `1`              |
+| `NODE_RANK`     | `NODE_RANK` → `SLURM_NODEID` → `SLURM_PROCID` → `0`                  |
+| `MASTER_ADDR`   | `MASTER_ADDR` (if not empty / not `localhost`) → first hostname from `scontrol show hostnames "$SLURM_NODELIST"` |
+| `MASTER_PORT`   | `MASTER_PORT` → `1234`                                               |
+| `GPUS_PER_NODE` | `GPUS_PER_NODE` → `8`                                                |
+
+Run it as a single task per node (the script invokes `torchrun` internally, which spawns one worker per GPU):
+
+> **Verify NCCL / network env first.** The script sets sensible `NCCL_`* defaults via `base_env.sh`, but auto-detection can pick the wrong device on multi-NIC nodes. Always confirm `NCCL_IB_HCA`, `NCCL_IB_GID_INDEX`, `NCCL_SOCKET_IFNAME`, and `GLOO_SOCKET_IFNAME` (set to the same value as `NCCL_SOCKET_IFNAME`) are correct for your fabric, and `export` overrides before running. See [§4](#4-cluster-specific-nccl-configuration) for cluster-specific values.
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# export NCCL_IB_HCA=rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
+# export NCCL_IB_GID_INDEX=3
+# export NCCL_SOCKET_IFNAME=eno0
+# export GLOO_SOCKET_IFNAME=eno0
+
+# Recommended: through the primus-cli SLURM wrapper.
+runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 \
+    --nodelist <nodes> --ntasks-per-node=1 \
+    -- direct -- preflight --perf-test
+
+# Or, equivalently, with bare srun:
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight --perf-test
+```
+
+
+
+#### Wrapper vs. bare-srun
+
+Both forms target the **same** `primus-cli-direct.sh` launcher and produce identical workloads. The difference is only in how the SLURM context is constructed:
+
+
+| Aspect                                   | `primus-cli slurm srun -- direct --` (recommended)                                                                                                             | Bare `srun ... primus-cli direct --`                                                                                                |
+| ---------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
+| `MASTER_ADDR` resolution                 | Resolved **once** on the launching node via `scontrol show hostnames "$SLURM_NODELIST" | head -n1`, then propagated to every rank via `--env MASTER_ADDR=...`. | Each rank re-derives it inside `primus-cli-direct.sh` STEP 4.7 from `SLURM_`* (same result, more `scontrol` calls).                 |
+| `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` | Set explicitly by `slurm-entry.sh` via `--env`.                                                                                                                | Derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_PROCID` inside `direct.sh`.                                                   |
+| `slurm.*` config defaults                | Honored (partition, time, ntasks-per-node, etc. from the active YAML).                                                                                         | Not consulted — you pass every flag explicitly to `srun`.                                                                           |
+| Default wall-time                        | `-t 4:00:00` is auto-added if you don't pass `--time`.                                                                                                         | None — `srun` uses the cluster default (may reject the job).                                                                        |
+| `direct` keyword                         | **Required**: `primus-cli slurm srun ... -- direct -- <cmd>`. Without `direct`, the wrapper routes through the **container** path.                             | N/A — there's only one path.                                                                                                        |
+| `--ntasks-per-node=1`                    | **Not auto-added**. Pass it on the CLI (before the first `--`) or set it in the `slurm.`* config.                                                              | **Not auto-added**. Pass it as an `srun` flag.                                                                                      |
+| Best for                                 | Production / repeatable runs. Same pattern as `train` / `benchmark` / `node_smoke`.                                                                            | Ad-hoc runs where you want to compose with arbitrary `srun` flags (`--nodelist=$(...)`, `--exclude=...` from a runtime file, etc.). |
+
+
+For the rest of this doc the examples use bare `srun` for brevity, but every example also works with the wrapper form by substituting `srun <flags> runner/primus-cli direct --` → `runner/primus-cli slurm srun <flags> -- direct --`.
+
+### Key `srun` flags
+
+
+| Flag                  | Why it's necessary                                                                                                                                                                                      |
+| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `-c 128`              | Allocate all CPU cores per task. Without this, SLURM may default to 1 core, which starves the RCCL network proxy threads and can cause >30× slowdown on perf tests. Set this to your node's core count. |
+| `--gpus-per-node=8`   | Grants GPU device access (`/dev/kfd`, `/dev/dri`). Required for non-container execution.                                                                                                                |
+| `--ntasks-per-node=1` | One launcher invocation per node; `primus-cli direct` then spawns 8 workers per node via `torchrun`.                                                                                                    |
+| `-t 00:45:00`         | Wall-clock limit. Full perf tests on 8N usually finish well under 10 min.                                                                                                                               |
+
+
+> Tip — check core count: `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`
+
+---
+
+## 4. Cluster-specific NCCL configuration
+
+`primus-cli direct` sources `runner/helpers/envs/base_env.sh`, which sets sensible defaults for `NCCL_`* and auto-detects `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME`. Pre-exported values from your shell take precedence, so the standard pattern is:
+
+```bash
+export VAR=value
+runner/primus-cli direct -- preflight ...
+```
+
+### Broadcom NICs (no AINIC)
+
+Most clusters fall here. The defaults from `base_env.sh` are usually fine, but the two values most commonly worth overriding are:
+
+```bash
+export NCCL_CROSS_NIC=1     # default in base_env.sh is 0
+export NCCL_PXN_DISABLE=0   # default in base_env.sh is 1
+
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight --perf-test
+```
+
+### Pensando Pollara (AINIC) RDMA
+
+```bash
+export USING_AINIC=1
+export NCCL_IB_GID_INDEX=1   # AINIC uses index 1 (default in base_env.sh is 3)
+export NCCL_PXN_DISABLE=0
+
+srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist <nodes> \
+    --ntasks-per-node=1 \
+    runner/primus-cli direct -- preflight
+```
+
+> `primus-cli direct` *does* accept `--env KEY=VALUE` on its own command line (placed before `--`), in addition to the conventional `export`/`srun --export=` approaches.
+
+---
+
+## 5. Launcher flags vs. preflight flags
+
+Anything you place **after** the `--` separator is forwarded verbatim to the `preflight` Python tool. The launcher (`primus-cli-direct.sh`) consumes a small set of flags **before** `--`. The one most users care about is `--silent`.
+
+### Launcher-only flags (before `--`)
+
+
+| Flag              | Effect                                                                                                                                                                                                                                                                                                                                                                      |
+| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `--silent`        | Back-pocket knob: redirect the launcher's and the Python tool's `stdout` to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN`, written to `stderr`) are preserved so real failures still surface; the log file under `logs/` captures everything. Exit code is propagated unchanged. **Not recommended** for normal use — you lose live progress; prefer the log file. |
+| `--debug`         | Verbose launcher logging (`PRIMUS_LOG_LEVEL=DEBUG`). Forwarded to the Python tool as `--debug` too.                                                                                                                                                                                                                                                                         |
+| `--dry-run`       | Print the resolved configuration and final `torchrun` / `python3` command without executing.                                                                                                                                                                                                                                                                                |
+| `--single`        | Force `python3` instead of `torchrun`. `node_smoke` auto-selects this; for `preflight` you usually want the default (`torchrun`).                                                                                                                                                                                                                                           |
+| `--env KEY=VALUE` | Inject an env var into the Python process (in addition to anything `export`-ed in the shell).                                                                                                                                                                                                                                                                               |
+| `--log_file PATH` | Redirect the captured tee log to a specific path (default: `logs/log_<timestamp>.txt`).                                                                                                                                                                                                                                                                                     |
+
+
+See `runner/primus-cli direct --help` for the full set.
+
+### Forwarded `preflight` flags (after `--`, most common)
+
+See [Preflight](./preflight.md) for the full list. The most common are:
+
+- Mode selection: `--host`, `--gpu`, `--network`, `--perf-test`, `--tests`, `--quick`
+- Test tuning: `--comm-sizes-mb`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--intra-group-sizes`, `--inter-group-sizes`, `--ring-p2p-sizes-mb`
+- Reporting: `--dump-path`, `--report-file-name`, `--disable-pdf`, `--plot`
+- Reliability: `--comm-cleanup-delay-sec`, `--dist-timeout-sec`
+
+If you do not pass `--report-file-name`, `preflight` auto-generates a unique one of the form:
+
+```
+preflight-${NNODES}N-YYYYMMDD-HHMMSS
+```
+
+This guarantees that each run lands in its own files and prevents stale leftovers from earlier runs from being mistaken for fresh output. The auto-name logic now lives in the Python tool itself, so every call site (host `srun ... primus-cli direct`, `primus-cli slurm ... -- direct`, `primus-cli slurm ... -- container`) gets the same fresh name.
+
+### Examples
+
+The examples below all assume one of the two equivalent shell-prefix conventions. Pick whichever matches your habits — every example block in this section works with either definition:
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+
+# Recommended: through the primus-cli SLURM wrapper. Auto-resolves
+# MASTER_ADDR/NNODES/NODE_RANK once on the launching node and propagates
+# them via --env; honors slurm.* config defaults.
+SRUN="runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist <nodes> --"
+# Then in every example below, replace `$SRUN runner/primus-cli direct --`
+# with just `$SRUN direct --`. (The wrapper expects the entry-mode keyword
+# `direct` as the first token after the inner `--`.)
+
+# Equivalent: bare srun. NNODES/NODE_RANK/MASTER_ADDR get derived inside
+# primus-cli-direct.sh's STEP 4.7 directly from SLURM_*; same net effect.
+SRUN="srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist <nodes>"
+```
+
+The examples in this section use the **bare-srun** form below for brevity (since `$SRUN runner/primus-cli direct -- preflight` reads naturally as one command line). To use the wrapper form instead, substitute `$SRUN runner/primus-cli direct --` → `$SRUN direct --` after exporting `SRUN` to the wrapper variant.
+
+#### A. Mode selection
+
+```bash
+# Default: info report + every perf test
+$SRUN runner/primus-cli direct -- preflight
+
+# Info-only (fast, no torch.distributed rendezvous)
+$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf
+
+# Perf-only, every test
+$SRUN runner/primus-cli direct -- preflight --perf-test
+
+# Fast pre-launch sanity preset (gemm + intra-AR + inter-AR @ 64,1024 MB,
+# full intra-node group, full N-node inter group, low warmup/iter)
+$SRUN runner/primus-cli direct -- preflight --quick
+```
+
+> **Note**: Mixing perf-mode flags (`--perf-test` / `--tests` / `--quick`) with info selectors (`--host` / `--gpu` / `--network`) makes preflight drop the info selectors with a `WARN`. Run two invocations if you want both reports.
+
+#### B. Test selection (`--tests`)
+
+```bash
+# Only GEMM
+$SRUN runner/primus-cli direct -- preflight --tests gemm
+
+# Only the inter-node bandwidth tests
+$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall
+
+# Only the inter-node ring P2P
+$SRUN runner/primus-cli direct -- preflight --tests inter-ring-p2p
+
+# Combine: GEMM + inter-AR with overridden sizes/groups
+$SRUN runner/primus-cli direct -- preflight \
+    --tests gemm,inter-allreduce \
+    --comm-sizes-mb 64,1024 \
+    --inter-group-sizes all
+```
+
+Valid `--tests` tokens: `gemm`, `intra-allreduce`, `intra-alltoall`, `inter-allreduce`, `inter-alltoall`, `inter-p2p`, `inter-ring-p2p`, `all`. Unknown tokens fail fast (before NCCL init).
+
+#### C. Message sizes
+
+```bash
+# One CSV applied to both intra and inter
+$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128
+
+# Different sizes for intra vs inter (override wins over --comm-sizes-mb)
+$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128 --intra-comm-sizes-mb 4,32
+
+# Inter-only override (also covers inter-p2p when enabled)
+$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-p2p \
+    --comm-sizes-mb 8,128 --inter-comm-sizes-mb 16,512
+```
+
+#### D. Group sizes
+
+```bash
+# Custom intra-node group sizes (each must divide LOCAL_WORLD_SIZE)
+$SRUN runner/primus-cli direct -- preflight \
+    --tests intra-allreduce \
+    --intra-group-sizes 4,8
+
+# Custom inter-node groups: 2-node pairs and the full N-node group
+$SRUN runner/primus-cli direct -- preflight \
+    --tests inter-allreduce \
+    --inter-group-sizes 2,all
+```
+
+> Note: for `inter-alltoall` only, every requested per-group node count is internally clamped to **16** (real-world MoE training rarely dispatches across more nodes; see `[preflight.md` §5.2](./preflight.md#52-group-sizes) for the rationale). The other inter-node tests use the requested sizes unchanged. So on a 128-node cluster, `--tests inter-alltoall --inter-group-sizes all` actually runs at 16-node sub-groups, while `--tests inter-allreduce --inter-group-sizes all` runs at 128 nodes as written.
+
+#### E. Ring P2P sizes
+
+```bash
+$SRUN runner/primus-cli direct -- preflight \
+    --tests inter-ring-p2p \
+    --ring-p2p-sizes-mb 5,20,80
+```
+
+#### F. Plotting
+
+```bash
+# Generate per-test bandwidth bar charts under <dump-path>/<test>/*.png
+$SRUN runner/primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce --plot
+```
+
+#### G. Reliability knobs
+
+```bash
+# Bump the per-phase cleanup delay. Default 2.0 is sufficient at every
+# cluster size for the comm shapes preflight exercises (inter-alltoall
+# is internally capped at 16 nodes; see preflight.md §5.2). Only bump
+# this on very flaky networks or unusual kernel TIME_WAIT settings.
+$SRUN runner/primus-cli direct -- preflight --quick --comm-cleanup-delay-sec 5
+
+# Fail fast if torch.distributed rendezvous can't complete in 30s
+$SRUN runner/primus-cli direct -- preflight --perf-test --dist-timeout-sec 30
+```
+
+> Operating clusters at ≥ 128 nodes? See `[preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the recommended OS-level tunings (`tcp_tw_reuse`, wider `ip_local_port_range`) and per-test invocation patterns. With the §5.2 inter-alltoall cap in place, a default invocation is safe at every cluster size; the §7.2 sysctls remain best-practice for any RDMA workload.
+
+#### H. Reporting & output layout
+
+```bash
+# Quick info-only check on 4 nodes, no PDF
+$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf \
+    --report-file-name info-4N
+
+# Perf test only, silenced (CI-friendly), explicit name. Note that --silent
+# is consumed by primus-cli-direct.sh and must appear BEFORE the `--`
+# separator; everything after `--` is forwarded to the preflight Python tool.
+$SRUN runner/primus-cli direct --silent -- preflight --perf-test \
+    --report-file-name nightly-4N-perf
+
+# Archive each run under its own directory
+$SRUN runner/primus-cli direct -- preflight --quick \
+    --dump-path /shared/preflight-archive/$(date +%Y%m%d-%H%M%S)
+```
+
+#### I. Backward-compat aliases
+
+These still work and are equivalent to flags above. Use them only when retrofitting older scripts.
+
+```bash
+# Same as --host --gpu --network
+$SRUN runner/primus-cli direct -- preflight --check-host --check-gpu --check-network
+
+# Same as --inter-group-sizes all AND drops inter-p2p
+$SRUN runner/primus-cli direct -- preflight --perf-test --no-split-nodes-subgroup
+```
+
+#### J. Combined "production-ready" pre-launch screen
+
+```bash
+# 1) Smoke first to prune broken nodes (note: --silent goes BEFORE `--`)
+srun -N "$SLURM_NNODES" --ntasks-per-node=1 \
+    runner/primus-cli direct --silent -- node_smoke --tier2-perf
+
+# 2) Quick perf sanity on the survivors
+srun -N <good-nnodes> -c 128 --gpus-per-node=8 --ntasks-per-node=1 \
+    --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \
+    runner/primus-cli direct --silent -- preflight --quick \
+        --comm-cleanup-delay-sec 5 --dist-timeout-sec 60 \
+        --report-file-name screen-$(date +%Y%m%d-%H%M%S)
+```
+
+---
+
+## 6. Outputs
+
+Reports are written to `--dump-path` (default: `output/preflight/`), with the basename from `--report-file-name` and a `_perf` suffix for performance reports:
+
+
+| File              | Produced by                                     | Notes                     |
+| ----------------- | ----------------------------------------------- | ------------------------- |
+| `<name>.md`       | `--host --gpu --network` (or default selection) | Info report               |
+| `<name>.pdf`      | same, unless `--disable-pdf`                    | Info report PDF           |
+| `<name>_perf.md`  | `--perf-test`                                   | Perf report (GEMM + comm) |
+| `<name>_perf.pdf` | same, unless `--disable-pdf`                    | Perf report PDF           |
+
+
+Only **rank 0** writes the report. After preflight completes, the Python tool prints the absolute path of every report file it produced to stdout. Under `--silent` these prints go to `/dev/null` along with everything else (one of the trade-offs of using `--silent`); without `--silent` the announcement is visible live. Sample output:
+
+```
+[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925.md
+[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925_perf.md
+```
+
+---
+
+## 7. Environment variable reference
+
+Variables read by `primus-cli direct` itself:
+
+
+| Variable        | Required | Default                                                                      | Purpose                                                                                                                |
+| --------------- | -------- | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------- |
+| `VENV_ACTIVATE` | no       | —                                                                            | Path to the venv `bin/activate` script. Unset = no-op (use system / container Python). Set + missing file = fail-fast. |
+| `NNODES`        | no       | `1` (or auto-derived from `SLURM_NNODES` / `SLURM_JOB_NUM_NODES`)            | Number of nodes. Pre-exported always wins.                                                                             |
+| `NODE_RANK`     | no       | `0` (or auto-derived from `SLURM_NODEID` / `SLURM_PROCID`)                   | This node's rank. Pre-exported always wins.                                                                            |
+| `GPUS_PER_NODE` | no       | `8`                                                                          | GPUs per node                                                                                                          |
+| `MASTER_ADDR`   | no       | `localhost` (or first host from `scontrol show hostnames "$SLURM_NODELIST"`) | Rendezvous host. Pre-exported always wins.                                                                             |
+| `MASTER_PORT`   | no       | `1234`                                                                       | Rendezvous port                                                                                                        |
+
+
+Variables consumed downstream by `primus-cli direct` / `base_env.sh` (set them via `export`):
+
+
+| Variable             | Default in `base_env.sh` | When to override                                  |
+| -------------------- | ------------------------ | ------------------------------------------------- |
+| `NCCL_SOCKET_IFNAME` | auto-detected            | Force a specific Ethernet interface for bootstrap |
+| `NCCL_IB_HCA`        | auto-detected            | Force specific RDMA HCAs                          |
+| `NCCL_IB_GID_INDEX`  | `3`                      | `1` on AINIC clusters                             |
+| `NCCL_CROSS_NIC`     | `0`                      | `1` for multi-rail IB fabrics                     |
+| `NCCL_PXN_DISABLE`   | `1`                      | `0` to enable PXN multi-hop NIC sharing           |
+| `USING_AINIC`        | unset                    | `1` on Pensando Pollara clusters                  |
+| `NCCL_DEBUG`         | unset                    | `INFO` for verbose NCCL logging                   |
+
+
+---
+
+## 8. Troubleshooting
+
+### `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...`
+
+`VENV_ACTIVATE` was set in the environment but the path it points at doesn't exist on this node. This is a fail-fast guard to prevent a silent fallback to system Python (which usually has the wrong `torch` / no ROCm). Either fix the path:
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+```
+
+… or unset it to fall back to the container / system Python:
+
+```bash
+unset VENV_ACTIVATE
+```
+
+If the path looks right but the file still appears missing, confirm the venv lives on a filesystem visible from the node SLURM scheduled you onto.
+
+### `[Primus:Preflight] FAIL: No GPUs detected`
+
+The Python process inside the venv can't find ROCm. Diagnose with:
+
+```bash
+srun --nodes=1 --nodelist=<node> bash -c '
+echo "=== PATH ==="; echo $PATH
+echo "=== LD_LIBRARY_PATH ==="; echo $LD_LIBRARY_PATH
+echo "=== rocm-smi ==="; rocm-smi --showid 2>&1
+echo "=== Python torch check ==="
+source ~/envs/preflight/.venv/bin/activate
+python3 -c "import torch; print(\"hip:\", torch.version.hip); print(\"available:\", torch.cuda.is_available()); print(\"count:\", torch.cuda.device_count())"
+'
+```
+
+If `LD_LIBRARY_PATH` is empty, set it explicitly:
+
+```bash
+export LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-}
+```
+
+### Report announcement points at stale files
+
+This shouldn't happen with the current Python tool — the auto-generated unique report name (`preflight-${NNODES}N-<timestamp>`) ensures every run gets a fresh path. If you explicitly pass `--report-file-name X`, you're responsible for choosing a name that doesn't collide with prior runs.
+
+### Slow perf tests (~30× expected)
+
+Almost always a symptom of insufficient CPU cores. Pass `-c <cores-per-node>` to `srun` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`.
+
+### Using `conda` instead of venv
+
+`primus-cli direct` does `source "$VENV_ACTIVATE"`, which works for venv/uv but not directly for conda. Two options:
+
+1. Create a venv inside the conda env and point `VENV_ACTIVATE` at that venv's activate script.
+2. Write a small shim activate script (e.g. `~/envs/conda-shim.sh`) that activates conda and the desired env, then point `VENV_ACTIVATE` at it:
+  ```bash
+   # ~/envs/conda-shim.sh
+   source "$HOME/miniconda3/etc/profile.d/conda.sh"
+   conda activate <env_name>
+  ```
+
+### "Address already in use" during perf tests
+
+This error occurs when peak simultaneous ESTAB sockets per node during an `ncclCommInit` exhausts the kernel ephemeral-port pool, so the next outgoing `bind()` walks the entire range without finding an allocatable port. (Despite the name and the `TIME_WAIT` framing in the kernel docs, accumulated `TIME_WAIT` count does *not* gate this for NCCL inter-node OOB — see `[preflight.md` §7.1](./preflight.md#71-why-address-already-in-use-used-to-surface-at-scale) for the mechanism and the empirical evidence.)
+
+Preflight has two complementary defenses:
+
+1. The **inter-node alltoall sub-group is internally capped at 16 nodes** (see `[preflight.md` §5.2](./preflight.md#52-group-sizes)) — the only test that, at large scale, can push peak ESTAB anywhere near the per-node ephemeral pool. The cap eliminates this failure mode by construction.
+2. A **global barrier + `--comm-cleanup-delay-sec` sleep** (default 2 s) is inserted after every comm destroy, primarily for cross-rank synchronization across the destroy → setup transition.
+
+If you still see `Address already in use` (e.g. on a network with an unusually narrow ephemeral-port range), the directly relevant **OS-level tuning** is widening that range — best-practice for any RDMA host:
+
+```bash
+# Widen the ephemeral port range from ~28k to ~64k. This is the only
+# OS knob that directly addresses the binding constraint.
+sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535"
+
+# General hygiene for hosts running mixed RDMA + repeated outgoing
+# TCP workloads (NCCL inter-node OOB by itself doesn't benefit from
+# this -- see preflight.md §7.2 for why).
+sudo sysctl -w net.ipv4.tcp_tw_reuse=1
+```
+
+As a fallback, raise the per-phase delay:
+
+```bash
+# Bump the per-phase delay (default 2 s) on a particularly stressed
+# network. Rarely needed in practice with the §5.2 alltoall cap.
+runner/primus-cli direct -- preflight --comm-cleanup-delay-sec 5
+```
+
+See `[preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the full explanation, persistence, and recommended large-cluster invocation patterns (split tests into separate runs, etc.).
+
+If the error occurs at `init_process_group` (before tests even start), it typically means a previous job left port 29500 in `TIME_WAIT`. Either wait ~60 s or use a different port:
+
+```bash
+export MASTER_PORT=29501
+```
+
+### Capturing full output
+
+The launcher already writes a complete log to `logs/log_<timestamp>.txt` (configurable via `--log_file PATH`), even under `--silent`. If you also want a copy at the call site, redirect there:
+
+```bash
+srun ... runner/primus-cli direct -- preflight --perf-test \
+    2>&1 | tee preflight-$(date +%Y%m%d-%H%M%S).log
+```
+
+---
+
+## 9. Automated node bisection (finding the bad node in an NCCL hang)
+
+When a cluster-wide preflight run hangs or fails, use
+`[tools/preflight_bisect/bisect.py](../tools/preflight_bisect/bisect.py)` to
+run `preflight --perf-test` on smaller Slurm node subsets until suspect nodes
+are isolated.
+
+### Prerequisites
+
+1. Working non-container preflight setup from the sections above, with
+  `VENV_ACTIVATE` exported from a shared filesystem path.
+2. Run from the SLURM login/head node, where both `scontrol` and `srun` are
+  available.
+3. Run from inside a Slurm allocation, or provide a Slurm nodelist explicitly.
+
+### Example from inside an allocation
+
+```bash
+export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+mkdir -p output
+
+python tools/preflight_bisect/bisect.py \
+    --nodelist "$SLURM_NODELIST" \
+    --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \
+    --trial-timeout-sec 600 \
+    --slurm-time 00:15:00 \
+    --preflight-env USING_AINIC=1 \
+    --preflight-env NCCL_IB_GID_INDEX=1 \
+    --preflight-env NCCL_CROSS_NIC=1 \
+    --preflight-env NCCL_PXN_DISABLE=0 \
+    2>&1 | tee output/bisect-latest.log
+```
+
+Adjust the `--preflight-env` lines to match your cluster. Per-trial logs and a
+final `summary.txt` are written under `--output-dir`.
+
+> Note: Set `--trial-timeout-sec` high enough for a healthy subset to finish.
+> Too small a timeout can turn slow-but-good trials into false failures, causing
+> the bisection to explore extra paths.
+>
+> Note: `--preflight-env KEY=VALUE` values are concatenated into a single
+> `srun --export=ALL,...` argument, so values must not contain commas or
+> whitespace. Keep comma-containing values as normal exported environment
+> variables.
+
+---
+
+## 10. See also
+
+- [Preflight](./preflight.md) — full reference for the `preflight` subcommand and its flags
+- [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) — container-based and `primus-cli slurm` workflows
+- `[runner/primus-cli-direct.sh](../runner/primus-cli-direct.sh)` — the direct launcher itself (`primus-cli direct` dispatches here)
+- `[primus/tools/preflight/](../primus/tools/preflight/)` — preflight implementation
+- `[tools/preflight_bisect/bisect.py](../tools/preflight_bisect/bisect.py)` — bisect wrapper for narrowing down failing nodes in multi-node preflight runs
diff --git a/docs_deprecated/preflight.md b/docs_deprecated/preflight.md
new file mode 100644
index 000000000..7f3e90c3f
--- /dev/null
+++ b/docs_deprecated/preflight.md
@@ -0,0 +1,522 @@
+# Preflight
+
+`preflight` is Primus' cluster diagnostic tool. It produces:
+
+- A **fast info report** (host / GPU / network configuration), and
+- A configurable suite of **performance tests** (GEMM TFLOPS, intra-node and inter-node communication bandwidth, P2P, ring P2P).
+
+Use it to spot misconfiguration, hardware degradation, or perf outliers **before** committing a large distributed training run to a global rendezvous.
+
+- **User-facing entry**: `primus-cli ... -- preflight [args]`
+- **No-container launcher**: `runner/primus-cli direct -- preflight ...` — see [`preflight-direct.md`](./preflight-direct.md).
+- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py` → `primus/tools/preflight/preflight_perf_test.py`.
+
+> Looking for a faster, distributed-rendezvous-free per-node screen? See [`node-smoke.md`](./node-smoke.md) (and the [quick-start guide](./node-smoke-test-instruction.md)). The recommended workflow is **smoke first, preflight second** — see [§10 Comparison with node-smoke](#10-comparison-with-node-smoke).
+
+---
+
+## 1. Two run modes (and how preflight picks one)
+
+Preflight has two report types, controlled by a single precedence rule:
+
+| Mode | Triggered by | What it does |
+|---|---|---|
+| **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. **No `torch.distributed` rendezvous.** Cannot hang on network misconfig. |
+| **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. |
+| **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. |
+
+### Mode precedence
+
+1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.**
+   If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations.
+2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.**
+   Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them.
+3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests.
+
+The default order ensures you always get a report even if `torch.distributed` initialization later hangs.
+
+---
+
+## 2. Quick start
+
+### Info report only (fast, no rendezvous)
+
+```bash
+primus-cli direct -- preflight --host --gpu --network
+```
+
+### Full preflight (info + every perf test)
+
+```bash
+primus-cli direct -- preflight
+```
+
+### Perf tests only
+
+```bash
+primus-cli direct -- preflight --perf-test
+```
+
+### Fast pre-launch sanity check
+
+```bash
+primus-cli direct -- preflight --quick
+```
+
+Equivalent on SLURM via `primus-cli slurm`:
+
+```bash
+primus-cli slurm srun -N 4 -- preflight --quick
+```
+
+Without a container, see [`preflight-direct.md`](./preflight-direct.md) for the equivalent `runner/primus-cli direct -- preflight ...` invocations.
+
+---
+
+## 3. Test selection (`--tests`)
+
+`--tests` takes a comma-separated list of canonical tokens (or `all`). Implies `--perf-test`.
+
+| Token | What it runs |
+|---|---|
+| `gemm` | Single-GPU square GEMM TFLOPS sweep. |
+| `intra-allreduce` | Intra-node `all_reduce` bandwidth at every selected `--intra-group-sizes` x `--intra-comm-sizes-mb`. |
+| `intra-alltoall` | Intra-node `all_to_all` bandwidth, same configuration matrix. |
+| `inter-allreduce` | Inter-node `all_reduce` bandwidth at every selected `--inter-group-sizes` x `--inter-comm-sizes-mb`. |
+| `inter-alltoall` | Inter-node `all_to_all` bandwidth, same configuration matrix. |
+| `inter-p2p` | Inter-node 2-rank P2P send/recv. Requires `--inter-group-sizes` to actually contain pair-able sizes. |
+| `inter-ring-p2p` | Inter-node ring-pattern P2P, sized by `--ring-p2p-sizes-mb`. |
+| `all` | Every token above. Default when `--tests` is omitted. |
+
+Examples:
+
+```bash
+# GEMM only
+primus-cli direct -- preflight --tests gemm
+
+# Just the inter-node bandwidth tests
+primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall
+
+# Combine with size overrides
+primus-cli direct -- preflight \
+    --tests gemm,inter-allreduce \
+    --comm-sizes-mb 64,1024 \
+    --inter-group-sizes all
+```
+
+Unknown tokens fail fast (before any rendezvous):
+
+```text
+[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'.
+Valid tokens: gemm, intra-allreduce, intra-alltoall, inter-allreduce,
+inter-alltoall, inter-p2p, inter-ring-p2p, all
+```
+
+---
+
+## 4. Quick preset (`--quick`)
+
+`--quick` is the recommended **pre-launch sanity** preset. Implies `--perf-test`. It substitutes:
+
+| Knob | `--quick` value |
+|---|---|
+| `--tests` | `gemm,intra-allreduce,inter-allreduce` |
+| `--comm-sizes-mb` | `64,1024` |
+| `--intra-group-sizes` | `LOCAL_WORLD_SIZE` (full intra-node group only) |
+| `--inter-group-sizes` | `all` (full N-node group only) |
+| `warmup` | `5` |
+| `iteration` | `20` |
+
+**User-supplied flags override the preset.** For example:
+
+```bash
+# Quick preset, but with a custom size set
+primus-cli direct -- preflight --quick --comm-sizes-mb 32,256
+```
+
+A full perf run with default knobs takes minutes; `--quick` typically finishes in <60s on healthy hardware.
+
+---
+
+## 5. Tuning the perf tests
+
+All perf tuning knobs default to `None` so preflight can tell whether you set them. When unset, the documented defaults below apply.
+
+### 5.1 Message sizes (collective + P2P)
+
+| Flag | Default | Applies to |
+|---|---|---|
+| `--comm-sizes-mb CSV` | `2,4,8,16,32,64,128,256,512,1024` | Default for both intra- and inter-node `allreduce` / `alltoall` and `inter-p2p` when no specific override is given. |
+| `--intra-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **intra-node** `allreduce` / `alltoall`. |
+| `--inter-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **inter-node** `allreduce` / `alltoall` / `inter-p2p`. |
+
+```bash
+# Smaller, focused sweep
+primus-cli direct -- preflight --comm-sizes-mb 8,128
+
+# Different sizes for intra vs inter
+primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce \
+    --comm-sizes-mb 8,128 \
+    --intra-comm-sizes-mb 4,32
+```
+
+### 5.2 Group sizes
+
+| Flag | Default | Notes |
+|---|---|---|
+| `--intra-group-sizes CSV` | `2,4,8` | Each value must divide `LOCAL_WORLD_SIZE`. |
+| `--inter-group-sizes CSV` | `2,4,all` | `all` means the full N-node group. Other values are subgroup sizes. **For `inter-alltoall` only**, every requested per-group node count is internally clamped to **16** before deduping (see "Inter-node alltoall is capped at 16 nodes" below). The other inter-node tests (`inter-allreduce`, `inter-p2p`, `inter-ring-p2p`) use the requested sizes unchanged. |
+
+```bash
+# All-GPU intra + full N-node inter only
+primus-cli direct -- preflight \
+    --tests intra-allreduce,inter-allreduce \
+    --intra-group-sizes 8 \
+    --inter-group-sizes all
+```
+
+Validation is gated by which tests are actually selected. For example, `--tests gemm --intra-group-sizes 3` does **not** abort on a host with `LOCAL_WORLD_SIZE=8`; the intra-group constraint is only checked when an intra test is enabled.
+
+#### Inter-node alltoall is capped at 16 nodes
+
+Regardless of the cluster size or what `--inter-group-sizes` requests, the `inter-alltoall` test always runs on per-group node counts of at most **16**. Concretely, every requested value `G` is replaced with `min(G, 16)`, and the resulting list is deduped. Examples:
+
+| Cluster | `--inter-group-sizes` | Requested (resolved) | `inter-alltoall` actually runs |
+|---|---|---|---|
+| 8 N | `all` | `[8]` | `[8]` (no change) |
+| 64 N | `2,4,all` | `[2, 4, 64]` | `[2, 4, 16]` |
+| 128 N | `2,4,16,32,all` | `[2, 4, 16, 32, 128]` | `[2, 4, 16]` |
+| 128 N | `64` | `[64]` | `[16]` |
+
+When the cap actually changes the list, preflight emits a single one-line WARN to stdout so the row labels in the report (e.g. `alltoall-16nodes` instead of `alltoall-128nodes`) are not surprising.
+
+Why a hard cap and why 16:
+
+- **Real-world MoE training rarely dispatches across more than ~8 nodes.** DeepSeek-V3's largest published configuration uses `EP=64` over 8 nodes with per-token dispatch capped at 4 nodes; NVIDIA Megatron-Core's published EP recipes follow the same shape. 16 covers every published configuration with comfortable headroom while staying well clear of the per-node ephemeral-port pressure described in §7.
+- **The cap eliminates the dominant source of `Address already in use` at large scale.** During each `ncclCommInit`, an inter-node alltoall sub-group of `K` nodes opens a near-full mesh of IB OOB sockets per local rank — empirically, peak simultaneous ESTAB sockets per node grow ~linearly with `K` at **~145 sockets per added node** (linear fit `peak_ESTAB ≈ 145·K + 683` over measurements at 24/32/48/56 N; see §7.1.2). Once peak ESTAB approaches the size of the kernel's ephemeral-port pool (default `28 232` ports), the next outgoing `bind()` walks the entire range without finding an allocatable port and fails. Capping the sub-group at 16 holds peak ESTAB at **~3.7 k** — measured directly on a 56 N cluster with the cap active — comfortably under any sensible pool, so the failure cannot occur regardless of cluster size or OS tuning.
+- **Other inter-node tests are unaffected.** `inter-allreduce` (ring/tree, ~`log K` peers per rank) and `inter-ring-p2p` (ring, 2 peers per rank) and `inter-p2p` (pairwise) all open far fewer simultaneous OOB sockets than alltoall and continue to honor `--inter-group-sizes` exactly as written. As a concrete reference point: at 56 nodes, a default-configured `inter-allreduce` peaks at ~1.8 k ESTAB; an `inter-alltoall` over the same 56 nodes peaks at ~8.8 k.
+- **Intentionally not configurable.** This is a known-safe ceiling for the comm shapes preflight is supposed to characterize, not a tuning knob; raising it would re-introduce the very failure mode preflight is meant to *detect* in the cluster, not *cause*.
+
+### 5.3 Ring P2P sizes
+
+| Flag | Default | Applies to |
+|---|---|---|
+| `--ring-p2p-sizes-mb CSV` | `10,20,40,80,160` | `inter-ring-p2p` only. |
+
+```bash
+primus-cli direct -- preflight \
+    --tests inter-ring-p2p \
+    --ring-p2p-sizes-mb 5,20,80
+```
+
+### 5.4 Plotting
+
+| Flag | Effect |
+|---|---|
+| `--plot` | After each perf test, write per-size bandwidth bar charts under `<dump-path>/<test>/` and reference them in the markdown report. |
+
+---
+
+## 6. Reliability knobs
+
+Two knobs that are inert under happy-path conditions but matter at scale or on flaky networks.
+
+### 6.1 `--comm-cleanup-delay-sec FLOAT` (default `2.0`)
+
+Delay (seconds) inserted between destroying NCCL/RCCL process groups and creating new ones. Provides cross-rank synchronization across the destroy → setup transition (so a rank doesn't try to connect to a peer whose listener hasn't finished closing) and gives the kernel a moment to unlink closed-socket bookkeeping. See §7.1 for why this knob is *not* primarily defending against `TIME_WAIT` pressure (which doesn't apply to NCCL's connection pattern) — the actual `Address already in use` defense is the §5.2 inter-alltoall cap.
+
+- Default `2.0` is essentially free and worth keeping as cheap insurance at every cluster size.
+- Set to `0` to disable the sleep entirely (barrier only).
+- Bump to e.g. `5` only on very flaky networks or pathological kernel scheduling; widening `ip_local_port_range` (§7.2) is a more direct fix when the binding constraint is genuinely hit.
+
+```bash
+# Small/medium clusters: defaults are fine. Override only if you see
+# port-reuse races on very flaky networks or unusual kernel TIME_WAIT
+# settings.
+primus-cli slurm srun -N 8 -- preflight --quick --comm-cleanup-delay-sec 5
+```
+
+See §7 ("Running on very large clusters") for the rationale behind why the default works at scale and for the OS-level best-practices that apply to any RDMA workload.
+
+### 6.2 `--dist-timeout-sec INT` (default `120`)
+
+Timeout (seconds) for `torch.distributed.init_process_group`. If init does not complete within this many seconds, preflight writes the info report (when applicable) plus a `Distributed Init` failure section to the markdown report, prints a clear error, and exits `2` — instead of hanging indefinitely.
+
+> Note: §6 used to also document `--comm-cleanup-large-threshold-nodes`, which forced a 60 s drain whenever a destroyed subgroup met or exceeded a size threshold (default 64 nodes). That flag was removed because the underlying failure mode it tried to paper over — peak simultaneous ESTAB exhausting the per-node ephemeral-port pool during a large inter-node alltoall `ncclCommInit` — is now prevented at the source by the §5.2 alltoall cap. The `--comm-cleanup-delay-sec` knob remains, but its primary role is now cross-rank synchronization across the destroy → setup transition rather than draining `TIME_WAIT`; the default 2 s is essentially free and worth keeping as cheap insurance.
+
+```bash
+# Fail fast if rendezvous does not work
+primus-cli direct -- preflight --perf-test --dist-timeout-sec 30
+```
+
+---
+
+## 7. Running on very large clusters (≥ 64 nodes)
+
+Beyond ~64 nodes, two practical concerns dominate that smaller runs never see. Read this section once if you operate clusters in this range; it explains the failure mode, the OS knobs that fix it at the system level, and the recommended preflight invocation patterns.
+
+### 7.1 Why "Address already in use" used to surface at scale
+
+The failure is a **per-node ephemeral port exhaustion** during a single inter-node alltoall `ncclCommInit`, not a chronic accumulation of `TIME_WAIT` sockets across phases. Understanding this distinction is what motivates both the §5.2 cap and the §7.2 OS tuning recommendations.
+
+#### 7.1.1 What actually consumes the per-node ephemeral pool
+
+Linux's outgoing-connection allocator (`__inet_hash_connect()`) does **not** reject ports just because some other socket is in any state on them. It rejects a port only when the new connection's full 4-tuple `(saddr, sport, daddr, dport)` collides with an existing socket's 4-tuple:
+
+- **Live ESTAB sockets** sit on a specific 4-tuple and prevent the kernel from reusing that exact 4-tuple. Each new outgoing connection that lands on a port already holding an ESTAB socket has to walk to the next port. As ESTAB density approaches one-socket-per-port across the entire ephemeral range, the walk takes longer and longer until eventually no port is allocatable — that's the EADDRINUSE.
+- **TIME_WAIT sockets** also live on specific 4-tuples but are governed by `tcp_tw_reuse`. Critically, they only block a new connection when the new connection's *desired* 4-tuple collides with the historical one — i.e. when the new connection is to the *same* `(daddr, dport)` from the *same* `(saddr, sport)`.
+
+For NCCL inter-node OOB traffic, the second case essentially never happens: each `ncclCommInit` connects to **fresh peer OOB listening ports** (the peer chooses an ephemeral listener per setup), so successive comm setups always have different `dport`. TIME_WAIT entries from a previous destroy sit on `(local, P_old, peer, dport_OLD)`; the next setup wants `(local, ?, peer, dport_NEW)`. Even when the new connection happens to land on `sport == P_old`, the dports differ → no 4-tuple collision → the TIME_WAIT entry is invisible to the allocator regardless of `tcp_tw_reuse`.
+
+The practical consequence: under NCCL's connection pattern, **the binding constraint reduces to peak simultaneous ESTAB ≤ size of the ephemeral pool**.
+
+#### 7.1.2 Why inter-node alltoall is the test that exhausts it
+
+Per-node peak ESTAB scales very differently across the inter-node test families:
+
+| Test (56 N, default knobs) | Peer topology per rank | Peak ESTAB per node |
+|---|---|---|
+| `inter-allreduce` | ring / tree, ~`log K` peers | ~1.8 k |
+| `inter-alltoall` | full mesh, `K-1` peers | ~8.8 k |
+
+The empirical scaling of `inter-alltoall` peak ESTAB in the default perf-test sweep. Four measurements at 24/32/48/56 N fit a near-perfect line:
+
+```
+peak_ESTAB(per node) ≈ 145.09 · K + 683
+```
+
+| K (nodes) | Measured peak ESTAB | Fit (145·K + 683) | Source |
+|---|---|---|---|
+| 16 (capped run) | **3 687** | 3 003 | direct measurement — fit slightly under-predicts at the low-N extrapolation |
+| 24 | 4 165 | 4 165 | calibration point |
+| 32 | 5 326 | 5 326 | calibration point |
+| 48 | 7 647 | 7 647 | calibration point |
+| 56 | 8 808 | 8 808 | calibration point |
+| 128 (uncapped, extrapolated) | — | **~19 240** | linear extrapolation |
+
+Two practical reads from this:
+
+- **At 56 N the workload already sits at ~31 % of the default 28 232-port pool**, with ~19 k ports of headroom. That's why every default-pool 56 N run in our experiment succeeded.
+- **An uncapped 128-N inter-alltoall would peak around ~19 k ESTAB** — still inside the default pool but with ~9 k ports of headroom, and *over the cliff* on any cluster that has narrowed `ip_local_port_range`, that runs additional outgoing TCP work concurrently, or that has the source-port allocator's random-walk hit a bad starting offset.
+
+#### 7.1.3 Empirical confirmation: the binding constraint really is `peak_ESTAB ≤ pool_size`
+
+Holding the workload constant (56 N, only inter-alltoall) and varying *only* the per-node ephemeral pool size:
+
+| Pool | Peak ESTAB | Headroom | Result |
+|---|---|---|---|
+| 28 231 (default) | 8 805 | +19 426 | OK |
+| 9 000 | 8 806 | +194 | OK (barely) |
+| 8 000 | (~8 800 expected) | −800 | **FAIL** |
+
+The transition is sharp and right at the predicted boundary. Note that peak `TIME_WAIT` count in the same runs ranged from ~16 k to ~24 k — **far above** the pool size in the 9 k and 8 k rows. If TIME_WAIT count drove the failure, both narrow-pool rows should fail. They don't; only the row where `peak_ESTAB > pool_size` does.
+
+The mechanism is further corroborated by a direct cap experiment on the same 56 N cluster, default pool: re-running the inter-node alltoall test with the per-group node count capped at 16 (matching the §5.2 cap) yields peak ESTAB **3 687** — well under both the default 28 k pool and the 9 k / 8 k stressed pools — and the run completes cleanly. The cap acts directly on the binding constraint by holding peak ESTAB low enough that the pool cannot be exhausted, regardless of cluster size.
+
+#### 7.1.4 The cap closes the failure mode at the source
+
+The §5.2 inter-node alltoall cap (16 nodes max) holds peak ESTAB at **~3.7 k** regardless of cluster size — measured directly in §7.1.3, comfortably under the default 28 k pool (~13 % utilization) and still safe at half-default pool widths. The OS-level tunings in §7.2 remain useful general hygiene for any RDMA / multi-NIC workload, but a default preflight invocation no longer needs them to avoid `Address already in use`.
+
+This is **not** a real training failure mode in any case — production training jobs create their TP / DP / PP / EP communicators *once* at startup and reuse them, and real-world MoE training rarely dispatches across more than ~8 nodes. The preflight tool is the only thing that builds many large communicators in a short window, which is why the issue was preflight-specific to begin with.
+
+### 7.2 OS-level tuning (best-practice for any large-cluster node)
+
+Given §7.1's mechanism (binding constraint = peak ESTAB ≤ pool size), the OS knobs split cleanly into "directly relevant" and "general hygiene":
+
+```bash
+# 1) DIRECTLY RELEVANT: widen the per-node ephemeral pool.
+#    The default range is 32768-60999 (~28 k ports). Widening it
+#    to 1024-65535 (~64 k ports) more than doubles the headroom
+#    for peak simultaneous ESTAB — and that is the only thing that
+#    can produce EADDRINUSE under the NCCL connection pattern.
+sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535"
+
+# 2) GENERAL HYGIENE: allow TIME_WAIT reuse for outgoing connections.
+#    For the NCCL inter-node OOB pattern this is largely a no-op
+#    (each ncclCommInit picks fresh peer destination ports, so
+#    historical TIME_WAIT 4-tuples never collide with what the
+#    next setup wants). It is still recommended for any host that
+#    runs additional outgoing TCP workloads where the SAME
+#    (daddr, dport) is hit repeatedly from the same source IP --
+#    the textbook scenario the kernel doc is written around.
+sudo sysctl -w net.ipv4.tcp_tw_reuse=1
+
+# Persist across reboots:
+cat <<'EOF' | sudo tee /etc/sysctl.d/99-large-cluster.conf
+net.ipv4.ip_local_port_range = 1024 65535
+net.ipv4.tcp_tw_reuse = 1
+EOF
+sudo sysctl --system
+```
+
+If you only have time to set one of these, pick `ip_local_port_range`. With the §5.2 inter-alltoall cap holding peak ESTAB at ~3.7 k even at 1024 N, neither knob is *required* for preflight — but the wider port range is the right insurance for any host that might also run other outgoing TCP traffic concurrently.
+
+Note on `tcp_tw_reuse=2`: as of Linux 4.12, value `2` enables TIME_WAIT reuse **only for loopback** (`127.0.0.0/8`, `::1`). For inter-node IB OOB connections this is equivalent to `tcp_tw_reuse=0`. The fact that several of our 56 N runs succeeded under `tcp_tw_reuse=2` with `peak_WAIT > 16 k` is direct evidence that `TIME_WAIT` *count* doesn't gate inter-node bind() — only `peak_ESTAB > pool_size` does (see §7.1.3).
+
+### 7.3 In-tool defenses
+
+Preflight has two complementary defenses:
+
+1. **The §5.2 inter-node alltoall cap (16 nodes max).** *This is the actual fix.* Regardless of cluster size or `--inter-group-sizes`, the alltoall test never builds a sub-group large enough to push peak ESTAB anywhere near the per-node ephemeral-port pool. This eliminates the historical EADDRINUSE failure mode by construction (see §7.1).
+2. **The §6.1 per-phase cleanup delay (`--comm-cleanup-delay-sec`, default `2.0`).** A global barrier + sleep inserted after every comm destroy. Its primary job is **cross-rank synchronization** across the destroy → setup transition (so a rank doesn't try to connect to a peer whose listener hasn't finished closing) and giving the kernel a moment to unlink closed-socket bookkeeping. It is *not* protecting against `TIME_WAIT` 4-tuple collisions — those don't occur in NCCL's connection pattern (see §7.1.1). The default 2 s is essentially free and worth keeping as cheap insurance.
+
+Together, a default preflight invocation is safe at every cluster size we test up to 1024 nodes. The only situation where you would consider raising `--comm-cleanup-delay-sec` is on a network with unusually narrow ephemeral-port ranges or pathological kernel scheduling — and even there, widening `ip_local_port_range` (§7.2) is the cleaner fix because it directly addresses the binding constraint.
+
+### 7.4 Recommended invocation patterns at very large scale
+
+For clusters at or beyond ~128 nodes, the most reliable and most informative way to use preflight is still to **split the run into one test family per invocation** rather than one big run — not for `Address already in use` reasons (the §7.3 defenses handle that), but because it keeps wall-clock per invocation small and makes it trivial to identify which specific comm shape is degraded if a metric looks off.
+
+```bash
+# 1) GPU + intra-node fabric first (cheap, no inter-node OOB churn).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests gemm,intra-allreduce,intra-alltoall
+
+# 2) Inter-node DP-style collectives, all-nodes group only.
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-allreduce \
+    --inter-group-sizes all
+# Note: --inter-group-sizes all is honored here for inter-allreduce.
+# For inter-alltoall it would be clamped to 16 (see §5.2).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-alltoall \
+    --inter-group-sizes all
+
+# 3) Inter-node PP-style ring P2P (the test that benefits most from
+#    isolation — it's the closest match to what real pipeline-parallel
+#    training actually exercises).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-ring-p2p
+
+# 4) Optional: pairwise inter-node P2P scan (useful for finding a
+#    single bad link, slower because it walks many pairs).
+primus-cli slurm srun -N 128 -- preflight \
+    --tests inter-p2p
+```
+
+Each invocation:
+
+- Tears down its own `WORLD` at exit, so the next invocation starts with a fresh per-node port pool.
+- Touches only one `--tests` value, so you get a per-test wall clock and can re-run a single phase if its numbers look off without paying for the others.
+- Carries its own report under `--report-file-name` (or the wrapper-generated default), which makes archiving and comparison across runs straightforward.
+
+### 7.5 Decision flow
+
+| Cluster size | Recommended approach |
+|---|---|
+| ≤ 32 nodes | Single command, default knobs. Nothing special. |
+| 33-127 nodes | Single command, default knobs. |
+| ≥ 128 nodes | Single command works with default knobs (the §5.2 alltoall cap holds peak ESTAB at ~3.7 k, ≈ 7.6× under the default 28 k ephemeral pool). Splitting per `--tests` token as in §7.4 is recommended for diagnostic clarity rather than for safety. Widening `ip_local_port_range` (§7.2) is best-practice for any RDMA workload but no longer required for preflight specifically. |
+| ≥ 256 nodes | Always split as in §7.4 — keeps every invocation snappy and makes regressions much easier to localize. |
+
+---
+
+## 8. Reporting
+
+| Flag | Default | Effect |
+|---|---|---|
+| `--dump-path DIR` | `output/preflight` | Output directory for reports + plots. |
+| `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. |
+| `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. |
+
+Output files:
+
+| File | Produced when | Notes |
+|---|---|---|
+| `<name>.md` / `<name>.pdf` | Info-only mode, or default mode | Info report. |
+| `<name>_perf.md` / `<name>_perf.pdf` | Perf-only mode, or default mode | Perf report (GEMM + comm). |
+
+Only **rank 0** writes the report.
+
+### Perf report layout
+
+A `<name>_perf.md` produced by a default run contains, in order:
+
+1. (Optional) `> Note:` line listing dropped info selectors.
+2. `# Nodes` legend — `Node N → Hostname` table, used by every subsequent table to keep host columns compact.
+3. `=======IB Bandwidth roofline (GB/s)=======` — bandwidth of the first IB device on Node 0.
+4. Per enabled test, in this order: `gemm`, `intra-comm`, `inter-comm`, `inter-p2p`, `inter-ring-p2p`. Each section has a configuration line, a results table (Node / Rank / hostname / per-size GB/s), optional plots, and a per-rank wall-clock summary.
+5. `[Primus:Preflight] <test> done in <T>s` lines on stdout for at-a-glance progress on the launching shell.
+
+---
+
+## 9. Backward-compat aliases
+
+| Flag | Equivalent | Notes |
+|---|---|---|
+| `--check-host`, `--check-gpu`, `--check-network` | `--host`, `--gpu`, `--network` | Same behavior. Keep working for older scripts. |
+| `--no-split-nodes-subgroup` | `--inter-group-sizes all` **and** drops `inter-p2p` | Pre-`--tests`/`--inter-group-sizes` alias. Use the new flags in new scripts. |
+
+---
+
+## 10. Comparison with node-smoke
+
+| Aspect | `node-smoke` | `preflight` |
+|---|---|---|
+| Rendezvous | None — every node independent | Global `torch.distributed` |
+| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests |
+| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) |
+| Inter-node bandwidth matrix | Not tested (intentionally) | Yes (allreduce/alltoall/p2p/ring-p2p) |
+| Drift detection | Yes (versions, NIC firmware, port count) | No |
+| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only |
+| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF |
+
+**Recommended workflow**: run `node-smoke` first to exclude broken nodes, then run `preflight` on the surviving set to get cross-node bandwidth measurements. See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) §3 ("Quick start") for the integration commands.
+
+---
+
+## 11. Validation & error handling
+
+Preflight resolves the perf config **before** any distributed rendezvous. This means typos and bad sizes/group-sizes fail in seconds, not after a 120s NCCL init:
+
+```text
+[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'.
+[Primus:Preflight] ERROR: invalid perf config:
+    --intra-group-sizes: [3] do not divide LOCAL_WORLD_SIZE=8
+[Primus:Preflight] ERROR: invalid perf config: --comm-sizes-mb: '0' must be positive
+```
+
+In info-only mode, perf-only tuning knobs trigger a single warning so you notice them but they don't abort:
+
+```text
+[Primus:Preflight] WARN: --comm-sizes-mb,--intra-group-sizes have no effect
+in info-only mode (no --perf-test/--tests/--quick).
+```
+
+In default mode where info selectors are dropped because perf intent was set, the preserved warning is also written into the perf report header:
+
+```text
+> Note: info selectors --host were dropped because perf mode
+> (--perf-test/--tests/--quick) takes precedence. Run them in a separate
+> invocation if you want both reports.
+```
+
+---
+
+## 12. Operational tips
+
+- **For multi-node runs, always use `primus-cli slurm` or `primus-cli direct` under `srun`** so distributed environment variables (`NNODES` / `NODE_RANK` / `MASTER_ADDR`) are set correctly.
+- **Insufficient CPU cores cause >30x perf slowdowns** — pass `srun -c <cores-per-node>` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`.
+- **For a quick environment snapshot**, prefer `--host --gpu --network` (no rendezvous, finishes in seconds even on broken networks).
+- **Between each communication test phase**, preflight performs a global barrier + `--comm-cleanup-delay-sec` sleep (default 2 s) for cross-rank sync across the destroy → setup transition. The default works at every cluster size we test up to 1024 nodes because the inter-node alltoall sub-group is internally capped at 16 (see §5.2), which holds peak simultaneous ESTAB sockets per node well under the kernel's ephemeral-port pool — the only constraint that actually produces `Address already in use` under NCCL's connection pattern (see §7.1). Widening `ip_local_port_range` (§7.2) is best-practice for any RDMA workload.
+- **For pre-launch screening of a large cluster**, the recommended sequence is:
+  1. `node-smoke` to prune broken nodes (`failing_nodes.txt`).
+  2. `preflight --quick` on the surviving nodes for the perf sanity numbers.
+  3. `preflight` (full) on the same set if the `--quick` numbers raise a flag.
+
+---
+
+## 13. Running preflight without a container
+
+If you cannot (or prefer not to) use a container, see [`preflight-direct.md`](./preflight-direct.md) for the step-by-step `runner/primus-cli direct -- preflight ...` walkthrough — Python virtual-environment setup, SLURM invocation patterns, NCCL configuration for Broadcom and Pensando (AINIC) clusters, and many configurable-knob examples.
+
+---
+
+## 14. See also
+
+- [`preflight-direct.md`](./preflight-direct.md) — quick-start guide for `primus-cli direct -- preflight` (no container).
+- [`node-smoke.md`](./node-smoke.md) — full reference for the per-node smoke test.
+- [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) — short quick-start for the smoke test.
+- [`runner/primus-cli-direct.sh`](../runner/primus-cli-direct.sh) — non-container launcher (`primus-cli direct` dispatches here).
+- [`primus/tools/preflight/`](../primus/tools/preflight/) — implementation.
+- [`primus/tools/preflight/preflight_args.py`](../primus/tools/preflight/preflight_args.py) — canonical CLI definition (single source of truth for flags + defaults).
diff --git a/docs/projection.md b/docs_deprecated/projection.md
similarity index 99%
rename from docs/projection.md
rename to docs_deprecated/projection.md
index 02596c6b3..3588c5efb 100644
--- a/docs/projection.md
+++ b/docs_deprecated/projection.md
@@ -626,7 +626,7 @@ Splitting the residual into these two terms is more robust than the older single
 
 With `--memory-mode both`, a side-by-side table shows the simulate vs. benchmark per-component deltas. A small delta (within ~10–20%) means the analytical model and the residual term are well-calibrated; a large positive delta (simulate ≫ benchmark) usually means simulate is over-counting unsharded components, while a large negative delta means the bench captured overhead the analytical model under-estimates.
 
-> The benchmark-based memory projection is what the [Tuning Agent](tuning_agent.md) uses for OOM-accurate feasibility filtering when a GPU is available, so its `tokens/s/GPU` rankings never include configs that would OOM on the real cluster.
+> The benchmark-based memory projection is what the [Tuning Agent](../docs/02-user-guide/tuning-agent.md) uses for OOM-accurate feasibility filtering when a GPU is available, so its `tokens/s/GPU` rankings never include configs that would OOM on the real cluster.
 
 ---
 
diff --git a/docs/quickstart.md b/docs_deprecated/quickstart.md
similarity index 96%
rename from docs/quickstart.md
rename to docs_deprecated/quickstart.md
index 4d6dd1e6c..6aa149624 100644
--- a/docs/quickstart.md
+++ b/docs_deprecated/quickstart.md
@@ -87,7 +87,7 @@ primus-cli [options] <mode> [mode-args] -- [command]
 
 **Learn More:**
 - [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) - Complete reference
-- [CLI Architecture](./cli/CLI-ARCHITECTURE.md) - Design deep dive
+- [CLI Architecture](../docs/06-developer-guide/cli-architecture.md) - Design deep dive
 - [Configuration Guide](./configuration.md) - YAML configuration
 - [Examples](../examples/README.md) - Real-world templates
 
diff --git a/docs/tech_blogs/primus_cli_unified_entry_rocm.md b/docs_deprecated/tech_blogs/primus_cli_unified_entry_rocm.md
similarity index 100%
rename from docs/tech_blogs/primus_cli_unified_entry_rocm.md
rename to docs_deprecated/tech_blogs/primus_cli_unified_entry_rocm.md
diff --git a/docs/tech_blogs/primus_pipeline/imgs/actual-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/actual-perf.png
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/imgs/actual-perf.png
rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/actual-perf.png
diff --git a/docs/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png
rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png
diff --git a/docs/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png
rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png
diff --git a/docs/tech_blogs/primus_pipeline/imgs/simulation.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/simulation.png
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/imgs/simulation.png
rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/simulation.png
diff --git a/docs/tech_blogs/primus_pipeline/imgs/simulator_shell.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/simulator_shell.png
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/imgs/simulator_shell.png
rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/simulator_shell.png
diff --git a/docs/tech_blogs/primus_pipeline/primus_pipeline.md b/docs_deprecated/tech_blogs/primus_pipeline/primus_pipeline.md
similarity index 100%
rename from docs/tech_blogs/primus_pipeline/primus_pipeline.md
rename to docs_deprecated/tech_blogs/primus_pipeline/primus_pipeline.md
diff --git a/docs/tech_blogs/projection/projection.md b/docs_deprecated/tech_blogs/projection/projection.md
similarity index 100%
rename from docs/tech_blogs/projection/projection.md
rename to docs_deprecated/tech_blogs/projection/projection.md
diff --git a/examples/README.md b/examples/README.md
index ce0e36460..5723b1168 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -36,10 +36,11 @@ It supports both **single-node** and **multi-node** training, and includes optio
 
 Primus supports multiple backends.
 
-| Backend    | Description                                                  |
-| ---------- | ------------------------------------------------------------ |
-| Megatron   | Open-source framework for large-scale transformer training   |
-| TorchTitan | PyTorch-compatible framework developed for training at scale |
+| Backend        | Description                                                  |
+| -------------- | ------------------------------------------------------------ |
+| Megatron       | Open-source framework for large-scale transformer training   |
+| TorchTitan     | PyTorch-compatible framework developed for training at scale |
+| NeMo AutoModel | NVIDIA-NeMo AutoModel (diffusion: Wan 2.2 T2V); `third_party/Automodel` submodule, installed editable on first run |
 
 
 ## 🖥️ Single Node Training
@@ -49,7 +50,7 @@ We recommend using the official [rocm/megatron-lm Docker image](https://hub.dock
 
 ```bash
 # Pull the latest Docker image
-docker pull docker.io/rocm/primus:v26.3
+docker pull docker.io/rocm/primus:v26.4
 
 ```
 
@@ -126,7 +127,7 @@ Multi-node training is launched via **SLURM**.
 Specify the number of nodes and the model config:
 
 ```bash
-export DOCKER_IMAGE="docker.io/rocm/primus:v26.3"
+export DOCKER_IMAGE="docker.io/rocm/primus:v26.4"
 export NNODES=8
 
 # Example for megatron llama3.1_8B
@@ -208,6 +209,13 @@ The following models are supported out of the box via provided configuration fil
 | Mixtral-8x7B-v0.1 | [mistralai/Mixtral-8x7B-v0.1 ](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)           | [mixtral_8x7B_v0.1-BF16-pretrain.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml)           | |
 | Mixtral-8x22B-v0.1 | [mistralai/Mixtral-8x22B-v0.1 ](https://huggingface.co/mistralai/Mixtral-8x22B-v0.1)           | [mixtral_8x22B_v0.1-BF16-pretrain.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml)           | |
 
+### Diffusion Models
+
+- **Flux** - Flow-based diffusion model for text-to-image generation
+  - Training guide: [examples/megatron/diffusion/README.md](megatron/diffusion/README.md) (Flux 535M and 12B)
+  - Architecture & developer docs: [docs/04-technical-guides/diffusion-models/README.md](../docs/04-technical-guides/diffusion-models/README.md)
+  - FP8 training: [docs/04-technical-guides/diffusion-models/fp8_training.md](../docs/04-technical-guides/diffusion-models/fp8_training.md)
+
 ---
 
 ### 🏃‍♂️ How to Run a Supported Model
@@ -285,7 +293,7 @@ When using the `create` command to start a new training workload, the following
 | `--gpu`        | Number of GPUs                                       | 8                                        |
 | `--exp`        | Path to experiment (training config) file (required) | —                                        |
 | `--data_path`  | Path to training data                                | —                                        |
-| `--image`      | Docker image to use                                  | `docker.io/rocm/primus:v26.3` |
+| `--image`      | Docker image to use                                  | `docker.io/rocm/primus:v26.4` |
 | `--hf_token`   | HuggingFace token                                    | Read from env var `HF_TOKEN`             |
 | `--workspace`  | Workspace name                                       | `primus-safe-pretrain`                   |
 | `--nodelist`   | Comma-separated list of node hostnames to run on     | —                                        |
diff --git a/examples/deepseek-v4/benchmark/bench_v4_attention.py b/examples/deepseek-v4/benchmark/bench_v4_attention.py
new file mode 100644
index 000000000..45cd76cde
--- /dev/null
+++ b/examples/deepseek-v4/benchmark/bench_v4_attention.py
@@ -0,0 +1,592 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 attention kernel benchmark — all backends in one table.
+
+Unified benchmark of the **forward** and **backward** V4 attention kernels for
+every backend, so there is a single place to compare them (one row per backend,
+``ms | TFLOP/s`` cells; TFLOP/s over the useful work so all rows are comparable):
+
+* ``triton``    — PRODUCTION Triton (separate K/V; pool/SWA/HCA launchers)
+* ``gluon``     — fused single-latent (K==V) sparse-MLA, hand-tuned gfx950
+* ``triton_v2`` — fused single-latent sparse-MLA in plain Triton (tl.dot / MFMA)
+* ``flydsl_v1`` — fused single-latent sparse-MLA in native FlyDSL MFMA (fwd + bwd)
+* ``turbo_flydsl`` — extracted Primus-Turbo sparse_mla_v2 native FlyDSL MFMA
+* ``_turbo_flydsl`` — the INTEGRATED in-tree Primus-Turbo backend via the turbo API
+  (``primus_turbo.flydsl.attention``); the ``turbo`` model backend
+  (``use_v4_attention_backend`` / ``use_v4_csa_attention_backend = "turbo"``)
+
+The legacy ``_flydsl_v0_deprecated`` gathered-CSA backend (scalarized GEMV) is
+NOT benchmarked: it has known correctness issues and depends on the
+/workspace/FlyDSL-amd source tree.
+
+The fused ``gluon`` / ``triton_v2`` / ``flydsl_v1`` backends share ONE kernel-pair
+API and are timed on IDENTICAL V4-form inputs (zero rope pad + [local ++ pool]
+kv + [SWA window ++ pool] topk). Each backend is guarded; unavailable ones are
+simply skipped.
+
+Benchmarks for the two production model sizes across all three layer kinds:
+
+* ``cr=0``   — dense / sliding-window (SWA-only) attention
+* ``cr=4``   — CSA (local SWA + sparse top-k from the compressed pool)
+* ``cr=128`` — HCA (local SWA + full compressed pool, joint softmax)
+
+at the real attention shapes (``seq_len=4096``, ``mbs=1``, bf16, sink on,
+``swa_window=128``):
+
+* V4-Flash:  H=64,  head_dim=512, index_topk=512
+* V4-Pro:    H=128, head_dim=512, index_topk=1024
+
+Backend detail:
+
+* ``triton`` (ALL crs): PRODUCTION launchers used by ``DeepseekV4Attention`` —
+    cr=0/128 ``_launch_v4_attention_fwd``/``_bwd`` (dense/HCA); cr=4
+    ``_launch_v4_csa_attention_pool_fwd``/``_pool_bwd`` (split FWD + segreduce
+    BWD in-kernel gather; NOT the legacy gathered API, ~30-260x slower).
+* ``gluon`` / ``triton_v2`` / ``flydsl_v1`` (ALL crs): fused single-latent
+    sparse-MLA (``sparse_mla_{fwd,bwd}_v4_*``); the layer kind is just a
+    different TOPK (cr=0: swa 128; cr=4: 128+sparse; cr=128: 128+pool).
+
+Effective TFLOP/s uses ``2*T*H*TOPK*(D_V+D_V)`` (useful work over head_dim=512),
+BWD = 2.5x FWD, the SAME formula for every backend so rows are directly
+comparable (the fused backends' zero-rope-pad overhead shows up in ms).
+
+Run inside the dev container (gfx950 / MI355X):
+
+    python examples/deepseek-v4/benchmark/bench_v4_attention.py
+    python examples/deepseek-v4/benchmark/bench_v4_attention.py --variant pro --cr 4
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import sys
+from typing import Tuple
+
+import torch
+
+# NOTE: the legacy `_flydsl_v0_deprecated` CSA backend (scalarized GEMV) is
+# intentionally NOT imported/benchmarked here — it has known correctness issues
+# and depends on the /workspace/FlyDSL-amd source tree. The native FlyDSL MFMA
+# backend is `_flydsl_v1` (benchmarked below as `flydsl_v1`).
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (
+    _launch_v4_attention_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import (
+    _launch_v4_attention_fwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_bwd import (
+    _launch_v4_csa_attention_pool_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_fwd import (
+    _launch_v4_csa_attention_pool_fwd,
+)
+
+# Fused single-latent (K==V) sparse-MLA backends. All share ONE kernel-pair
+# API (fwd(q, kv, topk, attn_sink, kv_lora_rank, scale) -> (o, lse); bwd ->
+# (dq, dkv, d_sink)) so they are timed on identical V4-form inputs. Each is
+# guarded so the benchmark still runs where a backend is unavailable.
+_SPARSE_MLA_BACKENDS = {}  # name -> (fwd_fn, bwd_fn)
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_dsa import (
+        sparse_mla_bwd_v4_gluon,
+        sparse_mla_fwd_v4_gluon,
+    )
+
+    _SPARSE_MLA_BACKENDS["gluon"] = (sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon)
+except Exception:  # noqa: BLE001
+    pass
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v2 import (
+        sparse_mla_bwd_v4_triton,
+        sparse_mla_fwd_v4_triton,
+    )
+
+    _SPARSE_MLA_BACKENDS["triton_v2"] = (sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton)
+except Exception:  # noqa: BLE001
+    pass
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v1 import (
+        sparse_mla_bwd_v4_flydsl,
+        sparse_mla_fwd_v4_flydsl,
+    )
+
+    _SPARSE_MLA_BACKENDS["flydsl_v1"] = (sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl)
+except Exception:  # noqa: BLE001
+    pass
+# _turbo_flydsl: the INTEGRATED Primus-Turbo sparse-MLA backend, via the turbo API
+# (primus_turbo.flydsl.attention). Not an agent/workspace extraction — this is the
+# in-tree `_turbo_flydsl` backend (enabled in the model via
+# use_v4_attention_backend / use_v4_csa_attention_backend = "turbo"). Requires an
+# installed primus_turbo carrying primus_turbo.flydsl.attention.
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._turbo_flydsl import (
+        sparse_mla_bwd_v4_turbo_flydsl,
+        sparse_mla_fwd_v4_turbo_flydsl,
+    )
+
+    _SPARSE_MLA_BACKENDS["_turbo_flydsl"] = (
+        sparse_mla_fwd_v4_turbo_flydsl,
+        sparse_mla_bwd_v4_turbo_flydsl,
+    )
+except Exception:  # noqa: BLE001
+    pass
+# gluon_v2: our aiter-gluon-inspired backend (guarded; skipped until present).
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2 import (
+        sparse_mla_bwd_v4_gluon_v2,
+        sparse_mla_fwd_v4_gluon_v2,
+    )
+
+    _SPARSE_MLA_BACKENDS["gluon_v2"] = (sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2)
+except Exception:  # noqa: BLE001
+    pass
+# gluon_v3: active gfx950 optimization campaign backend.
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import (
+        sparse_mla_bwd_v4_gluon_v3,
+        sparse_mla_fwd_v4_gluon_v3,
+    )
+
+    _SPARSE_MLA_BACKENDS["gluon_v3"] = (sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3)
+except Exception:  # noqa: BLE001
+    pass
+# aiter PR#3833 gluon sparse-MLA prefill (fwd-only) — OPTIONAL reference for
+# comparison. Lives under agent/workspace; set PRIMUS_V4_AITER_DIR to override.
+# Guarded so the benchmark is unchanged when the extracted adapter is absent.
+try:
+    _aiter_dir = os.environ.get(
+        "PRIMUS_V4_AITER_DIR",
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+            "..",
+            "agent",
+            "workspace",
+            "aiter_dsv4_prefill_20260706",
+        ),
+    )
+    _aiter_dir = os.path.abspath(_aiter_dir)
+    if _aiter_dir not in sys.path:
+        sys.path.insert(0, _aiter_dir)
+    from aiter_v4_adapter import sparse_mla_fwd_v4_aiter
+
+    _SPARSE_MLA_BACKENDS["aiter_gluon"] = (sparse_mla_fwd_v4_aiter, None)
+except Exception:  # noqa: BLE001
+    pass
+
+_GLUON_AVAIL = "gluon" in _SPARSE_MLA_BACKENDS
+
+_VARIANTS = {
+    "flash": dict(H=64, index_topk=512),
+    "pro": dict(H=128, index_topk=1024),
+}
+_HEAD_DIM = 512
+_SWA_WINDOW = 128
+_ROPE_DIM = 64  # sparse-MLA d_qk = kv_lora_rank (=_HEAD_DIM) + rope
+_CR_NAME = {0: "SWA/dense", 4: "CSA", 128: "HCA"}
+
+
+# ---------------------------------------------------------------------------
+# Input builders
+# ---------------------------------------------------------------------------
+
+
+def _common_inputs(B: int, H: int, S: int, D: int, seed: int = 0):
+    """q + MQA single-latent K/V (full + [.,1,.,.]) + sink + dout, all bf16."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    dev, dt = "cuda", torch.bfloat16
+    q = torch.randn(B, H, S, D, generator=g, device=dev, dtype=dt)
+    k_mqa = torch.randn(B, 1, S, D, generator=g, device=dev, dtype=dt)
+    v_mqa = torch.randn(B, 1, S, D, generator=g, device=dev, dtype=dt)
+    k_full = k_mqa.expand(B, H, S, D).contiguous()
+    v_full = v_mqa.expand(B, H, S, D).contiguous()
+    sink = torch.randn(H, generator=g, device=dev, dtype=torch.float32) * 0.1
+    dout = torch.randn(B, H, S, D, generator=g, device=dev, dtype=dt)
+    return q, k_mqa, v_mqa, k_full, v_full, sink, dout
+
+
+def _csa_sparse(B: int, H: int, S: int, D: int, K: int, P: int, g: torch.Generator):
+    """CSA sparse branch: compressed pool + per-query top-K indices, plus the
+    pre-gathered equivalent (FlyDSL) sharing the same valid/invalid pattern."""
+    dev, dt = "cuda", torch.bfloat16
+    pool = torch.randn(B, P, D, generator=g, device=dev, dtype=dt)
+    valid = torch.rand(B, S, K, generator=g, device=dev) > 0.25
+    topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device=dev, dtype=torch.int32)
+    topk_idxs = torch.where(valid, topk_idxs, torch.full_like(topk_idxs, -1))
+    safe = topk_idxs.clamp(min=0).long()
+    gathered = torch.gather(
+        pool.unsqueeze(1).expand(B, S, P, D), dim=2, index=safe.unsqueeze(-1).expand(B, S, K, D)
+    )
+    gathered = gathered * valid.unsqueeze(-1).to(dt)
+    sparse_mask = torch.where(
+        valid,
+        torch.zeros((), dtype=dt, device=dev),
+        torch.tensor(float("-inf"), dtype=dt, device=dev),
+    )
+    return pool, topk_idxs, gathered, sparse_mask
+
+
+def _hca_mask(S: int, P: int, ratio: int, device, dtype):
+    """HCA pool-only additive causal mask [S, P]: pool slot s visible to query t
+    iff (s+1)*ratio - 1 <= t (matches DeepseekV4Attention._hca_extra_mask)."""
+    t = torch.arange(S, device=device).unsqueeze(1)
+    s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1
+    return torch.where(s_end <= t, 0.0, float("-inf")).to(dtype)
+
+
+def _build_gluon_v4form(*, cr: int, H: int, S: int, D: int, K: int, P: int, W: int, seed: int = 0):
+    """Gluon inputs in the **V4 form** (matches the training adapter):
+
+    The 512 V4 latent (RoPE baked in-place) is the gluon "lora" with a zero rope
+    pad (kernel needs D_ROPE>0); kv = [local latent ++ pool] and topk = [SWA
+    window ++ (cr=4: sparse pool top-k | cr=128: full causal pool | cr=0: none)].
+    ``topk`` is padded to a multiple of 64 so the gluon bwd dKV tiling (TILE_K=64)
+    is valid (notably HCA: 128+32=160 -> 192). scale = 1/sqrt(D) (V4, over 512).
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    dev, dt = "cuda", torch.bfloat16
+    latent = torch.randn(S, D, generator=g, device=dev, dtype=dt)
+    q512 = torch.randn(S, H, D, generator=g, device=dev, dtype=dt)
+    z_q = torch.zeros(S, H, _ROPE_DIM, device=dev, dtype=dt)
+    q_g = torch.cat([q512, z_q], dim=-1).contiguous()  # [S, H, D+rope_pad]
+    sink = torch.randn(H, generator=g, device=dev, dtype=torch.float32) * 0.1
+    do = torch.randn(S, H, D, generator=g, device=dev, dtype=dt)
+
+    ti = torch.arange(S, device=dev).view(S, 1)
+    win = ti - W + 1 + torch.arange(W, device=dev).view(1, W)  # [S, W] local token idx
+    win = torch.where(win >= 0, win, torch.full_like(win, -1))
+
+    if cr == 0:
+        kv512 = latent.unsqueeze(1)  # [S, 1, D]
+        topk = win
+    else:
+        pool = torch.randn(P, D, generator=g, device=dev, dtype=dt)
+        kv512 = torch.cat([latent, pool], dim=0).unsqueeze(1)  # [S+P, 1, D]
+        if cr == 4:
+            sp = torch.randint(0, P, (S, K), generator=g, device=dev)
+            pool_topk = S + sp
+        else:  # cr == 128: HCA full causal pool
+            ps = torch.arange(P, device=dev).view(1, P)
+            vis = ((ps + 1) * cr - 1) <= ti  # [S, P]
+            pool_topk = torch.where(vis, S + ps, torch.full_like(ps.expand(S, P), -1))
+        topk = torch.cat([win, pool_topk], dim=1)
+
+    tk = topk.shape[1]
+    pad = ((tk + 63) // 64) * 64 - tk
+    if pad > 0:
+        topk = torch.cat([topk, torch.full((S, pad), -1, device=dev, dtype=topk.dtype)], dim=1)
+    topk_g = topk.to(torch.int32).contiguous()
+
+    z_kv = torch.zeros(kv512.shape[0], 1, _ROPE_DIM, device=dev, dtype=dt)
+    kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous()
+    return q_g, kv_g, topk_g, sink, do
+
+
+# ---------------------------------------------------------------------------
+# Timing helpers
+# ---------------------------------------------------------------------------
+
+
+def _time_ms(fn, *, warmup: int, iters: int) -> Tuple[float, float]:
+    """Return (median_ms, mean_ms) over ``iters`` timed launches."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for _ in range(iters):
+        start.record()
+        fn()
+        end.record()
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    times.sort()
+    return times[len(times) // 2], sum(times) / len(times)
+
+
+def _safe_time(fn, *, warmup: int, iters: int):
+    """Time ``fn``; on failure return (None, short_error_string)."""
+    if fn is None:
+        return None, None
+    try:
+        med, _ = _time_ms(fn, warmup=warmup, iters=iters)
+        return med, None
+    except Exception as exc:  # noqa: BLE001 - benchmark must survive one cell failing
+        msg = str(exc).strip().splitlines()[-1] if str(exc).strip() else type(exc).__name__
+        return None, f"{type(exc).__name__}: {msg}"
+
+
+def _call_fwd(fn):
+    """Call a fwd launcher once for the bwd's (out, lse); (None, None) on failure."""
+    if fn is None:
+        return None, None
+    try:
+        return fn()
+    except Exception:  # noqa: BLE001
+        return None, None
+
+
+def _tflops(flops: float, med_ms) -> float:
+    if med_ms is None or med_ms <= 0:
+        return float("nan")
+    return flops / (med_ms * 1e-3) / 1e12
+
+
+def _cell(med, flops) -> str:
+    if med is None:
+        return f"{'FAIL':>18s}"
+    return f"{med:9.2f} | {_tflops(flops, med):6.1f}"
+
+
+# ---------------------------------------------------------------------------
+# Per-(variant, cr) benchmark
+# ---------------------------------------------------------------------------
+
+
+def _bench_cr(variant: str, cr: int, *, B: int, S: int, warmup: int, iters: int):
+    cfg = _VARIANTS[variant]
+    H = cfg["H"]
+    D = _HEAD_DIM
+    scale = 1.0 / math.sqrt(D)
+    g = torch.Generator(device="cuda").manual_seed(11)
+    q, k_mqa, v_mqa, k_full, v_full, sink, dout = _common_inputs(B, H, S, D)
+
+    fwd_triton = bwd_triton = None
+
+    if cr == 0:
+        topk_eff = _SWA_WINDOW
+        glu_K, glu_P = 0, 0
+        extra = ""
+
+        def fwd_triton():
+            return _launch_v4_attention_fwd(
+                q,
+                k_full,
+                v_full,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                additive_mask=None,
+                scale=scale,
+                hca_local_seqlen=0,
+            )
+
+        out_t, lse_t = _call_fwd(fwd_triton)
+
+        def bwd_triton():
+            return _launch_v4_attention_bwd(
+                q,
+                k_full,
+                v_full,
+                out_t,
+                dout,
+                lse_t,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                additive_mask=None,
+                scale=scale,
+                hca_local_seqlen=0,
+            )
+
+    elif cr == 4:
+        P = max(S // 4, 1)
+        K = min(cfg["index_topk"], P)
+        topk_eff = _SWA_WINDOW + K
+        glu_K, glu_P = K, P
+        extra = f" K_topk={K} P={P}"
+        pool, topk_idxs, gathered, sparse_mask = _csa_sparse(B, H, S, D, K, P, g)
+
+        def fwd_triton():
+            return _launch_v4_csa_attention_pool_fwd(
+                q,
+                k_full,
+                v_full,
+                pool,
+                topk_idxs,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                scale=scale,
+            )
+
+        out_t, lse_t = _call_fwd(fwd_triton)
+
+        def bwd_triton():
+            return _launch_v4_csa_attention_pool_bwd(
+                q,
+                k_full,
+                v_full,
+                pool,
+                topk_idxs,
+                out_t,
+                dout,
+                lse_t,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                scale=scale,
+            )
+
+    elif cr == 128:
+        P = max(S // cr, 1)
+        topk_eff = _SWA_WINDOW + P
+        glu_K, glu_P = 0, P
+        extra = f" P={P}"
+        pool_bh = torch.randn(B, H, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+        k_hca = torch.cat([k_full, pool_bh], dim=2).contiguous()
+        v_hca = torch.cat([v_full, pool_bh], dim=2).contiguous()
+        hca_mask = _hca_mask(S, P, cr, "cuda", torch.bfloat16)
+
+        def fwd_triton():
+            return _launch_v4_attention_fwd(
+                q,
+                k_hca,
+                v_hca,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                additive_mask=hca_mask,
+                scale=scale,
+                hca_local_seqlen=S,
+            )
+
+        out_t, lse_t = _call_fwd(fwd_triton)
+
+        def bwd_triton():
+            return _launch_v4_attention_bwd(
+                q,
+                k_hca,
+                v_hca,
+                out_t,
+                dout,
+                lse_t,
+                sink=sink,
+                swa_window=_SWA_WINDOW,
+                additive_mask=hca_mask,
+                scale=scale,
+                hca_local_seqlen=S,
+            )
+
+    else:
+        raise ValueError(f"unsupported cr={cr}")
+
+    # Fused single-latent (K==V) sparse-MLA backends (gluon / triton_v2 /
+    # turbo_flydsl): all on the SAME V4-form inputs (zero rope pad + [local ++ pool]
+    # kv + [SWA window ++ pool] topk). Time each by swapping the kernel pair.
+    sparse_fb = {}  # name -> (fwd_closure, bwd_closure)
+    if _SPARSE_MLA_BACKENDS:
+        gq, gkv, gtopk_idx, gsink, gdo = _build_gluon_v4form(
+            cr=cr, H=H, S=S, D=D, K=glu_K, P=glu_P, W=_SWA_WINDOW
+        )
+        gscale = 1.0 / math.sqrt(D)  # V4 scale (score over 512)
+        for _name, (_fwd_k, _bwd_k) in _SPARSE_MLA_BACKENDS.items():
+            fwd_c = (
+                lambda fk: (lambda: fk(gq, gkv, gtopk_idx, attn_sink=gsink, kv_lora_rank=D, scale=gscale))
+            )(_fwd_k)
+            _out, _lse = _call_fwd(fwd_c)
+            if _out is None:  # fwd unavailable (e.g. flydsl_v2 kernel WIP) -> skip bwd
+                bwd_c = None
+            else:
+                bwd_c = (
+                    lambda bk, o, l: (
+                        lambda: bk(
+                            gq, gkv, o, gdo, gtopk_idx, l, attn_sink=gsink, kv_lora_rank=D, scale=gscale
+                        )
+                    )
+                )(_bwd_k, _out, _lse)
+            sparse_fb[_name] = (fwd_c, bwd_c)
+
+    # Effective FLOPs over the USEFUL work (score+value both over head_dim=512,
+    # TOPK = real key count); BWD = 2.5x FWD. Same formula for ALL backends so
+    # TFLOP/s is comparable (the sparse-MLA zero-rope-pad overhead shows up in
+    # ms, not in counted FLOPs).
+    T = B * S
+    fwd_flop = 2.0 * T * H * topk_eff * (D + D)
+    flops = {"fwd": fwd_flop, "bwd": 2.5 * fwd_flop}
+
+    # Backend table: native production Triton (separate K/V), then the fused
+    # sparse-MLA backends (legacy `_flydsl_v0` gathered CSA is excluded).
+    backends = [("triton", fwd_triton, bwd_triton)]
+    for _name in (
+        "gluon",
+        "triton_v2",
+        "gluon_v2",
+        "gluon_v3",
+        "flydsl_v1",
+        "_turbo_flydsl",
+        "aiter_gluon",
+    ):
+        if _name in sparse_fb:
+            backends.append((_name, sparse_fb[_name][0], sparse_fb[_name][1]))
+
+    print(
+        f"\n=== V4-{variant.upper()} cr={cr} ({_CR_NAME[cr]}) | B={B} H={H} S={S} D={D}"
+        f"{extra} TOPK_eff={topk_eff} swa={_SWA_WINDOW} sink=on bf16 ===\n"
+        f"  FWD GFLOP={fwd_flop / 1e9:.1f} (useful, over head_dim={D})  "
+        f"(cells: ms | TFLOP/s; *_v2/gluon = V4-form fused single-latent)",
+        flush=True,
+    )
+    print(
+        f"  {'backend':12s} {'fwd (ms|TF)':>18s} {'bwd (ms|TF)':>18s}",
+        flush=True,
+    )
+    rows = []
+    for name, fwd_fn, bwd_fn in backends:
+        fwd_med, fwd_err = _safe_time(fwd_fn, warmup=warmup, iters=iters)
+        bwd_med, bwd_err = _safe_time(bwd_fn, warmup=warmup, iters=iters)
+        print(f"  {name:12s} {_cell(fwd_med, flops['fwd'])} {_cell(bwd_med, flops['bwd'])}", flush=True)
+        for op, err in (("fwd", fwd_err), ("bwd", bwd_err)):
+            if err:
+                print(f"       {name} {op} error: {err}", flush=True)
+        rows.append((name, fwd_med, bwd_med))
+    return rows
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--variant",
+        choices=["flash", "pro", "both"],
+        default="both",
+        help="model size to benchmark (default: both)",
+    )
+    parser.add_argument(
+        "--cr",
+        choices=["0", "4", "128", "all"],
+        default="all",
+        help="compress ratio / layer kind to benchmark (default: all)",
+    )
+    parser.add_argument("--seq", type=int, default=4096, help="sequence length (default 4096)")
+    parser.add_argument("--mbs", type=int, default=1, help="micro batch size / B (default 1)")
+    parser.add_argument("--warmup", type=int, default=10, help="warmup launches (default 10)")
+    parser.add_argument("--iters", type=int, default=30, help="timed launches (default 30)")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA / HIP device required for this benchmark")
+
+    # Production-optimal Triton config (matches attention_perf.md P57 defaults):
+    # split CSA FWD (monolithic OFF), split + segreduce BWD.
+    os.environ.setdefault("PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC", "0")
+    os.environ.setdefault("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "1")
+    os.environ.setdefault("PRIMUS_V4_CSA_BWD_SEGREDUCE", "1")
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    variants = ["flash", "pro"] if args.variant == "both" else [args.variant]
+    crs = [0, 4, 128] if args.cr == "all" else [int(args.cr)]
+
+    print(
+        f"device={torch.cuda.get_device_name(0)} torch={torch.__version__} "
+        f"seq={args.seq} mbs={args.mbs} warmup={args.warmup} iters={args.iters}",
+        flush=True,
+    )
+    for v in variants:
+        for cr in crs:
+            _bench_cr(v, cr, B=args.mbs, S=args.seq, warmup=args.warmup, iters=args.iters)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/deepseek-v4/benchmark/bench_v4_attention_results.md b/examples/deepseek-v4/benchmark/bench_v4_attention_results.md
new file mode 100644
index 000000000..6fff4e42f
--- /dev/null
+++ b/examples/deepseek-v4/benchmark/bench_v4_attention_results.md
@@ -0,0 +1,97 @@
+# DeepSeek-V4 Attention — Backend Performance
+
+Full forward + backward benchmark of every V4 attention backend, produced by
+[`bench_v4_attention.py`](./bench_v4_attention.py).
+
+## Setup
+
+- **GPU**: AMD Instinct MI355X (gfx950), single GPU
+- **Container**: `dev_primus_wenx`
+- **Torch**: `2.10.0+git94c6e04`, Triton `3.7.0`, FlyDSL `0.2.2`
+- **Primus-Turbo**: `dev/kyle/flydsl_attn_deepseekv4` @ `350ec3f` (native-FlyDSL sparse-MLA v2)
+- **Config**: `seq_len=4096`, `mbs=1`, bf16, sink **on**, `swa_window=128`
+- **Models**: V4-Flash (`H=64`, index_topk=512), V4-Pro (`H=128`, index_topk=1024)
+- **Layer kinds**: `cr=0` dense/SWA, `cr=4` CSA, `cr=128` HCA
+- **Timing**: `--warmup 10 --iters 30`, median latency
+- **Cell format**: `latency ms | TFLOP/s`
+
+Raw log: `agent/workspace/_bench_all_final.log`
+
+## Backends
+
+| Backend | Description |
+|---------|-------------|
+| `triton` | Production separate-K/V Triton dense/SWA/HCA + split CSA pool kernels |
+| `gluon` | 1st-gen fused single-latent Gluon sparse-MLA (`_gluon_dsa`) |
+| `triton_v2` | Fused single-latent sparse-MLA in plain Triton |
+| `gluon_v2` | 2nd-gen Gluon sparse-MLA baseline |
+| `gluon_v3` | Optimized Gluon sparse-MLA: Round-9 CSA formula-pack + aiter Gluon LSE fwd route (benchmark-only; not wired for training) |
+| `flydsl_v1` | In-tree native FlyDSL MFMA sparse-MLA backend |
+| `turbo_flydsl` | Extracted Primus-Turbo `sparse_mla_v2` (June `optimize/...` branch) |
+| `_turbo_flydsl` | **Integrated** Primus-Turbo native-FlyDSL sparse-MLA via the turbo API (`primus_turbo.flydsl.attention`); the `turbo` model backend |
+| `aiter_gluon` | aiter Gluon sparse-MLA prefill reference, fwd-only |
+
+`aiter_gluon` has no backward implementation, so bwd is shown as `—`.
+
+## Forward
+
+`latency ms | TFLOP/s`; **bold** is fastest latency in the row.
+
+| variant | cr | triton | gluon | triton_v2 | gluon_v2 | gluon_v3 | flydsl_v1 | turbo_flydsl | _turbo_flydsl | aiter_gluon |
+|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
+| flash | 0 | 0.46 \| 151.0 | 0.31 \| 222.1 | 0.30 \| 230.0 | 0.28 \| 247.8 | 0.28 \| 248.3 | 0.45 \| 153.3 | 0.30 \| 228.7 | **0.20 \| 335.5** | 0.47 \| 145.2 |
+| flash | 4 | 1.47 \| 234.3 | 0.89 \| 386.1 | 0.87 \| 397.1 | 0.73 \| 469.0 | 0.66 \| 523.6 | 1.37 \| 250.9 | 0.72 \| 476.6 | **0.53 \| 651.7** | 0.83 \| 413.4 |
+| flash | 128 | 0.75 \| 114.7 | 0.38 \| 226.3 | 0.38 \| 223.9 | 0.33 \| 263.9 | 0.33 \| 263.2 | 0.58 \| 147.3 | 0.35 \| 245.5 | **0.22 \| 384.2** | 0.52 \| 164.7 |
+| pro | 0 | 0.86 \| 159.5 | 0.57 \| 239.8 | 0.58 \| 236.2 | 0.51 \| 269.1 | 0.51 \| 269.0 | 1.05 \| 131.0 | 0.55 \| 251.9 | **0.38 \| 357.9** | 0.79 \| 173.1 |
+| pro | 4 | 4.44 \| 278.3 | 2.91 \| 425.5 | 2.78 \| 444.3 | 2.36 \| 525.0 | 1.92 \| 645.1 | 4.82 \| 256.6 | 2.09 \| 591.0 | **1.41 \| 878.1** | 2.16 \| 573.4 |
+| pro | 128 | 1.46 \| 117.7 | 0.72 \| 239.4 | 0.72 \| 238.6 | 0.61 \| 281.2 | 0.61 \| 280.9 | 1.20 \| 143.5 | 0.63 \| 270.7 | **0.43 \| 395.5** | 0.88 \| 195.2 |
+
+## Backward
+
+`latency ms | TFLOP/s`; **bold** is fastest latency in the row.
+
+| variant | cr | triton | gluon | triton_v2 | gluon_v2 | gluon_v3 | flydsl_v1 | turbo_flydsl | _turbo_flydsl | aiter_gluon |
+|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
+| flash | 0 | 2.09 \| 82.2 | 1.20 \| 143.4 | 1.16 \| 148.4 | 1.13 \| 152.6 | 1.13 \| 152.0 | 2.20 \| 78.3 | 1.38 \| 124.6 | **0.67 \| 257.8** | — |
+| flash | 4 | 5.18 \| 165.8 | 5.18 \| 166.0 | 5.93 \| 144.9 | 4.81 \| 178.5 | 3.99 \| 215.1 | 6.16 \| 139.5 | 3.94 \| 218.1 | **2.55 \| 336.9** | — |
+| flash | 128 | 2.86 \| 75.0 | 1.63 \| 131.4 | 1.67 \| 128.8 | 1.54 \| 139.3 | 1.54 \| 139.2 | 2.77 \| 77.5 | 1.84 \| 117.0 | **0.78 \| 274.9** | — |
+| pro | 0 | 4.02 \| 85.4 | 1.82 \| 189.0 | 1.81 \| 190.1 | 1.71 \| 201.4 | 1.70 \| 202.2 | 5.69 \| 60.3 | 2.09 \| 164.5 | **1.29 \| 267.2** | — |
+| pro | 4 | 15.10 \| 204.8 | 13.48 \| 229.3 | 10.74 \| 287.8 | 8.52 \| 362.9 | 8.52 \| 362.9 | 30.45 \| 101.5 | 9.41 \| 328.7 | **6.32 \| 489.3** | — |
+| pro | 128 | 5.55 \| 77.3 | 2.42 \| 177.4 | 2.47 \| 174.2 | 2.27 \| 189.4 | 2.27 \| 189.4 | 6.94 \| 61.9 | 2.91 \| 147.6 | **1.49 \| 288.2** | — |
+
+## `_turbo_flydsl` (integrated turbo) vs `gluon_v3` (best in-tree)
+
+`_turbo_flydsl` is the fastest backend on **every** fwd and bwd cell. Speedup over
+`gluon_v3` (TFLOP/s ratio):
+
+| variant | cr | FWD speedup | BWD speedup |
+|---|---:|---:|---:|
+| flash | 0 | 1.35× | 1.70× |
+| flash | 4 | 1.24× | 1.57× |
+| flash | 128 | 1.46× | 1.98× |
+| pro | 0 | 1.33× | 1.32× |
+| pro | 4 | 1.36× | 1.35× |
+| pro | 128 | 1.41× | 1.52× |
+| **mean** | | **~1.36×** | **~1.57×** |
+
+## Summary
+
+- **`_turbo_flydsl`** (the integrated Primus-Turbo native-FlyDSL backend, selectable
+  in the model via `use_v4_attention_backend = turbo`) is the fastest backend across
+  all six shapes in both directions — **~1.36× fwd** and **~1.57× bwd** over the best
+  in-tree backend (`gluon_v3`), and larger margins over `triton_v2`/`gluon_v2`.
+- The biggest wins are the CSA `cr=4` forward (pro cr=4: `1.41 ms` / 878 TF vs
+  gluon_v3 `1.92 ms` / 645 TF) and the backward across the board (flash cr=128 bwd
+  `0.78 ms` / 275 TF ≈ **2×** gluon_v3's `1.54 ms` / 139 TF). The fully-native FlyDSL
+  backward is the headline improvement.
+- `_turbo_flydsl` also clearly beats the older extracted `turbo_flydsl` (June branch),
+  chiefly on the backward (e.g. pro cr=4 bwd `6.32 ms` vs `9.41 ms`).
+- `gluon_v3` remains the strongest **in-tree** backend but is benchmark-only (not wired
+  for training); `gluon_v2` is the wired gluon training backend.
+
+## Reproduce
+
+```bash
+PYTHONPATH=<repo> python examples/deepseek-v4/benchmark/bench_v4_attention.py \
+    --variant both --cr all --warmup 10 --iters 30
+```
diff --git a/examples/deepseek-v4/projection/README.md b/examples/deepseek-v4/projection/README.md
new file mode 100644
index 000000000..6c5b1319e
--- /dev/null
+++ b/examples/deepseek-v4/projection/README.md
@@ -0,0 +1,93 @@
+# DeepSeek-V4 Performance Projection
+
+A trace-driven performance projection toolkit + static website for DeepSeek-V4
+(Flash / Pro) training on AMD Instinct GPUs (MI355X measured, MI455X projected).
+
+## Idea in one paragraph
+
+We profile a **single transformer layer** of a given compression-ratio (`cr`)
+type on **MI355X**, extract a clean forward / backward **time + TFLOPs breakdown**
+per module (attention sub-modules, MoE sub-modules), plus the model's non-layer
+parts (embedding, output/logits, loss) and the optimizer step. We emit one JSON
+per model variant. A static website then loads that JSON and, given a target GPU
+and a distributed strategy (PP / VPP / EP / DP / CP), reconstructs the full model
+(real layer count + `cr` schedule), models the PP bubble, EP dispatch/combine,
+recompute, and the optimizer, and derives **iteration time, TFLOP/s/GPU and
+tokens/s/GPU** step by step. Page 1 is the MI355X projection; page 2 scales the
+breakdown to MI455X by theoretical ratios.
+
+## Pipeline
+
+```
+run profiling script (per cr)  ->  chrome trace JSON (rank 0)
+        |                                  |
+        |                                  v
+        |                       tools/parse_trace.py  (+ kernel/module map)
+        v                                  |
+  one trace per cr type {0,4,128}          v
+                                breakdown JSON (site/data/<model>.json)
+                                           |
+                                           v
+                                static site (site/index.html)
+                                  - model config view
+                                  - per-cr fwd/bwd breakdown tables
+                                  - GPU + parallelism controls
+                                  - step-by-step iter-time / TFLOPs / tok/s derivation
+                                  - iteration timeline (3 levels: layer / PP ranks / schedule)
+                                  - MI355X page + MI455X scaled page
+```
+
+## Directory layout
+
+```
+examples/deepseek-v4/projection/
+  README.md                      # this file
+  design/                        # methodology, assumptions, schema, math (the spec)
+    01-overview.md
+    02-assumptions.md
+    03-json-schema.md
+    04-projection-math.md
+    07-iteration-timeline.md       # 3-level iteration-time composition view
+  script/                        # profiling launchers (one trace per cr)
+    deepseek_v4_layer_trace-projection.sh
+  tools/                         # trace -> breakdown JSON
+    parse_trace.py
+    kernel_module_map.py
+  site/                          # static website (no build step)
+    index.html
+    assets/app.js
+    assets/style.css
+    data/<model>.json            # breakdown JSON consumed by the site
+```
+
+## Quick start
+
+1. Profile each `cr` on MI355X (run inside the training container):
+
+   ```bash
+   # one trace per cr type; 1 layer, seq 4096, adam + dist-opt, GA=2, no recompute
+   CR=0   ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+   CR=4   ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+   CR=128 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+   ```
+
+2. Build the breakdown JSON from the three traces:
+
+   ```bash
+   python3 examples/deepseek-v4/projection/tools/parse_trace.py \
+     --model pro \
+     --trace cr0=<trace_cr0.json> \
+     --trace cr4=<trace_cr4.json> \
+     --trace cr128=<trace_cr128.json> \
+     --out examples/deepseek-v4/projection/site/data/pro.json
+   ```
+
+3. Open the site locally:
+
+   ```bash
+   python3 -m http.server -d examples/deepseek-v4/projection/site 8000
+   # http://localhost:8000/?model=pro
+   ```
+
+See `design/` for the full methodology and the exact assumptions baked into the
+projection math.
diff --git a/examples/deepseek-v4/projection/design/01-overview.md b/examples/deepseek-v4/projection/design/01-overview.md
new file mode 100644
index 000000000..eed63aac8
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/01-overview.md
@@ -0,0 +1,91 @@
+# 01 — Overview & methodology
+
+## Goal
+
+Project DeepSeek-V4 (Flash / Pro) training throughput (iteration time,
+TFLOP/s/GPU, tokens/s/GPU) for a real multi-hundred-GPU job, from a small set of
+**single-layer** MI355X traces, then scale the result to MI455X by theoretical
+hardware ratios.
+
+## Why single-layer, per-cr traces
+
+DeepSeek-V4 layers come in three attention flavours selected by the per-layer
+`compress_ratio` (`cr`):
+
+| cr   | attention branch                              |
+|------|-----------------------------------------------|
+| 0    | dense + sliding-window attention (SWA)         |
+| 4    | CSA (compressed sparse attention, top-k via Indexer) |
+| 128  | HCA (compressed attention, full pool visibility) |
+
+The MoE block is **identical across all `cr`** (cr only changes attention). So
+three single-layer traces — one per cr — fully characterise the per-layer cost,
+and the full model is `Σ over layers (attention[cr_of_layer] + moe)`.
+
+Profiling a *single* layer per cr (instead of a full model) keeps:
+- **memory** within the MI355X budget at the production `seq=4096`, and
+- **attribution clean**: with one cr in the trace, the dense attention kernels
+  (which are shared across cr types in a multi-cr run) belong unambiguously to
+  that cr.
+
+## What we measure vs what we model
+
+**Measured on MI355X (from traces):**
+- per-module forward / backward time for one layer of each cr,
+- TFLOPs for the three compute-bound kernel classes (`gemm`, `grouped_gemm`,
+  `attn`); everything else is treated as memory-bound,
+- embedding, output/logits, loss (the non-layer parts), taken once,
+- the optimizer-step cost per unit parameter (for scaling).
+
+**Modeled on top (in the website):**
+- full layer count and exact `cr` schedule (Flash 43L, Pro 61L),
+- PP / VPP partitioning -> per-stage critical path + pipeline bubble,
+- EP dispatch/combine cost (no overlap in current stack),
+- gradient-accumulation (GA) and DP behaviour (DP/PP comm assumed hidden),
+- activation recompute (add a forward pass to the recomputed layers' backward),
+- optimizer step scaled to per-rank parameter count for the target sharding,
+- MI455X = MI355X breakdown scaled by compute / memory-bandwidth ratios.
+
+## Capture configuration (the projection trace)
+
+The profiling script (`script/deepseek_v4_layer_trace-projection.sh`) deliberately
+differs from `run_deepseek_v4_pro_muon.sh`:
+
+| knob              | projection value | why |
+|-------------------|------------------|-----|
+| `seq_length`      | 4096             | production per-microbatch token count |
+| `num_layers`      | 1                | clean per-layer attribution; fits memory at seq 4096 |
+| `compress_ratios` | `[CR]`           | one cr per trace |
+| optimizer         | `adam` + distributed optimizer | compute is optimizer-independent; dist-opt = zero1 |
+| overlap grad/param | off            | num_layers=1 breaks Megatron's chained param-sync; compute already clean |
+| `global_batch_size` | `2 * DP`       | GA=2 (see below) |
+| recompute         | off (`recompute_num_layers 0`) | capture pure fwd / pure bwd |
+| profiler          | on, `with_stack=True`, window iter 6->7 | map kernels -> nn.module |
+
+### Why overlap-off + GA=2 + take min
+
+We capture with the distributed optimizer's comm overlap **off**. Two reasons:
+
+- with `num_layers=1`, Megatron's chained param-gather sync trips an assertion
+  (`param_and_grad_buffer.start_param_sync`) when overlap is on; and
+- with overlap off, **every microbatch's compute is already overlap-free** —
+  exactly the clean per-kernel time we want. In a real run GA is large so the
+  vast majority of microbatches are clean anyway, and the projection assumes DP
+  comm is hidden (A2), so there is no need to measure the contaminated overlap.
+
+GA=2 (two microbatches) plus a late steady profiler window lets the parser
+compute per-microbatch time as `sum(kernel_durations) / num_microbatches`.
+This keeps scalar control-flow stalls (for example Indexer top-k syncs) that the
+full-model calibration shows are real per-layer costs, while avoiding warm-up
+and autotune iterations.
+
+## Non-overlap assumptions in the current stack
+
+- **No MoE deepep-comm + grouped-gemm overlap.** dispatch + grouped_gemm +
+  combine are summed directly.
+- **EP dispatch/combine has no overlap** with compute; counted in full.
+- **DP / PP comm assumed hidden** (only the PP bubble remains). Optimistic; first
+  thing to revisit at calibration time.
+
+See `02-assumptions.md` for the complete list, and `04-projection-math.md` for
+the formulas.
diff --git a/examples/deepseek-v4/projection/design/02-assumptions.md b/examples/deepseek-v4/projection/design/02-assumptions.md
new file mode 100644
index 000000000..c86e751d8
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/02-assumptions.md
@@ -0,0 +1,108 @@
+# 02 — Assumptions (single source of truth)
+
+Every assumption baked into the projection. When a projection number looks off,
+start here.
+
+## Optimizer / DP
+
+- **A1.** Production optimizer modeled = **AdamW + distributed optimizer (zero1)**.
+  Muon is out of scope. NOTE on **capture**: the trace itself is taken with
+  `use_distributed_optimizer=False` (+ fp32 states) because with dist-opt ON the
+  ROCm Kineto profiler drops the compute GPU kernels for pure dense(cr=0)/HCA
+  (cr=128) layers (CSA cr=4 is unaffected). dist-opt does not change the fwd/bwd
+  compute, so this only affects which kernels Kineto records; the optimizer step
+  is modeled analytically regardless (A3).
+- **A2.** DP communication (param all-gather / grad reduce-scatter) is **fully
+  hidden** behind compute at large GA. We do not add a DP comm term. *(Optimistic;
+  primary calibration target.)*
+- **A3.** The optimizer step is a per-iteration term, **not** multiplied by GA or
+  replicated per PP microbatch. It scales with **per-rank optimizer parameter
+  count**: full-model params are first averaged over PP/TP ownership, then
+  sharded over DP under ZeRO-1. CP does not shard parameters. The modeled Adam
+  traffic uses the full mixed-precision read/write cost per parameter and is
+  treated as memory-bound.
+
+## Pipeline / parallelism
+
+- **A4.** PP point-to-point comm is **hidden**; only the pipeline **bubble**
+  remains. CP/TP comm not modeled in v1 (TP=1, CP=1 in the V4 release configs).
+- **A5.** Pipeline bubble fraction uses 1F1B: `(PP-1)/GA`; interleaved VPP divides
+  it by the VPP degree: `(PP-1)/(GA*VPP)`.
+- **A6.** Per-stage time is the **sum of the specific cr-type layers** mapped to
+  that stage; the iteration critical path is driven by the **max** (slowest)
+  stage, plus embedding on stage 0 and output/loss on the last stage.
+
+## EP / MoE
+
+- **A7.** EP dispatch/combine has **no overlap** with compute (current stack) and
+  is counted in full, every microbatch.
+- **A8.** No MoE deepep-comm + grouped-gemm overlap; MoE = dispatch + grouped_gemm
+  + combine summed.
+- **A9.** EP is intra-node only (e.g. EP=8 within an 8-GPU node). Cross-node EP is
+  out of scope for v1; if EP spans nodes the dispatch/combine cost model must
+  change (RDMA, different bandwidth).
+- **A10.** MoE per-layer cost is `cr`-independent; the three single-cr traces
+  must agree on it (cross-check). The site uses one MoE breakdown for all layers.
+
+## Trace capture / attribution
+
+- **A11.** One trace per cr (`0`, `4`, `128`), **1 layer**, `seq=4096`,
+  `recompute off`, profiler window iter 6->7 (post warmup/autotune).
+- **A12.** Capture with comm-overlap **off** (`num_layers=1` breaks Megatron's
+  chained param sync, and compute is already clean without overlap). GA=2; clean
+  per-kernel time = `min` over launches grouped by `(module, phase, shape)`,
+  removing residual jitter.
+- **A13.** Kernel -> nn.module attribution uses `with_stack=True`: GPU kernels are
+  linked to their launching CPU op via trace flow events, and the module is read
+  from the CPU op's python call stack. **fwd/bwd phase** is determined by, in
+  priority order: (1) a `_fwd_`/`_bwd_` tag in the kernel name; (2) for linked
+  kernels, whether the launching CPU op's timestamp lies inside an
+  `autograd::engine::evaluate_function` interval (= backward); (3) for unlinked
+  kernels (no `External id`), whether the kernel's GPU timestamp lies in the
+  backward GPU-time window reconstructed from the linked backward kernels. The
+  old rule (default-to-forward when unlinked) leaked backward compute -- incl.
+  the MoE dgrad/wgrad grouped GEMMs -- into forward; see `design/06`. One-off
+  device stalls billed to a compute kernel (> `_MAX_PLAUSIBLE_LAUNCH_US`) are
+  dropped as artifacts and reported in `provenance.dropped_stall_us_per_mb`.
+- **A14.** Only `gemm`, `grouped_gemm`, `attn` kernels get a TFLOPs number; all
+  other kernels are memory-bound (TFLOPs = null) and contribute time only.
+- **A15.** Embedding / output-logits / loss are taken **once** (from any single
+  trace), not triple-counted across the three cr traces.
+
+## Recompute
+
+- **A16.** Traces are captured with recompute **off** (pure fwd, pure bwd). At
+  projection time, a recomputed layer's backward gets `+1 forward` of that layer
+  added back. Recompute selection (#layers / which) is a site control.
+
+## MTP
+
+- **A17.** MTP is modeled analytically when `mtp_num_layers > 0`. Each MTP
+  depth uses `mtp_compress_ratios` (default cr=4, matching the current Flash
+  Megatron FLOPs anchor), plus the per-depth `eh_proj`, extra logits/loss, and
+  HyperHead FLOPs. Timing is an approximation: the MTP inner layer reuses the
+  measured layer time for that cr, while `eh_proj` is scaled from output-GEMM
+  throughput. A dedicated MTP trace remains the next calibration step.
+
+## Scope exclusions (v1)
+
+- **A18.** No cross-node EP; no TP; no CP comm modeling.
+- **A19.** FP8 / MXFP8 not modeled; BF16 only.
+
+## MI355X -> MI455X scaling
+
+- **A20.** Compute-bound kernels (`gemm`, `grouped_gemm`, `attn`) scale by the
+  **peak-TFLOPs ratio** (BF16) `t_mi355 / t_mi455`.
+- **A21.** Memory-bound kernels (everything else, incl. optimizer) scale by the
+  **HBM-bandwidth ratio** `bw_mi355 / bw_mi455`.
+- **A22.** A single tunable **efficiency factor** (default 1.0) multiplies the
+  scaled compute time to account for MFU differences on new HW (flat peak ratio
+  is optimistic).
+- **A23.** Comm (EP/DP/PP) is not rescaled in v1 (intra-node EP only; DP/PP
+  hidden).
+
+## Validation
+
+- **A24.** Self-consistency: configured to the trace scenario (PP=1, EP=8,
+  measured GA, single node) the model must reproduce the measured single-node
+  iteration time. Multi-node calibration is deferred.
diff --git a/examples/deepseek-v4/projection/design/03-json-schema.md b/examples/deepseek-v4/projection/design/03-json-schema.md
new file mode 100644
index 000000000..3f1b7b8cb
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/03-json-schema.md
@@ -0,0 +1,117 @@
+# 03 — Breakdown JSON schema
+
+One JSON per model variant (`flash` / `pro`), written to
+`site/data/<model>.json`, consumed by the website. All times in **microseconds
+(us)**; all FLOP counts are per single layer / per single invocation at the
+captured `seq` and `micro_batch_size`.
+
+## Top level
+
+```jsonc
+{
+  "schema_version": 1,
+  "model": "pro",                       // "flash" | "pro"
+  "generated_at": "2026-06-18T11:00:00Z",
+  "provenance": {
+    "commit": "dac0a60c",
+    "host": "node001",
+    "container": "dev_primus_wenx",
+    "traces": { "cr0": "<path>", "cr4": "<path>", "cr128": "<path>" }
+  },
+
+  "capture": {                          // how the trace was taken (the unit)
+    "gpu": "MI355X",
+    "seq_length": 4096,
+    "micro_batch_size": 1,
+    "tokens_per_microbatch": 4096,      // seq_length * micro_batch_size
+    "ep": 8,
+    "ga_for_capture": 2,
+    "optimizer": "adam",
+    "distributed_optimizer": true,
+    "recompute": "off"
+  },
+
+  "model_config": {                     // shown on the site's config panel
+    "num_layers": 61,
+    "hidden_size": 7168,
+    "num_attention_heads": 128,
+    "kv_channels": 512,
+    "num_experts": 384,
+    "moe_router_topk": 6,
+    "moe_ffn_hidden_size": 3072,
+    "index_topk": 1024,
+    "vocab_size": 129280,
+    "compress_ratios": [128,128,4, /* ... */ ,0],
+    "cr_layer_counts": { "0": 1, "4": 29, "128": 31 }   // derived from compress_ratios
+  },
+
+  "hardware": {                         // for MI355->MI455 scaling (A20-A23)
+    "MI355X": { "peak_tflops_bf16": 2300, "hbm_bandwidth_gbps": 8000 },
+    "MI455X": { "peak_tflops_bf16": 4600, "hbm_bandwidth_gbps": 16000 }
+  },
+
+  "layers": {                           // per-cr breakdown
+    "0":   { "attention": <Breakdown>, "moe": <Breakdown> },
+    "4":   { "attention": <Breakdown>, "moe": <Breakdown> },
+    "128": { "attention": <Breakdown>, "moe": <Breakdown> }
+  },
+
+  "non_layer": {                        // taken once (A15)
+    "embedding": <Breakdown>,
+    "output":    <Breakdown>,           // final norm + lm_head/logits
+    "loss":      <Breakdown>
+  },
+
+  "optimizer": {                        // per-iteration term (A3)
+    "type": "adam",
+    "measured_params": 123456789,       // params updated in the 1-layer trace, this rank
+    "time_us": 850.0,                   // measured optimizer-step time for those params
+    "bytes_per_param": 18,              // bf16 master-param-remainder adam state bytes
+    "class": "memory_bound"
+  },
+
+  "comm": {                             // per-microbatch EP cost (A7-A9), per layer
+    "ep_dispatch_us": 0.0,
+    "ep_combine_us": 0.0
+  }
+}
+```
+
+## `<Breakdown>` object
+
+A phase-split list of modules. Each module is one row group; the site renders
+forward left-to-right and backward right-to-left.
+
+```jsonc
+{
+  "forward":  [ <Module>, ... ],
+  "backward": [ <Module>, ... ]
+}
+```
+
+## `<Module>` object
+
+```jsonc
+{
+  "module": "attn.core",          // logical module name (from kernel/module map)
+  "time_us": 740.0,               // clean min-grouped time for this module/phase
+  "class": "compute_bound",       // "compute_bound" | "memory_bound"
+  "flop_class": "attn",           // "gemm" | "grouped_gemm" | "attn" | null
+  "flops": 1.84e12,               // total FLOPs for this module/phase, or null
+  "tflops": 740.0,                // achieved TFLOP/s = flops / time_s / 1e12 (null if memory_bound)
+  "kernels": [                    // optional: contributing kernels (debug / drill-down)
+    { "name": "_v4_attention_fwd_kernel", "time_us": 500.0, "launches": 1 }
+  ]
+}
+```
+
+## Notes
+
+- `time_us` is always the **clean** (min-grouped, overlap-free) time per A12.
+- `tflops` is present only when `flop_class != null` (A14).
+- The site computes everything else (full model, PP/EP/DP, MI455) from this JSON;
+  the JSON itself is hardware-MI355X, single-layer, single-microbatch ground
+  truth + static config.
+- `cr_layer_counts` is derived from `compress_ratios` by the parser so the site
+  doesn't re-parse the schedule.
+- For Flash, `model` = `"flash"`, `cr_layer_counts` e.g. `{"0":3,"4":20,"128":20}`.
diff --git a/examples/deepseek-v4/projection/design/04-projection-math.md b/examples/deepseek-v4/projection/design/04-projection-math.md
new file mode 100644
index 000000000..fbe192bcf
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/04-projection-math.md
@@ -0,0 +1,193 @@
+# 04 — Projection math
+
+This is the exact derivation the website implements (`site/assets/app.js`). All
+inputs come from the breakdown JSON (`03-json-schema.md`); all assumptions are in
+`02-assumptions.md`. Times in seconds unless noted.
+
+Notation:
+- `cr ∈ {0, 4, 128}`, `n[cr]` = number of layers of that cr (`cr_layer_counts`).
+- A `<Breakdown>` has `forward` / `backward` module lists. Define
+  `T(bd, phase) = Σ_module bd[phase][m].time` (sum of clean module times).
+
+## Step 0 — per-layer and non-layer base times (MI355X)
+
+For each cr:
+```
+layer_fwd[cr] = T(attention[cr], forward) + T(moe, forward)
+layer_bwd[cr] = T(attention[cr], backward) + T(moe, backward)
+```
+EP dispatch/combine are included as memory-bound module rows inside `moe`
+(A7/A8), so they are already in these sums — do not add `comm` again (the `comm`
+field is informational for the UI).
+
+Non-layer (taken once, A15):
+```
+emb_fwd = T(embedding, forward);  emb_bwd = T(embedding, backward)
+out_fwd = T(output, forward) + T(loss, forward)
+out_bwd = T(output, backward) + T(loss, backward)
+```
+
+MTP (if `mtp_num_layers > 0`, A17):
+```
+mtp_inner_fwd = layer_fwd[mtp_cr] ; mtp_inner_bwd = layer_bwd[mtp_cr]
+mtp_out_fwd   = out_fwd          ; mtp_out_bwd   = out_bwd
+mtp_eh_time   ≈ output_time * (mtp_eh_proj_flops / output_flops)
+
+mtp_fwd = mtp_num_layers * (mtp_inner_fwd + mtp_out_fwd) + mtp_eh_time / 3
+mtp_bwd = mtp_num_layers * (mtp_inner_bwd + mtp_out_bwd) + 2 * mtp_eh_time / 3
+```
+The current Flash Megatron FLOPs anchor uses `mtp_cr=4`; a future dedicated MTP
+trace can replace the `eh_proj` and inner-layer approximation.
+
+## Step 0b — manual layer-time mode (optional, UI-only)
+
+The site exposes a **layer-timing mode** toggle in the controls panel:
+
+- `trace` (default): `layer_fwd/bwd[cr]` come from the breakdown JSON exactly as
+  in Step 0.
+- `manual`: the user types `layer_fwd[cr]` / `layer_bwd[cr]` directly (one
+  fwd/bwd pair per `cr ∈ {0,4,128}`, in µs) and those values replace the
+  trace-derived per-layer times. This is a what-if calculator: you supply the
+  per-layer cost and the site reuses Steps 1-5 unchanged to derive the full-model
+  iteration time, bubble, optimizer, tokens/s and TFLOP/s.
+
+Rules baked into the implementation:
+
+- **Granularity is per-`cr`** (same as the data model — `cr` only changes
+  attention, the MoE block is shared), not per physical layer. The `cr` schedule,
+  PP/VPP layout and recompute still expand these per-cr times to the full model.
+- **Per-GPU storage.** A hand-entered time already targets one GPU, so manual
+  values are stored separately for MI355X and MI455X and the
+  MI355→MI455 scaling (Step 6) is **bypassed** in manual mode. Switching the GPU
+  tab edits that GPU's own set.
+- **Prefill from trace.** Entering manual mode (or switching GPU within it)
+  seeds any unset field with the current trace-derived value, so toggling never
+  changes the result until you actually edit a number. Unset/blank fields keep
+  falling back to trace.
+- **Time only, FLOPs stay analytic.** Manual input overrides time but not FLOPs;
+  `TFLOP/s/GPU` continues to use the V4 analytic model FLOPs (Step 5), so it
+  stays meaningful.
+- **Scope.** Manual covers the three per-cr decoder layers **and** the non-layer
+  parts — embedding (PP stage 0), output / loss / MTP (last PP stage) — each as a
+  fwd/bwd pair. Any field left unset falls back to its trace-derived value, so you
+  can override just the parts you care about. FLOPs stay analytic regardless.
+
+## Step 1 — recompute (A16)
+
+If a layer is recomputed, its backward replays one forward:
+```
+layer_bwd_eff[cr] = layer_bwd[cr] + recompute_factor[cr] * layer_fwd[cr]
+```
+`recompute_factor` ∈ {0,1} per layer. The site exposes `none`, `full`, and
+`first-n`; `first-n` adds one forward replay to the first N decoder layers owned
+by each physical PP stage, matching the common Megatron
+`recompute_num_layers=N` block pattern.
+
+## Step 2 — map layers to PP stages / VPP chunks (A6)
+
+Inputs: `PP`, `VPP`, optional `pipeline_layout`. Total model chunks
+`C = PP * VPP`. If `pipeline_layout` is provided, parse Megatron-style `t` /
+`t*N` stage specs (for example `Et*10|t*11|t*11|t*11mL`) and assign virtual
+chunk `k` to device `k mod PP`. Otherwise build the ordered layer list from
+`compress_ratios`, slice it into `C` contiguous chunks, and use the same
+`k mod PP` mapping. The UI validates that an explicit layout has exactly
+`PP*VPP` stages and exactly `num_layers` decoder layers; invalid layouts block
+projection instead of silently falling back. For device `d`:
+```
+Df[d] = Σ_{chunks on d} Σ_{layer in chunk} layer_fwd[cr(layer)]
+Db[d] = Σ_{chunks on d} Σ_{layer in chunk} layer_bwd_eff[cr(layer)]
+```
+Add non-layer parts to their devices:
+```
+Df[0]      += emb_fwd ;  Db[0]      += emb_bwd
+Df[PP-1]   += out_fwd ;  Db[PP-1]   += out_bwd
+Df[PP-1]   += mtp_fwd ;  Db[PP-1]   += mtp_bwd
+```
+Critical device:
+```
+Df_crit = max_d Df[d] ;  Db_crit = max_d Db[d]
+```
+(Using per-device max is an upper-bound for imbalanced stages; for a balanced
+schedule it is exact.)
+
+## Step 3 — pipeline iteration time (A4/A5)
+
+With gradient accumulation `GA` microbatches and interleaved VPP, the steady
+1F1B time on the critical device is:
+```
+pipe_compute = (GA + (PP - 1) / VPP) * (Df_crit + Db_crit)
+```
+- `GA * (Df_crit + Db_crit)` is the steady throughput term;
+- `(PP-1)/VPP * (Df_crit + Db_crit)` is the bubble (fraction `(PP-1)/(GA*VPP)`).
+- PP=1 ⇒ `pipe_compute = GA * (Df_crit + Db_crit)` (no bubble).
+
+`GA = GBS / (DP * MBS)`. The site takes `GBS`, `MBS`, `DP` as inputs (or derives
+`DP = world_size / (PP * TP * CP)` with EP ⊆ DP).
+
+## Step 4 — optimizer step (A1/A3)
+
+Per-iteration, once, zero1-sharded, memory-bound:
+```
+local_model_params  = total_params / (PP * TP)
+per_rank_opt_params = local_model_params / DP
+opt_bytes           = per_rank_opt_params * bytes_per_param
+opt_time            = opt_bytes / hbm_bandwidth / opt_efficiency
+```
+`total_params` is computed from `model_config` for the full model (dense +
+expert params + untied embedding/output). PP/TP determine the average local
+model-parameter ownership; CP does not shard parameters. EP is represented in
+the full expert count and cancels with the data-replica count for ZeRO-1
+optimizer sharding, so the average optimizer shard is `total/(PP*TP*DP)`.
+`bytes_per_param` is the full Adam mixed-precision step traffic (default 30B:
+reads + writes), and `opt_efficiency` is tunable. The measured
+`optimizer.time_us` carried in the JSON is displayed as a sanity reference.
+
+## Step 5 — totals (A2/A4: DP & PP comm hidden)
+
+```
+iter_time = pipe_compute + opt_time
+```
+Throughput:
+```
+tokens_per_iter      = GBS * seq_length            (= GA * DP * MBS * seq)
+tokens_per_s         = tokens_per_iter / iter_time
+tokens_per_s_per_gpu = tokens_per_s / world_size
+```
+TFLOP/s/GPU (matmul-flops convention, A14): per-microbatch model compute FLOPs
+```
+F_mb = Σ_cr n[cr] * (flops_fwd[cr] + flops_bwd[cr]) + nonlayer_flops
+       where flops_*[cr] = Σ_module (module.flops or 0) over the cr breakdown
+       + mtp_inner + mtp_eh_proj + mtp_extra_logits + mtp_hc_head
+flops_per_iter   = F_mb * GA * DP            (all microbatches, all DP replicas)
+TFLOP_s_per_gpu  = flops_per_iter / iter_time / world_size / 1e12
+```
+`tokens_per_s_per_gpu` is the headline metric (independent of FLOP convention);
+`TFLOP/s/GPU` is reported for comparison with Primus' own logging.
+
+## Step 6 — MI455X scaling (A20-A23)
+
+Rescale every module time before re-running Steps 0-5:
+```
+ratio_compute = peak_tflops_bf16[MI355] / peak_tflops_bf16[MI455]
+ratio_memory  = hbm_bandwidth[MI355]   / hbm_bandwidth[MI455]
+
+time'(module) = module.time * ratio_compute / compute_efficiency   if compute_bound
+              = module.time * ratio_memory                          if memory_bound
+```
+`compute_efficiency` (default 1.0) is the MFU knob (A22). FLOPs are unchanged
+(same math), so MI455 `tflops` rises by `1/ratio_compute * compute_efficiency`.
+Optimizer scales by `ratio_memory`. Comm not rescaled (A23).
+
+## Step 7 — self-consistency check (A24)
+
+Configure `PP=1, VPP=1, EP=8, DP=1, MBS=1, GA=GA_capture` and confirm the model's
+`iter_time` matches the measured single-node iteration time within tolerance. The
+site shows this check on the MI355X page when capture metadata is present.
+
+## Worked control set (website inputs)
+
+GPU page (MI355X / MI455X), then: `world_size`, `PP`, `VPP`, `EP`, `DP` (or
+derive), `CP`, `TP`, `MBS`, `GBS` (or `GA`), recompute mode, and the tunables
+`opt_efficiency`, `compute_efficiency`, `bytes_per_param`. Every intermediate
+(`layer_fwd/bwd`, `Df/Db` per stage, `pipe_compute`, bubble %, `opt_time`,
+`iter_time`, `tokens/s/gpu`, `TFLOP/s/gpu`) is displayed step by step.
diff --git a/examples/deepseek-v4/projection/design/05-deployment.md b/examples/deepseek-v4/projection/design/05-deployment.md
new file mode 100644
index 000000000..3ca252ada
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/05-deployment.md
@@ -0,0 +1,58 @@
+# 05 — Deployment (GitHub Pages)
+
+The repository already publishes a single GitHub Pages site from `main` via
+`.github/workflows/deploy-backend-gap-dashboard.yml`, which builds a bundle with
+`tools/backend_gap_report/build_site_bundle.py` and deploys it with
+`actions/deploy-pages`. A repo can only serve one Pages site, so the projection
+site is published as a **subpath of that same bundle** rather than as a separate
+deployment.
+
+## How it is wired
+
+`build_site_bundle.py` copies `examples/deepseek-v4/projection/site/` into the bundle at
+`deepseek-v4-projection/` (after the backend-gap bundle is built, before
+validation). The projection site uses only relative asset/data paths
+(`./assets/...`, `./data/...`), so it works unchanged under a subpath.
+
+Result URL:
+
+```
+https://<user>.github.io/<repo>/deepseek-v4-projection/?model=pro
+https://<user>.github.io/<repo>/deepseek-v4-projection/?model=flash
+```
+
+## Triggering a deploy
+
+The Pages workflow triggers on pushes to `main` touching its `paths:` list
+(currently `docs/backend-gap/**`, `docs/weekly_reports/**`,
+`docs/monthly_reports/**`, `tools/backend_gap_report/**`, and the workflow file).
+
+- The change to `tools/backend_gap_report/build_site_bundle.py` in this work is
+  itself under a watched path, so the **first** merge to `main` will rebuild and
+  publish the projection site automatically.
+- To make **projection-only** changes (new `site/data/*.json`, site tweaks) also
+  auto-deploy, add the projection path to the workflow's `paths:` on `main`:
+
+  ```yaml
+  # .github/workflows/deploy-backend-gap-dashboard.yml  (on: push: paths:)
+      - "examples/deepseek-v4/projection/site/**"
+  ```
+
+- Or trigger manually: the workflow has `workflow_dispatch` (Run workflow button).
+
+No `.nojekyll` is needed: the bundle is uploaded as a Pages artifact and served
+directly (no Jekyll processing), and no asset path starts with `_`.
+
+## Local preview
+
+```bash
+python3 -m http.server -d examples/deepseek-v4/projection/site 8011
+# http://localhost:8011/?model=pro
+```
+
+## Bundle-build smoke (optional, needs pandoc/weasyprint for the backend-gap PDFs)
+
+```bash
+python3 tools/backend_gap_report/build_site_bundle.py --output-dir /tmp/primus-site
+ls /tmp/primus-site/deepseek-v4-projection/   # index.html, assets/, data/
+```
diff --git a/examples/deepseek-v4/projection/design/06-calibration.md b/examples/deepseek-v4/projection/design/06-calibration.md
new file mode 100644
index 000000000..72373ef03
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/06-calibration.md
@@ -0,0 +1,73 @@
+# 06 — Calibration (single-node, measured)
+
+The projection is anchored to a real single-node run so its iteration time and
+TFLOP/s line up with what Megatron reports.
+
+## FLOPs: ported V4 closed-form (exact)
+
+`tools/v4_flops.py` ports Megatron's `deepseek_v4_flops_patches` closed form.
+Self-test against the measured flash 16-layer run (GBS64, seq4096):
+
+| component | analytic (TFLOP/gb) | measured (TFLOP/gb) |
+|---|---:|---:|
+| attn_qkv_o | 10766.4 | 10766.4 |
+| attn_scores | 2476.0 | 2476.0 |
+| compressor | 527.8 | 527.8 |
+| indexer | 1650.9 | 1650.9 |
+| moe | 17838.5 | 17818.7 |
+| logits | 832.9 | 833.7 |
+| **TOTAL** | **34112** | **34093** (0.05%) |
+
+The site uses these analytic FLOPs (per cr-layer, B=1, capture seq) for TFLOP/s,
+matching Megatron's convention (fwd+bwd × FMA = 6×, recompute excluded). The
+breakdown JSON carries them in `analytic_flops`.
+
+## Iteration time: single-layer → full-model bias
+
+Measured single-node anchor (`script/_calibrate_flash.sh`, full flash):
+
+| knob | value |
+|---|---|
+| layers / cr | 16, cr=[0×3, 4×6, 128×7] |
+| parallel | PP1 / EP8 / DP8 (world 8), TP1/CP1 |
+| GBS / GA / MBS / seq | 64 / 8 / 1 / 4096 |
+| recompute | full (uniform, 1) |
+| optimizer | adam + distributed optimizer |
+
+**Measured**: iter ≈ 6665 ms, 636 TFLOP/s/GPU, ~4917 tokens/s/GPU.
+
+**Projection (raw, calibFactor=1.0)**: iter 7177 ms (+7.7%), 586 TFLOP/s/GPU
+(−7.9%). The per-layer time captured from the single-layer profile runs ~7-8%
+high vs a layer inside the full model (single-layer capture has no neighbour-
+layer overlap / cache reuse, and per-launch overhead is a larger share). This is
+a systematic, near-constant bias.
+
+**Calibration**: a single `calibFactor = 0.93` on the pipeline compute time
+brings it in line:
+
+| metric | measured | projection (calibFactor 0.93) |
+|---|---:|---:|
+| iter time | 6665 ms | ~6680 ms (+0.2%) |
+| TFLOP/s/GPU | 636 | ~630 (−1%) |
+| tokens/s/GPU | 4917 | ~4900 (−0.4%) |
+
+`calibFactor` is a site control (default 0.93).
+
+## Caveats
+
+- `calibFactor` is from one anchor (flash, 16L, PP1). Pro / other parallel
+  layouts may want a slightly different value; re-anchor with another
+  `_calibrate_*` run if precision matters.
+- analytic FLOPs are evaluated at the capture seq (4096); changing seq in the
+  UI does not re-derive them (attention FLOPs are seq-dependent).
+- Optimizer step is analytic (per-rank params / HBM-BW); DP/PP comm assumed
+  hidden (A2/A4).
+
+## Reproduce
+
+```bash
+# measured anchor (single node, full model)
+LAYERS=16 GBS=64 bash examples/deepseek-v4/projection/script/_calibrate_flash.sh   # (helper; not committed)
+# analytic flops self-test
+python3 examples/deepseek-v4/projection/tools/v4_flops.py
+```
diff --git a/examples/deepseek-v4/projection/design/07-iteration-timeline.md b/examples/deepseek-v4/projection/design/07-iteration-timeline.md
new file mode 100644
index 000000000..1dcc38386
--- /dev/null
+++ b/examples/deepseek-v4/projection/design/07-iteration-timeline.md
@@ -0,0 +1,165 @@
+# 07 — Iteration timeline (3-level composition view)
+
+A visual, drill-down view of **how one training iteration's time is composed**,
+layered from a single layer up to the whole pipeline schedule. It reuses the
+projection math in `04-projection-math.md` (same controls, same per-layer times)
+and adds only rendering + one pipeline-schedule simulator. Nothing here changes
+the headline `iteration time` / `tokens/s/GPU` numbers; it explains them.
+
+All times are microseconds (µs) unless noted, scaled to the active GPU tab by
+Step 6 (`rowScaledTime`) exactly like the rest of the site.
+
+## Why three levels
+
+The iteration time is built bottom-up:
+
+```
+module time  --(sum per phase)-->  layer fwd/bwd        (Level 1)
+layer times  --(map to PP/VPP)-->  per-device chunk cost (Level 2)
+device costs --(1F1B schedule)-->  iteration timeline    (Level 3)
+```
+
+Each level answers one question:
+
+- **Level 1 — where does a single layer's time go?** (attn / mlp / a2a)
+- **Level 2 — how is work distributed across pipeline ranks?** (layer granularity)
+- **Level 3 — how do the ranks overlap in time, and where is the bubble?**
+
+## Module → category mapping (Level 1 granularity)
+
+The default minimum granularity is three categories plus an explicit
+"unattributed" bucket, derived from the `module` field (`03-json-schema.md`):
+
+| category | modules | meaning |
+|----------|---------|---------|
+| `attn`   | `attn.norm`, `attn.proj`, `attn.core`, `attn.indexer`, `attn.misc` | attention (varies by `cr`) |
+| `mlp`    | `moe.grouped_gemm`, `moe.shared_expert`, `moe.router` | expert / shared-expert compute (cr-independent) |
+| `a2a`    | `moe.dispatch`, `moe.combine` | EP all-to-all (captured at EP=8) |
+| `misc`   | any `*.misc` / unattributed | casts, scalar, control-flow; kept visible |
+
+Rules:
+
+- `attn.misc` counts as `attn` (it is attention-local unattributed time), while a
+  standalone `misc` category is only used if a non-attn/non-moe module is
+  unmapped. In practice every row maps to `attn` or `mlp`/`a2a`; the `misc`
+  bucket surfaces `*.misc` share the same way `unattributedShare` does today.
+- Because MoE is identical across `cr` (A10), only three representative layers are
+  shown: `cr=0`, `cr=4`, `cr=128`. "Same params/category → show one" reduces to
+  "one per cr".
+- Optional drill-down expands `attn` into `core/indexer/proj/norm` and `mlp` into
+  `grouped_gemm/shared_expert/router`.
+
+## Level 2 — per-device chunk composition
+
+Granularity is **one physical decoder layer**. The projection already maps every
+layer to a `PP*VPP` chunk and a device (`04` Step 2). Level 2 exposes that map:
+
+- Each device (PP rank) is a row; within it, chunks (VPP virtual chunks) are laid
+  out in schedule order, each chunk a run of layers coloured by `cr`.
+- Recomputed layers (their backward replays one forward, `04` Step 1) are marked
+  (hatched / outlined) because they cost extra in `Db`.
+- Non-layer parts are drawn on their owning device: `embedding` on device 0,
+  `output`+`loss`+`MTP` on the last device.
+- **Dedup:** devices with an identical ordered signature
+  `(cr-list, recompute-flags, hasEmb, hasOut, hasMtp)` are drawn once, annotated
+  `×N ranks (d..d)`.
+- The **critical device** (`max Df` / `max Db`, the one that sets the pipeline
+  critical path) is highlighted; each row shows its `Df`/`Db` and share of the
+  critical stage, making load imbalance (the bubble's root cause) visible.
+
+## Level 3 — pipeline schedule (Megatron-2 Figure 4 style)
+
+A Gantt chart: y-axis = device, x-axis = time. Forward cells one colour, backward
+another; VPP virtual chunks distinguished by lightness; bubbles are gaps.
+
+Reference: Narayanan et al., "Efficient Large-Scale Language Model Training on
+GPU Clusters Using Megatron-LM", arXiv:2104.04473, Figure 4 (default 1F1B on top,
+interleaved below).
+
+### Colour scheme
+
+- forward = `--accent` (#4f8cff), backward = `--accent-2` (#36c08f).
+- VPP chunk index shifts lightness (chunk 0 lightest → deeper for higher chunks),
+  mirroring the paper's light/dark model-chunk shading.
+- bubble = empty (optionally a faint diagonal hatch) with the fraction labelled.
+
+### Schedule simulator
+
+The site's analytic pipe time is
+`pipe_compute = (GA + (PP-1)/VPP) * (Df_crit + Db_crit)` (`04` Step 3). To *draw*
+the schedule we simulate 1F1B (optionally interleaved) microbatch ordering and
+let bubbles emerge from the gaps.
+
+Inputs: `PP`, `VPP`, `GA`, and per-virtual-chunk forward/backward cost. Each
+virtual chunk `k` (device `k % PP`, vpp `⌊k/PP⌋`) uses the **exact** per-chunk
+sums from `schedule.chunks` (`f_chunk[k]`, `b_chunk[k]`), with the non-layer
+parts folded into the first (`embedding`) and last (`output`/`loss`/`MTP`) virtual
+chunk so the drawn per-device length stays consistent with `Df`/`Db`. The
+non-interleaved view collapses to one chunk per device with `f=Df[d]`, `b=Db[d]`.
+Large `GA` is capped to `TL_VIS_GA_CAP` (48) microbatches for display only.
+
+Algorithm (interleaved 1F1B, `VPP` model chunks per device, `k mod PP` device
+map):
+
+```
+num_chunks   = PP * VPP
+num_warmup   = min(GA, (PP - 1 - device) * ... )   # standard interleaved warmup
+events[device] = ordered list of {kind:'F'|'B', mb, chunk, start, dur}
+- time advances per device; a device starts an op when both its own timeline and
+  the producing/consuming neighbour dependency are satisfied (F flows forward
+  along devices, B flows backward), with p2p comm assumed zero (A4).
+```
+
+The simulator returns per-device event lists with `start`/`dur`; the drawn
+iteration length `max_device(last_end)` is compared against the analytic
+`pipe_compute` (pre-`calibFactor`). For a balanced schedule they agree; a
+mismatch beyond tolerance is surfaced as a warning rather than hidden. The
+analytic number remains the official iteration time; the Gantt chart is the
+visualization and is scaled so its total width equals the analytic pipe time.
+
+### Interleaving follows VPP
+
+There is no separate interleaved/non-interleaved toggle: the schedule is driven
+directly by the `VPP` control. `VPP=1` renders plain 1F1B (Figure 4 top);
+`VPP>1` renders interleaved 1F1B (Figure 4 bottom), which shrinks the bubble from
+`(PP-1)/GA` to `(PP-1)/(GA*VPP)` (A5). Change `VPP` in the projection controls to
+compare.
+
+## UI: layout + cross-level linkage
+
+- **Layout toggle.** *Tabbed* shows one level at a time (L1/L2/L3 buttons);
+  *Stacked (all + link)* renders all three top-to-bottom with section headings.
+- **Drill-down linkage** (works in both layouts, most visible when stacked):
+  - click a **cr** in Level 1 (or a layer cell in Level 2) → highlights that cr's
+    layers in Level 2 and dims the rest; Level 1 emphasises the matching row.
+  - click a **PP rank** in Level 2 (or a device row in Level 3) → highlights the
+    linked rank in Levels 2 and 3, and Level 1 highlights the cr types present on
+    that rank.
+  - a "Clear" bar removes the selection.
+- **Export.** Level 3's Gantt has *Export SVG* / *Export PNG* buttons; the
+  serializer resolves the theme CSS variables and paints a background so the file
+  is self-contained (good for slides / the paper-style figure).
+- **Fit + zoom.** Level 3 fits the whole schedule in the panel at 1× (no
+  scrollbar) for an at-a-glance overview. A zoom slider (1×–10×) widens the time
+  axis so the per-cell microbatch numbers become readable; above 1× the chart
+  scrolls horizontally. Zoom stretches only the time axis (font/row height fixed).
+- **Cell tooltip.** Hovering a Gantt cell shows `compute <dur> µs` (the op's
+  duration) and `starts @ <t> ms` (its start time measured from the iteration
+  start). The x-axis is that same wall-clock time; gaps between cells are bubble.
+
+## Consistency / self-checks
+
+- Level 1 category sums per cr == `layerTimes(cr).fwd/bwd` (no time lost in
+  categorization).
+- Level 2 per-device `Df/Db` == `project().Df/Db` (same mapping, just detailed).
+- Level 3 simulated iteration length ≈ analytic `pipe_compute` (balanced case);
+  otherwise warn.
+- All three levels react live to the shared projection controls (GPU tab, PP/VPP,
+  GA/GBS/MBS/DP, recompute, manual layer-timing mode).
+
+## Scope / caveats (inherit from 02-assumptions)
+
+- `a2a` is captured at EP=8 intra-node; other EP values are not re-modeled (A7-A9).
+- p2p PP comm is hidden; only the bubble is shown (A4).
+- seq is fixed at the capture value (4096).
+- MI455X tab rescales module times per Step 6 before all three levels are drawn.
diff --git a/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh b/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
new file mode 100755
index 000000000..9a3d16a29
--- /dev/null
+++ b/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
@@ -0,0 +1,252 @@
+#!/bin/bash
+###############################################################################
+# DeepSeek-V4 single-layer, single-cr profiling launcher for the perf
+# *projection* pipeline (examples/deepseek-v4/projection).
+#
+# Produces ONE chrome trace for ONE compression-ratio (cr) layer type, captured
+# under production-representative conditions so the trace can be turned into a
+# clean per-module forward/backward breakdown (see design/01-overview.md):
+#
+#   * seq_length = 4096            (production per-microbatch token count)
+#   * num_layers = 1              (clean per-layer attribution; fits memory @4096)
+#   * compress_ratios = [CR]      (one cr per trace; CR in {0,4,128})
+#   * optimizer = adam, NON-distributed, fp32 states     (dist-opt ON makes the
+#                                                         ROCm Kineto profiler
+#                                                         drop dense/HCA compute
+#                                                         kernels; NOT muon)
+#   * overlap_grad_reduce/param_gather = False           (num_layers=1 breaks
+#                                                         Megatron's chained
+#                                                         param-sync assert; and
+#                                                         with overlap off every
+#                                                         microbatch's compute is
+#                                                         already clean)
+#   * GA = 2  => GBS = 2 * DP * MBS                       (parser averages the
+#                                                         steady profiler window
+#                                                         per microbatch)
+#   * recompute = OFF             (capture pure fwd / pure bwd)
+#   * profiler ON, with_stack=True, window iter 6 -> 7   (kernel -> nn.module)
+#
+# Usage (inside the training container, repo root):
+#   CR=0   ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+#   CR=4   ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+#   CR=128 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+#   MODEL=flash CR=4 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh
+#
+# The trace lands under:
+#   output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/tensorboard/*.pt.trace.json
+###############################################################################
+set -euo pipefail
+set -x
+
+export HF_TOKEN="${HF_TOKEN:-}"
+
+# ---------- What to profile -------------------------------------------------
+export MODEL=${MODEL:-pro}                 # pro | flash
+export CR=${CR:-4}                          # 0 | 4 | 128  (single cr per trace)
+
+# ---------- Model: DeepSeek-V4 (pro or flash) ------------------------------
+# Widths (hidden/heads/kv) come from the model yaml via PRIMUS_MODEL; we only
+# override the MoE / indexer shape knobs the runner exposes, per variant.
+export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml}
+case "$MODEL" in
+  pro)
+    export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro}
+    export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-384}
+    export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6}
+    export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072}
+    export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-1024}
+    ;;
+  flash)
+    export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_flash}
+    export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256}
+    export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6}
+    export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048}
+    export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512}
+    ;;
+  *)
+    echo "[ERROR] MODEL must be 'pro' or 'flash', got '$MODEL'"; exit 1 ;;
+esac
+# Megatron aux-loss-free expert bias needs sigmoid; V4 uses sqrtsoftplus -> off.
+export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False}
+
+# ---------- Single layer per cr (or 3-layer mix) at production seq ---------
+# Pure dense (cr=0) / HCA (cr=128) single layers get CUDA-graph/stream-captured
+# (compute hidden from the trace). CR=mix runs a 3-layer [0,4,128] block: the
+# dynamic CSA (cr=4) layer keeps the block out of graph capture, so all three
+# attention types are visible in one trace and split by kernel name.
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}}
+case "$CR" in
+  0|4|128) export PRIMUS_COMPRESS_RATIOS="[$CR]"; export PRIMUS_TOTAL_LAYERS=1 ;;
+  mix)     export PRIMUS_COMPRESS_RATIOS="[0,4,128]"; export PRIMUS_TOTAL_LAYERS=3 ;;
+  *) echo "[ERROR] CR must be 0, 4, 128 or mix, got '$CR'"; exit 1 ;;
+esac
+
+# ---------- Single-node EP=8 (intra-node; DP=8) ----------------------------
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-1}
+export PRIMUS_EP=${PRIMUS_EP:-8}
+export MBS=${MBS:-1}
+# DP = world/(TP*PP). On one 8-GPU node with TP=PP=1 => DP=8.
+export DP=${DP:-8}
+# GA = 2 so the schedule is F1 B1 F2 B2: the clean (overlap-free) forward is F2
+# and the clean backward is B1. The parser averages the steady profiler window
+# per microbatch, keeping scalar control-flow stalls that survive full-model
+# calibration.
+export GBS=${GBS:-$((2 * DP * MBS))}
+
+# ---------- Optimizer: adam + distributed optimizer (zero1) ----------------
+# Compute (fwd/bwd) is optimizer-independent; we use adam (a) to model the
+# zero1 DP-comm overlap that GA=2 isolates, and (b) to dodge Muon's fp32
+# optimizer-state memory blow-up so seq=4096 fits. The optimizer step itself is
+# modeled separately in the projection (design/04-projection-math.md Step 4).
+export OPTIMIZER=${OPTIMIZER:-adam}
+# CRITICAL: distributed optimizer (zero1) MUST be off. With it on, the ROCm
+# Kineto profiler silently drops the compute GPU kernels for pure dense (cr=0)
+# and HCA (cr=128) layers (only optimizer/comm/elementwise get recorded);
+# turning it off makes every cr's kernels visible. dist-opt has no bearing on
+# the captured fwd/bwd compute, which is what the projection needs (the
+# optimizer step is modeled analytically, design/04 Step 4).
+export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False}
+# Overlap OFF. With num_layers=1, Megatron's chained param-gather sync trips an
+# assertion (param_and_grad_buffer.start_param_sync). More importantly, with
+# overlap off every microbatch's compute is already overlap-free — exactly the
+# clean per-kernel time we want; the parser averages the steady profiler window
+# per microbatch. (Production DP comm is assumed hidden in the projection anyway; A2.)
+export PRIMUS_OVERLAP_GRAD_REDUCE=${PRIMUS_OVERLAP_GRAD_REDUCE:-False}
+export PRIMUS_OVERLAP_PARAM_GATHER=${PRIMUS_OVERLAP_PARAM_GATHER:-False}
+# Distributed optimizer (zero1) is the default. The Kineto trace drops the
+# compute GPU kernels for pure dense(cr=0)/HCA(cr=128) layers ONLY when the
+# distributed optimizer is on; turning it off makes all kernels visible (but
+# then precision-aware optimizer must be off too -> fp32 optimizer states).
+DISTOPT_ARGS=(--use_distributed_optimizer "$USE_DISTRIBUTED_OPTIMIZER")
+if [ "$USE_DISTRIBUTED_OPTIMIZER" = "False" ]; then
+  DISTOPT_ARGS+=(--use_precision_aware_optimizer False --main_grads_dtype fp32 --exp_avg_dtype fp32 --exp_avg_sq_dtype fp32)
+fi
+
+# ---------- Perf knobs (production V4 attn backends + Turbo MoE) ------------
+export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-True}
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True}
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v1}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v1}
+export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-True}
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True}
+export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True}
+export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1}
+export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1}
+export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1}
+export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1}
+export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1}
+export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1}
+export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1}
+export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0}
+export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1}
+
+TURBO_DEEPEP_CLI_ARGS=()
+if [ "$USE_TURBO_DEEPEP" = "True" ]; then
+  export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-80}
+  export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False}
+  export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32}
+  export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False}
+  TURBO_DEEPEP_CLI_ARGS=(
+    --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU"
+    --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM"
+    --moe_router_dtype "$MOE_ROUTER_DTYPE"
+    --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP"
+  )
+fi
+
+export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1}
+# Disable the loss NaN/Inf validation (check_for_nan_in_loss_and_grad). Its
+# torch.isnan/isinf checks in loss_func force a device->host sync once per
+# microbatch; in a 1-layer capture (nothing to overlap) the profiler bills that
+# sync wait as a multi-ms "stall" kernel under the loss/rerun frames, which then
+# pollutes per-layer attribution and gets multiplied by the layer count in the
+# projection. It is a per-step validation cost, not per-layer compute, so we
+# turn it off for a clean capture (the projection models it separately if needed).
+# Profiler window must land in the STEADY state: the first ~9 iters are warm-up
+# (kernel autotune / hipBLASLt + Triton compilation) and have noisy, inflated
+# per-iter times (and inflated comm-kernel durations as ranks desync on the
+# autotuning rank); from ~iter 10 onward the iteration time is stable. Default
+# to a late, multi-step window so the captured steps are clean; all three are
+# env-overridable.
+export TRAIN_ITERS=${TRAIN_ITERS:-22}
+export PROFILE_STEP_START=${PROFILE_STEP_START:-16}
+export PROFILE_STEP_END=${PROFILE_STEP_END:-19}
+
+# ---------- Profiler: trace with python stacks for kernel->module ----------
+export PROFILE=True
+# Optional GPU-only trace (drop CPU activity). With CPU activity on, pure dense
+# (cr=0) / HCA (cr=128) layers' compute kernels can vanish from the trace
+# (profiler/stream-capture interaction); GPU-only capture brings them back, at
+# the cost of CPU-side module/stack attribution (fall back to kernel-name rules).
+PROFILER_ACTIVITY_ARGS=()
+if [ "${DISABLE_PROFILER_CPU:-False}" = "True" ]; then
+  PROFILER_ACTIVITY_ARGS=(--disable_profiler_activity_cpu True)
+fi
+export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"}
+export PRIMUS_TEAM=${PRIMUS_TEAM:-amd}
+export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-projection_${MODEL}_cr${CR}_seq${PRIMUS_SEQ_LENGTH}_ep${PRIMUS_EP}}
+
+if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then
+  echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH"
+  echo "Run: git submodule update --init --recursive"
+  exit 1
+fi
+
+mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME"
+
+./primus-cli direct \
+  -- train pretrain --config "$EXP" \
+  --backend_path "$BACKEND_PATH" \
+  --num_layers "$PRIMUS_TOTAL_LAYERS" \
+  --train_iters "$TRAIN_ITERS" \
+  --lr_warmup_iters 0 \
+  --lr_decay_iters "$TRAIN_ITERS" \
+  --micro_batch_size "$MBS" \
+  --global_batch_size "$GBS" \
+  --seq_length "$PRIMUS_SEQ_LENGTH" \
+  --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \
+  --rope_type rope \
+  --tensor_model_parallel_size "$PRIMUS_TP" \
+  --pipeline_model_parallel_size "$PRIMUS_PP" \
+  --expert_model_parallel_size "$PRIMUS_EP" \
+  --num_experts "$PRIMUS_NUM_EXPERTS" \
+  --moe_router_topk "$PRIMUS_MOE_TOPK" \
+  --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \
+  --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \
+  --index_topk "$PRIMUS_INDEX_TOPK" \
+  --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \
+  --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \
+  --mtp_num_layers 0 \
+  --mock_data True \
+  --optimizer "$OPTIMIZER" \
+  "${DISTOPT_ARGS[@]}" \
+  --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \
+  --use_turbo_attention "$USE_TURBO_ATTENTION" \
+  --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \
+  --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \
+  --use_v4_fp8_indexer "$USE_V4_FP8_INDEXER" \
+  --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \
+  --use_turbo_deepep "$USE_TURBO_DEEPEP" \
+  "${TURBO_DEEPEP_CLI_ARGS[@]}" \
+  --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \
+  --moe_use_legacy_grouped_gemm False \
+  --fp8 null \
+  --fp8_recipe null \
+  --recompute_num_layers 0 \
+  --check_for_nan_in_loss_and_grad False \
+  --overlap_grad_reduce "$PRIMUS_OVERLAP_GRAD_REDUCE" \
+  --overlap_param_gather "$PRIMUS_OVERLAP_PARAM_GATHER" \
+  --disable_last_saving True \
+  --disable_wandb True \
+  --disable_tensorboard False \
+  --profile True \
+  --use_pytorch_profiler True \
+  "${PROFILER_ACTIVITY_ARGS[@]}" \
+  --profile_step_start "$PROFILE_STEP_START" \
+  --profile_step_end "$PROFILE_STEP_END" \
+  2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt"
diff --git a/examples/deepseek-v4/projection/site/assets/app.js b/examples/deepseek-v4/projection/site/assets/app.js
new file mode 100644
index 000000000..b1b0b32bd
--- /dev/null
+++ b/examples/deepseek-v4/projection/site/assets/app.js
@@ -0,0 +1,1567 @@
+"use strict";
+
+// DeepSeek-V4 performance projection — static, no-build. Implements the math in
+// design/04-projection-math.md. All breakdown times are microseconds (us) for
+// one microbatch (seq from capture); the projection scales to a full model run.
+
+const STATE = {
+  data: null,
+  gpu: "MI355X",
+  controls: null,
+  // iteration-timeline view (design/07): active level (1|2|3). Level 3
+  // interleaving follows the VPP control directly (no separate toggle).
+  tlLevel: 1,
+  // layout: "tabs" (one level via buttons) or "stacked" (all three + linkage).
+  tlView: "stacked",
+  // Level 3 Gantt horizontal zoom (1 = fit-to-view, no scrollbar; >1 widens the
+  // chart and reveals per-cell microbatch numbers, with horizontal scroll).
+  tlZoom: 1,
+  // cross-level selection for drill-down highlight. cr: focus a compression
+  // ratio; devices: focus one or more PP ranks. Mutually exclusive.
+  tlSel: { cr: null, devices: null },
+};
+
+const $ = (sel) => document.querySelector(sel);
+const el = (tag, attrs = {}, ...kids) => {
+  const n = document.createElement(tag);
+  for (const [k, v] of Object.entries(attrs)) {
+    if (k === "class") n.className = v;
+    else if (k === "html") n.innerHTML = v;
+    else n.setAttribute(k, v);
+  }
+  for (const kid of kids) n.append(kid?.nodeType ? kid : document.createTextNode(kid ?? ""));
+  return n;
+};
+const fmt = (x, d = 1) =>
+  x == null || !isFinite(x) ? "—" : Number(x).toLocaleString(undefined, { maximumFractionDigits: d, minimumFractionDigits: d });
+const fmtInt = (x) => (x == null || !isFinite(x) ? "—" : Math.round(x).toLocaleString());
+
+// ---------------------------------------------------------------------------
+// Load
+// ---------------------------------------------------------------------------
+function modelFromQuery() {
+  const m = new URLSearchParams(location.search).get("model");
+  return m === "flash" || m === "pro" ? m : "pro";
+}
+
+async function loadModel(model) {
+  const res = await fetch(`./data/${model}.json`, { cache: "no-store" });
+  if (!res.ok) throw new Error(`failed to load data/${model}.json (${res.status})`);
+  return res.json();
+}
+
+// ---------------------------------------------------------------------------
+// Controls
+// ---------------------------------------------------------------------------
+function defaultControls(data) {
+  const hw = data.hardware || {};
+  const m355 = hw.MI355X || { peak_tflops_bf16: 2500, hbm_bandwidth_gbps: 8000 };
+  const m455 = hw.MI455X || { peak_tflops_bf16: 5000, hbm_bandwidth_gbps: 16000 };
+  const isPro = data.model === "pro";
+  const optBytes = data.optimizer?.bytes_per_param && data.optimizer.bytes_per_param !== 18
+    ? data.optimizer.bytes_per_param : 30;
+  return {
+    world: isPro ? 256 : 32,
+    pp: isPro ? 16 : 4, vpp: 1, ep: 8, tp: 1, cp: 1,
+    mbs: 1, gbs: isPro ? 1024 : 256,
+    recompute: isPro ? "full" : "first-n",
+    recomputeLayers: isPro ? 0 : 3,
+    ppLayout: data.model_config?.pipeline_layout || "",
+    optEff: 0.7, computeEff: 1.0, calibFactor: 0.91, bytesPerParam: optBytes,
+    peak355: m355.peak_tflops_bf16, bw355: m355.hbm_bandwidth_gbps,
+    peak455: m455.peak_tflops_bf16, bw455: m455.hbm_bandwidth_gbps,
+    // Modeling mode: "trace" (derive per-layer fwd/bwd from the breakdown JSON)
+    // or "manual" (user types per-cr fwd/bwd directly). Manual values are stored
+    // per GPU because a hand-entered time already targets a specific GPU, so the
+    // MI355->MI455 scaling is bypassed in manual mode.
+    modelMode: "trace",
+    man: {
+      MI355X: emptyManual(),
+      MI455X: emptyManual(),
+    },
+  };
+}
+
+// Per-cr manual fwd/bwd holders (µs). null = "not set yet" -> falls back to the
+// trace-derived value, so toggling into manual mode never changes the result
+// until the user actually edits a field.
+function emptyManual() {
+  return {
+    f0: null, b0: null, f4: null, b4: null, f128: null, b128: null,
+    // non-layer + MTP overrides (per iteration, one device)
+    emb_f: null, emb_b: null, out_f: null, out_b: null, loss_f: null, loss_b: null, mtp_f: null, mtp_b: null,
+  };
+}
+
+const MANUAL_CR_KEYS = { "0": ["f0", "b0"], "4": ["f4", "b4"], "128": ["f128", "b128"] };
+// Manually-overridable non-layer parts: key prefix, label, and the JSON section
+// (mtp is synthesized, not a non_layer entry).
+const MANUAL_NONLAYER = [
+  { key: "emb", label: "embedding", which: "embedding" },
+  { key: "out", label: "output", which: "output" },
+  { key: "loss", label: "loss", which: "loss" },
+  { key: "mtp", label: "MTP", which: "mtp" },
+];
+const MANUAL_NONLAYER_KEYS = MANUAL_NONLAYER.flatMap((n) => [`${n.key}_f`, `${n.key}_b`]);
+
+const CONTROL_DEFS = [
+  { key: "world", label: "World size (GPUs)", kind: "int" },
+  { key: "pp", label: "PP (pipeline)", kind: "int" },
+  { key: "vpp", label: "VPP (interleave)", kind: "int" },
+  { key: "ep", label: "EP (expert)", kind: "int" },
+  { key: "dp", label: "DP (derived)", kind: "ro" },
+  { key: "tp", label: "TP (tensor)", kind: "int" },
+  { key: "cp", label: "CP (context)", kind: "int" },
+  { key: "mbs", label: "Micro batch size", kind: "int" },
+  { key: "gbs", label: "Global batch size", kind: "int" },
+  { key: "recompute", label: "Recompute", kind: "sel" },
+  { key: "recomputeLayers", label: "Recompute layers", kind: "int" },
+  { key: "ppLayout", label: "PP layout", kind: "txt", full: true },
+  { key: "bytesPerParam", label: "Optim bytes/param", kind: "int" },
+  { key: "calibFactor", label: "Calibration factor", kind: "f" },
+  { key: "optEff", label: "Optim efficiency", kind: "f" },
+  { key: "computeEff", label: "MI455 compute eff", kind: "f" },
+  { key: "peak355", label: "MI355 peak TFLOPs", kind: "int" },
+  { key: "bw355", label: "MI355 HBM GB/s", kind: "int" },
+  { key: "peak455", label: "MI455 peak TFLOPs", kind: "int" },
+  { key: "bw455", label: "MI455 HBM GB/s", kind: "int" },
+];
+
+// DP is derived from the user-set world size: DP = world / (PP*TP*CP). EP is a
+// sub-grouping of DP (EP <= DP) and does not multiply world size.
+const derivedDP = (c) => {
+  const denom = c.pp * c.tp * c.cp;
+  if (!Number.isFinite(denom) || denom <= 0 || !Number.isFinite(c.world)) return NaN;
+  return c.world / denom;
+};
+
+function parseControlValue(input, kind) {
+  if (kind === "sel" || kind === "txt") return input.value;
+  if (kind === "int") return input.value.trim() === "" ? NaN : Number(input.value);
+  if (kind === "f") return input.value.trim() === "" ? NaN : Number(input.value);
+  return input.value;
+}
+
+const isPositiveInt = (x) => Number.isInteger(x) && x > 0;
+const isNonNegativeInt = (x) => Number.isInteger(x) && x >= 0;
+const isPositiveNumber = (x) => Number.isFinite(x) && x > 0;
+const isNonNegativeNumber = (x) => Number.isFinite(x) && x >= 0;
+
+function renderControls() {
+  const grid = $("#controls-grid");
+  grid.innerHTML = "";
+  const c = STATE.controls;
+  for (const def of CONTROL_DEFS) {
+    const { key, label, kind } = def;
+    const field = el("div", { class: "field" });
+    if (def.full) field.classList.add("field--full");
+    field.append(el("span", {}, label));
+    let input;
+    if (kind === "sel") {
+      input = el("select", { id: `ctl-${key}` });
+      for (const opt of ["none", "full", "first-n"]) {
+        const o = el("option", { value: opt }, opt);
+        if (c[key] === opt) o.selected = true;
+        input.append(o);
+      }
+    } else if (kind === "ro") {
+      const dp = derivedDP(c);
+      input = el("input", { id: `ctl-${key}`, value: Number.isFinite(dp) ? dp : "—", disabled: "true" });
+    } else if (kind === "txt") {
+      input = el("input", { id: `ctl-${key}`, type: "text", value: c[key] || "" });
+    } else {
+      input = el("input", { id: `ctl-${key}`, type: "number", value: c[key], step: kind === "f" ? "0.05" : "1" });
+    }
+    if (kind !== "ro") {
+      input.addEventListener("change", () => {
+        c[key] = parseControlValue(input, kind);
+        renderAll();
+      });
+    }
+    field.append(input);
+    grid.append(field);
+  }
+}
+
+// Prefill the active GPU's manual fields from the trace-derived times so that
+// switching into manual mode (or switching GPU while in manual mode) starts from
+// the current baseline instead of empty boxes. Already-set fields are kept.
+function prefillManual(gpu) {
+  const c = STATE.controls;
+  const m = c.man[gpu];
+  const lt = {};
+  for (const cr of ["0", "4", "128"]) {
+    const [fk, bk] = MANUAL_CR_KEYS[cr];
+    const t = layerTimes(STATE.data, cr, gpu, c);
+    lt[cr] = t;
+    if (!isPositiveNumber(m[fk])) m[fk] = Math.round(t.fwd);
+    if (!isPositiveNumber(m[bk])) m[bk] = Math.round(t.bwd);
+  }
+  // non-layer parts (embedding / output / loss). Skip seeding fields whose trace
+  // value is 0 (e.g. output/loss backward) — leaving them unset shows the "0"
+  // placeholder and falls back to trace, instead of pinning an explicit 0.
+  const seed = (key, val) => { if (!isPositiveNumber(m[key]) && Math.round(val) > 0) m[key] = Math.round(val); };
+  for (const which of ["embedding", "output", "loss"]) {
+    const key = NONLAYER_KEY[which];
+    seed(`${key}_f`, nonLayer(STATE.data, which, "forward", gpu, c));
+    seed(`${key}_b`, nonLayer(STATE.data, which, "backward", gpu, c));
+  }
+  // MTP (only when the model uses it)
+  if ((STATE.data.model_config.mtp_num_layers || 0) > 0) {
+    const base = mtpTimes(STATE.data, gpu, c, lt);
+    seed("mtp_f", base.fwd);
+    seed("mtp_b", base.bwd);
+  }
+}
+
+function renderModeSwitch() {
+  document.querySelectorAll(".mode-tab").forEach((t) => {
+    t.classList.toggle("is-active", t.dataset.mode === STATE.controls.modelMode);
+  });
+}
+
+function renderManualGrid() {
+  const host = $("#manual-grid");
+  if (!host) return;
+  host.hidden = STATE.controls.modelMode !== "manual";
+  host.innerHTML = "";
+  if (host.hidden) return;
+  const c = STATE.controls, gpu = STATE.gpu;
+  const counts = STATE.data.model_config.cr_layer_counts || {};
+  const m = c.man[gpu];
+  const header = el("div", { class: "manual-grid__header" });
+  header.append(el("p", { class: "muted manual-grid__hint" },
+    `Per-layer and non-layer fwd/bwd (µs) for ${gpu}. Set values override the trace; placeholders show the current trace baseline. Embedding is on PP stage 0; output / loss / MTP are on the last PP stage.`));
+  const resetBtn = el("button", { class: "manual-reset" }, "Restore defaults");
+  resetBtn.addEventListener("click", () => {
+    c.man[gpu] = emptyManual();
+    prefillManual(gpu);
+    renderAll();
+  });
+  header.append(resetBtn);
+  host.append(header);
+
+  // one fwd/bwd input row
+  const makeRow = (labelNode, fk, bk, traceF, traceB) => {
+    const row = el("div", { class: "manual-row" });
+    row.append(labelNode);
+    for (const [label, key, traceVal] of [["fwd", fk, traceF], ["bwd", bk, traceB]]) {
+      const field = el("div", { class: "field" });
+      field.append(el("span", {}, `${label} µs`));
+      const input = el("input", {
+        id: `man-${gpu}-${key}`, type: "number", step: "1", min: "0",
+        placeholder: String(Math.round(traceVal || 0)),
+      });
+      if (isPositiveNumber(m[key])) input.value = m[key];
+      input.addEventListener("change", () => {
+        const v = input.value.trim();
+        m[key] = v === "" ? null : Number(v);
+        renderAll();
+      });
+      field.append(input);
+      row.append(field);
+    }
+    return row;
+  };
+
+  const rows = el("div", { class: "manual-rows" });
+  for (const cr of ["0", "4", "128"]) {
+    if (!(counts[cr] > 0)) continue;
+    const [fk, bk] = MANUAL_CR_KEYS[cr];
+    const t = layerTimes(STATE.data, cr, gpu, c);
+    const lab = el("span", { class: `manual-row__lab cr-tag cr-${cr}` }, `cr=${cr} ×${counts[cr] || 0}`);
+    rows.append(makeRow(lab, fk, bk, t.fwd, t.bwd));
+  }
+  host.append(rows);
+
+  // non-layer + MTP rows
+  host.append(el("p", { class: "muted manual-grid__subhead" }, "Non-layer (once per iteration)"));
+  const nlRows = el("div", { class: "manual-rows" });
+  const lt = {};
+  for (const cr of ["0", "4", "128"]) lt[cr] = layerTimes(STATE.data, cr, gpu, c);
+  for (const nl of MANUAL_NONLAYER) {
+    if (nl.which === "mtp" && !((STATE.data.model_config.mtp_num_layers || 0) > 0)) continue;
+    let tF, tB;
+    if (nl.which === "mtp") {
+      const base = mtpTimes(STATE.data, gpu, c, lt);
+      tF = base.fwd; tB = base.bwd;
+    } else {
+      tF = nonLayer(STATE.data, nl.which, "forward", gpu, c);
+      tB = nonLayer(STATE.data, nl.which, "backward", gpu, c);
+    }
+    const lab = el("span", { class: "manual-row__lab manual-row__lab--nl" }, nl.label);
+    nlRows.append(makeRow(lab, `${nl.key}_f`, `${nl.key}_b`, tF, tB));
+  }
+  host.append(nlRows);
+}
+
+// ---------------------------------------------------------------------------
+// Hardware scaling (Step 6)
+// ---------------------------------------------------------------------------
+function rowScaledTime(row, gpu, c) {
+  if (gpu === "MI355X") return row.time_us;
+  const compute = row.class === "compute_bound";
+  if (compute) return row.time_us * (c.peak355 / c.peak455) / c.computeEff;
+  return row.time_us * (c.bw355 / c.bw455);
+}
+function rowTflops(row, scaledTimeUs) {
+  if (!row.flops || !scaledTimeUs) return null;
+  return row.flops / (scaledTimeUs * 1e-6) / 1e12;
+}
+
+const sumTime = (list, gpu, c) => list.reduce((a, r) => a + rowScaledTime(r, gpu, c), 0);
+const sumFlops = (list) => list.reduce((a, r) => a + (r.flops || 0), 0);
+
+// ---------------------------------------------------------------------------
+// Per-layer base (Step 0/1)
+// ---------------------------------------------------------------------------
+function layerTimes(data, cr, gpu, c) {
+  const L = data.layers[cr];
+  if (!L) return { fwd: 0, bwd: 0, fFlops: 0, bFlops: 0 };
+  const aF = L.attention.forward, aB = L.attention.backward;
+  const mF = L.moe.forward, mB = L.moe.backward;
+  let fwd = sumTime(aF, gpu, c) + sumTime(mF, gpu, c);
+  let bwd = sumTime(aB, gpu, c) + sumTime(mB, gpu, c);
+  let fFlops = sumFlops(aF) + sumFlops(mF);
+  let bFlops = sumFlops(aB) + sumFlops(mB);
+  return { fwd, bwd, fFlops, bFlops };
+}
+
+// Effective per-layer time used by the projection. In "manual" mode a set field
+// overrides the trace-derived time for the active GPU; unset fields fall back to
+// trace. FLOPs always stay analytic/trace-derived (manual only overrides time),
+// so TFLOP/s/GPU remains meaningful.
+function effectiveLayerTimes(data, cr, gpu, c) {
+  const trace = layerTimes(data, cr, gpu, c);
+  if (c.modelMode !== "manual") return trace;
+  const m = (c.man && c.man[gpu]) || {};
+  const [fk, bk] = MANUAL_CR_KEYS[cr] || [];
+  return {
+    fwd: isPositiveNumber(m[fk]) ? m[fk] : trace.fwd,
+    bwd: isPositiveNumber(m[bk]) ? m[bk] : trace.bwd,
+    fFlops: trace.fFlops,
+    bFlops: trace.bFlops,
+  };
+}
+
+function expandLayoutRepeats(layout) {
+  let out = layout;
+  let prev;
+  do {
+    prev = out;
+    out = out.replace(/\(([^()]+)\)\*(\d+)/g, (_m, body, count) => body.repeat(Number(count)));
+  } while (out !== prev);
+  return out;
+}
+
+function parsePipelineLayout(layout, numLayers, chunks) {
+  const raw = String(layout ?? "").trim();
+  if (!raw) return { ok: true, stages: null, counts: [], normalized: "", message: "empty layout; using balanced fallback" };
+  const normalized = expandLayoutRepeats(raw.replace(/^['"]|['"]$/g, ""));
+  if (/[()]/.test(normalized)) {
+    return { ok: false, stages: null, counts: [], normalized, message: "unsupported nested or malformed repeat expression" };
+  }
+  const specs = normalized.split("|").map((x) => x.trim()).filter(Boolean);
+  if (specs.length !== chunks) {
+    return {
+      ok: false,
+      stages: null,
+      counts: specs.map((spec) => [...spec.matchAll(/[tT](?:\*(\d+))?/g)].reduce((a, m) => a + Number(m[1] || 1), 0)),
+      normalized,
+      message: `layout has ${specs.length} stages, expected PP*VPP=${chunks}`,
+    };
+  }
+  let nextLayer = 0;
+  const out = [];
+  const counts = [];
+  for (const spec of specs) {
+    const layers = [];
+    for (const m of spec.matchAll(/[tT](?:\*(\d+))?/g)) {
+      const n = m[1] ? Number(m[1]) : 1;
+      for (let i = 0; i < n && nextLayer < numLayers; i++) layers.push(nextLayer++);
+    }
+    counts.push(layers.length);
+    out.push(layers);
+  }
+  if (nextLayer !== numLayers) {
+    return {
+      ok: false,
+      stages: null,
+      counts,
+      normalized,
+      message: `layout maps ${nextLayer} decoder layers, expected ${numLayers}`,
+    };
+  }
+  return { ok: true, stages: out, counts, normalized, message: "layout applied" };
+}
+
+function recomputeLayer(c, stageOrdinal) {
+  if (c.recompute === "full") return true;
+  if (c.recompute === "first-n") return stageOrdinal < Math.max(0, c.recomputeLayers || 0);
+  return false;
+}
+
+function validateControls(data, c, gpu = STATE.gpu) {
+  const errors = [];
+  const warnings = [];
+  const ints = ["world", "pp", "vpp", "ep", "tp", "cp", "mbs", "gbs", "bytesPerParam", "peak355", "bw355", "peak455", "bw455"];
+  for (const key of ints) {
+    if (!isPositiveInt(c[key])) errors.push(`${key} must be a positive integer.`);
+  }
+  if (!isNonNegativeInt(c.recomputeLayers)) errors.push("recomputeLayers must be a non-negative integer.");
+  for (const key of ["calibFactor", "optEff", "computeEff"]) {
+    if (!isPositiveNumber(c[key])) errors.push(`${key} must be a positive number.`);
+  }
+  if (!["none", "full", "first-n"].includes(c.recompute)) errors.push(`unsupported recompute mode: ${c.recompute}`);
+
+  const denom = c.pp * c.tp * c.cp;
+  const dp = derivedDP(c);
+  if (isPositiveInt(c.world) && isPositiveInt(denom) && c.world % denom !== 0) {
+    errors.push(`world must be divisible by PP*TP*CP (${denom}); got world=${c.world}.`);
+  }
+  if (Number.isFinite(dp) && isPositiveInt(c.gbs) && isPositiveInt(c.mbs) && !Number.isInteger(c.gbs / (dp * c.mbs))) {
+    errors.push(`GA must be an integer: GBS / (DP*MBS) = ${c.gbs} / (${dp}*${c.mbs}).`);
+  }
+  if (Number.isFinite(dp) && isPositiveInt(c.ep) && c.ep > dp) {
+    errors.push(`EP must be <= DP; got EP=${c.ep}, DP=${dp}.`);
+  }
+
+  const chunks = c.pp * c.vpp;
+  const layout = parsePipelineLayout(c.ppLayout, data.model_config.compress_ratios.length, chunks);
+  if (!layout.ok) errors.push(`PP layout invalid: ${layout.message}.`);
+
+  if (c.modelMode === "manual") {
+    const man = (c.man && c.man[gpu]) || {};
+    for (const cr of ["0", "4", "128"]) {
+      const count = data.model_config.cr_layer_counts?.[cr] || 0;
+      if (!count) continue; // cr not used by this model; ignore its inputs
+      for (const k of MANUAL_CR_KEYS[cr]) {
+        const v = man[k];
+        if (v != null && v !== "" && !isNonNegativeNumber(Number(v))) {
+          errors.push(`manual ${k} (cr=${cr}) must be a non-negative number (µs).`);
+        }
+      }
+    }
+    for (const k of MANUAL_NONLAYER_KEYS) {
+      const v = man[k];
+      if (v != null && v !== "" && !isNonNegativeNumber(Number(v))) {
+        errors.push(`manual ${k} must be a non-negative number (µs).`);
+      }
+    }
+    warnings.push(`Manual mode for ${gpu}: per-cr layer, non-layer (embedding/output/loss) and MTP fwd/bwd you set override the trace; unset fields fall back to trace.`);
+  }
+
+  const captureEp = data.capture?.ep;
+  if (captureEp && c.ep !== captureEp) {
+    warnings.push(`EP=${c.ep} is only partially modeled; traces captured EP=${captureEp}, so MoE dispatch/combine are not re-estimated.`);
+  } else {
+    warnings.push(`EP is a consistency control only; captured MoE dispatch/combine are reused from EP=${captureEp || 8}.`);
+  }
+  if (c.tp !== 1 || c.cp !== 1) {
+    warnings.push("TP/CP values affect derived DP/GA/optimizer only; TP/CP layer compute and communication are not re-modeled.");
+  }
+  if (gpu === "MI355X") warnings.push("MI455 peak/bandwidth/compute-eff controls affect only the MI455X tab.");
+  warnings.push(`Sequence length is fixed at captured seq=${data.capture.seq_length}; changing seq is not currently exposed.`);
+
+  return { errors, warnings, layout, dp };
+}
+
+function nonLayer(data, which, phase, gpu, c) {
+  const bd = data.non_layer[which];
+  return bd ? sumTime(bd[phase], gpu, c) : 0;
+}
+
+// Non-layer time with manual override (embedding / output / loss). Falls back to
+// the trace time when the field is unset or not in manual mode.
+const NONLAYER_KEY = { embedding: "emb", output: "out", loss: "loss" };
+function effectiveNonLayer(data, which, phase, gpu, c) {
+  const trace = nonLayer(data, which, phase, gpu, c);
+  if (c.modelMode !== "manual") return trace;
+  const m = (c.man && c.man[gpu]) || {};
+  const v = m[`${NONLAYER_KEY[which]}_${phase === "forward" ? "f" : "b"}`];
+  return isPositiveNumber(v) ? v : trace;
+}
+function nonLayerFlops(data, which, phase) {
+  const bd = data.non_layer[which];
+  return bd ? sumFlops(bd[phase]) : 0;
+}
+
+function mtpTimes(data, gpu, c, lt) {
+  const cfg = data.model_config;
+  const depth = cfg.mtp_num_layers || 0;
+  if (!depth) return { fwd: 0, bwd: 0, ehUs: 0 };
+
+  const cr = String((cfg.mtp_compress_ratios && cfg.mtp_compress_ratios[0]) || 0);
+  const inner = lt[cr] || lt["0"];
+  const innerBwd = inner.bwd + (c.recompute === "full" ? inner.fwd : 0);
+  const outF = nonLayer(data, "output", "forward", gpu, c) + nonLayer(data, "loss", "forward", gpu, c);
+  const outB = nonLayer(data, "output", "backward", gpu, c) + nonLayer(data, "loss", "backward", gpu, c);
+
+  // No MTP trace row exists yet. Approximate eh_proj as GEMM-like work scaled
+  // from the measured output projection time, then split fwd/bwd as 1:2.
+  const af = data.analytic_flops || {};
+  const mtp = af.mtp || {};
+  const outFlops = af.output_flops || 0;
+  const outUs = outF + outB;
+  const ehUs = outFlops > 0 ? outUs * ((mtp.eh_proj_flops || 0) / outFlops) : 0;
+
+  return {
+    fwd: depth * (inner.fwd + outF) + ehUs / 3,
+    bwd: depth * (innerBwd + outB) + (2 * ehUs) / 3,
+    ehUs,
+  };
+}
+
+// MTP time with manual override. Falls back to the analytic estimate.
+function effectiveMtp(data, gpu, c, lt) {
+  const base = mtpTimes(data, gpu, c, lt);
+  if (c.modelMode !== "manual" || !((data.model_config.mtp_num_layers || 0) > 0)) return base;
+  const m = (c.man && c.man[gpu]) || {};
+  return {
+    fwd: isPositiveNumber(m.mtp_f) ? m.mtp_f : base.fwd,
+    bwd: isPositiveNumber(m.mtp_b) ? m.mtp_b : base.bwd,
+    ehUs: base.ehUs,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Param estimate (Step 4)
+// ---------------------------------------------------------------------------
+function estimateParams(cfg) {
+  // Prefer the exact V4 count emitted by parse_trace (MLA low-rank attention +
+  // MoE + tied-free embedding/output). Fall back to a crude estimate.
+  if (cfg.total_params) return cfg.total_params;
+  const h = cfg.hidden_size, exp = cfg.num_experts, mff = cfg.moe_ffn_hidden_size;
+  const sff = cfg.moe_shared_expert_intermediate_size || mff, V = cfg.vocab_size, L = cfg.num_layers;
+  const perLayer = 4 * h * h + exp * 3 * h * mff + 3 * h * sff;
+  return L * perLayer + 2 * V * h;
+}
+
+function estimateOptimizerParams(data, c, dp) {
+  const cfg = data.model_config;
+  const totalParams = estimateParams(cfg);
+  const pp = Math.max(1, c.pp);
+  const tp = Math.max(1, c.tp);
+  const dpSize = Math.max(1, dp);
+
+  // total_params is a full-model count. CP does not shard weights; EP is already
+  // represented in the full expert count and cancels out with the DP replica
+  // count for ZeRO-1 optimizer sharding, so the average rank owns total/(PP*TP)
+  // model params and updates total/(PP*TP*DP) params per optimizer step.
+  const localModelParams = totalParams / (pp * tp);
+  const perRankParams = localModelParams / dpSize;
+  const measuredUs = data.optimizer?.time_us ?? null;
+  return { totalParams, localModelParams, perRankParams, measuredUs };
+}
+
+// ---------------------------------------------------------------------------
+// Pipeline mapping + projection (Steps 2-5)
+// ---------------------------------------------------------------------------
+function project(data, gpu, c, validation = validateControls(data, c, gpu)) {
+  if (validation.errors.length) return null;
+  const cfg = data.model_config;
+  const crs = cfg.compress_ratios;
+  const L = crs.length;
+  const dp = validation.dp;
+  const world = c.world;
+  const lt = {};
+  for (const cr of ["0", "4", "128"]) lt[cr] = effectiveLayerTimes(data, cr, gpu, c);
+
+  // Step 2: assign layers to PP*VPP chunks -> devices
+  const C = c.pp * c.vpp;
+  const perChunk = Math.ceil(L / C);
+  const Df = new Array(c.pp).fill(0), Db = new Array(c.pp).fill(0);
+  const layoutInfo = validation.layout;
+  const layoutStages = layoutInfo.ok ? layoutInfo.stages : null;
+  const stageOrdinals = new Array(c.pp).fill(0);
+  // Per virtual-chunk detail (k = 0..C-1): device = k % PP, vpp index = floor(k/PP).
+  // Kept for the iteration-timeline view (design/07); does not affect Df/Db math.
+  const chunks = [];
+  for (let k = 0; k < C; k++) {
+    chunks.push({ chunk: k, device: k % c.pp, vpp: Math.floor(k / c.pp), layers: [], fwd: 0, bwd: 0 });
+  }
+  const addLayer = (chunkIdx, i) => {
+    const dev = chunkIdx % c.pp;
+    const t = lt[String(crs[i])] || { fwd: 0, bwd: 0 };
+    const recompute = recomputeLayer(c, stageOrdinals[dev]++);
+    const effBwd = t.bwd + (recompute ? t.fwd : 0);
+    Df[dev] += t.fwd;
+    Db[dev] += effBwd;
+    const ch = chunks[chunkIdx];
+    ch.layers.push({ globalIdx: i, cr: crs[i], recompute, fwd: t.fwd, bwd: effBwd });
+    ch.fwd += t.fwd;
+    ch.bwd += effBwd;
+  };
+  if (layoutStages) {
+    layoutStages.forEach((layers, chunk) => {
+      for (const i of layers) addLayer(chunk, i);
+    });
+  } else {
+    for (let i = 0; i < L; i++) addLayer(Math.floor(i / perChunk), i);
+  }
+  // non-layer parts on first / last device (manual-overridable)
+  const embFwd = effectiveNonLayer(data, "embedding", "forward", gpu, c);
+  const embBwd = effectiveNonLayer(data, "embedding", "backward", gpu, c);
+  Df[0] += embFwd;
+  Db[0] += embBwd;
+  const last = c.pp - 1;
+  const outFwd = effectiveNonLayer(data, "output", "forward", gpu, c) + effectiveNonLayer(data, "loss", "forward", gpu, c);
+  const outBwd = effectiveNonLayer(data, "output", "backward", gpu, c) + effectiveNonLayer(data, "loss", "backward", gpu, c);
+  Df[last] += outFwd;
+  Db[last] += outBwd;
+  const mtp = effectiveMtp(data, gpu, c, lt);
+  Df[last] += mtp.fwd;
+  Db[last] += mtp.bwd;
+
+  const critF = Math.max(...Df), critB = Math.max(...Db);
+
+  // Assemble per-device schedule detail (design/07). Non-layer parts are attached
+  // to their owning device so the timeline can render them; Df/Db already include
+  // them above.
+  const critFdev = Df.indexOf(critF), critBdev = Db.indexOf(critB);
+  const devices = [];
+  for (let d = 0; d < c.pp; d++) {
+    devices.push({
+      device: d,
+      chunks: chunks.filter((ch) => ch.device === d).sort((a, b) => a.vpp - b.vpp),
+      Df: Df[d], Db: Db[d],
+      hasEmb: d === 0, hasOut: d === last, hasMtp: d === last && mtp.fwd > 0,
+      embFwd: d === 0 ? embFwd : 0, embBwd: d === 0 ? embBwd : 0,
+      outFwd: d === last ? outFwd : 0, outBwd: d === last ? outBwd : 0,
+      mtpFwd: d === last ? mtp.fwd : 0, mtpBwd: d === last ? mtp.bwd : 0,
+      isCritF: d === critFdev, isCritB: d === critBdev,
+    });
+  }
+  const schedule = { chunks, devices, C, critFdev, critBdev };
+
+  // Step 3: pipeline compute time (us) ; GA = gbs/(dp*mbs)
+  // calibFactor corrects the single-layer-capture -> full-model bias (~0.93,
+  // from the flash 16-layer single-node calibration; see design/06).
+  const ga = c.gbs / (dp * c.mbs);
+  const pipeUs = (ga + (c.pp - 1) / c.vpp) * (critF + critB) * c.calibFactor;
+  const bubbleFrac = (c.pp - 1) / c.vpp / (ga + (c.pp - 1) / c.vpp);
+
+  // Step 4: optimizer (memory-bound, zero1-sharded over DP; CP does not shard params)
+  const optParams = estimateOptimizerParams(data, c, dp);
+  const totalParams = optParams.totalParams;
+  const perRankParams = optParams.perRankParams;
+  const bw = (gpu === "MI355X" ? c.bw355 : c.bw455) * 1e9; // bytes/s
+  const optTimeS = (perRankParams * c.bytesPerParam) / bw / c.optEff;
+  const optUs = optTimeS * 1e6;
+
+  // Step 5: totals
+  const iterUs = pipeUs + optUs;
+  const iterS = iterUs * 1e-6;
+  const seq = data.capture.seq_length;
+  const tokIter = c.gbs * seq;
+  const tokS = tokIter / iterS;
+  const tokSgpu = tokS / world;
+
+  // FLOPs/iter — Megatron-convention V4 analytic model FLOPs (independent of
+  // recompute; recompute adds time, not model flops). Falls back to breakdown
+  // gemm flops if analytic_flops is absent.
+  const counts = cfg.cr_layer_counts;
+  let fMb = 0;
+  const af = data.analytic_flops;
+  if (af && af.per_cr_layer_flops) {
+    for (const cr of ["0", "4", "128"]) fMb += (counts[cr] || 0) * (af.per_cr_layer_flops[cr] || 0);
+    fMb += af.output_flops || 0;
+    if (af.mtp) {
+      fMb += (af.mtp.inner_layer_flops || 0)
+        + (af.mtp.eh_proj_flops || 0)
+        + (af.mtp.extra_logits_flops || 0)
+        + (af.mtp.hc_head_flops || 0);
+    }
+  } else {
+    for (const cr of ["0", "4", "128"]) fMb += (counts[cr] || 0) * (lt[cr].fFlops + lt[cr].bFlops);
+    fMb += nonLayerFlops(data, "output", "forward") + nonLayerFlops(data, "output", "backward");
+  }
+  const flopsIter = fMb * ga * dp;
+  const tflopsGpu = flopsIter / iterS / world / 1e12;
+
+  return {
+    lt, Df, Db, critF, critB, ga, pipeUs, bubbleFrac, world, dp, totalParams,
+    layoutApplied: Boolean(layoutStages),
+    layoutCounts: layoutInfo.counts,
+    layoutMessage: layoutInfo.message,
+    localModelParams: optParams.localModelParams, perRankParams, measuredOptUs: optParams.measuredUs,
+    mtp, optUs, iterUs, tokIter, tokS, tokSgpu, flopsIter, tflopsGpu, seq,
+    schedule,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Rendering
+// ---------------------------------------------------------------------------
+function renderConfig() {
+  const cfg = STATE.data.model_config;
+  const grid = $("#config-grid");
+  grid.innerHTML = "";
+  const items = [
+    ["Model", STATE.data.model.toUpperCase()],
+    ["Layers", cfg.num_layers],
+    ["Hidden", cfg.hidden_size],
+    ["Attn heads", cfg.num_attention_heads],
+    ["Experts", cfg.num_experts],
+    ["Router top-k", cfg.moe_router_topk],
+    ["MoE FFN", cfg.moe_ffn_hidden_size],
+    ["Index top-k", cfg.index_topk],
+    ["Vocab", cfg.vocab_size],
+    ["MTP depths", cfg.mtp_num_layers || 0],
+    ["Capture seq", STATE.data.capture.seq_length],
+    ["cr=0 layers", cfg.cr_layer_counts["0"]],
+    ["cr=4 layers", cfg.cr_layer_counts["4"]],
+    ["cr=128 layers", cfg.cr_layer_counts["128"]],
+  ];
+  for (const [k, v] of items) {
+    const kv = el("div", { class: "kv" });
+    kv.append(el("b", {}, k), el("span", {}, String(v)));
+    grid.append(kv);
+  }
+  // cr schedule strip
+  const strip = $("#cr-schedule");
+  strip.innerHTML = "";
+  const row = el("div", { class: "cr-schedule" });
+  for (const cr of cfg.compress_ratios) row.append(el("span", { class: `cr-cell cr-${cr}`, title: `cr=${cr}` }, "x"));
+  strip.append(row);
+  strip.append(el("div", { class: "cr-legend", html:
+    '<span><i class="cr-0"></i>cr=0 dense+SWA</span><span><i class="cr-4"></i>cr=4 CSA</span><span><i class="cr-128"></i>cr=128 HCA</span>' }));
+}
+
+function breakdownBlock(title, fwdList, bwdList) {
+  const c = STATE.controls, gpu = STATE.gpu;
+  const block = el("div", { class: "bd-block" });
+  block.append(el("h3", {}, title));
+  const scroll = el("div", { class: "bd-scroll" });
+  const table = el("table", { class: "bd" });
+
+  const fwd = fwdList.map((r) => ({ r, t: rowScaledTime(r, gpu, c), phase: "F" }));
+  const bwd = bwdList.map((r) => ({ r, t: rowScaledTime(r, gpu, c), phase: "B" })).reverse();
+  const cols = [...fwd, ...bwd];
+  const dividerIdx = fwd.length;
+
+  const head = el("tr");
+  head.append(el("th", { class: "rowlab" }, "Module"));
+  cols.forEach((col, i) => {
+    const th = el("th", { class: i === dividerIdx ? "divider" : "" });
+    th.append(el("div", {}, col.r.module.replace(/^(attn|moe)\./, "")));
+    th.append(el("div", { class: "phase-tag" }, col.phase));
+    head.append(th);
+  });
+  table.append(head);
+
+  const timeRow = el("tr");
+  timeRow.append(el("td", { class: "rowlab" }, "Time µs"));
+  cols.forEach((col, i) => {
+    const cls = (col.r.class === "compute_bound" ? "cell-compute" : "cell-memory") + (i === dividerIdx ? " divider" : "");
+    timeRow.append(el("td", { class: cls }, fmt(col.t, 0)));
+  });
+  table.append(timeRow);
+
+  const tfRow = el("tr");
+  tfRow.append(el("td", { class: "rowlab" }, "TFLOP/s (kernel)"));
+  cols.forEach((col, i) => {
+    const tf = rowTflops(col.r, col.t);
+    tfRow.append(el("td", { class: i === dividerIdx ? "divider" : "" }, tf ? fmt(tf, 0) : "—"));
+  });
+  table.append(tfRow);
+
+  scroll.append(table);
+  block.append(scroll);
+  return block;
+}
+
+function renderValidation(validation) {
+  const host = $("#validation-panel");
+  if (!host) return;
+  host.innerHTML = "";
+  if (!validation.errors.length && !validation.warnings.length) {
+    host.hidden = true;
+    return;
+  }
+  host.hidden = false;
+  if (validation.errors.length) {
+    const box = el("div", { class: "validation validation--error" });
+    box.append(el("b", {}, "Errors"));
+    const ul = el("ul");
+    for (const msg of validation.errors) ul.append(el("li", {}, msg));
+    box.append(ul);
+    host.append(box);
+  }
+  if (validation.warnings.length) {
+    const box = el("div", { class: "validation validation--warn" });
+    box.append(el("b", {}, "Warnings"));
+    const ul = el("ul");
+    for (const msg of validation.warnings) ul.append(el("li", {}, msg));
+    box.append(ul);
+    host.append(box);
+  }
+}
+
+function renderBreakdown() {
+  $("#breakdown-gpu").textContent = `· ${STATE.gpu}`;
+  const panel = $("#breakdown-panel");
+  if (panel) panel.classList.toggle("is-muted", STATE.controls?.modelMode === "manual");
+  const note = $("#breakdown-manual-note");
+  if (note) note.hidden = STATE.controls?.modelMode !== "manual";
+  const host = $("#breakdown-blocks");
+  host.innerHTML = "";
+  const d = STATE.data;
+  host.append(breakdownBlock("Embedding", d.non_layer.embedding.forward, d.non_layer.embedding.backward));
+  for (const cr of ["0", "4", "128"]) {
+    const L = d.layers[cr];
+    host.append(breakdownBlock(
+      `Layer cr=${cr} — attention`, L.attention.forward, L.attention.backward));
+    host.append(breakdownBlock(
+      `Layer cr=${cr} — MoE`, L.moe.forward, L.moe.backward));
+  }
+  host.append(breakdownBlock("Output (norm + lm_head)", d.non_layer.output.forward, d.non_layer.output.backward));
+  host.append(breakdownBlock("Loss", d.non_layer.loss.forward, d.non_layer.loss.backward));
+}
+
+function step(label, value, sub) {
+  const s = el("div", { class: "step" });
+  s.append(el("span", { class: "num" }, value));
+  s.append(el("b", {}, label));
+  if (sub) { s.append(document.createElement("br")); s.append(el("small", {}, sub)); }
+  return s;
+}
+
+function renderResults(validation) {
+  const c = STATE.controls, gpu = STATE.gpu, d = STATE.data;
+  $("#results-gpu").textContent = `· ${gpu}`;
+  const p = project(d, gpu, c, validation);
+
+  const head = $("#results-headline");
+  head.innerHTML = "";
+  const steps = $("#results-steps");
+  steps.innerHTML = "";
+  if (!p) {
+    const errs = (validation && validation.errors) || [];
+    const summary = errs.length === 0 ? "Invalid controls" : `${errs.length} error${errs.length > 1 ? "s" : ""}`;
+    head.append(el("div", { class: "metric metric--error" }, el("b", {}, "Projection blocked"), el("span", {}, summary)));
+    if (errs.length) {
+      for (const msg of errs) {
+        const s = el("div", { class: "step step--error" });
+        s.append(el("b", {}, "⚠ Validation error"));
+        s.append(document.createElement("br"));
+        s.append(el("small", {}, msg));
+        steps.append(s);
+      }
+    } else {
+      steps.append(step("Fix validation errors", "", "Projection is not recomputed while errors are present."));
+    }
+    return;
+  }
+  const mk = (label, val, primary) => {
+    const m = el("div", { class: "metric" + (primary ? " metric--primary" : "") });
+    m.append(el("b", {}, label), el("span", {}, val));
+    return m;
+  };
+  head.append(mk("tokens/s/GPU", fmtInt(p.tokSgpu), true));
+  head.append(mk("TFLOP/s/GPU", fmt(p.tflopsGpu, 0)));
+  head.append(mk("Iteration time", `${fmt(p.iterUs / 1000, 1)} ms`));
+  head.append(mk("Pipeline bubble", `${fmt(p.bubbleFrac * 100, 1)} %`));
+
+  const ltDesc = ["0", "4", "128"].map((cr) =>
+    `cr${cr}: F ${fmt(p.lt[cr].fwd, 0)} / B ${fmt(p.lt[cr].bwd, 0)} µs`).join("  ·  ");
+  const recomputeDesc = c.recompute === "first-n" ? `first ${c.recomputeLayers} layers/stage` : (c.recompute === "full" ? "recompute on" : "no recompute");
+  const sourceTag = c.modelMode === "manual" ? "manual" : "trace";
+  steps.append(step(`Per-layer fwd/bwd (µs, ${sourceTag}, ${recomputeDesc})`, "", ltDesc));
+  steps.append(step("Critical PP stage (µs)", `F ${fmt(p.critF, 0)} + B ${fmt(p.critB, 0)}`,
+    `max over ${c.pp} stages; layout ${p.layoutApplied ? "applied" : "balanced fallback"} (${p.layoutMessage}); counts=[${p.layoutCounts.join(", ")}]; per-device fwd=[${p.Df.map((x) => fmt(x / 1000, 1)).join(", ")}] ms`));
+  if ((d.model_config.mtp_num_layers || 0) > 0) {
+    steps.append(step("MTP on last stage", `F ${fmt(p.mtp.fwd, 0)} + B ${fmt(p.mtp.bwd, 0)} µs`,
+      `${d.model_config.mtp_num_layers} depth(s), inner cr=${(d.model_config.mtp_compress_ratios || [0])[0]}, eh_proj estimated from output GEMM throughput`));
+  }
+  steps.append(step("GA (microbatches)", fmt(p.ga, 0), `GA = GBS ${c.gbs} / (DP ${p.dp} × MBS ${c.mbs})`));
+  steps.append(step("Pipeline compute / iter", `${fmt(p.pipeUs / 1000, 2)} ms`,
+    `(GA + (PP−1)/VPP) × (F+B)_crit ; bubble ${fmt(p.bubbleFrac * 100, 1)}%`));
+  const optHint = p.measuredOptUs
+    ? `; one-layer trace optimizer ref ${fmt(p.measuredOptUs / 1000, 2)} ms`
+    : "";
+  steps.append(step("Optimizer step / iter", `${fmt(p.optUs / 1000, 2)} ms`,
+    `zero1: ${fmtInt(p.perRankParams / 1e6)}M optim params/rank (${fmtInt(p.localModelParams / 1e6)}M local model params) × ${c.bytesPerParam}B / HBM-BW / eff ${c.optEff}${optHint}`));
+  steps.append(step("Iteration time", `${fmt(p.iterUs / 1000, 2)} ms`, "pipeline compute + optimizer (DP/PP comm assumed hidden)"));
+  steps.append(step("World size", fmtInt(p.world), `PP ${c.pp} × TP ${c.tp} × CP ${c.cp} × DP ${p.dp} (derived); EP ${c.ep} ≤ DP`));
+  steps.append(step("Tokens / iter", fmtInt(p.tokIter), `GBS ${c.gbs} × seq ${p.seq}`));
+  steps.append(step("tokens/s/GPU", fmtInt(p.tokSgpu), `${fmtInt(p.tokS)} tok/s ÷ ${p.world} GPUs`));
+  steps.append(step("TFLOP/s/GPU", fmt(p.tflopsGpu, 0), "V4 analytic model FLOPs (Megatron convention)"));
+
+  // self-consistency hint
+  const cap = d.capture;
+  if (cap && cap.measured_iter_time_ms) {
+    steps.append(step("Self-consistency (measured)", `${fmt(cap.measured_iter_time_ms, 1)} ms`,
+      "set PP=1,VPP=1,DP=1,EP=8,MBS=1,GBS=2 to compare against capture"));
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Iteration timeline (design/07): 3-level composition view
+// ---------------------------------------------------------------------------
+const TL_CATS = ["attn", "mlp", "a2a", "misc"];
+const TL_CAT_LABEL = { attn: "attn", mlp: "mlp", a2a: "a2a (dispatch+combine)", misc: "misc/unattrib" };
+const CR_HEX = { "0": "#6b4a2d", "4": "#2d6b4a", "128": "#2d4a6b" };
+
+// Map a module name to one of the Level-1 categories (design/07).
+function moduleCategory(module) {
+  const m = String(module || "");
+  if (m.startsWith("attn.")) return "attn";
+  if (m === "moe.dispatch" || m === "moe.combine") return "a2a";
+  if (/(^|\.)misc$|unattrib/.test(m)) return "misc";
+  return "mlp"; // moe.grouped_gemm / shared_expert / router and any other moe.*
+}
+
+// Per-cr forward/backward time split into categories, scaled to the active GPU.
+// In manual mode the trace composition is rescaled so the bar total matches the
+// effective (manual-overridden) per-layer time used by the projection.
+function categoryBreakdown(data, cr, gpu, c) {
+  const out = { forward: { attn: 0, mlp: 0, a2a: 0, misc: 0 }, backward: { attn: 0, mlp: 0, a2a: 0, misc: 0 } };
+  const L = data.layers[cr];
+  if (!L) return out;
+  for (const bucket of [L.attention, L.moe]) {
+    for (const phase of ["forward", "backward"]) {
+      for (const r of bucket[phase]) out[phase][moduleCategory(r.module)] += rowScaledTime(r, gpu, c);
+    }
+  }
+  if (c.modelMode === "manual") {
+    const eff = effectiveLayerTimes(data, cr, gpu, c);
+    for (const phase of ["forward", "backward"]) {
+      const sum = TL_CATS.reduce((a, k) => a + out[phase][k], 0);
+      const target = phase === "forward" ? eff.fwd : eff.bwd;
+      if (sum > 0 && target > 0) for (const k of TL_CATS) out[phase][k] *= target / sum;
+    }
+  }
+  return out;
+}
+
+// True when, in manual mode, the user has actually changed this cr's fwd/bwd away
+// from the trace-derived baseline (comparison is rounded to µs so the prefilled
+// baseline itself does not count as an edit).
+function layerManualEdited(data, cr, gpu, c) {
+  if (c.modelMode !== "manual") return false;
+  const trace = layerTimes(data, cr, gpu, c);
+  const eff = effectiveLayerTimes(data, cr, gpu, c);
+  return Math.round(eff.fwd) !== Math.round(trace.fwd) || Math.round(eff.bwd) !== Math.round(trace.bwd);
+}
+
+// Small hex lighten/darken for VPP chunk shading (Level 3).
+function shadeHex(hex, amt) {
+  const n = parseInt(hex.slice(1), 16);
+  const clamp = (x) => Math.max(0, Math.min(255, Math.round(x)));
+  const r = clamp(((n >> 16) & 255) + amt), g = clamp(((n >> 8) & 255) + amt), b = clamp((n & 255) + amt);
+  return `#${((r << 16) | (g << 8) | b).toString(16).padStart(6, "0")}`;
+}
+
+// Group PP devices with an identical ordered layer/recompute/non-layer signature
+// so identical middle stages are drawn once (design/07 Level 2 dedup).
+function dedupDevices(devices) {
+  const groups = [];
+  const byKey = new Map();
+  for (const dev of devices) {
+    const sig = JSON.stringify({
+      layers: dev.chunks.map((ch) => ch.layers.map((l) => `${l.cr}${l.recompute ? "r" : ""}`)),
+      emb: dev.hasEmb, out: dev.hasOut, mtp: dev.hasMtp,
+    });
+    if (byKey.has(sig)) byKey.get(sig).members.push(dev.device);
+    else {
+      const g = { rep: dev, members: [dev.device] };
+      byKey.set(sig, g);
+      groups.push(g);
+    }
+  }
+  return groups;
+}
+
+// 1F1B (optionally interleaved) schedule simulator (design/07 Level 3). Returns
+// per-device event lists with start/dur (µs) plus the drawn iteration length.
+const TL_VIS_GA_CAP = 48;
+function simulateSchedule(p, c, { interleaved }) {
+  const PP = c.pp;
+  const VPP = interleaved ? Math.max(1, c.vpp) : 1;
+  const C = PP * VPP;
+  const gaFull = Math.round(p.ga);
+  const GA = Math.min(gaFull, TL_VIS_GA_CAP);
+  const capped = gaFull > GA;
+
+  // per-virtual-chunk fwd/bwd durations (k = 0..C-1, device = k % PP)
+  const fdur = new Array(C).fill(0), bdur = new Array(C).fill(0);
+  if (VPP === 1) {
+    for (let d = 0; d < PP; d++) { fdur[d] = p.Df[d]; bdur[d] = p.Db[d]; }
+  } else {
+    for (const ch of p.schedule.chunks) { fdur[ch.chunk] = ch.fwd; bdur[ch.chunk] = ch.bwd; }
+    // fold non-layer parts into first/last virtual chunk so the drawn length is
+    // consistent with the per-device Df/Db (which include them).
+    const dev0 = p.schedule.devices[0], devL = p.schedule.devices[PP - 1];
+    fdur[0] += dev0.embFwd; bdur[0] += dev0.embBwd;
+    fdur[C - 1] += devL.outFwd + devL.mtpFwd; bdur[C - 1] += devL.outBwd + devL.mtpBwd;
+  }
+
+  const total = GA * VPP;
+  const group = PP * VPP;
+  const fwdChunkMb = (i) => {
+    const inGroup = i % group;
+    const v = Math.floor(inGroup / PP);
+    const m = Math.floor(i / group) * PP + (inGroup % PP);
+    return { v, m };
+  };
+
+  const ops = []; // per device ordered op list
+  for (let d = 0; d < PP; d++) {
+    const warmup = Math.min(
+      VPP === 1 ? PP - 1 - d : (PP - d - 1) * 2 + (VPP - 1) * PP,
+      total,
+    );
+    const list = [];
+    for (let i = 0; i < warmup; i++) {
+      const { v, m } = fwdChunkMb(i);
+      list.push({ kind: "F", m, k: v * PP + d });
+    }
+    let fptr = warmup, bptr = 0;
+    const n1f1b = total - warmup;
+    for (let i = 0; i < n1f1b; i++) {
+      const f = fwdChunkMb(fptr++);
+      list.push({ kind: "F", m: f.m, k: f.v * PP + d });
+      const b = fwdChunkMb(bptr++);
+      list.push({ kind: "B", m: b.m, k: (VPP - 1 - b.v) * PP + d });
+    }
+    for (let i = 0; i < warmup; i++) {
+      const b = fwdChunkMb(bptr++);
+      list.push({ kind: "B", m: b.m, k: (VPP - 1 - b.v) * PP + d });
+    }
+    ops.push(list);
+  }
+
+  // ASAP scheduling: forward flows k-1 -> k, backward flows k+1 -> k (p2p hidden).
+  const endF = new Map(), endB = new Map();
+  const kf = (m, k) => `${m}:${k}`;
+  const ptr = new Array(PP).fill(0);
+  const free = new Array(PP).fill(0);
+  const events = Array.from({ length: PP }, () => []);
+  let remaining = ops.reduce((a, l) => a + l.length, 0);
+  let guard = remaining * 4 + 16;
+  while (remaining > 0 && guard-- > 0) {
+    let progressed = false;
+    for (let d = 0; d < PP; d++) {
+      if (ptr[d] >= ops[d].length) continue;
+      const op = ops[d][ptr[d]];
+      let dep = 0, ready = true;
+      if (op.kind === "F") {
+        if (op.k > 0) {
+          const e = endF.get(kf(op.m, op.k - 1));
+          if (e == null) ready = false; else dep = e;
+        }
+      } else {
+        if (op.k < C - 1) {
+          const e = endB.get(kf(op.m, op.k + 1));
+          if (e == null) ready = false; else dep = e;
+        } else {
+          const e = endF.get(kf(op.m, op.k));
+          if (e == null) ready = false; else dep = e;
+        }
+      }
+      if (!ready) continue;
+      const dur = op.kind === "F" ? fdur[op.k] : bdur[op.k];
+      const start = Math.max(free[d], dep);
+      const end = start + dur;
+      events[d].push({ kind: op.kind, m: op.m, k: op.k, vpp: Math.floor(op.k / PP), start, dur });
+      free[d] = end;
+      (op.kind === "F" ? endF : endB).set(kf(op.m, op.k), end);
+      ptr[d]++; remaining--; progressed = true;
+    }
+    if (!progressed) break; // dependency stall guard
+  }
+  const drawnUs = Math.max(0, ...free);
+  return { events, drawnUs, GA, capped, PP, VPP, C };
+}
+
+// --- cross-level selection (drill-down linkage) ---
+function tlSelActive() {
+  return STATE.tlSel.cr != null || (STATE.tlSel.devices && STATE.tlSel.devices.length);
+}
+function tlSelectCr(cr) {
+  STATE.tlSel = STATE.tlSel.cr === cr ? { cr: null, devices: null } : { cr, devices: null };
+  renderTimeline();
+}
+function tlSelectDevices(devices) {
+  const same = STATE.tlSel.devices && STATE.tlSel.devices.length === devices.length && STATE.tlSel.devices.every((d, i) => d === devices[i]);
+  STATE.tlSel = same ? { cr: null, devices: null } : { cr: null, devices };
+  renderTimeline();
+}
+function tlClearSel() {
+  STATE.tlSel = { cr: null, devices: null };
+  renderTimeline();
+}
+// Compression ratios "active" under the current selection (drives L1 highlight).
+function tlActiveCrSet(p) {
+  if (STATE.tlSel.cr != null) return new Set([String(STATE.tlSel.cr)]);
+  if (STATE.tlSel.devices && STATE.tlSel.devices.length) {
+    const s = new Set();
+    for (const d of STATE.tlSel.devices) {
+      for (const ch of p.schedule.devices[d].chunks) for (const l of ch.layers) s.add(String(l.cr));
+    }
+    return s;
+  }
+  return null;
+}
+const tlDevSet = () => (STATE.tlSel.devices && STATE.tlSel.devices.length ? new Set(STATE.tlSel.devices) : null);
+
+function renderTimeline() {
+  const host = $("#tl-body");
+  if (!host) return;
+  $("#timeline-gpu").textContent = `· ${STATE.gpu}`;
+  const stacked = STATE.tlView === "stacked";
+  document.querySelectorAll(".tl-view").forEach((t) => t.classList.toggle("is-active", (t.dataset.view === "stacked") === stacked));
+  document.querySelectorAll(".tl-tab").forEach((t) => t.classList.toggle("is-active", Number(t.dataset.level) === STATE.tlLevel));
+  const tabs = $("#tl-tabs");
+  if (tabs) tabs.style.display = stacked ? "none" : "";
+  host.innerHTML = "";
+  const validation = validateControls(STATE.data, STATE.controls, STATE.gpu);
+  const p = project(STATE.data, STATE.gpu, STATE.controls, validation);
+  if (!p) {
+    host.append(el("p", { class: "tl-warn" }, "Timeline unavailable: fix the validation errors above."));
+    return;
+  }
+
+  // selection / linkage indicator: a fixed-position floating card appended to
+  // <body> (NOT into #tl-body) so it is fully outside the timeline's flow and
+  // toggling a selection never shifts the page layout at all.
+  document.getElementById("tl-selbar-float")?.remove();
+  if (tlSelActive()) {
+    const bar = el("div", { id: "tl-selbar-float", class: "tl-selbar" });
+    const txt = el("div", { class: "tl-selbar__txt" });
+    const desc = STATE.tlSel.cr != null
+      ? `Focused on cr=${STATE.tlSel.cr}`
+      : `Focused on PP rank(s) ${STATE.tlSel.devices.join(", ")}`;
+    txt.append(el("b", {}, desc));
+    txt.append(el("span", { class: "muted" }, stacked ? " — highlighted across all levels" : " — highlighted; switch to Stacked to see all levels"));
+    bar.append(txt);
+    const btn = el("button", { class: "tl-clear" }, "Clear");
+    btn.addEventListener("click", tlClearSel);
+    bar.append(btn);
+    document.body.append(bar);
+  }
+
+  if (stacked) {
+    const section = (title, sub, fn) => {
+      const sec = el("div", { class: "tl-section" });
+      const h = el("h3", { class: "tl-section__title" }, title);
+      if (sub) h.append(el("span", { class: "tl-section__sub" }, sub));
+      sec.append(h);
+      const body = el("div", {});
+      fn(body);
+      sec.append(body);
+      host.append(sec);
+    };
+    section("Level 1 · single layer", "attn / mlp / a2a — click a cr to link", (b) => renderTimelineL1(b, p));
+    section("Level 2 · pipeline ranks", "layer granularity — click a rank or a layer to link", (b) => renderTimelineL2(b, p));
+    section("Level 3 · pipeline schedule", "1F1B / interleaved — click a device to link", (b) => renderTimelineL3(b, p));
+  } else if (STATE.tlLevel === 1) renderTimelineL1(host, p);
+  else if (STATE.tlLevel === 2) renderTimelineL2(host, p);
+  else renderTimelineL3(host, p);
+}
+
+// --- gantt export (SVG / PNG) ---
+function downloadBlob(blob, filename) {
+  const url = URL.createObjectURL(blob);
+  const a = el("a", { href: url, download: filename });
+  document.body.append(a);
+  a.click();
+  a.remove();
+  setTimeout(() => URL.revokeObjectURL(url), 2000);
+}
+function serializeGantt(svg) {
+  const clone = svg.cloneNode(true);
+  const cs = getComputedStyle(document.documentElement);
+  const bg = (cs.getPropertyValue("--bg").trim() || "#0f1218");
+  let markup = new XMLSerializer().serializeToString(clone);
+  for (const v of ["--panel-2", "--border", "--muted", "--text", "--bg"]) {
+    markup = markup.split(`var(${v})`).join(cs.getPropertyValue(v).trim() || "#888");
+  }
+  if (!markup.includes("xmlns=")) markup = markup.replace("<svg", '<svg xmlns="http://www.w3.org/2000/svg"');
+  const vb = svg.viewBox.baseVal;
+  markup = markup.replace(/(<svg[^>]*>)/, `$1<rect x="0" y="0" width="${vb.width}" height="${vb.height}" fill="${bg}"/>`);
+  return markup;
+}
+function exportGanttSvg(svg) {
+  downloadBlob(new Blob([serializeGantt(svg)], { type: "image/svg+xml" }), `dsv4-pp-schedule-${STATE.gpu}.svg`);
+}
+function exportGanttPng(svg) {
+  const markup = serializeGantt(svg);
+  const vb = svg.viewBox.baseVal, scale = 2;
+  const img = new Image();
+  img.onload = () => {
+    const canvas = el("canvas");
+    canvas.width = vb.width * scale;
+    canvas.height = vb.height * scale;
+    const ctx = canvas.getContext("2d");
+    ctx.scale(scale, scale);
+    ctx.drawImage(img, 0, 0);
+    canvas.toBlob((b) => downloadBlob(b, `dsv4-pp-schedule-${STATE.gpu}.png`));
+  };
+  img.src = "data:image/svg+xml;base64," + btoa(unescape(encodeURIComponent(markup)));
+}
+
+function catLegend() {
+  return el("div", { class: "tl-legend", html:
+    '<span><i class="tl-c-attn"></i>attn</span>' +
+    '<span><i class="tl-c-mlp"></i>mlp (experts)</span>' +
+    '<span><i class="tl-c-a2a"></i>a2a (dispatch+combine)</span>' +
+    '<span><i class="tl-c-misc"></i>misc/unattributed</span>' });
+}
+
+function stackedTrack(phaseObj, maxTotal) {
+  const track = el("div", { class: "tl-bar__track" });
+  const total = TL_CATS.reduce((a, k) => a + phaseObj[k], 0);
+  for (const cat of TL_CATS) {
+    const t = phaseObj[cat];
+    if (t <= 0) continue;
+    const w = maxTotal > 0 ? (t / maxTotal) * 100 : 0;
+    const seg = el("div", {
+      class: `tl-seg tl-c-${cat}`,
+      style: `width:${w}%`,
+      title: `${TL_CAT_LABEL[cat]}: ${fmt(t, 0)} µs (${fmt(total > 0 ? (t / total) * 100 : 0, 0)}%)`,
+    }, w > 6 ? cat : "");
+    track.append(seg);
+  }
+  return { track, total };
+}
+
+// Single-segment bar used when a cr's layer time is manually set: no module split
+// is possible, so show only the total fwd/bwd.
+function aggregateTrack(timeUs, maxTotal, label) {
+  const track = el("div", { class: "tl-bar__track" });
+  const w = maxTotal > 0 ? (timeUs / maxTotal) * 100 : 0;
+  track.append(el("div", {
+    class: "tl-seg tl-c-manual", style: `width:${w}%`,
+    title: `${label}: ${fmt(timeUs, 0)} µs (manual per-layer)`,
+  }, w > 6 ? "manual" : ""));
+  return { track, total: timeUs };
+}
+
+function renderTimelineL1(host, p) {
+  const d = STATE.data, gpu = STATE.gpu, c = STATE.controls;
+  host.append(el("p", { class: "tl-note" },
+    "One representative layer per compression ratio (MoE is cr-independent, so only attention differs). Forward and backward split into attn / mlp / a2a; bar length is comparable across cr."));
+  const crs = ["0", "4", "128"].filter((cr) => (d.model_config.cr_layer_counts?.[cr] || 0) > 0);
+  const crSet = tlActiveCrSet(p);
+  // Per-cr: whether the layer time is manually overridden (then no module split),
+  // the module breakdown, and the fwd/bwd totals used for the shared scale.
+  const info = {};
+  let maxTotal = 0;
+  let anyManual = false;
+  for (const cr of crs) {
+    const edited = layerManualEdited(d, cr, gpu, c);
+    const bd = edited ? null : categoryBreakdown(d, cr, gpu, c);
+    const eff = effectiveLayerTimes(d, cr, gpu, c);
+    const fT = edited ? eff.fwd : TL_CATS.reduce((a, k) => a + bd.forward[k], 0);
+    const bT = edited ? eff.bwd : TL_CATS.reduce((a, k) => a + bd.backward[k], 0);
+    info[cr] = { edited, bd, fT, bT };
+    anyManual = anyManual || edited;
+    maxTotal = Math.max(maxTotal, fT, bT);
+  }
+  for (const cr of crs) {
+    const { edited, bd, fT, bT } = info[cr];
+    const isSel = STATE.tlSel.cr === cr;
+    const dim = crSet && !crSet.has(cr);
+    const row = el("div", { class: "tl-l1-row tl-clickable" + (isSel ? " is-sel" : "") + (dim ? " tl-dim" : "") });
+    row.addEventListener("click", () => tlSelectCr(cr));
+    const lab = el("div", { class: "tl-l1-lab" });
+    lab.append(el("b", {}, `cr=${cr}`), el("small", {}, `×${d.model_config.cr_layer_counts[cr]} layers${edited ? " · manual" : ""}`));
+    row.append(lab);
+    const bars = el("div", { class: "tl-bars" });
+    for (const [tag, phase, tot] of [["fwd", "forward", fT], ["bwd", "backward", bT]]) {
+      const bar = el("div", { class: "tl-bar" });
+      bar.append(el("span", { class: "tl-bar__tag" }, tag));
+      const { track, total } = edited ? aggregateTrack(tot, maxTotal, `${tag} layer`) : stackedTrack(bd[phase], maxTotal);
+      bar.append(track);
+      bar.append(el("span", { class: "tl-bar__total" }, `${fmt(total, 0)} µs`));
+      bars.append(bar);
+    }
+    row.append(bars);
+    host.append(row);
+  }
+  if (anyManual) {
+    host.append(el("div", { class: "tl-legend", html:
+      '<span><i class="tl-c-manual"></i>manual per-layer total (no module split)</span>' }));
+    host.append(el("p", { class: "tl-note" },
+      "A cr marked “manual” uses a hand-entered whole-layer fwd/bwd time, so it cannot be split into attn / mlp / a2a — only the total is shown. Switch layer timing back to “Trace-derived”, or click “Restore defaults”, to see the per-module breakdown again."));
+  }
+  host.append(catLegend());
+}
+
+function renderTimelineL2(host, p) {
+  const c = STATE.controls;
+  host.append(el("p", { class: "tl-note" },
+    "Each pipeline rank's layers (coloured by cr; hatched = recomputed). Identical ranks are drawn once. The critical stage (max fwd/bwd, sets the pipeline critical path) is outlined."));
+  const groups = dedupDevices(p.schedule.devices);
+  const maxTotal = Math.max(1, ...p.schedule.devices.map((dv) => dv.Df + dv.Db));
+  const selDevices = tlDevSet();
+  const selCr = STATE.tlSel.cr != null ? String(STATE.tlSel.cr) : null;
+  for (const g of groups) {
+    const dev = g.rep;
+    const isCrit = dev.isCritF || dev.isCritB;
+    const isSelGroup = selDevices && g.members.some((m) => selDevices.has(m));
+    const dimRow = (selDevices && !isSelGroup) || (selCr && !g.members.some((m) => p.schedule.devices[m].chunks.some((ch) => ch.layers.some((l) => String(l.cr) === selCr))));
+    const row = el("div", { class: "tl-l2-row tl-clickable" + (isCrit ? " is-critical" : "") + (isSelGroup ? " is-sel" : "") + (dimRow ? " tl-dim" : "") });
+    row.addEventListener("click", () => tlSelectDevices(g.members));
+    const lab = el("div", { class: "tl-l2-lab" });
+    const members = g.members;
+    const rangeTxt = members.length > 1 ? `PP ranks ${members[0]}–${members[members.length - 1]} (×${members.length})` : `PP rank ${members[0]}`;
+    lab.append(el("b", {}, rangeTxt));
+    lab.append(el("small", {}, `${dev.chunks.reduce((a, ch) => a + ch.layers.length, 0)} layers · ${dev.chunks.length} vpp chunk(s)${isCrit ? " · critical" : ""}`));
+    row.append(lab);
+
+    const total = dev.Df + dev.Db;
+    const strip = el("div", { class: "tl-strip", style: `width:${(total / maxTotal) * 100}%` });
+    if (dev.hasEmb) {
+      const t = dev.embFwd + dev.embBwd;
+      strip.append(el("div", { class: "tl-cell tl-cell--emb", style: `width:${(t / total) * 100}%`, title: `embedding F ${fmt(dev.embFwd, 0)} / B ${fmt(dev.embBwd, 0)} µs` }));
+    }
+    dev.chunks.forEach((ch, ci) => {
+      if (ci > 0 || dev.hasEmb) strip.append(el("div", { class: "tl-chunk-gap" }));
+      for (const l of ch.layers) {
+        const t = l.fwd + l.bwd;
+        const cellSel = selCr && String(l.cr) === selCr;
+        const cellDim = selCr && String(l.cr) !== selCr;
+        const cell = el("div", {
+          class: "tl-cell tl-clickable" + (l.recompute ? " tl-cell--recompute" : "") + (cellSel ? " tl-cell--sel" : "") + (cellDim ? " tl-dim" : ""),
+          style: `width:${(t / total) * 100}%; background:${CR_HEX[String(l.cr)] || "#555"}`,
+          title: `layer #${l.globalIdx} cr=${l.cr}${l.recompute ? " (recompute)" : ""} · F ${fmt(l.fwd, 0)} / B ${fmt(l.bwd, 0)} µs`,
+        });
+        cell.addEventListener("click", (e) => { e.stopPropagation(); tlSelectCr(String(l.cr)); });
+        strip.append(cell);
+      }
+    });
+    if (dev.hasOut) {
+      const t = dev.outFwd + dev.outBwd + dev.mtpFwd + dev.mtpBwd;
+      strip.append(el("div", { class: "tl-chunk-gap" }));
+      strip.append(el("div", { class: "tl-cell tl-cell--out", style: `width:${(t / total) * 100}%`, title: `output/loss${dev.hasMtp ? "+MTP" : ""} F ${fmt(dev.outFwd + dev.mtpFwd, 0)} / B ${fmt(dev.outBwd + dev.mtpBwd, 0)} µs` }));
+    }
+    row.append(strip);
+    row.append(el("div", { class: "tl-l2-meta" }, `Df ${fmt(dev.Df / 1000, 2)} / Db ${fmt(dev.Db / 1000, 2)} ms`));
+    host.append(row);
+  }
+  host.append(el("div", { class: "cr-legend", html:
+    '<span><i class="cr-0"></i>cr=0</span><span><i class="cr-4"></i>cr=4</span><span><i class="cr-128"></i>cr=128</span>' +
+    '<span><i class="tl-c-emb"></i>embedding</span><span><i class="tl-c-out"></i>output/loss/MTP</span>' }));
+}
+
+function renderTimelineL3(host, p) {
+  const c = STATE.controls;
+  // Interleaving is driven directly by the VPP control (no separate toggle):
+  // VPP=1 is plain 1F1B, VPP>1 is interleaved 1F1B.
+  const interleaved = c.vpp > 1;
+
+  const controls = el("div", { class: "tl-controls" });
+  controls.append(el("span", { class: "mode-switch__label" }, `Schedule · 1F1B${interleaved ? ` (interleaved, VPP=${c.vpp})` : ""}`));
+
+  // Zoom slider: 1x fits the whole schedule in view (no scrollbar); zooming in
+  // widens the chart so the per-cell microbatch numbers become readable.
+  const zoomWrap = el("span", { class: "tl-zoom-wrap" });
+  zoomWrap.append(el("span", { class: "mode-switch__label" }, "Zoom"));
+  const zoom = el("input", { type: "range", min: "1", max: "10", step: "0.5", value: String(STATE.tlZoom), class: "tl-zoom" });
+  const zlab = el("span", { class: "tl-zoom-val" }, `${STATE.tlZoom}×`);
+  zoom.addEventListener("input", () => { zlab.textContent = `${zoom.value}×`; });
+  zoom.addEventListener("change", () => { STATE.tlZoom = Number(zoom.value); renderTimeline(); });
+  zoomWrap.append(zoom, zlab);
+  controls.append(zoomWrap);
+  host.append(controls);
+  if (!interleaved) {
+    host.append(el("p", { class: "tl-note" }, "VPP=1 → plain 1F1B. Set VPP>1 in the controls to interleave the schedule and shrink the pipeline bubble."));
+  }
+
+  const sim = simulateSchedule(p, c, { interleaved });
+  const PP = sim.PP;
+  const rowH = 30, gap = 6, padL = 54, padT = 8, padB = 26;
+  // Zoom widens the coordinate system itself (more px per µs; font size stays
+  // fixed so cells become readable) instead of CSS-scaling the SVG. At 1x the
+  // chart is sized to the actual right-content width so it fits with no
+  // scrollbar; >1x overflows and scrolls horizontally, undistorted.
+  const avail = Math.max(600, (($("#tl-body")?.clientWidth) || 1100) - 16);
+  const plotW = Math.round((avail - padL) * STATE.tlZoom);
+  const width = padL + plotW + 8;
+  const scale = sim.drawnUs > 0 ? plotW / sim.drawnUs : 0;
+  const height = padT + PP * (rowH + gap) + padB;
+
+  const svg = document.createElementNS("http://www.w3.org/2000/svg", "svg");
+  svg.setAttribute("class", "tl-gantt");
+  svg.setAttribute("viewBox", `0 0 ${width} ${height}`);
+  svg.setAttribute("width", String(width));
+  svg.setAttribute("height", String(height));
+  const mkEl = (tag, attrs, text) => {
+    const n = document.createElementNS("http://www.w3.org/2000/svg", tag);
+    for (const [k, v] of Object.entries(attrs)) n.setAttribute(k, v);
+    if (text != null) n.textContent = text;
+    return n;
+  };
+
+  const selDevices = tlDevSet();
+  for (let d = 0; d < PP; d++) {
+    const y = padT + d * (rowH + gap);
+    const isSel = selDevices ? selDevices.has(d) : false;
+    const g = mkEl("g", { class: "tl-devrow" });
+    if (selDevices) g.setAttribute("opacity", isSel ? "1" : "0.3");
+    g.append(mkEl("text", { x: padL - 8, y: y + rowH / 2 + 4, "text-anchor": "end", fill: "#e6eaf2", "font-size": "11" }, `dev ${d}`));
+    const bgRect = mkEl("rect", {
+      x: padL, y, width: plotW, height: rowH, rx: "3",
+      fill: isSel ? "var(--panel)" : "var(--panel-2)",
+      stroke: isSel ? "#36c08f" : "var(--border)", "stroke-width": isSel ? "2" : "1",
+      style: "cursor:pointer",
+    });
+    bgRect.addEventListener("click", () => tlSelectDevices([d]));
+    g.append(bgRect);
+    for (const ev of sim.events[d]) {
+      const x = padL + ev.start * scale;
+      const w = Math.max(1, ev.dur * scale);
+      const base = ev.kind === "F" ? "#4f8cff" : "#36c08f";
+      const fill = shadeHex(base, ev.vpp * -26);
+      const rect = mkEl("rect", {
+        x, y: y + 2, width: w, height: rowH - 4, rx: "2",
+        fill, class: ev.kind === "F" ? "tl-fwd" : "tl-bwd",
+      });
+      rect.append(mkEl("title", {}, `${ev.kind === "F" ? "Forward" : "Backward"} · microbatch ${ev.m}${sim.VPP > 1 ? ` · vpp chunk ${ev.vpp}` : ""} · compute ${fmt(ev.dur, 0)} µs · starts @ ${fmt(ev.start / 1000, 2)} ms (from iter start)`));
+      g.append(rect);
+      if (w >= 8) g.append(mkEl("text", { x: x + w / 2, y: y + rowH / 2 + 3, "text-anchor": "middle", fill: "#000000", "font-size": "9" }, String(ev.m)));
+    }
+    svg.append(g);
+  }
+  // time axis ticks
+  const yb = padT + PP * (rowH + gap);
+  for (let i = 0; i <= 4; i++) {
+    const tx = padL + (plotW * i) / 4;
+    svg.append(mkEl("text", { x: tx, y: yb + 16, "text-anchor": "middle", fill: "#8d97a8", "font-size": "10" }, `${fmt((sim.drawnUs * i) / 4 / 1000, 1)} ms`));
+  }
+
+  const toolbar = el("div", { class: "tl-export" });
+  const svgBtn = el("button", { class: "mode-tab" }, "Export SVG");
+  const pngBtn = el("button", { class: "mode-tab" }, "Export PNG");
+  svgBtn.addEventListener("click", () => exportGanttSvg(svg));
+  pngBtn.addEventListener("click", () => exportGanttPng(svg));
+  toolbar.append(svgBtn, pngBtn);
+  host.append(toolbar);
+
+  const wrap = el("div", { class: "tl-gantt-wrap" });
+  wrap.append(svg);
+  host.append(wrap);
+
+  // legend + axis/tooltip explanation + self-check vs analytic pipe time
+  host.append(el("div", { class: "tl-legend", html:
+    '<span><i class="tl-c-attn"></i>forward</span><span><i class="tl-c-mlp"></i>backward</span>' +
+    (sim.VPP > 1 ? '<span>lighter→darker = VPP chunk 0→' + (sim.VPP - 1) + '</span>' : '') +
+    '<span>gaps = pipeline bubble (idle)</span>' }));
+  host.append(el("p", { class: "tl-note tl-axis-note" },
+    "X-axis = wall-clock time measured from the start of the pipeline (ms). Each cell is one microbatch's forward or backward on that device; the number inside is the microbatch index. Hover a cell to read its compute duration in µs (the “… µs” before @) and its start time in ms from the iteration start (the “… ms” after @)."));
+
+  const analyticPipeUs = (p.ga + (c.pp - 1) / sim.VPP) * (p.critF + p.critB); // pre-calibFactor, matches the drawn schedule's VPP
+  const diff = analyticPipeUs > 0 ? Math.abs(sim.drawnUs - analyticPipeUs) / analyticPipeUs : 0;
+  const summary = el("div", { class: "tl-summary" });
+  const gaTxt = sim.capped ? `${sim.GA} of ${Math.round(p.ga)} microbatches (capped for display)` : `${sim.GA} microbatches`;
+  summary.append(el("div", {}, `Drawn iteration (pipeline compute): ${fmt(sim.drawnUs / 1000, 2)} ms over ${gaTxt}; PP=${c.pp}, VPP=${interleaved ? c.vpp : 1}. Analytic bubble fraction ${fmt(p.bubbleFrac * 100, 1)}%.`));
+  if (!sim.capped) {
+    const cls = diff > 0.08 ? "tl-warn" : "";
+    summary.append(el("div", { class: cls },
+      `Self-check vs analytic (GA + (PP−1)/VPP)·(F+B)_crit = ${fmt(analyticPipeUs / 1000, 2)} ms → ${fmt(diff * 100, 1)}% ${diff > 0.08 ? "difference (imbalanced stages; analytic uses per-device max)" : "match"}. Official iteration time uses the analytic value × calibFactor.`));
+  }
+  host.append(summary);
+}
+
+function renderAll() {
+  const validation = validateControls(STATE.data, STATE.controls, STATE.gpu);
+  // keep the derived-DP readonly box in sync
+  const w = STATE.controls;
+  const roDp = document.getElementById("ctl-dp");
+  if (roDp) roDp.value = Number.isFinite(validation.dp) ? validation.dp : "—";
+  renderValidation(validation);
+  renderModeSwitch();
+  renderManualGrid();
+  renderConfig();
+  renderBreakdown();
+  renderResults(validation);
+  renderTimeline();
+}
+
+// ---------------------------------------------------------------------------
+// Bootstrap
+// ---------------------------------------------------------------------------
+async function init(model) {
+  try {
+    STATE.data = await loadModel(model);
+    STATE.controls = defaultControls(STATE.data);
+    $("#mock-badge").hidden = !(STATE.data.provenance && STATE.data.provenance.mock);
+    $("#model-select").value = STATE.data.model;
+    renderControls();
+    renderAll();
+  } catch (e) {
+    const err = $("#error-state");
+    err.hidden = false;
+    err.textContent = String(e);
+  }
+}
+
+globalThis.DSV4Projection = {
+  defaultControls,
+  derivedDP,
+  expandLayoutRepeats,
+  parsePipelineLayout,
+  validateControls,
+  effectiveLayerTimes,
+  project,
+  moduleCategory,
+  categoryBreakdown,
+  dedupDevices,
+  simulateSchedule,
+};
+
+if (typeof document !== "undefined") {
+  $("#model-select").addEventListener("change", (e) => {
+    const m = e.target.value;
+    const u = new URL(location);
+    u.searchParams.set("model", m);
+    history.replaceState(null, "", u);
+    init(m);
+  });
+
+  document.querySelectorAll(".tab").forEach((tab) => {
+    tab.addEventListener("click", () => {
+      document.querySelectorAll(".tab").forEach((t) => t.classList.remove("is-active"));
+      tab.classList.add("is-active");
+      STATE.gpu = tab.dataset.gpu;
+      if (STATE.controls?.modelMode === "manual") prefillManual(STATE.gpu);
+      renderAll();
+    });
+  });
+
+  document.querySelectorAll(".mode-tab").forEach((tab) => {
+    tab.addEventListener("click", () => {
+      const mode = tab.dataset.mode;
+      if (!STATE.controls || STATE.controls.modelMode === mode) return;
+      STATE.controls.modelMode = mode;
+      if (mode === "manual") prefillManual(STATE.gpu);
+      renderAll();
+    });
+  });
+
+  document.querySelectorAll(".tl-tab").forEach((tab) => {
+    tab.addEventListener("click", () => {
+      STATE.tlLevel = Number(tab.dataset.level);
+      renderTimeline();
+    });
+  });
+
+  document.querySelectorAll(".tl-view").forEach((tab) => {
+    tab.addEventListener("click", () => {
+      STATE.tlView = tab.dataset.view;
+      renderTimeline();
+    });
+  });
+
+  // Re-fit the schedule Gantt to the content width when the window is resized.
+  let _tlResizeTimer;
+  window.addEventListener("resize", () => {
+    clearTimeout(_tlResizeTimer);
+    _tlResizeTimer = setTimeout(() => { if (STATE.data) renderTimeline(); }, 150);
+  });
+
+  init(modelFromQuery());
+}
diff --git a/examples/deepseek-v4/projection/site/assets/style.css b/examples/deepseek-v4/projection/site/assets/style.css
new file mode 100644
index 000000000..8d9d9d805
--- /dev/null
+++ b/examples/deepseek-v4/projection/site/assets/style.css
@@ -0,0 +1,333 @@
+:root {
+  --bg: #0f1218;
+  --panel: #171c26;
+  --panel-2: #1e2532;
+  --accent: #4f8cff;
+  --accent-2: #36c08f;
+  --text: #e6eaf2;
+  --muted: #8d97a8;
+  --border: #2a3342;
+  --warn: #e0a03a;
+  --compute: #2d4a6b;
+  --memory: #3a3346;
+}
+
+* { box-sizing: border-box; }
+
+body {
+  margin: 0;
+  font-family: ui-sans-serif, system-ui, -apple-system, "Segoe UI", Roboto, sans-serif;
+  background: var(--bg);
+  color: var(--text);
+  line-height: 1.45;
+}
+
+.hero {
+  background: linear-gradient(160deg, #1b2330, #11151d);
+  border-bottom: 1px solid var(--border);
+  padding: 28px 0;
+}
+.hero__inner, .layout { max-width: 1600px; margin: 0 auto; padding: 0 24px; }
+.eyebrow { color: var(--accent); font-weight: 600; letter-spacing: .08em; text-transform: uppercase; font-size: 12px; margin: 0 0 4px; }
+.hero h1 { margin: 0 0 8px; font-size: 28px; }
+.hero__summary { color: var(--muted); max-width: 760px; margin: 0 0 16px; }
+.hero__row { display: flex; align-items: center; gap: 14px; margin-bottom: 14px; }
+
+.tabs { display: flex; gap: 8px; }
+.tab {
+  background: var(--panel); color: var(--text); border: 1px solid var(--border);
+  padding: 9px 16px; border-radius: 8px 8px 0 0; cursor: pointer; font-size: 14px;
+}
+.tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; }
+
+/* Sticky switcher bar: model select + GPU tabs stay pinned to the top so the
+   user can switch without scrolling back up. */
+.switcher {
+  position: sticky;
+  top: 0;
+  z-index: 50;
+  background: rgba(15, 18, 24, .9);
+  backdrop-filter: blur(10px);
+  border-bottom: 1px solid var(--border);
+}
+.switcher__inner {
+  max-width: 1600px;
+  margin: 0 auto;
+  padding: 10px 24px;
+  display: flex;
+  align-items: center;
+  gap: 14px;
+  flex-wrap: wrap;
+}
+.switcher .tabs { margin-left: auto; }
+.switcher .tab { border-radius: 8px; }
+
+.layout { padding-top: 22px; padding-bottom: 60px; display: flex; flex-direction: column; gap: 18px; }
+
+/* Split layout: sticky Projection-controls sidebar on the left (scrolls on its
+   own), everything else in a scrolling content column on the right. Lets you
+   tweak a control and watch the relevant panel on the right update in place. */
+.layout--split { flex-direction: row; align-items: flex-start; }
+.sidebar {
+  flex: 0 0 300px;
+  position: sticky;
+  top: 72px;
+  max-height: calc(100vh - 88px);
+  overflow-y: auto;
+  overscroll-behavior: contain; /* wheel inside the panel does not scroll the page */
+}
+.content { flex: 1 1 auto; min-width: 0; display: flex; flex-direction: column; gap: 18px; }
+
+.panel { background: var(--panel); border: 1px solid var(--border); border-radius: 12px; padding: 18px 20px; }
+.panel--accent { border-color: #2d4f7a; }
+.panel h2 { margin: 0 0 12px; font-size: 18px; }
+.panel__head { margin-bottom: 12px; }
+.muted { color: var(--muted); font-weight: 400; font-size: 13px; }
+.two-col { display: grid; grid-template-columns: 360px 1fr; gap: 18px; align-items: start; }
+
+.badge { font-size: 11px; padding: 3px 8px; border-radius: 999px; font-weight: 700; letter-spacing: .04em; }
+.badge--warn { background: var(--warn); color: #1a1206; }
+
+select, input {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  border-radius: 7px; padding: 7px 9px; font-size: 14px;
+}
+/* Drop the native number spinner buttons: in the narrow sidebar they eat ~17px
+   and clip long values (e.g. 15824 shown as "1582"). Values are typed, not
+   stepped, so the arrows add no value. */
+input[type="number"] { -moz-appearance: textfield; appearance: textfield; }
+input[type="number"]::-webkit-outer-spin-button,
+input[type="number"]::-webkit-inner-spin-button { -webkit-appearance: none; margin: 0; }
+/* Slightly tighter type + padding for the sidebar grids so 5-6 digit numbers
+   fit their column without truncation. */
+.controls-grid input, .manual-row input { font-size: 13px; padding-left: 8px; padding-right: 8px; }
+.field { display: flex; flex-direction: column; gap: 4px; }
+.field--inline { flex-direction: row; align-items: center; gap: 8px; }
+.field > span { font-size: 12px; color: var(--muted); }
+
+.kv-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(180px, 1fr)); gap: 10px; }
+.kv { background: var(--panel-2); border: 1px solid var(--border); border-radius: 8px; padding: 8px 10px; }
+.kv b { display: block; font-size: 12px; color: var(--muted); font-weight: 500; }
+.kv span { font-size: 15px; font-variant-numeric: tabular-nums; }
+
+.cr-schedule { margin-top: 12px; display: flex; flex-wrap: wrap; gap: 3px; }
+.cr-cell { width: 16px; height: 16px; border-radius: 3px; font-size: 0; }
+.cr-0 { background: #6b4a2d; } .cr-4 { background: #2d6b4a; } .cr-128 { background: #2d4a6b; }
+.cr-legend { display: flex; gap: 14px; margin-top: 8px; font-size: 12px; color: var(--muted); }
+.cr-legend i { display: inline-block; width: 12px; height: 12px; border-radius: 3px; margin-right: 5px; vertical-align: -1px; }
+
+.bd-block { margin-bottom: 16px; }
+.bd-block h3 { font-size: 14px; margin: 0 0 6px; }
+.bd-scroll { overflow-x: auto; }
+table.bd { border-collapse: collapse; font-size: 12px; min-width: 100%; }
+table.bd th, table.bd td { border: 1px solid var(--border); padding: 5px 8px; text-align: right; white-space: nowrap; font-variant-numeric: tabular-nums; }
+table.bd th { background: var(--panel-2); color: var(--muted); font-weight: 600; }
+table.bd td.rowlab, table.bd th.rowlab { text-align: left; color: var(--muted); position: sticky; left: 0; background: var(--panel); }
+.cell-compute { background: rgba(79,140,255,.10); }
+.cell-memory { background: rgba(160,140,200,.07); }
+.divider { border-left: 2px solid var(--accent) !important; }
+.phase-tag { font-size: 10px; color: var(--accent-2); }
+
+/* Layer-timing mode switch (trace vs manual), styled like the GPU tabs. */
+.mode-switch { display: flex; align-items: center; gap: 10px; margin-bottom: 12px; flex-wrap: wrap; }
+.mode-switch__label { font-size: 12px; color: var(--muted); }
+.mode-switch__tabs { display: flex; gap: 6px; }
+.mode-tab {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  padding: 6px 12px; border-radius: 7px; cursor: pointer; font-size: 13px;
+}
+.mode-tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; }
+
+.manual-grid { margin-bottom: 14px; }
+.manual-grid__header { display: flex; align-items: flex-start; justify-content: space-between; gap: 12px; }
+.manual-grid__hint { margin: 0 0 10px; flex: 1; }
+.manual-reset {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  padding: 5px 12px; border-radius: 7px; cursor: pointer; font-size: 12px; white-space: nowrap;
+}
+.manual-reset:hover { border-color: var(--accent); }
+.manual-rows { display: flex; flex-direction: column; gap: 8px; }
+.manual-grid__subhead { margin: 12px 0 6px; font-size: 11px; text-transform: uppercase; letter-spacing: .05em; }
+.manual-row__lab--nl {
+  display: inline-block; padding: 3px 7px; border-radius: 6px; align-self: center;
+  background: var(--panel-2); border: 1px solid var(--border); color: var(--text); font-size: 12px;
+}
+.manual-row {
+  display: grid; grid-template-columns: 92px 1fr 1fr; gap: 10px; align-items: end;
+}
+.manual-row__lab { font-size: 12px; padding-bottom: 8px; }
+.cr-tag {
+  display: inline-block; padding: 3px 7px; border-radius: 6px; color: #e6eaf2;
+  font-variant-numeric: tabular-nums; align-self: center; padding-bottom: 3px;
+}
+.manual-row .cr-tag { margin-bottom: 0; }
+.manual-row .field { min-width: 0; }
+.manual-row input { width: 100%; min-width: 0; }
+
+#breakdown-panel.is-muted #breakdown-blocks { opacity: .5; }
+.manual-note { margin: 8px 0 0; color: var(--warn); }
+
+.controls-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; }
+.controls-grid .field--full { grid-column: 1 / -1; }
+/* number inputs have an intrinsic min-width that overflows the 1fr column and
+   overlaps the results panel; force them to fit their grid cell. */
+.controls-grid .field { min-width: 0; }
+.controls-grid input, .controls-grid select { width: 100%; min-width: 0; }
+
+.validation {
+  margin-top: 10px;
+  border-radius: 8px;
+  padding: 9px 11px;
+  font-size: 12px;
+}
+.validation + .validation { margin-top: 8px; }
+.validation b { display: block; margin-bottom: 4px; }
+.validation ul { margin: 0; padding-left: 18px; }
+.validation--error { background: #3a1c1c; border: 1px solid #6b2d2d; color: #ffb4b4; }
+.validation--warn { background: #352913; border: 1px solid #6a4d1c; color: #ffd28a; }
+
+.headline { display: flex; flex-wrap: wrap; gap: 14px; margin-bottom: 16px; }
+.metric { background: var(--panel-2); border: 1px solid var(--border); border-radius: 10px; padding: 12px 16px; min-width: 150px; }
+.metric b { display: block; font-size: 12px; color: var(--muted); }
+.metric span { font-size: 22px; font-weight: 700; font-variant-numeric: tabular-nums; }
+.metric.metric--primary { border-color: var(--accent-2); }
+.metric.metric--primary span { color: var(--accent-2); }
+.metric.metric--error { border-color: #6b2d2d; }
+.metric.metric--error span { color: #ffb4b4; font-size: 18px; }
+
+.steps { display: flex; flex-direction: column; gap: 8px; }
+.step { background: var(--panel-2); border: 1px solid var(--border); border-radius: 8px; padding: 9px 12px; font-size: 13px; }
+.step b { color: var(--accent); }
+.step .num { float: right; font-variant-numeric: tabular-nums; color: var(--text); }
+.step small { color: var(--muted); }
+.step--error { background: #2a1616; border-color: #6b2d2d; }
+.step--error b { color: #ffb4b4; }
+.step--error small { color: #ffd0d0; }
+
+.state--error { background: #3a1c1c; border: 1px solid #6b2d2d; color: #ffb4b4; padding: 12px 16px; border-radius: 8px; }
+.site-footer { border-top: 1px solid var(--border); padding: 20px 24px; color: var(--muted); font-size: 13px; max-width: 1600px; margin: 0 auto; }
+code { background: var(--panel-2); padding: 1px 5px; border-radius: 4px; }
+
+@media (max-width: 880px) {
+  .two-col { grid-template-columns: 1fr; }
+  .layout--split { flex-direction: column; }
+  .sidebar { position: static; max-height: none; flex-basis: auto; width: 100%; overflow: visible; }
+}
+
+/* ---------------------------------------------------------------------------
+   Iteration timeline (design/07): 3-level composition view
+   --------------------------------------------------------------------------- */
+.tl-tabs { display: flex; gap: 6px; margin-top: 10px; flex-wrap: wrap; }
+.tl-tab {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  padding: 6px 14px; border-radius: 7px; cursor: pointer; font-size: 13px;
+}
+.tl-tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; }
+
+.tl-note { color: var(--muted); font-size: 12px; margin: 0 0 12px; }
+.tl-controls { display: flex; align-items: center; gap: 10px; margin-bottom: 12px; flex-wrap: wrap; }
+
+/* category colours (Level 1 / shared legend) */
+.tl-legend { display: flex; gap: 16px; flex-wrap: wrap; margin: 10px 0 0; font-size: 12px; color: var(--muted); }
+.tl-legend i { display: inline-block; width: 12px; height: 12px; border-radius: 3px; margin-right: 5px; vertical-align: -1px; }
+.tl-c-attn { background: #4f8cff; }
+.tl-c-mlp  { background: #36c08f; }
+.tl-c-a2a  { background: #e0a03a; }
+.tl-c-misc { background: #8d97a8; }
+.tl-c-emb  { background: #a06cd5; }
+.tl-c-out  { background: #d5675a; }
+.tl-c-manual { background: #9a6cff; }
+
+/* Level 1: per-cr bidirectional stacked bars */
+.tl-l1-row { display: grid; grid-template-columns: 120px 1fr; gap: 12px; align-items: center; margin-bottom: 14px; }
+.tl-l1-lab { font-size: 13px; }
+.tl-l1-lab small { display: block; color: var(--muted); font-size: 11px; }
+.tl-bars { display: flex; flex-direction: column; gap: 6px; }
+.tl-bar { display: flex; align-items: center; gap: 8px; }
+.tl-bar__tag { width: 34px; font-size: 11px; color: var(--muted); text-align: right; }
+.tl-bar__track { position: relative; flex: 1; height: 26px; background: var(--panel-2); border: 1px solid var(--border); border-radius: 6px; overflow: hidden; display: flex; }
+.tl-seg { height: 100%; display: flex; align-items: center; justify-content: center; font-size: 10px; color: #08101c; overflow: hidden; white-space: nowrap; cursor: default; }
+.tl-seg.tl-c-misc, .tl-seg.tl-c-out, .tl-seg.tl-c-manual { color: #f5f7fb; }
+.tl-bar__total { width: 70px; font-size: 12px; font-variant-numeric: tabular-nums; color: var(--text); }
+
+/* Level 2: per-device layer strips */
+.tl-l2-row { display: grid; grid-template-columns: 150px 1fr 120px; gap: 12px; align-items: center; margin-bottom: 8px; }
+.tl-l2-lab { font-size: 12px; }
+.tl-l2-lab small { display: block; color: var(--muted); font-size: 11px; }
+.tl-strip { display: flex; height: 22px; border: 1px solid var(--border); border-radius: 5px; overflow: hidden; background: var(--panel-2); }
+.tl-cell { height: 100%; min-width: 2px; border-right: 1px solid rgba(0,0,0,.25); }
+.tl-cell:last-child { border-right: none; }
+.tl-cell--recompute { background-image: repeating-linear-gradient(45deg, transparent, transparent 3px, rgba(255,255,255,.35) 3px, rgba(255,255,255,.35) 5px); }
+.tl-chunk-gap { width: 3px; background: var(--bg); }
+.tl-cell--emb { background: #a06cd5 !important; }
+.tl-cell--out { background: #d5675a !important; }
+.tl-l2-meta { font-size: 11px; color: var(--muted); font-variant-numeric: tabular-nums; }
+.tl-l2-row.is-critical .tl-strip { outline: 2px solid var(--accent-2); outline-offset: 1px; }
+.tl-l2-row.is-critical .tl-l2-lab { color: var(--accent-2); }
+
+/* Level 3: pipeline schedule Gantt */
+.tl-gantt-wrap { overflow-x: auto; }
+.tl-gantt { display: block; }
+/* text colours are set via inline `fill` attributes in app.js (a stylesheet
+   `fill` here would override those SVG presentation attributes, e.g. forcing the
+   in-cell microbatch numbers grey instead of black). */
+.tl-gantt rect.tl-fwd { stroke: rgba(0,0,0,.35); stroke-width: .5; }
+.tl-gantt rect.tl-bwd { stroke: rgba(0,0,0,.35); stroke-width: .5; }
+.tl-gantt rect.tl-bubble { fill: transparent; }
+.tl-summary { font-size: 12px; color: var(--muted); margin-top: 10px; font-variant-numeric: tabular-nums; }
+.tl-warn { color: var(--warn); }
+
+/* layout switch + level tabs on one row */
+.tl-head-row { display: flex; align-items: center; gap: 18px; flex-wrap: wrap; margin-top: 10px; }
+.tl-viewswitch { display: flex; align-items: center; gap: 8px; }
+.tl-view {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  padding: 6px 12px; border-radius: 7px; cursor: pointer; font-size: 13px;
+}
+.tl-view.is-active { background: var(--accent-2); border-color: var(--accent-2); color: #06231a; font-weight: 600; }
+
+/* stacked sections */
+.tl-section { border-top: 1px solid var(--border); padding-top: 14px; margin-top: 14px; }
+.tl-section:first-of-type { border-top: none; padding-top: 0; margin-top: 4px; }
+.tl-section__title { font-size: 14px; margin: 0 0 10px; }
+.tl-section__sub { font-size: 12px; color: var(--muted); font-weight: 400; margin-left: 8px; }
+
+/* linkage selection indicator — floating card, does not affect page layout */
+.tl-selbar {
+  position: fixed;
+  right: 24px;
+  bottom: 24px;
+  z-index: 60;
+  display: flex; align-items: center; gap: 14px;
+  background: rgba(23, 28, 38, .97);
+  border: 1px solid var(--accent-2);
+  box-shadow: 0 8px 26px rgba(0, 0, 0, .5);
+  border-radius: 10px;
+  padding: 10px 14px; font-size: 12px; color: var(--text);
+  max-width: 360px;
+  backdrop-filter: blur(6px);
+}
+.tl-selbar__txt { line-height: 1.35; }
+.tl-clear {
+  background: var(--panel-2); color: var(--text); border: 1px solid var(--border);
+  padding: 5px 14px; border-radius: 7px; cursor: pointer; font-size: 12px; white-space: nowrap;
+}
+.tl-clear:hover { border-color: var(--accent-2); }
+@media (max-width: 640px) { .tl-selbar { left: 16px; right: 16px; bottom: 16px; max-width: none; } }
+
+/* linkage highlight states (shared) */
+.tl-clickable { cursor: pointer; }
+.tl-dim { opacity: .32; transition: opacity .12s; }
+.tl-l1-row.is-sel, .tl-l2-row.is-sel .tl-strip { outline: 2px solid var(--accent); outline-offset: 2px; border-radius: 6px; }
+.tl-l1-row.is-sel { background: rgba(79,140,255,.08); border-radius: 8px; }
+.tl-cell--sel { outline: 2px solid #fff; outline-offset: -2px; }
+
+/* gantt export toolbar */
+.tl-export { display: flex; gap: 8px; margin-bottom: 8px; }
+.tl-export .mode-tab { background: var(--panel-2); }
+
+/* gantt zoom slider */
+.tl-zoom-wrap { display: inline-flex; align-items: center; gap: 8px; margin-left: 6px; }
+.tl-zoom { width: 160px; accent-color: var(--accent); }
+.tl-zoom-val { font-size: 12px; color: var(--muted); font-variant-numeric: tabular-nums; min-width: 30px; }
+.tl-axis-note { margin: 8px 0 0; }
diff --git a/examples/deepseek-v4/projection/site/data/flash.json b/examples/deepseek-v4/projection/site/data/flash.json
new file mode 100644
index 000000000..31a744e44
--- /dev/null
+++ b/examples/deepseek-v4/projection/site/data/flash.json
@@ -0,0 +1,1727 @@
+{
+  "schema_version": 1,
+  "model": "flash",
+  "analytic_flops": {
+    "per_cr_layer_flops": {
+      "0": 28775232307200,
+      "4": 36583482851328,
+      "128": 29596695134208
+    },
+    "output_flops": 13013750906880,
+    "mtp": {
+      "num_layers": 1,
+      "compress_ratio": 4,
+      "inner_layer_flops": 36583482851328,
+      "eh_proj_flops": 824633720832,
+      "extra_logits_flops": 13013750906880,
+      "hc_head_flops": 3221225472
+    },
+    "seq": 4096,
+    "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP"
+  },
+  "generated_at": "2026-06-30T13:59:46Z",
+  "provenance": {
+    "traces": {
+      "0": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr0_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr0_seq4096_ep8]-rank[0].1782149271209635673.pt.trace.json",
+      "4": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr4_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr4_seq4096_ep8]-rank[0].1782148745954636194.pt.trace.json",
+      "128": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr128_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr128_seq4096_ep8]-rank[0].1782149357925690324.pt.trace.json"
+    },
+    "graphed_crs_estimated_from_cr4": [],
+    "dropped_stall_us_per_mb": {
+      "0": 0.0,
+      "4": 19280.7,
+      "128": 0.0
+    },
+    "note": "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured (compute not visible in trace); their breakdown is copied from cr=4 as an estimate. Re-run those cr with graph capture disabled for exact numbers. dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)."
+  },
+  "capture": {
+    "gpu": "MI355X",
+    "seq_length": 4096,
+    "micro_batch_size": 1,
+    "tokens_per_microbatch": 4096,
+    "ep": 8,
+    "ga_for_capture": 2,
+    "optimizer": "adam",
+    "distributed_optimizer": false,
+    "recompute": "off",
+    "measured_iter_time_ms": 4076.0,
+    "measured_anchor": {
+      "config": "PP8/VPP1/EP8/TP1, world 64, MBS1/GBS128, seq4096, full recompute, layout Et*4|t*5|(t*6|)*5,t*4mL (8-node MI355X)",
+      "iter_ms": 4076.0,
+      "tflops_gpu": 722,
+      "tok_s_gpu": 2010,
+      "calib_factor": 0.87,
+      "tolerance": "<0.1%"
+    }
+  },
+  "model_config": {
+    "num_layers": 43,
+    "hidden_size": 4096,
+    "num_attention_heads": 64,
+    "kv_channels": 512,
+    "num_experts": 256,
+    "moe_router_topk": 6,
+    "moe_ffn_hidden_size": 2048,
+    "moe_shared_expert_intermediate_size": 2048,
+    "index_topk": 512,
+    "vocab_size": 129280,
+    "mtp_num_layers": 1,
+    "mtp_compress_ratios": [
+      4
+    ],
+    "pipeline_layout": "Et*10|t*11|t*11|t*11mL",
+    "compress_ratios": [
+      0,
+      0,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      0
+    ],
+    "cr_layer_counts": {
+      "0": 3,
+      "4": 20,
+      "128": 20
+    },
+    "total_params": 290419900416
+  },
+  "hardware": {
+    "MI355X": {
+      "peak_tflops_bf16": 2500.0,
+      "hbm_bandwidth_gbps": 8000.0
+    },
+    "MI455X": {
+      "peak_tflops_bf16": 10000.0,
+      "hbm_bandwidth_gbps": 19600.0
+    }
+  },
+  "layers": {
+    "0": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 2768.835,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 600.524
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 488.376
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 252.844
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 190.953
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 151.259
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 136.028
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 732.222,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 449360953344.0,
+            "tflops": 613.7,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 260.835
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 226.171
+              },
+              {
+                "name": "Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950",
+                "time_us": 96.207
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 56.9
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 36.157
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 29.943
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 496.029,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 93.8
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 92.809
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 83.298
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 63.337
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 62.797
+              },
+              {
+                "name": "_apply_rope_fwd_kernel",
+                "time_us": 52.313
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 213.574,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 209.635
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 3.939
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.misc",
+            "time_us": 6266.462,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 1253.002
+              },
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 897.244
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 837.772
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 746.936
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 341.962
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 280.544
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 5292.432,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 2874943733760.0,
+            "tflops": 543.2,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 2713.086
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 1783.792
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 306.172
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 145.373
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 106.597
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 106.263
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 919.753,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_dkv_kernel",
+                "time_us": 437.837
+              },
+              {
+                "name": "_v4_attention_bwd_dq_kernel",
+                "time_us": 394.517
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 83.214
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 4.186
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 52.13,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 52.13
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1470.924,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 950.222
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 308.935
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 88.64
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 29.89
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 24.733
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 22.62
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 363.919,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 233.282
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 111.084
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 19.553
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 321.999,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 231.632
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 90.367
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 112.521,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 68719476736.0,
+            "tflops": 610.7,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 78.11
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.219
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 5.396
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 4.426
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 3.369
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 74.758,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 24.829
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 10.252
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 8.02
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 6.869
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 5.616
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 4.652
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1916.666,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 1465.403
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 327.882
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 115.553
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 3.986
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 3.842
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 298.379,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 238.449
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 59.93
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 242.331,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 221.765
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 20.566
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 44.476,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 44.476
+              }
+            ]
+          }
+        ]
+      }
+    },
+    "4": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 2773.299,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 603.359
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 493.853
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 247.801
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 192.373
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 150.529
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 136.255
+              }
+            ]
+          },
+          {
+            "module": "attn.indexer",
+            "time_us": 797.919,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 152.083
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 98.513
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 86.043
+              },
+              {
+                "name": "void at::native::reduce_kernel<128, 4, at::native::ReduceOp<c10::BFloat16, at::n",
+                "time_us": 68.023
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 61.763
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 45.115
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 736.399,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 449360953344.0,
+            "tflops": 610.2,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 266.672
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 223.295
+              },
+              {
+                "name": "Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950",
+                "time_us": 95.721
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 57.19
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 36.907
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 30.46
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 713.905,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_csa_attention_pool_sparse_merge_fwd_kernel",
+                "time_us": 503.845
+              },
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 206.032
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 4.029
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 576.738,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 119.475
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 92.917
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 92.574
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 64.33
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 63.18
+              },
+              {
+                "name": "_apply_rope_fwd_kernel",
+                "time_us": 55.569
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.misc",
+            "time_us": 6906.916,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 1292.801
+              },
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 894.834
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 862.656
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 754.642
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 499.523
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 469.792
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 6885.908,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_kernel",
+                "time_us": 5194.447
+              },
+              {
+                "name": "_v4_csa_attention_pool_sparse_bwd_partial_kernel",
+                "time_us": 1369.16
+              },
+              {
+                "name": "_v4_csa_attention_pool_segreduce_kernel",
+                "time_us": 238.882
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 79.314
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 4.106
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 5356.593,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 2943663210496.0,
+            "tflops": 549.5,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 2691.996
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 1819.353
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 305.419
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 144.986
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 106.44
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 106.05
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 54.462,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 54.462
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1461.438,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 945.722
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 305.939
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 86.43
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 28.903
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 26.583
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 21.893
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 330.683,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 230.235
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 100.447
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 279.761,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 232.332
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 27.653
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 19.776
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 112.854,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 68719476736.0,
+            "tflops": 608.9,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 78.307
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.372
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 5.393
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 4.353
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 3.429
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 74.189,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 24.839
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 10.072
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 8.086
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 6.905
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 5.669
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 4.696
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1896.364,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 1464.913
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 308.679
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 114.644
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 4.076
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 4.052
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 399.573,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 221.485
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 178.088
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 357.633,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 237.535
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 120.097
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 44.726,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 44.726
+              }
+            ]
+          }
+        ]
+      }
+    },
+    "128": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 2784.297,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 605.113
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 493.93
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 248.877
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 191.946
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 151.909
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 136.028
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 851.58,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::(anonymous namespace)::CatArrayBatchedCopy<at::native::(anonymo",
+                "time_us": 322.669
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 96.885
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 94.427
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 86.567
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 63.153
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 63.077
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 741.746,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 449360953344.0,
+            "tflops": 605.8,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 262.638
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 231.228
+              },
+              {
+                "name": "Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950",
+                "time_us": 96.854
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 56.606
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 36.38
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 31.55
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 327.905,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 323.406
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 4.499
+              }
+            ]
+          },
+          {
+            "module": "attn.indexer",
+            "time_us": 116.287,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::(anonymous namespace)::cunn_SpatialSoftMaxForward<float, float,",
+                "time_us": 43.017
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x64x64_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 40.393
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 5.909
+              },
+              {
+                "name": "void at::native::reduce_kernel<128, 4, at::native::ReduceOp<c10::BFloat16, at::n",
+                "time_us": 5.543
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 4.436
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 4.359
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.misc",
+            "time_us": 6427.921,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 1276.619
+              },
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 900.44
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 860.911
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 759.74
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 346.078
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 286.755
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 5346.296,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 2909303472128.0,
+            "tflops": 544.2,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 2699.575
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 1790.855
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 304.992
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 144.849
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 106.241
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 106.113
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 1333.44,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_dq_kernel",
+                "time_us": 486.324
+              },
+              {
+                "name": "_v4_attention_bwd_dkv_kernel",
+                "time_us": 448.724
+              },
+              {
+                "name": "_v4_attention_bwd_dkv_pool_mha_kernel",
+                "time_us": 313.569
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 80.704
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 4.119
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 58.646,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 58.646
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1476.627,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 950.715
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 312.572
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 88.707
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 30.373
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 25.726
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 22.426
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 310.355,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 229.722
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 60.447
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 20.186
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 309.929,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 227.382
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 82.547
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 112.514,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 68719476736.0,
+            "tflops": 610.8,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 77.993
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.266
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 5.43
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 4.356
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 3.47
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 83.718,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 25.606
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 14.203
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 10.482
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 7.482
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 6.356
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 4.886
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 1888.822,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 1458.196
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 309.245
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 113.277
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 4.056
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 4.049
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 343.592,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024>(hip_bfloat",
+                "time_us": 233.195
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 110.397
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 225.508,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type<int, 4u",
+                "time_us": 221.218
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 4.29
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 44.606,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 44.606
+              }
+            ]
+          }
+        ]
+      }
+    }
+  },
+  "non_layer": {
+    "embedding": {
+      "forward": [
+        {
+          "module": "embedding",
+          "time_us": 7.84,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "void at::native::vectorized_gather_kernel<16, long>(char*, char*, long*, int, lo",
+              "time_us": 7.84
+            }
+          ]
+        }
+      ],
+      "backward": [
+        {
+          "module": "embedding",
+          "time_us": 86.591,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "void at::native::(anonymous namespace)::compute_grad_weight<c10::BFloat16, long>",
+              "time_us": 33.993
+            },
+            {
+              "name": "void rocprim::ROCPRIM_400200_NS::detail::trampoline_kernel<rocprim::ROCPRIM_4002",
+              "time_us": 21.41
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::sum_and_scatter<c10::BFloat16, long>(lon",
+              "time_us": 21.21
+            },
+            {
+              "name": "void rocprim::ROCPRIM_400200_NS::detail::init_lookback_scan_state_kernel<rocprim",
+              "time_us": 3.929
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::krn_partial_segment_offset<long>(long*, ",
+              "time_us": 2.08
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::krn_partials_per_segment<long>(long*, lo",
+              "time_us": 2.023
+            }
+          ]
+        }
+      ]
+    },
+    "output": {
+      "forward": [
+        {
+          "module": "output",
+          "time_us": 1553.935,
+          "class": "compute_bound",
+          "flop_class": "gemm",
+          "flops": 2171374403584.0,
+          "tflops": 1397.3,
+          "kernels": [
+            {
+              "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+              "time_us": 1529.095
+            },
+            {
+              "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB1_AFC1_AF",
+              "time_us": 24.84
+            }
+          ]
+        }
+      ],
+      "backward": []
+    },
+    "loss": {
+      "forward": [
+        {
+          "module": "loss",
+          "time_us": 368.303,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "cross_entropy_kernel",
+              "time_us": 235.982
+            },
+            {
+              "name": "online_softmax_kernel",
+              "time_us": 132.321
+            }
+          ]
+        }
+      ],
+      "backward": [
+        {
+          "module": "loss",
+          "time_us": 197.392,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "element_mul_kernel",
+              "time_us": 197.392
+            }
+          ]
+        }
+      ]
+    }
+  },
+  "optimizer": {
+    "type": "adam",
+    "measured_params": null,
+    "time_us": 18459.0,
+    "bytes_per_param": 30,
+    "class": "memory_bound",
+    "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference"
+  },
+  "comm": {
+    "ep_dispatch_us": null,
+    "ep_combine_us": null,
+    "note": "EP dispatch/combine are memory_bound rows inside moe; informational"
+  }
+}
diff --git a/examples/deepseek-v4/projection/site/data/pro.json b/examples/deepseek-v4/projection/site/data/pro.json
new file mode 100644
index 000000000..8cfd34cb4
--- /dev/null
+++ b/examples/deepseek-v4/projection/site/data/pro.json
@@ -0,0 +1,1750 @@
+{
+  "schema_version": 1,
+  "model": "pro",
+  "analytic_flops": {
+    "per_cr_layer_flops": {
+      "0": 76885870510080,
+      "4": 95420801875968,
+      "128": 78425716948992
+    },
+    "output_flops": 22774064087040,
+    "mtp": {
+      "num_layers": 1,
+      "compress_ratio": 4,
+      "inner_layer_flops": 95420801875968,
+      "eh_proj_flops": 2525440770048,
+      "extra_logits_flops": 22774064087040,
+      "hc_head_flops": 5637144576
+    },
+    "seq": 4096,
+    "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP"
+  },
+  "generated_at": "2026-06-30T13:59:43Z",
+  "provenance": {
+    "traces": {
+      "0": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr0_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr0_seq4096_ep8]-rank[0].1781792532762897620.pt.trace.json",
+      "4": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr4_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr4_seq4096_ep8]-rank[0].1781792606762968569.pt.trace.json",
+      "128": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr128_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr128_seq4096_ep8]-rank[0].1781792680203624431.pt.trace.json"
+    },
+    "graphed_crs_estimated_from_cr4": [],
+    "dropped_stall_us_per_mb": {
+      "0": 36705.7,
+      "4": 0.0,
+      "128": 0.0
+    },
+    "note": "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured (compute not visible in trace); their breakdown is copied from cr=4 as an estimate. Re-run those cr with graph capture disabled for exact numbers. dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)."
+  },
+  "capture": {
+    "gpu": "MI355X",
+    "seq_length": 4096,
+    "micro_batch_size": 1,
+    "tokens_per_microbatch": 4096,
+    "ep": 8,
+    "ga_for_capture": 2,
+    "optimizer": "adam",
+    "distributed_optimizer": false,
+    "recompute": "off",
+    "measured_iter_time_ms": null,
+    "measured_anchor": null
+  },
+  "model_config": {
+    "num_layers": 61,
+    "hidden_size": 7168,
+    "num_attention_heads": 128,
+    "kv_channels": 512,
+    "num_experts": 384,
+    "moe_router_topk": 6,
+    "moe_ffn_hidden_size": 3072,
+    "moe_shared_expert_intermediate_size": 3072,
+    "index_topk": 1024,
+    "vocab_size": 129280,
+    "mtp_num_layers": 1,
+    "mtp_compress_ratios": [
+      4
+    ],
+    "pipeline_layout": "",
+    "compress_ratios": [
+      128,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      4,
+      128,
+      0
+    ],
+    "cr_layer_counts": {
+      "0": 1,
+      "4": 29,
+      "128": 31
+    },
+    "total_params": 1597579198464
+  },
+  "hardware": {
+    "MI355X": {
+      "peak_tflops_bf16": 2500.0,
+      "hbm_bandwidth_gbps": 8000.0
+    },
+    "MI455X": {
+      "peak_tflops_bf16": 10000.0,
+      "hbm_bandwidth_gbps": 19600.0
+    }
+  },
+  "layers": {
+    "0": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 7402.198,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 1437.355
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::N",
+                "time_us": 1366.235
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 1344.257
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 709.616
+              },
+              {
+                "name": "void at::native::unrolled_elementwise_kernel<at::native::AUnaryFunctor<float, fl",
+                "time_us": 374.819
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 324.583
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 1623.031,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 1191719206912.0,
+            "tflops": 734.3,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 905.959
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 476.296
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 117.74
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 93.017
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 21.693
+              },
+              {
+                "name": "Cijk_SS_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4",
+                "time_us": 4.239
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 929.729,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 177.157
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 176.451
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 148.509
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 127.904
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 117.211
+              },
+              {
+                "name": "_apply_rope_fwd_kernel",
+                "time_us": 105.693
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 449.499,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 445.103
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 4.396
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.misc",
+            "time_us": 13148.636,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 3024.544
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 2081.72
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 1555.489
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 1328.534
+              },
+              {
+                "name": "void transformer_engine::normalization::rmsnorm_bwd_general_kernel<transformer_e",
+                "time_us": 641.61
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 616.003
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 10065.546,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 5279185895424.0,
+            "tflops": 524.5,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 4938.364
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 3769.201
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 489.093
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 201.328
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 168.756
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB0_AFC1",
+                "time_us": 161.16
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 1753.653,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_dkv_kernel",
+                "time_us": 867.357
+              },
+              {
+                "name": "_v4_attention_bwd_dq_kernel",
+                "time_us": 757.883
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 123.454
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 4.959
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 105.86,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 105.86
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 4225.655,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 2683.892
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 1208.776
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 135.577
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 44.34
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 43.473
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 41.19
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 427.833,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 386.653
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 41.18
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 422.925,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 382.809
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 20.72
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 19.396
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 252.402,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC",
+                "time_us": 141.417
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 71.463
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.462
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 7.096
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 6.093
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 4.869
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 91.185,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 25.516
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 16.113
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 14.06
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 7.365
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 6.666
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 5.559
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 5933.968,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 4446.509
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 1055.078
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 423.883
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 4.526
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 3.972
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 586.867,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 394.596
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 192.271
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 449.116,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 375.466
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 73.65
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 44.656,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 44.656
+              }
+            ]
+          }
+        ]
+      }
+    },
+    "4": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 7425.221,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 1447.247
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::N",
+                "time_us": 1378.227
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 1342.16
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 709.102
+              },
+              {
+                "name": "void at::native::unrolled_elementwise_kernel<at::native::AUnaryFunctor<float, fl",
+                "time_us": 378.692
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 323.907
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 2128.316,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_csa_attention_pool_sparse_merge_fwd_kernel",
+                "time_us": 1717.604
+              },
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 405.95
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 4.762
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 1622.73,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 1191719206912.0,
+            "tflops": 734.4,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 914.619
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 468.363
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 117.723
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 92.227
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 21.286
+              },
+              {
+                "name": "Cijk_SS_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4",
+                "time_us": 4.346
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 1021.024,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 206.375
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 177.044
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 160.454
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 127.904
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 123.664
+              },
+              {
+                "name": "_apply_rope_fwd_kernel",
+                "time_us": 107.629
+              }
+            ]
+          },
+          {
+            "module": "attn.indexer",
+            "time_us": 874.391,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 152.603
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 96.286
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 96.272
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 84.363
+              },
+              {
+                "name": "void at::native::reduce_kernel<128, 4, at::native::ReduceOp<c10::BFloat16, at::n",
+                "time_us": 66.402
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 48.821
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.core",
+            "time_us": 15675.167,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_kernel",
+                "time_us": 10309.148
+              },
+              {
+                "name": "_v4_csa_attention_pool_sparse_bwd_partial_kernel",
+                "time_us": 4963.764
+              },
+              {
+                "name": "_v4_csa_attention_pool_segreduce_kernel",
+                "time_us": 276.325
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 120.904
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 5.026
+              }
+            ]
+          },
+          {
+            "module": "attn.misc",
+            "time_us": 14359.617,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 3023.79
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 2219.844
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 1570.513
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 1332.678
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 898.568
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 890.739
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 10271.383,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 5399444979712.0,
+            "tflops": 525.7,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 4991.694
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 3854.896
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 487.807
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 200.238
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 168.36
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB0_AFC1",
+                "time_us": 159.567
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 108.736,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 108.736
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 4214.341,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 2677.845
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 1206.016
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 135.8
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 43.96
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 43.02
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 40.16
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 486.166,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 384.236
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 82.547
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 19.383
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 470.04,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 385.316
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 84.724
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 254.469,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC",
+                "time_us": 143.571
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 71.73
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.249
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 6.96
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 6.136
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 4.823
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 91.678,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 25.176
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 16.116
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 13.973
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 7.535
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 6.562
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 5.852
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 5908.193,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 4400.005
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 1074.128
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 425.279
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 4.636
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 4.146
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 580.63,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 376.306
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 204.325
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 503.433,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 394.109
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 109.324
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 44.866,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 44.866
+              }
+            ]
+          }
+        ]
+      }
+    },
+    "128": {
+      "attention": {
+        "forward": [
+          {
+            "module": "attn.misc",
+            "time_us": 7410.422,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 1443.783
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::N",
+                "time_us": 1371.851
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<floa",
+                "time_us": 1343.05
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<float>",
+                "time_us": 708.356
+              },
+              {
+                "name": "void at::native::unrolled_elementwise_kernel<at::native::AUnaryFunctor<float, fl",
+                "time_us": 377.102
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 327.05
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 1639.461,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 1191719206912.0,
+            "tflops": 726.9,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 919.326
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 478.756
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 117.167
+              },
+              {
+                "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF",
+                "time_us": 94.187
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 21.563
+              },
+              {
+                "name": "Cijk_SS_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4",
+                "time_us": 4.369
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 1623.015,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::(anonymous namespace)::CatArrayBatchedCopy<at::native::(anonymo",
+                "time_us": 653.885
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 181.199
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 177.308
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 156.111
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 127.617
+              },
+              {
+                "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp<float, at::native::M",
+                "time_us": 119.317
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 643.754,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_fwd_kernel",
+                "time_us": 639.152
+              },
+              {
+                "name": "_hc_compute_tail_fwd_kernel",
+                "time_us": 4.602
+              }
+            ]
+          },
+          {
+            "module": "attn.indexer",
+            "time_us": 127.799,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::(anonymous namespace)::cunn_SpatialSoftMaxForward<float, float,",
+                "time_us": 45.123
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC",
+                "time_us": 40.536
+              },
+              {
+                "name": "Cijk_SB_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4",
+                "time_us": 8.172
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 5.979
+              },
+              {
+                "name": "void at::native::reduce_kernel<128, 4, at::native::ReduceOp<c10::BFloat16, at::n",
+                "time_us": 5.803
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 4.566
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "attn.misc",
+            "time_us": 13440.13,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc",
+                "time_us": 3030.941
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 2198.898
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 1584.346
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor<floa",
+                "time_us": 1330.558
+              },
+              {
+                "name": "void transformer_engine::normalization::rmsnorm_bwd_general_kernel<transformer_e",
+                "time_us": 641.923
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<fl",
+                "time_us": 621.813
+              }
+            ]
+          },
+          {
+            "module": "attn.proj",
+            "time_us": 10203.677,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": 5339315437568.0,
+            "tflops": 523.3,
+            "kernels": [
+              {
+                "name": "Cijk_Ailk_Bjlk_BSS_BH_Bias_S_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB0_A",
+                "time_us": 4952.69
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 3794.003
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1",
+                "time_us": 487.776
+              },
+              {
+                "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 201.361
+              },
+              {
+                "name": "Cijk_Ailk_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT256x256x16_MI16x16x1_SN_LDSB0_AFC1_A",
+                "time_us": 169.553
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB0_AFC1",
+                "time_us": 160.704
+              }
+            ]
+          },
+          {
+            "module": "attn.core",
+            "time_us": 2569.999,
+            "class": "compute_bound",
+            "flop_class": "attn",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_v4_attention_bwd_dq_kernel",
+                "time_us": 928.591
+              },
+              {
+                "name": "_v4_attention_bwd_dkv_kernel",
+                "time_us": 892.0
+              },
+              {
+                "name": "_v4_attention_bwd_dkv_pool_mha_kernel",
+                "time_us": 619.438
+              },
+              {
+                "name": "_v4_attention_bwd_preprocess_kernel",
+                "time_us": 124.994
+              },
+              {
+                "name": "_hc_compute_tail_bwd_kernel",
+                "time_us": 4.976
+              }
+            ]
+          },
+          {
+            "module": "attn.norm",
+            "time_us": 108.186,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_apply_rope_bwd_kernel",
+                "time_us": 108.186
+              }
+            ]
+          }
+        ]
+      },
+      "moe": {
+        "forward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 4170.341,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 2687.575
+              },
+              {
+                "name": "_stack_grouped_weight_fwd_kernel",
+                "time_us": 1151.912
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 135.587
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern",
+                "time_us": 44.276
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous",
+                "time_us": 42.447
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 40.733
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 514.457,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 387.153
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 127.304
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 448.916,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 384.249
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ",
+                "time_us": 45.283
+              },
+              {
+                "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ",
+                "time_us": 19.383
+              }
+            ]
+          },
+          {
+            "module": "moe.shared_expert",
+            "time_us": 254.795,
+            "class": "compute_bound",
+            "flop_class": "gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC",
+                "time_us": 143.427
+              },
+              {
+                "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+                "time_us": 71.993
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel",
+                "time_us": 21.433
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa",
+                "time_us": 7.026
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor<c10:",
+                "time_us": 6.246
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_",
+                "time_us": 4.67
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 90.989,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_fwd_kernel",
+                "time_us": 25.629
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_",
+                "time_us": 16.009
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor<bool>",
+                "time_us": 13.92
+              },
+              {
+                "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc",
+                "time_us": 7.152
+              },
+              {
+                "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa",
+                "time_us": 6.402
+              },
+              {
+                "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel",
+                "time_us": 5.779
+              }
+            ]
+          }
+        ],
+        "backward": [
+          {
+            "module": "moe.grouped_gemm",
+            "time_us": 5906.347,
+            "class": "compute_bound",
+            "flop_class": "grouped_gemm",
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel<ck_tile::GemmSpatiallyLocalTi",
+                "time_us": 4397.548
+              },
+              {
+                "name": "_stack_grouped_weight_bwd_kernel",
+                "time_us": 1076.558
+              },
+              {
+                "name": "void primus_turbo::grouped_gemm_variable_k_postprocess<std::bfloat16_t>(std::bfl",
+                "time_us": 423.59
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_variable_k_args<std::bfloat16_t, std::bf",
+                "time_us": 4.516
+              },
+              {
+                "name": "void primus_turbo::compute_grouped_gemm_args<std::bfloat16_t, std::bfloat16_t, s",
+                "time_us": 4.136
+              }
+            ]
+          },
+          {
+            "module": "moe.combine",
+            "time_us": 479.096,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::combine<hip_bfloat16, 8, 1024, true>(hip_",
+                "time_us": 394.309
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in",
+                "time_us": 84.787
+              }
+            ]
+          },
+          {
+            "module": "moe.dispatch",
+            "time_us": 408.969,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type<i",
+                "time_us": 375.616
+              },
+              {
+                "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int",
+                "time_us": 33.353
+              }
+            ]
+          },
+          {
+            "module": "moe.router",
+            "time_us": 45.489,
+            "class": "memory_bound",
+            "flop_class": null,
+            "flops": null,
+            "tflops": null,
+            "kernels": [
+              {
+                "name": "_sinkhorn_bwd_kernel",
+                "time_us": 45.489
+              }
+            ]
+          }
+        ]
+      }
+    }
+  },
+  "non_layer": {
+    "embedding": {
+      "forward": [
+        {
+          "module": "embedding",
+          "time_us": 14.566,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "void at::native::vectorized_gather_kernel<16, long>(char*, char*, long*, int, lo",
+              "time_us": 14.566
+            }
+          ]
+        }
+      ],
+      "backward": [
+        {
+          "module": "embedding",
+          "time_us": 129.501,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "void at::native::(anonymous namespace)::compute_grad_weight<c10::BFloat16, long>",
+              "time_us": 60.12
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::sum_and_scatter<c10::BFloat16, long>(lon",
+              "time_us": 38.046
+            },
+            {
+              "name": "void rocprim::ROCPRIM_400200_NS::detail::trampoline_kernel<rocprim::ROCPRIM_4002",
+              "time_us": 21.477
+            },
+            {
+              "name": "void rocprim::ROCPRIM_400200_NS::detail::init_lookback_scan_state_kernel<rocprim",
+              "time_us": 3.886
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::krn_partials_per_segment<long>(long*, lo",
+              "time_us": 2.026
+            },
+            {
+              "name": "void at::native::(anonymous namespace)::krn_partial_segment_offset<long>(long*, ",
+              "time_us": 2.023
+            }
+          ]
+        }
+      ]
+    },
+    "output": {
+      "forward": [
+        {
+          "module": "output",
+          "time_us": 2688.195,
+          "class": "compute_bound",
+          "flop_class": "gemm",
+          "flops": 3799905206272.0,
+          "tflops": 1413.6,
+          "kernels": [
+            {
+              "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0",
+              "time_us": 2634.481
+            },
+            {
+              "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB1_AFC1_AF",
+              "time_us": 53.713
+            }
+          ]
+        }
+      ],
+      "backward": []
+    },
+    "loss": {
+      "forward": [
+        {
+          "module": "loss",
+          "time_us": 350.465,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "cross_entropy_kernel",
+              "time_us": 222.888
+            },
+            {
+              "name": "online_softmax_kernel",
+              "time_us": 127.577
+            }
+          ]
+        }
+      ],
+      "backward": [
+        {
+          "module": "loss",
+          "time_us": 195.848,
+          "class": "memory_bound",
+          "flop_class": null,
+          "flops": null,
+          "tflops": null,
+          "kernels": [
+            {
+              "name": "element_mul_kernel",
+              "time_us": 195.848
+            }
+          ]
+        }
+      ]
+    }
+  },
+  "optimizer": {
+    "type": "adam",
+    "measured_params": null,
+    "time_us": 87037.8,
+    "bytes_per_param": 30,
+    "class": "memory_bound",
+    "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference"
+  },
+  "comm": {
+    "ep_dispatch_us": null,
+    "ep_combine_us": null,
+    "note": "EP dispatch/combine are memory_bound rows inside moe; informational"
+  }
+}
diff --git a/examples/deepseek-v4/projection/site/index.html b/examples/deepseek-v4/projection/site/index.html
new file mode 100644
index 000000000..47c29c364
--- /dev/null
+++ b/examples/deepseek-v4/projection/site/index.html
@@ -0,0 +1,134 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>DeepSeek-V4 Performance Projection</title>
+    <link rel="stylesheet" href="./assets/style.css" />
+  </head>
+  <body>
+    <header class="hero">
+      <div class="hero__inner">
+        <p class="eyebrow">Primus · DeepSeek-V4</p>
+        <h1>Training Performance Projection</h1>
+        <p class="hero__summary">
+          Trace-driven single-layer breakdown on MI355X, scaled to a full-model,
+          multi-GPU projection. Page 1 is measured-MI355X; page 2 scales to MI455X
+          by theoretical hardware ratios.
+        </p>
+      </div>
+    </header>
+
+    <div class="switcher" id="switcher">
+      <div class="switcher__inner">
+        <label class="field field--inline">
+          <span>Model</span>
+          <select id="model-select">
+            <option value="pro">DeepSeek-V4 Pro</option>
+            <option value="flash">DeepSeek-V4 Flash</option>
+          </select>
+        </label>
+        <span id="mock-badge" class="badge badge--warn" hidden>MOCK DATA</span>
+        <nav class="tabs" role="tablist">
+          <button class="tab is-active" data-gpu="MI355X" role="tab">MI355X (measured)</button>
+          <button class="tab" data-gpu="MI455X" role="tab">MI455X (projected)</button>
+        </nav>
+      </div>
+    </div>
+
+    <main class="layout layout--split">
+      <aside class="sidebar">
+        <section class="panel" id="controls-panel">
+          <h2>Projection controls</h2>
+          <div class="mode-switch" id="mode-switch">
+            <span class="mode-switch__label">Layer timing</span>
+            <div class="mode-switch__tabs" role="tablist">
+              <button class="mode-tab is-active" data-mode="trace" role="tab">Trace-derived</button>
+              <button class="mode-tab" data-mode="manual" role="tab">Manual per-layer</button>
+            </div>
+          </div>
+          <div id="manual-grid" class="manual-grid" hidden></div>
+          <div id="controls-grid" class="controls-grid"></div>
+          <div id="validation-panel" hidden></div>
+          <p class="muted">
+            DP is derived from World size / (PP·TP·CP); EP is a sub-group of DP.
+            Note: the captured MoE dispatch/combine is intra-node EP=8 and is not
+            re-modeled for other EP values; sequence length is fixed at the
+            capture value (4096).
+          </p>
+        </section>
+      </aside>
+
+      <div class="content">
+      <div id="error-state" class="state state--error" hidden></div>
+
+      <section class="panel" id="config-panel">
+        <h2>Model configuration</h2>
+        <div id="config-grid" class="kv-grid"></div>
+        <div id="cr-schedule" class="cr-schedule"></div>
+      </section>
+
+      <section class="panel" id="breakdown-panel">
+        <div class="panel__head">
+          <h2>Per-layer breakdown <span id="breakdown-gpu" class="muted"></span></h2>
+          <p class="muted">
+            Forward (left→right) then backward (right→left). Time in µs / one
+            microbatch (seq 4096). The TFLOP/s row is per-kernel <em>achieved</em>
+            (gemm / grouped_gemm / attn only); the headline TFLOP/s/GPU below uses
+            the V4 analytic model FLOPs instead.
+          </p>
+          <p id="breakdown-manual-note" class="muted manual-note" hidden>
+            Manual layer-timing mode is active — the projected iteration time is
+            driven by the values you set in “Manual per-layer”; the per-module
+            tables below are shown for reference only.
+          </p>
+        </div>
+        <div id="breakdown-blocks"></div>
+      </section>
+
+      <section class="panel panel--accent" id="results-panel">
+        <h2>Projected throughput <span id="results-gpu" class="muted"></span></h2>
+        <div id="results-headline" class="headline"></div>
+        <div id="results-steps" class="steps"></div>
+      </section>
+
+      <section class="panel" id="timeline-panel">
+        <div class="panel__head">
+          <h2>Iteration timeline <span id="timeline-gpu" class="muted"></span></h2>
+          <p class="muted">
+            How one iteration's time is composed, bottom-up: a single layer
+            (attn / mlp / a2a) → each pipeline rank's chunks (layer granularity) →
+            the whole 1F1B pipeline schedule. Reacts live to the controls on the left.
+            See <code>design/07-iteration-timeline.md</code>.
+          </p>
+          <div class="tl-head-row">
+            <div class="tl-viewswitch" id="tl-viewswitch">
+              <span class="mode-switch__label">Layout</span>
+              <div class="mode-switch__tabs" role="tablist">
+                <button class="tl-view" data-view="tabs" role="tab">Tabbed</button>
+                <button class="tl-view is-active" data-view="stacked" role="tab">Stacked (all + link)</button>
+              </div>
+            </div>
+            <div class="tl-tabs" id="tl-tabs" role="tablist">
+              <button class="tl-tab is-active" data-level="1" role="tab">L1 · Layer</button>
+              <button class="tl-tab" data-level="2" role="tab">L2 · PP ranks</button>
+              <button class="tl-tab" data-level="3" role="tab">L3 · Schedule</button>
+            </div>
+          </div>
+        </div>
+        <div id="tl-body"></div>
+      </section>
+      </div>
+    </main>
+
+    <footer class="site-footer">
+      <p>
+        Generated from <code>examples/deepseek-v4/projection/</code>. Methodology &amp;
+        assumptions in <code>design/</code>. Numbers are only as good as the input
+        trace; see the assumptions list before citing.
+      </p>
+    </footer>
+
+    <script src="./assets/app.js"></script>
+  </body>
+</html>
diff --git a/examples/deepseek-v4/projection/tools/gen_mock_data.py b/examples/deepseek-v4/projection/tools/gen_mock_data.py
new file mode 100755
index 000000000..1b223a619
--- /dev/null
+++ b/examples/deepseek-v4/projection/tools/gen_mock_data.py
@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""Generate MOCK breakdown JSON for site development / demo.
+
+Numbers are placeholders in the right *order of magnitude*, seeded from the
+published P57 single-layer attention micro-bench (V4-Flash widths: B=1, H=64,
+Sq=4096, D=512) and the P40 EP=8 MoE/kernel attribution. They are NOT measured
+ground truth — replace with `parse_trace.py` output once real traces exist
+(every file is marked provenance.mock = true).
+
+Usage:
+    python3 examples/deepseek-v4/projection/tools/gen_mock_data.py
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+
+OUT_DIR = Path(__file__).resolve().parent.parent / "site" / "data"
+
+PRO_COMPRESS = [128, 128] + [4 if i % 2 == 0 else 128 for i in range(2, 60)] + [0]
+FLASH_COMPRESS = [0, 0] + [4 if i % 2 == 0 else 128 for i in range(2, 42)] + [0]
+
+MODEL_CONFIGS = {
+    "pro": {
+        "num_layers": 61,
+        "hidden_size": 7168,
+        "num_attention_heads": 128,
+        "kv_channels": 512,
+        "num_experts": 384,
+        "moe_router_topk": 6,
+        "moe_ffn_hidden_size": 3072,
+        "moe_shared_expert_intermediate_size": 3072,
+        "index_topk": 1024,
+        "vocab_size": 129280,
+        "compress_ratios": PRO_COMPRESS,
+    },
+    "flash": {
+        "num_layers": 43,
+        "hidden_size": 4096,
+        "num_attention_heads": 64,
+        "kv_channels": 512,
+        "num_experts": 256,
+        "moe_router_topk": 6,
+        "moe_ffn_hidden_size": 2048,
+        "moe_shared_expert_intermediate_size": 2048,
+        "index_topk": 512,
+        "vocab_size": 129280,
+        "compress_ratios": FLASH_COMPRESS,
+    },
+}
+
+# MI355X: BF16 matrix 2.5 PFLOPS, HBM3E 8 TB/s (AMD product page).
+# MI455X (MI400): HBM4 19.6 TB/s; BF16 dense not officially published —
+# estimated ~10 PFLOPS (half of the 20 PFLOPS FP8 spec).
+HARDWARE = {
+    "MI355X": {"peak_tflops_bf16": 2500.0, "hbm_bandwidth_gbps": 8000.0},
+    "MI455X": {"peak_tflops_bf16": 10000.0, "hbm_bandwidth_gbps": 19600.0},
+}
+
+
+def row(module, time_us, flop_class=None, tflops=None):
+    flops = (tflops * time_us * 1e6) if (flop_class and tflops) else None
+    return {
+        "module": module,
+        "time_us": round(time_us, 1),
+        "class": "compute_bound" if flop_class else "memory_bound",
+        "flop_class": flop_class,
+        "flops": flops,
+        "tflops": tflops,
+        "kernels": [],
+    }
+
+
+# Base (flash) attention core fwd/bwd by cr, from P57 micro-bench (ms -> us).
+ATTN_CORE = {
+    "0": {"fwd": 500.0, "bwd": 2080.0},
+    "4": {"fwd": 1430.0, "bwd": 5110.0},
+    "128": {"fwd": 570.0, "bwd": 2810.0},
+}
+
+
+def attention_breakdown(cr, s):
+    """s = linear scale factor vs flash widths."""
+    fwd = [
+        row("attn.qkv_proj", 220 * s, "gemm", 480),
+        row("attn.core", ATTN_CORE[cr]["fwd"] * s, "attn", 210),
+        row("attn.rope", 35 * s),
+        row("attn.o_proj", 160 * s, "gemm", 470),
+        row("attn.norm", 25 * s),
+    ]
+    bwd = [
+        row("attn.qkv_proj", 440 * s, "gemm", 480),
+        row("attn.core", ATTN_CORE[cr]["bwd"] * s, "attn", 180),
+        row("attn.rope", 45 * s),
+        row("attn.o_proj", 320 * s, "gemm", 470),
+        row("attn.norm", 30 * s),
+    ]
+    if cr == "4":  # CSA uses the Indexer/Compressor
+        fwd.insert(2, row("attn.indexer", 300 * s, "gemm", 300))
+        bwd.insert(2, row("attn.indexer", 600 * s, "gemm", 300))
+    return {"forward": fwd, "backward": bwd}
+
+
+def moe_breakdown(sg, sc):
+    """sg = grouped-gemm scale, sc = comm/act scale vs flash."""
+    fwd = [
+        row("moe.router", 55 * sc),
+        row("moe.dispatch", 820 * sc),
+        row("moe.grouped_gemm", 1850 * sg, "grouped_gemm", 430),
+        row("moe.act", 110 * sc),
+        row("moe.shared_expert", 210 * sg, "gemm", 460),
+        row("moe.combine", 990 * sc),
+    ]
+    bwd = [
+        row("moe.router", 75 * sc),
+        row("moe.dispatch", 900 * sc),
+        row("moe.grouped_gemm", 3700 * sg, "grouped_gemm", 430),
+        row("moe.act", 150 * sc),
+        row("moe.shared_expert", 420 * sg, "gemm", 460),
+        row("moe.combine", 1050 * sc),
+    ]
+    return {"forward": fwd, "backward": bwd}
+
+
+def non_layer(s):
+    return {
+        "embedding": {"forward": [row("embedding", 120 * s)], "backward": [row("embedding", 60 * s)]},
+        "output": {
+            "forward": [row("output", 900 * s, "gemm", 300)],
+            "backward": [row("output", 1800 * s, "gemm", 300)],
+        },
+        "loss": {"forward": [row("loss", 80)], "backward": [row("loss", 60)]},
+    }
+
+
+def cr_counts(compress):
+    c = {"0": 0, "4": 0, "128": 0}
+    for x in compress:
+        c[str(x)] += 1
+    return c
+
+
+def build(model):
+    cfg = MODEL_CONFIGS[model]
+    # scale factors vs flash baseline
+    if model == "pro":
+        s_attn, s_grouped, s_comm, s_out = 1.75, 2.0, 1.3, 1.75
+        opt_params, opt_time = 2_300_000_000, 1600.0
+    else:
+        s_attn, s_grouped, s_comm, s_out = 1.0, 1.0, 1.0, 1.0
+        opt_params, opt_time = 1_000_000_000, 850.0
+
+    layers = {
+        cr: {"attention": attention_breakdown(cr, s_attn), "moe": moe_breakdown(s_grouped, s_comm)}
+        for cr in ("0", "4", "128")
+    }
+    return {
+        "schema_version": 1,
+        "model": model,
+        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "provenance": {
+            "mock": True,
+            "note": "MOCK data (P57/P40 order-of-magnitude); replace with parse_trace.py output",
+        },
+        "capture": {
+            "gpu": "MI355X",
+            "seq_length": 4096,
+            "micro_batch_size": 1,
+            "tokens_per_microbatch": 4096,
+            "ep": 8,
+            "ga_for_capture": 2,
+            "optimizer": "adam",
+            "distributed_optimizer": True,
+            "recompute": "off",
+            "measured_iter_time_ms": None,
+        },
+        "model_config": {**cfg, "cr_layer_counts": cr_counts(cfg["compress_ratios"])},
+        "hardware": HARDWARE,
+        "layers": layers,
+        "non_layer": non_layer(s_out),
+        "optimizer": {
+            "type": "adam",
+            "measured_params": opt_params,
+            "time_us": opt_time,
+            "bytes_per_param": 18,
+            "class": "memory_bound",
+        },
+        "comm": {
+            "ep_dispatch_us": None,
+            "ep_combine_us": None,
+            "note": "EP dispatch/combine included in moe rows; informational",
+        },
+    }
+
+
+def main():
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    for model in ("pro", "flash"):
+        doc = build(model)
+        path = OUT_DIR / f"{model}.json"
+        path.write_text(json.dumps(doc, indent=2))
+        print(f"[gen_mock_data] wrote {path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/deepseek-v4/projection/tools/kernel_module_map.py b/examples/deepseek-v4/projection/tools/kernel_module_map.py
new file mode 100644
index 000000000..fd6ddbcea
--- /dev/null
+++ b/examples/deepseek-v4/projection/tools/kernel_module_map.py
@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+"""Kernel / nn.module -> logical-module + flop-class mapping for the V4
+projection breakdown.
+
+Two independent classifications are provided:
+
+1. ``module_from_stack(stack)`` — primary: derive the logical module from the
+   python call stack captured by ``with_stack=True`` (matches nn.module class
+   names appearing in the stack frames). This is the accurate path (A13).
+
+2. ``module_from_kernel(name)`` — fallback: derive the logical module purely from
+   the GPU kernel name when no usable stack is attached.
+
+``flop_class_from_kernel(name)`` returns the compute-bound FLOP class
+(``gemm`` / ``grouped_gemm`` / ``attn``) or ``None`` (memory-bound, A14).
+
+The rules are intentionally data-driven so they can be extended as kernels are
+renamed. Order matters: the first matching rule wins.
+"""
+
+from __future__ import annotations
+
+# --- logical module taxonomy (kept small per the design) --------------------
+# attention sub-modules:
+#   attn.qkv_proj, attn.rope, attn.indexer, attn.core, attn.o_proj, attn.norm
+# moe sub-modules:
+#   moe.router, moe.dispatch, moe.grouped_gemm, moe.act, moe.combine,
+#   moe.shared_expert
+# non-layer: embedding, output, loss
+
+# (substring, logical_module) — matched against the python call stack text.
+# nn.module class names / function names seen in V4 forward stacks.
+STACK_MODULE_RULES: list[tuple[str, str]] = [
+    ("Indexer", "attn.indexer"),
+    ("Compressor", "attn.indexer"),
+    ("apply_rotary", "attn.rope"),
+    ("rope", "attn.rope"),
+    ("linear_qkv", "attn.qkv_proj"),
+    ("q_layernorm", "attn.norm"),
+    ("k_layernorm", "attn.norm"),
+    ("linear_proj", "attn.o_proj"),
+    ("o_proj", "attn.o_proj"),
+    ("core_attention", "attn.core"),
+    ("DeepseekV4Attention", "attn.core"),
+    ("MLASelfAttention", "attn.core"),
+    ("SelfAttention", "attn.core"),
+    ("input_layernorm", "attn.norm"),
+    ("pre_mlp_layernorm", "moe.router"),
+    ("TopKRouter", "moe.router"),
+    ("router", "moe.router"),
+    ("sinkhorn", "moe.router"),
+    ("shared_expert", "moe.shared_expert"),
+    ("token_dispatch", "moe.dispatch"),
+    ("dispatch", "moe.dispatch"),
+    ("combine", "moe.combine"),
+    ("GroupedMLP", "moe.grouped_gemm"),
+    ("SequentialMLP", "moe.grouped_gemm"),
+    ("grouped", "moe.grouped_gemm"),
+    ("experts", "moe.grouped_gemm"),
+    ("activation", "moe.act"),
+    ("swiglu", "moe.act"),
+    ("MoELayer", "moe.grouped_gemm"),
+    ("word_embeddings", "embedding"),
+    ("embedding", "embedding"),
+    ("output_layer", "output"),
+    ("lm_head", "output"),
+    ("loss", "loss"),
+    ("cross_entropy", "loss"),
+]
+
+# (substring, logical_module) — matched against the GPU kernel name (fallback).
+KERNEL_MODULE_RULES: list[tuple[str, str]] = [
+    ("_v4_csa_attention", "attn.core"),
+    ("_v4_attention", "attn.core"),
+    ("_hc_compute", "attn.core"),
+    ("hc_compute", "attn.core"),
+    ("_indexer", "attn.indexer"),
+    ("indexer", "attn.indexer"),
+    ("_compressor", "attn.indexer"),
+    ("compressor", "attn.indexer"),
+    ("apply_rope", "attn.rope"),
+    ("rotary", "attn.rope"),
+    ("rope", "attn.rope"),
+    ("_sinkhorn", "moe.router"),
+    ("sinkhorn", "moe.router"),
+    ("_v4_router", "moe.router"),
+    ("deep_ep::", "moe.dispatch"),  # refined to dispatch/combine below
+    ("dispatch", "moe.dispatch"),
+    ("combine", "moe.combine"),
+    ("GroupedGemm", "moe.grouped_gemm"),
+    ("_stack_grouped_weight", "moe.grouped_gemm"),
+    ("group_gemm", "moe.grouped_gemm"),
+    ("swiglu", "moe.act"),
+    ("embedding", "embedding"),
+    ("cross_entropy", "loss"),
+    ("nll_loss", "loss"),
+]
+
+# (substring, flop_class) — matched against the GPU kernel name.
+# Order matters: grouped GEMM must be checked before generic GEMM.
+FLOP_CLASS_RULES: list[tuple[str, str]] = [
+    ("GroupedGemmKernel", "grouped_gemm"),
+    ("grouped_gemm", "grouped_gemm"),
+    ("group_gemm", "grouped_gemm"),
+    ("_v4_csa_attention", "attn"),
+    ("_v4_attention", "attn"),
+    ("attention_fwd", "attn"),
+    ("attention_bwd", "attn"),
+    # generic dense GEMM kernels (hipBLASLt / rocBLAS / CK tile / Triton matmul)
+    ("GemmKernel", "gemm"),
+    ("Cijk_", "gemm"),
+    ("gemm", "gemm"),
+    ("matmul", "gemm"),
+]
+
+
+def module_from_stack(stack: str | None) -> str | None:
+    """Return the logical module from a python call-stack string, or None."""
+    if not stack:
+        return None
+    for needle, module in STACK_MODULE_RULES:
+        if needle in stack:
+            return module
+    return None
+
+
+def module_from_kernel(name: str) -> str:
+    """Return the logical module from a kernel name (fallback)."""
+    lowered = name
+    for needle, module in KERNEL_MODULE_RULES:
+        if needle in lowered:
+            if module == "moe.dispatch" and "combine" in lowered:
+                return "moe.combine"
+            return module
+    return "other"
+
+
+def flop_class_from_kernel(name: str) -> str | None:
+    """Return 'gemm' | 'grouped_gemm' | 'attn', or None for memory-bound."""
+    for needle, klass in FLOP_CLASS_RULES:
+        if needle in name:
+            return klass
+    return None
+
+
+def is_compute_bound(name: str) -> bool:
+    return flop_class_from_kernel(name) is not None
diff --git a/examples/deepseek-v4/projection/tools/parse_trace.py b/examples/deepseek-v4/projection/tools/parse_trace.py
new file mode 100755
index 000000000..dad6291f7
--- /dev/null
+++ b/examples/deepseek-v4/projection/tools/parse_trace.py
@@ -0,0 +1,651 @@
+#!/usr/bin/env python3
+"""Turn per-cr chrome traces into a projection breakdown JSON.
+
+Input: one rank-0 PyTorch/Kineto chrome trace per compression-ratio (cr), each
+captured by ``script/deepseek_v4_layer_trace-projection.sh`` (1 layer, seq 4096,
+GA=2, recompute off, overlap off, profiler window iter 6->7).
+
+Output: a single ``<model>.json`` matching ``design/03-json-schema.md``.
+
+Attribution (validated against the real ROCm/Kineto trace):
+  * GPU kernels (cat=="kernel") link to their launching CPU op via the shared
+    ``External id`` arg. Optimizer (``multi_tensor_apply``) and DP-comm
+    (``nccl``) kernels carry no External id and are classified by name.
+  * The module comes from the enclosing ``nn.Module: <Class>_n`` python_function
+    events (with_stack) on the CPU op's thread; fwd/bwd from "Backward"/
+    "autograd" in the CPU op name or an enclosing frame.
+  * Clean per-call time = ``min`` over launches grouped by
+    ``(phase, module, kernel, input-dims)`` (overlap is off, so this just
+    removes warm-up/jitter); calls-per-microbatch is treated as 1 for a single
+    captured layer.
+  * Compute-bound FLOP class from the kernel name; GEMM FLOPs from input dims.
+"""
+
+from __future__ import annotations
+
+import argparse
+import bisect
+import json
+from collections import defaultdict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+# Merge gap (us) used when reconstructing the backward GPU-time window from the
+# linked backward kernels, so an unlinked kernel sitting in a small stall
+# between two backward kernels is still counted as backward.
+_BWD_WINDOW_GAP_US = 150.0
+
+# A single layer-compute GPU kernel launch at seq=4096/B=1 realistically tops
+# out around ~10 ms (e.g. the V4 attention bwd or a grouped GEMM). Some captures
+# bill a one-off device-side stall / busy-wait to an ordinary elementwise kernel
+# (observed: 3 launches of ~145 ms attributed to aten::mul/add_/div in the cr=0
+# pro trace, 17x the cr=4/128 value). Those are capture artifacts, not per-layer
+# cost, so layer-compute launches above this cap are dropped (and reported in
+# provenance). Optimizer/comm kernels are routed out before this cap applies.
+_MAX_PLAUSIBLE_LAUNCH_US = 50_000.0
+
+from kernel_module_map import flop_class_from_kernel, module_from_kernel
+from v4_flops import FB_FMA, layer_fmac, model_total_params, mtp_flops, nonlayer_fmac
+
+PRO_COMPRESS = [128, 128] + [4 if i % 2 == 0 else 128 for i in range(2, 60)] + [0]
+FLASH_COMPRESS = [0, 0] + [4 if i % 2 == 0 else 128 for i in range(2, 42)] + [0]
+
+MODEL_CONFIGS: dict[str, dict[str, Any]] = {
+    "pro": {
+        "num_layers": 61,
+        "hidden_size": 7168,
+        "num_attention_heads": 128,
+        "kv_channels": 512,
+        "num_experts": 384,
+        "moe_router_topk": 6,
+        "moe_ffn_hidden_size": 3072,
+        "moe_shared_expert_intermediate_size": 3072,
+        "index_topk": 1024,
+        "vocab_size": 129280,
+        "mtp_num_layers": 1,
+        "mtp_compress_ratios": [4],
+        "pipeline_layout": "",
+        "compress_ratios": PRO_COMPRESS,
+    },
+    "flash": {
+        "num_layers": 43,
+        "hidden_size": 4096,
+        "num_attention_heads": 64,
+        "kv_channels": 512,
+        "num_experts": 256,
+        "moe_router_topk": 6,
+        "moe_ffn_hidden_size": 2048,
+        "moe_shared_expert_intermediate_size": 2048,
+        "index_topk": 512,
+        "vocab_size": 129280,
+        "mtp_num_layers": 1,
+        "mtp_compress_ratios": [4],
+        "pipeline_layout": "Et*10|t*11|t*11|t*11mL",
+        "compress_ratios": FLASH_COMPRESS,
+    },
+}
+
+# MI355X: BF16 matrix 2.5 PFLOPS, HBM3E 8 TB/s (AMD product page). MI455X
+# (MI400): HBM4 19.6 TB/s; BF16 dense not officially published — estimated
+# ~10 PFLOPS (half of the 20 PFLOPS FP8 spec). The site can override these.
+DEFAULT_HARDWARE = {
+    "MI355X": {"peak_tflops_bf16": 2500.0, "hbm_bandwidth_gbps": 8000.0},
+    "MI455X": {"peak_tflops_bf16": 10000.0, "hbm_bandwidth_gbps": 19600.0},
+}
+
+# Measured multi-node anchors used to set/validate the site's calibFactor. The
+# projection, configured to the anchor's parallel layout, must reproduce these
+# (see design/06-calibration.md). flash: real 8-node MI355X run (image pr-768,
+# commit 2c9b). pro: production multi-node anchor still TODO (needs >=8 idle
+# nodes); until then pro reuses flash's cross-model calibFactor.
+MODEL_ANCHORS: dict[str, dict[str, Any]] = {
+    "flash": {
+        "measured_iter_time_ms": 4076.0,
+        "measured_anchor": {
+            "config": "PP8/VPP1/EP8/TP1, world 64, MBS1/GBS128, seq4096, full recompute, "
+            "layout Et*4|t*5|(t*6|)*5,t*4mL (8-node MI355X)",
+            "iter_ms": 4076.0,
+            "tflops_gpu": 722,
+            "tok_s_gpu": 2010,
+            "calib_factor": 0.87,
+            "tolerance": "<0.1%",
+        },
+    },
+}
+
+ALL_MODULES = (
+    "attn.proj",
+    "attn.core",
+    "attn.indexer",
+    "attn.norm",
+    "moe.router",
+    "moe.dispatch",
+    "moe.grouped_gemm",
+    "moe.shared_expert",
+    "moe.combine",
+    "embedding",
+    "output",
+    "loss",
+    "other",
+)
+
+
+def _arg(ev: dict, *keys: str) -> Any:
+    args = ev.get("args") or {}
+    for k in keys:
+        if k in args:
+            return args[k]
+    return None
+
+
+def _dims_key(dims: Any) -> str:
+    if dims is None:
+        return ""
+    try:
+        return json.dumps(dims, separators=(",", ":"))
+    except TypeError:
+        return str(dims)
+
+
+def _merge_intervals(intervals: list[tuple[float, float]], gap: float = 0.0) -> list[tuple[float, float]]:
+    """Merge [start, end] intervals; bridge neighbours separated by <= gap."""
+    out: list[tuple[float, float]] = []
+    for s, e in sorted(intervals):
+        if out and s <= out[-1][1] + gap:
+            out[-1] = (out[-1][0], max(out[-1][1], e))
+        else:
+            out.append((s, e))
+    return out
+
+
+def _make_membership(merged: list[tuple[float, float]]):
+    """Return a fn(ts)->bool: is ts inside one of the merged intervals."""
+    starts = [s for s, _ in merged]
+
+    def _in(ts: float | None) -> bool:
+        if ts is None or not merged:
+            return False
+        i = bisect.bisect_right(starts, ts) - 1
+        return i >= 0 and merged[i][0] <= ts <= merged[i][1]
+
+    return _in
+
+
+def _is_opt_or_comm_kernel(name: str) -> bool:
+    low = name.lower()
+    return "multi_tensor_apply" in name or "fusedadam" in low or "adamw" in low or "nccl" in low
+
+
+def _gemm_flops(dims_key: str) -> float | None:
+    if not dims_key:
+        return None
+    try:
+        dims = json.loads(dims_key)
+    except json.JSONDecodeError:
+        return None
+    shapes = [d for d in dims if isinstance(d, list) and len(d) >= 2 and all(isinstance(x, int) for x in d)]
+    if len(shapes) >= 2:
+        a, b = shapes[0], shapes[1]
+        m, k = a[-2], a[-1]
+        k2, n = b[-2], b[-1]
+        if k == k2:
+            batch = 1
+            for x in a[:-2]:
+                batch *= x
+            return 2.0 * batch * m * n * k
+    return None
+
+
+def _resolve_module(kname: str, cpu_name: str, anc_classes: set[str], flop_class: str | None) -> str:
+    """Logical module for a kernel given its name, CPU op name, and enclosing
+    nn.Module classes. Returns '__optimizer__' / '__dpcomm__' for non-layer
+    buckets handled separately. Priority: kernel name > cpu-op name > enclosing
+    nn.Module (forward only; backward ops aren't inside module forward ranges)."""
+    n = kname
+    low = n.lower()
+    cn = cpu_name or ""
+
+    def has(*xs: str) -> bool:
+        return any(any(x in a for a in anc_classes) for x in xs)
+
+    def coarse(module: str) -> str:
+        if module in ("attn.qkv_proj", "attn.o_proj"):
+            return "attn.proj"
+        if module in ("attn.rope",):
+            return "attn.norm"
+        if module == "moe.act":
+            return "moe.grouped_gemm"
+        return module
+
+    # 1) kernel-name rules (most reliable; present for fwd and bwd)
+    if "multi_tensor_apply" in n or "fusedadam" in low or "adamw" in low:
+        return "__optimizer__"
+    if "nccl" in low:
+        return "__dpcomm__"
+    if "cross_entropy" in low or "online_softmax" in low:
+        return "loss"
+    if "deep_ep" in n or "deepep" in low:
+        return "moe.combine" if "combine" in low else "moe.dispatch"
+    if "_v4_csa" in n or "_v4_attention" in n or "_hc_" in n:
+        return "attn.core"
+    if "_sinkhorn" in n or "_v4_router" in n:
+        return "moe.router"
+    if (
+        "GroupedGemm" in n
+        or "_grouped" in n
+        or "group_gemm" in low
+        or "grouped_gemm" in low
+        or "grouped_variable" in n
+    ):
+        return "moe.grouped_gemm"
+
+    # 2) cpu-op (autograd Function / aten) name rules — needed for backward
+    if "Attention" in cn or "CSAPool" in cn or "MLA" in cn:
+        return "attn.core"
+    if "Indexer" in cn or "Compressor" in cn:
+        return "attn.indexer"
+    if "Sinkhorn" in cn or "Router" in cn:
+        return "moe.router"
+    if "RMSNorm" in cn or "LayerNorm" in cn or "layer_norm" in cn.lower():
+        return "attn.norm"
+    if "crossentropy" in cn.lower() or "cross_entropy" in cn.lower() or "nll_loss" in cn.lower():
+        return "loss"
+    if "embedding" in cn.lower():
+        return "output" if flop_class == "gemm" else "embedding"
+    if "LinearWithGradAccumulation" in cn or cn in ("aten::mm", "aten::addmm", "aten::matmul", "aten::bmm"):
+        if has("Embedding") or (
+            has("DeepseekV4Model")
+            and not has(
+                "DeepseekV4Attention", "DeepseekV4HybridLayer", "Compressor", "Indexer", "MLP", "Expert"
+            )
+        ):
+            return "output"
+        if has("Compressor", "Indexer"):
+            return "attn.indexer"
+        if has("MLP", "Expert"):
+            return "moe.grouped_gemm"
+        return "attn.proj"
+
+    # 3) enclosing nn.Module (forward only)
+    if has("Compressor", "Indexer"):
+        return "attn.indexer"
+    if has("SharedExpert"):
+        return "moe.shared_expert"
+    if has("GroupedMLP", "SequentialMLP", "GroupedExperts", "Experts"):
+        return "moe.grouped_gemm"
+    if has("Router"):
+        return "moe.router"
+    if has("DeepseekV4Attention", "MLASelfAttention", "SelfAttention"):
+        return "attn.proj" if flop_class == "gemm" else "attn.norm"
+    if has("Embedding"):
+        return "output" if flop_class == "gemm" else "embedding"
+    fallback = coarse(module_from_kernel(kname))
+    return fallback if fallback != "other" else "attn.misc"
+
+
+def parse_trace(path: Path, ga: int = 2):
+    payload = json.loads(path.read_text())
+    events = payload.get("traceEvents", [])
+
+    # index cpu ops by External id
+    cpu_by_extid: dict[Any, dict] = {}
+    # interesting python_function intervals per tid: (ts, end, name)
+    pf_by_tid: dict[Any, list[tuple[float, float, str]]] = defaultdict(list)
+    # `autograd::engine::evaluate_function` CPU intervals precisely bracket the
+    # backward pass; a kernel whose launching CPU op falls inside one is bwd.
+    eval_intervals: list[tuple[float, float]] = []
+    kernels: list[dict] = []
+    num_steps = 0  # real captured training steps (ProfilerStep events, dur>10ms)
+
+    for ev in events:
+        cat = (ev.get("cat") or "").lower()
+        ph = ev.get("ph")
+        if ph == "X" and ev.get("name", "").startswith("ProfilerStep") and (ev.get("dur") or 0) > 10000:
+            num_steps += 1
+        if cat == "cpu_op" and ph == "X":
+            ext = _arg(ev, "External id")
+            if ext is not None and ext not in cpu_by_extid:
+                cpu_by_extid[ext] = ev
+            if ev.get("name", "").startswith("autograd::engine::evaluate_function"):
+                ets = ev.get("ts")
+                if ets is not None:
+                    eval_intervals.append((ets, ets + (ev.get("dur") or 0)))
+        elif cat == "python_function" and ph == "X":
+            name = ev.get("name", "")
+            if name.startswith("nn.Module:") or "ackward" in name or "autograd" in name:
+                ts = ev.get("ts")
+                dur = ev.get("dur") or 0
+                if ts is not None:
+                    pf_by_tid[ev.get("tid")].append((ts, ts + dur, name))
+        elif cat == "kernel" and ph == "X" and ev.get("dur") is not None:
+            kernels.append(ev)
+
+    for tid in pf_by_tid:
+        pf_by_tid[tid].sort(key=lambda t: t[0])
+
+    def enclosing(cpu: dict) -> set[str]:
+        """Return enclosing nn.Module class names for a cpu op (forward only;
+        backward ops live under the autograd engine, not module forward ranges)."""
+        tid = cpu.get("tid")
+        ts = cpu.get("ts")
+        end = ts + (cpu.get("dur") or 0)
+        classes: set[str] = set()
+        for pts, pend, name in pf_by_tid.get(tid, ()):
+            if pts > ts:
+                break
+            if pend >= end and name.startswith("nn.Module:"):
+                classes.add(name.split(":", 1)[1].strip().rsplit("_", 1)[0])
+        return classes
+
+    # Per-microbatch time = (sum of all launches over the whole profiler window)
+    # / num_mb, where num_mb = num_steps * GA. We keep the SUM (not a min over
+    # launches): the single-layer capture serializes each layer's work, and the
+    # measured flash-16L anchor (design/06) shows this serial per-layer time is
+    # the right estimate (calibFactor ~0.93 against 6665 ms). A min-over-launches
+    # rule would assume the scalar control-flow stalls (e.g. the Indexer top-k
+    # device syncs) are fully hidden in the full model; the anchor refutes that
+    # (it would need calibFactor ~1.3), so they are kept as real per-layer cost.
+    #
+    # Subgroups are keyed by (kernel, input-dims) only so we can (a) classify a
+    # row as compute_bound iff FLOP-classed kernels dominate its time (>50%) --
+    # a stray flop-classed kernel can no longer flip a memory-bound aggregate --
+    # and (b) sum dim-derived GEMM FLOPs correctly.
+    num_mb = max(1, num_steps * ga)
+
+    # ---- phase classification (forward vs backward) ----------------------
+    # The backward pass is identified by the autograd engine: a kernel is
+    # backward iff its launching CPU op was issued inside an
+    # `autograd::engine::evaluate_function` interval (these precisely bracket
+    # the backward). A `_fwd_`/`_bwd_` tag in the kernel name (V4 Triton kernels
+    # encode it) wins. Unlinked kernels (no External id -> no CPU op; common for
+    # fused/elementwise launches Kineto fails to flow-link) are assigned by
+    # whether their GPU timestamp lands inside the backward GPU-time window
+    # rebuilt from the linked backward kernels. This replaces the old "default
+    # to forward" rule, which systematically leaked backward compute -- incl.
+    # the MoE dgrad/wgrad grouped GEMMs -- into the forward breakdown.
+    bwd_eval = _merge_intervals(eval_intervals)
+    in_bwd_eval = _make_membership(bwd_eval)
+
+    phase_by_idx: list[str | None] = [None] * len(kernels)
+    _bwd_windows: list[tuple[float, float]] = []
+    for i, ev in enumerate(kernels):
+        name = ev.get("name", "")
+        ln = name.lower()
+        if "_fwd" in ln and "_bwd" not in ln:
+            phase_by_idx[i] = "forward"
+            continue
+        if "_bwd" in ln:
+            phase_by_idx[i] = "backward"
+            if not _is_opt_or_comm_kernel(name) and ev["dur"] <= _MAX_PLAUSIBLE_LAUNCH_US:
+                _bwd_windows.append((ev["ts"], ev["ts"] + ev["dur"]))
+            continue
+        ext = _arg(ev, "External id")
+        cpu = cpu_by_extid.get(ext) if ext is not None else None
+        if cpu is not None:
+            ph = "backward" if in_bwd_eval(cpu.get("ts")) else "forward"
+            phase_by_idx[i] = ph
+            if (
+                ph == "backward"
+                and not _is_opt_or_comm_kernel(name)
+                and ev["dur"] <= _MAX_PLAUSIBLE_LAUNCH_US
+            ):
+                _bwd_windows.append((ev["ts"], ev["ts"] + ev["dur"]))
+        # else: unlinked -> decided below from the GPU-side backward window
+
+    in_bwd_gpu = _make_membership(_merge_intervals(_bwd_windows, _BWD_WINDOW_GAP_US))
+    for i, ev in enumerate(kernels):
+        if phase_by_idx[i] is None:
+            phase_by_idx[i] = "backward" if in_bwd_gpu(ev.get("ts")) else "forward"
+
+    # (phase, module) -> { (kernel_name, dims_key): {"durs": [...], "flop_class", "flop_per_launch"} }
+    groups: dict[tuple[str, str], dict] = defaultdict(dict)
+    optimizer_us = 0.0
+    dpcomm_us = 0.0
+    dropped_stall_us = 0.0
+
+    for i, ev in enumerate(kernels):
+        name = ev.get("name", "")
+        dur = float(ev["dur"])
+        ext = _arg(ev, "External id")
+        cpu = cpu_by_extid.get(ext) if ext is not None else None
+        cpu_name = cpu.get("name", "") if cpu else ""
+        anc = enclosing(cpu) if cpu else set()
+        flop_class = flop_class_from_kernel(name)
+        module = _resolve_module(name, cpu_name, anc, flop_class)
+        if module == "__optimizer__":
+            optimizer_us += dur
+            continue
+        if module == "__dpcomm__":
+            dpcomm_us += dur
+            continue
+        # Drop one-off device-side stalls billed to a layer-compute kernel (see
+        # _MAX_PLAUSIBLE_LAUNCH_US); they are capture artifacts, not per-layer cost.
+        if dur > _MAX_PLAUSIBLE_LAUNCH_US:
+            dropped_stall_us += dur
+            continue
+        phase = phase_by_idx[i]
+        dims_key = _dims_key(_arg(cpu, "Input Dims")) if cpu else ""
+        flop_per_launch = _gemm_flops(dims_key) if (cpu and flop_class == "gemm") else None
+        sub = groups[(phase, module)].setdefault(
+            (name, dims_key),
+            {"durs": [], "flop_class": flop_class, "flop_per_launch": flop_per_launch},
+        )
+        sub["durs"].append(dur)
+
+    optimizer_us /= max(1, num_steps)  # optimizer runs once per training iteration
+    dpcomm_us /= num_mb
+    dropped_stall_us /= num_mb  # report per-microbatch, like the breakdown rows
+
+    out = {b: {"forward": {}, "backward": {}} for b in ("attention", "moe", "embedding", "output", "loss")}
+    for (phase, module), subs in groups.items():
+        time_us = 0.0
+        compute_us = 0.0
+        flops = 0.0
+        has_flops = False
+        kern: dict[str, float] = defaultdict(float)
+        class_time: dict[str, float] = defaultdict(float)
+        for (name, _dims), sub in subs.items():
+            n = len(sub["durs"])
+            per_mb = sum(sub["durs"]) / num_mb
+            time_us += per_mb
+            kern[name[:80]] += per_mb
+            if sub["flop_class"]:
+                compute_us += per_mb
+                class_time[sub["flop_class"]] += per_mb
+            if sub["flop_per_launch"]:
+                flops += sub["flop_per_launch"] * n / num_mb
+                has_flops = True
+        compute = compute_us > 0.5 * time_us if time_us > 0 else False
+        flop_class = max(class_time, key=class_time.get) if (compute and class_time) else None
+        flops_out = flops if (compute and has_flops) else None
+        time_s = time_us * 1e-6
+        tflops = (flops_out / time_s / 1e12) if (flops_out and time_s > 0) else None
+        kern_list = sorted(
+            ({"name": n, "time_us": round(t, 3)} for n, t in kern.items()),
+            key=lambda k: -k["time_us"],
+        )[:6]
+        entry = {
+            "module": module,
+            "time_us": round(time_us, 3),
+            "class": "compute_bound" if compute else "memory_bound",
+            "flop_class": flop_class,
+            "flops": flops_out,
+            "tflops": round(tflops, 1) if tflops else None,
+            "kernels": kern_list,
+        }
+        # Logical bucket. Unattributed scalar kernels are labelled `attn.misc`
+        # by _resolve_module because the V4 traces show they are dominated by
+        # attention-side hyper-connection / Indexer control-flow work; this keeps
+        # the MoE bucket strictly the cr-independent router/dispatch/grouped_gemm
+        # /combine/shared_expert set (A10).
+        if module.startswith("attn."):
+            bucket = "attention"
+        elif module.startswith("moe."):
+            bucket = "moe"
+        elif module in ("embedding", "output", "loss"):
+            bucket = module
+        else:
+            bucket = "moe"
+        out[bucket][phase][module] = entry
+    return out, optimizer_us, dpcomm_us, dropped_stall_us
+
+
+def _lists(bd: dict) -> dict:
+    return {p: sorted(bd[p].values(), key=lambda r: -r["time_us"]) for p in ("forward", "backward")}
+
+
+def cr_layer_counts(compress: list[int]) -> dict[str, int]:
+    counts = {"0": 0, "4": 0, "128": 0}
+    for c in compress:
+        counts[str(c)] = counts.get(str(c), 0) + 1
+    return counts
+
+
+def _fwd_total(buckets: dict) -> float:
+    return sum(r["time_us"] for r in buckets["attention"]["forward"].values()) + sum(
+        r["time_us"] for r in buckets["moe"]["forward"].values()
+    )
+
+
+def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]:
+    cfg = MODEL_CONFIGS[model]
+    per_cr, opt_us = {}, []
+    dropped_stall = {}
+    for cr, path in traces.items():
+        buckets, o, _dp, stall = parse_trace(path, ga)
+        per_cr[cr] = buckets
+        opt_us.append(o)
+        dropped_stall[cr] = round(stall, 1)
+
+    # Some cr layers (pure dense cr=0 / HCA cr=128) get CUDA-graph / stream-
+    # captured, so their compute kernels are not individually visible in the
+    # trace (only optimizer/comm/elementwise appear). Fall back to the eager
+    # cr=4 sample for those so the full-model projection isn't zeroed; flag it.
+    ref = (
+        "4"
+        if "4" in per_cr and _fwd_total(per_cr["4"]) >= 1000
+        else max(per_cr, key=lambda c: _fwd_total(per_cr[c]))
+    )
+    graphed = []
+    for cr in list(per_cr):
+        if cr != ref and _fwd_total(per_cr[cr]) < 1000:
+            graphed.append(cr)
+            per_cr[cr] = per_cr[ref]
+
+    src = per_cr[ref]
+    layers = {cr: {"attention": _lists(b["attention"]), "moe": _lists(b["moe"])} for cr, b in per_cr.items()}
+    optimizer_us = round(sum(opt_us) / len(opt_us), 1) if opt_us else None
+
+    # Analytic V4 closed-form FLOPs (ported from Megatron's flops patch) so the
+    # site's TFLOP/s matches a real run. Per layer, per microbatch (B=1), at the
+    # capture seq; Megatron-convention (fwd+bwd, FMA -> x6).
+    cap_seq = 4096
+    mtp_num_layers = int(cfg.get("mtp_num_layers", 0) or 0)
+    mtp_cr = int((cfg.get("mtp_compress_ratios") or [0])[0])
+    mtp = mtp_flops(model, cap_seq, mtp_num_layers, mtp_cr)
+    analytic_flops = {
+        "per_cr_layer_flops": {
+            cr: sum(layer_fmac(model, int(cr), cap_seq).values()) * FB_FMA for cr in ("0", "4", "128")
+        },
+        "output_flops": nonlayer_fmac(model, cap_seq)["logits"] * FB_FMA,
+        "mtp": {
+            "num_layers": mtp_num_layers,
+            "compress_ratio": mtp_cr,
+            "inner_layer_flops": mtp["inner_layer"],
+            "eh_proj_flops": mtp["eh_proj"],
+            "extra_logits_flops": mtp["extra_logits"],
+            "hc_head_flops": mtp["hc_head"],
+        },
+        "seq": cap_seq,
+        "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP",
+    }
+
+    return {
+        "schema_version": 1,
+        "model": model,
+        "analytic_flops": analytic_flops,
+        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "provenance": {
+            "traces": {cr: str(p) for cr, p in traces.items()},
+            "graphed_crs_estimated_from_cr4": graphed,
+            "dropped_stall_us_per_mb": dropped_stall,
+            "note": (
+                "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured "
+                "(compute not visible in trace); their breakdown is copied from cr=4 as an "
+                "estimate. Re-run those cr with graph capture disabled for exact numbers. "
+                "dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as "
+                "implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)."
+            ),
+        },
+        "capture": {
+            "gpu": "MI355X",
+            "seq_length": 4096,
+            "micro_batch_size": 1,
+            "tokens_per_microbatch": 4096,
+            "ep": 8,
+            "ga_for_capture": 2,
+            "optimizer": "adam",
+            "distributed_optimizer": False,
+            "recompute": "off",
+            "measured_iter_time_ms": MODEL_ANCHORS.get(model, {}).get("measured_iter_time_ms"),
+            "measured_anchor": MODEL_ANCHORS.get(model, {}).get("measured_anchor"),
+        },
+        "model_config": {
+            **cfg,
+            "cr_layer_counts": cr_layer_counts(cfg["compress_ratios"]),
+            "total_params": model_total_params(model, cfg["num_layers"], mtp_num_layers),
+        },
+        "hardware": DEFAULT_HARDWARE,
+        "layers": layers,
+        "non_layer": {k: _lists(src[k]) for k in ("embedding", "output", "loss")},
+        "optimizer": {
+            "type": "adam",
+            "measured_params": None,
+            "time_us": optimizer_us,
+            "bytes_per_param": 30,
+            "class": "memory_bound",
+            "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference",
+        },
+        "comm": {
+            "ep_dispatch_us": None,
+            "ep_combine_us": None,
+            "note": "EP dispatch/combine are memory_bound rows inside moe; informational",
+        },
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Build projection breakdown JSON from per-cr traces.")
+    p.add_argument("--model", required=True, choices=sorted(MODEL_CONFIGS))
+    p.add_argument("--trace", action="append", default=[], metavar="cr=PATH")
+    p.add_argument(
+        "--ga", type=int, default=2, help="gradient-accumulation at capture (GBS/(DP*MBS)); default 2"
+    )
+    p.add_argument("--out", required=True, type=Path)
+    return p.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    traces: dict[str, Path] = {}
+    for spec in args.trace:
+        if "=" not in spec:
+            raise SystemExit(f"--trace must be cr=PATH, got: {spec}")
+        cr, path = spec.split("=", 1)
+        traces[cr.replace("cr", "")] = Path(path)
+    if not traces:
+        raise SystemExit("at least one --trace cr=PATH is required")
+    for cr, path in traces.items():
+        if not path.exists():
+            raise SystemExit(f"trace not found for cr={cr}: {path}")
+
+    doc = build(args.model, traces, args.ga)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(doc, indent=2) + "\n")
+    print(f"[parse_trace] wrote {args.out} (model={args.model}, crs={sorted(traces)})")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/examples/deepseek-v4/projection/tools/v4_flops.py b/examples/deepseek-v4/projection/tools/v4_flops.py
new file mode 100644
index 000000000..34329866e
--- /dev/null
+++ b/examples/deepseek-v4/projection/tools/v4_flops.py
@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""DeepSeek-V4 closed-form analytic FLOPs, ported from
+``primus/backends/megatron/patches/deepseek_v4_flops_patches.py`` so the
+projection's TFLOP/s matches what Megatron reports on a real run.
+
+Returns per-component FMAC (multiply-only, pre fwd+bwd/FMA expansion) for ONE
+layer of a given cr at batch_size=1. Multiply by FB_FMA (=6) for Megatron-
+convention FLOPs (fwd 1 + bwd 2, times FMA 2).
+
+Validated against the measured flash 16-layer run (TOTAL 34093 TFLOP/global-
+batch; per-component within rounding) — see __main__ self-test.
+"""
+
+from __future__ import annotations
+
+FB_FMA = 6  # _FORWARD_BACKWARD_FACTOR(3) * _FMA_FACTOR(2)
+SWIGLU = 3  # gate+up+down collapsed expansion factor
+DEFAULT_MTP_LAYERS = {"pro": 1, "flash": 1}
+
+# Per-model architecture params (from primus/configs/models/megatron/*.yaml +
+# deepseek_v4_base.yaml). Shared: index_head_dim=128, index_n_heads=64,
+# attn_sliding_window=128, hc_mult=4.
+MODEL_PARAMS = {
+    "pro": dict(
+        hidden=7168,
+        heads=128,
+        head_dim=512,
+        q_lora=1536,
+        o_lora=1024,
+        o_groups=16,
+        moe_ffn=3072,
+        shared_ffn=3072,
+        topk=6,
+        experts=384,
+        index_topk=1024,
+        vocab=129280,
+    ),
+    "flash": dict(
+        hidden=4096,
+        heads=64,
+        head_dim=512,
+        q_lora=1024,
+        o_lora=1024,
+        o_groups=8,
+        moe_ffn=2048,
+        shared_ffn=2048,
+        topk=6,
+        experts=256,
+        index_topk=512,
+        vocab=129280,
+    ),
+}
+SHARED = dict(index_head_dim=128, index_n_heads=64, swa_window=128, hc_mult=4)
+
+
+def _local_visible_pairs(swa, s):
+    if swa <= 0 or swa >= s:
+        return s * (s + 1) // 2
+    return swa * s - swa * (swa - 1) // 2
+
+
+def _pool_visible_pairs(cr, s):
+    if cr <= 0 or s <= 0:
+        return 0
+    n = s // cr
+    if n == 0:
+        return 0
+    return cr * n * (n - 1) // 2 + n * (s - cr * n + 1)
+
+
+def _visible_pairs(swa, cr, index_topk, s):
+    local = _local_visible_pairs(swa, s)
+    if cr == 0:
+        return local
+    pool = max(1, s // cr)
+    if cr == 128:
+        return local + _pool_visible_pairs(cr, s)
+    if cr == 4:
+        sparse = min(index_topk if index_topk else pool, pool)
+        return local + sparse * s
+    return local + pool * s
+
+
+def _attn_qkv_o(s_eff, p):
+    n_d = p["heads"] * p["head_dim"]
+    qkv = p["hidden"] * p["q_lora"] + p["q_lora"] * n_d + p["hidden"] * p["head_dim"]
+    if p["o_lora"] > 0:
+        o_proj = n_d * p["o_lora"] + (p["o_groups"] * p["o_lora"]) * p["hidden"]
+    else:
+        o_proj = n_d * p["hidden"]
+    return s_eff * (qkv + o_proj)
+
+
+def _attn_scores(s_eff, cr, p):
+    pairs = _visible_pairs(SHARED["swa_window"], cr, p["index_topk"], s_eff)
+    return 2 * p["heads"] * p["head_dim"] * pairs
+
+
+def _compressor(s_eff, cr, p):
+    if cr == 0:
+        return 0
+    coff = 2 if cr == 4 else 1
+    return 2 * s_eff * p["hidden"] * (coff * p["head_dim"])
+
+
+def _indexer(s_eff, cr, p):
+    if cr != 4:
+        return 0
+    ihd, inh = SHARED["index_head_dim"], SHARED["index_n_heads"]
+    pool = max(1, s_eff // cr)
+    dq_rank = ihd
+    proj = p["hidden"] * dq_rank + dq_rank * (inh * ihd) + p["hidden"] * inh
+    proj += 2 * p["hidden"] * (2 * ihd)  # mini-compressor
+    return s_eff * proj + s_eff * inh * pool * ihd
+
+
+def _moe(s_eff, p):
+    router = p["hidden"] * p["experts"]
+    routed = p["topk"] * SWIGLU * p["hidden"] * p["moe_ffn"]
+    shared = SWIGLU * p["hidden"] * p["shared_ffn"] if p["shared_ffn"] > 0 else 0
+    return s_eff * (router + routed + shared)
+
+
+def _hc_mixer(s, p):
+    hc = SHARED["hc_mult"]
+    n_d = hc * p["hidden"]
+    return 2 * s * n_d * ((2 + hc) * hc)
+
+
+def _hc_head(s, p, mtp_num_layers):
+    hc = SHARED["hc_mult"]
+    n_d = hc * p["hidden"]
+    return (1 + mtp_num_layers) * s * n_d * hc
+
+
+def _mtp_eh_proj(s, p, mtp_num_layers):
+    return mtp_num_layers * s * (2 * p["hidden"]) * p["hidden"]
+
+
+def layer_fmac(model: str, cr: int, seq: int) -> dict[str, float]:
+    """Per-layer FMAC components (batch_size=1) for one cr layer."""
+    p = MODEL_PARAMS[model]
+    s_eff = seq * SHARED["hc_mult"]
+    return {
+        "attn_qkv_o": _attn_qkv_o(s_eff, p),
+        "attn_scores": _attn_scores(s_eff, cr, p),
+        "compressor": _compressor(s_eff, cr, p),
+        "indexer": _indexer(s_eff, cr, p),
+        "moe": _moe(s_eff, p),
+        "hc": _hc_mixer(seq, p),
+    }
+
+
+def nonlayer_fmac(model: str, seq: int, mtp_num_layers: int = 0) -> dict[str, float]:
+    p = MODEL_PARAMS[model]
+    return {"logits": (1 + mtp_num_layers) * seq * p["hidden"] * p["vocab"]}
+
+
+def mtp_fmac(model: str, seq: int, mtp_num_layers: int = 1, mtp_cr: int = 4) -> dict[str, float]:
+    """Extra FMAC components for MTP depths, batch_size=1.
+
+    V4 MTP reuses a full V4 inner layer per depth; the current Flash Megatron
+    FLOPs anchor reports a CSA-style MTP inner layer (cr=4).
+    """
+    if mtp_num_layers <= 0:
+        return {
+            "inner_layer": 0,
+            "eh_proj": 0,
+            "extra_logits": 0,
+            "hc_head": _hc_head(seq, MODEL_PARAMS[model], 0),
+        }
+    p = MODEL_PARAMS[model]
+    inner = sum(layer_fmac(model, mtp_cr, seq).values()) * mtp_num_layers
+    main_logits = seq * p["hidden"] * p["vocab"]
+    return {
+        "inner_layer": inner,
+        "eh_proj": _mtp_eh_proj(seq, p, mtp_num_layers),
+        "extra_logits": main_logits * mtp_num_layers,
+        "hc_head": _hc_head(seq, p, mtp_num_layers),
+    }
+
+
+def model_total_params(model: str, num_layers: int, mtp_num_layers: int = 0) -> int:
+    """Approximate total parameter count (for optimizer-step sizing). Uses the
+    same V4 MLA low-rank attention shapes as the FLOPs formula (q/o LoRA + single
+    latent KV) instead of the crude 4*h^2, plus MoE experts + shared + router and
+    the tied-free embedding/output. MTP adds one full V4 inner layer plus the
+    2H->H eh_proj per depth; logits reuse the output layer weights."""
+    p = MODEL_PARAMS[model]
+    n_d = p["heads"] * p["head_dim"]
+    attn = (
+        p["hidden"] * p["q_lora"]
+        + p["q_lora"] * n_d
+        + p["hidden"] * p["head_dim"]
+        + n_d * p["o_lora"]
+        + p["o_groups"] * p["o_lora"] * p["hidden"]
+    )
+    moe = (
+        p["experts"] * SWIGLU * p["hidden"] * p["moe_ffn"]
+        + SWIGLU * p["hidden"] * p["shared_ffn"]
+        + p["hidden"] * p["experts"]
+    )
+    return int(
+        (num_layers + mtp_num_layers) * (attn + moe)
+        + mtp_num_layers * 2 * p["hidden"] * p["hidden"]
+        + 2 * p["vocab"] * p["hidden"]
+    )
+
+
+# Map analytic components to projection module names (per layer).
+def module_flops(model: str, cr: int, seq: int) -> dict[str, float]:
+    """Megatron-convention FLOPs (×FB_FMA) per module, per layer, batch_size=1."""
+    f = layer_fmac(model, cr, seq)
+    return {
+        "attn.proj": f["attn_qkv_o"] * FB_FMA,
+        "attn.core": f["attn_scores"] * FB_FMA,
+        "attn.indexer": (f["compressor"] + f["indexer"]) * FB_FMA,
+        "attn.norm": f["hc"] * FB_FMA,
+        "moe.grouped_gemm": f["moe"] * FB_FMA,
+    }
+
+
+def output_flops(model: str, seq: int) -> float:
+    return nonlayer_fmac(model, seq)["logits"] * FB_FMA
+
+
+def mtp_flops(model: str, seq: int, mtp_num_layers: int = 1, mtp_cr: int = 4) -> dict[str, float]:
+    f = mtp_fmac(model, seq, mtp_num_layers, mtp_cr)
+    return {k: v * FB_FMA for k, v in f.items()}
+
+
+def _self_test() -> None:
+    """Self-test against measured flash 16L (GBS64): cr [0x3,4x6,128x7]."""
+    sched = [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]
+    seq_t, batch = 4096, 64
+    comp = {k: 0.0 for k in ("attn_qkv_o", "attn_scores", "compressor", "indexer", "moe", "hc")}
+    for cr_t in sched:
+        for k, v in layer_fmac("flash", cr_t, seq_t).items():
+            comp[k] += v
+    logits = nonlayer_fmac("flash", seq_t)["logits"]
+    tot = (sum(comp.values()) + logits) * FB_FMA * batch / 1e12
+    print("flash 16L analytic vs measured (TFLOP/global-batch):")
+    for k, v in comp.items():
+        print(f"  {k:12s} = {v*FB_FMA*batch/1e12:9.1f}")
+    print(f"  {'logits':12s} = {logits*FB_FMA*batch/1e12:9.1f}")
+    print(f"  TOTAL        = {tot:9.1f}   (measured 34093.4)")
+
+
+if __name__ == "__main__":
+    _self_test()
diff --git a/examples/deepseek-v4/rccl_avg_workaround/.gitignore b/examples/deepseek-v4/rccl_avg_workaround/.gitignore
new file mode 100644
index 000000000..7a60b85e1
--- /dev/null
+++ b/examples/deepseek-v4/rccl_avg_workaround/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc
diff --git a/examples/deepseek-v4/rccl_avg_workaround/sitecustomize.py b/examples/deepseek-v4/rccl_avg_workaround/sitecustomize.py
new file mode 100644
index 000000000..bf10ac87e
--- /dev/null
+++ b/examples/deepseek-v4/rccl_avg_workaround/sitecustomize.py
@@ -0,0 +1,143 @@
+"""gfx1250 single-GPU bring-up workarounds, auto-imported in every Python
+worker via sitecustomize (this dir is on PYTHONPATH). Two independent fixes:
+
+1. RCCL AVG hang: torch.distributed.all_reduce(op=AVG) HANGS on this build for
+   (at least) single-rank process groups, while SUM works fine (verified by
+   collective microbench). Megatron's MoE aux-loss metric reduction
+   (moe_utils.reduce_aux_losses_tracker_across_ranks) uses op=AVG and
+   deadlocks. Replace AVG with SUM + divide-by-world-size, which is
+   mathematically identical for any world size.
+
+2. primus_turbo import shim: the MI355X production containers bundle the
+   `primus_turbo` package; the gfx1250 therock container does not. Most Primus
+   call-sites guard the import (try/except -> HAVE_TURBO=False), but the V4
+   model path imports it unconditionally:
+     deepseek_v4_layer_specs.py -> transformer_engine_spec_provider.py
+       (DeepSeekV4SpecProvider subclasses PrimusTurboSpecProvider)
+       -> extensions/primus_turbo.py -> `import primus_turbo.pytorch`
+   and backends/megatron/core/utils.py -> primus_turbo...attention_utils.
+   With every use_turbo_* flag False the turbo classes are never SELECTED
+   (the spec provider returns the TE classes), so a pure import-shim is safe:
+   install a meta-path finder that fabricates stub modules for primus_turbo.*
+   whose attributes are auto-generated dummy classes. Attribute chains
+   evaluated at class-definition time (e.g. the ScalingGranularity.TENSORWISE
+   default arg in PrimusTurboQuantConfig) resolve fine; actually CALLING or
+   instantiating any stub raises RuntimeError, so a misrouted turbo path fails
+   loudly instead of computing garbage. The shim only installs when the real
+   package is absent, so it can never shadow a real primus_turbo install.
+"""
+
+import sys
+
+
+def _install_rccl_avg_workaround():
+    import torch.distributed as dist
+
+    orig_all_reduce = dist.all_reduce
+    avg_op = dist.ReduceOp.AVG
+
+    def all_reduce_avg_safe(tensor, op=dist.ReduceOp.SUM, group=None, async_op=False):
+        is_avg = False
+        try:
+            is_avg = op == avg_op
+        except Exception:
+            is_avg = str(op) == str(avg_op)
+        if is_avg:
+            work = orig_all_reduce(tensor, op=dist.ReduceOp.SUM, group=group, async_op=async_op)
+            try:
+                ws = dist.get_world_size(group)
+            except Exception:
+                ws = 1
+            if ws and ws > 1:
+                # Enqueued on the same stream after the all_reduce, so ordering holds.
+                tensor.div_(ws)
+            return work
+        return orig_all_reduce(tensor, op=op, group=group, async_op=async_op)
+
+    dist.all_reduce = all_reduce_avg_safe
+    print(
+        "[rccl_avg_workaround] patched torch.distributed.all_reduce (AVG -> SUM/ws)",
+        file=sys.stderr,
+        flush=True,
+    )
+
+
+def _install_primus_turbo_stub():
+    import importlib.abc
+    import importlib.machinery
+    import importlib.util
+    import types
+
+    # Never shadow a real install.
+    if importlib.util.find_spec("primus_turbo") is not None:
+        return
+
+    class _StubMeta(type):
+        """Dummy-class metaclass: any attribute access mints another dummy
+        class (covers enum-style chains like ScalingGranularity.TENSORWISE)."""
+
+        def __getattr__(cls, name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            dummy = _make_dummy(f"{cls._stub_qual}.{name}")
+            setattr(cls, name, dummy)
+            return dummy
+
+    def _make_dummy(qual):
+        def _raise(self, *args, **kwargs):
+            raise RuntimeError(
+                f"primus_turbo stub: {qual} was invoked, but primus_turbo is NOT "
+                "installed in this container. A turbo code path ran despite all "
+                "use_turbo_*/enable_primus_turbo flags being False — fix the flags "
+                "instead of installing primus_turbo."
+            )
+
+        return _StubMeta(qual.rsplit(".", 1)[-1], (), {"__init__": _raise, "_stub_qual": qual})
+
+    class _StubModule(types.ModuleType):
+        def __getattr__(self, name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            dummy = _make_dummy(f"{self.__name__}.{name}")
+            setattr(self, name, dummy)
+            return dummy
+
+    class _Finder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
+        def find_spec(self, fullname, path=None, target=None):
+            if fullname == "primus_turbo" or fullname.startswith("primus_turbo."):
+                return importlib.machinery.ModuleSpec(fullname, self, is_package=True)
+            return None
+
+        def create_module(self, spec):
+            return _StubModule(spec.name)
+
+        def exec_module(self, module):
+            module.__path__ = []
+            module._primus_turbo_stub = True
+
+    sys.meta_path.append(_Finder())
+    print(
+        "[primus_turbo_stub] primus_turbo not installed -> import shim active "
+        "(turbo classes import as raising stubs; all use_turbo_* must stay False)",
+        file=sys.stderr,
+        flush=True,
+    )
+
+
+# Only patch the actual training worker. Importing torch here is heavy, so we
+# must NOT do it for pip/offload-arch/build helper invocations (they stall and
+# block setup). Gate on the training entrypoint appearing in argv.
+def _is_training_worker():
+    argv = " ".join(sys.argv)
+    return ("primus/cli/main.py" in argv) or ("run_pretrain" in argv) or ("pretrain" in argv)
+
+
+if _is_training_worker():
+    try:
+        _install_rccl_avg_workaround()
+    except Exception as e:  # noqa: BLE001
+        print(f"[rccl_avg_workaround] install FAILED: {e}", file=sys.stderr, flush=True)
+    try:
+        _install_primus_turbo_stub()
+    except Exception as e:  # noqa: BLE001
+        print(f"[primus_turbo_stub] install FAILED: {e}", file=sys.stderr, flush=True)
diff --git a/examples/deepseek-v4/run_deepseek_v4.sh b/examples/deepseek-v4/run_deepseek_v4.sh
new file mode 100755
index 000000000..0cd3797ec
--- /dev/null
+++ b/examples/deepseek-v4/run_deepseek_v4.sh
@@ -0,0 +1,393 @@
+#!/bin/bash
+set -euo pipefail
+set -x
+
+_RUN_START_SEC=$(date +%s)
+_RUN_START_TS=$(date '+%Y-%m-%d %H:%M:%S')
+_print_run_elapsed() {
+  local _end_sec _end_ts _elapsed _exit=$1
+  _end_sec=$(date +%s)
+  _end_ts=$(date '+%Y-%m-%d %H:%M:%S')
+  _elapsed=$((_end_sec - _RUN_START_SEC))
+  echo "----------------------------------------"
+  echo "run_deepseek_v4.sh wall time"
+  echo "  start:   ${_RUN_START_TS}"
+  echo "  end:     ${_end_ts}"
+  echo "  elapsed: ${_elapsed}s ($((_elapsed / 60))m $((_elapsed % 60))s)"
+  echo "  exit:    ${_exit}"
+}
+trap '_print_run_elapsed $?' EXIT
+
+export HF_TOKEN="${HF_TOKEN:-}"
+export WANDB_API_KEY="${WANDB_API_KEY:-your_wandb_api_key}"
+
+export NNODES=${NNODES:-1}
+export TRAIN_ITERS=${TRAIN_ITERS:-20}
+
+export DOCKER_IMAGE=${DOCKER_IMAGE:?set DOCKER_IMAGE to a Primus container image}
+export SLURM_PARTITION=${SLURM_PARTITION:-}
+export SLURM_NODELIST=${SLURM_NODELIST:-}
+export MASTER_PORT=${MASTER_PORT:-29500}
+
+export USING_AINIC=${USING_AINIC:-1}
+export NCCL_IB_HCA="ionic_0:1,ionic_1:1,ionic_2:1,ionic_3:1,ionic_4:1,ionic_5:1,ionic_6:1,ionic_7:1"
+# Default socket interface to loopback (single-node fallback). Override with
+# GLOO_SOCKET_IFNAME / NCCL_SOCKET_IFNAME for multi-node (e.g. ens3).
+export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-lo}
+export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-lo}
+export NCCL_IB_GID_INDEX=1
+export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1}
+export NVTE_CK_USES_BWD_V3=${NVTE_CK_USES_BWD_V3:-1}
+
+# Phase-7 fixed knobs for single-node bring-up.
+export MBS=${MBS:-1}
+export GBS=${GBS:-$((16 * NNODES * MBS))}
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-1}
+export PRIMUS_EP=${PRIMUS_EP:-8}
+
+# Keep this smoke config lightweight for quick bring-up.
+export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-8}
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-128}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-128}
+export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-8}
+export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-2}
+export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-512}
+export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-8}
+export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-"[0,0,4,4,4,4,4,0]"}
+export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False}
+export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True}
+export PROFILE=${PROFILE:-False}
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-False}
+export LEGACY_GG=${LEGACY_GG:-False}
+# MegaMoE: FlyDSL-based fused MoE layer replacing Megatron's MoELayer (see
+# docs/04-technical-guides/mega-moe.md). It owns the whole expert path --
+# dispatch/combine all-to-all is fused into the grouped GEMMs -- so it is
+# mutually exclusive with turbo DeepEP; force DeepEP off rather than letting
+# both patch the MoE layer. Requires EP-only (TP=1) + bf16 + EP>1.
+export USE_TURBO_MEGA_MOE=${USE_TURBO_MEGA_MOE:-False}
+if [ "$USE_TURBO_MEGA_MOE" = "True" ]; then
+  export USE_TURBO_DEEPEP=False
+fi
+# Plan-3 P22 / P23: PrimusTurbo gate (must be on for turbo attention /
+# turbo deepep to take effect; enable_primus_turbo gates the
+# `before_train` patches that re-bind the spec provider).
+export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-False}
+if [ "$USE_TURBO_ATTENTION" = "True" ] || [ "${USE_TURBO_DEEPEP:-False}" = "True" ] ||
+  [ "$USE_TURBO_MEGA_MOE" = "True" ]; then
+  ENABLE_PRIMUS_TURBO=True
+fi
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-False}
+
+if [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then
+  export PRIMUS_BIAS_SWIGLU_FUSION=True
+fi
+
+# Plan-3 P23: Turbo DeepEP-related knobs.  Only emit these CLI flags
+# when USE_TURBO_DEEPEP=True so non-deepep runs don't carry unrelated
+# overrides.  Best-practice CU count: 64 (or 80) for EP=8, 32 for
+# EP>=16 — the EP>=16 cap is asserted by
+# `primus/modules/trainer/megatron/utils.py:527`.  DeepEP itself
+# requires `moe_router_dtype=fp32` and forbids
+# `moe_shared_expert_overlap=True` (both are already V4-Flash YAML
+# defaults; we pin them via CLI defensively so a stray YAML override
+# or future config edit cannot flip them out from under the Turbo
+# path mid-run).
+TURBO_DEEPEP_CLI_ARGS=()
+if [ "$USE_TURBO_DEEPEP" = "True" ]; then
+  if [ "${PRIMUS_EP:-1}" -ge 16 ]; then
+    _DEFAULT_TURBO_DEEPEP_NUM_CU=32
+  else
+    _DEFAULT_TURBO_DEEPEP_NUM_CU=80
+  fi
+  export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-$_DEFAULT_TURBO_DEEPEP_NUM_CU}
+  export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False}
+  export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32}
+  export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False}
+  TURBO_DEEPEP_CLI_ARGS=(
+    --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU"
+    --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM"
+    --moe_router_dtype "$MOE_ROUTER_DTYPE"
+    --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP"
+  )
+fi
+
+# TransformerEngine full-scope CUDA graph capture. `--external_cuda_graph True`
+# maps to cuda_graph_impl="transformer_engine"; scope `full` is normalized to []
+# by Megatron, i.e. capture the whole layer. Pairs well with MegaMoE, which is
+# sync-free (no device-to-host sync in the expert path) and captures cleanly.
+export ENABLE_CUDA_GRAPH=${ENABLE_CUDA_GRAPH:-False}
+CUDA_GRAPH_CLI_ARGS=()
+if [ "$ENABLE_CUDA_GRAPH" = "True" ]; then
+  CUDA_GRAPH_CLI_ARGS=(
+    --external_cuda_graph True
+    --cuda_graph_scope "${CUDA_GRAPH_SCOPE:-full}"
+    --cuda_graph_warmup_steps "${CUDA_GRAPH_WARMUP_STEPS:-3}"
+  )
+fi
+
+export PRECISION_TYPE=${PRECISION_TYPE:-BF16}
+# Honor an incoming FP8 / FP8_RECIPE env (e.g. FP8_RECIPE=mxfp8); default null
+# so non-FP8 runs are unchanged. (Previously these were hard-set to null,
+# which silently clobbered a caller-provided recipe.)
+export FP8=${FP8:-null}
+export FP8_RECIPE=${FP8_RECIPE:-null}
+
+# ---------- Optimizer selection (adam default; muon = DeepSeek-V4 recipe) ----
+# OPTIMIZER=adam (default): unchanged behaviour (BF16 precision-aware AdamW
+#   from the EXP yaml); overlap_grad_reduce / overlap_param_gather stay ON.
+# OPTIMIZER=muon: Primus distributed-Muon path (primus .../optimizer/moun.py).
+#   Megatron asserts plain `muon` is incompatible with distributed optimizer +
+#   grad/param overlap, so we force them OFF and switch optimizer states to
+#   fp32 (Muon does not support the precision-aware optimizer). The
+#   Newton-Schulz coefficient set auto-selects 'deepseekv4' (8 aggressive + 2
+#   stable) for V4 configs inside get_megatron_muon_optimizer. Requires the
+#   emerging_optimizers package -> we set PRIMUS_INSTALL_EMERGING_OPTIMIZERS so
+#   the in-container install hook (runner/.../01_install_emerging_optimizers.sh)
+#   provisions it.
+export OPTIMIZER=${OPTIMIZER:-adam}
+export PRIMUS_OVERLAP_GRAD_REDUCE=${PRIMUS_OVERLAP_GRAD_REDUCE:-True}
+export PRIMUS_OVERLAP_PARAM_GATHER=${PRIMUS_OVERLAP_PARAM_GATHER:-True}
+OPTIMIZER_CLI_ARGS=()
+if [ "$OPTIMIZER" = "muon" ] || [ "$OPTIMIZER" = "dist_muon" ]; then
+  export PRIMUS_INSTALL_EMERGING_OPTIMIZERS=${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-1}
+  export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95}
+  export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18}
+  # Both plain muon (Megatron asserts) and dist_muon (LayerWiseDistributed-
+  # Optimizer docstring: "keep all megatron distributed-optimizer related
+  # options OFF"; it manages its own param all-gather, so DDP
+  # overlap_param_gather double-drives start_param_sync -> crash) need the
+  # DDP grad/param overlap OFF.
+  PRIMUS_OVERLAP_GRAD_REDUCE=False
+  PRIMUS_OVERLAP_PARAM_GATHER=False
+  OPTIMIZER_CLI_ARGS=(
+    --optimizer "$OPTIMIZER"
+    --muon_momentum "$MUON_MOMENTUM"
+    --muon_extra_scale_factor "$MUON_EXTRA_SCALE_FACTOR"
+    --use_distributed_optimizer False
+    --use_precision_aware_optimizer False
+    --main_grads_dtype fp32
+    --exp_avg_dtype fp32
+    --exp_avg_sq_dtype fp32
+  )
+fi
+
+# DeepSeek-V4 attention backend selection (unified string selectors). Default
+# triton_v2 (production default; fastest V4 sparse-MLA path). These are
+# V4-only; no effect on other model types.
+#   USE_V4_ATTENTION_BACKEND     (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon
+#   USE_V4_CSA_ATTENTION_BACKEND (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0
+# gluon is gfx950/CDNA4-only (lazily imported; asserts arch when selected).
+# use_turbo_attention (when core_attention is built) still wins for the dense path.
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-turbo}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-turbo}
+
+# Plan-9: FP8 (E4M3) Indexer QK path (CSA selector). Default OFF; flip with
+# USE_V4_FP8_INDEXER=True. Passed as a CLI override so it reliably reaches the
+# in-container config regardless of env propagation.
+export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-False}
+
+# Indexer distillation loss coefficient (CSA selector training). 0 keeps the
+# loss off and the indexer frozen -- correct when loading an already-trained
+# indexer. A from-scratch pretrain needs it ON (1e-2 is a reasonable starting
+# value), which also unfreezes the indexer params.
+export PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF=${PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF:-0.0}
+
+# Plan-5 P29 (RESCOPED): wrap sinkhorn_normalize in HyperMixer with a
+# cached torch.compile build. Default OFF here; the proxy script
+# (run_deepseek_v4_flash_proxy.sh) flips it ON. After G32 + G33b are
+# green, the default flips to True for the V4-Flash configs.
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False}
+
+# Plan-4 P27: TP-side guard for the V4 Triton kernels.
+# The dense / HCA / CSA kernels operate on the local head slice (each
+# rank only sees H/TP query heads) so TP-sharded execution is correct
+# by construction (no in-kernel collective comm needed).  Plan-4 unit
+# tests / smoke gates exercise TP=1 only; emit a soft warning when a
+# user enables the kernels at TP>1 so any TP-related regression is
+# easy to attribute.  TP=1 is the V4-Flash / V4-Pro release default
+# (release configs use PP+EP for parallelism, never TP).
+if echo "$USE_V4_ATTENTION_BACKEND $USE_V4_CSA_ATTENTION_BACKEND" | grep -q "triton" && [ "${PRIMUS_TP:-1}" -gt 1 ]; then
+  echo "[WARN] Plan-4 V4 Triton kernels enabled at PRIMUS_TP=${PRIMUS_TP}>1; this combination is not covered by Plan-4 unit tests / smoke gates (G28..G30 ran TP=1 only). Functionally the kernels operate per-rank on the local H/TP head slice, so this should work, but treat any TP>1 regression as a Plan-4 follow-up."
+fi
+
+if [ "$PRECISION_TYPE" = "FP8" ]; then
+  # Default to the paper's ue8m0 microscaling (e4m3 + mxfp8); honor explicit
+  # FP8 / FP8_RECIPE overrides. Sentinel-aware because "null" is non-empty, so
+  # a plain ${FP8:-...} would keep the off-sentinel instead of defaulting.
+  [ "$FP8" = "null" ] && export FP8=e4m3
+  [ "$FP8_RECIPE" = "null" ] && export FP8_RECIPE=mxfp8
+fi
+
+# ---------- MXFP8 + FP8 param-gather (Muon path; Megatron #4987 analogue) ----
+# Plan-9: combine the distributed-Muon (LayerWise) path with an MXFP8 forward
+# recipe + FP8 parameter all-gather. Enable with FP8_PARAM_GATHER=True (best
+# paired with OPTIMIZER=dist_muon + PRECISION_TYPE=FP8 FP8_RECIPE=mxfp8).
+# MXFP8 on ROCm/TE requires NVTE_ROCM_ENABLE_MXFP8=1; the mxfp8 param-AG path
+# is most memory-efficient with --reuse-grad-buf-for-mxfp8-param-ag. NOTE:
+# Megatron auto-disables --fp8-param-gather on TE>=2.0.0 (falls back to a
+# bf16/all_gather), so on such containers this exercises the MXFP8 forward +
+# dist-Muon path with param-gather requested-but-possibly-downgraded.
+export FP8_PARAM_GATHER=${FP8_PARAM_GATHER:-False}
+FP8_PARAM_GATHER_CLI_ARGS=()
+if [ "$FP8_PARAM_GATHER" = "True" ]; then
+  export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1}
+  export REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG=${REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG:-True}
+  FP8_PARAM_GATHER_CLI_ARGS=(--fp8_param_gather True)
+  if [ "$REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG" = "True" ] && [ "$FP8_RECIPE" = "mxfp8" ]; then
+    FP8_PARAM_GATHER_CLI_ARGS+=(--reuse_grad_buf_for_mxfp8_param_ag True)
+  fi
+fi
+
+# Force load balancing discards the real router decision so every expert receives
+# a similar number of tokens, removing run-to-run imbalance noise. This selects
+# *how*: `even` (Primus default) gives exactly equal, step-invariant per-expert
+# counts, so grouped-GEMM shapes never change; `uniform` balances only
+# statistically and keeps the step-to-step shape variation of real routing.
+# docs/04-technical-guides/mega-moe.md recommends `uniform` for benchmarking,
+# because `even` disproportionately favours the non-fused grouped-GEMM path.
+# Empty (default) leaves the config value alone.
+export MOE_FORCE_LB_TYPE=${MOE_FORCE_LB_TYPE:-}
+MOE_FORCE_LB_ARGS=()
+if [ -n "$MOE_FORCE_LB_TYPE" ]; then
+  MOE_FORCE_LB_ARGS=(--moe_router_force_load_balancing_type "$MOE_FORCE_LB_TYPE")
+fi
+
+PP_LAYOUT_ARGS=()
+if [ -n "${PRIMUS_PP_LAYOUT:-}" ]; then
+  PP_LAYOUT_ARGS=(--pipeline_model_parallel_layout "$PRIMUS_PP_LAYOUT")
+fi
+
+PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-0}
+
+export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml}
+export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"}
+export PRIMUS_TEAM=${PRIMUS_TEAM:-amd}
+export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_smoke_${PRECISION_TYPE}_MBS${MBS}_GBS${GBS}_PP${PRIMUS_PP}_EP${PRIMUS_EP}}
+# Host-side directory for the launcher's aggregated log. Defaults to the
+# canonical "output" tree; override when that tree is not writable by the
+# invoking user (e.g. it was created by an earlier root/sudo run).
+export PRIMUS_OUTPUT_ROOT=${PRIMUS_OUTPUT_ROOT:-output}
+
+if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then
+  echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH"
+  echo "Run: git submodule update --init --recursive"
+  exit 1
+fi
+
+mkdir -p "$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME"
+
+# On spur, direct the sbatch job's aggregated stdout+stderr (the actual per-node
+# training log) into the experiment output dir. sbatch returns right after
+# submission, so we can't `tee` it here; spur's sbatch reads SBATCH_OUTPUT/ERROR.
+if command -v spur >/dev/null 2>&1; then
+    export SBATCH_OUTPUT="$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/train_sbatch_output.log"
+    export SBATCH_ERROR="$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/train_sbatch_error.log"
+fi
+
+export PRIMUS_EXIT_FAST=1
+
+# Launcher: slurm (default, multi-node cluster) or direct (single-node, already
+# inside the container — e.g. local smoke on one box). PRIMUS_LAUNCHER=direct
+# drops the SLURM/srun + docker-image wrap that 'direct' doesn't use.
+export PRIMUS_LAUNCHER=${PRIMUS_LAUNCHER:-slurm}
+if [ "$PRIMUS_LAUNCHER" = "direct" ]; then
+  LAUNCHER_ARGS=(direct)
+  if [ "${PRIMUS_NUMA_BIND:-1}" = "1" ]; then
+    LAUNCHER_ARGS+=(--numa)
+  fi
+else
+  LAUNCHER_ARGS=(slurm "${SLURM_LAUNCH_CMD:-srun}" -N "$NNODES")
+  if [ -n "${SLURM_ATTACH_JOBID:-}" ]; then
+    # Run as a step inside an allocation that is already held (e.g. a long-lived
+    # `sbatch --exclusive --wrap "sleep ..."` holder), so a sweep of runs lands on
+    # the same nodes without re-queueing per run. partition / qos / account /
+    # nodelist / exclusive belong to the holder job -- passing them again on an
+    # attached step is redundant and spur rejects some of them.
+    LAUNCHER_ARGS+=(--jobid="${SLURM_ATTACH_JOBID}" --overlap)
+  else
+    [ -n "${SLURM_PARTITION:-}" ] && LAUNCHER_ARGS+=(--partition="${SLURM_PARTITION}")
+    [ -n "${SLURM_NODELIST:-}" ] && LAUNCHER_ARGS+=(--nodelist="${SLURM_NODELIST}")
+    [ -n "${SLURM_QOS:-}" ] && LAUNCHER_ARGS+=(--qos="${SLURM_QOS}")
+    [ -n "${SLURM_ACCOUNT:-}" ] && LAUNCHER_ARGS+=(--account="${SLURM_ACCOUNT}")
+    # --exclusive = whole-node allocation. On spur only sbatch accepts it (srun does
+    # not), so add it only in sbatch mode. Set SLURM_EXCLUSIVE=0 to disable.
+    # Spelled as an `if` rather than an `&&` chain because this is the last command
+    # of the else branch: under `set -e` a false `&&` chain there would make the
+    # whole compound command fail and abort the script.
+    if [ "${SLURM_LAUNCH_CMD:-srun}" = "sbatch" ] && [ "${SLURM_EXCLUSIVE:-1}" != "0" ]; then
+      LAUNCHER_ARGS+=(--exclusive)
+    fi
+  fi
+  # Each patch self-skips (exit 2) when its PRIMUS_* env gate is unset, so both
+  # can be passed unconditionally.
+  LAUNCHER_ARGS+=(-- --image "${DOCKER_IMAGE}" --clean --
+    --numa
+    --patch runner/helpers/patches/10_fix_libionic_abi4.sh
+    --patch runner/helpers/patches/11_fix_lld_stub.sh)
+fi
+
+./primus-cli "${LAUNCHER_ARGS[@]}" \
+  -- train pretrain --config "$EXP" \
+  --manual_gc True \
+  --manual_gc_interval 100 \
+  --pp_warmup "${PP_WARMUP:-True}" \
+  "${PP_LAYOUT_ARGS[@]}" \
+  --moe_router_force_load_balancing True \
+  "${MOE_FORCE_LB_ARGS[@]}" \
+  --log_avg_skip_iterations 3 \
+  --backend_path "$BACKEND_PATH" \
+  --num_layers "$PRIMUS_TOTAL_LAYERS" \
+  --train_iters "$TRAIN_ITERS" \
+  --lr_warmup_iters 0 \
+  --lr_decay_iters "$TRAIN_ITERS" \
+  --micro_batch_size "$MBS" \
+  --global_batch_size "$GBS" \
+  --seq_length "$PRIMUS_SEQ_LENGTH" \
+  --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \
+  --rope_type rope \
+  --tensor_model_parallel_size "$PRIMUS_TP" \
+  --pipeline_model_parallel_size "$PRIMUS_PP" \
+  --expert_model_parallel_size "$PRIMUS_EP" \
+  --num_experts "$PRIMUS_NUM_EXPERTS" \
+  --moe_router_topk "$PRIMUS_MOE_TOPK" \
+  --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \
+  --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \
+  --index_topk "$PRIMUS_INDEX_TOPK" \
+  --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \
+  --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \
+  --mtp_num_layers "${MTP_NUM_LAYERS:-0}" \
+  --mock_data True \
+  --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \
+  --use_turbo_attention "$USE_TURBO_ATTENTION" \
+  --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \
+  --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \
+  --use_v4_fp8_indexer "$USE_V4_FP8_INDEXER" \
+  --v4_indexer_distill_loss_coeff "$PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF" \
+  --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \
+  --use_turbo_deepep "$USE_TURBO_DEEPEP" \
+  "${TURBO_DEEPEP_CLI_ARGS[@]}" \
+  --use_turbo_mega_moe "$USE_TURBO_MEGA_MOE" \
+  "${CUDA_GRAPH_CLI_ARGS[@]}" \
+  --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \
+  --moe_use_legacy_grouped_gemm "$LEGACY_GG" \
+  "${OPTIMIZER_CLI_ARGS[@]}" \
+  --fp8 "$FP8" \
+  --fp8_recipe "$FP8_RECIPE" \
+  "${FP8_PARAM_GATHER_CLI_ARGS[@]}" \
+  --recompute_num_layers "$PRIMUS_RECOMPUTE_LAYERS" \
+  --recompute_granularity full \
+  --recompute_method block \
+  --overlap_grad_reduce "$PRIMUS_OVERLAP_GRAD_REDUCE" \
+  --overlap_param_gather "$PRIMUS_OVERLAP_PARAM_GATHER" \
+  --disable_last_saving True \
+  --disable_wandb True \
+  --disable_tensorboard True \
+  --profile "$PROFILE" \
+  --use_pytorch_profiler "$PROFILE" \
+  --profile_step_end 7 \
+  --profile_step_start 6 \
+  --bias_swiglu_fusion "$PRIMUS_BIAS_SWIGLU_FUSION" \
+  2>&1 | tee "$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt"
diff --git a/examples/deepseek-v4/run_deepseek_v4_flash.sh b/examples/deepseek-v4/run_deepseek_v4_flash.sh
new file mode 100755
index 000000000..4194086a8
--- /dev/null
+++ b/examples/deepseek-v4/run_deepseek_v4_flash.sh
@@ -0,0 +1,408 @@
+#!/bin/bash
+#
+# DeepSeek-V4-Flash pretraining on AMD Instinct MI355X.
+#
+# Six families of optimization, each behind one switch, all on by default, so
+# running this script with nothing but an image reproduces the best measured
+# configuration:
+#
+#   PRIMUS_OPT_FUSION     (1) small kernel fusions (RMSNorm, RoPE, Sinkhorn,
+#                             hyper-connections, compressor pooling, indexer
+#                             tail, router tail, grouped-weight stack, plus the
+#                             Megatron permute / CE / grad-accum fusions)
+#   PRIMUS_OPT_ATTENTION  (2) attention backend: FlyDSL sparse-MLA ("turbo")
+#   PRIMUS_OPT_DEEPEP     (3) DeepEP dispatch/combine + Turbo grouped GEMM
+#   PRIMUS_OPT_SYNC_FREE  (4) sync-free MoE stages
+#   PRIMUS_OPT_MEGA_MOE   (5) MegaMoE -- a REPLACEMENT for (3) and (4)
+#   PRIMUS_OPT_LAYOUT     (6) pipeline layout + recompute depth
+#
+# Set any of them to 0 to turn that family off. Individual knobs inside a family
+# can still be overridden one at a time: every export below is
+# `${VAR:-<value from the switch>}`, so the environment always wins.
+#
+# -----------------------------------------------------------------------------
+# Reproducing the speedup curve
+# -----------------------------------------------------------------------------
+# Each step keeps everything from the steps above it and adds one change. Set
+# these once, then run any step below unchanged:
+#
+#   export DOCKER_IMAGE=<primus-image>
+#   export SLURM_ALLOC_JOB_ID=<jobid>   # optional: join an allocation you hold
+#   export TRAIN_ITERS=10
+#   cd <primus-root>
+#
+# The throughput after each command is what a 4-node MI355X run measured at
+# 10 iterations with router load balancing forced to uniform (TFLOP/s per GPU).
+#
+# step 0 -- baseline, every optimization off
+#   PRIMUS_OPT_FUSION=0 PRIMUS_OPT_ATTENTION=0 PRIMUS_OPT_DEEPEP=0 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 1 -- + small kernel fusions
+#   PRIMUS_OPT_FUSION=1 PRIMUS_OPT_ATTENTION=0 PRIMUS_OPT_DEEPEP=0 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 2 -- attention to Gluon. The ATTENTION switch means FlyDSL, so the
+#           Gluon backend is named explicitly on both paths.
+#   PRIMUS_OPT_FUSION=1 PRIMUS_OPT_ATTENTION=0 PRIMUS_OPT_DEEPEP=0 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#   USE_V4_ATTENTION_BACKEND=gluon_v3 USE_V4_CSA_ATTENTION_BACKEND=gluon_v3 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 3 -- attention to FlyDSL sparse-MLA
+#   PRIMUS_OPT_FUSION=1 PRIMUS_OPT_ATTENTION=1 PRIMUS_OPT_DEEPEP=0 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 4 -- + DeepEP alone. The DEEPEP switch also turns on the Turbo grouped
+#           GEMM, so the GEMM is held off to measure the two apart.
+#   PRIMUS_OPT_FUSION=1 PRIMUS_OPT_ATTENTION=1 PRIMUS_OPT_DEEPEP=1 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#   TURBO_USE_GROUPED_MLP=False YAML_TURBO_GROUPED_GEMM=false \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 5 -- + Turbo grouped GEMM
+#   PRIMUS_OPT_FUSION=1 PRIMUS_OPT_ATTENTION=1 PRIMUS_OPT_DEEPEP=1 \
+#   PRIMUS_OPT_SYNC_FREE=0 PRIMUS_OPT_MEGA_MOE=0 PRIMUS_OPT_LAYOUT=0 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 6 -- MegaMoE, which replaces DeepEP and the grouped GEMM
+#   PRIMUS_OPT_MEGA_MOE=1 PRIMUS_OPT_LAYOUT=0 \
+#     bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# step 7 -- + tuned pipeline layout and no recompute. This is the default,
+#           so it needs no switches at all.
+#   bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+#
+# Steps 6 and 7 need no explicit PRIMUS_OPT_DEEPEP / PRIMUS_OPT_SYNC_FREE:
+# MegaMoE turns both off itself, since it replaces rather than stacks on them.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PRIMUS_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+
+# =============================================================================
+# Optimization switches
+# =============================================================================
+PRIMUS_OPT_FUSION=${PRIMUS_OPT_FUSION:-1}
+PRIMUS_OPT_ATTENTION=${PRIMUS_OPT_ATTENTION:-1}
+PRIMUS_OPT_DEEPEP=${PRIMUS_OPT_DEEPEP:-1}
+PRIMUS_OPT_SYNC_FREE=${PRIMUS_OPT_SYNC_FREE:-1}
+PRIMUS_OPT_MEGA_MOE=${PRIMUS_OPT_MEGA_MOE:-1}
+PRIMUS_OPT_LAYOUT=${PRIMUS_OPT_LAYOUT:-1}
+
+# (5) replaces (3) and (4) rather than stacking on them: MegaMoE builds its own
+# experts, brings its own all-to-all, and never reaches token_dispatcher or
+# grouped_experts, so DeepEP, the Turbo grouped GEMM and the sync-free stages
+# have nothing left to accelerate. run_deepseek_v4.sh already forces
+# USE_TURBO_DEEPEP=False when MegaMoE is on; the rest is turned off here so the
+# summary printed at the end matches what actually runs.
+if [ "$PRIMUS_OPT_MEGA_MOE" = 1 ]; then
+    PRIMUS_OPT_DEEPEP=0
+    PRIMUS_OPT_SYNC_FREE=0
+fi
+
+# =============================================================================
+# Cluster wiring -- site-specific, not an optimization
+# =============================================================================
+if command -v spur >/dev/null 2>&1; then
+    export PRIMUS_LAUNCHER=slurm
+    export SLURM_LAUNCH_CMD="${SLURM_LAUNCH_CMD:-srun}"
+    # Partition / QOS / account are site-specific and only forwarded when
+    # non-empty, so set them in the environment for your cluster.
+    export SLURM_PARTITION="${SLURM_PARTITION:-}"
+    export SLURM_QOS="${SLURM_QOS:-}"
+    export SLURM_ACCOUNT="${SLURM_ACCOUNT:-}"
+    # Empty = let the scheduler allocate nodes. An incoming SLURM_NODELIST is
+    # honored so callers can pin to specific known-good nodes.
+    export SLURM_NODELIST="${SLURM_NODELIST:-}"
+    # Path to an ABI-4 libionic provider .so to swap into the container at launch
+    # (fixes ionic RDMA on images whose bundled libionic only advertises uverbs
+    # ABI 1). tools/patches/fix_libionic_abi4.sh reads it; empty disables it.
+    export PRIMUS_LIBIONIC_SRC_ABI4_SO="${PRIMUS_LIBIONIC_SRC_ABI4_SO:-}"
+    export NCCL_DEBUG="${NCCL_DEBUG:-}"
+    export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-ens3}"
+    export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-ens3}"
+    # Registry login, needed only when the image is not public. Each node runs
+    # `docker login` before the pull, and sbatch --export=ALL propagates these.
+    # Empty means no login is attempted, which keeps credentials out of this
+    # git-tracked file.
+    export DOCKER_LOGIN_USER="${DOCKER_LOGIN_USER:-}"
+    export DOCKER_LOGIN_KEY="${DOCKER_LOGIN_KEY:-}"
+else
+    # dccs cluster. Partition / nodelist are not pinned here; export
+    # SLURM_PARTITION / SLURM_NODELIST to target specific hardware.
+    #
+    # Socket interface: run_deepseek_v4.sh falls back to `lo`, which leaves
+    # multi-node rendezvous hanging. The dccs front-end NIC is `fenic` (the RDMA
+    # devices are benic1p1..benic8p1 / ionic_0..7). runner/helpers/hooks/
+    # 10_auto_nccl_net.sh would auto-detect it, but only when these are unset, and
+    # the `lo` fallback sets them -- so pin them here.
+    export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-fenic}"
+    export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-fenic}"
+
+    # Optionally reuse one held allocation instead of queueing per run, so a sweep
+    # keeps landing on the same (known-good) nodes:
+    #   salloc --no-shell -N 4 --exclusive --partition=<partition> --mem=0
+    #   SLURM_ALLOC_JOB_ID=<granted id> bash examples/deepseek-v4/run_deepseek_v4_flash.sh
+    # srun then joins that allocation rather than asking for new nodes;
+    # --no-shell keeps it alive independently of any terminal, and
+    # `scancel <id>` releases it. Left unset by default: a job id is only valid
+    # for the life of its allocation.
+    SLURM_ALLOC_JOB_ID="${SLURM_ALLOC_JOB_ID:-}"
+    if [ -n "$SLURM_ALLOC_JOB_ID" ]; then
+        export SLURM_JOB_ID="${SLURM_JOB_ID:-$SLURM_ALLOC_JOB_ID}"
+        export SLURM_JOBID="${SLURM_JOBID:-$SLURM_ALLOC_JOB_ID}"
+        # Let the step share nodes the allocation already holds.
+        export SLURM_OVERLAP="${SLURM_OVERLAP:-1}"
+        # srun rejects a --nodelist that is not inside the allocation, so take the
+        # allocation's own list. Empty (allocation still pending) is fine; srun
+        # then just uses every node it was given.
+        export SLURM_NODELIST="${SLURM_NODELIST:-$(squeue -h -j "$SLURM_ALLOC_JOB_ID" -o '%N' 2>/dev/null || true)}"
+    fi
+    export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-fenic}"
+    export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-fenic}"
+fi
+
+# =============================================================================
+# Model shape and parallelism -- must stay identical across configurations
+# being compared
+# =============================================================================
+export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-43}
+export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256}
+export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6}
+export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048}
+export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512}
+export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-'[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]'}
+export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1}
+export NNODES=${NNODES:-4}
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-4}
+export PRIMUS_EP=${PRIMUS_EP:-8}
+
+export MBS=${MBS:-1}
+export GBS=${GBS:-$((64 * NNODES * MBS))}
+export TRAIN_ITERS=${TRAIN_ITERS:-10}
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}}
+
+# Force load balancing so the expert GEMM shapes are step-invariant. This keeps
+# a comparison between two configurations from being polluted by routing jitter;
+# `uniform` is closer to real routing than `even`, which hands every expert an
+# identical token count and flatters the unfused grouped-GEMM path.
+export MOE_FORCE_LB_TYPE=${MOE_FORCE_LB_TYPE:-uniform}
+
+# =============================================================================
+# (1) Small kernel fusions
+# =============================================================================
+# Each of these replaces a chain of small elementwise ops -- and the HBM round
+# trip per op -- with a single Triton kernel. Individually none is dramatic;
+# together they are the largest single step in the speedup curve.
+if [ "$PRIMUS_OPT_FUSION" = 1 ]; then
+    _F=1; _FY=true; _FB=True
+else
+    _F=0; _FY=false; _FB=False
+fi
+
+export PRIMUS_RMSNORM_TRITON=${PRIMUS_RMSNORM_TRITON:-$_F}                          # RMSNorm
+export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-$_F}                                # interleaved partial RoPE
+export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-$_F}                        # Sinkhorn-Knopp
+export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-$_F}                                    # hyper-connection glue
+export PRIMUS_HC_COLLAPSE_TRITON=${PRIMUS_HC_COLLAPSE_TRITON:-$_F}                  # hyper-connection collapse
+export PRIMUS_HC_EXPAND_TRITON=${PRIMUS_HC_EXPAND_TRITON:-$_F}                      # hyper-connection expand
+export PRIMUS_COMPRESS_POOL_TRITON=${PRIMUS_COMPRESS_POOL_TRITON:-$_F}              # compressor pooling
+export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-$_F}                          # indexer scoring tail
+export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-$_F}                      # MoE router tail
+export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-$_F}  # grouped weight stack
+export PRIMUS_INDEXER_FUSE_PROJ=${PRIMUS_INDEXER_FUSE_PROJ:-$_F}                    # indexer q/k projection fuse
+export PRIMUS_INDEXER_MASK_CACHE=${PRIMUS_INDEXER_MASK_CACHE:-$_F}                  # indexer causal-mask cache
+export PRIMUS_COMPRESS_ROPE_CACHE=${PRIMUS_COMPRESS_ROPE_CACHE:-$_F}                # compressed RoPE cache
+# Backward-side fusions inside the Triton/FlyDSL attention kernels.
+export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-$_F}
+export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-$_F}
+# An alternative, wider indexer fusion. Off by default even with (1) on: it
+# supersedes PRIMUS_INDEXER_TRITON and has not been the faster of the two here.
+export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0}
+# torch.compile path for Sinkhorn, an alternative to the Triton kernel above.
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-$_FB}
+
+# Fusions that live in the experiment yaml (see the derivation at the end).
+# use_turbo_rms_norm additionally needs the primus_turbo master gate, which only
+# comes on with (3) or (5).
+YAML_MOE_PERMUTE_FUSION=${YAML_MOE_PERMUTE_FUSION:-$_FY}
+YAML_CE_LOSS_FUSION=${YAML_CE_LOSS_FUSION:-$_FY}
+YAML_GRAD_ACC_FUSION=${YAML_GRAD_ACC_FUSION:-$_FY}
+YAML_TURBO_RMS_NORM=${YAML_TURBO_RMS_NORM:-$_FY}
+
+# =============================================================================
+# (2) Attention backend
+# =============================================================================
+# The dense/HCA path and the CSA path are selected separately because CSA's
+# indexer and top-k selection make it a different kernel problem. Backends, in
+# increasing order of speed: eager, triton_v1, triton_v2, gluon_v2, gluon_v3,
+# turbo (Primus-Turbo native FlyDSL sparse-MLA).
+#
+# `eager` cannot run at this scale: it materialises a [B, H, S, S] tensor, about
+# 16 GiB per microbatch per layer at V4-Flash dimensions, and fails to allocate
+# at any recompute setting. The runnable zero point is triton_v1.
+if [ "$PRIMUS_OPT_ATTENTION" = 1 ]; then _ATTN=turbo; else _ATTN=triton_v1; fi
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-$_ATTN}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-$_ATTN}
+# A separate flash-attention path that takes dispatch precedence over the V4
+# backend on dense layers; measured far slower here because it cannot do SWA.
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+# Keep the indexer QK in high precision. The indexer decides which compressed KV
+# entries each query attends to, so quantization error there changes the
+# selection itself rather than merely perturbing a value. It is also a fake
+# quant (quantize/dequantize around a BF16 GEMM), so leaving it off removes pure
+# overhead. Set True to opt back in for QAT experiments.
+export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-False}
+
+# =============================================================================
+# (3) DeepEP + Turbo grouped GEMM
+# =============================================================================
+# DeepEP moves the expert-parallel dispatch and combine into dedicated kernels
+# instead of PyTorch permutation around two all-to-alls. The Turbo grouped GEMM
+# issues the local experts' GEMMs as one ragged-batch kernel rather than a loop
+# of small ones -- measured separately, the grouped GEMM is worth several times
+# what DeepEP is.
+if [ "$PRIMUS_OPT_DEEPEP" = 1 ]; then _EP=True; _EPY=true; else _EP=False; _EPY=false; fi
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-$_EP}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-$_EP}
+YAML_TURBO_GROUPED_GEMM=${YAML_TURBO_GROUPED_GEMM:-$_EPY}
+export LEGACY_GG=${LEGACY_GG:-False}
+# run_deepseek_v4.sh only sets this when TURBO_USE_GROUPED_MLP=True but then
+# dereferences it unconditionally under `set -u`, so it must exist either way.
+export PRIMUS_BIAS_SWIGLU_FUSION=${PRIMUS_BIAS_SWIGLU_FUSION:-False}
+
+# =============================================================================
+# (4) Sync-free MoE
+# =============================================================================
+# Stage 1 fuses the router with the aux score and enables permutation fusion;
+# stage 2 additionally implies DeepEP and the grouped GEMM. Note that the fused
+# router drives Megatron's TopKRouter, while V4 uses its own learned router, so
+# on this model stage 1 reduces to permutation fusion alone.
+#
+# Stage > 1 with the grouped GEMM off raises ValueError before Megatron can
+# auto-enable it, which is why (4) implies (3) below.
+if [ "$PRIMUS_OPT_SYNC_FREE" = 1 ]; then _SF=1; _SFY=true; else _SF=0; _SFY=false; fi
+YAML_SYNC_FREE_STAGE=${YAML_SYNC_FREE_STAGE:-$_SF}
+YAML_FUSED_ROUTER=${YAML_FUSED_ROUTER:-$_SFY}
+
+# =============================================================================
+# (5) MegaMoE
+# =============================================================================
+# Fuses the expert-parallel all-to-all into the grouped GEMM itself, so the
+# ideal cost becomes max(comm, gemm) rather than their sum. EP-only: needs TP=1,
+# BF16 and an EP process group. Mutually exclusive with (3), which
+# run_deepseek_v4.sh enforces by forcing USE_TURBO_DEEPEP=False.
+if [ "$PRIMUS_OPT_MEGA_MOE" = 1 ]; then _MM=True; else _MM=False; fi
+export USE_TURBO_MEGA_MOE=${USE_TURBO_MEGA_MOE:-$_MM}
+
+# Master gate for every primus_turbo patch (DeepEP dispatcher, Turbo RMSNorm,
+# MegaMoE, Turbo grouped GEMM, sync-free auto-enable). run_deepseek_v4.sh turns
+# it on automatically when any turbo feature is requested; pinned here so the
+# state is explicit when every MoE optimization is off.
+if [ "$PRIMUS_OPT_DEEPEP" = 1 ] || [ "$PRIMUS_OPT_MEGA_MOE" = 1 ]; then
+    export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-True}
+else
+    export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-False}
+fi
+
+# =============================================================================
+# (6) Pipeline layout + recompute
+# =============================================================================
+# An even split is not balanced: the last stage also carries the MTP module and
+# the loss, while 1F1B leaves stage 0 holding the most microbatches in flight.
+# The tuned layout moves layers off the last stage onto the middle ones and
+# keeps stage 0 small. Recompute can then go to zero, but only because the
+# optimizations above have freed the memory for it -- dropping it first will
+# run out of memory.
+if [ "$PRIMUS_OPT_LAYOUT" = 1 ]; then
+    export PRIMUS_PP_LAYOUT="${PRIMUS_PP_LAYOUT:-Et*10|t*12|t*12|t*9mL}"
+    export PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-0}
+else
+    export PRIMUS_PP_LAYOUT="${PRIMUS_PP_LAYOUT:-Et*10|t*11|t*11|t*11mL}"
+    export PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-3}
+fi
+export RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full}
+export RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-block}
+
+# =============================================================================
+# Experiment config: derived yaml
+# =============================================================================
+# Seven of the switches above land on fields the shipped experiment yaml
+# hardcodes, and run_deepseek_v4.sh passes no CLI override for them, so no
+# environment variable can reach them. Derive a yaml rather than editing the
+# original, tagged so concurrent configurations never share a file.
+YAML_TAG=${YAML_TAG:-flash}
+SRC_YAML="$PRIMUS_ROOT/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml"
+GEN_YAML="$PRIMUS_ROOT/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-${YAML_TAG}.yaml"
+
+if [ ! -f "$SRC_YAML" ]; then
+    echo "[flash] ERROR: source config not found: $SRC_YAML" >&2
+    exit 1
+fi
+
+# Always regenerate: the same tag must never carry values from a previous run.
+{
+    echo "# GENERATED by examples/deepseek-v4/run_deepseek_v4_flash.sh -- do not edit."
+    echo "# switches: fusion=$PRIMUS_OPT_FUSION attention=$PRIMUS_OPT_ATTENTION deepep=$PRIMUS_OPT_DEEPEP"
+    echo "#           sync_free=$PRIMUS_OPT_SYNC_FREE mega_moe=$PRIMUS_OPT_MEGA_MOE layout=$PRIMUS_OPT_LAYOUT"
+    sed \
+        -e "s/^\( *\)turbo_sync_free_moe_stage: *[0-9]*/\1turbo_sync_free_moe_stage: $YAML_SYNC_FREE_STAGE/" \
+        -e "s/^\( *\)use_turbo_rms_norm: *\(true\|false\)/\1use_turbo_rms_norm: $YAML_TURBO_RMS_NORM/" \
+        -e "s/^\( *\)use_turbo_grouped_gemm: *\(true\|false\)/\1use_turbo_grouped_gemm: $YAML_TURBO_GROUPED_GEMM/" \
+        -e "s/^\( *\)moe_use_fused_router_with_aux_score: *\(true\|false\)/\1moe_use_fused_router_with_aux_score: $YAML_FUSED_ROUTER/" \
+        -e "s/^\( *\)moe_permute_fusion: *\(true\|false\)/\1moe_permute_fusion: $YAML_MOE_PERMUTE_FUSION/" \
+        -e "s/^\( *\)cross_entropy_loss_fusion: *\(true\|false\)/\1cross_entropy_loss_fusion: $YAML_CE_LOSS_FUSION/" \
+        -e "s/^\( *\)gradient_accumulation_fusion: *\(true\|false\)/\1gradient_accumulation_fusion: $YAML_GRAD_ACC_FUSION/" \
+        "$SRC_YAML"
+} > "$GEN_YAML"
+
+# Fail loudly rather than silently training a config that is not the one asked
+# for: a rename upstream would otherwise pass through unnoticed.
+for kv in "turbo_sync_free_moe_stage: $YAML_SYNC_FREE_STAGE" \
+          "use_turbo_rms_norm: $YAML_TURBO_RMS_NORM" \
+          "use_turbo_grouped_gemm: $YAML_TURBO_GROUPED_GEMM" \
+          "moe_use_fused_router_with_aux_score: $YAML_FUSED_ROUTER" \
+          "moe_permute_fusion: $YAML_MOE_PERMUTE_FUSION" \
+          "cross_entropy_loss_fusion: $YAML_CE_LOSS_FUSION" \
+          "gradient_accumulation_fusion: $YAML_GRAD_ACC_FUSION"; do
+    if ! grep -q "$kv" "$GEN_YAML"; then
+        echo "[flash] ERROR: '$kv' missing from $GEN_YAML -- did the source yaml change shape?" >&2
+        exit 1
+    fi
+done
+
+# Sync-free stage > 1 needs the grouped GEMM already on, or Megatron raises
+# before it gets a chance to auto-enable it.
+if [ "$YAML_SYNC_FREE_STAGE" -gt 1 ] && [ "$YAML_TURBO_GROUPED_GEMM" != "true" ] \
+   && [ "$TURBO_USE_GROUPED_MLP" != "True" ]; then
+    echo "[flash] ERROR: sync-free stage $YAML_SYNC_FREE_STAGE requires the grouped GEMM" >&2
+    exit 1
+fi
+
+export EXP="${EXP:-$GEN_YAML}"
+
+export PROFILE=${PROFILE:-False}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_nodes${NNODES}_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}}
+
+echo "[flash] (1) fusion=$PRIMUS_OPT_FUSION  (2) attention=$PRIMUS_OPT_ATTENTION  (3) deepep=$PRIMUS_OPT_DEEPEP  (4) sync_free=$PRIMUS_OPT_SYNC_FREE  (5) mega_moe=$PRIMUS_OPT_MEGA_MOE  (6) layout=$PRIMUS_OPT_LAYOUT"
+echo "[flash] attention=$USE_V4_ATTENTION_BACKEND/$USE_V4_CSA_ATTENTION_BACKEND turbo_gate=$ENABLE_PRIMUS_TURBO"
+echo "[flash] deepep=$USE_TURBO_DEEPEP grouped_gemm=$TURBO_USE_GROUPED_MLP mega_moe=$USE_TURBO_MEGA_MOE sync_free_stage=$YAML_SYNC_FREE_STAGE"
+echo "[flash] layout=${PRIMUS_PP_LAYOUT:-<even split>} recompute=$PRIMUS_RECOMPUTE_LAYERS"
+echo "[flash] nodes=$NNODES tp=$PRIMUS_TP pp=$PRIMUS_PP ep=$PRIMUS_EP gbs=$GBS seq=$PRIMUS_SEQ_LENGTH iters=$TRAIN_ITERS"
+echo "[flash] exp=$EXP"
+
+# Resolve the configuration and print it without launching anything, so a switch
+# combination can be checked before an allocation is spent on it.
+if [ "${PRIMUS_DRY_RUN:-0}" = 1 ]; then
+    echo "[flash] dry run: configuration resolved, not launching"
+    exit 0
+fi
+
+exec "${SCRIPT_DIR}/run_deepseek_v4.sh" 2>&1 | tee "train_flash_${YAML_TAG}.log"
diff --git a/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh b/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh
new file mode 100755
index 000000000..1552fd6ff
--- /dev/null
+++ b/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh
@@ -0,0 +1,280 @@
+#!/bin/bash
+###############################################################################
+# DeepSeek-V4 Flash perf PROXY runner (latest config: plan-5 P32 final).
+#
+# Wraps `run_deepseek_v4.sh` with a V4-Flash production-shape proxy:
+#
+#   - num_layers       8                        (vs production 43; PROXY)
+#   - hidden_size      4096                     (full V4-Flash; from yaml)
+#   - num_heads        64                       (full V4-Flash; from yaml)
+#   - head_dim         512                      (full V4-Flash; from yaml)
+#   - num_experts      256                      (full V4-Flash; PROXY-friendly
+#                                                32 experts/rank at EP=8)
+#   - moe_router_topk  6                        (full V4-Flash)
+#   - moe_ffn_hidden   2048                     (full V4-Flash)
+#   - index_topk       512                      (full V4-Flash CSA top-K)
+#   - compress_ratios  [0,0,4,128,4,128,4,0]    (8-layer slice exercising
+#                                                every layer kind: 3 cr=0,
+#                                                3 cr=4, 2 cr=128)
+#   - parallel         TP=1 PP=1 EP=8           (single-node 8 GPU)
+#   - seq_length       4096 (default)           (V4 pretrain target;
+#                                                set PRIMUS_SEQ_LENGTH=2048
+#                                                / 1024 / 512 on the
+#                                                command line if OOM)
+#
+# Plan-5 perf knobs default ON:
+#   - USE_V4_ATTENTION_BACKEND      (cr ∈ {0, 128} dense/HCA backend; default triton_v2)
+#   - USE_V4_CSA_ATTENTION_BACKEND  (cr == 4 CSA backend; default triton_v2)
+#   - USE_TURBO_DEEPEP              (PrimusTurboDeepEPTokenDispatcher)
+#   - TURBO_USE_GROUPED_MLP         (Turbo grouped-GEMM MoE expert path)
+#   - USE_V4_COMPILED_SINKHORN      (P29: torch.compile-fused Sinkhorn,
+#                                    kills the 7.6 s aten::sum fp32 reduce
+#                                    that dominated the P28 baseline)
+#
+# Plan-5 P32 final attention-kernel knobs (also default ON in code; surfaced
+# here for visibility / easy A/B):
+#   - PRIMUS_V4_ATTN_BWD_USE_SPLIT  (atomic-free split V4 attention BWD:
+#                                    dQ kernel + dK/dV kernel, each writes
+#                                    its own disjoint tiles via tl.store
+#                                    instead of atomic_add on a shared buf)
+#   - PRIMUS_V4_CSA_BWD_SEGREDUCE   (atomic-free CSA pool BWD via per-visit
+#                                    partial buffer + sorted inverse-index
+#                                    segmented reduction into dpool)
+#
+# These two relied on the **plan-5 P32 dual-RoPE bf16 cast fix** in
+# `apply_interleaved_partial_rope` (`primus/backends/megatron/core/transformer/
+# dual_rope.py`) to actually win in the proxy: pre-fix, cos/sin from
+# `position_ids.float() * inv_freq` was fp32, so `bf16 * fp32 = fp32`
+# silently upcast Q / K leaving RoPE — every V4 attention kernel paid 2x
+# HBM traffic and ran the slow fp32-specialised Triton binary, inflating
+# kernel times 1.8-7x in the proxy and masking the split / segreduce wins.
+# The one-line cast of cos/sin to `x.dtype` after the unsqueeze lets the
+# microbench-optimal kernels also win end-to-end. See
+# `deepseek-v4/develop/progress/p32/p32-summary.md` for the full
+# diagnostic walk-through.
+#
+# USE_TURBO_ATTENTION stays OFF — Turbo would take precedence over the V4
+# Triton dense path in `DeepseekV4Attention.forward` (plan-4 P27 dispatch
+# precedence: turbo > v4_triton > eager for cr ∈ {0, 128}).
+#
+# Steady-state perf (P32 final, mi355-gpu-8 / dev_primus_wenx_693,
+# iter 10 of 10, ${VAR:-DEFAULT} only):
+#
+#   iter time       :  603 ms / iter   (vs P28 baseline 8837 ms; 14.64x)
+#   TFLOP/s/GPU     :  1134            (vs P28 baseline 77.5)
+#   HBM peak / rank :  ~170 GiB
+#
+# Every override is `${VAR:-DEFAULT}`-guarded, so the caller can flip any
+# knob via `PRIMUS_SEQ_LENGTH=2048 ./run_deepseek_v4_flash_proxy.sh` etc.
+# without editing the script.
+#
+# Usage:
+#   ./run_deepseek_v4_flash_proxy.sh                                # 10-iter smoke
+#   TRAIN_ITERS=20 ./run_deepseek_v4_flash_proxy.sh                 # longer warmup pass
+#   PRIMUS_V4_ATTN_BWD_USE_SPLIT=0 ./run_deepseek_v4_flash_proxy.sh # fall back to
+#                                                                   #   monolithic V4 BWD
+#   PRIMUS_V4_CSA_BWD_SEGREDUCE=0 ./run_deepseek_v4_flash_proxy.sh  # fall back to
+#                                                                   #   gather+atomic CSA BWD
+#   PRIMUS_V4_DIAG_TIME=1 ./run_deepseek_v4_flash_proxy.sh          # dump per-call
+#                                                                   #   cuda.Event timings for
+#                                                                   #   v4_attention (rank 0)
+#
+# Profile is intentionally OFF in this script (this is the SMOKE / perf
+# runner, not the trace capture). For chrome-trace capture use
+# `deepseek-v4/develop/progress/p32/run_baseline_trace_ep8_p32_final.sh`
+# (mirrors the plan-4 P25 / plan-3 P23 profile-script pattern).
+###############################################################################
+set -euo pipefail
+
+# ---------- V4-Flash production widths (8-layer proxy slice) ----------------
+export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-8}
+export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256}
+export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6}
+export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048}
+export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512}
+# 8-layer slice — every V4 attention layer kind exercised:
+#   layer 0 / 1 : cr=0  (dense + SWA + sink)
+#   layer 2     : cr=4  (CSA)
+#   layer 3     : cr=128 (HCA)
+#   layer 4     : cr=4  (CSA)
+#   layer 5     : cr=128 (HCA)
+#   layer 6     : cr=4  (CSA)
+#   layer 7     : cr=0  (dense + SWA + sink)  -- V4-Flash production has
+#                                                cr=0 first/last layer
+export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-"[0,0,4,128,4,128,4,0]"}
+
+# ---------- Single-node EP=8 ------------------------------------------------
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-1}
+export PRIMUS_EP=${PRIMUS_EP:-8}
+
+# DP=8 with TP=1 PP=1 EP=8 on 8 GPUs (EP shards experts within DP group).
+# GBS=8, MBS=1 -> 1 microbatch / DP rank / iter.  Profiling-friendly:
+# minimises iter-to-iter variance + keeps activation memory bounded.
+export MBS=${MBS:-1}
+export GBS=${GBS:-8}
+
+# ---------- Production seq length target ------------------------------------
+# The CSA wrapper-side gather (plan-4 P26) materialises
+#   [B, H, Sq, K_topk, D] = [1, 64, Sq, 512, 512] * 2 bytes per microbatch
+# in HBM.  At Sq=4096 that is 64 GiB / microbatch on top of the 256-expert
+# MoE state (~12 GiB / rank for 8 layers) + KV cache + activations +
+# optimizer state — likely OOMs at MI355X (192 GiB HBM).  The plan-5 P28
+# task is to CALIBRATE this value (try 4096 -> 2048 -> 1024 -> 512) and
+# document the chosen value in `develop/profile/profile-baseline-ep8-*`.
+# Plan-5 P31 (in-kernel `topk_idxs` gather) is the structural fix that
+# eventually lets this default reach 4096.
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}}
+
+# ---------- Plan-5 perf knobs (all five ON) ---------------------------------
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2}
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True}
+# Plan-5 P29 (RESCOPED): torch.compile-fused HyperMixer Sinkhorn.  Kills
+# the 7.6 s aten::sum fp32 reduce (87.3 % of step time in the P28
+# baseline trace).  Default ON in the proxy after G32 + G33b are green.
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True}
+
+# Turbo attention OFF — would take precedence over V4 Triton dense path
+# in DeepseekV4Attention.forward (plan-4 P27 dispatch precedence:
+#   turbo > v4_triton > eager       for cr ∈ {0, 128}
+#   v4_triton_csa > eager           for cr == 4 ).
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+
+# ---------- Plan-5 P32 final attention-kernel knobs (split + segreduce) -----
+# Both default ON in the kernel code post-RoPE-fix; surface them here so
+# the proxy script self-documents the P32 final perf recipe and so a quick
+# A/B fallback is a single env-var flip. See header for the full root-cause
+# write-up.
+export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1}
+export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1}
+
+# ---------- Plan-6 elemwise-fusion knobs (default ON; A/B with =0) ----------
+# Plan-6 P40 close-out (2026-05-15): each plan-6 phase that wins the EP=8
+# proxy A/B adds its env knob here as default ON, mirroring the plan-5 P32
+# final precedent above.  Phases that microbench-win but proxy-noise-lose
+# (P38, P39) ship as default OFF with the kernel checked in for future
+# tuning.  Cumulative plan-6 win at this composition (P34..P37 ON, P38/P39
+# OFF): **-92.7 ms / iter (-15.4 %) vs plan-5 P32 final**; steady-state
+# iter time 510.6 ms / 524.9 TFLOP/s/GPU (peak HBM 172.3 GiB / rank).
+# See `deepseek-v4/develop/perf/proxy_ep8.md` row `P40 final` and
+# `progress/p40/p40-summary.md` for the full close-out write-up.
+#
+# The kernel code already defaults each to "1" / "0" appropriately; this
+# block makes the runner script self-document the recipe and lets users
+# flip individual fusions for A/B without editing source.
+#
+# P34 — stack_grouped_weight Triton FWD/BWD fusion in
+#   PrimusTurboGroupedMLP._stack_grouped_linear_weight.
+#   EP=8 proxy A/B win: 580.65 -> 530.85 ms / iter, -49.8 ms (-8.6%);
+#   TFLOP/s/GPU 463.2 -> 507.2, +9.5%; lm_loss bit-identical (pure
+#   layout transform). Default ON since 29baf151 (2026-05-14).
+export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1}
+
+# P35 — apply_interleaved_partial_rope Triton FWD/BWD fusion in
+#   dual_rope.py. Collapses the 9-op eager chain (slice + reshape +
+#   four broadcast muls + stack + reshape + cat) into one Triton
+#   kernel that does a single contiguous write with the rotation
+#   baked in.
+#   EP=8 proxy A/B win: 531.7 -> 526.7 ms / iter, -5.0 ms (-0.94%);
+#   TFLOP/s/GPU 507.1 -> 513.3, +1.2%; lm_loss bit-identical (pure
+#   analytic rotation). Default ON since landing (2026-05-14).
+export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1}
+
+# P36 — sinkhorn_normalize Triton FWD/BWD fusion in
+#   hyper_connection.py.  Replaces the plan-5 P29 ``torch.compile``
+#   cached Sinkhorn body with a hand-rolled Triton kernel that runs
+#   the 1 + 2*(n_iters - 1) alternating row/col normalize trajectory
+#   in registers per row of the leading axis (V4-Flash uses K=4).
+#   Microbench at V4-Flash K=4 (B=1, S=4096):
+#     FWD 0.045 ms (vs eager 0.600 ms = 13.4x; vs P29 compiled 0.270
+#                   ms = 6.0x)
+#     BWD 0.105 ms (vs eager 1.520 ms = 14.5x; vs P29 compiled 0.628
+#                   ms = 6.0x)
+#   The compiled-region overhead (`Torch-Compiled Region` ~21 ms / 16
+#   calls + `CompiledFunctionBackward` ~41 ms / 16 calls) is removed
+#   entirely.  Default ON since landing (2026-05-14).
+export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1}
+
+# P37 — HyperConnection compute_weights elemwise tail Triton fusion in
+#   hyper_connection.HyperMixer.compute_weights.  Fuses the 3 slices +
+#   3 fused-multiply-adds + 2 sigmoid + 1 softmax + 2 eps adds (the
+#   post-_packed_logits, pre-Sinkhorn chain) into one FWD + one BWD
+#   Triton kernel.  Microbench at V4-Flash K=4 (B=1, S=4096):
+#     FWD 0.044 ms (vs eager 0.102 ms = 2.34x)
+#     BWD 0.276 ms (vs eager 0.405 ms = 1.47x)
+#   The matmul inside _packed_logits stays as F.linear; collapse / expand
+#   (matmul-adjacent) stay eager too -- they are not net wins as
+#   separate Triton kernels.  Default ON since landing (2026-05-14).
+export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1}
+
+# P41 — Indexer.forward post-einsum tail Triton fusion.  Re-attempt
+#   of P38 that keeps the cuBLAS / hipBLASLt einsum eager and fuses
+#   only the bandwidth-bound tail (`relu + mul(w_i) + sum(H) +
+#   causal_mask`).
+#
+#   Plan-8 P57 close-out 2 (2026-05-15): default flipped to ON.
+#   Microbench at V4-Flash widths is a clear positive
+#   (FWD 4.30x / BWD 1.63x); the EP=8 proxy A/B (10-iter smoke)
+#   showed ~0.2 ms / iter aggregate gain within the ±1 ms noise band
+#   (small but consistently positive).  We default ON so the
+#   bandwidth-bound tail is fused by default; set
+#   PRIMUS_INDEXER_TRITON=0 to revert to the eager body.
+#
+#   The env knob was re-purposed at P41: it now controls the
+#   post-einsum tail path.  Legacy P38 full-fuse path lives behind
+#   PRIMUS_INDEXER_TRITON_FULL (default OFF, kept in tree for small-
+#   shape paths and future tuning).
+export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1}
+export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0}
+
+# P39 — V4 Router post-logits Triton FWD/BWD fusion (shared by topk +
+#   hash router).
+#
+#   Plan-8 P57 close-out 2 (2026-05-15): default flipped to ON.
+#   Microbench at V4-Flash widths (N=4096, E=256, K=8) wins on V4's
+#   production `sqrtsoftplus` score function (1.56x FWD / 1.22x BWD).
+#   P39 / P43 EP=8 proxy A/B was inside the proxy noise band, so the
+#   conservative landing posture left it default-OFF; for P57 R2 we
+#   default ON to keep the microbench-positive kernel on the production
+#   path.  Set PRIMUS_V4_ROUTER_TRITON=0 to revert to the eager body.
+export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1}
+
+# ---------- Precision: FP8 training (paper ue8m0 microscaling) --------------
+# V4-Flash trains in FP8 by default. PRECISION_TYPE=FP8 makes run_deepseek_v4.sh
+# emit --fp8 e4m3 --fp8_recipe mxfp8 (mxfp8 = the paper's ue8m0 microscaling,
+# E8M0 block scale, native on MI355X/CDNA4). The FP4 expert / FP4-Indexer path
+# is not yet wired in the Primus V4 integration ("Phase 2"), so experts run FP8
+# here (FP8-everywhere) rather than FP4 — the closest supported step. FP8 is
+# outlier-sensitive, hence the clamped SwiGLU (swiglu_limit) in the EXP yaml.
+# A/B back to BF16 with PRECISION_TYPE=BF16; override the recipe via FP8_RECIPE.
+export PRECISION_TYPE=${PRECISION_TYPE:-FP8}
+export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml}
+
+# ---------- Profile OFF in the proxy smoke runner ---------------------------
+# This script is the steady-state perf / smoke runner — kineto profiling
+# stays OFF to avoid contaminating the iter timer with profiler-collection
+# overhead. For chrome-trace capture use
+# `deepseek-v4/develop/progress/p32/run_baseline_trace_ep8_p32_final.sh`.
+export PROFILE=${PROFILE:-False}
+
+# ---------- Bookkeeping -----------------------------------------------------
+# Distinguish the proxy run output dir from the smoke run output dir so the
+# trace-capture script + the smoke run land side-by-side without clobbering
+# each other's logs.
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}}
+
+# ---------- Launcher: single-node in-container by default -------------------
+# This proxy is the single-node 8-GPU smoke/perf runner, normally invoked from
+# INSIDE the training container. run_deepseek_v4.sh defaults PRIMUS_LAUNCHER to
+# `slurm` (which needs `srun` from a SLURM allocation and would fail in a bare
+# container). Default to `direct` here so the proxy just torchruns locally;
+# override with PRIMUS_LAUNCHER=slurm when launching from a cluster login node.
+export PRIMUS_LAUNCHER=${PRIMUS_LAUNCHER:-direct}
+
+# Defer to run_deepseek_v4.sh for the actual training launch — every
+# CLI flag and the primus-cli invocation lives there.
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "${SCRIPT_DIR}/run_deepseek_v4.sh"
diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh
new file mode 100755
index 000000000..cea44a134
--- /dev/null
+++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh
@@ -0,0 +1,361 @@
+#!/bin/bash
+###############################################################################
+# DeepSeek-V4 *Pro* single-node bring-up with the Muon optimizer.
+#
+# Follows the DeepSeek-V4 paper Pro recipe (§4.2.1 architecture + §4.2.2
+# training setup) as closely as a single 8x288GB node allows:
+#   1. Model      : deepseek_v4_pro  (61L / d7168 / 384 experts — paper §4.2.1).
+#                   Selected via PRIMUS_MODEL, consumed by the
+#                   `model: ${PRIMUS_MODEL:...}.yaml` line in the EXP yaml.
+#                   Widths come from primus/configs/models/megatron/
+#                   deepseek_v4_pro.yaml; we only override the shape knobs the
+#                   runner exposes (layers/experts/topk/ffn/index_topk/
+#                   compress_ratios) so the runner's smoke defaults don't win.
+#   2. Reduced    : 61 layers + seq 4096 do NOT fit one node, so cut depth
+#      to fit       (PRIMUS_TOTAL_LAYERS) and seq (PRIMUS_SEQ_LENGTH); full/
+#                   uniform recompute on. (CSA-layer gather scales with seq.)
+#   3. Optimizer  : Muon (paper §4.2.2: Muon for matrices, AdamW for embedding/
+#                   pred-head/RMSNorm — in-tree ChainedOptimizer). Plain `muon`
+#                   requires overlap_{grad_reduce,param_gather}=False and
+#                   use_distributed_optimizer=False (Megatron asserts).
+#   4. Training   : paper §4.2.2 hyperparameters — momentum 0.95, update-RMS
+#      setup        0.18 (muon_extra_scale_factor), AdamW eps 1e-20, LR
+#                   2.0e-4→2.0e-5, balance-loss 1e-4, and (crucially) a LARGE
+#                   batch via gradient accumulation toward the paper's 94.4M
+#                   tokens/step. The batch is what amortizes the fixed Muon
+#                   Newton-Schulz cost: at GBS=8 (accum 1) NS looks like ~97%
+#                   of GEMM (a starved-batch artifact, NOT a Muon bug); at the
+#                   paper batch it falls to the reported ~1-3%.
+#
+# Single-node integration gaps vs the paper (not config — Primus V4 TODO):
+#   - MTP depth 1 (MultiTokenPredictionLayer unsupported) -> MTP_NUM_LAYERS=0
+#   - expert-bias/noaux_tc (Megatron needs sigmoid; V4 uses sqrtsoftplus) -> off
+#   - muon_weight_decay (0.01 vs paper 0.1) / mtp_loss_scaling are yaml-only.
+#
+# Usage:
+#   # paper-faithful single-node run (validated: ~890 TFLOP/s/GPU, Muon ~1% GPU):
+#   PRIMUS_TOTAL_LAYERS=2 PRIMUS_COMPRESS_RATIOS="[128,128]" \
+#       PRIMUS_SEQ_LENGTH=4096 GBS=256 ./run_deepseek_v4_pro_muon.sh
+#   PRIMUS_TOTAL_LAYERS=4 PRIMUS_SEQ_LENGTH=512 GBS=8 \
+#       ./run_deepseek_v4_pro_muon.sh                              # cheap validation
+#   OPTIMIZER=adam ./run_deepseek_v4_pro_muon.sh                   # A/B vs AdamW
+#   PRECISION_TYPE=BF16 ./run_deepseek_v4_pro_muon.sh             # A/B vs BF16 (fp8 is default-on)
+#   PROFILE=True DISABLE_TENSORBOARD=False ...                     # capture 1-step trace
+#
+# Precision: FP8 training is ON by default (FP8=e4m3, FP8_RECIPE=tensorwise).
+# Paper recipe is ue8m0/mxfp8 but it's not runnable on this gfx950 build (see the
+# "FP8 training" block below); tensorwise gives the paper's fp8 layout on the
+# weight GEMMs. PRECISION_TYPE=BF16 to A/B back to bf16.
+###############################################################################
+set -euo pipefail
+set -x
+
+export HF_TOKEN="${HF_TOKEN:-}"
+
+export NNODES=${PET_NNODES:-1}
+export TRAIN_ITERS=${TRAIN_ITERS:-10}
+export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1}
+
+# ---------- Model: DeepSeek-V4 Pro -----------------------------------------
+export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro}
+export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml}
+
+# ---------- Pro production widths (paper §4.2.1) ----------------------------
+export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-384}
+export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6}
+# Megatron only supports aux-loss-free expert bias with the sigmoid score
+# function; V4 uses sqrtsoftplus, so disable expert bias (matches the working
+# run_deepseek_v4.sh smoke; balancing falls back to seq_aux_loss).
+export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False}
+export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072}
+export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-1024}
+
+# ---------- Reduced depth + seq to fit single node -------------------------
+# MEASURED CEILING (chi2774, 8x288GB, 2026-05-20): with full Pro width
+# (384 experts) + Muon, *4 layers @ seq 512 = 268.8 GB/rank (93%)* is about
+# the single-node max.  The binding cost is weights + Muon's fp32 optimizer
+# states (Muon forces use_precision_aware_optimizer=False, so states cannot
+# be bf16) — NOT activations, so lowering seq does not buy more layers.
+# 5+ layers OOMs; for more depth use fewer experts or multi-node.
+export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-4}
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-512}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}}
+
+# Per-layer compression schedule, length == PRIMUS_TOTAL_LAYERS, mirroring the
+# Pro pattern: first two HCA(128), then CSA(4)/HCA(128) interleaved, last
+# dense+SWA(0).  (Pro full yaml: idx0,1=128; idx>=2 even=4 / odd=128; last=0.)
+export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-$(python3 - "$PRIMUS_TOTAL_LAYERS" <<'PY'
+import sys
+n=int(sys.argv[1])
+r=[]
+for i in range(n):
+    if i<2: r.append(128)
+    elif i==n-1: r.append(0)
+    else: r.append(4 if i%2==0 else 128)
+print("["+",".join(map(str,r))+"]")
+PY
+)}
+
+# ---------- Single-node EP=8 -----------------------------------------------
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-1}
+export PRIMUS_EP=${PRIMUS_EP:-8}
+export MBS=${MBS:-1}
+# Paper Pro batch = 94.4M tokens/step (batch-size schedule). On one node we
+# approach that regime via GRADIENT ACCUMULATION: grad_accum = GBS/(MBS*DP),
+# DP=8 here. Large GBS is what amortizes the fixed Muon Newton-Schulz cost over
+# many fwd/bwd microbatches (GBS=8 ⇒ accum=1 ⇒ NS runs every step ⇒ NS looks
+# like ~97% of GEMM; that was a starved-batch artifact, NOT a Muon bug). At
+# GBS=512 (accum 64) × seq 4096 = ~2.1M tokens/step the optimizer share drops
+# toward the paper's reported 1-3%.
+export GBS=${GBS:-512}
+
+# ---------- Optimizer: Muon (paper §4.2.2, values for BOTH Flash & Pro) ------
+export OPTIMIZER=${OPTIMIZER:-muon}
+# The Primus Muon path (primus/backends/megatron/core/optimizer/moun.py) needs
+# the emerging_optimizers package, which is NOT bundled in the container. Set
+# PRIMUS_INSTALL_EMERGING_OPTIMIZERS so the in-container install hook
+# (runner/.../01_install_emerging_optimizers.sh) provisions the pinned commit.
+# Gated on a Muon optimizer so an OPTIMIZER=adam A/B run pays nothing.
+if [ "$OPTIMIZER" = "muon" ] || [ "$OPTIMIZER" = "dist_muon" ]; then
+  export PRIMUS_INSTALL_EMERGING_OPTIMIZERS=${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-1}
+fi
+export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95}            # paper momentum 0.95
+# Paper: "rescale the RMS of each update matrix to 0.18 for reutilization of the
+# AdamW learning rate."  With scale_mode=spectral the realized update RMS ≈
+# extra_scale_factor (orth_grad RMS≈1/√max(m,n), times spectral scale √max(m,n)),
+# so 0.18 maps directly here.  Megatron default is 1.0 (≈5.5× too large vs paper).
+export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18}
+# Newton-Schulz hardening knobs (matter for mxfp8, where NS can diverge on the
+# quant-noised gradient). num_ns_steps = NS iterations (more = better convergence
+# on ill-conditioned input); fp32_matmul_prec = precision of the NS matmuls
+# ("medium" = tf32-ish, "high" = full fp32 — full precision keeps a near-σ=1
+# input from being pushed past the quintic's stable region by rounding error).
+export MUON_NUM_NS_STEPS=${MUON_NUM_NS_STEPS:-5}
+export MUON_FP32_MATMUL_PREC=${MUON_FP32_MATMUL_PREC:-medium}
+# NOTE: muon_weight_decay is yaml-only (trainer_base.yaml=0.01); paper=0.1.
+# No CLI flag, so it stays 0.01 here unless overridden in an EXP yaml.
+# Muon hard requirements (Megatron arguments.py:1422):
+export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False}
+export USE_PRECISION_AWARE_OPTIMIZER=${USE_PRECISION_AWARE_OPTIMIZER:-False}
+
+# ---------- Paper §4.2.2 Pro training hyperparameters -----------------------
+export LR=${LR:-2.0e-4}                       # Pro peak LR (Flash exp yaml had 1e-5)
+export MIN_LR=${MIN_LR:-2.0e-5}               # Pro end LR
+export ADAM_EPS=${ADAM_EPS:-1.0e-20}          # paper AdamW eps (NOTE: needs decimal point — Primus parses "1e-20" as a string)
+export MOE_AUX_LOSS_COEFF=${MOE_AUX_LOSS_COEFF:-0.0001}  # paper balance-loss weight
+# Paper MTP depth = 1, but the Primus V4 integration does NOT yet support the
+# MTP layer ("Unsupported mtp_model_layer submodules type ... when instantiating
+# MultiTokenPredictionLayer"), so default 0 here. Set =1 once V4 MTP lands.
+export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-0}
+
+# ---------- Perf knobs (V4 Triton attn + Turbo MoE; same family as proxy) --
+export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-True}
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True}
+# Primus Sync-Free MoE (eliminates the DeepEP host busy-wait on the variable
+# per-expert token counts): 0=off, 1=fused router/permute, 2=+no CPU busy-wait
+# (turbo deepep + grouped mlp), 3=fully sync-free (+fused act). Stage >=2 needs
+# use_turbo_grouped_gemm=True. Auto-enables the required sub-flags. Default 0.
+export TURBO_SYNC_FREE_MOE_STAGE=${TURBO_SYNC_FREE_MOE_STAGE:-0}
+# Phase 1b: route the dense/attention projections (q_down/kv/o_a etc.) through
+# Primus-Turbo linears so they pick up the mxfp8 (CK) path under the fp8 context.
+# Default OFF (attention stays bf16, the validated baseline). Set True to enable
+# fp8 attention/dense projections. Requires TP=1 and fp8_recipe in {tensorwise,
+# blockwise,mxfp8}; the MLA monkey-patch is auto-skipped for V4 (see mla_patches).
+export USE_TURBO_PARALLEL_LINEAR=${USE_TURBO_PARALLEL_LINEAR:-False}
+# Per-module recipe (paper): routed experts in MXFP4 while the rest of the layer
+# stays FP8. Works under the global FP8 recipe (no --fp4/--fp8 conflict): the
+# PrimusTurbo grouped MLP routes expert GEMMs through native FP4 (hipBLASLt).
+# Default OFF. When on, force the hipBLASLt FP4 backend (no AITER).
+export MOE_EXPERTS_FP4=${MOE_EXPERTS_FP4:-False}
+if [ "$MOE_EXPERTS_FP4" = "True" ]; then
+  export PRIMUS_TURBO_GEMM_BACKEND=${PRIMUS_TURBO_GEMM_BACKEND:-FP4:HIPBLASLT}
+fi
+# Phase 5 (paper): CSA-indexer QK score in FP4. Rounds q_i and K^{IComp} to
+# MXFP4 before the QK product (STE backward); w_i + ReLU/sum tail stay BF16.
+# Read directly by the Indexer via PRIMUS_INDEXER_FP4. Default OFF.
+export INDEXER_FP4=${INDEXER_FP4:-False}
+if [ "$INDEXER_FP4" = "True" ]; then
+  export PRIMUS_INDEXER_FP4=1
+  # The indexer QK now runs a REAL MXFP4 gemm (pt.ops.gemm_fp4) — force the
+  # hipBLASLt FP4 backend (no AITER), same as the MXFP4 expert path.
+  export PRIMUS_TURBO_GEMM_BACKEND=${PRIMUS_TURBO_GEMM_BACKEND:-FP4:HIPBLASLT}
+fi
+# MXFP8 expert-weight caching: expert weights are constant within an optimizer
+# step, so re-quantizing them every microbatch + recompute forward (the large
+# _mxfp8_quant_weight_fwd kernel) is redundant. When on, PrimusTurboGroupedMLP
+# prequantizes once per step and reuses the fp8 buffers (loss-neutral, faster).
+# Costs extra bytes/param resident — watch HBM at depth. Only affects the
+# mxfp8 (MX_BLOCKWISE) grouped path. Default OFF.
+export CACHE_MXFP8_WEIGHT=${CACHE_MXFP8_WEIGHT:-False}
+if [ "$CACHE_MXFP8_WEIGHT" = "True" ]; then
+  export PRIMUS_TURBO_CACHE_MXFP8_WEIGHT=1
+fi
+# FP8 attention projections (paper recipe): route q-up / o-proj through the fp8
+# turbo linear instead of the bf16 gather/scatter native path. Only valid at
+# TP=1 (gather/scatter are no-ops there); for TP>1 the turbo linear rejects
+# gather_output/scatter-input and it stays bf16. Default OFF.
+export V4_FP8_ATTN_PROJ=${V4_FP8_ATTN_PROJ:-False}
+if [ "$V4_FP8_ATTN_PROJ" = "True" ]; then
+  export PRIMUS_V4_FP8_ATTN_PROJ=1
+fi
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2}
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False}
+export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True}
+
+# ---------- FP8 training (paper §4.x quantization / techblog §9.6) ----------
+# Paper recipe: "FP4 + FP8 Mixed" — MoE experts + CSA-Indexer QK in FP4 (MXFP4),
+# EVERYTHING ELSE in FP8, all with the **ue8m0** microscaling scale format.
+# On this stack the ue8m0 path is `fp8_recipe=mxfp8` → TE MXFP8BlockScaling →
+# Primus-Turbo MX_BLOCKWISE granularity with scale_dtype=E8M0 (fp8_utils.py:148),
+# i.e. the paper's exact scaling format, native on MI355X/CDNA4.
+#
+# Integration gap vs the paper (NOT config — Primus V4 TODO, develop techblog
+# item 10 "Phase 2 FP4/FP8 Mixed"): the FP4 expert / FP4-Indexer path is not yet
+# wired in V4, so experts run at FP8 here (FP8 everywhere) rather than FP4. This
+# is the closest supported step toward the paper recipe. FP8 is highly outlier-
+# sensitive, which is why the paper pairs it with clamped SwiGLU (swiglu_limit,
+# already on above via PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=True).
+# Precision toggle — shared interface with run_deepseek_v4.sh / the flash proxy:
+#   PRECISION_TYPE=FP8 (default) -> e4m3 + tensorwise; BF16 -> fp8 off.
+# FP8 / FP8_RECIPE still override directly (e.g. FP8_RECIPE=blockwise, or FP8=null).
+export PRECISION_TYPE=${PRECISION_TYPE:-FP8}
+if [ "$PRECISION_TYPE" = "FP8" ]; then
+  export FP8=${FP8:-e4m3}                      # forward fp8 format (paper E4M3); "hybrid" = E4M3 fwd / E5M2 bwd
+  # Paper recipe is ue8m0 microscaling (mxfp8), but mxfp8 is NOT runnable on this
+  # gfx950 build (turbo grouped-GEMM has no MX path; TE ROCm MXFP8 needs K%128==0,
+  # V4 has a K=224 proj). `tensorwise` is the working recipe — paper fp8 layout,
+  # non-ue8m0 scale. (other: blockwise [TE-ROCm unsupported] / delayed)
+  export FP8_RECIPE=${FP8_RECIPE:-tensorwise}
+  # mxfp8 (paper ue8m0) + Muon: TRAINS, but shows a transient early-training
+  # grad-norm spike. Earlier 8-iter runs only caught the spike and mislabelled it
+  # a divergence; a 40-iter run shows it SELF-HEALS and the loss descends cleanly.
+  #   What's happening: mxfp8 quant noise ill-conditions the gradient early (random
+  #   init), so Muon's Newton-Schulz update norm spikes for ~10-20 iters, then
+  #   settles as the model organizes. RAW grads stay ~0.99 throughout; only the
+  #   NS-orthogonalized UPDATE norm spikes (Muon-specific: Adam@same-config is flat).
+  #   It is NOT a kernel bug (both MX GEMMs ~4% correct on REAL E2E inputs via
+  #   capture-replay; error zero-mean).
+  #   FIX (verified, L4/64-expert/GBS512/seq128, 40 iters):
+  #     no warmup  -> loss 12->0.80, grad-norm peak ~2.4e5 (settles to ~35)
+  #     warmup=10  -> loss 12->0.44, grad-norm peak ~2.1e4 (12x lower), no NaN
+  #   So LR warmup tames the transient AND improves the loss — and it is what the
+  #   paper does. We therefore AUTO-ENABLE warmup for mxfp8 (default 10; override
+  #   LR_WARMUP_ITERS). Two more NS-hardening knobs are exposed if needed:
+  #   MUON_NUM_NS_STEPS (more iters) and MUON_FP32_MATMUL_PREC=high. tensorwise
+  #   stays the conservative default (smooth per-tensor scale, no transient).
+  #   NOTE: validated at reduced depth/width; full 61-layer/384-expert is a
+  #   separate multi-GPU confirmation.
+  if [ "$FP8_RECIPE" = "mxfp8" ]; then
+    export LR_WARMUP_ITERS=${LR_WARMUP_ITERS:-10}
+    echo "[INFO] FP8_RECIPE=mxfp8 + Muon: expect a transient early grad-norm spike" >&2
+    echo "       (self-heals; LR warmup auto-set to ${LR_WARMUP_ITERS} to damp it)." >&2
+  fi
+else
+  export FP8=${FP8:-null}                  # PRECISION_TYPE=BF16 -> disable fp8
+  export FP8_RECIPE=${FP8_RECIPE:-null}
+fi
+
+TURBO_DEEPEP_CLI_ARGS=()
+if [ "$USE_TURBO_DEEPEP" = "True" ]; then
+  export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-80}
+  export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False}
+  export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32}
+  export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False}
+  TURBO_DEEPEP_CLI_ARGS=(
+    --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU"
+    --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM"
+    --moe_router_dtype "$MOE_ROUTER_DTYPE"
+    --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP"
+  )
+fi
+
+export PROFILE=${PROFILE:-False}
+# Profiler writes the chrome trace via tensorboard_trace_handler(args.tensorboard_dir),
+# so tensorboard must be enabled for a trace run.  Default True (smoke); set
+# DISABLE_TENSORBOARD=False together with PROFILE=True to capture a trace.
+export DISABLE_TENSORBOARD=${DISABLE_TENSORBOARD:-True}
+export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"}
+export PRIMUS_TEAM=${PRIMUS_TEAM:-amd}
+export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_pro_muon_L${PRIMUS_TOTAL_LAYERS}_seq${PRIMUS_SEQ_LENGTH}_ep${PRIMUS_EP}}
+
+if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then
+  echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH"
+  echo "Run: git submodule update --init --recursive"
+  exit 1
+fi
+
+mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME"
+
+./primus-cli direct \
+  -- train pretrain --config "$EXP" \
+  --backend_path "$BACKEND_PATH" \
+  --manual_gc True \
+  --manual_gc_interval 100 \
+  --num_layers "$PRIMUS_TOTAL_LAYERS" \
+  --train_iters "$TRAIN_ITERS" \
+  --lr_warmup_iters "${LR_WARMUP_ITERS:-0}" \
+  --lr_decay_iters "$TRAIN_ITERS" \
+  --micro_batch_size "$MBS" \
+  --global_batch_size "$GBS" \
+  --lr "$LR" \
+  --min_lr "$MIN_LR" \
+  --adam_eps "$ADAM_EPS" \
+  --moe_aux_loss_coeff "$MOE_AUX_LOSS_COEFF" \
+  --seq_length "$PRIMUS_SEQ_LENGTH" \
+  --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \
+  --rope_type rope \
+  --tensor_model_parallel_size "$PRIMUS_TP" \
+  --pipeline_model_parallel_size "$PRIMUS_PP" \
+  --expert_model_parallel_size "$PRIMUS_EP" \
+  --num_experts "$PRIMUS_NUM_EXPERTS" \
+  --moe_router_topk "$PRIMUS_MOE_TOPK" \
+  --moe_router_force_load_balancing "${MOE_FORCE_LOAD_BALANCE:-False}" \
+  --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \
+  --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \
+  --index_topk "$PRIMUS_INDEX_TOPK" \
+  --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \
+  --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \
+  --mtp_num_layers "$MTP_NUM_LAYERS" \
+  --mock_data True \
+  --optimizer "$OPTIMIZER" \
+  --muon_momentum "$MUON_MOMENTUM" \
+  --muon_extra_scale_factor "$MUON_EXTRA_SCALE_FACTOR" \
+  --muon_num_ns_steps "$MUON_NUM_NS_STEPS" \
+  --muon_fp32_matmul_prec "$MUON_FP32_MATMUL_PREC" \
+  --use_distributed_optimizer "$USE_DISTRIBUTED_OPTIMIZER" \
+  --use_precision_aware_optimizer "$USE_PRECISION_AWARE_OPTIMIZER" \
+  --main_grads_dtype fp32 \
+  --exp_avg_dtype fp32 \
+  --exp_avg_sq_dtype fp32 \
+  --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \
+  --use_turbo_attention "$USE_TURBO_ATTENTION" \
+  --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \
+  --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \
+  --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \
+  --use_turbo_deepep "$USE_TURBO_DEEPEP" \
+  --turbo_sync_free_moe_stage "$TURBO_SYNC_FREE_MOE_STAGE" \
+  "${TURBO_DEEPEP_CLI_ARGS[@]}" \
+  --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \
+  --use_turbo_gemm "$USE_TURBO_PARALLEL_LINEAR" \
+  --moe_experts_fp4 "$MOE_EXPERTS_FP4" \
+  --moe_use_legacy_grouped_gemm False \
+  --fp8 "$FP8" \
+  --fp8_recipe "$FP8_RECIPE" \
+  --recompute_num_layers 1 \
+  --recompute_granularity full \
+  --recompute_method uniform \
+  --overlap_grad_reduce False \
+  --overlap_param_gather False \
+  --disable_last_saving True \
+  --disable_wandb True \
+  --disable_tensorboard "$DISABLE_TENSORBOARD" \
+  --profile "$PROFILE" \
+  --use_pytorch_profiler "$PROFILE" \
+  --profile_step_end 7 \
+  --profile_step_start 6 \
+  2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt"
diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh
new file mode 100755
index 000000000..d30f05069
--- /dev/null
+++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh
@@ -0,0 +1,497 @@
+#!/bin/bash
+###############################################################################
+# DeepSeek-V4 *Pro* + Muon single-GPU bring-up on mi455 / gfx1250 (1 GPU).
+#
+# Single-GPU sibling of run_deepseek_v4_flash_proxy_1gpu.sh, for the Pro model.
+# The upstream run_deepseek_v4_pro_muon.sh uses `primus-cli direct`, which
+# assumes it is ALREADY inside the 8x288GB MI355X container at EP=8. This host
+# is one gfx1250 box with no SLURM and one GPU, so this script instead wraps
+# the SAME examples/run_pretrain.sh entrypoint in a local `docker run` (the
+# validated gfx1250 docker + TransformerEngine recipe from
+# ../../mi450/Primus/run_dsv3_proxy_4L.sh), selects the Pro model via
+# PRIMUS_MODEL, and scales it down to a MINIMUM single-GPU proxy:
+#
+#   - model           deepseek_v4_pro      hidden 7168 / 128 heads / head_dim
+#                                          512 / o_groups 16 (full Pro widths
+#                                          from deepseek_v4_pro.yaml)
+#   - parallel        TP=1 PP=1 EP=1       (single GPU; no SLURM, no DeepEP)
+#   - num_layers      4                    (vs production 61; MINIMUM slice
+#                                          that still exercises every V4
+#                                          attention layer kind)
+#   - compress_ratios [128,128,4,0]        Pro pattern (first two HCA, then
+#                                          CSA, last dense+SWA) -> covers
+#                                          HCA cr=128 / CSA cr=4 / dense cr=0
+#   - num_experts     48  topk 1           (production 384/topk6 div by 8 =
+#                                          the per-rank shape of the EP=8
+#                                          production run; topk ceil(6/8)=1.
+#                                          ~48 experts x 66M x 4L + Muon fp32
+#                                          states + seq-4096 activations -> 273
+#                                          GB peak (measured), fits the 432 GiB
+#                                          card; full 384 will NOT fit.)
+#   - moe_ffn_hidden  3072                 (full Pro MoE width; from yaml)
+#   - seq_length      4096                 (raised from 512: at GBS=8 the fixed
+#                                          Muon Newton-Schulz cost otherwise
+#                                          dominates GPU time; 4096
+#                                          tokens/microbatch amortizes it to a
+#                                          representative profile. CSA gather
+#                                          scales w/ seq. Set 512 for a fast smoke.)
+#   - index_topk      64                   (CSA top-K; <= cr=4 pool seq/4, i.e.
+#                                          4096/4=1024 at the default seq)
+#   - precision       FP8 e4m3 + tensorwise (paper fp8 LAYOUT, per-tensor scale;
+#                                          mxfp8/ue8m0 is GUARDED off — diverges
+#                                          on this build. CK-free TE path.
+#                                          PRECISION_TYPE=BF16 for a BF16 A/B.)
+#
+# Optimizer: Muon (paper §4.2.2), same recipe as the upstream pro_muon runner:
+#   momentum 0.95, update-RMS scale 0.18, AdamW eps 1e-20, LR 2.0e-4->2.0e-5,
+#   balance-loss 1e-4. Muon hard-requires use_distributed_optimizer=False +
+#   use_precision_aware_optimizer=False (so optimizer states stay fp32 — this
+#   is the binding memory cost, NOT activations). Set OPTIMIZER=adam to A/B.
+#   NOTE: at the tiny single-GPU GBS the Newton-Schulz cost looks huge as a %
+#   of GEMM (starved-batch artifact, not a Muon bug); raise GBS to amortize.
+#
+# Correctness-first defaults (this is a "does V4-Pro train at all on 1 gfx1250
+# GPU" bring-up, not a perf push). Eager attention; Turbo/DeepEP/tilelang/
+# compiled-Sinkhorn/plan-6 Triton fusions all OFF; stock hipBLASLt; profiler
+# OFF. Every knob is ${VAR:-DEFAULT}-guarded for command-line A/B.
+#
+# REQUIRED gfx1250 fix kept ON (single-GPU-safe): RCCL all_reduce(op=AVG) hangs
+# even at world_size=1 on this build and Megatron's MoE aux-loss reduce uses
+# AVG; sitecustomize on PYTHONPATH rewrites AVG -> SUM/world_size. See
+# rccl_avg_workaround/sitecustomize.py.
+#
+# Usage:
+#   ./run_deepseek_v4_pro_muon_1gpu.sh                                  # 10-iter smoke
+#   OPTIMIZER=adam ./run_deepseek_v4_pro_muon_1gpu.sh                   # A/B vs AdamW
+#   PRIMUS_TOTAL_LAYERS=6 ./run_deepseek_v4_pro_muon_1gpu.sh            # deeper slice (watch HBM)
+#   PRIMUS_NUM_EXPERTS=8 PRIMUS_MOE_TOPK=2 ./run_deepseek_v4_pro_muon_1gpu.sh  # tiny MoE
+#   HIP_VISIBLE_DEVICES=3 ./run_deepseek_v4_pro_muon_1gpu.sh            # pin a card
+###############################################################################
+set -eo pipefail
+
+export DOCKER_IMAGE=${DOCKER_IMAGE:-registry-sc-harbor.amd.com/framework/therock-npi@sha256:feba897e2a32a2465b8b296ed2662b2ad6136b5f1cf6f6c2716a3674aafc30f3}
+# Repo root: this script lives under examples/deepseek-v4/, so resolve two levels
+# up. All paths below (TE_DIR, rccl_avg_workaround, PRIMUS_PATH) are repo-root-relative.
+SCRIPT_DIR=$(realpath -m "$(dirname "$0")/../..")
+export TE_DIR=${TE_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/TransformerEngine")}
+export TE_WHEEL_DIR=${TE_WHEEL_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/dist/feba897")}
+
+# ---------- Attention backend env (TE side) --------------------------------
+export NVTE_FUSED_ATTN=1
+export NVTE_FUSED_ATTN_CK=0
+export NVTE_FUSED_ATTN_AOTRITON=1
+export NVTE_USE_CK_GEMM=0
+export NVTE_FLASH_ATTN=0
+
+# ---------- hipBLASLt: STOCK by default; opt-in TUNED (PRIMUS_TUNED_HIPBLASLT=1) -
+# Stock is the safe default: the older feba897 tuned bundle DEADLOCKED on a
+# backward-FP8 GSU split-K kernel on this host (GPU wedge -> node reboot,
+# 2026-06-10). PRIMUS_TUNED_HIPBLASLT=1 opts into a freshly built GridBased
+# gfx1250 tuned library (qwen3/dsv3 tuned; swept clean on dsv4 fwd shapes) via
+# LD_PRELOAD + HIPBLASLT_TENSILE_LIBPATH. This is a GUARDED EXPERIMENT: run with
+# a watchdog and expect a possible node reboot if the backward path still wedges.
+# NOTE: never export an EMPTY HIPBLASLT_TENSILE_LIBPATH into the container — a
+# missing path breaks even stock hipBLASLt ("Cannot read TensileLibrary...").
+export PRIMUS_TUNED_HIPBLASLT=${PRIMUS_TUNED_HIPBLASLT:-0}
+export HBL_TUNED_RELEASE=${HBL_TUNED_RELEASE:-/home/yanyuqin/hipblaslt/rocm-libraries/projects/hipblaslt/build/release}
+if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then
+    if [ ! -f "$HBL_TUNED_RELEASE/library/libhipblaslt.so.1" ]; then
+        echo "[hipblaslt] ERROR: tuned lib not found at $HBL_TUNED_RELEASE/library/libhipblaslt.so.1" >&2
+        exit 1
+    fi
+    # LD_PRELOAD / LD_LIBRARY_PATH are injected INSIDE the container (below) so the
+    # image's own rocm/torch lib paths are preserved (prepend, not override).
+    export HIPBLASLT_TENSILE_LIBPATH="$HBL_TUNED_RELEASE/Tensile/library/gfx1250"
+    echo "[hipblaslt] TUNED (opt-in): libpath=$HIPBLASLT_TENSILE_LIBPATH, LD_PRELOAD=libhipblaslt.so.1 (GUARDED: watch for wedge)"
+else
+    unset HIPBLASLT_DIR HIPBLASLT_LD_PRELOAD HIPBLASLT_TENSILE_LIBPATH
+    echo "[hipblaslt] STOCK (container built-in gfx1250 catalog)"
+fi
+
+# ---------- REQUIRED gfx1250 RCCL AVG->SUM workaround -----------------------
+export PYTHONPATH="$SCRIPT_DIR/examples/deepseek-v4/rccl_avg_workaround:${PYTHONPATH:-}"
+# Real primus_turbo imports flydsl at import time; put FLYDSL_PKG_DIR on PYTHONPATH.
+export FLYDSL_PKG_DIR=${FLYDSL_PKG_DIR:-}
+if [ -n "$FLYDSL_PKG_DIR" ] && [ -d "$FLYDSL_PKG_DIR/flydsl" ]; then
+    export PYTHONPATH="$FLYDSL_PKG_DIR:$PYTHONPATH"
+fi
+
+# SDMA OFF on this host (run 3 debugging, 2026-06-10): an SDMA H2D copy
+# intermittently never signals completion — py-spy --native showed the trainer
+# pinned in rocr BusyWaitSignal under a trivial `torch.tensor(n, device=dev)`
+# in Megatron get_batch, GPU 0%, dmesg clean. Killing the stuck proc then
+# leaves MES unrecoverable (recovery disabled) -> node reboot. Blit-kernel
+# copies are slower but don't use the flaky SDMA queues. This may ALSO be the
+# true cause of the run-1 "permute autotune wedge" (same stuck-queue
+# signature; permute fusion possibly innocent).
+export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-0}
+
+# ---------- Distributed / NCCL: single GPU, loopback only -------------------
+export HSA_NO_SCRATCH_RECLAIM=1
+export NCCL_IB_DISABLE=1
+export NCCL_P2P_DISABLE=1
+export NCCL_IB_HCA=
+export NCCL_SOCKET_IFNAME=lo
+export GLOO_SOCKET_IFNAME=lo
+export RCCL_DISABLE_AMDSMI=1
+export NCCL_AMDSMI_DISABLE=1
+export USING_AINIC=0
+
+export GPUS_PER_NODE=1
+export NNODES=1
+export PYTHONUNBUFFERED=1
+
+# REQUIRED gfx1250 workaround for V4-Pro (default ON). The Pro model build
+# (hidden_size 7168, NOT a multiple of 4096) leaves a memory layout that wedges
+# the process's first high-priority MES queue creation at iter-1 get_batch
+# -> deadlock -> node reboot (debugged 2026-06-11; root cause = MES queue
+# creation vs non-4096-aligned allocation layout). AMD_SERIALIZE_COPY=3 alone
+# prevents it (kernels stay async; small iter-time cost on the eager
+# proxy). Bisected: KERNEL serialize + LAUNCH_BLOCKING are NOT needed, so they
+# default off. Set AMD_SERIALIZE_COPY=0 only to re-demonstrate the hang.
+export AMD_SERIALIZE_COPY=${AMD_SERIALIZE_COPY:-3}
+export AMD_SERIALIZE_KERNEL=${AMD_SERIALIZE_KERNEL:-0}
+export HIP_LAUNCH_BLOCKING=${HIP_LAUNCH_BLOCKING:-0}
+
+# ---------- Model: DeepSeek-V4 Pro (selected via the EXP yaml model: line) --
+export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro}
+
+# ---------- Pro MINIMUM single-GPU proxy shape ------------------------------
+export PRIMUS_TP=${PRIMUS_TP:-1}
+export PRIMUS_PP=${PRIMUS_PP:-1}
+export PRIMUS_EP=${PRIMUS_EP:-1}
+# 3 layers (down from 4): the 4-layer model sits near the 432GB gfx1250's
+# capacity and OOMs at iter 2 (Muon keeps fp32 optimizer states). Dropping
+# the last layer frees the headroom for a warm step.
+export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-3}
+export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-48}
+export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-1}
+export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072}
+export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-64}
+export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096}
+export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}}
+export MBS=${MBS:-1}
+export GBS=${GBS:-8}
+export TRAIN_ITERS=${TRAIN_ITERS:-10}
+
+# Per-layer compression schedule, length == PRIMUS_TOTAL_LAYERS, mirroring the
+# Pro pattern (first two HCA(128), then CSA(4)/HCA(128) interleaved, last
+# dense+SWA(0)). Pure-bash generator (no host python needed).
+gen_pro_compress_ratios() {
+    local n=$1 i r=()
+    for ((i = 0; i < n; i++)); do
+        if   (( i < 2 ));     then r+=(128)
+        elif (( i == n-1 ));  then r+=(0)
+        elif (( i % 2 == 0)); then r+=(4)
+        else                       r+=(128)
+        fi
+    done
+    local IFS=,
+    echo "[${r[*]}]"
+}
+export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-$(gen_pro_compress_ratios "$PRIMUS_TOTAL_LAYERS")}
+
+# ---------- Optimizer: Muon (paper §4.2.2) ---------------------------------
+export OPTIMIZER=${OPTIMIZER:-muon}
+export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95}
+export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18}
+export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False}
+export USE_PRECISION_AWARE_OPTIMIZER=${USE_PRECISION_AWARE_OPTIMIZER:-False}
+export LR=${LR:-2.0e-4}
+export MIN_LR=${MIN_LR:-2.0e-5}
+export ADAM_EPS=${ADAM_EPS:-1.0e-20}   # needs a decimal point — Primus parses "1e-20" as a string
+export MOE_AUX_LOSS_COEFF=${MOE_AUX_LOSS_COEFF:-0.0001}
+export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-0}   # V4 MTP layer not yet supported in-tree
+# Pro uses sqrtsoftplus; Megatron only supports aux-loss-free expert bias with
+# sigmoid, so disable expert bias (balancing falls back to seq_aux_loss).
+export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False}
+
+# Muon needs fp32 optimizer states (precision-aware off forces this anyway).
+OPT_DTYPE_ARGS="--main_grads_dtype fp32 --exp_avg_dtype fp32 --exp_avg_sq_dtype fp32"
+
+# ---------- Perf knobs: V4 attention backends ON; turbo paths OFF ----------
+# V4 attention backend (replaces the unfused/eager path). Covers the dense +
+# HCA layers (compress_ratio in {0, 128}) via USE_V4_ATTENTION_BACKEND and the
+# CSA layers (compress_ratio == 4) via USE_V4_CSA_ATTENTION_BACKEND.
+# Validated on gfx1250 after the WMMA tile-floor fix (06ae5214).
+export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2}
+export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2}
+export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False}
+export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-False}
+export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-False}
+# Projections: the FP8 yaml sets use_turbo_gemm=true (PrimusTurboLinear);
+# turbo-free here -> TELinear. Override off so no turbo GEMM is invoked.
+export USE_TURBO_PARALLEL_LINEAR=${USE_TURBO_PARALLEL_LINEAR:-False}
+export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False}
+export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-0}
+# RoPE Triton: default ON. Trace (2026-06-25, L3) attributed 960 kernels / 513 tiny
+# (<5us) / 38.6 ms to the eager rotary-embedding path — a launch-bound fusion target.
+export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1}
+# Sinkhorn Triton fused FWD/BWD: default ON. The eager Sinkhorn-Knopp loop
+# (n_iters=20) launches ~18,600 tiny sum/add/div kernels per step (5,616 on the
+# fwd side alone); the Triton path emits exactly 1 fwd + 1 bwd kernel per call.
+# Measured 2026-06-25 (0612, L3, FP8): total GPU events 80,962 -> 62,340, sinkhorn
+# GPU kernels 5,616 -> 48, warm step ~2,890 -> ~2,797 ms (+3.2%), 0 NaN / loss
+# bit-identical. Falls back to eager when the shape/device is unsupported. Set =0 to A/B.
+export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1}
+# HyperConnection mHC Triton: default ON. The mHC HyperMixer glue (pre/post/comb
+# projections + scales), separate from the already-fused HC-expand and sinkhorn.
+# Trace (2026-06-25, L3): 1,320 kernels / 872 tiny (<5us) / 52 ms — top remaining
+# launch-bound target after sinkhorn.
+export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1}
+# CSA indexer Triton: kept OFF — inert at L3 (compress_ratios [128,128,0] has NO CSA
+# layer, so the indexer never runs). Enable only with a CSA layer (>=4 layers).
+export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-0}
+export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0}
+# V4 MoE router Triton: default ON. Trace (2026-06-25, L3): 432 kernels / 208 tiny /
+# 5.3 ms — marginal, but launch-bound and correctness-neutral.
+export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1}
+
+export ENABLE_PRIMUS_TURBO=False
+if [ "$USE_TURBO_ATTENTION" = "True" ] || [ "$USE_TURBO_DEEPEP" = "True" ] || [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then
+    ENABLE_PRIMUS_TURBO=True
+fi
+
+# MoE permute fusion OFF for Pro on gfx1250: the Triton permute_with_mask_map
+# BACKWARD autotune wedges the GPU stream at Pro shapes (48 experts / hidden
+# 7168) — cuda.synchronize inside triton do_bench never returns (debugged
+# 2026-06-10 via py-spy; flash shapes 32 experts / hidden 4096 autotune fine).
+# Eager permute is the safe path; flip =True to retry after a triton fix.
+export MOE_PERMUTE_FUSION=${MOE_PERMUTE_FUSION:-False}
+
+export PROFILE=${PROFILE:-False}
+# PyTorch profiler writes the chrome trace via tensorboard_trace_handler(
+# args.tensorboard_dir), so the tensorboard dir MUST be enabled to get a trace.
+# Default tensorboard off, but auto-enable it whenever PROFILE=True so a
+# profiled run actually produces a trace. profile window = steps [START,END);
+# need TRAIN_ITERS > PROFILE_STEP_END.
+export DISABLE_TENSORBOARD=${DISABLE_TENSORBOARD:-True}
+if [ "$PROFILE" = "True" ]; then export DISABLE_TENSORBOARD=False; fi
+export PROFILE_STEP_START=${PROFILE_STEP_START:-6}
+export PROFILE_STEP_END=${PROFILE_STEP_END:-7}
+export PRIMUS_TEAM=${PRIMUS_TEAM:-amd}
+export PRIMUS_USER=${PRIMUS_USER:-gfx1250-1gpu}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_pro_muon_1gpu_L${PRIMUS_TOTAL_LAYERS}_E${PRIMUS_NUM_EXPERTS}_seq${PRIMUS_SEQ_LENGTH}}
+
+PRIMUS_PATH="$SCRIPT_DIR"
+DATA_PATH="${PRIMUS_PATH}/data"
+mkdir -p "$DATA_PATH"
+
+EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml}
+LOG=${LOG:-deepseek-v4-pro-muon-1gpu.log}
+
+# ---------- FP8 training (matches upstream run_deepseek_v4_pro_fp8_paper.sh) --
+# PRECISION_TYPE=FP8 (default) -> FP8=e4m3, FP8_RECIPE=tensorwise. This is the
+# paper's fp8 LAYOUT (all weight GEMMs in fp8: MoE expert GEMMs via TEGroupedMLP,
+# attention QKV/O + dense proj via TELinear; attention core QK^T/softmax*V stays
+# BF16; mHC/Sinkhorn fp32; embedding/head/RMSNorm/router BF16; optimizer fp32) —
+# but with a per-TENSOR scale instead of the paper's ue8m0 microscale.
+#
+# CK-free by construction: this 1gpu launcher already has TURBO_USE_GROUPED_MLP=
+# False + USE_TURBO_DEEPEP=False, so experts route to TEGroupedMLP (hipBLASLt),
+# not PrimusTurboGroupedMLP (ck_grouped_gemm). NVTE_ROCM_ENABLE_MXFP8=1 is set
+# by examples/run_pretrain.sh.
+#
+# WHY NOT mxfp8 (paper ue8m0): two blockers on this build, both upstream-
+# root-caused. (1) TE-ROCm MXFP8 asserts GEMM K % 128 == 0 (rocm_gemm.hip:1529)
+# and V4 has non-128 K dims (e.g. K=224, K=32) -> errors out. (2) Even if it ran,
+# mxfp8's e8m0 per-block quant noise AMPLIFIES MULTIPLICATIVELY through backward
+# depth -> divergence. tensorwise's smooth per-tensor fp32 scale is stable
+# (upstream: loss 12 -> 0.82 at full depth). So mxfp8 is GUARDED below.
+#
+# The earlier FP8 no-op (decoder skipped the fp8 context) is fixed upstream
+# (commit b662c40b) and lives in the mounted repo, so FP8 now actually engages.
+# A/B back to BF16 with PRECISION_TYPE=BF16 (or FP8=null).
+export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1}
+# TURBO-FREE FP8: primus_turbo is only an import-shim in this gfx1250 container,
+# so the turbo FP8 path (PrimusTurboQuantConfig / primus_turbo_fp8_autocast)
+# can't run. Force the TE-native fp8_autocast branch (fp8_utils.py honors this).
+export PRIMUS_FP8_DISABLE_TURBO=${PRIMUS_FP8_DISABLE_TURBO:-1}
+export PRECISION_TYPE=${PRECISION_TYPE:-FP8}
+if [ "$PRECISION_TYPE" = "FP8" ]; then
+    export FP8=${FP8:-e4m3}
+    export FP8_RECIPE=${FP8_RECIPE:-tensorwise}
+    # GUARD: mxfp8 (paper ue8m0) diverges for V4 on this build. Refuse it unless
+    # explicitly forced, matching upstream run_deepseek_v4_pro_muon.sh.
+    if [ "$FP8_RECIPE" = "mxfp8" ] && [ "${MXFP8_I_KNOW_ITS_BROKEN:-0}" != "1" ]; then
+        echo "[FATAL] FP8_RECIPE=mxfp8 diverges for V4 on this build (TE K%128 assert +" >&2
+        echo "        e8m0 depth-amplified instability). Use FP8_RECIPE=tensorwise," >&2
+        echo "        or set MXFP8_I_KNOW_ITS_BROKEN=1 to force it anyway." >&2
+        exit 1
+    fi
+else
+    export FP8=${FP8:-null}
+    export FP8_RECIPE=${FP8_RECIPE:-null}
+fi
+
+if [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then
+  export PRIMUS_BIAS_SWIGLU_FUSION=True
+fi
+
+if [ ! -d "$PRIMUS_PATH/third_party/Megatron-LM" ] || \
+   [ -z "$(ls -A "$PRIMUS_PATH/third_party/Megatron-LM" 2>/dev/null)" ]; then
+    echo "[ERROR] third_party/Megatron-LM missing/empty -> run: git submodule update --init --recursive" >&2
+    exit 1
+fi
+
+echo "[pro] model=$PRIMUS_MODEL layers=$PRIMUS_TOTAL_LAYERS experts=$PRIMUS_NUM_EXPERTS seq=$PRIMUS_SEQ_LENGTH optimizer=$OPTIMIZER compress_ratios=$PRIMUS_COMPRESS_RATIOS"
+
+# V4-Pro single-GPU overrides (trailing args -> run_pretrain.sh -> primus cli
+# train pretrain --config $EXP ...). Mirrors run_deepseek_v4_pro_muon.sh's CLI
+# set, minus the DeepEP wiring, scaled to one GPU / minimum layers.
+# overlap_grad_reduce/param_gather stay OFF: upstream enabled them for multi-
+# node DP scaling (needs the distributed optimizer + the indexer-param freeze),
+# but at single-GPU DP=1 they are no-ops, and Muon requires them off anyway.
+PROXY_OVERRIDES="\
+    --backend_path $PRIMUS_PATH/third_party/Megatron-LM \
+    --train_iters $TRAIN_ITERS \
+    --lr_warmup_iters 0 \
+    --lr_decay_iters $TRAIN_ITERS \
+    --num_layers $PRIMUS_TOTAL_LAYERS \
+    --compress_ratios $PRIMUS_COMPRESS_RATIOS \
+    --micro_batch_size $MBS \
+    --global_batch_size $GBS \
+    --lr $LR \
+    --min_lr $MIN_LR \
+    --adam_eps $ADAM_EPS \
+    --moe_aux_loss_coeff $MOE_AUX_LOSS_COEFF \
+    --seq_length $PRIMUS_SEQ_LENGTH \
+    --max_position_embeddings $PRIMUS_MAX_POSITION_EMBEDDINGS \
+    --rope_type rope \
+    --tensor_model_parallel_size $PRIMUS_TP \
+    --pipeline_model_parallel_size $PRIMUS_PP \
+    --expert_model_parallel_size $PRIMUS_EP \
+    --num_experts $PRIMUS_NUM_EXPERTS \
+    --moe_router_topk $PRIMUS_MOE_TOPK \
+    --moe_router_enable_expert_bias $PRIMUS_MOE_ENABLE_EXPERT_BIAS \
+    --moe_ffn_hidden_size $PRIMUS_MOE_FFN_HIDDEN_SIZE \
+    --index_topk $PRIMUS_INDEX_TOPK \
+    --v4_grouped_experts_support_clamped_swiglu True \
+    --mtp_num_layers $MTP_NUM_LAYERS \
+    --mock_data True \
+    --moe_router_force_load_balancing True \
+    --log_avg_skip_iterations 3 \
+    --optimizer $OPTIMIZER \
+    --muon_momentum $MUON_MOMENTUM \
+    --muon_extra_scale_factor $MUON_EXTRA_SCALE_FACTOR \
+    --use_distributed_optimizer $USE_DISTRIBUTED_OPTIMIZER \
+    --use_precision_aware_optimizer $USE_PRECISION_AWARE_OPTIMIZER \
+    $OPT_DTYPE_ARGS \
+    --enable_primus_turbo $ENABLE_PRIMUS_TURBO \
+    --use_turbo_attention $USE_TURBO_ATTENTION \
+    --use_turbo_deepep $USE_TURBO_DEEPEP \
+    --use_turbo_grouped_gemm $TURBO_USE_GROUPED_MLP \
+    --use_turbo_gemm $USE_TURBO_PARALLEL_LINEAR \
+    --use_v4_attention_backend $USE_V4_ATTENTION_BACKEND \
+    --use_v4_csa_attention_backend $USE_V4_CSA_ATTENTION_BACKEND \
+    --use_v4_compiled_sinkhorn $USE_V4_COMPILED_SINKHORN \
+    --moe_use_legacy_grouped_gemm False \
+    --moe_permute_fusion $MOE_PERMUTE_FUSION \
+    --fp8 $FP8 \
+    --fp8_recipe $FP8_RECIPE \
+    --recompute_num_layers 0 \
+    --recompute_granularity full \
+    --recompute_method block \
+    --gradient_accumulation_fusion False \
+    --overlap_grad_reduce False \
+    --overlap_param_gather False \
+    --disable_last_saving True \
+    --disable_wandb True \
+    --disable_tensorboard $DISABLE_TENSORBOARD \
+    --profile $PROFILE \
+    --use_pytorch_profiler $PROFILE \
+    --profile_step_start $PROFILE_STEP_START \
+    --profile_step_end $PROFILE_STEP_END \
+    --bias_swiglu_fusion $PRIMUS_BIAS_SWIGLU_FUSION \
+    --torch_profiler_use_gzip True"
+
+ENV_ARGS=()
+for v in DOCKER_IMAGE NVTE_FUSED_ATTN NVTE_FUSED_ATTN_CK NVTE_FUSED_ATTN_AOTRITON \
+         PRIMUS_TURBO_GEMM_BACKEND PRIMUS_TURBO_GROUPED_GEMM_BACKEND TURBO_WHEEL_DIR FLYDSL_PKG_DIR \
+         NVTE_FLASH_ATTN NVTE_USE_CK_GEMM NVTE_ROCM_ENABLE_MXFP8 PRIMUS_FP8_DISABLE_TURBO PYTHONPATH HSA_ENABLE_SDMA HSA_NO_SCRATCH_RECLAIM \
+         TORCH_COMPILE_DISABLE TORCHINDUCTOR_COMPILE_THREADS TRITON_CACHE_DIR \
+         HSA_SIGNAL_ABORT_TIMEOUT HSA_ENABLE_INTERRUPT \
+         HIP_LAUNCH_BLOCKING AMD_SERIALIZE_KERNEL AMD_SERIALIZE_COPY \
+         AMD_LOG_LEVEL AMD_LOG_MASK MASTER_PORT TORCH_NCCL_HIGH_PRIORITY \
+         NCCL_IB_DISABLE NCCL_P2P_DISABLE NCCL_IB_HCA NCCL_SOCKET_IFNAME \
+         GLOO_SOCKET_IFNAME RCCL_DISABLE_AMDSMI NCCL_AMDSMI_DISABLE USING_AINIC \
+         GPUS_PER_NODE NNODES PYTHONUNBUFFERED TE_DIR TE_WHEEL_DIR PRIMUS_MODEL \
+         PRIMUS_SEQ_LENGTH PRIMUS_MAX_POSITION_EMBEDDINGS \
+         PRIMUS_TEAM PRIMUS_USER PRIMUS_EXP_NAME \
+         PRIMUS_STACK_GROUPED_WEIGHT_TRITON PRIMUS_ROPE_TRITON \
+         PRIMUS_SINKHORN_TRITON PRIMUS_HC_TRITON PRIMUS_INDEXER_TRITON \
+         PRIMUS_INDEXER_TRITON_FULL PRIMUS_V4_ROUTER_TRITON \
+         PRIMUS_MUON_BATCHED_NS PRIMUS_COMPRESS_ROPE_CACHE PRIMUS_COMPRESS_POOL_TRITON; do
+    ENV_ARGS+=("--env" "$v")
+done
+[[ -n "${HIP_VISIBLE_DEVICES:-}" ]] && ENV_ARGS+=("--env" "HIP_VISIBLE_DEVICES")
+# EXTRA_CLI: extra trailing --flag value overrides appended after PROXY_OVERRIDES
+# (argparse last-wins), for dimension bisects etc.
+[[ -n "${EXTRA_CLI:-}" ]] && ENV_ARGS+=("--env" "EXTRA_CLI")
+
+# Persistent Triton compile cache. The --rm container makes TRITON_CACHE_DIR
+# ephemeral, so every run recompiles ALL kernels from scratch (the slow CPU-bound
+# LLVM step that dominates iteration 1, esp. under this node's MCE storm). Mount a
+# host dir so compiled kernels (hsaco) are reused across runs -> iter-1 of every
+# later run with the same shapes skips the cold compile. Triton keys cache entries
+# by kernel-source + arch + constexpr hash, so a wheel/arch/shape change auto-
+# invalidates (safe to keep warm). Disable with PRIMUS_TRITON_CACHE_DIR="".
+export PRIMUS_TRITON_CACHE_DIR=${PRIMUS_TRITON_CACHE_DIR:-$PRIMUS_PATH/.triton_cache_shared}
+if [ -n "$PRIMUS_TRITON_CACHE_DIR" ]; then
+    mkdir -p "$PRIMUS_TRITON_CACHE_DIR"
+    export TRITON_CACHE_DIR="$PRIMUS_TRITON_CACHE_DIR"
+    echo "[triton] persistent compile cache: $TRITON_CACHE_DIR ($(find "$TRITON_CACHE_DIR" -maxdepth 1 -type d 2>/dev/null | wc -l) entries)"
+fi
+
+VOLUME_ARGS=(-v "$PRIMUS_PATH":"$PRIMUS_PATH" -v "$DATA_PATH":"$DATA_PATH")
+[[ -d "$TE_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TE_WHEEL_DIR":"$TE_WHEEL_DIR")
+[[ -d "$TE_DIR" ]] && VOLUME_ARGS+=(-v "$TE_DIR":"$TE_DIR")
+[[ -n "${TURBO_WHEEL_DIR:-}" && -d "$TURBO_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TURBO_WHEEL_DIR":"$TURBO_WHEEL_DIR")
+[[ -n "${FLYDSL_PKG_DIR:-}" && -d "$FLYDSL_PKG_DIR/flydsl" ]] && VOLUME_ARGS+=(-v "$FLYDSL_PKG_DIR":"$FLYDSL_PKG_DIR")
+[[ -n "${TRITON_CACHE_DIR:-}" ]] && VOLUME_ARGS+=(-v "$TRITON_CACHE_DIR":"$TRITON_CACHE_DIR")
+# Opt-in tuned hipBLASLt: mount the built library at the same path and pass the
+# loader env into the container (only when enabled, to keep stock runs untouched).
+if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then
+    VOLUME_ARGS+=(-v "$HBL_TUNED_RELEASE":"$HBL_TUNED_RELEASE")
+    ENV_ARGS+=("--env" "PRIMUS_TUNED_HIPBLASLT" "--env" "HIPBLASLT_TENSILE_LIBPATH" \
+               "--env" "HBL_TUNED_RELEASE")
+fi
+# Container-side loader injection for the tuned lib (prepends to the image's paths).
+HBL_PRELOAD_PREFIX=""
+if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then
+    HBL_PRELOAD_PREFIX="export LD_LIBRARY_PATH=\"\$HBL_TUNED_RELEASE/library:\${LD_LIBRARY_PATH:-}\" && export LD_PRELOAD=\"\$HBL_TUNED_RELEASE/library/libhipblaslt.so.1\${LD_PRELOAD:+:\$LD_PRELOAD}\" && echo \"[hipblaslt] container LD_PRELOAD=\$LD_PRELOAD\" && "
+fi
+
+TE_INSTALL_PREFIX="\
+    if ls ${TE_WHEEL_DIR}/transformer_engine-*.whl >/dev/null 2>&1; then \
+        echo '[TE] installing prebuilt wheel from ${TE_WHEEL_DIR}' && \
+        pip install --quiet --force-reinstall --no-deps ${TE_WHEEL_DIR}/transformer_engine-*.whl && \
+        pip install --quiet einops nvdlfw-inspect onnxscript onnx pydantic importlib-metadata packaging transformers pybind11; \
+    else \
+        echo '[TE] WARNING: no TE wheel found at ${TE_WHEEL_DIR}; run will likely fail'; \
+    fi && \
+    echo '[deps] installing Primus requirements' && \
+    pip install --quiet -r requirements.txt && \
+    if [ -n \"${TURBO_WHEEL_DIR:-}\" ] && ls ${TURBO_WHEEL_DIR}/primus_turbo-*.whl >/dev/null 2>&1; then \
+        echo '[turbo] installing real primus_turbo wheel from ${TURBO_WHEEL_DIR}' && \
+        pip install --quiet --force-reinstall --no-deps ${TURBO_WHEEL_DIR}/primus_turbo-*.whl && \
+        python -c 'import primus_turbo, primus_turbo.pytorch as _; print(\"[turbo] primus_turbo\", primus_turbo.__version__, \"imported OK\")'; \
+    fi && "
+
+docker run --rm \
+    "${ENV_ARGS[@]}" \
+    --ipc=host --network=host \
+    --device=/dev/kfd --device=/dev/dri \
+    --cap-add=SYS_PTRACE --cap-add=CAP_SYS_ADMIN \
+    --security-opt seccomp=unconfined --group-add video \
+    --privileged \
+    --name primus-v4-pro-muon-1gpu \
+    "${VOLUME_ARGS[@]}" \
+    "$DOCKER_IMAGE" /bin/bash -c "\
+        set -e && cd $PRIMUS_PATH && \
+        ${HBL_PRELOAD_PREFIX}\
+        ${TE_INSTALL_PREFIX}\
+        echo '==================== V4-PRO + MUON 1-GPU PROXY (gfx1250, BF16, eager, no profiler) ====================' && \
+        EXP=$EXP PRIMUS_MODEL=$PRIMUS_MODEL GPUS_PER_NODE=1 NNODES=1 bash examples/run_pretrain.sh \
+            ${PROXY_OVERRIDES} ${EXTRA_CLI:-}" \
+    2>&1 | tee "$LOG"
diff --git a/examples/deepseek-v4/run_dsv4_projection_1gpu.sh b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh
new file mode 100755
index 000000000..fe0ec6886
--- /dev/null
+++ b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh
@@ -0,0 +1,158 @@
+#!/bin/bash
+###############################################################################
+# Primus PROJECTION (memory / performance) for DeepSeek-V4 on one gfx1250 GPU.
+#
+# Sibling of run_deepseek_v4_pro_muon_1gpu.sh. Same validated gfx1250 docker
+# recipe and the same required workarounds, but instead of pretraining it runs
+# the Primus projection tool (docs/projection.md): benchmark a couple of layers
+# on this single GPU and analytically project memory + training performance to
+# a multi-node target cluster.
+#
+# Usage:
+#   ./run_dsv4_projection_1gpu.sh                       # performance, pro, ->8 nodes
+#   MODE=memory ./run_dsv4_projection_1gpu.sh           # memory projection only
+#   PRIMUS_MODEL=deepseek_v4_flash ./run_dsv4_projection_1gpu.sh   # flash model
+#   TARGET_NODES=16 ./run_dsv4_projection_1gpu.sh       # project to 16 nodes
+#   PROFILING_MODE=simulate GPU_ARCH=mi355x MODE=performance ./run_dsv4_projection_1gpu.sh  # CPU-only
+###############################################################################
+set -eo pipefail
+
+# Repo root: this script lives under examples/deepseek-v4/, so resolve two levels
+# up. All paths below (TE_WHEEL_DIR, third_party, VOLUME_ARGS, cd) are repo-root-relative.
+SCRIPT_DIR=$(realpath -m "$(dirname "$0")/../..")
+export DOCKER_IMAGE=${DOCKER_IMAGE:-registry-sc-harbor.amd.com/framework/therock-npi@sha256:feba897e2a32a2465b8b296ed2662b2ad6136b5f1cf6f6c2716a3674aafc30f3}
+export TE_WHEEL_DIR=${TE_WHEEL_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/dist/feba897")}
+
+# ---------- What to project -------------------------------------------------
+export MODE=${MODE:-performance}                       # memory | performance
+export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro}   # deepseek_v4_pro | deepseek_v4_flash
+export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml}
+export BENCHMARK_GPUS=${BENCHMARK_GPUS:-1}             # benchmark on this many GPUs (1 here)
+# This box has ONE physical GPU. We invoke `primus-cli direct --single`
+# (ONE python3 process, NOT torchrun) —
+# the same trick the validated run_dsv3_projection script uses. In --single mode
+# torchrun is not used, so GPUS_PER_NODE no longer drives --nproc_per_node; it
+# serves ONLY as the TARGET node size (8 GPUs/node), while the projection spawns
+# its own nproc=1 benchmark subprocess pinned to the single physical GPU. This
+# gives correct intra-/inter-node comm modeling for the real 8-GPU-node cluster.
+export GPUS_PER_NODE=8                                  # TARGET node size (8 nodes x 8 = 64 GPUs)
+export TARGET_NODES=${TARGET_NODES:-8}                  # production TP1*PP8*EP8 = 64 GPUs = 8 nodes
+export PROFILING_MODE=${PROFILING_MODE:-benchmark}     # benchmark | simulate | both
+export GPU_ARCH=${GPU_ARCH:-}                          # e.g. mi355x for --profiling-mode simulate
+
+# ---------- Required gfx1250 workarounds (see the pretrain launcher) ---------
+export HSA_NO_SCRATCH_RECLAIM=1
+# gfx1250 MES async-queue hang workaround — matched EXACTLY to the training
+# launcher (run_deepseek_v4_pro_muon_1gpu.sh): AMD_SERIALIZE_COPY=3 alone, which
+# was bisected sufficient for the V4-Pro proxy (KERNEL serialize + LAUNCH_BLOCKING
+# found unnecessary, default off). If the projection still hangs with this, the
+# cause is the full-model build (all 61 layers on one rank), not these knobs.
+export AMD_SERIALIZE_COPY=${AMD_SERIALIZE_COPY:-3}
+export AMD_SERIALIZE_KERNEL=${AMD_SERIALIZE_KERNEL:-0}
+export HIP_LAUNCH_BLOCKING=${HIP_LAUNCH_BLOCKING:-0}
+export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-0}           # flaky SDMA completion-signal workaround
+# Turbo-free TE-native FP8 (tensorwise / Float8CurrentScaling), matching the
+# training launcher. Without this the model uses TE DelayedScaling, which asserts
+# against the V4 attention's save_original_input. fp8_utils.py reads this env.
+export PRIMUS_FP8_DISABLE_TURBO=${PRIMUS_FP8_DISABLE_TURBO:-1}
+export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1}
+# RCCL all_reduce(AVG) hangs even at world_size=1 -> sitecustomize rewrites AVG->SUM/ws.
+# Also put the vendored Emerging-Optimizers on PYTHONPATH so the muon optimizer
+# (emerging_optimizers.*) imports — required for OPTIMIZER=muon.
+export PYTHONPATH_IN="$SCRIPT_DIR/examples/deepseek-v4/rccl_avg_workaround:$SCRIPT_DIR/third_party/Emerging-Optimizers"
+
+LOG=${LOG:-dsv4-projection-${MODE}.log}
+
+# Config overrides appended as trailing CLI key/value pairs. V4 uses its OWN
+# attention (multi_latent_attention=false + yarn via dual_rope), but Megatron's
+# stock validate_args rejects rope_type=yarn unless MLA is on. The projection
+# benchmarks a stock-Megatron layer (not the V4 custom attention), so force
+# rope_type=rope at the Megatron-arg level — V4 applies yarn internally and
+# rope-vs-yarn is a cheap elementwise op (negligible for timing).
+# (moe_router_score_function: V4 uses sqrtsoftplus, but Megatron's stock
+# validate requires sigmoid for expert-bias aux-loss-free routing; the score
+# function is a pointwise on router logits and does not change GEMM timing.)
+# (moe_token_dispatcher_type: V4 uses the turbo "flex" dispatcher, which asserts
+# TPxEP>1; the single-GPU benchmark runs at EP=1. Force "alltoall" — dispatcher
+# type only affects MoE *communication* (modeled analytically), not expert-GEMM
+# compute, so benchmarked layer time is unchanged.)
+# (enable_primus_turbo / use_turbo_deepep: a Primus patch [moe_dispatcher_patches]
+# force-replaces the dispatcher with the turbo DeepEP "flex" one when BOTH are
+# true, and flex asserts TPxEP>1 (can't benchmark on 1 GPU). gfx1250 runs
+# turbo-free anyway (training disables it), so force them off -> standard
+# alltoall dispatcher, single-GPU-benchmarkable.)
+# (gradient_accumulation_fusion: needs APEX fused_weight_gradient_mlp_cuda, not
+# in this container; off here exactly as in the training launcher.)
+# (optimizer: use the yaml default adam — the projection only benchmarks layer
+# fwd/bwd, so the optimizer choice doesn't affect timing, and adam avoids muon's
+# "Emerging Optimizers" package dependency that isn't in this container. The
+# trainer.py adam/muon get_*_optimizer call sites were patched to match the
+# bundled Megatron signature [dropped the removed no_wd_decay_cond/scale_lr_cond/
+# lr_mult positionals that collided with use_gloo_process_groups].)
+# (tokenizer NullTokenizer: the benchmark builds a mock dataset; Megatron's
+# MockGPTDataset JSON-serializes the tokenizer via .unique_identifiers, which the
+# DeepSeekV4 HuggingFace tokenizer lacks -> crash. NullTokenizer has it, needs no
+# HF download, and preserves vocab_size (129280) so embedding/LM-head GEMM dims
+# are unchanged. Layer compute is tokenizer-independent.)
+# (use_v4_triton_attention/csa: enable the fused flash-style V4 attention kernels
+# instead of eager attention so the benchmarked attention time is representative
+# of the real training config — eager materializes [B,H,S,S] and hugely inflates
+# attention at seq 4096. Verified working on gfx1250.)
+# V4-specific flags mirrored from the training launcher (now that the V4 builder
+# is used, the real DeepseekV4MoE/HybridLayer is built and needs these): clamped
+# SwiGLU support on the grouped backend, turbo off, legacy/permute-fusion off.
+EXTRA_OVERRIDES=${EXTRA_OVERRIDES:---rope_type rope --moe_router_score_function sigmoid --moe_token_dispatcher_type alltoall --enable_primus_turbo false --use_turbo_deepep false --use_turbo_grouped_gemm false --use_turbo_gemm false --use_v4_compiled_sinkhorn false --moe_use_legacy_grouped_gemm false --moe_permute_fusion false --gradient_accumulation_fusion false --mtp_num_layers 0 --tokenizer_type NullTokenizer --use_v4_triton_attention true --use_v4_triton_csa_attention true}
+
+# ---------- Build the projection CLI args -----------------------------------
+PROJ_ARGS="projection $MODE --config $EXP"
+if [ "$MODE" = "performance" ]; then
+    PROJ_ARGS="$PROJ_ARGS --benchmark-gpus $BENCHMARK_GPUS --target-nodes $TARGET_NODES --profiling-mode $PROFILING_MODE"
+    [ -n "$GPU_ARCH" ] && PROJ_ARGS="$PROJ_ARGS --gpu-arch $GPU_ARCH"
+fi
+PROJ_ARGS="$PROJ_ARGS $EXTRA_OVERRIDES"
+
+# TE wheel install prefix (same as pretrain launcher)
+TE_INSTALL="true"
+if ls "${TE_WHEEL_DIR}"/transformer_engine-*.whl >/dev/null 2>&1; then
+    TE_INSTALL="pip install --quiet --force-reinstall --no-deps ${TE_WHEEL_DIR}/transformer_engine-*.whl && \
+        pip install --quiet einops nvdlfw-inspect onnxscript onnx pydantic importlib-metadata packaging transformers pybind11"
+fi
+# simulate mode (CPU-only, no model instantiation) needs the Origami GEMM model.
+ORIGAMI_INSTALL="true"
+if [ "$PROFILING_MODE" = "simulate" ] || [ "$PROFILING_MODE" = "both" ]; then
+    ORIGAMI_INSTALL="pip install --quiet 'git+https://github.com/ROCm/rocm-libraries.git#subdirectory=shared/origami/python' || echo '[warn] origami install failed'"
+fi
+
+VOLUME_ARGS=(-v "$SCRIPT_DIR":"$SCRIPT_DIR")
+[[ -d "$TE_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TE_WHEEL_DIR":"$TE_WHEEL_DIR")
+
+echo "[projection] mode=$MODE model=$PRIMUS_MODEL target_nodes=$TARGET_NODES profiling_mode=$PROFILING_MODE config=$EXP"
+
+# Same docker invocation as run_deepseek_v4_pro_muon_1gpu.sh (validated gfx1250 recipe).
+docker run --rm \
+    --ipc=host --network=host \
+    --device=/dev/kfd --device=/dev/dri \
+    --cap-add=SYS_PTRACE --cap-add=CAP_SYS_ADMIN \
+    --security-opt seccomp=unconfined --group-add video \
+    --privileged \
+    --name primus-v4-projection \
+    -e NNODES=1 -e GPUS_PER_NODE="$GPUS_PER_NODE" \
+    -e MASTER_ADDR=localhost -e MASTER_PORT=1234 \
+    -e GLOO_SOCKET_IFNAME=lo -e NCCL_SOCKET_IFNAME=lo \
+    -e NCCL_IB_DISABLE=1 -e NCCL_P2P_DISABLE=1 \
+    -e HSA_NO_SCRATCH_RECLAIM="$HSA_NO_SCRATCH_RECLAIM" \
+    -e AMD_SERIALIZE_COPY="$AMD_SERIALIZE_COPY" -e AMD_SERIALIZE_KERNEL="$AMD_SERIALIZE_KERNEL" \
+    -e HIP_LAUNCH_BLOCKING="$HIP_LAUNCH_BLOCKING" -e HSA_ENABLE_SDMA="$HSA_ENABLE_SDMA" \
+    -e PRIMUS_FP8_DISABLE_TURBO="$PRIMUS_FP8_DISABLE_TURBO" -e NVTE_ROCM_ENABLE_MXFP8="$NVTE_ROCM_ENABLE_MXFP8" \
+    -e PRIMUS_PROJ_MAX_LAYERS="${PRIMUS_PROJ_MAX_LAYERS:-1}" -e PRIMUS_PROJ_COMPRESS_RATIOS="${PRIMUS_PROJ_COMPRESS_RATIOS:-}" \
+    -e PRIMUS_MODEL="$PRIMUS_MODEL" -e PYTHONUNBUFFERED=1 \
+    "${VOLUME_ARGS[@]}" \
+    "$DOCKER_IMAGE" /bin/bash -c "\
+        set -e && cd $SCRIPT_DIR && \
+        export PYTHONPATH=$PYTHONPATH_IN:\${PYTHONPATH:-} && \
+        ${TE_INSTALL} && \
+        ${ORIGAMI_INSTALL} && \
+        pip install --quiet -r requirements.txt && \
+        echo '==================== DSV4 PROJECTION ($MODE) ====================' && \
+        bash runner/primus-cli direct --single -- $PROJ_ARGS" \
+    2>&1 | tee "$LOG"
diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md
new file mode 100644
index 000000000..42ad9a61a
--- /dev/null
+++ b/examples/diffusion/README.md
@@ -0,0 +1,131 @@
+# Diffusion Examples
+
+This directory contains launch examples for the in-tree `diffusion` backend.
+
+## Common Launch Env
+
+```bash
+export NNODES=${NNODES:-1}
+export NODE_RANK=${NODE_RANK:-0}
+export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
+export MASTER_PORT=${MASTER_PORT:-29500}
+export GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+```
+
+## FLUX.1-schnell Raw Image-Text
+
+Raw mode loads image-text samples and runs frozen T5, CLIP, and FLUX AE online.
+The default `DATASET=cc12m-test` uses the Hugging Face dataset
+`zirui3/cc12m-test`, so no dataset preprocessing is required for a smoke test.
+
+Download the encoders and autoencoder before launching training:
+
+```bash
+huggingface-cli download google/t5-v1_1-xxl \
+  --local-dir /models/t5-v1_1-xxl
+huggingface-cli download openai/clip-vit-large-patch14 \
+  --local-dir /models/clip-vit-large-patch14
+huggingface-cli download black-forest-labs/FLUX.1-dev ae.safetensors \
+  --local-dir /models/FLUX.1-dev
+```
+
+Launch raw training:
+
+```bash
+T5_ENCODER=/models/t5-v1_1-xxl \
+CLIP_ENCODER=/models/clip-vit-large-patch14 \
+VAE_CHECKPOINT=/models/FLUX.1-dev/ae.safetensors \
+MAX_STEPS=10 \
+torchrun \
+  --nnodes="$NNODES" --node_rank="$NODE_RANK" \
+  --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
+  --nproc_per_node="$GPUS_PER_NODE" \
+  -m primus.cli.main train pretrain \
+  --config examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml
+```
+
+To use a local WebDataset directory instead, set `DATASET_PATH=/path/to/tars`.
+To use the full Hugging Face dataset directly, add `DATASET=cc12m-wds` to the
+launch command and omit `DATASET_PATH`.
+
+To run FLUX.1-dev, use the same training example shape and set the model preset
+to `flux.1_dev_t2i.yaml`. FLUX.1-dev has a guidance embedding module;
+FLUX.1-schnell does not.
+
+## Wan Data
+
+Wan examples use a JSONL metadata file plus a media directory:
+
+```bash
+huggingface-cli download zirui3/tiny-video-samples \
+  --repo-type dataset \
+  --local-dir /data/tiny-video-samples
+```
+
+Expected layout:
+
+```text
+/data/tiny-video-samples/
+  meta.jsonl
+  data/*.mp4
+```
+
+Download Wan checkpoints separately and set the model paths used by the selected
+config. For Wan2.2 TI2V 5B, the default paths can be overridden with:
+
+```bash
+export PRETRAINED_PATH=/models/Wan2.2-TI2V-5B
+export INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B
+export TEXT_TOKENIZER=/models/Wan2.2-TI2V-5B/google/umt5-xxl
+export TEXT_ENCODER=/models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth
+export VAE_CHECKPOINT=/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+```
+
+## Wan Pretrain
+
+```bash
+DATASET_PATH=/data/tiny-video-samples/meta.jsonl \
+DATA_FOLDER=/data/tiny-video-samples/data \
+ATTENTION_BACKEND=flash_attn_aiter \
+SP_SIZE=1 \
+MAX_STEPS=10 \
+torchrun \
+  --nnodes="$NNODES" --node_rank="$NODE_RANK" \
+  --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
+  --nproc_per_node="$GPUS_PER_NODE" \
+  -m primus.cli.main train pretrain \
+  --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml
+```
+
+Use `SP_SIZE=4` or `SP_SIZE=8` when the model head count supports it.
+
+## Wan Posttrain
+
+```bash
+INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \
+DATASET_PATH=/data/tiny-video-samples/meta.jsonl \
+DATA_FOLDER=/data/tiny-video-samples/data \
+ATTENTION_BACKEND=flash_attn_aiter \
+SP_SIZE=1 \
+MAX_STEPS=10 \
+torchrun \
+  --nnodes="$NNODES" --node_rank="$NODE_RANK" \
+  --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
+  --nproc_per_node="$GPUS_PER_NODE" \
+  -m primus.cli.main train posttrain \
+  --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml
+```
+
+## Prepare Check
+
+Validate configured paths before launching:
+
+```bash
+python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \
+  --config examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml
+
+python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \
+  --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml
+```
+
+On success the hook prints `env.PREPARED=1`.
diff --git a/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml
new file mode 100644
index 000000000..9429fb1f7
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:flux.1_schnell_t2i-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    model: flux.1_schnell_t2i.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: ${LOCAL_BATCH_SIZE:1}
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/flux.1_schnell_t2i-pretrain}
+        save_steps: 0
+        save_strategy: ${SAVE_STRATEGY:dit_only}
+        run_name: flux.1_schnell_t2i-pretrain
+
+      data:
+        dataset_path: ${DATASET_PATH:}
+        empty_encodings_path: ${EMPTY_ENCODINGS_PATH:}
+        prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1}
+        img_size: ${IMG_SIZE:256}
+
+      parallelism:
+        sp_size: 1
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: ${LR:2.0e-4}
+        weight_decay: ${WEIGHT_DECAY:0.1}
+        adam_beta2: ${ADAM_BETA2:0.95}
+
+      lr_scheduler:
+        lr_scheduler_type: constant_with_warmup
+        warmup_steps: ${WARMUP_STEPS:1600}
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false}
+        compile_transformer_blocks: ${COMPILE_TRANSFORMER_BLOCKS:true}
+        fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true}
+        report_to: none
diff --git a/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml
new file mode 100644
index 000000000..3b8798bc8
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml
@@ -0,0 +1,69 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:flux.1_schnell_t2i-raw-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    model: flux.1_schnell_t2i.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      model:
+        config:
+          encoder:
+            t5_encoder: ${T5_ENCODER:google/t5-v1_1-xxl}
+            clip_encoder: ${CLIP_ENCODER:openai/clip-vit-large-patch14}
+            autoencoder: ${VAE_CHECKPOINT:black-forest-labs/FLUX.1-dev/ae.safetensors}
+            max_t5_length: ${MAX_T5_LENGTH:256}
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: ${LOCAL_BATCH_SIZE:1}
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/flux.1_schnell_t2i-raw-pretrain}
+        save_steps: 0
+        save_strategy: ${SAVE_STRATEGY:dit_only}
+        run_name: flux.1_schnell_t2i-raw-pretrain
+
+      data:
+        dataset_type: raw
+        dataset: ${DATASET:cc12m-test}
+        dataset_format: ${DATASET_FORMAT:webdataset}
+        dataset_path: ${DATASET_PATH:}
+        prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1}
+        img_size: ${IMG_SIZE:256}
+        skip_low_resolution: false
+
+      parallelism:
+        sp_size: 1
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: ${LR:2.0e-4}
+        weight_decay: ${WEIGHT_DECAY:0.1}
+        adam_beta2: ${ADAM_BETA2:0.95}
+
+      lr_scheduler:
+        lr_scheduler_type: constant_with_warmup
+        warmup_steps: ${WARMUP_STEPS:1600}
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false}
+        compile_transformer_blocks: ${COMPILE_TRANSFORMER_BLOCKS:true}
+        fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true}
+        report_to: none
diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml
new file mode 100644
index 000000000..cbd877eb8
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml
@@ -0,0 +1,53 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-posttrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  post_trainer:
+    framework: diffusion
+    config: post_trainer.yaml
+
+    # Model preset to fine-tune from INIT_CHECKPOINT.
+    model: wan2.1_t2v_1.3b_sft.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: 1
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-posttrain}
+        save_steps: 0
+        run_name: wan2.1_t2v_1.3b-posttrain
+
+      data:
+        dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl}
+        data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data}
+        frame_num: 81
+        video_backend: imageio
+        text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl}
+        height: 480
+        width: 832
+
+      parallelism:
+        sp_size: ${SP_SIZE:1}
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: 5.0e-6
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        report_to: none
diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml
new file mode 100644
index 000000000..c13a2cf0b
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml
@@ -0,0 +1,53 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    # Model preset to run.
+    model: wan2.1_t2v_1.3b.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: 1
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-pretrain}
+        save_steps: 0
+        run_name: wan2.1_t2v_1.3b-pretrain
+
+      data:
+        dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl}
+        data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data}
+        frame_num: 81
+        video_backend: imageio
+        text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl}
+        height: 480
+        width: 832
+
+      parallelism:
+        sp_size: ${SP_SIZE:1}
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: 1.0e-5
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        report_to: none
diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml
new file mode 100644
index 000000000..94ee0d09e
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml
@@ -0,0 +1,53 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-posttrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  post_trainer:
+    framework: diffusion
+    config: post_trainer.yaml
+
+    # Model preset to fine-tune from INIT_CHECKPOINT.
+    model: wan2.2_ti2v_5b_sft.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: 1
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-posttrain}
+        save_steps: 0
+        run_name: wan2.2_ti2v_5b-posttrain
+
+      data:
+        dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl}
+        data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data}
+        frame_num: 81
+        video_backend: imageio
+        text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl}
+        height: 480
+        width: 832
+
+      parallelism:
+        sp_size: ${SP_SIZE:1}
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: 5.0e-6
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        report_to: none
diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml
new file mode 100644
index 000000000..765a454b5
--- /dev/null
+++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml
@@ -0,0 +1,53 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    # Model preset to run.
+    model: wan2.2_ti2v_5b.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: INFO
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: 1
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-pretrain}
+        save_steps: 0
+        run_name: wan2.2_ti2v_5b-pretrain
+
+      data:
+        dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl}
+        data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data}
+        frame_num: 81
+        video_backend: imageio
+        text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl}
+        height: 480
+        width: 832
+
+      parallelism:
+        sp_size: ${SP_SIZE:1}
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: 1.0e-5
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        report_to: none
diff --git a/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml b/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml
new file mode 100644
index 000000000..3db838b20
--- /dev/null
+++ b/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml
@@ -0,0 +1,41 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:mixtral_8x22B-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: maxtext
+    config: pre_trainer.yaml
+
+    # model to run
+    model: mixtral_8x22B.yaml
+    overrides:
+      run_name: "mixtral_8x22b_training"
+      base_output_directory: "./output"
+      steps: 50
+      profiler: ""
+
+      # data
+      dataset_type: "synthetic"
+      hf_access_token: ${HF_TOKEN:""}
+
+      # checkpoint
+      enable_checkpointing: false
+      async_checkpointing: false
+
+      # inter-node parallelism strategy
+      dcn_data_parallelism: -1
+      dcn_fsdp_parallelism: 1
+
+      # intra-node parallelism strategy
+      ici_fsdp_parallelism: 1
+      ici_data_parallelism: 1
+      ici_expert_parallelism: -1
+
+      sparse_matmul: false
+      megablox: false
+      capacity_factor: 1
+      max_target_length: 4096
+      per_device_batch_size: 4
+      remat_policy: "save_dot_with_context_except_mlp"
diff --git a/examples/maxtext/configs/MI300X/mixtral_8x7B-pretrain.yaml b/examples/maxtext/configs/MI300X/mixtral_8x7B-pretrain.yaml
index d983c794e..0d1806150 100644
--- a/examples/maxtext/configs/MI300X/mixtral_8x7B-pretrain.yaml
+++ b/examples/maxtext/configs/MI300X/mixtral_8x7B-pretrain.yaml
@@ -39,3 +39,4 @@ modules:
       max_target_length: 4096
       per_device_batch_size: 12
       remat_policy: "save_dot_with_context_except_mlp"
+      moe_dispatch_no_expert_sharding: true
diff --git a/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml b/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml
new file mode 100644
index 000000000..d4ea7fd6b
--- /dev/null
+++ b/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml
@@ -0,0 +1,41 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:mixtral_8x22B-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: maxtext
+    config: pre_trainer.yaml
+
+    # model to run
+    model: mixtral_8x22B.yaml
+    overrides:
+      run_name: "mixtral_8x22b_training"
+      base_output_directory: "./output"
+      steps: 50
+      profiler: ""
+
+      # data
+      dataset_type: "synthetic"
+      hf_access_token: ${HF_TOKEN:""}
+
+      # checkpoint
+      enable_checkpointing: false
+      async_checkpointing: false
+
+      # inter-node parallelism strategy
+      dcn_data_parallelism: -1
+      dcn_fsdp_parallelism: 1
+
+      # intra-node parallelism strategy
+      ici_fsdp_parallelism: 1
+      ici_data_parallelism: 1
+      ici_expert_parallelism: -1
+
+      sparse_matmul: false
+      megablox: false
+      capacity_factor: 1
+      max_target_length: 4096
+      per_device_batch_size: 8
+      remat_policy: "save_dot_with_context_except_mlp"
diff --git a/examples/maxtext/configs/MI355X/mixtral_8x7B-pretrain.yaml b/examples/maxtext/configs/MI355X/mixtral_8x7B-pretrain.yaml
index 342f0f454..e0da812df 100644
--- a/examples/maxtext/configs/MI355X/mixtral_8x7B-pretrain.yaml
+++ b/examples/maxtext/configs/MI355X/mixtral_8x7B-pretrain.yaml
@@ -39,3 +39,4 @@ modules:
       max_target_length: 4096
       per_device_batch_size: 11
       remat_policy: "minimal"
+      moe_dispatch_no_expert_sharding: true
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml
new file mode 100644
index 000000000..dbd76aa26
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml
@@ -0,0 +1,209 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MI300X)
+#
+# Combines Megatron DDP + distributed optimizer with the local spec provider
+# (PrimusTurboFloat8LocalSpecProvider) for FP8 tensorwise training.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - PrimusTurboFloat8LocalSpecProvider (NO TransformerEngine dependency)
+#   - FP8 hybrid tensorwise (per-module FP8 via Primus Turbo)
+#   - Primus Turbo attention
+#   - torch.compile enabled (per_block strategy, compatible with local spec + overlap)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      # Short LR warmup ramp over the first optimizer steps for FP8 stability.
+      nemo_aligned_lr_warmup: true
+      warmup_train_steps: 2
+
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboFloat8LocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration (scaled down for MI300X 192GB; tune to your hardware)
+      micro_batch_size: 32
+      global_batch_size: 256
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: true
+      main_params_dtype: fp32
+      main_grads_dtype: bf16
+      exp_avg_dtype: bf16
+      exp_avg_sq_dtype: bf16
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # FP8 Configuration — Tensorwise + Delayed FP8 via Primus Turbo
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: "hybrid"
+      # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects
+      # tensorwise FP8 with delayed scaling. Resolution path:
+      # primus/backends/megatron/core/extensions/primus_turbo_float8_local.py
+      # :: Float8{Column,Row}ParallelLinear._use_delayed_scaling.
+      fp8_recipe: "delayed"
+      fp8_margin: 0
+      fp8_amax_history_len: 1024
+      fp8_amax_compute_algo: "max"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+      fp8_reduce_amax: true
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 180
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_local_fp8
+      wandb_project: flux_12b_ddp_local_fp8
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+      use_dual_fp8_output_projection: false
+
+      seed: 2025
+      per_step_rng_reseed: false
+      nemo_chimera_init: false
+
+      # Torch Compile — compatible with Float8 local spec (per-module FP8)
+      torch_compile:
+        enable: true
+        strategy: "per_block"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        emulate_precision_casts: false
+        fused_ln_modulate: true
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml
new file mode 100644
index 000000000..12b11d34a
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml
@@ -0,0 +1,187 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TransformerEngine Spec + BF16 (MI300X)
+#
+# BF16 training using Megatron DDP + distributed optimizer with
+# TransformerEngine modules (TEColumnParallelLinear / TERowParallelLinear /
+# TEDotProductAttention / TENorm). This is the BF16 baseline on the
+# TransformerEngine path; see the *_te_spec_fp8 variant for FP8.
+#
+# Key settings:
+#   bf16: true
+#   params_dtype: bfloat16
+#   micro_batch_size: 32 / global_batch_size: 256
+#
+# Batch sizes are scaled down from the MI355X recipe (MBS=64/GBS=512) for
+# MI300X (192GB) headroom; tune to your hardware.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_bf16}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b_rope_fusion.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # TransformerEngine Spec (default)
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # ENERGON DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration (scaled down for MI300X; tune to your hardware)
+      micro_batch_size: 32
+      global_batch_size: 256
+      seq_length: 512
+
+      # ==========================================
+      # BF16 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings (identical to BF16 baseline)
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: true
+
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # Memory Optimizations
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: null
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_te_bf16
+      wandb_project: flux_12b_ddp_te_bf16
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # Torch Compile — stack strategy (compiles the double/single DiT block
+      # stacks as inductor regions). Matches the MI355X te_spec recipe and the
+      # MLPerf NeMo reference (COMPILE_DIT strategy=stack); reduces activation
+      # memory and improves throughput on the eager BF16 TE path.
+      torch_compile:
+        enable: true
+        strategy: "stack"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml
new file mode 100644
index 000000000..2df61605f
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml
@@ -0,0 +1,210 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MI300X)
+#
+# FP8 training using Megatron DDP + distributed optimizer with TransformerEngine
+# modules and FP8 hybrid delayed scaling:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm
+#   - FP8 hybrid (E4M3 fwd, E5M2 bwd) with delayed scaling (amax history 1024)
+#   - RoPE fusion via apply_rope_fusion: true
+#   - Energon pre-encoded dataset
+#
+# Required environment variables for MI300X (set before launch):
+#   export NVTE_FUSED_ATTN=1
+#   export NVTE_FUSED_ATTN_CK=1
+#   export NVTE_FP8_DPA_BWD=1
+#   export NVTE_USE_HIPBLASLT=1
+#   export USE_HIPBLASLT=1
+#   export TORCH_BLAS_PREFER_HIPBLASLT=1
+#   export NVTE_USE_CAST_TRANSPOSE_TRITON=1
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # ==========================================
+      # RoPE Fusion
+      # ==========================================
+      rotary_interleaved: true
+      apply_rope_fusion: true
+      position_embedding_type: rope
+
+      # ==========================================
+      # TransformerEngine Spec
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # FP8 — hybrid delayed scaling
+      # ==========================================
+      fp8: "hybrid"
+      fp8_recipe: "delayed"
+      fp8_margin: 0
+      fp8_amax_history_len: 1024
+      fp8_amax_compute_algo: "max"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+
+      # ==========================================
+      # Energon Dataset
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # ==========================================
+      # Batch Configuration
+      # ==========================================
+      # MI300X has 192GB HBM3 (vs 256GB on MI355X), so MBS reduced
+      # from 64 to 32 and GBS from 512 to 256 to avoid OOM with FP8.
+      micro_batch_size: 32
+      global_batch_size: 256
+      seq_length: 512
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # ==========================================
+      # Optimizer
+      # ==========================================
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # Memory / Misc
+      # ==========================================
+      use_flash_attn: true
+      empty_unused_memory_level: 0
+
+      # Manual GC — align GC timing across ranks to avoid stragglers
+      manual_gc: true
+      manual_gc_interval: 1000
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_te_fp8
+      wandb_project: flux_12b_ddp_te_fp8
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo — disabled for pure TE path
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # torch.compile — selective stack compilation for TE spec
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        strategy: "stack"
+        replace_qk_rmsnorm: true
+        disable_inductor_cudagraphs: false
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml
new file mode 100644
index 000000000..a4375f9a6
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml
@@ -0,0 +1,177 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — FSDP2 + Local Spec + BF16, VAE Resample Mode (MI300X)
+#
+# FSDP2 (ZeRO-2) BF16 training with the local spec provider and
+# vae_latent_mode: resample.
+#
+# In resample mode, latents are re-drawn from stored mean+logvar via
+# reparameterization (mean + exp(0.5*logvar) * randn) at every training step.
+# This introduces per-step stochasticity in the VAE latents.
+#
+# Dataset must be an Energon pre-encoded dataset containing mean.pth and
+# logvar.pth per sample.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_bf16}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboLocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # ENERGON DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 32
+      global_batch_size: 256
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # Mixed precision
+      bf16: true
+      fp16: false
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # PyTorch FSDP2 Configuration (ZeRO-2)
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+
+      torch_fsdp2_reshard_after_forward: false  # ZeRO-2
+
+      use_fsdp2_fp32_param_optimizer: true
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: false
+
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # Memory Optimizations
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: null
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_bf16
+      wandb_project: flux_12b_fsdp2_local_bf16
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+
+      seed: 2025
+
+      # Torch Compile
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: true
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml
new file mode 100644
index 000000000..64c2a6d1d
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml
@@ -0,0 +1,190 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — FSDP2 + Local Spec + FP8 + FP8 All-Gather (MI300X)
+#
+# FP8 training on the FSDP2 (ZeRO-3) path with:
+#   - FP8 training (tensorwise, local spec, dual FP8 output projection)
+#   - FP8 all-gather (keep weight in FP8 after FSDP2 all-gather)
+#   - BF16 master weight optimizer (BF16 params, FP32 optimizer states)
+#   - torch.compile enabled
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_fp8}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboLocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # ENERGON DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 52
+      global_batch_size: 416
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # Mixed precision
+      bf16: true
+      fp16: false
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # PyTorch FSDP2 Configuration (ZeRO-2)
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+
+      torch_fsdp2_reshard_after_forward: true  # ZeRO-3
+      use_fsdp2_fp8_all_gather: true
+      # fp8_all_gather_deq_requant: true  # Dequant FP8->BF16 after AG, fresh dynamic requant downstream
+      # use_triton_ops: true  # Use Triton @triton_op modulate/LN+modulate (compile-transparent)
+      fsdp_prefetch_depth: 1
+      fp8_precompute_data_cache: false
+      optimizer_foreach: false
+      use_cpp_fp8_quantize: true
+
+      # Optimizer mode: BF16 params + FP32 master weights
+      use_fsdp2_fp32_param_optimizer: false
+      use_fsdp2_bf16_master_weight_optimizer: true
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: false
+
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # FP8 Configuration — Tensorwise via Primus Turbo
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: "e4m3"
+      # Dynamic (tensorwise) scaling: the FSDP2 path does not yet exercise the
+      # delayed amax allreduce, so use tensorwise here. Switch to
+      # `fp8_recipe: "delayed"` once FSDP2 + delayed is wired up.
+      fp8_recipe: "tensorwise"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_fp8
+      wandb_project: flux_12b_fsdp2_local_fp8
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+      use_dual_fp8_output_projection: true
+
+      seed: 2025
+
+      # Torch Compile — compatible with Float8 local spec (per-module FP8)
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: true
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml
new file mode 100644
index 000000000..865cfe832
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml
@@ -0,0 +1,170 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M Pretraining Configuration (Pre-encoded Data Mode)
+#
+# This config demonstrates Flux 535M training with pre-encoded features.
+# Pre-encoded mode is faster and recommended for production training.
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+#   bash examples/run_pretrain.sh
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_535m.yaml
+
+    # Trainer class for diffusion models
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      # ============================================================================
+      # Model Configuration
+      # ============================================================================
+
+      model_type: diffusion_model
+
+      # ============================================================================
+      # Dataset Configuration — Synthetic / Mock Data (default)
+      # ============================================================================
+      #
+      # This 535M config defaults to in-memory synthetic data so it runs
+      # standalone for sanity checks, CI, and training-pipeline validation
+      # without a prepared Energon dataset. Samples are random tensors with
+      # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768).
+      #
+      # To train on real data instead: set `mock_data: false` and point
+      # `data_path` at a prepared Energon dataset (see notes at the bottom).
+      mock_data: true
+      mock_dataset:
+        class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+        params:
+          num_samples: 256         # In-memory synthetic samples (iterated cyclically)
+          image_size: 256          # 256 -> 32x32 latents (light/fast for testing)
+      # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Batch configuration
+      micro_batch_size: 2          # Per-GPU batch size (adjust based on VRAM)
+      global_batch_size: 16        # Total batch size across all GPUs
+      seq_length: 4096             # Sequence length for latent features
+
+      # DataLoader settings (used only with real Energon data, i.e. mock_data: false)
+      num_workers: 4               # Number of data loading workers per GPU
+      dataloader_type: external    # Required for Energon dataloaders
+
+      # ============================================================================
+      # Training Parameters
+      # ============================================================================
+
+      # Total training steps
+      train_iters: 100000          # Total training iterations
+      eval_interval: 1000          # Evaluate every N steps
+      eval_iters: 50               # Number of evaluation iterations
+
+      # Logging
+      log_interval: 10             # Log every N steps
+      tensorboard_dir: output/tensorboard/flux_535m_pretrain
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+
+      # Checkpointing (disabled)
+      save_interval: 5000          # Save checkpoint every N steps
+      save: null
+      load: null                   # Path to checkpoint for resuming (optional)
+      finetune: false
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Precision
+      bf16: true                   # Use bfloat16 (recommended for AMD MI300X)
+      fp16: false
+
+      # ============================================================================
+      # Optimizer Configuration
+      # ============================================================================
+
+      # Optimizer
+      optimizer: adam
+      lr: 1.0e-4                   # Peak learning rate
+      min_lr: 1.0e-5               # Minimum learning rate
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1.0e-8
+
+      # Gradient clipping
+      clip_grad: 1.0
+
+      # ============================================================================
+      # Learning Rate Scheduler
+      # ============================================================================
+
+      lr_decay_style: cosine
+      lr_warmup_iters: 1000        # Warmup iterations
+      lr_decay_iters: 100000       # Total decay steps (typically = train_iters)
+
+      # ============================================================================
+      # Distributed Training Configuration
+      # ============================================================================
+
+      # Parallelism settings (Flux 535M fits on 1 GPU, but can use DP for speed)
+      tensor_model_parallel_size: 1    # Tensor parallelism (no need for 535M)
+      pipeline_model_parallel_size: 1  # Pipeline parallelism
+
+      # Advanced settings
+      overlap_grad_reduce: true         # Overlap gradient communication
+      use_flash_attn: true              # Use Flash Attention 2
+      distributed_timeout_minutes: 60
+
+      # ============================================================================
+      # Seed and Reproducibility
+      # ============================================================================
+
+      seed: 42
+
+      # ============================================================================
+      # Monitoring and Logging
+      # ============================================================================
+
+      wandb_project: flux_535m_pretrain
+      wandb_exp_name: flux_535m_preencoded
+
+# ============================================================================
+# Notes
+# ============================================================================
+#
+# Dataset Preparation:
+#   1. Prepare pre-encoded dataset:
+#      tools/docker/primus data diffusion-encoded \
+#        --source-type directory --input-dir /data/raw \
+#        --output-dir /data/encoded --model-path black-forest-labs/FLUX.1-dev
+#   2. Copy dataset template to output directory:
+#      cp primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml \
+#         /data/encoded/dataset.yaml
+#   3. Run Energon indexing:
+#      energon prepare /data/encoded --num-workers 8
+#   4. Update data_path above to: /data/encoded/dataset.yaml
+#
+# For more information, see:
+#   - primus/configs/data/megatron/diffusion/README.md
+#   - examples/megatron/diffusion/README.md
+#
+# Single-node training (8 GPUs):
+#   EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+#   GPUS_PER_NODE=8 bash examples/run_pretrain.sh
+#
+# Multi-node training (4 nodes, 8 GPUs each):
+#   EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+#   NNODES=4 bash examples/run_slurm_pretrain.sh
+#
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml
new file mode 100644
index 000000000..fa4aa8da7
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml
@@ -0,0 +1,199 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M FP8 Pretraining Configuration (Testing/Development)
+#
+# This config is for testing FP8 functionality with the minimal Flux 535M model
+# before scaling to the full 12B model. Use this to validate:
+#   - FP8 setup and configuration
+#   - Numerical stability
+#   - Memory and speed improvements
+#   - Transformer Engine compatibility
+#
+# Target Hardware: Single AMD MI300X GPU with ROCm 6.0+
+# Requires: Transformer Engine 2.1.0+ with ROCm backend
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \
+#   GPUS_PER_NODE=1 bash examples/run_pretrain.sh
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain_fp8}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_535m_fp8.yaml  # Use FP8-enabled 535M config
+
+    # Trainer class for diffusion models
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      # ============================================================================
+      # Model Configuration
+      # ============================================================================
+
+      model_type: diffusion_model
+
+      # ============================================================================
+      # Dataset Configuration — Synthetic / Mock Data (default)
+      # ============================================================================
+      #
+      # This 535M config defaults to in-memory synthetic data so it runs
+      # standalone for sanity checks, CI, and training-pipeline validation
+      # without a prepared Energon dataset. Samples are random tensors with
+      # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768).
+      #
+      # To train on real data instead: set `mock_data: false` and point
+      # `data_path` at a prepared Energon dataset (see notes at the bottom).
+      mock_data: true
+      mock_dataset:
+        class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+        params:
+          num_samples: 256         # In-memory synthetic samples (iterated cyclically)
+          image_size: 256          # 256 -> 32x32 latents (light/fast for testing)
+      # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Batch configuration (can use larger batch with FP8)
+      micro_batch_size: 4          # Can increase from 2 to 4 with FP8 on 535M
+      global_batch_size: 32        # Small batch for quick testing
+      seq_length: 4096             # Sequence length for latent features
+
+      # DataLoader settings (used only with real Energon data, i.e. mock_data: false)
+      num_workers: 4               # Number of data loading workers per GPU
+      dataloader_type: external    # Required for Energon dataloaders
+
+      # ============================================================================
+      # Training Parameters (Quick Testing)
+      # ============================================================================
+
+      # Short training for validation
+      train_iters: 1000            # Just 1K steps for FP8 validation
+      eval_interval: 100           # Evaluate every 100 steps
+      eval_iters: 10               # Quick evaluation
+
+      # Logging
+      log_interval: 10             # Log every N steps
+      tensorboard_dir: output/tensorboard/flux_535m_pretrain_fp8
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+
+      # Checkpointing (disabled)
+      save_interval: 500           # Save more frequently for testing
+      save: null
+      load: null                   # Path to checkpoint for resuming (optional)
+      finetune: false
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Precision
+      bf16: true                   # Use bfloat16 for non-FP8 ops
+      fp16: false
+
+      # ============================================================================
+      # Optimizer Configuration
+      # ============================================================================
+
+      # Optimizer
+      optimizer: adam
+      lr: 1.0e-4                   # Peak learning rate
+      min_lr: 1.0e-5               # Minimum learning rate
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1.0e-8
+
+      # Gradient clipping (important for FP8 stability)
+      clip_grad: 1.0
+
+      # ============================================================================
+      # Learning Rate Scheduler
+      # ============================================================================
+
+      lr_decay_style: cosine
+      lr_warmup_iters: 100         # Short warmup for testing
+      lr_decay_iters: 1000         # Match train_iters
+
+      # ============================================================================
+      # Distributed Training Configuration
+      # ============================================================================
+
+      # Single GPU configuration
+      tensor_model_parallel_size: 1    # No TP needed for 535M
+      pipeline_model_parallel_size: 1  # No PP needed
+      context_parallel_size: 1         # No CP needed
+
+      # Distributed settings
+      use_distributed_optimizer: false  # Not needed for single GPU
+      overlap_grad_reduce: false        # Not applicable for single GPU
+      use_flash_attn: true              # Use Flash Attention 2
+      distributed_timeout_minutes: 60
+
+      # ============================================================================
+      # Memory Optimization
+      # ============================================================================
+
+      # Activation checkpointing (not needed for 535M with FP8)
+      recompute_granularity: null       # No recompute needed
+      recompute_method: null
+      recompute_num_layers: null
+
+      # Sequence parallelism
+      sequence_parallel: false          # Not needed for single GPU
+
+      # ============================================================================
+      # Seed and Reproducibility
+      # ============================================================================
+
+      seed: 42
+
+      # ============================================================================
+      # Monitoring and Logging
+      # ============================================================================
+
+      wandb_project: flux_535m_pretrain_fp8
+      wandb_exp_name: flux_535m_fp8_test
+
+# ============================================================================
+# Notes - FP8 Validation with 535M
+# ============================================================================
+#
+# Hardware Requirements (with FP8):
+#   - Minimum: 1× MI300X 192GB
+#   - Memory per GPU: ~3-5GB (vs ~7-10GB BF16)
+#   - Training time: Minutes
+#
+# Validation Checklist:
+#   [ ] Setup FP8 environment (see docs/04-technical-guides/diffusion-models/fp8_training.md)
+#   [ ] Verify TE FP8 support is available
+#   [ ] Run this config to validate FP8 training
+#   [ ] Check logs for NaN/Inf (should be none)
+#   [ ] Verify memory usage is ~50% of BF16
+#   [ ] Verify training speed is 1.5-2x faster than BF16
+#   [ ] Check loss decreases normally
+#
+# Expected Results:
+#   - Training completes 1000 steps in 5-15 minutes
+#   - No NaN/Inf in losses
+#   - Memory usage: ~3-5GB
+#   - Speed: ~10-50 steps/sec (depending on hardware)
+#   - Loss should decrease normally
+#
+# If validation passes, proceed to one of the 12B FP8 configs, e.g.:
+#   - flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml (TransformerEngine FP8)
+#   - flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml (local-spec FP8)
+#
+# Troubleshooting:
+#   - If NaN/Inf: Check Transformer Engine FP8 support
+#   - If OOM: Reduce micro_batch_size
+#   - If slow: Verify ROCm FP8 tensor cores are being used
+#   - If unstable: Try fp8_wgrad: false in model config
+#
+# For more information: See docs/04-technical-guides/diffusion-models/fp8_training.md
diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml
new file mode 100644
index 000000000..38f10f88d
--- /dev/null
+++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml
@@ -0,0 +1,58 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M with Guidance Embedding (Advanced Configuration)
+#
+# This config demonstrates Flux training with guidance embedding enabled.
+# This is an OPTIONAL advanced feature that allows for faster single-pass CFG
+# during inference, but requires training with guidance embedding enabled.
+#
+# IMPORTANT: Most users should use the standard flux_535m_pretrain.yaml config.
+# Only use this if you specifically need guidance embedding support.
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml \
+#   bash examples/run_pretrain.sh
+
+# Extend standard 535M config
+extends:
+  - flux_535m_pretrain.yaml
+
+modules:
+  pre_trainer:
+    overrides:
+      # ============================================================================
+      # Guidance Embedding Configuration (ADVANCED)
+      # ============================================================================
+
+      # Enable guidance embedding for single-pass CFG
+      # This adds a learned MLPEmbedder layer that conditions on guidance scale
+      guidance_embed: true
+
+      # Guidance scale used during training
+      # Model learns to adapt its predictions based on this scale
+      guidance_scale: 3.5
+
+      # ============================================================================
+      # Notes
+      # ============================================================================
+      #
+      # Training with guidance embedding:
+      # - Adds ~1-2% more parameters (guidance MLPEmbedder)
+      # - Allows single-pass CFG during inference (faster)
+      # - Requires more training data/iterations to converge
+      # - Model learns guidance as a conditioning signal
+      #
+      # Inference with guidance embedding:
+      # - Pipeline automatically detects guidance_embed layer
+      # - Uses single forward pass instead of batch doubling
+      # - ~2x faster CFG compared to explicit CFG
+      # - Guidance scale can be varied at inference time
+      #
+      # Standard approach (guidance_embed: false):
+      # - Default for most Primus training
+      # - Uses explicit CFG (batch doubling) at inference
+      # - More compatible with existing checkpoints
+      # - Slightly slower but more flexible
+      #
+      # See examples/megatron/diffusion/README.md for more details.
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml
index ca65bb754..f741c08fe 100644
--- a/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml
@@ -20,7 +20,7 @@ modules:
       log_avg_skip_iterations: 2
       log_avg_reset_interval: 50
 
-      train_iters: 100
+      train_iters: 50
       micro_batch_size: 8
       global_batch_size: 64
 
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_gdn-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn-pretrain.yaml
new file mode 100644
index 000000000..a45aed5bb
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn-pretrain.yaml
@@ -0,0 +1,83 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_gdn-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: zebra_llama_1B_gdn.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_Zebra_Llama_1B_GDN_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 0
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      train_iters: 50
+      micro_batch_size: 4
+      global_batch_size: 32
+
+      seq_length: 8192
+      max_position_embeddings: 8192
+      original_max_position_embeddings: 8192
+
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 38147
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+
+      # Pure GDN hybrid spec
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec']
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: fla-hub/gla-1.3B-100B
+
+      # parallel
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: true
+
+      # data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      save: ./output/zebra_llama_1B_gdn-pretrain
+      save_interval: 1000
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure-pretrain.yaml
new file mode 100644
index 000000000..56b242bd7
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure-pretrain.yaml
@@ -0,0 +1,93 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_gdn_pure-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_1B_gdn_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_1B_GDN_Pure_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 0
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # Training schedule — matched to FLA (4 GPUs):
+      #   batch=16, update=1, gpus=4, context=2048
+      #   global_batch = 16 * 4 = 64, tokens/step = 64 * 2048 = 131,072
+      #   total tokens ~= 76294 * 131072 ≈ 10B
+      train_iters: 50
+      micro_batch_size: 16
+      global_batch_size: 128
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — matched to FLA:
+      #   lr=2e-4, cosine (min_lr_rate=0.1 → min_lr=2e-5)
+      #   AdamW, beta1=0.9, beta2=0.95, weight_decay=0.01
+      #   max_grad_norm=1.0, warmup_steps=200
+      # Both DeepSpeed ZeRO-2 and Megatron clip on ||avg_grad||.
+      # DeepSpeed pre-divides by dp_size before reduce_scatter(SUM),
+      # so its grad norm is on averaged gradients. Match FLA's 1.0 directly.
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 76294
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Pure GDN hybrid spec (hybrid_attention_ratio=0.0 → all GDN, no MLA)
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec']
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: true
+
+      # Data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints
+      finetune: false
+      auto_continue_train: true
+      load: null
+      save: ./output/zebra_llama_1B_gdn_pure-pretrain
+      save_interval: 2048
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_100B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_100B-pretrain.yaml
new file mode 100644
index 000000000..2f2e2048b
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_100B-pretrain.yaml
@@ -0,0 +1,498 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_gdn_pure_100B-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Pure GDN 1B — 100B tokens on FineWeb-Edu (matched to FLA's
+# `setup_and_train_gdn_pure_1B_100B.sh` so the loss curves can be compared
+# iter-for-iter once we drive Primus from FLA's token order via PRIMUS_FLA_DATA).
+#
+# FLA's run:
+#   model:       configs/gated_deltanet_1B_pure_100B.json (same arch as the
+#                10B config we already replicate in zebra_llama_1B_gdn_pure.yaml)
+#   lr:          3e-4   (vs the 10B run's 2e-4 — bigger LR for the longer schedule)
+#   scheduler:   cosine_with_min_lr  (min_lr ≈ 0.1 × peak = 3e-5)
+#   warmup:      2000 iters         (vs 200 on the 10B run)
+#   batch:       64 per GPU         (vs 16 on the 10B run)
+#   update:      1                  (no grad-accum)
+#   gpus:        8
+#     → global_batch_size = 64 × 8 = 512
+#     → tokens/iter      = 512 × 2048 = 1,048,576
+#   steps:       95368              (95368 × 1,048,576 ≈ 100B tokens)
+#   data:        HuggingFaceFW/fineweb-edu, sample-100BT
+#                cache: .../data/HuggingFaceFW/fineweb-edu/sample-100BT/train
+#                (~364 GB on disk — pre-tokenized by FLA)
+#
+# Wall-time estimate on 8×MI300X with the full FLA-parity stack (~1.8 s/iter
+# based on the GDN-hybrid 300M @ 2.2 s/iter rescaled for the larger model):
+#   95368 × ~1.8 s ≈ 48 h (~2 days)
+# ─────────────────────────────────────────────────────────────────────────────
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_1B_gdn_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_1B_GDN_Pure_100B"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      # FLA's `train.sh` defaults `--dataloader_num_workers 32` and the 100B
+      # run does NOT override it. At b=64 mbs, ~131k tokens/GPU/iter, 32
+      # workers (vs Primus's previous 8 carried over from the 300M run) keeps
+      # the DataLoader queue full so the GPU never starves between iters.
+      # PyTorch's default `prefetch_factor=2` is already what FLA uses.
+      num_workers: 32
+      create_attention_mask_in_dataloader: false
+
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # Perf: Megatron defaults to inserting dist.barrier() before every L1 timer
+      # measurement (~5-10/iter). Costs throughput, no correctness impact.
+      barrier_with_L1_time: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # SPEED RECOVERY (lifted from 300M YAML which matched FLA at +0.27%):
+      #
+      # The profile of iters 15-16 showed 487 ms/iter spent in `aten::item`
+      # (CPU↔GPU sync points). The 300M run avoided this by disabling NaN
+      # checking and bumping log_interval to 100. Both are loss-identical
+      # — only affect logging and sync, not training math.
+      #
+      # check_for_nan_in_loss_and_grad: false
+      #   • Megatron default = true, which fires `loss.item()` and an
+      #     `isnan`/`isinf` validate_result on the loss EVERY iter
+      #     (pretrain_mamba.py:178). Each is a CPU-GPU sync. Saves ~50 ms.
+      #   • If loss ever spikes to NaN you'll still see it in the log
+      #     output, just won't auto-rerun the step.
+      #
+      # log_interval: 32
+      #   • Default is 1, meaning every step we run timers, compute
+      #     throughput, `loss.item()`, and print to stdout/log files —
+      #     ~200 ms of sync+formatting per iter. Logging every 32 iters
+      #     EXACTLY matches FLA's `--logging_steps 32` (legacy/training/
+      #     train.sh:55) so loss curves line up tick-for-tick with the
+      #     FLA reference log when overlaid for parity verification.
+      #     At ~7.5 s/iter that's a log line every 4 min — plenty.
+      #
+      # tensorboard_log_interval: 32
+      #   • Same logic for TB/wandb logging. Default is 1 = sync every
+      #     iter. ~100 ms recovered.
+      #
+      # Combined expected speedup: ~300-400 ms/iter (7.6 s → ~7.2 s),
+      # all loss-identical. This is what the 300M YAML had.
+      # ─────────────────────────────────────────────────────────────────────
+      check_for_nan_in_loss_and_grad: false
+      # FLA's 100B launcher uses `logging=10` → `--logging_steps 10`. Matching
+      # exactly so the Primus loss log lines up tick-for-tick with FLA's
+      # reference `train_gdn_pure_1B_100B.log`.
+      log_interval: 10
+      tensorboard_log_interval: 10
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA-parity numerics (lifted verbatim from the validated 300M parity
+      # YAML — these are the *exact* settings that produced the iter-1
+      # bit-perfect match and the < 0.5% late-training residual documented
+      # in docs/hybrid_models/GDN_FLA_PARITY.md).  Skipping any of them re-introduces a known
+      # source of drift:
+      #
+      #   layernorm_epsilon  — Megatron's TransformerConfig defaults to 1e-5;
+      #                        FLA uses 1e-6.  ~1% per-layer divergence.
+      #   hidden_dropout     — Megatron's TransformerConfig defaults these to
+      #   attention_dropout    0.1 each (transformer_config.py L152).  Even
+      #                        though the model YAML asks for 0, the EXP-level
+      #                        defaults leak through and we'd train with 10%
+      #                        dropout while FLA trains with 0.  CRITICAL.
+      #   no_persist_layer_norm — disables Apex's persistent buffer kernel
+      #                        which has slightly different rounding.
+      # ─────────────────────────────────────────────────────────────────────
+      layernorm_epsilon: 1.0e-6
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+      no_persist_layer_norm: true
+
+      # ── Training schedule ─ matches FLA EXACTLY (bit-perfect every knob):
+      #   batch=64, grad_accum=1, gpus=8, seq=2048
+      #   global_batch_size = 64 × 8 = 512
+      #   tokens/iter       = 512 × 2048 = 1,048,576
+      #   total tokens      = 95368 × 1,048,576 ≈ 100 B
+      train_iters: 50
+      micro_batch_size: 64
+      global_batch_size: 512
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # ── Optimizer ─ exact bit-match to FLA's 100B CLI:
+      #   --learning_rate 3.0e-4
+      #   --lr_scheduler_type cosine_with_min_lr
+      #   --warmup_steps 2000
+      #   --optim adamw_torch_fused
+      #   --adam_beta1=0.9 --adam_beta2=0.95
+      #   --weight_decay 0.01
+      #   --max_grad_norm 1.0
+      #   --seed 42 --bf16
+      #
+      # CRITICAL: although FLA's train.sh does NOT pass --lr_scheduler_kwargs,
+      # FLA's run.py OVERRIDES it programmatically — see
+      # flash-linear-attention/legacy/training/run.py:96-97 :
+      #     if args.lr_scheduler_type == 'cosine_with_min_lr':
+      #         args.lr_scheduler_kwargs = {'min_lr_rate': 0.1}
+      # So FLA's lr decays from 3e-4 to 0.1*3e-4 = 3e-5 over 95368 iters,
+      # NOT to 0 as the HF default would imply.  Megatron's cosine is
+      #     lr = min_lr + 0.5*(1+cos(π*p))*(max_lr - min_lr)
+      # HF's cosine_with_min_lr is the algebraically identical form
+      #     lr = peak * (factor*(1-min_lr_rate) + min_lr_rate)
+      # so setting `min_lr: 3.0e-5` here matches FLA's lr at every step,
+      # bit-exact.  (Setting `min_lr: 0.0` — what we had earlier — would
+      # cause Primus's lr to undershoot FLA's in late training; the loss
+      # curves would diverge after ~iter 50000.)
+      # (Warmup is `init_lr + (max_lr-init_lr)*step/warmup` in Megatron and
+      # `peak*step/warmup` in HF — identical when init_lr=0, which is the
+      # Megatron default.)
+      clip_grad: 1.0
+      lr: 3.0e-4
+      min_lr: 3.0e-5
+      lr_warmup_iters: 2000
+      lr_decay_iters: 95368
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Match FLA's --seed 42 exactly.  Primus's default is 1234; using FLA's
+      # seed minimises init-RNG drift between the two runs (still not bit-
+      # identical because Megatron and HF transformers walk parameters in
+      # different orders, but it eliminates the 12 % of variance that comes
+      # from the seed itself).
+      seed: 42
+
+      # ─────────────────────────────────────────────────────────────────────
+      # NO-TE spec — the validated 300M parity run uses this variant because
+      # FLA's reference is built on native PyTorch nn.Linear / RMSNorm.  TE's
+      # ColumnParallelLinear / TENorm wrappers introduce small numerical
+      # differences (cast ordering, persistent buffers) that drift the loss
+      # by ~0.1% per layer.  Using the no-TE spec aligns Megatron's layer
+      # numerics with FLA's.  (See docs/hybrid_models/GDN_FLA_PARITY.md §A.)
+      # ─────────────────────────────────────────────────────────────────────
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+
+      # Tokenizer (cached locally — same path FLA's run uses)
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # ── Parallelism ─
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+
+      # ─────────────────────────────────────────────────────────────────────
+      # DDP / optimizer sharding — Megatron-FSDP (ZeRO-2 equivalent)
+      # ─────────────────────────────────────────────────────────────────────
+      #
+      # WHY we switched from ZeRO-1 to Megatron-FSDP (this run only):
+      #
+      # Profile of iter 15-16 (captured 2026-05-25) showed:
+      #
+      #   GPU busy:                   ~43 %  (TFLOP/s = 261, peak 1300)
+      #   aten::empty                 1500 calls × 2.3 ms = 3390 ms / iter
+      #   aten::empty_like            1206 calls × 2.6 ms = 3114 ms / iter
+      #   aten::empty_strided         1100 calls × 2.3 ms = 2575 ms / iter
+      #   aten::mm dispatch            576 calls × 9.5 ms = 5485 ms / iter
+      #
+      # Normal aten::empty is *microseconds*. Ours is **1000× slower**
+      # because the PyTorch HIP allocator is constantly hitting hipMalloc/
+      # hipFree at 81 % VRAM (156 GB / 192 GB). For comparison the 300M
+      # run runs at 30 % VRAM, the allocator cache stays warm, every
+      # aten::empty is sub-µs, and it achieves 643 TFLOP/s (49 % of peak).
+      #
+      # Megatron's own FSDP docs explicitly call out this exact failure
+      # mode (third_party/Megatron-LM/docs/user-guide/features/
+      # custom_fsdp.md §5 line 187):
+      #
+      #   "FSDP can lead to crashes of the PyTorch memory allocator cache,
+      #    and a large number of cudaMalloc and cudaFree calls. This
+      #    problem is challenging and can only be mitigated by avoiding
+      #    frequent hits on the GPU memory limit."
+      #
+      # FLA's reference run uses DeepSpeed ZeRO-Stage-2 (ds_config.json
+      # "stage": 2, contiguous_gradients=true) which shards BOTH optimizer
+      # state AND gradients. Our prior config used Megatron's
+      # `use_distributed_optimizer=true` which is ZeRO-Stage-1 (only opt
+      # state sharded — full grads stay on every rank, ≈ 4.8 GB / rank).
+      # Megatron-FSDP with `data_parallel_sharding_strategy: optim_grads`
+      # is the bit-for-bit equivalent of FLA's setup: shards both opt-state
+      # AND grads, leaves params un-sharded. Memory savings ≈ 4.2 GB/rank
+      # (from 156 → ~152 GB, 81 % → 79 %), which along with removing the
+      # per-iter `empty_cache()` (see empty_unused_memory_level below)
+      # should give the allocator enough headroom to stop thrashing.
+      #
+      # Why optim_grads (ZeRO-2) NOT optim_grads_params (ZeRO-3):
+      #   • optim_grads_params shards params too, but Megatron-FSDP only
+      #     auto-wraps `TransformerLayer` instances as FSDP units. Our
+      #     GDN layers are `MambaLayer` (a separate Megatron class), so
+      #     they would NOT be FSDP-wrapped and would silently fall back
+      #     to no-sharding for the GDN half of the model.
+      #   • FLA itself uses ZeRO-2, not ZeRO-3.
+      #   • Loss numerics are identical to FLA only with ZeRO-2 (ZeRO-3
+      #     introduces a tiny float-rounding difference in the per-shard
+      #     param reconstruction).
+      #
+      # Compatibility caveats verified:
+      #   ✓ Megatron-FSDP requires use_distributed_optimizer=true (kept)
+      #   ✓ Requires gradient_accumulation_fusion=false (kept)
+      #   ✓ Auto-handles CUDA_DEVICE_MAX_CONNECTIONS via Primus's
+      #     env_patches.py:set_cuda_device_max_connections() — sets to
+      #     8 when use_megatron_fsdp=true (vs 1 for non-FSDP)
+      #   ✓ GDN layer has reset_parameters() (megatron/core/ssm/
+      #     gated_delta_net.py:227), so meta-device init would work too
+      #     (but we leave init_model_with_meta_device=false for safety)
+      #   ✓ Prior segfault (overlap_grad_reduce=true + ZeRO-1) was NOT
+      #     about FSDP — it was a separate Megatron bug in the ZeRO-1
+      #     bucket layout for heterogeneous Mamba params. FSDP uses its
+      #     own (different) buffer code path.
+      #
+      # If FSDP causes any crash, revert these 2 lines to restore ZeRO-1:
+      #   use_megatron_fsdp: false
+      #   data_parallel_sharding_strategy: no_shard
+      # (Rest of config is FSDP-safe AND ZeRO-1-safe.)
+      # ─────────────────────────────────────────────────────────────────────
+      use_distributed_optimizer: true
+      # ── Megatron-FSDP ATTEMPTS (2026-05-25 + 2026-05-26) ─────────────────
+      # Take #1 (2026-05-25): `use_megatron_fsdp: true` +
+      #   `data_parallel_sharding_strategy: optim_grads` crashed at iter-1
+      #   post-step NaN-check all_reduce with `Failed to CUDA calloc
+      #   268435456 bytes` (256 MiB).
+      #
+      #   Root cause (diagnosed 2026-05-26): RCCL's lazy comm-init allocates
+      #   BUFFSIZE × #channels per NCCL communicator. Defaults =
+      #   4 MiB × 64 channels = 256 MiB per comm. Megatron-FSDP creates
+      #   extra comm groups (HSDP outer + DP inner) on top of the one
+      #   ZeRO-1 needs, and at 99% VRAM there is no contiguous 256 MiB hole.
+      #
+      # Take #2 (2026-05-26): added NCCL channel clamp via launcher env vars
+      #   (NCCL_MIN_NCHANNELS=1 NCCL_MAX_NCHANNELS=4 NCCL_NCHANNELS_PER_PEER=1
+      #    + ckpt_format: fsdp_dtensor since Megatron asserts it).
+      #   Result: 500 iters complete, exit 0, ZERO crashes.
+      #
+      #   Steady-state iter time: 2.84 s/iter (vs ZeRO-1's 2.75 s/iter).
+      #   FSDP wins on non-flush iters (2.67 s vs 2.75 s) but loses on
+      #   empty_cache_interval=32 flushes (3.18 s vs ~2.98 s) because FSDP
+      #   has to re-allgather sharded params after the cache drop, while
+      #   ZeRO-1 keeps full param replicas resident.
+      #
+      #   NET (Take #2): FSDP 2.84 s/iter vs ZeRO-1 2.75 s/iter (90 ms slower).
+      #
+      # Take #3 (2026-05-26): Controlled 50-iter experiment EXP5 with
+      #   `empty_cache_interval` raised from 32 to 128 (one flush per 128
+      #   iters instead of one per 32):
+      #     - ZeRO-1 baseline (EXP0):  2768 ms steady-state
+      #     - FSDP + flush@128 (EXP5): 2684 ms steady-state  (-83 ms, -3.0%)
+      #     - 100B run savings: 95k iters × 83 ms = 2.2 hours
+      #     - Loss bit-identical to ZeRO-1 at iter 50: 11.7538
+      #
+      #   Why this works: between flushes, FSDP's per-iter allgather work is
+      #   pipelined into the backward pass (Megatron-FSDP enables this
+      #   internally even though `overlap_param_gather` is off). On flush
+      #   iters FSDP still pays the re-allgather cost, but at 1-in-128 the
+      #   amortized penalty is ~6 ms/iter vs the ~85 ms/iter saved.
+      #
+      # EXP6 (ZeRO-3 / optim_grads_params) FAILED at iter 1 with:
+      #   "RuntimeError: Pointer argument cannot be accessed from Triton
+      #    (cpu tensor?)"
+      # Because ZeRO-3 shards parameters across ranks, and the FLA Triton
+      # kernel cannot dereference a sharded DTensor. ZeRO-2 (optim_grads)
+      # is the maximum sharding compatible with our Triton GDN.
+      #
+      # FSDP working config + full experiment results:
+      #   experiments/results/SYNTHESIS.md
+      #   experiments/run_perf_exp.sh
+      #   examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_exp5-fsdp-rareflush.yaml
+      # ──────────────────────────────────────────────────────────────────
+      use_megatron_fsdp: true
+      data_parallel_sharding_strategy: optim_grads
+      # 2026-05-26 EXP7: FSDP path's overlap is a DIFFERENT code path than
+      # ZeRO-1's (which segfaulted in EXP4). With both flags on:
+      #   ZeRO-1 baseline:        2768 ms/iter (+19.8% vs FLA)
+      #   FSDP+rare flush (EXP5): 2684 ms/iter (+16.2% vs FLA)
+      #   FSDP+overlap (EXP7):   2414 ms/iter (+4.5% vs FLA)  ← winner
+      # 100B run savings: 95k iters × (2768-2414) ms = ~9.3 hours vs old prod
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: false
+      ddp_average_in_collective: true
+      use_torch_fsdp2: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # PyTorch profiler — already captured iter 15-17 trace which showed
+      # 487 ms/iter in `aten::item` (the 300M-parity log_interval/NaN-check
+      # fixes above target exactly this). Disabling for this run.
+      #
+      # To re-enable: set `profile: true`, `use_pytorch_profiler: true`
+      # and pick `profile_step_start`/`profile_step_end`.
+      # ─────────────────────────────────────────────────────────────────────
+      profile: false
+      use_pytorch_profiler: false
+      tensorboard_dir: output/amd/root/zebra_llama_1B_gdn_pure_100B-pretrain/tensorboard
+
+      # ─────────────────────────────────────────────────────────────────────
+      # MEMORY: empty_unused_memory_level + empty_cache_interval
+      # ─────────────────────────────────────────────────────────────────────
+      #
+      # empty_unused_memory_level=1 → calls `torch.cuda.empty_cache()` at
+      # training.py:1759, IMMEDIATELY before optimizer.step() (which
+      # internally calls get_grad_norm_fp32() → torch.distributed.all_reduce
+      # in clip_grads.py:130).
+      #
+      # That all_reduce is where NCCL crashes if we disable empty_cache
+      # entirely:
+      #   "Failed to CUDA calloc 4194304 bytes" inside ProcessGroupNCCL.
+      # NCCL's calloc bypasses PyTorch's allocator and goes straight to
+      # hipMalloc — at 81% VRAM there is no contiguous 4 MiB block free
+      # unless PyTorch first returns its cached pool to the driver.
+      #
+      # 2026-05-26 — PROFILE + FIX:
+      # PyTorch profiler attribution on iter 21 of the live 100B run
+      # measured 6,937 ms of the 7,650 ms iter (91%) in hipMalloc + hipFree,
+      # because the per-iter empty_cache() was thrashing ~250 cached
+      # blocks per iter against the ROCm driver (~17 ms per hipMalloc).
+      #
+      # FIX: keep `empty_unused_memory_level: 1` (Megatron's gate) but
+      # only fire empty_cache() every N iters via the Primus patch
+      # primus/backends/megatron/patches/empty_cache_interval_patches.py.
+      # The patch reads `empty_cache_interval` from this YAML (preferred),
+      # then falls back to the PRIMUS_EMPTY_CACHE_INTERVAL env var, then
+      # to the default 1 (passthrough).
+      #
+      # Iter 0 always fires empty_cache() so NCCL's first workspace alloc
+      # happens against a clean cache.  Iters 1..N-1 skip it, iter N fires
+      # again, etc.  NCCL workspace persists across the gap because once
+      # allocated it is not freed by PyTorch's allocator (NCCL owns the
+      # block).
+      #
+      # Measured impact (50-iter diag at empty_cache_interval=32,
+      # 2026-05-26):
+      #   before fix:  7.65 s/iter  ( 262 TFLOP/s/GPU,  20% of peak)
+      #   after fix:   2.79 s/iter  ( 727 TFLOP/s/GPU,  56% of peak)
+      #   speedup:     2.74×
+      #   loss curve:  iter-1 bit-identical to baseline (no math impact)
+      #   100B wall:   was 8.4 days → now 3.08 days  (~5.3 days saved)
+      #
+      # Knob choice:
+      #   - 1   = original per-iter behaviour (safest, slowest)
+      #   - 32  = our default (best tested point; ~2.8× speedup)
+      #   - 64+ = even fewer flushes; tested briefly, slightly faster
+      #           between flushes but bigger spike on the trigger iter
+      #   - 0   = NEVER flush (CRASHES at iter-1 NCCL alloc; do not use)
+      # ─────────────────────────────────────────────────────────────────────
+      empty_unused_memory_level: 1
+      empty_cache_interval: 128   # 2026-05-26 EXP5: 32→128 saves 83 ms/iter (-3% wall, -2.2 h on 100B run)
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA runtime knobs — declarative YAML surface for the PRIMUS_FLA_* /
+      # PRIMUS_FUSED_CE* env vars (consumed by the
+      # primus.backends.megatron.patches.fla_runtime_patches patch which
+      # re-exports them as env vars at phase="build_args").  These
+      # exactly mirror the values the legacy launcher script sets via
+      # `export PRIMUS_FLA_*` — keeping them here makes the YAML the
+      # single source of truth.  Env vars set on the launcher still win
+      # over the YAML (backward compat).
+      # ─────────────────────────────────────────────────────────────────────
+      use_fla_fused_swiglu: true       # FLA Triton SwiGLU (matches fla.modules.swiglu)
+      use_fla_fused_rmsnorm: true      # FLA Triton RMSNorm + FusedRMSNormGated
+      use_fla_fused_gated_norm: true   # Same env var (PRIMUS_FLA_NORM); kept explicit for clarity
+      use_fla_short_conv: true         # FLA Triton causal_conv1d for GDN short-conv
+      # ── Dataset source selector ───────────────────────────────────────
+      # use_fla_data + fla_cache_dir together pick the data path:
+      #   use_fla_data=false (or fla_cache_dir empty): vanilla Megatron
+      #     GPTDataset reading train_data_path (.bin/.idx).  Megatron's
+      #     native sampler order — fine for standalone runs, NOT bit-
+      #     comparable to FLA's loss curve.
+      #   use_fla_data=true  AND fla_cache_dir=<path>: replace
+      #     GPTDataset with tools/fla_order_dataset.FLAOrderGPTDataset
+      #     which emits tokens in the exact same order as FLA's HF
+      #     DistributedSampler (eliminates the data-ordering drift).
+      # The launcher script sets PRIMUS_FLA_CACHE_DIR; if you launch
+      # without it, uncomment fla_cache_dir below to make this YAML
+      # self-contained.
+      use_fla_data: true               # PRIMUS_FLA_DATA
+      # fla_cache_dir: /home/<user>/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-100BT/train
+      fused_ce_mode: 1                 # 1=chunked FLA FusedLinearCrossEntropyLoss
+      fused_ce_chunks: 32              # Chunk count for FLA fused CE (32 = launcher default)
+
+      # ─────────────────────────────────────────────────────────────────────
+      # SwiGLU path — disable Megatron's bias-fused SwiGLU so the MLP
+      # actually routes through FLA's Triton SwiGLU (the PRIMUS_FLA_SWIGLU=1
+      # env var's intent).
+      #
+      # `primus/configs/models/megatron/language_model.yaml` sets
+      # `bias_swiglu_fusion: true`, which Megatron translates to
+      # `bias_activation_fusion=true` (see arguments.py L1610-1613) and
+      # forces the MLP forward through `bias_swiglu_impl` (mlp.py L308).
+      # That code path NEVER checks `self._use_fla_swiglu` — meaning
+      # PRIMUS_FLA_SWIGLU=1 has been silently inert ever since.  The
+      # backward for `bias_swiglu_impl` materialises a ~4 GB grad
+      # temporary at b=64, s=2048, intermediate=8192 — exactly the
+      # iter-1 OOM we hit (HIPBLAS_STATUS_ALLOC_FAILED → 4.00 GiB request).
+      #
+      # Setting this to false:
+      #   • routes MLP through the `_use_fla_swiglu` branch (mlp.py L323)
+      #   • uses FLA's `swiglu` Triton kernel (the kernel we actually want
+      #     for parity — same one FLA's reference run uses)
+      #   • saves ~4 GB of backward-pass memory
+      # All bit-perfect parity guarantees are preserved (this is in fact
+      # what we were claiming to do all along — see docs/hybrid_models/GDN_FLA_PARITY.md §B
+      # patch 03 "mlp-fla-swiglu").
+      # ─────────────────────────────────────────────────────────────────────
+      bias_swiglu_fusion: false
+
+      # NOTE: `recompute_granularity` is intentionally NOT set here. It's a
+      # no-op for pure GDN — MambaBlock.forward()/HybridStack.forward() are
+      # plain `for layer in self.layers:` loops with zero recompute logic
+      # (verified in third_party/Megatron-LM/megatron/core/ssm/mamba_block.py
+      # and primus/backends/megatron/core/models/hybrid/hybrid_block.py).
+      # The actual fix for the iter-1 OOM is PRIMUS_FUSED_CE_CHUNKS=32 set
+      # by the launcher — see launch_gdn_pure_1B_100B.sh and
+      # third_party/Megatron-LM/megatron/core/models/mamba/mamba_model.py.
+
+      # Data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # ── Checkpoints ─
+      finetune: false
+      auto_continue_train: true
+      load: null
+      # Checkpoint save path
+      save: ./output/zebra_llama_1B_gdn_pure_100B-pretrain
+      # FLA's 100B launcher passes `save=99999` (HF `--save_steps 99999`),
+      # meaning over 95368 iters HF never writes an intermediate ckpt — only
+      # the final one at end-of-training (then `--save_total_limit 3` keeps
+      # the last 3 across all runs in the dir).
+      #
+      # For Primus we cannot match `99999` literally because
+      # `auto_continue_train: true` needs SOME periodic checkpoint to be able
+      # to resume mid-flight from a crash on a 48 h run. Set save_interval to
+      # 10000 (~10 B tokens, ~10 ckpts over the run) — closer in spirit to
+      # FLA's "almost never save" and ~2× lighter than the previous 5000
+      # setting. `disable_last_saving: false` still guarantees the very last
+      # iter is saved exactly like FLA does at end-of-training.
+      save_interval: 10000
+      disable_last_saving: false
+      ckpt_format: fsdp_dtensor   # Megatron-FSDP requires this when use_megatron_fsdp=true
+
+      # Turbo (kept off for the comparison run; can flip on later)
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_exp7-fsdp-overlap.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_exp7-fsdp-overlap.yaml
new file mode 100644
index 000000000..09d67ba9d
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_exp7-fsdp-overlap.yaml
@@ -0,0 +1,511 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_gdn_pure_exp7-fsdp-overlap}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Pure GDN 1B — 100B tokens on FineWeb-Edu (matched to FLA's
+# `setup_and_train_gdn_pure_1B_100B.sh` so the loss curves can be compared
+# iter-for-iter once we drive Primus from FLA's token order via PRIMUS_FLA_DATA).
+#
+# FLA's run:
+#   model:       configs/gated_deltanet_1B_pure_100B.json (same arch as the
+#                10B config we already replicate in zebra_llama_1B_gdn_pure.yaml)
+#   lr:          3e-4   (vs the 10B run's 2e-4 — bigger LR for the longer schedule)
+#   scheduler:   cosine_with_min_lr  (min_lr ≈ 0.1 × peak = 3e-5)
+#   warmup:      2000 iters         (vs 200 on the 10B run)
+#   batch:       64 per GPU         (vs 16 on the 10B run)
+#   update:      1                  (no grad-accum)
+#   gpus:        8
+#     → global_batch_size = 64 × 8 = 512
+#     → tokens/iter      = 512 × 2048 = 1,048,576
+#   steps:       95368              (95368 × 1,048,576 ≈ 100B tokens)
+#   data:        HuggingFaceFW/fineweb-edu, sample-100BT
+#                cache: .../data/HuggingFaceFW/fineweb-edu/sample-100BT/train
+#                (~364 GB on disk — pre-tokenized by FLA)
+#
+# Wall-time estimate on 8×MI300X with the full FLA-parity stack (~1.8 s/iter
+# based on the GDN-hybrid 300M @ 2.2 s/iter rescaled for the larger model):
+#   95368 × ~1.8 s ≈ 48 h (~2 days)
+# ─────────────────────────────────────────────────────────────────────────────
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_1B_gdn_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_1B_GDN_Pure_100B"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      # FLA's `train.sh` defaults `--dataloader_num_workers 32` and the 100B
+      # run does NOT override it. At b=64 mbs, ~131k tokens/GPU/iter, 32
+      # workers (vs Primus's previous 8 carried over from the 300M run) keeps
+      # the DataLoader queue full so the GPU never starves between iters.
+      # PyTorch's default `prefetch_factor=2` is already what FLA uses.
+      num_workers: 32
+      create_attention_mask_in_dataloader: false
+
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # Perf: Megatron defaults to inserting dist.barrier() before every L1 timer
+      # measurement (~5-10/iter). Costs throughput, no correctness impact.
+      barrier_with_L1_time: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # SPEED RECOVERY (lifted from 300M YAML which matched FLA at +0.27%):
+      #
+      # The profile of iters 15-16 showed 487 ms/iter spent in `aten::item`
+      # (CPU↔GPU sync points). The 300M run avoided this by disabling NaN
+      # checking and bumping log_interval to 100. Both are loss-identical
+      # — only affect logging and sync, not training math.
+      #
+      # check_for_nan_in_loss_and_grad: false
+      #   • Megatron default = true, which fires `loss.item()` and an
+      #     `isnan`/`isinf` validate_result on the loss EVERY iter
+      #     (pretrain_mamba.py:178). Each is a CPU-GPU sync. Saves ~50 ms.
+      #   • If loss ever spikes to NaN you'll still see it in the log
+      #     output, just won't auto-rerun the step.
+      #
+      # log_interval: 32
+      #   • Default is 1, meaning every step we run timers, compute
+      #     throughput, `loss.item()`, and print to stdout/log files —
+      #     ~200 ms of sync+formatting per iter. Logging every 32 iters
+      #     EXACTLY matches FLA's `--logging_steps 32` (legacy/training/
+      #     train.sh:55) so loss curves line up tick-for-tick with the
+      #     FLA reference log when overlaid for parity verification.
+      #     At ~7.5 s/iter that's a log line every 4 min — plenty.
+      #
+      # tensorboard_log_interval: 32
+      #   • Same logic for TB/wandb logging. Default is 1 = sync every
+      #     iter. ~100 ms recovered.
+      #
+      # Combined expected speedup: ~300-400 ms/iter (7.6 s → ~7.2 s),
+      # all loss-identical. This is what the 300M YAML had.
+      # ─────────────────────────────────────────────────────────────────────
+      check_for_nan_in_loss_and_grad: false
+      # FLA's 100B launcher uses `logging=10` → `--logging_steps 10`. Matching
+      # exactly so the Primus loss log lines up tick-for-tick with FLA's
+      # reference `train_gdn_pure_1B_100B.log`.
+      log_interval: 10
+      tensorboard_log_interval: 10
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA-parity numerics (lifted verbatim from the validated 300M parity
+      # YAML — these are the *exact* settings that produced the iter-1
+      # bit-perfect match and the < 0.5% late-training residual documented
+      # in docs/hybrid_models/GDN_FLA_PARITY.md).  Skipping any of them re-introduces a known
+      # source of drift:
+      #
+      #   layernorm_epsilon  — Megatron's TransformerConfig defaults to 1e-5;
+      #                        FLA uses 1e-6.  ~1% per-layer divergence.
+      #   hidden_dropout     — Megatron's TransformerConfig defaults these to
+      #   attention_dropout    0.1 each (transformer_config.py L152).  Even
+      #                        though the model YAML asks for 0, the EXP-level
+      #                        defaults leak through and we'd train with 10%
+      #                        dropout while FLA trains with 0.  CRITICAL.
+      #   no_persist_layer_norm — disables Apex's persistent buffer kernel
+      #                        which has slightly different rounding.
+      # ─────────────────────────────────────────────────────────────────────
+      layernorm_epsilon: 1.0e-6
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+      no_persist_layer_norm: true
+
+      # ── Training schedule ─ matches FLA EXACTLY (bit-perfect every knob):
+      #   batch=64, grad_accum=1, gpus=8, seq=2048
+      #   global_batch_size = 64 × 8 = 512
+      #   tokens/iter       = 512 × 2048 = 1,048,576
+      #   total tokens      = 95368 × 1,048,576 ≈ 100 B
+      train_iters: 50    # PERF EXP: 50 iters to measure steady-state
+      micro_batch_size: 64
+      global_batch_size: 512
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # ── Optimizer ─ exact bit-match to FLA's 100B CLI:
+      #   --learning_rate 3.0e-4
+      #   --lr_scheduler_type cosine_with_min_lr
+      #   --warmup_steps 2000
+      #   --optim adamw_torch_fused
+      #   --adam_beta1=0.9 --adam_beta2=0.95
+      #   --weight_decay 0.01
+      #   --max_grad_norm 1.0
+      #   --seed 42 --bf16
+      #
+      # CRITICAL: although FLA's train.sh does NOT pass --lr_scheduler_kwargs,
+      # FLA's run.py OVERRIDES it programmatically — see
+      # flash-linear-attention/legacy/training/run.py:96-97 :
+      #     if args.lr_scheduler_type == 'cosine_with_min_lr':
+      #         args.lr_scheduler_kwargs = {'min_lr_rate': 0.1}
+      # So FLA's lr decays from 3e-4 to 0.1*3e-4 = 3e-5 over 95368 iters,
+      # NOT to 0 as the HF default would imply.  Megatron's cosine is
+      #     lr = min_lr + 0.5*(1+cos(π*p))*(max_lr - min_lr)
+      # HF's cosine_with_min_lr is the algebraically identical form
+      #     lr = peak * (factor*(1-min_lr_rate) + min_lr_rate)
+      # so setting `min_lr: 3.0e-5` here matches FLA's lr at every step,
+      # bit-exact.  (Setting `min_lr: 0.0` — what we had earlier — would
+      # cause Primus's lr to undershoot FLA's in late training; the loss
+      # curves would diverge after ~iter 50000.)
+      # (Warmup is `init_lr + (max_lr-init_lr)*step/warmup` in Megatron and
+      # `peak*step/warmup` in HF — identical when init_lr=0, which is the
+      # Megatron default.)
+      clip_grad: 1.0
+      lr: 3.0e-4
+      min_lr: 3.0e-5
+      lr_warmup_iters: 2000
+      lr_decay_iters: 95368
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Match FLA's --seed 42 exactly.  Primus's default is 1234; using FLA's
+      # seed minimises init-RNG drift between the two runs (still not bit-
+      # identical because Megatron and HF transformers walk parameters in
+      # different orders, but it eliminates the 12 % of variance that comes
+      # from the seed itself).
+      seed: 42
+
+      # ─────────────────────────────────────────────────────────────────────
+      # NO-TE spec — the validated 300M parity run uses this variant because
+      # FLA's reference is built on native PyTorch nn.Linear / RMSNorm.  TE's
+      # ColumnParallelLinear / TENorm wrappers introduce small numerical
+      # differences (cast ordering, persistent buffers) that drift the loss
+      # by ~0.1% per layer.  Using the no-TE spec aligns Megatron's layer
+      # numerics with FLA's.  (See docs/hybrid_models/GDN_FLA_PARITY.md §A.)
+      # ─────────────────────────────────────────────────────────────────────
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+
+      # Tokenizer (cached locally — same path FLA's run uses)
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # ── Parallelism ─
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+
+      # ─────────────────────────────────────────────────────────────────────
+      # DDP / optimizer sharding — Megatron-FSDP (ZeRO-2 equivalent)
+      # ─────────────────────────────────────────────────────────────────────
+      #
+      # WHY we switched from ZeRO-1 to Megatron-FSDP (this run only):
+      #
+      # Profile of iter 15-16 (captured 2026-05-25) showed:
+      #
+      #   GPU busy:                   ~43 %  (TFLOP/s = 261, peak 1300)
+      #   aten::empty                 1500 calls × 2.3 ms = 3390 ms / iter
+      #   aten::empty_like            1206 calls × 2.6 ms = 3114 ms / iter
+      #   aten::empty_strided         1100 calls × 2.3 ms = 2575 ms / iter
+      #   aten::mm dispatch            576 calls × 9.5 ms = 5485 ms / iter
+      #
+      # Normal aten::empty is *microseconds*. Ours is **1000× slower**
+      # because the PyTorch HIP allocator is constantly hitting hipMalloc/
+      # hipFree at 81 % VRAM (156 GB / 192 GB). For comparison the 300M
+      # run runs at 30 % VRAM, the allocator cache stays warm, every
+      # aten::empty is sub-µs, and it achieves 643 TFLOP/s (49 % of peak).
+      #
+      # Megatron's own FSDP docs explicitly call out this exact failure
+      # mode (third_party/Megatron-LM/docs/user-guide/features/
+      # custom_fsdp.md §5 line 187):
+      #
+      #   "FSDP can lead to crashes of the PyTorch memory allocator cache,
+      #    and a large number of cudaMalloc and cudaFree calls. This
+      #    problem is challenging and can only be mitigated by avoiding
+      #    frequent hits on the GPU memory limit."
+      #
+      # FLA's reference run uses DeepSpeed ZeRO-Stage-2 (ds_config.json
+      # "stage": 2, contiguous_gradients=true) which shards BOTH optimizer
+      # state AND gradients. Our prior config used Megatron's
+      # `use_distributed_optimizer=true` which is ZeRO-Stage-1 (only opt
+      # state sharded — full grads stay on every rank, ≈ 4.8 GB / rank).
+      # Megatron-FSDP with `data_parallel_sharding_strategy: optim_grads`
+      # is the bit-for-bit equivalent of FLA's setup: shards both opt-state
+      # AND grads, leaves params un-sharded. Memory savings ≈ 4.2 GB/rank
+      # (from 156 → ~152 GB, 81 % → 79 %), which along with removing the
+      # per-iter `empty_cache()` (see empty_unused_memory_level below)
+      # should give the allocator enough headroom to stop thrashing.
+      #
+      # Why optim_grads (ZeRO-2) NOT optim_grads_params (ZeRO-3):
+      #   • optim_grads_params shards params too, but Megatron-FSDP only
+      #     auto-wraps `TransformerLayer` instances as FSDP units. Our
+      #     GDN layers are `MambaLayer` (a separate Megatron class), so
+      #     they would NOT be FSDP-wrapped and would silently fall back
+      #     to no-sharding for the GDN half of the model.
+      #   • FLA itself uses ZeRO-2, not ZeRO-3.
+      #   • Loss numerics are identical to FLA only with ZeRO-2 (ZeRO-3
+      #     introduces a tiny float-rounding difference in the per-shard
+      #     param reconstruction).
+      #
+      # Compatibility caveats verified:
+      #   ✓ Megatron-FSDP requires use_distributed_optimizer=true (kept)
+      #   ✓ Requires gradient_accumulation_fusion=false (kept)
+      #   ✓ Auto-handles CUDA_DEVICE_MAX_CONNECTIONS via Primus's
+      #     env_patches.py:set_cuda_device_max_connections() — sets to
+      #     8 when use_megatron_fsdp=true (vs 1 for non-FSDP)
+      #   ✓ GDN layer has reset_parameters() (megatron/core/ssm/
+      #     gated_delta_net.py:227), so meta-device init would work too
+      #     (but we leave init_model_with_meta_device=false for safety)
+      #   ✓ Prior segfault (overlap_grad_reduce=true + ZeRO-1) was NOT
+      #     about FSDP — it was a separate Megatron bug in the ZeRO-1
+      #     bucket layout for heterogeneous Mamba params. FSDP uses its
+      #     own (different) buffer code path.
+      #
+      # If FSDP causes any crash, revert these 2 lines to restore ZeRO-1:
+      #   use_megatron_fsdp: false
+      #   data_parallel_sharding_strategy: no_shard
+      # (Rest of config is FSDP-safe AND ZeRO-1-safe.)
+      # ─────────────────────────────────────────────────────────────────────
+      use_distributed_optimizer: true
+      # ── Megatron-FSDP ATTEMPTS (2026-05-25 + 2026-05-26) ─────────────────
+      # Take #1 (2026-05-25): `use_megatron_fsdp: true` +
+      #   `data_parallel_sharding_strategy: optim_grads` crashed at iter-1
+      #   post-step NaN-check all_reduce with `Failed to CUDA calloc
+      #   268435456 bytes` (256 MiB).
+      #
+      #   Root cause (diagnosed 2026-05-26): RCCL's lazy comm-init allocates
+      #   BUFFSIZE × #channels per NCCL communicator. Defaults =
+      #   4 MiB × 64 channels = 256 MiB per comm. Megatron-FSDP creates
+      #   extra comm groups (HSDP outer + DP inner) on top of the one
+      #   ZeRO-1 needs, and at 99% VRAM there is no contiguous 256 MiB hole.
+      #
+      # Take #2 (2026-05-26): added NCCL channel clamp via launcher env vars
+      #   (NCCL_MIN_NCHANNELS=1 NCCL_MAX_NCHANNELS=4 NCCL_NCHANNELS_PER_PEER=1
+      #    + ckpt_format: fsdp_dtensor since Megatron asserts it).
+      #   Result: 500 iters complete, exit 0, ZERO crashes.
+      #
+      #   Steady-state iter time: 2.84 s/iter (vs ZeRO-1's 2.75 s/iter).
+      #   FSDP wins on non-flush iters (2.67 s vs 2.75 s) but loses on
+      #   empty_cache_interval=32 flushes (3.18 s vs ~2.98 s) because FSDP
+      #   has to re-allgather sharded params after the cache drop, while
+      #   ZeRO-1 keeps full param replicas resident.
+      #
+      #   NET (Take #2): FSDP 2.84 s/iter vs ZeRO-1 2.75 s/iter (90 ms slower).
+      #
+      # Take #3 (2026-05-26): Controlled 50-iter experiment EXP5 with
+      #   `empty_cache_interval` raised from 32 to 128 (one flush per 128
+      #   iters instead of one per 32):
+      #     - ZeRO-1 baseline (EXP0):  2768 ms steady-state
+      #     - FSDP + flush@128 (EXP5): 2684 ms steady-state  (-83 ms, -3.0%)
+      #     - 100B run savings: 95k iters × 83 ms = 2.2 hours
+      #     - Loss bit-identical to ZeRO-1 at iter 50: 11.7538
+      #
+      #   Why this works: between flushes, FSDP's per-iter allgather work is
+      #   pipelined into the backward pass (Megatron-FSDP enables this
+      #   internally even though `overlap_param_gather` is off). On flush
+      #   iters FSDP still pays the re-allgather cost, but at 1-in-128 the
+      #   amortized penalty is ~6 ms/iter vs the ~85 ms/iter saved.
+      #
+      # EXP6 (ZeRO-3 / optim_grads_params) FAILED at iter 1 with:
+      #   "RuntimeError: Pointer argument cannot be accessed from Triton
+      #    (cpu tensor?)"
+      # Because ZeRO-3 shards parameters across ranks, and the FLA Triton
+      # kernel cannot dereference a sharded DTensor. ZeRO-2 (optim_grads)
+      # is the maximum sharding compatible with our Triton GDN.
+      #
+      # FSDP working config + full experiment results:
+      #   experiments/results/SYNTHESIS.md
+      #   experiments/run_perf_exp.sh
+      #   examples/megatron/configs/MI300X/zebra_llama_1B_gdn_pure_exp5-fsdp-rareflush.yaml
+      # ──────────────────────────────────────────────────────────────────
+      use_megatron_fsdp: true
+      data_parallel_sharding_strategy: optim_grads
+      # EXP7: FSDP path's overlap_grad_reduce is DIFFERENT code than ZeRO-1's
+      # (which segfaulted in EXP4). FSDP uses ReduceScatter into its own buffer.
+      # If this works, we save ~152 ms/iter (NCCL all-reduce → overlapped with backward).
+      overlap_grad_reduce: true        # EXP7: enable ReduceScatter overlap (FSDP path, not ZeRO-1 path)
+      overlap_param_gather: true       # EXP7: enable AllGather overlap (free with overlap_grad_reduce)
+      gradient_accumulation_fusion: false
+      ddp_average_in_collective: true
+      use_torch_fsdp2: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # PyTorch profiler — already captured iter 15-17 trace which showed
+      # 487 ms/iter in `aten::item` (the 300M-parity log_interval/NaN-check
+      # fixes above target exactly this). Disabling for this run.
+      #
+      # To re-enable: set `profile: true`, `use_pytorch_profiler: true`
+      # and pick `profile_step_start`/`profile_step_end`.
+      # ─────────────────────────────────────────────────────────────────────
+      profile: false
+      use_pytorch_profiler: false
+      tensorboard_dir: /tmp/primus_perf_exp7-fsdp-overlap_tb
+
+      # ─────────────────────────────────────────────────────────────────────
+      # MEMORY: empty_unused_memory_level + empty_cache_interval
+      # ─────────────────────────────────────────────────────────────────────
+      #
+      # empty_unused_memory_level=1 → calls `torch.cuda.empty_cache()` at
+      # training.py:1759, IMMEDIATELY before optimizer.step() (which
+      # internally calls get_grad_norm_fp32() → torch.distributed.all_reduce
+      # in clip_grads.py:130).
+      #
+      # That all_reduce is where NCCL crashes if we disable empty_cache
+      # entirely:
+      #   "Failed to CUDA calloc 4194304 bytes" inside ProcessGroupNCCL.
+      # NCCL's calloc bypasses PyTorch's allocator and goes straight to
+      # hipMalloc — at 81% VRAM there is no contiguous 4 MiB block free
+      # unless PyTorch first returns its cached pool to the driver.
+      #
+      # 2026-05-26 — PROFILE + FIX:
+      # PyTorch profiler attribution on iter 21 of the live 100B run
+      # measured 6,937 ms of the 7,650 ms iter (91%) in hipMalloc + hipFree,
+      # because the per-iter empty_cache() was thrashing ~250 cached
+      # blocks per iter against the ROCm driver (~17 ms per hipMalloc).
+      #
+      # FIX: keep `empty_unused_memory_level: 1` (Megatron's gate) but
+      # only fire empty_cache() every N iters via the Primus patch
+      # primus/backends/megatron/patches/empty_cache_interval_patches.py.
+      # The patch reads `empty_cache_interval` from this YAML (preferred),
+      # then falls back to the PRIMUS_EMPTY_CACHE_INTERVAL env var, then
+      # to the default 1 (passthrough).
+      #
+      # Iter 0 always fires empty_cache() so NCCL's first workspace alloc
+      # happens against a clean cache.  Iters 1..N-1 skip it, iter N fires
+      # again, etc.  NCCL workspace persists across the gap because once
+      # allocated it is not freed by PyTorch's allocator (NCCL owns the
+      # block).
+      #
+      # Measured impact (50-iter diag at empty_cache_interval=32,
+      # 2026-05-26):
+      #   before fix:  7.65 s/iter  ( 262 TFLOP/s/GPU,  20% of peak)
+      #   after fix:   2.79 s/iter  ( 727 TFLOP/s/GPU,  56% of peak)
+      #   speedup:     2.74×
+      #   loss curve:  iter-1 bit-identical to baseline (no math impact)
+      #   100B wall:   was 8.4 days → now 3.08 days  (~5.3 days saved)
+      #
+      # Knob choice:
+      #   - 1   = original per-iter behaviour (safest, slowest)
+      #   - 32  = our default (best tested point; ~2.8× speedup)
+      #   - 64+ = even fewer flushes; tested briefly, slightly faster
+      #           between flushes but bigger spike on the trigger iter
+      #   - 0   = NEVER flush (CRASHES at iter-1 NCCL alloc; do not use)
+      # ─────────────────────────────────────────────────────────────────────
+      empty_unused_memory_level: 1
+      empty_cache_interval: 128   # 2026-05-26 EXP5: 32→128 saves 83 ms/iter (-3% wall, -2.2 h on 100B run)
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA runtime knobs — declarative YAML surface for the PRIMUS_FLA_* /
+      # PRIMUS_FUSED_CE* env vars (consumed by the
+      # primus.backends.megatron.patches.fla_runtime_patches patch which
+      # re-exports them as env vars at phase="build_args").  These
+      # exactly mirror the values the legacy launcher script sets via
+      # `export PRIMUS_FLA_*` — keeping them here makes the YAML the
+      # single source of truth.  Env vars set on the launcher still win
+      # over the YAML (backward compat).
+      # ─────────────────────────────────────────────────────────────────────
+      use_fla_fused_swiglu: true       # FLA Triton SwiGLU (matches fla.modules.swiglu)
+      use_fla_fused_rmsnorm: true      # FLA Triton RMSNorm + FusedRMSNormGated
+      use_fla_fused_gated_norm: true   # Same env var (PRIMUS_FLA_NORM); kept explicit for clarity
+      use_fla_short_conv: true         # FLA Triton causal_conv1d for GDN short-conv
+      # ── Dataset source selector ───────────────────────────────────────
+      # use_fla_data + fla_cache_dir together pick the data path:
+      #   use_fla_data=false (or fla_cache_dir empty): vanilla Megatron
+      #     GPTDataset reading train_data_path (.bin/.idx).  This is
+      #     Megatron's native sampler order — fine for standalone runs
+      #     but NOT bit-comparable to FLA's loss curve.
+      #   use_fla_data=true  AND fla_cache_dir=<path>: replace
+      #     GPTDataset with tools/fla_order_dataset.FLAOrderGPTDataset
+      #     which emits tokens in the exact same order as FLA's HF
+      #     DistributedSampler (eliminates the data-ordering drift that
+      #     caused the +2.4 nat warm-up spike we debugged for GDN/KDA).
+      # The launcher script sets PRIMUS_FLA_CACHE_DIR; if you launch
+      # without it, uncomment fla_cache_dir below to make this YAML
+      # self-contained.
+      use_fla_data: true               # PRIMUS_FLA_DATA
+      # fla_cache_dir: /home/<user>/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-100BT/train
+      fused_ce_mode: 1                 # 1=chunked FLA FusedLinearCrossEntropyLoss
+      fused_ce_chunks: 32              # Chunk count for FLA fused CE (32 = launcher default)
+
+      # ─────────────────────────────────────────────────────────────────────
+      # SwiGLU path — disable Megatron's bias-fused SwiGLU so the MLP
+      # actually routes through FLA's Triton SwiGLU (the PRIMUS_FLA_SWIGLU=1
+      # env var's intent).
+      #
+      # `primus/configs/models/megatron/language_model.yaml` sets
+      # `bias_swiglu_fusion: true`, which Megatron translates to
+      # `bias_activation_fusion=true` (see arguments.py L1610-1613) and
+      # forces the MLP forward through `bias_swiglu_impl` (mlp.py L308).
+      # That code path NEVER checks `self._use_fla_swiglu` — meaning
+      # PRIMUS_FLA_SWIGLU=1 has been silently inert ever since.  The
+      # backward for `bias_swiglu_impl` materialises a ~4 GB grad
+      # temporary at b=64, s=2048, intermediate=8192 — exactly the
+      # iter-1 OOM we hit (HIPBLAS_STATUS_ALLOC_FAILED → 4.00 GiB request).
+      #
+      # Setting this to false:
+      #   • routes MLP through the `_use_fla_swiglu` branch (mlp.py L323)
+      #   • uses FLA's `swiglu` Triton kernel (the kernel we actually want
+      #     for parity — same one FLA's reference run uses)
+      #   • saves ~4 GB of backward-pass memory
+      # All bit-perfect parity guarantees are preserved (this is in fact
+      # what we were claiming to do all along — see docs/hybrid_models/GDN_FLA_PARITY.md §B
+      # patch 03 "mlp-fla-swiglu").
+      # ─────────────────────────────────────────────────────────────────────
+      bias_swiglu_fusion: false
+
+      # NOTE: `recompute_granularity` is intentionally NOT set here. It's a
+      # no-op for pure GDN — MambaBlock.forward()/HybridStack.forward() are
+      # plain `for layer in self.layers:` loops with zero recompute logic
+      # (verified in third_party/Megatron-LM/megatron/core/ssm/mamba_block.py
+      # and primus/backends/megatron/core/models/hybrid/hybrid_block.py).
+      # The actual fix for the iter-1 OOM is PRIMUS_FUSED_CE_CHUNKS=32 set
+      # by the launcher — see launch_gdn_pure_1B_100B.sh and
+      # third_party/Megatron-LM/megatron/core/models/mamba/mamba_model.py.
+
+      # ── Data ─ FLA-aligned via PRIMUS_FLA_DATA=1 at runtime.
+      #
+      # `train_data_path` below is just a Megatron config-parser placeholder;
+      # FLAOrderGPTDataset overrides it when PRIMUS_FLA_DATA=1 + a valid
+      # PRIMUS_FLA_CACHE_DIR are exported.  Point the launcher's
+      # PRIMUS_FLA_CACHE_DIR at the sample-100BT cache (FLA has it on disk
+      # at /home/<user>/flash-linear-attention/legacy/training/
+      # data/HuggingFaceFW/fineweb-edu/sample-100BT/train).
+      #
+      # The 10BT .bin/.idx is reused as a placeholder so Megatron's index
+      # builder doesn't choke on a missing file — its actual indices and
+      # tokens are never read when FLAOrderGPTDataset is active.
+      mock_data: false
+      train_data_path: >
+        data/fla_aligned/fla_fineweb_edu_10BT_text_sentence
+      valid_data_path: null
+      test_data_path: null
+
+      # ── Checkpoints ─
+      finetune: false
+      auto_continue_train: true
+      load: null
+      # FLA's `path=` for this run is /home/<user>/checkpoints/
+      # gdn_pure_1B_100B. We save to a sibling dir prefixed with `primus_` so
+      # the two runs' checkpoints can coexist on disk and downstream eval
+      # tools (tools/hybrid/eval_gdn_lm_eval.py) can A/B them side-by-side.
+      save: /tmp/primus_perf_exp_ckpt_unused
+      # FLA's 100B launcher passes `save=99999` (HF `--save_steps 99999`),
+      # meaning over 95368 iters HF never writes an intermediate ckpt — only
+      # the final one at end-of-training (then `--save_total_limit 3` keeps
+      # the last 3 across all runs in the dir).
+      #
+      # For Primus we cannot match `99999` literally because
+      # `auto_continue_train: true` needs SOME periodic checkpoint to be able
+      # to resume mid-flight from a crash on a 48 h run. Set save_interval to
+      # 10000 (~10 B tokens, ~10 ckpts over the run) — closer in spirit to
+      # FLA's "almost never save" and ~2× lighter than the previous 5000
+      # setting. `disable_last_saving: false` still guarantees the very last
+      # iter is saved exactly like FLA does at end-of-training.
+      save_interval: 100000   # PERF EXP: never save
+      disable_last_saving: false
+      ckpt_format: fsdp_dtensor   # Megatron-FSDP requires this when use_megatron_fsdp=true
+
+      # Turbo (kept off for the comparison run; can flip on later)
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml
new file mode 100644
index 000000000..c228cf00c
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_kda-pretrain.yaml
@@ -0,0 +1,84 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_kda-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: zebra_llama_1B.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_Zebra_Llama_1B_KDA_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 0
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      train_iters: 50
+      micro_batch_size: 4
+      global_batch_size: 32 # micro_batch_size * num_gpus
+
+      seq_length: 8192
+      max_position_embeddings: 8192
+      original_max_position_embeddings: 8192
+
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 38147
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+
+      # Use KDA (Kimi Delta Attention) hybrid spec
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec']
+      use_fla_triton_kda: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # parallel
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: true
+
+      # data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: true
+      load: null
+      save: ./output/zebra_llama_1B_kda-pretrain
+      save_interval: 1000
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B_kda_pure-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B_kda_pure-pretrain.yaml
new file mode 100644
index 000000000..2750f2e6f
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_1B_kda_pure-pretrain.yaml
@@ -0,0 +1,94 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B_kda_pure-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_1B_kda_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_1B_KDA_Pure_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 0
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # Training schedule — matched to FLA (4 GPUs):
+      #   batch=16, update=1, gpus=4, context=2048
+      #   global_batch = 16 * 4 = 64, tokens/step = 64 * 2048 = 131,072
+      #   total tokens ~= 76294 * 131072 ≈ 10B
+      train_iters: 50
+      micro_batch_size: 16
+      global_batch_size: 128
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — matched to FLA:
+      #   lr=2e-4, cosine (min_lr_rate=0.1 → min_lr=2e-5)
+      #   AdamW, beta1=0.9, beta2=0.95, weight_decay=0.01
+      #   max_grad_norm=1.0, warmup_steps=200
+      # Both DeepSpeed ZeRO-2 and Megatron clip on ||avg_grad||.
+      # DeepSpeed pre-divides by dp_size before reduce_scatter(SUM),
+      # so its grad norm is on averaged gradients. Match FLA's 1.0 directly.
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 76294
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Use KDA hybrid spec with hybrid_attention_ratio=0.0 (pure KDA)
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec']
+      use_fla_triton_kda: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: true
+
+      # Data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints
+      finetune: false
+      auto_continue_train: true
+      load: null
+      save: ./output/zebra_llama_1B_kda_pure-pretrain
+      save_interval: 2048
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_hybrid-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_hybrid-pretrain.yaml
new file mode 100644
index 000000000..8fa3e3230
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_hybrid-pretrain.yaml
@@ -0,0 +1,171 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_300M_gdn_hybrid-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_300M_gdn_hybrid.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_300M_GDN_Hybrid_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      # FLA parity: HF Trainer prefetches lightly from a memmap'd parquet, so
+      # 2 dataloader workers per rank (= 16 forked subprocesses on 8 GPUs) is
+      # enough to keep the dataloader pipeline full without bloating host RSS.
+      # `num_workers: 8` (= 64 forked subprocesses) was OOM-killing the host
+      # at ~iter 1200 on this 125 GB-RAM box.
+      num_workers: 2
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+      log_interval: 100
+      check_for_nan_in_loss_and_grad: false
+
+      # Perf: disable barrier-with-L1-time (serializes ranks per L1 timer call)
+      barrier_with_L1_time: false
+
+      # Match FLA seed for reproducibility comparison
+      seed: 42
+
+      # norm_epsilon in the model YAML maps to args.norm_epsilon, but
+      # TransformerConfig uses layernorm_epsilon (default 1e-5).
+      # Set explicitly to match FLA's 1e-6.
+      layernorm_epsilon: 1.0e-6
+
+      # FLA does no dropout; force off (Megatron defaults are 0.1)
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # Training schedule — matched to FLA (8 GPUs, FineWeb-Edu sample-10BT):
+      #   FLA: per_device_train_batch_size=128 × 8 GPUs → global=1024
+      #   tokens/step = 1024 × 2048 = 2,097,152
+      #   4768 steps × 2.097M ≈ 10B tokens
+      train_iters: 50
+      micro_batch_size: 8
+      global_batch_size: 64
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — matched to FLA:
+      #   lr=2e-4, cosine (min_lr_rate=0.1 → min_lr=2e-5)
+      #   AdamW, beta1=0.9, beta2=0.95, weight_decay=0.01
+      #   max_grad_norm=1.0, warmup_steps=200
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 4768
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA runtime knobs — declarative YAML surface for the
+      # PRIMUS_FLA_*/PRIMUS_FUSED_CE* env vars.  Consumed by
+      # primus.backends.megatron.patches.fla_runtime_patches (phase=
+      # "build_args") which re-exports them as env vars so the existing
+      # consumers in fla_flash_attention.py / gated_delta_net.py /
+      # hybrid_block.py / mamba_model.py / mlp.py / pretrain_mamba.py
+      # see identical values.  Env vars set on the launcher still win
+      # over the YAML (backward compat).
+      # ─────────────────────────────────────────────────────────────────────
+      use_fla_fused_swiglu: true       # FLA Triton SwiGLU
+      use_fla_fused_rmsnorm: true      # FLA Triton RMSNorm + FusedRMSNormGated
+      use_fla_fused_gated_norm: true   # Same env var (PRIMUS_FLA_NORM)
+      use_fla_short_conv: true         # FLA Triton causal_conv1d for GDN short-conv
+      # ── Dataset source selector ───────────────────────────────────────
+      # use_fla_data + fla_cache_dir together pick the data path:
+      #   use_fla_data=false (or fla_cache_dir empty): vanilla Megatron
+      #     GPTDataset reading train_data_path (.bin/.idx).
+      #   use_fla_data=true  AND fla_cache_dir=<path>: replace
+      #     GPTDataset with tools/fla_order_dataset.FLAOrderGPTDataset
+      #     for bit-identical token order to FLA's HF DistributedSampler.
+      # The launcher script sets PRIMUS_FLA_CACHE_DIR; if you launch
+      # without it, uncomment fla_cache_dir below.
+      use_fla_data: true               # PRIMUS_FLA_DATA
+      # fla_cache_dir: /home/<user>/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train
+      fused_ce_mode: 1                 # 1=chunked FLA FusedLinearCrossEntropyLoss
+      fla_mla_attn: "1"                # FLA flash-attn for the MLA blocks
+
+      # Hybrid GDN+MLA stack (no-TE spec to match FLA layers; same spec as pure GDN)
+      # With hybrid_attention_ratio=0.25, the allocator places MLA at mixer blocks
+      # [0, 4, 8] — identical to FLA's `attn.layers: [0, 4, 8]`.
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+      no_persist_layer_norm: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism.
+      #
+      # Distributed optimizer (ZeRO-1) shards optimizer state across DP ranks
+      # and amortizes the per-iter Adam step over the DP group.  Previously
+      # disabled because pairing it with `overlap_param_gather: true` skewed
+      # per-rank GPU memory by ~30 GB (gathered param buffer aggregated on
+      # ranks 0-1 → 189 GB used, ranks 2-7 → 155 GB), OOMing FLA's LCE
+      # 7.83 GB chunk.  With the new FLA-fusion stack (PRIMUS_FLA_NORM=1 →
+      # FusedRMSNormGated, in-kernel gate fusion) every rank now sits at
+      # ~171 GB / 192 GB, leaving ~21 GB headroom — enough for the
+      # distributed optimizer's gather buffer as long as `overlap_param_gather`
+      # stays off.
+      #
+      # Expected gain: ~25 ms/iter (closes the bulk of the remaining gap vs
+      # FLA's DeepSpeed-ZeRO-2 optimizer step amortization).
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      use_distributed_optimizer: true   # ZeRO-1 — shards optimizer state
+      overlap_grad_reduce: false
+      overlap_param_gather: false       # leave OFF — the skew that caused OOM
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      ddp_average_in_collective: true
+
+      # Data — FLA-aligned FineWeb-Edu sample-10BT (same indexed binary the
+      # pure GDN/KDA 300M runs consume; uses FLAOrderGPTDataset for identical
+      # token ordering to FLA)
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints — train from scratch by default. To start from an
+      # FLA-initialized weights snapshot for iter-1 parity, set
+      #   finetune: true
+      #   load: /path/to/fla_init_hybrid_ckpt
+      #   no_load_optim: true
+      #   no_load_rng: true
+      finetune: false
+      auto_continue_train: false
+      load: null
+      save: ./output/zebra_llama_300M_gdn_hybrid-pretrain
+      # FLA parity: FLA's setup_and_train_gdn_hybrid_300M.sh uses `save=99999`,
+      # i.e. it only writes a checkpoint at the very end of the 4768-step run.
+      # The mid-training save we previously did at iter 1024 spiked host RSS
+      # (each rank materialises its full state-dict in CPU memory during the
+      # save) and the OS OOM-killer took us out around iter 1200 on this
+      # 125 GB-RAM host.  Match FLA: skip mid-run saves, keep the end save.
+      save_interval: 99999
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo (kept off to match FLA's vanilla PyTorch path)
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml
new file mode 100644
index 000000000..178cf0381
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml
@@ -0,0 +1,123 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_300M_gdn_pure-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_300M_gdn_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_300M_GDN_Pure_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 8
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+      log_interval: 100
+      check_for_nan_in_loss_and_grad: false
+
+      # Perf: Megatron defaults to inserting dist.barrier() before every L1
+      # timer measurement (~5-10/iter) to make per-stage timings comparable.
+      # This serializes ranks. Disable for production training.
+      barrier_with_L1_time: false
+
+      # Match FLA seed for reproducibility comparison
+      seed: 42
+
+      # Fix: norm_epsilon in model YAML maps to args.norm_epsilon, but
+      # TransformerConfig uses layernorm_epsilon (default 1e-5).
+      # Must explicitly set layernorm_epsilon to match FLA's 1e-6.
+      layernorm_epsilon: 1.0e-6
+
+      # CRITICAL: Megatron's TransformerConfig defaults hidden_dropout=0.1 and
+      # attention_dropout=0.1 (transformer_config.py L152). Even though
+      # mamba_base.yaml sets these to 0.0, the YAML inheritance is being
+      # overridden by language_model.yaml (which sets 0.1) and the override is
+      # not propagating to args. This caused embeddings to be dropout-perturbed
+      # at every iteration in prior runs (verified empirically: same token
+      # produced DIFFERENT embedding vectors with 1/(1-0.1)=1.111 scaling).
+      # FLA does NOT apply any dropout. Force these to 0 here to match.
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # Training schedule — matched to FLA (8 GPUs):
+      #   FLA: per_device_train_batch_size=128, 8 GPUs → global=1024
+      #   tokens/step = 1024 * 2048 = 2,097,152
+      #   total tokens ~= 4768 * 2,097,152 ≈ 10B
+      train_iters: 50
+      micro_batch_size: 8
+      global_batch_size: 64
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — matched to FLA:
+      #   lr=2e-4, cosine (min_lr_rate=0.1 → min_lr=2e-5)
+      #   AdamW, beta1=0.9, beta2=0.95, weight_decay=0.01
+      #   max_grad_norm=1.0, warmup_steps=200
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 4768
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Pure GDN hybrid spec (hybrid_attention_ratio=0.0 → all GDN, no MLA)
+      # Use no-TE spec to match FLA's native PyTorch layers for loss curve alignment
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
+      no_persist_layer_norm: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      # Perf: For a small 300M model, distributed optimizer (ZeRO-1) adds
+      # REDUCE-SCATTER + ALL-GATHER overhead with no memory savings (full
+      # optimizer state ~3.6GB easily fits per-rank). Switch to plain DDP
+      # ALL-REDUCE to match FLA. Memory cost: +3GB/rank, expected: ~10-15% faster.
+      # NOTE: overlap_param_gather requires distributed optimizer, so disable it.
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: false
+      ddp_average_in_collective: true
+
+      # Data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints
+      finetune: false
+      auto_continue_train: false
+      save: ./output/zebra_llama_300M_gdn_pure-pretrain
+      save_interval: 1024
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml
new file mode 100644
index 000000000..f769fc9cd
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_300M_kda_pure-pretrain.yaml
@@ -0,0 +1,161 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_300M_kda_pure-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_300M_kda_pure.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_300M_KDA_Pure_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      num_workers: 8
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+      log_interval: 100
+      check_for_nan_in_loss_and_grad: false
+
+      # Perf: Megatron defaults to inserting dist.barrier() before every L1
+      # timer measurement (~5-10/iter) to make per-stage timings comparable.
+      # This serializes ranks. Disable for production training.
+      barrier_with_L1_time: false
+
+      # Match FLA seed for reproducibility comparison
+      seed: 42
+
+      # Fix: norm_epsilon in model YAML maps to args.norm_epsilon, but
+      # TransformerConfig uses layernorm_epsilon (default 1e-5).
+      # Must explicitly set layernorm_epsilon to match FLA's 1e-6.
+      layernorm_epsilon: 1.0e-6
+
+      # CRITICAL: Megatron's TransformerConfig defaults hidden_dropout=0.1 and
+      # attention_dropout=0.1. Even though mamba_base.yaml sets these to 0.0,
+      # the YAML inheritance is being overridden by language_model.yaml (which
+      # sets 0.1) and the override is not propagating to args. FLA does NOT
+      # apply any dropout. Force these to 0 here to match.
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # Training schedule — matched to FLA (8 GPUs):
+      #   FLA: per_device_train_batch_size=128, 8 GPUs → global=1024
+      #   tokens/step = 1024 * 2048 = 2,097,152
+      #   total tokens ~= 4768 * 2,097,152 ≈ 10B
+      train_iters: 50
+      micro_batch_size: 8
+      global_batch_size: 64
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — matched to FLA:
+      #   lr=2e-4, cosine (min_lr_rate=0.1 → min_lr=2e-5)
+      #   AdamW, beta1=0.9, beta2=0.95, weight_decay=0.01
+      #   max_grad_norm=1.0, warmup_steps=200
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 4768
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # Pure KDA hybrid spec — no-TE variant matches FLA's KDABlock layout
+      # exactly (fla/models/kda/modeling_kda.py):
+      #   - Pre-norm computed ONCE per layer in the wrapper
+      #     (KimiDeltaAttentionLayer.norm = WrappedTorchNorm).
+      #   - Mixer in_proj is plain ColumnParallelLinear (no fused norm).
+      #   - Mixer gate_norm = IdentityOp (side projections reuse the
+      #     wrapper-normed tensor — no second norm pass).
+      # The TE variant (kda_hybrid_stack_spec) re-normalizes hidden_states
+      # inside the mixer via gate_norm, costing ~12-15 GiB activation memory
+      # per layer × 12 = ~40 GiB peak, plus ~50 ms/iter for the extra norm
+      # launches. The no-TE variant saves both, mirroring how GDN matched
+      # FLA in docs/hybrid_models/GDN_FLA_PARITY.md.
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec_no_te']
+      use_fla_triton_kda: true
+
+      # Full FLA-exact kernel fusion (matches fla/layers/kda.py exactly):
+      #   use_fla_kda_in_kernel_gate=True  → chunk_kda(use_gate_in_kernel=True,
+      #     use_qk_l2norm_in_kernel=True); gate `-exp(A_log)*softplus(g+dt_bias)
+      #     +cumsum` fused inside the Triton kernel and recomputed in backward,
+      #     so the [B,T,H,K] fp32 activated-gate tensor is never materialized.
+      #   use_fla_fused_norm_gated=True   → out_norm = FusedRMSNormGated
+      #     (RMSNorm + sigmoid-gate + multiply in one Triton kernel); avoids
+      #     the post-norm fp32 tensor and the fp32-upcast gate save-for-backward.
+      # Speed: ~+15 % vs the unfused path (matches FLA's ~1480 ms/iter).
+      # Memory: ~-20 GiB peak (FusedRMSNormGated drops ~10 GiB activation +
+      #   in-kernel gate drops ~3 GiB + fewer DDP buckets ≈ ~20 GiB total).
+      # Loss-curve note: on ROCm the in-kernel `-exp(A_log)*softplus` and the
+      # fused norm+sigmoid+multiply accumulators run in bf16; ±1 ulp drift
+      # compounds across 12 layers and gives ~+0.2-0.4 lm-loss above FLA
+      # *unless* we also load FLA's init checkpoint (the drift cancels when
+      # both runs start from identical weights — proven on GDN). Without the
+      # init ckpt expect loss curve to track FLA shape but offset; with it,
+      # expect bit-perfect parity (GDN-style).
+      use_fla_kda_in_kernel_gate: true
+      use_fla_fused_norm_gated: true
+
+      no_persist_layer_norm: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      # Plain DDP (no distributed optimizer) — required for FLA loss-curve
+      # parity per docs/hybrid_models/GDN_FLA_PARITY.md line 143-149. Megatron's
+      # DistributedOptimizer (ZeRO-1) applies the optimizer update per-shard
+      # then all-gathers, which is mathematically equivalent in fp64 but
+      # NOT bit-identical to plain AdamW in bf16 → those ±1 ulp ordering
+      # differences compound into a measurable loss drift over ~100 iters
+      # even when iter-1 forward and gradient match FLA exactly.
+      #
+      # Memory cost: ~3.6 GiB / rank extra (un-sharded optim state). With
+      # FLA-init checkpoint + both fusions on + expandable_segments we have
+      # ~12 GiB free headroom at iter 200, so this fits.
+      overlap_grad_reduce: false
+      overlap_param_gather: false       # requires distributed optimizer
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      use_distributed_optimizer: false
+      ddp_average_in_collective: true   # divide gradients in NCCL collective
+
+      # Data — FLA-aligned FineWeb-Edu sample-10BT
+      # Converted from FLA's preprocessed Arrow dataset so both
+      # frameworks see the exact same tokens in the same order.
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints
+      finetune: false
+      auto_continue_train: false
+      save: ./output/zebra_llama_300M_kda_pure-pretrain
+      save_interval: 1024
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # Context parallel
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_300M_mamba_hybrid-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_300M_mamba_hybrid-pretrain.yaml
new file mode 100644
index 000000000..cd6ecaecf
--- /dev/null
+++ b/examples/megatron/configs/MI300X/zebra_llama_300M_mamba_hybrid-pretrain.yaml
@@ -0,0 +1,137 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:zebra_llama_300M_mamba_hybrid-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: zebra_llama_300M_mamba_hybrid.yaml
+    overrides:
+      wandb_project: "Primus_Zebra_Llama_300M_Mamba_Hybrid_Pretrain"
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      # 2 dataloader workers per rank — keeps the dataloader pipeline full
+      # without bloating host RSS (16 forked subprocs total at 8 GPUs).
+      num_workers: 2
+      create_attention_mask_in_dataloader: false
+
+      profile: false
+      use_pytorch_profiler: false
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+      log_interval: 100
+      check_for_nan_in_loss_and_grad: false
+      barrier_with_L1_time: false
+
+      seed: 42
+
+      # RMSNorm 1e-6 (override Megatron's default 1e-5)
+      layernorm_epsilon: 1.0e-6
+
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # Training schedule — same as the GDN hybrid 300M run for direct
+      # comparison:
+      #   8 GPUs, micro=128, global=1024
+      #   tokens/step = 1024 × 2048 = 2,097,152
+      #   4768 steps × 2.097M ≈ 10B tokens
+      train_iters: 50
+      micro_batch_size: 8
+      global_batch_size: 64
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+      original_max_position_embeddings: 2048
+
+      # Optimizer — AdamW with cosine decay (lr 2e-4 → 2e-5 over 4768 iters)
+      clip_grad: 1.0
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 200
+      lr_decay_iters: 4768
+      lr_decay_style: cosine
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # ─────────────────────────────────────────────────────────────────────
+      # FLA runtime knobs — declarative YAML surface for the
+      # PRIMUS_FLA_*/PRIMUS_FUSED_CE* env vars.  Mirror exactly the env
+      # vars `tools/hybrid/launch_mamba_hybrid_300M.sh` previously exported.
+      # Consumed by primus.backends.megatron.patches.fla_runtime_patches
+      # at phase="build_args".  Env vars on the launcher still win.
+      # ─────────────────────────────────────────────────────────────────────
+      use_fla_fused_swiglu: true       # FLA Triton SwiGLU
+      use_fla_fused_rmsnorm: true      # FLA Triton RMSNorm
+      use_fla_fused_gated_norm: true   # Same env var (PRIMUS_FLA_NORM)
+      # ── Dataset source selector ───────────────────────────────────────
+      # use_fla_data + fla_cache_dir together pick the data path:
+      #   use_fla_data=false (or fla_cache_dir empty): vanilla Megatron
+      #     GPTDataset reading train_data_path (.bin/.idx).
+      #   use_fla_data=true  AND fla_cache_dir=<path>: replace
+      #     GPTDataset with tools/fla_order_dataset.FLAOrderGPTDataset
+      #     for bit-identical token order to FLA's HF DistributedSampler.
+      # The launcher script sets PRIMUS_FLA_CACHE_DIR; if you launch
+      # without it, uncomment fla_cache_dir below.
+      use_fla_data: true               # PRIMUS_FLA_DATA
+      # fla_cache_dir: /home/<user>/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train
+      fused_ce_mode: 1                 # 1=chunked FLA FusedLinearCrossEntropyLoss
+      fla_mla_attn: "1"                # FLA flash-attn for the MLA blocks
+
+      # Mamba2 + MLA hybrid stack via the proven `HybridStack` (no-TE) path.
+      #
+      # Why not the upstream `hybrid_stack_spec` (MambaStack)?  Megatron's
+      # `MambaStack` builder passes `pp_layer_offset` to MLASelfAttention,
+      # but this version of MLASelfAttention.__init__ doesn't accept it
+      # (Megatron API drift).  The HybridStack path used by the GDN/KDA
+      # hybrids doesn't pass that kwarg, so we mirror it for Mamba2.
+      spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'mamba_hybrid_stack_spec_no_te']
+      no_persist_layer_norm: true
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: meta-llama/Llama-3.2-1B
+
+      # Parallelism — ZeRO-1 distributed optimizer for memory headroom.
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      use_distributed_optimizer: true
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+      use_torch_fsdp2: false
+      ddp_average_in_collective: true
+
+      # Data — FLA-aligned FineWeb-Edu 10BT binary (same dataset the GDN
+      # hybrid 300M consumed, so per-token loss curves are directly
+      # comparable).  No PRIMUS_FLA_DATA / PRIMUS_FLA_CACHE_DIR needed —
+      # this is the indexed mmap version Megatron consumes natively.
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # Checkpoints — train from scratch; save once at the end (matches the
+      # GDN hybrid run, avoids mid-training host-RAM OOM from the state-dict
+      # materialisation).
+      finetune: false
+      auto_continue_train: false
+      load: null
+      save: ./output/zebra_llama_300M_mamba_hybrid-pretrain
+      save_interval: 99999
+      disable_last_saving: false
+      ckpt_format: torch
+
+      # Turbo off (matches the GDN hybrid baseline)
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      context_parallel_size: 1
diff --git a/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml
index b8019be76..82473699a 100644
--- a/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml
+++ b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml
@@ -20,7 +20,7 @@ modules:
       log_avg_skip_iterations: 2
       log_avg_reset_interval: 50
 
-      train_iters: 100
+      train_iters: 50
       micro_batch_size: 4
       global_batch_size: 32
 
diff --git a/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml
index a7083c069..24962376f 100644
--- a/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml
+++ b/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml
@@ -20,7 +20,7 @@ modules:
       log_avg_skip_iterations: 2
       log_avg_reset_interval: 50
 
-      train_iters: 100
+      train_iters: 50
       micro_batch_size: 2
       global_batch_size: 16
 
diff --git a/examples/megatron/configs/MI325X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI325X/deepseek_v2_lite-BF16-pretrain.yaml
index d2c8d73f4..a8636a21f 100644
--- a/examples/megatron/configs/MI325X/deepseek_v2_lite-BF16-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/deepseek_v2_lite-BF16-pretrain.yaml
@@ -44,6 +44,9 @@ modules:
       expert_model_parallel_size: ${PRIMUS_EP:8}
       overlap_grad_reduce: true
       overlap_param_gather: true
+      # NOTE: gradient_accumulation_fusion=true raised throughput ~1% but caused
+      # an intermittent NaN in the grad-norm check (backward wgrad accumulation),
+      # so it is kept off for reliability.
       gradient_accumulation_fusion: false
 
       # data
@@ -73,10 +76,13 @@ modules:
       # Turbo
       enable_primus_turbo: true
       use_turbo_attention: true
-      use_turbo_grouped_gemm: false
+      # PrimusTurbo grouped GEMM for the MoE experts (replaces the removed
+      # use_turbo_grouped_gemm). Required for sync-free MoE stage >= 2.
+      use_turbo_grouped_gemm: true
 
       # deepep
       use_turbo_deepep: true
+      # DeepEP does not support moe_shared_expert_overlap; keep it off.
       moe_shared_expert_overlap: false
       moe_router_dtype: fp32
 
@@ -87,7 +93,7 @@ modules:
       # sync-free moe support stage 0-3, 0 means not use sync-free moe
       # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory
       # stage 2 is recommended for better performance
-      turbo_sync_free_moe_stage: 1
+      turbo_sync_free_moe_stage: 2
 
       # remove once super flag is functional again
       moe_use_fused_router_with_aux_score: true
diff --git a/examples/megatron/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml b/examples/megatron/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml
index f4feaffdc..8eb3d0fcd 100644
--- a/examples/megatron/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml
@@ -110,3 +110,14 @@ modules:
       cross_entropy_loss_fusion: true
 
       fp8: hybrid
+      # FP8 delayed scaling (default) is fastest, but the delayed_fp8_scaling
+      # stream patches crash with a GPU memory fault (sum_and_scatter) at ~iter
+      # 6; disable them. Alternative: fp8_recipe=tensorwise (no patches).
+      fp8_recipe: delayed
+      disable_delayed_scaling_patches: true
+      # enable_primus_turbo: true
+      # use_turbo_attention: true
+      # use_turbo_grouped_gemm: false
+      # enable_primus_turbo: false
+      # enable_turbo_attention_float8 : false
+      # enable_turbo_gemm_float8 : false
diff --git a/examples/megatron/configs/MI325X/llama3.1_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI325X/llama3.1_70B-FP8-pretrain.yaml
index 55809a6d4..3ff515822 100644
--- a/examples/megatron/configs/MI325X/llama3.1_70B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/llama3.1_70B-FP8-pretrain.yaml
@@ -81,5 +81,10 @@ modules:
 
       # enable fp8 training
       fp8: hybrid
+      # FP8 delayed scaling (default) + the delayed_fp8_scaling stream patches
+      # crash with a GPU memory fault; disable the patches. (This also drops the
+      # data-prefetch double-buffer, freeing some device memory.)
+      fp8_recipe: delayed
+      disable_delayed_scaling_patches: true
       moe_use_legacy_grouped_gemm: false
       no_fp8_weight_transpose_cache: true
diff --git a/examples/megatron/configs/MI325X/llama3.1_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI325X/llama3.1_8B-FP8-pretrain.yaml
index 2228dcffa..13d35f0ca 100644
--- a/examples/megatron/configs/MI325X/llama3.1_8B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/llama3.1_8B-FP8-pretrain.yaml
@@ -77,4 +77,14 @@ modules:
 
       # enable fp8 training
       fp8: hybrid
+      # FP8 recipe: keep delayed scaling (the inherited default), which is the
+      # fastest FP8 path here (~3% higher throughput than tensorwise/current
+      # scaling because it reuses cached amax history instead of recomputing
+      # per-step amax). However, the delayed_fp8_scaling stream patches
+      # (secondary-stream grad-zero + data prefetch) have a use-after-free that
+      # crashes with a GPU memory fault (sum_and_scatter, "write to read-only
+      # page") at ~iter 6, so disable them. Alternative: fp8_recipe=tensorwise
+      # (Turbo-native, no patches, ~3% slower).
+      fp8_recipe: delayed
+      disable_delayed_scaling_patches: true
       moe_use_legacy_grouped_gemm: false
diff --git a/examples/megatron/configs/MI325X/llama3.3_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI325X/llama3.3_70B-FP8-pretrain.yaml
index 6337830f1..f79dc1bc9 100644
--- a/examples/megatron/configs/MI325X/llama3.3_70B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/llama3.3_70B-FP8-pretrain.yaml
@@ -17,7 +17,11 @@ modules:
       log_avg_reset_interval: 50
 
       train_iters: 10
-      micro_batch_size: 3
+      # NOTE: reduced from micro_batch_size=3 to 1. At MBS=3 this config both
+      # OOMs and, once memory is freed, hits a deterministic FP8 NaN loss at
+      # iter 2 (reproducible with both delayed and tensorwise scaling). MBS=1
+      # trains stably (matches the working llama3.1_70B FP8 config).
+      micro_batch_size: 1
       global_batch_size: 24
 
       seq_length: 8192
@@ -69,7 +73,7 @@ modules:
       # recompute
       recompute_granularity: full # full, selective
       recompute_method: block # uniform, block
-      recompute_num_layers: 80 # int
+      recompute_num_layers: 72 # int
 
       # Turbo
       enable_primus_turbo: true
@@ -77,10 +81,12 @@ modules:
       use_turbo_grouped_gemm: false
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid
+      fp8_recipe: delayed
+      disable_delayed_scaling_patches: true
       moe_use_legacy_grouped_gemm: false
       no_fp8_weight_transpose_cache: true
diff --git a/examples/megatron/configs/MI325X/qwen3_30B_A3B-BF16-pretrain.yaml b/examples/megatron/configs/MI325X/qwen3_30B_A3B-BF16-pretrain.yaml
index a9c45a614..c6772b255 100644
--- a/examples/megatron/configs/MI325X/qwen3_30B_A3B-BF16-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/qwen3_30B_A3B-BF16-pretrain.yaml
@@ -82,7 +82,13 @@ modules:
 
       # Turbo
       enable_primus_turbo: true
-      use_turbo_attention: false
+      # AMD-optimized attention kernel (faster and typically lower activation
+      # memory than the default backend).
+      use_turbo_attention: true
+      # NOTE: use_turbo_grouped_gemm (the replacement for the removed
+      # use_turbo_grouped_gemm) was tested but uses more device memory than the
+      # legacy grouped GEMM here and OOMs at MBS=8 (only 5 recompute layers), so
+      # this model stays on the legacy grouped-GEMM path.
       use_turbo_grouped_gemm: false
       use_turbo_rms_norm: false # bug
 
diff --git a/examples/megatron/configs/MI325X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI325X/qwen3_30B_A3B-FP8-pretrain.yaml
index ef8965ae6..47871d4fd 100644
--- a/examples/megatron/configs/MI325X/qwen3_30B_A3B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI325X/qwen3_30B_A3B-FP8-pretrain.yaml
@@ -39,9 +39,13 @@ modules:
       norm_epsilon: 1.0e-6
 
       # recompute
+      # Recomputing full MoE layers (router + permute + 128-expert grouped GEMM)
+      # is very expensive (~2k tokens/s per layer here), so recompute only 4 of
+      # 48 layers. FP8 weights leave enough memory headroom to keep the rest of
+      # the activations resident.
       recompute_granularity: full # full, selective
       recompute_method: block # uniform, block
-      recompute_num_layers: 5 # int
+      recompute_num_layers: 4 # int
 
       # parallel
       tensor_model_parallel_size: ${PRIMUS_TP:1}
@@ -82,6 +86,10 @@ modules:
 
       # Turbo
       enable_primus_turbo: true
+      # NOTE: turbo attention is intentionally OFF for FP8. The BF16 turbo kernel
+      # inside the FP8 pipeline is ~8% slower than the default TE attention, and
+      # the FP8 flash-attention kernel (enable_turbo_attention_float8) is not
+      # compatible with this model's GQA head layout.
       use_turbo_attention: false
       use_turbo_grouped_gemm: false
       use_turbo_rms_norm: false # bug
@@ -91,7 +99,9 @@ modules:
       moe_shared_expert_overlap: false
       moe_router_dtype: fp32
 
-      # 64 or 80 for ep8, 32 for ep16-64 is best practice
+      # 64 or 80 for ep8, 32 for ep16-64 is best practice. 80 is required for
+      # the fast regime here: dropping below ~78 CUs cliffs throughput ~13%
+      # (DeepEP comm stalls), so keep 80.
       turbo_deepep_num_cu: 80
       turbo_deepep_use_comm_stream: false
 
@@ -105,5 +115,10 @@ modules:
 
       # enable fp8 training
       fp8: hybrid
+      # FP8 delayed scaling (default) is fastest, but the delayed_fp8_scaling
+      # stream patches crash with a GPU memory fault (sum_and_scatter) at ~iter
+      # 6; disable them. Alternative: fp8_recipe=tensorwise (no patches).
+      fp8_recipe: delayed
+      disable_delayed_scaling_patches: true
       moe_use_legacy_grouped_gemm: true
       gradient_accumulation_fusion: true
diff --git a/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml b/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml
new file mode 100644
index 000000000..a0ab10937
--- /dev/null
+++ b/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml
@@ -0,0 +1,116 @@
+### DeepSeek-R1-Distill-Qwen-1.5B ODC LB-Mini SFT example (single-node).
+### Sequence-length load balancing (LB-Mini) with the "fit" cost model on the
+### torch-FSDP2 + ODC path; packed (thd) attention enabled.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:deepseek1.5B-odc-lbmini}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  post_trainer:
+    framework: megatron
+    config: sft_trainer.yaml
+
+    model: deepseek_r1_distill_qwen_1.5B.yaml
+
+    overrides:
+      stage: sft
+
+      hf_path: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
+      tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
+
+      sft_dataset_name: "zai-org/LongAlign-10k"
+      sft_conversation_format: "messages"
+
+      # ODC on-demand communication (rocSHMEM/XGMI P2P gradient reduction) on the
+      # torch-FSDP2 path. Master switch (was the ODC_ENABLE env var).
+      enable_odc: true
+      odc_phase: 2
+
+      # thd segmented (packed) attention.
+      enable_packed_sequences: true
+      use_packed_attention: true
+      # LB-Mini sequence-length load balancing (was the ODC_LB_MINI env var).
+      enable_odc_lb_mini: true
+      # LB-Mini cost model: fit (attention-aware a*s^2 + b*s, 1.5B coefficients).
+      lb_mini_cost_model: fit
+      lb_mini_max_token_len: 32768
+
+      # ODC P2P backend: rocshmem (validated single-node host/XGMI-IPC). trainer_base default is mori.
+      odc_p2p_backend: rocshmem
+
+      wandb_project: "Primus_ODC_DeepSeek1.5B"
+      stderr_sink_level: DEBUG
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 100
+
+      train_iters: 50
+      micro_batch_size: 1
+      global_batch_size: 16
+
+      seq_length: 32768
+      max_position_embeddings: 32768
+
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: 100
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+
+      seed: 1234
+      eod_mask_loss: false
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+
+      enable_fused_linear_ce: false
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+      use_distributed_optimizer: false
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+
+      finetune: true
+      load: null
+      save: null
+      save_interval: 1000
+      eval_interval: 1000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch_dist
+
+      bf16: true
+
+      enable_primus_turbo: false
+      use_turbo_attention: false
+      use_turbo_grouped_mlp: false
+      use_turbo_rms_norm: false
+
+      eval_iters: 0
+
+      profile: false
+      use_pytorch_profiler: true
+      use_nsys_profiler: false
+      profile_step_start: 10
+      profile_step_end: 11
+      profile_ranks: [0]
+      pytorch_profiler_collect_shapes: true
+      pytorch_profiler_collect_callstack: false
+      pytorch_profiler_collect_chakra: false
+      record_shapes: false
+      record_memory_history: false
+      nvtx_ranges: false
+
+      lora:
+        enabled: false
diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml
new file mode 100644
index 000000000..277c129f7
--- /dev/null
+++ b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml
@@ -0,0 +1,138 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:deepseek_v4_flash-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+# DeepSeek-V4 Flash BF16 pretraining config tuned for MI355X.
+# This is a smoke / scaffold config — values (especially parallelism) will be
+# revised once the V4 builder + hybrid attention land in Phase 3-5.
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_DeepSeek_V4_Pretrain"
+      stderr_sink_level: DEBUG
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # hyper parameters
+      train_iters: 50
+      micro_batch_size: 1
+      global_batch_size: 256
+      seq_length: ${PRIMUS_SEQ_LENGTH:4096}
+      max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096}
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+
+      # parallel — scaffold defaults; will be retuned in Phase 6.
+      tensor_model_parallel_size: ${PRIMUS_TP:1}
+      pipeline_model_parallel_size: ${PRIMUS_PP:8}
+      expert_model_parallel_size: ${PRIMUS_EP:8}
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: true
+
+      # data
+      mock_data: true
+      train_data_path: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      valid_data_path: null
+      test_data_path: null
+
+      # ---------- DeepSeek-V4 specific ----------
+      hybrid_attention_enabled: true
+      attn_sink: true
+      hc_use_sinkhorn: true
+      mtp_use_separate_hc_head: true
+      moe_router_score_function: sqrtsoftplus
+      swiglu_limit: 10.0
+
+      # ---------- Optimizer ----------
+      # NOTE: Phase 7 will add Muon for the latent / hidden projections.
+      #       Until then we run plain BF16 AdamW (precision-aware).
+      use_precision_aware_optimizer: true
+      main_grads_dtype: bf16
+      exp_avg_dtype: bf16
+      exp_avg_sq_dtype: bf16
+
+      # rope fusion
+      enable_experimental: true
+      apply_rope_fusion: false   # V4 uses partial RoPE on a 512-dim head
+
+      # recompute
+      recompute_granularity: full
+      recompute_method: uniform
+      recompute_num_layers: 1
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 20000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch
+      eval_iters: 0
+
+      # turbo / deepep — keep off by default; enable via env vars when
+      # benchmarking the Turbo path (see run_deepseek_v4.sh).  Plan-3 P22
+      # routes the dense (compress_ratio == 0) attention layers through
+      # PrimusTurboAttention when ``use_turbo_attention=true``; the V4
+      # builder auto-derives ``use_sink_attention`` /
+      # ``sink_sliding_window`` from ``attn_sink`` / ``attn_sliding_window``
+      # so the YAML stays free of Turbo-internal knobs.
+      enable_primus_turbo: ${PRIMUS_ENABLE_TURBO:true}
+      # use_turbo_attention stays OFF: the dense (cr=0) path must run the
+      # triton_v2 sparse-MLA backend below, not PrimusTurboAttention (which
+      # would take dispatch precedence).
+      use_turbo_attention: ${PRIMUS_USE_TURBO_ATTENTION:false}
+      use_turbo_grouped_gemm: true
+      use_turbo_rms_norm: true
+      use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true}
+      moe_shared_expert_overlap: false
+      moe_router_dtype: fp32
+
+      # deepep tuning (64 or 80 for ep8, 32 for ep16-64 is best practice)
+      turbo_deepep_num_cu: 80
+      turbo_deepep_use_comm_stream: false
+
+      # sync-free moe support (stage 1-2; 0 = off). stage 2 = best perf.
+      turbo_sync_free_moe_stage: 1
+
+      # V4 attention backend selection (unified string selectors; default triton_v2).
+      #   use_v4_attention_backend     (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon
+      #   use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0
+      use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:triton_v2}
+      use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:triton_v2}
+
+      # FP8 (E4M3) Indexer QK path (CSA selector); BF16 index-score preserved.
+      use_v4_fp8_indexer: ${PRIMUS_USE_V4_FP8_INDEXER:false}
+
+      # plan-5 P29 (RESCOPED): wrap sinkhorn_normalize with a cached
+      # torch.compile(fullgraph=True, dynamic=False) build. Default
+      # false until G32 (parity) + G33b (trace) flip it on.
+      use_v4_compiled_sinkhorn: ${PRIMUS_USE_V4_COMPILED_SINKHORN:false}
+
+      moe_use_fused_router_with_aux_score: true
+      moe_permute_fusion: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml
new file mode 100644
index 000000000..bc5026810
--- /dev/null
+++ b/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml
@@ -0,0 +1,166 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:deepseek_v4_flash-fp8-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+# DeepSeek-V4 Flash FP8 pretraining config tuned for MI355X.
+# Derived from deepseek_v4_flash-BF16-pretrain.yaml; the ONLY substantive
+# delta is the FP8 training block at the bottom.
+#
+# Precision recipe (paper §4.x quantization / techblog §9.6):
+#   Paper = "FP4 + FP8 Mixed": MoE experts + CSA-Indexer QK in FP4 (MXFP4),
+#   EVERYTHING ELSE in FP8, all with the **ue8m0** microscaling scale format.
+#   ue8m0 = OCP Microscaling (MX) block scale → on this stack that is
+#   `fp8_recipe: mxfp8` → TE MXFP8BlockScaling → Primus-Turbo MX_BLOCKWISE
+#   with scale_dtype=E8M0 (fp8_utils.py:148), native on MI355X/CDNA4.
+#
+# Integration gap vs the paper (NOT config — Primus V4 TODO, techblog item 10
+# "Phase 2 FP4/FP8 Mixed"): the FP4 expert / FP4-Indexer path is not yet wired
+# in V4, so experts run at FP8 here (FP8 everywhere) rather than FP4. This is
+# the closest supported step toward the paper recipe. FP8 is highly outlier-
+# sensitive, which is why the paper pairs it with clamped SwiGLU (swiglu_limit,
+# set below) — keep that on whenever FP8 is enabled.
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_DeepSeek_V4_Pretrain"
+      stderr_sink_level: DEBUG
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      # hyper parameters
+      train_iters: 50
+      micro_batch_size: 1
+      global_batch_size: 256
+      seq_length: ${PRIMUS_SEQ_LENGTH:4096}
+      max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096}
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+
+      # parallel — scaffold defaults; will be retuned in Phase 6.
+      tensor_model_parallel_size: ${PRIMUS_TP:1}
+      pipeline_model_parallel_size: ${PRIMUS_PP:8}
+      expert_model_parallel_size: ${PRIMUS_EP:8}
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: true
+
+      # data
+      mock_data: true
+      train_data_path: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      valid_data_path: null
+      test_data_path: null
+
+      # ---------- DeepSeek-V4 specific ----------
+      hybrid_attention_enabled: true
+      attn_sink: true
+      hc_use_sinkhorn: true
+      mtp_use_separate_hc_head: true
+      moe_router_score_function: sqrtsoftplus
+      swiglu_limit: 10.0
+
+      # ---------- Optimizer ----------
+      # NOTE: Phase 7 will add Muon for the latent / hidden projections.
+      #       Until then we run plain BF16 AdamW (precision-aware).
+      use_precision_aware_optimizer: true
+      main_grads_dtype: bf16
+      exp_avg_dtype: bf16
+      exp_avg_sq_dtype: bf16
+
+      # rope fusion
+      enable_experimental: true
+      apply_rope_fusion: false   # V4 uses partial RoPE on a 512-dim head
+
+      # recompute
+      recompute_granularity: full
+      recompute_method: uniform
+      recompute_num_layers: 1
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 20000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch
+      eval_iters: 0
+
+      # turbo / deepep — keep off by default; enable via env vars when
+      # benchmarking the Turbo path (see run_deepseek_v4.sh).  Plan-3 P22
+      # routes the dense (compress_ratio == 0) attention layers through
+      # PrimusTurboAttention when ``use_turbo_attention=true``; the V4
+      # builder auto-derives ``use_sink_attention`` /
+      # ``sink_sliding_window`` from ``attn_sink`` / ``attn_sliding_window``
+      # so the YAML stays free of Turbo-internal knobs.
+      enable_primus_turbo: ${PRIMUS_ENABLE_TURBO:true}
+      # use_turbo_attention stays OFF: the dense (cr=0) path must run the
+      # triton_v2 sparse-MLA backend below, not PrimusTurboAttention (which
+      # would take dispatch precedence).
+      use_turbo_attention: ${PRIMUS_USE_TURBO_ATTENTION:false}
+      use_turbo_grouped_gemm: true
+      use_turbo_rms_norm: true
+      use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true}
+      moe_shared_expert_overlap: false
+      moe_router_dtype: fp32
+
+      # deepep tuning (64 or 80 for ep8, 32 for ep16-64 is best practice)
+      turbo_deepep_num_cu: 80
+      turbo_deepep_use_comm_stream: false
+
+      # sync-free moe support (stage 1-2; 0 = off). stage 2 = best perf.
+      turbo_sync_free_moe_stage: 1
+
+      # V4 attention backend selection (unified string selectors; default triton_v2).
+      #   use_v4_attention_backend     (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon
+      #   use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0
+      use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:triton_v2}
+      use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:triton_v2}
+
+      # plan-5 P29 (RESCOPED): wrap sinkhorn_normalize with a cached
+      # torch.compile(fullgraph=True, dynamic=False) build. Default
+      # false until G32 (parity) + G33b (trace) flip it on.
+      use_v4_compiled_sinkhorn: ${PRIMUS_USE_V4_COMPILED_SINKHORN:false}
+
+      moe_use_fused_router_with_aux_score: true
+      moe_permute_fusion: true
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
+
+      # ---------- FP8 training ----------
+      # fp8 = number format (E4M3); fp8_recipe = scaling strategy.
+      # The paper uses ue8m0 microscaling (= mxfp8, 1x32 block / E8M0 scale), but
+      # mxfp8 is NOT runnable on this gfx950 build: turbo grouped-GEMM has no MX
+      # path, and TE's ROCm MXFP8 GEMM requires K%128==0 (V4 has a K=224 proj).
+      # So we use `tensorwise` (per-tensor scale) — the working recipe. This gives
+      # the paper's fp8 *layout* (all weight GEMMs fp8) with a non-ue8m0 scale.
+      # Override via FP8 / FP8_RECIPE env knobs (CLI wins); FP8=null => BF16.
+      fp8: e4m3
+      fp8_recipe: tensorwise
+      # Route the attention/dense linear projections through PrimusTurboLinear so
+      # they quantize to fp8 too (not just the MoE experts) — matches the paper's
+      # "fp8 on all weight GEMMs". Without this the projections stay bf16.
+      use_turbo_gemm: true
+
+      moe_router_padding_for_quantization: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml
new file mode 100644
index 000000000..8ba4c13fb
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml
@@ -0,0 +1,209 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MI355X)
+#
+# Combines Megatron DDP + distributed optimizer with the local spec provider
+# (PrimusTurboFloat8LocalSpecProvider) for FP8 tensorwise training.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - PrimusTurboFloat8LocalSpecProvider (NO TransformerEngine dependency)
+#   - FP8 hybrid tensorwise (per-module FP8 via Primus Turbo)
+#   - Primus Turbo attention
+#   - torch.compile enabled (per_block strategy, compatible with local spec + overlap)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      # Short LR warmup ramp over the first optimizer steps for FP8 stability.
+      nemo_aligned_lr_warmup: true
+      warmup_train_steps: 2
+
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboFloat8LocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: true
+      main_params_dtype: fp32
+      main_grads_dtype: bf16
+      exp_avg_dtype: bf16
+      exp_avg_sq_dtype: bf16
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # FP8 Configuration — Tensorwise + Delayed FP8 via Primus Turbo
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: "hybrid"
+      # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects
+      # tensorwise FP8 with delayed scaling. Resolution path:
+      # primus/backends/megatron/core/extensions/primus_turbo_float8_local.py
+      # :: Float8{Column,Row}ParallelLinear._use_delayed_scaling.
+      fp8_recipe: "delayed"
+      fp8_margin: 0
+      fp8_amax_history_len: 1024
+      fp8_amax_compute_algo: "max"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+      fp8_reduce_amax: true
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 180
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_local_fp8
+      wandb_project: flux_12b_ddp_local_fp8
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+      use_dual_fp8_output_projection: false
+
+      seed: 2025
+      per_step_rng_reseed: false
+      nemo_chimera_init: false
+
+      # Torch Compile — compatible with Float8 local spec (per-module FP8)
+      torch_compile:
+        enable: true
+        strategy: "per_block"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        emulate_precision_casts: false
+        fused_ln_modulate: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml
new file mode 100644
index 000000000..d67662f76
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml
@@ -0,0 +1,202 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MLPerf Mode)
+#
+# NOTE: This is a benchmark-reproduction config for the MLPerf Training Flux.1
+# benchmark. It mirrors MLPerf logging/convergence conventions and is intended
+# for reproducing benchmark results, not as a general-purpose training starting
+# point. For everyday FP8 training use
+# flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml.
+#
+# MLPerf-compliant variant of the DDP + local spec FP8 config:
+#   - mlperf_mode: true — enables MLPerf logging via mlperf_logging.mllog
+#   - warmup_train_steps: 2 — synthetic data warmup for torch.compile + FP8
+#   - target_val_loss: 0.586 — convergence target for early stopping
+#   - DDP + distributed optimizer without precision-aware optimizer
+#   - PrimusTurboFloat8LocalSpecProvider with FP8 tensorwise
+#   - Suppresses TensorBoard/WandB/print_rank_last during training
+#   - Emits structured MLPerf events (INIT_START, RUN_START, EVAL, etc.)
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8_mlperf}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # MLPerf Mode
+      # ==========================================
+      mlperf_mode: true
+      warmup_train_steps: 2
+      target_val_loss: 0.586
+
+      # ==========================================
+      # MLPerf Training v5.1 Alignment
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboFloat8LocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # DATA — Real Energon data (not mock)
+      # ==========================================
+      mock_data: false
+      dataloader_type: external
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+      num_workers: 16
+      prefetch_factor: 4
+      max_samples_per_sequence: null
+
+      # Training iterations
+      train_iters: 5000
+      eval_interval: 512
+      eval_iters: 10
+      log_interval: 10
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings (MLPerf v5.1)
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      # Use FP32 optimizer states (m,v): strictly better time-to-train with
+      # no throughput trade-off vs BF16 on this FP8 + emulate_precision_casts recipe.
+      use_precision_aware_optimizer: false
+      main_params_dtype: fp32
+      main_grads_dtype: fp32
+      exp_avg_dtype: fp32
+      exp_avg_sq_dtype: fp32
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+      check_for_nan_in_loss_and_grad: false
+
+      # ==========================================
+      # FP8 — Tensorwise via Primus Turbo
+      # ==========================================
+      use_flash_attn: true
+
+      fp8: "hybrid"
+      # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects
+      # tensorwise FP8 with delayed scaling.
+      fp8_recipe: "delayed"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+      fp8_amax_history_len: 1
+      fp8_amax_compute_algo: "most_recent"
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled for MLPerf)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging — suppressed by mlperf_mode
+      tensorboard_dir: null
+      wandb_project: null
+      log_throughput: false
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: false
+      log_batch_size_to_tensorboard: false
+      log_learning_rate_to_tensorboard: false
+      log_memory_to_tensorboard: false
+
+      # Profiler — disabled for MLPerf
+      profile: false
+
+      # Primus Turbo
+      enable_primus_turbo: true
+      use_turbo_attention: true
+      use_dual_fp8_output_projection: false
+
+      seed: 42
+      # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf
+      # reproduction must opt in for run-to-run determinism).
+      per_step_rng_reseed: true
+
+      # Torch Compile — per_block: compile each transformer layer individually
+      torch_compile:
+        enable: true
+        strategy: "per_block"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        emulate_precision_casts: true
+        fused_ln_modulate: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml
new file mode 100644
index 000000000..10283f997
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml
@@ -0,0 +1,198 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + Local Spec + MXFP4 (MI355X)
+#
+# Combines Megatron DDP + distributed optimizer with the local spec provider
+# (PrimusTurboMXFP4LocalSpecProvider) for MXFP4 block-scaled training.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - PrimusTurboMXFP4LocalSpecProvider (NO TransformerEngine dependency)
+#   - MXFP4 (E2M1 + E8M0 block-of-32 scales) via Primus Turbo + AITER
+#   - Primus Turbo attention
+#   - torch.compile enabled (per_block strategy, compatible with local spec)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp4}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      # Short LR warmup ramp over the first optimizer steps for stability.
+      nemo_aligned_lr_warmup: true
+      warmup_train_steps: 2
+
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboMXFP4LocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # ==========================================
+      # BF16 + MXFP4 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: true
+      main_params_dtype: fp32
+      main_grads_dtype: bf16
+      exp_avg_dtype: bf16
+      exp_avg_sq_dtype: bf16
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # MXFP4 Configuration — Block-scaled via Primus Turbo + AITER
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp4: "mxfp4"
+      fp4_recipe: "mxfp4"
+      mxfp4_backward_precision: "mxfp4"  # "mxfp4" (pure) or "fp8" (hybrid)
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 180
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_local_mxfp4
+      wandb_project: flux_12b_ddp_local_mxfp4
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+
+      seed: 2025
+      per_step_rng_reseed: false
+      nemo_chimera_init: false
+
+      # Torch Compile — compatible with MXFP4 local spec (per-module FP4)
+      torch_compile:
+        enable: true
+        strategy: "per_block"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        emulate_precision_casts: false
+        fused_ln_modulate: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml
new file mode 100644
index 000000000..d4bddef34
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml
@@ -0,0 +1,184 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TransformerEngine Spec + BF16 (MI355X)
+#
+# BF16 training using Megatron DDP + distributed optimizer with
+# TransformerEngine modules (TEColumnParallelLinear / TERowParallelLinear /
+# TEDotProductAttention / TENorm). This is the BF16 baseline on the
+# TransformerEngine path; see the *_te_spec_fp8 variant for FP8.
+#
+# Key settings:
+#   bf16: true
+#   params_dtype: bfloat16
+#   micro_batch_size: 64 / global_batch_size: 512
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_bf16}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b_rope_fusion.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # TransformerEngine Spec (default)
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # ENERGON DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512
+
+      # ==========================================
+      # BF16 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # Optimizer settings (identical to BF16 baseline)
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: true
+
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # Memory Optimizations
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: null
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_te_bf16
+      wandb_project: flux_12b_ddp_te_bf16
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # Torch Compile — stack strategy (compiles the double/single DiT block
+      # stacks as inductor regions). Required to fit micro_batch_size 64 in
+      # BF16 on the TE path; matches the MLPerf NeMo reference (COMPILE_DIT
+      # strategy=stack).
+      torch_compile:
+        enable: true
+        strategy: "stack"
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml
new file mode 100644
index 000000000..4cb6ad043
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml
@@ -0,0 +1,209 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MI355X)
+#
+# FP8 training using Megatron DDP + distributed optimizer with TransformerEngine
+# modules and FP8 hybrid delayed scaling:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm
+#   - FP8 hybrid (E4M3 fwd, E5M2 bwd) with delayed scaling (amax history 1024)
+#   - RoPE fusion via apply_rope_fusion: true
+#   - Energon pre-encoded dataset
+#
+# Required environment variables for the TransformerEngine path (set before launch):
+#   export NVTE_FUSED_ATTN=1
+#   export NVTE_FUSED_ATTN_CK=1
+#   export NVTE_FP8_DPA_BWD=1
+#   export NVTE_USE_HIPBLASLT=1
+#   export USE_HIPBLASLT=1
+#   export TORCH_BLAS_PREFER_HIPBLASLT=1
+#   export NVTE_USE_CAST_TRANSPOSE_TRITON=1
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # ==========================================
+      # RoPE Fusion
+      # ==========================================
+      rotary_interleaved: true
+      apply_rope_fusion: true
+      position_embedding_type: rope
+
+      # ==========================================
+      # TransformerEngine Spec
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # FP8 — hybrid delayed scaling
+      # ==========================================
+      fp8: "hybrid"
+      fp8_recipe: "delayed"
+      fp8_margin: 0
+      fp8_amax_history_len: 1024
+      fp8_amax_compute_algo: "max"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+
+      # ==========================================
+      # Energon Dataset
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # ==========================================
+      # Batch Configuration
+      # ==========================================
+      # MI355X (256GB HBM3) fits the full MBS=64/GBS=512; tune to your hardware.
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # ==========================================
+      # Optimizer
+      # ==========================================
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # Memory / Misc
+      # ==========================================
+      use_flash_attn: true
+      empty_unused_memory_level: 0
+
+      # Manual GC — align GC timing across ranks to avoid stragglers
+      manual_gc: true
+      manual_gc_interval: 1000
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_ddp_te_fp8
+      wandb_project: flux_12b_ddp_te_fp8
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo — disabled for pure TE path
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # torch.compile — selective stack compilation for TE spec
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        strategy: "stack"
+        replace_qk_rmsnorm: true
+        disable_inductor_cudagraphs: false
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml
new file mode 100644
index 000000000..cd01468c8
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml
@@ -0,0 +1,202 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MLPerf Mode)
+#
+# NOTE: This is a benchmark-reproduction config for the MLPerf Training Flux.1
+# benchmark. It mirrors MLPerf logging/convergence conventions and is intended
+# for reproducing benchmark results, not as a general-purpose training starting
+# point. For everyday FP8 training on the TransformerEngine path use
+# flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml.
+#
+# MLPerf-compliant variant of the TE spec FP8 config:
+#   - mlperf_mode: true — enables MLPerf logging via mlperf_logging.mllog
+#   - warmup_train_steps: 2 — synthetic data warmup for torch.compile + FP8
+#   - target_val_loss: 0.586 — convergence target for early stopping
+#   - DDP + distributed optimizer
+#   - TransformerEngine modules with FP8 hybrid delayed scaling
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8_mlperf}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # MLPerf Mode
+      # ==========================================
+      mlperf_mode: true
+      warmup_train_steps: 2
+      target_val_loss: 0.586
+
+      # ==========================================
+      # MLPerf Training v5.1 Alignment
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # ==========================================
+      # RoPE Fusion
+      # ==========================================
+      rotary_interleaved: true
+      apply_rope_fusion: true
+      position_embedding_type: rope
+
+      # ==========================================
+      # TransformerEngine Spec
+      # ==========================================
+      transformer_impl: "transformer_engine"
+      adaln_plain_ops: true
+      adaln_always_jit_fuser: true
+
+      # ==========================================
+      # FP8 — hybrid delayed scaling
+      # ==========================================
+      fp8: "hybrid"
+      fp8_recipe: "delayed"
+      fp8_margin: 0
+      fp8_amax_history_len: 1024
+      fp8_amax_compute_algo: "max"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+
+      # ==========================================
+      # DATA — Real Energon data (not mock)
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+      mock_data: false
+      dataloader_type: external
+      num_workers: 8
+      max_samples_per_sequence: null
+
+      train_iters: 5000
+      eval_interval: 512
+      eval_iters: 10
+      log_interval: 10
+      save_interval: 10000
+
+      # ==========================================
+      # Batch Configuration
+      # ==========================================
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512
+
+      # ==========================================
+      # BF16 + FP8 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # ==========================================
+      # Optimizer (MLPerf v5.1)
+      # ==========================================
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+      check_for_nan_in_loss_and_grad: false
+
+      # ==========================================
+      # Memory / Misc
+      # ==========================================
+      use_flash_attn: true
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled for MLPerf)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging — suppressed by mlperf_mode
+      tensorboard_dir: null
+      wandb_project: null
+      log_throughput: false
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: false
+      log_batch_size_to_tensorboard: false
+      log_learning_rate_to_tensorboard: false
+      log_memory_to_tensorboard: false
+
+      # Profiler — disabled
+      profile: false
+
+      # Primus Turbo — disabled for pure TE path
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+      # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf
+      # reproduction must opt in for run-to-run determinism).
+      per_step_rng_reseed: true
+
+      # torch.compile — selective stack compilation for TE spec
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: false
+        strategy: "stack"
+        replace_qk_rmsnorm: true
+        disable_inductor_cudagraphs: true
+        emulate_precision_casts: false
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml
new file mode 100644
index 000000000..567b58aeb
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml
@@ -0,0 +1,175 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — FSDP2 + Local Spec + BF16, VAE Resample Mode (MI355X)
+#
+# FSDP2 (ZeRO-2) BF16 training with the local spec provider and
+# vae_latent_mode: resample.
+#
+# In resample mode, latents are re-drawn from stored mean+logvar via
+# reparameterization (mean + exp(0.5*logvar) * randn) at every training step.
+# This introduces per-step stochasticity in the VAE latents.
+#
+# Dataset must be an Energon pre-encoded dataset containing mean.pth and
+# logvar.pth per sample.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_bf16}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboLocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # ENERGON DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 2000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512  # 64 * 8 GPUs = 512
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # Mixed precision
+      bf16: true
+      fp16: false
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 4
+      dataloader_type: external
+
+      # ==========================================
+      # PyTorch FSDP2 Configuration (ZeRO-2)
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+
+      torch_fsdp2_reshard_after_forward: false  # ZeRO-2
+
+      use_fsdp2_fp32_param_optimizer: true
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: false
+
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # Memory Optimizations
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: null
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_bf16
+      wandb_project: flux_12b_fsdp2_local_bf16
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler (disabled)
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+
+      seed: 42
+
+      # Torch Compile
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml
new file mode 100644
index 000000000..1ee15a5fb
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml
@@ -0,0 +1,198 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — FSDP2 + Local Spec + FP8 + FP32 Param Optimizer (MI355X)
+#
+# FP8 training on the FSDP2 path with:
+#   - FP8 training (tensorwise, local spec)
+#   - BF16 all-gather (use_fsdp2_fp8_all_gather: false)
+#   - FP32 param optimizer (FP32 params + FP32 optimizer states)
+#   - No overlap grad norm (overlap_grad_norm: false)
+#   - torch.compile enabled
+#   - ZeRO-2 sharding (reshard_after_forward: false)
+#
+# Uses an Energon pre-encoded dataset with VAE resample mode.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_fp8}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      # VAE latent normalization (required for resample mode)
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+
+      # Resample mode: re-draw latents from mean+logvar each step
+      vae_latent_mode: resample
+
+      # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout
+      rotary_interleaved: true
+
+      # ==========================================
+      # PrimusTurboLocalSpecProvider
+      # ==========================================
+      transformer_impl: "local"
+
+      # ==========================================
+      # DATASET CONFIGURATION
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # Batch configuration
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # Mixed precision
+      bf16: true
+      fp16: false
+
+      # Optimizer settings
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader settings
+      num_workers: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # PyTorch FSDP2 Configuration (ZeRO-2)
+      # ==========================================
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+
+      torch_fsdp2_reshard_after_forward: false  # ZeRO-2
+      use_fsdp2_fp8_all_gather: false
+      fp8_all_gather_stochastic_rounding: true
+      # fp8_all_gather_deq_requant: true  # Dequant FP8->BF16 after AG, fresh dynamic requant downstream
+      use_triton_ops: true  # Triton @triton_op modulate/LN+modulate — eliminates ~510us dispatch overhead per graph
+      fsdp_prefetch_depth: 5
+      fp8_precompute_data_cache: false
+      optimizer_foreach: false
+      use_cpp_fp8_quantize: true
+      overlap_grad_norm: false
+
+      # Optimizer mode: FP32 params + FP32 optimizer states
+      use_fsdp2_fp32_param_optimizer: true
+      use_fsdp2_bf16_master_weight_optimizer: false
+
+      ckpt_format: torch_dist
+
+      use_distributed_optimizer: false
+
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      gradient_accumulation_fusion: false
+
+      # ==========================================
+      # FP8 Configuration — Tensorwise via Primus Turbo
+      # ==========================================
+
+      use_flash_attn: true
+
+      fp8: "hybrid"
+      # Dynamic (tensorwise) scaling: the FSDP2 path does not yet exercise the
+      # delayed amax allreduce, so use tensorwise here. Switch to
+      # `fp8_recipe: "delayed"` once FSDP2 + delayed is wired up.
+      fp8_recipe: "tensorwise"
+      fp8_wgrad: true
+      fp8_dot_product_attention: false
+      fp8_multi_head_attention: false
+
+      empty_unused_memory_level: 0
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_fp8
+      wandb_project: flux_12b_fsdp2_local_fp8
+      log_throughput: true
+      wall_clock_step_timer: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_memory_to_tensorboard: true
+
+      # PyTorch Profiler
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]
+      profile_step_start: 9
+      profile_step_end: 11
+      torch_profiler_record_shapes: true
+      torch_profiler_with_stack: true
+      torch_profiler_use_gzip: false
+      disable_profiler_activity_cpu: false
+
+      # Primus Turbo Configuration
+      enable_primus_turbo: true
+      use_turbo_attention: true
+      use_dual_fp8_output_projection: false
+
+      seed: 42
+
+      # Torch Compile — compatible with Float8 local spec (per-module FP8)
+      torch_compile:
+        enable: true
+        backend: "inductor"
+        mode: "default"
+        fullgraph: false
+        compile_optimizer: true
+        emulate_precision_casts: false
+        fused_ln_modulate: true
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml
new file mode 100644
index 000000000..0ec08c870
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml
@@ -0,0 +1,170 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M Pretraining Configuration (Pre-encoded Data Mode)
+#
+# This config demonstrates Flux 535M training with pre-encoded features.
+# Pre-encoded mode is faster and recommended for production training.
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \
+#   bash examples/run_pretrain.sh
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_535m.yaml
+
+    # Trainer class for diffusion models
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      # ============================================================================
+      # Model Configuration
+      # ============================================================================
+
+      model_type: diffusion_model
+
+      # ============================================================================
+      # Dataset Configuration — Synthetic / Mock Data (default)
+      # ============================================================================
+      #
+      # This 535M config defaults to in-memory synthetic data so it runs
+      # standalone for sanity checks, CI, and training-pipeline validation
+      # without a prepared Energon dataset. Samples are random tensors with
+      # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768).
+      #
+      # To train on real data instead: set `mock_data: false` and point
+      # `data_path` at a prepared Energon dataset (see notes at the bottom).
+      mock_data: true
+      mock_dataset:
+        class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+        params:
+          num_samples: 256         # In-memory synthetic samples (iterated cyclically)
+          image_size: 256          # 256 -> 32x32 latents (light/fast for testing)
+      # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Batch configuration
+      micro_batch_size: 2          # Per-GPU batch size (adjust based on VRAM)
+      global_batch_size: 16        # Total batch size across all GPUs
+      seq_length: 4096             # Sequence length for latent features
+
+      # DataLoader settings (used only with real Energon data, i.e. mock_data: false)
+      num_workers: 4               # Number of data loading workers per GPU
+      dataloader_type: external    # Required for Energon dataloaders
+
+      # ============================================================================
+      # Training Parameters
+      # ============================================================================
+
+      # Total training steps
+      train_iters: 100000          # Total training iterations
+      eval_interval: 1000          # Evaluate every N steps
+      eval_iters: 50               # Number of evaluation iterations
+
+      # Logging
+      log_interval: 10             # Log every N steps
+      tensorboard_dir: output/tensorboard/flux_535m_pretrain
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+
+      # Checkpointing (disabled)
+      save_interval: 5000          # Save checkpoint every N steps
+      save: null
+      load: null                   # Path to checkpoint for resuming (optional)
+      finetune: false
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Precision
+      bf16: true                   # Use bfloat16 (recommended for AMD MI355X)
+      fp16: false
+
+      # ============================================================================
+      # Optimizer Configuration
+      # ============================================================================
+
+      # Optimizer
+      optimizer: adam
+      lr: 1.0e-4                   # Peak learning rate
+      min_lr: 1.0e-5               # Minimum learning rate
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1.0e-8
+
+      # Gradient clipping
+      clip_grad: 1.0
+
+      # ============================================================================
+      # Learning Rate Scheduler
+      # ============================================================================
+
+      lr_decay_style: cosine
+      lr_warmup_iters: 1000        # Warmup iterations
+      lr_decay_iters: 100000       # Total decay steps (typically = train_iters)
+
+      # ============================================================================
+      # Distributed Training Configuration
+      # ============================================================================
+
+      # Parallelism settings (Flux 535M fits on 1 GPU, but can use DP for speed)
+      tensor_model_parallel_size: 1    # Tensor parallelism (no need for 535M)
+      pipeline_model_parallel_size: 1  # Pipeline parallelism
+
+      # Advanced settings
+      overlap_grad_reduce: true         # Overlap gradient communication
+      use_flash_attn: true              # Use Flash Attention 2
+      distributed_timeout_minutes: 60
+
+      # ============================================================================
+      # Seed and Reproducibility
+      # ============================================================================
+
+      seed: 42
+
+      # ============================================================================
+      # Monitoring and Logging
+      # ============================================================================
+
+      wandb_project: flux_535m_pretrain
+      wandb_exp_name: flux_535m_preencoded
+
+# ============================================================================
+# Notes
+# ============================================================================
+#
+# Dataset Preparation:
+#   1. Prepare pre-encoded dataset:
+#      tools/docker/primus data diffusion-encoded \
+#        --source-type directory --input-dir /data/raw \
+#        --output-dir /data/encoded --model-path black-forest-labs/FLUX.1-dev
+#   2. Copy dataset template to output directory:
+#      cp primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml \
+#         /data/encoded/dataset.yaml
+#   3. Run Energon indexing:
+#      energon prepare /data/encoded --num-workers 8
+#   4. Update data_path above to: /data/encoded/dataset.yaml
+#
+# For more information, see:
+#   - primus/configs/data/megatron/diffusion/README.md
+#   - examples/megatron/diffusion/README.md
+#
+# Single-node training (8 GPUs):
+#   EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \
+#   GPUS_PER_NODE=8 bash examples/run_pretrain.sh
+#
+# Multi-node training (4 nodes, 8 GPUs each):
+#   EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \
+#   NNODES=4 bash examples/run_slurm_pretrain.sh
+#
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml
new file mode 100644
index 000000000..1ec0e6433
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml
@@ -0,0 +1,199 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M FP8 Pretraining Configuration (Testing/Development)
+#
+# This config is for testing FP8 functionality with the minimal Flux 535M model
+# before scaling to the full 12B model. Use this to validate:
+#   - FP8 setup and configuration
+#   - Numerical stability
+#   - Memory and speed improvements
+#   - Transformer Engine compatibility
+#
+# Target Hardware: Single AMD MI355X GPU with ROCm 6.0+
+# Requires: Transformer Engine 2.1.0+ with ROCm backend
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml \
+#   GPUS_PER_NODE=1 bash examples/run_pretrain.sh
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain_fp8}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_535m_fp8.yaml  # Use FP8-enabled 535M config
+
+    # Trainer class for diffusion models
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      # ============================================================================
+      # Model Configuration
+      # ============================================================================
+
+      model_type: diffusion_model
+
+      # ============================================================================
+      # Dataset Configuration — Synthetic / Mock Data (default)
+      # ============================================================================
+      #
+      # This 535M config defaults to in-memory synthetic data so it runs
+      # standalone for sanity checks, CI, and training-pipeline validation
+      # without a prepared Energon dataset. Samples are random tensors with
+      # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768).
+      #
+      # To train on real data instead: set `mock_data: false` and point
+      # `data_path` at a prepared Energon dataset (see notes at the bottom).
+      mock_data: true
+      mock_dataset:
+        class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+        params:
+          num_samples: 256         # In-memory synthetic samples (iterated cyclically)
+          image_size: 256          # 256 -> 32x32 latents (light/fast for testing)
+      # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Batch configuration (can use larger batch with FP8)
+      micro_batch_size: 4          # Can increase from 2 to 4 with FP8 on 535M
+      global_batch_size: 32        # Small batch for quick testing
+      seq_length: 4096             # Sequence length for latent features
+
+      # DataLoader settings (used only with real Energon data, i.e. mock_data: false)
+      num_workers: 4               # Number of data loading workers per GPU
+      dataloader_type: external    # Required for Energon dataloaders
+
+      # ============================================================================
+      # Training Parameters (Quick Testing)
+      # ============================================================================
+
+      # Short training for validation
+      train_iters: 1000            # Just 1K steps for FP8 validation
+      eval_interval: 100           # Evaluate every 100 steps
+      eval_iters: 10               # Quick evaluation
+
+      # Logging
+      log_interval: 10             # Log every N steps
+      tensorboard_dir: output/tensorboard/flux_535m_pretrain_fp8
+      log_throughput: true
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+
+      # Checkpointing (disabled)
+      save_interval: 500           # Save more frequently for testing
+      save: null
+      load: null                   # Path to checkpoint for resuming (optional)
+      finetune: false
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Precision
+      bf16: true                   # Use bfloat16 for non-FP8 ops
+      fp16: false
+
+      # ============================================================================
+      # Optimizer Configuration
+      # ============================================================================
+
+      # Optimizer
+      optimizer: adam
+      lr: 1.0e-4                   # Peak learning rate
+      min_lr: 1.0e-5               # Minimum learning rate
+      weight_decay: 0.01
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1.0e-8
+
+      # Gradient clipping (important for FP8 stability)
+      clip_grad: 1.0
+
+      # ============================================================================
+      # Learning Rate Scheduler
+      # ============================================================================
+
+      lr_decay_style: cosine
+      lr_warmup_iters: 100         # Short warmup for testing
+      lr_decay_iters: 1000         # Match train_iters
+
+      # ============================================================================
+      # Distributed Training Configuration
+      # ============================================================================
+
+      # Single GPU configuration
+      tensor_model_parallel_size: 1    # No TP needed for 535M
+      pipeline_model_parallel_size: 1  # No PP needed
+      context_parallel_size: 1         # No CP needed
+
+      # Distributed settings
+      use_distributed_optimizer: false  # Not needed for single GPU
+      overlap_grad_reduce: false        # Not applicable for single GPU
+      use_flash_attn: true              # Use Flash Attention 2
+      distributed_timeout_minutes: 60
+
+      # ============================================================================
+      # Memory Optimization
+      # ============================================================================
+
+      # Activation checkpointing (not needed for 535M with FP8)
+      recompute_granularity: null       # No recompute needed
+      recompute_method: null
+      recompute_num_layers: null
+
+      # Sequence parallelism
+      sequence_parallel: false          # Not needed for single GPU
+
+      # ============================================================================
+      # Seed and Reproducibility
+      # ============================================================================
+
+      seed: 42
+
+      # ============================================================================
+      # Monitoring and Logging
+      # ============================================================================
+
+      wandb_project: flux_535m_pretrain_fp8
+      wandb_exp_name: flux_535m_fp8_test
+
+# ============================================================================
+# Notes - FP8 Validation with 535M
+# ============================================================================
+#
+# Hardware Requirements (with FP8):
+#   - Minimum: 1× MI355X 256GB
+#   - Memory per GPU: ~3-5GB (vs ~7-10GB BF16)
+#   - Training time: Minutes
+#
+# Validation Checklist:
+#   [ ] Setup FP8 environment (see docs/04-technical-guides/diffusion-models/fp8_training.md)
+#   [ ] Verify TE FP8 support is available
+#   [ ] Run this config to validate FP8 training
+#   [ ] Check logs for NaN/Inf (should be none)
+#   [ ] Verify memory usage is ~50% of BF16
+#   [ ] Verify training speed is 1.5-2x faster than BF16
+#   [ ] Check loss decreases normally
+#
+# Expected Results:
+#   - Training completes 1000 steps in 5-15 minutes
+#   - No NaN/Inf in losses
+#   - Memory usage: ~3-5GB
+#   - Speed: ~10-50 steps/sec (depending on hardware)
+#   - Loss should decrease normally
+#
+# If validation passes, proceed to one of the 12B FP8 configs, e.g.:
+#   - flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml (TransformerEngine FP8)
+#   - flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml (local-spec FP8)
+#
+# Troubleshooting:
+#   - If NaN/Inf: Check Transformer Engine FP8 support
+#   - If OOM: Reduce micro_batch_size
+#   - If slow: Verify ROCm FP8 tensor cores are being used
+#   - If unstable: Try fp8_wgrad: false in model config
+#
+# For more information: See docs/04-technical-guides/diffusion-models/fp8_training.md
diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml
new file mode 100644
index 000000000..49cb41194
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml
@@ -0,0 +1,58 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M with Guidance Embedding (Advanced Configuration)
+#
+# This config demonstrates Flux training with guidance embedding enabled.
+# This is an OPTIONAL advanced feature that allows for faster single-pass CFG
+# during inference, but requires training with guidance embedding enabled.
+#
+# IMPORTANT: Most users should use the standard flux_535m_pretrain.yaml config.
+# Only use this if you specifically need guidance embedding support.
+#
+# Usage:
+#   EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml \
+#   bash examples/run_pretrain.sh
+
+# Extend standard 535M config
+extends:
+  - flux_535m_pretrain.yaml
+
+modules:
+  pre_trainer:
+    overrides:
+      # ============================================================================
+      # Guidance Embedding Configuration (ADVANCED)
+      # ============================================================================
+
+      # Enable guidance embedding for single-pass CFG
+      # This adds a learned MLPEmbedder layer that conditions on guidance scale
+      guidance_embed: true
+
+      # Guidance scale used during training
+      # Model learns to adapt its predictions based on this scale
+      guidance_scale: 3.5
+
+      # ============================================================================
+      # Notes
+      # ============================================================================
+      #
+      # Training with guidance embedding:
+      # - Adds ~1-2% more parameters (guidance MLPEmbedder)
+      # - Allows single-pass CFG during inference (faster)
+      # - Requires more training data/iterations to converge
+      # - Model learns guidance as a conditioning signal
+      #
+      # Inference with guidance embedding:
+      # - Pipeline automatically detects guidance_embed layer
+      # - Uses single forward pass instead of batch doubling
+      # - ~2x faster CFG compared to explicit CFG
+      # - Guidance scale can be varied at inference time
+      #
+      # Standard approach (guidance_embed: false):
+      # - Default for most Primus training
+      # - Uses explicit CFG (batch doubling) at inference
+      # - More compatible with existing checkpoints
+      # - Slightly slower but more flexible
+      #
+      # See examples/megatron/diffusion/README.md for more details.
diff --git a/examples/megatron/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml
index 72baa5d20..731bf6fde 100644
--- a/examples/megatron/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml
@@ -46,7 +46,7 @@ modules:
       profile_step_start: 6
 
       # hyper parameters
-      train_iters: 10
+      train_iters: 50
       micro_batch_size: 8
       global_batch_size: 512
       seq_length: ${PRIMUS_SEQ_LENGTH:4096}
diff --git a/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml
new file mode 100644
index 000000000..5b62b5265
--- /dev/null
+++ b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml
@@ -0,0 +1,222 @@
+work_group: ${TEAM:amd}
+user_name: ${USER:root}
+exp_name: ${EXP_NAME:gpt_oss_20b}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: ${PRIMUS_MODEL:gpt_oss_20B}.yaml
+    overrides:
+
+      # Activate the migrated MLPerf pretrain trainer (mllog + MLPerf hooks).
+      stage: mlperf_pretrain
+
+      # tokenizer
+      tokenizer_type: Llama3Tokenizer
+      tokenizer_model: ${MODEL:meta-llama/Llama-3.1-8B}
+
+      # model
+      num_layers: 24
+      hidden_size: 2880
+      ffn_hidden_size: 2880
+      num_attention_heads: 64
+      num_query_groups: 8  # Group Query Attention (GQA) - matches HF num_key_value_heads
+      num_experts: 32
+      activation_func: swiglu  # SiLU activation (matches HF hidden_act: "silu")
+
+      # rotary
+      position_embedding_type: rope
+      rotary_base: 150000
+
+      # mixed-precision
+      attention_softmax_in_fp32: false
+      grad_reduce_in_bf16: ${PRIMUS_GRAD_REDUCE_IN_BF16:true}
+
+      # log
+      wandb_project: "Primus_GPT_OSS_20B"
+      stderr_sink_level: DEBUG
+      log_interval: ${LOG_INTERVAL:10}
+
+      # debug
+      # moe_router_force_load_balancing: true
+      # log_avg_skip_iterations: 2
+      # log_avg_reset_interval: 50
+
+      # profile
+      profile: ${PRIMUS_PROFILE:false}
+      use_pytorch_profiler: ${PRIMUS_PROFILE:false}
+      profile_step_end: ${PRIMUS_PROFILE_STEP_END:32}
+      profile_step_start: ${PRIMUS_PROFILE_STEP_START:16}
+      profile_ranks: [0,1,2,3,4,5,6,7]
+
+      # enable fp8 training
+      fp8: e4m3
+      fp8_recipe: tensorwise
+      clip_grad: 1.0  # Gradient clipping (already default, but explicit)
+      check_for_nan_in_loss_and_grad: false
+
+      # hyper parameters
+      train_iters: ${PRIMUS_TRAIN_ITERS:1200000}
+      micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2}
+      global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:16}
+      seq_length: ${PRIMUS_SEQ_LENGTH:8192}
+      max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:131072}
+      seed: ${SEED:1234}  # Random seed for reproducibility
+      lr: ${PRIMUS_LR:8.0e-4}  # Reduced from 8e-4 for FP8 stability
+      min_lr: ${PRIMUS_MIN_LR:8.0e-5}  # Set to 10% of max LR
+      lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:128}
+      lr_decay_iters: ${PRIMUS_LR_DECAY_ITERS:1199872}
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      optimizer: adam
+      use_distributed_optimizer: true # use distributed optimizer
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-5
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+      layernorm_epsilon: 1.0e-05  # RMSNorm epsilon (matches HF rms_norm_eps)
+
+      # Dropout (disabled for training)
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # parallel
+      tensor_model_parallel_size: ${PRIMUS_TP:1}
+      pipeline_model_parallel_size: ${PRIMUS_PP:1}
+      expert_model_parallel_size: ${PRIMUS_EP:8}
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      ddp_num_buckets: 8
+      ddp_average_in_collective: true
+
+      # data
+      mock_data: false
+      num_workers: ${PRIMUS_NUM_WORKERS:0}
+      train_data_path: "10 /data/c4-train.en_6_text_document"
+      valid_data_path: "/data/c4-validation-91205-samples.en_text_document"
+      test_data_path: "/data/c4-validation-91205-samples.en_text_document"
+      # Avoid copying a dense (B, 1, S, S) CPU attention mask every step.
+      # TE receives causal/sliding-window metadata from attn_mask_type + window_size.
+      # Use numeric 0/1 because Primus env expansion only type-casts numbers.
+      create_attention_mask_in_dataloader: ${PRIMUS_CREATE_ATTENTION_MASK_IN_DATALOADER:0}
+
+      # fusion
+      moe_permute_fusion: true
+      gradient_accumulation_fusion: true
+      moe_use_legacy_grouped_gemm: false # Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedMLP, please set `moe_use_legacy_grouped_gemm=True
+      moe_use_fused_router_with_aux_score: true
+      multi_latent_attention: false # Flag config.ENABLE_EXPERIMENTAL not enabled
+      apply_rope_fusion: true
+
+
+      # sliding window attention (GPT-OSS-20B model definition; matches HF sliding_window: 128)
+      # use_turbo_attention is false so non-turbo attention (which supports sliding window) is used.
+      # Pattern: alternating sliding_attention (1) and full_attention (0) for 24 layers
+      # window_size must be a tuple (left_window, right_window) for Transformer Engine
+      # For causal attention: left = past tokens, right = 0 (no future tokens)
+      # HF sliding_window: 128 means 128 past tokens, so use (128, 0)
+      window_size: [128, 0]  # Left window: 128 past tokens, Right: 0 (causal)
+      window_attn_skip_freq: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]
+
+      # MoE settings
+      moe_apply_probs_on_input: false
+      moe_aux_loss_coeff: 0.0 #0.9
+      moe_deepep_num_sms: 20
+      moe_enable_deepep: false
+      moe_expert_capacity_factor: null
+      moe_extended_tp: false
+      moe_ffn_hidden_size: 2880
+      moe_flex_dispatcher_backend: deepep
+      moe_grouped_gemm: true
+      moe_hybridep_num_sms: 16
+      moe_input_jitter_eps: null
+      moe_latent_size: null
+      moe_layer_freq: 1
+      moe_layer_recompute: false
+      moe_pad_expert_input_to_capacity: false
+      moe_per_layer_logging: false
+      moe_router_bias_update_rate: 0.001
+      moe_router_dtype: fp32 # DeepEP only supports float32 probs
+      moe_router_enable_expert_bias: false
+      moe_router_force_load_balancing: false
+      moe_router_fusion: true
+      moe_router_group_topk: null
+      moe_router_load_balancing_type: none
+      moe_router_num_groups: null
+      moe_router_padding_for_fp8: false
+      moe_router_padding_for_quantization: false
+      moe_router_pre_softmax: false
+      moe_router_score_function: softmax
+      moe_router_topk: 4
+      moe_router_topk_limited_devices: null
+      moe_router_topk_scaling_factor: null
+      moe_shared_expert_gate: false
+      moe_shared_expert_intermediate_size: null
+      moe_shared_expert_overlap: false
+      moe_token_dispatcher_type: alltoall
+      moe_token_drop_policy: probs
+      moe_token_dropping: false
+      moe_z_loss_coeff: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 100000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      exit_on_missing_checkpoint: false
+      ckpt_format: torch
+      eval_iters: ${EVAL_ITERS:64}  # eval_samples = eval_iters * GBS = 1024; set EVAL_ITERS in config shell (1024/GBS).
+      eval_interval: ${PRIMUS_EVAL_INTERVAL:768}
+
+      # Turbo
+      enable_primus_turbo: true
+      use_turbo_attention: false
+      use_turbo_grouped_gemm: true
+      use_turbo_rms_norm: ${USE_TURBO_RMS_NORM:true}
+      use_turbo_fused_act_with_probs : true
+      # Pad tokens-per-expert so the fp8 grouped GEMM path skips the buggy
+      # quantization_padding branch in PrimusGroupedMLP.forward (experts.py:97-109),
+      # which yields NaN with recompute. Not auto-enabled here because
+      # turbo_sync_free_moe_stage=0 (it is only auto-set for sync-free stages 1-3).
+      use_turbo_permute_padding: true
+
+      # deepep
+      use_turbo_deepep: false
+
+      # 64 or 80 for ep8, 32 for ep16-64 is best practice
+      turbo_deepep_num_cu: 64
+      turbo_deepep_use_comm_stream: false
+
+      # sync-free moe support stage 0-3, 0 means not use sync-free moe
+      # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory
+      # stage 2 is recommended for better performance
+      turbo_sync_free_moe_stage: 0
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
+
+      # tensorboard logging, set 'disable_tensorboard: false' to enable tensorboard logging
+      disable_tensorboard: true
+      tensorboard_dir: /workspace/code/tensorboard
+      tensorboard_log_interval: 1
+      tensorboard_queue_size: 1000
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_validation_ppl_to_tensorboard: true
+      log_memory_to_tensorboard: true
+      log_world_size_to_tensorboard: true
+      log_loss_scale_to_tensorboard: true
diff --git a/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml
index 9315ccc55..8d614c817 100644
--- a/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml
@@ -46,9 +46,9 @@ modules:
       profile_step_start: 6
 
       # hyper parameters
-      train_iters: 10
-      micro_batch_size: 8
-      global_batch_size: 512
+      train_iters: 50
+      micro_batch_size: 6
+      global_batch_size: 48
       seq_length: ${PRIMUS_SEQ_LENGTH:4096}
       max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096}
       lr: 1.0e-5
@@ -106,6 +106,7 @@ modules:
       eval_iters: 0
 
       cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
 
       fp8: hybrid
       # enable_primus_turbo: true
diff --git a/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml
index ff4f67060..0be54d697 100755
--- a/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml
@@ -17,8 +17,8 @@ modules:
       log_avg_reset_interval: 50
 
       train_iters: 50
-      micro_batch_size: 14
-      global_batch_size: 224
+      micro_batch_size: 13
+      global_batch_size: 104
 
       seq_length: 4096
       max_position_embeddings: 4096
diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml
index 9975256f7..2c2fa0b68 100644
--- a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml
@@ -17,8 +17,8 @@ modules:
       log_avg_reset_interval: 50
 
       train_iters: 50
-      micro_batch_size: 4
-      global_batch_size: 32
+      micro_batch_size: 3
+      global_batch_size: 24
 
       seq_length: 8192
       max_position_embeddings: 8192
diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml
index f5072cb46..e5e920603 100644
--- a/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml
@@ -110,7 +110,6 @@ modules:
       # --- Primus Turbo Config ---
       enable_primus_turbo: true
       use_turbo_attention: true
-      use_turbo_fp4_autocast: false                 # TE mxfp4 recipe should set it to false
       use_turbo_gemm: false              # can't use together with delayed recipe
       use_turbo_grouped_gemm: false
       moe_use_fused_router_with_aux_score: false
diff --git a/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml b/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml
new file mode 100644
index 000000000..88d3242cd
--- /dev/null
+++ b/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml
@@ -0,0 +1,125 @@
+### DeepSeek-R1-Distill-Qwen-14B ODC dual-node SFT example.
+### Sequence-length load balancing (LB-Mini) with the "fit" cost model on the
+### torch-FSDP2 + ODC path; packed (thd) attention enabled.
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:qwen14B-odc-dn}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  post_trainer:
+    framework: megatron
+    config: sft_trainer.yaml
+
+    model: qwen2.5_14B.yaml
+
+    overrides:
+      stage: sft
+
+      hf_path: deepseek-ai/DeepSeek-R1-Distill-Qwen-14B
+      tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-14B
+
+      sft_dataset_name: "zai-org/LongAlign-10k"
+      sft_conversation_format: "messages"
+
+      # ODC on-demand communication (rocSHMEM/XGMI P2P gradient reduction) on the
+      # torch-FSDP2 path. Master switch (was the ODC_ENABLE env var).
+      enable_odc: true
+      odc_phase: 2
+
+      # thd segmented (packed) attention.
+      enable_packed_sequences: true
+      use_packed_attention: true
+      # LB-Mini sequence-length load balancing (was the ODC_LB_MINI env var).
+      enable_odc_lb_mini: true
+      # LB-Mini cost model: fit (attention-aware a*s^2 + b*s).
+      lb_mini_cost_model: fit
+
+      # ODC P2P backend + dual-node GDA (validated). trainer_base defaults are mori / false.
+      odc_p2p_backend: rocshmem
+      odc_rocshmem_gda: true
+
+      wandb_project: "Primus_ODC_Qwen14B"
+      stderr_sink_level: DEBUG
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 100
+
+      train_iters: 50
+      micro_batch_size: 1
+      global_batch_size: 16
+
+      seq_length: 65536
+      max_position_embeddings: 65536
+
+      lr: 2.0e-6
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: 100
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+
+      seed: 1234
+      eod_mask_loss: false
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-5
+      add_qkv_bias: true
+      apply_rope_fusion: false
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+      recompute_granularity: full
+      recompute_method: uniform
+      recompute_num_layers: 1
+
+      enable_fused_linear_ce: false
+
+      use_torch_fsdp2: true
+      use_megatron_fsdp: false
+      use_distributed_optimizer: false
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      gradient_accumulation_fusion: false
+
+      finetune: true
+      # Megatron checkpoint to warm-start from. Set PRETRAINED_CKPT to your own
+      # converted DeepSeek-R1-Distill-Qwen-14B checkpoint dir; null trains from
+      # the HF weights only.
+      pretrained_checkpoint: ${PRETRAINED_CKPT:null}
+      load: null
+      save: null
+      save_interval: 1000
+      eval_interval: 1000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch_dist
+
+      bf16: true
+
+      enable_primus_turbo: false
+      use_turbo_attention: false
+      use_turbo_grouped_mlp: false
+      use_turbo_rms_norm: false
+
+      eval_iters: 0
+
+      profile: false
+      use_pytorch_profiler: true
+      use_nsys_profiler: false
+      profile_step_start: 8
+      profile_step_end: 9
+      profile_ranks: [0]
+      pytorch_profiler_collect_shapes: true
+      pytorch_profiler_collect_callstack: false
+      pytorch_profiler_collect_chakra: false
+      record_shapes: false
+      record_memory_history: false
+      nvtx_ranges: false
+
+      lora:
+        enabled: false
diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-BF16-pretrain.yaml
index a9c45a614..1510d90f1 100644
--- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-BF16-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-BF16-pretrain.yaml
@@ -21,7 +21,7 @@ modules:
       log_avg_reset_interval: 50
 
       # hyper parameters
-      train_iters: 10
+      train_iters: 50
       micro_batch_size: 8
       global_batch_size: 512
       seq_length: ${PRIMUS_SEQ_LENGTH:4096}
diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml
index ef8965ae6..057fb0b6c 100644
--- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml
+++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml
@@ -21,9 +21,9 @@ modules:
       log_avg_reset_interval: 50
 
       # hyper parameters
-      train_iters: 10
-      micro_batch_size: 8
-      global_batch_size: 512
+      train_iters: 50
+      micro_batch_size: 6
+      global_batch_size: 48
       seq_length: ${PRIMUS_SEQ_LENGTH:4096}
       max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096}
       lr: 1.0e-5
diff --git a/examples/megatron/diffusion/README.md b/examples/megatron/diffusion/README.md
new file mode 100644
index 000000000..bfc6ada27
--- /dev/null
+++ b/examples/megatron/diffusion/README.md
@@ -0,0 +1,317 @@
+# Flux Diffusion Model Training Examples
+
+Training examples for Flux diffusion models with Primus-Megatron on AMD GPUs.
+
+## Related Documentation
+
+- **Architecture & Developer Guide:** [docs/04-technical-guides/diffusion-models/README.md](../../../docs/04-technical-guides/diffusion-models/README.md)
+- **API Reference:** [docs/04-technical-guides/diffusion-models/api_reference.md](../../../docs/04-technical-guides/diffusion-models/api_reference.md)
+- **FP8 Training Guide:** [docs/04-technical-guides/diffusion-models/fp8_training.md](../../../docs/04-technical-guides/diffusion-models/fp8_training.md)
+- **MXFP4 Training Guide:** [docs/04-technical-guides/diffusion-models/mxfp4_training.md](../../../docs/04-technical-guides/diffusion-models/mxfp4_training.md)
+- **Dataset Preparation:** [primus/configs/data/megatron/diffusion/README.md](../../../primus/configs/data/megatron/diffusion/README.md)
+- **Tests:** [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/)
+
+---
+
+## Quick Start
+
+### Prerequisites
+
+- AMD Instinct GPU(s) (MI300X, MI325X, MI355X)
+- Docker or Podman with ROCm support
+- Primus Docker image: `docker.io/rocm/primus:v26.1`
+- Prepared dataset (see [Dataset Preparation](../../../primus/configs/data/megatron/diffusion/README.md))
+
+### 5-Minute Test Run
+
+1. **Prepare a small test dataset:**
+
+```bash
+mkdir -p /tmp/flux_test_data/raw && cd /tmp/flux_test_data/raw
+
+for i in {000..099}; do
+    convert -size 512x512 xc:blue sample_${i}.jpg
+    echo "A blue square" > sample_${i}.txt
+done
+
+tar -cf train-000000.tar sample_*.jpg sample_*.txt
+
+cat > dataset.yaml << 'EOF'
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: raw
+EOF
+
+energon prepare . --num-workers 4
+```
+
+2. **Launch training:**
+
+```bash
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+DATA_PATH=/tmp/flux_test_data \
+GPUS_PER_NODE=1 \
+bash examples/run_pretrain.sh
+```
+
+---
+
+## Model Variants
+
+| Feature | Flux 535M | Flux 12B |
+|---------|-----------|----------|
+| Parameters | 535M | 12B |
+| Joint Layers | 1 | 19 |
+| Single Layers | 1 | 38 |
+| Min GPUs | 1 | 8 (FSDP2 / DDP) |
+| Recommended GPUs | 1-8 | 8-64 |
+| Sharding | DP | FSDP2 (ZeRO-2/3) or DDP + distributed optimizer |
+| Best For | Testing | Production |
+
+---
+
+## Available Configurations
+
+The same configs are provided for both MI300X (`examples/megatron/configs/MI300X/diffusion/`)
+and MI355X (`examples/megatron/configs/MI355X/diffusion/`). The only differences
+are hardware-tuned batch sizes (MI300X has 192GB HBM3, MI355X has 256GB), so MI300X
+uses smaller default micro/global batch sizes on the 12B DDP configs.
+
+### Shared (MI300X and MI355X)
+
+| Config | Description |
+|--------|-------------|
+| `flux_535m_pretrain.yaml` | Flux 535M, BF16, single/multi-GPU |
+| `flux_535m_pretrain_fp8.yaml` | Flux 535M with FP8 precision |
+| `flux_535m_with_guidance_embed.yaml` | Flux 535M with guidance embedding |
+| `flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml` | Flux 12B, FSDP2, BF16, local spec |
+| `flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml` | Flux 12B, FSDP2, FP8, local spec |
+| `flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml` | Flux 12B, DDP, FP8, local spec (delayed scaling) |
+| `flux_12b_ddp_energon_schnell_resample_te_spec.yaml` | Flux 12B, DDP, BF16, TransformerEngine spec |
+| `flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml` | Flux 12B, DDP, FP8, TransformerEngine spec |
+
+### MI355X only
+
+| Config | Description |
+|--------|-------------|
+| `flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml` | Flux 12B, DDP, MXFP4, local spec |
+| `flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml` | MLPerf benchmark reproduction (local spec FP8) |
+| `flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml` | MLPerf benchmark reproduction (TE spec FP8) |
+
+> The `*_mlperf.yaml` configs reproduce the MLPerf Training Flux.1 benchmark
+> (MLPerf logging + convergence target). Use the non-MLPerf configs above for
+> general training.
+
+---
+
+## Training Modes
+
+### Single-Node Training
+
+```bash
+# Flux 535M (1-8 GPUs)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+GPUS_PER_NODE=8 \
+bash examples/run_pretrain.sh
+
+# Flux 12B (FSDP2, BF16, 8 GPUs)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml \
+GPUS_PER_NODE=8 \
+bash examples/run_pretrain.sh
+```
+
+### Multi-Node Training (SLURM)
+
+```bash
+export DOCKER_IMAGE="docker.io/rocm/primus:v26.1"
+export NNODES=8
+export GPUS_PER_NODE=8
+
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml \
+bash examples/run_slurm_pretrain.sh
+```
+
+### MLPerf Benchmark Reproduction (MI355X)
+
+The `*_mlperf.yaml` configs reproduce the MLPerf Training Flux.1 benchmark and are
+intended for benchmark reproduction rather than general training.
+
+```bash
+# Step 1: Ingest MLPerf data
+primus-cli direct -- data diffusion-ingest \
+  --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
+
+# Step 2: Train (configs already set vae_latent_mode: resample, vae_scale: 0.3611, vae_shift: 0.1159)
+EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml \
+GPUS_PER_NODE=8 \
+bash examples/run_pretrain.sh
+```
+
+---
+
+## FP8 Training
+
+FP8 provides ~2x memory reduction and 1.5-2x training speedup on AMD MI300X/MI355X GPUs.
+
+```bash
+# Quick validation with 535M
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \
+GPUS_PER_NODE=1 \
+bash examples/run_pretrain.sh
+
+# Production with 12B (TransformerEngine FP8)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml \
+GPUS_PER_NODE=8 NNODES=4 \
+bash examples/run_slurm_pretrain.sh
+
+# Production with 12B (local-spec FP8, no TransformerEngine dependency)
+EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \
+GPUS_PER_NODE=8 NNODES=4 \
+bash examples/run_slurm_pretrain.sh
+```
+
+| Model | Precision | Memory/GPU | Batch Size | Speed |
+|-------|-----------|------------|------------|-------|
+| Flux 535M | BF16 | ~7-10GB | 2 | 1.0x |
+| Flux 535M | FP8 | ~3-5GB | 4 | 1.5-2x |
+| Flux 12B | BF16 | ~40-50GB | 1 | 1.0x |
+| Flux 12B | FP8 | ~20-25GB | 2 | 1.5-2x |
+
+For configuration details, tuning recipes, benchmarks, and troubleshooting, see the [FP8 Training Guide](../../../docs/04-technical-guides/diffusion-models/fp8_training.md).
+
+---
+
+## MXFP4 Training
+
+MXFP4 (E2M1 + E8M0 block-of-32 scales) Flux 12B training on MI355X is supported via the local-spec provider (`PrimusTurboMXFP4LocalSpecProvider`, no TransformerEngine dependency). Forward and weight GEMMs run in FP4 through Primus-Turbo + AITER; attention, the optimizer state, and inter-rank communication stay in BF16.
+
+```bash
+# Path to a checkout of the `tuned_gemm_configs` directory.
+# Set TUNED_GEMM_DIR to wherever you have the tuned configs available.
+export TUNED_GEMM_DIR=${TUNED_GEMM_DIR:-/path/to/tuned_gemm_configs}
+
+EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml \
+PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER \
+AITER_CONFIG_GEMM_A4W4=$TUNED_GEMM_DIR/mi355x/flux_12b.csv \
+AITER_LOG_TUNED_CONFIG=1 \
+bash examples/run_pretrain.sh
+```
+
+For configuration knobs, backend-selector semantics, tuned-GEMM verification, and troubleshooting, see the [MXFP4 Training Guide](../../../docs/04-technical-guides/diffusion-models/mxfp4_training.md).
+
+---
+
+## Converting HuggingFace Checkpoints
+
+Convert pre-trained HuggingFace Flux checkpoints to Primus/Megatron-Core format:
+
+```bash
+python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \
+    --input black-forest-labs/FLUX.1-dev \
+    --output checkpoints/primus_flux_12b.safetensors \
+    --variant flux_12b
+```
+
+Supported variants: `flux_535m`, `flux_12b`, `custom` (with `--num-joint-layers` / `--num-single-layers`).
+
+For gated models (FLUX.1-dev), set `export HF_TOKEN="your_token"` or run `huggingface-cli login`.
+
+Primus also auto-detects tokens from `.hf_token` at the project root or `~/.cache/huggingface/token`.
+
+---
+
+## Configuration Reference
+
+### Key Training Parameters
+
+```yaml
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_535m.yaml
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      train_iters: 100000
+      micro_batch_size: 2
+      global_batch_size: 16
+      lr: 1.0e-4
+      min_lr: 1.0e-5
+      weight_decay: 0.01
+      clip_grad: 1.0
+      lr_decay_style: cosine
+      lr_warmup_iters: 1000
+```
+
+### Parallelism
+
+The Flux 12B configs scale with FSDP2 (ZeRO-2/3) or DDP + distributed optimizer
+rather than tensor/pipeline parallelism.
+
+| Setting | Flux 535M | Flux 12B |
+|---------|-----------|----------|
+| `tensor_model_parallel_size` | 1 | 1 |
+| `pipeline_model_parallel_size` | 1 | 1 |
+| `context_parallel_size` | 1 | 1 |
+| Sharding | DP | FSDP2 (`use_torch_fsdp2: true`) or DDP (`use_distributed_optimizer: true`) |
+
+### Memory Optimization
+
+```yaml
+modules:
+  pre_trainer:
+    overrides:
+      recompute_granularity: selective  # or 'full'
+      recompute_method: block
+      sequence_parallel: false
+```
+
+---
+
+## Troubleshooting
+
+### Dataset Not Found
+
+Verify `data_path` in your config, ensure `energon prepare` was run, and check that `dataset.yaml` exists.
+
+### Out of Memory (OOM)
+
+Reduce `micro_batch_size`, increase `tensor_model_parallel_size`, enable `recompute_granularity: full`, or switch to pre-encoded data mode.
+
+### NaN Loss
+
+Reduce learning rate to `1.0e-5`, ensure `clip_grad: 1.0`, increase `lr_warmup_iters`, check dataset for corruption.
+
+### NaN Loss with MLPerf Data
+
+Ensure config includes the required normalization constants:
+
+```yaml
+vae_latent_mode: resample
+vae_scale: 0.3611
+vae_shift: 0.1159
+```
+
+### Encoder Download Fails
+
+Set `export HF_TOKEN=your_token` or use a local model path via `encoder_model_path` in config overrides.
+
+### Slow Training
+
+Use pre-encoded data (2-3x faster), increase `num_workers`, enable `use_flash_attn: true`.
+
+---
+
+## Source Code Pointers
+
+- **Model Architecture:** `primus/backends/megatron/core/models/diffusion/flux/`
+- **Trainer:** `primus/modules/trainer/megatron/diffusion/flux_pretrain_trainer.py`
+- **Data Pipeline:** `primus/backends/megatron/data/diffusion/`
+- **Model Configs:** `primus/configs/models/megatron/diffusion/`
+
+## Getting Help
+
+- [GitHub Issues](https://github.com/AMD-AGI/Primus/issues)
+- [GitHub Discussions](https://github.com/AMD-AGI/Primus/discussions)
diff --git a/examples/megatron/prepare.py b/examples/megatron/prepare.py
index bdad391c3..86725f311 100644
--- a/examples/megatron/prepare.py
+++ b/examples/megatron/prepare.py
@@ -282,6 +282,15 @@ def build_megatron_helper(primus_path: Path, patch_args: Path, backend_path: str
             f"if you actually need it). SFT does not require this dependency."
         )
         return
+    # This prepare step may run once PER RANK (e.g. 8-16 concurrent processes on a
+    # node under torchrun), and concurrent `pip install -e` writes race on the shared
+    # venv .pth file (OSError). If the package is already importable (one-time
+    # pre-install), take a fast, write-free path so all ranks agree without racing.
+    import importlib.util
+
+    if importlib.util.find_spec("emerging_optimizers") is not None:
+        log_info("Emerging-Optimizers already installed; skipping editable reinstall.")
+        return
     log_info(f"Building Emerging Optimizers in {emerging_optimizers_path}")
     ret = subprocess.run(
         ["pip", "install", "--no-build-isolation", "-e", str(emerging_optimizers_path)], check=True
diff --git a/examples/megatron/prepare_fineweb_edu.py b/examples/megatron/prepare_fineweb_edu.py
new file mode 100644
index 000000000..54afdd185
--- /dev/null
+++ b/examples/megatron/prepare_fineweb_edu.py
@@ -0,0 +1,180 @@
+###############################################################################
+# End-to-End FineWeb-Edu Preparation Script
+###############################################################################
+
+import argparse
+import os
+import subprocess
+import time
+from pathlib import Path
+
+
+def prepare_fineweb_edu_dataset(
+    primus_path: Path,
+    data_path: Path,
+    tokenizer_type: str,
+    tokenizer_model: str,
+    sample_size: str = "10BT",
+    workers: int = None,
+    overwrite: bool = False,
+):
+    """Prepare FineWeb-Edu dataset for Megatron training."""
+
+    dataset_name = f"fineweb-edu-{sample_size}"
+    dataset_path = data_path / dataset_name
+    output_path = dataset_path / tokenizer_type
+
+    # Set HuggingFace cache
+    hf_home = Path(os.environ.get("HF_HOME", data_path / "huggingface"))
+    os.environ["HF_HOME"] = str(hf_home)
+
+    # Output tokenized files
+    tokenized_prefix = output_path / f"fineweb_edu_{sample_size}"
+    tokenized_bin = tokenized_prefix.with_name(f"{tokenized_prefix.name}_text_sentence.bin")
+    tokenized_idx = tokenized_prefix.with_name(f"{tokenized_prefix.name}_text_sentence.idx")
+
+    # Check if already processed
+    if tokenized_bin.exists() and tokenized_idx.exists():
+        if overwrite:
+            print(f"[Info] Overwriting existing tokenized files.")
+            tokenized_bin.unlink()
+            tokenized_idx.unlink()
+        else:
+            print(f"[Info] Tokenized files exist, skipping preprocessing.")
+            print(f"  - {tokenized_bin}")
+            print(f"  - {tokenized_idx}")
+            print(f"[Hint] Use --overwrite to force re-tokenization with a different tokenizer.")
+            return tokenized_prefix.with_name(f"{tokenized_prefix.name}_text_sentence")
+
+    output_path.mkdir(parents=True, exist_ok=True)
+    dataset_json = dataset_path / f"fineweb_edu_{sample_size}_megatron.json"
+
+    # Step 1: Download dataset if not exists
+    if dataset_json.exists():
+        print(f"[Info] Found dataset file: {dataset_json}, skipping download.")
+    else:
+        print(f"[Info] Downloading FineWeb-Edu ({sample_size}) dataset...")
+        subprocess.run(
+            [
+                "python3",
+                str(primus_path / "examples/megatron/prepare_fineweb_edu_megatron_dataset.py"),
+                "--out-dir",
+                str(dataset_path),
+                "--sample-size",
+                sample_size,
+            ],
+            check=True,
+        )
+        print("[Info] Download completed.")
+
+    # Step 2: Tokenize and create binary files
+    print(f"[Info] Preprocessing dataset with tokenizer {tokenizer_type} / {tokenizer_model}")
+    start = time.time()
+
+    if workers is None:
+        workers = os.cpu_count()
+
+    env = os.environ.copy()
+    megatron_path = str(primus_path / "third_party" / "Megatron-LM")
+    env["PYTHONPATH"] = f"{primus_path}:{megatron_path}:{env.get('PYTHONPATH', '')}"
+
+    subprocess.run(
+        [
+            "python3",
+            str(primus_path / "examples/megatron/preprocess_data.py"),
+            "--input",
+            str(dataset_json),
+            "--tokenizer-type",
+            tokenizer_type,
+            "--tokenizer-model",
+            tokenizer_model,
+            "--output-prefix",
+            str(tokenized_prefix),
+            "--workers",
+            str(workers),
+            "--split-sentences",
+            "--append-eod",
+            "--partitions",
+            "4",  # Increase for larger datasets
+            "--log-interval",
+            "10000",
+        ],
+        env=env,
+        check=True,
+    )
+
+    elapsed = int(time.time() - start)
+    print(f"[Info] Preprocessing completed in {elapsed} seconds ({elapsed/60:.1f} minutes)")
+
+    return tokenized_prefix.with_name(f"{tokenized_prefix.name}_text_sentence")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare FineWeb-Edu dataset for Megatron-LM training")
+    parser.add_argument("--primus-path", type=str, required=True, help="Root path to the Primus project")
+    parser.add_argument("--data-path", type=str, required=True, help="Path to data directory")
+    parser.add_argument(
+        "--tokenizer-type",
+        type=str,
+        default="HuggingFaceTokenizer",
+        help="Tokenizer type (HuggingFaceTokenizer, GPT2BPETokenizer, etc.)",
+    )
+    parser.add_argument(
+        "--tokenizer-model", type=str, default="meta-llama/Llama-3.2-1B", help="Tokenizer model name or path"
+    )
+    parser.add_argument(
+        "--sample-size",
+        type=str,
+        default="10BT",
+        choices=["10BT", "100BT", "350BT", "sample-10BT", "sample-100BT", "sample-350BT"],
+        help="FineWeb-Edu dataset size",
+    )
+    parser.add_argument(
+        "--workers", type=int, default=None, help="Number of workers for preprocessing (default: CPU count)"
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite existing tokenized files (use when switching tokenizers)",
+    )
+
+    args = parser.parse_args()
+
+    primus_path = Path(args.primus_path).resolve()
+    data_path = Path(args.data_path).resolve()
+
+    print("=" * 70)
+    print("FineWeb-Edu Dataset Preparation for Megatron-LM")
+    print("=" * 70)
+    print(f"Primus Path:    {primus_path}")
+    print(f"Data Path:      {data_path}")
+    print(f"Tokenizer:      {args.tokenizer_type} / {args.tokenizer_model}")
+    print(f"Sample Size:    {args.sample_size}")
+    print(f"Workers:        {args.workers or os.cpu_count()}")
+    print("=" * 70)
+
+    # Check HF_TOKEN
+    if not os.environ.get("HF_TOKEN"):
+        print("[Warning] HF_TOKEN not set. You may need it for gated datasets.")
+
+    output_prefix = prepare_fineweb_edu_dataset(
+        primus_path=primus_path,
+        data_path=data_path,
+        tokenizer_type=args.tokenizer_type,
+        tokenizer_model=args.tokenizer_model,
+        sample_size=args.sample_size,
+        workers=args.workers,
+        overwrite=args.overwrite,
+    )
+
+    print("\n" + "=" * 70)
+    print("SUCCESS! Dataset is ready for training.")
+    print("=" * 70)
+    print(f"\nAdd this to your training config:\n")
+    print(f"  train_data_path: {output_prefix}")
+    print(f"  mock_data: false")
+    print("\n" + "=" * 70)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/megatron/prepare_fineweb_edu.sh b/examples/megatron/prepare_fineweb_edu.sh
new file mode 100644
index 000000000..8070b0c4d
--- /dev/null
+++ b/examples/megatron/prepare_fineweb_edu.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+# Set Python path
+megatron_path="$(pwd)/third_party/Megatron-LM"
+export PYTHONPATH="${megatron_path}:${PYTHONPATH}"
+
+# Verify the import works
+python3 -c "from megatron.core.datasets import indexed_dataset; print('Import successful')"
+
+# Then run your preparation script
+python examples/megatron/prepare_fineweb_edu.py \
+    --primus-path . \
+    --data-path ./data \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model meta-llama/Llama-3.2-1B \
+    --sample-size 10BT
diff --git a/examples/megatron/prepare_fineweb_edu_megatron_dataset.py b/examples/megatron/prepare_fineweb_edu_megatron_dataset.py
new file mode 100644
index 000000000..d7a591716
--- /dev/null
+++ b/examples/megatron/prepare_fineweb_edu_megatron_dataset.py
@@ -0,0 +1,65 @@
+###############################################################################
+# Prepare FineWeb-Edu Dataset for Megatron-LM
+###############################################################################
+
+import argparse
+from pathlib import Path
+
+from datasets import load_dataset
+
+
+def prepare_fineweb_edu_dataset(out_dir: Path, sample_size: str = "10BT"):
+    """
+    Download FineWeb-Edu dataset and convert to JSON format.
+
+    Args:
+        out_dir: Output directory for JSON file
+        sample_size: Size of dataset to download
+            - "10BT" (10B tokens, ~13GB) - Recommended for testing
+            - "100BT" (100B tokens, ~130GB)
+            - "350BT" (350B tokens, ~450GB)
+            - "sample-10BT", "sample-100BT", "sample-350BT" for samples
+    """
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # FineWeb-Edu dataset name format
+    dataset_name = f"HuggingFaceFW/fineweb-edu"
+    config_name = f"sample-{sample_size}" if not sample_size.startswith("sample-") else sample_size
+
+    print(f"[Info] Loading fineweb-edu dataset ({sample_size}) from Hugging Face...")
+    print(f"[Info] This may take a while depending on dataset size...")
+
+    # Load dataset - you can specify num_proc for faster loading
+    dataset = load_dataset(
+        dataset_name,
+        name=config_name,
+        split="train",
+        trust_remote_code=True,
+        # streaming=True,  # Uncomment for very large datasets
+    )
+
+    output_file = out_dir / f"fineweb_edu_{sample_size}_megatron.json"
+    print(f"[Info] Saving dataset to {output_file} ...")
+
+    # Convert to JSON format that Megatron expects
+    dataset.to_json(str(output_file))
+
+    print(f"[Info] Dataset preparation completed: {output_file}")
+    print(f"[Info] Total samples: {len(dataset)}")
+
+    return output_file
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download and prepare FineWeb-Edu dataset for Megatron")
+    parser.add_argument("--out-dir", type=str, default="./data/fineweb-edu", help="Path to output directory")
+    parser.add_argument(
+        "--sample-size",
+        type=str,
+        default="10BT",
+        choices=["10BT", "100BT", "350BT", "sample-10BT", "sample-100BT", "sample-350BT"],
+        help="Size of FineWeb-Edu dataset to download",
+    )
+    args = parser.parse_args()
+
+    prepare_fineweb_edu_dataset(Path(args.out_dir), args.sample_size)
diff --git a/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml b/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml
new file mode 100644
index 000000000..5ba988890
--- /dev/null
+++ b/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml
@@ -0,0 +1,54 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:mamba_130M_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron_bridge
+    config: pretrain_trainer.yaml
+
+    # Model to run
+    model: mamba_130M.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Recipe flavor: upstream Megatron-Bridge Mamba2 130M pretrain config
+      flavor: mamba2_130m_pretrain_config
+
+      # --- Parameters accepted by _mamba2_common() ---
+
+      # Training configuration
+      train_iters: 50
+      global_batch_size: 32
+      micro_batch_size: 4
+      seq_length: ${PRIMUS_SEQ_LENGTH:2048}
+
+      # Nested overrides (applied by _apply_nested_overrides)
+      log_interval: 1
+      eval_interval: 500
+      eval_iters: 0
+      skip_save: true
+
+      # Optimizer
+      lr: 3.0e-4
+      min_lr: 3.0e-5
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # Data
+      mock: true
+      # Real data path should point to a Megatron indexed dataset prefix:
+      #   <prefix>.bin + <prefix>.idx
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
diff --git a/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml
index 330c9e382..b1493908c 100644
--- a/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml
+++ b/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml
@@ -31,7 +31,7 @@ modules:
       # Training configuration
       train_iters: 200
       global_batch_size: 32
-      micro_batch_size: 2
+      micro_batch_size: 1
       seq_length: 8192
       eval_interval: 30
       save_interval: 50
diff --git a/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml
index 75c721ed2..f05718454 100644
--- a/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml
+++ b/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml
@@ -31,7 +31,7 @@ modules:
       # Training configuration
       train_iters: 200
       global_batch_size: 8
-      micro_batch_size: 2
+      micro_batch_size: 1
       seq_length: 8192
       eval_interval: 30
       save_interval: 50
diff --git a/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml b/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml
new file mode 100644
index 000000000..5ba988890
--- /dev/null
+++ b/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml
@@ -0,0 +1,54 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:mamba_130M_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron_bridge
+    config: pretrain_trainer.yaml
+
+    # Model to run
+    model: mamba_130M.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Recipe flavor: upstream Megatron-Bridge Mamba2 130M pretrain config
+      flavor: mamba2_130m_pretrain_config
+
+      # --- Parameters accepted by _mamba2_common() ---
+
+      # Training configuration
+      train_iters: 50
+      global_batch_size: 32
+      micro_batch_size: 4
+      seq_length: ${PRIMUS_SEQ_LENGTH:2048}
+
+      # Nested overrides (applied by _apply_nested_overrides)
+      log_interval: 1
+      eval_interval: 500
+      eval_iters: 0
+      skip_save: true
+
+      # Optimizer
+      lr: 3.0e-4
+      min_lr: 3.0e-5
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # Data
+      mock: true
+      # Real data path should point to a Megatron indexed dataset prefix:
+      #   <prefix>.bin + <prefix>.idx
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
diff --git a/examples/mlperf/gpt_oss_20b/.dockerignore b/examples/mlperf/gpt_oss_20b/.dockerignore
new file mode 100644
index 000000000..e0def13df
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/.dockerignore
@@ -0,0 +1,5 @@
+**
+!Dockerfile.runtime-v26.3
+!Dockerfile.runtime-v26.5
+!prewarm_attention.py
+!aiter_hd64_asm_override.py
diff --git a/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3 b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3
new file mode 100644
index 000000000..5400df304
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3
@@ -0,0 +1,88 @@
+# 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=e160182d6a7967592579be2640e85106bc8f421e
+ARG TE_REF=5235bae2cc683a0ad4bf15221c746ab3c1e229e7
+ARG AITER_REF=c4b33df03faae1c4e470420d950f8a9589e9634d
+ARG FWD_ATTN_ASM_REF=53d3dadc3f3b0ac35ae536f2d1d7864a3e07ba22
+ARG BWD_ATTN_ASM_REF=9b9fb6444f3fee388617f62432c3faea74079377
+ARG BWD_ATTN_SYMBOL=_ZN5aiter43fmha_bwd_hd64_bf16_causal_a16_rtz_recompileE
+ARG BWD_ATTN_SLOT=bwd_hd64_bf16_causal_a16_rtz.co
+ARG MAX_JOBS=96
+
+ENV MLPERF_RUNTIME_SERIES=v26.3 \
+    MLPERF_ENABLE_FWD_ATTN_ASM=1 \
+    FMHA_HD64_ASM_CO=/opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co \
+    FMHA_HD64_ASM_LOG=0
+
+WORKDIR /workspace/deps
+
+# The v26.3 base ships Triton 3.6, while this Turbo branch requires the 3.7 API.
+RUN python3 -m pip install --upgrade --force-reinstall --no-deps flydsl==0.2.4 && \
+    git clone https://github.com/triton-lang/triton.git && \
+    cd triton && git checkout "${TRITON_REF}" && \
+    python3 -m pip install ninja cmake && \
+    MAX_JOBS="${MAX_JOBS}" python3 -m pip install \
+      --no-build-isolation --no-deps -v . && \
+    cd .. && rm -rf triton
+
+RUN git clone --recursive https://github.com/AMD-AGI/Primus-Turbo.git && \
+    cd Primus-Turbo && git checkout "${PRIMUS_TURBO_REF}" && \
+    git submodule update --init --recursive && \
+    python3 -m pip install -r requirements.txt && \
+    PRIMUS_TURBO_FRAMEWORK=PYTORCH \
+    GPU_ARCHS=gfx950 \
+    MAX_JOBS="${MAX_JOBS}" \
+      python3 -m pip install --no-build-isolation --no-deps -v . && \
+    cd .. && rm -rf Primus-Turbo
+
+RUN git clone https://github.com/mawad-amd/fwd-attn-asm.git && \
+    cd fwd-attn-asm && git checkout "${FWD_ATTN_ASM_REF}" && \
+    amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \
+      -o /tmp/fwd_attn.co kernels/fwd_d64_opt128.s && \
+    cd .. && \
+    git clone https://github.com/mawad-amd/bwd-attn-asm.git && \
+    cd bwd-attn-asm && git checkout "${BWD_ATTN_ASM_REF}" && \
+    amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \
+      -o /tmp/bwd_attn.co kernels/bwd_d64_v3_causal_opt_16x32.s && \
+    cd .. && rm -rf fwd-attn-asm bwd-attn-asm
+
+RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git && \
+    cd TransformerEngine && git checkout "${TE_REF}" && \
+    git submodule update --init --recursive && \
+    AITER_DIR=3rdparty/QoLA/3rdparty/aiter && \
+    git -C "${AITER_DIR}" fetch origin "${AITER_REF}" && \
+    git -C "${AITER_DIR}" checkout "${AITER_REF}" && \
+    git -C "${AITER_DIR}" submodule update --init --recursive && \
+    sed -i "s|^aiter_commit = .*|aiter_commit = \"${AITER_REF}\"|" \
+      transformer_engine/common/ck_fused_attn/qola_manifest.toml && \
+    install -m 0644 /tmp/bwd_attn.co \
+      "${AITER_DIR}/hsa/gfx950/fmha_v3_bwd/${BWD_ATTN_SLOT}" && \
+    python3 -m pip uninstall -y \
+      transformer_engine transformer_engine_rocm7 transformer_engine_rocm_torch && \
+    NVTE_FUSED_ATTN_AOTRITON=0 \
+    NVTE_CK_FUSED_ATTN_PATH="" \
+    NVTE_BUILD_MAX_JOBS="${MAX_JOBS}" \
+    NVTE_FRAMEWORK=pytorch \
+    NVTE_ROCM_ARCH=gfx950 \
+    NVTE_USE_HIPBLASLT=1 \
+    PYTORCH_ROCM_ARCH=gfx950 \
+    CU_NUM=304 \
+    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \
+    MAX_JOBS="${MAX_JOBS}" \
+      python3 -m pip install --no-build-isolation --no-deps -v . && \
+    cd .. && rm -rf TransformerEngine
+
+COPY aiter_hd64_asm_override.py \
+     /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.py
+COPY prewarm_attention.py /opt/mlperf-gpt-oss-20b/prewarm_attention.py
+RUN printf 'import aiter_hd64_asm_override\n' \
+      > /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.pth && \
+    install -m 0644 /tmp/fwd_attn.co \
+      /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co && \
+    rm -f /tmp/fwd_attn.co /tmp/bwd_attn.co
+
+WORKDIR /workspace
diff --git a/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5 b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5
new file mode 100644
index 000000000..913dc14e0
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5
@@ -0,0 +1,89 @@
+# syntax=docker/dockerfile:1.7
+
+ARG BASE_IMAGE=rocm/primus:v26.5@sha256:3040bf42974d791dd42de2e36b3c919a00869a5754cfc57a06b96d004c55eed1
+FROM ${BASE_IMAGE}
+
+ARG PRIMUS_TURBO_REF=e160182d6a7967592579be2640e85106bc8f421e
+ARG TE_REF=a07e607f14a5330807ffdafeeb6224f2d7dffacc
+ARG AITER_REF=d32b0cb62ecc32bb1a858e8437d58eb9b3856af6
+ARG FWD_ATTN_ASM_REF=53d3dadc3f3b0ac35ae536f2d1d7864a3e07ba22
+ARG BWD_ATTN_ASM_REF=9b9fb6444f3fee388617f62432c3faea74079377
+ARG BWD_ATTN_SOURCE_SYMBOL=_ZN5aiter43fmha_bwd_hd64_bf16_causal_a16_rtz_recompileE
+ARG BWD_ATTN_SYMBOL=_ZN5aiter44fmha_bwd_hd64_bf16_causal_a16_rtne_recompileE
+ARG BWD_ATTN_SLOT=bwd_hd64_bf16_causal_a16_rtne.co
+ARG MAX_JOBS=96
+
+ENV MLPERF_RUNTIME_SERIES=v26.5 \
+    MLPERF_ENABLE_FWD_ATTN_ASM=1 \
+    FMHA_HD64_ASM_CO=/opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co \
+    FMHA_HD64_ASM_LOG=0
+
+WORKDIR /workspace/deps
+
+# The v26.5 base already contains the tested Triton 3.7 compiler.
+RUN git clone --recursive https://github.com/AMD-AGI/Primus-Turbo.git && \
+    cd Primus-Turbo && git checkout "${PRIMUS_TURBO_REF}" && \
+    git submodule update --init --recursive && \
+    python3 -m pip install -r requirements.txt && \
+    PRIMUS_TURBO_FRAMEWORK=PYTORCH \
+    GPU_ARCHS=gfx950 \
+    MAX_JOBS="${MAX_JOBS}" \
+      python3 -m pip install --no-build-isolation --no-deps -v . && \
+    cd .. && rm -rf Primus-Turbo
+
+# Compile the hand-tuned forward and numerically validated backward payloads.
+RUN git clone https://github.com/mawad-amd/fwd-attn-asm.git && \
+    cd fwd-attn-asm && git checkout "${FWD_ATTN_ASM_REF}" && \
+    amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \
+      -o /tmp/fwd_attn.co kernels/fwd_d64_opt128.s && \
+    cd .. && \
+    git clone https://github.com/mawad-amd/bwd-attn-asm.git && \
+    cd bwd-attn-asm && git checkout "${BWD_ATTN_ASM_REF}" && \
+    sed "s/${BWD_ATTN_SOURCE_SYMBOL}/${BWD_ATTN_SYMBOL}/g" \
+      kernels/bwd_d64_v3_causal_opt_16x32.s > /tmp/bwd_attn.s && \
+    amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \
+      -o /tmp/bwd_attn.co /tmp/bwd_attn.s && \
+    cd .. && rm -rf fwd-attn-asm bwd-attn-asm /tmp/bwd_attn.s
+
+# The base image already carries the matching TE/AITER Python stack. Rebuild
+# TE only because the validated backward ASM must be embedded in libmha_bwd.
+RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git && \
+    cd TransformerEngine && git checkout "${TE_REF}" && \
+    git submodule update --init --recursive && \
+    sed -i "s|^aiter_commit = .*|aiter_commit = \"${AITER_REF}\"|" \
+      transformer_engine/common/ck_fused_attn/qola_manifest.toml && \
+    PYTHONPATH="$PWD/3rdparty/QoLA:${PYTHONPATH}" \
+      python3 -m qola.cli checkout \
+        --manifest transformer_engine/common/ck_fused_attn/qola_manifest.toml \
+        --aiter-root /workspace/deps/te-aiter && \
+    install -m 0644 /tmp/bwd_attn.co \
+      "/workspace/deps/te-aiter/hsa/gfx950/fmha_v3_bwd/${BWD_ATTN_SLOT}" && \
+    python3 -m pip uninstall -y \
+      transformer_engine transformer_engine_rocm7 transformer_engine_rocm_torch && \
+    NVTE_FUSED_ATTN_AOTRITON=0 \
+    NVTE_FUSED_ATTN_CK=1 \
+    NVTE_CK_JIT=1 \
+    NVTE_CK_FUSED_ATTN_PATH="" \
+    NVTE_AITER_SOURCE_DIR=/workspace/deps/te-aiter \
+    NVTE_BUILD_MAX_JOBS="${MAX_JOBS}" \
+    NVTE_FRAMEWORK=pytorch \
+    NVTE_ROCM_ARCH=gfx950 \
+    NVTE_USE_HIPBLASLT=1 \
+    PYTORCH_ROCM_ARCH=gfx950 \
+    CU_NUM=304 \
+    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \
+    CMAKE_BUILD_PARALLEL_LEVEL="${MAX_JOBS}" \
+    MAX_JOBS="${MAX_JOBS}" \
+      python3 -m pip install --no-build-isolation --no-deps -v . && \
+    cd .. && rm -rf TransformerEngine te-aiter
+
+COPY aiter_hd64_asm_override.py \
+     /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.py
+COPY prewarm_attention.py /opt/mlperf-gpt-oss-20b/prewarm_attention.py
+RUN printf 'import aiter_hd64_asm_override\n' \
+      > /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.pth && \
+    install -m 0644 /tmp/fwd_attn.co \
+      /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co && \
+    rm -f /tmp/fwd_attn.co /tmp/bwd_attn.co
+
+WORKDIR /workspace
diff --git a/examples/mlperf/gpt_oss_20b/README.md b/examples/mlperf/gpt_oss_20b/README.md
new file mode 100644
index 000000000..0d979ce69
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/README.md
@@ -0,0 +1,77 @@
+# GPT-OSS-20B MLPerf pretraining
+
+GPT-OSS 20B on one MI355X node with 8 GPUs and global batch size 32.
+
+## Build
+
+The Dockerfile builds the complete runtime, including the GPT-OSS
+Primus-Turbo test branch, TransformerEngine, and the attention ASM kernels.
+The v26.5 image reuses the base image's Triton 3.7 compiler; v26.3 upgrades its
+older base Triton for compatibility with the same Turbo branch.
+
+```bash
+cd examples/mlperf/gpt_oss_20b
+docker build --network host \
+  -f Dockerfile.runtime-v26.5 \
+  -t primus:gpt-oss-20b-mlperf-v26.5 .
+```
+
+Use `Dockerfile.runtime-v26.3` and a v26.3 tag for the compatibility stack.
+Push a shared tag with `docker push <image>`.
+
+## Data
+
+```bash
+mkdir -p /data/gpt_oss_20b
+cd /data/gpt_oss_20b
+bash <(curl -s https://raw.githubusercontent.com/mlcommons/r2-downloader/refs/heads/main/mlc-r2-downloader.sh) \
+  -d data \
+  https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri
+```
+
+Training uses the `c4-train.en_6_text_document` prefix and validation uses
+`c4-validation-91205-samples.en_text_document`.
+
+## Run
+
+```bash
+docker run -it --rm \
+  --privileged --network host --ipc host --shm-size 128g \
+  --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
+  --device /dev/dri --device /dev/kfd --device /dev/infiniband \
+  -v /path/to/Primus:/workspace/Primus \
+  -v /path/to/data:/data \
+  -v /path/to/model:/model \
+  -v /path/to/results:/results \
+  primus:gpt-oss-20b-mlperf-v26.5 bash
+```
+
+Inside the container:
+
+```bash
+cd /workspace/Primus/examples/mlperf/gpt_oss_20b
+source config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh
+./run_and_time.sh
+```
+
+That is the complete long-run entry. The config is the single source of
+submission defaults: MLPerf trainer, 1.2M iteration ceiling, 128-step warmup,
+FP8 Triton grouped GEMM, fused wgrad accumulation, and disabled profiling.
+Short diagnostics and backend ablations should override environment variables
+outside the checked-in submission config.
+
+## v26.5 attention prewarm
+
+The v26.5 TE stack lazily compiles two attention variants. Starting eight
+torchrun ranks against an empty cache can race while writing the same blobs.
+`run_and_time.sh` therefore runs `prewarm_attention.py` once before timing; the
+helper only populates the sliding-window and full-attention cache entries.
+The v26.3 TE/AITER stack does not exhibit this cache race, so the prewarm is
+skipped automatically for v26.3.
+
+## Key files
+
+- `Dockerfile.runtime-v26.3`, `Dockerfile.runtime-v26.5`: complete runtime builds
+- `config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh`: submission defaults
+- `run_and_time.sh`: benchmark entry
+- `prewarm_attention.py`: v26.5-only attention cache prewarm
diff --git a/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py b/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py
new file mode 100644
index 000000000..2ab198eb3
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py
@@ -0,0 +1,802 @@
+"""Route eligible TransformerEngine FMHA forward calls to the pinned gfx950
+HD64 BF16 assembly kernel.
+
+The module is imported at Python startup through ``aiter_hd64_asm_override.pth``
+and remains inactive unless ``MLPERF_ENABLE_FWD_ATTN_ASM=1``.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import inspect
+import logging
+import math
+import os
+import struct
+import sys
+from typing import Tuple
+
+logger = logging.getLogger("fwd_attn_asm_override")
+# Per-dispatch logging is opt-in only: it fires once per attention call, so an
+# inherited INFO root level (verbose runs) must not switch it on.
+_DISPATCH_LOG = os.environ.get("FMHA_HD64_ASM_LOG", "0") == "1"
+if _DISPATCH_LOG:
+    logging.basicConfig(level=logging.INFO)
+    logger.setLevel(logging.INFO)
+
+
+_ENABLED = os.environ.get("MLPERF_ENABLE_FWD_ATTN_ASM", "0") == "1"
+_AITER_ROPE_ENABLED = os.environ.get("NVTE_USE_AITER_ROPE", "0") == "1"
+_DEFAULT_CO_PATH = os.path.join(
+    os.path.dirname(__file__),
+    "aiter_hd64_asm_fwd_d64_opt128.co",
+)
+_CO_PATH = os.environ.get("FMHA_HD64_ASM_CO", _DEFAULT_CO_PATH)
+_KERNEL_NAME = b"fmha_fwd_d64_bf16_causal"
+
+# Tile shape baked into the kernel (BlockFmhaPipelineQRKSVSAsync<128,64,...>).
+_BLOCK_M = 128
+_LDS_BYTES = 13056
+_BLOCK_THREADS = 256
+
+_HIP_LIB = None
+_CO_DATA: bytes | None = None
+# HIP modules are bound to a device's primary context, so each rank/device
+# needs its own handle.
+_KFUNC_BY_DEV: dict = {}
+_KMODULE_BY_DEV: dict = {}
+_DISPATCH_COUNT = 0
+
+# CK aux tensors are needed by TE's backward wrapper. Capture a compatible
+# template on the first eligible call and then replace only its LSE tensor.
+_AUX_CTX_TEMPLATES: dict = {}
+
+
+def get_dispatch_count() -> int:
+    """Return successful hand-tuned kernel launches in this process."""
+    return _DISPATCH_COUNT
+
+
+def _ensure_hip_lib():
+    global _HIP_LIB
+    if _HIP_LIB is not None:
+        return _HIP_LIB
+    # ROCm Python wheels ship runtime and development copies of libamdhip64.
+    # PyTorch is linked against the versioned runtime SONAME; opening the
+    # unversioned development symlink can create a second HIP runtime with a
+    # separate module/context registry, causing hipModuleGetFunction to return
+    # hipErrorNotFound even though the code object contains the symbol.
+    lib = ctypes.CDLL("libamdhip64.so.7")
+    lib.hipModuleLoadData.restype = ctypes.c_int
+    lib.hipModuleLoadData.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+    lib.hipModuleGetFunction.restype = ctypes.c_int
+    lib.hipModuleGetFunction.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+    ]
+    lib.hipModuleLaunchKernel.restype = ctypes.c_int
+    lib.hipModuleLaunchKernel.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_uint,
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+    ]
+    _HIP_LIB = lib
+    return _HIP_LIB
+
+
+def _try_load_kernel() -> bool:
+    """Stage the code object; bind it to a HIP context on first launch."""
+    global _CO_DATA
+    try:
+        _ensure_hip_lib()
+        if _CO_DATA is None:
+            with open(_CO_PATH, "rb") as file:
+                _CO_DATA = file.read()
+        return True
+    except Exception as error:  # noqa: BLE001
+        logger.warning(
+            "could not stage hand-tuned hd64 kernel path=%s exists=%s: %r",
+            _CO_PATH,
+            os.path.exists(_CO_PATH),
+            error,
+        )
+        return False
+
+
+def _get_kfunc_for_device(device) -> ctypes.c_void_p:
+    import torch
+
+    dev_id = (
+        device.index if hasattr(device, "index") and device.index is not None else torch.cuda.current_device()
+    )
+    cached = _KFUNC_BY_DEV.get(dev_id)
+    if cached is not None:
+        return cached
+
+    hip = _ensure_hip_lib()
+    if _CO_DATA is None:
+        with open(_CO_PATH, "rb") as file:
+            globals()["_CO_DATA"] = file.read()
+
+    with torch.cuda.device(dev_id):
+        module = ctypes.c_void_p()
+        rc = hip.hipModuleLoadData(
+            ctypes.byref(module),
+            ctypes.create_string_buffer(_CO_DATA),
+        )
+        if rc != 0:
+            raise RuntimeError(f"hipModuleLoadData failed for {_CO_PATH} on device {dev_id} " f"(rc={rc})")
+        func = ctypes.c_void_p()
+        rc = hip.hipModuleGetFunction(
+            ctypes.byref(func),
+            module,
+            _KERNEL_NAME,
+        )
+        if rc != 0:
+            raise RuntimeError(f"hipModuleGetFunction failed on device {dev_id} (rc={rc})")
+
+    _KMODULE_BY_DEV[dev_id] = module
+    _KFUNC_BY_DEV[dev_id] = func
+    logger.info(
+        "loaded hand-tuned hd64 kernel from %s on device %d",
+        _CO_PATH,
+        dev_id,
+    )
+    return func
+
+
+def _is_gfx950(device) -> bool:
+    import torch
+
+    try:
+        return torch.cuda.get_device_properties(device).gcnArchName.startswith("gfx950")
+    except Exception:
+        return False
+
+
+def _ck_window_args(window_size, attn_mask_type: str) -> Tuple[int, int]:
+    if window_size is None:
+        window_left, window_right = -1, -1
+    else:
+        window_left, window_right = int(window_size[0]), int(window_size[1])
+    if "causal" in (attn_mask_type or ""):
+        window_right = 0
+    return window_left, window_right
+
+
+def _eligible(
+    *,
+    max_seqlen_q,
+    max_seqlen_kv,
+    q,
+    k,
+    v,
+    attn_scale,
+    attn_bias_type,
+    attn_mask_type,
+    softmax_type,
+    window_size,
+    bottom_right_diagonal,
+    qkv_layout,
+    dropout,
+    attn_bias,
+    softmax_offset,
+    s_quantizer,
+    o_quantizer,
+    fp8,
+) -> bool:
+    import torch
+
+    if not _ENABLED or fp8:
+        return False
+    if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16:
+        return False
+    # TE 2.15 passes quantizer objects through its BF16 fused-attention API even
+    # when fp8=False. They are unused by this BF16 output path and must not make
+    # an otherwise eligible call fall back to CK.
+    if attn_bias is not None or softmax_offset is not None:
+        return False
+    if attn_bias_type != "no_bias" or softmax_type != "vanilla":
+        return False
+    if bottom_right_diagonal not in (None, False):
+        return False
+    if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"):
+        return False
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        return False
+    if q.shape[-1] != 64 or v.shape[-1] != 64:
+        return False
+    if "causal" not in (attn_mask_type or ""):
+        return False
+    if dropout != 0.0:
+        return False
+    if not _is_gfx950(q.device):
+        return False
+
+    if qkv_layout.startswith("bshd"):
+        batch, sequence_length, query_heads, _ = q.shape
+        key_value_sequence_length = k.shape[1]
+        key_value_heads = k.shape[2]
+    else:
+        sequence_length, batch, query_heads, _ = q.shape
+        key_value_sequence_length = k.shape[0]
+        key_value_heads = k.shape[2]
+
+    if sequence_length != key_value_sequence_length:
+        return False
+    if max_seqlen_q != sequence_length or max_seqlen_kv != sequence_length:
+        return False
+    if q.shape[:2] != k.shape[:2] or k.shape[:2] != v.shape[:2]:
+        return False
+    if query_heads % key_value_heads != 0:
+        return False
+    if sequence_length % _BLOCK_M != 0:
+        return False
+    # TE uses 0.0 as a sentinel for the default 1/sqrt(head_dim) scale.
+    if attn_scale not in (None, 0.0) and not math.isclose(
+        float(attn_scale),
+        1.0 / math.sqrt(q.shape[-1]),
+        rel_tol=1e-6,
+        abs_tol=0.0,
+    ):
+        return False
+    if batch <= 0:
+        return False
+    return True
+
+
+def _launch(
+    q,
+    k,
+    v,
+    qkv_layout: str,
+    *,
+    attn_scale: float,
+    attn_mask_type: str,
+    window_size,
+    lse_out,
+):
+    import torch
+
+    if qkv_layout.startswith("bshd"):
+        batch, sequence_length, query_heads, head_dim = q.shape
+        key_value_heads = k.shape[2]
+        stride_q_s, stride_q_h, stride_q_b = (
+            q.stride(1),
+            q.stride(2),
+            q.stride(0),
+        )
+        stride_k_s, stride_k_h, stride_k_b = (
+            k.stride(1),
+            k.stride(2),
+            k.stride(0),
+        )
+        stride_v_s, stride_v_h, stride_v_b = (
+            v.stride(1),
+            v.stride(2),
+            v.stride(0),
+        )
+    else:
+        sequence_length, batch, query_heads, head_dim = q.shape
+        key_value_heads = k.shape[2]
+        stride_q_s, stride_q_h, stride_q_b = (
+            q.stride(0),
+            q.stride(2),
+            q.stride(1),
+        )
+        stride_k_s, stride_k_h, stride_k_b = (
+            k.stride(0),
+            k.stride(2),
+            k.stride(1),
+        )
+        stride_v_s, stride_v_h, stride_v_b = (
+            v.stride(0),
+            v.stride(2),
+            v.stride(1),
+        )
+
+    output = torch.empty_like(q)
+    if qkv_layout.startswith("bshd"):
+        stride_o_s, stride_o_h, stride_o_b = (
+            output.stride(1),
+            output.stride(2),
+            output.stride(0),
+        )
+    else:
+        stride_o_s, stride_o_h, stride_o_b = (
+            output.stride(0),
+            output.stride(2),
+            output.stride(1),
+        )
+
+    # The assembly kernel expects the CK log2(e)-scaled convention.
+    del attn_scale
+    scale_s = (1.0 / math.sqrt(head_dim)) * math.log2(math.e)
+    window_left, window_right = _ck_window_args(
+        window_size,
+        attn_mask_type,
+    )
+
+    kargs = struct.pack(
+        "<QQQQQ iiii iif iiii iiii ii iiii xxxx Q ii iiii QQ",
+        q.data_ptr(),
+        k.data_ptr(),
+        v.data_ptr(),
+        output.data_ptr(),
+        0,
+        sequence_length,
+        sequence_length,
+        head_dim,
+        head_dim,
+        query_heads,
+        query_heads // key_value_heads,
+        scale_s,
+        stride_q_s,
+        stride_k_s,
+        stride_v_s,
+        stride_o_s,
+        stride_q_h,
+        stride_k_h,
+        stride_v_h,
+        stride_o_h,
+        query_heads,
+        0,
+        window_left,
+        window_right,
+        0,
+        2,
+        lse_out.data_ptr(),
+        lse_out.stride(1),
+        lse_out.stride(0),
+        stride_q_b,
+        stride_k_b,
+        stride_v_b,
+        stride_o_b,
+        0,
+        0,
+    )
+    karg_buf = (ctypes.c_ubyte * len(kargs))(*kargs)
+    karg_size = ctypes.c_size_t(len(kargs))
+    extra = (ctypes.c_void_p * 5)(
+        ctypes.c_void_p(0x01),
+        ctypes.cast(karg_buf, ctypes.c_void_p),
+        ctypes.c_void_p(0x02),
+        ctypes.cast(ctypes.pointer(karg_size), ctypes.c_void_p),
+        ctypes.c_void_p(0x03),
+    )
+    kernel = _get_kfunc_for_device(q.device)
+    rc = _HIP_LIB.hipModuleLaunchKernel(
+        kernel,
+        query_heads,
+        sequence_length // _BLOCK_M,
+        batch,
+        _BLOCK_THREADS,
+        1,
+        1,
+        _LDS_BYTES,
+        ctypes.c_void_p(0),
+        ctypes.c_void_p(0),
+        ctypes.cast(extra, ctypes.c_void_p),
+    )
+    if rc != 0:
+        raise RuntimeError(f"hipModuleLaunchKernel returned rc={rc} on device {q.device}")
+    return output
+
+
+def _install_fused_attn_override():
+    try:
+        from transformer_engine.pytorch.cpp_extensions import (
+            fused_attn as fused_attn_module,
+        )
+    except Exception as error:  # noqa: BLE001
+        logger.info("TE not importable yet, deferring install (%r)", error)
+        return False
+
+    if getattr(
+        fused_attn_module.fused_attn_fwd,
+        "_fwd_attn_asm_patched",
+        False,
+    ):
+        return True
+    if not _try_load_kernel():
+        return False
+
+    import torch
+
+    original_fused_attn_fwd = fused_attn_module.fused_attn_fwd
+    original_parameter_names = frozenset(inspect.signature(original_fused_attn_fwd).parameters)
+
+    def patched_fused_attn_fwd(
+        is_training,
+        max_seqlen_q,
+        max_seqlen_kv,
+        cu_seqlens_q,
+        cu_seqlens_kv,
+        q,
+        k,
+        v,
+        fake_dtype,
+        fused_attention_backend,
+        attn_bias=None,
+        cu_seqlens_q_padded=None,
+        cu_seqlens_kv_padded=None,
+        page_table_k=None,
+        page_table_v=None,
+        s_quantizer=None,
+        o_quantizer=None,
+        attn_scale=None,
+        dropout=0.0,
+        fast_zero_fill=True,
+        qkv_layout="sbh3d",
+        attn_bias_type="no_bias",
+        attn_mask_type="padding",
+        softmax_type="vanilla",
+        window_size=(-1, -1),
+        bottom_right_diagonal=None,
+        rng_gen=None,
+        softmax_offset=None,
+        return_max_logit=False,
+        cuda_graph=False,
+    ):
+        global _DISPATCH_COUNT
+
+        # TE 2.12 uses False as the "no softmax offset" sentinel, while its
+        # C++ binding and TE 2.15 expect None.
+        if softmax_offset is False:
+            softmax_offset = None
+
+        all_kwargs = dict(
+            is_training=is_training,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_kv=max_seqlen_kv,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_kv=cu_seqlens_kv,
+            q=q,
+            k=k,
+            v=v,
+            fake_dtype=fake_dtype,
+            fused_attention_backend=fused_attention_backend,
+            attn_bias=attn_bias,
+            cu_seqlens_q_padded=cu_seqlens_q_padded,
+            cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+            page_table_k=page_table_k,
+            page_table_v=page_table_v,
+            s_quantizer=s_quantizer,
+            o_quantizer=o_quantizer,
+            attn_scale=attn_scale,
+            dropout=dropout,
+            fast_zero_fill=fast_zero_fill,
+            qkv_layout=qkv_layout,
+            attn_bias_type=attn_bias_type,
+            attn_mask_type=attn_mask_type,
+            softmax_type=softmax_type,
+            window_size=window_size,
+            bottom_right_diagonal=bottom_right_diagonal,
+            rng_gen=rng_gen,
+            softmax_offset=softmax_offset,
+            return_max_logit=return_max_logit,
+            cuda_graph=cuda_graph,
+        )
+        # TE 2.15 (v26.5) adds bottom_right_diagonal; TE 2.12 (v26.3)
+        # does not accept it. Keep one superset wrapper, but only forward
+        # parameters supported by the installed TE revision.
+        original_kwargs = {
+            name: value for name, value in all_kwargs.items() if name in original_parameter_names
+        }
+
+        if not _eligible(
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_kv=max_seqlen_kv,
+            q=q,
+            k=k,
+            v=v,
+            attn_scale=attn_scale,
+            attn_bias_type=attn_bias_type,
+            attn_mask_type=attn_mask_type,
+            softmax_type=softmax_type,
+            window_size=window_size,
+            bottom_right_diagonal=bottom_right_diagonal,
+            qkv_layout=qkv_layout,
+            dropout=dropout,
+            attn_bias=attn_bias,
+            softmax_offset=softmax_offset,
+            s_quantizer=s_quantizer,
+            o_quantizer=o_quantizer,
+            fp8=False,
+        ):
+            return original_fused_attn_fwd(**original_kwargs)
+        if return_max_logit or page_table_k is not None:
+            return original_fused_attn_fwd(**original_kwargs)
+
+        if qkv_layout.startswith("bshd"):
+            batch, sequence_length, query_heads, _ = q.shape
+        else:
+            sequence_length, batch, query_heads, _ = q.shape
+
+        lse = torch.empty(
+            batch,
+            query_heads,
+            sequence_length,
+            1,
+            dtype=torch.float32,
+            device=q.device,
+        )
+        try:
+            output = _launch(
+                q,
+                k,
+                v,
+                qkv_layout,
+                attn_scale=attn_scale or 0.0,
+                attn_mask_type=attn_mask_type,
+                window_size=window_size,
+                lse_out=lse,
+            )
+        except Exception as error:  # noqa: BLE001
+            logger.warning(
+                "hand-tuned launch failed (%r); falling back to CK",
+                error,
+            )
+            return original_fused_attn_fwd(**original_kwargs)
+
+        _DISPATCH_COUNT += 1
+        if _DISPATCH_LOG:
+            logger.info(
+                "fwd-attn-asm dispatched: shape=%s layout=%s mask=%s window=%s",
+                tuple(q.shape),
+                qkv_layout,
+                attn_mask_type,
+                window_size,
+            )
+        if not is_training:
+            return output, []
+
+        cache_key = (qkv_layout, attn_mask_type, dropout)
+        template = _AUX_CTX_TEMPLATES.get(cache_key)
+        if template is None:
+            ck_output = None
+            ck_aux = None
+            try:
+                ck_output, ck_aux = original_fused_attn_fwd(**original_kwargs)
+                template = list(ck_aux)
+            except RuntimeError as error:
+                if "fused attn configs not supported" not in str(error):
+                    raise
+                # TE 2.15 CK rejects the exact S=8192 gfx950 configuration,
+                # although its backward contract only needs softmaxStats and
+                # an RNG state. Dropout is zero for GPT-OSS, so a zero state is
+                # a valid placeholder and is not consumed by the kernel.
+                template = [
+                    lse,
+                    torch.zeros(2, dtype=torch.int64, device=q.device),
+                ]
+                logger.info("CK aux template unavailable; synthesized LSE/RNG aux " "for dropout=0")
+            _AUX_CTX_TEMPLATES[cache_key] = template
+            template_lse = template[0] if template else None
+            if (
+                template_lse is None
+                or getattr(template_lse, "ndim", -1) != lse.ndim
+                or template_lse.dtype != lse.dtype
+            ):
+                logger.warning(
+                    "CK aux LSE shape=%s dtype=%s mismatches ASM LSE "
+                    "shape=%s dtype=%s; falling back to CK",
+                    getattr(template_lse, "shape", None),
+                    getattr(template_lse, "dtype", None),
+                    lse.shape,
+                    lse.dtype,
+                )
+                _AUX_CTX_TEMPLATES.pop(cache_key, None)
+                return ck_output, ck_aux
+
+        aux = [lse] + list(template[1:])
+        return output, aux
+
+    patched_fused_attn_fwd._fwd_attn_asm_patched = True
+    fused_attn_module.fused_attn_fwd = patched_fused_attn_fwd
+
+    original_fused_attn_bwd = fused_attn_module.fused_attn_bwd
+
+    def diagnostic_fused_attn_bwd(*args, **kwargs):
+        if _DISPATCH_LOG:
+            logger.info(
+                "fused-attn-bwd config: q_shape=%s q_stride=%s k_stride=%s "
+                "v_stride=%s out_stride=%s dout_stride=%s",
+                getattr(args[4], "shape", None) if len(args) > 4 else None,
+                args[4].stride() if len(args) > 4 else None,
+                args[5].stride() if len(args) > 5 else None,
+                args[6].stride() if len(args) > 6 else None,
+                args[7].stride() if len(args) > 7 else None,
+                args[8].stride() if len(args) > 8 else None,
+            )
+        try:
+            return original_fused_attn_bwd(*args, **kwargs)
+        except RuntimeError:
+            aux = args[11] if len(args) > 11 else kwargs.get("aux_ctx_tensors")
+            logger.error(
+                "fused-attn-bwd rejected config: max_q=%s max_kv=%s "
+                "q=%s backend=%s aux=%s layout=%s mask=%s window=%s "
+                "bottom_right=%s deterministic=%s",
+                args[0] if len(args) > 0 else None,
+                args[1] if len(args) > 1 else None,
+                getattr(args[4], "shape", None) if len(args) > 4 else None,
+                args[12] if len(args) > 12 else None,
+                [(tuple(t.shape), str(t.dtype)) for t in (aux or [])],
+                args[21] if len(args) > 21 else None,
+                args[23] if len(args) > 23 else None,
+                args[25] if len(args) > 25 else None,
+                args[26] if len(args) > 26 else None,
+                args[27] if len(args) > 27 else None,
+            )
+            raise
+
+    diagnostic_fused_attn_bwd._fwd_attn_asm_bwd_diagnostic = True
+    fused_attn_module.fused_attn_bwd = diagnostic_fused_attn_bwd
+
+    try:
+        from transformer_engine.pytorch.attention.dot_product_attention import backends
+
+        backends.fused_attn_fwd = patched_fused_attn_fwd
+        backends.fused_attn_bwd = diagnostic_fused_attn_bwd
+    except Exception:
+        pass
+
+    logger.info("fused_attn_fwd patched " "(D=64 BF16 [SWA-]causal -> hand-tuned hd64 kernel)")
+    return True
+
+
+def _install_aiter_rope_override(rope_module):
+    """Restore the AITER RoPE route removed by newer ROCm TE revisions."""
+    if not _AITER_ROPE_ENABLED:
+        return False
+
+    fused_rope = rope_module.FusedRoPEFunc
+    if getattr(fused_rope, "_mlperf_aiter_rope_patched", False):
+        return True
+    # Older TE revisions already provide the same dispatch natively.
+    if hasattr(fused_rope, "_can_use_aiter"):
+        return True
+
+    from aiter.ops.rope import rope_bwd as aiter_rope_bwd
+    from aiter.ops.rope import rope_fwd as aiter_rope_fwd
+
+    original_forward = fused_rope.forward
+    original_backward = fused_rope.backward
+
+    def aiter_aware_forward(
+        ctx,
+        tensor,
+        freqs,
+        start_positions=None,
+        tensor_format="sbhd",
+        interleaved=False,
+        cu_seqlens=None,
+        cp_size=1,
+        cp_rank=0,
+    ):
+        use_aiter = (
+            tensor_format == "sbhd"
+            and not interleaved
+            and cu_seqlens is None
+            and cp_size == 1
+            and start_positions is None
+        )
+        ctx._mlperf_use_aiter_rope = use_aiter
+        if not use_aiter:
+            return original_forward(
+                ctx,
+                tensor,
+                freqs,
+                start_positions,
+                tensor_format,
+                interleaved,
+                cu_seqlens,
+                cp_size,
+                cp_rank,
+            )
+
+        if freqs.dtype != torch.float32:
+            freqs = freqs.float()
+        output = aiter_rope_fwd(tensor, freqs, 0, False, False)
+        ctx.save_for_backward(freqs, cu_seqlens, start_positions)
+        return output
+
+    def aiter_aware_backward(ctx, grad_output):
+        if not getattr(ctx, "_mlperf_use_aiter_rope", False):
+            return original_backward(ctx, grad_output)
+        freqs, _, _ = ctx.saved_tensors
+        grad_input = aiter_rope_bwd(grad_output, freqs, 0, False, False)
+        return grad_input, None, None, None, None, None, None, None, None
+
+    import torch
+
+    fused_rope.forward = staticmethod(aiter_aware_forward)
+    fused_rope.backward = staticmethod(aiter_aware_backward)
+    fused_rope._mlperf_aiter_rope_patched = True
+    logger.info("restored AITER fused RoPE dispatch for this TE revision")
+    return True
+
+
+class _DeferredInstaller:
+    _TARGETS = {
+        "transformer_engine.pytorch.attention.rope",
+        "transformer_engine.pytorch.cpp_extensions.fused_attn",
+        "transformer_engine.pytorch.cpp_extensions",
+    }
+
+    def find_spec(self, fullname, path, target=None):
+        try:
+            if fullname not in self._TARGETS:
+                return None
+            for finder in sys.meta_path:
+                if finder is self:
+                    continue
+                try:
+                    spec = finder.find_spec(fullname, path, target)
+                except (AttributeError, ImportError):
+                    spec = None
+                if spec is None:
+                    continue
+                original_loader = spec.loader
+
+                class _WrappedLoader:
+                    def create_module(self, module_spec):
+                        if hasattr(original_loader, "create_module"):
+                            return original_loader.create_module(module_spec)
+                        return None
+
+                    def exec_module(self, module):
+                        original_loader.exec_module(module)
+                        if fullname == "transformer_engine.pytorch.cpp_extensions.fused_attn":
+                            try:
+                                _install_fused_attn_override()
+                            except Exception as error:  # noqa: BLE001
+                                logger.warning(
+                                    "deferred install failed: %r",
+                                    error,
+                                )
+                        elif fullname == "transformer_engine.pytorch.attention.rope":
+                            try:
+                                _install_aiter_rope_override(module)
+                            except Exception as error:  # noqa: BLE001
+                                logger.warning(
+                                    "deferred AITER RoPE install failed: %r",
+                                    error,
+                                )
+
+                spec.loader = _WrappedLoader()
+                return spec
+            return None
+        except Exception as error:  # noqa: BLE001
+            logger.warning(
+                "fwd-attn-asm find_spec error for %s: %r",
+                fullname,
+                error,
+            )
+            return None
+
+
+def _register_deferred_install():
+    if any(isinstance(finder, _DeferredInstaller) for finder in sys.meta_path):
+        return
+    sys.meta_path.insert(0, _DeferredInstaller())
+    logger.info("deferred installer registered; will patch on TE load")
+
+
+if _ENABLED or _AITER_ROPE_ENABLED:
+    try:
+        _register_deferred_install()
+    except Exception as error:  # noqa: BLE001
+        logger.warning(
+            "fwd-attn-asm deferred install failed at startup: %r",
+            error,
+        )
diff --git a/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh b/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh
new file mode 100644
index 000000000..2e7e1a556
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh
@@ -0,0 +1,248 @@
+#!/bin/bash
+# =============================================================================
+# MLPerf GPT-OSS-20B Configuration for MI355X (1 node, 8 GPUs)
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# System Configuration
+# -----------------------------------------------------------------------------
+export DGXSYSTEM=MI355X_1x8x1
+export GPUS_PER_NODE=8
+export NNODES=1
+export NODE_RANK=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29501
+
+# -----------------------------------------------------------------------------
+# Paths
+# -----------------------------------------------------------------------------
+export PRIMUS_PATH=/workspace/Primus
+export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-LM:${PYTHONPATH}"
+export EXP=${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml
+export DATA_PATH=/data
+
+# -----------------------------------------------------------------------------
+# Training Hyperparameters
+# -----------------------------------------------------------------------------
+export PRIMUS_MICRO_BATCH_SIZE=4
+export PRIMUS_GLOBAL_BATCH_SIZE=32
+export EVAL_ITERS=$((1024 / PRIMUS_GLOBAL_BATCH_SIZE))  # MLPerf closed: eval_iters * GBS = 1024 eval samples
+export PRIMUS_LR=8.0e-4
+export PRIMUS_MIN_LR=8.0e-5 # Set to 10% of max LR
+export PRIMUS_TRAIN_ITERS=1200000
+export PRIMUS_LR_WARMUP_ITERS=128
+export PRIMUS_LR_DECAY_ITERS=$((PRIMUS_TRAIN_ITERS-PRIMUS_LR_WARMUP_ITERS)) # 1200000 - 128 = 1199872
+# export SEED=30279
+
+# Evaluation frequency (sample-based, adjusts automatically with GBS)
+export EVAL_SAMPLES_INTERVAL=12288   # Evaluate every 12,288 samples
+export PRIMUS_EVAL_INTERVAL=$((EVAL_SAMPLES_INTERVAL / PRIMUS_GLOBAL_BATCH_SIZE))  # Auto-computed
+
+# -----------------------------------------------------------------------------
+# Parallelism
+# -----------------------------------------------------------------------------
+export PRIMUS_TP=1
+export PRIMUS_PP=1
+export PRIMUS_EP=1
+
+# -----------------------------------------------------------------------------
+# Primus Configuration
+# -----------------------------------------------------------------------------
+export PRIMUS_TURBO_GROUPED_GEMM_BACKEND="${PRIMUS_TURBO_GROUPED_GEMM_BACKEND:-triton}"
+export PRIMUS_TURBO_GEMM_BACKEND=triton
+export PRIMUS_TURBO_FUSED_WGRAD_ACCUM="${PRIMUS_TURBO_FUSED_WGRAD_ACCUM:-1}"
+export PRIMUS_NUM_WORKERS="${PRIMUS_NUM_WORKERS:-2}"
+export PRIMUS_GRAD_REDUCE_IN_BF16=true
+export USE_TURBO_RMS_NORM=true
+
+# -----------------------------------------------------------------------------
+# ROCm / System Runtime
+# -----------------------------------------------------------------------------
+export GPU_MAX_HW_QUEUES=2
+export HIP_FORCE_DEV_KERNARG=1
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export HSA_KERNARG_POOL_SIZE=12582912
+export TORCH_NCCL_HIGH_PRIORITY=1
+export ENABLE_NUMA_BINDING=1
+export PYTORCH_ALLOC_CONF=expandable_segments:False
+export HSA_NO_SCRATCH_RECLAIM=1
+export HSA_ENABLE_SDMA=1
+export HSA_ENABLE_INTERRUPT=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export OMP_NUM_THREADS=1
+export PYTHONWARNINGS=ignore
+export TOKENIZERS_PARALLELISM=false
+
+# -----------------------------------------------------------------------------
+# RCCL / NCCL Tuning
+# -----------------------------------------------------------------------------
+export NCCL_MIN_P2P_NCHANNELS=32
+export NCCL_MIN_CTAS=32
+export NCCL_NCHANNELS_PER_NET_PEER=32
+export NCCL_NVLS_ENABLE=0
+export NCCL_CHECKS_DISABLE=1
+
+# -----------------------------------------------------------------------------
+# hipBLASLt
+# -----------------------------------------------------------------------------
+export USE_HIPBLASLT=1
+export TORCH_BLAS_PREFER_HIPBLASLT=1
+
+# Solution indices are runtime-specific. For normal v26.3 runs, replay the
+# Stage 0 cache validated with this runtime instead of the legacy override table.
+# Setting TE_HIPBLASLT_ALGO_SAVE selects online tuning and suppresses the default
+# replay cache. HIPBLASLT_TUNING_OVERRIDE_FILE remains available as an explicit
+# diagnostic override.
+if [ -n "${TE_HIPBLASLT_ALGO_SAVE:-}" ]; then
+    export TE_HIPBLASLT_ALGO_SAVE
+    unset TE_HIPBLASLT_ALGO_LOAD
+    unset HIPBLASLT_TUNING_OVERRIDE_FILE
+elif [ -n "${TE_HIPBLASLT_ALGO_LOAD:-}" ]; then
+    if [ -f "${TE_HIPBLASLT_ALGO_LOAD}" ]; then
+        export TE_HIPBLASLT_ALGO_LOAD
+        unset HIPBLASLT_TUNING_OVERRIDE_FILE
+    else
+        unset TE_HIPBLASLT_ALGO_LOAD
+    fi
+elif [ -n "${HIPBLASLT_TUNING_OVERRIDE_FILE:-}" ]; then
+    if [ -f "${HIPBLASLT_TUNING_OVERRIDE_FILE}" ]; then
+        export HIPBLASLT_TUNING_OVERRIDE_FILE
+    else
+        unset HIPBLASLT_TUNING_OVERRIDE_FILE
+    fi
+else
+    if [ "${MLPERF_RUNTIME_SERIES:-v26.3}" = "v26.3" ]; then
+        TE_HIPBLASLT_ALGO_LOAD="${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/te_hipblaslt_algo-v26.3.csv"
+        if [ -f "${TE_HIPBLASLT_ALGO_LOAD}" ]; then
+            export TE_HIPBLASLT_ALGO_LOAD
+        else
+            unset TE_HIPBLASLT_ALGO_LOAD
+        fi
+    else
+        HIPBLASLT_TUNING_OVERRIDE_FILE="${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/tune_gemm_results-${MLPERF_RUNTIME_SERIES}.txt"
+        if [ -f "${HIPBLASLT_TUNING_OVERRIDE_FILE}" ]; then
+            export HIPBLASLT_TUNING_OVERRIDE_FILE
+        else
+            unset HIPBLASLT_TUNING_OVERRIDE_FILE
+        fi
+    fi
+fi
+
+# -----------------------------------------------------------------------------
+# NVTE — FP8 & Cast Transpose
+# -----------------------------------------------------------------------------
+export NVTE_ROCM_ENABLE_MXFP8=0
+export NVTE_USE_CAST_TRANSPOSE_TRITON=0
+export NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE=1
+
+# -----------------------------------------------------------------------------
+# NVTE — FMHA / CK Backend
+# -----------------------------------------------------------------------------
+export NVTE_FLASH_ATTN=0                  # Disable FlashAttention so FusedAttention (CK/ASM) is used
+export NVTE_CK_USES_FWD_V3=1              # Globally on; aiter selects v3 vs CK-tile internally
+export NVTE_CK_USES_BWD_V3=1              # Globally on; aiter selects v3 vs CK-tile internally
+export NVTE_USE_AITER_ROPE=1              # Route RoPE through aiter's fused kernel instead of TE's own CK kernel
+export NVTE_FMHA_USE_BSHD=0               # Use the native SBHD path provided by the v26.5 AITER stack
+export NVTE_CK_HOW_V3_BF16_CVT=2          # 0=RTNE, 1=RTNA, 2=RTZ (gfx942 selector; gfx950 is fixed)
+
+# Route eligible D=64 BF16 causal forward calls to the hand-tuned gfx950
+# kernel. FMHA_HD64_ASM_LOG=1 prints one line per successful launch.
+export MLPERF_ENABLE_FWD_ATTN_ASM="${MLPERF_ENABLE_FWD_ATTN_ASM:-1}"
+export FMHA_HD64_ASM_LOG="${FMHA_HD64_ASM_LOG:-0}"
+
+# The RTZ custom backward code object is registered under the RTNE-named slot
+# that v26.5 AITER actually selects on gfx950. Unlike the bundled v26.5 a16
+# kernel that overflowed at sequence length 8192, this injected kernel passed a
+# 200-step FP8 C4 convergence run. Set the flag to 0 for the a32 PSSK fallback.
+export MLPERF_ENABLE_BWD_ATTN_ASM="${MLPERF_ENABLE_BWD_ATTN_ASM:-1}"
+if [ "${MLPERF_ENABLE_BWD_ATTN_ASM}" = "1" ]; then
+    export NVTE_CK_IS_V3_ATOMIC_FP32=0
+elif [ "${MLPERF_ENABLE_BWD_ATTN_ASM}" = "0" ]; then
+    export NVTE_CK_IS_V3_ATOMIC_FP32=1
+else
+    echo "MLPERF_ENABLE_BWD_ATTN_ASM must be 0 or 1" >&2
+    if [ "${BASH_SOURCE[0]}" != "$0" ]; then
+        return 2
+    fi
+    exit 2
+fi
+
+# -----------------------------------------------------------------------------
+# NVTE — Debug
+# -----------------------------------------------------------------------------
+export NVTE_DEBUG=0
+export NVTE_DEBUG_LEVEL=0
+export NVTE_LOG_FUSED_ATTN_CONFIG=0
+export NVTE_LOG_CK_CONFIG=0
+export CK_FUSED_ATTN_LOG_CONFIG=0
+# export NVTE_FMHA_DEBUG=1                    # keep commented; debug-only knob
+
+# -----------------------------------------------------------------------------
+# MLPerf Logging
+# -----------------------------------------------------------------------------
+export LOG_INTERVAL=999999
+export MLLOG_TRAIN_LOSS_LOG_FREQ=0
+export MLLOG_TARGET_EVAL_LOSS=3.34
+export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out
+export MLLOG_SAVE_TO_FILE=0
+export MLLOG_SUBMISSION_BENCHMARK=gpt_oss_20b
+export MLLOG_SUBMISSION_DIVISION=closed
+export MLLOG_SUBMISSION_ORG=AMD
+export MLLOG_SUBMISSION_PLATFORM=MI355X
+
+export MLLOG_TENSOR_PARALLELISM=1
+export MLLOG_PIPELINE_PARALLELISM=1
+export MLLOG_CONTEXT_PARALLELISM=1
+export MLLOG_EXPERT_PARALLELISM=1
+export MLLOG_MICRO_BATCH_SIZE=4
+MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}")
+export MLLOG_CONFIG_FILENAME
+export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR='fp8'
+
+# -----------------------------------------------------------------------------
+# Synthetic Warmup (kernel pre-compilation)
+# -----------------------------------------------------------------------------
+export SYNTH_WARMUP_STEPS=3
+
+# -----------------------------------------------------------------------------
+# MoE Token Dispatcher
+# -----------------------------------------------------------------------------
+# Skip sort_chunks_by_idxs when the per-local-expert index is an identity
+# permutation (fires at EP=1/TP=1). Set to 0 to run the original path; useful
+# for A/B measurements. Implemented by skip_identity_sort_patches.py.
+export MOE_SKIP_IDENTITY_SORT=1
+
+# -----------------------------------------------------------------------------
+# DDP Parameter All-Gather (SDMA)
+# -----------------------------------------------------------------------------
+# v26.5 turns five SDMA workspace barriers into ~5 ms waits each. A same-node
+# A/B with the latest Turbo main measured 1004 ms/step with SDMA and 988
+# ms/step with RCCL, matching the v26.3 control at 989 ms/step. Keep SDMA
+# opt-in on v26.5 until its barrier regression is fixed; retain the validated
+# v26.3 default.
+if [ "${MLPERF_RUNTIME_SERIES:-v26.5}" = "v26.3" ]; then
+    DEFAULT_ENABLE_SDMA_ALLGATHER=1
+else
+    DEFAULT_ENABLE_SDMA_ALLGATHER=0
+fi
+export ENABLE_SDMA_ALLGATHER="${ENABLE_SDMA_ALLGATHER:-${DEFAULT_ENABLE_SDMA_ALLGATHER}}"
+# Optional: cap the per-call peer-copy stream count. Default is
+# min(world_size-1, 8); lower values reduce SDMA / memory-system pressure.
+# export MEGATRON_SDMA_PEER_COPY_STREAMS=8
+# When SDMA is explicitly enabled, keep the source-patch behavior: the first
+# two parameter bucket groups use RCCL and later groups use SDMA.
+export MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS=${MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS:-2}
+
+# -----------------------------------------------------------------------------
+# Run-log verbosity
+# -----------------------------------------------------------------------------
+# MLPerf run-log verbosity. Default 0 keeps only :::MLLOG + ``run_and_time.sh``
+# banners on stdout. Set to 1 to restore the full framework output (Primus
+# loguru banners / Megatron / TE / aiter / Gloo / torchrun / hipify / ...)
+# when debugging. See src/_log_suppression.py for the full strategy.
+export MLPERF_VERBOSE_LOGS=${MLPERF_VERBOSE_LOGS:-0}
+
+# fused rms and swiglu no cat
+export PRIMUS_FUSED_RESIDUAL_NORM=1
+export PRIMUS_MOE_SWIGLU_NOCAT=1
+export MLLOG_BLOCK_TPUT_LOG=0
diff --git a/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml b/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml
new file mode 100644
index 000000000..00efb9ebb
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml
@@ -0,0 +1,223 @@
+work_group: ${TEAM:amd}
+user_name: ${USER:root}
+exp_name: ${EXP_NAME:gpt_oss_20b}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: ${PRIMUS_MODEL:gpt_oss_20B}.yaml
+    overrides:
+
+      # Activate the migrated MLPerf pretrain trainer (mllog + MLPerf hooks).
+      stage: ${PRIMUS_STAGE:mlperf_pretrain}
+
+      # tokenizer
+      tokenizer_type: Llama3Tokenizer
+      tokenizer_model: ${MODEL:meta-llama/Llama-3.1-8B}
+
+      # model
+      num_layers: 24
+      hidden_size: 2880
+      ffn_hidden_size: 2880
+      num_attention_heads: 64
+      num_query_groups: 8  # Group Query Attention (GQA) - matches HF num_key_value_heads
+      num_experts: 32
+      activation_func: swiglu  # SiLU activation (matches HF hidden_act: "silu")
+
+      # rotary
+      position_embedding_type: rope
+      rotary_base: 150000
+
+      # mixed-precision
+      attention_softmax_in_fp32: false
+      grad_reduce_in_bf16: ${PRIMUS_GRAD_REDUCE_IN_BF16:true}
+
+      # log
+      wandb_project: "Primus_GPT_OSS_20B"
+      stderr_sink_level: ERROR
+      log_interval: ${LOG_INTERVAL:999999}
+
+      # debug
+      # moe_router_force_load_balancing: true
+      # log_avg_skip_iterations: 2
+      # log_avg_reset_interval: 50
+
+      # profile
+      profile: ${PRIMUS_PROFILE:False}
+      use_pytorch_profiler: ${PRIMUS_PROFILE:False}
+      profile_step_end: ${PRIMUS_PROFILE_STEP_END:32}
+      profile_step_start: ${PRIMUS_PROFILE_STEP_START:16}
+      profile_ranks: [0]
+
+      # enable fp8 training
+      fp8: ${PRIMUS_FP8:e4m3}
+      fp8_recipe: tensorwise
+      clip_grad: 1.0  # Gradient clipping (already default, but explicit)
+      check_for_nan_in_loss_and_grad: false
+
+      # hyper parameters
+      train_iters: ${PRIMUS_TRAIN_ITERS:1200000}
+      micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2}
+      global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:16}
+      seq_length: ${PRIMUS_SEQ_LENGTH:8192}
+      max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:131072}
+      seed: ${SEED:1234}  # Random seed for reproducibility
+      lr: ${PRIMUS_LR:8.0e-4}  # Reduced from 8e-4 for FP8 stability
+      min_lr: ${PRIMUS_MIN_LR:8.0e-5}  # Set to 10% of max LR
+      lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:128}
+      lr_decay_iters: ${PRIMUS_LR_DECAY_ITERS:1199872}
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      optimizer: adam
+      use_distributed_optimizer: true # use distributed optimizer
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-5
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
+      layernorm_epsilon: 1.0e-05  # RMSNorm epsilon (matches HF rms_norm_eps)
+
+      # Dropout (disabled for training)
+      hidden_dropout: 0.0
+      attention_dropout: 0.0
+
+      # parallel
+      tensor_model_parallel_size: ${PRIMUS_TP:1}
+      pipeline_model_parallel_size: ${PRIMUS_PP:1}
+      expert_model_parallel_size: ${PRIMUS_EP:8}
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      ddp_num_buckets: 8
+      ddp_average_in_collective: true
+
+      # data
+      mock_data: false
+      num_workers: ${PRIMUS_NUM_WORKERS:0}
+      train_data_path: "10 /data/c4-train.en_6_text_document"
+      valid_data_path: "/data/c4-validation-91205-samples.en_text_document"
+      test_data_path: "/data/c4-validation-91205-samples.en_text_document"
+      # Avoid copying a dense (B, 1, S, S) CPU attention mask every step.
+      # TE receives causal/sliding-window metadata from attn_mask_type + window_size.
+      # Use numeric 0/1 because Primus env expansion only type-casts numbers.
+      create_attention_mask_in_dataloader: ${PRIMUS_CREATE_ATTENTION_MASK_IN_DATALOADER:0}
+
+      # fusion
+      moe_permute_fusion: true
+      gradient_accumulation_fusion: true
+      moe_use_legacy_grouped_gemm: false # Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedMLP, please set `moe_use_legacy_grouped_gemm=True
+      moe_use_fused_router_with_aux_score: true
+      multi_latent_attention: false # Flag config.ENABLE_EXPERIMENTAL not enabled
+      apply_rope_fusion: true
+
+
+      # sliding window attention (GPT-OSS-20B model definition; matches HF sliding_window: 128)
+      # use_turbo_attention is false so non-turbo attention (which supports sliding window) is used.
+      # Pattern: alternating sliding_attention (1) and full_attention (0) for 24 layers
+      # window_size must be a tuple (left_window, right_window) for Transformer Engine
+      # For causal attention: left = past tokens, right = 0 (no future tokens)
+      # HF sliding_window: 128 means 128 past tokens, so use (128, 0)
+      window_size: [128, 0]  # Left window: 128 past tokens, Right: 0 (causal)
+      window_attn_skip_freq: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]
+
+      # MoE settings
+      moe_apply_probs_on_input: false
+      moe_aux_loss_coeff: 0.0 #0.9
+      moe_deepep_num_sms: 20
+      moe_enable_deepep: false
+      moe_expert_capacity_factor: null
+      moe_extended_tp: false
+      moe_ffn_hidden_size: 2880
+      moe_flex_dispatcher_backend: deepep
+      moe_grouped_gemm: true
+      moe_hybridep_num_sms: 16
+      moe_input_jitter_eps: null
+      moe_latent_size: null
+      moe_layer_freq: 1
+      moe_layer_recompute: false
+      moe_pad_expert_input_to_capacity: false
+      moe_per_layer_logging: false
+      moe_router_bias_update_rate: 0.001
+      moe_router_dtype: fp32 # DeepEP only supports float32 probs
+      moe_router_enable_expert_bias: false
+      moe_router_force_load_balancing: false
+      moe_router_fusion: true
+      moe_router_group_topk: null
+      moe_router_load_balancing_type: none
+      moe_router_num_groups: null
+      moe_router_padding_for_fp8: false
+      # Keep routing exact. PrimusTurbo consumes the ragged expert token counts
+      # directly, so no fake routes or TE multi_padding/unpadding are needed.
+      moe_router_padding_for_quantization: false
+      moe_router_pre_softmax: false
+      moe_router_score_function: softmax
+      moe_router_topk: 4
+      moe_router_topk_limited_devices: null
+      moe_router_topk_scaling_factor: null
+      moe_shared_expert_gate: false
+      moe_shared_expert_intermediate_size: null
+      moe_shared_expert_overlap: false
+      moe_token_dispatcher_type: alltoall
+      moe_token_drop_policy: probs
+      moe_token_dropping: false
+      moe_z_loss_coeff: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 100000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      exit_on_missing_checkpoint: false
+      ckpt_format: torch
+      eval_iters: ${EVAL_ITERS:64}  # eval_samples = eval_iters * GBS = 1024; set EVAL_ITERS in config shell (1024/GBS).
+      eval_interval: ${PRIMUS_EVAL_INTERVAL:768}
+
+      # Turbo
+      enable_primus_turbo: true
+      use_turbo_attention: false
+      use_turbo_grouped_gemm: true
+      use_turbo_ragged_grouped_gemm: ${USE_TURBO_RAGGED_GROUPED_GEMM:true}
+      use_turbo_rms_norm: ${USE_TURBO_RMS_NORM:true}
+      # Keeps the fused layernorm+QKV site's norm on Turbo and its GEMM on
+      # hipBLASLt, which is where each backend is faster on these shapes.
+      use_turbo_norm_te_linear: ${USE_TURBO_NORM_TE_LINEAR:true}
+      use_turbo_fused_act_with_probs : true
+
+      # deepep
+      use_turbo_deepep: false
+
+      # 64 or 80 for ep8, 32 for ep16-64 is best practice
+      turbo_deepep_num_cu: 64
+      turbo_deepep_use_comm_stream: false
+
+      # sync-free moe support stage 0-3, 0 means not use sync-free moe
+      # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory
+      # stage 2 is recommended for better performance
+      turbo_sync_free_moe_stage: 0
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
+
+      # tensorboard logging, set 'disable_tensorboard: false' to enable tensorboard logging
+      disable_tensorboard: ${PRIMUS_DISABLE_TENSORBOARD:true}
+      tensorboard_dir: ${PRIMUS_WORKSPACE:./output}/tensorboard
+      tensorboard_log_interval: 1
+      tensorboard_queue_size: 1000
+      log_timers_to_tensorboard: true
+      log_batch_size_to_tensorboard: true
+      log_learning_rate_to_tensorboard: true
+      log_validation_ppl_to_tensorboard: true
+      log_memory_to_tensorboard: true
+      log_world_size_to_tensorboard: true
+      log_loss_scale_to_tensorboard: true
diff --git a/examples/mlperf/gpt_oss_20b/prewarm_attention.py b/examples/mlperf/gpt_oss_20b/prewarm_attention.py
new file mode 100644
index 000000000..dc0cfcf47
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/prewarm_attention.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Populate the two GPT-OSS attention JIT cache entries before torchrun."""
+
+import argparse
+
+import aiter_hd64_asm_override
+import torch
+from transformer_engine.pytorch.attention import DotProductAttention
+
+
+def prewarm(sequence_length: int, micro_batch_size: int, window_left: int) -> None:
+    attention = DotProductAttention(
+        num_attention_heads=64,
+        kv_channels=64,
+        num_gqa_groups=8,
+        attention_dropout=0.0,
+        qkv_format="sbhd",
+        attn_mask_type="causal",
+        window_size=(window_left, 0),
+    ).cuda()
+    shapes = (
+        (sequence_length, micro_batch_size, 64, 64),
+        (sequence_length, micro_batch_size, 8, 64),
+        (sequence_length, micro_batch_size, 8, 64),
+    )
+    query, key, value = (
+        torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) for shape in shapes
+    )
+    output = attention(query, key, value)
+    if isinstance(output, tuple):
+        output = output[0]
+    output.float().square().mean().backward()
+    torch.cuda.synchronize()
+    print(f"attention_prewarm=PASS window_left={window_left}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sequence-length", type=int, default=8192)
+    parser.add_argument("--micro-batch-size", type=int, default=4)
+    parser.add_argument("--window-left", type=int, action="append")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("ROCm GPU is not available")
+    torch.cuda.set_device(0)
+    torch.manual_seed(1234)
+
+    windows = args.window_left or [128, -1]
+    for window_left in windows:
+        prewarm(args.sequence_length, args.micro_batch_size, window_left)
+    if aiter_hd64_asm_override.get_dispatch_count() < len(windows):
+        raise RuntimeError("Forward ASM override was not dispatched")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/mlperf/gpt_oss_20b/run_and_time.sh b/examples/mlperf/gpt_oss_20b/run_and_time.sh
new file mode 100755
index 000000000..1ef73b1ad
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/run_and_time.sh
@@ -0,0 +1,67 @@
+#!/bin/bash
+
+set -e
+
+# Create results directory
+mkdir -p /results
+
+cd "${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b"
+TRAIN_LOG_FILE="${TRAIN_LOG_FILE:-train.mlperfpretrain.exp.log}"
+
+# Under a multi-node scheduler wrapper, inherit rendezvous + node sizing from
+# SLURM so the same benchmark config scales without edits. Single-node jobs
+# fall through to the config defaults.
+if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then
+    NNODES="${SLURM_NNODES}"
+    NODE_RANK="${SLURM_NODEID:-0}"
+fi
+
+# TE 2.15 lazily compiles CK attention blobs. Populate both GPT-OSS windows
+# once before torchrun so eight ranks do not race while writing the cache.
+if [ "${MLPERF_RUNTIME_SERIES:-v26.3}" = "v26.5" ] \
+    && [ "${MLPERF_SKIP_ATTENTION_PREWARM:-0}" != "1" ]; then
+    python3 /opt/mlperf-gpt-oss-20b/prewarm_attention.py
+fi
+
+echo "============================================"
+echo "MLPerf GPT-OSS-20B Training"
+echo "============================================"
+echo "Config: ${EXP}"
+echo "Data:   ${DATA_PATH}"
+echo "GPUs:   ${GPUS_PER_NODE}"
+echo "Nodes:  ${NNODES}"
+echo "Rank:   ${NODE_RANK}"
+echo "Master: ${MASTER_ADDR}:${MASTER_PORT}"
+echo "============================================"
+
+# Start timing
+start=$(date +%s)
+start_fmt=$(date +%Y-%m-%d\ %r)
+echo "STARTING TIMING RUN AT $start_fmt"
+
+# Launch through Primus CLI and keep the real exit code even though output is
+# piped through tee.
+set +e
+"${PRIMUS_PATH}/primus-cli" direct -- \
+    train pretrain \
+    --config "${EXP}" \
+    2>&1 | tee "${TRAIN_LOG_FILE}"
+ret_code=${PIPESTATUS[0]}
+set -e
+
+# End timing
+end=$(date +%s)
+end_fmt=$(date +%Y-%m-%d\ %r)
+echo "ENDING TIMING RUN AT $end_fmt"
+
+# Report result
+result=$(( end - start ))
+result_name="GPT_OSS_20B"
+echo "RESULT,$result_name,,$result,AMD,$start_fmt"
+
+if [[ $ret_code != 0 ]]; then
+    echo "Training failed with exit code: $ret_code"
+    exit "$ret_code"
+fi
+
+exit 0
diff --git a/examples/mlperf/gpt_oss_20b/te_hipblaslt_algo-v26.3.csv b/examples/mlperf/gpt_oss_20b/te_hipblaslt_algo-v26.3.csv
new file mode 100644
index 000000000..1191c4552
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/te_hipblaslt_algo-v26.3.csv
@@ -0,0 +1,10 @@
+dev_cap,m,n,k,trans_a,trans_b,type_a,type_b,type_d,bias_type,aux_type,lda,ldb,ldd,scale_mode,epi,comp,scale_type,ws_min,ws_max,algo_id,aidx
+905,5120,32768,2880,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,2880,2880,5120,0,-,f32,float32,0,67108864,349458,30
+905,2880,32768,4096,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,4096,4096,2880,0,-,f32,float32,0,67108864,349558,7
+905,32,32768,2880,T,N,bfloat16,bfloat16,float32,-,-,2880,2880,32,0,-,f32,float32,0,67108864,337773,36
+905,2880,32768,32,N,N,float32,float32,float32,-,-,2880,32,2880,0,-,f32,float32,0,67108864,384788,31
+905,2880,32,32768,N,T,float32,float32,float32,-,-,2880,32,2880,0,-,f32,float32,5160960,67108864,383709,1
+905,4096,32768,2880,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,2880,2880,4096,0,-,f32,float32,0,67108864,349495,40
+905,4096,2880,32768,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,32768,32768,4096,0,-,f32,float32,0,67108864,349625,0
+905,2880,32768,5120,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,5120,5120,2880,0,-,f32,float32,0,67108864,349570,4
+905,2880,5120,32768,T,N,float8e4m3,float8e4m3,bfloat16,float32,-,32768,32768,2880,0,-,f32,float32,0,67108864,349555,1
diff --git a/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.3.txt b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.3.txt
new file mode 100644
index 000000000..9d987b63d
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.3.txt
@@ -0,0 +1,25 @@
+Git Version: de5c1aebb6-dirty
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,32,32768,2880,1,2880,92160,0,2880,94371840,32,1048576,32,1048576,bf16_r,bf16_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,151492,4511.27,39.8689,303082,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,32,32768,1,2880,94371840,0,32,1048576,2880,92160,2880,92160,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,81491.9,4800.79,74.1153,311532,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,32768,4096,1,4096,11796480,0,4096,134217728,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.69535e+06,1086.96,286.825,306398,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,4096,2880,32768,1,32768,134217728,0,32768,94371840,4096,11796480,4096,11796480,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.59496e+06,788.341,297.921,306251,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,4096,32768,2880,1,2880,11796480,0,2880,94371840,4096,134217728,4096,134217728,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.42444e+06,1094.09,318.875,306237,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,5120,32768,2880,1,2880,14745600,0,2880,94371840,5120,167772160,5120,167772160,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.39063e+06,1024.47,404.231,306248,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,2880,32768,32,1,2880,92160,0,32,1048576,2880,94371840,2880,94371840,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,64447.9,3796.71,93.716,311699,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,32768,5120,1,5120,14745600,0,5120,167772160,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.81411e+06,1006.88,343.401,306396,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,5120,32768,1,32768,94371840,0,32768,167772160,2880,14745600,2880,14745600,f8_r,f8_r,bf16_r,bf16_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.96197e+06,832.489,326.259,306396,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,2880,32768,128256,1,2880,369377280,0,128256,4202692608,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.32172e+06,474.576,18315.1,300658,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,128256,32768,2880,1,2880,369377280,0,2880,94371840,128256,4202692608,128256,4202692608,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.45958e+06,524.074,16585.3,302286,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,128256,32768,1,2880,94371840,1,128256,4202692608,2880,369377280,2880,369377280,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.31028e+06,544.947,18475.1,300155,gfx950:sramecc+:xnack-,256
diff --git a/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt
new file mode 100644
index 000000000..4e3711ba3
--- /dev/null
+++ b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt
@@ -0,0 +1,41 @@
+Git Version: fa9cdd18
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,32,32768,2880,1,2880,92160,0,2880,94371840,32,1048576,32,1048576,bf16_r,bf16_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,182699,5440.59,33.0588,24641,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,2880,32768,32,1,2880,92160,0,32,1048576,2880,94371840,2880,94371840,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,68985.3,4064.01,87.552,73624,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,4096,2880,32768,1,4096,134217728,0,2880,94371840,4096,11796480,4096,11796480,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.4418e+06,835.047,536.202,12468,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,4096,2880,32768,1,4096,134217728,1,2880,94371840,4096,11796480,4096,11796480,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.4024e+06,891.947,551.265,11292,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,32,32768,1,2880,94371840,0,32,1048576,2880,92160,2880,92160,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,86893.9,5119.03,69.5077,72656,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,5120,32768,1,2880,94371840,0,5120,167772160,2880,14745600,2880,14745600,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.42455e+06,760.278,678.366,11258,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,5120,32768,1,2880,94371840,1,5120,167772160,2880,14745600,2880,14745600,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.38626e+06,818.644,697.102,12554,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,32768,4096,1,4096,11796480,0,4096,134217728,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.43674e+06,832.117,538.09,19081,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,4096,32768,2880,1,4096,11796480,0,2880,94371840,4096,134217728,4096,134217728,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.47155e+06,852.281,525.36,14593,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,2880,32768,5120,1,2880,14745600,0,5120,167772160,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.46461e+06,781.659,659.811,14734,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,5120,32768,2880,1,2880,14745600,0,2880,94371840,5120,167772160,5120,167772160,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.53294e+06,818.125,630.401,21406,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,N,0,1,2880,32768,128256,1,2880,369377280,0,128256,4202692608,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.53955e+06,552.788,15723.8,13283,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    N,T,0,1,2880,128256,32768,1,2880,94371840,1,128256,4202692608,2880,369377280,2880,369377280,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.49105e+06,620.132,16235.2,12111,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,128256,32768,2880,1,2880,369377280,0,2880,94371840,128256,4202692608,128256,4202692608,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.46089e+06,524.544,16570.4,22416,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,4096,2880,32768,1,32768,134217728,0,32768,94371840,4096,11796480,4096,11796480,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.79595e+06,849.399,276.505,37702,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,32768,4096,1,4096,11796480,0,4096,134217728,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.81692e+06,1135.99,274.447,38550,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,5120,32768,1,32768,94371840,0,32768,167772160,2880,14745600,2880,14745600,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.97207e+06,835.326,325.15,36173,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,2880,32768,5120,1,5120,14745600,0,5120,167772160,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.87346e+06,1028.12,336.308,38461,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,4096,32768,2880,1,2880,11796480,0,2880,94371840,4096,134217728,4096,134217728,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.56755e+06,1158.67,301.102,36093,gfx950:sramecc+:xnack-,256
+    transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs
+    T,N,0,1,5120,32768,2880,1,2880,14745600,0,2880,94371840,5120,167772160,5120,167772160,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.54105e+06,1088.93,380.303,36094,gfx950:sramecc+:xnack-,256
diff --git a/examples/mlperf/llama2_70b/README.md b/examples/mlperf/llama2_70b/README.md
new file mode 100644
index 000000000..46ede190a
--- /dev/null
+++ b/examples/mlperf/llama2_70b/README.md
@@ -0,0 +1,377 @@
+# Llama2-70B LoRA MLPerf on MI355X (Primus)
+
+MLPerf Training 6.0 Llama2-70B LoRA on **MI355X** (8× GPU, 1 node) via Megatron-Bridge and `primus-cli`.
+
+Dataset: [GovReport](https://gov-report-data.github.io/) (SCROLLS `gov_report`), packed to **8192** tokens.
+Model: **meta-llama/Llama-2-70b-hf** with LoRA (rank 16, alpha 32).
+Precision: **MXFP4** + BF16; **FP8 delayed scaling** after healing at step 340.
+
+## Key files
+
+- `configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` — post-train overrides
+- `config_MI355X_1x8x1.sh` — system config and env vars (set `PRIMUS_PATH` to your Primus clone)
+- `run_with_docker.sh` — **recommended**: host-side Docker orchestration (N timed trials)
+- `runtime_tunables.sh` — host CPU/THP/cache tuning before each trial
+- `run_and_time.sh` — timed MLPerf run inside the container (called by `run_with_docker.sh`)
+- `a4w4_tuned_gemms.csv` — tuned AITER A4W4 GEMM configs
+
+---
+
+## Prerequisites
+
+- 8× MI355X GPUs on one node
+- Hugging Face access to `meta-llama/Llama-2-70b-hf` (`HF_TOKEN`)
+- ~300 GB disk for packed data + Megatron checkpoint
+- Docker with ROCm (`/dev/kfd`, `/dev/dri`)
+
+---
+
+## 1. Manual container (optional)
+
+Use this only if you are **not** using **`run_with_docker.sh`** and want an interactive shell inside the image.
+
+```bash
+docker pull rocm/primus:v26.5
+
+docker run -it \
+  --device=/dev/kfd \
+  --device=/dev/dri \
+  --security-opt seccomp=unconfined \
+  --group-add 44 \
+  --group-add 109 \
+  --cap-add=SYS_PTRACE \
+  --ipc=host \
+  --shm-size=32g \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  --memory=0 \
+  --memory-swap=0 \
+  --privileged \
+  --ulimit nofile=65535:65535 \
+  -v /path/to/Primus:/workspace/Primus \
+  -v /path/on/host/data:/data \
+  rocm/primus:v26.5
+```
+
+Repo is at **`/workspace/Primus`** inside the container. Then follow [§4](#4-run-training-inside-an-existing-container).
+
+## 2. Data layout (host `DATADIR` / container `/data`)
+
+```text
+${DATADIR}/
+├── train.npy
+├── validation.npy
+├── packed_metadata.jsonl
+├── megatron_checkpoints/Llama-2-70b-hf/   # checkpoint root (not iter_0000000/)
+└── .cache/huggingface/                    # optional HF cache
+```
+
+**`run_with_docker.sh`** sets **`PACKED_DATA_DIR=/data`**, **`PRETRAINED_CHECKPOINT=/data/megatron_checkpoints/Llama-2-70b-hf`**, and **`HF_HOME=/data/.cache/huggingface`** automatically. On first run, posttrain hooks populate missing data/checkpoints when **`HF_TOKEN`** is set.
+
+For a **manual** container (§1), export the same paths inside the shell before **`run_and_time.sh`**:
+
+```bash
+export HF_TOKEN=hf_...
+export PACKED_DATA_DIR=/data
+export PRETRAINED_CHECKPOINT=/data/megatron_checkpoints/Llama-2-70b-hf
+source /workspace/Primus/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
+bash /workspace/Primus/examples/mlperf/llama2_70b/run_and_time.sh
+```
+
+Checkpoint root must contain **`latest_train_state.pt`**, not only **`iter_0000000/`**.
+
+---
+
+## 3. Run with `run_with_docker.sh` (recommended)
+
+Use this when you want the **MLPerf submission-style flow** from the **host**: one command pulls up Docker, mounts data and results, applies host runtime tunables, runs **N** timed trials, and collects logs under **`LOGDIR`**. You do **not** need to `docker run -it` manually unless you are debugging inside the container (see [§4](#4-run-training-inside-an-existing-container)).
+
+### 3.1 Quick start (non-interactive)
+
+From your **Primus repo root** on an 8× MI355X node with Docker + ROCm devices:
+
+```bash
+cd /path/to/Primus
+
+export HF_TOKEN=hf_...                      # required: Llama-2-70b + GovReport hooks
+export DATADIR=/data/mlperf_llama2          # writable: dataset + checkpoint (host)
+export LOGDIR=/data/mlperf_llama2/results   # writable: trial logs + MLLOG (host)
+export CONT=rocm/primus:v26.5               # Primus ROCm image
+export DGXSYSTEM=MI355X_1x8x1               # selects config_${DGXSYSTEM}.sh
+export NEXP=1                               # use 10 for MLPerf submission runs
+
+# Optional: quiet stdout (:::MLLOG + timing banners only)
+export SUBMISSION_QUIET=1
+
+bash examples/mlperf/llama2_70b/run_with_docker.sh
+```
+
+`PRIMUS_HOST` defaults to the repo root (three levels above `examples/mlperf/llama2_70b/`). Override if your checkout is elsewhere:
+
+```bash
+export PRIMUS_HOST=/other/path/Primus
+bash examples/mlperf/llama2_70b/run_with_docker.sh
+```
+
+### 3.2 Interactive mode
+
+Prompts for **`HF_TOKEN`**, image, **`NEXP`**, **`DATADIR`**, **`LOGDIR`**, and **`DGXSYSTEM`** when values are missing:
+
+```bash
+cd /path/to/Primus
+export SUBMISSION_QUIET=1   # optional
+INTERACTIVE=1 bash examples/mlperf/llama2_70b/run_with_docker.sh
+```
+
+### 3.3 What the script does
+
+1. **Validates** `config_${DGXSYSTEM}.sh` exists and **`HF_TOKEN`** is set.
+2. **Removes** any existing container named **`CONT_NAME`** (default `mlperf_llama2_70b_lora_primus`), then starts a **detached** container (`sleep infinity`) with ROCm devices and mounts below.
+3. **Installs** editable Primus from the mount (`pip install -e /workspace/Primus --no-deps`); Python/torchrun come from the image venv **`/opt/venv/bin`** (not the host `PATH`).
+4. For each trial **`1..NEXP`**:
+   - Runs **`runtime_tunables.sh`** on the **host** (unless disabled).
+   - Optionally drops host page cache if **`CLEAR_CACHES=1`**.
+   - **`docker exec`** → **`bash /workspace/code/run_and_time.sh`** with a new **`SEED=$RANDOM`** and env from **`config_MI355X_1x8x1.sh`** (+ MLLOG paths under **`/results`**).
+   - Streams container stdout to **`${LOGDIR}/${DATESTAMP}_<trial>.log`** on the host.
+   - Copies **`mlperf_logging.out`** and **`train.mlperfposttrain.exp.log`** into **`${LOGDIR}/artifacts/`** after each successful trial.
+5. **Removes** the container on exit (success or failure).
+
+First run can take a long time: posttrain hooks may download the HF model, convert checkpoints, and pack GovReport under **`DATADIR`**. Later runs reuse **`${DATADIR}/train.npy`**, **`validation.npy`**, and **`megatron_checkpoints/`**.
+
+### 3.4 Volume and path map
+
+| Host | Container | Purpose |
+|------|-----------|---------|
+| **`PRIMUS_HOST`** (Primus repo) | `/workspace/Primus` | Code; editable install; `PRIMUS_PATH` |
+| **`examples/mlperf/llama2_70b/`** (this example dir) | `/workspace/code` | `run_and_time.sh`, `config_*.sh` |
+| **`DATADIR`** | `/data` | Packed `.npy`, HF cache, Megatron checkpoint |
+| **`LOGDIR`** | `/results` | MLLOG, primus-cli logs, timed run artifacts |
+
+Inside the container, **`run_and_time.sh`** uses **`/results`** for:
+
+| File | Description |
+|------|-------------|
+| `mlperf_logging.out` | `:::MLLOG` log (`ENABLE_MLLOG=1`) |
+| `train.mlperfposttrain.exp.log` | Full timed-run stdout |
+| `logs/log_*.txt` | `primus-cli direct` log |
+| `RESULT,LLAMA2_70B_LORA,,<sec>,AMD,<time>` | MLPerf wall-clock line |
+
+On the **host**, after trial *k*:
+
+| File | Description |
+|------|-------------|
+| `${LOGDIR}/${DATESTAMP}_k.log` | Tee of that trial’s container output |
+| `${LOGDIR}/artifacts/mlperf_logging_${DATESTAMP}_k.out` | Copy of MLLOG |
+| `${LOGDIR}/artifacts/train_${DATESTAMP}_k.log` | Copy of training log |
+
+### 3.5 Environment variables
+
+Variables you set most often:
+
+| Variable | Default | Meaning |
+|----------|---------|---------|
+| **`HF_TOKEN`** | *(none)* | **Required.** Hugging Face token for Llama-2-70b + dataset hooks |
+| **`DATADIR`** | `$HOME/data/mlperf_llama2` | Host data + checkpoint tree (mounted at `/data`) |
+| **`LOGDIR`** | `${DATADIR}/results` | Host results (mounted at `/results`) |
+| **`CONT`** | `rocm/primus:v26.5` | Docker image |
+| **`DGXSYSTEM`** | `MI355X_1x8x1` | Suffix for `config_${DGXSYSTEM}.sh` |
+| **`NEXP`** | `1` | Number of timed trials (submission often uses **10**) |
+| **`SUBMISSION_QUIET`** | `0` | `1` → sets `PRIMUS_LOG_SUPPRESSION=1`, `MLPERF_VERBOSE_LOGS=0`, `PRIMUS_LOG_GPU_MEM=0`, `VERBOSE_TRAINING_LOG=0` |
+| **`INTERACTIVE`** | `0` | `1` → prompt for missing settings |
+| **`PRIMUS_HOST`** | auto (repo root) | Host path to Primus checkout |
+
+Optional tuning:
+
+| Variable | Default | Meaning |
+|----------|---------|---------|
+| **`CONT_NAME`** | `mlperf_llama2_70b_lora_primus` | Docker container name |
+| **`RUN_RUNTIME_TUNABLES`** | `1` | `0` → skip host **`runtime_tunables.sh`** |
+| **`CLEAR_CACHES`** | `0` | `1` → extra host page-cache drop before each trial (needs sudo) |
+| **`CHECK_COMPLIANCE`** | `0` | `1` → run `mlperf_logging.compliance_checker` after each trial (non-blocking on failure) |
+| **`MLPERF_RULESET`** | `6.0.0` | Ruleset passed to compliance checker |
+| **`DATESTAMP`** | `date +%y%m%d%H%M%S` | Prefix for host trial log filenames |
+
+The script also sets **`PRIMUS_GPU_MODEL`** from **`DGXSYSTEM`** (e.g. `MI355X`) so **`primus-env.sh`** works when **`rocm-smi`** is missing inside **`docker exec`**.
+
+Config env from **`config_MI355X_1x8x1.sh`** (NCCL, MXFP4, MLLOG, 550 iters, etc.) is forwarded into the container; host **`PATH`** is **not** forwarded (container uses **`/opt/venv/bin`**).
+
+### 3.6 Submission-style run (10 trials, quiet)
+
+```bash
+export HF_TOKEN=hf_...
+export DATADIR=/data/mlperf_llama2
+export LOGDIR=/data/mlperf_llama2/results
+export NEXP=10
+export SUBMISSION_QUIET=1
+bash examples/mlperf/llama2_70b/run_with_docker.sh
+```
+
+Each trial gets a random **`SEED`**. Collect **`${LOGDIR}/${DATESTAMP}_*.log`** and **`${LOGDIR}/artifacts/`** for your submission bundle.
+
+### 3.7 Host runtime tunables
+
+Before each trial, **`runtime_tunables.sh`** (when **`RUN_RUNTIME_TUNABLES=1`**) adjusts cpupower governor, transparent huge pages, and may drop caches—matching the MLPerf reference host setup. Some steps need **root** (`sudo` for `cpupower`, `/proc/sys/vm/drop_caches`). If tunables fail on your cluster policy, set **`RUN_RUNTIME_TUNABLES=0`** and apply equivalent settings via your admin workflow.
+
+### 3.8 Troubleshooting (`run_with_docker.sh`)
+
+| Symptom | What to check |
+|---------|----------------|
+| `HF_TOKEN is required` | Export token or use **`INTERACTIVE=1`** |
+| `Config not found: config_*.sh` | Set **`DGXSYSTEM=MI355X_1x8x1`** (or add a matching config file) |
+| `torchrun: command not found` / `No module named 'yaml'` | Usually host **`PATH`** leaked into the container; use current script (container **`PATH=/opt/venv/bin:...`**) |
+| `rocm-smi not found` warning | Expected in minimal **`docker exec`**; set **`PRIMUS_GPU_MODEL=MI355X`** or rely on auto-map from **`DGXSYSTEM`** |
+| Permission denied under **`LOGDIR`** / **`DATADIR`** | Ensure host paths exist and are writable by your user |
+| Trial fails immediately | Inspect **`${LOGDIR}/${DATESTAMP}_1.log`** and **`/results/logs/log_*.txt`** inside the mount |
+| Stuck on first run | Hooks downloading/converting; verify **`HF_TOKEN`** and disk space (~300 GB under **`DATADIR`**) |
+
+To debug inside the same mounts without re-running the full script, start a shell in the running container (while a trial is not active):
+
+```bash
+docker exec -it mlperf_llama2_70b_lora_primus bash -lc 'source /opt/venv/bin/activate && cd /workspace/Primus && bash /workspace/code/run_and_time.sh'
+```
+
+(Use your **`CONT_NAME`** if overridden.)
+
+---
+
+## 4. Run training inside an existing container
+
+```bash
+export HF_TOKEN=hf_...
+source examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
+bash examples/mlperf/llama2_70b/run_and_time.sh
+```
+
+Timed run writes under **`/results`**:
+
+- `train.mlperfposttrain.exp.log` — full stdout from the timed wrapper
+- `logs/log_*.txt` — primus-cli direct log
+- `mlperf_logging.out` — `:::MLLOG` submission log (`ENABLE_MLLOG=1`)
+- `RESULT,LLAMA2_70B_LORA,,<seconds>,AMD,<start time>` — wall-clock line for MLPerf timing
+
+One-shot without sourcing config first (env + hooks still run):
+
+```bash
+export HF_TOKEN=hf_...
+bash examples/mlperf/llama2_70b/run_and_time.sh
+```
+
+Equivalent manual launch (same as `run_and_time.sh` without timing/`tee`):
+
+```bash
+cd /workspace/Primus   # or: cd "${PRIMUS_PATH}"
+
+export HF_TOKEN=hf_...
+source examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
+
+mkdir -p /results/logs
+cd "${PRIMUS_PATH}"
+
+./primus-cli direct \
+  --log_file "/results/logs/log_$(date +%Y%m%d_%H%M%S).txt" \
+  -- \
+  train posttrain \
+  --config "${EXP}"
+```
+
+Notes:
+
+- The **`--`** separates launcher flags from the Primus Python CLI (`train posttrain ...`).
+- Run from **`${PRIMUS_PATH}`** with **`--log_file` under `/results/logs`** so a read-only Primus bind mount does not fail on `logs/`.
+- `EXP` is set by `config_MI355X_1x8x1.sh`; override with `export EXP=...` if needed.
+
+---
+
+## 5. MLPerf experiment configuration
+
+### Config files
+
+| File | Role |
+|------|------|
+| `examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` | Post-train overrides |
+| `examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh` | MLPerf env (MXFP4, AITER, NCCL, MLLOG) |
+| `examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv` | Tuned AITER A4W4 GEMM configs |
+| `primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml` | Model recipe |
+| `primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py` | `llama2_70b_lora_mxfp4_config` |
+
+### Training schedule
+
+| Parameter | Value |
+|-----------|-------|
+| `train_iters` | 550 |
+| `global_batch_size` | 8 |
+| `micro_batch_size` | 1 |
+| `seq_length` | 8192 |
+| `lr` | 0.0006 |
+| `eval_interval` / `eval_iters` | 48 / 24 |
+| Quality target | eval loss **< 0.925** |
+
+### Precision
+
+**MXFP4 (steps 0–339):** `fp4=mxfp4`, `fp8=None`, `PRE_QUANTIZED_MODEL=True`, fused attention, AITER A4W4 GEMMs (`a4w4_tuned_gemms.csv`).
+
+**FP8 healing (step 340+):** `HEALING_ITER=340`, delayed scaling via `FP8_*` env vars in `config_MI355X_1x8x1.sh`.
+
+### LoRA
+
+Targets `linear_qkv`, `linear_proj` (dim 16, alpha 32). `stable_lora_with_te_op_fuser=True` (unfused `LoRALinear` adapters).
+
+### Parallelism
+
+TP=1, PP=1, CP=1, 8 GPUs data parallel.
+
+### MLPerf overrides (Primus-side, no third_party git patches)
+
+Runtime patches under `primus/backends/megatron_bridge/patches/mlperf_llama2_70b/` apply only when
+`llama2_70b_lora_mxfp4` / `llama2_70b_lora_mlperf_posttrain.yaml` is selected.
+
+Recipe code lives under `primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/`.
+
+| File | Role |
+|------|------|
+| `lora.py` | NeMo-stable LoRA (`use_te_fused_lora=False`) |
+| `resettable_data_iterator.py` | Deterministic validation iterator |
+| `bridge_patches.py` | Data loaders, eval reset, SFT mask cache, NeMo timing |
+| `megatron_patches.py` | MXFP4 recipe + optional TE SwiGLU |
+| `conditions.py` | Scopes patches to MLPerf Llama2-70B only |
+
+One-time cleanup if you previously applied git patches to submodules:
+
+```bash
+git -C third_party/Megatron-Bridge checkout -- .
+git -C third_party/Megatron-Bridge/3rdparty/Megatron-LM checkout -- .
+```
+
+---
+
+## 6. Logging
+
+Bring-up defaults (`config_MI355X_1x8x1.sh`): `log_interval=10`, `PRIMUS_LOG_GPU_MEM=1`, `VERBOSE_TRAINING_LOG=1`.
+
+MLPerf submission (quiet):
+
+```bash
+export PRIMUS_LOG_GPU_MEM=0
+export VERBOSE_TRAINING_LOG=0
+# Or use run_with_docker.sh: SUBMISSION_QUIET=1 (also sets PRIMUS_LOG_SUPPRESSION + MLPERF_VERBOSE_LOGS)
+# yaml: log_interval: 99999, stderr_sink_level: ERROR
+```
+
+### Common issues
+
+| Symptom | Fix |
+|---------|-----|
+| NCCL hang, 0% GPU | `NCCL_IB_DISABLE=1` (default in config) |
+| Invalid pretrained checkpoint | Point at checkpoint **root**, not `iter_0000000` |
+| Long silence at start | Pre-quantize + warmup + AITER JIT (normal) |
+
+---
+
+## 7. Optional overrides
+
+```bash
+export PRIMUS_TRAIN_ITERS=550
+export SEED=1234
+export SYNTH_WARMUP_STEPS=0
+export NCCL_IB_DISABLE=0    # if RDMA works on your system
+```
diff --git a/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv b/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv
new file mode 100644
index 000000000..8c4d18c33
--- /dev/null
+++ b/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv
@@ -0,0 +1,10 @@
+cu_num,M,N,K,kernelId,splitK,kernelName,errRatio
+256,8192,10240,8192,45,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_128x512E,0.0
+256,8192,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,8192,57344,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,8192,8192,28672,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,8192,8192,10240,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,8192,8192,57344,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,8192,28672,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,10240,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,57344,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
diff --git a/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh b/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
new file mode 100755
index 000000000..d29f3d838
--- /dev/null
+++ b/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
@@ -0,0 +1,172 @@
+#!/bin/bash
+# MLPerf 6.0 environment for Llama2-70B LoRA on MI355X (8 GPUs, 1 node).
+# Source before run_and_time.sh, then from ${PRIMUS_PATH}:
+#   ./primus-cli direct --log_file /results/logs/log_*.txt -- train posttrain --config "${EXP}"
+
+export DGXSYSTEM=MI355X_1x8x1
+# Used by runner/helpers/envs/primus-env.sh when rocm-smi is not in PATH (e.g. minimal Docker exec).
+export PRIMUS_GPU_MODEL="${PRIMUS_GPU_MODEL:-MI355X}"
+export GPUS_PER_NODE=8
+export NNODES=1
+export NODE_RANK=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29502
+
+# MLPerf timed runs set SEED=$RANDOM per trial; default here for single-shot primus-cli.
+export SEED="${SEED:-$RANDOM}"
+
+export PRIMUS_PATH="${PRIMUS_PATH:-/workspace/Primus}"
+export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-Bridge:${PYTHONPATH:-}"
+export EXP="${EXP:-${PRIMUS_PATH}/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml}"
+export DATA_PATH="${DATA_PATH:-/data}"
+
+export PACKED_TRAIN_DATA_PATH="${DATA_PATH}/train.npy"
+export PACKED_VAL_DATA_PATH="${DATA_PATH}/validation.npy"
+export PACKED_METADATA_PATH="${DATA_PATH}/packed_metadata.jsonl"
+export PACKED_DATA_DIR="${DATA_PATH}"
+
+export PRETRAINED_CHECKPOINT="${PRETRAINED_CHECKPOINT:-/data/megatron_checkpoints/Llama-2-70b-hf}"
+
+export LR=0.0006
+
+export HSA_NO_SCRATCH_RECLAIM=1
+export HSA_ENABLE_SDMA=1
+export HSA_ENABLE_INTERRUPT=0
+export GPU_MAX_HW_QUEUES=2
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HIGH_PRIORITY=1
+export NCCL_CHECKS_DISABLE=1
+# Single-node: use GPU P2P, not ionic IB. Broken libibverbs ABI (ionic kernel
+# mismatch) commonly hangs NCCL here with 0% GPU util. MLPerf systems with
+# working RDMA can override: NCCL_IB_DISABLE=0
+export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}"
+export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-lo}"
+export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-lo}"
+export OMP_NUM_THREADS=1
+
+export NVTE_USE_AITER_ROPE=1
+# Fused attention for MXFP4 training (must match recipe attention_backend=fused).
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=1
+export NVTE_UNFUSED_ATTN=0
+export NVTE_FUSED_ATTN_CK=1
+export NVTE_FUSED_ATTN_AOTRITON=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_USES_BWD_V3=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=0
+export NVTE_RS_STRIDED_ATOMIC=2
+export NVTE_FP8_DPA_BWD=1
+export NVTE_USE_HIPBLASLT=1
+export NVTE_USE_CAST_TRANSPOSE_TRITON=0
+export NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE=0
+export NVTE_USE_RMSNORM_TRITON=0
+export USE_TE_SWIGLU=1
+export ENABLE_TRANSPOSE_CACHE=0
+export NVTE_MXFP4_USE_HADAMARD=${NVTE_MXFP4_USE_HADAMARD:-1}
+export NVTE_DEBUG=0
+export NVTE_DEBUG_LEVEL=0
+
+export HEALING_ITER=${HEALING_ITER:-340}
+export HEALING_PRECISION=${HEALING_PRECISION:-FP8_DS}
+export PRE_QUANTIZED_MODEL=${PRE_QUANTIZED_MODEL:-True}
+export NCCL_MIN_P2P_NCHANNELS=${NCCL_MIN_P2P_NCHANNELS:-32}
+export NCCL_MIN_CTAS=${NCCL_MIN_CTAS:-32}
+export NCCL_NCHANNELS_PER_NET_PEER=${NCCL_NCHANNELS_PER_NET_PEER:-32}
+export NCCL_NVLS_ENABLE=${NCCL_NVLS_ENABLE:-0}
+
+export MEGATRON_BRIDGE_LOGGING_LEVEL=50
+export PYTHONWARNINGS=ignore
+export PRIMUS_LOG_LEVEL=ERROR
+
+# Submission-style logging: suppress framework noise; emit :::MLLOG on stdout.
+# Override for bring-up/debug: VERBOSE_TRAINING_LOG=1 PRIMUS_LOG_GPU_MEM=1 MLPERF_VERBOSE_LOGS=1
+export SUBMISSION_QUIET=1
+export PRIMUS_LOG_SUPPRESSION=1
+export MLPERF_VERBOSE_LOGS=0
+export VERBOSE_TRAINING_LOG=0
+export PRIMUS_LOG_GPU_MEM=0
+
+export SYNTH_WARMUP_STEPS=5
+export SYNTH_WARMUP_VALID_STEPS=5
+
+export ENABLE_MLLOG=1
+export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out
+export MLLOG_SAVE_TO_FILE=0
+export MLLOG_TRAIN_LOSS_LOG_FREQ=0
+export MLLOG_TARGET_EVAL_LOSS=0.925
+export TARGET_EVAL_LOSS=0.925
+export MLLOG_SUBMISSION_BENCHMARK=llama2_70b_lora
+export MLLOG_SUBMISSION_DIVISION=closed
+export MLLOG_SUBMISSION_ORG=AMD
+export MLLOG_SUBMISSION_PLATFORM=MI355X
+
+export MLLOG_TENSOR_PARALLELISM=1
+export MLLOG_PIPELINE_PARALLELISM=1
+export MLLOG_CONTEXT_PARALLELISM=1
+export MLLOG_EXPERT_PARALLELISM=1
+export MLLOG_MICRO_BATCH_SIZE=1
+MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}")
+export MLLOG_CONFIG_FILENAME
+export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR=mxfp4
+
+export TP_COMM_OVERLAP=False
+export MC_TP_OVERLAP_AG=False
+export MC_TP_OVERLAP_RS=False
+export MC_TP_OVERLAP_RS_DGRAD=False
+
+export CUBLAS_FORCE_XMMA_KERNEL_INIT=DEVICE
+
+export LORA_A2A=1
+export POSSIBLE_USER_WARNINGS=0
+export CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT=0
+
+export TP=1
+export PP=1
+export CP=1
+export SP=False
+export VBOOST_VALUE=1
+export MBS=1
+export MINIBS=1
+export SKIP_EVALS=3
+export VAL_CHECK_INTERVAL=384
+export HYDRA_FULL_ERROR=1
+
+export FP8_DPA=0
+# FP8 env flags below apply to healing (HEALING_ITER=340) and TE delayed scaling,
+# not Megatron model_cfg.fp8 during the MXFP4 phase (recipe sets fp8=None, fp4=mxfp4).
+export FP8=True
+export FP8_AMAX_ALGO=most_recent
+export FP8_REDUCE_AMAX=False
+export FP8_AMAX_HISTORY=4
+export FP8_ACTIVATION=True
+
+export FUSED_SOFTMAX=0
+export RMSNORM_CAST=0
+
+export PT_TENSOR_VALIDATION=0
+export PROFILE_RPD=0
+
+export USE_HIPBLASLT=1
+export TORCH_BLAS_PREFER_HIPBLASLT=1
+
+export LOGGING_INTERVAL=5000
+# FP4 weights (MXFP4 e2m1 linear layers); distinct from FP8_* healing flags above.
+export FP4=True
+export FP4_RECIPE=mxfp4
+export MAX_STEPS=550
+export NEXP=1
+
+export LOAD_CKPT=True
+export MCORE_CUDA_GRAPH=False
+export RESET_CG_AFTER_HEALING=False
+
+export RECOMPUTE_GRANULARITY=null
+export RECOMPUTE_METHOD=null
+export RECOMPUTE_NUM_LAYERS=null
+
+export FP8_ACT=0
+export AITER_CONFIG_GEMM_A4W4="${PRIMUS_PATH}/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv"
+export AITER_LOG_TUNED_CONFIG=0
+export NVTE_FP4_LOG_GEMM_SHAPES=0
+export AITER_LOG_LEVEL=ERROR
+export AITER_LOG_MORE=0
diff --git a/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml
new file mode 100644
index 000000000..474ccfbea
--- /dev/null
+++ b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml
@@ -0,0 +1,60 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:llama2_70b_lora_mlperf_posttrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+
+    # MXFP4 LoRA recipe (primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_70b_lora_mxfp4_config)
+    model: llama2_70b_lora_mxfp4.yaml
+
+    overrides:
+      pretrained_checkpoint: ${PRETRAINED_CHECKPOINT:/data/megatron_checkpoints/Llama-2-70b-hf}
+      stderr_sink_level: DEBUG
+      log_interval: 10
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      pipeline_dtype: null
+      virtual_pipeline_model_parallel_size: null
+      context_parallel_size: 1
+      sequence_parallel: false
+      use_megatron_fsdp: false
+
+      peft: lora
+      dataset_type: mlperf_dataset
+      packed_sequence: true
+      packed_train_data_path: ${PACKED_DATA_DIR:/data}/train.npy
+      packed_val_data_path: ${PACKED_DATA_DIR:/data}/validation.npy
+      packed_metadata_path: ${PACKED_DATA_DIR:/data}/packed_metadata.jsonl
+
+      train_iters: ${PRIMUS_TRAIN_ITERS:550}
+      global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:8}
+      micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:1}
+      seq_length: 8192
+      eval_interval: ${PRIMUS_EVAL_INTERVAL:48}
+      eval_iters: ${PRIMUS_EVAL_ITERS:24}
+      save_interval: null
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1.0e-8
+      weight_decay: 0.0001
+      min_lr: 0.0
+      lr: ${LR:0.0006}
+      lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:0}
+      lr_decay_iters: 550
+      clip_grad: 0.3
+
+      comm_overlap_config: null
+
+      seed: ${SEED:1234}
+
+      enable_primus_turbo: false
+      use_turbo_attention: false
+      use_turbo_rms_norm: false
+      use_turbo_gemm: false
+
+      check_for_nan_in_loss: false
diff --git a/examples/mlperf/llama2_70b/run_and_time.sh b/examples/mlperf/llama2_70b/run_and_time.sh
new file mode 100755
index 000000000..584b67530
--- /dev/null
+++ b/examples/mlperf/llama2_70b/run_and_time.sh
@@ -0,0 +1,103 @@
+#!/bin/bash
+# MLPerf Llama2-70B LoRA (Megatron-Bridge posttrain) — timed run with :::MLLOG output.
+#
+# Hooks run automatically before training:
+#   00_install_requirements.sh — Megatron-Bridge pip deps (mlperf-logging from image / requirements.txt)
+#   01_convert_checkpoints.sh    — HF → Megatron checkpoint (needs HF_TOKEN)
+#   02_prepare_mlperf_dataset.sh — SCROLLS gov-report .npy + metadata (needs HF_TOKEN)
+#
+# Usage (inside Primus container):
+#   export HF_TOKEN=...
+#   source examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh
+#   bash examples/mlperf/llama2_70b/run_and_time.sh
+#
+# Optional:
+#   PACKED_DATA_DIR=/data
+#   MLPERF_VERBOSE_LOGS=1
+#   PRIMUS_LOG_GPU_MEM=0
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+RESULTS_DIR=/results
+
+DATA_ROOT="${PACKED_DATA_DIR:-${DATA_PATH:-${PRIMUS_ROOT}/data/mlperf_llama2}}"
+export PACKED_DATA_DIR="${DATA_ROOT}"
+export DATA_PATH="${DATA_ROOT}"
+export HF_HOME="${HF_HOME:-${DATA_ROOT}/.cache/huggingface}"
+mkdir -p "${DATA_ROOT}" "${HF_HOME}"
+
+if [[ -z "${HF_TOKEN:-}" ]]; then
+    echo "[ERROR] HF_TOKEN is required (meta-llama/Llama-2-70b-hf + MLPerf dataset hub access)." >&2
+    exit 1
+fi
+export HF_TOKEN
+
+export SEED="${SEED:-$RANDOM}"
+
+# shellcheck disable=SC1091
+source "${SCRIPT_DIR}/config_MI355X_1x8x1.sh"
+
+# Create results directory (:::MLLOG + timed run logs)
+mkdir -p "${RESULTS_DIR}/logs"
+
+cd "${PRIMUS_PATH}"
+
+# Under multi-node SLURM (run_with_docker_slurm.sh), inherit rendezvous + node
+# sizing from SLURM env so we can scale to N nodes without editing the config
+# file. Single-node SLURM jobs (NNODES=1) fall through to the config defaults
+# so torchrun doesn't try to do c10d rdzv against MASTER_ADDR=localhost.
+if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then
+    NNODES="${SLURM_NNODES}"
+    NODE_RANK="${SLURM_NODEID:-0}"
+fi
+
+TRAIN_EXP_LOG="${RESULTS_DIR}/train.mlperfposttrain.exp.log"
+PRIMUS_CLI_LOG="${RESULTS_DIR}/logs/log_$(date +%Y%m%d_%H%M%S).txt"
+
+echo "============================================"
+echo "MLPerf Llama2-70B LoRA Post-Train (SFT)"
+echo "============================================"
+echo "Config:  ${EXP}"
+echo "Data:    ${DATA_PATH}"
+echo "MLLOG:   ${MLLOG_OUTPUT_FILE}"
+echo "Run log: ${TRAIN_EXP_LOG}"
+echo "GPUs:    ${GPUS_PER_NODE}"
+echo "Nodes:   ${NNODES}"
+echo "Rank:    ${NODE_RANK}"
+echo "Master:  ${MASTER_ADDR}:${MASTER_PORT}"
+echo "============================================"
+
+# Start timing
+start=$(date +%s)
+start_fmt=$(date +%Y-%m-%d\ %r)
+echo "STARTING TIMING RUN AT $start_fmt"
+
+# Launch through Primus CLI and keep the real exit code even though output is
+# piped through tee. --log_file keeps primus-cli from mkdir logs/ in cwd.
+set +e
+./primus-cli direct --log_file "${PRIMUS_CLI_LOG}" -- \
+    train posttrain \
+    --config "${EXP}" \
+    2>&1 | tee "${TRAIN_EXP_LOG}"
+ret_code=${PIPESTATUS[0]}
+set -e
+
+# End timing
+end=$(date +%s)
+end_fmt=$(date +%Y-%m-%d\ %r)
+echo "ENDING TIMING RUN AT $end_fmt"
+
+# Report result (wall-clock seconds for MLPerf timing scripts)
+result=$(( end - start ))
+result_name="LLAMA2_70B_LORA"
+echo "RESULT,$result_name,,$result,AMD,$start_fmt"
+
+if [[ $ret_code != 0 ]]; then
+    echo "Training failed with exit code: $ret_code"
+    exit "$ret_code"
+fi
+
+exit 0
diff --git a/examples/mlperf/llama2_70b/run_with_docker.sh b/examples/mlperf/llama2_70b/run_with_docker.sh
new file mode 100755
index 000000000..48498b59b
--- /dev/null
+++ b/examples/mlperf/llama2_70b/run_with_docker.sh
@@ -0,0 +1,279 @@
+#!/bin/bash
+# Host-side MLPerf Llama2-70B LoRA launcher (Primus + Docker).
+#
+# Orchestrates N timed training trials in a long-lived container, matching the
+# MLPerf submission flow in training_results_v6.0 (run_with_docker.sh).
+#
+# Prerequisites on the host:
+#   - Docker + ROCm devices (/dev/kfd, /dev/dri)
+#   - Primus repo (this script lives under examples/mlperf/llama2_70b/)
+#   - Writable DATADIR (dataset + checkpoint; hooks can populate on first run)
+#   - Writable LOGDIR (trial logs + /results artifacts inside the container)
+#
+# Quick start:
+#   export DATADIR=/data/mlperf_llama2
+#   export LOGDIR=/data/mlperf_llama2/results
+#   export CONT=rocm/primus:v26.4
+#   export DGXSYSTEM=MI355X_1x8x1
+#   bash examples/mlperf/llama2_70b/run_with_docker.sh
+#
+# Interactive (prompts for HF_TOKEN, image, NEXP, paths):
+#   INTERACTIVE=1 bash examples/mlperf/llama2_70b/run_with_docker.sh
+#
+# Quiet MLPerf submission logging (:::MLLOG + timing only):
+#   SUBMISSION_QUIET=1 bash examples/mlperf/llama2_70b/run_with_docker.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PRIMUS_HOST="${PRIMUS_HOST:-$(cd "${SCRIPT_DIR}/../../.." && pwd)}"
+
+# Defaults (override via env or INTERACTIVE prompts)
+: "${DGXSYSTEM:=MI355X_1x8x1}"
+: "${CONT:=rocm/primus:v26.4}"
+: "${DATADIR:=${HOME}/data/mlperf_llama2}"
+: "${LOGDIR:=${DATADIR}/results}"
+: "${NEXP:=1}"
+: "${CONT_NAME:=mlperf_llama2_70b_lora_primus}"
+: "${CLEAR_CACHES:=0}"
+: "${CHECK_COMPLIANCE:=0}"
+: "${MLPERF_RULESET:=6.0.0}"
+: "${DATESTAMP:=$(date +'%y%m%d%H%M%S')}"
+: "${INTERACTIVE:=0}"
+: "${SUBMISSION_QUIET:=0}"
+: "${RUN_RUNTIME_TUNABLES:=1}"
+
+# Map DGXSYSTEM (e.g. MI355X_1x8x1) -> PRIMUS_GPU_MODEL for primus-env when rocm-smi is unavailable in Docker.
+if [[ -z "${PRIMUS_GPU_MODEL:-}" && "${DGXSYSTEM}" =~ ^(MI[0-9]+[A-Z]*) ]]; then
+    export PRIMUS_GPU_MODEL="${BASH_REMATCH[1]}"
+fi
+
+readonly _runtime_tunables="${SCRIPT_DIR}/runtime_tunables.sh"
+
+prompt_if_interactive() {
+    local var_name="$1"
+    local prompt_text="$2"
+    local default_val="${3:-}"
+    local current_val="${!var_name:-}"
+
+    if [[ "${INTERACTIVE}" != "1" ]]; then
+        return 0
+    fi
+    if [[ -n "${current_val}" ]]; then
+        read -r -p "${prompt_text} [${current_val}]: " _input || true
+        if [[ -n "${_input}" ]]; then
+            printf -v "${var_name}" '%s' "${_input}"
+        fi
+    else
+        read -r -p "${prompt_text}${default_val:+ [${default_val}]}: " _input || true
+        if [[ -n "${_input}" ]]; then
+            printf -v "${var_name}" '%s' "${_input}"
+        elif [[ -n "${default_val}" ]]; then
+            printf -v "${var_name}" '%s' "${default_val}"
+        fi
+    fi
+}
+
+prompt_if_interactive "CONT" "Docker image (CONT)"
+prompt_if_interactive "NEXP" "Number of experiments (NEXP)" "1"
+prompt_if_interactive "DATADIR" "Data directory on host (mounted at /data)" "${DATADIR}"
+prompt_if_interactive "LOGDIR" "Log/results directory on host (mounted at /results)" "${LOGDIR}"
+prompt_if_interactive "DGXSYSTEM" "DGXSYSTEM (config file suffix)" "${DGXSYSTEM}"
+
+readonly _config_file="${SCRIPT_DIR}/config_${DGXSYSTEM}.sh"
+
+if [[ -z "${HF_TOKEN:-}" ]]; then
+    if [[ "${INTERACTIVE}" == "1" ]] || [[ -t 0 ]]; then
+        read -rsp "Hugging Face HF_TOKEN (required): " HF_TOKEN
+        echo
+        export HF_TOKEN
+    fi
+fi
+
+if [[ ! -f "${_config_file}" ]]; then
+    echo "[ERROR] Config not found: ${_config_file} (set DGXSYSTEM correctly)" >&2
+    exit 1
+fi
+
+if [[ -z "${HF_TOKEN:-}" ]]; then
+    echo "[ERROR] HF_TOKEN is required (meta-llama/Llama-2-70b-hf + MLPerf dataset)." >&2
+    exit 1
+fi
+
+if [[ ! -d "${PRIMUS_HOST}" ]]; then
+    echo "[ERROR] Primus repo not found: ${PRIMUS_HOST} (set PRIMUS_HOST)" >&2
+    exit 1
+fi
+
+mkdir -p "${DATADIR}" "${LOGDIR}" "${LOGDIR}/artifacts"
+
+readonly _logfile_base="${LOGDIR}/${DATESTAMP}"
+readonly _cont_name="${CONT_NAME}"
+readonly _primus_mount="/workspace/Primus"
+readonly _code_mount="/workspace/code"
+
+_cont_mounts=(
+    "--volume=${PRIMUS_HOST}:${_primus_mount}"
+    "--volume=${SCRIPT_DIR}:${_code_mount}"
+    "--volume=${DATADIR}:/data"
+    "--volume=${LOGDIR}:/results"
+)
+if [[ -d /opt/rocm ]]; then
+    _cont_mounts+=("--volume=/opt/rocm:/opt/rocm:ro")
+fi
+
+# Container Python/torchrun live in /opt/venv (rocm/primus image). Do not inject the host PATH.
+_container_path="/opt/venv/bin:/opt/rocm/bin:/usr/local/bin:/usr/bin:/bin"
+
+_extra_env=(
+    "--env=PATH=${_container_path}"
+    "--env=VENV_ACTIVATE=/opt/venv/bin/activate"
+    "--env=PRIMUS_AUTO_INSTALL=0"
+    "--env=PRIMUS_GPU_MODEL=${PRIMUS_GPU_MODEL:-MI355X}"
+    "--env=HF_TOKEN=${HF_TOKEN}"
+    "--env=PRIMUS_PATH=${_primus_mount}"
+    "--env=DATADIR=/data"
+    "--env=DATA_PATH=/data"
+    "--env=PACKED_DATA_DIR=/data"
+    "--env=PRETRAINED_CHECKPOINT=/data/megatron_checkpoints/Llama-2-70b-hf"
+    "--env=HF_HOME=/data/.cache/huggingface"
+    "--env=ENABLE_MLLOG=1"
+    "--env=DGXSYSTEM=${DGXSYSTEM}"
+    "--env=LOGDIR=/results"
+)
+
+if [[ "${SUBMISSION_QUIET}" == "1" ]]; then
+    _extra_env+=(
+        "--env=PRIMUS_LOG_SUPPRESSION=1"
+        "--env=MLPERF_VERBOSE_LOGS=0"
+        "--env=PRIMUS_LOG_GPU_MEM=0"
+        "--env=VERBOSE_TRAINING_LOG=0"
+    )
+fi
+
+# Export names from config_${DGXSYSTEM}.sh for docker exec (mirrors MLPerf reference).
+mapfile -t _config_env < <(
+    env -i HOME="${HOME}" PATH="${PATH}" bash -c "
+        set -a
+        source '${_config_file}'
+        set +a
+        compgen -e
+    " | grep -E -v '^(PWD|SHLVL|_|OLDPWD)$' || true
+)
+_config_docker_env=()
+for v in "${_config_env[@]}"; do
+    # Never override container PATH/PYTHONHOME from a host-side config source.
+    case "${v}" in PATH=*|PYTHONHOME=* ) continue ;; esac
+    _config_docker_env+=("--env=${v}")
+done
+
+_base_exec_env=("${_config_docker_env[@]}" "${_extra_env[@]}")
+
+cleanup_docker() {
+    if docker ps -a --format '{{.Names}}' | grep -qx "${_cont_name}"; then
+        docker container rm -f "${_cont_name}" >/dev/null 2>&1 || true
+    fi
+}
+cleanup_docker
+trap 'cleanup_docker' EXIT
+
+echo "[INFO] Primus host path:  ${PRIMUS_HOST}"
+echo "[INFO] Docker image:      ${CONT}"
+echo "[INFO] Data mount:        ${DATADIR} -> /data"
+echo "[INFO] Results mount:     ${LOGDIR} -> /results"
+echo "[INFO] Experiments:       ${NEXP}"
+echo "[INFO] Container name:    ${_cont_name}"
+
+if [[ "${DGXSYSTEM}" == MI* ]]; then
+    docker run --rm --init --detach \
+        --net=host --uts=host --ipc=host \
+        --device=/dev/dri --device=/dev/kfd \
+        --security-opt=seccomp=unconfined \
+        --group-add=video \
+        --cap-add=SYS_PTRACE \
+        --privileged \
+        --shm-size=32g \
+        --ulimit=memlock=-1 \
+        --ulimit=stack=67108864 \
+        --name="${_cont_name}" "${_cont_mounts[@]}" \
+        -e IMAGE_NAME="${CONT}" \
+        "${CONT}" sleep infinity
+else
+    docker run --rm --init --detach \
+        --net=host --uts=host --ipc=host \
+        --gpus=all \
+        --ulimit=memlock=-1 \
+        --ulimit=stack=67108864 \
+        --security-opt=seccomp=unconfined \
+        --name="${_cont_name}" "${_cont_mounts[@]}" \
+        -e IMAGE_NAME="${CONT}" \
+        "${CONT}" sleep infinity
+fi
+
+sleep 3
+docker exec "${_cont_name}" true
+
+# Editable Primus from the host mount (ML deps including mlperf-logging come from the image venv).
+_setup_env=("${_base_exec_env[@]}")
+docker exec "${_setup_env[@]}" "${_cont_name}" bash -lc "
+set -e
+source /opt/venv/bin/activate
+pip install -q -e '${_primus_mount}' --no-deps 2>/dev/null || true
+echo '[INFO] Container Primus editable install complete.'
+"
+
+_run_cmd="bash ${_code_mount}/run_and_time.sh"
+
+for _experiment_index in $(seq 1 "${NEXP}"); do
+    echo "============================================"
+    echo "Beginning trial ${_experiment_index} of ${NEXP}"
+    echo "============================================"
+
+    if [[ "${RUN_RUNTIME_TUNABLES}" == "1" && -x "${_runtime_tunables}" ]]; then
+        bash "${_runtime_tunables}"
+    fi
+
+    if [[ "${CLEAR_CACHES}" == "1" ]]; then
+        echo "[INFO] Clearing host page cache before trial ${_experiment_index}..."
+        sync
+        if [[ -w /proc/sys/vm/drop_caches ]]; then
+            echo 3 | sudo tee /proc/sys/vm/drop_caches >/dev/null || true
+        else
+            sudo sysctl vm.drop_caches=3 >/dev/null 2>&1 || true
+        fi
+    fi
+
+    _run_env=("${_base_exec_env[@]}" --env="SEED=${RANDOM}")
+
+    set +e
+    docker exec "${_run_env[@]}" "${_cont_name}" bash -lc "${_run_cmd}" \
+        2>&1 | grep --line-buffered -v "connected peer ranks" \
+        | tee "${_logfile_base}_${_experiment_index}.log"
+    _trial_exit=${PIPESTATUS[0]}
+    set -e
+
+    if [[ "${_trial_exit}" -ne 0 ]]; then
+        echo "[ERROR] Trial ${_experiment_index} failed with exit code ${_trial_exit}" >&2
+        exit "${_trial_exit}"
+    fi
+
+    if [[ -f "${LOGDIR}/mlperf_logging.out" ]]; then
+        cp -f "${LOGDIR}/mlperf_logging.out" "${LOGDIR}/artifacts/mlperf_logging_${DATESTAMP}_${_experiment_index}.out" || true
+    fi
+    if [[ -f "${LOGDIR}/train.mlperfposttrain.exp.log" ]]; then
+        cp -f "${LOGDIR}/train.mlperfposttrain.exp.log" \
+            "${LOGDIR}/artifacts/train_${DATESTAMP}_${_experiment_index}.log" || true
+    fi
+
+    if [[ "${CHECK_COMPLIANCE}" == "1" ]]; then
+        docker exec "${_run_env[@]}" "${_cont_name}" bash -lc \
+            "python3 -m mlperf_logging.compliance_checker --usage training \
+             --ruleset '${MLPERF_RULESET}' \
+             --log_output '/results/compliance_${DATESTAMP}_${_experiment_index}.out' \
+             '/results/train.mlperfposttrain.exp.log'" \
+            || echo "[WARN] Compliance check failed for trial ${_experiment_index} (non-blocking)"
+    fi
+done
+
+echo "[OK] Completed ${NEXP} trial(s). Logs: ${_logfile_base}_*.log"
+echo "[OK] MLLOG (last trial): ${LOGDIR}/mlperf_logging.out"
diff --git a/examples/mlperf/llama2_70b/runtime_tunables.sh b/examples/mlperf/llama2_70b/runtime_tunables.sh
new file mode 100755
index 000000000..f51f85cae
--- /dev/null
+++ b/examples/mlperf/llama2_70b/runtime_tunables.sh
@@ -0,0 +1,51 @@
+#!/bin/bash
+# MLPerf host runtime tunables (MI355X submission reference).
+# Run on the **host** before each timed trial (see run_with_docker.sh).
+# Steps are best-effort; missing tools or permissions are ignored.
+
+set +e
+
+echo "[runtime_tunables] Applying host performance settings..."
+
+if [[ -w /proc/sys/vm/drop_caches ]] 2>/dev/null; then
+    sync
+    echo 3 >/proc/sys/vm/drop_caches 2>/dev/null
+else
+    sync
+    sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches' 2>/dev/null
+fi
+
+sudo modprobe amdgpu 2>/dev/null
+if command -v rocm_agent_enumerator &>/dev/null; then
+    rocm_agent_enumerator >/dev/null 2>&1
+fi
+
+if command -v cpupower &>/dev/null; then
+    sudo cpupower idle-set -d 2 >/dev/null 2>&1
+    sudo cpupower frequency-set -g performance >/dev/null 2>&1
+fi
+
+for _sysctl in \
+    /proc/sys/kernel/nmi_watchdog \
+    /proc/sys/kernel/numa_balancing \
+    /proc/sys/kernel/randomize_va_space
+do
+    if [[ -w "${_sysctl}" ]] 2>/dev/null; then
+        echo 0 >"${_sysctl}" 2>/dev/null
+    else
+        sudo sh -c "echo 0 > ${_sysctl}" 2>/dev/null
+    fi
+done
+
+for _thp in \
+    /sys/kernel/mm/transparent_hugepage/enabled \
+    /sys/kernel/mm/transparent_hugepage/defrag
+do
+    if [[ -w "${_thp}" ]] 2>/dev/null; then
+        echo always >"${_thp}" 2>/dev/null
+    else
+        sudo sh -c "echo always > ${_thp}" 2>/dev/null
+    fi
+done
+
+echo "[runtime_tunables] Done."
diff --git a/examples/mlperf/llama3.1_8b/README.md b/examples/mlperf/llama3.1_8b/README.md
new file mode 100644
index 000000000..4a9289941
--- /dev/null
+++ b/examples/mlperf/llama3.1_8b/README.md
@@ -0,0 +1,53 @@
+# LLama3.1 8B MLPerf Pretraining
+
+MLPerf-compliant LLama3.1 8B pretraining using Primus
+
+## Setup
+
+### Start Docker Image
+
+```bash
+export MLPERF_PAT=<your_github_pat>
+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
+
+cd /workspace/Primus
+```
+
+
+### Configuration
+
+- **Model**: LLama3.1 8B (4096 hidden, 32 layers, 32 attention heads)
+- **Training**: 1.2M iterations, GBS=32, MBS=2, LR=8e-4
+- **Precision**: MXFP4
+- **Data**: C4 dataset (tokenized)
+
+## Key Files
+
+- `configs/MI355X/llama3.1_8B-pretrain-FP4.yaml` - Model and training config
+    - Update `train_data_path` and `train_data_path` to your local downloaded location
+- `config_MI355X_1x8x1.sh` - System config and env vars
+   - Update `PRIMUS_PATH` to clone Primus Repo
+   - Update `EXP`to `<PRIMUS_PATH>/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml`
+- `run_and_time.sh` - Run script
+
+### Data
+
+Download preprocessed C4 dataset:
+
+```bash
+mkdir -p /data/mlperf_llama31_8b
+cd /data/mlperf_llama31_8b
+bash <(curl -s https://raw.githubusercontent.com/mlcommons/r2-downloader/refs/heads/main/mlc-r2-downloader.sh) \
+    -d data https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri
+```
+
+### How to run
+
+```bash
+export HF_TOKEN=<your_huggingface_token>
+source config_MI355X_1x8x1.sh
+bash run_and_time.sh
+```
+## Notes
+
+- `log_interval: 99999999` suppresses regular Primus logs
diff --git a/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv b/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv
new file mode 100644
index 000000000..5cf68d18b
--- /dev/null
+++ b/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv
@@ -0,0 +1,12 @@
+cu_num,M,N,K,kernelId,splitK,kernelName,errRatio
+256,4096,4096,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,6144,4096,16384,50,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_192x256E,0.0
+256,16384,4096,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,16384,4096,6144,45,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_128x512E,0.0
+256,16384,4096,14336,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,16384,4096,28672,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,28672,4096,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,16384,6144,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,4096,14336,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,16384,14336,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
+256,16384,28672,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0
diff --git a/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh b/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh
new file mode 100755
index 000000000..f56e999d1
--- /dev/null
+++ b/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh
@@ -0,0 +1,115 @@
+#!/bin/bash
+# MLPerf LLama3.1 8B Configuration for MI355X (1x8x1)
+
+export DGXSYSTEM=MI355X_1x8x1
+export GPUS_PER_NODE=8
+export NNODES=1
+export NODE_RANK=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29502
+
+export PRIMUS_PATH=/workspace/Primus
+export PRIMUS_MLPERF=1
+export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-LM:${PYTHONPATH}"
+export EXP=${PRIMUS_PATH}/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml
+export DATA_PATH=/data
+
+export PRIMUS_MICRO_BATCH_SIZE=2
+export PRIMUS_GLOBAL_BATCH_SIZE=32
+export PRIMUS_LR=8e-4
+export PRIMUS_MIN_LR=8e-5
+export PRIMUS_TRAIN_ITERS=1200000
+export PRIMUS_LR_WARMUP_ITERS=64
+export EVAL_SAMPLES_INTERVAL=12288
+export PRIMUS_EVAL_INTERVAL=$((EVAL_SAMPLES_INTERVAL / PRIMUS_GLOBAL_BATCH_SIZE))  # Auto-computed
+
+export HSA_ENABLE_INTERRUPT=0
+export HSA_TOOLS_LIB=/opt/rocm/lib/libroctracer64.so
+export PRIMUS_APPLY_ROPE_FUSION=True
+export PRIMUS_FP8_RECIPE=hybrid
+
+export HSA_NO_SCRATCH_RECLAIM=1
+export HSA_ENABLE_SDMA=1
+export GPU_MAX_HW_QUEUES=2
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+export NVTE_FUSED_ATTN=1
+export NVTE_FUSED_ATTN_CK=1
+export NVTE_FUSED_ATTN_AOTRITON=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_USES_BWD_V3=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=0
+export NVTE_USE_AITER_ROPE=1
+export NVTE_MXFP4_USE_HADAMARD=1
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=1
+export NVTE_USE_CAST_TRANSPOSE_TRITON=0
+export NVTE_ASYNC_AMAX_REDUCTION=1
+export NVTE_DP_AMAX_REDUCE_INTERVAL=0
+export NVTE_USE_RMSNORM_TRITON=0
+export NVTE_LOG_CK_CONFIG=0
+export NVTE_LOG_FUSED_ATTN_CONFIG=0
+export USE_TE_SWIGLU=1
+
+export ENABLE_TRANSPOSE_CACHE=1
+export CK_FUSED_ATTN_LOG_CONFIG=0
+export CHECK_FOR_NAN_IN_GRAD=0
+
+export TOKENIZERS_PARALLELISM=false
+export NCCL_CHECKS_DISABLE=1
+export TORCH_NCCL_HIGH_PRIORITY=1
+
+export AITER_CONFIG_GEMM_A4W4=${PRIMUS_PATH}/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv
+export AITER_LOG_TUNED_CONFIG=0
+export LOG_AITER_GEMMS=0
+export AITER_LOG_LEVEL=ERROR
+export AITER_LOG_MORE=0
+
+export ENABLE_MLLOG=1
+export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out
+export MLLOG_TRAIN_LOSS_LOG_FREQ=0
+export MLLOG_TARGET_EVAL_LOSS=3.3
+export TARGET_EVAL_LOSS=3.3
+export MLLOG_SUBMISSION_BENCHMARK=llama31_8b
+export MLLOG_SUBMISSION_DIVISION=closed
+export MLLOG_SUBMISSION_ORG=AMD
+export MLLOG_SUBMISSION_PLATFORM=MI355X
+export MLLOG_TENSOR_PARALLELISM=1
+export MLLOG_PIPELINE_PARALLELISM=1
+export MLLOG_CONTEXT_PARALLELISM=1
+export MLLOG_EXPERT_PARALLELISM=1
+export MLLOG_MICRO_BATCH_SIZE=2
+MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}")
+export MLLOG_CONFIG_FILENAME
+export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR='mxfp4'
+
+export FP4=true
+export FP4_RECIPE=mxfp4
+
+export USE_HIPBLASLT=1
+export TORCH_BLAS_PREFER_HIPBLASLT=1
+
+export SYNTH_WARMUP_STEPS=20
+export SYNTH_WARMUP_VALID_STEPS=10
+export WARMUP_RECIPE=fp8_hybrid
+export SYNTH_WARMUP_EMPTY_CACHE=1
+export MLPERF_VERBOSE_LOGS=${MLPERF_VERBOSE_LOGS:-0}
+
+export NCCL_MIN_P2P_NCHANNELS=32
+export NCCL_MIN_CTAS=32
+export NCCL_NCHANNELS_PER_NET_PEER=32
+export NCCL_NVLS_ENABLE=0
+export NCCL_P2P_LEVEL=5
+export NCCL_SINGLE_RING_THRESHOLD=0
+export NCCL_BUFFSIZE=2097152
+
+export TP_COMM_OVERLAP=False
+export MC_TP_OVERLAP_AG=False
+export MC_TP_OVERLAP_RS=False
+export MC_TP_OVERLAP_RS_DGRAD=False
+
+export HIP_FORCE_DEV_KERNARG=1
+export HSA_ENABLE_SDMA_OPTIMIZATIONS=1
+export PYTORCH_ROC_ALLOC_CONF=expandable_segments:True
+export HIP_API_BLOCKING=0
+export PYTORCH_NO_CUDA_MEMORY_CACHING=0
diff --git a/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml
new file mode 100644
index 000000000..5679cac48
--- /dev/null
+++ b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml
@@ -0,0 +1,113 @@
+work_group: ${TEAM:amd}
+user_name: ${USER:root}
+exp_name: ${EXP_NAME:llama3.1_8B-pretrain-mlperf}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: llama3.1_8B.yaml
+    overrides:
+      # --- Logging Config ---
+      stage: mlperf_pretrain
+      disable_wandb: true
+      disable_tensorboard: true
+      stderr_sink_level: ERROR
+      log_interval: 999999
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+      eval_iters: ${PRIMUS_EVAL_ITERS:32}  # 32 * GBS = 1024 eval samples
+      eval_interval: ${PRIMUS_EVAL_INTERVAL:10} # 10 * GBS = 320 eval samples perform evaluation.
+
+      # --- Training Config ---
+      train_iters: ${PRIMUS_TRAIN_ITERS:200}
+      micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2}                # grad_acc = global_batch_size / (micro_batch_size * num_gpus) = 32 / (2 * 8) = 2
+      global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:32}
+      seq_length: 8192
+      max_position_embeddings: 8192
+      seed: ${SEED:1234}
+      lr: ${PRIMUS_LR:0.0008}          # 8e-4
+      min_lr: ${PRIMUS_MIN_LR:0.00008}     # 10% of lr
+      lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:64}
+      lr_decay_iters: 1199936
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+      init_method_std: 0.02
+      norm_epsilon: 1.0e-5
+      adam_eps: 1.0e-5
+      check_for_nan_in_loss_and_grad: false # default true
+      check_for_spiky_loss: false # default false, but setting it here explicitly
+      check_for_large_grads: false # default false, but setting it here explicitly
+
+      # --- Model Parallel Config ---
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: true
+      masked_softmax_fusion: true
+      ddp_average_in_collective: true
+      fused_single_qkv_rope: true
+
+      # --- Data Config ---
+      mock_data: false
+      train_data_path: "/data/c4-train.en_6_text_document"
+      valid_data_path: "/data/c4-validation-91205-samples.en_text_document"
+      test_data_path: null
+      mmap_bin_files: true
+      data_cache_path: null
+
+      # --- Profiling Config ---
+      profile: false
+      use_pytorch_profiler: true
+      profile_ranks: [0]           # Only profile rank 0 to save disk space
+      profile_step_start: 8        # Start after warmup (step 8)
+      profile_step_end: 13        # Profile 5 iterations (8-12)
+      disable_profiler_activity_cpu: false     # GPU kernels only (smaller files)
+      torch_profiler_record_shapes: false     # Disable for smaller traces
+      torch_profiler_with_stack: false        # Disable for smaller traces
+      torch_profiler_use_gzip: true           # Compress output
+
+      # --- Checkpointing Config ---
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 2000000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch
+
+      # --- FSDP Config ---
+      use_torch_fsdp2: false
+      use_distributed_optimizer: true # this is needed for fsdp2
+
+      # Cross entropy flags
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
+
+      # --- Mixed Precision Config ---
+      fp4: e2m1
+      fp4_recipe: mxfp4
+      accumulate_allreduce_grads_in_fp32: false
+      grad_reduce_in_bf16: true
+      attention_softmax_in_fp32: false
+
+
+      # --- Primus Turbo Config ---
+      enable_primus_turbo: false
+      use_turbo_attention: false
+      use_turbo_gemm: false              # can't use together with delayed recipe
+      use_turbo_grouped_gemm: false
+      moe_use_fused_router_with_aux_score: false
+      enable_turbo_attention_float8 : false
diff --git a/examples/mlperf/llama3.1_8b/run_and_time.sh b/examples/mlperf/llama3.1_8b/run_and_time.sh
new file mode 100755
index 000000000..d9c11f583
--- /dev/null
+++ b/examples/mlperf/llama3.1_8b/run_and_time.sh
@@ -0,0 +1,60 @@
+#!/bin/bash
+
+set -e
+
+# Create results directory
+mkdir -p /results
+
+cd "${PRIMUS_PATH}/examples/mlperf/llama3.1_8b"
+
+# Under multi-node SLURM (run_with_docker_slurm.sh), inherit rendezvous + node
+# sizing from SLURM env so we can scale to N nodes without editing the config
+# file. Single-node SLURM jobs (NNODES=1) fall through to the config defaults
+# so torchrun doesn't try to do c10d rdzv against MASTER_ADDR=localhost.
+if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then
+    NNODES="${SLURM_NNODES}"
+    NODE_RANK="${SLURM_NODEID:-0}"
+fi
+
+echo "============================================"
+echo "MLPerf LLama3.1 8B Training"
+echo "============================================"
+echo "Config: ${EXP}"
+echo "Data:   ${DATA_PATH}"
+echo "GPUs:   ${GPUS_PER_NODE}"
+echo "Nodes:  ${NNODES}"
+echo "Rank:   ${NODE_RANK}"
+echo "Master: ${MASTER_ADDR}:${MASTER_PORT}"
+echo "============================================"
+
+# Start timing
+start=$(date +%s)
+start_fmt=$(date +%Y-%m-%d\ %r)
+echo "STARTING TIMING RUN AT $start_fmt"
+
+# Launch through Primus CLI and keep the real exit code even though output is
+# piped through tee.
+set +e
+"${PRIMUS_PATH}/primus-cli" direct -- \
+    train pretrain \
+    --config "${EXP}" \
+    2>&1 | tee train.mlperfpretrain.exp.log
+ret_code=${PIPESTATUS[0]}
+set -e
+
+# End timing
+end=$(date +%s)
+end_fmt=$(date +%Y-%m-%d\ %r)
+echo "ENDING TIMING RUN AT $end_fmt"
+
+# Report result
+result=$(( end - start ))
+result_name="LLAMA3.1_8B"
+echo "RESULT,$result_name,,$result,AMD,$start_fmt"
+
+if [[ $ret_code != 0 ]]; then
+    echo "Training failed with exit code: $ret_code"
+    exit "$ret_code"
+fi
+
+exit 0
diff --git a/examples/moe_package/run_glm5_4layers_proxy.sh b/examples/moe_package/run_glm5_4layers_proxy.sh
index 1d3d26e0a..c491bb248 100644
--- a/examples/moe_package/run_glm5_4layers_proxy.sh
+++ b/examples/moe_package/run_glm5_4layers_proxy.sh
@@ -18,7 +18,7 @@ export NNODES=${NNODES:-1}
 export TRAIN_ITERS=10
 export SLURM_TIME=01:00:00
 export SLURM_PARTITION=amd-aig
-# export SLURM_NODELIST="uswslocpm2m-106-079"
+# export SLURM_NODELIST="<node>"  # pin the run to specific nodes
 
 # export NCCL_DEBUG=INFO
 export USING_AINIC=1
diff --git a/examples/moe_package/start_training_qwen_30B_a3B.sh b/examples/moe_package/start_training_qwen_30B_a3B.sh
index 5cd32e797..ca88f42b2 100644
--- a/examples/moe_package/start_training_qwen_30B_a3B.sh
+++ b/examples/moe_package/start_training_qwen_30B_a3B.sh
@@ -12,7 +12,7 @@ export NNODES=1
 export TRAIN_ITERS=10
 export SLURM_TIME=48:00:00
 export SLURM_PARTITION=amd-aig-2
-# export SLURM_NODELIST="uswslocpm2m-106-1962"
+# export SLURM_NODELIST="<node>"  # pin the run to specific nodes
 # export NCCL_DEBUG=INFO
 export USING_AINIC=1
 # export NCCL_IB_HCA="ionic_0:1,ionic_2:1,ionic_3:1,ionic_4:1,ionic_5:1,ionic_7:1,ionic_8:1,ionic_9:1"
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/_flux_1024_ddp.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_1024_ddp.yaml
new file mode 100644
index 000000000..168386b31
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_1024_ddp.yaml
@@ -0,0 +1,26 @@
+# SCRATCH (Phase-2 sharding sweep, delete after) — FLUX.1-dev 1024 hero DDP point.
+# Uses _flux_ddp_pre_trainer.yaml (ddp: block instead of fsdp:) so the recipe
+# takes the DistributedDataParallel path (full 12B replication, no sharding).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:_flux_1024_ddp}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: _flux_ddp_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 20
+        local_batch_size: 2
+        global_batch_size: 16
+        log_remote_every_steps: 1
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_1024_hero
+          base_resolution: [1024, 1024]
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/_flux_512_ab.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_512_ab.yaml
new file mode 100644
index 000000000..b0a12b41b
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_512_ab.yaml
@@ -0,0 +1,27 @@
+# TEMP scratch exp (delete after) — FLUX 512 A/B on a single fixed bucket (384x640).
+# Run once baseline, once with PRIMUS_ATTN_NONDETERMINISTIC=1.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:_flux_512_ab}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 60
+        local_batch_size: 1
+        global_batch_size: 8
+        log_remote_every_steps: 5
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_512_hero
+          base_resolution: [384, 640]
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/_flux_ac_smoke.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_ac_smoke.yaml
new file mode 100644
index 000000000..48016e7e4
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_ac_smoke.yaml
@@ -0,0 +1,28 @@
+# TEMP scratch exp (delete after) — FLUX 1024 hero AC smoke (832x1216).
+# Run once with PRIMUS_FLUX_REAL_AC=0 (silent no-op) and once with =1 (real AC);
+# compare peak mem at fixed mbs1 to confirm the strategy actually checkpoints.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:_flux_ac_smoke}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 12
+        local_batch_size: 1
+        global_batch_size: 8
+        log_remote_every_steps: 2
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_1024_hero
+          base_resolution: [832, 1216]
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/_flux_hero_ab.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_hero_ab.yaml
new file mode 100644
index 000000000..5a9023682
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/_flux_hero_ab.yaml
@@ -0,0 +1,27 @@
+# TEMP scratch exp (delete after) — FLUX 1024 hero A/B on a single fixed bucket
+# (832x1216, seqlen 4464). Run once baseline, once with PRIMUS_ATTN_NONDETERMINISTIC=1.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:_flux_hero_ab}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 60
+        local_batch_size: 1
+        global_batch_size: 8
+        log_remote_every_steps: 5
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_1024_hero
+          base_resolution: [832, 1216]
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/_wan_720p_ddp.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/_wan_720p_ddp.yaml
new file mode 100644
index 000000000..3ebcba8e1
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/_wan_720p_ddp.yaml
@@ -0,0 +1,25 @@
+# SCRATCH (Phase-2 sharding sweep, delete after) — Wan2.2 720p DDP point.
+# Uses the _ddp_pre_trainer.yaml module (ddp: block instead of fsdp:) so the
+# recipe takes the DistributedDataParallel path. Expected to OOM at A14B.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:_wan_720p_ddp}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: _ddp_pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 20
+        local_batch_size: 1
+        global_batch_size: 16
+        log_remote_every_steps: 1
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_720p
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain-1024-hero.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain-1024-hero.yaml
new file mode 100644
index 000000000..c0e0ee9ce
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain-1024-hero.yaml
@@ -0,0 +1,46 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 1024x1024 HERO.
+# =============================================================================
+# The headline attention A/B point: at 1024x1024 the FLUX joint attention runs on
+# 4096 image + 512 text = 4608 tokens, so attention is a large fraction of the
+# step and the nondeterministic bf16 flash-attention backward
+# (PRIMUS_ATTN_NONDETERMINISTIC=1) shows its throughput uplift. Activation
+# checkpointing is ON — the long-sequence hero needs it.
+#
+# Data cache (1024x1024):
+#   python -m tools.diffusion.preprocessing_multiprocess image \
+#     --image_dir /dataset/pseudo-camera-10k/train \
+#     --output_dir /dataset/pcam_flux_1024 \
+#     --processor flux --model_name /models/FLUX.1-dev --resolution 1024
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain-1024-hero.yaml \
+#   bash examples/run_pretrain.sh
+#
+# A/B: run once as-is (baseline) and once with PRIMUS_ATTN_NONDETERMINISTIC=1 (best).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_1_dev-pretrain-1024-hero}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 40
+        local_batch_size: 1
+        global_batch_size: 8     # 1 x dp8
+        log_remote_every_steps: 5
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pcam_flux_1024
+          base_resolution: [1024, 1024]
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain.yaml
new file mode 100644
index 000000000..01c5b000c
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain.yaml
@@ -0,0 +1,46 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 512x512.
+# =============================================================================
+# Config-only FLUX integration: runs through the exact same Primus nemo_automodel
+# path as Wan2.2 (examples/run_pretrain.sh). The only FLUX-specific pieces are the
+# model preset (flux_1_dev.yaml) and the image module config (flux_pre_trainer.yaml);
+# everything else (env/data/steps) is set here, just like the Wan exps.
+#
+# Data cache (built once; image latents + T5/CLIP embeds, no encoders at train time):
+#   python -m tools.diffusion.preprocessing_multiprocess image \
+#     --image_dir /dataset/pseudo-camera-10k/train \
+#     --output_dir /dataset/pcam_flux_512 \
+#     --processor flux --model_name /models/FLUX.1-dev --resolution 512
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/flux_1_dev-pretrain.yaml \
+#   bash examples/run_pretrain.sh
+#
+# Attention A/B: add PRIMUS_ATTN_NONDETERMINISTIC=1 to the launch env to enable the
+# nondeterministic bf16 flash-attention backward (numerically equivalent, faster).
+# Remove step_scheduler.max_steps to train full epochs instead of a benchmark cap.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_1_dev-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25            # benchmark cap; remove for a real pretrain run
+        local_batch_size: 4
+        global_batch_size: 32    # 4 x dp8
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pcam_flux_512
+          base_resolution: [512, 512]
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq1024.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq1024.yaml
new file mode 100644
index 000000000..8a48d829f
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq1024.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 1024x1024 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev_sq1024}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_1024_sq
+          base_resolution: [1024, 1024]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: true
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq256.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq256.yaml
new file mode 100644
index 000000000..2570d1291
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq256.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 256x256 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev_sq256}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_256_sq
+          base_resolution: [256, 256]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: true
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq512.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq512.yaml
new file mode 100644
index 000000000..8cca0f71d
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq512.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 512x512 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev_sq512}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_512_sq
+          base_resolution: [512, 512]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: true
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq720.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq720.yaml
new file mode 100644
index 000000000..64fa48876
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_dev-sq720.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-dev T2I PRETRAIN (from scratch), 720x720 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev_sq720}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_dev.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_720_sq
+          base_resolution: [720, 720]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: true
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq1024.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq1024.yaml
new file mode 100644
index 000000000..ffe8f6323
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq1024.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-schnell T2I PRETRAIN (from scratch), 1024x1024 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_schnell_sq1024}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_schnell.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_1024_sq
+          base_resolution: [1024, 1024]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: false
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq256.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq256.yaml
new file mode 100644
index 000000000..350c5cda2
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq256.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-schnell T2I PRETRAIN (from scratch), 256x256 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_schnell_sq256}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_schnell.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_256_sq
+          base_resolution: [256, 256]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: false
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq512.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq512.yaml
new file mode 100644
index 000000000..19e761804
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq512.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-schnell T2I PRETRAIN (from scratch), 512x512 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_schnell_sq512}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_schnell.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_512_sq
+          base_resolution: [512, 512]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: false
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq720.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq720.yaml
new file mode 100644
index 000000000..67a03794d
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/flux_schnell-sq720.yaml
@@ -0,0 +1,31 @@
+# Primus EXP — FLUX.1-schnell T2I PRETRAIN (from scratch), 720x720 SQUARE (T5-256).
+# Auto-generated for the FLUX square resolution ladder (256/512/720/1024).
+# The 256 rung matches the MLPerf FLUX-schnell operating point (256px, 256-token T5).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_schnell_sq720}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: flux_pre_trainer.yaml
+    model: flux_1_schnell.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 4
+        global_batch_size: 32
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /checkpoints/pcam_flux_720_sq
+          base_resolution: [720, 720]
+      flow_matching:
+        adapter_kwargs:
+          guidance_scale: 1
+          use_guidance_embeds: false
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-ddp.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-ddp.yaml
new file mode 100644
index 000000000..18a02c968
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-ddp.yaml
@@ -0,0 +1,43 @@
+# Primus EXP — Ideogram-4 single-stream T2I PRETRAIN, DDP + ZeRO-1 variant.
+# =============================================================================
+# Sibling of ideogram4-pretrain.yaml (the FSDP2 EXP). Uses the DDP module preset
+# (ideogram4_ddp_pre_trainer.yaml -> `ddp:` block) so the recipe replicates
+# params+grads and skips per-layer all-gather; pair with PRIMUS_IDEOGRAM_ZERO1=1
+# to shard the optimizer state (ZeRO-1) + whole-block-checkpoint the Ideogram
+# layers on the DDP path (no-fork zero1.py hook).
+#
+# CAVEAT: torch.compile is FSDP2-only and single-node comm is ~hidden, so this
+# path trails FSDP2+compile on one node. It is the multi-node lever / reference;
+# do NOT set fsdp.enable_compile here (no fsdp block).
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:ideogram4-pretrain-ddp}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: ideogram4_ddp_pre_trainer.yaml
+    model: ideogram4.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 20
+        local_batch_size: 2
+        global_batch_size: 16    # 2 x dp8 (grad_acc = 1 -> no no_sync needed)
+        log_remote_every_steps: 1
+      optimizer:
+        lr: 1.0e-4
+      ddp:
+        activation_checkpointing: true
+      model:
+        torch_dtype: bfloat16    # ZeRO-1/DDP baseline (bf16 master, matches compile A/B)
+      data:
+        dataloader:
+          num_samples: 64
+          grid_h: 16
+          grid_w: 16
+          max_text_tokens: 128
+          min_text_tokens: 128
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-pcam.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-pcam.yaml
new file mode 100644
index 000000000..41b99289b
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-pcam.yaml
@@ -0,0 +1,46 @@
+# Primus EXP — Ideogram-4 single-stream T2I PRETRAIN on REAL (pre-encoded) data.
+# =============================================================================
+# Same no-fork Primus path as `ideogram4-pretrain.yaml`, but trains on a REAL
+# pre-encoded cache instead of the synthetic one:
+#   - data.dataloader -> Ideogram4CacheDataloaderConfig
+#     (primus/backends/nemo_automodel/models/ideogram4/data/cache.py), reading a
+#     Flux-2 VAE + Qwen3-VL 13-layer cache produced by Ideogram4Processor
+#     (primus/backends/nemo_automodel/models/ideogram4/processor.py).
+#   - flow_matching.adapter_type: ideogram4 -> Ideogram4Adapter
+#   - fsdp.activation_checkpointing + PRIMUS_IDEOGRAM_REAL_AC=1 -> real AC
+# The DiT still builds from a weightless local config dir (random init, no gated
+# transformer weights); only the DATA is real here.
+#
+# Point `data.dataloader.cache_dir` at a cache built by Ideogram4Processor from an
+# image dir + caption dir. Building it needs 1 GPU and a Hugging Face token with
+# access to the gated VAE/text-encoder repos.
+#
+# Goal: a clearly decreasing loss on 8-GPU FSDP2 bf16 (real AC) over REAL Ideogram-4
+# conditioning — validating the processor + cache loader end to end.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:ideogram4-pretrain-pcam}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: ideogram4_pre_trainer_pcam.yaml
+    model: ideogram4.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 200           # real-data loss-decrease validation (remove for full run)
+        local_batch_size: 2
+        global_batch_size: 16    # 2 x dp8 (grad_acc = 1)
+        log_remote_every_steps: 1
+      optimizer:
+        lr: 1.0e-4
+      fsdp:
+        activation_checkpointing: true
+        enable_compile: false
+      data:
+        dataloader:
+          cache_dir: /dataset/pcam_ideogram4_256
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain.yaml
new file mode 100644
index 000000000..5f1ab7a1e
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain.yaml
@@ -0,0 +1,58 @@
+# Primus EXP — Ideogram-4 single-stream T2I PRETRAIN (from scratch, random init).
+# =============================================================================
+# Config-only Ideogram-4 integration: runs through the exact same Primus
+# nemo_automodel path as Wan2.2 / FLUX. The Ideogram-specific pieces are all
+# no-fork Primus hooks (installed by NemoAutomodelPretrainTrainer):
+#   - flow_matching.adapter_type: ideogram4  -> Ideogram4Adapter
+#     (primus/backends/nemo_automodel/models/ideogram4/adapter.py)
+#   - data.dataloader -> SyntheticIdeogram4DataloaderConfig
+#     (primus/backends/nemo_automodel/models/ideogram4/data/synthetic.py)
+#   - fsdp.activation_checkpointing + PRIMUS_IDEOGRAM_REAL_AC=1 -> real AC
+#     (primus/backends/nemo_automodel/models/ideogram4/parallelize.py)
+# The model builds from a weightless local config dir (no gated weights).
+#
+# One-time setup: `model.pretrained_model_name_or_path` below must point at a local
+# Ideogram-4 config dir (a `model_index.json` + transformer `config.json` describing
+# the architecture). No weights are needed — the DiT is randomly initialised.
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   PYTHONPATH=<primus>:<primus>/third_party/Automodel \
+#   HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PRIMUS_IDEOGRAM_REAL_AC=1 \
+#   torchrun --nproc_per_node=8 --nnodes=1 --master_port=29533 \
+#     primus/cli/main.py train pretrain \
+#     --config examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain.yaml
+#
+# This is a from-config OVERFIT smoke on a fixed synthetic cache: the goal is a
+# clearly decreasing loss on 8-GPU FSDP2 bf16 with real activation checkpointing,
+# not convergence. Remove step_scheduler.max_steps / swap the dataloader for the
+# real preprocessor cache (`ideogram4-pretrain-pcam.yaml`) for a real pretrain run.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:ideogram4-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: ideogram4_pre_trainer.yaml
+    model: ideogram4.yaml
+
+    overrides:
+      step_scheduler:
+        max_steps: 60            # overfit-smoke cap; remove for a real pretrain run
+        local_batch_size: 2
+        global_batch_size: 16    # 2 x dp8 (grad_acc = 1)
+        log_remote_every_steps: 1
+      optimizer:
+        lr: 1.0e-4
+      fsdp:
+        activation_checkpointing: true
+        enable_compile: false
+      data:
+        dataloader:
+          num_samples: 64
+          grid_h: 16
+          grid_w: 16
+          max_text_tokens: 32
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml
new file mode 100644
index 000000000..a871c327c
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml
@@ -0,0 +1,70 @@
+# Primus EXP — Wan 2.2 T2V A14B FINETUNE via the NeMo AutoModel backend on MI355X.
+# =============================================================================
+# Full finetune of the pretrained Wan 2.2-A14B weights (loaded from
+# model.pretrained_model_name_or_path via from_pretrained). Differs from the
+# pretrain EXPs by selecting AutoModel's finetune path:
+#   model.mode: finetune        -> build_model_and_optimizer(finetune_mode=True);
+#                                  pipeline_spec is ignored in this path.
+#   model.torch_dtype: bfloat16 -> bf16 master weights (canonical finetune
+#                                  precision; pretrain uses fp32 master).
+#   gentle LR (5e-6) + standard AdamW betas/decay (foreach, not fused).
+#
+# Runs through the SAME diffusion recipe and the SAME launcher as pretrain
+# (AutoModel selects pretrain vs finetune from model.mode), so it launches via
+# `examples/run_pretrain.sh` like the other diffusion EXPs.
+#
+# NOTE: the finetune recipe is AutoModel's original Wan 2.2 path, provided here
+# as a complete, first-class example. The throughput numbers documented in the
+# pretrain EXPs were measured on the pretrain operating points.
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-finetune.yaml \
+#   bash examples/run_pretrain.sh
+#
+# Run once per expert. The expert is chosen solely by model.stage (finetune loads
+# via from_pretrained with the matching active_transformer; pipeline_spec/subfolder
+# is ignored in this path -- it only applies to pretraining's from_config):
+#   high_noise (default) -> transformer;  model.stage=low_noise -> transformer_2.
+# Remove step_scheduler.max_steps to finetune for full epochs.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-finetune}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      model:
+        mode: finetune
+        torch_dtype: bfloat16    # bf16 master weights for finetuning
+      step_scheduler:
+        max_steps: 25            # cap for a quick test; remove for a real finetune
+        local_batch_size: 8
+        global_batch_size: 64    # 8 x dp8
+      optimizer:
+        lr: 5e-6                 # gentle finetune LR
+        weight_decay: 0.01
+        betas: [0.9, 0.999]
+        # finetune uses foreach AdamW (fused is the pretrain throughput knob)
+        foreach: true
+        fused: false
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: true
+        checkpoint_dir: /checkpoints/wan22_finetune_high
+        restore_from: null       # null -> finetune the base pretrained weights
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml
new file mode 100644
index 000000000..147275927
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml
@@ -0,0 +1,52 @@
+# Primus EXP — Wan 2.2 T2V A14B PRETRAIN, optimized AC-on bs32.
+# =============================================================================
+# This is the "best AC-on" benchmark point reproduced end-to-end through the
+# full Primus runtime (NeMo AutoModel backend) on 8x MI355X. It layers the
+# optimized knobs (already baked into the pre_trainer module config:
+# torch.compile + fused AdamW + bf16 grad-reduce + deeper FSDP2 prefetch, AITER
+# attention, fp32 master / bf16 compute) on top of the Wan 2.2-A14B preset.
+#
+# Measured (8x MI355X, 512x512, real PusaV1, fp32 master / bf16 compute):
+#   ~3.10-3.12 samples/s/GPU (~24.9 total) at mem ~84.9 GB
+#   == +~46% vs the un-optimized baseline (2.14 s/s/GPU). torch.compile runs
+#      cleanly with the AITER custom-op backend (no graph break).
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs32-acon.yaml \
+#   bash examples/run_pretrain.sh
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+# Remove step_scheduler.max_steps to train full epochs instead of a benchmark cap.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain-bs32-acon}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # Headline AC-on plateau point. activation_checkpointing stays TRUE
+      # (the pre_trainer module default) so bs32 fits with headroom.
+      step_scheduler:
+        max_steps: 25            # benchmark cap; remove for a real pretrain run
+        local_batch_size: 32
+        global_batch_size: 256   # 32 x dp8
+      fsdp:
+        activation_checkpointing: true
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml
new file mode 100644
index 000000000..e5b3eeb84
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml
@@ -0,0 +1,52 @@
+# Primus EXP — Wan 2.2 T2V A14B PRETRAIN, optimized AC-off bs8.
+# =============================================================================
+# Fastest measured operating point. With torch.compile cutting activation
+# memory, activation checkpointing can be turned OFF at bs8 (which OOM'd in the
+# un-optimized baseline), giving the best throughput. This mirrors the Wan 2.1
+# GB300 small-batch / AC-off reference setup.
+#
+# Measured (8x MI355X, 512x512, real PusaV1, fp32 master / bf16 compute):
+#   ~3.68-3.69 samples/s/GPU (~29.4 total) at mem ~236.8 GB (AC OFF, bs8)
+#   == new best point; un-optimized AC-off OOM'd at bs8 (ceiling was bs6 @ ~240 GB).
+# (bs6 AC-off, the more conservative GB300-comparable point, measures
+#  ~3.54 s/s/GPU @ ~183.5 GB — set local_batch_size=6 / global_batch_size=48.)
+#
+# One-time setup (like any Primus backend): initialize the submodule; the
+# launcher's prepare step installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (auto-skipped if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-bs8-acoff.yaml \
+#   bash examples/run_pretrain.sh
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+# Remove step_scheduler.max_steps to train full epochs instead of a benchmark cap.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain-bs8-acoff}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # AC-OFF best point: compile buys enough memory to drop activation
+      # checkpointing at bs8 (~237 GB of the 288 GB HBM3e).
+      step_scheduler:
+        max_steps: 25            # benchmark cap; remove for a real pretrain run
+        local_batch_size: 8
+        global_batch_size: 64    # 8 x dp8
+      fsdp:
+        activation_checkpointing: false
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false           # set true + checkpoint_dir for a real run
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-te-fp8.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-te-fp8.yaml
new file mode 100644
index 000000000..4ac21d1a1
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-te-fp8.yaml
@@ -0,0 +1,53 @@
+# Primus EXP config — Wan 2.2 T2V A14B pretrain, FP8 via TransformerEngine autocast,
+# NeMo AutoModel backend, MI355X.
+#
+# This is the FP8 counterpart of wan2_2_t2v_a14b-pretrain.yaml. It is CONFIG-ONLY
+# — no AutoModel/Diffusers code change, no fork:
+# the AutoModel diffusion recipe already reads model.transformer_engine_fp8 and wraps
+# each forward in transformer_engine.pytorch.quantization.autocast, swapping the
+# Diffusers nn.Linear modules for TE Linear (transformer_engine_linear is auto-forced
+# true when fp8 is on). Attention stays on AITER; only the Linear GEMMs run in FP8.
+#
+# FP8 via TransformerEngine autocast for the Diffusers + FSDP2 Wan stack.
+#
+# Launch (inside the Primus container, 8 GPUs):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain-te-fp8.yaml \
+#   bash examples/run_pretrain.sh
+#
+# If TE fp8 + torch.compile graph-breaks or errors, retry with compile off:
+#   ... bash examples/run_pretrain.sh fsdp.enable_compile=false
+# Alternate scaling recipe (default delayed):
+#   ... bash examples/run_pretrain.sh model.transformer_engine_fp8_recipe=current
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain-te-fp8}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      model:
+        # FP8 via TransformerEngine autocast (delayed tensorwise scaling, matching
+        # the FLUX local-spec/te-fp8 reference). transformer_engine_linear is forced
+        # true by the recipe when fp8 is enabled.
+        transformer_engine_fp8: true
+        transformer_engine_fp8_recipe: delayed
+        transformer_engine_fp8_amax_history_len: 1024
+        transformer_engine_fp8_amax_compute_algo: max
+      # Batch sizing matches the BF16 EXP so the FP8 vs BF16 comparison is
+      # apples-to-apples.
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 32
+        global_batch_size: 256
+      data:
+        dataloader:
+          cache_dir: /path/to/pusa_v1_512sq
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml
new file mode 100644
index 000000000..dd19f295b
--- /dev/null
+++ b/examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml
@@ -0,0 +1,43 @@
+# Primus EXP config — Wan 2.2 T2V A14B pretrain via the NeMo AutoModel backend
+# on MI355X. Selects framework: nemo_automodel, consistent
+# with how the megatron diffusion FLUX EXPs select framework: megatron.
+#
+# One-time setup (like any other Primus backend): initialize the submodule. The
+# launcher's prepare step then installs it editable (ROCm-pinned) on first run.
+#   git submodule update --init third_party/Automodel
+#   # (skipped automatically if the image already ships nemo_automodel)
+#
+# Launch (inside the Primus container, 8 GPUs free):
+#   USE_HIPBLASLT=1 TORCH_BLAS_PREFER_HIPBLASLT=1 HIP_FORCE_DEV_KERNARG=1 GPU_MAX_HW_QUEUES=6 \
+#   GPUS_PER_NODE=8 \
+#   EXP=examples/nemo_automodel/configs/MI355X/diffusion/wan2_2_t2v_a14b-pretrain.yaml \
+#   bash examples/run_pretrain.sh
+#   # AC-off operating points (Wan 2.1 GB300-comparable) via CLI overrides, e.g.:
+#   #   ... bash examples/run_pretrain.sh \
+#   #       fsdp.activation_checkpointing=false \
+#   #       step_scheduler.local_batch_size=8 step_scheduler.global_batch_size=64
+#
+# low-noise expert: override model.stage=low_noise and
+#   model.pipeline_spec.subfolder=transformer_2.
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2_2_t2v_a14b-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: nemo_automodel
+    config: pre_trainer.yaml
+    model: wan2_2_t2v_a14b.yaml
+
+    overrides:
+      # Throughput-benchmark sizing (optimized headline point: bs32 AC-on).
+      step_scheduler:
+        max_steps: 25
+        local_batch_size: 32
+        global_batch_size: 256
+      data:
+        dataloader:
+          cache_dir: /dataset/pusa_v1_512sq
+      checkpoint:
+        enabled: false
diff --git a/examples/nemo_automodel/prepare.py b/examples/nemo_automodel/prepare.py
new file mode 100644
index 000000000..72f9ef7dd
--- /dev/null
+++ b/examples/nemo_automodel/prepare.py
@@ -0,0 +1,326 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Backend prepare entry for the NeMo AutoModel framework.
+
+Mirrors ``examples/torchtitan/prepare.py``: it resolves the backend checkout
+(the ``third_party/Automodel`` submodule by default) and installs it editable so
+the ``nemo_automodel`` package is importable for the training phase. Two
+AutoModel-specific differences from the torchtitan flow:
+
+  1. **ROCm-safe install.** AutoModel's dependency graph would otherwise let pip
+     pull CUDA torch wheels on top of the image's ROCm build. We pin the native
+     ROCm packages (torch/torchvision/torchaudio/triton/aiter/flash-attn) to the
+     versions already installed and pass them as a pip *constraint*, so the
+     editable install never replaces the ROCm stack. The ``diffusion`` extras
+     are installed for the Wan diffusion recipe.
+  2. **Idempotent / skip-if-present.** If ``nemo_automodel`` is already importable
+     (e.g. a base image that ships it pre-installed) we skip the install unless
+     ``AUTOMODEL_REINSTALL=1`` is set.
+
+Unlike torchtitan there is no tokenizer asset to pre-download: the Wan diffusion
+recipe pulls model weights from HuggingFace / a local ``/models`` mount at
+runtime. The heavy lifting (config -> AutoModel ConfigNode -> TrainDiffusionRecipe)
+is done later by ``NemoAutomodelPretrainTrainer`` inside the Primus core runtime.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+from examples.scripts.utils import (
+    get_env_case_insensitive,
+    log_error_and_exit,
+    log_info,
+    write_patch_args,
+)
+from primus.core.launcher.parser import PrimusParser
+
+# Native ROCm packages to pin so the editable install never swaps them for CUDA
+# wheels.
+_ROCM_PINS = [
+    "torch",
+    "torchvision",
+    "torchaudio",
+    "triton",
+    "pytorch-triton-rocm",
+    "amd-aiter",
+    "flash-attn",
+    "flash_attn",
+]
+
+# diffusers gates its ``aiter`` attention backend behind a minimum amd-aiter
+# version. Some ROCm base images ship a functional but dev-versioned aiter (e.g.
+# ``0.1.1.dev*``) that fails that guard even though the kernel works (its
+# ``flash_attn_func`` supports ``return_lse``). The configs default to
+# ``model.attention_backend: aiter``, so without a fix diffusers would refuse it.
+_REQUIRED_AITER_VERSION = "0.1.5"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Prepare Primus NeMo AutoModel backend")
+    parser.add_argument("--primus_path", type=str, required=True, help="Root path to the Primus project")
+    parser.add_argument("--data_path", type=str, required=True, help="Path to data directory")
+    parser.add_argument("--config", type=str, required=True, help="Path to experiment YAML config")
+    parser.add_argument(
+        "--patch_args",
+        type=str,
+        default="/tmp/primus_patch_args.txt",
+        help="Path to write additional args (used during training phase)",
+    )
+    parser.add_argument(
+        "--backend_path",
+        type=str,
+        default=None,
+        help="Optional AutoModel checkout path; overrides AUTOMODEL_PATH/BACKEND_PATH and the default submodule.",
+    )
+    return parser.parse_args()
+
+
+def resolve_backend_path(cli_path, primus_path: Path) -> Path:
+    """CLI --backend_path > AUTOMODEL_PATH/BACKEND_PATH env > third_party/Automodel."""
+    if cli_path:
+        path = Path(cli_path).resolve()
+        log_info(f"Using AutoModel path from CLI: {path}")
+        return path
+    env_value = get_env_case_insensitive("AUTOMODEL_PATH") or get_env_case_insensitive("BACKEND_PATH")
+    if env_value:
+        path = Path(env_value).resolve()
+        log_info(f"Using AutoModel path from env: {path}")
+        return path
+    path = (primus_path / "third_party" / "Automodel").resolve()
+    log_info(f"Using default AutoModel submodule path: {path}")
+    return path
+
+
+def write_rocm_constraints() -> str:
+    """Pin currently-installed ROCm packages so pip treats them as satisfied."""
+    import importlib.metadata as md
+
+    lines = []
+    for name in _ROCM_PINS:
+        try:
+            lines.append(f"{name}=={md.version(name)}")
+        except md.PackageNotFoundError:
+            continue
+    fd, path = tempfile.mkstemp(prefix="primus_rocm_constraints.", suffix=".txt")
+    with os.fdopen(fd, "w") as f:
+        f.write("\n".join(lines) + "\n")
+    log_info(f"ROCm constraints ({len(lines)} pins) -> {path}")
+    return path
+
+
+def install_automodel_editable(automodel_path: Path):
+    """ROCm-safe editable install of AutoModel with the diffusion extras."""
+    if not (automodel_path / "pyproject.toml").exists() and not (automodel_path / "setup.py").exists():
+        log_error_and_exit(
+            f"AutoModel checkout not found at {automodel_path} (no pyproject.toml/setup.py).\n"
+            "Initialize the submodule first:\n"
+            "    git submodule update --init third_party/Automodel\n"
+            "or run `primus-cli deps sync`, or pass --backend_path / set AUTOMODEL_PATH."
+        )
+
+    extras = os.environ.get("AUTOMODEL_EXTRAS", "diffusion,diffusion-media")
+    spec = f"{automodel_path}[{extras}]" if extras else str(automodel_path)
+    constraints = write_rocm_constraints()
+
+    env = os.environ.copy()
+    env["PIP_CONSTRAINT"] = constraints
+    log_info(f"Installing AutoModel (editable, ROCm-pinned): pip install --no-build-isolation -e {spec!r}")
+    ret = None
+    try:
+        ret = subprocess.run(
+            [sys.executable, "-m", "pip", "install", "--no-build-isolation", "-e", spec, "-q"],
+            env=env,
+            check=False,
+        )
+    except OSError as e:
+        log_error_and_exit(f"Failed to invoke pip for AutoModel editable install: {e}")
+    finally:
+        try:
+            os.unlink(constraints)
+        except OSError as e:
+            # Best-effort cleanup of the temp constraints file; a failure here is
+            # non-fatal and must not mask the install result.
+            log_info(f"Non-fatal: could not remove temp ROCm constraints file {constraints!r}: {e}")
+    if ret is None or ret.returncode != 0:
+        rc = ret.returncode if ret is not None else "n/a"
+        log_error_and_exit(f"AutoModel editable install failed (exit {rc}).")
+    log_info("AutoModel editable install complete.")
+
+
+def maybe_shim_aiter_version():
+    """Make diffusers' ``aiter`` backend selectable on dev-versioned aiter images.
+
+    Rewrites only the ``Version:`` field in the installed ``amd-aiter`` dist
+    metadata so it satisfies diffusers' minimum-version guard. This touches
+    neither aiter, diffusers, nor AutoModel source; it reverts on image rebuild
+    and is idempotent (skipped when the installed version already satisfies the
+    guard, when aiter is absent, or when its metadata is read-only). Disable with
+    ``AITER_VERSION_SHIM=0``.
+
+    NOTE (single source of truth): the same metadata-rewrite is performed by the
+    NeMo-AutoModel dev-env ``entrypoint.sh`` (tiger-training-internal #208). The
+    two must stay in sync -- keep ``_REQUIRED_AITER_VERSION`` and the rewrite
+    logic identical, or fold both onto a shared helper if the dev-env starts
+    importing from this repo. Kept duplicated for now because the dev-env image
+    does not depend on Primus at container-build time.
+    """
+    if os.environ.get("AITER_VERSION_SHIM", "1") != "1":
+        log_info("AITER version shim disabled (AITER_VERSION_SHIM=0).")
+        return
+
+    import importlib.metadata as md
+    import re
+
+    try:
+        dist = md.distribution("amd-aiter")
+        cur = dist.version
+    except md.PackageNotFoundError:
+        log_info("amd-aiter not installed; AITER version shim skipped.")
+        return
+
+    try:
+        from packaging.version import Version
+
+        if Version(cur) >= Version(_REQUIRED_AITER_VERSION):
+            log_info(f"AITER shim: amd-aiter {cur} already satisfies >={_REQUIRED_AITER_VERSION}; no change.")
+            return
+    except Exception:
+        pass  # unparseable version -> attempt the bump below
+
+    # Locate the metadata file via the public API (RECORD-relative paths resolved
+    # through Distribution.locate_file). Fall back to the .dist-info dir only if
+    # the distribution ships no RECORD (dist.files is None).
+    meta_path = None
+    try:
+        for f in dist.files or []:
+            if f.name in ("METADATA", "PKG-INFO"):
+                located = dist.locate_file(f)
+                if os.path.exists(located):
+                    meta_path = str(located)
+                    break
+    except Exception:
+        meta_path = None
+    if not meta_path:
+        try:
+            base = dist._path  # PathDistribution .dist-info dir (fallback only)
+            for cand in ("METADATA", "PKG-INFO"):
+                p = os.path.join(str(base), cand)
+                if os.path.exists(p):
+                    meta_path = p
+                    break
+        except Exception:
+            meta_path = None
+
+    if not meta_path or not os.access(meta_path, os.W_OK):
+        log_info(
+            f"AITER shim: amd-aiter metadata not writable ({meta_path}); leaving {cur} as-is. "
+            "If the diffusers aiter guard blocks it, set model.attention_backend: flash."
+        )
+        return
+
+    with open(meta_path, encoding="utf-8") as f:
+        txt = f.read()
+    new, n = re.subn(r"(?im)^Version: .*$", f"Version: {_REQUIRED_AITER_VERSION}", txt, count=1)
+    if n == 0 or new == txt:
+        log_info(
+            f"AITER shim: no Version field found in amd-aiter metadata ({meta_path}); leaving {cur} as-is. "
+            "If the diffusers aiter guard blocks it, set model.attention_backend: flash."
+        )
+        return
+    with open(meta_path, "w", encoding="utf-8") as f:
+        f.write(new)
+    log_info(
+        f"AITER shim: bumped amd-aiter metadata Version {cur} -> {_REQUIRED_AITER_VERSION} "
+        "to satisfy the diffusers aiter guard."
+    )
+
+
+def ensure_automodel_installed(cli_path, primus_path: Path):
+    """Install AutoModel from the submodule unless it is already importable."""
+    reinstall = os.environ.get("AUTOMODEL_REINSTALL", "0") == "1"
+    # Only take the skip-if-present fast path when the user did not explicitly
+    # point us at a checkout (CLI --backend_path / AUTOMODEL_PATH / BACKEND_PATH).
+    # If they did, honor it and (re)install so an explicit override beats a
+    # nemo_automodel that the base image happens to ship (matches --backend_path help).
+    explicit_source = (
+        bool(cli_path)
+        or bool(get_env_case_insensitive("AUTOMODEL_PATH"))
+        or bool(get_env_case_insensitive("BACKEND_PATH"))
+    )
+    if not reinstall and not explicit_source:
+        try:
+            import nemo_automodel  # noqa: F401
+
+            log_info(
+                f"nemo_automodel already importable ({getattr(nemo_automodel, '__file__', '?')}); "
+                "skipping install. Set AUTOMODEL_REINSTALL=1 to force a reinstall from third_party."
+            )
+            return
+        except Exception:
+            log_info("nemo_automodel not importable; installing from the AutoModel checkout.")
+
+    automodel_path = resolve_backend_path(cli_path, primus_path)
+    install_automodel_editable(automodel_path)
+
+
+def main():
+    args = parse_args()
+
+    primus_path = Path(args.primus_path).resolve()
+    exp_path = Path(args.config).resolve()
+    patch_args_file = Path(args.patch_args).resolve()
+
+    log_info(f"PRIMUS_PATH: {primus_path}")
+    log_info(f"DATA_PATH: {Path(args.data_path).resolve()}")
+    log_info(f"EXP: {exp_path}")
+    log_info(f"BACKEND_PATH: {args.backend_path}")
+    log_info(f"PATCH-ARGS: {patch_args_file}")
+
+    if not exp_path.is_file():
+        log_error_and_exit(f"EXP file not found: {exp_path}")
+
+    # 1) Make the nemo_automodel package importable (submodule -> editable install).
+    ensure_automodel_installed(args.backend_path, primus_path)
+
+    # 1b) Make diffusers' aiter attention backend selectable on images that ship a
+    #     dev-versioned amd-aiter (configs default to attention_backend: aiter).
+    maybe_shim_aiter_version()
+
+    # 2) Validate the experiment parses and routes to the nemo_automodel backend.
+    primus_cfg = PrimusParser().parse(args)
+    try:
+        pre_trainer_cfg = primus_cfg.get_module_config("pre_trainer")
+    except Exception:
+        log_error_and_exit("Missing required module config: pre_trainer")
+
+    framework = getattr(pre_trainer_cfg, "framework", None)
+    if framework != "nemo_automodel":
+        log_error_and_exit(
+            f"pre_trainer.framework must be 'nemo_automodel' (got {framework!r}). "
+            "Check the experiment config."
+        )
+
+    if not getattr(pre_trainer_cfg, "model", None):
+        log_error_and_exit("Missing required field: pre_trainer.model (model preset)")
+
+    log_info(
+        f"NeMo AutoModel backend ready (framework={framework}, model={pre_trainer_cfg.model}). "
+        "Weights resolve via HF cache / /models at runtime."
+    )
+
+    # 3) Keep multi-GPU stdout to rank 0 only, consistent with other backends.
+    write_patch_args(patch_args_file, "torchrun_args", {"local-ranks-filter": "0"})
+
+
+if __name__ == "__main__":
+    log_info("========== Prepare NeMo AutoModel backend ==========")
+    main()
diff --git a/examples/odc/README.md b/examples/odc/README.md
new file mode 100644
index 000000000..2c5a9b498
--- /dev/null
+++ b/examples/odc/README.md
@@ -0,0 +1,340 @@
+# Reproducing ODC on AMD ROCm (Primus + Primus-Turbo)
+
+This is the companion, from-zero reproduction guide for the ROCm blog on using ODC to accelerate AMD SFT training: it walks through single-node 1.5B and dual-node 14B runs — including the `nccl_pad` fair baseline — so you can build the ODC rocSHMEM operators, run the arms, and measure the speedup yourself.
+
+This guide distills the experiments into a follow-along, copy-pasteable tutorial, so that **a reader who has never touched this codebase** can start from zero and fully run the `odc_nopad` experiments for **single-node 1.5B** and **dual-node 14B** (including the `nccl_pad` fair baseline) and compute the speedup.
+
+**Overall flow at a glance (6 steps; strongly recommended to get single-node working first, then move to dual-node):**
+
+1. **Start the container** (see "Get the code, start the container"): clone the PR #864 branch + Primus-Turbo, `docker run` to start the container.
+2. **Get the GDA operators in place** (see "Build Primus-Turbo with ODC operators"): **build Primus-Turbo from source** (one host + one GDA build, linking rocSHMEM at build time), then use `.image_bak` to disable the stock turbo shipped in the image that has no ODC operators, consume it via `PRIMUS_TURBO_PATH`, and verify `has_gda=True`.
+3. **Prepare offline data** (see "Prerequisites"): stage the HF model/data, set `HF_HOME` + `HF_HUB_OFFLINE=1`.
+4. **Turn on ODC via config** (see "Turning on ODC = changing config items"): the `odc_nopad` arm = `enable_odc:true`+`enable_odc_lb_mini:true`+`odc_p2p_backend:rocshmem`; the `nccl_pad` aligned baseline = `enable_odc:false`+`enable_odc_lb_mini:true`.
+5. **Start training**: for single-node see "Single-node 1.5B `odc_nopad`" (8 GPUs, no RoCE/GDA needed, **newcomers do this first**); once it works, move to dual-node per "Dual-node 14B `odc_nopad`" (16 GPUs, rocSHMEM GDA).
+6. **Validate + compute the speedup** (see "Judging success, reading results, computing the speedup"): confirm from the logs that both fixes are in effect, loss is decreasing, 0 nan, and `speedup = nccl_pad's ms/iter ÷ odc_nopad's ms/iter`.
+
+When stuck, first check the "Pitfalls checklist" section below.
+
+> **Important change — ODC is now driven by "config items," no longer by environment variables.** PR [#864](https://github.com/AMD-AGI/Primus/pull/864) (branch `feat/odc-consume-turbo`) changed all production feature switches into Primus yaml config items (`enable_odc`, `odc_phase`, `enable_odc_lb_mini`, `odc_p2p_backend`, `odc_rocshmem_gda`, `odc_gda_*` …), which are no longer read from `os.environ` at runtime. So all commands in this guide are based on "change config item / CLI override"; the legacy `ODC_*` environment variables have been deprecated.
+
+To make ODC **run, run correctly, and run fast** on ROCm, **two fixes are indispensable** (division of labor: **Primus** owns the algorithm layer + integration, pure Python under `primus/core/odc/`; **Primus-Turbo** owns the rocSHMEM comm operators):
+
+- **① Primus-Turbo comm operators (`odc_rocshmem_host` / `odc_rocshmem_gda`).** The body of ODC's P2P communication. The operator source has been merged into [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) official `main` (PR #409, including the `--lto-partitions=1` required by GDA), but is gated by `#ifndef DISABLE_ROCSHMEM` — a plain `pip install @main` **will not build the operators** (`has_gda=False`). You must **build from source, pointing `ROCSHMEM_HOME` at a prebuilt rocSHMEM at build time** (one for single-node host / one for dual-node GDA); see "Build Primus-Turbo with ODC operators".
+- **② hook into the correct FSDP2 class (PR [#808](https://github.com/AMD-AGI/Primus/pull/808) / [#864](https://github.com/AMD-AGI/Primus/pull/864)).** ODC must hook the new `PrimusTorchFullyShardedDataParallel`, otherwise `reduction_service` is `None` and **iter2 crashes** with `'NoneType' ... clear_accumulations`. PR #864 fixes this, taking effect only when `enable_odc: true` + `use_torch_fsdp2: true`.
+
+Both above are **gated by yaml config items** (no longer by `ODC_*` environment variables), with zero impact on `nccl_pad` / native FSDP2.
+
+## Prerequisites
+
+### Hardware
+
+- Single-node 1.5B: 1 AMD MI300X/MI355X (`gfx942`) ×8, intra-node XGMI interconnect suffices.
+- Dual-node 14B: 2 machines ×8 = 16 GPUs, requiring RoCE/RDMA (multiple mlx5 NICs) between nodes. **Be sure to pick adjacent nodes under the same leaf switch, mutually reachable over RoCE** — cross-topology node pairs often hang at the first cross-node communication (we measured multiple non-adjacent node pairs stalling at GDA warmup, going a long time without producing an `iteration`).
+
+### Software
+
+- ROCm 7.2.0 (`gfx942`), Docker, (optional for dual-node) SLURM.
+- Container image: `tasimage/primus-odc:v26.2` (ROCm 7.2.0 based; `primus/core/odc/README.md` recommends the same one).
+- **Primus-Turbo GDA dependency**: a Primus-Turbo build with the ODC rocSHMEM operators compiled in (including `odc_rocshmem_host` / `odc_rocshmem_gda`). Dual-node GDA must use the one built with **`GDA_MLX5` on and device-LTO single-partition** (`-Xoffload-linker --lto-partitions=1`, built into `setup.py` by PR #409), otherwise cross-node `getmem` reads 0 → **dual-node `grad_norm=0`**.
+- HF offline cache: stage the model/data locally in advance, and set `HF_HOME` + `HF_HUB_OFFLINE=1` for fully offline training. Single-node 1.5B uses `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B`, dual-node 14B uses `deepseek-ai/DeepSeek-R1-Distill-Qwen-14B`, and the SFT data for both is `zai-org/LongAlign-10k`.
+
+**Pre-download (run once on a machine with network access, after which training can be fully offline):**
+
+```bash
+export HF_HOME=$HOME/primus_packed/hf_home           # cache lands here
+pip install -U "huggingface_hub[cli]"
+# models (single-node 1.5B / dual-node 14B)
+huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
+huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Qwen-14B
+# SFT dataset
+huggingface-cli download zai-org/LongAlign-10k --repo-type dataset
+# at training time, also set: export HF_HOME=$HOME/primus_packed/hf_home HF_HUB_OFFLINE=1
+```
+
+(In newer versions `huggingface-cli` can also be written as `hf download ...`; the cache directory follows `HF_HOME`, and the training side hits it as long as it points at the same `HF_HOME`.)
+
+## Get the code, start the container
+
+```bash
+# 1) get the PR #864 branch (ODC algorithm layer + integration, config-item driven)
+git clone -b feat/odc-consume-turbo https://github.com/AMD-AGI/Primus.git
+# the comm operators come from Primus-Turbo: its source is on official main (PR #409), but a plain pip install won't expose the operators
+# (the operators are compiled out by -DDISABLE_ROCSHMEM; you must link rocSHMEM at build time). See "Build Primus-Turbo with ODC operators": build host + GDA
+# from source, consumed via PRIMUS_TURBO_PATH + .image_bak
+
+# 2) start one container per physical node (mounting the code and HF cache)
+docker run -d --name odc_dev \
+  --privileged --network host --ipc host --shm-size 64G \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  --group-add video --cap-add SYS_PTRACE --cap-add CAP_SYS_ADMIN \
+  --security-opt seccomp=unconfined --ulimit memlock=-1:-1 \
+  -v "$HOME":"$HOME" \
+  tasimage/primus-odc:v26.2 sleep infinity
+```
+
+> The mount/device/`--ipc host`/`--security-opt seccomp=unconfined` combination above matches the `odc_dev200` container we used to validate the run (`/dev/kfd`+`/dev/dri` for GPUs, `/dev/infiniband` for dual-node RoCE). **`-v` must be written as `src:dst`**; writing a bare path lets docker create an anonymous empty volume that shadows it.
+
+## Build Primus-Turbo with ODC operators (build from source + `.image_bak` consumption)
+
+> **Why you can't just `pip install @main` (the most important point in this section).** Although the operator source is on official `main` (PR #409), it is gated by `#ifndef DISABLE_ROCSHMEM`: `setup.py` only compiles the operators in when `ROCSHMEM_HOME` (+`MPI_HOME`) points at a **prebuilt rocSHMEM static library**, otherwise it adds `-DDISABLE_ROCSHMEM` and compiles the whole section out. The image has no rocSHMEM and pip does not set `ROCSHMEM_HOME`, so a plain `pip install @main` yields a hollow package with `has_host=False has_gda=False` (reproduced on three nodes). Switching branches with `pip … @feat/odc-rocshmem-dist` doesn't work either — that branch exists only on a personal fork, not in the official repo.
+
+**The correct approach, and the one used for all results in the blog: build Primus-Turbo from source, pointing `ROCSHMEM_HOME` at prebuilt rocSHMEM at build time, building one for single-node / one for dual-node.** rocSHMEM is a **non-pip external dependency**, whose `librocshmem.a` must be prebuilt separately for the host/XGMI-IPC and GDA/MLX5 paths (denoted `rocshmem_single` / `rocshmem_gda`). Run once **inside each node's container**:
+
+```bash
+# inside the container (using the image's own ROCm torch, hence --no-build-isolation, to avoid pip pulling a non-ROCm torch)
+export MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi     # openmpi shipped in the image
+export GPU_ARCHS="gfx942"                             # fill per card: gfx942 / gfx950
+
+git clone https://github.com/AMD-AGI/Primus-Turbo.git Primus-Turbo-single   # host/single-node build
+git clone https://github.com/AMD-AGI/Primus-Turbo.git Primus-Turbo          # GDA/dual-node build
+
+# host build (single-node XGMI-IPC): link the host rocSHMEM
+cd Primus-Turbo-single && ROCSHMEM_HOME=<path>/rocshmem_single \
+  pip3 install --no-build-isolation -e ".[pytorch]" -v && cd ..
+# GDA build (multi-node GPU-Direct): link the GDA rocSHMEM (--lto-partitions=1 already built into setup.py)
+cd Primus-Turbo && ROCSHMEM_HOME=<path>/rocshmem_gda \
+  pip3 install --no-build-isolation -e ".[pytorch]" -v && cd ..
+```
+
+After installation, the site-packages of the image `tasimage/primus-odc:v26.2` may still hold a **stock `primus_turbo` without ODC operators**, which will **import first and shadow** the dev tree you want to use. Use `.image_bak` to rename and disable it, letting the one on `PRIMUS_TURBO_PATH` win, then verify the operators are present:
+
+```bash
+# de-shadow: rename the stock primus_turbo (do this once in each node's container, not skippable)
+SP=$(python -c "import sysconfig;print(sysconfig.get_paths()['purelib'])")
+[ -e "$SP/primus_turbo" ] && [ ! -L "$SP/primus_turbo" ] && mv "$SP/primus_turbo" "$SP/primus_turbo.image_bak"
+for d in "$SP"/primus_turbo-*.dist-info; do [ -e "$d" ] && mv "$d" "$d.image_bak"; done
+
+# verify the ODC operators are present (host=single-node XGMI-IPC, gda=multi-node GPU-Direct) — both dev trees should be True True
+PYTHONPATH=<path>/Primus-Turbo-single python -c "import primus_turbo.pytorch._C as C; \
+  print('single:', hasattr(C,'odc_rocshmem_host'), hasattr(C,'odc_rocshmem_gda'))"
+PYTHONPATH=<path>/Primus-Turbo        python -c "import primus_turbo.pytorch._C as C; \
+  print('gda   :', hasattr(C,'odc_rocshmem_host'), hasattr(C,'odc_rocshmem_gda'))"
+# expected: both lines are: True True
+```
+
+- **Use a separate dev tree for single-node / dual-node** (`Primus-Turbo-single` / `Primus-Turbo`): the two link rocSHMEM with different symmetric-heap / transport backends (host/IPC vs GDA/MLX5), so **they cannot share the same package**; mixing them throws a symmetric-heap error (see the "Pitfalls checklist" below). At runtime, the single-node section below uses `PRIMUS_TURBO_PATH=<...>/Primus-Turbo-single`, and the dual-node section uses `<...>/Primus-Turbo`.
+- `PRIMUS_TURBO_PATH` **cannot be omitted**: `run_odc.sh` prepends it to `PYTHONPATH` so that `import primus_turbo` hits the dev tree with the ODC operators; relying on site-packages alone (the stock version or the pip @main version) yields `has_gda=False`.
+- The `.image_bak` de-shadowing **cannot be omitted either**: without renaming, the stock version in site-packages imports first and shadows the dev tree.
+
+## Turning on ODC = changing config items (the core of this refactor)
+
+**ODC's switches are now all in yaml config, no longer `export ODC_ENABLE=1`.** The relevant config items (defaults in `primus/configs/modules/megatron/trainer_base.yaml` and `sft_trainer.yaml`):
+
+| Config | Effect | Old env replaced | Default |
+|---|---|---|---|
+| `enable_odc` | ODC **master switch** (true routes gradients to P2P; false = pure FSDP2+RCCL, all ODC patches no-op) | `ODC_ENABLE` | `false` |
+| `odc_phase` | integration depth, `2` = full gradient routing (production value) | `ODC_PHASE` | `2` |
+| `enable_odc_lb_mini` | LB-Mini variable-length KK balancing of **data** (**decoupled** from `enable_odc`, can take effect independently): `enable_odc: true` → **decoupled** (ranks may have unequal microbatch counts = nopad, needs ODC P2P); `enable_odc: false` → **aligned** (`all_reduce(MAX)` forces equal microbatch counts across ranks, standard FSDP2+RCCL lockstep = the RCCL-safe "same-data" baseline) | `ODC_LB_MINI` / `LB_MINI_FORCE_DATA` | `false` |
+| `lb_mini_cost_model` / `lb_mini_max_token_len` | LB-Mini cost model / per-microbatch token cap | `LB_MINI_*` | `linear` / `0` |
+| `odc_p2p_backend` | P2P backend: `rocshmem` (primary in this post, validated) ｜ `mori` | `ODC_P2P_BACKEND` | `mori` |
+| `odc_rocshmem_gda` | rocSHMEM GPU-Direct (GDA) device path, **must be on for dual-node** | `ODC_ROCSHMEM_GDA` | `false` |
+| `odc_gda_defer_reduce` / `odc_gda_warmup_mode` / `odc_gda_stride_bytes` / `odc_gda_pipe` | GDA settle deferral / warmup / stride / pipeline depth | `ODC_GDA_*` | `auto` / `strided` / `65536` / `1` |
+| `odc_grad_spike_threshold` | step-skip threshold for occasional gradient spikes from async reduce | — | `1000.0` |
+
+Key points:
+
+- **The `odc_nopad` arm** = `enable_odc: true` + `enable_odc_lb_mini: true` + `odc_p2p_backend: rocshmem` (add `odc_rocshmem_gda: true` for dual-node). The two example configs `deepseek1.5B-odc-lbmini.yaml`, `qwen14B-odc-dn.yaml` **already include all of these** (`enable_odc: true` / `odc_phase: 2` / `enable_odc_lb_mini: true` / `odc_p2p_backend: rocshmem`, and the dual-node version also carries `odc_rocshmem_gda: true`) — **ready to use out of the box, no need to manually add the backend**. (Note that in `trainer_base.yaml` the **global default of `odc_p2p_backend` is still `mori`**, so if you write your own config remember to set `rocshmem`.)
+- **The `nccl_pad` aligned baseline arm** = `enable_odc: false` + **`enable_odc_lb_mini: true`** (standard FSDP2 + RCCL). ⚠️ **`enable_odc_lb_mini: true` cannot be omitted (key to a fair comparison)**: when `enable_odc: false`, it switches to **aligned mode** (`all_reduce(MAX)` forces equal microbatch counts across ranks so collectives can be safely called per-microbatch), letting the baseline consume **exactly the same LB-Mini variable-length data as `odc_nopad`** (the compliant replacement for the old env `LB_MINI_FORCE_DATA`).
+- **The `mori` backend has a known bug on dual-node; never use it for dual-node**; all ODC results in this post use `rocshmem`.
+
+Unified entry point `run_odc.sh` (same script for single-node / dual-node, used as a fallback path):
+
+```text
+run_odc.sh <mori|rocshmem> <pad|nopad> <exp_yaml_relpath> <exp_name> [KEY=VAL ...]
+```
+
+The 1st positional arg is only responsible for laying down the rocSHMEM runtime infra env (heap / bootstrap ifname / a fresh `TRITON_CACHE_DIR` per run); **the backend is actually selected by the yaml's `odc_p2p_backend`**; the 2nd positional arg `pad|nopad` is now a no-op (alignment/decoupling is now auto-derived from `enable_odc`); the remaining `KEY=VAL` only pass infrastructure env (ODC feature switches all go in the yaml). The script lays down `PYTHONPATH` (including the `odc_early` shim and `PRIMUS_TURBO_PATH`) and then calls `run_pretrain.sh`; it **does not pass through CLI overrides**, so change batch/switches in the yaml.
+
+## Single-node 1.5B `odc_nopad` (8 GPUs, host / XGMI-IPC)
+
+**① Config**: the example config `examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml` is **ready out of the box** — it includes `enable_odc: true` / `odc_phase: 2` / `enable_odc_lb_mini: true` / **`odc_p2p_backend: rocshmem`** / `lb_mini_cost_model: fit` / `global_batch_size: 16`, with no changes needed (`odc_p2p_backend: rocshmem` is built in; the global default in `trainer_base.yaml` is `mori`).
+
+For the aligned baseline arm, copy one and **only turn off `enable_odc`, keeping `enable_odc_lb_mini: true`** (so the baseline consumes the same variable-length data as odc_nopad, just with microbatch counts aligned):
+
+```bash
+cd Primus
+cp examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml \
+   examples/megatron/configs/MI355X/deepseek1.5B-nccl.yaml
+# in the overrides: of deepseek1.5B-nccl.yaml:
+#   enable_odc: false          # turn off ODC communication
+#   enable_odc_lb_mini: true   # keep it! => aligned mode, feed the same LB-Mini variable-length data (fair comparison)
+# (do not set odc_p2p_backend; the backend item is meaningless when enable_odc: false)
+```
+
+**② Start training** (inside the container, single-node 8 GPUs):
+
+```bash
+export PRIMUS_TURBO_PATH=/path/to/Primus-Turbo-single      # single-node IPC Turbo
+export NNODES=1 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR=localhost MASTER_PORT=29700
+export HF_HOME=$HOME/primus_packed/hf_home HF_HUB_OFFLINE=1 PRIMUS_SKIP_PIP=1
+cd Primus
+RUN=primus/core/odc/rocshmem_runtime/scripts/run_odc.sh
+
+# odc_nopad (rocSHMEM host/XGMI-IPC) — the 1st positional arg rocshmem lays down the rocSHMEM infra env,
+# the backend is actually selected by odc_p2p_backend: rocshmem in the yaml; the 2nd positional arg nopad is now a no-op
+bash "$RUN" rocshmem nopad examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml odc_nopad
+
+# nccl_pad aligned baseline (yaml: enable_odc: false + enable_odc_lb_mini: true
+#  → standard FSDP2 + RCCL, fed the same LB-Mini variable-length data as odc_nopad, with microbatch counts aligned)
+bash "$RUN" rocshmem nopad examples/megatron/configs/MI355X/deepseek1.5B-nccl.yaml nccl_pad
+```
+
+Single-node goes over intra-node XGMI + HIP-IPC, and reduce-scatter-accumulate is a device-side owner-pull (on-chip fp32 summation, no host watcher subprocess). To sweep different batches, change `global_batch_size` in the yaml (`run_odc.sh` does not pass through CLI overrides).
+
+## Dual-node 14B `odc_nopad` (16 GPUs, rocSHMEM GDA)
+
+Dual-node **starts one rank group per node** (NODE_RANK 0/1, 8 GPUs each → 16 ranks total), and rocSHMEM uses **uid-over-socket** bootstrap (after PR #864 it no longer uses MPI, launching purely via torchrun). The **preferred, and end-to-end validated in this post** dual-node path is `primus-cli slurm srun`: one line fans out to the two nodes, each `docker run --rm` starts a fresh container and then drives `torchrun`, **with no persistent container and no `rank_node.sh`** (if the cluster hostname maps to loopback and this can't be used, see the `srun --overlap` fallback at the end). Because the feature switches are now config items, the `--env` list is reduced to **only genuine infrastructure variables**: import paths, the NCCL control plane, the rocSHMEM runtime (heap/ifname/HCA/GID), the HF offline cache, and the SFT-skip trio for bypassing primus-cli's built-in hook.
+
+**① Config**: `qwen14B-odc-dn.yaml` is **ready out of the box** — it includes `enable_odc: true` / `odc_phase: 2` / `enable_odc_lb_mini: true` / **`odc_p2p_backend: rocshmem`** / **`odc_rocshmem_gda: true`**, with no changes needed (the `--odc_p2p_backend`/`--odc_rocshmem_gda` overrides in the command below are just redundant insurance). GDA's `warmup=strided` / `stride=65536` / `defer_reduce=auto` / `pipe=1` are all config defaults; `defer_reduce: auto` automatically defers the per-microbatch reduce-scatter to a single rendezvous at the end of the minibatch when `n_pes>local_world_size` (16>8), avoiding cross-node barrier deadlock for nopad variable-length microbatches.
+
+```bash
+# orchestration node. JOB = the Slurm job id already holding these two nodes (keepalive)
+JOB=<slurm_jobid>; PORT=29604
+ROOT=<PRIMUS_ROOT>                                       # the Primus worktree with #864
+TURBO=/path/to/Primus-Turbo                              # dual-node GDA Turbo (lto=1)
+PYDEPS=/path/to/pydeps                                   # flydsl, etc.
+SP=/opt/venv/lib/python3.12/site-packages
+HCA=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9   # adjust per cluster NIC
+TS=$(date +%Y%m%d_%H%M%S)
+
+# primus-cli's built-in pretrain hook lacks the SFT skip; pre-place a .done flag so it skips prep (see pitfall ③)
+touch /home/botahu/primus_packed/pcli_notok_bookcorpus.done
+
+cd "$ROOT"
+./primus-cli slurm srun --jobid="$JOB" --overlap -N2 --ntasks-per-node=1 \
+  -p amd-rccl --nodelist=<nodeA>,<nodeB> \
+  -- --image tasimage/primus-odc:v26.2 --clean \
+     --privileged --network host --ipc host --shm-size 64G \
+     --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+     --group-add video --cap-add SYS_PTRACE --cap-add CAP_SYS_ADMIN \
+     --security-opt seccomp=unconfined \
+     --volume /home/botahu:/home/botahu \
+     `# --- import paths / turbo / linker (primus-cli doesn't run setup_pythonpath, must fill in) ---` \
+     --env PYTHONPATH="$ROOT:$PYDEPS:$TURBO:$ROOT/primus/core/odc/odc_early:$ROOT/primus/core:$SP" \
+     --env PRIMUS_TURBO_PATH="$PYDEPS:$TURBO" \
+     --env LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu/openmpi/lib \
+     `# --- NCCL control plane (override the single-node lo / IB-off defaults) ---` \
+     --env NCCL_SOCKET_IFNAME=eth0 --env GLOO_SOCKET_IFNAME=eth0 \
+     --env NCCL_IB_DISABLE=0 --env NCCL_IB_GID_INDEX=3 \
+     `# --- rocSHMEM runtime infra (heap in raw bytes / bootstrap / provider / NIC / GID) ---` \
+     --env ROCSHMEM_HEAP_SIZE=8589934592 --env ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME=eth0 \
+     --env ROCSHMEM_GDA_PROVIDER=mlx5 --env ROCSHMEM_HCA_LIST="$HCA" \
+     --env ROCSHMEM_IB_GID_INDEX=3 --env ROCSHMEM_ROCE_GID_INDEX=3 --env ROCSHMEM_GID_INDEX=3 \
+     `# --- create a fresh triton cache per run (avoid reusing mismatched device kernels → NaN) ---` \
+     --env TRITON_CACHE_DIR=/tmp/tcache_rocshmem_pcli_$TS \
+     `# --- HF offline cache + node-local cache ---` \
+     --env HF_HOME=/home/botahu/primus_packed/hf_home \
+     --env HF_HUB_CACHE=/home/botahu/primus_packed/hf_models \
+     --env HF_DATASETS_CACHE=/home/botahu/primus_packed/hf_home/datasets \
+     --env HF_HUB_OFFLINE=1 --env DATA_PATH=/workspace \
+     --env PRIMUS_PACK_CACHE_DIR=/home/botahu/primus_packed \
+     --env PRIMUS_PACK_LOCK_DIR=/tmp/primus_lock \
+     --env PRIMUS_CACHE_ROOT=/tmp/primus_cache_pcli_$TS --env PRIMUS_SKIP_PIP=1 \
+     `# --- bypass the built-in pretrain hook's SFT prep (see pitfall ③) ---` \
+     --env TOKENIZED_DATA_PATH=/home/botahu/primus_packed/pcli_notok_bookcorpus \
+     --env HF_TOKEN=hf_offline_dummy \
+  -- -- train pretrain --config examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml \
+     --backend_path "$ROOT/third_party/Megatron-LM" \
+     `# --- ODC feature switches: config items, passed directly as CLI overrides (no longer via --env) ---` \
+     --odc_p2p_backend rocshmem --odc_rocshmem_gda true --enable_fused_linear_ce true \
+     --train_iters 50 --global_batch_size 128 --micro_batch_size 1 \
+     --use_torch_fsdp2 True --manual_gc True \
+     --profile false --use_pytorch_profiler false \
+     --disable_wandb True --disable_tensorboard True --disable_last_saving True
+```
+
+The three `--` in `primus-cli slurm srun` separate, in order: **slurm/srun args** ｜ **docker/container args (including all `--env`)** ｜ **`train pretrain` training args (including ODC config overrides)**. `slurm-entry` will automatically inject `MASTER_ADDR/MASTER_PORT/NNODES/NODE_RANK/GPUS_PER_NODE` into each node. `ROCSHMEM_HCA_LIST` is the RoCE NIC allowlist (this cluster has 10 mlx5 cards; `mlx5_1/_6` are eth management ports and must be excluded).
+
+**For the `nccl_pad` aligned baseline**: replace the ODC overrides in the last segment with `--enable_odc false --enable_odc_lb_mini true` (drop `--odc_p2p_backend`/`--odc_rocshmem_gda`, everything else unchanged), i.e., standard FSDP2 + RCCL. ⚠️ **`--enable_odc_lb_mini true` cannot be omitted**: when `enable_odc: false`, it automatically switches to **aligned mode**, feeding the baseline **exactly the same LB-Mini variable-length data as `odc_nopad`** (each rank's microbatch count aligned via `all_reduce(MAX)`, RCCL-safe).
+
+### Three iron rules of correctness (not skippable)
+
+- **RoCE-adjacent nodes**: be sure to pick adjacent nodes under the same leaf switch, mutually reachable over RoCE; non-adjacent node pairs often stall at the first cross-node GDA warmup (a long time with no `iteration`).
+- **device-LTO single-partition**: the GDA Turbo must be built with `-Xoffload-linker --lto-partitions=1` (built in by PR #409), otherwise cross-node `getmem` reads 0 → **dual-node `grad_norm=0`**.
+- **dedicated dual-node GDA Turbo**: use the `$TURBO` relinked to `rocshmem_gda`, kept separate from the single-node IPC build — the two have different symmetric-heap/connection backends, and mixing them clashes on the symmetric heap.
+
+### 3 pitfalls of the direct `primus-cli` path
+
+1. **`--volume` must be written as `src:dst`** (`--volume /home/botahu:/home/botahu`): writing a bare path lets docker create an anonymous empty volume that shadows it → the container can't see the code/cache.
+2. **`PYTHONPATH` must be injected in full**: the direct primus-cli path **skips `setup_pythonpath`**, so you must fill in repo + pydeps + **GDA Turbo** (`$TURBO`, placing it ahead of site-packages lets the GDA operators win, `has_gda=True`, without needing `.image_bak` in the ephemeral container) + `odc/odc_early` + `primus/core` + site-packages, otherwise `import odc`/`import primus_turbo` fails.
+3. **the built-in pretrain hook lacks the SFT skip**: the built-in runner hook (`train/pretrain/megatron/prepare.py`) has no `stage=='sft'` skip branch and will demand `HF_TOKEN` + download bookcorpus. Fix: pre-`touch` a `.done` flag and point `TOKENIZED_DATA_PATH` at it, give `HF_TOKEN` a placeholder value, so the hook skips prep (the `--train_data_path` it emits is harmless to the SFT provider).
+
+> **A `slurm-entry` MASTER_ADDR caveat**: `slurm-entry` forces `MASTER_ADDR=`short-hostname. In the cluster validated for this post, the `127.0.1.1` line in `/etc/hosts` is commented out and the short hostname resolves to a routable IP, so it works directly; if the cluster maps the hostname to loopback (`127.0.1.1 <hostname>`), rendezvous binds to the loopback and cross-node connections fail, in which case use the `srun --overlap` fallback below.
+
+### Fallback: `srun --overlap` + `docker exec` + `run_odc.sh`
+
+If the cluster hostname maps to loopback, or you want to reuse the environment preset in a persistent container, you can fall back to **`srun --overlap` + `docker exec` + `run_odc.sh`** (delivering one rank per node into the persistent container), running on each of the two nodes:
+
+```bash
+export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=<0|1> MASTER_ADDR=<nodeA-routable-IP> MASTER_PORT=<port>
+bash run_odc.sh rocshmem nopad examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml odc_nopad \
+  NCCL_SOCKET_IFNAME=eth0 GLOO_SOCKET_IFNAME=eth0 NCCL_IB_DISABLE=0 NCCL_IB_GID_INDEX=3 \
+  ROCSHMEM_GDA_PROVIDER=mlx5 ROCSHMEM_HCA_LIST=$HCA \
+  ROCSHMEM_IB_GID_INDEX=3 ROCSHMEM_ROCE_GID_INDEX=3 ROCSHMEM_GID_INDEX=3
+```
+
+Here `KEY=VAL` only passes infrastructure env, while ODC feature switches still come from the yaml (`qwen14B-odc-dn.yaml` already has them all); the baseline is run the same way with a copy that sets `enable_odc: false` + keeps `enable_odc_lb_mini: true`, via `... <baseline_yaml> nccl_pad`. It is more robust to hostname/loopback (the 0→1 first bring-up used this path), with a wrapping example in `full_exp/matrix_dual_harness/`.
+
+## Judging success, reading results, computing the speedup
+
+**Step one — confirm ODC really is running on rocSHMEM GDA.** In the training log, check these real markers (all below are taken from the actual log of this post's dual-node primus-cli validation run):
+
+1. **Config items in effect (enable_odc + backend)**:
+
+   ```text
+   [Primus:Runtime] Applying CLI overrides: {'odc_p2p_backend': 'rocshmem', 'odc_rocshmem_gda': True, ...}
+   ```
+
+2. **rocSHMEM GDA backend is up (fix ①)** — the runtime evidence that `has_gda=True` is truly in effect:
+
+   ```text
+   init_shmem (rocSHMEM GDA, uid bootstrap): my_pe=15 n_pes=16
+   [GDA] reduce-scatter warm-up mode=strided stride_bytes=65536
+   ```
+
+   You should also see a batch of `Rank N create tensor gda_acc_* / gda_in_* / gda_scr_*` GDA symmetric-heap tensors. If the backend runs as `mori` (no such `rocSHMEM GDA` / `gda_*` lines), it means `odc_p2p_backend` was not set correctly.
+
+3. **Hooked into the correct FSDP2 class (fix ②)** + **nopad decoupled data**:
+
+   ```text
+   [PRE-FORWARD canary] wrap-chain: PrimusTorchFullyShardedDataParallel.module=FSDPFloat16Module.module=GPTModel
+   [ODC.lb_mini] built LB-Mini train iterator: ... same_micro_num=False (LB-Mini decoupled (ODC))
+   [ODC.torch_fsdp2] hooked optimizer.step (pre_optimizer_step injected)
+   ```
+
+   The wrapper must be **`PrimusTorchFullyShardedDataParallel`** (not the old `TorchFullyShardedDataParallel`; hooking the old class crashes in iter2 with `NoneType ... clear_accumulations`); `same_micro_num=False` means nopad/decoupled.
+   > Note: the code also has `log_rank_0` markers like `[ODC.torch_fsdp2] runtime config populated ...`; they print at patch registration / early `before_train` and can corroborate but may not necessarily land in the torchrun capture stream, so rely on the lines above.
+
+4. **Training advances stably**: `lm loss` **decreases** with iter, **0 nan / 0 crash**, steady-state `ms/iter` is stable, and **`grad norm` is nonzero throughout** — `grad_norm=0` means cross-node GDA `getmem` reads 0 (device-LTO multi-partition), indicating the GDA Turbo was not built correctly.
+
+Quick grab:
+
+```bash
+grep -E "rocSHMEM GDA|gda_acc_|reduce-scatter warm-up|LB-Mini decoupled|PrimusTorchFullyShardedDataParallel" "$LOG"
+grep -E "lm loss|elapsed time per iteration|grad norm|nan iterations" "$LOG" | tail -40
+```
+
+**Step two — read the key metrics, compute the speedup.** Each step's log prints `iteration N/M | ... | elapsed time per iteration (ms): X | ... | lm loss: ...`, plus `compute per GPU (TFLOP/s/GPU)` and `tokens/s/GPU`. Take the median `ms/iter` of the stable segment (skip the first few warmup steps), then:
+
+> speedup = `nccl_pad's ms/iter` ÷ `odc_nopad's ms/iter` (>1 means faster than RCCL)
+
+**Success criteria**: 0 nan throughout, loss decreasing step by step, running the full target step count.
+
+**Final validation results** (config-driven: ODC via `enable_odc`, `nccl_pad` fed the same variable-length data via **aligned** `enable_odc_lb_mini`; strictly same-source apples-to-apples — same image v26.2, same nodes, differing only in the arm):
+
+| Scenario | odc_nopad (config) | aligned nccl_pad | Speedup | loss |
+|---|---|---|---|---|
+| Dual-node 14B, gbs128 | ≈ **70.5k** ms/iter (primus-cli measured median) | ≈ **80.2k** ms/iter | ≈ **1.15×** | 12.07→9.34, grad_norm nonzero throughout, 0 nan |
+| Single-node 1.5B, gbs16 | ≈ **2.1k** ms/iter | ≈ **2.5k** ms/iter | ≈ **1.19–1.21×** | converges to ~10.0, 0 nan |
+
+> The dual-node 14B row is end-to-end validated via this post's **`primus-cli slurm` path**: 0 nan throughout training, `grad_norm` nonzero, steady-state median ~70.5k ms/iter, consistent with the earlier `srun --overlap` fallback path (~71k) — the two paths behave equivalently.
+
+## Pitfalls checklist (strongly recommended reading first)
+
+- **`has_host=False has_gda=False` (fix ①)**: you're using stock turbo or the hollow package from `pip install @main` (operators compiled out by `-DDISABLE_ROCSHMEM`). The correct fix is in "Build Primus-Turbo with ODC operators": build from source, point `ROCSHMEM_HOME` at prebuilt rocSHMEM at build time, then confirm `has_gda==True` via `PRIMUS_TURBO_PATH` (or `.image_bak` de-shadowing in site-packages).
+- **iter2 crashes with `'NoneType' ... clear_accumulations` (fix ②)**: you hooked the old FSDP2 class. Confirm Primus is on the `feat/odc-consume-turbo` branch, and that `PrimusTorchFullyShardedDataParallel` appears in the wrap-chain.
+- **backend runs as mori**: the global default in `trainer_base.yaml` is `mori`, so a self-written config that forgets to set `rocshmem` will run mori (**dual-node mori has a known bug**); the absence of `rocSHMEM GDA`/`gda_*` lines in the log identifies it.
+- **NaN on the very first step**: reused a Triton cache from a heterogeneous toolchain. Use a **fresh timestamped** `TRITON_CACHE_DIR` per run.
+- **dual-node `grad_norm=0`**: device-LTO multi-partition; the dual-node GDA Turbo must use `-Xoffload-linker --lto-partitions=1` (built in by PR #409).
+- **rocSHMEM heap is a raw byte count**: `ROCSHMEM_HEAP_SIZE` does not accept K/M/G suffixes (8 GiB = `8589934592`).
+- **data packing is somewhat slow**: at large gbs (dual-node gbs128), LB-Mini packing may take about 10 minutes before the first `iteration`, which is normal, not a hang.
diff --git a/examples/run_k8s_pretrain.sh b/examples/run_k8s_pretrain.sh
index a57eda70d..ae07ab2bd 100644
--- a/examples/run_k8s_pretrain.sh
+++ b/examples/run_k8s_pretrain.sh
@@ -15,7 +15,7 @@ GPU="8"
 EXP_PATH=""
 DATA_PATH=""
 BACKEND="megatron"
-IMAGE="docker.io/rocm/primus:v26.3"
+IMAGE="docker.io/rocm/primus:v26.4"
 HF_TOKEN="${HF_TOKEN:-}"
 WORKSPACE="primus-safe-pretrain"
 NODELIST=""
@@ -38,7 +38,7 @@ Options for create:
     --backend <name>            Training backend, e.g. megatron | torchtitan(default: megatron)
     --exp <exp_path>            Path to EXP config (optional)
     --data_path <data_path>     Data path (optional)
-    --image <docker_image>      Docker image to use (default: docker.io/rocm/primus:v26.3)
+    --image <docker_image>      Docker image to use (default: docker.io/rocm/primus:v26.4)
     --hf_token <token>          HuggingFace token (default: from env HF_TOKEN)
     --workspace <workspace>     Workspace name (default: safe-cluster-dev)
     --nodelist <node1,node2>    Comma-separated list of node names to run on (optional)
diff --git a/examples/run_local_pretrain.sh b/examples/run_local_pretrain.sh
index f40292308..f0db4a280 100755
--- a/examples/run_local_pretrain.sh
+++ b/examples/run_local_pretrain.sh
@@ -16,7 +16,7 @@ Usage: bash run_local_pretrain.sh
 This script launches a Primus pretraining task inside a Docker/Podman container.
 
 Environment Variables:
-    DOCKER_IMAGE   Docker image to use [Default: docker.io/rocm/primus:v26.3]
+    DOCKER_IMAGE   Docker image to use [Default: docker.io/rocm/primus:v26.4]
     MASTER_ADDR    Master node IP or hostname [Default: localhost]
     MASTER_PORT    Master node port [Default: 1234]
     NNODES         Total number of nodes [Default: 1]
@@ -44,7 +44,7 @@ export EXP=${EXP:-"examples/megatron/exp_pretrain.yaml"}
 if [ "${BACKEND:-}" = "MaxText" ]; then
     DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/jax-training:maxtext-v26.2"}
 else
-    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"}
+    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
 fi
 
 # Project root
diff --git a/examples/run_local_pretrain_cli.sh b/examples/run_local_pretrain_cli.sh
index 7e6305c9b..df930b077 100755
--- a/examples/run_local_pretrain_cli.sh
+++ b/examples/run_local_pretrain_cli.sh
@@ -15,7 +15,7 @@ EXP=${EXP:-"examples/megatron/exp_pretrain.yaml"}
 if [ "${BACKEND:-}" = "MaxText" ]; then
     DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/jax-training:maxtext-v25.9"}
 else
-    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"}
+    DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
 fi
 
 # ------------------ Cluster Env Defaults ------------------
diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh
index d93d251cb..70edcc698 100755
--- a/examples/run_pretrain.sh
+++ b/examples/run_pretrain.sh
@@ -1,6 +1,6 @@
 #!/bin/bash
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -89,13 +89,21 @@ PRIMUS_PATH=$(realpath "$(dirname "$0")/..")
 export DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"}
 export HF_HOME=${HF_HOME:-"${DATA_PATH}/huggingface"}
 
+# shellcheck source=/dev/null
 source "${PRIMUS_PATH}/runner/helpers/envs/path_utils.sh"
 
-LOG_INFO_RANK0 "Pip installing required packages ..."
-if [ "${BACKEND:-}" != "MaxText" ]; then
-    pip install -r "$PRIMUS_PATH/requirements.txt"  --quiet
+# PRIMUS_SKIP_PIP=1 skips the per-run pip install (deps already ship in the base
+# image). Use it when many ranks/nodes share one venv and concurrent pip installs
+# would race, or simply to speed up a warm container. Not keyed on the launcher.
+if [ "${PRIMUS_SKIP_PIP:-0}" == "1" ]; then
+    LOG_INFO_RANK0 "PRIMUS_SKIP_PIP=1: skipping pip install (deps from image)"
 else
-    pip install -r "$PRIMUS_PATH/requirements-jax.txt"  --quiet
+    LOG_INFO_RANK0 "Pip installing required packages ..."
+    if [ "${BACKEND:-}" != "MaxText" ]; then
+        pip install -r "$PRIMUS_PATH/requirements.txt"  --quiet
+    else
+        pip install -r "$PRIMUS_PATH/requirements-jax.txt"  --quiet
+    fi
 fi
 
 
@@ -117,7 +125,27 @@ if [ ! -f "${EXP}" ]; then
     exit 1
 fi
 
-TRAIN_LOG=${TRAIN_LOG:-"output/log_mp_pretrain_$(basename "$EXP" .yaml).txt"}
+# -------------------- Hybrid model (GDN/KDA/Mamba) FLA Triton patch --------------------
+# The zebra_llama_* hybrid configs under examples/megatron/configs/MI300X (added in
+# PR #676) depend on flash-linear-attention (FLA) Triton kernels that hang during
+# autotuning with num_stages >= 3 on MI300X/ROCm. Auto-apply the workaround whenever
+# one of those configs is used. See tools/hybrid/patch_fla_triton_autotune_hang.sh for details.
+EXP_BASENAME=$(basename "${EXP}")
+if [[ "${EXP_BASENAME}" == zebra_llama_* ]]; then
+    LOG_INFO_RANK0 "Detected hybrid model config (${EXP_BASENAME}); applying FLA Triton autotune-hang patch ..."
+    if ! bash "${PRIMUS_PATH}/tools/hybrid/patch_fla_triton_autotune_hang.sh"; then
+        LOG_ERROR "FLA Triton autotune-hang patch failed; continuing, but training may hang during autotuning."
+    fi
+fi
+
+if [ -z "${TRAIN_LOG:-}" ]; then
+    RUN_FOLDER=$(python3 -c "
+import yaml, sys
+d = yaml.safe_load(open(sys.argv[1]))
+print(f\"{d.get('workspace','./output')}/{d.get('work_group','default')}/{d.get('user_name','unknown')}/{d.get('exp_name','experiment')}\")
+" "$EXP" 2>/dev/null || echo "output")
+    TRAIN_LOG="${RUN_FOLDER}/train.log"
+fi
 
 LOG_INFO_RANK0 "==========Training info=========="
 LOG_INFO_RANK0 "EXP: $EXP"
@@ -139,6 +167,9 @@ export NCCL_CHECKS_DISABLE=1
 
 if [ "$USING_AINIC" == "1" ]; then
     LOG_INFO_RANK0 "Using AINIC"
+    # ROCm 7.2.1+ builds ANP into RCCL; RCCL_AINIC_ROCE=1 enables the built-in ANP
+    # path. An explicitly-set NCCL_NET_PLUGIN below still takes precedence.
+    export RCCL_AINIC_ROCE="${RCCL_AINIC_ROCE:-1}"
     export NCCL_IB_GID_INDEX=1
     export NCCL_IB_TC=${NCCL_IB_TC:-104}
     export NCCL_IB_FIFO_TC=${NCCL_IB_FIFO_TC:-192}
@@ -158,7 +189,12 @@ if [ "$USING_AINIC" == "1" ]; then
         "${RCCL_HOME_DIR}/build/release" \
         "${ANP_HOME_DIR}/build" \
         "${MPI_HOME_DIR}/lib"
-    if [ -n "${NCCL_NET_PLUGIN:-}" ]; then
+    # With built-in ANP (RCCL_AINIC_ROCE=1) default to "none" so RCCL uses its
+    # built-in ANP transport instead of auto-loading an external plugin. An
+    # explicitly-set NCCL_NET_PLUGIN always wins.
+    if [ "${RCCL_AINIC_ROCE}" = "1" ]; then
+        export NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-none}"
+    elif [ -n "${NCCL_NET_PLUGIN:-}" ]; then
         export NCCL_NET_PLUGIN
     elif [ -f "${ANP_HOME_DIR}/build/librccl-anp.so" ]; then
         export NCCL_NET_PLUGIN=librccl-anp.so
@@ -195,7 +231,16 @@ LOG_INFO "GLOO_SOCKET_IFNAME: $GLOO_SOCKET_IFNAME"
 LOG_INFO ""
 
 # ----------------- AMD-specific GPU optimizations -----------------
-export HSA_ENABLE_SDMA=1
+
+# Enable system DMA engine (SDMA) on AMD GPUs for better IO throughput.
+# ${:-} guarded so a caller can set HSA_ENABLE_SDMA=0 to route copies through
+# blit/compute kernels instead — workaround for hosts where an SDMA H2D copy
+# intermittently never signals completion (hangs in BusyWaitSignal).
+export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-1}
+
+# Prevent scratch memory from being reclaimed to stabilize large memory usage patterns (e.g., KV cache, MoE experts)
+# NOTE: Must disable scratch reclaim to avoid MoE training crash on AMD GPUs
+# Setting this to 0 prevents core dumps when using Mixture-of-Experts (MoE) models
 export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-0}
 export RCCL_MSCCL_ENABLE=0
 export RCCL_MSCCLPP_ENABLE=0
diff --git a/examples/run_slurm_pretrain_cli.sh b/examples/run_slurm_pretrain_cli.sh
index 6d10d0231..bdfa8850a 100755
--- a/examples/run_slurm_pretrain_cli.sh
+++ b/examples/run_slurm_pretrain_cli.sh
@@ -35,7 +35,7 @@ mkdir -p "$LOG_DIR"
 # NOTE: The --env entries below are passed into the container and will be visible
 # to the Primus training process (and system hooks) inside the container.
 bash "$PRIMUS_PATH/runner/primus-cli" slurm "${SLURM_ARGS[@]}" \
--- --image "${DOCKER_IMAGE:-rocm/primus:v26.3}" \
+-- --image "${DOCKER_IMAGE:-rocm/primus:v26.4}" \
 -- \
   --env "USING_AINIC=${USING_AINIC:-0}" \
   --env "PATCH_TE_FLASH_ATTN=${PATCH_TE_FLASH_ATTN:-0}" \
diff --git a/examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml
index 489e0c5cd..14d8d684c 100644
--- a/examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
-        local_batch_size: 8
+        local_batch_size: 6
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml
index bca6beaf8..5af0eadb8 100644
--- a/examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI300X/deepseek_v3_236b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/deepseek_v3_236b-BF16-pretrain.yaml
index 61ac2a2b1..2a9a61ac9 100644
--- a/examples/torchtitan/configs/MI300X/deepseek_v3_236b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI300X/deepseek_v3_236b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI300X/deepseek_v3_236b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/deepseek_v3_236b-FP8-pretrain.yaml
index 1a4d9214b..8a357ac88 100644
--- a/examples/torchtitan/configs/MI300X/deepseek_v3_236b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI300X/deepseek_v3_236b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml b/examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml
index cfa6d5e27..48e9e5f73 100644
--- a/examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml
+++ b/examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 16
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml
new file mode 100644
index 000000000..3c923f249
--- /dev/null
+++ b/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml
@@ -0,0 +1,57 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 1
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml
new file mode 100644
index 000000000..ae3fe1c10
--- /dev/null
+++ b/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 1
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml
new file mode 100644
index 000000000..2e87bd229
--- /dev/null
+++ b/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 4
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "selective"        # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention is installed by the
+      # gptoss_sink_attention setup patch. BF16 (no fp8/mx).
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml
new file mode 100644
index 000000000..65593906e
--- /dev/null
+++ b/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 6
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "selective"        # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI325X/deepseek_v3_16b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/deepseek_v3_16b-BF16-pretrain.yaml
index 8a934578b..0a52da5e2 100644
--- a/examples/torchtitan/configs/MI325X/deepseek_v3_16b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI325X/deepseek_v3_16b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 20
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI325X/deepseek_v3_16b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/deepseek_v3_16b-FP8-pretrain.yaml
index bca6beaf8..5af0eadb8 100644
--- a/examples/torchtitan/configs/MI325X/deepseek_v3_16b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI325X/deepseek_v3_16b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI325X/deepseek_v3_236b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/deepseek_v3_236b-BF16-pretrain.yaml
index 61ac2a2b1..2a9a61ac9 100644
--- a/examples/torchtitan/configs/MI325X/deepseek_v3_236b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI325X/deepseek_v3_236b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI325X/deepseek_v3_236b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/deepseek_v3_236b-FP8-pretrain.yaml
index 1a4d9214b..8a357ac88 100644
--- a/examples/torchtitan/configs/MI325X/deepseek_v3_236b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI325X/deepseek_v3_236b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI325X/deepseek_v3_671b-pretrain.yaml b/examples/torchtitan/configs/MI325X/deepseek_v3_671b-pretrain.yaml
index cfa6d5e27..48e9e5f73 100644
--- a/examples/torchtitan/configs/MI325X/deepseek_v3_671b-pretrain.yaml
+++ b/examples/torchtitan/configs/MI325X/deepseek_v3_671b-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 16
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml
new file mode 100644
index 000000000..24810a1f4
--- /dev/null
+++ b/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml
@@ -0,0 +1,57 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 2
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml
new file mode 100644
index 000000000..ea03d4a13
--- /dev/null
+++ b/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 2
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml
new file mode 100644
index 000000000..c071aa0c3
--- /dev/null
+++ b/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 6
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "selective"        # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention is installed by the
+      # gptoss_sink_attention setup patch. BF16 (no fp8/mx).
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml
new file mode 100644
index 000000000..ade0bd8e3
--- /dev/null
+++ b/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 8
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "selective"        # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI355X/deepseek_v3_16b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/deepseek_v3_16b-BF16-pretrain.yaml
index d25fbc830..f3ecb0bbb 100644
--- a/examples/torchtitan/configs/MI355X/deepseek_v3_16b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI355X/deepseek_v3_16b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
-        local_batch_size: 13
+        local_batch_size: 9
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 50
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI355X/deepseek_v3_16b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/deepseek_v3_16b-FP8-pretrain.yaml
index 6ac6a979a..ee423f547 100644
--- a/examples/torchtitan/configs/MI355X/deepseek_v3_16b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI355X/deepseek_v3_16b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 14
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI355X/deepseek_v3_236b-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/deepseek_v3_236b-BF16-pretrain.yaml
index 61ac2a2b1..2a9a61ac9 100644
--- a/examples/torchtitan/configs/MI355X/deepseek_v3_236b-BF16-pretrain.yaml
+++ b/examples/torchtitan/configs/MI355X/deepseek_v3_236b-BF16-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI355X/deepseek_v3_236b-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/deepseek_v3_236b-FP8-pretrain.yaml
index 1a4d9214b..8a357ac88 100644
--- a/examples/torchtitan/configs/MI355X/deepseek_v3_236b-FP8-pretrain.yaml
+++ b/examples/torchtitan/configs/MI355X/deepseek_v3_236b-FP8-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 8
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI355X/deepseek_v3_671b-pretrain.yaml b/examples/torchtitan/configs/MI355X/deepseek_v3_671b-pretrain.yaml
index cfa6d5e27..48e9e5f73 100644
--- a/examples/torchtitan/configs/MI355X/deepseek_v3_671b-pretrain.yaml
+++ b/examples/torchtitan/configs/MI355X/deepseek_v3_671b-pretrain.yaml
@@ -38,13 +38,15 @@ modules:
         min_lr_factor: 0.1
 
       training:
-        debug_moe_force_load_balance: true
         local_batch_size: 16
         seq_len: 4096
         max_norm: 1.0            # grad norm clipping
         steps: 15
         dataset: "c4_test"            # supported datasets: c4_test (2K), c4 (177M)
 
+      debug:
+        moe_force_load_balance: true
+
       parallelism:
         data_parallel_replicate_degree: 1
         data_parallel_shard_degree: -1
diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml
new file mode 100644
index 000000000..24810a1f4
--- /dev/null
+++ b/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml
@@ -0,0 +1,57 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 2
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml
new file mode 100644
index 000000000..ea03d4a13
--- /dev/null
+++ b/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_120b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 1.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 20
+
+      training:
+        local_batch_size: 2
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "always"   # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "full"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml
new file mode 100644
index 000000000..e0db78d38
--- /dev/null
+++ b/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 8
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "none"             # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention is installed by the
+      # gptoss_sink_attention setup patch. BF16 (no fp8/mx).
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: false
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml
new file mode 100644
index 000000000..c271776db
--- /dev/null
+++ b/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml
@@ -0,0 +1,58 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain}
+workspace: ./output
+
+modules:
+  pre_trainer:
+    framework: torchtitan
+    config: pre_trainer.yaml
+
+    model: gpt_oss_20b-fp8.yaml
+    overrides:
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      optimizer:
+        name: "AdamW"
+        lr: 2.2e-4
+        eps: 1.0e-8
+
+      lr_scheduler:
+        warmup_steps: 10
+
+      training:
+        local_batch_size: 10
+        seq_len: 4096
+        max_norm: 1.0            # grad norm clipping
+        steps: 50
+
+      parallelism:
+        data_parallel_replicate_degree: 1
+        data_parallel_shard_degree: -1
+        fsdp_reshard_after_forward: "default"  # default / never / always
+        tensor_parallel_degree: 1
+        pipeline_parallel_degree: 1
+        expert_parallel_degree: 8
+        expert_tensor_parallel_degree: 1
+
+      activation_checkpoint:
+        mode: "selective"        # ["none", "selective", "full"]
+        selective_ac_option: "op"
+
+      compile:
+        enable: true
+        components: ["model", "loss"]
+
+      # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch.
+      # FP8 dense linears via the quantize.linear.float8 converter.
+      primus_turbo:
+        enable_primus_turbo: true
+        use_turbo_attention: true
+        enable_attention_float8: false
+        use_turbo_float8_linear: true
+        use_turbo_mx_linear: false
+        use_moe_fp8: false
+        use_turbo_grouped_mm: false
+        use_classic_attention: false
diff --git a/examples/torchtitan/prepare.py b/examples/torchtitan/prepare.py
index aeff9043a..1777ab5ff 100644
--- a/examples/torchtitan/prepare.py
+++ b/examples/torchtitan/prepare.py
@@ -43,8 +43,16 @@ def parse_args():
 
 
 def pip_install_editable(path: Path, name: str):
-    log_info(f"Installing {name} in editable mode via pip (path: {path})")
-    ret = subprocess.run(["pip", "install", "-e", ".", "-q"], cwd=path)
+    # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools,
+    # no __init__.py) that the default editable finder cannot resolve. Use compat
+    # mode (writes a .pth with the source root on sys.path) after uninstalling any
+    # stale torchtitan, so the training subprocess can import them.
+    log_info(f"Installing {name} in editable (compat) mode via pip (path: {path})")
+    subprocess.run(["pip", "uninstall", "-y", name.lower()], cwd=path)
+    ret = subprocess.run(
+        ["pip", "install", "-e", ".", "--config-settings", "editable_mode=compat", "-q"],
+        cwd=path,
+    )
     if ret.returncode != 0:
         log_error_and_exit(f"Failed to install {name} via pip.")
 
diff --git a/primus/__init__.py b/primus/__init__.py
index 3f7299d9d..aa2df6eb4 100644
--- a/primus/__init__.py
+++ b/primus/__init__.py
@@ -1,4 +1,4 @@
 from .core.launcher.initialize import init
 from .core.utils.global_vars import get_primus_config
 
-__version__ = "0.2.0"
+__version__ = "26.4"
diff --git a/primus/_thirdparty.lock b/primus/_thirdparty.lock
index 465d3a116..60471467b 100644
--- a/primus/_thirdparty.lock
+++ b/primus/_thirdparty.lock
@@ -4,7 +4,7 @@
       "name": "torchtitan",
       "path": "third_party/torchtitan",
       "url": "https://github.com/pytorch/torchtitan.git",
-      "commit": "5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021"
+      "commit": "73a0e6979dd10b6b1904098eb3c8f62c18ab87ce"
     },
     {
       "name": "Megatron-LM",
@@ -16,7 +16,7 @@
       "name": "maxtext",
       "path": "third_party/maxtext",
       "url": "https://github.com/ROCm/maxtext.git",
-      "commit": "851c0935a4e9c9e4cd0771603e3bdf75a6578ad1"
+      "commit": "a7c6c7e5377cce39be51cbe2a48fe0f86fbab8dc"
     },
     {
       "name": "Emerging-Optimizers",
@@ -30,11 +30,23 @@
       "url": "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git",
       "commit": "9577b1280eaadd60b9d7b0ce6df09ac80e87e323"
     },
+    {
+      "name": "mamba",
+      "path": "third_party/mamba",
+      "url": "https://github.com/AndreasKaratzas/mamba.git",
+      "commit": "4d67c534d10b7299194ede55b36718cc7a1dd472"
+    },
     {
       "name": "HummingbirdXT",
       "path": "third_party/HummingbirdXT",
       "url": "https://github.com/AMD-AGI/HummingbirdXT.git",
       "commit": "ed7b7bd0dd03ed084f645c5dda1bc7b333f326f0"
+    },
+    {
+      "name": "Automodel",
+      "path": "third_party/Automodel",
+      "url": "https://github.com/NVIDIA-NeMo/Automodel.git",
+      "commit": "44f2acdee1628b8a5a5acb35ada6ed8c539a22e4"
     }
   ]
 }
diff --git a/primus/agents/tuning_agent/README.md b/primus/agents/tuning_agent/README.md
index e4df0a730..ed7543714 100644
--- a/primus/agents/tuning_agent/README.md
+++ b/primus/agents/tuning_agent/README.md
@@ -6,7 +6,7 @@ provides two estimates — **memory** and **performance** — each of which runs
 **benchmark-anchored by default** (measure what fits on a sub-node run, scale
 the rest analytically) with a fully analytical **no-GPU `simulate`** fallback.
 
-See [`docs/tuning_agent.md`](../../../docs/tuning_agent.md) — the full
+See [`docs/02-user-guide/tuning-agent.md`](../../../docs/02-user-guide/tuning-agent.md) — the full
 user/operator guide (modes, configuration, troubleshooting, worked example)
 plus the design write-up, paper-ready problem statement, and the list of
 deferred future features (cluster-spec retrieval, persistent memory cache,
diff --git a/primus/agents/tuning_agent/evaluator.py b/primus/agents/tuning_agent/evaluator.py
index 4aa2330c9..2966545a4 100644
--- a/primus/agents/tuning_agent/evaluator.py
+++ b/primus/agents/tuning_agent/evaluator.py
@@ -227,7 +227,7 @@ def write_trial_yaml(arch: ArchitectureRecord, cfg: TrialConfig, out_dir: Path,
     # enable ``use_turbo_gemm `` inside Primus's projection. Two
     # things go wrong on the v26.2 container ``primus_turbo==0.2.0``:
     #   * The default ``fp8_recipe: delayed`` is incompatible with that path
-    #     (``primus/modules/trainer/megatron/utils.py:464`` asserts).
+    #     (``primus/backends/megatron/patches/args/rocm_arg_validation.py`` asserts).
     #   * The dense FP8 GEMM op (``primus_turbo.pytorch.ops.gemm_fp8``) on
     #     this version raises ``ValueError: Unsupported FP8 format: HYBRID``
     #     for ``fp8: hybrid`` (a common DSv3 / Kimi-K2 configuration).
diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md
new file mode 100644
index 000000000..88e29227b
--- /dev/null
+++ b/primus/backends/diffusion/README.md
@@ -0,0 +1,76 @@
+# Primus Diffusion Backend
+
+`diffusion` is an in-tree PyTorch backend for Wan video training and FLUX
+text-to-image training. Primus owns config loading and launch; this backend owns
+model construction, datasets, attention selection, FSDP2 wrapping, and the
+training loop.
+
+## Supported Scope
+
+- Models: `wan`, `flux.1-schnell`, and `flux.1-dev`.
+- Trainer: FSDP2.
+- Wan sequence parallelism: supported through `sp_size`.
+- FLUX sequence parallelism: not supported; keep `sp_size: 1`.
+
+Install backend dependencies with:
+
+```bash
+pip install -r runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt
+```
+
+## Model Presets
+
+Model presets live under `primus/configs/models/diffusion/`.
+
+```text
+wan2.1_t2v_1.3b.yaml
+wan2.1_t2v_1.3b_sft.yaml
+wan2.2_ti2v_5b.yaml
+wan2.2_ti2v_5b_sft.yaml
+flux.1_schnell_t2i.yaml
+flux.1_dev_t2i.yaml
+```
+
+FLUX.1-schnell and FLUX.1-dev are separate presets. They share the same main
+transformer shape, but `flux.1-dev` has a guidance embedding module while
+`flux.1-schnell` does not. Select the preset by choosing
+`flux.1_schnell_t2i.yaml` or `flux.1_dev_t2i.yaml`.
+
+This follows the TorchTitan-style separation where `flux_schnell()` is its own
+config/preset and architecture differences such as `guidance_embed=False` are
+part of that preset, not a launch-time switch.
+
+## Public Config Sections
+
+Diffusion examples use Primus override sections that the backend converts into
+model, dataset, and trainer args:
+
+```yaml
+training:
+  local_batch_size: 1
+  steps: 50
+  output_dir: ./output/run
+
+data:
+  dataset_path: /path/to/data
+
+parallelism:
+  sp_size: 1
+  dp_replicate: 1
+
+optimizer:
+  lr: 2.0e-4
+  weight_decay: 0.1
+
+lr_scheduler:
+  lr_scheduler_type: constant_with_warmup
+  warmup_steps: 1600
+
+runtime:
+  attention_backend: flash_attn_aiter
+  gradient_checkpointing: false
+  report_to: none
+```
+
+`examples/diffusion/README.md` contains data preparation notes and launch
+commands for the shipped examples.
diff --git a/primus/backends/diffusion/__init__.py b/primus/backends/diffusion/__init__.py
new file mode 100644
index 000000000..c47954e23
--- /dev/null
+++ b/primus/backends/diffusion/__init__.py
@@ -0,0 +1,12 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus-owned diffusion backend components for t2i/t2v training."""
+
+from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter
+from primus.core.backend.backend_registry import BackendRegistry
+
+BackendRegistry.register_adapter("diffusion", DiffusionAdapter)
diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py
new file mode 100644
index 000000000..92f6a0a70
--- /dev/null
+++ b/primus/backends/diffusion/argument_builder.py
@@ -0,0 +1,396 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import copy
+from types import SimpleNamespace
+from typing import Any
+
+from primus.core.utils.yaml_utils import nested_namespace_to_dict
+
+
+class DiffusionArgBuilder:
+    """Build the compact config object consumed by diffusion trainers."""
+
+    DEFAULT_DATASET: dict[str, Any] = {
+        "name": "wan",
+        "config": {
+            "dataset_type": "vision",
+            "dataset_format": "jsonl",
+            "dataset_path": "/path/to/meta.jsonl",
+            "data_folder": "/path/to/videos",
+            "video_sampling_strategy": "frame_num",
+            "frame_num": 81,
+            "shuffle": True,
+            "video_backend": "imageio",
+            "processor_config": {
+                "processor_name": "wanvideo",
+                "processor_type": "wanvideo",
+                "max_text_length": 512,
+                "text_tokenizer": "/path/to/umt5-xxl",
+                "extra_kwargs": {
+                    "do_resize": True,
+                    "size": {
+                        "height": 480,
+                        "width": 832,
+                    },
+                    "do_normalize": True,
+                    "image_mean": [0.5, 0.5, 0.5],
+                    "image_std": [0.5, 0.5, 0.5],
+                },
+            },
+        },
+    }
+    DEFAULT_WAN_TRAINER: dict[str, Any] = {
+        "name": "fsdp2",
+        "args": {
+            "output_dir": "./output/wan",
+            "per_device_train_batch_size": 1,
+            "per_device_eval_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "gradient_checkpointing": True,
+            "attention_backend": "flash_attn_aiter",
+            "learning_rate": 1.0e-5,
+            "lr_scheduler_type": "constant",
+            "warmup_steps": 0,
+            "weight_decay": 0.01,
+            "num_train_epochs": 1,
+            "max_steps": 100,
+            "logging_steps": 1,
+            "save_steps": 0,
+            "dataloader_num_workers": 4,
+            "report_to": "none",
+            "run_name": "wan-fsdp2",
+            "bf16": True,
+            "seed": 10007,
+            "optim": "adamw_torch",
+            "adam_beta1": 0.9,
+            "adam_beta2": 0.999,
+            "adam_epsilon": 1.0e-8,
+            "max_grad_norm": 1.0,
+            "fsdp2_wrap_target": "dit",
+            "fsdp_transformer_layer_cls_to_wrap": "DiTBlock",
+            "fsdp2_reshard_after_forward": True,
+            "save_strategy": "dit_only",
+            "sp_size": 1,
+            "dp_replicate": 1,
+            "flow_match_scheduler": {
+                "shift": 5,
+                "sigma_min": 0.0,
+                "extra_one_step": True,
+                "num_train_timesteps": 1000,
+            },
+        },
+    }
+    DEFAULT_FLUX_DATASET: dict[str, Any] = {
+        "name": "flux",
+        "config": {
+            "dataset_type": "precomputed",
+            "dataset_format": "hf_dataset",
+            "dataset_path": "/path/to/flux_precomputed_dataset",
+            "dataset": None,
+            "shuffle": True,
+            "processor_config": {
+                "processor_name": "flux_precomputed",
+                "processor_type": "flux_precomputed",
+                "prompt_dropout_prob": 0.0,
+                "empty_encodings_path": None,
+                "img_size": 256,
+                "skip_low_resolution": True,
+            },
+        },
+    }
+    DEFAULT_FLUX_TRAINER: dict[str, Any] = {
+        "name": "fsdp2",
+        "args": {
+            "output_dir": "./output/flux",
+            "per_device_train_batch_size": 1,
+            "per_device_eval_batch_size": 1,
+            "gradient_accumulation_steps": 1,
+            "gradient_checkpointing": False,
+            "attention_backend": "flash_attn_aiter",
+            "learning_rate": 2.0e-4,
+            "lr_scheduler_type": "constant_with_warmup",
+            "warmup_steps": 1600,
+            "weight_decay": 0.1,
+            "num_train_epochs": 1,
+            "max_steps": 100,
+            "logging_steps": 1,
+            "save_steps": 0,
+            "dataloader_num_workers": 4,
+            "report_to": "none",
+            "run_name": "flux-schnell-fsdp2",
+            "bf16": True,
+            "seed": 10007,
+            "optim": "adamw_torch",
+            "adam_beta1": 0.9,
+            "adam_beta2": 0.95,
+            "adam_epsilon": 1.0e-8,
+            "max_grad_norm": 1.0,
+            "fsdp2_wrap_target": "dit",
+            "fsdp_transformer_layer_cls_to_wrap": "DoubleStreamBlock,SingleStreamBlock",
+            "fsdp_module_paths_to_wrap": "img_in,time_in,vector_in,txt_in,final_layer",
+            "fsdp_module_paths_no_reshard": "final_layer",
+            "fsdp2_reshard_after_forward": True,
+            "compile_transformer_blocks": True,
+            "save_strategy": "dit_only",
+            "sp_size": 1,
+            "dp_replicate": 1,
+            "flow_match_scheduler": {
+                "shift": 1,
+                "sigma_min": 0.0,
+                "extra_one_step": False,
+                "num_train_timesteps": 1000,
+            },
+        },
+    }
+
+    def __init__(self) -> None:
+        self._params: dict[str, Any] = {}
+
+    def update(self, params: Any) -> None:
+        if isinstance(params, SimpleNamespace):
+            self._params = nested_namespace_to_dict(params)
+        elif isinstance(params, dict):
+            self._params = copy.deepcopy(params)
+        else:
+            raise TypeError(
+                f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}"
+            )
+
+    def finalize(self) -> SimpleNamespace:
+        params = copy.deepcopy(self._params)
+        if "model" not in params:
+            raise ValueError("Diffusion backend config requires a model preset.")
+        for legacy_section in ("dataset", "trainer"):
+            if legacy_section in params:
+                raise ValueError(
+                    f"Diffusion backend no longer accepts public `{legacy_section}` overrides. "
+                    "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, "
+                    "`runtime`, and `metrics` sections instead."
+                )
+
+        params = self._normalize_primus_style_sections(params)
+
+        return SimpleNamespace(
+            model=params["model"],
+            dataset=params["dataset"],
+            trainer=params["trainer"],
+            stage=params.get("stage", "pretrain"),
+            primus=params.get("primus", {}),
+        )
+
+    @staticmethod
+    def _set_nested(target: dict[str, Any], path: tuple[str, ...], value: Any) -> None:
+        cursor = target
+        for key in path[:-1]:
+            next_value = cursor.get(key)
+            if not isinstance(next_value, dict):
+                next_value = {}
+                cursor[key] = next_value
+            cursor = next_value
+        cursor[path[-1]] = value
+
+    @staticmethod
+    def _get_any(source: dict[str, Any], *keys: str) -> Any:
+        for key in keys:
+            if key in source:
+                return source[key]
+        return None
+
+    @staticmethod
+    def _coerce_bool(value: Any) -> bool:
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            normalized = value.strip().lower()
+            if normalized in {"1", "true", "yes", "on"}:
+                return True
+            if normalized in {"0", "false", "no", "off"}:
+                return False
+        return bool(value)
+
+    def _model_family(self, params: dict[str, Any]) -> str:
+        model = params.get("model") or {}
+        if not isinstance(model, dict):
+            raise TypeError(f"Diffusion model preset must be a dict, got {type(model).__name__}")
+        name = str(model.get("name", "")).strip().lower()
+        if not name:
+            raise ValueError("Diffusion model preset requires `model.name`.")
+        return name
+
+    def _defaults_for_model(self, model_name: str) -> tuple[dict[str, Any], dict[str, Any]]:
+        if model_name == "wan":
+            return self.DEFAULT_DATASET, self.DEFAULT_WAN_TRAINER
+        if model_name == "flux" or model_name.startswith("flux."):
+            return self.DEFAULT_FLUX_DATASET, self.DEFAULT_FLUX_TRAINER
+        raise ValueError(f"Unsupported diffusion model name: {model_name!r}")
+
+    def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]:
+        """Translate concise Primus-style sections into trainer arguments.
+
+        Public configs use high-level sections such as `training`, `data`,
+        `parallelism`, and `runtime`; internal defaults supply compact
+        `dataset` and `trainer` objects consumed by the diffusion runtime.
+        """
+        model_name = self._model_family(params)
+        default_dataset, default_trainer = self._defaults_for_model(model_name)
+
+        normalized = {
+            "model": params["model"],
+            "dataset": copy.deepcopy(default_dataset),
+            "trainer": copy.deepcopy(default_trainer),
+            "stage": params.get("stage", "pretrain"),
+            "primus": params.get("primus", {}),
+        }
+
+        dataset_cfg = normalized["dataset"]["config"]
+        trainer_args = normalized["trainer"]["args"]
+
+        training = params.get("training") or {}
+        data = params.get("data") or {}
+        parallelism = params.get("parallelism") or {}
+        optimizer = params.get("optimizer") or {}
+        lr_scheduler = params.get("lr_scheduler") or {}
+        runtime = params.get("runtime") or {}
+        metrics = params.get("metrics") or {}
+
+        training_map = {
+            ("steps",): ("max_steps",),
+            ("local_batch_size",): ("per_device_train_batch_size",),
+            ("global_batch_size",): ("global_batch_size",),
+            ("gradient_accumulation_steps",): ("gradient_accumulation_steps",),
+            ("output_dir",): ("output_dir",),
+            ("save_steps",): ("save_steps",),
+            ("save_strategy",): ("save_strategy",),
+            ("run_name",): ("run_name",),
+            ("num_train_epochs",): ("num_train_epochs",),
+            ("dataloader_num_workers",): ("dataloader_num_workers",),
+            ("resume_from_checkpoint",): ("resume_from_checkpoint",),
+        }
+        for source_path, target_path in training_map.items():
+            value = self._get_any(training, *source_path)
+            if value is not None:
+                self._set_nested(trainer_args, target_path, value)
+
+        data_map = {
+            ("dataset_path",): ("dataset_path",),
+            ("data_folder",): ("data_folder",),
+            ("frame_num",): ("frame_num",),
+            ("video_backend",): ("video_backend",),
+            ("text_tokenizer",): ("processor_config", "text_tokenizer"),
+            ("processor_name",): ("processor_config", "processor_name"),
+            ("processor_type",): ("processor_config", "processor_type"),
+            ("dataset_format",): ("dataset_format",),
+            ("dataset_type",): ("dataset_type",),
+            ("dataset",): ("dataset",),
+            ("shuffle",): ("shuffle",),
+            ("empty_encodings_path",): ("processor_config", "empty_encodings_path"),
+            ("prompt_dropout_prob",): ("processor_config", "prompt_dropout_prob"),
+            ("img_size",): ("processor_config", "img_size"),
+            ("skip_low_resolution",): ("processor_config", "skip_low_resolution"),
+        }
+        for source_path, target_path in data_map.items():
+            value = self._get_any(data, *source_path)
+            if value is not None:
+                self._set_nested(dataset_cfg, target_path, value)
+
+        height = data.get("height")
+        width = data.get("width")
+        if height is not None:
+            self._set_nested(
+                dataset_cfg,
+                ("processor_config", "extra_kwargs", "size", "height"),
+                height,
+            )
+        if width is not None:
+            self._set_nested(
+                dataset_cfg,
+                ("processor_config", "extra_kwargs", "size", "width"),
+                width,
+            )
+
+        parallelism_map = {
+            ("sp_size",): ("sp_size",),
+            ("dp_replicate",): ("dp_replicate",),
+        }
+        for source_path, target_path in parallelism_map.items():
+            value = self._get_any(parallelism, *source_path)
+            if value is not None:
+                self._set_nested(trainer_args, target_path, value)
+
+        optimizer_map = {
+            ("lr",): ("learning_rate",),
+            ("learning_rate",): ("learning_rate",),
+            ("weight_decay",): ("weight_decay",),
+            ("adam_beta1",): ("adam_beta1",),
+            ("adam_beta2",): ("adam_beta2",),
+            ("adam_epsilon",): ("adam_epsilon",),
+            ("max_grad_norm",): ("max_grad_norm",),
+        }
+        for source_path, target_path in optimizer_map.items():
+            value = self._get_any(optimizer, *source_path)
+            if value is not None:
+                self._set_nested(trainer_args, target_path, value)
+
+        lr_scheduler_map = {
+            ("lr_scheduler_type",): ("lr_scheduler_type",),
+            ("warmup_steps",): ("warmup_steps",),
+        }
+        for source_path, target_path in lr_scheduler_map.items():
+            value = self._get_any(lr_scheduler, *source_path)
+            if value is not None:
+                self._set_nested(trainer_args, target_path, value)
+
+        runtime_map = {
+            ("attention_backend",): ("attention_backend",),
+            ("report_to",): ("report_to",),
+            ("seed",): ("seed",),
+            ("bf16",): ("bf16",),
+            ("fp16",): ("fp16",),
+            ("gradient_checkpointing",): ("gradient_checkpointing",),
+            ("compile_transformer_blocks",): ("compile_transformer_blocks",),
+            ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",),
+        }
+        for source_path, target_path in runtime_map.items():
+            value = self._get_any(runtime, *source_path)
+            if value is not None:
+                if target_path in {
+                    ("bf16",),
+                    ("fp16",),
+                    ("gradient_checkpointing",),
+                    ("compile_transformer_blocks",),
+                    ("fsdp2_reshard_after_forward",),
+                }:
+                    value = self._coerce_bool(value)
+                self._set_nested(trainer_args, target_path, value)
+
+        log_freq = metrics.get("log_freq")
+        if log_freq is not None:
+            self._set_nested(trainer_args, ("logging_steps",), log_freq)
+
+        enable_wandb = metrics.get("enable_wandb")
+        if enable_wandb is False and runtime.get("report_to") is None:
+            self._set_nested(trainer_args, ("report_to",), "none")
+        elif enable_wandb is True and runtime.get("report_to") is None:
+            self._set_nested(trainer_args, ("report_to",), "wandb")
+
+        checkpoint = params.get("checkpoint") or {}
+        resume_from_checkpoint = checkpoint.get("resume_from_checkpoint")
+        if resume_from_checkpoint is not None:
+            self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint)
+
+        if (model_name == "flux" or model_name.startswith("flux.")) and int(
+            trainer_args.get("sp_size", 1)
+        ) != 1:
+            raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.")
+
+        return normalized
+
+
+# Backwards-compatible import path for existing Wan-only callers.
+WanArgBuilder = DiffusionArgBuilder
diff --git a/primus/backends/diffusion/attention/__init__.py b/primus/backends/diffusion/attention/__init__.py
new file mode 100644
index 000000000..a1fea3747
--- /dev/null
+++ b/primus/backends/diffusion/attention/__init__.py
@@ -0,0 +1,29 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from .attention import (
+    AITER_FLASH_ATTN_AVAILABLE,
+    FLASH_ATTN_2_AVAILABLE,
+    FLASH_ATTN_3_AVAILABLE,
+    attention,
+    attention_fused,
+    flash_attention,
+    get_attention_backend,
+    set_attention_backend,
+)
+from .flex import FLEX_ATTENTION_AVAILABLE
+
+__all__ = [
+    "AITER_FLASH_ATTN_AVAILABLE",
+    "FLASH_ATTN_2_AVAILABLE",
+    "FLASH_ATTN_3_AVAILABLE",
+    "FLEX_ATTENTION_AVAILABLE",
+    "attention",
+    "attention_fused",
+    "flash_attention",
+    "get_attention_backend",
+    "set_attention_backend",
+]
diff --git a/primus/backends/diffusion/attention/_flash_common.py b/primus/backends/diffusion/attention/_flash_common.py
new file mode 100644
index 000000000..c1bf0f7a4
--- /dev/null
+++ b/primus/backends/diffusion/attention/_flash_common.py
@@ -0,0 +1,130 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+
+def _unwrap_output(x: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+    return x[0] if isinstance(x, tuple) else x
+
+
+def run_flash_attention_backend(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    *,
+    fixed_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]],
+    varlen_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]],
+    window_size_adapter: Callable[[tuple[int, ...]], tuple[int, ...]] | None = None,
+    fixed_extra_kwargs: dict | None = None,
+    varlen_extra_kwargs: dict | None = None,
+):
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+    call_window_size = window_size_adapter(window_size) if window_size_adapter is not None else window_size
+    fixed_extra_kwargs = fixed_extra_kwargs or {}
+    varlen_extra_kwargs = varlen_extra_kwargs or {}
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    if q_lens is None and k_lens is None:
+        qh = half(q)
+        kh = half(k)
+        vh = half(v)
+        qh = qh.to(vh.dtype)
+        kh = kh.to(vh.dtype)
+        if q_scale is not None:
+            qh = qh * q_scale
+
+        x = fixed_attention(
+            qh,
+            kh,
+            vh,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=call_window_size,
+            deterministic=deterministic,
+            **fixed_extra_kwargs,
+        )
+        return _unwrap_output(x).type(out_dtype)
+
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens, strict=False)]))
+
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens, strict=False)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens, strict=False)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+    if q_scale is not None:
+        q = q * q_scale
+
+    cu_seqlens_q = (
+        torch.cat([q_lens.new_zeros([1]), q_lens])
+        .cumsum(0, dtype=torch.int32)
+        .to(q.device, non_blocking=True)
+    )
+    cu_seqlens_k = (
+        torch.cat([k_lens.new_zeros([1]), k_lens])
+        .cumsum(0, dtype=torch.int32)
+        .to(q.device, non_blocking=True)
+    )
+    max_sq = int(q_lens.max())
+    max_sk = int(k_lens.max())
+
+    x = varlen_attention(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_sq,
+        max_seqlen_k=max_sk,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size=call_window_size,
+        deterministic=deterministic,
+        **varlen_extra_kwargs,
+    )
+    # Pad output back to (b, lq) shape to match the input padded layout
+    out = _unwrap_output(x)
+    if max_sq == lq:
+        return out.unflatten(0, (b, lq)).type(out_dtype)
+    # Variable-length: need to scatter back into padded tensor
+    result = q.new_zeros(b, lq, *out.shape[1:])
+    offset = 0
+    for i in range(b):
+        sl = int(q_lens[i])
+        result[i, :sl] = out[offset : offset + sl]
+        offset += sl
+    return result.type(out_dtype)
diff --git a/primus/backends/diffusion/attention/aiter.py b/primus/backends/diffusion/attention/aiter.py
new file mode 100644
index 000000000..dfe3e655a
--- /dev/null
+++ b/primus/backends/diffusion/attention/aiter.py
@@ -0,0 +1,71 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import torch
+
+from ._flash_common import run_flash_attention_backend
+
+AITER_FLASH_ATTN_AVAILABLE = False
+aiter = None
+
+try:
+    import aiter as _aiter  # type: ignore
+
+    if hasattr(_aiter, "flash_attn_func") and hasattr(_aiter, "flash_attn_varlen_func"):
+        aiter = _aiter
+        AITER_FLASH_ATTN_AVAILABLE = True
+except Exception:
+    AITER_FLASH_ATTN_AVAILABLE = False
+
+
+def _normalize_window_size(window_size: tuple[int, ...]) -> tuple[int, int, int]:
+    if len(window_size) == 2:
+        left, right = window_size
+        return (left, right, 0)
+    if len(window_size) == 3:
+        left, right, sink = window_size
+        return (left, right, sink)
+    raise ValueError(f"window_size must have 2 or 3 items, got {window_size}")
+
+
+def aiter_flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+):
+    assert AITER_FLASH_ATTN_AVAILABLE and aiter is not None
+
+    need_lse = torch.is_grad_enabled()
+    return run_flash_attention_backend(
+        q=q,
+        k=k,
+        v=v,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        q_scale=q_scale,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+        dtype=dtype,
+        fixed_attention=aiter.flash_attn_func,
+        varlen_attention=aiter.flash_attn_varlen_func,
+        window_size_adapter=_normalize_window_size,
+        fixed_extra_kwargs={"return_lse": need_lse},
+        varlen_extra_kwargs={"return_lse": need_lse},
+    )
diff --git a/primus/backends/diffusion/attention/attention.py b/primus/backends/diffusion/attention/attention.py
new file mode 100644
index 000000000..cbe8467ef
--- /dev/null
+++ b/primus/backends/diffusion/attention/attention.py
@@ -0,0 +1,410 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unified attention entry points.
+
+Public API:
+- `attention(q, k, v, ...)`        : Wan2.2-style, q/k/v shape [B, L, N, D]
+- `attention_fused(q, k, v, ...)`  : Legacy Wan-style, q/k/v shape [B, S, N*D]
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+
+from ._flash_common import run_flash_attention_backend
+from .aiter import AITER_FLASH_ATTN_AVAILABLE, aiter_flash_attention
+
+FLASH_ATTN_3_AVAILABLE = False
+FLASH_ATTN_2_AVAILABLE = False
+flash_attn_interface = None
+flash_attn = None
+
+try:
+    import flash_attn_interface as _flash_attn_interface  # type: ignore
+
+    flash_attn_interface = _flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn as _flash_attn  # type: ignore
+
+    flash_attn = _flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+__all__ = [
+    "AITER_FLASH_ATTN_AVAILABLE",
+    "FLASH_ATTN_2_AVAILABLE",
+    "FLASH_ATTN_3_AVAILABLE",
+    "attention",
+    "attention_fused",
+    "flash_attention",
+    "get_attention_backend",
+    "set_attention_backend",
+]
+
+
+# ---------------------------------------------------------------------------
+# Global backend config (set once at startup, e.g. from YAML)
+# ---------------------------------------------------------------------------
+_ATTENTION_BACKEND: str = "auto"
+
+
+_VALID_BACKENDS = ("auto", "sdpa", "flex_attention", "flash_attn2", "flash_attn3", "flash_attn_aiter")
+
+
+def set_attention_backend(backend: str) -> None:
+    """
+    Set global attention backend.
+
+    Accepted values (case-insensitive):
+    - "auto" | "sdpa" | "flex_attention" | "flash_attn2" | "flash_attn3" | "flash_attn_aiter"
+    """
+    global _ATTENTION_BACKEND
+    backend = (backend or "").strip().lower()
+    if backend not in _VALID_BACKENDS:
+        hint = ""
+        if backend.startswith("flash_atten"):
+            hint = " (did you mean 'flash_attn2' / 'flash_attn3'?)"
+        raise ValueError(f"attention_backend must be one of {_VALID_BACKENDS}, got '{backend}'{hint}")
+    _ATTENTION_BACKEND = backend
+
+
+def get_attention_backend() -> str:
+    return _ATTENTION_BACKEND
+
+
+def _resolve_flash_version(device_type: str) -> int | None:
+    """
+    Returns:
+        - None: do not use flash attention (use SDPA)
+        - 2 or 3: use flash attention, prefer that version (3 may fall back to 2 in `flash_attention`)
+    """
+    backend = _ATTENTION_BACKEND
+    if backend in ("sdpa", "flex_attention", "flash_attn_aiter"):
+        return None
+
+    if device_type != "cuda":
+        # auto: CPU should use SDPA
+        if backend == "auto":
+            return None
+        # explicit flash backend: error (prevents silent slow CPU path)
+        raise RuntimeError(f"attention_backend='{backend}' requires CUDA tensors.")
+
+    # CUDA path
+    if backend == "auto":
+        if FLASH_ATTN_3_AVAILABLE:
+            return 3
+        if FLASH_ATTN_2_AVAILABLE:
+            return 2
+        return None
+
+    if backend == "flash_attn2":
+        if not FLASH_ATTN_2_AVAILABLE:
+            raise RuntimeError("attention_backend='flash_attn2' but flash-attn 2 is not available.")
+        return 2
+
+    # "flash_attn3"
+    if not FLASH_ATTN_3_AVAILABLE:
+        raise RuntimeError("attention_backend='flash_attn3' but flash-attn 3 is not available.")
+    return 3
+
+
+# ---------------------------------------------------------------------------
+# Flash Attention implementation (varlen + fixed-length fast path)
+# ---------------------------------------------------------------------------
+def _flash_attn_3_fixed(q, k, v, **kwargs):
+    return flash_attn_interface.flash_attn_func(q, k, v)
+
+
+def _flash_attn_3_varlen(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p,
+    softmax_scale,
+    causal,
+    window_size,
+    deterministic,
+):
+    return flash_attn_interface.flash_attn_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=None,
+        seqused_k=None,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        deterministic=deterministic,
+    )
+
+
+def _flash_attn_2_fixed(q, k, v, dropout_p, softmax_scale, causal, window_size, deterministic):
+    try:
+        return flash_attn.flash_attn_func(
+            q,
+            k,
+            v,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+        )
+    except TypeError:
+        return flash_attn.flash_attn_func(q, k, v)
+
+
+def _flash_attn_2_varlen(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p,
+    softmax_scale,
+    causal,
+    window_size,
+    deterministic,
+):
+    return flash_attn.flash_attn_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+    )
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            "Flash attention 3 is not available, use flash attention 2 instead.",
+            stacklevel=2,
+        )
+
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p and window_size are not supported in FA3 now.
+        return run_flash_attention_backend(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            fixed_attention=_flash_attn_3_fixed,
+            varlen_attention=_flash_attn_3_varlen,
+        )
+    assert FLASH_ATTN_2_AVAILABLE
+    return run_flash_attention_backend(
+        q=q,
+        k=k,
+        v=v,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        q_scale=q_scale,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+        dtype=dtype,
+        fixed_attention=_flash_attn_2_fixed,
+        varlen_attention=_flash_attn_2_varlen,
+    )
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    """
+    Unified attention for Wan2.2-style tensors.
+
+    q: [B, Lq, N, D]
+    k: [B, Lk, N, D]
+    v: [B, Lk, N, D]
+    """
+    if _ATTENTION_BACKEND == "flash_attn_aiter":
+        if q.device.type != "cuda":
+            raise RuntimeError("attention_backend='flash_attn_aiter' requires CUDA tensors.")
+        if not AITER_FLASH_ATTN_AVAILABLE:
+            raise RuntimeError("attention_backend='flash_attn_aiter' but aiter is not available.")
+        return aiter_flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+        )
+
+    resolved = _resolve_flash_version(q.device.type)
+    if resolved is not None:
+        version = resolved
+        # Only "auto" allows call-site override.
+        if _ATTENTION_BACKEND == "auto" and fa_version is not None:
+            version = fa_version
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=version,
+        )
+
+    # FlexAttention path (PyTorch, requires torch.compile for perf)
+    if _ATTENTION_BACKEND == "flex_attention":
+        from .flex import flex_attention_fn
+
+        return flex_attention_fn(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+        )
+
+    # SDPA fallback (also used for CPU).
+    # Currently assumes *uniform-length* training sequences, so we do NOT support padding masks on the SDPA path yet
+    if q_lens is not None or k_lens is not None:
+        warnings.warn("Padding mask is disabled when using scaled_dot_product_attention.", stacklevel=2)
+    out_dtype = q.dtype
+    q_ = q.transpose(1, 2).to(dtype)
+    k_ = k.transpose(1, 2).to(dtype)
+    v_ = v.transpose(1, 2).to(dtype)
+    if q_scale is not None:
+        q_ = q_ * q_scale
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q_, k_, v_, attn_mask=None, is_causal=causal, dropout_p=dropout_p
+    )
+    return out.transpose(1, 2).contiguous().to(out_dtype)
+
+
+def attention_fused(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    num_heads: int,
+    dropout_p: float = 0.0,
+    softmax_scale=None,
+    causal: bool = False,
+):
+    """
+    Legacy Wan fused-head attention.
+
+    q/k/v: [B, S, num_heads * head_dim]
+    returns: [B, S, num_heads * head_dim]
+    """
+    b, s, c = q.shape
+    if c % num_heads != 0:
+        raise ValueError(f"hidden dim {c} not divisible by num_heads {num_heads}")
+    d = c // num_heads
+
+    q4d = q.view(b, s, num_heads, d)
+    k4d = k.view(b, -1, num_heads, d)
+    v4d = v.view(b, -1, num_heads, d)
+
+    # For fused calls, q_lens/k_lens are not used (fixed length).
+    out = attention(
+        q=q4d,
+        k=k4d,
+        v=v4d,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16,
+    )
+    return out.flatten(2)
diff --git a/primus/backends/diffusion/attention/flex.py b/primus/backends/diffusion/attention/flex.py
new file mode 100644
index 000000000..c5f0ad091
--- /dev/null
+++ b/primus/backends/diffusion/attention/flex.py
@@ -0,0 +1,280 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+FlexAttention backend — pure PyTorch, relies on torch.compile for performance.
+
+This module provides a drop-in replacement for the SDPA/FlashAttention paths
+using ``torch.nn.attention.flex_attention``.  FlexAttention compiles a
+user-defined ``score_mod`` function into a fused kernel via ``torch.compile``,
+giving FlashAttention-level performance without third-party dependencies.
+
+Key capabilities over SDPA:
+    - Variable-length padding masks via ``create_block_mask``
+    - Sliding-window attention via ``score_mod``
+    - Composable attention score modifications
+"""
+
+from __future__ import annotations
+
+import torch
+
+# Import flex_attention from PyTorch (available since PyTorch 2.5+).
+# We keep the import at module level so that ``FLEX_ATTENTION_AVAILABLE``
+# reflects the actual environment at init time.
+FLEX_ATTENTION_AVAILABLE = False
+_flex_attention_compiled = None
+_create_block_mask = None
+
+try:
+    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+    # torch.compile is essential for FlexAttention performance — without it,
+    # score_mod/block_mask run as eager Python callbacks (separate kernel per op).
+    # We attempt compile lazily on first call; if Triton/Inductor fails (e.g.
+    # missing system libs on ROCm), we fall back to eager with a warning.
+    _flex_attention_eager = flex_attention
+    _flex_attention_compiled = None  # lazy-init on first call
+    _create_block_mask = create_block_mask
+    FLEX_ATTENTION_AVAILABLE = True
+except Exception:
+    FLEX_ATTENTION_AVAILABLE = False
+
+
+def _get_flex_attention():
+    """
+    Return compiled flex_attention, falling back to eager if compile fails.
+
+    torch.compile is lazy: wrapping succeeds immediately, but actual Triton
+    codegen happens on first invocation and can fail (e.g. missing system libs
+    on ROCm).  We handle this by catching runtime errors on first call and
+    permanently switching to eager mode.
+    """
+    global _flex_attention_compiled
+    if _flex_attention_compiled is not None:
+        return _flex_attention_compiled
+
+    import logging
+
+    logger = logging.getLogger(__name__)
+    try:
+        compiled = torch.compile(_flex_attention_eager)
+        logger.info("flex_attention: torch.compile wrapper created (lazy — actual compile on first call)")
+        _flex_attention_compiled = compiled
+    except Exception as e:
+        logger.warning("flex_attention: torch.compile wrapping failed (%s), using eager mode", e)
+        _flex_attention_compiled = _flex_attention_eager
+    return _flex_attention_compiled
+
+
+def _flex_attention_with_fallback(q, k, v, **kwargs):
+    """Call flex_attention with runtime fallback if compiled version fails."""
+    global _flex_attention_compiled
+    fn = _get_flex_attention()
+    try:
+        return fn(q, k, v, **kwargs)
+    except Exception as e:
+        if fn is _flex_attention_eager:
+            raise  # already in eager mode, real error
+        import logging
+
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "flex_attention: compiled call failed (%s: %s), falling back to eager mode",
+            type(e).__name__,
+            e,
+        )
+        _flex_attention_compiled = _flex_attention_eager
+        return _flex_attention_eager(q, k, v, **kwargs)
+
+
+def _build_score_mod(
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    softmax_scale: float | None = None,
+):
+    """
+    Compose a ``score_mod`` function for ``flex_attention``.
+
+    The returned function has signature ``(score, b, h, q_idx, kv_idx) -> score``
+    and is compiled by ``torch.compile`` into fused operations.
+    """
+    mods: list = []
+
+    if softmax_scale is not None:
+        # flex_attention already applies 1/sqrt(d) by default, so we only
+        # inject a custom scale if the caller wants a non-default value.
+        def scale_mod(score, b, h, q_idx, kv_idx):
+            return score * softmax_scale
+
+        mods.append(scale_mod)
+
+    if causal:
+
+        def causal_mod(score, b, h, q_idx, kv_idx):
+            return torch.where(q_idx >= kv_idx, score, float("-inf"))
+
+        mods.append(causal_mod)
+
+    if window_size != (-1, -1):
+        left, right = window_size
+
+        def window_mod(score, b, h, q_idx, kv_idx):
+            in_window = True
+            if left >= 0:
+                in_window = in_window & (q_idx - kv_idx <= left)
+            if right >= 0:
+                in_window = in_window & (kv_idx - q_idx <= right)
+            return torch.where(in_window, score, float("-inf"))
+
+        mods.append(window_mod)
+
+    if not mods:
+        return None
+
+    # Compose all mods into one function.
+    def composed_score_mod(score, b, h, q_idx, kv_idx):
+        for mod in mods:
+            score = mod(score, b, h, q_idx, kv_idx)
+        return score
+
+    return composed_score_mod
+
+
+def _build_block_mask(
+    B: int,
+    N: int,
+    Lq: int,
+    Lk: int,
+    q_lens: torch.Tensor | None,
+    k_lens: torch.Tensor | None,
+    device: torch.device,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+):
+    """
+    Build a ``BlockMask`` for ``flex_attention`` from variable-length seqs.
+
+    If both ``q_lens`` and ``k_lens`` are None (uniform-length), returns None
+    so that ``flex_attention`` uses a full dense mask.
+    """
+    if q_lens is None and k_lens is None and not causal and window_size == (-1, -1):
+        return None
+
+    # Ensure length tensors are on the same device as the mask indices
+    # (create_block_mask uses vmap which creates index tensors on `device`)
+    if q_lens is not None:
+        q_lens = q_lens.to(device)
+    if k_lens is not None:
+        k_lens = k_lens.to(device)
+
+    # Build the mask function. Closures capture the lengths.
+    def mask_fn(b, h, q_idx, kv_idx):
+        mask = True
+        if q_lens is not None:
+            mask = mask & (q_idx < q_lens[b])
+        if k_lens is not None:
+            mask = mask & (kv_idx < k_lens[b])
+        if causal:
+            mask = mask & (q_idx >= kv_idx)
+        if window_size != (-1, -1):
+            left, right = window_size
+            if left >= 0:
+                mask = mask & (q_idx - kv_idx <= left)
+            if right >= 0:
+                mask = mask & (kv_idx - q_idx <= right)
+        return mask
+
+    return _create_block_mask(mask_fn, B=B, H=N, Q_LEN=Lq, KV_LEN=Lk, device=device)
+
+
+def flex_attention_fn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: torch.Tensor | None = None,
+    k_lens: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    q_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    deterministic: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    fa_version=None,
+) -> torch.Tensor:
+    """
+    Attention using ``torch.nn.attention.flex_attention``.
+
+    Drop-in replacement for ``attention()`` — same signature, same tensor
+    layout (q/k/v: ``[B, L, N, D]``).
+
+    Args:
+        q: Query tensor  ``[B, Lq, N, D]``
+        k: Key tensor    ``[B, Lk, N, D]``
+        v: Value tensor  ``[B, Lk, N, D]``
+        q_lens: Optional valid lengths per batch for queries ``[B]``
+        k_lens: Optional valid lengths per batch for keys   ``[B]``
+        dropout_p: Dropout probability (passed through)
+        softmax_scale: Custom softmax scale (None = default 1/sqrt(D))
+        q_scale: Pre-multiply q by this scalar
+        causal: Whether to apply causal mask
+        window_size: (left, right) sliding window; (-1, -1) = disabled
+        deterministic: Ignored (for API compat with flash_attention)
+        dtype: Target dtype for computation
+        fa_version: Ignored (for API compat)
+    """
+    if not FLEX_ATTENTION_AVAILABLE:
+        raise RuntimeError(
+            "flex_attention backend requested but torch.nn.attention.flex_attention "
+            "is not available. Requires PyTorch >= 2.5."
+        )
+
+    out_dtype = q.dtype
+    B, Lq, N, D = q.shape
+    Lk = k.shape[1]
+
+    # Cast to compute dtype
+    half_dtypes = (torch.float16, torch.bfloat16)
+    q_ = q.to(dtype) if q.dtype not in half_dtypes else q
+    k_ = k.to(dtype) if k.dtype not in half_dtypes else k
+    v_ = v.to(dtype) if v.dtype not in half_dtypes else v
+    q_ = q_.to(v_.dtype)
+    k_ = k_.to(v_.dtype)
+
+    if q_scale is not None:
+        q_ = q_ * q_scale
+
+    # flex_attention expects [B, N, L, D] (heads-first)
+    q_ = q_.transpose(1, 2)  # [B, N, Lq, D]
+    k_ = k_.transpose(1, 2)  # [B, N, Lk, D]
+    v_ = v_.transpose(1, 2)  # [B, N, Lk, D]
+
+    # Build block mask (handles padding, causal, window in the mask itself)
+    block_mask = _build_block_mask(
+        B=B,
+        N=N,
+        Lq=Lq,
+        Lk=Lk,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        device=q.device,
+        causal=causal,
+        window_size=window_size,
+    )
+
+    # Build score_mod only for softmax_scale override
+    # (causal and window are handled in block_mask for better sparsity)
+    score_mod = None
+    if softmax_scale is not None:
+        score_mod = _build_score_mod(softmax_scale=softmax_scale)
+
+    # Call flex_attention (with runtime compile-error fallback)
+    out = _flex_attention_with_fallback(q_, k_, v_, score_mod=score_mod, block_mask=block_mask)
+
+    # Back to [B, L, N, D]
+    out = out.transpose(1, 2).contiguous()
+    return out.to(out_dtype)
diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py
new file mode 100644
index 000000000..b72452684
--- /dev/null
+++ b/primus/backends/diffusion/data/__init__.py
@@ -0,0 +1,27 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Data modules for the Primus diffusion backend."""
+
+from .config import DatasetConfig
+from .dataset import WanVideoDataset
+from .flux_precomputed import (
+    FluxPrecomputedDataset,
+    FluxPrecomputedProcessor,
+    FluxRawImageTextDataset,
+    FluxRawImageTextProcessor,
+)
+from .processor import WanVideoDataProcessor
+
+__all__ = [
+    "DatasetConfig",
+    "FluxPrecomputedDataset",
+    "FluxPrecomputedProcessor",
+    "FluxRawImageTextDataset",
+    "FluxRawImageTextProcessor",
+    "WanVideoDataProcessor",
+    "WanVideoDataset",
+]
diff --git a/primus/backends/diffusion/data/collator.py b/primus/backends/diffusion/data/collator.py
new file mode 100644
index 000000000..fb6d5fba1
--- /dev/null
+++ b/primus/backends/diffusion/data/collator.py
@@ -0,0 +1,98 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import collections
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+
+@dataclass
+class VisionCollator:
+    processor: Any
+
+    def pad_sequence(self, input_ids, batch_first, padding_value):
+        if self.processor.tokenizer.padding_side == "left":
+            input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids]
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=batch_first, padding_value=padding_value
+        )
+        if self.processor.tokenizer.padding_side == "left":
+            input_ids = torch.flip(input_ids, [1])
+        return input_ids
+
+    def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]:
+        if isinstance(instances[0], list):
+            instances = [inst for instance in instances for inst in instance]
+        inputs = collections.defaultdict(list)
+        for instance in instances:
+            for key, values in instance.items():
+                inputs[key].append(values)
+
+        batched_inputs = {}
+        input_ids = None
+        if "input_ids" in inputs:
+            input_ids = inputs.pop("input_ids")
+            input_ids = self.pad_sequence(
+                input_ids,
+                batch_first=True,
+                padding_value=self.processor.tokenizer.pad_token_id,
+            )
+            batched_inputs["input_ids"] = input_ids
+        if "labels" in inputs:
+            labels = inputs.pop("labels")
+            labels = self.pad_sequence(
+                labels,
+                batch_first=True,
+                padding_value=-100,
+            )
+            batched_inputs["labels"] = labels
+
+        attention_mask = None
+        if "attention_mask" in inputs:
+            attention_mask = inputs.pop("attention_mask")
+
+        if input_ids is not None:
+            batched_inputs["attention_mask"] = input_ids.ne(self.processor.tokenizer.pad_token_id).long()
+        elif attention_mask is not None:
+            batched_inputs["attention_mask"] = self.pad_sequence(
+                attention_mask,
+                batch_first=True,
+                padding_value=0,
+            )
+
+        # for the other keys
+        for key, values in inputs.items():
+            # Handle scalar/boolean values ( use_audio_in_video)
+            if isinstance(values[0], bool) or (
+                isinstance(values[0], (int, float)) and not isinstance(values[0], torch.Tensor)
+            ):
+                batched_inputs[key] = values[0]
+            else:
+                batched_inputs[key] = torch.stack(values, dim=0)
+        return batched_inputs
+
+    @property
+    def image_token_id(self):
+        return self.processor.tokenizer.convert_tokens_to_ids(self.processor.image_token)
+
+
+@dataclass
+class RawBatchCollator:
+    """
+    A minimal collator that returns raw samples as a list of dicts.
+
+    This is useful when model-specific padding/encoding happens in a separate
+    batch preparation step (e.g. processor.prepare_batch / trainer.prepare_batch),
+    keeping the DataLoader and trainer model-agnostic.
+    """
+
+    def __call__(self, instances: Sequence[dict]) -> list[dict]:
+        if isinstance(instances[0], list):
+            instances = [inst for instance in instances for inst in instance]
+        return list(instances)
diff --git a/primus/backends/diffusion/data/config.py b/primus/backends/diffusion/data/config.py
new file mode 100644
index 000000000..e425149ba
--- /dev/null
+++ b/primus/backends/diffusion/data/config.py
@@ -0,0 +1,89 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, field_validator
+
+
+class Args(BaseModel):
+    extra_kwargs: dict[str, Any] = {}
+
+    def to_dict(self):
+        return self.model_dump()
+
+    def to_json(self):
+        return self.model_dump_json()
+
+
+class ProcessorConfig(Args):
+    processor_name: str
+    processor_type: str
+
+
+class DatasetConfig(Args):
+    dataset_type: str
+    data_folder: str
+    dataset_format: Literal["json", "jsonl", "csv", "yaml", "hf_dataset", "arrow"]
+    processor_config: dict | ProcessorConfig
+
+    # Dataset configuration
+    dataset_path: str | None = None  # Optional - used for external files
+    datasets: list[dict] | None = None  # Optional - used for inline YAML definitions
+    shuffle: bool = True
+    data_seed: int | None = 42
+    eval_dataset_path: str | None = None
+
+    # Object storage configuration
+    object_storage: Literal["azure", "gcs", "none"] | None = "none"
+    bucket_name: str | None = None
+
+    # Packing configuration
+    packing: bool | None = False
+    packing_strategy: str | None = None
+    packing_length: int | None = 32000
+    filter_overlong: bool | None = True
+    filter_overlong_workers: int | None = 8
+    max_length: int | None = None
+
+    # Video configuration
+    video_sampling_strategy: Literal["fps", "frame_num"] | None = "fps"
+    video_max_pixels: int | None = 768 * 28 * 28
+    video_max_frames: int | None = 768
+    video_min_pixels: int | None = 3136
+    frame_num: int | None = 64
+    fps: int | None = 1
+    video_backend: Literal["decord", "qwen_vl_utils", "qwen_omni_utils", "imageio"] | None = "qwen_vl_utils"
+
+    @field_validator(
+        "video_max_pixels",
+        "video_max_frames",
+        "frame_num",
+        "fps",
+        "packing_length",
+        "max_length",
+        "filter_overlong_workers",
+    )
+    @classmethod
+    def validate_positive_values(cls, v, info):
+        """Validate that numeric video and packing parameters are positive."""
+        if v is not None and v <= 0:
+            field_name = info.field_name
+            raise ValueError(f"{field_name} must be positive, got {v}")
+        return v
+
+    @field_validator("video_backend")
+    @classmethod
+    def validate_video_backend_migration(cls, v):
+        """Provide migration warning for deprecated torchvision backend."""
+        if v == "torchvision":
+            raise ValueError(
+                "The 'torchvision' video backend has been removed. "
+                "Please use 'decord', 'qwen_vl_utils', or 'qwen_omni_utils' instead. "
+                "Migration guide: If you were using torchvision, 'decord' provides "
+                "similar functionality with better performance."
+            )
+        return v
diff --git a/primus/backends/diffusion/data/dataset.py b/primus/backends/diffusion/data/dataset.py
new file mode 100644
index 000000000..c5584dbdc
--- /dev/null
+++ b/primus/backends/diffusion/data/dataset.py
@@ -0,0 +1,291 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Dataset class for loading video data from JSONL/CSV files.
+"""
+
+import json
+import os
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from primus.backends.diffusion.utils.data_utils import smart_nframes
+from primus.backends.diffusion.utils.vision_process import strip_file_uri
+
+from .collator import RawBatchCollator
+
+
+def _get_decord_vr(
+    video_path: str,
+    *,
+    num_threads: int,
+):
+    """Construct a decord.VideoReader."""
+    from decord import VideoReader, cpu
+
+    return VideoReader(video_path, ctx=cpu(0), num_threads=num_threads)
+
+
+class BaseDataset(Dataset):
+    """Base dataset with minimal shared interface."""
+
+    def __init__(self, config, **kwargs) -> None:
+        super().__init__()
+        self.config = config
+        self.samples = []
+
+    def __len__(self):
+        return len(self.samples)
+
+
+class WanVideoDataset(BaseDataset):
+    """Dataset for WanVideo training from JSONL/CSV."""
+
+    def __init__(self, processor, config=None):
+        """
+        Initialize WanVideo dataset.
+
+        Args:
+            processor: WanVideoDataProcessor instance
+            video_backend: Backend for video loading ('qwen_vl_utils' or 'decord')
+        """
+        if config is None:
+            raise ValueError("WanVideoDataset requires a dataset config with dataset_path")
+        if getattr(config, "dataset_path", None) is None:
+            raise ValueError("WanVideoDataset requires config.dataset_path")
+
+        super().__init__(config)
+        self.config = config
+        self.data_path = Path(self.config.dataset_path)
+        self.processor = processor
+
+        # Load metadata
+        self.samples = self._load_metadata()
+        self._sync_processor_video_limits()
+
+    def _sync_processor_video_limits(self):
+        image_processor = None
+        if hasattr(self.processor, "processor") and hasattr(self.processor.processor, "image_processor"):
+            image_processor = self.processor.processor.image_processor
+        if image_processor is None:
+            return
+        if (
+            getattr(image_processor, "max_pixels", None) is None
+            and getattr(self.config, "video_max_pixels", None) is not None
+        ):
+            image_processor.max_pixels = self.config.video_max_pixels
+
+    def _load_metadata(self) -> list[dict]:
+        """Load metadata from JSONL or CSV file."""
+        samples = []
+
+        if self.data_path.suffix == ".jsonl":
+            with open(self.data_path) as f:
+                for line in f:
+                    samples.append(json.loads(line.strip()))
+        elif self.data_path.suffix == ".json":
+            with open(self.data_path) as f:
+                samples = json.load(f)
+        elif self.data_path.suffix == ".csv":
+            import pandas as pd
+
+            df = pd.read_csv(self.data_path)
+            samples = df.to_dict("records")
+        else:
+            raise ValueError(f"Unsupported file format: {self.data_path=}")
+
+        return samples
+
+    def _load_video_frames(self, video_path: str, data_folder=None, fps: int = 1) -> tuple[np.ndarray, float]:
+        """Load video frames using the specified backend."""
+        if self.config.data_folder is not None:
+            video_path = os.path.join(self.config.data_folder, video_path)
+
+        if self.config.video_backend == "decord":
+            return self.load_video_decord(video_path, fps)
+        elif self.config.video_backend == "qwen_vl_utils":
+            return self.load_video_qwen_vl_utils(video_path, fps)
+        elif self.config.video_backend == "imageio":
+            return self.load_video_imageio(video_path, fps)
+        else:
+            raise ValueError(f"Unsupported video backend: {self.config.video_backend}")
+
+    def load_video_imageio(self, video_path, fps):
+        import imageio
+
+        video_path = strip_file_uri(video_path)
+        reader = imageio.get_reader(video_path)
+
+        # Sampling Strategy
+        total_frames = reader.count_frames()
+        total_frames = int(total_frames)
+
+        if self.config.video_sampling_strategy == "frame_num":
+            nframes = self.config.frame_num
+            # Enforce VAE divisibility: (n - 1) % 4 == 0
+            actual_nframes = min(nframes, total_frames)
+
+            valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1
+
+            # DiffSynth Sequential Reading
+            frames = []
+            for i, frame in enumerate(reader):
+                if i >= valid_nframes:
+                    break
+                frames.append(frame)
+
+            # Stack to numpy (T, H, W, C)
+            frames = np.array(frames)
+            sample_fps = fps  # Simplification
+
+            reader.close()
+        else:
+            reader.close()
+            raise NotImplementedError("Only frame_num strategy implemented for imageio backend")
+
+        return frames, sample_fps
+
+    def load_video_decord(
+        self,
+        video_path: str | list[str] | BytesIO,
+        fps: int,
+    ) -> tuple[np.ndarray, float]:
+        """
+        Load video using Decord backend.
+
+        Args:
+            video_path: Path to video file or BytesIO object
+            fps: Target frames per second
+
+        Returns:
+            Tuple of (video frames, sample fps)
+        """
+        # Keep dataset logic simple: use a fixed thread count here.
+        from decord import VideoReader, cpu
+
+        num_threads = 4
+
+        if isinstance(video_path, BytesIO):
+            vr = VideoReader(video_path, ctx=cpu(0), num_threads=num_threads)
+        elif isinstance(video_path, list):
+            vr = _get_decord_vr(strip_file_uri(video_path[0]), num_threads=num_threads)
+        elif isinstance(video_path, str):
+            vr = _get_decord_vr(strip_file_uri(video_path), num_threads=num_threads)
+        else:
+            raise ValueError(f"Unsupported video path type: {type(video_path)}")
+
+        total_frames, video_fps = len(vr), vr.get_avg_fps()
+        if self.config.video_sampling_strategy == "fps":
+            nframes = smart_nframes(total_frames, video_fps=video_fps, fps=fps)
+            # Maintain uniform sampling for FPS strategy
+            uniform_sampled_frames = np.linspace(0, total_frames - 1, nframes, dtype=int)
+        elif self.config.video_sampling_strategy == "frame_num":
+            nframes = self.config.frame_num
+            # Enforce VAE divisibility: (n - 1) % 4 == 0
+            actual_nframes = min(nframes, total_frames)
+
+            valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1
+            uniform_sampled_frames = np.linspace(0, total_frames - 1, valid_nframes, dtype=int)
+        else:
+            raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}")
+
+        frame_idx = uniform_sampled_frames.tolist()
+        spare_frames = vr.get_batch(frame_idx).asnumpy()
+        # spare_frames = torch.tensor(spare_frames).permute(0, 3, 1, 2)  # Convert to TCHW format
+
+        # Calculate sample_fps
+        sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+
+        # Return HWC numpy array to match processor expectations
+        return spare_frames, sample_fps  # (frames, height, width, channels)
+
+    def load_video_qwen_vl_utils(
+        self,
+        video_path: str,
+        fps: int,
+    ) -> tuple[np.ndarray, float]:
+        """
+        Load video using Qwen VL utils.
+
+        Args:
+            video_path: Path to video file
+            fps: Target frames per second
+
+        Returns:
+            Tuple of (video frames, sample fps)
+        """
+        from primus.backends.diffusion.utils.vision_process import fetch_video
+
+        video_dict = {
+            "type": "video",
+            "video": f"file://{video_path}",
+            "min_frames": 1,
+            "max_pixels": self.config.video_max_pixels,
+            "max_frames": self.config.video_max_frames,
+            "min_pixels": self.config.video_min_pixels,
+        }
+
+        if self.config.video_sampling_strategy == "frame_num":
+            is_even = self.config.frame_num % 2 == 0
+            n_frames = self.config.frame_num if is_even else self.config.frame_num + 1
+            video_dict["nframes"] = n_frames
+
+            frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True)
+            frames = frames.numpy()
+
+            # Enforce VAE divisibility constraint
+            actual_n = len(frames)
+            if actual_n > 1:
+                valid_n = ((actual_n - 1) // 4) * 4 + 1
+                frames = frames[:valid_n]
+            # else: keep 1 frame (or handle error)
+
+            return frames, sample_fps
+        elif self.config.video_sampling_strategy == "fps":
+            video_dict["fps"] = fps
+            frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True)
+            frames = frames.numpy()
+
+            # Also enforce for fps strategy
+            actual_n = len(frames)
+            if actual_n > 1:
+                valid_n = ((actual_n - 1) // 4) * 4 + 1
+                frames = frames[:valid_n]
+
+            return frames, sample_fps
+        else:
+            raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        """Get a single sample."""
+        sample = self.samples[idx]
+
+        # Load video frames
+        video_path = sample["video"]
+        video_frames, fps = self._load_video_frames(video_path)
+
+        # Get prompt
+        prompt = sample.get("prompt", "")
+        # Return raw sample.
+        return {
+            "video_frames": video_frames,  # np.ndarray, typically (T, H, W, C)
+            "prompt": prompt,
+            "fps": fps,
+            "num_frames": int(getattr(self.config, "frame_num", 0) or 0),
+            "video_path": str(video_path),
+        }
+
+    def get_collator(self):
+        # Prefer raw collation; model-specific processing should happen in processor.prepare_batch.
+        return RawBatchCollator()
diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py
new file mode 100644
index 000000000..f2ea9ab0d
--- /dev/null
+++ b/primus/backends/diffusion/data/flux_precomputed.py
@@ -0,0 +1,340 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import io
+import json
+import math
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from PIL import Image, ImageFile
+from torch.utils.data import Dataset
+
+from primus.backends.diffusion.data.collator import RawBatchCollator
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def _tensor_from_serialized_numpy(value: bytes) -> torch.Tensor:
+    array = np.load(io.BytesIO(value))
+    tensor = torch.from_numpy(array)
+    if tensor.dtype == torch.uint16:
+        tensor = tensor.view(torch.bfloat16)
+    return tensor
+
+
+def _to_tensor(value: Any) -> torch.Tensor:
+    if isinstance(value, torch.Tensor):
+        return value
+    if isinstance(value, (bytes, bytearray)):
+        return _tensor_from_serialized_numpy(bytes(value))
+    if isinstance(value, np.ndarray):
+        tensor = torch.from_numpy(value)
+        if tensor.dtype == torch.uint16:
+            tensor = tensor.view(torch.bfloat16)
+        return tensor
+    return torch.as_tensor(value)
+
+
+class FluxPrecomputedDataset(Dataset):
+    """Map-style dataset for precomputed FLUX text and VAE encodings."""
+
+    required_fields = ("t5_encodings", "clip_encodings", "mean", "logvar")
+
+    def __init__(self, dataset_path: str):
+        if not dataset_path:
+            raise ValueError("FLUX precomputed dataset requires `dataset_path`.")
+        if not os.path.isdir(dataset_path):
+            raise FileNotFoundError(f"FLUX precomputed dataset directory not found: {dataset_path}")
+        from datasets import load_from_disk
+
+        self.dataset_path = dataset_path
+        self.dataset = load_from_disk(dataset_path)
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        sample = self.dataset[idx]
+        missing = [field for field in self.required_fields if field not in sample]
+        if missing:
+            raise KeyError(f"FLUX precomputed sample missing fields: {missing}")
+        return {field: _to_tensor(sample[field]) for field in self.required_fields}
+
+    def get_collator(self):
+        return RawBatchCollator()
+
+
+class FluxRawImageTextDataset(Dataset):
+    """Map-style raw image-text dataset for online FLUX encoding."""
+
+    def __init__(
+        self,
+        *,
+        dataset_path: str | None,
+        dataset_format: str = "webdataset",
+        dataset_name: str | None = None,
+        data_folder: str | None = None,
+    ):
+        dataset_name = dataset_name or None
+        dataset_path, dataset_format = self._resolve_dataset(dataset_name, dataset_path, dataset_format)
+        self.dataset_path = dataset_path
+        self.dataset_format = dataset_format
+        self.dataset_name = dataset_name
+        self.data_folder = data_folder
+        self._records: list[dict[str, Any]] | None = None
+
+        if dataset_format == "jsonl":
+            path = Path(dataset_path)
+            if not path.is_file():
+                raise FileNotFoundError(f"FLUX raw jsonl metadata not found: {dataset_path}")
+            self._records = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
+            self.dataset = None
+        elif dataset_format == "hf_dataset":
+            if not os.path.isdir(dataset_path):
+                raise FileNotFoundError(f"FLUX raw HF dataset directory not found: {dataset_path}")
+            from datasets import load_from_disk
+
+            self.dataset = load_from_disk(dataset_path)
+        elif dataset_format == "hf_repo":
+            from datasets import load_dataset
+
+            self.dataset = load_dataset(dataset_path, split="train")
+        elif dataset_format == "webdataset":
+            if not os.path.isdir(dataset_path):
+                raise FileNotFoundError(f"FLUX raw webdataset directory not found: {dataset_path}")
+            from datasets import load_dataset
+
+            self.dataset = load_dataset(
+                "webdataset",
+                split="train",
+                data_dir=dataset_path,
+                data_files={"train": "*.tar"},
+            )
+        else:
+            raise ValueError("FLUX raw dataset_format must be one of: jsonl, hf_dataset, hf_repo, webdataset")
+
+    @staticmethod
+    def _resolve_dataset(
+        dataset_name: str | None,
+        dataset_path: str | None,
+        dataset_format: str,
+    ) -> tuple[str, str]:
+        if dataset_name == "cc12m-test":
+            return (dataset_path, dataset_format) if dataset_path else ("zirui3/cc12m-test", "hf_repo")
+        if dataset_name == "cc12m-wds":
+            return dataset_path or "pixparse/cc12m-wds", "hf_repo"
+        if not dataset_path:
+            raise ValueError("FLUX raw dataset requires either `dataset` or `dataset_path`.")
+        return dataset_path, dataset_format
+
+    def __len__(self) -> int:
+        if self._records is not None:
+            return len(self._records)
+        return len(self.dataset)
+
+    @staticmethod
+    def _prompt_from_sample(sample: dict[str, Any]) -> str:
+        for key in ("txt", "caption", "prompt", "text"):
+            if key in sample:
+                value = sample[key]
+                if isinstance(value, list):
+                    value = value[0]
+                return str(value)
+        raise KeyError("FLUX raw sample requires one of: txt, caption, prompt, text")
+
+    def _image_from_jsonl_sample(self, sample: dict[str, Any]) -> Image.Image:
+        image_key = "image" if "image" in sample else "jpg" if "jpg" in sample else "png"
+        image_path = Path(str(sample[image_key]))
+        if self.data_folder and not image_path.is_absolute():
+            image_path = Path(self.data_folder) / image_path
+        return Image.open(image_path).convert("RGB")
+
+    @staticmethod
+    def _image_from_dataset_sample(sample: dict[str, Any]) -> Image.Image:
+        if "image" in sample:
+            image = sample["image"]
+        elif "jpg" in sample:
+            image = sample["jpg"]
+        elif "png" in sample:
+            image = sample["png"]
+        else:
+            raise KeyError("FLUX raw dataset sample requires image, jpg, or png")
+        if isinstance(image, Image.Image):
+            return image.convert("RGB")
+        if isinstance(image, (bytes, bytearray)):
+            return Image.open(io.BytesIO(image)).convert("RGB")
+        return image.convert("RGB")
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        if self._records is not None:
+            sample = self._records[idx]
+            image = self._image_from_jsonl_sample(sample)
+        else:
+            sample = self.dataset[idx]
+            image = self._image_from_dataset_sample(sample)
+        return {"image": image, "prompt": self._prompt_from_sample(sample)}
+
+    def get_collator(self):
+        return RawBatchCollator()
+
+
+@dataclass
+class FluxPrecomputedProcessor:
+    config: dict[str, Any]
+
+    def __post_init__(self):
+        self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0)
+        self.empty_encodings_path = self.config.get("empty_encodings_path")
+        self.img_size = int(self.config.get("img_size", 256) or 256)
+        self._empty_t5: torch.Tensor | None = None
+        self._empty_clip: torch.Tensor | None = None
+
+    def build(self):
+        if self.prompt_dropout_prob > 0.0:
+            self._load_empty_encodings()
+
+    @staticmethod
+    def _normalize_empty_encoding(tensor: torch.Tensor) -> torch.Tensor:
+        if tensor.ndim >= 1 and tensor.shape[0] == 1:
+            return tensor[0]
+        return tensor
+
+    def _check_empty_encoding_shapes(self, t5_encodings: torch.Tensor, clip_encodings: torch.Tensor) -> None:
+        # The empty encodings are broadcast-assigned into the per-sample encodings
+        # for dropped prompts, so their trailing (non-batch) shape must match the
+        # dataset encodings exactly. Fail fast with a clear message instead of a
+        # cryptic in-place assignment shape error mid-training.
+        assert self._empty_t5 is not None and self._empty_clip is not None
+        if self._empty_t5.shape != t5_encodings.shape[1:]:
+            raise ValueError(
+                "FLUX prompt dropout: empty T5 encoding shape "
+                f"{tuple(self._empty_t5.shape)} does not match the per-sample T5 encoding shape "
+                f"{tuple(t5_encodings.shape[1:])}. Regenerate t5_empty.npy with a matching "
+                "sequence length / hidden size."
+            )
+        if self._empty_clip.shape != clip_encodings.shape[1:]:
+            raise ValueError(
+                "FLUX prompt dropout: empty CLIP encoding shape "
+                f"{tuple(self._empty_clip.shape)} does not match the per-sample CLIP encoding shape "
+                f"{tuple(clip_encodings.shape[1:])}. Regenerate clip_empty.npy with a matching shape."
+            )
+
+    def _load_empty_encodings(self):
+        if self._empty_t5 is not None and self._empty_clip is not None:
+            return
+        if not self.empty_encodings_path:
+            raise ValueError("FLUX prompt dropout requires `data.empty_encodings_path`.")
+        t5_path = os.path.join(self.empty_encodings_path, "t5_empty.npy")
+        clip_path = os.path.join(self.empty_encodings_path, "clip_empty.npy")
+        if not os.path.isfile(t5_path) or not os.path.isfile(clip_path):
+            raise FileNotFoundError(
+                "FLUX empty encodings must contain t5_empty.npy and clip_empty.npy "
+                f"under {self.empty_encodings_path}"
+            )
+        self._empty_t5 = self._normalize_empty_encoding(torch.from_numpy(np.load(t5_path)))
+        self._empty_clip = self._normalize_empty_encoding(torch.from_numpy(np.load(clip_path)))
+
+    @staticmethod
+    def _collate_raw(batch: Any) -> dict[str, torch.Tensor]:
+        if isinstance(batch, dict):
+            return {key: _to_tensor(value) for key, value in batch.items()}
+        if not isinstance(batch, Sequence) or not batch:
+            raise ValueError(f"FLUX prepare_batch expected non-empty sequence, got {type(batch).__name__}")
+
+        keys = ("t5_encodings", "clip_encodings", "mean", "logvar")
+        collated: dict[str, torch.Tensor] = {}
+        for key in keys:
+            collated[key] = torch.stack([_to_tensor(sample[key]) for sample in batch], dim=0)
+        return collated
+
+    def prepare_batch(
+        self, *, batch: Any, device: torch.device, dtype: torch.dtype
+    ) -> dict[str, torch.Tensor]:
+        tensors = self._collate_raw(batch)
+        for key, value in tensors.items():
+            tensors[key] = value.to(device=device, dtype=dtype, non_blocking=True)
+
+        if self.prompt_dropout_prob > 0.0:
+            self._load_empty_encodings()
+            assert self._empty_t5 is not None and self._empty_clip is not None
+            self._check_empty_encoding_shapes(tensors["t5_encodings"], tensors["clip_encodings"])
+            bsz = tensors["t5_encodings"].shape[0]
+            drop_mask = torch.rand((bsz,), device=device) < self.prompt_dropout_prob
+            if drop_mask.any():
+                tensors["t5_encodings"][drop_mask] = self._empty_t5.to(device=device, dtype=dtype)
+                tensors["clip_encodings"][drop_mask] = self._empty_clip.to(device=device, dtype=dtype)
+
+        return tensors
+
+
+@dataclass
+class FluxRawImageTextProcessor:
+    config: dict[str, Any]
+
+    def __post_init__(self):
+        self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0)
+        self.img_size = int(self.config.get("img_size", 256) or 256)
+        self.skip_low_resolution = bool(self.config.get("skip_low_resolution", True))
+
+    def build(self):
+        return
+
+    def _process_image(self, image: Image.Image) -> torch.Tensor | None:
+        width, height = image.size
+        if self.skip_low_resolution and (width < self.img_size or height < self.img_size):
+            return None
+
+        if width == self.img_size and height == self.img_size:
+            resized = image
+        elif width >= height:
+            new_width, new_height = math.ceil(self.img_size / height * width), self.img_size
+            image = image.resize((new_width, new_height), Image.Resampling.BICUBIC)
+            left = int(torch.randint(0, new_width - self.img_size + 1, (1,)).item())
+            resized = image.crop((left, 0, left + self.img_size, self.img_size))
+        else:
+            new_width, new_height = self.img_size, math.ceil(self.img_size / width * height)
+            image = image.resize((new_width, new_height), Image.Resampling.BICUBIC)
+            lower = int(torch.randint(0, new_height - self.img_size + 1, (1,)).item())
+            resized = image.crop((0, lower, self.img_size, lower + self.img_size))
+
+        if resized.mode != "RGB":
+            resized = resized.convert("RGB")
+        np_img = np.array(resized).transpose((2, 0, 1))
+        return torch.tensor(np_img).float() / 255.0 * 2.0 - 1.0
+
+    def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]:
+        if isinstance(batch, dict):
+            batch = [batch]
+        if not isinstance(batch, Sequence) or not batch:
+            raise ValueError(
+                f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}"
+            )
+
+        images: list[torch.Tensor] = []
+        prompts: list[str] = []
+        for sample in batch:
+            image = self._process_image(sample["image"])
+            if image is None:
+                continue
+            prompt = str(sample.get("prompt", ""))
+            if self.prompt_dropout_prob > 0.0 and torch.rand(1).item() < self.prompt_dropout_prob:
+                prompt = ""
+            images.append(image)
+            prompts.append(prompt)
+
+        if not images:
+            raise ValueError("FLUX raw batch contained no usable images after preprocessing")
+        return {
+            "image": torch.stack(images, dim=0).to(device=device, dtype=dtype, non_blocking=True),
+            "prompts": prompts,
+        }
diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py
new file mode 100644
index 000000000..ce1be8a32
--- /dev/null
+++ b/primus/backends/diffusion/data/processing_wanvideo.py
@@ -0,0 +1,401 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import AutoTokenizer
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.image_utils import ImageInput
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import TensorType, logging
+
+logger = logging.get_logger(__name__)
+
+
+class WanVideoImageProcessor(BaseImageProcessor):
+    """
+    Image/Video processor for WanVideo models.
+
+    Args:
+        do_resize: Whether to resize the image/video frames.
+        size: Target size for resizing.
+        do_center_crop: Whether to center crop.
+        crop_size: Size for center cropping.
+        do_normalize: Whether to normalize pixel values.
+        image_mean: Mean values for normalization.
+        image_std: Standard deviation values for normalization.
+        do_convert_rgb: Whether to convert to RGB.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[Dict[str, int]] = None,
+        do_center_crop: bool = True,
+        crop_size: Optional[Dict[str, int]] = None,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = True,
+        max_pixels: Optional[int] = None,
+        height_division_factor: int = 16,
+        width_division_factor: int = 16,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.do_resize = do_resize
+        self.size = size
+        self.do_center_crop = do_center_crop
+        self.crop_size = crop_size
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean or [0.5, 0.5, 0.5]
+        self.image_std = image_std or [0.5, 0.5, 0.5]
+        self.do_convert_rgb = do_convert_rgb
+        self.max_pixels = max_pixels
+        self.height_division_factor = height_division_factor
+        self.width_division_factor = width_division_factor
+
+    def resize(
+        self,
+        image: np.ndarray,
+        size: Dict[str, int],
+        **kwargs,
+    ) -> np.ndarray:
+        """Resize image or video frame."""
+        from PIL import Image as PILImage
+
+        image = PILImage.fromarray(image.astype(np.uint8))
+        image = image.resize((size["width"], size["height"]), PILImage.LANCZOS)
+        return np.array(image)
+
+    def _extract_hw(self, image: Union[np.ndarray, Image.Image]) -> Tuple[int, int]:
+        if isinstance(image, Image.Image):
+            width, height = image.size
+            return height, width
+        if isinstance(image, np.ndarray):
+            if image.ndim == 3:
+                return image.shape[0], image.shape[1]
+            if image.ndim == 2:
+                return image.shape[0], image.shape[1]
+        raise ValueError("Unsupported image type for size extraction.")
+
+    def _dynamic_target_size(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]:
+        height, width = self._extract_hw(image)
+        if self.max_pixels is not None and width * height > self.max_pixels:
+            scale = (width * height / self.max_pixels) ** 0.5
+            height = int(height / scale)
+            width = int(width / scale)
+        height = height // self.height_division_factor * self.height_division_factor
+        width = width // self.width_division_factor * self.width_division_factor
+        height = max(self.height_division_factor, height)
+        width = max(self.width_division_factor, width)
+        return {"height": height, "width": width}
+
+    def _resolve_sizes(
+        self,
+        image: Union[np.ndarray, Image.Image],
+        size: Optional[Dict[str, int]] = None,
+        crop_size: Optional[Dict[str, int]] = None,
+    ) -> Tuple[Dict[str, int], Dict[str, int]]:
+        size = size if size is not None else self.size
+        crop_size = crop_size if crop_size is not None else self.crop_size
+        if size is None and crop_size is None:
+            size = self._dynamic_target_size(image)
+            crop_size = dict(size)
+        elif size is None:
+            size = dict(crop_size)
+        elif crop_size is None:
+            crop_size = dict(size)
+        size = {
+            "height": self._ceil_to_factor(size["height"], self.height_division_factor),
+            "width": self._ceil_to_factor(size["width"], self.width_division_factor),
+        }
+        crop_size = {
+            "height": self._ceil_to_factor(crop_size["height"], self.height_division_factor),
+            "width": self._ceil_to_factor(crop_size["width"], self.width_division_factor),
+        }
+        return size, crop_size
+
+    @staticmethod
+    def _ceil_to_factor(value: int, factor: int) -> int:
+        if value % factor == 0:
+            return value
+        return (value + factor - 1) // factor * factor
+
+    def center_crop(
+        self,
+        image: np.ndarray,
+        crop_size: Dict[str, int],
+        **kwargs,
+    ) -> np.ndarray:
+        """Center crop image or video frame."""
+        h, w = image.shape[:2]
+        crop_h, crop_w = crop_size["height"], crop_size["width"]
+
+        top = (h - crop_h) // 2
+        left = (w - crop_w) // 2
+        if image.ndim == 3:
+            return image[top : top + crop_h, left : left + crop_w, :]
+        else:
+            return image[top : top + crop_h, left : left + crop_w]
+
+    def normalize(
+        self,
+        image: np.ndarray,
+        mean: List[float],
+        std: List[float],
+        **kwargs,
+    ) -> np.ndarray:
+        """Normalize image or video frame."""
+        image = image.astype(np.float32) / 255.0
+        mean = np.array(mean).reshape(1, 1, -1)
+        std = np.array(std).reshape(1, 1, -1)
+        return (image - mean) / std
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[Dict[str, int]] = None,
+        do_center_crop: Optional[bool] = None,
+        crop_size: Optional[Dict[str, int]] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        num_frames: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Preprocess images or video frames.
+
+        Args:
+            images: Input images or video frames.
+            return_tensors: Type of tensors to return ("pt" for PyTorch).
+
+        Returns:
+            Dictionary with preprocessed pixel values.
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
+        crop_size = crop_size if crop_size is not None else self.crop_size
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        # Handle single image or list of images (video frames)
+        if not isinstance(images, list):
+            images = [images]
+
+        processed_images = []
+        for image in images:
+            if isinstance(image, Image.Image):
+                image = np.array(image)
+
+            # Handle 4D tensor case (batch of video frames)
+            if isinstance(image, np.ndarray) and image.ndim == 4:
+                # Process each frame in the batch
+                batch_processed_frames = []
+                for i in range(image.shape[0]):
+                    frame = image[i]  # Get single frame
+                    if frame.ndim == 3 and (frame.shape[0] == 3 or frame.shape[0] == 1):
+                        frame = np.transpose(frame, (1, 2, 0))  # (C, H, W) -> (H, W, C)
+
+                    # Convert to RGB if needed
+                    if do_convert_rgb and frame.shape[-1] != 3:
+                        if len(frame.shape) == 2:  # Grayscale
+                            frame = np.stack([frame] * 3, axis=-1)
+                        elif frame.shape[-1] == 4:  # RGBA
+                            frame = frame[..., :3]
+
+                    size_, crop_size_ = self._resolve_sizes(frame, size=size, crop_size=crop_size)
+                    if do_resize:
+                        frame = self.resize(frame, size_)
+                    if do_center_crop:
+                        frame = self.center_crop(frame, crop_size_)
+
+                    # Normalize
+                    if do_normalize:
+                        frame = self.normalize(frame, image_mean, image_std)
+
+                    batch_processed_frames.append(frame)
+
+                # Stack frames back into 4D tensor
+                processed_image = np.stack(batch_processed_frames, axis=0)
+            else:
+                # Handle single image case (3D or 2D)
+                # Convert to RGB if needed
+                if do_convert_rgb and image.shape[-1] != 3:
+                    if len(image.shape) == 2:  # Grayscale
+                        image = np.stack([image] * 3, axis=-1)
+                    elif image.shape[-1] == 4:  # RGBA
+                        image = image[..., :3]
+
+                size_, crop_size_ = self._resolve_sizes(image, size=size, crop_size=crop_size)
+                if do_resize:
+                    image = self.resize(image, size_)
+                if do_center_crop:
+                    image = self.center_crop(image, crop_size_)
+
+                # Normalize
+                if do_normalize:
+                    image = self.normalize(image, image_mean, image_std)
+
+                processed_image = image
+
+            processed_images.append(processed_image)
+
+        # Stack frames for video
+        processed_images = np.stack(processed_images, axis=0)  # B, T, H, W, C
+
+        # Temporal Handling (Interpolate or Truncate)
+        if num_frames is not None:
+            current_frames = processed_images.shape[1]
+            if current_frames > num_frames:
+                logger.info(f"Truncating video frames from {current_frames} to {num_frames}")
+                processed_images = processed_images[:, :num_frames, ...]
+            elif current_frames < num_frames:
+                logger.info(f"Interpolating video frames from {current_frames} to {num_frames}")
+                # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C)
+                # Permute to (B, C, T, H, W) for interpolate
+                vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3)
+
+                # Interpolate
+                vid_tensor = torch.nn.functional.interpolate(
+                    vid_tensor,
+                    size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]),
+                    mode="trilinear",
+                    align_corners=False,
+                )
+
+                # Permute back to (B, T, H, W, C) and convert to numpy
+                processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy()
+
+        # Convert to tensor if requested
+        if return_tensors == "pt":
+            processed_images = torch.from_numpy(processed_images)
+            # Rearrange to (B, C, T, H, W) for video (since input was B, T, H, W, C)
+            if processed_images.ndim == 5:
+                processed_images = processed_images.permute(0, 4, 1, 2, 3)
+            elif processed_images.ndim == 4:
+                # (T, H, W, C) -> (C, T, H, W) if it was list of frames
+                processed_images = processed_images.permute(3, 0, 1, 2)
+            # Add batch dimension
+            # processed_images = processed_images.unsqueeze(0)
+
+        return {"pixel_values": processed_images}
+
+
+class WanVideoProcessor:
+    """
+    Processor for WanVideo models, combining image/video processing and text tokenization.
+
+    Args:
+        image_processor: Image/video processor instance.
+        tokenizer: Text tokenizer instance.
+    """
+
+    attributes = ["tokenizer", "image_processor"]
+    valid_kwargs = [
+        "chat_template",
+    ]
+    tokenizer_class = "AutoTokenizer"
+    image_processor_class = "WanVideoImageProcessor"
+
+    def __init__(self, image_processor=None, tokenizer=None, max_text_length: Optional[int] = None, **kwargs):
+        if image_processor is None:
+            image_processor = WanVideoImageProcessor(**kwargs)
+        if tokenizer is None:
+            # Default to T5 tokenizer for text encoding
+            try:
+                logger.info("Loading default tokenizer: google/umt5-xxl")
+                tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl")
+            except Exception:
+                logger.warning("Could not load default tokenizer, using None")
+                tokenizer = None
+
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+        self.max_text_length = max_text_length
+
+    def __call__(
+        self,
+        text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None,
+        images: Optional[ImageInput] = None,
+        videos: Optional[ImageInput] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Process text and image/video inputs.
+
+        Args:
+            text: Text input(s) to tokenize.
+            images: Image input(s) to process.
+            videos: Video frames to process.
+            return_tensors: Type of tensors to return.
+
+        Returns:
+            Dictionary with processed inputs.
+        """
+        if text is None and images is None and videos is None:
+            raise ValueError("You must provide either text, images, or videos.")
+
+        data = {}
+
+        # Process text
+        if text is not None and self.tokenizer is not None:
+            max_length = kwargs.pop("max_length", None)
+            if max_length is None:
+                max_length = self.max_text_length or getattr(self.tokenizer, "model_max_length", 256)
+            text_inputs = self.tokenizer(
+                text,
+                return_tensors=return_tensors,
+                padding=True,
+                truncation=True,
+                max_length=max_length,
+                **kwargs,
+            )
+            data.update(text_inputs)
+
+        # Process images or video
+        if images is not None or videos is not None:
+            image_inputs = self.image_processor(
+                images or videos,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+            data.update(image_inputs)
+
+        return data
+
+    def batch_decode(self, *args, **kwargs):
+        """Delegate to tokenizer's batch_decode method."""
+        if self.tokenizer is None:
+            raise ValueError("No tokenizer available for decoding.")
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """Delegate to tokenizer's decode method."""
+        if self.tokenizer is None:
+            raise ValueError("No tokenizer available for decoding.")
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """Get model input names from components."""
+        tokenizer_input_names = self.tokenizer.model_input_names if self.tokenizer else []
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(set(tokenizer_input_names + image_processor_input_names))
diff --git a/primus/backends/diffusion/data/processor.py b/primus/backends/diffusion/data/processor.py
new file mode 100644
index 000000000..6434d24a6
--- /dev/null
+++ b/primus/backends/diffusion/data/processor.py
@@ -0,0 +1,182 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Processes video and text data for training.
+"""
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+
+class WanVideoDataProcessor:
+    """Standalone data processor for WanVideo training."""
+
+    def __init__(self, config, model_id=None):
+        self.config = config
+        self.model_id = model_id
+        self.processor = None
+        self.tokenizer = None
+
+    def apply_prompt_template(self, hf_messages: str) -> str:
+        """Apply prompt template for WanVideo."""
+        # WanVideo uses direct prompts without special formatting.
+        # Keep backward compatibility:
+        # - old path passed `hf_messages` list[dict]
+        # - new path passes plain prompt string
+        if isinstance(hf_messages, str):
+            return hf_messages
+        try:
+            return hf_messages[0]["content"][1]["text"]
+        except Exception:
+            # Best-effort fallback
+            return str(hf_messages)
+
+    def save_pretrained(self, save_directory: str):
+        pass
+
+    def build(self):
+        """Initialize the processor and tokenizer."""
+        if self.processor is not None and self.tokenizer is not None:
+            return
+
+        from transformers import AutoTokenizer
+
+        from primus.backends.diffusion.data.processing_wanvideo import (
+            WanVideoProcessor as WanVideoModelProcessor,
+        )
+
+        wanvideo_kwargs = self.config.get("extra_kwargs", {})
+        max_text_length = self.config.get("max_text_length")
+        if max_text_length is not None:
+            wanvideo_kwargs.setdefault("max_text_length", max_text_length)
+
+        # Load tokenizer if specified
+        if self.config.get("text_tokenizer", None) is not None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("text_tokenizer"))
+        else:
+            self.tokenizer = None
+        self.processor = WanVideoModelProcessor(**wanvideo_kwargs, tokenizer=self.tokenizer)
+
+        if self.tokenizer is None:
+            self.tokenizer = self.processor.tokenizer
+
+    def _normalize_raw_batch(self, batch: Any) -> tuple[list[str], list[Any], int | None]:
+        """
+        Normalize dataloader output into:
+          - prompts: list[str]
+          - frames_list: list[Any]  (each item is typically np.ndarray (T,H,W,C))
+          - num_frames: Optional[int] (only if consistent across samples)
+        """
+        if isinstance(batch, dict):
+            prompts = batch.get("prompt")
+            frames_list = batch.get("video_frames")
+            num_frames = batch.get("num_frames", None)
+
+            if prompts is None or frames_list is None:
+                raise ValueError("prepare_batch(dict) requires keys: 'prompt' and 'video_frames'")
+            if isinstance(prompts, str):
+                prompts = [prompts]
+            if not isinstance(prompts, (list, tuple)):
+                raise TypeError(f"'prompt' must be str or list[str], got {type(prompts)}")
+            if not isinstance(frames_list, (list, tuple)):
+                raise TypeError(f"'video_frames' must be list, got {type(frames_list)}")
+
+            if isinstance(num_frames, (list, tuple)):
+                nfs = [int(x) for x in num_frames if x]
+                num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None
+            elif num_frames:
+                num_frames = int(num_frames)
+            else:
+                num_frames = None
+
+            return list(prompts), list(frames_list), num_frames
+
+        # RawBatchCollator returns list[dict]
+        if not isinstance(batch, (list, tuple)) or not batch:
+            raise ValueError(f"prepare_batch expected non-empty list/tuple, got {type(batch)}")
+
+        prompts = [str(ex.get("prompt", "")) for ex in batch]
+        frames_list = [ex.get("video_frames") for ex in batch]
+
+        # Prefer per-example num_frames if present and consistent; else None.
+        nfs = [ex.get("num_frames") for ex in batch if ex.get("num_frames")]
+        num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None
+        num_frames = int(num_frames) if num_frames else None
+
+        return prompts, frames_list, num_frames
+
+    def _tokenize_prompts(self, prompts: Sequence[str]) -> dict[str, torch.Tensor]:
+        formatted_prompts = [self.apply_prompt_template(p) for p in prompts]
+        return self.tokenizer(
+            formatted_prompts,
+            return_tensors="pt",
+            padding=self.config.get("padding_strategy", "max_length"),
+            truncation=True,
+            max_length=self.config.get("max_text_length", 512),
+        )
+
+    def _preprocess_videos(self, frames_list: Sequence[Any], num_frames: int | None) -> torch.Tensor:
+        if any(v is None for v in frames_list):
+            raise ValueError("prepare_batch got None in 'video_frames'")
+        video_inputs = self.processor.image_processor.preprocess(
+            list(frames_list),
+            num_frames=num_frames,
+            return_tensors="pt",
+        )
+        if "pixel_values" not in video_inputs:
+            raise KeyError("image_processor.preprocess must return dict with key 'pixel_values'")
+        pixel_values = video_inputs["pixel_values"]
+        if not isinstance(pixel_values, torch.Tensor):
+            raise TypeError(f"pixel_values must be torch.Tensor, got {type(pixel_values)}")
+        if pixel_values.ndim != 5:
+            raise ValueError(f"Expected pixel_values shape [B,C,T,H,W], got {tuple(pixel_values.shape)}")
+        return pixel_values
+
+    def _assemble_model_batch(
+        self, *, pixel_values: torch.Tensor, text_inputs: dict[str, torch.Tensor]
+    ) -> dict[str, Any]:
+        if "input_ids" not in text_inputs or "attention_mask" not in text_inputs:
+            raise KeyError("tokenizer output must include 'input_ids' and 'attention_mask'")
+
+        out: dict[str, Any] = {
+            "video": pixel_values,  # [B,C,T,H,W]
+            "input_ids": text_inputs["input_ids"],
+            "attention_mask": text_inputs["attention_mask"],
+        }
+
+        # Legacy convenience fields (some models expect these keys).
+        out.setdefault("num_frames", int(pixel_values.shape[2]))
+        out.setdefault("height", int(pixel_values.shape[3]))
+        out.setdefault("width", int(pixel_values.shape[4]))
+        return out
+
+    def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]:
+        """
+        Convert raw dataloader batch into model inputs.
+
+        Expected raw batch from `RawBatchCollator`:
+          - batch: list[dict] with keys:
+              - video_frames: np.ndarray (T,H,W,C) or equivalent
+              - prompt: str
+              - num_frames: int (optional)
+
+        Output matches current training pipelines:
+          - video: Tensor [B, C, T, H, W]
+          - input_ids: Tensor [B, L]
+          - attention_mask: Tensor [B, L]
+          - (optional) height/width/num_frames for legacy wan/wan_new paths
+        """
+        if self.processor is None or self.tokenizer is None:
+            # Lazy init to keep behavior robust across call sites.
+            self.build()
+
+        prompts, frames_list, num_frames = self._normalize_raw_batch(batch)
+        text_inputs = self._tokenize_prompts(prompts)
+        pixel_values = self._preprocess_videos(frames_list, num_frames=num_frames)
+        return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs)
diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py
new file mode 100644
index 000000000..4573e0a65
--- /dev/null
+++ b/primus/backends/diffusion/data/registrations/__init__.py
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Dataset registrations."""
diff --git a/primus/backends/diffusion/data/registrations/flux.py b/primus/backends/diffusion/data/registrations/flux.py
new file mode 100644
index 000000000..234868970
--- /dev/null
+++ b/primus/backends/diffusion/data/registrations/flux.py
@@ -0,0 +1,38 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from primus.backends.diffusion.data.flux_precomputed import (
+    FluxPrecomputedDataset,
+    FluxPrecomputedProcessor,
+    FluxRawImageTextDataset,
+    FluxRawImageTextProcessor,
+)
+from primus.backends.diffusion.utils.log import logger
+
+
+def build_flux_dataset(dataset_config: dict):
+    processor_config = dataset_config.get("processor_config", {}) or {}
+    dataset_type = str(dataset_config.get("dataset_type", "precomputed")).lower()
+    if dataset_type == "raw":
+        processor = FluxRawImageTextProcessor(processor_config)
+        processor.build()
+        dataset = FluxRawImageTextDataset(
+            dataset_path=dataset_config.get("dataset_path"),
+            dataset_format=dataset_config.get("dataset_format", "webdataset"),
+            dataset_name=dataset_config.get("dataset"),
+            data_folder=dataset_config.get("data_folder"),
+        )
+        logger.info(f"Built FLUX raw image-text dataset with {len(dataset)} samples")
+    elif dataset_type == "precomputed":
+        processor = FluxPrecomputedProcessor(processor_config)
+        processor.build()
+        dataset = FluxPrecomputedDataset(dataset_config["dataset_path"])
+        logger.info(f"Built FLUX precomputed dataset with {len(dataset)} samples")
+    else:
+        raise ValueError("FLUX dataset_type must be either 'precomputed' or 'raw'")
+    return dataset, processor
diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py
new file mode 100644
index 000000000..f832f352a
--- /dev/null
+++ b/primus/backends/diffusion/data/registrations/wan.py
@@ -0,0 +1,28 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Register Wan dataset builder."""
+
+from primus.backends.diffusion.data import (
+    DatasetConfig,
+    WanVideoDataProcessor,
+    WanVideoDataset,
+)
+from primus.backends.diffusion.utils.log import logger
+
+
+def build_wan_dataset(dataset_config: dict):
+    processor_config = dataset_config["processor_config"]
+    processor = WanVideoDataProcessor(processor_config)
+    processor.build()
+    logger.info("Built data processor")
+
+    dataset = WanVideoDataset(
+        processor=processor,
+        config=DatasetConfig(**dataset_config),
+    )
+    logger.info(f"Built dataset with {len(dataset)} samples")
+    return dataset, processor
diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py
new file mode 100644
index 000000000..878ea3fea
--- /dev/null
+++ b/primus/backends/diffusion/diffusion_adapter.py
@@ -0,0 +1,68 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from primus.backends.diffusion.argument_builder import DiffusionArgBuilder
+from primus.core.backend.backend_adapter import BackendAdapter
+from primus.core.utils.module_utils import log_rank_0
+
+
+class DiffusionAdapter(BackendAdapter):
+    """Primus adapter for PyTorch diffusion training."""
+
+    def __init__(self, framework: str = "diffusion"):
+        super().__init__(framework)
+
+    def setup_backend_path(self, backend_path=None) -> str:
+        """Validate the Primus-owned in-tree diffusion package."""
+        if backend_path:
+            raise ValueError(
+                "The diffusion backend is built into Primus and does not support "
+                f"external backend_path overrides. Got: {backend_path}"
+            )
+
+        resolved = Path(__file__).resolve().parent
+        if not resolved.exists():
+            raise FileNotFoundError(f"[Primus:Diffusion] backend package does not exist: {resolved}")
+
+        resolved_str = str(resolved)
+        try:
+            log_rank_0(f"[Primus:Diffusion] using in-tree backend package -> {resolved_str}")
+        except Exception:
+            # Best-effort startup logging should not block backend setup.
+            pass
+
+        return resolved_str
+
+    def convert_config(self, params: Any):
+        builder = DiffusionArgBuilder()
+        builder.update(params)
+        diffusion_args = builder.finalize()
+        # convert_config is also called by the standalone prepare hook, where the
+        # Primus logger may not be initialized yet; guard the informational log.
+        try:
+            model_name = getattr(getattr(diffusion_args, "model", {}), "get", lambda *_: None)("name")
+            log_rank_0(f"[Primus:DiffusionAdapter] Converted Primus module params -> {model_name} args")
+        except Exception:
+            # Standalone prepare hooks may run before the Primus logger is bound.
+            pass
+        return diffusion_args
+
+    def load_trainer_class(self, stage: str = "pretrain"):
+        if stage in ("pretrain", "posttrain", "sft"):
+            from primus.backends.diffusion.diffusion_pretrain_trainer import (
+                DiffusionPretrainTrainer,
+            )
+
+            return DiffusionPretrainTrainer
+        raise ValueError(f"Invalid stage for Diffusion backend: {stage}")
+
+    def detect_backend_version(self) -> str:
+        return "in-tree"
diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py
new file mode 100644
index 000000000..218fc7019
--- /dev/null
+++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py
@@ -0,0 +1,147 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import importlib.util
+from typing import Any
+
+from primus.core.base_module import BaseModule
+from primus.core.trainer.base_trainer import BaseTrainer
+from primus.core.utils.module_utils import log_rank_0
+from primus.core.utils.yaml_utils import nested_namespace_to_dict
+
+
+class DiffusionPretrainTrainer(BaseTrainer, BaseModule):
+    """Primus lifecycle wrapper for diffusion backend training."""
+
+    def __init__(self, backend_args: Any, *args, **kwargs):
+        super().__init__(backend_args=backend_args, *args, **kwargs)
+        self.diffusion_trainer = None
+
+    @staticmethod
+    def _as_dict(value: Any) -> dict:
+        if isinstance(value, dict):
+            return value
+        return nested_namespace_to_dict(value)
+
+    def setup(self):
+        trainer_cfg = self._as_dict(self.backend_args.trainer)
+        dataset_cfg = self._as_dict(self.backend_args.dataset)
+        trainer_args = trainer_cfg.get("args", {})
+        attention_backend = trainer_args.get("attention_backend")
+
+        missing = [
+            package
+            for package in (
+                "torch",
+                "loguru",
+                "safetensors",
+                "transformers",
+                "PIL",
+                "torchvision",
+                "requests",
+                "packaging",
+            )
+            if importlib.util.find_spec(package) is None
+        ]
+        dataset_config = dataset_cfg.get("config", {}) or {}
+        video_backend = dataset_config.get("video_backend")
+        if dataset_cfg.get("name") == "flux":
+            for package in ("datasets", "huggingface_hub", "sentencepiece"):
+                if importlib.util.find_spec(package) is None:
+                    missing.append(package)
+            if (
+                dataset_config.get("dataset_type") == "raw"
+                and dataset_config.get("dataset_format") == "webdataset"
+                and importlib.util.find_spec("webdataset") is None
+            ):
+                missing.append("webdataset")
+        if video_backend == "imageio" and importlib.util.find_spec("imageio") is None:
+            missing.append("imageio")
+        if video_backend == "decord" and importlib.util.find_spec("decord") is None:
+            missing.append("decord")
+        if missing:
+            raise RuntimeError(
+                "Diffusion backend missing required Python packages: "
+                f"{', '.join(missing)}. Install the diffusion training extras first."
+            )
+
+        if attention_backend:
+            from primus.backends.diffusion.attention import set_attention_backend
+
+            set_attention_backend(attention_backend)
+            log_rank_0(f"[Primus:Diffusion] attention_backend={attention_backend}")
+
+        if attention_backend == "flash_attn_aiter":
+            from primus.backends.diffusion.attention.aiter import (
+                AITER_FLASH_ATTN_AVAILABLE,
+            )
+
+            if not AITER_FLASH_ATTN_AVAILABLE:
+                raise RuntimeError(
+                    "attention_backend=flash_attn_aiter was requested, but AITER flash attention "
+                    "is unavailable in this environment."
+                )
+
+    def init(self):
+        from primus.backends.diffusion.registry import (
+            get_dataset_builder,
+            get_model_builder,
+            get_trainer_builder,
+        )
+
+        model_cfg = self._as_dict(self.backend_args.model)
+        dataset_cfg = self._as_dict(self.backend_args.dataset)
+        trainer_cfg = self._as_dict(self.backend_args.trainer)
+
+        model_name = model_cfg["name"]
+        dataset_name = dataset_cfg["name"]
+        trainer_name = trainer_cfg["name"]
+
+        model_config = model_cfg["config"]
+        dataset_config = dataset_cfg["config"]
+        trainer_args = trainer_cfg["args"]
+
+        log_rank_0(
+            f"[Primus:Diffusion] Building model={model_name}, dataset={dataset_name}, trainer={trainer_name}"
+        )
+        model = get_model_builder(model_name)(model_config)
+        dataset, processor = get_dataset_builder(dataset_name)(dataset_config)
+        self.diffusion_trainer = get_trainer_builder(trainer_name)(
+            model=model,
+            dataset=dataset,
+            processor=processor,
+            trainer_args=trainer_args,
+        )
+
+    def train(self):
+        if self.diffusion_trainer is None:
+            raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().")
+
+        self.diffusion_trainer.train()
+        self.diffusion_trainer.save_model()
+
+    def run(self, *args, **kwargs):
+        """Compatibility hook for BaseModule; TrainRuntime drives lifecycle phases."""
+        self.train()
+
+    def cleanup(self, on_error: bool = False):
+        try:
+            import wandb
+
+            if getattr(wandb, "run", None) is not None:
+                wandb.finish(exit_code=1 if on_error else 0)
+        except Exception as exc:
+            log_rank_0(f"[Primus:Diffusion] wandb cleanup failed: {exc}")
+
+        try:
+            import torch.distributed as dist
+
+            if dist.is_available() and dist.is_initialized():
+                dist.destroy_process_group()
+        except Exception as exc:
+            log_rank_0(f"[Primus:Diffusion] distributed cleanup failed: {exc}")
diff --git a/primus/backends/diffusion/distributed/__init__.py b/primus/backends/diffusion/distributed/__init__.py
new file mode 100644
index 000000000..7fea79b62
--- /dev/null
+++ b/primus/backends/diffusion/distributed/__init__.py
@@ -0,0 +1,32 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus diffusion distributed training utilities.
+
+Modules:
+  mesh        - Device mesh creation, distributed setup
+  checkpoint  - DTCP sharded checkpoint save/load
+  ulysses     - Ulysses Sequence Parallel primitives
+"""
+
+from .checkpoint import load_checkpoint_dtcp, save_checkpoint_dtcp
+from .mesh import create_device_mesh, setup_distributed
+from .ulysses import distributed_attention, sp_gather, sp_split, sp_unpad
+
+__all__ = [
+    # mesh
+    "setup_distributed",
+    "create_device_mesh",
+    # checkpoint
+    "save_checkpoint_dtcp",
+    "load_checkpoint_dtcp",
+    # ulysses
+    "distributed_attention",
+    "sp_split",
+    "sp_gather",
+    "sp_unpad",
+]
diff --git a/primus/backends/diffusion/distributed/checkpoint.py b/primus/backends/diffusion/distributed/checkpoint.py
new file mode 100644
index 000000000..5ca2f32ae
--- /dev/null
+++ b/primus/backends/diffusion/distributed/checkpoint.py
@@ -0,0 +1,114 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Distributed Tensor Checkpointing (DTCP) save / load utilities.
+
+Handles FSDP2 sharded checkpoints without gathering to rank 0.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter, load, save
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_model_state_dict,
+    get_optimizer_state_dict,
+    set_model_state_dict,
+    set_optimizer_state_dict,
+)
+
+from primus.backends.diffusion.utils.log import logger
+
+from .mesh import _ensure_process_group
+
+_META_FILENAME = "meta.json"
+
+
+def save_checkpoint_dtcp(
+    model: torch.nn.Module,
+    optimizer: Optional[torch.optim.Optimizer],
+    path: str,
+    epoch: int,
+    step: int,
+    additional_data: Dict[str, Any] = None,
+    *,
+    model_state_options: Optional[StateDictOptions] = None,
+    optim_state_options: Optional[StateDictOptions] = None,
+):
+    """Save checkpoint using DTCP (sharded, no gather to rank 0)."""
+    os.makedirs(path, exist_ok=True)
+    _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")
+
+    model_state_options = model_state_options or StateDictOptions(full_state_dict=False)
+    optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False)
+
+    model_state = get_model_state_dict(model, options=model_state_options)
+    optim_state = None
+    if optimizer is not None:
+        optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options)
+
+    meta = {"epoch": epoch, "step": step, **(additional_data or {})}
+    state_dict = {
+        "model": model_state,
+        **({"optimizer": optim_state} if optim_state is not None else {}),
+        "meta": meta,
+    }
+
+    save(state_dict, FileSystemWriter(path))
+    if (not dist.is_initialized()) or dist.get_rank() == 0:
+        with open(os.path.join(path, _META_FILENAME), "w", encoding="utf-8") as f:
+            json.dump(meta, f, indent=2, sort_keys=True)
+        logger.info(f"Saved DTCP checkpoint to {path}")
+
+
+def load_checkpoint_dtcp(
+    model: torch.nn.Module,
+    optimizer: Optional[torch.optim.Optimizer],
+    path: str,
+    *,
+    model_state_options: Optional[StateDictOptions] = None,
+    optim_state_options: Optional[StateDictOptions] = None,
+) -> Dict[str, Any]:
+    """Load checkpoint using DTCP. Updates model/optimizer in-place. Returns metadata."""
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Checkpoint not found at {path}")
+
+    _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")
+
+    model_state_options = model_state_options or StateDictOptions(full_state_dict=False)
+    optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False)
+
+    model_state = get_model_state_dict(model, options=model_state_options)
+    optim_state = None
+    if optimizer is not None:
+        optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options)
+
+    state_dict = {
+        "model": model_state,
+        **({"optimizer": optim_state} if optim_state is not None else {}),
+        "meta": {},
+    }
+
+    load(state_dict, FileSystemReader(path))
+
+    set_model_state_dict(model, model_state, options=model_state_options)
+    if optimizer is not None and optim_state is not None:
+        set_optimizer_state_dict(model, optimizer, optim_state, options=optim_state_options)
+
+    meta = state_dict.get("meta", {})
+    if not meta:
+        meta_path = os.path.join(path, _META_FILENAME)
+        if os.path.exists(meta_path):
+            with open(meta_path, encoding="utf-8") as f:
+                meta = json.load(f)
+    return meta
diff --git a/primus/backends/diffusion/distributed/mesh.py b/primus/backends/diffusion/distributed/mesh.py
new file mode 100644
index 000000000..b9ee3b7b4
--- /dev/null
+++ b/primus/backends/diffusion/distributed/mesh.py
@@ -0,0 +1,126 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Device mesh creation and distributed setup.
+
+Supports:
+  - Pure FSDP2 (1D mesh: dp_shard)
+  - FSDP2 + HSDP (2D mesh: dp_replicate × dp_shard)
+  - FSDP2 + Ulysses SP (mesh: dp_shard × ulysses, flattened into dp_shard_sp)
+  - FSDP2 + HSDP + Ulysses SP (mesh: dp_replicate × dp_shard × ulysses)
+
+SP ranks are included in the FSDP sharding mesh so that parameters are
+sharded across both DP and SP ranks — this reduces per-rank memory.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+
+from primus.backends.diffusion.utils.log import logger
+
+
+def _ensure_process_group(*, backend: str) -> None:
+    """Best-effort init_process_group (lazy, avoids single-GPU hangs)."""
+    if dist.is_initialized():
+        return
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "12345")
+    dist.init_process_group(backend, timeout=timedelta(minutes=60))
+
+
+def setup_distributed() -> tuple[int, int, int]:
+    """
+    Initialize distributed process group.
+
+    Returns: (rank, world_size, local_rank)
+    """
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+
+    if world_size > 1 and not dist.is_initialized():
+        _ensure_process_group(backend="nccl")
+
+    return rank, world_size, local_rank
+
+
+def create_device_mesh(
+    world_size: int,
+    sp_size: int = 1,
+    dp_replicate: int = 1,
+) -> Optional[DeviceMesh]:
+    """
+    Create a DeviceMesh for FSDP2 + optional Ulysses Sequence Parallel.
+
+    Mesh layout (innermost → outermost):
+        [dp_replicate?] × [dp_shard] × [ulysses?]
+
+    Flattened sub-meshes created automatically:
+        - "dp_shard_sp":  dp_shard × ulysses  (used for FSDP2 fully_shard)
+        - "dp":           dp_replicate × dp_shard  (used for DistributedSampler)
+
+    Args:
+        world_size: total number of ranks
+        sp_size:    Ulysses sequence parallel size (must divide world_size)
+        dp_replicate: HSDP replicate dimension (1 = no HSDP)
+
+    Returns:
+        DeviceMesh, or None if world_size <= 1
+    """
+    if world_size <= 1:
+        return None
+
+    _ensure_process_group(backend="nccl")
+
+    dp_shard = world_size // (sp_size * dp_replicate)
+    if dp_shard * sp_size * dp_replicate != world_size:
+        raise ValueError(
+            f"world_size={world_size} is not divisible by " f"sp_size={sp_size} * dp_replicate={dp_replicate}"
+        )
+
+    # Build mesh dimensions
+    dims = []
+    names = []
+    if dp_replicate > 1:
+        dims.append(dp_replicate)
+        names.append("dp_replicate")
+    dims.append(dp_shard)
+    names.append("dp_shard")
+    if sp_size > 1:
+        dims.append(sp_size)
+        names.append("ulysses")
+
+    mesh = init_device_mesh("cuda", tuple(dims), mesh_dim_names=tuple(names))
+
+    # Flatten composite sub-meshes for convenient access
+    if sp_size > 1:
+        mesh["dp_shard", "ulysses"]._flatten("dp_shard_sp")
+
+    if dp_replicate > 1:
+        mesh["dp_replicate", "dp_shard"]._flatten("dp")
+
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    if rank == 0:
+        logger.info(
+            f"DeviceMesh created: dims={dict(zip(names, dims))}, "
+            f"sp_size={sp_size}, dp_replicate={dp_replicate}"
+        )
+
+    return mesh
diff --git a/primus/backends/diffusion/distributed/ulysses.py b/primus/backends/diffusion/distributed/ulysses.py
new file mode 100644
index 000000000..b61ae80e2
--- /dev/null
+++ b/primus/backends/diffusion/distributed/ulysses.py
@@ -0,0 +1,211 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Ulysses Sequence Parallel primitives.
+
+Follows the Wan2.2 official design: a self-contained ``distributed_attention``
+function that wraps all-to-all + attention + all-to-all, plus ``sp_split`` /
+``sp_gather`` / ``sp_unpad`` for slicing model inputs and gathering outputs.
+
+The model code only needs to:
+  1. ``sp_split`` inputs before the blocks
+  2. call ``distributed_attention`` instead of ``attention`` in self-attn
+  3. ``sp_gather`` + ``sp_unpad`` the output after the head
+
+Gradient scaling (÷ sp_size in sp_slice backward, × sp_size in sp_gather
+backward) ensures correctness with FSDP2 gradient averaging across the
+combined DP+SP mesh.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Core autograd primitives
+# ---------------------------------------------------------------------------
+
+
+def _require_divisible_dim(x: torch.Tensor, dim: int, divisor: int, op_name: str) -> None:
+    size = x.shape[dim]
+    if size % divisor != 0:
+        raise ValueError(
+            f"{op_name} requires tensor dimension {dim} (size={size}) to be divisible by "
+            f"sequence parallel size {divisor}. Got shape={tuple(x.shape)}."
+        )
+
+
+class _SeqAllToAll(torch.autograd.Function):
+    """All-to-all with autograd.  Backward is the inverse (swap dims)."""
+
+    @staticmethod
+    def forward(ctx, group, x, scatter_dim, gather_dim):
+        ctx.group = group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        sp_size = dist.get_world_size(group)
+        _require_divisible_dim(x, scatter_dim, sp_size, "_SeqAllToAll")
+        input_list = [t.contiguous() for t in x.tensor_split(sp_size, scatter_dim)]
+        output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)]
+        dist.all_to_all(output_list, input_list, group=group)
+        return torch.cat(output_list, dim=gather_dim).contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return (
+            None,
+            _SeqAllToAll.apply(ctx.group, grad_output, ctx.gather_dim, ctx.scatter_dim),
+            None,
+            None,
+        )
+
+
+class _SliceWithGather(torch.autograd.Function):
+    """Forward: slice.  Backward: all-gather ÷ sp_size (FSDP2 compat)."""
+
+    @staticmethod
+    def forward(ctx, x, dim, group):
+        ctx.dim = dim
+        ctx.group = group
+        sp_size = dist.get_world_size(group)
+        sp_rank = dist.get_rank(group)
+        ctx.sp_size = sp_size
+        _require_divisible_dim(x, dim, sp_size, "_SliceWithGather")
+        chunk_size = x.shape[dim] // sp_size
+        return x.narrow(dim, sp_rank * chunk_size, chunk_size).contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        gathered = [torch.empty_like(grad_output) for _ in range(ctx.sp_size)]
+        dist.all_gather(gathered, grad_output.contiguous(), group=ctx.group)
+        return torch.cat(gathered, dim=ctx.dim) / ctx.sp_size, None, None
+
+
+class _GatherWithSlice(torch.autograd.Function):
+    """Forward: all-gather.  Backward: slice × sp_size (FSDP2 compat)."""
+
+    @staticmethod
+    def forward(ctx, x, dim, group):
+        ctx.dim = dim
+        ctx.group = group
+        sp_size = dist.get_world_size(group)
+        sp_rank = dist.get_rank(group)
+        ctx.sp_size = sp_size
+        ctx.sp_rank = sp_rank
+        ctx.chunk_size = x.shape[dim]
+        gathered = [torch.empty_like(x) for _ in range(sp_size)]
+        dist.all_gather(gathered, x.contiguous(), group=group)
+        return torch.cat(gathered, dim=dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        chunk = grad_output.narrow(ctx.dim, ctx.sp_rank * ctx.chunk_size, ctx.chunk_size).contiguous()
+        return chunk * ctx.sp_size, None, None
+
+
+# ---------------------------------------------------------------------------
+# Public API — high-level
+# ---------------------------------------------------------------------------
+
+
+def distributed_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    group: dist.ProcessGroup,
+    attention_fn: Callable,
+    **attention_kwargs,
+) -> torch.Tensor:
+    """
+    Ulysses distributed attention (DeepSpeed Ulysses, arXiv:2309.14509).
+
+    Wraps ``attention_fn`` with all-to-all communication so that each rank
+    computes attention on the **full sequence** with a **subset of heads**.
+
+    Input / output shapes: ``[B, S/P, H, D]`` (sharded-seq, full-heads).
+
+    Args:
+        q, k, v:  query / key / value  ``[B, S/P, H, D]``
+        group:    Ulysses SP process group
+        attention_fn:  e.g. ``primus.backends.diffusion.attention.attention``
+        **attention_kwargs:  forwarded to ``attention_fn(q=, k=, v=, ...)``
+    """
+    # scatter heads, gather seq: [B, S/P, H, D] → [B, S, H/P, D]
+    q = _SeqAllToAll.apply(group, q, 2, 1)
+    k = _SeqAllToAll.apply(group, k, 2, 1)
+    v = _SeqAllToAll.apply(group, v, 2, 1)
+    # standard attention on full sequence with partial heads
+    out = attention_fn(q=q, k=k, v=v, **attention_kwargs)
+    # scatter seq, gather heads: [B, S, H/P, D] → [B, S/P, H, D]
+    return _SeqAllToAll.apply(group, out, 1, 2)
+
+
+def sp_split(
+    tensors: List[torch.Tensor],
+    dim: int,
+    group: dist.ProcessGroup,
+) -> Tuple[List[torch.Tensor], int]:
+    """
+    Pad + slice multiple tensors for sequence parallelism.
+
+    Returns ``(sliced_tensors, original_size)`` where *original_size* is
+    the size along *dim* before padding (needed by :func:`sp_unpad`).
+    """
+    sp_size = dist.get_world_size(group)
+    original_size: Optional[int] = None
+    results: List[torch.Tensor] = []
+    for t in tensors:
+        t, orig = _sp_pad(t, dim, sp_size)
+        if original_size is None:
+            original_size = orig
+        results.append(_SliceWithGather.apply(t, dim, group))
+    assert original_size is not None
+    return results, original_size
+
+
+# ---------------------------------------------------------------------------
+# Public API — low-level (used directly for gathering the output)
+# ---------------------------------------------------------------------------
+
+
+def sp_slice(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor:
+    """Slice ``x`` along *dim* for this SP rank (with autograd). Low-level."""
+    return _SliceWithGather.apply(x, dim, group)
+
+
+def sp_gather(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor:
+    """All-gather ``x`` along *dim* from all SP ranks (with autograd)."""
+    return _GatherWithSlice.apply(x, dim, group)
+
+
+def sp_unpad(x: torch.Tensor, dim: int, original_size: int) -> torch.Tensor:
+    """Remove padding added by :func:`sp_split`."""
+    if x.shape[dim] == original_size:
+        return x
+    return x.narrow(dim, 0, original_size)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _sp_pad(x: torch.Tensor, dim: int, sp_size: int) -> Tuple[torch.Tensor, int]:
+    """Pad *x* along *dim* so its size is divisible by *sp_size*."""
+    original_size = x.shape[dim]
+    remainder = original_size % sp_size
+    if remainder == 0:
+        return x, original_size
+    pad_amount = sp_size - remainder
+    pad_config = [0] * (2 * x.dim())
+    pos = 2 * (x.dim() - 1 - dim)
+    pad_config[pos + 1] = pad_amount
+    return F.pad(x, pad_config), original_size
diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py
new file mode 100644
index 000000000..7503f134c
--- /dev/null
+++ b/primus/backends/diffusion/models/__init__.py
@@ -0,0 +1,13 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Model exports for the Primus diffusion backend."""
+
+from .wan import WanForTraining
+
+__all__ = [
+    "WanForTraining",
+]
diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py
new file mode 100644
index 000000000..04ec5ff51
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/__init__.py
@@ -0,0 +1,25 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from primus.backends.diffusion.models.flux.adapter import FluxForTraining
+from primus.backends.diffusion.models.flux.model import (
+    Flux,
+    FluxParams,
+    flux_1_dev_params,
+    flux_1_schnell_params,
+)
+from primus.backends.diffusion.models.flux.train_pipeline import (
+    FluxFlowMatchTrainPipeline,
+)
+
+__all__ = [
+    "Flux",
+    "FluxForTraining",
+    "FluxFlowMatchTrainPipeline",
+    "FluxParams",
+    "flux_1_dev_params",
+    "flux_1_schnell_params",
+]
diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py
new file mode 100644
index 000000000..428306ce7
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/adapter.py
@@ -0,0 +1,113 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from primus.backends.diffusion.models.flux.train_pipeline import (
+    FluxFlowMatchTrainPipeline,
+)
+from primus.backends.diffusion.models.interface import GenAIModel
+
+
+@dataclass
+class FluxConfigShim:
+    raw: dict
+
+    def save_pretrained(self, save_directory: str):
+        import json
+        import os
+
+        os.makedirs(save_directory, exist_ok=True)
+        with open(os.path.join(save_directory, "flux_config.json"), "w") as f:
+            json.dump(self.raw, f, indent=2, sort_keys=True)
+
+    def to_dict(self):
+        return self.raw
+
+
+class FluxForTraining(GenAIModel, nn.Module):
+    """Thin Primus adapter around the FLUX DiT backbone."""
+
+    def __init__(
+        self,
+        *,
+        dit: nn.Module,
+        train_pipeline: FluxFlowMatchTrainPipeline,
+        model_config: Any,
+        autoencoder: nn.Module | None = None,
+        t5_encoder: nn.Module | None = None,
+        clip_encoder: nn.Module | None = None,
+        raw_config: dict | None = None,
+        trainable_modules: str | None = None,
+    ):
+        super().__init__()
+        self.dit = dit
+        self.autoencoder = autoencoder
+        self.t5_encoder = t5_encoder
+        self.clip_encoder = clip_encoder
+        self.train_pipeline = train_pipeline
+        self.model_config = model_config
+        self.trainable_modules = trainable_modules
+        self.config = FluxConfigShim(raw=raw_config or {})
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def freeze_except(self):
+        mode = (
+            self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit"
+        ).lower()
+
+        def freeze(module: nn.Module):
+            for param in module.parameters():
+                param.requires_grad_(False)
+
+        def unfreeze(module: nn.Module):
+            for param in module.parameters():
+                param.requires_grad_(True)
+
+        freeze(self)
+        if mode in ("dit", "diffusion", "backbone"):
+            unfreeze(self.dit)
+        elif mode == "all":
+            unfreeze(self)
+        else:
+            unfreeze(self.dit)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        if hasattr(self.dit, "gradient_checkpointing"):
+            self.dit.gradient_checkpointing = True
+
+    def forward(self, *args, **kwargs):
+        if len(args) >= 1 and isinstance(args[0], dict):
+            scheduler = args[1] if len(args) >= 2 else kwargs.get("scheduler")
+            return self.forward_train(args[0], scheduler=scheduler)
+        raise TypeError("FluxForTraining.forward expects (batch_dict, scheduler)")
+
+    def forward_train(self, batch: dict[str, Any], scheduler: Any = None) -> dict[str, torch.Tensor]:
+        del scheduler
+        return self.train_pipeline.compute_loss(
+            dit=self.dit,
+            autoencoder=self.autoencoder,
+            t5_encoder=self.t5_encoder,
+            clip_encoder=self.clip_encoder,
+            batch=batch,
+            model_config=self.model_config,
+        )
+
+    def forward_inference(self, batch: dict[str, Any], **kwargs):
+        raise NotImplementedError("FLUX inference pipeline is not wired for Primus training yet")
diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py
new file mode 100644
index 000000000..9ff6f3e56
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/autoencoder.py
@@ -0,0 +1,315 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Adapted from Black Forest Labs FLUX official implementation.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from safetensors.torch import load_file as load_sft
+from torch import Tensor, nn
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int = 256
+    in_channels: int = 3
+    ch: int = 128
+    out_ch: int = 3
+    ch_mult: list[int] | None = None
+    num_res_blocks: int = 2
+    z_channels: int = 16
+    scale_factor: float = 0.3611
+    shift_factor: float = 0.1159
+
+    def __post_init__(self):
+        if self.ch_mult is None:
+            self.ch_mult = [1, 2, 4, 4]
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        bsz, channels, height, width = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=height, w=width, c=channels, b=bsz)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(nn.functional.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        return self.conv(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        upscale_dtype = next(self.up.parameters()).dtype
+        h = self.conv_in(z)
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h).to(upscale_dtype)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if not self.sample:
+            return mean
+        std = torch.exp(0.5 * logvar)
+        return mean + std * torch.randn_like(mean)
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams, sample_z: bool = True):
+        super().__init__()
+        self.params = params
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult or [1, 2, 4, 4],
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult or [1, 2, 4, 4],
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian(sample=sample_z)
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        return self.scale_factor * (z - self.shift_factor)
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))
+
+
+def load_autoencoder(
+    ckpt_path: str | None,
+    params: AutoEncoderParams | None = None,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    sample_z: bool = True,
+) -> AutoEncoder:
+    ae = AutoEncoder(params or AutoEncoderParams(), sample_z=sample_z)
+    if ckpt_path:
+        ckpt_path = _resolve_checkpoint_path(ckpt_path, default_filename="ae.safetensors")
+        state = load_sft(ckpt_path, device="cpu")
+        missing, unexpected = ae.load_state_dict(state, strict=False)
+        if missing:
+            raise ValueError(f"FLUX autoencoder checkpoint missing {len(missing)} keys; first={missing[:3]}")
+        if unexpected:
+            raise ValueError(
+                f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}"
+            )
+    return ae.to(dtype=dtype)
+
+
+def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) -> str:
+    path = Path(path_or_repo_file).expanduser()
+    if path.exists():
+        return str(path)
+    if path_or_repo_file.startswith(("/", "./", "../", "~")):
+        raise FileNotFoundError(f"FLUX checkpoint path not found: {path}")
+
+    parts = path_or_repo_file.split("/")
+    if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")):
+        raise FileNotFoundError(f"FLUX checkpoint path not found: {path}")
+    if len(parts) >= 3:
+        repo_id = "/".join(parts[:2])
+        filename = "/".join(parts[2:])
+    elif len(parts) == 2:
+        repo_id = path_or_repo_file
+        filename = default_filename
+    else:
+        raise FileNotFoundError(
+            f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}"
+        )
+
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=repo_id, filename=filename)
diff --git a/primus/backends/diffusion/models/flux/conditioner.py b/primus/backends/diffusion/models/flux/conditioner.py
new file mode 100644
index 000000000..736bf225d
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/conditioner.py
@@ -0,0 +1,54 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Adapted from Black Forest Labs FLUX official implementation.
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+
+
+class HFEmbedder(nn.Module):
+    def __init__(self, version: str, max_length: int, **hf_kwargs):
+        super().__init__()
+        self.is_clip = "clip" in version.lower() or version.startswith("openai")
+        self.max_length = max_length
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+
+        if self.is_clip:
+            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
+        else:
+            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+        for param in self.hf_module.parameters():
+            if not param.is_contiguous():
+                param.data = param.data.contiguous()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.hf_module.parameters()).device
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(self.device),
+            attention_mask=batch_encoding["attention_mask"].to(self.device),
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key]
diff --git a/primus/backends/diffusion/models/flux/configuration_flux.py b/primus/backends/diffusion/models/flux/configuration_flux.py
new file mode 100644
index 000000000..6808fcac7
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/configuration_flux.py
@@ -0,0 +1,14 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+
+class FluxTrainingConfig(SimpleNamespace):
+    def to_dict(self) -> dict:
+        return dict(self.__dict__)
diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py
new file mode 100644
index 000000000..475be057b
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/layers.py
@@ -0,0 +1,211 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Adapted from Black Forest Labs FLUX official implementation.
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from primus.backends.diffusion.models.flux.math import attention, rope
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat([rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> Tensor:
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        return self.proj(x)
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+        return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__()
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        img_modulated = (1 + img_mod1.scale) * self.img_norm1(img) + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        txt_modulated = (1 + txt_mod1.scale) * self.txt_norm1(txt) + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+        attn = attention(q, k, v, pe=pe)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+    def __init__(
+        self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+        self.norm = QKNorm(head_dim)
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        attn = attention(q, k, v, pe=pe)
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        return self.linear(x)
diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py
new file mode 100644
index 000000000..829d2f395
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/math.py
@@ -0,0 +1,44 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Adapted from Black Forest Labs FLUX official implementation.
+
+from __future__ import annotations
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+from primus.backends.diffusion.attention import attention as backend_attention
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+    x = backend_attention(
+        q=rearrange(q, "B H L D -> B L H D"),
+        k=rearrange(k, "B H L D -> B L H D"),
+        v=rearrange(v, "B H L D -> B L H D"),
+        dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16,
+    )
+    return rearrange(x, "B L H D -> B L (H D)")
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    if dim % 2 != 0:
+        raise ValueError(f"RoPE dimension must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py
new file mode 100644
index 000000000..ffa4fa108
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/model.py
@@ -0,0 +1,167 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Adapted from Black Forest Labs FLUX official implementation.
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+
+import torch
+from torch import Tensor, nn
+
+from primus.backends.diffusion.models.flux.layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    LastLayer,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    out_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def flux_1_dev_params(**overrides) -> FluxParams:
+    values = {
+        "in_channels": 64,
+        "out_channels": 64,
+        "vec_in_dim": 768,
+        "context_in_dim": 4096,
+        "hidden_size": 3072,
+        "mlp_ratio": 4.0,
+        "num_heads": 24,
+        "depth": 19,
+        "depth_single_blocks": 38,
+        "axes_dim": [16, 56, 56],
+        "theta": 10000,
+        "qkv_bias": True,
+        "guidance_embed": True,
+    }
+    values.update(overrides)
+    return FluxParams(**values)
+
+
+def flux_1_schnell_params(**overrides) -> FluxParams:
+    values = flux_1_dev_params(guidance_embed=False).to_dict()
+    values.update(overrides)
+    return FluxParams(**values)
+
+
+class Flux(nn.Module):
+    """Transformer model for flow matching on packed latent sequences."""
+
+    def __init__(self, params: FluxParams):
+        super().__init__()
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got axes_dim={params.axes_dim}, expected sum={pe_dim}")
+
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.gradient_checkpointing = False
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def _checkpoint_double(self, block: nn.Module, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor):
+        import torch.utils.checkpoint as checkpoint_utils
+
+        return checkpoint_utils.checkpoint(block, img, txt, vec, pe, use_reentrant=False)
+
+    def _checkpoint_single(self, block: nn.Module, img: Tensor, vec: Tensor, pe: Tensor):
+        import torch.utils.checkpoint as checkpoint_utils
+
+        return checkpoint_utils.checkpoint(block, img, vec, pe, use_reentrant=False)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("FLUX guidance-distilled models require a guidance tensor.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        use_checkpoint = self.training and self.gradient_checkpointing
+        for block in self.double_blocks:
+            if use_checkpoint:
+                img, txt = self._checkpoint_double(block, img, txt, vec, pe)
+            else:
+                img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+        img = torch.cat((txt, img), 1)
+        for block in self.single_blocks:
+            if use_checkpoint:
+                img = self._checkpoint_single(block, img, vec, pe)
+            else:
+                img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+        return self.final_layer(img, vec)
diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py
new file mode 100644
index 000000000..98ac2d077
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/train_pipeline.py
@@ -0,0 +1,184 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from primus.backends.diffusion.models.flux.utils import (
+    create_position_encoding_for_latents,
+    generate_latent_from_mean_logvar,
+    pack_latents,
+)
+
+
+@dataclass
+class FluxFlowMatchTrainPipelineConfig:
+    autoencoder_scale_factor: float = 0.3611
+    autoencoder_shift_factor: float = 0.1159
+    guidance: float | None = None
+
+
+class FluxFlowMatchTrainPipeline:
+    """Flow-matching loss for FLUX using precomputed or online encodings."""
+
+    def __init__(self, cfg: FluxFlowMatchTrainPipelineConfig | None = None):
+        self.cfg = cfg or FluxFlowMatchTrainPipelineConfig()
+
+    @staticmethod
+    def _require_module(module: torch.nn.Module | None, name: str) -> torch.nn.Module:
+        if module is None:
+            raise ValueError(f"FLUX raw image-text training requires `{name}` to be configured.")
+        return module
+
+    @staticmethod
+    def _align_module_dtype(module: torch.nn.Module, *, device: torch.device, dtype: torch.dtype) -> None:
+        """Move a frozen encoder to the target device/dtype only when it differs.
+
+        `nn.Module.to(dtype=...)` casts floating-point params/buffers only, leaving
+        integer buffers (e.g. token position ids) untouched. We inspect the first
+        floating-point parameter to avoid re-casting on every step.
+        """
+        current = next((p for p in module.parameters() if p.is_floating_point()), None)
+        if current is None:
+            module.to(device=device)
+            return
+        if current.dtype != dtype or current.device != device:
+            module.to(device=device, dtype=dtype)
+
+    def _prepare_precomputed(
+        self,
+        *,
+        batch: dict[str, Any],
+        device: torch.device,
+        dtype: torch.dtype,
+        model_config: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        t5_encodings = batch["t5_encodings"].to(device=device, dtype=dtype, non_blocking=True)
+        clip_encodings = batch["clip_encodings"].to(device=device, dtype=dtype, non_blocking=True)
+        mean = batch["mean"].to(device=device, dtype=dtype, non_blocking=True)
+        logvar = batch["logvar"].to(device=device, dtype=dtype, non_blocking=True)
+
+        latents = generate_latent_from_mean_logvar(mean, logvar)
+        scale = float(getattr(model_config, "autoencoder_scale_factor", self.cfg.autoencoder_scale_factor))
+        shift = float(getattr(model_config, "autoencoder_shift_factor", self.cfg.autoencoder_shift_factor))
+        labels = (latents - shift) * scale
+        return labels, t5_encodings, clip_encodings
+
+    def _prepare_raw(
+        self,
+        *,
+        batch: dict[str, Any],
+        autoencoder: torch.nn.Module | None,
+        t5_encoder: torch.nn.Module | None,
+        clip_encoder: torch.nn.Module | None,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        ae = self._require_module(autoencoder, "model.config.encoder.autoencoder")
+        t5 = self._require_module(t5_encoder, "model.config.encoder.t5_encoder")
+        clip = self._require_module(clip_encoder, "model.config.encoder.clip_encoder")
+        image = batch["image"].to(device=device, dtype=dtype, non_blocking=True)
+        prompts = batch.get("prompts")
+        if not isinstance(prompts, list):
+            raise ValueError("FLUX raw batch requires `prompts` as list[str].")
+
+        # The frozen encoders are built in bf16 regardless of the training compute
+        # dtype; align them with the DiT dtype so running them on `image`/inputs of
+        # `dtype` does not raise a dtype-mismatch error (e.g. for fp32 runs).
+        self._align_module_dtype(ae, device=device, dtype=dtype)
+        self._align_module_dtype(t5, device=device, dtype=dtype)
+        self._align_module_dtype(clip, device=device, dtype=dtype)
+        ae.eval()
+        t5.eval()
+        clip.eval()
+        with torch.no_grad():
+            labels = ae.encode(image).to(device=device, dtype=dtype)
+            t5_encodings = t5(prompts).to(device=device, dtype=dtype)
+            clip_encodings = clip(prompts).to(device=device, dtype=dtype)
+        return labels, t5_encodings, clip_encodings
+
+    def compute_loss(
+        self,
+        *,
+        dit: torch.nn.Module,
+        batch: dict[str, Any],
+        model_config: Any,
+        autoencoder: torch.nn.Module | None = None,
+        t5_encoder: torch.nn.Module | None = None,
+        clip_encoder: torch.nn.Module | None = None,
+    ) -> dict[str, torch.Tensor]:
+        if batch.get("sp_group") is not None:
+            raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.")
+
+        device = next(dit.parameters()).device
+        dtype = next(dit.parameters()).dtype
+        if "image" in batch:
+            labels, t5_encodings, clip_encodings = self._prepare_raw(
+                batch=batch,
+                autoencoder=autoencoder,
+                t5_encoder=t5_encoder,
+                clip_encoder=clip_encoder,
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            required = ("t5_encodings", "clip_encodings", "mean", "logvar")
+            missing = [key for key in required if key not in batch]
+            if missing:
+                raise ValueError(f"FLUX precomputed batch is missing required keys: {missing}")
+            labels, t5_encodings, clip_encodings = self._prepare_precomputed(
+                batch=batch,
+                device=device,
+                dtype=dtype,
+                model_config=model_config,
+            )
+
+        bsz = labels.shape[0]
+        noise = torch.randn_like(labels)
+        timesteps = torch.rand((bsz,), device=device, dtype=dtype)
+        sigmas = timesteps.view(-1, 1, 1, 1)
+        noisy_latents = (1 - sigmas) * labels + sigmas * noise
+        target = noise - labels
+
+        _, _, latent_height, latent_width = noisy_latents.shape
+        # Position ids are integer grid indices consumed by RoPE; build them in
+        # float32 (independent of the model compute dtype) so that indices remain
+        # exactly representable. bf16 only represents integers up to 256 exactly,
+        # which would silently corrupt positions for larger latent grids.
+        img_ids = create_position_encoding_for_latents(
+            bsz,
+            latent_height,
+            latent_width,
+            position_dim=3,
+            device=device,
+            dtype=torch.float32,
+        )
+        txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=torch.float32)
+        noisy_latents = pack_latents(noisy_latents)
+        target = pack_latents(target)
+
+        guidance_value = getattr(model_config, "guidance", self.cfg.guidance)
+        guidance = (
+            None
+            if guidance_value is None
+            else torch.full((bsz,), float(guidance_value), device=device, dtype=dtype)
+        )
+        pred = dit(
+            img=noisy_latents,
+            img_ids=img_ids,
+            txt=t5_encodings,
+            txt_ids=txt_ids,
+            y=clip_encodings,
+            timesteps=timesteps,
+            guidance=guidance,
+        )
+        loss = F.mse_loss(pred.float(), target.float(), reduction="sum") / target.numel()
+        return {"loss": loss, "log_metrics": {"mse": loss.detach()}}
diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py
new file mode 100644
index 000000000..00d0a2acf
--- /dev/null
+++ b/primus/backends/diffusion/models/flux/utils.py
@@ -0,0 +1,49 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+PATCH_HEIGHT = 2
+PATCH_WIDTH = 2
+LATENT_CHANNELS = 16
+IMAGE_LATENT_SIZE_RATIO = 8
+
+
+def generate_latent_from_mean_logvar(mean: Tensor, logvar: Tensor) -> Tensor:
+    return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean)
+
+
+def create_position_encoding_for_latents(
+    bsz: int,
+    latent_height: int,
+    latent_width: int,
+    position_dim: int = 3,
+    *,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+) -> Tensor:
+    height = latent_height // PATCH_HEIGHT
+    width = latent_width // PATCH_WIDTH
+    position_encoding = torch.zeros(height, width, position_dim, device=device, dtype=dtype)
+    position_encoding[:, :, 1] = torch.arange(height, device=device, dtype=dtype).unsqueeze(1)
+    position_encoding[:, :, 2] = torch.arange(width, device=device, dtype=dtype).unsqueeze(0)
+    return position_encoding.view(1, height * width, position_dim).repeat(bsz, 1, 1)
+
+
+def pack_latents(x: Tensor) -> Tensor:
+    bsz, channels, latent_height, latent_width = x.shape
+    if latent_height % PATCH_HEIGHT != 0 or latent_width % PATCH_WIDTH != 0:
+        raise ValueError(
+            "FLUX latents must have height and width divisible by 2, " f"got shape={tuple(x.shape)}"
+        )
+    height = latent_height // PATCH_HEIGHT
+    width = latent_width // PATCH_WIDTH
+    x = x.unfold(2, PATCH_HEIGHT, PATCH_HEIGHT).unfold(3, PATCH_WIDTH, PATCH_WIDTH)
+    x = x.permute(0, 2, 3, 1, 4, 5)
+    return x.reshape(bsz, height * width, channels * PATCH_HEIGHT * PATCH_WIDTH)
diff --git a/primus/backends/diffusion/models/interface.py b/primus/backends/diffusion/models/interface.py
new file mode 100644
index 000000000..1355c3658
--- /dev/null
+++ b/primus/backends/diffusion/models/interface.py
@@ -0,0 +1,48 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+import torch
+import torch.nn as nn
+
+
+class GenAIModel(nn.Module, ABC):
+    """
+    Unified interface for Generative Models (DiT, AR, Hybrid).
+    Wraps the Backbone (DiT/Transformer), VAE, and TextEncoder.
+    """
+
+    @abstractmethod
+    def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]:
+        """
+        The single entry point for training.
+
+        Responsibilities:
+        1. Process raw batch (Tokenization, VAE Encoding if needed).
+        2. Apply Training Recipe (e.g., Add Noise for DiT, Shift Tokens for AR).
+        3. Forward pass through Backbone.
+        4. Calculate Loss.
+
+        Args:
+            batch: Raw batch from data loader.
+            scheduler: Optional scheduler (e.g., FlowMatchScheduler) passed from Trainer.
+
+        Returns:
+            {
+                "loss": torch.Tensor,       # Main optimization objective
+                "log_metrics": Dict,        # Metrics for WandB (MSE, Accuracy, etc.)
+            }
+        """
+
+    @abstractmethod
+    def forward_inference(self, batch: Dict[str, Any], **kwargs):
+        """
+        Entry point for validation/inference.
+        For DiT: Runs the diffusion sampling loop.
+        For AR: Runs autoregressive generation.
+        """
diff --git a/primus/backends/diffusion/models/registrations/__init__.py b/primus/backends/diffusion/models/registrations/__init__.py
new file mode 100644
index 000000000..98d014572
--- /dev/null
+++ b/primus/backends/diffusion/models/registrations/__init__.py
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Model registrations."""
diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py
new file mode 100644
index 000000000..93a57be60
--- /dev/null
+++ b/primus/backends/diffusion/models/registrations/flux.py
@@ -0,0 +1,234 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import glob
+import os
+from typing import Any
+
+import torch
+from safetensors.torch import load_file as safe_load_file
+
+from primus.backends.diffusion.models.flux.adapter import FluxForTraining
+from primus.backends.diffusion.models.flux.autoencoder import (
+    AutoEncoderParams,
+    load_autoencoder,
+)
+from primus.backends.diffusion.models.flux.conditioner import HFEmbedder
+from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig
+from primus.backends.diffusion.models.flux.model import (
+    Flux,
+    flux_1_dev_params,
+    flux_1_schnell_params,
+)
+from primus.backends.diffusion.models.flux.train_pipeline import (
+    FluxFlowMatchTrainPipeline,
+    FluxFlowMatchTrainPipelineConfig,
+)
+from primus.backends.diffusion.utils.log import logger
+from primus.backends.diffusion.utils.train_utils import count_parameters
+
+_FLUX_PRESET_ALIASES = {
+    "flux-schnell": "flux-schnell",
+    "flux.1-schnell": "flux-schnell",
+    "flux1-schnell": "flux-schnell",
+    "flux-dev": "flux-dev",
+    "flux.1-dev": "flux-dev",
+    "flux1-dev": "flux-dev",
+}
+
+
+def _strip_known_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    prefixes = ("module.", "dit.", "model.")
+    out: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        stripped = key
+        changed = True
+        while changed:
+            changed = False
+            for prefix in prefixes:
+                if stripped.startswith(prefix):
+                    stripped = stripped[len(prefix) :]
+                    changed = True
+        out[stripped] = value
+    return out
+
+
+def _load_state_dict(path: str) -> dict[str, torch.Tensor]:
+    if path.endswith(".safetensors"):
+        return dict(safe_load_file(path))
+    obj = torch.load(path, map_location="cpu")
+    if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict):
+        obj = obj["model"]
+    if not isinstance(obj, dict):
+        raise ValueError(f"Unsupported checkpoint format at {path}")
+    return obj
+
+
+def _candidate_weight_files(path: str, *, default_filename: str) -> list[str]:
+    if os.path.isfile(path):
+        return [path]
+    if not os.path.exists(path):
+        resolved = _resolve_hf_checkpoint(path, default_filename=default_filename)
+        if resolved:
+            return [resolved]
+    candidates: list[str] = []
+    for fname in (
+        "flux1-schnell.safetensors",
+        "flux1-dev.safetensors",
+        "dit_model.safetensors",
+        "model.safetensors",
+    ):
+        candidate = os.path.join(path, fname)
+        if os.path.exists(candidate):
+            candidates.append(candidate)
+    if not candidates:
+        candidates = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+    if not candidates:
+        candidates = sorted(glob.glob(os.path.join(path, "*.bin")))
+    return candidates
+
+
+def _resolve_hf_checkpoint(path_or_repo_file: str, *, default_filename: str) -> str | None:
+    if path_or_repo_file.startswith(("/", "./", "../", "~")):
+        return None
+    parts = path_or_repo_file.split("/")
+    if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")):
+        return None
+    if len(parts) < 2:
+        return None
+    if len(parts) >= 3:
+        repo_id = "/".join(parts[:2])
+        filename = "/".join(parts[2:])
+    else:
+        repo_id = path_or_repo_file
+        filename = default_filename
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=repo_id, filename=filename)
+
+
+def _load_flux_weights(dit: torch.nn.Module, pretrained_path: str, *, default_filename: str) -> None:
+    candidates = _candidate_weight_files(pretrained_path, default_filename=default_filename)
+    if not candidates:
+        raise FileNotFoundError(f"No FLUX DiT weights found under {pretrained_path}")
+
+    merged: dict[str, torch.Tensor] = {}
+    for ckpt in candidates:
+        merged.update(_strip_known_prefixes(_load_state_dict(ckpt)))
+
+    result = dit.load_state_dict(merged, strict=False)
+    logger.info(
+        "Loaded FLUX DiT weights. "
+        f"files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}"
+    )
+
+
+def _build_flux_dit(params) -> Flux:
+    local_rank = os.environ.get("LOCAL_RANK")
+    use_cuda = local_rank is not None and torch.cuda.is_available()
+    device = torch.device(f"cuda:{local_rank}") if use_cuda else torch.device("cpu")
+    old_dtype = torch.get_default_dtype()
+    try:
+        if use_cuda:
+            torch.set_default_dtype(torch.bfloat16)
+        with torch.device(device):
+            dit = Flux(params)
+    finally:
+        torch.set_default_dtype(old_dtype)
+    return dit
+
+
+def build_flux_model(model_config: dict[str, Any]):
+    """
+    Build a FLUX model from the selected model preset.
+
+    `model_preset` is injected by the registry from `model.name` for Primus
+    configs such as `flux.1-dev` and `flux.1-schnell`.
+    """
+    cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {})
+    preset_name = str(model_config.get("model_preset") or cfg_dict.get("model_preset") or "flux.1-schnell")
+    preset = _FLUX_PRESET_ALIASES.get(preset_name.lower(), preset_name)
+
+    params_overrides = dict(cfg_dict.get("params", {}) or {})
+    if preset == "flux-dev":
+        params = flux_1_dev_params(**params_overrides)
+    elif preset == "flux-schnell":
+        params = flux_1_schnell_params(**params_overrides)
+    else:
+        raise ValueError(
+            "Unsupported FLUX model_preset="
+            f"{preset_name!r}; expected one of: 'flux.1-dev', 'flux.1-schnell'"
+        )
+    dit = _build_flux_dit(params)
+
+    pretrained_path = model_config.get("load_from_pretrained_path") or model_config.get("pretrained_path")
+    if pretrained_path:
+        logger.info(f"Loading FLUX DiT weights from {pretrained_path}")
+        default_filename = "flux1-dev.safetensors" if preset == "flux-dev" else "flux1-schnell.safetensors"
+        _load_flux_weights(dit, pretrained_path, default_filename=default_filename)
+
+    encoder_cfg = dict(model_config.get("encoder", {}) or cfg_dict.get("encoder", {}) or {})
+    dtype = torch.bfloat16
+    t5_encoder = None
+    clip_encoder = None
+    autoencoder = None
+    if encoder_cfg.get("t5_encoder"):
+        t5_encoder = HFEmbedder(
+            str(encoder_cfg["t5_encoder"]),
+            max_length=int(encoder_cfg.get("max_t5_length", 256)),
+            torch_dtype=dtype,
+        )
+    if encoder_cfg.get("clip_encoder"):
+        clip_encoder = HFEmbedder(
+            str(encoder_cfg["clip_encoder"]),
+            max_length=int(encoder_cfg.get("max_clip_length", 77)),
+            torch_dtype=dtype,
+        )
+    if encoder_cfg.get("autoencoder"):
+        ae_params = AutoEncoderParams(
+            resolution=int(encoder_cfg.get("resolution", 256)),
+            scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)),
+            shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)),
+        )
+        autoencoder = load_autoencoder(
+            str(encoder_cfg["autoencoder"]),
+            ae_params,
+            dtype=dtype,
+            sample_z=bool(encoder_cfg.get("sample_z", True)),
+        )
+
+    training_cfg = FluxTrainingConfig(
+        model_preset=preset,
+        trainable_modules=cfg_dict.get("trainable_modules", "dit"),
+        guidance=None if not params.guidance_embed else float(cfg_dict.get("guidance", 1.0)),
+        autoencoder_scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)),
+        autoencoder_shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)),
+    )
+    pipeline = FluxFlowMatchTrainPipeline(
+        FluxFlowMatchTrainPipelineConfig(
+            autoencoder_scale_factor=training_cfg.autoencoder_scale_factor,
+            autoencoder_shift_factor=training_cfg.autoencoder_shift_factor,
+            guidance=training_cfg.guidance,
+        )
+    )
+    model = FluxForTraining(
+        dit=dit,
+        train_pipeline=pipeline,
+        model_config=training_cfg,
+        autoencoder=autoencoder,
+        t5_encoder=t5_encoder,
+        clip_encoder=clip_encoder,
+        raw_config={
+            "model_config": model_config,
+            "flux_params": params.to_dict(),
+        },
+        trainable_modules=training_cfg.trainable_modules,
+    )
+    total_params, trainable_params = count_parameters(model)
+    logger.info(f"Built FLUX model: total={total_params:,} trainable={trainable_params:,}")
+    return model
diff --git a/primus/backends/diffusion/models/registrations/wan.py b/primus/backends/diffusion/models/registrations/wan.py
new file mode 100644
index 000000000..491284ff7
--- /dev/null
+++ b/primus/backends/diffusion/models/registrations/wan.py
@@ -0,0 +1,226 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Register the Wan model builder."""
+
+from __future__ import annotations
+
+import glob
+import os
+from typing import Any
+
+import torch
+from safetensors.torch import load_file as safe_load_file
+
+from primus.backends.diffusion.models.wan.adapter import WanForTraining
+from primus.backends.diffusion.models.wan.components import WanComponents
+from primus.backends.diffusion.models.wan.configuration_wanvideo import WanVideoConfig
+from primus.backends.diffusion.models.wan.t5 import umt5_xxl_encoder_from_checkpoint
+from primus.backends.diffusion.models.wan.train_pipeline import (
+    WanFlowMatchTrainPipeline,
+)
+from primus.backends.diffusion.models.wan.vae2_1 import Wan2_1_VAE
+from primus.backends.diffusion.models.wan.vae2_2 import Wan2_2_VAE
+from primus.backends.diffusion.models.wan.wan_dit import WanModel as WanDiT
+from primus.backends.diffusion.utils.log import logger
+from primus.backends.diffusion.utils.train_utils import count_parameters
+
+
+def _strip_module_prefix(
+    state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    out: dict[str, torch.Tensor] = {}
+    for k, v in state_dict.items():
+        if k.startswith("module."):
+            out[k[len("module.") :]] = v
+        else:
+            out[k] = v
+    return out
+
+
+def _load_state_dict(path: str) -> dict[str, torch.Tensor]:
+    if path.endswith(".safetensors"):
+        return dict(safe_load_file(path))
+    obj = torch.load(path, map_location="cpu")
+    if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict):
+        obj = obj["model"]
+    if not isinstance(obj, dict):
+        raise ValueError(f"Unsupported checkpoint format at {path}")
+    return obj
+
+
+def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str):
+    """
+    Support two common cases:
+    1) Wan training export: `dit_model.safetensors` with keys like `blocks.0...`
+    2) Official-ish export: `*model*.safetensors|bin` possibly with keys like `dit.blocks.0...` or `blocks.0...`
+    """
+    # Direct file
+    if os.path.isfile(pretrained_path):
+        state = _strip_module_prefix(_load_state_dict(pretrained_path))
+        # Accept either `dit.*` or plain keys.
+        if any(k.startswith("dit.") for k in state):
+            state = {k[len("dit.") :]: v for k, v in state.items() if k.startswith("dit.")}
+        result = dit.load_state_dict(state, strict=False)
+        logger.info(
+            f"Loaded DiT from file. missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}"
+        )
+        return
+
+    # Directory: prefer explicit Wan trainer file name, else fall back to pattern search
+    candidates: list[str] = []
+    for fname in (
+        "dit_model.safetensors",
+        "diffusion_pytorch_model.safetensors",
+        "model.safetensors",
+    ):
+        p = os.path.join(pretrained_path, fname)
+        if os.path.exists(p):
+            candidates.append(p)
+    if not candidates:
+        candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.safetensors")))
+    if not candidates:
+        candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.bin")))
+    if not candidates:
+        raise FileNotFoundError(f"No DiT weights found under {pretrained_path}")
+
+    merged: dict[str, torch.Tensor] = {}
+    for ckpt in candidates:
+        part = _strip_module_prefix(_load_state_dict(ckpt))
+        merged.update(part)
+
+    if any(k.startswith("dit.") for k in merged):
+        merged = {k[len("dit.") :]: v for k, v in merged.items() if k.startswith("dit.")}
+
+    result = dit.load_state_dict(merged, strict=False)
+    logger.info(
+        f"Loaded DiT from dir. files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}"
+    )
+
+
+def build_wan_model(model_config: dict):
+    """
+    YAML compatibility:
+      model_config:
+        name: wan
+        load_from_pretrained_path: /path/to/Wan2.2-TI2V-5B   (optional)
+        config: {...}                                       (optional overrides)
+        encoder:
+          t5_encoder: ...
+          autoencoder: ...
+    """
+    encoder_cfg = model_config.get("encoder", {}) if isinstance(model_config.get("encoder"), dict) else {}
+    cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {})
+    if encoder_cfg:
+        cfg_dict.setdefault("encoder", encoder_cfg)
+
+    # 1. Try to load config.json from pretrained path if available
+    pretrained_path = model_config.get("load_from_pretrained_path")
+    if pretrained_path:
+        import json
+
+        cfg_path = os.path.join(pretrained_path, "config.json")
+        if os.path.exists(cfg_path):
+            logger.info(f"Loading config from {cfg_path}")
+            with open(cfg_path) as f:
+                loaded_cfg = json.load(f)
+
+            # Map WanModel checkpoint config keys to WanVideoConfig keys.
+            mapping = {
+                "dim": "dit_hidden_size",
+                "num_layers": "dit_num_layers",
+                "num_heads": "dit_num_heads",
+                "ffn_dim": "dit_intermediate_size",
+                "in_dim": "dit_in_channels",
+                "out_dim": "dit_out_channels",
+                "freq_dim": "dit_freq_dim",
+                "text_len": "text_len",
+            }
+
+            for k, v in loaded_cfg.items():
+                target_key = mapping.get(k)
+                # Only infer values the user did not set explicitly in YAML.
+                if target_key is not None and target_key not in cfg_dict:
+                    cfg_dict[target_key] = v
+
+    # Build a WanVideoConfig-compatible object to maximize reuse of existing fields.
+    model_cfg = WanVideoConfig(**cfg_dict)
+
+    # Build DiT (close to official Wan2.2 modules/model.py signature).
+    # model_config.config.model_type: one of {"t2v","i2v","ti2v","s2v"}
+    dit_task_type = getattr(model_cfg, "model_type", None)
+    if dit_task_type not in ("t2v", "i2v", "ti2v", "s2v"):
+        raise ValueError(
+            "wan requires explicit `model_config.config.model_type` in YAML "
+            "(one of: t2v / i2v / ti2v / s2v). "
+            f"Got: {dit_task_type!r}"
+        )
+    dit = WanDiT(
+        model_type=dit_task_type,
+        patch_size=tuple(model_cfg.dit_patch_size),
+        text_len=int(getattr(model_cfg, "text_len", 512)),
+        in_dim=int(model_cfg.dit_in_channels),
+        dim=int(model_cfg.dit_hidden_size),
+        ffn_dim=int(model_cfg.dit_intermediate_size),
+        freq_dim=int(model_cfg.dit_freq_dim),
+        text_dim=int(model_cfg.dit_text_dim),
+        out_dim=int(model_cfg.dit_out_channels),
+        num_heads=int(model_cfg.dit_num_heads),
+        num_layers=int(model_cfg.dit_num_layers),
+        window_size=tuple(getattr(model_cfg, "dit_window_size", (-1, -1))),
+        qk_norm=bool(getattr(model_cfg, "dit_qk_norm", True)),
+        cross_attn_norm=bool(getattr(model_cfg, "dit_cross_attn_norm", True)),
+        eps=float(model_cfg.dit_eps),
+    )
+
+    # Build VAE
+    vae_type = getattr(model_cfg, "vae_type", "wan_video_vae_38")
+    vae_ckpt = encoder_cfg.get("vae_checkpoint") or encoder_cfg.get("autoencoder")
+    if vae_type in ("wan2.2", "wan_video_vae_38"):
+        if not vae_ckpt:
+            raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.2 VAE checkpoint path)")
+        vae = Wan2_2_VAE(z_dim=48, vae_pth=vae_ckpt)
+        vae.upsampling_factor = 16
+    else:
+        if not vae_ckpt:
+            raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.1 VAE checkpoint path)")
+        vae = Wan2_1_VAE(z_dim=16, vae_pth=vae_ckpt)
+        vae.upsampling_factor = 8
+
+    # Build text encoder (UMT5-XXL, encoder-only). Tokenization is handled by the dataset processor.
+    t5_ckpt = encoder_cfg.get("t5_encoder")
+    if not t5_ckpt:
+        raise ValueError("wan requires `model_config.encoder.t5_encoder` (UMT5 encoder checkpoint path)")
+    text_encoder = umt5_xxl_encoder_from_checkpoint(t5_ckpt, dtype=torch.bfloat16, device="cpu")
+
+    # Optionally load pretrained DiT
+    pretrained_path = model_config.get("load_from_pretrained_path")
+    if pretrained_path:
+        logger.info(f"Loading DiT weights from {pretrained_path}")
+        _load_dit_weights_into_module(dit, pretrained_path)
+
+    components = WanComponents(dit=dit, vae=vae, text_encoder=text_encoder, image_encoder=None)
+    pipeline = WanFlowMatchTrainPipeline()
+
+    model = WanForTraining(
+        components=components,
+        train_pipeline=pipeline,
+        model_config=model_cfg,
+        raw_config={
+            "model_config": model_config,
+            "resolved_model_cfg": getattr(model_cfg, "__dict__", cfg_dict),
+        },
+        trainable_modules=getattr(model_cfg, "trainable_modules", None),
+    )
+
+    total_params, trainable_params = count_parameters(model)
+    logger.info(f"wan parameters: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B")
+    if hasattr(model, "freeze_except"):
+        model.freeze_except()
+        total_params, trainable_params = count_parameters(model)
+        logger.info(f"wan after freeze: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B")
+
+    return model
diff --git a/primus/backends/diffusion/models/wan/__init__.py b/primus/backends/diffusion/models/wan/__init__.py
new file mode 100644
index 000000000..86c01a615
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/__init__.py
@@ -0,0 +1,20 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+wan: Components + TaskPipeline design for Wan PyTorch/FSDP training.
+
+Goal:
+- Keep modeling close to official Wan2.2 (pure torch modules).
+- Keep trainer generic: trainer calls model(batch, scheduler) -> {"loss": ...}.
+- Decouple training/inference workflow (pipeline) from modeling (modules).
+"""
+
+from .adapter import WanForTraining
+
+__all__ = [
+    "WanForTraining",
+]
diff --git a/primus/backends/diffusion/models/wan/adapter.py b/primus/backends/diffusion/models/wan/adapter.py
new file mode 100644
index 000000000..1d453c5f9
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/adapter.py
@@ -0,0 +1,154 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+
+from primus.backends.diffusion.models.interface import GenAIModel
+
+from .components import WanComponents
+from .train_pipeline import WanFlowMatchTrainPipeline
+
+
+@dataclass
+class WanConfigShim:
+    """
+    A tiny config shim to keep trainer saving behavior happy.
+    We intentionally do not implement a full HF/Diffusers config system.
+    """
+
+    raw: dict
+
+    def save_pretrained(self, save_directory: str):
+        # Best-effort: keep minimal JSON for reproducibility.
+        import json
+        import os
+
+        os.makedirs(save_directory, exist_ok=True)
+        path = os.path.join(save_directory, "wan_config.json")
+        with open(path, "w") as f:
+            json.dump(self.raw, f, indent=2, sort_keys=True)
+
+    def to_dict(self):
+        return self.raw
+
+
+class WanForTraining(GenAIModel, nn.Module):
+    """
+    A thin adapter that exposes the call pattern expected by Wan trainers.
+
+    Trainers call:
+      outputs = model(batch, scheduler)  -> {"loss": ...}
+
+    Internally we delegate workflow to the pipeline.
+    """
+
+    def __init__(
+        self,
+        *,
+        components: WanComponents,
+        train_pipeline: WanFlowMatchTrainPipeline,
+        model_config: Any,
+        raw_config: Optional[dict] = None,
+        trainable_modules: Optional[str] = None,
+    ):
+        super().__init__()
+        self.components = components
+        self.train_pipeline = train_pipeline
+        self.model_config = model_config
+        self.trainable_modules = trainable_modules
+
+        # Expose common attribute names expected by trainers/FSDP ignore regex.
+        self.dit = components.dit
+        self.vae = components.vae
+        self.text_encoder = components.text_encoder
+        self.image_encoder = components.image_encoder
+
+        self.config = WanConfigShim(raw=raw_config or {})
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def to(self, *args, **kwargs):
+        """
+        Override to propagate .to() to non-nn.Module components.
+
+        Wan VAE wrappers (Wan2_1_VAE, Wan2_2_VAE) are plain Python classes
+        with a custom .to() method, not nn.Module subclasses.  PyTorch's
+        nn.Module.to() only recurses into registered submodules, so the VAE
+        would be silently skipped — causing device mismatches on multi-GPU.
+        """
+        result = super().to(*args, **kwargs)
+        for component in (self.vae, self.image_encoder):
+            if component is not None and not isinstance(component, nn.Module) and hasattr(component, "to"):
+                component.to(*args, **kwargs)
+        return result
+
+    def freeze_except(self):
+        """
+        Keep the Wan training behavior: freeze non-trainable modules.
+        Default: train only DiT.
+        """
+        mode = (
+            self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit"
+        ).lower()
+
+        def freeze(m: nn.Module):
+            for p in m.parameters():
+                p.requires_grad_(False)
+
+        def unfreeze(m: nn.Module):
+            for p in m.parameters():
+                p.requires_grad_(True)
+
+        # Freeze everything first
+        freeze(self)
+
+        # Unfreeze requested parts
+        if mode in ("dit", "diffusion", "backbone"):
+            unfreeze(self.dit)
+        elif mode in ("all",):
+            unfreeze(self)
+        else:
+            # Conservative default: DiT only
+            unfreeze(self.dit)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        """
+        Activated by HF Trainer when `gradient_checkpointing=True`.
+        """
+        if self.dit and hasattr(self.dit, "gradient_checkpointing"):
+            self.dit.gradient_checkpointing = True
+
+    def forward(self, *args, **kwargs):
+        # Match existing trainer call convention: model(batch, scheduler)
+        if len(args) >= 1 and isinstance(args[0], dict):
+            scheduler = args[1] if len(args) > 1 else kwargs.get("scheduler", None)
+            return self.forward_train(args[0], scheduler=scheduler)
+        raise TypeError("WanForTraining.forward expects (batch_dict, scheduler)")
+
+    def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]:
+        if scheduler is None:
+            raise ValueError("scheduler must be provided by trainer")
+        return self.train_pipeline.compute_loss(
+            components=self.components,
+            batch=batch,
+            scheduler=scheduler,
+            model_config=self.model_config,
+        )
+
+    def forward_inference(self, batch: Dict[str, Any], **kwargs):
+        raise NotImplementedError("wan inference pipeline not wired yet")
diff --git a/primus/backends/diffusion/models/wan/attention_backend.py b/primus/backends/diffusion/models/wan/attention_backend.py
new file mode 100644
index 000000000..306f8afe4
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/attention_backend.py
@@ -0,0 +1,228 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+#
+# This file is vendored from `Wan2.2/wan/modules/attention.py` with minimal changes:
+# - keep PyTorch implementation
+# - allow CPU fallback when flash-attn is unavailable or tensors are on CPU
+
+import warnings
+
+import torch
+
+try:
+    import flash_attn_interface
+
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+__all__ = [
+    "flash_attention",
+    "attention",
+]
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # Fast path (fixed-length, no padding mask):
+    # Prefer flash_attn_func over flash_attn_varlen_func to reduce overhead.
+    if q_lens is None and k_lens is None:
+        qh = half(q)
+        kh = half(k)
+        vh = half(v)
+        qh = qh.to(vh.dtype)
+        kh = kh.to(vh.dtype)
+        if q_scale is not None:
+            qh = qh * q_scale
+
+        # apply attention
+        if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+            x = flash_attn_interface.flash_attn_func(qh, kh, vh)
+            if isinstance(x, tuple):
+                x = x[0]
+        else:
+            assert FLASH_ATTN_2_AVAILABLE
+            # flash_attn_func signature differs across versions; be conservative
+            try:
+                x = flash_attn.flash_attn_func(
+                    qh,
+                    kh,
+                    vh,
+                    dropout_p=dropout_p,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                    window_size=window_size,
+                )
+            except TypeError:
+                x = flash_attn.flash_attn_func(qh, kh, vh)
+        return x.type(out_dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.")
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        )[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    # Minimal change vs upstream: allow CPU fallback.
+    if (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE) and q.device.type == "cuda":
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+
+    if q_lens is not None or k_lens is not None:
+        warnings.warn(
+            "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
+        )
+    attn_mask = None
+
+    out_dtype = q.dtype
+    q = q.transpose(1, 2).to(dtype)
+    k = k.transpose(1, 2).to(dtype)
+    v = v.transpose(1, 2).to(dtype)
+
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+    )
+
+    out = out.transpose(1, 2).contiguous().to(out_dtype)
+    return out
diff --git a/primus/backends/diffusion/models/wan/components.py b/primus/backends/diffusion/models/wan/components.py
new file mode 100644
index 000000000..3e5a854d2
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/components.py
@@ -0,0 +1,26 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch.nn as nn
+
+
+@dataclass
+class WanComponents:
+    """
+    A thin container for model components.
+
+    Keep this intentionally simple: it's just a bundle of modules that a pipeline can use.
+    """
+
+    dit: nn.Module
+    vae: nn.Module
+    text_encoder: nn.Module
+    image_encoder: Optional[nn.Module] = None
diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py
new file mode 100644
index 000000000..0df88d0ec
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py
@@ -0,0 +1,84 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# coding=utf-8
+# Copyright 2024 WanVideo team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict
+
+
+class WanVideoConfig:
+    model_type = "wanvideo"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        dit_hidden_size: int = 3072,
+        dit_num_layers: int = 30,
+        dit_num_heads: int = 24,
+        dit_intermediate_size: int = 14336,
+        dit_patch_size: tuple = (1, 2, 2),
+        dit_in_channels: int = 48,
+        dit_out_channels: int = 48,
+        dit_freq_dim: int = 256,
+        dit_text_dim: int = 4096,
+        dit_eps: float = 1e-6,
+        dit_has_image_input: bool = False,
+        dit_has_image_pos_emb: bool = False,
+        dit_has_ref_conv: bool = False,
+        trainable_modules=None,
+        separated_timestep: bool = True,
+        require_clip_embedding: bool = False,
+        require_vae_embedding: bool = False,
+        fuse_vae_embedding_in_latents: bool = True,
+        tie_word_embeddings: bool = False,
+        vae_type: str = "wan_video_vae_38",
+        **kwargs,
+    ):
+        # DiT configuration
+        self.vae_type = vae_type
+        self.dit_hidden_size = dit_hidden_size
+        self.dit_num_layers = dit_num_layers
+        self.dit_num_heads = dit_num_heads
+        self.dit_intermediate_size = dit_intermediate_size
+        self.dit_patch_size = dit_patch_size
+        self.dit_in_channels = dit_in_channels
+        self.dit_out_channels = dit_out_channels
+        self.dit_freq_dim = dit_freq_dim
+        self.dit_text_dim = dit_text_dim
+        self.dit_eps = dit_eps
+        self.dit_has_image_input = dit_has_image_input
+
+        self.dit_has_image_pos_emb = dit_has_image_pos_emb
+        self.dit_has_ref_conv = dit_has_ref_conv
+
+        self.separated_timestep = bool(separated_timestep)
+        self.require_clip_embedding = require_clip_embedding
+        self.require_vae_embedding = require_vae_embedding
+        self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
+
+        self.trainable_modules = trainable_modules
+
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert configuration to dictionary"""
+        output = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+        output["model_type"] = self.model_type
+        return output
diff --git a/primus/backends/diffusion/models/wan/t5.py b/primus/backends/diffusion/models/wan/t5.py
new file mode 100644
index 000000000..d9cdfab71
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/t5.py
@@ -0,0 +1,299 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+#
+# Notes for Wan training:
+# - We keep this file focused on the PyTorch encoder only.
+# - We do NOT vendor the tokenizer wrapper (ftfy/regex deps). Training uses
+#   `input_ids/attention_mask` produced by the existing processor instead.
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = [
+    "T5Encoder",
+    "umt5_xxl_encoder",
+    "umt5_xxl_encoder_from_checkpoint",
+]
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        if self.weight.dtype in (torch.float16, torch.bfloat16):
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C]
+        context:    [B, L2, C] or None
+        mask:       [B, L2] or [B, L1, L2] or None
+        """
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in (2, 3)
+            mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # T5 does not use scaling
+        attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum("bnij,bjnc->binc", attn, v)
+
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super().__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        max_exact = num_buckets // 2
+        rel_pos_large = (
+            max_exact
+            + (
+                torch.log(rel_pos.float() / max_exact)
+                / math.log(self.max_dist / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5SelfAttention(nn.Module):
+    def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos=True, dropout=0.1):
+        super().__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = (
+            None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True)
+        )
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn) ** -0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn) ** -0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(m.embedding.weight, std=(2 * m.num_buckets * m.num_heads) ** -0.5)
+
+
+class T5Encoder(nn.Module):
+    def __init__(
+        self, vocab, dim, dim_attn, dim_ffn, num_heads, num_layers, num_buckets, shared_pos=True, dropout=0.1
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) else nn.Embedding(vocab, dim)
+        self.pos_embedding = (
+            T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList(
+            [
+                T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = T5LayerNorm(dim)
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+def umt5_xxl_encoder(**kwargs) -> T5Encoder:
+    """
+    Official UMT5-XXL encoder-only config.
+    """
+    cfg = dict(
+        vocab=256384,
+        dim=4096,
+        dim_attn=4096,
+        dim_ffn=10240,
+        num_heads=64,
+        num_layers=24,
+        num_buckets=32,
+        shared_pos=False,
+        dropout=0.1,
+    )
+    cfg.update(**kwargs)
+    model = T5Encoder(**cfg)
+    model.eval().requires_grad_(False)
+    return model
+
+
+def umt5_xxl_encoder_from_checkpoint(
+    checkpoint_path: str,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str | torch.device = "cpu",
+) -> T5Encoder:
+    """
+    Build UMT5-XXL encoder-only model and load weights without allocating
+    full parameter tensors upfront.
+
+    This mirrors the official pattern of meta-init + `assign=True` loading.
+    """
+    if not checkpoint_path:
+        raise ValueError("checkpoint_path is required")
+
+    # Build on meta to avoid allocating a ~1B parameter embedding upfront.
+    with torch.device("meta"):
+        model = umt5_xxl_encoder()
+
+    state = torch.load(checkpoint_path, map_location="cpu")
+    if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict):
+        state = state["model"]
+    if not isinstance(state, dict):
+        raise ValueError(f"Unsupported checkpoint format at {checkpoint_path}")
+
+    # Strip potential DDP prefix.
+    if any(k.startswith("module.") for k in state.keys()):
+        state = {k[len("module.") :]: v for k, v in state.items()}
+
+    try:
+        model.load_state_dict(state, strict=False, assign=True)
+    except TypeError as exc:
+        raise RuntimeError(
+            "Your PyTorch build does not support `assign=True` for loading into meta-initialized modules. "
+            "Please upgrade torch (>=2.0) or adjust loading strategy."
+        ) from exc
+
+    model = model.to(device=device, dtype=dtype).eval().requires_grad_(False)
+    return model
diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py
new file mode 100644
index 000000000..db755638e
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/train_pipeline.py
@@ -0,0 +1,290 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn.functional as F
+
+from .components import WanComponents
+
+
+@dataclass
+class WanFlowMatchTrainPipelineConfig:
+    """
+    Minimal training-pipeline config.
+
+    We intentionally keep this small and derive most behavior from the model config
+    to maximize YAML reuse.
+    """
+
+    # For Wan VAEs, temporal downsample is typically (4, 1): 4n+1 frames.
+    time_division_factor: int = 4
+    time_division_remainder: int = 1
+
+
+class WanFlowMatchTrainPipeline:
+    """
+    Flow-Matching training pipeline for Wan-style DiT models.
+
+    Contract:
+      compute_loss(components, batch, scheduler) -> {"loss": Tensor, ...}
+    """
+
+    def __init__(self, cfg: Optional[WanFlowMatchTrainPipelineConfig] = None):
+        self.cfg = cfg or WanFlowMatchTrainPipelineConfig()
+
+    @staticmethod
+    def _encode_prompt(
+        text_encoder: torch.nn.Module,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ):
+        # Match existing wan_new behavior: zero-out padded embeddings explicitly.
+        seq_lens = attention_mask.gt(0).sum(dim=1).long()
+        prompt_emb = text_encoder(input_ids, attention_mask)
+        for i, v in enumerate(seq_lens):
+            prompt_emb[i, v:] = 0
+        return prompt_emb
+
+    @staticmethod
+    def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]:
+        seed = batch.get("seed", None)
+        if seed is None:
+            return None
+        if isinstance(seed, torch.Tensor):
+            flat_seed = seed.detach().reshape(-1).cpu()
+            if flat_seed.numel() == 0:
+                return None
+            if flat_seed.numel() > 1 and not bool((flat_seed == flat_seed[0]).all()):
+                raise ValueError(
+                    "Per-sample dataset seeds are not supported; got different seeds in one batch."
+                )
+            seed = flat_seed[0].item()
+        elif isinstance(seed, (list, tuple)):
+            if not seed:
+                return None
+            if len(seed) > 1 and any(value != seed[0] for value in seed):
+                raise ValueError(
+                    "Per-sample dataset seeds are not supported; got different seeds in one batch."
+                )
+            seed = seed[0]
+        return int(seed)
+
+    @staticmethod
+    def _randn_like_on_cpu_then_to(
+        x: torch.Tensor,
+        *,
+        seed: Optional[int],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device="cpu")
+            generator.manual_seed(seed)
+        noise = torch.randn(x.shape, generator=generator, device="cpu", dtype=torch.float32)
+        return noise.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _vae_encode(vae: torch.nn.Module, videos_bcthw: torch.Tensor) -> torch.Tensor:
+        # Wan VAEs in this repo use list-of-tensors interface: List[[C,T,H,W]] -> List[[C,F,H',W']]
+        videos_list = [videos_bcthw[i] for i in range(videos_bcthw.shape[0])]
+        latents_list = vae.encode(videos_list)
+        return torch.stack(latents_list)
+
+    @staticmethod
+    def _get_dit_patch_size(model_config: Any) -> tuple[int, int, int]:
+        patch_size = tuple(getattr(model_config, "dit_patch_size", (1, 2, 2)))
+        if len(patch_size) != 3:
+            raise ValueError(f"Expected 3D dit_patch_size, got {patch_size!r}")
+        return patch_size
+
+    @staticmethod
+    def _pad_latents_for_dit(
+        latents: torch.Tensor, *, patch_size: tuple[int, int, int]
+    ) -> tuple[torch.Tensor, tuple[int, int]]:
+        _, _, _, height, width = latents.shape
+        _, patch_h, patch_w = patch_size
+        pad_h = (-height) % patch_h
+        pad_w = (-width) % patch_w
+        if pad_h == 0 and pad_w == 0:
+            return latents, (height, width)
+
+        # Pad only the latent-space bottom/right edges so DiT patchify works on any
+        # VAE output shape. We crop predictions back before computing loss.
+        latents = F.pad(latents, (0, pad_w, 0, pad_h))
+        return latents, (height, width)
+
+    @staticmethod
+    def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> torch.Tensor:
+        height, width = spatial_size
+        return latents[..., :height, :width]
+
+    @staticmethod
+    def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor:
+        """Sample one training timestep uniformly."""
+        timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device)
+        # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor.
+        timestep = scheduler.timesteps[timestep_id.cpu()].float()
+        return timestep.to(device=device)
+
+    @staticmethod
+    def _maybe_expand_separated_timestep(
+        *,
+        timestep: torch.Tensor,
+        x_list: list[torch.Tensor],
+        patch_size: tuple[int, int, int],
+        enabled: bool,
+    ) -> torch.Tensor:
+        """
+        Match `wan_new.forward_dit` separated-timestep behavior (only used when enabled).
+        For each sample, build per-token timestep [L] and set first-frame patches to 0.
+        Returns:
+          - t: [B] if not enabled
+          - t: [B, L] if enabled
+        """
+        if not enabled:
+            if timestep.ndim == 0:
+                return timestep.unsqueeze(0).repeat(len(x_list))
+            if timestep.ndim == 1 and timestep.numel() == 1 and len(x_list) > 1:
+                return timestep.repeat(len(x_list))
+            if timestep.ndim == 1 and timestep.shape[0] != len(x_list):
+                return timestep.repeat(len(x_list))
+            return timestep
+
+        d_f, d_h, d_w = patch_size
+        if timestep.ndim == 0:
+            timestep_b = timestep.unsqueeze(0).repeat(len(x_list))
+        elif timestep.ndim == 1 and timestep.shape[0] != len(x_list):
+            timestep_b = timestep.repeat(len(x_list))
+        else:
+            timestep_b = timestep
+
+        t_expand_list: list[torch.Tensor] = []
+        for i, x in enumerate(x_list):
+            # x: [C, F, H, W] in latent space
+            f, h, w = x.shape[1], x.shape[2], x.shape[3]
+            seq_len = (f // d_f) * (h // d_h) * (w // d_w)
+            t_seq = torch.full(
+                (seq_len,),
+                timestep_b[i],
+                device=timestep_b.device,
+                dtype=timestep_b.dtype,
+            )
+            spatial_patches = (h // d_h) * (w // d_w)
+            t_seq[:spatial_patches] = 0
+            t_expand_list.append(t_seq)
+        return torch.stack(t_expand_list)
+
+    def compute_loss(
+        self,
+        *,
+        components: WanComponents,
+        batch: Dict[str, Any],
+        scheduler: Any,
+        model_config: Any,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        batch requirements (from current dataset/processor):
+          - video: Tensor [B, C, T, H, W]
+          - input_ids: Tensor [B, L]
+          - attention_mask: Tensor [B, L]
+        """
+        video = batch.get("video")
+        if video is None:
+            raise ValueError("Batch must contain 'video'")
+        if not isinstance(video, torch.Tensor) or video.ndim != 5:
+            raise ValueError(
+                f"Expected batch['video'] as 5D tensor [B,C,T,H,W], got {type(video)} shape={getattr(video, 'shape', None)}"
+            )
+
+        input_ids = batch.get("input_ids")
+        attention_mask = batch.get("attention_mask")
+        if input_ids is None or attention_mask is None:
+            raise ValueError("Batch must contain 'input_ids' and 'attention_mask'")
+
+        device = next(components.dit.parameters()).device
+        dtype = next(components.dit.parameters()).dtype
+
+        video = video.to(device=device, dtype=dtype, non_blocking=True)
+        input_ids = input_ids.to(device=device, non_blocking=True)
+        attention_mask = attention_mask.to(device=device, non_blocking=True)
+
+        # 1) Encode to latents
+        # Keep VAE/text encoder in eval() by default (consistent with wan_new practice).
+        components.text_encoder.eval()
+        try:
+            components.vae.eval()
+        except AttributeError:
+            # Some VAE compatibility wrappers may not expose eval(); real eval
+            # failures should still surface.
+            pass
+
+        with torch.no_grad():
+            input_latents = self._vae_encode(components.vae, video).to(device=device, dtype=dtype)
+        patch_size = self._get_dit_patch_size(model_config)
+        input_latents, original_latent_spatial_size = self._pad_latents_for_dit(
+            input_latents, patch_size=patch_size
+        )
+
+        # 2) Sample timestep + noise (match `wan_new`: noise on CPU, timestep from scheduler.timesteps)
+        timestep = self._select_timestep(scheduler, device=device)  # [1] (float)
+        seed = self._get_seed_from_env_or_batch(batch)
+        noise = self._randn_like_on_cpu_then_to(input_latents, seed=seed, dtype=dtype, device=device)
+
+        # 3) Diffusion target + noise injection (match `wan_new`)
+        # Note: in this repo's FlowMatchScheduler, `training_target(sample, noise, t)` uses `noise - sample`.
+        # `wan_new` passes `input_latents` (not noisy latents).
+        target = scheduler.training_target(input_latents, noise, timestep)
+        noisy_latents = scheduler.add_noise(input_latents, noise, timestep=timestep)
+
+        # 4) Text embeddings
+        with torch.no_grad():
+            context = self._encode_prompt(components.text_encoder, input_ids, attention_mask)
+
+        # 5) DiT forward (official interface: List[Tensor] per-sample)
+        x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])]
+        context_list = [context[i] for i in range(context.shape[0])]
+
+        # seq_len matches DiT Conv3d patchification over [F, H, W].
+        d_f, d_h, d_w = patch_size
+        max_seq_len = 0
+        for x in x_list:
+            seq_len = (x.shape[1] // d_f) * (x.shape[2] // d_h) * (x.shape[3] // d_w)
+            max_seq_len = max(max_seq_len, seq_len)
+
+        separated = bool(getattr(model_config, "separated_timestep", False))
+        fuse_flag = bool(getattr(model_config, "fuse_vae_embedding_in_latents", False))
+        t = self._maybe_expand_separated_timestep(
+            timestep=timestep,
+            x_list=x_list,
+            patch_size=(d_f, d_h, d_w),
+            enabled=bool(separated and fuse_flag),
+        )
+
+        sp_group = batch.get("sp_group", None)
+        noise_pred_list = components.dit(
+            x=x_list,
+            t=t,
+            context=context_list,
+            seq_len=max_seq_len,
+            y=None,
+            sp_group=sp_group,
+        )
+        noise_pred = torch.stack(noise_pred_list)
+        noise_pred = self._crop_latents(noise_pred, spatial_size=original_latent_spatial_size)
+        target = self._crop_latents(target, spatial_size=original_latent_spatial_size)
+
+        # 6) Loss
+        loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+        # `wan_new` uses scalar `timestep` for weighting (not expanded per-token)
+        weight = scheduler.training_weight(timestep)
+        loss = loss * weight
+        return {"loss": loss}
diff --git a/primus/backends/diffusion/models/wan/vae2_1.py b/primus/backends/diffusion/models/wan/vae2_1.py
new file mode 100644
index 000000000..ac6b44760
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/vae2_1.py
@@ -0,0 +1,646 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Wan2.1 VAE (vendored from official Wan2.2 repo).
+
+Source: `Wan2.2/wan/modules/vae2_1.py`
+
+Minimal changes:
+- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage
+  (HF trainer sometimes calls `model.vae.to(dtype=...)`).
+"""
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+from typing import List
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    "Wan2_1_VAE",
+    "WanVAE",
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+    def __init__(self, dim, mode):
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        cache_x = torch.cat(
+                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1, c2)
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        )
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+            if i != len(dim_mult) - 1:
+                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 3, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        self.encoder = Encoder3d(
+            dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
+        )
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx
+                )
+                out = torch.cat([out, out_], 2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device="cpu", **kwargs):
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+    return model
+
+
+class Wan2_1_VAE:
+    def __init__(self, z_dim=16, vae_pth="cache/vae_step_411000.pth", dtype=torch.float, device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        mean = [
+            -0.7571,
+            -0.7089,
+            -0.9113,
+            0.1075,
+            -0.1745,
+            0.9653,
+            -0.1517,
+            1.5508,
+            0.4134,
+            -0.0715,
+            0.5517,
+            -0.3632,
+            -0.1922,
+            -0.9497,
+            0.2503,
+            -0.2921,
+        ]
+        std = [
+            2.8184,
+            1.4541,
+            2.3275,
+            2.6558,
+            1.2196,
+            1.7708,
+            2.6052,
+            2.0743,
+            3.2687,
+            2.1526,
+            2.8652,
+            1.5579,
+            1.6382,
+            1.1253,
+            2.8251,
+            1.9160,
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        if vae_pth is None:
+            raise ValueError("Wan2_1_VAE requires `vae_pth` (checkpoint path).")
+        self.model = _video_vae(pretrained_path=vae_pth, z_dim=z_dim).eval().requires_grad_(False).to(device)
+
+    # --- Wan trainer compatibility helpers (non-upstream) ---
+    def to(self, *args, **kwargs):
+        self.model.to(*args, **kwargs)
+        if "dtype" in kwargs and kwargs["dtype"] is not None:
+            self.dtype = kwargs["dtype"]
+        device = kwargs.get("device", None)
+        if device is None and len(args) > 0:
+            device = args[0]
+        if device is not None:
+            self.device = device
+            self.mean = self.mean.to(device)
+            self.std = self.std.to(device)
+            self.scale = [self.mean, 1.0 / self.std]
+        return self
+
+    def eval(self):
+        self.model.eval()
+        return self
+
+    def train(self, mode: bool = True):
+        self.model.train(mode)
+        return self
+
+    def encode(self, videos: List[torch.Tensor]):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+        with amp.autocast(dtype=self.dtype):
+            # Avoid forcing fp32 outputs: it increases memory and slows training.
+            return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos]
+
+    def decode(self, zs: List[torch.Tensor]):
+        with amp.autocast(dtype=self.dtype):
+            return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs]
+
+
+# Backward-compat alias (older Wan naming)
+WanVAE = Wan2_1_VAE
diff --git a/primus/backends/diffusion/models/wan/vae2_2.py b/primus/backends/diffusion/models/wan/vae2_2.py
new file mode 100644
index 000000000..05998db45
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/vae2_2.py
@@ -0,0 +1,873 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Wan2.2 VAE (vendored from official Wan2.2 repo).
+
+Source: `Wan2.2/wan/modules/vae2_2.py`
+
+Minimal changes:
+- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage
+  (HF trainer sometimes calls `model.vae.to(dtype=...)`).
+"""
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+from typing import List
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    "Wan2_2_VAE",
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+    def __init__(self, dim, mode):
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        cache_x = torch.cat(
+                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1, c2)
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        )
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size)
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size)
+    return x
+
+
+class AvgDown3D(nn.Module):
+    def __init__(self, in_channels, out_channels, factor_t, factor_s=1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        b, c, t, h, w = x.shape
+        x = x.view(
+            b,
+            c,
+            t // self.factor_t,
+            self.factor_t,
+            h // self.factor_s,
+            self.factor_s,
+            w // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(b, c * self.factor, t // self.factor_t, h // self.factor_s, w // self.factor_s)
+        x = x.view(
+            b, self.out_channels, self.group_size, t // self.factor_t, h // self.factor_s, w // self.factor_s
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1 :, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False):
+        super().__init__()
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False):
+        super().__init__()
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False):
+        if feat_idx is None:
+            feat_idx = [0]
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk)
+            return x_main + x_shortcut
+        return x_main
+
+
+class Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        dims = [dim * u for u in [1] + dim_mult]
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = temperal_downsample[i] if i < len(temperal_downsample) else False
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                )
+            )
+        self.downsamples = nn.Sequential(*downsamples)
+
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=None):
+        if feat_idx is None:
+            feat_idx = [0]
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                )
+            )
+        self.upsamples = nn.Sequential(*upsamples)
+
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False):
+        if feat_idx is None:
+            feat_idx = [0]
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+    # NOTE: Keep this identical to upstream for numerical parity.
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        self.encoder = Encoder3d(
+            dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
+        )
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        x = patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx
+                )
+                out = torch.cat([out, out_], 2)
+        out = unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+    return model
+
+
+class Wan2_2_VAE:
+    def __init__(
+        self,
+        z_dim=48,
+        c_dim=160,
+        vae_pth=None,
+        dim_mult=[1, 2, 4, 4],
+        temperal_downsample=[False, True, True],
+        dtype=torch.bfloat16,
+        device="cuda",
+    ):
+        self.dtype = dtype
+        self.device = device
+        mean = torch.tensor(
+            [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        std = torch.tensor(
+            [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        self.scale = [mean, 1.0 / std]
+
+        if vae_pth is None:
+            raise ValueError("Wan2_2_VAE requires `vae_pth` (checkpoint path).")
+        self.model = (
+            _video_vae(
+                pretrained_path=vae_pth,
+                z_dim=z_dim,
+                dim=c_dim,
+                dim_mult=dim_mult,
+                temperal_downsample=temperal_downsample,
+            )
+            .eval()
+            .requires_grad_(False)
+            .to(device)
+        )
+
+    # --- Wan trainer compatibility helpers (non-upstream) ---
+    def to(self, *args, **kwargs):
+        dtype = kwargs.get("dtype", None)
+        if dtype is None:
+            dtype = self.dtype
+        else:
+            self.dtype = dtype
+        device = kwargs.get("device", None)
+        if device is None and len(args) > 0:
+            device = args[0]
+        if device is not None:
+            self.model.to(device=device, dtype=dtype)
+        else:
+            self.model.to(*args, **kwargs)
+        if device is not None:
+            self.device = device
+            dev = next(self.model.parameters()).device
+            self.scale = [t.to(device=dev, dtype=dtype) for t in self.scale]
+        return self
+
+    def eval(self):
+        self.model.eval()
+        return self
+
+    def train(self, mode: bool = True):
+        self.model.train(mode)
+        return self
+
+    def encode(self, videos: List[torch.Tensor]):
+        try:
+            if not isinstance(videos, list):
+                raise TypeError("videos should be a list")
+            with amp.autocast(dtype=self.dtype):
+                # Avoid forcing fp32 outputs: it increases memory and slows training.
+                # Keep outputs in the model/autocast dtype (typically bf16).
+                return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos]
+        except TypeError as e:
+            logging.info(e)
+            return None
+
+    def decode(self, zs: List[torch.Tensor]):
+        try:
+            if not isinstance(zs, list):
+                raise TypeError("zs should be a list")
+            with amp.autocast(dtype=self.dtype):
+                return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs]
+        except TypeError as e:
+            logging.info(e)
+            return None
diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py
new file mode 100644
index 000000000..0d975b250
--- /dev/null
+++ b/primus/backends/diffusion/models/wan/wan_dit.py
@@ -0,0 +1,580 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+#
+# Vendored from `Wan2.2/wan/modules/model.py` with minimal changes:
+# - remove diffusers dependencies (ModelMixin/ConfigMixin/register_to_config)
+# - keep pure PyTorch modeling as close as possible
+# - provide a `DiTBlock` alias class so existing Wan configs can keep
+#   `fsdp_transformer_layer_cls_to_wrap: "DiTBlock"`
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+from primus.backends.diffusion.attention import attention
+from primus.backends.diffusion.distributed.ulysses import (
+    distributed_attention,
+    sp_gather,
+    sp_split,
+    sp_unpad,
+)
+
+__all__ = ["WanModel", "DiTBlock"]
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)),
+    )
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def rope_apply(x, grid_sizes, freqs, sp_group=None):
+    """
+    Apply 3-D RoPE.  When *sp_group* is given the input is assumed to hold
+    only this rank's local token chunk (S/P) and the correct positional
+    frequencies are selected automatically (following Wan2.2 official).
+    """
+    n, c = x.size(2), x.size(3) // 2
+    s = x.size(1)  # local token count (= full seq when no SP)
+
+    if sp_group is not None:
+        sp_size = dist.get_world_size(sp_group)
+        sp_rank = dist.get_rank(sp_group)
+    else:
+        sp_size, sp_rank = 1, 0
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # tokens to process: local chunk when SP, valid tokens when single-GPU
+        t = s if sp_size > 1 else seq_len
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :t].to(torch.float64).reshape(t, n, -1, 2))
+        freqs_i = torch.cat(
+            [
+                freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(seq_len, 1, -1)
+
+        # SP: pad freqs to padded-full-length, select this rank's range
+        if sp_size > 1:
+            full_len = s * sp_size
+            if seq_len < full_len:
+                freqs_i = torch.cat(
+                    [
+                        freqs_i,
+                        torch.ones(
+                            full_len - seq_len,
+                            1,
+                            freqs_i.size(-1),
+                            dtype=freqs_i.dtype,
+                            device=freqs_i.device,
+                        ),
+                    ],
+                    dim=0,
+                )
+            freqs_i = freqs_i[sp_rank * s : (sp_rank + 1) * s]
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, t:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).to(x.dtype)
+
+
+class WanRMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        # Keep numerical stability by normalizing in fp32, but do NOT rely on
+        # module parameters being fp32 (FSDP/bf16 can cast weights/bias).
+        weight = self.weight.float() if self.weight is not None else None
+        bias = self.bias.float() if self.bias is not None else None
+        y = F.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps)
+        return y.type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs, sp_group: Optional[dist.ProcessGroup] = None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C] (L = S/P when SP enabled)
+            seq_lens(Tensor): Shape [B] — valid lengths in the *full* sequence
+            grid_sizes(Tensor): Shape [B, 3], (F, H, W) — full spatial grid
+            freqs(Tensor): Rope freqs
+            sp_group: Ulysses SP process group (None = no SP)
+        """
+        b, s = x.shape[:2]
+        n = self.num_heads
+        d = self.head_dim
+
+        x_ = x.to(dtype=self.q.weight.dtype)
+        q = self.norm_q(self.q(x_)).view(b, s, n, d)
+        k = self.norm_k(self.k(x_)).view(b, s, n, d)
+        v = self.v(x_).view(b, s, n, d)
+
+        # RoPE on local tokens (SP-aware: uses rank-offset freqs when sp_group is set)
+        q = rope_apply(q, grid_sizes, freqs, sp_group=sp_group)
+        k = rope_apply(k, grid_sizes, freqs, sp_group=sp_group)
+
+        attn_kwargs = {
+            "q_lens": seq_lens,
+            "k_lens": seq_lens,
+            "window_size": self.window_size,
+            "dtype": self.q.weight.dtype,
+        }
+        if sp_group is not None:
+            x = distributed_attention(q, k, v, group=sp_group, attention_fn=attention, **attn_kwargs)
+        else:
+            x = attention(q=q, k=k, v=v, **attn_kwargs)
+
+        x = x.to(dtype=self.o.weight.dtype)
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanCrossAttention(WanSelfAttention):
+    def forward(
+        self,
+        x,
+        seq_lens=None,
+        grid_sizes=None,
+        freqs=None,
+        sp_group: Optional[dist.ProcessGroup] = None,
+        context=None,
+        context_lens=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            seq_lens/grid_sizes/freqs/sp_group: accepted for WanSelfAttention
+                signature compatibility; cross-attention uses text context only.
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        if context is None:
+            # Backward-compatible positional form: forward(x, context, context_lens).
+            context = seq_lens
+            context_lens = grid_sizes if context_lens is None else context_lens
+
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        x = x.to(dtype=self.q.weight.dtype)
+        context = context.to(dtype=self.k.weight.dtype)
+
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        x = attention(q, k, v, k_lens=context_lens, dtype=self.q.weight.dtype)
+
+        x = x.to(dtype=self.o.weight.dtype)
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        ffn_dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=False,
+        eps=1e-6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)
+        )
+
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+        sp_group: Optional[dist.ProcessGroup] = None,
+    ):
+        # Memory-critical:
+        # `e` can be very large ([B, seq_len, 6, dim]). Keeping it in fp32 and
+        # doing broadcast add in fp32 easily OOMs for Wan2.1 (more tokens).
+        # Align to x.dtype (bf16/fp16) like `wan_new` implementation.
+        e = e.to(dtype=x.dtype)
+        modulation = self.modulation.to(dtype=x.dtype)
+        e = (modulation.unsqueeze(0) + e).chunk(6, dim=2)
+
+        # self-attention (Ulysses all-to-all happens inside self_attn)
+        y = self.self_attn(
+            self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2),
+            seq_lens,
+            grid_sizes,
+            freqs,
+            sp_group=sp_group,
+        )
+        x = x + y * e[2].squeeze(2)
+
+        # cross-attention + FFN (no SP communication needed:
+        # each rank's visual queries attend to full text keys independently)
+        def cross_attn_ffn(x_, context_, context_lens_, e_):
+            x_ = x_ + self.cross_attn(
+                self.norm3(x_),
+                context=context_,
+                context_lens=context_lens_,
+            )
+            ffn_in = self.norm2(x_) * (1 + e_[4].squeeze(2)) + e_[3].squeeze(2)
+            y_ = self.ffn(ffn_in)
+            x_ = x_ + y_ * e_[5].squeeze(2)
+            return x_
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+# Wan config compatibility: FSDP auto-wrap often uses "DiTBlock".
+class DiTBlock(WanAttentionBlock):
+    pass
+
+
+class Head(nn.Module):
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        # Align modulation embedding dtype to x (saves memory vs fp32 broadcast).
+        e = e.to(dtype=x.dtype)
+        modulation = self.modulation.to(dtype=x.dtype)
+        e = (modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2)
+        x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))
+        return x
+
+
+class WanModel(nn.Module):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    def __init__(
+        self,
+        model_type="t2v",
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=16,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=16,
+        num_heads=16,
+        num_layers=32,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        eps=1e-6,
+    ):
+        super().__init__()
+
+        assert model_type in ["t2v", "i2v", "ti2v", "s2v"]
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        # Used by the Wan trainer: it sets `model.dit.gradient_checkpointing = True`
+        # when `trainer_args.gradient_checkpointing: true`.
+        self.gradient_checkpointing = False
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)
+        )
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks (use DiTBlock for YAML compatibility)
+        self.blocks = nn.ModuleList(
+            [
+                DiTBlock(dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps)
+                for _ in range(num_layers)
+            ]
+        )
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        # rope buffers (avoid register_buffer to keep dtype stable under .to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.freqs = torch.cat(
+            [
+                rope_params(1024, d - 4 * (d // 6)),
+                rope_params(1024, 2 * (d // 6)),
+                rope_params(1024, 2 * (d // 6)),
+            ],
+            dim=1,
+        )
+
+        self.init_weights()
+
+    def forward(self, x, t, context, seq_len, y=None, sp_group: Optional[dist.ProcessGroup] = None):
+        if self.model_type == "i2v":
+            assert y is not None
+
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x])
+
+        # time embeddings
+        if t.dim() == 1:
+            # [B] -> [B, seq_len]: per-sample timestep broadcast across the sequence.
+            # Must unsqueeze before expand, otherwise the batch dim (B) is wrongly
+            # aligned with seq_len and breaks for batch_size > 1.
+            t = t.unsqueeze(1).expand(t.size(0), seq_len)
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            bt = t.size(0)
+            t_flat = t.flatten()
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t_flat).unflatten(0, (bt, seq_len)).float()
+            )
+            e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # Cast large modulation tensors back to model dtype early to reduce peak memory.
+        e = e.to(dtype=x.dtype)
+        e0 = e0.to(dtype=x.dtype)
+
+        # --- Ulysses SP: split sequence-dim tensors across SP ranks ---
+        original_seq_len = None
+        if sp_group is not None:
+            (x, e, e0), original_seq_len = sp_split([x, e, e0], dim=1, group=sp_group)
+
+        # context (text embeddings — NOT sliced, full text on every SP rank)
+        context_lens = None
+        context_in = torch.stack(
+            [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]
+        )
+        context_in = context_in.to(dtype=self.text_embedding[0].weight.dtype)
+        context = self.text_embedding(context_in)
+
+        freqs = self.freqs
+        seq_lens_ = seq_lens
+        grid_sizes_ = grid_sizes
+        e0_ = e0
+        context_ = context
+        context_lens_ = context_lens
+
+        for block in self.blocks:
+            if self.training and getattr(self, "gradient_checkpointing", False):
+
+                def create_custom_forward(module, _sp_group):
+                    def custom_forward(x_in, e_in, ctx_in):
+                        return module(
+                            x_in,
+                            e=e_in,
+                            seq_lens=seq_lens_,
+                            grid_sizes=grid_sizes_,
+                            freqs=freqs,
+                            context=ctx_in,
+                            context_lens=context_lens_,
+                            sp_group=_sp_group,
+                        )
+
+                    return custom_forward
+
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block, sp_group),
+                    x,
+                    e0_,
+                    context_,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(
+                    x,
+                    e=e0_,
+                    seq_lens=seq_lens_,
+                    grid_sizes=grid_sizes_,
+                    freqs=freqs,
+                    context=context_,
+                    context_lens=context_lens_,
+                    sp_group=sp_group,
+                )
+
+        # head operates on local tokens (like Wan2.2 official)
+        x = self.head(x, e)
+
+        # --- Ulysses SP: gather output back to full sequence ---
+        if sp_group is not None:
+            x = sp_unpad(sp_gather(x, dim=1, group=sp_group), dim=1, original_size=original_seq_len)
+
+        x = self.unpatchify(x, grid_sizes)
+        # IMPORTANT:
+        # Returning fp32 here dramatically increases activation/grad memory and slows
+        # training (DiffSynth returns bf16 here). Keep dtype consistent with model.
+        return x
+
+    def unpatchify(self, x, grid_sizes):
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[: math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum("fhwpqrc->cfphqwr", u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)
diff --git a/primus/backends/diffusion/optim/adamw_fp32_state.py b/primus/backends/diffusion/optim/adamw_fp32_state.py
new file mode 100644
index 000000000..1403043f2
--- /dev/null
+++ b/primus/backends/diffusion/optim/adamw_fp32_state.py
@@ -0,0 +1,111 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+AdamW that keeps optimizer state in FP32 even when parameters are BF16/FP16.
+
+Why:
+- DeepSpeed bf16 commonly keeps FP32 master weights / FP32 optimizer states.
+- Vanilla torch.optim.AdamW will create exp_avg/exp_avg_sq with the same dtype as the parameter,
+  so bf16 params -> bf16 states, which changes early-step behavior and can destabilize.
+
+This optimizer:
+- Stores exp_avg/exp_avg_sq in FP32
+- Applies the AdamW update in FP32 on a FP32 view of the param
+- Writes the updated value back to the original parameter dtype
+
+Scope:
+- Intended for this repo's bf16 diffusion training (single GPU or FSDP world_size=1).
+- Supports common AdamW args: lr, betas, eps, weight_decay.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable, Optional
+
+import torch
+
+
+class AdamWFP32State(torch.optim.Optimizer):
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        betas: tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid lr: {lr}")
+        if eps < 0.0:
+            raise ValueError(f"Invalid eps: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta1: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta2: {betas[1]}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure: Optional[callable] = None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr: float = group["lr"]
+            beta1, beta2 = group["betas"]
+            eps: float = group["eps"]
+            weight_decay: float = group["weight_decay"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                if p.grad.is_sparse:
+                    raise RuntimeError("AdamWFP32State does not support sparse gradients")
+
+                grad = p.grad
+                state = self.state[p]
+
+                # Initialize master weights in state if not present
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(grad, dtype=torch.float32)
+                    state["exp_avg_sq"] = torch.zeros_like(grad, dtype=torch.float32)
+                    state["master_param"] = p.detach().clone().float()
+
+                # Always use the persistent master weight for updates
+                p_fp32 = state["master_param"]
+
+                exp_avg: torch.Tensor = state["exp_avg"]
+                exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
+                state["step"] += 1
+                step: int = state["step"]
+
+                # Decoupled weight decay (AdamW)
+                if weight_decay != 0.0:
+                    p_fp32.add_(p_fp32, alpha=-lr * weight_decay)
+
+                # Adam moments
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                bias_correction1 = 1.0 - beta1**step
+                bias_correction2 = 1.0 - beta2**step
+
+                denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5)).add_(eps)
+                step_size = lr / bias_correction1
+
+                p_fp32.addcdiv_(exp_avg, denom, value=-step_size)
+
+                # Write back to original dtype
+                p.copy_(p_fp32.to(dtype=p.dtype))
+
+        return loss
diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py
new file mode 100644
index 000000000..03bb08f3c
--- /dev/null
+++ b/primus/backends/diffusion/registry.py
@@ -0,0 +1,93 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Explicit factory registry for diffusion model, dataset, and trainer builders."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Tuple
+
+
+def _build_wan_model(model_config: dict):
+    from primus.backends.diffusion.models.registrations.wan import build_wan_model
+
+    return build_wan_model(model_config)
+
+
+def _build_wan_dataset(dataset_config: dict):
+    from primus.backends.diffusion.data.registrations.wan import build_wan_dataset
+
+    return build_wan_dataset(dataset_config)
+
+
+def _build_flux_model(model_config: dict):
+    from primus.backends.diffusion.models.registrations.flux import build_flux_model
+
+    return build_flux_model(model_config)
+
+
+def _build_flux_preset_model(preset: str):
+    def _builder(model_config: dict):
+        from primus.backends.diffusion.models.registrations.flux import build_flux_model
+
+        config = dict(model_config)
+        config["model_preset"] = preset
+        return build_flux_model(config)
+
+    return _builder
+
+
+def _build_flux_dataset(dataset_config: dict):
+    from primus.backends.diffusion.data.registrations.flux import build_flux_dataset
+
+    return build_flux_dataset(dataset_config)
+
+
+def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict):
+    from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer
+
+    return build_fsdp2_trainer(
+        model=model,
+        dataset=dataset,
+        processor=processor,
+        trainer_args=trainer_args,
+    )
+
+
+MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = {
+    "flux": _build_flux_model,
+    "flux.1-dev": _build_flux_preset_model("flux.1-dev"),
+    "flux.1-schnell": _build_flux_preset_model("flux.1-schnell"),
+    "wan": _build_wan_model,
+}
+DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = {
+    "flux": _build_flux_dataset,
+    "wan": _build_wan_dataset,
+}
+TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = {
+    "fsdp2": _build_fsdp2_trainer,
+}
+
+
+def get_model_builder(name: str) -> Callable[[dict], Any]:
+    try:
+        return MODEL_BUILDERS[name]
+    except KeyError:
+        raise KeyError(f"Unknown model name: {name}")
+
+
+def get_dataset_builder(name: str) -> Callable[[dict], Tuple[Any, Any]]:
+    try:
+        return DATASET_BUILDERS[name]
+    except KeyError:
+        raise KeyError(f"Unknown dataset name: {name}")
+
+
+def get_trainer_builder(name: str) -> Callable[..., Any]:
+    try:
+        return TRAINER_BUILDERS[name]
+    except KeyError:
+        raise KeyError(f"Unknown trainer name: {name}")
diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py
new file mode 100644
index 000000000..cdd3f2f9b
--- /dev/null
+++ b/primus/backends/diffusion/schedulers/flow_match.py
@@ -0,0 +1,135 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import math
+
+import torch
+
+
+class FlowMatchScheduler:
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.exponential_shift = exponential_shift
+        self.exponential_shift_mu = exponential_shift_mu
+        self.shift_terminal = shift_terminal
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(
+        self,
+        num_inference_steps=100,
+        denoising_strength=1.0,
+        training=False,
+        shift=None,
+        dynamic_shift_len=None,
+    ):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        if self.exponential_shift:
+            mu = (
+                self.calculate_shift(dynamic_shift_len)
+                if dynamic_shift_len is not None
+                else self.exponential_shift_mu
+            )
+            if mu is None:
+                raise ValueError(
+                    "`exponential_shift=True` requires either `dynamic_shift_len` "
+                    "or `exponential_shift_mu`."
+                )
+            self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1))
+        else:
+            self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.shift_terminal is not None:
+            one_minus_z = 1 - self.sigmas
+            scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+            self.sigmas = 1 - (one_minus_z / scale_factor)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+            self.training = True
+        else:
+            self.training = False
+
+    def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def return_to_timestep(self, timestep, sample, sample_stablized):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        model_output = (sample - sample_stablized) / sigma
+        return model_output
+
+    def add_noise(self, original_samples, noise, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+
+    def calculate_shift(
+        self,
+        image_seq_len,
+        base_seq_len: int = 256,
+        max_seq_len: int = 8192,
+        base_shift: float = 0.5,
+        max_shift: float = 0.9,
+    ):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu
diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py
new file mode 100644
index 000000000..f1fbeeb37
--- /dev/null
+++ b/primus/backends/diffusion/trainers/__init__.py
@@ -0,0 +1,11 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Trainer registrations for the Primus diffusion backend."""
+
+from .fsdp2 import build_fsdp2_trainer
+
+__all__ = ["build_fsdp2_trainer"]
diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py
new file mode 100644
index 000000000..f47537096
--- /dev/null
+++ b/primus/backends/diffusion/trainers/base.py
@@ -0,0 +1,577 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Base trainer with shared logic for Wan PyTorch trainers.
+
+This module holds functionality shared across FSDP-style Wan trainers:
+config parsing, optimizer creation, LR scheduling, training loop, logging,
+and W&B integration. Concrete trainers (e.g. FSDP2Trainer) subclass it.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import time
+from contextlib import contextmanager
+
+import torch
+
+from primus.backends.diffusion.optim.adamw_fp32_state import AdamWFP32State
+from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler
+from primus.backends.diffusion.utils.log import logger
+from primus.backends.diffusion.utils.train_utils import (
+    get_memory,
+    resolve_dtype,
+    set_seed,
+)
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+def create_lr_scheduler(optimizer, scheduler_type, warmup_steps, total_steps):
+    """
+    Create LR scheduler with warmup support.
+
+    Supports: constant, constant_with_warmup, linear, cosine, polynomial.
+    Shared between FSDP and FSDP2 trainers for consistency.
+    """
+    if total_steps <= 0:
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0)
+
+    warmup_steps = int(warmup_steps)
+    if warmup_steps > total_steps:
+        logger.warning(
+            f"Warmup steps ({warmup_steps}) exceed total steps ({total_steps}). "
+            f"Adjusting warmup steps to {total_steps}."
+        )
+        warmup_steps = total_steps
+
+    def linear_warmup(step):
+        if warmup_steps == 0:
+            return 1.0
+        return min(1.0, float(step) / float(max(1, warmup_steps)))
+
+    def linear_decay(step):
+        if step <= warmup_steps:
+            return linear_warmup(step)
+        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return max(0.0, 1.0 - progress)
+
+    def cosine_decay(step):
+        if step <= warmup_steps:
+            return linear_warmup(step)
+        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    def constant_with_warmup(step):
+        return linear_warmup(step)
+
+    def polynomial_decay(step, power=1.0):
+        if step <= warmup_steps:
+            return linear_warmup(step)
+        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return max(0.0, (1.0 - progress) ** power)
+
+    scheduler_type = (scheduler_type or "constant").lower()
+    lambdas = {
+        "constant": lambda step: 1.0,
+        "constant_with_warmup": constant_with_warmup,
+        "linear": linear_decay,
+        "cosine": cosine_decay,
+        "cosine_with_restarts": cosine_decay,
+        "polynomial": polynomial_decay,
+    }
+    lr_lambda = lambdas.get(scheduler_type)
+    if lr_lambda is None:
+        logger.warning(f"Unknown lr_scheduler_type={scheduler_type}, falling back to constant.")
+
+        def lr_lambda(step):
+            return 1.0
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+class BaseWanTrainer:
+    """
+    Shared base class for Wan PyTorch FSDP-style trainers.
+
+    Subclasses must implement:
+      - _apply_parallelism(): set up distributed wrapping / sharding
+      - save_model(): save final model
+
+    Subclasses may override:
+      - _grad_sync_context(is_update_step): gradient sync during accumulation
+      - _clip_grad_norm(): gradient clipping strategy
+      - _save_checkpoint(): periodic checkpoint saving
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        args: dict,
+        train_dataset,
+        data_collator,
+        processing_class,
+        rank: int,
+        world_size: int,
+        local_rank: int,
+    ):
+        self.model = model
+        self.args = args
+        self.rank = rank
+        self.world_size = world_size
+        self.local_rank = local_rank
+        self.device = torch.device(f"cuda:{local_rank}")
+
+        # --- Config extraction ---
+        self.output_dir = self.args.get("output_dir", "./output")
+        self.logging_steps = int(self.args.get("logging_steps", 1))
+        self.save_steps = int(self.args.get("save_steps", 0))
+        self.max_steps = int(self.args.get("max_steps", -1) if self.args.get("max_steps") is not None else -1)
+        self.grad_accum_steps = int(self.args.get("gradient_accumulation_steps", 1))
+        self.max_grad_norm = float(self.args.get("max_grad_norm", 1.0))
+        self.num_train_epochs = int(self.args.get("num_train_epochs", 1))
+
+        if self.rank == 0:
+            os.makedirs(self.output_dir, exist_ok=True)
+
+        # --- Seeding ---
+        seed = self.args.get("seed")
+        if seed is not None:
+            set_seed(int(seed))
+
+        # --- Gradient Checkpointing ---
+        if self.args.get("gradient_checkpointing", False):
+            if hasattr(self.model, "gradient_checkpointing_enable"):
+                self.model.gradient_checkpointing_enable()
+            elif hasattr(self.model, "dit") and hasattr(self.model.dit, "gradient_checkpointing"):
+                self.model.dit.gradient_checkpointing = True
+            if self.rank == 0:
+                logger.info("Gradient checkpointing enabled")
+
+        # --- W&B ---
+        self._setup_wandb()
+
+        # --- Parallelism (subclass hook) ---
+        # Subclass sets self.sp_group (Ulysses SP group) if SP is enabled.
+        self.sp_group = None
+        self._apply_parallelism()
+
+        # --- DataLoader ---
+        self.train_dataset = train_dataset
+        self.processing_class = processing_class
+        self.data_collator = data_collator
+
+        # When SP is enabled, all ranks in the same SP group process the same sample.
+        # DistributedSampler should use DP-only rank/size so SP peers get identical data.
+        self.sp_size = 1
+        dp_world_size = world_size
+        dp_rank = rank
+        if self.sp_group is not None:
+            import torch.distributed as dist
+
+            self.sp_size = dist.get_world_size(self.sp_group)
+            dp_world_size = world_size // self.sp_size
+            dp_rank = rank // self.sp_size
+
+        self.data_parallel_world_size = dp_world_size
+        self.per_device_train_batch_size = int(self.args.get("per_device_train_batch_size", 1))
+
+        self.sampler = torch.utils.data.distributed.DistributedSampler(
+            train_dataset,
+            num_replicas=dp_world_size,
+            rank=dp_rank,
+            shuffle=self.args.get("shuffle", True),
+        )
+
+        num_workers = int(self.args.get("dataloader_num_workers", 4) or 0)
+        self.dataloader = torch.utils.data.DataLoader(
+            train_dataset,
+            batch_size=self.per_device_train_batch_size,
+            sampler=self.sampler,
+            num_workers=num_workers,
+            collate_fn=data_collator,
+            pin_memory=True,
+            persistent_workers=num_workers > 0,
+            prefetch_factor=2 if num_workers > 0 else None,
+        )
+
+        # --- Optimizer ---
+        self.optimizer = self._create_optimizer()
+
+        # --- LR Scheduler ---
+        steps_per_epoch = math.ceil(len(self.dataloader) / max(1, self.grad_accum_steps))
+        self.total_steps = self.max_steps if self.max_steps > 0 else self.num_train_epochs * steps_per_epoch
+        self.lr_scheduler = create_lr_scheduler(
+            self.optimizer,
+            self.args.get("lr_scheduler_type", "constant"),
+            int(self.args.get("warmup_steps", 0)),
+            self.total_steps,
+        )
+
+        # --- Diffusion Scheduler (configurable from YAML) ---
+        scheduler_cfg = self.args.get("flow_match_scheduler", {}) or {}
+        self.scheduler = FlowMatchScheduler(
+            shift=float(scheduler_cfg.get("shift", 5)),
+            sigma_min=float(scheduler_cfg.get("sigma_min", 0.0)),
+            extra_one_step=bool(scheduler_cfg.get("extra_one_step", True)),
+        )
+        self.scheduler.set_timesteps(
+            int(scheduler_cfg.get("num_train_timesteps", 1000)),
+            training=True,
+        )
+
+        self.global_step = 0
+
+    # ------------------------------------------------------------------ #
+    #                       Subclass hooks                                 #
+    # ------------------------------------------------------------------ #
+
+    def _apply_parallelism(self):
+        """Set up distributed parallelism. Called during __init__."""
+        raise NotImplementedError
+
+    @contextmanager
+    def _grad_sync_context(self, is_update_step: bool):
+        """Context manager for gradient sync control during accumulation."""
+        yield
+
+    def _clip_grad_norm(self) -> float:
+        """Clip gradient norm. Returns the total norm value."""
+        if self.max_grad_norm <= 0:
+            return 0.0
+
+        parameters = [p for p in self.model.parameters() if p.grad is not None]
+        if not parameters:
+            return 0.0
+
+        grads = [p.grad for p in parameters]
+        norm = torch.nn.utils.get_total_norm(grads, norm_type=2.0, foreach=True)
+        dtensor_cls = None
+        try:
+            from torch.distributed.tensor import DTensor
+
+            dtensor_cls = DTensor
+        except (ImportError, RuntimeError) as exc:
+            logger.debug(f"Skipping DTensor grad-norm conversion: {exc}")
+        if dtensor_cls is not None and isinstance(norm, dtensor_cls):
+            norm = norm.full_tensor()
+
+        torch.nn.utils.clip_grads_with_norm_(parameters, self.max_grad_norm, norm, foreach=True)
+        return norm.item() if isinstance(norm, torch.Tensor) else float(norm)
+
+    def _save_checkpoint(self):
+        """Save checkpoint at save_steps intervals. Override for custom strategies."""
+
+    # ------------------------------------------------------------------ #
+    #                       Common methods                                 #
+    # ------------------------------------------------------------------ #
+
+    def _setup_wandb(self):
+        self.use_wandb = False
+        if self.rank != 0:
+            return
+        if str(self.args.get("report_to", "")).lower() != "wandb":
+            return
+        if self.args.get("use_wandb") is False:
+            return
+        if wandb is None:
+            logger.warning("W&B requested but wandb is not installed.")
+            return
+
+        project = self.args.get("wandb_project") or os.environ.get("WANDB_PROJECT", "primus-diffusion")
+        run_name = self.args.get("wandb_name") or self.args.get("run_name")
+        wandb_dir = self.args.get("wandb_dir") or os.environ.get("WANDB_DIR")
+        wandb.init(project=project, name=run_name, dir=wandb_dir, config=self.args)
+        self.use_wandb = True
+
+    def _resolve_dtype(self) -> torch.dtype:
+        return resolve_dtype(self.args)
+
+    def _create_optimizer(self):
+        lr = float(self.args.get("learning_rate", 1e-4))
+        wd = float(self.args.get("weight_decay", 0.01))
+        betas = (
+            float(self.args.get("adam_beta1", 0.9)),
+            float(self.args.get("adam_beta2", 0.999)),
+        )
+        eps = float(self.args.get("adam_epsilon", 1e-8))
+
+        params = [p for p in self.model.parameters() if p.requires_grad]
+        optimizer_kwargs = {
+            "params": params,
+            "lr": lr,
+            "betas": betas,
+            "eps": eps,
+            "weight_decay": wd,
+        }
+
+        optimizer = None
+        try:
+            optimizer = torch.optim.AdamW(**optimizer_kwargs, fused=True)
+            if self.rank == 0:
+                logger.info("Optimizer: torch.optim.AdamW(fused=True)")
+        except (TypeError, RuntimeError):
+            try:
+                optimizer = torch.optim.AdamW(**optimizer_kwargs, foreach=True)
+                if self.rank == 0:
+                    logger.info("Optimizer: torch.optim.AdamW(foreach=True)")
+            except (TypeError, RuntimeError):
+                optimizer = torch.optim.AdamW(**optimizer_kwargs)
+                if self.rank == 0:
+                    logger.info("Optimizer: torch.optim.AdamW(default)")
+
+        if (self.args.get("bf16", False) or self.args.get("fp16", False)) and os.getenv(
+            "FP32_MASTER_WEIGHTS", "0"
+        ) == "1":
+            if self.rank == 0:
+                logger.info(
+                    "FP32_MASTER_WEIGHTS=1: using AdamWFP32State (fp32 master weights + fp32 moments)."
+                )
+            optimizer = AdamWFP32State(
+                optimizer.param_groups,
+                lr=lr,
+                betas=betas,
+                eps=eps,
+                weight_decay=wd,
+            )
+
+        return optimizer
+
+    def compute_loss(self, batch):
+        """Prepare batch and compute training loss."""
+        prepare_batch = getattr(self.processing_class, "prepare_batch", None)
+        if callable(prepare_batch):
+            batch = prepare_batch(
+                batch=batch,
+                device=self.device,
+                dtype=self._resolve_dtype(),
+            )
+        # Ensure all tensors are on the correct device
+        for k, v in batch.items():
+            if isinstance(v, torch.Tensor):
+                batch[k] = v.to(self.device, non_blocking=True)
+
+        # Pass SP group so model can shard sequences across SP ranks
+        if self.sp_group is not None:
+            batch["sp_group"] = self.sp_group
+
+        # Use explicit training entry point if available (GenAIModel interface)
+        forward_train = getattr(self.model, "forward_train", None)
+        if callable(forward_train):
+            outputs = forward_train(batch, scheduler=self.scheduler)
+        else:
+            outputs = self.model(batch, self.scheduler)
+        return outputs["loss"]
+
+    def _infer_batch_size_from_tensors(self, value) -> int | None:
+        if isinstance(value, torch.Tensor):
+            return int(value.shape[0]) if value.ndim > 0 else 1
+        if isinstance(value, dict):
+            for item in value.values():
+                batch_size = self._infer_batch_size_from_tensors(item)
+                if batch_size is not None:
+                    return batch_size
+        if isinstance(value, (list, tuple)):
+            for item in value:
+                batch_size = self._infer_batch_size_from_tensors(item)
+                if batch_size is not None:
+                    return batch_size
+        return None
+
+    def _infer_batch_size_from_sequences(self, value) -> int | None:
+        if isinstance(value, dict):
+            for item in value.values():
+                batch_size = self._infer_batch_size_from_sequences(item)
+                if batch_size is not None:
+                    return batch_size
+            return None
+        if isinstance(value, (list, tuple)):
+            if not value:
+                return 0
+            first = value[0]
+            if isinstance(first, (dict, list, tuple, torch.Tensor)):
+                for item in value:
+                    batch_size = self._infer_batch_size_from_sequences(item)
+                    if batch_size is not None:
+                        return batch_size
+                return None
+            return len(value)
+        return None
+
+    def _infer_local_batch_size(self, batch) -> int:
+        if isinstance(batch, (list, tuple)):
+            return len(batch)
+
+        tensor_batch_size = self._infer_batch_size_from_tensors(batch)
+        if tensor_batch_size is not None:
+            return tensor_batch_size
+
+        sequence_batch_size = self._infer_batch_size_from_sequences(batch)
+        if sequence_batch_size is not None:
+            return sequence_batch_size
+
+        return self.per_device_train_batch_size
+
+    def _compute_samples_per_gpu_per_second(
+        self,
+        local_samples: int,
+        interval_seconds: float | None,
+    ) -> float | None:
+        if interval_seconds is None or interval_seconds <= 0 or local_samples <= 0 or self.world_size <= 0:
+            return None
+
+        global_samples = float(local_samples) * float(self.data_parallel_world_size)
+        return global_samples / float(self.world_size) / float(interval_seconds)
+
+    def _log_step(
+        self,
+        loss_value: float,
+        grad_norm: float = 0.0,
+        step_time: float | None = None,
+        elapsed: float | None = None,
+        eta_seconds: float | None = None,
+        throughput_samples_per_gpu_s: float | None = None,
+    ):
+        """Log training metrics. Format matches test regex expectations."""
+        if self.rank != 0:
+            return
+        alloc, res, max_mem = get_memory()
+        lr = self.optimizer.param_groups[0]["lr"]
+
+        # NOTE: The "step=... loss=... mem=.../...GB" line format is relied on by
+        # downstream log parsers; keep it stable when editing.
+        msg = (
+            f"step={self.global_step} loss={loss_value:.4f} "
+            f"mem={alloc:.2f}/{res:.2f}GB peak_mem={max_mem:.2f}GB "
+            f"gnorm={grad_norm:.4f}"
+        )
+        if step_time is not None:
+            msg += f" step_time={step_time:.2f}s"
+        if throughput_samples_per_gpu_s is not None:
+            msg += f" throughput={throughput_samples_per_gpu_s:.4f}samples/gpu/s"
+        if elapsed is not None:
+            msg += f" elapsed={elapsed / 60:.2f}m"
+        if eta_seconds is not None:
+            msg += f" eta={eta_seconds / 60:.2f}m"
+        logger.info(msg)
+
+        if self.use_wandb:
+            payload = {
+                "train/loss": loss_value,
+                "train/step": self.global_step,
+                "train/grad_norm": grad_norm,
+                "train/lr": lr,
+                "mem/allocated_gb": alloc,
+                "mem/reserved_gb": res,
+                "mem/max_alloc_gb": max_mem,
+            }
+            if step_time is not None:
+                payload["time/step_s"] = step_time
+            if throughput_samples_per_gpu_s is not None:
+                payload["perf/samples_per_gpu_s"] = throughput_samples_per_gpu_s
+            if elapsed is not None:
+                payload["time/elapsed_s"] = elapsed
+            if eta_seconds is not None:
+                payload["time/eta_s"] = eta_seconds
+            wandb.log(payload, step=self.global_step)
+
+    def train(self):
+        if self.rank == 0:
+            logger.info("Starting training...")
+
+        # Ensure frozen state (idempotent)
+        core = getattr(self.model, "module", self.model)
+        if hasattr(core, "freeze_except"):
+            core.freeze_except()
+
+        self.model.train()
+        self.optimizer.zero_grad(set_to_none=True)
+        torch.cuda.reset_peak_memory_stats()
+
+        start_time = time.time()
+        last_log_time = start_time
+        local_samples_in_update = 0
+        local_samples_since_log = 0
+        update_steps_since_log = 0
+        update_loss_sum = 0.0
+        update_loss_count = 0
+
+        for epoch in range(self.num_train_epochs):
+            self.sampler.set_epoch(epoch)
+
+            for batch_idx, batch in enumerate(self.dataloader):
+                is_update_step = ((batch_idx + 1) % max(1, self.grad_accum_steps)) == 0
+                local_samples_in_update += self._infer_local_batch_size(batch)
+
+                with self._grad_sync_context(is_update_step):
+                    raw_loss = self.compute_loss(batch)
+                    update_loss_sum += raw_loss.detach().float().item()
+                    update_loss_count += 1
+                    loss = raw_loss / max(1, self.grad_accum_steps)
+                    loss.backward()
+
+                if is_update_step:
+                    loss_val = update_loss_sum / max(1, update_loss_count)
+                    update_loss_sum = 0.0
+                    update_loss_count = 0
+                    grad_norm = self._clip_grad_norm()
+
+                    self.optimizer.step()
+                    self.lr_scheduler.step()
+                    self.optimizer.zero_grad(set_to_none=True)
+                    self.global_step += 1
+                    update_steps_since_log += 1
+                    local_samples_since_log += local_samples_in_update
+                    local_samples_in_update = 0
+
+                    # Logging
+                    if self.global_step % self.logging_steps == 0:
+                        now = time.time()
+                        log_interval = now - last_log_time
+                        step_time = log_interval / max(1, update_steps_since_log)
+                        last_log_time = now
+                        elapsed = now - start_time
+                        steps_left = max(0, self.total_steps - self.global_step)
+                        eta_seconds = step_time * steps_left
+                        throughput_samples_per_gpu_s = self._compute_samples_per_gpu_per_second(
+                            local_samples=local_samples_since_log,
+                            interval_seconds=log_interval,
+                        )
+                        self._log_step(
+                            loss_val,
+                            grad_norm=grad_norm,
+                            step_time=step_time,
+                            elapsed=elapsed,
+                            eta_seconds=eta_seconds,
+                            throughput_samples_per_gpu_s=throughput_samples_per_gpu_s,
+                        )
+                        local_samples_since_log = 0
+                        update_steps_since_log = 0
+
+                    # Periodic save
+                    if self.save_steps > 0 and self.global_step % self.save_steps == 0:
+                        self._save_checkpoint()
+
+                    # Early termination
+                    if self.max_steps > 0 and self.global_step >= self.max_steps:
+                        return
+
+            if self.max_steps > 0 and self.global_step >= self.max_steps:
+                break
+
+        if self.rank == 0:
+            elapsed = time.time() - start_time
+            logger.info(f"Training finished in {elapsed / 60:.2f} min")
+
+    def save_model(self):
+        """Save final model. Override in subclass."""
+        raise NotImplementedError
diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py
new file mode 100644
index 000000000..958ea4132
--- /dev/null
+++ b/primus/backends/diffusion/trainers/fsdp2.py
@@ -0,0 +1,440 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Wan PyTorch FSDP2 trainer
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from importlib import import_module
+
+import torch
+from safetensors.torch import load_file as safe_load_file
+from safetensors.torch import save_file as safe_save_file
+from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_model_state_dict,
+)
+
+from primus.backends.diffusion.distributed import (
+    create_device_mesh,
+    load_checkpoint_dtcp,
+    save_checkpoint_dtcp,
+    setup_distributed,
+)
+from primus.backends.diffusion.utils.log import logger
+
+from .base import BaseWanTrainer
+
+
+class FSDP2Trainer(BaseWanTrainer):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        args: dict,
+        train_dataset,
+        data_collator,
+        processing_class,
+        rank: int,
+        world_size: int,
+        local_rank: int,
+    ):
+        super().__init__(
+            model=model,
+            args=args,
+            train_dataset=train_dataset,
+            data_collator=data_collator,
+            processing_class=processing_class,
+            rank=rank,
+            world_size=world_size,
+            local_rank=local_rank,
+        )
+
+        # FSDP2-specific: checkpoint strategy
+        # - "dit_only": save `dit_model.safetensors` (default for T2V training)
+        # - "dtcp_full": save model + optimizer via DTCP
+        # - "dtcp_model_only": save only model via DTCP
+        # - "dtcp_trainable": DTCP save only trainable params + optimizer
+        self.save_strategy = str(self.args.get("save_strategy", "dit_only")).lower()
+
+        # Checkpoint loading (after optimizer & scheduler are fully initialized)
+        resume_from = self.args.get("resume_from_checkpoint")
+        if resume_from:
+            self._load_checkpoint(resume_from)
+
+    # ------------------------------------------------------------------ #
+    #                       Parallelism                                    #
+    # ------------------------------------------------------------------ #
+
+    def _apply_parallelism(self):
+        """Set up FSDP2 composable sharding, optionally with Ulysses SP."""
+        sp_size = int(self.args.get("sp_size", 1))
+        dp_replicate = int(self.args.get("dp_replicate", 1))
+
+        self.mesh = create_device_mesh(self.world_size, sp_size=sp_size, dp_replicate=dp_replicate)
+        self.sp_group = self.mesh.get_group("ulysses") if (self.mesh is not None and sp_size > 1) else None
+        self.model.to(self.device)
+
+        # Freeze non-trainable params BEFORE FSDP and optimizer creation
+        if hasattr(self.model, "freeze_except"):
+            self.model.freeze_except()
+            if self.rank == 0:
+                logger.info("FSDP2: Applied freeze_except (frozen non-trainable params)")
+
+        self._apply_fsdp2()
+
+    def _apply_fsdp2(self):
+        """Apply torch.distributed._composable.fsdp.fully_shard to the model."""
+        mp_dtype = self._resolve_dtype()
+
+        # ---- Mixed-precision policy (aligned with DeepSpeed bf16) ----
+        #   1. Pre-cast params to bf16 BEFORE FSDP wrapping  (bf16 storage)
+        #   2. reduce_dtype = bf16  (gradient reduce matches DeepSpeed bf16 all-reduce)
+        #   3. AdamWFP32State maintains fp32 master weights and writes back to bf16
+        if mp_dtype != torch.float32:
+            mp_policy = MixedPrecisionPolicy(
+                param_dtype=mp_dtype,
+                reduce_dtype=mp_dtype,  # bf16 reduce (matches DeepSpeed bf16)
+            )
+        else:
+            mp_policy = None
+
+        # When SP is enabled, FSDP2 shards across dp_shard_sp (DP + SP combined).
+        # This reduces per-rank parameter memory.
+        fsdp_mesh = self.mesh
+        if self.mesh is not None and self.sp_group is not None:
+            try:
+                fsdp_mesh = self.mesh["dp_shard_sp"]
+            except KeyError:
+                pass  # fallback to full mesh
+
+        wrap_target = str(self.args.get("fsdp2_wrap_target", "") or "").strip()
+        wrap_root = (
+            self._get_module_by_path(self.model, wrap_target, option_name="fsdp2_wrap_target")
+            if wrap_target
+            else self.model
+        )
+
+        # Pre-cast to bf16 before FSDP wrapping so parameters are stored in bf16,
+        # matching DiffSynth/DeepSpeed bf16 behavior.
+        if mp_dtype != torch.float32:
+            wrap_root.to(dtype=mp_dtype)
+            if self.rank == 0:
+                logger.info(f"FSDP2: pre-cast '{wrap_target or '<model>'}' to {mp_dtype} (DiffSynth-aligned)")
+
+        if self.world_size == 1:
+            if self.rank == 0:
+                logger.info("FSDP2: world_size=1; skipping composable FSDP wrapping.")
+            return
+
+        if bool(self.args.get("compile_transformer_blocks", False)):
+            self._compile_transformer_blocks(wrap_root)
+
+        reshard_after_forward = bool(self.args.get("fsdp2_reshard_after_forward", True))
+        no_reshard_paths = {
+            item.strip()
+            for item in str(self.args.get("fsdp_module_paths_no_reshard", "") or "").split(",")
+            if item.strip()
+        }
+
+        module_paths_spec = str(self.args.get("fsdp_module_paths_to_wrap", "") or "")
+        module_paths = [item.strip() for item in module_paths_spec.split(",") if item.strip()]
+        wrapped_paths = 0
+        for module_path in module_paths:
+            module = self._get_module_by_path(wrap_root, module_path, option_name="fsdp_module_paths_to_wrap")
+            fully_shard(
+                module,
+                mesh=fsdp_mesh,
+                reshard_after_forward=False if module_path in no_reshard_paths else reshard_after_forward,
+                mp_policy=mp_policy,
+            )
+            wrapped_paths += 1
+        if self.rank == 0 and wrapped_paths:
+            logger.info(
+                "FSDP2: wrapped explicit module paths under " f"'{wrap_target or '<model>'}': {module_paths}"
+            )
+
+        # Wrap transformer blocks first for optimal memory management
+        layer_cls_spec = self.args.get("fsdp_transformer_layer_cls_to_wrap")
+        if layer_cls_spec:
+            if isinstance(layer_cls_spec, str):
+                layer_cls_items = [x.strip() for x in layer_cls_spec.split(",") if x.strip()]
+            else:
+                layer_cls_items = [str(x).strip() for x in layer_cls_spec if str(x).strip()]
+
+            cls_objs = set()
+            cls_names = set()
+            for item in layer_cls_items:
+                if "." in item:
+                    mod_path, cls_name = item.rsplit(".", 1)
+                    cls_objs.add(getattr(import_module(mod_path), cls_name))
+                else:
+                    cls_names.add(item)
+
+            wrapped_count = 0
+            seen = set()
+            for _, module in wrap_root.named_modules():
+                if id(module) in seen:
+                    continue
+                seen.add(id(module))
+                if module is wrap_root:
+                    continue
+                if (cls_objs and isinstance(module, tuple(cls_objs))) or (
+                    cls_names and module.__class__.__name__ in cls_names
+                ):
+                    fully_shard(
+                        module,
+                        mesh=fsdp_mesh,
+                        reshard_after_forward=reshard_after_forward,
+                        mp_policy=mp_policy,
+                    )
+                    wrapped_count += 1
+
+            if self.rank == 0:
+                logger.info(f"FSDP2: wrapped {wrapped_count} submodules under '{wrap_target or '<model>'}'")
+
+        fully_shard(
+            wrap_root,
+            mesh=fsdp_mesh,
+            reshard_after_forward=reshard_after_forward,
+            mp_policy=mp_policy,
+        )
+        if self.rank == 0:
+            logger.info(f"FSDP2: applied fully_shard to '{wrap_target or '<model>'}' with mp={mp_dtype}")
+
+    def _compile_transformer_blocks(self, root: torch.nn.Module) -> None:
+        compiled = 0
+        for attr in ("double_blocks", "single_blocks"):
+            blocks = getattr(root, attr, None)
+            if blocks is None:
+                continue
+            for idx, block in enumerate(blocks):
+                blocks[idx] = torch.compile(block, fullgraph=True)
+                compiled += 1
+        if self.rank == 0:
+            logger.info(f"FSDP2: compiled {compiled} FLUX transformer blocks with torch.compile")
+
+    @staticmethod
+    def _get_module_by_path(root: torch.nn.Module, path: str, *, option_name: str) -> torch.nn.Module:
+        """Resolve a dot-separated attribute path on a module."""
+        cur = root
+        if not path:
+            return cur
+        for part in path.split("."):
+            if not hasattr(cur, part):
+                raise ValueError(f"{option_name}='{path}' is invalid: missing attribute '{part}'")
+            cur = getattr(cur, part)
+        if not isinstance(cur, torch.nn.Module):
+            raise ValueError(f"{option_name}='{path}' did not resolve to a torch.nn.Module")
+        return cur
+
+    # ------------------------------------------------------------------ #
+    #                       Gradient sync                                  #
+    # ------------------------------------------------------------------ #
+
+    def _set_requires_gradient_sync(self, enabled: bool) -> None:
+        """
+        Best-effort gradient sync control for composable FSDP2.
+        `fully_shard()` turns modules into FSDPM which exposes
+        `set_requires_gradient_sync`. For grad accumulation, we disable
+        sync on non-update micro-steps.
+        """
+        seen = set()
+        for _, m in self.model.named_modules():
+            if id(m) in seen:
+                continue
+            seen.add(id(m))
+            setter = getattr(m, "set_requires_gradient_sync", None)
+            if callable(setter):
+                setter(bool(enabled))
+
+    @contextmanager
+    def _grad_sync_context(self, is_update_step: bool):
+        if self.world_size > 1 and self.grad_accum_steps > 1:
+            self._set_requires_gradient_sync(is_update_step)
+        yield
+        if is_update_step and self.world_size > 1 and self.grad_accum_steps > 1:
+            self._set_requires_gradient_sync(True)
+
+    # ------------------------------------------------------------------ #
+    #                       Checkpointing                                  #
+    # ------------------------------------------------------------------ #
+
+    def _dtcp_save_args(self):
+        """Compute DTCP save parameters based on save_strategy."""
+        if self.save_strategy == "dtcp_model_only":
+            return None, None
+        if self.save_strategy == "dtcp_trainable":
+            opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True)
+            return self.optimizer, opts
+        # dtcp_full (default non-dit strategy)
+        return self.optimizer, None
+
+    def _save_dtcp(self, path):
+        """Save checkpoint via Distributed Tensor Checkpointing."""
+        optimizer, opts = self._dtcp_save_args()
+        kwargs = {}
+        if opts is not None:
+            kwargs["model_state_options"] = opts
+            kwargs["optim_state_options"] = opts
+        save_checkpoint_dtcp(
+            self.model,
+            optimizer,
+            path,
+            epoch=0,
+            step=self.global_step,
+            additional_data={
+                "lr_scheduler": (self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None)
+            },
+            **kwargs,
+        )
+
+    def _save_checkpoint(self):
+        path = os.path.join(self.output_dir, f"checkpoint-{self.global_step}")
+        if self.save_strategy == "dit_only":
+            self._save_dit(os.path.join(path, "dit_model.safetensors"))
+        else:
+            self._save_dtcp(path)
+
+    def _load_checkpoint(self, path):
+        if self.rank == 0:
+            logger.info(f"Loading checkpoint from {path}")
+
+        if self.save_strategy == "dit_only":
+            core_model = self.model
+            if not hasattr(core_model, "dit"):
+                raise ValueError("save_strategy=dit_only requires model.dit to exist for checkpoint loading")
+            candidate = path
+            if os.path.isdir(path):
+                candidate = os.path.join(path, "dit_model.safetensors")
+            if not os.path.exists(candidate):
+                raise FileNotFoundError(f"dit_only checkpoint not found at: {candidate}")
+            state = safe_load_file(candidate)
+            missing, unexpected = core_model.dit.load_state_dict(state, strict=False)
+            if self.rank == 0:
+                logger.info(
+                    f"Loaded DiT weights from {candidate}. "
+                    f"Missing keys: {len(missing)}, "
+                    f"Unexpected keys: {len(unexpected)}"
+                )
+            return
+
+        if self.save_strategy == "dtcp_trainable":
+            opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True)
+            meta = load_checkpoint_dtcp(
+                self.model,
+                self.optimizer,
+                path,
+                model_state_options=opts,
+                optim_state_options=opts,
+            )
+        elif self.save_strategy == "dtcp_model_only":
+            meta = load_checkpoint_dtcp(self.model, None, path)
+        else:
+            meta = load_checkpoint_dtcp(self.model, self.optimizer, path)
+
+        if "step" in meta:
+            self.global_step = meta["step"]
+            if self.rank == 0:
+                logger.info(f"Resumed from step {self.global_step}")
+        if self.lr_scheduler is not None and isinstance(meta, dict) and meta.get("lr_scheduler") is not None:
+            try:
+                self.lr_scheduler.load_state_dict(meta["lr_scheduler"])
+                if self.rank == 0:
+                    logger.info("Resumed lr_scheduler state from checkpoint meta")
+            except Exception as exc:
+                if self.rank == 0:
+                    logger.warning(f"Failed to restore lr_scheduler state: {exc}")
+
+    def _save_dit(self, save_path: str) -> None:
+        """
+        Save `dit` weights as a single safetensors file.
+
+        Notes:
+        - On composable FSDP (world_size>1), ``get_model_state_dict`` with
+          ``full_state_dict=True`` returns the full dict **only on rank 0**;
+          other ranks receive an empty dict.  This is expected – do NOT
+          early-return on the empty-dict check, or non-rank-0 processes will
+          race ahead and desynchronize from rank 0 (which still needs to
+          write to disk).
+        - A ``dist.barrier()`` at the end keeps all ranks in lockstep so
+          that back-to-back saves (e.g. periodic checkpoint + final save)
+          never overlap.
+        """
+        import torch.distributed as dist
+
+        core_model = self.model
+        if not hasattr(core_model, "dit"):
+            logger.warning("save_model: model has no `dit` attribute; skipping save.")
+            return
+
+        if self.world_size > 1:
+            # FSDP-wrapped: use get_model_state_dict to unshard DTensors
+            full_state = get_model_state_dict(
+                core_model,
+                options=StateDictOptions(full_state_dict=True, cpu_offload=True),
+            )
+            dit_state_dict = {k[len("dit.") :]: v for k, v in full_state.items() if k.startswith("dit.")}
+        else:
+            # Single GPU (no FSDP wrapping): direct state_dict
+            dit_state_dict = core_model.dit.state_dict()
+            if not dit_state_dict:
+                logger.warning("save_model: DiT state dict is empty; skipping save.")
+                return
+
+        # Only rank 0 has the full dict in multi-GPU; write from rank 0 only.
+        if self.rank == 0:
+            if not dit_state_dict:
+                logger.warning("save_model: DiT state dict is empty on rank 0; skipping save.")
+            else:
+                os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                safe_save_file(dit_state_dict, save_path)
+                logger.info(f"Saved DiT weights to {save_path}")
+
+            if hasattr(core_model, "config"):
+                try:
+                    core_model.config.save_pretrained(os.path.dirname(save_path))
+                except Exception as exc:
+                    logger.warning(f"save_model: failed to save config: {exc}")
+
+        # Barrier: ensure all ranks wait for rank 0 to finish writing before
+        # any rank proceeds to the next operation (e.g. another save or exit).
+        if self.world_size > 1:
+            dist.barrier()
+
+    def save_model(self):
+        """Save final model using the configured strategy."""
+        if self.save_strategy in ("none", "skip", "disabled"):
+            if self.rank == 0:
+                logger.info(f"Skipping final checkpoint save (strategy={self.save_strategy})")
+            return
+        if self.save_strategy == "dit_only":
+            save_path = os.path.join(self.output_dir, "dit_model.safetensors")
+            self._save_dit(save_path)
+            return
+
+        path = os.path.join(self.output_dir, "checkpoint-final")
+        if self.rank == 0:
+            logger.info(f"Saving final checkpoint to {path} (strategy={self.save_strategy})")
+        self._save_dtcp(path)
+
+
+def build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict):
+    rank, world_size, local_rank = setup_distributed()
+    return FSDP2Trainer(
+        model=model,
+        args=trainer_args,
+        train_dataset=dataset,
+        data_collator=dataset.get_collator(),
+        processing_class=processor,
+        rank=rank,
+        world_size=world_size,
+        local_rank=local_rank,
+    )
diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py
new file mode 100644
index 000000000..b4ecc5389
--- /dev/null
+++ b/primus/backends/diffusion/utils/__init__.py
@@ -0,0 +1,10 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Utility package for the Primus diffusion backend.
+
+Heavy vision helpers are imported lazily by the video dataset path.
+"""
diff --git a/primus/backends/diffusion/utils/data_utils.py b/primus/backends/diffusion/utils/data_utils.py
new file mode 100644
index 000000000..4e78671da
--- /dev/null
+++ b/primus/backends/diffusion/utils/data_utils.py
@@ -0,0 +1,56 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+
+import math
+
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+
+def smart_nframes(
+    total_frames: int,
+    video_fps: int | float,
+    fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+        fps (int | float): the target sampling fps for model inputs.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR)
+    max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR)
+    nframes = total_frames / video_fps * fps
+    nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
+    nframes = floor_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
diff --git a/primus/backends/diffusion/utils/log.py b/primus/backends/diffusion/utils/log.py
new file mode 100644
index 000000000..2193ddebc
--- /dev/null
+++ b/primus/backends/diffusion/utils/log.py
@@ -0,0 +1,42 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Logger proxy for the diffusion backend.
+
+Primus configures loguru sinks whose format references bound ``extra`` fields
+(``rank``, ``world_size``, ``user``, ``team``, ``module_name``, ``node_ip``).
+Those fields are only present on the Primus-bound logger created in
+``primus.core.utils.logger.setup_logger``. Emitting from the raw, unbound
+global loguru logger (``from loguru import logger``) produces records without
+those extras, which makes every sink raise "Logging error in Handler".
+
+This module exposes a ``logger`` proxy that forwards attribute access to the
+live Primus-bound logger, so diffusion code keeps the familiar
+``logger.info(...)`` style while inheriting the bound extras. Before the Primus
+logger is initialized (e.g. standalone tooling), it falls back to the raw
+loguru logger, which at that point still uses loguru's default sink.
+"""
+
+from __future__ import annotations
+
+from primus.core.utils import logger as _primus_logger_module
+
+
+class _LoggerProxy:
+    @staticmethod
+    def _resolve():
+        bound = getattr(_primus_logger_module, "_logger", None)
+        if bound is not None:
+            return bound
+        from loguru import logger as _raw_logger
+
+        return _raw_logger
+
+    def __getattr__(self, name):
+        return getattr(self._resolve(), name)
+
+
+logger = _LoggerProxy()
diff --git a/primus/backends/diffusion/utils/train_utils.py b/primus/backends/diffusion/utils/train_utils.py
new file mode 100644
index 000000000..219fb6957
--- /dev/null
+++ b/primus/backends/diffusion/utils/train_utils.py
@@ -0,0 +1,189 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+
+"""
+utils for training
+"""
+
+import hashlib
+import os
+import random
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+from safetensors import safe_open
+
+
+def resolve_dtype(config_or_args) -> torch.dtype:
+    """
+    Resolve mixed-precision dtype from either:
+    - a dict-like config (e.g. trainer_args)
+    - an object with attributes (e.g. HF TrainingArguments)
+
+    Priority:
+      bf16 -> fp16 -> fp32
+    """
+    # Dict-like
+    if isinstance(config_or_args, dict):
+        if config_or_args.get("bf16", False):
+            return torch.bfloat16
+        if config_or_args.get("fp16", False):
+            return torch.float16
+        return torch.float32
+
+    # Attribute-like (HF TrainingArguments etc.)
+    bf16 = bool(getattr(config_or_args, "bf16", False))
+    fp16 = bool(getattr(config_or_args, "fp16", False))
+    if bf16:
+        return torch.bfloat16
+    if fp16:
+        return torch.float16
+    return torch.float32
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def count_parameters(model):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total, trainable
+
+
+def get_memory(unit=1e9):
+    torch.cuda.synchronize()
+    allocated = torch.cuda.memory_allocated()
+    reserved = torch.cuda.memory_reserved()
+    max_alloc = torch.cuda.max_memory_allocated()
+    return allocated / unit, reserved / unit, max_alloc / unit
+
+
+@contextmanager
+def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False):
+
+    old_register_parameter = torch.nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = torch.nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            kwargs["requires_grad"] = param.requires_grad
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+    def register_empty_buffer(module, name, buffer, persistent=True):
+        old_register_buffer(module, name, buffer, persistent=persistent)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = device
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    if include_buffers:
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
+    else:
+        tensor_constructors_to_patch = {}
+
+    try:
+        torch.nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+        yield
+    finally:
+        torch.nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = old_register_buffer
+        for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)
+
+
+def load_state_dict_from_folder(file_path, torch_dtype=None):
+    state_dict = {}
+    for file_name in os.listdir(file_path):
+        if "." in file_name and file_name.split(".")[-1] in ["safetensors", "bin", "ckpt", "pth", "pt"]:
+            state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype))
+    return state_dict
+
+
+def load_state_dict(file_path, torch_dtype=None, device="cpu"):
+    if file_path.endswith(".safetensors"):
+        return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device)
+    else:
+        return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device)
+
+
+def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"):
+    state_dict = {}
+    with safe_open(file_path, framework="pt", device=str(device)) as f:
+        for k in f.keys():
+            state_dict[k] = f.get_tensor(k)
+            if torch_dtype is not None:
+                state_dict[k] = state_dict[k].to(torch_dtype)
+    return state_dict
+
+
+def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"):
+    state_dict = torch.load(file_path, map_location=device, weights_only=True)
+    if torch_dtype is not None:
+        for i in state_dict:
+            if isinstance(state_dict[i], torch.Tensor):
+                state_dict[i] = state_dict[i].to(torch_dtype)
+    return state_dict
+
+
+def convert_state_dict_keys_to_single_str(state_dict, with_shape=True):
+    keys = []
+    for key, value in state_dict.items():
+        if isinstance(key, str):
+            if isinstance(value, torch.Tensor):
+                if with_shape:
+                    shape = "_".join(map(str, list(value.shape)))
+                    keys.append(key + ":" + shape)
+                else:
+                    keys.append(key)
+            elif isinstance(value, dict):
+                keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape))
+    keys.sort()
+    keys_str = ",".join(keys)
+    return keys_str
+
+
+def split_state_dict_with_prefix(state_dict):
+    keys = sorted([key for key in state_dict if isinstance(key, str)])
+    prefix_dict = {}
+    for key in keys:
+        prefix = key if "." not in key else key.split(".")[0]
+        if prefix not in prefix_dict:
+            prefix_dict[prefix] = []
+        prefix_dict[prefix].append(key)
+    state_dicts = []
+    for prefix, keys in prefix_dict.items():
+        sub_state_dict = {key: state_dict[key] for key in keys}
+        state_dicts.append(sub_state_dict)
+    return state_dicts
+
+
+def hash_state_dict_keys(state_dict, with_shape=True):
+    keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
+    keys_str = keys_str.encode(encoding="UTF-8")
+    return hashlib.md5(keys_str).hexdigest()
diff --git a/primus/backends/diffusion/utils/vision_process.py b/primus/backends/diffusion/utils/vision_process.py
new file mode 100644
index 000000000..2058179be
--- /dev/null
+++ b/primus/backends/diffusion/utils/vision_process.py
@@ -0,0 +1,583 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+# Adapted from qwen-vl-utils
+"""
+
+import base64
+import copy
+import logging
+import math
+import os
+import sys
+import time
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+MAX_RATIO = 200
+SPATIAL_MERGE_SIZE = 2
+IMAGE_MIN_TOKEN_NUM = 4
+IMAGE_MAX_TOKEN_NUM = 16384
+VIDEO_MIN_TOKEN_NUM = 128
+VIDEO_MAX_TOKEN_NUM = 768
+
+FPS = 2.0
+FRAME_FACTOR = 2
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+MAX_NUM_WORKERS_FETCH_VIDEO = 8
+REMOTE_FETCH_TIMEOUT = (5, 30)
+
+MODEL_SEQ_LEN = int(float(os.environ.get("MODEL_SEQ_LEN", 128000)))
+logger = logging.getLogger(__name__)
+
+VideoMetadata = Dict[str, Any]
+VideoReaderOutput = Tuple[torch.Tensor, VideoMetadata, float]
+FetchVideoOutput = Union[
+    torch.Tensor,
+    Tuple[torch.Tensor, VideoMetadata],
+    Tuple[torch.Tensor, float],
+    Tuple[Tuple[torch.Tensor, VideoMetadata], float],
+]
+
+
+def strip_file_uri(path: str) -> str:
+    """Return a filesystem path for local file URIs."""
+    return path[7:] if path.startswith("file://") else path
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int, width: int, factor: int, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None
+) -> Tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    max_pixels = max_pixels if max_pixels is not None else (IMAGE_MAX_TOKEN_NUM * factor**2)
+    min_pixels = min_pixels if min_pixels is not None else (IMAGE_MIN_TOKEN_NUM * factor**2)
+    assert max_pixels >= min_pixels, "The max_pixels of image must be greater than or equal to min_pixels."
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def to_rgb(pil_image: Image.Image) -> Image.Image:
+    if pil_image.mode == "RGBA":
+        white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
+        white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
+        return white_background
+    else:
+        return pil_image.convert("RGB")
+
+
+def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = 14) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+
+    image_obj = None
+    patch_factor = int(image_patch_size * SPATIAL_MERGE_SIZE)
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        with requests.get(image, stream=True, timeout=REMOTE_FETCH_TIMEOUT) as response:
+            response.raise_for_status()
+            with BytesIO(response.content) as bio:
+                image_obj = copy.deepcopy(Image.open(bio))
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            with BytesIO(data) as bio:
+                image_obj = copy.deepcopy(Image.open(bio))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(
+            f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
+        )
+    image = to_rgb(image_obj)
+
+    ## resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=patch_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", IMAGE_MIN_TOKEN_NUM * patch_factor**2)
+        max_pixels = ele.get("max_pixels", IMAGE_MAX_TOKEN_NUM * patch_factor**2)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=patch_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+    return image
+
+
+def smart_nframes(
+    ele: Dict[str, Any],
+    total_frames: int,
+    video_fps: Union[int, float],
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR)
+        nframes = total_frames / video_fps * fps
+        if nframes > total_frames:
+            logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]")
+        nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
+        nframes = floor_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: Dict[str, Any],
+) -> VideoReaderOutput:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        Tuple of video tensor (T, C, H, W), metadata, and sampled FPS.
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn(
+                "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0."
+            )
+    video_path = strip_file_uri(video_path)
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    video = video[idx]
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="torchvision",
+    )
+    return video, video_metadata, sample_fps
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def calculate_video_frame_range(
+    ele: Dict[str, Any],
+    total_frames: int,
+    video_fps: float,
+) -> Tuple[int, int, int]:
+    """
+    Calculate the start and end frame indices based on the given time range.
+
+    Args:
+        ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds).
+        total_frames (int): Total number of frames in the video.
+        video_fps (float): Frames per second of the video.
+
+    Returns:
+        tuple: A tuple containing (start_frame, end_frame, frame_count).
+
+    Raises:
+        ValueError: If input parameters are invalid or the time range is inconsistent.
+    """
+    # Validate essential parameters
+    if video_fps <= 0:
+        raise ValueError("video_fps must be a positive number")
+    if total_frames <= 0:
+        raise ValueError("total_frames must be a positive integer")
+
+    # Get start and end time in seconds
+    video_start = ele.get("video_start", None)
+    video_end = ele.get("video_end", None)
+    if video_start is None and video_end is None:
+        return 0, total_frames - 1, total_frames
+
+    max_duration = total_frames / video_fps
+    # Process start frame
+    if video_start is not None:
+        video_start_clamped = max(0.0, min(video_start, max_duration))
+        start_frame = math.ceil(video_start_clamped * video_fps)
+    else:
+        start_frame = 0
+    # Process end frame
+    if video_end is not None:
+        video_end_clamped = max(0.0, min(video_end, max_duration))
+        end_frame = math.floor(video_end_clamped * video_fps)
+        end_frame = min(end_frame, total_frames - 1)
+    else:
+        end_frame = total_frames - 1
+
+    # Validate frame order
+    if start_frame >= end_frame:
+        raise ValueError(
+            f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) "
+            f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). "
+            f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)"
+        )
+
+    logger.info(
+        f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}"
+    )
+    return start_frame, end_frame, end_frame - start_frame + 1
+
+
+def _read_video_decord(
+    ele: Dict[str, Any],
+) -> VideoReaderOutput:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        Tuple of video tensor (T, C, H, W), metadata, and sampled FPS.
+    """
+    import decord
+
+    video_path = strip_file_uri(ele["video"])
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    start_frame, end_frame, total_frames = calculate_video_frame_range(
+        ele,
+        total_frames,
+        video_fps,
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="decord",
+    )
+    return video, video_metadata, sample_fps
+
+
+def is_torchcodec_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("torchcodec") is not None
+
+
+def _read_video_torchcodec(
+    ele: Dict[str, Any],
+) -> VideoReaderOutput:
+    """read video using torchcodec.decoders.VideoDecoder
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        Tuple of video tensor (T, C, H, W), metadata, and sampled FPS.
+    """
+    from torchcodec.decoders import VideoDecoder
+
+    TORCHCODEC_NUM_THREADS = int(os.environ.get("TORCHCODEC_NUM_THREADS", 8))
+    logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}")
+    video_path = strip_file_uri(ele["video"])
+    st = time.time()
+    decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS)
+    video_fps = decoder.metadata.average_fps
+    total_frames = decoder.metadata.num_frames
+    start_frame, end_frame, total_frames = calculate_video_frame_range(
+        ele,
+        total_frames,
+        video_fps,
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    video = decoder.get_frames_at(indices=idx).data
+    logger.info(f"torchcodec:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="torchcodec",
+    )
+    return video, video_metadata, sample_fps
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+    "torchcodec": _read_video_torchcodec,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_torchcodec_available():
+        video_reader_backend = "torchcodec"
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(
+    ele: Dict[str, Any],
+    image_patch_size: int = 14,
+    return_video_sample_fps: bool = False,
+    return_video_metadata: bool = False,
+) -> FetchVideoOutput:
+    image_factor = image_patch_size * SPATIAL_MERGE_SIZE
+    VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor
+    VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        try:
+            video, video_metadata, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        except Exception as e:
+            # logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}")
+            logger.warning(
+                f"video_reader_backend {video_reader_backend} error, for {ele=} use torchvision as default, msg: {e}"
+            )
+            video, video_metadata, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele)
+    else:
+        # The input is a list of frames
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        # use ThreadPoolExecutor to parallel process frames
+        max_workers = min(MAX_NUM_WORKERS_FETCH_VIDEO, len(ele["video"]))
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(fetch_image, {"image": video_element, **process_info}, image_factor)
+                for video_element in ele["video"]
+            ]
+            image_list = [future.result() for future in futures]
+
+        nframes = ceil_by_factor(len(image_list), FRAME_FACTOR)
+        if len(image_list) < nframes:
+            image_list.extend([image_list[-1]] * (nframes - len(image_list)))
+
+        sample_fps = ele.get("sample_fps", 2.0)
+        video = torch.stack([torch.from_numpy(np.array(image).transpose(2, 0, 1)) for image in image_list])
+
+        # fake video metadata
+        raw_fps = process_info.pop("raw_fps", sample_fps)
+        video_metadata = dict(
+            fps=raw_fps,
+            frames_indices=[i for i in range(len(video))],
+            total_num_frames=(nframes / sample_fps) * raw_fps,
+        )
+
+    nframes, _, height, width = video.shape
+    min_pixels = ele.get("min_pixels", VIDEO_FRAME_MIN_PIXELS)
+    total_pixels = ele.get("total_pixels", MODEL_SEQ_LEN * image_factor * image_factor * 0.9)
+    max_pixels = max(
+        min(VIDEO_FRAME_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05)
+    )
+    max_pixels_supposed = ele.get("max_pixels", max_pixels)
+    if max_pixels_supposed > max_pixels:
+        logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].")
+    max_pixels = min(max_pixels_supposed, max_pixels)
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=image_factor,
+        )
+    else:
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=image_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    video = transforms.functional.resize(
+        video,
+        [resized_height, resized_width],
+        interpolation=InterpolationMode.BICUBIC,
+        antialias=True,
+    ).float()
+
+    final_video = (video, video_metadata) if return_video_metadata else video
+    if return_video_sample_fps:
+        return final_video, sample_fps
+    return final_video
+
+
+def extract_vision_info(
+    conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]
+) -> List[Dict[str, Any]]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele.get("type", "text") in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]],
+    return_video_kwargs: bool = False,
+    return_video_metadata: bool = False,
+    image_patch_size: int = 14,
+) -> Tuple[
+    Optional[List[Image.Image]],
+    Optional[List[Union[torch.Tensor, List[Image.Image]]]],
+    Optional[Dict[str, Any]],
+]:
+
+    vision_infos = extract_vision_info(conversations)
+    ## Read images or videos
+    image_inputs = []
+    video_inputs = []
+    video_sample_fps_list = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info, image_patch_size=image_patch_size))
+        elif "video" in vision_info:
+            video_input, video_sample_fps = fetch_video(
+                vision_info,
+                return_video_sample_fps=True,
+                image_patch_size=image_patch_size,
+                return_video_metadata=return_video_metadata,
+            )
+            video_sample_fps_list.append(video_sample_fps)
+            video_inputs.append(video_input)
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+
+    video_kwargs = {"do_sample_frames": False}
+    if not return_video_metadata:  # BC for qwen2.5vl
+        video_kwargs.update({"fps": video_sample_fps_list})
+
+    if return_video_kwargs:
+        return image_inputs, video_inputs, video_kwargs
+    return image_inputs, video_inputs
diff --git a/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py b/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py
index 8b7156e40..5970dd135 100644
--- a/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py
+++ b/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py
@@ -9,7 +9,7 @@
 from primus.backends.hummingbirdxt.argument_builder import HummingbirdXTArgBuilder
 from primus.core.backend.backend_adapter import BackendAdapter
 from primus.core.backend.backend_registry import BackendRegistry
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class HummingbirdXTAdapter(BackendAdapter):
diff --git a/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py b/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py
index 2356fe114..aed8315b8 100644
--- a/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py
+++ b/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py
@@ -11,8 +11,9 @@
 
 
 class HummingbirdXTPosttrainTrainer(BaseTrainer):
-    def __init__(self, backend_args: Any):
-        super().__init__(backend_args=backend_args)
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # Accept and forward runtime context kwargs so BaseTrainer can filter them.
+        super().__init__(backend_args=backend_args, **kwargs)
 
     def setup(self):
         pass
diff --git a/primus/backends/maxtext/maxtext_adapter.py b/primus/backends/maxtext/maxtext_adapter.py
index 7e44f7c49..417070633 100644
--- a/primus/backends/maxtext/maxtext_adapter.py
+++ b/primus/backends/maxtext/maxtext_adapter.py
@@ -24,7 +24,7 @@
 import primus.backends.maxtext.patches  # noqa: F401  # Register patches
 from primus.backends.maxtext.argument_builder import MaxTextConfigBuilder
 from primus.core.backend.backend_adapter import BackendAdapter
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 class MaxTextAdapter(BackendAdapter):
diff --git a/primus/backends/maxtext/maxtext_pretrain_trainer.py b/primus/backends/maxtext/maxtext_pretrain_trainer.py
index e521041a2..eb380a4cf 100644
--- a/primus/backends/maxtext/maxtext_pretrain_trainer.py
+++ b/primus/backends/maxtext/maxtext_pretrain_trainer.py
@@ -28,26 +28,65 @@
 from typing import Any, Dict, Optional
 
 from primus.core.trainer.base_trainer import BaseTrainer
-from primus.modules.module_utils import (
+from primus.core.utils.module_utils import (
     error_rank_0,
     log_rank_0,
     set_logging_rank,
     warning_rank_0,
 )
 
+# Primus-internal params that are not part of MaxText's config schema. MaxText
+# v26.4's pyconfig raises on unknown fields (v26.3 merely warns), so these must
+# be stripped before the config is handed to ``pyconfig.initialize``.
+_PRIMUS_ONLY_PARAMS = (
+    "file_sink_level",
+    "stderr_sink_level",
+    "sink_level",
+    "trainable",
+    "model",
+)
+
+
+def _resolve_maxtext_train():
+    """Resolve MaxText's train entrypoints across MaxText versions.
+
+    MaxText v26.4+ ships as the ``maxtext`` package with the training loop at
+    ``maxtext.trainers.pre_train.train``. MaxText v26.3 and earlier expose it
+    as ``MaxText.train``. Prefer the newer layout and fall back to the legacy
+    one so a single Primus checkout works against both images.
+
+    Returns:
+        Tuple of ``(initialize, run, module_name)`` where ``module_name`` is the
+        importable module string to use as ``argv[0]`` for ``initialize``.
+    """
+    try:
+        from maxtext.trainers.pre_train.train import initialize, run
+
+        return initialize, run, "maxtext.trainers.pre_train.train"
+    except ImportError:
+        from MaxText.train import initialize, run
+
+        return initialize, run, "MaxText.train"
+
 
 class MaxTextPretrainTrainer(BaseTrainer):
     """
     Trainer class for MaxText pre-training.
     """
 
-    def __init__(self, backend_args: Any):
-        super().__init__(backend_args=backend_args)
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # The core runtime instantiates every trainer with BaseModule-style
+        # context kwargs (module_name, primus_config, module_rank, ...). MaxText
+        # does not need them; accept and forward so BaseTrainer can filter them.
+        super().__init__(backend_args=backend_args, **kwargs)
 
         # Training state (populated in init())
         self.train_config: Optional[Any] = None
         self.recorder: Optional[Any] = None
         self.diagnostic_config: Optional[Any] = None
+        # Raw tuple returned by MaxText's initialize(); forwarded verbatim to
+        # run() so we stay agnostic to its return arity across MaxText versions.
+        self._init_result: tuple = ()
 
         log_rank_0("Initialized MaxTextPretrainTrainer")
 
@@ -72,7 +111,7 @@ def init(self):
         """
         log_rank_0("MaxTextPretrainTrainer.init() - initializing MaxText training")
 
-        from MaxText.train import initialize
+        initialize, _, module_name = _resolve_maxtext_train()
 
         from primus.backends.maxtext.argument_builder import (
             export_params_to_yaml,
@@ -83,16 +122,34 @@ def init(self):
         params_dict = namespace_to_dict(self.backend_args)
         params_dict.pop("override_model", None)
 
+        # Strip Primus-internal params that MaxText's config schema rejects.
+        for key in _PRIMUS_ONLY_PARAMS:
+            params_dict.pop(key, None)
+
         yaml_path = export_params_to_yaml(params_dict)
         try:
-            argv = ["MaxText.train", yaml_path]
-            self.train_config, self.recorder, self.diagnostic_config = initialize(argv, **override_model_args)
+            argv = [module_name, yaml_path]
+            init_result = initialize(argv, **override_model_args)
         finally:
             try:
                 os.unlink(yaml_path)
             except OSError as e:
                 error_rank_0(f"MaxTextPretrainTrainer: Failed to delete temporary YAML at {yaml_path}: {e}")
 
+        # MaxText's ``initialize()`` return arity changed across versions:
+        #   - v26.4 and earlier: ``(config, recorder, diagnostic_config)``
+        #   - v26.5+:            ``(config, recorder)`` (diagnostic_config dropped)
+        # Keep the raw tuple so ``train()`` can forward it verbatim to ``run()``,
+        # whose parameter count matches ``initialize()``'s arity in the same
+        # version. Unpacking a fixed number of names here is what previously
+        # crashed on v26.5 with "not enough values to unpack (expected 3, got 2)".
+        if not isinstance(init_result, tuple):
+            init_result = (init_result,)
+        self._init_result = init_result
+        self.train_config = init_result[0] if len(init_result) > 0 else None
+        self.recorder = init_result[1] if len(init_result) > 1 else None
+        self.diagnostic_config = init_result[2] if len(init_result) > 2 else None
+
         self._update_logger_rank()
         log_rank_0("MaxText training components initialized successfully")
 
@@ -170,8 +227,11 @@ def train(self):
 
         log_rank_0("Executing MaxText pretrain...")
 
-        from MaxText.train import run
+        _, run, _ = _resolve_maxtext_train()
 
-        run(self.train_config, self.recorder, self.diagnostic_config)
+        # Forward the exact tuple returned by initialize(); run()'s signature
+        # matches initialize()'s arity within a given MaxText version (3 args on
+        # v26.4 with diagnostic_config, 2 args on v26.5+).
+        run(*self._init_result)
 
         log_rank_0("MaxText pretrain execution completed.")
diff --git a/primus/backends/maxtext/patches/logger_patches.py b/primus/backends/maxtext/patches/logger_patches.py
index 8d7e50833..69a33ca02 100644
--- a/primus/backends/maxtext/patches/logger_patches.py
+++ b/primus/backends/maxtext/patches/logger_patches.py
@@ -15,7 +15,26 @@
 from primus.core.patches import PatchContext, register_patch
 from primus.core.utils import checker
 from primus.core.utils.logger import _logger as primus_logger
-from primus.modules.module_utils import error_rank_0, log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import error_rank_0, log_rank_0, warning_rank_0
+
+
+def _resolve_max_logging():
+    """Resolve MaxText's ``max_logging`` module across MaxText versions.
+
+    v26.4+ ships it at ``maxtext.utils.max_logging``; v26.3 and earlier expose
+    it as ``MaxText.max_logging``. Returns ``None`` if neither is importable.
+    """
+    try:
+        from maxtext.utils import max_logging as maxtext_logging
+
+        return maxtext_logging
+    except ImportError:
+        try:
+            import MaxText.max_logging as maxtext_logging
+
+            return maxtext_logging
+        except ImportError:
+            return None
 
 
 @register_patch(
@@ -31,16 +50,14 @@ def patch_maxtext_logger(ctx: PatchContext) -> None:
     """
     log_rank_0("[Patch:maxtext.logger] Patching MaxText logger...")
 
-    try:
-        import MaxText.max_logging as maxtext_logging
-
-        if hasattr(maxtext_logging, "log"):
-            maxtext_logging.log = primus_logger.info
-            warning_rank_0("[Patch:maxtext.logger] MaxText logger patched successfully.")
-        else:
-            error_rank_0("[Patch:maxtext.logger] MaxText logging module does not have a 'log' function.")
-    except ImportError:
+    maxtext_logging = _resolve_max_logging()
+    if maxtext_logging is None:
         error_rank_0("[Patch:maxtext.logger] Failed to import MaxText's logging module.")
+    elif hasattr(maxtext_logging, "log"):
+        maxtext_logging.log = primus_logger.info
+        warning_rank_0("[Patch:maxtext.logger] MaxText logger patched successfully.")
+    else:
+        error_rank_0("[Patch:maxtext.logger] MaxText logging module does not have a 'log' function.")
 
     # Configure JAX logger level based on module config
     level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
diff --git a/primus/backends/maxtext/patches/train_patches.py b/primus/backends/maxtext/patches/train_patches.py
index ca2ca1d1d..0e99db535 100644
--- a/primus/backends/maxtext/patches/train_patches.py
+++ b/primus/backends/maxtext/patches/train_patches.py
@@ -16,7 +16,26 @@
 from typing import Any, Sequence
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+
+def _resolve_train_and_pyconfig():
+    """Resolve MaxText's train module and pyconfig across MaxText versions.
+
+    MaxText v26.4+ exposes the training loop at
+    ``maxtext.trainers.pre_train.train`` and config at ``maxtext.configs.pyconfig``.
+    MaxText v26.3 and earlier expose ``MaxText.train`` and ``MaxText.pyconfig``.
+    """
+    try:
+        from maxtext.configs import pyconfig
+        from maxtext.trainers.pre_train import train as orig_train
+
+        return orig_train, pyconfig
+    except ImportError:
+        import MaxText.train as orig_train
+        from MaxText import pyconfig
+
+        return orig_train, pyconfig
 
 
 @register_patch(
@@ -28,13 +47,19 @@
 )
 def patch_train(ctx: PatchContext) -> None:
     """
-    Monkey-patch ``MaxText.train.initialize`` so that callers can pass
+    Monkey-patch MaxText's ``train.initialize`` so that callers can pass
     ``**kwargs`` which are transparently forwarded to ``pyconfig.initialize``.
     """
     log_rank_0("[Patch:maxtext.train] Patching MaxText train module...")
 
-    import MaxText.train as orig_train
-    from MaxText import pyconfig
+    try:
+        orig_train, pyconfig = _resolve_train_and_pyconfig()
+    except ImportError as e:
+        warning_rank_0(
+            f"[Patch:maxtext.train] Could not locate MaxText train/pyconfig module; "
+            f"skipping override-forwarding patch: {e}"
+        )
+        return
 
     _upstream_initialize = orig_train.initialize
 
diff --git a/primus/backends/megatron/__init__.py b/primus/backends/megatron/__init__.py
index 8cf9e639f..b1576a233 100644
--- a/primus/backends/megatron/__init__.py
+++ b/primus/backends/megatron/__init__.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -7,8 +7,29 @@
 from primus.backends.megatron.megatron_adapter import MegatronAdapter
 from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer
 from primus.backends.megatron.megatron_sft_trainer import MegatronSFTTrainer
+from primus.backends.megatron.mlperf import MLPerfMegatronPretrainTrainer
 from primus.core.backend.backend_registry import BackendRegistry
 
 BackendRegistry.register_adapter("megatron", MegatronAdapter)
 BackendRegistry.register_trainer_class(MegatronPretrainTrainer, "megatron")
 BackendRegistry.register_trainer_class(MegatronSFTTrainer, "megatron", "sft")
+BackendRegistry.register_trainer_class(MLPerfMegatronPretrainTrainer, "megatron", "mlperf_pretrain")
+
+# Export trainers for convenience
+# Use lazy import for FluxPretrainTrainer to avoid Megatron dependency
+# when importing data pipeline components
+__all__ = [
+    "MegatronAdapter",
+    "FluxPretrainTrainer",
+    "MegatronPretrainTrainer",
+    "MegatronSFTTrainer",
+]
+
+
+def __getattr__(name):
+    """Lazy import for trainer classes to avoid Megatron dependency on module import."""
+    if name == "FluxPretrainTrainer":
+        from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer
+
+        return FluxPretrainTrainer
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/primus/backends/megatron/argument_builder.py b/primus/backends/megatron/argument_builder.py
index 89b31e850..696e3c610 100644
--- a/primus/backends/megatron/argument_builder.py
+++ b/primus/backends/megatron/argument_builder.py
@@ -7,11 +7,13 @@
 from __future__ import annotations
 
 import argparse
+import enum
 from functools import lru_cache
 from types import SimpleNamespace
-from typing import Any, Dict, Mapping, Union
+from typing import Any, Callable, Dict, Mapping, Union
 
 from primus.core.utils.env import get_torchrun_env
+from primus.core.utils.module_utils import warning_rank_0
 
 
 # ------------------------------------------------------------
@@ -52,6 +54,48 @@ def _load_megatron_defaults() -> Dict[str, Any]:
     return vars(args).copy()  # Convert Namespace → dict and cache
 
 
+# ------------------------------------------------------------
+# Load Megatron's enum argparse type converters (cached)
+# ------------------------------------------------------------
+@lru_cache(maxsize=1)
+def _load_megatron_enum_types() -> Dict[str, Callable[[str], Any]]:
+    """Map each *enum-typed* Megatron arg ``dest`` to its argparse converter.
+
+    Primus feeds config/CLI values straight into the namespace instead of
+    through Megatron's argparse, so an enum arg like ``attention_backend``
+    arrives as a plain ``str`` and silently loses against every downstream
+    ``== SomeEnum.member`` comparison. We only touch enum args (the class of
+    bug this addresses); int/float/str args already arrive well-typed from
+    YAML. Reuse the parser's own converters so nothing is hand-maintained.
+    Returns ``{}`` if the parser cannot be built.
+    """
+    try:
+        parser = _build_megatron_parser()
+    except Exception:  # noqa: BLE001 - megatron may be unavailable (e.g. unit tests)
+        return {}
+    types: Dict[str, Callable[[str], Any]] = {}
+    for action in parser._actions:
+        convert = getattr(action, "type", None)
+        choices = getattr(action, "choices", None)
+        if callable(convert) and choices and all(isinstance(c, enum.Enum) for c in choices):
+            types[action.dest] = convert
+    return types
+
+
+def _coerce_value(convert: Callable[[str], Any], value: Any) -> Any:
+    """Apply an argparse ``type`` converter to a raw override value.
+
+    Mirrors what argparse does for a CLI string. Non-string values (already the
+    right type, e.g. from YAML) pass through untouched; list values are
+    converted element-wise to support ``nargs='+'`` args.
+    """
+    if isinstance(value, str):
+        return convert(value)
+    if isinstance(value, list):
+        return [convert(v) if isinstance(v, str) else v for v in value]
+    return value
+
+
 # ------------------------------------------------------------
 # MegatronArgBuilder: merge Primus → Megatron
 # ------------------------------------------------------------
@@ -93,9 +137,10 @@ def update(self, values: Union[Mapping[str, Any], SimpleNamespace]) -> "Megatron
         - None values are allowed and will override Megatron defaults.
         - Non-Megatron parameters are silently ignored.
         """
-        # Get Megatron's supported parameters
+        # Get Megatron's supported parameters and enum argparse converters
         megatron_defaults = _load_megatron_defaults()
         megatron_keys = set(megatron_defaults.keys())
+        enum_types = _load_megatron_enum_types()
 
         # Normalize input to a (key, value) iterable
         if isinstance(values, SimpleNamespace):
@@ -105,11 +150,26 @@ def update(self, values: Union[Mapping[str, Any], SimpleNamespace]) -> "Megatron
 
         for key, value in items:
             # Only accept parameters that Megatron recognizes (including None overrides)
-            if key in megatron_keys:
-                self.overrides[key] = value
+            if key not in megatron_keys:
+                continue
+
+            self.overrides[key] = self._coerce_enum(key, value, enum_types)
 
         return self
 
+    @staticmethod
+    def _coerce_enum(key: str, value: Any, enum_types: Mapping[str, Callable[[str], Any]]) -> Any:
+        """Coerce an enum override to its enum type (see _load_megatron_enum_types)."""
+        convert = enum_types.get(key)
+        if convert is None or value is None:
+            return value
+
+        try:
+            return _coerce_value(convert, value)
+        except Exception as exc:  # noqa: BLE001 - keep raw value, don't abort the build
+            warning_rank_0(f"[MegatronArgBuilder] could not coerce '{key}'={value!r}: {exc}")
+            return value
+
     # ------------------------------------------------------------------
     # Produce the final Megatron Namespace
     # ------------------------------------------------------------------
diff --git a/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py
index 38108e464..f37c7d4dd 100644
--- a/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py
+++ b/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py
@@ -11,7 +11,7 @@
     FileSystemWriterAsync,
 )
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class PrimusFileSystemWriterAsync(FileSystemWriterAsync):
diff --git a/primus/backends/megatron/core/distributed/finalize_model_grad.py b/primus/backends/megatron/core/distributed/finalize_model_grad.py
index 02743cf3a..320b48f33 100644
--- a/primus/backends/megatron/core/distributed/finalize_model_grad.py
+++ b/primus/backends/megatron/core/distributed/finalize_model_grad.py
@@ -17,7 +17,7 @@
 )
 from megatron.core.utils import get_model_config
 
-from primus.modules.trainer.megatron.utils import is_v_schedule_enabled
+from primus.backends.megatron.training.utils import is_v_schedule_enabled
 
 
 def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] = None):
diff --git a/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py b/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py
new file mode 100644
index 000000000..9ef79d67a
--- /dev/null
+++ b/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py
@@ -0,0 +1,654 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+FP8 all-gather tensor subclass for FSDP2.
+
+Wraps BF16 or FP32 parameters so that FSDP2 communicates FP8 data
+(1 byte/element) instead of BF16/FP32, reducing all-gather volume.
+
+Approach A: Keep Weight in FP8 After All-Gather.
+Uses a precomputed global scale (all-reduced amax across all shards) so that
+the unsharded FP8 weight has a single consistent scale, compatible with
+tensorwise GEMM on HIPBLASLT. After all-gather, the unsharded weight stays
+in FP8 via FP8UnshardedWeightTensor, saving ~50% memory vs BF16 baseline.
+
+Reference: torchao WeightWithDynamicFloat8CastTensor
+(ao/torchao/float8/fsdp_utils.py)
+"""
+
+import math
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.utils._pytree as pytree
+import triton
+import triton.language as tl
+from primus_turbo.pytorch.core.low_precision import (
+    Float8QuantConfig,
+    Format,
+    ScalingGranularity,
+    float8_e4m3,
+)
+from torch.distributed._tensor import DTensor
+from torch.distributed._tensor.placement_types import Partial, Replicate
+from torch.library import triton_op, wrap_triton
+
+_ops_to_preserve_subclass = {
+    torch.ops.aten.empty_like.default,
+    torch.ops.aten.new_zeros.default,
+    torch.ops.aten.slice.Tensor,
+    torch.ops.aten.copy_.default,
+    torch.ops.aten.view.default,
+    torch.ops.aten.as_strided.default,
+    torch.ops.aten._to_copy.default,
+    torch.ops.aten._pin_memory.default,
+    torch.ops.aten.split.Tensor,
+    torch.ops.aten.clone.default,
+}
+
+
+def _get_fp8_dtype(fmt: Format) -> torch.dtype:
+    if fmt == Format.E4M3:
+        return float8_e4m3
+    elif fmt == Format.HYBRID:
+        return float8_e4m3
+    else:
+        raise ValueError(f"Unsupported FP8 format for all-gather: {fmt}")
+
+
+@triton.jit
+def _quantize_fp8_prescaled_kernel(
+    input_ptr,
+    output_ptr,
+    scale_ptr,
+    n_elements,
+    FP8_MAX: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(input_ptr + offsets, mask=mask).to(tl.float32)
+    scale = tl.load(scale_ptr)
+    v = x * scale
+    v = tl.clamp(v, min=-FP8_MAX, max=FP8_MAX)
+    tl.store(output_ptr + offsets, v.to(output_ptr.dtype.element_ty), mask=mask)
+
+
+@triton_op("primus::quantize_fp8_prescaled", mutates_args=())
+def quantize_fp8_prescaled(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    scale_inv: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_max = torch.finfo(fp8_dtype).max
+    output = torch.empty_like(x, dtype=fp8_dtype)
+    n_elements = x.numel()
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+    wrap_triton(_quantize_fp8_prescaled_kernel)[grid](
+        x,
+        output,
+        scale,
+        n_elements,
+        fp8_max,
+        BLOCK_SIZE=1024,
+    )
+    return output, scale_inv.clone()
+
+
+@quantize_fp8_prescaled.register_fake
+def _quantize_fp8_prescaled_fake(x, fp8_dtype, scale, scale_inv):
+    return torch.empty_like(x, dtype=fp8_dtype), scale_inv.clone()
+
+
+# ---------------------------------------------------------------------------
+# Stochastic rounding variant: adds uniform [-0.5, 0.5) noise before
+# truncation to make quantization error unbiased in expectation.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _quantize_fp8_prescaled_stochastic_kernel(
+    input_ptr,
+    output_ptr,
+    scale_ptr,
+    seed,
+    n_elements,
+    FP8_MAX: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(input_ptr + offsets, mask=mask).to(tl.float32)
+    scale = tl.load(scale_ptr)
+    v = x * scale
+    v = tl.clamp(v, min=-FP8_MAX, max=FP8_MAX)
+    noise = tl.rand(seed, offsets) - 0.5
+    v_sr = v + noise
+    tl.store(output_ptr + offsets, v_sr.to(output_ptr.dtype.element_ty), mask=mask)
+
+
+@triton_op("primus::quantize_fp8_prescaled_stochastic", mutates_args=())
+def quantize_fp8_prescaled_stochastic(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    scale_inv: torch.Tensor,
+    seed: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_max = torch.finfo(fp8_dtype).max
+    output = torch.empty_like(x, dtype=fp8_dtype)
+    n_elements = x.numel()
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+    wrap_triton(_quantize_fp8_prescaled_stochastic_kernel)[grid](
+        x,
+        output,
+        scale,
+        seed,
+        n_elements,
+        fp8_max,
+        BLOCK_SIZE=1024,
+    )
+    return output, scale_inv.clone()
+
+
+@quantize_fp8_prescaled_stochastic.register_fake
+def _quantize_fp8_prescaled_stochastic_fake(x, fp8_dtype, scale, scale_inv, seed):
+    return torch.empty_like(x, dtype=fp8_dtype), scale_inv.clone()
+
+
+@torch.compile(mode="max-autotune-no-cudagraphs")
+def _foreach_fp8_quantize(
+    inner_tensors: list[torch.Tensor],
+    fp8_outputs: list[torch.Tensor],
+    scales: torch.Tensor,
+    fp8_max: float,
+):
+    """Batch-quantize all FP8 parameters in a single fused operation.
+
+    Uses _foreach_* ops fused by torch.compile to replace 456 individual
+    kernel launches with ~1-2 fused kernels. Float32 intermediate ensures
+    bitwise equivalence with the per-tensor quantize_fp8_prescaled path.
+    """
+    scales_list = list(scales.unbind())
+    f32 = [t.float() for t in inner_tensors]
+    scaled = torch._foreach_mul(f32, scales_list)
+    clamped = torch._foreach_clamp_min(torch._foreach_clamp_max(scaled, fp8_max), -fp8_max)
+    torch._foreach_copy_(fp8_outputs, clamped)
+
+
+@torch.compile(mode="max-autotune-no-cudagraphs")
+def _foreach_fp8_quantize_stochastic(
+    inner_tensors: list[torch.Tensor],
+    fp8_outputs: list[torch.Tensor],
+    scales: torch.Tensor,
+    fp8_max: float,
+):
+    """Batch-quantize with stochastic rounding for unbiased FP8 quantization.
+
+    Same as _foreach_fp8_quantize but adds uniform [-0.5, 0.5) noise before
+    the truncating cast to FP8, making quantization error zero in expectation.
+    """
+    scales_list = list(scales.unbind())
+    f32 = [t.float() for t in inner_tensors]
+    scaled = torch._foreach_mul(f32, scales_list)
+    noise = [torch.rand_like(s) - 0.5 for s in scaled]
+    noisy = torch._foreach_add(scaled, noise)
+    clamped = torch._foreach_clamp_min(torch._foreach_clamp_max(noisy, fp8_max), -fp8_max)
+    torch._foreach_copy_(fp8_outputs, clamped)
+
+
+class WeightWithFP8AllGatherTensor(torch.Tensor):
+    """Tensor subclass that quantizes to FP8 before FSDP2 all-gather.
+
+    Wraps a BF16 or FP32 parameter tensor. FSDP2 calls fsdp_pre_all_gather
+    to get FP8-quantized data for the collective, then fsdp_post_all_gather
+    to return an FP8UnshardedWeightTensor (keeping weight in FP8).
+
+    Uses a precomputed global scale (set by precompute_fp8_scales_for_fsdp)
+    so all shards quantize with the same scale, producing a single consistent
+    scale for the unsharded weight.
+    """
+
+    @staticmethod
+    def __new__(cls, tensor: torch.Tensor, fp8_config: Float8QuantConfig):
+        if fp8_config.granularity != ScalingGranularity.TENSORWISE:
+            raise ValueError(
+                f"FP8 all-gather only supports TENSORWISE granularity, " f"got {fp8_config.granularity}"
+            )
+        return torch.Tensor._make_wrapper_subclass(
+            cls,
+            tensor.size(),
+            strides=tensor.stride(),
+            storage_offset=tensor.storage_offset(),
+            dtype=tensor.dtype,
+            layout=tensor.layout,
+            device=tensor.device,
+            requires_grad=tensor.requires_grad,
+        )
+
+    def __init__(self, tensor: torch.Tensor, fp8_config: Float8QuantConfig):
+        self._fp8_config = fp8_config
+        self._tensor = tensor
+        # Transient: set by precompute_fp8_scales_for_fsdp(), NOT in
+        # __tensor_flatten__ (must be hashable/static for torch.compile guards).
+        self._precomputed_scale = None
+        self._precomputed_scale_inv = None
+        self._cached_fp8_data = None
+        self._use_cpp_quantize = False
+        self._stochastic_rounding = False
+        self._sr_counter = 0
+        self._deq_after_ag = False
+
+    def inner_data(self) -> torch.Tensor:
+        """Return the underlying data tensor (bypassing subclass dispatch)."""
+        return self._tensor
+
+    def fsdp_pre_all_gather(self, mesh):
+        if self._cached_fp8_data is not None:
+            return (self._cached_fp8_data,), (self._precomputed_scale_inv, self._tensor.numel())
+        fp8_dtype = _get_fp8_dtype(self._fp8_config.format)
+        if self._precomputed_scale is None:
+            raise RuntimeError(
+                "precompute_fp8_scales_for_fsdp() must be called before the first forward pass"
+            )
+        with torch.no_grad():
+            if self._use_cpp_quantize:
+                from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+                fp8_data, scale_inv = quantize_fp8(
+                    self._tensor,
+                    fp8_dtype,
+                    self._fp8_config.granularity,
+                    scale=self._precomputed_scale,
+                )
+            elif self._stochastic_rounding:
+                if self._precomputed_scale_inv is None:
+                    raise RuntimeError(
+                        "precompute_fp8_scales_for_fsdp() must set scale_inv before stochastic-rounding quantize"
+                    )
+                if not self._tensor.is_contiguous():
+                    raise RuntimeError("FP8 stochastic-rounding quantize requires contiguous input")
+                self._sr_counter += 1
+                fp8_data, scale_inv = quantize_fp8_prescaled_stochastic(
+                    self._tensor,
+                    fp8_dtype,
+                    self._precomputed_scale,
+                    self._precomputed_scale_inv,
+                    seed=self._sr_counter,
+                )
+            else:
+                if self._precomputed_scale_inv is None:
+                    raise RuntimeError("precompute_fp8_scales_for_fsdp() must set scale_inv for Triton path")
+                if not self._tensor.is_contiguous():
+                    raise RuntimeError("FP8 prescaled quantize requires contiguous input")
+                fp8_data, scale_inv = quantize_fp8_prescaled(
+                    self._tensor,
+                    fp8_dtype,
+                    self._precomputed_scale,
+                    self._precomputed_scale_inv,
+                )
+        return (fp8_data,), (scale_inv, self._tensor.numel())
+
+    def fsdp_post_all_gather(
+        self,
+        all_gather_outputs: Tuple[torch.Tensor, ...],
+        metadata: Any,
+        param_dtype: torch.dtype,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ):
+        scale_inv, shard_numel = metadata
+        (fp8_gathered,) = all_gather_outputs
+
+        if self._deq_after_ag:
+            bf16_weight = fp8_gathered.to(torch.bfloat16) * scale_inv
+            if out is not None:
+                out.data.copy_(bf16_weight)
+                return
+            return bf16_weight, (fp8_gathered,)
+
+        if out is not None:
+            # Reshard path: FSDP already filled the FP8 buffer via
+            # all-gather into tracked storage. Just update the scale.
+            target = out.data if isinstance(out, nn.Parameter) else out
+            if isinstance(target, FP8UnshardedWeightTensor):
+                target._scale_inv = scale_inv
+            return
+        return FP8UnshardedWeightTensor(fp8_gathered, scale_inv, torch.bfloat16, self._fp8_config), (
+            fp8_gathered,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        if func == torch.ops.aten.detach.default:
+            return WeightWithFP8AllGatherTensor(args[0]._tensor.detach(), args[0]._fp8_config)
+
+        if func == torch.ops.aten.copy_.default:
+            src = args[1]
+            if isinstance(src, WeightWithFP8AllGatherTensor):
+                src = src._tensor
+            args[0]._tensor.copy_(src)
+            return args[0]
+
+        fp8_config = None
+        inner_dtype = None
+
+        def unwrap(t):
+            nonlocal fp8_config, inner_dtype
+            if fp8_config is None:
+                fp8_config = t._fp8_config
+                inner_dtype = t._tensor.dtype
+            return t._tensor
+
+        args, kwargs = pytree.tree_map_only(WeightWithFP8AllGatherTensor, unwrap, (args, kwargs or {}))
+        out = func(*args, **kwargs)
+        if func not in _ops_to_preserve_subclass:
+            return out
+
+        if func == torch.ops.aten._to_copy.default:
+            target_dtype = (kwargs or {}).get("dtype", None)
+            if target_dtype is not None and target_dtype != inner_dtype:
+                return out
+
+        return pytree.tree_map_only(
+            torch.Tensor,
+            lambda x: WeightWithFP8AllGatherTensor(x, fp8_config),
+            out,
+        )
+
+    def __tensor_flatten__(self):
+        # Guard for torch.compile tracing: dynamo may inspect this tensor
+        # via __tensor_flatten__ during __init__ before attributes are set.
+        config = getattr(self, "_fp8_config", None)
+        if config is None or not hasattr(self, "_tensor"):
+            return [], {}
+        # Float8QuantConfig is a mutable dataclass (not hashable),
+        # so store fields individually as hashable metadata.
+        return ["_tensor"], {
+            "format": config.format,
+            "granularity": config.granularity,
+            "strategy": config.strategy,
+            "scale_dtype": config.scale_dtype,
+            "block_size": config.block_size,
+        }
+
+    @staticmethod
+    def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride):
+        config = Float8QuantConfig(
+            format=metadata["format"],
+            granularity=metadata["granularity"],
+            strategy=metadata["strategy"],
+            scale_dtype=metadata["scale_dtype"],
+            block_size=metadata["block_size"],
+        )
+        return WeightWithFP8AllGatherTensor(inner_tensors["_tensor"], config)
+
+    def __repr__(self):
+        return (
+            f"WeightWithFP8AllGatherTensor("
+            f"shape={list(self._tensor.shape)}, "
+            f"dtype={self._tensor.dtype}, "
+            f"device={self._tensor.device}, "
+            f"granularity={self._fp8_config.granularity})"
+        )
+
+
+_unsharded_ops_to_preserve = {
+    torch.ops.aten.as_strided.default,
+    torch.ops.aten.view.default,
+    torch.ops.aten.slice.Tensor,
+    torch.ops.aten.clone.default,
+}
+
+
+class FP8UnshardedWeightTensor(torch.Tensor):
+    """Lightweight wrapper for the unsharded FP8 weight after all-gather.
+
+    Holds raw FP8 data + scalar inverse scale. Declares dtype=orig_dtype
+    (bfloat16) so PyTorch shape/dtype inference sees BF16, but stores 1
+    byte/element. FP8 linear layers detect this subclass and extract the
+    pre-quantized data directly, skipping redundant quantization.
+    """
+
+    @staticmethod
+    def __new__(
+        cls,
+        fp8_data: torch.Tensor,
+        scale_inv: torch.Tensor,
+        orig_dtype: torch.dtype,
+        fp8_config: Float8QuantConfig,
+    ):
+        return torch.Tensor._make_wrapper_subclass(
+            cls,
+            fp8_data.size(),
+            strides=fp8_data.stride(),
+            storage_offset=fp8_data.storage_offset(),
+            dtype=orig_dtype,
+            layout=fp8_data.layout,
+            device=fp8_data.device,
+            requires_grad=False,
+        )
+
+    def __init__(
+        self,
+        fp8_data: torch.Tensor,
+        scale_inv: torch.Tensor,
+        orig_dtype: torch.dtype,
+        fp8_config: Float8QuantConfig,
+    ):
+        self._fp8_data = fp8_data
+        self._scale_inv = scale_inv
+        self._orig_dtype = orig_dtype
+        self._fp8_config = fp8_config
+
+    def get_fp8_data_and_scale_inv(self):
+        return self._fp8_data, self._scale_inv
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        if func == torch.ops.aten.detach.default:
+            # Must preserve subclass: nn.Parameter(data) calls detach()
+            # and asserts type(detach_result) == type(data).
+            self = args[0]
+            return FP8UnshardedWeightTensor(
+                self._fp8_data.detach(),
+                self._scale_inv.detach(),
+                self._orig_dtype,
+                self._fp8_config,
+            )
+
+        if func in _unsharded_ops_to_preserve:
+            self = args[0]
+            new_data = func(self._fp8_data, *args[1:], **(kwargs or {}))
+            return FP8UnshardedWeightTensor(new_data, self._scale_inv, self._orig_dtype, self._fp8_config)
+
+        raise NotImplementedError(
+            f"FP8UnshardedWeightTensor does not support {func}. "
+            f"Only FP8-aware module weights should be wrapped with FP8 all-gather."
+        )
+
+    def __tensor_flatten__(self):
+        config = getattr(self, "_fp8_config", None)
+        if config is None or not hasattr(self, "_fp8_data"):
+            return [], {}
+        return ["_fp8_data", "_scale_inv"], {
+            "orig_dtype": self._orig_dtype,
+            "format": config.format,
+            "granularity": config.granularity,
+            "strategy": config.strategy,
+            "scale_dtype": config.scale_dtype,
+            "block_size": config.block_size,
+        }
+
+    @staticmethod
+    def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride):
+        config = Float8QuantConfig(
+            format=metadata["format"],
+            granularity=metadata["granularity"],
+            strategy=metadata["strategy"],
+            scale_dtype=metadata["scale_dtype"],
+            block_size=metadata["block_size"],
+        )
+        return FP8UnshardedWeightTensor(
+            inner_tensors["_fp8_data"],
+            inner_tensors["_scale_inv"],
+            metadata["orig_dtype"],
+            config,
+        )
+
+    def __repr__(self):
+        return (
+            f"FP8UnshardedWeightTensor("
+            f"shape={list(self._fp8_data.shape)}, "
+            f"fp8_dtype={self._fp8_data.dtype}, "
+            f"orig_dtype={self._orig_dtype}, "
+            f"device={self._fp8_data.device})"
+        )
+
+
+def _wrap_fp8_weights_for_all_gather(
+    module: nn.Module,
+    fp8_config: Float8QuantConfig,
+    stochastic_rounding: bool = False,
+    deq_after_ag: bool = False,
+) -> int:
+    """Wrap FP8-eligible weight parameters with WeightWithFP8AllGatherTensor.
+
+    Must be called BEFORE fully_shard(). FSDP2 natively handles tensor
+    subclasses through fsdp_pre_all_gather/fsdp_post_all_gather.
+
+    Args:
+        stochastic_rounding: If True, use stochastic rounding during FP8
+            quantization to make quantization noise unbiased.
+        deq_after_ag: If True, dequantize FP8 back to BF16 after all-gather
+            and return a plain tensor. Downstream dynamic quantization will
+            produce a fresh scale from the actual assembled weight, avoiding
+            the numerical issues of keeping weight in FP8 with a global scale.
+
+    Returns the number of wrapped parameters.
+    """
+    wrapped_count = 0
+    for child in module.modules():
+        if not (hasattr(child, "_fp8_config") and hasattr(child, "weight")):
+            continue
+        w = child.weight
+        if w is not None and w.dtype in (torch.bfloat16, torch.float32) and w.requires_grad:
+            wrapper = WeightWithFP8AllGatherTensor(w.data, fp8_config)
+            wrapper._stochastic_rounding = stochastic_rounding
+            wrapper._deq_after_ag = deq_after_ag
+            child.weight = nn.Parameter(wrapper, requires_grad=True)
+            wrapped_count += 1
+    return wrapped_count
+
+
+@torch.no_grad()
+def precompute_fp8_scales_for_fsdp(
+    module: nn.Module,
+    cache_data: bool = True,
+    use_cpp_quantize: bool = False,
+    stochastic_rounding: bool = False,
+):
+    """Precompute global FP8 scales for all FP8-all-gather parameters.
+
+    Call after optimizer.step(), before next forward. Uses _foreach_norm
+    for fused amax computation and DTensor Partial("max") -> Replicate()
+    redistribution for a single batched all-reduce across all parameters.
+
+    Args:
+        module: The model module containing FP8-wrapped parameters.
+        cache_data: If True, batch-quantize all weights and cache the FP8 data
+            so fsdp_pre_all_gather skips per-layer quantization. If False, only
+            precompute scales; quantization happens on-demand per layer.
+        use_cpp_quantize: If True, use C++ quantize_fp8 from primus_turbo for
+            on-demand quantization (matching run_35). Skips fp8_outputs buffer
+            allocation and scale_inv computation when cache_data is False.
+        stochastic_rounding: If True, use stochastic rounding during batch
+            quantization to make quantization noise unbiased.
+    """
+    need_buffers = cache_data or not use_cpp_quantize
+    need_scale_inv = cache_data or not use_cpp_quantize
+
+    cache = getattr(module, "_fp8_scale_precompute_cache", None)
+    if cache is None:
+        fp8_dtensor_params = []
+        fp8_plain_params = []
+        inner_tensors = []
+        for param in module.parameters():
+            if isinstance(param, DTensor) and isinstance(param._local_tensor, WeightWithFP8AllGatherTensor):
+                fp8_dtensor_params.append(param)
+                inner_tensors.append(param._local_tensor._tensor)
+            elif isinstance(param, WeightWithFP8AllGatherTensor):
+                fp8_plain_params.append(param)
+                inner_tensors.append(param._tensor)
+
+        fp8_config = (
+            fp8_dtensor_params[0]._local_tensor._fp8_config
+            if fp8_dtensor_params
+            else fp8_plain_params[0]._fp8_config if fp8_plain_params else None
+        )
+        fp8_dtype = _get_fp8_dtype(fp8_config.format) if fp8_config else None
+
+        if need_buffers:
+            fp8_outputs = [torch.empty_like(t, dtype=fp8_dtype) for t in inner_tensors]
+        else:
+            fp8_outputs = []
+        all_local = [
+            p._local_tensor if isinstance(p, DTensor) else p for p in fp8_dtensor_params + fp8_plain_params
+        ]
+        cache = (
+            fp8_dtensor_params,
+            fp8_plain_params,
+            inner_tensors,
+            fp8_outputs,
+            all_local,
+            fp8_config,
+        )
+        module._fp8_scale_precompute_cache = cache
+
+    (
+        fp8_dtensor_params,
+        fp8_plain_params,
+        inner_tensors,
+        fp8_outputs,
+        all_local,
+        fp8_config,
+    ) = cache
+    if not inner_tensors:
+        return
+
+    local_amaxes = torch.stack(torch._foreach_norm(inner_tensors, ord=math.inf))
+
+    if fp8_dtensor_params:
+        mesh = fp8_dtensor_params[0].device_mesh
+        partial_amaxes = DTensor.from_local(local_amaxes, device_mesh=mesh, placements=[Partial("max")])
+        global_amaxes = partial_amaxes.redistribute(device_mesh=mesh, placements=[Replicate()]).to_local()
+    else:
+        global_amaxes = local_amaxes
+
+    fp8_max = torch.finfo(_get_fp8_dtype(fp8_config.format)).max
+    global_amaxes = global_amaxes.to(torch.float64).clamp(min=1e-12)
+    scales = (fp8_max / global_amaxes).to(torch.float32)
+    scale_invs = (1.0 / scales) if need_scale_inv else None
+
+    if cache_data:
+        if stochastic_rounding:
+            _foreach_fp8_quantize_stochastic(inner_tensors, fp8_outputs, scales, fp8_max)
+        else:
+            _foreach_fp8_quantize(inner_tensors, fp8_outputs, scales, fp8_max)
+
+    for i, local_t in enumerate(all_local):
+        local_t._precomputed_scale = scales[i]
+        local_t._precomputed_scale_inv = scale_invs[i] if scale_invs is not None else None
+        local_t._cached_fp8_data = fp8_outputs[i] if cache_data else None
+        local_t._use_cpp_quantize = use_cpp_quantize
+
+
+torch.serialization.add_safe_globals([WeightWithFP8AllGatherTensor])
diff --git a/primus/backends/megatron/core/distributed/sdma_param_gather.py b/primus/backends/megatron/core/distributed/sdma_param_gather.py
new file mode 100644
index 000000000..758e516d7
--- /dev/null
+++ b/primus/backends/megatron/core/distributed/sdma_param_gather.py
@@ -0,0 +1,243 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""SDMA (copy-engine) param all-gather helpers for the distributed optimizer.
+
+Migrated from the source patch ``megatron_sdma_allgather.patch`` (mpo branch).
+
+These helpers implement a ``torch.distributed.all_gather_into_tensor``-compatible
+all-gather that routes the gather through Primus-Turbo symmetric memory + HIP
+copy-engine (SDMA) memcpys instead of RCCL kernels. The intent is to free up
+CU/compute resources by performing the param all-gather purely on the copy
+engine, overlapping it with forward compute.
+
+The implementation is consumed by
+:mod:`primus.backends.megatron.patches.parallelism.sdma_param_all_gather_patches`,
+which swaps Megatron's ``_ParamAndGradBucketGroup.start_param_sync`` distributed
+optimizer path to dispatch per-bucket all-gathers through
+:func:`all_gather_into_tensor_sdma`.
+
+Activation:
+    ``ENABLE_SDMA_ALLGATHER=1`` (gated by the patch). When the required
+    Primus-Turbo primitives are unavailable, every call falls back to
+    ``torch.distributed.all_gather_into_tensor`` so behaviour is preserved.
+"""
+
+import importlib
+import os
+import warnings
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+
+# Minimum symmetric-memory workspace size (bytes) requested per group. The
+# workspace is reused across calls, so it is sized for the largest expected
+# bucket. Overridable via env for A/B sizing.
+_SDMA_SYMM_MEM_MIN_BYTES = int(os.getenv("MEGATRON_SDMA_SYMM_MEM_MIN_BYTES", str(749887296)))
+
+
+class _WaitableHandle:
+    """Lightweight waitable handle that mimics ``torch.distributed.Work``."""
+
+    def __init__(self, wait_fn: Optional[Callable[[], None]] = None, work=None):
+        self._wait_fn = wait_fn
+        self._work = work
+        self._done = False
+
+    def wait(self):
+        if self._done:
+            return True
+        if self._work is not None:
+            self._work.wait()
+        if self._wait_fn is not None:
+            self._wait_fn()
+        self._done = True
+        return True
+
+
+def _all_gather_into_tensor_waitable_fallback(
+    output_tensor: torch.Tensor,
+    input_tensor: torch.Tensor,
+    group: Optional[torch.distributed.ProcessGroup] = None,
+    async_op: bool = False,
+):
+    """Fallback to ``torch.distributed.all_gather_into_tensor`` with a waitable handle."""
+    work = torch.distributed.all_gather_into_tensor(
+        output_tensor, input_tensor, group=group, async_op=async_op
+    )
+    if async_op and work is not None:
+        return _WaitableHandle(work=work)
+    return _WaitableHandle()
+
+
+def _get_sdma_peer_copy_stream_count(world_size: int) -> int:
+    """Resolve the number of peer-copy streams for SDMA param all-gather."""
+    max_streams = max(world_size - 1, 1)
+    default_streams = min(max_streams, 8)
+    env_value = os.getenv("MEGATRON_SDMA_PEER_COPY_STREAMS")
+    if env_value is None:
+        return default_streams
+
+    try:
+        configured_streams = int(env_value)
+    except ValueError:
+        warnings.warn(
+            f"Invalid MEGATRON_SDMA_PEER_COPY_STREAMS={env_value!r}; using default " f"{default_streams}."
+        )
+        return default_streams
+
+    if configured_streams < 1:
+        warnings.warn(
+            f"MEGATRON_SDMA_PEER_COPY_STREAMS must be >= 1, got {configured_streams}; "
+            f"using default {default_streams}."
+        )
+        return default_streams
+
+    return min(configured_streams, max_streams)
+
+
+class _SDMAGroupRuntime:
+    """Shared SDMA runtime for a process group (comm + peer-copy streams)."""
+
+    def __init__(self, world_size: int):
+        # Barrier/publish/local-copy run on one shared communication stream.
+        self.comm_stream = torch.cuda.Stream()
+        # Limit peer-copy fan-out by default to reduce SDMA/memory-system
+        # pressure and improve overlap with forward compute.
+        peer_copy_stream_count = _get_sdma_peer_copy_stream_count(world_size)
+        self.peer_copy_streams = [torch.cuda.Stream() for _ in range(peer_copy_stream_count)]
+
+
+_sdma_group_runtime_cache: Dict[Tuple[str, int], _SDMAGroupRuntime] = {}
+
+
+def _get_sdma_group_runtime(group_name: str, world_size: int) -> _SDMAGroupRuntime:
+    cache_key = (group_name, world_size)
+    runtime = _sdma_group_runtime_cache.get(cache_key)
+    if runtime is None:
+        runtime = _SDMAGroupRuntime(world_size=world_size)
+        _sdma_group_runtime_cache[cache_key] = runtime
+    return runtime
+
+
+def all_gather_into_tensor_sdma(
+    output_tensor: torch.Tensor,
+    input_tensor: torch.Tensor,
+    group: Optional[torch.distributed.ProcessGroup] = None,
+    async_op: bool = False,
+):
+    """SDMA all-gather with ``all_gather_into_tensor``-compatible signature.
+
+    Follows the Primus-Turbo async-tp style (symmetric memory + DMA copies) and
+    always returns a ``.wait()``-able handle. Falls back to
+    ``torch.distributed.all_gather_into_tensor`` when SDMA primitives are
+    unavailable.
+    """
+    if not torch.distributed.is_initialized():
+        raise RuntimeError("torch.distributed must be initialized before all_gather_into_tensor_sdma")
+
+    if group is None:
+        group = torch.distributed.distributed_c10d._get_default_group()
+
+    world_size = torch.distributed.get_world_size(group=group)
+    rank = torch.distributed.get_rank(group=group)
+
+    if output_tensor.numel() != input_tensor.numel() * world_size:
+        raise ValueError(
+            "output_tensor.numel() must equal input_tensor.numel() * world_size " "for all_gather_into_tensor"
+        )
+    if output_tensor.dtype != input_tensor.dtype:
+        raise ValueError("output_tensor.dtype must match input_tensor.dtype for SDMA all_gather_into_tensor")
+
+    # SDMA path is CUDA-only and expects a contiguous flattened layout.
+    if (
+        not output_tensor.is_cuda
+        or not input_tensor.is_cuda
+        or output_tensor.device != input_tensor.device
+        or output_tensor.device.index != torch.cuda.current_device()
+        or not output_tensor.is_contiguous()
+    ):
+        return _all_gather_into_tensor_waitable_fallback(
+            output_tensor, input_tensor, group=group, async_op=async_op
+        )
+    assert input_tensor.is_contiguous(), "SDMA all_gather_into_tensor requires contiguous input_tensor"
+
+    try:
+        symm_mem_module = importlib.import_module("primus_turbo.pytorch.core.symm_mem")
+        hip_runtime_module = importlib.import_module("primus_turbo.pytorch.core.pyhip_runtime_wrapper")
+        get_symm_mem_workspace = symm_mem_module.get_symm_mem_workspace
+        hip_runtime = hip_runtime_module.get_hip_runtime_lib()
+        memcpy_comm_kind = hip_runtime_module.hipMemcpyKindEnum.hipMemcpyDeviceToDeviceNoCU
+    except Exception:
+        return _all_gather_into_tensor_waitable_fallback(
+            output_tensor, input_tensor, group=group, async_op=async_op
+        )
+
+    group_name = getattr(group, "group_name", None)
+    if group_name is None:
+        return _all_gather_into_tensor_waitable_fallback(
+            output_tensor, input_tensor, group=group, async_op=async_op
+        )
+
+    input_nbytes = input_tensor.nbytes
+    output_flat = output_tensor.view(-1)
+
+    # Allocate/reuse symmetric workspace and expose each rank's local shard buffer.
+    symm_mem = get_symm_mem_workspace(
+        group,
+        min_size=max(input_nbytes, _SDMA_SYMM_MEM_MIN_BYTES),
+    )
+    gather_buffers = [
+        symm_mem.get_buffer(r, input_tensor.shape, input_tensor.dtype) for r in range(world_size)
+    ]
+    runtime = _get_sdma_group_runtime(group_name=group_name, world_size=world_size)
+
+    current_stream = torch.cuda.current_stream()
+    # Per-call barrier event so each handle tracks its own all-gather point.
+    barrier_event = torch.cuda.Event()
+
+    # Ensure input is ready on the communication stream.
+    runtime.comm_stream.wait_stream(current_stream)
+
+    with torch.cuda.stream(runtime.comm_stream):
+        # Barrier (workspace safety), publish local shard, post-publish barrier.
+        symm_mem.barrier()
+        gather_buffers[rank].copy_(input_tensor, True)
+        symm_mem.barrier()
+        # Peer copies may start after the publish barrier completes.
+        barrier_event.record(runtime.comm_stream)
+        # Copy local shard into its slot in the output.
+        hip_runtime.hipMemcpyAsync(
+            output_flat.data_ptr() + rank * input_nbytes,
+            input_tensor.data_ptr(),
+            input_nbytes,
+            memcpy_comm_kind,
+            runtime.comm_stream.cuda_stream,
+        )
+
+    for peer_idx, src_rank in enumerate(r for r in range(world_size) if r != rank):
+        src_buf = gather_buffers[src_rank]
+        copy_stream = runtime.peer_copy_streams[peer_idx % len(runtime.peer_copy_streams)]
+        with torch.cuda.stream(copy_stream):
+            copy_stream.wait_event(barrier_event)
+            hip_runtime.hipMemcpyAsync(
+                output_flat.data_ptr() + src_rank * input_nbytes,
+                src_buf.data_ptr(),
+                input_nbytes,
+                memcpy_comm_kind,
+                copy_stream.cuda_stream,
+            )
+
+    def _wait_impl():
+        wait_stream = torch.cuda.current_stream()
+        wait_stream.wait_stream(runtime.comm_stream)
+        for copy_stream in runtime.peer_copy_streams:
+            wait_stream.wait_stream(copy_stream)
+
+    handle = _WaitableHandle(wait_fn=_wait_impl)
+    if not async_op:
+        handle.wait()
+    return handle
diff --git a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py
index 40e6dafe5..593f119a0 100644
--- a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py
+++ b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py
@@ -1,32 +1,69 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
-from typing import List, Optional
+from typing import Optional, Set
 
 import torch
-from megatron.core import tensor_parallel
+import torch.distributed as dist
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.distributed.data_parallel_base import _BaseDataParallel
 from megatron.core.distributed.distributed_data_parallel_config import (
     DistributedDataParallelConfig,
 )
-from megatron.core.distributed.torch_fully_sharded_data_parallel import (
-    TorchFullyShardedDataParallel,
-)
+from megatron.core.fp8_utils import is_float8tensor
 from megatron.core.models.common.embeddings.language_model_embedding import (
     LanguageModelEmbedding,
 )
 from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.ssm.mamba_layer import MambaLayer
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.transformer_layer import TransformerLayer
+from torch.distributed import ProcessGroup
+
+from primus.core.utils.module_utils import warning_rank_0
+
+try:
+    from torch.distributed import DeviceMesh
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    HAVE_FSDP = True
+except ImportError:
+    HAVE_FSDP = False
 
-from primus.modules.module_utils import warning_rank_0
 
+def _validate_mesh_ranks(rank_tensor, shard_group, replicate_group):
+    """Validate that mesh tensor slices match the actual process group ranks."""
+    current_rank = dist.get_rank()
 
-class PrimusTorchFullyShardedDataParallel(TorchFullyShardedDataParallel):
+    shard_group_ranks = dist.get_process_group_ranks(shard_group)
+    for row in rank_tensor.tolist():
+        if current_rank in row:
+            if sorted(row) != sorted(shard_group_ranks):
+                raise RuntimeError(f"Mesh shard ranks {row} != group ranks {shard_group_ranks}")
+
+    replicate_group_ranks = dist.get_process_group_ranks(replicate_group)
+    for col_idx in range(rank_tensor.shape[1]):
+        col = rank_tensor[:, col_idx].tolist()
+        if current_rank in col:
+            if sorted(col) != sorted(replicate_group_ranks):
+                raise RuntimeError(f"Mesh replicate ranks {col} != group ranks {replicate_group_ranks}")
+
+
+class PrimusTorchFullyShardedDataParallel(_BaseDataParallel):
     """
-    Customized FSDP implementation for Primus framework, with pre-defined submodules to wrap.
+    Customized FSDP implementation for Primus framework with support for TransformerBlock.
+
+    For models using TransformerBlock (e.g., Flux with heterogeneous layers), FSDP wraps
+    the TransformerBlock itself rather than individual TransformerLayer subclasses to avoid
+    duplicate mesh_dim_names errors.
+
+    Key difference from base Megatron class:
+    - Prevents wrapping of modules that are descendants of already-wrapped modules
+    - This avoids the "Invalid mesh_dim_names ('dp', 'dp')" error when using TransformerBlock
+      with heterogeneous TransformerLayer subclasses.
     """
 
     def __init__(
@@ -34,20 +71,380 @@ def __init__(
         config: TransformerConfig,
         ddp_config: DistributedDataParallelConfig,
         module: torch.nn.Module,
-        sub_modules_to_wrap: Optional[List[torch.nn.Module]] = None,
+        sub_modules_to_wrap: Optional[Set[torch.nn.Module]] = None,
+        disable_bucketing: bool = False,
+        process_group: Optional[ProcessGroup] = None,
         **kwargs,
     ):
+        if not HAVE_FSDP:
+            raise RuntimeError("TorchFullyShardedDataParallel requires PyTorch >= 2.4.0 with FSDP 2 support.")
+
+        super().__init__(config=config, module=module)
+
+        # Store ddp_config for later access
+        self.ddp_config = ddp_config
+
+        if process_group is None:
+            self.process_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        else:
+            self.process_group = process_group
+
         if sub_modules_to_wrap is None:
-            sub_modules_to_wrap = [
+            sub_modules_to_wrap = {
                 TransformerLayer,
+                MambaLayer,
                 LanguageModelEmbedding,
                 RotaryEmbedding,
                 tensor_parallel.ColumnParallelLinear,
-            ]
+            }
 
         if kwargs:
-            warning_rank_0(f"PrimusTorchFullyShardedDataParallel: not use args: {kwargs}")
+            warning_rank_0(f"PrimusTorchFullyShardedDataParallel: unused args: {kwargs}")
+
+        # Build DeviceMesh from Megatron's process groups
+        from megatron.training import get_args
+
+        from primus.core.utils.module_utils import log_rank_0
+
+        args = get_args()
+        replicate_degree = getattr(args, "data_parallel_replicate_degree", 1)
+        dp_size = dist.get_world_size(self.process_group)
+
+        if replicate_degree > 1:
+            shard_group = parallel_state.get_data_parallel_group(
+                with_context_parallel=True, partial_data_parallel=True
+            )
+            replicate_group = parallel_state.get_inter_distributed_optimizer_instance_group()
+
+            shard_size = dist.get_world_size(shard_group)
+
+            full_dp_ranks = dist.get_process_group_ranks(self.process_group)
+            rank_tensor = torch.tensor(full_dp_ranks).reshape(replicate_degree, shard_size)
+
+            _validate_mesh_ranks(rank_tensor, shard_group, replicate_group)
+
+            mesh = DeviceMesh.from_group(
+                [replicate_group, shard_group],
+                device_type="cuda",
+                mesh=rank_tensor.tolist(),
+                mesh_dim_names=("dp_replicate", "dp_shard"),
+            )
+        else:
+            mesh = DeviceMesh.from_group(
+                self.process_group,
+                device_type="cuda",
+                mesh_dim_names=("dp",),
+            )
+
+        reshard_after_forward = getattr(self.ddp_config, "reshard_after_forward", True)
+
+        kwargs = {
+            "mesh": mesh,
+            "reshard_after_forward": reshard_after_forward,
+        }
+
+        # When params_dtype is FP32 but training is BF16 (FSDP2 FP32 param optimizer),
+        # add MixedPrecisionPolicy so FSDP2 casts to BF16 for forward/backward.
+        # param_dtype=bf16 is always required: it ensures forward inputs
+        # (activations) get cast to BF16. For params with FP8 AG extensions,
+        # FSDP2 uses fsdp_post_all_gather (unaffected by param_dtype).
+        # This matches TorchTitan where param_dtype is always set from config.
+        use_fp8_all_gather = getattr(args, "use_fsdp2_fp8_all_gather", False)
+        if config.params_dtype == torch.float32 and config.bf16:
+            mp_policy = MixedPrecisionPolicy(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.bfloat16,
+            )
+            kwargs["mp_policy"] = mp_policy
+            log_rank_0(
+                "FSDP2: MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16) "
+                "[FP32 param optimizer: FSDP casts activations to BF16, BF16 reduce]"
+            )
+        elif (
+            config.params_dtype == torch.bfloat16
+            and config.bf16
+            and getattr(args, "use_fsdp2_bf16_master_weight_optimizer", False)
+        ):
+            # BF16 master weight optimizer: params are already BF16, so no
+            # param_dtype casting is needed. Setting param_dtype=None skips
+            # cast_forward_inputs entirely (avoids FP32→BF16 activation cast
+            # at every FSDP module boundary). reduce_dtype=bf16 halves
+            # ReduceScatter bandwidth; FP32 master weights in the optimizer
+            # still ensure full-precision parameter updates.
+            mp_policy = MixedPrecisionPolicy(
+                param_dtype=None,
+                reduce_dtype=torch.bfloat16,
+            )
+            kwargs["mp_policy"] = mp_policy
+            log_rank_0(
+                "FSDP2: MixedPrecisionPolicy(param_dtype=None, reduce_dtype=bf16) "
+                "[BF16 master weight optimizer: no cast_forward_inputs, BF16 reduce]"
+            )
+
+        if replicate_degree > 1:
+            log_rank_0(f"FSDP2 Configuration:")
+            log_rank_0(f"  mode: HSDP (replicate={replicate_degree}, shard={dp_size // replicate_degree})")
+            log_rank_0(
+                f"  reshard_after_forward: {reshard_after_forward} "
+                f"(ZeRO-{'3' if reshard_after_forward else '2'})"
+            )
+            log_rank_0(
+                f"  Data parallel size: {dp_size} "
+                f"(replicate={replicate_degree} x shard={dp_size // replicate_degree})"
+            )
+        else:
+            log_rank_0(f"FSDP2 Configuration:")
+            log_rank_0(f"  mode: FSDP")
+            log_rank_0(
+                f"  reshard_after_forward: {reshard_after_forward} "
+                f"(ZeRO-{'3' if reshard_after_forward else '2'})"
+            )
+            log_rank_0(f"  Data parallel size: {dp_size}")
+
+        # Helper functions to save/restore custom parameter attributes
+        def save_custom_attrs(module):
+            custom_attrs = {}
+            for name, param in module.named_parameters():
+                attrs = vars(param)
+                if is_float8tensor(param):
+                    # disable fp8 transpose cache and perform transposing fp8 weights
+                    # at each micro-batch because torch-FSDP doesn't recognize the
+                    # micro-batch id, thus removing unnecessary memory stores
+                    attrs["_fp8_attrs"]["transpose_invalid"] = False
+                    del attrs["_fp8_attrs"]["transpose"]
+                custom_attrs[name] = {k: v for k, v in attrs.items()}
+            return custom_attrs
+
+        def restore_custom_attrs(module, custom_attrs):
+            for name, param in module.named_parameters():
+                if name in custom_attrs:
+                    for attr_name, attr_value in custom_attrs[name].items():
+                        setattr(param, attr_name, attr_value)
+
+        # Save custom attributes that might be removed by FSDP
+        attrs = save_custom_attrs(self.module)
+
+        # FP8 all-gather validation
+        if (
+            use_fp8_all_gather
+            and isinstance(reshard_after_forward, int)
+            and not isinstance(reshard_after_forward, bool)
+        ):
+            raise ValueError(
+                "FP8 all-gather is incompatible with reshard_after_forward=int (partial reshard)"
+            )
+
+        # FSDP2 FP8 all-gather is incompatible with delayed scaling: the delayed
+        # forward re-quantizes the weight and does not handle the all-gathered
+        # FP8UnshardedWeightTensor subclass produced by the all-gather path, so
+        # the combination would miscompute. Reject it explicitly.
+        if use_fp8_all_gather:
+            from megatron.core.enums import Fp8Recipe
+
+            uses_delayed = (
+                getattr(config, "fp8_scaling_strategy", "dynamic") == "delayed"
+                or getattr(config, "fp8_recipe", None) == Fp8Recipe.delayed
+            )
+            if uses_delayed:
+                raise ValueError(
+                    "use_fsdp2_fp8_all_gather is incompatible with delayed FP8 scaling "
+                    "(fp8_recipe='delayed' / fp8_scaling_strategy='delayed'). Use a "
+                    "non-delayed recipe (e.g. tensorwise) with FSDP2 FP8 all-gather, or "
+                    "disable use_fsdp2_fp8_all_gather."
+                )
+
+        # Local transformer implementation does not support ColumnParallelLinear.
+        if config.transformer_impl == "local":
+            sub_modules_to_wrap = {
+                sub_module
+                for sub_module in sub_modules_to_wrap
+                if sub_module != tensor_parallel.ColumnParallelLinear
+            }
+        sub_modules_to_wrap = set(sub_modules_to_wrap)
+
+        # Process _fsdp_modules attribute
+        for sub_module in self.module.modules():
+            fsdp_modules = getattr(sub_module, "_fsdp_modules", [])
+            for f in fsdp_modules:
+                sub_modules_to_wrap.add(f)
+
+        fp8_ag_stochastic_rounding = getattr(args, "fp8_all_gather_stochastic_rounding", False)
+        fp8_ag_deq_after_ag = getattr(args, "fp8_all_gather_deq_requant", False)
+        if use_fp8_all_gather:
+            from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+                _wrap_fp8_weights_for_all_gather,
+            )
+            from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+                _build_fp8_config,
+            )
+
+            fp8_ag_config = _build_fp8_config(config)
+            wrapped_count = _wrap_fp8_weights_for_all_gather(
+                self.module,
+                fp8_ag_config,
+                stochastic_rounding=fp8_ag_stochastic_rounding,
+                deq_after_ag=fp8_ag_deq_after_ag,
+            )
+            sr_tag = " [stochastic rounding]" if fp8_ag_stochastic_rounding else ""
+            deq_tag = " [deq+requant]" if fp8_ag_deq_after_ag else ""
+            log_rank_0(
+                f"FSDP2: Wrapped {wrapped_count} params with FP8 all-gather "
+                f"(granularity={fp8_ag_config.granularity}){sr_tag}{deq_tag}"
+            )
+
+        # ============================================================================
+        # CUSTOM WRAPPING LOGIC: Skip descendants of already-wrapped modules
+        # ============================================================================
+        wrapped_modules_set = set()
+        wrapped_list = []
+
+        for sub_module in self.module.modules():
+            should_wrap = any(
+                isinstance(sub_module, sub_module_to_wrap) for sub_module_to_wrap in sub_modules_to_wrap
+            )
+            if not should_wrap:
+                continue
+
+            is_descendant = False
+            for wrapped_module in wrapped_modules_set:
+                if sub_module is not wrapped_module:
+                    for child in wrapped_module.modules():
+                        if child is sub_module:
+                            is_descendant = True
+                            break
+                    if is_descendant:
+                        break
+
+            if is_descendant:
+                continue
+
+            fully_shard(sub_module, **kwargs)
+            wrapped_modules_set.add(sub_module)
+            wrapped_list.append(sub_module)
+
+        log_rank_0(f"FSDP2: wrapped {len(wrapped_list)} inner modules + root")
+
+        prefetch_depth = getattr(self.config, "fsdp_prefetch_depth", 1)
+        recompute_on = getattr(self.config, "recompute_granularity", None) is not None
+        for i, mod in enumerate(wrapped_list):
+            # With activation recompute enabled, the recomputed forward pass triggers
+            # wrong-direction forward-prefetch all-gathers that expose communication and
+            # inflate peak memory (matches Megatron upstream, which only sets backward
+            # prefetch when recompute is on). Skip explicit forward prefetch in that case.
+            if not recompute_on:
+                fwd_targets = wrapped_list[i + 1 : i + 1 + prefetch_depth]
+                if fwd_targets:
+                    mod.set_modules_to_forward_prefetch(fwd_targets)
+
+            bwd_start = max(0, i - prefetch_depth)
+            bwd_targets = list(reversed(wrapped_list[bwd_start:i]))
+            if bwd_targets:
+                mod.set_modules_to_backward_prefetch(bwd_targets)
+
+        # Wrap the root module as required by the FSDP API
+        fully_shard(self.module, **kwargs)
+
+        if use_fp8_all_gather:
+            from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+                precompute_fp8_scales_for_fsdp,
+            )
+
+            cache_data = getattr(self.config, "fp8_precompute_data_cache", True)
+            use_cpp = getattr(self.config, "use_cpp_fp8_quantize", False)
+            precompute_fp8_scales_for_fsdp(
+                self.module,
+                cache_data=cache_data,
+                use_cpp_quantize=use_cpp,
+                stochastic_rounding=fp8_ag_stochastic_rounding,
+            )
+
+        restore_custom_attrs(self.module, attrs)
+
+        if getattr(args, "overlap_grad_norm", False):
+            from primus.backends.megatron.core.optimizer.incremental_grad_norm import (
+                IncrementalGradNormAccumulator,
+            )
+
+            # Use the shard group for the norm all-reduce (correct for both
+            # FSDP and HSDP). For HSDP, replicate-group ranks hold identical
+            # gradients after AR, so the full DP group would over-count.
+            if replicate_degree > 1:
+                norm_reduce_group = shard_group
+            else:
+                norm_reduce_group = self.process_group
+
+            accumulator = IncrementalGradNormAccumulator(
+                shard_process_group=norm_reduce_group,
+                device=torch.device("cuda"),
+            )
+            n_hooked = 0
+            for param in self.module.parameters():
+                if param.requires_grad:
+                    param.register_post_accumulate_grad_hook(accumulator.make_hook())
+                    n_hooked += 1
+            args._grad_norm_accumulator = accumulator
+            log_rank_0(f"FSDP2: Registered incremental grad norm hooks on {n_hooked} params")
+
+    def compile_model(self):
+        """
+        Delegate torch.compile to the underlying model's compile_model() method.
+
+        Traverses the FSDP2 module hierarchy (up to 3 levels deep) to find
+        the actual model (e.g., Flux) and calls its compile_model() if present.
+        Skipped if enable_torch_compile is False or no compile_model method is found.
+        """
+        from primus.core.utils.module_utils import log_rank_0
+
+        try:
+            from megatron.training import get_args
+
+            args = get_args()
+
+            # Check if compilation is enabled
+            if not getattr(args, "enable_torch_compile", False):
+                return
+
+        except Exception:
+            # If args not available, skip compilation
+            log_rank_0("  ℹ  Cannot access args for torch.compile settings, skipping")
+            return
+
+        # FSDP2 wraps models in FSDPFloat16Module, which itself wraps the actual model
+        # We need to traverse the hierarchy to find the actual model (e.g., Flux)
+        current_module = self.module
+        module_path = []
+
+        # Traverse up to 3 levels deep to find a module with compile_model
+        for depth in range(3):
+            module_type = type(current_module).__name__
+            module_path.append(module_type)
+
+            if hasattr(current_module, "compile_model") and callable(
+                getattr(current_module, "compile_model")
+            ):
+                log_rank_0(f"  Found compile_model at depth {depth}: {' -> '.join(module_path)}")
+                log_rank_0(f"  Calling compile_model on {module_type}...")
+                current_module.compile_model()
+                log_rank_0(f"    ✓ Model compilation complete")
+                return
+
+            # Try to go deeper if there's a 'module' attribute
+            if hasattr(current_module, "module"):
+                current_module = current_module.module
+            else:
+                break
+
+        # No compile_model found - log that we can't compile directly
+        # Note: We can't reassign self.module after FSDP wrapping, so if the underlying
+        # model doesn't have compile_model, we skip compilation. The model should implement
+        # compile_model if torch.compile is desired.
+        log_rank_0(f"  ℹ  No compile_model method found in module hierarchy: {' -> '.join(module_path)}")
+        log_rank_0(f"    (Model should implement compile_model() if torch.compile is desired)")
+
+    def finish_grad_sync(self, *args, **kwargs):
+        """No-op for FSDP2: gradient sync is handled by the FSDP runtime."""
 
-        super().__init__(
-            config=config, ddp_config=ddp_config, module=module, sub_modules_to_wrap=sub_modules_to_wrap
-        )
+    def load_state_dict(self, state_dict, strict=True):
+        """
+        No-op because tensors are already loaded in-place by
+        `_load_base_checkpoint` with FSDP2."""
diff --git a/primus/backends/megatron/core/enums.py b/primus/backends/megatron/core/enums.py
index 1dd3113a8..c88650fb9 100644
--- a/primus/backends/megatron/core/enums.py
+++ b/primus/backends/megatron/core/enums.py
@@ -8,7 +8,7 @@
 
 
 class Fp4Recipe(str, enum.Enum):
-    """FP4 recipe names: nvfp4."""
+    """FP4 recipe names: nvfp4, mxfp4."""
 
     nvfp4 = "nvfp4"
     mxfp4 = "mxfp4"
diff --git a/primus/backends/megatron/core/extensions/_triton/__init__.py b/primus/backends/megatron/core/extensions/_triton/__init__.py
new file mode 100644
index 000000000..7f8d5a0bb
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/_triton/__init__.py
@@ -0,0 +1,8 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton kernels for the Primus Megatron extensions package.
+"""
diff --git a/primus/backends/megatron/core/extensions/_triton/inplace_add.py b/primus/backends/megatron/core/extensions/_triton/inplace_add.py
new file mode 100644
index 000000000..210095790
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/_triton/inplace_add.py
@@ -0,0 +1,63 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""In-place ``dst += src`` Triton kernel for the Primus Megatron extensions.
+
+Used by the FP8/FP4 weight-gradient bridge in
+:mod:`primus.backends.megatron.core.extensions.primus_turbo` to accumulate a
+freshly produced weight gradient into the persistent ``main_grad`` buffer with
+a single Triton launch (fp32 accumulate), instead of Torch's ``add_`` which
+tiles a large consolidated grouped-expert ``main_grad`` of ``[E, N, K]`` into
+multiple ~528M-element ``vectorized_elementwise_kernel`` launches.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _inplace_add_kernel(dst_ptr, src_ptr, n_elements, BLOCK: tl.constexpr):
+    """In-place ``dst += src`` over a flat buffer, accumulating in fp32.
+
+    int64 offsets so a single launch covers tensors with > 2**31 elements
+    (e.g. the consolidated grouped-expert ``main_grad`` of [E, N, K]),
+    instead of Torch's ``add_`` which tiles into multiple ~528M-element
+    ``vectorized_elementwise_kernel`` launches.
+    """
+    pid = tl.program_id(axis=0).to(tl.int64)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_elements
+    d = tl.load(dst_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    s = tl.load(src_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    tl.store(dst_ptr + offs, (d + s).to(dst_ptr.dtype.element_ty), mask=mask)
+
+
+def inplace_add_triton_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
+    """``dst.add_(src)`` via a single Triton launch (fp32 accumulate).
+
+    Falls back to Torch's ``add_`` when the layout is unsupported (not on CUDA
+    / non-contiguous / shape mismatch). The write is in-place on ``dst``'s
+    storage, so ``dst`` must be contiguous.
+    """
+    if not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel():
+        return dst.add_(src)
+
+    dst_flat = dst.view(-1)
+    # reshape (not view) so a non-contiguous grad is materialized contiguously.
+    src_flat = src.reshape(-1)
+    n_elements = dst_flat.numel()
+    BLOCK = 8192
+    grid = (triton.cdiv(n_elements, BLOCK),)
+    _inplace_add_kernel[grid](dst_flat, src_flat, n_elements, BLOCK=BLOCK)
+    return dst
+
+
+__all__ = [
+    "inplace_add_triton_",
+]
diff --git a/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py b/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py
new file mode 100644
index 000000000..2d6224962
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py
@@ -0,0 +1,207 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-7 P45 — multi-tensor BF16 add Triton kernel (prototype).
+
+Replaces the ~743 separate ``vec_elem<add_bf16>`` launches that
+dominate the P40 trace's optimizer-step residual (170.99 ms / iter,
+32.7 % of step) with a single Triton kernel that processes a
+multi-tensor batch in one launch.
+
+This is a **prototype** designed to demonstrate the perf headroom
+of multi-tensor fusion against ``torch._foreach_add_`` (the
+PyTorch reference path) and the upstream Apex / TE
+``multi_tensor_apply`` (the V4-Flash production path).  Production
+integration with the Apex / TE optimizer call sites is plan-8
+scope — replacing those call sites bit-exactly requires
+coordinating with the master-param remainder accumulation logic
+that Apex / TE owns.
+
+Gating: import-time only; no env knob (this module is currently
+a microbench / unit-test target only).
+"""
+
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _multi_tensor_add_per_tensor_kernel(
+    OUT_PTR,
+    A_PTR,
+    B_PTR,
+    N,
+    SCALE,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Single-tensor variant.  Each grid program handles ``BLOCK_SIZE``
+    elements of one tensor.
+
+    Caller is responsible for issuing the per-tensor launch grid
+    and serialising the loops, but the launches are CONCURRENT on
+    the GPU (Triton stream stacking).
+    """
+
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < N
+    a = tl.load(A_PTR + offs, mask=mask, other=0.0).to(tl.float32)
+    b = tl.load(B_PTR + offs, mask=mask, other=0.0).to(tl.float32)
+    out = a + SCALE * b
+    tl.store(OUT_PTR + offs, out, mask=mask)
+
+
+def multi_tensor_add_triton_per_tensor(
+    out_list: List[torch.Tensor],
+    a_list: List[torch.Tensor],
+    b_list: List[torch.Tensor],
+    scale: float = 1.0,
+    block_size: int = 8192,
+) -> None:
+    """Compute ``out_i = a_i + scale * b_i`` for each ``i`` in place
+    (per-tensor variant).
+
+    All tensors must be contiguous + on CUDA / HIP, same dtype, same
+    shape per ``(out_i, a_i, b_i)`` triple.
+    """
+    assert len(out_list) == len(a_list) == len(b_list)
+    for out, a, b in zip(out_list, a_list, b_list):
+        assert out.shape == a.shape == b.shape, "shape mismatch"
+        n = out.numel()
+        grid = (triton.cdiv(n, block_size),)
+        _multi_tensor_add_per_tensor_kernel[grid](
+            out,
+            a,
+            b,
+            n,
+            scale,
+            BLOCK_SIZE=block_size,
+        )
+
+
+@triton.jit
+def _multi_tensor_add_packed_kernel(
+    OUT_PTR_BUF,  # int64 [N_TENSORS] pointer array
+    A_PTR_BUF,  # int64 [N_TENSORS]
+    B_PTR_BUF,  # int64 [N_TENSORS]
+    SIZE_BUF,  # int32 [N_TENSORS] -- element counts
+    PID_TO_TID_BUF,  # int32 [N_PROGRAMS] -- program-id -> tensor-id
+    PID_TO_LOCAL_BUF,  # int32 [N_PROGRAMS] -- program-id -> local block index
+    SCALE,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Single-kernel multi-tensor add (CPU-side dispatch table variant).
+
+    One program per ``BLOCK_SIZE`` chunk of exactly one tensor.
+    The CPU side builds a sorted dispatch table mapping each
+    ``program_id`` to ``(tensor_idx, local_block_idx)`` so a single
+    grid launch absorbs all N_TENSORS tensors without per-element
+    tensor lookups.
+
+    This avoids the cross-tensor-boundary aliasing that a naive
+    "concatenated range" partition would create.
+    """
+
+    pid = tl.program_id(0)
+    tensor_idx = tl.load(PID_TO_TID_BUF + pid)
+    local_block = tl.load(PID_TO_LOCAL_BUF + pid)
+
+    offs = local_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    size_i = tl.load(SIZE_BUF + tensor_idx)
+    mask = offs < size_i
+
+    out_ptr = tl.load(OUT_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16))
+    a_ptr = tl.load(A_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16))
+    b_ptr = tl.load(B_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16))
+
+    a = tl.load(a_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    out = a + SCALE * b
+    tl.store(out_ptr + offs, out.to(tl.bfloat16), mask=mask)
+
+
+def _build_multi_tensor_dispatch_table(
+    out_list: List[torch.Tensor],
+    a_list: List[torch.Tensor],
+    b_list: List[torch.Tensor],
+    block_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build the GPU-side dispatch tables for the packed kernel.
+
+    Returns
+    -------
+    out_ptrs, a_ptrs, b_ptrs : int64 [N_TENSORS]
+    sizes : int32 [N_TENSORS]
+    pid_to_tid : int32 [N_PROGRAMS]  -- block-id -> tensor-id
+    pid_to_local : int32 [N_PROGRAMS] -- block-id -> local-block-index
+        within the tensor it belongs to.
+    """
+    len(out_list)
+    device = out_list[0].device
+    out_ptrs = torch.tensor([t.data_ptr() for t in out_list], dtype=torch.int64, device=device)
+    a_ptrs = torch.tensor([t.data_ptr() for t in a_list], dtype=torch.int64, device=device)
+    b_ptrs = torch.tensor([t.data_ptr() for t in b_list], dtype=torch.int64, device=device)
+    sizes = torch.tensor([t.numel() for t in out_list], dtype=torch.int32, device=device)
+    # Build pid -> (tid, local) on the host then copy.
+    pid_to_tid: List[int] = []
+    pid_to_local: List[int] = []
+    for t_idx, t in enumerate(out_list):
+        n_chunks = (t.numel() + block_size - 1) // block_size
+        for local in range(n_chunks):
+            pid_to_tid.append(t_idx)
+            pid_to_local.append(local)
+    pid_to_tid_t = torch.tensor(pid_to_tid, dtype=torch.int32, device=device)
+    pid_to_local_t = torch.tensor(pid_to_local, dtype=torch.int32, device=device)
+    return out_ptrs, a_ptrs, b_ptrs, sizes, pid_to_tid_t, pid_to_local_t
+
+
+def multi_tensor_add_triton_packed(
+    out_list: List[torch.Tensor],
+    a_list: List[torch.Tensor],
+    b_list: List[torch.Tensor],
+    scale: float = 1.0,
+    block_size: int = 8192,
+) -> None:
+    """Compute ``out_i = a_i + scale * b_i`` for each ``i`` in place
+    (packed multi-tensor variant; ONE kernel launch total).
+
+    All tensors must be bf16, contiguous, on the same CUDA / HIP
+    device, with matching ``(out_i, a_i, b_i)`` shapes.
+    """
+    assert len(out_list) == len(a_list) == len(b_list) > 0
+    for out, a, b in zip(out_list, a_list, b_list):
+        assert out.dtype == torch.bfloat16
+        assert a.dtype == torch.bfloat16
+        assert b.dtype == torch.bfloat16
+        assert out.is_contiguous() and a.is_contiguous() and b.is_contiguous()
+        assert out.shape == a.shape == b.shape
+
+    out_ptrs, a_ptrs, b_ptrs, sizes, pid_to_tid, pid_to_local = _build_multi_tensor_dispatch_table(
+        out_list, a_list, b_list, block_size
+    )
+
+    grid = (pid_to_tid.numel(),)
+    _multi_tensor_add_packed_kernel[grid](
+        out_ptrs,
+        a_ptrs,
+        b_ptrs,
+        sizes,
+        pid_to_tid,
+        pid_to_local,
+        scale,
+        BLOCK_SIZE=block_size,
+    )
+
+
+__all__ = [
+    "multi_tensor_add_triton_per_tensor",
+    "multi_tensor_add_triton_packed",
+]
diff --git a/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py b/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py
new file mode 100644
index 000000000..0d3b883bc
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py
@@ -0,0 +1,343 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Fused FP8 cast Triton kernels (vendored into Primus).
+
+Two self-contained @triton_op kernels used by the FP8 local-spec path:
+
+  - ``cast_fp8_triton``           : fused FP8 cast + amax (no transpose), for the
+                                    native NN/TN arm. Single quantize + amax HBM
+                                    pass, no [N, M] transpose write.
+  - ``cast_transpose_fp8_triton`` : fused FP8 cast + transpose + amax, for the
+                                    forced-NT / delayed-scaling arm. One kernel
+                                    replaces a quantize kernel + .t().contiguous()
+                                    copy.
+
+Both follow TE's _cast_transpose_triton pattern with 2D grouped tiling.
+
+Vendored from Primus-Turbo ``integration/fp8-native-main`` @ 022f2b2b
+(``primus_turbo/pytorch/kernels/quantization/cast_fp8.py`` and
+``cast_transpose_fp8.py``). They are kept in Primus so the FP8 local-spec path
+builds against stock public Primus-Turbo ``main``, which never carried these
+kernels. The op namespaces were renamed ``primus_turbo::`` -> ``primus::`` to
+match the other Primus ``@triton_op``s.
+
+Registered as @triton_ops so that:
+  - torch.compile / Inductor can see through the op
+  - Output tensors are standard torch.empty (no triton.reinterpret metadata)
+  - register_fake provides correct shapes/strides for compile-time validation
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from torch.library import triton_op, wrap_triton
+
+
+def _fp8_max(dtype: torch.dtype) -> float:
+    return torch.finfo(dtype).max
+
+
+# ---------------------------------------------------------------------------
+# Cast + transpose + amax
+# ---------------------------------------------------------------------------
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 1}, num_warps=4),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 8}, num_warps=4),
+        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=8),
+    ],
+    key=["M", "N"],
+)
+@triton.jit
+def _cast_transpose_amax_kernel(
+    X_ptr,
+    C_ptr,
+    T_ptr,
+    stride_xm,
+    stride_xn,
+    stride_cm,
+    stride_cn,
+    stride_tm,
+    stride_tn,
+    M,
+    N,
+    scale_ptr,
+    amax_ptr,
+    scale_inv_ptr,
+    max_fp8: tl.constexpr,
+    COMPUTE_AMAX: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    scale = tl.load(scale_ptr)
+
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = tl.minimum(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // group_size
+
+    rm = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+
+    x_ptrs = X_ptr + rm[:, None] * stride_xm + rn[None, :] * stride_xn
+    a = tl.load(x_ptrs, mask=mask)
+    val = a.to(tl.float32)
+
+    scaled = val * scale
+    scaled = tl.clamp(scaled, -max_fp8, max_fp8)
+    fp8_val = scaled.to(C_ptr.dtype.element_ty)
+
+    c_ptrs = C_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
+    tl.store(c_ptrs, fp8_val, mask=mask)
+
+    # Transpose the tile so the store to T is coalesced (stride-1 in the fast dim).
+    # Without this, the transpose store scatters with stride=M which kills
+    # write throughput on MI355X under memory pressure.
+    fp8_val_t = tl.trans(fp8_val)
+    rn2 = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N)
+    rm2 = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
+    mask_t = (rn2 < N)[:, None] & (rm2 < M)[None, :]
+    t_ptrs = T_ptr + rn2[:, None] * stride_tm + rm2[None, :] * stride_tn
+    tl.store(t_ptrs, fp8_val_t, mask=mask_t)
+
+    if COMPUTE_AMAX:
+        tile_amax = tl.max(tl.abs(val))
+        tl.atomic_max(amax_ptr, tile_amax, sem="relaxed")
+
+    if pid == 0:
+        tl.store(scale_inv_ptr, tl.fdiv(1.0, scale))
+
+
+@triton_op("primus::cast_transpose_fp8_triton", mutates_args=())
+def cast_transpose_fp8_triton(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fused FP8 cast + transpose + optional amax.
+
+    Args:
+        x: 2D input tensor [M, N] (bf16 or f16), must be contiguous.
+        fp8_dtype: Target FP8 dtype (e.g. torch.float8_e4m3fn).
+        scale: Scalar float32 tensor with the quantization scale.
+        amax_out: Optional scalar float32 tensor. If provided, the kernel
+                  writes the abs-max of x into it (atomically reduced).
+
+    Returns:
+        (cast_out, transpose_out, scale_inv) where:
+          cast_out:      [M, N] FP8 tensor (same layout as input)
+          transpose_out: [N, M] FP8 tensor (contiguous transpose)
+          scale_inv:     scalar float32, 1/scale
+    """
+    if x.ndim != 2:
+        raise ValueError(f"Expected 2D input, got {x.ndim}D")
+    if not x.is_contiguous():
+        x = x.contiguous()
+
+    M, N = x.shape
+    max_fp8 = _fp8_max(fp8_dtype)
+
+    cast_out = torch.empty((M, N), dtype=fp8_dtype, device=x.device)
+    transpose_out = torch.empty((N, M), dtype=fp8_dtype, device=x.device)
+    scale_inv = torch.empty((), dtype=torch.float32, device=x.device)
+
+    compute_amax = amax_out is not None
+    if not compute_amax:
+        amax_out = torch.empty((), dtype=torch.float32, device=x.device)
+    else:
+        amax_out.zero_()
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    wrap_triton(_cast_transpose_amax_kernel)[grid](
+        x,
+        cast_out,
+        transpose_out,
+        x.stride(0),
+        x.stride(1),
+        cast_out.stride(0),
+        cast_out.stride(1),
+        transpose_out.stride(0),
+        transpose_out.stride(1),
+        M,
+        N,
+        scale,
+        amax_out,
+        scale_inv,
+        max_fp8,
+        compute_amax,
+    )
+
+    return cast_out, transpose_out, scale_inv
+
+
+@cast_transpose_fp8_triton.register_fake
+def _cast_transpose_fp8_triton_meta(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    M, N = x.shape
+    return (
+        torch.empty((M, N), dtype=fp8_dtype, device=x.device),
+        torch.empty((N, M), dtype=fp8_dtype, device=x.device),
+        torch.empty((), dtype=torch.float32, device=x.device),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Cast + amax (no transpose)
+# ---------------------------------------------------------------------------
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 1}, num_warps=4),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 8}, num_warps=4),
+        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=8),
+    ],
+    key=["M", "N"],
+)
+@triton.jit
+def _cast_amax_kernel(
+    X_ptr,
+    C_ptr,
+    stride_xm,
+    stride_xn,
+    stride_cm,
+    stride_cn,
+    M,
+    N,
+    scale_ptr,
+    amax_ptr,
+    scale_inv_ptr,
+    max_fp8: tl.constexpr,
+    COMPUTE_AMAX: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    scale = tl.load(scale_ptr)
+
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = tl.minimum(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // group_size
+
+    rm = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+
+    x_ptrs = X_ptr + rm[:, None] * stride_xm + rn[None, :] * stride_xn
+    a = tl.load(x_ptrs, mask=mask)
+    val = a.to(tl.float32)
+
+    scaled = val * scale
+    scaled = tl.clamp(scaled, -max_fp8, max_fp8)
+    fp8_val = scaled.to(C_ptr.dtype.element_ty)
+
+    c_ptrs = C_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
+    tl.store(c_ptrs, fp8_val, mask=mask)
+
+    if COMPUTE_AMAX:
+        tile_amax = tl.max(tl.abs(val))
+        tl.atomic_max(amax_ptr, tile_amax, sem="relaxed")
+
+    if pid == 0:
+        tl.store(scale_inv_ptr, tl.fdiv(1.0, scale))
+
+
+@triton_op("primus::cast_fp8_triton", mutates_args=())
+def cast_fp8_triton(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fused FP8 cast + optional amax (no transpose).
+
+    Args:
+        x: 2D input tensor [M, N] (bf16 or f16), must be contiguous.
+        fp8_dtype: Target FP8 dtype (e.g. torch.float8_e4m3fn).
+        scale: Scalar float32 tensor with the quantization scale.
+        amax_out: Optional scalar float32 tensor. If provided, the kernel
+                  writes the abs-max of the *unscaled* x into it (atomically).
+
+    Returns:
+        (cast_out, scale_inv) where:
+          cast_out:  [M, N] FP8 tensor (same layout as input)
+          scale_inv: scalar float32, 1 / scale
+    """
+    if x.ndim != 2:
+        raise ValueError(f"Expected 2D input, got {x.ndim}D")
+    if not x.is_contiguous():
+        x = x.contiguous()
+
+    M, N = x.shape
+    max_fp8 = _fp8_max(fp8_dtype)
+
+    cast_out = torch.empty((M, N), dtype=fp8_dtype, device=x.device)
+    scale_inv = torch.empty((), dtype=torch.float32, device=x.device)
+
+    compute_amax = amax_out is not None
+    if not compute_amax:
+        amax_out = torch.empty((), dtype=torch.float32, device=x.device)
+    else:
+        amax_out.zero_()
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    wrap_triton(_cast_amax_kernel)[grid](
+        x,
+        cast_out,
+        x.stride(0),
+        x.stride(1),
+        cast_out.stride(0),
+        cast_out.stride(1),
+        M,
+        N,
+        scale,
+        amax_out,
+        scale_inv,
+        max_fp8,
+        compute_amax,
+    )
+
+    return cast_out, scale_inv
+
+
+@cast_fp8_triton.register_fake
+def _cast_fp8_triton_meta(
+    x: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    M, N = x.shape
+    return (
+        torch.empty((M, N), dtype=fp8_dtype, device=x.device),
+        torch.empty((), dtype=torch.float32, device=x.device),
+    )
diff --git a/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py b/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py
new file mode 100644
index 000000000..d9778d352
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py
@@ -0,0 +1,453 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Runtime monkeypatch — fuse residual+add into RMSNorm at two sites.
+
+This is the runtime install hook (callable from a trainer entry point such
+as ``small_llm_moe_pretraining/primus/src/train.py``). The Triton RMSNorm
+kernels live in Primus-Turbo (``primus_turbo.pytorch.ops.normalization``,
+which exposes ``rmsnorm`` and the fused ``rmsnorm_residual`` variant).
+
+V1 — IN-LAYER fuse  (``PRIMUS_FUSED_RESIDUAL_NORM=1``)
+    self_attn_bda(residual + attn_out) → pre_mlp_layernorm
+    Kills 1 of 2 ``vectorized_elementwise_kernel<CUDAFunctor_add<bf16>>``
+    launches per layer (the in-layer ADD#1).
+
+V2 — CROSS-LAYER fuse  (``PRIMUS_FUSED_RESIDUAL_NORM_V2=1``, implies V1)
+    layer N: skip mlp_bda(mlp_out + residual_post_attn), instead stash
+             ``(mlp_out, residual_post_attn)`` as a "carry" on layer N+1
+    layer N+1: input_layernorm(carry) does the fused add+norm
+    last layer: carry goes to final_layernorm via ``_v2_pending_carry``.
+    Kills the remaining ADD#2 in the bf16 add tax (24 → 1 launches per step;
+    the last layer keeps its add when final_layernorm isn't a
+    ``PrimusTurboRMSNorm``, otherwise that one is fused too).
+
+Activation
+----------
+* V1: ``PRIMUS_FUSED_RESIDUAL_NORM=1`` (default 0).
+* V2: ``PRIMUS_FUSED_RESIDUAL_NORM_V2=1`` (default 0). Implies V1.
+* Requires ``use_turbo_rms_norm=true`` so the norms are ``PrimusTurboRMSNorm``
+  instances we can extend with a ``residual=`` arg.
+
+Falls back silently to the original ``TransformerLayer.forward`` whenever
+any precondition fails (recompute paths, fp32 residual, inference fused TP,
+non-zero hidden_dropout, cross-attention layers, non-PrimusTurboRMSNorm
+pre-norm).
+
+Safety
+------
+* Both gates default off; A/B is always reproducible.
+* On any exception in the fused path the patch reverts permanently and the
+  rest of the run uses the original Megatron forward.
+* V2 keeps the original mlp_bda when the next consumer cannot fuse (e.g.
+  final_layernorm is not a PrimusTurboRMSNorm, or the next layer doesn't
+  pass ``_can_fuse``), so partial coverage degrades gracefully.
+"""
+from __future__ import annotations
+
+import os
+import sys
+from typing import Any
+
+
+def _enabled() -> bool:
+    """V1 gate (or implied by V2)."""
+    if _v2_enabled():
+        return True
+    v = os.environ.get("PRIMUS_FUSED_RESIDUAL_NORM", "0").strip().lower()
+    return v in ("1", "true", "yes", "on")
+
+
+def _v2_enabled() -> bool:
+    v = os.environ.get("PRIMUS_FUSED_RESIDUAL_NORM_V2", "0").strip().lower()
+    return v in ("1", "true", "yes", "on")
+
+
+def _log(msg: str) -> None:
+    rank = os.environ.get("RANK", "0")
+    if rank == "0":
+        print(f"[fused_residual_rmsnorm] {msg}", file=sys.stderr, flush=True)
+
+
+_INSTALLED = False
+
+
+def install() -> bool:
+    """Install the monkeypatch. Returns True on success."""
+    global _INSTALLED
+    if _INSTALLED:
+        return True
+    if not _enabled():
+        _log("disabled (set PRIMUS_FUSED_RESIDUAL_NORM=1 or " "PRIMUS_FUSED_RESIDUAL_NORM_V2=1 to enable)")
+        return False
+
+    try:
+        from primus_turbo.pytorch.ops.normalization import (
+            rmsnorm_residual as triton_rmsnorm_residual,
+        )
+    except ImportError as exc:
+        _log(f"could not import primus_turbo rmsnorm_residual: {exc}; abort install")
+        return False
+
+    try:
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboRMSNorm,
+        )
+    except ImportError as exc:
+        _log(f"PrimusTurboRMSNorm not importable: {exc}; abort install")
+        return False
+
+    try:
+        from megatron.core.transformer.transformer_block import TransformerBlock
+        from megatron.core.transformer.transformer_layer import TransformerLayer
+    except ImportError as exc:
+        _log(f"Megatron transformer modules not importable: {exc}; abort install")
+        return False
+
+    # ----- 1. Extend PrimusTurboRMSNorm to accept residual ------------------
+    if not getattr(PrimusTurboRMSNorm, "_fused_residual_patched", False):
+        _orig_norm_forward = PrimusTurboRMSNorm.forward
+
+        def _new_norm_forward(self, x, residual=None):
+            # V2 path: an upstream layer stashed a carry for *this* norm to
+            # consume on its next forward call. The carry takes priority
+            # over an explicit ``residual=`` arg (which is V1-style and is
+            # only used when the layer's own _do_fused_forward already
+            # picked the right pair).
+            pending = getattr(self, "_v2_pending_carry", None)
+            if pending is not None:
+                self._v2_pending_carry = None
+                x_pending, r_pending = pending
+                gamma = self.weight
+                if getattr(self, "zero_centered_gamma", False):
+                    gamma = gamma + 1
+                norm_out, _xpr = triton_rmsnorm_residual(x_pending, r_pending, gamma, self.eps)
+                return norm_out
+            if residual is None:
+                return _orig_norm_forward(self, x)
+            gamma = self.weight
+            if getattr(self, "zero_centered_gamma", False):
+                gamma = gamma + 1
+            return triton_rmsnorm_residual(x, residual, gamma, self.eps)
+
+        PrimusTurboRMSNorm.forward = _new_norm_forward
+        PrimusTurboRMSNorm._fused_residual_patched = True
+        _log("patched PrimusTurboRMSNorm.forward (residual= arg + V2 _pending_carry)")
+
+    # ----- 2. Patch TransformerBlock.__init__ to wire V2 layer links --------
+    if _v2_enabled() and not getattr(TransformerBlock, "_fused_residual_v2_init_patched", False):
+        _orig_block_init = TransformerBlock.__init__
+
+        def _v2_init(self, *args, **kwargs):
+            _orig_block_init(self, *args, **kwargs)
+            try:
+                _wire_v2_layer_links(self)
+            except Exception as exc:
+                _log(f"V2 layer-link wiring raised {type(exc).__name__}: {exc}; " "block will run V1-only")
+
+        TransformerBlock.__init__ = _v2_init
+        TransformerBlock._fused_residual_v2_init_patched = True
+        _log("patched TransformerBlock.__init__ (V2 layer-link wiring active)")
+
+    # ----- 3. Replace TransformerLayer.forward with fused variant -----------
+    if not getattr(TransformerLayer, "_fused_residual_patched", False):
+        _orig_layer_forward = TransformerLayer.forward
+
+        def _fused_layer_forward(self, *args, **kwargs):
+            if not _can_fuse(self):
+                # Per-layer fallback. If a carry was queued for *this* layer
+                # but the layer itself can't fuse, drop it back to a real
+                # add so the layer sees a correct hidden_states.
+                _drain_carry_into_hidden_states(self, args, kwargs)
+                return _orig_layer_forward(self, *args, **kwargs)
+            try:
+                return _do_fused_forward(self, *args, **kwargs)
+            except Exception as exc:
+                _log(
+                    f"fused forward raised {type(exc).__name__}: {exc}; "
+                    f"falling back to original forward for all layers"
+                )
+                TransformerLayer.forward = _orig_layer_forward
+                return _orig_layer_forward(self, *args, **kwargs)
+
+        TransformerLayer.forward = _fused_layer_forward
+        TransformerLayer._fused_residual_patched = True
+        if _v2_enabled():
+            _log("patched TransformerLayer.forward (V1 ADD#1 + V2 cross-layer ADD#2 fusion active)")
+        else:
+            _log("patched TransformerLayer.forward (V1 in-layer ADD#1+norm fusion active)")
+
+    _INSTALLED = True
+    return True
+
+
+# ---------------------------------------------------------------------------
+# V2 layer-link wiring
+# ---------------------------------------------------------------------------
+def _wire_v2_layer_links(block: Any) -> None:
+    """Annotate each layer in ``block`` with V2 navigation pointers.
+
+    Sets per-layer attributes:
+      * ``_v2_block``           -> back-reference to the owning block
+      * ``_v2_next_layer``      -> the next TransformerLayer or None
+      * ``_v2_is_last_layer``   -> True for the final layer in the block
+      * ``_v2_final_layernorm`` -> reference to block.final_layernorm if any
+                                   (only set on the last layer)
+    """
+    layers = getattr(block, "layers", None)
+    if layers is None:
+        return
+    n = len(layers)
+    if n == 0:
+        return
+    final_ln = getattr(block, "final_layernorm", None)
+    # CRITICAL: nn.Module.__setattr__ auto-registers any nn.Module-valued
+    # attribute as a child module. Setting layer._v2_next_layer = next_layer
+    # would create cycles (each layer becomes a child of the previous one)
+    # and blow up tree-walking ops like .cuda() / .children() with
+    # RecursionError. Bypass with object.__setattr__ so these are plain
+    # Python attributes invisible to nn.Module bookkeeping.
+    for i, layer in enumerate(layers):
+        nxt = layers[i + 1] if (i + 1) < n else None
+        is_last = (i + 1) == n
+        object.__setattr__(layer, "_v2_block", block)
+        object.__setattr__(layer, "_v2_next_layer", nxt)
+        object.__setattr__(layer, "_v2_is_last_layer", is_last)
+        object.__setattr__(layer, "_v2_final_layernorm", final_ln if is_last else None)
+        object.__setattr__(layer, "_v2_carry", None)
+
+
+def _drain_carry_into_hidden_states(layer: Any, args: tuple, kwargs: dict) -> None:
+    """If a V2 carry was left for ``layer`` but we are about to take the
+    *original* (unpatched) forward, materialise the carry as a regular add
+    into ``hidden_states`` so the original forward still gets the right
+    activations.
+    """
+    carry = getattr(layer, "_v2_carry", None)
+    if carry is None:
+        return
+    layer._v2_carry = None
+    mlp_out, residual = carry
+    new_hs = mlp_out + residual
+    if args:
+        # python tuples are immutable; the caller (TransformerBlock loop)
+        # passes hidden_states via kwargs. We can't mutate the caller's
+        # args tuple from here, so route via kwargs (the kwargs path is
+        # always taken for hidden_states in current Megatron).
+        kwargs["hidden_states"] = new_hs
+        return
+    kwargs["hidden_states"] = new_hs
+
+
+# ---------------------------------------------------------------------------
+# Per-layer "can we fuse?" guard
+# ---------------------------------------------------------------------------
+def _can_fuse(layer: Any) -> bool:
+    """Return True when ``layer`` matches the assumptions of the fused path."""
+    cfg = getattr(layer, "config", None)
+    if cfg is None:
+        return False
+    if getattr(cfg, "fp32_residual_connection", False):
+        return False
+    if getattr(cfg, "inference_fuse_tp_communication", False):
+        return False
+    if getattr(layer, "recompute_input_layernorm", False):
+        return False
+    if getattr(layer, "recompute_pre_mlp_layernorm", False):
+        return False
+    if getattr(layer, "hidden_dropout", 0.0) != 0.0:
+        return False
+    if getattr(layer, "offload_attn_norm", False):
+        return False
+    if getattr(layer, "offload_mlp_norm", False):
+        return False
+
+    pre_mlp = getattr(layer, "pre_mlp_layernorm", None)
+    if pre_mlp is None:
+        return False
+    try:
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboRMSNorm,
+        )
+    except ImportError:
+        return False
+    if not isinstance(pre_mlp, PrimusTurboRMSNorm):
+        return False
+
+    from megatron.core.transformer.identity_op import IdentityOp
+
+    cross_attn = getattr(layer, "cross_attention", None)
+    if cross_attn is not None and not isinstance(cross_attn, IdentityOp):
+        return False
+
+    return True
+
+
+def _v2_next_can_consume_carry(layer: Any) -> bool:
+    """Return True when layer N can stash an unfused carry instead of
+    running mlp_bda. Requires that the next consumer (next layer's
+    input_layernorm OR this block's final_layernorm) is a
+    ``PrimusTurboRMSNorm`` so it can absorb the residual at no extra cost.
+    """
+    if not _v2_enabled():
+        return False
+
+    try:
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboRMSNorm,
+        )
+    except ImportError:
+        return False
+
+    is_last = getattr(layer, "_v2_is_last_layer", False)
+    if not is_last:
+        nxt = getattr(layer, "_v2_next_layer", None)
+        if nxt is None:
+            return False
+        nxt_in_ln = getattr(nxt, "input_layernorm", None)
+        if not isinstance(nxt_in_ln, PrimusTurboRMSNorm):
+            return False
+        # If the next layer can't fuse for its own reasons, we shouldn't
+        # leave a carry it has to drain — drain it ourselves via the
+        # _drain_carry_into_hidden_states fallback in the wrapper.
+        return _can_fuse(nxt)
+    # Last layer: route through final_layernorm if it's a PrimusTurboRMSNorm.
+    final_ln = getattr(layer, "_v2_final_layernorm", None)
+    if final_ln is None:
+        return False
+    return isinstance(final_ln, PrimusTurboRMSNorm)
+
+
+# ---------------------------------------------------------------------------
+# The fused forward
+# ---------------------------------------------------------------------------
+def _do_fused_forward(layer: Any, hidden_states=None, *args, **kwargs):
+    """Mirror of ``TransformerLayer.forward`` with V1 + V2 fusion paths.
+
+    V1: pre_mlp_layernorm receives ``(attn_out, residual)`` and emits both
+        the normed activation AND ``x_plus_r`` for the next bda.
+    V2: input_layernorm consumes a ``_v2_carry`` left by the previous
+        layer, and at exit either stashes a new carry (next layer / final
+        layernorm) or falls back to the explicit mlp_bda add.
+    """
+    from megatron.core.utils import (
+        deprecate_inference_params,
+        make_viewless_tensor,
+        nvtx_range_pop,
+        nvtx_range_push,
+    )
+
+    if hidden_states is None:
+        hidden_states = kwargs.pop("hidden_states")
+    else:
+        kwargs.pop("hidden_states", None)
+
+    inference_context = deprecate_inference_params(
+        kwargs.get("inference_context"), kwargs.get("inference_params")
+    )
+
+    v2_active = _v2_enabled()
+
+    # ---- input_layernorm (with optional V2 carry consume) ----------------
+    nvtx_range_push(suffix="input_layernorm")
+    carry = getattr(layer, "_v2_carry", None) if v2_active else None
+    if carry is not None:
+        # Fused path: input_layernorm absorbs the previous layer's deferred
+        # mlp_bda add. The returned x_plus_r becomes our residual base.
+        layer._v2_carry = None
+        prev_mlp_out, prev_residual = carry
+        # We bypass PrimusTurboRMSNorm.forward to access x_plus_r directly,
+        # which is needed as the residual for ADD#1.
+        from primus_turbo.pytorch.ops.normalization import (
+            rmsnorm_residual as triton_rmsnorm_residual,
+        )
+
+        in_ln = layer.input_layernorm
+        gamma = in_ln.weight
+        if getattr(in_ln, "zero_centered_gamma", False):
+            gamma = gamma + 1
+        input_layernorm_output, hidden_states = triton_rmsnorm_residual(
+            prev_mlp_out, prev_residual, gamma, in_ln.eps
+        )
+    else:
+        input_layernorm_output = layer.input_layernorm(hidden_states)
+    nvtx_range_pop(suffix="input_layernorm")
+
+    residual = hidden_states  # base for ADD#1 (post input_layernorm)
+
+    # ---- self attention --------------------------------------------------
+    nvtx_range_push(suffix="self_attention")
+    attention_output_with_bias = layer.self_attention(
+        input_layernorm_output,
+        attention_mask=kwargs.get("attention_mask"),
+        inference_context=inference_context,
+        rotary_pos_emb=kwargs.get("rotary_pos_emb"),
+        rotary_pos_cos=kwargs.get("rotary_pos_cos"),
+        rotary_pos_sin=kwargs.get("rotary_pos_sin"),
+        rotary_pos_cos_sin=kwargs.get("rotary_pos_cos_sin"),
+        attention_bias=kwargs.get("attention_bias"),
+        packed_seq_params=kwargs.get("packed_seq_params"),
+        sequence_len_offset=kwargs.get("sequence_len_offset"),
+    )
+    nvtx_range_pop(suffix="self_attention")
+
+    attn_out, attn_bias = attention_output_with_bias
+    if attn_bias is not None:
+        attn_out = attn_out + attn_bias.to(attn_out.dtype)
+
+    # ---- V1 fuse: bda(attn_out, residual) + pre_mlp_layernorm ------------
+    nvtx_range_push(suffix="fused_residual_pre_mlp_layernorm")
+    pre_mlp_layernorm_output, hidden_states = layer.pre_mlp_layernorm(attn_out, residual=residual)
+    nvtx_range_pop(suffix="fused_residual_pre_mlp_layernorm")
+    # hidden_states now == attn_out + residual (= residual_post_attn).
+
+    # ---- mlp -------------------------------------------------------------
+    padding_mask = kwargs.get("padding_mask", None)
+    try:
+        mlp_output_with_bias = layer.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
+    except TypeError:
+        mlp_output_with_bias = layer.mlp(pre_mlp_layernorm_output)
+
+    # ---- mlp_bda OR V2 carry stash ---------------------------------------
+    if v2_active and _v2_next_can_consume_carry(layer):
+        # Skip mlp_bda. Stash carry on the next consumer, return
+        # residual_post_attn as the layer's "output" tensor (used only by
+        # TransformerBlock for offload commit / make_viewless plumbing;
+        # the next layer / final_layernorm reads the carry instead).
+        mlp_out, mlp_bias = mlp_output_with_bias
+        if mlp_bias is not None:
+            mlp_out = mlp_out + mlp_bias.to(mlp_out.dtype)
+
+        is_last = getattr(layer, "_v2_is_last_layer", False)
+        if is_last:
+            final_ln = layer._v2_final_layernorm
+            final_ln._v2_pending_carry = (mlp_out, hidden_states)
+        else:
+            layer._v2_next_layer._v2_carry = (mlp_out, hidden_states)
+
+        nvtx_range_push(suffix="v2_skip_mlp_bda")
+        output = make_viewless_tensor(
+            inp=hidden_states,
+            requires_grad=hidden_states.requires_grad,
+            keep_graph=True,
+        )
+        nvtx_range_pop(suffix="v2_skip_mlp_bda")
+        return output, None
+
+    # V1-only path (or V2 last-layer-without-norm-final): explicit mlp_bda.
+    nvtx_range_push(suffix="mlp_bda")
+    with layer.bias_dropout_add_exec_handler():
+        hidden_states = layer.mlp_bda(layer.training, layer.config.bias_dropout_fusion)(
+            mlp_output_with_bias, hidden_states, layer.hidden_dropout
+        )
+    nvtx_range_pop(suffix="mlp_bda")
+
+    output = make_viewless_tensor(
+        inp=hidden_states,
+        requires_grad=hidden_states.requires_grad,
+        keep_graph=True,
+    )
+    return output, None
diff --git a/primus/backends/megatron/core/extensions/mega_moe.py b/primus/backends/megatron/core/extensions/mega_moe.py
new file mode 100644
index 000000000..5d0ab7a02
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/mega_moe.py
@@ -0,0 +1,266 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MegaMoE layer, drop-in for Megatron MoELayer (EP-only, bf16)."""
+
+import contextlib
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.random import (
+    get_cuda_rng_tracker,
+    get_expert_parallel_rng_tracker_name,
+)
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group
+from primus_turbo.pytorch.ops.moe.fused_mega_moe import (
+    fused_mega_moe_stage1,
+    fused_mega_moe_stage2,
+)
+
+
+class MegaMoEWeightModule(MegatronModule):
+    """Callable expert-weight module used as a DDP overlap boundary."""
+
+    def __init__(self, config: TransformerConfig, weight_shape) -> None:
+        super().__init__(config)
+        device = torch.device("cpu") if config.use_cpu_initialization else torch.cuda.current_device()
+        self.weight = torch.nn.Parameter(torch.empty(weight_shape, device=device, dtype=config.params_dtype))
+
+    def forward(self) -> torch.Tensor:
+        return self.weight
+
+    def backward_dw(self) -> None:
+        # wgrad produced in custom autograd backward
+        return None
+
+
+class MegaMoEExperts(MegatronModule):
+    """Two-stage expert with separately wrapped w1/w2 parameters."""
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        experts_per_rank: int,
+        hidden_size: int,
+        intermediate_size: int,
+        ep_group,
+    ) -> None:
+        super().__init__(config)
+        self.ep_group = ep_group
+        self.experts_per_rank = experts_per_rank
+        # w1 [g, 2I, H] gate+up; w2 [g, H, I] down
+        self.fc1_weight = MegaMoEWeightModule(config, (experts_per_rank, 2 * intermediate_size, hidden_size))
+        self.fc2_weight = MegaMoEWeightModule(config, (experts_per_rank, hidden_size, intermediate_size))
+        # Expert weights are already EP-sharded, so the DP allreduce hook must
+        # not touch them.
+        expert_parallel = config.expert_model_parallel_size > 1
+        for p in (self.fc1_weight.weight, self.fc2_weight.weight):
+            setattr(p, "allreduce", not expert_parallel)
+
+    def reset_parameters(self, ep_rank: int) -> None:
+        """Init expert weights exactly like Megatron TEGroupedLinear."""
+        init_fc1 = self.config.init_method
+        init_fc2 = self.config.output_layer_init_method
+        assert init_fc1 is not None and init_fc2 is not None, "config init methods are unset"
+        # fc1 <- init_method, fc2 <- output_layer_init_method
+        weights = (
+            (self.fc1_weight.weight, init_fc1),
+            (self.fc2_weight.weight, init_fc2),
+        )
+
+        if self.config.use_cpu_initialization:
+            # cpu rng is rank-identical: draw all experts, keep this rank's shard
+            first_expert = ep_rank * self.experts_per_rank
+            for weight, init_method in weights:
+                master = torch.empty(weight.shape[1:], dtype=torch.float32)
+                for e in range(self.config.num_moe_experts):
+                    init_method(master)
+                    if first_expert <= e < first_expert + self.experts_per_rank:
+                        weight.data[e - first_expert].copy_(master)
+            return
+
+        tracker = get_cuda_rng_tracker()
+        rng_fork = (
+            tracker.fork(get_expert_parallel_rng_tracker_name())
+            if tracker.is_initialized()
+            else contextlib.nullcontext()
+        )
+        with rng_fork:
+            # init each expert as a 2D [out, in] slice, like TE per-gemm weights
+            for weight, init_method in weights:
+                for expert_weight in weight.data:
+                    init_method(expert_weight)
+
+    def forward(self, x, topk_idx, topk_weights):
+        w1 = self.fc1_weight()
+        l1_out, dwib, handle = fused_mega_moe_stage1(x, topk_idx, topk_weights, w1, self.ep_group)
+        w2 = self.fc2_weight()
+        return fused_mega_moe_stage2(l1_out, dwib, handle, topk_idx, topk_weights, w2, self.ep_group)
+
+    def backward_dw(self) -> None:
+        # match native fc2-then-fc1 order
+        self.fc2_weight.backward_dw()
+        self.fc1_weight.backward_dw()
+
+
+class PrimusTurboMegaMoELayer(MegatronModule):
+    """EP MoE layer: Megatron router -> two-stage experts -> shared expert."""
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: Optional[object] = None,
+        layer_number: Optional[int] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
+    ) -> None:
+        super().__init__(config)
+
+        assert (
+            pg_collection is not None and pg_collection.ep is not None
+        ), "MegaMoE requires an expert-parallel process group"
+        assert submodules is not None, "MegaMoE requires MoESubmodules (router/shared_experts)"
+        self._assert_supported_config(config)
+
+        self.config = config
+        self.layer_number = layer_number
+        self.is_mtp_layer = is_mtp_layer
+
+        # experts sharded evenly across EP group
+        self.ep_group = pg_collection.ep
+        self.ep_size = self.ep_group.size()
+        self.ep_rank = self.ep_group.rank()
+        assert config.num_moe_experts % self.ep_size == 0, "num_experts must be divisible by EP size"
+        self.experts_per_rank = config.num_moe_experts // self.ep_size
+
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.moe_ffn_hidden_size
+
+        # router owns all routing logic
+        self.router = submodules.router(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
+
+        # separate w1/w2 modules give DDP overlap boundaries
+        self.experts = MegaMoEExperts(
+            config,
+            self.experts_per_rank,
+            self.hidden_size,
+            self.intermediate_size,
+            self.ep_group,
+        )
+        if config.perform_initialization:
+            self.reset_parameters()
+
+        # optional shared expert
+        self.use_shared_expert = config.moe_shared_expert_intermediate_size is not None
+        if self.use_shared_expert:
+            self.shared_experts = build_module(
+                submodules.shared_experts,
+                config=config,
+                pg_collection=pg_collection,
+                gate=config.moe_shared_expert_gate,
+            )
+        else:
+            self.shared_experts = None
+
+    def reset_parameters(self) -> None:
+        """Init expert weights exactly like Megatron TEGroupedLinear."""
+        self.experts.reset_parameters(self.ep_rank)
+
+    @staticmethod
+    def _assert_supported_config(config: TransformerConfig) -> None:
+        """Only kernel-level constraints; routing features are handled by the router."""
+        assert config.tensor_model_parallel_size == 1, "MegaMoE is EP-only (TP==1)"
+        assert config.params_dtype == torch.bfloat16, "MegaMoE only supports bf16"
+        assert config.gated_linear_unit, "MegaMoE hardcodes a gated SwiGLU MLP"
+        assert config.activation_func in (F.silu, torch.nn.SiLU), "MegaMoE hardcodes SiLU"
+        assert (
+            config.moe_expert_capacity_factor is None
+        ), "MegaMoE is dropless; moe_expert_capacity_factor must be None"
+        assert not config.add_bias_linear, "MegaMoE fused expert has no bias; set add_bias_linear=False"
+        assert not config.init_model_with_meta_device, "MegaMoE does not support meta-device init"
+        load_balancing_type = config.moe_router_load_balancing_type
+        load_balancing_types = (
+            load_balancing_type if isinstance(load_balancing_type, (list, tuple)) else [load_balancing_type]
+        )
+        assert "sinkhorn" not in load_balancing_types, "MegaMoE does not support sinkhorn load balancing"
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        in_shape = hidden_states.shape
+        # transpose [bsz, seq] -> [seq, bsz]
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
+
+        # dense probs [T, E], aux-loss grad baked in
+        probs, _ = self.router(hidden_states, padding_mask)
+        probs = probs.reshape(-1, self.config.num_moe_experts)
+
+        # dense -> sparse top-k
+        topk_weights, topk_idx = probs.topk(self.router.topk, dim=-1)
+
+        x = hidden_states.reshape(-1, self.hidden_size).to(torch.bfloat16)
+        y = self.experts(
+            x,
+            topk_idx,
+            topk_weights.to(torch.float32),
+        )
+        y = y.reshape(in_shape).to(hidden_states.dtype)
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(hidden_states)
+        return y, None
+
+    def sharded_state_dict(
+        self,
+        prefix: str = "",
+        sharded_offsets: Tuple[Tuple[int, int, int], ...] = (),
+        metadata: Optional[dict] = None,
+    ):
+        # delegates below need metadata["dp_cp_group"]
+        metadata = ensure_metadata_has_dp_cp_group(metadata)
+        prepend_axis_num = len(sharded_offsets)
+        edp_rank = parallel_state.get_expert_data_parallel_rank()
+        expert_replica_id = (0, 0, edp_rank)
+
+        sharded_sd: dict = {}
+        # expert weights: EP-sharded on axis 0
+        for name, weight in (
+            ("fc1_weight", self.experts.fc1_weight.weight),
+            ("fc2_weight", self.experts.fc2_weight.weight),
+        ):
+            key = f"{prefix}experts.{name}.weight"
+            sharded_sd[key] = ShardedTensor.from_rank_offsets(
+                key,
+                weight,
+                *sharded_offsets,
+                (prepend_axis_num, self.ep_rank, self.ep_size),
+                replica_id=expert_replica_id,
+                prepend_axis_num=prepend_axis_num,
+            )
+        # router + shared expert: delegate to their own sharded_state_dict
+        sharded_sd.update(self.router.sharded_state_dict(f"{prefix}router.", sharded_offsets, metadata))
+        if self.shared_experts is not None:
+            sharded_sd.update(
+                self.shared_experts.sharded_state_dict(f"{prefix}shared_experts.", sharded_offsets, metadata)
+            )
+        return sharded_sd
+
+    def set_layer_number(self, layer_number: int) -> None:
+        self.layer_number = layer_number
+        self.router.set_layer_number(layer_number)
+
+    def backward_dw(self, *args: object, **kwargs: object) -> None:
+        return None
diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py
index 8b787044d..ffa81e9b4 100644
--- a/primus/backends/megatron/core/extensions/primus_turbo.py
+++ b/primus/backends/megatron/core/extensions/primus_turbo.py
@@ -1,10 +1,11 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 import gc
 from contextlib import contextmanager, nullcontext
+from functools import lru_cache
 from typing import Callable, Iterable, List, Optional, Tuple, Union
 
 import primus_turbo.pytorch as primus_turbo_torch
@@ -23,6 +24,7 @@
     condition_init_method,
 )
 from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import (
     get_context_parallel_group,
@@ -77,14 +79,44 @@
     float8_e4m3,
 )
 from torch import Tensor
+
+# Imported from .constants (not .fp8) for TransformerEngine >= 2.12 compat;
+# the symbol moved out of transformer_engine.pytorch.fp8 in that release.
 from transformer_engine.pytorch.constants import dist_group_type
-from transformer_engine.pytorch.fp8 import DelayedScaling, FP8GlobalStateManager, Recipe
+from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, Recipe
 
 from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER
 
+try:
+    pass
+
+    _HAVE_TRITON = True
+except (ImportError, ModuleNotFoundError):
+    _HAVE_TRITON = False
+
+from primus.backends.megatron.core.extensions._triton.inplace_add import (
+    inplace_add_triton_,
+)
+
 _dummy_wgrads = {}
 
 
+@lru_cache(maxsize=1)
+def _is_gfx1250() -> bool:
+    """Return True iff the local GPU is a gfx1250 device.
+
+    Used to route ``main_grad`` accumulation through the single-launch Triton
+    in-place add (:func:`inplace_add_triton_`) only on gfx1250, keeping Torch's
+    ``add_`` on every other architecture.
+    """
+    if not torch.cuda.is_available():
+        return False
+    try:
+        return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:
+        return False
+
+
 def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor:
     """Returns a dummy tensor of given shape.
 
@@ -136,7 +168,10 @@ def backward(ctx, grad_x, grad_quantized_weight, grad_quantized_weight_trans):
                 weight, "grad_added_to_main_grad"
             ), "weight.grad_added_to_main_grad don't have grad_added_to_main_grad attribute."
 
-            weight.main_grad.add_(grad_quantized_weight)
+            if _is_gfx1250():
+                inplace_add_triton_(weight.main_grad, grad_quantized_weight)
+            else:
+                weight.main_grad.add_(grad_quantized_weight)
             weight.grad_added_to_main_grad = True
 
             return grad_x, _get_dummy_wgrad(list(weight.shape), weight.dtype), None, None
@@ -310,6 +345,7 @@ def __init__(
         strategy: ScalingStrategy = ScalingStrategy.DYNAMIC,
         scale_dtype: ScaleDtype = ScaleDtype.FP32,
         block_size: int = None,
+        use_gradient_sr: bool = True,
     ):
         self._is_fp4 = False
         self._is_fp8 = False
@@ -321,6 +357,7 @@ def __init__(
                 strategy=strategy,
                 scale_dtype=scale_dtype,
                 block_size=block_size,
+                use_gradient_sr=use_gradient_sr,
             )
             self._is_fp4 = True
         else:
@@ -440,7 +477,7 @@ def get_turbo_quant_config(cls) -> PrimusTurboQuantConfig:
     @classmethod
     def get_fp8_autocast_state(
         cls,
-    ) -> Tuple[bool, bool, Recipe, dist_group_type, bool, bool, PrimusTurboQuantConfig]:
+    ) -> Tuple[bool, bool, Recipe, dist_group_type, bool, bool, bool, bool, PrimusTurboQuantConfig]:
         """FP8 autocast state getter"""
         return (
             FP8GlobalStateManager.FP8_ENABLED,
@@ -457,7 +494,7 @@ def get_fp8_autocast_state(
     @classmethod
     def set_fp8_autocast_state(
         cls,
-        fp8_state: Tuple[bool, bool, DelayedScaling, dist_group_type, bool, bool, PrimusTurboQuantConfig],
+        fp8_state: Tuple[bool, bool, Recipe, dist_group_type, bool, bool, bool, bool, PrimusTurboQuantConfig],
     ) -> None:
         """FP8 autocast state setter"""
         (
@@ -964,60 +1001,83 @@ def forward_internal(
                     or quant_config.block_scaling()
                 ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
                         weight,
-                        float8_e4m3,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float8_e4m3,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                            or quant_config.current_scaling(),
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp8(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
                 quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
                 assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
 
-                if is_first_microbatch:
-                    need_cache_colwise = not self.disable_parameter_transpose_cache
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
-                        weight, float4_e2m1fn_x2, quant_config, need_cache_colwise=need_cache_colwise
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
+                        weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float4_e2m1fn_x2,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp4(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             else:
                 out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None)
 
@@ -1133,62 +1193,83 @@ def forward_internal(
                     or quant_config.block_scaling()
                 ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
                         weight,
-                        float8_e4m3,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float8_e4m3,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                            or quant_config.current_scaling(),
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp8(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
                 quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
                 assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
                         weight,
-                        float4_e2m1fn_x2,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float4_e2m1fn_x2,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp4(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             else:
                 out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None)
 
@@ -1297,62 +1378,83 @@ def forward_internal(
                     or quant_config.block_scaling()
                 ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
                         weight,
-                        float8_e4m3,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float8_e4m3,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                            or quant_config.current_scaling(),
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp8(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
                 quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
                 assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
                         weight,
-                        float4_e2m1fn_x2,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float4_e2m1fn_x2,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        )
+
+                    x, quantized_weight = _bridge_weight_grad(
+                        x,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        x,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                x, quantized_weight = _bridge_weight_grad(
-                    x,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp4(
-                    x,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             else:
                 out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None)
 
@@ -1474,62 +1576,83 @@ def forward_internal(self, x, is_first_microbatch: bool = False):
                     or quant_config.block_scaling()
                 ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        inp,
                         weight,
-                        float8_e4m3,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float8_e4m3,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                            or quant_config.current_scaling(),
+                        )
+
+                    inp, quantized_weight = _bridge_weight_grad(
+                        inp,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp8(
+                        inp,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                inp, quantized_weight = _bridge_weight_grad(
-                    inp,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp8(
-                    inp,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
                 quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
                 assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
 
-                if is_first_microbatch:
-                    (
-                        self.quantized_weight_buffer,
-                        self.quantized_weight_t_buffer,
-                    ) = _maybe_create_quantized_weight_buffers(
+                if get_num_microbatches() == 1:
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        inp,
                         weight,
-                        float4_e2m1fn_x2,
-                        quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
+                    )
+                else:
+                    if is_first_microbatch:
+                        (
+                            self.quantized_weight_buffer,
+                            self.quantized_weight_t_buffer,
+                        ) = _maybe_create_quantized_weight_buffers(
+                            weight,
+                            float4_e2m1fn_x2,
+                            quant_config,
+                            disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        )
+
+                    inp, quantized_weight = _bridge_weight_grad(
+                        inp,
+                        weight,
+                        PrimusTurboQuantizedTensorPair(
+                            data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                        ),
+                    )
+                    out = primus_turbo_torch.ops.gemm_fp4(
+                        inp,
+                        quantized_weight,
+                        trans_a=False,
+                        trans_b=True,
+                        out_dtype=None,
+                        config=quant_config.data(),
                     )
-
-                inp, quantized_weight = _bridge_weight_grad(
-                    inp,
-                    weight,
-                    PrimusTurboQuantizedTensorPair(
-                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                    ),
-                )
-                out = primus_turbo_torch.ops.gemm_fp4(
-                    inp,
-                    quantized_weight,
-                    trans_a=False,
-                    trans_b=True,
-                    out_dtype=None,
-                    config=quant_config.data(),
-                )
             else:
                 out = primus_turbo_torch.ops.gemm(inp, weight, trans_a=False, trans_b=True, out_dtype=None)
 
@@ -1541,6 +1664,129 @@ def forward_internal(self, x, is_first_microbatch: bool = False):
         return out, None
 
 
+def _make_primus_turbo_norm_te_column_parallel_linear():
+    """Build a layernorm+linear module that keeps the GEMM on Transformer Engine.
+
+    ``PrimusTurboLayerNormColumnParallelLinear`` above routes both the norm and
+    the GEMM through Turbo. On the GPT-OSS-20B shapes only the norm is worth
+    moving: Turbo's Triton rmsnorm is markedly cheaper than TE's ``general``
+    kernels (measured 21.3 ms/step against 26.9 with the dgamma fix, and 38.5
+    without it), while Turbo's dense FP8 GEMM is slower than hipBLASLt here
+    (69.6 ms/step against 59.7).
+
+    TE fuses norm and GEMM inside one autograd function, so there is no way to
+    substitute just the norm. This composes ``PrimusTurboRMSNorm`` with
+    ``TEColumnParallelLinear`` instead, trading the fused launch for the faster
+    norm. Built in a factory so the TE import stays lazy.
+    """
+    import torch.nn as nn
+    from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelLinear,
+        _get_extra_te_kwargs,
+    )
+
+    class PrimusTurboNormTEColumnParallelLinear(nn.Module):
+        """See ``_make_primus_turbo_norm_te_column_parallel_linear``."""
+
+        def __init__(
+            self,
+            input_size: int,
+            output_size: int,
+            *,
+            config: TransformerConfig,
+            init_method: Callable,
+            gather_output: bool,
+            bias: bool,
+            skip_bias_add: bool,
+            is_expert: bool,
+            skip_weight_param_allocation: bool = False,
+            tp_comm_buffer_name: Optional[str] = None,
+            tp_group: Optional[torch.distributed.ProcessGroup] = None,
+            stride: int = 1,
+        ):
+            super().__init__()
+            self.config = config
+            # Mirror the kwargs TELayerNormColumnParallelLinear feeds into
+            # te.pytorch.LayerNormLinear, so the norm behaves identically.
+            extra = _get_extra_te_kwargs(config)
+            self.layernorm = PrimusTurboRMSNorm(
+                normalized_shape=input_size,
+                eps=config.layernorm_epsilon,
+                sequence_parallel=config.sequence_parallel,
+                zero_centered_gamma=config.layernorm_zero_centered_gamma,
+                params_dtype=extra.get("params_dtype", torch.float32),
+                device=extra.get("device", torch.cuda.current_device()),
+            )
+            self.linear = TEColumnParallelLinear(
+                input_size=input_size,
+                output_size=output_size,
+                config=config,
+                init_method=init_method,
+                gather_output=gather_output,
+                bias=bias,
+                skip_bias_add=skip_bias_add,
+                is_expert=is_expert,
+                skip_weight_param_allocation=skip_weight_param_allocation,
+                tp_comm_buffer_name=tp_comm_buffer_name,
+                tp_group=tp_group,
+            )
+
+        # Callers reach for these names on the fused TE module directly
+        # (te_op_fuser, modelopt state-dict hooks, checkpoint conversion), so
+        # expose them rather than forcing every caller to know the layout.
+        @property
+        def layer_norm_weight(self):
+            return self.layernorm.weight
+
+        @property
+        def normalization(self):
+            return self.config.normalization
+
+        @property
+        def weight(self):
+            return self.linear.weight
+
+        @property
+        def bias(self):
+            return getattr(self.linear, "bias", None)
+
+        @property
+        def in_features(self):
+            return self.linear.in_features
+
+        @property
+        def out_features(self):
+            return self.linear.out_features
+
+        def forward(self, x):
+            return self.linear(self.layernorm(x))
+
+        def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+            """Flatten to the names a checkpoint from the fused class carries."""
+            sd = {}
+            sd.update(
+                self.linear.sharded_state_dict(
+                    prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata
+                )
+            )
+            norm_sd = self.layernorm.state_dict(prefix="", keep_vars=True)
+            sd.update(
+                make_sharded_tensors_for_checkpoint(
+                    {"layer_norm_weight": norm_sd["weight"]}, prefix, {}, sharded_offsets
+                )
+            )
+            return sd
+
+        def __repr__(self):
+            return (
+                f"{type(self).__name__}(in_features={self.in_features}, "
+                f"out_features={self.out_features}, "
+                f"norm={type(self.layernorm).__name__})"
+            )
+
+    return PrimusTurboNormTEColumnParallelLinear
+
+
 def fused_bias_act_with_probs(
     intermediate_parallel: torch.Tensor,
     bias_parallel: torch.Tensor,
@@ -1625,9 +1871,7 @@ def __init__(
                 weight = getattr(self, f"weight{i}")
                 buffer[i].copy_(weight)
 
-            weights = buffer.clone()
-
-        self.register_parameter("weights", torch.nn.Parameter(weights))
+        self.register_parameter("weights", torch.nn.Parameter(buffer))
 
         # Capture the per-expert weights' extra attributes BEFORE deleting them.
         saved_weight_attrs = [dict(getattr(self, f"weight{i}").__dict__) for i in range(self.num_gemms)]
@@ -1643,31 +1887,41 @@ def __init__(
             name = f"weight{i}"
             if name in self._parameters:
                 del self._parameters[name]
-        del buffer
 
         gc.collect()
         torch.cuda.empty_cache()
 
-        # Re-expose each expert's slice as a zero-copy weight{i} Parameter view
-        # of self.weights, so existing code paths and checkpoints that look up
-        # weight{i} by name keep working without allocating a new buffer.
-        # ``requires_grad=False`` is required: self.weights is the canonical
-        # trainable Parameter and these views share its storage, so leaving
-        # them trainable would make the optimizer (and DDP) update / sync the
-        # same memory twice.  ``.detach()`` strips the view's autograd graph
-        # so each Parameter ends up as a leaf with ``_base is None`` (which is
-        # what Megatron's distributed-optimizer param-bucket re-mapping
-        # expects), while still aliasing the same underlying storage.
-        # We also restore each weight{i}'s saved extra attributes so checkpoint /
-        # state-dict code that inspects them keeps seeing the right markers.
+        # Defer weight{i} view registration until after DDP has remapped
+        # self.weights into the distributed-optimizer param buffer. Registering
+        # views here would pin the pre-remap storage and leave a duplicate copy
+        # of the consolidated weights resident on GPU.
+        self._saved_weight_attrs = saved_weight_attrs
+        self._weight_views_registered = False
+        self.register_forward_pre_hook(self._forward_pre_hook_ensure_weight_views)
+
+        self.register_buffer("quantized_weight_buffer", None, persistent=False)
+        self.register_buffer("quantized_weight_t_buffer", None, persistent=False)
+
+    def _ensure_weight_views(self) -> None:
+        """Register per-expert weight{i} views after DDP param-buffer remap."""
+        if self._weight_views_registered:
+            return
+
         for i in range(self.num_gemms):
-            weight_i = torch.nn.Parameter(self.weights[i].detach(), requires_grad=False)
-            for attr_name, attr_val in saved_weight_attrs[i].items():
+            weight_i = torch.nn.Parameter(self.weights[i], requires_grad=False)
+            for attr_name, attr_val in self._saved_weight_attrs[i].items():
                 setattr(weight_i, attr_name, attr_val)
             self.register_parameter(f"weight{i}", weight_i)
 
-        self.register_buffer("quantized_weight_buffer", None, persistent=False)
-        self.register_buffer("quantized_weight_t_buffer", None, persistent=False)
+        self._weight_views_registered = True
+
+    @staticmethod
+    def _forward_pre_hook_ensure_weight_views(module, _inputs):
+        module._ensure_weight_views()
+
+    def state_dict(self, *args, **kwargs):
+        self._ensure_weight_views()
+        return super().state_dict(*args, **kwargs)
 
     def forward(self, x: torch.Tensor, m_splits: torch.Tensor):
         _is_first_microbatch = self.is_first_microbatch
@@ -1689,7 +1943,8 @@ def forward_internal(
         """Forward step of the legacy PrimusTurbo grouped-gemm MLP."""
         weights = self.weights
         # NOTE: keep x and m_splits on the same device
-        m_splits = m_splits.to(x.device)
+        if m_splits.device != x.device:
+            m_splits = m_splits.to(x.device)
 
         if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled():
             quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
@@ -1697,34 +1952,81 @@ def forward_internal(
                 quant_config.mxfp8_scaling() or quant_config.current_scaling() or quant_config.block_scaling()
             ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling."
 
-            if is_first_microbatch:
-                (
-                    self.quantized_weight_buffer,
-                    self.quantized_weight_t_buffer,
-                ) = _maybe_create_quantized_weight_buffers(
+            if get_num_microbatches() == 1:
+                out = primus_turbo_torch.ops.grouped_gemm_fp8(
+                    x,
                     weights,
-                    float8_e4m3,
-                    quant_config,
-                    disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                    m_splits,
+                    trans_b=True,
+                    config=quant_config.data(),
                 )
+            else:
+                if is_first_microbatch:
+                    (
+                        self.quantized_weight_buffer,
+                        self.quantized_weight_t_buffer,
+                    ) = _maybe_create_quantized_weight_buffers(
+                        weights,
+                        float8_e4m3,
+                        quant_config,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                        or quant_config.current_scaling(),
+                    )
 
-            x, quantized_weights = _bridge_weight_grad(
-                x,
-                weights,
-                PrimusTurboQuantizedTensorPair(
-                    data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
-                ),
-            )
+                x, quantized_weights = _bridge_weight_grad(
+                    x,
+                    weights,
+                    PrimusTurboQuantizedTensorPair(
+                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                    ),
+                )
 
-            out = primus_turbo_torch.ops.grouped_gemm_fp8(
-                x,
-                quantized_weights,
-                m_splits,
-                trans_b=True,
-                config=quant_config.data(),
-            )
+                out = primus_turbo_torch.ops.grouped_gemm_fp8(
+                    x,
+                    quantized_weights,
+                    m_splits,
+                    trans_b=True,
+                    config=quant_config.data(),
+                )
         elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
-            assert False, "FP4 is not supported in PrimusTurboGroupedLinear"
+            quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
+            assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
+
+            if get_num_microbatches() == 1:
+                out = primus_turbo_torch.ops.grouped_gemm_fp4(
+                    x,
+                    weights,
+                    m_splits,
+                    trans_b=True,
+                    config=quant_config.data(),
+                )
+            else:
+                if is_first_microbatch:
+                    (
+                        self.quantized_weight_buffer,
+                        self.quantized_weight_t_buffer,
+                    ) = _maybe_create_quantized_weight_buffers(
+                        weights,
+                        float4_e2m1fn_x2,
+                        quant_config,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                    )
+
+                x, quantized_weights = _bridge_weight_grad(
+                    x,
+                    weights,
+                    PrimusTurboQuantizedTensorPair(
+                        data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                    ),
+                )
+
+                out = primus_turbo_torch.ops.grouped_gemm_fp4(
+                    x,
+                    quantized_weights,
+                    m_splits,
+                    trans_b=True,
+                    config=quant_config.data(),
+                )
         else:
             out = primus_turbo_torch.ops.grouped_gemm(x, weights, m_splits, trans_b=True)
 
@@ -1820,7 +2122,7 @@ def __init__(
         pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         """
-        Initialize the Flex token dispatcher.
+        Initialize the DeepEP token dispatcher.
 
         Args:
             num_local_experts (int): Number of local experts on the current device.
@@ -1830,13 +2132,14 @@ def __init__(
         """
         super().__init__(config=config, pg_collection=pg_collection)
 
-        assert self.tp_size * self.ep_size > 1, "Flex token dispatcher requires TPxEP > 1"
+        if self.tp_size * self.ep_size <= 1:
+            raise ValueError("DeepEP token dispatcher requires TPxEP > 1")
         assert (
             self.config.moe_enable_deepep
         ), "DeepEP is not enabled. Please set --moe-enable-deepep to use DeepEP backend."
         assert (
             self.config.moe_pad_expert_input_to_capacity is False
-        ), "Flex token dispatcher does not support --moe-pad-expert-input-to-capacity"
+        ), "DeepEP token dispatcher does not support --moe-pad-expert-input-to-capacity"
 
         args = get_args()
 
@@ -1888,6 +2191,12 @@ def __init__(
         self._comm_manager = self.deepep_dispatcher
 
         self.moe_router_force_load_balancing = args.moe_router_force_load_balancing
+        # "even" -> deterministic round-robin token assignment (constant per-expert
+        # counts / constant M_total); "uniform" -> Megatron-LM original random-logits
+        # balancing (handled upstream in the router, token_indices stays None here).
+        self.moe_router_force_load_balancing_type = getattr(
+            args, "moe_router_force_load_balancing_type", "uniform"
+        )
 
     def dispatch_preprocess(
         self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor
@@ -1911,9 +2220,12 @@ def dispatch_preprocess(
         hidden_states = hidden_states.view(-1, self.config.hidden_size)
         num_tokens = hidden_states.shape[0]
 
-        # when force_load_balancing, we use even token_indices to make sure each expert get same number of tokens
+        # when force_load_balancing with type "even", we use round-robin token_indices
+        # to make sure each expert gets the same (deterministic) number of tokens.
+        # type "uniform" keeps token_indices=None and relies on the router's upstream
+        # random-logits balancing.
         token_indices = None
-        if self.moe_router_force_load_balancing:
+        if self.moe_router_force_load_balancing and self.moe_router_force_load_balancing_type == "even":
             token_indices = (
                 torch.arange(num_tokens * self.config.moe_router_topk, device=hidden_states.device).view(
                     num_tokens, self.config.moe_router_topk
diff --git a/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py b/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py
new file mode 100644
index 000000000..23cb767f0
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py
@@ -0,0 +1,1716 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Compile-friendly FP8 linear layers for Megatron local spec.
+
+Self-contained autograd Functions that call Primus Turbo's low-level building
+blocks (quantize_fp8 + gemm_fp8_impl) directly, bypassing the higher-level
+pt.ops.gemm_fp8 wrappers entirely.
+
+Key properties:
+- Tensorwise uses the setup_context pattern with primitive-only args so
+  torch.compile can trace through without graph breaks. FP8 weight data is
+  extracted by the caller (Float8*ParallelLinear._forward_impl) to avoid
+  tensor subclass tracing inside the autograd.Function.
+- Rowwise and blockwise still use @allow_in_graph (separate feasibility work).
+- gemm_fp8_impl is already a torch.library.custom_op with register_fake
+- Zero TransformerEngine dependencies
+- Builds against stock public Primus-Turbo main: the only kernels not on main
+  (the fused FP8 cast / cast+transpose Triton ops) are vendored locally in
+  fp8_cast_kernels_triton.py; everything else still comes from Turbo main.
+- Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from megatron.core.enums import Fp8Recipe
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from primus_turbo.pytorch.core.backend import BackendType
+from primus_turbo.pytorch.core.low_precision import (
+    Float8QuantConfig,
+    Format,
+    ScaleDtype,
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl
+from primus_turbo.pytorch.kernels.quantization.quantization_impl import (
+    quant_fp8_blockwise_for_weight_impl,
+    quant_fp8_blockwise_impl,
+)
+from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+from primus.backends.megatron.core.fp8_utils import (
+    MXFP8_SCALING_BLOCK_SIZE,
+    SCALING_BLOCK_SIZE,
+)
+
+# ---------------------------------------------------------------------------
+# torch.compile-friendly wrappers for C++ quantization ops.
+#
+# The raw C++ ops (primus_turbo_cpp_extension::quantize_fp8_tensorwise etc.)
+# lack an Autograd dispatch key, which triggers warnings and can break
+# torch.compile graph tracing. Wrapping them with @torch.library.custom_op
+# (the same pattern used by gemm_fp8_impl) automatically handles Autograd
+# dispatch and provides register_fake for shape inference during tracing.
+# ---------------------------------------------------------------------------
+
+_custom_op = torch.library.custom_op
+
+
+@_custom_op("primus::quantize_fp8_tensorwise", mutates_args=(), device_types="cuda")
+def _quantize_fp8_tensorwise_op(
+    x: torch.Tensor, out_dtype: torch.dtype, scale: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(x, out_dtype, scale)
+
+
+@_quantize_fp8_tensorwise_op.register_fake
+def _quantize_fp8_tensorwise_fake(
+    x: torch.Tensor, out_dtype: torch.dtype, scale: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    x_fp8 = torch.empty_like(x, dtype=out_dtype)
+    scale_inv = torch.empty((), dtype=torch.float32, device=x.device)
+    return x_fp8, scale_inv
+
+
+def _quantize_fp8_tensorwise_setup_context(ctx, inputs, output):
+    pass
+
+
+def _quantize_fp8_tensorwise_backward(ctx, grad_x_fp8, grad_scale_inv):
+    return None, None, None
+
+
+_quantize_fp8_tensorwise_op.register_autograd(
+    _quantize_fp8_tensorwise_backward,
+    setup_context=_quantize_fp8_tensorwise_setup_context,
+)
+
+
+def _cast_transpose_fp8_fused_with_amax(
+    x: torch.Tensor,
+    out_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fused FP8 quantize + transpose + optional amax via the Triton @triton_op.
+
+    Returns (fp8_out, fp8_transpose, scale_inv). Inductor-transparent.
+
+    Uses the cast+transpose Triton kernel vendored in Primus
+    (fp8_cast_kernels_triton), so this arm builds against stock Turbo main.
+    """
+    from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import (
+        cast_transpose_fp8_triton,
+    )
+
+    return cast_transpose_fp8_triton(x, out_dtype, scale, amax_out)
+
+
+def _cast_fp8_fused_with_amax(
+    x: torch.Tensor,
+    out_dtype: torch.dtype,
+    scale: torch.Tensor,
+    amax_out: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Native-layout fused FP8 quantize (no transpose) + optional amax.
+
+    Analog of _cast_transpose_fp8_fused_with_amax for the natural NN/TN arm:
+    applies the delayed scale and captures the current abs-max in a single HBM
+    pass, with no [N, M] transpose write. Returns (fp8_out, scale_inv).
+
+    Uses the no-transpose cast Triton kernel vendored in Primus
+    (fp8_cast_kernels_triton), so this arm builds against stock Turbo main.
+    """
+    from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import (
+        cast_fp8_triton,
+    )
+
+    return cast_fp8_triton(x, out_dtype, scale, amax_out)
+
+
+def _quantize_fp8_tw(
+    x: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Dispatch tensorwise FP8 quantize via the standard C++ @custom_op."""
+    return _quantize_fp8_tensorwise_op(x, out_dtype)
+
+
+def _get_fp8_dtype(format: Format, is_fwd: bool):
+    if format == Format.E4M3:
+        return float8_e4m3
+    elif format == Format.E5M2:
+        return float8_e5m2
+    elif format == Format.HYBRID:
+        return float8_e4m3 if is_fwd else float8_e5m2
+    else:
+        raise ValueError(f"Unsupported FP8 format: {format}")
+
+
+def _build_fp8_config(config):
+    """Build Float8QuantConfig from TransformerConfig without TE dependency."""
+    FORMAT_MAP = {"e4m3": Format.E4M3, "hybrid": Format.HYBRID}
+    fmt = FORMAT_MAP[config.fp8]
+
+    if config.fp8_recipe in (Fp8Recipe.tensorwise, Fp8Recipe.delayed):
+        # Fp8Recipe.delayed is accepted as a TE-compatible alias: it implies
+        # tensorwise granularity (the only granularity that makes sense with
+        # delayed/amax-history scaling).  The actual delayed-vs-dynamic
+        # strategy is resolved separately via fp8_scaling_strategy or the
+        # recipe itself (see Float8*ParallelLinear._use_delayed_scaling).
+        return Float8QuantConfig(format=fmt, granularity=ScalingGranularity.TENSORWISE)
+    elif config.fp8_recipe == Fp8Recipe.blockwise:
+        return Float8QuantConfig(
+            format=fmt,
+            granularity=ScalingGranularity.BLOCKWISE,
+            block_size=SCALING_BLOCK_SIZE,
+        )
+    elif config.fp8_recipe == Fp8Recipe.mxfp8:
+        return Float8QuantConfig(
+            format=fmt,
+            granularity=ScalingGranularity.MX_BLOCKWISE,
+            block_size=MXFP8_SCALING_BLOCK_SIZE,
+            scale_dtype=ScaleDtype.E8M0,
+        )
+    else:
+        raise ValueError(
+            f"Float8 local spec does not support fp8_recipe={config.fp8_recipe}. "
+            f"Supported: tensorwise, delayed, blockwise, mxfp8."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Decomposed FP8 quantize — native aten ops for Inductor fusion
+# ---------------------------------------------------------------------------
+
+
+def _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max):
+    """Tensorwise FP8 quantize using native aten ops (traceable by Inductor).
+
+    Scale is computed in FP32 for precision, then narrowed to x.dtype (BF16) so
+    the pointwise chain (mul, clamp, to_fp8) stays in BF16 registers inside the
+    fused Triton kernel.  scale_inv is derived from the FP32 scale before the
+    narrowing to preserve hipBLASLt's float32 precision requirement.
+    """
+    amax = x.abs().amax().float()
+    scale_f32 = fp8_max / amax.clamp(min=1e-12)
+    scale_inv = 1.0 / scale_f32
+    scale = scale_f32.clamp(max=torch.finfo(x.dtype).max).to(x.dtype)
+    x_fp8 = (x * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
+    return x_fp8, scale_inv
+
+
+# ---------------------------------------------------------------------------
+# Autograd Functions — one per FP8 scaling granularity.
+# ---------------------------------------------------------------------------
+
+
+class DecomposedFP8LinearTensorwiseFunction(torch.autograd.Function):
+    """FP8 linear (Y = X @ W^T) with tensorwise scaling, using decomposed
+    quantize for forward-pass Inductor fusion.
+
+    Experimental: designed to enable Inductor fusion of quantize with
+    surrounding kernels. Currently slower than OpaqueFP8LinearTensorwiseFunction
+    due to Inductor's internal FP32 promotion and amax reduction barriers.
+    Kept for future work (delayed scaling, Compiled Autograd).
+
+    Uses setup_context pattern with primitive-only arguments (no config objects)
+    for clean Dynamo tracing.
+    """
+
+    @staticmethod
+    def forward(
+        input, weight, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value
+    ):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8, a_scale_inv = _quantize_fp8_tensorwise(input_2d, fp8_fwd_dtype, fp8_fwd_max)
+        b_fp8, b_scale_inv = _quantize_fp8_tensorwise(weight, fp8_fwd_dtype, fp8_fwd_max)
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        return output, a_fp8, a_scale_inv, b_fp8, b_scale_inv
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        _, _, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value = inputs
+        output_val, a_fp8, a_scale_inv, b_fp8, b_scale_inv = output
+
+        ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.mark_non_differentiable(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.out_dtype = inputs[0].dtype
+        ctx.orig_shape = inputs[0].shape
+        ctx.fp8_bwd_dtype = fp8_bwd_dtype
+        ctx.fp8_bwd_max = fp8_bwd_max
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        if not grad_2d.is_contiguous():
+            grad_2d = grad_2d.contiguous()
+
+        grad_fp8, grad_scale_inv = _quantize_fp8_tensorwise(grad_2d, ctx.fp8_bwd_dtype, ctx.fp8_bwd_max)
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            True,
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+
+        return grad_input, grad_weight, None, None, None, None, None, None
+
+
+class OpaqueFP8LinearTensorwiseFunction(torch.autograd.Function):
+    """FP8 linear (Y = X @ W^T) with tensorwise scaling using Primus Turbo's
+    opaque C++ quantize_fp8 kernel.
+
+    Uses the setup_context pattern with primitive-only saved state so that
+    torch.compile / TorchDynamo can trace through without graph breaks.
+    The caller must pre-extract FP8 weight data (b_fp8, b_scale_inv) from
+    FP8UnshardedWeightTensor *before* calling .apply().
+
+    Default path for tensorwise -- faster than the decomposed variant due to
+    optimized C++ quantize kernels.
+
+    Backward GEMMs are normalized to NT layout (transA=F, transB=T, transC=F)
+    so that hipBLASLt always selects the fast TN Tensile kernel on MI355X.
+    Pre-transposed FP8 copies of input and weight are saved for this purpose.
+    """
+
+    @staticmethod
+    def forward(
+        input,
+        weight,
+        weight_fp8,
+        weight_scale_inv,
+        fp8_fwd_dtype,
+        fp8_bwd_dtype,
+        gran_value,
+        backend_value,
+        force_nt=True,
+    ):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, fp8_fwd_dtype)
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            weight_fp8,
+            weight_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        if force_nt:
+            # Pre-transpose so backward dgrad/wgrad both run as NT GEMMs.
+            a_t_fp8 = a_fp8.t().contiguous()
+            w_t_fp8 = weight_fp8.t().contiguous()
+            return (output, a_t_fp8, a_scale_inv, w_t_fp8)
+
+        # Native: keep operands in their natural layout (no transpose). weight_fp8
+        # is not returned -- setup_context saves it straight from inputs.
+        return (output, a_fp8, a_scale_inv)
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        # Tolerant unpack: legacy 8-arg callers default to forced NT; only the
+        # 9th positional (force_nt) toggles the native arm.
+        force_nt = inputs[8] if len(inputs) > 8 else True
+        (
+            input,
+            weight,
+            weight_fp8,
+            weight_scale_inv,
+            fp8_fwd_dtype,
+            fp8_bwd_dtype,
+            gran_value,
+            backend_value,
+        ) = inputs[:8]
+
+        if force_nt:
+            output_val, a_t_fp8, a_scale_inv, w_t_fp8 = output
+            ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, weight_scale_inv)
+            ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8)
+        else:
+            output_val, a_fp8, a_scale_inv = output
+            ctx.save_for_backward(a_fp8, a_scale_inv, weight_fp8, weight_scale_inv)
+            ctx.mark_non_differentiable(a_fp8, a_scale_inv)
+
+        ctx.force_nt = force_nt
+        ctx.out_dtype = input.dtype
+        ctx.orig_shape = input.shape
+        ctx.fp8_bwd_dtype = fp8_bwd_dtype
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        grad_fp8, grad_scale_inv = _quantize_fp8_tw(grad_2d, ctx.fp8_bwd_dtype)
+
+        if ctx.force_nt:
+            a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv = ctx.saved_tensors
+
+            grad_input = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                w_t_fp8,
+                w_scale_inv,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            grad_t_fp8 = grad_fp8.t().contiguous()
+            grad_weight = gemm_fp8_impl(
+                grad_t_fp8,
+                grad_scale_inv,
+                False,
+                a_t_fp8,
+                a_scale_inv,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+        else:
+            a_fp8, a_scale_inv, w_fp8, w_scale_inv = ctx.saved_tensors
+
+            # dgrad NN: grad @ weight  (transA=F, transB=F)
+            grad_input = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                w_fp8,
+                w_scale_inv,
+                False,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            # wgrad TN: grad^T @ a  (transA=T, transB=F)
+            grad_weight = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                True,
+                a_fp8,
+                a_scale_inv,
+                False,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+
+        return grad_input, grad_weight, None, None, None, None, None, None, None
+
+
+class DelayedFP8LinearTensorwiseFunction(torch.autograd.Function):
+    """FP8 linear with delayed tensorwise scaling -- 3 independent scale tracks.
+
+    Weight is quantized inline during the forward pass (inside the compiled
+    graph where Inductor eliminates CPU dispatcher overhead).  Input and
+    gradient quantization also happen inline with fused amax capture.
+    Backward GEMM layout follows ``force_nt`` (set from ``fp8_force_nt_layout``):
+    native by default (dgrad=NN, wgrad=TN, no transpose), or forced-NT (operands
+    pre-transposed so both backward GEMMs run as NT) when opted in.
+    """
+
+    @staticmethod
+    def forward(
+        input,
+        weight,
+        scale_input,
+        scale_weight,
+        scale_grad,
+        staged_input_amax,
+        staged_weight_amax,
+        staged_grad_amax,
+        fp8_fwd_dtype,
+        fp8_bwd_dtype,
+        gran_value,
+        backend_value,
+        force_nt=True,
+    ):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        if force_nt:
+            a_fp8, a_t_fp8, a_scale_inv = _cast_transpose_fp8_fused_with_amax(
+                input_2d, fp8_fwd_dtype, scale_input, staged_input_amax
+            )
+            torch._assert(a_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose")
+
+            w_fp8, w_t_fp8, w_scale_inv = _cast_transpose_fp8_fused_with_amax(
+                weight, fp8_fwd_dtype, scale_weight, staged_weight_amax
+            )
+            torch._assert(w_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose")
+
+            output = gemm_fp8_impl(
+                a_fp8,
+                a_scale_inv,
+                False,
+                w_fp8,
+                w_scale_inv,
+                True,
+                out_dtype,
+                False,
+                granularity=gran_value,
+                default_backend=backend_value,
+            )
+            output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+            return (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv)
+
+        # Native: apply the delayed scale and capture the current amax in a
+        # single fused pass via the no-transpose Triton kernel. No transpose;
+        # backward runs dgrad=NN / wgrad=TN (see backward).
+        a_fp8, a_scale_inv = _cast_fp8_fused_with_amax(
+            input_2d, fp8_fwd_dtype, scale_input, staged_input_amax
+        )
+
+        w_fp8, w_scale_inv = _cast_fp8_fused_with_amax(
+            weight, fp8_fwd_dtype, scale_weight, staged_weight_amax
+        )
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            w_fp8,
+            w_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        return (output, a_fp8, a_scale_inv, w_fp8, w_scale_inv)
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        # Tolerant unpack: legacy 12-arg callers default to forced NT; only the
+        # 13th positional (force_nt) toggles the native arm.
+        force_nt = inputs[12] if len(inputs) > 12 else True
+        (
+            input,
+            weight,
+            scale_input,
+            scale_weight,
+            scale_grad,
+            staged_input_amax,
+            staged_weight_amax,
+            staged_grad_amax,
+            fp8_fwd_dtype,
+            fp8_bwd_dtype,
+            gran_value,
+            backend_value,
+        ) = inputs[:12]
+
+        if force_nt:
+            (output_val, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv) = output
+            ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad, staged_grad_amax)
+            ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv)
+        else:
+            (output_val, a_fp8, a_scale_inv, w_fp8, w_scale_inv) = output
+            ctx.save_for_backward(a_fp8, a_scale_inv, w_fp8, w_scale_inv, scale_grad, staged_grad_amax)
+            ctx.mark_non_differentiable(a_fp8, a_scale_inv, w_fp8, w_scale_inv)
+
+        ctx.force_nt = force_nt
+        ctx.out_dtype = input.dtype
+        ctx.orig_shape = input.shape
+        ctx.fp8_bwd_dtype = fp8_bwd_dtype
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        op_a, a_scale_inv, op_w, w_scale_inv, scale_grad, staged_grad_amax = ctx.saved_tensors
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+
+        if ctx.force_nt:
+            # op_a = a_t_fp8, op_w = w_t_fp8
+            grad_fp8, grad_t_fp8, grad_scale_inv = _cast_transpose_fp8_fused_with_amax(
+                grad_2d, ctx.fp8_bwd_dtype, scale_grad, staged_grad_amax
+            )
+            torch._assert(grad_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose")
+
+            grad_input = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                op_w,
+                w_scale_inv,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            grad_weight = gemm_fp8_impl(
+                grad_t_fp8,
+                grad_scale_inv,
+                False,
+                op_a,
+                a_scale_inv,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+        else:
+            # Native: op_a = a_fp8, op_w = w_fp8. Fused quantize + amax, no transpose.
+            grad_fp8, grad_scale_inv = _cast_fp8_fused_with_amax(
+                grad_2d, ctx.fp8_bwd_dtype, scale_grad, staged_grad_amax
+            )
+
+            # dgrad NN: grad @ weight  (transA=F, transB=F)
+            grad_input = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                op_w,
+                w_scale_inv,
+                False,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            # wgrad TN: grad^T @ a  (transA=T, transB=F)
+            grad_weight = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                True,
+                op_a,
+                a_scale_inv,
+                False,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.gran_value,
+                default_backend=ctx.backend_value,
+            )
+
+        # 13 inputs to forward -> 13 grads returned (2 real + 11 None).
+        return (
+            grad_input,
+            grad_weight,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+@torch._dynamo.allow_in_graph
+class DualFP8LinearTensorwiseFunction(torch.autograd.Function):
+    """Two independent FP8 linears in a single autograd node.
+
+    Y_a = X_a @ W_a^T,  Y_b = X_b @ W_b^T
+
+    Uses the setup_context pattern with primitive-only saved state so that
+    torch.compile can trace through without graph breaks. The caller must
+    pre-extract FP8 weight data before calling .apply().
+
+    Backward GEMMs are normalized to NT layout (transA=F, transB=T, transC=F)
+    so that hipBLASLt always selects the fast TN Tensile kernel on MI355X.
+
+    @torch._dynamo.allow_in_graph is required on top of the modern
+    forward/setup_context split: under PyTorch 2.12's AOTAutograd, letting
+    Dynamo trace *into* this Function's forward (the default for the
+    forward/setup_context API) silently produces wrong gradients for the 6
+    outputs that exist only to be captured by setup_context/save_for_backward
+    and are otherwise unused by the outer graph (matches
+    https://github.com/pytorch/pytorch/issues/131794, regressed by the fix in
+    https://github.com/pytorch/pytorch/pull/186355). allow_in_graph makes
+    Dynamo treat .apply() as a single opaque node instead of inlining into
+    forward, which sidesteps the bug while still keeping everything in one
+    fused graph (verified: graph_count=1, graph_break_count=0).
+    """
+
+    @staticmethod
+    def forward(
+        input_a,
+        weight_a,
+        weight_fp8_a,
+        weight_scale_a,
+        input_b,
+        weight_b,
+        weight_fp8_b,
+        weight_scale_b,
+        fp8_fwd_dtype,
+        fp8_bwd_dtype,
+        gran_value,
+        backend_value,
+    ):
+        out_dtype = input_a.dtype
+
+        orig_shape_a = input_a.shape
+        orig_shape_b = input_b.shape
+        input_a_2d = input_a.reshape(-1, input_a.shape[-1])
+        input_b_2d = input_b.reshape(-1, input_b.shape[-1])
+
+        a_fp8_a, a_scale_a = _quantize_fp8_tw(input_a_2d, fp8_fwd_dtype)
+        a_fp8_b, a_scale_b = _quantize_fp8_tw(input_b_2d, fp8_fwd_dtype)
+
+        output_a = gemm_fp8_impl(
+            a_fp8_a,
+            a_scale_a,
+            False,
+            weight_fp8_a,
+            weight_scale_a,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+        output_b = gemm_fp8_impl(
+            a_fp8_b,
+            a_scale_b,
+            False,
+            weight_fp8_b,
+            weight_scale_b,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+
+        output_a = output_a.reshape(*orig_shape_a[:-1], output_a.shape[-1])
+        output_b = output_b.reshape(*orig_shape_b[:-1], output_b.shape[-1])
+
+        a_t_fp8_a = a_fp8_a.t().contiguous()
+        w_t_fp8_a = weight_fp8_a.t().contiguous()
+        a_t_fp8_b = a_fp8_b.t().contiguous()
+        w_t_fp8_b = weight_fp8_b.t().contiguous()
+
+        return (output_a, output_b, a_t_fp8_a, a_scale_a, w_t_fp8_a, a_t_fp8_b, a_scale_b, w_t_fp8_b)
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        (
+            input_a,
+            weight_a,
+            weight_fp8_a,
+            weight_scale_a,
+            input_b,
+            weight_b,
+            weight_fp8_b,
+            weight_scale_b,
+            fp8_fwd_dtype,
+            fp8_bwd_dtype,
+            gran_value,
+            backend_value,
+        ) = inputs
+        (output_a, output_b, a_t_fp8_a, a_scale_a, w_t_fp8_a, a_t_fp8_b, a_scale_b, w_t_fp8_b) = output
+
+        ctx.save_for_backward(
+            a_t_fp8_a,
+            a_scale_a,
+            w_t_fp8_a,
+            weight_scale_a,
+            a_t_fp8_b,
+            a_scale_b,
+            w_t_fp8_b,
+            weight_scale_b,
+        )
+        ctx.mark_non_differentiable(
+            a_t_fp8_a,
+            a_scale_a,
+            w_t_fp8_a,
+            a_t_fp8_b,
+            a_scale_b,
+            w_t_fp8_b,
+        )
+        ctx.out_dtype = input_a.dtype
+        ctx.orig_shape_a = input_a.shape
+        ctx.orig_shape_b = input_b.shape
+        ctx.fp8_bwd_dtype = fp8_bwd_dtype
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output_a, grad_output_b, *_):
+        if not grad_output_a.is_contiguous():
+            grad_output_a = grad_output_a.contiguous()
+        if not grad_output_b.is_contiguous():
+            grad_output_b = grad_output_b.contiguous()
+
+        (a_t_fp8_a, a_scale_a, w_t_fp8_a, w_scale_a, a_t_fp8_b, a_scale_b, w_t_fp8_b, w_scale_b) = (
+            ctx.saved_tensors
+        )
+
+        # --- Stream A ---
+        grad_a_2d = grad_output_a.reshape(-1, grad_output_a.shape[-1])
+        grad_fp8_a, grad_scale_a = _quantize_fp8_tw(grad_a_2d, ctx.fp8_bwd_dtype)
+
+        grad_input_a = gemm_fp8_impl(
+            grad_fp8_a,
+            grad_scale_a,
+            False,
+            w_t_fp8_a,
+            w_scale_a,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        grad_input_a = grad_input_a.reshape(ctx.orig_shape_a)
+
+        grad_t_fp8_a = grad_fp8_a.t().contiguous()
+        grad_weight_a = gemm_fp8_impl(
+            grad_t_fp8_a,
+            grad_scale_a,
+            False,
+            a_t_fp8_a,
+            a_scale_a,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+
+        # --- Stream B ---
+        grad_b_2d = grad_output_b.reshape(-1, grad_output_b.shape[-1])
+        grad_fp8_b, grad_scale_b = _quantize_fp8_tw(grad_b_2d, ctx.fp8_bwd_dtype)
+
+        grad_input_b = gemm_fp8_impl(
+            grad_fp8_b,
+            grad_scale_b,
+            False,
+            w_t_fp8_b,
+            w_scale_b,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        grad_input_b = grad_input_b.reshape(ctx.orig_shape_b)
+
+        grad_t_fp8_b = grad_fp8_b.t().contiguous()
+        grad_weight_b = gemm_fp8_impl(
+            grad_t_fp8_b,
+            grad_scale_b,
+            False,
+            a_t_fp8_b,
+            a_scale_b,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+
+        return (
+            grad_input_a,
+            grad_weight_a,
+            None,
+            None,
+            grad_input_b,
+            grad_weight_b,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+@torch._dynamo.allow_in_graph
+class FP8LinearRowwiseFunction(torch.autograd.Function):
+    """FP8 linear (Y = X @ W^T) with rowwise scaling. Compile-friendly."""
+
+    @staticmethod
+    def forward(ctx, input, weight, config):
+        a_dtype = _get_fp8_dtype(config.format, is_fwd=True)
+        b_dtype = _get_fp8_dtype(config.format, is_fwd=True)
+        out_dtype = input.dtype
+        gran = config.granularity
+
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8_row, a_scale_inv_row = quantize_fp8(input_2d, a_dtype, gran, axis=-1)
+        b_fp8_row, b_scale_inv_row = quantize_fp8(weight, b_dtype, gran, axis=-1)
+
+        output = gemm_fp8_impl(
+            a_fp8_row,
+            a_scale_inv_row,
+            False,
+            b_fp8_row,
+            b_scale_inv_row,
+            True,
+            out_dtype,
+            False,
+            granularity=gran.value,
+            default_backend=BackendType.CK.value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        a_fp8_col, a_scale_inv_col = quantize_fp8(input_2d, a_dtype, gran, axis=-2)
+        b_fp8_col, b_scale_inv_col = quantize_fp8(weight, b_dtype, gran, axis=-2)
+
+        ctx.save_for_backward(a_fp8_col, a_scale_inv_col, b_fp8_col, b_scale_inv_col)
+        ctx.out_dtype = out_dtype
+        ctx.config = config
+        ctx.orig_shape = orig_shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        a_fp8_col, a_scale_inv_col, b_fp8_col, b_scale_inv_col = ctx.saved_tensors
+        grad_dtype = _get_fp8_dtype(ctx.config.format, is_fwd=False)
+        gran = ctx.config.granularity
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+
+        grad_fp8_row, grad_scale_inv_row = quantize_fp8(grad_2d, grad_dtype, gran, axis=-1)
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8_row,
+            grad_scale_inv_row,
+            False,
+            b_fp8_col,
+            b_scale_inv_col,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=gran.value,
+            default_backend=BackendType.CK.value,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_fp8_col, grad_scale_inv_col = quantize_fp8(grad_2d, grad_dtype, gran, axis=-2)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8_col,
+            a_scale_inv_col,
+            True,
+            grad_fp8_col,
+            grad_scale_inv_col,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=gran.value,
+            default_backend=BackendType.CK.value,
+        )
+
+        return grad_input, grad_weight, None
+
+
+@torch._dynamo.allow_in_graph
+class FP8LinearBlockwiseFunction(torch.autograd.Function):
+    """FP8 linear (Y = X @ W^T) with blockwise scaling. Compile-friendly.
+
+    Note: saves BF16 input AND BF16 weight for backward (must re-quantize
+    with different axis/dtype), so no activation memory savings compared to
+    BF16 baseline.
+
+    Forward uses CK backend (NT layout). Backward uses Triton backend for
+    NN (grad_input) and TN (grad_weight) layouts because CK's blockwise
+    GEMM produces NaN on these layouts.  To satisfy Triton's same-dtype
+    constraint, the weight is re-quantized to the backward dtype (E5M2 for
+    hybrid format) from the saved BF16 copy.
+    """
+
+    @staticmethod
+    def forward(ctx, input, weight, config):
+        a_dtype = _get_fp8_dtype(config.format, is_fwd=True)
+        b_dtype = _get_fp8_dtype(config.format, is_fwd=True)
+        out_dtype = input.dtype
+        gran = config.granularity
+        bs = config.block_size
+
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8_row, a_scale_inv_row = quant_fp8_blockwise_impl(input_2d, a_dtype, axis=1, block_size=bs)
+        b_fp8, b_scale_inv = quant_fp8_blockwise_for_weight_impl(weight, b_dtype, block_size=bs)
+
+        output = gemm_fp8_impl(
+            a_fp8_row,
+            a_scale_inv_row,
+            False,
+            b_fp8,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran.value,
+            default_backend=BackendType.CK.value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        ctx.save_for_backward(input_2d, weight)
+        ctx.out_dtype = out_dtype
+        ctx.config = config
+        ctx.orig_shape = orig_shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        input_saved, weight_saved = ctx.saved_tensors
+        bwd_dtype = _get_fp8_dtype(ctx.config.format, is_fwd=False)
+        gran = ctx.config.granularity
+        bs = ctx.config.block_size
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+
+        grad_fp8_row, grad_scale_inv_row = quant_fp8_blockwise_impl(
+            grad_2d, bwd_dtype, axis=-1, block_size=bs
+        )
+        grad_fp8_col, grad_scale_inv_col = quant_fp8_blockwise_impl(
+            grad_2d, bwd_dtype, axis=-2, block_size=bs
+        )
+        a_fp8_col, a_scale_inv_col = quant_fp8_blockwise_impl(input_saved, bwd_dtype, axis=0, block_size=bs)
+        b_fp8_bwd, b_scale_inv_bwd = quant_fp8_blockwise_for_weight_impl(
+            weight_saved, bwd_dtype, block_size=bs
+        )
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8_row,
+            grad_scale_inv_row,
+            False,
+            b_fp8_bwd,
+            b_scale_inv_bwd,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=gran.value,
+            default_backend=BackendType.TRITON.value,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8_col,
+            a_scale_inv_col,
+            True,
+            grad_fp8_col,
+            grad_scale_inv_col,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=gran.value,
+            default_backend=BackendType.TRITON.value,
+        )
+
+        return grad_input, grad_weight, None
+
+
+# ---------------------------------------------------------------------------
+# Delayed FP8 scaling helpers — shared by Float8ColumnParallelLinear and
+# Float8RowParallelLinear.
+# ---------------------------------------------------------------------------
+
+
+def _update_fp8_scale(scale_buf, scale_inv_buf, amax, fp8_max):
+    """TE-aligned scale computation with edge case handling."""
+    sf = fp8_max / amax.clamp(min=1e-12)
+    sf = torch.where(amax > 0.0, sf, scale_buf)
+    sf = torch.where(torch.isfinite(amax), sf, scale_buf)
+    sf = torch.where(
+        torch.isinf(sf),
+        torch.tensor(torch.finfo(torch.float32).max, device=sf.device),
+        sf,
+    )
+    scale_buf.fill_(sf)
+    if scale_inv_buf is not None:
+        scale_inv_buf.fill_(1.0 / sf)
+
+
+def _init_delayed_scaling_state(module):
+    """Initialize delayed FP8 scaling buffers on a Float8*ParallelLinear."""
+    config = module.config
+    history_len = getattr(config, "fp8_amax_history_len", 1)
+    fp8_fwd_max = torch.finfo(module._fp8_fwd_dtype).max
+    fp8_bwd_max = torch.finfo(module._fp8_bwd_dtype).max
+
+    for name in ["amax_history_input", "amax_history_weight", "amax_history_grad"]:
+        module.register_buffer(
+            name,
+            torch.zeros(history_len),
+            persistent=False,
+        )
+
+    for name in ["scale_input", "scale_weight", "scale_grad"]:
+        module.register_buffer(
+            name,
+            torch.tensor(1.0, dtype=torch.float32),
+            persistent=False,
+        )
+
+    for name in ["staged_input_amax", "staged_grad_amax", "staged_weight_amax"]:
+        module.register_buffer(
+            name,
+            torch.tensor(0.0),
+            persistent=False,
+        )
+
+    module._fp8_fwd_max = fp8_fwd_max
+    module._fp8_bwd_max = fp8_bwd_max
+    module._amax_compute_algo = getattr(
+        config,
+        "fp8_amax_compute_algo",
+        "most_recent",
+    )
+    module._history_idx = 0
+    module._first_delayed_step = True
+
+
+def _batch_compute_scales(amaxes, fp8_max, old_scales):
+    """Vectorized scale computation across all modules at once."""
+    sf = fp8_max / amaxes.clamp(min=1e-12)
+    sf = torch.where(amaxes > 0.0, sf, old_scales)
+    sf = torch.where(torch.isfinite(amaxes), sf, old_scales)
+    sf = torch.where(torch.isinf(sf), torch.finfo(torch.float32).max, sf)
+    return sf
+
+
+# ---------------------------------------------------------------------------
+# Fused Triton kernel: replaces Python-loop-based _batched_update per step
+# with a single GPU dispatch over all (module, track) pairs.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_delayed_scale_update_kernel(
+    amax_history_ptr,
+    staged_amaxes_ptr,
+    scales_ptr,
+    fp8_maxes_ptr,
+    history_idx,
+    N: tl.constexpr,
+    H: tl.constexpr,
+    use_max_algo: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    FP32_MAX: tl.constexpr,
+    FILTER_ZEROS: tl.constexpr = False,
+):
+    track = tl.program_id(1)
+    mod = tl.program_id(0)
+
+    base = track * N * H + mod * H
+
+    new_amax = tl.load(staged_amaxes_ptr + track * N + mod)
+    tl.store(amax_history_ptr + base + history_idx, new_amax)
+
+    if use_max_algo:
+        amax = tl.zeros([], dtype=tl.float32)
+        has_nonzero = tl.zeros([], dtype=tl.int32)
+        for off in range(0, H, BLOCK_H):
+            h_idx = off + tl.arange(0, BLOCK_H)
+            mask = h_idx < H
+            vals = tl.load(amax_history_ptr + base + h_idx, mask=mask, other=0.0)
+            if FILTER_ZEROS:
+                nonzero_mask = vals > 0.0
+                has_nonzero = has_nonzero | tl.sum(nonzero_mask.to(tl.int32), axis=0)
+                filtered = tl.where(nonzero_mask, vals, 0.0)
+                amax = tl.maximum(amax, tl.max(filtered, axis=0))
+            else:
+                amax = tl.maximum(amax, tl.max(vals, axis=0))
+    else:
+        amax = new_amax
+        has_nonzero = tl.where(new_amax > 0.0, 1, 0).to(tl.int32)
+
+    fp8_max = tl.load(fp8_maxes_ptr + track)
+    old_scale = tl.load(scales_ptr + track * N + mod)
+    sf = fp8_max / tl.maximum(amax, 1e-12)
+    if FILTER_ZEROS:
+        sf = tl.where(has_nonzero > 0, sf, old_scale)
+    else:
+        sf = tl.where(amax > 0.0, sf, old_scale)
+    sf = tl.where(amax == amax, sf, old_scale)
+    sf = tl.minimum(sf, FP32_MAX)
+    tl.store(scales_ptr + track * N + mod, sf)
+
+
+# ---------------------------------------------------------------------------
+# Global registry: replaces per-module scalar buffers with views into
+# contiguous (N,) tensors to eliminate Python iteration in the preamble.
+# ---------------------------------------------------------------------------
+
+
+class _DelayedScalingRegistry:
+    __slots__ = (
+        "n",
+        "modules",
+        "fwd_max",
+        "bwd_max",
+        "algo",
+        "history_len",
+        "_first_step",
+        "amax_history",
+        "_history_idx",
+        "fp8_maxes",
+        "staged_amaxes_3n",
+        "scales_3n",
+        "reduce_amax",
+        "amax_reduce_group",
+        "skip_bootstrap",
+        "filter_zeros",
+    )
+
+    def __init__(self, modules):
+        self.n = len(modules)
+        self.modules = modules
+        m0 = modules[0]
+        self.fwd_max = m0._fp8_fwd_max
+        self.bwd_max = m0._fp8_bwd_max
+        self.algo = m0._amax_compute_algo
+        self.history_len = m0.amax_history_input.shape[0]
+        device = m0.weight.device
+
+        self._first_step = True
+
+        H = self.history_len
+        N = self.n
+
+        self.scales_3n = torch.ones(3, N, dtype=torch.float32, device=device)
+        self.staged_amaxes_3n = torch.zeros(3, N, dtype=torch.float32, device=device)
+
+        self.amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=device)
+        self._history_idx = 0
+        self.fp8_maxes = torch.tensor(
+            [self.fwd_max, self.fwd_max, self.bwd_max],
+            dtype=torch.float32,
+            device=device,
+        )
+
+        _config = getattr(m0, "config", None)
+        self.reduce_amax = getattr(_config, "fp8_reduce_amax", False) if _config else False
+        self.skip_bootstrap = getattr(_config, "fp8_skip_first_step_bootstrap", False) if _config else False
+        self.filter_zeros = getattr(_config, "fp8_filter_zeros_in_history", False) if _config else False
+        self.amax_reduce_group = None
+        if self.reduce_amax:
+            from megatron.core import parallel_state
+
+            if parallel_state.model_parallel_is_initialized():
+                self.amax_reduce_group = parallel_state.get_amax_reduction_group(
+                    with_context_parallel=True,
+                    tp_only_amax_red=getattr(_config, "tp_only_amax_red", False),
+                )
+
+        for i, m in enumerate(modules):
+            m._buffers["scale_input"] = torch.ones((), dtype=torch.float32, device=device)
+            m._buffers["scale_weight"] = torch.ones((), dtype=torch.float32, device=device)
+            m._buffers["scale_grad"] = torch.ones((), dtype=torch.float32, device=device)
+            m._buffers["staged_input_amax"] = torch.zeros((), dtype=torch.float32, device=device)
+            m._buffers["staged_weight_amax"] = torch.zeros((), dtype=torch.float32, device=device)
+            m._buffers["staged_grad_amax"] = torch.zeros((), dtype=torch.float32, device=device)
+            m._buffers["amax_history_input"] = self.amax_history[0, i, :]
+            m._buffers["amax_history_weight"] = self.amax_history[1, i, :]
+            m._buffers["amax_history_grad"] = self.amax_history[2, i, :]
+
+
+@torch.no_grad()
+def _fast_update_scales(registry):
+    """Ultra-fast scale update for most_recent + history_len=1.
+
+    Gathers per-module scalar amaxes into batched tensors, computes scales,
+    then scatters back.  Each module has independent storage for its scale
+    and amax buffers so that torch.compile version tracking is correct.
+    """
+    r = registry
+
+    # Detect device migration AND per-module buffer pointer breakage (e.g.
+    # after _reset_fp8_local_spec during MLPerf warmup teardown). Mirrors the
+    # same check in _fast_update_scales_with_history so the warmup-reset
+    # docstring contract holds for history_len==1 configs too: a fresh
+    # registry.__init__(modules) call sets _first_step=True, which bootstraps
+    # weight amaxes from the restored (post-warmup) weights below.
+    if (
+        r.scales_3n.device != r.modules[0].weight.device
+        or r.amax_history.untyped_storage().data_ptr()
+        != r.modules[0].amax_history_input.untyped_storage().data_ptr()
+    ):
+        r.__init__(r.modules)
+
+    if r._first_step and not r.skip_bootstrap:
+        weights = [m.weight.data for m in r.modules]
+        try:
+            weight_amaxes = torch.stack(torch._foreach_norm(weights, ord=float("inf"))).float()
+        except (AttributeError, RuntimeError):
+            weight_amaxes = torch.stack([w.abs().amax().float() for w in weights])
+        for i, m in enumerate(r.modules):
+            m.staged_weight_amax.fill_(weight_amaxes[i].item())
+    if r._first_step:
+        r._first_step = False
+
+    staged_input = torch.stack([m.staged_input_amax for m in r.modules])
+    staged_weight = torch.stack([m.staged_weight_amax for m in r.modules])
+    staged_grad = torch.stack([m.staged_grad_amax for m in r.modules])
+
+    if r.reduce_amax and r.amax_reduce_group is not None:
+        amaxes_3n = torch.stack([staged_input, staged_weight, staged_grad])
+        torch.distributed.all_reduce(
+            amaxes_3n,
+            op=torch.distributed.ReduceOp.MAX,
+            group=r.amax_reduce_group,
+        )
+        staged_input, staged_weight, staged_grad = amaxes_3n[0], amaxes_3n[1], amaxes_3n[2]
+
+    scale_input = torch.stack([m.scale_input for m in r.modules])
+    scale_weight = torch.stack([m.scale_weight for m in r.modules])
+    scale_grad = torch.stack([m.scale_grad for m in r.modules])
+
+    new_input = _batch_compute_scales(staged_input, r.fwd_max, scale_input)
+    new_weight = _batch_compute_scales(staged_weight, r.fwd_max, scale_weight)
+    new_grad = _batch_compute_scales(staged_grad, r.bwd_max, scale_grad)
+
+    # Batched device-only scatter: unbinding a (N,) float32 tensor into N
+    # 0-d views and copying with torch._foreach_copy_ keeps the per-module
+    # scale_* buffers independent (required for torch.compile version
+    # tracking) without the 3*N .item() host syncs the loop version
+    # incurred.
+    #
+    # NOTE (perf): the per-module scalar layout was chosen so that
+    # torch.compile sees each module's scale_* as an independent buffer
+    # version. A shared (N,) buffer with module-level views would let us
+    # skip the unbind+foreach_copy entirely, at the cost of more frequent
+    # compile cache invalidation. Worth benchmarking once the compile
+    # convergence work settles.
+    torch._foreach_copy_(
+        [m.scale_input for m in r.modules],
+        list(new_input.unbind()),
+    )
+    torch._foreach_copy_(
+        [m.scale_weight for m in r.modules],
+        list(new_weight.unbind()),
+    )
+    torch._foreach_copy_(
+        [m.scale_grad for m in r.modules],
+        list(new_grad.unbind()),
+    )
+
+
+@torch.no_grad()
+def _fast_update_scales_with_history(registry):
+    """Fused scale update for delayed scaling with history_len > 1.
+
+    Single Triton kernel dispatch replaces the prior Python-loop-based
+    per-module scale update (~200+ tiny GPU ops) with one launch.
+    """
+    r = registry
+
+    if (
+        r.amax_history.device != r.modules[0].weight.device
+        or r.amax_history.untyped_storage().data_ptr()
+        != r.modules[0].amax_history_input.untyped_storage().data_ptr()
+    ):
+        r.__init__(r.modules)
+
+    if r._first_step and not r.skip_bootstrap:
+        weights = [m.weight.data for m in r.modules]
+        try:
+            weight_amaxes = torch.stack(torch._foreach_norm(weights, ord=float("inf"))).float()
+        except (AttributeError, RuntimeError):
+            weight_amaxes = torch.stack([w.abs().amax().float() for w in weights])
+        for i, m in enumerate(r.modules):
+            m.staged_weight_amax.fill_(weight_amaxes[i].item())
+    if r._first_step:
+        r._first_step = False
+
+    r.staged_amaxes_3n[0].copy_(torch.stack([m.staged_input_amax for m in r.modules]))
+    r.staged_amaxes_3n[1].copy_(torch.stack([m.staged_weight_amax for m in r.modules]))
+    r.staged_amaxes_3n[2].copy_(torch.stack([m.staged_grad_amax for m in r.modules]))
+
+    if r.reduce_amax and r.amax_reduce_group is not None:
+        torch.distributed.all_reduce(
+            r.staged_amaxes_3n,
+            op=torch.distributed.ReduceOp.MAX,
+            group=r.amax_reduce_group,
+        )
+
+    N = r.n
+    H = r.history_len
+    BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024
+
+    _fused_delayed_scale_update_kernel[(N, 3)](
+        r.amax_history,
+        r.staged_amaxes_3n,
+        r.scales_3n,
+        r.fp8_maxes,
+        r._history_idx,
+        N=N,
+        H=H,
+        use_max_algo=(r.algo == "max"),
+        BLOCK_H=BLOCK_H,
+        FP32_MAX=torch.finfo(torch.float32).max,
+        FILTER_ZEROS=r.filter_zeros,
+    )
+
+    # See _fast_update_scales for the rationale behind _foreach_copy_ +
+    # unbind here. r.scales_3n is the (3, N) registry-batched scale tensor
+    # produced by the Triton kernel; rows 0/1/2 = input/weight/grad.
+    torch._foreach_copy_(
+        [m.scale_input for m in r.modules],
+        list(r.scales_3n[0].unbind()),
+    )
+    torch._foreach_copy_(
+        [m.scale_weight for m in r.modules],
+        list(r.scales_3n[1].unbind()),
+    )
+    torch._foreach_copy_(
+        [m.scale_grad for m in r.modules],
+        list(r.scales_3n[2].unbind()),
+    )
+
+    r._history_idx = (r._history_idx + 1) % H
+    for m in r.modules:
+        m._history_idx = r._history_idx
+
+
+# ---------------------------------------------------------------------------
+# Async allreduce split: stage+launch / wait+compute
+# ---------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def _stage_and_launch_async_allreduce(registry):
+    """Stage per-module amaxes into r.staged_amaxes_3n and launch async allreduce.
+
+    Returns the async work handle (or None if allreduce is not needed).
+    Must be called after backward completes so staged_*_amax buffers are
+    populated.
+    """
+    r = registry
+    r.staged_amaxes_3n[0].copy_(torch.stack([m.staged_input_amax for m in r.modules]))
+    r.staged_amaxes_3n[1].copy_(torch.stack([m.staged_weight_amax for m in r.modules]))
+    r.staged_amaxes_3n[2].copy_(torch.stack([m.staged_grad_amax for m in r.modules]))
+
+    if r.reduce_amax and r.amax_reduce_group is not None:
+        handle = torch.distributed.all_reduce(
+            r.staged_amaxes_3n,
+            op=torch.distributed.ReduceOp.MAX,
+            group=r.amax_reduce_group,
+            async_op=True,
+        )
+        return handle
+    return None
+
+
+@torch.no_grad()
+def _wait_and_compute_scales(registry, handle):
+    """Wait on async allreduce handle and compute new scales from reduced amaxes.
+
+    For most_recent + history_len=1, computes scales directly from
+    staged_amaxes_3n. For history-based, dispatches the fused Triton kernel.
+    Handles weight caching if enabled.
+    """
+    r = registry
+
+    if handle is not None:
+        handle.wait()
+
+    if r.algo == "most_recent" and r.history_len == 1:
+        staged_input = r.staged_amaxes_3n[0]
+        staged_weight = r.staged_amaxes_3n[1]
+        staged_grad = r.staged_amaxes_3n[2]
+
+        scale_input = torch.stack([m.scale_input for m in r.modules])
+        scale_weight = torch.stack([m.scale_weight for m in r.modules])
+        scale_grad = torch.stack([m.scale_grad for m in r.modules])
+
+        new_input = _batch_compute_scales(staged_input, r.fwd_max, scale_input)
+        new_weight = _batch_compute_scales(staged_weight, r.fwd_max, scale_weight)
+        new_grad = _batch_compute_scales(staged_grad, r.bwd_max, scale_grad)
+
+        for i, m in enumerate(r.modules):
+            m.scale_input.fill_(new_input[i].item())
+            m.scale_weight.fill_(new_weight[i].item())
+            m.scale_grad.fill_(new_grad[i].item())
+    else:
+        N = r.n
+        H = r.history_len
+        BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024
+
+        _fused_delayed_scale_update_kernel[(N, 3)](
+            r.amax_history,
+            r.staged_amaxes_3n,
+            r.scales_3n,
+            r.fp8_maxes,
+            r._history_idx,
+            N=N,
+            H=H,
+            use_max_algo=(r.algo == "max"),
+            BLOCK_H=BLOCK_H,
+            FP32_MAX=torch.finfo(torch.float32).max,
+            FILTER_ZEROS=r.filter_zeros,
+        )
+
+        for i, m in enumerate(r.modules):
+            m.scale_input.fill_(r.scales_3n[0, i].item())
+            m.scale_weight.fill_(r.scales_3n[1, i].item())
+            m.scale_grad.fill_(r.scales_3n[2, i].item())
+
+        r._history_idx = (r._history_idx + 1) % H
+        for m in r.modules:
+            m._history_idx = r._history_idx
+
+
+# ---------------------------------------------------------------------------
+# Weight extraction helper — called outside autograd.Function to avoid
+# tensor-subclass isinstance checks inside the compiled graph.
+# ---------------------------------------------------------------------------
+
+
+def _extract_fp8_weight(weight, fp8_dtype):
+    """Return (fp8_data, scale_inv) from weight, handling FP8UnshardedWeightTensor."""
+    from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+        FP8UnshardedWeightTensor,
+    )
+
+    if isinstance(weight, FP8UnshardedWeightTensor):
+        return weight.get_fp8_data_and_scale_inv()
+    return _quantize_fp8_tw(weight, fp8_dtype)
+
+
+# ---------------------------------------------------------------------------
+# Granularity -> Function dispatch map
+# ---------------------------------------------------------------------------
+
+_FP8_FN_MAP = {
+    ScalingGranularity.TENSORWISE: OpaqueFP8LinearTensorwiseFunction,
+    ScalingGranularity.ROWWISE: FP8LinearRowwiseFunction,
+    ScalingGranularity.BLOCKWISE: FP8LinearBlockwiseFunction,
+}
+
+# Public alias
+FP8LinearTensorwiseFunction = OpaqueFP8LinearTensorwiseFunction
+
+
+# ---------------------------------------------------------------------------
+# FP8-aware parallel linear layers
+# ---------------------------------------------------------------------------
+
+
+class _Float8LinearMixin:
+    """Shared per-module FP8 behavior for Float8{Column,Row}ParallelLinear.
+
+    This is a mixin (not an nn.Module). Each concrete class also inherits a
+    Megatron ``*ParallelLinear`` and MUST list this mixin FIRST in its bases so
+    the ``_apply`` / ``_forward_impl`` overrides below take precedence over the
+    Megatron base (otherwise FP8 would be silently disabled). The concrete
+    ``__init__`` calls ``self._init_fp8_state()`` after ``super().__init__()``.
+
+    Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False,
+    sequence_parallel=False. Fails fast at construction if violated.
+
+    Tensorwise uses the setup_context autograd pattern with pre-extracted FP8
+    weight data for graph-break-free torch.compile tracing.
+    Rowwise and blockwise still use @allow_in_graph.
+    """
+
+    def _init_fp8_state(self):
+        cls = type(self).__name__
+        if self.config.tensor_model_parallel_size != 1:
+            raise ValueError(
+                f"{cls} requires tensor_model_parallel_size=1. "
+                f"Got {self.config.tensor_model_parallel_size}."
+            )
+        if self.gradient_accumulation_fusion:
+            raise ValueError(f"{cls} requires gradient_accumulation_fusion=False.")
+        if self.sequence_parallel:
+            raise ValueError(f"{cls} requires sequence_parallel=False.")
+        if self.config.fp8 is None:
+            raise ValueError(f"{cls} requires config.fp8 to be set (e.g. 'e4m3').")
+
+        self._fp8_config = _build_fp8_config(self.config)
+        self._fp8_fn = _FP8_FN_MAP[self._fp8_config.granularity]
+
+        if self._fp8_config.granularity == ScalingGranularity.TENSORWISE:
+            self._fp8_fwd_dtype = _get_fp8_dtype(self._fp8_config.format, is_fwd=True)
+            self._fp8_bwd_dtype = _get_fp8_dtype(self._fp8_config.format, is_fwd=False)
+            self._fp8_gran_value = ScalingGranularity.TENSORWISE.value
+            self._fp8_backend_value = BackendType.HIPBLASLT.value
+            # Native layout (NN/TN) by default; forced-NT is opt-in per config.
+            # Read once at layer init so it stays a compile-time constant.
+            self._force_nt = getattr(self.config, "fp8_force_nt_layout", False)
+            self._use_delayed_scaling = (
+                getattr(self.config, "fp8_scaling_strategy", "dynamic") == "delayed"
+                or self.config.fp8_recipe == Fp8Recipe.delayed
+            )
+            if self._use_delayed_scaling:
+                _init_delayed_scaling_state(self)
+
+    def _apply(self, fn, recurse=True):
+        result = super()._apply(fn, recurse)
+        if getattr(self, "_use_delayed_scaling", False):
+            for name in [
+                "scale_input",
+                "scale_weight",
+                "scale_grad",
+                "amax_history_input",
+                "amax_history_weight",
+                "amax_history_grad",
+                "staged_input_amax",
+                "staged_grad_amax",
+                "staged_weight_amax",
+            ]:
+                buf = getattr(self, name, None)
+                if buf is not None and buf.dtype != torch.float32:
+                    self._buffers[name] = buf.data.float()
+        return result
+
+    def _forward_impl(self, input, weight, *args, **kwargs):
+        bias = kwargs.get("bias", None)
+
+        if self._fp8_config.granularity == ScalingGranularity.TENSORWISE:
+            if self._use_delayed_scaling:
+                result = DelayedFP8LinearTensorwiseFunction.apply(
+                    input,
+                    weight,
+                    self.scale_input,
+                    self.scale_weight,
+                    self.scale_grad,
+                    self.staged_input_amax,
+                    self.staged_weight_amax,
+                    self.staged_grad_amax,
+                    self._fp8_fwd_dtype,
+                    self._fp8_bwd_dtype,
+                    self._fp8_gran_value,
+                    self._fp8_backend_value,
+                    self._force_nt,
+                )
+                output = result[0]
+            else:
+                weight_fp8, weight_scale_inv = _extract_fp8_weight(
+                    weight,
+                    self._fp8_fwd_dtype,
+                )
+                result = self._fp8_fn.apply(
+                    input,
+                    weight,
+                    weight_fp8,
+                    weight_scale_inv,
+                    self._fp8_fwd_dtype,
+                    self._fp8_bwd_dtype,
+                    self._fp8_gran_value,
+                    self._fp8_backend_value,
+                    self._force_nt,
+                )
+                output = result[0]
+        else:
+            output = self._fp8_fn.apply(input, weight, self._fp8_config)
+
+        if bias is not None:
+            output = output + bias
+        return output
+
+
+class Float8ColumnParallelLinear(_Float8LinearMixin, ColumnParallelLinear):
+    """ColumnParallelLinear with per-module FP8. torch.compile friendly.
+
+    Shared FP8 behavior lives in :class:`_Float8LinearMixin`, which is listed
+    first so its ``_apply`` / ``_forward_impl`` overrides take effect.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._init_fp8_state()
+
+
+class Float8RowParallelLinear(_Float8LinearMixin, RowParallelLinear):
+    """RowParallelLinear with per-module FP8. torch.compile friendly.
+
+    Shared FP8 behavior lives in :class:`_Float8LinearMixin`, which is listed
+    first so its ``_apply`` / ``_forward_impl`` overrides take effect.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._init_fp8_state()
diff --git a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py
new file mode 100644
index 000000000..c800ac54b
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py
@@ -0,0 +1,318 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Primus Turbo Local Spec Provider
+
+Clean implementation using native Megatron modules with Primus Turbo
+optimizations. NO TransformerEngine dependencies.
+
+Key Features:
+- PrimusTurboLocalAttention: New class, inherits from MegatronModule (not TE)
+- Native Megatron linear layers (ColumnParallelLinear, RowParallelLinear)
+- Native LayerNorm (WrappedTorchNorm or FusedLayerNorm)
+- Maximum torch.compile compatibility
+- 52ms attention advantage preserved via pt.ops.flash_attn_func
+
+Performance Target (vs Inductor baseline):
+- Attention: Same (52ms advantage over Inductor's flash_attn)
+- GEMM: -7ms (native PyTorch proven within 0.8% of TE)
+- Elementwise: -157ms (torch.compile fusion)
+- Reduce: -17ms (torch.compile fusion)
+- Framework: -107ms (no TE overhead)
+- Total: ~1,131ms vs Inductor's 1,415ms = 20% faster
+"""
+
+import math
+import os
+from typing import Optional
+
+import primus_turbo.pytorch as pt
+import torch
+from primus_turbo.pytorch.ops.attention.flash_attn_interface import AiterFlashAttnFunc
+
+torch._dynamo.allow_in_graph(AiterFlashAttnFunc)
+
+
+@torch._dynamo.disable
+def _advance_model_parallel_rng(batch_size: int, num_heads: int) -> None:
+    """Simulate TE's CUDA RNG consumption inside
+    PrimusTurboLocalAttention.
+
+    TEDotProductAttention is instantiated with
+    get_rng_state_tracker=get_cuda_rng_tracker, so TE's fused_attn_fwd kernel
+    call runs inside tracker.fork() (the 'model-parallel-rng' generator state).
+    The ASM kernel pulls a fresh philox state per call regardless of dropout,
+    advancing model-parallel-rng by counter_offset = B*H*warp_size each
+    attention call. Aiter's _ndropout path does not touch RNG, so PT and
+    aiter attention see different downstream RNG sequences and diverge.
+
+    This helper replays the same RNG advance. Decorated with @dynamo.disable
+    so AOTAutograd never tries to trace through the generator manipulation.
+    """
+    warp_size = 64
+    counter_offset = batch_size * num_heads * warp_size
+    dev_idx = torch.cuda.current_device()
+    gens = torch.cuda.default_generators
+    if dev_idx >= len(gens):
+        return
+    gen = gens[dev_idx]
+
+    def _advance_with(g):
+        try:
+            g.set_offset(g.get_offset() + counter_offset)
+        except AttributeError:
+            torch.empty(counter_offset, device="cuda", dtype=torch.int32).random_(generator=g)
+
+    try:
+        from megatron.core.tensor_parallel.random import get_cuda_rng_tracker
+
+        tracker = get_cuda_rng_tracker()
+    except Exception:
+        tracker = None
+
+    if tracker is not None and getattr(tracker, "is_initialized", lambda: False)():
+        with tracker.fork():
+            _advance_with(gen)
+    else:
+        _advance_with(gen)
+
+
+from megatron.core.models.backends import LocalSpecProvider
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+    get_context_parallel_group,
+    get_tensor_model_parallel_group,
+)
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training.global_vars import get_args
+from torch import Tensor
+
+
+class PrimusTurboLocalAttention(MegatronModule):
+    """
+    Primus Turbo Flash Attention - Native Megatron Implementation.
+
+    This is a COMPLETELY NEW class with ZERO TransformerEngine dependencies.
+
+    Key differences from PrimusTurboAttention (in primus_turbo.py):
+    - Inherits from MegatronModule (NOT te.pytorch.DotProductAttention)
+    - No TE infrastructure or overhead
+    - torch.compile friendly (no @no_torch_dynamo decorator)
+    - Same 52ms performance advantage via pt.ops.flash_attn_func
+
+    Used exclusively by: PrimusTurboLocalSpecProvider
+
+    Performance:
+    - 52ms faster than Inductor's flash_attn implementation
+    - Same as existing PrimusTurboAttention but without TE overhead
+    - torch.compile compatible for surrounding operations
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        layer_number: int,
+        attn_mask_type: AttnMaskType,
+        attention_type: str,
+        attention_dropout: Optional[float] = None,
+        softmax_scale: Optional[float] = None,
+        k_channels: Optional[int] = None,
+        v_channels: Optional[int] = None,
+        cp_comm_type: str = "p2p",
+        pg_collection: Optional[ProcessGroupCollection] = None,
+    ):
+        super().__init__(config=config)
+
+        self.config = config
+        self.layer_number = layer_number
+        self.attn_mask_type = attn_mask_type
+        self.attention_type = attention_type
+
+        # Calculate softmax scale
+        kv_channels = k_channels if k_channels is not None else config.kv_channels
+        self.softmax_scale = softmax_scale or (1.0 / math.sqrt(kv_channels))
+
+        # Setup process groups
+        if pg_collection is None:
+            pg_collection = ProcessGroupCollection(
+                tp=get_tensor_model_parallel_group(check_initialized=False),
+                cp=get_context_parallel_group(check_initialized=False),
+            )
+
+        # Select Primus Turbo flash attention variant
+        args = get_args()
+        if args.enable_turbo_attention_float8:
+            self.attn_func = (
+                pt.ops.flash_attn_fp8_usp_func
+                if config.context_parallel_size > 1
+                else pt.ops.flash_attn_fp8_func
+            )
+        else:
+            self.attn_func = (
+                pt.ops.flash_attn_usp_func if config.context_parallel_size > 1 else pt.ops.flash_attn_func
+            )
+
+        # The transpose in forward() produces a non-contiguous, SBHD-strided view.
+        # aiter's v3 flash-attention backward mishandles that strided layout on
+        # gfx942 (MI300X), corrupting gradients and causing grad-norm divergence
+        # ~step 31; gfx950 (MI355X) handles it correctly. Feed contiguous BSHD on
+        # gfx942 only, so gfx950 keeps the faster zero-copy strided path unchanged.
+        # Computed here (not in forward) so torch.compile sees a constant guard
+        # rather than a device query in the traced region.
+        self.force_contiguous_qkv = torch.cuda.get_device_capability() < (9, 5)
+
+        # Setup context parallel arguments
+        self.attn_kwargs = {}
+        if config.context_parallel_size > 1:
+            self.attn_kwargs["ulysses_group"] = pg_collection.cp
+
+        # Validate configuration
+        if config.window_size is not None:
+            raise ValueError("PrimusTurboLocalAttention does not support sliding window attention")
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        attention_mask: Tensor,
+        attn_mask_type: AttnMaskType,
+        attention_bias: Optional[Tensor] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> Tensor:
+        """
+        Forward pass using Primus Turbo flash attention.
+
+        Args:
+            query: Query tensor [seq_len, batch, num_heads, head_dim] (sbhd)
+            key: Key tensor [seq_len, batch, num_heads, head_dim] (sbhd)
+            value: Value tensor [seq_len, batch, num_heads, head_dim] (sbhd)
+            attention_mask: Attention mask (not used by flash attention)
+            attn_mask_type: Type of attention mask (causal, no_mask, etc.)
+            attention_bias: Attention bias (not used in this implementation)
+            packed_seq_params: Packed sequence parameters (optional)
+
+        Returns:
+            Attention output [seq_len, batch, num_heads * head_dim] (merged heads)
+        """
+        query, key, value = [x.transpose(0, 1) for x in (query, key, value)]
+
+        # gfx942: avoid aiter's broken strided-sbhd backward (see __init__).
+        if self.force_contiguous_qkv:
+            query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+
+        causal = attn_mask_type == AttnMaskType.causal
+
+        if os.environ.get("PRIMUS_PT_MIMIC_TE_RNG", "0") == "1":
+            B = query.size(0)
+            H = query.size(2)
+            _advance_model_parallel_rng(B, H)
+
+        output = self.attn_func(
+            query,
+            key,
+            value,
+            dropout_p=0.0,
+            softmax_scale=self.softmax_scale,
+            causal=causal,
+            window_size=(-1, -1),
+            bias=None,
+            alibi_slopes=None,
+            deterministic=False,
+            return_lse=False,
+            return_attn_probs=False,
+            **self.attn_kwargs,
+        )
+
+        # Transpose back to Megatron format (bshd -> sbhd) and merge heads
+        output = output.transpose(0, 1)
+        output = output.reshape(output.shape[0], output.shape[1], -1)
+
+        return output
+
+
+class PrimusTurboMXFP4LocalSpecProvider(LocalSpecProvider):
+    """
+    Compile-friendly MXFP4 spec: Primus Turbo attention + MXFP4 linear layers.
+    NO TransformerEngine. NO global FP4 state in forward path.
+    Requires tensor_model_parallel_size=1.
+    """
+
+    def column_parallel_linear(self) -> type:
+        from .primus_turbo_mxfp4_local import MXFP4ColumnParallelLinear
+
+        return MXFP4ColumnParallelLinear
+
+    def row_parallel_linear(self) -> type:
+        from .primus_turbo_mxfp4_local import MXFP4RowParallelLinear
+
+        return MXFP4RowParallelLinear
+
+    def core_attention(self) -> type:
+        return PrimusTurboLocalAttention
+
+
+class PrimusTurboFloat8LocalSpecProvider(LocalSpecProvider):
+    """
+    Compile-friendly FP8 spec: Primus Turbo attention + FP8 linear layers.
+    NO TransformerEngine. NO global FP8 state in forward path.
+    NO changes to Primus-Turbo repo. Requires tensor_model_parallel_size=1.
+    """
+
+    def column_parallel_linear(self) -> type:
+        from .primus_turbo_float8_local import Float8ColumnParallelLinear
+
+        return Float8ColumnParallelLinear
+
+    def row_parallel_linear(self) -> type:
+        from .primus_turbo_float8_local import Float8RowParallelLinear
+
+        return Float8RowParallelLinear
+
+    def core_attention(self) -> type:
+        return PrimusTurboLocalAttention
+
+
+class PrimusTurboLocalSpecProvider(LocalSpecProvider):
+    """
+    Spec provider that extends LocalSpecProvider with Primus Turbo attention.
+
+    Uses native Megatron modules everywhere except attention, where it uses
+    PrimusTurboLocalAttention for the 52ms performance advantage.
+
+    Key features:
+    - NO TransformerEngine dependencies
+    - Maximum torch.compile compatibility
+    - Minimal custom code (~120 lines total)
+    - Native PyTorch GEMM (proven within 0.8% of TE)
+    - Primus Turbo flash attention (52ms advantage)
+
+    Configuration:
+        config.use_primus_turbo_local_spec = True
+
+    Performance Target:
+    - Attention: -52ms (Primus Turbo advantage over Inductor)
+    - Operator fusion: -180ms (torch.compile)
+    - Framework overhead: -107ms (no TE)
+    - GEMM: ~same (native PyTorch excellent)
+    - Total: ~1,131ms vs Inductor's 1,415ms (20% faster)
+
+    vs Current TE implementation:
+    - Total: ~1,131ms vs 1,588ms (28.8% faster)
+    """
+
+    def core_attention(self) -> type:
+        """Return Primus Turbo local attention (NO TE dependency)"""
+        return PrimusTurboLocalAttention
+
+    # Everything else inherited from LocalSpecProvider:
+    # - column_parallel_linear() -> ColumnParallelLinear (native Megatron)
+    # - row_parallel_linear() -> RowParallelLinear (native Megatron)
+    # - layer_norm() -> WrappedTorchNorm/FusedLayerNorm (compile-friendly)
+    # - fuse_layernorm_and_linear() -> False (no fusion overhead)
+    # - grouped_mlp_modules() -> SequentialMLP (native modules)
+    # - activation_func() -> None (standard activation)
diff --git a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py
new file mode 100644
index 000000000..73f5ed68b
--- /dev/null
+++ b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py
@@ -0,0 +1,586 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Compile-friendly MXFP4 linear layers for Megatron local spec.
+
+Self-contained autograd Functions that call Primus Turbo's low-level
+quantize_mxfp4_dual / quantize_mxfp4 C++ ops and gemm_fp4_impl directly,
+bypassing the higher-level wrappers that construct ScalingRecipe objects
+and call check_mxfp4_support() global state.
+
+Key properties:
+- Uses setup_context pattern with primitive-only args so torch.compile
+  can trace through without graph breaks.
+- Two backward modes: pure MXFP4 or hybrid (FP4 fwd / FP8 bwd).
+- gemm_fp4_impl and gemm_fp8_impl are torch.library.custom_op with register_fake.
+- Zero TransformerEngine dependencies.
+- Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel.
+"""
+
+from typing import Tuple
+
+import torch
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from primus_turbo.pytorch.core.backend import (
+    BackendType,
+    GlobalBackendManager,
+    PrecisionType,
+)
+from primus_turbo.pytorch.core.low_precision import (
+    ScalingGranularity,
+    check_mxfp4_support,
+)
+from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import gemm_fp4_impl
+from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl
+
+from .primus_turbo_float8_local import _quantize_fp8_tw
+
+# The AITER MXFP4 preshuffle fast path used to be implemented here as a
+# monkey patch on GEMMFP4AITERBackend.execute. It now lives in Primus-Turbo
+# (commit 683a7de, "feat(gemm): add AITER MXFP4 preshuffle fast path") and
+# is opted into per-call via the gemm_fp4_impl(preshuffled=...) kwarg.
+#
+# Primus-Turbo PR #383 ("refactor preshuffle ...") removed the public
+# enable_preshuffle() helper and moved per-call preshuffle control onto
+# Float4QuantConfig.use_preshuffle. MXFP4LinearFunction passes a plain bool
+# into its custom ops (not a Float4QuantConfig), so we keep the original
+# runtime probe here as a module-local helper, _enable_preshuffle(), which
+# reproduces the removed upstream logic verbatim.
+#
+# We resolve _enable_preshuffle() once per MXFP4 linear at __init__ time and
+# cache the result on self._preshuffle. If the dispatcher state changes
+# after model construction (env unset, set_gemm_backend(None), autotune
+# flipped on) the cached flag is stale and the forward path will
+# mis-dispatch -- the same caching pattern as before, just now made loud by
+# the module-init contract check below (_assert_preshuffle_contract, which
+# raises RuntimeError). Resolving it at __init__ (not inside forward)
+# also keeps it out of the torch.compile traced region.
+
+
+def _enable_preshuffle() -> bool:
+    """True iff the AITER FP4 preshuffle fast path is safe.
+
+    Requires: (1) the FP4 GEMM backend is explicitly pinned to AITER (the only
+    backend that understands the shuffled layout), and (2) autotune is disabled
+    (AITER opts out of tuning, so the tuner cannot select a backend for
+    preshuffled inputs). Reproduces the removed Primus-Turbo enable_preshuffle()
+    (pre Primus-Turbo PR #383) so MXFP4LinearFunction can keep passing a bool
+    into its custom ops.
+    """
+    return (
+        GlobalBackendManager.get_gemm_backend(PrecisionType.FP4) == BackendType.AITER
+        and not GlobalBackendManager.auto_tune_enabled()
+    )
+
+
+def _assert_preshuffle_contract(config, preshuffle: bool) -> None:
+    """Fail loudly when MXFP4 was explicitly requested but the upstream
+    AITER preshuffle fast path is unavailable.
+
+    Pre-cleanup, the local monkey patch silently consumed preshuffled
+    inputs on the AITER backend regardless of dispatcher state, masking
+    misconfigurations as either a ~5% step-time regression (preshuffle
+    inputs landed on AITER but every GEMM paid the shuffle cost) or, on
+    dispatch paths 2/3/4, silent numerical corruption (HipBLASLt received
+    preshuffled bytes and ran with them). Post-cleanup, both modes become
+    a hard init-time failure with an actionable message.
+
+    Gated on ``config.fp4 == "mxfp4"`` so any speculative-instantiation
+    flow that constructs the class without actually opting into MXFP4 is
+    unaffected (mirrors the precedent of the other ``__init__`` asserts
+    in this file).
+    """
+    if getattr(config, "fp4", None) != "mxfp4":
+        return
+    if not preshuffle:
+        raise RuntimeError(
+            "MXFP4 linear requested (config.fp4='mxfp4') but the AITER "
+            "preshuffle fast path is unavailable. Set "
+            "PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER (or call "
+            "GlobalBackendManager.set_gemm_backend(BackendType.AITER, "
+            "PrecisionType.FP4)) and leave PRIMUS_TURBO_AUTO_TUNE unset. "
+            "Without this, every FP4 GEMM pays ~3 extra shuffle kernel "
+            "launches (~5% Flux 12B step-time regression vs the tuned "
+            "baseline) -- and on non-AITER dispatch the call now raises."
+        )
+
+
+# ---------------------------------------------------------------------------
+# torch.compile-friendly wrappers for MXFP4 C++ quantization ops.
+#
+# The raw C++ ops lack an Autograd dispatch key. Wrapping them with
+# @torch.library.custom_op provides Autograd dispatch and register_fake
+# for shape inference during tracing.
+# ---------------------------------------------------------------------------
+
+_custom_op = torch.library.custom_op
+
+MXFP4_BLOCK_SIZE = 32
+# Bumped from 16 -> 128 to match Primus-Turbo PR #335 ("feat: add quantized
+# tensor support"), which made `padding_align_size` an explicit positional arg
+# of quantize_mxfp4{_dual} and asserts it equals MXFP4_PADDING_ALIGN_SIZE (=128)
+# in csrc/include/primus_turbo/quantization.h.
+MXFP4_PADDING_ALIGN_SIZE = 128
+
+
+def _cdiv(a, b):
+    return (a + b - 1) // b
+
+
+@_custom_op("primus::quantize_mxfp4_dual", mutates_args=(), device_types="cuda")
+def _quantize_mxfp4_dual_op(
+    x: torch.Tensor,
+    out_dtype: torch.dtype,
+    padding_align_size: int,
+    rowwise_use_2d_block: bool,
+    rowwise_use_sr: bool,
+    rowwise_use_rht: bool,
+    colwise_use_2d_block: bool,
+    colwise_use_sr: bool,
+    colwise_use_rht: bool,
+    shuffle_rowwise_scale: bool,
+    shuffle_rowwise: bool,
+    shuffle_colwise_scale: bool,
+    shuffle_colwise: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    return torch.ops.primus_turbo_cpp_extension.quantize_mxfp4_dual(
+        x,
+        out_dtype,
+        padding_align_size,
+        rowwise_use_2d_block,
+        rowwise_use_sr,
+        rowwise_use_rht,
+        colwise_use_2d_block,
+        colwise_use_sr,
+        colwise_use_rht,
+        shuffle_rowwise_scale,
+        shuffle_rowwise,
+        shuffle_colwise_scale,
+        shuffle_colwise,
+    )
+
+
+@_quantize_mxfp4_dual_op.register_fake
+def _quantize_mxfp4_dual_fake(
+    x: torch.Tensor,
+    out_dtype: torch.dtype,
+    padding_align_size: int,
+    rowwise_use_2d_block: bool,
+    rowwise_use_sr: bool,
+    rowwise_use_rht: bool,
+    colwise_use_2d_block: bool,
+    colwise_use_sr: bool,
+    colwise_use_rht: bool,
+    shuffle_rowwise_scale: bool,
+    shuffle_rowwise: bool,
+    shuffle_colwise_scale: bool,
+    shuffle_colwise: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    M, N = x.shape
+    M_pad = _cdiv(M, MXFP4_PADDING_ALIGN_SIZE) * MXFP4_PADDING_ALIGN_SIZE
+    N_pad = _cdiv(N, MXFP4_PADDING_ALIGN_SIZE) * MXFP4_PADDING_ALIGN_SIZE
+
+    if shuffle_rowwise_scale:
+        rs_M = _cdiv(M, 256) * 256
+        rs_N = _cdiv(_cdiv(N_pad, MXFP4_BLOCK_SIZE), 8) * 8
+        rowwise_scale = torch.empty(rs_M, rs_N, dtype=torch.uint8, device=x.device)
+    else:
+        rowwise_scale = torch.empty(M, _cdiv(N_pad, MXFP4_BLOCK_SIZE), dtype=torch.uint8, device=x.device)
+
+    rowwise_output = torch.empty(M, N_pad // 2, dtype=torch.uint8, device=x.device)
+
+    if shuffle_colwise_scale:
+        cs_M = _cdiv(N, 256) * 256
+        cs_N = _cdiv(_cdiv(M_pad, MXFP4_BLOCK_SIZE), 8) * 8
+        colwise_scale = torch.empty(cs_M, cs_N, dtype=torch.uint8, device=x.device)
+    else:
+        colwise_scale = torch.empty(N, _cdiv(M_pad, MXFP4_BLOCK_SIZE), dtype=torch.uint8, device=x.device)
+
+    colwise_output = torch.empty(N, M_pad // 2, dtype=torch.uint8, device=x.device)
+
+    return (
+        rowwise_output.view(torch.float4_e2m1fn_x2),
+        rowwise_scale.view(torch.float8_e8m0fnu),
+        colwise_output.view(torch.float4_e2m1fn_x2),
+        colwise_scale.view(torch.float8_e8m0fnu),
+    )
+
+
+def _quantize_mxfp4_dual_setup_context(ctx, inputs, output):
+    pass
+
+
+def _quantize_mxfp4_dual_backward(ctx, *grad_outputs):
+    return (None,) * 13
+
+
+_quantize_mxfp4_dual_op.register_autograd(
+    _quantize_mxfp4_dual_backward,
+    setup_context=_quantize_mxfp4_dual_setup_context,
+)
+
+
+# ---------------------------------------------------------------------------
+# MXFP4 Autograd Function with setup_context (torch.compile friendly)
+# ---------------------------------------------------------------------------
+
+
+_FP4_DTYPE = torch.float4_e2m1fn_x2
+_GRAN_VALUE = ScalingGranularity.MX_BLOCKWISE.value
+_DEFAULT_BACKEND = BackendType.HIPBLASLT.value
+
+
+def _quantize_input_dual(input_2d, preshuffle):
+    """Quantize input (activation) with dual rowwise + colwise."""
+    return _quantize_mxfp4_dual_op(
+        input_2d,
+        _FP4_DTYPE,
+        MXFP4_PADDING_ALIGN_SIZE,
+        False,
+        False,
+        False,  # rowwise: no 2d_block, no sr, no rht
+        False,
+        False,
+        True,  # colwise: no 2d_block, no sr, yes rht
+        preshuffle,
+        False,  # shuffle_rowwise_scale, shuffle_rowwise
+        preshuffle,
+        preshuffle,  # shuffle_colwise_scale, shuffle_colwise
+    )
+
+
+def _quantize_weight_dual(weight, preshuffle):
+    """Quantize weight with dual rowwise + colwise."""
+    return _quantize_mxfp4_dual_op(
+        weight,
+        _FP4_DTYPE,
+        MXFP4_PADDING_ALIGN_SIZE,
+        True,
+        False,
+        False,  # rowwise: 2d_block, no sr, no rht
+        True,
+        False,
+        False,  # colwise: 2d_block, no sr, no rht
+        preshuffle,
+        preshuffle,  # shuffle_rowwise_scale, shuffle_rowwise
+        preshuffle,
+        preshuffle,  # shuffle_colwise_scale, shuffle_colwise
+    )
+
+
+def _quantize_grad_dual(grad_2d, preshuffle, use_sr=True):
+    """Quantize gradient (activation recipe) with dual rowwise + colwise."""
+    return _quantize_mxfp4_dual_op(
+        grad_2d,
+        _FP4_DTYPE,
+        MXFP4_PADDING_ALIGN_SIZE,
+        False,
+        use_sr,
+        False,  # rowwise: no 2d_block, SR configurable, no rht
+        False,
+        use_sr,
+        True,  # colwise: no 2d_block, SR configurable, yes rht
+        preshuffle,
+        False,  # shuffle_rowwise_scale, shuffle_rowwise
+        preshuffle,
+        False,  # shuffle_colwise_scale, shuffle_colwise
+    )
+
+
+class MXFP4LinearFunction(torch.autograd.Function):
+    """MXFP4 linear (Y = X @ W^T) with MX block-of-32 scaling.
+
+    Two modes via backward_is_fp8 bool primitive:
+    - Pure MXFP4: forward + backward both use FP4 quantization + gemm_fp4_impl
+    - Hybrid: forward uses FP4, backward re-quantizes saved BF16 to FP8 tensorwise
+
+    Uses setup_context pattern with primitive-only args for torch.compile.
+    """
+
+    @staticmethod
+    def forward(
+        input,
+        weight,
+        preshuffle,
+        backward_is_fp8,
+        fp8_bwd_dtype,
+        fp8_gran_value,
+        fp8_backend_value,
+        use_gradient_sr,
+    ):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp4, a_scale, a_t_fp4, a_t_scale = _quantize_input_dual(input_2d, preshuffle)
+        b_fp4, b_scale, b_t_fp4, b_t_scale = _quantize_weight_dual(weight, preshuffle)
+
+        output = gemm_fp4_impl(
+            a_fp4,
+            a_scale,
+            False,
+            b_fp4,
+            b_scale,
+            True,
+            out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_DEFAULT_BACKEND,
+            preshuffled=preshuffle,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        if backward_is_fp8:
+            return output, input_2d.view_as(input_2d), weight.view_as(weight)
+        else:
+            # Return FP4/FP8 tensors as uint8 views to avoid
+            # "fill_cuda not implemented for Float4_e2m1fn_x2" when
+            # the autograd engine creates zero gradients for
+            # non-differentiable outputs.
+            return (
+                output,
+                a_t_fp4.view(torch.uint8),
+                a_t_scale.view(torch.uint8),
+                b_t_fp4.view(torch.uint8),
+                b_t_scale.view(torch.uint8),
+            )
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        (
+            _,
+            _,
+            preshuffle,
+            backward_is_fp8,
+            fp8_bwd_dtype,
+            fp8_gran_value,
+            fp8_backend_value,
+            use_gradient_sr,
+        ) = inputs
+
+        ctx.preshuffle = preshuffle
+        ctx.backward_is_fp8 = backward_is_fp8
+        ctx.use_gradient_sr = use_gradient_sr
+        ctx.out_dtype = inputs[0].dtype
+        ctx.orig_shape = inputs[0].shape
+
+        if backward_is_fp8:
+            _, input_2d_saved, weight_saved = output
+            ctx.save_for_backward(input_2d_saved, weight_saved)
+            ctx.fp8_bwd_dtype = fp8_bwd_dtype
+            ctx.fp8_gran_value = fp8_gran_value
+            ctx.fp8_backend_value = fp8_backend_value
+        else:
+            _, a_t_u8, as_u8, b_t_u8, bs_u8 = output
+            ctx.save_for_backward(
+                a_t_u8.view(_FP4_DTYPE),
+                as_u8.view(torch.float8_e8m0fnu),
+                b_t_u8.view(_FP4_DTYPE),
+                bs_u8.view(torch.float8_e8m0fnu),
+            )
+            ctx.mark_non_differentiable(a_t_u8, as_u8, b_t_u8, bs_u8)
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+
+        if ctx.backward_is_fp8:
+            input_2d, weight = ctx.saved_tensors
+
+            grad_fp8, grad_scale_inv = _quantize_fp8_tw(grad_2d, ctx.fp8_bwd_dtype)
+            a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, ctx.fp8_bwd_dtype)
+            b_fp8, b_scale_inv = _quantize_fp8_tw(weight, ctx.fp8_bwd_dtype)
+
+            grad_input = gemm_fp8_impl(
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                b_fp8,
+                b_scale_inv,
+                False,
+                ctx.out_dtype,
+                False,
+                granularity=ctx.fp8_gran_value,
+                default_backend=ctx.fp8_backend_value,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            grad_weight = gemm_fp8_impl(
+                a_fp8,
+                a_scale_inv,
+                True,
+                grad_fp8,
+                grad_scale_inv,
+                False,
+                ctx.out_dtype,
+                True,
+                granularity=ctx.fp8_gran_value,
+                default_backend=ctx.fp8_backend_value,
+            )
+        else:
+            a_t_fp4, a_t_scale, b_t_fp4, b_t_scale = ctx.saved_tensors
+            preshuffle = ctx.preshuffle
+
+            g_fp4, g_scale, g_t_fp4, g_t_scale = _quantize_grad_dual(
+                grad_2d, preshuffle, use_sr=ctx.use_gradient_sr
+            )
+
+            grad_input = gemm_fp4_impl(
+                g_fp4,
+                g_scale,
+                False,
+                b_t_fp4,
+                b_t_scale,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=_GRAN_VALUE,
+                default_backend=_DEFAULT_BACKEND,
+                preshuffled=preshuffle,
+            )
+            grad_input = grad_input.reshape(ctx.orig_shape)
+
+            grad_weight = gemm_fp4_impl(
+                g_t_fp4,
+                g_t_scale,
+                False,
+                a_t_fp4,
+                a_t_scale,
+                True,
+                ctx.out_dtype,
+                False,
+                granularity=_GRAN_VALUE,
+                default_backend=_DEFAULT_BACKEND,
+                preshuffled=preshuffle,
+            )
+
+        return grad_input, grad_weight, None, None, None, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# MXFP4-aware parallel linear layers
+# ---------------------------------------------------------------------------
+
+
+class MXFP4ColumnParallelLinear(ColumnParallelLinear):
+    """ColumnParallelLinear with per-module MXFP4. torch.compile friendly.
+
+    Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False,
+    sequence_parallel=False.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self.config.tensor_model_parallel_size != 1:
+            raise ValueError(
+                "MXFP4ColumnParallelLinear requires tensor_model_parallel_size=1. "
+                f"Got {self.config.tensor_model_parallel_size}."
+            )
+        if self.gradient_accumulation_fusion:
+            raise ValueError("MXFP4ColumnParallelLinear requires gradient_accumulation_fusion=False.")
+        if self.sequence_parallel:
+            raise ValueError("MXFP4ColumnParallelLinear requires sequence_parallel=False.")
+
+        supported, reason = check_mxfp4_support()
+        if not supported:
+            raise RuntimeError(f"MXFP4 not supported on this device: {reason}")
+
+        self._preshuffle = _enable_preshuffle()
+        _assert_preshuffle_contract(self.config, self._preshuffle)
+        self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8"
+        self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False)
+
+        if self._backward_is_fp8:
+            from primus_turbo.pytorch.core.low_precision import float8_e5m2
+
+            self._fp8_bwd_dtype = float8_e5m2
+            self._fp8_gran_value = ScalingGranularity.TENSORWISE.value
+            self._fp8_backend_value = BackendType.HIPBLASLT.value
+        else:
+            self._fp8_bwd_dtype = None
+            self._fp8_gran_value = 0
+            self._fp8_backend_value = 0
+
+    def _forward_impl(self, input, weight, *args, **kwargs):
+        bias = kwargs.get("bias", None)
+
+        result = MXFP4LinearFunction.apply(
+            input,
+            weight,
+            self._preshuffle,
+            self._backward_is_fp8,
+            self._fp8_bwd_dtype,
+            self._fp8_gran_value,
+            self._fp8_backend_value,
+            self._use_gradient_sr,
+        )
+        output = result[0]
+
+        if bias is not None:
+            output = output + bias
+        return output
+
+
+class MXFP4RowParallelLinear(RowParallelLinear):
+    """RowParallelLinear with per-module MXFP4. torch.compile friendly.
+
+    Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False,
+    sequence_parallel=False.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self.config.tensor_model_parallel_size != 1:
+            raise ValueError(
+                "MXFP4RowParallelLinear requires tensor_model_parallel_size=1. "
+                f"Got {self.config.tensor_model_parallel_size}."
+            )
+        if self.gradient_accumulation_fusion:
+            raise ValueError("MXFP4RowParallelLinear requires gradient_accumulation_fusion=False.")
+        if self.sequence_parallel:
+            raise ValueError("MXFP4RowParallelLinear requires sequence_parallel=False.")
+
+        supported, reason = check_mxfp4_support()
+        if not supported:
+            raise RuntimeError(f"MXFP4 not supported on this device: {reason}")
+
+        self._preshuffle = _enable_preshuffle()
+        _assert_preshuffle_contract(self.config, self._preshuffle)
+        self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8"
+        self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False)
+
+        if self._backward_is_fp8:
+            from primus_turbo.pytorch.core.low_precision import float8_e5m2
+
+            self._fp8_bwd_dtype = float8_e5m2
+            self._fp8_gran_value = ScalingGranularity.TENSORWISE.value
+            self._fp8_backend_value = BackendType.HIPBLASLT.value
+        else:
+            self._fp8_bwd_dtype = None
+            self._fp8_gran_value = 0
+            self._fp8_backend_value = 0
+
+    def _forward_impl(self, input, weight, *args, **kwargs):
+        bias = kwargs.get("bias", None)
+
+        result = MXFP4LinearFunction.apply(
+            input,
+            weight,
+            self._preshuffle,
+            self._backward_is_fp8,
+            self._fp8_bwd_dtype,
+            self._fp8_gran_value,
+            self._fp8_backend_value,
+            self._use_gradient_sr,
+        )
+        output = result[0]
+
+        if bias is not None:
+            output = output + bias
+        return output
diff --git a/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py b/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py
deleted file mode 100644
index a529aa7d2..000000000
--- a/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py
+++ /dev/null
@@ -1,736 +0,0 @@
-# This file was modified for portability to AMDGPU
-# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
-# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
-# See LICENSE for license information.
-
-"""Linear API"""
-import functools
-from functools import reduce
-from operator import mul as multiply_op
-from typing import Optional, Tuple, Union
-
-import torch
-import transformer_engine_torch as tex
-from transformer_engine.pytorch.constants import dist_group_type
-from transformer_engine.pytorch.cpp_extensions import general_gemm
-from transformer_engine.pytorch.cpu_offload import set_offloading_param
-from transformer_engine.pytorch.distributed import (
-    _fsdp_gather_tensors,
-    _fsdp_scatter_tensors,
-    allreduce,
-    gather_along_first_dim,
-    get_distributed_world_size,
-    in_fp8_activation_recompute_phase,
-    is_fp8_activation_recompute_enabled,
-    reduce_scatter_along_first_dim,
-)
-from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
-from transformer_engine.pytorch.graph import is_graph_capturing
-from transformer_engine.pytorch.module._common import _fix_gathered_fp8_transpose
-from transformer_engine.pytorch.module.base import (
-    _2X_ACC_DGRAD,
-    _2X_ACC_FPROP,
-    _2X_ACC_WGRAD,
-    TransformerEngineBaseModule,
-    get_ub,
-    get_workspace,
-)
-from transformer_engine.pytorch.rocm_utils import (
-    clear_fp8_weight_transpose_cache,
-    create_fp8_weight_transpose_cache,
-)
-
-try:  # TE >= 2.12: base classes moved to tensor.storage and renamed *Storage
-    from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import (
-        MXFP8TensorStorage as MXFP8TensorBase,
-    )
-except ModuleNotFoundError:  # TE <= 2.8
-    from transformer_engine.pytorch.tensor._internal.mxfp8_tensor_base import (
-        MXFP8TensorBase,
-    )
-
-from transformer_engine.pytorch.utils import (
-    assert_dim_for_fp8_exec,
-    cast_if_needed,
-    clear_tensor_data,
-    non_tn_fp8_gemm_supported,
-    nvtx_range_pop,
-    nvtx_range_push,
-    requires_grad,
-)
-
-from primus.backends.megatron.core.pipeline_parallel.wgrad_adapter import (
-    insert_wgrad_func_into_cache,
-)
-
-try:
-    from transformer_engine.pytorch.tensor.quantized_tensor import (
-        QuantizedTensor,
-        Quantizer,
-        prepare_for_saving,
-        restore_from_saved,
-    )
-except ModuleNotFoundError:
-    from transformer_engine.pytorch.quantized_tensor import (
-        QuantizedTensor,
-        Quantizer,
-        prepare_for_saving,
-        restore_from_saved,
-    )
-
-
-class _LinearWithWGradSplit(torch.autograd.Function):
-    """Linear semi-top level module
-    Calls custom cuda extensions.
-    """
-
-    @staticmethod
-    def forward(
-        ctx,
-        weight: torch.Tensor,
-        inp: torch.Tensor,
-        bias: Optional[torch.Tensor],
-        is_first_microbatch: Union[bool, None],
-        fp8: bool,
-        fp8_calibration: bool,
-        input_quantizer: Optional[Quantizer],
-        weight_quantizer: Optional[Quantizer],
-        output_quantizer: Optional[Quantizer],
-        grad_output_quantizer: Optional[Quantizer],
-        grad_input_quantizer: Optional[Quantizer],
-        fuse_wgrad_accumulation: bool,
-        cpu_offloading: bool,
-        tp_group: Union[dist_group_type, None],
-        tp_size: int,
-        sequence_parallel: bool,
-        tensor_parallel: bool,
-        activation_dtype: torch.dtype,
-        parallel_mode: Union[str, None],
-        is_grad_enabled: bool,
-        ub_overlap_rs_fprop: bool,
-        ub_overlap_ag_dgrad: bool,
-        ub_overlap_ag_fprop: bool,
-        ub_overlap_rs_dgrad: bool,
-        ub_bulk_dgrad: bool,
-        ub_bulk_wgrad: bool,
-        ub_name: str,
-        fp8_output: bool,  # pylint: disable=unused-argument
-        fsdp_group: Union[dist_group_type, None],
-        module: torch.nn.Module,
-        skip_fp8_weight_update: bool,
-        keep_fp8_weight_transpose_cache: bool,
-    ) -> torch.Tensor:
-        # pylint: disable=missing-function-docstring
-        assert not ub_bulk_wgrad, "not support for ZeroBubble"
-        assert bias is None, "not support bias yet"
-
-        # NVTX label for profiling
-        nvtx_label = "transformer_engine._Linear.forward"
-        if ub_name is not None:
-            nvtx_label = f"{nvtx_label}.{ub_name}"
-
-        # Make sure input dimensions are compatible
-        out_features, in_features = weight.shape
-        inp_shape = inp.shape
-        assert inp_shape[-1] == in_features, "GEMM not possible"
-
-        tp_world_size = get_distributed_world_size(tp_group)
-        backward_needs_input = is_grad_enabled and weight.requires_grad
-
-        # Prepare input tensor
-        # Note: Cast to expected dtype and perform tensor-parallel communication
-        nvtx_range_push(f"{nvtx_label}.input_cast_comm")
-        inputmat = inp.view(-1, in_features)
-        inputmat_total = None
-        with_input_all_gather_nccl = (
-            parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop
-        )
-        own_quantized_input = False
-        if fp8:
-            assert_dim_for_fp8_exec(inputmat, weight)
-            if any([ub_overlap_ag_fprop, ub_overlap_rs_fprop]) and not (
-                FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling()
-            ):
-                raise NotImplementedError(
-                    "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor"
-                    " current scaling"
-                )
-
-            if input_quantizer is None:
-                raise ValueError("Missing quantizer for input tensor")
-            if with_input_all_gather_nccl:
-                assert not isinstance(inputmat, QuantizedTensor), "All gather of fp8 input is not supported"
-                input_quantizer.set_usage(rowwise=True, columnwise=False)
-                inputmat_total, _ = gather_along_first_dim(
-                    inputmat,
-                    tp_group,
-                    quantizer=input_quantizer,
-                )
-            else:
-                if FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() and ub_bulk_dgrad:
-                    # reduce duplicated transpose in `_fix_gathered_fp8_transpose`
-                    input_quantizer.set_usage(rowwise=True, columnwise=False)
-                else:
-                    input_quantizer.set_usage(
-                        rowwise=True,
-                        columnwise=backward_needs_input,
-                    )
-                if not isinstance(inputmat, QuantizedTensor):
-                    inputmat = input_quantizer(inputmat)
-                    own_quantized_input = True
-                elif backward_needs_input:
-                    inputmat.update_usage(rowwise_usage=True, columnwise_usage=True)
-                inputmat_total = inputmat
-        else:
-            inputmat = cast_if_needed(inp, activation_dtype)
-            if with_input_all_gather_nccl:
-                inputmat_total, _ = gather_along_first_dim(inputmat, tp_group)
-            else:
-                inputmat_total = inputmat
-        nvtx_range_pop(f"{nvtx_label}.input_cast_comm")
-
-        # Cast weight to expected dtype
-        weightmat = weight
-        if not fp8:
-            weightmat = cast_if_needed(weightmat, activation_dtype)
-        else:
-            if not isinstance(weight, QuantizedTensor):
-                # Configure quantizer
-                if weight_quantizer is not None:
-                    columnwise_usage = is_grad_enabled and inp.requires_grad
-                    if not columnwise_usage:
-                        columnwise_usage = (
-                            is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase()
-                        )
-                    weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
-
-                # FP8 cast to workspace buffer
-                update_workspace = is_first_microbatch is None or is_first_microbatch
-                weightmat = module.get_weight_workspace(
-                    tensor=weight,
-                    quantizer=weight_quantizer,
-                    cache_name=(None if is_first_microbatch is None else "weight"),
-                    update_workspace=update_workspace,
-                    skip_update_flag=skip_fp8_weight_update,
-                    fsdp_group=fsdp_group,
-                    create_transpose_cache=keep_fp8_weight_transpose_cache,
-                )
-
-        # Cast bias to expected dtype
-        bias_dtype = activation_dtype
-        if fp8 and activation_dtype == torch.float32:
-            bias_dtype = torch.bfloat16
-        bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
-
-        # Configure output quantizer
-        if output_quantizer is not None:
-            output_quantizer.set_usage(rowwise=True, columnwise=False)
-
-        # Calibrate quantizers if needed
-        if not fp8 and fp8_calibration:
-            if input_quantizer is not None:
-                input_quantizer.calibrate(inputmat_total)
-            if weight_quantizer is not None:
-                weight_quantizer.calibrate(weight)
-
-        ub_obj = None
-        ub_type = None
-        rs_out = None
-        out_dtype = activation_dtype
-        if ub_overlap_rs_fprop:
-            ub_obj = get_ub(ub_name + "_fprop")
-            ub_type = tex.CommOverlapType.RS
-            out_shape = [reduce(multiply_op, inp_shape[:-1]) // tp_world_size, out_features]
-            rs_out = torch.empty(out_shape, dtype=activation_dtype, device=inputmat_total.device)
-
-        elif ub_overlap_ag_fprop:
-            ub_obj = get_ub(ub_name + "_fprop")
-            ub_type = tex.CommOverlapType.AG
-            if fp8:
-                assert ub_obj.is_fp8_ubuf(), "AG overlap with FP8 GEMM inputs requires FP8 buffer."
-            ub_obj.copy_into_buffer(inputmat_total, input_quantizer, local_chunk=True)
-            inputmat_total = ub_obj.get_buffer(input_quantizer)
-
-        nvtx_range_push(f"{nvtx_label}.gemm")
-        fprop_gemm_use_split_accumulator = _2X_ACC_FPROP
-        if fp8:
-            recipe = FP8GlobalStateManager.get_fp8_recipe()
-            if hasattr(recipe, "fp8_gemm_fprop"):
-                fprop_gemm_use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator
-
-        out, *_, rs_out = general_gemm(
-            weightmat,
-            inputmat_total,
-            get_workspace(),
-            quantization_params=output_quantizer,
-            out_dtype=out_dtype,
-            bias=bias,
-            use_split_accumulator=fprop_gemm_use_split_accumulator,
-            ub=ub_obj,
-            ub_type=ub_type,
-            extra_output=rs_out,
-        )
-        nvtx_range_pop(f"{nvtx_label}.gemm")
-
-        if is_grad_enabled:
-            saved_inputmat = None
-
-            ctx.backward_input_needs_gather = (
-                weight.requires_grad and parallel_mode == "column" and sequence_parallel
-            )
-
-            if backward_needs_input:
-                if own_quantized_input and isinstance(inputmat, QuantizedTensor):
-                    # For sequence parallel in vanilla FP8, rowwise data is
-                    # to gather the input. For MXFP8, columnwise only data
-                    # can be allgathered.
-                    if isinstance(inputmat, MXFP8TensorBase) or not ctx.backward_input_needs_gather:
-                        inputmat.update_usage(rowwise_usage=False)
-                saved_inputmat = inputmat
-
-            # Weight with column-wise usage is needed for dgrad GEMM while keeping fp8 weight transpose cache.
-            if inp.requires_grad and keep_fp8_weight_transpose_cache:
-                if isinstance(weightmat, QuantizedTensor):
-                    weightmat.update_usage(columnwise_usage=True)
-
-            if cpu_offloading:
-                set_offloading_param(weight, "weight_offloading", True)
-                set_offloading_param(weightmat, "weight_offloading", True)
-                if saved_inputmat is not None:
-                    set_offloading_param(saved_inputmat, "activation_offloading", True)
-
-            # Scatter intermediate/activation tensors saved for the backward pass
-            # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights
-            nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
-            ctx.fsdp_group = fsdp_group
-            ctx.fsdp_shapes = _fsdp_scatter_tensors(
-                fsdp_group,
-                saved_inputmat,
-                weightmat if fp8 and not isinstance(weight, QuantizedTensor) else None,
-            )
-            nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
-
-            if cpu_offloading:
-                ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad")
-
-                if ctx.grad_added_to_main_grad:
-                    # If you are passing torch.nn.Parameter through the Torch hooks, you will
-                    # get back torch.Tensor. Torch rips off the Parameter wrapper.
-                    # You need to preserve the weight object to have all the attributes user
-                    # sets for the weights. Because of this, it is not recommended to offload
-                    # weights if weights are externally touched outside this module
-                    ctx.weight_object = weight
-
-            # TODO(ksivamani): Check memory usage
-            tensors_to_save, tensor_objects = prepare_for_saving(
-                saved_inputmat,
-                weightmat,
-                weight,
-                bias,
-            )
-            ctx.save_for_backward(*tensors_to_save)
-            ctx.tensor_objects = tensor_objects
-
-            ctx.activation_dtype = activation_dtype
-            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
-            ctx.fp8 = fp8
-            ctx.input_quantizer = input_quantizer
-            ctx.grad_output_quantizer = grad_output_quantizer
-            ctx.grad_input_quantizer = grad_input_quantizer
-            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
-            if fuse_wgrad_accumulation and weight.requires_grad:
-                ctx.main_grad = weight.main_grad
-
-            ctx.cpu_offloading = cpu_offloading
-            ctx.is_first_microbatch = is_first_microbatch
-            ctx.use_bias = bias is not None
-            ctx.sequence_parallel = sequence_parallel
-            ctx.tensor_parallel = tensor_parallel
-            ctx.inp_shape = inp_shape
-            ctx.parallel_mode = parallel_mode
-            ctx.tp_group = tp_group
-            ctx.ub_overlap_ag = ub_overlap_ag_dgrad
-            ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad
-            ctx.ub_bulk_dgrad = ub_bulk_dgrad
-            ctx.ub_bulk_wgrad = ub_bulk_wgrad
-            ctx.ub_name = ub_name
-            ctx.tp_size = tp_size
-            ctx.requires_dgrad = inp.requires_grad
-            ctx.requires_wgrad = weight.requires_grad
-            ctx.reduce_and_update_bwd_fp8_tensors = False
-            ctx.owns_input = saved_inputmat is not inp
-            ctx.keep_fp8_weight_transpose_cache = keep_fp8_weight_transpose_cache
-            if ctx.fp8 and requires_grad(inp, weight, bias):
-                _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE
-                ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
-                if in_fp8_activation_recompute_phase():
-                    FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module
-
-        # Row Parallel Linear
-        if ub_overlap_rs_fprop:
-            out = rs_out
-        elif parallel_mode == "row":
-            nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
-            if sequence_parallel:
-                out, _ = reduce_scatter_along_first_dim(out, tp_group)
-            elif tensor_parallel:
-                out, _ = allreduce(out, tp_group)
-            nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
-
-        out = out.view(-1, *inp_shape[1:-1], out_features)
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
-        # pylint: disable=missing-function-docstring
-
-        # NVTX label for profiling
-        nvtx_label = "transformer_engine._Linear.backward"
-        if ctx.ub_name is not None:
-            nvtx_label = f"{nvtx_label}.{ctx.ub_name}"
-
-        with torch.cuda.nvtx.range("_Linear_backward"):
-            if (
-                ctx.fp8
-                and any(
-                    [
-                        ctx.ub_overlap_ag,
-                        ctx.ub_overlap_rs_dgrad,
-                        ctx.ub_bulk_dgrad,
-                        ctx.ub_bulk_wgrad,
-                    ]
-                )
-                and (ctx.fp8_recipe is not None)
-            ):
-                if not ctx.fp8_recipe.float8_per_tensor_scaling():
-                    raise NotImplementedError(
-                        "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor"
-                        " current scaling"
-                    )
-
-            saved_tensors = ctx.saved_tensors
-            inputmat, weight_fp8, weight, bias = (  # pylint: disable=unbalanced-tuple-unpacking
-                restore_from_saved(ctx.tensor_objects, saved_tensors)
-            )
-            # Delete the references to tensor objects once they've been consumed
-            # by the `restore_from_saved` method to construct back the actual tensors.
-            ctx.tensor_objects = None
-
-            # Since main_grad can be modified inplace, it should not be a part of saved_tensors
-            main_grad = (
-                ctx.main_grad
-                if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad
-                else None
-            )
-
-            if ctx.cpu_offloading:
-                if ctx.grad_added_to_main_grad:
-                    weight = ctx.weight_object
-                if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation:
-                    weight.main_grad = main_grad
-
-            # Gather intermediate/activation tensors if needed
-            # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
-            #       shards/unshards the base weights so we don't do it ourselves
-            nvtx_range_push(f"{nvtx_label}.fsdp_gather")
-            _fsdp_gather_tensors(
-                ctx.fsdp_group,
-                ctx.fsdp_shapes,
-                inputmat,
-                weight_fp8,
-            )
-            nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
-
-            ctx.ub_obj_gradout = None
-            ub_obj_dgrad = None
-            ub_obj_wgrad = None
-            ub_type_dgrad = None
-            ub_type_wgrad = None
-            dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]]
-            rs_out = None
-            dgrad_bulk = None
-            if ctx.ub_overlap_ag:
-                # Overlap grad_output all-gather with dgrad compute
-                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
-                ub_obj_dgrad = ctx.ub_obj_gradout
-                ub_type_dgrad = tex.CommOverlapType.AG
-
-            elif ctx.ub_overlap_rs_dgrad:
-                # Overlap dgrad reduce-scatter with dgrad compute
-                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
-                ub_obj_dgrad = ctx.ub_obj_gradout
-                ub_type_dgrad = tex.CommOverlapType.RS
-                rs_out = torch.empty(dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device)
-
-            else:
-                if ctx.ub_bulk_dgrad:
-                    # Overlap inputmat all-gather with dgrad compute
-                    # NOTE: Copying into communication buffer will always prefer rowwise data,
-                    #       and will copy columnwise data if rowwise does not exist. In that case,
-                    #       the all-gather will apply to the leading dimension of the transpose,
-                    #       which then needs to be interleaved correctly before WGRAD.
-                    ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad")
-                    ub_obj_dgrad = ctx.ub_obj_gradout
-                    ub_type_dgrad = tex.CommOverlapType.AG
-                    ub_obj_dgrad.copy_into_buffer(inputmat, ctx.input_quantizer, local_chunk=True)
-
-                if ctx.ub_bulk_wgrad:
-                    # Overlap dgrad reduce-scatter with wgrad compute
-                    ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad")
-                    ub_type_wgrad = tex.CommOverlapType.RS
-                    ub_obj_wgrad.set_buffer_params(ctx.grad_input_quantizer)
-                    dgrad_bulk = ub_obj_wgrad.get_buffer(ctx.grad_input_quantizer)
-
-            # Prepare grad output tensor
-            # Note: Cast to expected dtype and perform tensor-parallel communication
-            if ctx.grad_output_quantizer is not None:
-                # Reduce duplicated transpose, which is performed in grad_output.update_usage
-                if ctx.ub_overlap_ag and ctx.fp8_recipe.float8_per_tensor_scaling():
-                    ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False)
-                else:
-                    ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=True)
-            nvtx_range_push(f"{nvtx_label}.grad_output_preprocess")
-            (
-                grad_output,
-                grad_bias,
-            ) = TransformerEngineBaseModule.grad_output_preprocess(
-                ctx,
-                grad_output,
-                ctx.parallel_mode == "row",
-                ctx.grad_output_quantizer,
-            )
-            nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
-
-            # Prepare input tensor
-            # Note: Perform tensor-parallel communication if needed
-            inputmat_total = None
-            inputmat_total_work = None
-            if ctx.backward_input_needs_gather and not ctx.ub_bulk_dgrad:
-                quantizer = None
-                if ctx.fp8:
-                    quantizer = ctx.input_quantizer
-                    quantizer.set_usage(rowwise=True, columnwise=True)
-                nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
-                inputmat_total, inputmat_total_work = gather_along_first_dim(
-                    inputmat,
-                    ctx.tp_group,
-                    async_op=True,
-                    quantizer=quantizer,
-                )
-                nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
-            else:
-                inputmat_total = inputmat
-
-            # Check whether to output wgrad GEMM directly into main grad
-            if ctx.is_first_microbatch is not None:
-                accumulate_wgrad_into_param_main_grad = (
-                    ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
-                )
-            else:
-                accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation
-
-            # Compute grad input tensor
-            dgrad = None
-            dgrad_work = None
-            if ctx.requires_dgrad:
-
-                # Update quantizer
-                if ctx.grad_input_quantizer is not None:
-                    ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
-
-                if ctx.fp8 and not ctx.keep_fp8_weight_transpose_cache:
-                    create_fp8_weight_transpose_cache(weight_fp8)
-
-                # dgrad GEMM
-                nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
-                dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD
-                if ctx.fp8:
-                    recipe = ctx.fp8_recipe
-                    if hasattr(recipe, "fp8_gemm_dgrad"):
-                        dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
-
-                dgrad, *_, rs_out = general_gemm(
-                    weight_fp8,
-                    grad_output,
-                    get_workspace(),
-                    layout="NN",
-                    grad=True,
-                    quantization_params=ctx.grad_input_quantizer,
-                    out=dgrad_bulk,
-                    out_dtype=ctx.activation_dtype,
-                    use_split_accumulator=dgrad_gemm_use_split_accumulator,
-                    ub=ub_obj_dgrad,
-                    ub_type=ub_type_dgrad,
-                    extra_output=rs_out,
-                    bulk_overlap=ctx.ub_bulk_dgrad,
-                )
-                nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
-
-                if ctx.fp8 and not ctx.keep_fp8_weight_transpose_cache:
-                    clear_fp8_weight_transpose_cache(weight_fp8)
-
-                # Launch tensor-parallel communication
-                if ctx.ub_overlap_rs_dgrad:
-                    dgrad = rs_out
-                elif ctx.parallel_mode == "column" and not ctx.ub_bulk_wgrad:
-                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
-                    if ctx.sequence_parallel:
-                        dgrad, dgrad_work = reduce_scatter_along_first_dim(
-                            dgrad,
-                            ctx.tp_group,
-                            async_op=True,
-                        )
-                    else:
-                        dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
-                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
-
-            # Compute grad weight tensor
-            wgrad = None
-            if ctx.requires_wgrad:
-                if ctx.ub_bulk_dgrad:
-                    inputmat_total = ub_obj_dgrad.get_buffer(ctx.input_quantizer)
-                    if ctx.fp8:
-                        if inputmat._data is None:
-                            # All-gather executed on columnwise data and result is in rowwise data,
-                            # so we need to fix the interleaving before WGRAD.
-                            inputmat_total = _fix_gathered_fp8_transpose(inputmat_total, ctx.tp_size)
-                        elif not non_tn_fp8_gemm_supported():
-                            # FP8 GEMM on Hopper only supports TN layout so the gathered input must
-                            # have a valid transpose.
-                            inputmat_total._create_transpose()
-
-                else:
-                    if inputmat_total_work is not None:
-                        # Synchronize tensor-parallel communication
-                        inputmat_total_work.wait()
-                        inputmat_total_work = None
-
-                if isinstance(grad_output, QuantizedTensor):
-                    # This is a no-op if platform supports non-TN FP8 GEMM or the transpose
-                    # already exists.
-                    grad_output.update_usage(rowwise_usage=True, columnwise_usage=True)
-
-                if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf():
-                    rs_out = torch.empty(dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device)
-
-                # wgrad GEMM
-                # Note: Fuse with bgrad computation if needed
-                def pre_process(_grad_output_, _input_, async_op=True):
-                    return _grad_output_, _input_, None
-
-                def process_wgrad(main_grad, grad_output, inputmat_total, handle=None):
-                    nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
-
-                    wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD
-                    if ctx.fp8:
-                        recipe = ctx.fp8_recipe
-                        if hasattr(recipe, "fp8_gemm_wgrad"):
-                            wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator
-                    # print(f"debug acc {accumulate_wgrad_into_param_main_grad}")
-                    wgrad, grad_bias_, _, _ = general_gemm(
-                        inputmat_total,
-                        grad_output,
-                        get_workspace(),
-                        layout="NT",
-                        grad=True,
-                        out_dtype=(main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype),
-                        bias=None,
-                        out=main_grad if ctx.fuse_wgrad_accumulation else None,
-                        use_split_accumulator=wgrad_gemm_use_split_accumulator,
-                        accumulate=accumulate_wgrad_into_param_main_grad,
-                        ub=ub_obj_wgrad,
-                        ub_type=ub_type_wgrad,
-                        extra_output=None,
-                        bulk_overlap=False,
-                    )
-
-                    nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
-
-                    # Deallocate input tensor
-                    if ctx.owns_input:
-                        clear_tensor_data(inputmat_total)
-                    # Handle custom DDP from mcore.
-                    if (
-                        ctx.fuse_wgrad_accumulation
-                        and weight is not None
-                        and hasattr(weight, "grad_added_to_main_grad")
-                    ):
-                        weight.grad_added_to_main_grad = True
-                        if getattr(weight, "zero_out_wgrad", False):
-                            wgrad = torch.zeros(
-                                weight.main_grad.shape,
-                                dtype=weight.dtype,
-                                device=torch.cuda.current_device(),
-                                requires_grad=False,
-                            )
-                        else:
-                            wgrad = torch.empty(
-                                weight.main_grad.shape,
-                                dtype=weight.dtype,
-                                device=torch.cuda.current_device(),
-                                requires_grad=False,
-                            )
-                    elif ctx.fuse_wgrad_accumulation:
-                        pass
-
-                insert_wgrad_func_into_cache(
-                    main_grad,
-                    functools.partial(pre_process, grad_output, inputmat_total),
-                    functools.partial(process_wgrad, main_grad),
-                )
-
-            # Synchronize tensor parallel communication
-            if inputmat_total_work is not None:
-                assert False
-                inputmat_total_work.wait()
-                inputmat_total_work = None
-            if dgrad_work is not None:
-                assert False
-                dgrad_work.wait()
-                dgrad_work = None
-
-        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
-            nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
-            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
-            nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
-
-        # Scatter fp8 weight buffers
-        if ctx.fp8 and not isinstance(weight, QuantizedTensor):
-            _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
-
-        wgrad = None
-        return (
-            wgrad,
-            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
-            grad_bias,
-            None,  # is_first_microbatch
-            None,  # fp8
-            None,  # fp8_calibration
-            None,  # input_quantizer
-            None,  # weight_quantizer
-            None,  # output_quantizer
-            None,  # grad_output_quantizer
-            None,  # grad_input_quantizer
-            None,  # fuse_wgrad_accumulation
-            None,  # cpu_offloading
-            None,  # tp_group
-            None,  # tp_size
-            None,  # sequence_parallel
-            None,  # tensor_parallel
-            None,  # activation_dtype
-            None,  # parallel_mode
-            None,  # is_grad_enabled
-            None,  # ub_overlap_rs_fprop
-            None,  # ub_overlap_ag_dgrad
-            None,  # ub_overlap_ag_fprop
-            None,  # ub_overlap_rs_dgrad
-            None,  # ub_bulk_dgrad
-            None,  # ub_bulk_wgrad
-            None,  # ub_name
-            None,  # fp8_output
-            None,  # fsdp_group
-            None,  # module
-            None,  # skip_fp8_weight_update
-            None,  # keep_fp8_weight_transpose_cache
-        )
diff --git a/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py b/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py
deleted file mode 100644
index d7b77ce31..000000000
--- a/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py
+++ /dev/null
@@ -1,372 +0,0 @@
-# This file was modified for portability to AMDGPU
-# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
-# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
-# See LICENSE for license information.
-
-import functools
-from typing import List, Tuple, Union
-
-import torch
-import transformer_engine_torch as tex
-from transformer_engine.pytorch.constants import TE_DType
-from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm
-from transformer_engine.pytorch.distributed import (
-    in_fp8_activation_recompute_phase,
-    is_fp8_activation_recompute_enabled,
-)
-from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
-from transformer_engine.pytorch.graph import is_graph_capturing
-from transformer_engine.pytorch.module.base import (
-    _2X_ACC_DGRAD,
-    _2X_ACC_FPROP,
-    _2X_ACC_WGRAD,
-    get_multi_stream_cublas_workspace,
-)
-from transformer_engine.pytorch.utils import (
-    assert_dim_for_fp8_exec,
-    cast_if_needed,
-    clear_tensor_data,
-    requires_grad,
-)
-
-from primus.backends.megatron.core.pipeline_parallel.wgrad_adapter import (
-    insert_wgrad_func_into_cache,
-)
-
-try:
-    from transformer_engine.pytorch.tensor.quantized_tensor import (
-        QuantizedTensor,
-        Quantizer,
-        prepare_for_saving,
-        restore_from_saved,
-    )
-except ModuleNotFoundError:
-    from transformer_engine.pytorch.quantized_tensor import (
-        QuantizedTensor,
-        Quantizer,
-        prepare_for_saving,
-        restore_from_saved,
-    )
-
-
-class _GroupedLinearWithWGradSplit(torch.autograd.Function):
-    """GroupedLinear semi-top level module
-    Calls custom cuda extensions.
-    """
-
-    @staticmethod
-    def forward(
-        ctx,
-        inp: torch.Tensor,
-        m_splits: List[int],
-        use_bias: bool,
-        is_first_microbatch: Union[bool, None],
-        fp8: bool,
-        fp8_calibration: bool,
-        input_quantizers: List[Quantizer],
-        weight_quantizers: List[Quantizer],
-        output_quantizers: List[Quantizer],
-        grad_output_quantizers: List[Quantizer],
-        fuse_wgrad_accumulation: bool,
-        cpu_offloading: bool,
-        sequence_parallel: bool,
-        activation_dtype: torch.dtype,
-        is_grad_enabled: bool,
-        module,
-        skip_fp8_weight_update,
-        *weights_and_biases,
-    ) -> torch.Tensor:
-        assert fuse_wgrad_accumulation, "fuse_wgrad_accumulation need to be true"
-        assert not use_bias, "zero bubble not support bias yet"
-        # pylint: disable=missing-function-docstring
-        num_gemms = len(m_splits)
-        weights = weights_and_biases[:num_gemms]
-        biases = weights_and_biases[num_gemms:]
-        device = inp.device
-
-        # TODO Support MXFP8  # pylint: disable=fixme
-        if fp8 and FP8GlobalStateManager.get_fp8_recipe().mxfp8():
-            raise NotImplementedError("GroupedLinear does not yet support MXFP8")
-        # TODO Support Float8 Current Scaling  # pylint: disable=fixme
-        if fp8 and FP8GlobalStateManager.get_fp8_recipe().float8_current_scaling():
-            raise NotImplementedError("GroupedLinear does not yet support Float8 Current Scaling")
-
-        # Make sure input dimensions are compatible
-        in_features = weights[0].shape[-1]
-        assert inp.shape[-1] == in_features, "GEMM not possible"
-        inputmats = torch.split(inp.view(-1, in_features), m_splits)
-        if fp8:
-            assert_dim_for_fp8_exec(*inputmats, *weights)
-
-        # Cast input to expected dtype
-        inputmats_no_fp8 = [cast_if_needed(mat, activation_dtype) for mat in inputmats]
-        inputmats = []
-
-        weight_requires_grad = weights[0].requires_grad
-
-        if input_quantizers[0] is not None:
-            for input_quantizer in input_quantizers:
-                input_quantizer.set_usage(
-                    rowwise=True,
-                    columnwise=(is_grad_enabled and weight_requires_grad),
-                )
-            columnwise_usage = is_grad_enabled and inp.requires_grad
-            if not columnwise_usage:
-                columnwise_usage = (
-                    is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase()
-                )
-            if weight_quantizers[0] is not None:
-                for weight_quantizer in weight_quantizers:
-                    weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
-        if output_quantizers[0] is not None:
-            for output_quantizer in output_quantizers:
-                output_quantizer.set_usage(rowwise=True, columnwise=False)
-
-        if fp8:
-            inputmats = tex.fused_multi_quantize(
-                inputmats_no_fp8, None, input_quantizers, TE_DType[activation_dtype]
-            )
-            weights_fp8 = []
-            bias_dtype = torch.bfloat16 if activation_dtype == torch.float32 else activation_dtype
-            if not isinstance(weights[0], QuantizedTensor):
-                # FP8 cast to workspace buffer
-                update_workspace = is_first_microbatch is None or is_first_microbatch
-                for i in range(num_gemms):
-                    weight_fp8 = module.get_weight_workspace(
-                        tensor=weights[i],
-                        quantizer=weight_quantizers[i],
-                        cache_name=(None if is_first_microbatch is None else f"weight{i}"),
-                        update_workspace=update_workspace,
-                        skip_update_flag=skip_fp8_weight_update,
-                    )
-                    weights_fp8.append(weight_fp8)
-            else:
-                weights_fp8 = weights
-
-        else:
-            inputmats = inputmats_no_fp8
-            bias_dtype = activation_dtype
-            weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]
-
-        biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases
-
-        out = torch.empty(
-            [sum(m_splits), weights_fp8[0].size(0)],
-            dtype=activation_dtype,
-            device=device,
-        )
-
-        _ = general_grouped_gemm(
-            weights_fp8,
-            inputmats,
-            [out],
-            activation_dtype,
-            get_multi_stream_cublas_workspace(),
-            single_output=True,
-            m_splits=m_splits,
-            bias=biases,
-            use_bias=use_bias,
-            use_split_accumulator=_2X_ACC_FPROP,
-        )
-
-        if fp8_calibration:
-            for i in range(num_gemms):
-                # amax of input
-                for i in range(num_gemms):
-                    input_quantizers[i].calibrate(inputmats[i])
-                for i in range(num_gemms):
-                    weight_quantizers[i].calibrate(weights[i])
-
-        if is_grad_enabled:
-
-            ctx.weights_shape_1 = weights[0].shape[1]
-
-            tensors_to_save, tensor_objects = prepare_for_saving(*inputmats, *weights_fp8, *biases)
-            ctx.save_for_backward(*tensors_to_save)
-            ctx.tensor_objects = tensor_objects
-
-            ctx.weights_requires_grad = weights[0].requires_grad
-            if fuse_wgrad_accumulation and ctx.weights_requires_grad:
-                ctx.main_grads = [weights[i].main_grad for i in range(num_gemms)]
-            else:
-                ctx.main_grads = [None] * num_gemms
-            ctx.device = device
-            ctx.grad_output_quantizers = grad_output_quantizers
-            ctx.m_splits = m_splits
-            ctx.num_gemms = num_gemms
-            ctx.activation_dtype = activation_dtype
-            ctx.fp8 = fp8
-            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
-            ctx.cpu_offloading = cpu_offloading
-            ctx.is_first_microbatch = is_first_microbatch
-            ctx.use_bias = use_bias
-            ctx.sequence_parallel = sequence_parallel
-            ctx.inp_shape = inp.shape
-            ctx.requires_dgrad = inp.requires_grad
-            ctx.reduce_and_update_bwd_fp8_tensors = False
-            if ctx.fp8 and requires_grad(inp, weights[0], biases[0]):
-                ctx.reduce_and_update_bwd_fp8_tensors = (
-                    ctx.reduce_and_update_bwd_fp8_tensors or FP8GlobalStateManager.is_first_fp8_module()
-                )
-
-        # [*, in_features] -> [*, out_features] except first dimension changes for SP
-        return out.view(-1, *inp.shape[1:-1], out.shape[-1])
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
-        # pylint: disable=missing-function-docstring
-        with torch.cuda.nvtx.range("_GroupedLinear_backward"):
-            saved_tensors = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors)
-            N = ctx.num_gemms
-            inputmats = saved_tensors[:N]
-            weights = saved_tensors[N : 2 * N]
-            biases = saved_tensors[2 * N : 3 * N]
-            main_grads = ctx.main_grads
-
-            # preprocess grad_output
-
-            grad_output = grad_output.contiguous()
-            grad_output_mats = torch.split(grad_output.view(-1, grad_output.shape[-1]), ctx.m_splits)
-            grad_output = [None] * ctx.num_gemms
-            grad_biases = [None] * ctx.num_gemms
-            if ctx.fp8:
-                if ctx.use_bias:
-                    for i in range(ctx.num_gemms):
-                        grad_biases[i], grad_output[i] = tex.bgrad_quantize(
-                            grad_output_mats[i], ctx.grad_output_quantizers[i]
-                        )
-                else:
-                    grad_output = tex.fused_multi_quantize(
-                        grad_output_mats,
-                        None,
-                        ctx.grad_output_quantizers,
-                        TE_DType[ctx.activation_dtype],
-                    )
-            else:
-                grad_output = grad_output_mats
-
-            if ctx.is_first_microbatch is not None:
-                accumulate_wgrad_into_param_main_grad = (
-                    ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
-                )
-            else:
-                accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation
-
-            if ctx.requires_dgrad:
-                dgrad = torch.empty(
-                    (sum(ctx.m_splits), ctx.weights_shape_1),
-                    dtype=ctx.activation_dtype,
-                    device=ctx.device,
-                )
-
-                general_grouped_gemm(
-                    weights,
-                    grad_output,
-                    [dgrad],
-                    ctx.activation_dtype,
-                    get_multi_stream_cublas_workspace(),
-                    single_output=True,
-                    layout="NN",
-                    m_splits=ctx.m_splits,
-                    grad=True,
-                    use_split_accumulator=_2X_ACC_DGRAD,
-                )
-
-            if ctx.weights_requires_grad:
-                if ctx.fuse_wgrad_accumulation:
-                    wgrad_list = main_grads
-                else:
-                    wgrad_list = [
-                        torch.empty(w.size(), dtype=ctx.activation_dtype, device=ctx.device) for w in weights
-                    ]
-
-                def handle_custom_ddp_from_mcore(w, wgrad):
-                    if ctx.weights_requires_grad:
-                        if ctx.fuse_wgrad_accumulation and hasattr(w, "grad_added_to_main_grad"):
-                            w.grad_added_to_main_grad = True
-                            if getattr(w, "zero_out_wgrad", False):
-                                wgrad = torch.zeros(
-                                    w.main_grad.shape,
-                                    dtype=w.dtype,
-                                    device=torch.cuda.current_device(),
-                                    requires_grad=False,
-                                )
-                            else:
-                                wgrad = torch.empty(
-                                    w.main_grad.shape,
-                                    dtype=w.dtype,
-                                    device=torch.cuda.current_device(),
-                                    requires_grad=False,
-                                )
-                        elif ctx.fuse_wgrad_accumulation:
-                            wgrad = None
-                    else:
-                        wgrad = None
-                    return wgrad
-
-                def pre_process(_grad_output_, _input_, async_op=True):
-                    return _grad_output_, _input_, None
-
-                kargs_dict = {
-                    "out_dtype": ctx.activation_dtype,
-                    "workspaces": get_multi_stream_cublas_workspace(),
-                    "layout": "NT",
-                    "grad": True,
-                    "m_splits": ctx.m_splits,
-                    "use_bias": ctx.use_bias if grad_biases[0] is None else None,
-                    "bias": biases,
-                    "use_split_accumulator": _2X_ACC_WGRAD,
-                    "accumulate": accumulate_wgrad_into_param_main_grad,
-                }
-
-                def process_wgrad(wgrad_list, kargs_dict, grad_output, inputmats, handle=None):
-                    _, grad_biases_, _ = general_grouped_gemm(
-                        inputmats, grad_output, wgrad_list, **kargs_dict
-                    )
-                    for i in range(ctx.num_gemms):
-                        if grad_biases[i] is None:
-                            grad_biases[i] = grad_biases_[i]
-                    del grad_biases_
-                    clear_tensor_data(*inputmats)
-                    wgrad_list = [
-                        handle_custom_ddp_from_mcore(w, wgrad) for w, wgrad in zip(weights, wgrad_list)
-                    ]
-
-                insert_wgrad_func_into_cache(
-                    wgrad_list,
-                    functools.partial(pre_process, grad_output, inputmats),
-                    functools.partial(process_wgrad, wgrad_list, kargs_dict),
-                )
-
-                wgrad_list_return = [None] * ctx.num_gemms
-            else:
-                wgrad_list_return = [None] * ctx.num_gemms
-
-            if not ctx.use_bias:
-                grad_biases = [None] * ctx.num_gemms
-
-        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
-            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
-
-        return (
-            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,  # is_grad_enabled
-            None,  # is_grad_enabled
-            *wgrad_list_return,
-            *grad_biases,
-        )
diff --git a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py
index 9b762b721..4bb006e5a 100644
--- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py
+++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py
@@ -5,6 +5,7 @@
 ###############################################################################
 
 import warnings
+from types import SimpleNamespace
 from typing import Optional
 
 from megatron.core.extensions.transformer_engine import (
@@ -21,6 +22,7 @@
 from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
 from megatron.core.models.backends import BackendSpecProvider
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.dot_product_attention import DotProductAttention
 from megatron.core.transformer.mlp import MLPSubmodules
 from megatron.core.transformer.moe.experts import (
     SequentialMLP,
@@ -41,6 +43,7 @@
         PrimusTurboLinear,
         PrimusTurboRowParallelGroupedLinear,
         PrimusTurboRowParallelLinear,
+        _make_primus_turbo_norm_te_column_parallel_linear,
     )
 except (ImportError, ModuleNotFoundError):
     PrimusTurboAttention = None
@@ -50,6 +53,7 @@
     PrimusTurboLinear = None
     PrimusTurboRowParallelGroupedLinear = None
     PrimusTurboRowParallelLinear = None
+    _make_primus_turbo_norm_te_column_parallel_linear = None
 
 _LEGACY_GROUPED_MLP_CLS = None
 
@@ -101,11 +105,27 @@ def __init__(
     return PrimusLegacyGroupedMLP
 
 
+def _build_default_primus_args() -> SimpleNamespace:
+    """Fallback args for environments without initialized Primus globals."""
+    return SimpleNamespace(
+        enable_primus_turbo=False,
+        use_turbo_gemm=False,
+        use_turbo_norm_te_linear=False,
+        use_turbo_attention=False,
+        use_turbo_grouped_gemm=False,
+        moe_use_legacy_grouped_gemm=False,
+    )
+
+
 class PrimusTurboSpecProvider(BackendSpecProvider):
     """A protocol for providing the submodules used in Spec building."""
 
-    def __init__(self):
-        self.cfg = get_primus_args()
+    def __init__(self, fallback_to_eager_attn: bool = False):
+        try:
+            self.cfg = get_primus_args()
+        except AssertionError:
+            self.cfg = _build_default_primus_args()
+        self.fallback_to_eager_attn = fallback_to_eager_attn
 
     def linear(self) -> type:
         """Which linear module TE backend uses"""
@@ -131,19 +151,65 @@ def row_parallel_linear(self) -> type:
             else TERowParallelLinear
         )
 
+    def column_parallel_linear_with_gather_output(self) -> type:
+        """Non-TE column-parallel linear that supports ``gather_output=True``.
+
+        TE / Turbo column-parallel wrappers explicitly raise
+        ``ValueError("Transformer Engine linear layers do not support
+        gather_output = True")`` (see
+        ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:747``
+        and ``:972``).  Callers that need a column-parallel layer
+        whose output dim is gathered back to full width across TP
+        ranks (so downstream math stays TP-agnostic) must use the
+        upstream Megatron-native :class:`ColumnParallelLinear`.
+
+        Plan-3 P21: V4's ``linear_q_up_proj`` is the canonical
+        consumer — it shards ``q_lora_rank -> n_heads * head_dim``
+        across TP and gathers the heads at forward time.
+        """
+        return ColumnParallelLinear
+
+    def row_parallel_linear_with_scatter_input(self) -> type:
+        """Non-TE row-parallel linear that supports ``input_is_parallel=False``.
+
+        TE / Turbo row-parallel wrappers explicitly raise
+        ``ValueError("Transformer Engine linear layers do not support
+        input_is_parallel = False")`` (see
+        ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:1081``).
+        Callers that hand the layer a full-width (non-sharded) input
+        and want it scattered internally + the output all-reduced
+        must use the upstream Megatron-native :class:`RowParallelLinear`.
+
+        Plan-3 P21: V4's grouped-O ``linear_o_b`` is the canonical
+        consumer — after the inner ``[..., n_per_group, o_lora_rank]
+        -> [..., o_groups * o_lora_rank]`` reshape, the input is
+        full-width across TP and the row-parallel layer's
+        weight-sharding + reduce is what produces the correct sum.
+        """
+        return RowParallelLinear
+
     def fuse_layernorm_and_linear(self) -> bool:
         """TE backend chooses a single module for layernorm and linear"""
         return True
 
     def column_parallel_layer_norm_linear(self) -> Optional[type]:
-        """Which module for sequential layernorm and linear"""
-        return (
-            _require_primus_turbo(
+        """Which module for sequential layernorm and linear.
+
+        Three choices, because the norm and the GEMM want different backends on
+        these shapes: Turbo's rmsnorm beats TE's, while TE's hipBLASLt GEMM beats
+        Turbo's. ``use_turbo_norm_te_linear`` takes only the norm.
+        """
+        if self.cfg.use_turbo_gemm:
+            return _require_primus_turbo(
                 PrimusTurboLayerNormColumnParallelLinear, "layernorm column parallel linear"
             )
-            if self.cfg.use_turbo_gemm
-            else TELayerNormColumnParallelLinear
-        )
+        if getattr(self.cfg, "use_turbo_norm_te_linear", False):
+            factory = _require_primus_turbo(
+                _make_primus_turbo_norm_te_column_parallel_linear,
+                "turbo-norm layernorm column parallel linear",
+            )
+            return factory()
+        return TELayerNormColumnParallelLinear
 
     def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type:
         """Which module to use for layer norm"""
@@ -156,6 +222,8 @@ def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type:
 
     def core_attention(self) -> type:
         """Which module to use for attention"""
+        if self.fallback_to_eager_attn:
+            return DotProductAttention
         return (
             _require_primus_turbo(PrimusTurboAttention, "attention")
             if self.cfg.use_turbo_attention
@@ -170,8 +238,7 @@ def grouped_mlp_modules(
             # Megatron callers only pass ``moe_use_grouped_gemm`` here, so when Primus
             # args do not expose the legacy switch we must match upstream TESpecProvider
             # and prefer TEGroupedMLP by default.
-            # let it raise an error if cfg does not have moe_use_legacy_grouped_gemm
-            moe_use_legacy_grouped_gemm = self.cfg.moe_use_legacy_grouped_gemm
+            moe_use_legacy_grouped_gemm = getattr(self.cfg, "moe_use_legacy_grouped_gemm", False)
 
         use_turbo_grouped_gemm = self.cfg.use_turbo_grouped_gemm
         assert not (
@@ -220,3 +287,157 @@ def grouped_mlp_modules(
     def activation_func(self) -> type:
         """Which module to use for activation function"""
         return TEActivationOp
+
+
+class DeepSeekV4SpecProvider(PrimusTurboSpecProvider):
+    """DeepSeek-V4 provider rooted on PrimusTurboSpecProvider."""
+
+    def __init__(self, config=None):
+        super().__init__()
+        self.config = config
+
+    def v4_norm_module(self):
+        """Norm module used by V4 specs (block / layer / final)."""
+        return self.layer_norm(rms_norm=True)
+
+    def v4_q_layernorm(self) -> type:
+        """Norm module for V4's `q_norm` (RMSNorm on `q_lora_rank`).
+
+        Same as MLA's `q_layernorm`; we use the for_qk path so the TE
+        version selection picks Apex / FusedLayerNorm on older TE.
+        """
+        return self.layer_norm(rms_norm=True, for_qk=True)
+
+    def v4_kv_layernorm(self) -> type:
+        """Norm module for V4's `kv_norm` (RMSNorm on `head_dim`).
+
+        Single-latent KV: V4 normalizes the `wkv` output (one shared head)
+        BEFORE broadcasting to all query heads.
+        """
+        return self.layer_norm(rms_norm=True, for_qk=True)
+
+    def v4_mlp_activation_func(self) -> Optional[type]:
+        """Activation-func selection for V4 MLP / shared-expert specs.
+
+        Plan-2 P18 (D2 audit): the parent ``activation_func()`` returns
+        the TE module **type** (``TEActivationOp``); but Megatron's
+        ``MLP.__init__`` only consumes ``submodules.activation_func`` when
+        ``config.use_te_activation_func == True`` (otherwise it falls
+        back to the callable in ``config.activation_func``).
+
+        Returning the TE class unconditionally caused a silent
+        contract mismatch in V4 yamls (which keep
+        ``use_te_activation_func: false`` by default — V4 wants the
+        clamped-SwiGLU eager path so the activation-clamp gets
+        applied).
+
+        Behavior:
+
+        * ``config.use_te_activation_func`` is True → return the TE
+          activation class (instantiated by Megatron MLP at build).
+        * Otherwise → return ``None`` so the spec leaves the
+          ``MLPSubmodules.activation_func`` slot empty and Megatron
+          MLP uses ``config.activation_func`` (V4's clamped SwiGLU).
+
+        This keeps the spec self-consistent: if the V4 yaml opts into
+        the TE path, the spec carries the TE class; otherwise the
+        spec carries ``None`` instead of a class that would be
+        silently ignored.
+        """
+        cfg = getattr(self, "config", None)
+        if cfg is not None and bool(getattr(cfg, "use_te_activation_func", False)):
+            return self.activation_func()
+        return None
+
+    def v4_grouped_mlp_modules(
+        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: Optional[bool] = None
+    ):
+        """Grouped-MLP module selection for V4 MoE expert path."""
+        return self.grouped_mlp_modules(
+            moe_use_grouped_gemm=moe_use_grouped_gemm,
+            moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        )
+
+    # ---- V4 spec factories (plan-2 P14 §5/§6) -------------------------
+
+    def v4_grouped_mlp_spec(
+        self,
+        *,
+        swiglu_limit: float,
+        moe_use_grouped_gemm: bool = True,
+        moe_use_legacy_grouped_gemm: Optional[bool] = None,
+    ):
+        """Return a ready-to-use ``ModuleSpec`` for V4 grouped MoE experts.
+
+        The V4 pre-multiplication clamp itself is applied through
+        ``config.activation_func_clamp_value`` (Megatron's MLP eager
+        ``glu()`` already clamps gate (max=alpha) and up (+/- alpha)
+        before SiLU + multiply, which is bit-equal to the HF reference
+        ``Expert.forward`` math). This spec only commits to the right
+        grouped module + the column / row-parallel linears; the runtime
+        config carries the clamp value.
+
+        Args:
+            swiglu_limit: V4 ``alpha`` value (the released ``DeepSeek-V4-Flash``
+                checkpoint uses ``7.0``). Recorded on the returned spec
+                so the caller can assert it lines up with the runtime
+                config; not consumed by the grouped-MLP module directly.
+            moe_use_grouped_gemm: prefer the TE / Turbo grouped-gemm
+                module over the local SequentialMLP fallback (default
+                True in production).
+            moe_use_legacy_grouped_gemm: optional override for the
+                legacy code path; when ``None`` the provider reads the
+                Primus arg of the same name.
+
+        Returns:
+            A ``ModuleSpec`` whose ``module`` is the grouped MoE module
+            and whose ``submodules`` carry the linear modules. Caller
+            wires this into :class:`DeepseekV4MoESubmodules.grouped_experts`.
+
+        Notes:
+            * Plan-2 P14 §5 calls for "downgrade to local experts with
+              explicit warning" when the grouped backend cannot apply
+              the clamp; that downgrade lives in
+              :class:`DeepseekV4MoE` (which already builds
+              :class:`ClampedSwiGLUMLP` local experts when
+              ``pg_collection is None`` or the backend declares no
+              clamp support).
+        """
+        del swiglu_limit  # documented but not consumed by the grouped MLP itself
+        from megatron.core.transformer.spec_utils import ModuleSpec
+
+        module, submodules = self.v4_grouped_mlp_modules(
+            moe_use_grouped_gemm=moe_use_grouped_gemm,
+            moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        )
+        if submodules is None:
+            return ModuleSpec(module=module)
+        return ModuleSpec(module=module, submodules=submodules)
+
+    def v4_router_spec(self, *, learned: bool = True):
+        """Return a ``ModuleSpec`` for the V4 hash / learned router.
+
+        Args:
+            learned: when True (default) returns the learned router
+                spec (``layer_idx >= num_hash_layers``); when False
+                returns the hash-router spec (``layer_idx <
+                num_hash_layers``).
+
+        Returns:
+            A bare-module ``ModuleSpec`` suitable for
+            :class:`DeepseekV4MoESubmodules.{learned_router, hash_router}`.
+            Both routers are ``nn.Module`` standalones (not
+            ``TopKRouter`` subclasses) so they instantiate cleanly on
+            CPU; aux-loss / z-loss / RouterReplay inheritance is
+            tracked as a P19 follow-up.
+        """
+        from megatron.core.transformer.spec_utils import ModuleSpec
+
+        from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+            DeepseekV4HashRouter,
+        )
+        from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+            DeepseekV4LearnedRouter,
+        )
+
+        return ModuleSpec(module=(DeepseekV4LearnedRouter if learned else DeepseekV4HashRouter))
diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py
index e141b4f8e..12602d698 100644
--- a/primus/backends/megatron/core/fp4_utils.py
+++ b/primus/backends/megatron/core/fp4_utils.py
@@ -1,5 +1,8 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from Megatron-LM
+# (https://github.com/NVIDIA/Megatron-LM), megatron/core/fp4_utils.py.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -14,7 +17,7 @@
 from megatron.core.utils import is_te_min_version
 
 from primus.backends.megatron.core.enums import Fp4Recipe
-from primus.modules.module_utils import warning_rank_0
+from primus.core.utils.module_utils import warning_rank_0
 
 # Check if Transformer Engine is installed
 HAVE_TE = False
@@ -45,8 +48,30 @@ def _primus_turbo_enabled() -> bool:
 
         args = get_args()
         enable_primus_turbo = bool(getattr(args, "enable_primus_turbo", False))
-        use_turbo_fp4_autocast = bool(getattr(args, "use_turbo_fp4_autocast", False))
-        return enable_primus_turbo and use_turbo_fp4_autocast
+        return enable_primus_turbo
+    except Exception:
+        return False
+
+
+_UNSET = object()
+
+
+def _mxfp4_gradient_sr_enabled(config: TransformerConfig) -> bool:
+    """Resolve gradient stochastic rounding from model config or global args.
+
+    Megatron's ``TransformerConfig`` only receives declared dataclass fields,
+    while this Primus option lives on the global args namespace. Diffusion
+    configs declare the field directly, so an explicit config value takes
+    precedence over the args fallback.
+    """
+    value = getattr(config, "mxfp4_gradient_stochastic_rounding", _UNSET)
+    if value is not _UNSET:
+        return bool(value)
+
+    try:
+        from megatron.training.global_vars import get_args
+
+        return bool(getattr(get_args(), "mxfp4_gradient_stochastic_rounding", False))
     except Exception:
         return False
 
@@ -68,7 +93,8 @@ def get_fp4_recipe(config: TransformerConfig):
                     fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
                 except AttributeError:
                     fp4_recipe_none_reason = (
-                        "NVFP4BlockScaling recipe is not available in this version of Transformer Engine."
+                        "NVFP4BlockScaling recipe is not available in this version of "
+                        "Transformer Engine. Please make sure you are using TE version >= 2.7.0.dev0."
                     )
             elif config.fp4_recipe == Fp4Recipe.mxfp4:
                 try:
@@ -78,10 +104,14 @@ def get_fp4_recipe(config: TransformerConfig):
                     fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1"
                 except AttributeError:
                     fp4_recipe_none_reason = (
-                        "MXFP4BlockScaling recipe is not available in this version of Transformer Engine."
+                        "MXFP4BlockScaling recipe is not available in this version of "
+                        "Transformer Engine. MXFP4 requires ROCm TE with AITER support."
                     )
             else:
-                fp4_recipe_none_reason = f"Unsupported FP4 recipe: {config.fp4_recipe}."
+                fp4_recipe_none_reason = (
+                    f"Unsupported fp4_recipe '{config.fp4_recipe}'. "
+                    "Supported recipes: 'nvfp4' (NVIDIA), 'mxfp4' (AMD ROCm with AITER)."
+                )
         else:
             fp4_recipe_none_reason = "FP4 support requires TransformerEngine version >= 2.7.0.dev0."
 
@@ -116,6 +146,7 @@ def get_fp4_quant_config(config: TransformerConfig):
             format=Format.E2M1_X2,
             block_size=MXFP4_SCALING_BLOCK_SIZE,
             scale_dtype=ScaleDtype.E8M0,
+            use_gradient_sr=_mxfp4_gradient_sr_enabled(config),
         )
         return fp4_quant_config, ""
 
@@ -133,62 +164,67 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
         elif layer_no >= 0 and config.first_last_layers_bf16 and (is_first_layer or is_last_layer):
             fp4_context = nullcontext()
         else:
-            fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config)
-            turbo_enabled = _primus_turbo_enabled()
-
-            global WARN_ONCE
-            if WARN_ONCE:
-                if fp4_recipe is None:
-                    warning_rank_0(
-                        f"TransformerEngine FP4 {config.fp4_recipe} not work since {fp4_recipe_none_reason}"
-                    )
-                if is_init:
-                    warning_rank_0(
-                        f"Primus-Turbo FP4 {config.fp4_recipe} not work since Primus-Turbo not support fp4 model init."
-                    )
-                WARN_ONCE = False
-
-            fp4_group = None
-            if parallel_state.model_parallel_is_initialized():
-                fp4_group = parallel_state.get_amax_reduction_group(
-                    with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red
-                )
+            # Local spec bypasses TE autocast -- quantization is handled
+            # internally by PrimusTurboMXFP4LocalSpecProvider's custom ops.
+            if getattr(config, "transformer_impl", "transformer_engine") == "local":
+                fp4_context = nullcontext()
+            else:
+                fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config)
+                turbo_enabled = _primus_turbo_enabled()
 
-            if not is_init:
-                # Only touch the Primus-Turbo extension when the Turbo FP4
-                # autocast path is explicitly enabled; otherwise use TE directly.
-                if turbo_enabled:
-                    fp4_quant_config, fp4_quant_config_none_reason = get_fp4_quant_config(config)
-                    if WARN_ONCE and fp4_quant_config is None:
+                global WARN_ONCE
+                if WARN_ONCE:
+                    if fp4_recipe is None:
+                        warning_rank_0(
+                            f"TransformerEngine FP4 {config.fp4_recipe} not work since {fp4_recipe_none_reason}"
+                        )
+                    if is_init:
                         warning_rank_0(
-                            f"Primus-Turbo FP4 {config.fp4_recipe} not work since {fp4_quant_config_none_reason}"
+                            f"Primus-Turbo FP4 {config.fp4_recipe} not work since Primus-Turbo not support fp4 model init."
                         )
+                    WARN_ONCE = False
 
-                    from primus.backends.megatron.core.extensions.primus_turbo import (
-                        primus_turbo_fp4_autocast,
+                fp4_group = None
+                if parallel_state.model_parallel_is_initialized():
+                    fp4_group = parallel_state.get_amax_reduction_group(
+                        with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red
                     )
 
-                    fp4_context = primus_turbo_fp4_autocast(
-                        enabled=True if fp4_recipe is not None else False,
-                        fp4_recipe=fp4_recipe,
-                        fp4_group=fp4_group,
-                        enabled_turbo=True if fp4_quant_config is not None else False,
-                        turbo_quant_config=fp4_quant_config,
-                    )
+                if not is_init:
+                    # Only touch the Primus-Turbo extension when the Turbo FP4
+                    # autocast path is explicitly enabled; otherwise use TE directly.
+                    if turbo_enabled:
+                        fp4_quant_config, fp4_quant_config_none_reason = get_fp4_quant_config(config)
+                        if WARN_ONCE and fp4_quant_config is None:
+                            warning_rank_0(
+                                f"Primus-Turbo FP4 {config.fp4_recipe} not work since {fp4_quant_config_none_reason}"
+                            )
+
+                        from primus.backends.megatron.core.extensions.primus_turbo import (
+                            primus_turbo_fp4_autocast,
+                        )
+
+                        fp4_context = primus_turbo_fp4_autocast(
+                            enabled=True if fp4_recipe is not None else False,
+                            fp4_recipe=fp4_recipe,
+                            fp4_group=fp4_group,
+                            enabled_turbo=True if fp4_quant_config is not None else False,
+                            turbo_quant_config=fp4_quant_config,
+                        )
+                    else:
+                        # TE currently uses fp8_autocast for fp8 and fp4 quantization.
+                        fp4_context = transformer_engine.pytorch.fp8_autocast(
+                            enabled=True if fp4_recipe is not None else False,
+                            fp8_recipe=fp4_recipe,
+                            fp8_group=fp4_group,
+                        )
                 else:
-                    # TE currently uses fp8_autocast for fp8 and fp4 quantization.
-                    fp4_context = transformer_engine.pytorch.fp8_autocast(
-                        enabled=True if fp4_recipe is not None else False,
-                        fp8_recipe=fp4_recipe,
-                        fp8_group=fp4_group,
-                    )
-            else:
-                import inspect
+                    import inspect
 
-                context_args = {"enabled": True}
-                if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters:
-                    context_args["recipe"] = fp4_recipe
-                fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args)
+                    context_args = {"enabled": True}
+                    if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters:
+                        context_args["recipe"] = fp4_recipe
+                    fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args)
 
         return fp4_context
 
@@ -196,22 +232,34 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
 
     def get_fp4_recipe(config: TransformerConfig):
         """Return fp4 recipe."""
-        if config.fp4_recipe == Fp4Recipe.nvfp4:
-            if not is_te_min_version("2.7.0.dev0"):
-                raise ValueError("NVFP4BlockScaling requires TransformerEngine >= 2.7.0.dev0.")
-            fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
-        elif config.fp4_recipe == Fp4Recipe.mxfp4:
-            try:
-                import os
-
-                fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling()
-                fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1"
-            except AttributeError:
+        if is_te_min_version("2.7.0.dev0"):
+            if config.fp4_recipe == Fp4Recipe.nvfp4:
+                try:
+                    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+                except AttributeError:
+                    raise ValueError(
+                        "NVFP4BlockScaling recipe is not available in this version of "
+                        "Transformer Engine. Please make sure you are using TE version "
+                        ">= 2.7.0.dev0."
+                    )
+            elif config.fp4_recipe == Fp4Recipe.mxfp4:
+                try:
+                    import os
+
+                    fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling()
+                    fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1"
+                except AttributeError:
+                    raise ValueError(
+                        "MXFP4BlockScaling recipe is not available in this version of "
+                        "Transformer Engine. MXFP4 requires ROCm TE with AITER support."
+                    )
+            else:
                 raise ValueError(
-                    "MXFP4BlockScaling recipe is not available in this version of Transformer Engine."
+                    f"Unsupported fp4_recipe '{config.fp4_recipe}'. "
+                    "Supported recipes: 'nvfp4' (NVIDIA), 'mxfp4' (AMD ROCm with AITER)."
                 )
         else:
-            raise ValueError(f"Unsupported FP4 recipe: {config.fp4_recipe}. " "Supported: nvfp4, mxfp4.")
+            raise ValueError("FP4 support requires TransformerEngine version >= 2.7.0.dev0.")
         return fp4_recipe
 
     def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
@@ -228,25 +276,29 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
         elif layer_no >= 0 and config.first_last_layers_bf16 and (is_first_layer or is_last_layer):
             fp4_context = nullcontext()
         else:
-            fp4_recipe = get_fp4_recipe(config)
-            fp4_group = None
-            if parallel_state.model_parallel_is_initialized():
-                fp4_group = parallel_state.get_amax_reduction_group(
-                    with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red
-                )
-
-            if not is_init:
-                # TE currently uses fp8_autocast for fp8 and fp4 quantization.
-                fp4_context = transformer_engine.pytorch.fp8_autocast(
-                    enabled=True, fp8_recipe=fp4_recipe, fp8_group=fp4_group
-                )
+            # Local spec bypasses TE autocast -- quantization is handled
+            # internally by PrimusTurboMXFP4LocalSpecProvider's custom ops.
+            if getattr(config, "transformer_impl", "transformer_engine") == "local":
+                fp4_context = nullcontext()
             else:
-                import inspect
+                fp4_recipe = get_fp4_recipe(config)
+                fp4_group = None
+                if parallel_state.model_parallel_is_initialized():
+                    fp4_group = parallel_state.get_amax_reduction_group(
+                        with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red
+                    )
+
+                if not is_init:
+                    fp4_context = transformer_engine.pytorch.fp8_autocast(
+                        enabled=True, fp8_recipe=fp4_recipe, fp8_group=fp4_group
+                    )
+                else:
+                    import inspect
 
-                context_args = {"enabled": True}
-                if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters:
-                    context_args["recipe"] = fp4_recipe
-                fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args)
+                    context_args = {"enabled": True}
+                    if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters:
+                        context_args["recipe"] = fp4_recipe
+                    fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args)
 
         return fp4_context
 
diff --git a/primus/backends/megatron/core/fp8_utils.py b/primus/backends/megatron/core/fp8_utils.py
index 7fa662c28..66e4d2200 100644
--- a/primus/backends/megatron/core/fp8_utils.py
+++ b/primus/backends/megatron/core/fp8_utils.py
@@ -12,7 +12,7 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 
-from primus.modules.module_utils import warning_rank_0
+from primus.core.utils.module_utils import warning_rank_0
 
 # Check if Transformer Engine is installed
 HAVE_TE = False
@@ -43,6 +43,16 @@
     # Primus-Turbo not importable (not installed, or a transitive dep is broken).
     pass
 
+# gfx1250 / turbo-free FP8: on builds where primus_turbo is only an import shim
+# (not a real install), the turbo FP8 path (PrimusTurboQuantConfig +
+# primus_turbo_fp8_autocast) cannot run. Set PRIMUS_FP8_DISABLE_TURBO=1 to force
+# the TE-native FP8 branch (transformer_engine.pytorch.fp8_autocast), which needs
+# no turbo. Opt-in; the production turbo path is unaffected when unset.
+import os as _os
+
+if _os.environ.get("PRIMUS_FP8_DISABLE_TURBO", "0") == "1":
+    HAVE_TURBO = False
+
 
 SCALING_BLOCK_SIZE = 128
 MXFP8_SCALING_BLOCK_SIZE = 32
@@ -73,14 +83,19 @@ def te_fp8_format_mapping(te_format):
         return format_mapping[te_format]
 
     def get_fp8_recipe(config: TransformerConfig):
-        """Return fp8 recipe.
+        """Return fp8 recipe (Primus Turbo + TE variant).
 
         Arguments:
             config (TransformerConfig): Configuration object.
 
         Returns:
-            FP8 recipe: Transformer Engine FP8 recipe.
-            FP8 None reason: reason why the fp8 recipe is None.
+            Tuple of (fp8_recipe, fp8_recipe_none_reason):
+                - fp8_recipe: Transformer Engine FP8 recipe, or None.
+                - fp8_recipe_none_reason: reason why the fp8 recipe is None.
+
+        Note:
+            This variant (HAVE_TE + HAVE_TURBO) returns a tuple. The TE-only
+            variant returns a single recipe value.
         """
         if config.fp8 == "e4m3":
             fp8_format = transformer_engine.common.recipe.Format.E4M3
diff --git a/primus/backends/megatron/core/fusions/fused_bias_swiglu.py b/primus/backends/megatron/core/fusions/fused_bias_swiglu.py
new file mode 100644
index 000000000..19a689c7f
--- /dev/null
+++ b/primus/backends/megatron/core/fusions/fused_bias_swiglu.py
@@ -0,0 +1,342 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus override for clamped weighted SwiGLU MoE activation fusion.
+
+Fuses clamp(gate), clamp(up), SiLU, multiply, and router-prob weighting into
+a single ``@jit_fuser`` forward/backward pair instead of separate ``chunk`` /
+``clamp`` / ``clamp`` / ``cat`` ATen ops followed by ``WeightedSwiGLUFunction``.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from megatron.core.fusions.fused_bias_swiglu import WeightedSwiGLUFunction
+
+# Column tile size for the elementwise kernels / reduction inner loop.
+_BLOCK_N = 1024
+
+
+@triton.jit
+def _clamped_weighted_swiglu_fwd_kernel(
+    y_ptr,  # [M, 2*half] row-major: gate = [:, :half], up = [:, half:]
+    w_ptr,  # [M] per-token weights (only read if HAS_WEIGHTS)
+    out_ptr,  # [M, half] row-major output
+    half,
+    K,  # = 2 * half (row stride of y)
+    stride_w,
+    clamp_value,
+    HAS_WEIGHTS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = cols < half
+
+    y_row = y_ptr + pid_m * K
+    gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32)
+
+    gate_c = tl.minimum(gate, clamp_value)
+    up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value)
+    out = (gate_c * tl.sigmoid(gate_c)) * up_c
+    if HAS_WEIGHTS:
+        out = out * tl.load(w_ptr + pid_m * stride_w).to(tl.float32)
+
+    tl.store(out_ptr + pid_m * half + cols, out, mask=mask)
+
+
+@triton.jit
+def _clamped_swiglu_bwd_kernel(
+    g_ptr,  # [M, half] grad w.r.t. output
+    y_ptr,  # [M, 2*half] pre-clamp input
+    w_ptr,  # [M] per-token weights (only read if HAS_WEIGHTS)
+    dy_ptr,  # [M, 2*half] grad w.r.t. input
+    half,
+    K,
+    stride_gm,
+    stride_w,
+    clamp_value,
+    HAS_WEIGHTS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = cols < half
+
+    g = tl.load(g_ptr + pid_m * stride_gm + cols, mask=mask, other=0.0).to(tl.float32)
+    if HAS_WEIGHTS:
+        g = g * tl.load(w_ptr + pid_m * stride_w).to(tl.float32)
+
+    y_row = y_ptr + pid_m * K
+    gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32)
+
+    gate_c = tl.minimum(gate, clamp_value)
+    up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value)
+    sig = tl.sigmoid(gate_c)
+    silu = gate_c * sig
+
+    # d/d gate_c [SiLU(gate_c)] = sig * (1 + gate_c * (1 - sig))
+    dy_glu_c = g * sig * (1.0 + gate_c * (1.0 - sig)) * up_c
+    dy_linear_c = g * silu
+
+    # Straight-through the clamp: gradient passes only inside the (un)clamped region.
+    gate_keep = (gate <= clamp_value).to(tl.float32)
+    up_keep = ((up >= -clamp_value) & (up <= clamp_value)).to(tl.float32)
+
+    dy_row = dy_ptr + pid_m * K
+    tl.store(dy_row + cols, dy_glu_c * gate_keep, mask=mask)
+    tl.store(dy_row + half + cols, dy_linear_c * up_keep, mask=mask)
+
+
+@triton.jit
+def _clamped_swiglu_weights_grad_kernel(
+    g_ptr,  # [M, half] grad w.r.t. output
+    y_ptr,  # [M, 2*half] pre-clamp input
+    wg_ptr,  # [M] per-token weights grad
+    half,
+    K,
+    stride_gm,
+    clamp_value,
+    NUM_TILES: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(axis=0)
+    y_row = y_ptr + pid_m * K
+    g_row = g_ptr + pid_m * stride_gm
+
+    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
+    for t in tl.static_range(NUM_TILES):
+        cols = t * BLOCK_N + tl.arange(0, BLOCK_N)
+        mask = cols < half
+        g = tl.load(g_row + cols, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32)
+        up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32)
+        gate_c = tl.minimum(gate, clamp_value)
+        up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value)
+        activation = (gate_c * tl.sigmoid(gate_c)) * up_c
+        acc += activation * g  # masked lanes contribute 0 (g == 0)
+
+    tl.store(wg_ptr + pid_m, tl.sum(acc, axis=0))
+
+
+def clamped_weighted_swiglu(y: torch.Tensor, weights: torch.Tensor, clamp_value: float) -> torch.Tensor:
+    """Clamped SwiGLU with per-token weights in one fused Triton kernel.
+
+    Semantics match the eager path:
+        gate_c = clamp(gate, max=alpha)
+        up_c   = clamp(up, min=-alpha, max=alpha)
+        out    = SiLU(gate_c) * up_c * weights
+    """
+    lead = y.shape[:-1]
+    K = y.shape[-1]
+    half = K // 2
+    y2 = y.reshape(-1, K).contiguous()
+    M = y2.shape[0]
+    w = weights.reshape(M).contiguous()
+
+    out = torch.empty((M, half), dtype=y.dtype, device=y.device)
+    grid = (M, triton.cdiv(half, _BLOCK_N))
+    _clamped_weighted_swiglu_fwd_kernel[grid](
+        y2,
+        w,
+        out,
+        half,
+        K,
+        w.stride(0),
+        float(clamp_value),
+        HAS_WEIGHTS=True,
+        BLOCK_N=_BLOCK_N,
+    )
+    return out.reshape(*lead, half)
+
+
+def clamped_swiglu(y: torch.Tensor, clamp_value: float) -> torch.Tensor:
+    """Clamped SwiGLU without per-token weights in one fused Triton kernel.
+
+    Semantics match the eager path used by ``mlp.MLP`` shared experts:
+        gate_c = clamp(gate, max=alpha)
+        up_c   = clamp(up, min=-alpha, max=alpha)
+        out    = SiLU(gate_c) * up_c
+    """
+    lead = y.shape[:-1]
+    K = y.shape[-1]
+    half = K // 2
+    y2 = y.reshape(-1, K).contiguous()
+    M = y2.shape[0]
+
+    out = torch.empty((M, half), dtype=y.dtype, device=y.device)
+    grid = (M, triton.cdiv(half, _BLOCK_N))
+    # w_ptr is unused when HAS_WEIGHTS=False; pass y2 as a valid placeholder.
+    _clamped_weighted_swiglu_fwd_kernel[grid](
+        y2,
+        y2,
+        out,
+        half,
+        K,
+        0,
+        float(clamp_value),
+        HAS_WEIGHTS=False,
+        BLOCK_N=_BLOCK_N,
+    )
+    return out.reshape(*lead, half)
+
+
+def clamped_swiglu_back(g: torch.Tensor, y: torch.Tensor, clamp_value: float) -> torch.Tensor:
+    """Backward for clamped SwiGLU w.r.t. the pre-clamp concatenated input."""
+    lead = y.shape[:-1]
+    K = y.shape[-1]
+    half = K // 2
+    y2 = y.reshape(-1, K).contiguous()
+    g2 = g.reshape(-1, half).contiguous()
+    M = y2.shape[0]
+
+    dy = torch.empty((M, K), dtype=g.dtype, device=g.device)
+    grid = (M, triton.cdiv(half, _BLOCK_N))
+    _clamped_swiglu_bwd_kernel[grid](
+        g2,
+        y2,
+        g2,
+        dy,
+        half,
+        K,
+        g2.stride(0),
+        0,
+        float(clamp_value),
+        HAS_WEIGHTS=False,
+        BLOCK_N=_BLOCK_N,
+    )
+    return dy.reshape(*lead, K)
+
+
+def clamped_weighted_swiglu_back(g: torch.Tensor, y: torch.Tensor, weights: torch.Tensor, clamp_value: float):
+    """Backward for :func:`clamped_weighted_swiglu`."""
+    input_dtype = y.dtype
+    w_dtype = weights.dtype
+    lead = y.shape[:-1]
+    K = y.shape[-1]
+    half = K // 2
+    y2 = y.reshape(-1, K).contiguous()
+    g2 = g.reshape(-1, half).contiguous()
+    M = y2.shape[0]
+    w = weights.reshape(M).contiguous()
+
+    # input grad: reuse the clamped-swiglu backward kernel with grad pre-scaled by weights.
+    input_grad = torch.empty((M, K), dtype=input_dtype, device=y.device)
+    grid = (M, triton.cdiv(half, _BLOCK_N))
+    _clamped_swiglu_bwd_kernel[grid](
+        g2,
+        y2,
+        w,
+        input_grad,
+        half,
+        K,
+        g2.stride(0),
+        w.stride(0),
+        float(clamp_value),
+        HAS_WEIGHTS=True,
+        BLOCK_N=_BLOCK_N,
+    )
+
+    # weights grad: sum over the feature dim of SiLU(gate_c) * up_c * g.
+    weights_grad = torch.empty((M,), dtype=w_dtype, device=y.device)
+    _clamped_swiglu_weights_grad_kernel[(M,)](
+        g2,
+        y2,
+        weights_grad,
+        half,
+        K,
+        g2.stride(0),
+        float(clamp_value),
+        NUM_TILES=triton.cdiv(half, _BLOCK_N),
+        BLOCK_N=_BLOCK_N,
+    )
+
+    return input_grad.reshape(*lead, K), weights_grad.reshape(*weights.shape)
+
+
+class ClampedWeightedSwiGLUFunction(torch.autograd.Function):
+    """Autograd wrapper for clamped token-wise weighted SwiGLU."""
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weights: torch.Tensor, fp8_input_store: bool, clamp_value: float):
+        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        ctx.save_for_backward(input_for_backward, weights)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = float(clamp_value)
+        return clamped_weighted_swiglu(input, weights, ctx.clamp_value)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weights = ctx.saved_tensors
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        input_grad, weights_grad = clamped_weighted_swiglu_back(grad_output, input, weights, ctx.clamp_value)
+        return input_grad, weights_grad, None, None
+
+
+def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp_value=None):
+    """Token-wise-weighted SwiGLU fusion with optional clamped gate/up."""
+    ori_shape = input.shape
+    assert len(ori_shape) in [2, 3]
+    input = input.view(-1, ori_shape[-1])
+    if bias is not None:
+        raise NotImplementedError("Bias is not supported for weighted swiglu fusion")
+
+    if clamp_value is not None:
+        output = ClampedWeightedSwiGLUFunction.apply(input, weights, fp8_input_store, float(clamp_value))
+    else:
+        output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store)
+
+    return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
+
+
+class ClampedSwiGLUFunction(torch.autograd.Function):
+    """Autograd wrapper for clamped (non-weighted) SwiGLU."""
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, fp8_input_store: bool, clamp_value: float):
+        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        ctx.save_for_backward(input_for_backward)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = float(clamp_value)
+        return clamped_swiglu(input, ctx.clamp_value)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (input,) = ctx.saved_tensors
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        input_grad = clamped_swiglu_back(grad_output, input, ctx.clamp_value)
+        return input_grad, None, None
+
+
+def swiglu_impl(input, bias, fp8_input_store=False, clamp_value=None):
+    """Non-weighted SwiGLU fusion with optional clamped gate/up.
+
+    Mirrors ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` but
+    adds the DeepSeek-V4 pre-multiplication clamp when ``clamp_value`` is set.
+    Used by the shared-expert MLP path, which has no per-token weights.
+    """
+    ori_shape = input.shape
+    assert len(ori_shape) in [2, 3]
+    input = input.view(-1, ori_shape[-1])
+    if bias is not None:
+        raise NotImplementedError("Bias is not supported for clamped swiglu fusion")
+
+    if clamp_value is not None:
+        output = ClampedSwiGLUFunction.apply(input, fp8_input_store, float(clamp_value))
+    else:
+        from megatron.core.fusions.fused_bias_swiglu import SwiGLUFunction
+
+        output = SwiGLUFunction.apply(input, fp8_input_store, False)
+
+    return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
diff --git a/primus/backends/megatron/core/fusions/fused_pad_routing_map.py b/primus/backends/megatron/core/fusions/fused_pad_routing_map.py
new file mode 100644
index 000000000..c61c58d58
--- /dev/null
+++ b/primus/backends/megatron/core/fusions/fused_pad_routing_map.py
@@ -0,0 +1,114 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus override for the fused routing-map padding used by MoE quantization.
+
+Rewrites Megatron's ``fused_pad_routing_map`` as a self-contained Triton kernel
+that operates directly on the native ``[num_tokens, num_experts]`` layout, so no
+``transpose`` / ``contiguous`` / intermediate copy is needed. It is intentionally
+*not* wrapped with ``@jit_fuser`` (``torch.compile``): the kernel is already the
+fused region, and wrapping user-written Triton kernels with ``torch.compile``
+triggers a functionalization failure on some torch/triton combos.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import torch
+from megatron.core.utils import null_decorator
+from packaging import version
+
+try:
+    import triton
+    import triton.language as tl
+
+    if version.parse(triton.__version__) < version.parse("3.4.0") and not torch.cuda.is_available():
+        HAVE_TRITON = False
+    else:
+        HAVE_TRITON = tl.constexpr(version.parse(triton.__version__) >= version.parse("2.0.0"))
+except ImportError:
+    HAVE_TRITON = False
+
+if not HAVE_TRITON:
+    triton = MagicMock()
+    triton.jit = null_decorator
+    triton.autotune = null_decorator
+    triton.heuristics = null_decorator
+    tl = MagicMock()
+
+
+@triton.jit
+def _pad_routing_map_kernel(
+    routing_map_ptr,  # *Pointer* to [num_tokens, num_experts] row-major routing map
+    output_ptr,  # *Pointer* to [num_tokens, num_experts] row-major output
+    num_tokens,
+    num_experts,
+    pad_multiple: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    # Each program instance handles one expert, i.e. one *column* of the
+    # [num_tokens, num_experts] routing map. This lets us operate directly on the
+    # native (row-major) layout, without transposing/copying to [num_experts, num_tokens].
+    expert_idx = tl.program_id(axis=0)
+
+    # Token indices for this block
+    token_indices = tl.arange(0, BLOCK_SIZE)
+    token_mask = token_indices < num_tokens
+
+    # Column-strided access: element (t, expert_idx) lives at t * num_experts + expert_idx.
+    offsets = token_indices * num_experts + expert_idx
+
+    # Load this expert's column; out-of-bounds tokens read as 0 and are masked on store.
+    row = tl.load(routing_map_ptr + offsets, mask=token_mask, other=0).to(tl.int32)
+
+    # 1. Number of tokens currently routed to this expert.
+    num_ones = tl.sum(row, axis=0)
+
+    # 2. How many zeros must be flipped to 1 to reach the next multiple of pad_multiple.
+    remainder = num_ones % pad_multiple
+    num_to_pad = tl.where(remainder != 0, pad_multiple - remainder, 0)
+
+    # 3. 1-based cumulative rank of each zero within the column.
+    is_zero = row == 0
+    zero_ranks = tl.cumsum(is_zero.to(tl.int32), axis=0)
+
+    # 4. Flip only the first `num_to_pad` zeros to 1.
+    mask_to_flip = (zero_ranks <= num_to_pad) & is_zero
+    output_row = tl.where(mask_to_flip, 1, row)
+
+    # 5. Store back in the same native layout, masking out-of-bounds tokens.
+    tl.store(output_ptr + offsets, output_row, mask=token_mask)
+
+
+def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
+    """Fused version of pad_routing_map.
+    Args:
+        routing_map (torch.Tensor): A boolean or integer tensor of shape [num_tokens,
+            num_experts] indicating which tokens are routed to which experts.
+        pad_multiple (int): The multiple to pad each expert's token count to.
+
+    Returns:
+        torch.Tensor: The padded routing map of shape [num_tokens, num_experts].
+    """
+    num_tokens, num_experts = routing_map.shape
+    if num_tokens == 0:
+        return routing_map
+
+    # Operate directly on the native [num_tokens, num_experts] layout: the kernel reads
+    # each expert's column with a stride of num_experts, so no transpose/copy is needed.
+    routing_map = routing_map.contiguous()
+    output_map = torch.empty_like(routing_map, dtype=torch.int32)
+
+    # One program instance per expert (column).
+    grid = (num_experts,)
+    BLOCK_SIZE = triton.next_power_of_2(num_tokens)
+
+    _pad_routing_map_kernel[grid](
+        routing_map, output_map, num_tokens, num_experts, pad_multiple, BLOCK_SIZE=BLOCK_SIZE
+    )
+
+    return output_map  # [num_tokens, num_experts]
diff --git a/primus/backends/megatron/core/models/common/diffusion_module/__init__.py b/primus/backends/megatron/core/models/common/diffusion_module/__init__.py
new file mode 100644
index 000000000..399364b6a
--- /dev/null
+++ b/primus/backends/megatron/core/models/common/diffusion_module/__init__.py
@@ -0,0 +1,12 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Base diffusion module with Megatron-Core integration.
+"""
+
+from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import (
+    DiffusionModule,
+)
+
+__all__ = ["DiffusionModule"]
diff --git a/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py b/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py
new file mode 100644
index 000000000..fd5d6cec1
--- /dev/null
+++ b/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py
@@ -0,0 +1,176 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Base diffusion module with Megatron-Core integration.
+
+This module provides common infrastructure for all diffusion models (Flux, DiT, etc.),
+following the architecture pattern established by LanguageModule in Megatron-LM.
+"""
+
+import os
+from abc import abstractmethod
+from typing import Any, Dict, Optional
+
+import torch.nn as nn
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.enums import AttnBackend
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import get_tensor_model_parallel_group_if_none
+
+
+class DiffusionModule(MegatronModule):
+    """
+    Base diffusion module with Megatron-Core integration.
+
+    Provides common infrastructure for all diffusion models (Flux, DiT, MovieGen, etc.):
+    - Process group management (TP, PP, CP, DP)
+    - Attention backend configuration
+    - Distributed checkpointing support
+    - Encoder management (VAE, T5, CLIP, etc.)
+    - Common loss computation utilities
+
+    This class follows the architecture pattern of LanguageModule from Megatron-LM,
+    adapted for diffusion model requirements.
+
+    Args:
+        config (TransformerConfig): Transformer config with diffusion-specific parameters
+        pg_collection (Optional[ProcessGroupCollection]): Model communication process groups.
+            Defaults to None (uses MPU process groups).
+        encoder_configs (Optional[Dict[str, Any]]): Optional encoder configurations.
+            Defaults to None.
+            Format: {'encoder_name': EncoderConfig, ...}
+            Example: {'vae': VAEConfig(...), 't5': T5XXLConfig(...), 'clip': CLIPLConfig(...)}
+
+    Example:
+        >>> from primus.backends.megatron.core.models.diffusion.flux import FluxConfig
+        >>> config = FluxConfig.flux_12b()
+        >>> model = Flux(config=config)
+
+    Reference:
+        Megatron-LM LanguageModule: megatron/core/models/common/language_module/language_module.py
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        encoder_configs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        super().__init__(config=config)
+
+        # Configure attention backend
+        self._set_attention_backend()
+
+        # Setup process groups for distributed training
+        if pg_collection is None:
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        self.pg_collection = pg_collection
+        self.cp_group = pg_collection.cp
+        self.tp_group = get_tensor_model_parallel_group_if_none(pg_collection.tp)
+        self.pp_group = pg_collection.pp
+
+        # Note: Diffusion models typically don't need embd_group since they don't share
+        # embeddings across pipeline stages like language models do
+
+        # Virtual pipeline parallelism
+        self.vp_stage = None
+        self.vp_size = self.config.virtual_pipeline_model_parallel_size
+
+        # Encoder management (VAE, T5, CLIP, etc.)
+        self.encoder_configs = encoder_configs or {}
+        self.encoders = nn.ModuleDict()
+
+    def _set_attention_backend(self):
+        """
+        Set attention backend for TransformerEngine.
+
+        Configures environment variables to control which attention implementation
+        TransformerEngine uses. Options:
+        - flash: Flash Attention (fastest, requires Ampere+ GPUs)
+        - fused: Fused attention kernels
+        - unfused: Standard unfused attention
+        - auto: Let TE choose automatically
+        - local: Disable TE attention, use PyTorch native
+
+        TransformerEngine works on opt-out basis. By default all three attention
+        backend flags are set to 1. If the user chooses a particular attention
+        backend, we set the other two to 0. If the user chooses local, we set
+        all 3 TE env variables to 0.
+
+        Reference:
+            LanguageModule._set_attention_backend() in Megatron-LM
+        """
+
+        def check_and_set_env_variable(
+            env_variable_name: str, expected_value: int, attn_type: AttnBackend
+        ) -> None:
+            current_value = os.getenv(env_variable_name)
+            if current_value is not None and current_value != str(expected_value):
+                raise ValueError(
+                    f"{env_variable_name} set to {current_value}, but expected {expected_value} "
+                    f"for attention backend type {attn_type.name}. Unset NVTE_FLASH_ATTN, "
+                    f"NVTE_FUSED_ATTN and NVTE_UNFUSED_ATTN. Use the --attention-backend argument "
+                    f"if you want to choose between (flash/fused/unfused/auto/local). Default is auto."
+                )
+            os.environ[env_variable_name] = str(expected_value)
+
+        if self.config.attention_backend == AttnBackend.local:
+            check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.local)
+            check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.local)
+            check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.local)
+        elif self.config.attention_backend == AttnBackend.flash:
+            check_and_set_env_variable("NVTE_FLASH_ATTN", 1, AttnBackend.flash)
+            check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.flash)
+            check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.flash)
+        elif self.config.attention_backend == AttnBackend.fused:
+            check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.fused)
+            check_and_set_env_variable("NVTE_FUSED_ATTN", 1, AttnBackend.fused)
+            check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.fused)
+        elif self.config.attention_backend == AttnBackend.unfused:
+            check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.unfused)
+            check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.unfused)
+            check_and_set_env_variable("NVTE_UNFUSED_ATTN", 1, AttnBackend.unfused)
+        elif self.config.attention_backend == AttnBackend.auto:
+            check_and_set_env_variable("NVTE_FLASH_ATTN", 1, AttnBackend.auto)
+            check_and_set_env_variable("NVTE_FUSED_ATTN", 1, AttnBackend.auto)
+            check_and_set_env_variable("NVTE_UNFUSED_ATTN", 1, AttnBackend.auto)
+
+    @abstractmethod
+    def forward(self, *args, **kwargs):
+        """
+        Forward pass through diffusion model.
+
+        This method must be implemented by all subclasses to define the
+        model's forward computation.
+
+        Returns:
+            Model prediction (noise, velocity, or other target type depending on model)
+        """
+
+    def sharded_state_dict(
+        self,
+        prefix: str = "",
+        sharded_offsets: tuple = (),
+        metadata: Optional[dict] = None,
+    ) -> ShardedStateDict:
+        """
+        Generate sharded state dictionary for distributed checkpointing.
+
+        Subclasses should override this to handle model-specific checkpointing logic
+        (e.g., weight tying, custom sharding strategies).
+
+        Args:
+            prefix: Prefix for state dict keys (e.g., 'module.')
+            sharded_offsets: Pipeline parallel offsets
+            metadata: Optional metadata for checkpoint conversion
+
+        Returns:
+            Dictionary mapping state dict keys to ShardedTensor objects
+
+        Reference:
+            LanguageModule.sharded_state_dict() in Megatron-LM
+        """
+        return super().sharded_state_dict(prefix, sharded_offsets, metadata)
diff --git a/primus/backends/megatron/core/models/deepseek_v4/__init__.py b/primus/backends/megatron/core/models/deepseek_v4/__init__.py
new file mode 100644
index 000000000..4ffc1f8a4
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/__init__.py
@@ -0,0 +1,52 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 model package.
+
+Plan-2 P17 surface (post-cleanup):
+
+    DeepseekV4Model               # top-level model (LanguageModule)
+    DeepseekV4TransformerBlock    # decoder block (TransformerBlock subclass)
+    DeepseekV4HybridLayer         # decoder layer (TransformerLayer subclass)
+    get_v4_mtp_block_spec         # spec helper for upstream MultiTokenPredictionBlock
+    deepseek_v4_builder           # builder used by model_provider
+    model_provider                # Megatron pretrain() entry point
+
+Retired in plan-2 P17:
+
+    DeepseekV4MTPBlock            # legacy primus-owned MTP head; replaced
+                                  # by the spec-based path above. The
+                                  # ``v4_use_custom_mtp_block`` config
+                                  # toggle is also gone.
+"""
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import (
+    DeepseekV4HybridLayer,
+    DeepseekV4HybridLayerSubmodules,
+    DeepseekV4TransformerBlock,
+    DeepseekV4TransformerBlockSubmodules,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+    deepseek_v4_builder,
+    model_provider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_model import (
+    DeepseekV4Model,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_specs import (
+    get_v4_mtp_block_spec,
+)
+
+__all__ = [
+    "DeepseekV4Model",
+    "DeepseekV4TransformerBlock",
+    "DeepseekV4TransformerBlockSubmodules",
+    "DeepseekV4HybridLayer",
+    "DeepseekV4HybridLayerSubmodules",
+    "get_v4_mtp_block_spec",
+    "deepseek_v4_builder",
+    "model_provider",
+]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/build_context.py b/primus/backends/megatron/core/models/deepseek_v4/build_context.py
new file mode 100644
index 000000000..1f132a049
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/build_context.py
@@ -0,0 +1,102 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-2 P18 — DeepSeek-V4 build context / provider singleton.
+
+This module hosts the helpers that make sure a single
+:class:`DeepSeekV4SpecProvider` is constructed per builder call and
+threaded down to every spec helper. Without this, the provider was
+re-instantiated inside ``_build_projection`` (block.py),
+``DeepseekV4TransformerBlock.__init__`` (block.py), the layer-spec
+factory (``deepseek_v4_layer_specs.py``), and the MTP spec helper
+(``deepseek_v4_mtp_specs.py``) — every call paid the
+``BackendSpecProvider`` setup cost and there was no single place to
+audit which provider was actually wiring the V4 modules.
+
+The implementation deliberately avoids module-level globals: the
+provider is cached **on the config object** under a private attribute
+name. This keeps each ``DeepSeekV4TransformerConfig`` instance
+self-contained (different configs get different providers) and avoids
+leaking state across unit tests.
+
+Usage:
+
+.. code-block:: python
+
+    from .build_context import resolve_v4_provider
+
+    def some_builder(*, config):
+        provider = resolve_v4_provider(config)
+        ...
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # avoid eager torch import in this lightweight module
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+    )
+
+_PROVIDER_ATTR = "_v4_spec_provider_singleton"
+
+
+def resolve_v4_provider(config) -> "DeepSeekV4SpecProvider":
+    """Return a cached :class:`DeepSeekV4SpecProvider` for ``config``.
+
+    The provider is cached on the ``config`` object itself (under a
+    private attribute name) so:
+
+    * repeated calls during a single builder invocation reuse the
+      same provider instance,
+    * different configs always get fresh providers,
+    * the cache is naturally garbage-collected when the config is
+      released,
+    * the helper is fully thread-safe-by-construction (each builder
+      thread holds its own config).
+
+    Args:
+        config: a :class:`DeepSeekV4TransformerConfig` instance. The
+            helper accesses ``config.__dict__`` directly so it works
+            on dataclasses without additional setattr support.
+
+    Returns:
+        The cached :class:`DeepSeekV4SpecProvider`. The first call
+        constructs it; later calls reuse it.
+    """
+    cached = getattr(config, _PROVIDER_ATTR, None)
+    if cached is not None:
+        return cached
+
+    # Lazy import: this module must be lightweight enough to import
+    # from the dataclass module without cyclic risks.
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+    )
+
+    provider = DeepSeekV4SpecProvider(config=config)
+    try:
+        setattr(config, _PROVIDER_ATTR, provider)
+    except Exception:
+        # Some MagicMock-style configs in unit tests reject setattr; that
+        # is OK — we just don't cache and pay the cost on each call.
+        pass
+    return provider
+
+
+def reset_v4_provider_cache(config) -> None:
+    """Drop the cached provider on ``config``. Intended for unit tests
+    that need to force a re-build (e.g. after monkey-patching the
+    provider class)."""
+    if hasattr(config, _PROVIDER_ATTR):
+        try:
+            delattr(config, _PROVIDER_ATTR)
+        except AttributeError:
+            pass
+
+
+__all__ = ["resolve_v4_provider", "reset_v4_provider_cache"]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py
new file mode 100644
index 000000000..b06eac1ea
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py
@@ -0,0 +1,1190 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DeepSeek-V4 transformer block (multi-stream HC + per-layer attention dispatch).
+
+Reference:
+* techblog §1 ("Hybrid Attention") — per-layer attention selected by
+  ``compress_ratios[layer_id]``.
+* techblog §2 ("mHC: Manifold-Constrained Hyper-Connections") — ``hc_mult``
+  parallel hidden streams, mixed via per-layer ``HyperMixer`` and final
+  ``HyperHead`` collapse.
+
+Plan-2 P15 (this commit) — Megatron parent-class integration:
+
+* :class:`DeepseekV4HybridLayer` now subclasses
+  :class:`megatron.core.transformer.transformer_layer.TransformerLayer`
+  (via ``MegatronModule`` ``__init__`` bypass) and reuses upstream
+  submodule names (``input_layernorm`` / ``self_attention`` /
+  ``pre_mlp_layernorm`` / ``mlp``) plus V4-specific ``attn_hc`` / ``ffn_hc``
+  hooks. The submodules dataclass extends
+  :class:`TransformerLayerSubmodules` so Megatron spec lifecycle code that
+  inspects the layer's ``submodules_config`` works without bespoke V4
+  branches.
+* :class:`DeepseekV4TransformerBlock` now subclasses
+  :class:`megatron.core.transformer.transformer_block.TransformerBlock`
+  for type identity / sharded-state-dict integration. ``HyperHead`` is now
+  built **only on the post_process stage** (it was previously built on
+  every PP rank, which wasted memory and risked correctness drift).
+* PP K-stream packing: ``[B, S, K, D] <-> [S*K, B, D]`` lift / lower
+  helpers (:func:`_lift_streams_in` / :func:`_lower_streams_out`) carry
+  the K dimension across PP P2P boundaries by folding it into the
+  sequence axis. The first PP stage lifts the embedded ``[S, B, D]`` to K
+  streams; intermediate stages preserve K (``[S*K, B, D]`` send/recv);
+  the final stage collapses with ``HyperHead`` to ``[B, S, D]`` and
+  transposes back to ``[S, B, D]`` for the output layer.
+* ``token_ids`` is now a real forward kwarg threaded through
+  ``DeepseekV4Model.forward -> DeepseekV4TransformerBlock.forward ->
+  DeepseekV4HybridLayer.forward -> DeepseekV4MoE.forward ->
+  DeepseekV4HashRouter.forward``. The legacy ``decoder._v4_token_ids``
+  attribute stash is gone.
+* ``position_ids`` is now consumed from the caller (forward kwarg) when
+  provided; the legacy ``arange(S)`` shortcut is kept only as a fallback
+  for callers that omit it (e.g. tiny CPU smokes).
+
+Forward shape contract:
+* Input ``hidden_states`` is ``[S, B, D]`` on the first PP stage and
+  ``[S*K, B, D]`` on subsequent stages (where K = ``hc_mult``). The block
+  reshapes to the K-stream form for HC math and packs back before
+  returning to the next PP stage.
+* Output is ``[S*K, B, D]`` on non-final PP stages and ``[S, B, D]`` on
+  the final (post_process) stage after HyperHead collapse.
+"""
+
+from __future__ import annotations
+
+import ast
+import logging
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from megatron.core.transformer.mlp import MLPSubmodules
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_layer import (
+    TransformerLayer,
+    TransformerLayerSubmodules,
+)
+from megatron.core.utils import make_viewless_tensor
+
+from primus.backends.megatron.core.models.deepseek_v4.build_context import (
+    resolve_v4_provider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import DualRoPE
+from primus.backends.megatron.core.transformer.hyper_connection import (
+    HyperHead,
+    HyperMixer,
+)
+from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+    DeepseekV4HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_moe import (
+    DeepseekV4MoE,
+    DeepseekV4MoESubmodules,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    DeepseekV4LearnedRouter,
+)
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Provider-aware projection helpers
+# ---------------------------------------------------------------------------
+
+
+def _default_init_method(_weight: torch.Tensor) -> None:
+    return None
+
+
+def _build_projection(
+    in_features: int,
+    out_features: int,
+    *,
+    config: DeepSeekV4TransformerConfig,
+) -> nn.Module:
+    """Build a duplicated linear via the V4 provider.
+
+    When ``config`` is ``None`` (CPU unit tests that build the dense
+    SwiGLU MLP without a provider) we instantiate a plain
+    :class:`nn.Linear` with the same shape.  Otherwise we delegate to
+    ``provider.linear()`` and let any constructor failure bubble up —
+    Plan-3 P21 retired the ``try/except/return nn.Linear`` fallback
+    because it produced an unsharded layer that masked real provider
+    bugs at TP=1 and would diverge at TP>1.
+    """
+    if config is None:
+        return nn.Linear(in_features, out_features, bias=False)
+
+    provider = resolve_v4_provider(config)
+    linear_module_cls = provider.linear()
+    init_method: Callable = config.init_method or _default_init_method
+    return linear_module_cls(
+        input_size=in_features,
+        output_size=out_features,
+        parallel_mode="duplicated",
+        config=config,
+        init_method=init_method,
+        bias=False,
+        skip_bias_add=False,
+        skip_weight_param_allocation=False,
+        tp_comm_buffer_name=None,
+        is_expert=False,
+    )
+
+
+def _projection_forward(proj: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    out = proj(x)
+    if isinstance(out, tuple):
+        return out[0]
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Pieces used by every layer
+# ---------------------------------------------------------------------------
+
+
+def _parse_int_sequence(value, *, field_name: str) -> Optional[List[int]]:
+    """Parse config-provided sequence fields into ``List[int]``.
+
+    YAML values may arrive as:
+    - actual list/tuple: ``[0, 4, 128, ...]``
+    - stringified list: ``"[0, 4, 128, ...]"``
+    """
+    if value is None:
+        return None
+
+    parsed = value
+    if isinstance(parsed, str):
+        try:
+            parsed = ast.literal_eval(parsed)
+        except (SyntaxError, ValueError) as exc:
+            raise ValueError(f"{field_name} must be a list-like value, got {value!r}") from exc
+
+    if isinstance(parsed, torch.Tensor):
+        parsed = parsed.tolist()
+
+    if not isinstance(parsed, (list, tuple)):
+        raise TypeError(f"{field_name} must be list/tuple, got {type(parsed).__name__}")
+
+    out: List[int] = []
+    for i, item in enumerate(parsed):
+        try:
+            out.append(int(item))
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"{field_name}[{i}]={item!r} is not int-castable") from exc
+    return out
+
+
+def _normalize_compress_ratios(
+    compress_ratios,
+    *,
+    num_layers: int,
+    mtp_num_layers: int,
+) -> List[int]:
+    """Normalize ``compress_ratios`` to exactly ``num_layers`` entries."""
+    ratios = _parse_int_sequence(compress_ratios, field_name="compress_ratios")
+    if ratios is None:
+        return [0] * num_layers
+
+    if len(ratios) == num_layers:
+        return ratios
+
+    # Common DeepSeek layout: decoder ratios + mtp ratios in one list.
+    if len(ratios) == num_layers + mtp_num_layers:
+        logger.warning(
+            "compress_ratios has decoder+MTP length (%s), truncating to decoder num_layers (%s).",
+            len(ratios),
+            num_layers,
+        )
+        return ratios[:num_layers]
+
+    if len(ratios) > num_layers:
+        logger.warning(
+            "compress_ratios length (%s) > num_layers (%s), truncating.",
+            len(ratios),
+            num_layers,
+        )
+        return ratios[:num_layers]
+
+    # len(ratios) < num_layers: extend with last ratio (or 0 if empty).
+    pad_value = ratios[-1] if ratios else 0
+    logger.warning(
+        "compress_ratios length (%s) < num_layers (%s), padding with %s.",
+        len(ratios),
+        num_layers,
+        pad_value,
+    )
+    return ratios + [pad_value] * (num_layers - len(ratios))
+
+
+# Plan-2 P17 dedup: the duplicated ``_RMSNorm`` here was retired.
+# Block-level RMSNorm fallbacks now use the canonical
+# :class:`LocalRMSNorm` from
+# ``primus.backends.megatron.core.transformer.local_rmsnorm`` (imported
+# at the top of this file). Spec-driven norms still come from
+# ``DeepSeekV4SpecProvider.v4_norm_module()``.
+
+
+class _DenseSwiGLUMLP(nn.Module):
+    """Plain dense SwiGLU FFN with V4 pre-multiplication clamp.
+
+    Used for non-MoE layers (or as a fallback when ``num_routed_experts``
+    is 0). V4-Flash has a tiny number of dense head/tail layers; the bulk
+    of layers are MoE (see :class:`DeepseekV4MoE`).
+
+    The activation matches V4's released ``Expert.forward``:
+    ``SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha)`` with
+    ``alpha = config.swiglu_limit`` (``0`` disables clamping).
+    """
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        *,
+        pg_collection=None,
+    ) -> None:
+        del pg_collection
+        super().__init__()
+        hidden_size = int(config.hidden_size)
+        ffn_hidden_size = int(config.ffn_hidden_size)
+        self.w_gate = _build_projection(
+            hidden_size,
+            ffn_hidden_size,
+            config=config,
+        )
+        self.w_up = _build_projection(
+            hidden_size,
+            ffn_hidden_size,
+            config=config,
+        )
+        self.w_down = _build_projection(
+            ffn_hidden_size,
+            hidden_size,
+            config=config,
+        )
+        self.swiglu_limit = float(getattr(config, "swiglu_limit", 0.0) or 0.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = _projection_forward(self.w_gate, x)
+        up = _projection_forward(self.w_up, x)
+        if self.swiglu_limit > 0.0:
+            gate = gate.clamp(max=self.swiglu_limit)
+            up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit)
+        return _projection_forward(self.w_down, F.silu(gate) * up)
+
+
+@dataclass
+class DeepseekV4HybridLayerSubmodules(TransformerLayerSubmodules):
+    """Spec tree for one DeepSeek-V4 hybrid layer.
+
+    Plan-2 P15: extends :class:`TransformerLayerSubmodules` so the V4
+    layer slots into Megatron's spec lifecycle without bespoke field
+    plumbing. The four core fields use upstream-canonical names:
+
+    * ``input_layernorm``     (was ``attn_norm``)
+    * ``self_attention``      (was ``attention``)
+    * ``pre_mlp_layernorm``   (was ``ffn_norm``)
+    * ``mlp``                 (was ``ffn``)
+
+    V4 adds two HC mixer hooks that the hybrid forward consumes; both
+    are optional (``None`` when ``hc_mult == 1`` because there is only
+    one stream to mix).
+
+    The cross-attention / BDA fields inherited from the parent stay at
+    their defaults (V4 has no cross-attention and no BDA wrapper — the
+    HC residual path replaces both).
+    """
+
+    attn_hc: Optional[Union[ModuleSpec, type]] = None
+    ffn_hc: Optional[Union[ModuleSpec, type]] = None
+
+
+@dataclass
+class DeepseekV4TransformerBlockSubmodules:
+    """Spec tree for the DeepSeek-V4 decoder block.
+
+    Plan-2 P15: ``hyper_head`` is built **only** on the post_process
+    stage (the final HC collapse `[B, S, K, D] -> [B, S, D]`). On
+    non-final PP stages the block keeps the K-stream form intact and
+    relies on the lift / lower helpers for P2P shape compatibility.
+    """
+
+    layer_specs: Optional[List[ModuleSpec]] = None
+    hyper_head: Optional[Union[ModuleSpec, type]] = None
+    final_layernorm: Optional[Union[ModuleSpec, type]] = None
+
+
+# ---------------------------------------------------------------------------
+# K-stream <-> sequence-axis packing helpers (PP P2P shape carrier)
+# ---------------------------------------------------------------------------
+
+
+def _lift_streams_in(
+    hidden_states: torch.Tensor,
+    *,
+    pre_process: bool,
+    hc_mult: int,
+) -> torch.Tensor:
+    """Reshape the block's input to ``[B, S, K, D]`` for HC math.
+
+    Args:
+        hidden_states: incoming tensor.
+            * On the first PP stage (``pre_process=True``): ``[S, B, D]``
+              (Megatron's sequence-first convention) — we expand to K
+              streams.
+            * On subsequent PP stages (``pre_process=False``):
+              ``[S*K, B, D]`` (K folded into the sequence axis by the
+              previous stage's :func:`_lower_streams_out`) — we unfold.
+        pre_process: ``True`` on the first PP stage.
+        hc_mult: number of HC streams ``K``.
+
+    Returns:
+        ``[B, S, K, D]`` if ``hc_mult > 1`` else ``[B, S, D]``.
+    """
+    if hc_mult <= 1:
+        # Single-stream: just transpose [S, B, D] -> [B, S, D].
+        return hidden_states.transpose(0, 1).contiguous()
+
+    if pre_process:
+        # [S, B, D] -> [B, S, D] -> [B, S, K, D] (broadcast across K).
+        x = hidden_states.transpose(0, 1).contiguous()
+        B, S, D = x.shape
+        return x.unsqueeze(2).expand(B, S, hc_mult, D).contiguous()
+
+    # Non-first stage: [S*K, B, D] -> [B, S, K, D].
+    SK, B, D = hidden_states.shape
+    if SK % hc_mult != 0:
+        raise ValueError(
+            f"PP boundary tensor first-dim {SK} not divisible by hc_mult={hc_mult}; "
+            "previous stage did not pack K via _lower_streams_out."
+        )
+    S = SK // hc_mult
+    # [S*K, B, D] -> [S, K, B, D] -> [B, S, K, D]
+    x = hidden_states.view(S, hc_mult, B, D).permute(2, 0, 1, 3).contiguous()
+    return x
+
+
+def _lower_streams_out(
+    x: torch.Tensor,
+    *,
+    post_process: bool,
+    hc_mult: int,
+) -> torch.Tensor:
+    """Reshape the block's output back to a P2P-compatible 3D tensor.
+
+    Args:
+        x: ``[B, S, K, D]`` (multi-stream) or ``[B, S, D]`` (single).
+            On the final stage callers pass the post-HyperHead
+            ``[B, S, D]``; on non-final stages they pass the
+            still-multi-stream ``[B, S, K, D]``.
+        post_process: ``True`` on the final PP stage.
+        hc_mult: number of HC streams ``K``.
+
+    Returns:
+        ``[S, B, D]`` on the final stage (matching Megatron's
+        sequence-first output convention) or ``[S*K, B, D]`` on
+        non-final stages (K folded into the sequence axis so PP P2P
+        kernels see a 3D tensor of the expected shape).
+    """
+    if hc_mult <= 1:
+        return x.transpose(0, 1).contiguous()
+
+    if post_process:
+        # x is [B, S, D] (already collapsed by HyperHead on this stage).
+        if x.dim() != 3:
+            raise ValueError(
+                f"_lower_streams_out: post_process expects [B, S, D] after HyperHead, "
+                f"got shape {tuple(x.shape)}."
+            )
+        return x.transpose(0, 1).contiguous()
+
+    # Non-final stage: pack [B, S, K, D] -> [S*K, B, D].
+    if x.dim() != 4:
+        raise ValueError(
+            f"_lower_streams_out: non-final stage expects [B, S, K, D], " f"got shape {tuple(x.shape)}."
+        )
+    B, S, K, D = x.shape
+    if K != hc_mult:
+        raise ValueError(f"_lower_streams_out: K dim of input ({K}) does not match hc_mult ({hc_mult}).")
+    # [B, S, K, D] -> [S, K, B, D] -> [S*K, B, D]
+    return x.permute(1, 2, 0, 3).contiguous().view(S * K, B, D)
+
+
+# ---------------------------------------------------------------------------
+# Per-layer attention factory
+# ---------------------------------------------------------------------------
+
+
+def _build_attention(
+    *,
+    compress_ratio: int,
+    rope: DualRoPE,
+    config: Optional[DeepSeekV4TransformerConfig] = None,
+):
+    """No-spec fallback used when the layer is built without an
+    ``attention`` :class:`ModuleSpec`.
+
+    Plan-2 P13: all three V4 layer types (``compress_ratio in {0, 4, 128}``)
+    construct through the single faithful :class:`DeepseekV4Attention`
+    class, which carries its own compressor / indexer (built locally
+    when no spec is provided). Production paths use
+    :func:`get_deepseek_v4_runtime_decoder_spec` which provides full
+    spec submodules; this fallback is for configs / unit tests that
+    construct a layer without an explicit attention spec.
+    """
+    return DeepseekV4Attention(
+        config=config,
+        rope=rope,
+        compress_ratio=int(compress_ratio),
+    )
+
+
+# ---------------------------------------------------------------------------
+# A single V4 block layer (attention sub-block + FFN sub-block, both wrapped
+# by HyperMixer for the K-stream residual)
+# ---------------------------------------------------------------------------
+
+
+class DeepseekV4HybridLayer(TransformerLayer):
+    """One layer of the V4 decoder.
+
+    Holds (using upstream-canonical submodule names):
+
+      * :attr:`input_layernorm`     — pre-attention RMSNorm.
+      * :attr:`self_attention`      — V4 attention (Dense / HCA / CSA, picked
+        from ``compress_ratio``).
+      * :attr:`pre_mlp_layernorm`   — pre-MLP RMSNorm.
+      * :attr:`mlp`                 — V4 MoE (when ``num_moe_experts > 0``)
+        or :class:`_DenseSwiGLUMLP` fallback for non-MoE layers.
+      * :attr:`attn_hc` / :attr:`ffn_hc`  — :class:`HyperMixer` instances
+        per sub-block when ``hc_mult > 1`` (else ``None`` and the residual
+        collapses to a vanilla ``x + sub(x)`` add).
+
+    Plan-2 P15 inherits from :class:`TransformerLayer` for type identity
+    (so Megatron's ``isinstance(layer, TransformerLayer)`` checks work
+    and ``BaseTransformerLayer`` -derived utilities apply). The parent
+    ``__init__`` is *not* called because V4's submodule set differs from
+    upstream (no cross-attention, no BDA, V4-specific attention
+    signature) — we initialize via :class:`MegatronModule` directly.
+    """
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        *,
+        layer_idx: int,
+        compress_ratio: int,
+        rope: Optional[DualRoPE] = None,
+        pg_collection=None,
+        submodules: Optional[DeepseekV4HybridLayerSubmodules] = None,
+        layer_number: Optional[int] = None,
+        is_mtp_layer: bool = False,
+        vp_stage: Optional[int] = None,
+        **_mtp_build_kwargs,
+    ) -> None:
+        # Bypass TransformerLayer.__init__ (it expects the upstream
+        # submodule contract — cross-attn, BDA, and a self_attention
+        # signature that takes layer_number / cp_comm_type). We build
+        # the V4 attribute set directly via MegatronModule.
+        #
+        # ``is_mtp_layer`` / ``vp_stage`` / ``**_mtp_build_kwargs`` are
+        # accepted (and ignored) so this layer can be built by the
+        # upstream :class:`MultiTokenPredictionLayer`, whose
+        # ``build_module(self.submodules.mtp_model_layer, config=...,
+        # vp_stage=..., layer_number=..., is_mtp_layer=True)`` call passes
+        # a richer kwarg set than the main decoder block does.
+        del is_mtp_layer, vp_stage, _mtp_build_kwargs
+        MegatronModule.__init__(self, config=config)
+
+        # ``rope`` is normally injected by ``DeepseekV4TransformerBlock``
+        # (one shared DualRoPE for the whole stack). When this layer is
+        # built standalone — e.g. as the inner layer of a V4 MTP depth via
+        # the upstream MTP ``build_module`` path, which does NOT thread
+        # ``rope`` — fall back to constructing a private DualRoPE from the
+        # config. DualRoPE holds only RoPE cos/sin caches (no trainable
+        # params), so a per-MTP-depth instance is numerically a no-op vs.
+        # sharing the trunk's.
+        if rope is None:
+            rope = DualRoPE(
+                rotary_dim=int(config.qk_pos_emb_head_dim),
+                rope_theta=float(config.rotary_base),
+                compress_rope_theta=float(config.compress_rope_theta),
+                yarn_factor=float(config.rotary_scaling_factor),
+                yarn_beta_fast=32.0,
+                yarn_beta_slow=1.0,
+                original_max_position_embeddings=int(config.original_max_position_embeddings),
+            )
+            # Register the privately-built RoPE as a real submodule so its
+            # (non-persistent) cos/sin ``inv_freq`` buffers follow the layer
+            # onto the GPU under ``module.to(device)`` / ``.cuda()``. The main
+            # decoder path instead receives the block-owned ``self.rope`` (a
+            # registered submodule of the block) and keeps it in the
+            # attention's ``self._rope`` *list* to avoid double-registration;
+            # here there is no other owner, so this layer must register it.
+            self._fallback_rope = rope
+
+        self.layer_idx = int(layer_idx)
+        self.compress_ratio = int(compress_ratio)
+        self.hc_mult = int(config.hc_mult)
+        # 1-based layer_number for Megatron's sharded_state_dict + recompute.
+        self.layer_number = int(layer_number) if layer_number is not None else (self.layer_idx + 1)
+
+        hidden_size = int(config.hidden_size)
+        norm_eps = float(config.norm_epsilon)
+        hc_eps = float(config.hc_eps)
+        hc_sinkhorn_iters = int(config.hc_sinkhorn_iters)
+        # Plan-5 P29 (RESCOPED): forward the V4 config flag to HyperMixer
+        # so its sinkhorn_normalize dispatches to the torch.compile fast
+        # path when the run script flips USE_V4_COMPILED_SINKHORN=True.
+        # Default is False — see deepseek_v4_transformer_config.py for
+        # the dispatch contract.
+        use_v4_compiled_sinkhorn = bool(getattr(config, "use_v4_compiled_sinkhorn", False))
+
+        use_spec_submodules = submodules is not None
+        submodules = submodules or DeepseekV4HybridLayerSubmodules()
+        # Cache for the parent (TransformerLayer) sharded_state_dict / recompute paths.
+        self.submodules_config = submodules
+
+        if use_spec_submodules and submodules.input_layernorm is not None:
+            self.input_layernorm = build_module(
+                submodules.input_layernorm,
+                config=config,
+                hidden_size=hidden_size,
+                eps=norm_eps,
+            )
+        else:
+            self.input_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps)
+
+        if use_spec_submodules and submodules.self_attention is not None:
+            self.self_attention = build_module(
+                submodules.self_attention,
+                config=config,
+                rope=rope,
+            )
+        else:
+            self.self_attention = _build_attention(
+                compress_ratio=self.compress_ratio,
+                rope=rope,
+                config=config,
+            )
+
+        if use_spec_submodules and submodules.pre_mlp_layernorm is not None:
+            self.pre_mlp_layernorm = build_module(
+                submodules.pre_mlp_layernorm,
+                config=config,
+                hidden_size=hidden_size,
+                eps=norm_eps,
+            )
+        else:
+            self.pre_mlp_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps)
+
+        self.is_moe = int(config.num_moe_experts) > 0
+        if use_spec_submodules and submodules.mlp is not None:
+            self.mlp = build_module(
+                submodules.mlp,
+                config=config,
+                pg_collection=pg_collection,
+            )
+            self.is_moe = isinstance(self.mlp, DeepseekV4MoE)
+        elif self.is_moe:
+            moe_use_grouped_gemm = bool(config.moe_grouped_gemm)
+            moe_use_legacy_grouped_gemm = bool(config.moe_use_legacy_grouped_gemm)
+
+            provider = resolve_v4_provider(config)
+            grouped_mlp_module, grouped_mlp_submodules = provider.v4_grouped_mlp_modules(
+                moe_use_grouped_gemm=moe_use_grouped_gemm,
+                moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+            )
+            assert (
+                grouped_mlp_module is not None
+            ), "DeepSeek-V4 grouped MLP module must be provided by DeepSeekV4SpecProvider."
+
+            shared_expert_spec = ModuleSpec(
+                module=SharedExpertMLP,
+                submodules=MLPSubmodules(
+                    linear_fc1=provider.column_parallel_linear(),
+                    linear_fc2=provider.row_parallel_linear(),
+                    # P18 D2: pull from the V4-aware helper so the slot
+                    # is None when ``use_te_activation_func`` is False
+                    # (the default in V4 yamls; clamped-SwiGLU lives on
+                    # the eager path via ``config.activation_func``).
+                    activation_func=provider.v4_mlp_activation_func(),
+                ),
+            )
+            moe_submodules = DeepseekV4MoESubmodules(
+                hash_router=ModuleSpec(module=DeepseekV4HashRouter),
+                learned_router=ModuleSpec(module=DeepseekV4LearnedRouter),
+                grouped_experts=ModuleSpec(
+                    module=grouped_mlp_module,
+                    submodules=grouped_mlp_submodules,
+                ),
+                shared_expert=shared_expert_spec,
+            )
+
+            self.mlp = DeepseekV4MoE(
+                config=config,
+                layer_idx=self.layer_idx,
+                pg_collection=pg_collection,
+                submodules=moe_submodules,
+            )
+        else:
+            self.mlp = _DenseSwiGLUMLP(
+                config=config,
+            )
+
+        if self.hc_mult > 1:
+            if use_spec_submodules and submodules.attn_hc is not None:
+                self.attn_hc = build_module(
+                    submodules.attn_hc,
+                    hidden_size=hidden_size,
+                    hc_mult=self.hc_mult,
+                    eps=hc_eps,
+                    sinkhorn_iters=hc_sinkhorn_iters,
+                    use_compiled_sinkhorn=use_v4_compiled_sinkhorn,
+                )
+            else:
+                self.attn_hc = HyperMixer(
+                    hidden_size=hidden_size,
+                    hc_mult=self.hc_mult,
+                    eps=hc_eps,
+                    sinkhorn_iters=hc_sinkhorn_iters,
+                    use_compiled_sinkhorn=use_v4_compiled_sinkhorn,
+                )
+            if use_spec_submodules and submodules.ffn_hc is not None:
+                self.ffn_hc = build_module(
+                    submodules.ffn_hc,
+                    hidden_size=hidden_size,
+                    hc_mult=self.hc_mult,
+                    eps=hc_eps,
+                    sinkhorn_iters=hc_sinkhorn_iters,
+                    use_compiled_sinkhorn=use_v4_compiled_sinkhorn,
+                )
+            else:
+                self.ffn_hc = HyperMixer(
+                    hidden_size=hidden_size,
+                    hc_mult=self.hc_mult,
+                    eps=hc_eps,
+                    sinkhorn_iters=hc_sinkhorn_iters,
+                    use_compiled_sinkhorn=use_v4_compiled_sinkhorn,
+                )
+        else:
+            self.attn_hc = None
+            self.ffn_hc = None
+
+    # ------------------------------------------------------------------
+
+    def _hc_apply(self, mixer: Optional[HyperMixer], x: torch.Tensor, sub_block, *args):
+        """Run a sub-block under HC.
+
+        ``x`` shape: ``[B, S, K, D]`` if ``hc_mult > 1``, else ``[B, S, D]``.
+        ``sub_block`` is ``Callable[[Tensor, *Any], Tensor]`` whose first
+        positional arg is the (collapsed) hidden in ``[B, S, D]``.
+        """
+        if mixer is None:
+            # Single-stream: classic residual; x already has shape [B, S, D].
+            out = sub_block(x, *args)
+            return x + out
+
+        pre, post, comb = mixer.compute_weights(x)  # [..., K], [..., K], [..., K, K]
+        collapsed = HyperMixer.collapse(x, pre)  # [B, S, D]
+        out = sub_block(collapsed, *args)  # sub-block first positional = collapsed
+        return HyperMixer.expand(x, out, post, comb)  # [B, S, K, D]
+
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        *,
+        position_ids: Optional[torch.Tensor] = None,
+        token_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        """Run one V4 layer.
+
+        Args:
+            hidden_states: ``[B, S, K, D]`` (multi-stream) or
+                ``[B, S, D]`` (single-stream). The block calls us with
+                the K-stream form when ``hc_mult > 1`` and with the
+                collapsed form otherwise.
+            attention_mask: ignored — V4 manages its own SWA / sink mask
+                inside :class:`DeepseekV4Attention`. Accepted for
+                upstream :class:`TransformerLayer` API compatibility.
+            position_ids: ``[B, S]`` or ``[S]``. Forwarded to attention.
+            token_ids: ``[B, S]`` integer tensor; required when this is
+                a hash-routed MoE layer
+                (``layer_idx < num_hash_layers``). Ignored for non-MoE /
+                non-hash layers.
+            **kwargs: ignored — accepted so this layer can be placed
+                inside upstream :class:`MultiTokenPredictionLayer`,
+                which forwards a richer kwargs set (rotary buffers,
+                inference params, etc.).
+
+        Returns:
+            ``(hidden_states, context)`` where ``context`` is always
+            ``None``. The upstream :class:`TransformerLayer` returns a
+            tuple of this shape (cross-attention context pass-through);
+            V4 has no cross-attention, so the second element is always
+            ``None``. Returning a tuple keeps V4 layers compatible with
+            :class:`MultiTokenPredictionLayer._proj_and_transformer_layer`
+            which unpacks ``hidden_states, _ = self.mtp_model_layer(...)``.
+        """
+        del attention_mask, kwargs
+
+        if position_ids is None:
+            # Tiny CPU smokes / unit tests may omit position_ids; fall
+            # back to the seq-only arange. The block always provides
+            # them in production.
+            S = hidden_states.shape[1]
+            position_ids = torch.arange(S, device=hidden_states.device)
+
+        # Attention sub-block. The collapse passes a [B, S, D] hidden, then
+        # the attention runs and returns [B, S, D]; HC expand writes back.
+        def _attn_sub(collapsed: torch.Tensor) -> torch.Tensor:
+            return self.self_attention(self.input_layernorm(collapsed), position_ids)
+
+        x = self._hc_apply(self.attn_hc, hidden_states, _attn_sub)
+
+        # MLP / MoE sub-block. MoE needs token_ids when the layer is
+        # hash-routed; plain SwiGLU ignores it.
+        if self.is_moe:
+
+            def _ffn_sub(collapsed: torch.Tensor) -> torch.Tensor:
+                return self.mlp(self.pre_mlp_layernorm(collapsed), token_ids=token_ids)
+
+        else:
+
+            def _ffn_sub(collapsed: torch.Tensor) -> torch.Tensor:
+                return self.mlp(self.pre_mlp_layernorm(collapsed))
+
+        x = self._hc_apply(self.ffn_hc, x, _ffn_sub)
+        return x, None
+
+
+# ---------------------------------------------------------------------------
+# Top-level V4 transformer block
+# ---------------------------------------------------------------------------
+
+
+class DeepseekV4TransformerBlock(TransformerBlock):
+    """Multi-stream HC decoder for DeepSeek-V4.
+
+    Plan-2 P15 subclasses Megatron's
+    :class:`megatron.core.transformer.transformer_block.TransformerBlock`
+    for type identity (so any upstream ``isinstance(block,
+    TransformerBlock)`` checks light up) and to inherit its
+    sharded-state-dict / debug surface. The parent ``__init__`` is
+    bypassed because V4's submodule contract differs (HyperHead on
+    post-process only, K-stream lift / lower at PP boundaries, no
+    upstream-style layer-norm impl override). We initialize via
+    :class:`MegatronModule` directly.
+
+    PP K-stream packing: between stages we send ``[S*K, B, D]`` (K folded
+    into the sequence axis), letting standard 3D PP P2P kernels carry
+    the multi-stream tensor unchanged. The first stage lifts ``[S, B, D]``
+    to ``[B, S, K, D]``; the final stage collapses with HyperHead and
+    transposes back to ``[S, B, D]``.
+
+    Phase 6 update:
+    - respects PP / VP layer partitioning by constructing only local layers
+      for this pipeline rank;
+    - supports ``set_input_tensor`` so non-first PP stages consume P2P input.
+    """
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        spec=None,
+        post_layer_norm: bool = True,
+        pre_process: bool = True,
+        post_process: bool = True,
+        pg_collection=None,
+        vp_stage=None,
+        submodules: Optional[DeepseekV4TransformerBlockSubmodules] = None,
+    ) -> None:
+        # Bypass TransformerBlock.__init__: it requires a real
+        # pg_collection (or pulls one from parallel_state) and runs
+        # upstream-specific layer construction. V4's lift / lower path
+        # plus the spec provider give us equivalent functionality with
+        # CPU instantiability.
+        MegatronModule.__init__(self, config=config)
+        # Save arguments matching the parent's interface for compatibility.
+        self.spec = spec
+        self.submodules = submodules
+        self.post_layer_norm = post_layer_norm
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.vp_stage = vp_stage
+        self.pg_collection = pg_collection
+        # Required by pipeline schedules (same contract as TransformerBlock).
+        self.input_tensor = None
+        logger.info(
+            "[DeepSeek-V4] decoder block initialized (pre_process=%s post_process=%s).",
+            pre_process,
+            post_process,
+        )
+
+        # ---- shape / model fields ----
+        hidden_size = config.hidden_size
+        rotary_dim = config.qk_pos_emb_head_dim
+        num_layers = config.num_layers
+        norm_eps = config.norm_epsilon
+
+        # ---- V4-specific fields ----
+        hc_mult = config.hc_mult
+        hc_eps = config.hc_eps
+        config.hc_sinkhorn_iters
+        compress_ratios = _normalize_compress_ratios(
+            config.compress_ratios,
+            num_layers=num_layers,
+            mtp_num_layers=int(config.mtp_num_layers),
+        )
+        self.compress_ratios: List[int] = compress_ratios
+
+        rope_theta = config.rotary_base
+        compress_rope_theta = config.compress_rope_theta
+        yarn_factor = config.rotary_scaling_factor
+        original_max_pos = config.original_max_position_embeddings
+
+        self.num_hash_layers = int(config.num_hash_layers)
+
+        # ---- shared dual-RoPE for the whole stack ----
+        self.rope = DualRoPE(
+            rotary_dim=rotary_dim,
+            rope_theta=rope_theta,
+            compress_rope_theta=compress_rope_theta,
+            yarn_factor=yarn_factor,
+            yarn_beta_fast=32.0,
+            yarn_beta_slow=1.0,
+            original_max_position_embeddings=original_max_pos,
+        )
+
+        # ---- stage-local layer specs (always provided by runtime spec) ----
+        provided_layer_specs = submodules.layer_specs if submodules is not None else None
+        assert provided_layer_specs, "DeepSeek-V4 requires non-empty submodules.layer_specs."
+        self.layers = nn.ModuleList()
+        self.global_layer_indices = []
+        for local_idx, layer_spec in enumerate(provided_layer_specs):
+            layer = build_module(
+                layer_spec,
+                config=config,
+                pg_collection=pg_collection,
+                rope=self.rope,
+            )
+            self.layers.append(layer)
+            self.global_layer_indices.append(int(getattr(layer, "layer_idx", local_idx)))
+        self.layer_offset = self.global_layer_indices[0] if self.global_layer_indices else 0
+        self.hc_mult = hc_mult
+
+        # Final HC collapse: built only on the post_process stage.
+        # Earlier PP stages forward the K-stream form via
+        # ``_lower_streams_out`` (no HyperHead per stage).
+        if hc_mult > 1 and self.post_process:
+            if submodules is not None and submodules.hyper_head is not None:
+                self.hyper_head = build_module(
+                    submodules.hyper_head,
+                    hidden_size=hidden_size,
+                    hc_mult=hc_mult,
+                    eps=hc_eps,
+                )
+            else:
+                self.hyper_head = HyperHead(hidden_size=hidden_size, hc_mult=hc_mult, eps=hc_eps)
+        else:
+            self.hyper_head = None
+
+        # Final RMSNorm placement follows Megatron semantics:
+        # - no MTP: on post_process stage
+        # - with MTP: on the stage containing decoder's final layer
+        if self._has_final_layernorm_in_this_stage(total_decoder_layers=num_layers):
+            if submodules is not None and submodules.final_layernorm is not None:
+                self.final_layernorm = build_module(
+                    submodules.final_layernorm,
+                    config=self.config,
+                    hidden_size=hidden_size,
+                    eps=norm_eps,
+                )
+            else:
+                self.final_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps)
+        else:
+            self.final_layernorm = None
+
+    # ------------------------------------------------------------------
+
+    def _has_final_layernorm_in_this_stage(self, *, total_decoder_layers: int) -> bool:
+        if not self.post_layer_norm:
+            return False
+
+        mtp_num_layers = self.config.mtp_num_layers
+        if mtp_num_layers is None:
+            return self.post_process
+
+        if not self.global_layer_indices:
+            return False
+        return self.global_layer_indices[-1] == (total_decoder_layers - 1)
+
+    def set_input_tensor(self, input_tensor: torch.Tensor):
+        """Pipeline-parallel hook: stash tensor from previous PP stage."""
+        self.input_tensor = input_tensor
+
+    @property
+    def num_layers_per_pipeline_rank(self) -> int:
+        """Compatibility shim used by upstream debug / recompute code."""
+        return len(self.layers)
+
+    # ------------------------------------------------------------------
+
+    def _recompute_local_layer_indices(self) -> Optional[set]:
+        """Local indices of layers to activation-checkpoint on this stage.
+
+        Returns ``None`` when no recompute should be applied (the caller
+        then runs the plain forward loop). Mirrors Megatron semantics:
+
+        * only active when ``recompute_granularity == 'full'`` and training;
+        * ``recompute_method == 'block'``   -> first ``recompute_num_layers``
+          layers of this PP stage;
+        * ``recompute_method == 'uniform'`` -> every layer on this stage
+          (chunk size ``recompute_num_layers``; for full-layer checkpointing
+          the chunk boundary does not change which layers are recomputed);
+        * Primus ``recompute_layer_ids`` (explicit GLOBAL indices) takes
+          precedence when set.
+        """
+        if not self.training:
+            return None
+
+        cfg = self.config
+        if getattr(cfg, "recompute_granularity", None) != "full":
+            return None
+
+        n_local = len(self.layers)
+        if n_local == 0:
+            return None
+
+        recompute_layer_ids = getattr(cfg, "recompute_layer_ids", None)
+        if recompute_layer_ids:
+            wanted = {int(i) for i in recompute_layer_ids}
+            local = {
+                local_idx
+                for local_idx, global_idx in enumerate(self.global_layer_indices)
+                if int(global_idx) in wanted
+            }
+            return local or None
+
+        num = int(getattr(cfg, "recompute_num_layers", 0) or 0)
+        if num <= 0:
+            return None
+
+        method = getattr(cfg, "recompute_method", None) or "block"
+        if method == "block":
+            return set(range(min(num, n_local)))
+        if method == "uniform":
+            return set(range(n_local))
+        raise ValueError(f"Invalid recompute_method for DeepSeek-V4: {method!r}")
+
+    def _layer_fp8_context(self, global_idx: int):
+        """FP8 autocast for a single V4 layer (no-op when fp8 is off).
+
+        DeepseekV4TransformerBlock.forward overrides Megatron's
+        TransformerBlock.forward with a custom layer loop, so the per-layer
+        ``get_fp8_context`` wrapping that Megatron's stock forward uses to
+        activate fp8 — and, via Primus's ``fp8_patches``, Primus-Turbo fp8 —
+        is not inherited and must be re-applied here. Without it the
+        ``primus_turbo_fp8_autocast`` context is never entered, so
+        ``PRIMUS_TURBO_FP8_ENABLED`` stays False and every GEMM silently runs
+        bf16 regardless of ``--fp8`` / ``--fp8_recipe``.
+
+        ``get_fp8_context`` is resolved as a module attribute at call time so
+        the ``before_train`` fp8 patch (which rebinds
+        ``megatron.core.fp8_utils.get_fp8_context``) is honored. ``global_idx``
+        is the 0-based global layer index (matches the function's ``layer_no``
+        contract / ``is_first_last_bf16_layer`` gating).
+        """
+        # FP4 (mxfp4) path: enter the Primus-Turbo fp4 autocast so is_turbo_fp4_enabled()
+        # is set for the layer (drives the grouped-MLP MXFP4 expert path and the
+        # fp4 turbo linears). fp4 and fp8 are mutually exclusive (global recipe).
+        if getattr(self.config, "fp4", None):
+            from megatron.core import fp4_utils
+
+            return fp4_utils.get_fp4_context(self.config, global_idx)
+        if not self.config.fp8:
+            return nullcontext()
+        from megatron.core import fp8_utils
+
+        return fp8_utils.get_fp8_context(self.config, global_idx)
+
+    def _forward_layer_checkpointed(self, layer, x, position_ids, token_ids, global_idx):
+        """Run one V4 layer under activation checkpointing.
+
+        Only the hidden-state tensor ``x`` is passed as the checkpointed
+        arg (``args[0]``); ``position_ids`` / ``token_ids`` are integer
+        tensors that need no grad and are captured by closure. The closure
+        returns just the hidden tensor (V4 layers return ``(hidden, None)``)
+        so the checkpoint backward sees a single grad-requiring output.
+
+        The fp8 context is entered *inside* the checkpointed closure so it is
+        re-established on recompute (backward), keeping the recomputed forward
+        in fp8 to match the original pass.
+        """
+        from megatron.core import tensor_parallel
+
+        def _run(hidden):
+            with self._layer_fp8_context(global_idx):
+                out, _ = layer(hidden, position_ids=position_ids, token_ids=token_ids)
+            return out
+
+        return tensor_parallel.checkpoint(
+            _run,
+            self.config.distribute_saved_activations,
+            x,
+        )
+
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        inference_context=None,
+        rotary_pos_emb=None,
+        rotary_pos_cos=None,
+        rotary_pos_sin=None,
+        rotary_pos_cos_sin=None,
+        packed_seq_params=None,
+        sequence_len_offset=None,
+        position_ids: Optional[torch.Tensor] = None,
+        token_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run the V4 decoder.
+
+        Shape contract:
+
+        * On the first PP stage (``pre_process=True``):
+          ``hidden_states`` arrives as ``[S, B, D]`` (sequence-first
+          embedding output). The lift helper expands it to
+          ``[B, S, K, D]`` for HC math.
+        * On subsequent PP stages (``pre_process=False``): the previous
+          stage packed K into the sequence axis, so
+          ``hidden_states`` arrives as ``[S*K, B, D]``. The lift helper
+          unfolds back to ``[B, S, K, D]``.
+        * On the final stage (``post_process=True``): HyperHead
+          collapses ``[B, S, K, D] -> [B, S, D]``, then the lower helper
+          transposes to the sequence-first ``[S, B, D]`` output.
+        * On non-final stages: the lower helper packs
+          ``[B, S, K, D] -> [S*K, B, D]`` so PP P2P kernels see a 3D
+          tensor of the expected rank.
+
+        ``attention_mask`` and the various ``rotary_pos_*`` kwargs are
+        ignored — V4 manages its own dual-RoPE and SWA mask internally.
+
+        ``position_ids`` is the caller-provided token-position tensor.
+        When omitted (e.g. unit tests), we fall back to ``arange(S)``;
+        production callers (:class:`DeepseekV4Model.forward`) always pass
+        it explicitly.
+
+        ``token_ids`` (``[B, S]`` long tensor) is required only when one
+        or more layers on this stage run hash routing
+        (``layer_idx < num_hash_layers``). The legacy
+        ``decoder._v4_token_ids`` attribute stash has been removed; the
+        model forwards ``input_ids`` here directly.
+        """
+        del (
+            inference_context,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            rotary_pos_cos_sin,
+            packed_seq_params,
+            sequence_len_offset,
+            attention_mask,
+            kwargs,
+        )
+
+        if not self.pre_process:
+            hidden_states = self.input_tensor if self.input_tensor is not None else hidden_states
+        if hidden_states is None:
+            raise ValueError("DeepseekV4TransformerBlock.forward received no hidden_states tensor")
+
+        needs_hash_token_ids = self.num_hash_layers > 0 and any(
+            layer_idx < self.num_hash_layers for layer_idx in self.global_layer_indices
+        )
+        if needs_hash_token_ids and token_ids is None:
+            raise ValueError(
+                "token_ids is required on this PP stage because it owns hash-routed MoE layers "
+                f"(global layer idx < num_hash_layers={self.num_hash_layers})."
+            )
+
+        # Lift incoming P2P tensor to the K-stream form.
+        x = _lift_streams_in(
+            hidden_states,
+            pre_process=self.pre_process,
+            hc_mult=self.hc_mult,
+        )
+        # x is [B, S, K, D] when hc_mult > 1, else [B, S, D].
+        seq_len = x.shape[1]
+
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=x.device)
+
+        # Run the layers. Each V4 layer returns ``(hidden, None)`` for
+        # upstream-tuple compatibility (see DeepseekV4HybridLayer.forward).
+        recompute_local = self._recompute_local_layer_indices()
+        for local_idx, layer in enumerate(self.layers):
+            global_idx = self.global_layer_indices[local_idx]
+            if recompute_local is not None and local_idx in recompute_local:
+                x = self._forward_layer_checkpointed(layer, x, position_ids, token_ids, global_idx)
+            else:
+                with self._layer_fp8_context(global_idx):
+                    x, _ = layer(
+                        x,
+                        position_ids=position_ids,
+                        token_ids=token_ids,
+                    )
+
+        # Final HC collapse on post_process stage; non-final stages
+        # forward the multi-stream form through PP P2P.
+        if self.post_process and self.hc_mult > 1 and self.hyper_head is not None:
+            x = self.hyper_head(x)  # [B, S, D]
+
+        if self.final_layernorm is not None:
+            x = self.final_layernorm(x)
+
+        if not self.pre_process and len(self.layers) == 0 and self.final_layernorm is None:
+            x = x.clone()
+
+        # Lower to a P2P-compatible 3D tensor.
+        out = _lower_streams_out(
+            x,
+            post_process=self.post_process,
+            hc_mult=self.hc_mult,
+        )
+        return make_viewless_tensor(inp=out, requires_grad=out.requires_grad, keep_graph=True)
+
+
+__all__ = [
+    "DeepseekV4HybridLayerSubmodules",
+    "DeepseekV4TransformerBlockSubmodules",
+    "DeepseekV4HybridLayer",
+    "DeepseekV4TransformerBlock",
+    "_lift_streams_in",
+    "_lower_streams_out",
+]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py
new file mode 100644
index 000000000..6041f755a
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py
@@ -0,0 +1,301 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 model builder + model_provider entry points.
+
+In the upstream Megatron-LM convention there are two pieces:
+
+* ``model_provider``: a thin wrapper that handles ``args.record_memory_history``,
+  ModelOpt etc. and delegates to a ``model_builder`` callable. Defined once
+  in ``Megatron-LM/model_provider.py``.
+* ``<modeltype>_builders.py``: contains ``<modeltype>_builder(args, ...)`` --
+  the actual model-class instantiation logic.
+
+For DeepSeek-V4 we keep both in a single primus-owned module so the dispatch
+in ``primus/core/utils/import_utils.py`` doesn't have to chase symbols across
+``third_party/Megatron-LM``.
+
+Phase 8 contract:
+- Resolve V4 runtime decoder spec externally in builder.
+- Pass runtime spec as ``transformer_layer_spec`` into ``DeepseekV4Model``.
+- DeepseekV4Model is rooted at ``LanguageModule`` and has no GPT placeholder
+  spec dependency.
+"""
+
+from typing import Optional
+
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training import get_args, print_rank_0
+from megatron.training.arguments import core_transformer_config_from_args
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import (
+    get_deepseek_v4_runtime_decoder_spec,
+    is_v4_turbo_deepep_active,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_model import (
+    DeepseekV4Model,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+
+
+def _resolve_runtime_decoder_spec(
+    args,
+    config: DeepSeekV4TransformerConfig,
+    vp_stage,
+):
+    """Resolve effective runtime decoder spec for DeepSeek-V4 decoder path."""
+    if args.spec is not None:
+        return import_module(args.spec)
+    return get_deepseek_v4_runtime_decoder_spec(config=config, vp_stage=vp_stage)
+
+
+def _maybe_plumb_v4_sink_attention_args(args) -> None:
+    """Plan-3 P22: derive Turbo sink-attention args from V4 config.
+
+    :class:`PrimusTurboAttention` reads ``use_sink_attention``,
+    ``sink_sliding_window`` and ``sink_window_even_layers_only`` directly
+    from the global ``args`` namespace at module-init time.  For V4 we
+    can derive all three from the V4 attention configuration:
+
+    * ``use_sink_attention`` follows ``args.attn_sink`` (V4's per-head
+      learnable softmax sink — set to ``true`` in the V4-Flash recipe).
+    * ``sink_sliding_window`` is derived from ``args.attn_sliding_window``,
+      *capped to seq_length*.  V4's dense layers use SWA with window 128
+      in the released checkpoint; the released aiter Triton flash-attn
+      backend (``aiter/ops/triton/attention/mha.py``) does not yet
+      accept ``window_size != (-1, -1)`` and raises ``ValueError:
+      Sliding Window is not supported yet in the Triton Backend``.
+      We therefore zero out the window when it equals or exceeds
+      ``seq_length`` (a window that covers the full causal triangle is
+      mathematically equivalent to no window).  When the window is
+      strictly shorter than the sequence length we *still* zero it out
+      and emit a warning — Turbo can't honor it, and falling back to
+      eager-Python every step would defeat the perf goal.  Long-context
+      V4-Flash (seq > 128) needs an aiter SWA upgrade before Turbo can
+      claim full V4 fidelity; tracked as a P22 follow-up.
+    * ``sink_window_even_layers_only`` is hard-set to ``False`` because
+      V4 applies SWA on every dense layer (unlike gpt-oss, which only
+      windows even-numbered layers).
+
+    The plumbing only fires when all of (a) Turbo is enabled, (b)
+    ``use_turbo_attention=True``, (c) the user has not set the sink
+    fields explicitly.  This keeps the V4-Flash YAML free of
+    Turbo-internal knobs while still producing a correct attention
+    forward when the user flips ``use_turbo_attention=true``.
+    """
+    if not getattr(args, "enable_primus_turbo", False):
+        return
+    if not getattr(args, "use_turbo_attention", False):
+        return
+    if not getattr(args, "attn_sink", False):
+        # Without the V4 sink, Turbo would not honor SWA either; defer
+        # to ``DeepseekV4Attention``'s eager-Python fallback for this
+        # configuration (it asserts ``_use_core_attention=False`` when
+        # ``attn_sink`` is off and ``attn_sliding_window > 0``).
+        return
+
+    if getattr(args, "use_sink_attention", None) in (None, False):
+        args.use_sink_attention = True
+        print_rank_0(
+            "[Primus:DeepSeek-V4][P22] derived args.use_sink_attention=True "
+            "from args.attn_sink=True (Turbo flash-attn sink path)."
+        )
+
+    if getattr(args, "sink_sliding_window", None) in (None, 0):
+        attn_sw = int(getattr(args, "attn_sliding_window", 0) or 0)
+        seq_len = int(getattr(args, "seq_length", 0) or 0)
+        if attn_sw <= 0:
+            args.sink_sliding_window = 0
+        elif seq_len > 0 and attn_sw >= seq_len:
+            # Window covers the full causal triangle — equivalent to no
+            # window.  Drop it so Turbo's flash kernel runs without
+            # window_size (avoids the aiter Triton SWA gap).
+            args.sink_sliding_window = 0
+            print_rank_0(
+                f"[Primus:DeepSeek-V4][P22] attn_sliding_window={attn_sw} "
+                f">= seq_length={seq_len}; dropping window for Turbo "
+                "(mathematically equivalent to full causal attention)."
+            )
+        else:
+            args.sink_sliding_window = 0
+            print_rank_0(
+                f"[Primus:DeepSeek-V4][P22] WARNING: "
+                f"attn_sliding_window={attn_sw} < seq_length={seq_len} "
+                "but aiter Triton flash-attn does not support sliding "
+                "window yet (raises ValueError).  Setting "
+                "sink_sliding_window=0 — V4 dense layers will attend "
+                "to *all* causal-prior tokens instead of the windowed "
+                "subset.  This deviates from V4-Flash math; for "
+                "checkpoint-fidelity training, keep "
+                "use_turbo_attention=False until aiter adds SWA support."
+            )
+
+    # gpt-oss windows only even layers; V4 windows all dense layers.
+    args.sink_window_even_layers_only = False
+
+
+def _maybe_plumb_v4_turbo_deepep_args(args) -> None:
+    """Plan-3 P23: derive Turbo DeepEP MoE args from V4 config.
+
+    When :func:`is_v4_turbo_deepep_active` returns True, the V4 MoE
+    layers will be built with :class:`PrimusTurboDeepEPTokenDispatcher`
+    (resolved by ``_pick_v4_dispatcher_cls`` in
+    ``deepseek_v4_layer_specs.py``).  Two ``args`` fields control
+    downstream behaviour:
+
+    * ``moe_enable_deepep`` — asserted by
+      :class:`PrimusTurboDeepEPTokenDispatcher.__init__`; without it
+      the Turbo dispatcher refuses to build.
+    * ``moe_token_dispatcher_type`` — read by
+      :class:`DeepseekV4TransformerConfig` to populate the V4 config's
+      ``moe_token_dispatcher_type``.  V4 spec build's
+      ``_pick_v4_dispatcher_cls`` then sees ``"flex"`` and pairs the
+      Turbo class with the right dispatcher-type label so
+      ``DeepseekV4MoE._resolve_dispatcher_type_from_spec`` reports
+      ``"flex"`` (instead of falling through to the
+      ``"unsupported dispatcher module"`` warning + ``"alltoall"``
+      fallback).
+
+    The ``before_train`` Turbo MoE patch
+    (``primus.backends.megatron.patches.turbo.moe_dispatcher_patches``)
+    sets the same two fields, but it fires AFTER ``deepseek_v4_builder``
+    has already built ``config`` from ``args``, so V4 needs to plumb
+    these values BEFORE ``core_transformer_config_from_args`` runs.
+
+    The plumbing is a no-op when:
+
+    * Turbo is not enabled / DeepEP not requested / TP > 1, OR
+    * the user has already set ``args.moe_token_dispatcher_type`` to
+      something other than ``"alltoall"`` (we only override the V4
+      base.yaml default, never an explicit user opt-in like
+      ``"allgather"``).
+    """
+    if not is_v4_turbo_deepep_active(args):
+        return
+
+    if not bool(getattr(args, "moe_enable_deepep", False)):
+        args.moe_enable_deepep = True
+        print_rank_0(
+            "[Primus:DeepSeek-V4][P23] derived args.moe_enable_deepep=True "
+            "from args.use_turbo_deepep=True (Turbo DeepEP dispatcher path)."
+        )
+
+    current = str(getattr(args, "moe_token_dispatcher_type", "") or "").lower()
+    # Only override when the V4 base default ("alltoall") is in effect
+    # or the user already opted into "flex".  Respect explicit
+    # "allgather" — that is a different parallelism scheme (tp+ep
+    # all-gather/scatter), not a deepep flavour.
+    if current in ("", "alltoall", "flex"):
+        if current != "flex":
+            args.moe_token_dispatcher_type = "flex"
+            print_rank_0(
+                f"[Primus:DeepSeek-V4][P23] override "
+                f"args.moe_token_dispatcher_type='{current or 'unset'}' -> "
+                f"'flex' (Turbo DeepEP dispatcher path)."
+            )
+    else:
+        print_rank_0(
+            "[Primus:DeepSeek-V4][P23] WARNING: "
+            f"args.use_turbo_deepep=True but moe_token_dispatcher_type="
+            f"'{current}' (not 'flex' / 'alltoall').  Keeping the user's "
+            "dispatcher type; the Turbo DeepEP path will not engage for V4 "
+            "MoE layers."
+        )
+
+
+def deepseek_v4_builder(
+    args,
+    pre_process,
+    post_process,
+    vp_stage=None,
+    config: Optional[DeepSeekV4TransformerConfig] = None,
+    pg_collection=None,
+):
+    """Build a DeepSeek-V4 model.
+
+    Phase 8: build from a DeepSeek runtime spec tree only.
+    """
+    print_rank_0("[Primus:DeepSeek-V4] building DeepseekV4Model...")
+
+    # Plan-3 P22: plumb V4 attn_sink + attn_sliding_window into the
+    # Turbo flash-attn sink-attention args before any V4 attention
+    # module is constructed (PrimusTurboAttention reads these in
+    # ``__init__`` from get_args()).
+    _maybe_plumb_v4_sink_attention_args(args)
+    # Plan-3 P23: plumb args.moe_enable_deepep + moe_token_dispatcher_type
+    # for the Turbo DeepEP MoE dispatcher path BEFORE config construction
+    # so the resulting ``config`` carries the right dispatcher type into
+    # ``_pick_v4_dispatcher_cls``.
+    _maybe_plumb_v4_turbo_deepep_args(args)
+
+    if config is None:
+        config = core_transformer_config_from_args(
+            args,
+            config_class=DeepSeekV4TransformerConfig,
+        )
+
+    assert not args.use_legacy_models, "DeepSeek-V4 requires use_legacy_models=False (Mcore-only)."
+
+    runtime_decoder_spec = _resolve_runtime_decoder_spec(args, config, vp_stage)
+
+    model = DeepseekV4Model(
+        config=config,
+        transformer_layer_spec=runtime_decoder_spec,
+        vocab_size=args.padded_vocab_size,
+        max_sequence_length=args.max_position_embeddings,
+        pre_process=pre_process,
+        post_process=post_process,
+        fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+        parallel_output=True,
+        share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+        position_embedding_type=args.position_embedding_type,
+        rotary_percent=args.rotary_percent,
+        rotary_base=args.rotary_base,
+        rope_scaling=args.use_rope_scaling,
+        pg_collection=pg_collection,
+        vp_stage=vp_stage,
+    )
+    return model
+
+
+def model_provider(
+    model_builder=None,
+    pre_process: bool = True,
+    post_process: bool = True,
+    vp_stage: Optional[int] = None,
+    config: Optional[DeepSeekV4TransformerConfig] = None,
+    pg_collection=None,
+):
+    """``model_provider`` entry point used by Megatron's ``pretrain()``.
+
+    ``MegatronPretrainTrainer`` will pass ``deepseek_v4_builder`` as the
+    first arg via ``functools.partial`` so the upstream ``pretrain()`` can
+    call this with the standard ``(pre_process, post_process, vp_stage)``
+    signature.
+    """
+    if model_builder is None:
+        model_builder = deepseek_v4_builder
+
+    args = get_args()
+    if args.record_memory_history:
+        import torch
+
+        torch.cuda.memory._record_memory_history(
+            True,
+            trace_alloc_max_entries=100000,
+            trace_alloc_record_context=True,
+        )
+
+    return model_builder(
+        args,
+        pre_process,
+        post_process,
+        vp_stage,
+        config=config,
+        pg_collection=pg_collection,
+    )
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py
new file mode 100644
index 000000000..07a9b0013
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py
@@ -0,0 +1,697 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DeepSeek-V4 spec entry points.
+
+This module only defines DeepSeek-native runtime specs.
+"""
+
+import importlib
+import importlib.util
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.mlp import MLPSubmodules
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.moe.token_dispatcher import (
+    MoEAllGatherTokenDispatcher,
+    MoEAlltoAllTokenDispatcher,
+    MoEFlexTokenDispatcher,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import get_num_layers_to_build
+from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+
+from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+    DeepSeekV4SpecProvider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.build_context import (
+    resolve_v4_provider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import (
+    DeepseekV4HybridLayer,
+    DeepseekV4HybridLayerSubmodules,
+    DeepseekV4TransformerBlock,
+    DeepseekV4TransformerBlockSubmodules,
+    _DenseSwiGLUMLP,
+    _normalize_compress_ratios,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.compressor import Compressor
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+    DeepseekV4Attention,
+    DeepseekV4AttentionSubmodules,
+)
+from primus.backends.megatron.core.transformer.hyper_connection import (
+    HyperHead,
+    HyperMixer,
+)
+from primus.backends.megatron.core.transformer.indexer import Indexer
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+    DeepseekV4HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_moe import (
+    DeepseekV4MoE,
+    DeepseekV4MoESubmodules,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    DeepseekV4LearnedRouter,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Plan-3 P23 — Turbo DeepEP dispatcher gating
+# ---------------------------------------------------------------------------
+#
+# Both ``deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args`` (which
+# fires BEFORE config construction so the V4 config inherits the right
+# ``moe_token_dispatcher_type`` / ``moe_enable_deepep``) and
+# ``_pick_v4_dispatcher_cls`` below need the same gating predicate.
+# Centralised here so the two stay in sync; mirrors
+# ``primus.backends.megatron.patches.turbo.moe_dispatcher_patches._is_turbo_deepep_enabled``
+# (the patch itself fires too late for V4 spec build, but its gating
+# is the canonical one).
+#
+# ``deepseek_v4_builders.py`` imports this helper, which is one-way
+# (builders → layer_specs); putting the helper in builders would
+# create a circular import because layer_specs would need to call
+# back.
+
+# Public name so unit tests can patch / probe it directly.
+PRIMUS_TURBO_DEEPEP_DISPATCHER_NAME = "PrimusTurboDeepEPTokenDispatcher"
+
+
+def is_v4_turbo_deepep_active(args) -> bool:
+    """Plan-3 P23: V4-side gate for the Turbo DeepEP MoE dispatcher.
+
+    Returns ``True`` only when **all** of the following hold (matches
+    the conditions enforced by the upstream
+    ``megatron.turbo.moe_dispatcher`` patch):
+
+    * ``primus_turbo`` package is importable;
+    * ``args.enable_primus_turbo`` is True;
+    * ``args.use_turbo_deepep`` is True;
+    * ``args.tensor_model_parallel_size == 1``
+      (PrimusTurboDeepEPTokenDispatcher requires TPxEP > 1; we keep
+      TP == 1 here because TP > 1 paths haven't been validated against
+      the Turbo dispatcher and the existing patch enforces the same
+      gate).
+    """
+    if importlib.util.find_spec("primus_turbo") is None:
+        return False
+    if not bool(getattr(args, "enable_primus_turbo", False)):
+        return False
+    if not bool(getattr(args, "use_turbo_deepep", False)):
+        return False
+    tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1)
+    if tp_size != 1:
+        return False
+    return True
+
+
+def _import_primus_turbo_deepep_dispatcher_cls():
+    """Plan-3 P23: lazy import of PrimusTurboDeepEPTokenDispatcher.
+
+    Returns ``None`` when the import fails (callers must fall back to
+    the upstream MoEFlexTokenDispatcher in that case + emit a warning).
+    The import is lazy because the wider primus build is allowed to
+    run without ``primus_turbo`` installed (e.g. CPU unit tests for
+    non-Turbo paths).
+    """
+    if importlib.util.find_spec("primus_turbo") is None:
+        return None
+    try:
+        module = importlib.import_module("primus.backends.megatron.core.extensions.primus_turbo")
+    except Exception as exc:  # pragma: no cover - defensive
+        logger.warning(
+            "[DeepSeek-V4][P23] Failed to import "
+            "primus.backends.megatron.core.extensions.primus_turbo: %s",
+            exc,
+        )
+        return None
+    return getattr(module, PRIMUS_TURBO_DEEPEP_DISPATCHER_NAME, None)
+
+
+def _pick_v4_dispatcher_cls(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    args=None,
+) -> Tuple[type, str]:
+    """Plan-3 P23: pick the MoE token-dispatcher class for V4 spec build.
+
+    Returns ``(dispatcher_cls, dispatcher_type)``.  The
+    ``dispatcher_type`` is the string label V4 stores in its config /
+    annotation (used by ``DeepseekV4MoE._resolve_dispatcher_type_from_spec``);
+    ``dispatcher_cls`` is the class that ``ModuleSpec`` will hand to
+    ``build_module``.
+
+    Selection order:
+
+    1. If ``config.moe_token_dispatcher_type == "allgather"`` →
+       :class:`MoEAllGatherTokenDispatcher` / ``"allgather"``.  This is
+       a different parallelism scheme (TPxEP all-gather/scatter), not
+       a deepep flavour; never overridden by the Turbo path.
+    2. If ``config.moe_token_dispatcher_type == "flex"`` AND
+       :func:`is_v4_turbo_deepep_active` returns True →
+       :class:`PrimusTurboDeepEPTokenDispatcher` / ``"flex"``.  When
+       the package isn't importable we log a one-shot warning and
+       fall back to the upstream :class:`MoEFlexTokenDispatcher`.
+    3. If ``config.moe_token_dispatcher_type == "flex"`` (Turbo
+       inactive) → :class:`MoEFlexTokenDispatcher` / ``"flex"``.
+    4. Anything else (default V4 base.yaml: ``"alltoall"``) →
+       :class:`MoEAlltoAllTokenDispatcher` / ``"alltoall"``.  Unknown
+       values emit a one-shot warning and fall through to alltoall.
+
+    The ``args=`` keyword is for unit tests; production callers pass
+    ``None`` and the helper consults ``megatron.training.get_args()``
+    lazily.  ``args=None`` with no Megatron args available is
+    interpreted as "Turbo not active" (i.e. take the non-turbo
+    branch).  This keeps ``_build_ffn_spec`` callable from CPU unit
+    tests that haven't called ``initialize_megatron``.
+    """
+    dispatcher_type = str(getattr(config, "moe_token_dispatcher_type", None) or "").lower()
+
+    if dispatcher_type == "allgather":
+        return MoEAllGatherTokenDispatcher, "allgather"
+
+    if dispatcher_type == "flex":
+        if args is None:
+            try:
+                from megatron.training import get_args as _megatron_get_args
+
+                args = _megatron_get_args()
+            except Exception:
+                args = None
+        if args is not None and is_v4_turbo_deepep_active(args):
+            turbo_cls = _import_primus_turbo_deepep_dispatcher_cls()
+            if turbo_cls is not None:
+                logger.info(
+                    "[DeepSeek-V4][P23] MoE dispatcher class resolved to "
+                    "PrimusTurboDeepEPTokenDispatcher (Turbo DeepEP path)."
+                )
+                return turbo_cls, "flex"
+            logger.warning(
+                "[DeepSeek-V4][P23] use_turbo_deepep=True but "
+                "PrimusTurboDeepEPTokenDispatcher is not importable; "
+                "falling back to MoEFlexTokenDispatcher."
+            )
+        return MoEFlexTokenDispatcher, "flex"
+
+    if dispatcher_type not in ("", "alltoall"):
+        logger.warning(
+            "[DeepSeek-V4] unsupported moe_token_dispatcher_type=%s; fallback to alltoall.",
+            dispatcher_type,
+        )
+    return MoEAlltoAllTokenDispatcher, "alltoall"
+
+
+def _default_init_method(_weight) -> None:
+    return None
+
+
+def _v4_fp8_attn_proj(config: "DeepSeekV4TransformerConfig") -> bool:
+    """FP8-ify the attention projections (q-up / o-proj) that otherwise fall
+    back to bf16 because the TE/Turbo fp8 linear rejects gather_output /
+    scatter-input. Only safe at TP=1, where gather/scatter are no-ops — so
+    we keep the bf16 gather/scatter native path for TP>1. Opt-in via
+    PRIMUS_V4_FP8_ATTN_PROJ=1.
+    """
+    return (
+        os.environ.get("PRIMUS_V4_FP8_ATTN_PROJ", "0") == "1"
+        and getattr(config, "tensor_model_parallel_size", 1) == 1
+    )
+
+
+def _build_linear_projection_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+    in_features: int,
+    out_features: int,
+) -> ModuleSpec:
+    """Default projection spec — a duplicated TE linear (no TP sharding).
+
+    Used for projections that V4's grouped-low-rank O does not natively
+    shard along TP (``linear_q_down_proj``, ``linear_kv``, ``linear_o_a``).
+    Keep these duplicated for now; full TP sharding of the grouped O
+    projection is tracked in P14.
+    """
+    return ModuleSpec(
+        module=provider.linear(),
+        params={
+            "input_size": in_features,
+            "output_size": out_features,
+            "parallel_mode": "duplicated",
+            "config": config,
+            "init_method": config.init_method or _default_init_method,
+            "bias": False,
+            "skip_bias_add": False,
+            "skip_weight_param_allocation": False,
+            "tp_comm_buffer_name": None,
+            "is_expert": False,
+        },
+    )
+
+
+def _build_column_parallel_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+    in_features: int,
+    out_features: int,
+    gather_output: bool = True,
+) -> ModuleSpec:
+    """Column-parallel projection spec.
+
+    With ``gather_output=True`` the output dim is gathered back to full
+    width across TP ranks, so downstream attention math (which assumes
+    ``H * head_dim`` per rank) does not need to know about TP at all.
+    Memory of the projection's weight matrix is sharded across TP ranks
+    even at ``gather_output=True``.
+
+    Plan-3 P21: TE / Turbo column-parallel wrappers explicitly reject
+    ``gather_output=True`` (see
+    ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:747/972``).
+    When the caller asks for the gather variant we route to the
+    upstream Megatron-native :class:`ColumnParallelLinear` via
+    ``provider.column_parallel_linear_with_gather_output()``; the
+    standard TE path stays for ``gather_output=False``.
+
+    Plan-2 P13 follow-up: this is used for ``linear_q_up_proj``. The
+    gather-then-shard variant for full sharded heads is tracked in P14
+    once the grouped-O TP plan lands.
+    """
+    # FP8 attention projections (paper recipe): the TE/Turbo fp8 column linear
+    # rejects gather_output=True, so q-up normally falls back to bf16 native.
+    # But at TP=1 the gather is a no-op, so we can route q-up through the fp8
+    # turbo linear (gather_output=False ≡ True) to capture it in mxfp8.
+    # Gated by PRIMUS_V4_FP8_ATTN_PROJ=1 and only when TP==1. Default off.
+    if gather_output and _v4_fp8_attn_proj(config):
+        module_cls = provider.column_parallel_linear()
+        gather_output = False
+    elif gather_output:
+        module_cls = provider.column_parallel_linear_with_gather_output()
+    else:
+        module_cls = provider.column_parallel_linear()
+    return ModuleSpec(
+        module=module_cls,
+        params={
+            "input_size": in_features,
+            "output_size": out_features,
+            "config": config,
+            "init_method": config.init_method or _default_init_method,
+            "gather_output": gather_output,
+            "bias": False,
+            "skip_bias_add": False,
+            "skip_weight_param_allocation": False,
+            "tp_comm_buffer_name": None,
+            "is_expert": False,
+        },
+    )
+
+
+def _build_row_parallel_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+    in_features: int,
+    out_features: int,
+    input_is_parallel: bool = False,
+) -> ModuleSpec:
+    """Row-parallel projection spec.
+
+    With ``input_is_parallel=False`` the linear scatters the input across
+    TP ranks internally and all-reduces the output, so the caller can
+    pass a full-width input tensor and get a full-width output tensor.
+    Weight memory is sharded across TP ranks. Used for ``linear_o_b``
+    and the flat-O fallback ``linear_proj``.
+
+    Plan-3 P21: TE / Turbo row-parallel wrappers explicitly reject
+    ``input_is_parallel=False`` (see
+    ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:1081``).
+    When the caller asks for scatter-input we route to the upstream
+    Megatron-native :class:`RowParallelLinear` via
+    ``provider.row_parallel_linear_with_scatter_input()``; the
+    standard TE path stays for ``input_is_parallel=True``.
+    """
+    # FP8 attention projections: the TE/Turbo fp8 row linear rejects
+    # input_is_parallel=False, so o-proj normally falls back to bf16 native.
+    # At TP=1 the scatter is a no-op, so route through the fp8 turbo linear
+    # (input_is_parallel=True ≡ False). Gated by PRIMUS_V4_FP8_ATTN_PROJ + TP==1.
+    if not input_is_parallel and _v4_fp8_attn_proj(config):
+        module_cls = provider.row_parallel_linear()
+        input_is_parallel = True
+    elif not input_is_parallel:
+        module_cls = provider.row_parallel_linear_with_scatter_input()
+    else:
+        module_cls = provider.row_parallel_linear()
+    return ModuleSpec(
+        module=module_cls,
+        params={
+            "input_size": in_features,
+            "output_size": out_features,
+            "config": config,
+            "init_method": config.init_method or _default_init_method,
+            "input_is_parallel": input_is_parallel,
+            "bias": False,
+            "skip_bias_add": False,
+            "tp_comm_buffer_name": None,
+            "is_expert": False,
+        },
+    )
+
+
+def _build_v4_attention_submodules(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+    compress_ratio: int,
+) -> DeepseekV4AttentionSubmodules:
+    """V4-canonical submodules for :class:`DeepseekV4Attention`.
+
+    Field names match the released V4-Flash checkpoint layout (and MLA's
+    canonical names where they overlap):
+
+    * ``linear_q_down_proj``  : ``hidden -> q_lora_rank``  (= ``wq_a``)
+    * ``q_layernorm``        : RMSNorm(``q_lora_rank``)    (= ``q_norm``)
+    * ``linear_q_up_proj``    : ``q_lora_rank -> n_heads * head_dim`` (= ``wq_b``)
+      — built as **column-parallel** so the projection's weight is
+      sharded across TP at ``tp > 1``. ``gather_output=True`` keeps
+      downstream math TP-agnostic.
+    * ``linear_kv``           : ``hidden -> head_dim``     (= ``wkv``,
+      single-latent KV)
+    * ``kv_layernorm``       : RMSNorm(``head_dim``)       (= ``kv_norm``)
+    * ``linear_o_a``         : grouped low-rank O down-proj (duplicated;
+      grouped-O TP plan is P14).
+    * ``linear_o_b``         : grouped low-rank O up-proj  (-> ``hidden``)
+      — built as **row-parallel** so its weight is sharded across TP.
+    * ``linear_proj``        : flat-O fallback (``o_lora_rank == 0``,
+      e.g. unit tests) — also row-parallel.
+    * ``compressor``         : :class:`Compressor` (compressed branches)
+    * ``indexer``            : :class:`Indexer` (CSA branch only)
+
+    The per-head learnable softmax sink lives directly on the attention
+    module as ``self.attn_sink: nn.Parameter`` — there is no separate
+    submodule slot (Plan-3 P21 dropped the ``attn_sink`` field; the
+    inline softmax-with-sink path in ``_attention_forward`` is canonical).
+    """
+    hidden_size = int(config.hidden_size)
+    num_heads = int(config.num_attention_heads)
+    head_dim = int(config.kv_channels or (hidden_size // num_heads))
+    q_lora_rank = int(config.q_lora_rank or 0)
+    o_groups = max(int(getattr(config, "o_groups", 1) or 1), 1)
+    o_lora_rank = int(getattr(config, "o_lora_rank", 0) or 0)
+
+    if q_lora_rank <= 0:
+        raise ValueError(
+            "DeepSeek-V4 requires q_lora_rank > 0; the released checkpoint "
+            "always low-rank-projects Q via wq_a / wq_b."
+        )
+
+    q_out = num_heads * head_dim
+    submods = DeepseekV4AttentionSubmodules(
+        linear_q_down_proj=_build_linear_projection_spec(
+            config=config,
+            provider=provider,
+            in_features=hidden_size,
+            out_features=q_lora_rank,
+        ),
+        linear_q_up_proj=_build_column_parallel_spec(
+            config=config,
+            provider=provider,
+            in_features=q_lora_rank,
+            out_features=q_out,
+        ),
+        linear_kv=_build_linear_projection_spec(
+            config=config,
+            provider=provider,
+            in_features=hidden_size,
+            out_features=head_dim,  # single-latent: K = V = wkv(hidden)
+        ),
+        q_layernorm=ModuleSpec(module=provider.v4_q_layernorm()),
+        kv_layernorm=ModuleSpec(module=provider.v4_kv_layernorm()),
+    )
+
+    if o_lora_rank > 0:
+        n_per_group = q_out // o_groups
+        submods.linear_o_a = _build_linear_projection_spec(
+            config=config,
+            provider=provider,
+            in_features=n_per_group,
+            out_features=o_groups * o_lora_rank,
+        )
+        submods.linear_o_b = _build_row_parallel_spec(
+            config=config,
+            provider=provider,
+            in_features=o_groups * o_lora_rank,
+            out_features=hidden_size,
+        )
+    else:
+        submods.linear_proj = _build_row_parallel_spec(
+            config=config,
+            provider=provider,
+            in_features=q_out,
+            out_features=hidden_size,
+        )
+
+    if compress_ratio > 0:
+        submods.compressor = ModuleSpec(module=Compressor)
+        if compress_ratio == 4:
+            submods.indexer = ModuleSpec(module=Indexer)
+    else:
+        # Plan-3 P22: dense layers route their softmax-and-attend through
+        # provider.core_attention() (PrimusTurboAttention when
+        # ``use_turbo_attention=True``, TEDotProductAttention otherwise).
+        # HCA + CSA layers do not get this slot — see
+        # ``DeepseekV4AttentionSubmodules`` docstring for why.
+        submods.core_attention = ModuleSpec(module=provider.core_attention())
+
+    return submods
+
+
+def _build_norm_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+):
+    del config
+    norm_module = provider.v4_norm_module()
+    assert norm_module is not None, "DeepSeek-V4 norm module must be provided by DeepSeekV4SpecProvider."
+    return ModuleSpec(module=norm_module)
+
+
+def _build_attention_spec(
+    *,
+    compress_ratio: int,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+) -> ModuleSpec:
+    """Plan-2 P13 attention spec — single :class:`DeepseekV4Attention`
+    class for all three V4 layer types (dense / HCA / CSA), dispatched
+    inside the class on ``compress_ratio``.
+
+    Plan-2 P16: ``attn_mask_type=AttnMaskType.causal`` is declared on the
+    spec params for upstream :class:`MultiTokenPredictionLayer`
+    compatibility. The value is functionally inert for V4 (the V4
+    attention forward manages its own SWA / sink mask internally) but
+    the upstream MTP layer's pre-build validator requires the field to
+    be one of ``{padding, causal, no_mask, padding_causal}`` when the
+    inner layer's submodules are :class:`TransformerLayerSubmodules` —
+    which they are, since :class:`DeepseekV4HybridLayerSubmodules` now
+    extends that dataclass.
+    """
+    return ModuleSpec(
+        module=DeepseekV4Attention,
+        params={
+            "compress_ratio": int(compress_ratio),
+            "attn_mask_type": AttnMaskType.causal,
+        },
+        submodules=_build_v4_attention_submodules(
+            config=config,
+            provider=provider,
+            compress_ratio=int(compress_ratio),
+        ),
+    )
+
+
+def _build_ffn_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    provider: DeepSeekV4SpecProvider,
+    layer_idx: int,
+) -> ModuleSpec:
+    num_routed_experts = int(config.num_moe_experts)
+    moe_use_grouped_gemm = bool(config.moe_grouped_gemm)
+    moe_use_legacy_grouped_gemm = bool(config.moe_use_legacy_grouped_gemm)
+    grouped_mlp_module, grouped_mlp_submodules = provider.v4_grouped_mlp_modules(
+        moe_use_grouped_gemm=moe_use_grouped_gemm,
+        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+    )
+    # Plan-3 P23: V4-side dispatcher selection.  Avoids the module-attr
+    # timing race where ``before_train`` patches the upstream
+    # ``MoEFlexTokenDispatcher`` symbol AFTER V4 spec build has captured
+    # it; by resolving the class locally we always get the right one.
+    dispatcher_cls, dispatcher_type = _pick_v4_dispatcher_cls(config)
+
+    assert (
+        grouped_mlp_module is not None
+    ), "DeepSeek-V4 grouped MLP module must be provided by DeepSeekV4SpecProvider."
+
+    grouped_experts_spec = ModuleSpec(
+        module=grouped_mlp_module,
+        submodules=grouped_mlp_submodules,
+    )
+
+    shared_expert_submodules = MLPSubmodules(
+        linear_fc1=provider.column_parallel_linear(),
+        linear_fc2=provider.row_parallel_linear(),
+        # P18 D2: V4-aware activation_func selection — see
+        # ``DeepSeekV4SpecProvider.v4_mlp_activation_func`` for why
+        # this is None on the eager (clamped-SwiGLU) path.
+        activation_func=provider.v4_mlp_activation_func(),
+    )
+    shared_expert_spec = ModuleSpec(
+        module=SharedExpertMLP,
+        submodules=shared_expert_submodules,
+    )
+
+    moe_submodules = DeepseekV4MoESubmodules(
+        hash_router=ModuleSpec(module=DeepseekV4HashRouter),
+        learned_router=ModuleSpec(module=DeepseekV4LearnedRouter),
+        token_dispatcher=ModuleSpec(module=dispatcher_cls),
+        grouped_experts=grouped_experts_spec,
+        shared_expert=shared_expert_spec,
+    )
+
+    if num_routed_experts > 0:
+        return ModuleSpec(
+            module=DeepseekV4MoE,
+            params={
+                "layer_idx": layer_idx,
+            },
+            submodules=moe_submodules,
+        )
+    return ModuleSpec(module=_DenseSwiGLUMLP)
+
+
+def _build_hybrid_layer_spec(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    provider: DeepSeekV4SpecProvider,
+    layer_idx: int,
+    compress_ratio: int,
+) -> ModuleSpec:
+    hc_mult = int(config.hc_mult)
+
+    layer_submodules = DeepseekV4HybridLayerSubmodules(
+        input_layernorm=_build_norm_spec(config=config, provider=provider),
+        self_attention=_build_attention_spec(
+            compress_ratio=compress_ratio,
+            config=config,
+            provider=provider,
+        ),
+        pre_mlp_layernorm=_build_norm_spec(config=config, provider=provider),
+        mlp=_build_ffn_spec(
+            config=config,
+            provider=provider,
+            layer_idx=layer_idx,
+        ),
+        attn_hc=ModuleSpec(module=HyperMixer) if hc_mult > 1 else None,
+        ffn_hc=ModuleSpec(module=HyperMixer) if hc_mult > 1 else None,
+    )
+
+    return ModuleSpec(
+        module=DeepseekV4HybridLayer,
+        params={
+            "layer_idx": layer_idx,
+            "compress_ratio": int(compress_ratio),
+        },
+        submodules=layer_submodules,
+    )
+
+
+def _build_stage_hybrid_layer_specs(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    provider: DeepSeekV4SpecProvider,
+    vp_stage: Optional[int],
+) -> List[ModuleSpec]:
+    """Build the current stage's decoder layer specs.
+
+    DeepSeek-V4 runtime always materializes a concrete stage-local
+    `layer_specs` list for `DeepseekV4TransformerBlock`.
+    """
+    num_layers = int(config.num_layers)
+    mtp_num_layers = int(config.mtp_num_layers)
+    compress_ratios = _normalize_compress_ratios(
+        config.compress_ratios,
+        num_layers=num_layers,
+        mtp_num_layers=mtp_num_layers,
+    )
+
+    try:
+        local_layer_count = int(get_num_layers_to_build(config, vp_stage=vp_stage))
+        layer_offset = int(get_transformer_layer_offset(config, vp_stage=vp_stage))
+    except Exception:
+        local_layer_count = num_layers
+        layer_offset = 0
+
+    local_start = max(0, layer_offset)
+    local_end = min(num_layers, local_start + max(0, local_layer_count))
+    local_layer_indices = range(local_start, local_end)
+
+    return [
+        _build_hybrid_layer_spec(
+            config,
+            provider=provider,
+            layer_idx=layer_idx,
+            compress_ratio=int(compress_ratios[layer_idx]),
+        )
+        for layer_idx in local_layer_indices
+    ]
+
+
+def get_deepseek_v4_runtime_decoder_spec(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
+) -> ModuleSpec:
+    """Return the effective V4 runtime decoder spec tree.
+
+    The returned block submodules always include a non-empty `layer_specs`.
+    """
+    del pp_rank
+
+    provider = resolve_v4_provider(config)
+    logger.info("[DeepSeek-V4] resolve spec provider=%s", type(provider).__name__)
+
+    hc_mult = int(config.hc_mult)
+    stage_layer_specs = _build_stage_hybrid_layer_specs(
+        config,
+        provider=provider,
+        vp_stage=vp_stage,
+    )
+    assert stage_layer_specs, "DeepSeek-V4 requires non-empty stage layer specs."
+    block_submodules = DeepseekV4TransformerBlockSubmodules(
+        layer_specs=stage_layer_specs,
+        hyper_head=ModuleSpec(module=HyperHead) if hc_mult > 1 else None,
+        # DeepseekV4TransformerBlock decides whether this stage owns final norm.
+        final_layernorm=_build_norm_spec(config=config, provider=provider),
+    )
+    return ModuleSpec(module=DeepseekV4TransformerBlock, submodules=block_submodules)
+
+
+__all__ = [
+    "get_deepseek_v4_runtime_decoder_spec",
+]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py
new file mode 100644
index 000000000..8964bb5c2
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py
@@ -0,0 +1,323 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DeepSeek-V4 top-level model.
+
+This model intentionally subclasses :class:`LanguageModule` (not GPTModel)
+so DeepSeek-V4 no longer depends on GPT's internal TransformerBlock
+construction path.
+"""
+
+from typing import Literal, Optional, Union
+
+from megatron.core import tensor_parallel
+from megatron.core.models.common.embeddings.language_model_embedding import (
+    LanguageModelEmbedding,
+)
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlock,
+    mtp_on_this_rank,
+    process_mtp_loss,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from torch import Tensor
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_specs import (
+    get_v4_mtp_block_spec,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+
+
+class DeepseekV4Model(LanguageModule):
+    """DeepSeek-V4 language model rooted on LanguageModule."""
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        transformer_layer_spec: Union[ModuleSpec, type],
+        vocab_size: int,
+        max_sequence_length: int,
+        pre_process: bool = True,
+        post_process: bool = True,
+        fp16_lm_cross_entropy: bool = False,
+        parallel_output: bool = True,
+        share_embeddings_and_output_weights: bool = False,
+        position_embedding_type: Literal[
+            "learned_absolute",
+            "rope",
+            "mrope",
+            "yarn",
+            "none",
+        ] = "none",
+        rotary_percent: float = 1.0,
+        rotary_base: int = 10000,
+        rope_scaling: bool = False,
+        scatter_embedding_sequence_parallel: bool = True,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        vp_stage: Optional[int] = None,
+        **_kwargs,
+    ) -> None:
+        del rotary_percent, rotary_base, rope_scaling
+        super().__init__(config=config, pg_collection=pg_collection)
+
+        self.transformer_layer_spec = transformer_layer_spec
+        self.vocab_size = vocab_size
+        self.max_sequence_length = max_sequence_length
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+        self.parallel_output = parallel_output
+        self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+        self.vp_stage = vp_stage
+        self.model_type = ModelType.encoder_or_decoder
+
+        if hasattr(self.config, "position_embedding_type"):
+            self.position_embedding_type = self.config.position_embedding_type
+        else:
+            self.position_embedding_type = position_embedding_type
+
+        # Compute ``mtp_process`` (and build the MTP block spec) *before* the
+        # embedding so we can mirror upstream GPTModel: when MTP layers live on
+        # a non-pre_process PP stage, that stage still needs a (tied) copy of
+        # the input embedding. Without this, ``setup_embeddings_and_output_layer``
+        # -> ``shared_embedding_or_output_weight`` asserts because ``self.embedding``
+        # was never created on the MTP stage.
+        #
+        # Plan-2 P16/P17: V4 wires multi-token prediction exclusively via
+        # the spec-based upstream :class:`MultiTokenPredictionBlock`,
+        # built from :func:`get_v4_mtp_block_spec`. The legacy
+        # primus-owned ``DeepseekV4MTPBlock`` (gated by
+        # ``v4_use_custom_mtp_block`` in plan-2 P16) was retired in plan-2
+        # P17; only the spec-based path remains.
+        mtp_num_layers = int(getattr(self.config, "mtp_num_layers", 0) or 0)
+        self.mtp_process = False
+        self.mtp_block_spec = None
+        if mtp_num_layers > 0:
+            self.mtp_block_spec = get_v4_mtp_block_spec(
+                self.config,
+                transformer_layer_spec=transformer_layer_spec,
+                vp_stage=vp_stage,
+            )
+            # ``mtp_on_this_rank`` reads ``parallel_state`` and
+            # :class:`MultiTokenPredictionBlock` walks ``pg_collection.cp``;
+            # both require a real distributed init. On CPU smokes (no
+            # ``torch.distributed``) we leave ``self.mtp`` as ``None`` and
+            # surface the spec via ``self.mtp_block_spec`` so callers can
+            # still inspect the MTP wiring (the spec helper itself is fully
+            # CPU-testable).
+            try:
+                self.mtp_process = mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=vp_stage)
+            except (AssertionError, RuntimeError, AttributeError):
+                self.mtp_process = False
+
+        # The embedding is needed on pre_process stages and, when MTP is
+        # enabled, also on MTP-process stages (which keep a tied copy).
+        if self.pre_process or self.mtp_process:
+            self.embedding = LanguageModelEmbedding(
+                config=self.config,
+                vocab_size=self.vocab_size,
+                max_sequence_length=self.max_sequence_length,
+                position_embedding_type=self.position_embedding_type,
+                scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
+                tp_group=self.pg_collection.tp,
+            )
+
+        self.decoder = build_module(
+            transformer_layer_spec,
+            config=self.config,
+            pre_process=self.pre_process,
+            post_process=self.post_process,
+            pg_collection=self.pg_collection,
+            vp_stage=vp_stage,
+        )
+
+        # NOTE: The cross-PP broadcast of ``input_ids`` for V4 hash-routed
+        # MoE layers (needed when middle PP stages own a hash-routed layer
+        # but Megatron's ``pretrain_gpt.get_batch`` returns ``None`` tokens
+        # for non-(first|last) PP stages) lives in a Primus patch on
+        # ``pretrain_gpt.get_batch``:
+        #   ``primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py``
+        # That hook runs once per ``forward_step`` (i.e. once per
+        # (chunk, microbatch) for both 1F1B and interleaved-1F1B), which is
+        # VPP-safe. An earlier in-``forward`` broadcast deadlocked the
+        # interleaved schedule because PP rank 0's broadcast wait blocked
+        # before its ``send_forward``, while PP rank 1 was simultaneously
+        # waiting for that ``send_forward`` in ``recv_forward``.
+
+        # ----- MTP block ---------------------------------------------------
+        # NOTE: do NOT pre-assign ``self.mtp = None``. Megatron's
+        # ``set_current_microbatch`` (third_party/Megatron-LM/megatron/core/
+        # transformer/cuda_graphs.py) probes ``hasattr(model, 'mtp')`` and
+        # unconditionally iterates ``model.mtp.layers``. We only create the
+        # attribute when MTP is actually live (matching upstream GPTModel).
+        if self.mtp_process:
+            self.mtp = MultiTokenPredictionBlock(
+                config=self.config,
+                spec=self.mtp_block_spec,
+                vp_stage=vp_stage,
+                pg_collection=self.pg_collection,
+            )
+
+        if self.post_process:
+            if getattr(self.config, "defer_embedding_wgrad_compute", False):
+                self.embedding_activation_buffer = []
+                self.grad_output_buffer = []
+            else:
+                self.embedding_activation_buffer = None
+                self.grad_output_buffer = None
+
+            self.output_layer = tensor_parallel.ColumnParallelLinear(
+                self.config.hidden_size,
+                self.vocab_size,
+                config=self.config,
+                init_method=(
+                    self.config.embedding_init_method
+                    if getattr(self.config, "use_mup", False) and not self.share_embeddings_and_output_weights
+                    else self.config.init_method
+                ),
+                bias=False,
+                skip_bias_add=False,
+                gather_output=not self.parallel_output,
+                skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights,
+                embedding_activation_buffer=self.embedding_activation_buffer,
+                grad_output_buffer=self.grad_output_buffer,
+                tp_group=self.pg_collection.tp,
+            )
+
+        if self.pre_process or self.post_process or self.mtp_process:
+            self.setup_embeddings_and_output_layer()
+
+    def set_input_tensor(self, input_tensor: Tensor) -> None:
+        """Pipeline-parallel hook to set decoder input tensor."""
+        if not isinstance(input_tensor, list):
+            input_tensor = [input_tensor]
+        assert len(input_tensor) == 1, "input_tensor should only be length 1 for decoder-only models"
+        self.decoder.set_input_tensor(input_tensor[0])
+
+    def forward(
+        self,
+        input_ids: Optional[Tensor],
+        position_ids: Optional[Tensor],
+        attention_mask: Optional[Tensor],
+        decoder_input: Optional[Tensor] = None,
+        labels: Optional[Tensor] = None,
+        loss_mask: Optional[Tensor] = None,
+        runtime_gather_output: Optional[bool] = None,
+        packed_seq_params=None,
+        **kwargs,
+    ):
+        """Forward pass for DeepSeek-V4.
+
+        Plan-2 P15: ``input_ids`` are passed to the decoder as the
+        ``token_ids`` forward kwarg (replacing the ``decoder._v4_token_ids``
+        attribute stash). Hash-routed MoE layers consume them directly via
+        the standard kwargs propagation chain
+        ``model.forward -> decoder.forward -> layer.forward -> mlp.forward
+        -> hash_router.forward``.
+
+        Plan-2 P16: when ``mtp_num_layers > 0`` and the spec-based MTP
+        path is enabled (default), this method runs the upstream
+        :class:`MultiTokenPredictionBlock` on the post_process stage and
+        feeds its concatenated output through :func:`process_mtp_loss`,
+        which adds the auxiliary MTP loss term to the main LM loss.
+        """
+        if decoder_input is None:
+            if self.pre_process:
+                if input_ids is None:
+                    raise ValueError("input_ids must be provided when pre_process=True.")
+                if position_ids is None:
+                    batch, seq = input_ids.shape
+                    position_ids = (
+                        input_ids.new_arange(seq, dtype=input_ids.dtype).unsqueeze(0).expand(batch, -1)
+                    )
+                decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+            else:
+                decoder_input = None
+
+        # ``input_ids`` arrives on every PP stage that owns hash-routed MoE
+        # layers because the Primus patch
+        # ``primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py``
+        # broadcasts the source-of-truth tokens from PP rank 0 inside
+        # ``pretrain_gpt.get_batch`` before this ``forward`` runs. So
+        # ``input_ids`` is non-``None`` here on every middle PP stage even
+        # though Megatron's data loader would otherwise feed it ``None``.
+
+        hidden_states = self.decoder(
+            hidden_states=decoder_input,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            token_ids=input_ids,
+            packed_seq_params=packed_seq_params,
+            **kwargs,
+        )
+
+        # Run the spec-based MTP block on stages that own MTP layers.
+        # Mirrors GPTModel's mtp_in_postprocess gating.
+        if self.mtp_process and getattr(self, "mtp", None) is not None:
+            hidden_states = self.mtp(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                packed_seq_params=packed_seq_params,
+                # MTP needs the (tied) word-embedding module to embed the
+                # shifted input_ids. On a dedicated MTP PP stage pre_process is
+                # False, but the stage still owns a tied embedding copy created
+                # in __init__ (gated on ``pre_process or mtp_process``), so pass
+                # whichever embedding module exists on this rank.
+                embedding=getattr(self, "embedding", None),
+            )
+
+        if not self.post_process:
+            return hidden_states
+
+        output_weight = None
+        if self.share_embeddings_and_output_weights:
+            output_weight = self.shared_embedding_or_output_weight()
+
+        # Plan-2 P16: when MTP is on, ``hidden_states`` arrives as the
+        # concatenation of the main-decoder hidden state plus
+        # ``mtp_num_layers`` shifted MTP hidden states (along the
+        # sequence axis). :func:`process_mtp_loss` splits the chunks,
+        # computes the per-depth MTP loss, and returns the main hidden
+        # state for the standard LM-head path below.
+        mtp_num_layers = int(getattr(self.config, "mtp_num_layers", 0) or 0)
+        if mtp_num_layers > 0 and getattr(self, "mtp", None) is not None:
+            cp_group = getattr(self.pg_collection, "cp", None)
+            hidden_states = process_mtp_loss(
+                hidden_states=hidden_states,
+                labels=labels,
+                loss_mask=loss_mask,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+                is_training=self.training,
+                compute_language_model_loss=self.compute_language_model_loss,
+                config=self.config,
+                cp_group=cp_group,
+                packed_seq_params=packed_seq_params,
+                scale_logits_fn=self._scale_logits if getattr(self.config, "use_mup", False) else None,
+            )
+
+        logits, _ = self.output_layer(
+            hidden_states,
+            weight=output_weight,
+            runtime_gather_output=runtime_gather_output,
+        )
+        logits = self._scale_logits(logits)
+
+        if labels is None:
+            return logits.transpose(0, 1).contiguous()
+        return self.compute_language_model_loss(labels, logits)
+
+
+__all__ = ["DeepseekV4Model"]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py
new file mode 100644
index 000000000..8c0aa7e82
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py
@@ -0,0 +1,189 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 multi-token-prediction (MTP) layer.
+
+V4's MTP differs from the V3 / upstream :class:`MultiTokenPredictionLayer`
+in two ways that the upstream layer cannot express on its own:
+
+1. **Multi-stream (mHC) inner layer.** The MTP inner transformer layer is a
+   :class:`DeepseekV4HybridLayer`, which — when ``hc_mult > 1`` — operates on
+   the K-stream form ``[B, S, K, D]`` (see ``deepseek_v4_block``). The upstream
+   MTP layer feeds its inner layer a single-stream ``[S, B, D]`` tensor, so we
+   must lift to K streams before the inner layer and collapse back after it.
+
+2. **Per-depth ``hc_head_fn``.** The released V4 checkpoint gives each MTP
+   depth its *own* small :class:`HyperHead` (``hc_head_fn``) to collapse the K
+   streams — it does **not** reuse the main trunk's HyperHead (techblog
+   §8 / DeepSeek-V4 report). This is gated by
+   ``config.mtp_use_separate_hc_head``.
+
+The class subclasses upstream :class:`MultiTokenPredictionLayer` so it keeps
+the shared embedding-roll / ``enorm`` / ``hnorm`` / ``eh_proj`` /
+final-layernorm machinery and plugs straight into
+:class:`MultiTokenPredictionBlock` (which builds each depth via
+``build_module(layer_spec, ...)`` and therefore honours our ``module`` slot).
+Only the inner-layer call site (``_proj_and_transformer_layer``) is overridden
+to insert the lift / collapse, and ``forward`` is wrapped to capture
+``position_ids`` (V4 attention derives its dual-RoPE from absolute positions,
+and the upstream layer does not thread ``position_ids`` into
+``_proj_and_transformer_layer``).
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import Optional
+
+import torch
+from megatron.core import tensor_parallel
+from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer
+from megatron.core.transformer.spec_utils import build_module
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import (
+    _lift_streams_in,
+    _lower_streams_out,
+)
+from primus.backends.megatron.core.transformer.hyper_connection import HyperHead
+
+try:  # get_fp8_context lives in megatron.core.fp8_utils across recent versions
+    from megatron.core.fp8_utils import get_fp8_context
+except Exception:  # pragma: no cover - defensive: keep BF16 path working
+
+    def get_fp8_context(config, *args, **kwargs):  # type: ignore[misc]
+        return nullcontext()
+
+
+class DeepseekV4MTPLayer(MultiTokenPredictionLayer):
+    """One V4 MTP depth (mHC-aware, per-depth HyperHead)."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        # Upstream MultiTokenPredictionLayer builds ``self.mtp_model_layer``
+        # WITHOUT threading ``pg_collection`` (see
+        # ``multi_token_prediction.py``: ``build_module(self.submodules.
+        # mtp_model_layer, config=..., vp_stage=..., layer_number=...,
+        # is_mtp_layer=True)``). For a V4 :class:`DeepseekV4HybridLayer` whose
+        # MLP is a :class:`DeepseekV4MoE`, ``pg_collection=None`` selects the
+        # *local-experts* (non-expert-parallel) path, which (a) instantiates
+        # ALL routed experts on every rank and (b) breaks the DDP grad-bucket
+        # invariant (``len(per_param_grad_ready_counts) == len(params)`` in
+        # ``param_and_grad_buffer.reset``) because the local path's per-expert
+        # modules do not match the EP dispatcher's grad-ready bookkeeping that
+        # the main decoder relies on. Rebuild the inner layer WITH the
+        # ``pg_collection`` the MTP block already holds so the MTP MoE uses the
+        # exact same expert-parallel dispatcher path as the main decoder. Free
+        # the throwaway first build before rebuilding to bound init memory.
+        pg_collection = kwargs.get("pg_collection", None)
+        vp_stage = kwargs.get("vp_stage", None)
+        if pg_collection is not None and getattr(self.submodules, "mtp_model_layer", None) is not None:
+            self.mtp_model_layer = None  # drop the pg_collection-less build
+            self.mtp_model_layer = build_module(
+                self.submodules.mtp_model_layer,
+                config=self.config,
+                vp_stage=vp_stage,
+                layer_number=self.layer_number,
+                is_mtp_layer=True,
+                pg_collection=pg_collection,
+            )
+
+        self.hc_mult = int(getattr(self.config, "hc_mult", 1) or 1)
+        use_separate_head = bool(getattr(self.config, "mtp_use_separate_hc_head", True))
+
+        self.mtp_hyper_head: Optional[HyperHead] = None
+        if self.hc_mult > 1:
+            if not use_separate_head:
+                # A single-stream MTP inner layer (no per-depth head) would
+                # require the inner DeepseekV4HybridLayer to be built with
+                # hc_mult=1, but it inherits the model's hc_mult and emits the
+                # K-stream form. Fail loud instead of silently shape-crashing.
+                raise NotImplementedError(
+                    "DeepseekV4MTPLayer with hc_mult>1 requires "
+                    "config.mtp_use_separate_hc_head=True (per-depth HyperHead); "
+                    "single-stream MTP-with-mHC is not wired."
+                )
+            self.mtp_hyper_head = HyperHead(
+                hidden_size=int(self.config.hidden_size),
+                hc_mult=self.hc_mult,
+                eps=float(getattr(self.config, "hc_eps", 1.0e-6)),
+            )
+
+        # Stash for ``_proj_and_transformer_layer`` (set in ``forward``).
+        self._mtp_position_ids: Optional[torch.Tensor] = None
+
+    # ------------------------------------------------------------------
+
+    def forward(self, *args, **kwargs):
+        """Capture ``position_ids`` then defer to the upstream forward.
+
+        The model calls this with keyword args
+        (``input_ids=``, ``position_ids=``, ``hidden_states=`` ...), so we
+        read ``position_ids`` from ``kwargs``. V4 attention consumes absolute
+        positions; the MTP token/label shift is handled by the upstream
+        embedding roll, so we thread the *unrolled* ``position_ids`` into the
+        inner hybrid layer.
+        """
+        if "position_ids" in kwargs:
+            self._mtp_position_ids = kwargs["position_ids"]
+        elif len(args) >= 2:
+            self._mtp_position_ids = args[1]
+        return super().forward(*args, **kwargs)
+
+    # ------------------------------------------------------------------
+
+    def _proj_and_transformer_layer(
+        self,
+        hidden_states: torch.Tensor,
+        decoder_input: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        **_kwargs,
+    ) -> torch.Tensor:
+        """eh_proj -> (lift -> V4 hybrid layer -> per-depth HyperHead) -> norm.
+
+        Mirrors the upstream method's fp8 / rng context handling but inserts
+        the K-stream lift / collapse around the inner
+        :class:`DeepseekV4HybridLayer` so the mHC math matches the main
+        decoder block exactly (same ``_lift_streams_in`` / ``_lower_streams_out``
+        helpers).
+        """
+        if self.config.sequence_parallel:
+            rng_context = tensor_parallel.get_cuda_rng_tracker().fork()
+        else:
+            rng_context = nullcontext()
+
+        if self.config.fp8:
+            fp8_context = get_fp8_context(self.config)
+            transformer_layer_fp8_context = get_fp8_context(self.config)
+        else:
+            fp8_context = nullcontext()
+            transformer_layer_fp8_context = nullcontext()
+
+        with rng_context:
+            with fp8_context:
+                # [S, B, D] single-stream after eh_proj.
+                hidden_states = self._concat_embeddings(hidden_states, decoder_input)
+
+            with transformer_layer_fp8_context:
+                # Lift to the K-stream form the V4 hybrid layer expects.
+                x = _lift_streams_in(
+                    hidden_states, pre_process=True, hc_mult=self.hc_mult
+                )  # [B, S, K, D] (or [B, S, D] when hc_mult == 1)
+                x, _ = self.mtp_model_layer(
+                    x,
+                    position_ids=self._mtp_position_ids,
+                    token_ids=None,
+                )
+                # Per-depth collapse K streams -> single stream.
+                if self.hc_mult > 1 and self.mtp_hyper_head is not None:
+                    x = self.mtp_hyper_head(x)  # [B, S, D]
+                hidden_states = _lower_streams_out(x, post_process=True, hc_mult=self.hc_mult)  # [S, B, D]
+
+        hidden_states = self._postprocess(hidden_states)
+        return hidden_states
+
+
+__all__ = ["DeepseekV4MTPLayer"]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py
new file mode 100644
index 000000000..b0ae7e583
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py
@@ -0,0 +1,201 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 MTP (Multi-Token Prediction) block spec.
+
+Plan-2 P16 wires V4 onto Megatron's upstream
+:class:`MultiTokenPredictionBlock` (multi-token prediction with the
+classic eh_proj + per-depth transformer layer + final layernorm
+recipe). The V4 specialization is:
+
+* Each MTP-depth's inner layer is a :class:`DeepseekV4HybridLayer` (with
+  HC, hash routing, and clamped-SwiGLU all wired through the same spec
+  tree as the main decoder).
+* The two pre-projection norms (``enorm`` over the embedding,
+  ``hnorm`` over the prior hidden state) and the post-MTP final
+  layernorm reuse the V4 RMSNorm provider.
+* The eh_proj is a column-parallel linear (also from the provider).
+
+What this file does *not* include:
+
+* The MTP block forward path itself — that's owned by upstream
+  :class:`MultiTokenPredictionBlock` once we hand it the spec.
+* The ``HyperHead`` / loss-aware shifting / RouterReplay snapshot
+  matching (those live inside :func:`process_mtp_loss` and the V4
+  layer's ``forward``).
+
+Reference: techblog §7 ("MTP V4 head") and
+``DeepSeek-V4-Flash/inference/model.py:MTPBlock``.
+
+This is the **only** MTP path in plan-2 (P17 retired the legacy
+primus-owned :class:`DeepseekV4MTPBlock` and the
+``v4_use_custom_mtp_block`` config flag). Set ``mtp_num_layers > 0`` in
+the config to enable MTP; the model wires this helper into
+:class:`MultiTokenPredictionBlock` automatically.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlock,
+    MultiTokenPredictionBlockSubmodules,
+    MultiTokenPredictionLayerSubmodules,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+    DeepSeekV4SpecProvider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.build_context import (
+    resolve_v4_provider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import (
+    DeepseekV4TransformerBlock,
+    DeepseekV4TransformerBlockSubmodules,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_layer import (
+    DeepseekV4MTPLayer,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+
+
+def _extract_v4_inner_layer_spec(transformer_layer_spec: ModuleSpec) -> ModuleSpec:
+    """Return a single :class:`DeepseekV4HybridLayer` spec for an MTP depth.
+
+    :class:`DeepseekV4Model` resolves the decoder as a
+    :class:`DeepseekV4TransformerBlock` ``ModuleSpec`` whose
+    ``submodules.layer_specs`` is the stage-local list of hybrid-layer specs.
+    Upstream :class:`MultiTokenPredictionLayer` needs a *single*
+    :class:`~megatron.core.transformer.transformer_layer.TransformerLayer`-style
+    spec as its inner layer (its ``__init__`` validates
+    ``mtp_model_layer.submodules`` against ``TransformerLayerSubmodules``),
+    so when handed the block spec we extract the last hybrid layer spec
+    (mirroring upstream GPT's ``spec.layer_specs[-1]`` convention; the last
+    decoder layer is ``cr=0`` dense in both V4-Flash and V4-Pro). When already
+    handed a layer spec (e.g. CPU unit tests) we thread it through unchanged.
+    """
+    submods = getattr(transformer_layer_spec, "submodules", None)
+    is_block_spec = getattr(
+        transformer_layer_spec, "module", None
+    ) is DeepseekV4TransformerBlock or isinstance(submods, DeepseekV4TransformerBlockSubmodules)
+    if is_block_spec:
+        layer_specs = getattr(submods, "layer_specs", None)
+        if not layer_specs:
+            raise ValueError(
+                "Cannot build a V4 MTP inner layer: the decoder block spec has "
+                "no stage-local layer_specs to extract from."
+            )
+        return layer_specs[-1]
+    return transformer_layer_spec
+
+
+def _v4_mtp_layer_spec(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    transformer_layer_spec: ModuleSpec,
+    provider: DeepSeekV4SpecProvider,
+) -> ModuleSpec:
+    """One MTP depth's :class:`MultiTokenPredictionLayer` spec.
+
+    Args:
+        config: V4 transformer config (carries ``hidden_size``, etc.).
+        transformer_layer_spec: the V4 hybrid-layer spec used as the
+            inner layer for this depth. Plan-2 wires it through
+            unchanged so the MTP layer reuses HC / hash routing /
+            clamped-SwiGLU exactly the way the main decoder does.
+        provider: V4 spec provider (resolves RMSNorm + column-parallel
+            linear modules).
+    """
+    del config
+    norm_module = provider.v4_norm_module()
+    column_parallel = provider.column_parallel_linear()
+
+    inner_layer_spec = _extract_v4_inner_layer_spec(transformer_layer_spec)
+
+    return ModuleSpec(
+        module=DeepseekV4MTPLayer,
+        submodules=MultiTokenPredictionLayerSubmodules(
+            enorm=norm_module,
+            hnorm=norm_module,
+            eh_proj=column_parallel,
+            mtp_model_layer=inner_layer_spec,
+            layer_norm=norm_module,
+        ),
+    )
+
+
+def get_v4_mtp_block_spec(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    transformer_layer_spec: ModuleSpec,
+    vp_stage: Optional[int] = None,
+) -> ModuleSpec:
+    """Return a :class:`ModuleSpec` for the V4 MTP block.
+
+    The returned spec wraps :class:`MultiTokenPredictionBlock` with one
+    :class:`MultiTokenPredictionLayer` spec per MTP depth (V4-Flash
+    uses ``mtp_num_layers=1``; larger variants may use more depths).
+    The block submodules carry the per-depth specs; the block's
+    ``__init__`` walks them and instantiates :class:`MultiTokenPredictionLayer`
+    instances that reuse the V4 hybrid layer for inner attention / MLP
+    math.
+
+    Args:
+        config: V4 transformer config. Must have
+            ``mtp_num_layers >= 1`` (caller checks before invoking
+            this helper).
+        transformer_layer_spec: the V4 hybrid-layer spec for the main
+            decoder. The same spec is reused for each MTP depth so the
+            inner attention + MoE math matches the main decoder
+            exactly. Plan-2 §16 confirms V4 uses a single
+            (non-repeated) shape for MTP layers.
+        vp_stage: optional virtual-pipeline stage index. Passed through
+            to upstream MTP code; ignored on non-VP runs.
+
+    Returns:
+        A ``ModuleSpec`` that builds a fully-wired
+        :class:`MultiTokenPredictionBlock` when handed to
+        :func:`build_module` with ``config`` + ``pg_collection``
+        kwargs.
+
+    Notes:
+        * V4 has no ``mtp_use_repeated_layer`` carry-over; the spec
+          replicates the inner layer ``mtp_num_layers`` times. (V4
+          checkpoints store one set of MTP weights per depth.)
+        * The ``HyperHead`` per-depth collapse lives inside
+          :class:`DeepseekV4HybridLayer` itself (via the layer's HC
+          mixers); the upstream MTP block does not need V4-specific
+          HyperHead awareness.
+    """
+    del vp_stage  # currently only forwarded by callers; not needed here
+
+    if int(config.mtp_num_layers or 0) < 1:
+        raise ValueError(
+            "get_v4_mtp_block_spec requires mtp_num_layers >= 1; "
+            f"got mtp_num_layers={config.mtp_num_layers!r}."
+        )
+
+    provider = resolve_v4_provider(config)
+    mtp_layer_specs = [
+        _v4_mtp_layer_spec(
+            config=config,
+            transformer_layer_spec=transformer_layer_spec,
+            provider=provider,
+        )
+        for _ in range(int(config.mtp_num_layers))
+    ]
+
+    return ModuleSpec(
+        module=MultiTokenPredictionBlock,
+        submodules=MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs),
+    )
+
+
+__all__ = ["get_v4_mtp_block_spec"]
diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py
new file mode 100644
index 000000000..837aad636
--- /dev/null
+++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py
@@ -0,0 +1,229 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 specific transformer config.
+
+This config extends Megatron's ``MLATransformerConfig`` with DeepSeek-V4
+runtime fields that are referenced by V4 modules but are not part of the
+upstream ``TransformerConfig``/``MLATransformerConfig`` schema.
+"""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+# ``mtp_compress_ratios`` and ``v4_use_custom_mtp_block`` lived here in
+# plan-1 / plan-2 P12-P16 as escape hatches for the legacy primus-owned
+# :class:`DeepseekV4MTPBlock`. Plan-2 P17 retired that block (the MTP path
+# is now exclusively the spec-based upstream
+# :class:`MultiTokenPredictionBlock` route via
+# :func:`get_v4_mtp_block_spec`); both fields are deliberately gone here
+# and their references in YAML configs / training scripts must be removed.
+
+
+def _normalize_compress_ratios_field(
+    value: Optional[Union[str, List[int], Tuple[int, ...]]],
+    *,
+    field_name: str = "compress_ratios",
+) -> Optional[Tuple[int, ...]]:
+    """Plan-2 P18 (D4 audit): normalize ``compress_ratios`` to a tuple.
+
+    YAML loaders deliver this field as a *string* (e.g.
+    ``"[0, 0, 4, 128, ...]"``) when wrapped in quotes, or as an actual
+    list when written without quotes. The dataclass stored both as
+    ``Optional[Union[str, List[int], Tuple[int, ...]]]``, and runtime
+    helpers (``_parse_int_sequence`` / ``_normalize_compress_ratios``
+    in ``deepseek_v4_block.py``) had to ``ast.literal_eval`` the string
+    on every consumer path.
+
+    With this helper running once in ``__post_init__``, every consumer
+    sees a single canonical type — ``tuple[int, ...]`` — and the runtime
+    parsing layer is reduced to a length-fitting check.
+
+    Args:
+        value: raw config value (string, list, tuple, or ``None``).
+        field_name: only used for error messages.
+
+    Returns:
+        ``None`` (when ``value`` is ``None``) or a tuple of ints.
+    """
+    if value is None:
+        return None
+    parsed = value
+    if isinstance(parsed, str):
+        try:
+            parsed = ast.literal_eval(parsed)
+        except (SyntaxError, ValueError) as exc:
+            raise ValueError(f"{field_name} must be a list-like value, got {value!r}") from exc
+    if isinstance(parsed, (list, tuple)):
+        try:
+            return tuple(int(x) for x in parsed)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"{field_name} entries must be int-castable; got {parsed!r}") from exc
+    raise TypeError(f"{field_name} must be a list/tuple/str, got {type(parsed).__name__}")
+
+
+@dataclass
+class DeepSeekV4TransformerConfig(MLATransformerConfig):
+    # ---- DeepSeek-V4 hybrid attention / HC ----
+    hc_mult: int = 1
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1.0e-6
+
+    # ---- DeepSeek-V4 MTP (multi-token prediction) ----
+    # The released V4 checkpoint gives each MTP depth its OWN small
+    # ``hc_head_fn`` (a per-depth :class:`HyperHead`) instead of reusing the
+    # main trunk's HyperHead. With ``hc_mult > 1`` the MTP inner
+    # :class:`DeepseekV4HybridLayer` runs on the K-stream form and must be
+    # collapsed back to a single stream by this per-depth head before the
+    # MTP final-layernorm + shared LM head. Set ``False`` to fall back to a
+    # single-stream MTP inner layer (no per-depth head); see
+    # ``deepseek_v4_mtp_layer.DeepseekV4MTPLayer``.
+    mtp_use_separate_hc_head: bool = True
+
+    compress_ratios: Optional[Union[str, List[int], Tuple[int, ...]]] = None
+    compress_rope_theta: float = 160000.0
+
+    # ---- DeepSeek-V4 attention extras ----
+    attn_sliding_window: int = 0
+    attn_sink: bool = False
+    index_topk: int = 0
+    index_head_dim: int = 128
+    index_n_heads: int = 64
+
+    # ---- DeepSeek-V4 FP8 Indexer (CSA selector QK path) ----
+    # When True, the CSA :class:`Indexer` fake-quantizes its query / compressed
+    # key activations to FP8 (E4M3) before the QK scoring einsum, simulating
+    # the released V4 low-precision indexer QK path (the report runs the
+    # indexer QK in FP4/FP8 while keeping the BF16 index-score / top-k path).
+    # The score reduction (ReLU + per-head weight + sum + causal mask) and the
+    # top-k selection stay in the activation dtype (BF16). The indexer is a
+    # non-differentiable, frozen top-k selector, so this is an inference-side
+    # precision reduction (no straight-through estimator needed). Default False
+    # so existing BF16 runs are bit-unchanged; enable via
+    # ``PRIMUS_USE_V4_FP8_INDEXER`` (see run_deepseek_v4.sh / flash yaml).
+    use_v4_fp8_indexer: bool = False
+
+    # ---- DeepSeek-V4 indexer distillation loss (CSA selector training) ----
+    # ``topk`` is not differentiable, so the CSA lightning indexer only learns
+    # through an auxiliary objective: KL between the indexer's score
+    # distribution and the distribution the real attention places over the same
+    # compressed entries (DeepSeek-V3.2 section 2.1). ``1e-2`` is a reasonable
+    # starting value.
+    #
+    # ``0.0`` (the default) disables the loss AND keeps the indexer parameters
+    # frozen, which is the right setting when loading an already-trained
+    # indexer: frozen params stay out of the distributed-optimizer grad buckets
+    # so ``overlap_grad_reduce`` keeps working. Set > 0 to unfreeze the indexer
+    # and train it -- required for a from-scratch pretrain, where a frozen
+    # randomly-initialised indexer selects essentially at random.
+    v4_indexer_distill_loss_coeff: float = 0.0
+
+    # ---- DeepSeek-V4 attention backend selection (unified string selectors) ----
+    # ``use_v4_attention_backend`` selects the dense (cr=0) / HCA (cr=128) kernel;
+    # ``use_v4_csa_attention_backend`` selects the CSA (cr=4) kernel:
+    #   dense/HCA: eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1 | turbo
+    #   CSA:       eager | triton_v0 | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v0 | flydsl_v1 | turbo
+    # (triton_v0 = deprecated gathered; flydsl_v0 = deprecated legacy FlyDSL,
+    #  fwd-only; ``turbo`` = Primus-Turbo native-FlyDSL sparse-MLA via the turbo
+    #  API, primus_turbo.flydsl.attention). ``use_turbo_attention`` (when a
+    # ``core_attention`` module is built) still takes precedence for the dense path.
+    use_v4_attention_backend: str = "triton_v1"
+    use_v4_csa_attention_backend: str = "triton_v1"
+
+    # ---- Per-module precision (paper recipe): routed experts in MXFP4 while the
+    # rest of the layer runs FP8. Decoupled from the global --fp4/--fp8 recipe
+    # (which are mutually exclusive): with FP8 on, the PrimusTurbo grouped MLP
+    # routes the expert GEMMs through the native FP4 (hipBLASLt) path. Default OFF.
+    moe_experts_fp4: bool = False
+
+    # ---- DeepSeek-V4 plan-5 P29: torch.compile-fused Sinkhorn ----
+    # Plan-5 P29 (RESCOPED from "small-op fusion" — see plan-5 02-phase-
+    # details.md). The Sinkhorn-Knopp doubly-stochastic projection inside
+    # ``HyperMixer.compute_weights`` (see ``hyper_connection.py``) issues
+    # 1 + 2 * (n_iters - 1) separate fp32 ``aten::sum`` reductions per
+    # call. At V4-Flash production widths the per-call shape is
+    # ``[B, S, K, K] = [1, 4096, 4, 4]`` and HIP's default
+    # ``reduce_kernel<512, 1, ...>`` runs the kernel ~250x over the
+    # memory-bound floor. The P28 baseline trace pinned this kernel as
+    # 87.3 % of step time. Setting this flag to ``True`` collapses every
+    # such call into a single ``torch.compile(fullgraph=True,
+    # dynamic=False)`` Inductor-fused Triton kernel; AOT autograd handles
+    # the BWD. Default ``False`` until G32 (FWD + BWD parity) and G33b
+    # (post-P29 trace) flip it on.
+    use_v4_compiled_sinkhorn: bool = False
+
+    # ---- DeepSeek-V4 grouped low-rank output projection ----
+    # Mirrors the released checkpoint's `wo_a` / `wo_b` layout.
+    # When ``o_lora_rank == 0`` the attention falls back to a flat O proj
+    # (Megatron's ``linear_proj``); set it >0 to use the grouped low-rank
+    # form (``linear_o_a`` + ``linear_o_b``) with ``o_groups`` groups.
+    o_groups: int = 8
+    o_lora_rank: int = 0
+
+    # ---- DeepSeek-V4 MoE routing / expert extras ----
+    num_hash_layers: int = 0
+    hash_routing_seed: int = 0
+
+    moe_intermediate_size: Optional[int] = None
+    moe_use_legacy_grouped_gemm: bool = False
+
+    swiglu_limit: float = 0.0
+    v4_grouped_experts_support_clamped_swiglu: bool = False
+
+    # ---- Vocab helpers used by hash router ----
+    vocab_size: Optional[int] = None
+    padded_vocab_size: Optional[int] = None
+
+    # ---- Compat aliases for V4 code paths ----
+    norm_epsilon: Optional[float] = None
+    position_embedding_type: str = "none"
+
+    def __post_init__(self) -> None:
+        # P18 D4: normalize compress_ratios once so downstream helpers
+        # always see ``tuple[int, ...]`` (or None). YAML strings like
+        # ``"[0, 0, 4, ...]"`` are evaluated here.
+        if self.compress_ratios is not None:
+            self.compress_ratios = _normalize_compress_ratios_field(
+                self.compress_ratios, field_name="compress_ratios"
+            )
+
+        # Keep V4's ``norm_epsilon`` alias consistent with MCore's
+        # ``layernorm_epsilon`` before parent validation runs.
+        if self.norm_epsilon is None:
+            self.norm_epsilon = float(self.layernorm_epsilon)
+        self.layernorm_epsilon = float(self.norm_epsilon)
+
+        # DeepSeek naming compatibility for MoE hidden size.
+        if self.moe_ffn_hidden_size is None and self.moe_intermediate_size is not None:
+            self.moe_ffn_hidden_size = int(self.moe_intermediate_size)
+
+        # Keep DeepSeek clamp name aligned with MCore clamp field.
+        clamp_from_activation = self.activation_func_clamp_value
+        clamp_from_swiglu = float(self.swiglu_limit)
+        if clamp_from_activation is None:
+            if clamp_from_swiglu > 0.0:
+                self.activation_func_clamp_value = clamp_from_swiglu
+        elif clamp_from_swiglu <= 0.0:
+            self.swiglu_limit = float(clamp_from_activation)
+
+        # Ensure hash-router vocab lookups always have a concrete size.
+        if self.padded_vocab_size is None and self.vocab_size is not None:
+            self.padded_vocab_size = int(self.vocab_size)
+        if self.vocab_size is None and self.padded_vocab_size is not None:
+            self.vocab_size = int(self.padded_vocab_size)
+
+        super().__post_init__()
+
+        if self.moe_intermediate_size is None and self.moe_ffn_hidden_size is not None:
+            self.moe_intermediate_size = int(self.moe_ffn_hidden_size)
+
+
+__all__ = ["DeepSeekV4TransformerConfig"]
diff --git a/primus/backends/megatron/core/models/diffusion/__init__.py b/primus/backends/megatron/core/models/diffusion/__init__.py
new file mode 100644
index 000000000..95a3c358f
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/__init__.py
@@ -0,0 +1,36 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Diffusion models for Primus-Megatron.
+
+This module contains implementations of various diffusion model architectures
+following Megatron-Core conventions.
+
+Supported models:
+    - Flux: Flow-based diffusion with MMDiT architecture
+    - (Future) DiT: Diffusion Transformer
+    - (Future) MovieGen: Video diffusion
+"""
+
+
+# Lazy import for Flux to avoid early dependencies. Import common components and
+# other models from their submodules (e.g. ``...diffusion.common``,
+# ``...diffusion.flux``) directly.
+def __getattr__(name):
+    """Lazy import for the Flux model and config."""
+    if name == "Flux":
+        from primus.backends.megatron.core.models.diffusion.flux import Flux
+
+        return Flux
+    elif name == "FluxConfig":
+        from primus.backends.megatron.core.models.diffusion.flux import FluxConfig
+
+        return FluxConfig
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+__all__ = [
+    "Flux",
+    "FluxConfig",
+]
diff --git a/primus/backends/megatron/core/models/diffusion/common/__init__.py b/primus/backends/megatron/core/models/diffusion/common/__init__.py
new file mode 100644
index 000000000..090f51849
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/common/__init__.py
@@ -0,0 +1,38 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Common components for diffusion models in Primus.
+
+This module provides shared components used across different diffusion model
+architectures (Flux, DiT, MovieGen, etc.).
+"""
+
+from primus.backends.megatron.core.models.diffusion.common.config import (
+    BaseDiffusionConfig,
+)
+from primus.backends.megatron.core.models.diffusion.common.embeddings import (
+    MLPEmbedder,
+    TimeStepEmbedder,
+    Timesteps,
+    get_timestep_embedding,
+)
+from primus.backends.megatron.core.models.diffusion.common.normalization import (
+    AdaLN,
+    AdaLNContinuous,
+    RMSNorm,
+)
+
+__all__ = [
+    # Base classes
+    "BaseDiffusionConfig",
+    # Embeddings
+    "TimeStepEmbedder",
+    "MLPEmbedder",
+    "Timesteps",
+    "get_timestep_embedding",
+    # Normalization
+    "RMSNorm",
+    "AdaLN",
+    "AdaLNContinuous",
+]
diff --git a/primus/backends/megatron/core/models/diffusion/common/config.py b/primus/backends/megatron/core/models/diffusion/common/config.py
new file mode 100644
index 000000000..de4074744
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/common/config.py
@@ -0,0 +1,190 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Configuration classes for diffusion models.
+
+This module defines configuration dataclasses that extend Megatron-Core's
+TransformerConfig to include diffusion-specific parameters.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from megatron.core.enums import Fp8Recipe
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+@dataclass
+class BaseDiffusionConfig(TransformerConfig):
+    """
+    Base configuration for all diffusion models in Primus.
+
+    This class extends Megatron-Core's TransformerConfig to add common
+    diffusion model parameters. Model-specific configurations (FluxConfig,
+    DiTConfig, etc.) should inherit from this class.
+
+    Attributes:
+        model_type: Type of diffusion model (e.g., 'flux', 'dit', 'moviegen')
+        in_channels: Number of input channels in latent space
+        out_channels: Number of output channels (default: same as in_channels)
+        patch_size: Patch size for patchification (if applicable)
+        fp8_scaling_strategy: FP8 scaling strategy for local spec provider (default: 'dynamic')
+        fp8_force_nt_layout: FP8 backward GEMM layout (default: False)
+        fp8_reduce_amax: Whether to allreduce amax across ranks (default: False)
+        mxfp4_backward_precision: MXFP4 backward precision, 'mxfp4' or 'fp8' (default: 'mxfp4')
+        mxfp4_gradient_stochastic_rounding: Stochastic rounding on gradients (default: False)
+        sensitive_layers_enabled: Enable sensitive layer configuration (default: False)
+        sensitive_layers_start: Number of sensitive layers at start (default: 0)
+        sensitive_layers_end: Number of sensitive layers at end (default: 0)
+        sensitive_layer_precision: Precision for sensitive layers (default: 'bf16')
+
+    Inherited from TransformerConfig:
+        hidden_size: Hidden dimension size
+        num_attention_heads: Number of attention heads
+        ffn_hidden_size: FFN intermediate dimension
+        layernorm_epsilon: LayerNorm epsilon value
+        bf16, fp16, params_dtype: Precision settings
+        And many more Megatron-Core transformer parameters...
+    """
+
+    # Model identification
+    model_type: str = "base"
+
+    # Input/output dimensions
+    in_channels: int = 64
+    out_channels: Optional[int] = None  # Defaults to in_channels if None
+
+    # Patchification
+    patch_size: int = 1
+
+    # FP8 scaling strategy for local spec provider
+    fp8_scaling_strategy: str = "dynamic"
+
+    # FP8 backward GEMM layout for the local spec provider (tensorwise path only).
+    # False (default) = native layouts (dgrad=NN, wgrad=TN), the validated 0-NaN path
+    # on hipBLASLt 1.3. True = forced-NT (every GEMM normalized to NT via pre-transposed
+    # operands); faster on some stacks but NaN-prone on hipBLASLt 1.3 (gfx950).
+    # Only affects ScalingGranularity.TENSORWISE; rowwise/blockwise ignore it.
+    fp8_force_nt_layout: bool = False
+
+    # Whether to allreduce amax across DP/TP ranks for delayed FP8 scaling
+    fp8_reduce_amax: bool = False
+
+    # MXFP4 backward precision: "mxfp4" (pure) or "fp8" (hybrid)
+    mxfp4_backward_precision: str = "mxfp4"
+
+    # Stochastic rounding on MXFP4 gradients (paper Section 4.4)
+    mxfp4_gradient_stochastic_rounding: bool = False
+
+    # Sensitive layer configuration (clean naming, maps to Megatron internals)
+    sensitive_layers_enabled: bool = False
+    sensitive_layers_start: int = 0
+    sensitive_layers_end: int = 0
+    sensitive_layer_precision: str = "bf16"  # "bf16", "tw_fp8", or "mxfp8" (future)
+
+    def __post_init__(self):
+        """Post-initialization processing."""
+        # Pipeline parallelism is not implemented for diffusion models: the
+        # forward path runs embeddings/output head on every rank and does not
+        # relay activations between stages, so PP > 1 would silently
+        # miscompute. Reject it explicitly (before TransformerConfig validation)
+        # rather than producing wrong results.
+        if self.pipeline_model_parallel_size > 1:
+            raise ValueError(
+                "Diffusion models do not support pipeline parallelism; "
+                f"got pipeline_model_parallel_size={self.pipeline_model_parallel_size}. "
+                "Set pipeline_model_parallel_size=1."
+            )
+
+        if self.sensitive_layers_enabled:
+            if self.num_layers <= 1:
+                raise ValueError(
+                    "sensitive_layers_enabled=True requires num_layers to be set by the child config "
+                    "BEFORE calling super().__post_init__(). Set self.num_layers in your model config's "
+                    "__post_init__ before the super() call."
+                )
+            if self.sensitive_layers_start + self.sensitive_layers_end <= 0:
+                raise ValueError("sensitive_layers_enabled=True but both start and end counts are 0")
+            if self.sensitive_layers_start + self.sensitive_layers_end > self.num_layers:
+                raise ValueError(
+                    f"sensitive_layers_start ({self.sensitive_layers_start}) + "
+                    f"sensitive_layers_end ({self.sensitive_layers_end}) exceeds "
+                    f"num_layers ({self.num_layers})"
+                )
+            self.first_last_layers_bf16 = True
+            self.num_layers_at_start_in_bf16 = self.sensitive_layers_start
+            self.num_layers_at_end_in_bf16 = self.sensitive_layers_end
+
+        if self.sensitive_layers_enabled and self.sensitive_layer_precision == "tw_fp8":
+            _deferred_fp8 = "e4m3" if self.fp8 is None else None
+            _deferred_fp8_recipe = (
+                Fp8Recipe.tensorwise
+                if self.fp8_recipe is None or self.fp8_recipe == Fp8Recipe.delayed
+                else None
+            )
+        else:
+            _deferred_fp8 = None
+            _deferred_fp8_recipe = None
+
+        super().__post_init__()
+
+        # Apply deferred FP8 settings for sensitive layers (set after super to
+        # avoid Megatron's "fp4 and fp8 cannot coexist" validation).
+        if _deferred_fp8 is not None:
+            self.fp8 = _deferred_fp8
+        if _deferred_fp8_recipe is not None:
+            self.fp8_recipe = _deferred_fp8_recipe
+
+        # Re-run the FP8 validations that Megatron skipped because self.fp8 was
+        # None during super().__post_init__() (TransformerConfig lines 988-1017).
+        if self.fp8 and self.sensitive_layers_enabled:
+            if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed:
+                raise ValueError("Delayed scaling does not support first / last layer in BF16.")
+            max_bf16 = self.num_layers // self.pipeline_model_parallel_size
+            if self.first_last_layers_bf16:
+                if not (0 <= self.num_layers_at_start_in_bf16 <= max_bf16):
+                    raise ValueError(
+                        f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) "
+                        f"must be between 0 and {max_bf16}."
+                    )
+                if not (0 <= self.num_layers_at_end_in_bf16 <= max_bf16):
+                    raise ValueError(
+                        f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) "
+                        f"must be between 0 and {max_bf16}."
+                    )
+
+        if self.out_channels is None:
+            self.out_channels = self.in_channels
+
+        # Run configuration validation on construction. (Subclass fields used by
+        # validate() are plain dataclass fields, so they are already populated.)
+        self.validate()
+
+    def validate(self):
+        """
+        Validate configuration parameters.
+
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        if self.in_channels <= 0:
+            raise ValueError(f"in_channels must be positive, got {self.in_channels}")
+
+        if self.out_channels <= 0:
+            raise ValueError(f"out_channels must be positive, got {self.out_channels}")
+
+        if self.patch_size <= 0:
+            raise ValueError(f"patch_size must be positive, got {self.patch_size}")
+
+        if self.hidden_size <= 0:
+            raise ValueError(f"hidden_size must be positive, got {self.hidden_size}")
+
+        if self.num_attention_heads <= 0:
+            raise ValueError(f"num_attention_heads must be positive, got {self.num_attention_heads}")
+
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) must be divisible by "
+                f"num_attention_heads ({self.num_attention_heads})"
+            )
diff --git a/primus/backends/megatron/core/models/diffusion/common/embeddings.py b/primus/backends/megatron/core/models/diffusion/common/embeddings.py
new file mode 100644
index 000000000..8ef927730
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/common/embeddings.py
@@ -0,0 +1,296 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Embedding layers for diffusion models.
+
+Provides timestep and vector conditioning embeddings used across
+Flux and other diffusion architectures.
+
+This module implements:
+    - TimeStepEmbedder: Sinusoidal timestep embedding with MLP projection
+    - MLPEmbedder: Simple 2-layer MLP for vector conditioning (e.g., CLIP pooled)
+    - Timesteps: Helper module for sinusoidal timestep encoding
+    - get_timestep_embedding: Standalone function for sinusoidal embeddings
+
+Reference:
+    - Based on Denoising Diffusion Probabilistic Models (Ho et al., 2020)
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class TimeStepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations using sinusoidal encoding.
+
+    This module converts scalar timestep values into high-dimensional embeddings
+    using sinusoidal position encoding (similar to transformers), followed by
+    an MLP projection.
+
+    Architecture:
+        1. Sinusoidal position encoding (creates periodic features)
+        2. Linear projection to hidden_dim
+        3. SiLU activation
+        4. Linear projection to hidden_dim
+
+    Args:
+        embedding_dim: Number of channels for sinusoidal encoding (typically 256)
+        hidden_dim: Hidden dimension for MLP projection (model hidden size, e.g., 3072)
+        flip_sin_to_cos: Whether to flip sine and cosine order (default: True)
+        downscale_freq_shift: Controls delta between frequencies (default: 0)
+        scale: Scaling factor for timesteps (default: 1.0)
+        max_period: Maximum period for sinusoidal encoding (default: 10000)
+
+    Input:
+        t: Timestep scalars [B] or [B, 1], typically in range [0, 1000]
+
+    Output:
+        Timestep embeddings [B, hidden_dim]
+
+    Example:
+        >>> embedder = TimeStepEmbedder(embedding_dim=256, hidden_dim=3072)
+        >>> timesteps = torch.randn(4)  # Batch of 4 timesteps
+        >>> t_emb = embedder(timesteps)
+        >>> assert t_emb.shape == (4, 3072)
+
+    Reference:
+        - Denoising Diffusion Probabilistic Models (Ho et al., 2020)
+        - Adapted from NeMo's Flux implementation
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        hidden_dim: int,
+        flip_sin_to_cos: bool = True,
+        downscale_freq_shift: float = 0,
+        scale: float = 1.0,
+        max_period: int = 10000,
+    ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+        self.max_period = max_period
+
+        # MLP for projecting sinusoidal embeddings to hidden_dim
+        self.time_proj = Timesteps(
+            embedding_dim=embedding_dim,
+            flip_sin_to_cos=flip_sin_to_cos,
+            downscale_freq_shift=downscale_freq_shift,
+            scale=scale,
+            max_period=max_period,
+        )
+
+        self.time_embedding = MLPEmbedder(in_dim=embedding_dim, hidden_dim=hidden_dim)
+
+    def forward(self, t: Tensor) -> Tensor:
+        """
+        Forward pass: Convert timesteps to embeddings.
+
+        Args:
+            t: Timesteps [B] or [B, 1], typically in range [0, 1000]
+
+        Returns:
+            Timestep embeddings [B, hidden_dim]
+        """
+        # Get sinusoidal embeddings
+        t_emb = self.time_proj(t)  # [B, embedding_dim]
+
+        # Project through MLP
+        t_emb = self.time_embedding(t_emb)  # [B, hidden_dim]
+
+        return t_emb
+
+
+class Timesteps(nn.Module):
+    """
+    Converts timesteps to sinusoidal embeddings.
+
+    This is a helper module that creates sinusoidal position encodings
+    from scalar timestep values.
+
+    Args:
+        embedding_dim: Dimension of output embeddings
+        flip_sin_to_cos: Whether to order as [cos, sin] instead of [sin, cos]
+        downscale_freq_shift: Frequency shift parameter
+        scale: Scaling factor for embeddings
+        max_period: Maximum period for sinusoidal functions
+
+    Input:
+        timesteps: Scalar timesteps [B]
+
+    Output:
+        Sinusoidal embeddings [B, embedding_dim]
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        flip_sin_to_cos: bool = True,
+        downscale_freq_shift: float = 0,
+        scale: float = 1.0,
+        max_period: int = 10000,
+    ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+        self.max_period = max_period
+
+    def forward(self, timesteps: Tensor) -> Tensor:
+        """
+        Create sinusoidal timestep embeddings.
+
+        Args:
+            timesteps: Timesteps [B]
+
+        Returns:
+            Sinusoidal embeddings [B, embedding_dim]
+        """
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.embedding_dim,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+            max_period=self.max_period,
+        )
+        return t_emb
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = True,
+    downscale_freq_shift: float = 0,
+    scale: float = 1.0,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """
+    Create sinusoidal timestep embeddings.
+
+    This matches the implementation in Denoising Diffusion Probabilistic Models.
+    It creates position encodings using sine and cosine functions at different
+    frequencies.
+
+    Args:
+        timesteps: 1-D Tensor of N indices, one per batch element [B]
+        embedding_dim: Dimension of the output embeddings
+        flip_sin_to_cos: Whether embedding order should be [cos, sin] (True) or [sin, cos] (False)
+        downscale_freq_shift: Controls delta between frequencies
+        scale: Scaling factor applied to embeddings
+        max_period: Controls maximum frequency of embeddings
+
+    Returns:
+        Tensor of positional embeddings [B, embedding_dim]
+
+    Reference:
+        - Denoising Diffusion Probabilistic Models (Ho et al., 2020)
+        - Adapted from NeMo's Flux implementation
+    """
+    if len(timesteps.shape) != 1:
+        raise ValueError("Timesteps should be a 1d-array")
+
+    half_dim = embedding_dim // 2
+
+    # Remember input dtype to preserve it
+    input_dtype = timesteps.dtype
+
+    # Compute frequencies
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # Scale embeddings
+    emb = scale * emb
+
+    # Concatenate sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # Flip sine and cosine embeddings if requested
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # Zero pad if embedding_dim is odd
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+
+    # Convert back to input dtype
+    emb = emb.to(input_dtype)
+
+    return emb
+
+
+class MLPEmbedder(nn.Module):
+    """
+    Simple 2-layer MLP for vector conditioning (e.g., CLIP pooled embeddings).
+
+    This module provides a learnable projection from an input vector space
+    to the model's hidden dimension space. It's commonly used for:
+    - CLIP pooled text embeddings: [B, 768] -> [B, hidden_dim]
+    - Other vector-level conditioning signals
+
+    Architecture:
+        1. Linear projection: in_dim -> hidden_dim
+        2. SiLU activation
+        3. Linear projection: hidden_dim -> hidden_dim
+
+    Args:
+        in_dim: Input dimension (e.g., 768 for CLIP-L pooled embeddings)
+        hidden_dim: Hidden dimension (model hidden size, e.g., 3072)
+
+    Input:
+        x: Vector conditioning [B, in_dim]
+
+    Output:
+        Embedded conditioning [B, hidden_dim]
+
+    Example:
+        >>> # For CLIP-L pooled embeddings
+        >>> embedder = MLPEmbedder(in_dim=768, hidden_dim=3072)
+        >>> clip_pooled = torch.randn(4, 768)
+        >>> embedded = embedder(clip_pooled)
+        >>> assert embedded.shape == (4, 3072)
+
+    Reference:
+        - Adapted from NeMo's Flux implementation
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_dim = in_dim
+        self.hidden_dim = hidden_dim
+
+        # Two-layer MLP with SiLU activation
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass: Project input vectors to hidden dimension.
+
+        Args:
+            x: Input vectors [B, in_dim]
+
+        Returns:
+            Embedded vectors [B, hidden_dim]
+        """
+        x = self.in_layer(x)
+        x = self.silu(x)
+        x = self.out_layer(x)
+        return x
diff --git a/primus/backends/megatron/core/models/diffusion/common/normalization.py b/primus/backends/megatron/core/models/diffusion/common/normalization.py
new file mode 100644
index 000000000..c4c14edaf
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/common/normalization.py
@@ -0,0 +1,1226 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Normalization layers for diffusion models.
+
+Provides adaptive normalization layers that condition on timesteps and other
+conditioning signals. These are essential for diffusion model architectures
+like Flux and DiT.
+
+This module implements:
+    - RMSNorm: Root Mean Square Layer Normalization
+    - AdaLN: Adaptive Layer Normalization for DiT
+    - AdaLNContinuous: Continuous variant of AdaLN for Flux
+
+Reference:
+    - DiT Paper: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023)
+    - Flux Paper: "Flux: A Scalable Diffusion Model"
+"""
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+from megatron.core.jit import jit_fuser
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from torch import Tensor
+from torch.library import triton_op, wrap_triton
+
+# ---------------------------------------------------------------------------
+# Opaque LayerNorm custom op — prevents Inductor from decomposing
+# native_layer_norm into a Triton Welford reduction (which uses a different
+# FP32 accumulation order than the eager CUDA kernel, causing numerical
+# divergence amplified by FP8 quantisation).
+#
+# Unlike @torch.compiler.disable, this stays inside the compiled graph
+# (no graph break) while remaining opaque to fusion/decomposition.
+# ---------------------------------------------------------------------------
+
+_custom_op = torch.library.custom_op
+
+
+# ---------------------------------------------------------------------------
+# Opaque modulate custom op — prevents Inductor from fusing the
+# x * (1 + scale) + shift modulation with surrounding ops.
+#
+# Uses hand-written Triton kernels that fuse all pointwise ops into single
+# kernel launches.  The backward kernel accumulates dscale/dshift in a
+# deterministic sequential loop over the sequence dimension (float32
+# registers), avoiding the non-deterministic parallel-reduction order that
+# Inductor's auto-generated Triton kernels would use.
+# ---------------------------------------------------------------------------
+
+_MODULATE_BLOCK_H = 1024
+
+
+@triton.jit
+def _modulate_fwd_kernel(
+    X_ptr,
+    Scale_ptr,
+    Shift_ptr,
+    Out_ptr,
+    S,
+    B,
+    H,
+    stride_x_s,
+    stride_x_b,
+    stride_sc_b,
+    BLOCK_H: tl.constexpr,
+):
+    h_block = tl.program_id(0)
+    sb_idx = tl.program_id(1)
+    s_idx = sb_idx // B
+    b_idx = sb_idx % B
+
+    offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+
+    x_off = s_idx * stride_x_s + b_idx * stride_x_b + offs_h
+    sc_off = b_idx * stride_sc_b + offs_h
+
+    x = tl.load(X_ptr + x_off, mask=mask)
+    sc = tl.load(Scale_ptr + sc_off, mask=mask)
+    sh = tl.load(Shift_ptr + sc_off, mask=mask)
+    out = x * (1.0 + sc) + sh
+    tl.store(Out_ptr + x_off, out, mask=mask)
+
+
+@triton.jit
+def _modulate_bwd_kernel(
+    Grad_ptr,
+    X_ptr,
+    Scale_ptr,
+    DX_ptr,
+    DScale_ptr,
+    DShift_ptr,
+    S,
+    B,
+    H,
+    stride_g_s,
+    stride_g_b,
+    stride_x_s,
+    stride_x_b,
+    stride_sc_b,
+    BLOCK_H: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    h_block = tl.program_id(0)
+    b_idx = tl.program_id(1)
+
+    offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+
+    sc_off = b_idx * stride_sc_b + offs_h
+    scale_val = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32)
+    one_plus_scale = 1.0 + scale_val
+
+    acc_dscale = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc_dshift = tl.zeros([BLOCK_H], dtype=tl.float32)
+
+    for s_idx in range(S):
+        g_off = s_idx * stride_g_s + b_idx * stride_g_b + offs_h
+        x_off = s_idx * stride_x_s + b_idx * stride_x_b + offs_h
+        g = tl.load(Grad_ptr + g_off, mask=mask).to(tl.float32)
+        x = tl.load(X_ptr + x_off, mask=mask).to(tl.float32)
+        dx = g * one_plus_scale
+        tl.store(DX_ptr + x_off, dx.to(OUT_DTYPE), mask=mask)
+        acc_dscale += g * x
+        acc_dshift += g
+
+    tl.store(DScale_ptr + sc_off, acc_dscale.to(OUT_DTYPE), mask=mask)
+    tl.store(DShift_ptr + sc_off, acc_dshift.to(OUT_DTYPE), mask=mask)
+
+
+_TORCH_TO_TRITON_DTYPE = {
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+    torch.float32: tl.float32,
+}
+
+
+@_custom_op("primus::modulate", mutates_args=(), device_types="cuda")
+def _opaque_modulate(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+) -> torch.Tensor:
+    if x.dim() != 3 or scale.dim() != 2:
+        return x * (1 + scale) + shift
+    x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous()
+    out = torch.empty_like(x)
+    S, B, H = x.shape
+    BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H)
+    num_h_blocks = triton.cdiv(H, BLOCK_H)
+    _modulate_fwd_kernel[(num_h_blocks, S * B)](
+        x,
+        scale,
+        shift,
+        out,
+        S,
+        B,
+        H,
+        x.stride(0),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+    )
+    return out
+
+
+@_opaque_modulate.register_fake
+def _opaque_modulate_fake(x, scale, shift):
+    return torch.empty_like(x)
+
+
+@_custom_op("primus::modulate_backward", mutates_args=(), device_types="cuda")
+def _opaque_modulate_backward_op(
+    grad_output: torch.Tensor,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    need_dx: bool,
+    need_dscale: bool,
+    need_dshift: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2:
+        grad_output = grad_output.to(x.dtype)
+        dx = grad_output * (1 + scale) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype)
+        reduce_dims = list(range(grad_output.dim() - scale.dim()))
+        dscale = (
+            (grad_output * x).sum(dim=reduce_dims)
+            if need_dscale
+            else torch.empty(0, device=x.device, dtype=x.dtype)
+        )
+        dshift = (
+            grad_output.sum(dim=reduce_dims)
+            if need_dshift
+            else torch.empty(0, device=x.device, dtype=x.dtype)
+        )
+        return dx, dscale, dshift
+
+    grad_output = grad_output.to(x.dtype).contiguous()
+    x = x.contiguous()
+    scale = scale.contiguous()
+    S, B, H = x.shape
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H)
+    num_h_blocks = triton.cdiv(H, BLOCK_H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    _modulate_bwd_kernel[(num_h_blocks, B)](
+        grad_output,
+        x,
+        scale,
+        dx,
+        dscale,
+        dshift,
+        S,
+        B,
+        H,
+        grad_output.stride(0),
+        grad_output.stride(1),
+        x.stride(0),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    if not need_dx:
+        dx = torch.empty(0, device=x.device, dtype=x.dtype)
+    if not need_dscale:
+        dscale = torch.empty(0, device=x.device, dtype=x.dtype)
+    if not need_dshift:
+        dshift = torch.empty(0, device=x.device, dtype=x.dtype)
+    return dx, dscale, dshift
+
+
+@_opaque_modulate_backward_op.register_fake
+def _opaque_modulate_backward_fake(grad_output, x, scale, need_dx, need_dscale, need_dshift):
+    dx = torch.empty_like(x) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype)
+    dscale = torch.empty_like(scale) if need_dscale else torch.empty(0, device=x.device, dtype=x.dtype)
+    dshift = torch.empty_like(scale) if need_dshift else torch.empty(0, device=x.device, dtype=x.dtype)
+    return dx, dscale, dshift
+
+
+def _opaque_modulate_setup_context(ctx, inputs, output):
+    x, scale, _shift = inputs
+    ctx.save_for_backward(x, scale)
+
+
+def _opaque_modulate_backward(ctx, grad_output):
+    x, scale = ctx.saved_tensors
+    dx, dscale, dshift = _opaque_modulate_backward_op(
+        grad_output,
+        x,
+        scale,
+        True,
+        True,
+        True,
+    )
+    return dx, dscale, dshift
+
+
+_opaque_modulate.register_autograd(
+    _opaque_modulate_backward,
+    setup_context=_opaque_modulate_setup_context,
+)
+
+
+# ---------------------------------------------------------------------------
+# @triton_op versions — transparent to Inductor for cross-op fusion.
+# Enabled via config.use_triton_ops=True.
+# ---------------------------------------------------------------------------
+
+
+@triton_op("primus::modulate_v2", mutates_args=())
+def _triton_modulate(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+) -> torch.Tensor:
+    if x.dim() != 3 or scale.dim() != 2:
+        return x * (1 + scale) + shift
+    x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous()
+    out = torch.empty_like(x)
+    S, B, H = x.shape
+    BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H)
+    num_h_blocks = triton.cdiv(H, BLOCK_H)
+    wrap_triton(_modulate_fwd_kernel)[(num_h_blocks, S * B)](
+        x,
+        scale,
+        shift,
+        out,
+        S,
+        B,
+        H,
+        x.stride(0),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+    )
+    return out
+
+
+@_triton_modulate.register_fake
+def _triton_modulate_fake(x, scale, shift):
+    return torch.empty_like(x)
+
+
+@triton_op("primus::modulate_backward_v2", mutates_args=())
+def _triton_modulate_backward_op(
+    grad_output: torch.Tensor,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    need_dx: bool,
+    need_dscale: bool,
+    need_dshift: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2:
+        grad_output = grad_output.to(x.dtype)
+        dx = grad_output * (1 + scale) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype)
+        reduce_dims = list(range(grad_output.dim() - scale.dim()))
+        dscale = (
+            (grad_output * x).sum(dim=reduce_dims)
+            if need_dscale
+            else torch.empty(0, device=x.device, dtype=x.dtype)
+        )
+        dshift = (
+            grad_output.sum(dim=reduce_dims)
+            if need_dshift
+            else torch.empty(0, device=x.device, dtype=x.dtype)
+        )
+        return dx, dscale, dshift
+
+    grad_output = grad_output.to(x.dtype).contiguous()
+    x = x.contiguous()
+    scale = scale.contiguous()
+    S, B, H = x.shape
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H)
+    num_h_blocks = triton.cdiv(H, BLOCK_H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    wrap_triton(_modulate_bwd_kernel)[(num_h_blocks, B)](
+        grad_output,
+        x,
+        scale,
+        dx,
+        dscale,
+        dshift,
+        S,
+        B,
+        H,
+        grad_output.stride(0),
+        grad_output.stride(1),
+        x.stride(0),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    if not need_dx:
+        dx = torch.empty(0, device=x.device, dtype=x.dtype)
+    if not need_dscale:
+        dscale = torch.empty(0, device=x.device, dtype=x.dtype)
+    if not need_dshift:
+        dshift = torch.empty(0, device=x.device, dtype=x.dtype)
+    return dx, dscale, dshift
+
+
+@_triton_modulate_backward_op.register_fake
+def _triton_modulate_backward_fake(grad_output, x, scale, need_dx, need_dscale, need_dshift):
+    dx = torch.empty_like(x) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype)
+    dscale = torch.empty_like(scale) if need_dscale else torch.empty(0, device=x.device, dtype=x.dtype)
+    dshift = torch.empty_like(scale) if need_dshift else torch.empty(0, device=x.device, dtype=x.dtype)
+    return dx, dscale, dshift
+
+
+def _triton_modulate_setup_context(ctx, inputs, output):
+    x, scale, _shift = inputs
+    ctx.save_for_backward(x, scale)
+
+
+def _triton_modulate_backward(ctx, grad_output):
+    x, scale = ctx.saved_tensors
+    dx, dscale, dshift = _triton_modulate_backward_op(
+        grad_output,
+        x,
+        scale,
+        True,
+        True,
+        True,
+    )
+    return dx, dscale, dshift
+
+
+_triton_modulate.register_autograd(
+    _triton_modulate_backward,
+    setup_context=_triton_modulate_setup_context,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fused LayerNorm + Modulate Triton kernel — computes
+#   (x - mean) / sqrt(var + eps) * (1 + scale) + shift
+# in a single kernel launch, eliminating the intermediate ln_out tensor
+# from DRAM.  Only works for elementwise_affine=False LayerNorm (no
+# learnable weight/bias), which is the case for all AdaLN variants in Flux.
+#
+# The backward kernel fuses the LN backward formula with the modulate
+# backward, using deterministic sequential accumulation over S for
+# d_scale / d_shift (same pattern as _modulate_bwd_kernel).
+# ---------------------------------------------------------------------------
+
+_FUSED_LN_MOD_MAX_H = 8192
+
+
+@triton.jit
+def _fused_ln_modulate_fwd_kernel(
+    X_ptr,
+    Scale_ptr,
+    Shift_ptr,
+    Out_ptr,
+    Mean_ptr,
+    Rstd_ptr,
+    S,
+    B,
+    H,
+    eps,
+    stride_x_sb,
+    stride_sc_b,
+    BLOCK_H: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    b_idx = row % B
+    offs_h = tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+
+    x_off = row * stride_x_sb + offs_h
+    sc_off = b_idx * stride_sc_b + offs_h
+
+    x = tl.load(X_ptr + x_off, mask=mask, other=0.0).to(tl.float32)
+    mean = tl.sum(x, axis=0) / H
+
+    x_centered = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(x_centered * x_centered, axis=0) / H
+    rstd = 1.0 / tl.sqrt(var + eps)
+    x_hat = x_centered * rstd
+
+    sc = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32)
+    sh = tl.load(Shift_ptr + sc_off, mask=mask).to(tl.float32)
+    out = x_hat * (1.0 + sc) + sh
+
+    tl.store(Out_ptr + x_off, out.to(OUT_DTYPE), mask=mask)
+    tl.store(Mean_ptr + row, mean)
+    tl.store(Rstd_ptr + row, rstd)
+
+
+@triton.jit
+def _fused_ln_modulate_bwd_dscale_dshift_kernel(
+    Grad_ptr,
+    X_ptr,
+    Mean_ptr,
+    Rstd_ptr,
+    DScale_ptr,
+    DShift_ptr,
+    S,
+    B,
+    H,
+    stride_g_sb,
+    stride_x_sb,
+    XBLOCK: tl.constexpr,
+    RBLOCK: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
+    xmask = xindex < (B * H)
+    x_bh = xindex
+    x_b = xindex // H
+
+    r_base = tl.arange(0, RBLOCK)[None, :]
+
+    acc_dscale = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
+    acc_dshift = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
+
+    for r_off in range(0, S, RBLOCK):
+        r_idx = r_off + r_base
+        r_mask = r_idx < S
+        mask = xmask & r_mask
+
+        g_off = x_bh + (r_idx * B * H)
+        x_off = x_bh + (r_idx * B * H)
+        row = x_b + r_idx * B
+
+        g = tl.load(Grad_ptr + g_off, mask=mask, other=0.0).to(tl.float32)
+        x = tl.load(X_ptr + x_off, mask=mask, other=0.0).to(tl.float32)
+        mean_val = tl.load(Mean_ptr + row, mask=mask, other=0.0)
+        rstd_val = tl.load(Rstd_ptr + row, mask=mask, other=0.0)
+
+        x_hat = (x - mean_val) * rstd_val
+        acc_dscale = tl.where(mask, acc_dscale + g * x_hat, acc_dscale)
+        acc_dshift = tl.where(mask, acc_dshift + g, acc_dshift)
+
+    dscale_val = tl.sum(acc_dscale, 1)[:, None]
+    dshift_val = tl.sum(acc_dshift, 1)[:, None]
+    tl.store(DScale_ptr + x_bh, dscale_val.to(OUT_DTYPE), mask=xmask)
+    tl.store(DShift_ptr + x_bh, dshift_val.to(OUT_DTYPE), mask=xmask)
+
+
+@triton.jit
+def _fused_ln_modulate_bwd_dx_kernel(
+    Grad_ptr,
+    X_ptr,
+    Mean_ptr,
+    Rstd_ptr,
+    Scale_ptr,
+    DX_ptr,
+    S,
+    B,
+    H,
+    stride_g_sb,
+    stride_x_sb,
+    stride_sc_b,
+    BLOCK_H: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    sb_idx = tl.program_id(0)
+    b_idx = sb_idx % B
+    offs_h = tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+
+    sc_off = b_idx * stride_sc_b + offs_h
+    scale_val = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32)
+    one_plus_scale = 1.0 + scale_val
+
+    g_off = sb_idx * stride_g_sb + offs_h
+    x_off = sb_idx * stride_x_sb + offs_h
+
+    g = tl.load(Grad_ptr + g_off, mask=mask).to(tl.float32)
+    x = tl.load(X_ptr + x_off, mask=mask).to(tl.float32)
+    mean = tl.load(Mean_ptr + sb_idx)
+    rstd = tl.load(Rstd_ptr + sb_idx)
+
+    x_hat = (x - mean) * rstd
+    d_x_hat = g * one_plus_scale
+
+    inv_H = 1.0 / H
+    d_x_hat_masked = tl.where(mask, d_x_hat, 0.0)
+    xhat_masked = tl.where(mask, x_hat, 0.0)
+    c1 = tl.sum(xhat_masked * d_x_hat_masked, axis=0) * inv_H
+    c2 = tl.sum(d_x_hat_masked, axis=0) * inv_H
+    d_x = rstd * (d_x_hat - c2 - x_hat * c1)
+
+    tl.store(DX_ptr + x_off, d_x.to(OUT_DTYPE), mask=mask)
+
+
+@_custom_op("primus::fused_ln_modulate", mutates_args=(), device_types="cuda")
+def _opaque_fused_ln_modulate(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H:
+        out, mean, rstd = torch.ops.aten.native_layer_norm(
+            x,
+            [x.shape[-1]],
+            None,
+            None,
+            eps,
+        )
+        out = out * (1 + scale) + shift
+        return out, mean.flatten(), rstd.flatten()
+
+    x = x.contiguous()
+    scale, shift = scale.contiguous(), shift.contiguous()
+    S, B, H = x.shape
+    out = torch.empty_like(x)
+    M = S * B
+    mean = torch.empty(M, dtype=torch.float32, device=x.device)
+    rstd = torch.empty(M, dtype=torch.float32, device=x.device)
+    BLOCK_H = triton.next_power_of_2(H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    _fused_ln_modulate_fwd_kernel[(M,)](
+        x,
+        scale,
+        shift,
+        out,
+        mean,
+        rstd,
+        S,
+        B,
+        H,
+        eps,
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    return out, mean, rstd
+
+
+@_opaque_fused_ln_modulate.register_fake
+def _opaque_fused_ln_modulate_fake(x, scale, shift, eps):
+    out = torch.empty_like(x)
+    M = x.numel() // x.shape[-1]
+    mean = torch.empty(M, dtype=torch.float32, device=x.device)
+    rstd = torch.empty(M, dtype=torch.float32, device=x.device)
+    return out, mean, rstd
+
+
+@_custom_op("primus::fused_ln_modulate_backward", mutates_args=(), device_types="cuda")
+def _opaque_fused_ln_modulate_backward_op(
+    grad_output: torch.Tensor,
+    x: torch.Tensor,
+    mean: torch.Tensor,
+    rstd: torch.Tensor,
+    scale: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H:
+        grad_output = grad_output.to(x.dtype)
+        mean_r = mean.view(*x.shape[:-1], 1)
+        rstd_r = rstd.view(*x.shape[:-1], 1)
+        x_hat = (x.float() - mean_r) * rstd_r
+        d_x_hat = grad_output.float() * (1 + scale.float())
+        c1 = (x_hat * d_x_hat).mean(dim=-1, keepdim=True)
+        c2 = d_x_hat.mean(dim=-1, keepdim=True)
+        d_x = (rstd_r * (d_x_hat - c2 - x_hat * c1)).to(x.dtype)
+        reduce_dims = list(range(grad_output.dim() - scale.dim()))
+        dscale = (grad_output.float() * x_hat).sum(dim=reduce_dims).to(scale.dtype)
+        dshift = grad_output.sum(dim=reduce_dims).to(scale.dtype)
+        return d_x, dscale, dshift
+
+    grad_output = grad_output.to(x.dtype).contiguous()
+    x = x.contiguous()
+    scale = scale.contiguous()
+    S, B, H = x.shape
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    BLOCK_H = triton.next_power_of_2(H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    XBLOCK, RBLOCK = 256, 8
+    _fused_ln_modulate_bwd_dscale_dshift_kernel[(triton.cdiv(B * H, XBLOCK),)](
+        grad_output,
+        x,
+        mean,
+        rstd,
+        dscale,
+        dshift,
+        S,
+        B,
+        H,
+        grad_output.stride(1),
+        x.stride(1),
+        XBLOCK=XBLOCK,
+        RBLOCK=RBLOCK,
+        OUT_DTYPE=out_dtype,
+    )
+    _fused_ln_modulate_bwd_dx_kernel[(S * B,)](
+        grad_output,
+        x,
+        mean,
+        rstd,
+        scale,
+        dx,
+        S,
+        B,
+        H,
+        grad_output.stride(1),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    return dx, dscale, dshift
+
+
+@_opaque_fused_ln_modulate_backward_op.register_fake
+def _opaque_fused_ln_modulate_backward_fake(grad_output, x, mean, rstd, scale):
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    return dx, dscale, dshift
+
+
+def _opaque_fused_ln_modulate_setup_context(ctx, inputs, output):
+    x, scale, _shift, _eps = inputs
+    _out, mean, rstd = output
+    ctx.save_for_backward(x, mean, rstd, scale)
+
+
+def _opaque_fused_ln_modulate_backward(ctx, grad_output, _grad_mean, _grad_rstd):
+    x, mean, rstd, scale = ctx.saved_tensors
+    dx, dscale, dshift = _opaque_fused_ln_modulate_backward_op(
+        grad_output,
+        x,
+        mean,
+        rstd,
+        scale,
+    )
+    return dx, dscale, dshift, None
+
+
+_opaque_fused_ln_modulate.register_autograd(
+    _opaque_fused_ln_modulate_backward,
+    setup_context=_opaque_fused_ln_modulate_setup_context,
+)
+
+
+# ---------------------------------------------------------------------------
+# @triton_op fused LN+modulate — transparent to Inductor.
+# ---------------------------------------------------------------------------
+
+
+@triton_op("primus::fused_ln_modulate_v2", mutates_args=())
+def _triton_fused_ln_modulate(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H:
+        out, mean, rstd = torch.ops.aten.native_layer_norm(
+            x,
+            [x.shape[-1]],
+            None,
+            None,
+            eps,
+        )
+        out = out * (1 + scale) + shift
+        return out, mean.flatten(), rstd.flatten()
+
+    x = x.contiguous()
+    scale, shift = scale.contiguous(), shift.contiguous()
+    S, B, H = x.shape
+    out = torch.empty_like(x)
+    M = S * B
+    mean = torch.empty(M, dtype=torch.float32, device=x.device)
+    rstd = torch.empty(M, dtype=torch.float32, device=x.device)
+    BLOCK_H = triton.next_power_of_2(H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    wrap_triton(_fused_ln_modulate_fwd_kernel)[(M,)](
+        x,
+        scale,
+        shift,
+        out,
+        mean,
+        rstd,
+        S,
+        B,
+        H,
+        eps,
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    return out, mean, rstd
+
+
+@_triton_fused_ln_modulate.register_fake
+def _triton_fused_ln_modulate_fake(x, scale, shift, eps):
+    out = torch.empty_like(x)
+    M = x.numel() // x.shape[-1]
+    mean = torch.empty(M, dtype=torch.float32, device=x.device)
+    rstd = torch.empty(M, dtype=torch.float32, device=x.device)
+    return out, mean, rstd
+
+
+@triton_op("primus::fused_ln_modulate_backward_v2", mutates_args=())
+def _triton_fused_ln_modulate_backward_op(
+    grad_output: torch.Tensor,
+    x: torch.Tensor,
+    mean: torch.Tensor,
+    rstd: torch.Tensor,
+    scale: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H:
+        grad_output = grad_output.to(x.dtype)
+        mean_r = mean.view(*x.shape[:-1], 1)
+        rstd_r = rstd.view(*x.shape[:-1], 1)
+        x_hat = (x.float() - mean_r) * rstd_r
+        d_x_hat = grad_output.float() * (1 + scale.float())
+        c1 = (x_hat * d_x_hat).mean(dim=-1, keepdim=True)
+        c2 = d_x_hat.mean(dim=-1, keepdim=True)
+        d_x = (rstd_r * (d_x_hat - c2 - x_hat * c1)).to(x.dtype)
+        reduce_dims = list(range(grad_output.dim() - scale.dim()))
+        dscale = (grad_output.float() * x_hat).sum(dim=reduce_dims).to(scale.dtype)
+        dshift = grad_output.sum(dim=reduce_dims).to(scale.dtype)
+        return d_x, dscale, dshift
+
+    grad_output = grad_output.to(x.dtype).contiguous()
+    x = x.contiguous()
+    scale = scale.contiguous()
+    S, B, H = x.shape
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    BLOCK_H = triton.next_power_of_2(H)
+    out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype]
+    XBLOCK, RBLOCK = 256, 8
+    wrap_triton(_fused_ln_modulate_bwd_dscale_dshift_kernel)[(triton.cdiv(B * H, XBLOCK),)](
+        grad_output,
+        x,
+        mean,
+        rstd,
+        dscale,
+        dshift,
+        S,
+        B,
+        H,
+        grad_output.stride(1),
+        x.stride(1),
+        XBLOCK=XBLOCK,
+        RBLOCK=RBLOCK,
+        OUT_DTYPE=out_dtype,
+    )
+    wrap_triton(_fused_ln_modulate_bwd_dx_kernel)[(S * B,)](
+        grad_output,
+        x,
+        mean,
+        rstd,
+        scale,
+        dx,
+        S,
+        B,
+        H,
+        grad_output.stride(1),
+        x.stride(1),
+        scale.stride(0),
+        BLOCK_H=BLOCK_H,
+        OUT_DTYPE=out_dtype,
+    )
+    return dx, dscale, dshift
+
+
+@_triton_fused_ln_modulate_backward_op.register_fake
+def _triton_fused_ln_modulate_backward_fake(grad_output, x, mean, rstd, scale):
+    dx = torch.empty_like(x)
+    dscale = torch.empty_like(scale)
+    dshift = torch.empty_like(scale)
+    return dx, dscale, dshift
+
+
+def _triton_fused_ln_modulate_setup_context(ctx, inputs, output):
+    x, scale, _shift, _eps = inputs
+    _out, mean, rstd = output
+    ctx.save_for_backward(x, mean, rstd, scale)
+
+
+def _triton_fused_ln_modulate_backward(ctx, grad_output, _grad_mean, _grad_rstd):
+    x, mean, rstd, scale = ctx.saved_tensors
+    dx, dscale, dshift = _triton_fused_ln_modulate_backward_op(
+        grad_output,
+        x,
+        mean,
+        rstd,
+        scale,
+    )
+    return dx, dscale, dshift, None
+
+
+_triton_fused_ln_modulate.register_autograd(
+    _triton_fused_ln_modulate_backward,
+    setup_context=_triton_fused_ln_modulate_setup_context,
+)
+
+
+def _fused_eager_ln_modulate(
+    norm_module: nn.Module,
+    x: Tensor,
+    scale: Tensor,
+    shift: Tensor,
+    use_triton_ops: bool = False,
+) -> Tensor:
+    """Fused LayerNorm + modulate via a single Triton kernel.
+
+    Computes norm(x) * (1 + scale) + shift in one pass, avoiding the
+    intermediate ln_out write/read to DRAM.  Only valid for
+    elementwise_affine=False norms (no learnable weight/bias).
+    """
+    _ln_mod_fn = _triton_fused_ln_modulate if use_triton_ops else _opaque_fused_ln_modulate
+    out, _mean, _rstd = _ln_mod_fn(
+        x,
+        scale,
+        shift,
+        norm_module.eps,
+    )
+    return out
+
+
+class RMSNorm(nn.Module):
+    """
+    Root Mean Square Layer Normalization.
+
+    Formula: RMSNorm(x) = (x / sqrt(mean(x^2) + eps)) * weight
+
+    Args:
+        hidden_size: Size of the normalized dimension
+        config: Transformer configuration (for compatibility, not used)
+        eps: Small constant for numerical stability (default: 1e-6)
+
+    Reference:
+        - "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019)
+        - Adapted from NeMo's DiT implementation
+    """
+
+    def __init__(self, hidden_size: int, config=None, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def _norm(self, x: Tensor) -> Tensor:
+        """
+        Compute RMS normalization.
+
+        Args:
+            x: Input tensor
+
+        Returns:
+            Normalized tensor (before scaling)
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass: Normalize input.
+
+        Args:
+            x: Input tensor [..., hidden_size]
+
+        Returns:
+            Normalized and scaled tensor [..., hidden_size]
+        """
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class AdaLN(MegatronModule):
+    """
+    Adaptive Layer Normalization for DiT (Diffusion Transformer).
+
+    Conditions layer normalization on timestep embeddings via learned scale,
+    shift, and gate parameters. Projects conditioning through SiLU + Linear
+    into n_adaln_chunks modulation parameters.
+
+    Args:
+        config: Transformer configuration
+        n_adaln_chunks: Number of modulation chunks (default: 9, i.e. (shift, scale, gate) x 3)
+        norm: Normalization layer class (default: nn.LayerNorm)
+        modulation_bias: Whether to use bias in modulation projection (default: False)
+        use_second_norm: Whether to use a second normalization layer (default: False)
+        init_method: Initialization function for the modulation projection
+            weight (default: zero-init via ``nn.init.zeros_``). Pass
+            ``nn.init.normal_`` only when matching NeMo's exact RNG draw
+            sequence -- Flux's ``init_weights()`` re-zeroes these weights
+            after construction, so the normal_ draw only serves to advance
+            the CUDA RNG by the same amount as NeMo's init for
+            cross-framework convergence comparison. Other consumers
+            (test harnesses, downstream DiT models that reuse AdaLN) get
+            the same observable zero-init either way and should leave
+            this at the default.
+
+    Reference:
+        - "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023)
+        - Adapted from NeMo's DiT implementation
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        n_adaln_chunks: int = 9,
+        norm: type = nn.LayerNorm,
+        modulation_bias: bool = False,
+        use_second_norm: bool = False,
+        init_method=nn.init.zeros_,
+    ):
+        super().__init__(config)
+
+        # Layer normalization (without affine parameters - scale/shift come from conditioning)
+        self.ln = norm(config.hidden_size, elementwise_affine=False, eps=config.layernorm_epsilon)
+
+        self.n_adaln_chunks = n_adaln_chunks
+
+        # Modulation network: conditioning -> (scale, shift, gate, ...)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            ColumnParallelLinear(
+                config.hidden_size,
+                n_adaln_chunks * config.hidden_size,
+                config=config,
+                init_method=init_method,
+                bias=modulation_bias,
+                gather_output=True,
+            ),
+        )
+
+        self.use_second_norm = use_second_norm
+        if use_second_norm:
+            self.ln2 = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
+
+        # Mark weight as sequence parallel if needed
+        setattr(self.adaLN_modulation[-1].weight, "sequence_parallel", config.sequence_parallel)
+
+        self._adaln_plain_ops = getattr(config, "adaln_plain_ops", False)
+        self._use_triton_ops = getattr(config, "use_triton_ops", False)
+
+        if self._adaln_plain_ops:
+            self.use_fused_ln_modulate = False
+        else:
+            self.use_fused_ln_modulate = True
+
+        adaln_always_jit = getattr(config, "adaln_always_jit_fuser", False)
+        if adaln_always_jit or not getattr(config, "enable_torch_compile", False):
+            AdaLN._apply_jit_fuser()
+
+    @classmethod
+    def _apply_jit_fuser(cls):
+        """Apply @jit_fuser to class methods once (shared across all instances)."""
+        if getattr(cls, "_jit_fuser_applied", False):
+            return
+        cls.forward = jit_fuser(cls.forward)
+        cls.modulate = jit_fuser(cls.modulate)
+        cls.scale_add = jit_fuser(cls.scale_add)
+        cls.modulated_layernorm = jit_fuser(cls.modulated_layernorm)
+        cls.scaled_modulated_layernorm = jit_fuser(cls.scaled_modulated_layernorm)
+        cls._jit_fuser_applied = True
+
+    def forward(self, timestep_emb: Tensor) -> Tuple[Tensor, ...]:
+        """
+        Generate modulation parameters from timestep embedding.
+
+        Args:
+            timestep_emb: Timestep embeddings [B, hidden_size]
+
+        Returns:
+            Tuple of n_adaln_chunks tensors, each [B, hidden_size]
+        """
+        output, bias = self.adaLN_modulation(timestep_emb)
+        if bias is not None:
+            output = output + bias
+        return output.chunk(self.n_adaln_chunks, dim=-1)
+
+    def modulate(self, x: Tensor, shift: Tensor, scale: Tensor) -> Tensor:
+        """
+        Apply adaptive modulation: x * (1 + scale) + shift.
+
+        Args:
+            x: Input tensor [B, ..., hidden_size]
+            shift: Shift parameter [B, hidden_size]
+            scale: Scale parameter [B, hidden_size]
+
+        Returns:
+            Modulated tensor [B, ..., hidden_size]
+        """
+        return x * (1 + scale) + shift
+
+    def scale_add(self, residual: Tensor, x: Tensor, gate: Tensor) -> Tensor:
+        """
+        Gated residual addition: residual + gate * x.
+
+        Args:
+            residual: Residual connection [B, ..., hidden_size]
+            x: Input to add [B, ..., hidden_size]
+            gate: Gate parameter [B, hidden_size]
+
+        Returns:
+            Combined tensor [B, ..., hidden_size]
+        """
+        return residual + gate * x
+
+    def modulated_layernorm(self, x: Tensor, shift: Tensor, scale: Tensor, layernorm_idx: int = 0) -> Tensor:
+        """
+        Apply layer normalization followed by adaptive modulation.
+
+        Args:
+            x: Input tensor [B, ..., hidden_size]
+            shift: Shift parameter [B, hidden_size]
+            scale: Scale parameter [B, hidden_size]
+            layernorm_idx: Which layer norm to use (0 or 1, if use_second_norm=True)
+
+        Returns:
+            Normalized and modulated tensor [B, ..., hidden_size]
+        """
+        # Select appropriate layer norm
+        if self.use_second_norm and layernorm_idx == 1:
+            layernorm = self.ln2
+        else:
+            layernorm = self.ln
+
+        if self._adaln_plain_ops:
+            input_layernorm_output = layernorm(x).type_as(x)
+            return self.modulate(input_layernorm_output, shift, scale)
+
+        if self.use_fused_ln_modulate:
+            return _fused_eager_ln_modulate(layernorm, x, scale, shift, self._use_triton_ops)
+
+        input_layernorm_output = layernorm(x).type_as(x)
+        return self.modulate(input_layernorm_output, shift, scale)
+
+    def scaled_modulated_layernorm(
+        self,
+        residual: Tensor,
+        x: Tensor,
+        gate: Tensor,
+        shift: Tensor,
+        scale: Tensor,
+        layernorm_idx: int = 0,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Combined operation: gated residual addition + modulated layer norm.
+
+        This is a common pattern in DiT: add gated residual, then normalize and modulate.
+
+        Args:
+            residual: Residual connection [B, ..., hidden_size]
+            x: Input to add [B, ..., hidden_size]
+            gate: Gate parameter [B, hidden_size]
+            shift: Shift parameter [B, hidden_size]
+            scale: Scale parameter [B, hidden_size]
+            layernorm_idx: Which layer norm to use
+
+        Returns:
+            Tuple of (hidden_states, shifted_pre_mlp_layernorm_output)
+        """
+        # Gated residual addition
+        hidden_states = self.scale_add(residual, x, gate)
+
+        # Apply modulated layer normalization
+        shifted_pre_mlp_layernorm_output = self.modulated_layernorm(
+            hidden_states, shift, scale, layernorm_idx
+        )
+
+        return hidden_states, shifted_pre_mlp_layernorm_output
+
+
+class AdaLNContinuous(MegatronModule):
+    """
+    Continuous Adaptive Layer Normalization for Flux.
+
+    Simpler variant of AdaLN that only produces scale and shift (no gating).
+    Formula: norm(x) * (1 + scale) + shift
+
+    Args:
+        config: Transformer configuration
+        conditioning_embedding_dim: Dimension of conditioning embedding input
+        modulation_bias: Whether to use bias in modulation layers (default: True)
+        norm_type: 'layer_norm' (default: 'layer_norm')
+
+    Note:
+        Uses Megatron's sequence-first format [S, B, D].
+
+    Reference:
+        - Flux Paper: "Flux: A Scalable Diffusion Model"
+        - Adapted from NeMo's DiT implementation
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        conditioning_embedding_dim: int,
+        modulation_bias: bool = True,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__(config)
+
+        # Modulation network: conditioning -> (scale, shift)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(conditioning_embedding_dim, config.hidden_size * 2, bias=modulation_bias),
+        )
+
+        # Normalization layer
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(
+                config.hidden_size, elementwise_affine=False, eps=1e-6, bias=modulation_bias
+            )
+        else:
+            raise ValueError(f"Unknown normalization type: {norm_type}")
+
+        self._adaln_plain_ops = getattr(config, "adaln_plain_ops", False)
+
+        if self._adaln_plain_ops:
+            self.use_fused_ln_modulate = False
+        else:
+            self.use_fused_ln_modulate = True
+
+        self._use_triton_ops = getattr(config, "use_triton_ops", False)
+
+    def forward(self, x: Tensor, conditioning_embedding: Tensor) -> Tensor:
+        """
+        Forward pass: Apply continuous adaptive normalization.
+
+        Args:
+            x: Input tensor [seq_len, B, hidden_size] (sequence-first format)
+            conditioning_embedding: Conditioning signal [B, conditioning_embedding_dim]
+
+        Returns:
+            Normalized and modulated tensor [seq_len, B, hidden_size]
+        """
+        # Generate scale and shift from conditioning
+        emb = self.adaLN_modulation(conditioning_embedding)  # [B, 2 * hidden_size]
+        # NeMo convention: first half is scale, second half is shift.
+        # NOTE: Checkpoints trained before this fix had (shift, scale) order and
+        # need the two halves of norm_out.adaLN_modulation weight/bias swapped.
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        if self._adaln_plain_ops:
+            x = self.norm(x) * (1 + scale) + shift
+            return x
+
+        if self.use_fused_ln_modulate:
+            return _fused_eager_ln_modulate(self.norm, x, scale, shift, self._use_triton_ops)
+
+        ln_out = self.norm(x)
+        _mod_fn = _triton_modulate if self._use_triton_ops else _opaque_modulate
+        return _mod_fn(ln_out, scale, shift)
diff --git a/primus/backends/megatron/core/models/diffusion/flux/__init__.py b/primus/backends/megatron/core/models/diffusion/flux/__init__.py
new file mode 100644
index 000000000..5ea3e2236
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/__init__.py
@@ -0,0 +1,96 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Flux diffusion model components.
+
+This module provides all components needed for the Flux architecture:
+- Model: Flux (main model class)
+- Configuration: FluxConfig
+- Layers: EmbedND (3D RoPE position embedding)
+- Attention: JointSelfAttention, FluxSingleAttention
+- Layer specs: MMDiTLayer, FluxSingleTransformerBlock
+
+Quick Start - Model Creation:
+    >>> from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig
+    >>>
+    >>> # Create and configure model
+    >>> config = FluxConfig.flux_12b()
+    >>> model = Flux(config=config)
+    >>>
+    >>> # Load checkpoint (native Primus format)
+    >>> model.load_checkpoint("flux_12b.safetensors")
+
+Checkpoint Conversion:
+    >>> from primus.backends.megatron.core.models.diffusion.flux import convert_hf_checkpoint, FluxConfig
+    >>>
+    >>> # Convert HuggingFace checkpoint to Primus format
+    >>> config = FluxConfig.flux_12b()
+    >>> primus_sd = convert_hf_checkpoint(
+    ...     "black-forest-labs/FLUX.1-dev/transformer",
+    ...     flux_config=config,
+    ...     save_to="primus_flux_12b.safetensors"
+    ... )
+"""
+
+from primus.backends.megatron.core.models.diffusion.flux.attention import (
+    FluxSingleAttention,
+    JointSelfAttention,
+    JointSelfAttentionSubmodules,
+)
+from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+    convert_hf_checkpoint,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    FluxSingleTransformerBlock,
+    MMDiTLayer,
+    get_flux_double_transformer_spec_for_backend,
+    get_flux_layer_spec,
+    get_flux_single_transformer_spec_for_backend,
+)
+from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND, rope
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    generate_text_position_ids,
+    pack_latents,
+    unpack_latents,
+)
+
+
+# LAZY IMPORT: Don't import Flux here to avoid early TransformerBlock import
+def __getattr__(name):
+    """Lazy import for Flux to avoid early TransformerBlock import."""
+    if name == "Flux":
+        from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+
+        return Flux
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+__all__ = [
+    # Model
+    "Flux",
+    # Configuration
+    "FluxConfig",
+    # Layers
+    "EmbedND",
+    "rope",
+    # Attention
+    "JointSelfAttention",
+    "FluxSingleAttention",
+    "JointSelfAttentionSubmodules",
+    # Layer specs
+    "MMDiTLayer",
+    "FluxSingleTransformerBlock",
+    "get_flux_double_transformer_spec_for_backend",
+    "get_flux_single_transformer_spec_for_backend",
+    "get_flux_layer_spec",
+    # Utils
+    "pack_latents",
+    "unpack_latents",
+    "generate_image_position_ids",
+    "generate_text_position_ids",
+    # Checkpoint conversion
+    "convert_hf_checkpoint",
+]
diff --git a/primus/backends/megatron/core/models/diffusion/flux/attention.py b/primus/backends/megatron/core/models/diffusion/flux/attention.py
new file mode 100644
index 000000000..b444b527c
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/attention.py
@@ -0,0 +1,597 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Flux attention mechanisms.
+
+This module implements specialized attention for Flux's MMDiT architecture:
+    - JointSelfAttention: Processes concatenated image + text tokens
+    - FluxSingleAttention: Processes image tokens only
+    - JointSelfAttentionSubmodules: Configuration for joint attention
+
+These implementations follow Megatron-Core's attention patterns with
+customizations for diffusion model conditioning.
+
+Reference:
+    - MMDiT: "Scaling Rectified Flow Transformers"
+    - Megatron-Core: megatron.core.transformer.attention
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from megatron.core.models.common.embeddings.rotary_pos_embedding import (
+    apply_rotary_pos_emb,
+)
+from megatron.core.transformer.attention import (
+    Attention,
+    SelfAttention,
+    SelfAttentionSubmodules,
+)
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from torch import Tensor
+
+try:
+    from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
+except ImportError:
+    SplitAlongDim = None
+
+
+@dataclass
+class JointSelfAttentionSubmodules:
+    """
+    Submodules configuration for Joint Self-Attention layer.
+
+    Joint attention processes both image and text tokens together (MMDiT architecture).
+    It requires separate QKV projections for image and text (context) streams.
+
+    Attributes:
+        linear_qkv: QKV projection for main stream (image tokens)
+        added_linear_qkv: QKV projection for added stream (text/context tokens)
+        core_attention: Core attention computation module
+        linear_proj: Output projection for main stream
+        q_layernorm: Optional layer norm for queries (main stream)
+        k_layernorm: Optional layer norm for keys (main stream)
+        added_q_layernorm: Optional layer norm for queries (added stream)
+        added_k_layernorm: Optional layer norm for keys (added stream)
+
+    Note:
+        Flux uses RMSNorm for Q/K normalization to improve training stability.
+
+    Reference:
+        - Paper: "Scaling Rectified Flow Transformers"
+    """
+
+    linear_qkv: Union[ModuleSpec, type] = None
+    added_linear_qkv: Union[ModuleSpec, type] = None
+    core_attention: Union[ModuleSpec, type] = None
+    linear_proj: Union[ModuleSpec, type] = None
+    q_layernorm: Union[ModuleSpec, type] = None
+    k_layernorm: Union[ModuleSpec, type] = None
+    added_q_layernorm: Union[ModuleSpec, type] = None
+    added_k_layernorm: Union[ModuleSpec, type] = None
+
+
+class JointSelfAttention(Attention):
+    """
+    Joint Self-Attention for MMDiT (Multimodal Diffusion Transformer).
+
+    Processes two token streams jointly -- main (image) and added (text) --
+    by projecting each through separate QKV layers, concatenating, computing
+    joint attention, then splitting back. This enables cross-modal interaction
+    in Flux's "double blocks".
+
+    Args:
+        config: Transformer configuration
+        submodules: JointSelfAttentionSubmodules with layer specifications
+        layer_number: Layer index in the model
+        attn_mask_type: Type of attention mask (default: padding)
+        context_pre_only: If True, only compute Q/K/V for context (default: False)
+
+    Input/Output:
+        hidden_states [seq_main, B, H] + additional_hidden_states [seq_added, B, H]
+        -> (main_output, added_output) with same shapes
+
+    Reference:
+        - "Scaling Rectified Flow Transformers"
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: JointSelfAttentionSubmodules,
+        layer_number: int,
+        attn_mask_type: AttnMaskType = AttnMaskType.padding,
+        context_pre_only: bool = False,
+        **kwargs,
+    ):
+        # Use RMSNorm for Q/K normalization (improves stability)
+        config.normalization = "RMSNorm"
+
+        super().__init__(
+            config=config,
+            submodules=submodules,
+            layer_number=layer_number,
+            attn_mask_type=attn_mask_type,
+            attention_type="self",
+            **kwargs,
+        )
+
+        # QKV projection for main stream (image tokens)
+        self.linear_qkv = build_module(
+            submodules.linear_qkv,
+            self.config.hidden_size,
+            self.query_projection_size + 2 * self.kv_projection_size,
+            config=self.config,
+            init_method=self.config.init_method,
+            gather_output=False,
+            bias=self.config.add_bias_linear or self.config.add_qkv_bias,
+            skip_bias_add=False,
+            is_expert=False,
+            tp_comm_buffer_name="qkv",
+        )
+
+        # QKV projection for added stream (text tokens)
+        if submodules.added_linear_qkv is not None:
+            self.added_linear_qkv = build_module(
+                submodules.added_linear_qkv,
+                self.config.hidden_size,
+                self.query_projection_size + 2 * self.kv_projection_size,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=self.config.add_qkv_bias,
+                skip_bias_add=False,
+                is_expert=False,
+                tp_comm_buffer_name="qkv",
+            )
+
+        # Output projection for added stream (text tokens)
+        if not context_pre_only:
+            self.added_linear_proj = build_module(
+                submodules.linear_proj,
+                self.query_projection_size,
+                self.config.hidden_size,
+                config=self.config,
+                init_method=self.config.output_layer_init_method,
+                bias=self.config.add_bias_linear,
+                input_is_parallel=True,
+                skip_bias_add=True,
+                is_expert=False,
+                tp_comm_buffer_name="proj",
+            )
+
+        if (
+            not context_pre_only
+            and getattr(self.config, "use_dual_fp8_output_projection", False)
+            and hasattr(self.linear_proj, "_fp8_config")
+        ):
+            from primus_turbo.pytorch.core.low_precision import ScalingGranularity
+
+            if self.linear_proj._fp8_config.granularity == ScalingGranularity.TENSORWISE:
+                from primus_turbo.pytorch.core.backend import BackendType
+
+                from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+                    DualFP8LinearTensorwiseFunction,
+                    _get_fp8_dtype,
+                )
+
+                self._dual_fp8_fn = DualFP8LinearTensorwiseFunction
+                cfg = self.linear_proj._fp8_config
+                self._dual_fp8_fwd_dtype = _get_fp8_dtype(cfg.format, is_fwd=True)
+                self._dual_fp8_bwd_dtype = _get_fp8_dtype(cfg.format, is_fwd=False)
+                self._dual_fp8_gran_value = ScalingGranularity.TENSORWISE.value
+                self._dual_fp8_backend_value = BackendType.HIPBLASLT.value
+
+        # Optional Q/K layer normalization for main stream
+        if submodules.q_layernorm is not None:
+            self.q_layernorm = build_module(
+                submodules.q_layernorm,
+                hidden_size=self.hidden_size_per_attention_head,
+                config=self.config,
+                eps=self.config.layernorm_epsilon,
+            )
+        else:
+            self.q_layernorm = None
+
+        if submodules.k_layernorm is not None:
+            self.k_layernorm = build_module(
+                submodules.k_layernorm,
+                hidden_size=self.hidden_size_per_attention_head,
+                config=self.config,
+                eps=self.config.layernorm_epsilon,
+            )
+        else:
+            self.k_layernorm = None
+
+        # Optional Q/K layer normalization for added stream
+        if submodules.added_q_layernorm is not None:
+            self.added_q_layernorm = build_module(
+                submodules.added_q_layernorm,
+                hidden_size=self.hidden_size_per_attention_head,
+                config=self.config,
+                eps=self.config.layernorm_epsilon,
+            )
+        else:
+            self.added_q_layernorm = None
+
+        if submodules.added_k_layernorm is not None:
+            self.added_k_layernorm = build_module(
+                submodules.added_k_layernorm,
+                hidden_size=self.hidden_size_per_attention_head,
+                config=self.config,
+                eps=self.config.layernorm_epsilon,
+            )
+        else:
+            self.added_k_layernorm = None
+
+    def _split_qkv(self, mixed_qkv: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Split mixed QKV tensor into separate Q, K, V tensors.
+
+        Args:
+            mixed_qkv: Combined QKV tensor [seq, batch, hidden]
+
+        Returns:
+            Tuple of (query, key, value) tensors
+        """
+        # Reshape: [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn]
+        new_tensor_shape = mixed_qkv.size()[:-1] + (
+            self.num_query_groups_per_partition,
+            (
+                (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2)
+                * self.hidden_size_per_attention_head
+            ),
+        )
+        mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+        # Define split sizes for Q, K, V
+        split_arg_list = [
+            (
+                self.num_attention_heads_per_partition
+                // self.num_query_groups_per_partition
+                * self.hidden_size_per_attention_head
+            ),
+            self.hidden_size_per_attention_head,
+            self.hidden_size_per_attention_head,
+        ]
+
+        # Split tensor
+        if SplitAlongDim is not None:
+            # Use Transformer Engine's optimized split if available
+            (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list)
+        else:
+            # Fallback to PyTorch split
+            (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3)
+
+        # Reshape query: [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
+        query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
+
+        return query, key, value
+
+    def get_query_key_value_tensors(
+        self, hidden_states: Tensor, key_value_states: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Derive Q, K, V tensors from main stream hidden states.
+
+        Args:
+            hidden_states: Main stream tokens [seq, batch, hidden]
+            key_value_states: Not used for self-attention
+
+        Returns:
+            Tuple of (query, key, value) tensors
+        """
+        # Project to QKV: [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)]
+        mixed_qkv, _ = self.linear_qkv(hidden_states)
+
+        # Split into Q, K, V
+        query, key, value = self._split_qkv(mixed_qkv)
+
+        # Apply optional Q/K normalization
+        if self.q_layernorm is not None:
+            query = self.q_layernorm(query)
+
+        if self.k_layernorm is not None:
+            key = self.k_layernorm(key)
+
+        return query, key, value
+
+    def get_added_query_key_value_tensors(
+        self, added_hidden_states: Tensor, key_value_states: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Derive Q, K, V tensors from added stream (text) hidden states.
+
+        Args:
+            added_hidden_states: Added stream tokens [seq, batch, hidden]
+            key_value_states: Not used for self-attention
+
+        Returns:
+            Tuple of (query, key, value) tensors
+        """
+        # Project to QKV
+        mixed_qkv, _ = self.added_linear_qkv(added_hidden_states)
+
+        # Split into Q, K, V
+        query, key, value = self._split_qkv(mixed_qkv)
+
+        # Apply optional Q/K normalization
+        if self.added_q_layernorm is not None:
+            query = self.added_q_layernorm(query)
+
+        if self.added_k_layernorm is not None:
+            key = self.added_k_layernorm(key)
+
+        return query, key, value
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor],
+        key_value_states: Optional[Tensor] = None,
+        inference_params=None,
+        rotary_pos_emb=None,
+        packed_seq_params=None,
+        additional_hidden_states: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Forward pass: Joint attention over image and text tokens.
+
+        Args:
+            hidden_states: Main stream (image) tokens [seq_main, batch, hidden]
+            attention_mask: Attention mask
+            key_value_states: Not used for self-attention
+            inference_params: Parameters for inference (e.g., KV cache)
+            rotary_pos_emb: RoPE position embeddings
+            packed_seq_params: Parameters for packed sequences
+            additional_hidden_states: Added stream (text) tokens [seq_added, batch, hidden]
+
+        Returns:
+            Tuple of (main_output, added_output):
+                - main_output: Processed main stream [seq_main, batch, hidden]
+                - added_output: Processed added stream [seq_added, batch, hidden]
+        """
+        # Ensure rotary_pos_emb is a tuple for Q and K
+        if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):
+            rotary_pos_emb = (rotary_pos_emb,) * 2
+
+        # Get Q, K, V for both streams
+        query, key, value = self.get_query_key_value_tensors(hidden_states)
+        added_query, added_key, added_value = self.get_added_query_key_value_tensors(additional_hidden_states)
+
+        # Concatenate streams: [added; main]
+        query = torch.cat([added_query, query], dim=0)
+        key = torch.cat([added_key, key], dim=0)
+        value = torch.cat([added_value, value], dim=0)
+
+        # Adjust for inference (KV caching, etc.)
+        query, key, value, rotary_pos_emb, attn_mask_type, *_ = self._adjust_key_value_for_inference(
+            inference_params, query, key, value, rotary_pos_emb
+        )
+
+        # Handle packed sequences
+        if packed_seq_params is not None:
+            query = query.squeeze(1)
+            key = key.squeeze(1)
+            value = value.squeeze(1)
+
+        # Apply RoPE position embeddings
+        if rotary_pos_emb is not None:
+            q_pos_emb, k_pos_emb = rotary_pos_emb
+
+            cu_seqlens_q = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None
+            cu_seqlens_kv = packed_seq_params.cu_seqlens_kv if packed_seq_params is not None else None
+
+            query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
+            key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
+
+        # Core attention computation
+        if self.checkpoint_core_attention and self.training:
+            core_attn_out = self._checkpointed_attention_forward(
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type=attn_mask_type,
+                packed_seq_params=packed_seq_params,
+            )
+        else:
+            core_attn_out = self.core_attention(
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type=attn_mask_type,
+                packed_seq_params=packed_seq_params,
+            )
+
+        # Handle packed sequences output
+        if packed_seq_params is not None:
+            # Reshape: (t, np, hn) -> (t, b=1, h=np*hn)
+            core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+        # Split output back into added and main streams
+        encoder_attention_output = core_attn_out[: additional_hidden_states.shape[0], :, :]
+        attention_output = core_attn_out[additional_hidden_states.shape[0] :, :, :]
+
+        # Project outputs
+        if hasattr(self, "_dual_fp8_fn"):
+            from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+                _extract_fp8_weight,
+            )
+
+            w_fp8_a, w_scale_a = _extract_fp8_weight(
+                self.linear_proj.weight,
+                self._dual_fp8_fwd_dtype,
+            )
+            w_fp8_b, w_scale_b = _extract_fp8_weight(
+                self.added_linear_proj.weight,
+                self._dual_fp8_fwd_dtype,
+            )
+            result = self._dual_fp8_fn.apply(
+                attention_output,
+                self.linear_proj.weight,
+                w_fp8_a,
+                w_scale_a,
+                encoder_attention_output,
+                self.added_linear_proj.weight,
+                w_fp8_b,
+                w_scale_b,
+                self._dual_fp8_fwd_dtype,
+                self._dual_fp8_bwd_dtype,
+                self._dual_fp8_gran_value,
+                self._dual_fp8_backend_value,
+            )
+            output, encoder_output = result[0], result[1]
+            if self.linear_proj.bias is not None:
+                output = output + self.linear_proj.bias
+            if self.added_linear_proj.bias is not None:
+                encoder_output = encoder_output + self.added_linear_proj.bias
+        else:
+            output, bias = self.linear_proj(attention_output)
+            encoder_output, encoder_bias = self.added_linear_proj(encoder_attention_output)
+            output = output + bias
+            encoder_output = encoder_output + encoder_bias
+
+        return output, encoder_output
+
+
+class FluxSingleAttention(SelfAttention):
+    """
+    Single-stream Self-Attention for Flux (image tokens only).
+
+    Standard self-attention without cross-modal interaction. Used in Flux's
+    "single blocks" after the joint MMDiT blocks.
+
+    Args:
+        config: Transformer configuration
+        submodules: SelfAttentionSubmodules with layer specifications
+        layer_number: Layer index in the model
+        attn_mask_type: Type of attention mask (default: padding)
+        cp_comm_type: Communication type for context parallelism
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: SelfAttentionSubmodules,
+        layer_number: int,
+        attn_mask_type: AttnMaskType = AttnMaskType.padding,
+        cp_comm_type: Optional[str] = None,
+        **kwargs,
+    ):
+        # Use RMSNorm for Q/K normalization
+        config.normalization = "RMSNorm"
+
+        super().__init__(
+            config=config,
+            submodules=submodules,
+            layer_number=layer_number,
+            attn_mask_type=attn_mask_type,
+            cp_comm_type=cp_comm_type,
+            **kwargs,
+        )
+
+        # The original Flux proj_out (Diffusers) / linear2 (TorchTitan) is a single fused
+        # projection with one bias. Megatron splits it into linear_proj + linear_fc2, so the
+        # bias only needs to be on one path (linear_fc2) to preserve mathematical equivalence.
+        self.linear_proj = build_module(
+            submodules.linear_proj,
+            self.query_projection_size,
+            self.config.hidden_size,
+            config=self.config,
+            init_method=self.config.output_layer_init_method,
+            bias=False,
+            input_is_parallel=True,
+            skip_bias_add=True,
+            is_expert=False,
+            tp_comm_buffer_name="proj",
+        )
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor],
+        key_value_states: Optional[Tensor] = None,
+        inference_params=None,
+        rotary_pos_emb=None,
+        packed_seq_params=None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """
+        Forward pass: Self-attention on image tokens.
+
+        Args:
+            hidden_states: Image tokens [seq, batch, hidden]
+            attention_mask: Attention mask
+            key_value_states: Not used for self-attention
+            inference_params: Parameters for inference (e.g., KV cache)
+            rotary_pos_emb: RoPE position embeddings
+            packed_seq_params: Parameters for packed sequences
+
+        Returns:
+            Tuple of (output, bias):
+                - output: Projected attention output [seq, batch, hidden]
+                - bias: Projection bias (None when linear_proj has bias=False)
+        """
+        # Ensure rotary_pos_emb is a tuple for Q and K
+        if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):
+            rotary_pos_emb = (rotary_pos_emb,) * 2
+
+        # Get Q, K, V
+        query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
+
+        # Adjust for inference
+        query, key, value, rotary_pos_emb, attn_mask_type, *_ = self._adjust_key_value_for_inference(
+            inference_params, query, key, value, rotary_pos_emb
+        )
+
+        # Handle packed sequences
+        if packed_seq_params is not None:
+            query = query.squeeze(1)
+            key = key.squeeze(1)
+            value = value.squeeze(1)
+
+        # Apply RoPE position embeddings
+        if rotary_pos_emb is not None:
+            q_pos_emb, k_pos_emb = rotary_pos_emb
+
+            cu_seqlens_q = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None
+            cu_seqlens_kv = packed_seq_params.cu_seqlens_kv if packed_seq_params is not None else None
+
+            query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
+            key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
+
+        # Core attention computation
+        if self.checkpoint_core_attention and self.training:
+            core_attn_out = self._checkpointed_attention_forward(
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type=attn_mask_type,
+                packed_seq_params=packed_seq_params,
+            )
+        else:
+            core_attn_out = self.core_attention(
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type=attn_mask_type,
+                packed_seq_params=packed_seq_params,
+            )
+
+        # Handle packed sequences output
+        if packed_seq_params is not None:
+            # Reshape: (t, np, hn) -> (t, b=1, h=np*hn)
+            core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+        # Project output (return both output and bias for skip_bias_add pattern)
+        output, bias = self.linear_proj(core_attn_out)
+
+        return output, bias
diff --git a/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py b/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py
new file mode 100644
index 000000000..a2715cb7b
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py
@@ -0,0 +1,881 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+"""
+HuggingFace to Primus Flux checkpoint converter.
+
+Converts HuggingFace Diffusers Flux transformer checkpoints to
+Primus/Megatron-Core compatible format.
+
+Key Conversion:
+    - HuggingFace: Separate double_blocks, single_blocks
+    - Primus: Unified transformer.layers.{0-N} with TransformerBlock
+
+    This reflects Primus's architectural enhancement using heterogeneous
+    layer specifications in a single TransformerBlock container.
+
+Usage:
+    from primus.backends.megatron.core.models.diffusion.flux import convert_hf_checkpoint
+
+    primus_state_dict = convert_hf_checkpoint(
+        checkpoint_path="black-forest-labs/FLUX.1-dev",
+        flux_config=config,
+        save_to="primus_flux_12b.safetensors"
+    )
+
+Reference:
+    - HuggingFace Diffusers checkpoint format
+    - Megatron-Core TransformerBlock architecture
+"""
+
+import logging
+import os
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import torch
+from safetensors.torch import load_file as load_safetensors
+from safetensors.torch import save_file as save_safetensors
+
+logger = logging.getLogger(__name__)
+
+
+def _fuse_qkv_weights(
+    config,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    v_weight: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Fuse separate Q, K, V weight matrices into Megatron's fused QKV format.
+
+    Megatron-Core uses grouped-query attention (GQA) format where Q, K, V
+    are interleaved per attention group:
+        [Q_group0, K_group0, V_group0, Q_group1, K_group1, V_group1, ...]
+
+    Args:
+        config: FluxConfig with num_attention_heads, num_query_groups
+        q_weight: Query weights [hidden_size, hidden_size]
+        k_weight: Key weights [hidden_size, hidden_size]
+        v_weight: Value weights [hidden_size, hidden_size]
+
+    Returns:
+        Fused QKV weights [(heads_per_group + 2) * num_query_groups * head_size, hidden_size]
+        in GQA interleaved format
+
+    Reference:
+        - Megatron-Core: Grouped Query Attention patterns
+    """
+    head_num = config.num_attention_heads
+    num_query_groups = getattr(config, "num_query_groups", head_num)
+    heads_per_group = head_num // num_query_groups
+    hidden_size = config.hidden_size
+    head_size = hidden_size // head_num
+
+    # Get input shape from q_weight
+    old_tensor_shape = q_weight.size()
+
+    # Reshape to [num_heads, head_size, ...]
+    new_q_tensor_shape = (head_num, head_size) + old_tensor_shape[1:]
+    new_kv_tensor_shape = (num_query_groups, head_size) + old_tensor_shape[1:]
+
+    q = q_weight.view(*new_q_tensor_shape)
+    k = k_weight.view(*new_kv_tensor_shape)
+    v = v_weight.view(*new_kv_tensor_shape)
+
+    # Interleave by group: [Q_heads_for_group, K_group, V_group, ...]
+    qkv_list = []
+    for i in range(num_query_groups):
+        qkv_list.append(q[i * heads_per_group : (i + 1) * heads_per_group, :, :])
+        qkv_list.append(k[i : i + 1, :, :])
+        qkv_list.append(v[i : i + 1, :, :])
+
+    qkv = torch.cat(qkv_list, dim=0)
+
+    # Validate shape
+    if qkv.ndim != 3:
+        raise ValueError(f"Expected 3D QKV tensor, got shape {qkv.shape}")
+    expected_dim0 = (heads_per_group + 2) * num_query_groups
+    if qkv.shape[0] != expected_dim0:
+        raise ValueError(f"Expected QKV dim0 {expected_dim0}, got shape {qkv.shape}")
+    if qkv.shape[1] != head_size:
+        raise ValueError(f"Expected QKV head_size {head_size}, got shape {qkv.shape}")
+    if qkv.shape[2] != old_tensor_shape[1]:
+        raise ValueError(f"Expected QKV hidden dim {old_tensor_shape[1]}, got shape {qkv.shape}")
+
+    # Reshape to [fused_dim, hidden_size]
+    qkv = qkv.reshape(head_size * (head_num + 2 * num_query_groups), hidden_size)
+
+    return qkv
+
+
+def _fuse_qkv_bias(
+    config,
+    q_bias: torch.Tensor,
+    k_bias: torch.Tensor,
+    v_bias: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Fuse Q, K, V bias terms into Megatron's fused format.
+
+    Args:
+        config: FluxConfig with num_attention_heads, num_query_groups
+        q_bias: Query bias [hidden_size]
+        k_bias: Key bias [hidden_size]
+        v_bias: Value bias [hidden_size]
+
+    Returns:
+        Fused QKV bias [(heads_per_group + 2) * num_query_groups * head_size]
+        in GQA interleaved format
+
+    Reference:
+        - Megatron-Core: Grouped Query Attention patterns
+    """
+    head_num = config.num_attention_heads
+    num_query_groups = getattr(config, "num_query_groups", head_num)
+    heads_per_group = head_num // num_query_groups
+    hidden_size = config.hidden_size
+    head_size = hidden_size // head_num
+
+    # Reshape to [num_heads, head_size]
+    new_q_bias_shape = (head_num, head_size)
+    new_kv_bias_shape = (num_query_groups, head_size)
+
+    q = q_bias.view(*new_q_bias_shape)
+    k = k_bias.view(*new_kv_bias_shape)
+    v = v_bias.view(*new_kv_bias_shape)
+
+    # Interleave by group
+    qkv_bias_list = []
+    for i in range(num_query_groups):
+        qkv_bias_list.append(q[i * heads_per_group : (i + 1) * heads_per_group, :])
+        qkv_bias_list.append(k[i : i + 1, :])
+        qkv_bias_list.append(v[i : i + 1, :])
+
+    qkv_bias = torch.cat(qkv_bias_list, dim=0)
+    qkv_bias = qkv_bias.reshape(head_size * (head_num + 2 * num_query_groups))
+
+    return qkv_bias
+
+
+# Key mapping from HuggingFace to Primus
+# After TransformerBlock refactor: double_blocks and single_blocks are now transformer.layers[0-56]
+FLUX_KEY_MAPPING = {
+    "double_blocks": {
+        "norm1.linear.weight": "adaln.adaLN_modulation.1.weight",
+        "norm1.linear.bias": "adaln.adaLN_modulation.1.bias",
+        "norm1_context.linear.weight": "adaln_context.adaLN_modulation.1.weight",
+        "norm1_context.linear.bias": "adaln_context.adaLN_modulation.1.bias",
+        "attn.norm_q.weight": "self_attention.q_layernorm.weight",
+        "attn.norm_k.weight": "self_attention.k_layernorm.weight",
+        "attn.norm_added_q.weight": "self_attention.added_q_layernorm.weight",
+        "attn.norm_added_k.weight": "self_attention.added_k_layernorm.weight",
+        "attn.to_out.0.weight": "self_attention.linear_proj.weight",
+        "attn.to_out.0.bias": "self_attention.linear_proj.bias",
+        "attn.to_add_out.weight": "self_attention.added_linear_proj.weight",
+        "attn.to_add_out.bias": "self_attention.added_linear_proj.bias",
+        "ff.net.0.proj.weight": "mlp.linear_fc1.weight",
+        "ff.net.0.proj.bias": "mlp.linear_fc1.bias",
+        "ff.net.2.weight": "mlp.linear_fc2.weight",
+        "ff.net.2.bias": "mlp.linear_fc2.bias",
+        "ff_context.net.0.proj.weight": "context_mlp.linear_fc1.weight",
+        "ff_context.net.0.proj.bias": "context_mlp.linear_fc1.bias",
+        "ff_context.net.2.weight": "context_mlp.linear_fc2.weight",
+        "ff_context.net.2.bias": "context_mlp.linear_fc2.bias",
+    },
+    "single_blocks": {
+        "norm.linear.weight": "adaln.adaLN_modulation.1.weight",
+        "norm.linear.bias": "adaln.adaLN_modulation.1.bias",
+        "proj_mlp.weight": "mlp.linear_fc1.weight",
+        "proj_mlp.bias": "mlp.linear_fc1.bias",
+        "attn.norm_q.weight": "self_attention.q_layernorm.weight",
+        "attn.norm_k.weight": "self_attention.k_layernorm.weight",
+    },
+    # Root-level mappings
+    "norm_out.linear.bias": "norm_out.adaLN_modulation.1.bias",
+    "norm_out.linear.weight": "norm_out.adaLN_modulation.1.weight",
+    "proj_out.bias": "proj_out.bias",
+    "proj_out.weight": "proj_out.weight",
+    "time_text_embed.guidance_embedder.linear_1.bias": "guidance_embedding.in_layer.bias",
+    "time_text_embed.guidance_embedder.linear_1.weight": "guidance_embedding.in_layer.weight",
+    "time_text_embed.guidance_embedder.linear_2.bias": "guidance_embedding.out_layer.bias",
+    "time_text_embed.guidance_embedder.linear_2.weight": "guidance_embedding.out_layer.weight",
+    "x_embedder.bias": "img_embed.bias",
+    "x_embedder.weight": "img_embed.weight",
+    "time_text_embed.timestep_embedder.linear_1.bias": "timestep_embedding.time_embedding.in_layer.bias",
+    "time_text_embed.timestep_embedder.linear_1.weight": "timestep_embedding.time_embedding.in_layer.weight",
+    "time_text_embed.timestep_embedder.linear_2.bias": "timestep_embedding.time_embedding.out_layer.bias",
+    "time_text_embed.timestep_embedder.linear_2.weight": "timestep_embedding.time_embedding.out_layer.weight",
+    "context_embedder.bias": "txt_embed.bias",
+    "context_embedder.weight": "txt_embed.weight",
+    "time_text_embed.text_embedder.linear_1.bias": "vector_embedding.in_layer.bias",
+    "time_text_embed.text_embedder.linear_1.weight": "vector_embedding.in_layer.weight",
+    "time_text_embed.text_embedder.linear_2.bias": "vector_embedding.out_layer.bias",
+    "time_text_embed.text_embedder.linear_2.weight": "vector_embedding.out_layer.weight",
+}
+
+
+# Key mapping from Black Forest Labs native format to Primus format
+# Reference: diffusers Flux checkpoint conversion
+# This is the format used in official BFL releases (e.g., flux1-dev.safetensors, flux1-schnell.sft)
+BFL_KEY_MAPPING = {
+    # Root-level embeddings
+    "img_in.weight": "img_embed.weight",
+    "img_in.bias": "img_embed.bias",
+    "txt_in.weight": "txt_embed.weight",
+    "txt_in.bias": "txt_embed.bias",
+    # Timestep embedding (time_in → timestep_embedding.time_embedding)
+    "time_in.in_layer.weight": "timestep_embedding.time_embedding.in_layer.weight",
+    "time_in.in_layer.bias": "timestep_embedding.time_embedding.in_layer.bias",
+    "time_in.out_layer.weight": "timestep_embedding.time_embedding.out_layer.weight",
+    "time_in.out_layer.bias": "timestep_embedding.time_embedding.out_layer.bias",
+    # Vector embedding (vector_in → vector_embedding)
+    "vector_in.in_layer.weight": "vector_embedding.in_layer.weight",
+    "vector_in.in_layer.bias": "vector_embedding.in_layer.bias",
+    "vector_in.out_layer.weight": "vector_embedding.out_layer.weight",
+    "vector_in.out_layer.bias": "vector_embedding.out_layer.bias",
+    # Guidance embedding (guidance_in → guidance_embedding) - optional
+    "guidance_in.in_layer.weight": "guidance_embedding.in_layer.weight",
+    "guidance_in.in_layer.bias": "guidance_embedding.in_layer.bias",
+    "guidance_in.out_layer.weight": "guidance_embedding.out_layer.weight",
+    "guidance_in.out_layer.bias": "guidance_embedding.out_layer.bias",
+    # Final layer
+    "final_layer.linear.weight": "proj_out.weight",
+    "final_layer.linear.bias": "proj_out.bias",
+    "final_layer.adaLN_modulation.1.weight": "norm_out.adaLN_modulation.1.weight",
+    "final_layer.adaLN_modulation.1.bias": "norm_out.adaLN_modulation.1.bias",
+}
+
+# Block-level mappings for BFL double_blocks
+BFL_DOUBLE_BLOCK_MAPPING = {
+    # Note: In BFL format, QKV are already FUSED (img_attn.qkv.weight contains Q+K+V concatenated)
+    # We need to UNFUSE them first, then REFUSE in Primus/Megatron GQA format
+    "img_attn.qkv.weight": "self_attention.linear_qkv.weight",  # Will need special handling
+    "img_attn.qkv.bias": "self_attention.linear_qkv.bias",
+    "img_attn.proj.weight": "self_attention.linear_proj.weight",
+    "img_attn.proj.bias": "self_attention.linear_proj.bias",
+    "txt_attn.qkv.weight": "self_attention.added_linear_qkv.weight",  # Will need special handling
+    "txt_attn.qkv.bias": "self_attention.added_linear_qkv.bias",
+    "txt_attn.proj.weight": "self_attention.added_linear_proj.weight",
+    "txt_attn.proj.bias": "self_attention.added_linear_proj.bias",
+    # QK LayerNorms
+    "img_attn.norm.query_norm.scale": "self_attention.q_layernorm.weight",
+    "img_attn.norm.key_norm.scale": "self_attention.k_layernorm.weight",
+    "txt_attn.norm.query_norm.scale": "self_attention.added_q_layernorm.weight",
+    "txt_attn.norm.key_norm.scale": "self_attention.added_k_layernorm.weight",
+    # Image MLPs
+    "img_mlp.0.weight": "mlp.linear_fc1.weight",
+    "img_mlp.0.bias": "mlp.linear_fc1.bias",
+    "img_mlp.2.weight": "mlp.linear_fc2.weight",
+    "img_mlp.2.bias": "mlp.linear_fc2.bias",
+    # Text MLPs
+    "txt_mlp.0.weight": "context_mlp.linear_fc1.weight",
+    "txt_mlp.0.bias": "context_mlp.linear_fc1.bias",
+    "txt_mlp.2.weight": "context_mlp.linear_fc2.weight",
+    "txt_mlp.2.bias": "context_mlp.linear_fc2.bias",
+    # Modulation (AdaLN)
+    "img_mod.lin.weight": "adaln.adaLN_modulation.1.weight",
+    "img_mod.lin.bias": "adaln.adaLN_modulation.1.bias",
+    "txt_mod.lin.weight": "adaln_context.adaLN_modulation.1.weight",
+    "txt_mod.lin.bias": "adaln_context.adaLN_modulation.1.bias",
+}
+
+# Block-level mappings for BFL single_blocks
+BFL_SINGLE_BLOCK_MAPPING = {
+    # Note: linear1 in BFL is FUSED [Q, K, V, MLP] - need to split
+    "linear1.weight": None,  # Special handling: split into QKV + MLP
+    "linear1.bias": None,
+    # Note: linear2 in BFL is FUSED [MLP_out, proj_out] - need to split
+    "linear2.weight": None,  # Special handling: split into proj_out only (simplified in some versions)
+    "linear2.bias": None,
+    # Modulation
+    "modulation.lin.weight": "adaln.adaLN_modulation.1.weight",
+    "modulation.lin.bias": "adaln.adaLN_modulation.1.bias",
+    # QK LayerNorms (renamed 'norm' in BFL single blocks)
+    "norm.query_norm.scale": "self_attention.q_layernorm.weight",
+    "norm.key_norm.scale": "self_attention.k_layernorm.weight",
+}
+
+
+def detect_checkpoint_format(state_dict: Dict[str, torch.Tensor]) -> str:
+    """
+    Detect checkpoint format by inspecting keys.
+
+    Args:
+        state_dict: Loaded checkpoint state dictionary
+
+    Returns:
+        'bfl_native': Black Forest Labs native format (img_in, txt_in, time_in)
+        'hf_diffusers': HuggingFace Diffusers format (x_embedder, context_embedder, time_text_embed)
+
+    Reference: diffusers Flux checkpoint conversion
+    The BFL native format is what Black Forest Labs releases directly.
+    The HF Diffusers format is what you get after saving from diffusers library.
+    """
+    sample_keys = list(state_dict.keys())
+
+    # Check for BFL native format markers
+    # BFL uses: img_in, txt_in, time_in, vector_in, guidance_in
+    if any(k.startswith("img_in.") for k in sample_keys):
+        return "bfl_native"
+    if any(k.startswith("txt_in.") for k in sample_keys):
+        return "bfl_native"
+    if any("time_in.in_layer" in k for k in sample_keys):
+        return "bfl_native"
+
+    # Check for HF Diffusers format markers
+    # Diffusers uses: x_embedder, context_embedder, time_text_embed
+    if any(k.startswith("x_embedder.") for k in sample_keys):
+        return "hf_diffusers"
+    if any("time_text_embed.timestep_embedder" in k for k in sample_keys):
+        return "hf_diffusers"
+
+    # Default to HF Diffusers (original behavior)
+    return "hf_diffusers"
+
+
+def _get_hf_token(token_file: Optional[str] = None) -> Optional[str]:
+    """
+    Get HuggingFace token with fallback options.
+
+    Thin wrapper around the shared
+    :func:`...preprocessing.auth.setup_hf_authentication` so the token-resolution
+    priority chain (file with permission check -> HF_TOKEN env -> HF CLI login ->
+    None) lives in one place. Returns None instead of raising so the converter
+    can fall back to public-only access.
+
+    Args:
+        token_file: Optional path to token file
+
+    Returns:
+        Token string if found, None otherwise
+    """
+    from primus.backends.megatron.data.diffusion.preprocessing.auth import (
+        HFAuthError,
+        setup_hf_authentication,
+    )
+
+    try:
+        return setup_hf_authentication(token_file=token_file, use_env=True)
+    except HFAuthError:
+        # Preserve fallback-to-public behavior for the converter rather than
+        # hard-failing on a bad/insecure token file.
+        return None
+
+
+def convert_hf_checkpoint(
+    checkpoint_path: Union[str, Path],
+    flux_config,
+    save_to: Optional[Union[str, Path]] = None,
+) -> Dict[str, torch.Tensor]:
+    """
+    Convert HuggingFace Flux checkpoint to Primus format.
+
+    Supports both local paths and HuggingFace repo IDs. If a repo ID is provided
+    (e.g., "black-forest-labs/FLUX.1-dev"), the checkpoint will be
+    automatically downloaded from HuggingFace Hub.
+
+    Args:
+        checkpoint_path: Path to HF checkpoint OR HuggingFace repo ID
+                        (e.g., "black-forest-labs/FLUX.1-dev/transformer")
+        flux_config: FluxConfig instance with model architecture info
+        save_to: Optional path to save converted checkpoint
+
+    Returns:
+        Dictionary of converted state dict (Primus format)
+
+    Example:
+        >>> from primus.backends.megatron.core.models.diffusion.flux import FluxConfig
+        >>> config = FluxConfig.flux_12b()
+        >>> # Auto-download from HuggingFace
+        >>> primus_sd = convert_hf_checkpoint(
+        ...     "black-forest-labs/FLUX.1-dev/transformer",
+        ...     flux_config=config,
+        ...     save_to="primus_flux_12b.safetensors"
+        ... )
+    """
+    checkpoint_path_str = str(checkpoint_path)
+    checkpoint_path_obj = Path(checkpoint_path)
+
+    # Check if path exists locally first
+    if checkpoint_path_obj.exists():
+        # Local file/directory exists, use it directly
+        pass
+    else:
+        # Path doesn't exist - check if it looks like a HuggingFace repo ID
+        # HF repo IDs: don't start with /, ./, ../, and don't contain ..
+        looks_like_hf_repo = (
+            "/" in checkpoint_path_str
+            and not checkpoint_path_str.startswith("/")
+            and not checkpoint_path_str.startswith("./")
+            and not checkpoint_path_str.startswith("../")
+            and ".." not in checkpoint_path_str
+            and not checkpoint_path_str.startswith(".")  # Avoid hidden files/dirs
+        )
+
+        if looks_like_hf_repo:
+            logger.info("Detected HuggingFace repo ID: %s", checkpoint_path_str)
+            logger.info("Downloading from HuggingFace Hub...")
+
+            try:
+                from huggingface_hub import snapshot_download
+            except ImportError:
+                raise ImportError(
+                    "huggingface_hub is required for downloading checkpoints. "
+                    "Install with: pip install huggingface_hub"
+                )
+
+            # Setup authentication: .hf_token file → HF_TOKEN env → HF CLI login
+            token_file = Path(__file__).parents[6] / ".hf_token"
+            hf_token = None
+            if token_file.exists():
+                logger.info("Using HuggingFace token from project root: %s", token_file)
+                hf_token = _get_hf_token(token_file=str(token_file))
+            else:
+                # Falls back to HF_TOKEN env var and HF CLI login
+                hf_token = _get_hf_token(token_file=None)
+
+            # Parse repo_id and subfolder
+            parts = checkpoint_path_str.split("/")
+            if len(parts) >= 2:
+                repo_id = "/".join(parts[:2])  # e.g., "black-forest-labs/FLUX.1-dev"
+                subfolder = "/".join(parts[2:]) if len(parts) > 2 else None  # e.g., "transformer"
+            else:
+                repo_id = checkpoint_path_str
+                subfolder = None
+
+            # Download to cache
+            cache_dir = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface/hub"))
+
+            try:
+                local_dir = snapshot_download(
+                    repo_id=repo_id,
+                    allow_patterns=[f"{subfolder}/*"] if subfolder else None,
+                    cache_dir=cache_dir,
+                    token=hf_token,
+                    resume_download=True,
+                )
+
+                # Construct path to actual checkpoint
+                checkpoint_path = Path(local_dir) / subfolder if subfolder else Path(local_dir)
+                logger.info("Downloaded to: %s", checkpoint_path)
+
+            except Exception as e:
+                logger.error("Download failed: %s", e)
+                if "401" in str(e) or "403" in str(e):
+                    raise RuntimeError(
+                        f"Authentication failed. Please set HuggingFace token using one of:\n"
+                        f"  1. Create .hf_token file: echo 'your_token' > .hf_token && chmod 600 .hf_token\n"
+                        f"  2. Set environment variable: export HF_TOKEN=your_token\n"
+                        f"  3. Login via CLI: huggingface-cli login\n"
+                        f"Get your token from: https://huggingface.co/settings/tokens\n"
+                        f"Accept the model license at: https://huggingface.co/{repo_id}"
+                    ) from e
+                raise
+        else:
+            # Looks like a local path that doesn't exist - raise FileNotFoundError
+            raise FileNotFoundError(
+                f"Checkpoint file or directory not found: {checkpoint_path_str}\n"
+                f"If this is a HuggingFace repo ID, ensure it follows the format 'org/repo' or 'org/repo/subfolder'"
+            )
+
+    logger.info("Loading HuggingFace checkpoint from: %s", checkpoint_path)
+
+    # Load HF checkpoint
+    hf_state_dict = {}
+    checkpoint_path = Path(checkpoint_path)
+
+    if checkpoint_path.is_dir():
+        # Load all .safetensors files in directory
+        safetensor_files = list(checkpoint_path.glob("*.safetensors"))
+        if not safetensor_files:
+            raise FileNotFoundError(f"No .safetensors files found in {checkpoint_path}")
+
+        for file in safetensor_files:
+            logger.info("  Loading %s...", file.name)
+            hf_state_dict.update(load_safetensors(str(file)))
+    elif checkpoint_path.is_file():
+        hf_state_dict = load_safetensors(str(checkpoint_path))
+    else:
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    logger.info("Loaded %d keys from HuggingFace checkpoint", len(hf_state_dict))
+
+    # Detect checkpoint format
+    checkpoint_format = detect_checkpoint_format(hf_state_dict)
+    logger.info("Detected checkpoint format: %s", checkpoint_format)
+
+    # Branch based on format
+    if checkpoint_format == "bfl_native":
+        return _convert_bfl_checkpoint(hf_state_dict, flux_config, save_to)
+    else:
+        return _convert_hf_diffusers_checkpoint(hf_state_dict, flux_config, save_to)
+
+
+def _convert_hf_diffusers_checkpoint(
+    hf_state_dict: Dict[str, torch.Tensor],
+    flux_config,
+    save_to: Optional[Union[str, Path]] = None,
+) -> Dict[str, torch.Tensor]:
+    """Convert HuggingFace Diffusers format checkpoint to Primus format."""
+    # Convert to Primus format
+    primus_state_dict = {}
+    num_double_blocks = -1
+    num_single_blocks = -1
+
+    # First pass: Convert simple key mappings (ONLY double blocks and root keys)
+    logger.info("Converting simple key mappings (double blocks and root keys)...")
+    for hf_key, value in hf_state_dict.items():
+        # Skip QKV weights - will handle separately
+        if any(
+            x in hf_key
+            for x in [
+                "attn.to_q",
+                "attn.to_k",
+                "attn.to_v",
+                "attn.add_q_proj",
+                "attn.add_k_proj",
+                "attn.add_v_proj",
+            ]
+        ):
+            continue
+
+        # Skip ALL single block keys - will handle in separate pass AFTER determining num_double_blocks
+        if hf_key.startswith("single_transformer_blocks"):
+            continue
+
+        # Map double blocks -> transformer.layers[0-18]
+        if hf_key.startswith("transformer_blocks"):
+            parts = hf_key.split(".")
+            idx = int(parts[1])
+            sub_key = ".".join(parts[2:])
+            num_double_blocks = max(idx, num_double_blocks)
+
+            if sub_key in FLUX_KEY_MAPPING["double_blocks"]:
+                # New key format: transformer.layers.{idx} instead of double_blocks.{idx}
+                primus_key = f"transformer.layers.{idx}.{FLUX_KEY_MAPPING['double_blocks'][sub_key]}"
+                primus_state_dict[primus_key] = value
+
+        # Map root-level keys
+        elif hf_key in FLUX_KEY_MAPPING:
+            # Special handling for norm_out: swap scale/shift halves
+            # HF Diffusers stores [SCALE; SHIFT], but Primus/BFL native expects [SHIFT; SCALE]
+            if hf_key == "norm_out.linear.weight":
+                half_size = value.shape[0] // 2
+                scale_half = value[:half_size, :]  # HF first half = SCALE
+                shift_half = value[half_size:, :]  # HF second half = SHIFT
+                # Swap to BFL native order: [SHIFT; SCALE]
+                value = torch.cat([shift_half, scale_half], dim=0)
+            elif hf_key == "norm_out.linear.bias":
+                half_size = value.shape[0] // 2
+                scale_half = value[:half_size]  # HF first half = SCALE
+                shift_half = value[half_size:]  # HF second half = SHIFT
+                # Swap to BFL native order: [SHIFT; SCALE]
+                value = torch.cat([shift_half, scale_half], dim=0)
+
+            primus_state_dict[FLUX_KEY_MAPPING[hf_key]] = value
+
+    # Detect number of single blocks
+    for hf_key in hf_state_dict.keys():
+        if hf_key.startswith("single_transformer_blocks"):
+            parts = hf_key.split(".")
+            idx = int(parts[1])
+            num_single_blocks = max(idx, num_single_blocks)
+
+    logger.info("Found %d double blocks, %d single blocks", num_double_blocks + 1, num_single_blocks + 1)
+
+    # Second pass: Convert single block simple keys (NOW num_double_blocks is known!)
+    logger.info("Converting single block simple keys...")
+    for hf_key, value in hf_state_dict.items():
+        if not hf_key.startswith("single_transformer_blocks"):
+            continue
+
+        # Skip QKV and proj_out - will handle separately in later passes
+        if any(x in hf_key for x in ["attn.to_q", "attn.to_k", "attn.to_v", "proj_out"]):
+            continue
+
+        parts = hf_key.split(".")
+        idx = int(parts[1])
+        sub_key = ".".join(parts[2:])
+
+        if sub_key in FLUX_KEY_MAPPING["single_blocks"]:
+            layer_idx = num_double_blocks + 1 + idx  # NOW num_double_blocks is determined!
+            primus_key = f"transformer.layers.{layer_idx}.{FLUX_KEY_MAPPING['single_blocks'][sub_key]}"
+            primus_state_dict[primus_key] = value
+
+    # Third pass: Fuse QKV weights for double blocks (now transformer.layers[0-18])
+    logger.info("Fusing QKV weights for double blocks...")
+    for i in range(num_double_blocks + 1):
+        # Main attention QKV
+        q_key = f"transformer_blocks.{i}.attn.to_q.weight"
+        k_key = f"transformer_blocks.{i}.attn.to_k.weight"
+        v_key = f"transformer_blocks.{i}.attn.to_v.weight"
+
+        fused_qkv = _fuse_qkv_weights(
+            flux_config, hf_state_dict[q_key], hf_state_dict[k_key], hf_state_dict[v_key]
+        )
+        # New key format: transformer.layers.{i} instead of double_blocks.{i}
+        primus_state_dict[f"transformer.layers.{i}.self_attention.linear_qkv.weight"] = fused_qkv
+
+        # QKV bias
+        q_bias_key = f"transformer_blocks.{i}.attn.to_q.bias"
+        k_bias_key = f"transformer_blocks.{i}.attn.to_k.bias"
+        v_bias_key = f"transformer_blocks.{i}.attn.to_v.bias"
+
+        fused_qkv_bias = _fuse_qkv_bias(
+            flux_config, hf_state_dict[q_bias_key], hf_state_dict[k_bias_key], hf_state_dict[v_bias_key]
+        )
+        primus_state_dict[f"transformer.layers.{i}.self_attention.linear_qkv.bias"] = fused_qkv_bias
+
+        # Context (added) attention QKV
+        add_q_key = f"transformer_blocks.{i}.attn.add_q_proj.weight"
+        add_k_key = f"transformer_blocks.{i}.attn.add_k_proj.weight"
+        add_v_key = f"transformer_blocks.{i}.attn.add_v_proj.weight"
+
+        fused_add_qkv = _fuse_qkv_weights(
+            flux_config, hf_state_dict[add_q_key], hf_state_dict[add_k_key], hf_state_dict[add_v_key]
+        )
+        primus_state_dict[f"transformer.layers.{i}.self_attention.added_linear_qkv.weight"] = fused_add_qkv
+
+        # Added QKV bias
+        add_q_bias_key = f"transformer_blocks.{i}.attn.add_q_proj.bias"
+        add_k_bias_key = f"transformer_blocks.{i}.attn.add_k_proj.bias"
+        add_v_bias_key = f"transformer_blocks.{i}.attn.add_v_proj.bias"
+
+        fused_add_qkv_bias = _fuse_qkv_bias(
+            flux_config,
+            hf_state_dict[add_q_bias_key],
+            hf_state_dict[add_k_bias_key],
+            hf_state_dict[add_v_bias_key],
+        )
+        primus_state_dict[f"transformer.layers.{i}.self_attention.added_linear_qkv.bias"] = fused_add_qkv_bias
+
+    # Fourth pass: Fuse QKV and split proj_out for single blocks (now transformer.layers[19+])
+    logger.info("Fusing QKV and splitting proj_out for single blocks...")
+    for i in range(num_single_blocks + 1):
+        # Calculate layer index with offset
+        layer_idx = num_double_blocks + 1 + i
+
+        # QKV
+        q_key = f"single_transformer_blocks.{i}.attn.to_q.weight"
+        k_key = f"single_transformer_blocks.{i}.attn.to_k.weight"
+        v_key = f"single_transformer_blocks.{i}.attn.to_v.weight"
+
+        fused_qkv = _fuse_qkv_weights(
+            flux_config, hf_state_dict[q_key], hf_state_dict[k_key], hf_state_dict[v_key]
+        )
+        # New key format with offset
+        primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_qkv.weight"] = fused_qkv
+
+        # QKV bias
+        q_bias_key = f"single_transformer_blocks.{i}.attn.to_q.bias"
+        k_bias_key = f"single_transformer_blocks.{i}.attn.to_k.bias"
+        v_bias_key = f"single_transformer_blocks.{i}.attn.to_v.bias"
+
+        fused_qkv_bias = _fuse_qkv_bias(
+            flux_config, hf_state_dict[q_bias_key], hf_state_dict[k_bias_key], hf_state_dict[v_bias_key]
+        )
+        primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_qkv.bias"] = fused_qkv_bias
+
+        # Split proj_out (combined MLP+Attention output in HF)
+        # HF format: [out_dim, hidden*2] where [:, :3072] is attention, [:, 3072:] is MLP
+        proj_out_weight = hf_state_dict[f"single_transformer_blocks.{i}.proj_out.weight"]
+        proj_out_bias = hf_state_dict[f"single_transformer_blocks.{i}.proj_out.bias"]
+
+        # Split weight at hidden_size (3072 for Flux)
+        hidden_size = flux_config.hidden_size
+        attn_proj = proj_out_weight[:, :hidden_size].clone()
+        mlp_proj = proj_out_weight[:, hidden_size:].clone()
+
+        primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_proj.weight"] = attn_proj
+        primus_state_dict[f"transformer.layers.{layer_idx}.mlp.linear_fc2.weight"] = mlp_proj
+
+        primus_state_dict[f"transformer.layers.{layer_idx}.mlp.linear_fc2.bias"] = proj_out_bias.clone()
+
+    logger.info("Conversion complete! Primus state dict has %d keys", len(primus_state_dict))
+
+    # Save if requested
+    if save_to:
+        save_to = Path(save_to)
+        save_to.parent.mkdir(parents=True, exist_ok=True)
+        logger.info("Saving to: %s", save_to)
+        save_safetensors(primus_state_dict, str(save_to))
+        logger.info("Saved successfully!")
+
+    return primus_state_dict
+
+
+def _convert_bfl_checkpoint(
+    bfl_state_dict: Dict[str, torch.Tensor],
+    flux_config,
+    save_to: Optional[Union[str, Path]] = None,
+) -> Dict[str, torch.Tensor]:
+    """
+    Convert Black Forest Labs native format checkpoint to Primus format.
+
+    Reference: diffusers Flux checkpoint conversion
+
+    Key differences from HF Diffusers format:
+    - QKV are already FUSED in BFL (img_attn.qkv.weight contains Q+K+V concatenated)
+    - Single blocks have fused linear1 (Q, K, V, MLP) and linear2 (proj_out)
+    """
+    logger.info("Converting BFL native format to Primus...")
+    primus_state_dict = {}
+    num_double_blocks = -1
+    num_single_blocks = -1
+
+    # First pass: Convert simple root-level mappings
+    logger.info("Converting root-level keys...")
+    for bfl_key, value in bfl_state_dict.items():
+        # Skip block-level keys - will handle in second pass
+        if bfl_key.startswith("double_blocks.") or bfl_key.startswith("single_blocks."):
+            continue
+
+        # Map root-level keys
+        if bfl_key in BFL_KEY_MAPPING:
+            primus_key = BFL_KEY_MAPPING[bfl_key]
+            primus_state_dict[primus_key] = value
+
+    # Detect number of blocks
+    for key in bfl_state_dict.keys():
+        if key.startswith("double_blocks."):
+            idx = int(key.split(".")[1])
+            num_double_blocks = max(idx, num_double_blocks)
+        elif key.startswith("single_blocks."):
+            idx = int(key.split(".")[1])
+            num_single_blocks = max(idx, num_single_blocks)
+
+    logger.info("Found %d double blocks, %d single blocks", num_double_blocks + 1, num_single_blocks + 1)
+
+    # Second pass: Convert double blocks -> transformer.layers[0-18]
+    logger.info("Converting double blocks (unfusing and refusing QKV)...")
+    for i in range(num_double_blocks + 1):
+        block_prefix_bfl = f"double_blocks.{i}"
+        block_prefix_primus = f"transformer.layers.{i}"  # New key format
+
+        # Convert simple mappings first
+        for bfl_suffix, primus_suffix in BFL_DOUBLE_BLOCK_MAPPING.items():
+            if primus_suffix is None:
+                continue  # Skip special handling keys
+
+            bfl_key = f"{block_prefix_bfl}.{bfl_suffix}"
+            if bfl_key in bfl_state_dict:
+                primus_key = f"{block_prefix_primus}.{primus_suffix}"
+
+                # Special handling for QKV (need to unfuse BFL QKV, then refuse in Primus format)
+                if "qkv" in bfl_suffix and ("weight" in bfl_suffix or "bias" in bfl_suffix):
+                    continue  # Handle separately below
+                else:
+                    primus_state_dict[primus_key] = bfl_state_dict[bfl_key]
+
+        # Handle img_attn QKV (unfuse from BFL concat format, then refuse in Primus GQA format)
+        # BFL format: [Q; K; V] concatenated along dim 0
+        # Reference: diffusers Flux checkpoint conversion
+        img_qkv_weight = bfl_state_dict[f"{block_prefix_bfl}.img_attn.qkv.weight"]
+        img_qkv_bias = bfl_state_dict[f"{block_prefix_bfl}.img_attn.qkv.bias"]
+
+        # Unfuse: split into Q, K, V
+        img_q, img_k, img_v = torch.chunk(img_qkv_weight, 3, dim=0)
+        img_q_bias, img_k_bias, img_v_bias = torch.chunk(img_qkv_bias, 3, dim=0)
+
+        # Refuse in Primus GQA format
+        fused_img_qkv = _fuse_qkv_weights(flux_config, img_q, img_k, img_v)
+        fused_img_qkv_bias = _fuse_qkv_bias(flux_config, img_q_bias, img_k_bias, img_v_bias)
+
+        primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.weight"] = fused_img_qkv
+        primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.bias"] = fused_img_qkv_bias
+
+        # Handle txt_attn QKV (same process)
+        txt_qkv_weight = bfl_state_dict[f"{block_prefix_bfl}.txt_attn.qkv.weight"]
+        txt_qkv_bias = bfl_state_dict[f"{block_prefix_bfl}.txt_attn.qkv.bias"]
+
+        txt_q, txt_k, txt_v = torch.chunk(txt_qkv_weight, 3, dim=0)
+        txt_q_bias, txt_k_bias, txt_v_bias = torch.chunk(txt_qkv_bias, 3, dim=0)
+
+        fused_txt_qkv = _fuse_qkv_weights(flux_config, txt_q, txt_k, txt_v)
+        fused_txt_qkv_bias = _fuse_qkv_bias(flux_config, txt_q_bias, txt_k_bias, txt_v_bias)
+
+        primus_state_dict[f"{block_prefix_primus}.self_attention.added_linear_qkv.weight"] = fused_txt_qkv
+        primus_state_dict[f"{block_prefix_primus}.self_attention.added_linear_qkv.bias"] = fused_txt_qkv_bias
+
+    # Third pass: Convert single blocks -> transformer.layers[19+]
+    # BFL single blocks have:
+    # - linear1: fused [Q, K, V, MLP] - need to split and handle separately
+    # - linear2: just proj_out (simpler than HF Diffusers which has MLP+proj)
+    # Reference: diffusers Flux checkpoint conversion
+    logger.info("Converting single blocks (splitting linear1)...")
+    for i in range(num_single_blocks + 1):
+        block_prefix_bfl = f"single_blocks.{i}"
+        # Calculate layer index with offset (num_double_blocks + 1 + i)
+        layer_idx = num_double_blocks + 1 + i
+        block_prefix_primus = f"transformer.layers.{layer_idx}"  # New key format
+
+        # Convert simple mappings (modulation, norms)
+        for bfl_suffix, primus_suffix in BFL_SINGLE_BLOCK_MAPPING.items():
+            if primus_suffix is None:
+                continue  # Skip special handling
+
+            bfl_key = f"{block_prefix_bfl}.{bfl_suffix}"
+            if bfl_key in bfl_state_dict:
+                primus_key = f"{block_prefix_primus}.{primus_suffix}"
+                primus_state_dict[primus_key] = bfl_state_dict[bfl_key]
+
+        # Handle linear1: fused [Q, K, V, MLP]
+        linear1_weight = bfl_state_dict[f"{block_prefix_bfl}.linear1.weight"]
+        linear1_bias = bfl_state_dict[f"{block_prefix_bfl}.linear1.bias"]
+
+        # Split along dim 0: [Q, K, V, MLP]
+        hidden_size = flux_config.hidden_size
+        mlp_hidden_dim = int(hidden_size * 4.0)  # mlp_ratio = 4.0
+        split_sizes = (hidden_size, hidden_size, hidden_size, mlp_hidden_dim)
+
+        q, k, v, mlp = torch.split(linear1_weight, split_sizes, dim=0)
+        q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_sizes, dim=0)
+
+        # Fuse Q, K, V in Primus GQA format
+        fused_qkv = _fuse_qkv_weights(flux_config, q, k, v)
+        fused_qkv_bias = _fuse_qkv_bias(flux_config, q_bias, k_bias, v_bias)
+
+        primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.weight"] = fused_qkv
+        primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.bias"] = fused_qkv_bias
+
+        # MLP goes to linear_fc1
+        primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc1.weight"] = mlp
+        primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc1.bias"] = mlp_bias
+
+        # Handle linear2: In BFL, this is just proj_out (not fused with MLP like in HF Diffusers)
+        linear2_weight = bfl_state_dict[f"{block_prefix_bfl}.linear2.weight"]
+        linear2_bias = bfl_state_dict[f"{block_prefix_bfl}.linear2.bias"]
+
+        # In BFL format, linear2 is [out_channels, hidden_size*2] where first half is attn proj, second is mlp proj
+        # Split at hidden_size
+        attn_proj = linear2_weight[:, :hidden_size].clone()
+        mlp_proj = linear2_weight[:, hidden_size:].clone()
+
+        primus_state_dict[f"{block_prefix_primus}.self_attention.linear_proj.weight"] = attn_proj
+        primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc2.weight"] = mlp_proj
+
+        primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc2.bias"] = linear2_bias.clone()
+
+    logger.info("Conversion complete! Primus state dict has %d keys", len(primus_state_dict))
+
+    # Save if requested
+    if save_to:
+        save_to = Path(save_to)
+        save_to.parent.mkdir(parents=True, exist_ok=True)
+        logger.info("Saving to: %s", save_to)
+        save_safetensors(primus_state_dict, str(save_to))
+        logger.info("Saved successfully!")
+
+    return primus_state_dict
+
+
+__all__ = [
+    "convert_hf_checkpoint",
+    "FLUX_KEY_MAPPING",
+    "BFL_KEY_MAPPING",
+    "BFL_DOUBLE_BLOCK_MAPPING",
+    "BFL_SINGLE_BLOCK_MAPPING",
+    "detect_checkpoint_format",
+    "_fuse_qkv_weights",
+    "_fuse_qkv_bias",
+]
diff --git a/primus/backends/megatron/core/models/diffusion/flux/config.py b/primus/backends/megatron/core/models/diffusion/flux/config.py
new file mode 100644
index 000000000..83446615b
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/config.py
@@ -0,0 +1,327 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Configuration for Flux diffusion model.
+
+Flux is a flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer)
+architecture that uses separate "joint" and "single" transformer blocks.
+
+Reference:
+    - https://github.com/black-forest-labs/flux
+    - NeMo's Flux implementation
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+
+from ..common.config import BaseDiffusionConfig
+
+# Try to import erf_gelu from megatron, fallback to custom implementation
+try:
+    from megatron.core.transformer.utils import erf_gelu
+except ImportError:
+    # Fallback used when Megatron's erf_gelu is unavailable
+    def erf_gelu(x):
+        """GELU activation using error function approximation."""
+        return 0.5 * x * (1.0 + torch.erf(x / 1.4142135623730951))
+
+
+__all__ = ["FluxConfig", "openai_gelu_no_jit", "erf_gelu"]
+
+
+# Custom non-JIT compiled openai_gelu to avoid ROCm bugs
+def openai_gelu_no_jit(x):
+    """
+    OpenAI's GELU implementation without JIT compilation (for ROCm compatibility).
+
+    This is the tanh-based approximation of GELU used in the original Flux model.
+    We use a non-JIT version to avoid ROCm compilation bugs that cause NaN values.
+    """
+    return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
+
+
+@dataclass
+class FluxConfig(BaseDiffusionConfig):
+    """
+    Flux-specific configuration.
+
+    Configuration for Flux flow-based diffusion model. Flux uses a dual-stream
+    architecture with joint (multimodal) and single (image-only) transformer layers.
+
+    Standard Configurations:
+        - Flux.1 [dev/schnell]: 12B parameters (19 joint + 38 single layers)
+        - Flux 535M: Minimal config for testing (1 joint + 1 single layer)
+
+    Key Configuration Groups:
+        - Architecture: num_joint_layers, num_single_layers, hidden_size
+        - Context: context_dim (T5), vec_in_dim (CLIP), model_channels (timestep)
+        - Position Encoding: theta, axes_dim, rotary_interleaved
+        - Guidance: guidance_embed, guidance_scale
+
+    For architecture details, see Flux class documentation.
+
+    Attributes:
+        num_joint_layers: Number of joint (multimodal) transformer layers
+        num_single_layers: Number of single (image-only) transformer layers
+        context_dim: Dimension of text context embeddings (T5-XXL: 4096)
+        vec_in_dim: Dimension of pooled text embeddings (CLIP-L: 768)
+        model_channels: Channels for timestep embedding (default: 256)
+        guidance_embed: Whether to use guidance embedding for CFG
+        guidance_scale: Guidance scale for classifier-free guidance (default: 3.5)
+        theta: Base for RoPE position embeddings (default: 10000)
+        axes_dim: Dimension for each axis in 3D RoPE (default: [16, 56, 56])
+        rotary_interleaved: Whether RoPE dimensions are interleaved (default: True)
+        apply_rope_fusion: Whether to apply RoPE fusion optimization (default: False)
+        add_qkv_bias: Whether to add bias to QKV projections (default: True)
+        single_block_bias: Whether to add bias to single block linear layers (default: True)
+        activation_func: Activation function (default: openai_gelu_no_jit)
+        use_te_rng_tracker: Whether to use Transformer Engine RNG tracker (default: False)
+    """
+
+    # Model identification
+    model_type: str = "flux"
+
+    # Dummy layer for compatibility (actual layers defined below)
+    num_layers: int = 1  # Not used in Flux, kept for compatibility
+
+    # Architecture: Number of layers
+    num_joint_layers: int = 19  # Default: Flux 12B
+    num_single_layers: int = 38  # Default: Flux 12B
+
+    # Architecture: Dimensions (Flux standard: 3072)
+    hidden_size: int = 3072
+    num_attention_heads: int = 24
+
+    # Input dimensions
+    in_channels: int = 64  # Packed latent channels (16 VAE channels * 4 from 2x2 patch packing)
+
+    # Context dimensions
+    context_dim: int = 4096  # T5-XXL hidden dimension
+    vec_in_dim: int = 768  # CLIP-L pooled dimension
+    model_channels: int = 256  # Channels for timestep embedding
+
+    # Guidance (for classifier-free guidance)
+    guidance_embed: bool = False  # Set True for CFG support
+    guidance_scale: float = 3.5  # CFG guidance scale
+    cfg_dropout_prob: float = (
+        0.0  # Probability of replacing text embeddings with empty encodings (MLPerf: 0.1)
+    )
+
+    # Training: timestep sampling strategy
+    # Options: "logit_normal" (SD3 default), "direct_uniform" (NVIDIA MLPerf), "uniform", "mode"
+    timestep_sampling_strategy: str = "logit_normal"
+
+    # Position embeddings (RoPE)
+    theta: int = 10000  # Base for RoPE frequencies
+    axes_dim: tuple = (16, 56, 56)  # 3D RoPE: (channels, height, width)
+    rotary_interleaved: bool = True  # Whether RoPE dimensions are interleaved
+    apply_rope_fusion: bool = False  # Whether to apply RoPE fusion optimization
+
+    # Patchification (Flux uses 1x1 patches by default)
+    patch_size: int = 1
+
+    # Attention configuration
+    add_qkv_bias: bool = True  # Whether to add bias to QKV projections
+
+    # Single block configuration
+    single_block_bias: bool = True  # Whether to add bias to single block linear layers
+
+    # Default: non-JIT tanh GELU (ROCm-safe). Override via YAML activation_func:
+    # "openai_gelu" -> fused F.gelu(approximate="tanh"); "erf_gelu" -> erf-based GELU.
+    activation_func: Callable = openai_gelu_no_jit
+
+    # Dropout (Flux typically uses 0)
+    hidden_dropout: float = 0.0
+    attention_dropout: float = 0.0
+
+    # Normalization
+    layernorm_epsilon: float = 1e-6
+
+    # Initialization
+    use_cpu_initialization: bool = True
+
+    # Optimization flags
+    gradient_accumulation_fusion: bool = False
+    use_dual_fp8_output_projection: bool = False
+    use_triton_ops: bool = False
+    adaln_plain_ops: bool = False
+    adaln_always_jit_fuser: bool = False
+
+    # FSDP2 prefetch depth: number of layers to prefetch ahead for all-gather overlap.
+    # 1 = prefetch next layer (default), 2 = prefetch next 2 layers, 0 = no prefetch.
+    fsdp_prefetch_depth: int = 1
+
+    # FP8 all-gather data cache: batch-quantize all weights upfront and cache
+    # the FP8 data so fsdp_pre_all_gather skips per-layer quantization.
+    # When false, only scales are precomputed; quantization happens on-demand.
+    fp8_precompute_data_cache: bool = True
+
+    # Optimizer foreach batching: use _foreach_copy_ for batched grad/weight
+    # copies instead of per-parameter loops.
+    optimizer_foreach: bool = True
+
+    # Overlap grad norm with reduce-scatter: accumulate squared norms
+    # incrementally inside FSDP2's post-reduce stream via
+    # register_post_accumulate_grad_hook, replacing the full recompute
+    # in clip_grad_norm with a single all-reduce + sqrt + clip.
+    overlap_grad_norm: bool = False
+
+    # Use the C++ quantize_fp8 kernel from primus_turbo for tensorwise FP8
+    # quantization instead of the default inline implementation.
+    use_cpp_fp8_quantize: bool = False
+
+    # CUDA graph support
+    enable_cuda_graph: bool = False
+    cuda_graph_scope: Optional[str] = None  # Options: "full", "full_iteration"
+    cuda_graph_warmup_steps: int = 2
+
+    # Transformer Engine
+    use_te_rng_tracker: bool = False
+
+    # Torch compile configuration (applied AFTER distributed setup)
+    enable_torch_compile: bool = False  # Enable compilation (applied after FSDP/DDP wrapping)
+    torch_compile_backend: str = "inductor"  # Compilation backend
+    torch_compile_mode: str = "default"  # "default", "reduce-overhead", "max-autotune"
+    torch_compile_fullgraph: bool = False  # Allow graph breaks (recommended for distributed)
+    torch_compile_optimizer: bool = False  # Whether to compile optimizer step
+    torch_compile_optimizer_scope: str = (
+        "full"  # "full" = compile FSDP2FP32Optimizer.step(), "inner_only" = compile only inner AdamW.step()
+    )
+
+    # Selective stack compilation (TE spec only — local spec falls back to per_block)
+    torch_compile_strategy: str = "per_block"
+    # "per_block": compile each transformer layer individually (default, recommended)
+    # "whole_model": compile entire forward (incompatible with overlap_param_gather)
+    # "double_stack": compile only the double (joint) block loop
+    # "single_stack": compile only the single block loop
+    # "stack": compile both double and single block loops
+    # "full_dit": compile the entire DiT (double + cat + single + output) as one region
+    torch_compile_replace_qk_rmsnorm: bool = False
+    torch_compile_disable_inductor_cudagraphs: bool = True
+    torch_compile_emulate_precision_casts: bool = True  # Preserve eager BF16 precision in Triton kernels
+    torch_compile_fused_ln_modulate: bool = True  # Use fused LN+modulate Triton kernel in AdaLN
+
+    def __post_init__(self):
+        """Post-initialization processing."""
+        # MUST set num_layers BEFORE super().__post_init__() -- convention for BaseDiffusionConfig
+        # Required because BaseDiffusionConfig maps sensitive_layer_* fields using num_layers,
+        # and TransformerConfig validates num_layers_at_start_in_bf16 against num_layers.
+        self.num_layers = self.num_joint_layers + self.num_single_layers
+
+        super().__post_init__()  # BaseDiffusionConfig (mapping) -> TransformerConfig (validation)
+
+        # Xavier uniform for Megatron parallel linear layers (common Flux reference default).
+        self.init_method = nn.init.xavier_uniform_
+        self.output_layer_init_method = nn.init.xavier_uniform_
+
+        # Ensure axes_dim is a tuple
+        if isinstance(self.axes_dim, list):
+            self.axes_dim = tuple(self.axes_dim)
+
+    def validate(self):
+        """
+        Validate Flux-specific configuration.
+
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        # Call parent validation
+        super().validate()
+
+        # Flux-specific validations
+        if self.num_joint_layers <= 0:
+            raise ValueError(f"num_joint_layers must be positive, got {self.num_joint_layers}")
+
+        if self.num_single_layers <= 0:
+            raise ValueError(f"num_single_layers must be positive, got {self.num_single_layers}")
+
+        if self.context_dim <= 0:
+            raise ValueError(f"context_dim must be positive, got {self.context_dim}")
+
+        if self.vec_in_dim <= 0:
+            raise ValueError(f"vec_in_dim must be positive, got {self.vec_in_dim}")
+
+        if self.theta <= 0:
+            raise ValueError(f"theta must be positive, got {self.theta}")
+
+        if len(self.axes_dim) != 3:
+            raise ValueError(f"axes_dim must have 3 elements for 3D RoPE, got {len(self.axes_dim)}")
+
+        if any(d <= 0 for d in self.axes_dim):
+            raise ValueError(f"All axes_dim values must be positive, got {self.axes_dim}")
+
+        # Validate RoPE fusion constraints
+        if self.apply_rope_fusion:
+            import warnings
+
+            warnings.warn(
+                "\n"
+                "=" * 80 + "\n"
+                "RoPE Fusion Enabled: Same-Resolution Batch Requirement\n"
+                "=" * 80 + "\n"
+                "RoPE fusion optimization requires ALL images in each training batch to have\n"
+                "the SAME resolution (height and width). Variable-resolution batches will\n"
+                "produce incorrect positional encodings and degrade model quality.\n"
+                "\n"
+                "Recommended data pipeline configurations:\n"
+                "  1. Fixed resolution: All images resized to same size (e.g., 512x512)\n"
+                "  2. Resolution bucketing: Group images by resolution in separate batches\n"
+                "  3. Aspect ratio bucketing: Use bucketing with consistent dimensions\n"
+                "\n"
+                "=" * 80,
+                UserWarning,
+                stacklevel=2,
+            )
+
+    def get_num_layers(self):
+        """
+        Get total number of transformer layers.
+
+        Returns:
+            Total number of layers (joint + single)
+        """
+        return self.num_joint_layers + self.num_single_layers
+
+    @classmethod
+    def flux_535m(cls, **kwargs):
+        """
+        Create configuration for Flux 535M (minimal model for testing).
+
+        Args:
+            **kwargs: Override default parameters
+
+        Returns:
+            FluxConfig instance
+        """
+        defaults = {
+            "num_joint_layers": 1,
+            "num_single_layers": 1,
+            "hidden_size": 3072,
+            "num_attention_heads": 24,
+        }
+        defaults.update(kwargs)
+        return cls(**defaults)
+
+    @classmethod
+    def flux_12b(cls, **kwargs):
+        """
+        Create configuration for Flux 12B (standard model).
+
+        Args:
+            **kwargs: Override default parameters
+
+        Returns:
+            FluxConfig instance
+        """
+        defaults = {
+            "num_joint_layers": 19,
+            "num_single_layers": 38,
+            "hidden_size": 3072,
+            "num_attention_heads": 24,
+        }
+        defaults.update(kwargs)
+        return cls(**defaults)
diff --git a/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py
new file mode 100644
index 000000000..9a1932b09
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py
@@ -0,0 +1,594 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Flux layer specifications for Megatron-Core integration.
+
+This module implements transformer layers for Flux's MMDiT architecture:
+    - MMDiTLayer: Joint image-text transformer block
+    - FluxSingleTransformerBlock: Image-only transformer block
+    - Factory functions for creating layer specs with Transformer Engine
+
+Key Innovation: Heterogeneous layer support via TransformerBlock
+    Instead of separate ModuleLists, Primus uses a unified TransformerBlock
+    with heterogeneous layer specifications, enabling better pipeline
+    parallelism and cleaner checkpoint management.
+
+Reference:
+    - Flux Paper: "Flux: A Scalable Diffusion Model"
+    - MMDiT Paper: "Scaling Rectified Flow Transformers"
+    - Megatron-Core: Heterogeneous layer patterns
+"""
+
+import copy
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from megatron.core.models.backends import BackendSpecProvider
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import (
+    TransformerLayer,
+    TransformerLayerSubmodules,
+)
+from megatron.core.utils import make_viewless_tensor
+from torch import Tensor
+
+from primus.backends.megatron.core.models.diffusion.common.normalization import (
+    AdaLN,
+    AdaLNContinuous,
+)
+from primus.backends.megatron.core.models.diffusion.flux.attention import (
+    FluxSingleAttention,
+    JointSelfAttention,
+    JointSelfAttentionSubmodules,
+)
+
+try:
+    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+    from megatron.core.transformer.attention import SelfAttentionSubmodules
+
+    HAVE_TE_SPEC_PROVIDER = True
+except ImportError:
+    HAVE_TE_SPEC_PROVIDER = False
+    SelfAttentionSubmodules = None
+
+# Try to import PrimusTurboLocalSpecProvider (may not exist in all versions)
+try:
+    from primus.backends.megatron.core.extensions.primus_turbo_local_spec import (
+        PrimusTurboFloat8LocalSpecProvider,
+        PrimusTurboLocalSpecProvider,
+    )
+
+    HAVE_PRIMUS_TURBO_LOCAL = True
+except ImportError:
+    HAVE_PRIMUS_TURBO_LOCAL = False
+    PrimusTurboLocalSpecProvider = None
+    PrimusTurboFloat8LocalSpecProvider = None
+
+# MXFP4 provider in a separate guard so a missing primus_turbo_mxfp4_local
+# doesn't break FP8 imports.
+try:
+    from primus.backends.megatron.core.extensions.primus_turbo_local_spec import (
+        PrimusTurboMXFP4LocalSpecProvider,
+    )
+except ImportError:
+    PrimusTurboMXFP4LocalSpecProvider = None
+
+
+class MMDiTLayer(TransformerLayer):
+    """
+    Multimodal Diffusion Transformer (MMDiT) Layer.
+
+    Processes image and text streams jointly via AdaLN-conditioned attention
+    and separate MLPs with gated residual connections. Used in Flux's "double
+    blocks" for cross-modal image-text interaction.
+
+    Args:
+        config: Transformer configuration
+        submodules: Submodule specifications for attention, MLP, etc.
+        layer_number: Layer index in the model (default: 1)
+        context_pre_only: If True, context stream only computes pre-attention norm
+
+    Input/Output:
+        (hidden_states [S_img, B, H], context [S_txt, B, H], timestep_emb [B, H])
+        -> (hidden_states, context) with same shapes
+
+    Reference:
+        - "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis"
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: TransformerLayerSubmodules,
+        layer_number: int = 1,
+        context_pre_only: bool = False,
+        **kwargs,  # Accept additional kwargs from TransformerBlock (e.g., pg_collection)
+    ):
+        hidden_size = config.hidden_size
+        super().__init__(config=config, submodules=submodules, layer_number=layer_number, **kwargs)
+
+        # Enable per-layer CUDA graph if configured
+        if config.enable_cuda_graph and config.cuda_graph_scope != "full_iteration":
+            self.cudagraph_manager = CudaGraphManager(config, share_cudagraph_io_buffers=False)
+
+        # Adaptive layer normalization for main stream (image).
+        # init_method=nn.init.normal_ draws RNG matching NeMo's DiT init
+        # sequence so cross-framework convergence comparisons line up.
+        # init_weights() (model.py) immediately re-zeroes these weights,
+        # so the only observable effect is RNG-sequence alignment.
+        self.adaln = AdaLN(
+            config,
+            modulation_bias=True,
+            n_adaln_chunks=6,
+            use_second_norm=True,
+            init_method=nn.init.normal_,
+        )
+
+        # Adaptive layer normalization for context stream (text)
+        self.context_pre_only = context_pre_only
+        context_norm_type = "ada_norm_continuous" if context_pre_only else "ada_norm_zero"
+
+        if context_norm_type == "ada_norm_continuous":
+            # Continuous AdaLN for context (simpler, used when context is pre-only)
+            self.adaln_context = AdaLNContinuous(
+                config, hidden_size, modulation_bias=True, norm_type="layer_norm"
+            )
+        elif context_norm_type == "ada_norm_zero":
+            # Full AdaLN for context (used when context has full processing).
+            # See note on self.adaln above re: init_method choice.
+            self.adaln_context = AdaLN(
+                config,
+                modulation_bias=True,
+                n_adaln_chunks=6,
+                use_second_norm=True,
+                init_method=nn.init.normal_,
+            )
+        else:
+            raise ValueError(
+                f"Unknown context_norm_type: {context_norm_type}, "
+                f"currently only support `ada_norm_continuous`, `ada_norm_zero`"
+            )
+
+        # Context MLP (only if not pre-only)
+        if not context_pre_only:
+            # Disable context parallelism for context MLP
+            cp_override_config = copy.deepcopy(config)
+            cp_override_config.context_parallel_size = 1
+            cp_override_config.tp_comm_overlap = False
+
+            from megatron.core.transformer.spec_utils import build_module
+
+            self.context_mlp = build_module(submodules.mlp, config=cp_override_config)
+        else:
+            self.context_mlp = None
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        context: Optional[Tensor] = None,
+        context_mask: Optional[Tensor] = None,
+        rotary_pos_emb: Optional[Tensor] = None,
+        timestep_emb: Optional[Tensor] = None,
+        packed_seq_params=None,
+        **kwargs,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Forward pass: Joint processing of image and text tokens.
+
+        Args:
+            hidden_states: Image tokens [seq_img, batch, hidden]
+            context: Text tokens [seq_txt, batch, hidden]
+            timestep_emb: Timestep conditioning [batch, hidden] (required)
+
+        Returns:
+            Tuple of (hidden_states, context) - both updated
+        """
+        # Map TransformerBlock's 'context' to internal 'encoder_hidden_states'
+        encoder_hidden_states = context
+
+        if timestep_emb is None:
+            raise ValueError("MMDiTLayer requires timestep_emb for AdaLN conditioning.")
+
+        emb = timestep_emb
+
+        # Get modulation parameters for main stream (image)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln(emb)
+
+        # Apply modulated layer norm to image tokens
+        norm_hidden_states = self.adaln.modulated_layernorm(
+            hidden_states, shift=shift_msa, scale=scale_msa, layernorm_idx=0
+        )
+
+        # Apply modulated layer norm to text tokens
+        if self.context_pre_only:
+            # Continuous AdaLN (simpler)
+            norm_encoder_hidden_states = self.adaln_context(encoder_hidden_states, emb)
+        else:
+            # Full AdaLN with modulation parameters
+            c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.adaln_context(
+                emb
+            )
+            norm_encoder_hidden_states = self.adaln_context.modulated_layernorm(
+                encoder_hidden_states, shift=c_shift_msa, scale=c_scale_msa, layernorm_idx=0
+            )
+
+        # Joint self-attention
+        attn_output, context_attn_output = self.self_attention(
+            norm_hidden_states,
+            attention_mask=attention_mask,
+            rotary_pos_emb=rotary_pos_emb,
+            additional_hidden_states=norm_encoder_hidden_states,
+        )
+
+        # MLP for main stream (image)
+        # Fused operation: gated residual + modulated layernorm for MLP input
+        hidden_states, norm_hidden_states = self.adaln.scaled_modulated_layernorm(
+            residual=hidden_states,
+            x=attn_output,
+            gate=gate_msa,
+            shift=shift_mlp,
+            scale=scale_mlp,
+            layernorm_idx=1,
+        )
+        mlp_output, mlp_bias = self.mlp(norm_hidden_states)
+        hidden_states = self.adaln.scale_add(hidden_states, x=(mlp_output + mlp_bias), gate=gate_mlp)
+
+        # MLP for context stream (text) - only if not pre-only
+        if not self.context_pre_only:
+            # Fused operation: gated residual + modulated layernorm for context MLP input
+            encoder_hidden_states, norm_encoder_hidden_states = self.adaln_context.scaled_modulated_layernorm(
+                residual=encoder_hidden_states,
+                x=context_attn_output,
+                gate=c_gate_msa,
+                shift=c_shift_mlp,
+                scale=c_scale_mlp,
+                layernorm_idx=1,
+            )
+            context_mlp_output, context_mlp_bias = self.context_mlp(norm_encoder_hidden_states)
+            encoder_hidden_states = self.adaln_context.scale_add(
+                encoder_hidden_states, x=(context_mlp_output + context_mlp_bias), gate=c_gate_mlp
+            )
+
+        # Make output viewless for MPU checkpoint compatibility
+        hidden_states = make_viewless_tensor(
+            inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+        )
+        encoder_hidden_states = make_viewless_tensor(
+            inp=encoder_hidden_states, requires_grad=encoder_hidden_states.requires_grad, keep_graph=True
+        )
+
+        return hidden_states, encoder_hidden_states
+
+    def __call__(self, *args, **kwargs):
+        """Override call to support CUDA graphs."""
+        if hasattr(self, "cudagraph_manager"):
+            return self.cudagraph_manager(self, args, kwargs)
+        return super(MegatronModule, self).__call__(*args, **kwargs)
+
+
+class FluxSingleTransformerBlock(TransformerLayer):
+    """
+    Flux Single Transformer Block (image-only processing).
+
+    Processes image tokens with parallel attention + MLP paths, AdaLN
+    conditioning, and gated residual connections. Used in Flux's "single
+    blocks" after the joint MMDiT layers.
+
+    Args:
+        config: Transformer configuration
+        submodules: Submodule specifications for attention, MLP, etc.
+        layer_number: Layer index in the model (default: 1)
+        mlp_ratio: MLP hidden size ratio (default: 4)
+        n_adaln_chunks: AdaLN modulation chunks (default: 3 for shift, scale, gate)
+        modulation_bias: Whether to use bias in modulation layers (default: True)
+
+    Input/Output:
+        (hidden_states [S, B, H], timestep_emb [B, H]) -> (hidden_states, None)
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: TransformerLayerSubmodules,
+        layer_number: int = 1,
+        mlp_ratio: int = 4,
+        n_adaln_chunks: int = 3,
+        modulation_bias: bool = True,
+        **kwargs,  # Accept additional kwargs from TransformerBlock (e.g., pg_collection)
+    ):
+        # Override add_bias_linear with single_block_bias for this block
+        # This allows independent control of bias in single blocks vs joint blocks
+        original_add_bias_linear = config.add_bias_linear
+        if hasattr(config, "single_block_bias"):
+            config.add_bias_linear = config.single_block_bias
+
+        super().__init__(config=config, submodules=submodules, layer_number=layer_number, **kwargs)
+
+        # Restore original value
+        config.add_bias_linear = original_add_bias_linear
+
+        # Enable per-layer CUDA graph if configured
+        if config.enable_cuda_graph and config.cuda_graph_scope != "full_iteration":
+            self.cudagraph_manager = CudaGraphManager(config, share_cudagraph_io_buffers=False)
+
+        # Adaptive layer normalization
+        # n_adaln_chunks=3 for (shift, scale, gate)
+        # init_method=nn.init.normal_: see note on FluxTransformerBlock's
+        # self.adaln above -- NeMo-aligned RNG draw, re-zeroed in init_weights().
+        self.adaln = AdaLN(
+            config=config,
+            n_adaln_chunks=n_adaln_chunks,
+            modulation_bias=modulation_bias,
+            use_second_norm=False,  # Single block only uses one norm
+            init_method=nn.init.normal_,
+        )
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        context: Optional[Tensor] = None,
+        context_mask: Optional[Tensor] = None,
+        rotary_pos_emb: Optional[Tensor] = None,
+        timestep_emb: Optional[Tensor] = None,
+        packed_seq_params=None,
+        **kwargs,
+    ) -> Tuple[Tensor, None]:
+        """
+        Single-stream block: first single layer may receive separate image and text streams
+        (concatenates them); later layers see one concatenated stream with context=None.
+
+        Args:
+            hidden_states: Image tokens [seq_img, B, H] or concatenated [seq_total, B, H]
+            context: Text tokens [seq_txt, B, H] for the first single block only
+            timestep_emb: Timestep conditioning [B, H] (required)
+
+        Returns:
+            Tuple of (hidden_states, None)
+        """
+        if timestep_emb is None:
+            raise ValueError("FluxSingleTransformerBlock requires timestep_emb for AdaLN conditioning.")
+
+        emb = timestep_emb
+
+        # TRANSITION HANDLING: Concatenate if this is first single block
+        if context is not None:
+            # This is layer 19 - concatenate text and image
+            hidden_states = torch.cat([context, hidden_states], dim=0)
+
+        residual = hidden_states
+
+        # Get modulation parameters (shift, scale, gate)
+        shift, scale, gate = self.adaln(emb)
+
+        # Apply modulated layer normalization
+        norm_hidden_states = self.adaln.modulated_layernorm(hidden_states, shift=shift, scale=scale)
+
+        # MLP path
+        mlp_hidden_states, mlp_bias = self.mlp(norm_hidden_states)
+
+        # Attention path
+        attention_output, attention_bias = self.self_attention(
+            norm_hidden_states,
+            attention_mask=attention_mask,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+
+        # Combine MLP and attention (parallel paths)
+        hidden_states = mlp_hidden_states + mlp_bias + attention_output
+        if attention_bias is not None:
+            hidden_states = hidden_states + attention_bias
+
+        # Gated residual connection
+        hidden_states = self.adaln.scale_add(residual, x=hidden_states, gate=gate)
+
+        return hidden_states, None
+
+    def __call__(self, *args, **kwargs):
+        """Override call to support CUDA graphs."""
+        if hasattr(self, "cudagraph_manager"):
+            return self.cudagraph_manager(self, args, kwargs)
+        return super(MegatronModule, self).__call__(*args, **kwargs)
+
+
+def get_flux_single_transformer_spec_for_backend(
+    backend: BackendSpecProvider,
+) -> ModuleSpec:
+    """
+    Get ModuleSpec for Flux single transformer block with backend support.
+
+    This factory function creates the specification for a Flux single block
+    using the provided backend spec provider for layer implementations.
+
+    Args:
+        backend: BackendSpecProvider (e.g., TESpecProvider, PrimusTurboSpecProvider)
+
+    Returns:
+        ModuleSpec for FluxSingleTransformerBlock with backend submodules
+    """
+    return ModuleSpec(
+        module=FluxSingleTransformerBlock,
+        submodules=TransformerLayerSubmodules(
+            self_attention=ModuleSpec(
+                module=FluxSingleAttention,
+                params={"attn_mask_type": AttnMaskType.no_mask},
+                submodules=SelfAttentionSubmodules(
+                    linear_qkv=backend.column_parallel_linear(),
+                    core_attention=backend.core_attention(),
+                    q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    linear_proj=backend.row_parallel_linear(),
+                ),
+            ),
+            mlp=ModuleSpec(
+                module=MLP,
+                submodules=MLPSubmodules(
+                    linear_fc1=backend.column_parallel_linear(),
+                    linear_fc2=backend.row_parallel_linear(),
+                ),
+            ),
+        ),
+    )
+
+
+def get_flux_double_transformer_spec_for_backend(
+    backend: BackendSpecProvider,
+) -> ModuleSpec:
+    """
+    Get ModuleSpec for Flux double (joint) transformer block with backend support.
+
+    This factory function creates the specification for a Flux MMDiT layer
+    using the provided backend spec provider for layer implementations.
+
+    Args:
+        backend: BackendSpecProvider (e.g., TESpecProvider, PrimusTurboSpecProvider)
+
+    Returns:
+        ModuleSpec for MMDiTLayer with backend submodules
+    """
+    return ModuleSpec(
+        module=MMDiTLayer,
+        submodules=TransformerLayerSubmodules(
+            self_attention=ModuleSpec(
+                module=JointSelfAttention,
+                params={"attn_mask_type": AttnMaskType.no_mask},
+                submodules=JointSelfAttentionSubmodules(
+                    q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    added_q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    added_k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True),
+                    linear_qkv=backend.column_parallel_linear(),
+                    added_linear_qkv=backend.column_parallel_linear(),
+                    core_attention=backend.core_attention(),
+                    linear_proj=backend.row_parallel_linear(),
+                ),
+            ),
+            mlp=ModuleSpec(
+                module=MLP,
+                submodules=MLPSubmodules(
+                    linear_fc1=backend.column_parallel_linear(),
+                    linear_fc2=backend.row_parallel_linear(),
+                ),
+            ),
+        ),
+    )
+
+
+def get_flux_layer_spec(
+    config,
+    backend: Optional[BackendSpecProvider] = None,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
+):
+    """
+    Create heterogeneous layer specifications for Flux TransformerBlock.
+
+    Builds a list of ModuleSpec objects for a unified TransformerBlock:
+    - Layers 0 to num_joint_layers-1: MMDiTLayer specs (joint image-text processing)
+    - Layers num_joint_layers to (num_joint_layers + num_single_layers - 1):
+      FluxSingleTransformerBlock specs (concatenated processing)
+
+    The layer specs are automatically sliced for pipeline parallelism based on
+    pp_rank and vp_stage, following the pattern from get_gpt_heterogeneous_layer_spec.
+
+    Args:
+        config: FluxConfig with num_joint_layers, num_single_layers, hidden_size, etc.
+        backend: BackendSpecProvider (auto-selected based on config.transformer_impl
+            and available providers if None)
+        vp_stage: Virtual pipeline stage number (for interleaved PP, optional)
+        pp_rank: Pipeline parallel rank (optional, auto-detected if None)
+
+    Returns:
+        TransformerBlockSubmodules containing layer_specs and layer_norm
+
+    Example:
+        >>> from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+        >>> config = FluxConfig.flux_12b()
+        >>> backend = TESpecProvider()
+        >>> spec = get_flux_layer_spec(config, backend=backend)
+        >>> transformer = TransformerBlock(config=config, spec=spec)
+
+    Reference:
+        Similar to megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py
+    """
+    from megatron.core import parallel_state
+    from megatron.core.transformer.transformer_block import (
+        TransformerBlockSubmodules,
+        get_num_layers_to_build,
+    )
+    from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+
+    # Default backend selection based on config
+    sensitive_backend = None
+
+    if backend is None:
+        if config.transformer_impl == "local":
+            if config.fp4 is not None and PrimusTurboMXFP4LocalSpecProvider is not None:
+                backend = PrimusTurboMXFP4LocalSpecProvider()
+
+                # Resolve sensitive layer backend
+                sensitive_precision = getattr(config, "sensitive_layer_precision", "bf16")
+                if sensitive_precision == "tw_fp8":
+                    sensitive_backend = PrimusTurboFloat8LocalSpecProvider()
+                elif sensitive_precision == "bf16":
+                    sensitive_backend = PrimusTurboLocalSpecProvider()
+            elif (
+                config.fp8 is not None
+                and HAVE_PRIMUS_TURBO_LOCAL
+                and PrimusTurboFloat8LocalSpecProvider is not None
+            ):
+                backend = PrimusTurboFloat8LocalSpecProvider()
+            elif HAVE_PRIMUS_TURBO_LOCAL and PrimusTurboLocalSpecProvider is not None:
+                backend = PrimusTurboLocalSpecProvider()
+            else:
+                from megatron.core.models.backends import LocalSpecProvider
+
+                backend = LocalSpecProvider()
+        elif HAVE_TE_SPEC_PROVIDER:
+            # Use TransformerEngine (current default)
+            backend = TESpecProvider()
+        else:
+            # Fallback to pure native Megatron
+            from megatron.core.models.backends import LocalSpecProvider
+
+            backend = LocalSpecProvider()
+
+    # Build per-layer specs with optional sensitive-layer heterogeneity
+    sensitive_enabled = getattr(config, "sensitive_layers_enabled", False)
+    num_start = getattr(config, "sensitive_layers_start", 0) if sensitive_enabled else 0
+    num_end = getattr(config, "sensitive_layers_end", 0) if sensitive_enabled else 0
+    total = config.num_joint_layers + config.num_single_layers
+
+    layer_specs = []
+    for i in range(total):
+        is_sensitive = sensitive_backend is not None and ((i < num_start) or (i >= total - num_end))
+        layer_backend = sensitive_backend if is_sensitive else backend
+
+        if i < config.num_joint_layers:
+            layer_specs.append(get_flux_double_transformer_spec_for_backend(layer_backend))
+        else:
+            layer_specs.append(get_flux_single_transformer_spec_for_backend(layer_backend))
+
+    # Slice for pipeline parallelism (only if parallel state is initialized)
+    try:
+        if parallel_state.model_parallel_is_initialized():
+            offset = get_transformer_layer_offset(config, vp_stage=vp_stage, pp_rank=pp_rank)
+            num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage, pp_rank=pp_rank)
+            layer_specs = layer_specs[offset : offset + num_layers_to_build]
+    except (AssertionError, RuntimeError):
+        # Parallel state not initialized - use all layers (for single-device testing)
+        pass
+
+    # Get layer norm from backend
+    layer_norm = backend.layer_norm(rms_norm=False, for_qk=False)
+
+    return TransformerBlockSubmodules(layer_specs=layer_specs, layer_norm=layer_norm)
diff --git a/primus/backends/megatron/core/models/diffusion/flux/layers.py b/primus/backends/megatron/core/models/diffusion/flux/layers.py
new file mode 100644
index 000000000..62ae3eb7c
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/layers.py
@@ -0,0 +1,130 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Flux-specific layers and components.
+
+This module implements components unique to the Flux architecture:
+    - EmbedND: Multi-dimensional Rotary Position Embedding (3D RoPE)
+    - Helper functions for position encoding
+
+Reference:
+    - Flux Paper: "Flux: A Scalable Diffusion Model"
+    - RoPE Paper: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+    - Adapted from NeMo's Flux layers
+"""
+
+from typing import List
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """
+    Generate RoPE (Rotary Position Embedding) frequencies for given positions.
+
+    This is adapted for Megatron-Core's attention implementation, which
+    calculates sin/cos internally. We only generate the frequency matrix here.
+
+    Args:
+        pos: Position indices [..., n] where n is the number of positions
+        dim: Dimension of the embedding (must be even)
+        theta: Base for frequency computation (typically 10000)
+
+    Returns:
+        Frequency matrix [..., n, dim/2] for RoPE computation
+
+    Note:
+        This differs from standard RoPE implementations because Megatron
+        attention applies sin/cos internally, so we only provide frequencies.
+
+    Reference:
+        - RoPE: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+        - Adapted from NeMo's Flux layers
+    """
+    if dim % 2 != 0:
+        raise ValueError("The dimension must be even for RoPE.")
+
+    # Compute scaling factors for each dimension pair
+    # scale: [0, 2, 4, ..., dim-2] / dim = [0, 1/dim, 2/dim, ..., (dim-2)/dim]
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+
+    # Compute base frequencies: 1 / (theta ^ scale)
+    omega = 1.0 / (theta**scale)
+
+    # Outer product of positions and frequencies
+    # pos: [..., n], omega: [dim/2] -> out: [..., n, dim/2]
+    out = torch.einsum("...n,d->...nd", pos, omega)
+
+    return out.float()
+
+
+class EmbedND(nn.Module):
+    """
+    Multi-dimensional Rotary Position Embedding (RoPE) for images.
+
+    Flux uses 3D RoPE with three axes:
+        - Axis 0: Always 0 for 2D images (reserved for video/temporal)
+        - Axis 1: Height positions (y-coordinates)
+        - Axis 2: Width positions (x-coordinates)
+
+    Each axis gets independent sinusoidal frequencies, concatenated to form
+    the complete position embedding.
+
+    Args:
+        dim: Model hidden dimension (stored for reference, not used in forward)
+        theta: Base for frequency computation (default: 10000)
+        axes_dim: Dimensions per axis, e.g., [16, 56, 56]
+
+    Input:
+        ids: Position IDs [B, seq, num_axes] -- for images: [B, H*W, 3]
+
+    Output:
+        RoPE frequency matrix [seq, B, 1, dim] for Megatron attention
+
+    Reference:
+        - Flux: 3D RoPE for spatial + channel position encoding
+        - RoPE: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+        - Adapted from NeMo's Flux layers
+    """
+
+    def __init__(self, dim: int, theta: int, axes_dim: List[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        """
+        Compute 3D RoPE frequencies from position IDs.
+
+        Args:
+            ids: Position IDs [B, seq, num_axes] where num_axes = len(axes_dim)
+
+        Returns:
+            RoPE frequency matrix [seq, B, 1, dim] for Megatron attention
+        """
+        n_axes = ids.shape[-1]
+
+        # Generate RoPE frequencies for each axis and concatenate
+        # For each axis i:
+        #   - ids[..., i]: positions for that axis [B, seq]
+        #   - rope(): generates frequencies [B, seq, axes_dim[i]/2]
+        # Concatenate along last dimension to get [B, seq, dim/2]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-1,
+        )
+
+        # Reshape for Megatron attention format:
+        # 1. Add sequence dimension: [B, seq, dim/2] -> [B, 1, seq, dim/2]
+        # 2. Permute to [seq, B, 1, dim/2]
+        emb = emb.unsqueeze(1).permute(2, 0, 1, 3)
+
+        # Stack [cos, sin] pairs and reshape to final format
+        # torch.stack([emb, emb], dim=-1): [seq, B, 1, dim/2, 2]
+        # reshape: [seq, B, 1, dim]
+        return torch.stack([emb, emb], dim=-1).reshape(*emb.shape[:-1], -1)
diff --git a/primus/backends/megatron/core/models/diffusion/flux/model.py b/primus/backends/megatron/core/models/diffusion/flux/model.py
new file mode 100644
index 000000000..b451045fa
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/model.py
@@ -0,0 +1,857 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Flux diffusion model implementation.
+
+This module implements the complete Flux model, integrating all components:
+- Input embeddings and position encodings
+- MMDiT (joint image-text) transformer blocks
+- Single (image-only) transformer blocks
+- Output projection
+
+API Note:
+    Primus's Flux model uses Megatron-Core's TransformerBlock for unified
+    layer management, enabling heterogeneous layer types (MMDiT + Single)
+    in a single container.
+
+    Architecture:
+        - Input: Packed latents [S, B, C*4] (pre-processed)
+        - Processing: Unified TransformerBlock with heterogeneous specs
+        - Output: Packed predictions [S, B, C*4]
+        - S = H*W/4 (2x2 spatial patches grouped into sequence tokens)
+
+    Key Enhancement (vs traditional ModuleList approach):
+        - Unified TransformerBlock with layer_specs for heterogeneous layers
+        - Automatic pipeline parallelism splitting
+        - Cleaner distributed checkpoint format
+
+Reference:
+    - Flux Paper: "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis"
+    - Megatron-Core: Heterogeneous TransformerBlock patterns
+"""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.utils import sharded_state_dict_default
+from torch import Tensor
+
+from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import (
+    DiffusionModule,
+)
+from primus.backends.megatron.core.models.diffusion.common.embeddings import (
+    MLPEmbedder,
+    TimeStepEmbedder,
+)
+from primus.backends.megatron.core.models.diffusion.common.normalization import (
+    AdaLNContinuous,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    get_flux_layer_spec,
+)
+from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND
+from primus.backends.megatron.core.transformer.diffusion_transformer_block import (
+    DiffusionTransformerBlock,
+)
+
+_QK_RMSNORM_PARAM_ATTRS = (
+    "sequence_parallel",
+    "shared",
+    "allreduce",
+    "tensor_model_parallel",
+    "partition_dim",
+    "partition_stride",
+)
+
+
+class _TorchRMSNorm(nn.Module):
+    """Compile-friendly RMSNorm replacement for TE's QK LayerNorm.
+
+    TE's TENorm is decorated with @no_torch_dynamo, causing graph breaks
+    inside compiled regions. This pure-PyTorch version is fully traceable.
+    Matches NeMo's _TorchRMSNorm (custom_flux.py).
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6, *, device=None, dtype=None):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size, device=device, dtype=dtype))
+        self.bias = None
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+
+    def forward(self, x: Tensor) -> Tensor:
+        mean_sq = x.square().mean(dim=-1, keepdim=True, dtype=torch.float32)
+        inv_rms = torch.rsqrt(mean_sq + self.eps).to(dtype=x.dtype)
+        return x * inv_rms * self.weight
+
+
+def _build_local_qk_rmsnorm(norm: nn.Module) -> "_TorchRMSNorm":
+    """Build a _TorchRMSNorm from an existing TE norm, preserving weights and attrs."""
+    weight = norm.weight
+    local_norm = _TorchRMSNorm(
+        weight.numel(),
+        eps=getattr(norm, "eps", 1e-6),
+        device=weight.device,
+        dtype=weight.dtype,
+    )
+    local_norm.weight.data.copy_(weight.data)
+    local_norm.weight.requires_grad_(weight.requires_grad)
+    for attr_name in _QK_RMSNORM_PARAM_ATTRS:
+        if hasattr(weight, attr_name):
+            setattr(local_norm.weight, attr_name, getattr(weight, attr_name))
+    return local_norm
+
+
+class Flux(DiffusionModule):
+    """
+    Flux: Flow-based diffusion model with MMDiT architecture.
+
+    Uses Megatron-Core's TransformerBlock with heterogeneous layer specs,
+    combining MMDiT "double blocks" (joint image-text) with "single blocks"
+    (image-only) for text-to-image generation.
+
+    Model Variants:
+        - Flux 535M: 1 joint + 1 single layer (~535M params, testing)
+        - Flux 12B: 19 joint + 38 single layers (~12B params, production)
+
+    Args:
+        config: FluxConfig with all model parameters
+        encoder_configs: Optional encoder configurations (VAE, T5, CLIP)
+        pg_collection: ProcessGroupCollection for distributed training
+        backend: Optional BackendSpecProvider for layer implementations (default: auto-selected)
+
+    Forward args:
+        img: Packed image latents [S_img, B, C*4] from VAE
+        txt: Text embeddings [S_txt, B, D_txt] from T5-XXL
+        y: CLIP pooled embeddings [B, D_pool]
+        timesteps: Diffusion timesteps [B] in [0, 1]
+        img_ids, txt_ids: Position IDs [B, S, 3]
+        guidance: Optional guidance scale [B]
+
+    Returns:
+        Predicted velocity [S_img, B, C*4] in packed format
+
+    torch.compile:
+        Set enable_torch_compile=True in config. Compilation is applied AFTER
+        FSDP/DDP wrapping by the trainer -- do NOT call compile_model() manually.
+
+    Reference:
+        - "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis"
+        - Megatron-Core TransformerBlock
+    """
+
+    def __init__(
+        self,
+        config: FluxConfig,
+        encoder_configs: Optional[Dict[str, Any]] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        backend: Optional["BackendSpecProvider"] = None,
+    ):
+        super().__init__(config, pg_collection=pg_collection, encoder_configs=encoder_configs)
+
+        self.out_channels = config.in_channels
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.patch_size = config.patch_size
+        self.in_channels = config.in_channels
+        self.guidance_embed = config.guidance_embed
+
+        # Position embedding (3D RoPE for image patches)
+        # axes_dim=[16, 56, 56] for 1024x1024 images with 64 channels
+        self.pos_embed = EmbedND(
+            dim=self.hidden_size,
+            theta=config.theta,
+            axes_dim=list(config.axes_dim),
+        )
+
+        # Input embeddings
+        # img_embed accepts packed format with in_channels already set to C*4 in config
+        # Note: config.in_channels = 64 = 16 (VAE channels) * 4 (packing factor)
+        self.img_embed = nn.Linear(config.in_channels, self.hidden_size)
+        self.txt_embed = nn.Linear(config.context_dim, self.hidden_size)
+
+        # Conditioning embeddings
+        self.timestep_embedding = TimeStepEmbedder(
+            embedding_dim=config.model_channels,
+            hidden_dim=self.hidden_size,
+        )
+        self.vector_embedding = MLPEmbedder(
+            in_dim=config.vec_in_dim,
+            hidden_dim=self.hidden_size,
+        )
+
+        # Optional guidance embedding for classifier-free guidance
+        if config.guidance_embed:
+            self.guidance_embedding = MLPEmbedder(
+                in_dim=config.model_channels,
+                hidden_dim=self.hidden_size,
+            )
+        else:
+            self.guidance_embedding = nn.Identity()
+
+        # Create unified DiffusionTransformerBlock with heterogeneous layers
+        # Layers 0-18: MMDiTLayer (joint processing)
+        # Layers 19-56: FluxSingleTransformerBlock (single processing)
+        self.transformer = DiffusionTransformerBlock(
+            config=config,
+            spec=get_flux_layer_spec(config, backend=backend),
+            post_layer_norm=False,  # Flux uses AdaLN, not standard final layernorm
+            pre_process=True,
+            post_process=True,
+        )
+
+        # Output layers
+        self.norm_out = AdaLNContinuous(
+            config=config,
+            conditioning_embedding_dim=self.hidden_size,
+        )
+        self.proj_out = nn.Linear(
+            self.hidden_size,
+            self.patch_size * self.patch_size * self.out_channels,
+            bias=True,
+        )
+
+        # NOTE: torch.compile is NOT applied here. It must be applied AFTER
+        # distributed wrapping (FSDP/DDP). The training framework will call
+        # compile_model() at the appropriate time.
+
+        # Stack runner references (replaced with compiled versions by compile_model)
+        self._double_block_stack_runner = self._run_double_block_stack
+        self._single_block_stack_runner = self._run_single_block_stack
+        self._output_head_runner = self._run_output_head
+        self._full_dit_runner = self._run_full_dit
+
+        # Replace TE QK RMSNorm BEFORE DDP wrapping so that DDP's forward
+        # pre-hook registration sees the final module tree.
+        if config.torch_compile_replace_qk_rmsnorm and config.transformer_impl == "transformer_engine":
+            try:
+                from primus.core.utils.module_utils import log_rank_0 as _log
+            except Exception:
+                _log = print
+            self._replace_qk_rmsnorm(_log)
+
+        self.init_weights()
+
+    def init_weights(self):
+        """
+        Custom weight initialization matching NVIDIA MLPerf Training v5.1.
+
+        - img_embed, txt_embed: Xavier uniform
+        - timestep_embedding, vector_embedding: Normal(std=0.02) for MLP layers
+        - Per-block AdaLN modulations: re-zeroed after construction (construction
+          now uses normal_ init_method to match NeMo's RNG sequence)
+        - norm_out.adaLN_modulation: zero-init (last linear layer)
+        - norm_out.norm: reset to default
+        - proj_out: zero-init weight and bias
+        """
+        from primus.core.utils.module_utils import log_rank_0
+
+        # Embedders: Xavier uniform
+        nn.init.xavier_uniform_(self.img_embed.weight)
+        nn.init.constant_(self.img_embed.bias, 0)
+        nn.init.xavier_uniform_(self.txt_embed.weight)
+        nn.init.constant_(self.txt_embed.bias, 0)
+
+        # Timestep embedder MLP: Normal(std=0.02)
+        self._init_mlpembedder(self.timestep_embedding.time_embedding)
+        self._init_mlpembedder(self.vector_embedding)
+
+        if self.guidance_embed and not isinstance(self.guidance_embedding, nn.Identity):
+            self._init_mlpembedder(self.guidance_embedding)
+
+        # Per-block AdaLN: zero modulation weights and reset LayerNorms.
+        # Order matches NeMo: single blocks first, then double blocks.
+        # These are deterministic ops (no RNG consumption).
+        for layer in self.transformer.layers[self.config.num_joint_layers :]:
+            nn.init.constant_(layer.adaln.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(layer.adaln.adaLN_modulation[-1].bias, 0)
+            layer.adaln.ln.reset_parameters()
+
+        for layer in self.transformer.layers[: self.config.num_joint_layers]:
+            nn.init.constant_(layer.adaln.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(layer.adaln.adaLN_modulation[-1].bias, 0)
+            layer.adaln.ln.reset_parameters()
+            layer.adaln.ln2.reset_parameters()
+            nn.init.constant_(layer.adaln_context.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(layer.adaln_context.adaLN_modulation[-1].bias, 0)
+            layer.adaln_context.ln.reset_parameters()
+            layer.adaln_context.ln2.reset_parameters()
+
+        # Output layers: zero-init for stable training start
+        nn.init.constant_(self.proj_out.weight, 0)
+        nn.init.constant_(self.proj_out.bias, 0)
+
+        # norm_out (AdaLNContinuous): zero-init modulation, reset norm
+        nn.init.constant_(self.norm_out.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.norm_out.adaLN_modulation[-1].bias, 0)
+        self.norm_out.norm.reset_parameters()
+
+        log_rank_0("Applied custom weight initialization (MLPerf v5.1 aligned, NeMo RNG-matched)")
+
+    @staticmethod
+    def _init_mlpembedder(module, init_std: float = 0.02):
+        """Initialize MLPEmbedder with Normal(std) for both linear layers."""
+        nn.init.normal_(module.in_layer.weight, std=init_std)
+        nn.init.constant_(module.in_layer.bias, 0)
+        nn.init.normal_(module.out_layer.weight, std=init_std)
+        nn.init.constant_(module.out_layer.bias, 0)
+
+    def compile_model(self):
+        """
+        Apply torch.compile to the model using the configured strategy.
+
+        CRITICAL: This must be called AFTER distributed wrapping (FSDP/DDP) is complete
+        AND after ddp_config is set. Do NOT call this in __init__. The training
+        framework will call this at the appropriate time.
+
+        Strategies:
+            per_block:   compile each transformer layer individually (default, recommended)
+            whole_model: compile entire forward (incompatible with overlap_param_gather)
+            double_stack: compile the double (joint) block loop
+            single_stack: compile the single block loop
+            stack:        compile both double and single block loops
+            full_dit:     compile entire DiT as one region
+        """
+        import os
+
+        if not self.config.enable_torch_compile:
+            return
+
+        try:
+            from primus.core.utils.module_utils import log_rank_0
+
+            log = log_rank_0
+        except Exception:
+            log = print
+
+        strategy = self.config.torch_compile_strategy
+        valid_strategies = {"whole_model", "per_block", "double_stack", "single_stack", "stack", "full_dit"}
+        if strategy not in valid_strategies:
+            raise ValueError(
+                f"Invalid torch_compile_strategy='{strategy}'. " f"Must be one of: {sorted(valid_strategies)}"
+            )
+
+        is_te = self.config.transformer_impl == "transformer_engine"
+
+        if not is_te and strategy not in ("whole_model", "per_block"):
+            log(
+                f"WARNING: torch_compile_strategy='{strategy}' is designed for TE spec. "
+                f"Local spec modules are already torch.compile-friendly; "
+                f"falling back to 'per_block' strategy."
+            )
+            strategy = "per_block"
+
+        if (
+            strategy not in ("whole_model", "per_block")
+            and getattr(self.config, "recompute_granularity", None) == "full"
+        ):
+            raise ValueError(
+                f"Activation checkpointing (recompute_granularity='full') is incompatible "
+                f"with torch_compile_strategy='{strategy}'. Stack runners bypass the "
+                f"checkpointed forward path."
+            )
+
+        if strategy == "whole_model":
+            from megatron.training import get_args
+
+            args = get_args()
+            if getattr(args, "overlap_param_gather", False):
+                raise ValueError(
+                    "whole_model compile strategy is incompatible with overlap_param_gather. "
+                    "DDP hooks are traced inside the compiled graph, causing ~20% convergence "
+                    "degradation. Use 'per_block' strategy instead, or disable overlap_param_gather."
+                )
+
+        if strategy == "per_block" and getattr(self.config, "enable_cuda_graph", False):
+            log(
+                "WARNING: per_block compile strategy may conflict with per-layer "
+                "CUDA graph __call__ overrides on MMDiTLayer/FluxSingleTransformerBlock."
+            )
+
+        compile_kwargs = {
+            "backend": self.config.torch_compile_backend,
+            "mode": self.config.torch_compile_mode,
+            "fullgraph": self.config.torch_compile_fullgraph,
+        }
+
+        if is_te and self.config.torch_compile_disable_inductor_cudagraphs:
+            torch._inductor.config.triton.cudagraphs = False
+            torch._inductor.config.triton.cudagraph_trees = False
+            os.environ["TORCHINDUCTOR_CUDAGRAPHS"] = "0"
+            log("  Disabled Inductor CUDA graphs (TE FP8 compatibility)")
+
+        if self.config.torch_compile_emulate_precision_casts:
+            torch._inductor.config.emulate_precision_casts = True
+            log(
+                "  Enabled emulate_precision_casts "
+                "(preserves eager BF16 precision semantics in fused Triton kernels)"
+            )
+
+        if not self.config.torch_compile_fused_ln_modulate:
+            for m in self.modules():
+                if hasattr(m, "use_fused_ln_modulate") and not getattr(m, "_adaln_plain_ops", False):
+                    m.use_fused_ln_modulate = False
+            log("  Disabled fused LN+modulate Triton kernels (using separate opaque ops)")
+
+        if strategy == "whole_model":
+            self._compile_whole_model(compile_kwargs, log)
+        elif strategy == "per_block":
+            log("=" * 80)
+            log("Applying PER-BLOCK torch.compile on Flux...")
+            for i, layer in enumerate(self.transformer.layers):
+                layer.forward = torch.compile(layer.forward, **compile_kwargs)
+            log(f"  Compiled {len(self.transformer.layers)} layers individually")
+            log("=" * 80)
+        elif strategy == "double_stack":
+            log("=" * 80)
+            log("Applying DOUBLE-STACK torch.compile on Flux...")
+            self._double_block_stack_runner = torch.compile(self._run_double_block_stack, **compile_kwargs)
+            log(f"  Compiled double block stack ({self.config.num_joint_layers} layers)")
+            log("=" * 80)
+        elif strategy == "single_stack":
+            log("=" * 80)
+            log("Applying SINGLE-STACK torch.compile on Flux...")
+            self._single_block_stack_runner = torch.compile(self._run_single_block_stack, **compile_kwargs)
+            log(f"  Compiled single block stack ({self.config.num_single_layers} layers)")
+            log("=" * 80)
+        elif strategy == "stack":
+            log("=" * 80)
+            log("Applying STACK torch.compile on Flux (double + single)...")
+            self._double_block_stack_runner = torch.compile(self._run_double_block_stack, **compile_kwargs)
+            self._single_block_stack_runner = torch.compile(self._run_single_block_stack, **compile_kwargs)
+            log(
+                f"  Compiled double ({self.config.num_joint_layers}) "
+                f"+ single ({self.config.num_single_layers}) block stacks"
+            )
+            log("=" * 80)
+        elif strategy == "full_dit":
+            log("=" * 80)
+            log("Applying FULL-DIT torch.compile on Flux...")
+            self._full_dit_runner = torch.compile(self._run_full_dit, **compile_kwargs)
+            log("  Compiled entire DiT as one region")
+            log("=" * 80)
+
+    def _compile_whole_model(self, compile_kwargs, log):
+        """Apply a single torch.compile on self.forward (whole-model strategy)."""
+        import torch
+
+        log("=" * 80)
+        log("Applying WHOLE-MODEL torch.compile on Flux...")
+        log(f"  Backend: {compile_kwargs['backend']}")
+        log(f"  Mode: {compile_kwargs['mode']}")
+        log(f"  Fullgraph: {compile_kwargs['fullgraph']}")
+        log("=" * 80)
+        self.forward = torch.compile(self.forward, **compile_kwargs)
+        log("  Done (actual compilation deferred to first forward pass)")
+        log("=" * 80)
+
+    def _replace_qk_rmsnorm(self, log):
+        """Replace TE QK RMSNorm modules with compile-friendly _TorchRMSNorm.
+
+        Only meaningful for transformer_impl="transformer_engine" where QK norms
+        are TENorm instances (opaque to torch.compile). Local spec already uses
+        torch.nn.RMSNorm which is fully traceable.
+        """
+        count = 0
+        for layer in self.transformer.layers:
+            attn = getattr(layer, "self_attention", None)
+            if attn is None:
+                continue
+            for attr in ("q_layernorm", "k_layernorm", "added_q_layernorm", "added_k_layernorm"):
+                old_norm = getattr(attn, attr, None)
+                if old_norm is not None and not isinstance(old_norm, _TorchRMSNorm):
+                    setattr(attn, attr, _build_local_qk_rmsnorm(old_norm))
+                    count += 1
+        log(f"  Replaced {count} TE QK RMSNorm modules with compile-friendly _TorchRMSNorm")
+
+    # ------------------------------------------------------------------
+    # Embedding computation: compilable helper that converts raw inputs
+    # into hidden states, conditioning vectors, and RoPE frequencies.
+    # ------------------------------------------------------------------
+
+    def _compute_embeddings(self, img, txt, timesteps, img_ids, txt_ids, guidance, y):
+        """Compute all embeddings and RoPE in a single compilable region."""
+        hidden_states = self.img_embed(img)
+        encoder_hidden_states = self.txt_embed(txt)
+
+        if hidden_states.dim() == 3 and hidden_states.shape[0] < hidden_states.shape[1]:
+            hidden_states = hidden_states.transpose(0, 1)
+        if (
+            encoder_hidden_states.dim() == 3
+            and encoder_hidden_states.shape[0] < encoder_hidden_states.shape[1]
+        ):
+            encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
+
+        txt_seq_len = encoder_hidden_states.shape[0]
+
+        timesteps = timesteps.to(hidden_states.dtype) * 1000.0
+        vec_emb = self.timestep_embedding(timesteps)
+
+        if guidance is not None:
+            guidance_emb = self.guidance_embedding(self.timestep_embedding.time_proj(guidance * 1000.0))
+            vec_emb = vec_emb + guidance_emb
+
+        vec_emb = vec_emb + self.vector_embedding(y)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        rotary_pos_emb = self.pos_embed(ids)
+
+        return hidden_states, encoder_hidden_states, vec_emb, rotary_pos_emb, txt_seq_len
+
+    # ------------------------------------------------------------------
+    # Stack runners: self-contained compilable callables that iterate
+    # transformer layers directly (bypassing DiffusionTransformerBlock).
+    # ------------------------------------------------------------------
+
+    def _run_double_block_stack(self, hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb):
+        """Run all double (joint MMDiT) blocks as one compilable region."""
+        from megatron.core.utils import make_viewless_tensor
+
+        hidden_states = make_viewless_tensor(
+            inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+        )
+        for layer in self.transformer.layers[: self.config.num_joint_layers]:
+            hidden_states, encoder_hidden_states = layer(
+                hidden_states=hidden_states,
+                attention_mask=None,
+                context=encoder_hidden_states,
+                context_mask=None,
+                rotary_pos_emb=rotary_pos_emb,
+                timestep_emb=vec_emb,
+            )
+        return hidden_states, encoder_hidden_states
+
+    def _run_single_block_stack(self, hidden_states, rotary_pos_emb, vec_emb):
+        """Run all single blocks as one compilable region.
+
+        The caller must concatenate [context, hidden] before entering this runner.
+        All single blocks receive context=None (concat already done).
+        """
+        for layer in self.transformer.layers[self.config.num_joint_layers :]:
+            hidden_states, _ = layer(
+                hidden_states=hidden_states,
+                attention_mask=None,
+                context=None,
+                context_mask=None,
+                rotary_pos_emb=rotary_pos_emb,
+                timestep_emb=vec_emb,
+            )
+        return hidden_states
+
+    def _run_output_head(self, hidden_states, txt_seq_len, vec_emb):
+        """Run the output head (slice + norm_out + proj_out)."""
+        hidden_states = hidden_states[txt_seq_len:, ...]
+        hidden_states = self.norm_out(hidden_states, vec_emb)
+        return self.proj_out(hidden_states)
+
+    def _run_full_dit(self, hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb, txt_seq_len):
+        """Run the entire DiT (double + cat + single + output) as one compiled region."""
+        hidden_states, encoder_hidden_states = self._run_double_block_stack(
+            hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb
+        )
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0)
+        hidden_states = self._run_single_block_stack(hidden_states, rotary_pos_emb, vec_emb)
+        if self.transformer.final_layernorm is not None:
+            hidden_states = self.transformer.final_layernorm(hidden_states)
+        return self._run_output_head(hidden_states, txt_seq_len, vec_emb)
+
+    def set_input_tensor(self, input_tensor):
+        """
+        Set the input tensor on the underlying TransformerBlock.
+
+        Required by the Megatron model interface (the scheduler calls this on
+        every model). Pipeline parallelism is not supported for diffusion
+        models (rejected at config construction; PP > 1 raises), so only the
+        single-stage case is handled here.
+
+        Args:
+            input_tensor: Union[Tensor, List[Tensor], None]
+        """
+        if input_tensor is None:
+            return
+
+        # Handle list or single tensor
+        if not isinstance(input_tensor, list):
+            input_tensor = [input_tensor]
+
+        # Pass image hidden states to TransformerBlock
+        self.transformer.set_input_tensor(input_tensor[0])
+
+    def sharded_state_dict(
+        self,
+        prefix: str = "",
+        sharded_offsets: tuple = (),
+        metadata: Optional[dict] = None,
+    ) -> Dict[str, Any]:
+        """
+        Generate sharded state dictionary for distributed checkpointing.
+
+        With TransformerBlock, this delegates to the transformer which handles
+        all layer numbering automatically for heterogeneous layers.
+
+        Flux is architecturally heterogeneous: joint MMDiT blocks carry params
+        (e.g. ``self_attention.added_linear_qkv.*``, ``context_mlp.*``) that the
+        single DiT blocks intentionally omit (FluxSingleAttention rebuilds
+        ``linear_proj`` with ``bias=False``). Megatron's heterogeneous, per-layer
+        indexed sharded keys (``transformer.layers.<i>.*``) are required here:
+        the homogeneous layer-stacked path would leave unclaimed slots for the
+        params absent in single blocks and raise a CheckpointingException at
+        save time. TransformerBlock.sharded_state_dict provides this
+        automatically, so no config toggle is needed.
+
+        Args:
+            prefix: Prefix for state dict keys (e.g., 'module.')
+            sharded_offsets: Pipeline parallel offsets
+            metadata: Optional metadata for checkpoint conversion
+
+        Returns:
+            Dictionary mapping state dict keys to ShardedTensor objects
+        """
+        sharded_state_dict = {}
+
+        # Delegate transformer layers to TransformerBlock
+        # Handles heterogeneous layers automatically
+        sharded_state_dict.update(
+            self.transformer.sharded_state_dict(f"{prefix}transformer.", sharded_offsets, metadata)
+        )
+
+        # Handle other submodules (embeddings, projections, norms)
+        for name, module in self.named_children():
+            if module is not self.transformer:
+                sharded_state_dict.update(
+                    sharded_state_dict_default(module, f"{prefix}{name}.", sharded_offsets, metadata)
+                )
+
+        return sharded_state_dict
+
+    def get_fp8_context(self):
+        """
+        Return FP8 context manager for FP8 training.
+
+        When transformer_impl="local", FP8 is handled per-module at init time
+        (inside Float8ColumnParallelLinear / Float8RowParallelLinear), so no
+        global FP8 context manager is needed. This avoids mutable global state
+        that would cause torch.compile graph breaks.
+
+        Returns:
+            Context manager for FP8 operations (or nullcontext if FP8 disabled
+            or using local spec)
+        """
+        from contextlib import nullcontext
+
+        if self.config.transformer_impl == "local":
+            return nullcontext()
+
+        from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util
+
+        return get_fp8_context_util(self.config)
+
+    def forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        y: Tensor,
+        timesteps: Tensor,
+        img_ids: Tensor,
+        txt_ids: Tensor,
+        guidance: Optional[Tensor] = None,
+    ) -> Tensor:
+        """
+        Forward pass through Flux model.
+
+        Args:
+            img: Image latents [S_img, B, C*4] or [B, S_img, C*4] - PACKED format from prepare_flux_latents
+                 where S_img = H*W/4 (2x2 patches grouped)
+            txt: Text embeddings [S_txt, B, D_txt] or [B, S_txt, D_txt] from T5-XXL
+            y: CLIP pooled embeddings [B, D_pool] from CLIP-L
+            timesteps: Diffusion timesteps [B] in range [0, 1]
+            img_ids: Position IDs for image patches [B, S_img, 3]
+            txt_ids: Position IDs for text tokens [B, S_txt, 3]
+            guidance: Optional guidance scale [B] for classifier-free guidance
+
+        Returns:
+            Predicted velocity [S_img, B, C*4] or [B, S_img, C*4] in packed format (matches input format)
+
+        Note:
+            The model accepts pre-packed
+            latents from the wrapper/training code. The packing groups 2x2 spatial patches,
+            reducing spatial dimensions by 4x and increasing channels by 4x.
+        """
+        hidden_states, encoder_hidden_states, vec_emb, rotary_pos_emb, txt_seq_len = self._compute_embeddings(
+            img, txt, timesteps, img_ids, txt_ids, guidance, y
+        )
+
+        strategy = self.config.torch_compile_strategy
+
+        if strategy in ("whole_model", "per_block"):
+            # Original path: delegate to DiffusionTransformerBlock
+            with self.get_fp8_context():
+                hidden_states = self.transformer(
+                    hidden_states=hidden_states,
+                    attention_mask=None,
+                    context=encoder_hidden_states,
+                    context_mask=None,
+                    rotary_pos_emb=rotary_pos_emb,
+                    rotary_pos_cos=None,
+                    rotary_pos_sin=None,
+                    attention_bias=None,
+                    timestep_emb=vec_emb,
+                    packed_seq_params=None,
+                )
+
+            hidden_states = hidden_states[txt_seq_len:, ...]
+            hidden_states = self.norm_out(hidden_states, vec_emb)
+            output = self.proj_out(hidden_states)
+
+        elif strategy == "full_dit":
+            # One large compiled region covering double + cat + single + output
+            with self.get_fp8_context():
+                output = self._full_dit_runner(
+                    hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb, txt_seq_len
+                )
+
+        else:
+            # stack / double_stack / single_stack: compiled stack runners
+            # with eager transitions between them
+            with self.get_fp8_context():
+                hidden_states, encoder_hidden_states = self._double_block_stack_runner(
+                    hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb
+                )
+
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0)
+
+            with self.get_fp8_context():
+                hidden_states = self._single_block_stack_runner(hidden_states, rotary_pos_emb, vec_emb)
+
+            if self.transformer.final_layernorm is not None:
+                hidden_states = self.transformer.final_layernorm(hidden_states)
+
+            output = self._output_head_runner(hidden_states, txt_seq_len, vec_emb)
+
+        return output
+
+    def load_checkpoint(
+        self,
+        checkpoint_path: Union[str, Path],
+        convert_from_hf: bool = False,
+        save_converted_to: Optional[Union[str, Path]] = None,
+        strict: bool = False,
+    ) -> Tuple[List[str], List[str]]:
+        """
+        Load checkpoint into Flux model.
+
+        Supports both Primus native checkpoints and HuggingFace checkpoints
+        with automatic conversion.
+
+        Note: HuggingFace checkpoint conversion is handled by the checkpoint_converter module.
+
+        Args:
+            checkpoint_path: Path to checkpoint file or directory
+            convert_from_hf: If True, convert from HuggingFace format
+            save_converted_to: Optional path to save converted checkpoint
+            strict: Whether to use strict state dict loading
+
+        Returns:
+            Tuple of (missing_keys, unexpected_keys)
+
+        Example:
+            >>> config = FluxConfig.flux_12b()
+            >>> model = Flux(config=config)
+            >>> # Load native Primus checkpoint
+            >>> missing, unexpected = model.load_checkpoint("primus_flux_12b.safetensors")
+        """
+        from safetensors.torch import load_file as load_safetensors
+
+        from primus.core.utils.module_utils import (
+            error_rank_0,
+            log_rank_0,
+            warning_rank_0,
+        )
+
+        if convert_from_hf:
+            try:
+                from .checkpoint_converter import convert_hf_checkpoint
+
+                log_rank_0(f"Converting HuggingFace checkpoint: {checkpoint_path}")
+                state_dict = convert_hf_checkpoint(
+                    checkpoint_path,
+                    flux_config=self.config,
+                    save_to=save_converted_to,
+                )
+            except ImportError:
+                raise NotImplementedError(
+                    "HuggingFace checkpoint conversion requires the checkpoint_converter module. "
+                    "Please use a pre-converted Primus checkpoint."
+                )
+        else:
+            # Load native Primus checkpoint
+            log_rank_0(f"Loading Primus checkpoint: {checkpoint_path}")
+            checkpoint_path = Path(checkpoint_path)
+
+            if checkpoint_path.is_dir():
+                # Load all .safetensors files
+                safetensor_files = list(checkpoint_path.glob("*.safetensors"))
+                if not safetensor_files:
+                    raise FileNotFoundError(f"No .safetensors files in {checkpoint_path}")
+
+                state_dict = {}
+                for file in safetensor_files:
+                    state_dict.update(load_safetensors(str(file)))
+            else:
+                state_dict = load_safetensors(str(checkpoint_path))
+
+        # Load state dict into model
+        missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=strict)
+
+        # Filter out _extra_state keys (Megatron-specific metadata, expected to be missing)
+        missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")]
+
+        # Check for critical missing keys
+        critical_patterns = ["timestep_embedding", "img_embed", "txt_embed", "vector_embedding"]
+        critical_missing = [k for k in missing_keys if any(pattern in k for pattern in critical_patterns)]
+
+        if critical_missing:
+            variant = "flux_12b" if self.config.num_joint_layers == 19 else "flux_535m"
+            err_lines = [
+                "\n" + "=" * 80,
+                "CRITICAL ERROR: Key model weights are missing!",
+                "=" * 80,
+                f"Missing {len(critical_missing)} critical keys:",
+            ]
+            err_lines += [f"  - {key}" for key in critical_missing[:10]]
+            if len(critical_missing) > 10:
+                err_lines.append(f"  ... and {len(critical_missing) - 10} more")
+            err_lines += [
+                "\nThis usually means:",
+                "  1. Checkpoint format mismatch (HuggingFace vs Primus)",
+                "  2. Checkpoint was converted with an older version",
+                "  3. Checkpoint file is corrupted or incomplete",
+                "\nTo fix, reconvert the checkpoint:",
+                "    python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\",
+                "      --input black-forest-labs/FLUX.1-dev/transformer \\",
+                f"      --output {checkpoint_path} \\",
+                f"      --variant {variant}",
+                "=" * 80,
+            ]
+            error_rank_0("\n".join(err_lines))
+
+            raise RuntimeError(
+                f"Critical model weights missing ({len(critical_missing)} keys). "
+                f"Checkpoint may be corrupted or incompatible. See details above."
+            )
+
+        if missing_keys:
+            warning_rank_0(f"Missing keys ({len(missing_keys)}): {missing_keys[:5]}...")
+        if unexpected_keys:
+            warning_rank_0(f"Unexpected keys ({len(unexpected_keys)}): {unexpected_keys[:5]}...")
+
+        log_rank_0("Checkpoint loaded successfully!")
+        return missing_keys, unexpected_keys
diff --git a/primus/backends/megatron/core/models/diffusion/flux/utils.py b/primus/backends/megatron/core/models/diffusion/flux/utils.py
new file mode 100644
index 000000000..604b7d810
--- /dev/null
+++ b/primus/backends/megatron/core/models/diffusion/flux/utils.py
@@ -0,0 +1,197 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Core utilities shared between Flux training and inference.
+
+This module provides fundamental operations used by both training and inference:
+- Latent packing/unpacking (2x2 spatial grouping)
+- Position ID generation (3D RoPE for Flux transformer)
+
+Design: Pure functions with no side effects, fully type-annotated.
+"""
+
+
+import torch
+from torch import Tensor
+
+
+def pack_latents(latents: Tensor) -> Tensor:
+    """
+    Pack latents from (B, C, H, W) to (B, H*W/4, C*4) format.
+
+    Groups 2x2 spatial patches into sequence tokens for transformer processing.
+    Used by both training and inference pipelines.
+
+    Args:
+        latents: Input tensor of shape (B, C, H, W)
+
+    Returns:
+        Packed tensor of shape (B, H*W/4, C*4)
+
+    Example:
+        >>> latents = torch.randn(2, 16, 128, 128)
+        >>> packed = pack_latents(latents)
+        >>> packed.shape
+        torch.Size([2, 4096, 64])  # 128*128/4=4096, 16*4=64
+    """
+    batch_size, num_channels, height, width = latents.shape
+
+    # Reshape to group 2x2 patches: (B, C, H, W) -> (B, C, H//2, 2, W//2, 2)
+    latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2)
+
+    # Permute to bring patch dimensions together: (B, H//2, W//2, C, 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+
+    # Flatten patches: (B, H//2*W//2, C*4)
+    latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4)
+
+    return latents
+
+
+def unpack_latents(
+    latents: Tensor,
+    height: int,
+    width: int,
+    vae_scale_factor: int = 1,
+) -> Tensor:
+    """
+    Unpack latents from (B, N, C*4) to (B, C, H, W) format.
+
+    Reverses pack_latents operation. Supports optional VAE scaling for inference.
+
+    Args:
+        latents: Packed tensor of shape (B, N, C*4)
+        height: Target spatial height
+        width: Target spatial width
+        vae_scale_factor: Downsampling factor (default: 1 for training, 8 for inference)
+
+    Returns:
+        Unpacked tensor of shape (B, C, H, W)
+
+    Example:
+        >>> packed = torch.randn(2, 4096, 64)
+        >>> unpacked = unpack_latents(packed, 1024, 1024, vae_scale_factor=8)
+        >>> unpacked.shape
+        torch.Size([2, 16, 128, 128])  # 1024/8 = 128 (VAE downscaling)
+    """
+    batch_size, num_patches, channels = latents.shape
+
+    # Apply VAE downsampling if specified (inference path)
+    if vae_scale_factor > 1:
+        height = height // vae_scale_factor
+        width = width // vae_scale_factor
+
+    # For packed latents: height and width are the dimensions BEFORE packing
+    # So we need to use height//2 and width//2 for the reshaped view
+    h_packed = height // 2
+    w_packed = width // 2
+
+    # Reshape to restore patch structure: (B, H//2, W//2, C//4, 2, 2)
+    latents = latents.view(batch_size, h_packed, w_packed, channels // 4, 2, 2)
+
+    # Permute to restore spatial dimensions: (B, C//4, H//2, 2, W//2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+    # Flatten spatial dimensions: (B, C//4, H, W)
+    latents = latents.reshape(batch_size, channels // 4, height, width)
+
+    return latents
+
+
+def generate_image_position_ids(
+    batch_size: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> Tensor:
+    """
+    Generate 3D position IDs for image patches (RoPE).
+
+    Creates position encodings for Flux's 3D RoPE where:
+    - Dimension 0: Always 0 (reserved for future temporal/video models)
+    - Dimension 1: Row index (y-coordinate, 0 to height//2-1 after packing)
+    - Dimension 2: Column index (x-coordinate, 0 to width//2-1 after packing)
+
+    Follows NeMo conventions for Flux position encoding. For video models, dimension 0
+    would encode frame indices instead of being 0.
+
+    NOTE FOR ROPE FUSION: When using fused RoPE kernels (apply_rope_fusion=True),
+    this function should be called with batch_size=1 to satisfy Transformer Engine's
+    dimension constraints (freqs tensor must have shape [S, 1, 1, D]). The resulting
+    [1, H*W/4, 3] position IDs will broadcast across the actual batch dimension
+    during attention computation. This requires all images in the batch to have the
+    same resolution (same height and width).
+
+    Args:
+        batch_size: Batch size (use 1 for RoPE fusion, actual batch size otherwise)
+        height: Latent height in unpacked format (before 2x2 packing)
+        width: Latent width in unpacked format (before 2x2 packing)
+        device: Target device
+        dtype: Target dtype
+
+    Returns:
+        Position IDs tensor of shape (batch_size, H*W/4, 3)
+
+    Example:
+        >>> ids = generate_image_position_ids(2, 128, 128, torch.device('cpu'))
+        >>> ids.shape
+        torch.Size([2, 4096, 3])  # 128*128/4 = 4096 packed positions
+    """
+    # Generate for packed latents (2x2 grouping)
+    h_packed = height // 2
+    w_packed = width // 2
+
+    # Create position grid
+    img_ids = torch.zeros(h_packed, w_packed, 3, device=device, dtype=dtype)
+
+    # Fill row dimension (dim 1)
+    img_ids[..., 1] = torch.arange(h_packed, device=device, dtype=dtype)[:, None]
+
+    # Fill column dimension (dim 2)
+    img_ids[..., 2] = torch.arange(w_packed, device=device, dtype=dtype)[None, :]
+
+    # Flatten spatial dimensions and expand for batch
+    img_ids = img_ids.reshape(h_packed * w_packed, 3)
+    img_ids = img_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+    return img_ids
+
+
+def generate_text_position_ids(
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> Tensor:
+    """
+    Generate position IDs for text tokens.
+
+    For Flux, text position IDs are zeros (no explicit positional encoding).
+
+    Args:
+        batch_size: Batch size
+        seq_len: Text sequence length
+        device: Target device
+        dtype: Target dtype
+
+    Returns:
+        Position IDs tensor of shape (B, seq_len, 3) filled with zeros
+
+    Example:
+        >>> ids = generate_text_position_ids(2, 512, torch.device('cpu'))
+        >>> ids.shape
+        torch.Size([2, 512, 3])
+        >>> ids.sum().item()
+        0.0
+    """
+    return torch.zeros(batch_size, seq_len, 3, device=device, dtype=dtype)
+
+
+__all__ = [
+    "pack_latents",
+    "unpack_latents",
+    "generate_image_position_ids",
+    "generate_text_position_ids",
+]
diff --git a/primus/backends/megatron/core/models/hybrid/gated_delta_net.py b/primus/backends/megatron/core/models/hybrid/gated_delta_net.py
new file mode 100644
index 000000000..3525c3f77
--- /dev/null
+++ b/primus/backends/megatron/core/models/hybrid/gated_delta_net.py
@@ -0,0 +1,728 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh.
+
+# Some of this code was adopted from https://github.com/huggingface/transformers
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import math
+from dataclasses import dataclass, replace
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory
+from megatron.core.fp8_utils import get_fp8_align_size
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel import get_cuda_rng_tracker
+from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.utils import (  # ensure_metadata_has_dp_cp_group,
+    make_sharded_tensors_for_checkpoint,
+    sharded_state_dict_default,
+)
+from megatron.core.utils import (
+    deprecate_inference_params,
+    nvtx_range_pop,
+    nvtx_range_push,
+)
+from torch import Tensor
+
+# TODO: Implement GatedDeltaNetContextParallel
+# from .gated_delta_net_context_parallel import GatedDeltaNetContextParallel
+
+try:
+    from fla.modules.l2norm import l2norm
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+
+    HAVE_FLA = True
+except ImportError:
+    chunk_gated_delta_rule = None
+
+    HAVE_FLA = False
+
+from megatron.training import get_args as _get_args
+
+try:
+    from causal_conv1d import causal_conv1d_fn
+except ImportError:
+    causal_conv1d_fn = None
+    causal_conv1d_update = None
+
+try:
+    from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d
+except ImportError:
+    _fla_causal_conv1d = None
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GatedDeltaNetSubmodules:
+    """
+    Contains the module specs for the input linear, output norm, and output linear layers.
+    """
+
+    in_proj: Union[ModuleSpec, type] = IdentityOp
+    out_norm: Union[ModuleSpec, type] = IdentityOp
+    out_proj: Union[ModuleSpec, type] = IdentityOp
+
+
+class GatedDeltaNet(MegatronModule):
+    """Gated Delta Net (GDN) layer class
+
+    GDN layer takes input with size [s, b, h]
+    and returns output of the same size.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: GatedDeltaNetSubmodules,
+        layer_number: int = None,
+        bias: bool = False,
+        conv_bias: bool = False,
+        conv_init: Optional[float] = None,
+        use_qk_l2norm: bool = True,
+        A_init_range: Tuple[float, float] = (1, 16),
+        pg_collection: ProcessGroupCollection = None,
+    ):
+        """
+        Args:
+            config: The config of the model.
+            submodules: Contains the module specs for the input and output linear layers.
+            layer_number: The layer number of this GDN layer.
+            bias: Whether to use bias in the linear layers.
+            conv_bias: Whether to use bias in the causal convolution.
+            conv_init: The initialization range for the causal convolution weights.
+            use_qk_l2norm: Whether to use L2 normalization in the kernel of the gated delta rule.
+            A_init_range: The initialization range for the attention weights.
+            pg_collection: The required process groups to use for tensor model parallel and context
+                parallel.
+        """
+
+        if not HAVE_FLA:
+            raise ImportError(
+                "FLA is not installed. Please install it with `pip install flash-linear-attention`."
+            )
+
+        super().__init__(config)
+
+        # Attributes from arguments
+        self.layer_number = layer_number
+        self.bias = bias
+        self.conv_bias = conv_bias
+        self.conv_init = conv_init
+        assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0]
+        self.A_init_range = A_init_range
+        self.use_qk_l2norm = use_qk_l2norm
+        assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet"
+        self.pg_collection = pg_collection
+        self.tp_size = self.pg_collection.tp.size()
+        self.sp_size = self.tp_size if config.sequence_parallel else 1
+
+        # Attributes from config
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.act_fn = config.activation_func
+        self.activation = self.act_fn.__name__
+        self.use_short_conv = getattr(config, "use_short_conv", True)
+        self.conv_kernel_dim = config.linear_conv_kernel_dim
+        self.key_head_dim = config.linear_key_head_dim
+        self.value_head_dim = config.linear_value_head_dim
+        self.num_key_heads = config.linear_num_key_heads
+        self.num_value_heads = config.linear_num_value_heads
+        self.qk_dim = self.key_head_dim * self.num_key_heads
+        self.v_dim = self.value_head_dim * self.num_value_heads
+
+        # Input projection (hidden_states -> q, k, v, gate, beta, alpha)
+        # TODO: for now, output gate is forced for GDN.
+        # We may remove this restriction in the future.
+        self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.num_value_heads * 2
+        if self.config.fp8:
+            fp8_align_size = get_fp8_align_size(self.config.fp8_recipe)
+            assert self.in_proj_dim % fp8_align_size == 0, (
+                "For FP8, the innermost dimension of the GDN layer "
+                "input projection output tensor must be a multiple of 16."
+            )
+        self.in_proj = build_module(
+            submodules.in_proj,
+            self.hidden_size,
+            self.in_proj_dim,
+            config=self.config,
+            init_method=self.config.init_method,
+            gather_output=False,
+            bias=bias,
+            skip_bias_add=False,
+            is_expert=False,
+            tp_comm_buffer_name="fc1",
+            tp_group=self.pg_collection.tp,
+        )
+
+        # Conv1d for QKV (only when use_short_conv is enabled)
+        self.conv_dim = self.qk_dim * 2 + self.v_dim
+        self.conv_dim_local_tp = self.conv_dim // self.tp_size
+
+        if self.use_short_conv:
+            self.conv1d = nn.Conv1d(
+                in_channels=self.conv_dim_local_tp,
+                out_channels=self.conv_dim_local_tp,
+                bias=conv_bias,
+                kernel_size=self.conv_kernel_dim,
+                groups=self.conv_dim_local_tp,
+                padding=self.conv_kernel_dim - 1,
+                device=torch.cuda.current_device(),
+                dtype=config.params_dtype,
+            )
+            setattr(self.conv1d.weight, "tensor_model_parallel", True)
+            if conv_bias:
+                setattr(self.conv1d.bias, "tensor_model_parallel", True)
+
+        # Time step projection (discretization)
+        self.num_v_heads_local_tp = self.num_value_heads // self.tp_size
+        # dt_bias parameter (fp32 to avoid bf16 quantisation noise in the gate)
+        self.dt_bias = nn.Parameter(
+            torch.empty(
+                self.num_v_heads_local_tp,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+        )
+        setattr(self.dt_bias, "tensor_model_parallel", True)
+        # A_log parameter (fp32: log(A) is fed into exp() inside the kernel,
+        # and bf16 uniform_(0,16) can produce exact 0 → log(0) = -inf → NaN)
+        self.A_log = nn.Parameter(
+            torch.empty(
+                self.num_v_heads_local_tp,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+        )
+        setattr(self.A_log, "tensor_model_parallel", True)
+
+        # Output layernorm before projection
+        self._use_fla_fused_gated_norm = False
+        _use_fla_gated_norm = getattr(_get_args(), "use_fla_fused_gated_norm", False)
+        if _use_fla_gated_norm and config.normalization == "RMSNorm":
+            try:
+                from fla.modules.fused_norm_gate import FusedRMSNormGated
+
+                self.out_norm = FusedRMSNormGated(
+                    self.value_head_dim,
+                    eps=self.config.layernorm_epsilon,
+                )
+                self._use_fla_fused_gated_norm = True
+            except ImportError:
+                pass
+        if not self._use_fla_fused_gated_norm:
+            self.out_norm = build_module(
+                submodules.out_norm,
+                config=self.config,
+                hidden_size=self.value_head_dim,
+                eps=self.config.layernorm_epsilon,
+            )
+
+        self.out_proj = build_module(
+            submodules.out_proj,
+            self.v_dim,
+            self.hidden_size,
+            config=self.config,
+            init_method=self.config.output_layer_init_method,
+            bias=bias,
+            input_is_parallel=True,
+            skip_bias_add=True,
+            is_expert=False,
+            tp_comm_buffer_name="fc2",
+            tp_group=self.pg_collection.tp,
+        )
+
+        # TODO: support CP
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reset the parameters.
+
+        Matches FLA's GatedDeltaNet initialization exactly:
+        - A_log: log(uniform(0, 16))
+        - dt_bias: inverse_softplus(uniform(dt_min, dt_max))
+        """
+        if self.config.perform_initialization:
+            with get_cuda_rng_tracker().fork():
+                # conv1d.weight
+                if self.use_short_conv and self.conv_init is not None:
+                    nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+                # dt_bias: inverse softplus of log-uniform in [dt_min, dt_max]
+                dt_min, dt_max = 0.001, 0.1
+                dt = torch.exp(
+                    torch.rand(
+                        self.num_v_heads_local_tp,
+                        dtype=torch.float32,
+                        device=torch.cuda.current_device(),
+                    )
+                    * (math.log(dt_max) - math.log(dt_min))
+                    + math.log(dt_min)
+                ).clamp(min=1e-4)
+                inv_dt = dt + torch.log(-torch.expm1(-dt))
+                self.dt_bias.data.copy_(inv_dt)
+                # A_log (compute in fp32 to avoid log(0) = -inf from bf16 zeros)
+                A = torch.empty(
+                    self.num_v_heads_local_tp,
+                    dtype=torch.float32,
+                    device=torch.cuda.current_device(),
+                ).uniform_(*self.A_init_range)
+                self.A_log.data.copy_(torch.log(A))
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Tensor,
+        key_value_states: Optional[Tensor] = None,
+        inference_context: Optional[BaseInferenceContext] = None,
+        attention_bias: Optional[Tensor] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+        sequence_len_offset: Optional[int] = None,
+        *,
+        inference_params: Optional[BaseInferenceContext] = None,
+        **kwargs,
+    ):
+        """
+        Perform a forward pass through the GDN module.
+
+        Args:
+            hidden_states (Tensor): Hidden states.
+            attention_mask (Tensor): Attention mask.
+            key_value_states (Optional[Tensor]): Key/value states (for cross attention).
+            inference_context (Optional[BaseInferenceContext]): Inference context that manages
+                KV cache.
+            attention_bias (Optional[Tensor]): Attention bias.
+            packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format.
+            sequence_len_offset (Optional[int]): Sequence length offset used for
+                inference CUDA graphs.
+
+        Return:
+            (Tuple[Tensor, Tensor]) GDN output and bias.
+
+        """
+        # TODO: Deal with attention_mask
+
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        seq_len, batch, _ = hidden_states.shape
+        seq_len = seq_len * self.sp_size
+
+        if inference_context is not None:
+            assert (
+                inference_context.is_static_batching()
+            ), "GDN does not currently support dynamic inference batching."
+            assert not self.config.sequence_parallel
+            # TODO: support inference
+            raise NotImplementedError("GDN does not support inference for now.")
+
+        if packed_seq_params is not None:
+            # TODO: support packed sequence
+            raise NotImplementedError("GDN does not support packed sequence for now.")
+
+        # Input projection
+        nvtx_range_push(suffix="in_proj")
+        qkvzba, _ = self.in_proj(hidden_states)
+        nvtx_range_pop(suffix="in_proj")
+
+        # Transpose: s b x --> b s x
+        # From sbhd to bshd format
+        qkvzba = qkvzba.transpose(0, 1)
+
+        # Split, reorder, and reshape the tensor into q, k, v, gate, beta, alpha
+        qkv, gate, beta, alpha = torch.split(
+            qkvzba,
+            [
+                (self.qk_dim * 2 + self.v_dim) // self.tp_size,
+                self.v_dim // self.tp_size,
+                self.num_value_heads // self.tp_size,
+                self.num_value_heads // self.tp_size,
+            ],
+            dim=-1,
+        )
+        gate = gate.reshape(batch, seq_len, -1, self.value_head_dim)
+        beta = beta.reshape(batch, seq_len, -1)
+        alpha = alpha.reshape(batch, seq_len, -1)
+
+        # Convolution on qkv (or SiLU-only when use_short_conv=False)
+        nvtx_range_push(suffix="conv1d")
+        if self.use_short_conv:
+            _use_fla_conv = getattr(_get_args(), "use_fla_short_conv", False)
+            if _use_fla_conv and _fla_causal_conv1d is not None and not self.config.deterministic_mode:
+                # FLA's Triton causal_conv1d expects [B, T, D] (no transpose needed)
+                # and matches FLA reference run bit-for-bit.
+                assert self.activation in ["silu", "swish"]
+                qkv, _ = _fla_causal_conv1d(
+                    x=qkv,
+                    weight=self.conv1d.weight.squeeze(1),  # d, 1, w -> d, w
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    backend="triton",
+                )
+            else:
+                qkv = qkv.transpose(1, 2).contiguous()  # b, s, d -> b, d, s
+                if (causal_conv1d_fn is None) or self.config.deterministic_mode:
+                    qkv = self.act_fn(self.conv1d(qkv)[..., :seq_len])
+                else:
+                    assert self.activation in ["silu", "swish"]
+                    qkv = causal_conv1d_fn(
+                        x=qkv,
+                        weight=self.conv1d.weight.squeeze(1),  # d, 1, w -> d, w
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                    )
+                qkv = qkv.transpose(1, 2)  # b, d, s -> b, s, d
+        else:
+            qkv = self.act_fn(qkv)
+        nvtx_range_pop(suffix="conv1d")
+        # Split qkv into query, key, and value (qkv is already b, s, d)
+        query, key, value = torch.split(
+            qkv,
+            [self.qk_dim // self.tp_size, self.qk_dim // self.tp_size, self.v_dim // self.tp_size],
+            dim=-1,
+        )
+        query = query.reshape(batch, seq_len, -1, self.key_head_dim)
+        key = key.reshape(batch, seq_len, -1, self.key_head_dim)
+        value = value.reshape(batch, seq_len, -1, self.value_head_dim)
+        # GVA head expansion: pre-expand Q/K via repeat_interleave so the
+        # GDN kernel sees one K-head per V-head.
+        if self.num_value_heads // self.num_key_heads > 1:
+            query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2)
+            key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2)
+
+        # Make contiguous
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+        gate = gate.contiguous()
+        beta = beta.contiguous()
+        alpha = alpha.contiguous()
+
+        nvtx_range_push(suffix="g_and_beta")
+        beta = beta.sigmoid()
+        nvtx_range_pop(suffix="g_and_beta")
+
+        nvtx_range_push(suffix="gated_delta_rule")
+        # Pre-compute gate in fp32 to avoid NaN from FLA's in-kernel gate
+        # path on ROCm (use_gate_in_kernel=True triggers Triton NaN bugs).
+        g = -self.A_log.float().exp() * F.softplus(alpha.float() + self.dt_bias)
+        if self.config.deterministic_mode:
+            if self.use_qk_l2norm:
+                query = l2norm(query)
+                key = l2norm(key)
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g,
+                beta=beta,
+                initial_state=None,
+                output_final_state=False,
+                use_qk_l2norm_in_kernel=False,
+            )
+        else:
+            core_attn_out, last_recurrent_state = chunk_gated_delta_rule(
+                query,
+                key,
+                value,
+                g=g.to(query.dtype),
+                beta=beta,
+                initial_state=None,
+                output_final_state=False,
+                use_qk_l2norm_in_kernel=self.use_qk_l2norm,
+                use_gate_in_kernel=False,
+            )
+        nvtx_range_pop(suffix="gated_delta_rule")
+
+        # RMSNorm
+        nvtx_range_push(suffix="gated_norm")
+        norm_out = self._apply_gated_norm(core_attn_out, gate)
+        nvtx_range_pop(suffix="gated_norm")
+
+        # Transpose: b s x --> s b x
+        # From bshd back to sbhd format
+        norm_out = norm_out.reshape(batch, seq_len, -1)
+        norm_out = norm_out.transpose(0, 1).contiguous()
+
+        # Output projection
+        nvtx_range_push(suffix="out_proj")
+        out, out_bias = self.out_proj(norm_out)
+        nvtx_range_pop(suffix="out_proj")
+
+        return out, out_bias
+
+    def _apply_gated_norm(self, x, gate):
+        x_dtype = x.dtype
+        x = x.reshape(-1, x.shape[-1])
+        gate = gate.reshape(-1, gate.shape[-1])
+        if self._use_fla_fused_gated_norm:
+            y = self.out_norm(x, gate)
+        else:
+            y = self.out_norm(x)
+            y = y * self.act_fn(gate.float())
+            y = y.to(x_dtype)
+        return y
+
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None):
+        """Provide a sharded state dictionary for distributed checkpointing."""
+        # Guard for cases metadata is not provided
+        # metadata = ensure_metadata_has_dp_cp_group(metadata)
+
+        sharded_state_dict = {}
+        # Parameters
+        self._save_to_state_dict(sharded_state_dict, "", keep_vars=True)
+        sharded_state_dict = make_sharded_tensors_for_checkpoint(
+            sharded_state_dict,
+            prefix,
+            tensor_parallel_layers_axis_map={
+                "A_log": 0,
+                "dt_bias": 0,
+            },  # parameters sharded across TP
+            sharded_offsets=sharded_offsets,
+            tp_group=(tp_group if tp_group is not None else self.pg_collection.tp),
+            dp_cp_group=metadata["dp_cp_group"],
+        )
+        # Submodules
+        tp_group = tp_group if tp_group is not None else self.pg_collection.tp
+        for name, module in self.named_children():
+            if name == "conv1d" and self.use_short_conv:
+                # Add TP sharding for Conv1d
+                module_sd = module.state_dict(prefix="", keep_vars=True)
+                tp_sharding_map = {f"weight": 0}
+                if self.conv_bias:
+                    tp_sharding_map[f"bias"] = 0
+                module_sharded_sd = make_sharded_tensors_for_checkpoint(
+                    module_sd,
+                    f"{prefix}{name}.",
+                    tp_sharding_map,
+                    sharded_offsets,
+                    tp_group=tp_group,
+                    dp_cp_group=metadata["dp_cp_group"],
+                )
+            else:
+                module_sharded_sd = sharded_state_dict_default(
+                    module, f"{prefix}{name}.", sharded_offsets, metadata, tp_group=tp_group
+                )
+
+            sharded_state_dict.update(module_sharded_sd)
+
+        # At this point the TP sharding is correctly defined for each tensor, but some of the
+        # tensors must be additionally split into separate parts
+        in_proj_dim_local_tp = self.in_proj_dim // self.tp_size
+        assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim_local_tp, (
+            in_proj_dim_local_tp,
+            sharded_state_dict[f"{prefix}in_proj.weight"],
+        )
+
+        sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory(
+            sharded_state_dict[f"{prefix}in_proj.weight"],
+            [
+                self.qk_dim // self.tp_size,
+                self.qk_dim // self.tp_size,
+                self.v_dim // self.tp_size,
+                self.v_dim // self.tp_size,
+                self.num_value_heads // self.tp_size,
+                self.num_value_heads // self.tp_size,
+            ],
+            ["query", "key", "value", "z", "beta", "alpha"],
+            0,
+        )
+
+        if self.use_short_conv:
+            conv_layer_name_list = ["conv1d.weight"]
+            assert sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) == self.conv_dim_local_tp, (
+                self.conv_dim_local_tp,
+                sharded_state_dict[f"{prefix}conv1d.weight"],
+            )
+            if self.conv_bias:
+                conv_layer_name_list.append("conv1d.bias")
+                assert sharded_state_dict[f"{prefix}conv1d.bias"].data.size(0) == self.conv_dim_local_tp, (
+                    self.conv_dim_local_tp,
+                    sharded_state_dict[f"{prefix}conv1d.bias"],
+                )
+            for conv_layer_name in conv_layer_name_list:
+                sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory(
+                    sharded_state_dict[f"{prefix}{conv_layer_name}"],
+                    [
+                        self.qk_dim // self.tp_size,
+                        self.qk_dim // self.tp_size,
+                        self.v_dim // self.tp_size,
+                    ],
+                    ["query", "key", "value"],
+                    0,
+                )
+
+        return sharded_state_dict
+
+    def backward_dw(self):
+        """Execute weight gradient computation for all linear layers."""
+        self._backward_in_proj()
+        self._backward_out_proj()
+
+    def _backward_in_proj(self):
+        """Computes weight gradients of input projection layer."""
+        self.in_proj.backward_dw()
+
+    def _backward_out_proj(self):
+        """Computes weight gradients of output projection layer."""
+        self.out_proj.backward_dw()
+
+
+def _split_tensor_factory(
+    orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int
+) -> ShardedTensorFactory:
+    """Builds a factory that splits a given ShardedTensor into several independent chunks."""
+    assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten)
+    orig_sh_ten_no_data = orig_sh_ten.without_data()  # remove `data` reference
+
+    if sum(split_sections) != orig_sh_ten_no_data.local_shape[split_dim]:
+        raise ValueError(
+            f"Split sections must cover the whole dimension size, "
+            f"got {split_sections=} vs dimensions size "
+            f"{orig_sh_ten_no_data.local_shape[split_dim]}"
+        )
+
+    assert not isinstance(
+        split_sections, int
+    ), "Splitting into predefined section sizes is supported (`split_sections` must be a list)"
+    assert len(split_sections) == len(split_names), (len(split_sections), len(split_names))
+
+    @torch.no_grad()
+    def sh_ten_build_fn(key: str, t: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice]):
+        factory_sh_ten = replace(
+            orig_sh_ten_no_data,
+            key=key,
+            data=t,
+            dtype=t.dtype,
+            replica_id=replica_id,
+            flattened_range=flattened_range,
+        )
+
+        chunk_sh_tens = []
+        split_start = 0
+        for split_size, split_name in zip(split_sections, split_names):
+            split_chunks = factory_sh_ten.narrow(split_dim, split_start, split_size)
+            for sh_ten in split_chunks:
+                sh_ten.key = f"{sh_ten.key}.{split_name}"
+            chunk_sh_tens.extend(split_chunks)
+            split_start += split_size
+
+        assert split_start == orig_sh_ten_no_data.local_shape[split_dim], (
+            split_start,
+            orig_sh_ten_no_data.local_shape[split_dim],
+        )
+        assert sum(sh_ten.data.numel() for sh_ten in chunk_sh_tens) == t.numel(), (
+            chunk_sh_tens,
+            t.shape,
+        )
+        return chunk_sh_tens
+
+    @torch.no_grad()
+    def sh_ten_merge_fn(sub_state_dict):
+        return torch.cat(sub_state_dict)
+
+    return ShardedTensorFactory(
+        orig_sh_ten.key, orig_sh_ten.data, sh_ten_build_fn, sh_ten_merge_fn, orig_sh_ten.replica_id
+    )
+
+
+def torch_chunk_gated_delta_rule(
+    query,
+    key,
+    value,
+    g,
+    beta,
+    chunk_size=64,
+    initial_state=None,
+    output_final_state=False,
+    use_qk_l2norm_in_kernel=False,
+):
+    # pylint: disable=line-too-long
+    """
+    Torch-native implementation of chunked gated delta rule for deterministic mode.
+    Need this because FLA is not deterministic.
+
+    Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547
+    """
+
+    initial_dtype = query.dtype
+    if use_qk_l2norm_in_kernel:
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
+    query, key, value, beta, g = [
+        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
+    ]
+
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    query = F.pad(query, (0, 0, 0, pad_size))
+    key = F.pad(key, (0, 0, 0, pad_size))
+    value = F.pad(value, (0, 0, 0, pad_size))
+    beta = F.pad(beta, (0, pad_size))
+    g = F.pad(g, (0, pad_size))
+    total_sequence_length = sequence_length + pad_size
+    scale = 1 / (query.shape[-1] ** 0.5)
+    query = query * scale
+
+    v_beta = value * beta.unsqueeze(-1)
+    k_beta = key * beta.unsqueeze(-1)
+    # reshape to chunks
+    query, key, value, k_beta, v_beta = [
+        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
+        for x in (query, key, value, k_beta, v_beta)
+    ]
+    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
+
+    # chunk decay
+    g = g.cumsum(dim=-1)
+    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        row = attn[..., i, :i].clone()
+        sub = attn[..., :i, :i].clone()
+        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    value = attn @ v_beta
+    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+    last_recurrent_state = (
+        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
+        if initial_state is None
+        else initial_state.to(value)
+    )
+    core_attn_out = torch.zeros_like(value)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
+
+    # for each chunk
+    for i in range(0, total_sequence_length // chunk_size):
+        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
+        v_new = v_i - v_prime
+        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
+        core_attn_out[:, :, i] = attn_inter + attn @ v_new
+        last_recurrent_state = (
+            last_recurrent_state * g[:, :, i, -1, None, None].exp()
+            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
+        )
+
+    if not output_final_state:
+        last_recurrent_state = None
+    core_attn_out = core_attn_out.reshape(
+        core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]
+    )
+    core_attn_out = core_attn_out[:, :, :sequence_length]
+    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+    return core_attn_out, last_recurrent_state
diff --git a/primus/backends/megatron/core/models/hybrid/gated_delta_net_layer.py b/primus/backends/megatron/core/models/hybrid/gated_delta_net_layer.py
new file mode 100644
index 000000000..cf829d543
--- /dev/null
+++ b/primus/backends/megatron/core/models/hybrid/gated_delta_net_layer.py
@@ -0,0 +1,140 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Union
+
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import deprecate_inference_params
+from torch import Tensor
+
+
+@dataclass
+class GatedDeltaNetLayerSubmodules:
+    """
+    Configuration class for specifying the submodules of a GatedDeltaNet layer.
+
+    Unlike MambaLayerSubmodules, this does not pass d_model to the mixer and
+    forwards attention_mask through to the GatedDeltaNet mixer.
+    """
+
+    norm: Union[ModuleSpec, type] = IdentityOp
+    mixer: Union[ModuleSpec, type] = IdentityOp
+    gdn_bda: Union[ModuleSpec, type] = IdentityOp
+
+    sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
+
+
+class GatedDeltaNetLayer(MegatronModule):
+    """
+    A single GatedDeltaNet layer wrapping the GatedDeltaNet mixer.
+
+    This is analogous to MambaLayer but avoids the interface mismatches:
+    - Does not pass d_model to the mixer (GatedDeltaNet reads hidden_size from config)
+    - Forwards attention_mask to the mixer
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: GatedDeltaNetLayerSubmodules,
+        layer_number: int = 1,
+        residual_in_fp32: bool = False,
+        pg_collection: ProcessGroupCollection = None,
+    ):
+        super().__init__(config)
+        assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNetLayer"
+
+        self.config = config
+        self.submodules_config = submodules
+        self.layer_number = layer_number
+        self.residual_in_fp32 = residual_in_fp32
+        self.hidden_dropout = config.hidden_dropout
+
+        self.mixer = build_module(
+            submodules.mixer,
+            self.config,
+            layer_number=layer_number,
+            pg_collection=pg_collection,
+        )
+        # Pass eps from config.layernorm_epsilon. WrappedTorchNorm defaults
+        # eps to 1e-5 if not provided; FLA's GatedDeltaNet uses 1e-6, so we
+        # must forward the YAML value explicitly or the pre-norm silently
+        # diverges from FLA by ~1.1% per layer.
+        self.norm = build_module(
+            submodules.norm,
+            self.config,
+            self.config.hidden_size,
+            eps=self.config.layernorm_epsilon,
+        )
+        self.gdn_bda = build_module(submodules.gdn_bda)
+        self.bias_dropout_add_exec_handler = torch.enable_grad
+        self._fuse_prenorm_with_next = False
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        inference_context: Optional[BaseInferenceContext] = None,
+        rotary_pos_emb: Optional[Tensor] = None,
+        *,
+        inference_params: Optional[BaseInferenceContext] = None,
+    ):
+        """
+        Forward pass through the GatedDeltaNet layer.
+
+        Args:
+            hidden_states (Tensor): Input tensor of shape [s, b, h].
+            attention_mask (Tensor, optional): Attention mask forwarded to the mixer.
+            inference_context (BaseInferenceContext, optional): Inference context.
+            rotary_pos_emb (Tensor, optional): Rotary positional embeddings (unused).
+
+        Returns:
+            Tensor: Transformed hidden states of shape [s, b, h].
+        """
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        residual = hidden_states
+
+        hidden_states = hidden_states.to(dtype=self.config.params_dtype)
+        hidden_states = self.norm(hidden_states)
+
+        mixer_out_with_bias = self.mixer(hidden_states, attention_mask, inference_context=inference_context)
+
+        if self._fuse_prenorm_with_next:
+            mixer_out = (
+                mixer_out_with_bias[0] if isinstance(mixer_out_with_bias, tuple) else mixer_out_with_bias
+            )
+            return mixer_out, residual
+
+        if self.residual_in_fp32:
+            residual = residual.to(torch.float32)
+
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.gdn_bda(training=self.training, fused=self.config.bias_dropout_fusion)(
+                mixer_out_with_bias, residual, self.hidden_dropout
+            )
+
+        if self.residual_in_fp32:
+            hidden_states = hidden_states.to(self.config.params_dtype)
+
+        return hidden_states
+
+    def sharded_state_dict(
+        self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None
+    ) -> ShardedStateDict:
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        prefixed_map = {
+            f"{prefix}{k}": f"{prefix}{v}"
+            for k, v in self.submodules_config.sharded_state_dict_keys_map.items()
+        }
+        if prefixed_map:
+            apply_prefix_mapping(sharded_state_dict, prefixed_map)
+        return sharded_state_dict
diff --git a/primus/backends/megatron/core/models/hybrid/hybrid_block.py b/primus/backends/megatron/core/models/hybrid/hybrid_block.py
index 81ac3efe1..1b7ac30cc 100644
--- a/primus/backends/megatron/core/models/hybrid/hybrid_block.py
+++ b/primus/backends/megatron/core/models/hybrid/hybrid_block.py
@@ -87,6 +87,7 @@ def __init__(
         submodules: HybridStackSubmodules,
         residual_in_fp32=False,
         pre_process: bool = True,
+        layer_type_list=None,
         hybrid_attention_ratio: float = 0.0,
         hybrid_mlp_ratio: float = 0.0,
         hybrid_override_pattern: str = None,
@@ -94,10 +95,13 @@ def __init__(
         post_process: bool = True,
         device=None,
         dtype=None,
+        pp_layer_offset: int = 0,
         pg_collection: ProcessGroupCollection = None,
         **kwargs,
     ) -> None:
         super().__init__(config=config)
+        if residual_in_fp32 is False and config.fp32_residual_connection:
+            residual_in_fp32 = True
         self.residual_in_fp32 = residual_in_fp32
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
@@ -129,20 +133,30 @@ def __init__(
         self.hybrid_mlp_ratio = hybrid_mlp_ratio
         self.hybrid_override_pattern = hybrid_override_pattern
 
-        # Customized layer allocation
-        # hybrid_mlp_ratio is not used in this hybrid stack.
-        # It is by default to be always followed by mamba or mla (i.e., mamba + MLP or MLA + MLP)
-        # By setting hybrid_attention_ratio, attention layers are by default to be distributed uniformly.
-        self.layer_type_list = self.allocate_layers(
-            self.config.num_layers,
-            self.hybrid_attention_ratio,
-        )
-
-        pp_layer_offset = 0
-        if self.pp_group.size() > 1:
-            pp_layer_offset, self.layer_type_list = self._select_layers_for_pipeline_parallel(
-                self.layer_type_list
+        # Modern Megatron `MambaModel` parses `hybrid_layer_pattern` into a
+        # concrete `layer_type_list` (list of `Symbols` like 'M', '*', '-',
+        # 'E') and a `pp_layer_offset`, then passes them to
+        # `build_module(mamba_stack_spec, ..., layer_type_list=..., pp_layer_offset=...)`.
+        # If we received a non-empty pre-computed list, USE IT.  An *empty*
+        # list (or None) means the caller didn't specify a pattern (e.g. pure
+        # GDN configs, or hybrid configs whose YAML uses `hybrid_attention_ratio`
+        # which upstream silently dropped as an arg); in that case fall back
+        # to the legacy ratio-based allocation so the old behaviour is preserved.
+        if layer_type_list:
+            self.layer_type_list = list(layer_type_list)
+        else:
+            # Legacy path: caller didn't pre-compute the list, allocate from
+            # the ratio.  hybrid_mlp_ratio is intentionally ignored here --
+            # this hybrid stack always follows attention/mamba with an MLP.
+            self.layer_type_list = self.allocate_layers(
+                self.config.num_layers,
+                self.hybrid_attention_ratio,
             )
+            pp_layer_offset = 0
+            if self.pp_group.size() > 1:
+                pp_layer_offset, self.layer_type_list = self._select_layers_for_pipeline_parallel(
+                    self.layer_type_list
+                )
 
         print(f"layer_type_list: {self.layer_type_list}")
 
@@ -180,11 +194,22 @@ def __init__(
                     assert False, "unexpected layer_type"
             self.layers.append(layer)
 
+        from megatron.training import get_args as _get_args
+
+        self._fuse_prenorm = getattr(_get_args(), "use_fla_fused_rmsnorm", False)
+        if self._fuse_prenorm:
+            from primus.backends.megatron.core.models.hybrid.gated_delta_net_layer import (
+                GatedDeltaNetLayer,
+            )
+
+            for i, (lt, layer) in enumerate(zip(self.layer_type_list, self.layers)):
+                if lt == LayerSymbols.MAMBA and isinstance(layer, GatedDeltaNetLayer):
+                    layer._fuse_prenorm_with_next = True
+
         # Required for activation recomputation
         self.num_layers_per_pipeline_rank = len(self.layers)
 
         if self.post_process and self.post_layer_norm:
-            # Final layer norm before output.
             self.final_norm = TENorm(
                 config=self.config,
                 hidden_size=self.config.hidden_size,
@@ -195,6 +220,10 @@ def allocate_layers(self, num_layers, hybrid_attention_ratio):
         layer_type_list = []
         num_attention_layers = int(num_layers // 2 * hybrid_attention_ratio)
         num_mamba_layers = num_layers // 2 - num_attention_layers
+
+        if num_attention_layers == 0:
+            return [LayerSymbols.MAMBA, LayerSymbols.MLP] * (num_layers // 2)
+
         num_mamba_per_attention_layer = num_mamba_layers // num_attention_layers
 
         if hybrid_attention_ratio <= 0.5:
@@ -254,6 +283,8 @@ def forward(
         rotary_pos_emb: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
+        packed_seq_params=None,
+        padding_mask=None,
         **kwargs,
     ):
         """
@@ -322,15 +353,27 @@ def forward(
         use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed
         outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext()
 
+        _pending_fuse = None
+
         with outer_fp8_context:
-            for layer in self.layers:
+            for i_layer, layer in enumerate(self.layers):
                 inner_fp8_context = (
                     get_fp8_context(self.config, layer.layer_number - 1)
                     if use_inner_fp8_context
                     else nullcontext()
                 )
                 with inner_fp8_context:
-                    if isinstance(layer, TransformerLayer):
+                    if isinstance(layer, TransformerLayer) and _pending_fuse is not None:
+                        mixer_out, block_residual = _pending_fuse
+                        _pending_fuse = None
+                        normed, new_residual = layer.pre_mlp_layernorm(mixer_out, block_residual, True)
+                        mlp_output_with_bias = layer.mlp(normed)
+                        bda_fn = layer.mlp_bda(
+                            training=self.training,
+                            fused=self.config.bias_dropout_fusion,
+                        )
+                        hidden_states = bda_fn(mlp_output_with_bias, new_residual, layer.hidden_dropout)
+                    elif isinstance(layer, TransformerLayer):
                         hidden_states, _ = layer(
                             hidden_states=hidden_states,
                             attention_mask=attention_mask,
@@ -338,12 +381,16 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                         )
-                    else:  # MambaLayer
-                        hidden_states = layer(
+                    else:  # MambaLayer / GatedDeltaNetLayer
+                        result = layer(
                             hidden_states=hidden_states,
                             attention_mask=attention_mask,
                             inference_context=inference_context,
                         )
+                        if isinstance(result, tuple):
+                            _pending_fuse = result
+                        else:
+                            hidden_states = result
 
                 # The attention layer (currently a simplified transformer layer)
                 # outputs a tuple of (hidden_states, context). Context is intended
diff --git a/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py b/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py
index 32616f85d..c9b3cc9bc 100644
--- a/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py
+++ b/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py
@@ -19,11 +19,29 @@
     MLASelfAttentionSubmodules,
 )
 
-# Import HybridStack from relative path
+from primus.backends.megatron.core.models.hybrid.gated_delta_net import (
+    GatedDeltaNet,
+    GatedDeltaNetSubmodules,
+)
+from primus.backends.megatron.core.models.hybrid.gated_delta_net_layer import (
+    GatedDeltaNetLayer,
+    GatedDeltaNetLayerSubmodules,
+)
 from primus.backends.megatron.core.models.hybrid.hybrid_block import (
     HybridStack,
     HybridStackSubmodules,
 )
+from primus.backends.megatron.core.models.hybrid.kimi_delta_attention import (
+    KimiDeltaAttention,
+    KimiDeltaAttentionSubmodules,
+)
+from primus.backends.megatron.core.models.hybrid.kimi_delta_attention_layer import (
+    KimiDeltaAttentionLayer,
+    KimiDeltaAttentionLayerSubmodules,
+)
+from primus.backends.megatron.core.models.hybrid.mamba_layer_adapter import (
+    Mamba2HybridLayer,
+)
 
 # Inference layers may not be available in older Megatron versions
 # They're only used in hybrid_inference_stack_spec, not the training spec
@@ -40,14 +58,65 @@
     InferenceRowParallelLinear = TERowParallelLinear
     HAS_INFERENCE_LAYERS = False
 
+from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.torch_norm import WrappedTorchNorm
 from megatron.core.transformer.transformer_layer import (
     TransformerLayer,
     TransformerLayerSubmodules,
 )
 
+from primus.backends.megatron.core.transformer.fla_flash_attention import (
+    FLAFlashAttention,
+)
+from primus.backends.megatron.core.transformer.fla_flash_attention import (
+    is_enabled as _fla_mla_attn_enabled,
+)
+
+# Route MLA's `core_attention` through a direct `flash_attn_func` call
+# instead of TransformerEngine's `TEDotProductAttention` whenever the
+# installed flash-attn version is newer than TE supports (>2.8.1) — that's
+# the case where TE silently drops to its Composable-Kernel backend and
+# loses ~30 ms/MLA-block on MI300X.  Auto-enabled by default; override
+# with `fla_mla_attn: "0"` in YAML (or `PRIMUS_FLA_MLA_ATTN=0`) to force the TE path.
+_MLA_CORE_ATTENTION = FLAFlashAttention if _fla_mla_attn_enabled() else TEDotProductAttention
+
+
+# Module-load diagnostic: drop a per-rank marker file so we can verify
+# unambiguously which copy of this spec was actually imported by training.
+# This sidesteps Megatron's stdout filtering and any later monkey-patching.
+def _record_spec_import_marker() -> None:
+    import os
+    import sys
+    import time
+
+    try:
+        rank = int(os.environ.get("RANK", "-1"))
+        marker = f"/tmp/primus_hybrid_spec_imported.rank{rank}.txt"
+        with open(marker, "w") as fh:
+            fh.write(f"file        = {__file__}\n")
+            fh.write(f"_MLA_CORE_ATTENTION = {_MLA_CORE_ATTENTION!r}\n")
+            try:
+                from megatron.training import get_args as _ga
+
+                _mla_val = getattr(_ga(), "fla_mla_attn", "")
+            except Exception:
+                _mla_val = "(args unavailable)"
+            fh.write(f"args.fla_mla_attn   = {_mla_val!r}\n")
+            fh.write(f"is_enabled()        = {_fla_mla_attn_enabled()}\n")
+            fh.write(f"pid                 = {os.getpid()}\n")
+            fh.write(f"ts                  = {time.time()}\n")
+            fh.write("sys.path[:6]:\n")
+            for p in sys.path[:6]:
+                fh.write(f"  {p}\n")
+    except Exception:
+        pass
+
+
+_record_spec_import_marker()
+
 moe = get_moe_module_spec(
     use_te=True,
     num_experts=8,  # Can be any positive integer (must not be None).
@@ -86,10 +155,73 @@
                         linear_q_up_proj=TELayerNormColumnParallelLinear,
                         linear_kv_down_proj=TELinear,
                         linear_kv_up_proj=TELayerNormColumnParallelLinear,
-                        core_attention=TEDotProductAttention,
+                        core_attention=_MLA_CORE_ATTENTION,
+                        linear_proj=TERowParallelLinear,
+                        # FLA's MLA wraps every LoRA projection in
+                        # `nn.Sequential(Linear → RMSNorm(fp32) → Linear)`; with
+                        # IdentityOp we skip the intermediate norm and the model
+                        # plateaus ~0.12 above FLA's loss curve from iter 100
+                        # onwards (iter-1 still matches bit-perfect).  TENorm
+                        # matches FLA's `RMSNorm(dtype=fp32)` exactly.
+                        q_layernorm=TENorm,
+                        kv_layernorm=TENorm,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
+        mlp_layer=ModuleSpec(
+            module=MLPLayer,
+            submodules=TransformerLayerSubmodules(
+                mlp=ModuleSpec(
+                    module=MLP,
+                    submodules=MLPSubmodules(
+                        linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+                    ),
+                ),
+                mlp_bda=get_bias_dropout_add,
+            ),
+        ),
+    ),
+)
+
+
+gdn_hybrid_stack_spec = ModuleSpec(
+    module=HybridStack,
+    submodules=HybridStackSubmodules(
+        mamba_layer=ModuleSpec(
+            module=GatedDeltaNetLayer,
+            submodules=GatedDeltaNetLayerSubmodules(
+                mixer=ModuleSpec(
+                    module=GatedDeltaNet,
+                    submodules=GatedDeltaNetSubmodules(
+                        in_proj=TELayerNormColumnParallelLinear, out_norm=TENorm, out_proj=TERowParallelLinear
+                    ),
+                ),
+                gdn_bda=get_bias_dropout_add,
+            ),
+        ),
+        attention_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=TEColumnParallelLinear,
+                        linear_q_down_proj=TELinear,
+                        linear_q_up_proj=TELayerNormColumnParallelLinear,
+                        linear_kv_down_proj=TELinear,
+                        linear_kv_up_proj=TELayerNormColumnParallelLinear,
+                        core_attention=_MLA_CORE_ATTENTION,
                         linear_proj=TERowParallelLinear,
-                        q_layernorm=IdentityOp,
-                        kv_layernorm=IdentityOp,
+                        # FLA's MLA applies `RMSNorm(dtype=fp32)` between every
+                        # LoRA down/up projection (see fla/layers/mla.py).
+                        # IdentityOp skips it and the loss plateaus ~0.12 above
+                        # FLA from iter 100 onwards.  TENorm = fp32 RMSNorm.
+                        q_layernorm=TENorm,
+                        kv_layernorm=TENorm,
                     ),
                 ),
                 self_attn_bda=get_bias_dropout_add,
@@ -107,12 +239,275 @@
                 mlp_bda=get_bias_dropout_add,
             ),
         ),
-        moe_layer=ModuleSpec(
-            # TODO (rwaleffe): change this to be an "MoELayer" to work with CudaGraphs?
+        moe_layer=moe,
+    ),
+)
+
+
+gdn_hybrid_stack_spec_no_te = ModuleSpec(
+    module=HybridStack,
+    submodules=HybridStackSubmodules(
+        mamba_layer=ModuleSpec(
+            module=GatedDeltaNetLayer,
+            submodules=GatedDeltaNetLayerSubmodules(
+                norm=WrappedTorchNorm,
+                mixer=ModuleSpec(
+                    module=GatedDeltaNet,
+                    submodules=GatedDeltaNetSubmodules(
+                        in_proj=ColumnParallelLinear,
+                        out_norm=WrappedTorchNorm,
+                        out_proj=RowParallelLinear,
+                    ),
+                ),
+                gdn_bda=get_bias_dropout_add,
+            ),
+        ),
+        attention_layer=ModuleSpec(
             module=TransformerLayer,
             submodules=TransformerLayerSubmodules(
-                pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
+                input_layernorm=WrappedTorchNorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=ColumnParallelLinear,
+                        linear_q_down_proj=ColumnParallelLinear,
+                        linear_q_up_proj=ColumnParallelLinear,
+                        linear_kv_down_proj=ColumnParallelLinear,
+                        linear_kv_up_proj=ColumnParallelLinear,
+                        core_attention=_MLA_CORE_ATTENTION,
+                        linear_proj=RowParallelLinear,
+                        # FLA's MLA wraps every LoRA projection in
+                        # `nn.Sequential(Linear → RMSNorm(fp32) → Linear)`
+                        # (fla/layers/mla.py lines 99-112).  Skipping this norm
+                        # (IdentityOp) leaves the loss curve plateaued ~0.12
+                        # above FLA from iter 100 onward; iter-1 still matches
+                        # bit-perfect because both models start from identical
+                        # init and the missing norm only kicks in after the
+                        # LoRA weights diverge from their init.  Using
+                        # WrappedTorchNorm here makes the per-LoRA norm pick up
+                        # FLA's Triton RMSNorm when `PRIMUS_FLA_NORM=1`.
+                        q_layernorm=WrappedTorchNorm,
+                        kv_layernorm=WrappedTorchNorm,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
+        mlp_layer=ModuleSpec(
+            module=MLPLayer,
+            submodules=TransformerLayerSubmodules(
+                pre_mlp_layernorm=WrappedTorchNorm,
+                mlp=ModuleSpec(
+                    module=MLP,
+                    submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
+                ),
+                mlp_bda=get_bias_dropout_add,
+            ),
+        ),
+        moe_layer=moe,
+    ),
+)
+
+
+mamba_hybrid_stack_spec_no_te = ModuleSpec(
+    module=HybridStack,
+    submodules=HybridStackSubmodules(
+        # Mamba2 mixer wrapped in our `Mamba2HybridLayer` adapter (subclass of
+        # upstream MambaLayer that accepts `residual_in_fp32` so it slots into
+        # Primus's `HybridStack` builder — see mamba_layer_adapter.py).
+        # No TE column-parallel-LayerNorm fusion here — matches the no-TE GDN
+        # variant so the same `_MLA_CORE_ATTENTION` (FLA flash-attn or TE) path
+        # is used end-to-end without TE-norm folding.
+        mamba_layer=ModuleSpec(
+            module=Mamba2HybridLayer,
+            submodules=MambaLayerSubmodules(
+                mixer=ModuleSpec(
+                    module=MambaMixer,
+                    params={
+                        "expand": 2,
+                        "d_conv": 4,
+                    },
+                    submodules=MambaMixerSubmodules(
+                        in_proj=ColumnParallelLinear,
+                        out_proj=RowParallelLinear,
+                    ),
+                ),
+                mamba_bda=get_bias_dropout_add,
+            ),
+        ),
+        attention_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=WrappedTorchNorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=ColumnParallelLinear,
+                        linear_q_down_proj=ColumnParallelLinear,
+                        linear_q_up_proj=ColumnParallelLinear,
+                        linear_kv_down_proj=ColumnParallelLinear,
+                        linear_kv_up_proj=ColumnParallelLinear,
+                        core_attention=_MLA_CORE_ATTENTION,
+                        linear_proj=RowParallelLinear,
+                        # Same MLA LoRA-norm fix as gdn_hybrid_stack_spec_no_te:
+                        # WrappedTorchNorm = RMSNorm(fp32) between every LoRA
+                        # down/up projection (matches FLA's
+                        # `nn.Sequential(Linear → RMSNorm(fp32) → Linear)`).
+                        q_layernorm=WrappedTorchNorm,
+                        kv_layernorm=WrappedTorchNorm,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
+        mlp_layer=ModuleSpec(
+            module=MLPLayer,
+            submodules=TransformerLayerSubmodules(
+                pre_mlp_layernorm=WrappedTorchNorm,
+                mlp=ModuleSpec(
+                    module=MLP,
+                    submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
+                ),
+                mlp_bda=get_bias_dropout_add,
+            ),
+        ),
+        moe_layer=moe,
+    ),
+)
+
+
+kda_hybrid_stack_spec = ModuleSpec(
+    module=HybridStack,
+    submodules=HybridStackSubmodules(
+        mamba_layer=ModuleSpec(
+            module=KimiDeltaAttentionLayer,
+            submodules=KimiDeltaAttentionLayerSubmodules(
+                mixer=ModuleSpec(
+                    module=KimiDeltaAttention,
+                    submodules=KimiDeltaAttentionSubmodules(
+                        in_proj=TELayerNormColumnParallelLinear,
+                        gate_norm=TENorm,
+                        out_norm=TENorm,
+                        out_proj=TERowParallelLinear,
+                    ),
+                ),
+                kda_bda=get_bias_dropout_add,
+            ),
+        ),
+        attention_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=TEColumnParallelLinear,
+                        linear_q_down_proj=TELinear,
+                        linear_q_up_proj=TELayerNormColumnParallelLinear,
+                        linear_kv_down_proj=TELinear,
+                        linear_kv_up_proj=TELayerNormColumnParallelLinear,
+                        core_attention=_MLA_CORE_ATTENTION,
+                        linear_proj=TERowParallelLinear,
+                        # FLA's MLA applies `RMSNorm(dtype=fp32)` between LoRA
+                        # down/up projections (fla/layers/mla.py).  IdentityOp
+                        # here breaks training-dynamics parity even when the
+                        # init checkpoint matches bit-perfect at iter 1.
+                        q_layernorm=TENorm,
+                        kv_layernorm=TENorm,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
+        mlp_layer=ModuleSpec(
+            module=MLPLayer,
+            submodules=TransformerLayerSubmodules(
+                mlp=ModuleSpec(
+                    module=MLP,
+                    submodules=MLPSubmodules(
+                        linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+                    ),
+                ),
+                mlp_bda=get_bias_dropout_add,
+            ),
+        ),
+        moe_layer=moe,
+    ),
+)
+
+
+# No-TE KDA spec — mirrors gdn_hybrid_stack_spec_no_te. Replaces every TE
+# wrapper with plain WrappedTorchNorm / ColumnParallelLinear / RowParallelLinear.
+# On ROCm this removes TE's per-call dispatch indirection + dtype recasts that
+# account for the bulk of Megatron's per-iter overhead vs FLA's HF-Trainer loop.
+#
+# Architectural match to fla/models/kda/modeling_kda.py KDABlock:
+#   - Single pre-norm at the layer (KimiDeltaAttentionLayer.norm = WrappedTorchNorm)
+#   - Mixer in_proj is plain ColumnParallelLinear (no fused norm-and-project)
+#   - Mixer gate_norm = IdentityOp (FLA has no re-norm for the gate path; the
+#     pre-norm-once-and-reuse pattern saves 1 norm launch per layer)
+#   - Mixer out_norm stays as WrappedTorchNorm (becomes FusedRMSNormGated when
+#     PRIMUS_FLA_NORM=1 + use_fla_triton_kda=true, see KimiDeltaAttention.__init__)
+kda_hybrid_stack_spec_no_te = ModuleSpec(
+    module=HybridStack,
+    submodules=HybridStackSubmodules(
+        mamba_layer=ModuleSpec(
+            module=KimiDeltaAttentionLayer,
+            submodules=KimiDeltaAttentionLayerSubmodules(
+                norm=WrappedTorchNorm,
+                mixer=ModuleSpec(
+                    module=KimiDeltaAttention,
+                    submodules=KimiDeltaAttentionSubmodules(
+                        in_proj=ColumnParallelLinear,
+                        gate_norm=IdentityOp,
+                        out_norm=WrappedTorchNorm,
+                        out_proj=RowParallelLinear,
+                    ),
+                ),
+                kda_bda=get_bias_dropout_add,
+            ),
+        ),
+        attention_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=WrappedTorchNorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=ColumnParallelLinear,
+                        linear_q_down_proj=ColumnParallelLinear,
+                        linear_q_up_proj=ColumnParallelLinear,
+                        linear_kv_down_proj=ColumnParallelLinear,
+                        linear_kv_up_proj=ColumnParallelLinear,
+                        core_attention=_MLA_CORE_ATTENTION,
+                        linear_proj=RowParallelLinear,
+                        # FLA wraps every LoRA proj in
+                        # `nn.Sequential(Linear → RMSNorm(fp32) → Linear)`
+                        # (fla/layers/mla.py).  WrappedTorchNorm gives us the
+                        # equivalent and, under PRIMUS_FLA_NORM=1, swaps to
+                        # FLA's Triton `RMSNorm` for bit-exact parity.
+                        q_layernorm=WrappedTorchNorm,
+                        kv_layernorm=WrappedTorchNorm,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
+        mlp_layer=ModuleSpec(
+            module=MLPLayer,
+            submodules=TransformerLayerSubmodules(
+                pre_mlp_layernorm=WrappedTorchNorm,
+                mlp=ModuleSpec(
+                    module=MLP,
+                    submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
+                ),
+                mlp_bda=get_bias_dropout_add,
             ),
         ),
+        moe_layer=moe,
     ),
 )
diff --git a/primus/backends/megatron/core/models/hybrid/kimi_delta_attention.py b/primus/backends/megatron/core/models/hybrid/kimi_delta_attention.py
new file mode 100644
index 000000000..d7b99955b
--- /dev/null
+++ b/primus/backends/megatron/core/models/hybrid/kimi_delta_attention.py
@@ -0,0 +1,919 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Adapted from Kimi-Linear (Moonshot AI) delta attention architecture.
+# Reference: https://huggingface.co/moonshotai/Kimi-Linear-48B-A3B-Instruct
+
+import logging
+import math
+from dataclasses import dataclass, replace
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel import get_cuda_rng_tracker
+from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.utils import (
+    make_sharded_tensors_for_checkpoint,
+    sharded_state_dict_default,
+)
+from megatron.core.utils import (
+    deprecate_inference_params,
+    nvtx_range_pop,
+    nvtx_range_push,
+)
+from torch import Tensor
+
+try:
+    from fla.ops.kda import chunk_kda
+    from fla.ops.kda.gate import fused_kda_gate
+
+    HAVE_FLA_KDA = True
+except ImportError:
+    chunk_kda = None
+    fused_kda_gate = None
+    HAVE_FLA_KDA = False
+
+try:
+    from fla.modules import FusedRMSNormGated
+
+    HAVE_FUSED_RMS_NORM_GATED = True
+except ImportError:
+    FusedRMSNormGated = None
+    HAVE_FUSED_RMS_NORM_GATED = False
+
+try:
+    from causal_conv1d import causal_conv1d_fn
+except ImportError:
+    causal_conv1d_fn = None
+
+from megatron.training import get_args as _get_args
+
+try:
+    from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d
+except ImportError:
+    _fla_causal_conv1d = None
+
+from torch.utils.checkpoint import checkpoint as _grad_checkpoint
+
+logger = logging.getLogger(__name__)
+
+
+def torch_kda_gate(g, A_log, dt_bias=None):
+    """Pure-PyTorch KDA gate (used for backward pass on ROCm)."""
+    H = g.shape[-2]
+    g = g.float()
+    if dt_bias is not None:
+        g = g + dt_bias.view(H, -1)
+    return -A_log.view(H, 1).float().exp() * F.softplus(g)
+
+
+def _torch_chunk_kda_fwd(q, k, v, g, beta, scale=None, chunk_size=64, use_qk_l2norm_in_kernel=False):
+    """Pure-PyTorch chunked KDA forward -- used for backward recomputation
+    on ROCm where FLA Triton backward kernels hang."""
+    initial_dtype = q.dtype
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    BT = chunk_size
+
+    if scale is None:
+        scale = K**-0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6)
+        k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6)
+
+    pad_size = (BT - T % BT) % BT
+    if pad_size > 0:
+        q = F.pad(q, (0, 0, 0, 0, 0, pad_size))
+        k = F.pad(k, (0, 0, 0, 0, 0, pad_size))
+        v = F.pad(v, (0, 0, 0, 0, 0, pad_size))
+        g = F.pad(g, (0, 0, 0, 0, 0, pad_size))
+        beta = F.pad(beta, (0, 0, 0, pad_size))
+
+    total_T = T + pad_size
+    NT = total_T // BT
+
+    q, k, v, g, beta = [x.transpose(1, 2).contiguous().float() for x in (q, k, v, g, beta)]
+    q = q * scale
+
+    q = q.reshape(B, H, NT, BT, K)
+    k = k.reshape(B, H, NT, BT, K)
+    v = v.reshape(B, H, NT, BT, V)
+    g = g.reshape(B, H, NT, BT, K)
+    beta = beta.reshape(B, H, NT, BT)
+
+    g = g.cumsum(dim=-2)
+    k_eg = k * g.exp()
+
+    A = torch.zeros(B, H, NT, BT, BT, dtype=torch.float, device=q.device)
+    for j in range(BT):
+        k_j = k[..., j, :]
+        g_j = g[..., j : j + 1, :]
+        decay = (g - g_j).clamp(max=0).exp()
+        A[..., j] = (k * decay * k_j.unsqueeze(-2)).sum(-1)
+
+    A = A * beta.unsqueeze(-1)
+    mask_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    A = -A.masked_fill(mask_upper, 0)
+
+    for i in range(1, BT):
+        A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
+
+    A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta.unsqueeze(-2)
+    w = A @ k_eg
+    u = A @ v
+
+    mask_causal = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
+    S = q.new_zeros(B, H, K, V)
+    o = torch.zeros_like(v)
+
+    for i in range(NT):
+        q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i]
+        A_qk = torch.zeros(B, H, BT, BT, dtype=torch.float, device=q.device)
+        for j in range(BT):
+            k_j = k_i[..., j, :]
+            g_j = g_i[..., j : j + 1, :]
+            decay = (g_i - g_j).clamp(max=0).exp()
+            A_qk[..., j] = (q_i * decay * k_j.unsqueeze(-2)).sum(-1)
+        A_qk = A_qk.masked_fill(mask_causal, 0)
+        v_i = u_i - w_i @ S
+        o[:, :, i] = (q_i * g_i.exp()) @ S + A_qk @ v_i
+        g_last = g_i[:, :, -1]
+        S = S * g_last.unsqueeze(-1).exp()
+        k_dec = (g_last.unsqueeze(-2) - g_i).exp() * k_i
+        S = S + k_dec.transpose(-1, -2) @ v_i
+
+    o = o.reshape(B, H, -1, V)[:, :, :T]
+    o = o.transpose(1, 2).contiguous().to(initial_dtype)
+    return o
+
+
+class _HybridChunkKDA(torch.autograd.Function):
+    """Triton forward + PyTorch backward for chunk_kda on ROCm MI300X.
+
+    FLA Triton forward kernels work on ROCm but the backward kernels hang.
+    This uses Triton for the fast forward pass and recomputes with PyTorch
+    for the backward pass.
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, g, beta):
+        ctx.save_for_backward(q, k, v, g, beta)
+        with torch.no_grad():
+            o, _ = chunk_kda(q, k, v, g, beta, use_qk_l2norm_in_kernel=True)
+        return o
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        q, k, v, g, beta = ctx.saved_tensors
+        with torch.enable_grad():
+            q = q.detach().requires_grad_(True)
+            k = k.detach().requires_grad_(True)
+            v = v.detach().requires_grad_(True)
+            g = g.detach().requires_grad_(True)
+            beta = beta.detach().requires_grad_(True)
+            o = _torch_chunk_kda_fwd(q, k, v, g, beta, use_qk_l2norm_in_kernel=True)
+            o.backward(grad_output)
+        return q.grad, k.grad, v.grad, g.grad, beta.grad
+
+
+def hybrid_chunk_kda(q, k, v, g, beta):
+    """chunk_kda: Triton forward, PyTorch backward (for ROCm MI300X)."""
+    return _HybridChunkKDA.apply(q, k, v, g, beta)
+
+
+def _torch_chunk_kda_ckpt(q, k, v, g, beta):
+    """PyTorch-only fallback with gradient checkpointing."""
+    return _torch_chunk_kda_fwd(q, k, v, g, beta, use_qk_l2norm_in_kernel=True)
+
+
+@dataclass
+class KimiDeltaAttentionSubmodules:
+    """Module specs for the Kimi Delta Attention (KDA) layer.
+
+    Uses a single fused in_proj (with TELayerNormColumnParallelLinear) for
+    combined Q/K/V projection, mirroring GatedDeltaNet's approach.
+    A separate gate_norm is provided for the side projections (gate, beta,
+    output_gate) that cannot be included in the column-parallel in_proj.
+    """
+
+    in_proj: Union[ModuleSpec, type] = IdentityOp
+    gate_norm: Union[ModuleSpec, type] = IdentityOp
+    out_norm: Union[ModuleSpec, type] = IdentityOp
+    out_proj: Union[ModuleSpec, type] = IdentityOp
+
+
+class KimiDeltaAttention(MegatronModule):
+    """Kimi Delta Attention (KDA) layer class.
+
+    Based on the architecture from Kimi-Linear (Moonshot AI).
+    Key differences from GatedDeltaNet:
+    - Separate causal conv1d per Q, K, V (combined into one depthwise conv).
+    - Low-rank gate factorization (f_a -> f_b) instead of direct alpha projection.
+    - Low-rank output gate (g_a -> g_b) instead of direct gate projection.
+    - Uses chunk_kda / fused_kda_gate from fla.ops.kda.
+
+    KDA layer takes input with size [s, b, h] and returns output of the same size.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: KimiDeltaAttentionSubmodules,
+        layer_number: int = None,
+        bias: bool = False,
+        conv_bias: bool = False,
+        conv_init: Optional[float] = None,
+        A_init_range: Tuple[float, float] = (1, 16),
+        pg_collection: ProcessGroupCollection = None,
+    ):
+        if not HAVE_FLA_KDA and getattr(config, "use_fla_triton_kda", False):
+            raise ImportError(
+                "use_fla_triton_kda is set but FLA KDA ops are not installed. "
+                "Install fla-core or remove the flag to use the pure-PyTorch default."
+            )
+
+        super().__init__(config)
+
+        self.layer_number = layer_number
+        self.bias = bias
+        self.conv_bias = conv_bias
+        self.conv_init = conv_init
+        assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0]
+        self.A_init_range = A_init_range
+        assert pg_collection is not None, "pg_collection must be provided for KimiDeltaAttention"
+        self.pg_collection = pg_collection
+        self.tp_size = self.pg_collection.tp.size()
+        self.sp_size = self.tp_size if config.sequence_parallel else 1
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.act_fn = config.activation_func
+        self.activation = self.act_fn.__name__
+        self.conv_kernel_dim = config.linear_conv_kernel_dim
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_dim = config.linear_value_head_dim
+        self.num_k_heads = config.linear_num_key_heads
+        self.num_heads = config.linear_num_value_heads
+        self.qk_dim = self.head_k_dim * self.num_k_heads
+        self.v_dim = self.head_dim * self.num_heads
+        self.num_heads_local_tp = self.num_heads // self.tp_size
+        self.qk_dim_local_tp = self.qk_dim // self.tp_size
+        self.v_dim_local_tp = self.v_dim // self.tp_size
+
+        # --- Fused projection (column-parallel) ---
+        # Match GDN's parity recipe: pack EVERY `hidden_states → X` projection
+        # into a single big matmul, so the per-step kernel-launch overhead drops
+        # from ~6 launches/layer (FLA reference) or 4 (un-fused Primus) to 1.
+        # The fused output is split downstream into:
+        #
+        #   [ qkv  (qk_dim*2 + v_dim) | f_a (head_v_dim) | g_a (head_v_dim) | beta (num_v_heads) ]
+        #
+        # f_a and g_a are the low-rank-bottleneck *inputs* to the (cheap)
+        # f_b / g_b expansion projections — those stay separate because their
+        # input is the 64-dim bottleneck output, not hidden_states.
+        # b_proj's output is the raw beta gate (still sigmoid-applied later).
+        #
+        # FLA does these as six separate `nn.Linear` modules in
+        # `fla/layers/kda.py:142-189`; numerically identical, but on ROCm each
+        # separate matmul pays ~3-5 ms of HIP dispatch + autograd overhead, so
+        # GDN's parity work measured this fusion alone at ~250 ms/iter saved.
+        self.gate_dim_local_tp = self.num_heads_local_tp * self.head_k_dim  # used below
+        # NOTE on TP: the fused in_proj is a ColumnParallelLinear, which splits
+        # its output evenly across TP ranks. For the gate-bottleneck slices
+        # (f_a, g_a) this is incorrect — the low-rank gate REQUIRES each rank
+        # to see the FULL bottleneck output before f_b_proj / g_b_proj expand
+        # it (otherwise per-rank f_b would map a non-contiguous slice of the
+        # bottleneck to a chunk of the gate, producing the wrong output).
+        # Hard-assert tp_size=1 here; if you ever need TP>1 for KDA, the right
+        # recipe is to keep f_a_proj / g_a_proj as replicated nn.Linear modules
+        # (so each rank computes the full bottleneck independently) and only
+        # fuse q/k/v/beta into the column-parallel in_proj.
+        assert self.tp_size == 1, (
+            f"KDA fused in_proj currently requires tp_size=1 (got {self.tp_size}). "
+            "See KimiDeltaAttention.__init__ for the TP>1 recipe."
+        )
+        self.in_proj_dim = (
+            self.qk_dim * 2
+            + self.v_dim
+            + self.head_dim  # f_a output (bottleneck dim, = head_v_dim)
+            + self.head_dim  # g_a output (bottleneck dim, = head_v_dim)
+            + self.num_heads  # beta (per value-head scalar)
+        )
+        self.in_proj = build_module(
+            submodules.in_proj,
+            self.hidden_size,
+            self.in_proj_dim,
+            config=self.config,
+            init_method=self.config.init_method,
+            gather_output=False,
+            bias=bias,
+            skip_bias_add=False,
+            is_expert=False,
+            tp_comm_buffer_name="kda_in",
+            tp_group=self.pg_collection.tp,
+        )
+
+        # --- Combined causal conv1d for Q, K, V (depthwise) ---
+        self.conv_dim = self.qk_dim * 2 + self.v_dim
+        self.conv_dim_local_tp = self.conv_dim // self.tp_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim_local_tp,
+            out_channels=self.conv_dim_local_tp,
+            bias=conv_bias,
+            kernel_size=self.conv_kernel_dim,
+            groups=self.conv_dim_local_tp,
+            padding=self.conv_kernel_dim - 1,
+            device=torch.cuda.current_device(),
+            dtype=config.params_dtype,
+        )
+        setattr(self.conv1d.weight, "tensor_model_parallel", True)
+        if conv_bias:
+            setattr(self.conv1d.bias, "tensor_model_parallel", True)
+
+        # --- Norm for gate/beta/output_gate side projections ---
+        # Retained for backward compat with the TE-fused spec, which keeps
+        # `in_proj` as TELayerNormColumnParallelLinear and computes the
+        # side-projection inputs via this second norm.  In the no-TE spec
+        # the wrapper layer (`KimiDeltaAttentionLayer.forward`) already
+        # pre-norms hidden_states once, so submodules.gate_norm is
+        # `IdentityOp` here and the call is a no-op.
+        self.gate_norm = build_module(
+            submodules.gate_norm,
+            config=self.config,
+            hidden_size=self.hidden_size,
+            eps=self.config.layernorm_epsilon,
+        )
+
+        # --- Low-rank gate expansion: f_b (bottleneck → gate_dim) ---
+        # f_a (hidden → head_v_dim) is now FUSED into the big in_proj above
+        # so only the cheap 64→256 bottleneck-expand matmul remains here.
+        # Gate g has shape [B, T, H, K] (per-key-dim gating); output dim
+        # must match q/k after the optional repeat_interleave: num_heads * head_k_dim.
+        self.f_b_proj = nn.Linear(
+            self.head_dim,
+            self.gate_dim_local_tp,
+            bias=False,
+            device=torch.cuda.current_device(),
+            dtype=config.params_dtype,
+        )
+        setattr(self.f_b_proj.weight, "tensor_model_parallel", True)
+
+        # --- A_log and dt_bias (TP-split per head) ---
+        self.A_log = nn.Parameter(
+            torch.empty(
+                1,
+                1,
+                self.num_heads_local_tp,
+                1,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+        )
+        setattr(self.A_log, "tensor_model_parallel", True)
+
+        self.dt_bias = nn.Parameter(
+            torch.empty(
+                self.gate_dim_local_tp,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+        )
+        setattr(self.dt_bias, "tensor_model_parallel", True)
+
+        # b_proj (hidden → num_v_heads) is now FUSED into the big in_proj
+        # above; beta is read out of the in_proj split and sigmoid-activated
+        # in the forward pass.
+
+        # --- Low-rank output gate expansion: g_b (bottleneck → value_dim) ---
+        # g_a (hidden → head_v_dim) is FUSED into the big in_proj above; only
+        # the cheap 64→512 bottleneck-expand matmul remains here.
+        # Match FLA: g_proj's second linear has bias=True — the bias gives
+        # the output gate a learnable pre-sigmoid offset that compounds
+        # across all KDA layers (fla/layers/kda.py:189).
+        self.g_b_proj = nn.Linear(
+            self.head_dim,
+            self.v_dim_local_tp,
+            bias=True,
+            device=torch.cuda.current_device(),
+            dtype=config.params_dtype,
+        )
+        setattr(self.g_b_proj.weight, "tensor_model_parallel", True)
+        setattr(self.g_b_proj.bias, "tensor_model_parallel", True)
+
+        # --- Output norm and projection ---
+        # Match FLA: use FusedRMSNormGated (RMSNorm + sigmoid-gate + multiply in
+        # ONE Triton kernel). This avoids materializing the post-norm tensor and
+        # the fp32-upcast gate for backward — saves ~6.4 GiB of activation memory
+        # per rank at micro_batch=128 vs the unfused (out_norm + _apply_gated_norm)
+        # path. Enabled when fla.modules.FusedRMSNormGated is importable AND the
+        # config opts in via use_fla_fused_norm_gated (default: True when use_fla_triton_kda).
+        self._use_fla_fused_norm_gated = HAVE_FUSED_RMS_NORM_GATED and getattr(
+            self.config, "use_fla_fused_norm_gated", getattr(self.config, "use_fla_triton_kda", False)
+        )
+        if self._use_fla_fused_norm_gated:
+            self.out_norm = FusedRMSNormGated(
+                self.head_dim,
+                activation="sigmoid",
+                eps=self.config.layernorm_epsilon,
+                device=torch.cuda.current_device(),
+                dtype=config.params_dtype,
+            )
+        else:
+            self.out_norm = build_module(
+                submodules.out_norm,
+                config=self.config,
+                hidden_size=self.head_dim,
+                eps=self.config.layernorm_epsilon,
+            )
+        self.out_proj = build_module(
+            submodules.out_proj,
+            self.v_dim,
+            self.hidden_size,
+            config=self.config,
+            init_method=self.config.output_layer_init_method,
+            bias=bias,
+            input_is_parallel=True,
+            skip_bias_add=True,
+            is_expert=False,
+            tp_comm_buffer_name="kda_out",
+            tp_group=self.pg_collection.tp,
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.config.perform_initialization:
+            with get_cuda_rng_tracker().fork():
+                if self.conv_init is not None:
+                    nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+                A = torch.empty(
+                    self.num_heads_local_tp,
+                    dtype=torch.float32,
+                    device=torch.cuda.current_device(),
+                ).uniform_(*self.A_init_range)
+                self.A_log.data.copy_(torch.log(A).view(1, 1, -1, 1))
+                # Match FLA dt_bias init exactly (fla/layers/kda.py:180-184):
+                # log-uniform sample dt in [0.001, 0.1], then store its inverse-
+                # softplus so that softplus(0 + dt_bias) ≈ dt at iter 0. Replaces
+                # the naive ones_ init which gave dt ≈ 1.31 — a ~20x larger
+                # initial decay step that compounds across all KDA layers and
+                # produces a noticeable loss-curve drift from FLA by iter ~100.
+                dt = torch.exp(
+                    torch.rand(
+                        self.gate_dim_local_tp,
+                        dtype=torch.float32,
+                        device=torch.cuda.current_device(),
+                    )
+                    * (math.log(0.1) - math.log(0.001))
+                    + math.log(0.001)
+                ).clamp(min=1e-4)
+                inv_dt = dt + torch.log(-torch.expm1(-dt))
+                self.dt_bias.data.copy_(inv_dt)
+                # g_b_proj.bias = 0 (PyTorch nn.Linear default for bias=True
+                # initialises bias uniform in [-1/sqrt(in), +1/sqrt(in)]; FLA
+                # uses default nn.Linear which gives the same thing, so we
+                # leave it as nn.Linear default — no explicit init here).
+
+    # NOTE: GDN's matched-parity forward (`gated_delta_net.py:285`) does NOT
+    # carry `@torch.compiler.disable`. The decorator was added defensively
+    # for KDA because the chunk_kda Triton kernel doesn't trace through
+    # `torch.compile`, but Megatron does not currently wrap mixer forwards
+    # in `torch.compile` at all, so the decorator was only adding ~20-30 ms
+    # of per-call eager-dispatch overhead (12 layers × ~2 ms × 2 fwd+bwd).
+    # Removing it matches GDN's recipe exactly.
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Tensor,
+        key_value_states: Optional[Tensor] = None,
+        inference_context: Optional[BaseInferenceContext] = None,
+        attention_bias: Optional[Tensor] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+        sequence_len_offset: Optional[int] = None,
+        *,
+        inference_params: Optional[BaseInferenceContext] = None,
+        **kwargs,
+    ):
+        """Forward pass through the KDA module.
+
+        Args:
+            hidden_states (Tensor): [s, b, h] input tensor (un-normed).
+            attention_mask (Tensor): Attention mask (reserved for future use).
+
+        Returns:
+            Tuple[Tensor, Tensor]: KDA output and bias.
+        """
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        local_seq_len, batch, _ = hidden_states.shape
+        seq_len = local_seq_len * self.sp_size
+        use_fla_triton = (
+            (not self.config.deterministic_mode)
+            and HAVE_FLA_KDA
+            and getattr(self.config, "use_fla_triton_kda", False)
+        )
+
+        if not hasattr(self, "_kda_kernel_logged"):
+            self._kda_kernel_logged = True
+            use_hybrid = getattr(self.config, "use_fla_triton_kda_hybrid", False)
+            in_kernel_gate = getattr(self.config, "use_fla_kda_in_kernel_gate", True)
+            if use_fla_triton and use_hybrid:
+                mode = "Hybrid (Triton fwd, PyTorch bwd), gate materialized"
+            elif use_fla_triton and in_kernel_gate:
+                mode = "FLA Triton fwd+bwd, gate fused in kernel"
+            elif use_fla_triton:
+                mode = "FLA Triton fwd+bwd, explicit fused_kda_gate (tw006)"
+            else:
+                mode = "Pure PyTorch fallback (gradient checkpointed)"
+            norm_mode = (
+                "FusedRMSNormGated (FLA)"
+                if self._use_fla_fused_norm_gated
+                else "unfused (out_norm + sigmoid-multiply)"
+            )
+            logger.warning(
+                f"[KDA layer {self.layer_number}] kernel={mode} | norm={norm_mode} | "
+                f"HAVE_FLA_KDA={HAVE_FLA_KDA} use_fla_triton_kda={getattr(self.config, 'use_fla_triton_kda', False)} "
+                f"deterministic={self.config.deterministic_mode}"
+            )
+
+        if inference_context is not None:
+            raise NotImplementedError("KDA does not support inference yet.")
+        if packed_seq_params is not None:
+            raise NotImplementedError("KDA does not support packed sequences yet.")
+
+        # --- Single fused projection ---
+        # One big GEMM produces [q | k | v | f_a_out | g_a_out | beta_raw] for the
+        # whole layer — matches GDN's parity recipe and removes 3 extra hidden→X
+        # matmul launches per layer (~5 ms × 12 layers ≈ saved 60 ms/iter alone,
+        # plus the corresponding backward GEMMs, plus removed activation copies
+        # for h_normed).
+        nvtx_range_push(suffix="kda_in_proj")
+        fused, _ = self.in_proj(hidden_states)
+        nvtx_range_pop(suffix="kda_in_proj")
+
+        # s b d -> b s d  (output is full seq_len from column-parallel)
+        fused = fused.transpose(0, 1)
+
+        # Split into [qkv | f_a_out | g_a_out | beta_raw]. The slice sizes are
+        # the per-TP local dims; sum equals self.in_proj_dim // tp_size.
+        qkv_local = self.qk_dim_local_tp * 2 + self.v_dim_local_tp
+        head_dim_local_tp = self.head_dim // self.tp_size
+        qkv, f_a_out, g_a_out, beta_raw = torch.split(
+            fused,
+            [qkv_local, head_dim_local_tp, head_dim_local_tp, self.num_heads_local_tp],
+            dim=-1,
+        )
+
+        # --- Causal conv1d on combined QKV ---
+        # Three backends, chosen at runtime:
+        #   (1) PRIMUS_FLA_CONV=1  → FLA's Triton causal_conv1d. Accepts
+        #       [B, T, D] directly, matches FLA's reference run bit-for-bit,
+        #       and removes the two transposes + contiguous() copy below.
+        #       Saves ~3 ms/layer × 12 = ~35 ms/iter and ~1.6 GiB/iter peak
+        #       (the contiguous() copy was a full-qkv allocation).
+        #   (2) Tri-Dao's causal_conv1d_fn (CUDA package) — current default.
+        #   (3) Pure-PyTorch fallback when neither is available or
+        #       deterministic_mode is set.
+        nvtx_range_push(suffix="kda_conv")
+        _use_fla_conv = getattr(_get_args(), "use_fla_short_conv", False)
+        if _use_fla_conv and _fla_causal_conv1d is not None and not self.config.deterministic_mode:
+            assert self.activation in ["silu", "swish"]
+            qkv, _ = _fla_causal_conv1d(
+                x=qkv,
+                weight=self.conv1d.weight.squeeze(1),  # d, 1, w -> d, w
+                bias=self.conv1d.bias,
+                activation=self.activation,
+                backend="triton",
+            )
+        else:
+            qkv = qkv.transpose(1, 2).contiguous()  # b s d -> b d s
+            if (causal_conv1d_fn is None) or self.config.deterministic_mode:
+                qkv = self.act_fn(self.conv1d(qkv)[..., :seq_len])
+            else:
+                assert self.activation in ["silu", "swish"]
+                qkv = causal_conv1d_fn(
+                    x=qkv,
+                    weight=self.conv1d.weight.squeeze(1),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                )
+            # b d s -> b s d. Do NOT contiguous() here — the downstream
+            # torch.split produces non-contiguous views along dim=-1 regardless,
+            # so a copy is forced inside .reshape() anyway.
+            qkv = qkv.transpose(1, 2)  # b d s -> b s d
+        nvtx_range_pop(suffix="kda_conv")
+
+        # Split conv output into Q, K, V
+        q, k, v = torch.split(
+            qkv,
+            [self.qk_dim_local_tp, self.qk_dim_local_tp, self.v_dim_local_tp],
+            dim=-1,
+        )
+
+        # Reshape to (batch, seq, heads, head_dim)
+        q = q.reshape(batch, seq_len, -1, self.head_k_dim)
+        k = k.reshape(batch, seq_len, -1, self.head_k_dim)
+        v = v.reshape(batch, seq_len, -1, self.head_dim)
+
+        if self.num_heads // self.num_k_heads > 1:
+            q = q.repeat_interleave(self.num_heads // self.num_k_heads, dim=2)
+            k = k.repeat_interleave(self.num_heads // self.num_k_heads, dim=2)
+
+        # --- Gate (low-rank expansion only — f_a is already inside in_proj) ---
+        nvtx_range_push(suffix="kda_gate")
+        g = self.f_b_proj(f_a_out)
+        g = g.reshape(batch, seq_len, self.num_heads_local_tp, self.head_k_dim)
+
+        # Match FLA: when using the Triton fwd+bwd path, pass RAW g into the
+        # kernel and let it fuse `-exp(A_log) * softplus(g + dt_bias) + cumsum`
+        # internally (use_gate_in_kernel=True). The kernel recomputes the gate
+        # in backward, so the fp32 [B,T,H,K] activated-gate tensor is never
+        # materialized — saves ~3.2 GiB of activation memory at micro_batch=128.
+        # For non-Triton paths (deterministic/CPU/fallback) we still compute
+        # the gate up front in PyTorch.
+        # Two opt-outs:
+        #   - use_fla_triton_kda_hybrid: routes to hybrid_chunk_kda (Triton fwd,
+        #     PyTorch bwd) for debugging; gate is materialized in PyTorch.
+        #   - use_fla_kda_in_kernel_gate=False: keeps the standard chunk_kda
+        #     path but materializes the gate up front (via fused_kda_gate),
+        #     mirroring the pre-fusion tw006 numerics. The forward output is
+        #     bit-identical between paths in fp32 but the bf16 in-kernel fused
+        #     `-exp(A_log)*softplus(g+dt_bias)` accumulator has +/-1 ulp drift
+        #     vs the explicit-gate path; with 12 layers of compounding this
+        #     gives ~0.2 lm-loss above tw006 by iter 200. Set to False to
+        #     recover tw006-tight loss curve at the cost of ~3 GiB extra
+        #     activation memory.
+        _use_in_kernel_gate = (
+            use_fla_triton
+            and getattr(self.config, "use_fla_kda_in_kernel_gate", True)
+            and not getattr(self.config, "use_fla_triton_kda_hybrid", False)
+        )
+        if not _use_in_kernel_gate:
+            if use_fla_triton:
+                g = fused_kda_gate(g, self.A_log.view(-1), dt_bias=self.dt_bias)
+            else:
+                g = torch_kda_gate(g, self.A_log.view(-1), dt_bias=self.dt_bias)
+
+        # beta = sigmoid of the raw beta slice from the fused in_proj. Match
+        # FLA's `b_proj(h).sigmoid()` (fla/layers/kda.py:251) in bf16 exactly.
+        # The earlier fp32 upcast was a defensive measure against bf16 drift
+        # compounding across 12 layers, but with both in-kernel fusions enabled
+        # the drift was empirically <0.001 vs fp32 at iter 200 — within batch
+        # noise. Saves ~256 MiB/layer × 12 = ~6 GiB activation peak and ~5 ms.
+        beta = beta_raw.sigmoid()
+        nvtx_range_pop(suffix="kda_gate")
+
+        # --- Core KDA attention ---
+        # Q/K/V contiguity: only required when the in_proj is the TE-fused
+        # variant (TELayerNormColumnParallelLinear), which leaves a non-contig
+        # stride pattern that interacts pathologically with chunk_kda's bwd
+        # (~29 GiB extra activation memory in the TE-fused path). With the
+        # no-TE spec (plain ColumnParallelLinear → torch.split), strides are
+        # already clean and match FLA's `rearrange()` output, so the explicit
+        # .contiguous() calls just waste ~5-8 GiB on saved-for-backward
+        # duplicates. Auto-detected via gate_norm == IdentityOp.
+        if not isinstance(self.gate_norm, IdentityOp):
+            q = q.contiguous()
+            k = k.contiguous()
+            v = v.contiguous()
+        nvtx_range_push(suffix="kda_attn")
+        use_hybrid = getattr(self.config, "use_fla_triton_kda_hybrid", False)
+        if _use_in_kernel_gate:
+            # FLA's new (post-fusion) call site — fuses the gate compute and
+            # qk-l2norm inside chunk_kda. Smallest activation footprint, but
+            # the bf16 accumulator drifts ~+0.2 lm-loss vs the explicit-gate
+            # path on ROCm at 12 layers depth.
+            core_attn_out, _ = chunk_kda(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                A_log=self.A_log.view(-1),
+                dt_bias=self.dt_bias,
+                use_qk_l2norm_in_kernel=True,
+                use_gate_in_kernel=True,
+            )
+        elif use_fla_triton and use_hybrid:
+            core_attn_out = hybrid_chunk_kda(q, k, v, g, beta)
+        elif use_fla_triton:
+            # tw006 call site — gate was pre-computed by fused_kda_gate above,
+            # only qk-l2norm is fused. Bit-identical to FLA's pre-fusion code
+            # and to the explicit-gate Triton path; this is the configuration
+            # that hit loss=4.7281 @ iter 500 vs FLA/8=4.7350.
+            core_attn_out, _ = chunk_kda(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                use_qk_l2norm_in_kernel=True,
+            )
+        else:
+            core_attn_out = _grad_checkpoint(
+                _torch_chunk_kda_ckpt,
+                q,
+                k,
+                v,
+                g,
+                beta,
+                use_reentrant=False,
+            )
+        nvtx_range_pop(suffix="kda_attn")
+
+        # --- Output gate (g_b expansion only — g_a is already inside in_proj) + gated norm ---
+        nvtx_range_push(suffix="kda_out_gate")
+        gate = self.g_b_proj(g_a_out)
+        gate = gate.reshape(batch, seq_len, -1, self.head_dim)
+        if self._use_fla_fused_norm_gated:
+            # FusedRMSNormGated: RMSNorm(core_attn_out) * sigmoid(gate) in ONE
+            # Triton kernel — no intermediate post-norm tensor, no fp32 upcast.
+            # Matches fla/layers/kda.py exactly.
+            norm_out = self.out_norm(core_attn_out, gate)
+        else:
+            norm_out = self._apply_gated_norm(core_attn_out, gate)
+        nvtx_range_pop(suffix="kda_out_gate")
+
+        # b s (h*d) -> s b (h*d)
+        norm_out = norm_out.reshape(batch, seq_len, -1)
+        norm_out = norm_out.transpose(0, 1).contiguous()
+
+        # --- Output projection ---
+        nvtx_range_push(suffix="kda_out_proj")
+        out, out_bias = self.out_proj(norm_out)
+        nvtx_range_pop(suffix="kda_out_proj")
+
+        return out, out_bias
+
+    def _apply_gated_norm(self, x, gate):
+        x_dtype = x.dtype
+        x = x.reshape(-1, x.shape[-1])
+        y = self.out_norm(x)
+        gate = gate.reshape(-1, gate.shape[-1])
+        y = y * gate.float().sigmoid()
+        y = y.to(x_dtype)
+        return y
+
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None):
+        """Provide a sharded state dictionary for distributed checkpointing."""
+        sharded_state_dict = {}
+        self._save_to_state_dict(sharded_state_dict, "", keep_vars=True)
+        sharded_state_dict = make_sharded_tensors_for_checkpoint(
+            sharded_state_dict,
+            prefix,
+            tensor_parallel_layers_axis_map={
+                "A_log": 2,
+                "dt_bias": 0,
+                "f_b_proj.weight": 0,
+                "g_b_proj.weight": 0,
+                "g_b_proj.bias": 0,
+            },
+            sharded_offsets=sharded_offsets,
+            tp_group=(tp_group if tp_group is not None else self.pg_collection.tp),
+            dp_cp_group=metadata["dp_cp_group"],
+        )
+
+        tp_group = tp_group if tp_group is not None else self.pg_collection.tp
+        for name, module in self.named_children():
+            if name == "conv1d":
+                module_sd = module.state_dict(prefix="", keep_vars=True)
+                tp_sharding_map = {"weight": 0}
+                if self.conv_bias:
+                    tp_sharding_map["bias"] = 0
+                module_sharded_sd = make_sharded_tensors_for_checkpoint(
+                    module_sd,
+                    f"{prefix}{name}.",
+                    tp_sharding_map,
+                    sharded_offsets,
+                    tp_group=tp_group,
+                    dp_cp_group=metadata["dp_cp_group"],
+                )
+            else:
+                module_sharded_sd = sharded_state_dict_default(
+                    module, f"{prefix}{name}.", sharded_offsets, metadata, tp_group=tp_group
+                )
+            sharded_state_dict.update(module_sharded_sd)
+
+        # Split the combined in_proj into named chunks for checkpoint compatibility.
+        # in_proj output layout (after GDN-style fusion):
+        #   [ query | key | value | f_a | g_a | beta ]
+        # First three are TP-sharded along output dim (qk/v split per rank),
+        # f_a/g_a have dim 64 (head_v_dim, NOT TP-sharded since head_v_dim < tp_size
+        # in some configs — we still split per-rank as standard column-parallel),
+        # beta has num_heads dim which is TP-sharded the same way.
+        in_proj_dim_local_tp = self.in_proj_dim // self.tp_size
+        assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim_local_tp, (
+            in_proj_dim_local_tp,
+            sharded_state_dict[f"{prefix}in_proj.weight"],
+        )
+        sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory(
+            sharded_state_dict[f"{prefix}in_proj.weight"],
+            [
+                self.qk_dim // self.tp_size,
+                self.qk_dim // self.tp_size,
+                self.v_dim // self.tp_size,
+                self.head_dim,  # f_a (bottleneck dim, not sharded)
+                self.head_dim,  # g_a (bottleneck dim, not sharded)
+                self.num_heads // self.tp_size,  # beta
+            ],
+            ["query", "key", "value", "f_a", "g_a", "beta"],
+            0,
+        )
+
+        # Split conv1d (same ordering as in_proj: Q, K, V)
+        conv_layer_name_list = ["conv1d.weight"]
+        assert sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) == self.conv_dim_local_tp, (
+            self.conv_dim_local_tp,
+            sharded_state_dict[f"{prefix}conv1d.weight"],
+        )
+        if self.conv_bias:
+            conv_layer_name_list.append("conv1d.bias")
+            assert sharded_state_dict[f"{prefix}conv1d.bias"].data.size(0) == self.conv_dim_local_tp, (
+                self.conv_dim_local_tp,
+                sharded_state_dict[f"{prefix}conv1d.bias"],
+            )
+        for conv_layer_name in conv_layer_name_list:
+            sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory(
+                sharded_state_dict[f"{prefix}{conv_layer_name}"],
+                [
+                    self.qk_dim // self.tp_size,
+                    self.qk_dim // self.tp_size,
+                    self.v_dim // self.tp_size,
+                ],
+                ["query", "key", "value"],
+                0,
+            )
+
+        return sharded_state_dict
+
+    def backward_dw(self):
+        """Execute weight gradient computation for all linear layers."""
+        self.in_proj.backward_dw()
+        self.out_proj.backward_dw()
+
+
+def _split_tensor_factory(
+    orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int
+) -> ShardedTensorFactory:
+    """Builds a factory that splits a given ShardedTensor into several independent chunks."""
+    assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten)
+    orig_sh_ten_no_data = orig_sh_ten.without_data()
+
+    if sum(split_sections) != orig_sh_ten_no_data.local_shape[split_dim]:
+        raise ValueError(
+            f"Split sections must cover the whole dimension size, "
+            f"got {split_sections=} vs dimensions size "
+            f"{orig_sh_ten_no_data.local_shape[split_dim]}"
+        )
+
+    assert not isinstance(
+        split_sections, int
+    ), "Splitting into predefined section sizes is supported (`split_sections` must be a list)"
+    assert len(split_sections) == len(split_names), (len(split_sections), len(split_names))
+
+    @torch.no_grad()
+    def sh_ten_build_fn(key: str, t: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice]):
+        factory_sh_ten = replace(
+            orig_sh_ten_no_data,
+            key=key,
+            data=t,
+            dtype=t.dtype,
+            replica_id=replica_id,
+            flattened_range=flattened_range,
+        )
+
+        chunk_sh_tens = []
+        split_start = 0
+        for split_size, split_name in zip(split_sections, split_names):
+            split_chunks = factory_sh_ten.narrow(split_dim, split_start, split_size)
+            for sh_ten in split_chunks:
+                sh_ten.key = f"{sh_ten.key}.{split_name}"
+            chunk_sh_tens.extend(split_chunks)
+            split_start += split_size
+
+        assert split_start == orig_sh_ten_no_data.local_shape[split_dim], (
+            split_start,
+            orig_sh_ten_no_data.local_shape[split_dim],
+        )
+        assert sum(sh_ten.data.numel() for sh_ten in chunk_sh_tens) == t.numel(), (
+            chunk_sh_tens,
+            t.shape,
+        )
+        return chunk_sh_tens
+
+    @torch.no_grad()
+    def sh_ten_merge_fn(sub_state_dict):
+        return torch.cat(sub_state_dict)
+
+    return ShardedTensorFactory(
+        orig_sh_ten.key, orig_sh_ten.data, sh_ten_build_fn, sh_ten_merge_fn, orig_sh_ten.replica_id
+    )
diff --git a/primus/backends/megatron/core/models/hybrid/kimi_delta_attention_layer.py b/primus/backends/megatron/core/models/hybrid/kimi_delta_attention_layer.py
new file mode 100644
index 000000000..bf00614f2
--- /dev/null
+++ b/primus/backends/megatron/core/models/hybrid/kimi_delta_attention_layer.py
@@ -0,0 +1,136 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Union
+
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import deprecate_inference_params
+from torch import Tensor
+
+
+@dataclass
+class KimiDeltaAttentionLayerSubmodules:
+    """Configuration class for specifying the submodules of a KDA layer.
+
+    Two layouts are supported:
+    1. TE-fused (default): norm=IdentityOp; the mixer's in_proj is
+       TELayerNormColumnParallelLinear and absorbs the pre-norm.
+    2. No-TE / FLA-style: norm=WrappedTorchNorm; the mixer's in_proj is plain
+       ColumnParallelLinear and gate_norm=IdentityOp (single pre-norm matches
+       fla/models/kda/modeling_kda.py KDABlock pattern — saves one redundant
+       gate_norm launch per layer).
+    """
+
+    norm: Union[ModuleSpec, type] = IdentityOp
+    mixer: Union[ModuleSpec, type] = IdentityOp
+    kda_bda: Union[ModuleSpec, type] = IdentityOp
+
+    sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
+
+
+class KimiDeltaAttentionLayer(MegatronModule):
+    """A single Kimi Delta Attention layer wrapping the KDA mixer.
+
+    Analogous to GatedDeltaNetLayer. The pre-normalization is handled
+    inside the mixer via the fused TELayerNormColumnParallelLinear in_proj,
+    so this wrapper only manages residual + bias-dropout-add.
+
+    The forward interface matches what HybridStack expects for the
+    mamba_layer slot.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: KimiDeltaAttentionLayerSubmodules,
+        layer_number: int = 1,
+        residual_in_fp32: bool = False,
+        pg_collection: ProcessGroupCollection = None,
+    ):
+        super().__init__(config)
+        assert pg_collection is not None, "pg_collection must be provided for KimiDeltaAttentionLayer"
+
+        self.config = config
+        self.submodules_config = submodules
+        self.layer_number = layer_number
+        self.residual_in_fp32 = residual_in_fp32
+        self.hidden_dropout = config.hidden_dropout
+
+        self.mixer = build_module(
+            submodules.mixer,
+            self.config,
+            layer_number=layer_number,
+            pg_collection=pg_collection,
+        )
+        # Optional pre-norm. With TE-fused in_proj this is IdentityOp (norm
+        # is absorbed into TELayerNormColumnParallelLinear). With plain
+        # ColumnParallelLinear in_proj this is WrappedTorchNorm and matches
+        # fla/models/kda/modeling_kda.py:113 `hidden_states = self.attn_norm(...)`.
+        # eps forwarded explicitly because WrappedTorchNorm defaults to 1e-5
+        # while KDA configs (and FLA) use 1e-6.
+        self.norm = build_module(
+            submodules.norm,
+            self.config,
+            self.config.hidden_size,
+            eps=self.config.layernorm_epsilon,
+        )
+        self.kda_bda = build_module(submodules.kda_bda)
+        self.bias_dropout_add_exec_handler = torch.enable_grad
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        inference_context: Optional[BaseInferenceContext] = None,
+        rotary_pos_emb: Optional[Tensor] = None,
+        *,
+        inference_params: Optional[BaseInferenceContext] = None,
+    ):
+        """Forward pass through the KDA layer.
+
+        Args:
+            hidden_states (Tensor): Input tensor of shape [s, b, h].
+            attention_mask (Tensor, optional): Attention mask forwarded to the mixer.
+            inference_context (BaseInferenceContext, optional): Inference context.
+            rotary_pos_emb (Tensor, optional): Rotary positional embeddings (unused).
+
+        Returns:
+            Tensor: Transformed hidden states of shape [s, b, h].
+        """
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        residual = hidden_states
+        if self.residual_in_fp32:
+            residual = residual.to(torch.float32)
+
+        hidden_states = hidden_states.to(dtype=self.config.params_dtype)
+        hidden_states = self.norm(hidden_states)
+
+        mixer_out_with_bias = self.mixer(hidden_states, attention_mask, inference_context=inference_context)
+
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.kda_bda(training=self.training, fused=self.config.bias_dropout_fusion)(
+                mixer_out_with_bias, residual, self.hidden_dropout
+            )
+
+        return hidden_states
+
+    def sharded_state_dict(
+        self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None
+    ) -> ShardedStateDict:
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        prefixed_map = {
+            f"{prefix}{k}": f"{prefix}{v}"
+            for k, v in self.submodules_config.sharded_state_dict_keys_map.items()
+        }
+        if prefixed_map:
+            apply_prefix_mapping(sharded_state_dict, prefixed_map)
+        return sharded_state_dict
diff --git a/primus/backends/megatron/core/models/hybrid/mamba_layer_adapter.py b/primus/backends/megatron/core/models/hybrid/mamba_layer_adapter.py
new file mode 100644
index 000000000..1b216a1bf
--- /dev/null
+++ b/primus/backends/megatron/core/models/hybrid/mamba_layer_adapter.py
@@ -0,0 +1,56 @@
+"""Adapter that lets upstream Megatron `MambaLayer` plug into Primus's
+`HybridStack` (the one used by the GDN/KDA hybrids).
+
+Why this exists
+---------------
+`HybridStack.__init__` builds each mamba layer with:
+
+    build_module(
+        submodules.mamba_layer,
+        config=...,
+        residual_in_fp32=...,
+        layer_number=...,
+        pg_collection=...,
+    )
+
+`GatedDeltaNetLayer.__init__` and `KimiDeltaAttentionLayer.__init__` both
+accept the `residual_in_fp32` kwarg, so they slot in fine.  Upstream
+`MambaLayer.__init__`, however, does NOT accept it — it has
+`pp_layer_offset` in its place — so plugging `MambaLayer` directly into
+`HybridStack` raises:
+
+    TypeError: MambaLayer.__init__() got an unexpected keyword argument
+    'residual_in_fp32'
+
+We avoid touching upstream Megatron by wrapping `MambaLayer` in a thin
+adapter that accepts `residual_in_fp32` (stored on `self` for parity with
+the other hybrid layers), and never forwards it to `MambaLayer.__init__`.
+`pp_layer_offset` defaults to 0 (HybridStack doesn't have pipeline parallel
+plumbing for the mamba leg anyway — pp_offset is applied separately to the
+TransformerLayer attention/MLP branches).
+"""
+
+from __future__ import annotations
+
+from megatron.core.ssm.mamba_layer import MambaLayer
+
+
+class Mamba2HybridLayer(MambaLayer):
+    """`MambaLayer` that silently accepts `residual_in_fp32`.
+
+    Inherits the full forward/init path of upstream `MambaLayer`; only the
+    constructor is shimmed to filter the extra kwarg coming from `HybridStack`.
+    """
+
+    def __init__(
+        self,
+        *args,
+        residual_in_fp32: bool = False,
+        pp_layer_offset: int = 0,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, pp_layer_offset=pp_layer_offset, **kwargs)
+        # Persist for parity with GatedDeltaNetLayer/KimiDeltaAttentionLayer;
+        # MambaLayer's own forward path manages residual dtype internally, so
+        # nothing else needs to consume this flag at runtime.
+        self.residual_in_fp32 = residual_in_fp32
diff --git a/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py
new file mode 100644
index 000000000..73714fefc
--- /dev/null
+++ b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py
@@ -0,0 +1,686 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+BF16 Master Weight optimizer for FSDP2 mixed precision training.
+
+Model parameters are stored natively in BF16 (no FP32-to-BF16 cast during
+all-gather), while FP32 master weight copies are maintained internally for
+optimizer precision. This follows Megatron's Float16OptimizerWithFloat16Params
+pattern adapted for FSDP2 DTensor parameters.
+
+Compared to the FP32 optimizer (fsdp2_fp32_optimizer.py), this eliminates
+the per-layer CopyFunctor<BFloat16, float> kernel in every forward all-gather,
+at the cost of +2 bytes/param for the FP32 master copy.
+
+Memory per parameter per GPU (sharded):
+    FP32 optimizer: 4 (FP32 param) + 4 (exp_avg) + 4 (exp_avg_sq) = 12 bytes
+    This optimizer: 2 (BF16 param) + 4 (FP32 master) + 4 + 4     = 14 bytes
+"""
+
+from itertools import chain
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.optimizer import (
+    get_param_id_to_sharded_param_map,
+    make_sharded_optimizer_tensor,
+    optim_state_to_sharding_state,
+)
+from megatron.core.optimizer.optimizer import MegatronOptimizer, _zero_grad_group_helper
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+
+from primus.core.utils.module_utils import log_rank_0
+
+if TYPE_CHECKING:
+    from megatron.core.process_groups_config import ProcessGroupCollection
+
+
+def _safe_log_rank_0(msg: str):
+    try:
+        log_rank_0(msg)
+    except (AttributeError, TypeError):
+        import torch.distributed as dist
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(msg)
+
+
+class FSDP2BF16MasterWeightOptimizer(MegatronOptimizer):
+    """BF16 optimizer with FP32 master weights for FSDP2.
+
+    Model parameters live in BF16 (eliminating the FP32->BF16 cast in every
+    forward all-gather). FP32 master copies are maintained for the optimizer
+    step, matching Megatron's Float16OptimizerWithFloat16Params pattern.
+
+    Three parameter groups are tracked:
+        bf16_groups: original BF16 model parameters (what FSDP2 shards/all-gathers)
+        fp32_from_bf16_groups: FP32 master copies (what the optimizer steps on)
+        fp32_from_fp32_groups: natively FP32 params (e.g. LayerNorm), no master copy
+
+    Args:
+        optimizer: Base PyTorch optimizer (param_groups already point to FP32 masters).
+        config: OptimizerConfig from Megatron.
+        init_state_fn: Function to initialize optimizer state tensors.
+        bf16_groups: List of lists of original BF16 model parameters.
+        fp32_from_bf16_groups: List of lists of FP32 master weight copies.
+        fp32_from_fp32_groups: List of lists of natively FP32 parameters.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        config: OptimizerConfig,
+        init_state_fn: Callable,
+        bf16_groups: List[List[torch.nn.Parameter]],
+        fp32_from_bf16_groups: List[List[torch.Tensor]],
+        fp32_from_fp32_groups: List[List[torch.nn.Parameter]],
+        use_foreach: bool = True,
+    ):
+        super().__init__(optimizer, config, init_state_fn)
+        self._scale = torch.tensor([1.0], dtype=torch.float, device="cuda")
+        self.is_stub_optimizer = optimizer is None
+
+        self.bf16_groups = bf16_groups
+        self.fp32_from_bf16_groups = fp32_from_bf16_groups
+        self.fp32_from_fp32_groups = fp32_from_fp32_groups
+
+        self._use_foreach = use_foreach
+        self._foreach_ready = False
+
+        from megatron.training import get_args
+
+        args = get_args()
+        self._grad_norm_accumulator = getattr(args, "_grad_norm_accumulator", None)
+        self._validate_count = 0
+
+    @torch.compiler.disable
+    @torch.no_grad()
+    def warmup_foreach_cache(self):
+        """Eagerly initialize the foreach cache before torch.compile wraps step().
+
+        Creates temporary zero gradients to satisfy _init_foreach_cache's
+        requirement that grads exist (including DTensor placements for the
+        assertion at line ~164), then clears them.
+        Must be called after FSDP wrapping but before torch.compile wraps
+        optimizer.step.
+        """
+        if self._foreach_ready or not self._use_foreach:
+            return
+
+        from torch.distributed.tensor import DTensor
+
+        for group in self.bf16_groups:
+            for p in group:
+                if p.grad is not None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    local = torch.zeros(
+                        p.data.to_local().shape,
+                        dtype=p.dtype,
+                        device=p.device,
+                    )
+                    p.grad = DTensor.from_local(
+                        local,
+                        device_mesh=p.data.device_mesh,
+                        placements=p.data.placements,
+                    )
+                else:
+                    p.grad = torch.zeros(p.shape, dtype=p.dtype, device=p.device)
+        self._init_foreach_cache()
+        for group in self.bf16_groups:
+            for p in group:
+                p.grad = None
+
+    @torch.compiler.disable
+    @torch.no_grad()
+    def _init_foreach_cache(self):
+        """Build cached flat tensor lists for _foreach_copy_ batched operations.
+
+        Called lazily on first prepare_grads() invocation (grads must exist for
+        DTensor spec extraction). Uses public DTensor APIs only.
+        """
+        from torch.distributed.tensor import DTensor
+
+        try:
+            from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+                WeightWithFP8AllGatherTensor,
+            )
+
+            self._has_fp8_subclass = True
+        except ImportError:
+            self._has_fp8_subclass = False
+
+        bf16_flat = [p for group in self.bf16_groups for p in group]
+        fp32_flat = [p for group in self.fp32_from_bf16_groups for p in group]
+        self._bf16_flat = bf16_flat
+        self._fp32_flat = fp32_flat
+
+        bf16_inners = []
+        for p in bf16_flat:
+            t = p.data
+            if isinstance(t, DTensor):
+                t = t.to_local()
+            if self._has_fp8_subclass and isinstance(t, WeightWithFP8AllGatherTensor):
+                t = t.inner_data()
+            bf16_inners.append(t)
+        self._bf16_inners = bf16_inners
+
+        fp32_locals = []
+        for m in fp32_flat:
+            t = m
+            if isinstance(t, DTensor):
+                t = t.to_local()
+            fp32_locals.append(t)
+        self._fp32_locals = fp32_locals
+
+        fp32_grad_locals = [torch.empty_like(fl) for fl in fp32_locals]
+        self._fp32_grad_locals = fp32_grad_locals
+
+        fp32_grad_dtensors = []
+        for m, gl, bp in zip(fp32_flat, fp32_grad_locals, bf16_flat):
+            if isinstance(m, DTensor):
+                if m.placements != bp.grad.placements:
+                    raise RuntimeError(
+                        f"FP32 master placements {m.placements} != "
+                        f"BF16 grad placements {bp.grad.placements}"
+                    )
+                dg = DTensor.from_local(gl, device_mesh=m.device_mesh, placements=m.placements)
+            else:
+                dg = gl
+            fp32_grad_dtensors.append(dg)
+            m.grad = dg
+        self._fp32_grad_dtensors = fp32_grad_dtensors
+
+        self._foreach_ready = True
+        n_params = len(bf16_flat)
+        _safe_log_rank_0(
+            f"[FSDP2BF16MasterWeightOptimizer] foreach cache initialized: "
+            f"{n_params} params, fp8_subclass={self._has_fp8_subclass}"
+        )
+
+    def zero_grad(self, set_to_none=True):
+        if self.is_stub_optimizer:
+            return
+        if self._grad_norm_accumulator is not None:
+            self._grad_norm_accumulator.reset()
+        for group in self.bf16_groups:
+            _zero_grad_group_helper(group, set_to_none)
+        if not self._foreach_ready:
+            for group in self.fp32_from_bf16_groups:
+                _zero_grad_group_helper(group, set_to_none)
+        for group in self.fp32_from_fp32_groups:
+            _zero_grad_group_helper(group, set_to_none)
+
+    def get_loss_scale(self):
+        return self._scale
+
+    @torch.compiler.disable
+    @torch.no_grad()
+    def _prepare_grads_fallback(self) -> bool:
+        """Original per-parameter gradient copy (fallback path)."""
+        for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                if bf16_param.grad is not None:
+                    fp32_master.grad = bf16_param.grad.float()
+                    bf16_param.grad = None
+        return False
+
+    @torch.compiler.disable
+    def _validate_foreach_cache(self):
+        """Debug assertion: verify cached tensor refs haven't gone stale."""
+        from torch.distributed.tensor import DTensor
+
+        if self._has_fp8_subclass:
+            from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+                WeightWithFP8AllGatherTensor,
+            )
+        for i, p in enumerate(self._bf16_flat):
+            live = p.data
+            if isinstance(live, DTensor):
+                live = live.to_local()
+            if self._has_fp8_subclass and isinstance(live, WeightWithFP8AllGatherTensor):
+                live = live.inner_data()
+            if live.data_ptr() != self._bf16_inners[i].data_ptr():
+                raise RuntimeError(f"Cached tensor ref stale for param {i}")
+        self._validate_count += 1
+
+    @torch.no_grad()
+    def prepare_grads(self) -> bool:
+        """Copy BF16 gradients to FP32 master grads.
+
+        FSDP2 writes gradients to the BF16 model parameter's .grad attribute
+        (after FP32 reduce-scatter followed by cast to orig_dtype). We cast
+        them to FP32 for the optimizer step on master weights.
+
+        Falls back to per-parameter loop if foreach cache is not ready or
+        any grad is missing.
+        """
+        if not self._foreach_ready:
+            any_grad = any(p.grad is not None for group in self.bf16_groups for p in group)
+            if any_grad and self._use_foreach:
+                self._init_foreach_cache()
+            if not self._foreach_ready:
+                self._prepare_grads_fallback()
+                for fp32_group in self.fp32_from_fp32_groups:
+                    for fp32_param in fp32_group:
+                        if hasattr(fp32_param, "main_grad"):
+                            fp32_param.grad = fp32_param.main_grad
+                return False
+
+        from torch.distributed.tensor import DTensor
+
+        bf16_grad_locals = []
+        for p in self._bf16_flat:
+            if p.grad is None:
+                self._prepare_grads_fallback()
+                for fp32_group in self.fp32_from_fp32_groups:
+                    for fp32_param in fp32_group:
+                        if hasattr(fp32_param, "main_grad"):
+                            fp32_param.grad = fp32_param.main_grad
+                return False
+            g = p.grad
+            if isinstance(g, DTensor):
+                g = g.to_local()
+            bf16_grad_locals.append(g)
+
+        if self._validate_count < 3:
+            self._validate_foreach_cache()
+
+        torch._foreach_copy_(self._fp32_grad_locals, bf16_grad_locals)
+
+        for p in self._bf16_flat:
+            p.grad = None
+
+        for fp32_group in self.fp32_from_fp32_groups:
+            for fp32_param in fp32_group:
+                if hasattr(fp32_param, "main_grad"):
+                    fp32_param.grad = fp32_param.main_grad
+
+        return False
+
+    @torch.no_grad()
+    def clip_grad_norm(self, clip_grad: float) -> float | torch.Tensor:
+        """DTensor-native gradient clipping matching TorchTitan.
+
+        Operates on FP32 master grads (from prepare_grads) plus any natively
+        FP32 parameter grads.
+
+        When overlap_grad_norm is enabled, the squared norms have already been
+        accumulated in the RS stream via post_accumulate_grad_hooks on the
+        BF16 model params. The pre-computed norm is correct because
+        prepare_grads copies BF16->FP32 without scaling. The clipping is
+        applied to the FP32 master params.
+        """
+        all_params = list(chain.from_iterable(self.fp32_from_bf16_groups))
+        all_params.extend(chain.from_iterable(self.fp32_from_fp32_groups))
+
+        if self._grad_norm_accumulator is not None:
+            return self._grad_norm_accumulator.finalize(clip_grad, all_params)
+
+        from torch.distributed.tensor import DTensor
+
+        all_grads = []
+        for fp32_group in self.fp32_from_bf16_groups:
+            for p in fp32_group:
+                if p.grad is not None:
+                    all_grads.append(p.grad)
+        for fp32_group in self.fp32_from_fp32_groups:
+            for p in fp32_group:
+                if p.grad is not None:
+                    all_grads.append(p.grad)
+
+        if not all_grads:
+            return 0.0
+
+        total_norm = torch.nn.utils.get_total_norm(all_grads, norm_type=2.0, foreach=True)
+        if isinstance(total_norm, DTensor):
+            total_norm = total_norm.full_tensor()
+
+        torch.nn.utils.clip_grads_with_norm_(all_params, clip_grad, total_norm, foreach=True)
+        return total_norm
+
+    @torch.no_grad()
+    def step_with_ready_grads(self) -> bool:
+        """Step the optimizer on FP32 masters, then copy back to BF16 params."""
+        if self.is_stub_optimizer:
+            return True
+        timers = self.config.timers
+
+        if timers is not None:
+            timers("optimizer-inner-step", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        self.optimizer.step()
+        if timers is not None:
+            timers("optimizer-inner-step").stop()
+
+        if timers is not None:
+            timers("optimizer-copy-main-to-model", log_level=1).start(
+                barrier=self.config.barrier_with_L1_time
+            )
+        self._copy_main_params_to_model_params()
+        if timers is not None:
+            timers("optimizer-copy-main-to-model").stop()
+
+        return True
+
+    def _copy_main_params_to_model_params(self):
+        """Copy FP32 master weights back to BF16 model parameters."""
+        if self._foreach_ready:
+            self._copy_main_to_model_foreach()
+        else:
+            self._copy_main_to_model_fallback()
+
+    def _copy_main_to_model_foreach(self):
+        torch._foreach_copy_(self._bf16_inners, self._fp32_locals)
+
+    @torch.compiler.disable
+    def _copy_main_to_model_fallback(self):
+        for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                bf16_param.data.copy_(fp32_master.data)
+
+    @torch.compiler.disable
+    def _copy_model_params_to_main_params(self):
+        """Copy BF16 model parameters to FP32 masters (for checkpoint reload)."""
+        if self._foreach_ready:
+            torch._foreach_copy_(self._fp32_locals, self._bf16_inners)
+        else:
+            self._copy_model_to_main_fallback()
+
+    @torch.compiler.disable
+    def _copy_model_to_main_fallback(self):
+        for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                fp32_master.data.copy_(bf16_param.data)
+
+    @torch.no_grad()
+    def step(self):
+        """Clip gradients and step. Always succeeds (no overflow for BF16)."""
+        timers = self.config.timers
+
+        found_inf_flag = self.prepare_grads()
+        if found_inf_flag:
+            return False, None, None
+
+        if timers is not None:
+            timers("optimizer-clip-main-grad", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        grad_norm = None
+        if self.config.clip_grad > 0.0:
+            grad_norm = self.clip_grad_norm(self.config.clip_grad)
+        if timers is not None:
+            timers("optimizer-clip-main-grad").stop()
+
+        if timers is not None:
+            timers("optimizer-count-zeros", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None
+        if timers is not None:
+            timers("optimizer-count-zeros").stop()
+
+        success = self.step_with_ready_grads()
+
+        return success, grad_norm, num_zeros_in_grad
+
+    def reload_model_params(self, state_dict=None):
+        """After loading a checkpoint, copy BF16 model params to FP32 masters."""
+        self._copy_model_params_to_main_params()
+
+    def state_dict(self, is_loading: bool = False):
+        if is_loading:
+            self.init_state_fn(self.optimizer, self.config)
+
+        state_dict = {}
+        state_dict["optimizer"] = self.optimizer.state_dict()
+        state_dict["fp32_from_fp16_params"] = self.fp32_from_bf16_groups
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        optimizer_key = "optimizer"
+        if optimizer_key not in state_dict:
+            optimizer_key = "optimizer_state_dict"
+
+        if "common_step" in state_dict[optimizer_key].get("state", {}):
+            common_step = state_dict[optimizer_key]["state"].pop("common_step")
+            self._restore_common_per_param_step(state_dict[optimizer_key], common_step)
+
+        state_dict[optimizer_key]["param_groups"] = self._filter_and_reorder_param_groups(
+            self.optimizer.param_groups, state_dict[optimizer_key]["param_groups"]
+        )
+        self.optimizer.load_state_dict(state_dict[optimizer_key])
+
+        # Restore FP32 master weights
+        if "fp32_from_fp16_params" in state_dict:
+            for current_group, saved_group in zip(
+                self.fp32_from_bf16_groups, state_dict["fp32_from_fp16_params"]
+            ):
+                for current_param, saved_param in zip(current_group, saved_group):
+                    current_param.data.copy_(saved_param.data)
+
+    def sharded_state_dict(
+        self,
+        model_sharded_state_dict: ShardedStateDict,
+        is_loading: bool = False,
+        metadata: Optional[dict] = None,
+    ):
+        if is_loading:
+            self.init_state_fn(self.optimizer, self.config)
+
+        state_dict = self.state_dict()
+
+        # Key on BF16 model params (not FP32 masters) for checkpoint sharding
+        # alignment with the model's sharded state dict
+        id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+            model_sharded_state_dict,
+            chain.from_iterable(g for g in self.bf16_groups),
+        )
+
+        # Convert fp32_from_fp16_params to sharded tensors
+        if len(state_dict["fp32_from_fp16_params"]) != len(state_dict["optimizer"]["param_groups"]):
+            raise ValueError(
+                "state_dict fp32_from_fp16_params length does not match optimizer param_groups length"
+            )
+        state_dict["fp32_from_fp16_params"] = [
+            [
+                make_sharded_optimizer_tensor(
+                    id_to_sharded_param_map[param_id],
+                    fp32_param,
+                    prefix="optimizer.state.fp32_param",
+                )
+                for param_id, fp32_param in zip(state_group["params"], fp32_group)
+            ]
+            for fp32_group, state_group in zip(
+                state_dict["fp32_from_fp16_params"],
+                state_dict["optimizer"]["param_groups"],
+            )
+        ]
+
+        step = self._extract_common_per_param_step(state_dict["optimizer"])
+
+        optim_state_to_sharding_state(state_dict["optimizer"], id_to_sharded_param_map, exclude_keys="step")
+        if step:
+            state_dict["optimizer"]["state"]["common_step"] = step
+        return state_dict
+
+    def finalize_dist_ckpt_load(self, iteration):
+        """Restore optimizer step counter and sync FP32 masters to BF16 params.
+
+        After dist_checkpointing in-place load (skip_load_to_model_and_opt=True),
+        load_state_dict is not called, so we manually set steps and copy masters.
+        """
+        step_val = float(iteration)
+        for fp32_group in self.fp32_from_bf16_groups:
+            for p in fp32_group:
+                if p in self.optimizer.state and "step" in self.optimizer.state[p]:
+                    self.optimizer.state[p]["step"].fill_(step_val)
+        for fp32_group in self.fp32_from_fp32_groups:
+            for p in fp32_group:
+                if p in self.optimizer.state and "step" in self.optimizer.state[p]:
+                    self.optimizer.state[p]["step"].fill_(step_val)
+
+        self._copy_main_params_to_model_params()
+
+
+def get_fsdp2_bf16_master_weight_optimizer(
+    config: OptimizerConfig,
+    model_chunks: List[torch.nn.Module],
+    no_weight_decay_cond: Optional[Callable] = None,
+    scale_lr_cond: Optional[Callable] = None,
+    lr_mult: float = 1.0,
+    use_gloo_process_groups: bool = True,
+    default_skip_embedding_weight_decay: bool = False,
+    pg_collection: Optional["ProcessGroupCollection"] = None,
+    base_optimizer_cls=torch.optim.AdamW,
+    use_foreach: bool = True,
+    **optimizer_kwargs,
+) -> FSDP2BF16MasterWeightOptimizer:
+    """Factory function to create FSDP2 BF16 master weight optimizer.
+
+    Collects trainable parameters, creates FP32 master copies for BF16 params,
+    replaces param_group entries with FP32 masters, builds a fused AdamW, and
+    wraps in FSDP2BF16MasterWeightOptimizer.
+    """
+    all_params = []
+    for model_chunk in model_chunks:
+        for param in model_chunk.parameters():
+            if param.requires_grad:
+                all_params.append(param)
+
+    if not all_params:
+        raise ValueError("No trainable parameters found in model chunks!")
+
+    weight_decay = config.weight_decay
+    lr = config.lr
+    param_groups = []
+
+    if default_skip_embedding_weight_decay and no_weight_decay_cond is None:
+        embedding_params = []
+        non_embedding_params = []
+        for param in all_params:
+            is_embedding = False
+            for model_chunk in model_chunks:
+                for name, p in model_chunk.named_parameters():
+                    if p is param and "embed" in name.lower():
+                        is_embedding = True
+                        break
+                if is_embedding:
+                    break
+            if is_embedding:
+                embedding_params.append(param)
+            else:
+                non_embedding_params.append(param)
+
+        if embedding_params:
+            param_groups.append({"params": embedding_params, "weight_decay": 0.0, "lr": lr})
+        if non_embedding_params:
+            param_groups.append({"params": non_embedding_params, "weight_decay": weight_decay, "lr": lr})
+    elif no_weight_decay_cond is not None:
+        no_wd_params = []
+        wd_params = []
+        for param in all_params:
+            if no_weight_decay_cond(param):
+                no_wd_params.append(param)
+            else:
+                wd_params.append(param)
+        if no_wd_params:
+            param_groups.append({"params": no_wd_params, "weight_decay": 0.0, "lr": lr})
+        if wd_params:
+            param_groups.append({"params": wd_params, "weight_decay": weight_decay, "lr": lr})
+    else:
+        param_groups.append({"params": all_params, "weight_decay": weight_decay, "lr": lr})
+
+    if scale_lr_cond is not None and lr_mult != 1.0:
+        new_groups = []
+        for param_group in param_groups:
+            scaled = [p for p in param_group["params"] if scale_lr_cond(p)]
+            normal = [p for p in param_group["params"] if not scale_lr_cond(p)]
+            if normal:
+                new_groups.append({**param_group, "params": normal})
+            if scaled:
+                new_groups.append({**param_group, "params": scaled, "lr": param_group["lr"] * lr_mult})
+        param_groups = new_groups
+
+    # Create FP32 master copies and replace in param_groups
+    bf16_groups = []
+    fp32_from_bf16_groups = []
+    fp32_from_fp32_groups = []
+
+    for param_group in param_groups:
+        bf16_params_this_group = []
+        fp32_from_bf16_this_group = []
+        fp32_from_fp32_this_group = []
+
+        for i, param in enumerate(param_group["params"]):
+            if param.dtype in (torch.float16, torch.bfloat16):
+                bf16_params_this_group.append(param)
+                main_param = param.detach().clone().float()
+                param.main_param = main_param
+                param_group["params"][i] = main_param
+                fp32_from_bf16_this_group.append(main_param)
+            elif param.dtype == torch.float32:
+                fp32_from_fp32_this_group.append(param)
+            else:
+                _safe_log_rank_0(f"WARNING: unexpected param dtype {param.dtype}, treating as FP32")
+                fp32_from_fp32_this_group.append(param)
+
+        bf16_groups.append(bf16_params_this_group)
+        fp32_from_bf16_groups.append(fp32_from_bf16_this_group)
+        fp32_from_fp32_groups.append(fp32_from_fp32_this_group)
+
+    bf16_count = sum(len(g) for g in bf16_groups)
+    fp32_count = sum(len(g) for g in fp32_from_fp32_groups)
+    master_count = sum(len(g) for g in fp32_from_bf16_groups)
+
+    _safe_log_rank_0(
+        f"Creating FSDP2 BF16 master weight optimizer with "
+        f"{bf16_count + fp32_count:,} parameters in {len(param_groups)} param groups"
+    )
+
+    base_optimizer = base_optimizer_cls(
+        param_groups,
+        betas=(config.adam_beta1, config.adam_beta2),
+        eps=config.adam_eps,
+        fused=True,
+        **optimizer_kwargs,
+    )
+
+    for param_group in base_optimizer.param_groups:
+        param_group.setdefault("wd_mult", 1.0)
+        param_group.setdefault("lr_mult", 1.0)
+        param_group.setdefault("is_expert_parallel", False)
+        param_group.setdefault("is_decoupled_lr", False)
+        param_group.setdefault("default_config", True)
+
+    def init_state_fn(opt, config=None):
+        for group in opt.param_groups:
+            for p in group["params"]:
+                if len(opt.state[p]) == 0:
+                    opt.state[p]["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                    opt.state[p]["exp_avg"] = torch.zeros_like(p.data)
+                    opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data)
+
+    optimizer = FSDP2BF16MasterWeightOptimizer(
+        optimizer=base_optimizer,
+        config=config,
+        init_state_fn=init_state_fn,
+        bf16_groups=bf16_groups,
+        fp32_from_bf16_groups=fp32_from_bf16_groups,
+        fp32_from_fp32_groups=fp32_from_fp32_groups,
+        use_foreach=use_foreach,
+    )
+
+    _safe_log_rank_0("=" * 80)
+    _safe_log_rank_0("[FSDP2BF16MasterWeightOptimizer Initialized]")
+    _safe_log_rank_0("  BF16 model params + FP32 master weights")
+    _safe_log_rank_0("  No FP32->BF16 cast in forward all-gather")
+    _safe_log_rank_0("  FP32 optimizer step + copy-back to BF16 per iteration")
+    _safe_log_rank_0(f"  BF16 parameters: {bf16_count:,} ({master_count:,} FP32 master copies)")
+    _safe_log_rank_0(f"  FP32 parameters: {fp32_count:,} (no master copy needed)")
+    _safe_log_rank_0(f"  Total trainable parameters: {bf16_count + fp32_count:,}")
+    _safe_log_rank_0("=" * 80)
+
+    optimizer.warmup_foreach_cache()
+
+    return optimizer
diff --git a/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py
new file mode 100644
index 000000000..04766ac3d
--- /dev/null
+++ b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py
@@ -0,0 +1,368 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+FP32 optimizer for FSDP2 mixed precision training.
+
+Uses the TorchTitan approach: model parameters are initialized in FP32,
+FSDP2's MixedPrecisionPolicy casts to BF16 for forward/backward, and
+the optimizer operates on FP32 parameters with FP32 states.
+
+This eliminates the "stale weights" problem of BF16 optimizer states
+while avoiding the master-copy duplication of Float16OptimizerWithFloat16Params.
+
+Memory impact vs BF16 optimizer (Flux 12B, 8 GPUs):
+    +3 GB/GPU for FP32 parameters (vs BF16)
+    +6 GB/GPU for FP32 optimizer states (vs BF16)
+    = +9 GB/GPU total
+
+Gradient clipping uses PyTorch-native DTensor-aware APIs matching TorchTitan.
+"""
+
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.optimizer import (
+    get_param_id_to_sharded_param_map,
+    optim_state_to_sharding_state,
+)
+from megatron.core.optimizer.optimizer import MegatronOptimizer
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+
+from primus.core.utils.module_utils import log_rank_0
+
+if TYPE_CHECKING:
+    from megatron.core.process_groups_config import ProcessGroupCollection
+
+
+def _safe_log_rank_0(msg: str):
+    try:
+        log_rank_0(msg)
+    except (AttributeError, TypeError):
+        import torch.distributed as dist
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(msg)
+
+
+class FSDP2FP32Optimizer(MegatronOptimizer):
+    """FP32 optimizer for FSDP2 mixed precision training.
+
+    Extends MegatronOptimizer directly (not MixedPrecisionOptimizer).
+    Modeled on Megatron's own FP32Optimizer with two key differences:
+
+    1. prepare_grads is a no-op: FSDP2 writes gradients directly to param.grad
+       (no main_grad -> grad copy needed).
+    2. clip_grad_norm uses TorchTitan-style DTensor-native APIs:
+       torch.nn.utils.get_total_norm + clip_grads_with_norm_ for correct
+       norm computation across FSDP2's sharded DTensor parameters.
+
+    Args:
+        optimizer: Base PyTorch optimizer (e.g., AdamW with fused=True).
+        config: OptimizerConfig from Megatron.
+        init_state_fn: Function to initialize optimizer state tensors.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        config: OptimizerConfig,
+        init_state_fn: Callable,
+    ):
+        super().__init__(optimizer, config, init_state_fn)
+        self._scale = torch.tensor([1.0], dtype=torch.float, device="cuda")
+        self.is_stub_optimizer = optimizer is None
+
+        from megatron.training import get_args
+
+        args = get_args()
+        self._grad_norm_accumulator = getattr(args, "_grad_norm_accumulator", None)
+
+    def zero_grad(self, set_to_none=True):
+        if self.is_stub_optimizer:
+            return
+        if self._grad_norm_accumulator is not None:
+            self._grad_norm_accumulator.reset()
+        self.optimizer.zero_grad(set_to_none=set_to_none)
+
+    def get_loss_scale(self):
+        return self._scale
+
+    @torch.no_grad()
+    def prepare_grads(self) -> bool:
+        """No-op: FSDP2 writes gradients directly to param.grad."""
+        return False
+
+    @torch.no_grad()
+    def clip_grad_norm(self, clip_grad: float) -> float | torch.Tensor:
+        """DTensor-native gradient clipping matching TorchTitan.
+
+        Uses torch.nn.utils.get_total_norm which natively handles DTensor
+        gradients (returns a DTensor with _NormPartial placement that is
+        reduced via full_tensor()), then clips with foreach-optimized
+        clip_grads_with_norm_.
+
+        When overlap_grad_norm is enabled, the squared norms have already been
+        accumulated in the RS stream via post_accumulate_grad_hooks. Only a
+        single all-reduce + sqrt + clip is needed.
+        """
+        params = self.get_parameters()
+
+        if self._grad_norm_accumulator is not None:
+            return self._grad_norm_accumulator.finalize(clip_grad, params)
+
+        from torch.distributed.tensor import DTensor
+
+        grads = [p.grad for p in params if p.grad is not None]
+
+        if not grads:
+            return 0.0
+
+        total_norm = torch.nn.utils.get_total_norm(grads, norm_type=2.0, foreach=True)
+        if isinstance(total_norm, DTensor):
+            total_norm = total_norm.full_tensor()
+        torch.nn.utils.clip_grads_with_norm_(params, clip_grad, total_norm, foreach=True)
+        return total_norm
+
+    @torch.no_grad()
+    def step_with_ready_grads(self) -> bool:
+        if self.is_stub_optimizer:
+            return True
+        timers = self.config.timers
+
+        if timers is not None:
+            timers("optimizer-inner-step", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        self.optimizer.step()
+        if timers is not None:
+            timers("optimizer-inner-step").stop()
+
+        return True
+
+    @torch.no_grad()
+    def step(self):
+        """Clip gradients and step. Always succeeds (no overflow for FP32)."""
+        timers = self.config.timers
+
+        found_inf_flag = self.prepare_grads()
+        if found_inf_flag:
+            return False, None, None
+
+        if timers is not None:
+            timers("optimizer-clip-main-grad", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        grad_norm = None
+        if self.config.clip_grad > 0.0:
+            grad_norm = self.clip_grad_norm(self.config.clip_grad)
+        if timers is not None:
+            timers("optimizer-clip-main-grad").stop()
+
+        if timers is not None:
+            timers("optimizer-count-zeros", log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None
+        if timers is not None:
+            timers("optimizer-count-zeros").stop()
+
+        success = self.step_with_ready_grads()
+
+        return success, grad_norm, num_zeros_in_grad
+
+    def reload_model_params(self, state_dict=None):
+        pass
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        if "common_step" in state_dict.get("state", {}):
+            common_step = state_dict["state"].pop("common_step")
+            self._restore_common_per_param_step(state_dict, common_step)
+
+        state_dict["param_groups"] = self._filter_and_reorder_param_groups(
+            self.optimizer.param_groups, state_dict["param_groups"]
+        )
+        self.optimizer.load_state_dict(state_dict)
+
+    def sharded_state_dict(
+        self,
+        model_sharded_state_dict: ShardedStateDict,
+        is_loading: bool = False,
+        metadata: Optional[dict] = None,
+    ):
+        if is_loading:
+            self.init_state_fn(self.optimizer, self.config)
+
+        state_dict = self.state_dict()
+        id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+            model_sharded_state_dict, self.get_parameters()
+        )
+        step = self._extract_common_per_param_step(state_dict)
+
+        optim_state_to_sharding_state(state_dict, id_to_sharded_param_map, exclude_keys="step")
+        if step:
+            state_dict["state"]["common_step"] = step
+        return state_dict
+
+    def finalize_dist_ckpt_load(self, iteration):
+        """Restore optimizer step counter after dist_checkpointing in-place load.
+
+        When skip_load_to_model_and_opt=True (FSDP2), load_state_dict is not
+        called, so common_step is never fanned out to per-parameter step
+        entries.  This fills step from the training iteration.
+        """
+        step_val = float(iteration)
+        for p in self.get_parameters():
+            if p in self.optimizer.state and "step" in self.optimizer.state[p]:
+                self.optimizer.state[p]["step"].fill_(step_val)
+
+
+def get_fsdp2_fp32_optimizer(
+    config: OptimizerConfig,
+    model_chunks: List[torch.nn.Module],
+    no_weight_decay_cond: Optional[Callable] = None,
+    scale_lr_cond: Optional[Callable] = None,
+    lr_mult: float = 1.0,
+    use_gloo_process_groups: bool = True,
+    default_skip_embedding_weight_decay: bool = False,
+    pg_collection: Optional["ProcessGroupCollection"] = None,
+    base_optimizer_cls=torch.optim.AdamW,
+    use_foreach: bool = False,
+    **optimizer_kwargs,
+) -> FSDP2FP32Optimizer:
+    """Factory function to create FSDP2 FP32 param optimizer from model chunks.
+
+    Collects trainable FP32 parameters, builds param groups with weight decay
+    and LR scaling, creates AdamW (fused or foreach), and wraps in
+    FSDP2FP32Optimizer.
+
+    Args:
+        config: OptimizerConfig from Megatron.
+        model_chunks: List of model modules (FSDP2-wrapped).
+        no_weight_decay_cond: Optional predicate for zero weight decay.
+        scale_lr_cond: Optional predicate for scaled learning rate.
+        lr_mult: Learning rate multiplier for scaled params.
+        use_gloo_process_groups: Unused (kept for API compatibility).
+        default_skip_embedding_weight_decay: Skip weight decay for embeddings
+            if no_weight_decay_cond not provided.
+        pg_collection: Unused (kept for API compatibility).
+        base_optimizer_cls: PyTorch optimizer class (default: AdamW).
+        use_foreach: If True, use foreach mode; if False (default), use fused.
+        **optimizer_kwargs: Additional kwargs for base optimizer.
+
+    Returns:
+        FSDP2FP32Optimizer instance.
+    """
+    all_params = []
+    for model_chunk in model_chunks:
+        for param in model_chunk.parameters():
+            if param.requires_grad:
+                all_params.append(param)
+
+    if not all_params:
+        raise ValueError("No trainable parameters found in model chunks!")
+
+    weight_decay = config.weight_decay
+    lr = config.lr
+    param_groups = []
+
+    if default_skip_embedding_weight_decay and no_weight_decay_cond is None:
+        embedding_params = []
+        non_embedding_params = []
+        for param in all_params:
+            is_embedding = False
+            for model_chunk in model_chunks:
+                for name, p in model_chunk.named_parameters():
+                    if p is param and "embed" in name.lower():
+                        is_embedding = True
+                        break
+                if is_embedding:
+                    break
+            if is_embedding:
+                embedding_params.append(param)
+            else:
+                non_embedding_params.append(param)
+
+        if embedding_params:
+            param_groups.append({"params": embedding_params, "weight_decay": 0.0, "lr": lr})
+        if non_embedding_params:
+            param_groups.append({"params": non_embedding_params, "weight_decay": weight_decay, "lr": lr})
+    elif no_weight_decay_cond is not None:
+        no_wd_params = []
+        wd_params = []
+        for param in all_params:
+            if no_weight_decay_cond(param):
+                no_wd_params.append(param)
+            else:
+                wd_params.append(param)
+        if no_wd_params:
+            param_groups.append({"params": no_wd_params, "weight_decay": 0.0, "lr": lr})
+        if wd_params:
+            param_groups.append({"params": wd_params, "weight_decay": weight_decay, "lr": lr})
+    else:
+        param_groups.append({"params": all_params, "weight_decay": weight_decay, "lr": lr})
+
+    if scale_lr_cond is not None and lr_mult != 1.0:
+        new_groups = []
+        for param_group in param_groups:
+            scaled = [p for p in param_group["params"] if scale_lr_cond(p)]
+            normal = [p for p in param_group["params"] if not scale_lr_cond(p)]
+            if normal:
+                new_groups.append({**param_group, "params": normal})
+            if scaled:
+                new_groups.append({**param_group, "params": scaled, "lr": param_group["lr"] * lr_mult})
+        param_groups = new_groups
+
+    _safe_log_rank_0(
+        f"Creating FSDP2 FP32 optimizer with {len(all_params):,} parameters "
+        f"in {len(param_groups)} param groups"
+    )
+
+    base_optimizer = base_optimizer_cls(
+        param_groups,
+        betas=(config.adam_beta1, config.adam_beta2),
+        eps=config.adam_eps,
+        fused=not use_foreach,
+        foreach=use_foreach,
+        **optimizer_kwargs,
+    )
+
+    for param_group in base_optimizer.param_groups:
+        param_group.setdefault("wd_mult", 1.0)
+        param_group.setdefault("lr_mult", 1.0)
+        param_group.setdefault("is_expert_parallel", False)
+        param_group.setdefault("is_decoupled_lr", False)
+        param_group.setdefault("default_config", True)
+
+    def init_state_fn(opt, config=None):
+        for group in opt.param_groups:
+            for p in group["params"]:
+                if len(opt.state[p]) == 0:
+                    opt.state[p]["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                    opt.state[p]["exp_avg"] = torch.zeros_like(p.data)
+                    opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data)
+
+    optimizer = FSDP2FP32Optimizer(
+        optimizer=base_optimizer,
+        config=config,
+        init_state_fn=init_state_fn,
+    )
+
+    fp32_count = sum(1 for p in all_params if p.dtype == torch.float32)
+    bf16_count = sum(1 for p in all_params if p.dtype == torch.bfloat16)
+    other_count = len(all_params) - fp32_count - bf16_count
+
+    _safe_log_rank_0("=" * 80)
+    _safe_log_rank_0("[FSDP2FP32ParamOptimizer Initialized]")
+    _safe_log_rank_0("  TorchTitan-style: FP32 params + FSDP2 MixedPrecisionPolicy")
+    _safe_log_rank_0("  DTensor-native gradient clipping")
+    _safe_log_rank_0(f"  AdamW mode: {'foreach' if use_foreach else 'fused'}")
+    _safe_log_rank_0(f"  FP32 parameters: {fp32_count:,}")
+    _safe_log_rank_0(f"  BF16 parameters: {bf16_count:,}")
+    if other_count > 0:
+        _safe_log_rank_0(f"  Other dtype parameters: {other_count:,}")
+    _safe_log_rank_0(f"  Total trainable parameters: {len(all_params):,}")
+    _safe_log_rank_0("=" * 80)
+
+    return optimizer
diff --git a/primus/backends/megatron/core/optimizer/incremental_grad_norm.py b/primus/backends/megatron/core/optimizer/incremental_grad_norm.py
new file mode 100644
index 000000000..9320c3244
--- /dev/null
+++ b/primus/backends/megatron/core/optimizer/incremental_grad_norm.py
@@ -0,0 +1,90 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Incremental gradient norm accumulator for overlapping grad norm with FSDP2
+reduce-scatter.
+
+Registers a `register_post_accumulate_grad_hook` on each sharded parameter.
+FSDP2 fires these hooks inside the post-reduce stream (RS or AR stream)
+after the sharded gradient is assigned. Each hook accumulates ||grad||_2^2
+into a shared tensor. By the time the default stream synchronises
+(`_wait_for_post_backward`), the total squared norm is already computed.
+The optimizer's `clip_grad_norm` then only needs a single all-reduce + sqrt
++ clip — eliminating the per-parameter norm recomputation that would
+otherwise serialise after the last reduce-scatter.
+"""
+
+from typing import Callable
+
+import torch
+import torch.distributed as dist
+
+
+class IncrementalGradNormAccumulator:
+    """Accumulates squared L2 norms of sharded gradients in the RS stream.
+
+    Args:
+        shard_process_group: The FSDP shard process group used for the final
+            all-reduce of the accumulated squared norm. For HSDP this must be
+            the shard-only group (not the full DP group) because replicate-group
+            ranks hold identical gradients after the HSDP all-reduce.
+        device: CUDA device for the accumulator tensor.
+    """
+
+    def __init__(
+        self,
+        shard_process_group: dist.ProcessGroup,
+        device: torch.device,
+    ) -> None:
+        self._shard_pg = shard_process_group
+        self._total_norm_sq = torch.zeros((), dtype=torch.float32, device=device)
+
+    def reset(self) -> None:
+        """Zero the accumulator. Called from optimizer.zero_grad()."""
+        self._total_norm_sq.zero_()
+
+    def make_hook(self) -> Callable[[torch.Tensor], None]:
+        """Return a closure suitable for register_post_accumulate_grad_hook.
+
+        The hook extracts the local shard data (bypassing DTensor overhead)
+        and accumulates its squared L2 norm into ``_total_norm_sq``.
+        """
+        acc = self
+
+        def hook(param: torch.Tensor) -> None:
+            grad = param.grad
+            if grad is None:
+                return
+            local_grad = grad._local_tensor if hasattr(grad, "_local_tensor") else grad
+            acc._total_norm_sq.add_(local_grad.float().norm(2).pow_(2))
+
+        return hook
+
+    @torch.no_grad()
+    def finalize(
+        self,
+        clip_grad: float,
+        params,
+    ) -> torch.Tensor:
+        """All-reduce the accumulated norm, sqrt, clip, and return total_norm.
+
+        No additional stream synchronisation is needed here: by the time the
+        optimizer calls this method the default stream has already waited on
+        ``post_reduce_event`` via ``_wait_for_post_backward``, so the
+        accumulated value is visible.
+
+        Args:
+            clip_grad: Maximum gradient norm for clipping.
+            params: Iterable of parameters whose ``.grad`` will be clipped.
+
+        Returns:
+            The global L2 gradient norm (scalar tensor).
+        """
+        dist.all_reduce(self._total_norm_sq, op=dist.ReduceOp.SUM, group=self._shard_pg)
+        total_norm = self._total_norm_sq.sqrt_()
+        torch.nn.utils.clip_grads_with_norm_(params, clip_grad, total_norm, foreach=True)
+        return total_norm
diff --git a/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py b/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py
deleted file mode 100644
index cc3ea94c8..000000000
--- a/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py
+++ /dev/null
@@ -1,307 +0,0 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-import logging
-from typing import Callable, List, Optional
-
-import torch
-from megatron.core.dist_checkpointing.dict_utils import nested_values
-from megatron.core.dist_checkpointing.mapping import (
-    LocalNonpersistentObject,
-    ShardedStateDict,
-)
-from megatron.core.optimizer.clip_grads import count_zeros_fp32, get_grad_norm_fp32
-from megatron.core.optimizer.optimizer import (
-    ChainedOptimizer,
-    Float16OptimizerWithFloat16Params,
-    FP32Optimizer,
-    MegatronOptimizer,
-)
-from megatron.core.optimizer.optimizer_config import OptimizerConfig
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.utils import get_pg_rank, get_pg_size
-from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
-
-logger = logging.getLogger(__name__)
-
-
-class LayerWiseDistributedOptimizer(ChainedOptimizer):
-    """Layer-wise distributed optimizer for Megatron-core models.
-
-    Experimental distributed optimizer wrapper that distributes weight to DP ranks by layer.
-    Implemented as ChainedOptimizer to support multiple optimizers (e.g. muon + adamW)
-    When using, keep all megatron distributed-optimizer related options OFF.
-
-    How LayerWiseDistributedOptimizer work:
-    1. weights are splited into lists and each rank only keep its shard in its optimizer
-    2. Megatron DDP handle allreduce grad, note that each rank have full model and grad
-    3. optimizer is already modified so only param belong to this DP rank is updated
-    4. grad_norm and zero counting will reduce metrics globally in step function
-    5. Do regular update with chained optimizers, modified optimizer only update shard
-    6. allgather updated params to every rank
-    """
-
-    def __init__(
-        self,
-        optimizers: List[MegatronOptimizer],
-        config: OptimizerConfig,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-        init_state_fn_list: Optional[List[Callable]] = None,
-    ) -> None:
-        """
-        Initialize LayerWiseDistributedOptimizer.
-
-        Args:
-            optimizers: List of MegatronOptimizers.
-            config: OptimizerConfig.
-            pg_collection: ProcessGroupCollection.
-            init_state_fn_list: List of init state functions.
-        """
-
-        self.pg_collection = pg_collection
-        self.shard_params(optimizers)
-        if init_state_fn_list:
-            assert len(init_state_fn_list) == len(
-                optimizers
-            ), "init_state_fn_list must be the same length as optimizers if provided"
-
-        # wrap optimizer after sharding to avoid unnecessary master weight creation
-        # for higher precision, optimizers are wrapped with megatron already
-        if config.bf16:
-            # unwrap FP32 optimizer, possibly from reusing get_megatron_optimizer for adam
-            for i in range(len(optimizers)):
-                opt = optimizers[i]
-                if isinstance(opt, Float16OptimizerWithFloat16Params):
-                    raise TypeError("LayerWiseDistributedOptimizer received Float16 optimizer already.")
-                # unwrap FP32 optimizer from reusing get_megatron_optimizer for adam
-                if isinstance(opt, FP32Optimizer):
-                    opt = opt.optimizer
-                optimizers[i] = Float16OptimizerWithFloat16Params(
-                    opt, config, None, init_state_fn_list[i] if init_state_fn_list else None
-                )
-
-        super().__init__(optimizers)
-
-        # TODO(kunlun, deyuf): potential future perf optimization
-        # since allreduce is unchanged and handled by megatron DDP, they're already in
-        # contiguous gbuf. So instead of shard param by layer randomly, we can shard by
-        # buf range but keep some "extras" to keep boundary weight not sharded.
-        # This way each rank do some duplicated work but allgather_v is no longer needed
-        # All current distopt optimization can also be potentially applied
-
-    def shard_params(self, optimizers):
-        """Shard all params into lists by rank."""
-        # list of parameter are sorted by numel and assigned to ranks in ping-pong style
-        # example of 4 ranks and 10 parameters p0-p9 after sorting, then dp_cp_params_list will be
-        # [[p0, p7, p8], [p1, p6, p9], [p2, p5], [p3, p4]]
-
-        # simplify when dp_cp group size is 1
-        if get_pg_size(self.pg_collection.dp_cp) == 1:
-            self.dp_cp_params_list = None
-            self.expt_dp_params_list = None
-            return
-
-        dp_cp_idx, expt_dp_idx = 0, 0
-        dp_cp_size = get_pg_size(self.pg_collection.dp_cp)
-        expt_dp_size = get_pg_size(self.pg_collection.expt_dp)
-        # create ping-pong style loop so memory is more balanced
-        dp_cp_loop = list(range(dp_cp_size)) + list(range(dp_cp_size))[::-1]
-        expt_dp_loop = list(range(expt_dp_size)) + list(range(expt_dp_size))[::-1]
-        self.dp_cp_params_list = [[] for _ in range(dp_cp_size)]
-        self.expt_dp_params_list = [[] for _ in range(expt_dp_size)]
-        # get all param groups
-        param_groups = []
-        for optimizer in optimizers:
-            param_groups += optimizer.param_groups
-
-        # sort param in all groups by param numel and assign to each rank evenly
-        param_list = []
-        for group_index, group in enumerate(param_groups):
-            for p in group["params"]:
-                param_list.append((p, group_index))
-        param_list.sort(key=lambda x: x[0].numel())
-        param_groups_this_rank = [[] for g in param_groups]
-
-        # assign params to rank in ping-pong style loop
-        for p, group_index in param_list:
-            if param_groups[group_index].get("is_expert_parallel", False):
-                if expt_dp_loop[expt_dp_idx] == get_pg_rank(self.pg_collection.expt_dp):
-                    param_groups_this_rank[group_index].append(p)
-                self.expt_dp_params_list[expt_dp_loop[expt_dp_idx]].append(p)
-                expt_dp_idx = (expt_dp_idx + 1) % len(expt_dp_loop)
-            else:
-                if dp_cp_loop[dp_cp_idx] == get_pg_rank(self.pg_collection.dp_cp):
-                    param_groups_this_rank[group_index].append(p)
-                self.dp_cp_params_list[dp_cp_loop[dp_cp_idx]].append(p)
-                dp_cp_idx = (dp_cp_idx + 1) % len(dp_cp_loop)
-
-        # now we modify the group to only handle local params
-        for groups, params in zip(param_groups, param_groups_this_rank):
-            groups["params"] = params
-
-        # simplify when expt_dp group size is 1 or expert parallel is off
-        if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0:
-            self.expt_dp_params_list = None
-
-    @torch.no_grad()
-    def allgather_params(self) -> None:
-        """All-gather updated params from all ranks."""
-
-        # helper function to flatten local params, allgather, unflatten and copy to model params
-        def _allgather_helper(params_list, group):
-            # flatten this rank's params and create empty tensor output list
-            device = params_list[0][0].device
-            dtype = params_list[0][0].dtype
-            rank = get_pg_rank(group)
-            # for rank without params create empty tensor and participate in allgather
-            src = (
-                _flatten_dense_tensors(params_list[rank])
-                if len(params_list[rank]) > 0
-                else torch.empty(0, device=device, dtype=dtype)
-            )
-            output_list = [
-                torch.empty(sum([p.numel() for p in params]), device=device, dtype=dtype)
-                for params in params_list
-            ]
-            # single all_gather_v to collect all updated params
-            torch.distributed.all_gather(output_list, src, group=group)
-            # unflatten and copy gathered params for each rank i
-            for idx, (flat_params, params) in enumerate(zip(output_list, params_list)):
-                # skip local params and empty tensors
-                if len(params) == 0 or idx == rank:
-                    continue
-                updated_params = _unflatten_dense_tensors(flat_params, params)
-                for updated_p, model_p in zip(updated_params, params):
-                    model_p.data.copy_(updated_p)
-
-        if self.pg_collection is None:
-            return
-        if self.dp_cp_params_list:
-            _allgather_helper(self.dp_cp_params_list, self.pg_collection.dp_cp)
-        if self.expt_dp_params_list:
-            _allgather_helper(self.expt_dp_params_list, self.pg_collection.expt_dp)
-
-    @torch.no_grad()
-    def broadcast_params(self):
-        """All rank broadcast updated local params."""
-        # Broadcast linear layer weights to all other ranks. Kept as reference test.
-        if self.dp_cp_params_list is None:
-            return
-        for i, params in enumerate(self.dp_cp_params_list):
-            src_global_rank = torch.distributed.get_global_rank(self.pg_collection.dp_cp, i)
-            for p in params:
-                torch.distributed.broadcast(p, src_global_rank, self.pg_collection.dp_cp)
-        if self.expt_dp_params_list is None:
-            return
-        for i, params in enumerate(self.expt_dp_params_list):
-            src_global_rank = torch.distributed.get_global_rank(self.pg_collection.expt_dp, i)
-            for p in params:
-                torch.distributed.broadcast(p, src_global_rank, self.pg_collection.expt_dp)
-
-    @torch.no_grad()
-    def get_grad_norm(self):
-        # similar to dist opt, always aggregate globally
-        grads_for_norm = []
-        for optimizer in self.chained_optimizers:
-            grads_for_norm += optimizer.get_main_grads_for_grad_norm()
-        grad_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None)
-        return grad_norm
-
-    @torch.no_grad()
-    def count_zeros(self):
-        params = []
-        for optimizer in self.chained_optimizers:
-            params += optimizer.get_parameters()
-        return count_zeros_fp32(
-            params,
-            grad_stats_parallel_group=None,
-            use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8,
-        )
-
-    @torch.no_grad()
-    def step(self):  # type: ignore[no-untyped-def]
-        """step function for layer-wise optimizer."""
-        update_successful, grad_norm, num_zeros_in_grad = super().step()
-
-        # All gather updated params.
-        self.allgather_params()
-
-        return update_successful, grad_norm, num_zeros_in_grad
-
-    # TODO(deyuf): need to improve dist checkpointing design to properly handle this
-    # fp32_from_fp16_params is list, each sub list could be empty if group is empty
-    # this breaks dist checkpointing assumption since extract_sharded_base drop list structure
-    # for now, we convert it to dict with index as key and convert back in load_state_dict
-    def load_state_dict(self, state_dict):
-        if len(self.chained_optimizers) == 1:
-            wrapped_state_dict = {1: state_dict}
-        else:
-            wrapped_state_dict = state_dict
-        for sd in wrapped_state_dict.values():
-            if "fp32_from_fp16_params" in sd and isinstance(sd["fp32_from_fp16_params"], dict):
-                logger.info("[layerwise] converting fp32_from_fp16_params from dict to list")
-                sd["fp32_from_fp16_params"] = [v for k, v in sorted(sd["fp32_from_fp16_params"].items())]
-        super().load_state_dict(state_dict)
-
-    def sharded_state_dict(
-        self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False, **kwargs
-    ):
-        """
-        Sharded state dict for torch_dist format checkpointing.
-        For fixed DP usage only, set replica_id to 0 for all ShardedTensor.
-        """
-        sharded_state_dict = super().sharded_state_dict(model_sharded_state_dict, is_loading, **kwargs)
-
-        # for fixed DP usage only
-        for sh_base in nested_values(sharded_state_dict):
-            if hasattr(sh_base, "replica_id"):
-                assert (
-                    isinstance(sh_base.replica_id, int) or len(sh_base.replica_id) == 3
-                ), f"Expected replica_id as int or (PP, TP, DP), got: {sh_base}"
-                sh_base.replica_id = (
-                    0 if isinstance(sh_base.replica_id, int) else (*sh_base.replica_id[:2], 0)
-                )
-
-        # later code assume list but chained optimizer fallback to non-list if there's only one
-        if len(self.chained_optimizers) == 1:
-            wrapped_sharded_state_dict = {1: sharded_state_dict}
-        else:
-            wrapped_sharded_state_dict = sharded_state_dict
-
-        # Adjust dict rank 0 output correct global metadata into common_dict
-        for sd in wrapped_sharded_state_dict.values():
-            # wrap empty containers into LocalNonpersistentObject so it won't be saved/loaded
-            # params is already wrapped, we only need to handle fp32_from_fp16_params and state
-            # more details in load_state_dict comment
-            if "fp32_from_fp16_params" in sd:
-                sd["fp32_from_fp16_params"][:] = [
-                    group if group else LocalNonpersistentObject(group)
-                    for group in sd["fp32_from_fp16_params"]
-                ]
-                sd["fp32_from_fp16_params"] = {i: v for i, v in enumerate(sd["fp32_from_fp16_params"])}
-            # state is a single dict and will be empty if optimizer is fully empty
-            if not sd["optimizer"]["state"]:
-                sd["optimizer"]["state"] = LocalNonpersistentObject(sd["optimizer"]["state"])
-            # group keys(e.g. 'step') might be missing or not updated
-            for i, group in enumerate(sd["optimizer"]["param_groups"]):
-                # keep local param tensor so we only gather metadata
-                local_params = group.pop("params")
-                # save whether this group is empty, so we can use non-empty rank for metadata
-                group["params"] = bool(local_params.unwrap())
-                all_rank_groups = [None for _ in range(torch.distributed.get_world_size())]
-                torch.distributed.all_gather_object(all_rank_groups, group)
-                # find first non-empty group if it exists
-                nonempty_rank_group = next((g for g in all_rank_groups if g["params"]), group)
-                nonempty_rank_group["params"] = local_params
-                sd["optimizer"]["param_groups"][i] = nonempty_rank_group
-        return sharded_state_dict
-
-    def save_state_dict_to_file(self, filename: str) -> None:
-        """Save the parameter state of the optimizer. For torch format only.
-        Args:
-            filename: The filename to save the parameter state.
-        """
-        torch.save(super().state_dict(), filename)
-
-    def load_state_dict_from_file(self, filename: str) -> None:
-        """Load the parameter state of the optimizer. For torch format only."""
-        super().load_state_dict(torch.load(filename))
diff --git a/primus/backends/megatron/core/optimizer/moun.py b/primus/backends/megatron/core/optimizer/moun.py
deleted file mode 100644
index 5464d6f00..000000000
--- a/primus/backends/megatron/core/optimizer/moun.py
+++ /dev/null
@@ -1,353 +0,0 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-"""Megatron muon optimizer wrapper to handle tensor-parallel."""
-
-import logging
-from typing import Any, Callable, Dict, List, Literal, Optional
-
-import torch
-from megatron.core import parallel_state
-from megatron.core.optimizer import _get_param_groups, get_megatron_optimizer
-from megatron.core.optimizer.optimizer import (
-    ChainedOptimizer,
-    Float16OptimizerWithFloat16Params,
-    FP32Optimizer,
-    MegatronOptimizer,
-)
-from megatron.core.optimizer.optimizer_config import OptimizerConfig
-
-try:
-    from megatron.core.optimizer.optimizer_config import ParamKey
-except ImportError:
-    ParamKey = Any
-
-try:
-    from megatron.core.optimizer_param_scheduler import ParamGroupOverride
-except ImportError:
-    ParamGroupOverride = Any
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import get_pg_size, log_single_rank
-from torch.optim.optimizer import ParamsT
-
-from .layer_wise_optimizer import LayerWiseDistributedOptimizer
-
-try:
-    from emerging_optimizers.orthogonalized_optimizers import (
-        OrthogonalizedOptimizer,
-        get_muon_scale_factor,
-    )
-    from emerging_optimizers.orthogonalized_optimizers.muon_utils import (
-        newton_schulz_tp,
-    )
-
-    HAVE_EMERGING_OPTIMIZERS = True
-
-except ImportError:
-    HAVE_EMERGING_OPTIMIZERS = False
-    OrthogonalizedOptimizer = object
-
-
-logger = logging.getLogger(__name__)
-
-
-class TensorParallelMuon(OrthogonalizedOptimizer):
-    """Tensor Parallel Muon optimizer."""
-
-    def __init__(
-        self,
-        params: ParamsT,
-        lr: float = 3e-4,
-        momentum_beta: float = 0.95,
-        use_nesterov: bool = True,
-        weight_decay: float = 0.01,
-        use_decoupled_weight_decay: bool = True,
-        split_qkv: bool = False,
-        is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
-        qkv_split_shapes: tuple[int, int, int] | None = None,
-        fp32_matmul_prec: str = "medium",
-        coefficient_type: str = "quintic",
-        num_ns_steps: int = 5,
-        scale_mode: str = "spectral",
-        extra_scale_factor: float = 1.0,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-        mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated",
-    ) -> None:
-        if num_ns_steps < 1:
-            raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
-
-        def scaled_orthogonalize_fn(
-            grad: torch.Tensor,
-            tp_group: torch.distributed.ProcessGroup,
-            partition_dim: int | None = None,
-        ) -> torch.Tensor:
-            # log_single_rank(
-            #     logger,
-            #     logging.DEBUG,
-            #     f"Orthogonalizing grad with {num_ns_steps} steps, {coefficient_type} coefficient, "
-            #     f"{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}",
-            # )
-            size = [grad.size(-2), grad.size(-1)]
-            if partition_dim:
-                size[partition_dim] *= get_pg_size(tp_group)
-            orth_grad = newton_schulz_tp(
-                grad,
-                steps=num_ns_steps,
-                coefficient_type=coefficient_type,
-                tp_group=tp_group,
-                partition_dim=partition_dim,
-                mode="duplicated" if mode == "blockwise" else mode,
-            )
-            scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)
-            return orth_grad * scale_factor * extra_scale_factor
-
-        self.pg_collection = pg_collection
-        self.mode = mode
-        self.split_qkv = split_qkv
-        self.is_qkv_fn = is_qkv_fn
-        self.qkv_split_shapes = qkv_split_shapes
-
-        weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
-        super().__init__(
-            params,
-            lr,
-            momentum_beta,
-            use_nesterov=use_nesterov,
-            weight_decay=weight_decay,
-            weight_decay_method=weight_decay_method,
-            fp32_matmul_prec=fp32_matmul_prec,
-            scaled_orthogonalize_fn=scaled_orthogonalize_fn,
-        )
-
-    def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor:
-        """Orthogonalize the momentum.
-
-        Args:
-            p: The parameter tensor. i is necessary to pass param tensor in addition to momentum
-                because a lot of information is only available in the param tensor,
-                attributes for example.
-            grad: The momentum tensor.
-
-        Returns:
-            The orthogonalized gradient tensor.
-        """
-        # TODO(deyuf): switch to group
-        if self.pg_collection:
-            tp_group = self.pg_collection.expt_tp if getattr(p, "expert_tp", False) else self.pg_collection.tp
-        else:
-            tp_group = None
-        partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None)
-        if partition_dim == -1:
-            # llm-shower use different default value for partition_dim than TE.
-            # Because -1 is a valid index for ndarray, we decided to not overload it.
-            partition_dim = None
-
-        if self.split_qkv and self.is_qkv_fn(p):  # type: ignore[misc]
-            # split grouped attention parameters (e.g., QKV, GQA, etc.)
-            grad_shape = grad.shape
-            log_single_rank(
-                logger,
-                logging.DEBUG,
-                f"qkv split grad shape {grad_shape}, split shapes {self.qkv_split_shapes}",
-            )
-            num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes)
-            qkv_grads = torch.split(
-                grad.view(num_query_groups, sum(self.qkv_split_shapes), -1),
-                self.qkv_split_shapes,
-                dim=1,
-            )
-            qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads]
-
-            # Apply Newton-Schulz and scales to each component, concat back
-            qkv_grads = [
-                self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view(
-                    num_query_groups, -1, grad_shape[-1]
-                )
-                for g in qkv_grads
-            ]
-            grad = torch.cat(qkv_grads, dim=1).view(grad_shape)
-        else:
-            grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim)
-        return grad
-
-
-def get_megatron_muon_optimizer(
-    config: OptimizerConfig,
-    model_chunks: List[MegatronModule],
-    config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = None,
-    use_gloo_process_groups: bool = True,
-    layer_wise_distributed_optimizer: bool = False,
-    pg_collection: Optional[ProcessGroupCollection] = None,
-    dump_param_to_param_group_map: Optional[str] = None,
-) -> MegatronOptimizer:
-    """This function is used to get the muon optimizer for the model chunks.
-    It is used to get the muon optimizer for the model chunks.
-
-    Args:
-        config (OptimizerConfig): optimizer configuration object.
-        model_chunks (List[MegatronModule]): model chunks to get optimizer for.
-        config_overrides (Optional[Dict[ParamKey, ParamGroupOverride]]): optional dictionary
-            of optimizer/scheduler overrides for parameter subsets.
-        use_gloo_process_groups (bool): if false, disable use of Gloo process groups
-            in underlying Megatron optimizers.
-        layer_wise_distributed_optimizer (bool): if true, use layer-wise distributed optimizer.
-        Defaults to False.
-    """
-    assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed."
-
-    # dist-optim is not supported due to strong coupling with how DDP init grad buffer
-    # in thoery we can put some weight to use non-dist-muon and rest to dist-adam
-    # but there are strong dependency and assumption in DDP that prevent it
-    if config.use_distributed_optimizer:
-        raise Exception("muon with dist optimizer is not supported.")
-
-    # before this function receive properly created collection
-    if pg_collection is None:
-        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
-        pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True)
-        pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group()
-
-    log_single_rank(logger, logging.INFO, f"Setting up emerging optimizer with config {config}")
-
-    optimizers = []
-    # record list of non/linear params
-    linear_params = []
-    nonlinear_params = []
-
-    for model_chunk in model_chunks:
-        # use config to determine qkv split shapes.
-        # no need to check tp since tp splits by head and this is per head(group) dimension
-        num_attention_heads = model_chunk.config.num_attention_heads
-        num_query_groups = model_chunk.config.num_query_groups
-        kv_channels = model_chunk.config.kv_channels
-        qkv_split_shapes = [
-            num_attention_heads // num_query_groups * kv_channels,
-            kv_channels,
-            kv_channels,
-        ]
-        for name, param in model_chunk.named_parameters():
-            if not param.requires_grad:
-                continue
-            # add flag for expert weight so optimizer can figure which tp group it uses
-            # alternatively, create new param group and save tp_group. this require more
-            # change in optimizer
-            if "experts" in name and "shared" not in name:
-                param.expert_tp = True
-            # add flag for qkv parameter
-            # TODO(deyuf): support MLA
-            if "linear_qkv.weight" in name and len(param.shape) == 2:
-                param.is_qkv = True
-            # TODO(deyuf): might not be sufficient for future algorithm. revisit this conditioning
-            if not getattr(param, "is_embedding_or_output_parameter", False) and not (len(param.shape) == 1):
-                linear_params.append(param)
-            else:
-                nonlinear_params.append(param)
-
-    # freezing nonlinear params and get param groups for muon
-    for param in nonlinear_params:
-        param.requires_grad = False
-
-    linear_param_groups = _get_param_groups(
-        model_chunks=model_chunks,
-        config=config,
-        config_overrides=config_overrides,
-    )
-
-    optimizer = TensorParallelMuon(
-        linear_param_groups,
-        lr=config.lr,
-        momentum_beta=config.muon_momentum,
-        use_nesterov=config.muon_use_nesterov,
-        weight_decay=config.weight_decay,
-        fp32_matmul_prec=config.muon_fp32_matmul_prec,
-        num_ns_steps=config.muon_num_ns_steps,
-        scale_mode=config.muon_scale_mode,
-        split_qkv=config.muon_split_qkv,
-        is_qkv_fn=lambda p: getattr(p, "is_qkv", False),
-        qkv_split_shapes=qkv_split_shapes,
-        extra_scale_factor=config.muon_extra_scale_factor,
-        pg_collection=pg_collection,
-        mode=config.muon_tp_mode,
-    )
-
-    # set config here to:
-    # 1. get adam for rest of layer
-    # 2. avoid ChainedOptimizer check fail that assert all optimizers are same kind
-    # side effect is muon optimizer will have wrong name str, i.e. config.optimizer == 'adam'
-    # TODO(deyuf): allow user to select optimizer mix and relax ChainedOptimizer design
-    config.optimizer = "adam"
-
-    # Needed for torch_dist ckpt_format, unlike torch ckpt_format
-    # For other emerging optimizers, need to implement init_state_fn as well
-    # TODO(boxiangw): Improve usability after optimizer refactor
-    # TODO(boxiangw): support precision aware optimizer
-    def muon_init_state_fn(opt, config=None):
-        for group in opt.param_groups:
-            for p in group["params"]:
-                if len(opt.state[p]) == 0:
-                    opt.state[p]["momentum_buffer"] = torch.zeros_like(p.data)
-
-    def adam_init_state_fn(opt, config=None):
-        for group in opt.param_groups:
-            for p in group["params"]:
-                if len(opt.state[p]) == 0:
-                    if config is None or not config.use_precision_aware_optimizer:
-                        opt.state[p]["exp_avg"] = torch.zeros_like(p.data)
-                        opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data)
-                    else:
-                        opt.initialize_state(p)
-
-    # need to wrap into megatron mix precision optimizer. (only support bf16 w/o loss scale now)
-    if config.fp16:
-        raise Exception("muon with fp16 is not supported.")
-
-    reset_config_bf16 = False
-    if config.bf16:
-        if layer_wise_distributed_optimizer:
-            # creating master weight before layerwise sharding will lead to unnecessary master
-            # weight so here we delay master weight creation into layer_wise unset config.bf16
-            # will also result in all optimizers below(adam) to also not be wrapped
-            config.bf16 = False
-            reset_config_bf16 = True
-        else:
-            # if not using layer_wise wrapper, just create master weight here is fine
-            optimizer = Float16OptimizerWithFloat16Params(optimizer, config, None, muon_init_state_fn)
-    else:
-        optimizer = FP32Optimizer(optimizer, config, muon_init_state_fn)
-
-    optimizers.append(optimizer)
-
-    # done with muon, unfreeze nonlinear and freeze linear
-    for param in nonlinear_params:
-        param.requires_grad = True
-    for param in linear_params:
-        param.requires_grad = False
-
-    # call original get. linear params will be skipped since they're freezed
-    chained_adam = get_megatron_optimizer(
-        config,
-        model_chunks,
-        config_overrides=config_overrides,
-        use_gloo_process_groups=use_gloo_process_groups,
-        pg_collection=pg_collection,
-        dump_param_to_param_group_map=dump_param_to_param_group_map,
-    )
-
-    # unfreeze everything
-    for param in linear_params:
-        param.requires_grad = True
-
-    # chain everything together
-    optimizers += chained_adam.chained_optimizers
-
-    if layer_wise_distributed_optimizer:
-        log_single_rank(logger, logging.INFO, "Using LayerWiseDistributedOptimizer for Muon")
-        if reset_config_bf16:
-            config.bf16 = True
-        return LayerWiseDistributedOptimizer(
-            optimizers,
-            config,
-            pg_collection,
-            init_state_fn_list=[muon_init_state_fn, adam_init_state_fn],
-        )
-    return ChainedOptimizer(optimizers)
diff --git a/primus/backends/megatron/core/optimizer/moun_optimizer_config.py b/primus/backends/megatron/core/optimizer/moun_optimizer_config.py
deleted file mode 100644
index 101e997a9..000000000
--- a/primus/backends/megatron/core/optimizer/moun_optimizer_config.py
+++ /dev/null
@@ -1,287 +0,0 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-from dataclasses import dataclass
-from typing import Callable, Optional
-
-import torch
-from megatron.core.utils import is_te_min_version
-
-
-@dataclass
-class MounOptimizerConfig:
-    """Configuration for optimizer."""
-
-    ##############
-    # General
-    ##############
-    optimizer: str = "adam"
-    """Optimizer to use (one of Adam, SGD, or Muon)."""
-
-    lr: Optional[float] = None
-    """Initial learning rate. Depending on decay style and initial warmup, the learning rate at each
-       iteration would be different.
-    """
-
-    min_lr: Optional[float] = None
-    """Minumum value for learning rate. The scheduler clip values below this threshold."""
-
-    decoupled_lr: Optional[float] = None
-    """Separate learning rate for the input and output layer."""
-
-    decoupled_min_lr: Optional[float] = None
-    """Minimum value for learning rate for the input and output layer. The scheduler clip values
-       below this threshold.
-    """
-
-    weight_decay: float = 0.01
-    """Weight decay coefficient for L2 regularization."""
-
-    ##############
-    # Precision
-    ##############
-    fp8_recipe: Optional[str] = None
-    """The type of fp8 recipe will affect the processing logic inside distributed optimizer."""
-
-    fp16: bool = False
-    """If true, train with fp16 mixed precision training. Defaults to False."""
-
-    bf16: bool = False
-    """If true, train with bf16 mixed precision training. Defaults to False."""
-
-    reuse_grad_buf_for_mxfp8_param_ag: bool = False
-    """If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be
-       set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True."""
-
-    params_dtype: torch.dtype = torch.float32
-    """dtype used when intializing the weights. Defaults to torch.float32."""
-
-    use_precision_aware_optimizer: bool = False
-    """If true, allows optimizer-related tensors (master_param, gradients and optimizer states)
-    to be set to lower precision. Defaults to False.
-    """
-
-    store_param_remainders: bool = True
-    """If true, store the 16-bit FP32 parameter remainders in the optimizer state, excluding the
-        16 bits shared with the BF16 parameters. This lowers GPU memory usage. Defaults to True.
-    """
-
-    main_grads_dtype: torch.dtype = torch.float32
-    """dtype of main grads when enabling precision-aware-optimizer"""
-
-    main_params_dtype: torch.dtype = torch.float32
-    """dtype of main params when enabling precision-aware-optimizer"""
-
-    exp_avg_dtype: torch.dtype = torch.float32
-    """dtype of exp_avg when enabling precision-aware-optimizer"""
-
-    exp_avg_sq_dtype: torch.dtype = torch.float32
-    """dtype of exp_avg_sq when enabling precision-aware-optimizer"""
-
-    ###############
-    # Loss scaling
-    ###############
-    loss_scale: Optional[float] = None
-    """Static loss scaling, positive power of 2 values can improve fp16 convergence. If None,
-       dynamic loss scaling is used.
-    """
-
-    initial_loss_scale: float = 2**32
-    """Initial loss-scale for dynamic loss scaling."""
-
-    min_loss_scale: float = 1.0
-    """Minimum loss scale for dynamic loss scaling."""
-
-    loss_scale_window: float = 1000
-    """Window over which to raise/lower dynamic scale."""
-
-    hysteresis: int = 2
-    """Hysteresis for dynamic loss scaling."""
-
-    ##############
-    # Optimizer
-    ##############
-    # Adam
-    adam_beta1: float = 0.9
-    """First coefficient for computing running averages of gradient and its square in Adam
-    optimizer.
-    """
-
-    adam_beta2: float = 0.999
-    """Second coefficient for computing running averages of gradient and its square in Adam
-    optimizer.
-    """
-
-    adam_eps: float = 1e-08
-    """Term added to the denominator to improve numerical stability in Adam optimizer."""
-
-    decoupled_weight_decay: bool = True
-    """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false,
-    original Adam update rule will be used. Defaults to True.
-    """
-
-    # SGD.
-    sgd_momentum: float = 0.9
-    """Momentum factor for SGD optimizer."""
-
-    # Muon
-    muon_momentum: float = 0.95
-    """The momentum used by the internal SGD."""
-
-    muon_split_qkv: bool = True
-    """Whether to split QKV parameters for Muon optimizer."""
-
-    muon_use_nesterov: bool = False
-    """Whether to use Nesterov-style momentum in the internal SGD."""
-
-    muon_scale_mode: str = "spectral"
-    """The mode to use for the scale factor. Defaults to "spectral"."""
-
-    muon_fp32_matmul_prec: str = "medium"
-    """The precision to use for the fp32 matmul. Defaults to "medium"."""
-
-    muon_num_ns_steps: int = 5
-    """The number of iteration steps to use in the Newton-Schulz iteration."""
-
-    muon_tp_mode: str = "blockwise"
-    """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise"."""
-
-    muon_extra_scale_factor: float = 1.0
-    """Additional scale factor for the muon update."""
-
-    #######################
-    # Distributed optimizer
-    #######################
-    use_distributed_optimizer: bool = False
-    """Distribute optimizer state over data-parallel replicas."""
-
-    overlap_param_gather: bool = False
-    """If true, overlap param all-gather with forward compute.
-        This argument is intended to have the same value as the "overlap_param_gather" argument
-        in the "distributed_data_parallel_config.py" file. In the optimizer, this argument is
-        only used when "reuse_grad_buf_for_mxfp8_param_ag=True & fp8_param_gather=True".
-    """
-
-    overlap_param_gather_with_optimizer_step: bool = False
-    """If true, overlap param all-gather of first bucket with optimizer step."""
-
-    #######################
-    # Optimizer Offload
-    #######################
-
-    optimizer_cpu_offload: bool = False
-    """If True, offload optimizer states tensor and compute to CPU."""
-
-    optimizer_offload_fraction: float = 0.0
-    """Specifies the fraction of optimizer states to offload from GPU memory to CPU."""
-
-    use_torch_optimizer_for_cpu_offload: bool = False
-    """If True, use torch.optim.Optimizer for CPU offload."""
-
-    overlap_cpu_optimizer_d2h_h2d: bool = False
-    """
-    When set to `True`, this flag enables overlapping of the CPU optimizer
-    update process with the data transfer operations. This can help improve
-    overall training efficiency by reducing idle time during data movement,
-    allowing the optimizer to perform updates while gradients and parameters
-    are being transferred between devices.
-    """
-
-    pin_cpu_grads: bool = True
-    """If True, pin the optimizer gradients to CPU memory."""
-
-    pin_cpu_params: bool = True
-    """If True, pin the optimizer parameters to CPU memory."""
-
-    ################
-    # Miscellaneous
-    ################
-    clip_grad: float = 1.0
-    """Gradient clipping based on global L2 norm."""
-
-    log_num_zeros_in_grad: bool = False
-    """If true, calculate and log the number of zeros in gradient."""
-
-    barrier_with_L1_time: bool = False
-    """If true, use barrier with level 1 time measurements."""
-
-    timers: Optional[Callable] = None
-    """Function to get timers."""
-
-    config_logger_dir: str = ""
-    """When non-empty, dumps entry-point configs to config_logger_dir"""
-
-    def __post_init__(self):
-        """Check the validity of the config."""
-
-        # The following condition is used to avoid repetition in distrib_optimizer.py.
-        # This is because in distrib_optimizer.py, the process to handle parameters are
-        # different for different training precision settings. FP8 cases require different
-        # handling while FP8 delayed scaling is an exception because the Adam optimizer in
-        # TransformerEngine supports it in the kernel computation.
-        # This is also the flag to determine the usage of param.grad or param.decoupled_grad
-        self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = self.use_precision_aware_optimizer and (
-            self.main_params_dtype != torch.float32
-            or (self.fp8_recipe is None or self.fp8_recipe == "delayed")
-            or self.optimizer_cpu_offload
-        )
-
-        if self.fp8_recipe == "mxfp8":
-            if not self.reuse_grad_buf_for_mxfp8_param_ag:
-                import warnings
-
-                warnings.warn(
-                    "mxfp8 without using reuse_grad_buf_for_mxfp8_param_ag and fp8_param_gather"
-                    "will use significant amount additional GPU memory."
-                    "Setting --reuse-grad-buf-for-mxfp8-param-ag and --fp8-param-gather is "
-                    "recommended for mxfp8 training."
-                )
-
-        if self.use_precision_aware_optimizer:
-            assert self.optimizer == "adam", "--use-precision-aware-optimizer only supported with adam"
-            assert (
-                self.use_distributed_optimizer
-            ), "--use-precision-aware-optimizer only supported with distributed optimizer"
-
-            if not is_te_min_version("2.1.0"):
-                self.store_param_remainders = False
-
-            # Only the FusedAdam in TE and HybridDeviceOptimizer supports
-            # --use-precision-aware-optimizer.
-            # TODO: Remove this check when apex's FusedAdam is no longer used.
-            if self.optimizer_cpu_offload:
-                return
-            try:
-                import inspect
-
-                from transformer_engine.pytorch.optimizers import FusedAdam as Adam
-
-                adam_args = inspect.signature(Adam).parameters
-                arg_names = [
-                    "master_weight_dtype",
-                    "exp_avg_dtype",
-                    "exp_avg_sq_dtype",
-                    "use_decoupled_grad",
-                ]
-                for name in arg_names:
-                    assert name in adam_args, (
-                        "Current FusedAdam of TE doesn't support --use-precision-aware-optimizer, "
-                        "please update TE version."
-                    )
-            except ImportError:
-                raise RuntimeError(
-                    "--use-precision-aware-optimizer requires FusedAdam from TransformerEngine, "
-                    "but not found."
-                )
-        else:
-            assert (
-                self.main_grads_dtype == torch.float32
-            ), "main_grads_dtype can only be fp32 when not using precision-aware optimizer"
-            assert (
-                self.main_params_dtype == torch.float32
-            ), "main_params_dtype can only be fp32 when not using precision-aware optimizer"
-            assert (
-                self.exp_avg_dtype == torch.float32
-            ), "exp_avg_dtype can only be fp32 when not using precision-aware optimizer"
-            assert (
-                self.exp_avg_sq_dtype == torch.float32
-            ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer"
diff --git a/primus/backends/megatron/core/optimizer/zbpp_optimizer.py b/primus/backends/megatron/core/optimizer/zbpp_optimizer.py
index c554d3d28..2f50d2413 100644
--- a/primus/backends/megatron/core/optimizer/zbpp_optimizer.py
+++ b/primus/backends/megatron/core/optimizer/zbpp_optimizer.py
@@ -18,7 +18,7 @@
     multi_tensor_applier,
 )
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 
 class ZeroBubblePPChainedOptimizer(ChainedOptimizer):
diff --git a/primus/backends/megatron/core/parallel_state.py b/primus/backends/megatron/core/parallel_state.py
index 71b76d7cc..b7e6b7e4f 100644
--- a/primus/backends/megatron/core/parallel_state.py
+++ b/primus/backends/megatron/core/parallel_state.py
@@ -17,7 +17,7 @@
 )
 from megatron.core.utils import get_pg_rank, get_pg_size
 
-from primus.modules.trainer.megatron.utils import is_v_schedule_enabled
+from primus.backends.megatron.training.utils import is_v_schedule_enabled
 
 
 def is_pp_first_stage(pp_group: torch.distributed.ProcessGroup):
diff --git a/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py b/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py
new file mode 100644
index 000000000..85387b258
--- /dev/null
+++ b/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py
@@ -0,0 +1,193 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Pipeline-parallel visualization helpers.
+
+Records per-iteration forward/backward/weight-grad CUDA events so the offline
+``dump_pp_data`` step can emit a JSON timeline consumed by the PP visualizer.
+These wrappers are installed by the ``megatron.pp.dump_pp_data`` patch (and by
+the Primus-Pipe / ZeroBubble handlers) when ``--dump_pp_data`` is enabled.
+"""
+
+import json
+import os
+
+import torch
+from megatron.core import parallel_state
+
+_GLOBAL_PP_VIS_EVENTS = []
+_GLOBAL_PP_VIS_EVENTS_PER_ITER = None
+
+
+def schedule_wrapper(func):
+    def wrapper(*args, **kwargs):
+        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER = {
+            "start": None,
+            "end": None,
+            "memory": None,
+            "fwd_start": [],
+            "fwd_end": [],
+            "fwd_minibatch": [],
+            "fwd_chunk": [],
+            "bwd_start": [],
+            "bwd_end": [],
+            "bwd_minibatch": [],
+            "bwd_chunk": [],
+            "wgrad_start": [],
+            "wgrad_end": [],
+            "wgrad_minibatch": [],
+            "wgrad_chunk": [],
+        }
+
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"] = torch.cuda.Event(enable_timing=True)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"].record()
+        res = func(*args, **kwargs)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"] = torch.cuda.Event(enable_timing=True)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"].record()
+
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["memory"] = torch.cuda.max_memory_reserved() / 1024**3
+
+        global _GLOBAL_PP_VIS_EVENTS
+        _GLOBAL_PP_VIS_EVENTS.append(_GLOBAL_PP_VIS_EVENTS_PER_ITER)
+
+        return res
+
+    return wrapper
+
+
+def fwd_bwd_wrapper(func, mode, minibatch=None, chunk=None):
+    def wrapper(*args, **kwargs):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        res = func(*args, **kwargs)
+        end.record()
+
+        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_start"].append(start)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_end"].append(end)
+
+        if minibatch is not None:
+            _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_minibatch"].append(minibatch)
+        if chunk is not None:
+            _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_chunk"].append(chunk)
+        return res
+
+    return wrapper
+
+
+def combined_fwd_bwd_wrapper(func, fwd_minibatch, fwd_chunk, bwd_minibatch, bwd_chunk):
+    """Record a single combined forward+backward call as both an ``fwd`` event
+    and a ``bwd`` event sharing the same ``[start, end]`` interval.
+
+    Used by ``megatron_combined_fwd_bkwd_handler`` so that nodes collapsed into
+    a combined FB group still appear in the dump_pp_data output. Without this
+    the visualizer's per-rank F/B/W totals are heavily under-counted on ranks
+    that hit the steady state (the F and B halves are interleaved inside
+    ``combined_forward_backward_step`` and cannot be timed separately).
+    """
+
+    def wrapper(*args, **kwargs):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        res = func(*args, **kwargs)
+        end.record()
+
+        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_start"].append(start)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_end"].append(end)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_minibatch"].append(fwd_minibatch)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_chunk"].append(fwd_chunk)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_start"].append(start)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_end"].append(end)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_minibatch"].append(bwd_minibatch)
+        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_chunk"].append(bwd_chunk)
+        return res
+
+    return wrapper
+
+
+def set_dump_pp_data_patch():
+    from megatron.core.pipeline_parallel import schedules
+
+    schedules.forward_step = fwd_bwd_wrapper(schedules.forward_step, "fwd")
+    schedules.backward_step = fwd_bwd_wrapper(schedules.backward_step, "bwd")
+
+
+def dump_pp_data(args, num_mbs, pp_data_dir):
+    torch.cuda.synchronize()
+
+    global _GLOBAL_PP_VIS_EVENTS
+    all_iter_data = {}
+    for iter_idx, iter_events in enumerate(_GLOBAL_PP_VIS_EVENTS):
+        iter_data = {
+            "total": None,
+            "memory": None,
+            "fwd_start": [],
+            "fwd_end": [],
+            "fwd_minibatch": [],
+            "fwd_chunk": [],
+            "bwd_start": [],
+            "bwd_end": [],
+            "bwd_minibatch": [],
+            "bwd_chunk": [],
+            "wgrad_start": [],
+            "wgrad_end": [],
+            "wgrad_minibatch": [],
+            "wgrad_chunk": [],
+        }
+        iter_data["total"] = iter_events["start"].elapsed_time(iter_events["end"])
+        iter_data["memory"] = iter_events["memory"]
+
+        for i in range(len(iter_events["fwd_start"])):
+            for key in ["fwd_start", "fwd_end", "bwd_start", "bwd_end", "wgrad_start", "wgrad_end"]:
+                if i >= len(iter_events[key]):
+                    continue
+                event_time = iter_events["start"].elapsed_time(iter_events[key][i])
+                iter_data[key].append(event_time)
+            for key in [
+                "fwd_minibatch",
+                "fwd_chunk",
+                "bwd_minibatch",
+                "bwd_chunk",
+                "wgrad_minibatch",
+                "wgrad_chunk",
+            ]:
+                if i >= len(iter_events[key]):
+                    continue
+                iter_data[key].append(iter_events[key][i])
+
+        all_iter_data[iter_idx + 1] = iter_data
+
+    rank = torch.distributed.get_rank()
+    dp_rank = parallel_state.get_data_parallel_rank()
+    pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+    os.makedirs(pp_data_dir, exist_ok=True)
+    if dp_rank == 0:
+        log_path = os.path.join(pp_data_dir, f"pp_rank_{pp_rank}.json")
+        with open(log_path, "w") as f:
+            json.dump(all_iter_data, f, indent=2)
+
+    if rank == 0:
+        vp_size = args.virtual_pipeline_model_parallel_size
+        vp_size = 1 if vp_size is None else vp_size
+        config_dict = {
+            "world_size": args.world_size,
+            "dp_size": args.data_parallel_size,
+            "tp_size": args.tensor_model_parallel_size,
+            "ep_size": args.expert_model_parallel_size,
+            "pp_size": args.pipeline_model_parallel_size,
+            "vp_size": vp_size,
+            "num_mbs": num_mbs,
+            "train_iters": args.train_iters,
+        }
+        log_path = os.path.join(pp_data_dir, f"config.json")
+        with open(log_path, "w") as f:
+            json.dump(config_dict, f, indent=2)
diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py
index 628f78dc4..8b6eec8b7 100644
--- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py
+++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py
@@ -8,6 +8,9 @@
 from megatron.core.pipeline_parallel.schedules import backward_step
 from megatron.training.global_vars import get_args
 
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    fwd_bwd_wrapper,
+)
 from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER
 from primus.core.pipeline_parallel.handler.wgrad_handler import WGRAD_RUNNING_CACHE
 from primus.core.pipeline_parallel.scheduler.scheduler_node import (
@@ -15,7 +18,6 @@
     SchedulerNode,
 )
 from primus.core.pipeline_parallel.utils import find_prev_node_with_type
-from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper
 
 
 def megatron_check_bwd_node_valid(node: SchedulerNode):
diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py
index 4a61cfd1e..d3ea91072 100644
--- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py
+++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py
@@ -8,6 +8,9 @@
 from megatron.core.pipeline_parallel.schedules import deallocate_output_tensor
 from megatron.training.global_vars import get_args
 
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    combined_fwd_bwd_wrapper,
+)
 from primus.backends.megatron.core.pipeline_parallel.primuspipe.handlers.communication_handler import (
     batch_p2p_communication_handler,
 )
@@ -17,7 +20,6 @@
     SchedulerNode,
 )
 from primus.core.pipeline_parallel.utils import find_prev_node_with_type
-from primus.modules.trainer.megatron.utils import combined_fwd_bwd_wrapper
 
 
 def megatron_check_combined_fwd_bkwd_node_valid(node: SchedulerNode):
diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py
index 374610fd9..b876119c6 100644
--- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py
+++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py
@@ -11,13 +11,15 @@
 )
 from megatron.training.global_vars import get_args
 
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    fwd_bwd_wrapper,
+)
 from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER
 from primus.core.pipeline_parallel.scheduler.scheduler_node import (
     FuncType,
     SchedulerNode,
 )
 from primus.core.pipeline_parallel.utils import find_prev_node_with_type
-from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper
 
 
 def megatron_check_fwd_node_valid(node: SchedulerNode):
diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py
index d196d8420..a0a09eb85 100644
--- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py
+++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py
@@ -33,7 +33,7 @@
 )
 from primus.core.pipeline_parallel.scheduler.scheduler import ScheduleRunner
 from primus.core.pipeline_parallel.scheduler.scheduler_node import FuncType
-from primus.modules.module_utils import warning_rank_0
+from primus.core.utils.module_utils import warning_rank_0
 
 
 class PrimusPipelineParallelLauncher:
@@ -309,7 +309,9 @@ def enable_grad_sync():
                 node.args["pp_group"] = pg_collection.pp
 
         if args.dump_pp_data:
-            from primus.modules.trainer.megatron.utils import schedule_wrapper
+            from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+                schedule_wrapper,
+            )
 
             schedule_wrapper(self.schedule_runner.run)(self.schedule_table, self.pp_rank)
         else:
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py
index d730f74ad..eb94068e0 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py
@@ -15,7 +15,7 @@
 import torch
 from torch.autograd.graph import saved_tensors_hooks
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 
 def checksum(tensor):
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py
index ff67048cb..296304fc6 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py
@@ -35,10 +35,12 @@
 from megatron.core.utils import get_model_config, get_model_type, get_model_xattn
 from megatron.training import get_args, print_rank_0
 
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    fwd_bwd_wrapper,
+)
 from primus.backends.megatron.training.training import RollbackDataIteratorWrapper
 from primus.backends.megatron.training.utils import is_second_last_pipeline_stage
-from primus.modules.module_utils import log_rank_0, log_rank_all
-from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper
+from primus.core.utils.module_utils import log_rank_0, log_rank_all
 
 from .offload import ActivationStorePool, FakeActivationStore, partial_recompute
 from .scheduler import (
@@ -454,7 +456,7 @@ def pre_load_batch(self, idx):
         cnt = (int(offload_time) + 1) * 3
         multi_chunks = get_virtual_pipeline_number() > 1
         conf = self.iteration_config
-        from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore
+        from primus.backends.megatron.data_loader_store import DataLoaderStore
 
         count = len(DataLoaderStore.cache)
         for i in range(cnt):
@@ -478,7 +480,7 @@ def pre_load_batch(self, idx):
     def load_all_batch(self):
         conf = self.iteration_config
         multi_chunks = get_virtual_pipeline_number() > 1
-        from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore
+        from primus.backends.megatron.data_loader_store import DataLoaderStore
 
         assert len(DataLoaderStore.cache) == 0
         for scheduled_node in conf.schedules:
@@ -531,7 +533,7 @@ def schedule_f_impl(self, scheduled_node: ScheduledNode):
             mem_before = torch.cuda.memory_allocated()
 
         set_seq_split_idx(scheduled_node.seq_split_idx)
-        from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore
+        from primus.backends.megatron.data_loader_store import DataLoaderStore
 
         if len(DataLoaderStore.cache) == 0:
             DataLoaderStore.push(conf.data_iterator[scheduled_node.chunk], vp_stage=scheduled_node.chunk)
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py
index a02982606..d405c2b79 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py
@@ -7,7 +7,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 from .graph import BW, F, GraphConfig, ScheduledNode
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py
index c28b679f0..e59ad7c04 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py
@@ -14,7 +14,7 @@
 
 from megatron.training.global_vars import get_args
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .graph import BW, B, CommDirection, F, FuncType, GraphConfig, ScheduledNode
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py
index 112419319..b66ac5e25 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py
@@ -13,7 +13,7 @@
 from enum import Enum
 from typing import List
 
-from primus.modules.module_utils import log_rank_0, log_rank_all
+from primus.core.utils.module_utils import log_rank_0, log_rank_all
 
 
 class PassType(Enum):
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py
index 4758ed9c6..a5049e12f 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py
@@ -11,7 +11,7 @@
 from dataclasses import dataclass
 from typing import List
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .graph import BW, B, F, FuncType, GraphConfig, NodeKey, ScheduledNode, W
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py
index 9edaf49ba..180a216de 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py
@@ -12,7 +12,7 @@
 
 from megatron.training.global_vars import get_args
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .communication import (
     add_communication_nodes_without_sorting,
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py
index e35e9a9b4..c6388cfbe 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py
@@ -11,7 +11,7 @@
 from collections import Counter, deque
 from dataclasses import dataclass
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 
 @dataclass(eq=True, frozen=True)
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py
index 2ca912c43..9af46992b 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py
@@ -7,7 +7,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .graph import BW, F, GraphConfig, ScheduledNode
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py
index b0ea0c47d..2a450220b 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py
@@ -13,7 +13,7 @@
 import pulp
 import torch
 
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .graph import FuncType, GraphConfig, ScheduledNode
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py
index 6ebfc15f3..203dae691 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py
@@ -9,7 +9,7 @@
 
 from collections import deque
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 from .graph import B, F, ScheduledNode, W
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py
index ad68b5e97..aea98f32c 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py
@@ -10,7 +10,7 @@
 # Implementation of vhalf and vmin schedules of Pipeline Parallelism
 # with Controllable Memory (https://arxiv.org/abs/2405.15362)
 # The reordering is based on a greedy algorithm.
-from primus.modules.module_utils import log_rank_all
+from primus.core.utils.module_utils import log_rank_all
 
 from .graph import B, F, ScheduledNode, W
 
diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py
index 7ef4315e0..6d7a36ddf 100644
--- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py
+++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py
@@ -15,7 +15,9 @@
 from megatron.core import parallel_state
 from megatron.training import get_args
 
-from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    fwd_bwd_wrapper,
+)
 
 
 def add_zero_bubble_args(parser):
diff --git a/primus/backends/megatron/core/transformer/clamped_swiglu.py b/primus/backends/megatron/core/transformer/clamped_swiglu.py
new file mode 100644
index 000000000..168986513
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/clamped_swiglu.py
@@ -0,0 +1,192 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Clamped SwiGLU activation for DeepSeek-V4 (pre-mul layout).
+
+Reference: techblog §3 ("Activation: clamped SwiGLU") and the inference
+reference at ``DeepSeek-V4-Flash/inference/model.py:Expert.forward``.
+
+DeepSeek-V4 replaces the standard ``SwiGLU(gate, up) = SiLU(gate) * up``
+with a **pre-multiplication** clamp:
+
+    gate_c = clamp(gate, max=alpha)             # one-sided (top only)
+    up_c   = clamp(up,   min=-alpha, max=alpha) # two-sided
+    out    = SiLU(gate_c) * up_c
+
+This matches the released checkpoint and Megatron's
+``mlp.MLP``-side clamp path (``activation_func_clamp_value``); both
+clamp the gate / up *before* the activation and multiply, so the
+post-multiply value is bounded by ``alpha * max(SiLU)`` and stays
+well-behaved in bf16 expert summations.
+
+Plan-2 P14 contract:
+
+* :func:`clamped_swiglu_pre_mul` — split-input pointwise activation. Used
+  by the eager :class:`ClampedSwiGLUMLP` and as the canonical reference
+  for unit tests.
+* :func:`clamped_swiglu_pre_mul_fused` — Megatron-fused-input
+  ``[..., 2I]`` gate-concat-up form. Lets a grouped-gemm expert backend
+  call a single function on the post-GEMM output.
+* :class:`ClampedSwiGLUMLP` — eager MLP using **separate** ``w1`` (gate)
+  and ``w3`` (up) Linears so the parameter layout matches the
+  ``DeepSeek-V4-Flash`` checkpoint (``w1.weight`` / ``w3.weight``).
+  An optional ``fused_gate_up`` flag lets callers fuse the two GEMMs at
+  forward time without changing the state-dict layout — the saved /
+  loaded keys are always ``w1.weight`` / ``w3.weight``.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _resolve_alpha(alpha: Optional[float]) -> Optional[float]:
+    """Return a clamp bound or ``None`` when clamping is disabled."""
+    if alpha is None:
+        return None
+    if alpha <= 0.0:
+        return None
+    return float(alpha)
+
+
+def clamped_swiglu_pre_mul(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    alpha: float = 7.0,
+) -> torch.Tensor:
+    """Pre-multiplication clamped SwiGLU on split inputs.
+
+    Args:
+        gate: ``[..., I]`` gate stream (output of ``w1``).
+        up:   ``[..., I]`` up stream (output of ``w3``).
+        alpha: clamp bound; V4-Flash default is ``7.0``. Pass ``0`` /
+            ``None`` to fall back to vanilla ``SiLU(gate) * up``.
+
+    Returns:
+        ``[..., I]`` activation output, ``SiLU(clamp(gate, max=alpha))
+        * clamp(up, +/- alpha)``.
+    """
+    if gate.shape != up.shape:
+        raise ValueError(
+            "clamped_swiglu_pre_mul expects matching gate / up shapes; "
+            f"got {tuple(gate.shape)} vs {tuple(up.shape)}."
+        )
+    bound = _resolve_alpha(alpha)
+    if bound is not None:
+        gate_c = gate.clamp(max=bound)
+        up_c = up.clamp(min=-bound, max=bound)
+    else:
+        gate_c = gate
+        up_c = up
+    return F.silu(gate_c) * up_c
+
+
+def clamped_swiglu_pre_mul_fused(
+    x: torch.Tensor,
+    *,
+    alpha: float = 7.0,
+) -> torch.Tensor:
+    """Pre-multiplication clamped SwiGLU on a fused ``[..., 2I]`` input.
+
+    The input is the Megatron-convention concatenated ``[gate | up]``
+    along the last dimension (matching
+    :func:`megatron.core.fusions.fused_bias_swiglu.bias_swiglu` and the
+    eager glu path in :class:`megatron.core.transformer.mlp.MLP`).
+
+    Args:
+        x: ``[..., 2 * I]`` — ``[gate | up]`` halves concatenated along
+            the last dim.
+        alpha: clamp bound; same semantics as :func:`clamped_swiglu_pre_mul`.
+
+    Returns:
+        ``[..., I]`` activation output.
+    """
+    if x.shape[-1] % 2 != 0:
+        raise ValueError(
+            "clamped_swiglu_pre_mul_fused expects a [gate | up] last dim; "
+            f"got shape {tuple(x.shape)} (last dim must be even)."
+        )
+    gate, up = x.chunk(2, dim=-1)
+    return clamped_swiglu_pre_mul(gate, up, alpha=alpha)
+
+
+class ClampedSwiGLUMLP(nn.Module):
+    """Eager SwiGLU MLP with V4's pre-multiplication clamp.
+
+    Computes::
+
+        gate = w1(x)                            # [..., I]
+        up   = w3(x)                            # [..., I]
+        h    = SiLU(clamp(gate, max=alpha))
+                * clamp(up, +/- alpha)
+        y    = w2(h)                            # [..., D]
+
+    The parameter layout (``w1`` / ``w2`` / ``w3``) mirrors the released
+    DeepSeek-V4-Flash ``Expert`` checkpoint exactly. The ``fused_gate_up``
+    knob fuses the gate / up GEMMs at *forward time only* by stacking
+    ``w1.weight`` and ``w3.weight`` on the fly; the saved / loaded
+    ``state_dict`` keys remain ``w1.weight`` / ``w3.weight`` so released
+    checkpoints can be loaded without remapping.
+
+    Notes:
+        * Bias is omitted by default to match V4 reference checkpoints.
+        * ``alpha=0`` (or ``None``) disables clamping → vanilla SwiGLU.
+        * This module is the canonical *eager* reference. Production
+          training uses Megatron's grouped-MLP path with
+          ``activation_func_clamp_value`` set; the math is the same.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        *,
+        alpha: float = 7.0,
+        bias: bool = False,
+        dtype: Optional[torch.dtype] = None,
+        fused_gate_up: bool = False,
+    ) -> None:
+        super().__init__()
+        if intermediate_size <= 0:
+            raise ValueError(f"intermediate_size must be > 0, got {intermediate_size}")
+        self.hidden_size = int(hidden_size)
+        self.intermediate_size = int(intermediate_size)
+        self.alpha = float(alpha)
+        self.fused_gate_up = bool(fused_gate_up)
+
+        kw = {} if dtype is None else {"dtype": dtype}
+        # w1 = gate, w3 = up (V4 convention; matches DeepSeek-V4-Flash checkpoint).
+        self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias, **kw)
+        self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias, **kw)
+        self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias, **kw)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.fused_gate_up:
+            # Build the fused weight on the fly so the state_dict layout
+            # is unchanged (w1.weight / w3.weight). Bias is None by default.
+            w_gu = torch.cat([self.w1.weight, self.w3.weight], dim=0)
+            b_gu = None
+            if self.w1.bias is not None and self.w3.bias is not None:
+                b_gu = torch.cat([self.w1.bias, self.w3.bias], dim=0)
+            gate_up = F.linear(x, w_gu, b_gu)
+            h = clamped_swiglu_pre_mul_fused(gate_up, alpha=self.alpha)
+        else:
+            gate = self.w1(x)
+            up = self.w3(x)
+            h = clamped_swiglu_pre_mul(gate, up, alpha=self.alpha)
+        return self.w2(h)
+
+
+__all__ = [
+    "clamped_swiglu_pre_mul",
+    "clamped_swiglu_pre_mul_fused",
+    "ClampedSwiGLUMLP",
+]
diff --git a/primus/backends/megatron/core/transformer/compressor.py b/primus/backends/megatron/core/transformer/compressor.py
new file mode 100644
index 000000000..3a8f21e94
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/compressor.py
@@ -0,0 +1,215 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DeepSeek-V4 Compressor.
+
+Reference: techblog §1.3 ("Compressor: the Long-Range Compression Branch") and
+the diagrams in ``deepseek-v4/develop/techblog/diagrams/csa.png`` /
+``hca.png``.
+
+Two configurations:
+
+* ``ratio == 4`` (CSA branch) — overlap mode, ``coff == 2``: each compressed
+  token sees an effective window of ``2*ratio`` raw tokens (current window
+  plus the previous window's "leftover-half" channels). This smooths
+  boundary effects between adjacent compressed positions.
+* ``ratio == 128`` (HCA branch) — non-overlap mode, ``coff == 1``: each
+  compressed token covers exactly ``ratio`` raw tokens.
+
+Compressor returns the pooled KV after a final ``kv_norm`` (RMSNorm). RoPE
+at the compress branch theta is applied **outside** this module by the
+caller (the dual-RoPE module produced in P4.3 + the CSA / HCA modules in
+P4.4 will consume the output).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.keep_in_fp32 import (
+    KeepInFp32Mixin,
+    mark_keep_in_fp32,
+)
+from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm
+
+
+class Compressor(KeepInFp32Mixin, nn.Module):
+    """V4 Compressor block.
+
+    Args:
+        hidden_size: input feature dim ``D``.
+        head_dim: output channel dim per compressed position.
+        ratio: compression ratio ``m``. Must be a divisor of the runtime
+            sequence length (``S % ratio == 0``).
+        overlap: whether to use the overlap-stitched mode. If ``None``,
+            defaults to ``ratio == 4`` (the V4 convention).
+        rmsnorm_eps: RMSNorm stability eps.
+
+    Shapes:
+        Forward input  ``hidden``: ``[B, S, D]``.
+        Forward output ``pooled``: ``[B, S // ratio, head_dim]``.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        head_dim: int,
+        ratio: int,
+        overlap: Optional[bool] = None,
+        rmsnorm_eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if ratio < 1:
+            raise ValueError(f"ratio must be >= 1, got {ratio}")
+
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+        self.ratio = ratio
+        self.overlap = bool(ratio == 4 if overlap is None else overlap)
+        # coff is the projection multiplier — overlap mode needs 2x the
+        # channels because half goes to the "current window" half and half to
+        # the "previous window" half.
+        self.coff = 2 if self.overlap else 1
+
+        proj_out = self.coff * head_dim
+        self._proj_out = proj_out
+        # Plain ``nn.Linear``, deliberately: the compressor has no low-precision
+        # path. It builds the KV entries every query reads and the pooling
+        # weights that mix them, and the open-source reference keeps it out of
+        # FP8 for that reason -- it wraps exactly this projection plus the
+        # indexer in an fp8-disabled context. ``nn.Linear`` is also immune to the
+        # enclosing TE / Turbo quantization context, so "high precision" here
+        # needs no guard beyond not adding a quantized branch.
+        #
+        # Fuse the kv + gate projections into ONE [hidden -> 2*proj_out] GEMM
+        # (default-on): ~1.5x on the projection and one launch instead of two.
+        # PRIMUS_COMPRESS_FUSE_PROJ=0 restores the two separate linears.
+        self._fuse_proj = os.environ.get("PRIMUS_COMPRESS_FUSE_PROJ", "1") != "0"
+        if self._fuse_proj:
+            self.wkv_gate = nn.Linear(hidden_size, 2 * proj_out, bias=False)
+        else:
+            self.wkv = nn.Linear(hidden_size, proj_out, bias=False)
+            self.wgate = nn.Linear(hidden_size, proj_out, bias=False)
+
+        # Learnable absolute position embedding (APE) added on top of the
+        # softmax score. After overlap, the effective window length is
+        # ``2*ratio`` slots of size ``head_dim``; in non-overlap mode it's
+        # ``ratio`` slots of size ``head_dim``.
+        # FP32 in the released checkpoint. The pooling path promotes with
+        # ``score.float()`` before its softmax, so following the model dtype
+        # costs stored resolution only; pinning it is opt-in via
+        # ``PRIMUS_V4_KEEP_FP32`` (see ``keep_in_fp32``).
+        ape_len = 2 * ratio if self.overlap else ratio
+        self.ape = nn.Parameter(torch.zeros(ape_len, head_dim, dtype=torch.float32))
+        nn.init.normal_(self.ape, std=0.02)
+        mark_keep_in_fp32(self.ape)
+
+        self.kv_norm = LocalRMSNorm(head_dim, eps=rmsnorm_eps)
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Bridge checkpoints across the fused/unfused projection layouts.
+
+        Old checkpoints store ``wkv.weight`` + ``wgate.weight``; the fused path
+        wants ``wkv_gate.weight`` = ``cat([wkv, wgate])`` (and vice-versa). Remap
+        in-place so either layout loads under either runtime setting.
+        """
+        wkv_k, wgate_k, fused_k = prefix + "wkv.weight", prefix + "wgate.weight", prefix + "wkv_gate.weight"
+        if self._fuse_proj and wkv_k in state_dict and fused_k not in state_dict:
+            state_dict[fused_k] = torch.cat([state_dict.pop(wkv_k), state_dict.pop(wgate_k)], dim=0)
+        elif (not self._fuse_proj) and fused_k in state_dict and wkv_k not in state_dict:
+            w = state_dict.pop(fused_k)
+            state_dict[wkv_k], state_dict[wgate_k] = w[: self._proj_out], w[self._proj_out :]
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    # ------------------------------------------------------------------
+    # internals
+    # ------------------------------------------------------------------
+
+    def _reshape_into_windows(self, t: torch.Tensor) -> torch.Tensor:
+        """``[B, S, coff*head_dim]`` → ``[B, N, ratio, coff*head_dim]``,
+        where ``N = S // ratio``.
+        """
+        B, S, C = t.shape
+        assert S % self.ratio == 0, f"Compressor: sequence length {S} not divisible by ratio {self.ratio}"
+        N = S // self.ratio
+        return t.reshape(B, N, self.ratio, C)
+
+    def _overlap_transform(self, t: torch.Tensor) -> torch.Tensor:
+        """``[B, N, ratio, 2*head_dim]`` → ``[B, N, 2*ratio, head_dim]``.
+
+        For window ``i``, the augmented sequence is
+        ``[half_a[i], half_b[i-1]]`` concatenated along the per-window
+        axis. Window 0's "previous half" is filled with zeros (causal
+        padding).
+        """
+        # Split channels.
+        half_a, half_b = torch.chunk(t, 2, dim=-1)  # each [B, N, ratio, head_dim]
+        # Roll along the window dim so half_b[i] becomes "previous-window's b" of i+1.
+        half_b_prev = torch.cat(
+            [torch.zeros_like(half_b[:, :1]), half_b[:, :-1]],
+            dim=1,
+        )
+        # Concat along the per-window token axis.
+        return torch.cat([half_a, half_b_prev], dim=2)  # [B, N, 2*ratio, head_dim]
+
+    # ------------------------------------------------------------------
+    # public API
+    # ------------------------------------------------------------------
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Pool ``hidden[B, S, D]`` to ``[B, S/ratio, head_dim]``."""
+        if self._fuse_proj:
+            kv_proj, score_proj = self.wkv_gate(hidden).split(self._proj_out, dim=-1)
+        else:
+            kv_proj = self.wkv(hidden)  # [B, S, coff*head_dim]
+            score_proj = self.wgate(hidden)  # [B, S, coff*head_dim]
+
+        kv = self._reshape_into_windows(kv_proj)  # [B, N, ratio, coff*head_dim]
+        score = self._reshape_into_windows(score_proj)  # [B, N, ratio, coff*head_dim]
+
+        if self.overlap:
+            kv = self._overlap_transform(kv)  # [B, N, 2*ratio, head_dim]
+            score = self._overlap_transform(score)  # [B, N, 2*ratio, head_dim]
+        # else: kv / score already at [B, N, ratio, head_dim]
+
+        # Per-window-softmax pool: APE bias + softmax over the window axis (dim=2)
+        # + weighted sum -- each compressed token is a softmax-weighted average of
+        # its window members. The forward burst (add + cast + softmax + cast + mul
+        # + reduce) is fused into one Triton launch on CUDA fp16/bf16/fp32 inputs;
+        # PRIMUS_COMPRESS_POOL_TRITON=0 (or non-CUDA / unsupported dtype) falls back
+        # to eager.
+        if (
+            os.environ.get("PRIMUS_COMPRESS_POOL_TRITON", "1") != "0"
+            and kv.is_cuda
+            and kv.dtype
+            in (
+                torch.float16,
+                torch.bfloat16,
+                torch.float32,
+            )
+        ):
+            from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.compressor_pool import (
+                fused_softmax_weighted_pool,
+            )
+
+            pooled = fused_softmax_weighted_pool(kv, score, self.ape)  # [B, N, head_dim]
+        else:
+            score = score + self.ape  # [B, N, win, head_dim]
+            weights = F.softmax(score.float(), dim=2).to(kv.dtype)
+            pooled = (kv * weights).sum(dim=2)  # [B, N, head_dim]
+
+        pooled = self.kv_norm(pooled)
+        return pooled
+
+
+__all__ = ["Compressor"]
diff --git a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py
new file mode 100644
index 000000000..d05ab3248
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py
@@ -0,0 +1,1819 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DeepSeek-V4 attention.
+
+Plan-2 P13 — *faithful* attention rooted on Megatron's
+``MLASelfAttention``. The released ``DeepSeek-V4-Flash`` checkpoint is
+reproduced for **all three** layer types (``compress_ratio in {0, 4, 128}``)
+inside a single attention class:
+
+* Single-latent KV: a single ``linear_kv`` projection ``hidden -> head_dim``
+  produces both K and V, broadcast across all query heads.
+* Per-head ``q_rms``: a parameter-less RMS normalization on ``head_dim``
+  applied AFTER ``linear_q_up_proj`` and BEFORE partial RoPE — matches
+  the ``inference/model.py`` reference exactly.
+* Grouped low-rank O projection: ``linear_o_a`` / ``linear_o_b`` (when
+  ``config.o_lora_rank > 0``) replace the standard flat ``linear_proj``.
+* Learnable per-head ``attn_sink``: an extra "virtual key" column with
+  zero value, joined into the softmax. Drops the column after softmax
+  so the value-weighted sum is unaffected; the head can still spend mass
+  on the sink as a "no attention" fallback.
+* Compressed branches (``compress_ratio > 0``) fold their compressor
+  (and indexer for CSA) in as :class:`Compressor` / :class:`Indexer`
+  spec submodules; the dense local SWA branch and the compressed branch
+  are softmax-joined together so the attention sink is shared across
+  both paths.
+* Field names mirror MLA's canonical layout (``linear_q_down_proj``,
+  ``linear_q_up_proj``, ``q_layernorm``, ``kv_layernorm``) plus the V4
+  extras (``linear_kv``, ``linear_o_a``, ``linear_o_b``, ``compressor``,
+  ``indexer``) so the state-dict adapter (P17) can map the released
+  safetensors keys
+  (``layers.{i}.attn.{wq_a,wq_b,wkv,q_norm,kv_norm,wo_a,wo_b,attn_sink,
+  compressor.*,indexer.*}``) in one straightforward table.  The
+  per-head learnable softmax sink lives directly on the attention module
+  as ``self.attn_sink: nn.Parameter`` (no submodule slot — Plan-3 P21
+  dropped the dead ``attn_sink`` field; the inline softmax-with-sink
+  path in :meth:`_attention_forward` is canonical).
+
+Forward signature:
+
+.. code-block:: python
+
+    out = attn(
+        hidden,                  # [B, S, D]
+        position_ids,            # [B, S] or [S]
+    )
+    # out: [B, S, D]
+"""
+
+from __future__ import annotations
+
+import atexit
+import collections
+import logging
+import math
+import os
+import statistics
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.multi_latent_attention import MLASelfAttention
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.compressor import Compressor
+from primus.backends.megatron.core.transformer.dual_rope import (
+    DualRoPE,
+    apply_interleaved_partial_rope,
+)
+from primus.backends.megatron.core.transformer.indexer import Indexer
+from primus.backends.megatron.core.transformer.indexer_distill_loss import (
+    V4IndexerLossAutoScaler,
+    compute_indexer_distill_loss,
+    log_indexer_distill_loss,
+)
+from primus.backends.megatron.core.transformer.keep_in_fp32 import (
+    KeepInFp32Mixin,
+    mark_keep_in_fp32,
+    unmark_keep_in_fp32,
+)
+from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm
+from primus.backends.megatron.core.transformer.sliding_window_kv import (
+    sliding_window_causal_mask,
+)
+
+# All attention backend entries come from the kernels package __init__ (the
+# single entry point): eager, triton v1/v2. Naming: v4_attention_<be>
+# (dense/HCA) and v4_csa_attention_<be> (CSA). The gluon backend is NOT imported
+# here — it hard-depends on triton.experimental.gluon (gfx950 only) and is loaded
+# lazily via load_gluon_attention_backends() only when a layer selects it.
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (
+    eager_v4_attention,
+    eager_v4_csa_attention,
+    load_flydsl_attention_backends,
+    load_gluon_attention_backends,
+    load_gluon_v2_attention_backends,
+    load_gluon_v3_attention_backends,
+    load_turbo_attention_backends,
+    v4_attention_v1,
+    v4_attention_v2,
+    v4_csa_attention_v0,
+    v4_csa_attention_v1,
+    v4_csa_attention_v2,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import (
+    fused_rms_norm,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import (
+    apply_rope_from_positions,
+)
+
+_SUPPORTED_COMPRESS_RATIOS = (0, 4, 128)
+
+logger = logging.getLogger(__name__)
+
+
+def _require_gfx950() -> None:
+    """Assert the current device is gfx950 / CDNA4 before using the gluon backend.
+
+    The gluon sparse-MLA kernels are hand-tuned for gfx950 (MI350/MI355X); running
+    them on any other arch is unsupported. Called only when a layer selects
+    ``use_v4_attention_backend`` / ``use_v4_csa_attention_backend = 'gluon'``.
+    """
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' requires a "
+            "CUDA/HIP gfx950 (CDNA4) device, but no accelerator is available. Select "
+            "eager | triton_v1 | triton_v2 instead."
+        )
+    arch = str(getattr(torch.cuda.get_device_properties(0), "gcnArchName", ""))
+    if "gfx950" not in arch:
+        raise RuntimeError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' targets gfx950 / "
+            f"CDNA4 (MI350/MI355X); got device arch {arch!r}. Select eager | triton_v1 | triton_v2 "
+            "instead, or run on gfx950."
+        )
+
+
+# ---------------------------------------------------------------------------
+# P32 diagnostic: collect in-context cuda.Event timings of v4_attention_v1
+# ---------------------------------------------------------------------------
+
+
+class _DeepseekV4AttentionDiag:
+    """Accumulator for ``PRIMUS_V4_DIAG_TIME=1`` per-call timings."""
+
+    _per_mode: dict[str, list[float]] = collections.defaultdict(list)
+    _registered: bool = False
+    shape_logged: dict[str, bool] = {}
+
+    @classmethod
+    def record(cls, *, mode: str, ms: float, swa: int) -> None:
+        cls._per_mode[mode].append(ms)
+        if not cls._registered:
+            cls._registered = True
+            atexit.register(cls.dump)
+
+    @classmethod
+    def dump(cls) -> None:
+        if not cls._per_mode:
+            return
+        rank = os.environ.get("RANK", "0")
+        try:
+            local_rank = int(rank)
+        except (TypeError, ValueError):
+            local_rank = 0
+        if local_rank != 0:
+            return
+        print("[PRIMUS_V4_DIAG_TIME] v4_attention_v1 inline cuda.Event timings:", flush=True)
+        for mode, vs in cls._per_mode.items():
+            if not vs:
+                continue
+            # Drop first 3 to skip warmup.
+            stable = vs[3:] if len(vs) > 3 else vs
+            print(
+                f"  mode={mode:<6s}  n={len(vs):4d}  "
+                f"all_med={statistics.median(vs):7.3f}ms  "
+                f"warm_med={statistics.median(stable):7.3f}ms  "
+                f"warm_min={min(stable):7.3f}ms  "
+                f"warm_max={max(stable):7.3f}ms",
+                flush=True,
+            )
+
+
+# ---------------------------------------------------------------------------
+# Spec submodules — V4 (plan-2 / MLA-canonical)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DeepseekV4AttentionSubmodules:
+    """Spec submodules for the plan-2 :class:`DeepseekV4Attention`.
+
+    The names follow MLA's canonical layout where they overlap (so that
+    Megatron's standard tensor-parallel / sequence-parallel / TE machinery
+    can apply unchanged), plus V4-specific extras for the single-latent KV
+    and grouped low-rank O.
+
+    Provider-built shapes:
+
+    * ``linear_q_down_proj``  : ``hidden -> q_lora_rank``    (= ``wq_a``)
+    * ``q_layernorm``        : RMSNorm on ``q_lora_rank``    (= ``q_norm``)
+    * ``linear_q_up_proj``    : ``q_lora_rank -> n_heads * head_dim`` (= ``wq_b``)
+    * ``linear_kv``           : ``hidden -> head_dim``       (= ``wkv``,
+      single latent — broadcast to all heads)
+    * ``kv_layernorm``       : RMSNorm on ``head_dim``       (= ``kv_norm``)
+    * ``linear_o_a``         : ``(n_heads * head_dim / o_groups) -> o_groups * o_lora_rank``
+    * ``linear_o_b``         : ``o_groups * o_lora_rank -> hidden``
+    * ``compressor``         : :class:`Compressor` (compress_ratio > 0 only)
+    * ``indexer``            : :class:`Indexer`    (compress_ratio == 4 only)
+
+    When the spec provider supplies ``linear_proj`` (instead of grouped
+    ``linear_o_a`` / ``linear_o_b``) the attention falls back to MLA's
+    standard flat output projection — useful for unit tests and the
+    ``o_lora_rank == 0`` fast-path config.
+
+    Plan-3 P21: there is no ``attn_sink`` submodule slot.  The per-head
+    learnable sink is :class:`torch.nn.Parameter` ``self.attn_sink``
+    on the attention module itself (key ``layers.{i}.attn.attn_sink``
+    in the released checkpoint), and the softmax-with-sink combine is
+    inlined in :meth:`DeepseekV4Attention._attention_forward`.
+
+    Plan-3 P22: ``core_attention`` is the Turbo / TE flash-attention
+    kernel.  Only the dense layer kind (``compress_ratio == 0``) emits
+    a spec for this slot — HCA / CSA cannot use a stock flash-attn
+    kernel (HCA needs a joint sink across two key streams which would
+    require an LSE-returning flash kernel; CSA needs per-query top-K
+    indexed keys which is not a flash pattern).  When the dense path
+    receives a ``core_attention`` it bypasses the eager-Python softmax
+    and runs through ``provider.core_attention()`` instead.  When
+    ``provider.core_attention()`` returns
+    :class:`PrimusTurboAttention` (i.e. ``use_turbo_attention=True``)
+    and V4's ``attn_sink`` is on, the attention module aliases
+    ``core_attention.sinks`` to ``self.attn_sink`` so the released
+    checkpoint key path is preserved.
+    """
+
+    linear_q_down_proj: Optional[Union[ModuleSpec, type]] = None
+    linear_q_up_proj: Optional[Union[ModuleSpec, type]] = None
+    linear_kv: Optional[Union[ModuleSpec, type]] = None
+    linear_o_a: Optional[Union[ModuleSpec, type]] = None
+    linear_o_b: Optional[Union[ModuleSpec, type]] = None
+    linear_proj: Optional[Union[ModuleSpec, type]] = None  # fallback flat O
+    q_layernorm: Optional[Union[ModuleSpec, type]] = None
+    kv_layernorm: Optional[Union[ModuleSpec, type]] = None
+    compressor: Optional[Union[ModuleSpec, type]] = None
+    indexer: Optional[Union[ModuleSpec, type]] = None
+    # Plan-3 P22: dense (compress_ratio == 0) layers only.
+    core_attention: Optional[Union[ModuleSpec, type]] = None
+
+
+# ---------------------------------------------------------------------------
+# Build helpers
+# ---------------------------------------------------------------------------
+
+
+def _build_projection(
+    submodule: Optional[Union[ModuleSpec, type]],
+    *,
+    in_features: int,
+    out_features: int,
+) -> nn.Module:
+    """Build a linear projection from a spec submodule.
+
+    When the spec is ``None`` (CPU unit tests that exercise the
+    forward pass without a TP group) we instantiate a plain
+    :class:`nn.Linear` with the same shape.  When a spec is supplied
+    we delegate to :func:`build_module` and let any constructor
+    failure bubble up — Plan-3 P21 retired the ``try/except/return
+    nn.Linear`` fallback because it produced an unsharded model
+    (vanilla ``nn.Linear`` instead of column / row parallel shards)
+    that silently masked spec bugs at TP=1 and would diverge at TP>1.
+    """
+    if submodule is None:
+        return nn.Linear(in_features, out_features, bias=False)
+    return build_module(submodule)
+
+
+def _projection_forward(proj: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    """Run a projection and unwrap Megatron's ``(out, bias)`` tuple."""
+    out = proj(x)
+    if isinstance(out, tuple):
+        return out[0]
+    return out
+
+
+def _v4_o_a_fp8_enabled(config) -> bool:
+    """Whether to run the grouped-O ``o_a`` down-projection in MXFP8.
+
+    ``o_a`` is a per-group (batched) matmul done as a manual einsum on
+    ``linear_o_a.weight``, so it bypasses the fp8 linear path and stays bf16.
+    When PRIMUS_V4_FP8_ATTN_PROJ is set (and TP=1, where the surrounding
+    projections are already routed to fp8) and the layer is in turbo-fp8, run
+    it as per-group fp8 GEMMs instead. Default off.
+    """
+    if os.environ.get("PRIMUS_V4_FP8_ATTN_PROJ", "0") != "1":
+        return False
+    if getattr(config, "tensor_model_parallel_size", 1) != 1:
+        return False
+    try:
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboLowPrecisionGlobalStateManager as _M,
+        )
+
+        return _M.is_turbo_fp8_enabled()
+    except Exception:
+        return False
+
+
+def _fp8_grouped_o_a(attn_g: torch.Tensor, wo_a_w: torch.Tensor) -> torch.Tensor:
+    """Fused MXFP8 grouped-O ``o_a`` down-projection (replaces the bf16 einsum).
+
+    ``attn_g`` [B,S,G,d], ``wo_a_w`` [G,r,d] -> [B,S,G,r]. The G groups are an
+    independent batched matmul ``[B*S,d] @ [d,r]`` per group; we run them as a
+    SINGLE ``grouped_gemm_fp8`` (one fused Triton launch) instead of a per-group
+    ``gemm_fp8`` loop (G launches + G× quant). Stack tokens group-major into
+    ``[G*B*S, d]`` with ``group_lens=[B*S]*G``; weight ``[G,d,r]`` (trans_b=False).
+    d=(H*head_dim)/G and B*S are multiples of 32, so the MX block scale is clean.
+    """
+    import primus_turbo.pytorch as pt
+
+    from primus.backends.megatron.core.extensions.primus_turbo import (
+        PrimusTurboLowPrecisionGlobalStateManager as _M,
+    )
+
+    cfg = _M.get_turbo_quant_config().data()
+    B, S, G, d = attn_g.shape
+    r = wo_a_w.shape[1]
+    a = attn_g.permute(2, 0, 1, 3).reshape(G * B * S, d).contiguous()  # [G*BS, d], group-major
+    b = wo_a_w.transpose(1, 2).contiguous()  # [G, d, r]  (K=d, N=r, trans_b=False)
+    group_lens = torch.full((G,), B * S, dtype=torch.int64, device=a.device)
+    out = pt.ops.grouped_gemm_fp8(a, b, group_lens, trans_b=False, config=cfg)  # [G*BS, r]
+    return out.reshape(G, B, S, r).permute(1, 2, 0, 3)  # [B, S, G, r]
+
+
+def _coerce_optional_bool_flag(value: object, *, field_name: str) -> bool:
+    """Coerce a possibly-stringified yaml flag to a clean ``bool``.
+
+    Yaml interpolation like ``${PRIMUS_FOO:false}`` resolves to the
+    STRING ``"false"`` when the env var is unset, and the naive
+    ``bool("false") is True`` would silently flip a default-off knob
+    to on.  Accept the common string spellings explicitly and treat
+    everything else as truthy/falsy via the normal ``bool(...)`` rule.
+
+    Plan-8 P57 close-out 2 added this helper for the new
+    ``use_v4_tilelang_*`` flags; existing flags
+    (``use_v4_triton_*`` / ``use_v4_compiled_sinkhorn``) avoid the
+    issue because the V4 run scripts always pass them via
+    ``--<flag> "False"`` and the override parser coerces to a Python
+    ``False`` before the config ever sees a string.
+    """
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in ("0", "false", "no", "off", ""):
+            return False
+        if lowered in ("1", "true", "yes", "on"):
+            return True
+        raise ValueError(
+            f"Unrecognised string value for boolean config flag "
+            f"{field_name!r}: {value!r}; expected one of "
+            "'true' / 'false' / '1' / '0' / 'yes' / 'no' / 'on' / 'off'."
+        )
+    return bool(value)
+
+
+def _per_head_rms_norm(x: torch.Tensor, *, eps: float) -> torch.Tensor:
+    """Parameter-less per-head RMSNorm.
+
+    Mirrors the released ``inference/model.py`` reference:
+
+    .. code-block:: python
+
+        q_rms = torch.rsqrt(q.float().square().mean(-1, keepdim=True) + eps)
+        q     = (q.float() * q_rms).to(q.dtype)
+
+    There is no learnable ``gamma`` — the per-head scale is "absorbed"
+    into the surrounding ``linear_q_up_proj`` weights at training time.
+    The check confirmed the released checkpoint has no separate
+    ``q_rms.weight`` parameter.
+
+    Small-kernel-fusion (2026-07-03): the eager chain (bf16->fp32 cast +
+    square + mean + rsqrt + mul + fp32->bf16 cast, ~6 kernels / call ×
+    8 attention layers) is collapsed into one Triton FWD + one BWD kernel
+    via :func:`fused_rms_norm` (parameter-less, ``out_dtype = in_dtype``).
+    Gated by ``PRIMUS_RMSNORM_TRITON`` (default on); the dispatcher falls
+    back to the bit-identical eager body on CPU / when the knob is off.
+    """
+    return fused_rms_norm(x, None, eps=eps, mid_cast=False, out_dtype=x.dtype)
+
+
+def _build_local_rms_norm(dim: int, *, eps: float) -> nn.Module:
+    """Tiny CPU-friendly RMSNorm used as a fallback when no spec is given.
+
+    Plan-2 P17 retired the closure-built ``_RMSNorm`` helper here; the
+    canonical implementation lives in
+    :class:`primus.backends.megatron.core.transformer.local_rmsnorm.LocalRMSNorm`
+    so the same code path is shared with
+    :mod:`...deepseek_v4_block` and :mod:`...compressor`.
+    """
+    return LocalRMSNorm(dim=dim, eps=eps)
+
+
+# ---------------------------------------------------------------------------
+# DeepseekV4Attention (faithful, MLA-rooted, dense + CSA + HCA)
+# ---------------------------------------------------------------------------
+
+
+class DeepseekV4Attention(KeepInFp32Mixin, MLASelfAttention):
+    """V4 attention faithful to the released ``DeepSeek-V4-Flash`` checkpoint.
+
+    Subclasses :class:`MLASelfAttention` for type identity (so downstream
+    Megatron isinstance checks treat V4 attention as an MLA variant) but
+    overrides ``__init__`` and ``forward`` because V4's parameter layout
+    differs from MLA's compressed-KV form:
+
+    * V4 has **no** ``linear_kv_down_proj`` / ``linear_kv_up_proj`` — the
+      KV is single-latent (``wkv``) and shared as both K and V.
+    * V4's ``linear_proj`` is replaced by grouped low-rank
+      ``linear_o_a`` / ``linear_o_b`` (when ``config.o_lora_rank > 0``).
+    * V4 adds a per-head parameter-less ``q_rms`` and a learnable
+      ``attn_sink``.
+    * V4 layers come in three flavours selected by ``compress_ratio``:
+
+      * ``0``   — dense / SWA over local KV.
+      * ``128`` — HCA: local SWA *plus* a fully-visible compressed pool
+        (Compressor in non-overlap mode).
+      * ``4``   — CSA: local SWA *plus* a per-query top-K selection over
+        a compressed pool (Compressor in overlap mode + Indexer).
+
+    Because the parent's ``__init__`` builds modules we don't want, we
+    skip the MLA / Attention init chain and call ``nn.Module.__init__``
+    directly. V4-shape modules are built from the spec submodules.
+
+    **Plan-4 P27 — kernel dispatch precedence.**
+
+    The softmax-and-attend kernel each layer fires through is selected
+    in :meth:`forward` / :meth:`_csa_forward` based on three independent
+    config flags (``use_turbo_attention``, ``use_v4_triton_attention``,
+    ``use_v4_triton_csa_attention``).  The layer-kind-specific
+    precedence is:
+
+    .. code-block:: text
+
+        compress_ratio == 0  (dense / SWA, single key axis):
+            use_turbo_attention      > use_v4_triton_attention > eager
+            (-> self.core_attention)   (-> v4_attention_v1)         (-> _attention_forward)
+
+        compress_ratio == 128  (HCA: local SWA + full compressed pool):
+            use_v4_triton_attention  > eager
+            (-> v4_attention_v1,          (-> _attention_forward
+             HCA path with joint        with cat([local, pool])
+             [local | pool] mask)       additive mask)
+
+            ``use_turbo_attention`` does NOT route HCA — Turbo's
+            flash-attn returns no LSE so the joint local+pool softmax
+            cannot be decomposed into two flash calls.
+
+        compress_ratio == 4  (CSA: local SWA + per-query top-K gather):
+            use_v4_triton_csa_attention  > eager
+            (-> v4_csa_attention_v0)          (-> _csa_forward eager)
+
+            Neither ``use_turbo_attention`` nor
+            ``use_v4_triton_attention`` applies to CSA — the per-query
+            top-K gather (``gathered = pool[..., topk_idxs, :]``) is
+            sparse-per-row indexed attention with no flash-attn
+            equivalent.
+
+    Auto-disable rules (init-side, fail-loud):
+
+    * ``use_v4_triton_attention=True`` + ``compress_ratio == 4`` →
+      auto-disabled (CSA layers must opt in via the separate flag).
+    * ``use_v4_triton_csa_attention=True`` + ``compress_ratio != 4`` →
+      auto-disabled (the dense / HCA flag is ``use_v4_triton_attention``).
+
+    On rank 0 each ``__init__`` emits one ``INFO`` log line through
+    :meth:`_log_kernel_choice` summarising the dispatch outcome for
+    the layer (e.g. ``[V4-attn] Layer 17: cr=128, kernel = v4_attention_v1
+    (Triton, HCA path)``) so smoke / training logs unambiguously show
+    which kernel each layer is firing through.
+    """
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        *,
+        rope: DualRoPE,
+        compress_ratio: int = 0,
+        submodules: Optional[DeepseekV4AttentionSubmodules] = None,
+        layer_number: Optional[int] = None,
+        pg_collection=None,
+        attn_mask_type=None,
+        **kwargs,
+    ) -> None:
+        # We deliberately bypass the MLA / Attention parent __init__ chain
+        # because V4's KV layout differs from MLA's compressed-KV form.
+        # The class still subclasses MLASelfAttention for type identity so
+        # that ``isinstance(layer.self_attention, MLASelfAttention)`` keeps
+        # working in the Megatron stack.
+        #
+        # Plan-2 P16: ``attn_mask_type`` is accepted (and ignored) so the
+        # attention spec can declare a value that satisfies upstream
+        # :class:`MultiTokenPredictionLayer`'s pre-build validator; V4
+        # manages its own SWA / sink mask internally. ``**kwargs`` swallows
+        # any forward-compatible kwargs upstream may add (e.g.
+        # ``cp_comm_type``) so the spec lifecycle keeps working.
+        del attn_mask_type, kwargs
+        nn.Module.__init__(self)
+
+        if compress_ratio not in _SUPPORTED_COMPRESS_RATIOS:
+            raise ValueError(
+                f"DeepseekV4Attention supports compress_ratio in "
+                f"{_SUPPORTED_COMPRESS_RATIOS} (got {compress_ratio})."
+            )
+
+        hidden_size = int(config.hidden_size)
+        num_heads = int(config.num_attention_heads)
+        head_dim = int(config.kv_channels)
+        rotary_dim = int(config.qk_pos_emb_head_dim)
+        attn_sliding_window = int(config.attn_sliding_window)
+        attn_sink_enabled = bool(config.attn_sink)
+        attn_dropout = float(config.attention_dropout)
+        norm_eps = float(getattr(config, "norm_epsilon", None) or config.layernorm_epsilon)
+        q_lora_rank = int(config.q_lora_rank or 0)
+        o_groups = int(getattr(config, "o_groups", 1))
+        o_lora_rank = int(getattr(config, "o_lora_rank", 0))
+
+        if q_lora_rank <= 0:
+            # V4 always uses a Q LoRA path; drop the no-LoRA branch to
+            # keep the math aligned with the checkpoint.
+            raise ValueError(
+                "DeepseekV4Attention requires config.q_lora_rank > 0; "
+                "V4 always low-rank-projects Q via wq_a / wq_b."
+            )
+
+        if num_heads * head_dim % max(o_groups, 1) != 0:
+            raise ValueError(
+                f"num_heads * head_dim ({num_heads * head_dim}) must be divisible "
+                f"by o_groups ({o_groups})"
+            )
+
+        self.config = config
+        self.compress_ratio = int(compress_ratio)
+        self.layer_number = int(layer_number) if layer_number is not None else 0
+        self.pg_collection = pg_collection
+
+        # ---- shape fields (read by helpers in this class) ----
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_attention_heads_per_partition = num_heads
+        self.num_query_groups_per_partition = 1  # single-latent KV
+        self.head_dim = head_dim
+        self.rotary_dim = rotary_dim
+        self.q_head_dim = head_dim  # MLA convention; here qk_head_dim + qk_pos_emb_head_dim == head_dim
+        self.attn_sliding_window = attn_sliding_window
+        self.attn_dropout = attn_dropout
+        self.q_lora_rank = q_lora_rank
+        self.o_groups = max(o_groups, 1)
+        self.o_lora_rank = o_lora_rank
+        self.norm_eps = norm_eps
+
+        # Shared dual-RoPE (held by reference; not registered to avoid
+        # double-counting parameters across attention layers).
+        self._rope = [rope]
+
+        submodules = submodules or DeepseekV4AttentionSubmodules()
+        self._submodules = submodules
+
+        # ---- Q branch: hidden -> q_lora_rank -> n_heads * head_dim ----
+        self.linear_q_down_proj = _build_projection(
+            submodules.linear_q_down_proj,
+            in_features=hidden_size,
+            out_features=q_lora_rank,
+        )
+        if submodules.q_layernorm is None:
+            self.q_layernorm = _build_local_rms_norm(q_lora_rank, eps=norm_eps)
+        else:
+            self.q_layernorm = build_module(
+                submodules.q_layernorm,
+                hidden_size=q_lora_rank,
+                config=config,
+                eps=norm_eps,
+            )
+        self.linear_q_up_proj = _build_projection(
+            submodules.linear_q_up_proj,
+            in_features=q_lora_rank,
+            out_features=num_heads * head_dim,
+        )
+
+        # ---- KV branch: single-latent ``wkv`` ----
+        self.linear_kv = _build_projection(
+            submodules.linear_kv,
+            in_features=hidden_size,
+            out_features=head_dim,
+        )
+        if submodules.kv_layernorm is None:
+            self.kv_layernorm = _build_local_rms_norm(head_dim, eps=norm_eps)
+        else:
+            self.kv_layernorm = build_module(
+                submodules.kv_layernorm,
+                hidden_size=head_dim,
+                config=config,
+                eps=norm_eps,
+            )
+
+        # ---- O projection ----
+        # Two paths:
+        #   - Grouped low-rank (V4 release): linear_o_a + linear_o_b
+        #   - Flat MLA-style: linear_proj (used when o_lora_rank == 0)
+        if o_lora_rank > 0:
+            n_per_group = num_heads * head_dim // self.o_groups
+            self.linear_o_a = _build_projection(
+                submodules.linear_o_a,
+                in_features=n_per_group,
+                out_features=self.o_groups * o_lora_rank,
+            )
+            self.linear_o_b = _build_projection(
+                submodules.linear_o_b,
+                in_features=self.o_groups * o_lora_rank,
+                out_features=hidden_size,
+            )
+            self.linear_proj = None
+        else:
+            self.linear_o_a = None
+            self.linear_o_b = None
+            self.linear_proj = _build_projection(
+                submodules.linear_proj,
+                in_features=num_heads * head_dim,
+                out_features=hidden_size,
+            )
+
+        # ---- attention sink ----
+        # The released checkpoint stores ``attn_sink`` as a [num_heads]
+        # learnable parameter directly on the attention module (key
+        # ``layers.{i}.attn.attn_sink`` — no wrapping submodule).
+        # We register it as ``self.attn_sink`` so the state-dict key
+        # matches the released checkpoint exactly; the softmax-with-sink
+        # combine is inlined in :meth:`_attention_forward`.
+        #
+        # Plan-3 P21 retired the optional ``self.attn_sink_module``
+        # build branch (and the ``submodules.attn_sink`` slot) — the
+        # branch was never exercised in the forward path and its
+        # ``try/except`` masked AttentionSink build failures.  A future
+        # TE-fused sink primitive can land as a new spec field once it
+        # actually replaces the inline path.
+        #
+        # FP32 in the released checkpoint, and the gluon / flydsl_v1 kernels
+        # assert ``attn_sink.dtype == torch.float32`` at their entry. That
+        # assertion is satisfied by the promotion the callers already do
+        # (``sink.float()`` in the sparse-MLA adapter and the eager reference),
+        # so the parameter itself may follow the model dtype. Pinning it to FP32
+        # is therefore opt-in via ``PRIMUS_V4_KEEP_FP32`` -- see
+        # ``keep_in_fp32`` for why holding a second parameter dtype is not free.
+        if attn_sink_enabled:
+            self.attn_sink = nn.Parameter(torch.zeros(num_heads, dtype=torch.float32))
+            mark_keep_in_fp32(self.attn_sink)
+        else:
+            self.register_parameter("attn_sink", None)
+
+        # ---- compressor / indexer (compressed branches only) ----
+        self.compressor: Optional[nn.Module] = None
+        self.indexer: Optional[nn.Module] = None
+        # Last indexer distillation loss (detached), for logging. None until a
+        # training step runs with the loss enabled.
+        self.last_indexer_distill_loss: Optional[torch.Tensor] = None
+        if self.compress_ratio > 0:
+            self.compressor = self._build_compressor(submodules.compressor)
+            if self.compress_ratio == 4:
+                self.indexer = self._build_indexer(submodules.indexer)
+                # ``topk`` is not differentiable and the forward only consumes
+                # ``topk_idxs``, so the Indexer can only learn through the
+                # distillation loss (see ``indexer_distill_loss``). Without it,
+                # leaving the params trainable inserts permanently-dead params
+                # into the distributed-optimizer grad buckets: that wastes grad /
+                # optimizer state plus cross-node grad-sync bandwidth AND trips
+                # Megatron's overlap_grad_reduce invariant (every bucket param
+                # must fire its grad-ready backward hook), which is what forced
+                # overlap_grad_reduce / param_gather OFF and crippled cross-node
+                # DP scaling. So freeze unless the loss is actually enabled.
+                # PRIMUS_V4_INDEXER_TRAINABLE=1 still forces trainable.
+                if not self.indexer_distill_enabled and (
+                    os.environ.get("PRIMUS_V4_INDEXER_TRAINABLE", "0") != "1"
+                ):
+                    for _indexer_param in self.indexer.parameters():
+                        _indexer_param.requires_grad_(False)
+
+        # ---- core attention (Turbo / TE flash) — dense layers only ----
+        # Plan-3 P22: when the spec emits a ``core_attention`` submodule
+        # (only on dense ``compress_ratio == 0`` layers), build it now and
+        # use it as the softmax-and-attend kernel instead of the
+        # eager-Python ``_attention_forward``.  HCA + CSA always run
+        # eager-Python because their joint softmax / per-query top-K
+        # gather can't be expressed as a stock flash-attn call (see
+        # comments in ``forward`` / ``_csa_forward``).
+        #
+        # The constant ``softmax_scale`` is precomputed via
+        # ``_attention_scale()`` (the YaRN ``m_scale`` is a layer-static
+        # constant set at RoPE init time, so this matches the eager-path
+        # scale exactly for ``compress_ratio == 0``).
+        # Plan-4 P25: in-tree Primus Triton kernel for cr ∈ {0, 128}.
+        # Read the config flag once at __init__ so ``forward`` only does
+        # a cheap attribute load. Precedence in ``forward`` is
+        # ``use_turbo_attention > use_v4_triton_attention > eager``.
+        # ---- attention backend selection (unified string selectors) ----
+        # ``use_v4_attention_backend`` selects the dense (cr=0) / HCA (cr=128)
+        # kernel; ``use_v4_csa_attention_backend`` selects the CSA (cr=4) kernel.
+        # ``use_turbo_attention`` (built as ``core_attention`` below) still takes
+        # precedence for the dense path when it can be built.
+        _ATTN_BACKENDS = (
+            "eager",
+            "triton_v1",
+            "triton_v2",
+            "gluon",
+            "gluon_v2",
+            "gluon_v3",
+            "flydsl_v1",
+            "turbo",
+        )
+        _CSA_BACKENDS = (
+            "eager",
+            "triton_v0",
+            "triton_v1",
+            "triton_v2",
+            "gluon",
+            "gluon_v2",
+            "gluon_v3",
+            "flydsl_v0",
+            "flydsl_v1",
+            "turbo",
+        )
+        self._attn_backend: str = str(getattr(config, "use_v4_attention_backend", "triton_v1") or "triton_v1")
+        self._csa_backend: str = str(
+            getattr(config, "use_v4_csa_attention_backend", "triton_v1") or "triton_v1"
+        )
+        if self._attn_backend not in _ATTN_BACKENDS:
+            raise ValueError(
+                f"use_v4_attention_backend={self._attn_backend!r} is not a valid dense/HCA backend; "
+                f"expected one of {_ATTN_BACKENDS}"
+            )
+        if self._csa_backend not in _CSA_BACKENDS:
+            raise ValueError(
+                f"use_v4_csa_attention_backend={self._csa_backend!r} is not a valid CSA backend; "
+                f"expected one of {_CSA_BACKENDS}"
+            )
+
+        # gluon is a gfx950/CDNA4-only backend with a hard triton.experimental.gluon
+        # dependency. Load it (and validate the arch) ONLY when a layer actually
+        # selects it, so non-gluon backends never pay the import and never crash on
+        # unsupported hardware / Triton builds. The loader raises a clear ImportError
+        # if the dependency is missing; ``_require_gfx950`` raises if the arch is wrong.
+        self._v4_attention_gluon = None
+        self._v4_csa_attention_gluon = None
+        if "gluon" in (self._attn_backend, self._csa_backend):
+            _require_gfx950()
+            self._v4_attention_gluon, self._v4_csa_attention_gluon = load_gluon_attention_backends()
+
+        # gluon_v2 (2nd-gen gluon fwd+bwd) — same gfx950-only lazy-load contract as gluon.
+        self._v4_attention_gluon_v2 = None
+        self._v4_csa_attention_gluon_v2 = None
+        if "gluon_v2" in (self._attn_backend, self._csa_backend):
+            _require_gfx950()
+            self._v4_attention_gluon_v2, self._v4_csa_attention_gluon_v2 = load_gluon_v2_attention_backends()
+
+        # gluon_v3 (3rd-gen optimized gluon fwd+bwd) — same gfx950-only lazy-load contract.
+        self._v4_attention_gluon_v3 = None
+        self._v4_csa_attention_gluon_v3 = None
+        if "gluon_v3" in (self._attn_backend, self._csa_backend):
+            _require_gfx950()
+            self._v4_attention_gluon_v3, self._v4_csa_attention_gluon_v3 = load_gluon_v3_attention_backends()
+
+        # flydsl_v1 (native FlyDSL MFMA) is likewise a gfx950/CDNA4-only backend with
+        # a hard `flydsl` pip dependency; load + arch-validate only when selected.
+        self._v4_attention_flydsl = None
+        self._v4_csa_attention_flydsl = None
+        if "flydsl_v1" in (self._attn_backend, self._csa_backend):
+            _require_gfx950()
+            self._v4_attention_flydsl, self._v4_csa_attention_flydsl = load_flydsl_attention_backends()
+
+        # turbo (Primus-Turbo native-FlyDSL sparse-MLA via the turbo API) — same gfx950-only
+        # lazy-load contract; hard-depends on the installed primus_turbo (flydsl attention) + flydsl.
+        self._v4_attention_turbo = None
+        self._v4_csa_attention_turbo = None
+        if "turbo" in (self._attn_backend, self._csa_backend):
+            _require_gfx950()
+            self._v4_attention_turbo, self._v4_csa_attention_turbo = load_turbo_attention_backends()
+
+        self.core_attention: Optional[nn.Module] = None
+        self._use_core_attention: bool = False
+        if submodules.core_attention is not None and self.compress_ratio == 0:
+            softmax_scale = self._attention_scale()
+            try:
+                self.core_attention = build_module(
+                    submodules.core_attention,
+                    config=config,
+                    layer_number=self.layer_number,
+                    attn_mask_type=AttnMaskType.causal,
+                    attention_type="self",
+                    softmax_scale=softmax_scale,
+                    k_channels=head_dim,
+                    v_channels=head_dim,
+                    cp_comm_type="p2p",
+                    pg_collection=self.pg_collection,
+                )
+            except TypeError:
+                # Some core-attention classes (e.g. local CPU stubs in
+                # unit tests) don't accept the full TE / Turbo kwarg set.
+                # Retry with the minimal kwargs Megatron ships everywhere.
+                self.core_attention = build_module(
+                    submodules.core_attention,
+                    config=config,
+                    layer_number=self.layer_number,
+                    attn_mask_type=AttnMaskType.causal,
+                    attention_type="self",
+                    softmax_scale=softmax_scale,
+                )
+
+            # Sink alias: when V4's per-head learnable sink is on AND the
+            # core-attention class supports learned sinks (Turbo only —
+            # the TE class does not), tie ``core_attention.sinks`` to
+            # ``self.attn_sink`` so the released-checkpoint key
+            # ``layers.{i}.attn.attn_sink`` keeps loading.  TE classes
+            # that don't expose ``use_sink_attention`` get ``False`` here
+            # and we fall back to eager-Python so the inline
+            # softmax-with-sink path still produces the right math.
+            core_use_sink = bool(getattr(self.core_attention, "use_sink_attention", False))
+            if attn_sink_enabled and core_use_sink:
+                # Aliasing means Turbo and V4 share one Parameter object, so the
+                # sink has to carry the dtype Turbo allocated for its own
+                # ``sinks``. That is the one configuration where the sink cannot
+                # stay FP32, so drop the keep-in-FP32 mark as well -- otherwise
+                # the model-wide dtype conversion would silently undo the cast
+                # and hand Turbo a tensor it did not allocate for. The eager
+                # path promotes to float32 inside the softmax either way, so
+                # only the stored resolution differs.
+                turbo_sinks = getattr(self.core_attention, "sinks", None)
+                if turbo_sinks is not None and turbo_sinks.dtype != self.attn_sink.dtype:
+                    unmark_keep_in_fp32(self.attn_sink)
+                    self.attn_sink.data = self.attn_sink.data.to(turbo_sinks.dtype)
+                self.core_attention.sinks = self.attn_sink
+                self._use_core_attention = True
+            elif not attn_sink_enabled and self.core_attention is not None:
+                # No-sink V4 still uses core_attention (e.g. unit tests,
+                # ablations).  SWA is honored by Turbo only when sinks are
+                # on, so we accept this only when SWA is off too.
+                if self.attn_sliding_window <= 0:
+                    self._use_core_attention = True
+
+        # Plan-4 P27: surface the dispatch outcome in the training log
+        # so smoke / debug logs unambiguously show which kernel each
+        # layer is firing through (precedence is documented in the
+        # class docstring).  Rank-0 only — every rank's own
+        # per-rank-file already captures the right entries.
+        self._log_kernel_choice()
+
+    # ------------------------------------------------------------------
+    # construction helpers (compressed branches)
+    # ------------------------------------------------------------------
+
+    def _log_kernel_choice(self) -> None:
+        """Emit one ``INFO`` log line summarising this layer's kernel choice.
+
+        Plan-4 P27.  Resolves the precedence outcome captured by
+        :meth:`forward` / :meth:`_csa_forward` (see class docstring) and
+        logs it once at ``__init__`` time so smoke / training logs
+        unambiguously show which kernel is firing for each layer.
+
+        Rank-0 only when distributed; in single-process unit tests the
+        log fires unconditionally so ``caplog.at_level(logging.INFO)``
+        captures it.  We cannot use Megatron's ``print_rank_0`` here
+        because this module is also imported in CPU-only unit tests
+        where Megatron's parallel-state isn't initialised.
+        """
+        try:
+            dist_initialized = torch.distributed.is_available() and torch.distributed.is_initialized()
+        except Exception:
+            dist_initialized = False
+        if dist_initialized and torch.distributed.get_rank() != 0:
+            return
+
+        if self.compress_ratio == 0:
+            if self._use_core_attention:
+                kernel = "core_attention (Turbo / TE flash)"
+            else:
+                kernel = f"dense attention backend = {self._attn_backend}"
+        elif self.compress_ratio == 128:
+            kernel = f"HCA attention backend = {self._attn_backend}"
+        elif self.compress_ratio == 4:
+            kernel = f"CSA attention backend = {self._csa_backend}"
+        else:
+            # Defensive: __init__ already raises ValueError for unsupported
+            # compress_ratio, so this branch should be unreachable.
+            kernel = f"<unknown compress_ratio={self.compress_ratio}>"
+
+        logger.info(
+            "[V4-attn] Layer %s: cr=%s, kernel = %s",
+            self.layer_number,
+            self.compress_ratio,
+            kernel,
+        )
+
+    def _build_compressor(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module:
+        """Build the V4 :class:`Compressor` for compressed branches.
+
+        Plan-1 conventions (kept under V4): ``ratio=4`` → overlap mode
+        (CSA), ``ratio=128`` → non-overlap mode (HCA). The released
+        checkpoint hard-codes ``coff=2`` for overlap (CSA) and ``coff=1``
+        for non-overlap (HCA); :class:`Compressor` enforces this through
+        its own ``overlap`` argument.
+
+        When the spec is ``None`` (CPU unit tests, ``DeepseekV4Attention``
+        constructed without a layer spec) we instantiate the local
+        :class:`Compressor` directly.  Otherwise we delegate to
+        :func:`build_module` and let any constructor failure bubble up —
+        Plan-3 P21 retired the ``try/except/local Compressor`` fallback
+        because the spec passes the same :class:`Compressor` class and
+        the fallback handler was dead code that masked real spec bugs.
+        """
+        kwargs = dict(
+            hidden_size=self.hidden_size,
+            head_dim=self.head_dim,
+            ratio=self.compress_ratio,
+            overlap=(self.compress_ratio == 4),
+        )
+        if spec is None:
+            return Compressor(**kwargs)
+        return build_module(spec, **kwargs)
+
+    def _build_indexer(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module:
+        """Build the V4 :class:`Indexer` for the CSA branch.
+
+        See :meth:`_build_compressor` for the spec-vs-fallback contract.
+        Plan-3 P21 retired the ``try/except/local Indexer`` fallback for
+        the same reason.
+        """
+        index_topk = int(self.config.index_topk)
+        index_head_dim = int(self.config.index_head_dim)
+        index_n_heads = int(self.config.index_n_heads)
+        kwargs = dict(
+            hidden_size=self.hidden_size,
+            index_head_dim=index_head_dim,
+            index_n_heads=index_n_heads,
+            index_topk=index_topk,
+            compress_ratio=self.compress_ratio,
+            use_fp8_qk=bool(getattr(self.config, "use_v4_fp8_indexer", False)),
+            # The indexer rotates its own Q / K with this layer's compressed
+            # RoPE before scoring, so it needs the same cache the compressed
+            # pool uses. Passed by reference; the Indexer does not register it.
+            rope=self.rope.get_rope(compress_ratio=self.compress_ratio),
+            rotary_dim=self.rotary_dim,
+        )
+        if spec is None:
+            return Indexer(**kwargs)
+        return build_module(spec, **kwargs)
+
+    # ------------------------------------------------------------------
+    # internals
+    # ------------------------------------------------------------------
+
+    @property
+    def rope(self) -> DualRoPE:
+        return self._rope[0]
+
+    @property
+    def indexer_distill_coeff(self) -> float:
+        """Coefficient of the indexer distillation loss (``0`` = disabled)."""
+        return float(getattr(self.config, "v4_indexer_distill_loss_coeff", 0.0) or 0.0)
+
+    @property
+    def indexer_distill_enabled(self) -> bool:
+        """True when this layer trains its indexer via the distillation loss."""
+        return self.compress_ratio == 4 and self.indexer_distill_coeff > 0.0
+
+    def _indexer_loss_head_group(self):
+        """Group the attention heads are sharded over, or ``None`` if they are not.
+
+        ``linear_q_up_proj`` is column-parallel with ``gather_output=True``, so
+        every TP rank holds all heads and the distillation loss's head sum is
+        already complete -- unlike the open-source reference, whose attention
+        keeps them sharded and therefore all-reduces. Resolve the group anyway
+        when the shapes say the heads *are* sharded, so dropping the gather
+        cannot silently turn the target distribution into a partial sum.
+        """
+        if self.num_attention_heads_per_partition == int(self.config.num_attention_heads):
+            return None
+        group = getattr(self.pg_collection, "tp", None)
+        if group is not None:
+            return group
+        from megatron.core import parallel_state
+
+        return parallel_state.get_tensor_model_parallel_group()
+
+    def _log_indexer_distill_loss(self, device: torch.device) -> None:
+        """Report this layer's indexer loss, zero on layers without an indexer.
+
+        Every V4 layer reports so the tracker key exists on every pipeline rank
+        -- see :func:`log_indexer_distill_loss`.
+        """
+        if self.indexer_distill_coeff <= 0.0:
+            return
+        if not (self.training and torch.is_grad_enabled()):
+            # Matches the condition the loss is computed under, so a stale
+            # value from an earlier step can never be reported.
+            return
+        mtp_layers = getattr(self.config, "mtp_num_layers", 0) or 0
+        log_indexer_distill_loss(
+            self.last_indexer_distill_loss,
+            layer_number=self.layer_number,
+            num_layers=int(self.config.num_layers) + int(mtp_layers),
+            device=device,
+        )
+
+    def _attention_scale(self) -> float:
+        """Softmax temperature for every branch: plain ``1 / sqrt(head_dim)``.
+
+        V4 keeps attention logits in range via the Q / KV RMSNorms, so the
+        YaRN magnitude factor (``m_scale``) is deliberately NOT folded in here
+        -- ``inference/model.py`` uses ``self.softmax_scale = head_dim ** -0.5``
+        for both the core attention and the indexer. The YaRN *frequency*
+        interpolation on ``inv_freq`` is unaffected.
+        """
+        return 1.0 / math.sqrt(self.head_dim)
+
+    def _apply_q(self, hidden: torch.Tensor) -> torch.Tensor:
+        """``[B, S, D]`` → ``[B, S, H, head_dim]`` (Q after q_norm + q_rms)."""
+        q_compressed = _projection_forward(self.linear_q_down_proj, hidden)
+        q_compressed = self.q_layernorm(q_compressed)
+        q = _projection_forward(self.linear_q_up_proj, q_compressed)
+        B, S, _ = q.shape
+        q = q.view(B, S, self.num_heads, self.head_dim)
+        # Per-head parameter-less RMS (matches `inference/model.py`).
+        q = _per_head_rms_norm(q, eps=self.norm_eps)
+        return q
+
+    def _apply_kv(self, hidden: torch.Tensor) -> torch.Tensor:
+        """``[B, S, D]`` → ``[B, S, 1, head_dim]`` (single-latent K = V)."""
+        kv = _projection_forward(self.linear_kv, hidden)
+        kv = self.kv_layernorm(kv)
+        B, S, _ = kv.shape
+        return kv.view(B, S, 1, self.head_dim)
+
+    def _apply_rope_q_k(self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor):
+        """Apply partial RoPE (last ``rotary_dim`` channels) to Q and K
+        using the LAYER's compress_ratio (so CSA/HCA use the compress base
+        + YaRN; dense uses the main base)."""
+        q = self.rope.apply_rope(q, position_ids=position_ids, compress_ratio=self.compress_ratio)
+        k = self.rope.apply_rope(k, position_ids=position_ids, compress_ratio=self.compress_ratio)
+        return q, k
+
+    def _local_mask(self, S: int, *, device, dtype) -> torch.Tensor:
+        """Mask for the local (SWA or full causal) branch.
+
+        ``attn_sliding_window > 0`` enables sliding-window; ``0`` (the
+        default for unit tests / configs without SWA) gives full causal.
+        """
+        window = self.attn_sliding_window if self.attn_sliding_window > 0 else 0
+        if window > 0:
+            return sliding_window_causal_mask(S, window, device=device, dtype=dtype)
+        return sliding_window_causal_mask(S, S, device=device, dtype=dtype)
+
+    def _append_sink_softmax(self, logits: torch.Tensor) -> torch.Tensor:
+        """Numerically-stable softmax with optional virtual-sink column.
+
+        ``logits`` shape is ``[B, H, ..., Sk]`` — the head axis is at
+        ``dim=1``. Returns probabilities on the *real* keys (sink column
+        dropped) of the same shape as ``logits``.
+        """
+        if self.attn_sink is None:
+            logits = logits - logits.amax(dim=-1, keepdim=True).detach()
+            return logits.softmax(dim=-1)
+
+        # Build a sink column that broadcasts over all dims except the
+        # head axis (dim=1) and the key axis (dim=-1).
+        ndim = logits.dim()
+        view_shape = [1] * ndim
+        view_shape[1] = self.num_heads
+        view_shape[-1] = 1
+        target_shape = list(logits.shape[:-1]) + [1]
+        sink_col = self.attn_sink.float().view(*view_shape).expand(*target_shape)
+        logits_aug = torch.cat([logits, sink_col], dim=-1)
+        logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach()
+        probs = logits_aug.softmax(dim=-1)
+        return probs[..., :-1]
+
+    def _attention_forward(
+        self,
+        q: torch.Tensor,  # [B, H, Sq, head_dim]
+        k: torch.Tensor,  # [B, H, Sk, head_dim]
+        v: torch.Tensor,  # [B, H, Sk, head_dim]
+        attn_mask: torch.Tensor,  # [Sq, Sk] additive (broadcasts over B,H)
+    ) -> torch.Tensor:
+        """Eager scaled-dot-product attention with optional attn_sink
+        for the dense / HCA paths (single key axis).
+
+        Plan-4 P24: math lives in
+        :func:`primus...v4_attention_kernels._eager.reference.eager_v4_attention`
+        so the dense / HCA path, the plan-4 Triton kernel (P25), and the
+        plan-4 unit-test harness share one definition. The caller has
+        already pre-built the ``[Sq, Sk]`` additive mask (SWA-causal
+        for dense, ``cat([local_mask, hca_mask])`` for HCA) so we pass
+        ``swa_window=0`` and let the reference op use the supplied
+        ``additive_mask`` directly — bit-identical to the pre-P24
+        inline implementation.
+        """
+        return eager_v4_attention(
+            q,
+            k,
+            v,
+            sink=self.attn_sink,
+            swa_window=0,
+            additive_mask=attn_mask,
+            attn_dropout=self.attn_dropout,
+            training=self.training,
+            scale=self._attention_scale(),
+        )
+
+    def _attention_forward_via_v4_triton(
+        self,
+        q: torch.Tensor,  # [B, H, Sq, head_dim]
+        k: torch.Tensor,  # [B, H, Sk, head_dim]
+        v: torch.Tensor,  # [B, H, Sk, head_dim]
+        attn_mask: Optional[torch.Tensor],  # [Sq, Sk] additive (broadcasts over B, H)
+        *,
+        swa_window: int = 0,
+        hca_local_seqlen: int = 0,
+    ) -> torch.Tensor:
+        """Run the dense / HCA softmax-and-attend through the plan-4
+        in-tree :func:`v4_attention_v1` Triton kernel.
+
+        Numerically equivalent to :meth:`_attention_forward` (same eager
+        ``q @ k^T * scale + mask + sink → softmax → @ v`` math) but
+        executes in a single fused kernel that re-materialises ``P``
+        from the saved LSE during the BWD instead of storing the
+        ``[Sq, Sk]`` ``P`` tensor — important at full V4-Flash dims
+        (``S=4096`` ⇒ ``P`` is 32 MiB / microbatch).
+
+        Plan-5 P30 flips the dense path to ``attn_mask=None`` +
+        ``swa_window > 0`` so the kernel can skip K tiles that are
+        guaranteed outside the sliding window. HCA uses the same pruning
+        for its local prefix by passing a pool-only mask plus
+        ``hca_local_seqlen``; the kernel then runs local SWA and pool
+        visibility as two loops under one joint softmax.
+        """
+        # Plan-5 P32: opt-in microbench-vs-proxy timing harness, gated
+        # by ``PRIMUS_V4_DIAG_TIME=1``. Adds a synchronous cuda.Event
+        # span around the kernel call and dumps per-mode median/min/max
+        # at process exit (rank 0 only). Used to root-cause the dual-RoPE
+        # bf16 -> fp32 upcast bug that made every V4 attention kernel
+        # run 1.8-7x slower in the proxy than in the standalone bench;
+        # left in-tree for future microbench-vs-proxy regressions.
+        if os.environ.get("PRIMUS_V4_DIAG_TIME", "0") == "1":
+            mode = "hca" if hca_local_seqlen > 0 else "dense"
+            if not _DeepseekV4AttentionDiag.shape_logged.get(mode, False):
+                _DeepseekV4AttentionDiag.shape_logged[mode] = True
+                print(
+                    f"[PRIMUS_V4_DIAG_TIME] mode={mode}  "
+                    f"q={tuple(q.shape)}/{q.dtype}/contig={q.is_contiguous()}  "
+                    f"k={tuple(k.shape)}/{k.dtype}/contig={k.is_contiguous()}  "
+                    f"v={tuple(v.shape)}/{v.dtype}/contig={v.is_contiguous()}  "
+                    f"swa={swa_window} hca_local={hca_local_seqlen}",
+                    flush=True,
+                )
+            torch.cuda.synchronize()
+            ev_s = torch.cuda.Event(enable_timing=True)
+            ev_e = torch.cuda.Event(enable_timing=True)
+            ev_s.record()
+            out = v4_attention_v1(
+                q,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(swa_window) if (attn_mask is None or hca_local_seqlen > 0) else 0,
+                additive_mask=attn_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=int(hca_local_seqlen),
+            )
+            ev_e.record()
+            torch.cuda.synchronize()
+            _DeepseekV4AttentionDiag.record(mode=mode, ms=ev_s.elapsed_time(ev_e), swa=swa_window)
+            return out
+        return v4_attention_v1(
+            q,
+            k,
+            v,
+            sink=self.attn_sink,
+            swa_window=int(swa_window) if (attn_mask is None or hca_local_seqlen > 0) else 0,
+            additive_mask=attn_mask,
+            attn_dropout=self.attn_dropout,
+            training=self.training,
+            scale=self._attention_scale(),
+            hca_local_seqlen=int(hca_local_seqlen),
+        )
+
+    def _attention_forward_via_core(
+        self,
+        q: torch.Tensor,  # [B, S, H, head_dim] (post-RoPE)
+        kv: torch.Tensor,  # [B, S, 1, head_dim] (post-RoPE, single-latent)
+    ) -> torch.Tensor:
+        """Run the dense (compress_ratio == 0) softmax-and-attend through
+        ``self.core_attention`` (Turbo flash-attn / TE flash-attn).
+
+        Plan-3 P22.  Avoids materialising the eager
+        ``[B, H, S, S] fp32`` logits tensor — at full V4-Flash dims
+        (``H=64, S=4096, hc_mult=4``) that's 16 GiB / microbatch, and
+        the dominant activation cost.
+
+        Inputs use V4's local-frame layout (Q has all H heads, KV is
+        single-latent with 1 head).  We forward as Turbo's required
+        ``qkv_format="sbhd"`` and let the underlying flash kernel
+        broadcast the 1-head KV across H query heads (MQA).  Causal
+        masking + (optional) sliding window are honored by the kernel
+        directly — the eager ``local_mask`` is not used here.
+
+        Returns ``[B, H, S, head_dim]`` to match the contract of
+        :meth:`_attention_forward`.
+        """
+        B, S, H, Dh = q.shape
+        # [B, S, H, D] -> [S, B, H, D] (qkv_format="sbhd").
+        q_sbhd = q.transpose(0, 1).contiguous()
+        kv_sbhd = kv.transpose(0, 1).contiguous()  # [S, B, 1, D]
+
+        # Turbo / TE flash-attn forward.  ``attention_mask=None`` is
+        # legal for causal+SWA (the kernel builds the mask internally
+        # from ``attn_mask_type`` + the layer's ``window_size``).
+        out = self.core_attention(
+            q_sbhd,
+            kv_sbhd,
+            kv_sbhd,
+            None,
+            attn_mask_type=AttnMaskType.causal,
+        )  # -> [S, B, H * head_dim]
+
+        # [S, B, H*D] -> [B, S, H, D] -> [B, H, S, D].
+        out = out.view(S, B, H, Dh).permute(1, 2, 0, 3).contiguous()
+        return out
+
+    def _grouped_o_projection(self, attn: torch.Tensor) -> torch.Tensor:
+        """Apply the V4 grouped low-rank O projection.
+
+        Input ``attn`` shape: ``[B, S, H, head_dim]``.
+        Output shape: ``[B, S, hidden_size]``.
+
+        Math (from ``inference/model.py``):
+
+        .. code-block:: python
+
+            # attn  : [B, S, G, (H*head_dim)/G]
+            # wo_a.weight : [G * o_lora_rank, (H*head_dim)/G]
+            wo_a_w = self.linear_o_a.weight.view(G, o_lora_rank, -1)
+            o      = einsum("bsgd,grd->bsgr", attn, wo_a_w)
+            o      = self.linear_o_b(o.flatten(2))
+
+        We use the Linear's stored ``weight`` directly so the per-group
+        einsum semantics are exact. (Megatron's parallel linears expose
+        ``.weight`` after ``build_module``.)
+        """
+        B, S, H, Dh = attn.shape
+        G = self.o_groups
+        attn_g = attn.reshape(B, S, G, (H * Dh) // G)  # [B, S, G, H*Dh/G]
+
+        wo_a = self.linear_o_a
+        weight = wo_a.weight if hasattr(wo_a, "weight") else None
+        if weight is None:
+            # Fall back to a dense linear apply (Megatron parallel linears
+            # without a directly accessible weight attribute).
+            o = _projection_forward(wo_a, attn_g.reshape(B, S, -1))
+            o = o.view(B, S, G * self.o_lora_rank)
+        else:
+            wo_a_w = weight.view(G, self.o_lora_rank, (H * Dh) // G)
+            if _v4_o_a_fp8_enabled(self.config):
+                o = _fp8_grouped_o_a(attn_g, wo_a_w)  # per-group MXFP8
+            else:
+                o = torch.einsum("bsgd,grd->bsgr", attn_g, wo_a_w)
+            o = o.flatten(2)
+        return _projection_forward(self.linear_o_b, o)
+
+    def _flat_o_projection(self, attn: torch.Tensor) -> torch.Tensor:
+        """MLA-style flat output projection (``o_lora_rank == 0`` fast path)."""
+        B, S, H, Dh = attn.shape
+        return _projection_forward(self.linear_proj, attn.reshape(B, S, H * Dh))
+
+    def _apply_inverse_rope(self, attn: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+        """De-rotate the core-attention output (RoPE at ``position = -t``).
+
+        V4 shares a single latent between K and V, so the values entering the
+        softmax are already rotated and the naive output
+
+            ``o_t = sum_s P[t,s] * R(s) v_s``
+
+        carries *absolute* positions. Rotating it by ``-t`` turns every
+        contribution into ``R(s - t) v_s``, i.e. the relative encoding the
+        model expects. The released ``inference/model.py`` does exactly this
+        with ``apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True)`` right
+        before ``wo_a``, and the inverse rotation must use the *same* per-layer
+        RoPE base that Q / KV were rotated with.
+
+        Routed through the same fused ``apply_rope_from_positions`` path Q / KV
+        use, with ``inverse=True`` negating the angle in-kernel
+        (``R(-t) == R(t)^T``). Generating cos / sin inside the kernel matters
+        here: this runs once per layer, so the alternative -- building cos / sin
+        eagerly and broadcasting them to ``[B, S, rd/2]`` in HBM -- adds three
+        small launches and a materialised tensor per layer, on all 43 of them.
+
+        Returns a new tensor: ``attn`` is what the attention backward saved and
+        must not be mutated.
+        """
+        if self.rotary_dim == 0:
+            return attn
+        rope = self.rope.get_rope(compress_ratio=self.compress_ratio)
+        return apply_rope_from_positions(
+            attn,
+            position_ids,
+            rope.inv_freq,
+            rotary_dim=self.rotary_dim,
+            inverse=True,
+        )
+
+    def _project_output(self, attn: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+        """Inverse-RoPE then O projection -- the single exit for every branch.
+
+        Kept as one helper so no attention path can reach the O projection
+        without de-rotating first. Being the single exit also makes it the one
+        place that reports the layer's indexer distillation loss, which has to
+        happen on every layer (see :meth:`_log_indexer_distill_loss`).
+        """
+        self._log_indexer_distill_loss(attn.device)
+        attn = self._apply_inverse_rope(attn, position_ids)
+        if self.linear_o_a is not None:
+            return self._grouped_o_projection(attn)
+        return self._flat_o_projection(attn)
+
+    # ------------------------------------------------------------------
+    # compressed branches (HCA / CSA)
+    # ------------------------------------------------------------------
+
+    def _build_compressed_pool(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Run the compressor + compress-base partial RoPE.
+
+        Returns ``[B, P, head_dim]`` where ``P = S // compress_ratio``.
+        """
+        device = hidden.device
+        pooled = self.compressor(hidden)  # [B, P, head_dim]
+        B, P = pooled.shape[0], pooled.shape[1]
+
+        # Compress-base partial RoPE. Compressed entry ``s`` stands for the
+        # window starting at original token ``s * compress_ratio``, so it is
+        # rotated at that position -- not at the bare block index ``s``. The
+        # queries are rotated at their own original positions, so using block
+        # indices here would put the two sides on different coordinate systems.
+        # Matches inference/model.py, which slices ``freqs_cis[:cutoff:ratio]``
+        # for prefill and indexes ``start_pos + 1 - ratio`` for decode -- both
+        # land on the window's first token. Positions stay deterministic, so the
+        # cached table is still reused every forward.
+        cos, sin = self.rope.compress_rope.forward_arange(P, device, stride=self.compress_ratio)
+        cos = cos[..., : self.rotary_dim // 2]
+        sin = sin[..., : self.rotary_dim // 2]
+        cos = cos.unsqueeze(0).expand(B, -1, -1)
+        sin = sin.unsqueeze(0).expand(B, -1, -1)
+        pool_kv = pooled.unsqueeze(2)  # [B, P, 1, head_dim]
+        pool_kv = apply_interleaved_partial_rope(pool_kv, cos, sin, rotary_dim=self.rotary_dim)
+        return pool_kv.squeeze(2)  # [B, P, head_dim]
+
+    def _hca_extra_kv(
+        self,
+        hidden: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build the HCA (compress_ratio == 128) compressed branch.
+
+        Returns ``(extra_k_bh, extra_v_bh, extra_mask)`` where the
+        compressed pool is broadcast across H heads (single-latent
+        compressor output) and the additive mask is shape ``[S, P]``
+        (broadcasts over B, H).
+
+        Per the techblog: pool position ``s`` covers raw tokens
+        ``[s*ratio, (s+1)*ratio)``; query at raw token ``t`` may attend
+        to ``s`` iff ``(s+1)*ratio - 1 <= t``.
+        """
+        B, S, _ = hidden.shape
+        device, dtype = hidden.device, hidden.dtype
+        pool = self._build_compressed_pool(hidden)  # [B, P, head_dim]
+        P = pool.shape[1]
+
+        # Broadcast pool across all H query-heads: [B, P, head_dim] -> [B, P, H, head_dim].
+        pool_h = pool.unsqueeze(2).expand(B, P, self.num_heads, self.head_dim)
+        # Move heads dim to dim=1: [B, H, P, head_dim].
+        pool_bh = pool_h.transpose(1, 2)
+
+        extra_mask = self._hca_extra_mask_cached(S, P, device, dtype)
+        return pool_bh, pool_bh, extra_mask  # K = V = compressed pool
+
+    def _hca_extra_mask_cached(self, S: int, P: int, device, dtype):
+        """HCA additive causal mask ``[S, P]``, cached (data-independent).
+
+        Pool slot ``s`` is visible to query ``t`` iff ``(s+1)*ratio - 1 <= t``;
+        the mask depends only on ``(S, P, compress_ratio, dtype)`` — all fixed
+        per run — so build it once instead of rebuilding arange + where every
+        compressed-layer forward. Bit-identical. PRIMUS_COMPRESS_MASK_CACHE=0
+        forces the eager rebuild.
+        """
+        if os.environ.get("PRIMUS_COMPRESS_MASK_CACHE", "1") == "0":
+            t = torch.arange(S, device=device).unsqueeze(1)
+            s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1
+            return torch.where(s_end <= t, 0.0, float("-inf")).to(dtype)
+        cache = getattr(self, "_hca_mask_cache", None)
+        if cache is None:
+            cache = self._hca_mask_cache = {}
+        key = (S, P, device, dtype)
+        m = cache.get(key)
+        if m is None:
+            t = torch.arange(S, device=device).unsqueeze(1)  # [S, 1]
+            s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1  # [1, P]
+            m = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype)
+            cache[key] = m
+        return m
+
+    def _csa_forward(
+        self,
+        hidden: torch.Tensor,
+        q_bh: torch.Tensor,  # [B, H, S, head_dim]
+        k_local_bh: torch.Tensor,  # [B, H, S, head_dim]
+        v_local_bh: torch.Tensor,  # [B, H, S, head_dim]
+        local_mask: torch.Tensor,  # [S, S] — built by caller; unused here, see below
+        position_ids: Optional[torch.Tensor] = None,  # [B, S] or [S]; rotates the indexer Q
+    ) -> torch.Tensor:
+        """CSA (compress_ratio == 4) joint local-SWA + sparse-compressed attention.
+
+        The compressor produces a per-batch pool ``[B, P, head_dim]``,
+        the indexer picks ``index_topk`` pool positions per query, and
+        the attention runs softmax JOINTLY over ``[local_keys, sparse_keys]``
+        so the optional ``attn_sink`` is shared across both branches.
+
+        Plan-4 P24: the compressor / indexer / per-query top-K gather
+        stay here (they are V4-specific side-paths that the kernel does
+        not own); the joint-softmax math is delegated to
+        :func:`primus...v4_attention_kernels._eager.reference.eager_v4_csa_attention`
+        so the CSA path, the plan-4 CSA Triton kernel (P26), and the
+        plan-4 unit-test harness share one definition. ``local_mask`` is
+        retained in the signature for back-compat but unused — the
+        reference op rebuilds the local SWA mask deterministically from
+        ``swa_window`` (same call to
+        :func:`sliding_window_causal_mask` as :meth:`_local_mask` makes,
+        so the result is bit-identical).
+
+        Plan-5 P31: when ``use_v4_triton_csa_attention=True`` the sparse
+        top-K pool gather moves into the Triton kernel. The eager fallback
+        still materialises ``gathered`` here so it remains the reference
+        implementation and keeps the old P26 API covered by unit tests.
+        """
+        del local_mask  # see docstring
+        B, H, S, Dh = q_bh.shape
+        dtype = hidden.dtype
+
+        # 1) Compressed pool with compress-base RoPE.
+        pool = self._build_compressed_pool(hidden)  # [B, P, head_dim]
+        P = pool.shape[1]
+
+        # 2) Indexer top-K per query. The indexer is fed a *detached* hidden
+        # state: it is a selector, so the only thing that should learn from its
+        # distillation loss is the indexer itself. Detaching here keeps the KL
+        # from leaking into the layers below (the open-source reference detaches
+        # the same way), and when the indexer is frozen it also stops autograd
+        # from building a subgraph whose output is discarded.
+        topk_idxs, topk_scores = self.indexer(hidden.detach(), position_ids)  # [B, S, K]
+
+        # 2b) Indexer distillation. argTopK is not differentiable and the
+        # scores are otherwise discarded, so this loss is the indexer's only
+        # learning signal. It is attached to ``pool`` -- which every CSA
+        # backend consumes -- so the aux gradient is seeded no matter which
+        # kernel the layer dispatches to, without threading a second return
+        # value through all of them.
+        if self.indexer_distill_enabled and self.training and torch.is_grad_enabled():
+            indexer_loss = compute_indexer_distill_loss(
+                index_topk_scores=topk_scores,
+                topk_idxs=topk_idxs,
+                query=q_bh,
+                pool=pool,
+                softmax_scale=self._attention_scale(),
+                loss_coeff=self.indexer_distill_coeff,
+                head_reduce_group=self._indexer_loss_head_group(),
+            )
+            self.last_indexer_distill_loss = indexer_loss.detach()
+            pool = V4IndexerLossAutoScaler.apply(pool, indexer_loss)
+
+        # Dispatch on ``use_v4_csa_attention_backend``. gluon / triton_v2 /
+        # triton_v1 consume (pool, topk) directly; eager / triton_v0 / flydsl_v0
+        # use the per-query gathered [B, S, K, Dh] representation.
+        be = self._csa_backend
+        if be == "gluon":
+            return self._v4_csa_attention_gluon(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "gluon_v2":
+            return self._v4_csa_attention_gluon_v2(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "gluon_v3":
+            return self._v4_csa_attention_gluon_v3(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "turbo":
+            return self._v4_csa_attention_turbo(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "triton_v2":
+            return v4_csa_attention_v2(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "flydsl_v1":
+            return self._v4_csa_attention_flydsl(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+        if be == "triton_v1":
+            return v4_csa_attention_v1(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                pool,
+                topk_idxs=topk_idxs,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+            )
+
+        # eager / triton_v0 / flydsl_v0: build the per-query gathered slices.
+        K = topk_idxs.shape[-1]
+        valid = topk_idxs >= 0  # [B, S, K]
+        safe_idx = topk_idxs.clamp(min=0)
+        idx_expand = safe_idx.unsqueeze(-1).expand(B, S, K, Dh)
+        pool_expand = pool.unsqueeze(1).expand(B, S, P, Dh)
+        gathered = torch.gather(pool_expand, dim=2, index=idx_expand)
+        gathered = gathered * valid.unsqueeze(-1).to(gathered.dtype)
+        sparse_mask = torch.where(valid, 0.0, float("-inf")).to(dtype)  # [B, S, K]
+
+        if be in ("triton_v0", "flydsl_v0"):
+            return v4_csa_attention_v0(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                gathered,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                sparse_mask=sparse_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                use_flydsl=(be == "flydsl_v0"),
+            )
+        return eager_v4_csa_attention(
+            q_bh,
+            k_local_bh,
+            v_local_bh,
+            gathered,
+            sink=self.attn_sink,
+            swa_window=int(self.attn_sliding_window),
+            sparse_mask=sparse_mask,
+            attn_dropout=self.attn_dropout,
+            training=self.training,
+            scale=self._attention_scale(),
+        )
+
+    # ------------------------------------------------------------------
+    # public forward
+    # ------------------------------------------------------------------
+
+    def _attention_backend_forward(self, q_bh, k, v, *, additive_mask, hca_local_seqlen, S, device, dtype):
+        """Dense (cr=0) / HCA (cr=128) dispatch on ``use_v4_attention_backend``."""
+        be = self._attn_backend
+        if be == "gluon":
+            return self._v4_attention_gluon(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "gluon_v2":
+            return self._v4_attention_gluon_v2(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "gluon_v3":
+            return self._v4_attention_gluon_v3(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "turbo":
+            return self._v4_attention_turbo(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "triton_v2":
+            return v4_attention_v2(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "flydsl_v1":
+            return self._v4_attention_flydsl(
+                q_bh,
+                k,
+                v,
+                sink=self.attn_sink,
+                swa_window=int(self.attn_sliding_window),
+                additive_mask=additive_mask,
+                attn_dropout=self.attn_dropout,
+                training=self.training,
+                scale=self._attention_scale(),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        if be == "triton_v1":
+            return self._attention_forward_via_v4_triton(
+                q_bh,
+                k,
+                v,
+                additive_mask,
+                swa_window=int(self.attn_sliding_window),
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        # eager
+        local_mask = self._local_mask(S, device=device, dtype=dtype)
+        mask = local_mask if additive_mask is None else torch.cat([local_mask, additive_mask], dim=-1)
+        return self._attention_forward(q_bh, k, v, mask)
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """``[B, S, D] -> [B, S, D]``.
+
+        Dispatches on ``self.compress_ratio``:
+
+        * ``0``   — dense / SWA over local KV (single key axis).
+        * ``128`` — HCA: concat compressed pool to local KV, joint softmax.
+        * ``4``   — CSA: per-query top-K from compressed pool, joint softmax.
+        """
+        B, S, _ = hidden.shape
+        device, dtype = hidden.device, hidden.dtype
+
+        q = self._apply_q(hidden)  # [B, S, H, head_dim]
+        kv = self._apply_kv(hidden)  # [B, S, 1, head_dim]
+
+        # Partial RoPE on Q and K. K is post-RoPE; V uses the SAME tensor
+        # (V4's single-latent design: K and V share the rope-applied kv).
+        q, kv = self._apply_rope_q_k(q, kv, position_ids)
+
+        if self.compress_ratio == 0 and self._use_core_attention:
+            # Plan-3 P22: dense layer, Turbo / TE flash path.  Causal + SWA
+            # are handled inside the kernel; the eager ``local_mask`` is
+            # not consulted here.  KV is forwarded as ``[S, B, 1, D]``
+            # and broadcast across H query heads via MQA.
+            out_bh = self._attention_forward_via_core(q, kv)
+            out = out_bh.transpose(1, 2).contiguous()  # [B, S, H, head_dim]
+            out = out.to(dtype=dtype)
+            return self._project_output(out, position_ids)
+
+        # Broadcast K / V across the H query-head axis.
+        k_h = kv.expand(B, S, self.num_heads, self.head_dim)
+        v_h = kv.expand(B, S, self.num_heads, self.head_dim)
+
+        # Move heads dim before sequence: [B, S, H, head_dim] -> [B, H, S, head_dim]
+        q_bh = q.transpose(1, 2)
+        k_local_bh = k_h.transpose(1, 2)
+        v_local_bh = v_h.transpose(1, 2)
+
+        if self.compress_ratio == 0:
+            out_bh = self._attention_backend_forward(
+                q_bh,
+                k_local_bh,
+                v_local_bh,
+                additive_mask=None,
+                hca_local_seqlen=0,
+                S=S,
+                device=device,
+                dtype=dtype,
+            )
+        elif self.compress_ratio == 128:
+            # HCA: the local SWA branch and the compressed-pool branch share ONE
+            # softmax with ONE sink column; concatenate the pool to the local
+            # keys and pass the pool-only additive mask.
+            extra_k_bh, extra_v_bh, extra_mask = self._hca_extra_kv(hidden)
+            k_full = torch.cat([k_local_bh, extra_k_bh], dim=2)  # along Sk
+            v_full = torch.cat([v_local_bh, extra_v_bh], dim=2)
+            out_bh = self._attention_backend_forward(
+                q_bh,
+                k_full,
+                v_full,
+                additive_mask=extra_mask,
+                hca_local_seqlen=S,
+                S=S,
+                device=device,
+                dtype=dtype,
+            )
+        elif self.compress_ratio == 4:
+            # CSA cannot use ``core_attention``: the per-query top-K
+            # gather (``gathered = pool[..., topk_idxs, :]``, shape
+            # ``[B, H, S, K, head_dim]``) is sparse-per-row indexed
+            # attention — there is no flash-attn kernel that reads a
+            # different per-query subset of keys from a pool.  Stays on
+            # eager-Python under plan-3 (a custom kernel is required).
+            local_mask = self._local_mask(S, device=device, dtype=dtype)
+            out_bh = self._csa_forward(hidden, q_bh, k_local_bh, v_local_bh, local_mask, position_ids)
+        else:
+            # Guarded by __init__; included for static-analysis completeness.
+            raise ValueError(f"Unsupported compress_ratio {self.compress_ratio}")
+
+        out = out_bh.transpose(1, 2).contiguous()  # [B, S, H, head_dim]
+        out = out.to(dtype=dtype)
+
+        return self._project_output(out, position_ids)
+
+
+__all__ = [
+    "DeepseekV4Attention",
+    "DeepseekV4AttentionSubmodules",
+]
diff --git a/primus/backends/megatron/core/transformer/diffusion_transformer_block.py b/primus/backends/megatron/core/transformer/diffusion_transformer_block.py
new file mode 100644
index 000000000..9e21033ca
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/diffusion_transformer_block.py
@@ -0,0 +1,352 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+DiffusionTransformerBlock for diffusion models with timestep conditioning.
+
+This module provides a specialized TransformerBlock for diffusion models that properly
+handles timestep embeddings and other conditioning parameters through gradient checkpointing.
+"""
+
+from typing import Optional, Union
+
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.utils import WrappedTensor
+from torch import Tensor
+
+
+class DiffusionTransformerBlock(TransformerBlock):
+    """
+    TransformerBlock for diffusion models with timestep conditioning.
+
+    Extends upstream Megatron TransformerBlock with explicit timestep_emb and
+    guidance_emb parameters, routed through gradient checkpointing in a
+    thread-safe manner (no instance attribute storage for conditioning).
+
+    Call chain:
+        Flux.forward() -> transformer(timestep_emb=...) ->
+        DiffusionTransformerBlock.forward() -> checkpoint(custom_forward, timestep_emb) ->
+        layer(timestep_emb=...) -> adaln(timestep_emb)
+    """
+
+    def forward(
+        self,
+        hidden_states: Union[Tensor, WrappedTensor],
+        attention_mask: Optional[Tensor],
+        context: Optional[Tensor] = None,
+        context_mask: Optional[Tensor] = None,
+        rotary_pos_emb: Optional[Tensor] = None,
+        rotary_pos_cos: Optional[Tensor] = None,
+        rotary_pos_sin: Optional[Tensor] = None,
+        rotary_pos_cos_sin: Optional[Tensor] = None,
+        attention_bias: Optional[Tensor] = None,
+        inference_context: Optional[BaseInferenceContext] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+        sequence_len_offset: Optional[Tensor] = None,
+        # Diffusion-specific conditioning parameters
+        timestep_emb: Optional[Tensor] = None,
+        guidance_emb: Optional[Tensor] = None,
+        **kwargs,
+    ):
+        """
+        Forward pass with explicit diffusion conditioning.
+
+        Timestep and guidance embeddings are passed as tensors and participate in
+        gradient checkpointing like other forward inputs.
+
+        Args:
+            hidden_states: Input tensor [seq, batch, hidden]
+            attention_mask: Attention mask tensor
+            context: Context tensor for cross-attention (e.g., text embeddings)
+            context_mask: Mask for context in cross-attention
+            timestep_emb: Timestep conditioning [batch, hidden] for diffusion models
+            guidance_emb: Guidance conditioning [batch, hidden] for classifier-free guidance
+            All other args: Same as TransformerBlock.forward()
+
+        Returns:
+            Output tensor [seq, batch, hidden], or a tuple (hidden_states, context)
+            when cross-attention context is present.
+
+        Note:
+            This implementation is thread-safe because no conditioning parameters are
+            stored as instance attributes. All conditioning flows through function arguments.
+        """
+        # Validate diffusion parameters if model type requires them
+        if hasattr(self.config, "model_type") and self.config.model_type in ["flux", "diffusion"]:
+            if timestep_emb is None:
+                raise ValueError(
+                    f"DiffusionTransformerBlock requires timestep_emb for model_type={self.config.model_type}"
+                )
+
+        # Build conditioning_kwargs for passing to layers
+        # Only include non-None values to avoid passing unused parameters
+        conditioning_kwargs = {}
+        if timestep_emb is not None:
+            conditioning_kwargs["timestep_emb"] = timestep_emb
+        if guidance_emb is not None:
+            conditioning_kwargs["guidance_emb"] = guidance_emb
+
+        # Extract kwargs that need to be passed to parent (excluding conditioning)
+        parent_kwargs = {
+            "hidden_states": hidden_states,
+            "attention_mask": attention_mask,
+            "context": context,
+            "context_mask": context_mask,
+            "rotary_pos_emb": rotary_pos_emb,
+            "rotary_pos_cos": rotary_pos_cos,
+            "rotary_pos_sin": rotary_pos_sin,
+            "rotary_pos_cos_sin": rotary_pos_cos_sin,
+            "attention_bias": attention_bias,
+            "inference_context": inference_context,
+            "packed_seq_params": packed_seq_params,
+            "sequence_len_offset": sequence_len_offset,
+        }
+        # Filter out conditioning parameters from kwargs before updating
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ("timestep_emb", "guidance_emb")}
+        parent_kwargs.update(filtered_kwargs)
+
+        # Check if checkpointing is needed (based on parent's logic)
+        if self.config.recompute_granularity == "full" and self.training:
+            # Checkpointed path: Need to override _checkpointed_forward to handle conditioning
+            # Delete the obsolete reference to the initial input tensor if necessary
+            if isinstance(hidden_states, WrappedTensor):
+                hidden_states = hidden_states.unwrap()
+
+            if not self.pre_process:
+                # See set_input_tensor()
+                hidden_states = self.input_tensor
+
+            # Determine inner quantization context usage (from parent logic)
+            from megatron.core.enums import Fp8Recipe
+
+            if self.config.fp8:
+                use_inner_quantization_context = self.config.fp8_recipe != Fp8Recipe.delayed
+            elif self.config.fp4:
+                use_inner_quantization_context = True
+            else:
+                use_inner_quantization_context = False
+
+            # Call custom checkpointed forward with conditioning
+            return self._checkpointed_forward_with_conditioning(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                context=context,
+                context_mask=context_mask,
+                rotary_pos_emb=rotary_pos_emb,
+                attention_bias=attention_bias,
+                packed_seq_params=packed_seq_params,
+                use_inner_quantization_context=use_inner_quantization_context,
+                **conditioning_kwargs,  # Pass conditioning to checkpoint
+            )
+        else:
+            # Normal forward path - manually iterate layers with conditioning
+            # Parent TransformerBlock doesn't accept timestep_emb/guidance_emb
+            # So we need to manually process layers when conditioning is present
+            if conditioning_kwargs:
+                # Process layers manually with conditioning
+                current_hidden = hidden_states
+                current_context = context
+
+                for layer in self.layers:
+                    layer_kwargs = {
+                        "hidden_states": current_hidden,
+                        "attention_mask": attention_mask,
+                        "context": current_context,
+                        "context_mask": context_mask,
+                        "rotary_pos_emb": rotary_pos_emb,
+                        "attention_bias": attention_bias,
+                        "inference_context": inference_context,
+                        "packed_seq_params": packed_seq_params,
+                    }
+                    # Add conditioning
+                    layer_kwargs.update(conditioning_kwargs)
+
+                    layer_output = layer(**layer_kwargs)
+                    if isinstance(layer_output, tuple):
+                        current_hidden, current_context = layer_output
+                    else:
+                        current_hidden = layer_output
+                        current_context = None
+
+                return current_hidden if current_context is None else (current_hidden, current_context)
+            else:
+                # No conditioning - use parent's forward
+                return super().forward(**parent_kwargs)
+
+    def _checkpointed_forward_with_conditioning(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Tensor,
+        context: Tensor,
+        context_mask: Tensor,
+        rotary_pos_emb: Tensor,
+        attention_bias: Tensor,
+        packed_seq_params: PackedSeqParams,
+        use_inner_quantization_context: bool,
+        timestep_emb: Optional[Tensor] = None,
+        guidance_emb: Optional[Tensor] = None,
+    ):
+        """
+        Checkpointed forward with conditioning support.
+
+        This method extends the parent's _checkpointed_forward to handle
+        diffusion conditioning parameters (timestep_emb, guidance_emb).
+
+        Note: Unlike the non-checkpointed path in forward(), this method
+        only returns hidden_states and discards context.
+        """
+        from contextlib import nullcontext
+
+        from megatron.core import tensor_parallel
+        from megatron.core.fp4_utils import get_fp4_context
+        from megatron.core.fp8_utils import get_fp8_context
+        from megatron.core.transformer.transformer_layer import (
+            get_transformer_layer_offset,
+        )
+
+        try:
+            from megatron.core.extensions.transformer_engine import te_checkpoint
+        except ImportError:
+            te_checkpoint = None
+
+        def custom(start: int, end: int):
+            def custom_forward(
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                timestep_emb=None,
+                guidance_emb=None,
+            ):
+                for index in range(start, end):
+                    layer = self._get_layer(index)
+
+                    # Get appropriate inner quantization context
+                    if use_inner_quantization_context:
+                        if self.config.fp8:
+                            inner_quantization_context = get_fp8_context(self.config, layer.layer_number - 1)
+                        elif self.config.fp4:
+                            inner_quantization_context = get_fp4_context(self.config, layer.layer_number - 1)
+                        else:
+                            inner_quantization_context = nullcontext()
+                    else:
+                        inner_quantization_context = nullcontext()
+
+                    with inner_quantization_context:
+                        # Build layer kwargs with conditioning
+                        layer_kwargs = {
+                            "hidden_states": hidden_states,
+                            "attention_mask": attention_mask,
+                            "context": context,
+                            "context_mask": context_mask,
+                            "rotary_pos_emb": rotary_pos_emb,
+                            "attention_bias": attention_bias,
+                            "inference_context": None,
+                            "packed_seq_params": packed_seq_params,
+                        }
+                        # Add conditioning if present
+                        if timestep_emb is not None:
+                            layer_kwargs["timestep_emb"] = timestep_emb
+                        if guidance_emb is not None:
+                            layer_kwargs["guidance_emb"] = guidance_emb
+
+                        hidden_states, context = layer(**layer_kwargs)
+                return hidden_states, context
+
+            return custom_forward
+
+        def checkpoint_handler(forward_func):
+            """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
+            if self.config.fp8 or self.config.fp4:
+                return te_checkpoint(
+                    forward_func,
+                    self.config.distribute_saved_activations,
+                    tensor_parallel.random.get_cuda_rng_tracker,
+                    self.pg_collection.tp,
+                    hidden_states,
+                    attention_mask,
+                    context,
+                    context_mask,
+                    rotary_pos_emb,
+                    (
+                        timestep_emb if timestep_emb is not None else context
+                    ),  # Use context as placeholder if None
+                    (
+                        guidance_emb if guidance_emb is not None else context
+                    ),  # Use context as placeholder if None
+                )
+            else:
+                return tensor_parallel.checkpoint(
+                    forward_func,
+                    self.config.distribute_saved_activations,
+                    hidden_states,
+                    attention_mask,
+                    context,
+                    context_mask,
+                    rotary_pos_emb,
+                    (
+                        timestep_emb if timestep_emb is not None else context
+                    ),  # Use context as placeholder if None
+                    (
+                        guidance_emb if guidance_emb is not None else context
+                    ),  # Use context as placeholder if None
+                )
+
+        recompute_layer_ids = getattr(self.config, "recompute_layer_ids", None)
+
+        if self.config.recompute_method == "uniform":
+            layer_idx = 0
+            while layer_idx < self.num_layers_per_pipeline_rank:
+                hidden_states, context = checkpoint_handler(
+                    custom(layer_idx, layer_idx + self.config.recompute_num_layers)
+                )
+                layer_idx += self.config.recompute_num_layers
+
+        elif self.config.recompute_method == "block":
+            recompute_skip_num_layers = 0
+            for layer_idx in range(self.num_layers_per_pipeline_rank):
+                if (self.config.fp8 or self.config.fp4) and not hidden_states.requires_grad:
+                    recompute_skip_num_layers += 1
+                if (
+                    layer_idx >= recompute_skip_num_layers
+                    and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers
+                ):
+                    hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1))
+                else:
+                    # Build args for non-checkpointed path
+                    forward_args = (
+                        hidden_states,
+                        attention_mask,
+                        context,
+                        context_mask,
+                        rotary_pos_emb,
+                        timestep_emb,
+                        guidance_emb,
+                    )
+                    hidden_states, context = custom(layer_idx, layer_idx + 1)(*forward_args)
+
+        elif recompute_layer_ids is not None:
+            for block_layer_idx in range(self.num_layers_per_pipeline_rank):
+                layer_idx = block_layer_idx + get_transformer_layer_offset(self.config, self.vp_stage)
+                if layer_idx not in recompute_layer_ids or (
+                    (self.config.fp8 or self.config.fp4) and not hidden_states.requires_grad
+                ):
+                    forward_args = (
+                        hidden_states,
+                        attention_mask,
+                        context,
+                        context_mask,
+                        rotary_pos_emb,
+                        timestep_emb,
+                        guidance_emb,
+                    )
+                    hidden_states, context = custom(block_layer_idx, block_layer_idx + 1)(*forward_args)
+                else:
+                    hidden_states, context = checkpoint_handler(custom(block_layer_idx, block_layer_idx + 1))
+        else:
+            raise ValueError("Invalid activation recompute method.")
+
+        return hidden_states
diff --git a/primus/backends/megatron/core/transformer/dual_rope.py b/primus/backends/megatron/core/transformer/dual_rope.py
new file mode 100644
index 000000000..2dad193c0
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/dual_rope.py
@@ -0,0 +1,384 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Dual-RoPE for DeepSeek-V4.
+
+Reference: techblog §6 ("RoPE: dual base + YaRN details").
+
+V4 layers fall into two RoPE regimes, decided per-layer by
+``compress_ratios[i]``:
+
+* ``compress_ratio == 0`` (dense / SWA layers): use the **main** RoPE base
+  (``rope_theta = 10000``), **no YaRN**.
+* ``compress_ratio != 0`` (CSA / HCA layers): use the **compress** RoPE
+  base (``compress_rope_theta``; 40000 for V4-Flash, 160000 for V4-Pro)
+  **with YaRN scaling**
+  (``factor=16, beta_fast=32, beta_slow=1, original_max_position_embeddings=65536``).
+
+Compressed KV entries are rotated at their **original-sequence** positions
+(entry ``s`` covers the window starting at token ``s * compress_ratio``), so
+they share the query coordinate system -- see ``forward_arange(..., stride=)``.
+
+Two important corrections over the HF Llama-style RoPE that V4 inherits
+from the released weights:
+
+1. **Interleaved RoPE** (pairs ``(2k, 2k+1)``), **not** rotate-half pairs
+   ``(d, d+rd/2)``. NeMo's port did this; HF PR 45616 originally did not.
+2. **Partial RoPE**: only the last ``qk_pos_emb_head_dim`` (= 64) channels of
+   each head are rotated. The first ``head_dim - 64`` channels stay nope.
+
+This module exposes:
+
+* :class:`DualRoPE` — owns two ``RoPECache`` instances (main + compress).
+  ``forward(layer_compress_ratio, x, position_ids, *, key=False)`` applies
+  the right partial-RoPE.
+* :class:`RoPECache` — single-base RoPE, with optional YaRN scaling.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.nn as nn
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import (
+    RoPEInterleavedPartialFn,
+    apply_rope_from_positions,
+)
+
+# ---------------------------------------------------------------------------
+# YaRN scaling
+# ---------------------------------------------------------------------------
+
+
+def _yarn_freq_scaling(
+    inv_freq: torch.Tensor,
+    *,
+    factor: float,
+    beta_fast: float,
+    beta_slow: float,
+    original_max_position_embeddings: int,
+) -> torch.Tensor:
+    """Apply YaRN frequency-band scaling to ``inv_freq``.
+
+    Each frequency component ``inv_freq[i] = 1 / theta**(2i/D)``. YaRN groups
+    the frequency band by wavelength ``2π / inv_freq[i]``:
+
+    * High-freq band (wavelength < ``original_seq / beta_fast``): keep as-is.
+    * Low-freq band  (wavelength > ``original_seq / beta_slow``): scale
+      ``inv_freq[i]`` by ``1/factor`` (i.e. lengthen wavelength).
+    * Mid-freq band: linear interpolation between the two.
+
+    Standard YaRN reference; see e.g. NeMo's ``YarnRotaryPositionEmbedding``.
+    """
+    if factor == 1.0 or original_max_position_embeddings <= 0:
+        return inv_freq
+
+    # Wavelengths corresponding to each inv_freq component.
+    wavelens = 2.0 * math.pi / inv_freq  # same shape as inv_freq
+
+    low_thresh = original_max_position_embeddings / beta_fast
+    high_thresh = original_max_position_embeddings / beta_slow
+
+    # Linear blend factor: 0 at high-freq side (no scaling), 1 at low-freq side (full /factor scaling).
+    smooth = ((wavelens - low_thresh) / max(high_thresh - low_thresh, 1e-12)).clamp(min=0.0, max=1.0)
+
+    # Scaled inv_freq.
+    inv_freq_scaled = inv_freq / factor
+    return inv_freq * (1.0 - smooth) + inv_freq_scaled * smooth
+
+
+def _yarn_attn_scale(factor: float) -> float:
+    """The standard YaRN attention magnitude scale ``m_scale``.
+
+    ``m_scale = 0.1 * log(factor) + 1`` — used to scale the attention
+    softmax temperature when YaRN is on.
+    """
+    if factor <= 1.0:
+        return 1.0
+    return 0.1 * math.log(factor) + 1.0
+
+
+# ---------------------------------------------------------------------------
+# RoPE cache (single base)
+# ---------------------------------------------------------------------------
+
+
+class RoPECache(nn.Module):
+    """Builds ``cos`` / ``sin`` tables on demand for a single RoPE base.
+
+    Stores ``inv_freq`` as a (non-trainable) buffer so it follows the model's
+    device / dtype movements. ``cos`` / ``sin`` are computed lazily on
+    ``forward``.
+    """
+
+    def __init__(
+        self,
+        rotary_dim: int,
+        *,
+        theta: float,
+        yarn_factor: float = 1.0,
+        yarn_beta_fast: float = 32.0,
+        yarn_beta_slow: float = 1.0,
+        original_max_position_embeddings: int = 0,
+    ) -> None:
+        super().__init__()
+        if rotary_dim % 2 != 0:
+            raise ValueError(f"rotary_dim must be even, got {rotary_dim}")
+
+        self.rotary_dim = rotary_dim
+        self.theta = float(theta)
+        self.yarn_factor = float(yarn_factor)
+        self.yarn_beta_fast = float(yarn_beta_fast)
+        self.yarn_beta_slow = float(yarn_beta_slow)
+        self.original_max_position_embeddings = int(original_max_position_embeddings)
+
+        # inv_freq[i] = 1 / theta**(2i/rotary_dim), i in [0, rotary_dim/2)
+        i = torch.arange(0, rotary_dim, 2, dtype=torch.float32)
+        inv_freq = 1.0 / (theta ** (i / rotary_dim))
+        inv_freq = _yarn_freq_scaling(
+            inv_freq,
+            factor=self.yarn_factor,
+            beta_fast=self.yarn_beta_fast,
+            beta_slow=self.yarn_beta_slow,
+            original_max_position_embeddings=self.original_max_position_embeddings,
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # YaRN's m_scale, kept for reference / diagnostics only. V4 does NOT
+        # fold it into the attention softmax temperature -- the released
+        # ``inference/model.py`` uses a plain ``head_dim ** -0.5`` and relies on
+        # the Q / KV RMSNorms to keep logits in range. Only the ``inv_freq``
+        # interpolation above is part of the model contract.
+        self.attn_scale: float = _yarn_attn_scale(self.yarn_factor)
+
+        # Memo of (cos, sin) tables keyed by (n, stride, device) for
+        # ``arange(n) * stride`` positions; see forward_arange. Not a buffer
+        # (recomputable, device-keyed).
+        self._arange_cache: dict = {}
+
+    def forward(self, position_ids: torch.Tensor) -> tuple:
+        """Return (cos, sin) for the given positions.
+
+        ``position_ids``: any integer/float tensor (typically ``[B, S]``
+        or ``[S]``).
+
+        Output ``(cos, sin)`` shape is ``position_ids.shape + (rotary_dim/2,)``;
+        the caller broadcasts over the head/batch axes.
+        """
+        # outer product positions × inv_freq
+        freqs = position_ids.float().unsqueeze(-1) * self.inv_freq  # [..., rotary_dim/2]
+        cos = freqs.cos()
+        sin = freqs.sin()
+        return cos, sin
+
+    def forward_arange(self, n: int, device, *, stride: int = 1) -> tuple:
+        """``(cos, sin)`` for positions ``torch.arange(n) * stride`` -- cached.
+
+        Equivalent to ``self.forward(torch.arange(n, device=device) * stride)``,
+        but memoised by ``(n, stride, device)``.
+
+        The compressed branch evaluates its RoPE at deterministic positions
+        (``P = S // compress_ratio`` entries, fixed per run), so the table is
+        identical every forward; caching it skips the ``arange ->
+        outer-product -> cos/sin`` recompute each step (and per compressed
+        layer). ``stride`` carries the compression ratio so the compressed
+        entries land on their *original-sequence* coordinates
+        (``0, ratio, 2*ratio, ...``) rather than on bare block indices. Set
+        ``PRIMUS_COMPRESS_ROPE_CACHE=0`` to disable the cache.
+        """
+        stride = int(stride)
+
+        def _build():
+            positions = torch.arange(n, device=device)
+            if stride != 1:
+                positions = positions * stride
+            return self.forward(positions)
+
+        if os.environ.get("PRIMUS_COMPRESS_ROPE_CACHE", "1") == "0":
+            return _build()
+        key = (int(n), stride, str(device))
+        hit = self._arange_cache.get(key)
+        if hit is None:
+            hit = _build()
+            self._arange_cache[key] = hit
+        return hit
+
+
+# ---------------------------------------------------------------------------
+# Partial interleaved RoPE application
+# ---------------------------------------------------------------------------
+
+
+def apply_interleaved_partial_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    rotary_dim: int,
+) -> torch.Tensor:
+    """Apply RoPE to the **last** ``rotary_dim`` channels of ``x`` using the
+    **interleaved** pairing convention: pairs ``(2k, 2k+1)``.
+
+    Args:
+        x: ``[..., head_dim]``. Typical layouts are ``[B, S, H, head_dim]``
+            or ``[S, B, H, head_dim]``.
+        cos, sin: must broadcast to ``x[..., :rotary_dim//2]`` after the
+            internal pair-reshape (i.e. shape ``[..., rotary_dim/2]`` matching
+            x's leading dims with a singleton heads axis inserted as needed).
+            The simplest contract: pass shape ``[..., rotary_dim/2]`` where
+            the leading dims are exactly the ``position_ids`` shape; this
+            function will insert a singleton "heads" dim at position ``-2``
+            so it broadcasts against the ``H`` axis of ``x``.
+        rotary_dim: the partial RoPE size; must be ``<= x.shape[-1]`` and
+            even.
+
+    Returns:
+        ``x`` with the last ``rotary_dim`` channels rotated; first
+        ``head_dim - rotary_dim`` channels untouched.
+    """
+    head_dim = x.shape[-1]
+    if rotary_dim > head_dim or rotary_dim % 2 != 0:
+        raise ValueError(f"rotary_dim must be even and <= head_dim ({head_dim}), got {rotary_dim}")
+    if rotary_dim == 0:
+        return x
+
+    # Plan-6 P35: route through the fused Triton kernel when on CUDA / HIP
+    # and the env knob is not "0".  The Triton path collapses the 9-op
+    # eager chain below (slice / reshape / four broadcast muls / stack /
+    # reshape / cat) into one kernel that does a single contiguous write
+    # with the rotation baked in -- removing the `CatArrayBatchedCopy_contig`
+    # bucket (~10 ms / 24 calls in the plan-5 P32 final trace) and the
+    # share of `elementwise_kernel_manual_unroll` that comes from the
+    # broadcast muls.  Eager body kept in tree as the `PRIMUS_ROPE_TRITON=0`
+    # fallback and as the G38 unit-test reference.
+    if x.is_cuda and os.environ.get("PRIMUS_ROPE_TRITON", "1") != "0":
+        return RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim)
+
+    orig_dtype = x.dtype
+    nope = head_dim - rotary_dim
+    x_nope = x[..., :nope]
+    x_rope = x[..., nope:]
+
+    # interleaved pairs: reshape last dim to (rotary_dim/2, 2)
+    x_pairs = x_rope.reshape(*x_rope.shape[:-1], rotary_dim // 2, 2)
+    even = x_pairs[..., 0]
+    odd = x_pairs[..., 1]
+
+    # Always insert a singleton "heads" axis at position -2 so cos/sin
+    # broadcast across H. cos starts as ``position_ids.shape + [rd/2]``;
+    # after unsqueeze it becomes ``position_ids.shape + [1, rd/2]`` which
+    # broadcasts naturally against ``even`` of shape ``[..., H, rd/2]``.
+    # Cast cos/sin to x's dtype so the rotation does not promote bf16
+    # inputs to fp32 -- otherwise the entire Q/K leaving RoPE doubles
+    # in HBM size and the downstream attention kernel runs 6-7x slower
+    # (see plan-5 P32 microbench-vs-proxy gap diagnostic).  Math is done
+    # in bf16 because the freqs come from
+    # ``position_ids.float() * inv_freq`` which is already a
+    # single-precision rounding of the integer position; casting the
+    # final cos/sin to bf16 keeps numerics indistinguishable from the
+    # bf16-only reference flash-attn path used at every other site.
+    cos = cos.unsqueeze(-2).to(orig_dtype)
+    sin = sin.unsqueeze(-2).to(orig_dtype)
+
+    rot_even = even * cos - odd * sin
+    rot_odd = even * sin + odd * cos
+
+    rotated = torch.stack([rot_even, rot_odd], dim=-1).reshape(*x_rope.shape[:-1], rotary_dim)
+    return torch.cat([x_nope, rotated], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# DualRoPE — main + compress bases together
+# ---------------------------------------------------------------------------
+
+
+class DualRoPE(nn.Module):
+    """Holds the two RoPE caches V4 needs and routes per-layer applications.
+
+    Args:
+        rotary_dim: partial-RoPE dim ``qk_pos_emb_head_dim`` (V4 = 64).
+        rope_theta: base for dense / SWA layers.
+        compress_rope_theta: base for CSA / HCA layers (longer base).
+        yarn_factor / yarn_beta_fast / yarn_beta_slow /
+            original_max_position_embeddings: YaRN config; applied **only**
+            to the compress base. Set ``yarn_factor=1.0`` to disable.
+    """
+
+    def __init__(
+        self,
+        *,
+        rotary_dim: int,
+        rope_theta: float,
+        compress_rope_theta: float,
+        yarn_factor: float = 1.0,
+        yarn_beta_fast: float = 32.0,
+        yarn_beta_slow: float = 1.0,
+        original_max_position_embeddings: int = 0,
+    ) -> None:
+        super().__init__()
+        self.rotary_dim = rotary_dim
+
+        self.main_rope = RoPECache(
+            rotary_dim=rotary_dim,
+            theta=rope_theta,
+        )
+        self.compress_rope = RoPECache(
+            rotary_dim=rotary_dim,
+            theta=compress_rope_theta,
+            yarn_factor=yarn_factor,
+            yarn_beta_fast=yarn_beta_fast,
+            yarn_beta_slow=yarn_beta_slow,
+            original_max_position_embeddings=original_max_position_embeddings,
+        )
+
+    def get_rope(self, *, compress_ratio: int) -> RoPECache:
+        """Pick the right cache for a layer.
+
+        ``compress_ratio == 0`` → main (dense / SWA). Anything else →
+        compress (CSA / HCA).
+        """
+        return self.main_rope if compress_ratio == 0 else self.compress_rope
+
+    def apply_rope(
+        self,
+        x: torch.Tensor,
+        *,
+        position_ids: torch.Tensor,
+        compress_ratio: int,
+    ) -> torch.Tensor:
+        """Convenience: pick the right rope and apply partial RoPE.
+
+        ``position_ids`` shape ``[B, S]`` or ``[S]``. cos/sin broadcast over
+        the heads dim of ``x``.
+
+        Small-kernel-fusion (2026-07-03): route through
+        :func:`apply_rope_from_positions`, which generates cos/sin *inside*
+        the RoPE Triton kernel from ``(position_ids, inv_freq)`` — removing
+        the separate ``position_ids.float() * inv_freq`` / ``cos`` / ``sin``
+        launches, the ``.to(x.dtype)`` cast, and the cos/sin HBM tensors that
+        the eager ``rope(position_ids)`` path materialised (and recomputed
+        identically for Q and K). YaRN is already baked into ``inv_freq``.
+        Gated by ``PRIMUS_ROPE_TRITON`` (falls back to the eager cos/sin path
+        on CPU / when off).
+        """
+        rope = self.get_rope(compress_ratio=compress_ratio)
+        return apply_rope_from_positions(x, position_ids, rope.inv_freq, rotary_dim=self.rotary_dim)
+
+    # Diagnostic accessor for the YaRN m_scale. Not used by the attention
+    # softmax scale -- see the note on ``RoPECache.attn_scale``.
+    def attn_scale(self, *, compress_ratio: int) -> float:
+        return self.get_rope(compress_ratio=compress_ratio).attn_scale
+
+
+__all__ = [
+    "RoPECache",
+    "DualRoPE",
+    "apply_interleaved_partial_rope",
+]
diff --git a/primus/backends/megatron/core/transformer/experts.py b/primus/backends/megatron/core/transformer/experts.py
index eb2d51e05..469639a9a 100644
--- a/primus/backends/megatron/core/transformer/experts.py
+++ b/primus/backends/megatron/core/transformer/experts.py
@@ -5,6 +5,12 @@
 import torch
 import torch.nn.functional as F
 from megatron.core import tensor_parallel
+from megatron.core.activations import squared_relu
+from megatron.core.fusions.fused_bias_geglu import (
+    quick_gelu,
+    weighted_bias_quick_geglu_impl,
+)
+from megatron.core.fusions.fused_weighted_squared_relu import weighted_squared_relu_impl
 from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
     FineGrainedActivationOffloadingInterface as off_interface,
 )
@@ -40,6 +46,94 @@ def __init__(
         # NOTE: use_turbo_fused_act_with_probs is prioritized over use_te_activation_func and bias_activation_fusion
         self.use_turbo_fused_act_with_probs = args.use_turbo_fused_act_with_probs
         self.moe_router_padding_for_quantization = args.moe_router_padding_for_quantization
+        self.use_turbo_ragged_grouped_gemm = getattr(args, "use_turbo_ragged_grouped_gemm", False)
+        # PrimusTurbo's tensorwise FP8 grouped GEMM consumes the original GPU
+        # tokens_per_expert tensor, including non-aligned (ragged) group sizes.
+        # Keep TE's explicit zero-padding fallback for other recipes/backends.
+
+    def _use_explicit_quantization_padding(self) -> bool:
+        """Whether this forward must pad expert groups before quantization."""
+        if not (self.config.fp8 or self.config.fp4):
+            return False
+        if self.moe_router_padding_for_quantization:
+            return False
+        if not self.use_turbo_ragged_grouped_gemm:
+            return True
+
+        # Check the active autocast state at forward time. The CLI recipe alone
+        # does not prove that this grouped linear is actually using Turbo FP8.
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboLowPrecisionGlobalStateManager,
+        )
+
+        if not PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled():
+            raise RuntimeError(
+                "use_turbo_ragged_grouped_gemm=True requires an active PrimusTurbo FP8 autocast context."
+            )
+        quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
+        if not quant_config.current_scaling():
+            raise RuntimeError(
+                "use_turbo_ragged_grouped_gemm=True currently requires tensorwise dynamic FP8 scaling."
+            )
+        return False
+
+    def bias_act_func(self, intermediate_parallel, bias_parallel, permuted_probs):
+        """
+        Applies bias and activation function to the output of linear_fc1.
+        """
+        if self.config.use_te_activation_func:
+            if bias_parallel is not None:
+                intermediate_parallel = intermediate_parallel + bias_parallel
+            intermediate_parallel = self.activation_func(intermediate_parallel)
+            if permuted_probs is not None:
+                original_dtype = intermediate_parallel.dtype
+                intermediate_parallel = intermediate_parallel * permuted_probs
+                intermediate_parallel = intermediate_parallel.to(original_dtype)
+        elif self.config.bias_activation_fusion:
+            if self.activation_func == F.silu and self.config.gated_linear_unit:
+                from primus.backends.megatron.core.fusions.fused_bias_swiglu import (
+                    weighted_bias_swiglu_impl,
+                )
+
+                # dtype is handled inside the fused kernel
+                intermediate_parallel = weighted_bias_swiglu_impl(
+                    intermediate_parallel,
+                    bias_parallel,
+                    permuted_probs,
+                    self.config.activation_func_fp8_input_store,
+                    self.config.activation_func_clamp_value,
+                )
+            elif self.activation_func == quick_gelu and self.config.gated_linear_unit:
+                intermediate_parallel = weighted_bias_quick_geglu_impl(
+                    intermediate_parallel,
+                    bias_parallel,
+                    permuted_probs,
+                    self.config.activation_func_fp8_input_store,
+                    self.config.glu_linear_offset,
+                    self.config.activation_func_clamp_value,
+                )
+            else:
+                raise ValueError("Only support fusion of swiglu and quick_gelu in TEGroupedMLP.")
+        elif self.activation_func == squared_relu and self.config.use_fused_weighted_squared_relu:
+            assert bias_parallel is None, "Bias is not supported with fused weighted squared relu."
+            intermediate_parallel = weighted_squared_relu_impl(intermediate_parallel, permuted_probs)
+        else:
+            if self.config.gated_linear_unit:
+
+                def glu(x):
+                    x_glu, x_linear = torch.chunk(x, 2, dim=-1)
+                    if (val := self.config.activation_func_clamp_value) is not None:
+                        x_glu = x_glu.clamp(min=None, max=val)
+                        x_linear = x_linear.clamp(min=-val, max=val)
+                    return self.config.activation_func(x_glu) * (x_linear + self.config.glu_linear_offset)
+
+                intermediate_parallel = glu(intermediate_parallel)
+            else:
+                intermediate_parallel = self.activation_func(intermediate_parallel)
+            original_dtype = intermediate_parallel.dtype
+            intermediate_parallel = intermediate_parallel * permuted_probs
+            intermediate_parallel = intermediate_parallel.to(original_dtype)
+        return intermediate_parallel
 
     def bias_act_func_with_mask(
         self,
@@ -77,6 +171,18 @@ def bias_act_func_with_mask(
             # use the original bias_act_func from TEGroupedMLP, ignore the tokens_per_experts
             return self.bias_act_func(intermediate_parallel, bias_parallel, permuted_probs)
 
+    @staticmethod
+    def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs):
+        if bias_parallel is None:
+            return intermediate_parallel
+
+        # NOTE: tokens_per_expert is on GPU, so we need to convert it to a list of ints.
+        tokens_per_expert_cpu = tokens_per_expert.tolist()
+
+        return super()._apply_bias(
+            intermediate_parallel, bias_parallel, tokens_per_expert_cpu, permuted_probs
+        )
+
     def forward(
         self,
         permuted_local_hidden_states: torch.Tensor,
@@ -94,7 +200,8 @@ def forward(
         Return:
             output (torch.Tensor): The output of the local experts.
         """
-        if not self.moe_router_padding_for_quantization and (self.config.fp8 or self.config.fp4):
+        use_explicit_quantization_padding = self._use_explicit_quantization_padding()
+        if use_explicit_quantization_padding:
             # NOTE: When moe_router_padding_for_quantization is true the token is padded. So we can skip the padding here to reduce cpu sync.
             tokens_per_expert_cpu: list[int] = tokens_per_expert.tolist()
             actual_tokens_per_expert_cpu: list[int] = tokens_per_expert_cpu
@@ -142,7 +249,14 @@ def forward(
                 )
         else:
             with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output:
-                bias_act_output = self.bias_act_func(fc1_output, bias_parallel, permuted_probs)
+                # Honor use_turbo_fused_act_with_probs independently of activation
+                # recomputation. The helper falls back to bias_act_func when disabled.
+                bias_act_output = self.bias_act_func_with_mask(
+                    fc1_output,
+                    bias_parallel,
+                    permuted_probs,
+                    tokens_per_expert,
+                )
         output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert)
         if self.activation_recompute:
             self.activation_checkpoint.discard_output_and_register_recompute(output)
@@ -151,11 +265,10 @@ def forward(
         # to make sure the fc1_output is reloaded to GPU before recomputing moe_act.
         if self.offload_moe_act:
             output = off_interface.group_commit(output, name="moe_act", forced_released_tensors=[fc1_output])
-        # NOTE: tokens_per_expert is on GPU, so we need to convert it to a list of ints.
-        output = self._apply_bias(output, output_bias, tokens_per_expert.tolist(), permuted_probs)
+        output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs)
 
         # upad and concat the output
-        if not self.moe_router_padding_for_quantization and (self.config.fp8 or self.config.fp4):
+        if use_explicit_quantization_padding:
             output = self.quantization_unpadding(output, actual_tokens_per_expert_cpu)
 
         output_bias = None
diff --git a/primus/backends/megatron/core/transformer/fla_flash_attention.py b/primus/backends/megatron/core/transformer/fla_flash_attention.py
new file mode 100644
index 000000000..704d12623
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/fla_flash_attention.py
@@ -0,0 +1,266 @@
+###############################################################################
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Drop-in `core_attention` replacement for Megatron MLA that calls
+`flash_attn_func` directly, bypassing TransformerEngine's version check
+and CK/SDPA fallback.
+
+This module exists because on ROCm/MI300X Transformer-Engine's
+`TEDotProductAttention` performs a strict version range check on the
+installed `flash-attn` package (currently ``2.1.1 <= x <= 2.8.1``).  Newer
+builds (e.g. 2.8.3 shipped in the production container) fail that check
+and TE silently falls back to its Composable-Kernel backend, which on
+small MLA dims (n_heads × head_dim = 16 × 64) loses ~30 ms per layer
+versus calling flash-attn directly.
+
+The wrapper matches FLA's reference path in
+``fla/layers/mla.py`` -- a single call to
+``flash_attn_func(q, k, v, causal=True, softmax_scale=…)`` with no other
+overhead.
+
+Activation
+----------
+The wrapper auto-enables whenever the installed ``flash_attn`` version is
+outside Transformer-Engine's supported range (currently ``<= 2.8.1``).
+Override with:
+
+* ``PRIMUS_FLA_MLA_ATTN=1`` -- force-enable (use wrapper).
+* ``PRIMUS_FLA_MLA_ATTN=0`` -- force-disable (use TE's
+  ``TEDotProductAttention``, which on newer flash-attn drops to CK fallback).
+
+When enabled, the hybrid GDN/KDA spec swaps
+``core_attention=TEDotProductAttention`` for
+``core_attention=FLAFlashAttention`` automatically (see
+``primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs``).
+
+Caveats
+-------
+* Only supports the unpacked training path (``packed_seq_params is None``).
+  Varlen / packed sequences are not implemented; the wrapper will raise.
+* Always uses causal masking (matches Megatron's MLA spec
+  ``params={"attn_mask_type": AttnMaskType.causal}`` and FLA's MLA path).
+* ``current_max_attn_logits`` is exposed as ``None`` so MLA's optional
+  qk_clip path (disabled in our configs) keeps importing without error.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import time
+from typing import Any, Optional
+
+import torch
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+
+# TransformerEngine's pinned flash-attn version range.  When the installed
+# `flash_attn` is newer (e.g. 2.8.3 in the production container), TE silently
+# falls back to its Composable-Kernel backend, which is ~30 ms/MLA-block
+# slower on MI300X.  We use the upper bound to decide whether the wrapper
+# should auto-enable.
+_TE_FLASH_ATTN_MAX_SUPPORTED = (2, 8, 1)
+
+
+def _parse_version(ver: str) -> tuple[int, int, int]:
+    parts: list[int] = []
+    for tok in ver.split("."):
+        digits = "".join(c for c in tok if c.isdigit())
+        if not digits:
+            break
+        parts.append(int(digits))
+        if len(parts) == 3:
+            break
+    while len(parts) < 3:
+        parts.append(0)
+    return tuple(parts)  # type: ignore[return-value]
+
+
+def _flash_attn_exceeds_te_range() -> bool:
+    """True if the installed flash-attn is newer than TE supports."""
+    try:
+        import flash_attn  # type: ignore
+    except Exception:
+        return False
+    ver = getattr(flash_attn, "__version__", "0.0.0")
+    return _parse_version(ver) > _TE_FLASH_ATTN_MAX_SUPPORTED
+
+
+def is_enabled() -> bool:
+    """Return True if MLA should use the direct flash-attn path.
+
+    Precedence (resolved by fla_runtime_patches → ``args.fla_mla_attn``):
+      1. ``fla_mla_attn="0"`` → force-disable (use TE).
+      2. ``fla_mla_attn="1"`` → force-enable (use wrapper).
+      3. ``""`` / unset → auto-enable whenever the installed flash-attn
+         version is outside TE's supported range (> 2.8.1).
+    """
+    try:
+        from megatron.training import get_args
+
+        val = getattr(get_args(), "fla_mla_attn", "")
+    except Exception:
+        val = ""
+    if val:
+        return val == "1"
+    return _flash_attn_exceeds_te_range()
+
+
+# Lazy import of flash_attn so the module can be imported on hosts where
+# flash-attn isn't installed (e.g. dev machines). The actual import only
+# happens when the layer is constructed inside the training container.
+_flash_attn_func = None
+
+# One-shot banner so it is obvious from the training log whether the
+# wrapper actually got instantiated.  Without this it can be ambiguous
+# because the FLA path and TE+CK path produce near-identical loss/speed
+# numbers on small attention dims.
+_BANNER_PRINTED = False
+
+
+def _load_flash_attn():
+    global _flash_attn_func
+    if _flash_attn_func is not None:
+        return _flash_attn_func
+    try:
+        from flash_attn import flash_attn_func  # type: ignore
+    except ImportError as exc:
+        raise RuntimeError(
+            "PRIMUS_FLA_MLA_ATTN=1 was set, but `flash_attn` is not "
+            "importable in this Python.  Install it with "
+            "`pip install flash-attn --no-build-isolation` and retry."
+        ) from exc
+    _flash_attn_func = flash_attn_func
+    return _flash_attn_func
+
+
+class FLAFlashAttention(MegatronModule):
+    """`core_attention` plug-in that routes through `flash_attn_func`.
+
+    Constructor signature is intentionally a superset of
+    ``TEDotProductAttention.__init__`` so ``build_module(...)`` from
+    ``MLASelfAttention`` can swap them without any other change.
+    """
+
+    def __init__(
+        self,
+        config,
+        layer_number: int,
+        attn_mask_type: AttnMaskType = AttnMaskType.causal,
+        attention_type: str = "self",
+        softmax_scale: Optional[float] = None,
+        k_channels: Optional[int] = None,
+        v_channels: Optional[int] = None,
+        cp_comm_type: Optional[str] = None,
+        pg_collection: Any = None,
+        **_unused_kwargs: Any,
+    ) -> None:
+        super().__init__(config=config)
+        self.layer_number = layer_number
+        self.attn_mask_type = attn_mask_type
+        self.attention_type = attention_type
+        self.softmax_scale = softmax_scale
+        self.k_channels = k_channels
+        self.v_channels = v_channels
+        # Exposed for MLA's optional qk_clip code path (disabled by default
+        # in our configs).  Keeping it on the instance silences AttributeError.
+        self.current_max_attn_logits = None
+
+        # Eager import so failures surface at model build time, not at
+        # the first forward.
+        _load_flash_attn()
+
+        global _BANNER_PRINTED
+        if not _BANNER_PRINTED:
+            try:
+                import flash_attn as _fa
+
+                _ver = getattr(_fa, "__version__", "unknown")
+            except Exception:
+                _ver = "unknown"
+            from megatron.training import get_args as _ga
+
+            _mla_val = getattr(_ga(), "fla_mla_attn", "")
+            if _mla_val == "1":
+                _reason = "args.fla_mla_attn='1'"
+            elif not _mla_val:
+                _reason = (
+                    f"auto-enabled (flash_attn {_ver} > TE max "
+                    f"{'.'.join(str(x) for x in _TE_FLASH_ATTN_MAX_SUPPORTED)})"
+                )
+            else:
+                _reason = f"args.fla_mla_attn={_mla_val!r}"
+            _msg = (
+                f"[PRIMUS_FLA_MLA_ATTN] FLAFlashAttention active "
+                f"(layer={layer_number}, softmax_scale={softmax_scale}, "
+                f"k_channels={k_channels}, v_channels={v_channels}, "
+                f"flash_attn={_ver}, reason={_reason}). "
+                f"This banner prints once per worker."
+            )
+            # Megatron silently consumes plain `print()` from rank-non-zero
+            # workers (and sometimes from rank-0 once its logger is set up),
+            # so emit to stderr -- which the run_pretrain.sh tee pipeline
+            # still captures -- AND drop a marker file so activation is
+            # provable even if all stdio gets eaten.
+            print(_msg, file=sys.stderr, flush=True)
+            try:
+                rank = int(os.environ.get("RANK", "-1"))
+                marker = f"/tmp/primus_fla_mla_attn_active.rank{rank}.txt"
+                with open(marker, "w") as fh:
+                    fh.write(_msg + "\n")
+                    fh.write(f"pid={os.getpid()} ts={time.time()}\n")
+            except Exception:
+                pass
+            _BANNER_PRINTED = True
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,  # ignored; causal=True
+        attn_mask_type: Optional[AttnMaskType] = None,
+        packed_seq_params: Any = None,
+        **_unused_kwargs: Any,
+    ) -> torch.Tensor:
+        # Megatron's MLA always builds with attn_mask_type=causal.  The
+        # `attention_mask` arg is unused for causal attention; we ignore it.
+        if packed_seq_params is not None:
+            raise NotImplementedError(
+                "FLAFlashAttention does not yet support packed_seq_params; "
+                "either disable PRIMUS_FLA_MLA_ATTN or run without "
+                "sequence packing."
+            )
+
+        mask_type = attn_mask_type if attn_mask_type is not None else self.attn_mask_type
+        if mask_type != AttnMaskType.causal:
+            raise NotImplementedError(
+                f"FLAFlashAttention only supports causal masking; got {mask_type}. "
+                "Disable PRIMUS_FLA_MLA_ATTN if you need a different mask."
+            )
+
+        flash_attn_func = _load_flash_attn()
+
+        # Megatron passes q/k/v as [s, b, h, d]; flash-attn expects [b, s, h, d].
+        # The `.contiguous()` is required after `.transpose(0, 1)` because
+        # flash-attn checks for contiguous last-dim-fastest layout.
+        q = query.transpose(0, 1).contiguous()
+        k = key.transpose(0, 1).contiguous()
+        v = value.transpose(0, 1).contiguous()
+
+        # FLA's MLA path:
+        #   o = flash_attn_func(q, k, v, causal=True, softmax_scale=…)
+        out = flash_attn_func(
+            q,
+            k,
+            v,
+            dropout_p=0.0,
+            softmax_scale=self.softmax_scale,
+            causal=True,
+        )
+        # out: [b, s, h, d_v]  ->  Megatron expects [s, b, h*d_v]
+        out = out.transpose(0, 1).contiguous()
+        s, b, h, d_v = out.shape
+        return out.view(s, b, h * d_v)
diff --git a/primus/backends/megatron/core/transformer/hadamard_rotation.py b/primus/backends/megatron/core/transformer/hadamard_rotation.py
new file mode 100644
index 000000000..30be35e1f
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/hadamard_rotation.py
@@ -0,0 +1,83 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Hadamard rotation for the DeepSeek-V4 lightning indexer.
+
+The indexer scores ``ReLU(q_h . k_s)`` over compressed keys, and the reference
+rotates both operands by a normalised Hadamard matrix first (DeepSeek-V3.2's
+``inference/model.py``, and the open-source training reference's
+``rotate_activation``). The rotation is orthogonal, so it preserves the inner
+product exactly in infinite precision -- what it buys is spreading each
+coordinate's energy across the whole vector, so no single channel dominates the
+low-precision QK product the indexer is designed to run in.
+
+The reference calls into the ``fast_hadamard_transform`` extension and
+hard-requires BF16. That extension is not guaranteed present here, so this
+module prefers it when available and otherwise multiplies by a cached Sylvester
+matrix: the same transform, one GEMM, and it works on CPU and in FP32 so the
+behaviour is testable without a GPU.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple
+
+import torch
+
+try:  # pragma: no cover - depends on the image
+    from fast_hadamard_transform import hadamard_transform as _fast_hadamard_transform
+except Exception:  # pragma: no cover
+    _fast_hadamard_transform = None
+
+__all__ = ["rotate_activation"]
+
+# Sylvester matrices keyed by (n, device, dtype). Tiny (128x128 at the V4
+# indexer width) and reused every layer, every step.
+_MATRIX_CACHE: Dict[Tuple[int, str, torch.dtype], torch.Tensor] = {}
+
+# The extension only accepts half precision.
+_FAST_PATH_DTYPES = (torch.float16, torch.bfloat16)
+
+
+def _sylvester(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Unnormalised ``n x n`` Sylvester Hadamard matrix (entries +-1).
+
+    Symmetric, so ``x @ H`` and ``x @ H.T`` agree and the caller does not have
+    to care which side it multiplies on.
+    """
+    key = (n, str(device), dtype)
+    cached = _MATRIX_CACHE.get(key)
+    if cached is not None:
+        return cached
+
+    h = torch.ones((1, 1), device=device, dtype=dtype)
+    while h.shape[0] < n:
+        top = torch.cat([h, h], dim=1)
+        bottom = torch.cat([h, -h], dim=1)
+        h = torch.cat([top, bottom], dim=0)
+    if h.shape[0] != n:
+        raise ValueError(f"Hadamard rotation needs a power-of-two width, got {n}")
+
+    _MATRIX_CACHE[key] = h
+    return h
+
+
+def rotate_activation(x: torch.Tensor, *, scale: Optional[float] = None) -> torch.Tensor:
+    """Apply the normalised Hadamard rotation over ``x``'s last dimension.
+
+    ``scale`` defaults to ``n ** -0.5``, which makes the transform orthonormal
+    (and matches ``rotate_activation`` in the reference).
+    """
+    n = x.shape[-1]
+    if n & (n - 1) != 0:
+        raise ValueError(f"Hadamard rotation needs a power-of-two width, got {n}")
+    norm = n**-0.5 if scale is None else scale
+
+    if _fast_hadamard_transform is not None and x.is_cuda and x.dtype in _FAST_PATH_DTYPES:
+        return _fast_hadamard_transform(x, scale=norm)
+
+    h = _sylvester(n, x.device, x.dtype)
+    return torch.matmul(x, h) * norm
diff --git a/primus/backends/megatron/core/transformer/hyper_connection.py b/primus/backends/megatron/core/transformer/hyper_connection.py
new file mode 100644
index 000000000..d3020e3f0
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/hyper_connection.py
@@ -0,0 +1,468 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Manifold-Constrained Hyper-Connections (mHC) for DeepSeek-V4.
+
+Reference: deepseek-v4/develop/techblog/01-deepseek-v4-architecture-deep-dive.md
+section 2 ("mHC: Manifold-Constrained Hyper-Connections") and section 9.3.
+
+Three modules live here:
+
+* :class:`HyperMixer` — per-layer mixer used twice per block (once before /
+  after the attention sub-block, once for the FFN sub-block). Produces
+  ``(pre, post, comb)`` triplets and exposes ``collapse`` / ``expand``
+  helpers for the surrounding block to drive the K parallel hidden streams.
+
+* :class:`HyperHead` — final collapse used once at the end of the main trunk
+  (and once *per MTP layer*, with its own copy). Sigmoid-weighted sum, no
+  Sinkhorn.
+
+* :func:`sinkhorn_normalize` — alternating row/column normalization, kept in
+  fp32 for stability (per RedNote slide 9 and NeMo port pitfall #3).
+
+Phase 4 contract:
+* Plain ``nn.Linear`` for the projection ``fn``. TP-friendly variants
+  (``ColumnParallelLinear``) come in Phase 6 once the rest of the V4 path
+  is correct end-to-end on a single device.
+* All HC parameters (``fn.weight``, ``scale``, ``base``) live in fp32.
+  The block is responsible for passing in/out tensors in whatever activation
+  dtype it uses; the module up-casts internally to fp32 around the Sinkhorn
+  region and casts back at the end.
+
+Plan-5 P29 (RESCOPED — see ``deepseek-v4/develop/plan-5/02-phase-details.md``)
+adds a ``torch.compile`` fast path for :func:`sinkhorn_normalize`. The eager
+loop issues 1 + 2*(n_iters - 1) separate fp32 ``aten::sum`` reductions per
+call; at V4-Flash production widths each reduction runs at ~250x over the
+memory-bound floor because HIP's default ``reduce_kernel<512, 1, ...>`` is
+sized for huge reductions and our ``[1, 4096, 4, 4] -> [1, 4096, 4, 1]``
+shape has only 4 elements per output. The compiled path collapses every
+sum / divide / broadcast into one Inductor-fused Triton kernel
+(``fullgraph=True, dynamic=False``). The compiled callable is cached
+module-globally on ``(n_iters, eps, dtype)`` so each combination compiles
+exactly once per process.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Callable, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_collapse import (
+    hc_collapse_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import (
+    hc_expand_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import (
+    is_triton_kernel_supported as _hc_expand_supported,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import (
+    is_triton_path_enabled as _hc_expand_enabled,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import (
+    hc_glue_compute_tail_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import (
+    is_triton_kernel_supported as _hc_glue_supported,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import (
+    is_triton_path_enabled as _hc_glue_enabled,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import (
+    fused_rms_norm,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import (
+    SinkhornNormalizeFn,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+# ---------------------------------------------------------------------------
+# Plan-5 P29 fast path: torch.compile-fused Sinkhorn-Knopp.
+#
+# Cache keyed on ``(n_iters, eps, dtype)``.  Each freshly-decorated callable
+# uses ``@torch.compile(fullgraph=True, dynamic=True)`` so a SINGLE compiled
+# artefact handles every input shape variation Inductor generates code that
+# accepts shape as a runtime argument.  The dynamic path is the right
+# trade-off for V4 because:
+#
+# * In production each rank sees exactly one Sinkhorn input shape
+#   (``[B, S, K, K] = [1, 4096, 4, 4]`` at V4-Flash) so ``dynamic=True``
+#   pays no specialization cost vs ``dynamic=False``.
+# * Multi-shape harnesses (unit tests, MTP heads with different leading
+#   dims) all hit the same compiled artefact, so we never thrash Dynamo's
+#   recompile_limit.  ``dynamic=False`` would force one Dynamo
+#   specialization per shape AND ALL CLOSURES SHARE THE SAME CODE OBJECT
+#   (closures from the same factory function inherit the factory's code
+#   object), which means Dynamo's cache_size_limit (default 8) is shared
+#   across every cached callable.  After ~8 distinct ``(shape, stride,
+#   requires_grad)`` combinations the limit is hit and ``fullgraph=True``
+#   raises ``FailOnRecompileLimitHit`` even though our cache is fine —
+#   the cache key just is not what Dynamo's internal cache is keyed on.
+#
+# ``in_dtype`` still participates in the cache key because the trailing
+# cast back to the input activation dtype changes the Inductor IR; running
+# bf16 and fp32 callers through the same compiled artefact would force a
+# miscompile or a redundant recompile.
+# ---------------------------------------------------------------------------
+
+_SinkhornCacheKey = tuple[int, float, torch.dtype]
+_compiled_sinkhorn_cache: dict[_SinkhornCacheKey, Callable[[torch.Tensor], torch.Tensor]] = {}
+
+
+def _build_compiled_sinkhorn(
+    n_iters: int,
+    eps: float,
+    in_dtype: torch.dtype,
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Build (and torch.compile) one Sinkhorn-Knopp implementation specialised
+    on ``(n_iters, eps, in_dtype)`` but generic over input shape.
+
+    The body is the same algorithm as the eager :func:`sinkhorn_normalize`
+    path but written so Inductor sees one straight-line graph of fp32
+    sums / divides / broadcasts (Python ``for`` is unrolled at compile time
+    because ``n_iters`` is a closure-captured Python int).
+    """
+
+    @torch.compile(fullgraph=True, dynamic=True)
+    def _impl(logits: torch.Tensor) -> torch.Tensor:
+        m = logits.float()
+        m = m / (m.sum(dim=-2, keepdim=True) + eps)
+        for _ in range(max(n_iters - 1, 0)):
+            m = m / (m.sum(dim=-1, keepdim=True) + eps)
+            m = m / (m.sum(dim=-2, keepdim=True) + eps)
+        return m.to(in_dtype)
+
+    return _impl
+
+
+def _get_compiled_sinkhorn(
+    n_iters: int,
+    eps: float,
+    in_dtype: torch.dtype,
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Cache-or-build accessor; returns the compiled callable for the given
+    Sinkhorn signature.  Shape is NOT part of the key — see module-level
+    cache notes (we use ``dynamic=True`` so one artefact handles every
+    shape).
+    """
+    key: _SinkhornCacheKey = (int(n_iters), float(eps), in_dtype)
+    fn = _compiled_sinkhorn_cache.get(key)
+    if fn is None:
+        fn = _build_compiled_sinkhorn(n_iters, eps, in_dtype)
+        _compiled_sinkhorn_cache[key] = fn
+    return fn
+
+
+def sinkhorn_normalize(
+    logits: torch.Tensor,
+    *,
+    n_iters: int = 20,
+    eps: float = 1e-6,
+    use_compiled: bool = False,
+    use_triton: bool = False,
+) -> torch.Tensor:
+    """Project a non-negative ``[..., K, K]`` matrix onto the doubly-stochastic
+    manifold via the Sinkhorn-Knopp algorithm.
+
+    The algorithm itself is the alternating row / column ``L1`` normalization:
+
+    .. code-block::
+
+        for _ in range(n_iters):
+            M /= M.sum(dim=-1, keepdim=True)   # rows  -> 1
+            M /= M.sum(dim=-2, keepdim=True)   # cols  -> 1
+
+    To match the released V4 reference (cf. ``compute_weights`` in the
+    techblog) we emit the **first** column-normalization once before the
+    standard alternating loop, so that the final iteration ends on a column
+    normalization — this matches the (pre, post, comb) projection convention
+    that the surrounding block consumes.
+
+    Stability:
+    * the input must already be non-negative (typically ``softmax`` output
+      plus an ``eps`` floor)
+    * runs in fp32 regardless of the input dtype, then casts back
+
+    Args:
+        logits: ``[..., K, K]`` non-negative.
+        n_iters: total Sinkhorn iterations (= 1 priming column step
+            + ``n_iters - 1`` row/col cycles).
+        eps: numerical floor to prevent divide-by-zero.
+        use_compiled: when ``True``, dispatch to the cached
+            ``torch.compile(fullgraph=True, dynamic=False)`` build of the
+            same loop (plan-5 P29 RESCOPED).  The compiled callable
+            collapses the 1 + 2*(n_iters - 1) fp32 ``aten::sum`` launches
+            (and their divides / broadcasts) into one Inductor-fused Triton
+            kernel; AOT autograd handles the BWD.  Numerical contract is
+            unchanged (algorithm is identical; only the kernel boundary
+            moves).  Default ``False`` until the V4 config flag
+            ``use_v4_compiled_sinkhorn`` is flipped on.
+        use_triton: plan-6 P36 hand-rolled Triton FWD/BWD path.  When
+            ``True`` (or when the ``PRIMUS_SINKHORN_TRITON`` env knob is
+            not ``"0"`` and the input is supported), dispatch to
+            :class:`primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn.SinkhornNormalizeFn`.
+            The Triton path runs the full 1 + 2*(n_iters - 1) normalize
+            trajectory in registers per row of the leading axis and emits
+            exactly **one** FWD kernel + **one** BWD kernel -- no Dynamo
+            bookkeeping per call.  Routing precedence at the call site:
+            ``use_triton or PRIMUS_SINKHORN_TRITON=1 > use_compiled >
+            eager``.  Defaults to ``False``; production turns it on via
+            ``PRIMUS_SINKHORN_TRITON=1`` (default ON in the proxy
+            launcher and in :func:`is_triton_path_enabled`).
+
+    Returns:
+        Approximately doubly-stochastic ``[..., K, K]`` (same dtype as input).
+    """
+    # Plan-6 P36: routing precedence is Triton > compiled > eager.  Env
+    # knob is default-ON; explicit `use_triton=True` forces the path
+    # even when the env is off (used by the G39 tests).  The Triton path
+    # gracefully falls back when the input shape / device is unsupported
+    # (`is_triton_kernel_supported`), so callers never have to special-
+    # case K out-of-range or CPU inputs.
+    if (use_triton or is_triton_path_enabled()) and is_triton_kernel_supported(logits):
+        return SinkhornNormalizeFn.apply(logits, n_iters, eps)
+    if use_compiled:
+        return _get_compiled_sinkhorn(n_iters, eps, logits.dtype)(logits)
+    in_dtype = logits.dtype
+    m = logits.float()
+    m = m / (m.sum(dim=-2, keepdim=True) + eps)
+    for _ in range(max(n_iters - 1, 0)):
+        m = m / (m.sum(dim=-1, keepdim=True) + eps)
+        m = m / (m.sum(dim=-2, keepdim=True) + eps)
+    return m.to(in_dtype)
+
+
+class HyperMixer(nn.Module):
+    """Per-layer mHC mixer.
+
+    Maintains the three projection scales / biases (``pre``, ``post``,
+    ``comb``) and a single packed ``Linear`` ``fn`` that produces
+    ``[..., (2 + K) * K]`` from ``[..., K * D]``.
+
+    Shapes (B and S are arbitrary; the mixer is agnostic to which is which):
+
+    * ``compute_weights(x)`` → ``pre [..., K]``, ``post [..., K]``,
+      ``comb [..., K, K]``
+    * ``collapse(x, pre)`` → ``[..., D]``
+    * ``expand(x, out, post, comb)`` → ``[..., K, D]``
+
+    Args:
+        hidden_size: per-stream feature dim ``D``.
+        hc_mult: number of parallel streams ``K``.
+        eps: floor for ``sigmoid(...) + eps`` and for Sinkhorn.
+        sinkhorn_iters: Sinkhorn iteration count (``20`` matches V4 release).
+        use_compiled_sinkhorn: plan-5 P29 (RESCOPED) fast path.  When
+            ``True`` the call to :func:`sinkhorn_normalize` inside
+            :meth:`compute_weights` dispatches to the
+            ``torch.compile(fullgraph=True, dynamic=False)`` build of the
+            same algorithm, kept in :data:`_compiled_sinkhorn_cache`.
+            Defaults to ``False`` so existing callers (and existing
+            checkpoints) are bit-equivalent.  The V4 block reads
+            ``config.use_v4_compiled_sinkhorn`` and forwards it here.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        eps: float = 1e-6,
+        sinkhorn_iters: int = 20,
+        use_compiled_sinkhorn: bool = False,
+    ) -> None:
+        super().__init__()
+        if hc_mult < 1:
+            raise ValueError(f"hc_mult must be >= 1, got {hc_mult}")
+
+        self.hidden_size = hidden_size
+        self.hc_mult = hc_mult
+        self.eps = eps
+        self.sinkhorn_iters = sinkhorn_iters
+        self.use_compiled_sinkhorn = bool(use_compiled_sinkhorn)
+
+        out_dim = (2 + hc_mult) * hc_mult
+        # All HC params kept in fp32; see techblog §2.2 pitfall #3.
+        self.fn = nn.Linear(hc_mult * hidden_size, out_dim, bias=False, dtype=torch.float32)
+        # Three independent scale scalars: one each for pre / post / comb.
+        self.scale = nn.Parameter(torch.ones(3, dtype=torch.float32))
+        # Bias terms — same partition as the linear output.
+        self.base = nn.Parameter(torch.zeros(out_dim, dtype=torch.float32))
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        # NormalInitGain-style init keeps the post-rms-scale logits well-behaved.
+        nn.init.normal_(self.fn.weight, std=1.0 / math.sqrt(self.hc_mult * self.hidden_size))
+        nn.init.zeros_(self.base)
+        nn.init.ones_(self.scale)
+
+    # ---- internals -------------------------------------------------------
+
+    def _packed_logits(self, x: torch.Tensor) -> torch.Tensor:
+        """Pack streams along the last dim, RMS-normalize, project to ``out_dim``.
+
+        ``x``: ``[..., K, D]`` → ``flat`` ``[..., K*D]`` → ``logits`` ``[..., out_dim]``.
+        Done in fp32 for stability (HC parameters are fp32 anyway).
+
+        Small-kernel-fusion (2026-07-03): the parameter-less RMS over the
+        packed ``K*D`` axis (bf16->fp32 cast + pow + mean + rsqrt + mul) is
+        fused into one Triton FWD + one BWD kernel producing the fp32
+        normalized tensor directly for ``F.linear``. Gated by
+        ``PRIMUS_RMSNORM_TRITON`` (default on).
+        """
+        flat = x.flatten(-2)
+        normed = fused_rms_norm(flat, None, eps=self.eps, mid_cast=False, out_dtype=torch.float32)
+        logits = F.linear(normed, self.fn.weight.to(dtype=normed.dtype))
+        return logits
+
+    # ---- public API ------------------------------------------------------
+
+    def compute_weights(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return ``(pre, post, comb)`` for the K parallel streams ``x``.
+
+        ``x``: ``[..., K, D]`` (typically ``[B, S, K, D]`` or ``[S, B, K, D]``).
+        """
+        K = self.hc_mult
+        logits = self._packed_logits(x)  # [..., (2+K)*K], fp32
+
+        out_dtype = x.dtype
+        # Plan-6 P37: fuse the elemwise tail (slice + scale + base +
+        # sigmoid/softmax + eps) into one Triton kernel when the env
+        # knob is on and the shape is supported.  Falls back to the
+        # eager chain for the unsupported case (CPU input, K not in
+        # the supported set, etc.).
+        if _hc_glue_enabled() and _hc_glue_supported(logits, K):
+            pre, post, comb = hc_glue_compute_tail_triton(
+                logits,
+                self.scale,
+                self.base,
+                K=K,
+                eps=self.eps,
+                out_dtype=torch.float32,
+            )
+        else:
+            pre_logit = logits[..., :K] * self.scale[0] + self.base[:K]
+            post_logit = logits[..., K : 2 * K] * self.scale[1] + self.base[K : 2 * K]
+            comb_logit = logits[..., 2 * K :].view(*logits.shape[:-1], K, K) * self.scale[2] + self.base[
+                2 * K :
+            ].view(K, K)
+            pre = torch.sigmoid(pre_logit) + self.eps  # (eps, 1+eps]
+            post = 2.0 * torch.sigmoid(post_logit)  # (0, 2)  no eps
+            comb = torch.softmax(comb_logit, dim=-1) + self.eps
+
+        comb = sinkhorn_normalize(
+            comb,
+            n_iters=self.sinkhorn_iters,
+            eps=self.eps,
+            use_compiled=self.use_compiled_sinkhorn,
+        )
+
+        # Cast back to the activation dtype to match downstream block compute.
+        return pre.to(out_dtype), post.to(out_dtype), comb.to(out_dtype)
+
+    @staticmethod
+    def collapse(x: torch.Tensor, pre: torch.Tensor) -> torch.Tensor:
+        """Collapse K streams into 1 via ``pre`` weights.
+
+        ``x``: ``[..., K, D]``;  ``pre``: ``[..., K]``;
+        returns ``[..., D]``.
+
+        Small-kernel-fusion (2026-07-03): the eager
+        ``(pre.unsqueeze(-1) * x).sum(-2)`` (broadcast-mul into a full
+        ``[..., K, D]`` temporary + reduce = 2 kernels + ``K*D`` extra HBM
+        traffic) is fused into one Triton FWD + one BWD kernel. Symmetric to
+        the already-shipped ``expand`` fusion. Gated by
+        ``PRIMUS_HC_COLLAPSE_TRITON`` (default on); eager fallback on
+        CPU / unsupported K.
+        """
+        return hc_collapse_triton(x, pre)
+
+    @staticmethod
+    def expand(
+        x: torch.Tensor,
+        out: torch.Tensor,
+        post: torch.Tensor,
+        comb: torch.Tensor,
+    ) -> torch.Tensor:
+        """Write the sub-block output ``out`` back to K streams.
+
+        ``new_stream[h] = post[h] * out + Σ_k comb[h, k] * x[k]``
+
+        Shapes:
+        * ``x``: ``[..., K, D]`` — current K streams
+        * ``out``: ``[..., D]`` — sub-block (attn or FFN) output
+        * ``post``: ``[..., K]``
+        * ``comb``: ``[..., K, K]``
+
+        Returns ``[..., K, D]``.
+        """
+        # Triton-fused expand; falls back to eager for unsupported configs.
+        if _hc_expand_enabled() and _hc_expand_supported(x, post, comb):
+            return hc_expand_triton(x, out, post, comb)
+        # post[..., K] * out[..., D] -> [..., K, D]
+        write = post.unsqueeze(-1) * out.unsqueeze(-2)
+        # comb[..., K, K] @ x[..., K, D] -> [..., K, D]
+        mix = torch.matmul(comb, x)
+        return write + mix
+
+
+class HyperHead(nn.Module):
+    """Final collapse from K streams → 1 stream.
+
+    Used once at the end of the main trunk, and once *per* MTP layer (each
+    with its own copy — ``num_nextn_predict_layers`` separate heads). Plain
+    sigmoid-weighted sum; no Sinkhorn.
+
+    Args:
+        hidden_size: per-stream feature dim ``D``.
+        hc_mult: number of input streams ``K``.
+        eps: floor for ``sigmoid(...) + eps``.
+    """
+
+    def __init__(self, *, hidden_size: int, hc_mult: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        if hc_mult < 1:
+            raise ValueError(f"hc_mult must be >= 1, got {hc_mult}")
+        self.hidden_size = hidden_size
+        self.hc_mult = hc_mult
+        self.eps = eps
+
+        self.fn = nn.Linear(hc_mult * hidden_size, hc_mult, bias=False, dtype=torch.float32)
+        self.scale = nn.Parameter(torch.ones((), dtype=torch.float32))
+        self.base = nn.Parameter(torch.zeros(hc_mult, dtype=torch.float32))
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.fn.weight, std=1.0 / math.sqrt(self.hc_mult * self.hidden_size))
+        nn.init.zeros_(self.base)
+        nn.init.ones_(self.scale)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """``x``: ``[..., K, D]`` → ``[..., D]``.
+
+        Small-kernel-fusion (2026-07-03): the parameter-less RMS over the
+        packed ``K*D`` axis is fused into one Triton FWD + one BWD kernel
+        (fp32 output for ``F.linear``). Gated by ``PRIMUS_RMSNORM_TRITON``.
+        """
+        flat = x.flatten(-2)
+        normed = fused_rms_norm(flat, None, eps=self.eps, mid_cast=False, out_dtype=torch.float32)
+        mixes = F.linear(normed, self.fn.weight.to(dtype=normed.dtype))  # [..., K]
+        pre = torch.sigmoid(mixes * self.scale + self.base) + self.eps
+        return (pre.unsqueeze(-1) * x).sum(dim=-2).to(x.dtype)
+
+
+__all__ = [
+    "HyperMixer",
+    "HyperHead",
+    "sinkhorn_normalize",
+]
diff --git a/primus/backends/megatron/core/transformer/indexer.py b/primus/backends/megatron/core/transformer/indexer.py
new file mode 100644
index 000000000..3ac67f7ec
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/indexer.py
@@ -0,0 +1,549 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+r"""
+DeepSeek-V4 Indexer (sparse position selector for CSA).
+
+Reference: techblog §1.4 ("Indexer: CSA's Sparse Selector").
+
+The Indexer is **only** used by CSA layers (``compress_ratio == 4``). For
+each query position ``t`` it picks ``index_topk`` compressed-KV positions
+``s`` (out of all compressed positions ``[0, P)`` where ``P = S // ratio``).
+
+Math (from the techblog):
+
+.. math::
+
+    q^Q_t = h_t W^{DQ},\quad q^I_{t,h} = q^Q_t W^{IUQ}_h,\quad
+    w^I_{t,h} = h_t W^w_h
+
+.. math::
+
+    I_{t,s} = \\sum_h w^I_{t,h}\\cdot \\mathrm{ReLU}(q^I_{t,h}\\cdot K^{IComp}_s)
+
+.. math::
+
+    \\mathrm{topk\\_idxs}_t = \\mathrm{argTopK}_s\\,I_{t,s}
+
+The Indexer carries its **own** mini-Compressor (``index_head_dim``,
+``index_n_heads``); the ``K^{IComp}`` it produces is independent of the
+main attention's compressed KV pool. It is only used to **select** top-k
+positions; the actual values fetched into main attention come from the
+main Compressor in the surrounding CSA layer.
+
+Both scoring operands are prepared the way the reference prepares them
+before the dot product:
+
+* **Partial RoPE**, at the compressed-branch base -- queries at their own
+  token positions, compressed keys at ``s * compress_ratio`` (the window's
+  first token), so the two live in one coordinate system.
+* **Hadamard rotation** (``rotate_activation``) on top. Orthogonal, so it
+  leaves the inner product alone in exact arithmetic; what it buys is
+  spreading each coordinate's energy across the vector so no single channel
+  dominates the low-precision QK product. Note the reference rotates *only*
+  here -- the main compressed KV pool is built with ``rotate=False``.
+
+Phase 4 contract:
+* Plain ``nn.Linear`` projections (TP integration in P6).
+* Causal masking: positions ``s`` whose start raw-token index exceeds the
+  query's raw-token index get a value of ``-inf`` so they cannot be
+  selected. Out-of-range positions are returned as ``-1`` in the output
+  ``topk_idxs`` so the caller can treat them as "no key".
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+def _is_rank0() -> bool:
+    try:
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return torch.distributed.get_rank() == 0
+    except Exception:
+        pass
+    return True
+
+
+from primus.backends.megatron.core.transformer.compressor import Compressor
+from primus.backends.megatron.core.transformer.dual_rope import (
+    apply_interleaved_partial_rope,
+)
+from primus.backends.megatron.core.transformer.hadamard_rotation import (
+    rotate_activation,
+)
+
+# E4M3 finite max magnitude (float8_e4m3fn): largest representable value.
+_FP8_E4M3_MAX = 448.0
+
+
+def fake_quantize_fp8_e4m3(x: torch.Tensor) -> torch.Tensor:
+    """Per-tensor dynamic FP8 (E4M3) fake-quantization.
+
+    Scales ``x`` so its max magnitude maps to the E4M3 finite range, rounds
+    through ``torch.float8_e4m3fn``, then dequantizes back to ``x.dtype``. This
+    simulates the precision of an FP8 QK GEMM input while keeping the matmul
+    itself in the activation dtype (QAT-style "simulated FP8" path). Returns
+    ``x`` unchanged when the platform/torch build lacks ``float8_e4m3fn`` or
+    when ``x`` is all-zero (degenerate scale).
+    """
+    if not hasattr(torch, "float8_e4m3fn"):
+        return x
+    orig_dtype = x.dtype
+    amax = x.detach().abs().amax()
+    if not torch.isfinite(amax) or float(amax) <= 0.0:
+        return x
+    scale = (_FP8_E4M3_MAX / amax).to(x.dtype)
+    x_scaled = torch.clamp(x * scale, -_FP8_E4M3_MAX, _FP8_E4M3_MAX)
+    x_fp8 = x_scaled.to(torch.float8_e4m3fn)
+    return x_fp8.to(orig_dtype) / scale
+
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import (
+    indexer_score_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import (
+    is_triton_kernel_supported as _indexer_triton_full_supported,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import (
+    is_triton_path_enabled as _indexer_triton_full_enabled,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import (
+    indexer_score_post_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import (
+    is_triton_kernel_supported as _indexer_tail_triton_supported,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import (
+    is_triton_path_enabled as _indexer_tail_triton_enabled,
+)
+
+# MXFP4 block size (E2M1 data + E8M0 per-32 block scales).
+_MXFP4_BLOCK = 32
+
+
+def _indexer_fp4_enabled() -> bool:
+    """True iff PRIMUS_INDEXER_FP4 == "1" (default off): run the CSA-indexer QK in MXFP4."""
+    return os.environ.get("PRIMUS_INDEXER_FP4", "0") == "1"
+
+
+def _fp4_qk_gemm(q_i: torch.Tensor, k_icomp: torch.Tensor) -> torch.Tensor:
+    """Real MXFP4 indexer QK: per-batch [S*H,Hd] @ [P,Hd]^T (NT, trans_b) -> [B,S,H,P].
+
+    hipBLASLt FP4 needs K=Hd%128, M,N%16; force PRIMUS_TURBO_GEMM_BACKEND=FP4:HIPBLASLT.
+    """
+    import primus_turbo.pytorch as pt
+    from primus_turbo.pytorch.core.low_precision import (
+        Float4QuantConfig,
+        Format,
+        ScaleDtype,
+        ScalingGranularity,
+    )
+
+    cfg = Float4QuantConfig(
+        format=Format.E2M1_X2,
+        granularity=ScalingGranularity.MX_BLOCKWISE,
+        block_size=_MXFP4_BLOCK,
+        scale_dtype=ScaleDtype.E8M0,
+    )
+    B, S, H, Hd = q_i.shape
+    P = k_icomp.shape[1]
+    outs = []
+    for b in range(B):
+        a = q_i[b].reshape(S * H, Hd).contiguous()  # [S*H, Hd]
+        bk = k_icomp[b].contiguous()  # [P, Hd]
+        o = pt.ops.gemm_fp4(a, bk, trans_b=True, config=cfg)  # [S*H, P]
+        outs.append(o.view(1, S, H, P))
+    return torch.cat(outs, dim=0)
+
+
+def _indexer_fp8_proj_enabled() -> bool:
+    """Run the indexer projections (w_dq/w_iuq/w_w) in MXFP8 (default off).
+
+    Gated by its own knob, deliberately **not** by the attention-projection flag
+    ``PRIMUS_V4_FP8_ATTN_PROJ``. The open-source reference wraps the indexer's
+    weight projection in an fp8-disabled context so it stays high precision even
+    when the enclosing layer runs FP8: the indexer decides *which* compressed KV
+    entries each query reads, so error there changes the selection instead of
+    perturbing a value. Quantizing the attention projections and quantizing the
+    selector are separate decisions, so one no longer implies the other.
+
+    Only fires inside turbo-fp8. The linears are duplicated (no TP shard), so
+    fp8 is safe at any TP.
+    """
+    if os.environ.get("PRIMUS_V4_FP8_INDEXER_PROJ", "0") != "1":
+        return False
+    try:
+        from primus.backends.megatron.core.extensions.primus_turbo import (
+            PrimusTurboLowPrecisionGlobalStateManager as _M,
+        )
+
+        return _M.is_turbo_fp8_enabled()
+    except Exception:
+        return False
+
+
+def _fp8_linear(lin: nn.Linear, x: torch.Tensor) -> torch.Tensor:
+    """MXFP8 apply of an ``nn.Linear`` (weight [out,in], no bias): y = x @ Wᵀ."""
+    import primus_turbo.pytorch as pt
+
+    from primus.backends.megatron.core.extensions.primus_turbo import (
+        PrimusTurboLowPrecisionGlobalStateManager as _M,
+    )
+
+    cfg = _M.get_turbo_quant_config().data()
+    orig = x.shape
+    x2 = x.reshape(-1, orig[-1]).contiguous()
+    out = pt.ops.gemm_fp8(x2, lin.weight, trans_b=True, config=cfg)  # [*, out]
+    return out.reshape(*orig[:-1], out.shape[-1])
+
+
+class Indexer(nn.Module):
+    """Sparse position selector for CSA.
+
+    Args:
+        hidden_size: input feature dim ``D`` (same as main attention).
+        index_head_dim: head dim used by the mini-Compressor and the
+            low-rank query projection.
+        index_n_heads: number of indexer "heads".
+        index_topk: number of compressed positions to select per query.
+        compress_ratio: ratio ``m`` of the mini-Compressor (matches the
+            main Compressor of the surrounding CSA layer; usually ``4``).
+        dq_rank: rank of the shared low-rank query projection ``W^{DQ}``.
+            Defaults to ``index_head_dim`` (the V4 reference doesn't expose
+            a separate setting; ``W^{IUQ}_h`` then projects from ``dq_rank``
+            to ``index_head_dim``).
+        rope: the surrounding CSA layer's compressed-branch
+            :class:`~primus.backends.megatron.core.transformer.dual_rope.RoPECache`.
+            Queries and compressed indexer keys are rotated with it before
+            scoring, exactly as in the reference. ``None`` disables the
+            rotation entirely and is only meant for isolated unit tests --
+            :meth:`DeepseekV4Attention._build_indexer` always supplies it.
+        rotary_dim: partial-RoPE width (``qk_pos_emb_head_dim``, V4 = 64),
+            applied to the trailing channels of the ``index_head_dim``-wide
+            queries and keys.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        index_head_dim: int,
+        index_n_heads: int,
+        index_topk: int,
+        compress_ratio: int = 4,
+        dq_rank: int = None,
+        use_fp8_qk: bool = False,
+        rope: Optional[nn.Module] = None,
+        rotary_dim: int = 0,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.index_head_dim = index_head_dim
+        self.index_n_heads = index_n_heads
+        self.index_topk = index_topk
+        self.compress_ratio = compress_ratio
+        self.dq_rank = dq_rank if dq_rank is not None else index_head_dim
+
+        # Held by reference (list-wrapped) so the shared RoPE cache is not
+        # re-registered as a submodule of every CSA layer's indexer.
+        self._rope = [rope] if rope is not None else []
+        self.rotary_dim = int(rotary_dim) if rope is not None else 0
+        if self.rotary_dim > self.index_head_dim:
+            raise ValueError(
+                f"rotary_dim ({self.rotary_dim}) must be <= index_head_dim ({self.index_head_dim})"
+            )
+        if self.rotary_dim and index_head_dim & (index_head_dim - 1) != 0:
+            # The Hadamard rotation spans the whole head, so fail here with the
+            # config name rather than deep inside the transform.
+            raise ValueError(
+                f"index_head_dim must be a power of two for the Hadamard rotation, got {index_head_dim}"
+            )
+        # FP8 (E4M3) fake-quant of the QK scoring inputs (V4 low-precision
+        # indexer QK path). See ``fake_quantize_fp8_e4m3`` / config flag
+        # ``use_v4_fp8_indexer``.
+        self.use_fp8_qk = bool(use_fp8_qk)
+        if self.use_fp8_qk and _is_rank0():
+            logger.info(
+                "[V4-Indexer] FP8 (E4M3) QK scoring path ENABLED "
+                "(query/compressed-key activations fake-quantized; "
+                "BF16 index-score + top-k preserved)."
+            )
+
+        # W^{DQ} (hidden->dq_rank) and W^w (hidden->n_heads) both consume `hidden`,
+        # so fuse them into ONE GEMM (default-on); split the output. W^{IUQ} stays
+        # separate (it consumes q_q, sequentially). PRIMUS_INDEXER_FUSE_PROJ=0 keeps
+        # the two separate linears.
+        self._fuse_qw_proj = os.environ.get("PRIMUS_INDEXER_FUSE_PROJ", "1") != "0"
+        if self._fuse_qw_proj:
+            self.w_dq_w = nn.Linear(hidden_size, self.dq_rank + index_n_heads, bias=False)
+        else:
+            # W^{DQ}: low-rank query down-projection.
+            self.w_dq = nn.Linear(hidden_size, self.dq_rank, bias=False)
+            # W^w_h: per-head scalar weight.
+            self.w_w = nn.Linear(hidden_size, index_n_heads, bias=False)
+        # W^{IUQ}_h: per-head up-projection from dq_rank → index_head_dim.
+        self.w_iuq = nn.Linear(self.dq_rank, index_n_heads * index_head_dim, bias=False)
+
+        # Temperature of the index score I_{t,s}. Folded into the per-head
+        # weights ``w_i`` so the head sum lands on the same scale the reference
+        # uses: it applies ``index_n_heads ** -0.5`` when building the per-head
+        # weights and ``index_head_dim ** -0.5`` (the indexer's own softmax
+        # scale) on the way into the loss.
+        #
+        # top-k is invariant under a positive constant, so this only matters
+        # once the scores enter a softmax -- which is exactly what the indexer
+        # distillation loss does. Without it the indexer distribution is
+        # sqrt(index_n_heads * index_head_dim) times too sharp (~90x at the V4
+        # widths) and the KL gradient is unusable.
+        self.score_scale: float = (index_n_heads**-0.5) * (index_head_dim**-0.5)
+
+        # Mini-Compressor producing K^{IComp}.
+        self.indexer_compressor = Compressor(
+            hidden_size=hidden_size,
+            head_dim=index_head_dim,
+            ratio=compress_ratio,
+        )
+
+    @property
+    def rope(self) -> Optional[nn.Module]:
+        """The compressed-branch RoPE cache, or ``None`` when not supplied."""
+        return self._rope[0] if self._rope else None
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Bridge checkpoints across the fused/unfused (w_dq, w_w) projection.
+
+        Old checkpoints store ``w_dq.weight`` + ``w_w.weight``; the fused path wants
+        ``w_dq_w.weight`` = ``cat([w_dq, w_w])`` (and vice-versa). Remap in-place so
+        either layout loads under either runtime setting.
+        """
+        dq_k, w_k, fused_k = prefix + "w_dq.weight", prefix + "w_w.weight", prefix + "w_dq_w.weight"
+        if self._fuse_qw_proj and dq_k in state_dict and fused_k not in state_dict:
+            state_dict[fused_k] = torch.cat([state_dict.pop(dq_k), state_dict.pop(w_k)], dim=0)
+        elif (not self._fuse_qw_proj) and fused_k in state_dict and dq_k not in state_dict:
+            w = state_dict.pop(fused_k)
+            state_dict[dq_k], state_dict[w_k] = w[: self.dq_rank], w[self.dq_rank :]
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    # ------------------------------------------------------------------
+
+    def _causal_mask(
+        self,
+        n_queries: int,
+        n_pool: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Return ``[n_queries, n_pool]`` mask: 0.0 if pool position ``s``
+        is allowed for query ``t``, ``-inf`` otherwise.
+
+        A compressed position ``s`` covers raw tokens ``[s*ratio, (s+1)*ratio)``;
+        a query at raw token ``t`` may attend to ``s`` iff its window
+        end ``(s+1)*ratio - 1 <= t``.
+        """
+        # The mask depends only on (n_queries, n_pool, compress_ratio, dtype) — all
+        # fixed per run — so cache it instead of rebuilding arange + where every
+        # call. PRIMUS_INDEXER_MASK_CACHE=0 forces the eager rebuild.
+        use_cache = os.environ.get("PRIMUS_INDEXER_MASK_CACHE", "1") != "0"
+        if use_cache:
+            cache = getattr(self, "_causal_mask_cache", None)
+            if cache is None:
+                cache = self._causal_mask_cache = {}
+            key = (n_queries, n_pool, device, dtype)
+            cached = cache.get(key)
+            if cached is not None:
+                return cached
+        t_idx = torch.arange(n_queries, device=device).unsqueeze(1)  # [t, 1]
+        s_end = (torch.arange(n_pool, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1  # [1, s]
+        allowed = s_end <= t_idx  # [t, s] bool
+        mask = torch.where(allowed, 0.0, float("-inf")).to(dtype)
+        if use_cache:
+            cache[key] = mask
+        return mask
+
+    # ------------------------------------------------------------------
+
+    def _rotate_keys(self, k_icomp: torch.Tensor) -> torch.Tensor:
+        """Partial RoPE + Hadamard on the compressed indexer keys.
+
+        ``k_icomp`` is ``[B, P, Hd]``. Compressed entry ``s`` covers the window
+        starting at original token ``s * compress_ratio``, so it is rotated
+        there -- the same coordinate system the queries use, and the same
+        stride sampling the main compressed pool uses.
+        """
+        rope = self.rope
+        if rope is None or self.rotary_dim == 0:
+            return k_icomp
+
+        P = k_icomp.shape[1]
+        cos, sin = rope.forward_arange(P, k_icomp.device, stride=self.compress_ratio)
+        # ``apply_interleaved_partial_rope`` inserts a singleton head axis into
+        # cos / sin, so give it a ``[B, P, 1, Hd]`` view to broadcast against.
+        rotated = apply_interleaved_partial_rope(
+            k_icomp.unsqueeze(2),
+            cos.unsqueeze(0).expand(k_icomp.shape[0], -1, -1),
+            sin.unsqueeze(0).expand(k_icomp.shape[0], -1, -1),
+            rotary_dim=self.rotary_dim,
+        )
+        return rotate_activation(rotated.squeeze(2))
+
+    def _rotate_queries(self, q_i: torch.Tensor, position_ids: Optional[torch.Tensor]) -> torch.Tensor:
+        """Partial RoPE + Hadamard on the indexer queries (``[B, S, H, Hd]``)."""
+        rope = self.rope
+        if rope is None or self.rotary_dim == 0:
+            return q_i
+
+        B, S = q_i.shape[0], q_i.shape[1]
+        if position_ids is None:
+            position_ids = torch.arange(S, device=q_i.device)
+        cos, sin = rope(position_ids)
+        if tuple(cos.shape[:-1]) != (B, S):
+            cos = cos.expand(B, S, -1)
+            sin = sin.expand(B, S, -1)
+        return rotate_activation(apply_interleaved_partial_rope(q_i, cos, sin, rotary_dim=self.rotary_dim))
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Select top-k compressed positions for each query.
+
+        Args:
+            hidden: ``[B, S, D]``.
+            position_ids: ``[B, S]`` or ``[S]`` absolute token positions, used
+                to rotate the indexer queries. Defaults to ``arange(S)``.
+
+        Returns:
+            ``(topk_idxs, topk_scores)`` where:
+            * ``topk_idxs`` ``[B, S, K]`` (long): selected pool positions
+              in ``[0, P)`` for valid slots, ``-1`` for masked / invalid.
+            * ``topk_scores`` ``[B, S, K]``: the selection scores ``I_{t,s}``
+              (``-inf`` for masked positions).
+        """
+        B, S, D = hidden.shape
+        assert S % self.compress_ratio == 0, (
+            f"Indexer: sequence length {S} not divisible by compress_ratio " f"{self.compress_ratio}"
+        )
+        K = self.index_topk
+        H = self.index_n_heads
+        Hd = self.index_head_dim
+
+        # 1) K^{IComp}: pool hidden via the mini-Compressor → [B, P, Hd], then
+        # rotate. The reference builds the indexer keys with a compressor
+        # configured as ``rotate=True`` (unlike the main one), and applies the
+        # compressed-position RoPE inside it.
+        k_icomp = self._rotate_keys(self.indexer_compressor(hidden))  # [B, P, Hd]
+        P = k_icomp.shape[1]
+        k_icomp = k_icomp.unsqueeze(2)  # [B, P, 1, Hd]
+
+        # 2) Per-head query and per-head weight.
+        # High precision by default, matching the reference: only the opt-in
+        # PRIMUS_V4_FP8_INDEXER_PROJ knob routes these through MXFP8. No TP
+        # gather/scatter (duplicated linears).
+        proj = _fp8_linear if _indexer_fp8_proj_enabled() else (lambda lin, x: lin(x))
+        if self._fuse_qw_proj:
+            dqw = proj(self.w_dq_w, hidden)  # [B, S, dq_rank + H] in one GEMM
+            q_q = dqw[..., : self.dq_rank]  # [B, S, dq_rank]
+            w_i = dqw[..., self.dq_rank :]  # [B, S, H]
+        else:
+            q_q = proj(self.w_dq, hidden)  # [B, S, dq_rank]
+            w_i = proj(self.w_w, hidden)  # [B, S, H]
+        q_i = proj(self.w_iuq, q_q).view(B, S, H, Hd)  # [B, S, H, Hd]
+        # Queries are rotated at their own token positions, then Hadamard-rotated
+        # like the keys, so the ReLU'd dot product sees both operands in the same
+        # basis the reference scores in.
+        q_i = self._rotate_queries(q_i, position_ids)
+
+        # Score temperature (see ``self.score_scale``). Applied on ``w_i``
+        # because every scoring branch below multiplies by it before summing
+        # over heads, so all of them inherit the scale from one place.
+        w_i = w_i * self.score_scale
+
+        # 3) Score I_{t,s} = Σ_h w_i[t,h] * ReLU(q_i[t,h] · k_icomp[s])
+        #    q_i [B,S,H,Hd] · k_icomp[B,P,Hd] → relu[B,S,H,P]; w_i[B,S,H,1] → sum over H
+        # 4) Causal mask + (effective topk capped at P).
+        #
+        # Dispatch precedence (P41 re-routing):
+        #   PRIMUS_INDEXER_TRITON=1       → post-einsum tail fused
+        #                                    (einsum stays eager / cuBLAS).
+        #   PRIMUS_INDEXER_TRITON_FULL=1  → legacy P38 full-fuse path
+        #                                    (einsum + tail in one kernel).
+        #   else                          → fully eager.
+        k_icomp_2d = k_icomp.squeeze(2)
+
+        # Indexer QK precision (both default OFF -> BF16 QK). FP8 (E4M3) fake-
+        # quantizes the operands before the normal score dispatch; FP4 (below)
+        # is a dedicated real-GEMM branch and takes precedence when both are set.
+        # The ReLU + per-head weight (``w_i``) + sum + causal mask + top-k stay
+        # in the activation dtype — only the QK operands are quantized.
+        if self.use_fp8_qk and not _indexer_fp4_enabled():
+            q_i = fake_quantize_fp8_e4m3(q_i)
+            k_icomp_2d = fake_quantize_fp8_e4m3(k_icomp_2d)
+
+        # Phase 5: FP4 CSA-indexer QK. Real MXFP4 GEMM for the QK product (paper
+        # §2.3.4/§5.2.1: "QK multiplied entirely in FP4"), then the eager
+        # ReLU/weight/sum tail (w_i + tail stay BF16/FP32 — only the QK is FP4).
+        if _indexer_fp4_enabled():
+            dot = _fp4_qk_gemm(q_i, k_icomp_2d)  # [B, S, H, P], real FP4 matmul
+            relu = F.relu(dot)
+            scores = (relu * w_i.unsqueeze(-1)).sum(dim=2)  # [B, S, P]
+            mask = self._causal_mask(S, P, scores.device, scores.dtype)  # [S, P]
+            scores = scores + mask.unsqueeze(0)  # [B, S, P]
+        elif _indexer_triton_full_enabled() and _indexer_triton_full_supported(q_i, k_icomp_2d, w_i):
+            scores = indexer_score_triton(
+                q_i,
+                k_icomp_2d,
+                w_i,
+                compress_ratio=self.compress_ratio,
+                out_dtype=hidden.dtype,
+            )
+        else:
+            dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp_2d)
+            if _indexer_tail_triton_enabled() and _indexer_tail_triton_supported(dot, w_i):
+                scores = indexer_score_post_triton(
+                    dot,
+                    w_i,
+                    compress_ratio=self.compress_ratio,
+                    out_dtype=hidden.dtype,
+                )
+            else:
+                relu = F.relu(dot)
+                scores = (relu * w_i.unsqueeze(-1)).sum(dim=2)  # [B, S, P]
+                mask = self._causal_mask(S, P, scores.device, scores.dtype)  # [S, P]
+                scores = scores + mask.unsqueeze(0)  # [B, S, P]
+
+        topk_eff = min(K, P)
+        topk_scores, topk_idxs = scores.topk(topk_eff, dim=-1)  # [B, S, topk_eff]
+
+        # 5) Replace selections that are still -inf (i.e. fewer than K valid
+        #    pool positions for very early queries) with sentinel ``-1`` so
+        #    callers can drop them.
+        invalid = torch.isneginf(topk_scores)
+        topk_idxs = torch.where(invalid, torch.full_like(topk_idxs, -1), topk_idxs)
+
+        # 6) Pad with -1 to exactly K columns if topk_eff < K (S smaller than
+        #    K * ratio in unit tests).
+        if topk_eff < K:
+            pad_idxs = torch.full((B, S, K - topk_eff), -1, dtype=topk_idxs.dtype, device=topk_idxs.device)
+            pad_scores = torch.full(
+                (B, S, K - topk_eff), float("-inf"), dtype=topk_scores.dtype, device=topk_scores.device
+            )
+            topk_idxs = torch.cat([topk_idxs, pad_idxs], dim=-1)
+            topk_scores = torch.cat([topk_scores, pad_scores], dim=-1)
+
+        return topk_idxs, topk_scores
+
+
+__all__ = ["Indexer"]
diff --git a/primus/backends/megatron/core/transformer/indexer_distill_loss.py b/primus/backends/megatron/core/transformer/indexer_distill_loss.py
new file mode 100644
index 000000000..839d3fb9f
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/indexer_distill_loss.py
@@ -0,0 +1,314 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Indexer distillation loss for DeepSeek-V4 CSA layers.
+
+The CSA lightning indexer picks which ``index_topk`` compressed KV entries each
+query attends to. ``topk`` is not differentiable, so without an auxiliary
+objective the indexer never receives a gradient and a from-scratch run selects
+essentially at random.
+
+DeepSeek-V3.2 (section 2.1) trains it by distillation: the indexer's score
+distribution is pulled towards the distribution the *real* attention places over
+the same entries, via ``KL(attention || indexer)``.
+
+This module implements the sparse variant -- the loss is evaluated only on the
+entries the indexer actually selected. That keeps the objective consistent with
+what the forward pass actually consumes, and because those entries are already
+gathered per query the computation stays in the ``[B, S, K]`` top-k space and
+never materialises the dense ``[B, H, S, P]`` score tensor.
+
+The gradient flows **one way**: into the indexer only. The target side (the
+attention queries and the compressed pool) is detached inside
+:func:`compute_indexer_distill_loss`, and the caller feeds the indexer from a
+detached hidden state, so the KL can neither reshape the attention distribution
+it is trying to imitate nor leak into the layers below.
+
+The loss is attached to the autograd graph with :class:`V4IndexerLossAutoScaler`,
+the same aux-loss autoscaler pattern this training stack already uses for the
+MoE auxiliary loss and the MTP loss: it passes a tensor through untouched in
+forward and seeds the auxiliary loss with a gradient of one in backward, so the
+aux objective backpropagates without having to be threaded through every forward
+return signature.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+
+__all__ = [
+    "INDEXER_DISTILL_LOSS_NAME",
+    "V4IndexerLossAutoScaler",
+    "compute_indexer_distill_loss",
+    "log_indexer_distill_loss",
+]
+
+# Guard for log(0) / division by zero.
+_EPS = 1e-10
+
+# Query rows per chunk when building the KL target; see ``_target_chunk_size``.
+_TARGET_CHUNK_ENV = "PRIMUS_V4_DISTILL_TARGET_CHUNK"
+_TARGET_CHUNK_DEFAULT = 512
+
+# Key under which the loss is reported. Shares the framework's MoE aux-loss
+# tracker, so it lands in the training log / TensorBoard / W&B next to the MoE
+# losses with no extra plumbing.
+INDEXER_DISTILL_LOSS_NAME = "indexer_distill_loss"
+
+
+def _moe_aux_loss_scale() -> Optional[torch.Tensor]:
+    """The scale the pipeline schedule installed for the MoE auxiliary loss.
+
+    ``forward_step_calc_loss`` sets it once per microbatch to
+    ``grad_scale * cp_size / num_microbatches`` (or just ``grad_scale`` under
+    ``calculate_per_token_loss``) whenever ``num_moe_experts`` is configured --
+    which is always true for V4. The indexer distillation loss needs exactly
+    that quantity, so it reads it rather than maintaining a second copy that
+    could silently drift. Returns ``None`` when the framework is not importable
+    (the torch-only unit tests).
+    """
+    try:
+        from megatron.core.transformer.moe.moe_utils import MoEAuxLossAutoScaler
+    except Exception:
+        return None
+    return getattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None)
+
+
+def log_indexer_distill_loss(
+    loss: Optional[torch.Tensor],
+    *,
+    layer_number: Optional[int],
+    num_layers: int,
+    device: torch.device,
+) -> None:
+    """Record this layer's indexer loss in the MoE aux-loss tracker.
+
+    Call this from **every** V4 attention layer whenever the loss is enabled,
+    passing ``None`` on the layers that do not own an indexer. The tracker is
+    reduced across pipeline ranks over whatever keys each rank happens to hold,
+    so a key that only appears on the ranks that own a CSA layer would make the
+    collective diverge. Non-CSA layers therefore contribute an explicit zero,
+    which keeps the key present everywhere and leaves the sum unchanged.
+
+    The reported value is the per-layer sum divided by the layer count (the
+    denominator ``track_moe_metrics`` applies to every tracked loss), not
+    divided by the number of CSA layers.
+    """
+    # ``layer_number`` is 1-based and indexes the tracker directly, so 0 (the
+    # "unnumbered" default on a standalone attention module) would silently
+    # write to the last slot.
+    if not layer_number:
+        return
+    try:
+        from megatron.core.transformer.moe.moe_utils import save_to_aux_losses_tracker
+    except Exception:
+        return
+
+    value = (
+        loss.detach().to(device=device, dtype=torch.float32)
+        if loss is not None
+        else torch.zeros((), device=device, dtype=torch.float32)
+    )
+    save_to_aux_losses_tracker(INDEXER_DISTILL_LOSS_NAME, value, layer_number, num_layers)
+
+
+class V4IndexerLossAutoScaler(torch.autograd.Function):
+    """Attach an auxiliary loss to an existing tensor's backward pass.
+
+    ``forward`` returns ``output`` unchanged; ``backward`` seeds ``aux_loss``
+    with the current aux-loss scale so its subgraph is differentiated as part
+    of the main backward. Same shape as the autoscalers already used for the MoE
+    auxiliary loss and the MTP loss.
+
+    The scale is not maintained here. Seeding a gradient of one would make the
+    effective coefficient ``num_microbatches`` times too large under gradient
+    accumulation and would ignore the grad scaler entirely, so by default this
+    follows :func:`_moe_aux_loss_scale` -- the same per-microbatch quantity the
+    schedule installs for the MoE auxiliary loss. :meth:`set_loss_scale`
+    overrides it for schedules that need to drive it explicitly.
+    """
+
+    # ``None`` means "follow the MoE aux-loss scale"; a tensor is an explicit
+    # override installed via ``set_loss_scale``.
+    main_loss_backward_scale: Optional[torch.Tensor] = None
+
+    @staticmethod
+    def current_loss_scale(reference: torch.Tensor) -> torch.Tensor:
+        """Resolve the scale to seed, on ``reference``'s device / dtype."""
+        scale = V4IndexerLossAutoScaler.main_loss_backward_scale
+        if scale is None:
+            scale = _moe_aux_loss_scale()
+        if scale is None:
+            return torch.ones((), device=reference.device, dtype=reference.dtype)
+        return scale.to(device=reference.device, dtype=reference.dtype)
+
+    @staticmethod
+    def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        ctx.save_for_backward(aux_loss)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):  # type: ignore[override]
+        (aux_loss,) = ctx.saved_tensors
+        scale = V4IndexerLossAutoScaler.current_loss_scale(aux_loss)
+        return grad_output, torch.ones_like(aux_loss) * scale
+
+    @staticmethod
+    def set_loss_scale(scale: Optional[torch.Tensor]) -> None:
+        """Override the gradient seeded into the auxiliary loss.
+
+        Pass ``None`` to go back to following the MoE auxiliary loss scale.
+        """
+        V4IndexerLossAutoScaler.main_loss_backward_scale = scale
+
+
+def _target_chunk_size(seq_len: int) -> int:
+    """Query rows per chunk when building the KL target.
+
+    The gathered pool is ``[B, S, K, head_dim]``, which at V4-Flash widths
+    (S=4096, K=512, head_dim=512) is 2.1 GB in BF16 for a single microbatch.
+    Materialising it whole costs more in HBM traffic than the GEMM that consumes
+    it, so slice the query axis. ``0`` disables chunking.
+    """
+    override = os.environ.get(_TARGET_CHUNK_ENV)
+    if override is not None:
+        try:
+            return max(int(override), 0)
+        except ValueError:
+            pass
+    return _TARGET_CHUNK_DEFAULT if seq_len > _TARGET_CHUNK_DEFAULT else 0
+
+
+def _target_distribution(
+    *,
+    query: torch.Tensor,
+    pool: torch.Tensor,
+    topk_idxs: torch.Tensor,
+    valid: torch.Tensor,
+    row_valid: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Head-summed attention distribution over the selected entries: ``[B, S, K]``.
+
+    This is the KL target, and it is fully detached, so nothing here needs to be
+    kept for backward. Three things follow from that and keep it cheap:
+
+    * The gather and the GEMM run at the **model dtype** instead of being
+      promoted to fp32 first. Promoting the gathered pool costs
+      ``head_dim`` times more traffic than promoting the result, and the logits
+      the main attention actually produces are computed at the model dtype too,
+      so staying there tracks the distribution being imitated more closely, not
+      less.
+    * The head sum happens **inside** the loop, so the ``[B, H, S, K]`` logits
+      never exist in full either -- only ``[B, H, chunk, K]``.
+    * The query axis is chunked, so neither does the ``[B, S, K, head_dim]``
+      gather, which at V4-Flash widths would be 2.1 GB per microbatch.
+    """
+    B, H, S, _ = query.shape
+    K = topk_idxs.shape[-1]
+    q = query.detach()
+    kv = pool.detach()
+    batch_idx = torch.arange(B, device=pool.device).view(B, 1, 1)
+
+    def block(sl: slice) -> torch.Tensor:
+        idx = topk_idxs[:, sl]
+        # Clamp the -1 sentinels to a legal row; masked back out right after.
+        gathered = kv[batch_idx, idx.clamp_min(0)]
+        logits = torch.einsum("bhsd,bskd->bhsk", q[:, :, sl], gathered).float()
+        logits *= softmax_scale
+        logits.masked_fill_(~valid[:, sl].unsqueeze(1), float("-inf"))
+        # An all -inf row would softmax to NaN; flatten it and drop it after.
+        rows = row_valid[:, sl].view(B, 1, -1, 1)
+        logits.masked_fill_(~rows, 0.0)
+        probs = torch.softmax(logits, dim=-1)
+        probs.mul_(rows)
+        return probs.sum(dim=1)  # [B, chunk, K]
+
+    chunk = _target_chunk_size(S)
+    if chunk <= 0 or chunk >= S:
+        return block(slice(0, S))
+
+    out = torch.empty((B, S, K), device=query.device, dtype=torch.float32)
+    for start in range(0, S, chunk):
+        stop = min(start + chunk, S)
+        out[:, start:stop] = block(slice(start, stop))
+    return out
+
+
+def compute_indexer_distill_loss(
+    *,
+    index_topk_scores: torch.Tensor,
+    topk_idxs: torch.Tensor,
+    query: torch.Tensor,
+    pool: torch.Tensor,
+    softmax_scale: float,
+    loss_coeff: float,
+    head_reduce_group: Optional["torch.distributed.ProcessGroup"] = None,
+) -> torch.Tensor:
+    """``KL(attention || indexer)`` over the selected compressed entries.
+
+    Distillation is **one-directional**: only the indexer learns. ``query`` and
+    ``pool`` are detached here so the KL cannot pull the main attention's Q
+    projection or the compressor towards whatever the indexer currently
+    predicts -- the open-source reference does the same, detaching both the
+    query and the key on the way into its indexer loss and feeding the indexer
+    itself from a detached hidden state.
+
+    Args:
+        index_topk_scores: indexer scores at the selected slots, ``[B, S, K]``.
+            The one tensor that stays attached -- it is the learning signal.
+            Invalid slots are ``-inf`` (fewer than K legal entries exist for
+            early queries).
+        topk_idxs: selected pool indices, ``[B, S, K]``; ``-1`` marks invalid.
+        query: post-RoPE queries, ``[B, H, S, head_dim]``. Detached internally.
+        pool: compressed KV pool, ``[B, P, head_dim]``. Detached internally.
+        softmax_scale: the attention softmax temperature (``1/sqrt(head_dim)``).
+        loss_coeff: scaling applied to the KL.
+        head_reduce_group: process group the attention heads are sharded over,
+            or ``None`` when every rank holds all of them. Only needed if the Q
+            projection stops gathering its output -- see
+            :meth:`DeepseekV4Attention._indexer_loss_head_group`.
+
+    Returns:
+        Scalar loss (fp32). Rows with no legal entry contribute nothing, so a
+        fully-masked row can never produce NaN.
+    """
+    B, _, S, _ = query.shape
+    valid = topk_idxs >= 0  # [B, S, K]
+    row_valid = valid.any(dim=-1)  # [B, S]
+
+    # The KL target: the head-summed attention distribution over the same
+    # entries. Fully detached (see the note above), so it is built under
+    # ``no_grad`` and never has to materialise the dense per-head tensors.
+    with torch.no_grad():
+        target = _target_distribution(
+            query=query,
+            pool=pool,
+            topk_idxs=topk_idxs,
+            valid=valid,
+            row_valid=row_valid,
+            softmax_scale=softmax_scale,
+        )
+        # The head sum must span all heads before the renormalisation below.
+        if head_reduce_group is not None and torch.distributed.get_world_size(head_reduce_group) > 1:
+            target = target.contiguous()
+            torch.distributed.all_reduce(target, group=head_reduce_group)
+        # Renormalise to a distribution; L1 is enough since softmax outputs are
+        # already non-negative.
+        target /= target.sum(dim=-1, keepdim=True).clamp(min=_EPS)
+
+    # The indexer side is the one that learns, so it stays attached. Rows with
+    # zero legal entries are neutralised before the softmax and dropped after,
+    # so a fully-masked row yields 0 rather than NaN and no gradient.
+    idx_row_mask = row_valid.view(B, S, 1)
+    idx_logits = index_topk_scores.float().masked_fill(~idx_row_mask, 0.0)
+    idx_probs = torch.softmax(idx_logits, dim=-1, dtype=torch.float32) * idx_row_mask.to(torch.float32)
+
+    kl_per_row = (target * (torch.log(target + _EPS) - torch.log(idx_probs + _EPS))).sum(dim=-1)
+    return kl_per_row.mean() * loss_coeff
diff --git a/primus/backends/megatron/core/transformer/keep_in_fp32.py b/primus/backends/megatron/core/transformer/keep_in_fp32.py
new file mode 100644
index 000000000..98a1ce135
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/keep_in_fp32.py
@@ -0,0 +1,122 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Keep selected parameters in FP32 across the framework's half-precision cast.
+
+``Float16Module`` wraps the model with a blanket ``module.bfloat16()`` (or
+``.half()``), which walks every floating-point parameter. Two DeepSeek-V4
+parameters are FP32 in the released checkpoint and both feed a softmax directly:
+the compressor's ``ape`` (an additive bias on the pooling scores) and the
+attention's ``attn_sink`` (an extra logit in the denominator). This module can
+hold them at FP32 through the cast: mark the parameter, mix
+:class:`KeepInFp32Mixin` into its owner, and the marked entries are restored to
+FP32 after any ``_apply`` (which is what ``.bfloat16()``, ``.half()``, ``.to()``
+and ``.cuda()`` all funnel through).
+
+**Off by default**, because holding a parameter at a dtype the rest of the model
+does not use is not free downstream. ``ParamAndGradBuffer`` does key its buffers
+on ``(param_dtype, grad_dtype)`` and allocate one per distinct combination, but
+the distributed optimizer on top of it does not follow: ``distrib_optimizer.py``
+carries five ``assert len(gbuf_range_maps) == 1, "single dtype supported, for
+now."`` guards. Enabling this on a 4-node / PP4 / EP8 run with
+``use_precision_aware_optimizer`` + ``store_param_remainders`` aborted every GPU
+with ``Memory access fault`` on the first training step, while the same build
+with the mechanism off completed cleanly. Single-node PP1 does not reproduce it,
+so this cannot be validated by the unit tests or a one-node soak.
+
+Leaving it off costs very little. ``store_param_remainders`` keeps FP32 master
+params in the optimizer either way, so update precision is unchanged; every
+consumer of these two parameters already promotes at the use site
+(``sink.float()`` in the sparse-MLA adapter and the eager reference,
+``score.float()`` before the pooling softmax), so forward precision is unchanged
+too. What the mark actually buys is the stored resolution matching the released
+checkpoint -- relevant to checkpoint parity, not to training correctness.
+
+Set ``PRIMUS_V4_KEEP_FP32=1`` to turn it on, but only with PP1 or after
+confirming the optimizer path tolerates two parameter dtypes.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+__all__ = [
+    "ENABLE_ENV_VAR",
+    "KeepInFp32Mixin",
+    "is_enabled",
+    "is_marked_keep_in_fp32",
+    "mark_keep_in_fp32",
+    "unmark_keep_in_fp32",
+]
+
+ENABLE_ENV_VAR = "PRIMUS_V4_KEEP_FP32"
+
+_MARK = "_primus_keep_in_fp32"
+
+
+def is_enabled() -> bool:
+    """Whether the keep-in-FP32 contract is active. Off unless opted in.
+
+    ``PRIMUS_V4_KEEP_FP32=1`` enables it; see the module docstring for why the
+    default is off and what the two parameters lose by following the model dtype
+    (stored resolution only -- not update or forward precision).
+    """
+    return os.environ.get(ENABLE_ENV_VAR, "0") == "1"
+
+
+def mark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor:
+    """Mark ``tensor`` so its owning :class:`KeepInFp32Mixin` keeps it FP32."""
+    if not is_enabled():
+        return tensor
+    setattr(tensor, _MARK, True)
+    return tensor
+
+
+def unmark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor:
+    """Drop the mark, letting ``tensor`` follow the model dtype again."""
+    if hasattr(tensor, _MARK):
+        delattr(tensor, _MARK)
+    return tensor
+
+
+def is_marked_keep_in_fp32(tensor: torch.Tensor) -> bool:
+    """Whether ``tensor`` is pinned to FP32."""
+    return bool(getattr(tensor, _MARK, False))
+
+
+class KeepInFp32Mixin:
+    """Restore marked parameters to FP32 after any ``_apply``.
+
+    Must come before ``nn.Module`` in the MRO so ``super()._apply`` reaches the
+    normal implementation. ``_apply`` runs once per module per conversion (and
+    the tensors involved here are tiny), so the save/restore cost is noise.
+    """
+
+    def _apply(self, fn, *args, **kwargs):  # type: ignore[no-untyped-def]
+        pinned = {
+            name: param.detach().to(torch.float32).clone()
+            for name, param in self._parameters.items()  # type: ignore[attr-defined]
+            if param is not None and is_marked_keep_in_fp32(param)
+        }
+
+        module = super()._apply(fn, *args, **kwargs)  # type: ignore[misc]
+
+        for name, original in pinned.items():
+            param = module._parameters.get(name)
+            if param is None:
+                continue
+            if param.dtype != torch.float32:
+                # ``fn`` may have produced a fresh Parameter, so restore from
+                # the saved FP32 copy rather than casting the downgraded values
+                # back (which would keep the BF16 rounding).
+                param.data = original.to(device=param.device)
+            # ``_apply`` can replace the Parameter object outright, dropping
+            # custom attributes; re-mark so the next conversion is protected.
+            mark_keep_in_fp32(param)
+
+        return module
diff --git a/primus/backends/megatron/core/transformer/local_rmsnorm.py b/primus/backends/megatron/core/transformer/local_rmsnorm.py
new file mode 100644
index 000000000..6efa92dad
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/local_rmsnorm.py
@@ -0,0 +1,103 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tiny RMSNorm fallback used by DeepSeek-V4 modules.
+
+Plan-2 P17 introduces this shared helper to retire three nearly
+identical ``_RMSNorm`` implementations that lived in:
+
+* ``primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py``
+* ``primus/backends/megatron/core/transformer/deepseek_v4_attention.py``
+  (a closure-built fallback for the attention's ``q_norm`` / ``kv_norm``
+  no-spec path)
+* ``primus/backends/megatron/core/transformer/compressor.py``
+
+All three computed the same RMSNorm with a learnable per-channel
+``weight`` and worked on the last dim. The goal of this module is to
+expose **one** implementation so:
+
+* dead-code audits stay clean,
+* state-dict round-trips treat the four call sites uniformly,
+* CPU-only unit tests can build V4 norms without dragging in
+  TransformerEngine / Megatron's TE-backed norm kernel.
+
+The spec-driven path (``DeepSeekV4SpecProvider.v4_norm_module()``) is
+unchanged — it returns Megatron's TE-backed RMSNorm or the local
+fallback, depending on the active runtime mode. This file is the local
+fallback only.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import (
+    fused_rms_norm,
+)
+
+
+class LocalRMSNorm(nn.Module):
+    """Single canonical RMSNorm fallback used across V4 modules.
+
+    Args:
+        dim: hidden dimension to normalize over (last axis).
+        eps: numerical stability epsilon.
+        hidden_size: alias for ``dim``; if both are provided ``dim``
+            wins. Provided for compatibility with Megatron-style norm
+            constructors that expect ``hidden_size=`` instead.
+        config: ignored (consumed for compatibility with Megatron's norm
+            factory signature).
+
+    Notes:
+        * The internal compute happens in fp32 to match the V4 reference
+          (``inference/model.py`` uses fp32 RMS even when activations
+          are bf16); the result is cast back to the input dtype.
+        * The ``weight`` parameter has shape ``(dim,)`` and is initialized
+          to ones so a freshly built norm is the identity transform —
+          this matches the V4 / HF reference checkpoint layout exactly,
+          so state-dict round-trips work without remapping.
+    """
+
+    def __init__(
+        self,
+        dim: Optional[int] = None,
+        eps: float = 1e-6,
+        *,
+        hidden_size: Optional[int] = None,
+        config: Optional[Any] = None,
+    ) -> None:
+        del config
+        super().__init__()
+        if dim is None:
+            dim = hidden_size
+        if dim is None:
+            raise ValueError("LocalRMSNorm requires `dim` (or the alias `hidden_size`).")
+        self.weight = nn.Parameter(torch.ones(int(dim)))
+        self.eps = float(eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Small-kernel-fusion (2026-07-03): collapse the eager RMS chain
+        # (bf16->fp32 cast + pow + mean + rsqrt + mul + fp32->bf16 cast +
+        # weight mul) into one Triton FWD + one BWD kernel. ``mid_cast=True``
+        # replicates the eager ``(x32*rstd).to(in_dtype)`` rounding BEFORE the
+        # weight multiply so numerics match bit-for-bit (within fp32 accum).
+        # ``out_dtype`` = the eager output dtype ``promote(in_dtype, weight)``.
+        # Gated by ``PRIMUS_RMSNORM_TRITON`` (default on); falls back to the
+        # eager body on CPU / when off.
+        out_dtype = torch.promote_types(x.dtype, self.weight.dtype)
+        return fused_rms_norm(
+            x,
+            self.weight,
+            eps=self.eps,
+            mid_cast=True,
+            out_dtype=out_dtype,
+        )
+
+
+__all__ = ["LocalRMSNorm"]
diff --git a/primus/modules/__init__.py b/primus/backends/megatron/core/transformer/moe/_triton/__init__.py
similarity index 100%
rename from primus/modules/__init__.py
rename to primus/backends/megatron/core/transformer/moe/_triton/__init__.py
diff --git a/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py
new file mode 100644
index 000000000..82a40c2fa
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py
@@ -0,0 +1,548 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused V4 router post-logits chain (plan-6 P39).
+
+Fuses the chain shared between
+:class:`primus.backends.megatron.core.transformer.moe.v4_topk_router.DeepseekV4LearnedRouter`
+and
+:class:`primus.backends.megatron.core.transformer.moe.v4_hash_router.DeepseekV4HashRouter`:
+
+.. code-block:: python
+
+    scores = score_fn(logits)            # softmax / sigmoid / sqrtsoftplus
+    weights = scores.gather(1, indices)  # [N, K]
+    if score_fn != softmax:
+        denom = weights.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        weights = weights / denom
+    weights *= topk_scaling_factor
+
+    probs = zeros(N, E); probs.scatter_(1, indices, weights)
+    routing_map = zeros(N, E); routing_map.scatter_(1, indices, True)
+
+`indices` is provided by the host (hash router pre-computes via
+``tid2eid[flat_ids]``; learned router pre-computes via
+``torch.topk(sel_score, K)``).  Keeping `topk` on the host lets us
+reuse the well-tuned cuBLAS / hip topk kernel and avoids stuffing
+two unrelated kernels into one register file.
+
+Gating: ``PRIMUS_V4_ROUTER_TRITON == "1"`` (**default-off**, P38 precedent).
+
+Microbench at V4-Flash widths (N=4096, E=256, K=8) shows the kernel
+wins on the V4 production score function:
+- ``sqrtsoftplus`` (V4 default): 1.56x FWD / 1.22x BWD
+- ``softmax`` (non-V4):           1.00x FWD / 0.73x BWD
+- ``sigmoid``:                    near-parity
+
+But the EP=8 proxy A/B (PRIMUS_V4_ROUTER_TRITON=1 vs =0, 10 iters
+each) shows the microbench gain does **not** surface end-to-end:
+~534 ms / iter both ways, lm_loss bit-identical iter-by-iter
+(parity confirmed).  The per-call savings (~0.06 ms FWD+BWD ×
+~16 router calls / iter ≈ 1 ms) are submerged in the ~500 ms /
+step variance of the EP=8 dispatch + grouped-MLP pipeline.
+
+P38-style descope: ship the kernel behind a knob, default OFF.
+Available for future tuning (e.g. when the broader graph compresses
+and exposes the per-call savings) and for small-shape paths where
+the FWD win is closer to 1.5-2x.  Bit-identity makes flipping the
+knob a safe operation.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# Score function enum (matches FWD/BWD constexpr).
+_SCORE_FN_SOFTMAX = 0
+_SCORE_FN_SIGMOID = 1
+_SCORE_FN_SQRTSOFTPLUS = 2
+_SCORE_FN_MAP = {
+    "softmax": _SCORE_FN_SOFTMAX,
+    "sigmoid": _SCORE_FN_SIGMOID,
+    "sqrtsoftplus": _SCORE_FN_SQRTSOFTPLUS,
+}
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_router_post_fwd_kernel(
+    LOGITS_PTR,  # [N, E] fp32
+    INDICES_PTR,  # [N, BLOCK_K] int64 (gather positions; padded to BLOCK_K)
+    PROBS_PTR,  # [N, E] OUT_DTYPE (must be zero-init'd by caller)
+    RMAP_PTR,  # [N, E] bool   (must be zero-init'd by caller)
+    SCORES_OUT_PTR,  # [N, E] fp32 (saved-for-backward; full row of scores)
+    WEIGHTS_OUT_PTR,  # [N, BLOCK_K] fp32 (saved-for-backward; gathered weights, post-denom, pre-scale)
+    DENOM_OUT_PTR,  # [N] fp32  (saved-for-backward; the clamped denom)
+    N,
+    E,  # runtime int (real expert count; NOT required to be a power of 2)
+    K,  # runtime int (real topk; NOT required to be a power of 2)
+    SCORE_FN: tl.constexpr,
+    SCALE: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_E: tl.constexpr,  # next_pow2(E) — column block over the expert axis
+    BLOCK_K: tl.constexpr,  # next_pow2(K) — column block over the topk axis
+    OUT_DTYPE: tl.constexpr,
+):
+    """One program tile = ``BLOCK_N`` rows of the post-logits chain.
+
+    Supports arbitrary (non-power-of-2) ``E`` / ``K``: the expert axis is
+    tiled by ``BLOCK_E = next_pow2(E)`` and the topk axis by
+    ``BLOCK_K = next_pow2(K)``, with masks zeroing the padded columns.
+    ``INDICES_PTR`` / ``WEIGHTS_OUT_PTR`` are laid out with row stride
+    ``BLOCK_K`` (host pads the indices tensor to ``[N, BLOCK_K]``).
+    """
+
+    pid = tl.program_id(0)
+    n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    n_mask = n_offs < N
+
+    e_idx = tl.arange(0, BLOCK_E)
+    e_mask = e_idx < E
+    k_idx = tl.arange(0, BLOCK_K)
+    k_mask = k_idx < K
+
+    # Load logits [BLOCK_N, BLOCK_E] in fp32 (padded cols masked out).
+    logits = tl.load(
+        LOGITS_PTR + n_offs[:, None] * E + e_idx[None, :],
+        mask=n_mask[:, None] & e_mask[None, :],
+        other=0.0,
+    ).to(tl.float32)
+
+    # Apply score_fn (compile-time specialised). Padded expert columns must
+    # not leak into the softmax reduction, so mask them to -inf first.
+    if SCORE_FN == 0:  # softmax
+        neg_inf = float("-inf")
+        logits_m = tl.where(e_mask[None, :], logits, neg_inf)
+        m = tl.max(logits_m, axis=1, keep_dims=True)
+        e = tl.where(e_mask[None, :], tl.exp(logits_m - m), 0.0)
+        s = tl.sum(e, axis=1, keep_dims=True)
+        scores = e / s
+    elif SCORE_FN == 1:  # sigmoid
+        scores = tl.sigmoid(logits)
+    else:  # sqrtsoftplus (SCORE_FN == 2)
+        softplus = tl.log(1.0 + tl.exp(logits))
+        # Stable: for large x, log(1+exp(x)) ≈ x.
+        scores = tl.sqrt(softplus)
+
+    # Zero padded expert columns so they never contribute to gather / denom.
+    scores = tl.where(e_mask[None, :], scores, 0.0)
+
+    # Save full scores row for backward (real columns only).
+    tl.store(
+        SCORES_OUT_PTR + n_offs[:, None] * E + e_idx[None, :],
+        scores,
+        mask=n_mask[:, None] & e_mask[None, :],
+    )
+
+    # Load indices [BLOCK_N, BLOCK_K] (int64; padded cols read as 0, masked below).
+    indices = tl.load(
+        INDICES_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :],
+        mask=n_mask[:, None] & k_mask[None, :],
+        other=0,
+    )
+
+    # Gather weights from scores at the K indices per row entirely in
+    # registers (avoid store-then-reload-via-HBM hazard). Build
+    # weights[BLOCK_N, BLOCK_K] via a static loop over BLOCK_K, extracting
+    # each column from the already-loaded ``indices`` tile.
+    weights = tl.zeros((BLOCK_N, BLOCK_K), dtype=tl.float32)
+    for k_off in tl.static_range(BLOCK_K):
+        k_is_off = (k_idx == k_off).to(tl.float32)
+        idx_col = tl.sum(tl.where(k_idx[None, :] == k_off, indices, 0), axis=1)
+        # scores_at_idx[n] = scores[n, idx_col[n]]
+        e_is_idx = e_idx[None, :] == idx_col[:, None]
+        scores_at_idx = tl.sum(scores * e_is_idx.to(tl.float32), axis=1)
+        weights += scores_at_idx[:, None] * k_is_off[None, :]
+
+    # Zero padded topk columns (k >= K).
+    weights = tl.where(k_mask[None, :], weights, 0.0)
+
+    # If non-softmax: denom + normalize.
+    if SCORE_FN != 0:
+        s = tl.sum(weights, axis=1, keep_dims=True)
+        denom = tl.maximum(s, EPS)
+        weights = weights / denom
+        denom_scalar = tl.reshape(denom, (BLOCK_N,))
+    else:
+        denom_scalar = tl.full((BLOCK_N,), 1.0, dtype=tl.float32)
+    tl.store(DENOM_OUT_PTR + n_offs, denom_scalar, mask=n_mask)
+
+    # Scale.
+    weights = weights * SCALE
+
+    # Save weights for backward.
+    tl.store(
+        WEIGHTS_OUT_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :],
+        weights,
+        mask=n_mask[:, None] & k_mask[None, :],
+    )
+
+    # Sparse scatter (probs[n, indices] = weights, rmap[n, indices] = True).
+    # No atomics needed because each (n, e) is written at most once
+    # per row (caller guarantees indices are unique within a row).  Padded
+    # topk columns are masked out so they never write to expert 0.
+    tl.store(
+        PROBS_PTR + n_offs[:, None] * E + indices,
+        weights.to(OUT_DTYPE),
+        mask=n_mask[:, None] & k_mask[None, :],
+    )
+    tl.store(
+        RMAP_PTR + n_offs[:, None] * E + indices,
+        tl.full((BLOCK_N, BLOCK_K), 1, dtype=tl.int1),
+        mask=n_mask[:, None] & k_mask[None, :],
+    )
+
+
+@triton.jit
+def _v4_router_post_bwd_kernel(
+    DPROBS_PTR,  # [N, E] OUT_DTYPE upstream grad
+    INDICES_PTR,  # [N, BLOCK_K] int64 (padded)
+    SCORES_PTR,  # [N, E] fp32 saved
+    WEIGHTS_PTR,  # [N, BLOCK_K] fp32 saved (post-scaled, padded)
+    DENOM_PTR,  # [N]   fp32 saved
+    DLOGITS_PTR,  # [N, E] fp32 OUT
+    N,
+    E,  # runtime int
+    K,  # runtime int
+    SCORE_FN: tl.constexpr,
+    SCALE: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_E: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """VJP through the chain.
+
+    FWD:
+        scores[n, e]  = score_fn(logits[n, e])
+        weights[n, k] = scores[n, indices[n, k]]
+        if not softmax: weights /= sum_k weights[n, k]
+        weights *= scale
+        probs[n, indices[n, k]] = weights
+
+    BWD (gather only the K positions touched by indices; others get 0):
+        dprobs_at[n, k] = dprobs[n, indices[n, k]]
+        dweights[n, k]  = dprobs_at[n, k] * SCALE   (chain through scale)
+        If not softmax:
+            saved weights = (gathered / denom) * SCALE
+            d_pre_denom[n, k] = dweights[n, k] / denom
+            d_denom[n] = -sum_k (dweights[n, k] * weights_scaled[n, k]) / (denom * scale)
+                (because weights_scaled = gathered * scale / denom, so
+                 ∂w/∂gathered = scale/denom, ∂w/∂denom = -gathered * scale / denom^2)
+            d_gathered[n, k] = dweights[n, k] * scale / denom + d_denom[n] * 1
+                Wait, denom = sum_k gathered, so ∂denom/∂gathered_k = 1.
+                Therefore d_gathered[n, k] = dweights * scale/denom + d_denom (each gathered contributes to denom)
+        else (softmax):
+            d_gathered[n, k] = dweights[n, k]
+        Scatter d_gathered back into d_scores at positions indices[n, k]
+        (rest of d_scores is 0).
+        Finally chain through score_fn:
+            d_logits[n, e] = vjp_score_fn(d_scores[n, e], scores[n, e])
+    """
+
+    pid = tl.program_id(0)
+    n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    n_mask = n_offs < N
+
+    e_idx = tl.arange(0, BLOCK_E)
+    e_mask = e_idx < E
+    k_idx = tl.arange(0, BLOCK_K)
+    k_mask = k_idx < K
+
+    # Load saved state.
+    scores = tl.load(
+        SCORES_PTR + n_offs[:, None] * E + e_idx[None, :],
+        mask=n_mask[:, None] & e_mask[None, :],
+        other=0.0,
+    )
+    indices = tl.load(
+        INDICES_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :],
+        mask=n_mask[:, None] & k_mask[None, :],
+        other=0,
+    )
+
+    # Gather upstream grad at the K indices per row.
+    # dprobs_at = dprobs[n, indices[n, k]]   shape [BLOCK_N, BLOCK_K], fp32.
+    # Padded topk columns are masked to 0 so they carry no gradient.
+    dprobs_at = tl.load(
+        DPROBS_PTR + n_offs[:, None] * E + indices,
+        mask=n_mask[:, None] & k_mask[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    # Chain through the SCALE multiply: forward was
+    #     weights_post = weights_pre * SCALE
+    # so d_weights_pre = d_weights_post * SCALE.
+    dweights_pre_scale = dprobs_at * SCALE
+
+    if SCORE_FN != 0:
+        # Non-softmax: forward applied
+        #     weights_pre = gathered / denom
+        # so d_gathered_k = d_weights_pre_k / denom + d_denom
+        # where d_denom = -sum_k d_weights_pre_k * gathered_k / denom^2.
+        # Recover gathered_k from saved state:
+        #     weights_saved_k = (gathered_k / denom) * SCALE
+        # so gathered_k = weights_saved_k * denom / SCALE.
+        denom = tl.load(DENOM_PTR + n_offs, mask=n_mask, other=1.0)
+        weights_saved = tl.load(
+            WEIGHTS_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :],
+            mask=n_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+        gathered_k = weights_saved * (denom[:, None] / SCALE)
+        d_denom_per_n = -tl.sum(dweights_pre_scale * gathered_k, axis=1) / (denom * denom)
+        d_gathered = dweights_pre_scale / denom[:, None] + d_denom_per_n[:, None]
+    else:
+        # Softmax: weights_pre is just gathered directly.
+        d_gathered = dweights_pre_scale
+
+    # Zero padded topk columns so they contribute no gradient.
+    d_gathered = tl.where(k_mask[None, :], d_gathered, 0.0)
+
+    # Build d_scores_full entirely in registers using an explicit static
+    # loop over BLOCK_K.  For each k, we add d_gathered[:, k] at the position
+    # indices[:, k] in the E-axis using a broadcast compare.  This avoids
+    # the round-trip-via-HBM hazard of a scatter-then-load pattern.  Padded
+    # columns carry d_gathered == 0, so they add nothing even though their
+    # index reads as 0.
+    d_scores_full = tl.zeros((BLOCK_N, BLOCK_E), dtype=tl.float32)
+    for k_off in tl.static_range(BLOCK_K):
+        idx_col = tl.sum(tl.where(k_idx[None, :] == k_off, indices, 0), axis=1)
+        # Slice d_gathered[:, k_off] -> shape [BLOCK_N]
+        grad_k = tl.sum(d_gathered * (k_idx[None, :] == k_off).to(tl.float32), axis=1)
+        # Contribution: grad_k[:, None] where e_idx == idx_col[:, None]
+        e_is_idx = e_idx[None, :] == idx_col[:, None]
+        d_scores_full += tl.where(e_is_idx, grad_k[:, None], 0.0)
+
+    if SCORE_FN == 0:  # softmax: d_logits = scores * (d_scores - sum(d_scores * scores))
+        dot = tl.sum(d_scores_full * scores, axis=1, keep_dims=True)
+        d_logits = scores * (d_scores_full - dot)
+    elif SCORE_FN == 1:  # sigmoid
+        d_logits = d_scores_full * scores * (1.0 - scores)
+    else:  # sqrtsoftplus: y = sqrt(softplus(x)); dy/dx = (1 / (2 * sqrt(softplus(x)))) * sigmoid(x)
+        # = sigmoid(x) / (2 * y); avoid div-by-zero by guarding small y.
+        # sigmoid(x) = exp(x) / (1 + exp(x)); when y is close to 0, x is very negative,
+        # and sigmoid(x) is also very small, so the limit is 0.
+        # Use: dy/dx = sigmoid(x) / (2 * scores)
+        # but we don't have x.  Recompute: scores = y = sqrt(softplus(x))
+        # so softplus(x) = y^2 = scores^2, and exp(x) = exp(scores^2) - 1? No:
+        # softplus(x) = log(1 + exp(x)) = y^2 -> exp(x) = exp(y^2) - 1 -> sigmoid(x) = (exp(y^2) - 1) / exp(y^2)
+        # = 1 - exp(-y^2).
+        sig_x = 1.0 - tl.exp(-scores * scores)
+        y_safe = tl.maximum(scores, EPS)
+        d_logits = d_scores_full * sig_x / (2.0 * y_safe)
+
+    tl.store(
+        DLOGITS_PTR + n_offs[:, None] * E + e_idx[None, :],
+        d_logits,
+        mask=n_mask[:, None] & e_mask[None, :],
+    )
+
+
+# ---------------------------------------------------------------------------
+# autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class V4RouterPostFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the FWD/BWD Triton kernels.
+
+    Inputs:
+        logits [N, E] fp32  -- gate output (pre-score-fn).
+        indices [N, K] long -- gather positions (from host-side topk
+                               or hash table).  Must be unique per row.
+        score_function: "softmax" | "sigmoid" | "sqrtsoftplus".
+        topk_scaling_factor: float multiplier applied after denom.
+        out_dtype: dtype of the returned `probs` tensor.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        logits: torch.Tensor,
+        indices: torch.Tensor,
+        score_function: str,
+        topk_scaling_factor: float,
+        out_dtype: torch.dtype,
+    ):
+        if logits.dim() != 2:
+            raise ValueError(f"logits must be [N, E], got shape {tuple(logits.shape)}")
+        if indices.dim() != 2:
+            raise ValueError(f"indices must be [N, K], got shape {tuple(indices.shape)}")
+        if logits.shape[0] != indices.shape[0]:
+            raise ValueError(f"logits N={logits.shape[0]} != indices N={indices.shape[0]}")
+        if score_function not in _SCORE_FN_MAP:
+            raise ValueError(
+                f"Unknown score_function: {score_function!r}; expected one of "
+                f"{sorted(_SCORE_FN_MAP.keys())}"
+            )
+        N, E = logits.shape
+        _, K = indices.shape
+        # Arbitrary (non-power-of-2) E / K are supported: the kernel tiles
+        # both axes by their next power of 2 and masks the padded columns.
+        block_e = triton.next_power_of_2(E)
+        block_k = triton.next_power_of_2(K)
+
+        score_fn_enum = _SCORE_FN_MAP[score_function]
+        logits_c = logits.contiguous().to(torch.float32)
+        indices_c = indices.contiguous().to(torch.int64)
+        # Pad the indices tensor to [N, BLOCK_K] so the kernel's BLOCK_K-wide
+        # loads stay in-bounds; padded columns (value 0) are masked out
+        # everywhere via k_mask so they never gather / scatter / grad.
+        if block_k != K:
+            indices_pad = torch.zeros((N, block_k), dtype=torch.int64, device=indices_c.device)
+            indices_pad[:, :K] = indices_c
+        else:
+            indices_pad = indices_c
+
+        device = logits_c.device
+        probs = torch.zeros((N, E), dtype=out_dtype, device=device)
+        routing_map = torch.zeros((N, E), dtype=torch.bool, device=device)
+        scores_saved = torch.empty((N, E), dtype=torch.float32, device=device)
+        weights_saved = torch.zeros((N, block_k), dtype=torch.float32, device=device)
+        denom_saved = torch.empty((N,), dtype=torch.float32, device=device)
+
+        # BLOCK_N heuristic: per-row state ~ BLOCK_E + BLOCK_K + few scalars.
+        if block_e <= 64:
+            block_n = 64
+        elif block_e <= 256:
+            block_n = 16
+        else:
+            block_n = 4
+        grid = (triton.cdiv(N, block_n),)
+
+        _v4_router_post_fwd_kernel[grid](
+            logits_c,
+            indices_pad,
+            probs,
+            routing_map,
+            scores_saved,
+            weights_saved,
+            denom_saved,
+            N,
+            E,
+            K,
+            SCORE_FN=score_fn_enum,
+            SCALE=float(topk_scaling_factor),
+            EPS=1e-12,
+            BLOCK_N=block_n,
+            BLOCK_E=block_e,
+            BLOCK_K=block_k,
+            OUT_DTYPE={
+                torch.float32: tl.float32,
+                torch.float16: tl.float16,
+                torch.bfloat16: tl.bfloat16,
+                torch.float64: tl.float64,
+            }[out_dtype],
+        )
+
+        ctx.save_for_backward(indices_pad, scores_saved, weights_saved, denom_saved)
+        ctx.score_fn_enum = score_fn_enum
+        ctx.scale = float(topk_scaling_factor)
+        ctx.E = E
+        ctx.K = K
+        ctx.block_e = block_e
+        ctx.block_k = block_k
+        ctx.block_n = block_n
+        ctx.in_dtype = logits.dtype
+        return probs, routing_map
+
+    @staticmethod
+    def backward(ctx, d_probs, d_routing_map):  # type: ignore[override]
+        indices_pad, scores_saved, weights_saved, denom_saved = ctx.saved_tensors
+        E = ctx.E
+        K = ctx.K
+        block_e = ctx.block_e
+        block_k = ctx.block_k
+        block_n = ctx.block_n
+        score_fn_enum = ctx.score_fn_enum
+        scale = ctx.scale
+        in_dtype = ctx.in_dtype
+
+        N = indices_pad.shape[0]
+        device = indices_pad.device
+
+        d_probs = d_probs.contiguous()
+        # d_routing_map ignored (bool tensor, no grad flow).
+        # Initialise d_logits to 0 so the scatter writes only at K positions per row.
+        d_logits_fp32 = torch.zeros((N, E), dtype=torch.float32, device=device)
+
+        grid = (triton.cdiv(N, block_n),)
+        _v4_router_post_bwd_kernel[grid](
+            d_probs,
+            indices_pad,
+            scores_saved,
+            weights_saved,
+            denom_saved,
+            d_logits_fp32,
+            N,
+            E,
+            K,
+            SCORE_FN=score_fn_enum,
+            SCALE=scale,
+            EPS=1e-12,
+            BLOCK_N=block_n,
+            BLOCK_E=block_e,
+            BLOCK_K=block_k,
+        )
+
+        return d_logits_fp32.to(in_dtype), None, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff ``PRIMUS_V4_ROUTER_TRITON != "0"`` (default ``"1"``).
+
+    Plan-8 P57 close-out 2 (2026-05-15): default flipped from ``"0"``
+    to ``"1"``.  Microbench at V4-Flash widths is a clear positive on
+    V4's production `sqrtsoftplus` score function (1.56x FWD / 1.22x
+    BWD); EP=8 proxy A/B (P39 / P43) sat inside the proxy noise band,
+    so we default the microbench-positive kernel ON to keep it on the
+    production path.  Set ``PRIMUS_V4_ROUTER_TRITON=0`` to revert to
+    the eager body.
+    """
+    return os.environ.get("PRIMUS_V4_ROUTER_TRITON", "1") != "0"
+
+
+def v4_router_post_triton(
+    logits: torch.Tensor,
+    indices: torch.Tensor,
+    *,
+    score_function: str,
+    topk_scaling_factor: float,
+    out_dtype: torch.dtype,
+):
+    """Run the fused V4 router post-logits kernel.
+
+    Returns ``(probs, routing_map)``.  Caller pre-computes ``indices``
+    (hash router via ``tid2eid[flat_ids]``; learned router via
+    ``torch.topk(sel_score, K).indices``).
+
+    Arbitrary (non-power-of-2) ``E`` / ``K`` are supported — the kernel
+    tiles both axes by the next power of 2 and masks the padded columns —
+    so the only hard requirement is that the tensors live on the GPU.
+    """
+    assert logits.is_cuda and indices.is_cuda, "v4_router_post_triton requires CUDA / HIP tensors"
+    return V4RouterPostFn.apply(logits, indices, score_function, topk_scaling_factor, out_dtype)
+
+
+__all__ = [
+    "V4RouterPostFn",
+    "v4_router_post_triton",
+    "is_triton_path_enabled",
+]
diff --git a/primus/backends/megatron/core/transformer/moe/router.py b/primus/backends/megatron/core/transformer/moe/router.py
index bec8a232b..a6b0d6bf7 100644
--- a/primus/backends/megatron/core/transformer/moe/router.py
+++ b/primus/backends/megatron/core/transformer/moe/router.py
@@ -115,4 +115,53 @@ def routing(self, logits: torch.Tensor, **kwargs):
         # so by the time we get here ``logits`` is already random when
         # ``args.moe_router_force_load_balancing`` is True. There is nothing
         # extra to do.
+        # ``moe_router_force_load_balancing_type`` selects the balancing mode
+        # (only relevant when ``moe_router_force_load_balancing`` is on):
+        #   * "even"    -> deterministic round-robin so per-expert counts are
+        #                  constant every step (constant M_total, no autotune/
+        #                  recompile churn); handled by ``_force_even_routing``.
+        #   * "uniform" -> Megatron-LM original random-logits balancing (already
+        #                  applied upstream in ``TopKRouter.forward``; nothing to
+        #                  do here).
+        force_load_balancing_type = getattr(args, "moe_router_force_load_balancing_type", "uniform")
+        if (
+            args.moe_router_force_load_balancing
+            and force_load_balancing_type == "even"
+            and not getattr(args, "moe_enable_deepep", False)
+        ):
+            scores, routing_map = self._force_even_routing(scores, routing_map)
+
         return scores, routing_map
+
+    def _force_even_routing(
+        self, scores: torch.Tensor, routing_map: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Deterministically assign each token's top-k slots to experts in a
+        round-robin cycle ``(token_idx * topk + k) % num_experts`` so per-expert
+        token counts are exactly balanced and step-invariant (constant M_total).
+
+        The original top-k probability magnitudes are carried over to the new
+        expert positions (per-slot), keeping ``scores`` and ``routing_map``
+        mutually consistent so the MoE combine weights stay correct.
+        """
+        num_tokens = routing_map.shape[0]
+        num_experts = self.config.num_moe_experts
+        topk = self.topk
+        device = routing_map.device
+
+        # slot[t, k] = deterministic destination expert for token t's k-th slot.
+        # topk consecutive experts per token -> distinct while topk <= num_experts.
+        slot = (
+            torch.arange(num_tokens * topk, device=device).view(num_tokens, topk) % num_experts
+        )  # [num_tokens, topk], int64
+
+        new_routing_map = torch.zeros_like(routing_map)
+        new_routing_map.scatter_(1, slot, torch.ones_like(slot, dtype=routing_map.dtype))
+
+        # Carry the real per-token top-k probability magnitudes to the new slots.
+        # scores is non-zero only on the real top-k, so topk() extracts exactly them.
+        topk_vals, _ = torch.topk(scores, topk, dim=1)  # [num_tokens, topk]
+        new_scores = torch.zeros_like(scores)
+        new_scores.scatter_(1, slot, topk_vals.to(scores.dtype))
+
+        return new_scores, new_routing_map
diff --git a/primus/backends/megatron/core/transformer/moe/shared_experts.py b/primus/backends/megatron/core/transformer/moe/shared_experts.py
new file mode 100644
index 000000000..c4eeadf83
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/shared_experts.py
@@ -0,0 +1,91 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus shared-expert MLP with fused (clamped) SwiGLU.
+
+Megatron's stock :class:`SharedExpertMLP` reaches the fused SwiGLU path via
+``mlp.MLP.forward`` only through ``bias_swiglu_impl``, which does **not**
+support DeepSeek-V4's pre-multiplication clamp. To keep the clamp correct,
+``v4_moe`` disables ``bias_activation_fusion`` for the shared expert, which
+forces the un-fused eager ``chunk``/``clamp``/``SiLU``/``mul`` path.
+
+:class:`PrimusSharedExpertMLP` overrides the activation to call Primus's fused
+clamped SwiGLU Triton kernel (:func:`swiglu_impl`), mirroring what
+``PrimusGroupedMLP`` does for the routed experts. Both the normal ``forward``
+and the ``--moe-shared-expert-overlap`` ``linear_fc1_forward_and_act`` paths
+are covered so the clamp semantics stay identical.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from megatron.core.transformer.moe.shared_experts import (
+    SharedExpertMLP,
+    set_tensor_grad_fn_sequence_sr,
+)
+from megatron.core.typed_torch import apply_module
+from megatron.core.utils import nvtx_range_pop, nvtx_range_push
+
+from primus.backends.megatron.core.fusions.fused_bias_swiglu import swiglu_impl
+
+
+class PrimusSharedExpertMLP(SharedExpertMLP):
+    """Shared-expert MLP that fuses the (clamped) SwiGLU activation."""
+
+    def _can_fuse_swiglu(self) -> bool:
+        return (
+            not self.config.use_te_activation_func
+            and self.config.gated_linear_unit
+            and self.activation_func == F.silu
+        )
+
+    def _fused_swiglu(self, intermediate_parallel, bias_parallel):
+        # dtype and the pre-mul clamp are handled inside the fused kernel.
+        return swiglu_impl(
+            intermediate_parallel,
+            bias_parallel,
+            self.config.activation_func_fp8_input_store,
+            self.config.activation_func_clamp_value,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Forward with fused clamped SwiGLU (non-overlap path)."""
+        if not self._can_fuse_swiglu():
+            return super().forward(hidden_states)
+
+        nvtx_range_push(suffix="linear_fc1")
+        intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)(hidden_states)
+        nvtx_range_pop(suffix="linear_fc1")
+
+        nvtx_range_push(suffix="activation")
+        intermediate_parallel = self._fused_swiglu(intermediate_parallel, bias_parallel)
+        nvtx_range_pop(suffix="activation")
+
+        nvtx_range_push(suffix="linear_fc2")
+        output, _ = apply_module(self.linear_fc2)(intermediate_parallel)
+        nvtx_range_pop(suffix="linear_fc2")
+
+        if self.use_shared_expert_gate:
+            logits = torch.nn.functional.linear(hidden_states, self.gate_weight)
+            gate_score = torch.nn.functional.sigmoid(logits)
+            output = output * gate_score
+        return output
+
+    def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
+        """Overlap-path FC1 + fused clamped SwiGLU activation."""
+        if not self._can_fuse_swiglu():
+            return super().linear_fc1_forward_and_act(overlapped_comm_output)
+
+        assert self.config.moe_shared_expert_overlap
+        assert self.cached_fc1_input is not None
+        if overlapped_comm_output is not None:
+            set_tensor_grad_fn_sequence_sr(overlapped_comm_output, torch.iinfo(torch.int).max)
+        with torch.cuda.stream(self.stream):
+            # [s, b, 4 * h/p]
+            intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)(self.cached_fc1_input)
+            self.cached_fc1_input = None
+            self.cached_fc2_input = self._fused_swiglu(intermediate_parallel, bias_parallel)
diff --git a/primus/backends/megatron/core/transformer/moe/v4_hash_router.py b/primus/backends/megatron/core/transformer/moe/v4_hash_router.py
new file mode 100644
index 000000000..eae8bca01
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/v4_hash_router.py
@@ -0,0 +1,270 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Hash router for DeepSeek-V4's first ``num_hash_layers`` MoE layers.
+
+Reference: techblog §4 ("MoE: hash routing for the first N layers") and
+the inference reference at ``DeepSeek-V4-Flash/inference/model.py:Gate``
+(the ``self.hash`` branch).
+
+For the first ``num_hash_layers`` of V4 (V4-Flash = 3), expert
+**selection** is static: each token id is permanently assigned to a
+fixed set of ``moe_router_topk`` experts via a deterministic
+``tid2eid`` table. Routing **weights**, however, are *not* uniform —
+they come from the same learned linear gate as the non-hash layers; we
+just gather the scores at the prescribed expert ids instead of running
+a top-K argmax.
+
+Released-checkpoint contract (per HF ``Gate.__init__``):
+
+* ``weight`` : ``nn.Parameter`` of shape ``[num_experts, hidden_size]``
+  — the learned gate. State-dict key matches the learned router so the
+  V4 state-dict adapter can map ``mlp.gate.weight`` uniformly.
+* ``tid2eid`` : ``nn.Parameter`` of shape ``[vocab_size, topk]``, dtype
+  ``int32``, ``requires_grad=False``. The released checkpoint stores
+  it as a parameter (not a buffer) so it is preserved across
+  ``state_dict`` round-trips.
+
+Forward semantics:
+
+    scores         = score_fn(linear(hidden, weight))   # fp32
+    indices        = tid2eid[token_ids]                 # static
+    weights        = scores.gather(1, indices)          # learned weights
+    if score_fn != softmax: weights /= weights.sum(-1)
+    weights       *= route_scale
+
+Plan-2 P14 contract:
+
+* :class:`DeepseekV4HashRouter` is a standalone ``nn.Module`` that
+  produces sparse ``(probs, routing_map)`` with the same ``[N,
+  num_experts]`` shape contract as :class:`DeepseekV4LearnedRouter`.
+* ``score_function`` ∈ ``{"softmax", "sigmoid", "sqrtsoftplus"}`` and
+  ``topk_scaling_factor`` are honored identically.
+
+Back-compat alias ``HashRouter`` is exposed but deprecated; new
+callers should use :class:`DeepseekV4HashRouter`.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (
+    is_triton_path_enabled as _v4_router_triton_enabled,
+)
+from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (
+    v4_router_post_triton,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    _VALID_SCORE_FUNCTIONS,
+    v4_score_fn,
+)
+
+
+def _build_default_tid2eid(*, vocab_size: int, num_experts: int, topk: int, seed: int) -> torch.Tensor:
+    """Build a deterministic ``[vocab_size, topk]`` int32 expert table.
+
+    Each token id gets ``topk`` distinct expert ids drawn uniformly
+    without replacement from ``[0, num_experts)``. The seed is fixed
+    across all ranks so PP / TP / EP shards see identical routing.
+
+    The table layout matches the HF reference: ``int32`` to keep on-disk
+    size small, indexed via long-cast at gather time.
+    """
+    # Build on CPU with a CPU generator so the table is bit-identical across
+    # ranks and devices, independent of the ambient ``torch.set_default_device``
+    # (the caller colocates the result with the module's weight).
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    rows = []
+    for _ in range(int(vocab_size)):
+        perm = torch.randperm(num_experts, generator=gen, device="cpu")[:topk]
+        rows.append(perm)
+    table = torch.stack(rows, dim=0).to(torch.int32)
+    return table
+
+
+class DeepseekV4HashRouter(nn.Module):
+    """Static hash-based MoE router with a *learned* score gate.
+
+    Args:
+        hidden_size: model dim ``D``; the gate is a single ``D ->
+            num_experts`` linear shared in shape with the learned router.
+        num_experts: total number of routed experts.
+        topk: number of experts each token is routed to.
+        vocab_size: tokenizer vocabulary size; controls the table length.
+        seed: deterministic seed for the hash table; same across all
+            ranks. Used only when ``tid2eid`` is built locally; if a
+            checkpoint provides ``tid2eid`` directly, the seed is
+            ignored at load time.
+        score_function: one of ``{"softmax", "sigmoid", "sqrtsoftplus"}``;
+            applied to the learned scores (matches the learned router).
+        topk_scaling_factor: scalar multiplier applied to the
+            renormalized routing weights (V3 ``moe_router_topk_scaling_factor``,
+            HF ``Gate.route_scale``).
+        dtype: dtype of the gate weight; defaults to fp32.
+
+    Parameters:
+        weight: ``[num_experts, hidden_size]`` learned gate.
+        tid2eid: ``[vocab_size, topk]`` int32 frozen mapping.
+            ``requires_grad=False``; this is a parameter (matching the
+            HF reference checkpoint) so that ``state_dict`` round-trips
+            preserve it.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_experts: int,
+        topk: int,
+        vocab_size: int,
+        seed: int = 0,
+        score_function: str = "sqrtsoftplus",
+        topk_scaling_factor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        super().__init__()
+        if num_experts <= 0:
+            raise ValueError(f"num_experts must be > 0, got {num_experts}")
+        if topk <= 0 or topk > num_experts:
+            raise ValueError(f"topk must be in [1, {num_experts}], got {topk}")
+        if vocab_size <= 0:
+            raise ValueError(f"vocab_size must be > 0, got {vocab_size}")
+        if score_function not in _VALID_SCORE_FUNCTIONS:
+            raise ValueError(
+                f"Unknown score_function: {score_function!r}. "
+                f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}."
+            )
+
+        self.hidden_size = int(hidden_size)
+        self.num_experts = int(num_experts)
+        self.topk = int(topk)
+        self.vocab_size = int(vocab_size)
+        self.seed = int(seed)
+        self.score_function = str(score_function)
+        self.topk_scaling_factor = float(topk_scaling_factor)
+
+        weight_dtype = dtype or torch.float32
+        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, dtype=weight_dtype))
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+
+        # tid2eid is a non-trainable parameter (matches HF reference layout
+        # so checkpoint round-trips include it). int32 to keep memory small.
+        tid2eid_init = _build_default_tid2eid(
+            vocab_size=self.vocab_size,
+            num_experts=self.num_experts,
+            topk=self.topk,
+            seed=self.seed,
+        )
+        # Colocate the CPU-built table with the module's weight (the ambient
+        # default device at construction, e.g. CUDA under set_default_device).
+        self.tid2eid = nn.Parameter(tid2eid_init.to(self.weight.device), requires_grad=False)
+
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        token_ids: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Route tokens whose expert ids are prescribed by ``tid2eid``.
+
+        Args:
+            hidden: ``[B, S, D]`` (or any shape with last dim ``D``).
+                The learned gate is evaluated against ``hidden``; the
+                resulting scores supply the *weights*.
+            token_ids: ``[B, S]`` (or any compatible shape) integer
+                tensor of token ids. Provides the *indices* via
+                ``tid2eid``.
+
+        Returns:
+            probs: ``[N, num_experts]`` float tensor, ``N = numel/D``.
+                Non-selected experts have probability 0; selected
+                experts hold the (renormalized + scaled) score.
+            routing_map: ``[N, num_experts]`` bool tensor; ``True`` at
+                ``(n, e)`` iff token ``n`` is routed to expert ``e``.
+        """
+        if (token_ids.dtype != torch.long) and (token_ids.dtype != torch.int):
+            raise TypeError(f"DeepseekV4HashRouter expects integer token_ids, got {token_ids.dtype}")
+        flat_hidden = hidden.reshape(-1, self.hidden_size)
+        flat_ids = token_ids.reshape(-1).long()
+        if flat_hidden.shape[0] != flat_ids.shape[0]:
+            raise ValueError(
+                "DeepseekV4HashRouter: hidden and token_ids must flatten to the same length; "
+                f"got hidden={flat_hidden.shape[0]} vs token_ids={flat_ids.shape[0]}."
+            )
+        if flat_ids.numel() == 0:
+            n_zero = 0
+            device = flat_hidden.device
+            probs = torch.zeros(n_zero, self.num_experts, dtype=torch.float32, device=device)
+            routing_map = torch.zeros(n_zero, self.num_experts, dtype=torch.bool, device=device)
+            return probs, routing_map
+        # NOTE: the token-id bounds check below is intentionally disabled — the
+        # ``.item()`` forces a device->host sync every hash-router forward, which
+        # stalls the CPU (shows up as a blocking cpu_op in the trace). token_ids
+        # come straight from the (already-validated) input pipeline, so the check
+        # is redundant on the hot path. Re-enable for debugging if needed.
+        # if int(flat_ids.max().item()) >= self.vocab_size:
+        #     raise ValueError(
+        #         f"token_ids has values >= vocab_size ({self.vocab_size}); "
+        #         f"max found = {int(flat_ids.max().item())}"
+        #     )
+
+        # Learned scores (fp32) over the full expert axis.
+        logits = F.linear(flat_hidden.to(torch.float32), self.weight.to(torch.float32))
+
+        # Static expert assignment from the table — cast to long for gather.
+        indices = self.tid2eid[flat_ids].long()  # [N, K]
+
+        # The V4 router is GPU-only: both the Triton and eager paths run on
+        # CUDA/HIP tensors (there is no CPU compute path).
+        assert logits.is_cuda, "V4 router requires CUDA / HIP tensors"
+        # Plan-6 P39: route the post-logits chain through one fused Triton
+        # kernel when PRIMUS_V4_ROUTER_TRITON=1 (default ON); else the eager body.
+        if _v4_router_triton_enabled():
+            probs, routing_map = v4_router_post_triton(
+                logits,
+                indices,
+                score_function=self.score_function,
+                topk_scaling_factor=self.topk_scaling_factor,
+                out_dtype=torch.float32,
+            )
+            return probs, routing_map
+
+        # Eager fallback (verbatim pre-P39 body).
+        scores = v4_score_fn(logits, score_function=self.score_function)  # [N, E]
+        weights = scores.gather(1, indices)  # [N, K]
+
+        if self.score_function != "softmax":
+            denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12)
+            weights = weights / denom
+
+        if self.topk_scaling_factor != 1.0:
+            weights = weights * float(self.topk_scaling_factor)
+
+        N = flat_hidden.shape[0]
+        device = flat_hidden.device
+
+        probs = torch.zeros(N, self.num_experts, dtype=weights.dtype, device=device)
+        probs.scatter_(1, indices, weights)
+
+        routing_map = torch.zeros(N, self.num_experts, dtype=torch.bool, device=device)
+        routing_map.scatter_(1, indices, True)
+
+        return probs, routing_map
+
+
+# Back-compat alias. New callers should use ``DeepseekV4HashRouter``.
+HashRouter = DeepseekV4HashRouter
+
+__all__ = [
+    "DeepseekV4HashRouter",
+    "HashRouter",
+]
diff --git a/primus/backends/megatron/core/transformer/moe/v4_moe.py b/primus/backends/megatron/core/transformer/moe/v4_moe.py
new file mode 100644
index 000000000..681769d1e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/v4_moe.py
@@ -0,0 +1,911 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 Mixture-of-Experts module.
+
+Reference: techblog §4 ("MoE: hash routing + sqrtsoftplus + shared experts")
+and ``DeepSeek-V4-Flash/inference/model.py:MoE``.
+
+V4's MoE block has three pieces:
+
+1. **Router** — either :class:`DeepseekV4HashRouter` (first
+   ``num_hash_layers`` layers) or :class:`DeepseekV4LearnedRouter` (the
+   rest). Both produce the same ``(probs, routing_map)`` shape contract:
+   ``[N, num_experts]``. The two routers share a learned gate weight;
+   only the *selection* differs (top-K argmax for the learned router,
+   ``tid2eid`` lookup for the hash router). Routing weights always come
+   from the same ``v4_score_fn(linear(hidden, weight))`` path.
+2. **Routed experts** — ``num_experts`` clamped-SwiGLU MLPs. Each token
+   contributes to ``moe_router_topk`` of them, weighted by the router
+   probability. The clamp is **pre-multiplication**:
+   ``SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha)``.
+3. **Shared expert(s)** — always-on MLP(s) whose output is added to every
+   token's contribution. V4-Flash has 1 shared expert with the same
+   ``moe_intermediate_size`` as the routed experts.
+
+Plan-2 P14 contract:
+
+P14 phase-1 (committed in 1a8bf32e) — math + parameter-layout
+faithfulness: pre-multiplication clamped SwiGLU activation, learned
+router rewritten with HF-aligned scoring + bias-only-for-selection
+semantics, hash router rewritten with a learnable gate weight + frozen
+``tid2eid`` Parameter.
+
+P14 phase-2 (this commit) — structural bring-up:
+* :class:`DeepseekV4MoE` now subclasses :class:`MegatronModule` (was
+  ``nn.Module``) so it integrates with Megatron's spec lifecycle and
+  shares config plumbing with the rest of the V4 stack.
+* CPU-friendly local-experts path: when ``pg_collection`` is ``None``
+  (or when the grouped backend does not declare clamped-SwiGLU support),
+  :class:`DeepseekV4MoE` builds a :class:`nn.ModuleList` of
+  :class:`ClampedSwiGLUMLP` routed experts plus a single
+  :class:`ClampedSwiGLUMLP` shared expert and runs a per-expert dispatch
+  loop in ``forward`` that mirrors the HF reference exactly. This makes
+  the MoE forward unit-testable on CPU at G5 (1L MoE forward agreement
+  vs HF reference within 1e-3 fp32) without requiring distributed init.
+* :meth:`set_layer_number` mirrors :class:`BaseMoELayer` so this module
+  slots into ``TransformerLayer`` via the spec lifecycle.
+* :attr:`local_expert_indices` exposed for compatibility with downstream
+  tooling that expects the ``BaseMoELayer`` public surface.
+
+Load balancing follows the paper: the auxiliary-loss-free expert bias plus a
+sequence-wise balance loss, the latter implemented directly on the learned
+router (see :mod:`...moe.v4_seq_balance_loss`) and driven by
+``moe_router_load_balancing_type: seq_aux_loss`` + ``moe_aux_loss_coeff``.
+It is not inherited from the framework's :class:`TopKRouter` for two reasons: the
+V4 routers are standalone ``nn.Module``\\ s (the parent registers CUDA buffers in
+``__init__`` and is impractical to instantiate on CPU), and the built-in aux-loss
+scores helper only handles ``softmax`` / ``sigmoid`` and rejects V4's
+``sqrtsoftplus``. z-loss remains unimplemented. The
+distributed re-validation phase (P19) will re-introduce that path
+behind a TopKRouter subclass once the CUDA-buffer init is gated by a
+device check upstream.
+"""
+
+from __future__ import annotations
+
+import logging
+from copy import copy
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from megatron.core import parallel_state
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.moe.token_dispatcher import (
+    MoEAllGatherTokenDispatcher,
+    MoEAlltoAllTokenDispatcher,
+    MoEFlexTokenDispatcher,
+    MoETokenDispatcher,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.clamped_swiglu import ClampedSwiGLUMLP
+from primus.backends.megatron.core.transformer.moe.shared_experts import (
+    PrimusSharedExpertMLP,
+)
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+    DeepseekV4HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    DeepseekV4LearnedRouter,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DeepseekV4MoESubmodules:
+    """Spec tree for V4 MoE construction."""
+
+    hash_router: Optional[Union[ModuleSpec, type]] = DeepseekV4HashRouter
+    learned_router: Optional[Union[ModuleSpec, type]] = DeepseekV4LearnedRouter
+    token_dispatcher: Optional[Union[ModuleSpec, type]] = MoEAlltoAllTokenDispatcher
+    grouped_experts: Optional[Union[ModuleSpec, type]] = None
+    shared_expert: Optional[Union[ModuleSpec, type]] = SharedExpertMLP
+
+
+class DeepseekV4MoE(MegatronModule):
+    """V4 MoE FFN sub-block.
+
+    Args:
+        config: runtime DeepSeek-V4 config. Core MoE dimensions and router
+            options are read directly from config.
+        layer_idx: 0-based decoder layer index. Used to pick router type
+            against ``num_hash_layers``.
+        pg_collection: Megatron process-group collection. When ``None``
+            (CPU unit tests), the module skips the distributed dispatcher
+            and builds a local :class:`nn.ModuleList` of
+            :class:`ClampedSwiGLUMLP` routed experts plus a single
+            :class:`ClampedSwiGLUMLP` shared expert; ``forward`` runs a
+            per-expert dispatch loop matching the HF reference math.
+        submodules: spec tree describing routers / dispatcher / experts /
+            shared expert. Must be provided.
+        layer_number: optional 1-based layer number used by Megatron's
+            spec lifecycle (mirrors :class:`BaseMoELayer.set_layer_number`).
+    """
+
+    def __init__(
+        self,
+        config: DeepSeekV4TransformerConfig,
+        *,
+        layer_idx: int,
+        pg_collection=None,
+        submodules: Optional[DeepseekV4MoESubmodules] = None,
+        layer_number: Optional[int] = None,
+    ) -> None:
+        if config is None:
+            raise ValueError("DeepSeek-V4 MoE requires config.")
+        super().__init__(config=config)
+        self.pg_collection = pg_collection
+        self.submodules = submodules
+        assert self.submodules is not None, "DeepSeek-V4 MoE requires explicit submodules."
+        self.layer_number = layer_number
+
+        hidden_size = int(config.hidden_size)
+        moe_intermediate_size = int(
+            config.moe_ffn_hidden_size or config.moe_intermediate_size or config.ffn_hidden_size
+        )
+        num_routed_experts = int(config.num_moe_experts)
+        moe_router_topk = int(config.moe_router_topk)
+        use_shared_expert = config.moe_shared_expert_intermediate_size is not None
+        layer_num_hash_layers = int(config.num_hash_layers)
+        layer_hash_vocab_size = config.padded_vocab_size or config.vocab_size
+        layer_hash_seed = int(config.hash_routing_seed)
+        score_function = str(config.moe_router_score_function)
+        enable_expert_bias = bool(config.moe_router_enable_expert_bias)
+        topk_scaling_factor = float(getattr(config, "moe_router_topk_scaling_factor", 1.0) or 1.0)
+        clamp_alpha = float(config.swiglu_limit)
+
+        if num_routed_experts <= 0:
+            raise ValueError(f"num_routed_experts must be > 0, got {num_routed_experts}")
+        if moe_router_topk <= 0 or moe_router_topk > num_routed_experts:
+            raise ValueError(f"moe_router_topk must be in [1, {num_routed_experts}], got {moe_router_topk}")
+
+        self.hidden_size = hidden_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_routed_experts = num_routed_experts
+        self.moe_router_topk = moe_router_topk
+        self.use_shared_expert = use_shared_expert
+        self.layer_idx = int(layer_idx)
+        self.num_hash_layers = layer_num_hash_layers
+        self.use_hash_router = self.layer_idx < self.num_hash_layers
+        self.clamp_alpha = clamp_alpha
+        self.moe_token_dispatcher_type = "alltoall"
+
+        # ---- EP placement ----
+        self.ep_group = getattr(pg_collection, "ep", None) if pg_collection is not None else None
+        self.ep_size = 1
+        self.ep_rank = 0
+        if self.ep_group is None and dist.is_available() and dist.is_initialized():
+            try:
+                self.ep_group = parallel_state.get_expert_model_parallel_group()
+            except Exception:
+                self.ep_group = None
+        if self.ep_group is not None and dist.is_available() and dist.is_initialized():
+            self.ep_size = int(self.ep_group.size())
+            self.ep_rank = int(self.ep_group.rank())
+
+        base = self.num_routed_experts // self.ep_size
+        remainder = self.num_routed_experts % self.ep_size
+        self.local_num_routed_experts = base + (1 if self.ep_rank < remainder else 0)
+        self.local_expert_start = (self.ep_rank * base) + min(self.ep_rank, remainder)
+        self.local_expert_end = self.local_expert_start + self.local_num_routed_experts
+        # BaseMoELayer-compatible public attribute.
+        self.local_expert_indices = list(range(self.local_expert_start, self.local_expert_end))
+
+        # ---- routers ----
+        self.router = None
+        self.learned_router = None
+        self._force_lb_mode = self._resolve_force_load_balancing_mode()
+        self._build_router_modules(
+            hash_vocab_size=layer_hash_vocab_size,
+            hash_seed=layer_hash_seed,
+            score_function=score_function,
+            enable_expert_bias=enable_expert_bias,
+            topk_scaling_factor=topk_scaling_factor,
+        )
+
+        # ---- experts ----
+        # Production path: full Megatron dispatcher + grouped-experts.
+        # CPU path: a local nn.ModuleList of ClampedSwiGLUMLP experts + a
+        # single ClampedSwiGLUMLP shared expert. The CPU path is used when
+        # ``pg_collection is None`` so unit tests can drive ``forward``
+        # without distributed initialization.
+        self.token_dispatcher: Optional[MoETokenDispatcher] = None
+        self.grouped_experts: Optional[nn.Module] = None
+        self.local_experts: Optional[nn.ModuleList] = None
+        self.shared_expert: Optional[nn.Module] = None
+        # True once the dispatcher owns the shared expert, which also means the
+        # dispatcher adds its output (see ``_enable_shared_expert_overlap``).
+        self.shared_expert_overlap = False
+        self.mega_moe_experts: Optional[nn.Module] = None
+
+        if pg_collection is None:
+            self.local_experts = self._build_local_experts(intermediate_size=self.moe_intermediate_size)
+            if self.use_shared_expert:
+                assert self.config.moe_shared_expert_intermediate_size is not None
+                self.shared_expert = ClampedSwiGLUMLP(
+                    hidden_size=self.hidden_size,
+                    intermediate_size=int(self.config.moe_shared_expert_intermediate_size),
+                    alpha=self.clamp_alpha,
+                    bias=False,
+                )
+        elif self._mega_moe_requested():
+            # MegaMoE fuses dispatch/combine into the grouped GEMMs, so it
+            # replaces the dispatcher *and* the grouped experts. The V4 routers
+            # and the shared expert are untouched.
+            self.mega_moe_experts = self._build_mega_moe_experts()
+            if self.use_shared_expert:
+                assert self.config.moe_shared_expert_intermediate_size is not None
+                self.shared_expert = self._build_shared_expert_module(
+                    intermediate_size=int(self.config.moe_shared_expert_intermediate_size)
+                )
+        else:
+            self.token_dispatcher = self._build_token_dispatcher()
+            self.grouped_experts = self._build_grouped_experts()
+            if self.use_shared_expert:
+                assert self.config.moe_shared_expert_intermediate_size is not None
+                self.shared_expert = self._build_shared_expert_module(
+                    intermediate_size=int(self.config.moe_shared_expert_intermediate_size)
+                )
+                self.shared_expert_overlap = self._enable_shared_expert_overlap()
+
+    # ------------------------------------------------------------------
+
+    def set_layer_number(self, layer_number: int) -> None:
+        """Mirror :class:`BaseMoELayer.set_layer_number` for spec lifecycle.
+
+        Megatron's :class:`TransformerLayer` walks every spec submodule
+        with a ``set_layer_number`` method to populate the 1-based layer
+        index. The V4 routers are intentionally standalone (CPU-clean),
+        but we still need to track ``layer_number`` here so future
+        TopKRouter-rooted upgrades plug in without spec changes.
+        """
+        self.layer_number = layer_number
+        # The router is built before this runs, so forward the index it needs
+        # to attribute its balance loss to the right layer in the tracker.
+        if self.learned_router is not None:
+            self.learned_router.layer_number = layer_number
+
+    def _mega_moe_requested(self) -> bool:
+        """Whether the fused MegaMoE expert path was asked for.
+
+        The upstream patch that swaps Megatron's ``MoELayer`` cannot reach V4 --
+        :class:`DeepseekV4MoE` is built directly by ``deepseek_v4_block`` and
+        never instantiates ``MoELayer`` -- so V4 drives
+        :class:`MegaMoEExperts` itself.
+
+        The turbo switches live on the Megatron ``args`` namespace, not on
+        :class:`TransformerConfig` (same as ``args.enable_primus_turbo`` in
+        :mod:`primus.backends.megatron.core.transformer.moe.router`), so read
+        them from there. ``get_args`` raises before Megatron is initialized,
+        which is the CPU unit-test path -- treat that as "not requested".
+        """
+        try:
+            from megatron.training import get_args
+
+            args = get_args()
+        except Exception:
+            return False
+        return bool(getattr(args, "enable_primus_turbo", False)) and bool(
+            getattr(args, "use_turbo_mega_moe", False)
+        )
+
+    @staticmethod
+    def _resolve_force_load_balancing_mode() -> Optional[str]:
+        """Return the requested force-load-balancing mode, or ``None`` if off.
+
+        Read from the Megatron ``args`` namespace like
+        :meth:`_mega_moe_requested`; ``get_args`` raises before Megatron is
+        initialized (the CPU unit-test path), which counts as "off".
+        """
+        try:
+            from megatron.training import get_args
+
+            args = get_args()
+        except Exception:
+            return None
+        if not bool(getattr(args, "moe_router_force_load_balancing", False)):
+            return None
+        mode = str(getattr(args, "moe_router_force_load_balancing_type", "uniform"))
+        if mode not in ("even", "uniform"):
+            raise ValueError(
+                f"moe_router_force_load_balancing_type must be 'even' or 'uniform', got {mode!r}"
+            )
+        return mode
+
+    def _apply_force_load_balancing(
+        self, probs: torch.Tensor, routing_map: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Rewrite the routed expert set so every expert gets a similar share.
+
+        Benchmark aid, mirroring ``moe_router_force_load_balancing`` on Megatron's
+        ``TopKRouter``. V4's routers are plain ``nn.Module``s rather than
+        ``TopKRouter`` subclasses, so neither upstream's ``apply_random_logits``
+        (``TopKRouter.forward``) nor ``PrimusTopKRouter._force_even_routing``
+        reaches this path -- this is the V4 equivalent of both, and it covers the
+        hash-routed layers too (force load balancing exists to discard the
+        router's decision, hash tables included).
+
+        Both modes carry the real per-token top-k probability magnitudes over to
+        the new slots, so ``probs`` and ``routing_map`` stay mutually consistent
+        and the combine weights keep their true scale:
+
+        * ``even``    -- round-robin ``(token_idx * topk + k) % num_experts``.
+          Per-expert counts are exactly equal and step-invariant, so
+          grouped-GEMM shapes never change. This is the same cycle
+          ``PrimusTurboDeepEPTokenDispatcher.dispatch_preprocess`` derives on its
+          own, so when Turbo DeepEP re-applies it the two agree instead of
+          dispatching to one expert set while weighting by another.
+        * ``uniform`` -- a fresh random k-subset per token every step, i.e. the
+          selection ``apply_random_logits`` + top-k produces upstream. Balanced
+          only statistically, so it keeps the step-to-step shape variation of
+          real routing.
+        """
+        if self._force_lb_mode is None:
+            return probs, routing_map
+
+        num_tokens = routing_map.shape[0]
+        topk = self.moe_router_topk
+        num_experts = self.num_routed_experts
+        device = routing_map.device
+
+        if self._force_lb_mode == "even":
+            # topk consecutive experts per token -> distinct while topk <= num_experts.
+            slot = torch.arange(num_tokens * topk, device=device).view(num_tokens, topk) % num_experts
+        else:
+            slot = torch.rand(num_tokens, num_experts, device=device).topk(topk, dim=-1).indices
+
+        new_routing_map = torch.zeros_like(routing_map)
+        new_routing_map.scatter_(1, slot, torch.ones_like(slot, dtype=routing_map.dtype))
+
+        # probs is non-zero only on the real top-k, so topk() extracts exactly them.
+        topk_vals, _ = torch.topk(probs, topk, dim=1)
+        new_probs = torch.zeros_like(probs)
+        new_probs.scatter_(1, slot, topk_vals.to(probs.dtype))
+
+        return new_probs, new_routing_map
+
+    def _build_mega_moe_experts(self) -> nn.Module:
+        """Build the fused MegaMoE expert path, or explain why it cannot run.
+
+        Every unmet requirement raises rather than falling back: a silent
+        fallback leaves the run with neither MegaMoE nor the Turbo DeepEP
+        dispatcher that ``use_turbo_mega_moe`` disables, which shows up only as
+        an unexplained slowdown.
+        """
+        from primus.backends.megatron.core.extensions.mega_moe import MegaMoEExperts
+
+        reasons = []
+        if self.config.tensor_model_parallel_size != 1:
+            reasons.append(f"MegaMoE is EP-only (TP==1), got TP={self.config.tensor_model_parallel_size}")
+        if self.config.params_dtype != torch.bfloat16:
+            reasons.append(f"MegaMoE only supports bf16, got params_dtype={self.config.params_dtype}")
+        if self.ep_group is None or self.ep_size <= 1:
+            reasons.append(f"MegaMoE needs an expert-parallel group, got EP={self.ep_size}")
+        if self.num_routed_experts % self.ep_size != 0:
+            reasons.append(
+                f"MegaMoE needs num_experts divisible by EP size, got "
+                f"{self.num_routed_experts} % {self.ep_size} != 0"
+            )
+        if reasons:
+            raise ValueError(
+                "[DeepSeek-V4] use_turbo_mega_moe=True but the fused MegaMoE path cannot run here: "
+                + "; ".join(reasons)
+            )
+
+        logger.warning(
+            "[DeepSeek-V4] MegaMoE expert path enabled: the fused kernel hardcodes unclamped "
+            "SwiGLU, so swiglu_limit=%s is NOT applied to routed experts.",
+            self.clamp_alpha,
+        )
+        experts = MegaMoEExperts(
+            self.config,
+            self.local_num_routed_experts,
+            self.hidden_size,
+            self.moe_intermediate_size,
+            self.ep_group,
+        )
+        if self.config.perform_initialization:
+            experts.reset_parameters(self.ep_rank)
+        logger.info(
+            "[DeepSeek-V4] MoE expert path resolved to MegaMoEExperts "
+            "(fused dispatch/combine; token dispatcher and grouped experts not built)."
+        )
+        return experts
+
+    def _mega_moe_forward(self, hidden: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
+        """Run the fused MegaMoE expert path.
+
+        ``probs`` is the dense ``[N, E]`` router output while MegaMoE wants the
+        compact ``(topk_idx, topk_weights)`` pair. ``probs`` is zero outside the
+        selected experts, so ``topk`` recovers exactly the routed set -- the
+        same conversion :class:`PrimusTurboMegaMoELayer` performs.
+        """
+        assert self.mega_moe_experts is not None
+        in_shape = hidden.shape
+        topk_weights, topk_idx = probs.topk(self.moe_router_topk, dim=-1)
+        x = hidden.reshape(-1, self.hidden_size).to(torch.bfloat16)
+        y = self.mega_moe_experts(x, topk_idx, topk_weights.to(torch.float32))
+        return y.reshape(in_shape).to(hidden.dtype)
+
+    def _build_local_experts(self, *, intermediate_size: int) -> nn.ModuleList:
+        """Build a local :class:`nn.ModuleList` of clamped-SwiGLU experts.
+
+        Used when ``pg_collection is None`` (CPU unit tests). Each module
+        in the list mirrors a single HF reference ``Expert`` (separate
+        ``w1`` / ``w2`` / ``w3`` Linears + V4 pre-mul clamp).
+        """
+        if self.local_num_routed_experts <= 0:
+            raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts.")
+        return nn.ModuleList(
+            [
+                ClampedSwiGLUMLP(
+                    hidden_size=self.hidden_size,
+                    intermediate_size=intermediate_size,
+                    alpha=self.clamp_alpha,
+                    bias=False,
+                )
+                for _ in range(self.local_num_routed_experts)
+            ]
+        )
+
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _resolve_dispatcher_type_from_spec(dispatcher_spec: Optional[Union[ModuleSpec, type]]) -> str:
+        module = dispatcher_spec.module if isinstance(dispatcher_spec, ModuleSpec) else dispatcher_spec
+        if module is MoEAllGatherTokenDispatcher:
+            return "allgather"
+        if module is MoEFlexTokenDispatcher:
+            return "flex"
+        if module is MoEAlltoAllTokenDispatcher or module is None:
+            return "alltoall"
+        # Plan-3 P23: PrimusTurboDeepEPTokenDispatcher is a "flex"
+        # variant — V4 spec build chose it explicitly when the user
+        # opted into Turbo DeepEP.  We recognise it by class name so
+        # this resolver does not require ``primus_turbo`` to be
+        # importable on hosts that never opt in (CPU unit tests,
+        # Megatron-only consumers).
+        module_name = getattr(module, "__name__", "")
+        if module_name == "PrimusTurboDeepEPTokenDispatcher":
+            return "flex"
+        logger.warning(
+            "[DeepSeek-V4] unsupported dispatcher module=%s; fallback type to alltoall.",
+            module_name or str(module),
+        )
+        return "alltoall"
+
+    def _build_router_modules(
+        self,
+        *,
+        hash_vocab_size: Optional[int],
+        hash_seed: int,
+        score_function: str,
+        enable_expert_bias: bool,
+        topk_scaling_factor: float,
+    ) -> None:
+        if self.use_hash_router:
+            if hash_vocab_size is None or hash_vocab_size <= 0:
+                raise ValueError(
+                    "hash_vocab_size must be provided (and > 0) when layer_idx < num_hash_layers"
+                )
+            hash_router_spec = self.submodules.hash_router or DeepseekV4HashRouter
+            self.router = build_module(
+                hash_router_spec,
+                hidden_size=self.hidden_size,
+                num_experts=self.num_routed_experts,
+                topk=self.moe_router_topk,
+                vocab_size=hash_vocab_size,
+                seed=hash_seed,
+                score_function=score_function,
+                topk_scaling_factor=topk_scaling_factor,
+            )
+            self.learned_router = None
+            return
+
+        learned_router_spec = self.submodules.learned_router or DeepseekV4LearnedRouter
+        self.router = None
+        self.learned_router = build_module(
+            learned_router_spec,
+            hidden_size=self.hidden_size,
+            num_experts=self.num_routed_experts,
+            topk=self.moe_router_topk,
+            score_function=score_function,
+            enable_expert_bias=enable_expert_bias,
+            topk_scaling_factor=topk_scaling_factor,
+            seq_balance_loss_coeff=self._seq_balance_loss_coeff(),
+            seq_balance_reduce_group=self._seq_balance_reduce_group(),
+            layer_number=self.layer_number,
+            num_layers=self._total_layers_for_loss_logging(),
+        )
+
+    def _seq_balance_loss_coeff(self) -> float:
+        """Coefficient of the sequence-wise balance loss, 0 when not selected.
+
+        V4 balances with the auxiliary-loss-free expert bias plus this
+        sequence-wise term (paper 2.1), which is what
+        ``moe_router_load_balancing_type: seq_aux_loss`` in the model yaml asks
+        for. Any other balancing type leaves the loss off -- the framework's
+        built-in aux-loss variants are not reachable from this router.
+        """
+        config = self.config
+        if str(getattr(config, "moe_router_load_balancing_type", "")) != "seq_aux_loss":
+            return 0.0
+        return float(getattr(config, "moe_aux_loss_coeff", 0.0) or 0.0)
+
+    def _seq_balance_reduce_group(self):
+        """Group the sequence axis is sharded over, or ``None`` if it is not.
+
+        The loss is per-sequence, so with sequence parallelism (TP) or context
+        parallelism a rank only sees part of each sequence and the expert counts
+        have to be summed across the shards.
+        """
+        tp_size = int(getattr(self.config, "tensor_model_parallel_size", 1) or 1)
+        cp_size = int(getattr(self.config, "context_parallel_size", 1) or 1)
+        sequence_parallel = bool(getattr(self.config, "sequence_parallel", False))
+        if cp_size == 1 and not (sequence_parallel and tp_size > 1):
+            return None
+        group = getattr(self.pg_collection, "tp_cp", None)
+        if group is not None:
+            return group
+        from megatron.core import parallel_state
+
+        return parallel_state.get_tensor_and_context_parallel_group()
+
+    def _total_layers_for_loss_logging(self) -> Optional[int]:
+        """Tracker width: decoder layers plus MTP layers, matching the framework."""
+        num_layers = getattr(self.config, "num_layers", None)
+        if not num_layers:
+            return None
+        return int(num_layers) + int(getattr(self.config, "mtp_num_layers", 0) or 0)
+
+    def _enable_shared_expert_overlap(self) -> bool:
+        """Hand the shared expert to the dispatcher so it runs under the A2A.
+
+        ``moe_shared_expert_overlap`` asks for the shared expert's GEMMs to be
+        interleaved with the dispatch / combine all-to-alls instead of running
+        after them. The token dispatcher already carries the hooks, and they sit
+        in exactly the decomposed calls :meth:`_dispatcher_forward` already
+        makes: ``dispatch_preprocess`` kicks off the input comm on a side stream,
+        ``dispatch_postprocess`` runs fc1 + activation, and
+        ``combine_postprocess`` runs fc2 and **adds the result into the output**.
+
+        That last part is why the return value matters to the caller: once the
+        dispatcher owns the shared expert, adding it again in ``forward`` would
+        double-count it.
+
+        Returns False (leaving the serial path in place) when the config asks for
+        no overlap or the dispatcher does not support it -- the flex dispatcher
+        raises ``NotImplementedError`` from ``set_shared_experts``.
+        """
+        if not bool(getattr(self.config, "moe_shared_expert_overlap", False)):
+            return False
+        if self.shared_expert is None or self.token_dispatcher is None:
+            return False
+        if not hasattr(self.shared_expert, "pre_forward_comm"):
+            # The CPU ClampedSwiGLUMLP has no overlap protocol.
+            return False
+        try:
+            self.token_dispatcher.set_shared_experts(self.shared_expert)
+        except NotImplementedError:
+            logger.warning(
+                "[V4-MoE] layer %s: %s does not support shared-expert overlap; "
+                "running the shared expert serially.",
+                self.layer_idx,
+                type(self.token_dispatcher).__name__,
+            )
+            return False
+        return True
+
+    def _build_shared_expert_module(self, *, intermediate_size: int) -> nn.Module:
+        shared_expert_spec = self.submodules.shared_expert
+        assert isinstance(
+            shared_expert_spec, ModuleSpec
+        ), "DeepSeek-V4 MoE requires shared_expert ModuleSpec in submodules."
+        shared_expert_module = shared_expert_spec.module
+        assert issubclass(
+            shared_expert_module, SharedExpertMLP
+        ), "DeepSeek-V4 shared_expert must be (a subclass of) SharedExpertMLP."
+        if self.config is None or self.pg_collection is None:
+            raise RuntimeError("DeepSeek-V4 MoE SharedExpertMLP requires config and pg_collection.")
+
+        # Shared experts run with clamped SwiGLU. PrimusSharedExpertMLP fuses the
+        # clamp+SiLU+mul into a single Triton kernel (matching the routed experts),
+        # so we no longer need to force the un-fused eager path.
+        shared_cfg = copy(self.config)
+        shared_cfg.add_bias_linear = False
+        shared_cfg.gated_linear_unit = True
+        shared_cfg.activation_func = F.silu
+        shared_cfg.bias_activation_fusion = False
+        shared_cfg.use_te_activation_func = False
+        if self.clamp_alpha > 0:
+            shared_cfg.activation_func_clamp_value = float(self.clamp_alpha)
+        else:
+            shared_cfg.activation_func_clamp_value = None
+        if int(shared_cfg.moe_shared_expert_intermediate_size or 0) <= 0:
+            setattr(
+                shared_cfg,
+                "moe_shared_expert_intermediate_size",
+                int(intermediate_size),
+            )
+
+        # Build the Primus fused-SwiGLU shared expert while keeping the spec's
+        # submodules (linear layers / activation). The state-dict layout is
+        # unchanged since PrimusSharedExpertMLP only overrides the activation.
+        fused_shared_expert_spec = ModuleSpec(
+            module=PrimusSharedExpertMLP,
+            submodules=shared_expert_spec.submodules,
+            params=shared_expert_spec.params,
+        )
+        try:
+            return build_module(
+                fused_shared_expert_spec,
+                config=shared_cfg,
+                pg_collection=self.pg_collection,
+                gate=bool(shared_cfg.moe_shared_expert_gate),
+            )
+        except Exception as exc:
+            raise RuntimeError(
+                f"DeepSeek-V4 MoE shared expert build failed with PrimusSharedExpertMLP: {exc}"
+            ) from exc
+
+    def _build_token_dispatcher(self) -> MoETokenDispatcher:
+        if self.config is None or self.pg_collection is None:
+            raise RuntimeError(
+                "DeepSeek-V4 MoE requires config and pg_collection for Megatron dispatcher path."
+            )
+        if self.local_num_routed_experts <= 0:
+            raise RuntimeError(
+                f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts for dispatcher path."
+            )
+
+        dispatcher_spec: Union[ModuleSpec, type, None] = self.submodules.token_dispatcher
+        assert dispatcher_spec is not None, "DeepSeek-V4 MoE requires token_dispatcher spec in submodules."
+        requested_dispatcher_type = self._resolve_dispatcher_type_from_spec(dispatcher_spec)
+        ep_group = getattr(self.pg_collection, "ep", None)
+        tp_ep_group = getattr(self.pg_collection, "tp_ep", None)
+        if requested_dispatcher_type == "alltoall" and ep_group is None:
+            logger.info(
+                "[DeepSeek-V4] MoE layer=%s alltoall dispatcher requires EP group.",
+                self.layer_idx,
+            )
+        if requested_dispatcher_type == "flex" and tp_ep_group is None:
+            logger.info(
+                "[DeepSeek-V4] MoE layer=%s flex dispatcher requires TPxEP group.",
+                self.layer_idx,
+            )
+        self.moe_token_dispatcher_type = requested_dispatcher_type
+
+        local_expert_indices = list(range(self.local_expert_start, self.local_expert_end))
+        try:
+            dispatcher = build_module(
+                dispatcher_spec,
+                num_local_experts=self.local_num_routed_experts,
+                local_expert_indices=local_expert_indices,
+                config=self.config,
+                pg_collection=self.pg_collection,
+            )
+            logger.info(
+                "[DeepSeek-V4] MoE layer=%s dispatcher active via %s.",
+                self.layer_idx,
+                type(dispatcher).__name__,
+            )
+            return dispatcher
+        except Exception as exc:
+            raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} dispatcher build failed: {exc}")
+
+    def _route(
+        self,
+        hidden: torch.Tensor,
+        token_ids: Optional[torch.Tensor],
+    ):
+        """Return ``(probs, routing_map)`` for the current router.
+
+        Hash-routed layers feed both ``hidden`` (for the learned routing
+        weights) AND ``token_ids`` (for the static expert ids from
+        ``tid2eid``); learned layers only consume ``hidden``.
+
+        The routed set then goes through :meth:`_apply_force_load_balancing`,
+        which is a no-op unless ``moe_router_force_load_balancing`` is set.
+        """
+        if self.use_hash_router:
+            assert self.router is not None
+            if token_ids is None:
+                raise ValueError(
+                    f"layer {self.layer_idx} uses DeepseekV4HashRouter; "
+                    "token_ids is required (shape [B, S])."
+                )
+            probs, routing_map = self.router(hidden, token_ids)
+        else:
+            assert self.learned_router is not None
+            probs, routing_map = self.learned_router(hidden)
+        return self._apply_force_load_balancing(probs, routing_map)
+
+    # ------------------------------------------------------------------
+
+    def _build_grouped_experts(self):
+        grouped_experts_spec: Optional[Union[ModuleSpec, type]] = self.submodules.grouped_experts
+        assert (
+            grouped_experts_spec is not None
+        ), "DeepSeek-V4 MoE requires grouped experts spec in submodules."
+        if self.local_num_routed_experts <= 0:
+            raise RuntimeError(
+                f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts for grouped backend."
+            )
+
+        if self.config is None or self.pg_collection is None:
+            raise RuntimeError("DeepSeek-V4 MoE requires config and pg_collection to build grouped experts.")
+
+        try:
+            module = build_module(
+                grouped_experts_spec,
+                num_local_experts=self.local_num_routed_experts,
+                config=self.config,
+                pg_collection=self.pg_collection,
+            )
+            if not self._grouped_backend_supports_clamped_swiglu(module):
+                raise RuntimeError(
+                    "DeepSeek-V4 MoE grouped backend "
+                    f"{type(module).__name__} does not declare clamped-SwiGLU support. "
+                    "Set `v4_grouped_experts_support_clamped_swiglu=True` only "
+                    "after backend parity is validated."
+                )
+            logger.info(
+                "[DeepSeek-V4] MoE layer=%s provider grouped-gemm active via %s.",
+                self.layer_idx,
+                type(module).__name__,
+            )
+            return module
+        except Exception as exc:
+            raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} grouped experts build failed: {exc}")
+
+    def _grouped_backend_supports_clamped_swiglu(self, module: nn.Module) -> bool:
+        if self.clamp_alpha <= 0:
+            return True
+        if bool(getattr(module, "supports_clamped_swiglu", False)):
+            return True
+        if self.config is not None and bool(self.config.v4_grouped_experts_support_clamped_swiglu):
+            return True
+        return False
+
+    def _dispatcher_expert_forward(
+        self,
+        permuted_hidden: torch.Tensor,
+        tokens_per_expert: torch.Tensor,
+        permuted_probs: torch.Tensor,
+        routing_map: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.grouped_experts is not None
+        try:
+            grouped_out = self.grouped_experts(
+                permuted_hidden,
+                tokens_per_expert,
+                permuted_probs,
+                routing_map=routing_map,
+            )
+        except TypeError:
+            grouped_out = self.grouped_experts(
+                permuted_hidden,
+                tokens_per_expert,
+                permuted_probs,
+            )
+        if isinstance(grouped_out, tuple):
+            return grouped_out[0]
+        return grouped_out
+
+    def _dispatcher_forward(
+        self,
+        hidden: torch.Tensor,
+        probs: torch.Tensor,
+        routing_map: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.token_dispatcher is not None
+        hidden_states, probs = self.token_dispatcher.dispatch_preprocess(hidden, routing_map, probs)
+        hidden_states, probs = self.token_dispatcher.token_dispatch(hidden_states, probs)
+        expert_input, tokens_per_expert, permuted_probs = self.token_dispatcher.dispatch_postprocess(
+            hidden_states, probs
+        )
+
+        expert_output = self._dispatcher_expert_forward(
+            expert_input,
+            tokens_per_expert,
+            permuted_probs,
+            routing_map,
+        )
+
+        combined = self.token_dispatcher.combine_preprocess(expert_output)
+        combined = self.token_dispatcher.token_combine(combined)
+        return self.token_dispatcher.combine_postprocess(combined)
+
+    def _local_experts_forward(
+        self,
+        hidden: torch.Tensor,
+        probs: torch.Tensor,
+        routing_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Per-expert dispatch loop matching the HF reference math.
+
+        Drives :attr:`local_experts` directly (no token dispatcher); used
+        on the CPU path when ``pg_collection is None``. The math mirrors
+        ``DeepSeek-V4-Flash/inference/model.py:MoE.forward`` exactly:
+
+            for i in local_experts:
+                idx = where(routing_map[:, i])
+                out[idx] += probs[idx, i] * expert_i(hidden[idx])
+
+        Args:
+            hidden: ``[N, D]`` flattened input.
+            probs: ``[N, num_experts]`` sparse routing weights (already
+                renormalized + scaled by the router).
+            routing_map: ``[N, num_experts]`` bool mask, True at
+                ``(n, e)`` iff token ``n`` is routed to expert ``e``.
+
+        Returns:
+            ``[N, D]`` routed-expert contribution (no shared expert).
+        """
+        assert self.local_experts is not None
+        out = torch.zeros_like(hidden, dtype=hidden.dtype)
+        for local_i, global_i in enumerate(self.local_expert_indices):
+            mask = routing_map[:, global_i]  # [N]
+            if not bool(mask.any()):
+                continue
+            idx = mask.nonzero(as_tuple=True)[0]  # [n_i]
+            weight = probs[idx, global_i].unsqueeze(-1).to(hidden.dtype)  # [n_i, 1]
+            expert = self.local_experts[local_i]
+            out_idx = expert(hidden[idx])
+            out[idx] = out[idx] + weight * out_idx
+        return out
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        *,
+        token_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Run V4 MoE FFN.
+
+        Args:
+            hidden: ``[B, S, D]`` input.
+            token_ids: ``[B, S]`` integer token ids, required only when
+                ``layer_idx < num_hash_layers``.
+
+        Returns:
+            ``[B, S, D]`` output. Sum of routed-expert and shared-expert
+            contributions.
+        """
+        probs, routing_map = self._route(hidden, token_ids)  # [N, E], bool
+
+        if self.local_experts is not None:
+            # CPU local-experts path. Reshape to flat then back; the
+            # router already returned [N, E] sparse outputs.
+            shape = hidden.shape
+            flat_hidden = hidden.reshape(-1, self.hidden_size)
+            flat_out = self._local_experts_forward(flat_hidden, probs, routing_map)
+            if self.shared_expert is not None:
+                flat_out = flat_out + self.shared_expert(flat_hidden)
+            return flat_out.view(*shape)
+
+        # Production path: fused MegaMoE experts, or Megatron dispatcher +
+        # grouped experts.
+        if self.mega_moe_experts is not None:
+            out = self._mega_moe_forward(hidden, probs)
+        else:
+            out = self._dispatcher_forward(hidden, probs, routing_map)
+        if self.shared_expert is not None and not self.shared_expert_overlap:
+            out = out + self.shared_expert(hidden)
+        return out
+
+
+__all__ = ["DeepseekV4MoE", "DeepseekV4MoESubmodules"]
diff --git a/primus/backends/megatron/core/transformer/moe/v4_seq_balance_loss.py b/primus/backends/megatron/core/transformer/moe/v4_seq_balance_loss.py
new file mode 100644
index 000000000..4498994fc
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/v4_seq_balance_loss.py
@@ -0,0 +1,152 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Sequence-wise load-balancing loss for the DeepSeek-V4 MoE router.
+
+V4 balances experts with the auxiliary-loss-free bias (``expert_bias`` /
+"noaux_tc"), "augmented by a slight sequence-wise balance loss that prevents
+extreme imbalance within individual sequences" (paper 2.1). The bias reacts over
+many steps and across the whole batch, so it cannot stop a single sequence from
+collapsing onto a handful of experts; this loss is what covers that.
+
+Following DeepSeek-V3, for each sequence of length ``T`` with ``E`` routed
+experts and top-``K`` routing:
+
+.. math::
+
+    s'_{i,t} = s_{i,t} / \\sum_j s_{j,t}
+
+    f_i = \\frac{E}{K T} \\sum_t \\mathbb{1}(\\text{token } t \\to \\text{expert } i)
+
+    P_i = \\frac{1}{T} \\sum_t s'_{i,t}
+
+    L = \\alpha \\sum_i f_i P_i
+
+``f`` sums to ``E`` and ``P`` sums to 1, so a perfectly balanced sequence gives
+exactly ``L = alpha`` -- which makes the logged value directly readable.
+
+Why this is not the framework's built-in ``seq_aux_loss``: that path runs
+through ``compute_routing_scores_for_aux_loss``, which only knows ``softmax``
+and ``sigmoid`` and raises on V4's ``sqrtsoftplus``. The normalisation it
+applies to ``sigmoid`` is the same one used here.
+
+Two deliberate choices about which quantities feed the formula:
+
+* ``f`` counts the **actual** routing, i.e. with ``expert_bias`` applied. The
+  paper defines it over the tokens the router really dispatched, and the point of
+  the loss is real load, not the router's unbiased preference.
+* ``P`` uses the **unbiased** normalised affinities, which is what ``s'_{i,t}``
+  denotes and the only part with a gradient.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+__all__ = [
+    "SEQ_BALANCE_LOSS_NAME",
+    "log_seq_balance_loss",
+    "normalise_affinities",
+    "sequence_balance_loss",
+]
+
+# Reported through the framework's MoE aux-loss tracker.
+SEQ_BALANCE_LOSS_NAME = "seq_load_balancing_loss"
+
+_EPS = 1e-20
+
+
+def normalise_affinities(scores: torch.Tensor) -> torch.Tensor:
+    """``s'_{i,t} = s_{i,t} / sum_j s_{j,t}`` over the expert axis.
+
+    V4's ``sqrtsoftplus`` score function is pointwise and non-negative but does
+    not sum to one, so the affinities need normalising before they can act as
+    the ``P_i`` distribution -- the same treatment the framework's aux-loss path
+    gives ``sigmoid``.
+    """
+    return scores / (scores.sum(dim=-1, keepdim=True) + _EPS)
+
+
+def sequence_balance_loss(
+    *,
+    scores: torch.Tensor,
+    routing_map: torch.Tensor,
+    topk: int,
+    coeff: float,
+    reduce_group: Optional["torch.distributed.ProcessGroup"] = None,
+) -> torch.Tensor:
+    """``alpha * mean_b sum_i f_i P_i`` over the sequences in the batch.
+
+    Args:
+        scores: ``[B, S, E]`` **normalised** affinities (see
+            :func:`normalise_affinities`); the differentiable input.
+        routing_map: ``[B, S, E]`` bool, the routing actually dispatched.
+        topk: experts per token.
+        coeff: ``alpha``.
+        reduce_group: group the sequence axis is sharded over (tensor +
+            context parallel when sequence parallelism is on), or ``None`` when
+            each rank holds whole sequences.
+
+    Returns:
+        Scalar loss in ``scores.dtype``. Under sharding this is the local
+        rank's share: the expert counts are reduced so ``f`` is global, while
+        ``P`` stays local so each rank differentiates only its own tokens --
+        the same split the framework's aux-loss path uses. The reported value
+        therefore needs the same group as its tracker reduction.
+    """
+    if scores.dim() != 3 or routing_map.shape != scores.shape:
+        raise ValueError(
+            f"scores and routing_map must both be [B, S, E]; got {tuple(scores.shape)} "
+            f"and {tuple(routing_map.shape)}"
+        )
+
+    _, seq_len, num_experts = scores.shape
+
+    counts = routing_map.sum(dim=1).to(scores.dtype)  # [B, E]
+    score_sums = scores.sum(dim=1)  # [B, E], differentiable
+
+    if reduce_group is not None and torch.distributed.get_world_size(reduce_group) > 1:
+        world_size = torch.distributed.get_world_size(reduce_group)
+        counts = counts.contiguous()
+        # No gradient flows through the counts (they come from a bool mask), so
+        # a plain all-reduce is enough.
+        torch.distributed.all_reduce(counts, group=reduce_group)
+        seq_len = seq_len * world_size
+
+    # f_i * P_i = (E / (K*T)) * C_i * (1/T) * S_i, summed over experts.
+    scale = num_experts / (float(topk) * float(seq_len) ** 2)
+    per_sequence = (counts * score_sums).sum(dim=-1) * scale  # [B]
+    return per_sequence.mean() * coeff
+
+
+def log_seq_balance_loss(
+    loss: torch.Tensor,
+    *,
+    layer_number: Optional[int],
+    num_layers: int,
+    reduce_group: Optional["torch.distributed.ProcessGroup"] = None,
+) -> None:
+    """Report the loss through the framework's MoE aux-loss tracker.
+
+    Uses the same group :func:`sequence_balance_loss` was given, so the logged
+    value is the global loss rather than one rank's share.
+    """
+    if not layer_number:
+        return
+    try:
+        from megatron.core.transformer.moe.moe_utils import save_to_aux_losses_tracker
+    except Exception:
+        return
+
+    save_to_aux_losses_tracker(
+        SEQ_BALANCE_LOSS_NAME,
+        loss.detach(),
+        layer_number,
+        num_layers,
+        reduce_group=reduce_group,
+    )
diff --git a/primus/backends/megatron/core/transformer/moe/v4_topk_router.py b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py
new file mode 100644
index 000000000..8e2893a4d
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py
@@ -0,0 +1,376 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Learned Top-K MoE router for DeepSeek-V4.
+
+Reference: techblog §4 ("MoE: routing scoring") and the inference
+reference at ``DeepSeek-V4-Flash/inference/model.py:Gate.forward``.
+
+For layers with ``layer_idx >= num_hash_layers`` V4 uses a learned
+top-K router. The HF-released ``Gate`` module computes a single
+``[D -> num_experts]`` linear and supports three score functions:
+
+* ``softmax`` — standard competitive normalization.
+* ``sigmoid`` — independent expert scoring (V3 fallback).
+* ``sqrtsoftplus`` — V4 default. ``sqrt(softplus(x))`` combines the
+  positive-only behavior of softplus with the sub-linear growth of
+  sqrt; sits between sigmoid (saturating) and softmax (competition)
+  and yields smoother routing gradients in long training runs.
+
+Optionally the router supports an **expert bias** correction
+(``moe_router_enable_expert_bias`` / "noaux_tc" — V3-style auxiliary-free
+balancing): a learnable per-expert bias is added to the score *only for
+top-K selection*, and the returned routing weights are gathered from the
+**un-biased** scores. This keeps gradient flow clean (probs flow back to
+the gate weight, not the bias) while still letting the bias term steer
+load balance.
+
+After top-K selection, the routing weights are renormalized to sum to 1
+**only when the score function is non-softmax** (matches HF; with
+softmax the sum is already 1 by construction). A final scalar
+``topk_scaling_factor`` ("route_scale" in the HF reference) is applied
+multiplicatively.
+
+Plan-2 P14 contract:
+
+* :class:`DeepseekV4LearnedRouter` — standalone ``nn.Module`` that
+  produces sparse ``(probs, routing_map)`` with the same ``[N, num_experts]``
+  shape contract as Megatron's :class:`TopKRouter`. The eager,
+  CPU-testable form is the canonical reference for G4 unit tests.
+* Parameter layout:
+   - ``weight``: ``nn.Parameter`` of shape ``[num_experts, hidden_size]``
+     (matches both Megatron's ``TopKRouter.weight`` and HF reference
+     ``Gate.weight``).
+   - ``expert_bias`` (optional): ``nn.Parameter`` of shape
+     ``[num_experts]`` (matches HF reference ``Gate.bias``).
+* ``score_function`` ∈ ``{"softmax", "sigmoid", "sqrtsoftplus"}``.
+
+Phase-2 of P14 will subclass Megatron's :class:`TopKRouter` directly so
+the router participates in aux-loss / z-loss / dispatcher lifecycle in
+production. The standalone form here remains the reference for unit
+tests and the state-dict adapter (P17).
+
+Back-compat alias ``V4TopKRouter`` is exposed but deprecated; new
+callers should use :class:`DeepseekV4LearnedRouter`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (
+    is_triton_path_enabled as _v4_router_triton_enabled,
+)
+from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (
+    v4_router_post_triton,
+)
+from primus.backends.megatron.core.transformer.moe.v4_seq_balance_loss import (
+    log_seq_balance_loss,
+    normalise_affinities,
+    sequence_balance_loss,
+)
+
+logger = logging.getLogger(__name__)
+
+_VALID_SCORE_FUNCTIONS = {"softmax", "sigmoid", "sqrtsoftplus"}
+
+
+def v4_score_fn(logits: torch.Tensor, *, score_function: str) -> torch.Tensor:
+    """Apply a V4-supported score function to gate logits.
+
+    Args:
+        logits: ``[..., num_experts]`` tensor of pre-score linear
+            outputs. Must be float (fp32 in the HF reference; we follow).
+        score_function: one of ``"softmax"``, ``"sigmoid"``,
+            ``"sqrtsoftplus"``.
+
+    Returns:
+        Tensor of the same shape, post score-function. ``softmax`` sums
+        to 1 along the expert axis; the other two are pointwise.
+    """
+    if score_function == "softmax":
+        return F.softmax(logits, dim=-1)
+    if score_function == "sigmoid":
+        return torch.sigmoid(logits)
+    if score_function == "sqrtsoftplus":
+        return F.softplus(logits).sqrt()
+    raise ValueError(
+        f"Unknown score_function: {score_function!r}. " f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}."
+    )
+
+
+def _compute_route(
+    *,
+    logits: torch.Tensor,
+    expert_bias: Optional[torch.Tensor],
+    score_function: str,
+    topk: int,
+    topk_scaling_factor: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Shared selection / renormalization core for V4 routers.
+
+    Returns sparse ``(probs[N, E], routing_map[N, E])``. The dense
+    ``(weights[N, K], indices[N, K])`` form follows the HF reference;
+    we expose only the sparse contract here so downstream Megatron
+    dispatchers consume it directly.
+
+    NOTE: this helper assumes ``logits`` is already shaped ``[N,
+    num_experts]`` and in fp32. Callers are responsible for the cast.
+    """
+    # Pre-compute topk indices on the host (heavy GPU compute that
+    # benefits from being its own kernel).  When PRIMUS_V4_ROUTER_TRITON
+    # is on (and supported), route the rest of the chain through one
+    # fused Triton kernel; otherwise fall back to the eager body.
+    scores_for_selection = v4_score_fn(logits, score_function=score_function)
+    if expert_bias is not None:
+        sel_score = scores_for_selection + expert_bias.to(scores_for_selection.dtype)
+    else:
+        sel_score = scores_for_selection
+    indices = sel_score.topk(topk, dim=-1).indices  # [N, K]
+
+    # The V4 router is GPU-only: both the Triton and eager paths run on
+    # CUDA/HIP tensors (there is no CPU compute path).
+    assert logits.is_cuda, "V4 router requires CUDA / HIP tensors"
+    if _v4_router_triton_enabled():
+        # The Triton kernel re-applies the score function inside (so
+        # it can save the full row for backward).  This keeps the
+        # eager path's bit-equivalent behaviour for the gathered
+        # weights at the cost of one extra pass over [N, E] of fp32
+        # logits -- negligible at V4-Flash widths.
+        probs, routing_map = v4_router_post_triton(
+            logits,
+            indices,
+            score_function=score_function,
+            topk_scaling_factor=topk_scaling_factor,
+            out_dtype=scores_for_selection.dtype,
+        )
+        return probs, routing_map
+
+    # Eager fallback (verbatim pre-P39 body).
+    original_scores = scores_for_selection
+    weights = original_scores.gather(1, indices)  # [N, K]
+
+    if score_function != "softmax":
+        denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12)
+        weights = weights / denom
+
+    if topk_scaling_factor != 1.0:
+        weights = weights * float(topk_scaling_factor)
+
+    num_experts = logits.shape[-1]
+    N = logits.shape[0]
+    device = logits.device
+
+    probs = torch.zeros(N, num_experts, dtype=weights.dtype, device=device)
+    probs.scatter_(1, indices, weights)
+
+    routing_map = torch.zeros(N, num_experts, dtype=torch.bool, device=device)
+    routing_map.scatter_(1, indices, True)
+
+    return probs, routing_map
+
+
+class DeepseekV4LearnedRouter(nn.Module):
+    """Learned top-K router for DeepSeek-V4 MoE layers (l >= num_hash_layers).
+
+    Args:
+        hidden_size: model dim ``D``; the gate is a single ``D -> num_experts``
+            linear.
+        num_experts: total number of routed experts.
+        topk: number of experts each token is routed to.
+        score_function: one of ``{"softmax", "sigmoid", "sqrtsoftplus"}``.
+            V4 default is ``"sqrtsoftplus"``.
+        enable_expert_bias: if True, allocate a learnable per-expert bias
+            used for selection only ("noaux_tc"). Probabilities are
+            re-read from the un-biased score so probs gradient flows
+            only into ``weight``, not ``expert_bias``.
+        topk_scaling_factor: scalar multiplier applied to the
+            renormalized probs (V3-style ``moe_router_topk_scaling_factor``,
+            HF reference ``Gate.route_scale``). Defaults to ``1.0``.
+        dtype: dtype of the gate weight; defaults to fp32 (matches HF
+            reference; the routing math runs in fp32 regardless).
+        seq_balance_loss_coeff: coefficient of the sequence-wise balance loss
+            (paper 2.1). ``0`` disables it. The caller is responsible for
+            checking ``moe_router_load_balancing_type``, so this single number
+            fully decides whether the loss runs.
+        seq_balance_reduce_group: group the sequence axis is sharded over, or
+            ``None`` when each rank holds whole sequences.
+        layer_number: 1-based layer index, and ``num_layers`` the total, both
+            only used to report the loss. Reporting is skipped without them.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_experts: int,
+        topk: int,
+        score_function: str = "sqrtsoftplus",
+        enable_expert_bias: bool = False,
+        topk_scaling_factor: float = 1.0,
+        dtype: Optional[torch.dtype] = None,
+        seq_balance_loss_coeff: float = 0.0,
+        seq_balance_reduce_group: Optional["torch.distributed.ProcessGroup"] = None,
+        layer_number: Optional[int] = None,
+        num_layers: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        if num_experts <= 0:
+            raise ValueError(f"num_experts must be > 0, got {num_experts}")
+        if topk <= 0 or topk > num_experts:
+            raise ValueError(f"topk must be in [1, {num_experts}], got {topk}")
+        if score_function not in _VALID_SCORE_FUNCTIONS:
+            raise ValueError(
+                f"Unknown score_function: {score_function!r}. "
+                f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}."
+            )
+
+        self.hidden_size = int(hidden_size)
+        self.num_experts = int(num_experts)
+        self.topk = int(topk)
+        self.score_function = str(score_function)
+        self.topk_scaling_factor = float(topk_scaling_factor)
+        self.seq_balance_loss_coeff = float(seq_balance_loss_coeff or 0.0)
+        # Held by reference so the group is not treated as module state.
+        self._seq_balance_reduce_group = seq_balance_reduce_group
+        self.layer_number = layer_number
+        self.num_layers = num_layers
+        # Last computed value (detached), for tests and ad-hoc inspection.
+        self.last_seq_balance_loss: Optional[torch.Tensor] = None
+
+        weight_dtype = dtype or torch.float32
+        # Gate weight: [num_experts, hidden_size] (matches Megatron TopKRouter
+        # AND the HF reference Gate.weight). State-dict key: ``weight``.
+        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, dtype=weight_dtype))
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+
+        if enable_expert_bias:
+            # Per-expert selection bias. State-dict key: ``expert_bias``.
+            self.expert_bias = nn.Parameter(torch.zeros(self.num_experts, dtype=weight_dtype))
+        else:
+            self.register_parameter("expert_bias", None)
+
+    # ------------------------------------------------------------------
+
+    @property
+    def seq_balance_loss_enabled(self) -> bool:
+        """Whether the sequence-wise balance loss is on for this router."""
+        return self.seq_balance_loss_coeff > 0.0
+
+    def _apply_seq_balance_loss(
+        self,
+        probs: torch.Tensor,
+        logits: torch.Tensor,
+        routing_map: torch.Tensor,
+        *,
+        batch_size: int,
+        seq_length: int,
+    ) -> torch.Tensor:
+        """Attach the sequence-wise balance loss to ``probs``.
+
+        ``probs`` is returned unchanged in value; the loss rides along in the
+        autograd graph via ``MoEAuxLossAutoScaler``, the same mechanism the
+        framework's own routers use, so the schedule's per-microbatch loss scale
+        applies automatically.
+        """
+        num_experts = logits.shape[-1]
+        # Unbiased, normalised affinities -- the ``s'_{i,t}`` of the paper.
+        affinities = normalise_affinities(v4_score_fn(logits, score_function=self.score_function))
+
+        loss = sequence_balance_loss(
+            scores=affinities.view(batch_size, seq_length, num_experts),
+            routing_map=routing_map.view(batch_size, seq_length, num_experts),
+            topk=self.topk,
+            coeff=self.seq_balance_loss_coeff,
+            reduce_group=self._seq_balance_reduce_group,
+        )
+        self.last_seq_balance_loss = loss.detach()
+
+        if self.num_layers:
+            log_seq_balance_loss(
+                loss,
+                layer_number=self.layer_number,
+                num_layers=int(self.num_layers),
+                reduce_group=self._seq_balance_reduce_group,
+            )
+
+        from megatron.core.transformer.moe.moe_utils import MoEAuxLossAutoScaler
+
+        return MoEAuxLossAutoScaler.apply(probs, loss)
+
+    def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Route ``hidden`` to top-K experts.
+
+        Args:
+            hidden: ``[B, S, D]`` (or any shape with last dim ``D``).
+
+        Returns:
+            probs: ``[N, num_experts]`` float tensor, ``N = numel/D``.
+                Non-selected experts have probability 0; selected experts
+                hold the (possibly renormalized + scaled) un-biased
+                score.
+            routing_map: ``[N, num_experts]`` bool tensor; ``True`` at
+                ``(n, e)`` iff token ``n`` is routed to expert ``e``.
+        """
+        flat = hidden.reshape(-1, self.hidden_size)
+        # Match HF reference: routing math runs in fp32 regardless of
+        # input dtype.
+        logits = F.linear(flat.to(torch.float32), self.weight.to(torch.float32))
+        probs, routing_map = _compute_route(
+            logits=logits,
+            expert_bias=self.expert_bias,
+            score_function=self.score_function,
+            topk=self.topk,
+            topk_scaling_factor=self.topk_scaling_factor,
+        )
+
+        # The balance loss is per-sequence, so it needs the [B, S] structure
+        # that ``flat`` just collapsed. V4 carries activations as [B, S, D]
+        # throughout (not the framework's usual [S, B, D]), so the sequence is
+        # axis -2.
+        if self.seq_balance_loss_enabled and self.training and torch.is_grad_enabled() and hidden.dim() >= 3:
+            seq_length = int(hidden.shape[-2])
+            num_tokens = flat.shape[0]
+            if seq_length > 0 and num_tokens % seq_length == 0:
+                probs = self._apply_seq_balance_loss(
+                    probs,
+                    logits,
+                    routing_map,
+                    batch_size=num_tokens // seq_length,
+                    seq_length=seq_length,
+                )
+            else:
+                # Grouping tokens into the wrong sequences would silently change
+                # what the loss measures, so skip rather than guess.
+                logger.warning(
+                    "[V4-router] skipping the sequence-wise balance loss: %d tokens do not "
+                    "divide into sequences of %d (input shape %s)",
+                    num_tokens,
+                    seq_length,
+                    tuple(hidden.shape),
+                )
+
+        return probs, routing_map
+
+
+# Back-compat alias. New callers should use ``DeepseekV4LearnedRouter``.
+V4TopKRouter = DeepseekV4LearnedRouter
+
+# Back-compat alias for the standalone score-function helper. The leading
+# underscore was dropped because the helper is part of the test surface.
+_score_fn = v4_score_fn
+
+__all__ = [
+    "DeepseekV4LearnedRouter",
+    "V4TopKRouter",
+    "v4_score_fn",
+    "_score_fn",
+]
diff --git a/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py b/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py
index f1155098e..589ce54f5 100644
--- a/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py
+++ b/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py
@@ -6,7 +6,7 @@
     PipelineParallelLayerLayout,
 )
 
-from primus.modules.trainer.megatron.utils import is_v_schedule_enabled
+from primus.backends.megatron.training.utils import is_v_schedule_enabled
 
 
 class PrimusPipelineParallelLayerLayout(PipelineParallelLayerLayout):
diff --git a/primus/backends/megatron/core/transformer/sliding_window_kv.py b/primus/backends/megatron/core/transformer/sliding_window_kv.py
new file mode 100644
index 000000000..68a336696
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/sliding_window_kv.py
@@ -0,0 +1,81 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Sliding-window mask helpers for DeepSeek-V4 dense / SWA attention layers.
+
+Reference: techblog §1 ("Hybrid Attention") — dense layers
+(``compress_ratio == 0``) attend only to the last ``attn_sliding_window``
+tokens (default 128) plus an optional ``attn_sink``. HCA layers also use
+SWA over the raw KV in addition to the compressed-KV pool.
+
+This module only generates the **mask** — the actual ``q @ k^T`` happens in
+the surrounding attention class. Returning a plain ``[Sq, Sk]`` additive
+mask keeps it composable with both eager and flash-style backends.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def sliding_window_causal_mask(
+    seq_len: int,
+    window: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return a ``[seq_len, seq_len]`` additive attention mask.
+
+    Position ``j`` is **allowed** for query ``i`` iff:
+
+    * ``j <= i`` (causal), and
+    * ``i - j < window`` (sliding window).
+
+    Disallowed positions get ``-inf``; allowed positions get ``0``.
+
+    A ``window`` of ``0`` or ``>= seq_len`` degenerates to the standard
+    causal mask.
+    """
+    q = torch.arange(seq_len, device=device).unsqueeze(1)
+    k = torch.arange(seq_len, device=device).unsqueeze(0)
+    dist = q - k
+    if window <= 0 or window >= seq_len:
+        allowed = dist >= 0
+    else:
+        allowed = (dist >= 0) & (dist < window)
+    return torch.where(allowed, 0.0, float("-inf")).to(dtype)
+
+
+def sliding_window_kv_indices(
+    seq_len: int,
+    window: int,
+    *,
+    device: torch.device,
+) -> torch.Tensor:
+    """For each query ``i`` return the ``window`` raw-KV indices it attends
+    to: ``[max(0, i-window+1), i]``.
+
+    Returned tensor has shape ``[seq_len, window]`` (long); positions before
+    the start of the sequence are filled with ``-1`` so the caller can drop
+    or zero-mask them.
+    """
+    if window <= 0:
+        # No sliding-window restriction: nothing to gather (caller should use full causal).
+        return torch.empty(seq_len, 0, dtype=torch.long, device=device)
+
+    i = torch.arange(seq_len, device=device).unsqueeze(1)  # [S, 1]
+    offset = torch.arange(window - 1, -1, -1, device=device).unsqueeze(0)  # [1, W] (descending)
+    j = i - offset  # [S, W]
+    j = torch.where(j >= 0, j, torch.full_like(j, -1))
+    return j
+
+
+__all__ = [
+    "sliding_window_causal_mask",
+    "sliding_window_kv_indices",
+]
diff --git a/primus/backends/megatron/core/transformer/transformer_layer.py b/primus/backends/megatron/core/transformer/transformer_layer.py
index e544845a8..8fe9f4f5c 100644
--- a/primus/backends/megatron/core/transformer/transformer_layer.py
+++ b/primus/backends/megatron/core/transformer/transformer_layer.py
@@ -13,7 +13,7 @@
 from megatron.core.transformer.enums import LayerType
 from megatron.core.transformer.transformer_config import TransformerConfig
 
-from primus.modules.trainer.megatron.utils import is_v_schedule_enabled
+from primus.backends.megatron.training.utils import is_v_schedule_enabled
 
 
 def get_transformer_layer_offset(
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md b/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md
new file mode 100644
index 000000000..68b6ef159
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md
@@ -0,0 +1,102 @@
+# DeepSeek-V4 attention kernels
+
+This package holds every attention backend used by `DeepseekV4Attention`
+(`primus/backends/megatron/core/transformer/deepseek_v4_attention.py`).
+`__init__.py` is the single entry point: it maps each backend to its functional
+entry, and the attention module imports everything from here.
+
+## Attention variants
+
+DeepSeek-V4 has three attention shapes, selected per layer by `compress_ratio`:
+
+| `compress_ratio` | Variant   | What it does                                                        |
+| ---------------- | --------- | ------------------------------------------------------------------- |
+| `0`              | dense/SWA | Causal + sliding-window attention over the local KV.                |
+| `128`            | HCA       | Hierarchical compressed attention: local KV concatenated with a compressed pool, joint softmax. |
+| `4`              | CSA       | Compressed sparse attention: per-query top-K gather from the compressed pool, joint softmax. |
+
+Two config selectors pick the kernel per group:
+
+- **`use_v4_attention_backend`** → dense (`cr=0`) and HCA (`cr=128`) layers.
+- **`use_v4_csa_attention_backend`** → CSA (`cr=4`) layers.
+
+Both default to `triton_v1`. `gluon` is a gfx950/CDNA4-only opt-in: it is imported
+lazily and only when selected, and selecting it asserts the device is gfx950.
+
+## Backends
+
+| Selector value | Applies to     | Folder                   | Entry point(s)                              | Notes                                                                 |
+| -------------- | -------------- | ------------------------ | ------------------------------------------- | --------------------------------------------------------------------- |
+| `eager`        | dense/HCA, CSA | `_eager/`                | `eager_v4_attention`, `eager_v4_csa_attention` | Pure-PyTorch reference. Bit-identical baseline shared with unit tests; slow, used for correctness. |
+| `triton_v0`    | CSA only       | `_triton_v0_deprecated/` | `v4_csa_attention_v0`                        | **Deprecated.** Gathered per-query CSA with scalar GEMV (~30–260× slower). Not for production. |
+| `triton_v1`    | dense/HCA, CSA | `_triton_v1/`            | `v4_attention_v1`, `v4_csa_attention_v1`     | **Production default.** Separate K/V, pool-based CSA + dense/HCA Triton kernels. |
+| `triton_v2`    | dense/HCA, CSA | `_triton_v2/`            | `v4_attention_v2`, `v4_csa_attention_v2`     | Fused single-latent sparse-MLA (K=V) using plain Triton `tl.dot` / MFMA. |
+| `gluon`        | dense/HCA, CSA | `_gluon_dsa/`            | `v4_attention_gluon`, `v4_csa_attention_gluon` | Hand-tuned fused single-latent sparse-MLA for **gfx950 (CDNA4) only**. Lazily imported; selecting it asserts the arch. |
+| `flydsl_v0`    | CSA only       | `_flydsl_v0_deprecated/` | routed via `v4_csa_attention_v0` (`use_flydsl=True`) | **Deprecated**, forward-only legacy FlyDSL scalar CSA.               |
+
+Support folders (not directly selectable):
+
+- `_triton_common/` — shared Triton helpers (indexer, compressor, sinkhorn, RoPE, HC).
+- `_flydsl_v1/` — WIP native FlyDSL MFMA sparse-MLA backend (forward kernel not yet implemented).
+- `_tilelang/` — experimental TileLang path (not wired into the selectors).
+- `v4_sparse_mla_adapter.py` — kernel-agnostic adapter mapping V4 tensors to the fused sparse-MLA interface (used by `triton_v2` / `gluon`).
+
+Valid selector values (enforced in `DeepseekV4Attention.__init__`):
+
+- dense/HCA: `eager | triton_v1 | triton_v2 | gluon`
+- CSA: `eager | triton_v0 | triton_v1 | triton_v2 | gluon | flydsl_v0`
+
+> Note: when a Turbo `core_attention` module is built, `use_turbo_attention`
+> still takes precedence for the dense (`cr=0`) path.
+
+## How to enable each backend
+
+Set the two selectors to any valid value. All three mechanisms below set the
+same config fields.
+
+### 1. Config YAML
+
+```yaml
+# primus/configs/models/megatron/deepseek_v4_*.yaml
+use_v4_attention_backend: triton_v1      # dense (cr=0) + HCA (cr=128)
+use_v4_csa_attention_backend: triton_v1  # CSA (cr=4)
+```
+
+### 2. CLI flags
+
+```bash
+--use_v4_attention_backend triton_v2 \
+--use_v4_csa_attention_backend gluon
+```
+
+### 3. Environment variables (root run scripts)
+
+The `run_deepseek_v4*.sh` scripts read these and forward them as CLI flags:
+
+```bash
+export USE_V4_ATTENTION_BACKEND=triton_v1       # dense + HCA
+export USE_V4_CSA_ATTENTION_BACKEND=triton_v1   # CSA
+```
+
+### Examples
+
+```bash
+# Default: production Triton v1 (separate K/V) everywhere
+export USE_V4_ATTENTION_BACKEND=triton_v1
+export USE_V4_CSA_ATTENTION_BACKEND=triton_v1
+
+# gfx950-only hand-tuned gluon backend (asserts arch when selected)
+export USE_V4_ATTENTION_BACKEND=gluon
+export USE_V4_CSA_ATTENTION_BACKEND=gluon
+
+# Fused single-latent sparse-MLA (Triton v2) for all groups
+export USE_V4_ATTENTION_BACKEND=triton_v2
+export USE_V4_CSA_ATTENTION_BACKEND=triton_v2
+
+# Eager reference (correctness / debugging)
+export USE_V4_ATTENTION_BACKEND=eager
+export USE_V4_CSA_ATTENTION_BACKEND=eager
+```
+
+The two selectors are independent, so you can mix backends per group, e.g.
+`USE_V4_ATTENTION_BACKEND=triton_v1` with `USE_V4_CSA_ATTENTION_BACKEND=gluon`.
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py
new file mode 100644
index 000000000..018e57a48
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py
@@ -0,0 +1,219 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 attention kernels — single entry point for every backend.
+
+``DeepseekV4Attention`` imports all attention entries from here, so this module
+is the one place that maps a backend to its functional entry. Naming:
+
+* dense (cr=0) / HCA (cr=128) entry: ``v4_attention_<backend>``
+* CSA (cr=4) entry:                  ``v4_csa_attention_<backend>``
+
+Backends:
+
+* ``eager``  — pure-Python reference (:mod:`_eager`): ``eager_v4_attention`` /
+  ``eager_v4_csa_attention``.
+* ``v0``     — Triton, DEPRECATED gathered CSA (:mod:`_triton_v0_deprecated`):
+  ``v4_csa_attention_v0`` (cr=4 only, ~30-260x slower; not used in production).
+* ``v1``     — Triton production, separate K/V (:mod:`_triton_v1`):
+  ``v4_attention_v1`` (dense/HCA) / ``v4_csa_attention_v1`` (pool CSA).
+* ``v2``     — Triton fused single-latent sparse-MLA (:mod:`_triton_v2`,
+  ``tl.dot`` / MFMA): ``v4_attention_v2`` / ``v4_csa_attention_v2``.
+* ``gluon``  — hand-tuned gfx950 fused single-latent sparse-MLA (:mod:`_gluon_dsa`):
+  loaded LAZILY via :func:`load_gluon_attention_backends` (NOT imported eagerly).
+
+``gluon`` hard-depends on ``triton.experimental.gluon`` (gfx950 / CDNA4 only), so
+importing it here unconditionally would make *any* ``import ...v4_attention_kernels``
+fail on a Triton build without gluon — even when the caller selected
+``eager`` / ``triton_v1`` / ``triton_v2``. It is therefore imported on demand only
+when a layer actually selects the ``gluon`` backend (see
+:func:`load_gluon_attention_backends`).
+
+The eager references share exactly one definition with the kernels + unit tests
+and keep the checkpoint-reproduction baseline bit-identical at the call sites.
+"""
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager import (
+    eager_v4_attention,
+    eager_v4_csa_attention,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated import (
+    V4CSAAttentionFn,
+    v4_csa_attention_v0,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1 import (
+    V4AttentionFn,
+    V4CSAPoolAttentionFn,
+    v4_attention_v1,
+    v4_csa_attention_v1,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_triton import (
+    v4_attention_v2,
+    v4_csa_attention_v2,
+)
+
+
+def load_gluon_attention_backends():
+    """Lazily import the gluon sparse-MLA attention entries.
+
+    The gluon backend (:mod:`_gluon_dsa`) hard-depends on
+    ``triton.experimental.gluon`` (gfx950 / CDNA4 only). This helper defers that
+    import so selecting any other backend (``eager`` / ``triton_v1`` /
+    ``triton_v2``) never pays it — and never crashes on a Triton build / GPU arch
+    without gluon support. Call it only when a layer actually selects ``gluon``.
+
+    Returns ``(v4_attention_gluon, v4_csa_attention_gluon)``. Raises
+    :class:`ImportError` with an actionable message when the gluon dependency is
+    unavailable.
+
+    NOTE: the import is intentionally inline (optional, hardware-specific
+    dependency); it must not be hoisted to module scope.
+    """
+    try:
+        from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon import (
+            v4_attention_gluon,
+            v4_csa_attention_gluon,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' requires "
+            "the gluon sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), "
+            f"which failed to import: {exc}. Select a different backend "
+            "(eager | triton_v1 | triton_v2), or run on a gfx950 build with Triton gluon support."
+        ) from exc
+    return v4_attention_gluon, v4_csa_attention_gluon
+
+
+def load_gluon_v2_attention_backends():
+    """Lazily import the gluon_v2 sparse-MLA attention entries (:mod:`_gluon_v2`).
+
+    Second-generation Gluon backend (gfx950 / CDNA4): Gluon forward (rope-skip + exp2 +
+    MFMA K=32) + Gluon backward (rope-skip + K=32 + single-chunk RMW). Same lazy-import
+    rationale as :func:`load_gluon_attention_backends` (hard ``triton.experimental.gluon``
+    dependency). Returns ``(v4_attention_gluon_v2, v4_csa_attention_gluon_v2)``.
+    """
+    try:
+        from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v2 import (
+            v4_attention_gluon_v2,
+            v4_csa_attention_gluon_v2,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon_v2' requires "
+            "the gluon_v2 sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), "
+            f"which failed to import: {exc}. Select a different backend "
+            "(eager | triton_v1 | triton_v2 | gluon), or run on a gfx950 build with Triton gluon support."
+        ) from exc
+    return v4_attention_gluon_v2, v4_csa_attention_gluon_v2
+
+
+def load_gluon_v3_attention_backends():
+    """Lazily import the gluon_v3 sparse-MLA attention entries (:mod:`_gluon_v3`).
+
+    Optimized 3rd-gen Gluon backend (gfx950 / CDNA4): Round-9 CSA formula-pack +
+    aiter Gluon LSE fwd route, gluon_v2/Round-2 bwd chunking. Same lazy-import
+    rationale as :func:`load_gluon_attention_backends`. Returns
+    ``(v4_attention_gluon_v3, v4_csa_attention_gluon_v3)``.
+    """
+    try:
+        from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v3 import (
+            v4_attention_gluon_v3,
+            v4_csa_attention_gluon_v3,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon_v3' requires "
+            "the gluon_v3 sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), "
+            f"which failed to import: {exc}. Select a different backend "
+            "(eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1 | turbo), "
+            "or run on a gfx950 build with Triton gluon support."
+        ) from exc
+    return v4_attention_gluon_v3, v4_csa_attention_gluon_v3
+
+
+def load_flydsl_attention_backends():
+    """Lazily import the native-FlyDSL sparse-MLA attention entries.
+
+    The flydsl_v1 backend (:mod:`_flydsl_v1`) hard-depends on the installed
+    ``flydsl`` pip package (gfx950 / CDNA4). This helper defers that import so
+    selecting any other backend never pays it — and never crashes on a build /
+    GPU arch without flydsl. Call it only when a layer actually selects
+    ``flydsl_v1``.
+
+    Returns ``(v4_attention_flydsl, v4_csa_attention_flydsl)``. Raises
+    :class:`ImportError` with an actionable message when flydsl is unavailable.
+    """
+    try:
+        from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_flydsl import (
+            v4_attention_flydsl,
+            v4_csa_attention_flydsl,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'flydsl_v1' requires "
+            "the native FlyDSL sparse-MLA backend (the `flydsl` pip package, gfx950 / CDNA4), "
+            f"which failed to import: {exc}. Select a different backend "
+            "(eager | triton_v1 | triton_v2 | gluon), or install flydsl on a gfx950 build."
+        ) from exc
+    return v4_attention_flydsl, v4_csa_attention_flydsl
+
+
+def load_turbo_attention_backends():
+    """Lazily import the Primus-Turbo native-FlyDSL sparse-MLA attention entries.
+
+    The ``turbo`` backend (:mod:`_turbo_flydsl`) binds to the installed
+    ``primus_turbo`` flydsl sparse-MLA v2 kernels (the "turbo API" integration),
+    which hard-depend on the installed ``primus_turbo`` (with the flydsl attention
+    submodule) and the ``flydsl`` pip package (gfx950 / CDNA4). Deferred so
+    selecting any other backend never pays that import — and never crashes on a
+    build / GPU arch without it. Call it only when a layer selects ``turbo``.
+
+    Returns ``(v4_attention_turbo, v4_csa_attention_turbo)``. Raises
+    :class:`ImportError` with an actionable message when the dependency is missing.
+    """
+    try:
+        from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_turbo_flydsl import (
+            v4_attention_turbo,
+            v4_csa_attention_turbo,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "use_v4_attention_backend / use_v4_csa_attention_backend = 'turbo' requires the "
+            "Primus-Turbo native-FlyDSL sparse-MLA backend (the installed `primus_turbo` with its "
+            "flydsl sparse-MLA attention, plus the `flydsl` pip package, gfx950 / CDNA4), "
+            f"which failed to import: {exc}. Select a different backend "
+            "(eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1), or install a "
+            "primus_turbo build carrying primus_turbo.flydsl.attention on a gfx950 build."
+        ) from exc
+    return v4_attention_turbo, v4_csa_attention_turbo
+
+
+__all__ = [
+    # eager reference
+    "eager_v4_attention",
+    "eager_v4_csa_attention",
+    # triton v0 (deprecated gathered CSA)
+    "v4_csa_attention_v0",
+    "V4CSAAttentionFn",
+    # triton v1 (production, separate K/V)
+    "v4_attention_v1",
+    "v4_csa_attention_v1",
+    "V4AttentionFn",
+    "V4CSAPoolAttentionFn",
+    # triton v2 (fused single-latent sparse-MLA)
+    "v4_attention_v2",
+    "v4_csa_attention_v2",
+    # gluon (fused single-latent sparse-MLA, gfx950) — lazily loaded
+    "load_gluon_attention_backends",
+    # gluon_v2 (2nd-gen gluon fwd+bwd, gfx950) — lazily loaded
+    "load_gluon_v2_attention_backends",
+    # gluon_v3 (3rd-gen optimized gluon fwd+bwd, gfx950) — lazily loaded
+    "load_gluon_v3_attention_backends",
+    # flydsl_v1 (native FlyDSL fused single-latent sparse-MLA, gfx950) — lazily loaded
+    "load_flydsl_attention_backends",
+    # turbo (Primus-Turbo native-FlyDSL sparse-MLA via the turbo API, gfx950) — lazily loaded
+    "load_turbo_attention_backends",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py
new file mode 100644
index 000000000..e72167d5d
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py
@@ -0,0 +1,18 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Eager-Python reference ops for DeepSeek-V4 attention.
+
+The single source of "eager truth" shared by ``DeepseekV4Attention``, every
+kernel backend (triton v0/v1/v2, gluon, flydsl_v2, tilelang) and the unit tests:
+
+* :func:`eager_v4_attention`     — dense (cr=0) / HCA (cr=128)
+* :func:`eager_v4_csa_attention` — CSA (cr=4)
+"""
+
+from .reference import eager_v4_attention, eager_v4_csa_attention
+
+__all__ = ["eager_v4_attention", "eager_v4_csa_attention"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py
new file mode 100644
index 000000000..6fa6598d4
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py
@@ -0,0 +1,329 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Eager-Python references for the V4 attention kernels (plan-4 P24).
+
+The math here is **bit-identical** to what previously lived inline in
+:meth:`DeepseekV4Attention._attention_forward` and
+:meth:`DeepseekV4Attention._csa_forward`. Plan-4 extracts it into pure
+functions so that:
+
+* ``DeepseekV4Attention`` itself, the plan-4 Triton kernels (P25 / P26),
+  and the plan-4 unit-test harness share exactly one definition;
+* the function signatures match the kernel signatures
+  (``v4_attention_v1`` / ``v4_csa_attention_v0``) so the test harness can
+  plug reference ↔ candidate interchangeably;
+* the compute kernel of V4 attention is decoupled from the
+  ``DeepseekV4Attention`` class (no ``self``-bound state) — the
+  per-call inputs are explicit.
+
+The two functions:
+
+* :func:`eager_v4_attention` — single-key-axis attention with optional
+  per-head learned softmax sink, optional sliding window, and optional
+  ``[Sq, Sk]`` additive bias. Covers ``compress_ratio == 0`` (dense +
+  SWA + sink, no bias) and ``compress_ratio == 128`` (HCA — caller
+  pre-concatenates the compressed pool to the local keys and supplies
+  the joint-softmax additive bias).
+* :func:`eager_v4_csa_attention` — fused local-SWA + per-query top-K
+  sparse attention with shared per-head sink and joint softmax.
+  Covers ``compress_ratio == 4`` (CSA). The caller is responsible for
+  the per-query top-K gather; the function takes the gathered
+  ``[B, Sq, K, head_dim]`` tensor directly.
+
+Both functions:
+
+* keep every matmul / einsum on tensor cores in the input dtype (bf16
+  in production); the matmul accumulator inside is fp32. The *only*
+  fp32 step is the softmax block (max-subtract + ``exp`` + sum +
+  divide), which :func:`_softmax_with_sink` upcasts internally and
+  returns in fp32; the caller-side ``probs.to(v.dtype)`` puts probs
+  back on the bf16 path before the V-matmul. This matches the
+  FlashAttention / Megatron de-facto layout (bf16 ``Q @ K^T``, fp32
+  softmax, bf16 ``probs @ V``);
+* honor ``attn_dropout`` only when ``training`` is also ``True`` (so
+  the eval / inference path is deterministic);
+* return ``[B, H, Sq, head_dim]`` in ``v.dtype`` (or ``v_local.dtype``
+  for the CSA path);
+* are autograd-friendly: fwd produces a graph; ``out.sum().backward()``
+  populates ``q.grad / k.grad / v.grad / sink.grad`` (and
+  ``gathered.grad`` for CSA).
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (
+    sliding_window_causal_mask,
+)
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _softmax_with_sink(
+    logits: torch.Tensor,
+    sink: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Numerically stable softmax with optional per-head learned sink column.
+
+    ``logits`` shape: ``[B, H, ..., Sk]`` — the head axis is at ``dim=1``.
+    When ``sink`` is given (shape ``[H]``) it is joined as a virtual key
+    column with notional value zero (so the V-weighted sum after
+    multiplying by ``v`` is unaffected by the sink slot), then dropped
+    after softmax. The head can still spend probability mass on the
+    sink as a "no attention" fallback.
+
+    **dtype contract.** This is the *only* fp32 step in V4 attention.
+    ``logits`` may arrive in the model dtype (bf16 in production —
+    coming straight out of a tensor-core ``Q @ K^T`` matmul); the
+    softmax block (max-subtract + ``exp`` + sum + divide) is the
+    numerically sensitive part, so we **upcast to fp32 here** and
+    return probabilities in **fp32**. The caller is responsible for
+    ``probs.to(v.dtype)`` before the value matmul so the V-matmul
+    stays on bf16 tensor cores.
+
+    Returns probabilities on the *real* keys of the same shape as
+    ``logits`` — but in **fp32** regardless of input dtype.
+    """
+    # logits: [B, H, ..., Sk] (any dtype) -> logits_fp32: [B, H, ..., Sk] in fp32
+    logits_fp32 = logits.float()
+    if sink is None:
+        logits_fp32 = logits_fp32 - logits_fp32.amax(dim=-1, keepdim=True).detach()
+        return logits_fp32.softmax(dim=-1)
+
+    # sink: [H] (any dtype) -> sink_col: [B, H, ..., 1] in fp32
+    ndim = logits_fp32.dim()
+    num_heads = sink.shape[0]
+    view_shape = [1] * ndim
+    view_shape[1] = num_heads
+    view_shape[-1] = 1
+    target_shape = list(logits_fp32.shape[:-1]) + [1]
+    sink_col = sink.float().view(*view_shape).expand(*target_shape)
+    # logits_fp32: [B, H, ..., Sk], sink_col: [B, H, ..., 1]
+    # -> logits_aug: [B, H, ..., Sk+1] in fp32
+    logits_aug = torch.cat([logits_fp32, sink_col], dim=-1)
+    logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach()
+    probs = logits_aug.softmax(dim=-1)
+    return probs[..., :-1]
+
+
+def _build_local_attention_mask(
+    seq_len: int,
+    swa_window: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Build the ``[seq_len, seq_len]`` local SWA-causal additive mask.
+
+    Equivalent to :meth:`DeepseekV4Attention._local_mask`:
+
+    * ``swa_window > 0`` — sliding-window causal (queries see the last
+      ``swa_window`` keys);
+    * otherwise — full causal (queries see all earlier keys).
+
+    The two-call contract is deterministic: same ``(seq_len,
+    swa_window, device, dtype)`` always returns the same tensor, so
+    rebuilding inside this function vs. passing a pre-built mask gives
+    bit-identical attention output.
+    """
+    window = swa_window if swa_window > 0 else seq_len
+    return sliding_window_causal_mask(seq_len, window, device=device, dtype=dtype)
+
+
+# ---------------------------------------------------------------------------
+# Public reference ops
+# ---------------------------------------------------------------------------
+
+
+def eager_v4_attention(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k: torch.Tensor,  # [B, H, Sk, D]
+    v: torch.Tensor,  # [B, H, Sk, D]
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],  # [Sq, Sk] or None
+    attn_dropout: float,
+    training: bool,
+    scale: float,
+) -> torch.Tensor:
+    """Eager-Python V4 dense / HCA attention.
+
+    Math::
+
+        logits = (q @ k.T) * scale + mask        # bf16 tensor-core matmul
+        probs  = softmax_with_sink(logits, sink) # softmax internally upcasts to fp32
+        if attn_dropout > 0 and training: probs = dropout(probs, attn_dropout)
+        out    = probs.to(v.dtype) @ v           # bf16 tensor-core matmul
+
+    **dtype contract.** All matmuls (``Q @ K^T`` and ``probs @ V``) run
+    on tensor cores in the input dtype (bf16 in production); the
+    accumulator inside the matmul is fp32. The *only* fp32 step is the
+    softmax block, which :func:`_softmax_with_sink` handles internally
+    (input may be bf16, output is fp32). The caller-side
+    ``probs.to(v.dtype)`` puts probs back on the bf16 path before the
+    V-matmul.
+
+    Mask resolution:
+
+    * ``additive_mask is not None`` — used directly. ``swa_window`` is
+      ignored. (HCA's ``compress_ratio == 128`` caller pre-builds the
+      joint-softmax mask via ``cat([local_mask, hca_mask])`` and
+      passes it here.)
+    * ``additive_mask is None`` and ``swa_window > 0`` — sliding-window
+      causal mask is built internally. Requires ``Sq == Sk``.
+    * ``additive_mask is None`` and ``swa_window <= 0`` — full causal
+      mask is built internally. Requires ``Sq == Sk``.
+
+    Returns ``[B, H, Sq, D]`` in ``v.dtype``.
+    """
+    # q: [B, H, Sq, D], k: [B, H, Sk, D] -> Sq: scalar, Sk: scalar
+    Sq = q.shape[2]
+    Sk = k.shape[2]
+
+    # mask: [Sq, Sk] in q.dtype   (broadcasts over B, H at the addition site below)
+    if additive_mask is None:
+        if Sq != Sk:
+            raise ValueError(
+                "eager_v4_attention requires `additive_mask` when Sq != Sk; "
+                f"got Sq={Sq}, Sk={Sk}. The HCA caller must pre-concatenate "
+                "the compressed pool to the local keys and supply the joint "
+                "additive mask."
+            )
+        # _build_local_attention_mask(Sq, swa_window) -> mask: [Sq, Sq] (== [Sq, Sk]) in q.dtype
+        mask = _build_local_attention_mask(Sq, swa_window, device=q.device, dtype=q.dtype)
+    else:
+        # additive_mask: [Sq, Sk] -> mask: [Sq, Sk] in caller's dtype
+        mask = additive_mask
+
+    # q: [B, H, Sq, D], k.transpose(-2,-1): [B, H, D, Sk] -> matmul(...): [B, H, Sq, Sk] in q.dtype
+    # bf16 tensor-core matmul (fp32 accumulator inside, output bf16)
+    # * scale: [B, H, Sq, Sk] in q.dtype
+    logits = torch.matmul(q, k.transpose(-2, -1)) * scale
+    # logits: [B, H, Sq, Sk], mask: [Sq, Sk] -> logits: [B, H, Sq, Sk]   (mask broadcasts over B, H)
+    logits = logits + mask
+    # logits: [B, H, Sq, Sk] (bf16), sink: [H] or None -> probs: [B, H, Sq, Sk] in fp32
+    # softmax block internally upcasts logits + sink to fp32 (numerical contract: only step in fp32)
+    probs = _softmax_with_sink(logits, sink)
+    # probs: [B, H, Sq, Sk] (fp32) -> probs: [B, H, Sq, Sk] (fp32)
+    if attn_dropout > 0.0 and training:
+        probs = torch.nn.functional.dropout(probs, p=attn_dropout)
+    # probs.to(v.dtype): [B, H, Sq, Sk] in v.dtype, v: [B, H, Sk, D] -> out: [B, H, Sq, D] in v.dtype
+    # bf16 tensor-core matmul (fp32 accumulator inside)
+    return torch.matmul(probs.to(v.dtype), v)
+
+
+def eager_v4_csa_attention(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    gathered: torch.Tensor,  # [B, Sq, K, D] — pre-gathered per-query top-K from compressed pool
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    sparse_mask: torch.Tensor,  # [B, Sq, K] additive (broadcasts over H)
+    attn_dropout: float,
+    training: bool,
+    scale: float,
+) -> torch.Tensor:
+    """Eager-Python V4 CSA fused attention (joint local SWA + sparse top-K).
+
+    Math::
+
+        local_logits  = (q @ k_local.T) * scale + local_mask                # bf16 matmul
+        sparse_logits = einsum("bhsd,bhskd->bhsk", q, gathered_h) * scale
+                       + sparse_mask.unsqueeze(1)                            # bf16 einsum
+        joint_logits  = cat([local_logits, sparse_logits], dim=-1)           # [B, H, Sq, Sq+K]
+        probs         = softmax_with_sink(joint_logits, sink)                # JOINT softmax in fp32
+        if attn_dropout > 0 and training: probs = dropout(probs, attn_dropout)
+        probs_local, probs_sparse = probs[..., :Sq], probs[..., Sq:]
+        out = probs_local @ v_local + einsum("bhsk,bhskd->bhsd", probs_sparse, gathered_h)
+                                                                             # bf16 matmul / einsum
+
+    where ``gathered_h = gathered.unsqueeze(1).expand(B, H, Sq, K, D)``
+    (heads broadcast across the single compressor output).
+
+    **dtype contract.** Same as :func:`eager_v4_attention`: every
+    matmul / einsum runs on tensor cores in the input dtype (bf16 in
+    production, fp32 accumulator inside); the *only* fp32 step is the
+    joint softmax inside :func:`_softmax_with_sink`. ``probs.to(v.dtype)``
+    puts the per-branch probabilities back on bf16 before the V-stage.
+
+    The local SWA mask is built internally from ``swa_window`` (see
+    :func:`_build_local_attention_mask`); the caller pre-builds
+    ``sparse_mask`` to flag indexer-dropped slots (``-inf`` for
+    ``topk_idx == -1``).
+
+    Returns ``[B, H, Sq, D]`` in ``v_local.dtype``.
+    """
+    # q: [B, H, Sq, D], gathered: [B, Sq, K, D] -> B, H, Sq, D, K: scalars
+    B, H, Sq, D = q.shape
+    K = gathered.shape[2]
+
+    # _build_local_attention_mask(Sq, swa_window) -> local_mask: [Sq, Sq] in q.dtype
+    local_mask = _build_local_attention_mask(Sq, swa_window, device=q.device, dtype=q.dtype)
+
+    # q: [B, H, Sq, D], k_local.transpose(-2,-1): [B, H, D, Sq]
+    # -> matmul(...): [B, H, Sq, Sq] in q.dtype   (bf16 tensor core, fp32 accumulator inside)
+    # * scale: local_logits: [B, H, Sq, Sq] in q.dtype
+    local_logits = torch.matmul(q, k_local.transpose(-2, -1)) * scale
+    # local_logits: [B, H, Sq, Sq], local_mask: [Sq, Sq]
+    # -> local_logits: [B, H, Sq, Sq] in q.dtype   (mask broadcasts over B, H)
+    local_logits = local_logits + local_mask
+
+    # gathered: [B, Sq, K, D] -> unsqueeze(1): [B, 1, Sq, K, D]
+    # -> expand: gathered_h: [B, H, Sq, K, D] in gathered.dtype   (view, no copy)
+    gathered_h = gathered.unsqueeze(1).expand(B, H, Sq, K, D)
+    # q: [B, H, Sq, D], gathered_h: [B, H, Sq, K, D]
+    # -> einsum("bhsd,bhskd->bhsk"): [B, H, Sq, K] in q.dtype   (bf16 tensor core)
+    # * scale: sparse_logits: [B, H, Sq, K] in q.dtype
+    sparse_logits = torch.einsum("bhsd,bhskd->bhsk", q, gathered_h) * scale
+    # sparse_logits: [B, H, Sq, K], sparse_mask.unsqueeze(1): [B, 1, Sq, K]
+    # -> sparse_logits: [B, H, Sq, K] in q.dtype   (mask broadcasts over H)
+    sparse_logits = sparse_logits + sparse_mask.unsqueeze(1)
+
+    # local_logits: [B, H, Sq, Sq], sparse_logits: [B, H, Sq, K]
+    # -> joint_logits: [B, H, Sq, Sq+K] in q.dtype
+    joint_logits = torch.cat([local_logits, sparse_logits], dim=-1)
+    # joint_logits: [B, H, Sq, Sq+K] (bf16), sink: [H] or None -> probs: [B, H, Sq, Sq+K] in fp32
+    # softmax block internally upcasts to fp32 (numerical contract: only step in fp32)
+    probs = _softmax_with_sink(joint_logits, sink)
+
+    # probs: [B, H, Sq, Sq+K] (fp32) -> probs: [B, H, Sq, Sq+K] (fp32)
+    if attn_dropout > 0.0 and training:
+        probs = torch.nn.functional.dropout(probs, p=attn_dropout)
+
+    # probs[..., :Sq]: [B, H, Sq, Sq] in fp32 -> probs_local: [B, H, Sq, Sq] in v_local.dtype
+    probs_local = probs[..., :Sq].to(v_local.dtype)
+    # probs[..., Sq:]: [B, H, Sq, K] in fp32 -> probs_sparse: [B, H, Sq, K] in v_local.dtype
+    probs_sparse = probs[..., Sq:].to(v_local.dtype)
+
+    # probs_local: [B, H, Sq, Sq], v_local: [B, H, Sq, D]
+    # -> matmul(...): out_local: [B, H, Sq, D] in v_local.dtype   (bf16 tensor core)
+    out_local = torch.matmul(probs_local, v_local)
+    # probs_sparse: [B, H, Sq, K], gathered_h.to(v_local.dtype): [B, H, Sq, K, D]
+    # -> einsum("bhsk,bhskd->bhsd"): out_sparse: [B, H, Sq, D] in v_local.dtype   (bf16 tensor core)
+    out_sparse = torch.einsum(
+        "bhsk,bhskd->bhsd",
+        probs_sparse,
+        gathered_h.to(v_local.dtype),
+    )
+
+    # out_local: [B, H, Sq, D], out_sparse: [B, H, Sq, D] -> out: [B, H, Sq, D] in v_local.dtype
+    return out_local + out_sparse
+
+
+__all__ = [
+    "eager_v4_attention",
+    "eager_v4_csa_attention",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py
new file mode 100644
index 000000000..5d2445b7c
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FlyDSL attention-kernel backend for DeepSeek-V4 (gfx950 / MI355X).
+
+A *soft-dependency* alternate backend, structured exactly like the sibling
+``_tilelang`` package: the V4 attention layer calls :func:`should_dispatch`
+with the ``enabled`` flag from a config knob, and this module lazily imports
+the FlyDSL kernel wrappers. If the ``flydsl`` runtime (or a kernel submodule)
+is not importable, :func:`should_dispatch` returns ``False`` and the caller
+transparently falls back to the in-tree Triton path -- so a container without
+FlyDSL never breaks and never changes behaviour.
+
+Kernels live under ``_flydsl/kernels/``. Forward-only for now (inference/eval);
+training autograd is a follow-up.
+"""
+from __future__ import annotations
+
+import os
+import warnings
+from typing import Any, Optional, Set
+
+__all__ = [
+    "should_dispatch",
+    "is_flydsl_available",
+    "v4_attention_fwd_flydsl",
+    "v4_csa_attention_fwd_flydsl",
+]
+
+# Kernel names the layer may ask for, mirroring the _tilelang registry.
+_KNOWN_KERNEL_NAMES: Set[str] = {
+    "v4_attention_fwd",  # SWA / HCA forward (MQA launcher)
+    "v4_csa_attention_fwd",  # CSA forward
+    "v4_attention_bwd",  # present in-package; training-autograd hook TBD
+    "v4_csa_attention_bwd",  # present in-package; training-autograd hook TBD
+}
+
+# Forward kernels that are actually wired (have a working adapter below).
+_WIRED_FWD: Set[str] = {"v4_attention_fwd", "v4_csa_attention_fwd"}
+
+_PROBE_DONE: bool = False
+_FLYDSL_AVAILABLE: bool = False
+
+
+def _probe_flydsl() -> bool:
+    """Import the ``flydsl`` runtime once (cached); warn once on rank 0 if absent."""
+    global _PROBE_DONE, _FLYDSL_AVAILABLE
+    if _PROBE_DONE:
+        return _FLYDSL_AVAILABLE
+    _PROBE_DONE = True
+    try:
+        # The kernel wrappers add the FlyDSL build dir to sys.path at import;
+        # allow an override for non-default install locations.
+        src = os.environ.get("PRIMUS_V4_FLYDSL_SRC", "/workspace/FlyDSL-amd")
+        import sys
+
+        if src and src not in sys.path and os.path.isdir(src):
+            sys.path.insert(0, src)
+        import flydsl  # noqa: F401
+
+        _FLYDSL_AVAILABLE = True
+    except Exception as exc:  # ImportError, or a runtime/arch probe failure
+        if int(os.environ.get("RANK", "0")) == 0:
+            warnings.warn(
+                f"[v4-flydsl] FlyDSL runtime unavailable ({exc!r}); FlyDSL "
+                f"attention backend disabled, falling back to Triton.",
+                RuntimeWarning,
+                stacklevel=3,
+            )
+        _FLYDSL_AVAILABLE = False
+    return _FLYDSL_AVAILABLE
+
+
+def is_flydsl_available() -> bool:
+    """True iff the FlyDSL runtime imported successfully (cached)."""
+    return _probe_flydsl()
+
+
+# --- lazy adapter cache -----------------------------------------------------
+_FWD_MQA = None  # _launch_v4_attention_fwd_flydsl_mqa
+_FWD_CSA = None  # _launch_v4_attention_fwd_csa
+
+
+def _load_fwd_mqa():
+    global _FWD_MQA
+    if _FWD_MQA is None:
+        from .kernels.v4_attention_fwd_flydsl_mqa import (
+            _launch_v4_attention_fwd_flydsl_mqa as fn,
+        )
+
+        _FWD_MQA = fn
+    return _FWD_MQA
+
+
+def _load_fwd_csa():
+    global _FWD_CSA
+    if _FWD_CSA is None:
+        from .kernels.v4_attention_fwd_flydsl_csa import (
+            _launch_v4_attention_fwd_csa as fn,
+        )
+
+        _FWD_CSA = fn
+    return _FWD_CSA
+
+
+def should_dispatch(kernel_name: str, *, enabled: bool) -> bool:
+    """True iff FlyDSL should handle ``kernel_name`` (else Triton fallback).
+
+    Short-circuits when ``enabled`` is False so the off path never imports FlyDSL.
+    """
+    if not enabled:
+        return False
+    if kernel_name not in _KNOWN_KERNEL_NAMES:
+        raise ValueError(
+            f"Unknown FlyDSL kernel {kernel_name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}"
+        )
+    if kernel_name not in _WIRED_FWD:
+        # bwd kernels are present in-package but lack a training-autograd hook
+        return False
+    if not _probe_flydsl():
+        return False
+    try:
+        _load_fwd_csa() if kernel_name == "v4_csa_attention_fwd" else _load_fwd_mqa()
+    except Exception as exc:
+        if int(os.environ.get("RANK", "0")) == 0:
+            warnings.warn(
+                f"[v4-flydsl] kernel {kernel_name!r} import failed ({exc!r}); " f"falling back to Triton.",
+                RuntimeWarning,
+                stacklevel=3,
+            )
+        return False
+    return True
+
+
+# --- forward adapters (signatures match the dispatch sites) -----------------
+def v4_attention_fwd_flydsl(
+    q,
+    k,
+    v,
+    *,
+    sink: Optional[Any] = None,
+    swa_window: int = 0,
+    additive_mask: Optional[Any] = None,
+    scale: float,
+    hca_local_seqlen: int = 0,
+):
+    """SWA/HCA forward via FlyDSL. Returns the attention output tensor."""
+    out, _lse = _load_fwd_mqa()(
+        q, k, v, sink, int(swa_window), additive_mask, float(scale), int(hca_local_seqlen)
+    )
+    return out
+
+
+def v4_csa_attention_fwd_flydsl(
+    q,
+    k_local,
+    v_local,
+    gathered,
+    *,
+    sparse_mask,
+    sink: Optional[Any] = None,
+    swa_window: int = 0,
+    scale: float,
+    attn_dropout: float = 0.0,
+    training: bool = False,
+):
+    """CSA forward via FlyDSL (the 2.79x kernel). Returns the output tensor."""
+    out, _lse = _load_fwd_csa()(
+        q, k_local, v_local, gathered, sink, int(swa_window), sparse_mask, float(scale)
+    )
+    return out
diff --git a/primus/modules/trainer/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py
similarity index 100%
rename from primus/modules/trainer/__init__.py
rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py
new file mode 100644
index 000000000..fc9faf040
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py
@@ -0,0 +1,1276 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 SWA attention backward FlyDSL launcher (Phase B STEP 1b, cr=0, SWA-only).
+
+API contract (matches the Triton ``_launch_v4_attention_bwd`` signature; the
+bwd_modes harness calls this function):
+
+    flydsl_v4_attention_bwd(
+        q, k, v, out, dout, lse,
+        *,
+        sink, swa_window, additive_mask, scale, hca_local_seqlen,
+    ) -> (dq, dk, dv, dsink)
+
+STEP 1b SCOPE
+-------------
+Current state of each compute stage:
+  * preprocess (D scalar) -> FlyDSL kernel (``v4_swa_bwd_preprocess_kernel``)
+  * dq                    -> FlyDSL kernel (``v4_swa_bwd_dq_kernel``)
+                              forked from kernels/sla_bwd_dq.py
+  * dk / dv               -> Triton kernel (``_v4_attention_bwd_dkv_kernel``)
+  * dsink                 -> FlyDSL kernel writes it (via atomic_fadd in dq)
+
+Env knobs:
+  V4_FLYDSL_BWD_FLY_PREPROCESS  default 1 (1=FlyDSL, 0=Triton)
+  V4_FLYDSL_BWD_FLY_DQ          default 1 (1=FlyDSL, 0=Triton)
+  V4_FLYDSL_BWD_VERBOSE         default 0 (1=print provenance line)
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+from typing import Optional, Tuple
+
+import torch
+
+_FLYDSL_SRC = "/workspace/FlyDSL-amd"
+if _FLYDSL_SRC not in sys.path:
+    sys.path.insert(0, _FLYDSL_SRC)
+
+os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2")
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+
+if "/workspace/Primus" not in sys.path:
+    sys.path.insert(0, "/workspace/Primus")
+import triton  # noqa: E402
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (  # noqa: E402
+    _v4_attention_bwd_dkv_kernel,
+    _v4_attention_bwd_dkv_pool_kernel,
+    _v4_attention_bwd_dq_kernel,
+    _v4_attention_bwd_preprocess_kernel,
+)
+
+# Built lazily on first call; module-level kernel build is expensive.
+_PREPROCESS_KERNEL_CACHE = {}
+_PREPROCESS_KERNEL_LOCK = threading.Lock()
+
+_DQ_KERNEL_CACHE = {}
+_DQ_KERNEL_LOCK = threading.Lock()
+_DKV_KERNEL_CACHE = {}
+_DKV_KERNEL_LOCK = threading.Lock()
+_DQ_POOL_KERNEL_CACHE = {}
+_DQ_POOL_KERNEL_LOCK = threading.Lock()
+_DKV_POOL_KERNEL_CACHE = {}
+_DKV_POOL_KERNEL_LOCK = threading.Lock()
+
+
+def _get_fly_preprocess(head_dim: int, dtype_str: str, block_rows: int):
+    key = (head_dim, dtype_str, block_rows)
+    with _PREPROCESS_KERNEL_LOCK:
+        if key in _PREPROCESS_KERNEL_CACHE:
+            return _PREPROCESS_KERNEL_CACHE[key]
+        from v4_sla_bwd_kernel import build_v4_swa_bwd_preprocess_module
+
+        launch = build_v4_swa_bwd_preprocess_module(
+            head_dim=head_dim,
+            dtype_str=dtype_str,
+            block_rows=block_rows,
+        )
+        _PREPROCESS_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _get_fly_dq(num_heads: int, head_dim: int, swa_window: int, dtype_str: str, mqa_kv: bool, has_sink: bool):
+    key = (num_heads, head_dim, swa_window, dtype_str, mqa_kv, has_sink)
+    with _DQ_KERNEL_LOCK:
+        if key in _DQ_KERNEL_CACHE:
+            return _DQ_KERNEL_CACHE[key]
+        from v4_sla_bwd_kernel import build_v4_swa_bwd_dq_module
+
+        launch = build_v4_swa_bwd_dq_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            swa_window=swa_window,
+            dtype_str=dtype_str,
+            mqa_kv=mqa_kv,
+            has_sink=has_sink,
+        )
+        _DQ_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _get_fly_dq_pool(
+    num_heads: int, head_dim: int, pool_size: int, hca_local_seqlen: int, dtype_str: str, mqa_kv: bool
+):
+    key = (num_heads, head_dim, pool_size, hca_local_seqlen, dtype_str, mqa_kv)
+    with _DQ_POOL_KERNEL_LOCK:
+        if key in _DQ_POOL_KERNEL_CACHE:
+            return _DQ_POOL_KERNEL_CACHE[key]
+        from v4_hca_bwd_dq_pool_kernel import build_v4_hca_bwd_dq_pool_module
+
+        launch = build_v4_hca_bwd_dq_pool_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            pool_size=pool_size,
+            hca_local_seqlen=hca_local_seqlen,
+            dtype_str=dtype_str,
+            mqa_kv=mqa_kv,
+        )
+        _DQ_POOL_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _get_fly_dkv_pool(
+    num_heads: int, head_dim: int, pool_size: int, hca_local_seqlen: int, dtype_str: str, mqa_kv: bool
+):
+    key = (num_heads, head_dim, pool_size, hca_local_seqlen, dtype_str, mqa_kv)
+    with _DKV_POOL_KERNEL_LOCK:
+        if key in _DKV_POOL_KERNEL_CACHE:
+            return _DKV_POOL_KERNEL_CACHE[key]
+        from v4_hca_bwd_dkv_pool_kernel import build_v4_hca_bwd_dkv_pool_module
+
+        launch = build_v4_hca_bwd_dkv_pool_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            pool_size=pool_size,
+            hca_local_seqlen=hca_local_seqlen,
+            dtype_str=dtype_str,
+            mqa_kv=mqa_kv,
+        )
+        _DKV_POOL_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _get_fly_dkv(num_heads: int, head_dim: int, swa_window: int, dtype_str: str, mqa_kv: bool):
+    key = (num_heads, head_dim, swa_window, dtype_str, mqa_kv)
+    with _DKV_KERNEL_LOCK:
+        if key in _DKV_KERNEL_CACHE:
+            return _DKV_KERNEL_CACHE[key]
+        from v4_sla_bwd_dkv_kernel import build_v4_swa_bwd_dkv_module
+
+        launch = build_v4_swa_bwd_dkv_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            swa_window=swa_window,
+            dtype_str=dtype_str,
+            mqa_kv=mqa_kv,
+        )
+        _DKV_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _run_preprocess(
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    *,
+    use_flydsl: bool,
+    block_m: int,
+    block_dmodel: int,
+) -> torch.Tensor:
+    """Compute D[b,h,m] = sum_d (out[b,h,m,d] * dout[b,h,m,d]) in fp32."""
+    B, HQ, Sq, D = out.shape
+    d_buf = torch.empty((B, HQ, Sq), device=out.device, dtype=torch.float32)
+    if use_flydsl:
+        out_f = out.contiguous().view(-1, D)
+        dout_f = dout.contiguous().view(-1, D)
+        delta_f = d_buf.view(-1)
+        n_rows = out_f.shape[0]
+        dtype_str = "bf16" if out.dtype == torch.bfloat16 else "f16"
+        max_block_threads = 256
+        threads_per_row = D // 8
+        for br in (8, 4, 2, 1):
+            if n_rows % br == 0 and br * threads_per_row <= max_block_threads:
+                block_rows = br
+                break
+        else:
+            block_rows = 1
+        launch = _get_fly_preprocess(D, dtype_str, block_rows)
+        launch(out_f, dout_f, delta_f, n_rows)
+        return d_buf
+    pre_grid = (triton.cdiv(Sq, block_m), B * HQ)
+    _v4_attention_bwd_preprocess_kernel[pre_grid](
+        out,
+        dout,
+        d_buf,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        Sq,
+        HEAD=HQ,
+        BLOCK_M=block_m,
+        BLOCK_DMODEL=block_dmodel,
+        num_warps=4,
+        num_stages=1,
+    )
+    return d_buf
+
+
+def _run_dq_triton(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dq_fp32,
+    dsink_fp32,
+    sink_arg,
+    mask_arg,
+    *,
+    scale,
+    swa_window_constexpr,
+    has_sink,
+    has_add_mask,
+    hca_local_seqlen,
+    use_causal,
+    block_m,
+    block_n,
+    block_dmodel,
+    stride_ms,
+    stride_mn,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    exact_tiles_m,
+    exact_tiles_n,
+):
+    dq_grid = (triton.cdiv(Sq, block_m), B * HQ)
+    _v4_attention_bwd_dq_kernel[dq_grid](
+        q,
+        k,
+        v,
+        dout,
+        lse,
+        d_buf,
+        dq_fp32,
+        dsink_fp32,
+        sink_arg,
+        mask_arg,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        k.stride(3),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        v.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        dq_fp32.stride(0),
+        dq_fp32.stride(1),
+        dq_fp32.stride(2),
+        dq_fp32.stride(3),
+        stride_ms,
+        stride_mn,
+        Sq,
+        Sk,
+        float(scale),
+        HEAD_Q=HQ,
+        HEAD_K=HK,
+        SWA_WINDOW=swa_window_constexpr,
+        HAS_SINK=has_sink,
+        HAS_ADD_MASK=has_add_mask,
+        HCA_LOCAL_SEQLEN=hca_local_seqlen,
+        USE_CAUSAL=use_causal,
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_DMODEL=block_dmodel,
+        EXACT_TILES_M=exact_tiles_m,
+        EXACT_TILES_N=exact_tiles_n,
+        num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_WARPS", "2")),
+        num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_STAGES", "1")),
+    )
+
+
+def _run_dq_flydsl(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dq_fp32,
+    dsink_fp32,
+    sink_arg,
+    *,
+    scale,
+    swa_window,
+    has_sink,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    D,
+):
+    """V4 SWA bwd dQ via FlyDSL. Writes dq_fp32 (fp32) and adds into
+    dsink_fp32 (fp32, atomic) when has_sink. Sink_arg is a fp32 dummy
+    buffer when has_sink=False.
+
+    Expects BHLD contiguous: q [B,HQ,Sq,D], k/v [B,HK,Sk,D], dout/lse/d_buf
+    same as Triton path. For MQA we ENFORCE HK==1 and pass k/v as-is
+    (the FlyDSL kernel uses stride_kh=0 via the mqa_kv codepath, which
+    expects K/V to be flat [B, Sk, D]-contiguous regardless of stride).
+    """
+    mqa_kv = HK == 1
+    # Sanity: V4 STEP 1 only supports MQA (HK==1). Keep this gate; if HK!=1
+    # we fall back to the Triton path upstream.
+    assert mqa_kv, f"FlyDSL dq STEP 1b only supports MQA (HK==1); got HK={HK}"
+    dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16"
+
+    launch = _get_fly_dq(
+        num_heads=HQ,
+        head_dim=D,
+        swa_window=int(swa_window),
+        dtype_str=dtype_str,
+        mqa_kv=True,
+        has_sink=has_sink,
+    )
+
+    # The FlyDSL kernel reads K/V via flat pointer ops, so they must be
+    # contiguous and rank-4 doesn't strictly matter. We pass the rank-4
+    # tensors directly; the BHLD path computes ((b*1 + 0)*Sk + n)*D + col
+    # for KV when mqa_kv=True. That matches the stride pattern of a
+    # [B, 1, Sk, D]-contiguous tensor.
+    # NOTE: K and V come in as [B, 1, Sk, D] from the harness MQA leaves
+    # (see bwd_modes _build_swa_inputs).
+    assert q.is_contiguous(), "q must be contiguous"
+    assert k.is_contiguous(), "k must be contiguous"
+    assert v.is_contiguous(), "v must be contiguous"
+    assert dout.is_contiguous(), "dout must be contiguous"
+    assert lse.is_contiguous(), "lse must be contiguous"
+    assert d_buf.is_contiguous(), "d_buf must be contiguous"
+    assert dq_fp32.is_contiguous(), "dq_fp32 must be contiguous"
+
+    launch(
+        q,  # Q [B, HQ, Sq, D]
+        k,  # K [B, 1, Sk, D] (MQA broadcast view)
+        v,  # V [B, 1, Sk, D]
+        dout,  # DOS [B, HQ, Sq, D]
+        lse,  # LSE [B, HQ, Sq] fp32
+        d_buf,  # DELTAS [B, HQ, Sq] fp32
+        dq_fp32,  # DQ [B, HQ, Sq, D] fp32 (OUTPUT)
+        dsink_fp32,  # DSINK [HQ] fp32 (OUTPUT, atomic)
+        sink_arg,  # SINK [HQ] fp32 (INPUT, dummy if !has_sink)
+        int(B),
+        int(Sq),
+        int(Sk),
+    )
+
+
+def _run_dkv_flydsl(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dk_fp32,
+    dv_fp32,
+    *,
+    scale,
+    swa_window,
+    has_sink,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    D,
+):
+    """V4 SWA bwd dKdV via FlyDSL. Writes dk_fp32 / dv_fp32 (fp32 buffers).
+    MQA only (HK==1). One program per (b, n_block), head-loop accumulator,
+    no atomics."""
+    mqa_kv = HK == 1
+    assert mqa_kv, f"FlyDSL dkv STEP 1c only supports MQA (HK==1); got HK={HK}"
+    print(
+        f"[provenance] FlyDSL dkv launcher invoked, B={B} H={HQ} Sk={Sk} D={D}",
+        flush=True,
+    )
+    dtype_str = "bf16" if q.dtype == __import__("torch").bfloat16 else "f16"
+    launch = _get_fly_dkv(
+        num_heads=HQ,
+        head_dim=D,
+        swa_window=int(swa_window),
+        dtype_str=dtype_str,
+        mqa_kv=True,
+    )
+    assert q.is_contiguous(), "q must be contiguous"
+    assert k.is_contiguous(), "k must be contiguous"
+    assert v.is_contiguous(), "v must be contiguous"
+    assert dout.is_contiguous(), "dout must be contiguous"
+    assert lse.is_contiguous(), "lse must be contiguous"
+    assert d_buf.is_contiguous(), "d_buf must be contiguous"
+    assert dk_fp32.is_contiguous(), "dk_fp32 must be contiguous"
+    assert dv_fp32.is_contiguous(), "dv_fp32 must be contiguous"
+    launch(
+        q,  # Q [B, HQ, Sq, D]
+        k,  # K [B, 1, Sk, D]
+        v,  # V [B, 1, Sk, D]
+        dout,  # DOS [B, HQ, Sq, D]
+        lse,  # LSE [B, HQ, Sq] fp32 (RAW domain)
+        d_buf,  # DELTAS [B, HQ, Sq] fp32
+        dk_fp32,  # DK [B, 1, Sk, D] fp32 (OUTPUT)
+        dv_fp32,  # DV [B, 1, Sk, D] fp32 (OUTPUT)
+        int(B),
+        int(Sq),
+        int(Sk),
+    )
+
+
+def _run_dq_pool_flydsl(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dq_fp32,
+    add_mask,
+    *,
+    pool_size,
+    hca_local_seqlen,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    D,
+):
+    """V4 HCA bwd dq POOL stream via FlyDSL. ACCUMULATES into dq_fp32.
+
+    Called AFTER ``_run_dq_flydsl`` has written the LOCAL stream dq into
+    dq_fp32. Race-free since each program owns a unique
+    (b, qhid, m_block) slice and the launch is sequenced after the local
+    dq launch.
+    """
+    mqa_kv = HK == 1
+    assert mqa_kv, f"FlyDSL dq_pool only supports MQA (HK==1); got HK={HK}"
+    assert pool_size <= 64, f"pool_size must fit in one BLOCK_N=64 block; got pool_size={pool_size}"
+    assert (
+        hca_local_seqlen % 64 == 0
+    ), f"hca_local_seqlen must be multiple of BLOCK_N=64; got {hca_local_seqlen}"
+    assert (
+        hca_local_seqlen + pool_size
+    ) == Sk, f"expected hca_local_seqlen+pool_size == Sk; got {hca_local_seqlen}+{pool_size}!={Sk}"
+    assert add_mask is not None and add_mask.shape == (
+        Sq,
+        pool_size,
+    ), f"add_mask shape mismatch: expected ({Sq},{pool_size}); got {tuple(add_mask.shape) if add_mask is not None else None}"
+    assert add_mask.dtype == q.dtype, f"add_mask dtype must match q.dtype; got {add_mask.dtype} vs {q.dtype}"
+    dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16"
+    launch = _get_fly_dq_pool(
+        num_heads=HQ,
+        head_dim=D,
+        pool_size=int(pool_size),
+        hca_local_seqlen=int(hca_local_seqlen),
+        dtype_str=dtype_str,
+        mqa_kv=True,
+    )
+    assert q.is_contiguous(), "q must be contiguous"
+    assert k.is_contiguous(), "k must be contiguous"
+    assert v.is_contiguous(), "v must be contiguous"
+    assert dout.is_contiguous(), "dout must be contiguous"
+    assert lse.is_contiguous(), "lse must be contiguous"
+    assert d_buf.is_contiguous(), "d_buf must be contiguous"
+    assert dq_fp32.is_contiguous(), "dq_fp32 must be contiguous"
+    assert add_mask.is_contiguous(), "add_mask must be contiguous"
+    launch(
+        q,
+        k,
+        v,
+        dout,
+        lse,
+        d_buf,
+        dq_fp32,
+        add_mask,
+        int(B),
+        int(Sq),
+        int(Sk),
+    )
+
+
+def _run_dkv_pool_flydsl(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dk_fp32,
+    dv_fp32,
+    add_mask,
+    *,
+    pool_size,
+    hca_local_seqlen,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    D,
+):
+    """V4 HCA bwd dKdV POOL stream via FlyDSL. WRITES into the POOL slice of
+    dk_fp32/dv_fp32 (which are zero-initialized by the wrapper and have the
+    LOCAL slice already populated by ``_run_dkv_flydsl``). Race-free:
+    LOCAL writes [..., :hca_local_seqlen, :] and POOL writes
+    [..., hca_local_seqlen:hca_local_seqlen+pool_size, :] -- disjoint.
+    """
+    mqa_kv = HK == 1
+    assert mqa_kv, f"FlyDSL dkv_pool only supports MQA (HK==1); got HK={HK}"
+    assert pool_size <= 32, f"pool_size must fit in one BLOCK_N=32 block; got pool_size={pool_size}"
+    assert (
+        hca_local_seqlen % 32 == 0
+    ), f"hca_local_seqlen must be multiple of BLOCK_N=32; got {hca_local_seqlen}"
+    assert (hca_local_seqlen + pool_size) == Sk, (
+        f"expected hca_local_seqlen+pool_size == Sk; got " f"{hca_local_seqlen}+{pool_size}!={Sk}"
+    )
+    assert add_mask is not None and add_mask.shape == (Sq, pool_size), (
+        f"add_mask shape mismatch: expected ({Sq},{pool_size}); got "
+        f"{tuple(add_mask.shape) if add_mask is not None else None}"
+    )
+    assert add_mask.dtype == q.dtype, f"add_mask dtype must match q.dtype; got {add_mask.dtype} vs {q.dtype}"
+    dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16"
+    launch = _get_fly_dkv_pool(
+        num_heads=HQ,
+        head_dim=D,
+        pool_size=int(pool_size),
+        hca_local_seqlen=int(hca_local_seqlen),
+        dtype_str=dtype_str,
+        mqa_kv=True,
+    )
+    assert q.is_contiguous(), "q must be contiguous"
+    assert k.is_contiguous(), "k must be contiguous"
+    assert v.is_contiguous(), "v must be contiguous"
+    assert dout.is_contiguous(), "dout must be contiguous"
+    assert lse.is_contiguous(), "lse must be contiguous"
+    assert d_buf.is_contiguous(), "d_buf must be contiguous"
+    assert dk_fp32.is_contiguous(), "dk_fp32 must be contiguous"
+    assert dv_fp32.is_contiguous(), "dv_fp32 must be contiguous"
+    assert add_mask.is_contiguous(), "add_mask must be contiguous"
+    launch(
+        q,
+        k,
+        v,
+        dout,
+        lse,
+        d_buf,
+        dk_fp32,
+        dv_fp32,
+        add_mask,
+        int(B),
+        int(Sq),
+        int(Sk),
+    )
+
+
+def _run_dkv_triton(
+    q,
+    k,
+    v,
+    dout,
+    lse,
+    d_buf,
+    dk_fp32,
+    dv_fp32,
+    mask_arg,
+    *,
+    scale,
+    swa_window_constexpr,
+    has_add_mask,
+    hca_local_seqlen,
+    use_causal,
+    block_m,
+    block_n,
+    block_dmodel,
+    stride_ms,
+    stride_mn,
+    Sq,
+    Sk,
+    HQ,
+    HK,
+    B,
+    exact_tiles_m,
+):
+    dkv_block_n = block_n
+    dkv_n_blocks = triton.cdiv(Sk, dkv_block_n)
+    num_head_groups = 1
+    if HQ > HK:
+        # MQA/HQ>=128 (V4-Pro): HG=2 on gfx950 mirrors the prod
+        # _launch_v4_attention_bwd default. Overridable via env.
+        _hg_default = "2" if (HQ >= 64 and HK == 1) else "1"
+        target = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS", _hg_default))
+        while target > 1 and HQ % target != 0:
+            target //= 2
+        num_head_groups = max(1, target)
+    dkv_grid = (dkv_n_blocks, B * HK, num_head_groups)
+    _v4_attention_bwd_dkv_kernel[dkv_grid](
+        q,
+        k,
+        v,
+        dout,
+        lse,
+        d_buf,
+        dk_fp32,
+        dv_fp32,
+        mask_arg,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        k.stride(3),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        v.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        dk_fp32.stride(0),
+        dk_fp32.stride(1),
+        dk_fp32.stride(2),
+        dk_fp32.stride(3),
+        dv_fp32.stride(0),
+        dv_fp32.stride(1),
+        dv_fp32.stride(2),
+        dv_fp32.stride(3),
+        stride_ms,
+        stride_mn,
+        Sq,
+        Sk,
+        float(scale),
+        HEAD_Q=HQ,
+        HEAD_K=HK,
+        SWA_WINDOW=swa_window_constexpr,
+        HAS_ADD_MASK=has_add_mask,
+        HCA_LOCAL_SEQLEN=hca_local_seqlen,
+        USE_CAUSAL=use_causal,
+        BLOCK_M=block_m,
+        BLOCK_N=dkv_block_n,
+        BLOCK_DMODEL=block_dmodel,
+        NUM_HEAD_GROUPS=num_head_groups,
+        EXACT_TILES_M=exact_tiles_m,
+        EXACT_TILES_N=(Sk % dkv_block_n) == 0,
+        num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_WARPS", "2")),
+        num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_STAGES", "1")),
+    )
+
+
+# ---------------------------------------------------------------------------
+# HCA pool-only dq accumulator (Triton, inline).
+#
+# This is a TEMPORARY Triton fallback for the POOL stream of HCA dq while
+# the FlyDSL pool kernel is being authored. It is structurally equivalent
+# to the pool branch of ``_v4_attention_bwd_dq_kernel`` (Triton ref) but
+# written as a standalone kernel so we can call it as a pure accumulator
+# (``DQ += pool_contrib``, no local loop). One program per (m_block, b*qhid).
+# ---------------------------------------------------------------------------
+
+
+import triton.language as tl  # noqa: E402
+
+
+@triton.jit
+def _hca_pool_dq_accumulator(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    ADD_MASK,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    pool_size,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+    if HEAD_K == HEAD_Q:
+        khid = qhid
+    else:
+        khid = 0
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    pool_n = tl.arange(0, BLOCK_N)
+    offs_n = HCA_LOCAL_SEQLEN + pool_n
+    NEG_INF: tl.constexpr = -1.0e30
+    pool_n_mask = pool_n < pool_size
+
+    q_ptrs = (
+        Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+    )
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + qhid * stride_doh
+        + offs_m[:, None] * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+    dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+    q_load_mask = offs_m[:, None] < seqlen_q
+    q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+    dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+    lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+    dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+    k_ptrs = (
+        K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+    )
+    v_ptrs = (
+        V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+    )
+    kv_load_mask = pool_n_mask[:, None]
+    k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+    v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+
+    qk = tl.dot(q, tl.trans(k)) * sm_scale
+    mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+    mask_load_mask = (offs_m[:, None] < seqlen_q) & pool_n_mask[None, :]
+    add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+    qk = qk + add_bias
+    qk = tl.where(pool_n_mask[None, :], qk, NEG_INF)
+    qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+    p = tl.exp(qk - lse[:, None])
+    dp = tl.dot(dout, tl.trans(v))
+    ds = p * (dp - dvec[:, None])
+    dq_contrib = tl.dot(ds.to(k.dtype), k) * sm_scale
+
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + qhid * stride_dqh
+        + offs_m[:, None] * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    dq_prev = tl.load(dq_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
+    tl.store(dq_ptrs, dq_prev + dq_contrib, mask=offs_m[:, None] < seqlen_q)
+
+
+def flydsl_v4_attention_bwd(
+    q: torch.Tensor,  # [B, HQ, Sq, D]
+    k: torch.Tensor,  # [B, HK, Sk, D]  (HK == 1 for MQA, HK == HQ for MHA)
+    v: torch.Tensor,  # [B, HK, Sk, D]
+    out: torch.Tensor,  # [B, HQ, Sq, D]
+    dout: torch.Tensor,  # [B, HQ, Sq, D]
+    lse: torch.Tensor,  # [B, HQ, Sq] fp32
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],
+    scale: float,
+    hca_local_seqlen: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """V4 SWA backward launcher (STEP 1b scope: SWA-only, no HCA)."""
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError("rank-4 q/k/v required")
+    if dout.shape != out.shape or out.shape != q.shape:
+        raise ValueError(
+            f"shape mismatch: q={tuple(q.shape)} out={tuple(out.shape)} dout={tuple(dout.shape)}"
+        )
+    B, HQ, Sq, D = q.shape
+    HK = k.shape[1]
+    Sk = k.shape[2]
+    if k.shape != (B, HK, Sk, D) or v.shape != k.shape:
+        raise ValueError(f"k/v shape mismatch: k={tuple(k.shape)} v={tuple(v.shape)}")
+    if q.dtype != torch.bfloat16:
+        raise NotImplementedError(f"bf16 only; got q.dtype={q.dtype}")
+    if D % 16 != 0:
+        raise NotImplementedError(f"head_dim must be multiple of 16; got D={D}")
+    # STEP 2 HCA: accept SWA-only (additive_mask=None, hca_local_seqlen=0)
+    # AND HCA (additive_mask is [Sq, P], hca_local_seqlen=Sq, Sk=Sq+P).
+    is_hca = (additive_mask is not None) and (int(hca_local_seqlen) > 0)
+    if (additive_mask is not None) != (int(hca_local_seqlen) > 0):
+        raise NotImplementedError(
+            "additive_mask and hca_local_seqlen must both be set (HCA) or both unset (SWA)"
+        )
+    if is_hca:
+        if int(hca_local_seqlen) != Sq:
+            raise NotImplementedError(
+                f"HCA requires hca_local_seqlen == Sq; got {int(hca_local_seqlen)} vs Sq={Sq}"
+            )
+        if int(swa_window) <= 0:
+            raise NotImplementedError("HCA requires swa_window > 0 for the local stream")
+        if Sk <= Sq:
+            raise NotImplementedError(f"HCA requires Sk>Sq; got Sk={Sk} Sq={Sq}")
+        pool_size = Sk - Sq
+        if additive_mask.shape != (Sq, pool_size):
+            raise NotImplementedError(
+                f"HCA additive_mask must be [Sq={Sq}, P={pool_size}]; got {tuple(additive_mask.shape)}"
+            )
+        if B != 1:
+            raise NotImplementedError(
+                f"HCA path currently only supports B=1 (got B={B}); B>1 would mis-stride K/V."
+            )
+    else:
+        if int(swa_window) <= 0:
+            raise NotImplementedError("SWA requires swa_window > 0")
+        if Sq != Sk:
+            raise NotImplementedError("SWA requires Sq == Sk")
+
+    has_sink = sink is not None
+    has_add_mask = False
+    use_causal = False
+    swa_window_constexpr = int(swa_window)
+
+    BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_M", "32"))
+    BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_N", "16"))
+    BLOCK_DMODEL = D
+    exact_tiles_m = (Sq % BLOCK_M) == 0
+    exact_tiles_n = (Sk % BLOCK_N) == 0
+
+    dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+    dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32)
+    dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32)
+    if has_sink:
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink
+    else:
+        # Dummy buffers; FlyDSL kernel still takes them but doesn't touch
+        # them when has_sink=False at build time. For Triton path we mirror
+        # the original behavior of passing q-shape sentinels.
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        sink_arg = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+
+    use_fly_preprocess = os.getenv("V4_FLYDSL_BWD_FLY_PREPROCESS", "1") == "1"
+    d_buf = _run_preprocess(
+        out,
+        dout,
+        use_flydsl=use_fly_preprocess,
+        block_m=BLOCK_M,
+        block_dmodel=BLOCK_DMODEL,
+    )
+
+    mask_arg = q
+    stride_ms = 0
+    stride_mn = 0
+
+    use_fly_dq = os.getenv("V4_FLYDSL_BWD_FLY_DQ", "1") == "1"
+    use_fly_dkv = os.getenv("V4_FLYDSL_BWD_FLY_DKV", "1") == "1"
+    verbose = os.getenv("V4_FLYDSL_BWD_VERBOSE", "0") == "1"
+
+    if verbose:
+        pp_tag = "flydsl" if use_fly_preprocess else "triton"
+        dq_tag = "flydsl" if use_fly_dq else "triton"
+        dkv_tag = "flydsl" if use_fly_dkv else "triton"
+        print(
+            f"[v4_bwd] preproc={pp_tag} dq={dq_tag} dkv={dkv_tag} "
+            f"has_sink={has_sink} Sq={Sq} Sk={Sk} HQ={HQ} HK={HK} D={D} swa={swa_window}",
+            flush=True,
+        )
+
+    if is_hca:
+        # ---- HCA mode (split-mask) ----
+        # 1. LOCAL stream (n < HCA_LOCAL_SEQLEN): existing FlyDSL dq + dkv,
+        #    with effective seq_len_k = HCA_LOCAL_SEQLEN. Kernel uses
+        #    seq_len_k for both n-loop bound and batch base; B=1 makes the
+        #    base zero so the reduced seq_len_k cleanly bounds the n-loop
+        #    without mis-striding K/V. (B=1-only; gated above.) The LOCAL
+        #    pass uses the JOINT lse (saved from fwd with both streams) and
+        #    JOINT delta (preprocess used full out, dout).
+        # 2. POOL stream (n in [HCA_LOCAL_SEQLEN, Sk)): Triton pool kernels
+        #    for this round; FlyDSL pool kernels are next-round work.
+        Sk_local = int(hca_local_seqlen)  # == Sq
+        pool_size = Sk - Sk_local
+        sink_dq_arg = sink_arg  # both fp32 [HQ]
+        use_fly_dq_pool = os.getenv("V4_FLYDSL_BWD_FLY_DQ_POOL", "0") == "1"
+        use_fly_dkv_pool = os.getenv("V4_FLYDSL_BWD_FLY_DKV_POOL", "0") == "1"
+        # Provenance: tag dq_pool / dkv_pool only based on knob (silent
+        # fallback would violate the strict-gate rule).
+        dq_pool_tag = "FlyDSL" if use_fly_dq_pool else "Triton"
+        dkv_pool_tag = "FlyDSL" if use_fly_dkv_pool else "Triton"
+        print(
+            f"[v4_bwd_hca] provenance: dq_local=FlyDSL dkv_local=FlyDSL "
+            f"dsink=FlyDSL dq_pool={dq_pool_tag} dkv_pool={dkv_pool_tag} "
+            f"B={B} HQ={HQ} HK={HK} Sq={Sq} Sk={Sk} pool={pool_size} D={D} "
+            f"swa={swa_window} hca_local_seqlen={Sk_local}",
+            flush=True,
+        )
+        # ---- LOCAL FlyDSL dq (computes dq_local + dsink) ----
+        _run_dq_flydsl(
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dq_fp32,
+            dsink_fp32,
+            sink_dq_arg,
+            scale=scale,
+            swa_window=swa_window_constexpr,
+            has_sink=has_sink,
+            Sq=Sq,
+            Sk=Sk_local,
+            HQ=HQ,
+            HK=HK,
+            B=B,
+            D=D,
+        )
+        # ---- LOCAL FlyDSL dkv ----
+        _run_dkv_flydsl(
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dk_fp32,
+            dv_fp32,
+            scale=scale,
+            swa_window=swa_window_constexpr,
+            has_sink=has_sink,
+            Sq=Sq,
+            Sk=Sk_local,
+            HQ=HQ,
+            HK=HK,
+            B=B,
+            D=D,
+        )
+        # ---- POOL dq accumulator: FlyDSL (knob ON) or Triton (default) ----
+        if use_fly_dq_pool:
+            _run_dq_pool_flydsl(
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dq_fp32,
+                additive_mask,
+                pool_size=int(pool_size),
+                hca_local_seqlen=Sk_local,
+                Sq=Sq,
+                Sk=Sk,
+                HQ=HQ,
+                HK=HK,
+                B=B,
+                D=D,
+            )
+        else:
+            pool_block_n_dq = max(16, triton.next_power_of_2(pool_size))
+            _hca_pool_dq_accumulator[(triton.cdiv(Sq, BLOCK_M), B * HQ)](
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dq_fp32,
+                additive_mask,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                k.stride(0),
+                k.stride(1),
+                k.stride(2),
+                k.stride(3),
+                v.stride(0),
+                v.stride(1),
+                v.stride(2),
+                v.stride(3),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_fp32.stride(0),
+                dq_fp32.stride(1),
+                dq_fp32.stride(2),
+                dq_fp32.stride(3),
+                additive_mask.stride(0),
+                additive_mask.stride(1),
+                Sq,
+                Sk,
+                pool_size,
+                float(scale),
+                HEAD_Q=HQ,
+                HEAD_K=HK,
+                HCA_LOCAL_SEQLEN=Sk_local,
+                BLOCK_M=BLOCK_M,
+                BLOCK_N=pool_block_n_dq,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=2,
+                num_stages=1,
+            )
+        # ---- POOL dkv: FlyDSL (knob ON) or Triton (default) ----
+        if use_fly_dkv_pool:
+            _run_dkv_pool_flydsl(
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dk_fp32,
+                dv_fp32,
+                additive_mask,
+                pool_size=int(pool_size),
+                hca_local_seqlen=Sk_local,
+                Sq=Sq,
+                Sk=Sk,
+                HQ=HQ,
+                HK=HK,
+                B=B,
+                D=D,
+            )
+        else:
+            pool_block_n_dkv = max(16, triton.next_power_of_2(pool_size))
+            pool_grid_m = triton.cdiv(Sq, BLOCK_M)
+            _v4_attention_bwd_dkv_pool_kernel[(pool_grid_m, B)](
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dk_fp32,
+                dv_fp32,
+                additive_mask,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                k.stride(0),
+                k.stride(1),
+                k.stride(2),
+                k.stride(3),
+                v.stride(0),
+                v.stride(1),
+                v.stride(2),
+                v.stride(3),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dk_fp32.stride(0),
+                dk_fp32.stride(1),
+                dk_fp32.stride(2),
+                dk_fp32.stride(3),
+                dv_fp32.stride(0),
+                dv_fp32.stride(1),
+                dv_fp32.stride(2),
+                dv_fp32.stride(3),
+                additive_mask.stride(0),
+                additive_mask.stride(1),
+                Sq,
+                Sk,
+                pool_size,
+                float(scale),
+                HEAD_Q=HQ,
+                HEAD_K=HK,
+                HCA_LOCAL_SEQLEN=Sk_local,
+                BLOCK_M=BLOCK_M,
+                BLOCK_N=pool_block_n_dkv,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=2,
+                num_stages=1,
+            )
+    elif use_fly_dq:
+        # SWA-only path (STEP 1b unchanged).
+        sink_dq_arg = sink_arg  # both are fp32 [HQ]
+        _run_dq_flydsl(
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dq_fp32,
+            dsink_fp32,
+            sink_dq_arg,
+            scale=scale,
+            swa_window=swa_window_constexpr,
+            has_sink=has_sink,
+            Sq=Sq,
+            Sk=Sk,
+            HQ=HQ,
+            HK=HK,
+            B=B,
+            D=D,
+        )
+    else:
+        _run_dq_triton(
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dq_fp32,
+            dsink_fp32,
+            sink_arg,
+            mask_arg,
+            scale=scale,
+            swa_window_constexpr=swa_window_constexpr,
+            has_sink=has_sink,
+            has_add_mask=has_add_mask,
+            hca_local_seqlen=0,
+            use_causal=use_causal,
+            block_m=BLOCK_M,
+            block_n=BLOCK_N,
+            block_dmodel=BLOCK_DMODEL,
+            stride_ms=stride_ms,
+            stride_mn=stride_mn,
+            Sq=Sq,
+            Sk=Sk,
+            HQ=HQ,
+            HK=HK,
+            B=B,
+            exact_tiles_m=exact_tiles_m,
+            exact_tiles_n=exact_tiles_n,
+        )
+
+    if not is_hca:
+        if use_fly_dkv:
+            _run_dkv_flydsl(
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dk_fp32,
+                dv_fp32,
+                scale=scale,
+                swa_window=swa_window_constexpr,
+                has_sink=has_sink,
+                Sq=Sq,
+                Sk=Sk,
+                HQ=HQ,
+                HK=HK,
+                B=B,
+                D=D,
+            )
+        else:
+            _run_dkv_triton(
+                q,
+                k,
+                v,
+                dout,
+                lse,
+                d_buf,
+                dk_fp32,
+                dv_fp32,
+                mask_arg,
+                scale=scale,
+                swa_window_constexpr=swa_window_constexpr,
+                has_add_mask=has_add_mask,
+                hca_local_seqlen=0,
+                use_causal=use_causal,
+                block_m=BLOCK_M,
+                block_n=BLOCK_N,
+                block_dmodel=BLOCK_DMODEL,
+                stride_ms=stride_ms,
+                stride_mn=stride_mn,
+                Sq=Sq,
+                Sk=Sk,
+                HQ=HQ,
+                HK=HK,
+                B=B,
+                exact_tiles_m=exact_tiles_m,
+            )
+
+    dq_out = dq_fp32.to(q.dtype)
+    dk_out = dk_fp32.to(k.dtype)
+    dv_out = dv_fp32.to(v.dtype)
+    dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+    return dq_out, dk_out, dv_out, dsink_out
+
+
+__all__ = ["flydsl_v4_attention_bwd"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py
new file mode 100644
index 000000000..ba88d7e76
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py
@@ -0,0 +1,225 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention forward FlyDSL launcher (Round 3 Step 2b).
+
+Stage A: correctness only. Per-row design forked from Triton monolithic CSA.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import sys
+import threading
+from typing import Optional, Tuple
+
+import torch
+
+_FLYDSL_SRC = "/workspace/FlyDSL-amd"
+if _FLYDSL_SRC not in sys.path:
+    sys.path.insert(0, _FLYDSL_SRC)
+
+os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2")
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+from v4_csa_fwd_kernel import build_v4_csa_fwd_module  # noqa: E402
+
+_KERNEL_CACHE = {}
+_KERNEL_CACHE_LOCK = threading.Lock()
+
+
+def _get_kernel(
+    num_heads_q,
+    head_dim,
+    swa_window,
+    dtype_str,
+    block_n,
+    block_k,
+    waves_per_eu,
+    has_sink,
+    has_sparse,
+    mqa_kv,
+    head_group=1,
+):
+    key = (
+        num_heads_q,
+        head_dim,
+        swa_window,
+        dtype_str,
+        block_n,
+        block_k,
+        waves_per_eu,
+        has_sink,
+        has_sparse,
+        mqa_kv,
+        head_group,
+    )
+    with _KERNEL_CACHE_LOCK:
+        if key in _KERNEL_CACHE:
+            return _KERNEL_CACHE[key]
+        launch = build_v4_csa_fwd_module(
+            num_heads=num_heads_q,
+            head_dim=head_dim,
+            swa_window=int(swa_window),
+            dtype_str=dtype_str,
+            waves_per_eu=waves_per_eu,
+            block_n=block_n,
+            block_k=block_k,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            mqa_kv=mqa_kv,
+            head_group=head_group,
+        )
+        _KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _broadcast_kv_mqa(k: torch.Tensor, v: torch.Tensor, head_q: int):
+    if k.shape[1] == head_q:
+        return k, v
+    if k.shape[1] != 1:
+        raise ValueError(f"MQA expects K_H=1; got {k.shape[1]}")
+    k_full = k.expand(-1, head_q, -1, -1).clone(memory_format=torch.contiguous_format)
+    v_full = v.expand(-1, head_q, -1, -1).clone(memory_format=torch.contiguous_format)
+    return k_full, v_full
+
+
+def _launch_v4_attention_fwd_csa(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H_KV, Sq, D]  (H_KV=1 for MQA)
+    v_local: torch.Tensor,  # [B, H_KV, Sq, D]
+    gathered: torch.Tensor,  # [B, Sq, K_topk, D]
+    *,
+    sink: Optional[torch.Tensor],  # [H] fp32 or None
+    swa_window: int,
+    sparse_mask: torch.Tensor,  # [B, Sq, K_topk] additive
+    scale: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """V4 CSA forward launcher.
+
+    Returns (out, lse) where out is bf16 [B, H, Sq, D] and lse is fp32 [B, H, Sq]
+    in raw-qk-scaled-domain (m + ln(l), where m absorbs sm_scale).
+    """
+    if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4:
+        raise ValueError("rank-4 q/k_local/v_local required")
+    if gathered.dim() != 4:
+        raise ValueError("rank-4 gathered [B,Sq,K_topk,D] required")
+    if sparse_mask.dim() != 3:
+        raise ValueError("rank-3 sparse_mask [B,Sq,K_topk] required")
+    B, HQ, Sq, D = q.shape
+    Bk, HK, Sk, Dk = k_local.shape
+    if (Bk, Sk, Dk) != (B, Sq, D) or v_local.shape != k_local.shape:
+        raise ValueError("k_local/v_local shape mismatch w.r.t. q")
+    if HK != 1 and HK != HQ:
+        raise ValueError(f"K_H must be 1 or {HQ}; got {HK}")
+    Bg, Sqg, K_topk, Dg = gathered.shape
+    if Bg != B or Sqg != Sq or Dg != D:
+        raise ValueError(f"gathered shape mismatch {tuple(gathered.shape)}")
+    Bm, Sqm, Km = sparse_mask.shape
+    if Bm != B or Sqm != Sq or Km != K_topk:
+        raise ValueError(f"sparse_mask shape mismatch {tuple(sparse_mask.shape)}")
+    if q.dtype != torch.bfloat16:
+        raise NotImplementedError(f"bf16 only; got {q.dtype}")
+    if D != 512:
+        raise NotImplementedError(f"head_dim=512 only; got {D}")
+    if int(swa_window) <= 0:
+        raise NotImplementedError("swa_window > 0 required")
+    expected_scale = 1.0 / math.sqrt(D)
+    if not math.isclose(float(scale), expected_scale, rel_tol=1e-4):
+        raise NotImplementedError(f"only scale=1/sqrt(D) supported; got {scale}")
+
+    has_sink = sink is not None
+    has_sparse = int(K_topk) > 0
+    mqa = (HK == 1) and (HQ > 1)
+
+    if has_sink:
+        sink_fp32 = sink.float().contiguous()
+        if sink_fp32.shape != (HQ,):
+            raise ValueError(f"sink shape must be ({HQ},); got {tuple(sink_fp32.shape)}")
+    else:
+        sink_fp32 = torch.zeros((max(HQ, 1),), dtype=torch.float32, device=q.device)
+
+    # MQA: avoid the .expand().clone() broadcast — pass [B,1,Sq,D] directly,
+    # kernel uses mqa_kv=True to drop head_idx from K/V indexing.
+    q_bhld = q.contiguous()
+    if mqa:
+        k_bhld = k_local.contiguous()
+        v_bhld = v_local.contiguous()
+    else:
+        k_bhld = k_local.contiguous()
+        v_bhld = v_local.contiguous()
+    o_bhld = torch.empty_like(q_bhld)
+    lse = torch.zeros((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    # Sparse mask & gathered must be contiguous fp32 / bf16.
+    if K_topk > 0:
+        gathered_c = gathered.contiguous()
+        # sparse_mask is bf16 (typically) -> upcast to fp32 here so kernel can
+        # just add it as f32. (Triton path takes the dtype of input.)
+        if sparse_mask.dtype != torch.float32:
+            sparse_mask_fp32 = sparse_mask.float().contiguous()
+        else:
+            sparse_mask_fp32 = sparse_mask.contiguous()
+    else:
+        gathered_c = torch.empty((B, Sq, 1, D), dtype=q.dtype, device=q.device)
+        sparse_mask_fp32 = torch.zeros((B, Sq, 1), dtype=torch.float32, device=q.device)
+
+    block_n = int(os.environ.get("PRIMUS_V4_CSA_BLOCK_N", "8"))
+    block_k = int(os.environ.get("PRIMUS_V4_CSA_BLOCK_K", "16"))
+    waves_per_eu = int(os.environ.get("FLYDSL_WAVES_PER_EU", "2"))
+    # HG=2 is the banked default for the sparse path (same VGPR=250 / spill=0
+    # footprint as HG=1). K_topk==0 (dense fallback) still forces HG=1 below.
+    head_group = int(os.environ.get("PRIMUS_V4_CSA_HEAD_GROUP", "2"))
+    # Round 5 fix: dense-fallback path (K_topk==0) has no payoff from head-group
+    # fusion AND triggered a flyc closure-cache collision when has_sparse differed
+    # between HG=1 dense and HG>1 sparse compiles. Force HG=1 when dense.
+    if not has_sparse:
+        head_group = 1
+    if HQ % head_group != 0:
+        # Silently fall back to HG=1 if shape doesn't divide.
+        head_group = 1
+    launch = _get_kernel(
+        HQ,
+        D,
+        int(swa_window),
+        "bf16",
+        block_n,
+        block_k,
+        waves_per_eu,
+        has_sink,
+        has_sparse,
+        mqa,
+        head_group=head_group,
+    )
+
+    # FlyDSL packs each tensor *shape* dim as int32 (jit_argument.py:
+    # `"i" * len(shape)`). A flat `.view(-1)` collapses a tensor to a single
+    # dim equal to its element count; for V4-Pro CSA `gathered` that is
+    # B*Sq*K_topk*D = 1*4096*1024*512 = 2**31, which overflows int32 ('i')
+    # and aborts the launch. The kernel only needs the aligned base pointer
+    # (it computes all offsets from the B / Sq / K_topk scalars), so passing
+    # `gathered` in its natural [B, Sq, K_topk, D] shape keeps every dim
+    # < 2**31 (strides are packed as int64) without changing kernel math.
+    launch(
+        q_bhld.view(-1),
+        k_bhld.view(-1),
+        v_bhld.view(-1),
+        gathered_c,
+        sparse_mask_fp32.view(-1),
+        sink_fp32.view(-1),
+        o_bhld.view(-1),
+        lse.view(-1),
+        B,
+        Sq,
+        int(K_topk),
+    )
+    return o_bhld, lse
+
+
+__all__ = ["_launch_v4_attention_fwd_csa"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py
new file mode 100644
index 000000000..c66a82be2
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py
@@ -0,0 +1,149 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 SWA attention forward FlyDSL launcher with MQA stride-trick (Round 3 Stage C).
+
+Eliminates the K/V .expand().clone() broadcast that allocates H_Q copies of
+K and V. Passes the un-broadcast [B, 1, Sk, D] view directly; the kernel reads
+K/V with stride_kh=0 via the mqa_kv compile-time flag.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import sys
+import threading
+from typing import Optional, Tuple
+
+import torch
+
+_FLYDSL_SRC = "/workspace/FlyDSL-amd"
+if _FLYDSL_SRC not in sys.path:
+    sys.path.insert(0, _FLYDSL_SRC)
+
+os.environ.setdefault("FLYDSL_SLA_FWD_ENABLE_DMA", "1")
+os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2")
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+from v4_sla_fwd_kernel import build_v4_swa_fwd_module  # noqa: E402
+
+_KERNEL_CACHE = {}
+_KERNEL_CACHE_LOCK = threading.Lock()
+
+
+def _get_kernel(
+    num_heads_q, head_dim, swa_window, dtype_str, block_m, block_n, waves_per_eu, mqa_kv, flat_work_group_size
+):
+    key = (
+        num_heads_q,
+        head_dim,
+        swa_window,
+        dtype_str,
+        block_m,
+        block_n,
+        waves_per_eu,
+        mqa_kv,
+        flat_work_group_size,
+    )
+    with _KERNEL_CACHE_LOCK:
+        if key in _KERNEL_CACHE:
+            return _KERNEL_CACHE[key]
+        launch = build_v4_swa_fwd_module(
+            num_heads=num_heads_q,
+            head_dim=head_dim,
+            swa_window=int(swa_window),
+            dtype_str=dtype_str,
+            waves_per_eu=waves_per_eu,
+            block_m=block_m,
+            block_n=block_n,
+            flat_work_group_size=flat_work_group_size,
+            layout_bhld=True,
+            mqa_kv=mqa_kv,
+        )
+        _KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _launch_v4_attention_fwd_flydsl_mqa(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],
+    scale: float,
+    hca_local_seqlen: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """SWA-only launcher that AVOIDS the MQA broadcast.
+
+    For K_H=1 (MQA), passes the [B, 1, Sk, D] tensor directly and uses
+    the mqa_kv kernel variant. For K_H=HQ (non-MQA), falls back to the
+    standard kernel.
+    """
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError("rank-4 q/k/v required")
+    B, HQ, Sq, D = q.shape
+    Bk, HK, Sk, Dk = k.shape
+    if (Bk, Sk, Dk) != (B, Sk, D) or v.shape != k.shape:
+        raise ValueError(f"shape mismatch")
+    if HK != 1 and HK != HQ:
+        raise ValueError(f"K_H must be 1 or {HQ}; got {HK}")
+    if q.dtype != torch.bfloat16:
+        raise NotImplementedError(f"bf16 only")
+    if D != 512:
+        raise NotImplementedError(f"head_dim=512 only")
+    if sink is not None:
+        raise NotImplementedError("sink not yet supported in MQA wrapper")
+    if additive_mask is not None or int(hca_local_seqlen) != 0:
+        raise NotImplementedError("HCA mode not in this wrapper")
+    if int(swa_window) <= 0:
+        raise NotImplementedError("swa_window > 0 required")
+    if Sq != Sk:
+        raise NotImplementedError("SWA: Sq==Sk required")
+    expected_scale = 1.0 / math.sqrt(D)
+    if not math.isclose(float(scale), expected_scale, rel_tol=1e-4):
+        raise NotImplementedError(f"scale=1/sqrt(D) only")
+
+    mqa = (HK == 1) and (HQ > 1)
+
+    q_bhld = q.contiguous()
+    if mqa:
+        # K/V are [B, 1, Sk, D]. We pass them flat-view but use mqa_kv kernel
+        # variant that drops head_idx from K/V indexing.
+        k_bhld = k.contiguous()
+        v_bhld = v.contiguous()
+    else:
+        k_bhld = k.contiguous()
+        v_bhld = v.contiguous()
+    o_bhld = torch.empty_like(q_bhld)
+    lse = torch.zeros((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    block_m = int(os.environ.get("PRIMUS_V4_FLYDSL_BLOCK_M", "128"))
+    block_n = int(os.environ.get("PRIMUS_V4_FLYDSL_BLOCK_N", "32"))
+    waves_per_eu = int(os.environ.get("FLYDSL_WAVES_PER_EU", "2"))
+    fwgs_env = os.environ.get("PRIMUS_V4_FLYDSL_FWGS", "")
+    flat_work_group_size = int(fwgs_env) if fwgs_env else None
+    launch = _get_kernel(
+        HQ, D, int(swa_window), "bf16", block_m, block_n, waves_per_eu, mqa, flat_work_group_size
+    )
+
+    launch(
+        q_bhld.view(-1),
+        k_bhld.view(-1),
+        v_bhld.view(-1),
+        o_bhld.view(-1),
+        lse.view(-1),
+        B,
+        Sq,
+    )
+    return o_bhld, lse
+
+
+__all__ = ["_launch_v4_attention_fwd_flydsl_mqa"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py
new file mode 100644
index 000000000..d01d0ebde
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py
@@ -0,0 +1,405 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention backward FlyDSL launcher (Phase B STEP 3a + 3b).
+
+Env knobs:
+  V4_FLYDSL_CSA_BWD_FLY_DQ   default 0 (1 = FlyDSL dq, 0 = Triton dq)
+  V4_FLYDSL_CSA_BWD_FLY_DKV  default 0 (1 = FlyDSL dk_local/dv_local/
+                              dgathered/dsink, 0 = Triton)
+  V4_FLYDSL_BWD_VERBOSE      default 0
+
+Wiring:
+  knob_dq=0, knob_dkv=0   -> all Triton.
+  knob_dq=1, knob_dkv=0   -> FlyDSL dq-only; Triton emits all others.
+  knob_dq=0, knob_dkv=1   -> Triton dq; FlyDSL emits dk/dv/dgathered/dsink.
+                            (rare path; uses the FULL FlyDSL kernel and
+                            discards its dq.)
+  knob_dq=1, knob_dkv=1   -> ONE FlyDSL launch produces all 5 grads
+                            (no Triton call needed; the cheapest path
+                            when both are on).
+
+The launcher always emits a provenance line so the harness can prove the
+expected backend was used.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+from typing import Optional, Tuple
+
+import torch
+
+_FLYDSL_SRC = "/workspace/FlyDSL-amd"
+if _FLYDSL_SRC not in sys.path:
+    sys.path.insert(0, _FLYDSL_SRC)
+
+os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2")
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+if "/workspace/Primus" not in sys.path:
+    sys.path.insert(0, "/workspace/Primus")
+
+import triton  # noqa: E402
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_bwd import (  # noqa: E402
+    _launch_v4_csa_attention_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (  # noqa: E402
+    _v4_attention_bwd_preprocess_kernel,
+)
+
+_DQ_KERNEL_CACHE = {}
+_DQ_KERNEL_LOCK = threading.Lock()
+_FULL_KERNEL_CACHE = {}
+_FULL_KERNEL_LOCK = threading.Lock()
+
+
+def _get_fly_csa_dq(num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k):
+    key = (num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k)
+    with _DQ_KERNEL_LOCK:
+        if key in _DQ_KERNEL_CACHE:
+            return _DQ_KERNEL_CACHE[key]
+        from v4_csa_bwd_dq_kernel import build_v4_csa_bwd_dq_module
+
+        launch = build_v4_csa_bwd_dq_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            swa_window=swa_window,
+            dtype_str=dtype_str,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            block_n=block_n,
+            block_k=block_k,
+            mqa_kv=True,
+        )
+        _DQ_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _get_fly_csa_full(num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k):
+    key = (num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k)
+    with _FULL_KERNEL_LOCK:
+        if key in _FULL_KERNEL_CACHE:
+            return _FULL_KERNEL_CACHE[key]
+        from v4_csa_bwd_full_kernel import build_v4_csa_bwd_full_module
+
+        launch = build_v4_csa_bwd_full_module(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            swa_window=swa_window,
+            dtype_str=dtype_str,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            block_n=block_n,
+            block_k=block_k,
+            mqa_kv=True,
+        )
+        _FULL_KERNEL_CACHE[key] = launch
+        return launch
+
+
+def _preprocess_deltas(out, dout, B, HQ, Sq, D):
+    BLOCK_M_PRE = 64
+    d_buf = torch.empty((B, HQ, Sq), device=out.device, dtype=torch.float32)
+    pre_grid = (triton.cdiv(Sq, BLOCK_M_PRE), B * HQ)
+    _v4_attention_bwd_preprocess_kernel[pre_grid](
+        out,
+        dout,
+        d_buf,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        Sq,
+        HEAD=HQ,
+        BLOCK_M=BLOCK_M_PRE,
+        BLOCK_DMODEL=D,
+        num_warps=4,
+        num_stages=1,
+    )
+    return d_buf
+
+
+def flydsl_v4_csa_attention_bwd(
+    q: torch.Tensor,
+    k_local: torch.Tensor,
+    v_local: torch.Tensor,
+    gathered: torch.Tensor,
+    sparse_mask: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    scale: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4:
+        raise ValueError("rank-4 q / k_local / v_local required")
+    if dout.shape != out.shape or out.shape != q.shape:
+        raise ValueError(
+            f"shape mismatch: q={tuple(q.shape)} out={tuple(out.shape)} dout={tuple(dout.shape)}"
+        )
+    B, HQ, Sq, D = q.shape
+    K_topk = gathered.shape[2]
+    has_sink = sink is not None
+
+    use_fly_dq = os.getenv("V4_FLYDSL_CSA_BWD_FLY_DQ", "0") == "1"
+    use_fly_dkv = os.getenv("V4_FLYDSL_CSA_BWD_FLY_DKV", "0") == "1"
+    verbose = os.getenv("V4_FLYDSL_BWD_VERBOSE", "0") == "1"
+
+    dq_tag = "FlyDSL" if use_fly_dq else "Triton"
+    dkv_tag = "FlyDSL" if use_fly_dkv else "Triton"
+    dsink_tag = (
+        "FlyDSL" if (use_fly_dkv and has_sink) else ("FlyDSL" if (use_fly_dq and has_sink) else "Triton")
+    )
+    print(
+        f"[v4_csa_bwd] provenance: dq={dq_tag} dk_local={dkv_tag} dv_local={dkv_tag} "
+        f"dgathered={dkv_tag} dsink={dsink_tag} "
+        f"B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} D={D} swa={swa_window} "
+        f"has_sink={has_sink}",
+        flush=True,
+    )
+
+    # ============================================================
+    # Path 1: nothing FlyDSL -> direct Triton pass-through.
+    # ============================================================
+    if not use_fly_dq and not use_fly_dkv:
+        return _launch_v4_csa_attention_bwd(
+            q,
+            k_local,
+            v_local,
+            gathered,
+            sparse_mask,
+            out,
+            dout,
+            lse,
+            sink=sink,
+            swa_window=int(swa_window),
+            scale=float(scale),
+        )
+
+    # ============================================================
+    # Validate the FlyDSL kernel's preconditions.
+    # ============================================================
+    if q.dtype != torch.bfloat16:
+        raise NotImplementedError(f"FlyDSL CSA bwd supports bf16 only; got {q.dtype}")
+    if D % 64 != 0:
+        raise NotImplementedError(f"FlyDSL CSA bwd requires D % 64 == 0; got D={D}")
+
+    # MQA view -- the Triton wrapper expands MQA->MHA by .expand().contiguous(),
+    # so head 0 has the original values.
+    k_mqa = k_local[:, :1, :, :].contiguous()
+    v_mqa = v_local[:, :1, :, :].contiguous()
+
+    q_c = q.contiguous()
+    dout_c = dout.contiguous()
+    lse_c = lse.contiguous()
+    gathered_c = gathered.contiguous()
+    sparse_mask_c = sparse_mask.contiguous()
+
+    d_buf = _preprocess_deltas(out, dout_c, B, HQ, Sq, D)
+
+    if has_sink:
+        sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink.contiguous()
+    else:
+        sink_arg = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+
+    block_n = 32
+    block_k = 32
+    dtype_str = "bf16"
+    has_sparse = K_topk > 0
+
+    # ============================================================
+    # Path 2: BOTH knobs on -- one FlyDSL full-kernel launch.
+    # ============================================================
+    if use_fly_dq and use_fly_dkv:
+        dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        # dk_local / dv_local are stored at [B, HQ, Sq, D] (MHA-shape buffer
+        # matching Triton's atomic-add target). Caller is free to reduce.
+        dkl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dvl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32)
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+
+        launch = _get_fly_csa_full(
+            num_heads=HQ,
+            head_dim=D,
+            swa_window=int(swa_window),
+            dtype_str=dtype_str,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            block_n=block_n,
+            block_k=block_k,
+        )
+
+        if verbose:
+            print(
+                f"[v4_csa_bwd] FlyDSL FULL launch: B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} "
+                f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}",
+                flush=True,
+            )
+
+        launch(
+            q_c,
+            k_mqa,
+            v_mqa,
+            gathered_c,
+            sparse_mask_c,
+            dout_c,
+            lse_c,
+            d_buf,
+            sink_arg,
+            dq_fp32,
+            dkl_fp32,
+            dvl_fp32,
+            dgathered_fp32,
+            dsink_fp32,
+            int(B),
+            int(Sq),
+            int(K_topk),
+        )
+
+        dq_out = dq_fp32.to(q.dtype)
+        # dk_local / dv_local return shape [B, HQ, Sq, D] to match the Triton
+        # output (it returns the broadcast-MHA buffer; bwd_modes._run_csa_pair
+        # does the autograd reduction itself when the leaves were created
+        # via .expand().contiguous()).
+        dkl_out = dkl_fp32.to(k_local.dtype)
+        dvl_out = dvl_fp32.to(v_local.dtype)
+        dg_out = dgathered_fp32.to(gathered.dtype)
+        dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+        return dq_out, dkl_out, dvl_out, dg_out, dsink_out
+
+    # ============================================================
+    # Path 3: partial overlap with Triton.
+    # ============================================================
+    # For partial paths we always run the Triton kernel and selectively
+    # override its outputs with FlyDSL values.
+    triton_out = _launch_v4_csa_attention_bwd(
+        q,
+        k_local,
+        v_local,
+        gathered,
+        sparse_mask,
+        out,
+        dout,
+        lse,
+        sink=sink,
+        swa_window=int(swa_window),
+        scale=float(scale),
+    )
+    dq_t, dkl_t, dvl_t, dg_t, dsink_t = triton_out
+
+    if use_fly_dq:
+        # Override dq (+ dsink for sink case) with FlyDSL dq-only kernel.
+        dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        launch = _get_fly_csa_dq(
+            num_heads=HQ,
+            head_dim=D,
+            swa_window=int(swa_window),
+            dtype_str=dtype_str,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            block_n=block_n,
+            block_k=block_k,
+        )
+        if verbose:
+            print(
+                f"[v4_csa_bwd] FlyDSL dq launch: B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} "
+                f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}",
+                flush=True,
+            )
+        launch(
+            q_c,
+            k_mqa,
+            v_mqa,
+            gathered_c,
+            sparse_mask_c,
+            dout_c,
+            lse_c,
+            d_buf,
+            sink_arg,
+            dq_fp32,
+            dsink_fp32,
+            int(B),
+            int(Sq),
+            int(K_topk),
+        )
+        dq_out = dq_fp32.to(q.dtype)
+        dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+        return dq_out, dkl_t, dvl_t, dg_t, dsink_out
+
+    if use_fly_dkv:
+        # Run the FULL kernel; we keep dk/dv/dgathered/dsink from FlyDSL and
+        # discard its dq (Triton's dq is the reference here).
+        dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dkl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dvl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32)
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+
+        launch = _get_fly_csa_full(
+            num_heads=HQ,
+            head_dim=D,
+            swa_window=int(swa_window),
+            dtype_str=dtype_str,
+            has_sink=has_sink,
+            has_sparse=has_sparse,
+            block_n=block_n,
+            block_k=block_k,
+        )
+
+        if verbose:
+            print(
+                f"[v4_csa_bwd] FlyDSL DKV launch (dq discarded): B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} "
+                f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}",
+                flush=True,
+            )
+
+        launch(
+            q_c,
+            k_mqa,
+            v_mqa,
+            gathered_c,
+            sparse_mask_c,
+            dout_c,
+            lse_c,
+            d_buf,
+            sink_arg,
+            dq_fp32,
+            dkl_fp32,
+            dvl_fp32,
+            dgathered_fp32,
+            dsink_fp32,
+            int(B),
+            int(Sq),
+            int(K_topk),
+        )
+
+        dkl_out = dkl_fp32.to(k_local.dtype)
+        dvl_out = dvl_fp32.to(v_local.dtype)
+        dg_out = dgathered_fp32.to(gathered.dtype)
+        dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+        return dq_t, dkl_out, dvl_out, dg_out, dsink_out
+
+    raise RuntimeError("unreachable")
+
+
+__all__ = ["flydsl_v4_csa_attention_bwd"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py
new file mode 100644
index 000000000..a3d72eb19
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py
@@ -0,0 +1,665 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_csa_bwd_dq: V4 CSA-backward dq kernel for FlyDSL (STEP 3a).
+
+Per-row design (mirrors v4_csa_fwd_kernel.py):
+  grid       = (Sq, B * HQ)
+  BLOCK_SIZE = 64 (one wave). Each lane owns D_PER_LANE = HEAD_DIM // 64
+                 contiguous d values (=8 for D=512).
+
+Computes the FULL dq for one (b, qhid, q_row) by iterating BOTH
+  - LOCAL SWA stream (k_local / v_local, SWA-window-causal-mask)
+  - GATHERED stream  (gathered / sparse_mask, top-K=K_topk per row)
+The two streams share the saved JOINT LSE from CSA fwd.
+
+Inputs:
+  Q          [B, HQ, Sq, D]   bf16
+  K_LOCAL    [B,  1, Sk, D]   bf16  (MQA, HK=1)
+  V_LOCAL    [B,  1, Sk, D]   bf16
+  GATHERED   [B, Sq, K_topk, D]  bf16
+  SPARSE_MASK[B, Sq, K_topk]     bf16  (additive bias)
+  DOUT       [B, HQ, Sq, D]   bf16
+  LSE        [B, HQ, Sq]      fp32  (JOINT LSE, RAW domain)
+  DELTAS     [B, HQ, Sq]      fp32
+  SINK       [HQ]             fp32  (dummy if !has_sink)
+
+Outputs:
+  DQ         [B, HQ, Sq, D]   fp32  (overwritten -- one program owns the row)
+  DSINK      [HQ]             fp32  (atomic_add)
+"""
+
+from __future__ import annotations
+
+import math
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME = "v4_csa_bwd_dq_kernel"
+_LOG2E = math.log2(math.e)
+_LLVM_GEP_DYNAMIC = -2147483648
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def build_v4_csa_bwd_dq_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    block_n=32,
+    block_k=32,
+    has_sink=True,
+    has_sparse=True,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    mqa_kv=True,
+):
+    """Build the V4 CSA backward dq launcher (one program per (b, h, q_row))."""
+    gpu_arch = get_hip_arch()
+    WARP_SIZE = 64
+    BLOCK_SIZE = WARP_SIZE
+    BLOCK_N = int(block_n)
+    BLOCK_K = int(block_k)
+    NUM_HEADS = int(num_heads)
+    HEAD_DIM = int(head_dim)
+    assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}"
+    D_PER_LANE = HEAD_DIM // WARP_SIZE
+    assert mqa_kv, "v4_csa_bwd_dq only supports MQA (HK=1)"
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM)
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_csa_bwd_dq_smem_N{BLOCK_N}_K{BLOCK_K}_S{int(has_sink)}_HS{int(has_sparse)}",
+    )
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_csa_bwd_dq_kernel(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        DOUT: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        SINK: fx.Tensor,
+        DQ: fx.Tensor,
+        DSINK: fx.Tensor,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        f16_ty = elem_type
+        f32_ty = T.f32
+        fm_fast = arith.FastMathFlags.fast
+
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL)
+        vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL)
+        g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED)
+        sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOUT)
+        dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+        if has_sink:
+            sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True)
+            dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True)
+
+        def _gep_load(base_ptr, elem_idx, vec_type, elem_t):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_t,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_f16_v(base_ptr, elem_idx, n):
+            vt = T.vec(n, f16_ty)
+            return _gep_load(base_ptr, elem_idx, vt, f16_ty)
+
+        # ---- Thread / program IDs ----
+        pid_m = arith.index_cast(T.index, gpu.block_idx.x)
+        pid_bh = arith.index_cast(T.index, gpu.block_idx.y)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+        lane = tid
+
+        seq_len_v = arith.index_cast(T.index, seq_len)
+        K_topk_v = arith.index_cast(T.index, K_topk)
+
+        bid = pid_bh // arith.index(NUM_HEADS)
+        qhid = pid_bh % arith.index(NUM_HEADS)
+
+        q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v)
+        pid_m_safe = arith.select(q_active, pid_m, arith.index(0))
+
+        NEG_INF_F = -1.0e30
+        c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty)
+        c_zero_f = arith.constant(0.0, type=f32_ty)
+        c_sm_scale = arith.constant(float(sm_scale), type=f32_ty)
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+
+        zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty))
+        q_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(HEAD_DIM)
+        q_lane_off = q_row_base + lane * arith.index(D_PER_LANE)
+        q_vec_raw = load_f16_v(q_ptr, q_lane_off, D_PER_LANE)
+        q_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec_raw)
+        q_f32 = arith.select(q_active, q_f32, zero_f32_vec)
+
+        do_lane_off = q_lane_off
+        do_vec_raw = load_f16_v(do_ptr, do_lane_off, D_PER_LANE)
+        do_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), do_vec_raw)
+        do_f32 = arith.select(q_active, do_f32, zero_f32_vec)
+
+        lse_delta_off = (bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe
+        lse_delta_off_i32 = arith.index_cast(T.i32, lse_delta_off)
+        lse_val = buffer_ops.buffer_load(
+            lse_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=f32_ty,
+        )
+        delta_val = buffer_ops.buffer_load(
+            deltas_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=f32_ty,
+        )
+
+        if has_sink:
+            qhid_i32 = arith.index_cast(T.i32, qhid)
+            sink_h = buffer_ops.buffer_load(
+                sink_rsrc,
+                qhid_i32,
+                vec_width=1,
+                dtype=f32_ty,
+            )
+            sink_h = rocdl.readfirstlane(f32_ty, sink_h)
+            sub_sh = arith.SubFOp(sink_h, lse_val, fastmath=fm_fast).result
+            p_sink = math_dialect.exp(sub_sh, fastmath=fm_fast)
+            neg_p_sink = arith.SubFOp(c_zero_f, p_sink, fastmath=fm_fast).result
+            dsink_contrib = arith.MulFOp(
+                neg_p_sink,
+                delta_val,
+                fastmath=fm_fast,
+            ).result
+            is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0))
+            do_sink = arith.AndIOp(is_lane0, q_active).result
+            _if_sink = scf.IfOp(do_sink, [], has_else=False)
+            with ir.InsertionPoint(_if_sink.then_block):
+                _dsink_byte_off = arith.MulIOp(
+                    qhid_i32,
+                    arith.constant(4, type=T.i32),
+                ).result
+                _zero_i32_atom = arith.constant(0, type=T.i32)
+                rocdl.raw_ptr_buffer_atomic_fadd(
+                    dsink_contrib,
+                    dsink_rsrc,
+                    _dsink_byte_off,
+                    _zero_i32_atom,
+                    _zero_i32_atom,
+                )
+                scf.YieldOp([])
+
+        def warp_reduce_sum_f32(v):
+            cur = v
+            for off in [32, 16, 8, 4, 2, 1]:
+                xor_amt = arith.constant(off, type=T.i32)
+                peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32)
+                cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result
+            return cur
+
+        def vec_dot_f32(a_vec, b_vec):
+            s = c_zero_f
+            for i in range_constexpr(D_PER_LANE):
+                av = vector.extract(a_vec, static_position=[i], dynamic_position=[])
+                bv = vector.extract(b_vec, static_position=[i], dynamic_position=[])
+                p = arith.MulFOp(av, bv, fastmath=fm_fast).result
+                s = arith.AddFOp(s, p, fastmath=fm_fast).result
+            return s
+
+        # ---- LOCAL SWA bounds ----
+        _pid_p1 = pid_m + arith.index(1)
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v)
+        n_loop_end_row = arith.select(_le_seq, _pid_p1, seq_len_v)
+        SWA = arith.index(int(swa_window))
+        _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA)
+        _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0))
+        BN_idx = arith.index(BLOCK_N)
+        n_loop_start = (_n_lo_raw // BN_idx) * BN_idx
+        n_loop_end_blk = ((n_loop_end_row + BN_idx - arith.index(1)) // BN_idx) * BN_idx
+
+        # Pad iter_args to >= 2 elements so scf.for_ always yields a tuple
+        # (the singleton path auto-unwraps to a scalar, breaking subscripting
+        # at D_PER_LANE==1 / D=64).
+        _PAD = 1 if D_PER_LANE == 1 else 0
+        init_local = []
+        for _ in range_constexpr(D_PER_LANE):
+            init_local.append(c_zero_f)
+        for _ in range_constexpr(_PAD):
+            init_local.append(c_zero_f)
+
+        pid_m_i32 = arith.index_cast(T.i32, pid_m)
+        seq_len_i32 = arith.index_cast(T.i32, seq_len_v)
+        w_i32 = arith.constant(int(swa_window), type=T.i32)
+        K_topk_i32 = arith.index_cast(T.i32, K_topk_v)
+
+        # ==== LOCAL SWA loop ====
+        for n_start, inner_args, loop_results_local in scf.for_(
+            n_loop_start,
+            n_loop_end_blk,
+            BN_idx,
+            iter_args=init_local,
+        ):
+            dq_accs = [inner_args[d] for d in range_constexpr(D_PER_LANE)]
+
+            n_start_i32 = arith.index_cast(T.i32, n_start)
+
+            kl_f32_cache = []
+            p_cache = []
+            dp_cache = []
+
+            for n_off in range_constexpr(BLOCK_N):
+                kv_col_i32 = arith.AddIOp(
+                    n_start_i32,
+                    arith.constant(n_off, type=T.i32),
+                ).result
+                _kv_plus_w = arith.AddIOp(kv_col_i32, w_i32).result
+                is_swa = arith.cmpi(
+                    arith.CmpIPredicate.sle,
+                    _kv_plus_w,
+                    pid_m_i32,
+                )
+                is_causal = arith.cmpi(
+                    arith.CmpIPredicate.sgt,
+                    kv_col_i32,
+                    pid_m_i32,
+                )
+                is_oob = arith.cmpi(
+                    arith.CmpIPredicate.sge,
+                    kv_col_i32,
+                    seq_len_i32,
+                )
+                bad = arith.OrIOp(
+                    arith.OrIOp(is_causal, is_swa).result,
+                    is_oob,
+                ).result
+
+                kv_col_idx = arith.index_cast(T.index, kv_col_i32)
+                kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx)
+                kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM)
+                kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE)
+                kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE)
+                kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec)
+                kl_f32_cache.append(kl_f32)
+
+                vl_vec = load_f16_v(vl_ptr, kl_lane_off, D_PER_LANE)
+                vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec)
+
+                lane_dot_qk = vec_dot_f32(q_f32, kl_f32)
+                qk_full = warp_reduce_sum_f32(lane_dot_qk)
+                qk_scaled = arith.MulFOp(
+                    qk_full,
+                    c_sm_scale,
+                    fastmath=fm_fast,
+                ).result
+                qk_masked = arith.select(bad, c_neg_inf, qk_scaled)
+                diff_qk = arith.SubFOp(
+                    qk_masked,
+                    lse_val,
+                    fastmath=fm_fast,
+                ).result
+                p = math_dialect.exp(diff_qk, fastmath=fm_fast)
+                p_cache.append(p)
+
+                lane_dot_dp = vec_dot_f32(do_f32, vl_f32)
+                dp_full = warp_reduce_sum_f32(lane_dot_dp)
+                dp_cache.append(dp_full)
+
+            for n_off in range_constexpr(BLOCK_N):
+                p = p_cache[n_off]
+                dp = dp_cache[n_off]
+                diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result
+                ds = arith.MulFOp(p, diff, fastmath=fm_fast).result
+                ds_scaled = arith.MulFOp(
+                    ds,
+                    c_sm_scale,
+                    fastmath=fm_fast,
+                ).result
+                kl_f32 = kl_f32_cache[n_off]
+                for d_off in range_constexpr(D_PER_LANE):
+                    klv = vector.extract(
+                        kl_f32,
+                        static_position=[d_off],
+                        dynamic_position=[],
+                    )
+                    contrib = arith.MulFOp(
+                        ds_scaled,
+                        klv,
+                        fastmath=fm_fast,
+                    ).result
+                    dq_accs[d_off] = arith.AddFOp(
+                        dq_accs[d_off],
+                        contrib,
+                        fastmath=fm_fast,
+                    ).result
+
+            _yield = list(dq_accs)
+            for _ in range_constexpr(_PAD):
+                _yield.append(c_zero_f)
+            yield _yield
+
+        dq_accs = [loop_results_local[d] for d in range_constexpr(D_PER_LANE)]
+
+        # ==== GATHERED branch ====
+        if has_sparse:
+            init_sparse = list(dq_accs)
+            for _ in range_constexpr(_PAD):
+                init_sparse.append(c_zero_f)
+            for k_start, inner_args_g, loop_results_g in scf.for_(
+                arith.index(0),
+                K_topk_v,
+                arith.index(BLOCK_K),
+                iter_args=init_sparse,
+            ):
+                dq_accs_g = [inner_args_g[d] for d in range_constexpr(D_PER_LANE)]
+
+                k_start_i32 = arith.index_cast(T.i32, k_start)
+
+                g_f32_cache = []
+                p_cache_g = []
+                dp_cache_g = []
+
+                for k_off in range_constexpr(BLOCK_K):
+                    k_pos_i32 = arith.AddIOp(
+                        k_start_i32,
+                        arith.constant(k_off, type=T.i32),
+                    ).result
+                    is_oob = arith.cmpi(
+                        arith.CmpIPredicate.sge,
+                        k_pos_i32,
+                        K_topk_i32,
+                    )
+                    k_pos_idx = arith.index_cast(T.index, k_pos_i32)
+                    k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx)
+
+                    g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index(
+                        HEAD_DIM
+                    )
+                    g_lane_off = g_row_base + lane * arith.index(D_PER_LANE)
+                    g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE)
+                    g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec)
+                    g_f32_cache.append(g_f32)
+
+                    sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe
+                    sm_raw_v1 = _gep_load(
+                        sm_ptr,
+                        sm_off,
+                        T.vec(1, elem_type),
+                        elem_type,
+                    )
+                    sm_raw = vector.extract(
+                        sm_raw_v1,
+                        static_position=[0],
+                        dynamic_position=[],
+                    )
+                    sm_val = arith.extf(f32_ty, sm_raw)
+                    sm_val = arith.select(is_oob, c_zero_f, sm_val)
+
+                    lane_dot_qk = vec_dot_f32(q_f32, g_f32)
+                    qk_full = warp_reduce_sum_f32(lane_dot_qk)
+                    qk_scaled = arith.MulFOp(
+                        qk_full,
+                        c_sm_scale,
+                        fastmath=fm_fast,
+                    ).result
+                    qk_biased = arith.AddFOp(
+                        qk_scaled,
+                        sm_val,
+                        fastmath=fm_fast,
+                    ).result
+                    bad = arith.OrIOp(
+                        is_oob,
+                        arith.cmpi(
+                            arith.CmpIPredicate.sge,
+                            pid_m_i32,
+                            seq_len_i32,
+                        ),
+                    ).result
+                    qk_masked = arith.select(bad, c_neg_inf, qk_biased)
+                    diff_qk = arith.SubFOp(
+                        qk_masked,
+                        lse_val,
+                        fastmath=fm_fast,
+                    ).result
+                    p = math_dialect.exp(diff_qk, fastmath=fm_fast)
+                    p_cache_g.append(p)
+
+                    lane_dot_dp = vec_dot_f32(do_f32, g_f32)
+                    dp_full = warp_reduce_sum_f32(lane_dot_dp)
+                    dp_cache_g.append(dp_full)
+
+                for k_off in range_constexpr(BLOCK_K):
+                    p = p_cache_g[k_off]
+                    dp = dp_cache_g[k_off]
+                    diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result
+                    ds = arith.MulFOp(p, diff, fastmath=fm_fast).result
+                    ds_scaled = arith.MulFOp(
+                        ds,
+                        c_sm_scale,
+                        fastmath=fm_fast,
+                    ).result
+                    g_f32 = g_f32_cache[k_off]
+                    for d_off in range_constexpr(D_PER_LANE):
+                        gv = vector.extract(
+                            g_f32,
+                            static_position=[d_off],
+                            dynamic_position=[],
+                        )
+                        contrib = arith.MulFOp(
+                            ds_scaled,
+                            gv,
+                            fastmath=fm_fast,
+                        ).result
+                        dq_accs_g[d_off] = arith.AddFOp(
+                            dq_accs_g[d_off],
+                            contrib,
+                            fastmath=fm_fast,
+                        ).result
+
+                _yield_g = list(dq_accs_g)
+                for _ in range_constexpr(_PAD):
+                    _yield_g.append(c_zero_f)
+                yield _yield_g
+
+            dq_accs = [loop_results_g[d] for d in range_constexpr(D_PER_LANE)]
+
+        # ==== Store dq[d] into DQ buffer (fp32 direct store) ====
+        _o_guard = scf.IfOp(q_active, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            dq_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(
+                HEAD_DIM
+            )
+            dq_lane_off = dq_row_base + lane * arith.index(D_PER_LANE)
+            for d_off in range_constexpr(D_PER_LANE):
+                elem_off = dq_lane_off + arith.index(d_off)
+                _gep_store_f32(dq_accs[d_off], dq_ptr, elem_off)
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_csa_bwd_dq(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        DOUT: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        SINK: fx.Tensor,
+        DQ: fx.Tensor,
+        DSINK: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len)
+        grid_x = sl_idx
+        grid_y = bs_idx * arith.index(NUM_HEADS)
+
+        launcher = v4_csa_bwd_dq_kernel(
+            Q,
+            K_LOCAL,
+            V_LOCAL,
+            GATHERED,
+            SPARSE_MASK,
+            DOUT,
+            LSE,
+            DELTAS,
+            SINK,
+            DQ,
+            DSINK,
+            seq_len,
+            K_topk,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(
+                            T.i32,
+                            _wpe,
+                        )
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, grid_y, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(compile_hints):
+            return launch_v4_csa_bwd_dq(*args, **kwargs)
+
+    def _compile(
+        Q,
+        K_LOCAL,
+        V_LOCAL,
+        GATHERED,
+        SPARSE_MASK,
+        DOUT,
+        LSE,
+        DELTAS,
+        SINK,
+        DQ,
+        DSINK,
+        batch_size,
+        seq_len,
+        K_topk,
+        stream=None,
+    ):
+        with CompilationContext.compile_hints(compile_hints):
+            return flyc.compile(
+                launch_v4_csa_bwd_dq,
+                Q,
+                K_LOCAL,
+                V_LOCAL,
+                GATHERED,
+                SPARSE_MASK,
+                DOUT,
+                LSE,
+                DELTAS,
+                SINK,
+                DQ,
+                DSINK,
+                batch_size,
+                seq_len,
+                K_topk,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+    return _launch
+
+
+build_v4_csa_bwd_dq_module_primary = build_v4_csa_bwd_dq_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py
new file mode 100644
index 000000000..010a2a64b
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py
@@ -0,0 +1,666 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_csa_bwd_full: V4 CSA backward kernel, full output set (STEP 3b).
+
+Emits dq + dk_local + dv_local + dgathered + dsink in one launch.
+Mirrors `_v4_csa_attention_bwd_kernel` (Triton, _triton/v4_csa_attention_bwd.py)
+1:1: grid=(Sq, B*HQ); each program owns one query row. dq is direct-stored;
+dk_local, dv_local, dgathered, dsink are accumulated via atomic_fadd.
+"""
+from __future__ import annotations
+
+import math
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import (
+    arith,
+    buffer_ops,
+    const_expr,
+    gpu,
+    range_constexpr,
+    rocdl,
+    vector,
+)
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME = "v4_csa_bwd_full_kernel"
+_LLVM_GEP_DYNAMIC = -2147483648
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def build_v4_csa_bwd_full_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    block_n=32,
+    block_k=32,
+    has_sink=True,
+    has_sparse=True,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    mqa_kv=True,
+):
+    gpu_arch = get_hip_arch()
+    WARP_SIZE = 64
+    BLOCK_SIZE = WARP_SIZE
+    BLOCK_N = int(block_n)
+    BLOCK_K = int(block_k)
+    NUM_HEADS = int(num_heads)
+    HEAD_DIM = int(head_dim)
+    assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}"
+    D_PER_LANE = HEAD_DIM // WARP_SIZE
+    assert mqa_kv, "v4_csa_bwd_full only supports MQA (HK=1)"
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM)
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_csa_bwd_full_smem_N{BLOCK_N}_K{BLOCK_K}_S{int(has_sink)}_HS{int(has_sparse)}",
+    )
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_csa_bwd_full_kernel(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        DOUT: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        SINK: fx.Tensor,
+        DQ: fx.Tensor,
+        DK_LOCAL: fx.Tensor,
+        DV_LOCAL: fx.Tensor,
+        DGATHERED: fx.Tensor,
+        DSINK: fx.Tensor,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        # FlyDSL >=0.2.2 compat: dtype_to_elem_type returns a Numeric meta
+        # (e.g. fx.BFloat16); the MLIR type-arg sites below (T.vec, GEPOp,
+        # SmemPtr, trunc_f) require an ir.Type. Coerce once here.
+        if hasattr(elem_type, "ir_type"):
+            elem_type = elem_type.ir_type
+        f16_ty = elem_type
+        f32_ty = T.f32
+        fm_fast = arith.FastMathFlags.fast
+
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL)
+        vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL)
+        g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED)
+        sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOUT)
+        dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+        dkl_rsrc = buffer_ops.create_buffer_resource(DK_LOCAL, max_size=True)
+        dvl_rsrc = buffer_ops.create_buffer_resource(DV_LOCAL, max_size=True)
+        dg_rsrc = buffer_ops.create_buffer_resource(DGATHERED, max_size=True)
+        if const_expr(has_sink):
+            sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True)
+            dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True)
+
+        def _gep_load(base_ptr, elem_idx, vec_type, elem_t):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_t,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_f16_v(base_ptr, elem_idx, n):
+            vt = T.vec(n, f16_ty)
+            return _gep_load(base_ptr, elem_idx, vt, f16_ty)
+
+        pid_m = arith.index_cast(T.index, gpu.block_idx.x)
+        pid_bh = arith.index_cast(T.index, gpu.block_idx.y)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+        lane = tid
+
+        seq_len_v = arith.index_cast(T.index, seq_len)
+        K_topk_v = arith.index_cast(T.index, K_topk)
+
+        bid = pid_bh // arith.index(NUM_HEADS)
+        qhid = pid_bh % arith.index(NUM_HEADS)
+
+        q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v)
+        pid_m_safe = arith.select(q_active, pid_m, arith.index(0))
+
+        NEG_INF_F = -1.0e30
+        c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty)
+        c_zero_f = arith.constant(0.0, type=f32_ty)
+        c_sm_scale = arith.constant(float(sm_scale), type=f32_ty)
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+        c_four_i32 = arith.constant(4, type=T.i32)
+        c_zero_i32 = arith.constant(0, type=T.i32)
+
+        zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty))
+        q_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(HEAD_DIM)
+        q_lane_off = q_row_base + lane * arith.index(D_PER_LANE)
+        q_vec_raw = load_f16_v(q_ptr, q_lane_off, D_PER_LANE)
+        q_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec_raw)
+        q_f32 = arith.select(q_active, q_f32, zero_f32_vec)
+
+        do_lane_off = q_lane_off
+        do_vec_raw = load_f16_v(do_ptr, do_lane_off, D_PER_LANE)
+        do_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), do_vec_raw)
+        do_f32 = arith.select(q_active, do_f32, zero_f32_vec)
+
+        lse_delta_off = (bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe
+        lse_delta_off_i32 = arith.index_cast(T.i32, lse_delta_off)
+        lse_val = buffer_ops.buffer_load(
+            lse_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=f32_ty,
+        )
+        delta_val = buffer_ops.buffer_load(
+            deltas_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=f32_ty,
+        )
+
+        qhid_i32 = arith.index_cast(T.i32, qhid)
+        bid_i32 = arith.index_cast(T.i32, bid)
+        lane_i32 = arith.index_cast(T.i32, lane)
+        pid_m_safe_i32 = arith.index_cast(T.i32, pid_m_safe)
+        head_dim_i32 = arith.constant(HEAD_DIM, type=T.i32)
+        d_per_lane_i32 = arith.constant(D_PER_LANE, type=T.i32)
+        num_heads_i32 = arith.constant(NUM_HEADS, type=T.i32)
+
+        if const_expr(has_sink):
+            sink_h = buffer_ops.buffer_load(
+                sink_rsrc,
+                qhid_i32,
+                vec_width=1,
+                dtype=f32_ty,
+            )
+            sink_h = rocdl.readfirstlane(f32_ty, sink_h)
+            sub_sh = arith.SubFOp(sink_h, lse_val, fastmath=fm_fast).result
+            p_sink = math_dialect.exp(sub_sh, fastmath=fm_fast)
+            neg_p_sink = arith.SubFOp(c_zero_f, p_sink, fastmath=fm_fast).result
+            dsink_contrib = arith.MulFOp(
+                neg_p_sink,
+                delta_val,
+                fastmath=fm_fast,
+            ).result
+            is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0))
+            do_sink = arith.AndIOp(is_lane0, q_active).result
+            _if_sink = scf.IfOp(do_sink, [], has_else=False)
+            with ir.InsertionPoint(_if_sink.then_block):
+                _dsink_byte_off = arith.MulIOp(qhid_i32, c_four_i32).result
+                rocdl.raw_ptr_buffer_atomic_fadd(
+                    dsink_contrib,
+                    dsink_rsrc,
+                    _dsink_byte_off,
+                    c_zero_i32,
+                    c_zero_i32,
+                )
+                scf.YieldOp([])
+
+        def warp_reduce_sum_f32(v):
+            cur = v
+            for off in [32, 16, 8, 4, 2, 1]:
+                xor_amt = arith.constant(off, type=T.i32)
+                peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32)
+                cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result
+            return cur
+
+        def vec_dot_f32(a_vec, b_vec):
+            s = c_zero_f
+            for i in range_constexpr(D_PER_LANE):
+                av = vector.extract(a_vec, static_position=[i], dynamic_position=[])
+                bv = vector.extract(b_vec, static_position=[i], dynamic_position=[])
+                p = arith.MulFOp(av, bv, fastmath=fm_fast).result
+                s = arith.AddFOp(s, p, fastmath=fm_fast).result
+            return s
+
+        _pid_p1 = pid_m + arith.index(1)
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v)
+        n_loop_end_row = arith.select(_le_seq, _pid_p1, seq_len_v)
+        SWA = arith.index(int(swa_window))
+        _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA)
+        _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0))
+        BN_idx = arith.index(BLOCK_N)
+        n_loop_start = (_n_lo_raw // BN_idx) * BN_idx
+        n_loop_end_blk = ((n_loop_end_row + BN_idx - arith.index(1)) // BN_idx) * BN_idx
+
+        _PAD = 1 if D_PER_LANE == 1 else 0
+        init_local = []
+        for _ in range_constexpr(D_PER_LANE):
+            init_local.append(c_zero_f)
+        for _ in range_constexpr(_PAD):
+            init_local.append(c_zero_f)
+
+        pid_m_i32 = arith.index_cast(T.i32, pid_m)
+        seq_len_i32 = arith.index_cast(T.i32, seq_len_v)
+        w_i32 = arith.constant(int(swa_window), type=T.i32)
+        K_topk_i32 = arith.index_cast(T.i32, K_topk_v)
+
+        # ==== LOCAL SWA loop ====
+        for n_start, inner_args, loop_results_local in scf.for_(
+            n_loop_start,
+            n_loop_end_blk,
+            BN_idx,
+            iter_args=init_local,
+        ):
+            dq_accs = [inner_args[d] for d in range_constexpr(D_PER_LANE)]
+            n_start_i32 = arith.index_cast(T.i32, n_start)
+
+            kl_f32_cache = []
+            p_cache = []
+            dp_cache = []
+            kv_col_i32_cache = []
+
+            for n_off in range_constexpr(BLOCK_N):
+                kv_col_i32 = arith.AddIOp(
+                    n_start_i32,
+                    arith.constant(n_off, type=T.i32),
+                ).result
+                kv_col_i32_cache.append(kv_col_i32)
+                _kv_plus_w = arith.AddIOp(kv_col_i32, w_i32).result
+                is_swa = arith.cmpi(arith.CmpIPredicate.sle, _kv_plus_w, pid_m_i32)
+                is_causal = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_i32, pid_m_i32)
+                is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32)
+                bad = arith.OrIOp(
+                    arith.OrIOp(is_causal, is_swa).result,
+                    is_oob,
+                ).result
+
+                kv_col_idx = arith.index_cast(T.index, kv_col_i32)
+                kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx)
+                kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM)
+                kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE)
+                kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE)
+                kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec)
+                kl_f32_cache.append(kl_f32)
+
+                vl_vec = load_f16_v(vl_ptr, kl_lane_off, D_PER_LANE)
+                vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec)
+
+                lane_dot_qk = vec_dot_f32(q_f32, kl_f32)
+                qk_full = warp_reduce_sum_f32(lane_dot_qk)
+                qk_scaled = arith.MulFOp(qk_full, c_sm_scale, fastmath=fm_fast).result
+                qk_masked = arith.select(bad, c_neg_inf, qk_scaled)
+                diff_qk = arith.SubFOp(qk_masked, lse_val, fastmath=fm_fast).result
+                p = math_dialect.exp(diff_qk, fastmath=fm_fast)
+                p_cache.append(p)
+
+                lane_dot_dp = vec_dot_f32(do_f32, vl_f32)
+                dp_full = warp_reduce_sum_f32(lane_dot_dp)
+                dp_cache.append(dp_full)
+
+            for n_off in range_constexpr(BLOCK_N):
+                p = p_cache[n_off]
+                dp = dp_cache[n_off]
+                diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result
+                ds = arith.MulFOp(p, diff, fastmath=fm_fast).result
+                ds_scaled = arith.MulFOp(ds, c_sm_scale, fastmath=fm_fast).result
+                kl_f32 = kl_f32_cache[n_off]
+
+                # dq accumulator
+                for d_off in range_constexpr(D_PER_LANE):
+                    klv = vector.extract(kl_f32, static_position=[d_off], dynamic_position=[])
+                    contrib = arith.MulFOp(ds_scaled, klv, fastmath=fm_fast).result
+                    dq_accs[d_off] = arith.AddFOp(dq_accs[d_off], contrib, fastmath=fm_fast).result
+
+                # dk_local / dv_local atomic_add
+                kv_col_i32 = kv_col_i32_cache[n_off]
+                in_range = arith.cmpi(arith.CmpIPredicate.slt, kv_col_i32, seq_len_i32)
+                do_atom = arith.AndIOp(in_range, q_active).result
+                _if_dkv = scf.IfOp(do_atom, [], has_else=False)
+                with ir.InsertionPoint(_if_dkv.then_block):
+                    _bh = arith.AddIOp(
+                        arith.MulIOp(bid_i32, num_heads_i32).result,
+                        qhid_i32,
+                    ).result
+                    _bh_n = arith.AddIOp(
+                        arith.MulIOp(_bh, seq_len_i32).result,
+                        kv_col_i32,
+                    ).result
+                    _row_d = arith.AddIOp(
+                        arith.MulIOp(_bh_n, head_dim_i32).result,
+                        arith.MulIOp(lane_i32, d_per_lane_i32).result,
+                    ).result
+                    for d_off in range_constexpr(D_PER_LANE):
+                        elem_i32 = arith.AddIOp(_row_d, arith.constant(d_off, type=T.i32)).result
+                        byte_off = arith.MulIOp(elem_i32, c_four_i32).result
+                        qv = vector.extract(q_f32, static_position=[d_off], dynamic_position=[])
+                        dk_val = arith.MulFOp(ds_scaled, qv, fastmath=fm_fast).result
+                        rocdl.raw_ptr_buffer_atomic_fadd(
+                            dk_val,
+                            dkl_rsrc,
+                            byte_off,
+                            c_zero_i32,
+                            c_zero_i32,
+                        )
+                        dov = vector.extract(do_f32, static_position=[d_off], dynamic_position=[])
+                        dv_val = arith.MulFOp(p, dov, fastmath=fm_fast).result
+                        rocdl.raw_ptr_buffer_atomic_fadd(
+                            dv_val,
+                            dvl_rsrc,
+                            byte_off,
+                            c_zero_i32,
+                            c_zero_i32,
+                        )
+                    scf.YieldOp([])
+
+            _yield = list(dq_accs)
+            for _ in range_constexpr(_PAD):
+                _yield.append(c_zero_f)
+            yield _yield
+
+        dq_accs = [loop_results_local[d] for d in range_constexpr(D_PER_LANE)]
+
+        # ==== GATHERED branch ====
+        if const_expr(has_sparse):
+            init_sparse = list(dq_accs)
+            for _ in range_constexpr(_PAD):
+                init_sparse.append(c_zero_f)
+            for k_start, inner_args_g, loop_results_g in scf.for_(
+                arith.index(0),
+                K_topk_v,
+                arith.index(BLOCK_K),
+                iter_args=init_sparse,
+            ):
+                dq_accs_g = [inner_args_g[d] for d in range_constexpr(D_PER_LANE)]
+                k_start_i32 = arith.index_cast(T.i32, k_start)
+
+                g_f32_cache = []
+                p_cache_g = []
+                dp_cache_g = []
+                k_pos_i32_cache = []
+
+                for k_off in range_constexpr(BLOCK_K):
+                    k_pos_i32 = arith.AddIOp(
+                        k_start_i32,
+                        arith.constant(k_off, type=T.i32),
+                    ).result
+                    k_pos_i32_cache.append(k_pos_i32)
+                    is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32)
+                    k_pos_idx = arith.index_cast(T.index, k_pos_i32)
+                    k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx)
+
+                    g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index(
+                        HEAD_DIM
+                    )
+                    g_lane_off = g_row_base + lane * arith.index(D_PER_LANE)
+                    g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE)
+                    g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec)
+                    g_f32_cache.append(g_f32)
+
+                    sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe
+                    sm_raw_v1 = _gep_load(sm_ptr, sm_off, T.vec(1, elem_type), elem_type)
+                    sm_raw = vector.extract(sm_raw_v1, static_position=[0], dynamic_position=[])
+                    sm_val = arith.extf(f32_ty, sm_raw)
+                    sm_val = arith.select(is_oob, c_zero_f, sm_val)
+
+                    lane_dot_qk = vec_dot_f32(q_f32, g_f32)
+                    qk_full = warp_reduce_sum_f32(lane_dot_qk)
+                    qk_scaled = arith.MulFOp(qk_full, c_sm_scale, fastmath=fm_fast).result
+                    qk_biased = arith.AddFOp(qk_scaled, sm_val, fastmath=fm_fast).result
+                    bad = arith.OrIOp(
+                        is_oob,
+                        arith.cmpi(arith.CmpIPredicate.sge, pid_m_i32, seq_len_i32),
+                    ).result
+                    qk_masked = arith.select(bad, c_neg_inf, qk_biased)
+                    diff_qk = arith.SubFOp(qk_masked, lse_val, fastmath=fm_fast).result
+                    p = math_dialect.exp(diff_qk, fastmath=fm_fast)
+                    p_cache_g.append(p)
+
+                    lane_dot_dp = vec_dot_f32(do_f32, g_f32)
+                    dp_full = warp_reduce_sum_f32(lane_dot_dp)
+                    dp_cache_g.append(dp_full)
+
+                for k_off in range_constexpr(BLOCK_K):
+                    p = p_cache_g[k_off]
+                    dp = dp_cache_g[k_off]
+                    diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result
+                    ds = arith.MulFOp(p, diff, fastmath=fm_fast).result
+                    ds_scaled = arith.MulFOp(ds, c_sm_scale, fastmath=fm_fast).result
+                    g_f32 = g_f32_cache[k_off]
+
+                    for d_off in range_constexpr(D_PER_LANE):
+                        gv = vector.extract(g_f32, static_position=[d_off], dynamic_position=[])
+                        contrib = arith.MulFOp(ds_scaled, gv, fastmath=fm_fast).result
+                        dq_accs_g[d_off] = arith.AddFOp(dq_accs_g[d_off], contrib, fastmath=fm_fast).result
+
+                    k_pos_i32 = k_pos_i32_cache[k_off]
+                    in_range_k = arith.cmpi(arith.CmpIPredicate.slt, k_pos_i32, K_topk_i32)
+                    do_atom_k = arith.AndIOp(in_range_k, q_active).result
+                    _if_dg = scf.IfOp(do_atom_k, [], has_else=False)
+                    with ir.InsertionPoint(_if_dg.then_block):
+                        _bm = arith.AddIOp(
+                            arith.MulIOp(bid_i32, seq_len_i32).result,
+                            pid_m_safe_i32,
+                        ).result
+                        _bm_k = arith.AddIOp(
+                            arith.MulIOp(_bm, K_topk_i32).result,
+                            k_pos_i32,
+                        ).result
+                        _row_d = arith.AddIOp(
+                            arith.MulIOp(_bm_k, head_dim_i32).result,
+                            arith.MulIOp(lane_i32, d_per_lane_i32).result,
+                        ).result
+                        for d_off in range_constexpr(D_PER_LANE):
+                            elem_i32 = arith.AddIOp(_row_d, arith.constant(d_off, type=T.i32)).result
+                            byte_off = arith.MulIOp(elem_i32, c_four_i32).result
+                            qv = vector.extract(q_f32, static_position=[d_off], dynamic_position=[])
+                            dov = vector.extract(do_f32, static_position=[d_off], dynamic_position=[])
+                            t1 = arith.MulFOp(ds_scaled, qv, fastmath=fm_fast).result
+                            t2 = arith.MulFOp(p, dov, fastmath=fm_fast).result
+                            dg_val = arith.AddFOp(t1, t2, fastmath=fm_fast).result
+                            rocdl.raw_ptr_buffer_atomic_fadd(
+                                dg_val,
+                                dg_rsrc,
+                                byte_off,
+                                c_zero_i32,
+                                c_zero_i32,
+                            )
+                        scf.YieldOp([])
+
+                _yield_g = list(dq_accs_g)
+                for _ in range_constexpr(_PAD):
+                    _yield_g.append(c_zero_f)
+                yield _yield_g
+
+            dq_accs = [loop_results_g[d] for d in range_constexpr(D_PER_LANE)]
+
+        # ==== Store dq direct ====
+        _o_guard = scf.IfOp(q_active, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            dq_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(
+                HEAD_DIM
+            )
+            dq_lane_off = dq_row_base + lane * arith.index(D_PER_LANE)
+            for d_off in range_constexpr(D_PER_LANE):
+                elem_off = dq_lane_off + arith.index(d_off)
+                _gep_store_f32(dq_accs[d_off], dq_ptr, elem_off)
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_csa_bwd_full(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        DOUT: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        SINK: fx.Tensor,
+        DQ: fx.Tensor,
+        DK_LOCAL: fx.Tensor,
+        DV_LOCAL: fx.Tensor,
+        DGATHERED: fx.Tensor,
+        DSINK: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len)
+        grid_x = sl_idx
+        grid_y = bs_idx * arith.index(NUM_HEADS)
+
+        launcher = v4_csa_bwd_full_kernel(
+            Q,
+            K_LOCAL,
+            V_LOCAL,
+            GATHERED,
+            SPARSE_MASK,
+            DOUT,
+            LSE,
+            DELTAS,
+            SINK,
+            DQ,
+            DK_LOCAL,
+            DV_LOCAL,
+            DGATHERED,
+            DSINK,
+            seq_len,
+            K_topk,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, grid_y, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    compile_hints = {"fast_fp_math": fast_fp_math, "unsafe_fp_math": unsafe_fp_math}
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(compile_hints):
+            return launch_v4_csa_bwd_full(*args, **kwargs)
+
+    def _compile(
+        Q,
+        K_LOCAL,
+        V_LOCAL,
+        GATHERED,
+        SPARSE_MASK,
+        DOUT,
+        LSE,
+        DELTAS,
+        SINK,
+        DQ,
+        DK_LOCAL,
+        DV_LOCAL,
+        DGATHERED,
+        DSINK,
+        batch_size,
+        seq_len,
+        K_topk,
+        stream=None,
+    ):
+        with CompilationContext.compile_hints(compile_hints):
+            return flyc.compile(
+                launch_v4_csa_bwd_full,
+                Q,
+                K_LOCAL,
+                V_LOCAL,
+                GATHERED,
+                SPARSE_MASK,
+                DOUT,
+                LSE,
+                DELTAS,
+                SINK,
+                DQ,
+                DK_LOCAL,
+                DV_LOCAL,
+                DGATHERED,
+                DSINK,
+                batch_size,
+                seq_len,
+                K_topk,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+    return _launch
+
+
+build_v4_csa_bwd_full_module_primary = build_v4_csa_bwd_full_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py
new file mode 100644
index 000000000..b225af233
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py
@@ -0,0 +1,748 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_csa_fwd: V4 CSA forward (FlyDSL per-row design).
+
+Round-3 Step 2b: ports the Triton monolithic CSA forward kernel to FlyDSL
+in a 1:1 per-row design. The Triton CSA monolithic kernel uses one program
+per (b, h, m), with per-key scalar dot-product. That structure ports cleanly:
+
+  grid = (Sq, B * HQ)
+  BLOCK_SIZE = 64 (one wave). Lane in 0..63 handles partial D=8.
+  Online softmax accumulator (m, l, acc) lives in fp32 in-register and is
+    distributed across lanes (each lane owns D/64 = 8 elements of the
+    accumulator's D dimension).
+
+This design does NOT use MFMA -- the per-row scalar dot-product matches
+Triton's monolithic kernel exactly (Triton CSA monolithic also uses
+``tl.sum(k * q, axis=1)`` not ``tl.dot``).
+
+Layout: BHLD (Q/K_local/V_local/O all [B, H, Sq, D]).
+  - Gathered: [B, Sq, K_topk, D] (no H dim -- shared across heads).
+  - sparse_mask: [B, Sq, K_topk] (no H dim -- broadcasts over H).
+  - sink: [H] fp32 or None.
+  - LSE: [B, H, Sq] fp32 (raw-domain: m_final + ln(l_final), since m_final
+    already includes sm_scale via the online softmax in raw-qk*sm_scale).
+
+Forward computes the joint online softmax over local_SWA (block-causal,
+window=swa_window, keys are k_local) + sparse (keys are gathered) + sink.
+
+K_topk == 0 supported (kernel skips the sparse loop).
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import (
+    arith,
+    buffer_ops,
+    const_expr,
+    gpu,
+    range_constexpr,
+    rocdl,
+    vector,
+)
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME = "v4_csa_fwd_kernel"
+
+_LOG2E = math.log2(math.e)
+
+_LLVM_GEP_DYNAMIC = -2147483648
+
+
+def _waitcnt_lgkm_0():
+    """s_waitcnt lgkmcnt(0): drain LDS/SMEM ops. vmcnt=63, expcnt=7."""
+    val = 0xF | (0x7 << 4) | (0 << 8) | (0x3 << 14)
+    rocdl.s_waitcnt(val)
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+def build_v4_csa_fwd_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    block_n=32,
+    block_k=32,
+    has_sink=False,
+    has_sparse=True,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    mqa_kv=False,
+    head_group=1,
+):
+    """Build the V4 CSA forward per-row launcher.
+
+    Parameters:
+      num_heads: int  -- H_Q
+      head_dim: int   -- D (must be divisible by 64)
+      swa_window: int -- SWA window (> 0)
+      block_n: int    -- local-branch tile width (default 32)
+      block_k: int    -- sparse-branch tile width (default 32)
+      has_sink: bool  -- include sink epilogue
+      has_sparse: bool -- include sparse branch loop (K_topk > 0)
+    """
+    gpu_arch = get_hip_arch()
+    WARP_SIZE = 64
+    BLOCK_SIZE = WARP_SIZE  # one wave per program
+    BLOCK_N = int(block_n)
+    BLOCK_K = int(block_k)
+    NUM_HEADS = int(num_heads)
+    HEAD_DIM = int(head_dim)
+    assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}"
+    D_PER_LANE = HEAD_DIM // WARP_SIZE  # 8 for D=512
+    HEAD_GROUP = int(head_group)
+    assert NUM_HEADS % HEAD_GROUP == 0, f"num_heads {NUM_HEADS} must be divisible by head_group {HEAD_GROUP}"
+    NUM_HEAD_GROUPS = NUM_HEADS // HEAD_GROUP
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM)
+
+    # ---- LDS cache for sparse-branch gathered K-tile ----
+    ENABLE_LDS_CACHE = bool(has_sparse) and (os.environ.get("PRIMUS_V4_CSA_LDS_CACHE", "0") == "1")
+    LDS_GATHER_TILE_ELEMS = BLOCK_K * HEAD_DIM
+    LDS_GATHER_TILE_BYTES = LDS_GATHER_TILE_ELEMS * 2
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_csa_fwd_smem_N{BLOCK_N}_K{BLOCK_K}_C{int(ENABLE_LDS_CACHE)}_HG{HEAD_GROUP}",
+    )
+    if ENABLE_LDS_CACHE:
+        lds_gather_offset = allocator._align(allocator.ptr, 16)
+        allocator.ptr = lds_gather_offset + LDS_GATHER_TILE_BYTES
+    else:
+        lds_gather_offset = 0
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_csa_fwd_kernel(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        Sink: fx.Tensor,
+        O: fx.Tensor,
+        LSE: fx.Tensor,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        # FlyDSL >=0.2.2 compat: dtype_to_elem_type returns a Numeric meta
+        # (e.g. fx.BFloat16); the MLIR type-arg sites below (T.vec, GEPOp,
+        # SmemPtr, trunc_f) require an ir.Type. Coerce once here.
+        if hasattr(elem_type, "ir_type"):
+            elem_type = elem_type.ir_type
+        T.f32
+        fm_fast = arith.FastMathFlags.fast
+
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL)
+        vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL)
+        g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED)
+        sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK)
+        o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        if const_expr(has_sink):
+            sink_rsrc = buffer_ops.create_buffer_resource(Sink, max_size=True)
+
+        f16_ty = elem_type
+        f32_ty = T.f32
+
+        # ---- Helpers ----
+        def _gep_load_scalar(base_ptr, elem_idx, vec_type, elem_t):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_t,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def load_f16_v(base_ptr, elem_idx, n):
+            vt = T.vec(n, f16_ty)
+            return _gep_load_scalar(base_ptr, elem_idx, vt, f16_ty)
+
+        def load_f32_scalar(base_ptr, elem_idx):
+            return _gep_load_scalar(base_ptr, elem_idx, f32_ty, f32_ty)
+
+        # ---- Thread / program ----
+        pid_m = arith.index_cast(T.index, gpu.block_idx.x)
+        pid_bh = arith.index_cast(T.index, gpu.block_idx.y)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+        lane = tid
+
+        seq_len_v = arith.index_cast(T.index, seq_len)
+        K_topk_v = arith.index_cast(T.index, K_topk)
+
+        bid = pid_bh // arith.index(NUM_HEAD_GROUPS)
+        qhid_group = pid_bh % arith.index(NUM_HEAD_GROUPS)
+        # qhid_base is the head index of the first head in this program's head group.
+        qhid_base = qhid_group * arith.index(HEAD_GROUP)
+
+        # ---- LDS view for sparse-branch gathered cache ----
+        if ENABLE_LDS_CACHE:
+            base_ptr = allocator.get_base()
+            lds_gather = SmemPtr(
+                base_ptr,
+                lds_gather_offset,
+                elem_type,
+                shape=(LDS_GATHER_TILE_ELEMS,),
+            ).get()
+
+        # ---- Q row in-bounds guard ----
+        q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v)
+        pid_m_safe = arith.select(q_active, pid_m, arith.index(0))
+
+        # ---- Load Q for all HEAD_GROUP heads in this program ----
+        zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty))
+        q_f32_vecs = []
+        for h_off in range_constexpr(HEAD_GROUP):
+            qhid_h = qhid_base + arith.index(h_off)
+            q_row_base = ((bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe) * arith.index(
+                HEAD_DIM
+            )
+            q_lane_off = q_row_base + lane * arith.index(D_PER_LANE)
+            q_vec = load_f16_v(q_ptr, q_lane_off, D_PER_LANE)
+            q_f32_vec = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec)
+            q_f32_vec = arith.select(q_active, q_f32_vec, zero_f32_vec)
+            q_f32_vecs.append(q_f32_vec)
+
+        # ---- Constants ----
+        NEG_INF_F = -1.0e30
+        c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty)
+        c_zero_f = arith.constant(0.0, type=f32_ty)
+        c_one_f = arith.constant(1.0, type=f32_ty)
+        c_sm_scale_f = arith.constant(float(sm_scale), type=f32_ty)
+        c_log2e_f = arith.constant(_LOG2E, type=f32_ty)
+
+        arith.index_cast(T.i32, lane)
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+
+        def warp_reduce_sum_f32(v):
+            cur = v
+            for off in [32, 16, 8, 4, 2, 1]:
+                xor_amt = arith.constant(off, type=T.i32)
+                peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32)
+                cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result
+            return cur
+
+        def vec_dot_f32(a_vec, b_vec):
+            s = c_zero_f
+            for i in range_constexpr(D_PER_LANE):
+                av = vector.extract(a_vec, static_position=[i], dynamic_position=[])
+                bv = vector.extract(b_vec, static_position=[i], dynamic_position=[])
+                p = arith.MulFOp(av, bv, fastmath=fm_fast).result
+                s = arith.AddFOp(s, p, fastmath=fm_fast).result
+            return s
+
+        # ---- Local SWA loop bounds ----
+        _pid_p1 = pid_m + arith.index(1)
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v)
+        n_loop_end = arith.select(_le_seq, _pid_p1, seq_len_v)
+        SWA = arith.index(int(swa_window))
+        _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA)
+        _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0))
+        BN_idx = arith.index(BLOCK_N)
+        n_loop_start = (_n_lo_raw // BN_idx) * BN_idx
+
+        # Init state: HEAD_GROUP copies of (m_i, l_i, acc[D_PER_LANE]).
+        # Layout: [m_0, l_0, acc_0_0..acc_0_{D-1}, m_1, l_1, acc_1_0..]
+        STATE_PER_HEAD = 2 + D_PER_LANE
+        init_args = []
+        for _h in range_constexpr(HEAD_GROUP):
+            init_args.append(c_neg_inf)  # m
+            init_args.append(c_zero_f)  # l
+            for _ in range_constexpr(D_PER_LANE):
+                init_args.append(c_zero_f)
+
+        # ==== LOCAL SWA loop ====
+        for n_start, inner_args, loop_results_local in scf.for_(
+            n_loop_start,
+            n_loop_end,
+            BN_idx,
+            iter_args=init_args,
+        ):
+            # Unpack HEAD_GROUP states from inner_args
+            m_is = [inner_args[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)]
+            l_is = [inner_args[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)]
+            accs = [
+                [inner_args[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)]
+                for h in range_constexpr(HEAD_GROUP)
+            ]
+
+            n_start_i32 = arith.index_cast(T.i32, n_start)
+            pid_m_i32 = arith.index_cast(T.i32, pid_m)
+            seq_len_i32 = seq_len
+            if hasattr(seq_len_i32, "ir_value"):
+                seq_len_i32 = seq_len_i32.ir_value()
+            w_i32 = arith.constant(int(swa_window), type=T.i32)
+
+            # Per-head QK values: [HEAD_GROUP][BLOCK_N]
+            qk_vals_per_head = [[] for _ in range_constexpr(HEAD_GROUP)]
+            kl_f32_cache = []
+            bad_lo_cache = []
+            for n_off in range_constexpr(BLOCK_N):
+                kv_col_i32 = arith.AddIOp(
+                    n_start_i32,
+                    arith.constant(n_off, type=T.i32),
+                ).result
+                _kv_plus_w_lo = arith.AddIOp(kv_col_i32, w_i32).result
+                is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, _kv_plus_w_lo, pid_m_i32)
+                is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_i32, pid_m_i32)
+                is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32)
+                bad_lo = arith.OrIOp(arith.OrIOp(is_causal_lo, is_swa_lo).result, is_oob).result
+                bad_lo_cache.append(bad_lo)
+
+                kv_col_idx = arith.index_cast(T.index, kv_col_i32)
+                kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx)
+                # When mqa_kv, K is shared across heads -> load once.
+                # When mqa_kv=False, K differs per head, so we'd need per-head loads.
+                # For now this kernel requires mqa_kv when head_group > 1.
+                if const_expr(mqa_kv):
+                    kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM)
+                else:
+                    # head_group > 1 with non-MQA: would need per-head load. Disallowed at setup.
+                    kl_row_base = (
+                        (bid * arith.index(NUM_HEADS) + qhid_base) * seq_len_v + kv_col_safe
+                    ) * arith.index(HEAD_DIM)
+                kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE)
+                kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE)
+                kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec)
+                kl_f32_cache.append(kl_f32)
+                # Compute qk for each head using the shared K.
+                for h_off in range_constexpr(HEAD_GROUP):
+                    lane_dot = vec_dot_f32(kl_f32, q_f32_vecs[h_off])
+                    qk_full = warp_reduce_sum_f32(lane_dot)
+                    qk_scaled = arith.MulFOp(qk_full, c_sm_scale_f, fastmath=fm_fast).result
+                    qk_masked = arith.select(bad_lo, c_neg_inf, qk_scaled)
+                    qk_vals_per_head[h_off].append(qk_masked)
+
+            # Per-head softmax update
+            new_m_is = []
+            new_l_is = []
+            new_accs = []
+            p_vals_per_head = []
+            for h_off in range_constexpr(HEAD_GROUP):
+                qk_vals_h = qk_vals_per_head[h_off]
+                m_tile = qk_vals_h[0]
+                for n_off in range_constexpr(BLOCK_N - 1):
+                    m_tile = arith.MaxNumFOp(m_tile, qk_vals_h[n_off + 1], fastmath=fm_fast).result
+                m_new = arith.MaxNumFOp(m_is[h_off], m_tile, fastmath=fm_fast).result
+
+                diff_m = arith.SubFOp(m_is[h_off], m_new, fastmath=fm_fast).result
+                diff_m_log2 = arith.MulFOp(diff_m, c_log2e_f, fastmath=fm_fast).result
+                alpha = arith.ArithValue(diff_m_log2).exp2(fastmath=fm_fast)
+
+                p_vals = []
+                tile_sum = c_zero_f
+                for n_off in range_constexpr(BLOCK_N):
+                    d = arith.SubFOp(qk_vals_h[n_off], m_new, fastmath=fm_fast).result
+                    dl = arith.MulFOp(d, c_log2e_f, fastmath=fm_fast).result
+                    p = arith.ArithValue(dl).exp2(fastmath=fm_fast)
+                    p_vals.append(p)
+                    tile_sum = arith.AddFOp(tile_sum, p, fastmath=fm_fast).result
+                p_vals_per_head.append(p_vals)
+
+                l_alpha = arith.MulFOp(l_is[h_off], alpha, fastmath=fm_fast).result
+                l_new = arith.AddFOp(l_alpha, tile_sum, fastmath=fm_fast).result
+
+                acc_h = accs[h_off]
+                new_acc_h = []
+                for d_off in range_constexpr(D_PER_LANE):
+                    new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha, fastmath=fm_fast).result)
+                new_m_is.append(m_new)
+                new_l_is.append(l_new)
+                new_accs.append(new_acc_h)
+
+            # AV phase: load V once (MQA), reuse across HEAD_GROUP heads.
+            for n_off in range_constexpr(BLOCK_N):
+                kv_col_i32 = arith.AddIOp(
+                    n_start_i32,
+                    arith.constant(n_off, type=T.i32),
+                ).result
+                is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32)
+                kv_col_idx = arith.index_cast(T.index, kv_col_i32)
+                kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx)
+                if const_expr(mqa_kv):
+                    vl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM)
+                else:
+                    vl_row_base = (
+                        (bid * arith.index(NUM_HEADS) + qhid_base) * seq_len_v + kv_col_safe
+                    ) * arith.index(HEAD_DIM)
+                vl_lane_off = vl_row_base + lane * arith.index(D_PER_LANE)
+                vl_vec = load_f16_v(vl_ptr, vl_lane_off, D_PER_LANE)
+                vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec)
+                for h_off in range_constexpr(HEAD_GROUP):
+                    p_vals_h = p_vals_per_head[h_off]
+                    new_acc_h = new_accs[h_off]
+                    for d_off in range_constexpr(D_PER_LANE):
+                        vv = vector.extract(vl_f32, static_position=[d_off], dynamic_position=[])
+                        contrib = arith.MulFOp(p_vals_h[n_off], vv, fastmath=fm_fast).result
+                        new_acc_h[d_off] = arith.AddFOp(new_acc_h[d_off], contrib, fastmath=fm_fast).result
+
+            # Pack yield args
+            yield_args = []
+            for h_off in range_constexpr(HEAD_GROUP):
+                yield_args.append(new_m_is[h_off])
+                yield_args.append(new_l_is[h_off])
+                for d in range_constexpr(D_PER_LANE):
+                    yield_args.append(new_accs[h_off][d])
+            yield yield_args
+
+        # Unpack HEAD_GROUP states from local SWA results
+        m_is = [loop_results_local[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)]
+        l_is = [loop_results_local[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)]
+        accs = [
+            [loop_results_local[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)]
+            for h in range_constexpr(HEAD_GROUP)
+        ]
+
+        # ==== SPARSE branch ====
+        if const_expr(has_sparse):
+            init_sparse = []
+            for h_off in range_constexpr(HEAD_GROUP):
+                init_sparse.append(m_is[h_off])
+                init_sparse.append(l_is[h_off])
+                for d in range_constexpr(D_PER_LANE):
+                    init_sparse.append(accs[h_off][d])
+            for k_start, inner_args, loop_results_sparse in scf.for_(
+                arith.index(0),
+                K_topk_v,
+                arith.index(BLOCK_K),
+                iter_args=init_sparse,
+            ):
+                m_is = [inner_args[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)]
+                l_is = [inner_args[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)]
+                accs = [
+                    [inner_args[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)]
+                    for h in range_constexpr(HEAD_GROUP)
+                ]
+
+                k_start_i32 = arith.index_cast(T.i32, k_start)
+                K_topk_i32 = K_topk
+                if hasattr(K_topk_i32, "ir_value"):
+                    K_topk_i32 = K_topk_i32.ir_value()
+
+                qk_vals_sparse_per_head = [[] for _ in range_constexpr(HEAD_GROUP)]
+                for k_off in range_constexpr(BLOCK_K):
+                    k_pos_i32 = arith.AddIOp(
+                        k_start_i32,
+                        arith.constant(k_off, type=T.i32),
+                    ).result
+                    is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32)
+                    k_pos_idx = arith.index_cast(T.index, k_pos_i32)
+                    k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx)
+
+                    g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index(
+                        HEAD_DIM
+                    )
+                    g_lane_off = g_row_base + lane * arith.index(D_PER_LANE)
+                    g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE)
+                    if ENABLE_LDS_CACHE:
+                        lds_idx = arith.index(k_off * HEAD_DIM) + lane * arith.index(D_PER_LANE)
+                        vector.store(g_vec, lds_gather, [lds_idx])
+                    g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec)
+
+                    sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe
+                    sm_val = load_f32_scalar(sm_ptr, sm_off)
+                    # For each head, compute QK using shared g_f32.
+                    for h_off in range_constexpr(HEAD_GROUP):
+                        lane_dot = vec_dot_f32(g_f32, q_f32_vecs[h_off])
+                        qk_full = warp_reduce_sum_f32(lane_dot)
+                        qk_scaled = arith.MulFOp(qk_full, c_sm_scale_f, fastmath=fm_fast).result
+                        qk_biased = arith.AddFOp(qk_scaled, sm_val, fastmath=fm_fast).result
+                        qk_masked = arith.select(is_oob, c_neg_inf, qk_biased)
+                        qk_vals_sparse_per_head[h_off].append(qk_masked)
+
+                new_m_is = []
+                new_l_is = []
+                new_accs = []
+                p_vals_per_head = []
+                for h_off in range_constexpr(HEAD_GROUP):
+                    qk_vals_h = qk_vals_sparse_per_head[h_off]
+                    m_tile = qk_vals_h[0]
+                    for k_off in range_constexpr(BLOCK_K - 1):
+                        m_tile = arith.MaxNumFOp(m_tile, qk_vals_h[k_off + 1], fastmath=fm_fast).result
+                    m_new = arith.MaxNumFOp(m_is[h_off], m_tile, fastmath=fm_fast).result
+
+                    diff_m = arith.SubFOp(m_is[h_off], m_new, fastmath=fm_fast).result
+                    diff_m_log2 = arith.MulFOp(diff_m, c_log2e_f, fastmath=fm_fast).result
+                    alpha = arith.ArithValue(diff_m_log2).exp2(fastmath=fm_fast)
+
+                    p_vals = []
+                    tile_sum = c_zero_f
+                    for k_off in range_constexpr(BLOCK_K):
+                        d = arith.SubFOp(qk_vals_h[k_off], m_new, fastmath=fm_fast).result
+                        dl = arith.MulFOp(d, c_log2e_f, fastmath=fm_fast).result
+                        p = arith.ArithValue(dl).exp2(fastmath=fm_fast)
+                        p_vals.append(p)
+                        tile_sum = arith.AddFOp(tile_sum, p, fastmath=fm_fast).result
+                    p_vals_per_head.append(p_vals)
+
+                    l_alpha = arith.MulFOp(l_is[h_off], alpha, fastmath=fm_fast).result
+                    l_new = arith.AddFOp(l_alpha, tile_sum, fastmath=fm_fast).result
+
+                    acc_h = accs[h_off]
+                    new_acc_h = []
+                    for d_off in range_constexpr(D_PER_LANE):
+                        new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha, fastmath=fm_fast).result)
+                    new_m_is.append(m_new)
+                    new_l_is.append(l_new)
+                    new_accs.append(new_acc_h)
+
+                # ---- AV phase: re-read gathered K-block once per k_off, reuse for all heads ----
+                if ENABLE_LDS_CACHE:
+                    _waitcnt_lgkm_0()
+                for k_off in range_constexpr(BLOCK_K):
+                    if ENABLE_LDS_CACHE:
+                        lds_idx = arith.index(k_off * HEAD_DIM) + lane * arith.index(D_PER_LANE)
+                        g_vec = vector.load(T.vec(D_PER_LANE, f16_ty), lds_gather, [lds_idx])
+                    else:
+                        k_pos_i32 = arith.AddIOp(
+                            k_start_i32,
+                            arith.constant(k_off, type=T.i32),
+                        ).result
+                        is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32)
+                        k_pos_idx = arith.index_cast(T.index, k_pos_i32)
+                        k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx)
+                        g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index(
+                            HEAD_DIM
+                        )
+                        g_lane_off = g_row_base + lane * arith.index(D_PER_LANE)
+                        g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE)
+                    g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec)
+                    for h_off in range_constexpr(HEAD_GROUP):
+                        p_vals_h = p_vals_per_head[h_off]
+                        new_acc_h = new_accs[h_off]
+                        for d_off in range_constexpr(D_PER_LANE):
+                            vv = vector.extract(g_f32, static_position=[d_off], dynamic_position=[])
+                            contrib = arith.MulFOp(p_vals_h[k_off], vv, fastmath=fm_fast).result
+                            new_acc_h[d_off] = arith.AddFOp(
+                                new_acc_h[d_off], contrib, fastmath=fm_fast
+                            ).result
+
+                # Pack yield args
+                yield_args = []
+                for h_off in range_constexpr(HEAD_GROUP):
+                    yield_args.append(new_m_is[h_off])
+                    yield_args.append(new_l_is[h_off])
+                    for d in range_constexpr(D_PER_LANE):
+                        yield_args.append(new_accs[h_off][d])
+                yield yield_args
+
+            m_is = [loop_results_sparse[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)]
+            l_is = [loop_results_sparse[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)]
+            accs = [
+                [loop_results_sparse[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)]
+                for h in range_constexpr(HEAD_GROUP)
+            ]
+
+        # ==== Sink epilogue (per-head) ====
+        if const_expr(has_sink):
+            for h_off in range_constexpr(HEAD_GROUP):
+                qhid_h = qhid_base + arith.index(h_off)
+                qhid_i32 = arith.index_cast(T.i32, qhid_h)
+                sink_h_val = buffer_ops.buffer_load(
+                    sink_rsrc,
+                    qhid_i32,
+                    vec_width=1,
+                    dtype=f32_ty,
+                )
+                m_i_h = m_is[h_off]
+                l_i_h = l_is[h_off]
+                acc_h = accs[h_off]
+                m_new = arith.MaxNumFOp(m_i_h, sink_h_val, fastmath=fm_fast).result
+                d_alpha = arith.SubFOp(m_i_h, m_new, fastmath=fm_fast).result
+                d_alpha_log2 = arith.MulFOp(d_alpha, c_log2e_f, fastmath=fm_fast).result
+                alpha_sink = arith.ArithValue(d_alpha_log2).exp2(fastmath=fm_fast)
+                d_beta = arith.SubFOp(sink_h_val, m_new, fastmath=fm_fast).result
+                d_beta_log2 = arith.MulFOp(d_beta, c_log2e_f, fastmath=fm_fast).result
+                beta_sink = arith.ArithValue(d_beta_log2).exp2(fastmath=fm_fast)
+                l_alpha = arith.MulFOp(l_i_h, alpha_sink, fastmath=fm_fast).result
+                l_is[h_off] = arith.AddFOp(l_alpha, beta_sink, fastmath=fm_fast).result
+                new_acc_h = []
+                for d_off in range_constexpr(D_PER_LANE):
+                    new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha_sink, fastmath=fm_fast).result)
+                accs[h_off] = new_acc_h
+                m_is[h_off] = m_new
+
+        # ==== Final divide and store (per-head) ====
+        _o_guard = scf.IfOp(q_active, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for h_off in range_constexpr(HEAD_GROUP):
+                qhid_h = qhid_base + arith.index(h_off)
+                l_i_h = l_is[h_off]
+                m_i_h = m_is[h_off]
+                acc_h = accs[h_off]
+                inv_l = arith.DivFOp(c_one_f, l_i_h, fastmath=fm_fast).result
+                o_row_base = ((bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe) * arith.index(
+                    HEAD_DIM
+                )
+                o_lane_off = o_row_base + lane * arith.index(D_PER_LANE)
+                for d_off in range_constexpr(D_PER_LANE):
+                    o_f32 = arith.MulFOp(acc_h[d_off], inv_l, fastmath=fm_fast).result
+                    o_f16 = arith.trunc_f(elem_type, o_f32)
+                    elem_off = o_lane_off + arith.index(d_off)
+                    idx_i64 = arith.index_cast(T.i64, elem_off)
+                    gep = _llvm.GEPOp(
+                        _llvm_ptr_ty(),
+                        o_ptr,
+                        [idx_i64],
+                        rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                        elem_type=elem_type,
+                        noWrapFlags=0,
+                    )
+                    _llvm.StoreOp(o_f16, gep.result)
+
+                is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0))
+                _lse_if = scf.IfOp(is_lane0, [], has_else=False)
+                with ir.InsertionPoint(_lse_if.then_block):
+                    ln_l = math_dialect.log(l_i_h, fastmath=fm_fast)
+                    lse_val = arith.AddFOp(m_i_h, ln_l, fastmath=fm_fast).result
+                    lse_off = (bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe
+                    lse_off_i32 = arith.index_cast(T.i32, lse_off)
+                    buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32)
+                    scf.YieldOp([])
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_csa_fwd(
+        Q: fx.Tensor,
+        K_LOCAL: fx.Tensor,
+        V_LOCAL: fx.Tensor,
+        GATHERED: fx.Tensor,
+        SPARSE_MASK: fx.Tensor,
+        Sink: fx.Tensor,
+        O: fx.Tensor,
+        LSE: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len: fx.Int32,
+        K_topk: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len)
+        grid_x = sl_idx
+        grid_y = bs_idx * arith.index(NUM_HEAD_GROUPS)
+
+        launcher = v4_csa_fwd_kernel(
+            Q,
+            K_LOCAL,
+            V_LOCAL,
+            GATHERED,
+            SPARSE_MASK,
+            Sink,
+            O,
+            LSE,
+            seq_len,
+            K_topk,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, grid_y, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(compile_hints):
+            return launch_v4_csa_fwd(*args, **kwargs)
+
+    def _compile(
+        Q, K_LOCAL, V_LOCAL, GATHERED, SPARSE_MASK, Sink, O, LSE, batch_size, seq_len, K_topk, stream=None
+    ):
+        with CompilationContext.compile_hints(compile_hints):
+            return flyc.compile(
+                launch_v4_csa_fwd,
+                Q,
+                K_LOCAL,
+                V_LOCAL,
+                GATHERED,
+                SPARSE_MASK,
+                Sink,
+                O,
+                LSE,
+                batch_size,
+                seq_len,
+                K_topk,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+build_v4_csa_fwd_module_primary = build_v4_csa_fwd_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py
new file mode 100644
index 000000000..528456c3d
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py
@@ -0,0 +1,926 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_hca_bwd_dkv_pool: V4 HCA backward dK/dV POOL-stream kernel for FlyDSL.
+
+Forked from v4_sla_bwd_dkv_kernel.py (the SLA MQA dKdV head-loop accumulator
+template). Computes the POOL-stream contribution to dk/dv for an HCA
+(Hybrid-Causal-Additive) attention backward and stores into the POOL slice
+of dk_fp32 / dv_fp32 buffers (which are zero-initialised by the wrapper).
+
+Key differences vs the SWA kernel:
+  - KV range is fixed at [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE).
+    POOL_SIZE is a build-time constexpr. For HCA shapes POOL_SIZE <= BLOCK_N
+    (BLOCK_N=32 here, POOL_SIZE in {4,32}), so each program owns the entire
+    pool slice for one batch.
+  - No SWA-window mask, no causal mask. Two element-wise predicates only:
+      (a) pool_n < POOL_SIZE      -> NEG_INF outside the pool
+      (b) q_row < seq_len_q       -> NEG_INF for OOB q rows
+  - qk + add_bias from ADD_MASK[Sq, POOL_SIZE]. Each element loaded as bf16/f16
+    then cast to f32 inside the inner mask loop (matches dq_pool sibling).
+  - LSE / DELTAS are JOINT (saved from HCA fwd) so the same q_row LSE governs
+    both local and pool streams.
+  - Grid: (B,) -- one program per batch, owning the entire pool slice.
+  - Head loop is the SAME head-loop accumulator pattern: dynamic ``scf.for_``
+    over HQ (NOT range_constexpr; constexpr-unroll at HQ=128 hangs MLIR).
+  - sm_scale applied to dK once post head-loop (same as SWA dkv, P57 cr=0).
+  - No atomics. Single store per (b, pool_n) slice at the end.
+
+LDS budget at BLOCK_N=32, BLOCK_M2=32, D=512:
+    K = 32 KB, V = 32 KB, DO = 32 KB, Q = 32 KB, pT = 2 KB
+    Total = 130 KB <= 160 KB.
+Auto-fallback: if predicted > 160 KB, drop DO/Q LDS scratches and read
+Q/DO directly to register packs from HBM (mirrors SWA dkv).
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import memref as _memref
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME = "v4_hca_bwd_dkv_pool_kernel"
+
+_LLVM_GEP_DYNAMIC = -2147483648
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+def build_v4_hca_bwd_dkv_pool_module(
+    num_heads,
+    head_dim,
+    pool_size,
+    hca_local_seqlen,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    flat_work_group_size=None,
+    block_n=None,
+    block_m2=None,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    layout_bhld=True,
+    mqa_kv=True,
+):
+    """Build the V4 HCA backward dK/dV POOL-stream launcher (MQA only)."""
+    gpu_arch = get_hip_arch()
+
+    if block_n is None:
+        BLOCK_N = 32
+    else:
+        BLOCK_N = int(block_n)
+    if block_m2 is None:
+        BLOCK_M2 = 32
+    else:
+        BLOCK_M2 = int(block_m2)
+    WARP_SIZE = 64
+    # R6d-A: MFMA 16x16x32 -> each wave covers 16 KV rows.
+    NUM_WAVES = max(1, BLOCK_N // 16)
+    ROWS_PER_WAVE = BLOCK_N // NUM_WAVES
+    if flat_work_group_size is None:
+        flat_work_group_size = NUM_WAVES * WARP_SIZE
+    BLOCK_SIZE = flat_work_group_size
+
+    ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1"
+    USE_K16 = gpu_arch.startswith("gfx950")
+    # R6d-A: MFMA 16x16x32 K-step = 32; A/B-frag = 8 bf16 per lane.
+    assert USE_K16, "R6d-A dkv_pool requires gfx950 (MFMA 16x16x32 bf16)."
+    K_STEP_QK = 32
+    K_STEPS_QK = head_dim // K_STEP_QK
+    # R6d-A: each MFMA 16x16x32 tile produces a 16-wide N-chunk; D_CHUNK = 16 cols.
+    D_CHUNK = 16
+    D_CHUNKS = head_dim // D_CHUNK
+    K_STEPS_PT = BLOCK_M2 // K_STEP_QK
+    # R6d-A: GEMM1/GEMM3 cover m_col [0..BLOCK_M2) with multiple 16-col MFMA-N tiles per ks.
+    assert BLOCK_M2 % 16 == 0, f"BLOCK_M2 must be a multiple of 16 (MFMA-N), got {BLOCK_M2}"
+    M_TILES = BLOCK_M2 // 16
+
+    assert BLOCK_N % NUM_WAVES == 0
+    assert ROWS_PER_WAVE == 16, f"ROWS_PER_WAVE must equal 16 for MFMA 16x16x32, got {ROWS_PER_WAVE}"
+    assert head_dim % 32 == 0
+    assert head_dim >= 64
+    assert flat_work_group_size in (64, 128, 256, 512)
+    assert dtype_str == "bf16", "R6d-A dkv_pool currently only supports bf16."
+    assert BLOCK_N % 16 == 0
+    assert BLOCK_M2 % K_STEP_QK == 0, (
+        f"BLOCK_M2 ({BLOCK_M2}) must be a multiple of MFMA K-step " f"({K_STEP_QK}) for the dV/dK GEMMs."
+    )
+    assert mqa_kv, "v4_hca_bwd_dkv_pool currently only supports MQA (HK=1)."
+    assert (
+        isinstance(pool_size, int) and 0 < pool_size <= BLOCK_N
+    ), f"pool_size must be int in (0, {BLOCK_N}], got {pool_size!r}"
+    assert isinstance(hca_local_seqlen, int) and hca_local_seqlen >= 0
+    assert (
+        hca_local_seqlen % BLOCK_N == 0
+    ), f"hca_local_seqlen must be multiple of BLOCK_N={BLOCK_N}, got {hca_local_seqlen}"
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(head_dim)
+
+    NUM_HEADS = num_heads
+    HEAD_DIM = head_dim
+    POOL_SIZE = int(pool_size)
+    HCA_LOCAL = int(hca_local_seqlen)
+
+    K_STRIDE = HEAD_DIM
+    K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1
+    assert K_SWZ_ROW_MASK >= 0
+    assert (K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)) == 0
+    V_STRIDE = HEAD_DIM
+
+    VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+
+    LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE
+    LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE
+    LDS_DO_STRIDE = HEAD_DIM
+    LDS_DO_ELEMS = BLOCK_M2 * LDS_DO_STRIDE
+    LDS_Q_STRIDE = HEAD_DIM
+    LDS_Q_ELEMS = BLOCK_M2 * LDS_Q_STRIDE
+    LDS_PT_STRIDE = BLOCK_M2
+    LDS_PT_ELEMS = BLOCK_N * LDS_PT_STRIDE
+
+    _LDS_LIMIT_BYTES = 160 * 1024
+
+    def _predict_lds_bytes(use_lds_for_q_do):
+        b = (LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + LDS_PT_ELEMS) * 2
+        if use_lds_for_q_do:
+            b += (LDS_DO_ELEMS + LDS_Q_ELEMS) * 2
+        return b
+
+    USE_LDS_FOR_Q_DO = True
+    _pred_full = _predict_lds_bytes(True)
+    if _pred_full > _LDS_LIMIT_BYTES:
+        USE_LDS_FOR_Q_DO = False
+        _pred_min = _predict_lds_bytes(False)
+        if _pred_min > _LDS_LIMIT_BYTES:
+            raise RuntimeError(
+                f"v4_hca_bwd_dkv_pool: minimal LDS {_pred_min}B > limit "
+                f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, "
+                f"BLOCK_M2={BLOCK_M2}."
+            )
+        import sys as _sys
+
+        print(
+            f"[v4_hca_bwd_dkv_pool] auto-fallback: LDS {_pred_full}B > "
+            f"{_LDS_LIMIT_BYTES}B; disabling Q/DO LDS "
+            f"({_pred_full} -> {_pred_min})",
+            file=_sys.stderr,
+            flush=True,
+        )
+    else:
+        import sys as _sys
+
+        print(
+            f"[v4_hca_bwd_dkv_pool] LDS OK: predicted {_pred_full}B <= "
+            f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, "
+            f"BLOCK_M2={BLOCK_M2}, USE_LDS_FOR_Q_DO=True",
+            file=_sys.stderr,
+            flush=True,
+        )
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=(
+            f"v4_hca_bwd_dkv_pool_smem_N{BLOCK_N}_M2_{BLOCK_M2}_P{POOL_SIZE}"
+            f"_L{HCA_LOCAL}_MQ{int(mqa_kv)}_QDO{int(USE_LDS_FOR_Q_DO)}"
+        ),
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    LDS_V_BASE = LDS_K_TILE_SIZE
+    LDS_KV_TOTAL_SIZE = LDS_K_TILE_SIZE + LDS_V_TILE_SIZE
+    allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2
+
+    lds_pt_offset = allocator._align(allocator.ptr, 16)
+    allocator.ptr = lds_pt_offset + LDS_PT_ELEMS * 2
+
+    if USE_LDS_FOR_Q_DO:
+        lds_do_offset = allocator._align(allocator.ptr, 16)
+        allocator.ptr = lds_do_offset + LDS_DO_ELEMS * 2
+        lds_q_offset = allocator._align(allocator.ptr, 16)
+        allocator.ptr = lds_q_offset + LDS_Q_ELEMS * 2
+    else:
+        lds_do_offset = None
+        lds_q_offset = None
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_hca_bwd_dkv_pool_kernel(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DK: fx.Tensor,
+        DV: fx.Tensor,
+        ADD_MASK: fx.Tensor,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K)
+        v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS)
+        _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DK)
+        _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DV)
+        add_mask_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), ADD_MASK)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+        # R6e: dk/dv pool-slice writes are atomic_fadd (multi-program collisions).
+        dk_rsrc = buffer_ops.create_buffer_resource(DK, max_size=True)
+        dv_rsrc = buffer_ops.create_buffer_resource(DV, max_size=True)
+
+        fm_fast = arith.FastMathFlags.fast
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        # R6d-A: MFMA 16x16x32 C-frag = 4 fp32 per lane.
+        v4f32_type = T.vec(4, compute_type)
+        mfma_pack_type = v8f16_type
+        MFMA_LANE_K = 8
+        v1_elem_type = T.vec(1, elem_type)
+        _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)
+
+        def mfma_acc(a, b, c):
+            # rocdl.mfma_f32_16x16x32_bf16 is the wrapped form: takes
+            # (result_type, [operands]) and returns the Value directly.
+            return rocdl.mfma_f32_16x16x32_bf16(
+                v4f32_type,
+                [a, b, c, _mfma_zero, _mfma_zero, _mfma_zero],
+            )
+
+        seq_len_q_v = arith.index_cast(T.index, seq_len_q)
+        seq_len_k_v = arith.index_cast(T.index, seq_len_k)
+
+        base_ptr = allocator.get_base()
+        lds_kv = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_TOTAL_SIZE,)).get()
+        lds_pt = SmemPtr(base_ptr, lds_pt_offset, elem_type, shape=(LDS_PT_ELEMS,)).get()
+        if USE_LDS_FOR_Q_DO:
+            lds_do = SmemPtr(base_ptr, lds_do_offset, elem_type, shape=(LDS_DO_ELEMS,)).get()
+            lds_q = SmemPtr(base_ptr, lds_q_offset, elem_type, shape=(LDS_Q_ELEMS,)).get()
+
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        # R6d-A MFMA 16x16x32 lane decomposition:
+        #   A-frag: A[lane_mod_16, ks*32 + lane_div_16*8 + 0..7]
+        #   B-frag: B[ks*32 + lane_div_16*8 + 0..7, lane_mod_16]
+        #   C-frag: C[lane_div_16*4 + ii, lane_mod_16]  for ii in 0..3
+        lane_mod_16 = lane % 16
+        lane_div_16 = lane // 16
+
+        wave_n_offset = wave_id * ROWS_PER_WAVE
+
+        # R6e: grid is (B * num_m_blocks,). Decompose into (batch_idx, m_tile_idx).
+        # Layout: m_tile fastest -> consecutive programs share Q/dO batch.
+        BM2_idx_grid = arith.index(BLOCK_M2)
+        _one_idx_grid = arith.index(1)
+        num_m_blocks = (seq_len_q_v + BM2_idx_grid - _one_idx_grid) // BM2_idx_grid
+        m_tile_idx = block_id % num_m_blocks
+        batch_idx = block_id // num_m_blocks
+        m_start = m_tile_idx * BM2_idx_grid
+        # Pool slice starts at HCA_LOCAL_SEQLEN and runs for POOL_SIZE keys.
+        kv_start = arith.index(HCA_LOCAL)
+
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        bh_base_tokens_kv = batch_idx * seq_len_k_v
+
+        def bh_base_tokens_q_of(qhid_index):
+            return (batch_idx * NUM_HEADS + qhid_index) * seq_len_q_v
+
+        def global_idx_q(qhid_index, token_idx, col):
+            return (bh_base_tokens_q_of(qhid_index) + token_idx) * arith.index(HEAD_DIM) + col
+
+        def global_idx_kv(token_idx, col):
+            return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col
+
+        def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=et,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr_, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, vxf16_type)
+
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        if ROWS_PER_BATCH_LOAD >= BLOCK_N:
+            NUM_BATCHES_KV = 1
+            KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N
+        else:
+            assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+            NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+            KV_NEEDS_GUARD = False
+
+        if ROWS_PER_BATCH_LOAD >= BLOCK_M2:
+            NUM_BATCHES_M = 1
+            M_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_M2
+        else:
+            assert BLOCK_M2 % ROWS_PER_BATCH_LOAD == 0
+            NUM_BATCHES_M = BLOCK_M2 // ROWS_PER_BATCH_LOAD
+            M_NEEDS_GUARD = False
+
+        c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        c_zero_elem = arith.constant(0.0, type=elem_type)
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)
+        c_sm_scale = arith.constant(sm_scale, type=compute_type)
+        # R6d-A: 4-elem fp32 acc per MFMA 16x16x32 op.
+        v4f32_zero = arith.constant_vector(0.0, v4f32_type)
+
+        # ---- PROLOGUE: cooperative load K, V into LDS ----
+        pool_end = kv_start + arith.index(POOL_SIZE)
+
+        def coop_load_k():
+            k_base = arith.index(0)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = kv_start + load_row_in_batch + row_offset
+                in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, pool_end)
+                row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                g_idx = global_idx_kv(row_safe, load_col_base)
+                vec = load_global_f16xN(k_ptr, g_idx)
+                vec_safe = arith.select(in_bounds, vec, c_zero_vxf16)
+                lds_row = load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N)
+                    )
+                    _if_k = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_k.then_block):
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                        vector.store(vec_safe, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    vector.store(vec_safe, lds_kv, [lds_idx])
+
+        def coop_load_v():
+            v_base = arith.index(LDS_V_BASE)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = kv_start + load_row_in_batch + row_offset
+                in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, pool_end)
+                row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                g_idx = global_idx_kv(row_safe, load_col_base)
+                vec = load_global_f16xN(v_ptr, g_idx)
+                vec_safe = arith.select(in_bounds, vec, c_zero_vxf16)
+                lds_row = load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N)
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = v_base + lds_row * V_STRIDE + swz_col
+                        vector.store(vec_safe, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = v_base + lds_row * V_STRIDE + swz_col
+                    vector.store(vec_safe, lds_kv, [lds_idx])
+
+        coop_load_k()
+        coop_load_v()
+        gpu.barrier()
+
+        # R6d-A: K/V LDS A-frag uses lane_mod_16 (16-row MFMA) and lane_div_16*8 for col.
+        # Per-call swizzle mask (folds wave_n_offset bits at D=512).
+        def _k_idx_wave(ks):
+            kv_row = wave_n_offset + lane_mod_16
+            col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+            mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return kv_row * arith.index(K_STRIDE) + (col ^ mask)
+
+        def _v_idx_wave(ks):
+            kv_row = wave_n_offset + lane_mod_16
+            col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+            mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return arith.index(LDS_V_BASE) + kv_row * arith.index(V_STRIDE) + (col ^ mask)
+
+        # R6e: m-loop replaced by per-program m-tile.
+        arith.index(BLOCK_N)
+        arith.index(BLOCK_M2)
+        arith.index(1)
+
+        outer_carry = [v4f32_zero for _ in range(D_CHUNKS)] + [v4f32_zero for _ in range(D_CHUNKS)]
+
+        seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v)
+        wave_n_off_i32 = arith.index_cast(T.i32, wave_n_offset)
+        lane_div_16_i32 = arith.index_cast(T.i32, lane_div_16)
+        pool_size_i32 = arith.constant(POOL_SIZE, type=T.i32)
+
+        # ADD_MASK row stride in elements (= POOL_SIZE; bf16/f16 contiguous).
+        ADD_MASK_STRIDE_M = arith.index(POOL_SIZE)
+
+        NUM_HEADS_idx = arith.index(NUM_HEADS)
+        for qhid_constexpr_idx, h_carry, h_loop_results in scf.for_(
+            arith.index(0),
+            NUM_HEADS_idx,
+            arith.index(1),
+            iter_args=outer_carry,
+        ):
+            qhid = qhid_constexpr_idx
+
+            # R6e: per-program m-tile (no inner m-loop). Accumulators are head-loop carry.
+            dv_accs = [h_carry[dc] for dc in range_constexpr(D_CHUNKS)]
+            dk_accs = [h_carry[D_CHUNKS + dc] for dc in range_constexpr(D_CHUNKS)]
+
+            # R6d-A: lane_mod_16 is the per-tile q-row coord; full per-tile loop builds row below.
+            q_row_abs = m_start + lane_mod_16
+            q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row_abs, seq_len_q_v)
+            arith.select(q_in_bounds, q_row_abs, arith.index(0))
+
+            if USE_LDS_FOR_Q_DO:
+                for batch in range_constexpr(NUM_BATCHES_M):
+                    row_offset = batch * ROWS_PER_BATCH_LOAD
+                    row_idx = m_start + load_row_in_batch + row_offset
+                    in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_q_v)
+                    row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                    g_idx_do = global_idx_q(qhid, row_safe, load_col_base)
+                    g_idx_q = global_idx_q(qhid, row_safe, load_col_base)
+                    vec_do = load_global_f16xN(do_ptr, g_idx_do)
+                    vec_q = load_global_f16xN(q_ptr, g_idx_q)
+                    vec_do_safe = arith.select(in_bounds, vec_do, c_zero_vxf16)
+                    vec_q_safe = arith.select(in_bounds, vec_q, c_zero_vxf16)
+                    lds_row = load_row_in_batch + row_offset
+                    if M_NEEDS_GUARD:
+                        row_valid = arith.cmpi(
+                            arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_M2)
+                        )
+                        _if_qd = scf.IfOp(row_valid)
+                        with ir.InsertionPoint(_if_qd.then_block):
+                            lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base
+                            lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base
+                            vector.store(vec_do_safe, lds_do, [lds_idx_do])
+                            vector.store(vec_q_safe, lds_q, [lds_idx_q])
+                            scf.YieldOp([])
+                    else:
+                        lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base
+                        lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base
+                        vector.store(vec_do_safe, lds_do, [lds_idx_do])
+                        vector.store(vec_q_safe, lds_q, [lds_idx_q])
+                gpu.barrier()
+
+                # R6d-A: B-frag for GEMM1 (qkT = K @ Q^T) per m-tile mt in [0, M_TILES).
+                #   Q[mt*16 + lane_mod_16, ks*32 + lane_div_16*8 + 0..7]
+                q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                for mt in range_constexpr(M_TILES):
+                    for ks in range_constexpr(K_STEPS_QK):
+                        m_row = arith.index(mt * 16) + lane_mod_16
+                        d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                        lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col
+                        pack = vector.load_op(mfma_pack_type, lds_q, [lds_idx])
+                        q_b_packs[mt][ks] = pack
+
+                # R6d-A: B-frag for GEMM3 (dp = V @ DO^T) per m-tile mt.
+                do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                for mt in range_constexpr(M_TILES):
+                    for ks in range_constexpr(K_STEPS_QK):
+                        m_row = arith.index(mt * 16) + lane_mod_16
+                        d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                        lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col
+                        pack = vector.load_op(mfma_pack_type, lds_do, [lds_idx])
+                        do_b_packs_gemm3[mt][ks] = pack
+            else:
+                # R6d-A fallback: HBM-direct B-frag per m-tile. row = m_start + mt*16 + lane_mod_16.
+                q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                for mt in range_constexpr(M_TILES):
+                    row_abs = m_start + arith.index(mt * 16) + lane_mod_16
+                    in_bnd = arith.cmpi(arith.CmpIPredicate.slt, row_abs, seq_len_q_v)
+                    row_safe = arith.select(in_bnd, row_abs, arith.index(0))
+                    for ks in range_constexpr(K_STEPS_QK):
+                        col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                        g_idx = global_idx_q(qhid, row_safe, col)
+                        q_raw = load_global_mfma_pack(q_ptr, g_idx)
+                        q_b_packs[mt][ks] = arith.select(in_bnd, q_raw, c_zero_mfma_pack)
+                        do_raw = load_global_mfma_pack(do_ptr, g_idx)
+                        do_b_packs_gemm3[mt][ks] = arith.select(in_bnd, do_raw, c_zero_mfma_pack)
+
+            # ---- GEMM1: qkT = K @ Q^T (MFMA 16x16x32, one tile per m-tile mt) ----
+            s_accs = [v4f32_zero for _ in range(M_TILES)]
+            for ks in range_constexpr(K_STEPS_QK):
+                k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)])
+                for mt in range_constexpr(M_TILES):
+                    s_accs[mt] = mfma_acc(k_pack, q_b_packs[mt][ks], s_accs[mt])
+
+            # R6d-A: per m-tile lse/delta, pool-only mask, ADD_MASK bias, softmax.
+            # m_row for tile mt = m_start + mt*16 + lane_mod_16.
+            # Resulting pT_vals_per_tile[mt][ii] for ii in 0..3.
+            pT_vals_per_tile = []
+            lse_per_tile = []
+            delta_per_tile = []
+            for mt in range_constexpr(M_TILES):
+                m_row_for_lse = m_start + arith.index(mt * 16) + lane_mod_16
+                m_row_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, m_row_for_lse, seq_len_q_v)
+                m_row_safe = arith.select(m_row_in_bounds, m_row_for_lse, arith.index(0))
+                lse_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q_of(qhid) + m_row_safe)
+                lse_v = buffer_ops.buffer_load(lse_rsrc, lse_off_i32, vec_width=1, dtype=T.f32)
+                delta_v = buffer_ops.buffer_load(deltas_rsrc, lse_off_i32, vec_width=1, dtype=T.f32)
+                lse_per_tile.append(lse_v)
+                delta_per_tile.append(delta_v)
+
+                m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse)
+                m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32)
+                # ADD_MASK row base (clamp to row 0 if OOB; lane is masked).
+                add_mask_row_base = m_row_safe * ADD_MASK_STRIDE_M
+
+                pT_tile = []
+                for ii in range_constexpr(4):
+                    ii_i32 = arith.constant(ii, type=T.i32)
+                    # R6d-A: 16x16x32 C-frag lane stride is 4 (NOT 8 like 32x32x16).
+                    pool_n_rel_i32 = arith.AddIOp(
+                        arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result, ii_i32
+                    ).result
+                    pool_n_i32 = arith.AddIOp(wave_n_off_i32, pool_n_rel_i32).result
+                    pool_n_oob = arith.cmpi(arith.CmpIPredicate.sge, pool_n_i32, pool_size_i32)
+                    bad = arith.OrIOp(pool_n_oob, m_oob).result
+
+                    # Load add_bias from ADD_MASK[m_row, pool_n].
+                    pool_n_safe_i32 = arith.select(bad, arith.constant(0, type=T.i32), pool_n_i32)
+                    pool_n_idx = arith.index_cast(T.index, pool_n_safe_i32)
+                    add_elem_idx = add_mask_row_base + pool_n_idx
+                    bias_raw_v1 = _gep_load(add_mask_ptr, add_elem_idx, T.vec(1, elem_type))
+                    bias_raw = vector.extract(bias_raw_v1, static_position=[0], dynamic_position=[])
+                    bias_f32 = arith.extf(compute_type, bias_raw)
+                    bias_safe = arith.select(bad, c_zero_f, bias_f32)
+
+                    s_ii = vector.extract(s_accs[mt], static_position=[ii], dynamic_position=[])
+                    scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result
+                    scaled_plus_bias = arith.AddFOp(scaled, bias_safe, fastmath=fm_fast).result
+                    scaled_m = arith.select(bad, c_neg_inf, scaled_plus_bias)
+                    diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result
+                    p = math_dialect.exp(diff, fastmath=fm_fast)
+                    pT_tile.append(p)
+                pT_vals_per_tile.append(pT_tile)
+
+            # ---- pT register -> LDS (per m-tile mt, write 4 C-frag elems/lane @ col mt*16 + lane_mod_16) ----
+            for mt in range_constexpr(M_TILES):
+                for ii in range_constexpr(4):
+                    kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                    kv_row = wave_n_offset + kv_row_rel
+                    pt_bf16 = arith.trunc_f(elem_type, pT_vals_per_tile[mt][ii])
+                    lds_pt_idx = kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                    v1 = vector.from_elements(v1_elem_type, [pt_bf16])
+                    vector.store(v1, lds_pt, [lds_pt_idx])
+
+            # ---- pT A-frag for GEMM2 (dV += pT @ DO):
+            #   A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7]
+            pt_a_packs = []
+            for m_step in range_constexpr(K_STEPS_PT):
+                kv_row_a = wave_n_offset + lane_mod_16
+                m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a
+                pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx])
+                pt_a_packs.append(pack)
+
+            # ---- GEMM2: dV += pT @ DO ----
+            # R6d-A GEMM2 (dV += pT @ DO): B-frag DO[m_step*32+lane_div_16*8+k, dc*16+lane_mod_16].
+            if USE_LDS_FOR_Q_DO:
+
+                def read_do_b_pack(m_step_idx, dc_idx):
+                    d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                    m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    vals = []
+                    for rk in range_constexpr(MFMA_LANE_K):
+                        m_row = m_base + arith.index(rk)
+                        lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col
+                        val = _memref.load(lds_do, [lds_idx])
+                        vals.append(val)
+                    return vector.from_elements(mfma_pack_type, vals)
+
+            else:
+
+                def read_do_b_pack(m_step_idx, dc_idx):
+                    d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                    m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    vals = []
+                    for rk in range_constexpr(MFMA_LANE_K):
+                        m_row_rel = m_base + arith.index(rk)
+                        m_row_abs = m_start + m_row_rel
+                        in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v)
+                        m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0))
+                        g_idx = global_idx_q(qhid, m_row_safe2, d_col)
+                        v1 = _gep_load(do_ptr, g_idx, T.vec(1, elem_type))
+                        v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[])
+                        v_safe = arith.select(in_b, v_scalar, c_zero_elem)
+                        vals.append(v_safe)
+                    return vector.from_elements(mfma_pack_type, vals)
+
+            new_dv_accs = list(dv_accs)
+            for dc in range_constexpr(D_CHUNKS):
+                for pks in range_constexpr(K_STEPS_PT):
+                    b_pack = read_do_b_pack(pks, dc)
+                    new_dv_accs[dc] = mfma_acc(pt_a_packs[pks], b_pack, new_dv_accs[dc])
+
+            # ---- GEMM3: dp = V @ DO^T (MFMA 16x16x32, one tile per m-tile mt) ----
+            dp_accs = [v4f32_zero for _ in range(M_TILES)]
+            for ks in range_constexpr(K_STEPS_QK):
+                v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)])
+                for mt in range_constexpr(M_TILES):
+                    dp_accs[mt] = mfma_acc(v_pack, do_b_packs_gemm3[mt][ks], dp_accs[mt])
+
+            # ---- dsT = pT * (dp - delta) per m-tile ----
+            dsT_vals_per_tile = []
+            for mt in range_constexpr(M_TILES):
+                ds_tile = []
+                for ii in range_constexpr(4):
+                    dp_ii = vector.extract(dp_accs[mt], static_position=[ii], dynamic_position=[])
+                    diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result
+                    ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result
+                    ds_tile.append(ds_ii)
+                dsT_vals_per_tile.append(ds_tile)
+
+            # ---- dsT register -> LDS (per m-tile, col = mt*16 + lane_mod_16) ----
+            for mt in range_constexpr(M_TILES):
+                for ii in range_constexpr(4):
+                    kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                    kv_row = wave_n_offset + kv_row_rel
+                    ds_bf16 = arith.trunc_f(elem_type, dsT_vals_per_tile[mt][ii])
+                    lds_pt_idx = kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                    v1_ds = vector.from_elements(v1_elem_type, [ds_bf16])
+                    vector.store(v1_ds, lds_pt, [lds_pt_idx])
+
+            # ---- dsT A-frag for GEMM4 (dK += dsT @ Q):
+            #   A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7]
+            ds_a_packs = []
+            for m_step in range_constexpr(K_STEPS_PT):
+                kv_row_a = wave_n_offset + lane_mod_16
+                m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a
+                pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx])
+                ds_a_packs.append(pack)
+
+            # ---- GEMM4: dK += dsT @ Q  (MFMA 16x16x32) ----
+            # B-frag: Q[m_step*32+lane_div_16*8+k, dc*16+lane_mod_16].
+            if USE_LDS_FOR_Q_DO:
+
+                def read_q_b_pack(m_step_idx, dc_idx):
+                    d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                    m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    vals = []
+                    for rk in range_constexpr(MFMA_LANE_K):
+                        m_row = m_base + arith.index(rk)
+                        lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col
+                        val = _memref.load(lds_q, [lds_idx])
+                        vals.append(val)
+                    return vector.from_elements(mfma_pack_type, vals)
+
+            else:
+
+                def read_q_b_pack(m_step_idx, dc_idx):
+                    d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                    m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    vals = []
+                    for rk in range_constexpr(MFMA_LANE_K):
+                        m_row_rel = m_base + arith.index(rk)
+                        m_row_abs = m_start + m_row_rel
+                        in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v)
+                        m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0))
+                        g_idx = global_idx_q(qhid, m_row_safe2, d_col)
+                        v1 = _gep_load(q_ptr, g_idx, T.vec(1, elem_type))
+                        v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[])
+                        v_safe = arith.select(in_b, v_scalar, c_zero_elem)
+                        vals.append(v_safe)
+                    return vector.from_elements(mfma_pack_type, vals)
+
+            new_dk_accs = list(dk_accs)
+            for dc in range_constexpr(D_CHUNKS):
+                for pks in range_constexpr(K_STEPS_PT):
+                    q_b_pack = read_q_b_pack(pks, dc)
+                    new_dk_accs[dc] = mfma_acc(ds_a_packs[pks], q_b_pack, new_dk_accs[dc])
+
+            gpu.barrier()
+
+            # R6e: yield (dv, dk) partial directly as head-loop carry; no inner m-loop.
+            yield list(new_dv_accs) + list(new_dk_accs)
+        outer_carry = list(h_loop_results)
+
+        # ---- Final: dK *= sm_scale, store ----
+        dv_finals = [outer_carry[dc] for dc in range(D_CHUNKS)]
+        dk_finals = [outer_carry[D_CHUNKS + dc] for dc in range(D_CHUNKS)]
+
+        # R6e: atomic_fadd into shared pool slice of dk/dv.
+        # Multiple m-tile programs collide on the same (b, pool_n, d) addresses.
+        # DK/DV are fp32 contiguous; byte_offset = global_idx_kv * 4.
+        _atom_zero_i32 = arith.constant(0, type=T.i32)
+        _four_i32 = arith.constant(4, type=T.i32)
+        for dc in range_constexpr(D_CHUNKS):
+            for ii in range_constexpr(4):
+                kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                kv_row_abs = kv_start + wave_n_offset + kv_row_rel
+                d_col_abs = arith.index(dc * D_CHUNK) + lane_mod_16
+                kv_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, kv_row_abs, pool_end)
+                _if_kv = scf.IfOp(kv_in_bounds)
+                with ir.InsertionPoint(_if_kv.then_block):
+                    dv_val = vector.extract(dv_finals[dc], static_position=[ii], dynamic_position=[])
+                    dk_val = vector.extract(dk_finals[dc], static_position=[ii], dynamic_position=[])
+                    dk_scaled = arith.MulFOp(dk_val, c_sm_scale, fastmath=fm_fast).result
+                    g_elem_idx = global_idx_kv(kv_row_abs, d_col_abs)
+                    g_elem_i32 = arith.index_cast(T.i32, g_elem_idx)
+                    byte_off_i32 = arith.MulIOp(g_elem_i32, _four_i32).result
+                    rocdl.raw_ptr_buffer_atomic_fadd(
+                        dv_val,
+                        dv_rsrc,
+                        byte_off_i32,
+                        _atom_zero_i32,
+                        _atom_zero_i32,
+                    )
+                    rocdl.raw_ptr_buffer_atomic_fadd(
+                        dk_scaled,
+                        dk_rsrc,
+                        byte_off_i32,
+                        _atom_zero_i32,
+                        _atom_zero_i32,
+                    )
+                    scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_hca_bwd_dkv_pool(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DK: fx.Tensor,
+        DV: fx.Tensor,
+        ADD_MASK: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        # R6e: grid_x = B * num_m_blocks. Parallelize over m-tiles for B=1.
+        sq_idx_h = arith.index_cast(T.index, seq_len_q)
+        BM2_idx_h = arith.index(BLOCK_M2)
+        _one_idx_h = arith.index(1)
+        num_m_blocks_v = (sq_idx_h + BM2_idx_h - _one_idx_h) // BM2_idx_h
+        grid_x = bs_idx * num_m_blocks_v
+
+        launcher = v4_hca_bwd_dkv_pool_kernel(
+            Q,
+            K,
+            V,
+            DOS,
+            LSE,
+            DELTAS,
+            DK,
+            DV,
+            ADD_MASK,
+            seq_len_q,
+            seq_len_k,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+        if flat_work_group_size is not None:
+            _fwgs = int(flat_work_group_size)
+            if _fwgs >= 1:
+                flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, 1, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    _fmha_compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {
+            "enable-post-misched": False,
+            "lsr-drop-solution": True,
+        },
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return launch_v4_hca_bwd_dkv_pool(*args, **kwargs)
+
+    def _compile(Q, K, V, DOS, LSE, DELTAS, DK, DV, ADD_MASK, batch_size, seq_len_q, seq_len_k, stream=None):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return flyc.compile(
+                launch_v4_hca_bwd_dkv_pool,
+                Q,
+                K,
+                V,
+                DOS,
+                LSE,
+                DELTAS,
+                DK,
+                DV,
+                ADD_MASK,
+                batch_size,
+                seq_len_q,
+                seq_len_k,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+# Convenience alias.
+build_v4_hca_bwd_dkv_pool_module_primary = build_v4_hca_bwd_dkv_pool_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py
new file mode 100644
index 000000000..f69001491
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py
@@ -0,0 +1,1404 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_hca_bwd_dq_pool: V4 HCA backward dQ POOL-stream kernel for FlyDSL.
+
+Forked from v4_sla_bwd_dq_kernel.py. Computes the POOL stream contribution
+to dq for an HCA (Hybrid-Causal-Additive) attention backward, then
+ACCUMULATES into an existing dq_fp32 buffer (which already contains the
+LOCAL stream dq from the SWA dq kernel).
+
+Differences vs the SWA kernel:
+  - KV range is fixed at [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE).
+    POOL_SIZE is a build-time constexpr. For our shapes POOL_SIZE <= BLOCK_N
+    so there is exactly ONE n-block iteration.
+  - No SWA-window mask, no causal mask. Two element-wise predicates only:
+      (a) pool_n < POOL_SIZE     -> NEG_INF outside the pool
+      (b) q_row < seq_len_q      -> NEG_INF for OOB q rows
+  - qk + add_bias from the ADD_MASK tensor (shape [Sq, POOL_SIZE]).
+    Each element loaded as f32 (after cast from bf16/f16) and added to qk
+    before -lse and exp.
+  - No SINK / DSINK. Sink is handled by the LOCAL FlyDSL dq kernel; the
+    pool stream does not touch the sink.
+  - Final store ACCUMULATES into DQ (load + add + store). Race-free
+    because each program owns a unique (b, qhid, m_block) slice and the
+    pool kernel runs AFTER the SWA dq has finished writing dq for that
+    same slice (sequential launches from the wrapper).
+  - sm_scale is applied INSIDE the loop on qk (matches Triton ref); after
+    the n-loop (which has only one iter), sm_scale is applied ONCE MORE
+    on dq before accumulating into DQ (P57 cr=0 BWD).
+
+Layout: BHLD. Q/DOUT/DQ flat from (B, HQ, Sq, D). K/V flat from
+        (B, HK, Sk, D); MQA means HK=1, no head stride applied to KV.
+LSE/DELTAS: (B, HQ, Sq) flat, fp32, raw-domain (JOINT lse from fwd).
+ADD_MASK:   (Sq, POOL_SIZE) bf16/f16; loaded as f32 inside the kernel.
+DQ:         (B, HQ, Sq, D) fp32; ACCUMULATED.
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import memref as _memref
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+# ---- Module-level constants ----
+
+KERNEL_NAME = "v4_hca_bwd_dq_pool_kernel"
+
+_LOG2E = math.log2(math.e)  # 1.4426950408889634
+
+_LLVM_GEP_DYNAMIC = -2147483648  # LLVM kDynamicIndex sentinel (0x80000000 as signed i32)
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+_VMCNT_LO_MASK = 0xF
+_LGKMCNT_EXPCNT_BASE = 0x3F70
+_VMCNT_HI_SHIFT = 14
+_VMCNT_HI_MASK = 0x3
+
+
+def _waitcnt_vm_n(n):
+    """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7)."""
+    val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT)
+    rocdl.s_waitcnt(val)
+
+
+def build_v4_hca_bwd_dq_pool_module(
+    num_heads,
+    head_dim,
+    pool_size,
+    hca_local_seqlen,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    flat_work_group_size=None,
+    block_m=None,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    layout_bhld=True,
+    mqa_kv=True,
+):
+    """Build the V4 HCA backward dQ POOL-stream launcher.
+
+    Args:
+        num_heads: HQ (number of Q heads).
+        head_dim: D (head dimension), must be % 32 == 0 and >= 64.
+        pool_size: int > 0; number of pool keys (<= BLOCK_N=64 in this
+            kernel; one n-block per program iter). Build-time constexpr.
+        hca_local_seqlen: int >= 0; offset into K/V at which the pool
+            keys start (== Sq for HCA split-mask). Build-time constexpr;
+            must be a multiple of BLOCK_N to keep the single-iter
+            n-block aligned.
+        sm_scale: defaults to 1/sqrt(head_dim).
+        mqa_kv: if True, K/V indexing drops head_idx (HK=1, stride_h=0).
+        layout_bhld: BHLD layout (True) or BLHD (False); V4 uses BHLD.
+
+    Returns: launch(Q, K, V, DOUT, LSE, DELTAS, DQ_FP32, ADD_MASK,
+                    batch_size, seq_len_q, seq_len_k, stream=None).
+            DQ_FP32 is ACCUMULATED into (load + add + store).
+    """
+    gpu_arch = get_hip_arch()
+
+    BLOCK_N = 64
+    K_SUB_N = 32
+    WARP_SIZE = 64
+
+    if block_m is not None:
+        BLOCK_M = block_m
+    else:
+        BLOCK_M = 128
+
+    if flat_work_group_size is None:
+        if BLOCK_M <= 128:
+            flat_work_group_size = 256
+        else:
+            flat_work_group_size = 512
+    NUM_WAVES = flat_work_group_size // WARP_SIZE
+    BLOCK_SIZE = flat_work_group_size
+    ROWS_PER_WAVE = BLOCK_M // NUM_WAVES
+    # V4 SWA: dense path. One outer iter = one BLOCK_N block.
+    BLOCK_N_OUT = BLOCK_N
+    ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1"
+    _has_lds_load_b128 = not gpu_arch.startswith("gfx942")
+    ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1")
+    ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1"
+    REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower()
+    if REDUCE_MODE not in ("xor", "ds_bpermute"):
+        REDUCE_MODE = "xor"
+    NUM_PREFETCH_K = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1)
+    NUM_PREFETCH_V = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1)
+    (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,)
+
+    USE_HW_TR = gpu_arch.startswith("gfx950")
+    USE_K16 = gpu_arch.startswith("gfx950")
+
+    # Auto-fallback: gfx950 LDS limit is 160 KB. K+V double-buffering at D=512
+    # easily exceeds that. If predicted LDS exceeds budget, force DMA off
+    # (NUM_PREFETCH = 1) and warn. Keeps the kernel correct regardless of
+    # the env knob.
+    _LDS_LIMIT_BYTES = 160 * 1024
+
+    def _predicted_lds_bytes(nk, nv, dma):
+        # USE_HW_TR & DMA  -> V_STRIDE = head_dim
+        # USE_HW_TR & !DMA -> V_STRIDE = head_dim + 4
+        # !USE_HW_TR       -> V stored transposed: LDS V = head_dim * (BLOCK_N + 2)
+        if USE_HW_TR:
+            v_str = head_dim if dma else head_dim + 4
+            return (nk * BLOCK_N * head_dim + nv * BLOCK_N * v_str) * 2
+        vt = BLOCK_N + 2
+        return (nk * BLOCK_N * head_dim + nv * head_dim * vt) * 2
+
+    _pred = _predicted_lds_bytes(NUM_PREFETCH_K, NUM_PREFETCH_V, ENABLE_DMA)
+    if _pred > _LDS_LIMIT_BYTES and ENABLE_DMA:
+        # Try DMA off (single-buffered).
+        ENABLE_DMA = False
+        NUM_PREFETCH_K = 1
+        NUM_PREFETCH_V = 1
+        _pred2 = _predicted_lds_bytes(1, 1, False)
+        if _pred2 > _LDS_LIMIT_BYTES:
+            raise RuntimeError(
+                f"v4_hca_bwd_dq_pool: predicted LDS {_pred2}B > limit {_LDS_LIMIT_BYTES}B "
+                f"even single-buffered at D={head_dim}, BLOCK_N={BLOCK_N}"
+            )
+        import sys as _sys
+
+        print(
+            f"[v4_hca_bwd_dq_pool] LDS overflow at D={head_dim}, BLOCK_N={BLOCK_N}: "
+            f"auto-disabled DMA (predicted {_pred} -> {_pred2} bytes)",
+            file=_sys.stderr,
+            flush=True,
+        )
+    K_STEP_QK = 16 if USE_K16 else 8
+    K_STEPS_QK = head_dim // K_STEP_QK
+    D_CHUNK = 32
+    D_CHUNKS = head_dim // D_CHUNK
+    PV_K_STEP = 16 if USE_K16 else 8
+    PV_K_STEPS = K_SUB_N // PV_K_STEP  # 2 steps per sub-tile (K=16) or 4 (K=8)
+
+    assert BLOCK_M % NUM_WAVES == 0
+    assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32"
+    assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64"
+    assert flat_work_group_size in (
+        128,
+        256,
+        512,
+    ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}"
+    assert dtype_str in ("f16", "bf16"), "v4_hca_bwd_dq_pool only supports f16 and bf16"
+    assert BLOCK_N % 32 == 0
+    assert BLOCK_N_OUT == BLOCK_N
+    assert (
+        isinstance(pool_size, int) and 0 < pool_size <= BLOCK_N
+    ), f"pool_size must be int in (0, {BLOCK_N}], got {pool_size!r}"
+    assert (
+        isinstance(hca_local_seqlen, int) and hca_local_seqlen >= 0
+    ), f"hca_local_seqlen must be int >= 0, got {hca_local_seqlen!r}"
+    assert (
+        hca_local_seqlen % BLOCK_N == 0
+    ), f"hca_local_seqlen must be multiple of BLOCK_N={BLOCK_N}, got {hca_local_seqlen}"
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(head_dim)
+
+    NUM_HEADS = num_heads
+    HEAD_DIM = head_dim
+    STRIDE_TOKEN = NUM_HEADS * HEAD_DIM
+
+    K_STRIDE = HEAD_DIM
+    # XOR swizzle mask must fit within the row stride. The swizzle is
+    # applied at 16-element granularity, so the maximum mask is
+    # `(K_STRIDE // 16 - 1) << 4`. For D=128 that's 7 (=0x7), giving max
+    # mask 112 < 128. For D=64 it's 3 (=0x3), giving max mask 48 < 64.
+    # SLA test only covers D=128 (mask=7); using a hardcoded `& 7` for
+    # D=64 wraps writes into adjacent rows -> silent corruption.
+    K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1
+    assert K_SWZ_ROW_MASK >= 0
+    assert (
+        K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)
+    ) == 0, f"K_SWZ_ROW_MASK must be 2^n-1, got {K_SWZ_ROW_MASK} (K_STRIDE={K_STRIDE})"
+    if USE_HW_TR:
+        V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4
+    else:
+        VT_STRIDE = BLOCK_N + 2
+        V_STRIDE = VT_STRIDE
+    # V swizzle: similarly bounded by V_STRIDE / 16.
+    V_SWZ_ROW_MASK = min(3, (V_STRIDE // 16) - 1)
+    assert V_SWZ_ROW_MASK >= 0
+
+    VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+
+    if ROWS_PER_BATCH_LOAD >= BLOCK_N:
+        NUM_BATCHES_KV = 1
+        KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N
+    else:
+        assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+        NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+        KV_NEEDS_GUARD = False
+
+    LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE
+    if USE_HW_TR:
+        LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE
+    else:
+        LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE
+    LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE
+    LDS_V_BASE = LDS_K_TOTAL_SIZE
+    LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE
+    LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_hca_bwd_dq_pool_smem_M{BLOCK_M}_P{pool_size}_L{hca_local_seqlen}_MQ{int(mqa_kv)}",
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_hca_bwd_dq_pool_kernel(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,  # grad of output (input)
+        LSE: fx.Tensor,  # log-sum-exp from fwd (input, f32, RAW domain, JOINT)
+        DELTAS: fx.Tensor,  # (o * do).sum(-1) preprocess (input, f32)
+        DQ: fx.Tensor,  # grad wrt Q (output, FP32, ACCUMULATED)
+        ADD_MASK: fx.Tensor,  # additive pool mask (input, bf16/f16, [Sq, POOL_SIZE])
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K)
+        v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS)
+        dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ)
+        add_mask_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), ADD_MASK)
+        # LSE / DELTAS: f32 scalar READS via buffer_load.
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+
+        # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation).
+        fm_fast = arith.FastMathFlags.fast
+        v4f16_type = T.vec(4, elem_type)
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        v16f32_type = T.vec(16, compute_type)
+        mfma_pack_type = v8f16_type if USE_K16 else v4f16_type
+        MFMA_LANE_K = 8 if USE_K16 else 4
+        _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)
+
+        def _mfma(ods_fn, a, b, c):
+            return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result
+
+        def mfma_acc(a, b, c):
+            if dtype_str == "bf16":
+                if USE_K16:
+                    return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c)
+                a = vector.bitcast(T.i16x4, a)
+                b = vector.bitcast(T.i16x4, b)
+                return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c)
+            if USE_K16:
+                return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c)
+            return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c)
+
+        seq_len_q_v = arith.index_cast(T.index, seq_len_q)
+        seq_len_k_v = arith.index_cast(T.index, seq_len_k)
+
+        # ---- LDS view ----
+        base_ptr = allocator.get_base()
+        lds_kv = SmemPtr(
+            base_ptr,
+            lds_kv_offset,
+            elem_type,
+            shape=(LDS_KV_TOTAL_SIZE,),
+        ).get()
+
+        # ---- Thread / block indices ----
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        lane_mod_32 = lane % 32
+        lane_div_32 = lane // 32  # 0/1
+
+        # ds_read_b64_tr_b16 lane decomposition
+        tr_k_group = (lane % 16) // 4
+        tr_col_sub = lane % 4
+        tr_col_half = (lane % 32) // 16
+
+        def ds_read_tr_v4f16(lds_elem_idx):
+            byte_offset = lds_elem_idx * 2 + lds_kv_offset
+            byte_i64 = arith.index_cast(T.i64, byte_offset)
+            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result
+
+        wave_q_offset = wave_id * ROWS_PER_WAVE
+
+        # ---- Decompose block_id: (batch, q_tile, head) ----
+        head_idx = block_id % NUM_HEADS
+        batch_q_tile_id = block_id // NUM_HEADS
+        num_q_tiles = (seq_len_q_v + BLOCK_M - 1) // BLOCK_M
+        q_tile_idx = batch_q_tile_id % num_q_tiles
+        batch_idx = batch_q_tile_id // num_q_tiles
+        q_start = q_tile_idx * BLOCK_M
+
+        # ---- V4 HCA POOL: fixed K-block range [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE) ----
+        # POOL_SIZE <= BLOCK_N and HCA_LOCAL_SEQLEN % BLOCK_N == 0 (asserted
+        # at build time), so this is exactly ONE n-block per program.
+        BN = arith.index(BLOCK_N)
+        arith.index(BLOCK_M)
+        arith.index(0)
+        _one_idx = arith.index(1)
+        n_block_start = arith.index(hca_local_seqlen // BLOCK_N)
+        n_block_end = n_block_start + _one_idx
+
+        # ---- Cooperative load decomposition ----
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        # ---- Helper: global flat index ----
+        if layout_bhld:
+            bh_base_tokens_q = (batch_idx * NUM_HEADS + head_idx) * seq_len_q_v
+            if mqa_kv:
+                bh_base_tokens_kv = batch_idx * seq_len_k_v
+            else:
+                bh_base_tokens_kv = (batch_idx * NUM_HEADS + head_idx) * seq_len_k_v
+
+            def global_idx_q(token_idx, col):
+                return (bh_base_tokens_q + token_idx) * arith.index(HEAD_DIM) + col
+
+            def global_idx_kv(token_idx, col):
+                return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col
+
+        else:
+            # BLHD path: kept for symmetry but V4 uses BHLD.
+            def global_idx_q(token_idx, col):
+                token = batch_idx * seq_len_q_v + token_idx
+                return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col
+
+            if mqa_kv:
+
+                def global_idx_kv(token_idx, col):
+                    token = batch_idx * seq_len_k_v + token_idx
+                    return token * HEAD_DIM + col
+
+            else:
+
+                def global_idx_kv(token_idx, col):
+                    token = batch_idx * seq_len_k_v + token_idx
+                    return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col
+
+        def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=et,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr_, elem_idx):
+            """Store a single f32 value via GEP into an fp32 buffer."""
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, vxf16_type)
+
+        def bf16_trunc_pack_v4(f32_vals):
+            _v2i32 = T.vec(2, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32)
+            b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32)
+            p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result
+            a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32)
+            b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32)
+            p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result
+            return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1]))
+
+        def bf16_trunc_pack_v8(f32_vals):
+            _v4i32 = T.vec(4, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            pairs = []
+            for j in range_constexpr(4):
+                a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32)
+                b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32)
+                p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result
+                pairs.append(p)
+            return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs))
+
+        def k_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(buf_id * LDS_K_TILE_SIZE)
+            return buf_id * arith.index(LDS_K_TILE_SIZE)
+
+        def v_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE)
+            return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE)
+
+        # ---- K XOR swizzle: col ^ ((row & K_SWZ_ROW_MASK) << 4) at 16-element granularity ----
+        # K_SWZ_ROW_MASK derived from K_STRIDE to stay within the row.
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        # ---- Cooperative K load (row-major, XOR-swizzled) ----
+        def coop_load_k(tile_start, buf_id=0):
+            """K row-bounds-aware cooperative load.
+
+            Pool kernel reads BLOCK_N rows starting at kv_start; for pool
+            sizes < BLOCK_N the top rows are past seq_len_k and would read
+            garbage. Gate each load with row_idx < seq_len_k_v and store
+            zero into the LDS slot for OOB rows.
+            """
+            k_base = k_buf_base(buf_id)
+            c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid_blk = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(k_ptr, g_idx)
+                    vec = arith.select(row_valid, raw_vec, c_zero_vxf16)
+                    _if_k = scf.IfOp(
+                        arith.cmpi(
+                            arith.CmpIPredicate.ult,
+                            load_row_in_batch,
+                            arith.index(BLOCK_N),
+                        )
+                    )
+                    with ir.InsertionPoint(_if_k.then_block):
+                        vector.store(vec, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(k_ptr, g_idx)
+                    vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16)
+                    vector.store(vec, lds_kv, [lds_idx])
+
+        # ---- Cooperative V-into-K-LDS-slot load (K-style XOR swizzle) ----
+        def coop_load_v_as_k(tile_start, buf_id=0):
+            """V-into-K-slot with row bounds check (zero OOB rows)."""
+            k_base = k_buf_base(buf_id)
+            c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid_blk = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(v_ptr, g_idx)
+                    vec = arith.select(row_valid, raw_vec, c_zero_vxf16)
+                    _if_v = scf.IfOp(
+                        arith.cmpi(
+                            arith.CmpIPredicate.ult,
+                            load_row_in_batch,
+                            arith.index(BLOCK_N),
+                        )
+                    )
+                    with ir.InsertionPoint(_if_v.then_block):
+                        vector.store(vec, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(v_ptr, g_idx)
+                    vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16)
+                    vector.store(vec, lds_kv, [lds_idx])
+
+        # ---- Cooperative V load (V LDS layout) ----
+        def _v_store_row_major(v_base, lds_row, vec):
+            lds_idx = v_base + lds_row * V_STRIDE + load_col_base
+            vector.store(vec, lds_kv, [lds_idx])
+
+        _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None
+
+        def _v_store_transposed(v_base, lds_row, vec):
+            for _e in range_constexpr(VEC_WIDTH):
+                elem = vector.extract(vec, static_position=[_e], dynamic_position=[])
+                vt_d = load_col_base + _e
+                vt_idx = v_base + vt_d * VT_STRIDE + lds_row
+                v1 = vector.from_elements(_v1_type, [elem])
+                vector.store(v1, lds_kv, [vt_idx])
+
+        _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed
+
+        def coop_load_v(tile_start, buf_id=0):
+            v_base = v_buf_base(buf_id)
+            c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid_blk = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(v_ptr, g_idx)
+                    vec = arith.select(row_valid, raw_vec, c_zero_vxf16)
+                    _if_v = scf.IfOp(
+                        arith.cmpi(
+                            arith.CmpIPredicate.ult,
+                            load_row_in_batch,
+                            arith.index(BLOCK_N),
+                        )
+                    )
+                    with ir.InsertionPoint(_if_v.then_block):
+                        lds_row = load_row_in_batch + row_offset
+                        _v_store_to_lds(v_base, lds_row, vec)
+                        scf.YieldOp([])
+                else:
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(v_ptr, g_idx)
+                    vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16)
+                    lds_row = load_row_in_batch + row_offset
+                    _v_store_to_lds(v_base, lds_row, vec)
+
+        # ---- Cooperative K-into-V-LDS-slot load ----
+        def coop_load_k_as_v(tile_start, buf_id=0):
+            """K-into-V-slot with row bounds check (zero OOB rows)."""
+            v_base = v_buf_base(buf_id)
+            c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid_blk = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(k_ptr, g_idx)
+                    vec = arith.select(row_valid, raw_vec, c_zero_vxf16)
+                    _if_v = scf.IfOp(
+                        arith.cmpi(
+                            arith.CmpIPredicate.ult,
+                            load_row_in_batch,
+                            arith.index(BLOCK_N),
+                        )
+                    )
+                    with ir.InsertionPoint(_if_v.then_block):
+                        lds_row = load_row_in_batch + row_offset
+                        _v_store_to_lds(v_base, lds_row, vec)
+                        scf.YieldOp([])
+                else:
+                    row_valid_seq = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        row_idx,
+                        seq_len_k_v,
+                    )
+                    g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0))
+                    g_idx = global_idx_kv(g_idx_safe_row, load_col_base)
+                    raw_vec = load_global_f16xN(k_ptr, g_idx)
+                    vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16)
+                    lds_row = load_row_in_batch + row_offset
+                    _v_store_to_lds(v_base, lds_row, vec)
+
+        # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ----
+        if ENABLE_DMA:
+            from flydsl._mlir.dialects import llvm
+
+            k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True)
+            _lds_ptr_ty = _llvm_lds_ptr_ty()
+            DMA_BYTES = 16
+            DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES
+            K_TILE_BYTES = BLOCK_N * K_STRIDE * 2
+            NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+            lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv)
+            _dma_size = arith.constant(DMA_BYTES, type=T.i32)
+            _dma_soff = arith.constant(0, type=T.i32)
+            _dma_off = arith.constant(0, type=T.i32)
+            _dma_aux = arith.constant(1, type=T.i32)
+
+            def _kv_global_byte(tile_start, row_in_tile, col_byte):
+                if layout_bhld:
+                    row_within = tile_start + row_in_tile
+                    return (bh_base_tokens_kv + row_within) * arith.index(HEAD_DIM * 2) + col_byte
+                else:
+                    global_row = batch_idx * seq_len_k_v + tile_start + row_in_tile
+                    if mqa_kv:
+                        return global_row * arith.index(HEAD_DIM * 2) + col_byte
+                    else:
+                        return (
+                            global_row * arith.index(STRIDE_TOKEN * 2)
+                            + head_idx * arith.index(HEAD_DIM * 2)
+                            + col_byte
+                        )
+
+            def coop_dma_k(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2)
+                else:
+                    k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2)
+                for d in range_constexpr(NUM_DMA_K):
+                    lds_addr = (
+                        k_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH)
+                    swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        k_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        def _v_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        if ENABLE_DMA:
+            v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True)
+            V_TILE_BYTES = BLOCK_N * V_STRIDE * 2
+            NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+
+            def coop_dma_v(tile_start, buf_id=0):
+                v_lds_byte_base = lds_kv_base_idx + arith.index((LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2)
+                for d in range_constexpr(NUM_DMA_V):
+                    lds_addr = (
+                        v_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V)
+                    swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        v_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+            # ---- Bwd dQ DMA variants (cross-pointer, matching swizzle) ----
+            def coop_dma_v_as_k(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2)
+                else:
+                    k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2)
+                for d in range_constexpr(NUM_DMA_K):
+                    lds_addr = (
+                        k_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH)
+                    swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        v_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+            def coop_dma_k_as_v(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    v_lds_byte_base = lds_kv_base_idx + arith.index(
+                        (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2
+                    )
+                else:
+                    v_lds_byte_base = (
+                        lds_kv_base_idx
+                        + arith.index(LDS_V_BASE * 2)
+                        + buf_id * arith.index(LDS_V_TILE_SIZE * 2)
+                    )
+                for d in range_constexpr(NUM_DMA_V):
+                    lds_addr = (
+                        v_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V)
+                    swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        k_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        # ---- Preload Q^T B-operand and DO^T B-operand packs ----
+        q_row = q_start + wave_q_offset + lane_mod_32
+        q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_q_v)
+        q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0))
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        q_b_packs = []
+        do_b_packs = []
+        for ks in range_constexpr(K_STEPS_QK):
+            col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+            g_idx = global_idx_q(q_row_safe, col)
+            q_raw = load_global_mfma_pack(q_ptr, g_idx)
+            q_b_packs.append(arith.select(q_in_bounds, q_raw, c_zero_mfma_pack))
+            do_raw = load_global_mfma_pack(do_ptr, g_idx)
+            do_b_packs.append(arith.select(q_in_bounds, do_raw, c_zero_mfma_pack))
+
+        # ---- Constants ----
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)  # finite NEG_INF (matches Triton)
+        c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type)
+        # V4 LSE is RAW-domain (qk*sm_scale + ln(l)). So use sm_scale, not sm_scale*log2e.
+        c_sm_scale = arith.constant(sm_scale, type=compute_type)
+
+        # ---- Per-q-row scalars (LSE, delta) ----
+        # bh_base_tokens_q is the Q LSE/DELTAS base too (LSE shape [B,HQ,Sq]).
+        lse_delta_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q + q_row_safe)
+        lse_val = buffer_ops.buffer_load(
+            lse_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=T.f32,
+        )
+        delta_val = buffer_ops.buffer_load(
+            deltas_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=T.f32,
+        )
+
+        # ---- (No SINK contribution: pool kernel does not touch sink.) ----
+
+        # ---- POOL: single n-block; double-buffered LDS optional ----
+        _use_dbuf = ENABLE_DMA
+
+        init_args = []
+        for _ in range_constexpr(D_CHUNKS):
+            init_args.append(c_zero_v16f32)
+        if _use_dbuf:
+            init_args.append(arith.index(0))  # cur_buf_id
+
+            # PROLOGUE: prefetch iter 0's K into BOTH slots of buf 0.
+            _init_kv_start = n_block_start * BN
+            coop_dma_k(_init_kv_start, buf_id=0)
+            coop_dma_k_as_v(_init_kv_start, buf_id=0)
+
+        for block_idx, inner_iter_args, loop_results in scf.for_(
+            n_block_start,
+            n_block_end,
+            arith.index(1),
+            iter_args=init_args,
+        ):
+            dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)]
+            if _use_dbuf:
+                cur_buf = inner_iter_args[D_CHUNKS]
+                next_buf = arith.index(1) - cur_buf
+
+            # V4 SWA: block_idx is directly the K-block index.
+            kv_block_start = block_idx * BN
+            kv_start = kv_block_start
+
+            if _use_dbuf:
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+                k_base = k_buf_base(cur_buf)
+            else:
+                coop_load_k(kv_start, buf_id=0)
+                coop_load_k_as_v(kv_start, buf_id=0)
+                gpu.barrier()
+                k_base = k_buf_base(0)
+
+            # ==== GEMM1: s = Q @ K^T ====
+            k_hi_offset = K_SUB_N * K_STRIDE
+            k_swz_mask = (lane_mod_32 & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+
+            def _k_idx_lo(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+            def _k_idx_hi(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+            s_acc_lo = c_zero_v16f32
+            s_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                k_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)])
+                k_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)])
+                s_acc_lo = mfma_acc(k_pack_lo, q_b_packs[ks], s_acc_lo)
+                s_acc_hi = mfma_acc(k_pack_hi, q_b_packs[ks], s_acc_hi)
+
+            # ==== Compute p[r] = exp((qk + add_bias) * sm_scale_inline - LSE) ====
+            #
+            # POOL-only mask. The s_acc registers map (per lane, per reg r) to
+            # the S = Q @ K^T cell at row = q_row (= lane_mod_32 + wave_q_offset
+            # + q_start) and column = kv_start + lane_div_32*4 + (r//4)*8 + r%4
+            # (lo half), or +K_SUB_N (hi half). Translate to pool_n by
+            # subtracting HCA_LOCAL_SEQLEN; gate with pool_n < POOL_SIZE and
+            # q_row < seq_len_q.
+            seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v)
+            q_row_i32_mask = arith.index_cast(T.i32, q_row)
+            lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32)
+            lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result
+            hca_local_i32 = arith.constant(hca_local_seqlen, type=T.i32)
+            pool_size_i32 = arith.constant(pool_size, type=T.i32)
+            kv_start_i32 = arith.index_cast(T.i32, kv_start)
+            # pool_n base for this lane = kv_start + lane_off - hca_local_seqlen
+            # (kv_start = hca_local_seqlen + 0 for the single pool block).
+            pool_n_base_i32 = arith.SubIOp(
+                arith.AddIOp(kv_start_i32, lane_off_i32).result,
+                hca_local_i32,
+            ).result
+
+            # Is this q_row out of bounds? -> entire row masked.
+            q_oob = arith.cmpi(
+                arith.CmpIPredicate.sge,
+                q_row_i32_mask,
+                seq_len_q_i32,
+            )
+
+            # --- Load add_bias from ADD_MASK[q_row, pool_n] via GEP ---
+            # ADD_MASK shape [Sq, POOL_SIZE] contiguous; element bytes = 2
+            # (bf16 / f16). For each (lane, r), compute pool_n (lo or hi) and
+            # load a single bf16/f16, then cast to f32.
+            ADD_MASK_STRIDE_M = arith.index(pool_size)  # row stride in elements
+            q_row_idx_for_mask = arith.select(q_oob, arith.index(0), q_row)
+            # Precompute add-mask row base offset in elements.
+            add_mask_row_base = q_row_idx_for_mask * ADD_MASK_STRIDE_M
+
+            def _load_add_bias(pool_n_i32, bad_pred):
+                """Load ADD_MASK[q_row, pool_n] as f32; return 0.0 if bad."""
+                # Convert pool_n to index for GEP.
+                pool_n_i32_safe = arith.select(
+                    bad_pred,
+                    arith.constant(0, type=T.i32),
+                    pool_n_i32,
+                )
+                pool_n_idx = arith.index_cast(T.index, pool_n_i32_safe)
+                elem_idx = add_mask_row_base + pool_n_idx
+                bias_raw_v1 = _gep_load(
+                    add_mask_ptr,
+                    elem_idx,
+                    T.vec(1, elem_type),
+                )
+                bias_raw = vector.extract(
+                    bias_raw_v1,
+                    static_position=[0],
+                    dynamic_position=[],
+                )
+                bias_f32 = arith.extf(compute_type, bias_raw)
+                # If bad (OOB pool_n or q_oob), zero out the bias (mask wins
+                # anyway via NEG_INF, but keep numerics tidy).
+                return arith.select(bad_pred, c_zero_f, bias_f32)
+
+            p_vals_lo = []
+            p_vals_hi = []
+            for r in range_constexpr(16):
+                r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32)
+
+                # --- lo half ---
+                pool_n_lo_i32 = arith.AddIOp(
+                    pool_n_base_i32,
+                    r_off_i32,
+                ).result
+                is_pool_oob_lo = arith.cmpi(
+                    arith.CmpIPredicate.sge,
+                    pool_n_lo_i32,
+                    pool_size_i32,
+                )
+                bad_lo = arith.OrIOp(is_pool_oob_lo, q_oob).result
+
+                s_lo_f32 = vector.extract(s_acc_lo, static_position=[r], dynamic_position=[])
+                bias_lo = _load_add_bias(pool_n_lo_i32, bad_lo)
+                qk_plus_bias_lo = arith.AddFOp(s_lo_f32, bias_lo, fastmath=fm_fast).result
+                # Triton: qk = qk * sm_scale; qk = qk + add_bias.
+                # We deferred sm_scale to here: (qk + bias) is wrong; need
+                # qk*scale + bias. Re-do as: s_scaled + bias.
+                scaled_lo = arith.MulFOp(s_lo_f32, c_sm_scale, fastmath=fm_fast).result
+                scaled_plus_bias_lo = arith.AddFOp(scaled_lo, bias_lo, fastmath=fm_fast).result
+                scaled_lo_masked = arith.select(
+                    bad_lo,
+                    c_neg_inf,
+                    scaled_plus_bias_lo,
+                )
+                diff_lo = arith.SubFOp(scaled_lo_masked, lse_val, fastmath=fm_fast).result
+                p_lo = math_dialect.exp(diff_lo, fastmath=fm_fast)
+                p_vals_lo.append(p_lo)
+
+                # --- hi half: pool_n_hi = pool_n_lo + K_SUB_N ---
+                pool_n_hi_i32 = arith.AddIOp(
+                    pool_n_lo_i32,
+                    arith.constant(K_SUB_N, type=T.i32),
+                ).result
+                is_pool_oob_hi = arith.cmpi(
+                    arith.CmpIPredicate.sge,
+                    pool_n_hi_i32,
+                    pool_size_i32,
+                )
+                bad_hi = arith.OrIOp(is_pool_oob_hi, q_oob).result
+
+                s_hi_f32 = vector.extract(s_acc_hi, static_position=[r], dynamic_position=[])
+                bias_hi = _load_add_bias(pool_n_hi_i32, bad_hi)
+                scaled_hi = arith.MulFOp(s_hi_f32, c_sm_scale, fastmath=fm_fast).result
+                scaled_plus_bias_hi = arith.AddFOp(scaled_hi, bias_hi, fastmath=fm_fast).result
+                scaled_hi_masked = arith.select(
+                    bad_hi,
+                    c_neg_inf,
+                    scaled_plus_bias_hi,
+                )
+                diff_hi = arith.SubFOp(scaled_hi_masked, lse_val, fastmath=fm_fast).result
+                p_hi = math_dialect.exp(diff_hi, fastmath=fm_fast)
+                p_vals_hi.append(p_hi)
+
+            # ==== Overwrite K LDS slot with V for GEMM2 ====
+            gpu.barrier()
+            if _use_dbuf:
+                coop_dma_v_as_k(kv_start, buf_id=cur_buf)
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+            elif ENABLE_DMA:
+                coop_dma_v_as_k(kv_start, buf_id=0)
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+            else:
+                coop_load_v_as_k(kv_start, buf_id=0)
+                gpu.barrier()
+
+            # ==== GEMM2: dP = DO @ V^T ====
+            dp_acc_lo = c_zero_v16f32
+            dp_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                v_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)])
+                v_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)])
+                dp_acc_lo = mfma_acc(v_pack_lo, do_b_packs[ks], dp_acc_lo)
+                dp_acc_hi = mfma_acc(v_pack_hi, do_b_packs[ks], dp_acc_hi)
+
+            # ==== Prefetch iter+1's K DMAs (async) ====
+            if _use_dbuf:
+                _next_block_idx = block_idx + arith.index(1)
+                _has_next = arith.cmpi(arith.CmpIPredicate.slt, _next_block_idx, n_block_end)
+                _pre_if = scf.IfOp(_has_next)
+                with ir.InsertionPoint(_pre_if.then_block):
+                    _next_kv_start = _next_block_idx * BN
+                    coop_dma_k(_next_kv_start, next_buf)
+                    coop_dma_k_as_v(_next_kv_start, next_buf)
+                    scf.YieldOp([])
+
+            # ==== Compute dS[r] = p[r] * (dp[r] - delta) ====
+            ds_vals_lo = []
+            ds_vals_hi = []
+            for r in range_constexpr(16):
+                dp_lo = vector.extract(dp_acc_lo, static_position=[r], dynamic_position=[])
+                dp_hi = vector.extract(dp_acc_hi, static_position=[r], dynamic_position=[])
+                diff_lo = arith.SubFOp(dp_lo, delta_val, fastmath=fm_fast).result
+                diff_hi = arith.SubFOp(dp_hi, delta_val, fastmath=fm_fast).result
+                ds_lo = arith.MulFOp(p_vals_lo[r], diff_lo, fastmath=fm_fast).result
+                ds_hi = arith.MulFOp(p_vals_hi[r], diff_hi, fastmath=fm_fast).result
+                ds_vals_lo.append(ds_lo)
+                ds_vals_hi.append(ds_hi)
+
+            # ==== Pack dS f32 -> mfma_pack_type (bf16/f16) ====
+            if dtype_str == "bf16" and USE_K16:
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    base = pks * 8
+                    ds_packs_lo.append(bf16_trunc_pack_v8(ds_vals_lo[base : base + 8]))
+                    ds_packs_hi.append(bf16_trunc_pack_v8(ds_vals_hi[base : base + 8]))
+            elif dtype_str == "bf16":
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    base = pks * 4
+                    ds_packs_lo.append(bf16_trunc_pack_v4(ds_vals_lo[base : base + 4]))
+                    ds_packs_hi.append(bf16_trunc_pack_v4(ds_vals_hi[base : base + 4]))
+            else:
+                ds_f16_lo = [arith.trunc_f(elem_type, ds_vals_lo[r]) for r in range_constexpr(16)]
+                ds_f16_hi = [arith.trunc_f(elem_type, ds_vals_hi[r]) for r in range_constexpr(16)]
+                _pack_ty = v8f16_type if USE_K16 else v4f16_type
+                _pw = 8 if USE_K16 else 4
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    b = pks * _pw
+                    ds_packs_lo.append(vector.from_elements(_pack_ty, [ds_f16_lo[b + i] for i in range(_pw)]))
+                    ds_packs_hi.append(vector.from_elements(_pack_ty, [ds_f16_hi[b + i] for i in range(_pw)]))
+
+            # ==== GEMM3: dQ += K @ dS  (uses fwd's V-read schedule with K substituted) ====
+            _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)]
+            TOTAL_PV = len(_steps)
+            v_base = v_buf_base(cur_buf) if _use_dbuf else v_buf_base(0)
+
+            def _read_k_as_v_pack(step_idx):
+                dc, pks = _steps[step_idx]
+                if USE_HW_TR:
+                    d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4
+                    k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group
+                    _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col
+                    lds_lo = v_base + k_row * V_STRIDE + _d_col_eff
+                    lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE)
+                    if USE_K16:
+                        vl_a = ds_read_tr_v4f16(lds_lo)
+                        vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE))
+                        vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                        vh_a = ds_read_tr_v4f16(lds_hi)
+                        vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE))
+                        vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                    else:
+                        vl = ds_read_tr_v4f16(lds_lo)
+                        vh = ds_read_tr_v4f16(lds_hi)
+                else:
+                    d_pos = arith.index(dc * D_CHUNK) + lane_mod_32
+                    kb = arith.index(pks * PV_K_STEP) + lane_div_32 * 4
+                    v_lo_idx = v_base + d_pos * VT_STRIDE + kb
+                    v_hi_idx = v_lo_idx + arith.index(K_SUB_N)
+                    vl = vector.load(v4f16_type, lds_kv, [v_lo_idx])
+                    vh = vector.load(v4f16_type, lds_kv, [v_hi_idx])
+                return vl, vh
+
+            k_lo_cur, k_hi_cur = _read_k_as_v_pack(0)
+            for si in range_constexpr(TOTAL_PV):
+                dc, pks = _steps[si]
+                if si + 1 < TOTAL_PV:
+                    k_lo_nxt, k_hi_nxt = _read_k_as_v_pack(si + 1)
+                dq_accs[dc] = mfma_acc(k_lo_cur, ds_packs_lo[pks], dq_accs[dc])
+                dq_accs[dc] = mfma_acc(k_hi_cur, ds_packs_hi[pks], dq_accs[dc])
+                if si + 1 < TOTAL_PV:
+                    k_lo_cur = k_lo_nxt
+                    k_hi_cur = k_hi_nxt
+
+            # End of iter: yield dq_accs and swapped cur_buf.
+            gpu.barrier()
+            _yield = list(dq_accs)
+            if _use_dbuf:
+                _yield.append(next_buf)
+            yield _yield
+
+        # ---- Final store: dQ_fp32 += dq_acc * sm_scale (load + add + store) ----
+        # Pool kernel ACCUMULATES into the existing LOCAL-stream dq buffer.
+        # Race-free because each program owns a unique (b, qhid, m_block)
+        # slice, and the launcher runs this kernel sequentially AFTER the
+        # SWA dq kernel has finished writing the LOCAL contribution.
+        dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)]
+        sm_scale_vec = vector.broadcast(v16f32_type, c_sm_scale)
+
+        def _gep_load_f32(base_ptr_, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(T.f32, gep.result).result
+
+        _o_guard = scf.IfOp(q_in_bounds, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for dc in range_constexpr(D_CHUNKS):
+                dq_scaled = arith.MulFOp(
+                    dq_finals[dc],
+                    sm_scale_vec,
+                    fastmath=fm_fast,
+                ).result
+                for r in range_constexpr(16):
+                    dq_val = vector.extract(
+                        dq_scaled,
+                        static_position=[r],
+                        dynamic_position=[],
+                    )
+                    d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4)
+                    d_col = arith.index(dc * D_CHUNK) + d_row_rel
+                    dq_global = global_idx_q(q_row, d_col)
+                    dq_prev = _gep_load_f32(dq_ptr, dq_global)
+                    dq_sum = arith.AddFOp(
+                        dq_val,
+                        dq_prev,
+                        fastmath=fm_fast,
+                    ).result
+                    _gep_store_f32(dq_sum, dq_ptr, dq_global)
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_hca_bwd_dq_pool(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DQ: fx.Tensor,
+        ADD_MASK: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len_q)
+        num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M
+        grid_x = bs_idx * num_q_tiles * NUM_HEADS
+
+        launcher = v4_hca_bwd_dq_pool_kernel(
+            Q,
+            K,
+            V,
+            DOS,
+            LSE,
+            DELTAS,
+            DQ,
+            ADD_MASK,
+            seq_len_q,
+            seq_len_k,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(
+                            T.i32,
+                            _wpe,
+                        )
+        if flat_work_group_size is not None:
+            _fwgs = int(flat_work_group_size)
+            if _fwgs >= 1:
+                flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, 1, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    _fmha_compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {
+            "enable-post-misched": False,
+            "lsr-drop-solution": True,
+        },
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return launch_v4_hca_bwd_dq_pool(*args, **kwargs)
+
+    def _compile(Q, K, V, DOS, LSE, DELTAS, DQ, ADD_MASK, batch_size, seq_len_q, seq_len_k, stream=None):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return flyc.compile(
+                launch_v4_hca_bwd_dq_pool,
+                Q,
+                K,
+                V,
+                DOS,
+                LSE,
+                DELTAS,
+                DQ,
+                ADD_MASK,
+                batch_size,
+                seq_len_q,
+                seq_len_k,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+# Convenience alias.
+build_v4_hca_bwd_dq_pool_module_primary = build_v4_hca_bwd_dq_pool_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py
new file mode 100644
index 000000000..eda441de5
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py
@@ -0,0 +1,1113 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_swa_bwd_dkv: V4 SWA-causal attention backward dK/dV kernel for FlyDSL.
+
+Strategy:
+  * One workgroup per (batch, n_block). K and V loaded once into LDS
+    and stay resident for the program's lifetime (no LUT, no atomics).
+  * Inner loop iterates qhid in [0, HQ); for each head, iterates m-blocks
+    bounded by the SWA window. dK and dV f32 accumulators live in VGPRs.
+  * MQA head-loop accumulator: one dk/dv slice per (b, n_block) gets
+    all HQ heads accumulated -- deterministic, no atomics.
+  * SWA + causal mask per element. LSE is RAW-domain
+    (qk*sm_scale + ln(l)); p = exp(qk*sm_scale - lse).
+  * sm_scale applied to dK exactly once post-loop (Triton P57 cr=0).
+  * dK / dV written as f32 (matches launcher's dk_fp32 / dv_fp32 buffers).
+
+LDS budget at BLOCK_N=32, BLOCK_M2=32, D=512:
+    K = 32 KB, V = 32 KB, DO = 32 KB, Q = 32 KB, pT = 2 KB
+    Total = 130 KB <= 160 KB.
+Auto-fallback: if predicted > 160 KB, drop DO/Q LDS scratches and
+read Q/DO directly to register packs from HBM (mirrors STEP-1b dq).
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME = "v4_swa_bwd_dkv_kernel"
+
+_LLVM_GEP_DYNAMIC = -2147483648
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+def build_v4_swa_bwd_dkv_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    flat_work_group_size=None,
+    block_n=None,
+    block_m2=None,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    layout_bhld=True,
+    mqa_kv=True,
+):
+    """Build the V4 SWA backward dK/dV launcher (MQA only)."""
+    gpu_arch = get_hip_arch()
+
+    if block_n is None:
+        # R6k-1: D-split wave remap default. BLOCK_N=16 halves per-lane
+        # fp32 accumulator footprint (256 -> 128 fp32) by partitioning D
+        # across waves instead of N. Both waves cover the same 16 KV
+        # rows; each wave owns half of the D-chunks of dk/dv.
+        BLOCK_N = 16
+    else:
+        BLOCK_N = int(block_n)
+    if block_m2 is None:
+        BLOCK_M2 = 32
+    else:
+        BLOCK_M2 = int(block_m2)
+    WARP_SIZE = 64
+    # R6k-1: D-split mode. When BLOCK_N == 16 (matches MFMA-N tile),
+    # force NUM_WAVES=2 and partition D between the two waves instead
+    # of partitioning N. Each wave then carries only half the dk/dv
+    # accumulator footprint. For BLOCK_N > 16 we keep the legacy N-split.
+    D_SPLIT_WAVES = BLOCK_N == 16
+    if D_SPLIT_WAVES:
+        NUM_WAVES = 2
+        ROWS_PER_WAVE = 16  # both waves work on the same 16 rows
+    else:
+        # Legacy N-split: each wave gets BLOCK_N/NUM_WAVES rows.
+        NUM_WAVES = max(1, BLOCK_N // 16)
+        ROWS_PER_WAVE = BLOCK_N // NUM_WAVES
+    if flat_work_group_size is None:
+        flat_work_group_size = NUM_WAVES * WARP_SIZE
+    BLOCK_SIZE = flat_work_group_size
+
+    ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1"
+    USE_K16 = gpu_arch.startswith("gfx950")
+    # R6b: MFMA 16x16x32 K-step = 32; A/B-frag = 8 bf16 per lane.
+    assert USE_K16, "R6b dkv requires gfx950 (MFMA 16x16x32 bf16)."
+    K_STEP_QK = 32
+    K_STEPS_QK = head_dim // K_STEP_QK
+    # R6b: each MFMA 16x16x32 tile produces a 16-wide N-chunk; D_CHUNK = 16 cols.
+    D_CHUNK = 16
+    D_CHUNKS = head_dim // D_CHUNK
+    # R6k-1: per-wave D-chunk count. In D-split mode each wave owns half.
+    if D_SPLIT_WAVES:
+        assert D_CHUNKS % NUM_WAVES == 0, (
+            f"D_CHUNKS ({D_CHUNKS}) must be divisible by NUM_WAVES " f"({NUM_WAVES}) for D-split mode."
+        )
+        D_CHUNKS_LOCAL = D_CHUNKS // NUM_WAVES
+    else:
+        D_CHUNKS_LOCAL = D_CHUNKS
+    K_STEPS_PT = BLOCK_M2 // K_STEP_QK
+    # R6b: GEMM1/GEMM3 cover m_col [0..BLOCK_M2) with multiple 16-col MFMA-N tiles per ks.
+    assert BLOCK_M2 % 16 == 0, f"BLOCK_M2 must be a multiple of 16 (MFMA-N), got {BLOCK_M2}"
+    M_TILES = BLOCK_M2 // 16
+
+    assert BLOCK_N % NUM_WAVES == 0
+    assert ROWS_PER_WAVE == 16, f"ROWS_PER_WAVE must equal 16 for MFMA 16x16x32, got {ROWS_PER_WAVE}"
+    assert head_dim % 32 == 0
+    assert head_dim >= 64
+    assert flat_work_group_size in (64, 128, 256, 512)
+    assert dtype_str == "bf16", "R6b dkv currently only supports bf16."
+    assert BLOCK_N % 16 == 0
+    assert BLOCK_M2 % K_STEP_QK == 0, (
+        f"BLOCK_M2 ({BLOCK_M2}) must be a multiple of MFMA K-step " f"({K_STEP_QK}) for the dV/dK GEMMs."
+    )
+    assert isinstance(swa_window, int) and swa_window > 0
+    assert mqa_kv, "v4_swa_bwd_dkv currently only supports MQA (HK=1)."
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(head_dim)
+
+    NUM_HEADS = num_heads
+    HEAD_DIM = head_dim
+
+    K_STRIDE = HEAD_DIM
+    K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1
+    assert K_SWZ_ROW_MASK >= 0
+    assert (K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)) == 0
+    V_STRIDE = HEAD_DIM
+
+    VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+
+    LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE
+    LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE
+    LDS_DO_STRIDE = HEAD_DIM
+    LDS_DO_ELEMS = BLOCK_M2 * LDS_DO_STRIDE
+    LDS_Q_STRIDE = HEAD_DIM
+    LDS_Q_ELEMS = BLOCK_M2 * LDS_Q_STRIDE
+    LDS_PT_STRIDE = BLOCK_M2
+    LDS_PT_ELEMS = BLOCK_N * LDS_PT_STRIDE
+
+    _LDS_LIMIT_BYTES = 160 * 1024
+
+    def _predict_lds_bytes(use_lds_for_q_do):
+        b = (LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + LDS_PT_ELEMS) * 2
+        if use_lds_for_q_do:
+            b += (LDS_DO_ELEMS + LDS_Q_ELEMS) * 2
+        return b
+
+    USE_LDS_FOR_Q_DO = True
+    _pred_full = _predict_lds_bytes(True)
+    if _pred_full > _LDS_LIMIT_BYTES:
+        USE_LDS_FOR_Q_DO = False
+        _pred_min = _predict_lds_bytes(False)
+        if _pred_min > _LDS_LIMIT_BYTES:
+            raise RuntimeError(
+                f"v4_swa_bwd_dkv: minimal LDS {_pred_min}B > limit "
+                f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, "
+                f"BLOCK_M2={BLOCK_M2}."
+            )
+        import sys as _sys
+
+        print(
+            f"[v4_swa_bwd_dkv] auto-fallback: LDS {_pred_full}B > "
+            f"{_LDS_LIMIT_BYTES}B; disabling Q/DO LDS "
+            f"({_pred_full} -> {_pred_min})",
+            file=_sys.stderr,
+            flush=True,
+        )
+    else:
+        import sys as _sys
+
+        print(
+            f"[v4_swa_bwd_dkv] LDS OK: predicted {_pred_full}B <= "
+            f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, "
+            f"BLOCK_M2={BLOCK_M2}, USE_LDS_FOR_Q_DO=True",
+            file=_sys.stderr,
+            flush=True,
+        )
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=(
+            f"v4_swa_bwd_dkv_smem_N{BLOCK_N}_M2_{BLOCK_M2}_W{swa_window}"
+            f"_MQ{int(mqa_kv)}_QDO{int(USE_LDS_FOR_Q_DO)}"
+            f"_DS{int(D_SPLIT_WAVES)}"
+        ),
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    LDS_V_BASE = LDS_K_TILE_SIZE
+    LDS_KV_TOTAL_SIZE = LDS_K_TILE_SIZE + LDS_V_TILE_SIZE
+    allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2
+
+    lds_pt_offset = allocator._align(allocator.ptr, 16)
+    allocator.ptr = lds_pt_offset + LDS_PT_ELEMS * 2
+
+    if USE_LDS_FOR_Q_DO:
+        lds_do_offset = allocator._align(allocator.ptr, 16)
+        allocator.ptr = lds_do_offset + LDS_DO_ELEMS * 2
+        lds_q_offset = allocator._align(allocator.ptr, 16)
+        allocator.ptr = lds_q_offset + LDS_Q_ELEMS * 2
+    else:
+        lds_do_offset = None
+        lds_q_offset = None
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_swa_bwd_dkv_kernel(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DK: fx.Tensor,
+        DV: fx.Tensor,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K)
+        v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS)
+        dk_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DK)
+        dv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DV)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+
+        fm_fast = arith.FastMathFlags.fast
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        # R6M: tr16 returns 4 bf16 per lane (v4f16-typed). Two calls + shuffle
+        # form the MFMA 16x16x32 B-frag (8 bf16) via HW transpose.
+        v4f16_type = T.vec(4, elem_type)
+        # R6b: MFMA 16x16x32 C-frag = 4 fp32 per lane.
+        v4f32_type = T.vec(4, compute_type)
+        mfma_pack_type = v8f16_type
+        MFMA_LANE_K = 8
+        v1_elem_type = T.vec(1, elem_type)
+        _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)
+
+        def mfma_acc(a, b, c):
+            # rocdl.mfma_f32_16x16x32_bf16 is the wrapped form: takes
+            # (result_type, [operands]) and returns the Value directly.
+            return rocdl.mfma_f32_16x16x32_bf16(
+                v4f32_type,
+                [a, b, c, _mfma_zero, _mfma_zero, _mfma_zero],
+            )
+
+        seq_len_q_v = arith.index_cast(T.index, seq_len_q)
+        seq_len_k_v = arith.index_cast(T.index, seq_len_k)
+
+        base_ptr = allocator.get_base()
+        lds_kv = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_TOTAL_SIZE,)).get()
+        lds_pt = SmemPtr(base_ptr, lds_pt_offset, elem_type, shape=(LDS_PT_ELEMS,)).get()
+        if USE_LDS_FOR_Q_DO:
+            lds_do = SmemPtr(base_ptr, lds_do_offset, elem_type, shape=(LDS_DO_ELEMS,)).get()
+            lds_q = SmemPtr(base_ptr, lds_q_offset, elem_type, shape=(LDS_Q_ELEMS,)).get()
+
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        # R6b MFMA 16x16x32 lane decomposition:
+        #   A-frag: A[lane_mod_16, ks*32 + lane_div_16*8 + 0..7]
+        #   B-frag: B[ks*32 + lane_div_16*8 + 0..7, lane_mod_16]
+        #   C-frag: C[lane_div_16*4 + ii, lane_mod_16]  for ii in 0..3
+        lane_mod_16 = lane % 16
+        lane_div_16 = lane // 16
+        # R6M: ds_read_tr16_b64 per-lane decomposition (matches proto).
+        tr_k_group = lane_mod_16 // arith.index(4)
+        tr_col_sub = lane_mod_16 % arith.index(4)
+
+        # R6k-1: in D-split mode both waves cover the same 16 KV rows
+        # (wave_n_offset == 0) and instead carry disjoint D-chunks of
+        # dk/dv. wave_d_col_offset is the wave's D-column base offset
+        # (0 for wave 0, D_CHUNKS_LOCAL*D_CHUNK for wave 1).
+        if D_SPLIT_WAVES:
+            wave_n_offset = arith.index(0)
+            wave_d_col_offset = wave_id * arith.index(D_CHUNKS_LOCAL * D_CHUNK)
+        else:
+            wave_n_offset = wave_id * ROWS_PER_WAVE
+            wave_d_col_offset = arith.index(0)
+
+        # R6S: Hoist a SINGLE per-lane base byte-pointer for DO/Q LDS
+        # tr16 reads. Each ds_read_tr16_b64 then uses
+        # `base + constexpr byte imm` so LLVM ISel folds the
+        # (m_step, dc, addr_0/1) deltas into the `offset:` immediate
+        # of the LDS instruction. Removes the per-call IntToPtrOp ->
+        # 32 AGPR spill chain that was throttling GEMM2/4 to 28.6%
+        # MFMA duty (vs Triton's 53.7%).
+        if USE_LDS_FOR_Q_DO:
+            assert LDS_DO_STRIDE == LDS_Q_STRIDE, "R6S tr16 base+imm assumes DO and Q share STRIDE"
+            tr_lane_base_elem = (
+                lane_div_16 * arith.index(MFMA_LANE_K * LDS_DO_STRIDE)
+                + tr_k_group * arith.index(LDS_DO_STRIDE)
+                + wave_d_col_offset
+                + tr_col_sub * arith.index(4)
+            )
+            tr_lane_base_byte = tr_lane_base_elem * arith.index(2)
+            tr_lane_base_i64 = arith.index_cast(T.i64, tr_lane_base_byte)
+            tr_lane_base_do_byte_i64 = arith.AddIOp(
+                tr_lane_base_i64,
+                arith.constant(lds_do_offset, type=T.i64),
+            ).result
+            tr_lane_base_q_byte_i64 = arith.AddIOp(
+                tr_lane_base_i64,
+                arith.constant(lds_q_offset, type=T.i64),
+            ).result
+            tr_lane_base_do_ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), tr_lane_base_do_byte_i64).result
+            tr_lane_base_q_ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), tr_lane_base_q_byte_i64).result
+
+        # 6N-2 iter2: de-dup helper (Q/DO B-frag loads also gated). Each m-tile mt is owned by wave (mt % NUM_WAVES).
+        # In D-split mode the owning wave runs GEMM1/softmax/pT-write/GEMM3/dsT
+        # for that mt; the other wave skips. GEMM2/4 use the SHARED pT/dsT LDS.
+        if D_SPLIT_WAVES:
+            owns_mt_preds = [
+                arith.cmpi(arith.CmpIPredicate.eq, wave_id, arith.index(mt % NUM_WAVES))
+                for mt in range(M_TILES)
+            ]
+        else:
+            owns_mt_preds = None
+
+        num_n_blocks = (seq_len_k_v + BLOCK_N - 1) // BLOCK_N
+        n_block_idx = block_id % num_n_blocks
+        batch_idx = block_id // num_n_blocks
+        kv_start = n_block_idx * arith.index(BLOCK_N)
+
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        # MQA: KV stride drops head.
+        bh_base_tokens_kv = batch_idx * seq_len_k_v
+
+        def bh_base_tokens_q_of(qhid_index):
+            return (batch_idx * NUM_HEADS + qhid_index) * seq_len_q_v
+
+        def global_idx_q(qhid_index, token_idx, col):
+            return (bh_base_tokens_q_of(qhid_index) + token_idx) * arith.index(HEAD_DIM) + col
+
+        def global_idx_kv(token_idx, col):
+            return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col
+
+        def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=et,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr_, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, vxf16_type)
+
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        if ROWS_PER_BATCH_LOAD >= BLOCK_N:
+            NUM_BATCHES_KV = 1
+            KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N
+        else:
+            assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+            NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+            KV_NEEDS_GUARD = False
+
+        if ROWS_PER_BATCH_LOAD >= BLOCK_M2:
+            NUM_BATCHES_M = 1
+            M_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_M2
+        else:
+            assert BLOCK_M2 % ROWS_PER_BATCH_LOAD == 0
+            NUM_BATCHES_M = BLOCK_M2 // ROWS_PER_BATCH_LOAD
+            M_NEEDS_GUARD = False
+
+        c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        c_zero_elem = arith.constant(0.0, type=elem_type)
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)
+        c_sm_scale = arith.constant(sm_scale, type=compute_type)
+        # R6b: 4-elem fp32 acc per MFMA 16x16x32 op.
+        v4f32_zero = arith.constant_vector(0.0, v4f32_type)
+
+        def coop_load_k(tile_start):
+            k_base = arith.index(0)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_k_v)
+                row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                g_idx = global_idx_kv(row_safe, load_col_base)
+                vec = load_global_f16xN(k_ptr, g_idx)
+                vec_safe = arith.select(in_bounds, vec, c_zero_vxf16)
+                lds_row = load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N)
+                    )
+                    _if_k = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_k.then_block):
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                        vector.store(vec_safe, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    vector.store(vec_safe, lds_kv, [lds_idx])
+
+        def coop_load_v(tile_start):
+            v_base = arith.index(LDS_V_BASE)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_k_v)
+                row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                g_idx = global_idx_kv(row_safe, load_col_base)
+                vec = load_global_f16xN(v_ptr, g_idx)
+                vec_safe = arith.select(in_bounds, vec, c_zero_vxf16)
+                lds_row = load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N)
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = v_base + lds_row * V_STRIDE + swz_col
+                        vector.store(vec_safe, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = v_base + lds_row * V_STRIDE + swz_col
+                    vector.store(vec_safe, lds_kv, [lds_idx])
+
+        # ---- PROLOGUE: K, V into LDS (persistent for program lifetime) ----
+        coop_load_k(kv_start)
+        coop_load_v(kv_start)
+        gpu.barrier()
+
+        # R6b: K/V LDS A-frag load uses lane_mod_16 (16-row MFMA) and lane_div_16*8 for col.
+        # Swizzle mask key is the per-wave KV row (wave_n_offset + lane_mod_16), but since
+        # NUM_WAVES <= 2 and (wave_n_offset & K_SWZ_ROW_MASK) is 0 at D=64 (and bit-aligned
+        # at D=512 once wave_n_offset is folded in below), we recompute the mask per call.
+        def _k_idx_wave(ks):
+            kv_row = wave_n_offset + lane_mod_16
+            col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+            mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return kv_row * arith.index(K_STRIDE) + (col ^ mask)
+
+        def _v_idx_wave(ks):
+            kv_row = wave_n_offset + lane_mod_16
+            col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+            mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return arith.index(LDS_V_BASE) + kv_row * arith.index(V_STRIDE) + (col ^ mask)
+
+        # ---- SWA + m-loop bounds (constant across head loop) ----
+        SWA = arith.index(swa_window)
+        BN_idx = arith.index(BLOCK_N)
+        BM2_idx = arith.index(BLOCK_M2)
+        _one_idx = arith.index(1)
+        n_block_lo = kv_start
+        n_block_hi = kv_start + BN_idx
+        m_loop_start = (n_block_lo // BM2_idx) * BM2_idx
+        _m_end_raw = n_block_hi + SWA - _one_idx
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _m_end_raw, seq_len_q_v)
+        _m_end_cl = arith.select(_le_seq, _m_end_raw, seq_len_q_v)
+        m_loop_end = ((_m_end_cl + BM2_idx - _one_idx) // BM2_idx) * BM2_idx
+
+        # R6k-1: each wave only carries D_CHUNKS_LOCAL chunks of dv and dk.
+        outer_carry = [v4f32_zero for _ in range(D_CHUNKS_LOCAL)] + [
+            v4f32_zero for _ in range(D_CHUNKS_LOCAL)
+        ]
+
+        seq_len_k_i32 = arith.index_cast(T.i32, seq_len_k_v)
+        seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v)
+        w_i32 = arith.constant(swa_window, type=T.i32)
+        wave_n_off_i32 = arith.index_cast(T.i32, wave_n_offset)
+        kv_start_i32 = arith.index_cast(T.i32, kv_start)
+        lane_div_16_i32 = arith.index_cast(T.i32, lane_div_16)
+
+        NUM_HEADS_idx = arith.index(NUM_HEADS)
+        for qhid_constexpr_idx, h_carry, h_loop_results in scf.for_(
+            arith.index(0),
+            NUM_HEADS_idx,
+            arith.index(1),
+            iter_args=outer_carry,
+        ):
+            qhid = qhid_constexpr_idx
+
+            m_init_args = [h_carry[i] for i in range(2 * D_CHUNKS_LOCAL)]
+            for m_start, m_carry, m_loop_results in scf.for_(
+                m_loop_start,
+                m_loop_end,
+                BM2_idx,
+                iter_args=m_init_args,
+            ):
+                # R6k-1: each wave's dv/dk slice has D_CHUNKS_LOCAL chunks.
+                dv_accs = [m_carry[dc] for dc in range_constexpr(D_CHUNKS_LOCAL)]
+                dk_accs = [m_carry[D_CHUNKS_LOCAL + dc] for dc in range_constexpr(D_CHUNKS_LOCAL)]
+
+                q_row_abs = m_start + lane_mod_16
+                q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row_abs, seq_len_q_v)
+                arith.select(q_in_bounds, q_row_abs, arith.index(0))
+
+                # ---- Q + DO loads ----
+                if USE_LDS_FOR_Q_DO:
+                    for batch in range_constexpr(NUM_BATCHES_M):
+                        row_offset = batch * ROWS_PER_BATCH_LOAD
+                        row_idx = m_start + load_row_in_batch + row_offset
+                        in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_q_v)
+                        row_safe = arith.select(in_bounds, row_idx, arith.index(0))
+                        g_idx_do = global_idx_q(qhid, row_safe, load_col_base)
+                        g_idx_q = global_idx_q(qhid, row_safe, load_col_base)
+                        vec_do = load_global_f16xN(do_ptr, g_idx_do)
+                        vec_q = load_global_f16xN(q_ptr, g_idx_q)
+                        vec_do_safe = arith.select(in_bounds, vec_do, c_zero_vxf16)
+                        vec_q_safe = arith.select(in_bounds, vec_q, c_zero_vxf16)
+                        lds_row = load_row_in_batch + row_offset
+                        if M_NEEDS_GUARD:
+                            row_valid = arith.cmpi(
+                                arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_M2)
+                            )
+                            _if_qd = scf.IfOp(row_valid)
+                            with ir.InsertionPoint(_if_qd.then_block):
+                                lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base
+                                lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base
+                                vector.store(vec_do_safe, lds_do, [lds_idx_do])
+                                vector.store(vec_q_safe, lds_q, [lds_idx_q])
+                                scf.YieldOp([])
+                        else:
+                            lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base
+                            lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base
+                            vector.store(vec_do_safe, lds_do, [lds_idx_do])
+                            vector.store(vec_q_safe, lds_q, [lds_idx_q])
+                    gpu.barrier()
+
+                    # 6N-2 iter2: when D-split, q_b_packs / do_b_packs_gemm3
+                    # are loaded INSIDE the per-mt scf.IfOp gates (see
+                    # GEMM1 / GEMM3 below) so un-owned waves do not waste
+                    # LDS read instructions. Build trivial sentinels.
+                    if D_SPLIT_WAVES:
+                        q_b_packs = None  # gated path: loads done in GEMM1 if-block
+                        do_b_packs_gemm3 = None  # gated path: loads done in GEMM3 if-block
+                    else:
+                        q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                        for mt in range_constexpr(M_TILES):
+                            for ks in range_constexpr(K_STEPS_QK):
+                                m_row = arith.index(mt * 16) + lane_mod_16
+                                d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                                lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col
+                                pack = vector.load_op(mfma_pack_type, lds_q, [lds_idx])
+                                q_b_packs[mt][ks] = pack
+
+                        do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                        for mt in range_constexpr(M_TILES):
+                            for ks in range_constexpr(K_STEPS_QK):
+                                m_row = arith.index(mt * 16) + lane_mod_16
+                                d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                                lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col
+                                pack = vector.load_op(mfma_pack_type, lds_do, [lds_idx])
+                                do_b_packs_gemm3[mt][ks] = pack
+                else:
+                    # R6b fallback: HBM-direct B-frag per m-tile. row = m_start + mt*16 + lane_mod_16.
+                    q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                    do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)]
+                    for mt in range_constexpr(M_TILES):
+                        row_abs = m_start + arith.index(mt * 16) + lane_mod_16
+                        in_bnd = arith.cmpi(arith.CmpIPredicate.slt, row_abs, seq_len_q_v)
+                        row_safe = arith.select(in_bnd, row_abs, arith.index(0))
+                        for ks in range_constexpr(K_STEPS_QK):
+                            col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                            g_idx = global_idx_q(qhid, row_safe, col)
+                            q_raw = load_global_mfma_pack(q_ptr, g_idx)
+                            q_b_packs[mt][ks] = arith.select(in_bnd, q_raw, c_zero_mfma_pack)
+                            do_raw = load_global_mfma_pack(do_ptr, g_idx)
+                            do_b_packs_gemm3[mt][ks] = arith.select(in_bnd, do_raw, c_zero_mfma_pack)
+
+                # 6N-2: de-dup GEMM1 + softmax + pT_write per mt-tile.
+                # Each mt is owned by one wave (wave_id == mt % NUM_WAVES).
+                # The owning wave runs everything; the other yields zeros.
+                # delta_per_tile must be valid in BOTH waves (used after barrier
+                # for the dsT compute), so we read it unconditionally.
+                s_accs = [None for _ in range(M_TILES)]
+                pT_vals_per_tile = [None for _ in range(M_TILES)]
+                delta_per_tile = []
+                for mt in range_constexpr(M_TILES):
+                    # delta is needed by every wave (dsT compute happens after
+                    # GEMM3 dp_accs[mt] is restored from LDS-broadcast / read).
+                    # In de-dup mode dsT is owned by mt-owner only, so delta
+                    # also only needs to be live in the owning wave. But the
+                    # buffer_load is a cheap SMEM op; we leave it unconditional
+                    # for simplicity. Same for lse below.
+                    m_row_for_lse = m_start + arith.index(mt * 16) + lane_mod_16
+                    m_row_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, m_row_for_lse, seq_len_q_v)
+                    m_row_safe = arith.select(m_row_in_bounds, m_row_for_lse, arith.index(0))
+                    lse_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q_of(qhid) + m_row_safe)
+                    lse_v = buffer_ops.buffer_load(lse_rsrc, lse_off_i32, vec_width=1, dtype=T.f32)
+                    delta_v = buffer_ops.buffer_load(deltas_rsrc, lse_off_i32, vec_width=1, dtype=T.f32)
+                    delta_per_tile.append(delta_v)
+
+                    if D_SPLIT_WAVES:
+                        owns = owns_mt_preds[mt]
+                        # ---- GATED block: GEMM1[mt] + softmax + pT LDS write ----
+                        # Yields (s_acc, pT0, pT1, pT2, pT3).
+                        if_g1sm = scf.IfOp(
+                            owns, results_=[v4f32_type, T.f32, T.f32, T.f32, T.f32], has_else=True
+                        )
+                        with ir.InsertionPoint(if_g1sm.then_block):
+                            # GEMM1[mt]: accumulate over K_STEPS_QK.
+                            # 6N-2 iter2: load Q B-frag inside the gate when LDS Q/DO.
+                            acc = v4f32_zero
+                            for ks in range_constexpr(K_STEPS_QK):
+                                k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)])
+                                if USE_LDS_FOR_Q_DO:
+                                    _m_row_b = arith.index(mt * 16) + lane_mod_16
+                                    _d_col_b = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                                    _lds_idx_b = _m_row_b * arith.index(LDS_Q_STRIDE) + _d_col_b
+                                    q_pack = vector.load_op(mfma_pack_type, lds_q, [_lds_idx_b])
+                                else:
+                                    q_pack = q_b_packs[mt][ks]
+                                acc = mfma_acc(k_pack, q_pack, acc)
+
+                            m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse)
+                            m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32)
+                            pT_pieces = []
+                            for ii in range_constexpr(4):
+                                ii_i32 = arith.constant(ii, type=T.i32)
+                                kv_row_rel_i32 = arith.AddIOp(
+                                    arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result,
+                                    ii_i32,
+                                ).result
+                                kv_abs_i32 = arith.AddIOp(
+                                    arith.AddIOp(kv_start_i32, wave_n_off_i32).result, kv_row_rel_i32
+                                ).result
+                                kv_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_abs_i32, seq_len_k_i32)
+                                is_causal = arith.cmpi(
+                                    arith.CmpIPredicate.sgt, kv_abs_i32, m_row_abs_for_mask_i32
+                                )
+                                kv_plus_w = arith.AddIOp(kv_abs_i32, w_i32).result
+                                is_swa = arith.cmpi(
+                                    arith.CmpIPredicate.sle, kv_plus_w, m_row_abs_for_mask_i32
+                                )
+                                bad = arith.OrIOp(
+                                    arith.OrIOp(arith.OrIOp(is_causal, is_swa).result, kv_oob).result, m_oob
+                                ).result
+                                s_ii = vector.extract(acc, static_position=[ii], dynamic_position=[])
+                                scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result
+                                scaled_m = arith.select(bad, c_neg_inf, scaled)
+                                diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result
+                                p = math_dialect.exp(diff, fastmath=fm_fast)
+                                pT_pieces.append(p)
+                                # Write pT to LDS for this mt.
+                                kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                                kv_row = wave_n_offset + kv_row_rel
+                                pt_bf16 = arith.trunc_f(elem_type, p)
+                                lds_pt_idx = (
+                                    kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                                )
+                                v1_pt = vector.from_elements(v1_elem_type, [pt_bf16])
+                                vector.store(v1_pt, lds_pt, [lds_pt_idx])
+                            scf.YieldOp([acc] + pT_pieces)
+                        with ir.InsertionPoint(if_g1sm.else_block):
+                            _zero_f32 = arith.constant(0.0, type=T.f32)
+                            scf.YieldOp([v4f32_zero, _zero_f32, _zero_f32, _zero_f32, _zero_f32])
+                        s_accs[mt] = if_g1sm.results[0]
+                        pT_vals_per_tile[mt] = [
+                            if_g1sm.results[1],
+                            if_g1sm.results[2],
+                            if_g1sm.results[3],
+                            if_g1sm.results[4],
+                        ]
+                    else:
+                        # Legacy path: every wave does GEMM1[mt] + softmax + pT_write.
+                        acc = v4f32_zero
+                        for ks in range_constexpr(K_STEPS_QK):
+                            k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)])
+                            acc = mfma_acc(k_pack, q_b_packs[mt][ks], acc)
+                        s_accs[mt] = acc
+
+                        m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse)
+                        m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32)
+                        pT_tile = []
+                        for ii in range_constexpr(4):
+                            ii_i32 = arith.constant(ii, type=T.i32)
+                            kv_row_rel_i32 = arith.AddIOp(
+                                arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result, ii_i32
+                            ).result
+                            kv_abs_i32 = arith.AddIOp(
+                                arith.AddIOp(kv_start_i32, wave_n_off_i32).result, kv_row_rel_i32
+                            ).result
+                            kv_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_abs_i32, seq_len_k_i32)
+                            is_causal = arith.cmpi(
+                                arith.CmpIPredicate.sgt, kv_abs_i32, m_row_abs_for_mask_i32
+                            )
+                            kv_plus_w = arith.AddIOp(kv_abs_i32, w_i32).result
+                            is_swa = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w, m_row_abs_for_mask_i32)
+                            bad = arith.OrIOp(
+                                arith.OrIOp(arith.OrIOp(is_causal, is_swa).result, kv_oob).result, m_oob
+                            ).result
+                            s_ii = vector.extract(acc, static_position=[ii], dynamic_position=[])
+                            scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result
+                            scaled_m = arith.select(bad, c_neg_inf, scaled)
+                            diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result
+                            p = math_dialect.exp(diff, fastmath=fm_fast)
+                            pT_tile.append(p)
+                            kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                            kv_row = wave_n_offset + kv_row_rel
+                            pt_bf16 = arith.trunc_f(elem_type, p)
+                            lds_pt_idx = (
+                                kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                            )
+                            v1_pt = vector.from_elements(v1_elem_type, [pt_bf16])
+                            vector.store(v1_pt, lds_pt, [lds_pt_idx])
+                        pT_vals_per_tile[mt] = pT_tile
+
+                # R6k-1: D-split barrier. Both waves write the SAME pT
+                # values (identical s_accs derived from full-D GEMM1) to
+                # the SAME LDS rows (wave_n_offset=0 for both waves), but
+                # the writer-lane != reader-lane mapping crosses wave
+                # boundaries: lane L in wave 0 reads cells written by
+                # lanes from wave 1 (and vice versa). Without a barrier
+                # wave 1 could observe wave 0's stale prior-iteration pT.
+                if D_SPLIT_WAVES:
+                    gpu.barrier()
+
+                # ---- pT A-frag for GEMM2 (dV += pT @ DO):
+                #   A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7]
+                pt_a_packs = []
+                for m_step in range_constexpr(K_STEPS_PT):
+                    kv_row_a = wave_n_offset + lane_mod_16
+                    m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a
+                    pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx])
+                    pt_a_packs.append(pack)
+
+                # ---- GEMM2: dV += pT @ DO ----
+                # R6b GEMM2: B-frag DO[m_step*32+lane_div_16*8+k, d_col_global+lane_mod_16].
+                # R6k-1: dc_idx is wave-local; d_col_global = wave_d_col_offset + dc_idx*D_CHUNK.
+                if USE_LDS_FOR_Q_DO:
+                    # R6S: tr16 LDS reads via `base + constexpr byte imm`
+                    # so LLVM ISel folds the per-call delta into the LDS
+                    # instruction's `offset:` immediate. Replaces the
+                    # per-call IntToPtrOp pattern that was spilling 32
+                    # absolute addresses across AGPRs/scratch.
+                    def _ds_read_tr_do_imm(byte_imm):
+                        gep = _llvm.GEPOp(
+                            _llvm_lds_ptr_ty(),
+                            tr_lane_base_do_ptr,
+                            [],
+                            rawConstantIndices=[byte_imm],
+                            elem_type=T.i8,
+                            noWrapFlags=0,
+                        )
+                        return rocdl.ds_read_tr16_b64(v4f16_type, gep.result).result
+
+                    def read_do_b_pack(m_step_idx, dc_idx):
+                        # All offsets are Python ints (constexpr unrolled).
+                        base_elem = m_step_idx * K_STEP_QK * LDS_DO_STRIDE + dc_idx * D_CHUNK
+                        byte_imm_0 = base_elem * 2
+                        byte_imm_1 = byte_imm_0 + 4 * LDS_DO_STRIDE * 2
+                        v_lo = _ds_read_tr_do_imm(byte_imm_0)
+                        v_hi = _ds_read_tr_do_imm(byte_imm_1)
+                        v_full = vector.shuffle(v_lo, v_hi, [0, 1, 2, 3, 4, 5, 6, 7])
+                        return vector.bitcast(mfma_pack_type, v_full)
+
+                else:
+
+                    def read_do_b_pack(m_step_idx, dc_idx):
+                        d_col = wave_d_col_offset + arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                        m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                        vals = []
+                        for rk in range_constexpr(MFMA_LANE_K):
+                            m_row_rel = m_base + arith.index(rk)
+                            m_row_abs = m_start + m_row_rel
+                            in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v)
+                            m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0))
+                            g_idx = global_idx_q(qhid, m_row_safe2, d_col)
+                            v1 = _gep_load(do_ptr, g_idx, T.vec(1, elem_type))
+                            v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[])
+                            v_safe = arith.select(in_b, v_scalar, c_zero_elem)
+                            vals.append(v_safe)
+                        return vector.from_elements(mfma_pack_type, vals)
+
+                # R6k-1: dc loops wave-local (D_CHUNKS_LOCAL); reader adds
+                # wave_d_col_offset internally.
+                new_dv_accs = list(dv_accs)
+                for dc in range_constexpr(D_CHUNKS_LOCAL):
+                    for pks in range_constexpr(K_STEPS_PT):
+                        b_pack = read_do_b_pack(pks, dc)
+                        new_dv_accs[dc] = mfma_acc(pt_a_packs[pks], b_pack, new_dv_accs[dc])
+
+                # 6N-2: de-dup GEMM3 + dsT + dsT_write per mt-tile.
+                # The owning wave for mt does GEMM3 MFMA, dsT compute, and LDS write.
+                # The other wave skips entirely. GEMM4 readers see the full
+                # dsT via LDS broadcast after the trailing barrier.
+                if D_SPLIT_WAVES:
+                    for mt in range_constexpr(M_TILES):
+                        owns = owns_mt_preds[mt]
+                        if_g3 = scf.IfOp(owns, results_=[], has_else=False)
+                        with ir.InsertionPoint(if_g3.then_block):
+                            # GEMM3[mt]
+                            # 6N-2 iter2: load DO B-frag inside the gate when LDS Q/DO.
+                            acc = v4f32_zero
+                            for ks in range_constexpr(K_STEPS_QK):
+                                v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)])
+                                if USE_LDS_FOR_Q_DO:
+                                    _m_row_b = arith.index(mt * 16) + lane_mod_16
+                                    _d_col_b = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                                    _lds_idx_b = _m_row_b * arith.index(LDS_DO_STRIDE) + _d_col_b
+                                    do_pack = vector.load_op(mfma_pack_type, lds_do, [_lds_idx_b])
+                                else:
+                                    do_pack = do_b_packs_gemm3[mt][ks]
+                                acc = mfma_acc(v_pack, do_pack, acc)
+                            # dsT + LDS write per ii.
+                            for ii in range_constexpr(4):
+                                dp_ii = vector.extract(acc, static_position=[ii], dynamic_position=[])
+                                diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result
+                                ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result
+                                kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                                kv_row = wave_n_offset + kv_row_rel
+                                ds_bf16 = arith.trunc_f(elem_type, ds_ii)
+                                lds_pt_idx = (
+                                    kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                                )
+                                v1_ds = vector.from_elements(v1_elem_type, [ds_bf16])
+                                vector.store(v1_ds, lds_pt, [lds_pt_idx])
+                            scf.YieldOp([])
+                else:
+                    # Legacy path
+                    dp_accs = [v4f32_zero for _ in range(M_TILES)]
+                    for ks in range_constexpr(K_STEPS_QK):
+                        v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)])
+                        for mt in range_constexpr(M_TILES):
+                            dp_accs[mt] = mfma_acc(v_pack, do_b_packs_gemm3[mt][ks], dp_accs[mt])
+                    for mt in range_constexpr(M_TILES):
+                        for ii in range_constexpr(4):
+                            dp_ii = vector.extract(dp_accs[mt], static_position=[ii], dynamic_position=[])
+                            diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result
+                            ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result
+                            kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                            kv_row = wave_n_offset + kv_row_rel
+                            ds_bf16 = arith.trunc_f(elem_type, ds_ii)
+                            lds_pt_idx = (
+                                kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16
+                            )
+                            v1_ds = vector.from_elements(v1_elem_type, [ds_bf16])
+                            vector.store(v1_ds, lds_pt, [lds_pt_idx])
+
+                # R6k-1: D-split barrier (see GEMM2 pT barrier comment).
+                if D_SPLIT_WAVES:
+                    gpu.barrier()
+
+                # ---- dsT A-frag for GEMM4 (dK += dsT @ Q):
+                #   A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7]
+                ds_a_packs = []
+                for m_step in range_constexpr(K_STEPS_PT):
+                    kv_row_a = wave_n_offset + lane_mod_16
+                    m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                    lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a
+                    pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx])
+                    ds_a_packs.append(pack)
+
+                # ---- GEMM4: dK += dsT @ Q  (MFMA 16x16x32) ----
+                # R6k-1: Q B-frag d_col = wave_d_col_offset + dc_idx*D_CHUNK + lane_mod_16.
+                if USE_LDS_FOR_Q_DO:
+                    # R6S: tr16 LDS reads via `base + constexpr byte imm`
+                    # (same recipe as GEMM2). Uses tr_lane_base_q_ptr.
+                    def _ds_read_tr_q_imm(byte_imm):
+                        gep = _llvm.GEPOp(
+                            _llvm_lds_ptr_ty(),
+                            tr_lane_base_q_ptr,
+                            [],
+                            rawConstantIndices=[byte_imm],
+                            elem_type=T.i8,
+                            noWrapFlags=0,
+                        )
+                        return rocdl.ds_read_tr16_b64(v4f16_type, gep.result).result
+
+                    def read_q_b_pack(m_step_idx, dc_idx):
+                        base_elem = m_step_idx * K_STEP_QK * LDS_Q_STRIDE + dc_idx * D_CHUNK
+                        byte_imm_0 = base_elem * 2
+                        byte_imm_1 = byte_imm_0 + 4 * LDS_Q_STRIDE * 2
+                        v_lo = _ds_read_tr_q_imm(byte_imm_0)
+                        v_hi = _ds_read_tr_q_imm(byte_imm_1)
+                        v_full = vector.shuffle(v_lo, v_hi, [0, 1, 2, 3, 4, 5, 6, 7])
+                        return vector.bitcast(mfma_pack_type, v_full)
+
+                else:
+
+                    def read_q_b_pack(m_step_idx, dc_idx):
+                        d_col = wave_d_col_offset + arith.index(dc_idx * D_CHUNK) + lane_mod_16
+                        m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K
+                        vals = []
+                        for rk in range_constexpr(MFMA_LANE_K):
+                            m_row_rel = m_base + arith.index(rk)
+                            m_row_abs = m_start + m_row_rel
+                            in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v)
+                            m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0))
+                            g_idx = global_idx_q(qhid, m_row_safe2, d_col)
+                            v1 = _gep_load(q_ptr, g_idx, T.vec(1, elem_type))
+                            v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[])
+                            v_safe = arith.select(in_b, v_scalar, c_zero_elem)
+                            vals.append(v_safe)
+                        return vector.from_elements(mfma_pack_type, vals)
+
+                # R6k-1: dc loops wave-local; reader adds wave_d_col_offset.
+                new_dk_accs = list(dk_accs)
+                for dc in range_constexpr(D_CHUNKS_LOCAL):
+                    for pks in range_constexpr(K_STEPS_PT):
+                        q_b_pack = read_q_b_pack(pks, dc)
+                        new_dk_accs[dc] = mfma_acc(ds_a_packs[pks], q_b_pack, new_dk_accs[dc])
+
+                gpu.barrier()
+
+                yield list(new_dv_accs) + list(new_dk_accs)
+
+            yield list(m_loop_results)
+        outer_carry = list(h_loop_results)
+
+        # ---- Final: dK *= sm_scale, store ----
+        # R6k-1: D-split mode -> each wave only owns D_CHUNKS_LOCAL chunks
+        # and writes its own slice of dK/dV (no inter-wave reduction needed:
+        # the D-axis is disjoint per wave, and the N-axis is shared so the
+        # accumulators ARE complete for each wave's D-slice).
+        dv_finals = [outer_carry[dc] for dc in range(D_CHUNKS_LOCAL)]
+        dk_finals = [outer_carry[D_CHUNKS_LOCAL + dc] for dc in range(D_CHUNKS_LOCAL)]
+
+        for dc in range_constexpr(D_CHUNKS_LOCAL):
+            for ii in range_constexpr(4):
+                kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii)
+                kv_row_abs = kv_start + wave_n_offset + kv_row_rel
+                d_col_abs = wave_d_col_offset + arith.index(dc * D_CHUNK) + lane_mod_16
+                kv_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, kv_row_abs, seq_len_k_v)
+                _if_kv = scf.IfOp(kv_in_bounds)
+                with ir.InsertionPoint(_if_kv.then_block):
+                    dv_val = vector.extract(dv_finals[dc], static_position=[ii], dynamic_position=[])
+                    dk_val = vector.extract(dk_finals[dc], static_position=[ii], dynamic_position=[])
+                    dk_scaled = arith.MulFOp(dk_val, c_sm_scale, fastmath=fm_fast).result
+                    g_idx = global_idx_kv(kv_row_abs, d_col_abs)
+                    _gep_store_f32(dv_val, dv_ptr, g_idx)
+                    _gep_store_f32(dk_scaled, dk_ptr, g_idx)
+                    scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_swa_bwd_dkv(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DK: fx.Tensor,
+        DV: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sk_idx = arith.index_cast(T.index, seq_len_k)
+        num_n_blocks = (sk_idx + BLOCK_N - 1) // BLOCK_N
+        grid_x = bs_idx * num_n_blocks
+
+        launcher = v4_swa_bwd_dkv_kernel(
+            Q,
+            K,
+            V,
+            DOS,
+            LSE,
+            DELTAS,
+            DK,
+            DV,
+            seq_len_q,
+            seq_len_k,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+        if flat_work_group_size is not None:
+            _fwgs = int(flat_work_group_size)
+            if _fwgs >= 1:
+                flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, 1, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    _fmha_compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {
+            "enable-post-misched": True,  # R6L iter5
+            "lsr-drop-solution": True,
+        },
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return launch_v4_swa_bwd_dkv(*args, **kwargs)
+
+    def _compile(Q, K, V, DOS, LSE, DELTAS, DK, DV, batch_size, seq_len_q, seq_len_k, stream=None):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return flyc.compile(
+                launch_v4_swa_bwd_dkv,
+                Q,
+                K,
+                V,
+                DOS,
+                LSE,
+                DELTAS,
+                DK,
+                DV,
+                batch_size,
+                seq_len_q,
+                seq_len_k,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+build_v4_swa_bwd_dkv_module_primary = build_v4_swa_bwd_dkv_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py
new file mode 100644
index 000000000..470223ae0
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py
@@ -0,0 +1,1284 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_swa_bwd_dq: V4 SWA-causal attention backward dQ kernel for FlyDSL.
+
+Forked from kernels/sla_bwd_dq.py (FlyDSL SLA backward dQ). Differences:
+  - Outer KV loop iterates a contiguous block range bounded by the SWA
+    window for the current Q tile (no LUT).
+  - Per-element SWA + causal mask + boundary mask (mirrors Triton ref).
+  - LSE is RAW-domain (Triton: lse = m + ln(l), in domain qk*sm_scale),
+    so p = exp(qk*sm_scale - lse) (NOT exp2).
+  - sm_scale is applied TWICE: once inside the loop on qk (matches
+    Triton: qk = qk * sm_scale), and once after the n-loop on dq
+    (matches Triton: dq = dq * sm_scale). The two multiplies are NOT
+    redundant: ds = p*(dp - dvec) carries the in-loop scale via P;
+    the post-loop multiply is the outer chain-rule scale on dq.
+  - MQA: K/V are [B, 1, Sk, D] - stride_kh = 0 - drop head_idx from
+    KV indexing when mqa_kv=True.
+  - Sink: optional SINK[HQ] fp32; dsink = -sum(p_sink * dvec) via
+    atomic_fadd one scalar per row-owner-lane into DSINK[qhid]. Sink
+    does NOT change dq (fwd already folded p_sink into lse).
+  - dq written as fp32 (matches the launcher's fp32 dq_fp32 buffer).
+
+Layout: BHLD. Q/DOUT/DQ flat from (B, HQ, Sq, D). K/V flat from
+        (B, HK, Sk, D); MQA means HK=1, no head stride applied to KV.
+LSE/DELTAS: (B, HQ, Sq) flat, fp32, raw-domain.
+SINK:   (HQ,) fp32 - dummy buffer when has_sink=False.
+DSINK:  (HQ,) fp32 - atomic accumulator - caller must zero-init.
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import memref as _memref
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+# ---- Module-level constants ----
+
+KERNEL_NAME = "v4_swa_bwd_dq_kernel"
+
+_LOG2E = math.log2(math.e)  # 1.4426950408889634
+
+_LLVM_GEP_DYNAMIC = -2147483648  # LLVM kDynamicIndex sentinel (0x80000000 as signed i32)
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+_VMCNT_LO_MASK = 0xF
+_LGKMCNT_EXPCNT_BASE = 0x3F70
+_VMCNT_HI_SHIFT = 14
+_VMCNT_HI_MASK = 0x3
+
+
+def _waitcnt_vm_n(n):
+    """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7)."""
+    val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT)
+    rocdl.s_waitcnt(val)
+
+
+def build_v4_swa_bwd_dq_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    flat_work_group_size=None,
+    block_m=None,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    layout_bhld=True,
+    mqa_kv=True,
+    has_sink=True,
+):
+    """Build the V4 SWA backward dQ launcher.
+
+    Args:
+        num_heads: HQ (number of Q heads).
+        head_dim: D (head dimension), must be % 32 == 0 and >= 64.
+        swa_window: int > 0; SWA window length.
+        sm_scale: defaults to 1/sqrt(head_dim).
+        mqa_kv: if True, K/V indexing drops head_idx (HK=1, stride_h=0).
+        has_sink: if True, SINK/DSINK tensors are used. If False, both
+            are dummy buffers (the kernel still takes them but doesn't
+            touch them).
+        layout_bhld: BHLD layout (True) or BLHD (False); V4 uses BHLD.
+
+    Returns: launch(Q, K, V, DOUT, LSE, DELTAS, DQ_FP32, DSINK, SINK,
+                    batch_size, seq_len_q, seq_len_k, stream=None).
+    """
+    gpu_arch = get_hip_arch()
+
+    BLOCK_N = 64
+    K_SUB_N = 32
+    WARP_SIZE = 64
+
+    if block_m is not None:
+        BLOCK_M = block_m
+    else:
+        BLOCK_M = 128
+
+    if flat_work_group_size is None:
+        if BLOCK_M <= 128:
+            flat_work_group_size = 256
+        else:
+            flat_work_group_size = 512
+    NUM_WAVES = flat_work_group_size // WARP_SIZE
+    BLOCK_SIZE = flat_work_group_size
+    ROWS_PER_WAVE = BLOCK_M // NUM_WAVES
+    # V4 SWA: dense path. One outer iter = one BLOCK_N block.
+    BLOCK_N_OUT = BLOCK_N
+    ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1"
+    _has_lds_load_b128 = not gpu_arch.startswith("gfx942")
+    ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1")
+    ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1"
+    REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower()
+    if REDUCE_MODE not in ("xor", "ds_bpermute"):
+        REDUCE_MODE = "xor"
+    NUM_PREFETCH_K = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1)
+    NUM_PREFETCH_V = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1)
+    (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,)
+
+    USE_HW_TR = gpu_arch.startswith("gfx950")
+    USE_K16 = gpu_arch.startswith("gfx950")
+
+    # Auto-fallback: gfx950 LDS limit is 160 KB. K+V double-buffering at D=512
+    # easily exceeds that. If predicted LDS exceeds budget, force DMA off
+    # (NUM_PREFETCH = 1) and warn. Keeps the kernel correct regardless of
+    # the env knob.
+    _LDS_LIMIT_BYTES = 160 * 1024
+
+    def _predicted_lds_bytes(nk, nv, dma):
+        # USE_HW_TR & DMA  -> V_STRIDE = head_dim
+        # USE_HW_TR & !DMA -> V_STRIDE = head_dim + 4
+        # !USE_HW_TR       -> V stored transposed: LDS V = head_dim * (BLOCK_N + 2)
+        if USE_HW_TR:
+            v_str = head_dim if dma else head_dim + 4
+            return (nk * BLOCK_N * head_dim + nv * BLOCK_N * v_str) * 2
+        vt = BLOCK_N + 2
+        return (nk * BLOCK_N * head_dim + nv * head_dim * vt) * 2
+
+    _pred = _predicted_lds_bytes(NUM_PREFETCH_K, NUM_PREFETCH_V, ENABLE_DMA)
+    if _pred > _LDS_LIMIT_BYTES and ENABLE_DMA:
+        # Try DMA off (single-buffered).
+        ENABLE_DMA = False
+        NUM_PREFETCH_K = 1
+        NUM_PREFETCH_V = 1
+        _pred2 = _predicted_lds_bytes(1, 1, False)
+        if _pred2 > _LDS_LIMIT_BYTES:
+            raise RuntimeError(
+                f"v4_swa_bwd_dq: predicted LDS {_pred2}B > limit {_LDS_LIMIT_BYTES}B "
+                f"even single-buffered at D={head_dim}, BLOCK_N={BLOCK_N}"
+            )
+        import sys as _sys
+
+        print(
+            f"[v4_swa_bwd_dq] LDS overflow at D={head_dim}, BLOCK_N={BLOCK_N}: "
+            f"auto-disabled DMA (predicted {_pred} -> {_pred2} bytes)",
+            file=_sys.stderr,
+            flush=True,
+        )
+    K_STEP_QK = 16 if USE_K16 else 8
+    K_STEPS_QK = head_dim // K_STEP_QK
+    D_CHUNK = 32
+    D_CHUNKS = head_dim // D_CHUNK
+    PV_K_STEP = 16 if USE_K16 else 8
+    PV_K_STEPS = K_SUB_N // PV_K_STEP  # 2 steps per sub-tile (K=16) or 4 (K=8)
+
+    assert BLOCK_M % NUM_WAVES == 0
+    assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32"
+    assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64"
+    assert flat_work_group_size in (
+        128,
+        256,
+        512,
+    ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}"
+    assert dtype_str in ("f16", "bf16"), "v4_swa_bwd_dq only supports f16 and bf16"
+    assert BLOCK_N % 32 == 0
+    assert BLOCK_N_OUT == BLOCK_N
+    assert isinstance(swa_window, int) and swa_window > 0, f"swa_window must be int > 0, got {swa_window!r}"
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(head_dim)
+
+    NUM_HEADS = num_heads
+    HEAD_DIM = head_dim
+    STRIDE_TOKEN = NUM_HEADS * HEAD_DIM
+
+    K_STRIDE = HEAD_DIM
+    # XOR swizzle mask must fit within the row stride. The swizzle is
+    # applied at 16-element granularity, so the maximum mask is
+    # `(K_STRIDE // 16 - 1) << 4`. For D=128 that's 7 (=0x7), giving max
+    # mask 112 < 128. For D=64 it's 3 (=0x3), giving max mask 48 < 64.
+    # SLA test only covers D=128 (mask=7); using a hardcoded `& 7` for
+    # D=64 wraps writes into adjacent rows -> silent corruption.
+    K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1
+    assert K_SWZ_ROW_MASK >= 0
+    assert (
+        K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)
+    ) == 0, f"K_SWZ_ROW_MASK must be 2^n-1, got {K_SWZ_ROW_MASK} (K_STRIDE={K_STRIDE})"
+    if USE_HW_TR:
+        V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4
+    else:
+        VT_STRIDE = BLOCK_N + 2
+        V_STRIDE = VT_STRIDE
+    # V swizzle: similarly bounded by V_STRIDE / 16.
+    V_SWZ_ROW_MASK = min(3, (V_STRIDE // 16) - 1)
+    assert V_SWZ_ROW_MASK >= 0
+
+    VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+
+    if ROWS_PER_BATCH_LOAD >= BLOCK_N:
+        NUM_BATCHES_KV = 1
+        KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N
+    else:
+        assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+        NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+        KV_NEEDS_GUARD = False
+
+    LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE
+    if USE_HW_TR:
+        LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE
+    else:
+        LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE
+    LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE
+    LDS_V_BASE = LDS_K_TOTAL_SIZE
+    LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE
+    LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_swa_bwd_dq_smem_M{BLOCK_M}_W{swa_window}_S{int(has_sink)}_MQ{int(mqa_kv)}",
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_swa_bwd_dq_kernel(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,  # grad of output (input)
+        LSE: fx.Tensor,  # log-sum-exp from fwd (input, f32, RAW domain)
+        DELTAS: fx.Tensor,  # (o * do).sum(-1) preprocess (input, f32)
+        DQ: fx.Tensor,  # grad wrt Q (output, FP32)
+        DSINK: fx.Tensor,  # grad wrt sink (output, FP32, [HQ])
+        SINK: fx.Tensor,  # sink param (input, FP32, [HQ]); dummy if !has_sink
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K)
+        v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS)
+        dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ)
+        # LSE / DELTAS: f32 scalar READS via buffer_load.
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+        sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True)
+        # DSINK: f32 atomic_fadd via buffer_atomic_fadd.
+        dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True)
+
+        # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation).
+        fm_fast = arith.FastMathFlags.fast
+        v4f16_type = T.vec(4, elem_type)
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        v16f32_type = T.vec(16, compute_type)
+        mfma_pack_type = v8f16_type if USE_K16 else v4f16_type
+        MFMA_LANE_K = 8 if USE_K16 else 4
+        _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)
+
+        def _mfma(ods_fn, a, b, c):
+            return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result
+
+        def mfma_acc(a, b, c):
+            if dtype_str == "bf16":
+                if USE_K16:
+                    return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c)
+                a = vector.bitcast(T.i16x4, a)
+                b = vector.bitcast(T.i16x4, b)
+                return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c)
+            if USE_K16:
+                return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c)
+            return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c)
+
+        seq_len_q_v = arith.index_cast(T.index, seq_len_q)
+        seq_len_k_v = arith.index_cast(T.index, seq_len_k)
+
+        # ---- LDS view ----
+        base_ptr = allocator.get_base()
+        lds_kv = SmemPtr(
+            base_ptr,
+            lds_kv_offset,
+            elem_type,
+            shape=(LDS_KV_TOTAL_SIZE,),
+        ).get()
+
+        # ---- Thread / block indices ----
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        lane_mod_32 = lane % 32
+        lane_div_32 = lane // 32  # 0/1
+
+        # ds_read_b64_tr_b16 lane decomposition
+        tr_k_group = (lane % 16) // 4
+        tr_col_sub = lane % 4
+        tr_col_half = (lane % 32) // 16
+
+        def ds_read_tr_v4f16(lds_elem_idx):
+            byte_offset = lds_elem_idx * 2 + lds_kv_offset
+            byte_i64 = arith.index_cast(T.i64, byte_offset)
+            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result
+
+        wave_q_offset = wave_id * ROWS_PER_WAVE
+
+        # ---- Decompose block_id: (batch, q_tile, head) ----
+        head_idx = block_id % NUM_HEADS
+        batch_q_tile_id = block_id // NUM_HEADS
+        num_q_tiles = (seq_len_q_v + BLOCK_M - 1) // BLOCK_M
+        q_tile_idx = batch_q_tile_id % num_q_tiles
+        batch_idx = batch_q_tile_id // num_q_tiles
+        q_start = q_tile_idx * BLOCK_M
+
+        # ---- V4 SWA: per-tile contiguous K-block range (wave-uniform) ----
+        SWA = arith.index(swa_window)
+        BN = arith.index(BLOCK_N)
+        BM = arith.index(BLOCK_M)
+        _zero_idx = arith.index(0)
+        _one_idx = arith.index(1)
+        # n_block_start = max(0, q_start - W + 1) // BLOCK_N
+        # n_block_end   = ceil(min(q_start + BLOCK_M, seq_len_k), BLOCK_N)
+        _q_plus_one = q_start + _one_idx
+        _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _q_plus_one, SWA)
+        _n_start_row = arith.select(_ge_w, _q_plus_one - SWA, _zero_idx)
+        n_block_start = _n_start_row // BN
+        _n_end_row_uncl = q_start + BM
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _n_end_row_uncl, seq_len_k_v)
+        n_end_row_cl = arith.select(_le_seq, _n_end_row_uncl, seq_len_k_v)
+        n_block_end = (n_end_row_cl + BN - _one_idx) // BN
+
+        # ---- Cooperative load decomposition ----
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        # ---- Helper: global flat index ----
+        if layout_bhld:
+            bh_base_tokens_q = (batch_idx * NUM_HEADS + head_idx) * seq_len_q_v
+            if mqa_kv:
+                bh_base_tokens_kv = batch_idx * seq_len_k_v
+            else:
+                bh_base_tokens_kv = (batch_idx * NUM_HEADS + head_idx) * seq_len_k_v
+
+            def global_idx_q(token_idx, col):
+                return (bh_base_tokens_q + token_idx) * arith.index(HEAD_DIM) + col
+
+            def global_idx_kv(token_idx, col):
+                return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col
+
+        else:
+            # BLHD path: kept for symmetry but V4 uses BHLD.
+            def global_idx_q(token_idx, col):
+                token = batch_idx * seq_len_q_v + token_idx
+                return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col
+
+            if mqa_kv:
+
+                def global_idx_kv(token_idx, col):
+                    token = batch_idx * seq_len_k_v + token_idx
+                    return token * HEAD_DIM + col
+
+            else:
+
+                def global_idx_kv(token_idx, col):
+                    token = batch_idx * seq_len_k_v + token_idx
+                    return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col
+
+        def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=et,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store_f32(val, base_ptr_, elem_idx):
+            """Store a single f32 value via GEP into an fp32 buffer."""
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr_,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=T.f32,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(base_ptr_, base_idx):
+            return _gep_load(base_ptr_, base_idx, vxf16_type)
+
+        def bf16_trunc_pack_v4(f32_vals):
+            _v2i32 = T.vec(2, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32)
+            b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32)
+            p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result
+            a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32)
+            b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32)
+            p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result
+            return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1]))
+
+        def bf16_trunc_pack_v8(f32_vals):
+            _v4i32 = T.vec(4, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            pairs = []
+            for j in range_constexpr(4):
+                a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32)
+                b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32)
+                p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result
+                pairs.append(p)
+            return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs))
+
+        def k_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(buf_id * LDS_K_TILE_SIZE)
+            return buf_id * arith.index(LDS_K_TILE_SIZE)
+
+        def v_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE)
+            return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE)
+
+        # ---- K XOR swizzle: col ^ ((row & K_SWZ_ROW_MASK) << 4) at 16-element granularity ----
+        # K_SWZ_ROW_MASK derived from K_STRIDE to stay within the row.
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        # ---- Cooperative K load (row-major, XOR-swizzled) ----
+        def coop_load_k(tile_start, buf_id=0):
+            k_base = k_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_k = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_k.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                        vec = load_global_f16xN(k_ptr, g_idx)
+                        vector.store(vec, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    vec = load_global_f16xN(k_ptr, g_idx)
+                    vector.store(vec, lds_kv, [lds_idx])
+
+        # ---- Cooperative V-into-K-LDS-slot load (K-style XOR swizzle) ----
+        def coop_load_v_as_k(tile_start, buf_id=0):
+            k_base = k_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                        vec = load_global_f16xN(v_ptr, g_idx)
+                        vector.store(vec, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    vec = load_global_f16xN(v_ptr, g_idx)
+                    vector.store(vec, lds_kv, [lds_idx])
+
+        # ---- Cooperative V load (V LDS layout) ----
+        def _v_store_row_major(v_base, lds_row, vec):
+            lds_idx = v_base + lds_row * V_STRIDE + load_col_base
+            vector.store(vec, lds_kv, [lds_idx])
+
+        _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None
+
+        def _v_store_transposed(v_base, lds_row, vec):
+            for _e in range_constexpr(VEC_WIDTH):
+                elem = vector.extract(vec, static_position=[_e], dynamic_position=[])
+                vt_d = load_col_base + _e
+                vt_idx = v_base + vt_d * VT_STRIDE + lds_row
+                v1 = vector.from_elements(_v1_type, [elem])
+                vector.store(v1, lds_kv, [vt_idx])
+
+        _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed
+
+        def coop_load_v(tile_start, buf_id=0):
+            v_base = v_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        vec = load_global_f16xN(v_ptr, g_idx)
+                        _v_store_to_lds(v_base, lds_row, vec)
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    vec = load_global_f16xN(v_ptr, g_idx)
+                    _v_store_to_lds(v_base, lds_row, vec)
+
+        # ---- Cooperative K-into-V-LDS-slot load ----
+        def coop_load_k_as_v(tile_start, buf_id=0):
+            v_base = v_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        vec = load_global_f16xN(k_ptr, g_idx)
+                        _v_store_to_lds(v_base, lds_row, vec)
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    vec = load_global_f16xN(k_ptr, g_idx)
+                    _v_store_to_lds(v_base, lds_row, vec)
+
+        # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ----
+        if ENABLE_DMA:
+            from flydsl._mlir.dialects import llvm
+
+            k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True)
+            _lds_ptr_ty = _llvm_lds_ptr_ty()
+            DMA_BYTES = 16
+            DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES
+            K_TILE_BYTES = BLOCK_N * K_STRIDE * 2
+            NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+            lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv)
+            _dma_size = arith.constant(DMA_BYTES, type=T.i32)
+            _dma_soff = arith.constant(0, type=T.i32)
+            _dma_off = arith.constant(0, type=T.i32)
+            _dma_aux = arith.constant(1, type=T.i32)
+
+            def _kv_global_byte(tile_start, row_in_tile, col_byte):
+                if layout_bhld:
+                    row_within = tile_start + row_in_tile
+                    return (bh_base_tokens_kv + row_within) * arith.index(HEAD_DIM * 2) + col_byte
+                else:
+                    global_row = batch_idx * seq_len_k_v + tile_start + row_in_tile
+                    if mqa_kv:
+                        return global_row * arith.index(HEAD_DIM * 2) + col_byte
+                    else:
+                        return (
+                            global_row * arith.index(STRIDE_TOKEN * 2)
+                            + head_idx * arith.index(HEAD_DIM * 2)
+                            + col_byte
+                        )
+
+            def coop_dma_k(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2)
+                else:
+                    k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2)
+                for d in range_constexpr(NUM_DMA_K):
+                    lds_addr = (
+                        k_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH)
+                    swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        k_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        def _v_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+            return col_idx ^ mask
+
+        if ENABLE_DMA:
+            v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True)
+            V_TILE_BYTES = BLOCK_N * V_STRIDE * 2
+            NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+
+            def coop_dma_v(tile_start, buf_id=0):
+                v_lds_byte_base = lds_kv_base_idx + arith.index((LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2)
+                for d in range_constexpr(NUM_DMA_V):
+                    lds_addr = (
+                        v_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V)
+                    swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        v_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+            # ---- Bwd dQ DMA variants (cross-pointer, matching swizzle) ----
+            def coop_dma_v_as_k(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2)
+                else:
+                    k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2)
+                for d in range_constexpr(NUM_DMA_K):
+                    lds_addr = (
+                        k_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH)
+                    swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        v_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+            def coop_dma_k_as_v(tile_start, buf_id=0):
+                if isinstance(buf_id, int):
+                    v_lds_byte_base = lds_kv_base_idx + arith.index(
+                        (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2
+                    )
+                else:
+                    v_lds_byte_base = (
+                        lds_kv_base_idx
+                        + arith.index(LDS_V_BASE * 2)
+                        + buf_id * arith.index(LDS_V_TILE_SIZE * 2)
+                    )
+                for d in range_constexpr(NUM_DMA_V):
+                    lds_addr = (
+                        v_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V)
+                    swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte)
+                    voffset = arith.index_cast(T.i32, global_byte)
+                    rocdl.raw_ptr_buffer_load_lds(
+                        k_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        # ---- Preload Q^T B-operand and DO^T B-operand packs ----
+        q_row = q_start + wave_q_offset + lane_mod_32
+        q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_q_v)
+        q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0))
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        q_b_packs = []
+        do_b_packs = []
+        for ks in range_constexpr(K_STEPS_QK):
+            col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+            g_idx = global_idx_q(q_row_safe, col)
+            q_raw = load_global_mfma_pack(q_ptr, g_idx)
+            q_b_packs.append(arith.select(q_in_bounds, q_raw, c_zero_mfma_pack))
+            do_raw = load_global_mfma_pack(do_ptr, g_idx)
+            do_b_packs.append(arith.select(q_in_bounds, do_raw, c_zero_mfma_pack))
+
+        # ---- Constants ----
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)  # finite NEG_INF (matches Triton)
+        c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type)
+        # V4 LSE is RAW-domain (qk*sm_scale + ln(l)). So use sm_scale, not sm_scale*log2e.
+        c_sm_scale = arith.constant(sm_scale, type=compute_type)
+
+        # ---- Per-q-row scalars (LSE, delta) ----
+        # bh_base_tokens_q is the Q LSE/DELTAS base too (LSE shape [B,HQ,Sq]).
+        lse_delta_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q + q_row_safe)
+        lse_val = buffer_ops.buffer_load(
+            lse_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=T.f32,
+        )
+        delta_val = buffer_ops.buffer_load(
+            deltas_rsrc,
+            lse_delta_off_i32,
+            vec_width=1,
+            dtype=T.f32,
+        )
+
+        # ---- SINK contribution to DSINK ----
+        # Per Triton ref:
+        #   sink_h = SINK[qhid]   (uniform across lanes of this program)
+        #   p_sink = exp(sink_h - lse)
+        #   dsink_contrib = sum_m -p_sink_masked * dvec_masked  (over BLOCK_M rows)
+        #   atomic_fadd(DSINK + qhid, dsink_contrib)
+        # We do the atomic per row-owner lane (lane_div_32==0 lane) to avoid
+        # the cross-wave reduction. That gives BLOCK_M atomics per program -
+        # slow but bulletproof. Each lane contributes -p_sink * dvec for its
+        # owned row only when q_in_bounds.
+        if has_sink:
+            head_idx_i32 = arith.index_cast(T.i32, head_idx)
+            sink_h_scalar = buffer_ops.buffer_load(
+                sink_rsrc,
+                head_idx_i32,
+                vec_width=1,
+                dtype=T.f32,
+            )
+            sink_h_uniform = rocdl.readfirstlane(T.f32, sink_h_scalar)
+            sub_val = arith.SubFOp(
+                sink_h_uniform,
+                lse_val,
+                fastmath=fm_fast,
+            ).result
+            p_sink = math_dialect.exp(sub_val, fastmath=fm_fast)
+            neg_p_sink = arith.SubFOp(
+                c_zero_f,
+                p_sink,
+                fastmath=fm_fast,
+            ).result
+            contrib = arith.MulFOp(
+                neg_p_sink,
+                delta_val,
+                fastmath=fm_fast,
+            ).result
+            # Gate: only the row-owner lane (lane_div_32==0) and only when
+            # q_row < seq_len_q, atomically adds.
+            is_row_owner = arith.cmpi(
+                arith.CmpIPredicate.eq,
+                lane_div_32,
+                arith.index(0),
+            )
+            do_sink_atomic = arith.AndIOp(is_row_owner, q_in_bounds).result
+            _if_sink = scf.IfOp(do_sink_atomic, [], has_else=False)
+            with ir.InsertionPoint(_if_sink.then_block):
+                # DSINK byte offset = head_idx * 4 (fp32).
+                _dsink_byte_off = arith.MulIOp(
+                    head_idx_i32,
+                    arith.constant(4, type=T.i32),
+                ).result
+                _zero_i32_atom = arith.constant(0, type=T.i32)
+                rocdl.raw_ptr_buffer_atomic_fadd(
+                    contrib,
+                    dsink_rsrc,
+                    _dsink_byte_off,
+                    _zero_i32_atom,
+                    _zero_i32_atom,
+                )
+                scf.YieldOp([])
+
+        # ---- MILESTONE: dense SWA bwd dQ with double-buffered LDS ----
+        _use_dbuf = ENABLE_DMA
+
+        init_args = []
+        for _ in range_constexpr(D_CHUNKS):
+            init_args.append(c_zero_v16f32)
+        if _use_dbuf:
+            init_args.append(arith.index(0))  # cur_buf_id
+
+            # PROLOGUE: prefetch iter 0's K into BOTH slots of buf 0.
+            _init_kv_start = n_block_start * BN
+            coop_dma_k(_init_kv_start, buf_id=0)
+            coop_dma_k_as_v(_init_kv_start, buf_id=0)
+
+        for block_idx, inner_iter_args, loop_results in scf.for_(
+            n_block_start,
+            n_block_end,
+            arith.index(1),
+            iter_args=init_args,
+        ):
+            dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)]
+            if _use_dbuf:
+                cur_buf = inner_iter_args[D_CHUNKS]
+                next_buf = arith.index(1) - cur_buf
+
+            # V4 SWA: block_idx is directly the K-block index.
+            kv_block_start = block_idx * BN
+            kv_start = kv_block_start
+
+            if _use_dbuf:
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+                k_base = k_buf_base(cur_buf)
+            else:
+                coop_load_k(kv_start, buf_id=0)
+                coop_load_k_as_v(kv_start, buf_id=0)
+                gpu.barrier()
+                k_base = k_buf_base(0)
+
+            # ==== GEMM1: s = Q @ K^T ====
+            k_hi_offset = K_SUB_N * K_STRIDE
+            k_swz_mask = (lane_mod_32 & arith.index(K_SWZ_ROW_MASK)) << arith.index(4)
+
+            def _k_idx_lo(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+            def _k_idx_hi(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+            s_acc_lo = c_zero_v16f32
+            s_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                k_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)])
+                k_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)])
+                s_acc_lo = mfma_acc(k_pack_lo, q_b_packs[ks], s_acc_lo)
+                s_acc_hi = mfma_acc(k_pack_hi, q_b_packs[ks], s_acc_hi)
+
+            # ==== Compute p[r] = exp(qk * sm_scale - LSE) with SWA + causal + boundary mask ====
+            kv_start_i32 = arith.index_cast(T.i32, kv_start)
+            seq_len_k_i32 = arith.index_cast(T.i32, seq_len_k_v)
+            seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v)
+            q_row_i32_mask = arith.index_cast(T.i32, q_row)
+            w_i32 = arith.constant(swa_window, type=T.i32)
+            lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32)
+            lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result
+
+            # Is this q_row out of bounds? In that case all mask elements = NEG_INF.
+            q_oob = arith.cmpi(
+                arith.CmpIPredicate.sge,
+                q_row_i32_mask,
+                seq_len_q_i32,
+            )
+
+            p_vals_lo = []
+            p_vals_hi = []
+            for r in range_constexpr(16):
+                r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32)
+
+                # lo half: col = kv_start + lane_off + r_off
+                kv_col_lo_i32 = arith.AddIOp(
+                    arith.AddIOp(kv_start_i32, lane_off_i32).result, r_off_i32
+                ).result
+                # boundary
+                is_oob_lo = arith.cmpi(arith.CmpIPredicate.sge, kv_col_lo_i32, seq_len_k_i32)
+                # causal: kv_col > q_row
+                is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_lo_i32, q_row_i32_mask)
+                # SWA: kv_col + W <= q_row
+                kv_plus_w_lo = arith.AddIOp(kv_col_lo_i32, w_i32).result
+                is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_lo, q_row_i32_mask)
+                bad_lo = arith.OrIOp(
+                    arith.OrIOp(
+                        arith.OrIOp(is_causal_lo, is_swa_lo).result,
+                        is_oob_lo,
+                    ).result,
+                    q_oob,
+                ).result
+
+                s_lo_f32 = vector.extract(s_acc_lo, static_position=[r], dynamic_position=[])
+                scaled_lo = arith.MulFOp(s_lo_f32, c_sm_scale, fastmath=fm_fast).result
+                scaled_lo_masked = arith.select(bad_lo, c_neg_inf, scaled_lo)
+                diff_lo = arith.SubFOp(scaled_lo_masked, lse_val, fastmath=fm_fast).result
+                p_lo = math_dialect.exp(diff_lo, fastmath=fm_fast)
+                p_vals_lo.append(p_lo)
+
+                # hi half: col = lo_col + K_SUB_N
+                kv_col_hi_i32 = arith.AddIOp(kv_col_lo_i32, arith.constant(K_SUB_N, type=T.i32)).result
+                is_oob_hi = arith.cmpi(arith.CmpIPredicate.sge, kv_col_hi_i32, seq_len_k_i32)
+                is_causal_hi = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_hi_i32, q_row_i32_mask)
+                kv_plus_w_hi = arith.AddIOp(kv_col_hi_i32, w_i32).result
+                is_swa_hi = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_hi, q_row_i32_mask)
+                bad_hi = arith.OrIOp(
+                    arith.OrIOp(
+                        arith.OrIOp(is_causal_hi, is_swa_hi).result,
+                        is_oob_hi,
+                    ).result,
+                    q_oob,
+                ).result
+
+                s_hi_f32 = vector.extract(s_acc_hi, static_position=[r], dynamic_position=[])
+                scaled_hi = arith.MulFOp(s_hi_f32, c_sm_scale, fastmath=fm_fast).result
+                scaled_hi_masked = arith.select(bad_hi, c_neg_inf, scaled_hi)
+                diff_hi = arith.SubFOp(scaled_hi_masked, lse_val, fastmath=fm_fast).result
+                p_hi = math_dialect.exp(diff_hi, fastmath=fm_fast)
+                p_vals_hi.append(p_hi)
+
+            # ==== Overwrite K LDS slot with V for GEMM2 ====
+            gpu.barrier()
+            if _use_dbuf:
+                coop_dma_v_as_k(kv_start, buf_id=cur_buf)
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+            elif ENABLE_DMA:
+                coop_dma_v_as_k(kv_start, buf_id=0)
+                rocdl.s_waitcnt(0)
+                gpu.barrier()
+            else:
+                coop_load_v_as_k(kv_start, buf_id=0)
+                gpu.barrier()
+
+            # ==== GEMM2: dP = DO @ V^T ====
+            dp_acc_lo = c_zero_v16f32
+            dp_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                v_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)])
+                v_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)])
+                dp_acc_lo = mfma_acc(v_pack_lo, do_b_packs[ks], dp_acc_lo)
+                dp_acc_hi = mfma_acc(v_pack_hi, do_b_packs[ks], dp_acc_hi)
+
+            # ==== Prefetch iter+1's K DMAs (async) ====
+            if _use_dbuf:
+                _next_block_idx = block_idx + arith.index(1)
+                _has_next = arith.cmpi(arith.CmpIPredicate.slt, _next_block_idx, n_block_end)
+                _pre_if = scf.IfOp(_has_next)
+                with ir.InsertionPoint(_pre_if.then_block):
+                    _next_kv_start = _next_block_idx * BN
+                    coop_dma_k(_next_kv_start, next_buf)
+                    coop_dma_k_as_v(_next_kv_start, next_buf)
+                    scf.YieldOp([])
+
+            # ==== Compute dS[r] = p[r] * (dp[r] - delta) ====
+            ds_vals_lo = []
+            ds_vals_hi = []
+            for r in range_constexpr(16):
+                dp_lo = vector.extract(dp_acc_lo, static_position=[r], dynamic_position=[])
+                dp_hi = vector.extract(dp_acc_hi, static_position=[r], dynamic_position=[])
+                diff_lo = arith.SubFOp(dp_lo, delta_val, fastmath=fm_fast).result
+                diff_hi = arith.SubFOp(dp_hi, delta_val, fastmath=fm_fast).result
+                ds_lo = arith.MulFOp(p_vals_lo[r], diff_lo, fastmath=fm_fast).result
+                ds_hi = arith.MulFOp(p_vals_hi[r], diff_hi, fastmath=fm_fast).result
+                ds_vals_lo.append(ds_lo)
+                ds_vals_hi.append(ds_hi)
+
+            # ==== Pack dS f32 -> mfma_pack_type (bf16/f16) ====
+            if dtype_str == "bf16" and USE_K16:
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    base = pks * 8
+                    ds_packs_lo.append(bf16_trunc_pack_v8(ds_vals_lo[base : base + 8]))
+                    ds_packs_hi.append(bf16_trunc_pack_v8(ds_vals_hi[base : base + 8]))
+            elif dtype_str == "bf16":
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    base = pks * 4
+                    ds_packs_lo.append(bf16_trunc_pack_v4(ds_vals_lo[base : base + 4]))
+                    ds_packs_hi.append(bf16_trunc_pack_v4(ds_vals_hi[base : base + 4]))
+            else:
+                ds_f16_lo = [arith.trunc_f(elem_type, ds_vals_lo[r]) for r in range_constexpr(16)]
+                ds_f16_hi = [arith.trunc_f(elem_type, ds_vals_hi[r]) for r in range_constexpr(16)]
+                _pack_ty = v8f16_type if USE_K16 else v4f16_type
+                _pw = 8 if USE_K16 else 4
+                ds_packs_lo = []
+                ds_packs_hi = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    b = pks * _pw
+                    ds_packs_lo.append(vector.from_elements(_pack_ty, [ds_f16_lo[b + i] for i in range(_pw)]))
+                    ds_packs_hi.append(vector.from_elements(_pack_ty, [ds_f16_hi[b + i] for i in range(_pw)]))
+
+            # ==== GEMM3: dQ += K @ dS  (uses fwd's V-read schedule with K substituted) ====
+            _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)]
+            TOTAL_PV = len(_steps)
+            v_base = v_buf_base(cur_buf) if _use_dbuf else v_buf_base(0)
+
+            def _read_k_as_v_pack(step_idx):
+                dc, pks = _steps[step_idx]
+                if USE_HW_TR:
+                    d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4
+                    k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group
+                    _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col
+                    lds_lo = v_base + k_row * V_STRIDE + _d_col_eff
+                    lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE)
+                    if USE_K16:
+                        vl_a = ds_read_tr_v4f16(lds_lo)
+                        vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE))
+                        vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                        vh_a = ds_read_tr_v4f16(lds_hi)
+                        vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE))
+                        vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                    else:
+                        vl = ds_read_tr_v4f16(lds_lo)
+                        vh = ds_read_tr_v4f16(lds_hi)
+                else:
+                    d_pos = arith.index(dc * D_CHUNK) + lane_mod_32
+                    kb = arith.index(pks * PV_K_STEP) + lane_div_32 * 4
+                    v_lo_idx = v_base + d_pos * VT_STRIDE + kb
+                    v_hi_idx = v_lo_idx + arith.index(K_SUB_N)
+                    vl = vector.load(v4f16_type, lds_kv, [v_lo_idx])
+                    vh = vector.load(v4f16_type, lds_kv, [v_hi_idx])
+                return vl, vh
+
+            k_lo_cur, k_hi_cur = _read_k_as_v_pack(0)
+            for si in range_constexpr(TOTAL_PV):
+                dc, pks = _steps[si]
+                if si + 1 < TOTAL_PV:
+                    k_lo_nxt, k_hi_nxt = _read_k_as_v_pack(si + 1)
+                dq_accs[dc] = mfma_acc(k_lo_cur, ds_packs_lo[pks], dq_accs[dc])
+                dq_accs[dc] = mfma_acc(k_hi_cur, ds_packs_hi[pks], dq_accs[dc])
+                if si + 1 < TOTAL_PV:
+                    k_lo_cur = k_lo_nxt
+                    k_hi_cur = k_hi_nxt
+
+            # End of iter: yield dq_accs and swapped cur_buf.
+            gpu.barrier()
+            _yield = list(dq_accs)
+            if _use_dbuf:
+                _yield.append(next_buf)
+            yield _yield
+
+        # ---- Final store: dQ_fp32 = dq_acc * sm_scale (NO trunc; fp32) ----
+        dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)]
+        sm_scale_vec = vector.broadcast(v16f32_type, c_sm_scale)
+
+        _o_guard = scf.IfOp(q_in_bounds, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for dc in range_constexpr(D_CHUNKS):
+                dq_scaled = arith.MulFOp(
+                    dq_finals[dc],
+                    sm_scale_vec,
+                    fastmath=fm_fast,
+                ).result
+                for r in range_constexpr(16):
+                    dq_val = vector.extract(
+                        dq_scaled,
+                        static_position=[r],
+                        dynamic_position=[],
+                    )
+                    d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4)
+                    d_col = arith.index(dc * D_CHUNK) + d_row_rel
+                    dq_global = global_idx_q(q_row, d_col)
+                    _gep_store_f32(dq_val, dq_ptr, dq_global)
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_swa_bwd_dq(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        DOS: fx.Tensor,
+        LSE: fx.Tensor,
+        DELTAS: fx.Tensor,
+        DQ: fx.Tensor,
+        DSINK: fx.Tensor,
+        SINK: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len_q: fx.Int32,
+        seq_len_k: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len_q)
+        num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M
+        grid_x = bs_idx * num_q_tiles * NUM_HEADS
+
+        launcher = v4_swa_bwd_dq_kernel(
+            Q,
+            K,
+            V,
+            DOS,
+            LSE,
+            DELTAS,
+            DQ,
+            DSINK,
+            SINK,
+            seq_len_q,
+            seq_len_k,
+        )
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(
+                            T.i32,
+                            _wpe,
+                        )
+        if flat_work_group_size is not None:
+            _fwgs = int(flat_work_group_size)
+            if _fwgs >= 1:
+                flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, 1, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    _fmha_compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {
+            "enable-post-misched": False,
+            "lsr-drop-solution": True,
+        },
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return launch_v4_swa_bwd_dq(*args, **kwargs)
+
+    def _compile(Q, K, V, DOS, LSE, DELTAS, DQ, DSINK, SINK, batch_size, seq_len_q, seq_len_k, stream=None):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return flyc.compile(
+                launch_v4_swa_bwd_dq,
+                Q,
+                K,
+                V,
+                DOS,
+                LSE,
+                DELTAS,
+                DQ,
+                DSINK,
+                SINK,
+                batch_size,
+                seq_len_q,
+                seq_len_k,
+                fx.Stream(stream),
+            )
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+# Convenience alias.
+build_v4_swa_bwd_dq_module_primary = build_v4_swa_bwd_dq_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py
new file mode 100644
index 000000000..9efb746a5
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py
@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_sla_bwd_kernel: V4 SWA attention backward kernels (FlyDSL, STEP 1).
+
+Forked from /workspace/FlyDSL-amd/kernels/sla_bwd_preprocess.py.
+
+STEP 1 SCOPE
+============
+Only the preprocess kernel is implemented in this file. dq / dkv are still
+handled by Triton in the v4_attention_bwd_flydsl_mqa wrapper (see that file
+for the rationale). The hooks defined here will be the landing spots for
+the dq / dkv kernels in STEP 1b / STEP 1c.
+
+PREPROCESS KERNEL (D scalar)
+---------------------------
+Computes ``delta[b, h, m] = sum_d (out[b, h, m, d] * dout[b, h, m, d])`` in
+fp32 for every query row. This is the standard FA-2 pre-pass and is dense
+(no SWA / sink dependency) so it is the simplest piece to lift to FlyDSL
+first.
+
+Inputs (flat views):
+    OS, DOS  shape (N_ROWS, D) where N_ROWS = B*HQ*Sq
+    DELTAS   shape (N_ROWS,)    fp32
+
+Grid: (N_ROWS / BLOCK_ROWS, 1, 1).
+Block: BLOCK_ROWS * THREADS_PER_ROW threads.
+
+Each row uses THREADS_PER_ROW = D // VEC_WIDTH threads. Cross-thread row
+reduction uses ``shuffle_xor`` at offsets THREADS_PER_ROW/2, /4, ..., 1
+which stays inside a THREADS_PER_ROW-lane subgroup (XOR never carries out
+of a power-of-2 subgroup).
+
+Requires: head_dim % VEC_WIDTH == 0, N_ROWS % BLOCK_ROWS == 0.
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr import arith, buffer_ops, range_constexpr
+from flydsl.expr.arith import ArithValue
+from flydsl.expr.numeric import Float32
+from flydsl.expr.typing import Int32, T
+from flydsl.expr.vector import ReductionOp
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from kernels.kernels_common import dtype_to_elem_type
+
+KERNEL_NAME_PRE = "v4_swa_bwd_preprocess_kernel"
+
+# Wider VEC_WIDTH than the SLA kernel (D=512 vs D=128). At D=512 with
+# VEC_WIDTH=8 we would need 64 threads per row (== a full warp), which is
+# fine but uses one wave per BLOCK_ROWS rows. We use VEC_WIDTH=16 so
+# THREADS_PER_ROW=32 and a 256-thread block processes 8 rows = 2 waves
+# (more parallel tile generators per CU at the cost of slightly wider
+# loads). Either is correct.
+VEC_WIDTH = 8
+WARP_SIZE = 64
+
+
+def build_v4_swa_bwd_preprocess_module(
+    head_dim,
+    dtype_str="bf16",
+    block_rows=None,
+):
+    """Build the V4 SWA backward preprocess launcher.
+
+    Inputs (flattened to 2D by the launcher):
+        OS, DOS  shape (N_ROWS, D) where N_ROWS = B*H*L
+        DELTAS   shape (N_ROWS,)    f32
+
+    Args:
+        head_dim: head dimension (must be divisible by VEC_WIDTH).
+        dtype_str: "bf16" or "f16" for O_S / DO_S; DELTAS is always f32.
+        block_rows: rows processed per block. Default 8.
+    """
+    get_hip_arch()
+
+    if block_rows is None:
+        BLOCK_ROWS = 8
+    else:
+        BLOCK_ROWS = block_rows
+
+    assert head_dim % VEC_WIDTH == 0, f"head_dim {head_dim} must be % {VEC_WIDTH}"
+    assert head_dim >= VEC_WIDTH, f"head_dim {head_dim} < VEC_WIDTH {VEC_WIDTH}"
+    assert dtype_str in ("bf16", "f16"), f"unsupported dtype {dtype_str}"
+
+    D = head_dim
+    THREADS_PER_ROW = D // VEC_WIDTH  # 32 for D=512, VEC_WIDTH=16
+    BLOCK_THREADS = BLOCK_ROWS * THREADS_PER_ROW
+    assert BLOCK_THREADS <= 1024, f"BLOCK_THREADS {BLOCK_THREADS} > 1024"
+    assert (
+        THREADS_PER_ROW & (THREADS_PER_ROW - 1) == 0
+    ), "THREADS_PER_ROW must be power of 2 for shfl_xor reduction"
+
+    elem_bits = 16  # bf16 / f16
+
+    @flyc.kernel
+    def v4_swa_bwd_preprocess_kernel(
+        OS: fx.Tensor,  # (N_ROWS, D)
+        DOS: fx.Tensor,  # (N_ROWS, D)
+        DELTAS: fx.Tensor,  # (N_ROWS,)
+    ):
+        bid = fx.block_idx.x
+        tid = fx.thread_idx.x
+
+        elem_type = dtype_to_elem_type(dtype_str)
+        T.f32
+        fm_fast = arith.FastMathFlags.fast
+
+        # ---- Thread decomposition ----
+        row_in_block_i32 = tid // Int32(THREADS_PER_ROW)
+        col_in_row_i32 = tid % Int32(THREADS_PER_ROW)
+
+        global_row_i32 = bid * Int32(BLOCK_ROWS) + row_in_block_i32
+        global_row_idx = ArithValue(global_row_i32).index_cast(T.index)
+
+        # ---- Buffer-backed 2D tensors for loads ----
+        OS_buf = fx.rocdl.make_buffer_tensor(OS)
+        DOS_buf = fx.rocdl.make_buffer_tensor(DOS)
+        delta_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True)
+
+        row_os = fx.slice(OS_buf, (global_row_i32, None))
+        row_dos = fx.slice(DOS_buf, (global_row_i32, None))
+
+        in_div_os = fx.logical_divide(row_os, fx.make_layout(VEC_WIDTH, 1))
+        in_div_dos = fx.logical_divide(row_dos, fx.make_layout(VEC_WIDTH, 1))
+
+        copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits)
+        vec_reg_ty = fx.MemRefType.get(elem_type, fx.LayoutType.get(VEC_WIDTH, 1), fx.AddressSpace.Register)
+        vec_reg_lay = fx.make_layout(VEC_WIDTH, 1)
+
+        def _load_vec(div_tensor, col_idx):
+            r = fx.memref_alloca(vec_reg_ty, vec_reg_lay)
+            fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, col_idx)), r)
+            return fx.memref_load_vec(r)
+
+        os_vec = _load_vec(in_div_os, col_in_row_i32)
+        dos_vec = _load_vec(in_div_dos, col_in_row_i32)
+
+        os_f32 = os_vec.to(Float32)
+        dos_f32 = dos_vec.to(Float32)
+        prod_f32 = os_f32 * dos_f32
+        local_sum = prod_f32.reduce(ReductionOp.ADD, fastmath=fm_fast)
+
+        width_i32 = Int32(WARP_SIZE)
+        val = local_sum
+        num_rounds = int(math.log2(THREADS_PER_ROW))
+        for _sh_exp in range_constexpr(num_rounds):
+            off = Int32(THREADS_PER_ROW // (2 << _sh_exp))
+            peer = val.shuffle_xor(off, width_i32)
+            val = val.addf(peer, fastmath=fm_fast)
+
+        if col_in_row_i32 == Int32(0):
+            delta_off_i32 = arith.index_cast(T.i32, global_row_idx)
+            val_ir = val.ir_value() if hasattr(val, "ir_value") else val
+            buffer_ops.buffer_store(val_ir, delta_rsrc, delta_off_i32)
+
+    @flyc.jit
+    def launch_v4_swa_bwd_preprocess(
+        OS: fx.Tensor,
+        DOS: fx.Tensor,
+        DELTAS: fx.Tensor,
+        n_rows: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        """Launch grid: (n_rows / BLOCK_ROWS, 1, 1).
+
+        Expects OS, DOS to be views of shape (n_rows, D). DELTAS is (n_rows,).
+        """
+        blocks_i32 = n_rows // Int32(BLOCK_ROWS)
+        blocks_idx = ArithValue(blocks_i32).index_cast(T.index)
+        launcher = v4_swa_bwd_preprocess_kernel(OS, DOS, DELTAS)
+        launcher.launch(
+            grid=(blocks_idx, arith.index(1), arith.index(1)),
+            block=(BLOCK_THREADS, 1, 1),
+            stream=stream,
+        )
+
+    return launch_v4_swa_bwd_preprocess
+
+
+# Convenience alias used by the wrapper.
+def build_v4_swa_bwd_preprocess(*args, **kwargs):
+    return build_v4_swa_bwd_preprocess_module(*args, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# STEP 1b: dq kernel re-export (forked from kernels/sla_bwd_dq.py).
+# Lives in v4_sla_bwd_dq_kernel.py for file-size reasons; re-exported here
+# so the wrapper can import a single module.
+# ---------------------------------------------------------------------------
+import os as _os
+import sys as _sys
+
+_HERE = _os.path.dirname(_os.path.abspath(__file__))
+if _HERE not in _sys.path:
+    _sys.path.insert(0, _HERE)
+from v4_sla_bwd_dq_kernel import (  # noqa: E402,F401
+    build_v4_swa_bwd_dq_module,
+    build_v4_swa_bwd_dq_module_primary,
+)
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py
new file mode 100644
index 000000000..5e331ec77
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py
@@ -0,0 +1,1387 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""v4_swa_fwd: V4 dense sliding-window-causal attention forward (FlyDSL).
+
+Forked from kernels/sla_fwd.py. Differences:
+  - No LUT: outer KV loop iterates a contiguous block range bounded by
+    the SWA window for the current Q tile (depends on q_tile_idx, but
+    wave-uniform).
+  - Per-element SWA mask: kv_col > q_row OR kv_col + swa_window <= q_row
+    OR kv_col >= seq_len -> NEG_INF, applied before softmax max-reduce.
+  - LSE is fp32 raw-domain `lse = m_final*scale + ln(l_final)` to match
+    V4 Triton reference (m_i + tl.log(l_i)).
+  - A-1 scope: no sink, no additive mask, no HCA. MQA broadcast at the
+    Python launcher (kernel sees full MHA K/V).
+
+Layout: BHLD. Q/K/V/O flattened from (B, H, L, D).
+LSE: (B, H, L) flat, fp32.
+Grid: (B * num_q_tiles * H,), num_q_tiles = L / BLOCK_M.
+"""
+
+import math
+import os
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import memref as _memref
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+from kernels.kernels_common import dtype_to_elem_type
+
+# ---- Module-level constants ----
+
+KERNEL_NAME = "v4_swa_fwd_kernel"
+
+_LOG2E = math.log2(math.e)  # 1.4426950408889634
+
+_LLVM_GEP_DYNAMIC = -2147483648  # LLVM kDynamicIndex sentinel (0x80000000 as signed i32)
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+_VMCNT_LO_MASK = 0xF
+_LGKMCNT_EXPCNT_BASE = 0x3F70
+_VMCNT_HI_SHIFT = 14
+_VMCNT_HI_MASK = 0x3
+
+
+def _waitcnt_vm_n(n):
+    """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7)."""
+    val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT)
+    rocdl.s_waitcnt(val)
+
+
+def build_v4_swa_fwd_module(
+    num_heads,
+    head_dim,
+    swa_window,
+    dtype_str="bf16",
+    sm_scale=None,
+    waves_per_eu=2,
+    flat_work_group_size=None,
+    block_m=None,
+    block_n=None,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+    layout_bhld=True,
+    mqa_kv=False,
+):
+    """Build the SLA sparse-fwd launcher. Single-variant only (no auto-dispatch).
+
+    mqa_kv: if True, K and V are MQA-shape [B, 1, Sk, D] (stride_kh=0) and
+    indexing into them drops head_idx. The kernel still sees Q at [B, H, Sq, D].
+
+    layout_bhld:
+      False → Q/K/V/O in (B, L, H, D) layout (inherited from flash_attn_func,
+              matches torch `(batch, seq, num_heads, head_dim)` natural form).
+      True  → Q/K/V/O in (B, H, L, D) layout (SLA's native form; avoids a
+              transpose at the integration boundary in `SparseLinearAttention`).
+      LUT and LSE layouts are unchanged either way:
+        LUT (B, H, M_BLOCKS, topk) flattened, dtype i32.
+        LSE (B, H, L) flattened, dtype f32.
+    """
+    gpu_arch = get_hip_arch()
+
+    if block_n is None:
+        BLOCK_N = 64
+    else:
+        BLOCK_N = int(block_n)
+    K_SUB_N = min(BLOCK_N, 32)
+    N_HALVES = BLOCK_N // 32  # 1 or 2; how many K_SUB_N halves per BLOCK_N
+    WARP_SIZE = 64
+    # SLA is always non-causal; the sparse map decides which blocks attend.
+
+    if block_m is not None:
+        BLOCK_M = block_m
+    else:
+        BLOCK_M = 128
+
+    if flat_work_group_size is None:
+        if BLOCK_M <= 128:
+            flat_work_group_size = 256
+        else:
+            flat_work_group_size = 512
+    NUM_WAVES = flat_work_group_size // WARP_SIZE
+    BLOCK_SIZE = flat_work_group_size
+    ROWS_PER_WAVE = BLOCK_M // NUM_WAVES
+    # V4 SWA: always N32 path. One outer iter = one BLOCK_N block.
+    BLOCK_N_OUT = BLOCK_N
+    N_SUBTILES = 1
+    ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1"
+    # buffer_load_dwordx4_lds (16B DMA-to-LDS) requires gfx950+; gfx94x only has dword (4B).
+    # For SLA we default-on DMA when hardware supports it — this is the whole point.
+    _has_lds_load_b128 = not gpu_arch.startswith("gfx942")
+    ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1")
+    ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1"
+    REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower()
+    if REDUCE_MODE not in ("xor", "ds_bpermute"):
+        REDUCE_MODE = "xor"
+    FORCE_SINGLE_BUF_DMA = os.getenv("FLYDSL_V4_SWA_SINGLE_BUF_DMA", "0") == "1"
+    if ENABLE_PREFETCH_3BUF:
+        NUM_PREFETCH_K = 3
+    elif ENABLE_DMA and not FORCE_SINGLE_BUF_DMA:
+        NUM_PREFETCH_K = 2
+    else:
+        NUM_PREFETCH_K = 1
+    # Lever A: double-buffer V in DMA-dbuf mode so V HBM latency overlaps
+    # with the QK MFMA of the next iteration (mirrors K-dbuf behavior).
+    if ENABLE_PREFETCH_3BUF:
+        NUM_PREFETCH_V = 3
+    elif ENABLE_DMA and not FORCE_SINGLE_BUF_DMA:
+        NUM_PREFETCH_V = 2
+    else:
+        NUM_PREFETCH_V = 1
+    CK_LDS_SEQ = (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,)
+
+    # gfx950+ has ds_read_tr16_b64 (HW transpose LDS read); gfx942 needs V^T stored in LDS.
+    USE_HW_TR = gpu_arch.startswith("gfx950")
+
+    # MFMA32 K-dimension: 16 on gfx950+ (CDNA4) for both GEMMs.
+    USE_K16 = gpu_arch.startswith("gfx950")
+    K_STEP_QK = 16 if USE_K16 else 8
+    K_STEPS_QK = head_dim // K_STEP_QK
+    D_CHUNK = 32
+    D_CHUNKS = head_dim // D_CHUNK
+    PV_K_STEP = 16 if USE_K16 else 8
+    PV_K_STEPS = K_SUB_N // PV_K_STEP  # 2 steps per sub-tile (K=16) or 4 (K=8)
+
+    assert BLOCK_M % NUM_WAVES == 0
+    assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32"
+    assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64"
+    assert flat_work_group_size in (
+        128,
+        256,
+        512,
+    ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}"
+    assert dtype_str in ("f16", "bf16"), "sla_fwd only supports f16 and bf16"
+    assert BLOCK_N % 32 == 0 or BLOCK_N == 32
+    assert BLOCK_N_OUT == BLOCK_N
+    assert N_HALVES in (1, 2)
+    assert isinstance(swa_window, int) and swa_window > 0, f"swa_window must be int > 0, got {swa_window!r}"
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(head_dim)
+
+    NUM_HEADS = num_heads
+    HEAD_DIM = head_dim
+    STRIDE_TOKEN = NUM_HEADS * HEAD_DIM
+
+    # Bank-conflict-free LDS strides.
+    # K uses XOR swizzle (col ^ ((row & 7) << 4)) at 16-element granularity
+    # instead of padding. This enables ds_read_b128 (stride is 256B-aligned).
+    K_STRIDE = HEAD_DIM
+    if USE_HW_TR:
+        V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4
+    else:
+        VT_STRIDE = BLOCK_N + 2
+        V_STRIDE = VT_STRIDE
+
+    # Vectorized cooperative load constants.
+    VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+
+    if ROWS_PER_BATCH_LOAD >= BLOCK_N:
+        NUM_BATCHES_KV = 1
+        KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N
+    else:
+        assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+        NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+        KV_NEEDS_GUARD = False
+
+    # K/V circular buffers; defaults to 1/1, optional 3/3 with CK-like LDS sequence.
+    LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE
+    if USE_HW_TR:
+        LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE
+    else:
+        LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE
+    LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE
+    LDS_V_BASE = LDS_K_TOTAL_SIZE
+    LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE
+    LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"v4_swa_fwd_smem_M{BLOCK_M}_N{BLOCK_N}_W{swa_window}",
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def v4_swa_fwd_kernel(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        O: fx.Tensor,
+        LSE: fx.Tensor,
+        seq_len: fx.Int32,
+    ):
+        elem_type = dtype_to_elem_type(dtype_str)
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K)
+        v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V)
+        o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O)
+        # LSE: f32 scalar writes via buffer_store.
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+
+        # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation).
+        # The unsafe_fp_math/fast_fp_math builder params control LLVM-level attributes only.
+        fm_fast = arith.FastMathFlags.fast
+        v4f16_type = T.vec(4, elem_type)
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        v16f32_type = T.vec(16, compute_type)
+        mfma_pack_type = v8f16_type if USE_K16 else v4f16_type
+        MFMA_LANE_K = 8 if USE_K16 else 4
+        _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0)
+
+        def _mfma(ods_fn, a, b, c):
+            return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result
+
+        def mfma_acc(a, b, c):
+            if dtype_str == "bf16":
+                if USE_K16:
+                    return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c)
+                a = vector.bitcast(T.i16x4, a)
+                b = vector.bitcast(T.i16x4, b)
+                return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c)
+            if USE_K16:
+                return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c)
+            return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c)
+
+        seq_len_v = arith.index_cast(T.index, seq_len)
+
+        # ---- LDS view ----
+        base_ptr = allocator.get_base()
+        lds_kv = SmemPtr(
+            base_ptr,
+            lds_kv_offset,
+            elem_type,
+            shape=(LDS_KV_TOTAL_SIZE,),
+        ).get()
+
+        # ---- Thread / block indices ----
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+
+        # ---- Wave decomposition ----
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        lane_mod_32 = lane % 32
+        lane_div_32 = lane // 32  # 0/1
+
+        # ---- ds_read_b64_tr_b16 lane decomposition ----
+        # Hardware does 4×4 transpose within blocks of 16 lanes.
+        # tr_k_group selects which of 4 K-rows within the block,
+        # tr_col_sub selects which 4-column sub-group within 16 columns.
+        tr_k_group = (lane % 16) // 4  # 0..3: K-row offset within 4-row group
+        tr_col_sub = lane % 4  # 0..3: 4-column sub-group
+        tr_col_half = (lane % 32) // 16  # 0 or 1: first/second 16-column half
+
+        # ---- ds_read_b64_tr_b16 helper ----
+
+        def ds_read_tr_v4f16(lds_elem_idx):
+            """Read v4f16 from LDS with hardware transpose.
+
+            Within each block of 16 lanes, the hardware performs a 4×4
+            transpose across 4 groups of 4 lanes.  After the transpose,
+            result[lane, elem_e] = Input[source_lane, lane%4] where
+            source_lane = e*4 + (lane%16)//4.  This naturally produces
+            the MFMA A-operand layout when per-lane addresses point to
+            the correct K-row and D-column sub-group.
+            """
+            byte_offset = lds_elem_idx * 2 + lds_kv_offset
+            byte_i64 = arith.index_cast(T.i64, byte_offset)
+            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result
+
+        # ---- Wave offsets ----
+        wave_q_offset = wave_id * ROWS_PER_WAVE
+
+        # ---- Decompose block_id ----
+        head_idx = block_id % NUM_HEADS
+        batch_q_tile_id = block_id // NUM_HEADS
+        num_q_tiles = (seq_len_v + BLOCK_M - 1) // BLOCK_M
+        q_tile_idx = batch_q_tile_id % num_q_tiles
+        batch_idx = batch_q_tile_id // num_q_tiles
+        q_start = q_tile_idx * BLOCK_M
+
+        # ---- V4 SWA: per-tile contiguous K-block range (wave-uniform) ----
+        # Q tile rows: [q_start, q_start + BLOCK_M). Union of visible K
+        # columns under SWA-causal: [q_start - W + 1, q_start + BLOCK_M).
+        # n_block_start = max(0, q_start - W + 1) // BLOCK_N
+        # n_block_end   = ceil(min(q_start + BLOCK_M, seq_len), BLOCK_N)
+        # Per-element causal/SWA/boundary mask still applied inside the loop.
+        SWA = arith.index(swa_window)
+        BN = arith.index(BLOCK_N)
+        BM = arith.index(BLOCK_M)
+        _zero_idx = arith.index(0)
+        _one_idx = arith.index(1)
+        # q_start - W + 1 -- compute as index arithmetic, then clamp.
+        _q_plus_one = q_start + _one_idx
+        # We need max(0, _q_plus_one - SWA). To avoid negative-index issues
+        # we test _q_plus_one >= SWA first.
+        _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _q_plus_one, SWA)
+        _n_start_row = arith.select(_ge_w, _q_plus_one - SWA, _zero_idx)
+        n_block_start = _n_start_row // BN
+        _n_end_row_uncl = q_start + BM
+        _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _n_end_row_uncl, seq_len_v)
+        n_end_row_cl = arith.select(_le_seq, _n_end_row_uncl, seq_len_v)
+        n_block_end = (n_end_row_cl + BN - _one_idx) // BN
+
+        # ---- Cooperative load decomposition ----
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        # ---- Helper: global flat index ----
+        # BLHD: stride = (L*H*D, H*D, D, 1) → (token=b*L + l) * (H*D) + h*D + col
+        # BHLD: stride = (H*L*D, L*D, D, 1) → ((b*H + h) * L + l) * D + col
+        if layout_bhld:
+            bh_base_tokens = (batch_idx * NUM_HEADS + head_idx) * seq_len_v
+            if mqa_kv:
+                # K/V are [B, 1, Sk, D] -- drop head_idx from KV indexing.
+                bh_base_tokens_kv = batch_idx * seq_len_v
+            else:
+                bh_base_tokens_kv = bh_base_tokens
+
+            def global_idx(token_idx, col):
+                return (bh_base_tokens + token_idx) * arith.index(HEAD_DIM) + col
+
+            def global_idx_kv(token_idx, col):
+                return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col
+
+        else:
+
+            def global_idx(token_idx, col):
+                token = batch_idx * seq_len_v + token_idx
+                return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col
+
+            if mqa_kv:
+                # BLHD MQA: K/V token = batch_idx * seq_len_v + token_idx (no h*D added).
+                def global_idx_kv(token_idx, col):
+                    token = batch_idx * seq_len_v + token_idx
+                    return token * HEAD_DIM + col  # no STRIDE_TOKEN, no head_idx
+
+            else:
+                global_idx_kv = global_idx
+
+        def _gep_load(base_ptr, elem_idx, vec_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store(val, base_ptr, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                base_ptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_f16x4(base_ptr, base_idx):
+            return _gep_load(base_ptr, base_idx, v4f16_type)
+
+        def load_global_mfma_pack(base_ptr, base_idx):
+            return _gep_load(base_ptr, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(base_ptr, base_idx):
+            return _gep_load(base_ptr, base_idx, vxf16_type)
+
+        def bf16_trunc_pack_v4(f32_vals):
+            """Pack 4 f32 values into v4bf16 via bitwise truncation (upper 16 bits).
+            ~2 fewer instructions/element vs arith.TruncFOp round-to-nearest."""
+            _v2i32 = T.vec(2, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32)
+            b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32)
+            p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result
+            a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32)
+            b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32)
+            p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result
+            return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1]))
+
+        def bf16_trunc_pack_v8(f32_vals):
+            """Pack 8 f32 values into v8bf16 via bitwise truncation (upper 16 bits)."""
+            _v4i32 = T.vec(4, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            pairs = []
+            for j in range_constexpr(4):
+                a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32)
+                b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32)
+                p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result
+                pairs.append(p)
+            return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs))
+
+        def k_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(buf_id * LDS_K_TILE_SIZE)
+            return buf_id * arith.index(LDS_K_TILE_SIZE)
+
+        def v_buf_base(buf_id):
+            if isinstance(buf_id, int):
+                return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE)
+            return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE)
+
+        # ---- K XOR swizzle: col ^ ((row & 7) << 4) at 16-element granularity ----
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(0x7)) << arith.index(4)
+            return col_idx ^ mask
+
+        # ---- Cooperative K load (row-major, XOR-swizzled) ----
+        def coop_load_k(tile_start, buf_id=0):
+            k_base = k_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_k = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_k.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        swz_col = _k_swizzle(lds_row, load_col_base)
+                        lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                        vec = load_global_f16xN(k_ptr, g_idx)
+                        vector.store(vec, lds_kv, [lds_idx])
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    swz_col = _k_swizzle(lds_row, load_col_base)
+                    lds_idx = k_base + lds_row * K_STRIDE + swz_col
+                    vec = load_global_f16xN(k_ptr, g_idx)
+                    vector.store(vec, lds_kv, [lds_idx])
+
+        # ---- Cooperative V load ----
+        def _v_store_row_major(v_base, lds_row, vec):
+            lds_idx = v_base + lds_row * V_STRIDE + load_col_base
+            vector.store(vec, lds_kv, [lds_idx])
+
+        _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None
+
+        def _v_store_transposed(v_base, lds_row, vec):
+            for _e in range_constexpr(VEC_WIDTH):
+                elem = vector.extract(vec, static_position=[_e], dynamic_position=[])
+                vt_d = load_col_base + _e
+                vt_idx = v_base + vt_d * VT_STRIDE + lds_row
+                v1 = vector.from_elements(_v1_type, [elem])
+                vector.store(v1, lds_kv, [vt_idx])
+
+        _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed
+
+        def coop_load_v(tile_start, buf_id=0):
+            v_base = v_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        g_idx = global_idx_kv(row_idx, load_col_base)
+                        lds_row = load_row_in_batch + row_offset
+                        vec = load_global_f16xN(v_ptr, g_idx)
+                        _v_store_to_lds(v_base, lds_row, vec)
+                        scf.YieldOp([])
+                else:
+                    g_idx = global_idx_kv(row_idx, load_col_base)
+                    lds_row = load_row_in_batch + row_offset
+                    vec = load_global_f16xN(v_ptr, g_idx)
+                    _v_store_to_lds(v_base, lds_row, vec)
+
+        def coop_load_v_global(tile_start):
+            """Issue global loads for V, return vectors (non-blocking)."""
+            vecs = []
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                row_idx = tile_start + load_row_in_batch + row_offset
+                g_idx = global_idx_kv(row_idx, load_col_base)
+                vecs.append(load_global_f16xN(v_ptr, g_idx))
+            return vecs
+
+        def coop_store_v_lds(vecs, buf_id=0):
+            """Write previously-loaded V vectors to LDS."""
+            v_base = v_buf_base(buf_id)
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                row_offset = batch * ROWS_PER_BATCH_LOAD
+                if KV_NEEDS_GUARD:
+                    row_valid = arith.cmpi(
+                        arith.CmpIPredicate.ult,
+                        load_row_in_batch,
+                        arith.index(BLOCK_N),
+                    )
+                    _if_v = scf.IfOp(row_valid)
+                    with ir.InsertionPoint(_if_v.then_block):
+                        lds_row = load_row_in_batch + row_offset
+                        _v_store_to_lds(v_base, lds_row, vecs[batch])
+                        scf.YieldOp([])
+                else:
+                    lds_row = load_row_in_batch + row_offset
+                    _v_store_to_lds(v_base, lds_row, vecs[batch])
+
+        # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ----
+        if ENABLE_DMA:
+            from flydsl._mlir.dialects import llvm
+
+            k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True)
+            _lds_ptr_ty = _llvm_lds_ptr_ty()
+            DMA_BYTES = 16  # buffer_load_dwordx4 = 16 bytes per lane
+            DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES
+            K_TILE_BYTES = BLOCK_N * K_STRIDE * 2
+            NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+            lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv)
+            _dma_size = arith.constant(DMA_BYTES, type=T.i32)
+            _dma_soff = arith.constant(0, type=T.i32)
+            _dma_off = arith.constant(0, type=T.i32)
+            _dma_aux = arith.constant(1, type=T.i32)
+
+            def coop_dma_k(tile_start, buf_id=0):
+                """Load K tile via DMA with XOR-swizzled global fetch."""
+                if isinstance(buf_id, int):
+                    k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2)
+                else:
+                    k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2)
+                for d in range_constexpr(NUM_DMA_K):
+                    lds_addr = (
+                        k_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH)
+                    swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(0x7)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    # Layout-aware global byte offset.
+                    if layout_bhld:
+                        # BHLD: ((b*H + h)*L + (tile_start + row)) * (D*2) + col_byte
+                        # MQA: ((b)*L + (tile_start + row)) * (D*2) + col_byte
+                        row_within_head = tile_start + row_in_tile
+                        global_byte = (bh_base_tokens_kv + row_within_head) * arith.index(
+                            HEAD_DIM * 2
+                        ) + col_byte
+                    else:
+                        # BLHD: ((b*L + row) * (H*D) + h*D) * 2 + col_byte
+                        global_row = batch_idx * seq_len_v + tile_start + row_in_tile
+                        if mqa_kv:
+                            global_byte = global_row * arith.index(HEAD_DIM * 2) + col_byte
+                        else:
+                            global_byte = (
+                                global_row * arith.index(STRIDE_TOKEN * 2)
+                                + head_idx * arith.index(HEAD_DIM * 2)
+                                + col_byte
+                            )
+                    voffset = arith.index_cast(T.i32, global_byte)
+
+                    rocdl.raw_ptr_buffer_load_lds(
+                        k_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        # ---- V XOR swizzle: col ^ ((row & 3) << 4) at 16-element granularity ----
+        def _v_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(0x3)) << arith.index(4)
+            return col_idx ^ mask
+
+        # ---- DMA loading for V (buffer_load_dwordx4 ... lds) ----
+        if ENABLE_DMA:
+            v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True)
+            V_TILE_BYTES = BLOCK_N * V_STRIDE * 2
+            NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES
+            LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES
+            ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2)
+
+            def coop_dma_v(tile_start, buf_id=0):
+                """Load V tile via DMA with XOR-swizzled global fetch."""
+                if isinstance(buf_id, int):
+                    v_lds_byte_base = lds_kv_base_idx + arith.index(
+                        (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2
+                    )
+                else:
+                    v_lds_byte_base = (
+                        lds_kv_base_idx
+                        + arith.index(LDS_V_BASE * 2)
+                        + buf_id * arith.index(LDS_V_TILE_SIZE * 2)
+                    )
+                for d in range_constexpr(NUM_DMA_V):
+                    lds_addr = (
+                        v_lds_byte_base
+                        + wave_id * arith.index(WARP_SIZE * DMA_BYTES)
+                        + arith.index(d * DMA_BATCH_BYTES)
+                    )
+                    lds_i64 = arith.index_cast(T.i64, lds_addr)
+                    lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64)
+                    lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result
+
+                    row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V)
+                    swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2)
+                    xor_mask = (row_in_tile & arith.index(0x3)) << arith.index(4)
+                    unsw_col_f16 = swiz_col_f16 ^ xor_mask
+                    col_byte = unsw_col_f16 * 2
+                    # Layout-aware global byte offset (see K DMA path above).
+                    if layout_bhld:
+                        row_within_head = tile_start + row_in_tile
+                        global_byte = (bh_base_tokens_kv + row_within_head) * arith.index(
+                            HEAD_DIM * 2
+                        ) + col_byte
+                    else:
+                        global_row = batch_idx * seq_len_v + tile_start + row_in_tile
+                        if mqa_kv:
+                            global_byte = global_row * arith.index(HEAD_DIM * 2) + col_byte
+                        else:
+                            global_byte = (
+                                global_row * arith.index(STRIDE_TOKEN * 2)
+                                + head_idx * arith.index(HEAD_DIM * 2)
+                                + col_byte
+                            )
+                    voffset = arith.index_cast(T.i32, global_byte)
+
+                    rocdl.raw_ptr_buffer_load_lds(
+                        v_rsrc,
+                        lds_ptr,
+                        _dma_size,
+                        voffset,
+                        _dma_soff,
+                        _dma_off,
+                        _dma_aux,
+                    )
+
+        # ---- Preload Q^T B-operand packs once (register-resident) ----
+        # B operand uses j = lane_mod_32, k-subblock = lane_div_32*MFMA_LANE_K.
+        q_row = q_start + wave_q_offset + lane_mod_32
+        arith.index_cast(T.i32, q_row)
+        q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_v)
+        q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0))
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        q_b_packs = []
+        for ks in range_constexpr(K_STEPS_QK):
+            q_col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+            g_idx = global_idx(q_row_safe, q_col)
+            raw = load_global_mfma_pack(q_ptr, g_idx)
+            q_b_packs.append(arith.select(q_in_bounds, raw, c_zero_mfma_pack))
+
+        # ---- Constants ----
+        c_neg_inf = arith.constant(
+            -1.0e30, type=compute_type
+        )  # finite -inf to keep diff_m_raw=0 in all-masked tiles
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_one_f = arith.constant(1.0, type=compute_type)
+        c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type)
+        c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type)
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+        shuf_32_i32 = arith.constant(32, type=T.i32)
+        c4_i32 = arith.constant(4, type=T.i32)
+        lane_i32 = arith.index_cast(T.i32, lane)
+        lane_xor_32_i32 = arith.XOrIOp(lane_i32, shuf_32_i32).result
+        lane_xor_32_byte = arith.MulIOp(lane_xor_32_i32, c4_i32).result
+
+        def reduction_peer(v_f32):
+            if REDUCE_MODE == "ds_bpermute":
+                v_i32 = arith.ArithValue(v_f32).bitcast(T.i32)
+                peer_i32 = rocdl.ds_bpermute(T.i32, lane_xor_32_byte, v_i32)
+                return arith.ArithValue(peer_i32).bitcast(compute_type)
+            return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32)
+
+        # ---- SLA sparse outer loop: iterate over `topk` LUT entries ----
+        # N_SUBTILES == 1 by construction: each LUT entry covers exactly one
+        # BLOCK_N block. The inner `kv_sub` loop collapses to one iteration,
+        # and the dense kernel's intra-outer "next kv_sub" prefetch path is
+        # unused — we prefetch the next OUTER iteration (block_idx+1) instead.
+        assert N_SUBTILES == 1
+
+        # Loop-carried: [m_old, l_old, o_acc_chunks..., (buf_id if DMA dbuf)]
+        _use_dma_dbuf = ENABLE_DMA and not ENABLE_PREFETCH_3BUF and NUM_PREFETCH_K >= 2
+        init_args = [c_neg_inf, c_zero_f]
+        for _ in range_constexpr(D_CHUNKS):
+            init_args.append(c_zero_v16f32)
+        if _use_dma_dbuf:
+            init_args.append(arith.index(0))
+            # Prefetch the first SWA-window block (K AND V — Lever A V dbuf).
+            _init_kv_start = n_block_start * BN
+            coop_dma_k(_init_kv_start, buf_id=0)
+            coop_dma_v(_init_kv_start, buf_id=0)
+
+        for block_idx, inner_iter_args, loop_results in scf.for_(
+            n_block_start,
+            n_block_end,
+            _one_idx,
+            iter_args=init_args,
+        ):
+            m_running = inner_iter_args[0]
+            l_running = inner_iter_args[1]
+            o_accs = [inner_iter_args[2 + i] for i in range_constexpr(D_CHUNKS)]
+            _cur_buf_id = inner_iter_args[2 + D_CHUNKS] if _use_dma_dbuf else None
+
+            # V4 SWA: block_idx is directly the K-block index.
+            kv_block_start = block_idx * BN
+            preload_k_count = 1  # N_SUBTILES == 1
+
+            if ENABLE_PREFETCH_3BUF:
+                # 3-buf prefetch for sparse: look up LUT[block_idx + pre_k] and
+                # fire DMA per slot. Bounded by (block_idx + pre_k) < topk.
+                for pre_k in range_constexpr(preload_k_count):
+                    pre_k_slot = CK_LDS_SEQ[pre_k % len(CK_LDS_SEQ)] % NUM_PREFETCH_K
+                    if pre_k == 0:
+                        pre_k_start = kv_block_start
+                    else:
+                        _pre_idx = block_idx + arith.index(pre_k)
+                        _pre_has = arith.cmpi(arith.CmpIPredicate.slt, _pre_idx, n_block_end)
+                        _pre_if = scf.IfOp(_pre_has)
+                        with ir.InsertionPoint(_pre_if.then_block):
+                            pre_k_start = _pre_idx * BN
+                            if ENABLE_DMA:
+                                coop_dma_k(pre_k_start, pre_k_slot)
+                            else:
+                                coop_load_k(pre_k_start, pre_k_slot)
+                            scf.YieldOp([])
+                        continue
+                    if ENABLE_DMA:
+                        coop_dma_k(pre_k_start, pre_k_slot)
+                    else:
+                        coop_load_k(pre_k_start, pre_k_slot)
+                if ENABLE_DMA:
+                    rocdl.s_waitcnt(0)
+                else:
+                    rocdl.sched_group_barrier(rocdl.mask_vmem_rd, 1, 0)
+                gpu.barrier()
+
+            for kv_sub in range_constexpr(N_SUBTILES):  # single iteration
+                kv_start = kv_block_start  # sparse: kv_sub == 0 always
+
+                if ENABLE_PREFETCH_3BUF:
+                    k_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K
+                elif _use_dma_dbuf:
+                    _k_buf_id = _cur_buf_id
+                    rocdl.s_waitcnt(0)
+                    gpu.barrier()
+                    _next_k_buf_id = arith.index(1) - _k_buf_id
+                    _next_block_idx = block_idx + _one_idx
+                    _has_next = arith.cmpi(
+                        arith.CmpIPredicate.slt,
+                        _next_block_idx,
+                        n_block_end,
+                    )
+                    _if_dma = scf.IfOp(_has_next)
+                    with ir.InsertionPoint(_if_dma.then_block):
+                        _next_kv = _next_block_idx * BN
+                        coop_dma_k(_next_kv, _next_k_buf_id)
+                        # Lever A: also fire next V DMA into the same buf_id.
+                        coop_dma_v(_next_kv, _next_k_buf_id)
+                        scf.YieldOp([])
+                    rocdl.sched_barrier(0)
+                    k_base = k_buf_base(_k_buf_id)
+                elif ENABLE_DMA:
+                    # Single-buf DMA: fire the DMA, then wait + barrier inline.
+                    k_slot = 0
+                    coop_dma_k(kv_start, k_slot)
+                    rocdl.s_waitcnt(0)
+                    gpu.barrier()
+                else:
+                    k_slot = 0
+                    coop_load_k(kv_start, k_slot)
+                    gpu.barrier()
+                if not _use_dma_dbuf:
+                    k_base = k_buf_base(k_slot)
+
+                if not USE_HW_TR or (not ENABLE_DMA and not ENABLE_PREFETCH_3BUF):
+                    _v_vecs_prefetch = coop_load_v_global(kv_start)
+
+                # ==== GEMM1: bulk-read all K packs, then pipeline MFMAs ====
+                k_hi_offset = K_SUB_N * K_STRIDE
+                # XOR swizzle: col ^ ((row & 0x7) << 4) avoids LDS bank conflicts
+                k_swz_mask = (lane_mod_32 & arith.index(0x7)) << arith.index(4)
+
+                def _k_idx_lo(ks):
+                    col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                    return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+                def _k_idx_hi(ks):
+                    col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                    return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask)
+
+                _QK_PREFETCH_DEPTH = 2
+                k_packs_lo = [None] * K_STEPS_QK
+                k_packs_hi = [None] * K_STEPS_QK
+                for p in range_constexpr(_QK_PREFETCH_DEPTH):
+                    k_packs_lo[p] = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(p)])
+                    if N_HALVES == 2:
+                        k_packs_hi[p] = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(p)])
+
+                if ENABLE_DMA and not ENABLE_PREFETCH_3BUF and not _use_dma_dbuf:
+                    # Single-buf DMA path: fire V into buf 0 inside the iter.
+                    coop_dma_v(kv_start, 0)
+                    rocdl.sched_barrier(0)
+
+                s_acc_lo = c_zero_v16f32
+                s_acc_hi = c_zero_v16f32
+                for ks in range_constexpr(K_STEPS_QK):
+                    s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo)
+                    if N_HALVES == 2:
+                        s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi)
+                    if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK:
+                        k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                            mfma_pack_type, lds_kv, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)]
+                        )
+                        if N_HALVES == 2:
+                            k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                                mfma_pack_type, lds_kv, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)]
+                            )
+
+                # ==== Online softmax over BLOCK_N KV positions ====
+                s_raw_lo = []
+                s_raw_hi = []
+                for r in range_constexpr(16):
+                    s_raw_lo.append(vector.extract(s_acc_lo, static_position=[r], dynamic_position=[]))
+                    if N_HALVES == 2:
+                        s_raw_hi.append(vector.extract(s_acc_hi, static_position=[r], dynamic_position=[]))
+
+                # SWA causal + boundary mask, per element. For the MFMA
+                # 32x32 C-layout (BLOCK_M=32 rows, BLOCK_N=64 cols arranged
+                # as two 32-col halves):
+                #   row owner: q_row = q_start + wave_q_offset + lane_mod_32
+                #              (preloaded above as `q_row`)
+                #   col index in tile: lane_div_32*4 + (r//4)*8 + (r%4)
+                #   tile column 0..31 = lo half, 32..63 = hi half
+                # The kernel `kv_start = block_idx * BLOCK_N` is the LEFT edge
+                # of the lo half (kv_col_lo). kv_col_hi = kv_col_lo + 32.
+                # Mask conditions (set element to NEG_INF):
+                #   1) causal:   kv_col > q_row
+                #   2) SWA:      kv_col + W <= q_row     (window length W)
+                #   3) boundary: kv_col >= seq_len
+                # NEG_INF == -inf (sla_fwd c_neg_inf). The all-masked-tile
+                # case is safe because m_running=-inf and l_running=0; the
+                # subsequent (m, l) update is identity:
+                #   m_new = max(-inf, -inf) = -inf
+                #   corr = exp(0) by convention but actually NaN here -- so
+                #   we must avoid running the masked tile.
+                # Mitigation: n_block_start/n_block_end above already prune
+                # entire tiles outside the SWA window; the only tiles we run
+                # have AT LEAST one in-window element per warp row. The
+                # per-element mask handles the remaining partial overlap.
+                kv_start_i32 = arith.index_cast(T.i32, kv_start)
+                lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32)
+                seq_len_i32 = arith.index_cast(T.i32, seq_len_v)
+                q_row_i32_mask = arith.index_cast(T.i32, q_row)
+                w_i32 = arith.constant(swa_window, type=T.i32)
+                # Always mask. Tile-level gate omitted (cost negligible vs MFMA).
+                _bool_ty = ir.IntegerType.get_signless(1)
+                tile_needs_mask = arith.constant(1, type=_bool_ty)
+                _MASK_N_OUT = 16 * N_HALVES
+                _mask_if = scf.IfOp(tile_needs_mask, [T.f32] * _MASK_N_OUT, has_else=True)
+                with ir.InsertionPoint(_mask_if.then_block):
+                    _m_lo = []
+                    _m_hi = []
+                    for r in range_constexpr(16):
+                        r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32)
+                        lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result
+                        kv_col_lo = arith.AddIOp(
+                            arith.AddIOp(kv_start_i32, lane_off_i32).result, r_off_i32
+                        ).result
+                        # Boundary
+                        is_oob_lo = arith.cmpi(arith.CmpIPredicate.sge, kv_col_lo, seq_len_i32)
+                        # Causal: kv_col > q_row
+                        is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_lo, q_row_i32_mask)
+                        # SWA: kv_col + W <= q_row
+                        kv_plus_w_lo = arith.AddIOp(kv_col_lo, w_i32).result
+                        is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_lo, q_row_i32_mask)
+                        bad_lo = arith.OrIOp(arith.OrIOp(is_causal_lo, is_swa_lo).result, is_oob_lo).result
+                        _m_lo.append(arith.select(bad_lo, c_neg_inf, s_raw_lo[r]))
+                        if N_HALVES == 2:
+                            kv_col_hi = arith.AddIOp(kv_col_lo, arith.constant(K_SUB_N, type=T.i32)).result
+                            is_oob_hi = arith.cmpi(arith.CmpIPredicate.sge, kv_col_hi, seq_len_i32)
+                            is_causal_hi = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_hi, q_row_i32_mask)
+                            kv_plus_w_hi = arith.AddIOp(kv_col_hi, w_i32).result
+                            is_swa_hi = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_hi, q_row_i32_mask)
+                            bad_hi = arith.OrIOp(
+                                arith.OrIOp(is_causal_hi, is_swa_hi).result, is_oob_hi
+                            ).result
+                            _m_hi.append(arith.select(bad_hi, c_neg_inf, s_raw_hi[r]))
+                    scf.YieldOp(_m_lo + _m_hi)
+                with ir.InsertionPoint(_mask_if.else_block):
+                    scf.YieldOp(s_raw_lo + s_raw_hi)
+                s_raw_lo = [_mask_if.results[i] for i in range(16)]
+                if N_HALVES == 2:
+                    s_raw_hi = [_mask_if.results[16 + i] for i in range(16)]
+                else:
+                    s_raw_hi = []
+
+                _max_fm = {"fastmath": fm_fast}
+                local_max = s_raw_lo[0]
+                for r in range_constexpr(15):
+                    local_max = arith.MaxNumFOp(local_max, s_raw_lo[r + 1], **_max_fm).result
+                if N_HALVES == 2:
+                    for r in range_constexpr(16):
+                        local_max = arith.MaxNumFOp(local_max, s_raw_hi[r], **_max_fm).result
+                peer_max = reduction_peer(local_max)
+                row_max = arith.MaxNumFOp(local_max, peer_max, **_max_fm).result
+                m_new_raw = arith.MaxNumFOp(m_running, row_max, **_max_fm).result
+
+                diff_m_raw = arith.SubFOp(m_running, m_new_raw, fastmath=fm_fast).result
+                diff_m_scaled = arith.MulFOp(diff_m_raw, c_sm_scale_log2e, fastmath=fm_fast).result
+                corr = arith.ArithValue(diff_m_scaled).exp2(fastmath=fm_fast)
+
+                scaled_max = arith.MulFOp(c_sm_scale_log2e, m_new_raw, fastmath=fm_fast).result
+                neg_scaled_max = arith.SubFOp(c_zero_f, scaled_max, fastmath=fm_fast).result
+
+                p_vals_lo = []
+                p_vals_hi = []
+                local_sum = c_zero_f
+                for r in range_constexpr(16):
+                    diff_lo = math_dialect.fma(s_raw_lo[r], c_sm_scale_log2e, neg_scaled_max)
+                    p_lo = arith.ArithValue(diff_lo).exp2(fastmath=fm_fast)
+                    p_vals_lo.append(p_lo)
+                    local_sum = arith.AddFOp(local_sum, p_lo, fastmath=fm_fast).result
+                if N_HALVES == 2:
+                    for r in range_constexpr(16):
+                        diff_hi = math_dialect.fma(s_raw_hi[r], c_sm_scale_log2e, neg_scaled_max)
+                        p_hi = arith.ArithValue(diff_hi).exp2(fastmath=fm_fast)
+                        p_vals_hi.append(p_hi)
+                        local_sum = arith.AddFOp(local_sum, p_hi, fastmath=fm_fast).result
+
+                peer_sum = reduction_peer(local_sum)
+                tile_sum = arith.AddFOp(local_sum, peer_sum, fastmath=fm_fast).result
+                l_corr = arith.MulFOp(corr, l_running, fastmath=fm_fast).result
+                l_new = arith.AddFOp(l_corr, tile_sum, fastmath=fm_fast).result
+
+                # ==== Rescale O accumulators ====
+                # Lever B: defer per-dc rescale to inside the PV loop so all 16
+                # corr-multiplied o_accs are not live simultaneously. This shrinks
+                # the register live range and (we hope) reduces VGPR spill.
+                corr_vec = vector.broadcast(v16f32_type, corr)
+                if not USE_HW_TR:
+                    o_accs[0] = arith.MulFOp(o_accs[0], corr_vec, fastmath=fm_fast).result
+                # USE_HW_TR: rescale happens inline in the PV loop below.
+
+                if ENABLE_PREFETCH_3BUF and (kv_sub + preload_k_count) < N_SUBTILES:
+                    next_k_sub = kv_sub + preload_k_count
+                    next_k_start = kv_block_start + next_k_sub * BLOCK_N
+                    next_k_slot = CK_LDS_SEQ[next_k_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K
+                    if ENABLE_DMA:
+                        coop_dma_k(next_k_start, next_k_slot)
+                    else:
+                        coop_load_k(next_k_start, next_k_slot)
+
+                if ENABLE_PREFETCH_3BUF:
+                    v_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_V
+                    v_base = v_buf_base(v_slot)
+                    coop_load_v(kv_start, v_slot)
+                    rocdl.sched_group_barrier(rocdl.mask_dswr, 1, 0)
+                    gpu.barrier()
+                elif _use_dma_dbuf:
+                    # Lever A: V is in the same buf_id as K (dbuf). The
+                    # s_waitcnt(0) + barrier issued at the K branch already
+                    # waited for the CURRENT V tile (fired in PREVIOUS iter,
+                    # or as part of the pre-loop prefetch for iter 0).
+                    v_base = v_buf_base(_k_buf_id)
+                elif ENABLE_DMA:
+                    v_base = v_buf_base(0)
+                    rocdl.s_waitcnt(0)
+                    gpu.barrier()
+                else:
+                    v_slot = 0
+                    v_base = v_buf_base(v_slot)
+                    _waitcnt_vm_n(0)
+                    coop_store_v_lds(_v_vecs_prefetch, v_slot)
+                    rocdl.sched_group_barrier(rocdl.mask_dswr, 1, 0)
+                    gpu.barrier()
+
+                # ==== Build P packs for lo and hi halves ====
+                if dtype_str == "bf16" and not USE_K16:
+                    p_packs_lo = []
+                    p_packs_hi = []
+                    for pks in range_constexpr(PV_K_STEPS):
+                        p_base = pks * 4
+                        p_packs_lo.append(bf16_trunc_pack_v4(p_vals_lo[p_base : p_base + 4]))
+                        if N_HALVES == 2:
+                            p_packs_hi.append(bf16_trunc_pack_v4(p_vals_hi[p_base : p_base + 4]))
+                elif dtype_str == "bf16" and USE_K16:
+                    p_packs_lo = []
+                    p_packs_hi = []
+                    for pks in range_constexpr(PV_K_STEPS):
+                        p_base = pks * 8
+                        p_packs_lo.append(bf16_trunc_pack_v8(p_vals_lo[p_base : p_base + 8]))
+                        if N_HALVES == 2:
+                            p_packs_hi.append(bf16_trunc_pack_v8(p_vals_hi[p_base : p_base + 8]))
+                else:
+                    p_f16_lo = []
+                    p_f16_hi = []
+                    for r in range_constexpr(16):
+                        p_f16_lo.append(arith.trunc_f(elem_type, p_vals_lo[r]))
+                        if N_HALVES == 2:
+                            p_f16_hi.append(arith.trunc_f(elem_type, p_vals_hi[r]))
+
+                    if USE_K16:
+                        p_packs_lo = []
+                        p_packs_hi = []
+                        for pks in range_constexpr(PV_K_STEPS):
+                            p_base = pks * 8
+                            p_packs_lo.append(
+                                vector.from_elements(
+                                    v8f16_type,
+                                    [
+                                        p_f16_lo[p_base + 0],
+                                        p_f16_lo[p_base + 1],
+                                        p_f16_lo[p_base + 2],
+                                        p_f16_lo[p_base + 3],
+                                        p_f16_lo[p_base + 4],
+                                        p_f16_lo[p_base + 5],
+                                        p_f16_lo[p_base + 6],
+                                        p_f16_lo[p_base + 7],
+                                    ],
+                                )
+                            )
+                            if N_HALVES == 2:
+                                p_packs_hi.append(
+                                    vector.from_elements(
+                                        v8f16_type,
+                                        [
+                                            p_f16_hi[p_base + 0],
+                                            p_f16_hi[p_base + 1],
+                                            p_f16_hi[p_base + 2],
+                                            p_f16_hi[p_base + 3],
+                                            p_f16_hi[p_base + 4],
+                                            p_f16_hi[p_base + 5],
+                                            p_f16_hi[p_base + 6],
+                                            p_f16_hi[p_base + 7],
+                                        ],
+                                    )
+                                )
+                    else:
+                        p_packs_lo = []
+                        p_packs_hi = []
+                        for pks in range_constexpr(PV_K_STEPS):
+                            p_base = pks * 4
+                            p_packs_lo.append(
+                                vector.from_elements(
+                                    v4f16_type,
+                                    [
+                                        p_f16_lo[p_base],
+                                        p_f16_lo[p_base + 1],
+                                        p_f16_lo[p_base + 2],
+                                        p_f16_lo[p_base + 3],
+                                    ],
+                                )
+                            )
+                            if N_HALVES == 2:
+                                p_packs_hi.append(
+                                    vector.from_elements(
+                                        v4f16_type,
+                                        [
+                                            p_f16_hi[p_base],
+                                            p_f16_hi[p_base + 1],
+                                            p_f16_hi[p_base + 2],
+                                            p_f16_hi[p_base + 3],
+                                        ],
+                                    )
+                                )
+
+                # Build flat (dc, pks) schedule for interleaved GEMM2.
+                _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)]
+                TOTAL_PV = len(_steps)
+
+                def _read_v_pack(step_idx):
+                    dc, pks = _steps[step_idx]
+                    vh = None
+                    if USE_HW_TR:
+                        d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4
+                        k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group
+                        _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col
+                        lds_lo = v_base + k_row * V_STRIDE + _d_col_eff
+                        if N_HALVES == 2:
+                            lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE)
+                        if USE_K16:
+                            vl_a = ds_read_tr_v4f16(lds_lo)
+                            vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE))
+                            vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                            if N_HALVES == 2:
+                                vh_a = ds_read_tr_v4f16(lds_hi)
+                                vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE))
+                                vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                        else:
+                            vl = ds_read_tr_v4f16(lds_lo)
+                            if N_HALVES == 2:
+                                vh = ds_read_tr_v4f16(lds_hi)
+                    else:
+                        d_pos = arith.index(dc * D_CHUNK) + lane_mod_32
+                        k_base = arith.index(pks * PV_K_STEP) + lane_div_32 * 4
+                        v_lo_idx = v_base + d_pos * VT_STRIDE + k_base
+                        vl = vector.load(v4f16_type, lds_kv, [v_lo_idx])
+                        if N_HALVES == 2:
+                            v_hi_idx = v_lo_idx + arith.index(K_SUB_N)
+                            vh = vector.load(v4f16_type, lds_kv, [v_hi_idx])
+                    return vl, vh
+
+                # Pre-read V for the first step.
+                v_lo_cur, v_hi_cur = _read_v_pack(0)
+
+                # ==== GEMM2: O += V^T_lo @ P_lo (+ V^T_hi @ P_hi if N_HALVES==2) ====
+                for si in range_constexpr(TOTAL_PV):
+                    dc, pks = _steps[si]
+                    if si + 1 < TOTAL_PV:
+                        v_lo_nxt, v_hi_nxt = _read_v_pack(si + 1)
+                    # Lever B: per-dc inline rescale (USE_HW_TR path).
+                    # On the first PV pack for each dc, rescale o_accs[dc] just
+                    # before its first accumulate; keeps live range small.
+                    if USE_HW_TR and pks == 0:
+                        o_accs[dc] = arith.MulFOp(
+                            o_accs[dc],
+                            corr_vec,
+                            fastmath=fm_fast,
+                        ).result
+                    o_accs[dc] = mfma_acc(v_lo_cur, p_packs_lo[pks], o_accs[dc])
+                    if N_HALVES == 2:
+                        o_accs[dc] = mfma_acc(v_hi_cur, p_packs_hi[pks], o_accs[dc])
+                    if not USE_HW_TR and dc == 0 and pks < D_CHUNKS - 1:
+                        o_accs[pks + 1] = arith.MulFOp(
+                            o_accs[pks + 1],
+                            corr_vec,
+                            fastmath=fm_fast,
+                        ).result
+                    if si + 1 < TOTAL_PV:
+                        v_lo_cur = v_lo_nxt
+                        v_hi_cur = v_hi_nxt
+
+                m_running = m_new_raw
+                l_running = l_new
+
+            _yield_args = [m_running, l_running] + o_accs
+            if _use_dma_dbuf:
+                if N_SUBTILES % 2 == 1:
+                    _yield_args.append(arith.index(1) - _cur_buf_id)
+                else:
+                    _yield_args.append(_cur_buf_id)
+            yield _yield_args
+
+        # ---- Normalize and store O (skip OOB rows for partial Q tiles) ----
+        m_final = loop_results[0]
+        l_final = loop_results[1]
+        o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)]
+
+        inv_l = arith.DivFOp(
+            c_one_f,
+            l_final,
+            fastmath=fm_fast,
+        ).result
+        inv_l_vec = vector.broadcast(v16f32_type, inv_l)
+
+        # V4 LSE in raw-e scaled domain to match Triton ref:
+        #   lse = m_final * sm_scale + ln(l_final)
+        # (Triton stores m_i + tl.log(l_i) where m_i is in qk*sm_scale domain
+        #  already; our m_final is raw qk so we scale here.)
+        c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type)
+        scaled_m_final = arith.MulFOp(
+            m_final,
+            c_sm_scale_f,
+            fastmath=fm_fast,
+        ).result
+        ln_l_final = math_dialect.log(l_final, fastmath=fm_fast)
+        lse_val = arith.AddFOp(
+            scaled_m_final,
+            ln_l_final,
+            fastmath=fm_fast,
+        ).result
+
+        # O and LSE stores share the Q-row in-bounds guard. Note: the MFMA
+        # 32x32 register layout has the four rows held by this lane at offsets
+        # (0, 8, 16, 24) from the base (q_row = q_start + wave_q_offset +
+        # lane_mod_32). The O store already uses these offsets; for LSE we
+        # need to pick the ROW owner (lane_div_32 == 0 && lane_mod_32 < 32)
+        # and write one f32 per Q row. Simpler approach: every lane with
+        # lane_div_32 == 0 writes its lane_mod_32 row to LSE (32 rows per
+        # wave, one wave per Q tile row group). Matches the way `q_row` was
+        # computed for the Q preload at line 592.
+        _o_guard = scf.IfOp(q_in_bounds, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for dc in range_constexpr(D_CHUNKS):
+                o_norm_vec = arith.MulFOp(
+                    o_finals[dc],
+                    inv_l_vec,
+                    fastmath=fm_fast,
+                ).result
+                for r in range_constexpr(16):
+                    o_val = vector.extract(
+                        o_norm_vec,
+                        static_position=[r],
+                        dynamic_position=[],
+                    )
+                    o_f16 = arith.trunc_f(elem_type, o_val)
+
+                    d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4)
+                    d_col = arith.index(dc * D_CHUNK) + d_row_rel
+                    o_global = global_idx(q_row, d_col)
+                    _gep_store(o_f16, o_ptr, o_global)
+
+            # LSE store: one f32 per Q row. Wave-lane layout for 32x32 MFMA
+            # holds one row per lane in lane_div_32==0. Gate on lane_div_32
+            # == 0 to avoid 2x duplicate stores (the lane_div_32==1 copy has
+            # the same q_row and the same lse_val).
+            _is_row_owner = arith.cmpi(
+                arith.CmpIPredicate.eq,
+                lane_div_32,
+                arith.index(0),
+            )
+            _lse_if = scf.IfOp(_is_row_owner, [], has_else=False)
+            with ir.InsertionPoint(_lse_if.then_block):
+                # Flat LSE index: (batch*NUM_HEADS + head)*L + q_row
+                lse_off = (batch_idx * NUM_HEADS + head_idx) * seq_len_v + q_row
+                lse_off_i32 = arith.index_cast(T.i32, lse_off)
+                buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32)
+                scf.YieldOp([])
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_v4_swa_fwd(
+        Q: fx.Tensor,
+        K: fx.Tensor,
+        V: fx.Tensor,
+        O: fx.Tensor,
+        LSE: fx.Tensor,
+        batch_size: fx.Int32,
+        seq_len: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        bs_idx = arith.index_cast(T.index, batch_size)
+        sl_idx = arith.index_cast(T.index, seq_len)
+        num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M
+        grid_x = bs_idx * num_q_tiles * NUM_HEADS
+
+        launcher = v4_swa_fwd_kernel(Q, K, V, O, LSE, seq_len)
+
+        if waves_per_eu is not None:
+            _wpe = int(waves_per_eu)
+            if _wpe >= 1:
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(
+                            T.i32,
+                            _wpe,
+                        )
+        if flat_work_group_size is not None:
+            _fwgs = int(flat_work_group_size)
+            if _fwgs >= 1:
+                flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+                for op in ctx.gpu_module_body.operations:
+                    if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                        op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("no-nans-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("unsafe-fp-math"),
+                        ir.StringAttr.get("true"),
+                    ]
+                )
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(
+            grid=(grid_x, 1, 1),
+            block=(BLOCK_SIZE, 1, 1),
+            stream=stream,
+        )
+
+    # Best MI355X FMHA numbers so far were measured with ROCm/llvm-project
+    # `felix/tune_fmha` at c8cf6da4367c010c7cbbb7789a9c4349e7407619.
+    # Other LLVM revisions can compile/run this kernel, but usually leave a
+    # few percent of peak throughput on the table.
+    _fmha_compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {
+            "enable-post-misched": False,
+            "lsr-drop-solution": True,
+        },
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return launch_v4_swa_fwd(*args, **kwargs)
+
+    def _compile(Q, K, V, O, LSE, batch_size, seq_len, stream=None):
+        with CompilationContext.compile_hints(_fmha_compile_hints):
+            return flyc.compile(launch_v4_swa_fwd, Q, K, V, O, LSE, batch_size, seq_len, fx.Stream(stream))
+
+    _launch.compile = _compile
+
+    return _launch
+
+
+build_v4_swa_fwd_module_primary = build_v4_swa_fwd_module
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py
new file mode 100644
index 000000000..39fc8c1fd
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py
@@ -0,0 +1,36 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""FlyDSL-v2 DeepSeek-V4 sparse-MLA attention backend (MFMA, gfx950 / CDNA4).
+
+A fresh re-implementation that mirrors the gluon backend (``_gluon_dsa``): same
+fused single-latent (K == V) sparse-MLA representation and the **same public
+kernel-pair API**, so it plugs straight into the kernel-agnostic V4 adapter
+(:mod:`v4_sparse_mla_adapter`) with zero adapter changes:
+
+* ``q``  : ``[T, H, d_qk]``     bf16  (``d_qk = kv_lora_rank + rope_rank``)
+* ``kv`` : ``[T, 1, d_qk]``     bf16  (single MQA latent; ``V = K[:kv_lora_rank]``)
+* ``topk_indices`` : ``[T, TOPK]`` int32 (SWA window ++ sparse pool, -1 = invalid)
+* ``attn_sink`` : ``[H]`` fp32 optional per-head softmax sink
+
+Unlike the legacy wired FlyDSL CSA path (scalarized GEMV, no MFMA), the v1
+kernels use FlyDSL MFMA (``rocdl.mfma_*`` matrix cores) over a top-k gather.
+
+* :func:`sparse_mla_fwd_v4_flydsl` -> ``(o, lse)``   (native FlyDSL MFMA)
+* :func:`sparse_mla_bwd_v4_flydsl` -> ``(dq, dkv, d_sink)``   (native FlyDSL MFMA dQ
+  + shared Triton dKV intermediate/scatter-gather)
+
+Depends only on the installed ``flydsl`` pip package (no /workspace/FlyDSL-amd
+source tree required).
+"""
+
+from .dsa_bwd_v4_flydsl import sparse_mla_bwd_v4_flydsl
+from .dsa_fwd_v4_flydsl import sparse_mla_fwd_v4_flydsl
+
+__all__ = [
+    "sparse_mla_fwd_v4_flydsl",
+    "sparse_mla_bwd_v4_flydsl",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py
new file mode 100644
index 000000000..c05b862b3
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py
@@ -0,0 +1,595 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions adapted from FlyDSL Project Contributors (Apache-2.0); see
+# primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/.
+
+"""dsa_fwd_v4_flydsl_kernel: DeepSeek-V4 sparse-MLA attention forward (FlyDSL MFMA).
+
+Native FlyDSL MFMA forward for the fused single-latent (K == V) sparse-MLA form
+used by the V4 attention adapter. Same public contract as the gluon / triton_v2
+backends:
+
+  * ``q``   : ``[T, H, D_QK]`` bf16   (D_QK = kv_lora_rank + rope; rope is a zero pad)
+  * ``kv``  : ``[num_kv, 1, D_QK]`` bf16  (single MQA latent; V == K[:kv_lora_rank])
+  * ``topk``: ``[T, TOPK]`` int32   (SWA window ++ sparse pool; -1 = invalid)
+  * ``sink``: ``[H]`` fp32          (optional per-head softmax sink)
+  * out ``o``   : ``[T, H, kv_lora_rank]`` bf16
+  * out ``lse`` : ``[T, H]`` fp32   (sink-inclusive)
+
+Design (adapted from the in-tree v0 SWA flash kernel v4_sla_fwd_kernel.py):
+  * Grid: one workgroup per query TOKEN. The MFMA "M" axis is the HEAD axis
+    (BLOCK_H heads, one head-group = all H), the "N" axis is the gathered key
+    axis (BLOCK_N per tile), the contraction "K" axis is kv_lora_rank (=512).
+  * Each outer tile gathers BLOCK_N latent rows kv[topk[t, tile]] into LDS
+    (invalid topk == -1 -> row zeroed + column masked), runs the QK MFMA, an
+    online (flash) softmax over the key axis, and the PV MFMA (V == the same
+    latent, read transposed via ds_read_tr16_b64). K == V: the gathered tile is
+    written to a K-LDS region (XOR swizzled, for the QK read) and a V-LDS region
+    (row-major, for the transposed PV read).
+  * Epilogue: fold the per-head sink into the denominator (V4), normalize, write
+    O and sink-inclusive LSE = m*scale + ln(l).
+
+Numerics mirror v0: raw-domain running max, exp2(scale*log2e*(s - m)) softmax,
+bf16 truncation pack for P, finite NEG_INF (-1e30). One addition vs v0: a
+running-max clamp (m<=-1e29 -> 0) so a fully-masked leading tile (common for the
+SWA window of early tokens) contributes nothing instead of exp2(0)=1.
+
+gfx950 / CDNA4 only (USE_HW_TR + K16 MFMA). Non-DMA cooperative-gather path first
+(correctness); DMA / double-buffer / single-latent LDS sharing are follow-ups.
+"""
+
+import math
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+
+_LOG2E = math.log2(math.e)  # 1.4426950408889634
+_LLVM_GEP_DYNAMIC = -2147483648  # LLVM kDynamicIndex sentinel
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+def build_dsa_bwd_dq_module(
+    num_heads,
+    kv_lora_rank,
+    d_qk,
+    topk,
+    dtype_str="bf16",
+    sm_scale=None,
+    has_sink=True,
+    block_n=64,
+    block_h=None,
+    single_latent=False,
+    waves_per_eu=2,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+):
+    """Build the sparse-MLA forward launcher (single-variant, gfx950).
+
+    num_heads    : H (must be a multiple of 32).
+    kv_lora_rank : D_V = 512 (contraction of QK, output dim of PV).
+    d_qk         : row stride of q / kv (kv_lora_rank + rope pad, e.g. 576).
+    topk         : TOPK (padded to a multiple of block_n by the launcher).
+    """
+    gpu_arch = get_hip_arch()
+    assert gpu_arch.startswith("gfx950"), "dsa_fwd_v4_flydsl targets gfx950 (CDNA4)"
+    assert dtype_str == "bf16", "bf16 only"
+
+    HEAD_DIM = int(kv_lora_rank)  # D_V, the MFMA contraction / output dim
+    D_QK = int(d_qk)  # q / kv row stride (includes rope pad)
+    NUM_HEADS = int(num_heads)
+    TOPK = int(topk)
+    assert HEAD_DIM % 32 == 0 and HEAD_DIM >= 64
+    assert NUM_HEADS % 32 == 0, f"num_heads ({NUM_HEADS}) must be a multiple of 32"
+
+    BLOCK_N = int(block_n)
+    assert BLOCK_N % 32 == 0
+    assert TOPK % BLOCK_N == 0, f"TOPK ({TOPK}) must be a multiple of BLOCK_N ({BLOCK_N})"
+    K_SUB_N = 32
+    N_HALVES = BLOCK_N // 32  # halves of 32 columns per BLOCK_N
+
+    WARP_SIZE = 64
+    BLOCK_H = int(block_h) if block_h else NUM_HEADS  # heads per workgroup (M tile)
+    assert BLOCK_H % 32 == 0 and NUM_HEADS % BLOCK_H == 0
+    NUM_HEAD_GROUPS = NUM_HEADS // BLOCK_H
+    NUM_WAVES = BLOCK_H // 32
+    BLOCK_SIZE = NUM_WAVES * WARP_SIZE
+    ROWS_PER_WAVE = 32  # each wave owns 32 heads (MFMA M = 32)
+    NUM_TILES = TOPK // BLOCK_N
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM)
+
+    # ---- MFMA / K-step config (gfx950 CDNA4) ----
+    K_STEP_QK = 16
+    K_STEPS_QK = HEAD_DIM // K_STEP_QK  # 32 MFMA K-steps for the QK GEMM
+    D_CHUNK = 32
+    D_CHUNKS = HEAD_DIM // D_CHUNK  # 16 output chunks for the PV accumulator
+    PV_K_STEP = 16
+    PV_K_STEPS = K_SUB_N // PV_K_STEP  # 2 PV K-steps per 32-col sub-tile
+    MFMA_LANE_K = 8
+
+    # ---- LDS layout ----
+    # SINGLE_LATENT (flash / small H, LDS-occupancy-bound): one row-major tile
+    # (no swizzle, +4 pad) serves both QK (K) and PV (V via ds_read_tr). Halves LDS
+    # (occ 1->2). DUAL (pro / large H, VGPR-bound): separate XOR-swizzled K tile
+    # (conflict-free QK read) + row-major V tile; smaller LDS doesn't help pro
+    # occupancy but the swizzle avoids QK bank conflicts.
+    SINGLE_LATENT = bool(single_latent)
+    if SINGLE_LATENT:
+        K_STRIDE = HEAD_DIM + 4
+        V_STRIDE = HEAD_DIM + 4
+        LDS_V_BASE = 0
+        LDS_KV_ELEMS = BLOCK_N * (HEAD_DIM + 4)
+    else:
+        K_STRIDE = HEAD_DIM
+        V_STRIDE = HEAD_DIM + 4
+        LDS_V_BASE = BLOCK_N * K_STRIDE
+        LDS_KV_ELEMS = BLOCK_N * K_STRIDE + BLOCK_N * V_STRIDE
+
+    # ---- Cooperative gather-load decomposition ----
+    VEC_WIDTH = 16
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH  # 32 threads per gathered row
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+    assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+    NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"dsa_fwd_smem_H{BLOCK_H}_N{BLOCK_N}_K{TOPK}",
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    lds_valid_offset = allocator._align(lds_kv_offset + LDS_KV_ELEMS * 2, 16)  # f32 region (bytes)
+    allocator.ptr = lds_valid_offset + BLOCK_N * 4
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def dsa_bwd_dq_kernel(
+        Q: fx.Tensor,  # [T, H, D_QK] bf16 flat
+        KV: fx.Tensor,  # [num_kv, D_QK] bf16 flat (single latent)
+        dO: fx.Tensor,  # [T, H, HEAD_DIM] bf16 flat
+        TopK: fx.Tensor,  # [T, TOPK] int32 flat
+        LSE: fx.Tensor,  # [T, H] fp32 flat (sink-inclusive)
+        Delta: fx.Tensor,  # [T, H] fp32 flat (rowsum(O*dO))
+        dQ: fx.Tensor,  # [T, H, HEAD_DIM] bf16 flat (output; rope cols discarded)
+        dS: fx.Tensor,  # [T, H, TOPK] bf16 flat (output for dKV kernel)
+        Pout: fx.Tensor,  # [T, H, TOPK] bf16 flat (output for dKV kernel)
+    ):
+        elem_type = T.bf16
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        kv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), KV)
+        do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dO)
+        dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dQ)
+        ds_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dS)
+        pout_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Pout)
+        topk_rsrc = buffer_ops.create_buffer_resource(TopK, max_size=True)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+        delta_rsrc = buffer_ops.create_buffer_resource(Delta, max_size=True)
+
+        fm_fast = arith.FastMathFlags.fast
+        v4f16_type = T.vec(4, elem_type)
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        v16f32_type = T.vec(16, compute_type)
+        mfma_pack_type = v8f16_type
+
+        def mfma_acc(a, b, c):
+            # bf16, K16 (gfx950): mfma_f32_32x32x16_bf16(result_type, [a, b, c])
+            return rocdl.mfma_f32_32x32x16_bf16(v16f32_type, [a, b, c])
+
+        # ---- LDS view ----
+        base_ptr = allocator.get_base()
+        lds = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_ELEMS,)).get()
+        lds_valid = SmemPtr(base_ptr, lds_valid_offset, compute_type, shape=(BLOCK_N,)).get()
+
+        # ---- Thread / block indices ----
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        # block_id = token * NUM_HEAD_GROUPS + hg  (hg=0, hg_offset=0 when 1 group)
+        token = block_id // arith.index(NUM_HEAD_GROUPS)
+        hg_offset = (block_id % arith.index(NUM_HEAD_GROUPS)) * arith.index(BLOCK_H)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        lane_mod_32 = lane % 32
+        lane_div_32 = lane // 32  # 0/1
+
+        # ds_read_b64_tr_b16 lane decomposition (hardware 4x4 transpose)
+        tr_k_group = (lane % 16) // 4
+        tr_col_sub = lane % 4
+        tr_col_half = (lane % 32) // 16
+
+        wave_h_offset = wave_id * ROWS_PER_WAVE  # this wave's head-row base
+
+        # ---- ds_read_tr helper ----
+        def ds_read_tr_v4f16(lds_elem_idx):
+            byte_offset = lds_elem_idx * 2 + lds_kv_offset
+            byte_i64 = arith.index_cast(T.i64, byte_offset)
+            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result
+
+        # ---- global index helpers (token-major sparse-MLA) ----
+        HD_IDX = arith.index(HEAD_DIM)
+        DQK_IDX = arith.index(D_QK)
+        H_IDX = arith.index(NUM_HEADS)
+        TOPK_IDX = arith.index(TOPK)
+
+        def q_global_idx(head, col):
+            # q[token, head, col] ; row stride = H * D_QK, head stride = D_QK
+            return (token * H_IDX + head) * DQK_IDX + col
+
+        def kv_global_idx(kv_row, col):
+            # kv[kv_row, col] ; row stride = D_QK
+            return kv_row * DQK_IDX + col
+
+        def o_global_idx(head, col):
+            # O[token, head, col] ; row stride = H * HEAD_DIM (no rope)
+            return (token * H_IDX + head) * HD_IDX + col
+
+        def _gep_load(bptr, elem_idx, vec_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                bptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store(val, bptr, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                bptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(bptr, base_idx):
+            return _gep_load(bptr, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(bptr, base_idx):
+            return _gep_load(bptr, base_idx, vxf16_type)
+
+        def bf16_trunc_pack_v8(f32_vals):
+            _v4i32 = T.vec(4, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            pairs = []
+            for j in range_constexpr(4):
+                a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32)
+                b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32)
+                p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result
+                pairs.append(p)
+            return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs))
+
+        # ---- cooperative decomposition ----
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_one_f = arith.constant(1.0, type=compute_type)
+        c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type)
+        c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+        c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type)
+        c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type)
+
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+        shuf_32_i32 = arith.constant(32, type=T.i32)
+
+        def reduction_peer(v_f32):
+            return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32)
+
+        # ---- K XOR swizzle (col ^ ((row & 7) << 4)) ----
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(0x7)) << arith.index(4)
+            return col_idx ^ mask
+
+        def _swz_none(row_idx, col_idx):
+            return col_idx
+
+        # build-time layout selection (no traced `if`)
+        _swz_k = _swz_none if SINGLE_LATENT else _k_swizzle
+
+        def _store_row_single(lds_row, col, vec):
+            vector.store(vec, lds, [lds_row * K_STRIDE + col])
+
+        def _store_row_dual(lds_row, col, vec):
+            vector.store(vec, lds, [lds_row * K_STRIDE + _k_swizzle(lds_row, col)])
+            vector.store(vec, lds, [arith.index(LDS_V_BASE) + lds_row * V_STRIDE + col])
+
+        _store_row = _store_row_single if SINGLE_LATENT else _store_row_dual
+
+        # ---- Preload Q + dO B-operand packs and per-head lse/delta ----
+        # head row = hg_offset + wave_h_offset + lane_mod_32 (MFMA M axis)
+        head_row = hg_offset + wave_h_offset + lane_mod_32
+        head_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, head_row, H_IDX)
+        head_row_safe = arith.select(head_in_bounds, head_row, arith.index(0))
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        q_b_packs = []
+        do_b_packs = []
+        for ks in range_constexpr(K_STEPS_QK):
+            col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+            q_b_packs.append(
+                arith.select(
+                    head_in_bounds,
+                    load_global_mfma_pack(q_ptr, q_global_idx(head_row_safe, col)),
+                    c_zero_mfma_pack,
+                )
+            )
+            do_b_packs.append(
+                arith.select(
+                    head_in_bounds,
+                    load_global_mfma_pack(do_ptr, o_global_idx(head_row_safe, col)),
+                    c_zero_mfma_pack,
+                )
+            )
+        c_log2e = arith.constant(_LOG2E, type=compute_type)
+        head_flat_i32 = arith.index_cast(T.i32, token * H_IDX + head_row_safe)
+        lse_val = buffer_ops.buffer_load(lse_rsrc, head_flat_i32, vec_width=1, dtype=T.f32)
+        delta_val = buffer_ops.buffer_load(delta_rsrc, head_flat_i32, vec_width=1, dtype=T.f32)
+        neg_lse_log2e = arith.MulFOp(
+            lse_val, arith.SubFOp(c_zero_f, c_log2e, fastmath=fm_fast).result, fastmath=fm_fast
+        ).result
+
+        # ---- outer loop over TOPK tiles: dQ += dS @ K ----
+        init_args = []
+        for _ in range_constexpr(D_CHUNKS):
+            init_args.append(c_zero_v16f32)
+
+        for tile_idx, inner_iter_args, loop_results in scf.for_(
+            arith.index(0),
+            arith.index(NUM_TILES),
+            arith.index(1),
+            iter_args=init_args,
+        ):
+            dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)]
+
+            tile_topk_start = tile_idx * arith.index(BLOCK_N)
+
+            coop_gather_tile_dyn = tile_topk_start
+            # gather this tile
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                lds_row = load_row_in_batch + batch * ROWS_PER_BATCH_LOAD
+                topk_pos = coop_gather_tile_dyn + lds_row
+                topk_flat = token * TOPK_IDX + topk_pos
+                topk_flat_i32 = arith.index_cast(T.i32, topk_flat)
+                idx_raw = buffer_ops.buffer_load(topk_rsrc, topk_flat_i32, vec_width=1, dtype=T.i32)
+                valid = arith.cmpi(arith.CmpIPredicate.sge, idx_raw, arith.constant(0, type=T.i32))
+                safe_i32 = arith.select(valid, idx_raw, arith.constant(0, type=T.i32))
+                kv_row = arith.index_cast(T.index, safe_i32)
+                g_idx = kv_global_idx(kv_row, load_col_base)
+                vec_raw = load_global_f16xN(kv_ptr, g_idx)
+                vec = arith.select(valid, vec_raw, c_zero_vxf16)
+                _store_row(lds_row, load_col_base, vec)
+                is_col0 = arith.cmpi(arith.CmpIPredicate.eq, load_col_base, arith.index(0))
+                _if_c0 = scf.IfOp(is_col0)
+                with ir.InsertionPoint(_if_c0.then_block):
+                    mask_add = arith.select(valid, c_zero_f, c_neg_inf)
+                    vector.store(
+                        vector.from_elements(T.vec(1, compute_type), [mask_add]),
+                        lds_valid,
+                        [lds_row],
+                    )
+                    scf.YieldOp([])
+            gpu.barrier()
+
+            # ==== GEMM1: QK. bulk-read K packs, pipelined MFMA ====
+            k_hi_offset = K_SUB_N * K_STRIDE
+
+            def _k_idx_lo(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col)
+
+            def _k_idx_hi(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_hi_offset + lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col)
+
+            _QK_PREFETCH_DEPTH = 2
+            k_packs_lo = [None] * K_STEPS_QK
+            k_packs_hi = [None] * K_STEPS_QK
+            for p in range_constexpr(_QK_PREFETCH_DEPTH):
+                k_packs_lo[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_lo(p)])
+                if N_HALVES == 2:
+                    k_packs_hi[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_hi(p)])
+
+            s_acc_lo = c_zero_v16f32
+            s_acc_hi = c_zero_v16f32
+            dp_acc_lo = c_zero_v16f32
+            dp_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo)
+                dp_acc_lo = mfma_acc(k_packs_lo[ks], do_b_packs[ks], dp_acc_lo)
+                if N_HALVES == 2:
+                    s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi)
+                    dp_acc_hi = mfma_acc(k_packs_hi[ks], do_b_packs[ks], dp_acc_hi)
+                if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK:
+                    k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                        mfma_pack_type, lds, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)]
+                    )
+                    if N_HALVES == 2:
+                        k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                            mfma_pack_type, lds, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)]
+                        )
+
+            # ==== P = exp(scale*S - lse); dS = P*(dP - delta)*scale (per element) ====
+            # Column mapping (v0 32x32 C-layout): tile col = lane_div_32*4 +
+            # (r//4)*8 + (r%4) (+ K_SUB_N for the hi half).
+            lane_off = lane_div_32 * arith.index(4)
+
+            def _ds_packs(s_acc, dp_acc, half):
+                ds_vals = []
+                for r in range_constexpr(16):
+                    s_r = vector.extract(s_acc, static_position=[r], dynamic_position=[])
+                    dp_r = vector.extract(dp_acc, static_position=[r], dynamic_position=[])
+                    col = lane_off + arith.index((r % 4) + (r // 4) * 8 + half * K_SUB_N)
+                    mv = vector.extract(
+                        vector.load_op(T.vec(1, compute_type), lds_valid, [col]),
+                        static_position=[0],
+                        dynamic_position=[],
+                    )
+                    s_m = arith.AddFOp(s_r, mv, fastmath=fm_fast).result  # invalid -> -inf
+                    # P = exp2(scale*log2e*s_m - lse*log2e)
+                    p_arg = math_dialect.fma(s_m, c_sm_scale_log2e, neg_lse_log2e)
+                    p_r = arith.ArithValue(p_arg).exp2(fastmath=fm_fast)
+                    # dS = P * (dP - delta) * scale
+                    dp_md = arith.SubFOp(dp_r, delta_val, fastmath=fm_fast).result
+                    ds_r = arith.MulFOp(
+                        arith.MulFOp(p_r, dp_md, fastmath=fm_fast).result, c_sm_scale_f, fastmath=fm_fast
+                    ).result
+                    ds_vals.append(ds_r)
+                    # store dS and P to [T, H, TOPK] for the dKV kernel (kv_pos in [0,TOPK))
+                    kv_pos = tile_topk_start + col
+                    dsp_idx = (token * H_IDX + head_row) * TOPK_IDX + kv_pos
+                    _gep_store(arith.trunc_f(elem_type, ds_r), ds_ptr, dsp_idx)
+                    _gep_store(arith.trunc_f(elem_type, p_r), pout_ptr, dsp_idx)
+                packs = []
+                for pks in range_constexpr(PV_K_STEPS):
+                    packs.append(bf16_trunc_pack_v8(ds_vals[pks * 8 : pks * 8 + 8]))
+                return packs
+
+            ds_packs_lo = _ds_packs(s_acc_lo, dp_acc_lo, 0)
+            ds_packs_hi = _ds_packs(s_acc_hi, dp_acc_hi, 1) if N_HALVES == 2 else []
+
+            # ==== GEMM2: dQ += dS @ K  (K == V, ds_read_tr) ====
+            v_base = arith.index(LDS_V_BASE)
+            _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)]
+            TOTAL_PV = len(_steps)
+
+            def _read_v_pack(step_idx):
+                dc, pks = _steps[step_idx]
+                d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4
+                k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group
+                lds_lo = v_base + k_row * V_STRIDE + d_col
+                vl_a = ds_read_tr_v4f16(lds_lo)
+                vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE))
+                vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                vh = None
+                if N_HALVES == 2:
+                    lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE)
+                    vh_a = ds_read_tr_v4f16(lds_hi)
+                    vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE))
+                    vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                return vl, vh
+
+            for si in range_constexpr(TOTAL_PV):
+                dc, pks = _steps[si]
+                v_lo_cur, v_hi_cur = _read_v_pack(si)
+                dq_accs[dc] = mfma_acc(v_lo_cur, ds_packs_lo[pks], dq_accs[dc])
+                if N_HALVES == 2:
+                    dq_accs[dc] = mfma_acc(v_hi_cur, ds_packs_hi[pks], dq_accs[dc])
+
+            # protect this tile's V/K LDS reads from the next tile's gather writes
+            gpu.barrier()
+
+            yield dq_accs
+
+        # ---- epilogue: store dQ_lora [token, head, :HEAD_DIM] ----
+        dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)]
+        _o_guard = scf.IfOp(head_in_bounds, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for dc in range_constexpr(D_CHUNKS):
+                for r in range_constexpr(16):
+                    dq_val = vector.extract(dq_finals[dc], static_position=[r], dynamic_position=[])
+                    dq_f16 = arith.trunc_f(elem_type, dq_val)
+                    d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4)
+                    d_col = arith.index(dc * D_CHUNK) + d_row_rel
+                    # dQ is [T, H, d_qk] (matches q); write the :HEAD_DIM lora cols
+                    _gep_store(dq_f16, dq_ptr, q_global_idx(head_row, d_col))
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_dsa_bwd_dq(
+        Q: fx.Tensor,
+        KV: fx.Tensor,
+        dO: fx.Tensor,
+        TopK: fx.Tensor,
+        LSE: fx.Tensor,
+        Delta: fx.Tensor,
+        dQ: fx.Tensor,
+        dS: fx.Tensor,
+        Pout: fx.Tensor,
+        total_tokens: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        grid_x = arith.index_cast(T.index, total_tokens) * arith.index(NUM_HEAD_GROUPS)
+        launcher = dsa_bwd_dq_kernel(Q, KV, dO, TopK, LSE, Delta, dQ, dS, Pout)
+
+        if waves_per_eu is not None and int(waves_per_eu) >= 1:
+            _wpe = int(waves_per_eu)
+            for op in ctx.gpu_module_body.operations:
+                if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                    op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+
+        _fwgs = int(BLOCK_SIZE)
+        flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get([ir.StringAttr.get("no-nans-fp-math"), ir.StringAttr.get("true")])
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get([ir.StringAttr.get("unsafe-fp-math"), ir.StringAttr.get("true")])
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(grid=(grid_x, 1, 1), block=(BLOCK_SIZE, 1, 1), stream=stream)
+
+    _compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {"enable-post-misched": False, "lsr-drop-solution": True},
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_compile_hints):
+            return launch_dsa_bwd_dq(*args, **kwargs)
+
+    return _launch
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py
new file mode 100644
index 000000000..cf3ff37d8
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py
@@ -0,0 +1,184 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""FlyDSL-v1 sparse-MLA backward — native FlyDSL dQ + shared dKV gather.
+
+``sparse_mla_bwd_v4_flydsl(q, kv, o, do, topk, lse, ...) -> (dq, dkv, d_sink)``.
+
+The **dQ** kernel is native FlyDSL MFMA (``dsa_bwd_dq_flydsl_kernel``): it
+recomputes S = Q·Kᵀ, P = exp(scale·S − lse), dP = dO·Kᵀ, dS = P·(dP − Δ)·scale,
+accumulates dQ = dS·K, and writes the per-token dS/P buffers. The dKV path
+(intermediate GEMM + CSR inverted-topk scatter-reduce) reuses the shared,
+proven Triton kernels — the dKV gather is a variable-length scatter-reduction
+with no MFMA content, so there is nothing to gain from a FlyDSL rewrite there.
+
+This runs the whole top-k as a single chunk (R_CHUNK = TOPK), so dQ needs no
+cross-chunk read-modify-write and the dKV intermediate is built in one pass.
+
+Depends only on the installed ``flydsl`` pip package (no /workspace source).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+
+import torch
+import triton
+
+from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice
+from .._triton_v2.dsa_bwd_kernels import (
+    _bwd_compute_dkv_intermediate,
+    _bwd_dkv_gather_acc,
+)
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+
+from dsa_bwd_dq_flydsl_kernel import build_dsa_bwd_dq_module  # noqa: E402
+
+_DQ_CACHE = {}
+_DQ_LOCK = threading.Lock()
+_BLOCK_N = int(os.environ.get("PRIMUS_DSA_FLYDSL_BWD_BLOCK_N", "64"))
+_BLOCK_H = int(os.environ.get("PRIMUS_DSA_FLYDSL_BWD_BLOCK_H", "64"))
+
+
+def _get_dq_kernel(num_heads, kv_lora_rank, d_qk, topk, single_latent, scale):
+    block_h = min(_BLOCK_H, num_heads)
+    while num_heads % block_h != 0:
+        block_h -= 32
+    key = (num_heads, kv_lora_rank, d_qk, topk, block_h, single_latent, round(float(scale), 8))
+    with _DQ_LOCK:
+        launch = _DQ_CACHE.get(key)
+        if launch is None:
+            launch = build_dsa_bwd_dq_module(
+                num_heads=num_heads,
+                kv_lora_rank=kv_lora_rank,
+                d_qk=d_qk,
+                topk=topk,
+                dtype_str="bf16",
+                sm_scale=float(scale),
+                block_n=_BLOCK_N,
+                block_h=block_h,
+                single_latent=single_latent,
+            )
+            _DQ_CACHE[key] = launch
+        return launch
+
+
+def sparse_mla_bwd_v4_flydsl(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA backward: native FlyDSL dQ + Triton dKV gather."""
+    assert q.is_contiguous() and o.is_contiguous() and do.is_contiguous()
+    assert topk_indices.is_contiguous() and lse.is_contiguous()
+    total_tokens, num_heads, d_qk = q.shape
+    D = int(kv_lora_rank)
+    rope_rank = d_qk - D
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    assert kv.is_contiguous()
+    num_kv = kv.shape[0]
+    assert q.dtype == torch.bfloat16
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,)
+
+    # pad topk to a BLOCK_N multiple (-1 = invalid); usually already padded upstream
+    topk = topk_indices.shape[1]
+    if topk % _BLOCK_N != 0:
+        pad = ((topk + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N - topk
+        topk_p = torch.cat(
+            [topk_indices, torch.full((total_tokens, pad), -1, dtype=torch.int32, device=q.device)], dim=1
+        ).contiguous()
+    else:
+        topk_p = topk_indices
+    TOPK = topk_p.shape[1]
+
+    # Delta = rowsum(O * dO)  (o is [T,H,D])
+    delta = (o[:, :, :D].float() * do.float()).sum(-1).contiguous()
+    lse_c = lse.contiguous()
+
+    # ---- native FlyDSL dQ (whole top-k) + dS/P buffers ----
+    dq = torch.zeros(total_tokens, num_heads, d_qk, dtype=q.dtype, device=q.device)  # rope cols stay 0
+    chunk_dS = torch.empty(total_tokens, num_heads, TOPK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, TOPK, dtype=torch.bfloat16, device=q.device)
+    single_latent = num_heads <= 64
+    dq_launch = _get_dq_kernel(int(num_heads), D, int(d_qk), int(TOPK), single_latent, scale)
+    dq_launch(
+        q.reshape(-1),
+        kv.reshape(-1),
+        do.reshape(-1),
+        topk_p.reshape(-1),
+        lse_c.reshape(-1),
+        delta.reshape(-1),
+        dq.reshape(-1),
+        chunk_dS.reshape(-1),
+        chunk_P.reshape(-1),
+        int(total_tokens),
+    )
+
+    # ---- dKV: intermediate GEMM + CSR inverted-topk scatter-reduce (Triton) ----
+    # BH_DKV=32/TK_DKV=64 whole-top-k (R_CHUNK=TOPK) in one dKV-intermediate pass.
+    # This config's LDS stays < 160 KB in both the standalone and training Triton
+    # contexts (the wider 64/128 config overflows under the training runtime).
+    HAS_ROPE = False
+    BH_DKV, TK_DKV = 32, 64
+    num_hg_dkv = triton.cdiv(num_heads, BH_DKV)
+
+    interm = torch.empty(total_tokens, TOPK, d_qk, dtype=torch.bfloat16, device=q.device)
+    _bwd_compute_dkv_intermediate[(total_tokens,)](
+        q,
+        do,
+        chunk_dS,
+        chunk_P,
+        interm,
+        q.stride(0),
+        q.stride(1),
+        do.stride(0),
+        do.stride(1),
+        chunk_dS.stride(0),
+        chunk_dS.stride(1),
+        interm.stride(0),
+        interm.stride(1),
+        num_heads,
+        R_CHUNK=TOPK,
+        TILE_K=TK_DKV,
+        BLOCK_H=BH_DKV,
+        NUM_HG=num_hg_dkv,
+        D_V=D,
+        D_ROPE=rope_rank,
+        HAS_ROPE=HAS_ROPE,
+        num_warps=4,
+    )
+
+    dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device)
+    inv_ptr, inv_data = _build_inverted_topk_slice(topk_p, 0, TOPK, num_kv=num_kv)
+    _bwd_dkv_gather_acc[(num_kv,)](
+        interm,
+        inv_ptr,
+        inv_data,
+        dkv_acc,
+        interm.stride(1),
+        dkv_acc.stride(0),
+        D_V=D,
+        D_ROPE=rope_rank,
+        HAS_ROPE=HAS_ROPE,
+        num_warps=4,
+    )
+
+    d_sink = None
+    if has_sink:
+        d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0)
+
+    dkv = dkv_acc.to(kv.dtype).unsqueeze(1)
+    return dq, dkv, d_sink
+
+
+__all__ = ["sparse_mla_bwd_v4_flydsl"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py
new file mode 100644
index 000000000..8366214ce
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py
@@ -0,0 +1,151 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""FlyDSL-v1 sparse-MLA forward (native FlyDSL MFMA) launcher.
+
+Public API mirrors :func:`sparse_mla_fwd_v4_gluon` / ``sparse_mla_fwd_v4_triton``:
+
+    sparse_mla_fwd_v4_flydsl(q, kv, topk_indices, attn_sink=None,
+                             kv_lora_rank=512, scale=None) -> (o, lse)
+
+The heavy lifting is the native FlyDSL MFMA kernel in
+``dsa_fwd_v4_flydsl_kernel.build_dsa_fwd_module``; this module just marshals the
+tensors, caches the built launcher per (H, D, TOPK, has_sink, block_n), and
+launches it. See the kernel module docstring for the design.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+
+import torch
+
+# flydsl_v1 uses only the installed `flydsl` pip package + its own local kernel
+# modules — it does NOT need the /workspace/FlyDSL-amd source tree.
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+
+from dsa_fwd_v4_flydsl_kernel import build_dsa_fwd_module  # noqa: E402
+
+_KERNEL_CACHE = {}
+_KERNEL_CACHE_LOCK = threading.Lock()
+
+_DEFAULT_BLOCK_N = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_BLOCK_N", "0"))  # 0 = shape-conditional
+_DEFAULT_BLOCK_H = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_BLOCK_H", "256"))
+_DEFAULT_WPE = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_WPE", "2"))
+
+
+def _get_kernel(
+    num_heads, kv_lora_rank, d_qk, topk, has_sink, block_n, block_h, single_latent, waves_per_eu, scale
+):
+    block_h = min(block_h, num_heads)
+    while num_heads % block_h != 0:
+        block_h -= 32
+    key = (
+        num_heads,
+        kv_lora_rank,
+        d_qk,
+        topk,
+        has_sink,
+        block_n,
+        block_h,
+        single_latent,
+        waves_per_eu,
+        round(float(scale), 8),
+    )
+    with _KERNEL_CACHE_LOCK:
+        launch = _KERNEL_CACHE.get(key)
+        if launch is None:
+            launch = build_dsa_fwd_module(
+                num_heads=num_heads,
+                kv_lora_rank=kv_lora_rank,
+                d_qk=d_qk,
+                topk=topk,
+                dtype_str="bf16",
+                sm_scale=float(scale),
+                has_sink=has_sink,
+                block_n=block_n,
+                block_h=block_h,
+                single_latent=single_latent,
+                waves_per_eu=waves_per_eu,
+            )
+            _KERNEL_CACHE[key] = launch
+        return launch
+
+
+def sparse_mla_fwd_v4_flydsl(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA forward (native FlyDSL MFMA, gfx950)."""
+    assert q.is_contiguous() and topk_indices.is_contiguous()
+    total_tokens, num_heads, d_qk = q.shape
+    kv_lora_rank = int(kv_lora_rank)
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    assert kv.is_contiguous()
+    assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk
+    assert q.dtype == torch.bfloat16, "bf16 only"
+
+    # Shape-conditional tile: flash (H<=64) is LDS/occupancy-bound -> small tile
+    # (BLOCK_N=32) doubles occupancy (occ 1->2). pro (H>=128) is VGPR-bound
+    # (occupancy stuck at 1 regardless of LDS) -> a smaller tile only adds
+    # barrier/softmax overhead, so keep the larger BLOCK_N=64 (fewer tiles).
+    # flash (H<=64) is LDS/occupancy-bound → single-latent (one shared tile) halves
+    # LDS (occ 1→2) at BLOCK_N=64 (single-latent already halves LDS, so no need for
+    # a smaller tile — smaller tiles only add barriers). pro (H>=128) is VGPR-bound
+    # → dual swizzled tile (conflict-free QK), extra LDS is free (occ stuck at 1).
+    single_latent = num_heads <= 64
+    block_n = _DEFAULT_BLOCK_N if _DEFAULT_BLOCK_N else 64
+    topk = topk_indices.shape[1]
+    # Pad TOPK up to a multiple of block_n with -1 (masked) if needed.
+    if topk % block_n != 0:
+        pad = ((topk + block_n - 1) // block_n) * block_n - topk
+        topk_indices = torch.cat(
+            [topk_indices, torch.full((total_tokens, pad), -1, dtype=torch.int32, device=q.device)],
+            dim=1,
+        ).contiguous()
+        topk = topk_indices.shape[1]
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.is_contiguous() and attn_sink.dtype == torch.float32
+        assert attn_sink.shape == (num_heads,)
+        sink = attn_sink
+    else:
+        # kernel always folds the sink; -inf makes it a no-op (af=1, sink_e=0)
+        sink = torch.full((num_heads,), float("-inf"), dtype=torch.float32, device=q.device)
+
+    o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device)
+    lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    launch = _get_kernel(
+        int(num_heads),
+        kv_lora_rank,
+        int(d_qk),
+        int(topk),
+        has_sink,
+        block_n,
+        _DEFAULT_BLOCK_H,
+        single_latent,
+        _DEFAULT_WPE,
+        scale,
+    )
+    launch(
+        q.reshape(-1),
+        kv.reshape(-1),
+        topk_indices.reshape(-1),
+        sink.reshape(-1),
+        o.reshape(-1),
+        lse.reshape(-1),
+        int(total_tokens),
+    )
+    return o, lse
+
+
+__all__ = ["sparse_mla_fwd_v4_flydsl"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py
new file mode 100644
index 000000000..d2b90af86
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py
@@ -0,0 +1,653 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions adapted from FlyDSL Project Contributors (Apache-2.0); see
+# primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/.
+
+"""dsa_fwd_v4_flydsl_kernel: DeepSeek-V4 sparse-MLA attention forward (FlyDSL MFMA).
+
+Native FlyDSL MFMA forward for the fused single-latent (K == V) sparse-MLA form
+used by the V4 attention adapter. Same public contract as the gluon / triton_v2
+backends:
+
+  * ``q``   : ``[T, H, D_QK]`` bf16   (D_QK = kv_lora_rank + rope; rope is a zero pad)
+  * ``kv``  : ``[num_kv, 1, D_QK]`` bf16  (single MQA latent; V == K[:kv_lora_rank])
+  * ``topk``: ``[T, TOPK]`` int32   (SWA window ++ sparse pool; -1 = invalid)
+  * ``sink``: ``[H]`` fp32          (optional per-head softmax sink)
+  * out ``o``   : ``[T, H, kv_lora_rank]`` bf16
+  * out ``lse`` : ``[T, H]`` fp32   (sink-inclusive)
+
+Design (adapted from the in-tree v0 SWA flash kernel v4_sla_fwd_kernel.py):
+  * Grid: one workgroup per query TOKEN. The MFMA "M" axis is the HEAD axis
+    (BLOCK_H heads, one head-group = all H), the "N" axis is the gathered key
+    axis (BLOCK_N per tile), the contraction "K" axis is kv_lora_rank (=512).
+  * Each outer tile gathers BLOCK_N latent rows kv[topk[t, tile]] into LDS
+    (invalid topk == -1 -> row zeroed + column masked), runs the QK MFMA, an
+    online (flash) softmax over the key axis, and the PV MFMA (V == the same
+    latent, read transposed via ds_read_tr16_b64). K == V: the gathered tile is
+    written to a K-LDS region (XOR swizzled, for the QK read) and a V-LDS region
+    (row-major, for the transposed PV read).
+  * Epilogue: fold the per-head sink into the denominator (V4), normalize, write
+    O and sink-inclusive LSE = m*scale + ln(l).
+
+Numerics mirror v0: raw-domain running max, exp2(scale*log2e*(s - m)) softmax,
+bf16 truncation pack for P, finite NEG_INF (-1e30). One addition vs v0: a
+running-max clamp (m<=-1e29 -> 0) so a fully-masked leading tile (common for the
+SWA window of early tokens) contributes nothing instead of exp2(0)=1.
+
+gfx950 / CDNA4 only (USE_HW_TR + K16 MFMA). Non-DMA cooperative-gather path first
+(correctness); DMA / double-buffer / single-latent LDS sharing are follow-ups.
+"""
+
+import math
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import fly as _fly
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl._mlir.dialects import math as math_dialect
+from flydsl._mlir.dialects import scf
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch as get_hip_arch
+from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
+
+_LOG2E = math.log2(math.e)  # 1.4426950408889634
+_LLVM_GEP_DYNAMIC = -2147483648  # LLVM kDynamicIndex sentinel
+
+
+def _llvm_ptr_ty():
+    return ir.Type.parse("!llvm.ptr")
+
+
+def _llvm_lds_ptr_ty():
+    return ir.Type.parse("!llvm.ptr<3>")
+
+
+def build_dsa_fwd_module(
+    num_heads,
+    kv_lora_rank,
+    d_qk,
+    topk,
+    dtype_str="bf16",
+    sm_scale=None,
+    has_sink=True,
+    block_n=64,
+    block_h=None,
+    single_latent=False,
+    waves_per_eu=2,
+    unsafe_fp_math=True,
+    fast_fp_math=True,
+    daz=True,
+):
+    """Build the sparse-MLA forward launcher (single-variant, gfx950).
+
+    num_heads    : H (must be a multiple of 32).
+    kv_lora_rank : D_V = 512 (contraction of QK, output dim of PV).
+    d_qk         : row stride of q / kv (kv_lora_rank + rope pad, e.g. 576).
+    topk         : TOPK (padded to a multiple of block_n by the launcher).
+    """
+    gpu_arch = get_hip_arch()
+    assert gpu_arch.startswith("gfx950"), "dsa_fwd_v4_flydsl targets gfx950 (CDNA4)"
+    assert dtype_str == "bf16", "bf16 only"
+
+    HEAD_DIM = int(kv_lora_rank)  # D_V, the MFMA contraction / output dim
+    D_QK = int(d_qk)  # q / kv row stride (includes rope pad)
+    NUM_HEADS = int(num_heads)
+    TOPK = int(topk)
+    assert HEAD_DIM % 32 == 0 and HEAD_DIM >= 64
+    assert NUM_HEADS % 32 == 0, f"num_heads ({NUM_HEADS}) must be a multiple of 32"
+
+    BLOCK_N = int(block_n)
+    assert BLOCK_N % 32 == 0
+    assert TOPK % BLOCK_N == 0, f"TOPK ({TOPK}) must be a multiple of BLOCK_N ({BLOCK_N})"
+    K_SUB_N = 32
+    N_HALVES = BLOCK_N // 32  # halves of 32 columns per BLOCK_N
+
+    WARP_SIZE = 64
+    BLOCK_H = int(block_h) if block_h else NUM_HEADS  # heads per workgroup (M tile)
+    assert BLOCK_H % 32 == 0 and NUM_HEADS % BLOCK_H == 0
+    NUM_HEAD_GROUPS = NUM_HEADS // BLOCK_H
+    NUM_WAVES = BLOCK_H // 32
+    BLOCK_SIZE = NUM_WAVES * WARP_SIZE
+    ROWS_PER_WAVE = 32  # each wave owns 32 heads (MFMA M = 32)
+    NUM_TILES = TOPK // BLOCK_N
+
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM)
+
+    # ---- MFMA / K-step config (gfx950 CDNA4) ----
+    K_STEP_QK = 16
+    K_STEPS_QK = HEAD_DIM // K_STEP_QK  # 32 MFMA K-steps for the QK GEMM
+    D_CHUNK = 32
+    D_CHUNKS = HEAD_DIM // D_CHUNK  # 16 output chunks for the PV accumulator
+    PV_K_STEP = 16
+    PV_K_STEPS = K_SUB_N // PV_K_STEP  # 2 PV K-steps per 32-col sub-tile
+    MFMA_LANE_K = 8
+
+    # ---- LDS layout ----
+    # SINGLE_LATENT (flash / small H, LDS-occupancy-bound): one row-major tile
+    # (no swizzle, +4 pad) serves both QK (K) and PV (V via ds_read_tr). Halves LDS
+    # (occ 1->2). DUAL (pro / large H, VGPR-bound): separate XOR-swizzled K tile
+    # (conflict-free QK read) + row-major V tile; smaller LDS doesn't help pro
+    # occupancy but the swizzle avoids QK bank conflicts.
+    SINGLE_LATENT = bool(single_latent)
+    if SINGLE_LATENT:
+        K_STRIDE = HEAD_DIM + 4
+        V_STRIDE = HEAD_DIM + 4
+        LDS_V_BASE = 0
+        LDS_KV_ELEMS = BLOCK_N * (HEAD_DIM + 4)
+    else:
+        K_STRIDE = HEAD_DIM
+        V_STRIDE = HEAD_DIM + 4
+        LDS_V_BASE = BLOCK_N * K_STRIDE
+        LDS_KV_ELEMS = BLOCK_N * K_STRIDE + BLOCK_N * V_STRIDE
+
+    # ---- Cooperative gather-load decomposition ----
+    VEC_WIDTH = 16
+    assert HEAD_DIM % VEC_WIDTH == 0
+    THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH  # 32 threads per gathered row
+    assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0
+    ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD
+    assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0
+    NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD
+
+    allocator = SmemAllocator(
+        None,
+        arch=gpu_arch,
+        global_sym_name=f"dsa_fwd_smem_H{BLOCK_H}_N{BLOCK_N}_K{TOPK}",
+    )
+    lds_kv_offset = allocator._align(allocator.ptr, 16)
+    lds_valid_offset = allocator._align(lds_kv_offset + LDS_KV_ELEMS * 2, 16)  # f32 region (bytes)
+    allocator.ptr = lds_valid_offset + BLOCK_N * 4
+
+    @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
+    def dsa_fwd_kernel(
+        Q: fx.Tensor,  # [T, H, D_QK] bf16 flat
+        KV: fx.Tensor,  # [num_kv, D_QK] bf16 flat (single latent)
+        TopK: fx.Tensor,  # [T, TOPK] int32 flat
+        Sink: fx.Tensor,  # [H] fp32 flat
+        O: fx.Tensor,  # [T, H, HEAD_DIM] bf16 flat
+        LSE: fx.Tensor,  # [T, H] fp32 flat
+    ):
+        elem_type = T.bf16
+        compute_type = T.f32
+        q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q)
+        kv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), KV)
+        o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O)
+        topk_rsrc = buffer_ops.create_buffer_resource(TopK, max_size=True)
+        sink_rsrc = buffer_ops.create_buffer_resource(Sink, max_size=True)
+        lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True)
+
+        fm_fast = arith.FastMathFlags.fast
+        v4f16_type = T.vec(4, elem_type)
+        vxf16_type = T.vec(VEC_WIDTH, elem_type)
+        v8f16_type = T.vec(8, elem_type)
+        v16f32_type = T.vec(16, compute_type)
+        mfma_pack_type = v8f16_type
+
+        def mfma_acc(a, b, c):
+            # bf16, K16 (gfx950): mfma_f32_32x32x16_bf16(result_type, [a, b, c])
+            return rocdl.mfma_f32_32x32x16_bf16(v16f32_type, [a, b, c])
+
+        # ---- LDS view ----
+        base_ptr = allocator.get_base()
+        lds = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_ELEMS,)).get()
+        lds_valid = SmemPtr(base_ptr, lds_valid_offset, compute_type, shape=(BLOCK_N,)).get()
+
+        # ---- Thread / block indices ----
+        block_id = arith.index_cast(T.index, gpu.block_idx.x)
+        # block_id = token * NUM_HEAD_GROUPS + hg  (hg=0, hg_offset=0 when 1 group)
+        token = block_id // arith.index(NUM_HEAD_GROUPS)
+        hg_offset = (block_id % arith.index(NUM_HEAD_GROUPS)) * arith.index(BLOCK_H)
+        tid = arith.index_cast(T.index, gpu.thread_idx.x)
+
+        wave_id = tid // WARP_SIZE
+        lane = tid % WARP_SIZE
+        lane_mod_32 = lane % 32
+        lane_div_32 = lane // 32  # 0/1
+
+        # ds_read_b64_tr_b16 lane decomposition (hardware 4x4 transpose)
+        tr_k_group = (lane % 16) // 4
+        tr_col_sub = lane % 4
+        tr_col_half = (lane % 32) // 16
+
+        wave_h_offset = wave_id * ROWS_PER_WAVE  # this wave's head-row base
+
+        # ---- ds_read_tr helper ----
+        def ds_read_tr_v4f16(lds_elem_idx):
+            byte_offset = lds_elem_idx * 2 + lds_kv_offset
+            byte_i64 = arith.index_cast(T.i64, byte_offset)
+            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result
+
+        # ---- global index helpers (token-major sparse-MLA) ----
+        HD_IDX = arith.index(HEAD_DIM)
+        DQK_IDX = arith.index(D_QK)
+        H_IDX = arith.index(NUM_HEADS)
+        TOPK_IDX = arith.index(TOPK)
+
+        def q_global_idx(head, col):
+            # q[token, head, col] ; row stride = H * D_QK, head stride = D_QK
+            return (token * H_IDX + head) * DQK_IDX + col
+
+        def kv_global_idx(kv_row, col):
+            # kv[kv_row, col] ; row stride = D_QK
+            return kv_row * DQK_IDX + col
+
+        def o_global_idx(head, col):
+            # O[token, head, col] ; row stride = H * HEAD_DIM (no rope)
+            return (token * H_IDX + head) * HD_IDX + col
+
+        def _gep_load(bptr, elem_idx, vec_type):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                bptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            return _llvm.LoadOp(vec_type, gep.result).result
+
+        def _gep_store(val, bptr, elem_idx):
+            idx_i64 = arith.index_cast(T.i64, elem_idx)
+            gep = _llvm.GEPOp(
+                _llvm_ptr_ty(),
+                bptr,
+                [idx_i64],
+                rawConstantIndices=[_LLVM_GEP_DYNAMIC],
+                elem_type=elem_type,
+                noWrapFlags=0,
+            )
+            _llvm.StoreOp(val, gep.result)
+
+        def load_global_mfma_pack(bptr, base_idx):
+            return _gep_load(bptr, base_idx, mfma_pack_type)
+
+        def load_global_f16xN(bptr, base_idx):
+            return _gep_load(bptr, base_idx, vxf16_type)
+
+        def bf16_trunc_pack_v8(f32_vals):
+            _v4i32 = T.vec(4, T.i32)
+            _c16 = arith.constant(16, type=T.i32)
+            _cmask = arith.constant(0xFFFF0000, type=T.i32)
+            pairs = []
+            for j in range_constexpr(4):
+                a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32)
+                b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32)
+                p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result
+                pairs.append(p)
+            return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs))
+
+        # ---- cooperative decomposition ----
+        load_row_in_batch = tid // THREADS_PER_ROW_LOAD
+        load_lane_in_row = tid % THREADS_PER_ROW_LOAD
+        load_col_base = load_lane_in_row * VEC_WIDTH
+
+        c_neg_inf = arith.constant(-1.0e30, type=compute_type)
+        c_zero_f = arith.constant(0.0, type=compute_type)
+        c_one_f = arith.constant(1.0, type=compute_type)
+        c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type)
+        c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type)
+        c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type)
+        c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type)
+
+        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+        shuf_32_i32 = arith.constant(32, type=T.i32)
+
+        def reduction_peer(v_f32):
+            return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32)
+
+        # ---- K XOR swizzle (col ^ ((row & 7) << 4)) ----
+        def _k_swizzle(row_idx, col_idx):
+            mask = (row_idx & arith.index(0x7)) << arith.index(4)
+            return col_idx ^ mask
+
+        def _swz_none(row_idx, col_idx):
+            return col_idx
+
+        # build-time layout selection (no traced `if`)
+        _swz_k = _swz_none if SINGLE_LATENT else _k_swizzle
+
+        def _store_row_single(lds_row, col, vec):
+            vector.store(vec, lds, [lds_row * K_STRIDE + col])
+
+        def _store_row_dual(lds_row, col, vec):
+            vector.store(vec, lds, [lds_row * K_STRIDE + _k_swizzle(lds_row, col)])
+            vector.store(vec, lds, [arith.index(LDS_V_BASE) + lds_row * V_STRIDE + col])
+
+        _store_row = _store_row_single if SINGLE_LATENT else _store_row_dual
+
+        # ---- Preload Q B-operand packs (register-resident) ----
+        # head row = hg_offset + wave_h_offset + lane_mod_32 (MFMA M axis)
+        head_row = hg_offset + wave_h_offset + lane_mod_32
+        head_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, head_row, H_IDX)
+        head_row_safe = arith.select(head_in_bounds, head_row, arith.index(0))
+        c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type)
+        q_b_packs = []
+        for ks in range_constexpr(K_STEPS_QK):
+            q_col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+            g_idx = q_global_idx(head_row_safe, q_col)
+            raw = load_global_mfma_pack(q_ptr, g_idx)
+            q_b_packs.append(arith.select(head_in_bounds, raw, c_zero_mfma_pack))
+
+        # ---- outer loop over TOPK tiles ----
+        init_args = [c_neg_inf, c_zero_f]
+        for _ in range_constexpr(D_CHUNKS):
+            init_args.append(c_zero_v16f32)
+
+        for tile_idx, inner_iter_args, loop_results in scf.for_(
+            arith.index(0),
+            arith.index(NUM_TILES),
+            arith.index(1),
+            iter_args=init_args,
+        ):
+            m_running = inner_iter_args[0]
+            l_running = inner_iter_args[1]
+            o_accs = [inner_iter_args[2 + i] for i in range_constexpr(D_CHUNKS)]
+
+            tile_topk_start = tile_idx * arith.index(BLOCK_N)
+
+            coop_gather_tile_dyn = tile_topk_start
+            # gather this tile
+            for batch in range_constexpr(NUM_BATCHES_KV):
+                lds_row = load_row_in_batch + batch * ROWS_PER_BATCH_LOAD
+                topk_pos = coop_gather_tile_dyn + lds_row
+                topk_flat = token * TOPK_IDX + topk_pos
+                topk_flat_i32 = arith.index_cast(T.i32, topk_flat)
+                idx_raw = buffer_ops.buffer_load(topk_rsrc, topk_flat_i32, vec_width=1, dtype=T.i32)
+                valid = arith.cmpi(arith.CmpIPredicate.sge, idx_raw, arith.constant(0, type=T.i32))
+                safe_i32 = arith.select(valid, idx_raw, arith.constant(0, type=T.i32))
+                kv_row = arith.index_cast(T.index, safe_i32)
+                g_idx = kv_global_idx(kv_row, load_col_base)
+                vec_raw = load_global_f16xN(kv_ptr, g_idx)
+                vec = arith.select(valid, vec_raw, c_zero_vxf16)
+                _store_row(lds_row, load_col_base, vec)
+                is_col0 = arith.cmpi(arith.CmpIPredicate.eq, load_col_base, arith.index(0))
+                _if_c0 = scf.IfOp(is_col0)
+                with ir.InsertionPoint(_if_c0.then_block):
+                    mask_add = arith.select(valid, c_zero_f, c_neg_inf)
+                    vector.store(
+                        vector.from_elements(T.vec(1, compute_type), [mask_add]),
+                        lds_valid,
+                        [lds_row],
+                    )
+                    scf.YieldOp([])
+            gpu.barrier()
+
+            # ==== GEMM1: QK. bulk-read K packs, pipelined MFMA ====
+            k_hi_offset = K_SUB_N * K_STRIDE
+
+            def _k_idx_lo(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col)
+
+            def _k_idx_hi(ks):
+                col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K
+                return k_hi_offset + lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col)
+
+            _QK_PREFETCH_DEPTH = 2
+            k_packs_lo = [None] * K_STEPS_QK
+            k_packs_hi = [None] * K_STEPS_QK
+            for p in range_constexpr(_QK_PREFETCH_DEPTH):
+                k_packs_lo[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_lo(p)])
+                if N_HALVES == 2:
+                    k_packs_hi[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_hi(p)])
+
+            s_acc_lo = c_zero_v16f32
+            s_acc_hi = c_zero_v16f32
+            for ks in range_constexpr(K_STEPS_QK):
+                s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo)
+                if N_HALVES == 2:
+                    s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi)
+                if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK:
+                    k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                        mfma_pack_type, lds, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)]
+                    )
+                    if N_HALVES == 2:
+                        k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op(
+                            mfma_pack_type, lds, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)]
+                        )
+
+            # ==== Online softmax over BLOCK_N KV positions ====
+            s_raw_lo = []
+            s_raw_hi = []
+            for r in range_constexpr(16):
+                s_raw_lo.append(vector.extract(s_acc_lo, static_position=[r], dynamic_position=[]))
+                if N_HALVES == 2:
+                    s_raw_hi.append(vector.extract(s_acc_hi, static_position=[r], dynamic_position=[]))
+
+            # Add validity additive mask (from LDS). Column mapping (v0 32x32
+            # C-layout): tile col_lo = lane_div_32*4 + (r//4)*8 + (r%4); hi = +32.
+            lane_off = lane_div_32 * arith.index(4)
+            _m_lo = []
+            _m_hi = []
+            for r in range_constexpr(16):
+                r_off = arith.index((r % 4) + (r // 4) * 8)
+                col_lo = lane_off + r_off
+                mv_lo = vector.load_op(T.vec(1, compute_type), lds_valid, [col_lo])
+                mval_lo = vector.extract(mv_lo, static_position=[0], dynamic_position=[])
+                _m_lo.append(arith.AddFOp(s_raw_lo[r], mval_lo, fastmath=fm_fast).result)
+                if N_HALVES == 2:
+                    col_hi = col_lo + arith.index(K_SUB_N)
+                    mv_hi = vector.load_op(T.vec(1, compute_type), lds_valid, [col_hi])
+                    mval_hi = vector.extract(mv_hi, static_position=[0], dynamic_position=[])
+                    _m_hi.append(arith.AddFOp(s_raw_hi[r], mval_hi, fastmath=fm_fast).result)
+            s_raw_lo = _m_lo
+            s_raw_hi = _m_hi
+
+            _max_fm = {"fastmath": fm_fast}
+            local_max = s_raw_lo[0]
+            for r in range_constexpr(15):
+                local_max = arith.MaxNumFOp(local_max, s_raw_lo[r + 1], **_max_fm).result
+            if N_HALVES == 2:
+                for r in range_constexpr(16):
+                    local_max = arith.MaxNumFOp(local_max, s_raw_hi[r], **_max_fm).result
+            peer_max = reduction_peer(local_max)
+            row_max = arith.MaxNumFOp(local_max, peer_max, **_max_fm).result
+            m_new_raw = arith.MaxNumFOp(m_running, row_max, **_max_fm).result
+            # clamp: fully-masked-so-far -> 0 (placeholder; shift-invariant)
+            _finite = arith.cmpf(
+                arith.CmpFPredicate.OGT, m_new_raw, arith.constant(-1.0e29, type=compute_type)
+            )
+            m_new_raw = arith.select(_finite, m_new_raw, c_zero_f)
+
+            diff_m_raw = arith.SubFOp(m_running, m_new_raw, fastmath=fm_fast).result
+            diff_m_scaled = arith.MulFOp(diff_m_raw, c_sm_scale_log2e, fastmath=fm_fast).result
+            corr = arith.ArithValue(diff_m_scaled).exp2(fastmath=fm_fast)
+
+            scaled_max = arith.MulFOp(c_sm_scale_log2e, m_new_raw, fastmath=fm_fast).result
+            neg_scaled_max = arith.SubFOp(c_zero_f, scaled_max, fastmath=fm_fast).result
+
+            p_vals_lo = []
+            p_vals_hi = []
+            local_sum = c_zero_f
+            for r in range_constexpr(16):
+                diff_lo = math_dialect.fma(s_raw_lo[r], c_sm_scale_log2e, neg_scaled_max)
+                p_lo = arith.ArithValue(diff_lo).exp2(fastmath=fm_fast)
+                p_vals_lo.append(p_lo)
+                local_sum = arith.AddFOp(local_sum, p_lo, fastmath=fm_fast).result
+            if N_HALVES == 2:
+                for r in range_constexpr(16):
+                    diff_hi = math_dialect.fma(s_raw_hi[r], c_sm_scale_log2e, neg_scaled_max)
+                    p_hi = arith.ArithValue(diff_hi).exp2(fastmath=fm_fast)
+                    p_vals_hi.append(p_hi)
+                    local_sum = arith.AddFOp(local_sum, p_hi, fastmath=fm_fast).result
+
+            peer_sum = reduction_peer(local_sum)
+            tile_sum = arith.AddFOp(local_sum, peer_sum, fastmath=fm_fast).result
+            l_corr = arith.MulFOp(corr, l_running, fastmath=fm_fast).result
+            l_new = arith.AddFOp(l_corr, tile_sum, fastmath=fm_fast).result
+
+            corr_vec = vector.broadcast(v16f32_type, corr)
+            # online-softmax rescale of all O accumulators before this tile's PV
+            for dc in range_constexpr(D_CHUNKS):
+                o_accs[dc] = arith.MulFOp(o_accs[dc], corr_vec, fastmath=fm_fast).result
+
+            # ==== Build P packs (bf16 truncation) ====
+            p_packs_lo = []
+            p_packs_hi = []
+            for pks in range_constexpr(PV_K_STEPS):
+                p_base = pks * 8
+                p_packs_lo.append(bf16_trunc_pack_v8(p_vals_lo[p_base : p_base + 8]))
+                if N_HALVES == 2:
+                    p_packs_hi.append(bf16_trunc_pack_v8(p_vals_hi[p_base : p_base + 8]))
+
+            # ==== GEMM2: PV. read V transposed (ds_read_tr), interleaved ====
+            v_base = arith.index(LDS_V_BASE)
+            _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)]
+            TOTAL_PV = len(_steps)
+
+            def _read_v_pack(step_idx):
+                dc, pks = _steps[step_idx]
+                d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4
+                k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group
+                lds_lo = v_base + k_row * V_STRIDE + d_col
+                vl_a = ds_read_tr_v4f16(lds_lo)
+                vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE))
+                vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                vh = None
+                if N_HALVES == 2:
+                    lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE)
+                    vh_a = ds_read_tr_v4f16(lds_hi)
+                    vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE))
+                    vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7])
+                return vl, vh
+
+            for si in range_constexpr(TOTAL_PV):
+                dc, pks = _steps[si]
+                v_lo_cur, v_hi_cur = _read_v_pack(si)
+                o_accs[dc] = mfma_acc(v_lo_cur, p_packs_lo[pks], o_accs[dc])
+                if N_HALVES == 2:
+                    o_accs[dc] = mfma_acc(v_hi_cur, p_packs_hi[pks], o_accs[dc])
+
+            # protect this tile's V/K LDS reads from the next tile's gather writes
+            gpu.barrier()
+
+            m_running = m_new_raw
+            l_running = l_new
+
+            yield [m_running, l_running] + o_accs
+
+        # ---- epilogue: sink fold + normalize + store ----
+        m_final = loop_results[0]
+        l_final = loop_results[1]
+        o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)]
+
+        # scaled-domain sink fold (matches triton_v2 / gluon):
+        #   M = m_final * sm_scale ; l_final = sum exp(scaled_s - M)
+        # Always fold the per-head sink into the denominator (V4). For the
+        # no-sink case the launcher fills Sink with -inf, so af=1, sink_e=0 ->
+        # lse=M+log(l), acc_scale=1/l (identical to the no-sink formula). This
+        # avoids a Python `if` in the traced body (the AST rewriter does not
+        # propagate branch-local rebindings out of a dispatched if).
+        _log2e = arith.constant(_LOG2E, type=compute_type)
+        M_scaled = arith.MulFOp(m_final, c_sm_scale_f, fastmath=fm_fast).result
+        head_row_i32 = arith.index_cast(T.i32, head_row_safe)
+        sink_val = buffer_ops.buffer_load(sink_rsrc, head_row_i32, vec_width=1, dtype=T.f32)
+        m_fin = arith.MaxNumFOp(M_scaled, sink_val, fastmath=fm_fast).result
+        _daf = arith.MulFOp(
+            arith.SubFOp(M_scaled, m_fin, fastmath=fm_fast).result, _log2e, fastmath=fm_fast
+        ).result
+        af = arith.ArithValue(_daf).exp2(fastmath=fm_fast)
+        l_af = arith.MulFOp(l_final, af, fastmath=fm_fast).result
+        _dse = arith.MulFOp(
+            arith.SubFOp(sink_val, m_fin, fastmath=fm_fast).result, _log2e, fastmath=fm_fast
+        ).result
+        sink_e = arith.ArithValue(_dse).exp2(fastmath=fm_fast)
+        l_total = arith.AddFOp(l_af, sink_e, fastmath=fm_fast).result
+        ln_l = math_dialect.log(l_total, fastmath=fm_fast)
+        lse_val = arith.AddFOp(m_fin, ln_l, fastmath=fm_fast).result
+        inv_l = arith.DivFOp(c_one_f, l_total, fastmath=fm_fast).result
+        acc_scale = arith.MulFOp(af, inv_l, fastmath=fm_fast).result
+
+        acc_scale_vec = vector.broadcast(v16f32_type, acc_scale)
+
+        _o_guard = scf.IfOp(head_in_bounds, [], has_else=False)
+        with ir.InsertionPoint(_o_guard.then_block):
+            for dc in range_constexpr(D_CHUNKS):
+                o_norm_vec = arith.MulFOp(o_finals[dc], acc_scale_vec, fastmath=fm_fast).result
+                for r in range_constexpr(16):
+                    o_val = vector.extract(o_norm_vec, static_position=[r], dynamic_position=[])
+                    o_f16 = arith.trunc_f(elem_type, o_val)
+                    d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4)
+                    d_col = arith.index(dc * D_CHUNK) + d_row_rel
+                    _gep_store(o_f16, o_ptr, o_global_idx(head_row, d_col))
+
+            _is_row_owner = arith.cmpi(arith.CmpIPredicate.eq, lane_div_32, arith.index(0))
+            _lse_if = scf.IfOp(_is_row_owner, [], has_else=False)
+            with ir.InsertionPoint(_lse_if.then_block):
+                lse_off = token * H_IDX + head_row
+                lse_off_i32 = arith.index_cast(T.i32, lse_off)
+                buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32)
+                scf.YieldOp([])
+            scf.YieldOp([])
+
+    @flyc.jit
+    def launch_dsa_fwd(
+        Q: fx.Tensor,
+        KV: fx.Tensor,
+        TopK: fx.Tensor,
+        Sink: fx.Tensor,
+        O: fx.Tensor,
+        LSE: fx.Tensor,
+        total_tokens: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        allocator.finalized = False
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            allocator.finalize()
+
+        grid_x = arith.index_cast(T.index, total_tokens) * arith.index(NUM_HEAD_GROUPS)
+        launcher = dsa_fwd_kernel(Q, KV, TopK, Sink, O, LSE)
+
+        if waves_per_eu is not None and int(waves_per_eu) >= 1:
+            _wpe = int(waves_per_eu)
+            for op in ctx.gpu_module_body.operations:
+                if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                    op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe)
+
+        _fwgs = int(BLOCK_SIZE)
+        flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}")
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr
+
+        passthrough_entries = []
+        if daz:
+            passthrough_entries.append(
+                ir.ArrayAttr.get(
+                    [
+                        ir.StringAttr.get("denormal-fp-math-f32"),
+                        ir.StringAttr.get("preserve-sign,preserve-sign"),
+                    ]
+                )
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get([ir.StringAttr.get("no-nans-fp-math"), ir.StringAttr.get("true")])
+            )
+            passthrough_entries.append(
+                ir.ArrayAttr.get([ir.StringAttr.get("unsafe-fp-math"), ir.StringAttr.get("true")])
+            )
+        for op in ctx.gpu_module_body.operations:
+            if getattr(op, "OPERATION_NAME", None) == "gpu.func":
+                op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries)
+
+        launcher.launch(grid=(grid_x, 1, 1), block=(BLOCK_SIZE, 1, 1), stream=stream)
+
+    _compile_hints = {
+        "fast_fp_math": fast_fp_math,
+        "unsafe_fp_math": unsafe_fp_math,
+        "llvm_options": {"enable-post-misched": False, "lsr-drop-solution": True},
+    }
+
+    def _launch(*args, **kwargs):
+        with CompilationContext.compile_hints(_compile_hints):
+            return launch_dsa_fwd(*args, **kwargs)
+
+    return _launch
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py
new file mode 100644
index 000000000..6a8ac57da
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py
@@ -0,0 +1,40 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this package are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (``aiter/ops/triton/_gluon_kernels/gfx950``).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon (hardware-controlled Triton) DeepSeek-V4 sparse-MLA attention backend.
+
+Ported from ROCm/aiter PR #2922 (``aiter/ops/triton/_gluon_kernels/gfx950``)
+for gfx950 / CDNA4 (MI350/MI355X). These kernels operate on the **sparse-MLA
+latent** representation used by the DeepSeek V4 paper / FlashMLA:
+
+* ``q``  : ``[T, H, d_qk]``  with ``d_qk = kv_lora_rank (512) + rope_rank (64)``
+* ``kv`` : ``[T, 1, d_qk]``  single MQA latent (K and V share it; ``V_lora`` is
+           the first ``kv_lora_rank`` channels of ``K_lora``)
+* ``topk_indices`` : ``[T, TOPK]`` int32 absolute KV-token indices (SWA window +
+           sparse top-k already concatenated by the caller; ``-1`` = invalid)
+* ``attn_sink`` : ``[H]`` fp32 optional per-head learnable softmax sink
+
+This is a different (latent + per-token-topk) representation than the in-tree
+CSA path (``v4_csa_attention_v0``: ``q / k_local / v_local / gathered /
+sparse_mask``); it is exposed here as a standalone ``gluon`` backend.
+
+Public API mirrors aiter's ``sparse_mla_fwd_v4`` / ``sparse_mla_bwd_v4`` with
+``backend="gluon"``:
+
+* :func:`sparse_mla_fwd_v4_gluon` -> ``(o, lse)``
+* :func:`sparse_mla_bwd_v4_gluon` -> ``(dq, dkv, d_sink)``
+"""
+
+from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon
+from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon
+
+__all__ = [
+    "sparse_mla_fwd_v4_gluon",
+    "sparse_mla_bwd_v4_gluon",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py
new file mode 100644
index 000000000..1e0777b07
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py
@@ -0,0 +1,626 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+import torch
+import triton
+import triton.language as tl
+
+
+# =====================================================================
+# Backward method="gather" — intermediate storage + inverted topk gather
+# =====================================================================
+@triton.jit
+def _bwd_compute_dkv_intermediate(
+    Q_T_ptr,  # [T, D_QK, H] bf16  (q transposed: stride_qt_t = D_QK * H)
+    dO_T_ptr,  # [T, D_V,  H] bf16
+    dS_ptr,  # [T, H, TOPK] bf16
+    P_ptr,  # [T, H, TOPK] bf16
+    TopK_ptr,  # [T, TOPK] int32
+    Interm_ptr,  # [T, TOPK, D_QK] bf16 — output, one writer per (q, topk_rank)
+    stride_qt_t: tl.int64,
+    stride_dot_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_interm_t: tl.int64,  # TOPK * D_QK
+    stride_interm_k: tl.int64,  # D_QK
+    num_heads: tl.int32,
+    TOPK: tl.constexpr,
+    TILE_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NUM_HG: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+):
+    """
+    Same compute as _bwd_dkv_hg_fused but writes to a private intermediate
+    [T, TOPK, D] bf16 instead of atomic_add to shared dKV — no atomics needed.
+
+    Grid: (total_tokens,) -- one program per query token q.
+    For each tile of TOPK, stores dKV_lora/rope for that (q, tile) block.
+    """
+    token_idx = tl.program_id(0)
+
+    NUM_TILES: tl.constexpr = (TOPK + TILE_K - 1) // TILE_K
+    token_idx * stride_topk_t
+    offs_tile = tl.arange(0, TILE_K)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    interm_base_t = token_idx * stride_interm_t  # base for this query token
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        tile_offs = tile_start + offs_tile
+        valid = tile_offs < TOPK
+
+        dKV_lora = tl.zeros([D_V, TILE_K], dtype=tl.float32)
+        dKV_rope = tl.zeros([D_ROPE, TILE_K], dtype=tl.float32)
+
+        for hg in range(NUM_HG):
+            offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H)
+            mask_h = offs_h < num_heads
+
+            qt_base = token_idx * stride_qt_t
+            Q_lora_T = tl.load(
+                Q_T_ptr + qt_base + offs_v[:, None] * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )
+            Q_rope_T = tl.load(
+                Q_T_ptr + qt_base + (D_V + offs_r[:, None]) * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )
+
+            dot_base = token_idx * stride_dot_t
+            dO_T = tl.load(
+                dO_T_ptr + dot_base + offs_v[:, None] * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )
+
+            ds_base = token_idx * stride_ds_t
+            dS_val = tl.load(
+                dS_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :],
+                mask=mask_h[:, None] & valid[None, :],
+                other=0.0,
+            )
+            P_val = tl.load(
+                P_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :],
+                mask=mask_h[:, None] & valid[None, :],
+                other=0.0,
+            )
+
+            dKV_lora += tl.dot(Q_lora_T, dS_val.to(Q_lora_T.dtype)).to(tl.float32)
+            dKV_lora += tl.dot(dO_T, P_val.to(dO_T.dtype)).to(tl.float32)
+            dKV_rope += tl.dot(Q_rope_T, dS_val.to(Q_rope_T.dtype)).to(tl.float32)
+
+        # Store to intermediate: Interm[token_idx, tile_start:tile_start+TILE_K, 0:D_V]
+        # Layout: [T, TOPK, D] so pointer = interm_base_t + tile_offs[None,:]*D + offs_v[:,None]
+        interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + offs_v[:, None]
+        tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[None, :])
+
+        interm_rope_ptrs = (
+            Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + D_V + offs_r[:, None]
+        )
+        tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[None, :])
+
+
+@triton.jit
+def _bwd_dkv_gather(
+    Interm_ptr,  # [T, TOPK, D] bf16, flattened as [T*TOPK, D]
+    InvPtr_ptr,  # [T+1] int32 — CSR row pointers (kv_token -> range in inv_data)
+    InvData_ptr,  # [T*TOPK] int32 — encoded as q*TOPK+r, sorted by KV token
+    dKV_ptr,  # [T, D] bf16 — output
+    stride_interm_k: tl.int64,  # D_V + D_ROPE
+    stride_dkv_t: tl.int64,
+    TOPK: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+):
+    """
+    Gather dKV from intermediate buffer using CSR-style inverted topk index.
+
+    Grid: (total_tokens,) -- one CTA per KV token k.
+    Accumulates in fp32, stores bf16. No atomics.
+    """
+    k = tl.program_id(0)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    start = tl.load(InvPtr_ptr + k)
+    end = tl.load(InvPtr_ptr + k + 1)
+
+    dkv_acc_lora = tl.zeros([D_V], dtype=tl.float32)
+    dkv_acc_rope = tl.zeros([D_ROPE], dtype=tl.float32)
+
+    n_entries = end - start
+    for i in range(n_entries):
+        # entry = q*TOPK + r, used directly as flat index into [T*TOPK, D] intermediate
+        entry = tl.load(InvData_ptr + start + i).to(tl.int64)
+        base = entry * stride_interm_k
+        lora_val = tl.load(Interm_ptr + base + offs_v)
+        rope_val = tl.load(Interm_ptr + base + D_V + offs_r)
+        dkv_acc_lora += lora_val.to(tl.float32)
+        dkv_acc_rope += rope_val.to(tl.float32)
+
+    dkv_base = k.to(tl.int64) * stride_dkv_t
+    tl.store(dKV_ptr + dkv_base + offs_v, dkv_acc_lora.to(tl.bfloat16))
+    tl.store(dKV_ptr + dkv_base + D_V + offs_r, dkv_acc_rope.to(tl.bfloat16))
+
+
+def _build_inverted_topk(topk_indices):
+    """
+    Build CSR-style inverted index from topk_indices [T, TOPK] int32.
+
+    Returns:
+        inv_ptr:  [T+1] int32 — row pointers (kv_token -> range in inv_data)
+        inv_data: [T*TOPK] int32 — encoded (q*TOPK+r) values, sorted by KV token
+    """
+    T, TOPK = topk_indices.shape
+    device = topk_indices.device
+
+    flat_kv = topk_indices.reshape(-1).long()  # [T*TOPK] KV token indices
+    # argsort by KV token to get the flat indices (q*TOPK+r) in sorted order
+    order = torch.argsort(flat_kv, stable=True)
+    inv_data = order.to(torch.int32)  # [T*TOPK]
+
+    counts = torch.zeros(T, dtype=torch.int32, device=device)
+    counts.scatter_add_(0, flat_kv, torch.ones(T * TOPK, dtype=torch.int32, device=device))
+
+    inv_ptr = torch.zeros(T + 1, dtype=torch.int32, device=device)
+    torch.cumsum(counts, dim=0, out=inv_ptr[1:])
+
+    return inv_ptr, inv_data
+
+
+def _build_inverted_topk_slice(topk_indices_slice, r_start, R_CHUNK, num_kv=None):
+    """
+    Build CSR-style inverted index for a topk slice, excluding invalid (-1) entries.
+
+    Args:
+        topk_indices_slice: [T, R_CHUNK] int32 — topk_indices[:, r_start:r_start+R_CHUNK]
+          May contain -1 for padding (when actual chunk < R_CHUNK at the last chunk).
+        r_start:  int — first rank index in this slice (unused, for documentation)
+        R_CHUNK:  int — number of ranks in this slice (constexpr width)
+        num_kv:   int or None — number of KV tokens. When the KV buffer holds
+          MORE rows than query tokens (V4 [local ++ pool], num_kv = S + P), pass
+          it so ``inv_ptr`` has length ``num_kv + 1`` even if the last KV tokens
+          are referenced by no query. Defaults to ``T`` (query tokens).
+
+    Returns:
+        inv_ptr:  [num_kv+1] int32 — row pointers (kv_token -> range in inv_data)
+        inv_data: [valid_entries] int32 — flat indices q*R_CHUNK+local_r, sorted by KV token
+    """
+    T, RC = topk_indices_slice.shape
+    n_kv = T if num_kv is None else int(num_kv)
+    flat_kv = topk_indices_slice.reshape(-1).long()  # [T*R_CHUNK]; -1 marks invalid
+
+    # Sort ALL entries by KV token. argsort returns the flat positions (q*RC+r) in
+    # KV-token order; invalid (-1) entries sort to the FRONT and are never referenced
+    # because inv_ptr[0] starts past them. This avoids the boolean-mask `nonzero`
+    # + per-element `scatter_add` + extra `index` gathers of the previous version
+    # (all slow torch ops, especially on gfx1250 where they dominated the bwd).
+    inv_data = torch.argsort(flat_kv, stable=True).to(torch.int32)  # [T*R_CHUNK]
+    # bincount of (kv+1): bin 0 = #invalid, bins 1..num_kv = per-KV counts.
+    counts = torch.bincount(flat_kv + 1, minlength=n_kv + 1)  # [num_kv+1]
+    inv_ptr = torch.cumsum(counts, dim=0).to(torch.int32)  # [num_kv+1]; inv_ptr[0]=#invalid
+
+    return inv_ptr, inv_data
+
+
+# =====================================================================
+# Backward method="chunked_gather" — three kernels per chunk (no atomics)
+# =====================================================================
+@triton.jit
+def _bwd_chunk_dq_store_ds(
+    Q_ptr,  # [T, H, D] bf16
+    KV_ptr,  # [T, 1, D] bf16
+    dO_ptr,  # [T, H, D_V] bf16
+    TopK_ptr,  # [T, TOPK] int32
+    LSE_ptr,  # [T, H] fp32
+    Delta_ptr,  # [T, H] fp32
+    dQ_ptr,  # [T, H, D] bf16 — read-modify-write across chunks
+    dS_ptr,  # [T, H, R_CHUNK] bf16 — output chunk dS
+    P_ptr,  # [T, H, R_CHUNK] bf16 — output chunk P
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    TILE_K: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    IS_FIRST_CHUNK: tl.constexpr,
+):
+    """
+    dQ accumulation for rank chunk [R_START, R_START+R_CHUNK), plus stores
+    chunk dS and P to [T, H, R_CHUNK] buffers for use by _bwd_compute_dkv_intermediate.
+    Grid: (total_tokens, num_hg).
+    """
+    token_idx = tl.program_id(0)
+    hg_idx = tl.program_id(1)
+    offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offs_h < num_heads
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    q_base = token_idx * stride_q_t
+    Q_lora = tl.load(
+        Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    Q_rope = tl.load(
+        Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]),
+        mask=mask_h[:, None],
+        other=0.0,
+    )
+    do_base = token_idx * stride_do_t
+    dO_val = tl.load(
+        dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+    delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+
+    dq_base = token_idx * stride_dq_t
+    if IS_FIRST_CHUNK:
+        dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32)
+        dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32)
+    else:
+        dQ_lora = tl.load(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+            mask=mask_h[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        dQ_rope = tl.load(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+            mask=mask_h[:, None],
+            other=0.0,
+        ).to(tl.float32)
+
+    NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K
+    topk_base = token_idx * stride_topk_t + R_START
+    offs_tile = tl.arange(0, TILE_K)
+    ds_base = token_idx * stride_ds_t + hg_idx * BLOCK_H * stride_ds_h
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        tile_offs = tile_start + offs_tile
+        valid = tile_offs < R_CHUNK
+        topk_pos = tl.load(TopK_ptr + topk_base + tile_offs, mask=valid, other=-1)
+        valid = valid & (topk_pos != -1)
+        safe_pos = tl.where(valid, topk_pos, 0)
+
+        K_lora_T = tl.load(
+            KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0
+        )
+        K_rope_T = tl.load(
+            KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]), mask=valid[None, :], other=0.0
+        )
+
+        S = tl.dot(Q_lora, K_lora_T) + tl.dot(Q_rope, K_rope_T)
+        S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf"))
+        P = tl.exp(S - lse[:, None])
+        P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0)
+        dP = tl.dot(dO_val, K_lora_T)
+        dS = P * (dP - delta[:, None]) * scale
+        dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0)
+
+        dQ_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(K_lora_T)).to(tl.float32)
+        dQ_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(K_rope_T)).to(tl.float32)
+
+        # Store chunk dS and P for this tile — use local head offsets (0..BLOCK_H-1)
+        # since ds_base already encodes hg_idx*BLOCK_H*stride_ds_h
+        local_h = tl.arange(0, BLOCK_H)
+        tl.store(
+            dS_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :],
+            dS.to(tl.bfloat16),
+            mask=mask_h[:, None] & valid[None, :],
+        )
+        tl.store(
+            P_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :],
+            P.to(tl.bfloat16),
+            mask=mask_h[:, None] & valid[None, :],
+        )
+
+    tl.store(
+        dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+        dQ_lora.to(Q_lora.dtype),
+        mask=mask_h[:, None],
+    )
+    tl.store(
+        dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+        dQ_rope.to(Q_rope.dtype),
+        mask=mask_h[:, None],
+    )
+
+
+@triton.jit
+def _bwd_chunk_dq(
+    Q_ptr,  # [T, H, D] bf16
+    KV_ptr,  # [T, 1, D] bf16
+    dO_ptr,  # [T, H, D_V] bf16
+    TopK_ptr,  # [T, TOPK] int32
+    LSE_ptr,  # [T, H] fp32
+    Delta_ptr,  # [T, H] fp32
+    dQ_ptr,  # [T, H, D] bf16 — read-modify-write across chunks
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    TILE_K: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    IS_FIRST_CHUNK: tl.constexpr,
+):
+    """
+    dQ accumulation for one rank chunk [R_START, R_START+R_CHUNK).
+    Grid: (total_tokens, num_hg). No writes to intermediate buffer.
+    IS_FIRST_CHUNK=True: initialises dQ to zero (avoids a global memory read).
+    """
+    token_idx = tl.program_id(0)
+    hg_idx = tl.program_id(1)
+    offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offs_h < num_heads
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    q_base = token_idx * stride_q_t
+    Q_lora = tl.load(
+        Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    Q_rope = tl.load(
+        Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]),
+        mask=mask_h[:, None],
+        other=0.0,
+    )
+    do_base = token_idx * stride_do_t
+    dO_val = tl.load(
+        dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+    delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+
+    dq_base = token_idx * stride_dq_t
+    if IS_FIRST_CHUNK:
+        dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32)
+        dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32)
+    else:
+        dQ_lora = tl.load(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+            mask=mask_h[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        dQ_rope = tl.load(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+            mask=mask_h[:, None],
+            other=0.0,
+        ).to(tl.float32)
+
+    NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K
+    topk_base = token_idx * stride_topk_t + R_START
+    offs_tile = tl.arange(0, TILE_K)
+    topk_pos = tl.load(TopK_ptr + topk_base + offs_tile, mask=offs_tile < R_CHUNK, other=-1)
+    topk_pos_next = topk_pos
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        valid = (tile_start + offs_tile) < R_CHUNK
+        valid = valid & (topk_pos != -1)
+        if t + 1 < NUM_TILES:
+            next_offs = (t + 1) * TILE_K + offs_tile
+            topk_pos_next = tl.load(TopK_ptr + topk_base + next_offs, mask=next_offs < R_CHUNK, other=-1)
+        safe_pos = tl.where(valid, topk_pos, 0)
+
+        K_lora_T = tl.load(
+            KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0
+        )
+        K_rope_T = tl.load(
+            KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]), mask=valid[None, :], other=0.0
+        )
+
+        S = tl.dot(Q_lora, K_lora_T) + tl.dot(Q_rope, K_rope_T)
+        S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf"))
+        P = tl.exp(S - lse[:, None])
+        P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0)
+        dP = tl.dot(dO_val, K_lora_T)
+        dS = P * (dP - delta[:, None]) * scale
+        dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0)
+
+        dQ_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(K_lora_T)).to(tl.float32)
+        dQ_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(K_rope_T)).to(tl.float32)
+
+        if t + 1 < NUM_TILES:
+            topk_pos = topk_pos_next
+
+    tl.store(
+        dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+        dQ_lora.to(Q_lora.dtype),
+        mask=mask_h[:, None],
+    )
+    tl.store(
+        dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+        dQ_rope.to(Q_rope.dtype),
+        mask=mask_h[:, None],
+    )
+
+
+@triton.jit
+def _bwd_chunk_dkv_interm(
+    Q_T_ptr,  # [T, D_QK, H] bf16  (transposed: stride = D_QK * H)
+    dO_T_ptr,  # [T, D_V,  H] bf16
+    TopK_ptr,  # [T, TOPK] int32
+    LSE_ptr,  # [T, H] fp32
+    Delta_ptr,  # [T, H] fp32
+    KV_ptr,  # [T, 1, D] bf16
+    Interm_ptr,  # [T, R_CHUNK, D] bf16 — output (plain store, one writer per slot)
+    stride_qt_t: tl.int64,
+    stride_dot_t: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_interm_t: tl.int64,  # R_CHUNK * D
+    stride_interm_r: tl.int64,  # D
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: tl.constexpr,
+    TILE_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NUM_HG: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+):
+    """
+    dKV intermediate for one rank chunk [R_START, R_START+R_CHUNK).
+    Grid: (total_tokens,) — one CTA per query token, inner loop over head groups.
+    Recomputes S/P/dS on-the-fly. Plain stores to bf16 interm — no atomics.
+    """
+    token_idx = tl.program_id(0)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+    offs_tile = tl.arange(0, TILE_K)
+
+    NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K
+    topk_base = token_idx * stride_topk_t + R_START
+    interm_base_t = token_idx * stride_interm_t
+
+    qt_base = token_idx * stride_qt_t
+    dot_base = token_idx * stride_dot_t
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        tile_offs = tile_start + offs_tile
+        valid_tile = tile_offs < R_CHUNK
+
+        topk_pos = tl.load(TopK_ptr + topk_base + tile_start + offs_tile, mask=valid_tile, other=-1)
+        valid = valid_tile & (topk_pos != -1)
+        safe_pos = tl.where(valid, topk_pos, 0)
+
+        K_lora = tl.load(
+            KV_ptr + safe_pos[:, None] * stride_kv_t + offs_v[None, :], mask=valid[:, None], other=0.0
+        )  # [TILE_K, D_V]
+        K_rope = tl.load(
+            KV_ptr + safe_pos[:, None] * stride_kv_t + (D_V + offs_r[None, :]), mask=valid[:, None], other=0.0
+        )  # [TILE_K, D_ROPE]
+
+        dKV_lora = tl.zeros([TILE_K, D_V], dtype=tl.float32)
+        dKV_rope = tl.zeros([TILE_K, D_ROPE], dtype=tl.float32)
+
+        for hg in range(NUM_HG):
+            offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H)
+            mask_h = offs_h < num_heads
+
+            Q_lora_T = tl.load(
+                Q_T_ptr + qt_base + offs_v[:, None] * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )  # [D_V, BLOCK_H]
+            Q_rope_T = tl.load(
+                Q_T_ptr + qt_base + (D_V + offs_r[:, None]) * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )  # [D_ROPE, BLOCK_H]
+            dO_T = tl.load(
+                dO_T_ptr + dot_base + offs_v[:, None] * num_heads + offs_h[None, :],
+                mask=mask_h[None, :],
+                other=0.0,
+            )  # [D_V, BLOCK_H]
+            lse_h = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+            delta_h = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+
+            # S = K @ Q^T  [TILE_K, BLOCK_H]
+            S = tl.dot(K_lora, Q_lora_T) + tl.dot(K_rope, Q_rope_T)
+            S = tl.where(valid[:, None] & mask_h[None, :], S * scale, float("-inf"))
+            P = tl.exp(S - lse_h[None, :])
+            P = tl.where(valid[:, None] & mask_h[None, :], P, 0.0)
+            dP = tl.dot(K_lora, dO_T)  # [TILE_K, BLOCK_H]
+            dS = P * (dP - delta_h[None, :]) * scale
+            dS = tl.where(valid[:, None] & mask_h[None, :], dS, 0.0)
+
+            dKV_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(Q_lora_T)).to(tl.float32)
+            dKV_lora += tl.dot(P.to(tl.bfloat16), tl.trans(dO_T)).to(tl.float32)
+            dKV_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(Q_rope_T)).to(tl.float32)
+
+        # Plain store to bf16 interm — one writer per (token, local_r) slot
+        interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[:, None] * stride_interm_r + offs_v[None, :]
+        tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[:, None])
+
+        interm_rope_ptrs = (
+            Interm_ptr + interm_base_t + tile_offs[:, None] * stride_interm_r + (D_V + offs_r[None, :])
+        )
+        tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[:, None])
+
+
+@triton.jit
+def _bwd_dkv_gather_acc(
+    Interm_ptr,  # [T, R_CHUNK, D] bf16 — chunk intermediate
+    InvPtr_ptr,  # [T+1] int32 — CSR row pointers
+    InvData_ptr,  # [T*R_CHUNK] int32 — encoded as q*R_CHUNK + local_r
+    dKV_acc_ptr,  # [T, D] fp32 — accumulator (read-modify-write across chunks)
+    stride_interm_r: tl.int64,  # D
+    stride_acc_t: tl.int64,  # D
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    BLOCK_K: tl.constexpr = 64,
+):
+    """
+    Gather one chunk's bf16 intermediate into the fp32 dKV accumulator.
+    Grid: (total_tokens,) — one CTA per KV token k, no atomics.
+
+    Tiled over the CSR segment in BLOCK_K-row blocks: each iteration vector-gathers
+    BLOCK_K interm rows ([BLOCK_K, D]) and reduces with tl.sum, instead of the old
+    scalar `for i in range(n_entries)` one-row-at-a-time loop (~6.5x faster, gfx1250;
+    memory-bound 327 -> 2130 GB/s at BLOCK_K=64). RMW semantics unchanged.
+    """
+    k = tl.program_id(0)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    start = tl.load(InvPtr_ptr + k)
+    end = tl.load(InvPtr_ptr + k + 1)
+
+    acc_base = k.to(tl.int64) * stride_acc_t
+    dkv_acc_lora = tl.load(dKV_acc_ptr + acc_base + offs_v).to(tl.float32)
+    dkv_acc_rope = tl.load(dKV_acc_ptr + acc_base + D_V + offs_r).to(tl.float32)
+
+    n_entries = end - start
+    for ti in range(0, tl.cdiv(n_entries, BLOCK_K)):
+        e_local = ti * BLOCK_K + offs_k
+        valid = e_local < n_entries
+        entry = tl.load(InvData_ptr + start + e_local, mask=valid, other=0).to(tl.int64)
+        rp = entry[:, None] * stride_interm_r
+        rows_v = tl.load(Interm_ptr + rp + offs_v[None, :], mask=valid[:, None], other=0.0).to(tl.float32)
+        rows_r = tl.load(Interm_ptr + rp + D_V + offs_r[None, :], mask=valid[:, None], other=0.0).to(
+            tl.float32
+        )
+        dkv_acc_lora += tl.sum(rows_v, axis=0)
+        dkv_acc_rope += tl.sum(rows_r, axis=0)
+
+    tl.store(dKV_acc_ptr + acc_base + offs_v, dkv_acc_lora)
+    tl.store(dKV_acc_ptr + acc_base + D_V + offs_r, dkv_acc_rope)
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py
new file mode 100644
index 000000000..12ef192d4
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py
@@ -0,0 +1,59 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+import triton
+import triton.language as tl
+
+
+# =====================================================================
+# Backward — preprocess kernel (Delta computation)
+# =====================================================================
+@triton.jit
+def _sparse_mla_bwd_preprocess(
+    O_ptr,  # [total_tokens, num_heads, D_V]
+    dO_ptr,  # [total_tokens, num_heads, D_V]
+    Delta_ptr,  # [total_tokens, num_heads]
+    stride_o_t: tl.int64,
+    stride_o_h: tl.int64,
+    num_heads: tl.int32,
+    D_V: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """
+    Delta[t, h] = sum_d(O[t, h, d] * dO[t, h, d])
+
+    Grid: (total_tokens, cdiv(num_heads, BLOCK_H))
+    """
+    token_idx = tl.program_id(0)
+    hg_idx = tl.program_id(1)
+
+    offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offs_h < num_heads
+    offs_d = tl.arange(0, D_V)
+
+    base = token_idx * stride_o_t
+
+    O = tl.load(
+        O_ptr + base + offs_h[:, None] * stride_o_h + offs_d[None, :],
+        mask=mask_h[:, None],
+        other=0.0,
+    )
+    dO = tl.load(
+        dO_ptr + base + offs_h[:, None] * stride_o_h + offs_d[None, :],
+        mask=mask_h[:, None],
+        other=0.0,
+    )
+
+    delta = tl.sum(O.to(tl.float32) * dO.to(tl.float32), axis=1)
+
+    tl.store(
+        Delta_ptr + token_idx * num_heads + offs_h,
+        delta,
+        mask=mask_h,
+    )
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py
new file mode 100644
index 000000000..aa964a07b
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py
@@ -0,0 +1,238 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton:
+the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`);
+this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing
+the external transpose copy.
+
+Per program: 1 query token. Grid: (total_tokens,).
+Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups:
+    dKV_lora[D_V, TILE_K]   = sum_hg ( Q_lora_T @ dS + dO_T @ P )    # contract over heads
+    dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS )
+    store interm[token, rank, :D_QK]
+
+Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in
+LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout
+-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dkv_interm_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16 (UNtransposed)
+    dO_ptr,  # [T, H, D_V]  bf16 (UNtransposed)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    Interm_ptr,  # [T, R_CHUNK, D_QK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    stride_interm_t: tl.int64,
+    stride_interm_r: tl.int64,
+    num_heads: tl.int32,
+    R_CHUNK: gl.constexpr,
+    TILE_K: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    NUM_HG: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads.
+    mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for HBM loads ----
+    # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read.
+    _q_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _q_tpw_m: gl.constexpr = 64 // _q_tpw_k
+    blk_q: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_q_tpw_m, _q_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose).
+    blk_ds: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 4],
+        threads_per_warp=[16, 4],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ----
+    sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0])
+
+    # ---- Dot operand layouts ----
+    dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+    dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+
+    token_idx = gl.program_id(axis=0)
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ---- LDS for Q/dO/Q_rope (single-buffered) ----
+    smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q)
+    smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do)
+    smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    ds_base = token_idx.to(tl.int64) * stride_ds_t
+    interm_base = token_idx.to(tl.int64) * stride_interm_t
+
+    # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d]
+    offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma))
+    offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma))
+    offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma))
+
+    for t in range(NUM_TILES):
+        dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma)
+        dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma)
+
+        for hg in range(NUM_HG):
+            hg_off = hg * BLOCK_H
+
+            # ---- stage Q/dO/Q_rope (this head group) into LDS ----
+            offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q))
+            offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q))
+            mask_h_q = offs_h_q < num_heads
+            q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64)
+            do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+
+            offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+            offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+            mask_h_qr = offs_h_qr < num_heads
+            qr_offs = (
+                q_base
+                + offs_h_qr[:, None].to(tl.int64) * stride_q_h
+                + (D_V + offs_r_qr[None, :]).to(tl.int64)
+            )
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None]
+            )
+            gl.amd.cdna4.async_copy.commit_group()
+
+            # ---- load dS / P (this tile, this head group) -> dot operands ----
+            offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds))
+            offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds))
+            mask_h_ds = offs_h_ds < num_heads
+            dsp_offs = (
+                ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64)
+            )
+            dS_blk = gl.amd.cdna4.buffer_load(
+                ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            P_blk = gl.amd.cdna4.buffer_load(
+                ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            dS_dot = gl.convert_layout(dS_blk, dot_ds_b)
+            P_dot = gl.convert_layout(P_blk, dot_p_b)
+
+            # ---- wait + transpose-read Q/dO/Q_rope ----
+            gl.amd.cdna4.async_copy.wait_group(0)
+            Q_T = smem_q.permute([1, 0]).load(dot_qT_a)  # [D_V, BLOCK_H]
+            dO_T = smem_do.permute([1, 0]).load(dot_doT_a)  # [D_V, BLOCK_H]
+            Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a)  # [D_ROPE, BLOCK_H]
+
+            dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora)
+            dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora)
+            dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope)
+
+        # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ----
+        col = t * TILE_K + offs_col_st
+        interm_lora_offs = (
+            interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64)
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty),
+            ptr=Interm_ptr,
+            offsets=interm_lora_offs.to(tl.int32),
+        )
+        interm_rope_offs = (
+            interm_base
+            + col[None, :].to(tl.int64) * stride_interm_r
+            + (D_V + offs_dr_st[:, None]).to(tl.int64)
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty),
+            ptr=Interm_ptr,
+            offsets=interm_rope_offs.to(tl.int32),
+        )
+
+
+def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64):
+    """
+    Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel).
+
+    Returns interm [T, R_CHUNK, D_QK] bf16.
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    assert R_CHUNK % TILE_K == 0
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+        q,
+        do,
+        chunk_dS,
+        chunk_P,
+        interm,
+        q.stride(0),
+        q.stride(1),
+        do.stride(0),
+        do.stride(1),
+        chunk_dS.stride(0),
+        chunk_dS.stride(1),
+        interm.stride(0),
+        interm.stride(1),
+        num_heads,
+        R_CHUNK=R_CHUNK,
+        TILE_K=TILE_K,
+        BLOCK_H=BLOCK_H,
+        NUM_HG=num_hg,
+        D_V=kv_lora_rank,
+        D_ROPE=rope_rank,
+        num_warps=4,
+    )
+    return interm
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py
new file mode 100644
index 000000000..5b0bf2c46
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py
@@ -0,0 +1,511 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the
+gluon hardware-control structure of Leon's V3.2 forward.
+
+Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK).
+Grid: (total_tokens, cdiv(num_heads, BLOCK_H)).  Dispatch loops chunks externally,
+dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits).
+
+Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K):
+    S       = Q_lora @ K_lora_T + Q_rope @ K_rope_T      # 2 MMAs (mfma_s),  contract over D
+    P       = exp(S*scale - lse)                          # lse is sink-inclusive (from fwd)
+    dP      = dO @ K_lora_T                               # 1 MMA (mfma_s),  reuses K_lora_T_dot
+    dS      = P * (dP - delta) * scale
+    dQ_lora += dS @ K_lora                                # 1 MMA (mfma_acc), K_lora = K_lora_T.T view
+    dQ_rope += dS @ K_rope                                # 1 MMA (mfma_acc), K_rope = K_rope_T.T view
+    store dS, P chunk -> HBM (consumed by dKV-intermediate kernel)
+
+Differences vs Leon's fwd (the structural template):
+  * dO added as a second stationary [BH, D_V] operand (async-loaded with Q).
+  * No online softmax: single P = exp(S - lse) (lse precomputed by fwd).
+  * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways.
+  * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write.
+  * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher),
+    so this kernel needs no atomics. lse from fwd already folds the sink in.
+
+M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA
+k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950).
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dq_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16
+    KV_ptr,  # [T, 1, D_QK] bf16
+    dO_ptr,  # [T, H, D_V]  bf16
+    TopK_ptr,  # [T, TOPK_padded] int32
+    LSE_ptr,  # [T, H] fp32 (sink-inclusive)
+    Delta_ptr,  # [T, H] fp32
+    dQ_ptr,  # [T, H, D_QK] bf16 (read-modify-write across chunks)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    IS_FIRST_CHUNK: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for global loads (per Leon's fwd) ----
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_V]  (Q_lora, dO, dQ_lora)
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    # ---- Shared layouts (only K is staged in LDS) ----
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1])
+
+    # ---- Dot operand layouts ----
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+    dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ===================== program ids =====================
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ===================== Q / dO offsets =====================
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+    offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+    mask_h_qrope = offs_h_qrope < num_heads
+    q_offs_rope = (
+        q_base + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+    )
+    q_mask_rope = mask_h_qrope[:, None]
+
+    offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_do = offs_h_do < num_heads
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64)
+    do_mask = mask_h_do[:, None]
+
+    # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) =====================
+    # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to
+    # the dot-operand layout once. The convert's LDS scratch is transient (freed before
+    # the K loop), unlike persistent staging, so only K occupies LDS during the loop.
+    q_lora_blk = gl.amd.cdna4.buffer_load(
+        ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0
+    )
+    q_rope_blk = gl.amd.cdna4.buffer_load(
+        ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0
+    )
+    do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0)
+    Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a)
+    Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a)
+    dO_dot = gl.convert_layout(do_blk, dot_do_a)
+
+    # ===================== topk / KV offsets =====================
+    topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START
+
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # ===================== shared mem for K loop (double-buffered) =====================
+    smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope)
+    smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora)
+
+    # ===================== dQ accumulators =====================
+    # Always zero-init; the read-modify-write across chunks is folded in at STORE
+    # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout
+    # whose LDS scratch would overflow the K/Q/dO buffers).
+    dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+    dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc)
+
+    # ===================== lse / delta (in mfma_s row-slice layout) =====================
+    offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_s = offs_h_s < num_heads
+    lse = gl.amd.cdna4.buffer_load(
+        ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+    delta = gl.amd.cdna4.buffer_load(
+        ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+
+    # ===================== prologue: K tile 0 (group B, buffer 0) =====================
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=(topk_base + offs_tile_klora).to(tl.int32),
+        mask=offs_tile_klora < R_CHUNK,
+        other=-1,
+    )
+    topk_pos_krope = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=(topk_base + offs_tile_krope).to(tl.int32),
+        mask=offs_tile_krope < R_CHUNK,
+        other=-1,
+    )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1
+    )
+
+    valid_klora = topk_pos_klora != -1
+    valid_krope = topk_pos_krope != -1
+    valid_mma = topk_pos_mma != -1
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+    safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :]
+    )
+    krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :]
+    )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # dS / P store offsets in the mfma_s layout (store directly from the compute
+    # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead).
+    ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h
+    offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads
+
+    # ===================== main loop: prefetch t+1, compute t =====================
+    cur_buf = 0
+    for t in range(NUM_TILES - 1):
+        next_offs_klora = (t + 1) * TILE_K + offs_tile_klora
+        next_offs_krope = (t + 1) * TILE_K + offs_tile_krope
+        next_offs_mma = (t + 1) * TILE_K + offs_tile_mma
+
+        topk_pos_klora_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_klora).to(tl.int32),
+            mask=next_offs_klora < R_CHUNK,
+            other=-1,
+        )
+        topk_pos_krope_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_krope).to(tl.int32),
+            mask=next_offs_krope < R_CHUNK,
+            other=-1,
+        )
+        topk_pos_mma_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_mma).to(tl.int32),
+            mask=next_offs_mma < R_CHUNK,
+            other=-1,
+        )
+
+        valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1)
+        valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1)
+        valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1)
+        safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0)
+        safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0)
+
+        next_buf = 1 - cur_buf
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(next_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+            D_V + offs_r_krope[:, None]
+        ).to(tl.int64)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(next_buf),
+            ptr=KV_ptr,
+            offsets=krope_offs_next.to(tl.int32),
+            mask=valid_krope_next[None, :],
+        )
+        gl.amd.cdna4.async_copy.commit_group()
+
+        gl.amd.cdna4.async_copy.wait_group(1)
+
+        # ----- read K views from current buffer -----
+        klora_smem_cur = smem_klora.index(cur_buf)
+        K_lora_T_dot = klora_smem_cur.load(dot_klora_b)  # [D_V, TILE_K]  opIdx1 mfma_s
+        K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)  # [TILE_K, D_V] opIdx1 mfma_acc
+        krope_smem_cur = smem_krope.index(cur_buf)
+        K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+        K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+
+        # ----- S = Q_lora@K_lora_T + Q_rope@K_rope_T -----
+        S = gl.amd.cdna4.mfma(
+            Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+        S = S * scale
+        offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+        S = gl.where(valid_mask, S, float("-inf"))
+
+        # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale -----
+        P = gl.exp(S - lse[:, None])
+        P = gl.where(valid_mask, P, 0.0)
+        dP = gl.amd.cdna4.mfma(
+            dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        dS = P * (dP - delta[:, None]) * scale
+        dS = gl.where(valid_mask, dS, 0.0)
+
+        # ----- dQ_lora += dS@K_lora ; dQ_rope += dS@K_rope -----
+        dS_bf = dS.to(KV_ptr.dtype.element_ty)
+        dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+        dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+        dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+        # ----- store dS, P chunk -----
+        col = t * TILE_K + offs_tile_dsp
+        dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+        gl.amd.cdna4.buffer_store(
+            stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=P.to(KV_ptr.dtype.element_ty),
+            ptr=P_ptr,
+            offsets=dsp_offs.to(tl.int32),
+            mask=mask_h_dsp[:, None],
+        )
+
+        # promote prefetch -> current
+        cur_buf = next_buf
+        valid_mma = valid_mma_next
+
+    # ===================== epilogue: last tile =====================
+    gl.amd.cdna4.async_copy.wait_group(0)
+    t = NUM_TILES - 1
+    klora_smem_cur = smem_klora.index(cur_buf)
+    K_lora_T_dot = klora_smem_cur.load(dot_klora_b)
+    K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)
+    krope_smem_cur = smem_krope.index(cur_buf)
+    K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+    K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+
+    S = gl.amd.cdna4.mfma(
+        Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+    )
+    S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+    S = S * scale
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+    S = gl.where(valid_mask, S, float("-inf"))
+
+    P = gl.exp(S - lse[:, None])
+    P = gl.where(valid_mask, P, 0.0)
+    dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s))
+    dS = P * (dP - delta[:, None]) * scale
+    dS = gl.where(valid_mask, dS, 0.0)
+
+    dS_bf = dS.to(KV_ptr.dtype.element_ty)
+    dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+    dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+    dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+    col = t * TILE_K + offs_tile_dsp
+    dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+    )
+    gl.amd.cdna4.buffer_store(
+        stored_value=P.to(KV_ptr.dtype.element_ty),
+        ptr=P_ptr,
+        offsets=dsp_offs.to(tl.int32),
+        mask=mask_h_dsp[:, None],
+    )
+
+    # ===================== store dQ (lora + rope) =====================
+    dq_base = token_idx.to(tl.int64) * stride_dq_t
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64)
+    dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora)
+    if not IS_FIRST_CHUNK:
+        prev_lora = gl.amd.cdna4.buffer_load(
+            ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0
+        )
+        dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None]
+    )
+
+    offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+    offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+    mask_h_or = offs_h_or < num_heads
+    dq_offs_rope = (
+        dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64)
+    )
+    dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope)
+    if not IS_FIRST_CHUNK:
+        prev_rope = gl.amd.cdna4.buffer_load(
+            ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0
+        )
+        dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None]
+    )
+
+
+# =====================================================================
+# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P.
+# d_sink (if needed) is a torch reduction handled by the caller.
+# =====================================================================
+def sparse_mla_bwd_dq_gl(
+    q,
+    kv,
+    do,
+    topk_indices_padded,
+    lse,
+    delta,
+    R_CHUNK,
+    topk,
+    kv_lora_rank=512,
+    scale=None,
+    BLOCK_H=64,
+    TILE_K=16,
+):
+    """
+    Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop.
+
+    Returns:
+        dq:        [T, H, D_QK] bf16 (fully accumulated across chunks)
+        chunk_dS:  [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation)
+        chunk_P:   [T, H, R_CHUNK] bf16 (LAST chunk's P)
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK"
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    grid = (total_tokens, num_hg)
+
+    for r_start in range(0, topk, R_CHUNK):
+        is_first = r_start == 0
+        _sparse_mla_bwd_dq_gl_kernel[grid](
+            q,
+            kv,
+            do,
+            topk_indices_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_indices_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BLOCK_H,
+            TILE_K=TILE_K,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,
+        )
+    return dq, chunk_dS, chunk_P
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py
new file mode 100644
index 000000000..3c381aad6
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py
@@ -0,0 +1,175 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+End-to-end gluon backward for DeepSeek V4 sparse MLA (M4).
+
+Wires: Triton preprocess (Delta) + gluon dQ + gluon dKV-intermediate +
+Triton gather (CSR inverted-topk) + torch d_sink reduction.
+
+Drop-in for `sparse_mla_bwd_v4` (chunked_gather). Differences vs the Triton path:
+  - dQ kernel = gluon `_sparse_mla_bwd_dq_gl_kernel` (BH=64, TK=16)
+  - dKV-interm kernel = gluon `_sparse_mla_bwd_dkv_interm_gl_kernel` (BH=32, TK=64),
+    takes UNtransposed q/do -> the `q.transpose(1,2).contiguous()` copies are GONE.
+  - d_sink computed in torch (no in-kernel atomics).
+  - gather + preprocess unchanged (Triton).
+"""
+
+import torch
+import triton
+
+from ._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc
+from ._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess
+from .dsa_bwd_dkv_interm import _sparse_mla_bwd_dkv_interm_gl_kernel
+from .dsa_bwd_dq import _sparse_mla_bwd_dq_gl_kernel
+
+
+def sparse_mla_bwd_v4_gluon(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None):
+    assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous()
+    assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    # num_kv may exceed total_tokens (V4 [local ++ compressed-pool] buffer):
+    # dKV is accumulated per KV-token, so size it by kv rows, not query tokens.
+    num_kv = kv.shape[0]
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,)
+
+    # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ----
+    delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+    BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads))
+    _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))](
+        O_ptr=o,
+        dO_ptr=do,
+        Delta_ptr=delta,
+        stride_o_t=o.stride(0),
+        stride_o_h=o.stride(1),
+        num_heads=num_heads,
+        D_V=kv_lora_rank,
+        BLOCK_H=BLOCK_H_PRE,
+    )
+
+    # ---- config ----
+    R_CHUNK = min(256, topk)
+    BH_DQ, TK_DQ = 64, 16
+    BH_DKV, TK_DKV = 32, 64
+    num_hg_dq = triton.cdiv(num_heads, BH_DQ)
+    num_hg_dkv = triton.cdiv(num_heads, BH_DKV)
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    # ---- pad topk to R_CHUNK multiple ----
+    topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK
+    if topk_padded_len != topk:
+        pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device)
+        topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous()
+    else:
+        topk_padded = topk_indices
+
+    all_csr = [
+        _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv)
+        for rs in range(0, topk, R_CHUNK)
+    ]
+
+    for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)):
+        is_first = r_start == 0
+
+        # gluon dQ (writes dq RMW, chunk_dS, chunk_P)
+        _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)](
+            q,
+            kv,
+            do,
+            topk_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BH_DQ,
+            TILE_K=TK_DQ,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,  # sweep: 1 is best (+3-7%); >=2 spills catastrophically
+        )
+
+        # gluon dKV-intermediate (untransposed q/do -> no external transpose)
+        _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+            q,
+            do,
+            chunk_dS,
+            chunk_P,
+            interm,
+            q.stride(0),
+            q.stride(1),
+            do.stride(0),
+            do.stride(1),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            interm.stride(0),
+            interm.stride(1),
+            num_heads,
+            R_CHUNK=R_CHUNK,
+            TILE_K=TK_DKV,
+            BLOCK_H=BH_DKV,
+            NUM_HG=num_hg_dkv,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            num_warps=4,
+        )
+
+        # Triton gather (CSR inverted-topk reduce interm -> dkv_acc).
+        # Grid is over KV tokens (num_kv), which may exceed query tokens.
+        inv_ptr, inv_data = all_csr[chunk_idx]
+        _bwd_dkv_gather_acc[(num_kv,)](
+            interm,
+            inv_ptr,
+            inv_data,
+            dkv_acc,
+            interm.stride(1),
+            dkv_acc.stride(0),
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            num_warps=4,
+        )
+
+    # ---- d_sink in torch: -sum_t exp(sink - lse) * delta ----
+    d_sink = None
+    if has_sink:
+        d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0)
+
+    dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1)
+    return dq, dkv_out, d_sink
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py
new file mode 100644
index 000000000..8a956f513
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py
@@ -0,0 +1,634 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), PR #2922 (aiter/ops/triton/_gluon_kernels/gfx950).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink.
+
+Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch
+`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware
+control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline,
+ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456.
+This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches
+`sparse_mla_fwd_v4`, and is exposed via `sparse_mla_fwd_v4(..., backend="gluon")`.
+
+Pipeline (Leon's):
+  Prologue:  Q_lora + Q_rope -> shared via async DMA (group A).
+             K_lora tile 0 + K_rope tile 0 -> shared via async DMA (group B, double-buffered).
+             wait_group(1) -> Q in shared; load Q dot operands once.
+  Loop tile t (0..N-2): prefetch K[t+1]; wait_group(1); read K_lora/V_lora(permute)/K_rope;
+             S = Q_lora @ K_lora_T + Q_rope @ K_rope_T; online softmax; acc += P @ V_lora.
+  Epilogue:  wait_group(0) -> last tile; fold sink into the denominator (V4); write O, LSE.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+# =====================================================================
+# Utility
+# =====================================================================
+def _get_lds_limit():
+    """Return the per-CU LDS limit in bytes for the current GPU.
+
+    gfx942 (MI300X): 64 KB = 65536 bytes
+    gfx950 (MI355X): 160 KB = 163840 bytes
+    """
+    if torch.cuda.is_available():
+        prop = torch.cuda.get_device_properties(0)
+        gcn_arch = getattr(prop, "gcnArchName", "")
+        if "gfx950" in gcn_arch:
+            return 163840
+    return 65536
+
+
+_LDS_LIMIT = _get_lds_limit()
+
+
+# =====================================================================
+# Forward — autotune configs and pruning
+# =====================================================================
+def _fwd_prune_configs(configs, named_args, **kwargs):
+    """Prune autotune configs that would exceed per-CU LDS."""
+    D_V = kwargs.get("D_V", named_args.get("D_V"))
+    D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE"))
+    pruned = []
+    for config in configs:
+        config.kwargs["BLOCK_H"]
+        tk = config.kwargs["TILE_K"]
+        ns = config.num_stages
+        kv_lds = (D_V + D_ROPE) * tk * 2 * ns
+        if kv_lds <= _LDS_LIMIT:
+            pruned.append(config)
+    if not pruned:
+        pruned.append(configs[0])
+    return pruned
+
+
+def _get_fwd_autotune_configs():
+    configs = [
+        triton.Config(
+            {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE},
+            num_warps=nw,
+        )
+        for BLOCK_H in [16, 32, 64]
+        for TILE_K in [16, 32, 64, 128]
+        for WPE in [0, 1, 2]
+        for nw in [4]  # num_warps must be 4 to align with kernel implementation
+    ]
+    # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),]
+    return configs
+
+
+@triton.autotune(
+    configs=_get_fwd_autotune_configs(),
+    key=["num_heads", "TOPK", "D_V", "D_ROPE"],
+    prune_configs_by={"early_config_prune": _fwd_prune_configs},
+)
+@gluon.jit
+def _sparse_mla_fwd_gl_kernel(
+    Q_ptr,  # [total_tokens, num_heads, D_QK] bf16
+    KV_ptr,  # [total_tokens, 1, D_QK]         bf16
+    TopK_ptr,  # [total_tokens, TOPK]            int32
+    Sink_ptr,  # [num_heads]                     fp32; ignored if HAS_SINK == False
+    O_ptr,  # [total_tokens, num_heads, D_V]  bf16
+    LSE_ptr,  # [total_tokens, num_heads]       fp32 (sink-inclusive if HAS_SINK)
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_o_t: tl.int64,
+    stride_o_h: tl.int64,
+    stride_topk_t: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    TOPK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_SINK: gl.constexpr,
+):
+    # ---------- constexpr layouts ----------
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # Blocked layouts for global loads.
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K] = [64, 16]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_topk: gl.constexpr = gl.BlockedLayout(  # [TILE_K] int32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+    blk_lse: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H] fp32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+
+    # Shared layouts.
+    sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [BLOCK_H, D_V],
+        [1, 0],
+    )
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[1, 0],
+    )
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[0, 1],
+    )
+
+    # Dot operand layouts
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ---------- program ids ----------
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    # ---------- offsets for Q ----------
+    # Q_lora  [BLOCK_H, D_V]
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    # Q_rope  [BLOCK_H, D_ROPE]
+    offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+    offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+    mask_h_qrope = offs_h_qrope < num_heads
+
+    q_offs_rope = (
+        q_base + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+    )
+    q_mask_rope = mask_h_qrope[:, None]
+
+    smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora)
+    smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_qlora,
+        ptr=Q_ptr,
+        offsets=q_offs_lora.to(tl.int32),
+        mask=q_mask_lora,
+    )
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_qrope,
+        ptr=Q_ptr,
+        offsets=q_offs_rope.to(tl.int32),
+        mask=q_mask_rope,
+    )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # ---------- topk and KV offsets ----------
+    NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K
+    topk_base = token_idx.to(tl.int64) * stride_topk_t
+
+    # offs_tile in three layouts (sliced from each of the three loaders)
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk)
+
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # (removed dead `topk_pos_reg` prologue load — was never consumed)
+
+    # ---------- shared mem allocations for the K loop ----------
+    smem_krope = gl.allocate_shared_memory(
+        KV_ptr.dtype.element_ty,
+        [2, D_ROPE, TILE_K],
+        layout=sh_krope,
+    )
+    smem_klora = gl.allocate_shared_memory(
+        KV_ptr.dtype.element_ty,
+        [2, D_V, TILE_K],
+        layout=sh_klora,
+    )
+
+    # ---------- accumulators ----------
+    m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+
+    # ---------- tile-0 prefetch (prologue) ----------
+    # Load K_lora and K_rope for tile 0.
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_klora,
+        mask=offs_tile_klora < TOPK,
+        other=-1,
+    )
+    topk_pos_krope = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_krope,
+        mask=offs_tile_krope < TOPK,
+        other=-1,
+    )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_mma,
+        mask=offs_tile_mma < TOPK,
+        other=-1,
+    )
+
+    # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a
+    # single register set; converted to the klora/krope/mma layouts at point of use, to
+    # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen
+    # regression -- see att_fwd_gluon_mi350/RESULTS.md).
+    p1_off_topk = TILE_K + offs_tile_topk
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + p1_off_topk,
+        mask=p1_off_topk < TOPK,
+        other=-1,
+    )
+
+    valid_klora = topk_pos_klora != -1  # tile_start=0 -> offs_tile<TOPK already true
+    valid_krope = topk_pos_krope != -1
+    valid_mma = topk_pos_mma != -1
+
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+    safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+
+    # K_lora async DMA into smem_klora[0]
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    klora_smem0 = smem_klora.index(0)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=klora_smem0,
+        ptr=KV_ptr,
+        offsets=klora_offs.to(tl.int32),
+        mask=valid_klora[None, :],
+    )
+
+    # K_rope async DMA into smem_krope[0] — same group as K_lora.
+    krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(tl.int64)
+    krope_smem0 = smem_krope.index(0)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=krope_smem0,
+        ptr=KV_ptr,
+        offsets=krope_offs.to(tl.int32),
+        mask=valid_krope[None, :],
+    )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    gl.amd.cdna4.async_copy.wait_group(1)
+    Q_lora_dot = smem_qlora.load(dot_qlora_a)
+    Q_rope_dot = smem_qrope.load(dot_qrope_a)
+
+    # ---------- main loop (3-stage): gather tile t+1 from ALREADY-loaded topk,
+    #            deep-prefetch topk for tile t+2, compute tile t ----------
+    # Carried across iterations: tkraw_{klora,krope,mma} = raw topk for the tile to be
+    # GATHERED this iter (loaded the previous iter / prologue); valid_mma = mask for the
+    # tile being COMPUTED this iter. This lifts the topk load off the KV-gather critical
+    # path so the async_copy is never gated by a fresh index load.
+    # ===== Software-pipelined: softmax+PV of tile t-1 overlaps the QK MFMA of tile t. =====
+    # Carry S_prev (the masked/scaled QK output, small). QK reads cur_buf; PV reads 1-cur_buf;
+    # gather the next tile into 1-cur_buf AFTER V_dot has been read into registers, so the gather
+    # overlaps the PV MFMA (and the gather is cache-served/short). 2-deep buffer suffices.
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_mma = offs_h_mma < num_heads
+
+    # WARM-UP: gather K[1] -> buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet).
+    tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+    tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+    tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+    valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+    valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+    valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+    safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+    safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+    klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(1),
+        ptr=KV_ptr,
+        offsets=klora_offs_next.to(tl.int32),
+        mask=valid_klora_next[None, :],
+    )
+    krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(
+        tl.int64
+    )
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_krope.index(1),
+        ptr=KV_ptr,
+        offsets=krope_offs_next.to(tl.int32),
+        mask=valid_krope_next[None, :],
+    )
+    gl.amd.cdna4.async_copy.commit_group()
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk),
+        mask=(2 * TILE_K + offs_tile_topk) < TOPK,
+        other=-1,
+    )
+    gl.amd.cdna4.async_copy.wait_group(1)
+    S_prev = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(0).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev)
+    S_prev = S_prev * scale
+    S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf"))
+    cur_buf = 1
+
+    for t in range(NUM_TILES - 2):
+        gl.amd.cdna4.async_copy.wait_group(0)  # drain K[t+1] (cur_buf) before QK reads it
+        # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer,
+        # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer).
+        # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range.
+        V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b)
+        tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+        tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+        tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+        valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+        valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+        valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+        safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+        safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(1 - cur_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+            D_V + offs_r_krope[:, None]
+        ).to(tl.int64)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(1 - cur_buf),
+            ptr=KV_ptr,
+            offsets=krope_offs_next.to(tl.int32),
+            mask=valid_krope_next[None, :],
+        )
+        gl.amd.cdna4.async_copy.commit_group()
+        tkraw_n = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk),
+            mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK,
+            other=-1,
+        )
+        # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below
+        S_cur = gl.amd.cdna4.mfma(
+            Q_lora_dot,
+            smem_klora.index(cur_buf).load(dot_klora_b),
+            gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+        )
+        S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+        S_cur = S_cur * scale
+        S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+        # softmax(S_prev = tile t) [VALU, overlaps QK]
+        m_j = gl.max(S_prev, axis=1)
+        m_new = gl.maximum(m_i, m_j)
+        m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+        alpha = gl.exp(m_i - m_new)
+        P = gl.exp(S_prev - m_new[:, None])
+        l_i = alpha * l_i + gl.sum(P, axis=1)
+        m_i = m_new
+        # PV tile t from registers -- matrix; overlaps the gather still in flight
+        alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_acc[:, None]
+        P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+        acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc)
+        # promote
+        S_prev = S_cur
+        valid_qk = valid_qk_next
+        tkraw = tkraw_n
+        cur_buf = 1 - cur_buf
+
+    # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ----------
+    gl.amd.cdna4.async_copy.wait_group(0)  # drain K[N-1] (last loop gather)
+    S_cur = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(cur_buf).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+    S_cur = S_cur * scale
+    S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp(m_i - m_new)
+    P = gl.exp(S_prev - m_new[:, None])
+    l_i = alpha * l_i + gl.sum(P, axis=1)
+    m_i = m_new
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    S_prev = S_cur
+
+    # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ----------
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp(m_i - m_new)
+    P = gl.exp(S_prev - m_new[:, None])
+    l_new = alpha * l_i + gl.sum(P, axis=1)
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    m_i = m_new
+    l_i = l_new
+
+    # ---------- epilogue: fold sink into the denominator (V4 delta) ----------
+    if HAS_SINK:
+        offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        sink = gl.amd.cdna4.buffer_load(
+            ptr=Sink_ptr,
+            offsets=offs_h_sink.to(tl.int32),
+            mask=offs_h_sink < num_heads,
+            other=float("-inf"),
+        )
+        m_final = gl.maximum(m_i, sink)
+        alpha_fix = gl.exp(m_i - m_final)
+        l_total = l_i * alpha_fix + gl.exp(sink - m_final)
+        alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_fix_acc[:, None]
+        l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_total_acc[:, None]
+        lse = m_final + gl.log(l_total)
+    else:
+        l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_i_acc[:, None]
+        lse = m_i + gl.log(l_i)
+
+    # Output O[token_idx, h, v]
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    o_base = token_idx.to(tl.int64) * stride_o_t
+    o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64)
+    acc_bf = acc.to(O_ptr.dtype.element_ty)
+    acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora)
+    gl.amd.cdna4.buffer_store(
+        stored_value=acc_bf_blk,
+        ptr=O_ptr,
+        offsets=o_offs.to(tl.int32),
+        mask=mask_h_o[:, None],
+    )
+
+    # LSE[token_idx, h]
+    offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse)
+    mask_h_lse = offs_h_lse < num_heads
+    lse_base = token_idx * num_heads
+    lse_offs = lse_base + offs_h_lse
+    lse_blk = gl.convert_layout(lse, blk_lse)
+    gl.amd.cdna4.buffer_store(
+        stored_value=lse_blk,
+        ptr=LSE_ptr,
+        offsets=lse_offs.to(tl.int32),
+        mask=mask_h_lse,
+    )
+
+
+# =====================================================================
+# Launcher
+# =====================================================================
+def sparse_mla_fwd_v4_gluon(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """
+    DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink.
+
+    Args:
+        q:             [total_tokens, num_heads, d_qk] bfloat16
+        kv:            [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk])
+        topk_indices:  [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid)
+        attn_sink:     [num_heads] fp32, optional per-head learnable sink logit.
+                       When None, behaves like the V3.2 forward.
+        kv_lora_rank:  int, default 512
+        scale:         float, default 1/sqrt(d_qk)
+
+    Returns:
+        o:   [total_tokens, num_heads, kv_lora_rank] same dtype as q
+        lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given)
+    """
+    assert q.is_contiguous()
+    assert kv.is_contiguous()
+    assert topk_indices.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    # kv may hold MORE rows than there are query tokens (V4 feeds a
+    # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens).
+    # The kernel only dereferences kv via topk indices (stride_kv_t), so any
+    # num_kv >= max(topk_index)+1 is valid.
+    assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.is_contiguous()
+        assert attn_sink.dtype == torch.float32
+        assert attn_sink.shape == (num_heads,)
+        sink_ptr = attn_sink
+    else:
+        sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device)  # guarded by HAS_SINK
+
+    o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device)
+    lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    # Grid is autotune-aware: BLOCK_H comes from the chosen config.
+    grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"]))
+
+    _sparse_mla_fwd_gl_kernel[grid](
+        Q_ptr=q,
+        KV_ptr=kv,
+        TopK_ptr=topk_indices,
+        Sink_ptr=sink_ptr,
+        O_ptr=o,
+        LSE_ptr=lse,
+        stride_q_t=q.stride(0),
+        stride_q_h=q.stride(1),
+        stride_kv_t=kv.stride(0),
+        stride_o_t=o.stride(0),
+        stride_o_h=o.stride(1),
+        stride_topk_t=topk_indices.stride(0),
+        scale=scale,
+        num_heads=num_heads,
+        TOPK=topk,
+        D_V=kv_lora_rank,
+        D_ROPE=rope_rank,
+        HAS_SINK=has_sink,
+    )
+
+    return o, lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py
new file mode 100644
index 000000000..e3c606370
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py
@@ -0,0 +1,31 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon DeepSeek-V4 sparse-MLA attention backend ("gluon v2").
+
+Second-generation Gluon (gfx950 / CDNA4) sparse-MLA backend, using the same fused
+single-latent (K == V) sparse-MLA representation and public API as the other V4
+backends. Both directions are explicit-layout Gluon kernels:
+
+* forward (``dsa_fwd_v4_gluon``): MFMA layouts, padded/swizzled shared, async
+  double-buffered pipeline, rope-skip, exp2 softmax, MFMA K=32.
+* backward (``dsa_bwd_v4_gluon``): Gluon dQ + dKV-intermediate kernels (rope-skip,
+  MFMA K=32, single-chunk dQ RMW) + Triton Delta preprocess + CSR inverted-topk gather.
+
+The Gluon kernels need a Gluon-capable (recompiled) triton whose CDNA4 async_copy
+accepts general offset layouts, and raise a clear install hint otherwise.
+
+* :func:`sparse_mla_fwd_v4_gluon_v2` -> ``(o, lse)``
+* :func:`sparse_mla_bwd_v4_gluon_v2` -> ``(dq, dkv, d_sink)``
+"""
+
+from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon_v2
+from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon_v2
+
+__all__ = [
+    "sparse_mla_fwd_v4_gluon_v2",
+    "sparse_mla_bwd_v4_gluon_v2",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py
new file mode 100644
index 000000000..f3bf49f25
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py
@@ -0,0 +1,247 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton:
+the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`);
+this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing
+the external transpose copy.
+
+Per program: 1 query token. Grid: (total_tokens,).
+Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups:
+    dKV_lora[D_V, TILE_K]   = sum_hg ( Q_lora_T @ dS + dO_T @ P )    # contract over heads
+    dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS )
+    store interm[token, rank, :D_QK]
+
+Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in
+LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout
+-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dkv_interm_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16 (UNtransposed)
+    dO_ptr,  # [T, H, D_V]  bf16 (UNtransposed)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    Interm_ptr,  # [T, R_CHUNK, D_QK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    stride_interm_t: tl.int64,
+    stride_interm_r: tl.int64,
+    num_heads: tl.int32,
+    R_CHUNK: gl.constexpr,
+    TILE_K: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    NUM_HG: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads.
+    mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for HBM loads ----
+    # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read.
+    _q_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _q_tpw_m: gl.constexpr = 64 // _q_tpw_k
+    blk_q: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_q_tpw_m, _q_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose).
+    blk_ds: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 4],
+        threads_per_warp=[16, 4],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ----
+    sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0])
+
+    # ---- Dot operand layouts ----
+    dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+    dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+
+    token_idx = gl.program_id(axis=0)
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ---- LDS for Q/dO/Q_rope (single-buffered) ----
+    smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q)
+    smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do)
+    if HAS_ROPE:
+        smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    ds_base = token_idx.to(tl.int64) * stride_ds_t
+    interm_base = token_idx.to(tl.int64) * stride_interm_t
+
+    # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d]
+    offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma))
+    offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma))
+    offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma))
+
+    for t in range(NUM_TILES):
+        dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma)
+        if HAS_ROPE:
+            dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma)
+
+        for hg in range(NUM_HG):
+            hg_off = hg * BLOCK_H
+
+            # ---- stage Q/dO/Q_rope (this head group) into LDS ----
+            offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q))
+            offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q))
+            mask_h_q = offs_h_q < num_heads
+            q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64)
+            do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+
+            if HAS_ROPE:
+                offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+                offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+                mask_h_qr = offs_h_qr < num_heads
+                qr_offs = (
+                    q_base
+                    + offs_h_qr[:, None].to(tl.int64) * stride_q_h
+                    + (D_V + offs_r_qr[None, :]).to(tl.int64)
+                )
+                gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                    dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None]
+                )
+            gl.amd.cdna4.async_copy.commit_group()
+
+            # ---- load dS / P (this tile, this head group) -> dot operands ----
+            offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds))
+            offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds))
+            mask_h_ds = offs_h_ds < num_heads
+            dsp_offs = (
+                ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64)
+            )
+            dS_blk = gl.amd.cdna4.buffer_load(
+                ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            P_blk = gl.amd.cdna4.buffer_load(
+                ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            dS_dot = gl.convert_layout(dS_blk, dot_ds_b)
+            P_dot = gl.convert_layout(P_blk, dot_p_b)
+
+            # ---- wait + transpose-read Q/dO/Q_rope ----
+            gl.amd.cdna4.async_copy.wait_group(0)
+            Q_T = smem_q.permute([1, 0]).load(dot_qT_a)  # [D_V, BLOCK_H]
+            dO_T = smem_do.permute([1, 0]).load(dot_doT_a)  # [D_V, BLOCK_H]
+
+            dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora)
+            dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora)
+            if HAS_ROPE:
+                Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a)  # [D_ROPE, BLOCK_H]
+                dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope)
+
+        # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ----
+        col = t * TILE_K + offs_col_st
+        interm_lora_offs = (
+            interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64)
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty),
+            ptr=Interm_ptr,
+            offsets=interm_lora_offs.to(tl.int32),
+        )
+        # dKV_rope is provably zero for the V4 zero-rope-pad (discarded downstream by the
+        # gather/adapter, which use dkv[..., :D_V]); skip its compute + store when HAS_ROPE=False.
+        if HAS_ROPE:
+            interm_rope_offs = (
+                interm_base
+                + col[None, :].to(tl.int64) * stride_interm_r
+                + (D_V + offs_dr_st[:, None]).to(tl.int64)
+            )
+            gl.amd.cdna4.buffer_store(
+                stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty),
+                ptr=Interm_ptr,
+                offsets=interm_rope_offs.to(tl.int32),
+            )
+
+
+def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64):
+    """
+    Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel).
+
+    Returns interm [T, R_CHUNK, D_QK] bf16.
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    assert R_CHUNK % TILE_K == 0
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+        q,
+        do,
+        chunk_dS,
+        chunk_P,
+        interm,
+        q.stride(0),
+        q.stride(1),
+        do.stride(0),
+        do.stride(1),
+        chunk_dS.stride(0),
+        chunk_dS.stride(1),
+        interm.stride(0),
+        interm.stride(1),
+        num_heads,
+        R_CHUNK=R_CHUNK,
+        TILE_K=TILE_K,
+        BLOCK_H=BLOCK_H,
+        NUM_HG=num_hg,
+        D_V=kv_lora_rank,
+        D_ROPE=rope_rank,
+        num_warps=4,
+    )
+    return interm
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py
new file mode 100644
index 000000000..8ad3853ce
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py
@@ -0,0 +1,533 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the
+gluon hardware-control structure of Leon's V3.2 forward.
+
+Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK).
+Grid: (total_tokens, cdiv(num_heads, BLOCK_H)).  Dispatch loops chunks externally,
+dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits).
+
+Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K):
+    S       = Q_lora @ K_lora_T + Q_rope @ K_rope_T      # 2 MMAs (mfma_s),  contract over D
+    P       = exp(S*scale - lse)                          # lse is sink-inclusive (from fwd)
+    dP      = dO @ K_lora_T                               # 1 MMA (mfma_s),  reuses K_lora_T_dot
+    dS      = P * (dP - delta) * scale
+    dQ_lora += dS @ K_lora                                # 1 MMA (mfma_acc), K_lora = K_lora_T.T view
+    dQ_rope += dS @ K_rope                                # 1 MMA (mfma_acc), K_rope = K_rope_T.T view
+    store dS, P chunk -> HBM (consumed by dKV-intermediate kernel)
+
+Differences vs Leon's fwd (the structural template):
+  * dO added as a second stationary [BH, D_V] operand (async-loaded with Q).
+  * No online softmax: single P = exp(S - lse) (lse precomputed by fwd).
+  * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways.
+  * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write.
+  * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher),
+    so this kernel needs no atomics. lse from fwd already folds the sink in.
+
+M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA
+k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950).
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dq_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16
+    KV_ptr,  # [T, 1, D_QK] bf16
+    dO_ptr,  # [T, H, D_V]  bf16
+    TopK_ptr,  # [T, TOPK_padded] int32
+    LSE_ptr,  # [T, H] fp32 (sink-inclusive)
+    Delta_ptr,  # [T, H] fp32
+    dQ_ptr,  # [T, H, D_QK] bf16 (read-modify-write across chunks)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+    IS_FIRST_CHUNK: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    # mfma_s drives S = Q@K_T and dP = dO@K_T, both reducing D_V=512 -> K=32 halves their
+    # MFMA instruction count vs K=16 (mirrors the fwd R15 win). mfma_acc (dQ += dS@K) stays
+    # K=16 since its reduction dim is TILE_K (=16).
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 32],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for global loads (per Leon's fwd) ----
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_V]  (Q_lora, dO, dQ_lora)
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    # ---- Shared layouts (only K is staged in LDS) ----
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1])
+
+    # ---- Dot operand layouts ----
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+    dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ===================== program ids =====================
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ===================== Q / dO offsets =====================
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    if HAS_ROPE:
+        offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_qrope = offs_h_qrope < num_heads
+        q_offs_rope = (
+            q_base
+            + offs_h_qrope[:, None].to(tl.int64) * stride_q_h
+            + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+        )
+        q_mask_rope = mask_h_qrope[:, None]
+
+    offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_do = offs_h_do < num_heads
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64)
+    do_mask = mask_h_do[:, None]
+
+    # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) =====================
+    # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to
+    # the dot-operand layout once. The convert's LDS scratch is transient (freed before
+    # the K loop), unlike persistent staging, so only K occupies LDS during the loop.
+    q_lora_blk = gl.amd.cdna4.buffer_load(
+        ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0
+    )
+    do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0)
+    Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a)
+    dO_dot = gl.convert_layout(do_blk, dot_do_a)
+    if HAS_ROPE:
+        q_rope_blk = gl.amd.cdna4.buffer_load(
+            ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0
+        )
+        Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a)
+
+    # ===================== topk / KV offsets =====================
+    topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START
+
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # ===================== shared mem for K loop (double-buffered) =====================
+    if HAS_ROPE:
+        smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope)
+    smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora)
+
+    # ===================== dQ accumulators =====================
+    # Always zero-init; the read-modify-write across chunks is folded in at STORE
+    # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout
+    # whose LDS scratch would overflow the K/Q/dO buffers).
+    dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+    if HAS_ROPE:
+        dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc)
+
+    # ===================== lse / delta (in mfma_s row-slice layout) =====================
+    offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_s = offs_h_s < num_heads
+    lse = gl.amd.cdna4.buffer_load(
+        ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+    delta = gl.amd.cdna4.buffer_load(
+        ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+
+    # ===================== prologue: K tile 0 (group B, buffer 0) =====================
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=(topk_base + offs_tile_klora).to(tl.int32),
+        mask=offs_tile_klora < R_CHUNK,
+        other=-1,
+    )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1
+    )
+
+    valid_klora = topk_pos_klora != -1
+    valid_mma = topk_pos_mma != -1
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :]
+    )
+    if HAS_ROPE:
+        topk_pos_krope = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + offs_tile_krope).to(tl.int32),
+            mask=offs_tile_krope < R_CHUNK,
+            other=-1,
+        )
+        valid_krope = topk_pos_krope != -1
+        safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+        krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :]
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # dS / P store offsets in the mfma_s layout (store directly from the compute
+    # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead).
+    ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h
+    offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads
+
+    # ===================== main loop: prefetch t+1, compute t =====================
+    cur_buf = 0
+    for t in range(NUM_TILES - 1):
+        next_offs_klora = (t + 1) * TILE_K + offs_tile_klora
+        next_offs_krope = (t + 1) * TILE_K + offs_tile_krope
+        next_offs_mma = (t + 1) * TILE_K + offs_tile_mma
+
+        topk_pos_klora_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_klora).to(tl.int32),
+            mask=next_offs_klora < R_CHUNK,
+            other=-1,
+        )
+        topk_pos_mma_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_mma).to(tl.int32),
+            mask=next_offs_mma < R_CHUNK,
+            other=-1,
+        )
+
+        valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1)
+        valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1)
+        safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0)
+
+        next_buf = 1 - cur_buf
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(next_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        if HAS_ROPE:
+            topk_pos_krope_next = gl.amd.cdna4.buffer_load(
+                ptr=TopK_ptr,
+                offsets=(topk_base + next_offs_krope).to(tl.int32),
+                mask=next_offs_krope < R_CHUNK,
+                other=-1,
+            )
+            valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1)
+            safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0)
+            krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+                D_V + offs_r_krope[:, None]
+            ).to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_krope.index(next_buf),
+                ptr=KV_ptr,
+                offsets=krope_offs_next.to(tl.int32),
+                mask=valid_krope_next[None, :],
+            )
+        gl.amd.cdna4.async_copy.commit_group()
+
+        gl.amd.cdna4.async_copy.wait_group(1)
+
+        # ----- read K views from current buffer -----
+        klora_smem_cur = smem_klora.index(cur_buf)
+        K_lora_T_dot = klora_smem_cur.load(dot_klora_b)  # [D_V, TILE_K]  opIdx1 mfma_s
+        K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)  # [TILE_K, D_V] opIdx1 mfma_acc
+
+        # ----- S = Q_lora@K_lora_T (+ Q_rope@K_rope_T when HAS_ROPE) -----
+        S = gl.amd.cdna4.mfma(
+            Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        if HAS_ROPE:
+            krope_smem_cur = smem_krope.index(cur_buf)
+            K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+            S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+        S = S * scale
+        offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+        S = gl.where(valid_mask, S, float("-inf"))
+
+        # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale -----
+        P = gl.exp(S - lse[:, None])
+        P = gl.where(valid_mask, P, 0.0)
+        dP = gl.amd.cdna4.mfma(
+            dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        dS = P * (dP - delta[:, None]) * scale
+        dS = gl.where(valid_mask, dS, 0.0)
+
+        # ----- dQ_lora += dS@K_lora (+ dQ_rope += dS@K_rope when HAS_ROPE) -----
+        dS_bf = dS.to(KV_ptr.dtype.element_ty)
+        dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+        dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+        if HAS_ROPE:
+            K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+            dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+        # ----- store dS, P chunk -----
+        col = t * TILE_K + offs_tile_dsp
+        dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+        gl.amd.cdna4.buffer_store(
+            stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=P.to(KV_ptr.dtype.element_ty),
+            ptr=P_ptr,
+            offsets=dsp_offs.to(tl.int32),
+            mask=mask_h_dsp[:, None],
+        )
+
+        # promote prefetch -> current
+        cur_buf = next_buf
+        valid_mma = valid_mma_next
+
+    # ===================== epilogue: last tile =====================
+    gl.amd.cdna4.async_copy.wait_group(0)
+    t = NUM_TILES - 1
+    klora_smem_cur = smem_klora.index(cur_buf)
+    K_lora_T_dot = klora_smem_cur.load(dot_klora_b)
+    K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)
+
+    S = gl.amd.cdna4.mfma(
+        Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+    )
+    if HAS_ROPE:
+        krope_smem_cur = smem_krope.index(cur_buf)
+        K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+        S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+    S = S * scale
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+    S = gl.where(valid_mask, S, float("-inf"))
+
+    P = gl.exp(S - lse[:, None])
+    P = gl.where(valid_mask, P, 0.0)
+    dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s))
+    dS = P * (dP - delta[:, None]) * scale
+    dS = gl.where(valid_mask, dS, 0.0)
+
+    dS_bf = dS.to(KV_ptr.dtype.element_ty)
+    dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+    dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+    if HAS_ROPE:
+        K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+        dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+    col = t * TILE_K + offs_tile_dsp
+    dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+    )
+    gl.amd.cdna4.buffer_store(
+        stored_value=P.to(KV_ptr.dtype.element_ty),
+        ptr=P_ptr,
+        offsets=dsp_offs.to(tl.int32),
+        mask=mask_h_dsp[:, None],
+    )
+
+    # ===================== store dQ (lora + rope) =====================
+    dq_base = token_idx.to(tl.int64) * stride_dq_t
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64)
+    dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora)
+    if not IS_FIRST_CHUNK:
+        prev_lora = gl.amd.cdna4.buffer_load(
+            ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0
+        )
+        dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None]
+    )
+
+    # dQ_rope is provably zero for the V4 zero-rope-pad (the adapter discards dq[..., D_V:]),
+    # so skip its accumulation + store entirely when HAS_ROPE is False.
+    if HAS_ROPE:
+        offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_or = offs_h_or < num_heads
+        dq_offs_rope = (
+            dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64)
+        )
+        dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope)
+        if not IS_FIRST_CHUNK:
+            prev_rope = gl.amd.cdna4.buffer_load(
+                ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0
+            )
+            dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+        gl.amd.cdna4.buffer_store(
+            stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None]
+        )
+
+
+# =====================================================================
+# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P.
+# d_sink (if needed) is a torch reduction handled by the caller.
+# =====================================================================
+def sparse_mla_bwd_dq_gl(
+    q,
+    kv,
+    do,
+    topk_indices_padded,
+    lse,
+    delta,
+    R_CHUNK,
+    topk,
+    kv_lora_rank=512,
+    scale=None,
+    BLOCK_H=64,
+    TILE_K=16,
+):
+    """
+    Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop.
+
+    Returns:
+        dq:        [T, H, D_QK] bf16 (fully accumulated across chunks)
+        chunk_dS:  [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation)
+        chunk_P:   [T, H, R_CHUNK] bf16 (LAST chunk's P)
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK"
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    grid = (total_tokens, num_hg)
+
+    for r_start in range(0, topk, R_CHUNK):
+        is_first = r_start == 0
+        _sparse_mla_bwd_dq_gl_kernel[grid](
+            q,
+            kv,
+            do,
+            topk_indices_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_indices_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BLOCK_H,
+            TILE_K=TILE_K,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,
+        )
+    return dq, chunk_dS, chunk_P
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py
new file mode 100644
index 000000000..62d6cd72a
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py
@@ -0,0 +1,175 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon DeepSeek-V4 sparse-MLA backward for the "gluon_v2" backend.
+
+Companion to the gluon_v2 forward (:func:`sparse_mla_fwd_v4_gluon_v2`). Wires the Gluon
+dQ + Gluon dKV-intermediate compute kernels with a Triton Delta preprocess and the
+backend-neutral CSR inverted-topk gather + torch d_sink reduction (non-atomic
+chunked-gather scheme).
+
+The dQ / dKV-intermediate Gluon kernels apply the forward campaign's accepted techniques
+to the backward: rope-skip (the V4 zero-rope-pad makes the rope gradients provably zero)
++ MFMA K=32 for the D_V=512-reduction matmuls, plus a single-chunk dQ read-modify-write
+for high head counts. Beats the plain-Triton backward ~1.12x geomean over the 6
+flash/pro x cr{0,4,128} shapes (eager-UT 9/9).
+"""
+
+import torch
+import triton
+
+from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc
+from .._gluon_dsa._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess
+from .dsa_bwd_dkv_interm_gluon import _sparse_mla_bwd_dkv_interm_gl_kernel
+from .dsa_bwd_dq_gluon import _sparse_mla_bwd_dq_gl_kernel
+
+
+def sparse_mla_bwd_v4_gluon_v2(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA backward (Gluon dQ/dKV). Returns ``(dq, dkv, d_sink)``."""
+    assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous()
+    assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    num_kv = kv.shape[0]
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,)
+
+    # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ----
+    delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+    BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads))
+    _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))](
+        O_ptr=o,
+        dO_ptr=do,
+        Delta_ptr=delta,
+        stride_o_t=o.stride(0),
+        stride_o_h=o.stride(1),
+        num_heads=num_heads,
+        D_V=kv_lora_rank,
+        BLOCK_H=BLOCK_H_PRE,
+    )
+
+    # ---- config ----
+    # R2: dQ is read-modify-written across chunks, so more chunks = more redundant dq
+    # reload passes + repeated CSR builds. For high head counts (H>=128) the dq RMW volume
+    # is large, so a single chunk over the whole topk (bounded for memory) is a big win at
+    # pro cr4 (mirrors the triton bwd). Low head counts keep the 256 cap.
+    if num_heads >= 128:
+        R_CHUNK = min(topk, 1536)
+    else:
+        R_CHUNK = min(256, topk)
+    BH_DQ, TK_DQ = 64, 16
+    BH_DKV, TK_DKV = 32, 64
+    num_hg_dq = triton.cdiv(num_heads, BH_DQ)
+    num_hg_dkv = triton.cdiv(num_heads, BH_DKV)
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    # ---- pad topk to R_CHUNK multiple ----
+    topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK
+    if topk_padded_len != topk:
+        pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device)
+        topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous()
+    else:
+        topk_padded = topk_indices
+
+    all_csr = [
+        _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv)
+        for rs in range(0, topk, R_CHUNK)
+    ]
+
+    for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)):
+        is_first = r_start == 0
+
+        _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)](
+            q,
+            kv,
+            do,
+            topk_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BH_DQ,
+            TILE_K=TK_DQ,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_ROPE=False,  # V4 zero-rope-pad: dQ_rope is provably zero + discarded by the adapter
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,
+        )
+
+        _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+            q,
+            do,
+            chunk_dS,
+            chunk_P,
+            interm,
+            q.stride(0),
+            q.stride(1),
+            do.stride(0),
+            do.stride(1),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            interm.stride(0),
+            interm.stride(1),
+            num_heads,
+            R_CHUNK=R_CHUNK,
+            TILE_K=TK_DKV,
+            BLOCK_H=BH_DKV,
+            NUM_HG=num_hg_dkv,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_ROPE=False,  # V4 zero-rope-pad: dKV_rope provably zero + discarded downstream
+            num_warps=4,
+        )
+
+        inv_ptr, inv_data = all_csr[chunk_idx]
+        _bwd_dkv_gather_acc[(num_kv,)](
+            interm,
+            inv_ptr,
+            inv_data,
+            dkv_acc,
+            interm.stride(1),
+            dkv_acc.stride(0),
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            num_warps=4,
+        )
+
+    d_sink = None
+    if has_sink:
+        d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0)
+
+    dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1)
+    return dq, dkv_out, d_sink
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py
new file mode 100644
index 000000000..3d6847dbb
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py
@@ -0,0 +1,732 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink.
+
+Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch
+`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware
+control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline,
+ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456.
+This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches
+`sparse_mla_fwd_v4`; it is the forward of the "gluon_v2" backend.
+
+Optimizations accepted over the base gluon forward (gfx950 campaign):
+  * rope-skip (HAS_ROPE=False): the V4 latent bakes RoPE in-place over the 512, so the
+    rope QK term is provably zero -- skip the 64-wide rope MFMA + K_rope loads entirely.
+  * exp2 softmax: fold log2(e) into the QK scale so the per-element exp is one hardware
+    exp2 (m_i/l_i in log2 units; LSE converted back to natural log for the backward).
+  * MFMA K=32 for the QK score matmul (instr_shape [16,16,32]): the score reduces D_V=512,
+    so K=32 halves the QK MFMA instruction count vs K=16 (PV/acc stays K=16, TILE_K-bound).
+  * register-prefetched topk index (off the KV-gather critical path) + async double-buffered
+    K gather overlapping the QK/softmax/PV MFMAs.
+
+Pipeline:
+  Prologue:  Q -> shared (async); K tile 0 -> shared (async, double-buffered); deep-prefetch topk.
+  Loop tile t: gather tile t+2 (async) while computing QK[t+1] + softmax(t) + PV(t); promote.
+  Epilogue:  drain; fold sink into the denominator (V4); write O, LSE.
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+# ---------------------------------------------------------------------------
+# Triton capability gate: the gluon_v2 forward is a Gluon kernel (the backend's backward
+# is currently plain-Triton, being migrated to Gluon). The Gluon fwd needs a Gluon-capable triton whose CDNA4
+# async_copy accepts arbitrary (DistributedLinearLayout) offsets. Released
+# triton 3.7.0/3.7.1 still restricts async_copy offsets to Blocked/Slice and
+# will NOT compile this path; build triton from the commit below.
+# ---------------------------------------------------------------------------
+_GLUON_V2_REQUIRED_COMMIT = "09500db9f0"
+_GLUON_V2_INSTALL_HINT = (
+    "gluon_v2 forward (Gluon) requires a Gluon-capable triton whose CDNA4 "
+    "async_copy accepts general offset layouts. The installed triton ({ver}) does not (released "
+    "3.7.0/3.7.1 restrict async_copy offsets to BlockedLayout/SliceLayout).\n"
+    "Build & install triton-lang/triton @ commit " + _GLUON_V2_REQUIRED_COMMIT + ":\n"
+    "  git clone https://github.com/triton-lang/triton.git third_party/triton\n"
+    "  cd third_party/triton && git checkout " + _GLUON_V2_REQUIRED_COMMIT + "\n"
+    "  pip install -r python/requirements.txt\n"
+    "  TRITON_CODEGEN_BACKENDS=amd MAX_JOBS=128 pip wheel --no-build-isolation --no-deps . -w dist\n"
+    "  pip install --force-reinstall --no-deps dist/triton-*.whl"
+)
+
+
+@functools.lru_cache(maxsize=1)
+def _gluon_available() -> bool:
+    """True iff triton exposes the experimental CDNA4 Gluon dialect this fwd uses."""
+    try:
+        from triton.experimental import gluon as _gl  # noqa: F401
+        from triton.experimental.gluon.language import amd as _amd
+
+        return hasattr(_amd, "cdna4")
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _require_gluon_v2_triton() -> None:
+    """Fail fast with a build hint when Gluon is unavailable. The kernel compile is
+    ALSO guarded at the launch site: a Gluon-capable-but-incompatible triton raises a
+    CompilationError there, which is re-wrapped with the same install hint (so the user
+    always gets the commit rather than a raw layout/compile error)."""
+    if not _gluon_available():
+        raise RuntimeError(_GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown")))
+
+
+def _wrap_compile_error(exc: Exception) -> RuntimeError:
+    return RuntimeError(
+        _GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown"))
+        + f"\n(triton failed to compile the Gluon fwd: {type(exc).__name__}: "
+        + (str(exc).splitlines()[0] if str(exc).strip() else "")
+        + ")"
+    )
+
+
+# =====================================================================
+# Utility
+# =====================================================================
+def _get_lds_limit():
+    """Return the per-CU LDS limit in bytes for the current GPU.
+
+    gfx942 (MI300X): 64 KB = 65536 bytes
+    gfx950 (MI355X): 160 KB = 163840 bytes
+    """
+    if torch.cuda.is_available():
+        prop = torch.cuda.get_device_properties(0)
+        gcn_arch = getattr(prop, "gcnArchName", "")
+        if "gfx950" in gcn_arch:
+            return 163840
+    return 65536
+
+
+_LDS_LIMIT = _get_lds_limit()
+
+
+# =====================================================================
+# Forward — autotune configs and pruning
+# =====================================================================
+def _fwd_prune_configs(configs, named_args, **kwargs):
+    """Prune autotune configs that would exceed per-CU LDS."""
+    D_V = kwargs.get("D_V", named_args.get("D_V"))
+    D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE"))
+    pruned = []
+    for config in configs:
+        config.kwargs["BLOCK_H"]
+        tk = config.kwargs["TILE_K"]
+        ns = config.num_stages
+        kv_lds = (D_V + D_ROPE) * tk * 2 * ns
+        if kv_lds <= _LDS_LIMIT:
+            pruned.append(config)
+    if not pruned:
+        pruned.append(configs[0])
+    return pruned
+
+
+def _get_fwd_autotune_configs():
+    configs = [
+        triton.Config(
+            {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE},
+            num_warps=nw,
+        )
+        for BLOCK_H in [16, 32, 64]
+        for TILE_K in [16, 32, 64, 128]
+        for WPE in [0, 1, 2]
+        for nw in [4]  # num_warps must be 4 to align with kernel implementation
+    ]
+    # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),]
+    return configs
+
+
+@triton.autotune(
+    configs=_get_fwd_autotune_configs(),
+    key=["num_heads", "TOPK", "D_V", "D_ROPE"],
+    prune_configs_by={"early_config_prune": _fwd_prune_configs},
+)
+@gluon.jit
+def _sparse_mla_fwd_gl_v2_kernel(
+    Q_ptr,  # [total_tokens, num_heads, D_QK] bf16
+    KV_ptr,  # [total_tokens, 1, D_QK]         bf16
+    TopK_ptr,  # [total_tokens, TOPK]            int32
+    Sink_ptr,  # [num_heads]                     fp32; ignored if HAS_SINK == False
+    O_ptr,  # [total_tokens, num_heads, D_V]  bf16
+    LSE_ptr,  # [total_tokens, num_heads]       fp32 (sink-inclusive if HAS_SINK)
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_o_t: tl.int64,
+    stride_o_h: tl.int64,
+    stride_topk_t: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    TOPK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_SINK: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+):
+    # ---------- constexpr layouts ----------
+    # QK MFMA uses K=32 (aiter): the score matmul reduces D_V=512, so instr_shape=[16,16,32]
+    # halves the MFMA instruction count vs [16,16,16]. mfma_acc (PV) stays K=16 since its
+    # reduction dim is TILE_K (autotuned down to 16).
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 32],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # Blocked layouts for global loads.
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K] = [64, 16]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_topk: gl.constexpr = gl.BlockedLayout(  # [TILE_K] int32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+    blk_lse: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H] fp32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+
+    # Shared layouts.
+    sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [BLOCK_H, D_V],
+        [1, 0],
+    )
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[1, 0],
+    )
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[0, 1],
+    )
+
+    # Dot operand layouts
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ---------- program ids ----------
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    # ---------- offsets for Q ----------
+    # Q_lora  [BLOCK_H, D_V]
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_qlora,
+        ptr=Q_ptr,
+        offsets=q_offs_lora.to(tl.int32),
+        mask=q_mask_lora,
+    )
+    # V4 zero-rope-pad: skip the rope Q load + rope MFMA entirely when HAS_ROPE is False
+    # (RoPE is baked in-place over the 512 latent, so the rope QK term is provably zero).
+    if HAS_ROPE:
+        # Q_rope  [BLOCK_H, D_ROPE]
+        offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_qrope = offs_h_qrope < num_heads
+        q_offs_rope = (
+            q_base
+            + offs_h_qrope[:, None].to(tl.int64) * stride_q_h
+            + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+        )
+        q_mask_rope = mask_h_qrope[:, None]
+        smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_qrope,
+            ptr=Q_ptr,
+            offsets=q_offs_rope.to(tl.int32),
+            mask=q_mask_rope,
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # ---------- topk and KV offsets ----------
+    NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K
+    topk_base = token_idx.to(tl.int64) * stride_topk_t
+
+    # offs_tile in three layouts (sliced from each of the three loaders)
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk)
+
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # (removed dead `topk_pos_reg` prologue load — was never consumed)
+
+    # ---------- shared mem allocations for the K loop ----------
+    if HAS_ROPE:
+        smem_krope = gl.allocate_shared_memory(
+            KV_ptr.dtype.element_ty,
+            [2, D_ROPE, TILE_K],
+            layout=sh_krope,
+        )
+    smem_klora = gl.allocate_shared_memory(
+        KV_ptr.dtype.element_ty,
+        [2, D_V, TILE_K],
+        layout=sh_klora,
+    )
+
+    # ---------- accumulators ----------
+    m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+    # exp2 softmax (aiter technique): fold log2(e) into the QK scale so the per-element
+    # softmax exp becomes a single hardware exp2 (no per-element *log2e). m_i / l_i are then
+    # in log2 units; lse is converted back to natural log at the epilogue (the bwd needs nat-log).
+    scale_log2 = scale * 1.4426950408889634
+
+    # ---------- tile-0 prefetch (prologue) ----------
+    # Load K_lora and K_rope for tile 0.
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_klora,
+        mask=offs_tile_klora < TOPK,
+        other=-1,
+    )
+    if HAS_ROPE:
+        topk_pos_krope = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=topk_base.to(tl.int32) + offs_tile_krope,
+            mask=offs_tile_krope < TOPK,
+            other=-1,
+        )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_mma,
+        mask=offs_tile_mma < TOPK,
+        other=-1,
+    )
+
+    # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a
+    # single register set; converted to the klora/krope/mma layouts at point of use, to
+    # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen
+    # regression -- see att_fwd_gluon_mi350/RESULTS.md).
+    p1_off_topk = TILE_K + offs_tile_topk
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + p1_off_topk,
+        mask=p1_off_topk < TOPK,
+        other=-1,
+    )
+
+    valid_klora = topk_pos_klora != -1  # tile_start=0 -> offs_tile<TOPK already true
+    valid_mma = topk_pos_mma != -1
+
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+
+    # K_lora async DMA into smem_klora[0]
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    klora_smem0 = smem_klora.index(0)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=klora_smem0,
+        ptr=KV_ptr,
+        offsets=klora_offs.to(tl.int32),
+        mask=valid_klora[None, :],
+    )
+
+    # K_rope async DMA into smem_krope[0] — same group as K_lora (skipped if HAS_ROPE=False).
+    if HAS_ROPE:
+        valid_krope = topk_pos_krope != -1
+        safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+        krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(
+            tl.int64
+        )
+        krope_smem0 = smem_krope.index(0)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=krope_smem0,
+            ptr=KV_ptr,
+            offsets=krope_offs.to(tl.int32),
+            mask=valid_krope[None, :],
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    gl.amd.cdna4.async_copy.wait_group(1)
+    Q_lora_dot = smem_qlora.load(dot_qlora_a)
+    if HAS_ROPE:
+        Q_rope_dot = smem_qrope.load(dot_qrope_a)
+
+    # ---------- main loop (3-stage): gather tile t+1 from ALREADY-loaded topk,
+    #            deep-prefetch topk for tile t+2, compute tile t ----------
+    # Carried across iterations: tkraw_{klora,krope,mma} = raw topk for the tile to be
+    # GATHERED this iter (loaded the previous iter / prologue); valid_mma = mask for the
+    # tile being COMPUTED this iter. This lifts the topk load off the KV-gather critical
+    # path so the async_copy is never gated by a fresh index load.
+    # ===== Software-pipelined: softmax+PV of tile t-1 overlaps the QK MFMA of tile t. =====
+    # Carry S_prev (the masked/scaled QK output, small). QK reads cur_buf; PV reads 1-cur_buf;
+    # gather the next tile into 1-cur_buf AFTER V_dot has been read into registers, so the gather
+    # overlaps the PV MFMA (and the gather is cache-served/short). 2-deep buffer suffices.
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_mma = offs_h_mma < num_heads
+
+    # WARM-UP: gather K[1] -> buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet).
+    tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+    tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+    valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+    valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+    safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+    klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(1),
+        ptr=KV_ptr,
+        offsets=klora_offs_next.to(tl.int32),
+        mask=valid_klora_next[None, :],
+    )
+    if HAS_ROPE:
+        tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+        valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+        safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+        krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+            D_V + offs_r_krope[:, None]
+        ).to(tl.int64)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(1),
+            ptr=KV_ptr,
+            offsets=krope_offs_next.to(tl.int32),
+            mask=valid_krope_next[None, :],
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk),
+        mask=(2 * TILE_K + offs_tile_topk) < TOPK,
+        other=-1,
+    )
+    gl.amd.cdna4.async_copy.wait_group(1)
+    S_prev = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(0).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    if HAS_ROPE:
+        S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev)
+    S_prev = S_prev * scale_log2
+    S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf"))
+    cur_buf = 1
+
+    for t in range(NUM_TILES - 2):
+        gl.amd.cdna4.async_copy.wait_group(0)  # drain K[t+1] (cur_buf) before QK reads it
+        # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer,
+        # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer).
+        # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range.
+        V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b)
+        tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+        tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+        valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+        valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+        safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(1 - cur_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        if HAS_ROPE:
+            tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+            valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+            safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+            krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+                D_V + offs_r_krope[:, None]
+            ).to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_krope.index(1 - cur_buf),
+                ptr=KV_ptr,
+                offsets=krope_offs_next.to(tl.int32),
+                mask=valid_krope_next[None, :],
+            )
+        gl.amd.cdna4.async_copy.commit_group()
+        tkraw_n = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk),
+            mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK,
+            other=-1,
+        )
+        # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below
+        S_cur = gl.amd.cdna4.mfma(
+            Q_lora_dot,
+            smem_klora.index(cur_buf).load(dot_klora_b),
+            gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+        )
+        if HAS_ROPE:
+            S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+        S_cur = S_cur * scale_log2
+        S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+        # softmax(S_prev = tile t) [VALU, overlaps QK]
+        m_j = gl.max(S_prev, axis=1)
+        m_new = gl.maximum(m_i, m_j)
+        m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+        alpha = gl.exp2(m_i - m_new)
+        P = gl.exp2(S_prev - m_new[:, None])
+        l_i = alpha * l_i + gl.sum(P, axis=1)
+        m_i = m_new
+        # PV tile t from registers -- matrix; overlaps the gather still in flight
+        alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_acc[:, None]
+        P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+        acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc)
+        # promote
+        S_prev = S_cur
+        valid_qk = valid_qk_next
+        tkraw = tkraw_n
+        cur_buf = 1 - cur_buf
+
+    # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ----------
+    gl.amd.cdna4.async_copy.wait_group(0)  # drain K[N-1] (last loop gather)
+    S_cur = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(cur_buf).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    if HAS_ROPE:
+        S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+    S_cur = S_cur * scale_log2
+    S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp2(m_i - m_new)
+    P = gl.exp2(S_prev - m_new[:, None])
+    l_i = alpha * l_i + gl.sum(P, axis=1)
+    m_i = m_new
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    S_prev = S_cur
+
+    # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ----------
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp2(m_i - m_new)
+    P = gl.exp2(S_prev - m_new[:, None])
+    l_new = alpha * l_i + gl.sum(P, axis=1)
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    m_i = m_new
+    l_i = l_new
+
+    # ---------- epilogue: fold sink into the denominator (V4 delta) ----------
+    if HAS_SINK:
+        offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        sink = gl.amd.cdna4.buffer_load(
+            ptr=Sink_ptr,
+            offsets=offs_h_sink.to(tl.int32),
+            mask=offs_h_sink < num_heads,
+            other=float("-inf"),
+        )
+        sink = sink * 1.4426950408889634  # natural-log sink -> log2 units (exp2 softmax)
+        m_final = gl.maximum(m_i, sink)
+        alpha_fix = gl.exp2(m_i - m_final)
+        l_total = l_i * alpha_fix + gl.exp2(sink - m_final)
+        alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_fix_acc[:, None]
+        l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_total_acc[:, None]
+        # lse back to natural log for the backward: m_final is log2, l_total is the natural denom.
+        lse = m_final * 0.6931471805599453 + gl.log(l_total)
+    else:
+        l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_i_acc[:, None]
+        lse = m_i * 0.6931471805599453 + gl.log(l_i)
+
+    # Output O[token_idx, h, v]
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    o_base = token_idx.to(tl.int64) * stride_o_t
+    o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64)
+    acc_bf = acc.to(O_ptr.dtype.element_ty)
+    acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora)
+    gl.amd.cdna4.buffer_store(
+        stored_value=acc_bf_blk,
+        ptr=O_ptr,
+        offsets=o_offs.to(tl.int32),
+        mask=mask_h_o[:, None],
+    )
+
+    # LSE[token_idx, h]
+    offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse)
+    mask_h_lse = offs_h_lse < num_heads
+    lse_base = token_idx * num_heads
+    lse_offs = lse_base + offs_h_lse
+    lse_blk = gl.convert_layout(lse, blk_lse)
+    gl.amd.cdna4.buffer_store(
+        stored_value=lse_blk,
+        ptr=LSE_ptr,
+        offsets=lse_offs.to(tl.int32),
+        mask=mask_h_lse,
+    )
+
+
+# =====================================================================
+# Launcher
+# =====================================================================
+def sparse_mla_fwd_v4_gluon_v2(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """
+    DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink.
+
+    Args:
+        q:             [total_tokens, num_heads, d_qk] bfloat16
+        kv:            [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk])
+        topk_indices:  [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid)
+        attn_sink:     [num_heads] fp32, optional per-head learnable sink logit.
+                       When None, behaves like the V3.2 forward.
+        kv_lora_rank:  int, default 512
+        scale:         float, default 1/sqrt(d_qk)
+
+    Returns:
+        o:   [total_tokens, num_heads, kv_lora_rank] same dtype as q
+        lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given)
+    """
+    _require_gluon_v2_triton()
+    assert q.is_contiguous()
+    assert kv.is_contiguous()
+    assert topk_indices.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    # kv may hold MORE rows than there are query tokens (V4 feeds a
+    # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens).
+    # The kernel only dereferences kv via topk indices (stride_kv_t), so any
+    # num_kv >= max(topk_index)+1 is valid.
+    assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.is_contiguous()
+        assert attn_sink.dtype == torch.float32
+        assert attn_sink.shape == (num_heads,)
+        sink_ptr = attn_sink
+    else:
+        sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device)  # guarded by HAS_SINK
+
+    o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device)
+    lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked in-place
+    # over the 512 latent), so the rope QK term is provably zero. Skip it — bit-identical
+    # to computing it, but avoids the wasteful 64-wide rope MFMA + K_rope loads every tile
+    # (this is the win triton_v2 already has that our ported gluon fwd lacked).
+    has_rope = False
+
+    # Grid is autotune-aware: BLOCK_H comes from the chosen config.
+    grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"]))
+
+    try:
+        _sparse_mla_fwd_gl_v2_kernel[grid](
+            Q_ptr=q,
+            KV_ptr=kv,
+            TopK_ptr=topk_indices,
+            Sink_ptr=sink_ptr,
+            O_ptr=o,
+            LSE_ptr=lse,
+            stride_q_t=q.stride(0),
+            stride_q_h=q.stride(1),
+            stride_kv_t=kv.stride(0),
+            stride_o_t=o.stride(0),
+            stride_o_h=o.stride(1),
+            stride_topk_t=topk_indices.stride(0),
+            scale=scale,
+            num_heads=num_heads,
+            TOPK=topk,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_SINK=has_sink,
+            HAS_ROPE=has_rope,
+        )
+    except Exception as exc:  # noqa: BLE001 - surface a build hint on Gluon compile failures
+        _n = type(exc).__name__.lower()
+        if "compil" in _n or "compil" in str(exc).lower() or "layout" in str(exc).lower():
+            raise _wrap_compile_error(exc) from exc
+        raise
+
+    return o, lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py
new file mode 100644
index 000000000..a736001ef
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py
@@ -0,0 +1,31 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon DeepSeek-V4 sparse-MLA attention backend ("gluon v3").
+
+Third-generation Gluon (gfx950 / CDNA4) sparse-MLA optimization campaign backend.
+Round 1 intentionally starts from the stable ``gluon_v2`` implementation so every
+future round can change exactly one kernel variable and be compared linearly.
+
+* forward (``dsa_fwd_v4_gluon``): MFMA layouts, padded/swizzled shared, async
+  double-buffered pipeline, rope-skip, exp2 softmax, MFMA K=32.
+* backward (``dsa_bwd_v4_gluon``): Gluon dQ + dKV-intermediate kernels (rope-skip,
+  MFMA K=32, single-chunk dQ RMW) + Triton Delta preprocess + CSR inverted-topk gather.
+
+The Gluon kernels need a Gluon-capable (recompiled) triton whose CDNA4 async_copy
+accepts general offset layouts, and raise a clear install hint otherwise.
+
+* :func:`sparse_mla_fwd_v4_gluon_v3` -> ``(o, lse)``
+* :func:`sparse_mla_bwd_v4_gluon_v3` -> ``(dq, dkv, d_sink)``
+"""
+
+from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon_v2 as sparse_mla_bwd_v4_gluon_v3
+from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon_v2 as sparse_mla_fwd_v4_gluon_v3
+
+__all__ = [
+    "sparse_mla_fwd_v4_gluon_v3",
+    "sparse_mla_bwd_v4_gluon_v3",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py
new file mode 100644
index 000000000..e3437d291
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py
@@ -0,0 +1,181 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter) gluon sparse-MLA kernels; see _gluon_v3/.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from .aiter_mla_gluon import mla_gluon
+
+
+@triton.jit
+def _count_valid_topk_kernel(topk_ptr, counts_ptr, stride_t, topk: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    offs = tl.arange(0, BLOCK)
+    vals = tl.load(topk_ptr + row * stride_t + offs, mask=offs < topk, other=-1)
+    valid = (offs < topk) & (vals >= 0)
+    count = tl.sum(valid.to(tl.int32), axis=0)
+    tl.store(counts_ptr + row, count)
+
+
+@triton.jit
+def _pack_valid_topk_kernel(
+    topk_ptr,
+    indptr_ptr,
+    flat_ptr,
+    stride_t,
+    topk: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offs = tl.arange(0, BLOCK)
+    vals = tl.load(topk_ptr + row * stride_t + offs, mask=offs < topk, other=-1)
+    valid = (offs < topk) & (vals >= 0)
+    pos = tl.cumsum(valid.to(tl.int32), axis=0) - 1
+    start = tl.load(indptr_ptr + row)
+    tl.store(flat_ptr + start + pos, vals.to(tl.int32), mask=valid)
+
+
+def _v4_topk_to_ragged_gpu(topk_indices: torch.Tensor):
+    total_tokens, topk = topk_indices.shape
+    block = triton.next_power_of_2(topk)
+    counts = torch.empty(total_tokens, dtype=torch.int32, device=topk_indices.device)
+    _count_valid_topk_kernel[(total_tokens,)](
+        topk_indices,
+        counts,
+        topk_indices.stride(0),
+        topk,
+        BLOCK=block,
+    )
+    indptr = torch.empty(total_tokens + 1, dtype=torch.int32, device=topk_indices.device)
+    indptr[0] = 0
+    torch.cumsum(counts, dim=0, out=indptr[1:])
+
+    # Allocate the max possible nnz to avoid a CPU sync on indptr[-1]. The Gluon
+    # kernel uses indptr for bounds, so unused tail elements are never read.
+    flat = torch.empty(total_tokens * topk, dtype=torch.int32, device=topk_indices.device)
+    _pack_valid_topk_kernel[(total_tokens,)](
+        topk_indices,
+        indptr,
+        flat,
+        topk_indices.stride(0),
+        topk,
+        BLOCK=block,
+    )
+    return flat, indptr
+
+
+@triton.jit
+def _pack_v4_csa_h64_kernel(
+    topk_ptr,
+    indptr_ptr,
+    flat_ptr,
+    stride_t,
+    total_tokens: tl.constexpr,
+    TOPK: tl.constexpr,
+    WINDOW: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    pool_k: tl.constexpr = TOPK - WINDOW
+    local_count = tl.minimum(row + 1, WINDOW)
+
+    if row < WINDOW:
+        local_prefix = row * (row + 1) // 2
+    else:
+        local_prefix = WINDOW * (WINDOW + 1) // 2 + (row - WINDOW) * WINDOW
+    start = row * pool_k + local_prefix
+    tl.store(indptr_ptr + row, start)
+
+    offs = tl.arange(0, BLOCK)
+    local_mask = offs < local_count
+    pool_offs = offs - local_count
+    pool_mask = (offs >= local_count) & (pool_offs < pool_k)
+    src = tl.where(local_mask, WINDOW - local_count + offs, WINDOW + pool_offs)
+    vals = tl.load(topk_ptr + row * stride_t + src, mask=local_mask | pool_mask, other=-1)
+    tl.store(flat_ptr + start + offs, vals.to(tl.int32), mask=local_mask | pool_mask)
+
+    if row == total_tokens - 1:
+        tl.store(indptr_ptr + total_tokens, start + local_count + pool_k)
+
+
+def _v4_csa_h64_topk_to_ragged(topk_indices: torch.Tensor):
+    total_tokens, topk = topk_indices.shape
+    block = triton.next_power_of_2(topk)
+    indptr = torch.empty(total_tokens + 1, dtype=torch.int32, device=topk_indices.device)
+    flat = torch.empty(total_tokens * topk, dtype=torch.int32, device=topk_indices.device)
+    _pack_v4_csa_h64_kernel[(total_tokens,)](
+        topk_indices,
+        indptr,
+        flat,
+        topk_indices.stride(0),
+        total_tokens,
+        topk,
+        WINDOW=128,
+        BLOCK=block,
+    )
+    return flat, indptr
+
+
+def sparse_mla_fwd_v4_aiter_lse(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """Aiter Gluon sparse-MLA fwd with LSE, adapted to the V4 dense-topk API."""
+    _, _, d_qk = q.shape
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+
+    q_nope = q[..., :kv_lora_rank].contiguous()
+    kv2 = kv[:, 0, :] if kv.dim() == 3 else kv
+    kv_c = kv2[:, :kv_lora_rank].contiguous()
+    flat, indptr = _v4_topk_to_ragged_gpu(topk_indices)
+    out = torch.empty(q.shape[0], q.shape[1], kv_lora_rank, dtype=q.dtype, device=q.device)
+
+    return mla_gluon(
+        q_nope,
+        None,
+        kv_c,
+        out,
+        page_table=flat,
+        seq_info=indptr,
+        sm_scale=float(scale),
+        has_pe=False,
+        min_kv_seq_len=float("inf"),
+        attn_sink=attn_sink,
+        return_lse=True,
+    )
+
+
+def sparse_mla_fwd_v4_aiter_lse_csa_formula(
+    q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None
+):
+    """CSA specialization using the V4 [SWA128 + pool topk] dense-topk layout."""
+    _, _, d_qk = q.shape
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+
+    q_nope = q[..., :kv_lora_rank].contiguous()
+    kv2 = kv[:, 0, :] if kv.dim() == 3 else kv
+    kv_c = kv2[:, :kv_lora_rank].contiguous()
+    flat, indptr = _v4_csa_h64_topk_to_ragged(topk_indices)
+    out = torch.empty(q.shape[0], q.shape[1], kv_lora_rank, dtype=q.dtype, device=q.device)
+
+    return mla_gluon(
+        q_nope,
+        None,
+        kv_c,
+        out,
+        page_table=flat,
+        seq_info=indptr,
+        sm_scale=float(scale),
+        has_pe=False,
+        min_kv_seq_len=float("inf"),
+        attn_sink=attn_sink,
+        return_lse=True,
+    )
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py
new file mode 100644
index 000000000..4f7230f0e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py
@@ -0,0 +1,1002 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions copied and modified from ROCm/aiter (https://github.com/ROCm/aiter) and
+# DeepSeek FlashMLA (https://github.com/deepseek-ai/FlashMLA); see comments below.
+
+# Gluon MLA decode kernel originated from FlashMLA triton kernel(https://github.com/deepseek-ai/FlashMLA/blob/main/benchmark/bench_flash_mla.py).
+# Stage-1 split-KV MLA attention using explicit Gluon layouts. Three regimes:
+#
+#   REGIME='bh64'      - bf16 Q + bf16 KV, BLOCK_H=64, BLOCK_N=64,
+#                        nhead in {64, 128}, batch_size in {64, 128, 256},
+#                        NUM_KV_SPLITS auto-picked to fill ~256 WGs (in {1,2,4}).
+#                        Fast path: when NUM_KV_SPLITS==1, stage-1 writes the
+#                        final output directly to O and stage-2 reduce is skipped.
+#   REGIME='bh16bn128' - bf16 Q + fp8 KV, BLOCK_H=16, BLOCK_N=128,
+#                        nhead <= 16, batch_size=1, NUM_KV_SPLITS=256.
+#                        2-D (batch, split) grid. Always splits + always
+#                        reduces. NHEAD < BLOCK_H masks OOB heads on Q load
+#                        and O store.
+#   REGIME='bh16bn64'  - bf16 Q + bf16 KV, BLOCK_H=16, BLOCK_N=64,
+#                        nhead <= 16, batch_size >= 1, 2-D (batch, split) grid,
+#                        NUM_KV_SPLITS = max(1, 256 // batch_size). Full decode
+#                        (stage-1 + stage-2 reduce into the final O).
+#                        NHEAD < BLOCK_H masks OOB heads on Q load and O store.
+#
+# The bh16 regimes support num_iter in {1, 2, ...} (no gl.assume(num_iter>=3));
+# only bh64 assumes >= 3. See epilogue-1 handling below.
+#
+# Wrapper dispatch: nhead in {64,128} -> bh64; nhead <= 16 routes by KV dtype
+# (bf16 -> bh16bn64, fp8 -> bh16bn128).
+#
+# Full decode for all regimes. For NUM_KV_SPLITS>1 stage-1 writes per-split acc +
+# fp32 lse; stage-2 (_mla_softmax_reducev_kernel) reduces into O. RETURN_LSE also
+# returns the merged fp32 lse [B, H] (stage-2 for splits>1, else stage-1).
+#
+# 3-stage software pipeline (double-buffered, BLOCK_N with 2x(BLOCK_N/2) KV slices):
+#   AC = async_copy (global->LDS), LL = load (LDS->reg), P = page, K = K-cache, V = V-cache
+#
+#                      iter i          iter i+1        iter i+2
+#   ACP(page):        [i+2]           [i+3]           [i+4]
+#   LLP+ACK(K):       [i+1]           [i+2]           [i+3]
+#   LLK+MFMA+LLV:     [i]             [i+1]           [i+2]
+#
+#   Within each loop iteration (operating on buf_idx=current, async_idx=next):
+#     ACP                                -- async_copy page numbers [i+2]
+#     LLP, ACK                           -- local_load pages [i+1], async_copy K/KPE [i+1]
+#     LLK, MFMA0, softmax, LLV, MFMA1   -- compute on [i]: QK dot, softmax, PV dot
+
+import aiter.ops.triton.utils._triton.arch_info as arch_info
+import torch
+import triton
+import triton.language as tl
+from aiter.ops.triton.utils.device_info import get_num_xcds
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+# fmt: off
+@gluon.jit
+def _mla_gluon(
+    Q_nope,
+    Q_pe,
+    Kv_c_cache,
+    K_pe_cache,
+    Req_to_tokens,
+    B_seq_len,
+    O,  # noqa: E741
+    Attn_sink,
+    sm_scale,
+    kv_scale,
+    stride_q_nope_bs,
+    stride_q_nope_h,
+    stride_q_pe_bs,
+    stride_q_pe_h,
+    stride_kv_c_bs,
+    stride_k_pe_bs,
+    stride_req_to_tokens_bs,
+    stride_o_b,
+    stride_o_h,
+    stride_o_s,
+    Mid_lse,  # split>1: per-split fp32 lse [B, H, NUM_KV_SPLITS] (else None)
+    stride_mid_lse_b,
+    stride_mid_lse_h,
+    stride_mid_lse_s,
+    Final_lse,  # RETURN_LSE only: merged fp32 lse [B, H] (else None)
+    stride_final_lse_b,
+    stride_final_lse_h,
+    BLOCK_H: gl.constexpr,
+    BLOCK_N: gl.constexpr,
+    NUM_KV_SPLITS: gl.constexpr,
+    PAGE_SIZE: gl.constexpr,
+    HEAD_DIM_CKV: gl.constexpr,
+    HEAD_DIM_KPE: gl.constexpr,
+    KV_PE_OFFSET: gl.constexpr,
+    USE_2D_VIEW: gl.constexpr,
+    WITHIN_2GB: gl.constexpr,
+    NUM_XCDS: gl.constexpr,
+    NHEAD: gl.constexpr,
+    REGIME: gl.constexpr,
+    RETURN_LSE: gl.constexpr,
+    # --- dsv4-prefill knobs ---
+    HAS_PE: gl.constexpr,
+    HAS_ATTN_SINK: gl.constexpr,
+):
+    # Grid mapping: bh64 uses 3-D XCD-aware multi-batch; bh16bn64 and bh16bn128
+    # use 2-D (batch, split) — for batch_size=1 this is (1, NUM_KV_SPLITS).
+    if REGIME == 'bh64':
+        cur_batch = gl.program_id(0) + (gl.program_id(2) // NUM_KV_SPLITS) * NUM_XCDS
+        cur_head_id = gl.program_id(1)
+        split_kv_id = gl.program_id(2) % NUM_KV_SPLITS
+    else:
+        cur_batch = gl.program_id(0)
+        cur_head_id = 0
+        split_kv_id = gl.program_id(1)
+
+    # USE_2D_VIEW=True: fixed len or max padded VarLen
+    # Req_to_tokens = block_table[batch, max_seqlen], B_seq_len = cache_seqlens[batch]
+    # USE_2D_VIEW=False: flattened VarLen
+    # Req_to_tokens = kv_indices[total_kv],           B_seq_len = kv_indptr[batch+1]
+    if USE_2D_VIEW:
+        batch_page_start = stride_req_to_tokens_bs * cur_batch
+        cur_batch_seq_len = gl.load(B_seq_len + cur_batch)
+    else:
+        batch_page_start = gl.load(B_seq_len + cur_batch)
+        cur_batch_seq_len = gl.load(B_seq_len + cur_batch + 1) - batch_page_start
+
+    # split-KV: each program covers [split_kv_start, split_kv_end).
+    # OLD: ceil-based per_split. The LAST split could be empty (num_iter=0),
+    # which breaks the unconditional epilogue-2 consume. Kept here as commented
+    # reference; remove in cleanup.
+    # kv_len_per_split = gl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    # split_kv_start = kv_len_per_split * split_kv_id
+    # split_kv_end = gl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+    #
+    # NEW: floor per_split with the last split absorbing the remainder
+    # (remainder = seq mod NUM_KV_SPLITS, in [0, NUM_KV_SPLITS)). Combined with
+    # the wrapper bound min_kv_seq_len >= NUM_KV_SPLITS this guarantees every
+    # split is non-empty (split_len >= floor >= 1, hence num_iter >= 1); bh64
+    # additionally bounds min_kv_seq_len so num_iter >= 3 for its gl.assume.
+    # Trade-off: at seqs just above the wrapper minimum the last CU does up to
+    # ~(floor + NUM_KV_SPLITS - 1)/floor more work than the others.
+    kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = split_kv_start + kv_len_per_split
+    if split_kv_id == NUM_KV_SPLITS - 1:
+        split_kv_end = cur_batch_seq_len
+    num_iter = gl.cdiv(split_kv_end - split_kv_start, BLOCK_N)
+    start_n = split_kv_start
+
+    # early return with empty kv slice to save compute
+    if split_kv_start >= split_kv_end:
+        return
+
+    ######### layout setting begin #########
+    # Q-side layouts + mfma_layout: switch by BLOCK_H.
+    # bh64 has BLOCK_H=64; bh16bn128 and bh16bn64 share BLOCK_H=16 (identical Q layouts + mfma orientation).
+    if BLOCK_H == 64:
+        # bh64: Q is [64, 512] / [64, 64]; warps tile M.
+        blocked_q_nope: gl.constexpr = gl.BlockedLayout(
+            size_per_thread=[1, 8],
+            threads_per_warp=[1, 64],
+            warps_per_cta=[4, 1],
+            order=[1, 0],
+        )
+        shared_q_nope: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[512, 16]],
+            offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 128], [0, 256], [1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0]],
+            cga_layout=[],
+            shape=[64, 512]
+        )
+        blocked_q_pe: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0, 1), (0, 2), (0, 4), (32, 0)),
+            lane_bases=((0, 8), (0, 16), (0, 32), (4, 0), (8, 0), (16, 0)),
+            warp_bases=((1, 0), (2, 0)),
+            block_bases=[],
+            shape=[64, 64],
+        )
+        shared_q_pe: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[512, 16]],
+            offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [4, 0], [8, 0], [16, 0], [1, 0], [2, 0], [32, 0]],
+            cga_layout=[],
+            shape=[64, 64]
+        )
+        mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
+            version=4,
+            instr_shape=[16, 16, 32],
+            transposed=True,
+            warps_per_cta=[4, 1],
+        )
+    else:
+        # BLOCK_H == 16: shared by bh16bn128 and bh16bn64. Q is [16, 512] / [16, 64]; warps tile K.
+        blocked_q_nope: gl.constexpr = gl.BlockedLayout(
+            size_per_thread=[1, 8],
+            threads_per_warp=[1, 64],
+            warps_per_cta=[4, 1],
+            order=[1, 0],
+        )
+        shared_q_nope: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[512, 16]],
+            offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 128], [0, 256], [1, 0], [2, 0], [4, 0], [8, 0]],
+            cga_layout=[],
+            shape=[16, 512]
+        )
+        blocked_q_pe: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0, 1), (0, 2), (0, 4)),
+            lane_bases=((0, 8), (0, 16), (0, 32), (1, 0), (2, 0), (4, 0)),
+            warp_bases=((8, 0), (0, 0)),
+            block_bases=[],
+            shape=[16, 64],
+        )
+        shared_q_pe: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0])
+        mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
+            version=4,
+            instr_shape=[16, 16, 32],
+            transposed=True,
+            warps_per_cta=[1, 4],
+        )
+
+    # KV-side layouts: switch by BLOCK_N.
+    # bh16bn128 (BLOCK_N=128, fp8 KV) needs distinct K layouts; bh64 and bh16bn64 share BLOCK_N=64 bf16 KV.
+    if BLOCK_N == 128:
+        # bh16bn128: K is [512, 128]fp8, KPE is [64, 128]fp8.
+        blocked_kv: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 4), (0, 32), (0, 64)),
+            lane_bases=((16, 0), (32, 0), (64, 0), (128, 0), (256, 0), (0, 16)),
+            warp_bases=((0, 1), (0, 2)),
+            block_bases=[],
+            shape=[512, 128],
+        )
+        shared_kv: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[1024, 32], [8192, 16]],
+            offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [64, 0], [128, 0], [256, 0], [0, 16], [0, 1], [0, 2], [0, 8], [0, 4], [0, 32], [0, 64]],
+            cga_layout=[],
+            shape=[512, 128]
+        )
+        blocked_kpe: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 2)),
+            lane_bases=((16, 0), (32, 0), (0, 4), (0, 8), (0, 16), (0, 32)),
+            warp_bases=((0, 64), (0, 1)),
+            block_bases=[],
+            shape=[64, 128],
+        )
+        shared_kpe: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[2048, 16]],
+            offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 1], [0, 2]],
+            cga_layout=[],
+            shape=[64, 128]
+        )
+        blocked_page: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0,),),
+            lane_bases=((1,), (2,), (4,), (8,), (16,), (32,)),
+            warp_bases=((64,), (0,)),
+            block_bases=[],
+            shape=[128],
+        )
+        blocked_kv_slice: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 4), (0, 32)),
+            lane_bases=((16, 0), (32, 0), (64, 0), (128, 0), (256, 0), (0, 16)),
+            warp_bases=((0, 1), (0, 2)),
+            block_bases=[],
+            shape=[512, 64],
+        )
+    else:
+        # BLOCK_N == 64: shared by bh64 and bh16bn64 (both bf16 KV).
+        # K is [512, 64]bf16, KPE is [64, 64]bf16.
+        blocked_kv: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (0, 8), (0, 4), (0, 16), (0, 32)),
+            lane_bases=((8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)),
+            warp_bases=((0, 1), (0, 2)),
+            block_bases=[],
+            shape=[512, 64],
+        )
+        shared_kv: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[512, 16]],
+            offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [64, 0], [128, 0], [256, 0], [0, 1], [0, 2], [0, 8], [0, 4], [0, 16], [0, 32]],
+            cga_layout=[],
+            shape=[512, 64]
+        )
+        blocked_kpe: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (0, 32)),
+            lane_bases=((8, 0), (16, 0), (32, 0), (0, 4), (0, 8), (0, 16)),
+            warp_bases=((0, 1), (0, 2)),
+            block_bases=[],
+            shape=[64, 64],
+        )
+        shared_kpe: gl.constexpr = gl.PaddedSharedLayout(
+            interval_padding_pairs=[[512, 16]],
+            offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [0, 4], [0, 8], [0, 16], [0, 1], [0, 2], [0, 32]],
+            cga_layout=[],
+            shape=[64, 64]
+        )
+        blocked_page: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0,),),
+            lane_bases=((1,), (2,), (4,), (8,), (16,), (32,)),
+            warp_bases=((0,), (0,)),
+            block_bases=[],
+            shape=[64],
+        )
+        blocked_kv_slice: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((1, 0), (2, 0), (4, 0), (0, 8), (0, 4), (0, 16)),
+            lane_bases=((8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)),
+            warp_bases=((0, 1), (0, 2)),
+            block_bases=[],
+            shape=[512, 32],
+        )
+
+    # linear_v: each regime has unique warp/reg mapping (bh64 has degenerate warp_bases,
+    # bh16bn128 has an extra K reg base for the 128-wide K, bh16bn64 has the bh16 warp layout at 64-wide K).
+    if REGIME == 'bh64':
+        linear_v: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)),
+            lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)),
+            warp_bases=((0, 0), (0, 0)),
+            block_bases=[],
+            shape=[512, 64],
+        )
+    elif REGIME == 'bh16bn128':
+        linear_v: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (0, 64), (64, 0), (128, 0), (256, 0)),
+            lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)),
+            warp_bases=((16, 0), (32, 0)),
+            block_bases=[],
+            shape=[512, 128],
+        )
+    else:
+        linear_v: gl.constexpr = gl.DistributedLinearLayout(
+            reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (64, 0), (128, 0), (256, 0)),
+            lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)),
+            warp_bases=((16, 0), (32, 0)),
+            block_bases=[],
+            shape=[512, 64],
+        )
+
+    mfma_layout_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_layout, k_width=8)
+    mfma_layout_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_layout, k_width=8)
+    dtype = Q_nope.type.element_ty
+    kvtype = Kv_c_cache.type.element_ty
+    ######### layout setting end #########
+
+    buf_q_nope = gl.allocate_shared_memory(dtype, shape=[BLOCK_H, HEAD_DIM_CKV], layout=shared_q_nope)
+    if HAS_PE:
+        buf_q_pe = gl.allocate_shared_memory(dtype, shape=[BLOCK_H, HEAD_DIM_KPE], layout=shared_q_pe)
+
+    # load q_nope
+    offs_d_ckv = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(0, blocked_q_nope))
+    cur_head = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blocked_q_nope))
+    offs_q_nope = cur_batch * stride_q_nope_bs + cur_head[:, None] * stride_q_nope_h + offs_d_ckv[None, :]
+    ### For nhead < BLOCK_H, mask OOB heads to zero on Q load and skip OOB O stores; wasted MFMA lanes are free (memory-bound).
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(buf_q_nope, Q_nope, offs_q_nope, mask = (cur_head < NHEAD)[:, None] if NHEAD < BLOCK_H else None)
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # load q_pe
+    if HAS_PE:
+        offs_d_kpe = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(0, blocked_q_pe))
+        cur_head_qpe = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blocked_q_pe))
+        offs_q_pe = cur_batch * stride_q_pe_bs + cur_head_qpe[:, None] * stride_q_pe_h + offs_d_kpe[None, :]
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(buf_q_pe, Q_pe, offs_q_pe, mask = (cur_head_qpe < NHEAD)[:, None] if NHEAD < BLOCK_H else None)
+        gl.amd.cdna4.async_copy.commit_group()
+
+    e_max = gl.zeros([BLOCK_H], dtype=gl.float32, layout=gl.SliceLayout(1, mfma_layout)) - float("inf")
+    e_sum = gl.zeros([BLOCK_H], dtype=gl.float32, layout=gl.SliceLayout(1, mfma_layout))
+    acc = gl.zeros([BLOCK_H, HEAD_DIM_CKV], dtype=gl.float32, layout=mfma_layout)
+
+    # Fold KV dequant scale into the QK temperature. For fp8 KV the real
+    # logits are (Q @ K_fp8^T) * kv_scale * sm_scale; softmax is shift- but
+    # not scale-invariant, so kv_scale must affect qk (not just acc).
+    # For bf16 KV the wrapper passes kv_scale=1.0, so this is a no-op.
+    qk_scale = sm_scale * kv_scale
+
+    ### bufs of page_number
+    shared_page: gl.constexpr = gl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0])
+    bufs_page = gl.allocate_shared_memory(gl.int32, shape=[2, BLOCK_N], layout=shared_page)
+    gl.static_assert(PAGE_SIZE == 1)
+
+    offs_page_raw = gl.arange(0, BLOCK_N, layout=blocked_page)
+
+    ################ prologue
+    #### global load page number
+    offs_n_page = start_n + offs_page_raw
+    offs_page = batch_page_start + offs_n_page // PAGE_SIZE
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(0), Req_to_tokens, offs_page, offs_n_page < split_kv_end)
+    gl.amd.cdna4.async_copy.commit_group()
+
+    start_n += BLOCK_N
+    #### global load page number
+    offs_n_page = start_n + offs_page_raw
+    offs_page = batch_page_start + offs_n_page // PAGE_SIZE
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(1), Req_to_tokens, offs_page, offs_n_page < split_kv_end)
+    gl.amd.cdna4.async_copy.commit_group()
+
+    #### local load Q
+    gl.amd.cdna4.async_copy.wait_group(2)
+    q_nope = gl.amd.cdna4.async_copy.load_shared_relaxed(buf_q_nope, mfma_layout_a)
+    if HAS_PE:
+        q_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(buf_q_pe, mfma_layout_a)
+
+    #################### move here to work around allocate_shared_memory bug
+    bufs_kv = gl.allocate_shared_memory(kvtype, shape=[2, HEAD_DIM_CKV, BLOCK_N], layout=shared_kv)
+    if HAS_PE:
+        bufs_kpe = gl.allocate_shared_memory(kvtype, shape=[2, HEAD_DIM_KPE, BLOCK_N], layout=shared_kpe)
+
+    #### global load K
+    # local load page number
+    gl.amd.cdna4.async_copy.wait_group(1)
+    if HAS_PE:
+        kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(0), gl.SliceLayout(0, blocked_kpe))
+        # simplify for page_size 1
+        kv_loc_pe = kv_page_number_pe
+
+    # local load page number for slice 0
+    bufs_page_0 = bufs_page.index(0).slice(0, BLOCK_N // 2, 0)
+    kv_page_number_0 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_0, gl.SliceLayout(0, blocked_kv_slice))
+    kv_loc0 = kv_page_number_0
+
+    # global load K_nope slice 0
+    offs_n_nope0 = split_kv_start + gl.arange(0, BLOCK_N // 2, layout=gl.SliceLayout(0, blocked_kv_slice))
+    offs_d_ckv_10 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv_slice))
+    offs_k_c0 = kv_loc0[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None]
+    bufs_kv0 = bufs_kv.index(0).slice(0, BLOCK_N // 2, 1)
+    if WITHIN_2GB:
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv0, Kv_c_cache, offs_k_c0, mask=offs_n_nope0[None, :] < split_kv_end)
+    else:
+        gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv0, Kv_c_cache + offs_k_c0)
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # global load K_pe
+    if HAS_PE:
+        offs_n_pe0 = split_kv_start + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe))
+        offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe))
+        offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET
+        if WITHIN_2GB:
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(0), K_pe_cache, offs_k_pe, mask=offs_n_pe0[None, :] < split_kv_end)
+        else:
+            gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(0), K_pe_cache + offs_k_pe)
+        gl.amd.cdna4.async_copy.commit_group()
+
+    # local load page number for slice 1
+    bufs_page_1 = bufs_page.index(0).slice(BLOCK_N // 2, BLOCK_N // 2, 0)
+    kv_page_number_1 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_1, gl.SliceLayout(0, blocked_kv_slice))
+    kv_loc1 = kv_page_number_1
+
+    # global load K_nope slice 1
+    offs_n_nope1 = offs_n_nope0 + BLOCK_N // 2
+    bufs_kv1 = bufs_kv.index(0).slice(BLOCK_N // 2, BLOCK_N // 2, 1)
+    offs_k_c1 = kv_loc1[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None]
+    if WITHIN_2GB:
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv1, Kv_c_cache, offs_k_c1, mask=offs_n_nope1[None, :] < split_kv_end)
+    else:
+        gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv1, Kv_c_cache + offs_k_c1)
+    gl.amd.cdna4.async_copy.commit_group()
+
+    if REGIME == 'bh64':
+        # bh64 guarantees >= 3 iters/split; this constant-folds the
+        # `if num_iter >= 2` epilogue-1 guard below so its codegen is unchanged.
+        gl.assume(num_iter >= 3)
+    buf_idx = 0
+    ################ loop
+    for i in range(num_iter - 2):
+        async_idx = (buf_idx + 1) % 2
+
+        gl.amd.cdna4.async_copy.wait_group(0)
+        #### global load page number
+        offs_n_page = start_n + BLOCK_N + offs_page_raw
+        offs_page = batch_page_start + offs_n_page // PAGE_SIZE
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(buf_idx), Req_to_tokens, offs_page, offs_n_page < split_kv_end)
+        gl.amd.cdna4.async_copy.commit_group()
+
+        #### global load K
+        bufs_kv0 = bufs_kv.index(async_idx).slice(0, BLOCK_N // 2, 1)
+        bufs_kv1 = bufs_kv.index(async_idx).slice(BLOCK_N // 2, BLOCK_N // 2, 1)
+        # local load page number for slice 0
+        bufs_page_0 = bufs_page.index(async_idx).slice(0, BLOCK_N // 2, 0)
+        kv_page_number_0 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_0, gl.SliceLayout(0, blocked_kv_slice))
+        kv_loc0 = kv_page_number_0
+        # global load K_nope slice 0
+        offs_n_nope0 = start_n + gl.arange(0, BLOCK_N // 2, layout=gl.SliceLayout(0, blocked_kv_slice))
+        offs_d_ckv_10 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv_slice))
+        offs_k_c0 = kv_loc0[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None]
+        if WITHIN_2GB:
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv0, Kv_c_cache, offs_k_c0, mask=offs_n_nope0[None, :] < split_kv_end)
+        else:
+            # No mask needed on global_load path in the loop body: all
+            # iterations are guaranteed in-bounds by num_iter arithmetic.
+            # Only the epilogue uses mask + other=0 for the last
+            # potentially-partial block.
+            gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv0, Kv_c_cache + offs_k_c0)
+        gl.amd.cdna4.async_copy.commit_group()
+
+        # local load page_number_pe
+        if HAS_PE:
+            kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kpe))
+            kv_loc_pe = kv_page_number_pe
+            # global load K_pe
+            offs_n_pe = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe))
+            offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe))
+            offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET
+            if WITHIN_2GB:
+                gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache, offs_k_pe, mask=offs_n_pe[None, :] < split_kv_end)
+            else:
+                # No mask needed: loop iterations are in-bounds (see KV slice 0 comment).
+                gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache + offs_k_pe)
+            gl.amd.cdna4.async_copy.commit_group()
+
+        #### dot, softmax, dot (part0)
+        k_c = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kv.index(buf_idx), mfma_layout_b)
+        zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout)
+        qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros)
+        if HAS_PE:
+            k_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kpe.index(buf_idx), mfma_layout_b)
+            qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk)
+
+        # local load page number for slice 1
+        bufs_page_1 = bufs_page.index(async_idx).slice(BLOCK_N // 2, BLOCK_N // 2, 0)
+        kv_page_number_1 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_1, gl.SliceLayout(0, blocked_kv_slice))
+        kv_loc1 = kv_page_number_1
+        # global load K_nope slice 1
+        offs_n1 = offs_n_nope0 + BLOCK_N // 2
+        offs_k_c1 = kv_loc1[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None]
+        if WITHIN_2GB:
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv1, Kv_c_cache, offs_k_c1, mask=offs_n1[None, :] < split_kv_end)
+        else:
+            # No mask needed: loop iterations are in-bounds (see KV slice 0 comment).
+            gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv1, Kv_c_cache + offs_k_c1)
+        gl.amd.cdna4.async_copy.commit_group()
+
+        #### dot, softmax, dot (part1)
+        qk *= qk_scale
+        offs_n_qk = split_kv_start + i * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout))
+        qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf"))
+        n_e_max = gl.maximum(gl.max(qk, 1), e_max)
+        LOG2E: gl.constexpr = 1.4426950408889634
+        re_scale = gl.exp2((e_max - n_e_max) * LOG2E)
+        p = gl.exp2((qk - n_e_max[:, None]) * LOG2E)
+        e_sum = e_sum * re_scale + gl.sum(p, 1)
+        e_max = n_e_max
+        p = p.to(dtype)
+        p = gl.convert_layout(p, mfma_layout_a)
+        acc *= re_scale[:, None]
+        v_c = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kv.index(buf_idx), linear_v)
+        v_c = v_c.to(dtype)
+        v_c = gl.permute(v_c, [1, 0])
+        v_c = gl.convert_layout(v_c, mfma_layout_b)
+        acc = gl.amd.cdna4.mfma(p, v_c, acc)
+
+        start_n += BLOCK_N
+        buf_idx = (buf_idx + 1) % 2
+
+    LOG2E: gl.constexpr = 1.4426950408889634
+
+    ################ epilogue 1
+    # Skip when num_iter < 2 (possible for bh16bn64 / bh16bn128 in either mode).
+    # bh64 has gl.assume(num_iter >= 3) above so the compiler folds this branch
+    # out there; for the bh16 regimes it stays a runtime branch.
+    if num_iter >= 2:
+        async_idx = (buf_idx + 1) % 2
+
+        #### global load K
+        # local load page number
+        gl.amd.cdna4.async_copy.wait_group(3 if HAS_PE else 2)
+        kv_page_number = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kv))
+        kv_loc = kv_page_number
+        if HAS_PE:
+            kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kpe))
+            kv_loc_pe = kv_page_number_pe
+        # global load K_nope
+        offs_n_nope = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kv))
+        offs_d_ckv_1 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv))
+        offs_k_c = kv_loc[None, :] * stride_kv_c_bs + offs_d_ckv_1[:, None]
+        if WITHIN_2GB:
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv.index(async_idx), Kv_c_cache, offs_k_c, mask=offs_n_nope[None, :] < split_kv_end)
+        else:
+            # No mask needed: out-of-range positions are discarded by the qk score mask
+            gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv.index(async_idx), Kv_c_cache + offs_k_c)
+        gl.amd.cdna4.async_copy.commit_group()
+        # global load K_pe
+        if HAS_PE:
+            offs_n_pe = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe))
+            offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe))
+            offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET
+            if WITHIN_2GB:
+                gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache, offs_k_pe, mask=offs_n_pe[None, :] < split_kv_end)
+            else:
+                gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache + offs_k_pe)
+            gl.amd.cdna4.async_copy.commit_group()
+
+        # dot, softmax, dot
+        gl.amd.cdna4.async_copy.wait_group(2 if HAS_PE else 1)
+        k_c = bufs_kv.index(buf_idx).load(layout=mfma_layout_b)
+        zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout)
+        qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros)
+
+        if HAS_PE:
+            k_pe = bufs_kpe.index(buf_idx).load(layout=mfma_layout_b)
+            qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk)
+        qk *= qk_scale
+        offs_n_qk = split_kv_start + (num_iter - 2) * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout))
+        qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf"))
+        n_e_max = gl.maximum(gl.max(qk, 1), e_max)
+        re_scale = gl.exp2((e_max - n_e_max) * LOG2E)
+        p = gl.exp2((qk - n_e_max[:, None]) * LOG2E)
+        e_sum = e_sum * re_scale + gl.sum(p, 1)
+        e_max = n_e_max
+        p = p.to(dtype)
+        p = gl.convert_layout(p, mfma_layout_a)
+        acc *= re_scale[:, None]
+        v_c = bufs_kv.index(buf_idx).load(layout=linear_v)
+        v_c = v_c.to(dtype)
+        v_c = gl.permute(v_c, [1, 0])
+        v_c = gl.convert_layout(v_c, mfma_layout_b)
+        acc = gl.amd.cdna4.mfma(p, v_c, acc)
+
+        start_n += BLOCK_N
+        buf_idx = (buf_idx + 1) % 2
+
+    ################ epilogue 2
+    #### dot, softmax, dot
+    gl.amd.cdna4.async_copy.wait_group(0)
+    k_c = bufs_kv.index(buf_idx).load(layout=mfma_layout_b)
+    zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout)
+    qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros)
+
+    if HAS_PE:
+        k_pe = bufs_kpe.index(buf_idx).load(layout=mfma_layout_b)
+        qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk)
+    qk *= qk_scale
+    offs_n_qk = split_kv_start + (num_iter - 1) * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout))
+    qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf"))
+    n_e_max = gl.maximum(gl.max(qk, 1), e_max)
+    re_scale = gl.exp2((e_max - n_e_max) * LOG2E)
+    p = gl.exp2((qk - n_e_max[:, None]) * LOG2E)
+    e_sum = e_sum * re_scale + gl.sum(p, 1)
+    e_max = n_e_max
+    p = p.to(dtype)
+    p = gl.convert_layout(p, mfma_layout_a)
+    acc *= re_scale[:, None]
+    v_c = bufs_kv.index(buf_idx).load(layout=linear_v)
+    v_c = v_c.to(dtype)
+    v_c = gl.permute(v_c, [1, 0])
+    v_c = gl.convert_layout(v_c, mfma_layout_b)
+    acc = gl.amd.cdna4.mfma(p, v_c, acc)
+
+    cur_head_o = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_layout))
+    offs_d_ckv_o = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(0, mfma_layout))
+    offs_o = cur_batch * stride_o_b + cur_head_o[:, None] * stride_o_h + split_kv_id * stride_o_s + offs_d_ckv_o[None, :]
+
+    if HAS_ATTN_SINK:
+        # Fold the optional per-head sink into the softmax denom (no V contribution).
+        # e_max/e_sum are natural-log units (the *LOG2E is inside exp2), so is sink.
+        if NHEAD < BLOCK_H:
+            sink = gl.load(Attn_sink + cur_head_o, mask=cur_head_o < NHEAD, other=float("-inf")).to(gl.float32)
+        else:
+            sink = gl.load(Attn_sink + cur_head_o).to(gl.float32)
+        n_e_max = gl.maximum(e_max, sink)
+        re_scale = gl.exp2((e_max - n_e_max) * LOG2E)
+        acc *= re_scale[:, None]
+        e_sum = e_sum * re_scale + gl.exp2((sink - n_e_max) * LOG2E)
+        e_max = n_e_max
+
+    acc *= kv_scale
+    rcp = 1.0 / e_sum
+    stored_value = (acc * rcp[:, None]).to(dtype)
+    if NHEAD < BLOCK_H:
+        gl.amd.cdna4.buffer_store(stored_value, ptr=O, offsets=offs_o, mask=(cur_head_o < NHEAD)[:, None])
+    else:
+        gl.amd.cdna4.buffer_store(stored_value, ptr=O, offsets=offs_o)
+
+    ### store lse
+    blocked_lse: gl.constexpr = gl.BlockedLayout(size_per_thread=[1], threads_per_warp=[64], warps_per_cta=[4], order=[0])
+    cur_head_lse = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=blocked_lse)
+    if RETURN_LSE and NUM_KV_SPLITS == 1:
+        # split==1: single split is the whole sequence, so its lse is the final lse.
+        offs_final_lse = cur_batch * stride_final_lse_b + cur_head_lse * stride_final_lse_h
+        lse = e_max + gl.log(e_sum)
+        lse = gl.convert_layout(lse, blocked_lse)
+        if NHEAD < BLOCK_H:
+            gl.amd.cdna4.buffer_store(lse, ptr=Final_lse, offsets=offs_final_lse, mask=(cur_head_lse < NHEAD))
+        else:
+            gl.amd.cdna4.buffer_store(lse, ptr=Final_lse, offsets=offs_final_lse)
+    elif NUM_KV_SPLITS > 1:
+        # per-split lse for stage-2 reduce.
+        offs_mid_lse = cur_batch * stride_mid_lse_b + cur_head_lse * stride_mid_lse_h + split_kv_id * stride_mid_lse_s
+        lse = e_max + gl.log(e_sum)
+        lse = gl.convert_layout(lse, blocked_lse)
+        if NHEAD < BLOCK_H:
+            gl.amd.cdna4.buffer_store(lse, ptr=Mid_lse, offsets=offs_mid_lse, mask=(cur_head_lse < NHEAD))
+        else:
+            gl.amd.cdna4.buffer_store(lse, ptr=Mid_lse, offsets=offs_mid_lse)
+# fmt: on
+
+
+# fmt: off
+@triton.jit
+def _mla_softmax_reducev_kernel(
+    Logits,
+    Mid_lse,
+    O,  # noqa: E741
+    Final_lse,
+    B_seq_len,  # same seq_info as the decode kernel to derive empty kv splits
+    stride_l_b,
+    stride_l_h,
+    stride_l_s,
+    stride_ml_b,
+    stride_ml_h,
+    stride_ml_s,
+    stride_o_b,
+    stride_o_h,
+    stride_fl_b,
+    stride_fl_h,
+    NUM_KV_SPLITS: tl.constexpr,
+    HEAD_DIM_CKV: tl.constexpr,
+    HAS_FINAL_LSE: tl.constexpr,
+    USE_2D_VIEW: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+
+    # Recompute this batch's seq len exactly as the decode kernel did, so we can
+    # rederive which splits are empty. Stage-1 early-returns on empty splits
+    # (num_iter == 0) and writes nothing, so their logits_buf / mid_lse slots hold
+    # raw, uninitialised memory - they cannot be loaded or reduced.
+    if USE_2D_VIEW:
+        cur_batch_seq_len = tl.load(B_seq_len + cur_batch)
+    else:
+        batch_page_start = tl.load(B_seq_len + cur_batch)
+        cur_batch_seq_len = tl.load(B_seq_len + cur_batch + 1) - batch_page_start
+    kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS
+
+    offs_d_ckv = tl.arange(0, HEAD_DIM_CKV)
+    offs_l = cur_batch * stride_l_b + cur_head * stride_l_h + offs_d_ckv
+    offs_ml = cur_batch * stride_ml_b + cur_head * stride_ml_h
+
+    e_sum = 0.0
+    e_max = -float("inf")
+    acc = tl.zeros([HEAD_DIM_CKV], dtype=tl.float32)
+
+    LOOP_START = NUM_KV_SPLITS - 1 if kv_len_per_split == 0 else 0
+    for split_kv_id in range(LOOP_START, NUM_KV_SPLITS):
+        logits = tl.load(Logits + offs_l + split_kv_id * stride_l_s)
+        logits_1 = tl.load(Mid_lse + offs_ml + split_kv_id * stride_ml_s)
+
+        n_e_max = tl.maximum(logits_1, e_max)
+        old_scale = tl.where(e_max == -float("inf"), 0.0, tl.exp(e_max - n_e_max))
+        acc *= old_scale
+        exp_logic = tl.where(logits_1 == -float("inf"), 0.0, tl.exp(logits_1 - n_e_max))
+        acc += exp_logic * logits
+
+        e_sum = e_sum * old_scale + exp_logic
+        e_max = n_e_max
+
+    out = acc / e_sum if e_sum > 0.0 else tl.zeros([HEAD_DIM_CKV], dtype=tl.float32)
+    tl.store(
+        O + cur_batch * stride_o_b + cur_head * stride_o_h + offs_d_ckv,
+        out,
+    )
+    if HAS_FINAL_LSE:
+        tl.store(
+            Final_lse + cur_batch * stride_fl_b + cur_head * stride_fl_h,
+            e_max + tl.log(e_sum),
+        )
+# fmt: on
+
+
+def mla_gluon(
+    q_nope,  # [batch, nhead, kv_lora_rank]
+    q_pe,  # [batch, nhead, qk_rope_head_dim]
+    # Shared: kv_c=[N, kv_lora_rank+qk_rope_head_dim], k_pe=None,              kv_pe_offset=kv_lora_rank
+    # Split:  kv_c=[N, kv_lora_rank],                k_pe=[N,qk_rope_head_dim], kv_pe_offset=0
+    kv_c,
+    # final output [batch, nhead, kv_lora_rank].
+    o,
+    page_table,  # 2D: block_table [batch, max_seqlen] | 1D: kv_indices [total_kv]
+    seq_info,  # 2D: cache_seqlens [batch]           | 1D: kv_indptr [batch+1]
+    sm_scale,
+    k_pe=None,
+    kv_pe_offset=512,
+    use_2d_view=True,
+    kv_scale=1.0,
+    min_kv_seq_len=1,
+    return_lse=False,
+    has_pe=True,
+    attn_sink=None,  # [nhead] fp32 per-head sink bias, None means no sink
+):
+    """Unified Gluon MLA entry (gfx950 / CDNA4) — decode and DeepSeek V4 sparse prefill.
+
+    `mla_gluon` supports the full decode (stage-1 + stage-2 reduce, or the stage-1-only
+    fast path when NUM_KV_SPLITS==1) and writes the final attention into the
+    caller's `o` ([batch, nhead, kv_lora_rank]).
+
+    return_lse=False (default): returns (o, None).
+
+    return_lse=True: additionally returns the merged log-sum-exp, a separate
+        fp32 tensor [batch, nhead]
+
+    DSv4 Sparse prefill packs NoPE and RoPE in to one contiguous row (448+64).
+    To run DSv4 prefill, it requires has_pe=False, prepares valid Q / K in q_nope / kv_c,
+    and attn_sink, q_pe / k_pe are unused placeholders.
+    """
+    if k_pe is None:
+        k_pe = kv_c
+
+    batch_size, nhead, head_dim_ckv = q_nope.shape
+    # Decode carries a real q_pe [.., 64];
+    # DSV4 prefill (HAS_PE=False) has no PE, so q_pe may be None and RoPE head_dim is the fixed 64.
+    head_dim_kpe = q_pe.shape[-1] if has_pe else 64
+    if not has_pe:
+        q_pe = q_nope
+        k_pe = kv_c
+        kv_pe_offset = 0
+        use_2d_view = False
+
+    assert arch_info.get_arch() == "gfx950", f"mla_gluon requires gfx950 (CDNA4), got {arch_info.get_arch()}"
+    assert head_dim_ckv == 512, f"mla_gluon requires head_dim_ckv=512, got {head_dim_ckv}"
+    assert head_dim_kpe == 64, f"mla_gluon requires head_dim_kpe=64, got {head_dim_kpe}"
+
+    # attn sink: decode never sets one; prefill may pass a per-head [H] bias.
+    has_attn_sink = attn_sink is not None
+    if attn_sink is None:
+        attn_sink = torch.empty(1, device=o.device, dtype=torch.float32)  # dummy ptr
+
+    # Pick regime by (nhead, kv dtype).
+    if nhead in (64, 128):
+        REGIME = "bh64"
+    elif 1 <= nhead <= 16:
+        if kv_c.dtype == torch.bfloat16:
+            REGIME = "bh16bn64"
+        elif kv_c.dtype == torch.float8_e4m3fn:  # gfx950 fp8 (e4m3fn, not e4m3fnuz)
+            REGIME = "bh16bn128"
+        else:
+            raise AssertionError(
+                f"mla_gluon[bh16*] requires kv_c.dtype in (bfloat16, float8_e4m3fn), got {kv_c.dtype}"
+            )
+    else:
+        raise AssertionError(
+            f"mla_gluon requires nhead <= 16 [bh16bn128/bh16bn64] or nhead in (64,128) [bh64], got {nhead}"
+        )
+
+    PAGE_SIZE = 1
+
+    if REGIME == "bh64":
+        BLOCK_H, BLOCK_N = 64, 64
+        NUM_XCDS = get_num_xcds()
+        # Auto-pick NUM_KV_SPLITS so the launch fills ~256 workgroups (one wave on
+        # MI350). For the supported (batch, nhead) matrix the result is in {1, 2, 4}.
+        base_grid = NUM_XCDS * triton.cdiv(nhead, BLOCK_H) * (batch_size // NUM_XCDS)
+        NUM_KV_SPLITS = max(1, triton.next_power_of_2(triton.cdiv(256, base_grid)))
+
+        assert batch_size % 64 == 0, f"mla_gluon[bh64] requires batch_size divisible by 64, got {batch_size}"
+        # gl.assume(num_iter > 3) inside the kernel requires every split to have
+        # > 3*BLOCK_N tokens. Smallest split (last) for batch length s is
+        # s - (k-1)*ceil(s/k); a sufficient bound is min_kv_seq_len > k*(3*BLOCK_N + k).
+        min_kv_seq_len_required = NUM_KV_SPLITS * (3 * BLOCK_N + NUM_KV_SPLITS)
+        assert (
+            min_kv_seq_len > min_kv_seq_len_required
+        ), f"mla_gluon[bh64] requires min_kv_seq_len > {min_kv_seq_len_required} (NUM_KV_SPLITS={NUM_KV_SPLITS}), got {min_kv_seq_len}"
+        assert (
+            q_nope.dtype == torch.bfloat16 and q_pe.dtype == torch.bfloat16
+        ), f"q_nope/q_pe must be bf16, got {q_nope.dtype}/{q_pe.dtype}"
+        assert (
+            kv_c.dtype == torch.bfloat16 and k_pe.dtype == torch.bfloat16
+        ), f"kv_c/k_pe must be bf16, got {kv_c.dtype}/{k_pe.dtype}"
+    else:  # bh16bn128 (fp8 KV) or bh16bn64 (bf16 KV)
+        BLOCK_H = 16
+        BLOCK_N = 128 if REGIME == "bh16bn128" else 64
+        kv_dtype = torch.float8_e4m3fn if REGIME == "bh16bn128" else torch.bfloat16
+        NUM_XCDS = 1  # unused by 2-D split grid mapping
+        # 2-D grid (batch, split). Both bh16 regimes support num_iter in {1, 2, ...}
+        # (no gl.assume(num_iter >= 3) in the kernel); the only correctness need is
+        # that every split is non-empty (floor split size = min_kv_seq_len //
+        # NUM_KV_SPLITS >= 1). Each clamp below keeps NUM_KV_SPLITS <= min_kv_seq_len,
+        if REGIME == "bh16bn128":
+            assert batch_size == 1, f"mla_gluon[bh16bn128] requires batch_size=1, got {batch_size}"
+            NUM_KV_SPLITS = max(1, min(256 // batch_size, min_kv_seq_len))
+        else:  # bh16bn64
+            # Fill ~256 WGs (total WGs = B * NUM_KV_SPLITS <= 256, one MI350 wave),
+            # but never split a sequence into more blocks than it has: bound by the
+            # shortest seq's block count so every split holds >= 1 block (no wasted
+            # partial-block MFMA). For min_kv_seq_len <= BLOCK_N this collapses to
+            # NUM_KV_SPLITS=1, i.e. one WG per batch computing the whole (short) seq.
+            NUM_KV_SPLITS = max(1, min(256 // batch_size, triton.cdiv(min_kv_seq_len, BLOCK_N)))
+        assert (
+            q_nope.dtype == torch.bfloat16 and q_pe.dtype == torch.bfloat16
+        ), f"q_nope/q_pe must be bf16, got {q_nope.dtype}/{q_pe.dtype}"
+        assert (
+            kv_c.dtype == kv_dtype and k_pe.dtype == kv_dtype
+        ), f"kv_c/k_pe must be {kv_dtype}, got {kv_c.dtype}/{k_pe.dtype}"
+
+    # buffer_load uses scalar base + 32-bit offsets, limiting addressable range.
+    # For KV caches > 2 GB the kernel falls back to global_load (64-bit pointers).
+    max_kv_bytes = kv_c.shape[0] * kv_c.stride(0) * kv_c.element_size()
+    within_2gb = max_kv_bytes <= 0x80000000  # 2 GB
+
+    if NUM_KV_SPLITS == 1:
+        # Fast path: stage-1 writes the final attention (and lse) directly to o.
+        logits_buf = o.view(batch_size, nhead, NUM_KV_SPLITS, head_dim_ckv)
+        mid_lse = None
+        stride_mid_lse_b, stride_mid_lse_h, stride_mid_lse_s = 0, 0, 0
+    else:
+        # stage-1 -> per-split (acc, lse); stage-2 reduces into o.
+        logits_buf = torch.empty(
+            (batch_size, nhead, NUM_KV_SPLITS, head_dim_ckv),
+            dtype=o.dtype,
+            device=o.device,
+        )
+        mid_lse = torch.empty(
+            (batch_size, nhead, NUM_KV_SPLITS),
+            dtype=torch.float32,
+            device=o.device,
+        )
+        stride_mid_lse_b, stride_mid_lse_h, stride_mid_lse_s = mid_lse.stride()
+
+    if return_lse:
+        final_lse = torch.empty((batch_size, nhead), dtype=torch.float32, device=q_nope.device)
+        stride_final_lse_b, stride_final_lse_h = final_lse.stride()
+    else:
+        final_lse = None
+        stride_final_lse_b, stride_final_lse_h = 0, 0
+
+    if REGIME == "bh64":
+        grid = (
+            NUM_XCDS,
+            triton.cdiv(nhead, BLOCK_H),
+            (batch_size // NUM_XCDS) * NUM_KV_SPLITS,
+        )
+    else:
+        grid = (batch_size, NUM_KV_SPLITS)
+    stride_page_bs = page_table.stride(0) if use_2d_view else 0
+
+    _mla_gluon[grid](
+        q_nope,
+        q_pe,
+        kv_c,
+        k_pe,
+        page_table,
+        seq_info,
+        logits_buf,
+        attn_sink,
+        sm_scale,
+        kv_scale,
+        q_nope.stride(0),
+        q_nope.stride(1),
+        q_pe.stride(0),
+        q_pe.stride(1),
+        kv_c.stride(-2),
+        k_pe.stride(-2),
+        stride_page_bs,
+        logits_buf.stride(0),
+        logits_buf.stride(1),
+        logits_buf.stride(2),
+        mid_lse,
+        stride_mid_lse_b,
+        stride_mid_lse_h,
+        stride_mid_lse_s,
+        final_lse,
+        stride_final_lse_b,
+        stride_final_lse_h,
+        BLOCK_H=BLOCK_H,
+        BLOCK_N=BLOCK_N,
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        PAGE_SIZE=PAGE_SIZE,
+        HEAD_DIM_CKV=head_dim_ckv,
+        HEAD_DIM_KPE=head_dim_kpe,
+        KV_PE_OFFSET=kv_pe_offset,
+        USE_2D_VIEW=use_2d_view,
+        WITHIN_2GB=within_2gb,
+        NUM_XCDS=NUM_XCDS,
+        NHEAD=nhead,
+        REGIME=REGIME,
+        RETURN_LSE=return_lse,
+        HAS_PE=has_pe,
+        HAS_ATTN_SINK=has_attn_sink,
+    )
+
+    if NUM_KV_SPLITS == 1:
+        # Fast path: stage-1 already wrote o (and lse) directly.
+        return o, final_lse
+
+    # Stage-2: reduce per-split (acc, lse) into o (and lse when return_lse).
+    grid_reduce = (batch_size, nhead)
+    _mla_softmax_reducev_kernel[grid_reduce](
+        logits_buf,
+        mid_lse,
+        o,
+        final_lse,
+        seq_info,
+        logits_buf.stride(0),
+        logits_buf.stride(1),
+        logits_buf.stride(2),
+        stride_mid_lse_b,
+        stride_mid_lse_h,
+        stride_mid_lse_s,
+        o.stride(0),
+        o.stride(1),
+        stride_final_lse_b,
+        stride_final_lse_h,
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        HEAD_DIM_CKV=head_dim_ckv,
+        HAS_FINAL_LSE=return_lse,
+        USE_2D_VIEW=use_2d_view,
+        num_warps=8,
+    )
+
+    return o, final_lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py
new file mode 100644
index 000000000..f3bf49f25
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py
@@ -0,0 +1,247 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton:
+the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`);
+this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing
+the external transpose copy.
+
+Per program: 1 query token. Grid: (total_tokens,).
+Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups:
+    dKV_lora[D_V, TILE_K]   = sum_hg ( Q_lora_T @ dS + dO_T @ P )    # contract over heads
+    dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS )
+    store interm[token, rank, :D_QK]
+
+Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in
+LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout
+-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dkv_interm_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16 (UNtransposed)
+    dO_ptr,  # [T, H, D_V]  bf16 (UNtransposed)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    Interm_ptr,  # [T, R_CHUNK, D_QK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    stride_interm_t: tl.int64,
+    stride_interm_r: tl.int64,
+    num_heads: tl.int32,
+    R_CHUNK: gl.constexpr,
+    TILE_K: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    NUM_HG: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads.
+    mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for HBM loads ----
+    # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read.
+    _q_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _q_tpw_m: gl.constexpr = 64 // _q_tpw_k
+    blk_q: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_q_tpw_m, _q_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose).
+    blk_ds: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 4],
+        threads_per_warp=[16, 4],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ----
+    sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0])
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0])
+
+    # ---- Dot operand layouts ----
+    dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8)
+    dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+    dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8)
+
+    token_idx = gl.program_id(axis=0)
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ---- LDS for Q/dO/Q_rope (single-buffered) ----
+    smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q)
+    smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do)
+    if HAS_ROPE:
+        smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    ds_base = token_idx.to(tl.int64) * stride_ds_t
+    interm_base = token_idx.to(tl.int64) * stride_interm_t
+
+    # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d]
+    offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma))
+    offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma))
+    offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma))
+
+    for t in range(NUM_TILES):
+        dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma)
+        if HAS_ROPE:
+            dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma)
+
+        for hg in range(NUM_HG):
+            hg_off = hg * BLOCK_H
+
+            # ---- stage Q/dO/Q_rope (this head group) into LDS ----
+            offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q))
+            offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q))
+            mask_h_q = offs_h_q < num_heads
+            q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64)
+            do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None]
+            )
+
+            if HAS_ROPE:
+                offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+                offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+                mask_h_qr = offs_h_qr < num_heads
+                qr_offs = (
+                    q_base
+                    + offs_h_qr[:, None].to(tl.int64) * stride_q_h
+                    + (D_V + offs_r_qr[None, :]).to(tl.int64)
+                )
+                gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                    dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None]
+                )
+            gl.amd.cdna4.async_copy.commit_group()
+
+            # ---- load dS / P (this tile, this head group) -> dot operands ----
+            offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds))
+            offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds))
+            mask_h_ds = offs_h_ds < num_heads
+            dsp_offs = (
+                ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64)
+            )
+            dS_blk = gl.amd.cdna4.buffer_load(
+                ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            P_blk = gl.amd.cdna4.buffer_load(
+                ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0
+            )
+            dS_dot = gl.convert_layout(dS_blk, dot_ds_b)
+            P_dot = gl.convert_layout(P_blk, dot_p_b)
+
+            # ---- wait + transpose-read Q/dO/Q_rope ----
+            gl.amd.cdna4.async_copy.wait_group(0)
+            Q_T = smem_q.permute([1, 0]).load(dot_qT_a)  # [D_V, BLOCK_H]
+            dO_T = smem_do.permute([1, 0]).load(dot_doT_a)  # [D_V, BLOCK_H]
+
+            dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora)
+            dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora)
+            if HAS_ROPE:
+                Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a)  # [D_ROPE, BLOCK_H]
+                dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope)
+
+        # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ----
+        col = t * TILE_K + offs_col_st
+        interm_lora_offs = (
+            interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64)
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty),
+            ptr=Interm_ptr,
+            offsets=interm_lora_offs.to(tl.int32),
+        )
+        # dKV_rope is provably zero for the V4 zero-rope-pad (discarded downstream by the
+        # gather/adapter, which use dkv[..., :D_V]); skip its compute + store when HAS_ROPE=False.
+        if HAS_ROPE:
+            interm_rope_offs = (
+                interm_base
+                + col[None, :].to(tl.int64) * stride_interm_r
+                + (D_V + offs_dr_st[:, None]).to(tl.int64)
+            )
+            gl.amd.cdna4.buffer_store(
+                stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty),
+                ptr=Interm_ptr,
+                offsets=interm_rope_offs.to(tl.int32),
+            )
+
+
+def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64):
+    """
+    Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel).
+
+    Returns interm [T, R_CHUNK, D_QK] bf16.
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    assert R_CHUNK % TILE_K == 0
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+        q,
+        do,
+        chunk_dS,
+        chunk_P,
+        interm,
+        q.stride(0),
+        q.stride(1),
+        do.stride(0),
+        do.stride(1),
+        chunk_dS.stride(0),
+        chunk_dS.stride(1),
+        interm.stride(0),
+        interm.stride(1),
+        num_heads,
+        R_CHUNK=R_CHUNK,
+        TILE_K=TILE_K,
+        BLOCK_H=BLOCK_H,
+        NUM_HG=num_hg,
+        D_V=kv_lora_rank,
+        D_ROPE=rope_rank,
+        num_warps=4,
+    )
+    return interm
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py
new file mode 100644
index 000000000..8ad3853ce
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py
@@ -0,0 +1,533 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X).
+
+M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the
+gluon hardware-control structure of Leon's V3.2 forward.
+
+Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK).
+Grid: (total_tokens, cdiv(num_heads, BLOCK_H)).  Dispatch loops chunks externally,
+dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits).
+
+Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K):
+    S       = Q_lora @ K_lora_T + Q_rope @ K_rope_T      # 2 MMAs (mfma_s),  contract over D
+    P       = exp(S*scale - lse)                          # lse is sink-inclusive (from fwd)
+    dP      = dO @ K_lora_T                               # 1 MMA (mfma_s),  reuses K_lora_T_dot
+    dS      = P * (dP - delta) * scale
+    dQ_lora += dS @ K_lora                                # 1 MMA (mfma_acc), K_lora = K_lora_T.T view
+    dQ_rope += dS @ K_rope                                # 1 MMA (mfma_acc), K_rope = K_rope_T.T view
+    store dS, P chunk -> HBM (consumed by dKV-intermediate kernel)
+
+Differences vs Leon's fwd (the structural template):
+  * dO added as a second stationary [BH, D_V] operand (async-loaded with Q).
+  * No online softmax: single P = exp(S - lse) (lse precomputed by fwd).
+  * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways.
+  * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write.
+  * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher),
+    so this kernel needs no atomics. lse from fwd already folds the sink in.
+
+M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA
+k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950).
+"""
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+
+@gluon.jit
+def _sparse_mla_bwd_dq_gl_kernel(
+    Q_ptr,  # [T, H, D_QK] bf16
+    KV_ptr,  # [T, 1, D_QK] bf16
+    dO_ptr,  # [T, H, D_V]  bf16
+    TopK_ptr,  # [T, TOPK_padded] int32
+    LSE_ptr,  # [T, H] fp32 (sink-inclusive)
+    Delta_ptr,  # [T, H] fp32
+    dQ_ptr,  # [T, H, D_QK] bf16 (read-modify-write across chunks)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+    IS_FIRST_CHUNK: gl.constexpr,
+):
+    # ===================== constexpr layouts =====================
+    # mfma_s drives S = Q@K_T and dP = dO@K_T, both reducing D_V=512 -> K=32 halves their
+    # MFMA instruction count vs K=16 (mirrors the fwd R15 win). mfma_acc (dQ += dS@K) stays
+    # K=16 since its reduction dim is TILE_K (=16).
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 32],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # ---- Blocked layouts for global loads (per Leon's fwd) ----
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_V]  (Q_lora, dO, dQ_lora)
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H, D_ROPE]
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    # ---- Shared layouts (only K is staged in LDS) ----
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1])
+
+    # ---- Dot operand layouts ----
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+    dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ===================== program ids =====================
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    NUM_TILES: gl.constexpr = R_CHUNK // TILE_K
+
+    # ===================== Q / dO offsets =====================
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    if HAS_ROPE:
+        offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_qrope = offs_h_qrope < num_heads
+        q_offs_rope = (
+            q_base
+            + offs_h_qrope[:, None].to(tl.int64) * stride_q_h
+            + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+        )
+        q_mask_rope = mask_h_qrope[:, None]
+
+    offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_do = offs_h_do < num_heads
+    do_base = token_idx.to(tl.int64) * stride_do_t
+    do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64)
+    do_mask = mask_h_do[:, None]
+
+    # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) =====================
+    # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to
+    # the dot-operand layout once. The convert's LDS scratch is transient (freed before
+    # the K loop), unlike persistent staging, so only K occupies LDS during the loop.
+    q_lora_blk = gl.amd.cdna4.buffer_load(
+        ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0
+    )
+    do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0)
+    Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a)
+    dO_dot = gl.convert_layout(do_blk, dot_do_a)
+    if HAS_ROPE:
+        q_rope_blk = gl.amd.cdna4.buffer_load(
+            ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0
+        )
+        Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a)
+
+    # ===================== topk / KV offsets =====================
+    topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START
+
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # ===================== shared mem for K loop (double-buffered) =====================
+    if HAS_ROPE:
+        smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope)
+    smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora)
+
+    # ===================== dQ accumulators =====================
+    # Always zero-init; the read-modify-write across chunks is folded in at STORE
+    # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout
+    # whose LDS scratch would overflow the K/Q/dO buffers).
+    dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+    if HAS_ROPE:
+        dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc)
+
+    # ===================== lse / delta (in mfma_s row-slice layout) =====================
+    offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_s = offs_h_s < num_heads
+    lse = gl.amd.cdna4.buffer_load(
+        ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+    delta = gl.amd.cdna4.buffer_load(
+        ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0
+    )
+
+    # ===================== prologue: K tile 0 (group B, buffer 0) =====================
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=(topk_base + offs_tile_klora).to(tl.int32),
+        mask=offs_tile_klora < R_CHUNK,
+        other=-1,
+    )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1
+    )
+
+    valid_klora = topk_pos_klora != -1
+    valid_mma = topk_pos_mma != -1
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :]
+    )
+    if HAS_ROPE:
+        topk_pos_krope = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + offs_tile_krope).to(tl.int32),
+            mask=offs_tile_krope < R_CHUNK,
+            other=-1,
+        )
+        valid_krope = topk_pos_krope != -1
+        safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+        krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :]
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # dS / P store offsets in the mfma_s layout (store directly from the compute
+    # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead).
+    ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h
+    offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads
+
+    # ===================== main loop: prefetch t+1, compute t =====================
+    cur_buf = 0
+    for t in range(NUM_TILES - 1):
+        next_offs_klora = (t + 1) * TILE_K + offs_tile_klora
+        next_offs_krope = (t + 1) * TILE_K + offs_tile_krope
+        next_offs_mma = (t + 1) * TILE_K + offs_tile_mma
+
+        topk_pos_klora_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_klora).to(tl.int32),
+            mask=next_offs_klora < R_CHUNK,
+            other=-1,
+        )
+        topk_pos_mma_next = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=(topk_base + next_offs_mma).to(tl.int32),
+            mask=next_offs_mma < R_CHUNK,
+            other=-1,
+        )
+
+        valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1)
+        valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1)
+        safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0)
+
+        next_buf = 1 - cur_buf
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(next_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        if HAS_ROPE:
+            topk_pos_krope_next = gl.amd.cdna4.buffer_load(
+                ptr=TopK_ptr,
+                offsets=(topk_base + next_offs_krope).to(tl.int32),
+                mask=next_offs_krope < R_CHUNK,
+                other=-1,
+            )
+            valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1)
+            safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0)
+            krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+                D_V + offs_r_krope[:, None]
+            ).to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_krope.index(next_buf),
+                ptr=KV_ptr,
+                offsets=krope_offs_next.to(tl.int32),
+                mask=valid_krope_next[None, :],
+            )
+        gl.amd.cdna4.async_copy.commit_group()
+
+        gl.amd.cdna4.async_copy.wait_group(1)
+
+        # ----- read K views from current buffer -----
+        klora_smem_cur = smem_klora.index(cur_buf)
+        K_lora_T_dot = klora_smem_cur.load(dot_klora_b)  # [D_V, TILE_K]  opIdx1 mfma_s
+        K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)  # [TILE_K, D_V] opIdx1 mfma_acc
+
+        # ----- S = Q_lora@K_lora_T (+ Q_rope@K_rope_T when HAS_ROPE) -----
+        S = gl.amd.cdna4.mfma(
+            Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        if HAS_ROPE:
+            krope_smem_cur = smem_krope.index(cur_buf)
+            K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+            S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+        S = S * scale
+        offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+        S = gl.where(valid_mask, S, float("-inf"))
+
+        # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale -----
+        P = gl.exp(S - lse[:, None])
+        P = gl.where(valid_mask, P, 0.0)
+        dP = gl.amd.cdna4.mfma(
+            dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+        )
+        dS = P * (dP - delta[:, None]) * scale
+        dS = gl.where(valid_mask, dS, 0.0)
+
+        # ----- dQ_lora += dS@K_lora (+ dQ_rope += dS@K_rope when HAS_ROPE) -----
+        dS_bf = dS.to(KV_ptr.dtype.element_ty)
+        dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+        dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+        if HAS_ROPE:
+            K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+            dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+        # ----- store dS, P chunk -----
+        col = t * TILE_K + offs_tile_dsp
+        dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+        gl.amd.cdna4.buffer_store(
+            stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+        )
+        gl.amd.cdna4.buffer_store(
+            stored_value=P.to(KV_ptr.dtype.element_ty),
+            ptr=P_ptr,
+            offsets=dsp_offs.to(tl.int32),
+            mask=mask_h_dsp[:, None],
+        )
+
+        # promote prefetch -> current
+        cur_buf = next_buf
+        valid_mma = valid_mma_next
+
+    # ===================== epilogue: last tile =====================
+    gl.amd.cdna4.async_copy.wait_group(0)
+    t = NUM_TILES - 1
+    klora_smem_cur = smem_klora.index(cur_buf)
+    K_lora_T_dot = klora_smem_cur.load(dot_klora_b)
+    K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b)
+
+    S = gl.amd.cdna4.mfma(
+        Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)
+    )
+    if HAS_ROPE:
+        krope_smem_cur = smem_krope.index(cur_buf)
+        K_rope_T_dot = krope_smem_cur.load(dot_krope_b)
+        S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S)
+    S = S * scale
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None]
+    S = gl.where(valid_mask, S, float("-inf"))
+
+    P = gl.exp(S - lse[:, None])
+    P = gl.where(valid_mask, P, 0.0)
+    dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s))
+    dS = P * (dP - delta[:, None]) * scale
+    dS = gl.where(valid_mask, dS, 0.0)
+
+    dS_bf = dS.to(KV_ptr.dtype.element_ty)
+    dS_dot = gl.convert_layout(dS_bf, dot_ds_a)
+    dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora)
+    if HAS_ROPE:
+        K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b)
+        dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope)
+
+    col = t * TILE_K + offs_tile_dsp
+    dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None]
+    )
+    gl.amd.cdna4.buffer_store(
+        stored_value=P.to(KV_ptr.dtype.element_ty),
+        ptr=P_ptr,
+        offsets=dsp_offs.to(tl.int32),
+        mask=mask_h_dsp[:, None],
+    )
+
+    # ===================== store dQ (lora + rope) =====================
+    dq_base = token_idx.to(tl.int64) * stride_dq_t
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64)
+    dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora)
+    if not IS_FIRST_CHUNK:
+        prev_lora = gl.amd.cdna4.buffer_load(
+            ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0
+        )
+        dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+    gl.amd.cdna4.buffer_store(
+        stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None]
+    )
+
+    # dQ_rope is provably zero for the V4 zero-rope-pad (the adapter discards dq[..., D_V:]),
+    # so skip its accumulation + store entirely when HAS_ROPE is False.
+    if HAS_ROPE:
+        offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_or = offs_h_or < num_heads
+        dq_offs_rope = (
+            dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64)
+        )
+        dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope)
+        if not IS_FIRST_CHUNK:
+            prev_rope = gl.amd.cdna4.buffer_load(
+                ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0
+            )
+            dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty)
+        gl.amd.cdna4.buffer_store(
+            stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None]
+        )
+
+
+# =====================================================================
+# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P.
+# d_sink (if needed) is a torch reduction handled by the caller.
+# =====================================================================
+def sparse_mla_bwd_dq_gl(
+    q,
+    kv,
+    do,
+    topk_indices_padded,
+    lse,
+    delta,
+    R_CHUNK,
+    topk,
+    kv_lora_rank=512,
+    scale=None,
+    BLOCK_H=64,
+    TILE_K=16,
+):
+    """
+    Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop.
+
+    Returns:
+        dq:        [T, H, D_QK] bf16 (fully accumulated across chunks)
+        chunk_dS:  [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation)
+        chunk_P:   [T, H, R_CHUNK] bf16 (LAST chunk's P)
+    """
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK"
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+
+    num_hg = triton.cdiv(num_heads, BLOCK_H)
+    grid = (total_tokens, num_hg)
+
+    for r_start in range(0, topk, R_CHUNK):
+        is_first = r_start == 0
+        _sparse_mla_bwd_dq_gl_kernel[grid](
+            q,
+            kv,
+            do,
+            topk_indices_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_indices_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BLOCK_H,
+            TILE_K=TILE_K,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,
+        )
+    return dq, chunk_dS, chunk_P
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py
new file mode 100644
index 000000000..458adaf9c
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py
@@ -0,0 +1,177 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon DeepSeek-V4 sparse-MLA backward for the "gluon_v2" backend.
+
+Companion to the gluon_v2 forward (:func:`sparse_mla_fwd_v4_gluon_v2`). Wires the Gluon
+dQ + Gluon dKV-intermediate compute kernels with a Triton Delta preprocess and the
+backend-neutral CSR inverted-topk gather + torch d_sink reduction (non-atomic
+chunked-gather scheme).
+
+The dQ / dKV-intermediate Gluon kernels apply the forward campaign's accepted techniques
+to the backward: rope-skip (the V4 zero-rope-pad makes the rope gradients provably zero)
++ MFMA K=32 for the D_V=512-reduction matmuls, plus a single-chunk dQ read-modify-write
+for high head counts. Beats the plain-Triton backward ~1.12x geomean over the 6
+flash/pro x cr{0,4,128} shapes (eager-UT 9/9).
+"""
+
+import torch
+import triton
+
+from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc
+from .._gluon_dsa._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess
+from .dsa_bwd_dkv_interm_gluon import _sparse_mla_bwd_dkv_interm_gl_kernel
+from .dsa_bwd_dq_gluon import _sparse_mla_bwd_dq_gl_kernel
+
+
+def sparse_mla_bwd_v4_gluon_v2(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA backward (Gluon dQ/dKV). Returns ``(dq, dkv, d_sink)``."""
+    assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous()
+    assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    num_kv = kv.shape[0]
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,)
+
+    # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ----
+    delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+    BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads))
+    _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))](
+        O_ptr=o,
+        dO_ptr=do,
+        Delta_ptr=delta,
+        stride_o_t=o.stride(0),
+        stride_o_h=o.stride(1),
+        num_heads=num_heads,
+        D_V=kv_lora_rank,
+        BLOCK_H=BLOCK_H_PRE,
+    )
+
+    # ---- config ----
+    # R2: dQ is read-modify-written across chunks, so more chunks = more redundant dq
+    # reload passes + repeated CSR builds. H=64 CSA has TOPK=640, so the old 256 cap
+    # split it into 3 chunks; a 320 cap tests a two-chunk schedule without changing
+    # high-head whole-topk behavior.
+    if num_heads >= 128:
+        R_CHUNK = min(topk, 1536)
+    elif num_heads >= 64:
+        R_CHUNK = min(topk, 320)
+    else:
+        R_CHUNK = min(256, topk)
+    BH_DQ, TK_DQ = 64, 16
+    BH_DKV, TK_DKV = 32, 64
+    num_hg_dq = triton.cdiv(num_heads, BH_DQ)
+    num_hg_dkv = triton.cdiv(num_heads, BH_DKV)
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    # ---- pad topk to R_CHUNK multiple ----
+    topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK
+    if topk_padded_len != topk:
+        pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device)
+        topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous()
+    else:
+        topk_padded = topk_indices
+
+    all_csr = [
+        _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv)
+        for rs in range(0, topk, R_CHUNK)
+    ]
+
+    for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)):
+        is_first = r_start == 0
+
+        _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)](
+            q,
+            kv,
+            do,
+            topk_padded,
+            lse,
+            delta,
+            dq,
+            chunk_dS,
+            chunk_P,
+            q.stride(0),
+            q.stride(1),
+            kv.stride(0),
+            do.stride(0),
+            do.stride(1),
+            dq.stride(0),
+            dq.stride(1),
+            topk_padded.stride(0),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            scale,
+            num_heads,
+            r_start,
+            R_CHUNK=R_CHUNK,
+            BLOCK_H=BH_DQ,
+            TILE_K=TK_DQ,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_ROPE=False,  # V4 zero-rope-pad: dQ_rope is provably zero + discarded by the adapter
+            IS_FIRST_CHUNK=is_first,
+            num_warps=4,
+            waves_per_eu=1,
+        )
+
+        _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)](
+            q,
+            do,
+            chunk_dS,
+            chunk_P,
+            interm,
+            q.stride(0),
+            q.stride(1),
+            do.stride(0),
+            do.stride(1),
+            chunk_dS.stride(0),
+            chunk_dS.stride(1),
+            interm.stride(0),
+            interm.stride(1),
+            num_heads,
+            R_CHUNK=R_CHUNK,
+            TILE_K=TK_DKV,
+            BLOCK_H=BH_DKV,
+            NUM_HG=num_hg_dkv,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_ROPE=False,  # V4 zero-rope-pad: dKV_rope provably zero + discarded downstream
+            num_warps=4,
+        )
+
+        inv_ptr, inv_data = all_csr[chunk_idx]
+        _bwd_dkv_gather_acc[(num_kv,)](
+            interm,
+            inv_ptr,
+            inv_data,
+            dkv_acc,
+            interm.stride(1),
+            dkv_acc.stride(0),
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            num_warps=4,
+        )
+
+    d_sink = None
+    if has_sink:
+        d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0)
+
+    dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1)
+    return dq, dkv_out, d_sink
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py
new file mode 100644
index 000000000..dce02e511
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py
@@ -0,0 +1,743 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from ROCm/aiter
+# (https://github.com/ROCm/aiter), leonling-ll/aiter branch liyang/dsa
+# (ROCm/aiter#3456); see module docstring.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink.
+
+Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch
+`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware
+control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline,
+ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456.
+This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches
+`sparse_mla_fwd_v4`; it is the forward of the "gluon_v2" backend.
+
+Optimizations accepted over the base gluon forward (gfx950 campaign):
+  * rope-skip (HAS_ROPE=False): the V4 latent bakes RoPE in-place over the 512, so the
+    rope QK term is provably zero -- skip the 64-wide rope MFMA + K_rope loads entirely.
+  * exp2 softmax: fold log2(e) into the QK scale so the per-element exp is one hardware
+    exp2 (m_i/l_i in log2 units; LSE converted back to natural log for the backward).
+  * MFMA K=32 for the QK score matmul (instr_shape [16,16,32]): the score reduces D_V=512,
+    so K=32 halves the QK MFMA instruction count vs K=16 (PV/acc stays K=16, TILE_K-bound).
+  * register-prefetched topk index (off the KV-gather critical path) + async double-buffered
+    K gather overlapping the QK/softmax/PV MFMAs.
+
+Pipeline:
+  Prologue:  Q -> shared (async); K tile 0 -> shared (async, double-buffered); deep-prefetch topk.
+  Loop tile t: gather tile t+2 (async) while computing QK[t+1] + softmax(t) + PV(t); promote.
+  Epilogue:  drain; fold sink into the denominator (V4); write O, LSE.
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+from .aiter_lse_fwd import sparse_mla_fwd_v4_aiter_lse_csa_formula
+
+# ---------------------------------------------------------------------------
+# Triton capability gate: the gluon_v2 forward is a Gluon kernel (the backend's backward
+# is currently plain-Triton, being migrated to Gluon). The Gluon fwd needs a Gluon-capable triton whose CDNA4
+# async_copy accepts arbitrary (DistributedLinearLayout) offsets. Released
+# triton 3.7.0/3.7.1 still restricts async_copy offsets to Blocked/Slice and
+# will NOT compile this path; build triton from the commit below.
+# ---------------------------------------------------------------------------
+_GLUON_V2_REQUIRED_COMMIT = "09500db9f0"
+_GLUON_V2_INSTALL_HINT = (
+    "gluon_v2 forward (Gluon) requires a Gluon-capable triton whose CDNA4 "
+    "async_copy accepts general offset layouts. The installed triton ({ver}) does not (released "
+    "3.7.0/3.7.1 restrict async_copy offsets to BlockedLayout/SliceLayout).\n"
+    "Build & install triton-lang/triton @ commit " + _GLUON_V2_REQUIRED_COMMIT + ":\n"
+    "  git clone https://github.com/triton-lang/triton.git third_party/triton\n"
+    "  cd third_party/triton && git checkout " + _GLUON_V2_REQUIRED_COMMIT + "\n"
+    "  pip install -r python/requirements.txt\n"
+    "  TRITON_CODEGEN_BACKENDS=amd MAX_JOBS=128 pip wheel --no-build-isolation --no-deps . -w dist\n"
+    "  pip install --force-reinstall --no-deps dist/triton-*.whl"
+)
+
+
+@functools.lru_cache(maxsize=1)
+def _gluon_available() -> bool:
+    """True iff triton exposes the experimental CDNA4 Gluon dialect this fwd uses."""
+    try:
+        from triton.experimental import gluon as _gl  # noqa: F401
+        from triton.experimental.gluon.language import amd as _amd
+
+        return hasattr(_amd, "cdna4")
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _require_gluon_v2_triton() -> None:
+    """Fail fast with a build hint when Gluon is unavailable. The kernel compile is
+    ALSO guarded at the launch site: a Gluon-capable-but-incompatible triton raises a
+    CompilationError there, which is re-wrapped with the same install hint (so the user
+    always gets the commit rather than a raw layout/compile error)."""
+    if not _gluon_available():
+        raise RuntimeError(_GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown")))
+
+
+def _wrap_compile_error(exc: Exception) -> RuntimeError:
+    return RuntimeError(
+        _GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown"))
+        + f"\n(triton failed to compile the Gluon fwd: {type(exc).__name__}: "
+        + (str(exc).splitlines()[0] if str(exc).strip() else "")
+        + ")"
+    )
+
+
+# =====================================================================
+# Utility
+# =====================================================================
+def _get_lds_limit():
+    """Return the per-CU LDS limit in bytes for the current GPU.
+
+    gfx942 (MI300X): 64 KB = 65536 bytes
+    gfx950 (MI355X): 160 KB = 163840 bytes
+    """
+    if torch.cuda.is_available():
+        prop = torch.cuda.get_device_properties(0)
+        gcn_arch = getattr(prop, "gcnArchName", "")
+        if "gfx950" in gcn_arch:
+            return 163840
+    return 65536
+
+
+_LDS_LIMIT = _get_lds_limit()
+
+
+# =====================================================================
+# Forward — autotune configs and pruning
+# =====================================================================
+def _fwd_prune_configs(configs, named_args, **kwargs):
+    """Prune autotune configs that would exceed per-CU LDS."""
+    D_V = kwargs.get("D_V", named_args.get("D_V"))
+    D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE"))
+    pruned = []
+    for config in configs:
+        config.kwargs["BLOCK_H"]
+        tk = config.kwargs["TILE_K"]
+        ns = config.num_stages
+        kv_lds = (D_V + D_ROPE) * tk * 2 * ns
+        if kv_lds <= _LDS_LIMIT:
+            pruned.append(config)
+    if not pruned:
+        pruned.append(configs[0])
+    return pruned
+
+
+def _get_fwd_autotune_configs():
+    configs = [
+        triton.Config(
+            {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE},
+            num_warps=nw,
+        )
+        for BLOCK_H in [16, 32, 64]
+        for TILE_K in [16, 32, 64, 128]
+        for WPE in [0, 1, 2]
+        for nw in [4]  # num_warps must be 4 to align with kernel implementation
+    ]
+    # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),]
+    return configs
+
+
+@triton.autotune(
+    configs=_get_fwd_autotune_configs(),
+    key=["num_heads", "TOPK", "D_V", "D_ROPE"],
+    prune_configs_by={"early_config_prune": _fwd_prune_configs},
+)
+@gluon.jit
+def _sparse_mla_fwd_gl_v2_kernel(
+    Q_ptr,  # [total_tokens, num_heads, D_QK] bf16
+    KV_ptr,  # [total_tokens, 1, D_QK]         bf16
+    TopK_ptr,  # [total_tokens, TOPK]            int32
+    Sink_ptr,  # [num_heads]                     fp32; ignored if HAS_SINK == False
+    O_ptr,  # [total_tokens, num_heads, D_V]  bf16
+    LSE_ptr,  # [total_tokens, num_heads]       fp32 (sink-inclusive if HAS_SINK)
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_o_t: tl.int64,
+    stride_o_h: tl.int64,
+    stride_topk_t: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    TOPK: gl.constexpr,
+    BLOCK_H: gl.constexpr,
+    TILE_K: gl.constexpr,
+    D_V: gl.constexpr,
+    D_ROPE: gl.constexpr,
+    HAS_SINK: gl.constexpr,
+    HAS_ROPE: gl.constexpr,
+):
+    # ---------- constexpr layouts ----------
+    # QK MFMA uses K=32 (aiter): the score matmul reduces D_V=512, so instr_shape=[16,16,32]
+    # halves the MFMA instruction count vs [16,16,16]. mfma_acc (PV) stays K=16 since its
+    # reduction dim is TILE_K (autotuned down to 16).
+    mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 32],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+    mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout(
+        version=4,
+        instr_shape=[16, 16, 16],
+        transposed=True,
+        warps_per_cta=[4, 1],
+    )
+
+    # Blocked layouts for global loads.
+    _qlora_tpw_k: gl.constexpr = min(64, D_V // 8)
+    _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k
+    blk_qlora: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+    blk_qrope: gl.constexpr = gl.BlockedLayout(
+        size_per_thread=[1, 8],
+        threads_per_warp=[8, 8],
+        warps_per_cta=[4, 1],
+        order=[1, 0],
+    )
+
+    _klora_tpw_m: gl.constexpr = min(64, D_V // 8)
+    _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m
+    blk_klora: gl.constexpr = gl.BlockedLayout(  # [D_V, TILE_K]
+        size_per_thread=[8, 1],
+        threads_per_warp=[_klora_tpw_m, _klora_tpw_n],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_krope: gl.constexpr = gl.BlockedLayout(  # [D_ROPE, TILE_K] = [64, 16]
+        size_per_thread=[2, 1],
+        threads_per_warp=[32, 2],
+        warps_per_cta=[1, 4],
+        order=[0, 1],
+    )
+    blk_topk: gl.constexpr = gl.BlockedLayout(  # [TILE_K] int32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+    blk_lse: gl.constexpr = gl.BlockedLayout(  # [BLOCK_H] fp32
+        size_per_thread=[1],
+        threads_per_warp=[64],
+        warps_per_cta=[4],
+        order=[0],
+    )
+
+    # Shared layouts.
+    sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [BLOCK_H, D_V],
+        [1, 0],
+    )
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[1, 0],
+    )
+    sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for(
+        [[512, 16]],
+        [D_V, TILE_K],
+        [0, 1],
+    )
+    sh_krope: gl.constexpr = gl.SwizzledSharedLayout(
+        vec=8,
+        per_phase=2,
+        max_phase=8,
+        order=[0, 1],
+    )
+
+    # Dot operand layouts
+    dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8)
+    dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8)
+    dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4)
+    dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4)
+
+    # ---------- program ids ----------
+    token_idx = gl.program_id(axis=0)
+    hg_idx = gl.program_id(axis=1)
+    hg_offset = hg_idx * BLOCK_H
+
+    # ---------- offsets for Q ----------
+    # Q_lora  [BLOCK_H, D_V]
+    offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_qlora = offs_h_qlora < num_heads
+
+    q_base = token_idx.to(tl.int64) * stride_q_t
+    q_offs_lora = (
+        q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64)
+    )
+    q_mask_lora = mask_h_qlora[:, None]
+
+    smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_qlora,
+        ptr=Q_ptr,
+        offsets=q_offs_lora.to(tl.int32),
+        mask=q_mask_lora,
+    )
+    # V4 zero-rope-pad: skip the rope Q load + rope MFMA entirely when HAS_ROPE is False
+    # (RoPE is baked in-place over the 512 latent, so the rope QK term is provably zero).
+    if HAS_ROPE:
+        # Q_rope  [BLOCK_H, D_ROPE]
+        offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope))
+        offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope))
+        mask_h_qrope = offs_h_qrope < num_heads
+        q_offs_rope = (
+            q_base
+            + offs_h_qrope[:, None].to(tl.int64) * stride_q_h
+            + (D_V + offs_r_qrope[None, :]).to(tl.int64)
+        )
+        q_mask_rope = mask_h_qrope[:, None]
+        smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_qrope,
+            ptr=Q_ptr,
+            offsets=q_offs_rope.to(tl.int32),
+            mask=q_mask_rope,
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    # ---------- topk and KV offsets ----------
+    NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K
+    topk_base = token_idx.to(tl.int64) * stride_topk_t
+
+    # offs_tile in three layouts (sliced from each of the three loaders)
+    offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora))
+    offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope))
+    offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s))
+    offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk)
+
+    offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora))
+    offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope))
+
+    # (removed dead `topk_pos_reg` prologue load — was never consumed)
+
+    # ---------- shared mem allocations for the K loop ----------
+    if HAS_ROPE:
+        smem_krope = gl.allocate_shared_memory(
+            KV_ptr.dtype.element_ty,
+            [2, D_ROPE, TILE_K],
+            layout=sh_krope,
+        )
+    smem_klora = gl.allocate_shared_memory(
+        KV_ptr.dtype.element_ty,
+        [2, D_V, TILE_K],
+        layout=sh_klora,
+    )
+
+    # ---------- accumulators ----------
+    m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s))
+    acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc)
+    # exp2 softmax (aiter technique): fold log2(e) into the QK scale so the per-element
+    # softmax exp becomes a single hardware exp2 (no per-element *log2e). m_i / l_i are then
+    # in log2 units; lse is converted back to natural log at the epilogue (the bwd needs nat-log).
+    scale_log2 = scale * 1.4426950408889634
+
+    # ---------- tile-0 prefetch (prologue) ----------
+    # Load K_lora and K_rope for tile 0.
+    topk_pos_klora = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_klora,
+        mask=offs_tile_klora < TOPK,
+        other=-1,
+    )
+    if HAS_ROPE:
+        topk_pos_krope = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=topk_base.to(tl.int32) + offs_tile_krope,
+            mask=offs_tile_krope < TOPK,
+            other=-1,
+        )
+    topk_pos_mma = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + offs_tile_mma,
+        mask=offs_tile_mma < TOPK,
+        other=-1,
+    )
+
+    # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a
+    # single register set; converted to the klora/krope/mma layouts at point of use, to
+    # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen
+    # regression -- see att_fwd_gluon_mi350/RESULTS.md).
+    p1_off_topk = TILE_K + offs_tile_topk
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + p1_off_topk,
+        mask=p1_off_topk < TOPK,
+        other=-1,
+    )
+
+    valid_klora = topk_pos_klora != -1  # tile_start=0 -> offs_tile<TOPK already true
+    valid_mma = topk_pos_mma != -1
+
+    safe_klora = gl.where(valid_klora, topk_pos_klora, 0)
+
+    # K_lora async DMA into smem_klora[0]
+    klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    klora_smem0 = smem_klora.index(0)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=klora_smem0,
+        ptr=KV_ptr,
+        offsets=klora_offs.to(tl.int32),
+        mask=valid_klora[None, :],
+    )
+
+    # K_rope async DMA into smem_krope[0] — same group as K_lora (skipped if HAS_ROPE=False).
+    if HAS_ROPE:
+        valid_krope = topk_pos_krope != -1
+        safe_krope = gl.where(valid_krope, topk_pos_krope, 0)
+        krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(
+            tl.int64
+        )
+        krope_smem0 = smem_krope.index(0)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=krope_smem0,
+            ptr=KV_ptr,
+            offsets=krope_offs.to(tl.int32),
+            mask=valid_krope[None, :],
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+
+    gl.amd.cdna4.async_copy.wait_group(1)
+    Q_lora_dot = smem_qlora.load(dot_qlora_a)
+    if HAS_ROPE:
+        Q_rope_dot = smem_qrope.load(dot_qrope_a)
+
+    # ---------- main loop (3-stage): gather tile t+1 from ALREADY-loaded topk,
+    #            deep-prefetch topk for tile t+2, compute tile t ----------
+    # Carried across iterations: tkraw_{klora,krope,mma} = raw topk for the tile to be
+    # GATHERED this iter (loaded the previous iter / prologue); valid_mma = mask for the
+    # tile being COMPUTED this iter. This lifts the topk load off the KV-gather critical
+    # path so the async_copy is never gated by a fresh index load.
+    # ===== Software-pipelined: softmax+PV of tile t-1 overlaps the QK MFMA of tile t. =====
+    # Carry S_prev (the masked/scaled QK output, small). QK reads cur_buf; PV reads 1-cur_buf;
+    # gather the next tile into 1-cur_buf AFTER V_dot has been read into registers, so the gather
+    # overlaps the PV MFMA (and the gather is cache-served/short). 2-deep buffer suffices.
+    offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+    mask_h_mma = offs_h_mma < num_heads
+
+    # WARM-UP: gather K[1] -> buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet).
+    tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+    tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+    valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+    valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+    safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+    klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        dest=smem_klora.index(1),
+        ptr=KV_ptr,
+        offsets=klora_offs_next.to(tl.int32),
+        mask=valid_klora_next[None, :],
+    )
+    if HAS_ROPE:
+        tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+        valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+        safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+        krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+            D_V + offs_r_krope[:, None]
+        ).to(tl.int64)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_krope.index(1),
+            ptr=KV_ptr,
+            offsets=krope_offs_next.to(tl.int32),
+            mask=valid_krope_next[None, :],
+        )
+    gl.amd.cdna4.async_copy.commit_group()
+    tkraw = gl.amd.cdna4.buffer_load(
+        ptr=TopK_ptr,
+        offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk),
+        mask=(2 * TILE_K + offs_tile_topk) < TOPK,
+        other=-1,
+    )
+    gl.amd.cdna4.async_copy.wait_group(1)
+    S_prev = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(0).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    if HAS_ROPE:
+        S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev)
+    S_prev = S_prev * scale_log2
+    S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf"))
+    cur_buf = 1
+
+    for t in range(NUM_TILES - 2):
+        gl.amd.cdna4.async_copy.wait_group(0)  # drain K[t+1] (cur_buf) before QK reads it
+        # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer,
+        # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer).
+        # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range.
+        V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b)
+        tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora))
+        tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s))
+        valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1)
+        valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1)
+        safe_klora_next = gl.where(valid_klora_next, tk_klora, 0)
+        klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(
+            tl.int64
+        )
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            dest=smem_klora.index(1 - cur_buf),
+            ptr=KV_ptr,
+            offsets=klora_offs_next.to(tl.int32),
+            mask=valid_klora_next[None, :],
+        )
+        if HAS_ROPE:
+            tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope))
+            valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1)
+            safe_krope_next = gl.where(valid_krope_next, tk_krope, 0)
+            krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (
+                D_V + offs_r_krope[:, None]
+            ).to(tl.int64)
+            gl.amd.cdna4.async_copy.buffer_load_to_shared(
+                dest=smem_krope.index(1 - cur_buf),
+                ptr=KV_ptr,
+                offsets=krope_offs_next.to(tl.int32),
+                mask=valid_krope_next[None, :],
+            )
+        gl.amd.cdna4.async_copy.commit_group()
+        tkraw_n = gl.amd.cdna4.buffer_load(
+            ptr=TopK_ptr,
+            offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk),
+            mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK,
+            other=-1,
+        )
+        # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below
+        S_cur = gl.amd.cdna4.mfma(
+            Q_lora_dot,
+            smem_klora.index(cur_buf).load(dot_klora_b),
+            gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+        )
+        if HAS_ROPE:
+            S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+        S_cur = S_cur * scale_log2
+        S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+        # softmax(S_prev = tile t) [VALU, overlaps QK]
+        m_j = gl.max(S_prev, axis=1)
+        m_new = gl.maximum(m_i, m_j)
+        m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+        alpha = gl.exp2(m_i - m_new)
+        P = gl.exp2(S_prev - m_new[:, None])
+        l_i = alpha * l_i + gl.sum(P, axis=1)
+        m_i = m_new
+        # PV tile t from registers -- matrix; overlaps the gather still in flight
+        alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_acc[:, None]
+        P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+        acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc)
+        # promote
+        S_prev = S_cur
+        valid_qk = valid_qk_next
+        tkraw = tkraw_n
+        cur_buf = 1 - cur_buf
+
+    # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ----------
+    gl.amd.cdna4.async_copy.wait_group(0)  # drain K[N-1] (last loop gather)
+    S_cur = gl.amd.cdna4.mfma(
+        Q_lora_dot,
+        smem_klora.index(cur_buf).load(dot_klora_b),
+        gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s),
+    )
+    if HAS_ROPE:
+        S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur)
+    S_cur = S_cur * scale_log2
+    S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf"))
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp2(m_i - m_new)
+    P = gl.exp2(S_prev - m_new[:, None])
+    l_i = alpha * l_i + gl.sum(P, axis=1)
+    m_i = m_new
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    S_prev = S_cur
+
+    # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ----------
+    m_j = gl.max(S_prev, axis=1)
+    m_new = gl.maximum(m_i, m_j)
+    m_new = gl.where(m_new > float("-inf"), m_new, 0.0)
+    alpha = gl.exp2(m_i - m_new)
+    P = gl.exp2(S_prev - m_new[:, None])
+    l_new = alpha * l_i + gl.sum(P, axis=1)
+    alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc))
+    acc = acc * alpha_acc[:, None]
+    P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a)
+    acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc)
+    m_i = m_new
+    l_i = l_new
+
+    # ---------- epilogue: fold sink into the denominator (V4 delta) ----------
+    if HAS_SINK:
+        offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s))
+        sink = gl.amd.cdna4.buffer_load(
+            ptr=Sink_ptr,
+            offsets=offs_h_sink.to(tl.int32),
+            mask=offs_h_sink < num_heads,
+            other=float("-inf"),
+        )
+        sink = sink * 1.4426950408889634  # natural-log sink -> log2 units (exp2 softmax)
+        m_final = gl.maximum(m_i, sink)
+        alpha_fix = gl.exp2(m_i - m_final)
+        l_total = l_i * alpha_fix + gl.exp2(sink - m_final)
+        alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc))
+        acc = acc * alpha_fix_acc[:, None]
+        l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_total_acc[:, None]
+        # lse back to natural log for the backward: m_final is log2, l_total is the natural denom.
+        lse = m_final * 0.6931471805599453 + gl.log(l_total)
+    else:
+        l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc))
+        acc = acc / l_i_acc[:, None]
+        lse = m_i * 0.6931471805599453 + gl.log(l_i)
+
+    # Output O[token_idx, h, v]
+    offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora))
+    offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora))
+    mask_h_o = offs_h_o < num_heads
+    o_base = token_idx.to(tl.int64) * stride_o_t
+    o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64)
+    acc_bf = acc.to(O_ptr.dtype.element_ty)
+    acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora)
+    gl.amd.cdna4.buffer_store(
+        stored_value=acc_bf_blk,
+        ptr=O_ptr,
+        offsets=o_offs.to(tl.int32),
+        mask=mask_h_o[:, None],
+    )
+
+    # LSE[token_idx, h]
+    offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse)
+    mask_h_lse = offs_h_lse < num_heads
+    lse_base = token_idx * num_heads
+    lse_offs = lse_base + offs_h_lse
+    lse_blk = gl.convert_layout(lse, blk_lse)
+    gl.amd.cdna4.buffer_store(
+        stored_value=lse_blk,
+        ptr=LSE_ptr,
+        offsets=lse_offs.to(tl.int32),
+        mask=mask_h_lse,
+    )
+
+
+# =====================================================================
+# Launcher
+# =====================================================================
+def sparse_mla_fwd_v4_gluon_v2(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """
+    DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink.
+
+    Args:
+        q:             [total_tokens, num_heads, d_qk] bfloat16
+        kv:            [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk])
+        topk_indices:  [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid)
+        attn_sink:     [num_heads] fp32, optional per-head learnable sink logit.
+                       When None, behaves like the V3.2 forward.
+        kv_lora_rank:  int, default 512
+        scale:         float, default 1/sqrt(d_qk)
+
+    Returns:
+        o:   [total_tokens, num_heads, kv_lora_rank] same dtype as q
+        lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given)
+    """
+    _require_gluon_v2_triton()
+    assert q.is_contiguous()
+    assert kv.is_contiguous()
+    assert topk_indices.is_contiguous()
+
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    # kv may hold MORE rows than there are query tokens (V4 feeds a
+    # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens).
+    # The kernel only dereferences kv via topk indices (stride_kv_t), so any
+    # num_kv >= max(topk_index)+1 is valid.
+    assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.is_contiguous()
+        assert attn_sink.dtype == torch.float32
+        assert attn_sink.shape == (num_heads,)
+        sink_ptr = attn_sink
+    else:
+        sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device)  # guarded by HAS_SINK
+
+    if num_heads == 128 and topk == 1152 and kv_lora_rank == 512 and rope_rank == 64:
+        return sparse_mla_fwd_v4_aiter_lse_csa_formula(
+            q, kv, topk_indices, attn_sink=attn_sink, kv_lora_rank=kv_lora_rank, scale=scale
+        )
+    if num_heads == 64 and topk == 640 and kv_lora_rank == 512 and rope_rank == 64:
+        return sparse_mla_fwd_v4_aiter_lse_csa_formula(
+            q, kv, topk_indices, attn_sink=attn_sink, kv_lora_rank=kv_lora_rank, scale=scale
+        )
+
+    o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device)
+    lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked in-place
+    # over the 512 latent), so the rope QK term is provably zero. Skip it — bit-identical
+    # to computing it, but avoids the wasteful 64-wide rope MFMA + K_rope loads every tile
+    # (this is the win triton_v2 already has that our ported gluon fwd lacked).
+    has_rope = False
+
+    # Grid is autotune-aware: BLOCK_H comes from the chosen config.
+    grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"]))
+
+    try:
+        _sparse_mla_fwd_gl_v2_kernel[grid](
+            Q_ptr=q,
+            KV_ptr=kv,
+            TopK_ptr=topk_indices,
+            Sink_ptr=sink_ptr,
+            O_ptr=o,
+            LSE_ptr=lse,
+            stride_q_t=q.stride(0),
+            stride_q_h=q.stride(1),
+            stride_kv_t=kv.stride(0),
+            stride_o_t=o.stride(0),
+            stride_o_h=o.stride(1),
+            stride_topk_t=topk_indices.stride(0),
+            scale=scale,
+            num_heads=num_heads,
+            TOPK=topk,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_SINK=has_sink,
+            HAS_ROPE=has_rope,
+        )
+    except Exception as exc:  # noqa: BLE001 - surface a build hint on Gluon compile failures
+        _n = type(exc).__name__.lower()
+        if "compil" in _n or "compil" in str(exc).lower() or "layout" in str(exc).lower():
+            raise _wrap_compile_error(exc) from exc
+        raise
+
+    return o, lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py
new file mode 100644
index 000000000..97ad10bf1
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py
@@ -0,0 +1,388 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-8 P49 — Tilelang V4 attention dispatcher (infra-only).
+
+This module is the dispatch point for the plan-8 tilelang-backed V4
+attention kernels (cr ∈ {0, 4, 128}).  At P49 it ships only the
+**infra**:
+
+* The pinned tilelang version probe (one-time module-import warning
+  if the installed tilelang ≠ the plan-8 pin).
+* The :func:`should_dispatch` predicate that takes the per-call
+  ``enabled`` flag (from the ``use_v4_tilelang_attention`` /
+  ``use_v4_tilelang_csa_attention`` config flags surfaced by
+  ``DeepSeekV4TransformerConfig``).  Plan-8 P57 close-out 2: the
+  dispatcher no longer reads ``PRIMUS_V4_TILELANG_ATTN`` — callers
+  pass the config flag explicitly so a container without tilelang
+  installed can simply leave the flag ``False`` and never trigger
+  any tilelang import.
+* The :func:`is_tilelang_kernel_available` predicate that lets each
+  plan-8 phase (P50..P55) register its kernel name as it lands.
+  Empty at P49 — every dispatcher call falls through to the
+  plan-4 P25 / P26 Triton path with a one-time rank-0 warning.
+* Four stub entry points that raise :class:`NotImplementedError`
+  until the corresponding plan-8 phase lands them:
+
+  - :func:`v4_attention_fwd_tilelang`           (P50 — cr=0 / cr=128 FWD)
+  - :func:`v4_attention_bwd_tilelang`           (P51 — cr=0 / cr=128 BWD)
+  - :func:`v4_csa_attention_fwd_tilelang`       (P54 — cr=4 FWD)
+  - :func:`v4_csa_attention_bwd_tilelang`       (P55 — cr=4 BWD)
+
+The dispatcher precedence (enforced inside the V4 attention
+functional wrappers ``v4_attention_v1`` / ``v4_csa_attention_v0``) is:
+
+.. code-block:: text
+
+    cr ∈ {0, 128}:
+      use_turbo_attention > use_v4_tilelang_attention
+                          > use_v4_triton_attention > eager
+    cr == 4:
+      use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager
+
+At P49 (default-off + empty available-kernels set) the dispatcher
+emits **no** behaviour change vs the plan-7 P48 anchor.  Plan-4..7
+ratchet stays green by construction.
+
+R6.2 — Tilelang is vendored at ``tilelang/`` and NEVER edited.  This
+module imports it lazily so a missing install does not raise at
+module import time; it only logs a one-time rank-0 warning the
+first time a caller actually asks the dispatcher to use tilelang.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any, Set
+
+# ---------------------------------------------------------------------------
+# Tilelang version pin (set by plan-8 P49; bump when upstream is updated)
+# ---------------------------------------------------------------------------
+
+# Read from ``tilelang/VERSION`` at plan-8 P49 land time.  Future plan-8
+# phases (P50..P55) MUST re-run the G50..G55 ratchets before bumping
+# this pin.
+TILELANG_VERSION_PIN: str = "0.1.9+cuda.gitbcb2da33"
+
+
+# ---------------------------------------------------------------------------
+# Dispatch-enable signal (replaces the PRIMUS_V4_TILELANG_ATTN env knob)
+# ---------------------------------------------------------------------------
+#
+# Plan-8 P57 close-out 2 (2026-05-15): the dispatcher no longer reads an
+# environment variable.  Callers pass ``enabled`` to :func:`should_dispatch`
+# from the V4 attention layer's config flag
+# (``use_v4_tilelang_attention`` / ``use_v4_tilelang_csa_attention``),
+# which is plumbed through :class:`DeepSeekV4TransformerConfig` from
+# the run-time CLI args.  Default-False everywhere — running on a
+# container without tilelang installed leaves the flag off and the
+# dispatcher never imports tilelang.
+
+
+# ---------------------------------------------------------------------------
+# Per-kernel availability registry
+# ---------------------------------------------------------------------------
+
+# Mutable set of kernel names that plan-8 phases register as they
+# land.  Empty at P49; populated by P50..P55 as each kernel ships
+# via ``register_available_kernel(<name>)``.  Names follow the
+# four-stub convention below.
+_AVAILABLE_KERNELS: Set[str] = set()
+
+_KNOWN_KERNEL_NAMES: Set[str] = {
+    "v4_attention_fwd",  # P50 — dense / HCA FWD
+    "v4_attention_bwd",  # P51 — dense / HCA BWD
+    "v4_csa_attention_fwd",  # P54 — CSA FWD
+    "v4_csa_attention_bwd",  # P55 — CSA BWD
+}
+
+
+def is_tilelang_kernel_available(name: str) -> bool:
+    """Return True iff the plan-8 kernel ``name`` has landed.
+
+    Empty at P49; each plan-8 phase (P50..P55) registers its
+    kernel name via :func:`register_available_kernel` once the
+    parity gate passes.
+    """
+    if name not in _KNOWN_KERNEL_NAMES:
+        raise ValueError(
+            f"Unknown tilelang kernel name {name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}"
+        )
+    return name in _AVAILABLE_KERNELS
+
+
+def register_available_kernel(name: str) -> None:
+    """Mark a plan-8 tilelang kernel as available.
+
+    Called at module-import time by the phase that lands the
+    corresponding kernel.  P50: ``register_available_kernel("v4_attention_fwd")``
+    inside ``v4_attention_fwd_tilelang.py``.
+    """
+    if name not in _KNOWN_KERNEL_NAMES:
+        raise ValueError(
+            f"Unknown tilelang kernel name {name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}"
+        )
+    _AVAILABLE_KERNELS.add(name)
+
+
+# Map kernel name -> submodule name to lazy-import.  Each submodule
+# registers its kernel name + overrides the stub in this namespace on
+# first import.
+_KERNEL_SUBMODULES: dict[str, str] = {
+    "v4_attention_fwd": "v4_attention_fwd_tilelang",
+    "v4_attention_bwd": "v4_attention_bwd_tilelang",
+    "v4_csa_attention_fwd": "v4_csa_attention_fwd_tilelang",
+    "v4_csa_attention_bwd": "v4_csa_attention_bwd_tilelang",
+}
+
+_LAZY_LOADED: Set[str] = set()
+
+
+def _lazy_load(name: str) -> bool:
+    """Import the submodule that implements kernel ``name`` (if landed).
+
+    Returns True iff the import succeeded.  Each plan-8 kernel
+    submodule's import-time code:
+
+    1. Calls :func:`register_available_kernel` to flip
+       :func:`is_tilelang_kernel_available` to True.
+    2. Replaces the module-level stub in this namespace via
+       ``setattr(<this module>, name + "_tilelang", real_fn)``.
+
+    Idempotent — subsequent calls no-op.
+    """
+    if name in _LAZY_LOADED:
+        return name in _AVAILABLE_KERNELS
+    _LAZY_LOADED.add(name)
+    sub = _KERNEL_SUBMODULES.get(name)
+    if sub is None:
+        return False
+    try:
+        from importlib import import_module
+
+        sub_mod = import_module(__name__ + "." + sub)
+    except ImportError as exc:
+        warnings.warn(
+            f"[plan-8] tilelang submodule {sub!r} import failed ({exc}); "
+            f"kernel {name!r} stays unavailable.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+        return False
+    # Override the stub: the submodule exposes a function with the same
+    # name as the submodule (e.g. `v4_attention_fwd_tilelang.py`
+    # exposes `v4_attention_fwd_tilelang`).  Importing the submodule
+    # would otherwise set `_tilelang.v4_attention_fwd_tilelang` to
+    # the SUBMODULE, shadowing the stub.  We restore the function
+    # attribute here.
+    real_fn = getattr(sub_mod, sub, None)
+    if real_fn is None:
+        warnings.warn(
+            f"[plan-8] tilelang submodule {sub!r} does not expose a "
+            f"function named {sub!r}; kernel {name!r} stays unavailable.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+        return False
+    import sys
+
+    sys.modules[__name__].__dict__[sub] = real_fn
+    return name in _AVAILABLE_KERNELS
+
+
+# ---------------------------------------------------------------------------
+# Tilelang import probe (lazy)
+# ---------------------------------------------------------------------------
+
+_TILELANG_PROBE_DONE: bool = False
+_TILELANG_AVAILABLE: bool = False
+
+
+def _probe_tilelang() -> bool:
+    """Lazily import tilelang on first dispatcher call.
+
+    Returns True iff the import succeeds at the pinned version.
+    Emits a one-time warning when the installed version differs from
+    the pin or when tilelang is not importable.
+    """
+    global _TILELANG_PROBE_DONE, _TILELANG_AVAILABLE
+    if _TILELANG_PROBE_DONE:
+        return _TILELANG_AVAILABLE
+    _TILELANG_PROBE_DONE = True
+    try:
+        import tilelang  # noqa: F401
+    except ImportError as exc:
+        warnings.warn(
+            f"[plan-8 P49] tilelang import failed ({exc}); falling back to "
+            "the plan-4 P25 / P26 Triton kernels.  Install the pinned "
+            f"tilelang {TILELANG_VERSION_PIN} to enable the plan-8 path.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        _TILELANG_AVAILABLE = False
+        return False
+    installed = getattr(tilelang, "__version__", "<unknown>")
+    if installed != TILELANG_VERSION_PIN:
+        warnings.warn(
+            f"[plan-8 P49] installed tilelang {installed!r} != pinned "
+            f"{TILELANG_VERSION_PIN!r}; the plan-8 G50..G55 parity gates were "
+            "run against the pin.  Proceeding anyway; please re-run the "
+            "ratchets if the version drifted.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+    _TILELANG_AVAILABLE = True
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Cache directory
+# ---------------------------------------------------------------------------
+
+
+def cache_dir() -> str:
+    """Tilelang autotune cache dir.
+
+    Defaults to ``output/.tilelang_cache/v4/`` (gitignored under
+    ``output/``).  Override via ``PRIMUS_V4_TILELANG_CACHE_DIR``.
+    """
+    import os as _os
+
+    override = _os.environ.get("PRIMUS_V4_TILELANG_CACHE_DIR")
+    if override:
+        return override
+    return _os.path.join("output", ".tilelang_cache", "v4")
+
+
+# ---------------------------------------------------------------------------
+# Stub entry points — replaced by P50 / P51 / P54 / P55
+# ---------------------------------------------------------------------------
+
+
+def _stub_raise(phase: str, name: str) -> None:
+    raise NotImplementedError(
+        f"plan-8 {phase} tilelang kernel {name!r} has not landed yet; "
+        "the dispatcher should fall back to the plan-4 / plan-5 Triton "
+        "path via `is_tilelang_kernel_available({name!r})` returning False."
+    )
+
+
+def v4_attention_fwd_tilelang(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+    """Stub for plan-8 P50 (dense / HCA FWD)."""
+    _stub_raise("P50", "v4_attention_fwd")
+
+
+def v4_attention_bwd_tilelang(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+    """Stub for plan-8 P51 (dense / HCA BWD)."""
+    _stub_raise("P51", "v4_attention_bwd")
+
+
+def v4_csa_attention_fwd_tilelang(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+    """Stub for plan-8 P54 (CSA FWD)."""
+    _stub_raise("P54", "v4_csa_attention_fwd")
+
+
+def v4_csa_attention_bwd_tilelang(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+    """Stub for plan-8 P55 (CSA BWD)."""
+    _stub_raise("P55", "v4_csa_attention_bwd")
+
+
+# ---------------------------------------------------------------------------
+# Dispatcher helpers — used by the v4_attention_v1 / v4_csa_attention_v0 wrappers
+# ---------------------------------------------------------------------------
+
+_FALLBACK_WARNED: Set[str] = set()
+
+
+def _maybe_warn_fallback(kernel_name: str) -> None:
+    """Emit a one-time rank-0 warning when a tilelang dispatch falls
+    back to Triton.
+
+    Hit conditions:
+
+    * Config flag ``use_v4_tilelang_attention`` /
+      ``use_v4_tilelang_csa_attention`` set but the kernel hasn't
+      landed yet (``is_tilelang_kernel_available`` returns False).
+    * Config flag set but tilelang import failed (e.g. missing
+      install / version drift).
+
+    Once per kernel name per process.  Banned-warning ratchet is
+    extended in `plan-8/03-test-strategy.md` to allow this string
+    (since it's the documented dispatcher fallback signal).
+    """
+    if kernel_name in _FALLBACK_WARNED:
+        return
+    _FALLBACK_WARNED.add(kernel_name)
+    # Best-effort rank-0 detection — the V4 attention wrappers may
+    # run in unit tests without parallel-state initialised.
+    rank = 0
+    try:
+        import torch.distributed as dist
+
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+    except Exception:
+        pass
+    if rank == 0:
+        warnings.warn(
+            f"[plan-8 P57] use_v4_tilelang_*=True but kernel "
+            f"{kernel_name!r} is not available; falling back to the "
+            "plan-4 / plan-5 Triton path.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+
+
+def should_dispatch(kernel_name: str, enabled: bool = False) -> bool:
+    """Single dispatcher predicate used by ``v4_attention_v1`` /
+    ``v4_csa_attention_v0`` wrappers.
+
+    Returns True iff all three conditions hold:
+
+    1. ``enabled`` is True (i.e. the caller passes the relevant
+       ``use_v4_tilelang_*`` config flag).
+    2. The named plan-8 kernel has been registered (P50..P55 land it).
+    3. The tilelang import succeeds at the pinned version (probed
+       lazily on the first dispatcher call).
+
+    Otherwise emits a one-time rank-0 warning + returns False so the
+    caller falls through to the Triton path.  ``enabled=False`` short-
+    circuits before any tilelang import attempt, so containers without
+    tilelang installed can leave the config flag off and the dispatcher
+    never touches tilelang.
+
+    Plan-8 P57 close-out 2 (2026-05-15): previously gated by the
+    ``PRIMUS_V4_TILELANG_ATTN`` env var; now driven by the caller's
+    config flag so default-off runs in tilelang-free containers do
+    not need to set or unset any env knob.
+    """
+    if not enabled:
+        return False
+    if not _probe_tilelang():
+        _maybe_warn_fallback(kernel_name)
+        return False
+    # Lazy-import the submodule that implements ``kernel_name``.
+    # The submodule's import-time code flips
+    # ``is_tilelang_kernel_available(...)`` to True and overrides
+    # the stub function in this namespace.  Cheap on subsequent calls.
+    _lazy_load(kernel_name)
+    if not is_tilelang_kernel_available(kernel_name):
+        _maybe_warn_fallback(kernel_name)
+        return False
+    return True
+
+
+__all__ = [
+    "TILELANG_VERSION_PIN",
+    "cache_dir",
+    "is_tilelang_kernel_available",
+    "register_available_kernel",
+    "should_dispatch",
+    "v4_attention_bwd_tilelang",
+    "v4_attention_fwd_tilelang",
+    "v4_csa_attention_bwd_tilelang",
+    "v4_csa_attention_fwd_tilelang",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py
new file mode 100644
index 000000000..10769790d
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py
@@ -0,0 +1,147 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-8 P51 — autograd wrapper that ties the tilelang FWD + BWD.
+
+Mirrors :class:`primus.backends.megatron.core.transformer.v4_attention_kernels.v4_attention.V4AttentionFn`
+(the Triton autograd Function) but routes through the tilelang
+FWD / BWD kernels.  The wrapper falls back to the Triton path
+inside the FWD / BWD wrapper functions when the tilelang kernels
+don't yet support the requested feature (e.g. additive_mask,
+hca_local_seqlen > 0).
+
+Implementation note: no ``from __future__ import annotations``
+because tilelang's annotation eval is eager.
+"""
+
+from typing import Optional
+
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_bwd_tilelang import (
+    v4_attention_bwd_tilelang,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_fwd_tilelang import (
+    v4_attention_fwd_tilelang_with_lse,
+)
+
+
+class V4AttentionTilelangFn(torch.autograd.Function):
+    """Autograd-aware V4 dense / SWA / sink attention via tilelang.
+
+    FWD: :func:`v4_attention_fwd_tilelang_with_lse` returns
+    ``(out, lse)``; we save the inputs + ``lse`` for backward.
+
+    BWD: :func:`v4_attention_bwd_tilelang` consumes the saved
+    tensors + the incoming ``dO``; returns ``(dq, dk, dv, dsink)``.
+    For unsupported features (additive_mask, hca_local_seqlen),
+    both wrappers fall back to the Triton kernels so this autograd
+    path remains correct even when tilelang's scope is limited.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        sink: Optional[torch.Tensor],
+        additive_mask: Optional[torch.Tensor],
+        swa_window: int,
+        attn_dropout: float,
+        training: bool,
+        scale: float,
+        hca_local_seqlen: int,
+    ) -> torch.Tensor:
+        out, lse = v4_attention_fwd_tilelang_with_lse(
+            q,
+            k,
+            v,
+            sink=sink,
+            additive_mask=additive_mask,
+            swa_window=int(swa_window),
+            attn_dropout=float(attn_dropout),
+            training=bool(training),
+            scale=float(scale),
+            hca_local_seqlen=int(hca_local_seqlen),
+        )
+        # Save tensors needed for BWD.  `additive_mask` and `sink`
+        # may be None; save_for_backward only stores tensors, so we
+        # stash None-ness on ctx and conditionally save.
+        ctx.save_for_backward(q, k, v, out, lse)
+        ctx.sink = sink
+        ctx.additive_mask = additive_mask
+        ctx.swa_window = int(swa_window)
+        ctx.scale = float(scale)
+        ctx.hca_local_seqlen = int(hca_local_seqlen)
+        return out
+
+    @staticmethod
+    def backward(ctx, d_out):  # type: ignore[override]
+        q, k, v, out, lse = ctx.saved_tensors
+        dq, dk, dv, dsink = v4_attention_bwd_tilelang(
+            q,
+            k,
+            v,
+            out,
+            lse,
+            d_out,
+            sink=ctx.sink,
+            additive_mask=ctx.additive_mask,
+            swa_window=ctx.swa_window,
+            scale=ctx.scale,
+            hca_local_seqlen=ctx.hca_local_seqlen,
+        )
+        # Match the input-arg count of `forward(...)`: q, k, v, sink,
+        # additive_mask, swa_window, attn_dropout, training, scale,
+        # hca_local_seqlen.  Non-tensor args get None.
+        return (
+            dq,
+            dk,
+            dv,
+            dsink,
+            None,  # additive_mask
+            None,  # swa_window
+            None,  # attn_dropout
+            None,  # training
+            None,  # scale
+            None,  # hca_local_seqlen
+        )
+
+
+def v4_attention_tilelang(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor] = None,
+    additive_mask: Optional[torch.Tensor] = None,
+    swa_window: int = 0,
+    attn_dropout: float = 0.0,
+    training: bool = False,
+    scale: Optional[float] = None,
+    hca_local_seqlen: int = 0,
+) -> torch.Tensor:
+    """Convenience wrapper matching the existing `v4_attention_v1()`
+    functional API but routing through the tilelang autograd
+    Function."""
+    if scale is None:
+        scale = 1.0 / (q.shape[-1] ** 0.5)
+    return V4AttentionTilelangFn.apply(
+        q,
+        k,
+        v,
+        sink,
+        additive_mask,
+        int(swa_window),
+        float(attn_dropout),
+        bool(training),
+        float(scale),
+        int(hca_local_seqlen),
+    )
+
+
+__all__ = ["V4AttentionTilelangFn", "v4_attention_tilelang"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py
new file mode 100644
index 000000000..937f8a70e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py
@@ -0,0 +1,427 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-8 P51 — V4 dense BWD tilelang kernels (compress_ratio == 0).
+
+Three-kernel pipeline matching the tilelang FlashAttention v2 BWD
+reference (`tilelang/examples/amd/example_amd_flash_attn_bwd.py`):
+
+* :func:`_make_preprocess_kernel` — ``Delta[b, h, m] =
+  sum_d O[b, h, m, d] * dO[b, h, m, d]``.  One tiny kernel.
+* :func:`_make_bwd_kernel` — main BWD pass.  One program per
+  ``(h, n_tile, b)``; loops over m_tile; re-materialises ``P``
+  from saved ``LSE``; accumulates ``dQ`` / ``dK`` / ``dV``
+  via ``tl.atomic_add``.  Sink BWD reuses the dq stride.
+* The sink gradient is computed in the Python wrapper as a small
+  reduction over the saved ``LSE`` + ``Delta`` (cheap; one ATen
+  launch).
+
+Scope at P51 (matches P50):
+
+* Q/K/V in BHSD layout with optional MQA broadcast (`K_H ∈ {1, H}`).
+* Optional per-head sink.
+* Optional sliding-window-causal mask.
+* No `additive_mask`, no `hca_local_seqlen > 0` — those route
+  through the plan-4 P25 / plan-5 P32 final Triton BWD kernel.
+
+dtype contract (matches :func:`reference.eager_v4_attention`
+backward): matmuls run in input dtype on tensor cores; the
+accumulator + softmax recompute live in fp32.  Output gradients
+are in `q.dtype` / `k.dtype` / `v.dtype` (typically all bf16);
+internal accumulators are fp32 buffers cast at the end.
+
+Implementation note: this module does NOT use ``from __future__
+import annotations`` because tilelang's ``@T.prim_func``
+decorator evaluates annotations eagerly via
+``typing.get_type_hints``.
+"""
+
+from typing import Optional, Tuple
+
+import tilelang
+import tilelang.language as T
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (
+    _launch_v4_attention_bwd,
+)
+
+_PREPROCESS_CACHE: dict[tuple, object] = {}
+_BWD_CACHE: dict[tuple, object] = {}
+
+
+# ---------------------------------------------------------------------------
+# Preprocess: Delta[b, h, m] = (O * dO).sum(-1)
+# ---------------------------------------------------------------------------
+
+
+@tilelang.jit(out_idx=[2])
+def _make_preprocess_kernel(
+    batch: int,
+    heads: int,
+    seq_q: int,
+    dim: int,
+    dtype: "T.dtype" = T.bfloat16,
+    block: int = 32,
+):
+    """Preprocess kernel: emit ``Delta [B, H, Sq]`` fp32.
+
+    Tiles `(b, h, m_tile)`; each program reduces over `dim` in
+    chunks of `block`.
+    """
+    accum_dtype = T.float32
+    bhsd_shape = [batch, heads, seq_q, dim]
+    delta_shape = [batch, heads, seq_q]
+
+    @T.prim_func
+    def main(
+        O: T.Tensor(bhsd_shape, dtype),
+        dO: T.Tensor(bhsd_shape, dtype),
+        Delta: T.Tensor(delta_shape, accum_dtype),
+    ):
+        with T.Kernel(batch, heads, T.ceildiv(seq_q, block)) as (bz, bx, by):
+            o = T.alloc_fragment([block, block], dtype)
+            do = T.alloc_fragment([block, block], dtype)
+            acc = T.alloc_fragment([block, block], accum_dtype)
+            delta = T.alloc_fragment([block], accum_dtype)
+            T.clear(acc)
+            for k in range(T.ceildiv(dim, block)):
+                T.copy(
+                    O[bz, bx, by * block : (by + 1) * block, k * block : (k + 1) * block],
+                    o,
+                )
+                T.copy(
+                    dO[bz, bx, by * block : (by + 1) * block, k * block : (k + 1) * block],
+                    do,
+                )
+                for i, j in T.Parallel(block, block):
+                    acc[i, j] += o[i, j] * do[i, j]
+            T.reduce_sum(acc, delta, 1)
+            T.copy(delta, Delta[bz, bx, by * block : (by + 1) * block])
+
+    return main
+
+
+# ---------------------------------------------------------------------------
+# Main BWD: dQ, dK, dV via re-materialised P from saved LSE
+# ---------------------------------------------------------------------------
+
+
+# Note: dK / dV are emitted as fp32 (not the input dtype) so the
+# MQA path can use `T.atomic_add` — tilelang's HIP runtime only
+# supports `AtomicAddx2(float*, ...)`.  The wrapper casts to
+# k.dtype / v.dtype before returning.
+@tilelang.jit(out_idx=[6, 7, 8])
+def _make_bwd_kernel(
+    batch: int,
+    heads_q: int,
+    heads_k: int,
+    seq_q: int,
+    seq_k: int,
+    dim: int,
+    has_sink: bool,
+    swa_window: int,
+    dtype: "T.dtype" = T.bfloat16,
+    block_M: int = 32,
+    block_N: int = 32,
+    threads: int = 64,
+):
+    """Plan-8 P51 main BWD: dQ / dK / dV via re-materialised softmax.
+
+    Grid: one program per `(b, h, n_tile)` (n_tile along K/V seq dim).
+    Each program loops over `m_tile` along Q.  Pattern mirrors the
+    tilelang AMD example BWD with three adaptations:
+
+    * BHSD layout instead of BSHD.
+    * MQA via `k_head = bx // groups`.
+    * SWA-aware m_tile loop bounds (skip tiles outside the window).
+
+    Output `dQ` is in fp32 (caller casts back to q.dtype after a
+    `dsink` host-side reduction); `dK`, `dV` are emitted in
+    `k.dtype` / `v.dtype` via the kernel epilogue.
+    """
+    groups = heads_q // heads_k
+    accum_dtype = T.float32
+    sm_scale = 1.0 / (dim**0.5)
+    q_shape = [batch, heads_q, seq_q, dim]
+    kv_shape = [batch, heads_k, seq_k, dim]
+    dq_shape = [batch, heads_q, seq_q, dim]
+    dkv_shape = [batch, heads_k, seq_k, dim]
+    lse_shape = [batch, heads_q, seq_q]
+    delta_shape = [batch, heads_q, seq_q]
+
+    @T.prim_func
+    def main(
+        Q: T.Tensor(q_shape, dtype),
+        K: T.Tensor(kv_shape, dtype),
+        V: T.Tensor(kv_shape, dtype),
+        dO: T.Tensor(q_shape, dtype),
+        LSE: T.Tensor(lse_shape, accum_dtype),
+        Delta: T.Tensor(delta_shape, accum_dtype),
+        dQ: T.Tensor(dq_shape, accum_dtype),
+        dK: T.Tensor(dkv_shape, accum_dtype),
+        dV: T.Tensor(dkv_shape, accum_dtype),
+    ):
+        with T.Kernel(heads_q, T.ceildiv(seq_k, block_M), batch, threads=threads) as (
+            bx,
+            by,
+            bz,
+        ):
+            K_shared = T.alloc_shared([block_M, dim], dtype)
+            V_shared = T.alloc_shared([block_M, dim], dtype)
+            q_shared = T.alloc_shared([block_N, dim], dtype)
+            do_shared = T.alloc_shared([block_N, dim], dtype)
+            lse_shared = T.alloc_shared([block_N], accum_dtype)
+            delta_shared = T.alloc_shared([block_N], accum_dtype)
+            ds_shared = T.alloc_shared([block_M, block_N], dtype)
+            p_cast = T.alloc_fragment([block_M, block_N], dtype)
+            qkT = T.alloc_fragment([block_M, block_N], accum_dtype)
+            P_acc = T.alloc_fragment([block_M, block_N], accum_dtype)
+            dP = T.alloc_fragment([block_M, block_N], accum_dtype)
+            dv = T.alloc_fragment([block_M, dim], accum_dtype)
+            dk = T.alloc_fragment([block_M, dim], accum_dtype)
+            dq_tile = T.alloc_fragment([block_N, dim], accum_dtype)
+
+            k_head = bx // groups
+
+            T.copy(K[bz, k_head, by * block_M : (by + 1) * block_M, :], K_shared)
+            T.copy(V[bz, k_head, by * block_M : (by + 1) * block_M, :], V_shared)
+            T.clear(dv)
+            T.clear(dk)
+
+            # Causal: only m_tiles that touch this n_tile contribute.
+            loop_st = by * block_M // block_N
+            loop_ed = T.ceildiv(seq_q, block_N)
+
+            for k in T.Pipelined(loop_st, loop_ed, num_stages=1):
+                T.copy(Q[bz, bx, k * block_N : (k + 1) * block_N, :], q_shared)
+                T.clear(qkT)
+                T.gemm(
+                    K_shared,
+                    q_shared,
+                    qkT,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullRow,
+                )
+                T.copy(LSE[bz, bx, k * block_N : (k + 1) * block_N], lse_shared)
+
+                for i, j in T.Parallel(block_M, block_N):
+                    P_acc[i, j] = T.exp(qkT[i, j] * sm_scale - lse_shared[j])
+
+                # Causal mask + optional SWA: keep entries where the
+                # query token (`k * block_N + j`) attends to the key
+                # token (`by * block_M + i`).  SWA: also require
+                # `q_idx - k_idx < swa_window`.
+                for i, j in T.Parallel(block_M, block_N):
+                    q_idx = k * block_N + j
+                    k_idx = by * block_M + i
+                    valid = q_idx >= k_idx
+                    if swa_window > 0:
+                        valid = valid and (q_idx - k_idx < swa_window)
+                    P_acc[i, j] = T.if_then_else(valid, P_acc[i, j], 0.0)
+
+                T.copy(dO[bz, bx, k * block_N : (k + 1) * block_N, :], do_shared)
+                T.clear(dP)
+                T.gemm(
+                    V_shared,
+                    do_shared,
+                    dP,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullRow,
+                )
+                T.copy(P_acc, p_cast)
+                T.gemm(p_cast, do_shared, dv, policy=T.GemmWarpPolicy.FullRow)
+
+                T.copy(Delta[bz, bx, k * block_N : (k + 1) * block_N], delta_shared)
+                for i, j in T.Parallel(block_M, block_N):
+                    p_cast[i, j] = P_acc[i, j] * (dP[i, j] - delta_shared[j]) * sm_scale
+                T.gemm(p_cast, q_shared, dk, policy=T.GemmWarpPolicy.FullRow)
+                T.copy(p_cast, ds_shared)
+
+                T.clear(dq_tile)
+                T.gemm(ds_shared, K_shared, dq_tile, transpose_A=True)
+                for i, j in T.Parallel(block_N, dim):
+                    T.atomic_add(dQ[bz, bx, k * block_N + i, j], dq_tile[i, j])
+
+            # MQA: every query head writes to the same K/V head;
+            # use atomic_add to merge across query-head programs.
+            # `dK` / `dV` are fp32 buffers (tilelang's HIP atomic-add
+            # only supports float*); the wrapper casts to k.dtype /
+            # v.dtype after the kernel returns.
+            if groups > 1:
+                for i, j in T.Parallel(block_M, dim):
+                    T.atomic_add(dV[bz, k_head, by * block_M + i, j], dv[i, j])
+                    T.atomic_add(dK[bz, k_head, by * block_M + i, j], dk[i, j])
+            else:
+                for i, j in T.Parallel(block_M, dim):
+                    dV[bz, k_head, by * block_M + i, j] = dv[i, j]
+                    dK[bz, k_head, by * block_M + i, j] = dk[i, j]
+
+    return main
+
+
+# ---------------------------------------------------------------------------
+# Wrapper API
+# ---------------------------------------------------------------------------
+
+
+def _kernel_supports(
+    *,
+    additive_mask: Optional[torch.Tensor],
+    hca_local_seqlen: int,
+) -> bool:
+    if additive_mask is not None:
+        return False
+    if hca_local_seqlen > 0:
+        return False
+    return True
+
+
+def _torch_dtype_to_tilelang(dtype: torch.dtype):
+    if dtype == torch.bfloat16:
+        return T.bfloat16
+    if dtype == torch.float16:
+        return T.float16
+    if dtype == torch.float32:
+        return T.float32
+    raise ValueError(f"unsupported dtype {dtype}")
+
+
+def _get_preprocess(batch, heads, seq_q, dim, dtype):
+    key = (batch, heads, seq_q, dim, str(dtype))
+    if key in _PREPROCESS_CACHE:
+        return _PREPROCESS_CACHE[key]
+    k = _make_preprocess_kernel(batch, heads, seq_q, dim, dtype=_torch_dtype_to_tilelang(dtype))
+    _PREPROCESS_CACHE[key] = k
+    return k
+
+
+def _get_bwd(batch, heads_q, heads_k, seq_q, seq_k, dim, has_sink, swa_window, dtype):
+    key = (
+        batch,
+        heads_q,
+        heads_k,
+        seq_q,
+        seq_k,
+        dim,
+        has_sink,
+        int(swa_window > 0) * swa_window,
+        str(dtype),
+    )
+    if key in _BWD_CACHE:
+        return _BWD_CACHE[key]
+    # Same SMEM-budget heuristic as P50.
+    if dim >= 256:
+        block_M, block_N, threads = 32, 32, 64
+    else:
+        block_M, block_N, threads = 64, 64, 128
+    k = _make_bwd_kernel(
+        batch,
+        heads_q,
+        heads_k,
+        seq_q,
+        seq_k,
+        dim,
+        has_sink,
+        swa_window,
+        dtype=_torch_dtype_to_tilelang(dtype),
+        block_M=block_M,
+        block_N=block_N,
+        threads=threads,
+    )
+    _BWD_CACHE[key] = k
+    return k
+
+
+def v4_attention_bwd_tilelang(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    do: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor] = None,
+    additive_mask: Optional[torch.Tensor] = None,
+    swa_window: int = 0,
+    scale: Optional[float] = None,
+    hca_local_seqlen: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Plan-8 P51 tilelang dense / SWA / sink BWD wrapper.
+
+    Drop-in replacement for ``_launch_v4_attention_bwd`` (plan-5
+    P32 final Triton split BWD) on the dense subset of the V4
+    attention contract.  Falls back to Triton for `additive_mask`
+    / `hca_local_seqlen > 0` paths (P53 territory).
+
+    Returns ``(dq, dk, dv, dsink)`` matching the Triton launcher
+    signature.
+    """
+    if not _kernel_supports(additive_mask=additive_mask, hca_local_seqlen=hca_local_seqlen):
+        return _launch_v4_attention_bwd(
+            q,
+            k,
+            v,
+            out,
+            lse,
+            do,
+            sink=sink,
+            additive_mask=additive_mask,
+            swa_window=swa_window,
+            scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5),
+            hca_local_seqlen=hca_local_seqlen,
+        )
+
+    B, HQ, Sq, D = q.shape
+    _, HK, Sk, _ = k.shape
+    assert HK in (1, HQ)
+    dtype = q.dtype
+
+    q_c = q.contiguous()
+    k_c = k.contiguous()
+    v_c = v.contiguous()
+    do_c = do.contiguous()
+    out_c = out.contiguous()
+    lse_c = lse.contiguous()
+
+    # Preprocess: Delta = (O * dO).sum(-1)
+    preprocess = _get_preprocess(B, HQ, Sq, D, dtype)
+    delta = preprocess(out_c, do_c)
+
+    has_sink = sink is not None
+    bwd = _get_bwd(B, HQ, HK, Sq, Sk, D, has_sink, int(swa_window), dtype)
+    dq_fp32, dk_fp32, dv_fp32 = bwd(q_c, k_c, v_c, do_c, lse_c, delta)
+    # All three grads are emitted in fp32 (MQA atomic-add target dtype);
+    # cast back to input dtypes.
+    dq = dq_fp32.to(dtype)
+    dk = dk_fp32.to(k.dtype)
+    dv = dv_fp32.to(v.dtype)
+
+    # Sink BWD: computed host-side (cheap).  When sink is present:
+    #   dsink[h] = sum_b sum_m exp(sink[h] - lse[b, h, m]) * delta[b, h, m]
+    if has_sink:
+        # Re-cast delta to fp32 for the reduce (it's already fp32 from
+        # preprocess).
+        dsink = (torch.exp(sink.float().unsqueeze(0).unsqueeze(-1) - lse_c.float()) * delta).sum(dim=(0, 2))
+        dsink = dsink.to(sink.dtype)
+    else:
+        dsink = None
+
+    return dq, dk, dv, dsink
+
+
+# ---------------------------------------------------------------------------
+# Module-level registration
+# ---------------------------------------------------------------------------
+
+_tilelang.register_available_kernel("v4_attention_bwd")
+# `_tilelang._lazy_load(...)` restores the parent-module attribute
+# (`_tilelang.v4_attention_bwd_tilelang`) after Python's import
+# machinery would otherwise shadow the stub with this submodule.
+
+
+__all__ = ["v4_attention_bwd_tilelang"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py
new file mode 100644
index 000000000..7340b03c9
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py
@@ -0,0 +1,480 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-8 P50 — V4 dense FWD tilelang kernel (compress_ratio == 0).
+
+Drop-in replacement for ``_launch_v4_attention_fwd`` (plan-4 P25
+Triton FWD) at the dense / SWA / sink subset of the V4 attention
+contract.  HCA (``hca_local_seqlen > 0``) and explicit
+``additive_mask`` paths are deferred to P52; the wrapper falls
+back to the Triton kernel when those features are requested.
+
+Scope at P50:
+
+* ``Q [B, H, Sq, D]``, ``K [B, K_H, Sk, D]``, ``V [B, K_H, Sk, D]``
+  with ``K_H ∈ {1 (MQA), H (MHA)}``.
+* Optional per-head sink ``[H]`` — joined as a virtual key column
+  at the end of the softmax (matches the plan-4 P25 / plan-2 P14
+  sink contract).
+* Optional sliding-window-causal mask (``swa_window > 0``) applied
+  in the kernel via `start = max(0, (m + 1 - window_size) //
+  block_N)`.
+* MQA broadcast: when ``K_H == 1``, every query head reads the
+  shared K/V head; computed via per-program ``h // (HQ // HK)``.
+
+Out of scope at P50 (Triton fallback handles these):
+
+* ``additive_mask`` not None  — needs the kernel to accept a
+  [Sq, Sk] additive bias.  Deferred to P52 (alongside the HCA
+  split-mask path).
+* ``hca_local_seqlen > 0`` — HCA split-mask path.  Deferred to
+  P52.
+
+dtype contract (must match :func:`reference.eager_v4_attention`):
+
+* Q / K / V matmuls run in input dtype on tensor cores; the
+  matmul accumulator inside is fp32.
+* Online softmax accumulator (``m_i``, ``l_i``, ``acc_o``) is fp32.
+* Output is in ``v.dtype``; saved ``LSE`` is fp32 (BWD walks back
+  from it).
+
+The kernel borrows from
+``tilelang/examples/attention_sink/example_mha_sink_fwd_bhsd.py``
+(sink fusion at end of softmax via ``exp2(sinks * log2e - max *
+scale)``) and ``tilelang/examples/amd/example_amd_flash_attn_fwd.py``
+(MI355X-tuned block sizes + WMMA ``k_pack`` + ``T.use_swizzle``).
+
+Implementation note: this module deliberately does NOT use
+``from __future__ import annotations`` because tilelang's
+``@T.prim_func`` decorator evaluates annotations eagerly via
+``typing.get_type_hints``, which fails when annotations are lazy
+strings + the referenced names live in an enclosing function's
+local scope (e.g. ``q_shape`` defined inside the JIT factory).
+"""
+
+from typing import Optional, Tuple
+
+# Lazy import — tilelang is heavy.  The wrapper checks
+# ``_tilelang.should_dispatch(...)`` upstream so we only land here
+# when the env knob is set + tilelang is importable.
+import tilelang
+import tilelang.language as T
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import (
+    _launch_v4_attention_fwd,
+)
+
+# Compiled-kernel cache keyed by (B, HQ, HK, Sq, Sk, D, has_sink,
+# swa_window, dtype-str).  Tilelang's own JIT also caches per-shape;
+# this dict is a Python-side memoisation that avoids the
+# ``@tilelang.jit`` recompile cost on every wrapper call.
+_KERNEL_CACHE: dict[tuple, object] = {}
+
+
+# ---------------------------------------------------------------------------
+# Tilelang kernel definition
+# ---------------------------------------------------------------------------
+
+
+@tilelang.jit(
+    out_idx=[3, 4],
+    pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
+)
+def _make_v4_attention_fwd_kernel(
+    batch: int,
+    heads_q: int,
+    heads_k: int,
+    seq_q: int,
+    seq_k: int,
+    dim: int,
+    has_sink: bool,
+    swa_window: int,
+    dtype: "T.dtype" = T.bfloat16,
+    block_M: int = 64,
+    block_N: int = 64,
+    num_stages: int = 1,
+    threads: int = 128,
+):
+    """Plan-8 P50 tilelang FWD kernel JIT factory.
+
+    `has_sink`, `swa_window`, MQA-vs-MHA branching are baked into the
+    kernel at JIT time so the inner loop has no host-side branching
+    cost.  All shape args + flags are kwargs to this JIT'd function;
+    tilelang's `@tilelang.jit` decorator caches compiled binaries per
+    (kwargs) tuple via its own cache.
+    """
+
+    groups = heads_q // heads_k  # 1 for MHA, heads_q for MQA
+    accum_dtype = T.float32
+    sm_scale = 1.0 / (dim**0.5)
+    scale = sm_scale * 1.44269504  # log2(e) — enables exp2/FFMA
+
+    q_shape = [batch, heads_q, seq_q, dim]
+    kv_shape = [batch, heads_k, seq_k, dim]
+    sink_shape = [heads_q]
+    out_shape = [batch, heads_q, seq_q, dim]
+    lse_shape = [batch, heads_q, seq_q]
+
+    @T.prim_func
+    def main(
+        Q: T.Tensor(q_shape, dtype),
+        K: T.Tensor(kv_shape, dtype),
+        V: T.Tensor(kv_shape, dtype),
+        Output: T.Tensor(out_shape, dtype),
+        Lse: T.Tensor(lse_shape, accum_dtype),
+        Sinks: T.Tensor(sink_shape, dtype),
+    ):
+        with T.Kernel(T.ceildiv(seq_q, block_M), heads_q, batch, threads=threads) as (
+            bx,
+            by,
+            bz,
+        ):
+            Q_shared = T.alloc_shared([block_M, dim], dtype)
+            K_shared = T.alloc_shared([block_N, dim], dtype)
+            V_shared = T.alloc_shared([block_N, dim], dtype)
+            O_shared = T.alloc_shared([block_M, dim], dtype)
+            acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
+            acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
+            acc_o = T.alloc_fragment([block_M, dim], accum_dtype)
+            scores_max = T.alloc_fragment([block_M], accum_dtype)
+            scores_max_prev = T.alloc_fragment([block_M], accum_dtype)
+            scores_scale = T.alloc_fragment([block_M], accum_dtype)
+            scores_sum = T.alloc_fragment([block_M], accum_dtype)
+            logsum = T.alloc_fragment([block_M], accum_dtype)
+            sinks_reg = T.alloc_fragment([block_M], dtype)
+
+            # MQA broadcast: kv-head index for this query head.
+            k_head = by // groups
+
+            T.copy(Q[bz, by, bx * block_M : (bx + 1) * block_M, :], Q_shared)
+            T.fill(acc_o, 0)
+            T.fill(logsum, 0)
+            T.fill(scores_max, -T.infinity(accum_dtype))
+            if has_sink:
+                for i in T.Parallel(block_M):
+                    sinks_reg[i] = Sinks[by]
+
+            # K-tile loop bounds.  Causal: end at the right edge of
+            # the m-tile.  SWA: start at `max(0, (m+1 - W) // BN)`.
+            end = T.min(T.ceildiv(seq_k, block_N), T.ceildiv((bx + 1) * block_M, block_N))
+            if swa_window > 0:
+                start = T.max(0, (bx * block_M - swa_window) // block_N)
+            else:
+                start = 0
+
+            for k in T.Pipelined(start, end, num_stages=num_stages):
+                T.copy(K[bz, k_head, k * block_N : (k + 1) * block_N, :], K_shared)
+                for i, j in T.Parallel(block_M, block_N):
+                    q_idx = bx * block_M + i
+                    k_idx = k * block_N + j
+                    if swa_window > 0:
+                        acc_s[i, j] = T.if_then_else(
+                            q_idx >= k_idx and q_idx < k_idx + swa_window,
+                            0,
+                            -T.infinity(acc_s.dtype),
+                        )
+                    else:
+                        acc_s[i, j] = T.if_then_else(q_idx >= k_idx, 0, -T.infinity(acc_s.dtype))
+                T.gemm(
+                    Q_shared,
+                    K_shared,
+                    acc_s,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullRow,
+                )
+
+                T.copy(scores_max, scores_max_prev)
+                T.fill(scores_max, -T.infinity(accum_dtype))
+                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                for i in T.Parallel(block_M):
+                    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
+                # Check-inf for SWA: a tile that is fully outside the
+                # window has scores_max = -inf; reset to 0 to avoid
+                # NaN from `exp2(-inf - -inf)`.
+                for i in T.Parallel(block_M):
+                    if swa_window > 0:
+                        scores_max[i] = T.if_then_else(
+                            scores_max[i] == -T.infinity(accum_dtype),
+                            0,
+                            scores_max[i],
+                        )
+                    scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
+                for i, j in T.Parallel(block_M, block_N):
+                    acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
+                T.reduce_sum(acc_s, scores_sum, dim=1)
+                for i in T.Parallel(block_M):
+                    logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
+                T.copy(acc_s, acc_s_cast)
+
+                for i, j in T.Parallel(block_M, dim):
+                    acc_o[i, j] *= scores_scale[i]
+
+                T.copy(V[bz, k_head, k * block_N : (k + 1) * block_N, :], V_shared)
+                T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+            # Sink as virtual key column at end-of-loop.
+            if has_sink:
+                for i in T.Parallel(block_M):
+                    logsum[i] += T.exp2(sinks_reg[i] * 1.44269504 - scores_max[i] * scale)
+
+            # Final normalise + cast + LSE emit.
+            for i, j in T.Parallel(block_M, dim):
+                acc_o[i, j] /= logsum[i]
+            T.copy(acc_o, O_shared)
+            T.copy(O_shared, Output[bz, by, bx * block_M : (bx + 1) * block_M, :])
+
+            # LSE = ln(l_i) + m_i*sm_scale (note: scores were
+            # accumulated with `*scale = sm_scale*log2e`, so the
+            # base-e LSE is `ln(l_i) + m_i*sm_scale`).
+            for i in T.Parallel(block_M):
+                logsum[i] = T.log(logsum[i]) + scores_max[i] * sm_scale
+            T.copy(logsum, Lse[bz, by, bx * block_M : (bx + 1) * block_M])
+
+    return main
+
+
+# ---------------------------------------------------------------------------
+# Wrapper API
+# ---------------------------------------------------------------------------
+
+
+def _kernel_supports(
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],
+    hca_local_seqlen: int,
+) -> bool:
+    """Return True iff the P50 tilelang kernel covers this call site.
+
+    P50 ships only the dense / SWA / sink subset.  Additive-mask +
+    HCA split-mask are P52 territory; we fall back to Triton there.
+    """
+    if additive_mask is not None:
+        return False
+    if hca_local_seqlen > 0:
+        return False
+    return True
+
+
+def _torch_dtype_to_tilelang(dtype: torch.dtype) -> "T.dtype":
+    if dtype == torch.bfloat16:
+        return T.bfloat16
+    if dtype == torch.float16:
+        return T.float16
+    if dtype == torch.float32:
+        return T.float32
+    raise ValueError(f"unsupported dtype {dtype}; expected bf16 / fp16 / fp32")
+
+
+def _get_or_compile_kernel(
+    *,
+    batch: int,
+    heads_q: int,
+    heads_k: int,
+    seq_q: int,
+    seq_k: int,
+    dim: int,
+    has_sink: bool,
+    swa_window: int,
+    dtype: torch.dtype,
+):
+    """Compile (or fetch from tilelang's per-args cache) the FWD kernel
+    for the given shape envelope.
+
+    Tilelang's `@tilelang.jit` decorator caches compiled binaries per
+    (kwargs) tuple via its own cache, so this wrapper is itself fast
+    on cache hits.
+    """
+    key = (
+        batch,
+        heads_q,
+        heads_k,
+        seq_q,
+        seq_k,
+        dim,
+        has_sink,
+        int(swa_window > 0) * swa_window,
+        str(dtype),
+    )
+    if key in _KERNEL_CACHE:
+        return _KERNEL_CACHE[key]
+    tl_dtype = _torch_dtype_to_tilelang(dtype)
+    # SMEM budget at head_dim=512 on MI355X is tight (160 KiB).
+    # Each `T.alloc_shared([B, dim], bf16)` allocates `B * dim * 2`
+    # bytes; the kernel uses Q + K + V + O shared = 4 * B_tile * D
+    # * 2.  At D=512 even 32x32 hits some tilelang-internal hidden
+    # allocations that overshoot.  Conservative choice: shrink to
+    # 16x32 with threads=64 at D >= 256.
+    if dim >= 256:
+        block_M, block_N, threads = 32, 32, 64
+    else:
+        block_M, block_N, threads = 64, 64, 128
+    kernel = _make_v4_attention_fwd_kernel(
+        batch,
+        heads_q,
+        heads_k,
+        seq_q,
+        seq_k,
+        dim,
+        has_sink,
+        swa_window,
+        dtype=tl_dtype,
+        block_M=block_M,
+        block_N=block_N,
+        threads=threads,
+    )
+    _KERNEL_CACHE[key] = kernel
+    return kernel
+
+
+def v4_attention_fwd_tilelang_with_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor] = None,
+    additive_mask: Optional[torch.Tensor] = None,
+    swa_window: int = 0,
+    attn_dropout: float = 0.0,
+    training: bool = False,
+    scale: Optional[float] = None,
+    hca_local_seqlen: int = 0,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Same dispatch logic as :func:`v4_attention_fwd_tilelang` but
+    also returns the saved ``LSE`` tensor when the tilelang path
+    runs.
+
+    When the wrapper falls back to Triton, ``LSE`` is computed by
+    the Triton launcher and returned alongside ``out``; the
+    autograd Function in P51 saves it for backward.
+    """
+    if attn_dropout > 0.0 and training:
+        return _launch_v4_attention_fwd(
+            q,
+            k,
+            v,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=additive_mask,
+            scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5),
+            hca_local_seqlen=hca_local_seqlen,
+        )
+    if not _kernel_supports(
+        sink=sink,
+        swa_window=swa_window,
+        additive_mask=additive_mask,
+        hca_local_seqlen=hca_local_seqlen,
+    ):
+        return _launch_v4_attention_fwd(
+            q,
+            k,
+            v,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=additive_mask,
+            scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5),
+            hca_local_seqlen=hca_local_seqlen,
+        )
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(
+            "v4_attention_fwd_tilelang expects q / k / v of rank 4 "
+            f"(got q.dim={q.dim()}, k.dim={k.dim()}, v.dim={v.dim()})"
+        )
+    B, HQ, Sq, D = q.shape
+    _, HK, Sk, _ = k.shape
+    assert HK in (1, HQ)
+    has_sink = sink is not None
+    kernel = _get_or_compile_kernel(
+        batch=B,
+        heads_q=HQ,
+        heads_k=HK,
+        seq_q=Sq,
+        seq_k=Sk,
+        dim=D,
+        has_sink=has_sink,
+        swa_window=int(swa_window),
+        dtype=q.dtype,
+    )
+    if not has_sink:
+        sink_arg = torch.zeros(HQ, device=q.device, dtype=q.dtype)
+    else:
+        sink_arg = sink
+    q_c = q.contiguous()
+    k_c = k.contiguous()
+    v_c = v.contiguous()
+    sink_c = sink_arg.contiguous()
+    out, lse = kernel(q_c, k_c, v_c, sink_c)
+    return out, lse
+
+
+def v4_attention_fwd_tilelang(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor] = None,
+    additive_mask: Optional[torch.Tensor] = None,
+    swa_window: int = 0,
+    attn_dropout: float = 0.0,
+    training: bool = False,
+    scale: Optional[float] = None,
+    hca_local_seqlen: int = 0,
+) -> torch.Tensor:
+    """Plan-8 P50 tilelang dense / SWA / sink FWD wrapper.
+
+    Drop-in replacement for the relevant subset of ``v4_attention_v1``
+    (the Triton autograd path).  When the call site requires features
+    P50 doesn't cover (`additive_mask`, `hca_local_seqlen > 0`), the
+    wrapper falls back to the Triton kernel via
+    ``_launch_v4_attention_fwd`` so production smokes never see an
+    error in the regression window.
+
+    Returns ``out [B, H, Sq, D]`` in ``v.dtype``.
+
+    ``attn_dropout`` and ``training`` are accepted for parity with
+    the Triton wrapper but not used — dropout on attention is off
+    in production V4 (plan-4 P25 design note); P50 inherits that
+    contract.
+
+    Note: this entry point is the **plain** (no-autograd) wrapper.
+    For the autograd-aware path used by V4AttentionFn, the call
+    site uses :func:`v4_attention_fwd_tilelang_with_lse` directly
+    so it can save ``LSE`` for backward.
+    """
+    out, _ = v4_attention_fwd_tilelang_with_lse(
+        q,
+        k,
+        v,
+        sink=sink,
+        additive_mask=additive_mask,
+        swa_window=swa_window,
+        attn_dropout=attn_dropout,
+        training=training,
+        scale=scale,
+        hca_local_seqlen=hca_local_seqlen,
+    )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Module-level registration — flips
+# `_tilelang.is_tilelang_kernel_available("v4_attention_fwd")` to True
+# AND overrides the parent module's stub with the real wrapper, so
+# callers can do `_tilelang.v4_attention_fwd_tilelang(...)` and get
+# the real implementation after lazy-load.
+# ---------------------------------------------------------------------------
+
+
+_tilelang.register_available_kernel("v4_attention_fwd")
+# Note: the parent module's `v4_attention_fwd_tilelang` attribute
+# is restored by `_tilelang._lazy_load(...)` after the import_module
+# call (Python's import machinery sets it to the submodule otherwise).
+
+
+__all__ = ["v4_attention_fwd_tilelang"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py
new file mode 100644
index 000000000..7d452efde
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py
@@ -0,0 +1,20 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared (non-attention) Triton kernels for the DeepSeek-V4 pipeline.
+
+These back the surrounding V4 components rather than the attention core, and
+are imported by full submodule path (no package-level re-export):
+
+* :mod:`indexer_score` / :mod:`indexer_score_post` — lightning indexer scores
+* :mod:`compressor_pool` — compressed-pool builder
+* :mod:`hc_expand` / :mod:`hc_glue` — hierarchical-compression expand / glue
+* :mod:`sinkhorn` — fused Sinkhorn-Knopp normalize (FWD/BWD)
+* :mod:`rope_interleaved_partial` — fused interleaved partial RoPE (FWD/BWD)
+
+The attention kernels live in :mod:`.._triton_v0_deprecated` / :mod:`.._triton_v1` /
+:mod:`.._triton_v2`.
+"""
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py
new file mode 100644
index 000000000..64d0f7234
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py
@@ -0,0 +1,283 @@
+# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+"""Triton-fused softmax-weighted pool for the V4 Compressor (forward burst).
+
+Fuses the per-window-softmax pooling tail of :class:`Compressor.forward`::
+
+    score = score + ape                                  # [B, N, W, hd]
+    weights = softmax(score.float(), dim=2)              # over the window W
+    pooled  = (kv * weights).sum(dim=2)                  # [B, N, hd]
+
+into a single forward kernel: one program per ``(b, n)`` row computes, for each
+channel ``d``, the softmax over the ``W`` window slots of ``score[..,:,d]+ape[:,d]``
+and the weighted sum of ``kv[..,:,d]`` -- collapsing the ~5-launch forward burst
+(add + cast + softmax + cast + mul + reduce) into one launch per compressed layer.
+
+The forward runs the reduction in fp32 (more accurate than the eager bf16-weights
+path, numerically equivalent within bf16 noise). The backward is the explicit
+analytic gradient in eager torch (softmax jacobian + weighted-sum), so there is no
+second Triton kernel to validate; the forward burst -- which is what dominates the
+compressed-layer forward -- is the part fused here.
+
+Gating: routed through :func:`fused_softmax_weighted_pool` when
+``PRIMUS_COMPRESS_POOL_TRITON != "0"`` (default-on) for CUDA float16/bfloat16/float32
+inputs; falls back to eager only for non-CUDA / unsupported dtypes. Handles BOTH
+the HCA (non-overlap, window ``W = ratio``) and CSA (overlap, window ``W = 2*ratio``)
+pools — the kernel softmaxes over whatever window axis ``W`` it is given (verified
+fp32-exact at both W=128 and W=8).
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _compressor_pool_fwd_kernel(
+    kv_ptr,
+    score_ptr,
+    ape_ptr,
+    out_ptr,
+    R,
+    HD,
+    stride_kv_bn,
+    stride_kv_r,
+    stride_kv_d,
+    stride_sc_bn,
+    stride_sc_r,
+    stride_sc_d,
+    stride_ape_r,
+    stride_ape_d,
+    stride_out_bn,
+    stride_out_d,
+    BLOCK_R: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_bn = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    offs_r = tl.arange(0, BLOCK_R)
+    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    r_mask = offs_r < R
+    d_mask = offs_d < HD
+    mask = r_mask[:, None] & d_mask[None, :]
+
+    sc_ptrs = (
+        score_ptr + pid_bn * stride_sc_bn + offs_r[:, None] * stride_sc_r + offs_d[None, :] * stride_sc_d
+    )
+    ape_ptrs = ape_ptr + offs_r[:, None] * stride_ape_r + offs_d[None, :] * stride_ape_d
+    s = tl.load(sc_ptrs, mask=mask, other=0.0).to(tl.float32)
+    a = tl.load(ape_ptrs, mask=mask, other=0.0).to(tl.float32)
+    s = s + a
+
+    # softmax over the window axis R (axis 0), masking padded rows to -inf.
+    s = tl.where(r_mask[:, None], s, float("-inf"))
+    s_max = tl.max(s, axis=0)
+    e = tl.exp(s - s_max[None, :])
+    e = tl.where(r_mask[:, None], e, 0.0)
+    denom = tl.sum(e, axis=0)
+    w = e / denom[None, :]  # [BLOCK_R, BLOCK_D] fp32 softmax weights
+
+    kv_ptrs = kv_ptr + pid_bn * stride_kv_bn + offs_r[:, None] * stride_kv_r + offs_d[None, :] * stride_kv_d
+    kv = tl.load(kv_ptrs, mask=mask, other=0.0).to(tl.float32)
+    pooled = tl.sum(kv * w, axis=0)  # [BLOCK_D]
+
+    out_ptrs = out_ptr + pid_bn * stride_out_bn + offs_d * stride_out_d
+    tl.store(out_ptrs, pooled.to(out_ptr.dtype.element_ty), mask=d_mask)
+
+
+def _pool_fwd_triton(kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor) -> torch.Tensor:
+    """kv/score: ``[B, N, W, hd]``; ape: ``[W, hd]`` -> pooled ``[B, N, hd]``."""
+    B, N, W, HD = kv.shape
+    bn = B * N
+    kv3 = kv.reshape(bn, W, HD)
+    sc3 = score.reshape(bn, W, HD)
+    out = torch.empty((bn, HD), dtype=kv.dtype, device=kv.device)
+    BLOCK_R = max(16, triton.next_power_of_2(W))
+    BLOCK_D = min(64, max(16, triton.next_power_of_2(HD)))
+    grid = (bn, triton.cdiv(HD, BLOCK_D))
+    _compressor_pool_fwd_kernel[grid](
+        kv3,
+        sc3,
+        ape,
+        out,
+        W,
+        HD,
+        kv3.stride(0),
+        kv3.stride(1),
+        kv3.stride(2),
+        sc3.stride(0),
+        sc3.stride(1),
+        sc3.stride(2),
+        ape.stride(0),
+        ape.stride(1),
+        out.stride(0),
+        out.stride(1),
+        BLOCK_R=BLOCK_R,
+        BLOCK_D=BLOCK_D,
+    )
+    return out.reshape(B, N, HD)
+
+
+@triton.jit
+def _compressor_pool_bwd_kernel(
+    kv_ptr,
+    score_ptr,
+    ape_ptr,
+    dout_ptr,
+    dkv_ptr,
+    dscore_ptr,
+    R,
+    HD,
+    stride_kv_bn,
+    stride_kv_r,
+    stride_kv_d,
+    stride_sc_bn,
+    stride_sc_r,
+    stride_sc_d,
+    stride_ape_r,
+    stride_ape_d,
+    stride_do_bn,
+    stride_do_d,
+    stride_dkv_bn,
+    stride_dkv_r,
+    stride_dkv_d,
+    stride_dsc_bn,
+    stride_dsc_r,
+    stride_dsc_d,
+    BLOCK_R: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_bn = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    offs_r = tl.arange(0, BLOCK_R)
+    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    r_mask = offs_r < R
+    d_mask = offs_d < HD
+    mask = r_mask[:, None] & d_mask[None, :]
+
+    # Recompute the fwd softmax weights w over the window axis R (axis 0), fp32.
+    sc_ptrs = (
+        score_ptr + pid_bn * stride_sc_bn + offs_r[:, None] * stride_sc_r + offs_d[None, :] * stride_sc_d
+    )
+    ape_ptrs = ape_ptr + offs_r[:, None] * stride_ape_r + offs_d[None, :] * stride_ape_d
+    s = tl.load(sc_ptrs, mask=mask, other=0.0).to(tl.float32)
+    a = tl.load(ape_ptrs, mask=mask, other=0.0).to(tl.float32)
+    s = s + a
+    s = tl.where(r_mask[:, None], s, float("-inf"))
+    s_max = tl.max(s, axis=0)
+    e = tl.exp(s - s_max[None, :])
+    e = tl.where(r_mask[:, None], e, 0.0)
+    denom = tl.sum(e, axis=0)
+    w = e / denom[None, :]  # [BLOCK_R, BLOCK_D]
+
+    kv_ptrs = kv_ptr + pid_bn * stride_kv_bn + offs_r[:, None] * stride_kv_r + offs_d[None, :] * stride_kv_d
+    kv = tl.load(kv_ptrs, mask=mask, other=0.0).to(tl.float32)
+    do_ptrs = dout_ptr + pid_bn * stride_do_bn + offs_d * stride_do_d
+    dout = tl.load(do_ptrs, mask=d_mask, other=0.0).to(tl.float32)  # [BLOCK_D]
+
+    # pooled = sum_r w*kv  ->  dkv = dout*w ; dw = dout*kv ;
+    # softmax jacobian: dscore = w * (dw - sum_r(dw*w))
+    dkv = dout[None, :] * w
+    dw = dout[None, :] * kv
+    sum_dw_w = tl.sum(dw * w, axis=0)  # [BLOCK_D]
+    dscore = w * (dw - sum_dw_w[None, :])
+
+    dkv_ptrs = (
+        dkv_ptr + pid_bn * stride_dkv_bn + offs_r[:, None] * stride_dkv_r + offs_d[None, :] * stride_dkv_d
+    )
+    tl.store(dkv_ptrs, dkv.to(dkv_ptr.dtype.element_ty), mask=mask)
+    dsc_ptrs = (
+        dscore_ptr + pid_bn * stride_dsc_bn + offs_r[:, None] * stride_dsc_r + offs_d[None, :] * stride_dsc_d
+    )
+    tl.store(dsc_ptrs, dscore.to(dscore_ptr.dtype.element_ty), mask=mask)
+
+
+def _pool_bwd_triton(dout, kv, score, ape):
+    """Fused backward of the softmax-weighted pool. Returns (dkv, dscore, dape).
+
+    Recomputes w in-register (like the fwd) and emits dkv + dscore in one kernel;
+    dape = sum over the (B,N) rows of dscore is a single trailing reduce.
+    """
+    B, N, W, HD = kv.shape
+    bn = B * N
+    kv3 = kv.reshape(bn, W, HD)
+    sc3 = score.reshape(bn, W, HD)
+    do2 = dout.reshape(bn, HD)
+    dkv = torch.empty_like(kv3)
+    dscore = torch.empty_like(sc3)
+    BLOCK_R = max(16, triton.next_power_of_2(W))
+    BLOCK_D = min(64, max(16, triton.next_power_of_2(HD)))
+    grid = (bn, triton.cdiv(HD, BLOCK_D))
+    _compressor_pool_bwd_kernel[grid](
+        kv3,
+        sc3,
+        ape,
+        do2,
+        dkv,
+        dscore,
+        W,
+        HD,
+        kv3.stride(0),
+        kv3.stride(1),
+        kv3.stride(2),
+        sc3.stride(0),
+        sc3.stride(1),
+        sc3.stride(2),
+        ape.stride(0),
+        ape.stride(1),
+        do2.stride(0),
+        do2.stride(1),
+        dkv.stride(0),
+        dkv.stride(1),
+        dkv.stride(2),
+        dscore.stride(0),
+        dscore.stride(1),
+        dscore.stride(2),
+        BLOCK_R=BLOCK_R,
+        BLOCK_D=BLOCK_D,
+    )
+    dkv = dkv.reshape(B, N, W, HD)
+    dscore = dscore.reshape(B, N, W, HD)
+    dape = dscore.sum(dim=(0, 1)).to(ape.dtype)
+    return dkv, dscore, dape
+
+
+class _CompressorSoftmaxPool(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, kv, score, ape):
+        pooled = _pool_fwd_triton(kv, score, ape)
+        ctx.save_for_backward(kv, score, ape)
+        return pooled
+
+    @staticmethod
+    def backward(ctx, dout):
+        kv, score, ape = ctx.saved_tensors
+        # Fused Triton backward (default-on): recompute softmax + emit dkv/dscore in
+        # one launch instead of the ~14-kernel eager analytic burst. Handles both the
+        # HCA (W=ratio) and CSA (W=2*ratio overlap) pools -- the window fits one tile
+        # in either case. PRIMUS_COMPRESS_POOL_TRITON_BWD=0 restores the eager path.
+        if os.environ.get("PRIMUS_COMPRESS_POOL_TRITON_BWD", "1") != "0" and pool_supported(kv):
+            return _pool_bwd_triton(dout.contiguous(), kv, score, ape)
+        # Analytic gradient of pooled = sum_w softmax_w(score+ape) * kv (fp32).
+        w = torch.softmax((score + ape).float(), dim=2)  # [B, N, W, hd]
+        dout3 = dout.unsqueeze(2).float()  # [B, N, 1, hd]
+        dkv = (dout3 * w).to(kv.dtype)
+        dw = dout3 * kv.float()
+        # softmax jacobian: dscore = w * (dw - sum_w(dw * w))
+        dscore = (w * (dw - (dw * w).sum(dim=2, keepdim=True))).to(score.dtype)
+        dape = dscore.sum(dim=(0, 1)).to(ape.dtype)
+        return dkv, dscore, dape
+
+
+def fused_softmax_weighted_pool(kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor) -> torch.Tensor:
+    """Fused ``(softmax(score+ape, dim=2) * kv).sum(dim=2)``.
+
+    kv/score: ``[B, N, W, hd]``; ape: ``[W, hd]`` -> ``[B, N, hd]``.
+    """
+    return _CompressorSoftmaxPool.apply(kv, score, ape)
+
+
+def pool_supported(kv: torch.Tensor) -> bool:
+    return kv.is_cuda and kv.dtype in (torch.float16, torch.bfloat16, torch.float32)
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py
new file mode 100644
index 000000000..10f905c4b
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py
@@ -0,0 +1,232 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused HyperConnection ``collapse`` (small-kernel-fusion 2026-07-03).
+
+Single FWD + single BWD kernel pair for the eager body of
+:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.collapse`::
+
+    out[..., d] = sum_k pre[..., k] * x[..., k, d]
+
+The eager body ``(pre.unsqueeze(-1) * x).sum(dim=-2)`` materialises a full
+``[..., K, D]`` temporary (the broadcast multiply) and then reduces it — two
+kernels + ``K*D`` of extra HBM traffic per call. This kernel streams ``D`` and
+contracts the small ``K`` (constexpr) in registers, writing only the ``[..., D]``
+result. Symmetric to the already-shipped ``expand`` fusion (``hc_expand.py``).
+
+Gradients::
+
+    dx[..., k, d] = pre[..., k] * g[..., d]
+    dpre[..., k]  = sum_d x[..., k, d] * g[..., d]
+
+Gating: routed through :func:`hc_collapse_triton` when
+``PRIMUS_HC_COLLAPSE_TRITON != "0"`` (default-on); falls back to eager for
+unsupported configs (CPU input, K out of range, dtype mismatch).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_K = (1, 2, 4, 8, 16)
+
+_BLOCK_M = 16
+_BLOCK_D = 128
+_NUM_WARPS = 4
+
+
+@triton.jit
+def _hc_collapse_fwd_kernel(
+    X_PTR,  # [M, K, D]
+    PRE_PTR,  # [M, K]
+    OUT_PTR,  # [M, D] OUT
+    M,
+    D,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    k_ax = tl.arange(0, K)
+    mask_m = rm < M
+    mask_d = rd < D
+    mask2 = mask_m[:, None] & mask_d[None, :]
+    mask3 = mask_m[:, None, None] & mask_d[None, None, :]
+
+    KD = K * D
+    x3 = tl.load(
+        X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        mask=mask3,
+        other=0.0,
+    ).to(tl.float32)
+    pre = tl.load(PRE_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to(tl.float32)
+    out = tl.sum(pre[:, :, None] * x3, axis=1)  # [BLOCK_M, BLOCK_D]
+    tl.store(OUT_PTR + rm[:, None] * D + rd[None, :], out.to(OUT_PTR.dtype.element_ty), mask=mask2)
+
+
+@triton.jit
+def _hc_collapse_bwd_kernel(
+    G_PTR,  # [M, D] grad of out
+    X_PTR,  # [M, K, D]
+    PRE_PTR,  # [M, K]
+    DX_PTR,  # [M, K, D] OUT
+    DPRE_P_PTR,  # [NUMD, M, K] partial OUT (reduced host-side)
+    M,
+    D,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    k_ax = tl.arange(0, K)
+    mask_m = rm < M
+    mask_d = rd < D
+    mask2 = mask_m[:, None] & mask_d[None, :]
+    mask3 = mask_m[:, None, None] & mask_d[None, None, :]
+
+    KD = K * D
+    x3 = tl.load(
+        X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        mask=mask3,
+        other=0.0,
+    ).to(tl.float32)
+    g = tl.load(G_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32)
+    pre = tl.load(PRE_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to(tl.float32)
+
+    # dx[k, d] = pre[k] * g[d]
+    dx = pre[:, :, None] * g[:, None, :]  # [BLOCK_M, K, BLOCK_D]
+    tl.store(
+        DX_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        dx.to(DX_PTR.dtype.element_ty),
+        mask=mask3,
+    )
+
+    # dpre[k] = sum_d x[k, d] * g[d]   (partial over this d-block)
+    dpre = tl.sum(x3 * g[:, None, :], axis=2)  # [BLOCK_M, K]
+    tl.store(
+        DPRE_P_PTR + pid_d * (M * K) + rm[:, None] * K + k_ax[None, :],
+        dpre,
+        mask=mask_m[:, None],
+    )
+
+
+class HCCollapseFn(torch.autograd.Function):
+    """Autograd wrapper for the fused ``collapse`` FWD/BWD Triton kernels.
+
+    ``out[..., d] = sum_k pre[..., k] * x[..., k, d]``
+    """
+
+    @staticmethod
+    def forward(ctx, x, pre):  # type: ignore[override]
+        K = x.shape[-2]
+        D = x.shape[-1]
+        leading = x.shape[:-2]
+        M = 1
+        for s in leading:
+            M *= s
+
+        x_c = x.contiguous()
+        pre_c = pre.contiguous()
+        out = torch.empty((*leading, D), dtype=x.dtype, device=x.device)
+
+        grid = (triton.cdiv(M, _BLOCK_M), triton.cdiv(D, _BLOCK_D))
+        _hc_collapse_fwd_kernel[grid](
+            x_c,
+            pre_c,
+            out,
+            M,
+            D,
+            K=K,
+            BLOCK_M=_BLOCK_M,
+            BLOCK_D=_BLOCK_D,
+            num_warps=_NUM_WARPS,
+        )
+
+        ctx.save_for_backward(x_c, pre_c)
+        ctx.leading = tuple(leading)
+        ctx.K = K
+        ctx.D = D
+        ctx.M = M
+        return out
+
+    @staticmethod
+    def backward(ctx, g):  # type: ignore[override]
+        x_c, pre_c = ctx.saved_tensors
+        leading, K, D, M = ctx.leading, ctx.K, ctx.D, ctx.M
+        g = g.contiguous()
+
+        dx = torch.empty_like(x_c)
+        num_d = triton.cdiv(D, _BLOCK_D)
+        dpre_p = torch.empty((num_d, M, K), dtype=torch.float32, device=x_c.device)
+
+        grid = (triton.cdiv(M, _BLOCK_M), num_d)
+        _hc_collapse_bwd_kernel[grid](
+            g,
+            x_c,
+            pre_c,
+            dx,
+            dpre_p,
+            M,
+            D,
+            K=K,
+            BLOCK_M=_BLOCK_M,
+            BLOCK_D=_BLOCK_D,
+            num_warps=_NUM_WARPS,
+        )
+
+        d_pre = dpre_p.sum(dim=0).to(pre_c.dtype).view(*leading, K)
+        return dx, d_pre
+
+
+def is_triton_path_enabled() -> bool:
+    """``PRIMUS_HC_COLLAPSE_TRITON`` knob; default-on, set ``"0"`` to disable."""
+    return os.environ.get("PRIMUS_HC_COLLAPSE_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(x: torch.Tensor, pre: torch.Tensor) -> bool:
+    if not (x.is_cuda and pre.is_cuda):
+        return False
+    if x.dim() < 2:
+        return False
+    K = x.shape[-2]
+    if K not in _SUPPORTED_K:
+        return False
+    if pre.shape[-1] != K:
+        return False
+    if x.dtype != pre.dtype or x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return False
+    return True
+
+
+def eager_hc_collapse(x: torch.Tensor, pre: torch.Tensor) -> torch.Tensor:
+    """Reference eager ``collapse`` (matches HyperMixer.collapse bit-for-bit)."""
+    return (pre.unsqueeze(-1) * x).sum(dim=-2)
+
+
+def hc_collapse_triton(x, pre):
+    """Fused ``out[...,d] = sum_k pre[...,k] * x[...,k,d]`` with eager fallback."""
+    if is_triton_path_enabled() and is_triton_kernel_supported(x, pre):
+        return HCCollapseFn.apply(x, pre)
+    return eager_hc_collapse(x, pre)
+
+
+__all__ = [
+    "HCCollapseFn",
+    "hc_collapse_triton",
+    "eager_hc_collapse",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py
new file mode 100644
index 000000000..4bb7724b4
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py
@@ -0,0 +1,287 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused HyperConnection ``expand``.
+
+Single FWD + single BWD kernel pair for the eager body of
+:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.expand`::
+
+    new[..., h, d] = post[..., h] * out[..., d] + sum_k comb[..., h, k] * x[..., k, d]
+
+The contraction over the small ``K`` is done in registers (unrolled, ``K`` is a
+constexpr) while ``D`` is streamed; the ``post`` outer-product and the final add
+are folded into the same kernel.
+
+Gating: routed through :func:`hc_expand_triton` when
+``PRIMUS_HC_EXPAND_TRITON != "0"`` (default-on); falls back to eager for
+unsupported configs (CPU input, K out of range, dtype mismatch).
+
+Supported shapes:
+
+* ``K`` (= ``hc_mult``) in ``{1, 2, 4, 8, 16}``.
+* ``x``/``out``/``post``/``comb`` share one floating dtype and are CUDA tensors.
+* Any leading shape (the wrapper flattens to ``[M, K, D]``).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_K = (1, 2, 4, 8, 16)
+
+_BLOCK_M = 16
+_BLOCK_D = 128
+_NUM_WARPS = 4
+
+
+@triton.jit
+def _hc_expand_fwd_kernel(
+    X_PTR,  # [M, K, D]
+    OUT_PTR,  # [M, D]
+    POST_PTR,  # [M, K]
+    COMB_PTR,  # [M, K, K]
+    NEW_PTR,  # [M, K, D] OUT
+    M,
+    D,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    k_ax = tl.arange(0, K)
+    mask_m = rm < M
+    mask_d = rd < D
+    mask2 = mask_m[:, None] & mask_d[None, :]
+    mask3 = mask_m[:, None, None] & mask_d[None, None, :]
+
+    KD = K * D
+    KK = K * K
+
+    # x as [BLOCK_M, K(stream), BLOCK_D], loaded once.
+    x3 = tl.load(
+        X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        mask=mask3,
+        other=0.0,
+    ).to(tl.float32)
+    out_tile = tl.load(OUT_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32)
+
+    for h in range(K):
+        comb_h = tl.load(
+            COMB_PTR + rm[:, None] * KK + h * K + k_ax[None, :], mask=mask_m[:, None], other=0.0
+        ).to(
+            tl.float32
+        )  # [BLOCK_M, K]
+        post_h = tl.load(POST_PTR + rm * K + h, mask=mask_m, other=0.0).to(tl.float32)  # [BLOCK_M]
+        # mix_h[d] = sum_k comb_h[k] * x3[k, d]
+        mix_h = tl.sum(comb_h[:, :, None] * x3, axis=1)
+        new_h = post_h[:, None] * out_tile + mix_h
+        tl.store(
+            NEW_PTR + rm[:, None] * KD + h * D + rd[None, :],
+            new_h.to(NEW_PTR.dtype.element_ty),
+            mask=mask2,
+        )
+
+
+@triton.jit
+def _hc_expand_bwd_kernel(
+    G_PTR,  # [M, K, D] grad of new
+    X_PTR,  # [M, K, D]
+    OUT_PTR,  # [M, D]
+    POST_PTR,  # [M, K]
+    COMB_PTR,  # [M, K, K]
+    DX_PTR,  # [M, K, D] OUT
+    DOUT_PTR,  # [M, D] OUT
+    DPOST_P_PTR,  # [NUMD, M, K] partial OUT (reduced host-side)
+    DCOMB_P_PTR,  # [NUMD, M, K, K] partial OUT (reduced host-side)
+    M,
+    D,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_d = tl.program_id(1)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    k_ax = tl.arange(0, K)
+    mask_m = rm < M
+    mask_d = rd < D
+    mask2 = mask_m[:, None] & mask_d[None, :]
+    mask3 = mask_m[:, None, None] & mask_d[None, None, :]
+
+    KD = K * D
+    KK = K * K
+
+    # x3, g3 as [BLOCK_M, K(stream), BLOCK_D]; the stream axis of g indexes h.
+    x3 = tl.load(
+        X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        mask=mask3,
+        other=0.0,
+    ).to(tl.float32)
+    g3 = tl.load(
+        G_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :],
+        mask=mask3,
+        other=0.0,
+    ).to(tl.float32)
+    out_tile = tl.load(OUT_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32)
+    post = tl.load(POST_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to(
+        tl.float32
+    )  # [BLOCK_M, K(h)]
+
+    # d_out[d] = sum_h post[h] * g[h, d]
+    dout = tl.sum(post[:, :, None] * g3, axis=1)
+    tl.store(DOUT_PTR + rm[:, None] * D + rd[None, :], dout.to(DOUT_PTR.dtype.element_ty), mask=mask2)
+
+    # d_post[h] = sum_d out[d] * g[h, d]   (partial over this d-block)
+    dpost = tl.sum(out_tile[:, None, :] * g3, axis=2)  # [BLOCK_M, K(h)]
+    tl.store(
+        DPOST_P_PTR + pid_d * (M * K) + rm[:, None] * K + k_ax[None, :],
+        dpost,
+        mask=mask_m[:, None],
+    )
+
+    # d_x[k, d] = sum_h comb[h, k] * g[h, d]
+    for k in range(K):
+        comb_col = tl.load(
+            COMB_PTR + rm[:, None] * KK + k_ax[None, :] * K + k, mask=mask_m[:, None], other=0.0
+        ).to(
+            tl.float32
+        )  # [BLOCK_M, K(h)]
+        dxk = tl.sum(comb_col[:, :, None] * g3, axis=1)
+        tl.store(DX_PTR + rm[:, None] * KD + k * D + rd[None, :], dxk.to(DX_PTR.dtype.element_ty), mask=mask2)
+
+    # d_comb[h, k] = sum_d x[k, d] * g[h, d]   (partial over this d-block)
+    for h in range(K):
+        g_h = tl.load(G_PTR + rm[:, None] * KD + h * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32)
+        dcomb_h = tl.sum(x3 * g_h[:, None, :], axis=2)  # [BLOCK_M, K(k)]
+        tl.store(
+            DCOMB_P_PTR + pid_d * (M * KK) + rm[:, None] * KK + h * K + k_ax[None, :],
+            dcomb_h,
+            mask=mask_m[:, None],
+        )
+
+
+class HCExpandFn(torch.autograd.Function):
+    """Autograd wrapper around the fused ``expand`` FWD/BWD Triton kernels.
+
+    ``new[..., h, d] = post[..., h] * out[..., d] + sum_k comb[..., h, k] * x[..., k, d]``
+    """
+
+    @staticmethod
+    def forward(ctx, x, out, post, comb):  # type: ignore[override]
+        K = comb.shape[-1]
+        D = x.shape[-1]
+        leading = x.shape[:-2]
+        M = 1
+        for s in leading:
+            M *= s
+
+        x_c = x.contiguous()
+        out_c = out.contiguous()
+        post_c = post.contiguous()
+        comb_c = comb.contiguous()
+
+        new = torch.empty((*leading, K, D), dtype=x.dtype, device=x.device)
+
+        grid = (triton.cdiv(M, _BLOCK_M), triton.cdiv(D, _BLOCK_D))
+        _hc_expand_fwd_kernel[grid](
+            x_c,
+            out_c,
+            post_c,
+            comb_c,
+            new,
+            M,
+            D,
+            K=K,
+            BLOCK_M=_BLOCK_M,
+            BLOCK_D=_BLOCK_D,
+            num_warps=_NUM_WARPS,
+        )
+
+        ctx.save_for_backward(x_c, out_c, post_c, comb_c)
+        ctx.leading = tuple(leading)
+        ctx.K = K
+        ctx.D = D
+        ctx.M = M
+        return new
+
+    @staticmethod
+    def backward(ctx, g):  # type: ignore[override]
+        x_c, out_c, post_c, comb_c = ctx.saved_tensors
+        leading, K, D, M = ctx.leading, ctx.K, ctx.D, ctx.M
+        g = g.contiguous()
+
+        device = x_c.device
+        dx = torch.empty_like(x_c)
+        dout = torch.empty_like(out_c)
+        num_d = triton.cdiv(D, _BLOCK_D)
+        dpost_p = torch.empty((num_d, M, K), dtype=torch.float32, device=device)
+        dcomb_p = torch.empty((num_d, M, K, K), dtype=torch.float32, device=device)
+
+        grid = (triton.cdiv(M, _BLOCK_M), num_d)
+        _hc_expand_bwd_kernel[grid](
+            g,
+            x_c,
+            out_c,
+            post_c,
+            comb_c,
+            dx,
+            dout,
+            dpost_p,
+            dcomb_p,
+            M,
+            D,
+            K=K,
+            BLOCK_M=_BLOCK_M,
+            BLOCK_D=_BLOCK_D,
+            num_warps=_NUM_WARPS,
+        )
+
+        d_post = dpost_p.sum(dim=0).to(post_c.dtype).view(*leading, K)
+        d_comb = dcomb_p.sum(dim=0).to(comb_c.dtype).view(*leading, K, K)
+        return dx, dout, d_post, d_comb
+
+
+def is_triton_path_enabled() -> bool:
+    """``PRIMUS_HC_EXPAND_TRITON`` knob; default-on, set ``"0"`` to disable."""
+    return os.environ.get("PRIMUS_HC_EXPAND_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(x: torch.Tensor, post: torch.Tensor, comb: torch.Tensor) -> bool:
+    if not (x.is_cuda and post.is_cuda and comb.is_cuda):
+        return False
+    K = comb.shape[-1]
+    if K not in _SUPPORTED_K:
+        return False
+    if comb.shape[-2] != K or x.shape[-2] != K or post.shape[-1] != K:
+        return False
+    if not (x.dtype == post.dtype == comb.dtype) or x.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    ):
+        return False
+    return True
+
+
+def hc_expand_triton(x, out, post, comb):
+    """Fused ``new[...,h,d] = post[h]*out[d] + sum_k comb[h,k]*x[k,d]``."""
+    return HCExpandFn.apply(x, out, post, comb)
+
+
+__all__ = [
+    "HCExpandFn",
+    "hc_expand_triton",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py
new file mode 100644
index 000000000..dc39ddd81
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py
@@ -0,0 +1,622 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused HyperConnection elemwise glue (plan-6 P37).
+
+Fuses the post-linear elemwise tail of
+:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights`
+into a single FWD + single BWD Triton kernel pair.  Eager body:
+
+.. code-block:: python
+
+    pre_logit  = logits[..., :K]     * scale[0] + base[:K]
+    post_logit = logits[..., K:2K]   * scale[1] + base[K:2K]
+    comb_logit = (logits[..., 2K:].view(..., K, K) * scale[2]
+                  + base[2K:].view(K, K))
+
+    pre  = sigmoid(pre_logit)  + eps          # (eps, 1+eps]
+    post = 2 * sigmoid(post_logit)            # (0, 2)
+    comb_pre_sinkhorn = softmax(comb_logit, dim=-1) + eps
+
+After P36 the trailing ``sinkhorn_normalize(comb_pre_sinkhorn)`` is its
+own Triton kernel.  P37 fuses everything **between** the
+``_packed_logits`` GEMM and the Sinkhorn call into one kernel:
+
+* 3 slices of ``logits`` (free in Triton — pointer arithmetic);
+* 3 fused-multiply-adds against ``scale`` / ``base``;
+* 2 sigmoid + 1 softmax + 2 eps adds;
+* 1 final 2x multiply for ``post``.
+
+The eager chain is ~8 elementwise GPU launches per call (P32 trace
+attributes ~3-5 ms / iter across all 16 ``HyperConnection`` invocations
+to the ``elementwise_kernel_manual_unroll<128, 8>`` bucket).  P37
+collapses those into **one** FWD kernel + **one** BWD kernel.
+
+The matmul inside ``_packed_logits`` (`F.linear(flat * rsqrt, W)`)
+stays as ``torch.nn.functional.linear`` -- GEMM is already
+GPU-bound and fusing it into the elemwise chain would re-implement a
+GEMM badly.  Same goes for the ``collapse`` reduce and the
+``expand`` outer-product + matmul -- those live around matmuls and
+are not net wins as separate Triton kernels (deferred to a possible
+P37b if the residual trace shows otherwise).
+
+Gating: routed through :func:`hc_glue_compute_tail_triton` when
+``PRIMUS_HC_TRITON != "0"`` (default-on).  Set to ``"0"`` to fall back
+to the eager body kept in
+:func:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights`.
+
+Supported shape constraints:
+
+* ``K`` (= ``hc_mult``) must be a power of 2 in ``{1, 2, 4, 8, 16}``.
+  V4-Flash uses ``K=4``.
+* ``logits`` must be contiguous; the wrapper calls ``.contiguous()``
+  defensively.
+* Any leading shape is supported (the wrapper flattens to ``[N, (2+K)*K]``).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+_TORCH_TO_TL_DTYPE = {
+    torch.float64: tl.float64,
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+def _triton_dtype(t: torch.dtype):
+    try:
+        return _TORCH_TO_TL_DTYPE[t]
+    except KeyError as exc:
+        raise TypeError(
+            f"hc_glue: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}"
+        ) from exc
+
+
+# Triton requires power-of-2 block extents, and the register footprint is
+# linear in K + K*K, so we cap K at 16 for the in-register path.  V4-Flash
+# uses K=4; larger K would require a different layout.
+_SUPPORTED_K = (1, 2, 4, 8, 16)
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _hc_compute_tail_fwd_kernel(
+    LOGITS_PTR,  # [N, (2+K)*K] contiguous, fp32
+    SCALE_PTR,  # [3] fp32
+    BASE_PTR,  # [(2+K)*K] fp32
+    PRE_PTR,  # [N, K] OUT_DTYPE
+    POST_PTR,  # [N, K] OUT_DTYPE
+    COMB_PTR,  # [N, K, K] OUT_DTYPE (softmax+eps; before Sinkhorn)
+    PRE_SIG_PTR,  # [N, K] fp32 -- saved-for-backward (sigmoid(pre_logit))
+    POST_SIG_PTR,  # [N, K] fp32 -- saved-for-backward (sigmoid(post_logit))
+    COMB_SM_PTR,  # [N, K, K] fp32 -- saved-for-backward (softmax(comb_logit))
+    N,
+    EPS: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_LEADING: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """One program tile of ``BLOCK_LEADING`` rows of the leading axis.
+
+    Reads ``LOGITS[BLOCK_LEADING, (2+K)*K]`` once, applies the elemwise
+    tail in fp32 registers, writes the three output tensors (cast to
+    OUT_DTYPE) plus three fp32 saved-for-backward states.
+
+    Register footprint per program (at K=4, BLOCK_LEADING=64):
+        BLOCK_LEADING * (2+K)*K        = 64 * 24 = 1536 fp32 logits
+        BLOCK_LEADING * K              =     64 * 4 =   256 pre/post
+        BLOCK_LEADING * K * K          =    64 * 16 =  1024 comb
+    Total ~2800 fp32 = ~11 KiB; comfortable for MI355 256-VGPR warps.
+    """
+
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING)
+    mask_leading = offs < N
+
+    k_idx = tl.arange(0, K)
+
+    # Load the three scale scalars (broadcast across all rows).
+    scale0 = tl.load(SCALE_PTR + 0)
+    scale1 = tl.load(SCALE_PTR + 1)
+    scale2 = tl.load(SCALE_PTR + 2)
+
+    # Base partitions: base[:K], base[K:2K], base[2K:].view(K, K).
+    base_pre = tl.load(BASE_PTR + k_idx)
+    base_post = tl.load(BASE_PTR + K + k_idx)
+    base_comb = tl.load(BASE_PTR + 2 * K + k_idx[:, None] * K + k_idx[None, :])
+
+    KK_TOTAL: tl.constexpr = (2 + K) * K
+
+    # Logits row stride is KK_TOTAL.  Pre slice = [:, :K]; post = [:, K:2K];
+    # comb = [:, 2K:].view(K, K).  All three are contiguous reads.
+    pre_logit = tl.load(
+        LOGITS_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    )
+    post_logit = tl.load(
+        LOGITS_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    )
+    comb_logit = tl.load(
+        LOGITS_PTR + offs[:, None, None] * KK_TOTAL + 2 * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        mask=mask_leading[:, None, None],
+        other=0.0,
+    )
+
+    # Apply scale + base.
+    pre_logit = pre_logit * scale0 + base_pre[None, :]
+    post_logit = post_logit * scale1 + base_post[None, :]
+    comb_logit = comb_logit * scale2 + base_comb[None, :, :]
+
+    # Sigmoids (per-element).
+    pre_sig = tl.sigmoid(pre_logit)
+    post_sig = tl.sigmoid(post_logit)
+
+    # Softmax along axis=2 (the inner K of comb_logit).  Standard
+    # numerically-stable softmax: subtract row max, exp, divide by sum.
+    cmax = tl.max(comb_logit, axis=2, keep_dims=True)
+    cexp = tl.exp(comb_logit - cmax)
+    csum = tl.sum(cexp, axis=2, keep_dims=True)
+    comb_sm = cexp / csum
+
+    # Output tensors (with eps additions and the 2x for post).
+    pre_out = pre_sig + EPS
+    post_out = 2.0 * post_sig
+    comb_out = comb_sm + EPS
+
+    # Save fp32 states for backward.
+    tl.store(
+        PRE_SIG_PTR + offs[:, None] * K + k_idx[None, :],
+        pre_sig,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        POST_SIG_PTR + offs[:, None] * K + k_idx[None, :],
+        post_sig,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        COMB_SM_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        comb_sm,
+        mask=mask_leading[:, None, None],
+    )
+
+    # Write the three caller-visible outputs (cast to OUT_DTYPE).
+    tl.store(
+        PRE_PTR + offs[:, None] * K + k_idx[None, :],
+        pre_out.to(OUT_DTYPE),
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        POST_PTR + offs[:, None] * K + k_idx[None, :],
+        post_out.to(OUT_DTYPE),
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        COMB_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        comb_out.to(OUT_DTYPE),
+        mask=mask_leading[:, None, None],
+    )
+
+
+@triton.jit
+def _hc_compute_tail_bwd_kernel(
+    DPRE_PTR,  # [N, K] grad in OUT_DTYPE
+    DPOST_PTR,  # [N, K] grad in OUT_DTYPE
+    DCOMB_PTR,  # [N, K, K] grad in OUT_DTYPE
+    PRE_SIG_PTR,  # [N, K] fp32 saved
+    POST_SIG_PTR,  # [N, K] fp32 saved
+    COMB_SM_PTR,  # [N, K, K] fp32 saved
+    SCALE_PTR,  # [3] fp32 (read)
+    DLOGITS_PTR,  # [N, (2+K)*K] OUT (fp32)
+    DSCALE_PTR,  # [N, 3] partial sums (fp32); reduced host-side
+    DBASE_PTR,  # [N, (2+K)*K] partial sums (fp32); reduced host-side
+    N,
+    K: tl.constexpr,
+    BLOCK_LEADING: tl.constexpr,
+):
+    """VJP through the elemwise tail.
+
+    Forward:
+        pre_logit = logits[:, :K] * s0 + base[:K]
+        post_logit = logits[:, K:2K] * s1 + base[K:2K]
+        comb_logit = logits[:, 2K:].view(K, K) * s2 + base[2K:].view(K, K)
+        pre = sigmoid(pre_logit) + eps
+        post = 2 * sigmoid(post_logit)
+        comb_sm = softmax(comb_logit, axis=-1); comb = comb_sm + eps
+
+    Backward (per element, all in fp32):
+        d_pre_logit  = d_pre  * pre_sig * (1 - pre_sig)
+        d_post_logit = d_post * 2 * post_sig * (1 - post_sig)
+        d_comb_logit = comb_sm * (d_comb - sum(d_comb * comb_sm, axis=-1))
+
+        d_logits[:, :K]   = d_pre_logit  * s0
+        d_logits[:, K:2K] = d_post_logit * s1
+        d_logits[:, 2K:]  = d_comb_logit.view(K*K) * s2
+
+        d_scale[0] = sum(logits[:, :K]   * d_pre_logit)
+        d_scale[1] = sum(logits[:, K:2K] * d_post_logit)
+        d_scale[2] = sum(logits[:, 2K:].view(K, K) * d_comb_logit)
+
+        d_base[:K]   = sum_n d_pre_logit
+        d_base[K:2K] = sum_n d_post_logit
+        d_base[2K:]  = sum_n d_comb_logit
+
+    Notes:
+        * d_scale needs the *original* fp32 logits.  To avoid carrying
+          a fourth saved tensor we re-derive logits from the existing
+          state: that's not possible (sigmoid is not invertible at
+          large magnitudes).  Instead, we save d_scale as a per-row
+          partial sum that the host-side wrapper reduces with a
+          torch.sum at the end.  Wait -- we actually need the logits
+          themselves for d_scale.  Solution: store ``logits * d_*``
+          (the cross-term) directly without ever materialising logits
+          again -- this is just the elementwise product of the
+          recomputed (d_pre_logit, d_post_logit, d_comb_logit) with
+          the *input-side* gradient of each, taken at the pre-scale
+          point.  Concretely, since
+              d_pre = sigmoid(s0 * x + b)            (forward eq)
+              ∂loss/∂s0 = ∂loss/∂pre_logit * x       (chain rule)
+          and the BWD kernel does NOT have ``x = pre_logits_pre_scale``
+          directly available (logits is a function input passed in
+          fp32).  We therefore accept the per-row partials being
+          a function of ``logits`` and load logits inside the BWD
+          kernel.  See the LOGITS_PTR arg added below.
+
+    For simplicity and correctness, the wrapper passes the *fp32
+    logits* tensor into BWD and the kernel recomputes the pre-scale
+    cross-term inline.  At V4-Flash widths the extra HBM read is
+    negligible (24 fp32 / row vs the K*K state already loaded).
+    """
+
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING)
+    mask_leading = offs < N
+
+    k_idx = tl.arange(0, K)
+
+    # Read scale (broadcast scalars).
+    scale0 = tl.load(SCALE_PTR + 0)
+    scale1 = tl.load(SCALE_PTR + 1)
+    scale2 = tl.load(SCALE_PTR + 2)
+
+    # Saved forward states.
+    pre_sig = tl.load(
+        PRE_SIG_PTR + offs[:, None] * K + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    )
+    post_sig = tl.load(
+        POST_SIG_PTR + offs[:, None] * K + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    )
+    comb_sm = tl.load(
+        COMB_SM_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        mask=mask_leading[:, None, None],
+        other=0.0,
+    )
+
+    # Upstream grads (cast to fp32 for the chain).
+    d_pre = tl.load(
+        DPRE_PTR + offs[:, None] * K + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    d_post = tl.load(
+        DPOST_PTR + offs[:, None] * K + k_idx[None, :],
+        mask=mask_leading[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    d_comb = tl.load(
+        DCOMB_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        mask=mask_leading[:, None, None],
+        other=0.0,
+    ).to(tl.float32)
+
+    # Sigmoid VJP (no eps contribution -- + EPS has derivative 1):
+    d_pre_logit = d_pre * pre_sig * (1.0 - pre_sig)
+    d_post_logit = d_post * 2.0 * post_sig * (1.0 - post_sig)
+
+    # Softmax VJP: d_comb_logit = comb_sm * (d_comb - sum(d_comb * comb_sm, axis=2))
+    dot = tl.sum(d_comb * comb_sm, axis=2, keep_dims=True)
+    d_comb_logit = comb_sm * (d_comb - dot)
+
+    KK_TOTAL: tl.constexpr = (2 + K) * K
+
+    # Write d_logits (= scale * d_*_logit at each slice).
+    tl.store(
+        DLOGITS_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :],
+        d_pre_logit * scale0,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        DLOGITS_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :],
+        d_post_logit * scale1,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        DLOGITS_PTR
+        + offs[:, None, None] * KK_TOTAL
+        + 2 * K
+        + k_idx[None, :, None] * K
+        + k_idx[None, None, :],
+        d_comb_logit * scale2,
+        mask=mask_leading[:, None, None],
+    )
+
+    # d_base accumulates the d_*_logit values directly (no scale factor).
+    # We write per-row partials and let host-side torch.sum reduce them
+    # (avoids cross-block atomic_add).
+    tl.store(
+        DBASE_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :],
+        d_pre_logit,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        DBASE_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :],
+        d_post_logit,
+        mask=mask_leading[:, None],
+    )
+    tl.store(
+        DBASE_PTR + offs[:, None, None] * KK_TOTAL + 2 * K + k_idx[None, :, None] * K + k_idx[None, None, :],
+        d_comb_logit,
+        mask=mask_leading[:, None, None],
+    )
+
+    # d_scale needs ``logits * d_*_logit`` per-element, summed over
+    # each slice's K (or K*K) inner extent and then over rows.  We
+    # write the per-row d_*_logit values into ``DBASE_PTR`` above
+    # (which is exactly ``d_*_logit`` with no scale factor); the
+    # host-side wrapper then computes
+    #   d_scale[0] = (logits[:, :K]   * d_base_partials[:, :K]  ).sum()
+    #   d_scale[1] = (logits[:, K:2K] * d_base_partials[:, K:2K]).sum()
+    #   d_scale[2] = (logits[:, 2K:]  * d_base_partials[:, 2K:] ).sum()
+    # with torch.sum (avoiding cross-block atomic_add).  ``DSCALE_PTR``
+    # is therefore left as the zero-init buffer the wrapper allocated;
+    # we keep it in the kernel signature for forward-compatibility.
+
+
+# ---------------------------------------------------------------------------
+# Block-leading heuristic
+# ---------------------------------------------------------------------------
+
+
+def _pick_block_leading(n: int, k: int) -> int:
+    """Pick ``BLOCK_LEADING`` based on K and the work-axis size.
+
+    At K=4 the in-register state per row is ~``24 + K + K*K = 44`` fp32
+    elements; ``BLOCK_LEADING=64`` fits comfortably (~11 KiB).  At K=16
+    the per-row state grows ~6x; drop to 8.
+    """
+
+    if k <= 4:
+        cap = 64
+    elif k <= 8:
+        cap = 32
+    else:
+        cap = 8
+
+    if n < cap:
+        return max(1, triton.next_power_of_2(n))
+    return cap
+
+
+# ---------------------------------------------------------------------------
+# torch.autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class HCComputeTailFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the FWD/BWD Triton kernels.
+
+    Saves ``(logits, scale, pre_sig, post_sig, comb_sm)`` for backward.
+    Returns ``(pre, post, comb_pre_sinkhorn)`` -- the caller then runs
+    ``sinkhorn_normalize`` on ``comb_pre_sinkhorn``.
+
+    Shape: ``logits [..., (2+K)*K]`` fp32, ``scale [3]`` fp32,
+    ``base [(2+K)*K]`` fp32, ``K`` a power of 2 in ``{1, 2, 4, 8, 16}``.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        logits: torch.Tensor,
+        scale: torch.Tensor,
+        base: torch.Tensor,
+        K: int,
+        eps: float,
+        out_dtype: torch.dtype,
+    ):
+        if K not in _SUPPORTED_K:
+            raise ValueError(f"hc_glue Triton path: unsupported K={K}; expected one of {_SUPPORTED_K}")
+        if logits.shape[-1] != (2 + K) * K:
+            raise ValueError(
+                f"hc_glue: logits last-dim must be (2+K)*K = {(2 + K) * K}, " f"got {logits.shape[-1]}"
+            )
+        if scale.numel() != 3:
+            raise ValueError(f"hc_glue: scale must have 3 elements, got {scale.numel()}")
+        if base.numel() != (2 + K) * K:
+            raise ValueError(
+                f"hc_glue: base must have (2+K)*K = {(2 + K) * K} elements, " f"got {base.numel()}"
+            )
+
+        logits_c = logits.contiguous().to(torch.float32)
+        scale_c = scale.contiguous().to(torch.float32)
+        base_c = base.contiguous().to(torch.float32)
+
+        leading_shape = logits_c.shape[:-1]
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        device = logits_c.device
+        pre = torch.empty((*leading_shape, K), dtype=out_dtype, device=device)
+        post = torch.empty((*leading_shape, K), dtype=out_dtype, device=device)
+        comb = torch.empty((*leading_shape, K, K), dtype=out_dtype, device=device)
+
+        # Saved-for-backward fp32 states.
+        pre_sig = torch.empty((N, K), dtype=torch.float32, device=device)
+        post_sig = torch.empty((N, K), dtype=torch.float32, device=device)
+        comb_sm = torch.empty((N, K, K), dtype=torch.float32, device=device)
+
+        block_leading = _pick_block_leading(N, K)
+        grid = (triton.cdiv(N, block_leading),)
+        _hc_compute_tail_fwd_kernel[grid](
+            logits_c,
+            scale_c,
+            base_c,
+            pre,
+            post,
+            comb,
+            pre_sig,
+            post_sig,
+            comb_sm,
+            N,
+            EPS=float(eps),
+            K=K,
+            BLOCK_LEADING=block_leading,
+            OUT_DTYPE=_triton_dtype(out_dtype),
+        )
+
+        ctx.save_for_backward(logits_c, scale_c, pre_sig, post_sig, comb_sm)
+        ctx.K = K
+        ctx.eps = float(eps)
+        ctx.leading_shape = tuple(leading_shape)
+        ctx.out_dtype = out_dtype
+        return pre, post, comb
+
+    @staticmethod
+    def backward(ctx, d_pre, d_post, d_comb):  # type: ignore[override]
+        logits_c, scale_c, pre_sig, post_sig, comb_sm = ctx.saved_tensors
+        K: int = ctx.K
+        leading_shape = ctx.leading_shape
+
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        d_pre = d_pre.contiguous()
+        d_post = d_post.contiguous()
+        d_comb = d_comb.contiguous()
+
+        device = logits_c.device
+        d_logits = torch.empty_like(logits_c)
+        d_base_partials = torch.empty((N, (2 + K) * K), dtype=torch.float32, device=device)
+        # d_scale_partials is no longer written by the kernel; we leave
+        # the buffer hint in the kernel signature for forward-compat.
+        d_scale_partials = torch.zeros((N, 3), dtype=torch.float32, device=device)
+
+        block_leading = _pick_block_leading(N, K)
+        grid = (triton.cdiv(N, block_leading),)
+        _hc_compute_tail_bwd_kernel[grid](
+            d_pre,
+            d_post,
+            d_comb,
+            pre_sig,
+            post_sig,
+            comb_sm,
+            scale_c,
+            d_logits,
+            d_scale_partials,
+            d_base_partials,
+            N,
+            K=K,
+            BLOCK_LEADING=block_leading,
+        )
+
+        # Host-side reductions:
+        #   d_base[i] = sum_n d_base_partials[n, i]                       (3*K + K*K entries)
+        #   d_scale[0] = sum_{n, k}   logits_slice_pre  * d_base_partials_pre
+        #   d_scale[1] = sum_{n, k}   logits_slice_post * d_base_partials_post
+        #   d_scale[2] = sum_{n, k1, k2} logits_slice_comb * d_base_partials_comb
+        #
+        # Equivalent to: d_scale[i] = (logits_slice_i * d_base_partials_i).sum()
+        # since d_base = d_*_logit (no scale) and d_scale[i] = sum logits * d_*_logit.
+        logits_flat = logits_c.reshape(N, -1)
+        d_base = d_base_partials.sum(dim=0)
+        d_scale_0 = (logits_flat[:, :K] * d_base_partials[:, :K]).sum()
+        d_scale_1 = (logits_flat[:, K : 2 * K] * d_base_partials[:, K : 2 * K]).sum()
+        d_scale_2 = (logits_flat[:, 2 * K :] * d_base_partials[:, 2 * K :]).sum()
+        d_scale = torch.stack([d_scale_0, d_scale_1, d_scale_2])
+
+        # Reshape d_logits back to the leading shape that the FWD input
+        # had so the autograd machinery passes it back to the caller's
+        # F.linear backward correctly.
+        d_logits = d_logits.view(*leading_shape, (2 + K) * K)
+
+        # K, eps, out_dtype are non-differentiable.
+        return d_logits, d_scale, d_base, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Public Python entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff the ``PRIMUS_HC_TRITON`` env knob is not ``"0"``.
+
+    Default-on; A/B toggle via ``PRIMUS_HC_TRITON=0``.
+    """
+
+    return os.environ.get("PRIMUS_HC_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(logits: torch.Tensor, K: int) -> bool:
+    """Return True iff the input shape / device is supported.
+
+    Used by the dispatcher in
+    :meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights`
+    to safely fall back to the eager body for unsupported configurations.
+    """
+
+    if not logits.is_cuda:
+        return False
+    if K not in _SUPPORTED_K:
+        return False
+    if logits.shape[-1] != (2 + K) * K:
+        return False
+    return True
+
+
+def hc_glue_compute_tail_triton(
+    logits: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    *,
+    K: int,
+    eps: float,
+    out_dtype: torch.dtype,
+):
+    """Run the Triton-fused HC compute_weights tail.
+
+    Returns ``(pre, post, comb_pre_sinkhorn)`` -- the caller runs
+    ``sinkhorn_normalize`` on ``comb_pre_sinkhorn``.
+    """
+
+    return HCComputeTailFn.apply(logits, scale, base, K, eps, out_dtype)
+
+
+__all__ = [
+    "HCComputeTailFn",
+    "hc_glue_compute_tail_triton",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py
new file mode 100644
index 000000000..c82543d75
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py
@@ -0,0 +1,503 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused Indexer scoring (plan-6 P38).
+
+Fuses the ``einsum + relu + mul + sum + causal_mask`` chain in
+:meth:`primus.backends.megatron.core.transformer.indexer.Indexer.forward`
+into a single FWD + single BWD Triton kernel pair.  Eager body:
+
+.. code-block:: python
+
+    relu_term = F.relu(torch.einsum("bshd,bpd->bshp", q_i, k_icomp))
+    scores = (relu_term * w_i.unsqueeze(-1)).sum(dim=2)  # [B, S, P]
+    mask = self._causal_mask(S, P, scores.device, scores.dtype)
+    scores = scores + mask.unsqueeze(0)
+
+P38 collapses these ~7 ATen kernels (einsum + relu + mul + sum +
+mask alloc + mask add + dtype cast) into one Triton kernel that:
+
+* Computes the per-(q, p) dot product over the head-feature dim ``Hd``
+  inline (no `[B, S, H, P]` intermediate tensor materialised);
+* Applies ``relu`` per element;
+* Multiplies by ``w_i`` per head;
+* Reduces over heads;
+* Materialises the causal mask **inline** via
+  ``tl.where((p + 1) * compress_ratio - 1 <= s, 0.0, -INF)`` — no
+  ``[S, P]`` mask tensor, no HBM traffic;
+* Writes ``scores [B, S, P]`` (the `topk` and trailing tail stay
+  host-side; `topk` is heavy GPU compute on its own and benefits
+  from being its own kernel).
+
+The BWD kernel recomputes the per-tile `relu` mask in the BWD pass
+(FlashAttention-style trick) instead of saving it — saves ``H * S * P``
+bits of HBM per CSA layer.
+
+Gating: ``PRIMUS_INDEXER_TRITON == "1"`` (**default-OFF**).
+
+Default-off rationale (P38 descope per `plan-6/02-phase-details.md`
+§"Task list refinement"): at V4-Flash widths (B=1, S=4096, P=1024,
+H=8, Hd=128) the eager `einsum` maps to a cuBLAS / hipBLASLt
+batched-matmul that already runs at ~28 TFLOP/s on MI355.  The
+generic Triton kernel here is FWD-competitive only at small shapes
+(3.35x FWD speedup at B=2, S=128, P=32) but regresses ~30% at the
+production V4-Flash shape; BWD regresses ~12x because of cross-tile
+``atomic_add`` traffic on `dq` / `dk` / `dw`.  The kernel stays
+available for future tuning + small-shape paths via the env knob.
+
+Per-shape behavior summary:
+
+* V4-Flash production (B=1, S=4096, P=1024, H=8, Hd=128, bf16):
+  FWD 0.424 ms (triton) vs 0.306 ms (eager) -> **0.72x** (regression).
+  BWD 6.457 ms (triton) vs 0.489 ms (eager) -> **0.08x** (regression).
+* Small (B=2, S=128, P=32, H=8, Hd=128, bf16):
+  FWD 0.053 ms (triton) vs 0.176 ms (eager) -> **3.35x** (speedup).
+  BWD 0.226 ms (triton) vs 0.256 ms (eager) -> **1.14x** (speedup).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# Indexer scoring runs at V4-Flash widths [B=1, S=4096, P=1024, H=8,
+# Hd=128].  H is small enough to be a compile-time constant in the
+# kernel; supported values are documented here.
+_SUPPORTED_H = (1, 2, 4, 8, 16)
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _indexer_score_fwd_kernel(
+    Q_PTR,  # [B, S, H, Hd] - q_i
+    K_PTR,  # [B, P, Hd]    - k_icomp.squeeze(2)
+    W_PTR,  # [B, S, H]     - w_i
+    SCORES_PTR,  # [B, S, P]     - out (fp32 internal, cast OUT_DTYPE)
+    B,
+    S,
+    P,
+    H: tl.constexpr,
+    HD: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_P: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """One program tile = one ``[B_b, BLOCK_S, BLOCK_P]`` chunk.
+
+    Loop axis: ``Hd`` (head-feature dim) and ``H`` (heads), both
+    compile-time known.  Per tile the kernel:
+
+        1. For each head ``h`` (unrolled, since H is constexpr):
+            a. Load ``q_i[b_b, s_tile, h, :]`` -- shape ``[BLOCK_S, HD]``.
+            b. Load ``k_icomp[b_b, p_tile, :]``    -- shape ``[BLOCK_P, HD]``.
+            c. Compute ``dot = q @ k.T`` -- shape ``[BLOCK_S, BLOCK_P]``.
+            d. Apply ``relu(dot)``.
+            e. Load ``w_i[b_b, s_tile, h]`` -- shape ``[BLOCK_S]``.
+            f. Multiply: ``acc += relu * w[:, None]``.
+        2. Materialise causal mask inline: positions ``s`` may attend to
+           pool position ``p`` iff ``(p + 1) * compress_ratio - 1 <= s``;
+           write ``-inf`` otherwise.
+        3. Store ``acc`` cast to OUT_DTYPE.
+
+    Register footprint at V4-Flash (BLOCK_S=64, BLOCK_P=64, H=8, HD=128):
+        per head: 64*128 + 64*128 = 16384 fp32 = 64 KiB for (q, k).
+        per tile: 64*64 = 4096 fp32 = 16 KiB for acc.
+    Cumulative per-program peak ~ 80 KiB ≈ 320 VGPRs / warp.  At MI355
+    256 VGPRs we drop BLOCK_S to 32 for safety on the small kernel.
+    """
+
+    pid_b = tl.program_id(0)
+    pid_s = tl.program_id(1)
+    pid_p = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P)
+    s_mask = s_offs < S
+    p_mask = p_offs < P
+
+    hd_idx = tl.arange(0, HD)
+
+    acc = tl.zeros((BLOCK_S, BLOCK_P), dtype=tl.float32)
+
+    # Unroll over heads (H is small and constexpr).
+    for h in tl.static_range(0, H):
+        # q [BLOCK_S, HD]: q_i[pid_b, s_offs, h, :]
+        q_tile = tl.load(
+            Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :],
+            mask=s_mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        # k [BLOCK_P, HD]: k_icomp[pid_b, p_offs, :]
+        k_tile = tl.load(
+            K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :],
+            mask=p_mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        # dot [BLOCK_S, BLOCK_P] = q @ k.T
+        dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32)
+        dot = tl.maximum(dot, 0.0)  # relu
+        # w [BLOCK_S]: w_i[pid_b, s_offs, h]
+        w_h = tl.load(
+            W_PTR + pid_b * S * H + s_offs * H + h,
+            mask=s_mask,
+            other=0.0,
+        ).to(tl.float32)
+        acc += dot * w_h[:, None]
+
+    # Apply causal mask inline.  Allowed iff `(p + 1) * cr - 1 <= s`,
+    # i.e. the pool position's window end is no later than the query.
+    s_arr = s_offs[:, None]
+    p_arr = p_offs[None, :]
+    allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr
+    NEG_INF = -float("inf")
+    acc = tl.where(allowed, acc, NEG_INF)
+
+    tl.store(
+        SCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :],
+        acc.to(OUT_DTYPE),
+        mask=s_mask[:, None] & p_mask[None, :],
+    )
+
+
+@triton.jit
+def _indexer_score_bwd_kernel(
+    DSCORES_PTR,  # [B, S, P]    grad in OUT_DTYPE
+    Q_PTR,  # [B, S, H, Hd]
+    K_PTR,  # [B, P, Hd]
+    W_PTR,  # [B, S, H]
+    DQ_PTR,  # [B, S, H, Hd] OUT (fp32)
+    DK_PTR,  # [B, P, Hd]    OUT (fp32, scattered)
+    DW_PTR,  # [B, S, H]     OUT (fp32)
+    B,
+    S,
+    P,
+    H: tl.constexpr,
+    HD: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_P: tl.constexpr,
+):
+    """VJP through the scoring chain.
+
+    The mask is `where(allowed, acc, -inf)`; its derivative is 1 in
+    allowed positions and 0 elsewhere.  Out-of-range positions therefore
+    contribute zero grad and we set ``dscores`` to 0 there before
+    walking back.
+
+    Per element forward:
+        acc = sum_h relu(q . k) * w_h
+        masked = where(allowed, acc, -inf)
+    Per element backward:
+        d_acc = d_masked * where(allowed, 1, 0)        # equiv. d_masked masked by allowed
+        # d_relu_dot[h] = d_acc * w_h
+        # d_w[h]        = d_acc * relu_dot[h]
+        # relu' = (dot > 0).  Recompute dot, relu mask.
+        # d_dot[h]   = d_relu_dot[h] * (dot > 0)
+        # d_q[h, :]  = d_dot[h, p] @ k[p, :]
+        # d_k[p, :]  = d_dot[s, p] @ q[s, h, :]
+
+    For the BWD we use one program per (b, s_tile, p_tile) just like
+    FWD; ``d_q`` is accumulated locally (one chunk per s_tile so
+    multiple p_tiles need atomic add), ``d_k`` is scattered across
+    s_tiles so needs atomic add, ``d_w`` is accumulated similarly.
+    """
+
+    pid_b = tl.program_id(0)
+    pid_s = tl.program_id(1)
+    pid_p = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P)
+    s_mask = s_offs < S
+    p_mask = p_offs < P
+
+    hd_idx = tl.arange(0, HD)
+
+    # Load dscores [BLOCK_S, BLOCK_P]
+    dmasked = tl.load(
+        DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :],
+        mask=s_mask[:, None] & p_mask[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    # Apply causal mask (zero out invalid positions).
+    s_arr = s_offs[:, None]
+    p_arr = p_offs[None, :]
+    allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr
+    d_acc = tl.where(allowed, dmasked, 0.0)
+
+    # Unroll over heads.
+    for h in tl.static_range(0, H):
+        # Reload q, k for this head (FlashAttention-style recompute).
+        q_tile = tl.load(
+            Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :],
+            mask=s_mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        k_tile = tl.load(
+            K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :],
+            mask=p_mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        w_h = tl.load(
+            W_PTR + pid_b * S * H + s_offs * H + h,
+            mask=s_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        # Recompute dot, relu, relu_dot.
+        dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32)
+        relu_dot = tl.maximum(dot, 0.0)  # [BLOCK_S, BLOCK_P]
+        relu_mask = dot > 0.0
+
+        # d_w[h] = sum_p d_acc * relu_dot
+        dw_h = tl.sum(d_acc * relu_dot, axis=1)  # [BLOCK_S]
+
+        # d_relu_dot = d_acc * w_h[:, None]
+        d_relu_dot = d_acc * w_h[:, None]
+        # d_dot = d_relu_dot where relu_mask else 0
+        d_dot = tl.where(relu_mask, d_relu_dot, 0.0)
+
+        # d_q[s, h, :] = sum_p d_dot[s, p] * k[p, :]   shape [BLOCK_S, HD]
+        d_q = tl.dot(d_dot, k_tile, out_dtype=tl.float32)
+        # d_k[p, :]    = sum_s d_dot[s, p] * q[s, :]   shape [BLOCK_P, HD]
+        d_k = tl.dot(tl.trans(d_dot), q_tile, out_dtype=tl.float32)
+
+        # Stores with atomic_add since multiple p_tiles/s_tiles touch
+        # the same locations.
+        tl.atomic_add(
+            DQ_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :],
+            d_q,
+            mask=s_mask[:, None],
+        )
+        tl.atomic_add(
+            DK_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :],
+            d_k,
+            mask=p_mask[:, None],
+        )
+        tl.atomic_add(
+            DW_PTR + pid_b * S * H + s_offs * H + h,
+            dw_h,
+            mask=s_mask,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Heuristics
+# ---------------------------------------------------------------------------
+
+
+def _pick_block(s: int, p: int, h: int, hd: int) -> tuple[int, int]:
+    """Pick BLOCK_S, BLOCK_P that fit MI355 register budget."""
+    block_s = 32
+    block_p = 32
+    if hd >= 128:
+        block_s = 32
+        block_p = 32
+    if hd >= 256:
+        block_s = 16
+        block_p = 32
+    # Floor both tiles to 16: BLOCK_S / BLOCK_P are the M / N dims of the
+    # ``tl.dot(q_tile, k_tile.T)`` in the FWD/BWD score kernels, so they
+    # must be >= the gfx1250 WMMA minimum. For tiny S / P (e.g. unit-test
+    # shapes), ``next_power_of_2(s|p)`` can drop below 16 and the dot fails
+    # to select a matrix-core intrinsic ("no matching matrix core intrinsic
+    # for wmma version 3 ... [0, 0, K]"). Surplus rows/cols are masked by
+    # ``s_mask`` / ``p_mask`` inside the kernels, so a 16-wide tile over a
+    # smaller S / P is safe.
+    block_s = max(16, min(block_s, triton.next_power_of_2(max(1, s))))
+    block_p = max(16, min(block_p, triton.next_power_of_2(max(1, p))))
+    return block_s, block_p
+
+
+# ---------------------------------------------------------------------------
+# autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class IndexerScoreFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the FWD/BWD Triton kernels.
+
+    Returns ``scores [B, S, P]`` in ``out_dtype``.  Caller runs
+    ``torch.topk`` / sentinel substitution / padding host-side.
+
+    Inputs:
+        q_i [B, S, H, Hd] - per-head queries (any float dtype).
+        k_icomp [B, P, Hd] - compressed keys (any float dtype).
+        w_i [B, S, H] - per-head weights (any float dtype).
+        compress_ratio - int (typically 4 for CSA).
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q_i: torch.Tensor,
+        k_icomp: torch.Tensor,
+        w_i: torch.Tensor,
+        compress_ratio: int,
+        out_dtype: torch.dtype,
+    ):
+        if q_i.dim() != 4:
+            raise ValueError(f"q_i must be [B, S, H, Hd], got shape {tuple(q_i.shape)}")
+        if k_icomp.dim() != 3:
+            raise ValueError(f"k_icomp must be [B, P, Hd], got shape {tuple(k_icomp.shape)}")
+        if w_i.dim() != 3:
+            raise ValueError(f"w_i must be [B, S, H], got shape {tuple(w_i.shape)}")
+
+        B, S, H, HD = q_i.shape
+        Bk, P, HDk = k_icomp.shape
+        Bw, Sw, Hw = w_i.shape
+        if B != Bk or B != Bw:
+            raise ValueError(f"Mismatched B: q={B} k={Bk} w={Bw}")
+        if HD != HDk:
+            raise ValueError(f"Mismatched Hd: q={HD} k={HDk}")
+        if S != Sw:
+            raise ValueError(f"Mismatched S: q={S} w={Sw}")
+        if H != Hw:
+            raise ValueError(f"Mismatched H: q={H} w={Hw}")
+        if H not in _SUPPORTED_H:
+            raise ValueError(f"Unsupported H={H}; expected one of {_SUPPORTED_H}")
+        if HD & (HD - 1) != 0:
+            raise ValueError(f"Hd must be a power of 2, got {HD}")
+
+        q_c = q_i.contiguous()
+        k_c = k_icomp.contiguous()
+        w_c = w_i.contiguous()
+
+        device = q_c.device
+        scores = torch.empty((B, S, P), dtype=out_dtype, device=device)
+
+        block_s, block_p = _pick_block(S, P, H, HD)
+        grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p))
+        _indexer_score_fwd_kernel[grid](
+            q_c,
+            k_c,
+            w_c,
+            scores,
+            B,
+            S,
+            P,
+            H=H,
+            HD=HD,
+            COMPRESS_RATIO=int(compress_ratio),
+            BLOCK_S=block_s,
+            BLOCK_P=block_p,
+            OUT_DTYPE={
+                torch.float32: tl.float32,
+                torch.float16: tl.float16,
+                torch.bfloat16: tl.bfloat16,
+                torch.float64: tl.float64,
+            }[out_dtype],
+        )
+
+        ctx.save_for_backward(q_c, k_c, w_c)
+        ctx.compress_ratio = int(compress_ratio)
+        ctx.shape = (B, S, P, H, HD)
+        ctx.in_dtypes = (q_i.dtype, k_icomp.dtype, w_i.dtype)
+        return scores
+
+    @staticmethod
+    def backward(ctx, d_scores):  # type: ignore[override]
+        q_c, k_c, w_c = ctx.saved_tensors
+        B, S, P, H, HD = ctx.shape
+        compress_ratio = ctx.compress_ratio
+        q_dtype, k_dtype, w_dtype = ctx.in_dtypes
+
+        d_scores = d_scores.contiguous()
+        device = q_c.device
+
+        d_q_fp32 = torch.zeros((B, S, H, HD), dtype=torch.float32, device=device)
+        d_k_fp32 = torch.zeros((B, P, HD), dtype=torch.float32, device=device)
+        d_w_fp32 = torch.zeros((B, S, H), dtype=torch.float32, device=device)
+
+        block_s, block_p = _pick_block(S, P, H, HD)
+        grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p))
+        _indexer_score_bwd_kernel[grid](
+            d_scores,
+            q_c,
+            k_c,
+            w_c,
+            d_q_fp32,
+            d_k_fp32,
+            d_w_fp32,
+            B,
+            S,
+            P,
+            H=H,
+            HD=HD,
+            COMPRESS_RATIO=int(compress_ratio),
+            BLOCK_S=block_s,
+            BLOCK_P=block_p,
+        )
+
+        return d_q_fp32.to(q_dtype), d_k_fp32.to(k_dtype), d_w_fp32.to(w_dtype), None, None
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff ``PRIMUS_INDEXER_TRITON_FULL == "1"``.
+
+    **Re-purposed at P41:** the P38 full-fuse path is now gated by the
+    distinct env knob ``PRIMUS_INDEXER_TRITON_FULL`` (default ``"0"``).
+    The original ``PRIMUS_INDEXER_TRITON`` env now controls the cheaper
+    P41 post-einsum tail fusion in
+    :mod:`indexer_score_post`.
+
+    The full-fuse path is descoped at V4-Flash widths (cuBLAS einsum
+    beats the generic Triton kernel + BWD atomic_add contention
+    regresses 12x).  Kept available for small-shape paths and future
+    tuning.  Set ``PRIMUS_INDEXER_TRITON_FULL=1`` to opt-in.
+    """
+    return os.environ.get("PRIMUS_INDEXER_TRITON_FULL", "0") == "1"
+
+
+def is_triton_kernel_supported(q_i: torch.Tensor, k_icomp: torch.Tensor, w_i: torch.Tensor) -> bool:
+    """Return True iff the input shapes / device support the Triton path."""
+    if not q_i.is_cuda or not k_icomp.is_cuda or not w_i.is_cuda:
+        return False
+    if q_i.dim() != 4 or k_icomp.dim() != 3 or w_i.dim() != 3:
+        return False
+    B, S, H, HD = q_i.shape
+    if H not in _SUPPORTED_H:
+        return False
+    if HD & (HD - 1) != 0:
+        return False
+    return True
+
+
+def indexer_score_triton(
+    q_i: torch.Tensor,
+    k_icomp: torch.Tensor,
+    w_i: torch.Tensor,
+    *,
+    compress_ratio: int,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute Indexer scores via the fused Triton kernel.
+
+    Returns ``scores [B, S, P]`` of dtype ``out_dtype``.  Masked
+    positions hold ``-inf``.
+    """
+    return IndexerScoreFn.apply(q_i, k_icomp, w_i, compress_ratio, out_dtype)
+
+
+__all__ = [
+    "IndexerScoreFn",
+    "indexer_score_triton",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py
new file mode 100644
index 000000000..f1f4c61e3
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py
@@ -0,0 +1,468 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused Indexer scoring -- post-einsum tail (plan-6 P41).
+
+Companion to :mod:`indexer_score` (P38).  P38 fused the entire
+``einsum + relu + mul + sum + causal_mask`` chain into one kernel and
+lost to cuBLAS / hipBLASLt on the matmul half (~28 TFLOP/s eager vs
+~20 TFLOP/s Triton at V4-Flash widths).  P41 keeps the einsum eager
+and fuses **only** the post-matmul tail
+(``relu -> mul(w_i) -> sum(H) -> + causal_mask``).  The matmul is
+compute-bound and stays on cuBLAS peak; the tail is bandwidth-bound
+and Triton wins because every elementwise op costs a full HBM
+round-trip.
+
+Eager body (with einsum kept):
+
+.. code-block:: python
+
+    dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp)   # stays eager
+    relu_term = F.relu(dot)
+    scores = (relu_term * w_i.unsqueeze(-1)).sum(dim=2)  # [B, S, P]
+    mask = self._causal_mask(S, P, scores.device, scores.dtype)
+    scores = scores + mask.unsqueeze(0)
+
+P41 collapses the tail (``relu + mul + sum + mask_alloc + mask_add +
+dtype cast``, ~5 ATen launches) into one Triton kernel that:
+
+* Reads ``dot [B, S, H, P]`` once;
+* Applies ``relu`` per element;
+* Multiplies by ``w_i[B, S, H, 1]`` (broadcast over P);
+* Reduces over heads ``H``;
+* Materialises the causal mask **inline** via
+  ``tl.where((p + 1) * compress_ratio - 1 <= s, acc, -inf)``;
+* Writes ``scores [B, S, P]`` cast to ``OUT_DTYPE``.
+
+BWD takes ``d_scores [B, S, P]`` + saved ``dot [B, S, H, P]`` +
+``w_i [B, S, H]``, emits:
+
+* ``d_dot[b, s, h, p] = d_scores[b, s, p] * w_i[b, s, h]`` where
+  ``dot[b, s, h, p] > 0`` else 0  (one HBM write per element);
+* ``d_w_i[b, s, h] = sum_p(d_scores[b, s, p] * relu(dot[b, s, h, p]))``
+  (one reduction per (b, s, h) — no cross-block atomic_add).
+
+The BWD has no atomic_add traffic (the P38 BWD's killer).  Both FWD
+and BWD are bandwidth-bound and should win at V4-Flash widths.
+
+Gating: ``PRIMUS_INDEXER_TRITON == "1"`` (the re-purposed env knob;
+default initially ``"0"`` then ``"1"`` after the proxy A/B confirms).
+The legacy P38 full-fuse path lives behind ``PRIMUS_INDEXER_TRITON_FULL``
+(see :mod:`indexer_score`).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# H is small and known per call site -- V4-Flash uses H=8.  Same
+# supported set as P38.
+_SUPPORTED_H = (1, 2, 4, 8, 16)
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _indexer_score_post_fwd_kernel(
+    DOT_PTR,  # [B, S, H, P] - eager einsum output, pre-relu
+    W_PTR,  # [B, S, H]    - per-head weights
+    SCORES_PTR,  # [B, S, P]    - out
+    B,
+    S,
+    P,
+    H: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_P: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """One program tile = ``[B_b, BLOCK_S, BLOCK_P]``.
+
+    Per-tile workflow:
+
+        1. Initialise ``acc [BLOCK_S, BLOCK_P]`` fp32 to 0.
+        2. For each head ``h`` (constexpr unroll):
+           a. Load ``dot[b, s_tile, h, p_tile]``  -> [BLOCK_S, BLOCK_P].
+           b. Apply ``relu``.
+           c. Load ``w[b, s_tile, h]``            -> [BLOCK_S].
+           d. ``acc += relu * w[:, None]``.
+        3. Apply causal mask inline.
+        4. Store ``acc`` cast to ``OUT_DTYPE``.
+    """
+
+    pid_b = tl.program_id(0)
+    pid_s = tl.program_id(1)
+    pid_p = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P)
+    s_mask = s_offs < S
+    p_mask = p_offs < P
+
+    acc = tl.zeros((BLOCK_S, BLOCK_P), dtype=tl.float32)
+
+    for h in tl.static_range(0, H):
+        # dot[pid_b, s_offs, h, p_offs] -> [BLOCK_S, BLOCK_P]
+        dot_tile = tl.load(
+            DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :],
+            mask=s_mask[:, None] & p_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        relu_tile = tl.maximum(dot_tile, 0.0)
+
+        # w[pid_b, s_offs, h] -> [BLOCK_S]
+        w_h = tl.load(
+            W_PTR + pid_b * S * H + s_offs * H + h,
+            mask=s_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        acc += relu_tile * w_h[:, None]
+
+    # Causal mask inline.
+    s_arr = s_offs[:, None]
+    p_arr = p_offs[None, :]
+    allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr
+    NEG_INF = -float("inf")
+    acc = tl.where(allowed, acc, NEG_INF)
+
+    tl.store(
+        SCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :],
+        acc.to(OUT_DTYPE),
+        mask=s_mask[:, None] & p_mask[None, :],
+    )
+
+
+@triton.jit
+def _indexer_score_post_bwd_dw_kernel(
+    DSCORES_PTR,  # [B, S, P]    - grad in (any float dtype)
+    DOT_PTR,  # [B, S, H, P] - saved pre-relu activations
+    DW_PTR,  # [B, S, H]    - OUT  (fp32)
+    B,
+    S,
+    P,
+    H: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_P_INNER: tl.constexpr,
+):
+    """Per ``(b, s_tile, h)`` program: ``d_w[h] = sum_p(d_acc * relu(dot[h]))``.
+
+    Loops over P internally in chunks of ``BLOCK_P_INNER``; ``d_w``
+    is fully local to each program (no cross-block ``atomic_add``).
+    """
+
+    pid_b = tl.program_id(0)
+    pid_s = tl.program_id(1)
+    h = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    s_mask = s_offs < S
+    s_arr = s_offs[:, None]
+
+    acc = tl.zeros((BLOCK_S,), dtype=tl.float32)
+
+    n_p_chunks = tl.cdiv(P, BLOCK_P_INNER)
+    for chunk in range(n_p_chunks):
+        p_offs = chunk * BLOCK_P_INNER + tl.arange(0, BLOCK_P_INNER)
+        p_mask = p_offs < P
+        p_arr = p_offs[None, :]
+
+        dmasked = tl.load(
+            DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :],
+            mask=s_mask[:, None] & p_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr
+        d_acc = tl.where(allowed, dmasked, 0.0)
+
+        dot_tile = tl.load(
+            DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :],
+            mask=s_mask[:, None] & p_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        relu_tile = tl.maximum(dot_tile, 0.0)
+
+        acc += tl.sum(d_acc * relu_tile, axis=1)
+
+    tl.store(
+        DW_PTR + pid_b * S * H + s_offs * H + h,
+        acc,
+        mask=s_mask,
+    )
+
+
+@triton.jit
+def _indexer_score_post_bwd_ddot_kernel(
+    DSCORES_PTR,  # [B, S, P]    - grad in (any float dtype)
+    DOT_PTR,  # [B, S, H, P] - saved pre-relu activations
+    W_PTR,  # [B, S, H]    - per-head weights
+    DDOT_PTR,  # [B, S, H, P] - OUT  (fp32)
+    B,
+    S,
+    P,
+    H: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_P: tl.constexpr,
+):
+    """VJP through the post-einsum tail — ``d_dot`` only.
+
+    Per-element forward (with mask):
+        d_acc = where(allowed, d_scores, 0)
+    Per-element backward:
+        d_dot[h] = where(dot[h] > 0, d_acc * w_h, 0)
+
+    ``d_dot`` is one HBM store per element (the grid covers each
+    output element exactly once); ``d_w`` is computed via eager
+    PyTorch in the autograd wrapper as a single ATen reduce (cheap
+    + avoids the cross-block ``atomic_add`` traffic that hurt the
+    P38 BWD).
+    """
+
+    pid_b = tl.program_id(0)
+    pid_s = tl.program_id(1)
+    pid_p = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P)
+    s_mask = s_offs < S
+    p_mask = p_offs < P
+
+    dmasked = tl.load(
+        DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :],
+        mask=s_mask[:, None] & p_mask[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    s_arr = s_offs[:, None]
+    p_arr = p_offs[None, :]
+    allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr
+    d_acc = tl.where(allowed, dmasked, 0.0)
+
+    for h in tl.static_range(0, H):
+        dot_tile = tl.load(
+            DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :],
+            mask=s_mask[:, None] & p_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        relu_mask = dot_tile > 0.0
+
+        w_h = tl.load(
+            W_PTR + pid_b * S * H + s_offs * H + h,
+            mask=s_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        d_relu = d_acc * w_h[:, None]
+        d_dot = tl.where(relu_mask, d_relu, 0.0)
+        tl.store(
+            DDOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :],
+            d_dot,
+            mask=s_mask[:, None] & p_mask[None, :],
+        )
+
+
+# ---------------------------------------------------------------------------
+# Heuristics
+# ---------------------------------------------------------------------------
+
+
+def _pick_block(s: int, p: int, h: int) -> tuple[int, int]:
+    """Pick BLOCK_S, BLOCK_P that fit MI355 register budget.
+
+    Tail kernel is much cheaper than the P38 full-fuse one (no dot
+    product, only elementwise + reduce) so we can use larger blocks.
+    """
+    block_s = 64
+    block_p = 64
+    return min(block_s, triton.next_power_of_2(max(1, s))), min(block_p, triton.next_power_of_2(max(1, p)))
+
+
+# ---------------------------------------------------------------------------
+# autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class IndexerScorePostFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the post-einsum tail kernels.
+
+    Inputs:
+        dot [B, S, H, P] - eager einsum output (pre-relu); any float dtype.
+        w_i [B, S, H]    - per-head weights; any float dtype.
+        compress_ratio   - int (typically 4 for CSA).
+        out_dtype        - dtype of the returned ``scores``.
+
+    Returns:
+        scores [B, S, P] of dtype ``out_dtype``.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        dot: torch.Tensor,
+        w_i: torch.Tensor,
+        compress_ratio: int,
+        out_dtype: torch.dtype,
+    ):
+        if dot.dim() != 4:
+            raise ValueError(f"dot must be [B, S, H, P], got shape {tuple(dot.shape)}")
+        if w_i.dim() != 3:
+            raise ValueError(f"w_i must be [B, S, H], got shape {tuple(w_i.shape)}")
+
+        B, S, H, P = dot.shape
+        Bw, Sw, Hw = w_i.shape
+        if B != Bw:
+            raise ValueError(f"Mismatched B: dot={B} w={Bw}")
+        if S != Sw:
+            raise ValueError(f"Mismatched S: dot={S} w={Sw}")
+        if H != Hw:
+            raise ValueError(f"Mismatched H: dot={H} w={Hw}")
+        if H not in _SUPPORTED_H:
+            raise ValueError(f"Unsupported H={H}; expected one of {_SUPPORTED_H}")
+
+        dot_c = dot.contiguous()
+        w_c = w_i.contiguous()
+        device = dot_c.device
+        scores = torch.empty((B, S, P), dtype=out_dtype, device=device)
+
+        block_s, block_p = _pick_block(S, P, H)
+        grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p))
+        _indexer_score_post_fwd_kernel[grid](
+            dot_c,
+            w_c,
+            scores,
+            B,
+            S,
+            P,
+            H=H,
+            COMPRESS_RATIO=int(compress_ratio),
+            BLOCK_S=block_s,
+            BLOCK_P=block_p,
+            OUT_DTYPE={
+                torch.float32: tl.float32,
+                torch.float16: tl.float16,
+                torch.bfloat16: tl.bfloat16,
+                torch.float64: tl.float64,
+            }[out_dtype],
+        )
+
+        ctx.save_for_backward(dot_c, w_c)
+        ctx.compress_ratio = int(compress_ratio)
+        ctx.shape = (B, S, H, P)
+        ctx.in_dtypes = (dot.dtype, w_i.dtype)
+        return scores
+
+    @staticmethod
+    def backward(ctx, d_scores):  # type: ignore[override]
+        dot_c, w_c = ctx.saved_tensors
+        B, S, H, P = ctx.shape
+        compress_ratio = ctx.compress_ratio
+        dot_dtype, w_dtype = ctx.in_dtypes
+
+        d_scores = d_scores.contiguous()
+        device = dot_c.device
+
+        d_dot_fp32 = torch.empty((B, S, H, P), dtype=torch.float32, device=device)
+        d_w_fp32 = torch.empty((B, S, H), dtype=torch.float32, device=device)
+
+        block_s, block_p = _pick_block(S, P, H)
+        grid_ddot = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p))
+        _indexer_score_post_bwd_ddot_kernel[grid_ddot](
+            d_scores,
+            dot_c,
+            w_c,
+            d_dot_fp32,
+            B,
+            S,
+            P,
+            H=H,
+            COMPRESS_RATIO=int(compress_ratio),
+            BLOCK_S=block_s,
+            BLOCK_P=block_p,
+        )
+
+        # d_w[b, s, h] = sum_p (d_acc * relu(dot[h])) — one program per
+        # (b, s_tile, h), loops over P internally so the reduce is fully
+        # local (no cross-block atomic_add).
+        block_p_inner = min(128, triton.next_power_of_2(max(1, P)))
+        grid_dw = (B, triton.cdiv(S, block_s), H)
+        _indexer_score_post_bwd_dw_kernel[grid_dw](
+            d_scores,
+            dot_c,
+            d_w_fp32,
+            B,
+            S,
+            P,
+            H=H,
+            COMPRESS_RATIO=int(compress_ratio),
+            BLOCK_S=block_s,
+            BLOCK_P_INNER=block_p_inner,
+        )
+
+        return d_dot_fp32.to(dot_dtype), d_w_fp32.to(w_dtype), None, None
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff ``PRIMUS_INDEXER_TRITON != "0"`` (default ``"1"``).
+
+    The env knob is **re-purposed** at P41 to mean "post-einsum tail
+    fusion" (cheap, bandwidth-bound).  Legacy P38 full-fuse path
+    lives behind ``PRIMUS_INDEXER_TRITON_FULL``.
+
+    Plan-8 P57 close-out 2 (2026-05-15): default flipped from ``"0"``
+    to ``"1"``.  Microbench at V4-Flash widths is consistently positive
+    (FWD 4.30x / BWD 1.63x) and the EP=8 proxy A/B shows a small but
+    positive ~0.2 ms / iter gain.  Set ``PRIMUS_INDEXER_TRITON=0`` to
+    revert to the eager Python body.
+    """
+    return os.environ.get("PRIMUS_INDEXER_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(dot: torch.Tensor, w_i: torch.Tensor) -> bool:
+    """Return True iff the input shapes / device support the Triton path."""
+    if not dot.is_cuda or not w_i.is_cuda:
+        return False
+    if dot.dim() != 4 or w_i.dim() != 3:
+        return False
+    _, _, H, _ = dot.shape
+    if H not in _SUPPORTED_H:
+        return False
+    return True
+
+
+def indexer_score_post_triton(
+    dot: torch.Tensor,
+    w_i: torch.Tensor,
+    *,
+    compress_ratio: int,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Compute Indexer scores from the eager einsum output via Triton.
+
+    Returns ``scores [B, S, P]`` of dtype ``out_dtype``.  Masked
+    positions hold ``-inf``.
+    """
+    return IndexerScorePostFn.apply(dot, w_i, compress_ratio, out_dtype)
+
+
+__all__ = [
+    "IndexerScorePostFn",
+    "indexer_score_post_triton",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py
new file mode 100644
index 000000000..ca8410dc6
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py
@@ -0,0 +1,366 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused RMSNorm FWD/BWD (small-kernel-fusion campaign 2026-07-03).
+
+Collapses the eager RMSNorm chain — ``x.float()`` (bf16→fp32 cast),
+``x.pow(2)`` / ``x.square()``, ``.mean(-1)``, ``+ eps``, ``rsqrt``,
+``* rstd``, ``.to(in_dtype)`` (fp32→bf16 cast), optional ``* weight`` —
+into ONE Triton kernel (FWD) + ONE Triton kernel (BWD).
+
+Every non-TE eager RMS site in the DeepSeek-V4 model body routes through
+this kernel:
+
+* ``_per_head_rms_norm`` (``deepseek_v4_attention``)   — no weight, out=in_dtype.
+* ``LocalRMSNorm``       (``compressor.kv_norm`` etc.) — weight (+grad),
+  mid-cast to in_dtype before the weight multiply, out=promote(in, weight).
+* ``HyperMixer._packed_logits`` RMS                    — no weight, out=fp32.
+* ``HyperHead.forward``  RMS                           — no weight, out=fp32.
+
+The eager reference (matched bit-for-bit modulo fp32 accumulation order):
+
+.. code-block:: python
+
+    x32 = x.float()
+    rstd = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    y = x32 * rstd                       # fp32 normalized value
+    if mid_cast: y = y.to(in_dtype)      # LocalRMSNorm rounds here first
+    if weight is not None: y = y * weight
+    out = y.to(out_dtype)
+
+Gradients (weight applied per channel, mid-cast treated as identity for
+autograd — matches ``Tensor.to`` grad semantics):
+
+.. code-block:: python
+
+    # forward:  y_k = w_k * rstd * x_k,   rstd = (mean(x^2)+eps)^-1/2
+    c   = sum_k( g_k * w_k * x_k )                 # reduce over D
+    dx_k = rstd * g_k * w_k - x_k * rstd**3 * c / D
+    dw_k = sum_over_rows( g_k * (x_k * rstd) )     # normalized value pre-weight
+
+Gating: routed through :func:`fused_rms_norm` when ``PRIMUS_RMSNORM_TRITON
+!= "0"`` (default-on) and the input is a supported CUDA/HIP float tensor;
+otherwise the eager reference runs.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+_TORCH_TO_TL_DTYPE = {
+    torch.float64: tl.float64,
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+def _triton_dtype(t: torch.dtype):
+    try:
+        return _TORCH_TO_TL_DTYPE[t]
+    except KeyError as exc:
+        raise TypeError(
+            f"rmsnorm: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}"
+        ) from exc
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _rmsnorm_fwd_kernel(
+    X_PTR,  # [N, D] contiguous
+    W_PTR,  # [D] contiguous (fp32) or dummy when HAS_WEIGHT is False
+    OUT_PTR,  # [N, D] contiguous (OUT_DTYPE)
+    RSTD_PTR,  # [N] fp32 (saved for backward)
+    N,
+    D,
+    EPS: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    MID_CAST: tl.constexpr,
+    IN_DTYPE: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """One program per row; two passes over D (sum-of-squares, then write)."""
+    row = tl.program_id(0)
+    if row >= N:
+        return
+    x_row = X_PTR + row * D
+
+    # Pass 1: sum of squares in fp32.
+    acc = tl.zeros((), dtype=tl.float32)
+    for off in range(0, D, BLOCK_D):
+        cols = off + tl.arange(0, BLOCK_D)
+        mask = cols < D
+        x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32)
+        acc += tl.sum(x * x, axis=0)
+
+    rstd = 1.0 / tl.sqrt(acc / D + EPS)
+    tl.store(RSTD_PTR + row, rstd)
+
+    # Pass 2: normalize (+ optional mid-cast, weight) and write.
+    for off in range(0, D, BLOCK_D):
+        cols = off + tl.arange(0, BLOCK_D)
+        mask = cols < D
+        x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32)
+        y = x * rstd
+        if MID_CAST:
+            y = y.to(IN_DTYPE).to(tl.float32)
+        if HAS_WEIGHT:
+            w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32)
+            y = y * w
+        tl.store(OUT_PTR + row * D + cols, y.to(OUT_DTYPE), mask=mask)
+
+
+@triton.jit
+def _rmsnorm_bwd_kernel(
+    X_PTR,  # [N, D] contiguous (original input)
+    W_PTR,  # [D] contiguous (fp32) or dummy
+    RSTD_PTR,  # [N] fp32
+    DY_PTR,  # [N, D] contiguous (upstream grad)
+    DX_PTR,  # [N, D] contiguous (output, IN_DTYPE)
+    DW_PTR,  # [D] fp32 accumulator (atomic) or dummy
+    N,
+    D,
+    BLOCK_D: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    IN_DTYPE: tl.constexpr,
+):
+    """One program per row. Computes dx (and atomic-accumulates dw)."""
+    row = tl.program_id(0)
+    if row >= N:
+        return
+    x_row = X_PTR + row * D
+    dy_row = DY_PTR + row * D
+    rstd = tl.load(RSTD_PTR + row)
+
+    # Pass 1: c = sum_k( dy_k * w_k * x_k ).
+    c = tl.zeros((), dtype=tl.float32)
+    for off in range(0, D, BLOCK_D):
+        cols = off + tl.arange(0, BLOCK_D)
+        mask = cols < D
+        x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32)
+        dy = tl.load(dy_row + cols, mask=mask, other=0.0).to(tl.float32)
+        if HAS_WEIGHT:
+            w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32)
+            c += tl.sum(dy * w * x, axis=0)
+        else:
+            c += tl.sum(dy * x, axis=0)
+
+    coef = rstd * rstd * rstd * c / D
+
+    # Pass 2: dx_k = rstd*dy_k*w_k - x_k*coef ; accumulate dw_k += dy_k*x_k*rstd.
+    for off in range(0, D, BLOCK_D):
+        cols = off + tl.arange(0, BLOCK_D)
+        mask = cols < D
+        x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32)
+        dy = tl.load(dy_row + cols, mask=mask, other=0.0).to(tl.float32)
+        if HAS_WEIGHT:
+            w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32)
+            dx = rstd * dy * w - x * coef
+            tl.atomic_add(DW_PTR + cols, dy * (x * rstd), mask=mask)
+        else:
+            dx = rstd * dy - x * coef
+        tl.store(DX_PTR + row * D + cols, dx.to(IN_DTYPE), mask=mask)
+
+
+def _pick_block_d(d: int) -> int:
+    """Tile the reduction axis; power-of-2, capped so registers stay bounded."""
+    if d <= 2048:
+        return triton.next_power_of_2(d)
+    return 2048
+
+
+# ---------------------------------------------------------------------------
+# torch.autograd.Function
+# ---------------------------------------------------------------------------
+
+
+class FusedRMSNormFn(torch.autograd.Function):
+    """Autograd wrapper for the fused RMSNorm FWD/BWD Triton kernels.
+
+    ``apply(x, weight, eps, mid_cast, out_dtype)``:
+
+    * ``x``: ``[..., D]`` float tensor (RMS over the last dim).
+    * ``weight``: ``[D]`` tensor or ``None`` (parameter-less RMS).
+    * ``eps``: numerical floor.
+    * ``mid_cast``: when ``True``, round the normalized value to ``x.dtype``
+      before the weight multiply (matches ``LocalRMSNorm``); no-op when
+      ``weight is None``.
+    * ``out_dtype``: dtype of the returned tensor.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        x: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        eps: float,
+        mid_cast: bool,
+        out_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        D = x.shape[-1]
+        x2 = x.reshape(-1, D)
+        x2 = x2.contiguous()
+        N = x2.shape[0]
+
+        has_weight = weight is not None
+        w32 = None
+        if has_weight:
+            if weight.shape[-1] != D:
+                raise ValueError(f"rmsnorm: weight dim {tuple(weight.shape)} != x last dim {D}")
+            w32 = weight.reshape(D).to(torch.float32).contiguous()
+
+        out = torch.empty((N, D), dtype=out_dtype, device=x.device)
+        rstd = torch.empty((N,), dtype=torch.float32, device=x.device)
+        block_d = _pick_block_d(D)
+        grid = (N,)
+        _rmsnorm_fwd_kernel[grid](
+            x2,
+            w32 if has_weight else x2,  # dummy ptr when no weight
+            out,
+            rstd,
+            N,
+            D,
+            EPS=float(eps),
+            BLOCK_D=block_d,
+            HAS_WEIGHT=has_weight,
+            MID_CAST=bool(mid_cast) and has_weight,
+            IN_DTYPE=_triton_dtype(x.dtype),
+            OUT_DTYPE=_triton_dtype(out_dtype),
+        )
+
+        ctx.save_for_backward(x2, w32, rstd)
+        ctx.has_weight = has_weight
+        ctx.in_dtype = x.dtype
+        ctx.weight_dtype = weight.dtype if has_weight else None
+        ctx.weight_shape = tuple(weight.shape) if has_weight else None
+        ctx.block_d = block_d
+        ctx.D = D
+        ctx.x_shape = tuple(x.shape)
+        return out.reshape(ctx.x_shape)
+
+    @staticmethod
+    def backward(ctx, dy: torch.Tensor):  # type: ignore[override]
+        x2, w32, rstd = ctx.saved_tensors
+        D = ctx.D
+        N = x2.shape[0]
+        has_weight = ctx.has_weight
+
+        # dy comes in out_dtype; the kernel upcasts to fp32 internally, so any
+        # float dtype is fine. Keep it as-is (contiguous) without a lossy cast.
+        dy2 = dy.reshape(-1, D).contiguous()
+
+        dx = torch.empty((N, D), dtype=ctx.in_dtype, device=dy.device)
+        dw_acc = (
+            torch.zeros((D,), dtype=torch.float32, device=dy.device)
+            if has_weight
+            else torch.empty((1,), dtype=torch.float32, device=dy.device)
+        )
+        grid = (N,)
+        _rmsnorm_bwd_kernel[grid](
+            x2,
+            w32 if has_weight else x2,
+            rstd,
+            dy2,
+            dx,
+            dw_acc,
+            N,
+            D,
+            BLOCK_D=ctx.block_d,
+            HAS_WEIGHT=has_weight,
+            IN_DTYPE=_triton_dtype(ctx.in_dtype),
+        )
+
+        dx = dx.reshape(ctx.x_shape)
+        dweight = None
+        if has_weight:
+            dweight = dw_acc.to(ctx.weight_dtype).reshape(ctx.weight_shape)
+        return dx, dweight, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Public entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Default-on; A/B toggle via ``PRIMUS_RMSNORM_TRITON=0``."""
+    return os.environ.get("PRIMUS_RMSNORM_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(x: torch.Tensor, weight: Optional[torch.Tensor]) -> bool:
+    """Supported iff CUDA/HIP float input (and matching-dtype-family weight)."""
+    if not x.is_cuda:
+        return False
+    if x.dtype not in _TORCH_TO_TL_DTYPE:
+        return False
+    if x.shape[-1] == 0:
+        return False
+    if weight is not None and not weight.is_cuda:
+        return False
+    return True
+
+
+def eager_rms_norm(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    *,
+    eps: float,
+    mid_cast: bool,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Reference eager RMSNorm matching the fused kernel's contract."""
+    in_dtype = x.dtype
+    x32 = x.float()
+    rstd = torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + eps)
+    y = x32 * rstd
+    if weight is not None:
+        if mid_cast:
+            y = y.to(in_dtype)
+        y = y.to(torch.float32) * weight.to(torch.float32)
+    return y.to(out_dtype)
+
+
+def fused_rms_norm(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    *,
+    eps: float,
+    mid_cast: bool = False,
+    out_dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    """Dispatch: Triton path when enabled + supported, else eager reference.
+
+    ``out_dtype`` defaults to the eager output dtype:
+    ``promote_types(x.dtype, weight.dtype)`` when weighted, else ``x.dtype``.
+    """
+    if out_dtype is None:
+        if weight is not None:
+            out_dtype = torch.promote_types(x.dtype, weight.dtype)
+        else:
+            out_dtype = x.dtype
+
+    if is_triton_path_enabled() and is_triton_kernel_supported(x, weight):
+        return FusedRMSNormFn.apply(x, weight, float(eps), bool(mid_cast), out_dtype)
+    return eager_rms_norm(x, weight, eps=eps, mid_cast=mid_cast, out_dtype=out_dtype)
+
+
+__all__ = [
+    "FusedRMSNormFn",
+    "fused_rms_norm",
+    "eager_rms_norm",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py
new file mode 100644
index 000000000..96f2e17ec
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py
@@ -0,0 +1,801 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused interleaved partial RoPE FWD/BWD (plan-6 P35).
+
+The eager body in :func:`apply_interleaved_partial_rope` is a 9-op chain:
+
+.. code-block:: python
+
+    x_nope = x[..., :nope]                           # slice  (1)
+    x_rope = x[..., nope:]                           # slice  (2)
+    x_pairs = x_rope.reshape(..., rd // 2, 2)        # reshape (3)
+    even, odd = x_pairs[..., 0], x_pairs[..., 1]
+    cos = cos.unsqueeze(-2).to(orig_dtype)           # unsqueeze + cast (4)
+    sin = sin.unsqueeze(-2).to(orig_dtype)           # unsqueeze + cast (5)
+    rot_even = even * cos - odd * sin                # 2 muls + 1 sub (6)
+    rot_odd  = even * sin + odd * cos                # 2 muls + 1 add (7)
+    rotated  = torch.stack([rot_even, rot_odd], -1).reshape(..., rd)  # stack + reshape (8)
+    return torch.cat([x_nope, rotated], -1)          # cat (9)
+
+The plan-5 P32 final EP=8 proxy trace attributes:
+
+* ``CatArrayBatchedCopy_contig`` ≈ **10.0 ms / 24 calls** to the closing
+  ``torch.cat`` (the nope-prefix copy + the rotated suffix into one
+  contiguous tensor), and
+* a non-trivial share of ``elementwise_kernel_manual_unroll<128, 8>``
+  (~61 ms / 693 calls) to the four broadcast muls.
+
+At 16 invocations per iter (q + k per ``DualRoPE`` call × 8 layers) the
+per-call cost is **~3-5 ms** at the V4-Flash widths.
+
+This module collapses the 9-op chain into one Triton kernel that:
+
+1. Flattens the input to ``[N, H, head_dim]`` where ``N`` is the product
+   of all leading axes (caller does the ``.reshape(-1, H, head_dim)``
+   plumbing; the kernel is shape-agnostic).
+2. Per program processes a ``[BLOCK_H]`` slice of one position's
+   ``H * head_dim`` row.  cos/sin are loaded **once per position**
+   (shared across the ``BLOCK_H`` heads in the program) so the per-call
+   HBM traffic for cos/sin is exactly ``N * rd / 2`` reads, not
+   ``N * H * rd / 2`` (which the broadcast-muls in the eager body would
+   imply if they hit memory).
+3. Writes ``out [N, H, head_dim]`` in one pass: nope channels copied
+   verbatim; trailing ``rotary_dim`` channels rotated using the
+   interleaved (2k, 2k+1) pairing.  No ``torch.cat``-style second
+   memcpy; the kernel writes the full output in a single contiguous
+   pass.
+
+The BWD kernel is the transpose of the FWD rotation matrix
+(``cos, sin / -sin, cos -> cos, -sin / sin, cos``) — analytically:
+
+.. code-block:: python
+
+    dx[..., 2k]     = dout[..., 2k] * cos + dout[..., 2k+1] * sin
+    dx[..., 2k+1]   = -dout[..., 2k] * sin + dout[..., 2k+1] * cos
+    dx[..., :nope]  = dout[..., :nope]                       # straight copy
+
+Cos / sin are buffers (not Parameters) so they have no gradient.
+
+Gating: routed through :func:`apply_interleaved_partial_rope` when
+``PRIMUS_ROPE_TRITON != "0"`` (default-on).  Set to ``"0"`` to fall back
+to the eager body (kept in tree for A/B and as the reference path for
+G38).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Triton dtype mapping
+# ---------------------------------------------------------------------------
+
+_TORCH_TO_TL_DTYPE = {
+    torch.float64: tl.float64,
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+def _triton_dtype(t: torch.dtype):
+    try:
+        return _TORCH_TO_TL_DTYPE[t]
+    except KeyError as exc:
+        raise TypeError(
+            f"rope_interleaved_partial: unsupported dtype {t}; " f"expected one of {list(_TORCH_TO_TL_DTYPE)}"
+        ) from exc
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _apply_rope_fwd_kernel(
+    X_PTR,  # [N, H, head_dim] contiguous
+    COS_PTR,  # [N, rd_half] contiguous (broadcast over H)
+    SIN_PTR,  # [N, rd_half] contiguous (broadcast over H)
+    OUT_PTR,  # [N, H, head_dim] contiguous
+    N,
+    H,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    RD_HALF: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_NOPE: tl.constexpr,  # next_pow2(NOPE) — block stride for nope copy
+    BLOCK_RD_HALF: tl.constexpr,  # next_pow2(RD_HALF)
+    DTYPE: tl.constexpr,
+):
+    """Apply interleaved partial RoPE FWD over a tile of ``BLOCK_H`` heads
+    for one position ``pid_n``.
+
+    Layout: x / out are ``[N, H, head_dim]``; the program processes
+    ``x[pid_n, pid_h*BLOCK_H : (pid_h+1)*BLOCK_H, :]`` and writes the
+    rotated result to the same slice of ``out``.
+
+    The nope prefix (``head_dim - rotary_dim`` channels) is copied
+    verbatim — kept inside the kernel so the eager-body
+    ``torch.cat([x_nope, rotated], -1)`` second-pass memcpy goes away.
+    """
+
+    pid_n = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = h_offs < H
+
+    # cos / sin for this position — load once, broadcast across BLOCK_H heads.
+    rd_half_offs = tl.arange(0, BLOCK_RD_HALF)
+    rd_half_mask = rd_half_offs < RD_HALF
+    cos = tl.load(
+        COS_PTR + pid_n * RD_HALF + rd_half_offs,
+        mask=rd_half_mask,
+        other=0.0,
+    )
+    sin = tl.load(
+        SIN_PTR + pid_n * RD_HALF + rd_half_offs,
+        mask=rd_half_mask,
+        other=0.0,
+    )
+    cos = cos.to(DTYPE)
+    sin = sin.to(DTYPE)
+
+    # Base pointers for this position × head block.
+    row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM  # [BLOCK_H, 1]
+
+    # 1) Copy the nope channels verbatim.
+    if NOPE > 0:
+        nope_offs = tl.arange(0, BLOCK_NOPE)
+        nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None]
+        x_nope = tl.load(
+            X_PTR + row_base + nope_offs[None, :],
+            mask=nope_mask,
+            other=0.0,
+        )
+        tl.store(
+            OUT_PTR + row_base + nope_offs[None, :],
+            x_nope,
+            mask=nope_mask,
+        )
+
+    # 2) Rotate the trailing rotary_dim channels (interleaved pairs).
+    even_offs = NOPE + 2 * rd_half_offs
+    odd_offs = NOPE + 2 * rd_half_offs + 1
+    pair_mask = h_mask[:, None] & rd_half_mask[None, :]
+    even = tl.load(
+        X_PTR + row_base + even_offs[None, :],
+        mask=pair_mask,
+        other=0.0,
+    )
+    odd = tl.load(
+        X_PTR + row_base + odd_offs[None, :],
+        mask=pair_mask,
+        other=0.0,
+    )
+    rot_even = even * cos[None, :] - odd * sin[None, :]
+    rot_odd = even * sin[None, :] + odd * cos[None, :]
+    tl.store(
+        OUT_PTR + row_base + even_offs[None, :],
+        rot_even,
+        mask=pair_mask,
+    )
+    tl.store(
+        OUT_PTR + row_base + odd_offs[None, :],
+        rot_odd,
+        mask=pair_mask,
+    )
+
+
+@triton.jit
+def _apply_rope_bwd_kernel(
+    DOUT_PTR,  # [N, H, head_dim] contiguous
+    COS_PTR,  # [N, rd_half] contiguous
+    SIN_PTR,  # [N, rd_half] contiguous
+    DX_PTR,  # [N, H, head_dim] contiguous
+    N,
+    H,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    RD_HALF: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_NOPE: tl.constexpr,
+    BLOCK_RD_HALF: tl.constexpr,
+    DTYPE: tl.constexpr,
+):
+    """Apply the transpose rotation for the BWD pass.
+
+    .. code-block:: python
+
+        dx[..., 2k]     = dout[..., 2k] * cos + dout[..., 2k+1] * sin
+        dx[..., 2k+1]   = -dout[..., 2k] * sin + dout[..., 2k+1] * cos
+        dx[..., :nope]  = dout[..., :nope]
+    """
+
+    pid_n = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = h_offs < H
+
+    rd_half_offs = tl.arange(0, BLOCK_RD_HALF)
+    rd_half_mask = rd_half_offs < RD_HALF
+    cos = tl.load(
+        COS_PTR + pid_n * RD_HALF + rd_half_offs,
+        mask=rd_half_mask,
+        other=0.0,
+    )
+    sin = tl.load(
+        SIN_PTR + pid_n * RD_HALF + rd_half_offs,
+        mask=rd_half_mask,
+        other=0.0,
+    )
+    cos = cos.to(DTYPE)
+    sin = sin.to(DTYPE)
+
+    row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM
+
+    # 1) Straight copy of the nope-prefix gradient.
+    if NOPE > 0:
+        nope_offs = tl.arange(0, BLOCK_NOPE)
+        nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None]
+        dout_nope = tl.load(
+            DOUT_PTR + row_base + nope_offs[None, :],
+            mask=nope_mask,
+            other=0.0,
+        )
+        tl.store(
+            DX_PTR + row_base + nope_offs[None, :],
+            dout_nope,
+            mask=nope_mask,
+        )
+
+    # 2) Transposed rotation for the rotary suffix.
+    even_offs = NOPE + 2 * rd_half_offs
+    odd_offs = NOPE + 2 * rd_half_offs + 1
+    pair_mask = h_mask[:, None] & rd_half_mask[None, :]
+    dout_even = tl.load(
+        DOUT_PTR + row_base + even_offs[None, :],
+        mask=pair_mask,
+        other=0.0,
+    )
+    dout_odd = tl.load(
+        DOUT_PTR + row_base + odd_offs[None, :],
+        mask=pair_mask,
+        other=0.0,
+    )
+    dx_even = dout_even * cos[None, :] + dout_odd * sin[None, :]
+    dx_odd = -dout_even * sin[None, :] + dout_odd * cos[None, :]
+    tl.store(
+        DX_PTR + row_base + even_offs[None, :],
+        dx_even,
+        mask=pair_mask,
+    )
+    tl.store(
+        DX_PTR + row_base + odd_offs[None, :],
+        dx_odd,
+        mask=pair_mask,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Block-size heuristic
+# ---------------------------------------------------------------------------
+
+
+def _pick_block_h(h: int) -> int:
+    """Pick BLOCK_H tiling the heads axis.
+
+    V4-Flash Q has ``H = 64`` heads, K has ``H = 1``.  For H=1 the block
+    must be 1; for H=64 a BLOCK_H of 8 keeps the per-program live tile
+    under ~16 KiB at head_dim=512 / bf16 while amortising the cos/sin
+    load across 8 heads.
+    """
+
+    if h <= 1:
+        return 1
+    if h <= 4:
+        return min(h, 4)
+    if h <= 16:
+        return 8
+    return 8  # default for V4-Flash Q (H=64) and K (H=1 handled above)
+
+
+# ---------------------------------------------------------------------------
+# torch.autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class RoPEInterleavedPartialFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the FWD/BWD Triton kernels.
+
+    The wrapper:
+
+    1. Flattens ``x`` to ``[N, H, head_dim]`` and ``cos / sin`` to
+       ``[N, rd_half]`` (callers may pass any leading shape; the wrapper
+       does the reshape).
+    2. Casts cos / sin to ``x.dtype`` if needed (matches the plan-5 P32
+       RoPE bf16 cast contract in :func:`apply_interleaved_partial_rope`).
+    3. Saves only ``cos, sin, rotary_dim, head_dim, nope`` for the
+       backward — the input ``x`` is not needed, and the saved tensors
+       are buffers (cos / sin) so they have no autograd graph node.
+
+    Layout assumptions (validated up front):
+
+    * ``x.is_contiguous()``  — call ``.contiguous()`` in the caller if needed.
+    * ``cos.shape == sin.shape == leading_shape + (rotary_dim // 2,)`` where
+      ``leading_shape == x.shape[:-2]`` (i.e. one cos/sin row per position).
+    * ``rotary_dim <= x.shape[-1]`` and ``rotary_dim % 2 == 0``.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        rotary_dim: int,
+    ) -> torch.Tensor:
+        if rotary_dim == 0:
+            ctx.save_for_backward(cos, sin)
+            ctx.rotary_dim = 0
+            ctx.head_dim = x.shape[-1]
+            return x.contiguous()
+
+        if rotary_dim % 2 != 0:
+            raise ValueError(f"rotary_dim must be even, got {rotary_dim}")
+        head_dim = x.shape[-1]
+        if rotary_dim > head_dim:
+            raise ValueError(f"rotary_dim ({rotary_dim}) must be <= head_dim ({head_dim})")
+
+        x = x.contiguous()
+        leading_shape = x.shape[:-2]
+        H = x.shape[-2]
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        rd_half = rotary_dim // 2
+        # cos / sin should have shape == leading_shape + (rd_half,)
+        if cos.shape[-1] != rd_half or sin.shape[-1] != rd_half:
+            raise ValueError(
+                f"cos/sin last dim must be rotary_dim // 2 ({rd_half}); "
+                f"got cos.shape={tuple(cos.shape)}, sin.shape={tuple(sin.shape)}"
+            )
+        cos_flat = cos.contiguous().reshape(N, rd_half).to(x.dtype)
+        sin_flat = sin.contiguous().reshape(N, rd_half).to(x.dtype)
+
+        out = torch.empty_like(x)
+
+        nope = head_dim - rotary_dim
+        block_h = _pick_block_h(H)
+        block_nope = max(triton.next_power_of_2(max(nope, 1)), 1)
+        block_rd_half = triton.next_power_of_2(rd_half)
+
+        grid = (N, triton.cdiv(H, block_h))
+        _apply_rope_fwd_kernel[grid](
+            x,
+            cos_flat,
+            sin_flat,
+            out,
+            N,
+            H,
+            HEAD_DIM=head_dim,
+            NOPE=nope,
+            RD_HALF=rd_half,
+            BLOCK_H=block_h,
+            BLOCK_NOPE=block_nope,
+            BLOCK_RD_HALF=block_rd_half,
+            DTYPE=_triton_dtype(x.dtype),
+        )
+
+        ctx.save_for_backward(cos_flat, sin_flat)
+        ctx.rotary_dim = rotary_dim
+        ctx.head_dim = head_dim
+        ctx.leading_shape = tuple(leading_shape)
+        ctx.H = H
+        return out
+
+    @staticmethod
+    def backward(ctx, dout: torch.Tensor):  # type: ignore[override]
+        rotary_dim = ctx.rotary_dim
+        cos_flat, sin_flat = ctx.saved_tensors
+
+        if rotary_dim == 0:
+            return dout.contiguous(), None, None, None
+
+        dout = dout.contiguous()
+        head_dim = ctx.head_dim
+        H = ctx.H
+        leading_shape = ctx.leading_shape
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        rd_half = rotary_dim // 2
+        nope = head_dim - rotary_dim
+        block_h = _pick_block_h(H)
+        block_nope = max(triton.next_power_of_2(max(nope, 1)), 1)
+        block_rd_half = triton.next_power_of_2(rd_half)
+
+        dx = torch.empty_like(dout)
+
+        grid = (N, triton.cdiv(H, block_h))
+        _apply_rope_bwd_kernel[grid](
+            dout,
+            cos_flat,
+            sin_flat,
+            dx,
+            N,
+            H,
+            HEAD_DIM=head_dim,
+            NOPE=nope,
+            RD_HALF=rd_half,
+            BLOCK_H=block_h,
+            BLOCK_NOPE=block_nope,
+            BLOCK_RD_HALF=block_rd_half,
+            DTYPE=_triton_dtype(dout.dtype),
+        )
+        return dx, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Fused variant: compute cos/sin IN-KERNEL from (position_ids, inv_freq).
+#
+# The plain RoPE kernel above consumes precomputed cos/sin tensors, which
+# means the caller still pays 3 small kernels per call
+# (``position_ids.float() * inv_freq`` -> ``cos`` -> ``sin``) plus a
+# ``.to(x.dtype)`` cast, and (in ``DeepseekV4Attention._apply_rope_q_k``)
+# recomputes them identically for Q and K.  This variant folds the cos/sin
+# generation into the rotation kernel: it loads the row's scalar position +
+# the ``[rd_half]`` inv_freq vector and computes ``cos``/``sin`` in registers,
+# so no cos/sin HBM tensor is ever materialised.  YaRN is already baked into
+# ``inv_freq`` (a buffer), so the compress base is handled transparently.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _rope_gen_fwd_kernel(
+    X_PTR,  # [N, H, head_dim] contiguous
+    POS_PTR,  # [N] (position per row; fp32)
+    INVFREQ_PTR,  # [rd_half] fp32
+    OUT_PTR,  # [N, H, head_dim] contiguous
+    N,
+    H,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    RD_HALF: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_NOPE: tl.constexpr,
+    BLOCK_RD_HALF: tl.constexpr,
+    DTYPE: tl.constexpr,
+    INVERSE: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = h_offs < H
+
+    rd_half_offs = tl.arange(0, BLOCK_RD_HALF)
+    rd_half_mask = rd_half_offs < RD_HALF
+
+    # cos/sin for this position, computed once, broadcast across BLOCK_H heads.
+    pos = tl.load(POS_PTR + pid_n).to(tl.float32)
+    inv_freq = tl.load(INVFREQ_PTR + rd_half_offs, mask=rd_half_mask, other=0.0).to(tl.float32)
+    angle = pos * inv_freq
+    # ``R(-pos) == R(pos)^T``, i.e. the same rotation with a negated sine, so
+    # the inverse costs a sign flip rather than a second kernel.
+    if INVERSE:
+        angle = -angle
+    cos = tl.cos(angle).to(DTYPE)
+    sin = tl.sin(angle).to(DTYPE)
+
+    row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM
+
+    if NOPE > 0:
+        nope_offs = tl.arange(0, BLOCK_NOPE)
+        nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None]
+        x_nope = tl.load(X_PTR + row_base + nope_offs[None, :], mask=nope_mask, other=0.0)
+        tl.store(OUT_PTR + row_base + nope_offs[None, :], x_nope, mask=nope_mask)
+
+    even_offs = NOPE + 2 * rd_half_offs
+    odd_offs = NOPE + 2 * rd_half_offs + 1
+    pair_mask = h_mask[:, None] & rd_half_mask[None, :]
+    even = tl.load(X_PTR + row_base + even_offs[None, :], mask=pair_mask, other=0.0)
+    odd = tl.load(X_PTR + row_base + odd_offs[None, :], mask=pair_mask, other=0.0)
+    rot_even = even * cos[None, :] - odd * sin[None, :]
+    rot_odd = even * sin[None, :] + odd * cos[None, :]
+    tl.store(OUT_PTR + row_base + even_offs[None, :], rot_even, mask=pair_mask)
+    tl.store(OUT_PTR + row_base + odd_offs[None, :], rot_odd, mask=pair_mask)
+
+
+@triton.jit
+def _rope_gen_bwd_kernel(
+    DOUT_PTR,  # [N, H, head_dim] contiguous
+    POS_PTR,  # [N] fp32
+    INVFREQ_PTR,  # [rd_half] fp32
+    DX_PTR,  # [N, H, head_dim] contiguous
+    N,
+    H,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    RD_HALF: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_NOPE: tl.constexpr,
+    BLOCK_RD_HALF: tl.constexpr,
+    DTYPE: tl.constexpr,
+    INVERSE: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    h_mask = h_offs < H
+
+    rd_half_offs = tl.arange(0, BLOCK_RD_HALF)
+    rd_half_mask = rd_half_offs < RD_HALF
+
+    pos = tl.load(POS_PTR + pid_n).to(tl.float32)
+    inv_freq = tl.load(INVFREQ_PTR + rd_half_offs, mask=rd_half_mask, other=0.0).to(tl.float32)
+    angle = pos * inv_freq
+    # Must mirror the forward's sign so the transposed rotation below is the
+    # actual adjoint of what was applied.
+    if INVERSE:
+        angle = -angle
+    cos = tl.cos(angle).to(DTYPE)
+    sin = tl.sin(angle).to(DTYPE)
+
+    row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM
+
+    if NOPE > 0:
+        nope_offs = tl.arange(0, BLOCK_NOPE)
+        nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None]
+        dout_nope = tl.load(DOUT_PTR + row_base + nope_offs[None, :], mask=nope_mask, other=0.0)
+        tl.store(DX_PTR + row_base + nope_offs[None, :], dout_nope, mask=nope_mask)
+
+    even_offs = NOPE + 2 * rd_half_offs
+    odd_offs = NOPE + 2 * rd_half_offs + 1
+    pair_mask = h_mask[:, None] & rd_half_mask[None, :]
+    dout_even = tl.load(DOUT_PTR + row_base + even_offs[None, :], mask=pair_mask, other=0.0)
+    dout_odd = tl.load(DOUT_PTR + row_base + odd_offs[None, :], mask=pair_mask, other=0.0)
+    dx_even = dout_even * cos[None, :] + dout_odd * sin[None, :]
+    dx_odd = -dout_even * sin[None, :] + dout_odd * cos[None, :]
+    tl.store(DX_PTR + row_base + even_offs[None, :], dx_even, mask=pair_mask)
+    tl.store(DX_PTR + row_base + odd_offs[None, :], dx_odd, mask=pair_mask)
+
+
+class RoPEFromPositionsFn(torch.autograd.Function):
+    """RoPE that computes cos/sin in-kernel from ``(position_ids, inv_freq)``.
+
+    Eliminates the separate ``cos``/``sin`` generation kernels and HBM
+    tensors.  ``position_ids`` / ``inv_freq`` are non-differentiable
+    (positions are integer, inv_freq is a buffer), so the backward returns
+    ``None`` for both and recomputes cos/sin in-kernel.
+
+    ``inverse=True`` rotates by ``-position`` instead, which is what the V4
+    attention output needs before the O projection. It is the same kernel with
+    a negated angle, so the de-rotation gets the same single-pass, no-cos/sin-in-
+    HBM treatment as Q / KV instead of paying for a separate cos/sin build.
+    """
+
+    @staticmethod
+    def forward(ctx, x, pos_flat, inv_freq, rotary_dim, inverse=False):  # type: ignore[override]
+        head_dim = x.shape[-1]
+        if rotary_dim % 2 != 0:
+            raise ValueError(f"rotary_dim must be even, got {rotary_dim}")
+        if rotary_dim > head_dim:
+            raise ValueError(f"rotary_dim ({rotary_dim}) must be <= head_dim ({head_dim})")
+
+        x = x.contiguous()
+        leading = x.shape[:-2]
+        H = x.shape[-2]
+        N = 1
+        for s in leading:
+            N *= s
+        rd_half = rotary_dim // 2
+        if inv_freq.shape[-1] != rd_half:
+            raise ValueError(
+                f"inv_freq last dim must be rotary_dim // 2 ({rd_half}); got {tuple(inv_freq.shape)}"
+            )
+
+        pos_c = pos_flat.contiguous().to(torch.float32)
+        inv_c = inv_freq.contiguous().to(torch.float32)
+        out = torch.empty_like(x)
+
+        nope = head_dim - rotary_dim
+        block_h = _pick_block_h(H)
+        block_nope = max(triton.next_power_of_2(max(nope, 1)), 1)
+        block_rd_half = triton.next_power_of_2(rd_half)
+        grid = (N, triton.cdiv(H, block_h))
+        _rope_gen_fwd_kernel[grid](
+            x,
+            pos_c,
+            inv_c,
+            out,
+            N,
+            H,
+            HEAD_DIM=head_dim,
+            NOPE=nope,
+            RD_HALF=rd_half,
+            BLOCK_H=block_h,
+            BLOCK_NOPE=block_nope,
+            BLOCK_RD_HALF=block_rd_half,
+            DTYPE=_triton_dtype(x.dtype),
+            INVERSE=bool(inverse),
+        )
+
+        ctx.save_for_backward(pos_c, inv_c)
+        ctx.rotary_dim = rotary_dim
+        ctx.head_dim = head_dim
+        ctx.H = H
+        ctx.N = N
+        ctx.inverse = bool(inverse)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):  # type: ignore[override]
+        pos_c, inv_c = ctx.saved_tensors
+        rotary_dim = ctx.rotary_dim
+        head_dim = ctx.head_dim
+        H = ctx.H
+        N = ctx.N
+        dout = dout.contiguous()
+        rd_half = rotary_dim // 2
+        nope = head_dim - rotary_dim
+        block_h = _pick_block_h(H)
+        block_nope = max(triton.next_power_of_2(max(nope, 1)), 1)
+        block_rd_half = triton.next_power_of_2(rd_half)
+        dx = torch.empty_like(dout)
+        grid = (N, triton.cdiv(H, block_h))
+        _rope_gen_bwd_kernel[grid](
+            dout,
+            pos_c,
+            inv_c,
+            dx,
+            N,
+            H,
+            HEAD_DIM=head_dim,
+            NOPE=nope,
+            RD_HALF=rd_half,
+            BLOCK_H=block_h,
+            BLOCK_NOPE=block_nope,
+            BLOCK_RD_HALF=block_rd_half,
+            DTYPE=_triton_dtype(dout.dtype),
+            INVERSE=ctx.inverse,
+        )
+        return dx, None, None, None, None
+
+
+def apply_rope_from_positions(
+    x: torch.Tensor,
+    position_ids: torch.Tensor,
+    inv_freq: torch.Tensor,
+    *,
+    rotary_dim: int,
+    inverse: bool = False,
+) -> torch.Tensor:
+    """Fused interleaved partial RoPE that generates cos/sin in-kernel.
+
+    ``x``: ``[..., H, head_dim]``; ``position_ids`` broadcastable to
+    ``x.shape[:-2]``; ``inv_freq``: ``[rotary_dim // 2]`` (YaRN pre-applied).
+    Dispatches to the Triton kernel when enabled + CUDA, else falls back to
+    the eager ``cos = pos*inv_freq -> cos/sin -> rotate`` path.
+
+    ``inverse=True`` applies ``R(-position)``, which the V4 attention output
+    needs before the O projection. Both paths are out-of-place.
+    """
+    if rotary_dim == 0:
+        return x
+    if is_triton_path_enabled() and x.is_cuda:
+        leading = x.shape[:-2]
+        pos_flat = position_ids.broadcast_to(leading).reshape(-1)
+        return RoPEFromPositionsFn.apply(x, pos_flat, inv_freq, rotary_dim, inverse)
+    # Eager fallback: build cos/sin then apply.
+    freqs = position_ids.float().unsqueeze(-1) * inv_freq
+    if inverse:
+        freqs = -freqs
+    return eager_apply_interleaved_partial_rope(x, freqs.cos(), freqs.sin(), rotary_dim=rotary_dim)
+
+
+# ---------------------------------------------------------------------------
+# Public Python entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff the ``PRIMUS_ROPE_TRITON`` env knob is not ``"0"``.
+
+    Default-on, A/B toggle via ``PRIMUS_ROPE_TRITON=0``.
+    """
+
+    return os.environ.get("PRIMUS_ROPE_TRITON", "1") != "0"
+
+
+def eager_apply_interleaved_partial_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    rotary_dim: int,
+) -> torch.Tensor:
+    """Reference eager implementation matching
+    :func:`primus.backends.megatron.core.transformer.dual_rope.apply_interleaved_partial_rope`.
+
+    Kept here so the test / bench code can A/B against the exact eager
+    body without depending on the consumer module.  Bit-for-bit
+    equivalent to the original eager body (same op order, same dtype
+    cast).
+    """
+
+    head_dim = x.shape[-1]
+    if rotary_dim > head_dim or rotary_dim % 2 != 0:
+        raise ValueError(f"rotary_dim must be even and <= head_dim ({head_dim}), got {rotary_dim}")
+    if rotary_dim == 0:
+        return x
+
+    orig_dtype = x.dtype
+    nope = head_dim - rotary_dim
+    x_nope = x[..., :nope]
+    x_rope = x[..., nope:]
+
+    x_pairs = x_rope.reshape(*x_rope.shape[:-1], rotary_dim // 2, 2)
+    even = x_pairs[..., 0]
+    odd = x_pairs[..., 1]
+
+    cos = cos.unsqueeze(-2).to(orig_dtype)
+    sin = sin.unsqueeze(-2).to(orig_dtype)
+
+    rot_even = even * cos - odd * sin
+    rot_odd = even * sin + odd * cos
+
+    rotated = torch.stack([rot_even, rot_odd], dim=-1).reshape(*x_rope.shape[:-1], rotary_dim)
+    return torch.cat([x_nope, rotated], dim=-1)
+
+
+def apply_rope_interleaved_partial(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    rotary_dim: int,
+) -> torch.Tensor:
+    """Dispatch: Triton path when ``PRIMUS_ROPE_TRITON != "0"`` (default),
+    else eager fallback.
+
+    Shape contract matches the dual-RoPE caller:
+
+    * ``x``: ``[..., H, head_dim]`` (any leading shape; flattened
+      internally).
+    * ``cos, sin``: ``[..., rotary_dim // 2]`` where the leading shape
+      is ``x.shape[:-2]`` (i.e. one cos/sin row per position).
+    * ``rotary_dim``: even and ``<= head_dim``.
+
+    Output shape matches ``x``.
+    """
+
+    if rotary_dim == 0:
+        return x
+
+    if is_triton_path_enabled() and x.is_cuda:
+        return RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim)
+    return eager_apply_interleaved_partial_rope(x, cos, sin, rotary_dim=rotary_dim)
+
+
+__all__ = [
+    "RoPEInterleavedPartialFn",
+    "RoPEFromPositionsFn",
+    "apply_rope_interleaved_partial",
+    "apply_rope_from_positions",
+    "eager_apply_interleaved_partial_rope",
+    "is_triton_path_enabled",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py
new file mode 100644
index 000000000..70c622eee
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py
@@ -0,0 +1,520 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton-fused Sinkhorn-Knopp FWD/BWD (plan-6 P36).
+
+Replaces the plan-5 P29 ``torch.compile`` Sinkhorn path with a
+hand-rolled Triton kernel that runs the entire alternating row/col
+normalize trajectory **in registers** per row of the leading axis.
+The plan-5 P32 final trace attributes:
+
+* ``Torch-Compiled Region``         ≈ **21 ms / 16 calls** (FWD-side)
+* ``CompiledFunctionBackward``      ≈ **41 ms / 16 calls** (BWD-side)
+
+to the cached compiled artefact built in
+:func:`primus.backends.megatron.core.transformer.hyper_connection._build_compiled_sinkhorn`.
+``torch.compile`` collapses the 1 + 2*(n_iters - 1) fp32 reductions
+into one Inductor-fused Triton kernel, but its Dynamo-side
+bookkeeping (the ``Torch-Compiled Region`` event itself) and
+Inductor's BWD path are still non-trivial per-call overhead.
+
+The Triton path emits **one** FWD kernel and **one** BWD kernel; no
+Dynamo bookkeeping per call.  At V4-Flash widths
+(``[..., K=4, K=4]``, ``n_iters=20``, fp32 internal compute) the
+``K*K = 16`` elements per row fit in fp32 registers; the 39 normalize
+steps (1 priming col-normalize + 19 row+col pairs) all run in
+registers; the BWD recomputes the FWD trajectory in registers and
+walks the analytic VJP backward step-by-step (the doubly-stochastic
+projection has a closed-form VJP per step that fits in the same
+register budget).
+
+Eager reference (see :func:`primus...hyper_connection.sinkhorn_normalize`):
+
+.. code-block:: python
+
+    in_dtype = logits.dtype
+    m = logits.float()
+    m = m / (m.sum(dim=-2, keepdim=True) + eps)     # initial col-normalize
+    for _ in range(n_iters - 1):
+        m = m / (m.sum(dim=-1, keepdim=True) + eps) # row-normalize
+        m = m / (m.sum(dim=-2, keepdim=True) + eps) # col-normalize
+    return m.to(in_dtype)
+
+Per-step VJP (closed form):
+
+.. code-block:: python
+
+    # y = x / s, where s = sum(x, axis=axis) + eps
+    # dx = (dy - sum(dy * y, axis=axis, keep_dims=True)) / s
+
+Both axes (row / col) have the same shape of VJP (Triton's ``tl.sum``
+handles either with the right ``axis=`` argument), so the BWD kernel
+is just the FWD-trajectory recompute followed by 39 of these VJP
+steps in reverse order.
+
+Gating: routed through :func:`apply_sinkhorn_normalize_triton` when
+``PRIMUS_SINKHORN_TRITON != "0"`` (default-on).  Set to ``"0"`` to
+fall back to either the plan-5 P29 compiled path (if
+``use_compiled=True`` on the call site) or the eager path.
+
+Supported shape constraints:
+
+* Last two dims must be square and ``K ∈ {1, 2, 4, 8, 16}`` (Triton's
+  ``tl.arange`` needs a power-of-2 block); V4-Flash uses ``K=4``.
+* ``logits`` must be contiguous (the wrapper calls ``.contiguous()``
+  defensively).
+* Any leading shape is supported (the wrapper flattens to
+  ``[N, K, K]``).
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Triton dtype mapping
+# ---------------------------------------------------------------------------
+
+_TORCH_TO_TL_DTYPE = {
+    torch.float64: tl.float64,
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+def _triton_dtype(t: torch.dtype):
+    try:
+        return _TORCH_TO_TL_DTYPE[t]
+    except KeyError as exc:
+        raise TypeError(
+            f"sinkhorn: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}"
+        ) from exc
+
+
+# Supported K (last two dim) values for the in-register Triton path.
+# Block-size constraints (Triton requires power-of-2 block extents) plus
+# register budget (~K*K*BLOCK_LEADING fp32 elements live per program)
+# practically limit K to {1, 2, 4, 8, 16}.  V4-Flash uses K=4.
+_SUPPORTED_K = (1, 2, 4, 8, 16)
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _sinkhorn_fwd_kernel(
+    X_PTR,  # [N, K, K] contiguous
+    Y_PTR,  # [N, K, K] contiguous
+    STATES_PTR,  # [N, 2 * N_ITERS, K, K] contiguous (FWD trajectory cache)
+    N,
+    K: tl.constexpr,
+    N_ITERS: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_LEADING: tl.constexpr,
+    DTYPE: tl.constexpr,
+    COMPUTE_DTYPE: tl.constexpr,
+):
+    """Run the 1 + 2*(N_ITERS - 1) alternating row/col normalize
+    trajectory for ``BLOCK_LEADING`` rows of the leading axis.
+
+    The internal ``m`` tensor stays in fp32 throughout (matches the
+    eager :func:`sinkhorn_normalize` fp32 contract); the trailing cast
+    to ``DTYPE`` happens at the final store.
+
+    Writes the full FWD state trajectory ``m_0, m_1, ..., m_{2*N_ITERS-1}``
+    to ``STATES_PTR``; the BWD kernel reads it back so it can walk the
+    analytic VJP backward step-by-step without re-running FWD.  At
+    ``K=4, N_ITERS=20`` the cache is ``40 * 16 * 4 = 2560`` bytes per
+    row (~10 MiB total at V4-Flash ``N=4096``) -- negligible HBM
+    overhead vs the 256 KiB FWD output of the same call.
+    """
+
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING)
+    mask_leading = offs < N
+
+    r_offs = tl.arange(0, K)
+    c_offs = tl.arange(0, K)
+
+    KK: tl.constexpr = K * K
+    base = offs[:, None, None] * KK + r_offs[None, :, None] * K + c_offs[None, None, :]
+    full_mask = mask_leading[:, None, None]
+
+    # State buffer layout: [N, N_STATES, K, K] flattened row-major.
+    # state_base[s] = offs * (N_STATES * K * K) + s * (K * K) + r*K + c.
+    n_states_total: tl.constexpr = 2 * N_ITERS
+    state_row_stride: tl.constexpr = n_states_total * KK
+    state_base = offs[:, None, None] * state_row_stride + r_offs[None, :, None] * K + c_offs[None, None, :]
+
+    m = tl.load(X_PTR + base, mask=full_mask, other=0.0).to(COMPUTE_DTYPE)
+    # Save m_0 = x.float() (state index 0).
+    tl.store(STATES_PTR + state_base, m, mask=full_mask)
+
+    # Priming step: col-normalize (sum over the row axis, dim=-2 in the
+    # caller's [...,K,K]; this is axis=1 in our [BLOCK,K,K] layout).
+    s = tl.sum(m, axis=1, keep_dims=True) + EPS
+    m = m / s
+    tl.store(STATES_PTR + state_base + 1 * KK, m, mask=full_mask)  # m_1
+
+    # Alternating loop -- N_ITERS - 1 (row, col) pairs.  After the
+    # priming step we are at state index 1; each loop iteration writes
+    # two more states (the row half then the col half).
+    for it in tl.static_range(N_ITERS - 1):
+        s = tl.sum(m, axis=2, keep_dims=True) + EPS  # row-normalize
+        m = m / s
+        tl.store(STATES_PTR + state_base + (2 + 2 * it) * KK, m, mask=full_mask)
+        s = tl.sum(m, axis=1, keep_dims=True) + EPS  # col-normalize
+        m = m / s
+        tl.store(STATES_PTR + state_base + (3 + 2 * it) * KK, m, mask=full_mask)
+
+    tl.store(Y_PTR + base, m.to(DTYPE), mask=full_mask)
+
+
+@triton.jit
+def _sinkhorn_bwd_kernel(
+    STATES_PTR,  # [N, 2 * N_ITERS, K, K] contiguous (FWD trajectory cache)
+    DY_PTR,  # [N, K, K] contiguous (upstream grad in caller dtype)
+    DX_PTR,  # [N, K, K] contiguous (output)
+    N,
+    K: tl.constexpr,
+    N_ITERS: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_LEADING: tl.constexpr,
+    DTYPE: tl.constexpr,
+    COMPUTE_DTYPE: tl.constexpr,
+):
+    """Apply the analytic VJP for the FWD trajectory.
+
+    Reads the cached FWD states ``m_0, m_1, ..., m_{2*N_ITERS-1}`` from
+    HBM (written by :func:`_sinkhorn_fwd_kernel`) and walks the
+    trajectory backward, applying the per-step closed-form VJP:
+
+    .. code-block:: python
+
+        # forward:  y = x / s, where s = sum(x, axis=axis) + eps
+        # backward: dx = (dy - sum(dy * y, axis=axis, keep_dims=True)) / s
+
+    The HBM round-trip for the cache is ~10 MiB total at V4-Flash
+    widths (``K=4, N=4096``, 40 states / row × 16 fp32 / state); on a
+    3 TB/s HBM device that's ~3 microseconds of read traffic --
+    negligible vs the BWD's actual arithmetic.  This sidesteps a
+    Triton AST-visitor restriction in our toolchain that rejects
+    runtime indexing of Python lists holding ``tl.tensor`` bundles,
+    which would otherwise let us recompute the trajectory in
+    registers.
+
+    Trajectory indexing convention:
+    * ``m_0`` (state 0) = ``x.float()`` (input, BEFORE the priming step);
+    * step ``s`` transforms ``m_{s-1}`` -> ``m_s`` via
+      ``axis = 1`` (col-normalize) when ``s`` is odd, ``axis = 2``
+      (row-normalize) when ``s`` is even;
+    * total steps ``n_steps = 1 + 2*(N_ITERS - 1) = 2*N_ITERS - 1``.
+    """
+
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING)
+    mask_leading = offs < N
+
+    r_offs = tl.arange(0, K)
+    c_offs = tl.arange(0, K)
+
+    KK: tl.constexpr = K * K
+    n_states_total: tl.constexpr = 2 * N_ITERS
+    state_row_stride: tl.constexpr = n_states_total * KK
+
+    base = offs[:, None, None] * KK + r_offs[None, :, None] * K + c_offs[None, None, :]
+    full_mask = mask_leading[:, None, None]
+    state_base = offs[:, None, None] * state_row_stride + r_offs[None, :, None] * K + c_offs[None, None, :]
+
+    dy = tl.load(DY_PTR + base, mask=full_mask, other=0.0).to(COMPUTE_DTYPE)
+    dm = dy
+
+    # Walk the FWD trajectory backward.  Trajectory parity is regular:
+    # step 1 (priming) is col-normalize, then alternating row, col, ...
+    # ending with step ``2*N_ITERS - 1`` (col).  Walking BACKWARD we
+    # always see a (col, row) pair, repeated ``N_ITERS - 1`` times,
+    # followed by the priming col-step.
+    #
+    # We pair the loop iterations explicitly instead of using a
+    # ``step % 2`` runtime check so the axis= argument to ``tl.sum``
+    # stays a compile-time Python int (otherwise the two branches'
+    # ``keep_dims=True`` outputs disagree on shape and Triton refuses
+    # to compile).
+    for i in range(N_ITERS - 1):
+        # Outer (col) step: indices 2*N_ITERS - 1, 2*N_ITERS - 3, ..., 3.
+        col_step = 2 * N_ITERS - 1 - 2 * i
+        m_before_c = tl.load(
+            STATES_PTR + state_base + (col_step - 1) * KK,
+            mask=full_mask,
+            other=0.0,
+        )
+        m_after_c = tl.load(
+            STATES_PTR + state_base + col_step * KK,
+            mask=full_mask,
+            other=0.0,
+        )
+        s_c = tl.sum(m_before_c, axis=1, keep_dims=True) + EPS
+        dot_c = tl.sum(dm * m_after_c, axis=1, keep_dims=True)
+        dm = (dm - dot_c) / s_c
+
+        # Inner (row) step: indices 2*N_ITERS - 2, 2*N_ITERS - 4, ..., 2.
+        row_step = col_step - 1
+        m_before_r = tl.load(
+            STATES_PTR + state_base + (row_step - 1) * KK,
+            mask=full_mask,
+            other=0.0,
+        )
+        m_after_r = tl.load(
+            STATES_PTR + state_base + row_step * KK,
+            mask=full_mask,
+            other=0.0,
+        )
+        s_r = tl.sum(m_before_r, axis=2, keep_dims=True) + EPS
+        dot_r = tl.sum(dm * m_after_r, axis=2, keep_dims=True)
+        dm = (dm - dot_r) / s_r
+
+    # Final priming col-normalize step (step 1: m_0 -> m_1).
+    m_before_p = tl.load(STATES_PTR + state_base, mask=full_mask, other=0.0)
+    m_after_p = tl.load(STATES_PTR + state_base + KK, mask=full_mask, other=0.0)
+    s_p = tl.sum(m_before_p, axis=1, keep_dims=True) + EPS
+    dot_p = tl.sum(dm * m_after_p, axis=1, keep_dims=True)
+    dm = (dm - dot_p) / s_p
+
+    tl.store(DX_PTR + base, dm.to(DTYPE), mask=full_mask)
+
+
+# ---------------------------------------------------------------------------
+# Block-leading heuristic
+# ---------------------------------------------------------------------------
+
+
+def _pick_block_leading(n: int, k: int) -> int:
+    """Pick ``BLOCK_LEADING`` based on K and the work-axis size.
+
+    At K=4 the in-register tensor is ``[BLOCK_LEADING, 4, 4] = 16
+    fp32 / row``; with ~256 VGPRs per warp on MI355 a BLOCK_LEADING of
+    ``128`` keeps live registers + VGPR-spilled state comfortable.
+
+    At K=8 (16x larger per-row footprint) drop to 32.
+
+    At K=16 drop to 8 -- the FWD trajectory keeps ~40 intermediate
+    matrices live during BWD recomputation, so per-program LDS / spill
+    needs to stay bounded.
+    """
+
+    if k <= 4:
+        cap = 128
+    elif k <= 8:
+        cap = 32
+    else:
+        cap = 8
+
+    if n < cap:
+        # Round to next power of 2 >= n (Triton requires power-of-2 block).
+        return max(1, triton.next_power_of_2(n))
+    return cap
+
+
+# ---------------------------------------------------------------------------
+# torch.autograd.Function wrapper
+# ---------------------------------------------------------------------------
+
+
+class SinkhornNormalizeFn(torch.autograd.Function):
+    """Autograd-aware wrapper around the FWD/BWD Triton kernels.
+
+    Saves only the input ``x`` for the backward; ``n_iters`` and
+    ``eps`` are static keys (compiled into the kernel binary cache).
+
+    Shape: any ``[..., K, K]`` where ``K`` is a power of 2 in
+    ``{1, 2, 4, 8, 16}``.  V4-Flash uses ``K=4, n_iters=20, eps=1e-6``.
+    """
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        logits: torch.Tensor,
+        n_iters: int,
+        eps: float,
+    ) -> torch.Tensor:
+        if logits.dim() < 2:
+            raise ValueError(
+                f"sinkhorn_normalize: input must be at least 2-D, got shape {tuple(logits.shape)}"
+            )
+        K = logits.shape[-1]
+        if logits.shape[-2] != K:
+            raise ValueError(
+                f"sinkhorn_normalize: input must be square in the last two dims, "
+                f"got shape {tuple(logits.shape)}"
+            )
+        if K not in _SUPPORTED_K:
+            raise ValueError(f"sinkhorn Triton path: unsupported K={K}; expected one of {_SUPPORTED_K}")
+        if int(n_iters) < 1:
+            raise ValueError(f"n_iters must be >= 1, got {n_iters}")
+
+        x = logits.contiguous()
+        leading_shape = x.shape[:-2]
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        out = torch.empty_like(x)
+        # FWD trajectory cache: 2 * N_ITERS states per row of K*K fp32
+        # elements.  Saved for backward.  At V4-Flash widths (K=4,
+        # N_ITERS=20, N=4096) this is 10 MiB per call -- negligible vs
+        # the 256 KiB FWD output and the ~170 GiB total HBM footprint
+        # of the proxy.
+        n_states_total = 2 * int(n_iters)
+        states_buf = torch.empty(
+            (N, n_states_total, K, K),
+            dtype=torch.float32,
+            device=x.device,
+        )
+
+        block_leading = _pick_block_leading(N, K)
+        grid = (triton.cdiv(N, block_leading),)
+        _sinkhorn_fwd_kernel[grid](
+            x,
+            out,
+            states_buf,
+            N,
+            K=K,
+            N_ITERS=int(n_iters),
+            EPS=float(eps),
+            BLOCK_LEADING=block_leading,
+            DTYPE=_triton_dtype(x.dtype),
+            COMPUTE_DTYPE=tl.float32,
+        )
+
+        ctx.save_for_backward(states_buf)
+        ctx.n_iters = int(n_iters)
+        ctx.eps = float(eps)
+        ctx.K = K
+        ctx.leading_shape = tuple(leading_shape)
+        return out
+
+    @staticmethod
+    def backward(ctx, dy: torch.Tensor):  # type: ignore[override]
+        (states_buf,) = ctx.saved_tensors
+        n_iters: int = ctx.n_iters
+        eps: float = ctx.eps
+        K: int = ctx.K
+        leading_shape = ctx.leading_shape
+
+        dy = dy.contiguous()
+        N = 1
+        for s in leading_shape:
+            N *= s
+
+        dx = torch.empty_like(dy)
+        block_leading = _pick_block_leading(N, K)
+        grid = (triton.cdiv(N, block_leading),)
+        _sinkhorn_bwd_kernel[grid](
+            states_buf,
+            dy,
+            dx,
+            N,
+            K=K,
+            N_ITERS=n_iters,
+            EPS=eps,
+            BLOCK_LEADING=block_leading,
+            DTYPE=_triton_dtype(dy.dtype),
+            COMPUTE_DTYPE=tl.float32,
+        )
+        # n_iters / eps gradients: not differentiable parameters.
+        return dx, None, None
+
+
+# ---------------------------------------------------------------------------
+# Public Python entry points
+# ---------------------------------------------------------------------------
+
+
+def is_triton_path_enabled() -> bool:
+    """Return True iff the ``PRIMUS_SINKHORN_TRITON`` env knob is not ``"0"``.
+
+    Default-on; A/B toggle via ``PRIMUS_SINKHORN_TRITON=0``.
+    """
+
+    return os.environ.get("PRIMUS_SINKHORN_TRITON", "1") != "0"
+
+
+def is_triton_kernel_supported(logits: torch.Tensor) -> bool:
+    """Return True iff the input shape / device is supported by the Triton path.
+
+    Used by the dispatcher in
+    :func:`primus.backends.megatron.core.transformer.hyper_connection.sinkhorn_normalize`
+    to safely fall back to the compiled / eager path when the kernel
+    can't handle the input (non-CUDA, non-square, K out-of-range).
+    """
+
+    if not logits.is_cuda:
+        return False
+    if logits.dim() < 2:
+        return False
+    K = logits.shape[-1]
+    if logits.shape[-2] != K:
+        return False
+    if K not in _SUPPORTED_K:
+        return False
+    return True
+
+
+def eager_sinkhorn_normalize(
+    logits: torch.Tensor,
+    *,
+    n_iters: int = 20,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Reference eager implementation matching
+    :func:`primus.backends.megatron.core.transformer.hyper_connection.sinkhorn_normalize`
+    eager body bit-for-bit (same op order, same fp32 cast contract).
+
+    Kept here so the unit tests / bench can A/B against the canonical
+    eager body without depending on the consumer module.
+    """
+
+    in_dtype = logits.dtype
+    m = logits.float()
+    m = m / (m.sum(dim=-2, keepdim=True) + eps)
+    for _ in range(max(n_iters - 1, 0)):
+        m = m / (m.sum(dim=-1, keepdim=True) + eps)
+        m = m / (m.sum(dim=-2, keepdim=True) + eps)
+    return m.to(in_dtype)
+
+
+def sinkhorn_normalize_triton(
+    logits: torch.Tensor,
+    *,
+    n_iters: int = 20,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Run the Triton-fused Sinkhorn-Knopp normalize.
+
+    Dispatcher: if the env knob is on AND the shape is supported, call
+    the Triton path; else fall back to the eager body in
+    :func:`eager_sinkhorn_normalize`.
+    """
+
+    if is_triton_path_enabled() and is_triton_kernel_supported(logits):
+        return SinkhornNormalizeFn.apply(logits, n_iters, eps)
+    return eager_sinkhorn_normalize(logits, n_iters=n_iters, eps=eps)
+
+
+__all__ = [
+    "SinkhornNormalizeFn",
+    "sinkhorn_normalize_triton",
+    "eager_sinkhorn_normalize",
+    "is_triton_path_enabled",
+    "is_triton_kernel_supported",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py
new file mode 100644
index 000000000..bbedd45f7
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py
@@ -0,0 +1,23 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton **v0** attention backend — DEPRECATED gathered CSA.
+
+The original ``compress_ratio == 4`` CSA path that consumes a pre-gathered
+``[B, Sq, K_topk, head_dim]`` tensor. It is ~30-260x slower than the v1 pool
+path (see ``deepseek-v4/develop/perf/attention_perf.md``) and is NOT used by the
+production dispatch. Kept for reference / correctness tests only.
+
+Prefer :mod:`.._triton_v1` (pool CSA + dense/HCA) or :mod:`.._triton_v2`
+(fused single-latent sparse-MLA).
+"""
+
+from .v4_csa_attention import V4CSAAttentionFn, v4_csa_attention_v0
+
+__all__ = [
+    "v4_csa_attention_v0",
+    "V4CSAAttentionFn",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py
new file mode 100644
index 000000000..5f7d2f853
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py
@@ -0,0 +1,238 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA (gathered) attention — Triton **v0** autograd entry point. DEPRECATED.
+
+The original ``gathered`` CSA path (takes a pre-gathered ``[B, Sq, K, D]``
+tensor). It is ~30-260x slower than the v1 pool path (see attention_perf.md)
+and is NOT used by the production dispatch — retained for reference/tests only.
+Prefer ``_triton_v1`` (pool) or ``_triton_v2`` (fused sparse-MLA).
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (
+    _flydsl_v0_deprecated as _flydsl,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_bwd import (
+    _launch_v4_csa_attention_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_fwd import (
+    _launch_v4_csa_attention_fwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention import (
+    v4_attention_v1,
+)
+
+
+class V4CSAAttentionFn(torch.autograd.Function):
+    """Triton FWD + Triton BWD for the CSA fused attention path."""
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q: torch.Tensor,  # [B, H, Sq, D]
+        k_local: torch.Tensor,  # [B, H, Sq, D]
+        v_local: torch.Tensor,  # [B, H, Sq, D]
+        gathered: torch.Tensor,  # [B, Sq, K_topk, D]
+        sparse_mask: torch.Tensor,  # [B, Sq, K_topk]
+        sink: Optional[torch.Tensor],  # [H] or None
+        swa_window: int,
+        attn_dropout: float,
+        training: bool,
+        scale: float,
+    ) -> torch.Tensor:
+        if attn_dropout > 0.0 and training:
+            # Plan-4 P26 does not implement dropout in the kernel — V4
+            # is trained with attn_dropout=0 so this branch is unreachable
+            # in production. We refuse explicitly so a stray non-zero
+            # dropout configuration raises rather than silently dropping
+            # the kernel path.
+            raise NotImplementedError(
+                "v4_csa_attention_v0 does not implement in-kernel attention "
+                "dropout (V4 trains with attn_dropout=0). Got "
+                f"attn_dropout={attn_dropout}, training={training}."
+            )
+
+        out, lse = _launch_v4_csa_attention_fwd(
+            q,
+            k_local,
+            v_local,
+            gathered,
+            sparse_mask,
+            sink=sink,
+            swa_window=swa_window,
+            scale=scale,
+        )
+        # Save tensors the BWD kernel needs. ``sink`` may be None; we
+        # stash that fact on ``ctx`` because ``save_for_backward`` does
+        # not accept ``None``.
+        ctx.save_for_backward(q, k_local, v_local, gathered, sparse_mask, out, lse, sink)
+        ctx.swa_window = int(swa_window)
+        ctx.attn_dropout = float(attn_dropout)
+        ctx.training_mode = bool(training)
+        ctx.scale = float(scale)
+        ctx.sink_was_none = sink is None
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):  # type: ignore[override]
+        """Triton BWD: re-materialises joint ``P`` from saved LSE; emits all five gradients."""
+        q, k_local, v_local, gathered, sparse_mask, out, lse, sink = ctx.saved_tensors
+
+        sink_arg = None if ctx.sink_was_none else sink
+
+        # Ensure dout is contiguous in the [B, H, Sq, D] layout the
+        # kernel expects.
+        if not grad_out.is_contiguous():
+            grad_out = grad_out.contiguous()
+
+        dq, dk_local, dv_local, dgathered, dsink = _launch_v4_csa_attention_bwd(
+            q,
+            k_local,
+            v_local,
+            gathered,
+            sparse_mask,
+            out,
+            grad_out,
+            lse,
+            sink=sink_arg,
+            swa_window=ctx.swa_window,
+            scale=ctx.scale,
+        )
+
+        # Honor needs_input_grad: zero-cost ``None`` for inputs that
+        # don't want gradients. The kernel still computed them (the
+        # main cost is the matmul, not the per-output cast), so this is
+        # purely cleanliness.
+        if not ctx.needs_input_grad[0]:
+            dq = None
+        if not ctx.needs_input_grad[1]:
+            dk_local = None
+        if not ctx.needs_input_grad[2]:
+            dv_local = None
+        if not ctx.needs_input_grad[3]:
+            dgathered = None
+        # sparse_mask (index 4) is built from the indexer's ``topk_idxs >=
+        # 0`` test — it is NOT a learnable parameter and never needs a
+        # gradient. The kernel does not produce one.
+        if not ctx.needs_input_grad[5] or ctx.sink_was_none:
+            dsink = None
+
+        # Forward signature: (q, k_local, v_local, gathered, sparse_mask,
+        # sink, swa_window, attn_dropout, training, scale).
+        return dq, dk_local, dv_local, dgathered, None, dsink, None, None, None, None
+
+
+def v4_csa_attention_v0(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    gathered: torch.Tensor,  # [B, Sq, K_topk, D]
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    sparse_mask: torch.Tensor,  # [B, Sq, K_topk]
+    attn_dropout: float,
+    training: bool,
+    scale: float,
+    use_tilelang: bool = False,
+    use_flydsl: bool = False,
+) -> torch.Tensor:
+    """Triton-backed V4 CSA fused attention.
+
+    Drop-in replacement for :func:`eager_v4_csa_attention` with
+    identical signature and dtype contract. Routes through
+    :class:`V4CSAAttentionFn` so autograd works.
+
+    When ``gathered.shape[2] == 0`` (degenerate Indexer state — no
+    valid top-K positions) the wrapper short-circuits to the dense
+    :func:`v4_attention_v1` kernel: the local SWA + sink path is exactly
+    what CSA reduces to in that limit, and the dense kernel handles it
+    natively. ``sparse_mask`` is unused on that path so it is allowed
+    to be empty too.
+
+    ``use_tilelang`` is plumbed by ``DeepseekV4Attention.forward``
+    from the ``use_v4_tilelang_csa_attention`` config flag and only
+    triggers a tilelang dispatch when the relevant plan-8 P54 / P55
+    kernels are registered (otherwise the dispatcher warns once and
+    falls back here).
+
+    Returns ``[B, H, Sq, D]`` in ``v_local.dtype``.
+    """
+    K_topk = gathered.shape[2]
+    if K_topk == 0:
+        # Degenerate sparse branch — fall through to the dense kernel.
+        # CSA's local SWA branch matches v4_attention_v1's dense+SWA+sink
+        # path bit-identically when K_topk == 0 (the joint softmax
+        # collapses to the local-only softmax).
+        return v4_attention_v1(
+            q,
+            k_local,
+            v_local,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=None,
+            attn_dropout=attn_dropout,
+            training=training,
+            scale=scale,
+        )
+
+    # Plan-8 P49 / P57 close-out 2: tilelang dispatcher hook for the
+    # CSA family.  Defaults OFF; only fires when the caller passes
+    # ``use_tilelang=True`` (i.e. the config flag is set).
+    if _tilelang.should_dispatch("v4_csa_attention_fwd", enabled=use_tilelang):
+        return _tilelang.v4_csa_attention_fwd_tilelang(
+            q,
+            k_local,
+            v_local,
+            gathered,
+            sparse_mask=sparse_mask,
+            sink=sink,
+            swa_window=swa_window,
+            attn_dropout=attn_dropout,
+            training=training,
+            scale=scale,
+        )
+    # FlyDSL CSA backend hook (forward-only; soft-dep, default off).
+    # Short-circuits on enabled=False; falls back to Triton if the
+    # runtime/kernel is unavailable. Inference/eval path -- training flows
+    # through V4CSAAttentionFn below.
+    if _flydsl.should_dispatch("v4_csa_attention_fwd", enabled=use_flydsl):
+        return _flydsl.v4_csa_attention_fwd_flydsl(
+            q,
+            k_local,
+            v_local,
+            gathered,
+            sparse_mask=sparse_mask,
+            sink=sink,
+            swa_window=swa_window,
+            scale=scale,
+            attn_dropout=attn_dropout,
+            training=training,
+        )
+    return V4CSAAttentionFn.apply(
+        q,
+        k_local,
+        v_local,
+        gathered,
+        sparse_mask,
+        sink,
+        swa_window,
+        attn_dropout,
+        training,
+        scale,
+    )
+
+
+__all__ = [
+    "V4CSAAttentionFn",
+    "v4_csa_attention_v0",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py
new file mode 100644
index 000000000..5e1355fdd
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py
@@ -0,0 +1,488 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention backward Triton kernel (plan-4 P26, ``compress_ratio == 4``).
+
+Two-kernel design (mirroring :mod:`v4_attention_bwd`):
+
+* Pre-pass: ``D[b, h, m] = sum_d (dout[b,h,m,d] * out[b,h,m,d])`` —
+  reuses :func:`_v4_attention_bwd_preprocess_kernel` from the dense
+  module since the contract is identical.
+* Main pass: one program per ``(b, qhid, m)`` query row; re-materialises
+  the joint softmax row from the saved LSE; emits
+
+  ::
+    dq         [B, H, Sq, D]      direct store (one program per row)
+    dk_local   [B, H, Sq, D]      atomic-add (multiple m's hit same n)
+    dv_local   [B, H, Sq, D]      atomic-add
+    dgathered  [B, Sq, K_topk, D] atomic-add (no H dim — broadcast in fwd
+                                  means all H heads contribute)
+    dsink      [H]                atomic-add per query
+
+dtype contract:
+
+* All inputs loaded in input dtype; per-row dot products reduce in fp32
+  via ``.to(tl.float32)`` upcast before the multiply (matches the FWD's
+  bf16-tensor-core / fp32-accumulator semantics).
+* The online ``P / dP / dS`` re-materialisation is fp32 (matches the
+  FWD's softmax-in-fp32 contract).
+* Output gradients are returned in input dtype (cast from fp32 buffers
+  by the launcher).
+
+Math derivation (per query (b, h, m), see plan-4 ``02-phase-details.md``
+Phase 26 section):
+
+  joint_logits = cat(qk_local, qk_sparse, sink_h)
+  P_j = exp(joint_logits[j] - lse)
+  out_d = sum_n P_local[n] * v_local[n,d] + sum_k P_sparse[k] * g[k,d]
+  D = sum_d (dout[d] * out[d])
+  dP_local[n]  = sum_d (dout[d] * v_local[n,d])
+  dP_sparse[k] = sum_d (dout[d] * g[k,d])
+  dS_local[n]  = P_local[n]  * (dP_local[n]  - D)
+  dS_sparse[k] = P_sparse[k] * (dP_sparse[k] - D)
+  dS_sink      = -P_sink * D                                    # sink val is 0
+
+  dq[d]               = sum_n dS_local[n]  * scale * k_local[n,d]
+                      + sum_k dS_sparse[k] * scale * g[k,d]
+  dk_local[n,d]      += dS_local[n]  * scale * q[d]
+  dv_local[n,d]      += P_local[n]   * dout[d]
+  dgathered[k,d]     += dS_sparse[k] * scale * q[d]
+                      + P_sparse[k]  * dout[d]                   # both branches
+  dsink_h            += dS_sink
+
+Edge cases:
+
+* ``K_topk == 0`` — the wrapper short-circuits to the dense
+  :func:`v4_attention_v1` BWD before reaching this kernel.
+* All-masked tile rows — uses ``NEG_INF = -1e30`` finite sentinel so
+  ``exp(NEG_INF - lse) = exp(-large) ≈ 0`` for fully-masked positions
+  (matches the FWD).
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (
+    _v4_attention_bwd_preprocess_kernel,
+)
+
+# ---------------------------------------------------------------------------
+# Main BWD kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_csa_attention_bwd_kernel(
+    Q,
+    K_LOCAL,
+    V_LOCAL,
+    GATHERED,
+    SPARSE_MASK,
+    DOUT,
+    LSE,
+    D,
+    DQ,  # fp32 buffer [B, H, Sq, D]
+    DK_LOCAL,  # fp32 buffer [B, H, Sq, D]
+    DV_LOCAL,  # fp32 buffer [B, H, Sq, D]
+    DGATHERED,  # fp32 buffer [B, Sq, K_topk, D]
+    DSINK,  # fp32 buffer [H] or sentinel
+    SINK,  # [H] or sentinel
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_klb,
+    stride_klh,
+    stride_kln,
+    stride_kld,
+    stride_vlb,
+    stride_vlh,
+    stride_vln,
+    stride_vld,
+    stride_gb,
+    stride_gm,
+    stride_gk,
+    stride_gd,
+    stride_smb,
+    stride_smm,
+    stride_smk,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dklb,
+    stride_dklh,
+    stride_dkln,
+    stride_dkld,
+    stride_dvlb,
+    stride_dvlh,
+    stride_dvln,
+    stride_dvld,
+    stride_dgb,
+    stride_dgm,
+    stride_dgk,
+    stride_dgd,
+    seqlen_q,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """V4 CSA fused-attention BWD (one program per (b, qhid, m) query row)."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    q_active = pid_m < seqlen_q
+
+    # ---- Load Q row, dout row, lse, D scalar ------------------------------
+    q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm
+    q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0)
+
+    do_row_offset = bid * stride_dob + qhid * stride_doh + pid_m * stride_dom
+    dout = tl.load(DOUT + do_row_offset + offs_d * stride_dod, mask=q_active, other=0.0)
+    q_f = q.to(tl.float32)
+    dout.to(tl.float32)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm,
+        mask=q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + qhid * stride_dh + pid_m * stride_dm,
+        mask=q_active,
+        other=0.0,
+    )
+
+    # ---- Sink contribution to dsink ---------------------------------------
+    # dS_sink = -P_sink * D; logit_sink = sink_h, so dsink_h += dS_sink.
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        p_sink = tl.exp(sink_h - lse)
+        # Mask boundary rows so they don't contribute.
+        dsink_contrib = tl.where(q_active, -p_sink * dvec, 0.0)
+        tl.atomic_add(DSINK + qhid, dsink_contrib)
+
+    # dq accumulator (fp32, kept in registers across the n-loop and k-loop)
+    dq = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    # ---- Local SWA branch -------------------------------------------------
+    n_loop_end = pid_m + 1
+    if n_loop_end > seqlen_q:
+        n_loop_end = seqlen_q
+
+    if SWA_WINDOW > 0:
+        n_lo_raw = pid_m - SWA_WINDOW + 1
+        if n_lo_raw < 0:
+            n_lo_raw = 0
+        n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N
+    else:
+        n_loop_start = 0
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        kl_ptrs = (
+            K_LOCAL
+            + bid * stride_klb
+            + qhid * stride_klh
+            + offs_n[:, None] * stride_kln
+            + offs_d[None, :] * stride_kld
+        )
+        kl_load_mask = offs_n[:, None] < seqlen_q
+        kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0)
+
+        vl_ptrs = (
+            V_LOCAL
+            + bid * stride_vlb
+            + qhid * stride_vlh
+            + offs_n[:, None] * stride_vln
+            + offs_d[None, :] * stride_vld
+        )
+        vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0)
+
+        # Re-materialise qk in fp32 (matches FWD).
+        kl_f = kl.to(tl.float32)
+        qk = tl.sum(kl_f * q_f[None, :], axis=1) * sm_scale
+
+        if SWA_WINDOW > 0:
+            in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m)
+        else:
+            in_window = offs_n <= pid_m
+        qk = tl.where(in_window, qk, NEG_INF)
+        qk = tl.where(offs_n < seqlen_q, qk, NEG_INF)
+        # Boundary: off-grid m rows have lse=0 already, but we additionally
+        # zero this whole tile's contribution by forcing qk to NEG_INF.
+        qk = tl.where(q_active, qk, NEG_INF)
+
+        # P = exp(qk - lse)  (joint softmax slice for the local branch)
+        p = tl.exp(qk - lse)
+
+        # dP[n] = sum_d (dout[d] * vl[n, d])
+        dp = tl.sum(dout[None, :].to(tl.float32) * vl.to(tl.float32), axis=1)
+
+        # dS[n] = P[n] * (dP[n] - D)
+        ds = p * (dp - dvec)
+
+        # dq += sum_n (ds[n] * scale * kl[n, d])
+        dq += tl.sum(ds[:, None] * kl.to(tl.float32), axis=0) * sm_scale
+
+        # dk_local[n, d] += ds[n] * scale * q[d]   — atomic-add into fp32 buf
+        dk_contrib = ds[:, None] * sm_scale * q[None, :].to(tl.float32)
+        dk_ptrs = (
+            DK_LOCAL
+            + bid * stride_dklb
+            + qhid * stride_dklh
+            + offs_n[:, None] * stride_dkln
+            + offs_d[None, :] * stride_dkld
+        )
+        tl.atomic_add(dk_ptrs, dk_contrib, mask=kl_load_mask, sem="relaxed")
+
+        # dv_local[n, d] += p[n] * dout[d]   — atomic-add into fp32 buf
+        dv_contrib = p[:, None] * dout[None, :].to(tl.float32)
+        dv_ptrs = (
+            DV_LOCAL
+            + bid * stride_dvlb
+            + qhid * stride_dvlh
+            + offs_n[:, None] * stride_dvln
+            + offs_d[None, :] * stride_dvld
+        )
+        tl.atomic_add(dv_ptrs, dv_contrib, mask=kl_load_mask, sem="relaxed")
+
+    # ---- Sparse branch ----------------------------------------------------
+    for k_start in range(0, K_topk, BLOCK_K):
+        offs_k = k_start + tl.arange(0, BLOCK_K)
+
+        g_ptrs = (
+            GATHERED
+            + bid * stride_gb
+            + pid_m * stride_gm
+            + offs_k[:, None] * stride_gk
+            + offs_d[None, :] * stride_gd
+        )
+        g_load_mask = offs_k[:, None] < K_topk
+        g = tl.load(g_ptrs, mask=g_load_mask, other=0.0)
+
+        sm_ptrs = SPARSE_MASK + bid * stride_smb + pid_m * stride_smm + offs_k * stride_smk
+        sm_load_mask = offs_k < K_topk
+        sm = tl.load(sm_ptrs, mask=sm_load_mask, other=0.0).to(tl.float32)
+
+        qk_sparse = tl.sum(g.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + sm
+        qk_sparse = tl.where(offs_k < K_topk, qk_sparse, NEG_INF)
+        qk_sparse = tl.where(q_active, qk_sparse, NEG_INF)
+
+        p = tl.exp(qk_sparse - lse)
+
+        # dP[k] = sum_d (dout[d] * g[k, d])
+        dp = tl.sum(dout[None, :].to(tl.float32) * g.to(tl.float32), axis=1)
+        ds = p * (dp - dvec)
+
+        # dq += sum_k (ds[k] * scale * g[k, d])
+        dq += tl.sum(ds[:, None] * g.to(tl.float32), axis=0) * sm_scale
+
+        # dgathered[k, d] += ds[k] * scale * q[d] + p[k] * dout[d]
+        # gathered is broadcast across H in the FWD, so this atomic-add
+        # accumulates contributions from every query head — matches the
+        # eager autograd semantics of ``gathered.unsqueeze(1).expand(B, H,
+        # Sq, K, D)``.
+        dg_contrib = ds[:, None] * sm_scale * q[None, :].to(tl.float32) + p[:, None] * dout[None, :].to(
+            tl.float32
+        )
+        dg_ptrs = (
+            DGATHERED
+            + bid * stride_dgb
+            + pid_m * stride_dgm
+            + offs_k[:, None] * stride_dgk
+            + offs_d[None, :] * stride_dgd
+        )
+        tl.atomic_add(dg_ptrs, dg_contrib, mask=g_load_mask, sem="relaxed")
+
+    # ---- Store dq (direct — no collisions across programs) ----------------
+    dq_offset = bid * stride_dqb + qhid * stride_dqh + pid_m * stride_dqm
+    tl.store(DQ + dq_offset + offs_d * stride_dqd, dq, mask=q_active)
+
+
+def _launch_v4_csa_attention_bwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    gathered: torch.Tensor,  # [B, Sq, K_topk, D]
+    sparse_mask: torch.Tensor,  # [B, Sq, K_topk]
+    out: torch.Tensor,  # [B, H, Sq, D] (FWD output)
+    dout: torch.Tensor,  # [B, H, Sq, D]
+    lse: torch.Tensor,  # [B, H, Sq] fp32
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Launch the V4 CSA attention backward kernel.
+
+    Returns ``(dq, dk_local, dv_local, dgathered, dsink)`` — gradients in
+    the input dtype, with ``dsink`` returned only when ``sink is not
+    None`` (else ``None``).
+    """
+    if not q.is_cuda:
+        raise ValueError("v4_csa_attention_v0 BWD requires CUDA / HIP tensors.")
+    if dout.shape != out.shape or out.shape != q.shape:
+        raise ValueError(
+            "v4_csa_attention_v0 BWD shape mismatch: "
+            f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}"
+        )
+
+    B, HQ, Sq, D = q.shape
+    K_topk = gathered.shape[2]
+
+    has_sink = sink is not None
+
+    BLOCK_N = 32
+    BLOCK_K = 32
+    BLOCK_DMODEL = D
+
+    # Allocate fp32 output buffers for atomic_add. Cast to input dtype
+    # before returning.
+    dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+    dk_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+    dv_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+    dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32)
+    if has_sink:
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink
+    else:
+        dsink_fp32 = q  # sentinel; HAS_SINK=False inside kernel
+        sink_arg = q
+
+    # D scalar = (dout * out).sum(-1) — reuse the dense module's pre-pass
+    d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+    pre_grid = (triton.cdiv(Sq, BLOCK_N), B * HQ)
+    _v4_attention_bwd_preprocess_kernel[pre_grid](
+        out,
+        dout,
+        d_buf,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        Sq,
+        HEAD=HQ,
+        BLOCK_M=BLOCK_N,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+
+    grid = (Sq, B * HQ)
+    _v4_csa_attention_bwd_kernel[grid](
+        q,
+        k_local,
+        v_local,
+        gathered,
+        sparse_mask,
+        dout,
+        lse,
+        d_buf,
+        dq_fp32,
+        dk_local_fp32,
+        dv_local_fp32,
+        dgathered_fp32,
+        dsink_fp32,
+        sink_arg,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k_local.stride(0),
+        k_local.stride(1),
+        k_local.stride(2),
+        k_local.stride(3),
+        v_local.stride(0),
+        v_local.stride(1),
+        v_local.stride(2),
+        v_local.stride(3),
+        gathered.stride(0),
+        gathered.stride(1),
+        gathered.stride(2),
+        gathered.stride(3),
+        sparse_mask.stride(0),
+        sparse_mask.stride(1),
+        sparse_mask.stride(2),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        dq_fp32.stride(0),
+        dq_fp32.stride(1),
+        dq_fp32.stride(2),
+        dq_fp32.stride(3),
+        dk_local_fp32.stride(0),
+        dk_local_fp32.stride(1),
+        dk_local_fp32.stride(2),
+        dk_local_fp32.stride(3),
+        dv_local_fp32.stride(0),
+        dv_local_fp32.stride(1),
+        dv_local_fp32.stride(2),
+        dv_local_fp32.stride(3),
+        dgathered_fp32.stride(0),
+        dgathered_fp32.stride(1),
+        dgathered_fp32.stride(2),
+        dgathered_fp32.stride(3),
+        Sq,
+        K_topk,
+        float(scale),
+        HEAD_Q=HQ,
+        SWA_WINDOW=int(swa_window) if swa_window > 0 else 0,
+        HAS_SINK=has_sink,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+
+    dq_out = dq_fp32.to(q.dtype)
+    dk_local_out = dk_local_fp32.to(k_local.dtype)
+    dv_local_out = dv_local_fp32.to(v_local.dtype)
+    dgathered_out = dgathered_fp32.to(gathered.dtype)
+    dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+    return dq_out, dk_local_out, dv_local_out, dgathered_out, dsink_out
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py
new file mode 100644
index 000000000..a1798758a
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py
@@ -0,0 +1,407 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention forward Triton kernel (plan-4 P26, ``compress_ratio == 4``).
+
+CSA fuses three branches into a single online softmax:
+
+* **Local SWA**: ``q @ k_local^T`` with sliding-window-causal masking.
+* **Sparse top-K**: ``q . gathered[m, :, :]`` where the wrapper has
+  pre-gathered ``[B, Sq, K, D]`` rows from the compressed pool (the
+  per-query top-K gather lives outside the kernel — see plan-4
+  ``02-phase-details.md`` Phase 26 design notes).
+* **Per-head learned sink**: a virtual key column with notional value
+  zero, joined as the last softmax candidate so its probability mass
+  is shared across local + sparse branches.
+
+The kernel produces one ``[BLOCK_DMODEL]`` output row per program; the
+grid is ``(seqlen_q, batch * head_q)`` so each program owns exactly one
+``(b, qhid, m)`` query row. The per-row design keeps the sparse-branch
+SMEM footprint inside the MI355 budget at ``head_dim=512``: the
+gathered tile is only ``[BLOCK_K, head_dim] * 2 bytes ≈ 32 KiB`` per
+program, while a multi-row tile would balloon to
+``[BLOCK_M, BLOCK_K, head_dim] * 2 bytes ≈ 1 MiB``.
+
+dtype contract (matches :func:`eager_v4_csa_attention`):
+
+* Q / K / V / gathered are loaded in input dtype (bf16 in production);
+  the per-row dot products (``sum(k * q[None, :], axis=-1)``) reduce in
+  fp32 because we ``.to(tl.float32)`` before the multiply.
+* The online-softmax accumulator (``m_running``, ``l_running``,
+  ``acc``) lives in fp32 — the *only* fp32 step inside the kernel.
+* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD
+  re-materialises ``P`` from it).
+
+Edge cases handled:
+
+* ``K_topk == 0`` — wrapper short-circuits to the dense
+  :func:`v4_attention_v1` kernel before reaching this file.
+* ``topk_idx == -1`` — wrapper sets the corresponding ``sparse_mask``
+  entry to ``-inf``; the kernel just adds the bias and the masked
+  position contributes ~0 to the softmax denominator.
+* All-masked tile rows — the running max and per-tile max are both
+  the finite ``NEG_INF`` sentinel (``-1e30``), so
+  ``exp(NEG_INF - NEG_INF) = exp(0) = 1`` algebraically but the
+  contribution to ``acc`` and ``l_running`` is gated by the
+  per-element ``exp(qk - m_new)`` which stays at exactly zero for
+  every ``-inf``-masked entry. This avoids the ``exp(-inf - -inf) =
+  exp(NaN)`` failure mode that ``-float("inf")`` would have.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Triton kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_csa_attention_fwd_kernel(
+    Q,
+    K_LOCAL,
+    V_LOCAL,
+    GATHERED,
+    SPARSE_MASK,
+    SINK,
+    OUT,
+    LSE,
+    # Q strides: [B, H, Sq, D] row-major (contiguous on D)
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    # K_local strides: [B, H, Sq, D] row-major (CSA always has K_H == HQ —
+    # the V4 forward broadcast-expanded MQA single-latent KV across the H
+    # query heads before this call)
+    stride_klb,
+    stride_klh,
+    stride_kln,
+    stride_kld,
+    # V_local strides: [B, H, Sq, D] row-major
+    stride_vlb,
+    stride_vlh,
+    stride_vln,
+    stride_vld,
+    # gathered strides: [B, Sq, K_topk, D] row-major (no H dim — gather
+    # is per-query but shared across heads)
+    stride_gb,
+    stride_gm,
+    stride_gk,
+    stride_gd,
+    # sparse_mask strides: [B, Sq, K_topk] row-major (broadcasts over H)
+    stride_smb,
+    stride_smm,
+    stride_smk,
+    # OUT strides: [B, H, Sq, D] row-major
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    # LSE strides: [B, H, Sq] row-major
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    seqlen_q,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,  # > 0 for V4; 0 falls back to full causal
+    HAS_SINK: tl.constexpr,
+    BLOCK_N: tl.constexpr,  # local-key tile size
+    BLOCK_K: tl.constexpr,  # sparse-key tile size
+    BLOCK_DMODEL: tl.constexpr,  # head_dim — must be a power of 2
+):
+    """V4 CSA fused-attention FWD (one program per output row)."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    # ---- Load Q row [BLOCK_DMODEL] -----------------------------------------
+    q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm
+    q_ptrs = Q + q_row_offset + offs_d * stride_qd
+    q_active = pid_m < seqlen_q
+    q = tl.load(q_ptrs, mask=q_active, other=0.0)
+
+    # Online-softmax running state (fp32). NEG_INF is a finite sentinel
+    # (-1e30) so all-masked tiles do not produce NaN through
+    # ``exp(-inf - -inf) = exp(NaN)``.
+    NEG_INF: tl.constexpr = -1.0e30
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full((), value=NEG_INF, dtype=tl.float32)
+    l_i = tl.zeros((), dtype=tl.float32)
+
+    # ---- Local SWA branch --------------------------------------------------
+    # Causal: keys n in [0, pid_m]. SWA: keys n in [pid_m - SWA_WINDOW + 1,
+    # pid_m]. We walk from the SWA window's lower bound (rounded down to
+    # BLOCK_N) up to (pid_m + 1). The in-kernel window check inside the
+    # tile loop handles the boundary cases exactly so the result matches
+    # ``sliding_window_causal_mask(...)``.
+    n_loop_end = pid_m + 1
+    if n_loop_end > seqlen_q:
+        n_loop_end = seqlen_q
+
+    # Lower bound of the SWA window (clamped to >= 0). When SWA_WINDOW <= 0
+    # this collapses to a full causal walk from 0.
+    if SWA_WINDOW > 0:
+        n_lo_raw = pid_m - SWA_WINDOW + 1
+        if n_lo_raw < 0:
+            n_lo_raw = 0
+        # Round down to a BLOCK_N multiple so tile-aligned loads stay aligned.
+        n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N
+    else:
+        n_loop_start = 0
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        # K_local tile: [BLOCK_N, BLOCK_DMODEL] in k_local.dtype
+        kl_ptrs = (
+            K_LOCAL
+            + bid * stride_klb
+            + qhid * stride_klh
+            + offs_n[:, None] * stride_kln
+            + offs_d[None, :] * stride_kld
+        )
+        kl_load_mask = offs_n[:, None] < seqlen_q
+        kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0)
+
+        # qk = sum_d (kl[n, d] * q[d]) -> [BLOCK_N], computed in fp32 by
+        # upcasting the operands. (Matches the eager reference's
+        # bf16-tensor-core matmul w/ fp32 accumulator semantics; a 1xD
+        # tl.dot is not portable on the HIP backend, see plan-4 P26 note.)
+        qk = tl.sum(kl.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale
+
+        # SWA-causal mask: keep n in [pid_m - SWA_WINDOW + 1, pid_m].
+        # When SWA_WINDOW <= 0 fall back to full causal.
+        if SWA_WINDOW > 0:
+            in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m)
+        else:
+            in_window = offs_n <= pid_m
+        qk = tl.where(in_window, qk, NEG_INF)
+        # Boundary mask for keys past seqlen_q.
+        qk = tl.where(offs_n < seqlen_q, qk, NEG_INF)
+
+        # Online softmax update (shared with sparse branch + sink).
+        m_tile = tl.max(qk, axis=0)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+
+        # V_local tile: [BLOCK_N, BLOCK_DMODEL] in v_local.dtype
+        vl_ptrs = (
+            V_LOCAL
+            + bid * stride_vlb
+            + qhid * stride_vlh
+            + offs_n[:, None] * stride_vln
+            + offs_d[None, :] * stride_vld
+        )
+        vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0)
+
+        # acc += sum_n (p[n] * vl[n, :]) — fp32 accumulator.
+        acc = acc * alpha + tl.sum(p[:, None] * vl.to(tl.float32), axis=0)
+        m_i = m_new
+
+    # ---- Sparse top-K branch ----------------------------------------------
+    # gathered is per-query (no H dim — broadcast across heads in the
+    # eager reference). We walk K_topk in BLOCK_K tiles.
+    for k_start in range(0, K_topk, BLOCK_K):
+        offs_k = k_start + tl.arange(0, BLOCK_K)
+
+        g_ptrs = (
+            GATHERED
+            + bid * stride_gb
+            + pid_m * stride_gm
+            + offs_k[:, None] * stride_gk
+            + offs_d[None, :] * stride_gd
+        )
+        g_load_mask = offs_k[:, None] < K_topk
+        g = tl.load(g_ptrs, mask=g_load_mask, other=0.0)
+
+        # qk_sparse = sum_d (g[k, d] * q[d]) -> [BLOCK_K]
+        qk_sparse = tl.sum(g.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale
+
+        # Caller-supplied sparse_mask: -inf for topk_idx == -1 entries.
+        sm_ptrs = SPARSE_MASK + bid * stride_smb + pid_m * stride_smm + offs_k * stride_smk
+        sm_load_mask = offs_k < K_topk
+        sm = tl.load(sm_ptrs, mask=sm_load_mask, other=0.0).to(tl.float32)
+        qk_sparse = qk_sparse + sm
+
+        # Boundary mask for offs_k past K_topk.
+        qk_sparse = tl.where(offs_k < K_topk, qk_sparse, NEG_INF)
+
+        # Online softmax update — shares m_i / l_i with the local branch.
+        m_tile = tl.max(qk_sparse, axis=0)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk_sparse - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+
+        # acc += sum_k (p[k] * g[k, :]) — fp32 accumulator.
+        acc = acc * alpha + tl.sum(p[:, None] * g.to(tl.float32), axis=0)
+        m_i = m_new
+
+    # ---- Sink (joint over both branches) ----------------------------------
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        m_new = tl.maximum(m_i, sink_h)
+        alpha = tl.exp(m_i - m_new)
+        beta = tl.exp(sink_h - m_new)
+        l_i = l_i * alpha + beta
+        acc = acc * alpha
+        m_i = m_new
+
+    # ---- Final divide + cast back to output dtype -------------------------
+    out = acc / l_i
+    lse = m_i + tl.log(l_i)
+
+    out_offset = bid * stride_ob + qhid * stride_oh + pid_m * stride_om
+    out_ptrs = OUT + out_offset + offs_d * stride_od
+    tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=q_active)
+
+    lse_ptr = LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm
+    tl.store(lse_ptr, lse, mask=q_active)
+
+
+def _launch_v4_csa_attention_fwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    gathered: torch.Tensor,  # [B, Sq, K_topk, D]
+    sparse_mask: torch.Tensor,  # [B, Sq, K_topk]
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch the V4 CSA attention forward kernel.
+
+    Returns ``(out, lse)`` where ``out`` matches ``v_local.dtype`` and
+    ``lse`` is fp32. ``lse`` is what the BWD kernel needs to
+    re-materialise the joint softmax without storing the
+    ``[Sq, Sq + K_topk]`` joint ``P`` matrix.
+    """
+    if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4:
+        raise ValueError(
+            "v4_csa_attention_v0 forward expects q / k_local / v_local of rank 4 "
+            f"(got {q.dim()} / {k_local.dim()} / {v_local.dim()})"
+        )
+    if gathered.dim() != 4:
+        raise ValueError(
+            f"v4_csa_attention_v0 forward expects gathered of rank 4 [B, Sq, K, D]; "
+            f"got rank {gathered.dim()}, shape {tuple(gathered.shape)}"
+        )
+    if sparse_mask.dim() != 3:
+        raise ValueError(
+            f"v4_csa_attention_v0 forward expects sparse_mask of rank 3 [B, Sq, K]; "
+            f"got rank {sparse_mask.dim()}, shape {tuple(sparse_mask.shape)}"
+        )
+
+    B, HQ, Sq, D = q.shape
+    if k_local.shape != q.shape or v_local.shape != q.shape:
+        raise ValueError(
+            "v4_csa_attention_v0 requires k_local.shape == v_local.shape == q.shape "
+            f"(got q={tuple(q.shape)}, k_local={tuple(k_local.shape)}, "
+            f"v_local={tuple(v_local.shape)})."
+        )
+
+    Bg, Sqg, K_topk, Dg = gathered.shape
+    if Bg != B or Sqg != Sq or Dg != D:
+        raise ValueError(
+            "v4_csa_attention_v0 gathered shape mismatch: expected "
+            f"[B, Sq, K, D] = [{B}, {Sq}, *, {D}]; got {tuple(gathered.shape)}."
+        )
+    Bm, Sqm, Km = sparse_mask.shape
+    if Bm != B or Sqm != Sq or Km != K_topk:
+        raise ValueError(
+            "v4_csa_attention_v0 sparse_mask shape mismatch: expected "
+            f"[B, Sq, K] = [{B}, {Sq}, {K_topk}]; got {tuple(sparse_mask.shape)}."
+        )
+
+    if not q.is_cuda:
+        raise ValueError("v4_csa_attention_v0 requires CUDA / HIP tensors.")
+    if q.dtype != k_local.dtype or q.dtype != v_local.dtype or q.dtype != gathered.dtype:
+        raise ValueError(
+            "v4_csa_attention_v0 requires q.dtype == k_local.dtype == v_local.dtype "
+            f"== gathered.dtype (got {q.dtype} / {k_local.dtype} / "
+            f"{v_local.dtype} / {gathered.dtype})."
+        )
+
+    has_sink = sink is not None
+
+    out = torch.empty_like(q)
+    lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    # Tile sizes: BLOCK_N / BLOCK_K conservative for SMEM at head_dim=512.
+    # Per-row design (one program per (b, qhid, m)) means the gathered
+    # tile is [BLOCK_K, D] only; multi-row would balloon SMEM.
+    BLOCK_N = 32
+    BLOCK_K = 32
+    BLOCK_DMODEL = D  # head_dim must be a power of 2 for tl.arange
+
+    grid = (Sq, B * HQ)
+
+    # Sentinel pointer when sink is absent. Triton requires a real tensor —
+    # we pass q (any tensor) and gate via the constexpr.
+    sink_ptr = sink if has_sink else q
+
+    _v4_csa_attention_fwd_kernel[grid](
+        q,
+        k_local,
+        v_local,
+        gathered,
+        sparse_mask,
+        sink_ptr,
+        out,
+        lse,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k_local.stride(0),
+        k_local.stride(1),
+        k_local.stride(2),
+        k_local.stride(3),
+        v_local.stride(0),
+        v_local.stride(1),
+        v_local.stride(2),
+        v_local.stride(3),
+        gathered.stride(0),
+        gathered.stride(1),
+        gathered.stride(2),
+        gathered.stride(3),
+        sparse_mask.stride(0),
+        sparse_mask.stride(1),
+        sparse_mask.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        Sq,
+        K_topk,
+        float(scale),
+        HEAD_Q=HQ,
+        SWA_WINDOW=int(swa_window) if swa_window > 0 else 0,
+        HAS_SINK=has_sink,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+    return out, lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py
new file mode 100644
index 000000000..dccccc15c
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py
@@ -0,0 +1,29 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Triton **v1** attention backend (production, separate K/V).
+
+The current production Triton attention kernels + autograd Functions:
+
+* dense (cr=0) / HCA (cr=128): :func:`v4_attention_v1` / :class:`V4AttentionFn`
+  (``v4_attention_fwd`` / ``v4_attention_bwd`` launchers).
+* CSA (cr=4), in-kernel pool gather + scatter-add:
+  :func:`v4_csa_attention_v1` / :class:`V4CSAPoolAttentionFn`
+  (``v4_csa_attention_fwd`` / ``v4_csa_attention_bwd`` pool launchers).
+
+See ``_triton_v0_deprecated`` for the deprecated gathered CSA path and ``_triton_v2`` for
+the fused single-latent sparse-MLA path.
+"""
+
+from .v4_attention import V4AttentionFn, v4_attention_v1
+from .v4_csa_attention import V4CSAPoolAttentionFn, v4_csa_attention_v1
+
+__all__ = [
+    "v4_attention_v1",
+    "V4AttentionFn",
+    "v4_csa_attention_v1",
+    "V4CSAPoolAttentionFn",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py
new file mode 100644
index 000000000..be72fad73
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py
@@ -0,0 +1,65 @@
+# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+"""Architecture-aware tuned defaults for the V4 Triton attention kernels.
+
+The V4 attention fwd/bwd/CSA kernels expose ~60 ``PRIMUS_V4_ATTN_*`` / ``PRIMUS_V4_CSA_*``
+env knobs whose hard-coded defaults are the **gfx950 / MI355X** R1/R2 sweep winners. On
+gfx1250 (MI450, CDNA-next) the optimum differs systematically -- larger ``BLOCK_M`` and
+fewer warps/stages, and far more dKV head-split -- so a gfx950-tuned default leaves a lot on
+the table. This module centralises the per-arch defaults so each launcher can ask for the
+right value for the GPU it is running on; the env knobs still override everything.
+
+gfx1250 values come from node-safe microbench sweeps (``ab_sweep/opt7b_*`` / ``opt7c_*``);
+untuned knobs/archs fall back to the historical gfx950 defaults so nothing regresses.
+"""
+from __future__ import annotations
+
+import functools
+
+
+@functools.lru_cache(maxsize=1)
+def gpu_arch() -> str:
+    """Lower-cased GPU gfx arch (e.g. ``gfx1250``); ``""`` if it can't be determined."""
+    try:
+        import torch
+
+        name = torch.cuda.get_device_properties(0).gcnArchName  # e.g. "gfx1250:sramecc+:xnack-"
+        return name.split(":")[0].strip().lower()
+    except Exception:
+        return ""
+
+
+def is_gfx1250() -> bool:
+    return gpu_arch() == "gfx1250"
+
+
+def fwd_attn_defaults(is_hca: bool):
+    """``(BLOCK_M, BLOCK_N, NUM_WARPS, NUM_STAGES)`` for the V4 attention FWD kernel.
+
+    gfx1250: ``BM=128, BN=32, W=4, S=1`` wins **both** shapes -- SWA **+62-70%** vs the gfx950
+    winner ``BM=64/BN=16/W=8/S=2`` (``ab_sweep/opt7c``), and HCA (cr=128) **+15-21%** vs the
+    gfx950 HCA winner ``BM=128/BN=16/W=8/S=1`` (``ab_sweep/opt7d``).
+    """
+    if is_gfx1250():
+        return 128, 32, 4, 1
+    # gfx950 / MI355X -- historical R2-sweep winners (shape-dependent BLOCK_M / stages).
+    if is_hca:
+        return 128, 16, 8, 1
+    return 64, 16, 8, 2
+
+
+def bwd_dkv_head_groups_default(hq: int, hk: int) -> int:
+    """dKV head-split groups for the V4 attention BWD -- ONLY the MQA (HQ>=64, HK==1) path.
+
+    gfx1250 sweep (``ab_sweep/opt7b``, MQA): ``HG=32`` is +37-53% vs the gfx950 default ``2``
+    (the kernel notes "HG=4/8 regress" -- true on gfx950, inverted on gfx1250).
+
+    NOTE: the current DeepSeek-V4-Pro attention runs **MHA at the kernel level** -- the
+    single-latent KV is expanded to all H heads (deepseek_v4_attention.py), so HK==HQ and the
+    caller's ``if HQ > HK`` guard means this function is **not reached** there (HG stays 1).
+    The MHA bwd is already gfx950-optimal on gfx1250 (``ab_sweep/opt7e`` -- every alt config
+    regresses), so there is no bwd retune for V4-Pro; this default only matters for a true-MQA
+    config. The real gfx1250 attention win is the FWD (see ``fwd_attn_defaults``).
+    """
+    if not (hq >= 64 and hk == 1):
+        return 1
+    return 32 if is_gfx1250() else 2
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py
new file mode 100644
index 000000000..3acbc6493
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py
@@ -0,0 +1,298 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention Triton autograd entry point (plan-4 P25).
+
+Public API:
+
+* :func:`v4_attention_v1` — functional API matching
+  :func:`eager_v4_attention`'s signature; routes through
+  :class:`V4AttentionFn` so autograd works.
+* :class:`V4AttentionFn` — :class:`torch.autograd.Function` wrapping
+  the Triton FWD + Triton BWD (re-materialises softmax from the saved
+  LSE; sink gradient atomic-added per query head).
+
+Dispatch contract (consumed by ``DeepseekV4Attention.forward``):
+
+* ``compress_ratio == 0``: caller passes ``swa_window > 0`` and
+  ``additive_mask=None`` so the kernel applies the SWA-causal mask
+  in-place.
+* ``compress_ratio == 128`` (HCA): caller pre-concatenates pool keys
+  to local keys, passes the pool-only ``[Sq, P]`` additive mask, sets
+  ``hca_local_seqlen=Sq``, and keeps ``swa_window > 0``. The kernel
+  splits the loop into a pruned local SWA branch plus the compressed-pool
+  suffix.
+
+dtype contract (must match :func:`eager_v4_attention`):
+
+* Q / K / V matmuls run on tensor cores in input dtype (bf16 in
+  production); the matmul accumulator inside is fp32.
+* The online-softmax accumulator (``m_running / l_running / acc``) is
+  fp32 — the *only* fp32 step.
+* Output is in ``v.dtype``.
+* Saved LSE is fp32 (BWD re-materialises ``P`` from it).
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (
+    _flydsl_v0_deprecated as _flydsl,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (
+    _launch_v4_attention_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import (
+    _launch_v4_attention_fwd,
+)
+
+# ---------------------------------------------------------------------------
+# Autograd Function
+# ---------------------------------------------------------------------------
+
+
+class V4AttentionFn(torch.autograd.Function):
+    """Triton FWD + (eager-recompute BWD until P25-stage-2 lands)."""
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q: torch.Tensor,  # [B, H, Sq, D]
+        k: torch.Tensor,  # [B, K_H, Sk, D]
+        v: torch.Tensor,  # [B, K_H, Sk, D]
+        sink: Optional[torch.Tensor],  # [H] or None
+        additive_mask: Optional[torch.Tensor],  # [Sq, Sk] or None
+        swa_window: int,
+        attn_dropout: float,
+        training: bool,
+        scale: float,
+        hca_local_seqlen: int,
+    ) -> torch.Tensor:
+        if attn_dropout > 0.0 and training:
+            # Plan-4 P25 does not implement dropout in the kernel — V4
+            # is trained with attn_dropout=0 so this branch is unreachable
+            # in production. We refuse explicitly so a stray non-zero
+            # dropout configuration raises rather than silently dropping
+            # the kernel path.
+            raise NotImplementedError(
+                "v4_attention_v1 does not implement in-kernel attention "
+                "dropout (V4 trains with attn_dropout=0). Got "
+                f"attn_dropout={attn_dropout}, training={training}."
+            )
+
+        out, lse = _launch_v4_attention_fwd(
+            q,
+            k,
+            v,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=additive_mask,
+            scale=scale,
+            hca_local_seqlen=hca_local_seqlen,
+        )
+        # Save tensors the BWD kernel needs. None-typed args
+        # (``sink`` / ``additive_mask``) are stashed on ``ctx`` because
+        # ``save_for_backward`` does not accept ``None``.
+        ctx.save_for_backward(q, k, v, out, lse, sink, additive_mask)
+        ctx.swa_window = int(swa_window)
+        ctx.attn_dropout = float(attn_dropout)
+        ctx.training_mode = bool(training)
+        ctx.scale = float(scale)
+        ctx.hca_local_seqlen = int(hca_local_seqlen)
+        ctx.sink_was_none = sink is None
+        ctx.mask_was_none = additive_mask is None
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):  # type: ignore[override]
+        """Triton BWD: re-materialises ``P`` from saved LSE; emits dq, dk, dv, dsink.
+
+        Routes through :func:`_launch_v4_attention_bwd`, which:
+
+        * runs a small fp32 pre-pass that computes the per-query
+          ``D = (dout * out).sum(-1)`` scalar,
+        * runs the main BWD kernel parallelized over query blocks,
+          accumulating ``dQ`` in registers and atomic-adding into
+          ``dK / dV / dsink``,
+        * casts the fp32 accumulator buffers back to the input dtype
+          for return.
+
+        The Function returns gradients in the same positional order as
+        the forward: ``(q, k, v, sink, additive_mask, swa_window,
+        attn_dropout, training, scale)``. Non-tensor / non-grad inputs
+        get ``None``.
+        """
+        q, k, v, out, lse, sink, additive_mask = ctx.saved_tensors
+
+        sink_arg = None if ctx.sink_was_none else sink
+        mask_arg = None if ctx.mask_was_none else additive_mask
+
+        # Ensure dout is contiguous in the [B, H, Sq, D] layout the
+        # kernel expects. The kernel reads strides from the tensor so
+        # any contiguous-in-its-strides tensor would work, but a
+        # ``.contiguous()`` here keeps the access pattern simple.
+        if not grad_out.is_contiguous():
+            grad_out = grad_out.contiguous()
+
+        dq, dk, dv, dsink = _launch_v4_attention_bwd(
+            q,
+            k,
+            v,
+            out,
+            grad_out,
+            lse,
+            sink=sink_arg,
+            swa_window=ctx.swa_window,
+            additive_mask=mask_arg,
+            scale=ctx.scale,
+            hca_local_seqlen=ctx.hca_local_seqlen,
+        )
+
+        # Honor needs_input_grad: zero-cost ``None`` for inputs that
+        # don't want gradients. (We still computed them — the main
+        # cost in the kernel is the matmul, not the per-output cast —
+        # so this is purely a cleanliness touch.)
+        if not ctx.needs_input_grad[0]:
+            dq = None
+        if not ctx.needs_input_grad[1]:
+            dk = None
+        if not ctx.needs_input_grad[2]:
+            dv = None
+        if not ctx.needs_input_grad[3] or ctx.sink_was_none:
+            dsink = None
+
+        # Forward signature: (q, k, v, sink, additive_mask, swa_window,
+        # attn_dropout, training, scale, hca_local_seqlen)
+        return dq, dk, dv, dsink, None, None, None, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Functional API
+# ---------------------------------------------------------------------------
+
+
+def v4_attention_v1(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k: torch.Tensor,  # [B, K_H, Sk, D]   K_H ∈ {1, H}
+    v: torch.Tensor,  # [B, K_H, Sk, D]
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],  # [Sq, Sk] or None
+    attn_dropout: float,
+    training: bool,
+    scale: float,
+    hca_local_seqlen: int = 0,
+    use_tilelang: bool = False,
+    use_flydsl: bool = False,
+) -> torch.Tensor:
+    """Triton- or tilelang-backed V4 dense / HCA attention.
+
+    Drop-in replacement for :func:`eager_v4_attention` with identical
+    signature and dtype contract.
+
+    Dispatch precedence (plan-8 P57 close-out 2):
+
+    * If ``use_tilelang=True`` AND the plan-8 P50 / P51 tilelang
+      kernels are registered + tilelang importable at the pinned
+      version, route through the tilelang FWD/BWD.
+    * Otherwise, route through :class:`V4AttentionFn` (plan-4 P25
+      Triton FWD + plan-5 P32 final split BWD).
+
+    ``use_tilelang`` is plumbed by ``DeepseekV4Attention.forward``
+    from the ``use_v4_tilelang_attention`` config flag.  Default-False
+    so containers without tilelang installed never trigger any
+    tilelang import.
+
+    The MQA case (``K_H == 1``) is detected from ``k.shape[1]`` and
+    each kernel internally broadcasts the single shared K / V head
+    across the query heads.
+
+    Returns ``[B, H, Sq, D]`` in ``v.dtype``.
+    """
+    # Plan-8 P49 / P57 close-out 2: tilelang dispatcher hook.
+    # ``should_dispatch`` short-circuits on ``enabled=False`` so the
+    # tilelang import never fires on a container that does not have
+    # the package installed.
+    if _tilelang.should_dispatch("v4_attention_fwd", enabled=use_tilelang):
+        # P51: route through the tilelang autograd Function when both
+        # FWD + BWD kernels are registered, so gradients flow through
+        # the tilelang path correctly.  When only the FWD is registered
+        # (P50-only state), fall through to the FWD-direct call — the
+        # Triton BWD will still get called via the autograd graph that
+        # the caller sets up (only the FWD wrapper is non-autograd in
+        # that mode).
+        # `_lazy_load` is idempotent + cheap on cache hit; calling it
+        # here ensures `is_tilelang_kernel_available("v4_attention_bwd")`
+        # reflects the post-import state.
+        _tilelang._lazy_load("v4_attention_bwd")
+        if _tilelang.is_tilelang_kernel_available("v4_attention_bwd"):
+            from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_autograd_tilelang import (
+                v4_attention_tilelang,
+            )
+
+            return v4_attention_tilelang(
+                q,
+                k,
+                v,
+                sink=sink,
+                additive_mask=additive_mask,
+                swa_window=swa_window,
+                attn_dropout=attn_dropout,
+                training=training,
+                scale=scale,
+                hca_local_seqlen=hca_local_seqlen,
+            )
+        return _tilelang.v4_attention_fwd_tilelang(
+            q,
+            k,
+            v,
+            sink=sink,
+            additive_mask=additive_mask,
+            swa_window=swa_window,
+            attn_dropout=attn_dropout,
+            training=training,
+            scale=scale,
+            hca_local_seqlen=hca_local_seqlen,
+        )
+    # FlyDSL backend hook (forward-only; soft-dep, default off). Mirrors the
+    # _tilelang hook: short-circuits on enabled=False so FlyDSL is never
+    # imported on the common path, and falls back to Triton if the runtime
+    # or the kernel is unavailable. Training (autograd) still flows through
+    # V4AttentionFn below; this path is for inference/eval.
+    if _flydsl.should_dispatch("v4_attention_fwd", enabled=use_flydsl):
+        return _flydsl.v4_attention_fwd_flydsl(
+            q,
+            k,
+            v,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=additive_mask,
+            scale=scale,
+            hca_local_seqlen=hca_local_seqlen,
+        )
+    return V4AttentionFn.apply(
+        q,
+        k,
+        v,
+        sink,
+        additive_mask,
+        swa_window,
+        attn_dropout,
+        training,
+        scale,
+        hca_local_seqlen,
+    )
+
+
+__all__ = [
+    "V4AttentionFn",
+    "v4_attention_v1",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py
new file mode 100644
index 000000000..ff8072ef9
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py
@@ -0,0 +1,2086 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention backward Triton kernel (plan-4 P25, ``compress_ratio in {0, 128}``).
+
+Two-kernel design:
+
+* :func:`_v4_attention_bwd_preprocess_kernel` — computes the per-query
+  diagonal scalar ``D = sum_d (dout[..., d] * out[..., d])`` in fp32.
+  This is the standard FlashAttention-2 BWD pre-pass; it lets the main
+  kernel skip materialising ``dP @ V^T`` as a tile-wise sum.
+* :func:`_v4_attention_bwd_kernel` — main BWD pass. Parallelizes over
+  query blocks (one program per ``m_block × batch × head_q``) and:
+
+    1. re-materialises ``P = exp(qk - LSE)`` in fp32 from saved
+       ``Q / K / V / LSE`` (so the [Sq, Sk] ``P`` matrix is never
+       stored),
+    2. computes ``dS = P * (dP - D)`` where ``dP = dout @ V^T``,
+    3. accumulates ``dQ`` in registers (no atomic — one program per
+       m-block writes its dQ tile straight to global),
+    4. atomic-adds ``dK = scale * dS^T @ Q`` and ``dV = P^T @ dout``
+       into the global dK / dV buffers,
+    5. atomic-adds the sink gradient
+       ``dsink_h += -sum_t (P_sink_t * D_t)`` per query head.
+
+The launcher allocates ``dQ / dK / dV / dsink`` as fp32 buffers (so
+``tl.atomic_add`` works regardless of input dtype) and casts back to
+the input dtype on return.
+
+dtype contract:
+
+* All matmuls run on tensor cores in input dtype with fp32 accumulator.
+* The online ``P / dP / dS`` re-materialisation is fp32 (matches the
+  FWD's softmax-in-fp32 contract).
+* Output gradients are returned in input dtype (cast from fp32
+  buffers).
+
+Tile choice: ``BLOCK_M = BLOCK_N = 32`` at ``head_dim = 512`` so the
+peak SMEM (Q + K + V + dout = 4 × 32 × 512 × 2 = 128 KiB) fits under
+MI355's 160 KiB SMEM budget with some headroom for register pressure.
+P25 perf follow-up may explore ``BLOCK_M = BLOCK_N = 64`` once the
+correctness baseline is locked.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Pre-pass kernel (D scalar)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_bwd_preprocess_kernel(
+    OUT,
+    DOUT,
+    D,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    seqlen_q,
+    HEAD: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """Compute ``D[b, h, m] = sum_d (dout[b, h, m, d] * out[b, h, m, d])``."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD
+    hid = pid_bh % HEAD
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    out_ptrs = (
+        OUT + bid * stride_ob + hid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
+    )
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + hid * stride_doh
+        + offs_m[:, None] * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    mask = offs_m[:, None] < seqlen_q
+    o = tl.load(out_ptrs, mask=mask, other=0.0).to(tl.float32)
+    do = tl.load(dout_ptrs, mask=mask, other=0.0).to(tl.float32)
+    d = tl.sum(o * do, 1)
+
+    d_ptrs = D + bid * stride_db + hid * stride_dh + offs_m * stride_dm
+    tl.store(d_ptrs, d, mask=offs_m < seqlen_q)
+
+
+# ---------------------------------------------------------------------------
+# Main BWD kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_bwd_kernel(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DQ,  # fp32 buffer [B, H, Sq, D]
+    DK,  # fp32 buffer [B, K_H, Sk, D]
+    DV,  # fp32 buffer [B, K_H, Sk, D]
+    DSINK,  # fp32 buffer [H] or sentinel
+    SINK,  # [H] fp32 or sentinel
+    ADD_MASK,  # [Sq, Sk] or sentinel
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dkd,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    HAS_ADD_MASK: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    USE_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """V4 attention BWD (single kernel, parallelize over m-blocks)."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+    if HEAD_K == HEAD_Q:
+        khid = qhid
+    else:
+        khid = 0
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    # Load Q, dout, LSE, D for this m-block.
+    q_ptrs = (
+        Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+    )
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + qhid * stride_doh
+        + offs_m[:, None] * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+    dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+    q_load_mask = offs_m[:, None] < seqlen_q
+    q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+    dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+    lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+    dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+    # dQ accumulator (fp32, kept in registers across the n-loop)
+    dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    # Sink gradient — accumulate per-query contribution then atomic_add
+    # into DSINK[qhid] at the end. The sink contributes
+    # ``dsink_h += sum_t -P_sink_t * D_t`` because dlogits_sink = P_sink * (0 - D).
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        # P_sink for each query in the m-block, masked for boundary rows
+        p_sink = tl.exp(sink_h - lse)
+        p_sink_masked = tl.where(offs_m < seqlen_q, p_sink, 0.0)
+        dvec_masked = tl.where(offs_m < seqlen_q, dvec, 0.0)
+        dsink_contrib = tl.sum(-p_sink_masked * dvec_masked)
+        tl.atomic_add(DSINK + qhid, dsink_contrib)
+
+    # Determine n-loop bounds (matches FWD's P30 SWA tile pruning).
+    n_loop_start = 0
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+        n_loop_end = seqlen_k
+    elif SWA_WINDOW > 0:
+        n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1
+        if n_loop_start < 0:
+            n_loop_start = 0
+        n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k
+        if n_loop_end > local_end:
+            n_loop_end = local_end
+    elif USE_CAUSAL:
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        if n_loop_end > seqlen_k:
+            n_loop_end = seqlen_k
+    else:
+        n_loop_end = seqlen_k
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        k_ptrs = (
+            K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+        )
+        v_ptrs = (
+            V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+        )
+        kv_load_mask = offs_n[:, None] < seqlen_k
+        k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+        v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+
+        # qk = Q @ K.T * scale + mask (re-materialise in fp32)
+        qk = tl.dot(q, tl.trans(k)) * sm_scale
+        if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+        else:
+            if SWA_WINDOW > 0:
+                in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & (
+                    offs_n[None, :] <= offs_m[:, None]
+                )
+                qk = tl.where(in_window, qk, NEG_INF)
+            elif USE_CAUSAL:
+                qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF)
+        qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+        qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+        # P = exp(qk - LSE) in fp32. For boundary rows / fully-masked
+        # rows, lse is 0 (loaded with mask) and qk is NEG_INF, so
+        # exp(NEG_INF - 0) = 0 — no contribution. ✓
+        p = tl.exp(qk - lse[:, None])
+
+        # dP = dout @ V.T (fp32 accumulator)
+        dp = tl.dot(dout, tl.trans(v))
+
+        # dS = P * (dP - D)
+        ds = p * (dp - dvec[:, None])
+
+        # dQ += dS @ K * scale
+        dq += tl.dot(ds.to(k.dtype), k) * sm_scale
+
+        # dK += scale * dS.T @ Q  (atomic_add into fp32 DK buffer)
+        dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale
+        dk_ptrs = (
+            DK
+            + bid * stride_dkb
+            + khid * stride_dkh
+            + offs_n[:, None] * stride_dkn
+            + offs_d[None, :] * stride_dkd
+        )
+        dk_mask = offs_n[:, None] < seqlen_k
+        tl.atomic_add(dk_ptrs, dk_contrib, mask=dk_mask, sem="relaxed")
+
+        # dV += P.T @ dout  (atomic_add into fp32 DV buffer)
+        dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout)
+        dv_ptrs = (
+            DV
+            + bid * stride_dvb
+            + khid * stride_dvh
+            + offs_n[:, None] * stride_dvn
+            + offs_d[None, :] * stride_dvd
+        )
+        dv_mask = offs_n[:, None] < seqlen_k
+        tl.atomic_add(dv_ptrs, dv_contrib, mask=dv_mask, sem="relaxed")
+
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0:
+        for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N):
+            offs_n = n_start + tl.arange(0, BLOCK_N)
+            pool_n = offs_n - HCA_LOCAL_SEQLEN
+
+            k_ptrs = (
+                K
+                + bid * stride_kb
+                + khid * stride_kh
+                + offs_n[:, None] * stride_kn
+                + offs_d[None, :] * stride_kd
+            )
+            v_ptrs = (
+                V
+                + bid * stride_vb
+                + khid * stride_vh
+                + offs_n[:, None] * stride_vn
+                + offs_d[None, :] * stride_vd
+            )
+            kv_load_mask = offs_n[:, None] < seqlen_k
+            k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+            v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+
+            qk = tl.dot(q, tl.trans(k)) * sm_scale
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+            qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+            qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+            p = tl.exp(qk - lse[:, None])
+            dp = tl.dot(dout, tl.trans(v))
+            ds = p * (dp - dvec[:, None])
+
+            dq += tl.dot(ds.to(k.dtype), k) * sm_scale
+
+            dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale
+            dk_ptrs = (
+                DK
+                + bid * stride_dkb
+                + khid * stride_dkh
+                + offs_n[:, None] * stride_dkn
+                + offs_d[None, :] * stride_dkd
+            )
+            dk_mask = offs_n[:, None] < seqlen_k
+            tl.atomic_add(dk_ptrs, dk_contrib, mask=dk_mask, sem="relaxed")
+
+            dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout)
+            dv_ptrs = (
+                DV
+                + bid * stride_dvb
+                + khid * stride_dvh
+                + offs_n[:, None] * stride_dvn
+                + offs_d[None, :] * stride_dvd
+            )
+            dv_mask = offs_n[:, None] < seqlen_k
+            tl.atomic_add(dv_ptrs, dv_contrib, mask=dv_mask, sem="relaxed")
+
+    # Store dQ (fp32 buffer; launcher casts back to input dtype on return).
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + qhid * stride_dqh
+        + offs_m[:, None] * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    tl.store(dq_ptrs, dq, mask=offs_m[:, None] < seqlen_q)
+
+
+# ---------------------------------------------------------------------------
+# Plan-5 P32 split-kernel BWD — dQ kernel + dK/dV kernel (no atomics)
+#
+# The monolithic ``_v4_attention_bwd_kernel`` parallelises over m-blocks
+# and ``tl.atomic_add``s into ``dK / dV``. With ``H=64`` heads ×
+# ``SWA_WINDOW / BLOCK_M = 4`` m-blocks per K position, every K position
+# is touched by ~256 concurrent atomics; even at "relaxed" semantics on
+# MI355 these serialise through the L2 atomic engine.
+#
+# P32 splits the BWD into two kernels that each write their own output
+# slice with NO atomics:
+#
+#  * ``_v4_attention_bwd_dq_kernel`` — one program per ``(b, qhid, m_block)``;
+#    accumulates ``dQ`` in registers (same as the monolithic), but also
+#    handles the per-head sink gradient (``dsink``) which only needs the
+#    saved ``LSE`` and ``D``. Writes ``dQ`` straight to global; ``dsink``
+#    is atomic-added (single counter per head — contention there is
+#    negligible).
+#  * ``_v4_attention_bwd_dkv_kernel`` — one program per
+#    ``(b, khid, n_block)``; iterates the *m*-blocks that contribute to
+#    this n-block (pruned by SWA / causal bounds), accumulates ``dK`` /
+#    ``dV`` in registers and writes straight to global. No atomics.
+#
+# Total compute roughly doubles (``Q @ K.T`` re-materialised once per
+# kernel instead of shared) but the atomic pressure goes to zero, which
+# is the dominant cost on MI355 at ``H=64`` × ``BLOCK_M=BLOCK_N=32``.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_bwd_dq_kernel(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    DSINK,
+    SINK,
+    ADD_MASK,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    HAS_ADD_MASK: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    USE_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACCUMULATE: tl.constexpr = False,
+    EXACT_TILES_M: tl.constexpr = False,
+    EXACT_TILES_N: tl.constexpr = False,
+):
+    """V4 attention BWD — dQ only (parallel over m-blocks, no atomics for dQ).
+
+    When ``ACCUMULATE`` is True, the kernel performs ``dQ += dq`` instead
+    of ``dQ = dq``. Each ``(b, qhid, m_block)`` program owns a unique
+    output slice, so the implicit read-modify-write is race-free across
+    programs within a single launch, and sequential launches of this
+    kernel against the same buffer are also race-free.
+    """
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+    if HEAD_K == HEAD_Q:
+        khid = qhid
+    else:
+        khid = 0
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    q_ptrs = (
+        Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+    )
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + qhid * stride_doh
+        + offs_m[:, None] * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+    dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+    q_load_mask = offs_m[:, None] < seqlen_q
+    q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+    dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+    lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+    dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+    dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        p_sink = tl.exp(sink_h - lse)
+        p_sink_masked = tl.where(offs_m < seqlen_q, p_sink, 0.0)
+        dvec_masked = tl.where(offs_m < seqlen_q, dvec, 0.0)
+        dsink_contrib = tl.sum(-p_sink_masked * dvec_masked)
+        tl.atomic_add(DSINK + qhid, dsink_contrib)
+
+    n_loop_start = 0
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+        n_loop_end = seqlen_k
+    elif SWA_WINDOW > 0:
+        n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1
+        if n_loop_start < 0:
+            n_loop_start = 0
+        n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k
+        if n_loop_end > local_end:
+            n_loop_end = local_end
+    elif USE_CAUSAL:
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        if n_loop_end > seqlen_k:
+            n_loop_end = seqlen_k
+    else:
+        n_loop_end = seqlen_k
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        k_ptrs = (
+            K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+        )
+        v_ptrs = (
+            V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+        )
+        kv_load_mask = offs_n[:, None] < seqlen_k
+        k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+        v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+
+        qk = tl.dot(q, tl.trans(k)) * sm_scale
+        if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+        else:
+            if SWA_WINDOW > 0:
+                in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & (
+                    offs_n[None, :] <= offs_m[:, None]
+                )
+                qk = tl.where(in_window, qk, NEG_INF)
+            elif USE_CAUSAL:
+                qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF)
+        # Plan-8 P57: EXACT_TILES_* skip the boundary masks when the
+        # launcher confirms ``seqlen_q % BLOCK_M == 0`` (resp. seqlen_k).
+        if not EXACT_TILES_N:
+            qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+        if not EXACT_TILES_M:
+            qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout, tl.trans(v))
+        ds = p * (dp - dvec[:, None])
+        # P57 cr=0 BWD: defer ``sm_scale`` to a single multiply after
+        # the n-loop, and fuse the inner ``dq += dot(ds, k)`` into an
+        # MFMA-acc form via ``tl.dot(..., acc=dq)``.  Numerics are
+        # bit-identical modulo fp32 associativity.
+        dq = tl.dot(ds.to(k.dtype), k, acc=dq)
+
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0:
+        for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N):
+            offs_n = n_start + tl.arange(0, BLOCK_N)
+            pool_n = offs_n - HCA_LOCAL_SEQLEN
+
+            k_ptrs = (
+                K
+                + bid * stride_kb
+                + khid * stride_kh
+                + offs_n[:, None] * stride_kn
+                + offs_d[None, :] * stride_kd
+            )
+            v_ptrs = (
+                V
+                + bid * stride_vb
+                + khid * stride_vh
+                + offs_n[:, None] * stride_vn
+                + offs_d[None, :] * stride_vd
+            )
+            kv_load_mask = offs_n[:, None] < seqlen_k
+            k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+            v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+
+            qk = tl.dot(q, tl.trans(k)) * sm_scale
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+            if not EXACT_TILES_N:
+                qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+            if not EXACT_TILES_M:
+                qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+            p = tl.exp(qk - lse[:, None])
+            dp = tl.dot(dout, tl.trans(v))
+            ds = p * (dp - dvec[:, None])
+            dq = tl.dot(ds.to(k.dtype), k, acc=dq)
+
+    # P57 cr=0 BWD: fold ``sm_scale`` once after both loops.
+    dq = dq * sm_scale
+
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + qhid * stride_dqh
+        + offs_m[:, None] * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    if ACCUMULATE:
+        dq_prev = tl.load(dq_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
+        tl.store(dq_ptrs, dq + dq_prev, mask=offs_m[:, None] < seqlen_q)
+    else:
+        tl.store(dq_ptrs, dq, mask=offs_m[:, None] < seqlen_q)
+
+
+@triton.jit
+def _v4_attention_bwd_dkv_kernel(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DK,  # fp32 buffer [B, K_H, Sk, D]
+    DV,  # fp32 buffer [B, K_H, Sk, D]
+    ADD_MASK,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dkd,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_ADD_MASK: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    USE_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    # Pool-suffix flag: when HCA split-mask, this program is one of
+    # the (HEAD_Q-block over n in pool range). The pool branch iterates
+    # all m's, so we run it inside the same kernel after the local
+    # branch finishes.
+    ATOMIC_REDUCE: tl.constexpr = False,
+    # Plan-8 P57 — MQA head-split parallelism. When ``NUM_HEAD_GROUPS > 1``
+    # and ``HEAD_K != HEAD_Q`` (MQA), the kernel adds an extra grid dim
+    # ``pid_h_group = program_id(2)`` and each program owns
+    # ``HEAD_Q / NUM_HEAD_GROUPS`` query heads. Multiple head_group
+    # programs collide on the same MQA dK / dV slice, so writes are
+    # ``tl.atomic_add`` instead of ``tl.store``.
+    NUM_HEAD_GROUPS: tl.constexpr = 1,
+    EXACT_TILES_M: tl.constexpr = False,
+    EXACT_TILES_N: tl.constexpr = False,
+):
+    """V4 attention BWD — dK / dV only (parallel over n-blocks, no atomics for dK/dV).
+
+    For MQA (``HEAD_K == 1``) every query head contributes to the same
+    shared K / V, so this kernel must iterate ``H`` query heads per
+    ``(b, n_block)``. We expose that via the ``HEAD_Q`` constexpr loop.
+
+    HCA pool n-blocks (``n_block_start >= HCA_LOCAL_SEQLEN``) are handled
+    by :func:`_v4_attention_bwd_dkv_pool_kernel`, which parallelises the
+    pool work across ``(m_block, b * qhid)``. We early-return here for
+    those blocks so we do not double-count the pool contribution.
+    """
+    pid_n = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    pid_h_group = tl.program_id(2)
+    bid = pid_bh // HEAD_K
+    khid = pid_bh % HEAD_K
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    # HCA mode: pool n-blocks are handled by the dedicated pool kernel.
+    if HCA_LOCAL_SEQLEN > 0:
+        if pid_n * BLOCK_N >= HCA_LOCAL_SEQLEN:
+            return
+        is_pool_block = False
+    else:
+        is_pool_block = False
+
+    # Load K, V tiles for this n-block once.
+    k_ptrs = (
+        K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+    )
+    v_ptrs = (
+        V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+    )
+    kv_load_mask = offs_n[:, None] < seqlen_k
+    k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+    v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+    # R2: hoist tl.trans outside the m-loop. Lets the JIT keep kt / vt
+    # in registers across iterations instead of re-transposing every step.
+    kt = tl.trans(k)
+    vt = tl.trans(v)
+
+    dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+
+    # Determine m-loop bounds for this n-block. SWA / causal restricts the
+    # range of m's that can see this n-block. For ``HAS_ADD_MASK`` with
+    # arbitrary additive bias we iterate the full m axis.
+    n_block_lo = pid_n * BLOCK_N
+    n_block_hi = n_block_lo + BLOCK_N
+
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+        m_loop_start = 0
+        m_loop_end = seqlen_q
+    elif is_pool_block:
+        # Every m sees every visible pool slot; pool mask drives the rest.
+        m_loop_start = 0
+        m_loop_end = seqlen_q
+    elif SWA_WINDOW > 0:
+        # n is seen by m where n_block_lo <= m <= n_block_hi + SWA_WINDOW - 1.
+        # Causal also requires m >= n_block_lo. Round to BLOCK_M tiles.
+        m_loop_start = (n_block_lo // BLOCK_M) * BLOCK_M
+        m_loop_end = n_block_hi + SWA_WINDOW - 1
+        if m_loop_end > seqlen_q:
+            m_loop_end = seqlen_q
+        # Round up to BLOCK_M boundary so the m-loop iterates whole tiles.
+        m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M
+    elif USE_CAUSAL:
+        m_loop_start = (n_block_lo // BLOCK_M) * BLOCK_M
+        m_loop_end = seqlen_q
+        m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M
+    else:
+        m_loop_start = 0
+        m_loop_end = seqlen_q
+        m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M
+
+    # MHA path (HEAD_K == HEAD_Q) — only one query head contributes to
+    # this khid, so we use ``qhid = khid`` and skip the head loop. For
+    # MQA (HEAD_K != HEAD_Q) we iterate every query head.
+    if HEAD_K == HEAD_Q:
+        qhid = khid
+        for m_start in range(m_loop_start, m_loop_end, BLOCK_M):
+            offs_m = m_start + tl.arange(0, BLOCK_M)
+
+            q_ptrs = (
+                Q
+                + bid * stride_qb
+                + qhid * stride_qh
+                + offs_m[:, None] * stride_qm
+                + offs_d[None, :] * stride_qd
+            )
+            dout_ptrs = (
+                DOUT
+                + bid * stride_dob
+                + qhid * stride_doh
+                + offs_m[:, None] * stride_dom
+                + offs_d[None, :] * stride_dod
+            )
+            lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+            dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+            q_load_mask = offs_m[:, None] < seqlen_q
+            q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+            dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+            lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+            dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+            qk = tl.dot(q, kt) * sm_scale
+
+            if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+                mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn
+                mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+                add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+                qk = qk + add_bias
+            elif is_pool_block:
+                pool_n = offs_n - HCA_LOCAL_SEQLEN
+                mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+                mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+                add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+                qk = qk + add_bias
+            else:
+                if SWA_WINDOW > 0:
+                    in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & (
+                        offs_n[None, :] <= offs_m[:, None]
+                    )
+                    qk = tl.where(in_window, qk, NEG_INF)
+                elif USE_CAUSAL:
+                    qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF)
+
+            if not EXACT_TILES_N:
+                qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+            if not EXACT_TILES_M:
+                qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+            p = tl.exp(qk - lse[:, None])
+            dp = tl.dot(dout, vt)
+            ds = p * (dp - dvec[:, None])
+
+            # P57 cr=0 BWD: fuse the inner ``dv += dot(p, dout)`` and
+            # ``dk += dot(ds, q)`` into MFMA-acc form; defer ``sm_scale``
+            # on dk to a single multiply after the m-loop.
+            dv = tl.dot(tl.trans(p.to(dout.dtype)), dout, acc=dv)
+            dk = tl.dot(tl.trans(ds.to(q.dtype)), q, acc=dk)
+    else:
+        # MQA path. With NUM_HEAD_GROUPS > 1 the program owns a slice of
+        # query heads; otherwise it iterates all HEAD_Q heads.
+        head_per_group: tl.constexpr = HEAD_Q // NUM_HEAD_GROUPS
+        h_start = pid_h_group * head_per_group
+        h_end = h_start + head_per_group
+        for h_iter in range(h_start, h_end):
+            qhid = h_iter
+
+            for m_start in range(m_loop_start, m_loop_end, BLOCK_M):
+                offs_m = m_start + tl.arange(0, BLOCK_M)
+
+                q_ptrs = (
+                    Q
+                    + bid * stride_qb
+                    + qhid * stride_qh
+                    + offs_m[:, None] * stride_qm
+                    + offs_d[None, :] * stride_qd
+                )
+                dout_ptrs = (
+                    DOUT
+                    + bid * stride_dob
+                    + qhid * stride_doh
+                    + offs_m[:, None] * stride_dom
+                    + offs_d[None, :] * stride_dod
+                )
+                lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+                dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+                q_load_mask = offs_m[:, None] < seqlen_q
+                q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+                dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+                lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+                dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+                qk = tl.dot(q, kt) * sm_scale
+
+                if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+                    mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn
+                    mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+                    add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+                    qk = qk + add_bias
+                elif is_pool_block:
+                    pool_n = offs_n - HCA_LOCAL_SEQLEN
+                    mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+                    mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+                    add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+                    qk = qk + add_bias
+                else:
+                    if SWA_WINDOW > 0:
+                        in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & (
+                            offs_n[None, :] <= offs_m[:, None]
+                        )
+                        qk = tl.where(in_window, qk, NEG_INF)
+                    elif USE_CAUSAL:
+                        qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF)
+
+                if not EXACT_TILES_N:
+                    qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+                if not EXACT_TILES_M:
+                    qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+                p = tl.exp(qk - lse[:, None])
+                dp = tl.dot(dout, vt)
+                ds = p * (dp - dvec[:, None])
+
+                # P57 cr=0 BWD: same scale-defer + tl.dot(acc=) trick
+                # as the non-MQA branch above.
+                dv = tl.dot(tl.trans(p.to(dout.dtype)), dout, acc=dv)
+                dk = tl.dot(tl.trans(ds.to(q.dtype)), q, acc=dk)
+
+    # P57 cr=0 BWD: fold ``sm_scale`` on dk once after the (head x m)
+    # loop.  dv carries no scale.
+    dk = dk * sm_scale
+
+    dk_ptrs = (
+        DK
+        + bid * stride_dkb
+        + khid * stride_dkh
+        + offs_n[:, None] * stride_dkn
+        + offs_d[None, :] * stride_dkd
+    )
+    dv_ptrs = (
+        DV
+        + bid * stride_dvb
+        + khid * stride_dvh
+        + offs_n[:, None] * stride_dvn
+        + offs_d[None, :] * stride_dvd
+    )
+    if ATOMIC_REDUCE or NUM_HEAD_GROUPS > 1:
+        # ``ATOMIC_REDUCE`` is set on the CSA pool BWD path (multiple
+        # ``(b, head, n_block)`` programs collapse into a single
+        # ``(b, n_block)`` global slice with ``stride_dkh = stride_dvh = 0``).
+        # ``NUM_HEAD_GROUPS > 1`` is the MQA head-split path: multiple
+        # head_group programs target the same MQA dK / dV slice. Both
+        # require fp32 atomic_add to merge.
+        tl.atomic_add(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k, sem="relaxed")
+        tl.atomic_add(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k, sem="relaxed")
+    else:
+        tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
+        tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
+
+
+# ---------------------------------------------------------------------------
+# Plan-8 P57 — HCA pool dK / dV kernel
+#
+# The HCA split-mask BWD has a single "pool" n-block (P=Sk-HCA_LOCAL_SEQLEN
+# keys; typically 32 at V4-Flash widths). Folding the pool into the main
+# dK/dV kernel makes that single n-block program iterate every m-block ×
+# every query head — at H=64 / Sq=4096 / BLOCK_M=32 that is 8192 inner
+# iterations versus ~256 for a local-SWA n-block program (32× more work).
+#
+# This kernel parallelises pool work over ``m_block`` (one program per
+# ``(b, m_block)``) and iterates ``HEAD_Q`` heads internally. Each program
+# accumulates its (BLOCK_N, BLOCK_DMODEL) tile contribution in registers
+# (no per-head atomic) and atomic-adds the final tile into the shared
+# pool slice of dK / dV exactly twice (once for dk, once for dv) per
+# program. At V4-Flash widths the contention is 128-way per pool cell
+# instead of 8192-way for the m_block × qhid full grid, while still
+# extracting 128× more parallelism than the original single-program pool
+# branch.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_bwd_dkv_pool_kernel(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DK,  # fp32 buffer [B, K_H, Sk, D]
+    DV,  # fp32 buffer [B, K_H, Sk, D]
+    ADD_MASK,  # [Sq, P] additive pool mask
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dkd,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    pool_size,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """V4 HCA BWD — dK / dV for the pool keys (parallel m-blocks, head-loop inside)."""
+    pid_m = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    bid = pid_b
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    pool_n = tl.arange(0, BLOCK_N)
+    offs_n = HCA_LOCAL_SEQLEN + pool_n
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    pool_n_mask = pool_n < pool_size
+    q_load_mask = offs_m[:, None] < seqlen_q
+
+    # Per-(m_block, b) program iterates all query heads, accumulating in
+    # registers. For MQA (HK=1) all heads share the same (b, khid=0) K/V,
+    # so we reload K/V once at the start. For MHA (HK=HQ) we reload K/V
+    # per head inside the loop.
+    if HEAD_K == 1:
+        k_ptrs = K + bid * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+        v_ptrs = V + bid * stride_vb + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+        kv_load_mask = pool_n_mask[:, None]
+        k_shared = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+        v_shared = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+        kt_shared = tl.trans(k_shared)
+        vt_shared = tl.trans(v_shared)
+
+    dk_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dv_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+
+    for qhid in range(0, HEAD_Q):
+        if HEAD_K == HEAD_Q:
+            khid = qhid
+        else:
+            khid = 0
+
+        q_ptrs = (
+            Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+        )
+        dout_ptrs = (
+            DOUT
+            + bid * stride_dob
+            + qhid * stride_doh
+            + offs_m[:, None] * stride_dom
+            + offs_d[None, :] * stride_dod
+        )
+        lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+        dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+        q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+        dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+        lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+        dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+        if HEAD_K == 1:
+            kt = kt_shared
+            vt = vt_shared
+        else:
+            k_ptrs = (
+                K
+                + bid * stride_kb
+                + khid * stride_kh
+                + offs_n[:, None] * stride_kn
+                + offs_d[None, :] * stride_kd
+            )
+            v_ptrs = (
+                V
+                + bid * stride_vb
+                + khid * stride_vh
+                + offs_n[:, None] * stride_vn
+                + offs_d[None, :] * stride_vd
+            )
+            kv_load_mask = pool_n_mask[:, None]
+            k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+            v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+            kt = tl.trans(k)
+            vt = tl.trans(v)
+
+        qk = tl.dot(q, kt) * sm_scale
+        mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+        mask_load_mask = (offs_m[:, None] < seqlen_q) & pool_n_mask[None, :]
+        add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+        qk = qk + add_bias
+        qk = tl.where(pool_n_mask[None, :], qk, NEG_INF)
+        qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout, vt)
+        ds = p * (dp - dvec[:, None])
+
+        if HEAD_K == 1:
+            dk_acc += tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale
+            dv_acc += tl.dot(tl.trans(p.to(dout.dtype)), dout)
+        else:
+            # MHA path — each qhid maps to its own khid slice. Accumulator is
+            # not shared across heads; flush per-head with atomic_add.
+            dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale
+            dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout)
+            dk_ptrs_h = (
+                DK
+                + bid * stride_dkb
+                + khid * stride_dkh
+                + offs_n[:, None] * stride_dkn
+                + offs_d[None, :] * stride_dkd
+            )
+            dv_ptrs_h = (
+                DV
+                + bid * stride_dvb
+                + khid * stride_dvh
+                + offs_n[:, None] * stride_dvn
+                + offs_d[None, :] * stride_dvd
+            )
+            tl.atomic_add(dk_ptrs_h, dk_contrib, mask=pool_n_mask[:, None], sem="relaxed")
+            tl.atomic_add(dv_ptrs_h, dv_contrib, mask=pool_n_mask[:, None], sem="relaxed")
+
+    if HEAD_K == 1:
+        khid_final = 0
+        dk_ptrs = (
+            DK
+            + bid * stride_dkb
+            + khid_final * stride_dkh
+            + offs_n[:, None] * stride_dkn
+            + offs_d[None, :] * stride_dkd
+        )
+        dv_ptrs = (
+            DV
+            + bid * stride_dvb
+            + khid_final * stride_dvh
+            + offs_n[:, None] * stride_dvn
+            + offs_d[None, :] * stride_dvd
+        )
+        write_mask = pool_n_mask[:, None]
+        tl.atomic_add(dk_ptrs, dk_acc, mask=write_mask, sem="relaxed")
+        tl.atomic_add(dv_ptrs, dv_acc, mask=write_mask, sem="relaxed")
+
+
+# ---------------------------------------------------------------------------
+# Plan-8 P57 R2 — atomic-free MHA pool dK / dV kernel
+#
+# The R1 ``_v4_attention_bwd_dkv_pool_kernel`` MHA branch parallelises over
+# m-blocks (one program per ``(b, m_block)``) and atomic-adds the per-head
+# dK / dV contribution INSIDE the head loop:
+#
+#     for qhid in range(HEAD_Q):  # 64 heads
+#         tl.atomic_add(dk_ptrs[qhid], dk_contrib)
+#         tl.atomic_add(dv_ptrs[qhid], dv_contrib)
+#
+# At V4-Flash widths (B=1 H=64 Sq=4096 BM=64) this fires
+#   m_blocks × heads × 2 = 64 × 64 × 2 = 8192
+# ``tl.atomic_add`` instructions per pool kernel invocation, each on a
+# (BLOCK_N=32, D=512) fp32 tile. Even at ``sem="relaxed"`` the L2 atomic
+# engine on MI355 serialises 64-way contention on each per-head pool
+# cache line — the pool kernel is dominated by atomic stalls.
+#
+# R2 rewrite: SWAP the loop nesting. Parallelise over ``(b, qhid)`` (one
+# program per query head, owning the unique ``DK / DV[bid, qhid, pool, :]``
+# slice). The kernel pre-loads ``K / V[qhid, pool]`` ONCE, then iterates
+# all m-blocks for that head, accumulating ``dk_acc / dv_acc`` in
+# registers. The final write is a single ``tl.store`` per program — no
+# atomic at all.
+#
+# Grid: ``(HEAD_Q, B)`` = (64, 1) at V4-Flash → 64 programs. That's only
+# 25% of MI355's 256 CUs, but each program now has ~64 m-iter × dense
+# matmul work (≈ heavy enough to saturate VALU+MFMA in its CU).
+#
+# Per-program HBM:
+#   * K / V[qhid, pool] = 2 × 32 × 512 × 2 B = 64 KiB (loaded once)
+#   * Q / dout[qhid, *] = 2 × Sq × D × 2 B = 8 MiB total per program
+#   * dK / dV[qhid, pool] = 2 × 32 × 512 × 4 B = 128 KiB (stored once)
+#
+# Compared to the R1 pool kernel:
+#   * Total Q / dout HBM reads UNCHANGED (64 programs × 64 m-iter ==
+#     R1's 64 programs × 64 head-iter).
+#   * Total K / V HBM reads DROP ~64× (each (head, pool) tile loaded
+#     exactly once instead of once per m-block).
+#   * Atomic_add instructions DROP from 8192 → 0.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_bwd_dkv_pool_mha_kernel(
+    Q,
+    K,
+    V,
+    DOUT,
+    LSE,
+    D,
+    DK,
+    DV,
+    ADD_MASK,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dkd,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvd,
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    pool_size,
+    sm_scale,
+    HCA_LOCAL_SEQLEN: tl.constexpr,
+    DK_POOL_OFFSET: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    POOL_EXACT: tl.constexpr,
+    EXACT_TILES_M: tl.constexpr,
+    M_SPLIT: tl.constexpr,
+):
+    """MHA-only atomic-free HCA pool dK / dV kernel.
+
+    Parallelism: ``(M_SPLIT * HEAD_Q, B)`` — ``M_SPLIT`` programs per
+    ``(b, qhid)``. ``M_SPLIT=1`` is the pure atomic-free design (one
+    program owns the full pool slice for its head); ``M_SPLIT>1`` shards
+    the m-loop into ``M_SPLIT`` chunks and merges with ``M_SPLIT``-way
+    ``tl.atomic_add`` per (b, qhid, pool, :) slice — the contention is
+    bounded by ``M_SPLIT`` and the extra programs lift HBM utilisation.
+
+    Constraints: ``HEAD_K == HEAD_Q`` (MHA only). The MQA path keeps
+    using the original :func:`_v4_attention_bwd_dkv_pool_kernel`.
+    """
+    pid_hm = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    bid = pid_b
+    qhid = pid_hm // M_SPLIT
+    pid_m_chunk = pid_hm % M_SPLIT
+    khid = qhid  # MHA invariant
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    pool_n = tl.arange(0, BLOCK_N)
+    offs_n = HCA_LOCAL_SEQLEN + pool_n
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    # Pre-load K / V for this (bid, khid) pool slice ONCE.
+    k_ptrs = (
+        K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+    )
+    v_ptrs = (
+        V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+    )
+    if POOL_EXACT:
+        k_tile = tl.load(k_ptrs)
+        v_tile = tl.load(v_ptrs)
+    else:
+        pool_n_mask = pool_n < pool_size
+        kv_load_mask = pool_n_mask[:, None]
+        k_tile = tl.load(k_ptrs, mask=kv_load_mask, other=0.0)
+        v_tile = tl.load(v_ptrs, mask=kv_load_mask, other=0.0)
+    kt = tl.trans(k_tile)
+    vt = tl.trans(v_tile)
+
+    dk_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dv_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+
+    # Shard the m-loop into M_SPLIT contiguous chunks; each program owns
+    # one chunk of m-blocks. Chunk size is ``ceil(num_m_blocks / M_SPLIT)``
+    # rounded up so the last chunk may be shorter.
+    num_m_blocks = (seqlen_q + BLOCK_M - 1) // BLOCK_M
+    m_chunk_size = (num_m_blocks + M_SPLIT - 1) // M_SPLIT
+    m_block_lo = pid_m_chunk * m_chunk_size
+    m_block_hi = m_block_lo + m_chunk_size
+    if m_block_hi > num_m_blocks:
+        m_block_hi = num_m_blocks
+    m_loop_start = m_block_lo * BLOCK_M
+    m_loop_end = m_block_hi * BLOCK_M
+    if m_loop_end > seqlen_q:
+        m_loop_end = seqlen_q
+
+    for m_start in range(m_loop_start, m_loop_end, BLOCK_M):
+        offs_m = m_start + tl.arange(0, BLOCK_M)
+
+        q_ptrs = (
+            Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+        )
+        dout_ptrs = (
+            DOUT
+            + bid * stride_dob
+            + qhid * stride_doh
+            + offs_m[:, None] * stride_dom
+            + offs_d[None, :] * stride_dod
+        )
+        lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+        dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm
+
+        if EXACT_TILES_M:
+            q_local = tl.load(q_ptrs)
+            dout_local = tl.load(dout_ptrs)
+            lse = tl.load(lse_ptrs)
+            dvec = tl.load(dvec_ptrs)
+        else:
+            q_load_mask = offs_m[:, None] < seqlen_q
+            q_local = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+            dout_local = tl.load(dout_ptrs, mask=q_load_mask, other=0.0)
+            lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0)
+            dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0)
+
+        # qk = Q @ K.T then scale + additive bias. We *cannot* defer the
+        # scale across the softmax (exp(scale*x) ≠ exp(x)), so scale stays
+        # on the qk path.
+        qk = tl.dot(q_local, kt) * sm_scale
+        mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+        if POOL_EXACT and EXACT_TILES_M:
+            add_bias = tl.load(mask_ptrs).to(tl.float32)
+        else:
+            if POOL_EXACT:
+                mask_load_mask = offs_m[:, None] < seqlen_q
+            elif EXACT_TILES_M:
+                mask_load_mask = pool_n[None, :] < pool_size
+            else:
+                mask_load_mask = (offs_m[:, None] < seqlen_q) & (pool_n[None, :] < pool_size)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+        qk = qk + add_bias
+        if not POOL_EXACT:
+            qk = tl.where(pool_n[None, :] < pool_size, qk, NEG_INF)
+        if not EXACT_TILES_M:
+            qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_local, vt)
+        ds = p * (dp - dvec[:, None])
+
+        # R2 scale-defer for dk: accumulate ds.T @ Q without sm_scale,
+        # fold the single multiply after the m-loop. dv carries no scale.
+        dv_acc = tl.dot(tl.trans(p.to(dout_local.dtype)), dout_local, acc=dv_acc)
+        dk_acc = tl.dot(tl.trans(ds.to(q_local.dtype)), q_local, acc=dk_acc)
+
+    dk_acc = dk_acc * sm_scale
+
+    # Plan-8 P57 R2: DK / DV may either be the full ``[B, HK, Sk, D]``
+    # fp32 buffer (then ``DK_POOL_OFFSET == HCA_LOCAL_SEQLEN`` so we
+    # land on the pool slice) or a pool-only ``[B, HK, pool_size, D]``
+    # fp32 sidecar (then ``DK_POOL_OFFSET == 0``). The latter is the
+    # default R2 path: the full ``dk_out / dv_out`` allocate in the
+    # input dtype and we cast the small sidecar to bf16 at the end —
+    # saves the big ``dk_fp32 -> bf16`` cast.
+    offs_n_dk = DK_POOL_OFFSET + pool_n
+    dk_ptrs = (
+        DK
+        + bid * stride_dkb
+        + khid * stride_dkh
+        + offs_n_dk[:, None] * stride_dkn
+        + offs_d[None, :] * stride_dkd
+    )
+    dv_ptrs = (
+        DV
+        + bid * stride_dvb
+        + khid * stride_dvh
+        + offs_n_dk[:, None] * stride_dvn
+        + offs_d[None, :] * stride_dvd
+    )
+    if M_SPLIT == 1:
+        if POOL_EXACT:
+            tl.store(dk_ptrs, dk_acc)
+            tl.store(dv_ptrs, dv_acc)
+        else:
+            write_mask = pool_n[:, None] < pool_size
+            tl.store(dk_ptrs, dk_acc, mask=write_mask)
+            tl.store(dv_ptrs, dv_acc, mask=write_mask)
+    else:
+        if POOL_EXACT:
+            tl.atomic_add(dk_ptrs, dk_acc, sem="relaxed")
+            tl.atomic_add(dv_ptrs, dv_acc, sem="relaxed")
+        else:
+            write_mask = pool_n[:, None] < pool_size
+            tl.atomic_add(dk_ptrs, dk_acc, mask=write_mask, sem="relaxed")
+            tl.atomic_add(dv_ptrs, dv_acc, mask=write_mask, sem="relaxed")
+
+
+# ---------------------------------------------------------------------------
+# Python launcher
+# ---------------------------------------------------------------------------
+
+
+def _launch_v4_attention_bwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k: torch.Tensor,  # [B, K_H, Sk, D]
+    v: torch.Tensor,  # [B, K_H, Sk, D]
+    out: torch.Tensor,  # [B, H, Sq, D] (FWD output)
+    dout: torch.Tensor,  # [B, H, Sq, D]
+    lse: torch.Tensor,  # [B, H, Sq] fp32
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],  # [Sq, Sk] or None
+    scale: float,
+    hca_local_seqlen: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Launch the V4 attention backward kernel.
+
+    Returns ``(dq, dk, dv, dsink)`` — gradients in the input dtype, with
+    ``dsink`` returned only when ``sink is not None`` (else ``None``).
+    """
+    if not q.is_cuda:
+        raise ValueError("v4_attention_v1 BWD requires CUDA / HIP tensors.")
+    if dout.shape != out.shape or out.shape != q.shape:
+        raise ValueError(
+            "v4_attention_v1 BWD shape mismatch: "
+            f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}"
+        )
+
+    B, HQ, Sq, D = q.shape
+    HK = k.shape[1]
+    Sk = k.shape[2]
+
+    has_sink = sink is not None
+    has_add_mask = additive_mask is not None
+    hca_local_seqlen = int(hca_local_seqlen)
+    if hca_local_seqlen:
+        if not has_add_mask:
+            raise ValueError("hca_local_seqlen requires a pool additive_mask.")
+        if hca_local_seqlen <= 0 or hca_local_seqlen >= Sk:
+            raise ValueError(
+                "hca_local_seqlen must split local and pool keys "
+                f"(got hca_local_seqlen={hca_local_seqlen}, Sk={Sk})."
+            )
+        expected_mask_shape = (Sq, Sk - hca_local_seqlen)
+        if tuple(additive_mask.shape) != expected_mask_shape:
+            raise ValueError(
+                "HCA split-mask mode expects additive_mask shape "
+                f"{expected_mask_shape}, got {tuple(additive_mask.shape)}."
+            )
+        if swa_window <= 0:
+            raise ValueError("HCA split-mask mode requires swa_window > 0.")
+    use_causal = (not has_add_mask) and (swa_window <= 0)
+    swa_window_constexpr = (
+        int(swa_window) if ((not has_add_mask or hca_local_seqlen) and swa_window > 0) else 0
+    )
+
+    # Plan-8 P57 R2 sweep at V4-Flash widths (B=1 H=64 Sq=4096 D=512 P=32):
+    # R1 settled on ``BLOCK_M=64 BLOCK_N=16``; the R2 sweep — once the
+    # MHA pool kernel went atomic-free with ``M_SPLIT=4`` — discovered
+    # ``BLOCK_M=32`` with ``dKV num_warps=2`` cuts cr=128 HCA BWD from
+    # 3.79 ms to ~3.44 ms AND cr=0 dense BWD from 3.01 ms to ~2.77 ms.
+    # Why BM=32 now wins:
+    #
+    #   * dq: smaller Q tile (32×D=512=32KiB vs 64×512=64KiB) drops VGPR
+    #     pressure enough to run ``num_warps=2`` (1 warp = 64 threads
+    #     per program) instead of nw=4 — same overall warp occupancy
+    #     across the grid but the per-warp register budget doubles and
+    #     k-axis matmul (D=512) pipelines without VGPR spilling.
+    #   * dkv (local SWA): the m-loop is 5 iter at BM=32 (vs 3 iter at
+    #     BM=64), but each iter loads half the Q/dout tile (32×512×2 B
+    #     each) so per-iter HBM is halved. Total per-program HBM is
+    #     roughly the same; the win is reduced per-iter register pressure
+    #     enabling ``num_warps=2``.
+    #   * pool (MHA atomic-free): more m-blocks per program (128 at
+    #     BM=32 vs 64 at BM=64) lets the ``M_SPLIT=4`` shard cleanly
+    #     into 32 m-iter per program (vs 16) and fills the SIMDs.
+    BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_M", "32"))
+    BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_N", "16"))
+    BLOCK_DMODEL = D
+
+    # Plan-8 P57 R2: dq/dk/dv output buffer strategy
+    # ----------------------------------------------
+    # dq is written by ``_v4_attention_bwd_dq_kernel`` via ``tl.store``
+    # (no atomic_add — each (b, qhid, m_block) program owns a unique
+    # slice). So we let the kernel write directly into the input dtype
+    # and skip the final fp32 -> bf16 cast (~ 0.15 ms saved at V4-Flash
+    # widths). The dq accumulator inside the kernel remains fp32 for the
+    # n-loop reduction; only the final store converts.
+    #
+    # dk / dv are similar for the MHA non-MQA case (NUM_HEAD_GROUPS=1):
+    # the dKV kernel ``tl.store``s the local-SWA slice with no atomic.
+    # The HCA pool kernel still uses atomic_add (M_SPLIT > 1) which
+    # requires fp32, so we keep a SEPARATE small fp32 pool sidecar
+    # buffer of shape ``[B, HK, pool_size, D]`` and merge it into the
+    # full bf16 dk / dv at the end.
+    #
+    # For MQA (HK != HQ) or NUM_HEAD_GROUPS > 1: keep the original
+    # fp32-everywhere path since dKV uses atomic_add.
+    dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=q.dtype)
+    use_native_dkv_dtype = (HQ == HK) and (q.dtype in (torch.bfloat16, torch.float16))
+    if use_native_dkv_dtype:
+        dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=k.dtype)
+        dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=v.dtype)
+    else:
+        dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32)
+        dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32)
+    # Pool-side fp32 sidecar (only allocated when HCA + MHA + native dtype dKV).
+    if use_native_dkv_dtype and hca_local_seqlen > 0:
+        pool_size_alloc = Sk - hca_local_seqlen
+        dk_pool_fp32 = torch.zeros((B, HK, pool_size_alloc, D), device=q.device, dtype=torch.float32)
+        dv_pool_fp32 = torch.zeros((B, HK, pool_size_alloc, D), device=q.device, dtype=torch.float32)
+    else:
+        dk_pool_fp32 = None
+        dv_pool_fp32 = None
+    if has_sink:
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink
+    else:
+        dsink_fp32 = q  # sentinel; HAS_SINK=False inside kernel
+        sink_arg = q
+
+    # D scalar = (dout * out).sum(-1)
+    d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+    pre_grid = (triton.cdiv(Sq, BLOCK_M), B * HQ)
+    _v4_attention_bwd_preprocess_kernel[pre_grid](
+        out,
+        dout,
+        d_buf,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        Sq,
+        HEAD=HQ,
+        BLOCK_M=BLOCK_M,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+
+    # Mask sentinel ptr for HAS_ADD_MASK=False
+    mask_arg = additive_mask if has_add_mask else q
+    if has_add_mask:
+        stride_ms = additive_mask.stride(0)
+        stride_mn = additive_mask.stride(1)
+    else:
+        stride_ms = 0
+        stride_mn = 0
+
+    # Plan-8 P57: skip the boundary masks (``offs_m < seqlen_q`` /
+    # ``offs_n < seqlen_k``) when the launcher confirms the seqlens are
+    # exact multiples of BLOCK_M / BLOCK_N. Saves a per-inner-iter
+    # broadcast + tl.where on production V4-Flash widths where Sq=4096
+    # (BM=64) and Sk=4128 (BN=16) both divide cleanly.
+    exact_tiles_m = (Sq % BLOCK_M) == 0
+    exact_tiles_n = (Sk % BLOCK_N) == 0
+
+    # Plan-5 P32: split BWD (dQ kernel + dK/dV kernel, no atomics for
+    # dQ / dK / dV) is now the default — wins both the operator microbench
+    # *and* the EP8 proxy after the dual-RoPE bf16-cast fix (P32 RoPE bug:
+    # ``apply_interleaved_partial_rope`` was upcasting Q/K to fp32 because
+    # cos/sin came from ``position_ids.float()`` and bf16 * fp32 = fp32,
+    # which 2x'd Q/K HBM traffic, inflated the kernel time 1.8-7x in the
+    # proxy and made the monolithic design *look* faster in A/B traces).
+    # ``PRIMUS_V4_ATTN_BWD_USE_SPLIT=0`` falls back to the monolithic
+    # design for kernel-level perf experiments / debugging.
+    if os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "1") == "1":
+        dq_grid = (triton.cdiv(Sq, BLOCK_M), B * HQ)
+        _v4_attention_bwd_dq_kernel[dq_grid](
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dq_fp32,
+            dsink_fp32,
+            sink_arg,
+            mask_arg,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            k.stride(0),
+            k.stride(1),
+            k.stride(2),
+            k.stride(3),
+            v.stride(0),
+            v.stride(1),
+            v.stride(2),
+            v.stride(3),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            d_buf.stride(0),
+            d_buf.stride(1),
+            d_buf.stride(2),
+            dq_fp32.stride(0),
+            dq_fp32.stride(1),
+            dq_fp32.stride(2),
+            dq_fp32.stride(3),
+            stride_ms,
+            stride_mn,
+            Sq,
+            Sk,
+            float(scale),
+            HEAD_Q=HQ,
+            HEAD_K=HK,
+            SWA_WINDOW=swa_window_constexpr,
+            HAS_SINK=has_sink,
+            HAS_ADD_MASK=has_add_mask,
+            HCA_LOCAL_SEQLEN=hca_local_seqlen,
+            USE_CAUSAL=use_causal,
+            BLOCK_M=BLOCK_M,
+            BLOCK_N=BLOCK_N,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            EXACT_TILES_M=exact_tiles_m,
+            EXACT_TILES_N=exact_tiles_n,
+            # Plan-8 P57 R2: ``num_warps=2 num_stages=1`` is the new sweet
+            # spot once BLOCK_M dropped from 64 -> 32. With a 32×D Q tile
+            # the dq inner loop fits in the per-warp register budget with
+            # only 1 warp/64-thread wave, doubling the per-warp register
+            # budget vs the R1 nw=4 default and removing the K/V double-
+            # buffer (num_stages=2) cost since the loop now runs short.
+            num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_WARPS", "2")),
+            num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_STAGES", "1")),
+        )
+        # Plan-8 P57: in HCA mode the pool n-block(s) are handled by the
+        # dedicated pool kernel that parallelises over (m_block, b * qhid)
+        # — see ``_v4_attention_bwd_dkv_pool_kernel``. Skip the pool grid
+        # rows here so we don't double-count.
+        # Plan-8 P57 R2: per-kernel BN override. Default uses the global
+        # BLOCK_N (BN=16 from R1 sweep). Setting ``PRIMUS_V4_ATTN_BWD_DKV_BLOCK_N``
+        # lets us test wider K-tiles in the dKV kernel without affecting
+        # the dq / pool paths.
+        dkv_block_n = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_BLOCK_N", str(BLOCK_N)))
+        if hca_local_seqlen > 0:
+            dkv_n_blocks = triton.cdiv(hca_local_seqlen, dkv_block_n)
+        else:
+            dkv_n_blocks = triton.cdiv(Sk, dkv_block_n)
+        # Plan-8 P57: ``NUM_HEAD_GROUPS`` controls MQA head-split parallelism
+        # for the local dKV kernel. Default (1) keeps the original head loop;
+        # >1 splits the head loop and uses ``tl.atomic_add`` for dKV.
+        # The dense-bench sweep showed head-split is essentially neutral on
+        # MI355 at H=64 (HBM/atomic-bound, not compute-bound), so default 1.
+        # BUT that result does NOT hold for the MQA/HQ>=128 (V4-Pro) SWA dkv
+        # shape: at HQ=128/HK=1 the HG=1 grid is 1 workgroup/CU (occ=1) with each
+        # program serially grinding all 128 heads -> tail/pipeline-fill bound.
+        # HG=2 doubles the grid + fills the idle MFMA cycles (grad cos ~1.0);
+        # HG=4/8 regress. So default the MQA/HQ>=128 case to 2; still
+        # overridable via PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS.
+        num_head_groups = 1
+        if HQ > HK:
+            # Arch-aware (see ._v4_attn_tuning): gfx1250 wants HG=32 (+37-53%, ab_sweep/opt7b),
+            # gfx950 wants 2. Env knob still overrides.
+            from ._v4_attn_tuning import bwd_dkv_head_groups_default
+
+            _hg_default = str(bwd_dkv_head_groups_default(HQ, HK))
+            target = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS", _hg_default))
+            while target > 1 and HQ % target != 0:
+                target //= 2
+            num_head_groups = max(1, target)
+        dkv_grid = (dkv_n_blocks, B * HK, num_head_groups)
+        # Plan-8 P57 R2: ``num_warps=2 num_stages=1`` paired with BM=32.
+        # The 32×D Q tile + persistent K/V tile fit in the 2-warp register
+        # budget; bumping to nw=4 doubles spills (m-loop is 5 iter at
+        # BM=32 so the wider warp budget hurts more than the SIMD-fill
+        # benefit). num_stages>1 also costs more than it pays because
+        # the K/V are loaded once before the m-loop.
+        dkv_num_warps = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_WARPS", "2"))
+        dkv_num_stages = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_STAGES", "1"))
+        _v4_attention_bwd_dkv_kernel[dkv_grid](
+            q,
+            k,
+            v,
+            dout,
+            lse,
+            d_buf,
+            dk_fp32,
+            dv_fp32,
+            mask_arg,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            k.stride(0),
+            k.stride(1),
+            k.stride(2),
+            k.stride(3),
+            v.stride(0),
+            v.stride(1),
+            v.stride(2),
+            v.stride(3),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            d_buf.stride(0),
+            d_buf.stride(1),
+            d_buf.stride(2),
+            dk_fp32.stride(0),
+            dk_fp32.stride(1),
+            dk_fp32.stride(2),
+            dk_fp32.stride(3),
+            dv_fp32.stride(0),
+            dv_fp32.stride(1),
+            dv_fp32.stride(2),
+            dv_fp32.stride(3),
+            stride_ms,
+            stride_mn,
+            Sq,
+            Sk,
+            float(scale),
+            HEAD_Q=HQ,
+            HEAD_K=HK,
+            SWA_WINDOW=swa_window_constexpr,
+            HAS_ADD_MASK=has_add_mask,
+            HCA_LOCAL_SEQLEN=hca_local_seqlen,
+            USE_CAUSAL=use_causal,
+            BLOCK_M=BLOCK_M,
+            BLOCK_N=dkv_block_n,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            NUM_HEAD_GROUPS=num_head_groups,
+            EXACT_TILES_M=exact_tiles_m,
+            EXACT_TILES_N=(Sk % dkv_block_n) == 0,
+            num_warps=dkv_num_warps,
+            num_stages=dkv_num_stages,
+        )
+        if hca_local_seqlen > 0 and os.getenv("PRIMUS_V4_ATTN_BWD_HCA_POOL", "1") == "1":
+            pool_size = Sk - hca_local_seqlen
+            pool_block_n = max(16, triton.next_power_of_2(pool_size))
+            pool_exact = pool_size == pool_block_n
+            # Plan-8 P57 R2: use the atomic-free MHA pool kernel when
+            # HEAD_K == HEAD_Q (production V4-Flash widths). Parallelism
+            # is over (qhid, b) instead of (m_block, b); each program
+            # owns the unique dK / dV[bid, qhid, pool, :] slice and writes
+            # it with ``tl.store`` after accumulating across all m-blocks
+            # in registers — no atomic_add at all. Eliminates the 8192
+            # per-launch atomic_add stalls of the R1 MHA pool path.
+            #
+            # For MQA (HK=1, multiple qhids share one khid slice) we still
+            # need to merge across heads; fall back to the original pool
+            # kernel which atomic-adds in the head loop.
+            use_mha_pool = (HQ == HK) and (os.getenv("PRIMUS_V4_ATTN_BWD_POOL_MHA", "1") == "1")
+            if use_mha_pool:
+                # Plan-8 P57 R2: pool M_SPLIT controls m-loop sharding.
+                # M_SPLIT=1 = pure atomic-free (64 progs); M_SPLIT=N adds
+                # N-way atomic_add per (b, qhid, pool) slice but lifts the
+                # program count to N*HQ (better HBM bandwidth utilization
+                # at the cost of contention). Default 4 = 256 progs.
+                pool_m_split = int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_M_SPLIT", "4"))
+                num_m_blocks = triton.cdiv(Sq, BLOCK_M)
+                while pool_m_split > num_m_blocks:
+                    pool_m_split //= 2
+                pool_m_split = max(1, pool_m_split)
+                pool_grid = (HQ * pool_m_split, B)
+                # Plan-8 P57 R2: pool kernel writes to fp32 sidecar (so
+                # ``M_SPLIT > 1`` atomic_add works). The kernel uses
+                # offset ``HCA_LOCAL_SEQLEN + pool_n`` along the seqlen
+                # axis, so we pass a "virtual" full-seqlen view of the
+                # pool buffer: a tensor of shape (B, HK, Sk, D) where the
+                # data starts ``HCA_LOCAL_SEQLEN`` BEFORE the pool
+                # buffer's actual base. We accomplish this by passing the
+                # pool buffer's base pointer with an *offset* base
+                # subtracted via the seqlen stride. PyTorch tensors don't
+                # support negative offsets cleanly, so instead we pass
+                # the pool buffer's storage with seqlen stride and tell
+                # the kernel to subtract HCA_LOCAL_SEQLEN — see the
+                # ``DK_SEQ_OFFSET`` constexpr below.
+                if dk_pool_fp32 is not None:
+                    pool_dk_buf = dk_pool_fp32
+                    pool_dv_buf = dv_pool_fp32
+                    pool_dk_strides = (
+                        dk_pool_fp32.stride(0),
+                        dk_pool_fp32.stride(1),
+                        dk_pool_fp32.stride(2),
+                        dk_pool_fp32.stride(3),
+                    )
+                    pool_dv_strides = (
+                        dv_pool_fp32.stride(0),
+                        dv_pool_fp32.stride(1),
+                        dv_pool_fp32.stride(2),
+                        dv_pool_fp32.stride(3),
+                    )
+                    pool_hca_arg = 0  # pool buffer is pool-only; no HCA_LOCAL offset
+                else:
+                    pool_dk_buf = dk_fp32
+                    pool_dv_buf = dv_fp32
+                    pool_dk_strides = (
+                        dk_fp32.stride(0),
+                        dk_fp32.stride(1),
+                        dk_fp32.stride(2),
+                        dk_fp32.stride(3),
+                    )
+                    pool_dv_strides = (
+                        dv_fp32.stride(0),
+                        dv_fp32.stride(1),
+                        dv_fp32.stride(2),
+                        dv_fp32.stride(3),
+                    )
+                    pool_hca_arg = hca_local_seqlen
+                _v4_attention_bwd_dkv_pool_mha_kernel[pool_grid](
+                    q,
+                    k,
+                    v,
+                    dout,
+                    lse,
+                    d_buf,
+                    pool_dk_buf,
+                    pool_dv_buf,
+                    additive_mask,
+                    q.stride(0),
+                    q.stride(1),
+                    q.stride(2),
+                    q.stride(3),
+                    k.stride(0),
+                    k.stride(1),
+                    k.stride(2),
+                    k.stride(3),
+                    v.stride(0),
+                    v.stride(1),
+                    v.stride(2),
+                    v.stride(3),
+                    dout.stride(0),
+                    dout.stride(1),
+                    dout.stride(2),
+                    dout.stride(3),
+                    lse.stride(0),
+                    lse.stride(1),
+                    lse.stride(2),
+                    d_buf.stride(0),
+                    d_buf.stride(1),
+                    d_buf.stride(2),
+                    pool_dk_strides[0],
+                    pool_dk_strides[1],
+                    pool_dk_strides[2],
+                    pool_dk_strides[3],
+                    pool_dv_strides[0],
+                    pool_dv_strides[1],
+                    dv_fp32.stride(2),
+                    dv_fp32.stride(3),
+                    stride_ms,
+                    stride_mn,
+                    Sq,
+                    pool_size,
+                    float(scale),
+                    HCA_LOCAL_SEQLEN=hca_local_seqlen,
+                    DK_POOL_OFFSET=pool_hca_arg,
+                    BLOCK_M=BLOCK_M,
+                    BLOCK_N=pool_block_n,
+                    BLOCK_DMODEL=BLOCK_DMODEL,
+                    POOL_EXACT=pool_exact,
+                    EXACT_TILES_M=exact_tiles_m,
+                    M_SPLIT=pool_m_split,
+                    # Plan-8 P57 R2: ``num_warps=4 num_stages=3`` for the
+                    # atomic-free MHA pool kernel. Pool m-loop is long
+                    # (Sq/(BM*M_SPLIT) = 32 iter at BM=32 M_SPLIT=4) so
+                    # 3-stage pipelining of Q+dout loads pays off.
+                    num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_WARPS", "4")),
+                    num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_STAGES", "3")),
+                )
+            else:
+                # Plan-8 P57: pool kernel block_m is decoupled from the dKV
+                # block_m. dKV benefits from BM=64 (fewer programs each loading
+                # a wider Q tile); pool kernel benefits from a smaller BM (more
+                # programs => more parallelism over the head loop).
+                # gfx950/MI355X: the pool kernel scales with program count, so
+                # BM=16 (vs the dKV default 32) speeds up the full HCA backward
+                # (cos ~1.0). Only the pool path reads this, so SWA bwd is
+                # unaffected. Env-overridable.
+                pool_block_m = int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_BLOCK_M", "16"))
+                pool_grid = (triton.cdiv(Sq, pool_block_m), B)
+                _v4_attention_bwd_dkv_pool_kernel[pool_grid](
+                    q,
+                    k,
+                    v,
+                    dout,
+                    lse,
+                    d_buf,
+                    dk_fp32,
+                    dv_fp32,
+                    additive_mask,
+                    q.stride(0),
+                    q.stride(1),
+                    q.stride(2),
+                    q.stride(3),
+                    k.stride(0),
+                    k.stride(1),
+                    k.stride(2),
+                    k.stride(3),
+                    v.stride(0),
+                    v.stride(1),
+                    v.stride(2),
+                    v.stride(3),
+                    dout.stride(0),
+                    dout.stride(1),
+                    dout.stride(2),
+                    dout.stride(3),
+                    lse.stride(0),
+                    lse.stride(1),
+                    lse.stride(2),
+                    d_buf.stride(0),
+                    d_buf.stride(1),
+                    d_buf.stride(2),
+                    dk_fp32.stride(0),
+                    dk_fp32.stride(1),
+                    dk_fp32.stride(2),
+                    dk_fp32.stride(3),
+                    dv_fp32.stride(0),
+                    dv_fp32.stride(1),
+                    dv_fp32.stride(2),
+                    dv_fp32.stride(3),
+                    stride_ms,
+                    stride_mn,
+                    Sq,
+                    Sk,
+                    pool_size,
+                    float(scale),
+                    HEAD_Q=HQ,
+                    HEAD_K=HK,
+                    HCA_LOCAL_SEQLEN=hca_local_seqlen,
+                    BLOCK_M=pool_block_m,
+                    BLOCK_N=pool_block_n,
+                    BLOCK_DMODEL=BLOCK_DMODEL,
+                    # Plan-8 P57 sweep with dKV BM=64 BN=16: pool kernel
+                    # ``num_warps=4 num_stages=1`` is the sweet spot. Higher
+                    # num_stages adds Q/dout double-buffering overhead that
+                    # doesn't pay off because the pool kernel keeps K/V in
+                    # registers (single load) and the head-loop dominates
+                    # the iteration count.
+                    num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_WARPS", "4")),
+                    num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_STAGES", "1")),
+                )
+        # Plan-8 P57 R2: cast outputs back to input dtype only when not
+        # already in the input dtype. dq is always allocated in q.dtype
+        # (no atomic), so this is a no-op. dk / dv may be either:
+        #   * native bf16 (HQ == HK MHA path) — no cast needed for the
+        #     local-SWA part; only the pool sidecar (if present) needs
+        #     a fp32 -> bf16 cast that we splice into the pool slice.
+        #   * fp32 (MQA or grouped paths) — full cast at the end.
+        dq_out = dq_fp32 if dq_fp32.dtype == q.dtype else dq_fp32.to(q.dtype)
+        if dk_fp32.dtype != k.dtype:
+            dk_out = dk_fp32.to(k.dtype)
+        else:
+            dk_out = dk_fp32
+        if dv_fp32.dtype != v.dtype:
+            dv_out = dv_fp32.to(v.dtype)
+        else:
+            dv_out = dv_fp32
+        if dk_pool_fp32 is not None and hca_local_seqlen > 0:
+            # Splice the pool sidecar (fp32) into the pool slice of the
+            # bf16 dk_out / dv_out buffers.
+            dk_out[..., hca_local_seqlen:, :] = dk_pool_fp32.to(dk_out.dtype)
+            dv_out[..., hca_local_seqlen:, :] = dv_pool_fp32.to(dv_out.dtype)
+        dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+        return dq_out, dk_out, dv_out, dsink_out
+
+    grid = (triton.cdiv(Sq, BLOCK_M), B * HQ)
+    _v4_attention_bwd_kernel[grid](
+        q,
+        k,
+        v,
+        dout,
+        lse,
+        d_buf,
+        dq_fp32,
+        dk_fp32,
+        dv_fp32,
+        dsink_fp32,
+        sink_arg,
+        mask_arg,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        k.stride(3),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        v.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        dq_fp32.stride(0),
+        dq_fp32.stride(1),
+        dq_fp32.stride(2),
+        dq_fp32.stride(3),
+        dk_fp32.stride(0),
+        dk_fp32.stride(1),
+        dk_fp32.stride(2),
+        dk_fp32.stride(3),
+        dv_fp32.stride(0),
+        dv_fp32.stride(1),
+        dv_fp32.stride(2),
+        dv_fp32.stride(3),
+        stride_ms,
+        stride_mn,
+        Sq,
+        Sk,
+        float(scale),
+        HEAD_Q=HQ,
+        HEAD_K=HK,
+        SWA_WINDOW=swa_window_constexpr,
+        HAS_SINK=has_sink,
+        HAS_ADD_MASK=has_add_mask,
+        HCA_LOCAL_SEQLEN=hca_local_seqlen,
+        USE_CAUSAL=use_causal,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=8,
+        num_stages=1,
+    )
+
+    # Cast fp32 buffers back to input dtype.
+    dq_out = dq_fp32.to(q.dtype)
+    dk_out = dk_fp32.to(k.dtype)
+    dv_out = dv_fp32.to(v.dtype)
+    dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+    return dq_out, dk_out, dv_out, dsink_out
+
+
+__all__ = [
+    "_v4_attention_bwd_preprocess_kernel",
+    "_v4_attention_bwd_kernel",
+    "_launch_v4_attention_bwd",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py
new file mode 100644
index 000000000..4d1bf440e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py
@@ -0,0 +1,494 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention forward Triton kernel (plan-4 P25, ``compress_ratio in {0, 128}``).
+
+FlashAttention-style block-wise online softmax that handles V4's exact
+shape envelope:
+
+* ``head_dim = 512`` (single tile — no partial-RoPE / NOPE split because
+  RoPE is applied outside the kernel);
+* MQA single-latent KV (``K.shape[1] == 1``) broadcast across query
+  heads, *or* full MHA (``K.shape[1] == H``);
+* optional per-head learned softmax sink (``[H]``), joined as a virtual
+  key column with notional value zero at the end of the K-loop;
+* optional sliding-window-causal mask (``swa_window > 0``) applied
+  in-kernel;
+* optional caller-supplied additive bias. The generic path accepts
+  ``[Sq, Sk]`` and ignores ``swa_window``. The HCA split-mask path
+  accepts pool-only ``[Sq, P]`` with ``HCA_LOCAL_SEQLEN=Sq`` and keeps
+  the local branch on kernel-native SWA.
+
+dtype contract (matches :func:`eager_v4_attention`):
+
+* Q / K / V matmuls run in input dtype on tensor cores; accumulators
+  inside the matmul are fp32.
+* The online-softmax accumulator (``m_running``, ``l_running``,
+  ``acc``) lives in fp32 — this is the *only* fp32 step inside the
+  kernel.
+* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD
+  re-materialises ``P`` from it).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Triton kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_attention_fwd_kernel(
+    Q,
+    K,
+    V,
+    OUT,
+    LSE,
+    SINK,
+    ADD_MASK,
+    # Q strides: [B, H, Sq, D] row-major (contiguous on D)
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    # K strides: [B, K_H, Sk, D] row-major (K_H == 1 for MQA, == H for MHA)
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    # V strides: [B, K_H, Sk, D] row-major
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    # OUT strides: [B, H, Sq, D] row-major
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    # LSE strides: [B, H, Sq] row-major
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    # ADD_MASK strides: [Sq, Sk] row-major (broadcasts over B, H)
+    stride_ms,
+    stride_mn,
+    seqlen_q,
+    seqlen_k,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,  # 1 for MQA, == HEAD_Q for MHA
+    SWA_WINDOW: tl.constexpr,  # 0 = off, > 0 = SWA window
+    HAS_SINK: tl.constexpr,
+    HAS_ADD_MASK: tl.constexpr,
+    HCA_LOCAL_SEQLEN: tl.constexpr,  # 0 = generic mask; >0 = [local SWA keys | pool keys]
+    USE_CAUSAL: tl.constexpr,  # only meaningful when HAS_ADD_MASK = False
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """V4 attention FWD.
+
+    Grid layout: ``(cdiv(seqlen_q, BLOCK_M), batch * HEAD_Q)``. Each
+    program (program-id) computes one ``[BLOCK_M, BLOCK_DMODEL]`` slice
+    of ``OUT`` and the matching ``[BLOCK_M]`` slice of ``LSE``.
+
+    Mask precedence (must match :func:`eager_v4_attention`):
+
+    1. ``HAS_ADD_MASK`` — load the ``[Sq, Sk]`` additive bias and add
+       to ``qk``. SWA / causal masks are NOT applied in-kernel (the
+       caller has embedded all structure in the bias).
+    2. ``SWA_WINDOW > 0`` — sliding-window causal:
+       keep ``offs_n in [offs_m - SWA_WINDOW + 1, offs_m]``.
+    3. ``USE_CAUSAL`` — full causal: keep ``offs_n <= offs_m``.
+
+    In all branches, keys at index ``>= seqlen_k`` are masked to ``-inf``
+    so the boundary rows of the K-loop tile do not contaminate the
+    softmax denominator.
+    """
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+    if HEAD_K == HEAD_Q:
+        khid = qhid
+    else:
+        khid = 0  # MQA: single shared K / V head
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    # Q tile: [BLOCK_M, BLOCK_DMODEL] in q.dtype
+    q_ptrs = (
+        Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q_load_mask = offs_m[:, None] < seqlen_q
+    q = tl.load(q_ptrs, mask=q_load_mask, other=0.0)
+
+    # Online-softmax running state (fp32). We use the finite sentinel
+    # NEG_INF (-1e30) instead of -float("inf") so that the all-masked-
+    # tile corner case (m_running == m_tile == NEG_INF, e.g. for early
+    # queries under SWA when n_start is entirely outside the window)
+    # does NOT produce NaN through ``exp(-inf - -inf) = exp(NaN)``.
+    # NEG_INF is far enough below any plausible logit that
+    # ``exp(NEG_INF) ≈ 0`` and the algebra is identical to using -inf.
+    NEG_INF: tl.constexpr = -1.0e30
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full([BLOCK_M], value=NEG_INF, dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+
+    # Determine k-loop bounds.
+    #
+    # Plan-5 P30: for SWA, skip K tiles that are guaranteed outside the
+    # sliding window for *every* row in this M block.  The old P25 loop
+    # started at zero and relied on the mask below, which means late
+    # rows at S=4096 spent most of their time multiplying all-masked
+    # tiles.  Rounding down to the BLOCK_N boundary preserves exact
+    # masking semantics while cutting the steady-state SWA tile count
+    # from O(S / BLOCK_N) to O(window / BLOCK_N).
+    n_loop_start = 0
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+        # Caller's additive_mask handles all masking; iterate the full
+        # key axis. (We still apply boundary mask for keys past seqlen_k.)
+        n_loop_end = seqlen_k
+    elif SWA_WINDOW > 0:
+        # This block's earliest row is pid_m * BLOCK_M.  Any key before
+        # earliest_row - SWA_WINDOW + 1 is invisible for every row in the
+        # block, so skip those tiles entirely.
+        n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1
+        if n_loop_start < 0:
+            n_loop_start = 0
+        n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N
+        # Keys with n > max(offs_m) are causal-masked for every row.
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k
+        if n_loop_end > local_end:
+            n_loop_end = local_end
+    elif USE_CAUSAL:
+        # Full causal: keys with n > max(offs_m) are -inf, so the loop
+        # only needs to cover keys up to (pid_m + 1) * BLOCK_M.
+        n_loop_end = (pid_m + 1) * BLOCK_M
+        if n_loop_end > seqlen_k:
+            n_loop_end = seqlen_k
+    else:
+        n_loop_end = seqlen_k
+
+    # K-loop: local branch. For HCA split-mask this covers only the
+    # pruned local SWA prefix; the pool suffix is handled by the second
+    # loop below.
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        # K tile: [BLOCK_N, BLOCK_DMODEL] in k.dtype
+        k_ptrs = (
+            K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
+        )
+        k_load_mask = offs_n[:, None] < seqlen_k
+        k = tl.load(k_ptrs, mask=k_load_mask, other=0.0)
+
+        # qk = q @ k.T : [BLOCK_M, BLOCK_N] in fp32 (tl.dot accumulator)
+        qk = tl.dot(q, tl.trans(k)) * sm_scale
+
+        # Mask: additive_mask OR SWA / causal (mutually exclusive — see
+        # the docstring's precedence rule).
+        if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0:
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+        else:
+            if SWA_WINDOW > 0:
+                # offs_n in [offs_m - SWA_WINDOW + 1, offs_m]
+                in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & (
+                    offs_n[None, :] <= offs_m[:, None]
+                )
+                qk = tl.where(in_window, qk, NEG_INF)
+            elif USE_CAUSAL:
+                qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF)
+
+        # Boundary: keys past seqlen_k were loaded as 0; mask them to
+        # NEG_INF (finite sentinel) so they do not contribute to the
+        # softmax denominator and do NOT produce NaN through
+        # ``exp(-inf - -inf)``.
+        qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+
+        # Online softmax update
+        m_ij = tl.max(qk, 1)
+        m_new = tl.maximum(m_i, m_ij)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, 1)
+
+        # V tile: [BLOCK_N, BLOCK_DMODEL] in v.dtype
+        v_ptrs = (
+            V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
+        )
+        v_load_mask = offs_n[:, None] < seqlen_k
+        v = tl.load(v_ptrs, mask=v_load_mask, other=0.0)
+
+        # P57 R2: fuse ``acc * alpha`` into the V-dot MFMA-acc input.
+        # AMD's Triton lowering threads ``acc`` directly into the MFMA
+        # C-tile, eliminating one fp32 register round-trip per K-tile.
+        # Same pattern as the cr=0 BWD's ``dq = tl.dot(ds, k, acc=dq)``.
+        acc = tl.dot(p.to(v.dtype), v, acc=acc * alpha[:, None])
+        m_i = m_new
+
+    # HCA split-mask pool branch. The local branch above prunes the SWA
+    # prefix; the pool suffix is short (P = S / 128 for V4-Flash) and
+    # uses the caller-provided pool-only visibility mask [Sq, P].
+    if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0:
+        for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N):
+            offs_n = n_start + tl.arange(0, BLOCK_N)
+            pool_n = offs_n - HCA_LOCAL_SEQLEN
+
+            k_ptrs = (
+                K
+                + bid * stride_kb
+                + khid * stride_kh
+                + offs_n[:, None] * stride_kn
+                + offs_d[None, :] * stride_kd
+            )
+            k_load_mask = offs_n[:, None] < seqlen_k
+            k = tl.load(k_ptrs, mask=k_load_mask, other=0.0)
+
+            qk = tl.dot(q, tl.trans(k)) * sm_scale
+
+            mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn
+            mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)
+            add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32)
+            qk = qk + add_bias
+            qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF)
+
+            m_ij = tl.max(qk, 1)
+            m_new = tl.maximum(m_i, m_ij)
+            alpha = tl.exp(m_i - m_new)
+            p = tl.exp(qk - m_new[:, None])
+            l_i = l_i * alpha + tl.sum(p, 1)
+
+            v_ptrs = (
+                V
+                + bid * stride_vb
+                + khid * stride_vh
+                + offs_n[:, None] * stride_vn
+                + offs_d[None, :] * stride_vd
+            )
+            v_load_mask = offs_n[:, None] < seqlen_k
+            v = tl.load(v_ptrs, mask=v_load_mask, other=0.0)
+
+            # P57 R2: MFMA-acc fusion -- see local-branch comment above.
+            acc = tl.dot(p.to(v.dtype), v, acc=acc * alpha[:, None])
+            m_i = m_new
+
+    # Sink: virtual key column with value 0. The max-subtract trick
+    # uses sink as a candidate row maximum; sink contributes to
+    # l_running but NOT to acc (its notional value is zero).
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        m_new = tl.maximum(m_i, sink_h)
+        alpha = tl.exp(m_i - m_new)
+        beta = tl.exp(sink_h - m_new)
+        l_i = l_i * alpha + beta
+        acc = acc * alpha[:, None]
+        m_i = m_new
+
+    # Final divide + cast back to output dtype.
+    out = acc / l_i[:, None]
+    lse = m_i + tl.log(l_i)
+
+    out_ptrs = (
+        OUT + bid * stride_ob + qhid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
+    )
+    tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=offs_m[:, None] < seqlen_q)
+
+    lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+    tl.store(lse_ptrs, lse, mask=offs_m < seqlen_q)
+
+
+# ---------------------------------------------------------------------------
+# Python launcher
+# ---------------------------------------------------------------------------
+
+
+def _launch_v4_attention_fwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k: torch.Tensor,  # [B, K_H, Sk, D]   K_H ∈ {1, H}
+    v: torch.Tensor,  # [B, K_H, Sk, D]
+    *,
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    additive_mask: Optional[torch.Tensor],  # [Sq, Sk] or None
+    scale: float,
+    hca_local_seqlen: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch the V4 attention forward kernel.
+
+    Returns ``(out, lse)`` where ``out`` matches ``v.dtype`` and
+    ``lse`` is fp32. ``lse`` is what the BWD kernel needs to
+    re-materialise the softmax without storing the ``[Sq, Sk]`` ``P``
+    matrix.
+
+    The launcher does NO autograd bookkeeping — it's a thin wrapper
+    around the kernel suitable for the
+    :class:`V4AttentionFn` autograd Function and for unit tests.
+    """
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(
+            "v4_attention_v1 forward expects q / k / v of rank 4 "
+            f"(got q.dim={q.dim()}, k.dim={k.dim()}, v.dim={v.dim()})"
+        )
+    B, HQ, Sq, D = q.shape
+    Bk, HK, Sk, Dk = k.shape
+    Bv, HKv, Skv, Dv = v.shape
+    if (Bk, Sk, Dk) != (B, Sk, D) or (Bv, HKv, Skv, Dv) != (Bk, HK, Sk, D):
+        raise ValueError(
+            "v4_attention_v1 shape mismatch: " f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}"
+        )
+    if HK != 1 and HK != HQ:
+        raise ValueError(f"v4_attention_v1 requires K_H ∈ {{1 (MQA), {HQ} (MHA)}}; got K_H={HK}.")
+    if not q.is_cuda or not k.is_cuda or not v.is_cuda:
+        raise ValueError("v4_attention_v1 requires CUDA / HIP tensors.")
+    if q.dtype != k.dtype or q.dtype != v.dtype:
+        raise ValueError(
+            "v4_attention_v1 requires q.dtype == k.dtype == v.dtype "
+            f"(got {q.dtype} / {k.dtype} / {v.dtype})."
+        )
+
+    has_sink = sink is not None
+    has_add_mask = additive_mask is not None
+    hca_local_seqlen = int(hca_local_seqlen)
+    if hca_local_seqlen:
+        if not has_add_mask:
+            raise ValueError("hca_local_seqlen requires a pool additive_mask.")
+        if hca_local_seqlen <= 0 or hca_local_seqlen >= Sk:
+            raise ValueError(
+                "hca_local_seqlen must split local and pool keys "
+                f"(got hca_local_seqlen={hca_local_seqlen}, Sk={Sk})."
+            )
+        expected_mask_shape = (Sq, Sk - hca_local_seqlen)
+        if tuple(additive_mask.shape) != expected_mask_shape:
+            raise ValueError(
+                "HCA split-mask mode expects additive_mask shape "
+                f"{expected_mask_shape}, got {tuple(additive_mask.shape)}."
+            )
+        if swa_window <= 0:
+            raise ValueError("HCA split-mask mode requires swa_window > 0.")
+
+    # Mask precedence: additive_mask wins over swa_window. When neither
+    # is set, USE_CAUSAL = True (eager / V4 default).
+    use_causal = (not has_add_mask) and (swa_window <= 0)
+    swa_window_constexpr = (
+        int(swa_window) if ((not has_add_mask or hca_local_seqlen) and swa_window > 0) else 0
+    )
+
+    out = torch.empty_like(q)
+    lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    # P57 R2: tile sweep on V4-Flash widths (B=1, H=64, S=4096, D=512,
+    # SWA=128) finds (BLOCK_M=64, BLOCK_N=16, num_warps=8, num_stages=2)
+    # wins by **~17 %** on cr=4 FWD (1.69 -> 1.43 ms) and **~38 %** on
+    # cr=0 FWD (0.79 -> 0.49 ms). The intuition:
+    #
+    # * BLOCK_M=64 halves the program grid (vs the R1 32x32 layout) so
+    #   the SWA K-tile prefix is shared across twice as many query rows.
+    #   Q is loaded once per program -- bigger M tiles amortise the
+    #   Q-load over more K-tile iterations.
+    # * BLOCK_N=16 keeps the per-stage K/V tile at 16x512x2 = 16 KiB so
+    #   num_stages=2 double-buffering fits the LDS budget:
+    #   64 (Q) + 16*2 (K) + 16*2 (V) = 128 KiB << 160 KiB MI355X limit.
+    #   The R1 32x32 layout already needed 96 KiB at num_stages=1 and
+    #   would have hit 160 KiB at num_stages=2 -- no headroom.
+    # * num_stages=2 software-pipelines K/V loads against the
+    #   ``tl.dot`` + softmax update, hiding the HBM gather latency
+    #   that previously serialised the inner loop.
+    #
+    # Env-overridable so future shape regressions can fall back without
+    # rebuilding. The defaults are the per-shape winner from the R2
+    # sweep (`p57/r2_sweep_local.sh`).
+    # Arch-aware defaults (see ._v4_attn_tuning); env knobs still override.
+    # gfx950/MI355X: HCA fwd wins at BLOCK_M=128, pure SWA wants 64 (regresses at 128).
+    # gfx1250: SWA wins at BM=128/BN=32/W4/S1 (+62-70% vs gfx950; ab_sweep/opt7c).
+    from ._v4_attn_tuning import fwd_attn_defaults
+
+    _bm, _bn, _w, _s = fwd_attn_defaults(is_hca=bool(hca_local_seqlen))
+    BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_M", str(_bm)))
+    BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_N", str(_bn)))
+    NUM_WARPS_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_WARPS", str(_w)))
+    NUM_STAGES_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_STAGES", str(_s)))
+    BLOCK_DMODEL = D  # head_dim must be a power of 2 for tl.dot
+
+    grid = (triton.cdiv(Sq, BLOCK_M), B * HQ)
+
+    # Sentinel pointers when sink / mask are absent. Triton requires a
+    # real tensor — we pass q (any tensor) and gate via the constexpr.
+    sink_ptr = sink if has_sink else q
+    mask_ptr = additive_mask if has_add_mask else q
+    if has_add_mask:
+        stride_ms = additive_mask.stride(0)
+        stride_mn = additive_mask.stride(1)
+    else:
+        stride_ms = 0
+        stride_mn = 0
+
+    _v4_attention_fwd_kernel[grid](
+        q,
+        k,
+        v,
+        out,
+        lse,
+        sink_ptr,
+        mask_ptr,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        k.stride(3),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        v.stride(3),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        stride_ms,
+        stride_mn,
+        Sq,
+        Sk,
+        float(scale),
+        HEAD_Q=HQ,
+        HEAD_K=HK,
+        SWA_WINDOW=swa_window_constexpr,
+        HAS_SINK=has_sink,
+        HAS_ADD_MASK=has_add_mask,
+        HCA_LOCAL_SEQLEN=hca_local_seqlen,
+        USE_CAUSAL=use_causal,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=NUM_WARPS_FWD,
+        num_stages=NUM_STAGES_FWD,
+    )
+    return out, lse
+
+
+__all__ = [
+    "_v4_attention_fwd_kernel",
+    "_launch_v4_attention_fwd",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py
new file mode 100644
index 000000000..8310d707a
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py
@@ -0,0 +1,175 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA (pool) attention — Triton **v1** autograd entry point (production, cr=4).
+
+In-kernel top-k gather + pool scatter-add. Moved from the former top-level
+``v4_csa_attention_v0.py`` during the triton v0/v1/v2 reorg; pairs with the
+dense/HCA ``v4_attention_v1`` (also v1). See ``_triton_v0_deprecated`` for the deprecated
+gathered path and ``_triton_v2`` for the fused sparse-MLA path.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention import (
+    v4_attention_v1,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_bwd import (
+    _launch_v4_csa_attention_pool_bwd,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_fwd import (
+    _launch_v4_csa_attention_pool_fwd,
+)
+
+
+class V4CSAPoolAttentionFn(torch.autograd.Function):
+    """Triton CSA attention with in-kernel topk gather and pool scatter-add."""
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q: torch.Tensor,
+        k_local: torch.Tensor,
+        v_local: torch.Tensor,
+        pool: torch.Tensor,
+        topk_idxs: torch.Tensor,
+        sink: Optional[torch.Tensor],
+        swa_window: int,
+        attn_dropout: float,
+        training: bool,
+        scale: float,
+    ) -> torch.Tensor:
+        if attn_dropout > 0.0 and training:
+            raise NotImplementedError(
+                "v4_csa_attention_v1 does not implement in-kernel attention "
+                "dropout (V4 trains with attn_dropout=0). Got "
+                f"attn_dropout={attn_dropout}, training={training}."
+            )
+
+        out, lse = _launch_v4_csa_attention_pool_fwd(
+            q,
+            k_local,
+            v_local,
+            pool,
+            topk_idxs,
+            sink=sink,
+            swa_window=swa_window,
+            scale=scale,
+        )
+        ctx.save_for_backward(q, k_local, v_local, pool, topk_idxs, out, lse, sink)
+        ctx.swa_window = int(swa_window)
+        ctx.attn_dropout = float(attn_dropout)
+        ctx.training_mode = bool(training)
+        ctx.scale = float(scale)
+        ctx.sink_was_none = sink is None
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):  # type: ignore[override]
+        q, k_local, v_local, pool, topk_idxs, out, lse, sink = ctx.saved_tensors
+        sink_arg = None if ctx.sink_was_none else sink
+
+        if not grad_out.is_contiguous():
+            grad_out = grad_out.contiguous()
+
+        dq, dk_local, dv_local, dpool, dsink = _launch_v4_csa_attention_pool_bwd(
+            q,
+            k_local,
+            v_local,
+            pool,
+            topk_idxs,
+            out,
+            grad_out,
+            lse,
+            sink=sink_arg,
+            swa_window=ctx.swa_window,
+            scale=ctx.scale,
+        )
+
+        if not ctx.needs_input_grad[0]:
+            dq = None
+        if not ctx.needs_input_grad[1]:
+            dk_local = None
+        if not ctx.needs_input_grad[2]:
+            dv_local = None
+        if not ctx.needs_input_grad[3]:
+            dpool = None
+        if not ctx.needs_input_grad[5] or ctx.sink_was_none:
+            dsink = None
+
+        # Forward signature: (q, k_local, v_local, pool, topk_idxs, sink,
+        # swa_window, attn_dropout, training, scale).
+        return dq, dk_local, dv_local, dpool, None, dsink, None, None, None, None
+
+
+def v4_csa_attention_v1(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    pool: torch.Tensor,  # [B, P, D]
+    *,
+    topk_idxs: torch.Tensor,  # [B, Sq, K_topk], -1 masks a slot
+    sink: Optional[torch.Tensor],  # [H] or None
+    swa_window: int,
+    attn_dropout: float,
+    training: bool,
+    scale: float,
+    use_tilelang: bool = False,
+    use_flydsl: bool = False,  # accepted for call-site parity; no from-pool FlyDSL kernel, Triton path
+) -> torch.Tensor:
+    """Triton-backed CSA attention that gathers sparse keys in-kernel.
+
+    ``pool`` is the compressed-pool tensor before per-query top-K gather.
+    ``topk_idxs`` drives the sparse branch directly; negative entries are
+    masked and contribute no probability mass. The backward kernel emits
+    ``dpool`` with atomic scatter-add, avoiding the materialised
+    ``[B, Sq, K_topk, D]`` gathered tensor and its autograd scatter.
+
+    ``use_tilelang`` is reserved for the plan-8 P54 / P55 from-pool
+    tilelang path (not landed); currently always falls through to
+    :class:`V4CSAPoolAttentionFn` (Triton).
+    """
+    K_topk = topk_idxs.shape[2]
+    if K_topk == 0:
+        return v4_attention_v1(
+            q,
+            k_local,
+            v_local,
+            sink=sink,
+            swa_window=swa_window,
+            additive_mask=None,
+            attn_dropout=attn_dropout,
+            training=training,
+            scale=scale,
+        )
+
+    # Plan-8 P57 close-out 2: from-pool tilelang path is not landed
+    # (P54 / P55 descoped); the gated dispatcher always returns False
+    # so this stays on the Triton autograd path.  We still consult
+    # the dispatcher so a future P54 / P55 landing can flip behavior
+    # without re-wiring this callsite.
+    del use_tilelang
+    return V4CSAPoolAttentionFn.apply(
+        q,
+        k_local,
+        v_local,
+        pool,
+        topk_idxs,
+        sink,
+        swa_window,
+        attn_dropout,
+        training,
+        scale,
+    )
+
+
+__all__ = [
+    "V4CSAPoolAttentionFn",
+    "v4_csa_attention_v1",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py
new file mode 100644
index 000000000..030e35f76
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py
@@ -0,0 +1,2321 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention backward Triton kernel (plan-4 P26, ``compress_ratio == 4``).
+
+Two-kernel design (mirroring :mod:`v4_attention_bwd`):
+
+* Pre-pass: ``D[b, h, m] = sum_d (dout[b,h,m,d] * out[b,h,m,d])`` —
+  reuses :func:`_v4_attention_bwd_preprocess_kernel` from the dense
+  module since the contract is identical.
+* Main pass: one program per ``(b, qhid, m)`` query row; re-materialises
+  the joint softmax row from the saved LSE; emits
+
+  ::
+    dq         [B, H, Sq, D]      direct store (one program per row)
+    dk_local   [B, H, Sq, D]      atomic-add (multiple m's hit same n)
+    dv_local   [B, H, Sq, D]      atomic-add
+    dgathered  [B, Sq, K_topk, D] atomic-add (no H dim — broadcast in fwd
+                                  means all H heads contribute)
+    dsink      [H]                atomic-add per query
+
+dtype contract:
+
+* All inputs loaded in input dtype; per-row dot products reduce in fp32
+  via ``.to(tl.float32)`` upcast before the multiply (matches the FWD's
+  bf16-tensor-core / fp32-accumulator semantics).
+* The online ``P / dP / dS`` re-materialisation is fp32 (matches the
+  FWD's softmax-in-fp32 contract).
+* Output gradients are returned in input dtype (cast from fp32 buffers
+  by the launcher).
+
+Math derivation (per query (b, h, m), see plan-4 ``02-phase-details.md``
+Phase 26 section):
+
+  joint_logits = cat(qk_local, qk_sparse, sink_h)
+  P_j = exp(joint_logits[j] - lse)
+  out_d = sum_n P_local[n] * v_local[n,d] + sum_k P_sparse[k] * g[k,d]
+  D = sum_d (dout[d] * out[d])
+  dP_local[n]  = sum_d (dout[d] * v_local[n,d])
+  dP_sparse[k] = sum_d (dout[d] * g[k,d])
+  dS_local[n]  = P_local[n]  * (dP_local[n]  - D)
+  dS_sparse[k] = P_sparse[k] * (dP_sparse[k] - D)
+  dS_sink      = -P_sink * D                                    # sink val is 0
+
+  dq[d]               = sum_n dS_local[n]  * scale * k_local[n,d]
+                      + sum_k dS_sparse[k] * scale * g[k,d]
+  dk_local[n,d]      += dS_local[n]  * scale * q[d]
+  dv_local[n,d]      += P_local[n]   * dout[d]
+  dgathered[k,d]     += dS_sparse[k] * scale * q[d]
+                      + P_sparse[k]  * dout[d]                   # both branches
+  dsink_h            += dS_sink
+
+Edge cases:
+
+* ``K_topk == 0`` — the wrapper short-circuits to the dense
+  :func:`v4_attention_v1` BWD before reaching this kernel.
+* All-masked tile rows — uses ``NEG_INF = -1e30`` finite sentinel so
+  ``exp(NEG_INF - lse) = exp(-large) ≈ 0`` for fully-masked positions
+  (matches the FWD).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import (
+    _v4_attention_bwd_dkv_kernel,
+    _v4_attention_bwd_dq_kernel,
+    _v4_attention_bwd_kernel,
+    _v4_attention_bwd_preprocess_kernel,
+)
+
+# ---------------------------------------------------------------------------
+# Main BWD kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_csa_attention_pool_bwd_kernel(
+    Q,
+    K_LOCAL,
+    V_LOCAL,
+    POOL,
+    TOPK_IDXS,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    DK_LOCAL,
+    DV_LOCAL,
+    DPOOL,
+    DSINK,
+    SINK,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_klb,
+    stride_klh,
+    stride_kln,
+    stride_kld,
+    stride_vlb,
+    stride_vlh,
+    stride_vln,
+    stride_vld,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dklb,
+    stride_dklh,
+    stride_dkln,
+    stride_dkld,
+    stride_dvlb,
+    stride_dvlh,
+    stride_dvln,
+    stride_dvld,
+    stride_dpb,
+    stride_dpp,
+    stride_dpd,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    STORE_DPOOL: tl.constexpr,
+    LOCAL_ONLY: tl.constexpr,
+):
+    """CSA BWD with in-kernel scatter-add into the compressed-pool gradient."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    NEG_INF: tl.constexpr = -1.0e30
+    q_active = pid_m < seqlen_q
+
+    q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm
+    q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0)
+
+    do_row_offset = bid * stride_dob + qhid * stride_doh + pid_m * stride_dom
+    dout = tl.load(DOUT + do_row_offset + offs_d * stride_dod, mask=q_active, other=0.0)
+    q_f = q.to(tl.float32)
+    dout_f = dout.to(tl.float32)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm,
+        mask=q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + qhid * stride_dh + pid_m * stride_dm,
+        mask=q_active,
+        other=0.0,
+    )
+
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        p_sink = tl.exp(sink_h - lse)
+        tl.atomic_add(DSINK + qhid, tl.where(q_active, -p_sink * dvec, 0.0))
+
+    dq = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    n_loop_end = pid_m + 1
+    if n_loop_end > seqlen_q:
+        n_loop_end = seqlen_q
+    if SWA_WINDOW > 0:
+        n_lo_raw = pid_m - SWA_WINDOW + 1
+        if n_lo_raw < 0:
+            n_lo_raw = 0
+        n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N
+    else:
+        n_loop_start = 0
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+
+        kl_ptrs = (
+            K_LOCAL
+            + bid * stride_klb
+            + qhid * stride_klh
+            + offs_n[:, None] * stride_kln
+            + offs_d[None, :] * stride_kld
+        )
+        kl_load_mask = offs_n[:, None] < seqlen_q
+        kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0)
+
+        vl_ptrs = (
+            V_LOCAL
+            + bid * stride_vlb
+            + qhid * stride_vlh
+            + offs_n[:, None] * stride_vln
+            + offs_d[None, :] * stride_vld
+        )
+        vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0)
+
+        kl_f = kl.to(tl.float32)
+        qk = tl.sum(kl_f * q_f[None, :], axis=1) * sm_scale
+        if SWA_WINDOW > 0:
+            in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m)
+        else:
+            in_window = offs_n <= pid_m
+        qk = tl.where(in_window & (offs_n < seqlen_q) & q_active, qk, NEG_INF)
+
+        p = tl.exp(qk - lse)
+        vl_f = vl.to(tl.float32)
+        dp = tl.sum(dout_f[None, :] * vl_f, axis=1)
+        ds = p * (dp - dvec)
+
+        dq += tl.sum(ds[:, None] * kl_f, axis=0) * sm_scale
+
+        dk_contrib = ds[:, None] * sm_scale * q_f[None, :]
+        dk_ptrs = (
+            DK_LOCAL
+            + bid * stride_dklb
+            + qhid * stride_dklh
+            + offs_n[:, None] * stride_dkln
+            + offs_d[None, :] * stride_dkld
+        )
+        tl.atomic_add(dk_ptrs, dk_contrib, mask=kl_load_mask, sem="relaxed")
+
+        dv_contrib = p[:, None] * dout_f[None, :]
+        dv_ptrs = (
+            DV_LOCAL
+            + bid * stride_dvlb
+            + qhid * stride_dvlh
+            + offs_n[:, None] * stride_dvln
+            + offs_d[None, :] * stride_dvld
+        )
+        tl.atomic_add(dv_ptrs, dv_contrib, mask=kl_load_mask, sem="relaxed")
+
+    if not LOCAL_ONLY:
+        for k_start in range(0, K_topk, BLOCK_K):
+            offs_k = k_start + tl.arange(0, BLOCK_K)
+            topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + offs_k * stride_tik
+            topk = tl.load(topk_ptrs, mask=offs_k < K_topk, other=-1)
+            valid = (offs_k < K_topk) & (topk >= 0) & (topk < pool_size)
+            safe_topk = tl.where(valid, topk, 0)
+
+            pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+            pool = tl.load(pool_ptrs, mask=valid[:, None], other=0.0)
+            pool_f = pool.to(tl.float32)
+
+            qk_sparse = tl.sum(pool_f * q_f[None, :], axis=1) * sm_scale
+            qk_sparse = tl.where(valid & q_active, qk_sparse, NEG_INF)
+
+            p = tl.exp(qk_sparse - lse)
+            dp = tl.sum(dout_f[None, :] * pool_f, axis=1)
+            ds = p * (dp - dvec)
+
+            dq += tl.sum(ds[:, None] * pool_f, axis=0) * sm_scale
+
+            if STORE_DPOOL:
+                dpool_contrib = ds[:, None] * sm_scale * q_f[None, :] + p[:, None] * dout_f[None, :]
+                dpool_ptrs = (
+                    DPOOL + bid * stride_dpb + safe_topk[:, None] * stride_dpp + offs_d[None, :] * stride_dpd
+                )
+                tl.atomic_add(dpool_ptrs, dpool_contrib, mask=valid[:, None], sem="relaxed")
+
+    dq_offset = bid * stride_dqb + qhid * stride_dqh + pid_m * stride_dqm
+    tl.store(DQ + dq_offset + offs_d * stride_dqd, dq, mask=q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_bwd_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    DPOOL,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dp_part,
+    stride_dpb,
+    stride_dpp,
+    stride_dpd,
+    partition_size,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    STORE_DPOOL: tl.constexpr = True,
+):
+    """Sparse CSA BWD using a head block so pool work maps to tl.dot.
+
+    Plan-5 P32: ``DPOOL`` is laid out as ``[N_PART, B, P, D]`` instead
+    of ``[B, P, D]``. Each program writes its dpool contributions into
+    ``DPOOL[pid_m // partition_size, bid, ...]``. With ``N_PART = Sq /
+    partition_size``, the per-cache-line atomic contention drops by a
+    factor of ``N_PART`` because writes from different m-partitions hit
+    different DRAM rows. The launcher reduces the partial axis on the
+    Python side, which is bandwidth-cheap (``2 * N_PART * P * D *
+    sizeof(fp32)``).
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+    partition_id = pid_m // partition_size
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + offs_h[:, None] * stride_doh
+        + pid_m * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+
+    dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        q_bf16 = q.to(pool.dtype)
+        dout_bf16 = dout.to(pool.dtype)
+        qk = tl.dot(q_bf16, tl.trans(pool)) * sm_scale
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_bf16, tl.trans(pool))
+        ds = p * (dp - dvec[:, None])
+
+        dq += tl.dot(ds.to(pool.dtype), pool) * sm_scale
+
+        if STORE_DPOOL:
+            dpool_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q_bf16) * sm_scale
+            dpool_contrib += tl.dot(tl.trans(p.to(dout.dtype)), dout_bf16)
+            dpool_ptrs = (
+                DPOOL
+                + partition_id * stride_dp_part
+                + bid * stride_dpb
+                + safe_topk[:, None] * stride_dpp
+                + offs_d[None, :] * stride_dpd
+            )
+            tl.atomic_add(dpool_ptrs, dpool_contrib, mask=valid_k[:, None], sem="relaxed")
+        else:
+            # Force the optimiser to keep this branch as a no-op so the
+            # ``STORE_DPOOL=False`` variant truly drops the two trailing
+            # ``tl.dot`` ops + atomic from the IR.
+            pass
+
+    # Plan-5 P32: each ``(pid_m, pid_h_block, bid)`` program writes to a
+    # disjoint slice of ``DQ`` (different ``m`` rows + different head
+    # blocks), so a plain ``tl.store`` is safe and ~2× cheaper than the
+    # atomic on MI355.
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + offs_h[:, None] * stride_dqh
+        + pid_m * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_bwd_dq_only_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """P57: dq-only sparse CSA BWD (no dpool_partial write).
+
+    Mirrors :func:`_v4_csa_attention_pool_sparse_bwd_partial_kernel`
+    but drops the per-visit ``dpool_partial`` write and the two
+    ``dpool_contrib`` matmuls. The freed register file lets us run
+    with a larger ``BLOCK_K`` / more ``num_stages`` so the per-program
+    latency drops sharply. The dpool partial is produced by a sibling
+    ``_v4_csa_attention_pool_sparse_bwd_dpool_only_kernel`` that
+    omits the ``dq`` accumulator instead, and the segreduce kernel
+    folds the partial into ``dpool[B, P, D]``.
+
+    The reason for splitting: the joint kernel's wall-clock at the
+    proxy shape is ~3.7 ms, dominated by the 4 GiB ``dpool_partial``
+    write *interleaved* with the ``[BLOCK_H, BLOCK_DMODEL]`` fp32
+    ``dq`` accumulator (~64 KB live across the K-loop). Splitting
+    drops the live-register footprint of each kernel ~2×, letting
+    the compiler use higher ``num_stages`` for prefetching.
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + offs_h[:, None] * stride_doh
+        + pid_m * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+
+    dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+
+    # P57 R2: scale-defer + acc-form MFMA (mirror of the joint
+    # partial kernel). See the comment in
+    # ``_v4_csa_attention_pool_sparse_bwd_partial_kernel`` for the
+    # math.
+    pool_dtype = POOL.dtype.element_ty
+    q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype)
+    dout_bf16 = dout.to(pool_dtype)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        qk = tl.dot(q_scaled, tl.trans(pool))
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_bf16, tl.trans(pool))
+        ds = p * (dp - dvec[:, None])
+
+        dq = tl.dot(ds.to(pool_dtype), pool, acc=dq)
+
+    dq = dq * sm_scale
+
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + offs_h[:, None] * stride_dqh
+        + pid_m * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    DOUT,
+    LSE,
+    D,
+    DPOOL_PARTIAL,  # [B, M, K_topk, D] fp32 — NO atomics, each program owns its slice
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dpb,
+    stride_dpm,
+    stride_dpk,
+    stride_dpd,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """P57: dpool-partial-only sparse CSA BWD (no dq accumulator).
+
+    Mirror of :func:`_v4_csa_attention_pool_sparse_bwd_dq_only_kernel`
+    that drops the ``dq`` accumulator (64 KB fp32 per program) so
+    the per-iter ``dpool_contrib`` matmul + write loop runs with a
+    tighter register footprint and can issue more in-flight HBM
+    writes per warp.
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + offs_h[:, None] * stride_doh
+        + pid_m * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+
+    # P57 R2: scale-defer (q-fold) — see partial kernel for math.
+    pool_dtype = POOL.dtype.element_ty
+    q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype)
+    dout_bf16 = dout.to(pool_dtype)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        qk = tl.dot(q_scaled, tl.trans(pool))
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_bf16, tl.trans(pool))
+        ds = p * (dp - dvec[:, None])
+
+        dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled)
+        dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib)
+        dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0)
+        dpool_partial_ptrs = (
+            DPOOL_PARTIAL
+            + bid * stride_dpb
+            + pid_m * stride_dpm
+            + sparse_k[:, None] * stride_dpk
+            + offs_d[None, :] * stride_dpd
+        )
+        tl.store(
+            dpool_partial_ptrs,
+            dpool_contrib,
+            mask=(sparse_k[:, None] < K_topk) & q_active,
+        )
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_bwd_partial_sorted_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    INV_PERM,  # [B, MK] int32 — sorted_position = inv_perm[orig_flat]
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    DPOOL_PARTIAL,  # [B, MK_sorted, D] partial buffer (bf16 / fp32)
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_ipb,
+    stride_ipi,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dpb,
+    stride_dpi,
+    stride_dpd,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """P57 sorted variant of the sparse CSA BWD partial kernel.
+
+    Writes the per-visit ``dpool_contrib`` tile into the sorted-order
+    position of ``DPOOL_PARTIAL`` (a flat ``[B, MK, D]`` buffer)
+    using a host-side ``INV_PERM`` (the inverse of the segreduce
+    ``perm``). The downstream segreduce kernel can then read
+    contiguous ``[bin_start:bin_end, :]`` slices instead of gathering
+    via ``perm[i]`` — turning a random-access reduction into a
+    streaming one (better L2 / Infinity Cache reuse on MI355).
+
+    The partial kernel writes are now scattered (one row per sparse_k
+    slot to a sorted position), but the HBM write bandwidth is
+    similar because Triton's vector store still coalesces 128 B
+    cache-line bursts as long as ``BLOCK_DMODEL`` is contiguous.
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + offs_h[:, None] * stride_doh
+        + pid_m * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+
+    dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+    flat_base = pid_m * K_topk
+
+    # P57 R2: scale-defer (q-fold) — see partial kernel for math.
+    pool_dtype = POOL.dtype.element_ty
+    q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype)
+    dout_bf16 = dout.to(pool_dtype)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        qk = tl.dot(q_scaled, tl.trans(pool))
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_bf16, tl.trans(pool))
+        ds = p * (dp - dvec[:, None])
+
+        dq = tl.dot(ds.to(pool_dtype), pool, acc=dq)
+
+        dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled)
+        dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib)
+        dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0)
+
+        # P57: look up the SORTED position for each ``(m, sparse_k)``
+        # visit. ``flat_base + sparse_k`` is the ORIGINAL flat index;
+        # ``inv_perm[bid, orig_flat]`` is the sorted slot to write.
+        orig_flat = flat_base + sparse_k
+        inv_perm_ptrs = INV_PERM + bid * stride_ipb + orig_flat * stride_ipi
+        sorted_idx = tl.load(inv_perm_ptrs, mask=sparse_k < K_topk, other=0)
+        dpool_partial_ptrs = (
+            DPOOL_PARTIAL + bid * stride_dpb + sorted_idx[:, None] * stride_dpi + offs_d[None, :] * stride_dpd
+        )
+        tl.store(
+            dpool_partial_ptrs,
+            dpool_contrib,
+            mask=(sparse_k[:, None] < K_topk) & q_active,
+        )
+
+    dq = dq * sm_scale
+
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + offs_h[:, None] * stride_dqh
+        + pid_m * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_segreduce_sequential_kernel(
+    DPOOL_PARTIAL,  # [B, MK_sorted, D] — written in sorted order by partial_sorted kernel
+    BIN_PTR,  # [B, P+1] int32 — prefix sum of count per pool slot (same as non-sorted variant)
+    DPOOL,  # [B, P, D]
+    stride_dpb,
+    stride_dpi,
+    stride_dpd,
+    stride_binb,
+    stride_binp,
+    stride_db,
+    stride_dp,
+    stride_dd,
+    P,
+    D_size,
+    BLOCK_D: tl.constexpr,
+    BLOCK_I: tl.constexpr,
+):
+    """P57 sequential segreduce — reads dpool_partial[bid, i:i+BI, :]
+    contiguously without a perm lookup since the partial kernel
+    already wrote in sorted order.
+    """
+    pid_p = tl.program_id(0)
+    pid_d_block = tl.program_id(1)
+    pid_b = tl.program_id(2)
+
+    offs_d = pid_d_block * BLOCK_D + tl.arange(0, BLOCK_D)
+    d_mask = offs_d < D_size
+
+    bin_start = tl.load(BIN_PTR + pid_b * stride_binb + pid_p * stride_binp)
+    bin_end = tl.load(BIN_PTR + pid_b * stride_binb + (pid_p + 1) * stride_binp)
+
+    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
+    i = bin_start
+    while i < bin_end:
+        offs_i = i + tl.arange(0, BLOCK_I)
+        valid_i = offs_i < bin_end
+        partial_ptrs = (
+            DPOOL_PARTIAL + pid_b * stride_dpb + offs_i[:, None] * stride_dpi + offs_d[None, :] * stride_dpd
+        )
+        partial = tl.load(
+            partial_ptrs,
+            mask=valid_i[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+        acc += tl.sum(partial, axis=0)
+        i += BLOCK_I
+
+    dpool_offset = pid_b * stride_db + pid_p * stride_dp + offs_d * stride_dd
+    tl.store(DPOOL + dpool_offset, acc, mask=d_mask)
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_bwd_partial_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    DOUT,
+    LSE,
+    D,
+    DQ,
+    DPOOL_PARTIAL,  # [B, M, K_topk, D] fp32 — NO atomics, each program owns its slice
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_dod,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_db,
+    stride_dh,
+    stride_dm,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_dqd,
+    stride_dpb,
+    stride_dpm,
+    stride_dpk,
+    stride_dpd,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """Sparse CSA BWD that writes dpool contributions to a compact
+    ``[B, M, K_topk, D]`` partial buffer **without** atomics.
+
+    Plan-5 P32: the atomic_add to the shared ``dpool[B, P, D]`` buffer
+    was the dominant cost (~15 ms out of ~24 ms — see
+    ``PRIMUS_V4_CSA_BWD_SKIP_DPOOL_ATOMIC`` profile). By emitting the
+    raw per-visit contributions into a compact partial buffer (4 GB
+    for the proxy shape, indexed by ``(b, m, k_slot)``), a follow-up
+    segmented reduction kernel can fold them into ``dpool[B, P, D]``
+    using a sorted inverse index — atomics-free and bandwidth-bound.
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    dout_ptrs = (
+        DOUT
+        + bid * stride_dob
+        + offs_h[:, None] * stride_doh
+        + pid_m * stride_dom
+        + offs_d[None, :] * stride_dod
+    )
+    dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    lse = tl.load(
+        LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+    dvec = tl.load(
+        D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm,
+        mask=h_mask & q_active,
+        other=0.0,
+    )
+
+    dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+
+    # P57 R2: hoist q/dout dtype-cast and SCALE-DEFER ``sm_scale`` out
+    # of the K-loop. Pre-scaling ``q_bf16`` folds the ``sm_scale``
+    # factor that the kernel previously applied to ``qk`` (line:
+    # ``tl.dot(q_bf16, K^T) * sm_scale``) and to the first
+    # ``dpool_contrib`` matmul (``tl.dot(ds^T, q_bf16) * sm_scale``)
+    # — both are linear in q so factoring sm_scale into q is exact.
+    # ``dq`` accumulates ``Σ_k ds @ pool`` without sm_scale; a single
+    # ``dq *= sm_scale`` after the loop replaces the per-iter
+    # multiply on the ``[BLOCK_H, BLOCK_DMODEL]`` fp32 accumulator.
+    # Combined with ``tl.dot(..., acc=dq)`` (MFMA in-place
+    # accumulator), this drops ~3 fp32 multiplies + 1 fp32 add per
+    # K_topk/BLOCK_K=16 iteration on the proxy shape.
+    pool_dtype = POOL.dtype.element_ty
+    dout_bf16 = dout.to(pool_dtype)
+    q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        qk = tl.dot(q_scaled, tl.trans(pool))
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30)
+
+        p = tl.exp(qk - lse[:, None])
+        dp = tl.dot(dout_bf16, tl.trans(pool))
+        ds = p * (dp - dvec[:, None])
+
+        dq = tl.dot(ds.to(pool_dtype), pool, acc=dq)
+
+        # Plan-5 P32: write the per-visit dpool contribution to its own
+        # compact slot in ``DPOOL_PARTIAL[b, m, k_slot, :]``. No atomic
+        # needed because each ``(b, m, k_slot)`` slot is owned by
+        # exactly one program × iteration.
+        # P57 R2: ``q_scaled`` already carries the ``sm_scale`` factor,
+        # so the first matmul's ``* sm_scale`` is folded in. The
+        # ``acc=`` form on the second matmul fuses the fp32 add.
+        dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled)
+        dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib)
+        # Zero out invalid k slots so the reduction can sum them in
+        # without first checking validity (the inverse index will skip
+        # invalid visits anyway via the sentinel sort key, but a clean
+        # buffer is friendlier to debugging and to scalar fallbacks).
+        dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0)
+        dpool_partial_ptrs = (
+            DPOOL_PARTIAL
+            + bid * stride_dpb
+            + pid_m * stride_dpm
+            + sparse_k[:, None] * stride_dpk
+            + offs_d[None, :] * stride_dpd
+        )
+        tl.store(
+            dpool_partial_ptrs,
+            dpool_contrib,
+            mask=(sparse_k[:, None] < K_topk) & q_active,
+        )
+
+    # P57 R2: deferred ``sm_scale`` on the dq accumulator. Single
+    # ``[BLOCK_H, BLOCK_DMODEL]`` fp32 multiply replaces the per-iter
+    # ``* sm_scale`` inside the loop.
+    dq = dq * sm_scale
+
+    dq_ptrs = (
+        DQ
+        + bid * stride_dqb
+        + offs_h[:, None] * stride_dqh
+        + pid_m * stride_dqm
+        + offs_d[None, :] * stride_dqd
+    )
+    tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_segreduce_kernel(
+    DPOOL_PARTIAL,  # [B, M, K_topk, D] fp32 — compact partial buffer
+    SORTED_PERM,  # [B, M*K_topk] int32 — sorted (m*K + k) per pool slot
+    BIN_PTR,  # [B, P+1] int32 — prefix sum of count per pool slot
+    DPOOL,  # [B, P, D] fp32 — output (single tl.store per p, d)
+    stride_dpb,
+    stride_dpmk,
+    stride_dpd,
+    stride_permb,
+    stride_permi,
+    stride_binb,
+    stride_binp,
+    stride_db,
+    stride_dp,
+    stride_dd,
+    P,
+    D_size,
+    BLOCK_D: tl.constexpr,
+    BLOCK_I: tl.constexpr,  # tiles over the visit indices to expose ILP
+):
+    """Segmented reduction: ``DPOOL[b, p, :] = Σ_i DPOOL_PARTIAL[b, perm[i], :]``
+    for ``i ∈ [bin_ptr[b, p], bin_ptr[b, p+1])``.
+
+    Plan-5 P32: one program per ``(b, p, d_block)``. Different
+    programs write disjoint output slices, so the writes are plain
+    ``tl.store``\\ s — no atomics. The visit indices for slot ``p``
+    are stored consecutively in ``SORTED_PERM`` (built once on the
+    host by sorting ``topk_idxs``), so loads from
+    ``DPOOL_PARTIAL[..., flat_idx, :]`` are reasonably coalescable
+    after the sort. ``BLOCK_D`` may exceed the actual head dim
+    ``D_size`` (the default proxy shape has ``D = 512`` but the unit
+    tests use ``D = 32``), so all dpool / dpool_partial accesses are
+    masked by ``offs_d < D_size``.
+    """
+    pid_p = tl.program_id(0)
+    pid_d_block = tl.program_id(1)
+    pid_b = tl.program_id(2)
+
+    offs_d = pid_d_block * BLOCK_D + tl.arange(0, BLOCK_D)
+    d_mask = offs_d < D_size
+
+    bin_start = tl.load(BIN_PTR + pid_b * stride_binb + pid_p * stride_binp)
+    bin_end = tl.load(BIN_PTR + pid_b * stride_binb + (pid_p + 1) * stride_binp)
+
+    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
+    i = bin_start
+    while i < bin_end:
+        # Load up to BLOCK_I visit indices and prefetch their fp32
+        # partial rows; the tl.dot-free vectorised path nets enough
+        # ILP to saturate HBM bandwidth on MI355.
+        offs_i = i + tl.arange(0, BLOCK_I)
+        valid_i = offs_i < bin_end
+        flat_idx = tl.load(
+            SORTED_PERM + pid_b * stride_permb + offs_i * stride_permi,
+            mask=valid_i,
+            other=0,
+        )
+        partial_ptrs = (
+            DPOOL_PARTIAL
+            + pid_b * stride_dpb
+            + flat_idx[:, None] * stride_dpmk
+            + offs_d[None, :] * stride_dpd
+        )
+        partial = tl.load(
+            partial_ptrs,
+            mask=valid_i[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+        acc += tl.sum(partial, axis=0)
+        i += BLOCK_I
+
+    dpool_offset = pid_b * stride_db + pid_p * stride_dp + offs_d * stride_dd
+    tl.store(DPOOL + dpool_offset, acc, mask=d_mask)
+
+
+# ---------------------------------------------------------------------------
+# Python launcher
+# ---------------------------------------------------------------------------
+
+
+def _launch_v4_csa_attention_pool_bwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    pool: torch.Tensor,  # [B, P, D]
+    topk_idxs: torch.Tensor,  # [B, Sq, K_topk]
+    out: torch.Tensor,  # [B, H, Sq, D]
+    dout: torch.Tensor,  # [B, H, Sq, D]
+    lse: torch.Tensor,  # [B, H, Sq]
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Launch CSA backward with in-kernel scatter-add into ``pool.grad``."""
+    if not q.is_cuda:
+        raise ValueError("v4_csa_attention_v0 pool BWD requires CUDA / HIP tensors.")
+    if dout.shape != out.shape or out.shape != q.shape:
+        raise ValueError(
+            "v4_csa_attention_v0 pool BWD shape mismatch: "
+            f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}"
+        )
+    if pool.dim() != 3:
+        raise ValueError(
+            f"v4_csa_attention_v0 pool BWD expects pool rank 3 [B, P, D], got {tuple(pool.shape)}."
+        )
+    if topk_idxs.dim() != 3:
+        raise ValueError(
+            f"v4_csa_attention_v0 pool BWD expects topk_idxs rank 3 [B, Sq, K], got {tuple(topk_idxs.shape)}."
+        )
+
+    B, HQ, Sq, D = q.shape
+    Bp, P, Dp = pool.shape
+    Bt, Sqt, K_topk = topk_idxs.shape
+    if Bp != B or Dp != D or Bt != B or Sqt != Sq:
+        raise ValueError(
+            "v4_csa_attention_v0 pool BWD shape mismatch: "
+            f"q={tuple(q.shape)}, pool={tuple(pool.shape)}, topk_idxs={tuple(topk_idxs.shape)}"
+        )
+    if topk_idxs.dtype not in (torch.int32, torch.int64):
+        raise ValueError(f"v4_csa_attention_v0 topk_idxs must be int32/int64, got {topk_idxs.dtype}.")
+
+    has_sink = sink is not None
+
+    BLOCK_N = 32
+    # Plan-5 P32: sweeping ``BLOCK_K ∈ {32, 64, 128}`` and ``num_warps
+    # ∈ {4, 8}`` on the EP8 microbench picked ``BLOCK_K=32`` /
+    # ``num_warps=4`` as the cheapest configuration (24.9 ms vs 26.0 ms
+    # at the previous defaults). Override via
+    # ``PRIMUS_V4_CSA_BWD_BLOCK_K`` for shape-specific tuning.
+    BLOCK_K = int(os.getenv("PRIMUS_V4_CSA_BWD_BLOCK_K", "32"))
+    # Plan-5 P32: the segreduce / partial path is markedly more
+    # register-pressured than the atomic-add gather kernel (the extra
+    # ``dpool_partial`` write tile pins 64 × 512 fp32 in flight per
+    # warp). A second sweep at ``BLOCK_K_PARTIAL ∈ {8, 16, 32, 64,
+    # 128}`` × ``num_warps ∈ {4, 8, 16}`` showed ``BLOCK_K=16`` /
+    # ``num_warps=8`` is ~2 ms faster than the gather path's
+    # ``BLOCK_K=32`` choice (16.2 ms vs 18.6 ms). Keep the gather
+    # path's ``BLOCK_K=32`` since it has its own register profile.
+    BLOCK_K_PARTIAL = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_BLOCK_K", "16"))
+    BLOCK_DMODEL = D
+    use_split_sparse = os.getenv("PRIMUS_V4_CSA_BWD_SPLIT_SPARSE", "1") != "0"
+
+    # Plan-5 P32: enable stream-level overlap between the local SWA
+    # BWD (dq + dk/dv kernels) and the sparse pool BWD kernel. The two
+    # paths write to disjoint output buffers (``dq_local_fp32`` /
+    # ``dk_local_fp32`` / ``dv_local_fp32`` / ``dsink_fp32`` vs
+    # ``dq_sparse_fp32`` / ``dpool_partial``) so the streams can
+    # progress concurrently. We sum ``dq_local + dq_sparse`` on the
+    # default stream after the join.
+    # Plan-5 P32: stream overlap defaults off — the segreduce path
+    # already places ``dq`` writes on disjoint buffers and the local
+    # kernel cost (~7.6 ms) is shorter than the sparse path (~3 ms in
+    # the segreduce variant), so there is little compute to hide. Set
+    # ``PRIMUS_V4_CSA_BWD_STREAM_OVERLAP=1`` to re-enable for
+    # experimentation. P57 R2 confirmed that enabling overlap
+    # actually doubles wall-clock to ~11 ms on the proxy shape — the
+    # two paths both saturate HBM read bandwidth, so concurrent
+    # execution serializes on memory traffic and adds stream
+    # synchronization overhead.
+    use_stream_overlap = use_split_sparse and os.getenv("PRIMUS_V4_CSA_BWD_STREAM_OVERLAP", "0") != "0"
+
+    # P57: when the input is bf16/fp16 and we use the split local SWA
+    # path (every slab overwritten by tl.store, no atomic_add), keep
+    # ``dq / dk_local / dv_local`` in the INPUT dtype so the kernels
+    # write directly at the final precision. Eliminates the trailing
+    # ~256 MB fp32 → bf16 cast on each of these 3 buffers (~0.5 ms
+    # total at the proxy shape) AND halves their HBM write traffic
+    # in-kernel (~1 GB → 512 MB for the three buffers). The
+    # monolithic local kernel and the gather sparse path still need
+    # fp32 because they ``atomic_add`` into the buffer and Triton's
+    # bf16 atomic_add is not supported on AMD MI355.
+    use_local_split_alloc = use_split_sparse and os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "0") == "1"
+    if use_local_split_alloc and q.dtype in (torch.bfloat16, torch.float16):
+        local_dq_dtype = q.dtype
+    else:
+        local_dq_dtype = torch.float32
+    if use_local_split_alloc:
+        dq_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype)
+        dk_local_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype)
+        dv_local_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype)
+    else:
+        dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dk_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+        dv_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32)
+    # ``dpool_fp32`` is small (2 MB at proxy) so the zero-fill cost
+    # is negligible (~20 us); keep ``zeros`` so the gather +
+    # atomic_add fallback path stays safe.
+    dpool_fp32 = torch.zeros((B, P, D), device=q.device, dtype=torch.float32)
+    # The split-sparse paths (both segreduce and gather) always write
+    # ``dq_sparse`` via ``tl.store`` from disjoint ``(pid_m,
+    # pid_h_block)`` programs, so we use a dedicated buffer and sum
+    # at the end. Aliasing to ``dq_fp32`` would clobber the
+    # local-SWA contribution. P57: match input dtype (bf16) when
+    # ``local_dq_dtype`` is bf16 so the final add and dtype-cast
+    # become free.
+    if use_split_sparse:
+        dq_sparse_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype)
+    else:
+        dq_sparse_fp32 = dq_fp32
+    if has_sink:
+        dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32)
+        sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink
+    else:
+        dsink_fp32 = q
+        sink_arg = q
+
+    d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+    pre_grid = (triton.cdiv(Sq, BLOCK_N), B * HQ)
+    _v4_attention_bwd_preprocess_kernel[pre_grid](
+        out,
+        dout,
+        d_buf,
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        dout.stride(0),
+        dout.stride(1),
+        dout.stride(2),
+        dout.stride(3),
+        d_buf.stride(0),
+        d_buf.stride(1),
+        d_buf.stride(2),
+        Sq,
+        HEAD=HQ,
+        BLOCK_M=BLOCK_N,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+
+    # Plan-5 P32: pre-compute the segreduce inverse index on the
+    # default stream BEFORE launching the local kernels. The sort +
+    # searchsorted are CPU-light but launch ~0.2 ms of kernels, and
+    # placing them BEFORE the local BWD launches lets the sparse
+    # stream's wait_event fire as soon as the index + d_buf are
+    # ready (instead of waiting until after all local kernels have
+    # been queued).
+    perm32 = None
+    bin_ptr = None
+    dpool_partial = None
+    use_segreduce = use_split_sparse and os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE", "1") == "1"
+    if use_segreduce:
+        # P57: ``dpool_partial`` defaults to the INPUT dtype when bf16
+        # / fp16 — halves HBM write+read traffic on the 4 GiB partial
+        # buffer (4 GB fp32 → 2 GB bf16) for the ~production-shape
+        # bench, saving ~0.6 ms / step. fp32 inputs keep an fp32
+        # partial so the parity tests' tight 1e-4 atol holds.
+        # The legacy P32 attempt failed parity because it used bf16
+        # unconditionally; gating on the input dtype keeps the
+        # bf16-only speed win without regressing fp32 numerics.
+        env_dtype = os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DTYPE", "")
+        if env_dtype == "bf16":
+            partial_dtype = torch.bfloat16
+        elif env_dtype == "fp16":
+            partial_dtype = torch.float16
+        elif env_dtype == "fp32":
+            partial_dtype = torch.float32
+        else:
+            # Default: match input dtype for bf16 / fp16, fp32
+            # otherwise. ``q.dtype`` is the most reliable proxy for
+            # the gradient-tolerance class.
+            if q.dtype in (torch.bfloat16, torch.float16):
+                partial_dtype = q.dtype
+            else:
+                partial_dtype = torch.float32
+        # P57: experimental ``use_sorted_partial`` flips the partial
+        # buffer layout from ``[B, M, K_topk, D]`` to
+        # ``[B, MK_sorted, D]``. The downstream segreduce kernel can
+        # then read contiguous slices, but the partial kernel writes
+        # become scattered (each ``sparse_k`` row goes to its
+        # ``inv_perm[orig_flat]`` sorted position). On the EP8 proxy
+        # the natural-order partial wins (~6.43 ms vs ~6.82 ms),
+        # because the partial kernel is write-bound and contiguous
+        # writes coalesce better than scattered ones. Keep the
+        # sorted-partial kernels in tree for shapes where the
+        # segreduce-read trade-off tips the other way (e.g., tiny
+        # ``D`` or very-large ``P`` workloads).
+        use_sorted_partial = os.getenv("PRIMUS_V4_CSA_BWD_SORTED_PARTIAL", "0") != "0"
+        with torch.no_grad():
+            MK = Sq * K_topk
+            flat_topk = topk_idxs.contiguous().view(B, MK).to(torch.int32)
+            sentinel = torch.full_like(flat_topk, P)
+            masked = torch.where((flat_topk >= 0) & (flat_topk < P), flat_topk, sentinel)
+            sorted_topk, perm = torch.sort(masked, dim=1, stable=True)
+            perm32 = perm.to(torch.int32)
+            queries = torch.arange(P + 1, device=q.device, dtype=torch.int32)
+            queries = queries.unsqueeze(0).expand(B, -1).contiguous()
+            bin_ptr = torch.searchsorted(sorted_topk, queries, right=False).to(torch.int32)
+            inv_perm32 = None
+            if use_sorted_partial:
+                # inv_perm[bid, orig_flat] = sorted_position. Built by
+                # scattering ``arange(MK)`` to the ``perm``-indexed
+                # positions. ~30 us at the proxy shape.
+                inv_perm32 = torch.empty_like(perm32)
+                idx_range = torch.arange(MK, device=q.device, dtype=torch.int32).unsqueeze(0).expand(B, -1)
+                inv_perm32.scatter_(1, perm.long(), idx_range)
+        dpool_partial = torch.empty((B, Sq, K_topk, D), device=q.device, dtype=partial_dtype)
+
+    # Lazily set up an extra CUDA stream for the sparse pool BWD so it
+    # can overlap with the local BWD launches (default-stream serial).
+    sparse_stream_ctx = None
+    if use_stream_overlap:
+        sparse_stream = torch.cuda.Stream(device=q.device)
+        d_buf_done = torch.cuda.current_stream(q.device).record_event()
+        sparse_stream.wait_event(d_buf_done)
+        sparse_stream_ctx = torch.cuda.stream(sparse_stream)
+
+    if use_split_sparse:
+        # P57: expose local-SWA tuning knobs (BLOCK_M, num_warps,
+        # num_stages). The local dq / dkv kernels live in
+        # ``v4_attention_bwd.py`` (outside the P57 file scope), but the
+        # LAUNCHER picks the block size + warp / stage count, and those
+        # parameters dominate the local-path wall clock at the
+        # production SWA=128 shape. The new defaults below are tuned
+        # specifically for ``B=1, H=64, Sq=4096, D=512, K_topk=512,
+        # swa_window=128`` on MI355 and cut the local dq+dkv from
+        # ~6.9 ms (P32 defaults: BM=32, BN=32, w=8, s=1) to ~4.1 ms
+        # (BM=64, BN=16 → fewer m programs and smaller n tiles to
+        # match the small SWA window).
+        local_block_m = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_BLOCK_M", "64"))
+        local_block_n = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_BLOCK_N", "16"))
+        # P57: the dq and dkv kernels prefer different per-axis block
+        # sizes (dq iterates ``n`` over a small SWA window so larger
+        # ``BLOCK_M`` packs more m-rows per program; dkv iterates
+        # ``m`` so its program-axis ``BLOCK_N`` and the inner-axis
+        # ``BLOCK_M`` can be tuned independently).
+        local_block_m_dq = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_BLOCK_M", str(local_block_m)))
+        local_block_n_dq = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_BLOCK_N", str(local_block_n)))
+        # P57 R2: dkv-specific tuning — ``BLOCK_M_dkv=16 warps=2 stages=1``
+        # wins post-R2 scale-defer (the dkv kernel's m-tile prefers
+        # smaller block sizes when the sparse partial path runs faster
+        # and exposes more concurrent CU slots). BM=16 is ~10 us faster
+        # than BM=32 at the proxy shape.
+        local_block_m_dkv = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_BLOCK_M", "16"))
+        local_block_n_dkv = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_BLOCK_N", str(local_block_n)))
+        # P57 R2: ``DQ_WARPS=4, DQ_STAGES=2`` wins ~140 us over the
+        # R1 ``W=8, S=1`` default. The dq kernel iterates ``n`` over
+        # a small SWA=128 window, so 4 warps fit the working set in
+        # registers and 2 stages prefetch the next K/V tile while the
+        # current MFMA chain executes.
+        local_dq_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_WARPS", "4"))
+        local_dq_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_STAGES", "2"))
+        local_dkv_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_WARPS", "2"))
+        local_dkv_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_STAGES", "1"))
+        swa_local = int(swa_window) if swa_window > 0 else 0
+        use_local_split = os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "0") == "1"
+        if use_local_split:
+            dq_grid = (triton.cdiv(Sq, local_block_m_dq), B * HQ)
+            _v4_attention_bwd_dq_kernel[dq_grid](
+                q,
+                k_local,
+                v_local,
+                dout,
+                lse,
+                d_buf,
+                dq_fp32,
+                dsink_fp32,
+                sink_arg,
+                q,  # ADD_MASK sentinel
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                k_local.stride(0),
+                k_local.stride(1),
+                k_local.stride(2),
+                k_local.stride(3),
+                v_local.stride(0),
+                v_local.stride(1),
+                v_local.stride(2),
+                v_local.stride(3),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_fp32.stride(0),
+                dq_fp32.stride(1),
+                dq_fp32.stride(2),
+                dq_fp32.stride(3),
+                0,
+                0,
+                Sq,
+                Sq,
+                float(scale),
+                HEAD_Q=HQ,
+                HEAD_K=HQ,
+                SWA_WINDOW=swa_local,
+                HAS_SINK=has_sink,
+                HAS_ADD_MASK=False,
+                HCA_LOCAL_SEQLEN=0,
+                USE_CAUSAL=True,
+                BLOCK_M=local_block_m_dq,
+                BLOCK_N=local_block_n_dq,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=local_dq_warps,
+                num_stages=local_dq_stages,
+            )
+            dkv_grid = (triton.cdiv(Sq, local_block_n_dkv), B * HQ)
+            _v4_attention_bwd_dkv_kernel[dkv_grid](
+                q,
+                k_local,
+                v_local,
+                dout,
+                lse,
+                d_buf,
+                dk_local_fp32,
+                dv_local_fp32,
+                q,  # ADD_MASK sentinel
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                k_local.stride(0),
+                k_local.stride(1),
+                k_local.stride(2),
+                k_local.stride(3),
+                v_local.stride(0),
+                v_local.stride(1),
+                v_local.stride(2),
+                v_local.stride(3),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dk_local_fp32.stride(0),
+                dk_local_fp32.stride(1),
+                dk_local_fp32.stride(2),
+                dk_local_fp32.stride(3),
+                dv_local_fp32.stride(0),
+                dv_local_fp32.stride(1),
+                dv_local_fp32.stride(2),
+                dv_local_fp32.stride(3),
+                0,
+                0,
+                Sq,
+                Sq,
+                float(scale),
+                HEAD_Q=HQ,
+                HEAD_K=HQ,
+                SWA_WINDOW=swa_local,
+                HAS_ADD_MASK=False,
+                HCA_LOCAL_SEQLEN=0,
+                USE_CAUSAL=True,
+                BLOCK_M=local_block_m_dkv,
+                BLOCK_N=local_block_n_dkv,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=local_dkv_warps,
+                num_stages=local_dkv_stages,
+            )
+        else:
+            local_grid = (triton.cdiv(Sq, local_block_m), B * HQ)
+            _v4_attention_bwd_kernel[local_grid](
+                q,
+                k_local,
+                v_local,
+                dout,
+                lse,
+                d_buf,
+                dq_fp32,
+                dk_local_fp32,
+                dv_local_fp32,
+                dsink_fp32,
+                sink_arg,
+                q,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                k_local.stride(0),
+                k_local.stride(1),
+                k_local.stride(2),
+                k_local.stride(3),
+                v_local.stride(0),
+                v_local.stride(1),
+                v_local.stride(2),
+                v_local.stride(3),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_fp32.stride(0),
+                dq_fp32.stride(1),
+                dq_fp32.stride(2),
+                dq_fp32.stride(3),
+                dk_local_fp32.stride(0),
+                dk_local_fp32.stride(1),
+                dk_local_fp32.stride(2),
+                dk_local_fp32.stride(3),
+                dv_local_fp32.stride(0),
+                dv_local_fp32.stride(1),
+                dv_local_fp32.stride(2),
+                dv_local_fp32.stride(3),
+                0,
+                0,
+                Sq,
+                Sq,
+                float(scale),
+                HEAD_Q=HQ,
+                HEAD_K=HQ,
+                SWA_WINDOW=swa_local,
+                HAS_SINK=has_sink,
+                HAS_ADD_MASK=False,
+                HCA_LOCAL_SEQLEN=0,
+                USE_CAUSAL=True,
+                BLOCK_M=local_block_m,
+                BLOCK_N=BLOCK_N,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=8,
+                num_stages=1,
+            )
+    else:
+        grid = (Sq, B * HQ)
+        _v4_csa_attention_pool_bwd_kernel[grid](
+            q,
+            k_local,
+            v_local,
+            pool,
+            topk_idxs,
+            dout,
+            lse,
+            d_buf,
+            dq_fp32,
+            dk_local_fp32,
+            dv_local_fp32,
+            dpool_fp32,
+            dsink_fp32,
+            sink_arg,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            k_local.stride(0),
+            k_local.stride(1),
+            k_local.stride(2),
+            k_local.stride(3),
+            v_local.stride(0),
+            v_local.stride(1),
+            v_local.stride(2),
+            v_local.stride(3),
+            pool.stride(0),
+            pool.stride(1),
+            pool.stride(2),
+            topk_idxs.stride(0),
+            topk_idxs.stride(1),
+            topk_idxs.stride(2),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            d_buf.stride(0),
+            d_buf.stride(1),
+            d_buf.stride(2),
+            dq_fp32.stride(0),
+            dq_fp32.stride(1),
+            dq_fp32.stride(2),
+            dq_fp32.stride(3),
+            dk_local_fp32.stride(0),
+            dk_local_fp32.stride(1),
+            dk_local_fp32.stride(2),
+            dk_local_fp32.stride(3),
+            dv_local_fp32.stride(0),
+            dv_local_fp32.stride(1),
+            dv_local_fp32.stride(2),
+            dv_local_fp32.stride(3),
+            dpool_fp32.stride(0),
+            dpool_fp32.stride(1),
+            dpool_fp32.stride(2),
+            Sq,
+            P,
+            K_topk,
+            float(scale),
+            HEAD_Q=HQ,
+            SWA_WINDOW=int(swa_window) if swa_window > 0 else 0,
+            HAS_SINK=has_sink,
+            BLOCK_N=BLOCK_N,
+            BLOCK_K=BLOCK_K,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            STORE_DPOOL=(os.getenv("PRIMUS_V4_CSA_BWD_SKIP_DPOOL", "0") != "1"),
+            LOCAL_ONLY=False,
+            num_warps=4,
+            num_stages=1,
+        )
+
+    # Plan-5 P32: replace the gather-based sparse pool BWD with a
+    # dense-pool BWD that re-uses the split dQ / dK_pool+dV_pool kernels
+    # against a ``[Sq, P]`` additive mask synthesised from
+    # ``topk_idxs``. The dense path eliminates the per-tile
+    # ``tl.atomic_add(DPOOL, ...)`` collisions (``H_blocks × ~2k m's``
+    # programs colliding on the shared ``[P, D]`` pool slots) at the
+    # cost of iterating the full pool dimension per m-block. With
+    # ``P=1024`` and ``K_topk=512`` (visibility ≈ 50%) the extra
+    # compute is small but the wall time win is large because atomics
+    # are the dominant cost at MI355 ``H=64``.
+    #
+    # Supported when ``B == 1`` (the proxy shape). For larger batches
+    # the kernel mask is per-(b, m), which the legacy gather kernel
+    # already handles natively — we fall back to it.
+    # Plan-5 P32: an alternative dense-pool path re-uses the split
+    # dQ + dK/dV kernels with a ``[Sq, P]`` ``log(count)`` additive
+    # mask, swapping ~1 B fp32 atomic adds for ``H * P/BLOCK_N``
+    # atomic_adds plus ~2× the compute (full ``P`` vs ``K_topk``).
+    # At the proxy shape (``H=64, P=1024, K_topk=512``) the extra
+    # compute outweighs the atomic savings (56 ms vs the gather
+    # path's 25 ms), so the dense path is kept in tree but **off**
+    # by default. Toggle on for shapes with ``K_topk / P`` near 1
+    # via ``PRIMUS_V4_CSA_BWD_DENSE_POOL=1``.
+    use_dense_pool_sparse = (
+        use_split_sparse and B == 1 and os.getenv("PRIMUS_V4_CSA_BWD_DENSE_POOL", "0") == "1"
+    )
+    # Enter the sparse stream just before launching the sparse pool
+    # kernels (default-stream kernels above have already been issued
+    # async on the default stream — the GPU can now overlap their
+    # execution with the sparse stream's work).
+    sparse_stream_entered = False
+    if sparse_stream_ctx is not None and use_split_sparse:
+        sparse_stream_ctx.__enter__()
+        sparse_stream_entered = True
+
+    if use_dense_pool_sparse:
+        # Build ``additive_mask[Sq, P]`` from ``topk_idxs[0]``. A finite
+        # ``NEG_INF`` is required so the bf16 ``Q @ pool.T`` matmul does
+        # not produce ``NaN`` after ``+ -inf``. Duplicate top-K slots
+        # (k1, k2 both pointing to pool position p) make the gather BWD
+        # accumulate ``2 *`` the contribution of a single key. We
+        # collapse the gather visibility into a *count*-weighted mask:
+        #
+        #   mask[m, p] = log(count[m, p])  if count[m, p] > 0
+        #                NEG_INF           otherwise
+        #
+        # With this mask the dense kernel's ``P[m, p] = exp(qk + log
+        # count - LSE) = count * exp(qk - LSE)`` so the dense-pool
+        # ``ds`` matches the gather BWD's ``sum_k ds[m, k]`` term-by-
+        # term and the dense ``dq`` / ``dpool`` agree with the
+        # reference. (Invalid ``-1`` slots contribute zero count.)
+        NEG_INF = -1.0e30
+        topk_b = topk_idxs[0]  # [Sq, K_topk]
+        valid_per_topk = (topk_b >= 0) & (topk_b < P)
+        safe_topk = torch.where(valid_per_topk, topk_b, torch.zeros_like(topk_b)).to(torch.int64)
+        count = torch.zeros((Sq, P), device=q.device, dtype=torch.float32)
+        count.scatter_add_(
+            1,
+            safe_topk,
+            valid_per_topk.to(torch.float32),
+        )
+        sparse_mask = torch.where(
+            count > 0,
+            count.clamp_min(1.0).log(),
+            torch.full((), NEG_INF, device=q.device, dtype=torch.float32),
+        )
+
+        # View pool as a per-head K/V with stride_kh=stride_vh=0 so the
+        # split dq/dkv kernels see ``HEAD_K == HEAD_Q`` and parallelise
+        # the dkv kernel over the full ``(P/BLOCK_N, B*HEAD_Q)`` grid
+        # (2 048 programs at proxy) instead of ``HEAD_K=1`` MQA (32
+        # programs) — that path was bandwidth-starved at ``D=512``.
+        # Use ``stride_dk = stride_dv = 0`` so each program's ``dk /
+        # dv`` write fans into the same shared ``dpool`` slice; the
+        # accumulation is via ``tl.atomic_add`` which costs ~4 096
+        # call ops total (single atomic per (n_block, head) plus dK +
+        # dV) — orders of magnitude less than the gather kernel's
+        # ~1 B fp32 atomic adds.
+        pool_4d = pool.unsqueeze(1).expand(B, HQ, P, D)  # stride_kh=0
+
+        sparse_block_m = 32
+        sparse_block_n = 32
+        dq_sparse_grid = (triton.cdiv(Sq, sparse_block_m), B * HQ)
+        _v4_attention_bwd_dq_kernel[dq_sparse_grid](
+            q,
+            pool_4d,
+            pool_4d,
+            dout,
+            lse,
+            d_buf,
+            dq_sparse_fp32,
+            dsink_fp32,  # sentinel only; HAS_SINK=False
+            q,
+            sparse_mask,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            pool_4d.stride(0),
+            pool_4d.stride(1),
+            pool_4d.stride(2),
+            pool_4d.stride(3),
+            pool_4d.stride(0),
+            pool_4d.stride(1),
+            pool_4d.stride(2),
+            pool_4d.stride(3),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            d_buf.stride(0),
+            d_buf.stride(1),
+            d_buf.stride(2),
+            dq_sparse_fp32.stride(0),
+            dq_sparse_fp32.stride(1),
+            dq_sparse_fp32.stride(2),
+            dq_sparse_fp32.stride(3),
+            sparse_mask.stride(0),
+            sparse_mask.stride(1),
+            Sq,
+            P,
+            float(scale),
+            HEAD_Q=HQ,
+            HEAD_K=HQ,
+            SWA_WINDOW=0,
+            HAS_SINK=False,
+            HAS_ADD_MASK=True,
+            HCA_LOCAL_SEQLEN=0,
+            USE_CAUSAL=False,
+            BLOCK_M=sparse_block_m,
+            BLOCK_N=sparse_block_n,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            ACCUMULATE=True,
+            num_warps=8,
+            num_stages=1,
+        )
+        # Single shared ``dpool`` slice that all ``(head, n_block)``
+        # programs atomic_add into. Reshape with an extra head axis of
+        # stride 0 so the dkv kernel can index it as ``[B, HQ, P, D]``
+        # without per-head storage; the actual underlying tensor is
+        # still ``[B, P, D]``.
+        dk_pool_shared = dpool_fp32.unsqueeze(1).expand(B, HQ, P, D)
+        dv_pool_shared = dpool_fp32.unsqueeze(1).expand(B, HQ, P, D)
+        dkv_sparse_grid = (triton.cdiv(P, sparse_block_n), B * HQ)
+        _v4_attention_bwd_dkv_kernel[dkv_sparse_grid](
+            q,
+            pool_4d,
+            pool_4d,
+            dout,
+            lse,
+            d_buf,
+            dk_pool_shared,
+            dv_pool_shared,
+            sparse_mask,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            pool_4d.stride(0),
+            pool_4d.stride(1),
+            pool_4d.stride(2),
+            pool_4d.stride(3),
+            pool_4d.stride(0),
+            pool_4d.stride(1),
+            pool_4d.stride(2),
+            pool_4d.stride(3),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            d_buf.stride(0),
+            d_buf.stride(1),
+            d_buf.stride(2),
+            dk_pool_shared.stride(0),
+            dk_pool_shared.stride(1),
+            dk_pool_shared.stride(2),
+            dk_pool_shared.stride(3),
+            dv_pool_shared.stride(0),
+            dv_pool_shared.stride(1),
+            dv_pool_shared.stride(2),
+            dv_pool_shared.stride(3),
+            sparse_mask.stride(0),
+            sparse_mask.stride(1),
+            Sq,
+            P,
+            float(scale),
+            HEAD_Q=HQ,
+            HEAD_K=HQ,
+            SWA_WINDOW=0,
+            HAS_ADD_MASK=True,
+            HCA_LOCAL_SEQLEN=0,
+            USE_CAUSAL=False,
+            BLOCK_M=sparse_block_m,
+            BLOCK_N=sparse_block_n,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            ATOMIC_REDUCE=True,
+            num_warps=8,
+            num_stages=1,
+        )
+        # ``dpool_fp32`` already holds the summed dk+dv contributions
+        # because both kernel views aliased the same ``[B, P, D]``
+        # storage via ``stride_kh=stride_vh=0`` and the kernel does
+        # both ``atomic_add(dk_ptrs, ...)`` and ``atomic_add(dv_ptrs,
+        # ...)``.
+    elif use_split_sparse:
+        # P57: BLOCK_H=32, num_warps=4, num_stages=2 wins the proxy
+        # sweep over the P32 default (BLOCK_H=64, num_warps=8, stages=1).
+        # Doubling the head-axis grid (HQ=64 -> 2 h-blocks per m) lifts
+        # MI355 occupancy, while fewer warps + 2 stages reduces register
+        # pressure per warp (the dq[BLOCK_H, BLOCK_DMODEL]=fp32
+        # accumulator dominates VGPR live range).
+        BLOCK_H = int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_BLOCK_H", "32"))
+        # Plan-5 P32: ``PRIMUS_V4_CSA_BWD_SEGREDUCE=1`` is now the
+        # default — wins both the standalone CSA BWD microbench
+        # (16.31 ms vs 24.83 ms gather/atomic) and the EP8 proxy
+        # (578 ms vs 665 ms / iter) after the P32 dual-RoPE bf16-cast
+        # fix.  Pre-fix, ``apply_interleaved_partial_rope`` was
+        # promoting Q/K to fp32 (cos/sin came from
+        # ``position_ids.float()`` and bf16*fp32=fp32), which 2x'd
+        # Q/K HBM traffic, inflated *every* attention kernel time
+        # 1.8-7x in the proxy trace, and made the gather + atomic
+        # path look faster end-to-end purely because the segreduce
+        # 4 GiB partial buffer competed against artificially-bloated
+        # attention traffic. Setting ``PRIMUS_V4_CSA_BWD_SEGREDUCE=0``
+        # falls back to gather + atomic for kernel-tuning. The
+        # segmented-reduction path writes per-visit dpool
+        # contributions to a compact ``[B, M, K_topk, D]`` partial
+        # buffer (no atomics) and then folds them into
+        # ``dpool[B, P, D]`` via a sorted inverse index. ``perm32``,
+        # ``bin_ptr`` and ``dpool_partial`` were built above on the
+        # default stream BEFORE the local kernels were launched, so
+        # they're ready by the time the sparse stream gets here.
+        # P57: experimental split-kernel path (opt-in).
+        # Sub-kernels:
+        #
+        #   _v4_csa_attention_pool_sparse_bwd_dq_only_kernel
+        #   _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel
+        #
+        # vs the joint kernel
+        # ``_v4_csa_attention_pool_sparse_bwd_partial_kernel``. The
+        # split path frees the live-VGPR footprint of each sub-kernel
+        # but pays a 2× read on Q / dout / lse / dvec / pool /
+        # topk_idxs. On the EP8 proxy shape the joint kernel still
+        # wins (~3.7 ms vs ~4.5 ms split), so split is opt-in via
+        # ``PRIMUS_V4_CSA_BWD_SPLIT_DQ_DPOOL=1`` and the joint kernel
+        # remains the default. Keeping the split sub-kernels in tree
+        # for future shapes where the read amplification is cheaper
+        # (smaller H or D, or compute-bound regimes).
+        use_split_dq_dpool = os.getenv("PRIMUS_V4_CSA_BWD_SPLIT_DQ_DPOOL", "0") != "0"
+        if use_segreduce and use_split_dq_dpool:
+            sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B)
+            dq_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DQ_WARPS", "4"))
+            dq_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DQ_STAGES", "2"))
+            dpool_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DPOOL_WARPS", "4"))
+            dpool_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DPOOL_STAGES", "2"))
+            _v4_csa_attention_pool_sparse_bwd_dq_only_kernel[sparse_grid](
+                q,
+                pool,
+                topk_idxs,
+                dout,
+                lse,
+                d_buf,
+                dq_sparse_fp32,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                pool.stride(0),
+                pool.stride(1),
+                pool.stride(2),
+                topk_idxs.stride(0),
+                topk_idxs.stride(1),
+                topk_idxs.stride(2),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_sparse_fp32.stride(0),
+                dq_sparse_fp32.stride(1),
+                dq_sparse_fp32.stride(2),
+                dq_sparse_fp32.stride(3),
+                Sq,
+                P,
+                K_topk,
+                float(scale),
+                HEAD_Q=HQ,
+                BLOCK_H=BLOCK_H,
+                BLOCK_K=BLOCK_K_PARTIAL,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=dq_warps,
+                num_stages=dq_stages,
+            )
+            _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel[sparse_grid](
+                q,
+                pool,
+                topk_idxs,
+                dout,
+                lse,
+                d_buf,
+                dpool_partial,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                pool.stride(0),
+                pool.stride(1),
+                pool.stride(2),
+                topk_idxs.stride(0),
+                topk_idxs.stride(1),
+                topk_idxs.stride(2),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dpool_partial.stride(0),
+                dpool_partial.stride(1),
+                dpool_partial.stride(2),
+                dpool_partial.stride(3),
+                Sq,
+                P,
+                K_topk,
+                float(scale),
+                HEAD_Q=HQ,
+                BLOCK_H=BLOCK_H,
+                BLOCK_K=BLOCK_K_PARTIAL,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=dpool_warps,
+                num_stages=dpool_stages,
+            )
+            # P57: segreduce reduction (shared between split and joint
+            # partial paths). Reduces ``dpool_partial[B, MK, D]`` into
+            # ``dpool[B, P, D]`` using the sorted inverse index.
+            dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D)
+            block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512"))
+            block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64"))
+            seg_grid = (P, triton.cdiv(D, block_d_seg), B)
+            _v4_csa_attention_pool_segreduce_kernel[seg_grid](
+                dpool_partial_flat,
+                perm32,
+                bin_ptr,
+                dpool_fp32,
+                dpool_partial_flat.stride(0),
+                dpool_partial_flat.stride(1),
+                dpool_partial_flat.stride(2),
+                perm32.stride(0),
+                perm32.stride(1),
+                bin_ptr.stride(0),
+                bin_ptr.stride(1),
+                dpool_fp32.stride(0),
+                dpool_fp32.stride(1),
+                dpool_fp32.stride(2),
+                P,
+                D,
+                BLOCK_D=block_d_seg,
+                BLOCK_I=block_i_seg,
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")),
+            )
+        elif use_segreduce and use_sorted_partial:
+            sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B)
+            dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D)
+            _v4_csa_attention_pool_sparse_bwd_partial_sorted_kernel[sparse_grid](
+                q,
+                pool,
+                topk_idxs,
+                inv_perm32,
+                dout,
+                lse,
+                d_buf,
+                dq_sparse_fp32,
+                dpool_partial_flat,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                pool.stride(0),
+                pool.stride(1),
+                pool.stride(2),
+                topk_idxs.stride(0),
+                topk_idxs.stride(1),
+                topk_idxs.stride(2),
+                inv_perm32.stride(0),
+                inv_perm32.stride(1),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_sparse_fp32.stride(0),
+                dq_sparse_fp32.stride(1),
+                dq_sparse_fp32.stride(2),
+                dq_sparse_fp32.stride(3),
+                dpool_partial_flat.stride(0),
+                dpool_partial_flat.stride(1),
+                dpool_partial_flat.stride(2),
+                Sq,
+                P,
+                K_topk,
+                float(scale),
+                HEAD_Q=HQ,
+                BLOCK_H=BLOCK_H,
+                BLOCK_K=BLOCK_K_PARTIAL,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_STAGES", "2")),
+            )
+            # Sequential segreduce — no perm lookup.
+            block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512"))
+            block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64"))
+            seg_grid = (P, triton.cdiv(D, block_d_seg), B)
+            _v4_csa_attention_pool_segreduce_sequential_kernel[seg_grid](
+                dpool_partial_flat,
+                bin_ptr,
+                dpool_fp32,
+                dpool_partial_flat.stride(0),
+                dpool_partial_flat.stride(1),
+                dpool_partial_flat.stride(2),
+                bin_ptr.stride(0),
+                bin_ptr.stride(1),
+                dpool_fp32.stride(0),
+                dpool_fp32.stride(1),
+                dpool_fp32.stride(2),
+                P,
+                D,
+                BLOCK_D=block_d_seg,
+                BLOCK_I=block_i_seg,
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")),
+            )
+        elif use_segreduce:
+            sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B)
+            _v4_csa_attention_pool_sparse_bwd_partial_kernel[sparse_grid](
+                q,
+                pool,
+                topk_idxs,
+                dout,
+                lse,
+                d_buf,
+                dq_sparse_fp32,
+                dpool_partial,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                pool.stride(0),
+                pool.stride(1),
+                pool.stride(2),
+                topk_idxs.stride(0),
+                topk_idxs.stride(1),
+                topk_idxs.stride(2),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_sparse_fp32.stride(0),
+                dq_sparse_fp32.stride(1),
+                dq_sparse_fp32.stride(2),
+                dq_sparse_fp32.stride(3),
+                dpool_partial.stride(0),
+                dpool_partial.stride(1),
+                dpool_partial.stride(2),
+                dpool_partial.stride(3),
+                Sq,
+                P,
+                K_topk,
+                float(scale),
+                HEAD_Q=HQ,
+                BLOCK_H=BLOCK_H,
+                BLOCK_K=BLOCK_K_PARTIAL,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_STAGES", "3")),
+            )
+            # Same segreduce reduction as the split path above.
+            dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D)
+            block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512"))
+            block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64"))
+            seg_grid = (P, triton.cdiv(D, block_d_seg), B)
+            _v4_csa_attention_pool_segreduce_kernel[seg_grid](
+                dpool_partial_flat,
+                perm32,
+                bin_ptr,
+                dpool_fp32,
+                dpool_partial_flat.stride(0),
+                dpool_partial_flat.stride(1),
+                dpool_partial_flat.stride(2),
+                perm32.stride(0),
+                perm32.stride(1),
+                bin_ptr.stride(0),
+                bin_ptr.stride(1),
+                dpool_fp32.stride(0),
+                dpool_fp32.stride(1),
+                dpool_fp32.stride(2),
+                P,
+                D,
+                BLOCK_D=block_d_seg,
+                BLOCK_I=block_i_seg,
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")),
+            )
+        else:
+            # Plan-5 P32: legacy gather-sparse path. Keeps the atomic_add
+            # to a shared ``dpool[B, P, D]`` buffer; this is the shipped
+            # default. The segreduce path is opt-in via
+            # ``PRIMUS_V4_CSA_BWD_SEGREDUCE=1`` for kernel-level perf
+            # experiments on the CSA microbench.
+            n_part_env = int(os.getenv("PRIMUS_V4_CSA_BWD_DPOOL_PARTITIONS", "1"))
+            n_part = max(1, min(n_part_env, Sq))
+            while Sq % n_part != 0 and n_part > 1:
+                n_part -= 1
+            partition_size = Sq // n_part
+            if n_part > 1:
+                dpool_partial = torch.zeros((n_part, B, P, D), device=q.device, dtype=torch.float32)
+            else:
+                dpool_partial = dpool_fp32.unsqueeze(0)
+            sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B)
+            _v4_csa_attention_pool_sparse_bwd_kernel[sparse_grid](
+                q,
+                pool,
+                topk_idxs,
+                dout,
+                lse,
+                d_buf,
+                dq_sparse_fp32,
+                dpool_partial,
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                q.stride(3),
+                pool.stride(0),
+                pool.stride(1),
+                pool.stride(2),
+                topk_idxs.stride(0),
+                topk_idxs.stride(1),
+                topk_idxs.stride(2),
+                dout.stride(0),
+                dout.stride(1),
+                dout.stride(2),
+                dout.stride(3),
+                lse.stride(0),
+                lse.stride(1),
+                lse.stride(2),
+                d_buf.stride(0),
+                d_buf.stride(1),
+                d_buf.stride(2),
+                dq_sparse_fp32.stride(0),
+                dq_sparse_fp32.stride(1),
+                dq_sparse_fp32.stride(2),
+                dq_sparse_fp32.stride(3),
+                dpool_partial.stride(0),
+                dpool_partial.stride(1),
+                dpool_partial.stride(2),
+                dpool_partial.stride(3),
+                partition_size,
+                Sq,
+                P,
+                K_topk,
+                float(scale),
+                HEAD_Q=HQ,
+                BLOCK_H=BLOCK_H,
+                BLOCK_K=BLOCK_K,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                STORE_DPOOL=(os.getenv("PRIMUS_V4_CSA_BWD_SKIP_DPOOL_ATOMIC", "0") != "1"),
+                num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_WARPS", "4")),
+                num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_STAGES", "1")),
+            )
+            if n_part > 1:
+                dpool_fp32 = dpool_partial.sum(dim=0)
+
+    # Exit the sparse stream context (if entered) and join it with the
+    # default stream before the final dtype casts so the dq accumulator
+    # below sees both the local and sparse contributions.
+    if sparse_stream_entered:
+        sparse_done = torch.cuda.current_stream(q.device).record_event()
+        sparse_stream_ctx.__exit__(None, None, None)
+        torch.cuda.current_stream(q.device).wait_event(sparse_done)
+
+    # The split-sparse paths write the sparse contribution to a
+    # dedicated ``dq_sparse_fp32`` buffer; combine it with the local
+    # SWA contribution before the dtype cast. In-place add saves the
+    # extra ~128 MB allocation that ``+`` would do.
+    if use_split_sparse:
+        dq_fp32.add_(dq_sparse_fp32)
+
+    # P57: if the local buffers are already in input dtype, ``.to``
+    # is a no-op view; only fp32 buffers actually do a cast.
+    dq_out = dq_fp32 if dq_fp32.dtype == q.dtype else dq_fp32.to(q.dtype)
+    dk_local_out = dk_local_fp32 if dk_local_fp32.dtype == k_local.dtype else dk_local_fp32.to(k_local.dtype)
+    dv_local_out = dv_local_fp32 if dv_local_fp32.dtype == v_local.dtype else dv_local_fp32.to(v_local.dtype)
+    dpool_out = dpool_fp32 if dpool_fp32.dtype == pool.dtype else dpool_fp32.to(pool.dtype)
+    dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None
+    return dq_out, dk_local_out, dv_local_out, dpool_out, dsink_out
+
+
+__all__ = [
+    "_v4_csa_attention_bwd_kernel",
+    "_v4_csa_attention_pool_bwd_kernel",
+    "_v4_csa_attention_pool_sparse_bwd_kernel",
+    "_launch_v4_csa_attention_bwd",
+    "_launch_v4_csa_attention_pool_bwd",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py
new file mode 100644
index 000000000..5ee75bbbe
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py
@@ -0,0 +1,1040 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 CSA attention forward Triton kernel (plan-4 P26, ``compress_ratio == 4``).
+
+CSA fuses three branches into a single online softmax:
+
+* **Local SWA**: ``q @ k_local^T`` with sliding-window-causal masking.
+* **Sparse top-K**: ``q . gathered[m, :, :]`` where the wrapper has
+  pre-gathered ``[B, Sq, K, D]`` rows from the compressed pool (the
+  per-query top-K gather lives outside the kernel — see plan-4
+  ``02-phase-details.md`` Phase 26 design notes).
+* **Per-head learned sink**: a virtual key column with notional value
+  zero, joined as the last softmax candidate so its probability mass
+  is shared across local + sparse branches.
+
+The kernel produces one ``[BLOCK_DMODEL]`` output row per program; the
+grid is ``(seqlen_q, batch * head_q)`` so each program owns exactly one
+``(b, qhid, m)`` query row. The per-row design keeps the sparse-branch
+SMEM footprint inside the MI355 budget at ``head_dim=512``: the
+gathered tile is only ``[BLOCK_K, head_dim] * 2 bytes ≈ 32 KiB`` per
+program, while a multi-row tile would balloon to
+``[BLOCK_M, BLOCK_K, head_dim] * 2 bytes ≈ 1 MiB``.
+
+dtype contract (matches :func:`eager_v4_csa_attention`):
+
+* Q / K / V / gathered are loaded in input dtype (bf16 in production);
+  the per-row dot products (``sum(k * q[None, :], axis=-1)``) reduce in
+  fp32 because we ``.to(tl.float32)`` before the multiply.
+* The online-softmax accumulator (``m_running``, ``l_running``,
+  ``acc``) lives in fp32 — the *only* fp32 step inside the kernel.
+* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD
+  re-materialises ``P`` from it).
+
+Edge cases handled:
+
+* ``K_topk == 0`` — wrapper short-circuits to the dense
+  :func:`v4_attention_v1` kernel before reaching this file.
+* ``topk_idx == -1`` — wrapper sets the corresponding ``sparse_mask``
+  entry to ``-inf``; the kernel just adds the bias and the masked
+  position contributes ~0 to the softmax denominator.
+* All-masked tile rows — the running max and per-tile max are both
+  the finite ``NEG_INF`` sentinel (``-1e30``), so
+  ``exp(NEG_INF - NEG_INF) = exp(0) = 1`` algebraically but the
+  contribution to ``acc`` and ``l_running`` is gated by the
+  per-element ``exp(qk - m_new)`` which stays at exactly zero for
+  every ``-inf``-masked entry. This avoids the ``exp(-inf - -inf) =
+  exp(NaN)`` failure mode that ``-float("inf")`` would have.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import (
+    _launch_v4_attention_fwd,
+)
+
+# ---------------------------------------------------------------------------
+# Triton kernel
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_csa_attention_pool_fwd_kernel(
+    Q,
+    K_LOCAL,
+    V_LOCAL,
+    POOL,
+    TOPK_IDXS,
+    SINK,
+    OUT,
+    LSE,
+    # Q strides: [B, H, Sq, D]
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    # K_local strides: [B, H, Sq, D]
+    stride_klb,
+    stride_klh,
+    stride_kln,
+    stride_kld,
+    # V_local strides: [B, H, Sq, D]
+    stride_vlb,
+    stride_vlh,
+    stride_vln,
+    stride_vld,
+    # Pool strides: [B, P, D] (shared across heads)
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    # topk_idxs strides: [B, Sq, K_topk]
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    # OUT strides: [B, H, Sq, D]
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    # LSE strides: [B, H, Sq]
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """CSA FWD with in-kernel topk gather from the compressed pool."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    q_active = pid_m < seqlen_q
+    NEG_INF: tl.constexpr = -1.0e30
+
+    q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm
+    q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0)
+
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full((), value=NEG_INF, dtype=tl.float32)
+    l_i = tl.zeros((), dtype=tl.float32)
+
+    n_loop_end = pid_m + 1
+    if n_loop_end > seqlen_q:
+        n_loop_end = seqlen_q
+    if SWA_WINDOW > 0:
+        n_lo_raw = pid_m - SWA_WINDOW + 1
+        if n_lo_raw < 0:
+            n_lo_raw = 0
+        n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N
+    else:
+        n_loop_start = 0
+
+    for n_start in range(n_loop_start, n_loop_end, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+        kl_ptrs = (
+            K_LOCAL
+            + bid * stride_klb
+            + qhid * stride_klh
+            + offs_n[:, None] * stride_kln
+            + offs_d[None, :] * stride_kld
+        )
+        kl_load_mask = offs_n[:, None] < seqlen_q
+        kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0)
+        qk = tl.sum(kl.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale
+
+        if SWA_WINDOW > 0:
+            in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m)
+        else:
+            in_window = offs_n <= pid_m
+        qk = tl.where(in_window & (offs_n < seqlen_q), qk, NEG_INF)
+
+        m_tile = tl.max(qk, axis=0)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+
+        vl_ptrs = (
+            V_LOCAL
+            + bid * stride_vlb
+            + qhid * stride_vlh
+            + offs_n[:, None] * stride_vln
+            + offs_d[None, :] * stride_vld
+        )
+        vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0)
+        acc = acc * alpha + tl.sum(p[:, None] * vl.to(tl.float32), axis=0)
+        m_i = m_new
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        offs_k = k_start + tl.arange(0, BLOCK_K)
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + offs_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=offs_k < K_topk, other=-1)
+        valid = (offs_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid[:, None], other=0.0)
+
+        qk_sparse = tl.sum(pool.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale
+        qk_sparse = tl.where(valid, qk_sparse, NEG_INF)
+
+        m_tile = tl.max(qk_sparse, axis=0)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk_sparse - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+
+        acc = acc * alpha + tl.sum(p[:, None] * pool.to(tl.float32), axis=0)
+        m_i = m_new
+
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+        m_new = tl.maximum(m_i, sink_h)
+        alpha = tl.exp(m_i - m_new)
+        beta = tl.exp(sink_h - m_new)
+        l_i = l_i * alpha + beta
+        acc = acc * alpha
+        m_i = m_new
+
+    out = acc / l_i
+    lse = m_i + tl.log(l_i)
+
+    out_offset = bid * stride_ob + qhid * stride_oh + pid_m * stride_om
+    tl.store(OUT + out_offset + offs_d * stride_od, out.to(OUT.dtype.element_ty), mask=q_active)
+
+    lse_ptr = LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm
+    tl.store(lse_ptr, lse, mask=q_active)
+
+
+# ---------------------------------------------------------------------------
+# Python launcher
+# ---------------------------------------------------------------------------
+
+
+def _launch_v4_csa_attention_pool_fwd(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k_local: torch.Tensor,  # [B, H, Sq, D]
+    v_local: torch.Tensor,  # [B, H, Sq, D]
+    pool: torch.Tensor,  # [B, P, D]
+    topk_idxs: torch.Tensor,  # [B, Sq, K_topk], -1 means masked
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch CSA forward with in-kernel topk gather from ``pool``."""
+    if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4:
+        raise ValueError(
+            "v4_csa_attention_v0 pool forward expects q / k_local / v_local of rank 4 "
+            f"(got {q.dim()} / {k_local.dim()} / {v_local.dim()})"
+        )
+    if pool.dim() != 3:
+        raise ValueError(
+            f"v4_csa_attention_v0 pool forward expects pool of rank 3 [B, P, D]; "
+            f"got rank {pool.dim()}, shape {tuple(pool.shape)}"
+        )
+    if topk_idxs.dim() != 3:
+        raise ValueError(
+            f"v4_csa_attention_v0 pool forward expects topk_idxs of rank 3 [B, Sq, K]; "
+            f"got rank {topk_idxs.dim()}, shape {tuple(topk_idxs.shape)}"
+        )
+
+    B, HQ, Sq, D = q.shape
+    if k_local.shape != q.shape or v_local.shape != q.shape:
+        raise ValueError(
+            "v4_csa_attention_v0 pool path requires k_local.shape == v_local.shape == q.shape "
+            f"(got q={tuple(q.shape)}, k_local={tuple(k_local.shape)}, "
+            f"v_local={tuple(v_local.shape)})."
+        )
+    Bp, P, Dp = pool.shape
+    if Bp != B or Dp != D:
+        raise ValueError(
+            "v4_csa_attention_v0 pool shape mismatch: expected "
+            f"[B, P, D] = [{B}, *, {D}]; got {tuple(pool.shape)}."
+        )
+    Bt, Sqt, K_topk = topk_idxs.shape
+    if Bt != B or Sqt != Sq:
+        raise ValueError(
+            "v4_csa_attention_v0 topk_idxs shape mismatch: expected "
+            f"[B, Sq, K] = [{B}, {Sq}, *]; got {tuple(topk_idxs.shape)}."
+        )
+    if topk_idxs.dtype not in (torch.int32, torch.int64):
+        raise ValueError(f"v4_csa_attention_v0 topk_idxs must be int32/int64, got {topk_idxs.dtype}.")
+    if not q.is_cuda:
+        raise ValueError("v4_csa_attention_v0 requires CUDA / HIP tensors.")
+    if q.dtype != k_local.dtype or q.dtype != v_local.dtype or q.dtype != pool.dtype:
+        raise ValueError(
+            "v4_csa_attention_v0 pool path requires q.dtype == k_local.dtype == "
+            f"v_local.dtype == pool.dtype (got {q.dtype} / {k_local.dtype} / "
+            f"{v_local.dtype} / {pool.dtype})."
+        )
+
+    has_sink = sink is not None
+
+    # Plan-5 P32: default to the split FWD (local SWA dense kernel +
+    # head-block sparse kernel + LSE merge). The monolithic per-row
+    # ``_v4_csa_attention_pool_fwd_kernel`` stays in tree as the
+    # ``PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC=1`` fallback so we can A/B
+    # the two designs from the proxy without rebuilding.
+    if os.getenv("PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC", "0") != "1":
+        return _launch_v4_csa_attention_pool_fwd_split(
+            q,
+            k_local,
+            v_local,
+            pool,
+            topk_idxs,
+            sink=sink,
+            swa_window=int(swa_window),
+            scale=float(scale),
+        )
+
+    out = torch.empty_like(q)
+    lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    BLOCK_N = 32
+    BLOCK_K = 32
+    BLOCK_DMODEL = D
+    grid = (Sq, B * HQ)
+    sink_ptr = sink if has_sink else q
+
+    _v4_csa_attention_pool_fwd_kernel[grid](
+        q,
+        k_local,
+        v_local,
+        pool,
+        topk_idxs,
+        sink_ptr,
+        out,
+        lse,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        k_local.stride(0),
+        k_local.stride(1),
+        k_local.stride(2),
+        k_local.stride(3),
+        v_local.stride(0),
+        v_local.stride(1),
+        v_local.stride(2),
+        v_local.stride(3),
+        pool.stride(0),
+        pool.stride(1),
+        pool.stride(2),
+        topk_idxs.stride(0),
+        topk_idxs.stride(1),
+        topk_idxs.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        Sq,
+        P,
+        K_topk,
+        float(scale),
+        HEAD_Q=HQ,
+        SWA_WINDOW=int(swa_window) if swa_window > 0 else 0,
+        HAS_SINK=has_sink,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+    return out, lse
+
+
+# ---------------------------------------------------------------------------
+# Plan-5 P32 split CSA FWD — sparse head-block tile + LSE merge
+#
+# The per-row design in ``_v4_csa_attention_pool_fwd_kernel`` cannot reach
+# tensor-core throughput because ``tl.sum(k * q, axis=1)`` is a per-row
+# reduction (one program per ``(b, qhid, m)``). FlashAttention-style
+# multi-row ``tl.dot`` tiles need ``BLOCK_M >= 16`` queries per program;
+# the sparse branch's per-query ``topk_idxs`` gather blocks that — adjacent
+# query rows have different sparse keys.
+#
+# P32 splits the FWD into three launches that ARE multi-row friendly:
+#
+#  1. ``_launch_v4_attention_fwd`` (already shipped) handles the local
+#     SWA branch with ``BLOCK_M=BLOCK_N=32`` ``tl.dot`` tiles. We call it
+#     with ``sink=None`` so the returned ``(out_local, lse_local)`` does
+#     NOT include the per-head sink.
+#  2. ``_v4_csa_attention_pool_sparse_fwd_kernel`` (new) handles the
+#     sparse pool branch with a **head-block** tile. The per-query
+#     ``topk_idxs`` gather is shared across all ``H`` query heads — the
+#     pool tensor itself has no head dimension. One program owns one
+#     ``(b, m, h_block)`` and runs ``tl.dot(Q[BLOCK_H, D],
+#     tl.trans(pool[BLOCK_K, D]))`` per top-K tile, online-softmax-
+#     updating per-head ``m_i / l_i / acc`` along the way. Output
+#     ``(out_sparse, lse_sparse)`` does NOT include the sink.
+#  3. ``_v4_csa_attention_lse_merge_kernel`` (new) combines the two
+#     ``(out, lse)`` pairs with the per-head sink under one final online
+#     softmax. Result: a joint ``out`` and joint ``lse`` mathematically
+#     identical to the monolithic kernel (modulo fp32 reduction order),
+#     plus the per-iteration BWD contract is unchanged because the joint
+#     ``lse`` is what the BWD already re-materialises ``P`` from.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_fwd_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    OUT,
+    LSE,
+    # Q strides: [B, H, Sq, D]
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    # Pool strides: [B, P, D]
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    # topk_idxs strides: [B, Sq, K_topk]
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    # OUT strides: [B, H, Sq, D]
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    # LSE strides: [B, H, Sq]
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """Sparse-branch CSA FWD with head-block tile + ``tl.dot``.
+
+    Grid: ``(seqlen_q, cdiv(HEAD_Q, BLOCK_H), B)``. Each program owns
+    one ``(b, h_block, m)`` and computes the sparse branch's normalized
+    output + LSE for ``BLOCK_H`` heads. The pool gather is shared
+    across heads because ``pool`` and ``topk_idxs`` have no H axis.
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    h_mask = offs_h < HEAD_Q
+    q_active = pid_m < seqlen_q
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    # Q tile: [BLOCK_H, BLOCK_DMODEL]
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0)
+
+    # Online-softmax state per head.
+    acc = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_H], dtype=tl.float32)
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik
+        topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+        valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        # qk = Q @ pool.T : [BLOCK_H, BLOCK_K] in fp32 accumulator
+        qk = tl.dot(q.to(pool.dtype), tl.trans(pool)) * sm_scale
+        qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, NEG_INF)
+
+        # Online softmax update per head row.
+        m_tile = tl.max(qk, axis=1)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(pool.dtype), pool)
+        m_i = m_new
+
+    # Detect "all invalid" rows (K_topk == 0 or every topk_idx == -1):
+    # l_i stays zero; we leave acc as zero and write a NEG_INF lse so the
+    # merge kernel treats the sparse branch as carrying zero softmax mass
+    # for that (b, h, m).
+    empty = l_i == 0.0
+    safe_l = tl.where(empty, 1.0, l_i)
+    out = acc / safe_l[:, None]
+    lse = tl.where(empty, NEG_INF, m_i + tl.log(safe_l))
+
+    out_ptrs = (
+        OUT + bid * stride_ob + offs_h[:, None] * stride_oh + pid_m * stride_om + offs_d[None, :] * stride_od
+    )
+    tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=h_mask[:, None] & q_active)
+
+    lse_ptrs = LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm
+    tl.store(lse_ptrs, lse, mask=h_mask & q_active)
+
+
+@triton.jit
+def _v4_csa_attention_pool_sparse_merge_fwd_kernel(
+    Q,
+    POOL,
+    TOPK_IDXS,
+    OUT_LOCAL,  # local-branch result (no sink) — input
+    LSE_LOCAL,  # local-branch lse (no sink) — input
+    SINK,
+    OUT,  # joint output
+    LSE,  # joint lse
+    # Q strides: [B, H, Sq, D]
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    # Pool strides: [B, P, D]
+    stride_pb,
+    stride_pp,
+    stride_pd,
+    # topk_idxs strides: [B, Sq, K_topk]
+    stride_tib,
+    stride_tim,
+    stride_tik,
+    # OUT_LOCAL strides: [B, H, Sq, D]
+    stride_olb,
+    stride_olh,
+    stride_olm,
+    stride_old,
+    # LSE_LOCAL strides: [B, H, Sq]
+    stride_llb,
+    stride_llh,
+    stride_llm,
+    # OUT (joint) strides: [B, H, Sq, D]
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    # LSE (joint) strides: [B, H, Sq]
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    seqlen_q,
+    pool_size,
+    K_topk,
+    sm_scale,
+    HEAD_Q: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    K_DIVISIBLE: tl.constexpr,  # True iff K_topk % BLOCK_K == 0
+    H_DIVISIBLE: tl.constexpr,  # True iff HEAD_Q  % BLOCK_H == 0
+):
+    """Sparse-branch CSA FWD + joint merge fused.
+
+    Plan-8 P57 cr=4 FWD speedup: eliminates the separate
+    ``_v4_csa_attention_lse_merge_kernel`` launch by reading
+    ``out_local`` / ``lse_local`` at the tail of this kernel and writing
+    the joint output directly. Saves one kernel launch (~10 us) plus the
+    intermediate ``[B, H, Sq, D]`` ``out_sparse`` HBM round-trip
+    (~250 MiB read + 250 MiB write = ~100 us at MI355X HBM speed).
+
+    Grid: ``(seqlen_q, cdiv(HEAD_Q, BLOCK_H), B)`` — same as the
+    non-fused sparse kernel. Each program owns one ``(b, h_block, m)``
+    and produces the joint ``out[BLOCK_H, BLOCK_DMODEL]`` row.
+
+    Math (joint softmax over local + sparse + sink, all with
+    ``-m_max`` rescaling):
+
+        m_max = max(lse_local, lse_sparse, sink_h)
+        denom = exp(lse_local - m_max)
+              + exp(lse_sparse - m_max)
+              + exp(sink_h    - m_max)     # 0 if no sink
+        joint_out = (out_local * exp(lse_local  - m_max)
+                  +  out_sparse * exp(lse_sparse - m_max)) / denom
+        joint_lse = m_max + log(denom)
+
+    where ``out_sparse = acc / l_sparse`` is the normalized sparse
+    output (computed inside this kernel, never materialised to HBM).
+    """
+    pid_m = tl.program_id(0)
+    pid_h_block = tl.program_id(1)
+    bid = tl.program_id(2)
+
+    offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    if H_DIVISIBLE:
+        h_mask = tl.full([BLOCK_H], True, dtype=tl.int1)
+    else:
+        h_mask = offs_h < HEAD_Q
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    # Q tile: [BLOCK_H, BLOCK_DMODEL]
+    q_ptrs = (
+        Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd
+    )
+    if H_DIVISIBLE:
+        q = tl.load(q_ptrs)
+    else:
+        q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0)
+
+    # Online-softmax state for the sparse branch.
+    acc = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_H], dtype=tl.float32)
+
+    # Hoist loop-invariant pointer offsets out of the K-loop. Triton
+    # usually does this but spelling it out keeps the inner body lean.
+    topk_base = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim
+    pool_base = POOL + bid * stride_pb
+    pool_d_offs = offs_d[None, :] * stride_pd
+
+    for k_start in range(0, K_topk, BLOCK_K):
+        sparse_k = k_start + offs_k
+        topk_ptrs = topk_base + sparse_k * stride_tik
+        if K_DIVISIBLE:
+            topk = tl.load(topk_ptrs)
+            valid_k = (topk >= 0) & (topk < pool_size)
+        else:
+            topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1)
+            valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size)
+        safe_topk = tl.where(valid_k, topk, 0)
+
+        pool_ptrs = pool_base + safe_topk[:, None] * stride_pp + pool_d_offs
+        pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0)
+
+        qk = tl.dot(q.to(pool.dtype), tl.trans(pool), out_dtype=tl.float32) * sm_scale
+        if H_DIVISIBLE:
+            qk = tl.where(valid_k[None, :], qk, NEG_INF)
+        else:
+            qk = tl.where(h_mask[:, None] & valid_k[None, :], qk, NEG_INF)
+
+        m_tile = tl.max(qk, axis=1)
+        m_new = tl.maximum(m_i, m_tile)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(qk - m_new[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(pool.dtype), pool, out_dtype=tl.float32)
+        m_i = m_new
+
+    # Detect "all invalid" sparse rows (every topk_idx == -1 or
+    # K_topk == 0): l_sparse stays zero so the sparse branch carries
+    # zero softmax mass for that (b, h, m).
+    empty_sparse = l_i == 0.0
+    safe_l_sparse = tl.where(empty_sparse, 1.0, l_i)
+    out_sparse_norm = acc / safe_l_sparse[:, None]  # [BLOCK_H, BLOCK_DMODEL] fp32
+    lse_sparse = tl.where(empty_sparse, NEG_INF, m_i + tl.log(safe_l_sparse))
+
+    # ---- Load local-branch output + lse and merge with sink -------------
+    out_local_ptrs = (
+        OUT_LOCAL
+        + bid * stride_olb
+        + offs_h[:, None] * stride_olh
+        + pid_m * stride_olm
+        + offs_d[None, :] * stride_old
+    )
+    out_local = tl.load(out_local_ptrs, mask=h_mask[:, None], other=0.0).to(tl.float32)
+
+    lse_local_ptrs = LSE_LOCAL + bid * stride_llb + offs_h * stride_llh + pid_m * stride_llm
+    lse_local = tl.load(lse_local_ptrs, mask=h_mask, other=NEG_INF)
+
+    if HAS_SINK:
+        sink_h = tl.load(SINK + offs_h, mask=h_mask, other=NEG_INF).to(tl.float32)
+    else:
+        sink_h = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32)
+
+    m_max = tl.maximum(lse_local, lse_sparse)
+    if HAS_SINK:
+        m_max = tl.maximum(m_max, sink_h)
+
+    alpha_local = tl.exp(lse_local - m_max)
+    alpha_sparse = tl.exp(lse_sparse - m_max)
+    if HAS_SINK:
+        alpha_sink = tl.exp(sink_h - m_max)
+    else:
+        alpha_sink = tl.zeros([BLOCK_H], dtype=tl.float32)
+
+    denom = alpha_local + alpha_sparse + alpha_sink
+    # Empty-row safety: if every branch is NEG_INF, denom == 0. Use a
+    # safe denom and rely on the alpha_* being zero so the numerator is
+    # zero too — output stays at 0 / 1 = 0, lse remains NEG_INF.
+    safe_denom = tl.where(denom == 0.0, 1.0, denom)
+
+    joint_out = (out_local * alpha_local[:, None] + out_sparse_norm * alpha_sparse[:, None]) / safe_denom[
+        :, None
+    ]
+    joint_lse = tl.where(denom == 0.0, NEG_INF, m_max + tl.log(safe_denom))
+
+    out_ptrs = (
+        OUT + bid * stride_ob + offs_h[:, None] * stride_oh + pid_m * stride_om + offs_d[None, :] * stride_od
+    )
+    tl.store(out_ptrs, joint_out.to(OUT.dtype.element_ty), mask=h_mask[:, None])
+
+    lse_ptrs = LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm
+    tl.store(lse_ptrs, joint_lse, mask=h_mask)
+
+
+@triton.jit
+def _v4_csa_attention_lse_merge_kernel(
+    OUT_LOCAL,
+    LSE_LOCAL,
+    OUT_SPARSE,
+    LSE_SPARSE,
+    SINK,
+    OUT,
+    LSE,
+    stride_olb,
+    stride_olh,
+    stride_olm,
+    stride_old,
+    stride_llb,
+    stride_llh,
+    stride_llm,
+    stride_osb,
+    stride_osh,
+    stride_osm,
+    stride_osd,
+    stride_lsb,
+    stride_lsh,
+    stride_lsm,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    seqlen_q,
+    HEAD_Q: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    """Merge ``(out_local, lse_local)`` and ``(out_sparse, lse_sparse)``
+    with the per-head sink under one final online softmax.
+
+    Grid: ``(cdiv(seqlen_q, BLOCK_M), B * HEAD_Q)``. Each program owns
+    one ``[BLOCK_M, BLOCK_DMODEL]`` slice of the joint output.
+    """
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    bid = pid_bh // HEAD_Q
+    qhid = pid_bh % HEAD_Q
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    m_mask = offs_m < seqlen_q
+
+    NEG_INF: tl.constexpr = -1.0e30
+
+    lse_local = tl.load(
+        LSE_LOCAL + bid * stride_llb + qhid * stride_llh + offs_m * stride_llm,
+        mask=m_mask,
+        other=NEG_INF,
+    )
+    lse_sparse = tl.load(
+        LSE_SPARSE + bid * stride_lsb + qhid * stride_lsh + offs_m * stride_lsm,
+        mask=m_mask,
+        other=NEG_INF,
+    )
+
+    out_local = tl.load(
+        OUT_LOCAL
+        + bid * stride_olb
+        + qhid * stride_olh
+        + offs_m[:, None] * stride_olm
+        + offs_d[None, :] * stride_old,
+        mask=m_mask[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    out_sparse = tl.load(
+        OUT_SPARSE
+        + bid * stride_osb
+        + qhid * stride_osh
+        + offs_m[:, None] * stride_osm
+        + offs_d[None, :] * stride_osd,
+        mask=m_mask[:, None],
+        other=0.0,
+    ).to(tl.float32)
+
+    if HAS_SINK:
+        sink_h = tl.load(SINK + qhid).to(tl.float32)
+    else:
+        sink_h = NEG_INF
+
+    m_max = tl.maximum(lse_local, lse_sparse)
+    if HAS_SINK:
+        m_max = tl.maximum(m_max, sink_h)
+
+    alpha_local = tl.exp(lse_local - m_max)
+    alpha_sparse = tl.exp(lse_sparse - m_max)
+    if HAS_SINK:
+        alpha_sink = tl.exp(sink_h - m_max)
+    else:
+        alpha_sink = tl.zeros_like(alpha_local)
+
+    denom = alpha_local + alpha_sparse + alpha_sink
+    # Empty-row safety: if every branch is NEG_INF, denom == 0. Use a
+    # safe denom and rely on the alpha_* being zero so the numerator is
+    # zero too — output stays at 0 / 1 = 0, lse remains NEG_INF.
+    safe_denom = tl.where(denom == 0.0, 1.0, denom)
+
+    out = (out_local * alpha_local[:, None] + out_sparse * alpha_sparse[:, None]) / safe_denom[:, None]
+    lse = tl.where(denom == 0.0, NEG_INF, m_max + tl.log(safe_denom))
+
+    out_ptrs = (
+        OUT + bid * stride_ob + qhid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
+    )
+    tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=m_mask[:, None])
+
+    lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm
+    tl.store(lse_ptrs, lse, mask=m_mask)
+
+
+def _launch_v4_csa_attention_pool_fwd_split(
+    q: torch.Tensor,
+    k_local: torch.Tensor,
+    v_local: torch.Tensor,
+    pool: torch.Tensor,
+    topk_idxs: torch.Tensor,
+    *,
+    sink: Optional[torch.Tensor],
+    swa_window: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """P32 split CSA FWD: local SWA via dense kernel + sparse head-block + LSE merge."""
+    B, HQ, Sq, D = q.shape
+    P = pool.shape[1]
+    K_topk = topk_idxs.shape[2]
+    has_sink = sink is not None
+
+    # Step 1: local SWA branch (no sink — applied in the merge).
+    out_local, lse_local = _launch_v4_attention_fwd(
+        q,
+        k_local,
+        v_local,
+        sink=None,
+        swa_window=int(swa_window) if swa_window > 0 else 0,
+        additive_mask=None,
+        scale=float(scale),
+        hca_local_seqlen=0,
+    )
+
+    # P57 attempt-11: tile sweep on the fused sparse+merge kernel found
+    # (BLOCK_H=64, BLOCK_K=16, num_warps=8, num_stages=3) wins on the
+    # V4-Flash widths (H=64 D=512). The intuition:
+    #
+    # * BLOCK_H=64 keeps the full head axis in a single program, so the
+    #   pool gather (which is per-(b, m) — no H dim) is read once and
+    #   reused across all H query heads.
+    # * BLOCK_K=16 halves the per-tile pool tile (16 × 512 × 2 B = 16 KiB)
+    #   so the AMD Triton software pipeliner can keep 3 stages live in
+    #   LDS (3 × 16 KiB = 48 KiB pool buffer) without crowding out the
+    #   `acc` register accumulator. With BLOCK_K=32 the pool tile is
+    #   2× larger and 3-stage pipelining ran out of resources.
+    # * num_stages=3 overlaps pool gather → `tl.dot` → softmax update
+    #   across 3 K-tile iterations, hiding the gather latency.
+    #
+    # P57 R2 re-confirmed the BLOCK_K/stages winner via a fresh BK x
+    # ST x NW sweep (`p57/r2_sweep.sh`); BLOCK_K=32 ran out of LDS at
+    # num_stages>=2 and regressed 2-3 x. The R2 win for cr=4 FWD comes
+    # from the *local-SWA* kernel re-tile, not this kernel.
+    # BLOCK_H must be >= 16 so the head axis maps to a valid gfx1250 WMMA
+    # M-tile. Without the floor, HQ in {1, 2, 4, 8} (a power of two below
+    # 16 — e.g. high tensor-parallel head sharding, or tiny test shapes)
+    # leaves ``next_power_of_2(HQ) == HQ`` and the old ``BLOCK_H > HQ``
+    # guard never fired, so BLOCK_H stayed < 16 and the sparse ``tl.dot``
+    # failed to select a matrix-core intrinsic ("no matching matrix core
+    # intrinsic for wmma version 3 ... instruction shape [0, 0, K]").
+    # Surplus heads in the 16-wide tile are masked by H_DIVISIBLE/h_mask.
+    BLOCK_H = 64 if HQ >= 64 else max(triton.next_power_of_2(HQ), 16)
+    BLOCK_K = 16
+    BLOCK_DMODEL = D
+    sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B)
+    sink_arg = sink if has_sink else q
+
+    # P57 attempt-3: fuse sparse + merge into one kernel. The legacy
+    # 2-kernel path (sparse → ``out_sparse``/``lse_sparse``; merge)
+    # required a ~500 MiB intermediate HBM round-trip plus a ~135 us
+    # merge launch. The fused kernel reads ``out_local`` / ``lse_local``
+    # at the tail of the sparse loop and writes joint ``out`` / ``lse``
+    # directly. Env-gated A/B fallback:
+    # ``PRIMUS_V4_CSA_FWD_SEPARATE_MERGE=1`` keeps the split layout for
+    # parity debugging.
+    out = torch.empty_like(q)
+    lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+
+    if os.getenv("PRIMUS_V4_CSA_FWD_SEPARATE_MERGE", "0") != "1":
+        _v4_csa_attention_pool_sparse_merge_fwd_kernel[sparse_grid](
+            q,
+            pool,
+            topk_idxs,
+            out_local,
+            lse_local,
+            sink_arg,
+            out,
+            lse,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            q.stride(3),
+            pool.stride(0),
+            pool.stride(1),
+            pool.stride(2),
+            topk_idxs.stride(0),
+            topk_idxs.stride(1),
+            topk_idxs.stride(2),
+            out_local.stride(0),
+            out_local.stride(1),
+            out_local.stride(2),
+            out_local.stride(3),
+            lse_local.stride(0),
+            lse_local.stride(1),
+            lse_local.stride(2),
+            out.stride(0),
+            out.stride(1),
+            out.stride(2),
+            out.stride(3),
+            lse.stride(0),
+            lse.stride(1),
+            lse.stride(2),
+            Sq,
+            P,
+            K_topk,
+            float(scale),
+            HEAD_Q=HQ,
+            HAS_SINK=has_sink,
+            BLOCK_H=BLOCK_H,
+            BLOCK_K=BLOCK_K,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            K_DIVISIBLE=(K_topk % BLOCK_K == 0),
+            H_DIVISIBLE=(HQ % BLOCK_H == 0),
+            num_warps=8,
+            num_stages=3,
+        )
+        return out, lse
+
+    # ---- legacy split-merge path (env-gated for A/B debugging) ----------
+    out_sparse = torch.empty_like(q)
+    lse_sparse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32)
+    _v4_csa_attention_pool_sparse_fwd_kernel[sparse_grid](
+        q,
+        pool,
+        topk_idxs,
+        out_sparse,
+        lse_sparse,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),
+        pool.stride(0),
+        pool.stride(1),
+        pool.stride(2),
+        topk_idxs.stride(0),
+        topk_idxs.stride(1),
+        topk_idxs.stride(2),
+        out_sparse.stride(0),
+        out_sparse.stride(1),
+        out_sparse.stride(2),
+        out_sparse.stride(3),
+        lse_sparse.stride(0),
+        lse_sparse.stride(1),
+        lse_sparse.stride(2),
+        Sq,
+        P,
+        K_topk,
+        float(scale),
+        HEAD_Q=HQ,
+        BLOCK_H=BLOCK_H,
+        BLOCK_K=BLOCK_K,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=8,
+        num_stages=1,
+    )
+    MERGE_BLOCK_M = 32
+    merge_grid = (triton.cdiv(Sq, MERGE_BLOCK_M), B * HQ)
+    _v4_csa_attention_lse_merge_kernel[merge_grid](
+        out_local,
+        lse_local,
+        out_sparse,
+        lse_sparse,
+        sink_arg,
+        out,
+        lse,
+        out_local.stride(0),
+        out_local.stride(1),
+        out_local.stride(2),
+        out_local.stride(3),
+        lse_local.stride(0),
+        lse_local.stride(1),
+        lse_local.stride(2),
+        out_sparse.stride(0),
+        out_sparse.stride(1),
+        out_sparse.stride(2),
+        out_sparse.stride(3),
+        lse_sparse.stride(0),
+        lse_sparse.stride(1),
+        lse_sparse.stride(2),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        out.stride(3),
+        lse.stride(0),
+        lse.stride(1),
+        lse.stride(2),
+        Sq,
+        HEAD_Q=HQ,
+        HAS_SINK=has_sink,
+        BLOCK_M=MERGE_BLOCK_M,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        num_warps=4,
+        num_stages=1,
+    )
+    return out, lse
+
+
+__all__ = [
+    "_v4_csa_attention_fwd_kernel",
+    "_v4_csa_attention_pool_fwd_kernel",
+    "_v4_csa_attention_pool_sparse_fwd_kernel",
+    "_v4_csa_attention_pool_sparse_merge_fwd_kernel",
+    "_v4_csa_attention_lse_merge_kernel",
+    "_launch_v4_csa_attention_fwd",
+    "_launch_v4_csa_attention_pool_fwd",
+    "_launch_v4_csa_attention_pool_fwd_split",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py
new file mode 100644
index 000000000..5aa001f2e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py
@@ -0,0 +1,24 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plain-Triton DeepSeek-V4 sparse-MLA attention backend ("triton v2").
+
+Same fused single-latent (K == V) sparse-MLA representation and public API as
+the gluon backend, but written in vanilla Triton so the QK / PV GEMMs lower to
+MFMA via ``tl.dot``. Distinct from the separate-KV Triton CSA/dense backends
+(``_triton_v0_deprecated`` gathered / ``_triton_v1`` pool), which keep K and V separate.
+
+* :func:`sparse_mla_fwd_v4_triton` -> ``(o, lse)``
+* :func:`sparse_mla_bwd_v4_triton` -> ``(dq, dkv, d_sink)``
+"""
+
+from .dsa_bwd_v4_triton import sparse_mla_bwd_v4_triton
+from .dsa_fwd_v4_triton import sparse_mla_fwd_v4_triton
+
+__all__ = [
+    "sparse_mla_fwd_v4_triton",
+    "sparse_mla_bwd_v4_triton",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py
new file mode 100644
index 000000000..2c910de7e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py
@@ -0,0 +1,44 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared AMD Triton compiler-knob helpers for the triton_v2 sparse-MLA kernels.
+
+primus_turbo globally enables ``TRITON_HIP_USE_BLOCK_PINGPONG`` /
+``TRITON_HIP_USE_ASYNC_COPY`` (via ``set_triton_knobs_gfx950``) as soon as it is
+imported by the training runtime. Those knobs double-buffer (ping-pong) a
+kernel's LDS operand tiles. Measured on gfx950 (bench_v4_attention, flash) they
+are a *pessimization* for the V4 sparse-MLA kernels — the forward is ~16-29%
+slower and the backward is slightly slower with them on — and they also overflow
+the 160 KB LDS limit for the wide (BH=64/TK=128) dKV tiling.
+
+These knobs are read at *compile time* and are NOT part of Triton's compile
+cache key (``HIPOptions.hash`` does not hash them), so compiling a kernel inside
+:func:`amd_pingpong_disabled` pins that kernel to the non-ping-pong schedule for
+the whole process, while restoring the knobs on exit leaves every other kernel
+(compiled outside the scope) exactly as primus_turbo configured it.
+"""
+
+import contextlib
+
+import triton
+
+
+@contextlib.contextmanager
+def amd_pingpong_disabled():
+    """Temporarily disable the AMD Triton ping-pong / async-copy LDS knobs."""
+    amd = getattr(getattr(triton, "knobs", None), "amd", None)
+    if amd is None:
+        yield
+        return
+    prev_pp = amd.use_block_pingpong
+    prev_ac = amd.use_async_copy
+    try:
+        amd.use_block_pingpong = False
+        amd.use_async_copy = False
+        yield
+    finally:
+        amd.use_block_pingpong = prev_pp
+        amd.use_async_copy = prev_ac
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py
new file mode 100644
index 000000000..8f2478801
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py
@@ -0,0 +1,337 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Owned plain-Triton backward compute kernels for the "triton v2" sparse-MLA
+backend. Forked from the shared ``_gluon_dsa/_dsa_bwd_gather.py`` plain-Triton
+kernels so this backend can be tuned independently of the gluon path.
+
+Three kernels implement the non-atomic chunked-gather backward:
+  * ``_bwd_chunk_dq_store_ds`` — dQ accumulation for one rank chunk, ALSO stores
+    per-tile dS and P to [T, H, R_CHUNK] buffers for reuse by the dKV kernel.
+  * ``_bwd_compute_dkv_intermediate`` — dKV intermediate for one chunk, REUSES
+    the stored dS/P (no S/P/dS recompute). Consumes q/do transposed to [T, D, H].
+  * ``_bwd_dkv_gather_acc`` — CSR inverted-topk reduce of the intermediate into
+    the fp32 dKV accumulator.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _bwd_chunk_dq_store_ds(
+    Q_ptr,  # [T, H, D] bf16
+    KV_ptr,  # [T, 1, D] bf16
+    dO_ptr,  # [T, H, D_V] bf16
+    TopK_ptr,  # [T, TOPK] int32
+    LSE_ptr,  # [T, H] fp32
+    Delta_ptr,  # [T, H] fp32 (computed here on the first chunk, read after)
+    O_ptr,  # [T, H, D_V] bf16 (fwd output, for Delta = rowsum(O*dO))
+    dQ_ptr,  # [T, H, D] bf16 — read-modify-write across chunks
+    dS_ptr,  # [T, H, R_CHUNK] bf16 — output chunk dS
+    P_ptr,  # [T, H, R_CHUNK] bf16 — output chunk P
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_kv_t: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_o_t: tl.int64,
+    stride_o_h: tl.int64,
+    stride_dq_t: tl.int64,
+    stride_dq_h: tl.int64,
+    stride_topk_t: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    scale: tl.float32,
+    num_heads: tl.int32,
+    R_START: tl.int32,
+    R_CHUNK: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    TILE_K: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    HAS_ROPE: tl.constexpr,
+    IS_FIRST_CHUNK: tl.constexpr,
+):
+    """dQ accumulation for rank chunk [R_START, R_START+R_CHUNK), plus stores
+    chunk dS and P to [T, H, R_CHUNK] buffers for _bwd_compute_dkv_intermediate.
+    Grid: (total_tokens, num_hg).
+
+    Delta = rowsum(O*dO) is fused here (computed + stored on the first chunk,
+    reloaded on later chunks) instead of a separate preprocess kernel — dO is
+    already resident, so this only adds the O load and drops a whole kernel + a
+    duplicate dO read.
+
+    HAS_ROPE=False (V4 zero-pad): skips all rope MMAs/loads and writes the dQ rope
+    columns as zero (they are discarded downstream but kept well-defined)."""
+    token_idx = tl.program_id(0)
+    hg_idx = tl.program_id(1)
+    offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offs_h < num_heads
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    q_base = token_idx * stride_q_t
+    Q_lora = tl.load(
+        Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    if HAS_ROPE:
+        Q_rope = tl.load(
+            Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]),
+            mask=mask_h[:, None],
+            other=0.0,
+        )
+    do_base = token_idx * stride_do_t
+    dO_val = tl.load(
+        dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0
+    )
+    lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+    if IS_FIRST_CHUNK:
+        # Delta = rowsum(O * dO): fold the preprocess in (dO already loaded).
+        O_val = tl.load(
+            O_ptr + token_idx * stride_o_t + offs_h[:, None] * stride_o_h + offs_v[None, :],
+            mask=mask_h[:, None],
+            other=0.0,
+        )
+        delta = tl.sum(O_val.to(tl.float32) * dO_val.to(tl.float32), axis=1)
+        tl.store(Delta_ptr + token_idx * num_heads + offs_h, delta, mask=mask_h)
+    else:
+        delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0)
+
+    dq_base = token_idx * stride_dq_t
+    if IS_FIRST_CHUNK:
+        dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32)
+    else:
+        dQ_lora = tl.load(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+            mask=mask_h[:, None],
+            other=0.0,
+        ).to(tl.float32)
+    if HAS_ROPE:
+        if IS_FIRST_CHUNK:
+            dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32)
+        else:
+            dQ_rope = tl.load(
+                dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+                mask=mask_h[:, None],
+                other=0.0,
+            ).to(tl.float32)
+
+    NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K
+    topk_base = token_idx * stride_topk_t + R_START
+    offs_tile = tl.arange(0, TILE_K)
+    ds_base = token_idx * stride_ds_t + hg_idx * BLOCK_H * stride_ds_h
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        tile_offs = tile_start + offs_tile
+        valid = tile_offs < R_CHUNK
+        topk_pos = tl.load(TopK_ptr + topk_base + tile_offs, mask=valid, other=-1)
+        valid = valid & (topk_pos != -1)
+        safe_pos = tl.where(valid, topk_pos, 0)
+
+        K_lora_T = tl.load(
+            KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0
+        )
+
+        S = tl.dot(Q_lora, K_lora_T)
+        if HAS_ROPE:
+            K_rope_T = tl.load(
+                KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]),
+                mask=valid[None, :],
+                other=0.0,
+            )
+            S += tl.dot(Q_rope, K_rope_T)
+        S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf"))
+        P = tl.exp(S - lse[:, None])
+        P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0)
+        dP = tl.dot(dO_val, K_lora_T)
+        dS = P * (dP - delta[:, None]) * scale
+        dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0)
+
+        dS_bf = dS.to(tl.bfloat16)
+        dQ_lora += tl.dot(dS_bf, tl.trans(K_lora_T)).to(tl.float32)
+        if HAS_ROPE:
+            dQ_rope += tl.dot(dS_bf, tl.trans(K_rope_T)).to(tl.float32)
+
+        local_h = tl.arange(0, BLOCK_H)
+        tl.store(
+            dS_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :],
+            dS_bf,
+            mask=mask_h[:, None] & valid[None, :],
+        )
+        tl.store(
+            P_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :],
+            P.to(tl.bfloat16),
+            mask=mask_h[:, None] & valid[None, :],
+        )
+
+    tl.store(
+        dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :],
+        dQ_lora.to(Q_lora.dtype),
+        mask=mask_h[:, None],
+    )
+    if HAS_ROPE:
+        tl.store(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+            dQ_rope.to(Q_lora.dtype),
+            mask=mask_h[:, None],
+        )
+    else:
+        tl.store(
+            dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]),
+            tl.zeros([BLOCK_H, D_ROPE], dtype=Q_lora.dtype),
+            mask=mask_h[:, None],
+        )
+
+
+@triton.jit
+def _bwd_compute_dkv_intermediate(
+    Q_ptr,  # [T, H, D_QK] bf16  (UNtransposed; loaded transposed via strided index)
+    dO_ptr,  # [T, H, D_V]  bf16 (UNtransposed)
+    dS_ptr,  # [T, H, R_CHUNK] bf16
+    P_ptr,  # [T, H, R_CHUNK] bf16
+    Interm_ptr,  # [T, R_CHUNK, D_QK] bf16 — output, one writer per (q, rank)
+    stride_q_t: tl.int64,
+    stride_q_h: tl.int64,
+    stride_do_t: tl.int64,
+    stride_do_h: tl.int64,
+    stride_ds_t: tl.int64,
+    stride_ds_h: tl.int64,
+    stride_interm_t: tl.int64,  # R_CHUNK * D_QK
+    stride_interm_k: tl.int64,  # D_QK
+    num_heads: tl.int32,
+    R_CHUNK: tl.constexpr,
+    TILE_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NUM_HG: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    HAS_ROPE: tl.constexpr,
+):
+    """dKV intermediate for one chunk, REUSING stored dS/P (no recompute).
+    Loads Q/dO UNtransposed with a [D, BLOCK_H] strided index pattern (contiguous
+    along D per head), removing the external q.transpose(1,2).contiguous() /
+    do.transpose(...).contiguous() copies. Grid: (total_tokens,).
+
+    HAS_ROPE=False (V4 zero-pad): skips Q_rope load, the dKV_rope MMA and the
+    interm rope store (interm rope columns are never read by the gather)."""
+    token_idx = tl.program_id(0)
+
+    NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K
+    offs_tile = tl.arange(0, TILE_K)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    interm_base_t = token_idx * stride_interm_t
+    q_base = token_idx * stride_q_t
+    do_base = token_idx * stride_do_t
+    ds_base = token_idx * stride_ds_t
+
+    for t in range(NUM_TILES):
+        tile_start = t * TILE_K
+        tile_offs = tile_start + offs_tile
+        valid = tile_offs < R_CHUNK
+
+        dKV_lora = tl.zeros([D_V, TILE_K], dtype=tl.float32)
+        if HAS_ROPE:
+            dKV_rope = tl.zeros([D_ROPE, TILE_K], dtype=tl.float32)
+
+        for hg in range(NUM_HG):
+            offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H)
+            mask_h = offs_h < num_heads
+
+            # [D_V, BLOCK_H] loaded transposed: rows=offs_v (stride 1, contiguous),
+            # cols=offs_h (stride stride_q_h). No HBM transpose copy needed.
+            Q_lora_T = tl.load(
+                Q_ptr + q_base + offs_h[None, :] * stride_q_h + offs_v[:, None],
+                mask=mask_h[None, :],
+                other=0.0,
+            )
+            dO_T = tl.load(
+                dO_ptr + do_base + offs_h[None, :] * stride_do_h + offs_v[:, None],
+                mask=mask_h[None, :],
+                other=0.0,
+            )
+
+            dS_val = tl.load(
+                dS_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :],
+                mask=mask_h[:, None] & valid[None, :],
+                other=0.0,
+            )
+            P_val = tl.load(
+                P_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :],
+                mask=mask_h[:, None] & valid[None, :],
+                other=0.0,
+            )
+
+            dKV_lora += tl.dot(Q_lora_T, dS_val.to(Q_lora_T.dtype)).to(tl.float32)
+            dKV_lora += tl.dot(dO_T, P_val.to(dO_T.dtype)).to(tl.float32)
+            if HAS_ROPE:
+                Q_rope_T = tl.load(
+                    Q_ptr + q_base + offs_h[None, :] * stride_q_h + (D_V + offs_r[:, None]),
+                    mask=mask_h[None, :],
+                    other=0.0,
+                )
+                dKV_rope += tl.dot(Q_rope_T, dS_val.to(Q_rope_T.dtype)).to(tl.float32)
+
+        interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + offs_v[:, None]
+        tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[None, :])
+
+        if HAS_ROPE:
+            interm_rope_ptrs = (
+                Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + D_V + offs_r[:, None]
+            )
+            tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[None, :])
+
+
+@triton.jit
+def _bwd_dkv_gather_acc(
+    Interm_ptr,  # [T, R_CHUNK, D] bf16 — chunk intermediate
+    InvPtr_ptr,  # [T+1] int32 — CSR row pointers
+    InvData_ptr,  # [T*R_CHUNK] int32 — encoded as q*R_CHUNK + local_r
+    dKV_acc_ptr,  # [T, D] fp32 — accumulator (read-modify-write across chunks)
+    stride_interm_r: tl.int64,  # D
+    stride_acc_t: tl.int64,  # D
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    HAS_ROPE: tl.constexpr,
+    BLOCK_K: tl.constexpr = 64,
+):
+    """Gather one chunk's bf16 intermediate into the fp32 dKV accumulator.
+    Grid: (num_kv,) — one CTA per KV token, no atomics. Tiled BLOCK_K rows/iter.
+
+    HAS_ROPE=False (V4 zero-pad): interm rope columns are never written, so skip
+    reading/accumulating them; dKV_acc rope stays at its zero-init value."""
+    k = tl.program_id(0)
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    start = tl.load(InvPtr_ptr + k)
+    end = tl.load(InvPtr_ptr + k + 1)
+
+    acc_base = k.to(tl.int64) * stride_acc_t
+    dkv_acc_lora = tl.load(dKV_acc_ptr + acc_base + offs_v).to(tl.float32)
+    if HAS_ROPE:
+        dkv_acc_rope = tl.load(dKV_acc_ptr + acc_base + D_V + offs_r).to(tl.float32)
+
+    n_entries = end - start
+    for ti in range(0, tl.cdiv(n_entries, BLOCK_K)):
+        e_local = ti * BLOCK_K + offs_k
+        valid = e_local < n_entries
+        entry = tl.load(InvData_ptr + start + e_local, mask=valid, other=0).to(tl.int64)
+        rp = entry[:, None] * stride_interm_r
+        rows_v = tl.load(Interm_ptr + rp + offs_v[None, :], mask=valid[:, None], other=0.0).to(tl.float32)
+        dkv_acc_lora += tl.sum(rows_v, axis=0)
+        if HAS_ROPE:
+            rows_r = tl.load(Interm_ptr + rp + D_V + offs_r[None, :], mask=valid[:, None], other=0.0).to(
+                tl.float32
+            )
+            dkv_acc_rope += tl.sum(rows_r, axis=0)
+
+    tl.store(dKV_acc_ptr + acc_base + offs_v, dkv_acc_lora)
+    if HAS_ROPE:
+        tl.store(dKV_acc_ptr + acc_base + D_V + offs_r, dkv_acc_rope)
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py
new file mode 100644
index 000000000..1256d7ade
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py
@@ -0,0 +1,229 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plain-Triton DeepSeek-V4 sparse-MLA backward (the "triton v2" backend).
+
+Companion to :func:`sparse_mla_fwd_v4_triton`; API mirrors
+:func:`sparse_mla_bwd_v4_gluon` -> ``(dq, dkv, d_sink)``.
+
+Uses the **non-atomic chunked-gather** scheme (the same one the gluon backward
+shares for its dKV gather): per rank-chunk it runs a pure-Triton dQ kernel and a
+dKV-intermediate kernel (both ``tl.dot`` / MFMA, plain stores — no atomics),
+then reduces the intermediate into ``dkv`` via a CSR inverted-topk gather. This
+is the fully-Triton analogue of the gluon dQ/dKV-interm kernels, so its dKV is
+not bottlenecked by global atomics. ``d_sink`` is the closed-form torch
+reduction ``-sum_t exp(sink - lse) * delta``.
+"""
+
+import contextlib
+import os
+
+import torch
+import triton
+
+# CSR-inverted-topk builder + Delta preprocess are backend-neutral torch/host
+# helpers; reuse them. The compute kernels are owned locally (dsa_bwd_kernels)
+# so this backend can be tuned independently of the gluon path.
+from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice
+from ._amd_knobs import amd_pingpong_disabled
+from .dsa_bwd_kernels import (
+    _bwd_chunk_dq_store_ds,
+    _bwd_compute_dkv_intermediate,
+    _bwd_dkv_gather_acc,
+)
+
+
+def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA backward (plain Triton / MFMA, non-atomic).
+
+    Returns ``(dq, dkv, d_sink)`` with ``dkv`` shaped like ``kv`` and ``d_sink``
+    ``[num_heads]`` fp32 (or ``None`` when ``attn_sink`` is None).
+    """
+    assert q.is_contiguous() and o.is_contiguous() and do.is_contiguous()
+    assert topk_indices.is_contiguous() and lse.is_contiguous()
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    assert kv.is_contiguous()
+    num_kv = kv.shape[0]
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,)
+
+    # Delta = rowsum(O*dO) is fused into the first dQ chunk (no preprocess kernel).
+    delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    # ---- config (mirror the gluon bwd chunking) ----
+    # R_CHUNK (rank-chunk width): dQ is read-modify-written across chunks, so more
+    # chunks = more redundant dq reload passes + repeated launches/CSR builds. For
+    # high head counts (H>=128) the dq RMW volume is large, so a single chunk over
+    # the whole topk (bounded for memory) is a big win (-22% on pro cr4). For low
+    # head counts (H<=64) the smaller dq RMW is outweighed by the larger per-chunk
+    # buffers/occupancy, and 256 stays best — so keep the cap there.
+    if num_heads >= 128:
+        R_CHUNK = min(topk, 1536)
+    else:
+        R_CHUNK = min(256, topk)
+    BH_DQ, TK_DQ = 64, 16
+    # dKV-intermediate tiling. The default (BH_DKV=32, TK_DKV=64) is best for high
+    # head counts (H>=128) and for chunk widths that are not 128-aligned. For low
+    # head counts (H<=64) with a 128-aligned chunk, a wider TILE_K=128 over a single
+    # head-group (BH_DKV=64, NUM_HG=1) reduces redundant Q/dO re-loads and issues
+    # fuller MMAs (measured ~6-7% faster on the full flash bwd: cr=0 1.21->1.14 ms,
+    # cr=4 5.44->5.09 ms in bench_v4_attention); it regresses for H>=128 (register
+    # pressure) and for non-128-aligned chunks (partial tiles), so it is guarded.
+    #
+    # The wide config's per-launch LDS overflows the 160 KB limit ONLY when the AMD
+    # ping-pong / async-copy knobs are on (primus_turbo's set_triton_knobs_gfx950()
+    # enables them globally in training), which double-buffer the LDS operand tiles
+    # and ~double this kernel's shared memory (BH64/TK128 R_CHUNK256: fits
+    # standalone, 347904 B in training). Since the whole bwd now compiles with those
+    # knobs disabled (see the amd_pingpong_disabled scope below), the wide config
+    # fits in both the benchmark and training, so it is the default where it applies
+    # (~6-7% faster on the full flash bwd). Set PRIMUS_DSA_DKV_SAFE=1 to force the
+    # narrow 32/64 (which fits regardless of the knobs).
+    use_wide_dkv = (
+        num_heads <= 64 and R_CHUNK % 128 == 0 and os.environ.get("PRIMUS_DSA_DKV_SAFE", "0") != "1"
+    )
+    BH_DKV, TK_DKV = (64, 128) if use_wide_dkv else (32, 64)
+    num_hg_dq = triton.cdiv(num_heads, BH_DQ)
+    num_hg_dkv = triton.cdiv(num_heads, BH_DKV)
+
+    # In the V4 single-latent form the D_ROPE block of q/kv is a zero pad (RoPE is
+    # baked in-place over the 512 latent) and its gradient is discarded by the
+    # adapter (dq[..., :D_V], dkv[..., :D_V]). So all rope MMAs / K_rope loads /
+    # interm-rope traffic compute a provably-zero result that is thrown away.
+    # Skip them: bit-identical non-rope outputs, zero rope outputs (== gluon).
+    HAS_ROPE = False
+
+    dq = torch.empty_like(q)
+    chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device)
+    dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device)
+    interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device)
+
+    # pad topk to an R_CHUNK multiple (-1 = invalid).
+    topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK
+    if topk_padded_len != topk:
+        pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device)
+        topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous()
+    else:
+        topk_padded = topk_indices
+
+    all_csr = [
+        _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv)
+        for rs in range(0, topk, R_CHUNK)
+    ]
+
+    # The AMD ping-pong / async-copy knobs primus_turbo enables globally are a
+    # pessimization for the whole triton_v2 backward (measured ~5-7% slower on the
+    # full bwd for both flash and pro in bench_v4_attention), and they overflow the
+    # 160 KB LDS limit for the wide dKV tiling. Compile the entire bwd (dQ /
+    # dKV-intermediate / gather) with them disabled; the knobs are read at compile
+    # time and are not in Triton's cache key, so this pins the faster non-ping-pong
+    # schedule for these kernels without touching any kernel compiled elsewhere.
+    # PRIMUS_DSA_BWD_PINGPONG_OFF=0 keeps the ambient knobs (the wide dKV still
+    # forces them off below, since it does not fit otherwise).
+    _bwd_ctx = (
+        amd_pingpong_disabled()
+        if os.environ.get("PRIMUS_DSA_BWD_PINGPONG_OFF", "1") == "1"
+        else contextlib.nullcontext()
+    )
+    with _bwd_ctx:
+        for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)):
+            is_first = r_start == 0
+
+            _bwd_chunk_dq_store_ds[(total_tokens, num_hg_dq)](
+                q,
+                kv,
+                do,
+                topk_padded,
+                lse,
+                delta,
+                o,
+                dq,
+                chunk_dS,
+                chunk_P,
+                q.stride(0),
+                q.stride(1),
+                kv.stride(0),
+                do.stride(0),
+                do.stride(1),
+                o.stride(0),
+                o.stride(1),
+                dq.stride(0),
+                dq.stride(1),
+                topk_padded.stride(0),
+                chunk_dS.stride(0),
+                chunk_dS.stride(1),
+                scale,
+                num_heads,
+                r_start,
+                R_CHUNK=R_CHUNK,
+                BLOCK_H=BH_DQ,
+                TILE_K=TK_DQ,
+                D_V=kv_lora_rank,
+                D_ROPE=rope_rank,
+                HAS_ROPE=HAS_ROPE,
+                IS_FIRST_CHUNK=is_first,
+                num_warps=4,
+                waves_per_eu=1,
+            )
+
+            # The wide dKV kernel MUST compile with ping-pong off to fit LDS, even
+            # if the outer bwd gate is disabled — force it off here regardless.
+            _dkv_ctx = amd_pingpong_disabled() if use_wide_dkv else contextlib.nullcontext()
+            with _dkv_ctx:
+                _bwd_compute_dkv_intermediate[(total_tokens,)](
+                    q,
+                    do,
+                    chunk_dS,
+                    chunk_P,
+                    interm,
+                    q.stride(0),
+                    q.stride(1),
+                    do.stride(0),
+                    do.stride(1),
+                    chunk_dS.stride(0),
+                    chunk_dS.stride(1),
+                    interm.stride(0),
+                    interm.stride(1),
+                    num_heads,
+                    R_CHUNK=R_CHUNK,
+                    TILE_K=TK_DKV,
+                    BLOCK_H=BH_DKV,
+                    NUM_HG=num_hg_dkv,
+                    D_V=kv_lora_rank,
+                    D_ROPE=rope_rank,
+                    HAS_ROPE=HAS_ROPE,
+                    num_warps=4,
+                )
+
+            inv_ptr, inv_data = all_csr[chunk_idx]
+            _bwd_dkv_gather_acc[(num_kv,)](
+                interm,
+                inv_ptr,
+                inv_data,
+                dkv_acc,
+                interm.stride(1),
+                dkv_acc.stride(0),
+                D_V=kv_lora_rank,
+                D_ROPE=rope_rank,
+                HAS_ROPE=HAS_ROPE,
+                num_warps=4,
+            )
+
+    d_sink = None
+    if has_sink:
+        d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0)
+
+    dkv = dkv_acc.to(kv.dtype).unsqueeze(1)
+    return dq, dkv, d_sink
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py
new file mode 100644
index 000000000..3f6c58b1e
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py
@@ -0,0 +1,206 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plain-Triton DeepSeek-V4 sparse-MLA forward (the "triton v2" backend).
+
+Same sparse-MLA latent representation and public API as the gluon backend
+(:func:`sparse_mla_fwd_v4_gluon`) — fused single MQA latent (K = V = the first
+``kv_lora_rank`` channels), per-token absolute top-k indices, optional per-head
+softmax sink — but written in vanilla Triton so the QK / PV GEMMs lower to MFMA
+through ``tl.dot`` (no hand-rolled gluon layouts). One program handles one query
+token and a block of heads; it gathers the selected KV latent rows tile-by-tile
+and runs an online (flash) softmax with a sink-augmented denominator.
+
+This contrasts with the in-tree Triton CSA backend (``_triton/v4_csa_*``) which
+keeps K and V separate to share kernels with the dense path; here K == V is a
+single latent (matching gluon / the V4 paper), which is the whole point of v2.
+"""
+
+import contextlib
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+from ._amd_knobs import amd_pingpong_disabled
+
+
+def _get_fwd_configs():
+    # Focused around the autotune winners from the wide sweep (TILE_K=16,
+    # num_stages=3 dominated -> latency-bound; deep pipelining is the lever), plus
+    # num_stages=4 to probe deeper pipelining. Keeps first-call autotune cheap.
+    return [
+        triton.Config({"BLOCK_H": bh, "TILE_K": tk, "waves_per_eu": wpe}, num_warps=4, num_stages=ns)
+        for bh in (32, 64)
+        for tk in (16, 32)
+        for ns in (2, 3, 4)
+        for wpe in (0, 1)
+    ]
+
+
+@triton.autotune(configs=_get_fwd_configs(), key=["num_heads", "TOPK", "D_V", "D_ROPE", "HAS_ROPE"])
+@triton.jit
+def _sparse_mla_fwd_tr_kernel(
+    Q_ptr,  # [total_tokens, num_heads, D_QK] bf16
+    KV_ptr,  # [num_kv, 1, D_QK]               bf16
+    TopK_ptr,  # [total_tokens, TOPK]            int32
+    Sink_ptr,  # [num_heads]                     fp32 (guarded by HAS_SINK)
+    O_ptr,  # [total_tokens, num_heads, D_V]  bf16
+    LSE_ptr,  # [total_tokens, num_heads]       fp32 (sink-inclusive)
+    stride_q_t,
+    stride_q_h,
+    stride_kv_t,
+    stride_o_t,
+    stride_o_h,
+    stride_topk_t,
+    scale,
+    num_heads,
+    TOPK: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    TILE_K: tl.constexpr,
+    D_V: tl.constexpr,
+    D_ROPE: tl.constexpr,
+    HAS_ROPE: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+):
+    tok = tl.program_id(0)
+    hg = tl.program_id(1)
+
+    offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = offs_h < num_heads
+    offs_v = tl.arange(0, D_V)
+    offs_r = tl.arange(0, D_ROPE)
+
+    q_base = tok.to(tl.int64) * stride_q_t + offs_h.to(tl.int64)[:, None] * stride_q_h
+    q_lora = tl.load(Q_ptr + q_base + offs_v[None, :], mask=mask_h[:, None], other=0.0)
+    if HAS_ROPE:
+        q_rope = tl.load(Q_ptr + q_base + (D_V + offs_r)[None, :], mask=mask_h[:, None], other=0.0)
+
+    m_i = tl.full([BLOCK_H], float("-inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, D_V], dtype=tl.float32)
+
+    topk_base = tok.to(tl.int64) * stride_topk_t
+    for kt in range(0, TOPK, TILE_K):
+        offs_k = kt + tl.arange(0, TILE_K)
+        idx = tl.load(TopK_ptr + topk_base + offs_k, mask=offs_k < TOPK, other=-1)
+        valid = idx >= 0
+        safe = tl.where(valid, idx, 0).to(tl.int64)
+
+        kv_base = safe[:, None] * stride_kv_t
+        k_lora = tl.load(KV_ptr + kv_base + offs_v[None, :], mask=valid[:, None], other=0.0)
+
+        # S = q @ k^T over [lora ++ rope]  -> [BLOCK_H, TILE_K]
+        s = tl.dot(q_lora, tl.trans(k_lora))
+        if HAS_ROPE:
+            k_rope = tl.load(KV_ptr + kv_base + (D_V + offs_r)[None, :], mask=valid[:, None], other=0.0)
+            s += tl.dot(q_rope, tl.trans(k_rope))
+        s = s * scale
+        s = tl.where(valid[None, :] & mask_h[:, None], s, float("-inf"))
+
+        m_ij = tl.max(s, axis=1)
+        m_new = tl.maximum(m_i, m_ij)
+        m_new = tl.where(m_new > float("-inf"), m_new, 0.0)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(s - m_new[:, None])
+        l_i = alpha * l_i + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(k_lora.dtype), k_lora)
+        m_i = m_new
+
+    if HAS_SINK:
+        sink = tl.load(Sink_ptr + offs_h, mask=mask_h, other=float("-inf"))
+        m_f = tl.maximum(m_i, sink)
+        af = tl.exp(m_i - m_f)
+        l_t = l_i * af + tl.exp(sink - m_f)
+        acc = acc * af[:, None]
+        acc = acc / l_t[:, None]
+        lse = m_f + tl.log(l_t)
+    else:
+        acc = acc / l_i[:, None]
+        lse = m_i + tl.log(l_i)
+
+    o_base = tok.to(tl.int64) * stride_o_t + offs_h.to(tl.int64)[:, None] * stride_o_h
+    tl.store(O_ptr + o_base + offs_v[None, :], acc.to(O_ptr.dtype.element_ty), mask=mask_h[:, None])
+    tl.store(LSE_ptr + tok.to(tl.int64) * num_heads + offs_h, lse, mask=mask_h)
+
+
+def sparse_mla_fwd_v4_triton(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None):
+    """DeepSeek-V4 sparse-MLA forward (plain Triton / MFMA). API mirrors the gluon path.
+
+    Args:
+        q:            [total_tokens, num_heads, d_qk] bf16
+        kv:           [num_kv, 1, d_qk] bf16 (or [num_kv, d_qk]); single MQA latent
+        topk_indices: [total_tokens, topk] int32 (SWA + sparse, -1 = invalid)
+        attn_sink:    [num_heads] fp32 optional per-head learnable sink
+        kv_lora_rank: int, default 512
+        scale:        float, default 1/sqrt(d_qk)
+
+    Returns:
+        o:   [total_tokens, num_heads, kv_lora_rank] (q.dtype)
+        lse: [total_tokens, num_heads] fp32 (sink-inclusive when attn_sink given)
+    """
+    assert q.is_contiguous() and topk_indices.is_contiguous()
+    total_tokens, num_heads, d_qk = q.shape
+    rope_rank = d_qk - kv_lora_rank
+    topk = topk_indices.shape[1]
+    if scale is None:
+        scale = 1.0 / (d_qk**0.5)
+    if kv.dim() == 2:
+        kv = kv.unsqueeze(1)
+    assert kv.is_contiguous()
+    assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk
+
+    has_sink = attn_sink is not None
+    if has_sink:
+        assert attn_sink.is_contiguous() and attn_sink.dtype == torch.float32
+        assert attn_sink.shape == (num_heads,)
+        sink_ptr = attn_sink
+    else:
+        sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device)
+
+    o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device)
+    lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device)
+
+    # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked
+    # in-place over the 512 latent), so the rope QK term is provably zero — skip it.
+    has_rope = False
+
+    grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"]))
+    # The AMD ping-pong / async-copy knobs primus_turbo enables globally are a
+    # pessimization for this fwd kernel (~16-29% slower on flash in
+    # bench_v4_attention). They are read at compile time and are not part of
+    # Triton's cache key, so autotuning/compiling this kernel with them disabled
+    # pins the faster (non-ping-pong) schedule for the process without touching
+    # any other kernel. PRIMUS_DSA_FWD_PINGPONG_OFF=0 keeps the ambient knobs.
+    _fwd_ctx = (
+        amd_pingpong_disabled()
+        if os.environ.get("PRIMUS_DSA_FWD_PINGPONG_OFF", "1") == "1"
+        else contextlib.nullcontext()
+    )
+    with _fwd_ctx:
+        _sparse_mla_fwd_tr_kernel[grid](
+            Q_ptr=q,
+            KV_ptr=kv,
+            TopK_ptr=topk_indices,
+            Sink_ptr=sink_ptr,
+            O_ptr=o,
+            LSE_ptr=lse,
+            stride_q_t=q.stride(0),
+            stride_q_h=q.stride(1),
+            stride_kv_t=kv.stride(0),
+            stride_o_t=o.stride(0),
+            stride_o_h=o.stride(1),
+            stride_topk_t=topk_indices.stride(0),
+            scale=scale,
+            num_heads=num_heads,
+            TOPK=topk,
+            D_V=kv_lora_rank,
+            D_ROPE=rope_rank,
+            HAS_ROPE=has_rope,
+            HAS_SINK=has_sink,
+        )
+    return o, lse
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py
new file mode 100644
index 000000000..c850286f4
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py
@@ -0,0 +1,62 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 sparse-MLA kernel-pair via the **Primus-Turbo** public API.
+
+Unlike the other in-tree fused single-latent backends (``_gluon_v2`` /
+``_triton_v2`` / ``_flydsl_v1``), which vendor their kernels inside Primus, this
+backend calls straight into the installed **``primus_turbo``** package — the
+native-FlyDSL sparse-MLA attention under ``primus_turbo.flydsl.attention``. This
+is the "turbo API" integration: Primus owns only the thin V4 adapter binding; the
+kernels live in Primus-Turbo.
+
+Public kernel-pair API (identical to the other sparse-MLA backends):
+
+* ``sparse_mla_fwd_v4_turbo_flydsl(q, kv, topk, attn_sink=None,
+  kv_lora_rank=512, scale=None) -> (o, lse)``
+* ``sparse_mla_bwd_v4_turbo_flydsl(q, kv, o, do, topk, lse, attn_sink=None,
+  kv_lora_rank=512, scale=None) -> (dq, dkv, d_sink)``
+
+The upstream module layout / function names changed between Primus-Turbo
+releases, so we bind to whichever is present (the call signatures are identical):
+
+* **new** (>= the ``flydsl_attn`` refactor, e.g. commit ``edc8d2c``):
+  ``primus_turbo.flydsl.attention.sparse_mla_fwd.sparse_mla_fwd_flydsl`` /
+  ``...sparse_mla_bwd.sparse_mla_bwd_flydsl`` (flat modules, no ``_v4`` infix).
+* **old** (<= ``0.3.2.dev*``):
+  ``primus_turbo.flydsl.attention.kernels.sparse_mla_v2.{sparse_mla_fwd_v4_flydsl,
+  sparse_mla_bwd_v4_flydsl}``.
+
+``primus_turbo`` (with the flydsl sparse-MLA attention) and the ``flydsl`` pip
+package are required; the import fails with a clear message otherwise (handled by
+the lazy loader in :mod:`..` / :func:`load_turbo_attention_backends`).
+"""
+
+from __future__ import annotations
+
+try:
+    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
+    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
+        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
+    )
+    from primus_turbo.flydsl.attention.sparse_mla_fwd import (
+        sparse_mla_fwd_flydsl as _sparse_mla_fwd,
+    )
+except ImportError:
+    # Older Primus-Turbo layout: kernels.sparse_mla_v2 with the _v4 infix.
+    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
+        sparse_mla_bwd_v4_flydsl as _sparse_mla_bwd,
+    )
+    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
+        sparse_mla_fwd_v4_flydsl as _sparse_mla_fwd,
+    )
+
+# Turbo-suffixed re-exports so the V4 wrapper / benchmark can bind them without
+# clashing with the in-tree ``_flydsl_v1`` (which has identically-named kernels).
+sparse_mla_fwd_v4_turbo_flydsl = _sparse_mla_fwd
+sparse_mla_bwd_v4_turbo_flydsl = _sparse_mla_bwd
+
+__all__ = ["sparse_mla_fwd_v4_turbo_flydsl", "sparse_mla_bwd_v4_turbo_flydsl"]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py
new file mode 100644
index 000000000..968992336
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py
@@ -0,0 +1,38 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the native FlyDSL sparse-MLA backend ("flydsl_v1").
+
+Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`)
+to the ``sparse_mla_{fwd,bwd}_v4_flydsl`` kernel pair (``_flydsl_v1``): the fused
+single-latent (K == V) sparse-MLA path implemented in native FlyDSL MFMA
+(``rocdl.mfma_*``) over a per-token top-k gather. The forward is fully native
+FlyDSL; the backward uses a native FlyDSL dQ kernel plus the shared Triton dKV
+intermediate/scatter-gather. Numerically equivalent to the eager V4 references.
+
+Depends only on the installed ``flydsl`` pip package (gfx950 / CDNA4); it is
+therefore loaded LAZILY (see ``load_flydsl_attention_backends``), never at
+package import time.
+"""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v1 import (
+    sparse_mla_bwd_v4_flydsl,
+    sparse_mla_fwd_v4_flydsl,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_flydsl = make_csa_from_pool(sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl)
+v4_attention_flydsl = make_attention(sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl)
+
+__all__ = [
+    "v4_csa_attention_flydsl",
+    "v4_attention_flydsl",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py
new file mode 100644
index 000000000..0f4bfb305
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py
@@ -0,0 +1,33 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the Gluon sparse-MLA backend (gfx950).
+
+Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`)
+to the gluon ``sparse_mla_{fwd,bwd}_v4_gluon`` kernel pair (``_gluon_dsa``).
+See the adapter module for the full V4 <-> sparse-MLA mapping (zero rope pad,
+``[local ++ pool]`` kv buffer, ``[SWA window ++ pool]`` topk, grad mapping).
+Numerically equivalent to :func:`eager_v4_csa_attention` / :func:`eager_v4_attention`.
+"""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_dsa import (
+    sparse_mla_bwd_v4_gluon,
+    sparse_mla_fwd_v4_gluon,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_gluon = make_csa_from_pool(sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon)
+v4_attention_gluon = make_attention(sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon)
+
+__all__ = [
+    "v4_csa_attention_gluon",
+    "v4_attention_gluon",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py
new file mode 100644
index 000000000..53c3a4d8b
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py
@@ -0,0 +1,38 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the Gluon sparse-MLA backend ("gluon_v2").
+
+``gluon_v2`` forward is the Gluon sparse-MLA kernel (gfx950 / CDNA4 hardware-controlled
+layouts + async double-buffered pipeline, rope-skip + exp2 + MFMA K=32); the backward is
+currently the plain-Triton chunked-gather kernel (shared with ``triton_v2``) and is being
+migrated to Gluon. Thin binding of the kernel-agnostic V4 adapters
+(:mod:`v4_sparse_mla_adapter`) to the ``sparse_mla_{fwd,bwd}_v4_gluon_v2`` kernel pair
+(``_gluon_v2``).
+
+The Gluon forward requires a Gluon-capable (recompiled) triton; the kernel raises a clear
+build hint otherwise. Numerically equivalent to the eager V4 references (validated in
+tests/.../test_v4_gluon_v2_attention.py).
+"""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2 import (
+    sparse_mla_bwd_v4_gluon_v2,
+    sparse_mla_fwd_v4_gluon_v2,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_gluon_v2 = make_csa_from_pool(sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2)
+v4_attention_gluon_v2 = make_attention(sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2)
+
+__all__ = [
+    "v4_csa_attention_gluon_v2",
+    "v4_attention_gluon_v2",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py
new file mode 100644
index 000000000..6fe1f022b
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py
@@ -0,0 +1,26 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the Gluon sparse-MLA backend ("gluon_v3")."""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import (
+    sparse_mla_bwd_v4_gluon_v3,
+    sparse_mla_fwd_v4_gluon_v3,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_gluon_v3 = make_csa_from_pool(sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3)
+v4_attention_gluon_v3 = make_attention(sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3)
+
+__all__ = [
+    "v4_csa_attention_gluon_v3",
+    "v4_attention_gluon_v3",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py
new file mode 100644
index 000000000..98f14e186
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py
@@ -0,0 +1,37 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the plain-Triton sparse-MLA backend ("triton_v2").
+
+The ``_v2`` denotes the second *kernel* implementation of the V4 sparse-MLA
+attention (NOT a new Triton release). Thin binding of the kernel-agnostic V4
+adapters (:mod:`v4_sparse_mla_adapter`) to the ``sparse_mla_{fwd,bwd}_v4_triton``
+kernel pair (``_triton_v2``) — the fused single-latent (K == V) sparse-MLA path
+whose GEMMs lower to MFMA via ``tl.dot``. Unlike the gluon backend (gfx950 /
+CDNA4 hardware-controlled layouts), this is vanilla Triton and runs on any
+MFMA-capable arch (gfx942 / gfx950). Distinct from the in-tree separate-KV CSA
+Triton backend (``v4_csa_attention_v0``). Numerically equivalent to the eager V4
+references.
+"""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v2 import (
+    sparse_mla_bwd_v4_triton,
+    sparse_mla_fwd_v4_triton,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_v2 = make_csa_from_pool(sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton)
+v4_attention_v2 = make_attention(sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton)
+
+__all__ = [
+    "v4_csa_attention_v2",
+    "v4_attention_v2",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py
new file mode 100644
index 000000000..9da7c4ea9
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py
@@ -0,0 +1,39 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4 attention via the Primus-Turbo native-FlyDSL sparse-MLA backend ("turbo").
+
+Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`) to
+the ``sparse_mla_{fwd,bwd}_v4_turbo_flydsl`` kernel pair (:mod:`_turbo_flydsl`),
+which re-exports the installed ``primus_turbo`` flydsl sparse-MLA v2 kernels. Same
+fused single-latent (K == V) sparse-MLA-with-sink math as the ``gluon_v2`` /
+``triton_v2`` / ``flydsl_v1`` backends, so it is numerically equivalent to the
+eager V4 references (validated in
+``tests/.../test_v4_turbo_flydsl_attention.py``).
+
+Requires the installed ``primus_turbo`` (with the flydsl sparse-MLA attention) and
+the ``flydsl`` pip package (gfx950 / CDNA4); the import raises a clear hint
+otherwise (see :func:`..load_turbo_attention_backends`).
+"""
+
+from __future__ import annotations
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels._turbo_flydsl import (
+    sparse_mla_bwd_v4_turbo_flydsl,
+    sparse_mla_fwd_v4_turbo_flydsl,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import (
+    make_attention,
+    make_csa_from_pool,
+)
+
+v4_csa_attention_turbo = make_csa_from_pool(sparse_mla_fwd_v4_turbo_flydsl, sparse_mla_bwd_v4_turbo_flydsl)
+v4_attention_turbo = make_attention(sparse_mla_fwd_v4_turbo_flydsl, sparse_mla_bwd_v4_turbo_flydsl)
+
+__all__ = [
+    "v4_csa_attention_turbo",
+    "v4_attention_turbo",
+]
diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py
new file mode 100644
index 000000000..f31815887
--- /dev/null
+++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py
@@ -0,0 +1,295 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Kernel-agnostic V4 attention adapters over a sparse-MLA fwd/bwd kernel pair.
+
+Every "fused single-latent" V4 backend (gluon, triton-v2, flydsl-v2) speaks the
+same sparse-MLA contract:
+
+* ``fwd(q[T,H,Dqk], kv[T,1,Dqk], topk[T,TOPK], attn_sink, kv_lora_rank, scale)``
+  -> ``(o[T,H,Dv], lse[T,H])``
+* ``bwd(q, kv, o, do, topk, lse, attn_sink, kv_lora_rank, scale)``
+  -> ``(dq, dkv, d_sink)``
+
+This module maps Primus's V4 attention representations (per-head q, single MQA
+latent K = V with RoPE baked in-place over ``head_dim = 512``, compressed pool,
+per-query top-K, joint local-SWA + sparse softmax with sink) onto that contract
+and maps gradients back — once — so each backend is just a kernel pair:
+
+* :func:`make_csa_from_pool(fwd, bwd)`  -> CSA (cr=4) wrapper
+* :func:`make_attention(fwd, bwd)`      -> dense (cr=0) / HCA (cr=128) wrapper
+
+The fwd/bwd kernels are passed to the autograd Function as non-tensor args, so
+the same Function serves all backends (backward returns ``None`` for them).
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+import torch
+
+_ROPE_PAD = 64  # dummy separate-rope block (zeros); the kernels need D_ROPE > 0
+
+
+def _pad_topk_64(topk: torch.Tensor) -> torch.Tensor:
+    """Pad the topk width to a multiple of 64 with -1 so a backend whose dKV
+    tiling is 64-wide (e.g. gluon) stays valid (HCA 128+32=160 -> 192)."""
+    tk = topk.shape[1]
+    pad = ((tk + 63) // 64) * 64 - tk
+    if pad > 0:
+        topk = torch.cat(
+            [topk, torch.full((topk.shape[0], pad), -1, device=topk.device, dtype=topk.dtype)], dim=1
+        )
+    return topk.contiguous()
+
+
+def _build_csa_topk(topk_idxs: torch.Tensor, S: int, P: int, W: int) -> torch.Tensor:
+    """Flat topk [B*S, W+K] over the per-batch [local ++ pool] buffer.
+
+    ``topk_idxs`` [B, S, K] holds pool indices in [0, P) (or -1). Batch ``b``
+    occupies rows ``[b*(S+P) : (b+1)*(S+P))`` (local 0..S-1, pool S..S+P-1).
+    """
+    B, _, K = topk_idxs.shape
+    device = topk_idxs.device
+    base = (torch.arange(B, device=device) * (S + P)).view(B, 1, 1)
+
+    win_pos = torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W)
+    win_valid = win_pos >= 0
+    win_idx = base + win_pos.view(1, S, W)
+    win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1))
+
+    pool_valid = topk_idxs >= 0
+    pool_idx = torch.where(pool_valid, base + S + topk_idxs, torch.full_like(topk_idxs, -1))
+
+    return torch.cat([win_idx, pool_idx], dim=2).reshape(B * S, W + K).to(torch.int32).contiguous()
+
+
+class _V4SparseMLACSAFn(torch.autograd.Function):
+    """Autograd wrapper: sparse-MLA FWD/BWD for the V4 CSA (cr=4) layer."""
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q_bh: torch.Tensor,  # [B, H, S, D]
+        k_local_bh: torch.Tensor,  # [B, H, S, D] (single MQA latent, head-broadcast)
+        v_local_bh: torch.Tensor,  # [B, H, S, D] (== k_local in V4)
+        pool: torch.Tensor,  # [B, P, D]
+        topk_idxs: torch.Tensor,  # [B, S, K] pool indices, -1 = invalid
+        sink: Optional[torch.Tensor],  # [H] fp32 or None
+        swa_window: int,
+        scale: float,
+        fwd_fn: Callable,
+        bwd_fn: Callable,
+    ) -> torch.Tensor:
+        B, H, S, D = q_bh.shape
+        P = pool.shape[1]
+        W = int(swa_window)
+        assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16"
+        assert W > 0, "sparse-MLA adapter requires swa_window > 0"
+
+        latent = k_local_bh[:, 0, :, :]  # [B, S, D]
+
+        z_q = torch.zeros(B * S, H, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype)
+        q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous()
+
+        kv512 = torch.cat([latent, pool], dim=1).reshape(B * (S + P), 1, D)
+        z_kv = torch.zeros(B * (S + P), 1, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype)
+        kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous()
+
+        topk_g = _pad_topk_64(_build_csa_topk(topk_idxs, S, P, W))
+
+        sink_arg = sink.float().contiguous() if sink is not None else None
+        o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale))
+
+        ctx.save_for_backward(q_g, kv_g, o_g, lse, topk_g, sink_arg if sink is not None else q_g.new_empty(0))
+        ctx.shapes = (B, H, S, D, P, W)
+        ctx.scale = float(scale)
+        ctx.sink_was_none = sink is None
+        ctx.bwd_fn = bwd_fn
+        return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_o_bh: torch.Tensor):  # type: ignore[override]
+        q_g, kv_g, o_g, lse, topk_g, sink_saved = ctx.saved_tensors
+        B, H, S, D, P, W = ctx.shapes
+        sink_arg = None if ctx.sink_was_none else sink_saved
+
+        grad_o_g = grad_o_bh.permute(0, 2, 1, 3).reshape(B * S, H, D).contiguous()
+        dq_g, dkv_g, dsink = ctx.bwd_fn(
+            q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale
+        )
+
+        dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous()
+        dkv512 = dkv_g[:, 0, :D].reshape(B, S + P, D)
+        dlatent = dkv512[:, :S, :]
+        dpool = dkv512[:, S:, :].contiguous()
+
+        dk_local = torch.zeros(B, H, S, D, device=dq_bh.device, dtype=dq_bh.dtype)
+        dk_local[:, 0, :, :] = dlatent.to(dq_bh.dtype)
+        # V4 is single-latent (K = V = kv): the kernel returns one combined
+        # ``dkv`` which we route entirely through ``dk_local``. The V branch
+        # gradient is structurally zero, so we return ``None`` for it instead
+        # of allocating (and zeroing) a full [B, H, S, D] tensor — this removes
+        # the largest ``Memset (Device)`` bucket in the trace (~268 MB / call).
+        # ``k_local_bh`` and ``v_local_bh`` are two ``kv.expand`` views of the
+        # same latent, so autograd accumulates ``dk_local + 0`` into ``kv`` —
+        # identical to before.
+        dv_local = None
+
+        dsink_out = None
+        if not ctx.sink_was_none and dsink is not None:
+            dsink_out = dsink.to(sink_saved.dtype)
+
+        # forward args: (q, k_local, v_local, pool, topk_idxs, sink, swa_window, scale, fwd_fn, bwd_fn)
+        return dq_bh, dk_local, dv_local, dpool.to(dq_bh.dtype), None, dsink_out, None, None, None, None
+
+
+class _V4SparseMLAAttnFn(torch.autograd.Function):
+    """Sparse-MLA FWD/BWD for the V4 dense (cr=0) and HCA (cr=128) layers."""
+
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx,
+        q_bh: torch.Tensor,  # [B, H, S, D]
+        k_bh: torch.Tensor,  # [B, H, Skv, D] (Skv = S for cr=0; S+P for HCA)
+        v_bh: torch.Tensor,  # [B, H, Skv, D] (== k_bh in V4)
+        sink: Optional[torch.Tensor],
+        swa_window: int,
+        additive_mask: Optional[torch.Tensor],  # [S, P] pool-only mask (HCA) or None
+        scale: float,
+        hca_local_seqlen: int,
+        fwd_fn: Callable,
+        bwd_fn: Callable,
+    ) -> torch.Tensor:
+        B, H, S, D = q_bh.shape
+        Skv = k_bh.shape[2]
+        W = int(swa_window)
+        assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16"
+        assert W > 0, "sparse-MLA adapter requires swa_window > 0"
+
+        device = q_bh.device
+        base = (torch.arange(B, device=device) * Skv).view(B, 1, 1)
+        win_pos = (
+            torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W)
+        )
+        win_valid = win_pos >= 0
+        win_idx = base + win_pos.view(1, S, W)
+        win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1))
+
+        if hca_local_seqlen > 0 and additive_mask is not None:
+            P = Skv - int(hca_local_seqlen)
+            vis = (additive_mask == 0).view(1, S, P)
+            ps = torch.arange(P, device=device).view(1, 1, P)
+            pool_idx = torch.where(
+                vis, base + hca_local_seqlen + ps, torch.full((B, S, P), -1, device=device)
+            )
+            topk = torch.cat([win_idx, pool_idx], dim=2)
+        else:
+            topk = win_idx
+        topk_g = _pad_topk_64(topk.reshape(B * S, -1).to(torch.int32))
+
+        z_q = torch.zeros(B * S, H, _ROPE_PAD, device=device, dtype=q_bh.dtype)
+        q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous()
+        kv512 = k_bh[:, 0, :, :].reshape(B * Skv, 1, D)
+        z_kv = torch.zeros(B * Skv, 1, _ROPE_PAD, device=device, dtype=q_bh.dtype)
+        kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous()
+
+        sink_arg = sink.float().contiguous() if sink is not None else None
+        o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale))
+
+        ctx.save_for_backward(q_g, kv_g, o_g, lse, topk_g, sink_arg if sink is not None else q_g.new_empty(0))
+        ctx.shapes = (B, H, S, D, Skv)
+        ctx.scale = float(scale)
+        ctx.sink_was_none = sink is None
+        ctx.bwd_fn = bwd_fn
+        return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_o_bh: torch.Tensor):  # type: ignore[override]
+        q_g, kv_g, o_g, lse, topk_g, sink_saved = ctx.saved_tensors
+        B, H, S, D, Skv = ctx.shapes
+        sink_arg = None if ctx.sink_was_none else sink_saved
+
+        grad_o_g = grad_o_bh.permute(0, 2, 1, 3).reshape(B * S, H, D).contiguous()
+        dq_g, dkv_g, dsink = ctx.bwd_fn(
+            q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale
+        )
+
+        dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous()
+        dkv = dkv_g[:, 0, :D].reshape(B, Skv, D)
+        dk_bh = torch.zeros(B, H, Skv, D, device=dq_bh.device, dtype=dq_bh.dtype)
+        dk_bh[:, 0, :, :] = dkv.to(dq_bh.dtype)
+        # Single-latent (K = V): route the combined ``dkv`` through ``dk_bh``;
+        # the V branch gradient is structurally zero, so return ``None`` and
+        # skip the big [B, H, Skv, D] memset (see the CSA branch note).
+        dv_bh = None
+
+        dsink_out = None
+        if not ctx.sink_was_none and dsink is not None:
+            dsink_out = dsink.to(sink_saved.dtype)
+
+        # forward args: (q, k, v, sink, swa_window, additive_mask, scale, hca_local_seqlen, fwd_fn, bwd_fn)
+        return dq_bh, dk_bh, dv_bh, dsink_out, None, None, None, None, None, None
+
+
+def make_csa_from_pool(fwd_fn: Callable, bwd_fn: Callable) -> Callable:
+    """Build a ``v4_csa_attention_v1``-style wrapper for a kernel pair."""
+
+    def _csa_from_pool(
+        q_bh,
+        k_local_bh,
+        v_local_bh,
+        pool,
+        *,
+        topk_idxs,
+        sink,
+        swa_window,
+        attn_dropout,
+        training,
+        scale,
+    ):
+        if attn_dropout > 0.0 and training:
+            raise NotImplementedError(
+                "sparse-MLA CSA adapter does not implement in-kernel attention dropout "
+                f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}."
+            )
+        return _V4SparseMLACSAFn.apply(
+            q_bh, k_local_bh, v_local_bh, pool, topk_idxs, sink, int(swa_window), float(scale), fwd_fn, bwd_fn
+        )
+
+    return _csa_from_pool
+
+
+def make_attention(fwd_fn: Callable, bwd_fn: Callable) -> Callable:
+    """Build a dense (cr=0) / HCA (cr=128) attention wrapper for a kernel pair."""
+
+    def _attention(
+        q,
+        k,
+        v,
+        *,
+        sink,
+        swa_window,
+        additive_mask,
+        attn_dropout,
+        training,
+        scale,
+        hca_local_seqlen=0,
+    ):
+        if attn_dropout > 0.0 and training:
+            raise NotImplementedError(
+                "sparse-MLA attention adapter does not implement in-kernel attention dropout "
+                f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}."
+            )
+        return _V4SparseMLAAttnFn.apply(
+            q, k, v, sink, int(swa_window), additive_mask, float(scale), int(hca_local_seqlen), fwd_fn, bwd_fn
+        )
+
+    return _attention
+
+
+__all__ = ["make_csa_from_pool", "make_attention"]
diff --git a/primus/backends/megatron/core/utils.py b/primus/backends/megatron/core/utils.py
index a5855ad58..be2ca21af 100644
--- a/primus/backends/megatron/core/utils.py
+++ b/primus/backends/megatron/core/utils.py
@@ -1,26 +1,34 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 from functools import lru_cache
+from typing import Any, List
 
+import torch
 from megatron.core import parallel_state
-from primus_turbo.pytorch.ops.attention.attention_utils import (
-    All2AllAttentionSharder,
-    AttentionSharder,
-)
+
+from primus.core.utils.module_utils import log_rank_0
 
 
 @lru_cache
 def produce_attention_sharder(cp_comm_type: str):
+    # Import only when needed to avoid import errors if primus_turbo is not available
+    try:
+        from primus_turbo.pytorch.ops.attention.attention_utils import (
+            All2AllAttentionSharder,
+        )
+    except ImportError:
+        raise ImportError("All2AllAttentionSharder not available. Ensure primus_turbo is properly installed.")
+
     if cp_comm_type == "a2a":
         return All2AllAttentionSharder()
     else:
         raise ValueError(f"Unsupported cp_comm_type: {cp_comm_type}")
 
 
-def shard_batch_on_this_cp_rank(sharder: AttentionSharder, batch):
+def shard_batch_on_this_cp_rank(sharder, batch):
     cp_size = parallel_state.get_context_parallel_world_size()
     cp_group = parallel_state.get_context_parallel_group()
     if cp_size > 1:
@@ -29,3 +37,118 @@ def shard_batch_on_this_cp_rank(sharder: AttentionSharder, batch):
                 seq_dim = 1 if key != "attention_mask" else 2
                 batch[key] = sharder.shard_cp_input([val], cp_group, seq_dim)[0]
     return batch
+
+
+def apply_torch_compile_if_enabled(model: List[Any], args: Any) -> None:
+    """
+    Apply torch.compile to model AFTER distributed wrapping and config setup.
+
+    This must be called:
+    1. AFTER get_model() completes (which does FSDP/DDP wrapping internally)
+    2. AFTER ddp_config is set on the wrapped model
+    3. BEFORE optimizer creation
+
+    Args:
+        model: List of model modules (already wrapped in DDP/FSDP)
+        args: Megatron args object
+
+    Raises:
+        Exception: If compilation fails, exception is raised (not caught)
+    """
+    from megatron.training.utils import unwrap_model
+
+    # Check if compilation is enabled
+    if not getattr(args, "enable_torch_compile", False):
+        return
+
+    log_rank_0("=" * 80)
+    log_rank_0("Applying torch.compile AFTER distributed wrapping...")
+    log_rank_0("=" * 80)
+
+    # Import torch_FSDP for type checking
+    try:
+        from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import (
+            PrimusTorchFullyShardedDataParallel as torch_FSDP,
+        )
+    except ImportError:
+        torch_FSDP = None
+
+    for idx, model_module in enumerate(model):
+        # Check what type of wrapper we have
+        wrapper_type = type(model_module).__name__
+        log_rank_0(f"  Model [{idx}] wrapper: {wrapper_type}")
+
+        # Check if wrapped with FSDP and has ddp_config
+        if torch_FSDP and isinstance(model_module, torch_FSDP):
+            has_config = hasattr(model_module, "ddp_config")
+            log_rank_0(f"  FSDP detected - has ddp_config: {has_config}")
+
+            # Check if FSDP wrapper has compile_model method
+            if hasattr(model_module, "compile_model"):
+                log_rank_0(f"  Calling compile_model on FSDP wrapper...")
+                model_module.compile_model()
+                log_rank_0(f"    ✓ FSDP wrapper compilation complete")
+                continue
+
+        # Unwrap to get the actual model (Flux, GPT, etc.)
+        unwrapped_models = unwrap_model([model_module])
+
+        for unwrapped in unwrapped_models:
+            model_type = type(unwrapped).__name__
+
+            # Check if the model has a compile_model method
+            if hasattr(unwrapped, "compile_model"):
+                log_rank_0(f"  Compiling {model_type}...")
+
+                # Apply compilation
+                unwrapped.compile_model()
+
+                log_rank_0(f"    ✓ {model_type} compilation complete")
+            else:
+                log_rank_0(f"  ℹ  {model_type} does not support torch.compile " f"(no compile_model method)")
+
+    log_rank_0("torch.compile application complete")
+    log_rank_0("=" * 80)
+
+
+def apply_torch_compile_to_optimizer_if_enabled(optimizer: Any, args: Any) -> None:
+    """
+    Apply torch.compile to the optimizer's step() when enable_torch_compile is True.
+    Replaces optimizer.step with a compiled version so the trainer needs no changes.
+    Uses fullgraph=False for the optimizer step (recommended; allows timers/conditionals).
+    """
+    if optimizer is None:
+        return
+    if not getattr(args, "enable_torch_compile", False):
+        return
+    if not getattr(args, "torch_compile_optimizer", False):
+        log_rank_0("  Optimizer compilation disabled (torch_compile_optimizer=False)")
+        return
+
+    backend = getattr(args, "torch_compile_backend", "inductor")
+    mode = getattr(args, "torch_compile_mode", "default")
+    fullgraph = False  # Optimizer step: allow graph breaks (timers, conditionals)
+
+    compile_kwargs = {
+        "backend": backend,
+        "mode": mode,
+        "fullgraph": fullgraph,
+    }
+
+    scope = getattr(args, "torch_compile_optimizer_scope", "full")
+    log_rank_0(f"Applying torch.compile to optimizer step (scope={scope})...")
+    log_rank_0(f"  backend={backend}, mode={mode}, fullgraph={fullgraph}")
+
+    if scope == "inner_only":
+        inner_opt = getattr(optimizer, "optimizer", None)
+        if inner_opt is None:
+            log_rank_0("  ⚠ No inner optimizer found, falling back to full scope")
+            scope = "full"
+        else:
+            inner_opt.step = torch.compile(inner_opt.step, **compile_kwargs)
+            log_rank_0("  ✓ Inner optimizer step compiled (clip_grad_norm/DTensor ops excluded)")
+
+    if scope == "full":
+        original_step = optimizer.step
+        optimizer.step = torch.compile(original_step, **compile_kwargs)
+        log_rank_0("  ✓ Optimizer step compiled (full scope)")
diff --git a/primus/backends/megatron/data/__init__.py b/primus/backends/megatron/data/__init__.py
new file mode 100644
index 000000000..48c0e16f5
--- /dev/null
+++ b/primus/backends/megatron/data/__init__.py
@@ -0,0 +1,39 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Megatron data loading infrastructure.
+
+This module provides:
+    - DatasetProvider abstraction for pluggable data pipelines
+    - EnergonDatasetProvider for multimodal/diffusion (Energon)
+    - SyntheticDatasetProvider for mock/synthetic data
+    - MegatronDataloaderWrapper for generic dataloader compatibility
+    - Synthetic datasets for testing and development
+    - Diffusion-specific task encoders and preprocessing
+
+Architecture:
+    The strategy pattern allows different trainers to use different
+    data sources while sharing the same training infrastructure.
+"""
+
+_LAZY_IMPORTS = {
+    "DatasetProvider": ".dataset_provider",
+    "EnergonDatasetProvider": ".energon_dataset_provider",
+    "SyntheticDatasetProvider": ".synthetic_dataset_provider",
+    "MegatronDataloaderWrapper": ".dataloader",
+}
+
+
+def __getattr__(name):
+    if name in _LAZY_IMPORTS:
+        import importlib
+
+        module = importlib.import_module(_LAZY_IMPORTS[name], __name__)
+        value = getattr(module, name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = list(_LAZY_IMPORTS.keys())
diff --git a/primus/backends/megatron/data/cuda_prefetch.py b/primus/backends/megatron/data/cuda_prefetch.py
new file mode 100644
index 000000000..3f4c35626
--- /dev/null
+++ b/primus/backends/megatron/data/cuda_prefetch.py
@@ -0,0 +1,84 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+CUDA stream-based data prefetcher for overlapping HtoD transfers with compute.
+
+Wraps an iterator (typically RerunDataIterator -> MegatronDataloaderWrapper ->
+Energon SavableDataLoader) and transfers the next batch to GPU on a dedicated
+secondary stream.
+
+Each ``__next__`` call waits for the in-flight HtoD to complete, returns the
+GPU batch, and immediately kicks off the *next* CPU fetch + HtoD dispatch so
+the transfer overlaps with forward/backward compute.
+
+Requirements:
+    - Upstream DataLoader must use pin_memory=True for truly async non_blocking
+      transfers. Energon's SavableDataLoader satisfies this.
+    - Batches must be dict[str, Tensor | Any].
+"""
+
+import torch
+
+
+class CudaPrefetchIterator:
+    """Prefetch data batches to GPU on a secondary CUDA stream.
+
+    On construction, eagerly fetches the first batch and dispatches its HtoD
+    transfer (cold start).  Each subsequent ``__next__`` waits for the previous
+    HtoD, grabs the GPU batch, kicks off the *next* prefetch, and returns.
+
+    ``wait_stream()`` (not CUDA events) is used because prefetch depth is 1.
+    ``record_stream()`` is not needed because ``self._next_batch`` holds GPU
+    tensor references until consumed, and the caller holds the returned batch
+    through forward/backward.
+    """
+
+    def __init__(self, iterator, compute_dtype=torch.bfloat16):
+        if isinstance(iterator, (list, tuple)):
+            # Virtual pipeline parallel hands Megatron a list of per-chunk
+            # iterators; this single-stream prefetcher only wraps one iterator.
+            raise TypeError(
+                "CudaPrefetchIterator expects a single data iterator, got "
+                f"{type(iterator).__name__}. Skip prefetch for virtual pipeline "
+                "parallel or wrap each per-chunk iterator separately."
+            )
+        self._iterator = iterator
+        self._stream = torch.cuda.Stream()
+        self._dtype = compute_dtype
+        self._next_batch = None
+        self._prefetch()
+
+    def _prefetch(self):
+        """Fetch next batch from upstream and dispatch HtoD on secondary stream."""
+        try:
+            batch = next(self._iterator)
+        except StopIteration:
+            self._next_batch = None
+            return
+        with torch.cuda.stream(self._stream):
+            gpu_batch = {}
+            for k, v in batch.items():
+                if isinstance(v, torch.Tensor):
+                    if v.is_floating_point():
+                        gpu_batch[k] = v.to(dtype=self._dtype, device="cuda", non_blocking=True)
+                    else:
+                        gpu_batch[k] = v.cuda(non_blocking=True)
+                else:
+                    gpu_batch[k] = v
+            self._next_batch = gpu_batch
+
+    def __next__(self):
+        torch.cuda.current_stream().wait_stream(self._stream)
+        batch = self._next_batch
+        self._next_batch = None
+        if batch is None:
+            raise StopIteration
+        self._prefetch()
+        return batch
+
+    def __iter__(self):
+        return self
diff --git a/primus/backends/megatron/data/dataloader.py b/primus/backends/megatron/data/dataloader.py
new file mode 100644
index 000000000..4a53ceb13
--- /dev/null
+++ b/primus/backends/megatron/data/dataloader.py
@@ -0,0 +1,163 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from Megatron-LM multimodal examples
+# Reference: examples/multimodal/dataloader_provider.py
+
+"""
+Generic Megatron-compatible dataloader wrapper.
+
+This module provides utilities to adapt any iterable (PyTorch DataLoader,
+Megatron Energon loader, synthetic data, etc.) to work with Megatron's
+training loop.
+
+Key Features:
+    - Cyclic iteration (never raises StopIteration)
+    - Optional checkpoint support via duck typing
+    - Works with any Python iterable
+    - No dependencies on specific data sources
+
+Historical Note:
+    Originally named "EnergonDataloader" but has no Energon dependencies.
+    Renamed to MegatronDataloaderWrapper for clarity.
+"""
+
+import logging
+from typing import Any, Iterator
+
+logger = logging.getLogger(__name__)
+
+
+def cyclic_iter(iterator: Iterator) -> Iterator:
+    """
+    Create a cyclic iterator that restarts when exhausted.
+
+    Megatron's training loop expects infinite iterators that never
+    raise StopIteration. This wrapper cycles any iterator infinitely.
+
+    Args:
+        iterator: Any Python iterator
+
+    Yields:
+        Items from iterator, cycling infinitely
+
+    Example:
+        >>> loader = DataLoader(dataset, batch_size=4)
+        >>> infinite_loader = cyclic_iter(loader)
+        >>> # Never raises StopIteration
+
+    Reference:
+        Megatron-LM examples/multimodal/dataloader_provider.py:cyclic_iter
+    """
+    while True:
+        for item in iterator:
+            yield item
+
+
+class MegatronDataloaderWrapper:
+    """
+    Generic wrapper to make any iterable compatible with Megatron training loop.
+
+    This wrapper is completely generic and has NO dependencies on specific
+    data sources. It works with:
+        - PyTorch DataLoader (for synthetic/mock data)
+        - Megatron Energon loaders (for WebDataset)
+        - Any Python iterable
+
+    Features:
+        1. Cyclic iteration - never raises StopIteration
+        2. Optional state management - duck-typed for compatibility
+        3. Megatron training loop compatibility
+
+    The wrapper uses duck typing for checkpoint methods. If your dataloader
+    has save_state_rank() or restore_state_rank(), they'll be used.
+    Otherwise, state operations are no-ops (perfect for synthetic data).
+
+    Example with PyTorch DataLoader (synthetic data):
+        >>> from torch.utils.data import DataLoader
+        >>> loader = DataLoader(mock_dataset, batch_size=4)
+        >>> megatron_loader = MegatronDataloaderWrapper(loader)
+        >>> for batch in megatron_loader:
+        ...     # Never exhausts, cycles infinitely
+        ...     pass
+
+    Example with Energon (real data):
+        >>> from megatron.energon import get_loader
+        >>> energon_loader = get_loader(...)
+        >>> megatron_loader = MegatronDataloaderWrapper(energon_loader)
+        >>> # Checkpoint support via save_state_rank() is available
+
+    Reference:
+        Megatron-LM examples/multimodal/dataloader_provider.py:EnergonDataloader
+    """
+
+    def __init__(self, dataloader):
+        """
+        Initialize wrapper for any iterable.
+
+        Args:
+            dataloader: Any iterable (PyTorch DataLoader, Energon loader, etc.)
+        """
+        self._dataloader = dataloader
+        self._iter = iter(cyclic_iter(dataloader))
+        logger.debug(f"Initialized MegatronDataloaderWrapper for {type(dataloader).__name__}")
+
+    def __next__(self):
+        """Get next batch (never raises StopIteration)."""
+        return next(self._iter)
+
+    def __iter__(self):
+        """Return iterator."""
+        return self._iter
+
+    def save_state(self) -> Any:
+        """
+        Save dataloader state for checkpointing (if supported).
+
+        Uses duck typing - checks for save_state_rank() method.
+        Returns None if not supported (e.g., for synthetic/mock data).
+
+        Returns:
+            Dataloader state dictionary, or None if not supported
+        """
+        if hasattr(self._dataloader, "save_state_rank"):
+            return self._dataloader.save_state_rank()
+        else:
+            logger.debug(
+                f"{type(self._dataloader).__name__} does not support save_state_rank() "
+                f"(OK for synthetic data)"
+            )
+            return None
+
+    def restore_state(self, state: Any):
+        """
+        Restore dataloader state from checkpoint (if supported).
+
+        Uses duck typing - checks for restore_state_rank() method.
+        No-op if not supported (e.g., for synthetic/mock data).
+
+        Args:
+            state: Dataloader state dictionary
+        """
+        if hasattr(self._dataloader, "restore_state_rank"):
+            self._dataloader.restore_state_rank(state)
+            # Recreate iterator after restore
+            self._iter = iter(cyclic_iter(self._dataloader))
+            logger.info("Restored dataloader state from checkpoint")
+        else:
+            logger.debug(
+                f"{type(self._dataloader).__name__} does not support restore_state_rank() "
+                f"(OK for synthetic data)"
+            )
+
+
+# Backwards compatibility alias (DEPRECATED)
+# NOTE: Deprecated alias kept for backward compatibility.
+EnergonDataloader = MegatronDataloaderWrapper
+
+
+__all__ = [
+    "MegatronDataloaderWrapper",
+    "cyclic_iter",
+    "EnergonDataloader",  # Deprecated, use MegatronDataloaderWrapper
+]
diff --git a/primus/backends/megatron/data/dataset_provider.py b/primus/backends/megatron/data/dataset_provider.py
new file mode 100644
index 000000000..4955338b9
--- /dev/null
+++ b/primus/backends/megatron/data/dataset_provider.py
@@ -0,0 +1,63 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Abstract DatasetProvider for Megatron trainers.
+
+This module provides a strategy pattern for pluggable data pipelines,
+allowing different trainers to use different data sources while sharing
+the same training infrastructure.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional, Tuple
+
+
+class DatasetProvider(ABC):
+    """
+    Abstract strategy for dataset/dataloader creation.
+
+    Implementations provide different data pipelines:
+        - EnergonDatasetProvider: Megatron Energon for multimodal/diffusion
+        - SyntheticDatasetProvider: synthetic/mock batches for smoke tests
+
+    The provider pattern allows MegatronTrainer to remain agnostic about
+    the underlying data source while maintaining compatibility with
+    Megatron's build_train_valid_test_data_iterators().
+    """
+
+    @abstractmethod
+    def create_dataloaders(
+        self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None
+    ) -> Tuple[Any, Any, Any]:
+        """
+        Create train/valid/test dataloaders.
+
+        Args:
+            trainer_config: Megatron args namespace (from megatron.training.get_args())
+            train_val_test_num_samples: [train_samples, valid_samples, test_samples]
+            vp_stage: Virtual pipeline stage (for VP parallelism)
+
+        Returns:
+            Tuple of (train_dataloader, valid_dataloaders, test_dataloader)
+            - train_dataloader: Training data iterator
+            - valid_dataloaders: List of validation data iterators (or None)
+            - test_dataloader: Test data iterator (or None)
+        """
+
+    @property
+    @abstractmethod
+    def is_distributed(self) -> bool:
+        """
+        Whether dataloaders are distributed across ranks.
+
+        This flag tells Megatron whether to bypass indexed dataset logic:
+            False: GPTDataset (uses BlendedMegatronDatasetBuilder, rank-specific slicing)
+            True: Energon (handles distribution internally via WorkerConfig)
+
+        Returns:
+            bool: True if dataloaders handle distribution internally
+        """
+
+
+__all__ = ["DatasetProvider"]
diff --git a/primus/backends/megatron/data/diffusion/__init__.py b/primus/backends/megatron/data/diffusion/__init__.py
new file mode 100644
index 000000000..57c8e928f
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Data pipeline components for diffusion models.
+
+This module contains diffusion-specific data handling:
+    - Encoders (VAE, T5, CLIP, etc.)
+    - TaskEncoders for Energon integration
+    - Preprocessing utilities
+"""
diff --git a/primus/backends/megatron/data/diffusion/encoders/__init__.py b/primus/backends/megatron/data/diffusion/encoders/__init__.py
new file mode 100644
index 000000000..c379365e1
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/__init__.py
@@ -0,0 +1,162 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Encoder implementations for diffusion models.
+
+This module contains encoders for various modalities:
+    - Image encoders (VAE, etc.)
+    - Text encoders (T5, CLIP, etc.)
+    - Encoder registry for config-driven selection
+    - Base encoder classes
+
+Hierarchical structure:
+    - encoders/image/autoencoder_kl.py - VAE (AutoencoderKL)
+    - encoders/text/t5_xxl.py - T5-XXL text encoder
+    - encoders/text/clip_l.py - CLIP-L text encoder
+
+Each encoder type can have multiple implementations registered
+in the ENCODER_REGISTRY for flexible configuration.
+"""
+
+import importlib
+import logging
+from typing import Dict, Optional, Type
+
+from .base import BaseEncoder, BaseTextEncoder, BaseVAE
+from .config import (
+    CLIPLConfig,
+    EncoderConfig,
+    FluxEncoderConfig,
+    T5XXLConfig,
+    TextEncoderConfig,
+    VAEConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+# Global encoder registry
+ENCODER_REGISTRY: Dict[str, Type[BaseEncoder]] = {}
+
+
+def register_encoder(name: str, encoder_class: Type[BaseEncoder]):
+    """
+    Register an encoder class in the global registry.
+
+    Args:
+        name: Unique identifier for the encoder (e.g., 'autoencoder_kl', 't5_xxl')
+        encoder_class: Encoder class to register
+
+    Example:
+        >>> register_encoder('my_encoder', MyEncoder)
+    """
+    if name in ENCODER_REGISTRY:
+        logger.warning(f"Encoder '{name}' already registered. Overwriting with {encoder_class}")
+    ENCODER_REGISTRY[name] = encoder_class
+    logger.debug(f"Registered encoder: {name} -> {encoder_class.__name__}")
+
+
+def get_encoder(config: EncoderConfig) -> BaseEncoder:
+    """
+    Factory function to get an encoder instance from config.
+
+    Args:
+        config: Encoder configuration object with 'type' field
+
+    Returns:
+        Instantiated encoder instance
+
+    Raises:
+        ValueError: If encoder type not found in registry
+
+    Example:
+        >>> config = VAEConfig(type='autoencoder_kl', model_path='...')
+        >>> encoder = get_encoder(config)
+    """
+    _ensure_encoders_discovered()
+
+    encoder_type = config.type
+
+    if encoder_type not in ENCODER_REGISTRY:
+        raise ValueError(
+            f"Encoder type '{encoder_type}' not found in registry. "
+            f"Available encoders: {list(ENCODER_REGISTRY.keys())}"
+        )
+
+    encoder_class = ENCODER_REGISTRY[encoder_type]
+    logger.info(f"Loading encoder: {encoder_type} from {config.model_path}")
+
+    # Use from_pretrained if available, otherwise use constructor
+    if hasattr(encoder_class, "from_pretrained") and callable(encoder_class.from_pretrained):
+        return encoder_class.from_pretrained(config.model_path, config=config)
+    else:
+        return encoder_class(config)
+
+
+def list_encoders() -> Dict[str, Type[BaseEncoder]]:
+    """
+    List all registered encoders.
+
+    Returns:
+        Dictionary mapping encoder names to encoder classes
+    """
+    _ensure_encoders_discovered()
+    return ENCODER_REGISTRY.copy()
+
+
+_ENCODERS_DISCOVERED = False
+
+
+def _auto_discover_encoders():
+    """
+    Auto-discover and register all encoder implementations.
+
+    This function imports encoder modules which triggers their registration
+    through decorators or explicit register_encoder() calls.
+    """
+    encoder_modules = [
+        "primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl",
+        "primus.backends.megatron.data.diffusion.encoders.text.t5_xxl",
+        "primus.backends.megatron.data.diffusion.encoders.text.clip_l",
+    ]
+
+    for module_path in encoder_modules:
+        try:
+            importlib.import_module(module_path)
+            logger.debug(f"Successfully imported encoder module: {module_path}")
+        except ImportError as e:
+            logger.debug(f"Could not import encoder module {module_path}: {e}")
+        except Exception as e:
+            logger.warning(f"Error importing encoder module {module_path}: {e}")
+
+
+def _ensure_encoders_discovered():
+    """Run encoder discovery once, lazily (on first get_encoder/list_encoders).
+
+    Avoids importing all encoder backends (and their heavy deps) at package
+    import time.
+    """
+    global _ENCODERS_DISCOVERED
+    if not _ENCODERS_DISCOVERED:
+        _auto_discover_encoders()
+        _ENCODERS_DISCOVERED = True
+
+
+__all__ = [
+    # Base classes
+    "BaseEncoder",
+    "BaseVAE",
+    "BaseTextEncoder",
+    # Config classes
+    "EncoderConfig",
+    "VAEConfig",
+    "TextEncoderConfig",
+    "T5XXLConfig",
+    "CLIPLConfig",
+    "FluxEncoderConfig",
+    # Registry functions
+    "register_encoder",
+    "get_encoder",
+    "list_encoders",
+    "ENCODER_REGISTRY",
+]
diff --git a/primus/backends/megatron/data/diffusion/encoders/base.py b/primus/backends/megatron/data/diffusion/encoders/base.py
new file mode 100644
index 000000000..7f676f323
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/base.py
@@ -0,0 +1,285 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Base encoder classes for diffusion models.
+
+This module defines abstract base classes for different encoder types.
+All concrete encoder implementations should inherit from these base classes.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+
+def get_torch_dtype(precision: str) -> torch.dtype:
+    """
+    Convert precision string to torch dtype.
+
+    Args:
+        precision: One of 'bf16', 'fp16', 'fp32'
+
+    Returns:
+        Corresponding torch dtype
+
+    Example:
+        >>> dtype = get_torch_dtype('bf16')
+        >>> dtype
+        torch.bfloat16
+    """
+    DTYPE_MAP = {
+        "bf16": torch.bfloat16,
+        "fp16": torch.float16,
+        "fp32": torch.float32,
+    }
+    return DTYPE_MAP.get(precision, torch.bfloat16)
+
+
+def load_pretrained_with_subfolder_fallback(
+    loader_class: Any, model_path: str, subfolder: Optional[str] = None, **kwargs
+) -> Any:
+    """
+    Load a pretrained model with automatic subfolder fallback.
+
+    Tries to load from subfolder first, falls back to root path if that fails.
+    This handles both HuggingFace repos (with subfolders) and local paths.
+
+    Args:
+        loader_class: Model class with from_pretrained method
+        model_path: Path to model (HuggingFace repo or local path)
+        subfolder: Optional subfolder within model_path
+        **kwargs: Additional arguments passed to from_pretrained
+                  (e.g., torch_dtype, cache_dir, token)
+
+    Returns:
+        Loaded model instance
+
+    Example:
+        >>> from transformers import T5EncoderModel
+        >>> model = load_pretrained_with_subfolder_fallback(
+        ...     T5EncoderModel,
+        ...     "black-forest-labs/FLUX.1-dev",
+        ...     subfolder="text_encoder_2",
+        ...     torch_dtype=torch.bfloat16,
+        ...     cache_dir="/custom/cache/path"
+        ... )
+    """
+    # Add token from environment if available and not already in kwargs
+    import os
+
+    if "token" not in kwargs and "HF_TOKEN" in os.environ:
+        kwargs["token"] = os.environ["HF_TOKEN"]
+
+    if subfolder:
+        try:
+            return loader_class.from_pretrained(model_path, subfolder=subfolder, **kwargs)
+        except Exception as e:
+            logger.debug(f"Failed loading from subfolder '{subfolder}': {e}")
+            logger.debug("Retrying without subfolder...")
+
+    return loader_class.from_pretrained(model_path, **kwargs)
+
+
+# ============================================================================
+# Base Encoder Classes
+# ============================================================================
+
+
+class BaseEncoder(ABC, nn.Module):
+    """Abstract base class for all encoders."""
+
+    def __init__(self, config):
+        """
+        Initialize encoder.
+
+        Args:
+            config: Encoder configuration object
+        """
+        super().__init__()
+        self.config = config
+        self._device = config.device if hasattr(config, "device") else "cuda"
+        self._dtype = get_torch_dtype(config.precision if hasattr(config, "precision") else "bf16")
+
+    @property
+    def device(self) -> torch.device:
+        """Get encoder device."""
+        if isinstance(self._device, str):
+            return torch.device(self._device)
+        return self._device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Get encoder dtype."""
+        return self._dtype
+
+    def _get_dtype(self, precision: str) -> torch.dtype:
+        """
+        Convert precision string to torch dtype.
+
+        Args:
+            precision: One of 'bf16', 'fp16', 'fp32'
+
+        Returns:
+            Corresponding torch dtype
+        """
+        return get_torch_dtype(precision)
+
+    @abstractmethod
+    def encode(self, *args, **kwargs):
+        """
+        Encode input data.
+
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement encode()")
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(cls, model_path: str, config=None):
+        """
+        Load encoder from pretrained weights.
+
+        Args:
+            model_path: Path to pretrained model
+            config: Optional encoder configuration
+
+        Returns:
+            Loaded encoder instance
+        """
+        raise NotImplementedError("Subclasses must implement from_pretrained()")
+
+    def freeze(self):
+        """Freeze all encoder parameters."""
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def unfreeze(self):
+        """Unfreeze all encoder parameters."""
+        for param in self.parameters():
+            param.requires_grad = True
+
+
+class BaseVAE(BaseEncoder):
+    """Abstract base class for VAE encoders."""
+
+    def __init__(self, config):
+        """
+        Initialize VAE encoder.
+
+        Args:
+            config: VAEConfig object
+        """
+        super().__init__(config)
+        self.scale_factor = config.scale_factor if hasattr(config, "scale_factor") else 1.0
+        self.shift_factor = config.shift_factor if hasattr(config, "shift_factor") else 0.0
+        self.in_channels = config.in_channels if hasattr(config, "in_channels") else 3
+        self.out_channels = config.out_channels if hasattr(config, "out_channels") else 16
+        self.latent_downsample_factor = (
+            config.latent_downsample_factor if hasattr(config, "latent_downsample_factor") else 8
+        )
+
+    @abstractmethod
+    def encode(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Encode images to latent representations.
+
+        Args:
+            images: Input images tensor of shape (B, C, H, W)
+
+        Returns:
+            Latent representations of shape (B, latent_channels, H/downsample, W/downsample)
+        """
+        raise NotImplementedError("Subclasses must implement encode()")
+
+    @abstractmethod
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent representations to images.
+
+        Args:
+            latents: Latent representations of shape (B, latent_channels, H/downsample, W/downsample)
+
+        Returns:
+            Reconstructed images of shape (B, C, H, W)
+        """
+        raise NotImplementedError("Subclasses must implement decode()")
+
+    def get_latent_shape(self, image_height: int, image_width: int) -> Tuple[int, int, int]:
+        """
+        Calculate latent shape from image dimensions.
+
+        Args:
+            image_height: Input image height
+            image_width: Input image width
+
+        Returns:
+            Tuple of (channels, latent_height, latent_width)
+        """
+        latent_h = image_height // self.latent_downsample_factor
+        latent_w = image_width // self.latent_downsample_factor
+        return (self.out_channels, latent_h, latent_w)
+
+
+class BaseTextEncoder(BaseEncoder):
+    """Abstract base class for text encoders."""
+
+    def __init__(self, config):
+        """
+        Initialize text encoder.
+
+        Args:
+            config: TextEncoderConfig object
+        """
+        super().__init__(config)
+        self.max_length = config.max_length if hasattr(config, "max_length") else 512
+        self.embedding_dim = config.embedding_dim if hasattr(config, "embedding_dim") else 4096
+        self.return_pooled = config.return_pooled if hasattr(config, "return_pooled") else False
+        self._tokenizer = None
+
+    @abstractmethod
+    def encode(self, texts: Union[str, List[str]], **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+        """
+        Encode text(s) to embeddings.
+
+        Args:
+            texts: Single text string or list of text strings
+            **kwargs: Additional encoding arguments
+
+        Returns:
+            Text embeddings tensor(s)
+            - For standard encoders: (B, seq_len, embedding_dim)
+            - For pooled encoders: tuple of (sequence_embeddings, pooled_embeddings)
+        """
+        raise NotImplementedError("Subclasses must implement encode()")
+
+    @property
+    def tokenizer(self):
+        """Get tokenizer instance."""
+        if self._tokenizer is None:
+            raise ValueError("Tokenizer not initialized. Call from_pretrained() first.")
+        return self._tokenizer
+
+    def _prepare_texts(self, texts: Union[str, List[str]]) -> List[str]:
+        """
+        Prepare text inputs for encoding.
+
+        Args:
+            texts: Single text string or list of text strings
+
+        Returns:
+            List of text strings
+        """
+        if isinstance(texts, str):
+            return [texts]
+        return texts
diff --git a/primus/backends/megatron/data/diffusion/encoders/config.py b/primus/backends/megatron/data/diffusion/encoders/config.py
new file mode 100644
index 000000000..a3b3a8cdf
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/config.py
@@ -0,0 +1,160 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Configuration dataclasses for diffusion encoders.
+
+This module defines configuration classes for various encoder types used in diffusion models.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Optional
+
+
+class EncoderType(Enum):
+    """Categories of encoders for diffusion models."""
+
+    IMAGE = "image"
+    TEXT = "text"
+
+
+@dataclass
+class EncoderConfig:
+    """Base configuration for encoders."""
+
+    model_path: str  # Path to pretrained model weights (REQUIRED)
+    type: str = None  # Encoder type identifier (e.g., 'autoencoder_kl', 't5_xxl', 'clip_l')
+    encoder_type: Optional[EncoderType] = None  # Category: IMAGE or TEXT
+    precision: str = "bf16"  # Model precision: 'bf16', 'fp32', 'fp16'
+    device: str = "cuda"  # Device to load model on
+    use_cached: bool = True  # Whether to use cached/pre-encoded data
+    freeze_weights: bool = True  # Whether to freeze encoder weights during training
+    cache_dir: Optional[str] = None  # Directory to cache downloaded models (if None, uses HF default)
+    subfolder: Optional[str] = None  # Subfolder containing model weights (e.g., 'vae', 'text_encoder_2')
+    # Opt-in only: allow executing custom modeling code from the HF repo. Defaults to False
+    # so a malicious/unexpected repo cannot run arbitrary code unless explicitly enabled.
+    trust_remote_code: bool = False
+
+    # Optional parameters for specific encoders
+    extra_config: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        """Validate configuration."""
+        if self.type is None:
+            raise ValueError("'type' must be specified in encoder config")
+        valid_precisions = ["bf16", "fp32", "fp16"]
+        if self.precision not in valid_precisions:
+            raise ValueError(f"precision must be one of {valid_precisions}, got {self.precision}")
+
+
+@dataclass
+class VAEConfig(EncoderConfig):
+    """Configuration for VAE encoders."""
+
+    type: str = "autoencoder_kl"  # VAE type identifier
+    scale_factor: float = 0.3611  # Flux default scale factor
+    shift_factor: float = 0.1159  # Flux default shift factor
+    in_channels: int = 3  # Input image channels
+    out_channels: int = 16  # Latent channels
+    latent_downsample_factor: int = 8  # Spatial downsampling factor (H/8, W/8)
+
+    def __post_init__(self):
+        """Set encoder type and validate."""
+        if self.encoder_type is None:
+            self.encoder_type = EncoderType.IMAGE
+        super().__post_init__()
+
+
+@dataclass
+class TextEncoderConfig(EncoderConfig):
+    """Configuration for text encoders."""
+
+    max_length: int = 512  # Maximum sequence length
+    embedding_dim: int = 4096  # Output embedding dimension
+    return_pooled: bool = False  # Whether to return pooled embeddings
+    tokenizer_path: Optional[str] = None  # Optional separate tokenizer path
+    tokenizer_subfolder: Optional[str] = (
+        None  # Subfolder containing tokenizer files (e.g., 'tokenizer', 'tokenizer_2')
+    )
+
+    def __post_init__(self):
+        """Validate configuration."""
+        if not self.tokenizer_path:
+            self.tokenizer_path = self.model_path  # Default to same as model path
+        if self.encoder_type is None:
+            self.encoder_type = EncoderType.TEXT
+        super().__post_init__()
+
+
+@dataclass
+class T5XXLConfig(TextEncoderConfig):
+    """Configuration for T5-XXL encoder."""
+
+    type: str = "t5_xxl"  # T5-XXL type identifier
+    max_length: int = 512  # Flux uses 512 for T5-XXL
+    embedding_dim: int = 4096  # T5-XXL hidden size
+
+    def __post_init__(self):
+        """Validate configuration."""
+        super().__post_init__()
+
+
+@dataclass
+class CLIPLConfig(TextEncoderConfig):
+    """Configuration for CLIP-L encoder."""
+
+    type: str = "clip_l"  # CLIP-L type identifier
+    max_length: int = 77  # CLIP uses 77 max length
+    embedding_dim: int = 768  # CLIP-L hidden size
+    return_pooled: bool = True  # CLIP returns both sequence and pooled embeddings
+    pooled_dim: int = 768  # CLIP-L pooled embedding dimension
+
+    def __post_init__(self):
+        """Validate configuration."""
+        super().__post_init__()
+
+
+@dataclass
+class FluxEncoderConfig:
+    """Configuration for all Flux encoders (VAE + T5-XXL + CLIP-L)."""
+
+    vae: VAEConfig
+    t5: T5XXLConfig
+    clip: CLIPLConfig
+    use_preencoded: bool = True  # Whether to use pre-encoded data
+
+    @classmethod
+    def from_pretrained_flux(
+        cls,
+        model_path: str = "black-forest-labs/FLUX.1-dev",
+        precision: str = "bf16",
+        device: str = "cuda",
+        use_preencoded: bool = True,
+        cache_dir: Optional[str] = None,
+    ):
+        """Create FluxEncoderConfig from pretrained Flux model path."""
+        return cls(
+            vae=VAEConfig(
+                type="autoencoder_kl",
+                model_path=f"{model_path}",
+                precision=precision,
+                device=device,
+                cache_dir=cache_dir,
+            ),
+            t5=T5XXLConfig(
+                type="t5_xxl",
+                model_path=f"{model_path}",
+                precision=precision,
+                device=device,
+                cache_dir=cache_dir,
+            ),
+            clip=CLIPLConfig(
+                type="clip_l",
+                model_path=f"{model_path}",
+                precision=precision,
+                device=device,
+                cache_dir=cache_dir,
+            ),
+            use_preencoded=use_preencoded,
+        )
diff --git a/primus/backends/megatron/data/diffusion/encoders/image/__init__.py b/primus/backends/megatron/data/diffusion/encoders/image/__init__.py
new file mode 100644
index 000000000..c11e6c777
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/image/__init__.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Image encoder implementations for diffusion models.
+"""
+
+from primus.backends.megatron.data.diffusion.encoders.base import BaseVAE
+
+from .autoencoder_kl import AutoencoderKL
+
+__all__ = [
+    "BaseVAE",
+    "AutoencoderKL",
+]
diff --git a/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py b/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py
new file mode 100644
index 000000000..58401b331
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py
@@ -0,0 +1,271 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from NVIDIA NeMo Flux encoder wrappers (https://github.com/NVIDIA-NeMo/NeMo).
+
+"""
+AutoencoderKL implementation for Flux VAE.
+
+This module provides a wrapper around the diffusers AutoencoderKL model
+for use with Flux diffusion models. It handles encoding images to latent
+representations and decoding latents back to images.
+"""
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+
+try:
+    from diffusers import AutoencoderKL as DiffusersAutoencoderKL
+except ImportError:
+    DiffusersAutoencoderKL = None
+
+from primus.backends.megatron.data.diffusion.encoders.base import (
+    BaseVAE,
+    get_torch_dtype,
+    load_pretrained_with_subfolder_fallback,
+)
+from primus.backends.megatron.data.diffusion.encoders.config import VAEConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AutoencoderKL(BaseVAE):
+    """
+    AutoencoderKL VAE for Flux diffusion models.
+
+    This encoder wraps the diffusers AutoencoderKL model and applies
+    Flux-specific scale and shift factors to the latent representations.
+
+    For Flux, the default configuration is:
+        - scale_factor: 0.3611
+        - shift_factor: 0.1159
+        - in_channels: 3 (RGB images)
+        - out_channels: 16 (latent channels)
+        - latent_downsample_factor: 8 (spatial downsampling)
+
+    This means a 1024x1024 image becomes a 16x128x128 latent.
+    """
+
+    def __init__(self, config: VAEConfig):
+        """
+        Initialize AutoencoderKL encoder.
+
+        Args:
+            config: VAEConfig with model_path, scale_factor, shift_factor, etc.
+        """
+        super().__init__(config)
+
+        if DiffusersAutoencoderKL is None:
+            raise ImportError(
+                "diffusers library is required for AutoencoderKL. " "Install with: pip install diffusers"
+            )
+
+        self.vae = None  # Will be loaded in from_pretrained
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        config: Optional[VAEConfig] = None,
+        subfolder: Optional[str] = None,
+    ) -> "AutoencoderKL":
+        """
+        Load AutoencoderKL from pretrained weights.
+
+        Args:
+            model_path: Path to pretrained model (local path or HuggingFace repo)
+            config: Optional VAEConfig. If None, uses defaults.
+            subfolder: Subfolder for VAE weights. Priority: param > config.subfolder
+
+        Returns:
+            Loaded AutoencoderKL instance
+
+        Examples:
+            >>> # Using config (recommended)
+            >>> config = VAEConfig(
+            ...     model_path="black-forest-labs/FLUX.1-dev",
+            ...     subfolder="vae"
+            ... )
+            >>> vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", config=config)
+
+            >>> # Using method parameter (backward compatible)
+            >>> vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae")
+        """
+        if config is None:
+            config = VAEConfig(
+                type="autoencoder_kl",
+                model_path=model_path,
+                precision="bf16",
+            )
+
+        instance = cls(config)
+
+        # Prepare kwargs for from_pretrained calls
+        pretrained_kwargs = {}
+        if config.cache_dir:
+            pretrained_kwargs["cache_dir"] = config.cache_dir
+            logger.info(f"Using cache directory: {config.cache_dir}")
+
+        # Resolve model subfolder with priority: param > config.subfolder > error
+        model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None)
+
+        # Require explicit configuration for model subfolder
+        if model_subfolder is None and not hasattr(config, "subfolder"):
+            raise ValueError(
+                f"subfolder must be specified for AutoencoderKL with model_path='{model_path}'. "
+                f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='vae'. "
+                f"Set it via config.subfolder or the subfolder parameter."
+            )
+
+        # Load VAE from diffusers
+        torch_dtype = get_torch_dtype(config.precision)
+
+        logger.info(f"Loading AutoencoderKL from {model_path} (subfolder={model_subfolder})")
+        instance.vae = load_pretrained_with_subfolder_fallback(
+            DiffusersAutoencoderKL,
+            model_path,
+            subfolder=model_subfolder,
+            torch_dtype=torch_dtype,
+            **pretrained_kwargs,
+        )
+
+        instance.vae.to(instance.device)
+
+        if config.freeze_weights:
+            instance.freeze()
+            instance.vae.eval()
+
+        logger.info(
+            f"Loaded AutoencoderKL: in_channels={instance.in_channels}, "
+            f"out_channels={instance.out_channels}, scale={instance.scale_factor:.4f}, "
+            f"shift={instance.shift_factor:.4f}"
+        )
+
+        return instance
+
+    @torch.no_grad()
+    def encode(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Encode images to latent representations.
+
+        Applies the Flux scale and shift factors:
+            latents = scale_factor * (vae_encode(images) - shift_factor)
+
+        Args:
+            images: Input images tensor of shape (B, C, H, W)
+                    Values should be in range [-1, 1] (normalized)
+
+        Returns:
+            Latent representations of shape (B, 16, H/8, W/8)
+
+        Example:
+            >>> images = torch.randn(2, 3, 1024, 1024)  # B=2, 1024x1024 RGB
+            >>> latents = vae.encode(images)
+            >>> latents.shape
+            torch.Size([2, 16, 128, 128])
+        """
+        if self.vae is None:
+            raise RuntimeError("VAE not loaded. Call from_pretrained() first.")
+
+        # Ensure images are on correct device and dtype
+        images = images.to(device=self.device, dtype=self.dtype)
+
+        # Encode using diffusers VAE
+        latent_dist = self.vae.encode(images).latent_dist
+        latents = latent_dist.sample()
+
+        # Apply Flux scale and shift factors
+        # Formula: z = scale_factor * (encoded - shift_factor)
+        latents = self.scale_factor * (latents - self.shift_factor)
+
+        return latents
+
+    @torch.no_grad()
+    def encode_for_resample(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Encode images returning posterior parameters for training-time resampling.
+
+        Unlike ``encode()`` which draws a single stochastic sample, this method
+        returns the raw posterior parameters (mean, logvar) so the training loop
+        can re-draw latents via reparameterization at every step.
+
+        The returned ``latents`` use the deterministic posterior mode (mean)
+        with scale/shift applied, for backward compatibility and debugging.
+
+        Args:
+            images: Input images tensor (B, C, H, W) in range [-1, 1]
+
+        Returns:
+            Tuple of (latents, mean, logvar):
+                - latents: scale * (mode - shift), shape (B, 16, H/8, W/8)
+                - mean: raw posterior mean, shape (B, 16, H/8, W/8)
+                - logvar: raw posterior log-variance, shape (B, 16, H/8, W/8)
+        """
+        if self.vae is None:
+            raise RuntimeError("VAE not loaded. Call from_pretrained() first.")
+
+        images = images.to(device=self.device, dtype=self.dtype)
+
+        latent_dist = self.vae.encode(images).latent_dist
+        latents = self.scale_factor * (latent_dist.mode() - self.shift_factor)
+
+        return latents, latent_dist.mean, latent_dist.logvar
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent representations to images.
+
+        Reverses the Flux scale and shift factors before decoding:
+            vae_decode(latents / scale_factor + shift_factor)
+
+        Args:
+            latents: Latent representations of shape (B, 16, H/8, W/8)
+
+        Returns:
+            Reconstructed images of shape (B, C, H, W)
+            Values in range [-1, 1]
+
+        Example:
+            >>> latents = torch.randn(2, 16, 128, 128)
+            >>> images = vae.decode(latents)
+            >>> images.shape
+            torch.Size([2, 3, 1024, 1024])
+        """
+        if self.vae is None:
+            raise RuntimeError("VAE not loaded. Call from_pretrained() first.")
+
+        # Ensure latents are on correct device and dtype
+        latents = latents.to(device=self.device, dtype=self.dtype)
+
+        # Reverse Flux scale and shift factors
+        # Formula: decoded_input = z / scale_factor + shift_factor
+        latents = latents / self.scale_factor + self.shift_factor
+
+        # Decode using diffusers VAE
+        images = self.vae.decode(latents, return_dict=False)[0]
+
+        return images
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass: encode then decode (reconstruction).
+
+        Args:
+            images: Input images tensor of shape (B, C, H, W)
+
+        Returns:
+            Reconstructed images of shape (B, C, H, W)
+        """
+        latents = self.encode(images)
+        reconstructed = self.decode(latents)
+        return reconstructed
+
+
+# Register encoder in registry
+from primus.backends.megatron.data.diffusion.encoders import register_encoder
+
+register_encoder("autoencoder_kl", AutoencoderKL)
diff --git a/primus/backends/megatron/data/diffusion/encoders/text/__init__.py b/primus/backends/megatron/data/diffusion/encoders/text/__init__.py
new file mode 100644
index 000000000..ce83875a3
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/text/__init__.py
@@ -0,0 +1,14 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Text encoder implementations for diffusion models.
+"""
+
+from .clip_l import CLIPLEncoder
+from .t5_xxl import T5XXLEncoder
+
+__all__ = [
+    "T5XXLEncoder",
+    "CLIPLEncoder",
+]
diff --git a/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py b/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py
new file mode 100644
index 000000000..459ffbe18
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py
@@ -0,0 +1,255 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from NVIDIA NeMo Flux encoder wrappers (https://github.com/NVIDIA-NeMo/NeMo).
+
+"""
+CLIP-L text encoder implementation for Flux.
+
+This module provides a wrapper around the transformers CLIPTextModel
+for encoding text prompts into embeddings for Flux diffusion models.
+"""
+
+import logging
+from typing import List, Optional, Tuple, Union
+
+import torch
+
+try:
+    from transformers import CLIPTextModel, CLIPTokenizer
+except ImportError:
+    CLIPTextModel = None
+    CLIPTokenizer = None
+
+from primus.backends.megatron.data.diffusion.encoders.base import (
+    BaseTextEncoder,
+    get_torch_dtype,
+    load_pretrained_with_subfolder_fallback,
+)
+from primus.backends.megatron.data.diffusion.encoders.config import CLIPLConfig
+
+logger = logging.getLogger(__name__)
+
+
+class CLIPLEncoder(BaseTextEncoder):
+    """
+    CLIP-L text encoder for Flux diffusion models.
+
+    This encoder uses the CLIP ViT-L/14 text encoder from HuggingFace to encode
+    text prompts into embeddings. For Flux, the default configuration is:
+        - max_length: 77 tokens (CLIP default)
+        - embedding_dim: 768 (CLIP-L hidden size)
+        - pooled_dim: 768 (CLIP-L pooled embedding size)
+        - precision: bf16
+
+    The encoder returns both sequence embeddings (B, seq_len, 768) and
+    pooled embeddings (B, 768) which are used differently in Flux.
+
+    Note: Sequence length is padded to a multiple of 8 for Flux compatibility.
+    """
+
+    def __init__(self, config: CLIPLConfig):
+        """
+        Initialize CLIP-L encoder.
+
+        Args:
+            config: CLIPLConfig with model_path, max_length, etc.
+        """
+        super().__init__(config)
+        self.pooled_dim = config.pooled_dim if hasattr(config, "pooled_dim") else 768
+
+        if CLIPTextModel is None or CLIPTokenizer is None:
+            raise ImportError(
+                "transformers library is required for CLIPLEncoder. " "Install with: pip install transformers"
+            )
+
+        self.transformer = None  # Will be loaded in from_pretrained
+        self._tokenizer = None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        config: Optional[CLIPLConfig] = None,
+        subfolder: Optional[str] = None,
+    ) -> "CLIPLEncoder":
+        """
+        Load CLIP-L encoder from pretrained weights.
+
+        Args:
+            model_path: Path to pretrained model (local path or HuggingFace repo)
+            config: Optional CLIPLConfig. If None, uses defaults.
+            subfolder: Subfolder for model weights. Priority: param > config.subfolder
+
+        Returns:
+            Loaded CLIPLEncoder instance
+
+        Examples:
+            >>> # Using config (recommended)
+            >>> config = CLIPLConfig(
+            ...     model_path="black-forest-labs/FLUX.1-dev",
+            ...     subfolder="text_encoder",
+            ...     tokenizer_subfolder="tokenizer"
+            ... )
+            >>> encoder = CLIPLEncoder.from_pretrained("black-forest-labs/FLUX.1-dev", config=config)
+
+            >>> # Using method parameter (backward compatible)
+            >>> encoder = CLIPLEncoder.from_pretrained("openai/clip-vit-large-patch14", subfolder=None)
+        """
+        if config is None:
+            config = CLIPLConfig(
+                type="clip_l",
+                model_path=model_path,
+                precision="bf16",
+            )
+
+        instance = cls(config)
+
+        # Prepare kwargs for from_pretrained calls
+        pretrained_kwargs = {}
+        if config.cache_dir:
+            pretrained_kwargs["cache_dir"] = config.cache_dir
+            logger.info(f"Using cache directory: {config.cache_dir}")
+
+        # Resolve model subfolder with priority: param > config.subfolder > error
+        model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None)
+
+        # Require explicit configuration for model subfolder
+        # NOTE: config.subfolder can legitimately be None for standalone models like openai/clip-vit-large-patch14
+        # but it must be explicitly set in the config. If neither param nor config specify it, that's an error.
+        if model_subfolder is None and not hasattr(config, "subfolder"):
+            raise ValueError(
+                f"subfolder must be specified for CLIPLEncoder with model_path='{model_path}'. "
+                f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='text_encoder'. "
+                f"For standalone CLIP models (e.g., openai/clip-vit-large-patch14), use subfolder=None. "
+                f"Set it via config.subfolder or the subfolder parameter."
+            )
+
+        # Resolve tokenizer subfolder with priority: config.tokenizer_subfolder > model_subfolder
+        tokenizer_subfolder = getattr(config, "tokenizer_subfolder", None)
+        if tokenizer_subfolder is None:
+            tokenizer_subfolder = model_subfolder  # Use model subfolder as fallback
+
+        # Load tokenizer
+        tokenizer_path = config.tokenizer_path or model_path
+        logger.info(f"Loading CLIP tokenizer from {tokenizer_path} (subfolder={tokenizer_subfolder})")
+        try:
+            instance._tokenizer = load_pretrained_with_subfolder_fallback(
+                CLIPTokenizer,
+                tokenizer_path,
+                subfolder=tokenizer_subfolder,
+                **pretrained_kwargs,
+            )
+        except Exception as e:
+            # Try standard CLIP tokenizer if neither subfolder nor root work
+            logger.info(f"Failed loading from {tokenizer_path}, trying standard CLIP tokenizer: {e}")
+            instance._tokenizer = CLIPTokenizer.from_pretrained(
+                "openai/clip-vit-large-patch14",
+                **pretrained_kwargs,
+            )
+
+        # Load model
+        torch_dtype = get_torch_dtype(config.precision)
+
+        logger.info(f"Loading CLIP-L model from {model_path} (subfolder={model_subfolder})")
+        instance.transformer = load_pretrained_with_subfolder_fallback(
+            CLIPTextModel,
+            model_path,
+            subfolder=model_subfolder,
+            torch_dtype=torch_dtype,
+            **pretrained_kwargs,
+        )
+
+        instance.transformer.to(instance.device)
+
+        if config.freeze_weights:
+            instance.freeze()
+            instance.transformer.eval()
+
+        logger.info(
+            f"Loaded CLIP-L: max_length={instance.max_length}, "
+            f"embedding_dim={instance.embedding_dim}, pooled_dim={instance.pooled_dim}, "
+            f"dtype={torch_dtype}"
+        )
+
+        return instance
+
+    @torch.no_grad()
+    def encode(
+        self,
+        texts: Union[str, List[str]],
+        max_sequence_length: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Encode text(s) to embeddings.
+
+        Args:
+            texts: Single text string or list of text strings
+            max_sequence_length: Optional max length override. If None, uses config max_length.
+
+        Returns:
+            Tuple of (sequence_embeddings, pooled_embeddings):
+                - sequence_embeddings: (batch_size, padded_seq_len, 768)
+                  Sequence length is padded to a multiple of 8 for Flux
+                - pooled_embeddings: (batch_size, 768)
+                  Pooled representation from CLIP's [CLS] token
+
+        Example:
+            >>> texts = ["A beautiful sunset over the ocean", "A cat sitting on a mat"]
+            >>> seq_embeds, pooled_embeds = encoder.encode(texts)
+            >>> seq_embeds.shape
+            torch.Size([2, 80, 768])  # Padded to multiple of 8
+            >>> pooled_embeds.shape
+            torch.Size([2, 768])
+        """
+        if self.transformer is None or self._tokenizer is None:
+            raise RuntimeError("CLIP encoder not loaded. Call from_pretrained() first.")
+
+        texts = self._prepare_texts(texts)
+        max_len = max_sequence_length if max_sequence_length is not None else self.max_length
+
+        # Tokenize
+        batch_encoding = self._tokenizer(
+            texts,
+            truncation=True,
+            max_length=max_len,
+            return_length=True,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        tokens = batch_encoding["input_ids"].to(self.device, non_blocking=True)
+
+        # Encode
+        outputs = self.transformer(input_ids=tokens, output_hidden_states=False)
+
+        # Get sequence embeddings (last hidden state)
+        sequence_embeddings = outputs.last_hidden_state
+
+        # Pad sequence length to multiple of 8 (required for Flux)
+        seq_len = sequence_embeddings.shape[1]
+        padded_seq_len = ((seq_len + 7) // 8) * 8  # Round up to nearest multiple of 8
+
+        if padded_seq_len > seq_len:
+            sequence_embeddings = torch.nn.functional.pad(
+                sequence_embeddings,
+                (0, 0, 0, padded_seq_len - seq_len),  # Pad on sequence dimension
+                value=0.0,
+            )
+
+        # Get pooled embeddings (from pooler output)
+        pooled_embeddings = outputs.pooler_output
+
+        return sequence_embeddings, pooled_embeddings
+
+    def forward(self, texts: Union[str, List[str]], **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass (alias for encode)."""
+        return self.encode(texts, **kwargs)
+
+
+# Register encoder in registry
+from primus.backends.megatron.data.diffusion.encoders import register_encoder
+
+register_encoder("clip_l", CLIPLEncoder)
diff --git a/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py b/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py
new file mode 100644
index 000000000..8ce80b08e
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py
@@ -0,0 +1,230 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from NVIDIA NeMo Flux encoder wrappers (https://github.com/NVIDIA-NeMo/NeMo).
+
+"""
+T5-XXL text encoder implementation for Flux.
+
+This module provides a wrapper around the transformers T5EncoderModel
+for encoding text prompts into embeddings for Flux diffusion models.
+"""
+
+import logging
+from typing import List, Optional, Union
+
+import torch
+
+try:
+    from transformers import T5EncoderModel, T5Tokenizer
+except ImportError:
+    T5EncoderModel = None
+    T5Tokenizer = None
+
+from primus.backends.megatron.data.diffusion.encoders.base import (
+    BaseTextEncoder,
+    get_torch_dtype,
+    load_pretrained_with_subfolder_fallback,
+)
+from primus.backends.megatron.data.diffusion.encoders.config import T5XXLConfig
+
+logger = logging.getLogger(__name__)
+
+
+class T5XXLEncoder(BaseTextEncoder):
+    """
+    T5-XXL text encoder for Flux diffusion models.
+
+    This encoder uses the T5-v1.1-XXL model from HuggingFace to encode
+    text prompts into embeddings. For Flux, the default configuration is:
+        - max_length: 512 tokens
+        - embedding_dim: 4096 (T5-XXL hidden size)
+        - precision: bf16
+
+    The output embeddings have shape (batch_size, seq_len, 4096).
+    """
+
+    def __init__(self, config: T5XXLConfig):
+        """
+        Initialize T5-XXL encoder.
+
+        Args:
+            config: T5XXLConfig with model_path, max_length, etc.
+        """
+        super().__init__(config)
+
+        if T5EncoderModel is None or T5Tokenizer is None:
+            raise ImportError(
+                "transformers library is required for T5XXLEncoder. " "Install with: pip install transformers"
+            )
+
+        self.transformer = None  # Will be loaded in from_pretrained
+        self._tokenizer = None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        config: Optional[T5XXLConfig] = None,
+        subfolder: Optional[str] = None,
+    ) -> "T5XXLEncoder":
+        """
+        Load T5-XXL encoder from pretrained weights.
+
+        Args:
+            model_path: Path to pretrained model (local path or HuggingFace repo)
+            config: Optional T5XXLConfig. If None, uses defaults.
+            subfolder: Subfolder for model weights. Priority: param > config.subfolder
+
+        Returns:
+            Loaded T5XXLEncoder instance
+
+        Examples:
+            >>> # Using config (recommended)
+            >>> config = T5XXLConfig(
+            ...     model_path="black-forest-labs/FLUX.1-dev",
+            ...     subfolder="text_encoder_2",
+            ...     tokenizer_subfolder="tokenizer_2"
+            ... )
+            >>> encoder = T5XXLEncoder.from_pretrained("black-forest-labs/FLUX.1-dev", config=config)
+
+            >>> # Using method parameter (backward compatible)
+            >>> encoder = T5XXLEncoder.from_pretrained("google/t5-v1_1-xxl", subfolder=None)
+        """
+        if config is None:
+            config = T5XXLConfig(
+                type="t5_xxl",
+                model_path=model_path,
+                precision="bf16",
+            )
+
+        instance = cls(config)
+
+        # Prepare kwargs for from_pretrained calls
+        pretrained_kwargs = {}
+        if config.cache_dir:
+            pretrained_kwargs["cache_dir"] = config.cache_dir
+            logger.info(f"Using cache directory: {config.cache_dir}")
+
+        # Resolve model subfolder with priority: param > config.subfolder > error
+        model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None)
+
+        # Require explicit configuration for model subfolder
+        # NOTE: config.subfolder can legitimately be None for standalone models like google/t5-v1_1-xxl
+        # but it must be explicitly set in the config. If neither param nor config specify it, that's an error.
+        if model_subfolder is None and not hasattr(config, "subfolder"):
+            raise ValueError(
+                f"subfolder must be specified for T5XXLEncoder with model_path='{model_path}'. "
+                f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='text_encoder_2'. "
+                f"For standalone T5 models (e.g., google/t5-v1_1-xxl), use subfolder=None. "
+                f"Set it via config.subfolder or the subfolder parameter."
+            )
+
+        # Resolve tokenizer subfolder with priority: config.tokenizer_subfolder > model_subfolder
+        tokenizer_subfolder = getattr(config, "tokenizer_subfolder", None)
+        if tokenizer_subfolder is None:
+            tokenizer_subfolder = model_subfolder  # Use model subfolder as fallback
+
+        # Load tokenizer
+        tokenizer_path = config.tokenizer_path or model_path
+        logger.info(f"Loading T5 tokenizer from {tokenizer_path} (subfolder={tokenizer_subfolder})")
+
+        # Load tokenizer from subfolder
+        import os
+
+        tokenizer_kwargs = {
+            "token": os.environ.get("HF_TOKEN"),
+            "trust_remote_code": getattr(config, "trust_remote_code", False),
+            **pretrained_kwargs,  # Merge cache_dir if provided
+        }
+        if tokenizer_subfolder:
+            tokenizer_kwargs["subfolder"] = tokenizer_subfolder
+
+        instance._tokenizer = T5Tokenizer.from_pretrained(
+            tokenizer_path,
+            **tokenizer_kwargs,
+        )
+
+        # Load model
+        torch_dtype = get_torch_dtype(config.precision)
+
+        logger.info(f"Loading T5-XXL model from {model_path} (subfolder={model_subfolder})")
+        instance.transformer = load_pretrained_with_subfolder_fallback(
+            T5EncoderModel,
+            model_path,
+            subfolder=model_subfolder,
+            torch_dtype=torch_dtype,
+            **pretrained_kwargs,
+        )
+
+        instance.transformer.to(instance.device)
+
+        if config.freeze_weights:
+            instance.freeze()
+            instance.transformer.eval()
+
+        logger.info(
+            f"Loaded T5-XXL: max_length={instance.max_length}, "
+            f"embedding_dim={instance.embedding_dim}, dtype={torch_dtype}"
+        )
+
+        return instance
+
+    @torch.no_grad()
+    def encode(
+        self,
+        texts: Union[str, List[str]],
+        max_sequence_length: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Encode text(s) to embeddings.
+
+        Args:
+            texts: Single text string or list of text strings
+            max_sequence_length: Optional max length override. If None, uses config max_length.
+
+        Returns:
+            Text embeddings tensor of shape (batch_size, seq_len, 4096)
+            Sequence length is padded to max_length.
+
+        Example:
+            >>> texts = ["A beautiful sunset over the ocean", "A cat sitting on a mat"]
+            >>> embeddings = encoder.encode(texts)
+            >>> embeddings.shape
+            torch.Size([2, 512, 4096])
+        """
+        if self.transformer is None or self._tokenizer is None:
+            raise RuntimeError("T5 encoder not loaded. Call from_pretrained() first.")
+
+        texts = self._prepare_texts(texts)
+        max_len = max_sequence_length if max_sequence_length is not None else self.max_length
+
+        # Tokenize
+        batch_encoding = self._tokenizer(
+            texts,
+            truncation=True,
+            max_length=max_len,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        tokens = batch_encoding["input_ids"].to(self.device, non_blocking=True)
+
+        # Encode
+        outputs = self.transformer(input_ids=tokens, output_hidden_states=None)
+        embeddings = outputs.last_hidden_state
+
+        return embeddings
+
+    def forward(self, texts: Union[str, List[str]], **kwargs) -> torch.Tensor:
+        """Forward pass (alias for encode)."""
+        return self.encode(texts, **kwargs)
+
+
+# Register encoder in registry
+from primus.backends.megatron.data.diffusion.encoders import register_encoder
+
+register_encoder("t5_xxl", T5XXLEncoder)
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/__init__.py b/primus/backends/megatron/data/diffusion/preprocessing/__init__.py
new file mode 100644
index 000000000..21f7888d0
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Data preprocessing for diffusion models.
+
+Pipeline implementations live in the ``pipelines`` subpackage; shared helpers
+live in ``utils``. Import those submodules directly.
+"""
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/auth.py b/primus/backends/megatron/data/diffusion/preprocessing/auth.py
new file mode 100644
index 000000000..8cb574f08
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/auth.py
@@ -0,0 +1,174 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Authentication utilities for HuggingFace dataset access.
+
+Provides secure token loading from files with permission checks.
+"""
+
+import logging
+import os
+import stat
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class HFAuthError(Exception):
+    """Exception raised for HuggingFace authentication errors."""
+
+
+def check_file_permissions(file_path: Path) -> bool:
+    """
+    Check if file has secure permissions (600 or 400).
+
+    Args:
+        file_path: Path to token file
+
+    Returns:
+        True if permissions are secure, False otherwise
+    """
+    try:
+        file_stat = os.stat(file_path)
+        mode = stat.S_IMODE(file_stat.st_mode)
+
+        # Check if file is readable by others or group
+        if mode & (stat.S_IRGRP | stat.S_IROTH):
+            return False
+
+        # Check if file is writable by others or group
+        if mode & (stat.S_IWGRP | stat.S_IWOTH):
+            return False
+
+        return True
+    except OSError as e:
+        logger.error(f"Failed to check file permissions: {e}")
+        return False
+
+
+def load_token_from_file(token_file: str) -> str:
+    """
+    Load HuggingFace token from file with security checks.
+
+    Args:
+        token_file: Path to file containing HuggingFace token
+
+    Returns:
+        Token string
+
+    Raises:
+        HFAuthError: If file has insecure permissions or cannot be read
+    """
+    token_path = Path(token_file).expanduser().resolve()
+
+    # Check if file exists
+    if not token_path.exists():
+        raise HFAuthError(
+            f"Token file not found: {token_path}\n" f"Please create the file or check the path."
+        )
+
+    # Check if it's a file (not directory)
+    if not token_path.is_file():
+        raise HFAuthError(f"Token path is not a file: {token_path}")
+
+    # Check file permissions
+    if not check_file_permissions(token_path):
+        raise HFAuthError(
+            f"Token file has insecure permissions: {token_path}\n"
+            f"Please set secure permissions:\n"
+            f"  chmod 600 {token_path}\n"
+            f"Current permissions allow read/write by group or others."
+        )
+
+    # Read token
+    try:
+        with open(token_path, "r") as f:
+            token = f.read().strip()
+
+        if not token:
+            raise HFAuthError(f"Token file is empty: {token_path}")
+
+        # Basic validation (HF tokens start with 'hf_')
+        if not token.startswith("hf_"):
+            logger.warning(
+                f"Token from {token_path} doesn't start with 'hf_' - "
+                f"this may not be a valid HuggingFace token"
+            )
+
+        logger.info(f"Loaded HuggingFace token from {token_path}")
+        return token
+
+    except IOError as e:
+        raise HFAuthError(f"Failed to read token file {token_path}: {e}")
+
+
+def setup_hf_authentication(token_file: Optional[str] = None, use_env: bool = True) -> Optional[str]:
+    """
+    Setup HuggingFace authentication with multiple fallback options.
+
+    Priority:
+    1. Token from file (if token_file provided)
+    2. Token from HF_TOKEN environment variable (if use_env=True)
+    3. Token from HF CLI login (~/.cache/huggingface/token)
+    4. No authentication (public datasets only)
+
+    Args:
+        token_file: Optional path to token file
+        use_env: Whether to check HF_TOKEN environment variable
+
+    Returns:
+        Token string if found, None otherwise
+
+    Side effects:
+        Sets HF_TOKEN environment variable if token is found
+    """
+    # Priority 1: Token file
+    if token_file:
+        try:
+            token = load_token_from_file(token_file)
+            os.environ["HF_TOKEN"] = token
+            logger.info("Using HuggingFace token from file")
+            return token
+        except HFAuthError as e:
+            logger.error(str(e))
+            raise
+
+    # Priority 2: Environment variable
+    if use_env and "HF_TOKEN" in os.environ:
+        token = os.environ["HF_TOKEN"]
+        if token:
+            logger.info("Using HuggingFace token from HF_TOKEN environment variable")
+            return token
+
+    # Priority 3: HF CLI login
+    hf_cache_token = Path.home() / ".cache" / "huggingface" / "token"
+    if hf_cache_token.exists():
+        if not check_file_permissions(hf_cache_token):
+            logger.warning(
+                f"HuggingFace CLI token file has insecure permissions: {hf_cache_token}. "
+                "Skipping it; run `chmod 600` on the file to use it."
+            )
+        else:
+            try:
+                with open(hf_cache_token, "r") as f:
+                    token = f.read().strip()
+                if token:
+                    os.environ["HF_TOKEN"] = token
+                    logger.info("Using HuggingFace token from CLI login (~/.cache/huggingface/token)")
+                    return token
+            except IOError as e:
+                logger.debug(f"Could not read HF CLI token: {e}")
+
+    # No authentication found
+    logger.info("No HuggingFace authentication found. " "Only public datasets will be accessible.")
+    return None
+
+
+__all__ = [
+    "HFAuthError",
+    "load_token_from_file",
+    "setup_hf_authentication",
+    "check_file_permissions",
+]
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/download.py b/primus/backends/megatron/data/diffusion/preprocessing/download.py
new file mode 100644
index 000000000..c58c2814e
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/download.py
@@ -0,0 +1,156 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Reusable download utilities for dataset preparation pipelines.
+
+Provides HTTP download with exponential backoff, MD5 verification,
+and MLCommons R2 manifest resolution (.uri / .md5 protocol).
+Uses urllib.request (stdlib) -- no external HTTP dependencies needed.
+"""
+
+import hashlib
+import logging
+import random
+import shutil
+import time
+import urllib.error
+import urllib.request
+from http.client import HTTPResponse
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+_DOWNLOAD_TIMEOUT = 300  # 5 min per file
+_MAX_RETRIES = 5
+_BASE_DELAY = 1.0
+_MAX_DELAY = 60.0
+_USER_AGENT = "Wget/1.21"
+
+
+class _MD5MismatchError(Exception):
+    """Raised when a downloaded file's MD5 doesn't match the expected value."""
+
+
+def _is_retryable(status_code: int) -> bool:
+    return status_code == 429 or 500 <= status_code < 600
+
+
+def download_with_backoff(
+    url: str,
+    dest: Path,
+    expected_md5: Optional[str] = None,
+    max_retries: int = _MAX_RETRIES,
+    base_delay: float = _BASE_DELAY,
+    timeout: int = _DOWNLOAD_TIMEOUT,
+) -> None:
+    """Download a file via HTTP with exponential backoff on 429/5xx.
+
+    Streams the response to disk to avoid holding large files in memory.
+    Verifies MD5 checksum if provided.
+    """
+    dest.parent.mkdir(parents=True, exist_ok=True)
+
+    last_error: Optional[Exception] = None
+    for attempt in range(max_retries + 1):
+        try:
+            req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
+            resp: HTTPResponse = urllib.request.urlopen(req, timeout=timeout)
+            with open(dest, "wb") as f:
+                shutil.copyfileobj(resp, f)
+
+            if expected_md5:
+                actual_md5 = hashlib.md5(dest.read_bytes()).hexdigest()
+                if actual_md5 != expected_md5:
+                    dest.unlink(missing_ok=True)
+                    raise _MD5MismatchError(
+                        f"MD5 mismatch for {dest.name}: " f"expected {expected_md5}, got {actual_md5}"
+                    )
+            return
+
+        except urllib.error.HTTPError as e:
+            last_error = e
+            if not _is_retryable(e.code) or attempt == max_retries:
+                dest.unlink(missing_ok=True)
+                raise RuntimeError(
+                    f"HTTP {e.code} downloading {url} " f"(attempt {attempt + 1}/{max_retries + 1})"
+                ) from e
+            delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY)
+            logger.warning(f"  HTTP {e.code} on {url}, retry {attempt + 1}/{max_retries} " f"in {delay:.1f}s")
+            time.sleep(delay)
+
+        except (urllib.error.URLError, TimeoutError, OSError) as e:
+            last_error = e
+            if attempt == max_retries:
+                dest.unlink(missing_ok=True)
+                raise RuntimeError(f"Download failed for {url} after {max_retries + 1} attempts: {e}") from e
+            delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY)
+            logger.warning(
+                f"  Network error on {url}: {e}, retry {attempt + 1}/{max_retries} " f"in {delay:.1f}s"
+            )
+            time.sleep(delay)
+
+        except _MD5MismatchError as e:
+            last_error = e
+            if attempt == max_retries:
+                dest.unlink(missing_ok=True)
+                raise RuntimeError(str(e)) from e
+            delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY)
+            logger.warning(f"  {e}, retry {attempt + 1}/{max_retries} in {delay:.1f}s")
+            time.sleep(delay)
+
+    dest.unlink(missing_ok=True)
+    raise RuntimeError(f"Download failed for {url}: {last_error}")
+
+
+def fetch_url_text(url: str, timeout: int = 30) -> str:
+    """Fetch a small text resource (manifest, etc.) via HTTP."""
+    req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
+    resp = urllib.request.urlopen(req, timeout=timeout)
+    return resp.read().decode("utf-8")
+
+
+def parse_md5_manifest(
+    manifest_text: str,
+    suffix_filter: Optional[str] = None,
+) -> List[Tuple[str, str]]:
+    """Parse an MLCommons .md5 manifest into (md5, filename) pairs.
+
+    Args:
+        manifest_text: Raw text content of the .md5 file.
+        suffix_filter: If provided, only include files ending with this
+            suffix (e.g. ".arrow"). None means include all files.
+
+    Returns:
+        Sorted list of (md5, filename) tuples, sorted by filename
+        for deterministic ordering.
+    """
+    entries = []
+    for line in manifest_text.strip().splitlines():
+        parts = line.strip().split(None, 1)
+        if len(parts) != 2:
+            continue
+        md5, fname = parts
+        if suffix_filter is not None and not fname.endswith(suffix_filter):
+            continue
+        entries.append((md5, fname))
+    entries.sort(key=lambda x: x[1])
+    return entries
+
+
+def fetch_manifest(manifest_url: str) -> Tuple[str, List[Tuple[str, str]]]:
+    """Fetch .uri and .md5 manifests, return (base_url, [(md5, filename)]).
+
+    The manifest_url should end with '.uri' or '.md5'. The function derives
+    the complementary URL by replacing the suffix.
+    """
+    uri_url = manifest_url.replace(".md5", ".uri")
+    md5_url = manifest_url.replace(".uri", ".md5")
+
+    base_url = fetch_url_text(uri_url).strip()
+    md5_text = fetch_url_text(md5_url)
+    entries = parse_md5_manifest(md5_text)
+
+    logger.info(f"Manifest: {len(entries)} files, base URL: {base_url}")
+    return base_url, entries
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/finalize.py b/primus/backends/megatron/data/diffusion/preprocessing/finalize.py
new file mode 100644
index 000000000..a2c88fbc2
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/finalize.py
@@ -0,0 +1,230 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Automatic Energon dataset finalization for Primus diffusion pipelines.
+
+This module automates the post-processing steps required after data preprocessing:
+1. Create .nv-meta/dataset.yaml with correct sample type configuration
+2. Run ``energon prepare`` non-interactively to index the dataset
+3. Verify the dataset is ready for training
+
+Supports two indexing strategies:
+- **Ratio-based** (default): splits shards by train/val/test ratio.
+  Works with both the programmatic ``BaseWebdatasetFactory`` API and
+  the ``energon prepare`` subprocess fallback.
+- **Pattern-based**: assigns shards to splits by regex patterns on
+  relative paths (e.g. ``"train/.*"``).  Requires the programmatic
+  API (``megatron-energon >= 7.x``).
+
+Usage:
+    from primus.backends.megatron.data.diffusion.preprocessing.finalize import finalize_energon_dataset
+
+    # Ratio-based (original behavior)
+    finalize_energon_dataset(
+        output_dir='/workspace/Primus/data/encoded_pokemon',
+        train_split=1.0,
+        encoding='preencoded',
+    )
+
+    # Pattern-based (MLPerf layout with separate train/ and val/ dirs)
+    finalize_energon_dataset(
+        output_dir='/workspace/Primus/data/mlperf_flux1',
+        encoding='preencoded_numpy',
+        split_parts_patterns=[("train", "train/.*"), ("val", "val/.*")],
+    )
+"""
+
+import logging
+import subprocess
+from pathlib import Path
+from typing import List, Literal, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+_ENCODING_TYPE = Literal["preencoded", "preencoded_numpy", "raw"]
+
+
+def _write_dataset_yaml(output_path: Path, encoding: str) -> Path:
+    """Write .nv-meta/dataset.yaml with the correct subflavor."""
+    meta_dir = output_path / ".nv-meta"
+    meta_dir.mkdir(exist_ok=True)
+
+    dataset_yaml_path = meta_dir / "dataset.yaml"
+    dataset_yaml_path.write_text(
+        f"__module__: megatron.energon\n"
+        f"__class__: CrudeWebdataset\n"
+        f"subflavors:\n"
+        f"  encoding: {encoding}\n"
+    )
+    logger.info(f"✓ Created {dataset_yaml_path}")
+    return dataset_yaml_path
+
+
+def _prepare_programmatic(
+    output_path: Path,
+    num_workers: int,
+    train_split: float,
+    split_parts_patterns: Optional[List[Tuple[str, str]]],
+) -> None:
+    """Index dataset using BaseWebdatasetFactory.prepare_dataset()."""
+    from megatron.energon import BaseWebdatasetFactory
+
+    shard_paths = sorted(str(p.relative_to(output_path)) for p in output_path.glob("**/*.tar"))
+    if not shard_paths:
+        raise FileNotFoundError(f"No .tar shards found under {output_path}")
+
+    kwargs: dict = {
+        "parent_path": str(output_path),
+        "paths": shard_paths,
+        "workers": num_workers,
+    }
+
+    if split_parts_patterns:
+        kwargs["split_parts_patterns"] = split_parts_patterns
+    else:
+        val_split = (1.0 - train_split) / 2
+        test_split = (1.0 - train_split) / 2
+        kwargs["split_parts_ratio"] = [
+            ("train", train_split),
+            ("val", val_split),
+            ("test", test_split),
+        ]
+
+    logger.info(f"  Indexing with BaseWebdatasetFactory ({len(shard_paths)} shards)")
+    BaseWebdatasetFactory.prepare_dataset(**kwargs)
+    logger.info("✓ Energon indexing complete (programmatic API)")
+
+
+def _prepare_subprocess(
+    output_path: Path,
+    num_workers: int,
+    train_split: float,
+    split_parts_patterns: Optional[List[Tuple[str, str]]],
+) -> None:
+    """Fallback: index dataset using ``energon prepare`` subprocess.
+
+    NOTE: The subprocess path only supports ratio-based splitting.
+    Pattern-based ``split_parts_patterns`` requires the programmatic API.
+    """
+    if split_parts_patterns:
+        raise RuntimeError(
+            "Pattern-based split_parts_patterns requires "
+            "megatron.energon.BaseWebdatasetFactory (programmatic API). "
+            "Install megatron-energon >= 7.x or use ratio-based splitting."
+        )
+
+    val_split = (1.0 - train_split) / 2
+    test_split = (1.0 - train_split) / 2
+    split_input = f"{train_split}, {val_split}, {test_split}\nn\n"
+
+    logger.info(f"  Split ratios: train={train_split:.2f}, " f"val={val_split:.2f}, test={test_split:.2f}")
+    logger.info("  Running energon prepare subprocess...")
+
+    try:
+        process = subprocess.Popen(
+            [
+                "energon",
+                "prepare",
+                str(output_path),
+                "--num-workers",
+                str(num_workers),
+            ],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+        )
+
+        stdout, _ = process.communicate(input=split_input, timeout=300)
+
+        if process.returncode != 0:
+            logger.error("energon prepare failed:")
+            logger.error(stdout)
+            raise RuntimeError(f"energon prepare exited with code {process.returncode}")
+
+        for line in stdout.splitlines():
+            if any(skip in line for skip in ["libibverbs", "Warning:", "UserWarning"]):
+                continue
+            if any(key in line for key in ["Indexing", "Done", "Found", "samples", "shards"]):
+                logger.info(f"  {line.strip()}")
+
+        logger.info("✓ Energon indexing complete (subprocess)")
+
+    except subprocess.TimeoutExpired:
+        process.kill()
+        raise RuntimeError(
+            "energon prepare timed out after 5 minutes. "
+            "This may indicate a problem with the dataset or system resources."
+        )
+    except FileNotFoundError:
+        raise RuntimeError(
+            "energon command not found. Make sure megatron-energon is installed:\n"
+            "  pip install megatron-energon"
+        )
+
+
+def finalize_energon_dataset(
+    output_dir: str,
+    train_split: float = 1.0,
+    encoding: _ENCODING_TYPE = "preencoded",
+    num_workers: int = 8,
+    split_parts_patterns: Optional[List[Tuple[str, str]]] = None,
+) -> None:
+    """
+    Finalize Energon WebDataset by creating dataset.yaml and running energon prepare.
+
+    This automates the manual steps typically required after data preprocessing:
+    1. Create .nv-meta/dataset.yaml with correct subflavor configuration
+    2. Run energon indexing (programmatic API with subprocess fallback)
+    3. Verify the dataset is ready for training
+
+    Args:
+        output_dir: Path to dataset directory containing tar files.
+        train_split: Fraction for training (rest split evenly between
+            val/test). Ignored when split_parts_patterns is provided.
+        encoding: Dataset encoding mode.
+        num_workers: Number of workers for energon prepare (default: 8).
+        split_parts_patterns: List of (split_name, pattern) tuples for
+            pattern-based splitting, e.g.
+            ``[("train", "train/.*"), ("val", "val/.*")]``.
+            When provided, train_split is ignored.
+
+    Raises:
+        RuntimeError: If energon prepare fails or energon command not found
+        FileNotFoundError: If output_dir doesn't exist or has no tar files
+    """
+    output_path = Path(output_dir)
+
+    if not output_path.exists():
+        raise FileNotFoundError(f"Output directory not found: {output_dir}")
+
+    tar_files = list(output_path.glob("**/*.tar"))
+    if not tar_files:
+        raise FileNotFoundError(
+            f"No .tar files found in {output_dir}. " "Make sure data preprocessing completed successfully."
+        )
+
+    logger.info("=" * 80)
+    logger.info("Finalizing Energon dataset for training...")
+    logger.info("=" * 80)
+    logger.info(f"Found {len(tar_files)} shard(s) to index")
+
+    _write_dataset_yaml(output_path, encoding)
+
+    logger.info("Running energon prepare (this may take a few minutes)...")
+
+    try:
+        _prepare_programmatic(output_path, num_workers, train_split, split_parts_patterns)
+    except ImportError:
+        logger.info("BaseWebdatasetFactory not available, falling back to subprocess")
+        _prepare_subprocess(output_path, num_workers, train_split, split_parts_patterns)
+
+    from .validate import validate_energon_dataset
+
+    validate_energon_dataset(output_dir, encoding=encoding)
+
+    logger.info("=" * 80)
+    logger.info(f"✓ Dataset finalized: {output_dir}")
+    logger.info(f"  To use in training, set: dataset_path: {output_dir}")
+    logger.info("=" * 80)
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py
new file mode 100644
index 000000000..15280d896
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py
@@ -0,0 +1,21 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Dataset preparation pipelines for Megatron diffusion models.
+
+Provides reusable pipeline classes for creating Energon WebDatasets
+from various input sources (HuggingFace, directories, WebDatasets).
+"""
+
+from .base import DatasetPipeline
+from .encoded import EncodedDatasetPipeline
+from .ingest import StreamingIngestPipeline
+from .raw import RawDatasetPipeline
+
+__all__ = [
+    "DatasetPipeline",
+    "RawDatasetPipeline",
+    "EncodedDatasetPipeline",
+    "StreamingIngestPipeline",
+]
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py
new file mode 100644
index 000000000..65d579cea
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py
@@ -0,0 +1,70 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""Abstract base class for dataset preparation pipelines."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from ..utils import load_from_directory, load_from_huggingface, load_from_webdataset
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetPipeline(ABC):
+    """Abstract base for all dataset preparation pipelines.
+
+    Subclasses implement run() which executes the full pipeline
+    and returns a dict of statistics (samples_processed, shards_written, etc.).
+    """
+
+    # Whether HuggingFace sources are loaded in streaming mode. Subclasses that
+    # need the full dataset materialized (e.g. for distributed total-count
+    # splitting) set this to False.
+    HF_STREAMING: bool = True
+
+    @abstractmethod
+    def run(self, **kwargs) -> Dict[str, Any]: ...
+
+    def load_data(self, **source_kwargs):
+        """Load data based on ``self.source_type``.
+
+        Args:
+            **source_kwargs: Source-specific arguments
+                - directory: input_dir
+                - huggingface: hf_dataset, hf_split, hf_data_files, image/caption keys
+                - webdataset: input_path
+
+        Returns:
+            Iterator over samples with 'image' and 'caption' keys.
+        """
+        if self.source_type == "directory":
+            logger.info(f"Loading from directory: {source_kwargs['input_dir']}")
+            return load_from_directory(source_kwargs["input_dir"])
+        elif self.source_type == "huggingface":
+            hf_split = source_kwargs.get("hf_split", "train")
+            hf_data_files = source_kwargs.get("hf_data_files")
+            if hf_data_files:
+                logger.info(
+                    f"Loading from HuggingFace: {source_kwargs['hf_dataset']} "
+                    f"(split: {hf_split}, data_files: {hf_data_files})"
+                )
+            else:
+                logger.info(f"Loading from HuggingFace: {source_kwargs['hf_dataset']} (split: {hf_split})")
+
+            return load_from_huggingface(
+                source_kwargs["hf_dataset"],
+                split=hf_split,
+                streaming=self.HF_STREAMING,
+                data_files=hf_data_files,
+                image_key=source_kwargs.get("image_key"),
+                caption_key=source_kwargs.get("caption_key"),
+                image_keys=source_kwargs.get("image_keys"),
+                caption_keys=source_kwargs.get("caption_keys"),
+            )
+        elif self.source_type == "webdataset":
+            logger.info(f"Loading from WebDataset: {source_kwargs['input_path']}")
+            return load_from_webdataset(source_kwargs["input_path"])
+        else:
+            raise ValueError(f"Unknown source type: {self.source_type}")
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py
new file mode 100644
index 000000000..ba4d171ba
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py
@@ -0,0 +1,608 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Pre-encoded dataset preparation pipeline for Megatron diffusion models.
+
+Creates pre-encoded Energon WebDatasets with VAE/T5/CLIP encoded tensors
+for faster training (no on-the-fly encoding overhead).
+"""
+
+import logging
+import time
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from primus.backends.megatron.data.diffusion.encoders import get_encoder
+from primus.backends.megatron.data.diffusion.encoders.config import (
+    CLIPLConfig,
+    T5XXLConfig,
+    VAEConfig,
+)
+
+from ..utils import (
+    get_distributed_info,
+    preprocess_image,
+    save_to_webdataset,
+    split_work_for_rank,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def set_reproducibility(seed: int = 42) -> None:
+    """Set global seeds and cuDNN flags for reproducible encoding.
+
+    Called explicitly from the pipeline entry point rather than at import time,
+    so importing this module does not mutate global torch/cuDNN state for
+    unrelated callers.
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+from .base import DatasetPipeline
+
+
+class EncodedDatasetPipeline(DatasetPipeline):
+    """
+    Pipeline for creating pre-encoded Energon WebDataset with VAE/T5/CLIP tensors.
+
+    This pipeline pre-encodes images with VAE and text with T5/CLIP, creating larger
+    datasets that enable faster training (no encoding overhead).
+
+    Args:
+        source_type: Type of input source ('directory', 'huggingface', 'webdataset')
+        output_dir: Output directory for Energon WebDataset
+        model_path: Pretrained model path (HF or local, default: black-forest-labs/FLUX.1-dev)
+        vae_path: Custom VAE path (overrides model_path)
+        t5_path: Custom T5 path (overrides model_path)
+        clip_path: Custom CLIP path (overrides model_path)
+        precision: Model precision ('bf16', 'fp16', 'fp32', default: 'bf16')
+        device: Device for encoding (default: 'cuda')
+        batch_size: Encoding batch size (default: 8)
+        t5_max_length: T5 max sequence length (default: 512, use 256 for schnell)
+        image_size: Target image size (default: 1024)
+        variable_size: If True, resize to nearest multiple of 16 instead of fixed size (default: False)
+        center_crop: Whether to center crop images (default: True)
+        max_size: Maximum dimension when variable_size is True (default: 1024)
+        shard_size: Samples per shard (default: 1000)
+        max_samples: Maximum samples to process (default: None for all)
+        compress: Whether to compress tar files with gzip (default: False)
+        hf_token_file: Path to HuggingFace token file (default: None)
+        vae_latent_mode: 'presampled' stores only latents; 'resample' additionally
+            stores mean and logvar for training-time reparameterization (default: 'presampled')
+
+    Example:
+        >>> pipeline = EncodedDatasetPipeline(
+        ...     source_type='huggingface',
+        ...     output_dir='/data/encoded_pokemon',
+        ...     model_path='black-forest-labs/FLUX.1-dev',
+        ...     batch_size=8,
+        ...     image_size=1024,
+        ... )
+        >>> results = pipeline.run(
+        ...     hf_dataset='diffusers/pokemon-gpt4-captions',
+        ...     hf_split='train'
+        ... )
+        >>> print(f"Processed {results['samples_processed']} samples")
+    """
+
+    # Encoded pipeline materializes the full HF dataset (non-streaming) so the
+    # distributed loader can split by total count.
+    HF_STREAMING: bool = False
+
+    def __init__(
+        self,
+        source_type: str,
+        output_dir: str,
+        model_path: str = "black-forest-labs/FLUX.1-dev",
+        vae_path: Optional[str] = None,
+        t5_path: Optional[str] = None,
+        clip_path: Optional[str] = None,
+        precision: str = "bf16",
+        device: str = "cuda",
+        batch_size: int = 8,
+        t5_max_length: int = 512,
+        variable_size: bool = False,
+        image_size: int = 1024,
+        center_crop: bool = True,
+        max_size: int = 1024,
+        shard_size: int = 1000,
+        max_samples: Optional[int] = None,
+        compress: bool = False,
+        hf_token_file: Optional[str] = None,
+        vae_latent_mode: str = "presampled",
+    ):
+        if vae_latent_mode not in ("presampled", "resample"):
+            raise ValueError(f"vae_latent_mode must be 'presampled' or 'resample', got '{vae_latent_mode}'")
+
+        self.source_type = source_type
+        self.output_dir = Path(output_dir)
+        self.model_path = model_path
+        self.vae_path = vae_path
+        self.t5_path = t5_path
+        self.clip_path = clip_path
+        self.precision = precision
+        self.device = device
+        self.batch_size = batch_size
+        self.t5_max_length = t5_max_length
+        self.variable_size = variable_size
+        self.image_size = image_size
+        self.center_crop = center_crop
+        self.max_size = max_size
+        self.shard_size = shard_size
+        self.max_samples = max_samples
+        self.compress = compress
+        self.vae_latent_mode = vae_latent_mode
+
+        # Setup HF authentication if token file provided
+        if hf_token_file:
+            from ..auth import HFAuthError, setup_hf_authentication
+
+            try:
+                setup_hf_authentication(token_file=hf_token_file)
+            except HFAuthError as e:
+                raise ValueError(f"HuggingFace authentication failed: {e}")
+
+        # Get distributed info
+        self.rank, self.world_size = get_distributed_info()
+
+        logger.info(f"Initialized EncodedDatasetPipeline (rank {self.rank}/{self.world_size})")
+        logger.info(f"Model: {self.model_path}")
+        logger.info(f"Precision: {self.precision}")
+        logger.info(f"Batch size: {self.batch_size}")
+        logger.info(f"T5 max sequence length: {self.t5_max_length}")
+        logger.info(f"VAE latent mode: {self.vae_latent_mode}")
+
+        # Load encoders
+        self.vae, self.t5, self.clip = self._load_encoders()
+
+    @staticmethod
+    def _raise_encoder_auth_error(encoder_name: str, model_path: str, original_error: Exception):
+        """Raise a clear error when encoder download fails due to authentication."""
+        err_str = str(original_error).lower()
+        is_auth = any(
+            kw in err_str
+            for kw in [
+                "token",
+                "permission",
+                "private repository",
+                "gated",
+                "401",
+                "403",
+                "authentication",
+                "login",
+            ]
+        )
+        if is_auth:
+            raise RuntimeError(
+                f"Failed to download {encoder_name} encoder from '{model_path}'.\n"
+                f"This model likely requires HuggingFace authentication.\n\n"
+                f"To fix, provide a token using one of:\n"
+                f"  1. --hf-token-file /path/to/.hf_token\n"
+                f"  2. export HF_TOKEN=hf_xxx\n"
+                f"  3. huggingface-cli login\n"
+            ) from original_error
+        raise RuntimeError(
+            f"Failed to load {encoder_name} encoder from '{model_path}': {original_error}"
+        ) from original_error
+
+    def _load_encoders(self):
+        """Load VAE, T5, and CLIP encoders."""
+        logger.info("Loading encoders...")
+
+        # Determine if we're using FLUX model (needs subfolders)
+        vae_model_path = self.vae_path or self.model_path
+        t5_model_path = self.t5_path or self.model_path
+        clip_model_path = self.clip_path or self.model_path
+
+        is_flux = "FLUX" in self.model_path or "flux" in self.model_path.lower()
+
+        # Create encoder configs with subfolders for FLUX models
+        vae_config = VAEConfig(
+            type="autoencoder_kl",
+            model_path=vae_model_path,
+            subfolder="vae" if (is_flux and not self.vae_path) else None,
+            precision=self.precision,
+            device=self.device,
+        )
+
+        t5_config = T5XXLConfig(
+            type="t5_xxl",
+            model_path=t5_model_path,
+            subfolder="text_encoder_2" if (is_flux and not self.t5_path) else None,
+            tokenizer_subfolder="tokenizer_2" if (is_flux and not self.t5_path) else None,
+            max_length=self.t5_max_length,
+            precision=self.precision,
+            device=self.device,
+        )
+
+        clip_config = CLIPLConfig(
+            type="clip_l",
+            model_path=clip_model_path,
+            subfolder="text_encoder" if (is_flux and not self.clip_path) else None,
+            tokenizer_subfolder="tokenizer" if (is_flux and not self.clip_path) else None,
+            precision=self.precision,
+            device=self.device,
+        )
+
+        # Load encoders with clear error messages on failure
+        try:
+            vae = get_encoder(vae_config)
+        except Exception as e:
+            self._raise_encoder_auth_error("VAE", vae_config.model_path, e)
+        logger.info(f"Loaded VAE from {vae_config.model_path}")
+
+        try:
+            t5 = get_encoder(t5_config)
+        except Exception as e:
+            self._raise_encoder_auth_error("T5-XXL", t5_config.model_path, e)
+        logger.info(f"Loaded T5-XXL from {t5_config.model_path}")
+
+        try:
+            clip = get_encoder(clip_config)
+        except Exception as e:
+            self._raise_encoder_auth_error("CLIP-L", clip_config.model_path, e)
+        logger.info(f"Loaded CLIP-L from {clip_config.model_path}")
+
+        # Set to eval mode
+        vae.eval()
+        t5.eval()
+        clip.eval()
+
+        return vae, t5, clip
+
+    def load_data_distributed(self, **source_kwargs):
+        """
+        Load only the data needed for this rank (distributed-aware loading).
+
+        This method loads data more efficiently by having each rank load only
+        its assigned portion of the dataset, rather than loading everything
+        and then splitting.
+
+        Args:
+            **source_kwargs: Source-specific arguments
+                - For directory: input_dir
+                - For huggingface: hf_dataset, hf_split
+                - For webdataset: input_path
+
+        Returns:
+            List of samples assigned to this rank
+        """
+        if self.source_type == "directory":
+            # Get file list (cheap operation, all ranks do this)
+            from pathlib import Path
+
+            from PIL import Image
+
+            input_path = Path(source_kwargs["input_dir"])
+            images_dir = input_path / "images"
+            captions_dir = input_path / "captions"
+
+            if not images_dir.exists():
+                raise ValueError(f"Images directory not found: {images_dir}")
+            if not captions_dir.exists():
+                raise ValueError(f"Captions directory not found: {captions_dir}")
+
+            # Get all image files
+            image_extensions = [".jpg", ".jpeg", ".png", ".webp"]
+            all_image_files = []
+            for ext in image_extensions:
+                all_image_files.extend(sorted(images_dir.glob(f"*{ext}")))
+
+            total_files = len(all_image_files)
+            logger.info(f"Found {total_files} total images")
+
+            # Apply max_samples limit BEFORE splitting across ranks so that
+            # max_samples refers to the total number of samples, not per-rank.
+            if self.max_samples and total_files > self.max_samples:
+                logger.info(f"Limiting to {self.max_samples} total samples (before rank split)")
+                all_image_files = all_image_files[: self.max_samples]
+                total_files = self.max_samples
+
+            # Split file list across ranks (before loading!)
+            start_idx, end_idx = split_work_for_rank(total_files, self.rank, self.world_size)
+            my_files = all_image_files[start_idx:end_idx]
+
+            logger.info(f"Rank {self.rank}: Loading {len(my_files)} images (indices {start_idx}-{end_idx})")
+
+            # Now load only this rank's files
+            items = []
+            for img_path in my_files:
+                caption_path = captions_dir / f"{img_path.stem}.txt"
+                if not caption_path.exists():
+                    logger.warning(f"Caption not found for {img_path.name}, skipping")
+                    continue
+
+                try:
+                    image = Image.open(img_path).convert("RGB")
+                    with open(caption_path, "r", encoding="utf-8") as f:
+                        caption = f.read().strip()
+                    items.append({"image": image, "caption": caption})
+                except Exception as e:
+                    logger.warning(f"Failed to load {img_path.name}: {e}")
+
+            return items
+        else:
+            # For non-directory sources, fall back to loading all data
+            logger.warning(
+                f"Distributed loading not yet implemented for {self.source_type}, loading all data"
+            )
+            data_iter = self.load_data(**source_kwargs)
+            all_items = list(data_iter)
+            total_items = len(all_items)
+
+            # Apply max_samples limit BEFORE splitting across ranks so that
+            # max_samples refers to the total number of samples, not per-rank.
+            if self.max_samples and total_items > self.max_samples:
+                logger.info(f"Limiting to {self.max_samples} total samples (before rank split)")
+                all_items = all_items[: self.max_samples]
+                total_items = self.max_samples
+
+            # Split work across ranks
+            start_idx, end_idx = split_work_for_rank(total_items, self.rank, self.world_size)
+            return all_items[start_idx:end_idx]
+
+    def _preprocess_images_batch(self, images):
+        """
+        Preprocess batch of PIL images for VAE encoding.
+
+        Args:
+            images: List of PIL Images
+
+        Returns:
+            Tensor [B, C, H, W] in range [-1, 1]
+        """
+        tensors = []
+        for image in images:
+            # Ensure RGB mode
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+
+            # Convert to numpy array and normalize to [-1, 1]
+            img_array = np.array(image).astype(np.float32) / 255.0
+            img_array = img_array * 2.0 - 1.0
+
+            # Convert to tensor (HWC -> CHW)
+            img_tensor = torch.from_numpy(np.transpose(img_array, (2, 0, 1)))
+            tensors.append(img_tensor)
+
+        return torch.stack(tensors)
+
+    def _encode_batch(self, images, captions):
+        """
+        Encode a batch of images and captions.
+
+        Position IDs are NOT generated during preprocessing - they will be
+        computed at runtime based on actual tensor shapes. This provides
+        flexibility for variable-resolution training.
+
+        Args:
+            images: List of PIL Images
+            captions: List of caption strings
+
+        Returns:
+            List of sample dicts with encoded tensors
+        """
+        with torch.no_grad():
+            # Preprocess images
+            images_tensor = self._preprocess_images_batch(images)
+            images_tensor = images_tensor.to(self.device)
+
+            # Encode images with VAE
+            if self.vae_latent_mode == "resample":
+                latents, mean, logvar = self.vae.encode_for_resample(images_tensor)
+            else:
+                latents = self.vae.encode(images_tensor)
+
+            # Encode text with T5
+            prompt_embeds = self.t5.encode(captions)
+
+            # Encode text with CLIP
+            _, pooled_prompt_embeds = self.clip.encode(captions)
+
+        # Move to CPU and create samples (NO position IDs)
+        samples = []
+        for i in range(len(images)):
+            sample = {
+                "latents.pth": latents[i].cpu(),
+                "prompt_embeds.pth": prompt_embeds[i].cpu(),
+                "pooled_prompt_embeds.pth": pooled_prompt_embeds[i].cpu(),
+            }
+            if self.vae_latent_mode == "resample":
+                sample["mean.pth"] = mean[i].cpu()
+                sample["logvar.pth"] = logvar[i].cpu()
+            samples.append(sample)
+
+        return samples
+
+    def run(self, **source_kwargs):
+        """
+        Execute the encoded dataset preparation pipeline.
+
+        Args:
+            **source_kwargs: Source-specific arguments (passed to load_data)
+
+        Returns:
+            Dictionary with processing statistics:
+                - samples_processed: Number of samples successfully processed
+                - samples_skipped: Number of samples that failed
+                - shards_written: Number of output shards created
+        """
+        import torch.distributed as dist
+
+        # Reproducible encoding: set seeds/cuDNN flags here (not at import time).
+        set_reproducibility()
+
+        start_time = time.time()
+        rank_prefix = f"[RANK {self.rank}]"
+
+        tqdm.write(f"{rank_prefix} Stage 1/4: Initialization")
+        tqdm.write(f"{rank_prefix} Source: {self.source_type}")
+        tqdm.write(
+            f"{rank_prefix} Image preprocessing: size={self.image_size}, crop={self.center_crop}, variable_size={self.variable_size}"
+        )
+
+        # Synchronization point: ensure all ranks start together
+        if dist.is_initialized():
+            dist.barrier()
+
+        # Load data using distributed-aware loading (each rank loads only its portion)
+        tqdm.write(f"{rank_prefix} Stage 2/4: Loading data")
+        items_to_process = self.load_data_distributed(**source_kwargs)
+        tqdm.write(f"{rank_prefix} Loaded {len(items_to_process)} items")
+
+        # Synchronization point: ensure all ranks finished loading
+        if dist.is_initialized():
+            dist.barrier()
+
+        # Sort items by dimensions AFTER splitting to ensure batches have uniform sizes
+        # This prevents tensor stacking errors while maintaining balanced workload across ranks
+        items_to_process = sorted(
+            items_to_process, key=lambda item: item["image"].size  # Returns (width, height)
+        )
+
+        total_to_process = len(items_to_process)
+
+        tqdm.write(f"{rank_prefix} Stage 3/4: Encoding samples")
+
+        # Process in batches
+        batch_images = []
+        batch_captions = []
+        encoded_samples = []
+        samples_processed = 0
+        samples_skipped = 0
+        shards_written = 0
+
+        pbar = tqdm(
+            total=total_to_process,
+            desc=f"{rank_prefix} Encoding",
+            unit="sample",
+            disable=(self.rank != 0 and self.world_size > 4),
+        )
+
+        def _flush_batch():
+            """Encode and accumulate the current batch, returning count."""
+            nonlocal batch_images, batch_captions, encoded_samples, samples_processed
+            nonlocal shards_written
+            batch_samples = self._encode_batch(batch_images, batch_captions)
+            encoded_samples.extend(batch_samples)
+            count = len(batch_samples)
+            samples_processed += count
+            pbar.update(count)
+            batch_images = []
+            batch_captions = []
+
+            if len(encoded_samples) >= self.shard_size:
+                shard_offset = shards_written * self.world_size + self.rank
+                num_shards = save_to_webdataset(
+                    encoded_samples,
+                    str(self.output_dir),
+                    self.shard_size,
+                    shard_offset=shard_offset,
+                    compress=self.compress,
+                )
+                shards_written += num_shards
+                encoded_samples = []
+
+        for idx, item in enumerate(items_to_process):
+            try:
+                image = item["image"]
+                if self.image_size or self.center_crop or self.variable_size:
+                    image = preprocess_image(
+                        image,
+                        variable_size=self.variable_size,
+                        size=self.image_size,
+                        center_crop=self.center_crop,
+                        max_size=self.max_size,
+                    )
+
+                current_size = image.size
+
+                if batch_images and batch_images[0].size != current_size:
+                    _flush_batch()
+
+                batch_images.append(image)
+                batch_captions.append(item["caption"])
+
+                if len(batch_images) >= self.batch_size:
+                    _flush_batch()
+
+            except Exception as e:
+                tqdm.write(f"{rank_prefix} Failed to encode sample {idx}: {e}")
+                logger.debug(f"Sample {idx} encoding error", exc_info=True)
+                samples_skipped += 1
+                batch_images = []
+                batch_captions = []
+                continue
+
+        # Process remaining batch
+        if batch_images:
+            try:
+                _flush_batch()
+            except Exception as e:
+                tqdm.write(f"{rank_prefix} Failed to encode final batch: {e}")
+
+        # Save remaining samples
+        if encoded_samples:
+            shard_offset = shards_written * self.world_size + self.rank
+            num_shards = save_to_webdataset(
+                encoded_samples,
+                str(self.output_dir),
+                self.shard_size,
+                shard_offset=shard_offset,
+                compress=self.compress,
+            )
+            shards_written += num_shards
+
+        pbar.close()
+
+        elapsed_time = time.time() - start_time
+        tqdm.write(f"{rank_prefix} Stage 4/4: Complete")
+        tqdm.write(
+            f"{rank_prefix} Finished in {elapsed_time:.1f}s — "
+            f"processed: {samples_processed}, skipped: {samples_skipped}, shards: {shards_written}"
+        )
+
+        # Generate empty encodings for CFG dropout (rank 0 only).
+        # Uses the same T5/CLIP models already loaded with the same t5_max_length,
+        # guaranteeing sequence length consistency with the encoded dataset.
+        if self.rank == 0:
+            self._generate_empty_encodings()
+
+        # Critical synchronization point: ensure all ranks finished encoding before returning
+        if dist.is_initialized():
+            dist.barrier()
+
+        return {
+            "samples_processed": samples_processed,
+            "samples_skipped": samples_skipped,
+            "shards_written": shards_written,
+        }
+
+    def _generate_empty_encodings(self):
+        """Generate and save T5/CLIP encodings for the empty string (rank 0 only)."""
+        empty_dir = self.output_dir / "empty_encodings"
+        empty_dir.mkdir(parents=True, exist_ok=True)
+
+        with torch.no_grad():
+            t5_empty = self.t5.encode([""])
+            _, clip_empty = self.clip.encode([""])
+
+        t5_np = t5_empty.cpu().float().numpy()
+        clip_np = clip_empty.cpu().float().numpy()
+
+        np.save(str(empty_dir / "t5_empty.npy"), t5_np)
+        np.save(str(empty_dir / "clip_empty.npy"), clip_np)
+
+        tqdm.write(
+            f"[RANK 0] Saved empty encodings to {empty_dir} " f"(t5={t5_np.shape}, clip={clip_np.shape})"
+        )
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py
new file mode 100644
index 000000000..5c68bc572
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py
@@ -0,0 +1,311 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Streaming Arrow-to-WebDataset pipeline for MLPerf Flux datasets.
+
+Downloads Apache Arrow IPC files from MLCommons R2 storage using a
+concurrent prefetch pool, converts numpy-serialized bfloat16 samples
+directly into WebDataset tar shards, then deletes each Arrow file to
+minimize peak disk usage.
+
+Reduces the 6 TB disk requirement to roughly the final dataset size
+(~1.2 TB for cc12m) plus a small prefetch buffer of Arrow files.
+"""
+
+import io
+import json
+import logging
+import queue
+import tarfile
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import pyarrow.ipc
+
+from ..download import download_with_backoff, fetch_manifest
+from .base import DatasetPipeline
+
+logger = logging.getLogger(__name__)
+
+ARROW_COLUMNS = ("t5_encodings", "clip_encodings", "mean", "logvar")
+WDS_KEYS = ("t5.bytes", "clip.bytes", "mean.bytes", "logvar.bytes")
+
+_SENTINEL = None
+
+
+def _arrow_to_tar(
+    arrow_path: Path,
+    tar_path: Path,
+    global_sample_offset: int,
+) -> int:
+    """Convert one Arrow IPC stream file into a WebDataset tar shard.
+
+    Reads the Arrow stream, writes each row as a set of compound-key
+    entries (e.g. ``00000000.t5.bytes``) into a tar archive. Uses the
+    ``__key__`` column from the Arrow file when available, otherwise
+    generates sequential keys.
+
+    Returns the number of samples written.
+    """
+    reader = pyarrow.ipc.open_stream(str(arrow_path))
+    table = reader.read_all()
+    num_rows = table.num_rows
+
+    has_key_col = "__key__" in table.schema.names
+
+    tar_path.parent.mkdir(parents=True, exist_ok=True)
+    with tarfile.open(str(tar_path), "w") as tar:
+        for row_idx in range(num_rows):
+            if has_key_col:
+                base_name = table.column("__key__")[row_idx].as_py()
+            else:
+                base_name = f"{global_sample_offset + row_idx:08d}"
+
+            for col_name, wds_key in zip(ARROW_COLUMNS, WDS_KEYS):
+                col = table.column(col_name)
+                raw_bytes = col[row_idx].as_py()
+                if raw_bytes is None:
+                    continue
+                data = bytes(raw_bytes)
+
+                entry_name = f"{base_name}.{wds_key}"
+                info = tarfile.TarInfo(name=entry_name)
+                info.size = len(data)
+                tar.addfile(info, io.BytesIO(data))
+
+            meta = json.dumps({"key": base_name}).encode("utf-8")
+            meta_info = tarfile.TarInfo(name=f"{base_name}.json")
+            meta_info.size = len(meta)
+            tar.addfile(meta_info, io.BytesIO(meta))
+
+    return num_rows
+
+
+class StreamingIngestPipeline(DatasetPipeline):
+    """Streaming pipeline: download Arrow files from R2 -> WebDataset tars.
+
+    Downloads are parallelized with a prefetch pool (``max_workers`` threads)
+    while conversion runs sequentially on the main thread, preserving
+    deterministic shard ordering.  A semaphore (``prefetch_depth``) limits
+    how many Arrow files are kept on disk at once.
+
+    Args:
+        manifest_url: URL to the .uri manifest on MLCommons R2.
+        input_dir: Local directory for temporary Arrow file storage.
+        output_dir: Output directory for WebDataset tar shards.
+        split_name: Subdirectory name for the split (e.g. 'train', 'val').
+        max_files: Maximum number of Arrow files to process (None = all).
+        max_workers: Number of concurrent download threads (default 4).
+        prefetch_depth: Max Arrow files buffered on disk ahead of
+            conversion (default 6).  Disk overhead ~ prefetch_depth x 200 MB.
+    """
+
+    def __init__(
+        self,
+        manifest_url: str,
+        input_dir: str,
+        output_dir: str,
+        split_name: str = "train",
+        max_files: Optional[int] = None,
+        max_workers: int = 4,
+        prefetch_depth: int = 6,
+    ):
+        self.manifest_url = manifest_url
+        self.input_dir = Path(input_dir)
+        self.output_dir = Path(output_dir) / split_name
+        self.split_name = split_name
+        self.max_files = max_files
+        self.max_workers = max_workers
+        self.prefetch_depth = prefetch_depth
+
+    def run(self, **kwargs) -> Dict[str, int]:
+        """Execute the streaming conversion pipeline.
+
+        Returns:
+            Dict with 'files_processed', 'samples_written', 'shards_created',
+            'shards_skipped', 'files_failed'.
+        """
+        base_url, entries = fetch_manifest(self.manifest_url)
+
+        if self.max_files is not None:
+            entries = entries[: self.max_files]
+
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.input_dir.mkdir(parents=True, exist_ok=True)
+
+        total_files = len(entries)
+
+        logger.info("=" * 80)
+        logger.info(f"Streaming Arrow->WebDataset conversion: {self.split_name}")
+        logger.info(f"  Arrow files to process: {total_files}")
+        logger.info(f"  Output: {self.output_dir}")
+        logger.info(f"  Download workers: {self.max_workers}, " f"prefetch depth: {self.prefetch_depth}")
+        logger.info("=" * 80)
+
+        prefetch_q: queue.Queue[Union[Tuple[int, str, Path], None]] = queue.Queue()
+        download_semaphore = threading.Semaphore(self.prefetch_depth)
+        producer_error: List[BaseException] = []
+        cancel_event = threading.Event()
+
+        failed_files: List[Dict] = []
+        failed_lock = threading.Lock()
+
+        def _download_one(file_idx: int, md5: str, fname: str) -> Tuple[int, str, Path]:
+            file_url = f"{base_url}/{fname}"
+            # fname comes from a remote manifest; never let it traverse outside
+            # input_dir (e.g. "../../etc/passwd"). Use the basename and verify
+            # the resolved path stays under input_dir before writing.
+            safe_name = Path(fname).name
+            base_dir = self.input_dir.resolve()
+            arrow_path = (base_dir / safe_name).resolve()
+            if not safe_name or arrow_path.parent != base_dir:
+                raise ValueError(f"Unsafe manifest filename rejected: {fname!r}")
+            logger.info(f"  [download {file_idx + 1}/{total_files}] {fname}")
+            download_with_backoff(file_url, arrow_path, expected_md5=md5)
+            return file_idx, fname, arrow_path
+
+        skipped_count = [0]
+
+        def _producer() -> None:
+            """Submit downloads and drain results concurrently.
+
+            A separate drain thread processes completed futures in submission
+            order and feeds the prefetch queue.  The semaphore acquire in the
+            submit loop blocks when ``prefetch_depth`` files are already
+            in-flight or queued, bounding temporary disk usage.
+            """
+            try:
+                with ThreadPoolExecutor(max_workers=self.max_workers) as pool:
+                    futures_q: queue.Queue = queue.Queue()
+
+                    def _drain() -> None:
+                        while True:
+                            item = futures_q.get()
+                            if item is _SENTINEL:
+                                break
+                            fut, file_idx, fname = item
+                            if cancel_event.is_set():
+                                download_semaphore.release()
+                                continue
+                            try:
+                                result = fut.result()
+                                prefetch_q.put(result)
+                            except Exception as exc:
+                                download_semaphore.release()
+                                logger.warning(
+                                    f"  [FAILED download " f"{file_idx + 1}/{total_files}] " f"{fname}: {exc}"
+                                )
+                                with failed_lock:
+                                    failed_files.append(
+                                        {
+                                            "file_idx": file_idx,
+                                            "filename": fname,
+                                            "error": str(exc),
+                                            "stage": "download",
+                                        }
+                                    )
+
+                    drain_thread = threading.Thread(target=_drain, daemon=True)
+                    drain_thread.start()
+
+                    for file_idx, (md5, fname) in enumerate(entries):
+                        if cancel_event.is_set():
+                            break
+                        tar_path = self.output_dir / f"shard_{file_idx:06d}.tar"
+                        if tar_path.exists():
+                            skipped_count[0] += 1
+                            logger.info(f"  [skip {file_idx + 1}/{total_files}] " f"{fname} (shard exists)")
+                            continue
+                        download_semaphore.acquire()
+                        fut = pool.submit(_download_one, file_idx, md5, fname)
+                        futures_q.put((fut, file_idx, fname))
+
+                    futures_q.put(_SENTINEL)
+                    drain_thread.join()
+            except BaseException as exc:
+                producer_error.append(exc)
+            finally:
+                prefetch_q.put(_SENTINEL)
+
+        producer_thread = threading.Thread(target=_producer, daemon=True)
+        producer_thread.start()
+
+        global_sample_offset = 0
+        total_samples = 0
+        files_processed = 0
+
+        try:
+            while True:
+                item = prefetch_q.get()
+                if item is _SENTINEL:
+                    break
+
+                file_idx, fname, arrow_path = item
+
+                shard_name = f"shard_{file_idx:06d}.tar"
+                tar_path = self.output_dir / shard_name
+
+                try:
+                    num_samples = _arrow_to_tar(arrow_path, tar_path, global_sample_offset)
+                except Exception as exc:
+                    logger.warning(f"  [FAILED convert {file_idx + 1}/{total_files}] " f"{fname}: {exc}")
+                    tar_path.unlink(missing_ok=True)
+                    arrow_path.unlink(missing_ok=True)
+                    download_semaphore.release()
+                    with failed_lock:
+                        failed_files.append(
+                            {
+                                "file_idx": file_idx,
+                                "filename": fname,
+                                "error": str(exc),
+                                "stage": "conversion",
+                            }
+                        )
+                    continue
+
+                global_sample_offset += num_samples
+                total_samples += num_samples
+                files_processed += 1
+
+                arrow_path.unlink(missing_ok=True)
+                download_semaphore.release()
+
+                logger.info(
+                    f"[{files_processed}/{total_files}] {fname} -> {shard_name}: "
+                    f"{num_samples} samples (cumulative: {total_samples})"
+                )
+        except BaseException:
+            cancel_event.set()
+            raise
+        finally:
+            producer_thread.join(timeout=10)
+
+        if producer_error:
+            raise RuntimeError(f"Download failed: {producer_error[0]}") from producer_error[0]
+
+        if failed_files:
+            manifest_path = self.output_dir / "failed_files.json"
+            manifest_path.write_text(json.dumps(failed_files, indent=2))
+            logger.warning(f"{len(failed_files)} file(s) failed. " f"See {manifest_path}")
+
+        skipped = skipped_count[0]
+        failed = len(failed_files)
+        logger.info("=" * 80)
+        logger.info(
+            f"Done {self.split_name}: {files_processed} new + {skipped} skipped"
+            f"{f' + {failed} failed' if failed else ''}"
+            f" -> {files_processed + skipped} shards, "
+            f"{total_samples} new samples"
+        )
+        logger.info("=" * 80)
+
+        return {
+            "files_processed": files_processed,
+            "samples_written": total_samples,
+            "shards_created": files_processed + skipped,
+            "shards_skipped": skipped,
+            "files_failed": failed,
+        }
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py
new file mode 100644
index 000000000..7902c1bb1
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py
@@ -0,0 +1,234 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Raw dataset preparation pipeline for Megatron diffusion models.
+
+Creates raw Energon WebDatasets with images and captions for on-the-fly
+encoding during training.
+"""
+
+import logging
+from pathlib import Path
+from typing import Optional
+
+from ..utils import (
+    encode_image_to_bytes,
+    get_distributed_info,
+    preprocess_image,
+    save_to_webdataset,
+)
+
+logger = logging.getLogger(__name__)
+
+
+from .base import DatasetPipeline
+
+
+class RawDatasetPipeline(DatasetPipeline):
+    """
+    Pipeline for creating raw Energon WebDataset with images and captions.
+
+    This pipeline creates smaller datasets where encoding (VAE/T5/CLIP) happens
+    on-the-fly during training. Suitable for experimentation and when storage
+    is limited.
+
+    Args:
+        source_type: Type of input source ('directory', 'huggingface', 'webdataset')
+        output_dir: Output directory for Energon WebDataset
+        image_size: Target image size (default: 1024)
+        center_crop: Whether to center crop images (default: True)
+        image_format: Output image format (default: 'JPEG')
+        image_quality: JPEG/WEBP quality 1-100 (default: 95)
+        shard_size: Samples per shard (default: 1000)
+        max_samples: Maximum samples to process (default: None for all)
+        compress: Whether to compress tar files with gzip (default: False)
+        hf_token_file: Path to HuggingFace token file (default: None)
+        variable_size: If True, resize to nearest multiple of 16 instead of fixed size (default: False)
+        max_size: Maximum dimension when variable_size is True (default: 1024)
+
+    Example:
+        >>> pipeline = RawDatasetPipeline(
+        ...     source_type='huggingface',
+        ...     output_dir='/data/raw_pokemon',
+        ...     image_size=1024,
+        ...     center_crop=True,
+        ... )
+        >>> results = pipeline.run(
+        ...     hf_dataset='diffusers/pokemon-gpt4-captions',
+        ...     hf_split='train'
+        ... )
+        >>> print(f"Processed {results['samples_processed']} samples")
+    """
+
+    def __init__(
+        self,
+        source_type: str,
+        output_dir: str,
+        variable_size: bool = False,
+        image_size: int = 1024,
+        center_crop: bool = True,
+        max_size: int = 1024,
+        image_format: str = "JPEG",
+        image_quality: int = 95,
+        shard_size: int = 1000,
+        max_samples: Optional[int] = None,
+        compress: bool = False,
+        hf_token_file: Optional[str] = None,
+    ):
+        self.source_type = source_type
+        self.output_dir = Path(output_dir)
+        self.variable_size = variable_size
+        self.image_size = image_size
+        self.center_crop = center_crop
+        self.max_size = max_size
+        self.image_format = image_format
+        self.image_quality = image_quality
+        self.shard_size = shard_size
+        self.max_samples = max_samples
+        self.compress = compress
+
+        # Setup HF authentication if token file provided
+        if hf_token_file:
+            from ..auth import HFAuthError, setup_hf_authentication
+
+            try:
+                setup_hf_authentication(token_file=hf_token_file)
+            except HFAuthError as e:
+                raise ValueError(f"HuggingFace authentication failed: {e}")
+
+        # Get distributed info
+        self.rank, self.world_size = get_distributed_info()
+
+        # Determine output format extension
+        self.format_ext = {
+            "JPEG": "jpg",
+            "PNG": "png",
+            "WEBP": "webp",
+        }[self.image_format]
+
+        logger.info(f"Initialized RawDatasetPipeline (rank {self.rank}/{self.world_size})")
+        logger.info(f"Output directory: {self.output_dir}")
+        logger.info(f"Image preprocessing: size={self.image_size}, crop={self.center_crop}")
+
+    def process_sample(self, item):
+        """
+        Process a single sample.
+
+        Args:
+            item: Sample dict with 'image' (PIL Image) and 'caption' (str)
+
+        Returns:
+            Sample dict with image bytes and caption text
+        """
+        # Preprocess image
+        image = preprocess_image(
+            item["image"],
+            variable_size=self.variable_size,
+            size=self.image_size,
+            center_crop=self.center_crop,
+            max_size=self.max_size,
+        )
+
+        # Encode to bytes
+        image_bytes = encode_image_to_bytes(image, format=self.image_format, quality=self.image_quality)
+
+        # Create sample with standard keys
+        return {
+            self.format_ext: image_bytes,
+            "txt": item["caption"],
+        }
+
+    def run(self, **source_kwargs):
+        """
+        Execute the raw dataset preparation pipeline.
+
+        Args:
+            **source_kwargs: Source-specific arguments (passed to load_data)
+
+        Returns:
+            Dictionary with processing statistics:
+                - samples_processed: Number of samples successfully processed
+                - samples_skipped: Number of samples that failed
+                - shards_written: Number of output shards created
+        """
+        logger.info(f"Starting raw dataset preparation (rank {self.rank}/{self.world_size})")
+        logger.info(f"Source: {self.source_type}")
+
+        # Load data as a lazy iterator (never materialized in full)
+        data_iter = self.load_data(**source_kwargs)
+
+        if self.world_size > 1:
+            logger.info(
+                f"Distributed mode: rank {self.rank}/{self.world_size} "
+                f"processing every {self.world_size}th item (round-robin by load index)"
+            )
+
+        # Process and accumulate samples
+        samples = []
+        samples_processed = 0
+        samples_skipped = 0
+        shards_written = 0
+
+        # Stream items and assign them to ranks round-robin by global load
+        # index, so the full dataset is never held in memory. max_samples caps
+        # the GLOBAL index (total across ranks, not per-rank), preserving the
+        # prior semantics; the rank split changes from contiguous ranges to
+        # round-robin (same overall dataset, more balanced load). The index is
+        # assigned at load time, before process_sample, so the rank assignment
+        # stays deterministic even when samples are skipped.
+        for global_idx, item in enumerate(data_iter):
+            # Truthy check (not `is not None`) mirrors the original semantics
+            # where max_samples=0 / None meant "no limit".
+            if self.max_samples and global_idx >= self.max_samples:
+                break
+            if global_idx % self.world_size != self.rank:
+                continue
+            try:
+                sample = self.process_sample(item)
+                samples.append(sample)
+                samples_processed += 1
+
+                # Log progress
+                if samples_processed % 100 == 0:
+                    logger.info(f"Processed {samples_processed} samples (skipped {samples_skipped})")
+
+                # Save shard when full
+                if len(samples) >= self.shard_size:
+                    # Use distributed-aware shard naming to prevent conflicts
+                    shard_offset = shards_written * self.world_size + self.rank
+                    if self.world_size > 1 and shards_written == 0:
+                        logger.info(f"Rank {self.rank}: First shard will be numbered {shard_offset:06d}.tar")
+
+                    num_shards = save_to_webdataset(
+                        samples,
+                        str(self.output_dir),
+                        self.shard_size,
+                        shard_offset=shard_offset,
+                        compress=self.compress,
+                    )
+                    shards_written += num_shards
+                    samples = []
+
+            except Exception as e:
+                logger.warning(f"Failed to process sample {global_idx}: {e}")
+                samples_skipped += 1
+                continue
+
+        # Save remaining samples
+        if samples:
+            shard_offset = shards_written * self.world_size + self.rank
+            num_shards = save_to_webdataset(
+                samples,
+                str(self.output_dir),
+                self.shard_size,
+                shard_offset=shard_offset,
+                compress=self.compress,
+            )
+            shards_written += num_shards
+
+        return {
+            "samples_processed": samples_processed,
+            "samples_skipped": samples_skipped,
+            "shards_written": shards_written,
+        }
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/utils.py b/primus/backends/megatron/data/diffusion/preprocessing/utils.py
new file mode 100644
index 000000000..cd57d6836
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/utils.py
@@ -0,0 +1,546 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Common utilities for dataset preprocessing.
+
+This module provides shared functionality for creating WebDataset shards
+from various input sources (HuggingFace Hub, directories, existing WebDatasets).
+"""
+
+import io
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+import torch
+import webdataset as wds
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+def get_distributed_info() -> Tuple[int, int]:
+    """
+    Get distributed processing information (rank, world_size).
+
+    Returns:
+        Tuple of (rank, world_size). Returns (0, 1) if not in distributed mode.
+
+    Example:
+        >>> rank, world_size = get_distributed_info()
+        >>> if rank == 0:
+        ...     print("This is the master process")
+    """
+    try:
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            rank = torch.distributed.get_rank()
+            world_size = torch.distributed.get_world_size()
+            return rank, world_size
+    except (ImportError, AttributeError):
+        pass
+
+    # Fallback to environment variables
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    return rank, world_size
+
+
+def split_work_for_rank(total: int, rank: int, world_size: int) -> Tuple[int, int]:
+    """
+    Calculate start/end indices for this rank in distributed processing.
+
+    Args:
+        total: Total number of items to process
+        rank: Current process rank (0 to world_size-1)
+        world_size: Total number of processes
+
+    Returns:
+        Tuple of (start_idx, end_idx) for this rank
+
+    Example:
+        >>> start, end = split_work_for_rank(1000, 0, 4)
+        >>> print(f"Process 0 handles items {start} to {end}")
+        Process 0 handles items 0 to 250
+    """
+    items_per_rank = total // world_size
+    start_idx = rank * items_per_rank
+
+    if rank == world_size - 1:
+        # Last rank gets all remaining items (including remainder)
+        end_idx = total
+    else:
+        end_idx = start_idx + items_per_rank
+
+    return start_idx, end_idx
+
+
+def save_to_webdataset(
+    samples: List[Dict[str, Any]],
+    output_dir: str,
+    shard_size: int,
+    shard_offset: int = 0,
+    compress: bool = False,
+) -> int:
+    """
+    Save samples to WebDataset tar shards.
+
+    Args:
+        samples: List of sample dictionaries with keys as extensions
+            Example: {'images': image_bytes, 'txt': caption_text}
+        output_dir: Directory to save shards
+        shard_size: Number of samples per shard
+        shard_offset: Starting shard number (for distributed processing)
+        compress: Whether to compress tar files (gzip)
+
+    Returns:
+        Number of shards created
+
+    Example:
+        >>> samples = [
+        ...     {'images': image_bytes, 'txt': 'a cat'},
+        ...     {'images': image_bytes2, 'txt': 'a dog'},
+        ... ]
+        >>> num_shards = save_to_webdataset(samples, '/path/to/output', shard_size=1000)
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Calculate number of shards needed
+    num_shards = (len(samples) + shard_size - 1) // shard_size
+
+    for shard_idx in range(num_shards):
+        start_idx = shard_idx * shard_size
+        end_idx = min(start_idx + shard_size, len(samples))
+        shard_samples = samples[start_idx:end_idx]
+
+        # Create shard filename
+        shard_num = shard_offset + shard_idx
+        ext = ".tar.gz" if compress else ".tar"
+        shard_path = os.path.join(output_dir, f"{shard_num:06d}{ext}")
+
+        # Write shard
+        with wds.TarWriter(shard_path) as sink:
+            for sample_idx, sample in enumerate(shard_samples):
+                # Generate unique key for this sample
+                key = f"{shard_num:06d}_{sample_idx:06d}"
+
+                # Create WebDataset sample dict
+                wds_sample = {"__key__": key}
+
+                for ext, data in sample.items():
+                    if ext.startswith("__"):
+                        continue  # Skip special keys
+
+                    # Handle different data types
+                    if isinstance(data, bytes):
+                        wds_sample[ext] = data
+                    elif isinstance(data, str):
+                        wds_sample[ext] = data.encode("utf-8")
+                    elif isinstance(data, torch.Tensor):
+                        # Save tensor to bytes
+                        buffer = io.BytesIO()
+                        torch.save(data, buffer)
+                        wds_sample[ext] = buffer.getvalue()
+                    elif isinstance(data, Image.Image):
+                        # Save PIL image to bytes
+                        buffer = io.BytesIO()
+                        data.save(buffer, format="JPEG")
+                        wds_sample[ext] = buffer.getvalue()
+                    else:
+                        logger.warning(f"Unknown data type for key {ext}: {type(data)}")
+                        continue
+
+                sink.write(wds_sample)
+
+        logger.info(f"Created shard {shard_path} with {len(shard_samples)} samples")
+
+    return num_shards
+
+
+# Default keys to try when loading from HuggingFace if not specified in config
+DEFAULT_IMAGE_KEYS = ["image", "jpg", "jpeg", "png", "img", "photo"]
+DEFAULT_CAPTION_KEYS = ["caption", "text", "txt", "description", "prompt"]
+
+
+def _extract_field(item: dict, field_keys: List[str], idx: int, field_type: str) -> Any:
+    """
+    Extract field from item using list of possible keys.
+    Supports dot notation for JSON paths (e.g., 'json.caption').
+
+    Args:
+        item: Dataset sample dictionary
+        field_keys: List of keys to try (e.g., ['caption', 'json.caption'])
+        idx: Sample index (for logging)
+        field_type: 'image' or 'caption' (currently unused)
+
+    Returns:
+        Field value if found, None otherwise
+
+    Example:
+        >>> item = {'jpg': image_bytes, 'json': {'caption': 'a cat'}}
+        >>> _extract_field(item, ['caption', 'json.caption'], 0, 'caption')
+        'a cat'
+    """
+    for key in field_keys:
+        if "." in key:
+            # Handle JSON path (e.g., 'json.caption')
+            value = _extract_json_path(item, key, idx)
+            if value is not None:
+                return value
+        elif key in item:
+            return item[key]
+    return None
+
+
+def _extract_json_path(item: dict, path: str, idx: int) -> Any:
+    """
+    Extract value from nested JSON using dot notation.
+
+    Args:
+        item: Dataset sample dictionary
+        path: Dot-separated path (e.g., 'json.metadata.caption')
+        idx: Sample index (for logging)
+
+    Returns:
+        Extracted value if found, None otherwise
+
+    Example:
+        >>> item = {'json': b'{"caption": "a cat"}'}
+        >>> _extract_json_path(item, 'json.caption', 0)
+        'a cat'
+    """
+    import json as json_module
+
+    parts = path.split(".")
+    current = item
+
+    for i, part in enumerate(parts):
+        if part not in current:
+            return None
+
+        current = current[part]
+
+        # If we hit a JSON string/bytes at any level, parse it
+        if isinstance(current, (bytes, str)):
+            try:
+                if isinstance(current, bytes):
+                    current = json_module.loads(current.decode("utf-8"))
+                elif isinstance(current, str) and (current.startswith("{") or current.startswith("[")):
+                    current = json_module.loads(current)
+            except Exception as e:
+                logger.debug(f"Sample {idx} failed to parse JSON at '{part}': {e}")
+                return None
+
+    return current if isinstance(current, (str, int, float)) else None
+
+
+def load_from_huggingface(
+    dataset_name: str,
+    split: str = "train",
+    streaming: bool = True,
+    data_files: Optional[Union[str, List[str]]] = None,
+    image_key: Optional[str] = None,
+    caption_key: Optional[str] = None,
+    image_keys: Optional[List[str]] = None,
+    caption_keys: Optional[List[str]] = None,
+) -> Iterator[Dict[str, Any]]:
+    """
+    Load dataset from HuggingFace Hub with configurable field mappings.
+
+    Args:
+        dataset_name: HF dataset identifier (e.g., 'laion/laion400m')
+        split: Dataset split to load
+        streaming: Whether to stream the dataset (recommended for large datasets)
+        data_files: Specific files/paths to load from the dataset (e.g., 'data_1024_10K/*.tar')
+        image_key: Single image field name (e.g., 'jpg', 'image')
+        caption_key: Single caption field or JSON path (e.g., 'caption', 'json.caption')
+        image_keys: List of image field names to try (fallback to defaults if None)
+        caption_keys: List of caption fields/paths to try (fallback to defaults if None)
+
+    Yields:
+        Dicts with 'image' (PIL Image) and 'caption' (str) keys
+
+    Example:
+        >>> # Load with auto-detection (uses defaults)
+        >>> for item in load_from_huggingface('diffusers/pokemon-gpt4-captions'):
+        ...     print(f"Caption: {item['caption']}")
+
+        >>> # Load with specific field mappings
+        >>> for item in load_from_huggingface(
+        ...     'jackyhate/text-to-image-2M',
+        ...     image_key='jpg',
+        ...     caption_key='json.caption',
+        ...     data_files='data_1024_10K/*.tar'
+        ... ):
+        ...     print(f"Caption: {item['caption']}")
+    """
+    try:
+        from datasets import load_dataset
+    except ImportError:
+        raise ImportError(
+            "The 'datasets' package is required for HuggingFace dataset loading. "
+            "Install with: pip install datasets"
+        )
+
+    if data_files:
+        logger.info(f"Loading HuggingFace dataset: {dataset_name} (split: {split}, data_files: {data_files})")
+    else:
+        logger.info(f"Loading HuggingFace dataset: {dataset_name} (split: {split})")
+
+    dataset = load_dataset(dataset_name, split=split, streaming=streaming, data_files=data_files)
+
+    # Determine which keys to try
+    if image_keys is None:
+        image_keys = [image_key] if image_key else DEFAULT_IMAGE_KEYS
+    if caption_keys is None:
+        caption_keys = [caption_key] if caption_key else DEFAULT_CAPTION_KEYS
+
+    for idx, item in enumerate(dataset):
+        # Try to extract image using configured keys
+        image = _extract_field(item, image_keys, idx, "image")
+
+        # Try to extract caption using configured keys (supports JSON paths)
+        caption = _extract_field(item, caption_keys, idx, "caption")
+
+        if isinstance(caption, list):
+            caption = caption[0] if caption else None
+
+        if image is None or caption is None:
+            logger.warning(f"Sample {idx} missing image or caption, skipping. Keys: {item.keys()}")
+            continue
+
+        # Ensure image is PIL Image
+        if not isinstance(image, Image.Image):
+            try:
+                if isinstance(image, bytes):
+                    image = Image.open(io.BytesIO(image))
+                else:
+                    logger.warning(f"Sample {idx} has unexpected image type: {type(image)}")
+                    continue
+            except Exception as e:
+                logger.warning(f"Sample {idx} failed to load image: {e}")
+                continue
+
+        yield {
+            "image": image,
+            "caption": caption,
+        }
+
+
+def load_from_directory(input_dir: str) -> Iterator[Dict[str, Any]]:
+    """
+    Load dataset from directory structure.
+
+    Expected structure:
+        input_dir/
+            images/
+                0000000.jpg
+                0000001.jpg
+            captions/
+                0000000.txt
+                0000001.txt
+
+    Args:
+        input_dir: Root directory containing 'images' and 'captions' subdirectories
+
+    Yields:
+        Dicts with 'image' (PIL Image) and 'caption' (str) keys
+
+    Example:
+        >>> for item in load_from_directory('/path/to/dataset'):
+        ...     print(f"Caption: {item['caption']}")
+    """
+    input_path = Path(input_dir)
+    images_dir = input_path / "images"
+    captions_dir = input_path / "captions"
+
+    if not images_dir.exists():
+        raise ValueError(f"Images directory not found: {images_dir}")
+
+    if not captions_dir.exists():
+        raise ValueError(f"Captions directory not found: {captions_dir}")
+
+    logger.info(f"Loading dataset from directory: {input_dir}")
+
+    # Find all image files
+    image_extensions = [".jpg", ".jpeg", ".png", ".webp"]
+    image_files = []
+    for ext in image_extensions:
+        image_files.extend(sorted(images_dir.glob(f"*{ext}")))
+
+    logger.info(f"Found {len(image_files)} image files")
+
+    for img_path in image_files:
+        # Find corresponding caption file
+        caption_path = captions_dir / f"{img_path.stem}.txt"
+
+        if not caption_path.exists():
+            logger.warning(f"Caption not found for {img_path.name}, skipping")
+            continue
+
+        try:
+            # Load image
+            image = Image.open(img_path).convert("RGB")
+
+            # Load caption
+            with open(caption_path, "r", encoding="utf-8") as f:
+                caption = f.read().strip()
+
+            yield {
+                "image": image,
+                "caption": caption,
+            }
+
+        except Exception as e:
+            logger.warning(f"Failed to load {img_path.name}: {e}")
+            continue
+
+
+def load_from_webdataset(input_path: str) -> Iterator[Dict[str, Any]]:
+    """
+    Load existing WebDataset with non-standard keys.
+
+    Converts to standard format {'image': PIL Image, 'caption': str}
+
+    Args:
+        input_path: Path or glob pattern to WebDataset tar files
+            Examples: '/path/to/shards/*.tar', '/path/to/shard_000000.tar'
+
+    Yields:
+        Dicts with 'image' (PIL Image) and 'caption' (str) keys
+
+    Example:
+        >>> for item in load_from_webdataset('/path/to/dataset/*.tar'):
+        ...     print(f"Caption: {item['caption']}")
+    """
+    logger.info(f"Loading WebDataset from: {input_path}")
+
+    dataset = wds.WebDataset(input_path)
+
+    for sample in dataset:
+        try:
+            # Try to find image with various keys
+            image = None
+            for img_key in ["jpg", "png", "jpeg", "webp", "image"]:
+                if img_key in sample:
+                    img_data = sample[img_key]
+                    if isinstance(img_data, bytes):
+                        image = Image.open(io.BytesIO(img_data)).convert("RGB")
+                    elif isinstance(img_data, Image.Image):
+                        image = img_data.convert("RGB")
+                    else:
+                        continue
+                    break
+
+            # Try to find caption with various keys
+            caption = None
+            for cap_key in ["txt", "caption.txt", "text", "caption"]:
+                if cap_key in sample:
+                    cap_data = sample[cap_key]
+                    if isinstance(cap_data, bytes):
+                        caption = cap_data.decode("utf-8")
+                    elif isinstance(cap_data, str):
+                        caption = cap_data
+                    else:
+                        continue
+                    break
+
+            if image is None or caption is None:
+                logger.warning(f"Sample missing image or caption, skipping. Keys: {sample.keys()}")
+                continue
+
+            yield {
+                "image": image,
+                "caption": caption,
+            }
+
+        except Exception as e:
+            logger.warning(f"Failed to load sample: {e}")
+            continue
+
+
+def preprocess_image(
+    image: Image.Image,
+    variable_size: bool = True,  # if True, then image is resized to the nearest multiple of 16, otherwise it is resized to the given size
+    size: int = 1024,  # only used if variable_size is False, then this is applied
+    center_crop: bool = False,  # only used if variable_size is False, then this is applied
+    max_size: int = 1024,  # maximum dimension when variable_size is True
+) -> Image.Image:
+    """
+    Preprocess image (resize, crop).
+
+    Args:
+        image: PIL Image
+        variable_size: If True, resize to nearest multiple of 16 up to max_size
+        size: Target size (square) - only used if variable_size is False
+        center_crop: Whether to center crop before resize - only used if variable_size is False
+        max_size: Maximum dimension when variable_size is True (default: 1024)
+
+    Returns:
+        Preprocessed PIL Image
+
+    Example:
+        >>> image = Image.open('photo.jpg')
+        >>> processed = preprocess_image(image, variable_size=True, max_size=2048)
+    """
+    # Center crop if requested
+    if center_crop and variable_size is False:
+        width, height = image.size
+        crop_size = min(width, height)
+        left = (width - crop_size) // 2
+        top = (height - crop_size) // 2
+        right = left + crop_size
+        bottom = top + crop_size
+        image = image.crop((left, top, right, bottom))
+
+    # Resize
+    if variable_size is False:
+        image = image.resize((size, size), Image.Resampling.LANCZOS)
+    else:
+        width, height = image.size
+        if max(width, height) > max_size:
+            scale = max_size / float(max(width, height))
+            new_w = int(round(width * scale))
+            new_h = int(round(height * scale))
+            # High-quality downsampling
+            image = image.resize((new_w, new_h), resample=Image.Resampling.LANCZOS)
+
+        width, height = image.size
+        new_size = (round(width / 16) * 16, round(height / 16) * 16)
+        image = image.resize(new_size, Image.LANCZOS)
+
+    return image
+
+
+def encode_image_to_bytes(image: Image.Image, format: str = "JPEG", quality: int = 95) -> bytes:
+    """
+    Encode PIL Image to bytes.
+
+    Args:
+        image: PIL Image
+        format: Image format ('JPEG', 'PNG', 'WEBP')
+        quality: JPEG/WEBP quality (1-100)
+
+    Returns:
+        Image bytes
+
+    Example:
+        >>> image = Image.open('photo.jpg')
+        >>> image_bytes = encode_image_to_bytes(image, format='JPEG', quality=95)
+    """
+    buffer = io.BytesIO()
+    image.save(buffer, format=format, quality=quality)
+    return buffer.getvalue()
+
+
+__all__ = [
+    "get_distributed_info",
+    "split_work_for_rank",
+    "save_to_webdataset",
+    "load_from_huggingface",
+    "load_from_directory",
+    "load_from_webdataset",
+    "preprocess_image",
+    "encode_image_to_bytes",
+]
diff --git a/primus/backends/megatron/data/diffusion/preprocessing/validate.py b/primus/backends/megatron/data/diffusion/preprocessing/validate.py
new file mode 100644
index 000000000..769cea725
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/preprocessing/validate.py
@@ -0,0 +1,302 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Energon dataset validation for Primus diffusion pipelines.
+
+Validates that a prepared Energon WebDataset is structurally correct and
+loadable by the training code. Replaces the broken ``energon info`` CLI
+tool which does not support CrudeWebdataset format.
+
+Standalone usage:
+    python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset
+
+Programmatic usage:
+    from primus.backends.megatron.data.diffusion.preprocessing.validate import (
+        validate_energon_dataset,
+    )
+    validate_energon_dataset('/path/to/dataset', encoding='preencoded')
+"""
+
+import json
+import logging
+import tarfile
+from pathlib import Path
+from typing import Dict, List, Literal, Optional
+
+import torch
+import yaml
+
+logger = logging.getLogger(__name__)
+
+EXPECTED_KEYS_PREENCODED = {"latents.pth", "prompt_embeds.pth", "pooled_prompt_embeds.pth"}
+EXPECTED_KEYS_PREENCODED_NUMPY = {"t5.bytes", "clip.bytes", "mean.bytes", "logvar.bytes"}
+EXPECTED_KEYS_RAW_IMAGE = {"jpg", "jpeg", "png", "webp"}
+
+
+def _check_metadata(output_path: Path) -> Optional[Dict]:
+    """Check that .nv-meta files exist and are valid. Returns info dict or None."""
+    meta_dir = output_path / ".nv-meta"
+
+    info_path = meta_dir / ".info.json"
+    if not info_path.exists():
+        info_path_yaml = meta_dir / ".info.yaml"
+        if info_path_yaml.exists():
+            info_path = info_path_yaml
+        else:
+            logger.error(f"Missing metadata: neither .info.json nor .info.yaml in {meta_dir}")
+            return None
+
+    try:
+        with open(info_path) as f:
+            if info_path.suffix == ".json":
+                info = json.load(f)
+            else:
+                info = yaml.safe_load(f)
+    except Exception as e:
+        logger.error(f"Failed to parse {info_path}: {e}")
+        return None
+
+    if "shard_counts" not in info:
+        logger.error(f"Missing 'shard_counts' in {info_path}")
+        return None
+
+    split_path = meta_dir / "split.yaml"
+    if not split_path.exists():
+        split_path = meta_dir / "split.json"
+    if not split_path.exists():
+        logger.error(f"Missing split config in {meta_dir}")
+        return None
+
+    try:
+        with open(split_path) as f:
+            splits = yaml.safe_load(f) if split_path.suffix == ".yaml" else json.load(f)
+    except Exception as e:
+        logger.error(f"Failed to parse {split_path}: {e}")
+        return None
+
+    train_parts = splits.get("split_parts", {}).get("train", [])
+    if not train_parts:
+        logger.warning("Train split is empty")
+
+    dataset_yaml = meta_dir / "dataset.yaml"
+    if not dataset_yaml.exists():
+        logger.error(f"Missing {dataset_yaml}")
+        return None
+
+    missing_files: List[str] = []
+    for shard_name in info["shard_counts"]:
+        if not (output_path / shard_name).exists():
+            missing_files.append(shard_name)
+        idx_name = shard_name + ".idx"
+        if not (output_path / idx_name).exists():
+            missing_files.append(idx_name)
+
+    if missing_files:
+        logger.error(f"Missing files on disk: {missing_files[:10]}")
+        return None
+
+    info["_splits"] = splits
+    return info
+
+
+def _check_sample_counts(output_path: Path, info: Dict) -> bool:
+    """Spot-check that the first shard's entry count matches .info.json."""
+    shard_counts = info["shard_counts"]
+    first_shard = next(iter(shard_counts))
+    expected = shard_counts[first_shard]
+
+    tar_path = output_path / first_shard
+    try:
+        with tarfile.open(str(tar_path), "r") as tar:
+            members = tar.getmembers()
+            sample_keys = set()
+            for m in members:
+                key = m.name.split(".", 1)[0]
+                sample_keys.add(key)
+            actual = len(sample_keys)
+    except Exception as e:
+        logger.error(f"Failed to read {tar_path}: {e}")
+        return False
+
+    if actual != expected:
+        logger.error(
+            f"Sample count mismatch in {first_shard}: " f".info.json says {expected}, tar contains {actual}"
+        )
+        return False
+
+    return True
+
+
+def _check_sample_load(
+    output_path: Path,
+    encoding: str,
+) -> Optional[Dict[str, str]]:
+    """
+    Load one sample through energon's Python API (same path as training).
+    Returns a dict of {key: description} for the sample, or None on failure.
+    """
+    try:
+        from megatron.energon import StandardWebdatasetFactory
+        from megatron.energon.dataset_config import load_config
+        from megatron.energon.epathlib import EPath
+        from megatron.energon.flavors.webdataset.config import MAIN_FOLDER_NAME
+        from megatron.energon.loader import get_loader
+        from megatron.energon.worker import WorkerConfig
+    except ImportError as e:
+        logger.warning(f"Cannot import megatron.energon for sample validation: {e}")
+        return None
+
+    try:
+        ds_path = EPath(str(output_path))
+        worker_config = WorkerConfig(rank=0, world_size=1, num_workers=0)
+        dataset = load_config(
+            ds_path / MAIN_FOLDER_NAME / "dataset.yaml",
+            default_kwargs=dict(
+                path=ds_path,
+                split_part="train",
+                training=False,
+                worker_config=worker_config,
+            ),
+            default_type=StandardWebdatasetFactory,
+        )
+        sample = next(iter(get_loader(dataset.build())))
+    except Exception as e:
+        logger.error(f"Failed to load sample via energon API: {e}")
+        return None
+
+    data_keys = [k for k in sample.keys() if not k.startswith("_")]
+
+    if encoding == "preencoded":
+        missing = EXPECTED_KEYS_PREENCODED - set(data_keys)
+        if missing:
+            logger.error(f"Sample missing expected keys for preencoded data: {missing}")
+            return None
+    elif encoding == "preencoded_numpy":
+        missing = EXPECTED_KEYS_PREENCODED_NUMPY - set(data_keys)
+        if missing:
+            logger.error(f"Sample missing expected keys for preencoded_numpy data: {missing}")
+            return None
+    elif encoding == "raw":
+        has_image = any(k in EXPECTED_KEYS_RAW_IMAGE for k in data_keys)
+        has_text = "txt" in data_keys
+        if not has_image or not has_text:
+            logger.error(f"Raw sample missing image or text key. Got: {data_keys}")
+            return None
+
+    descriptions = {}
+    for key in sorted(data_keys):
+        val = sample[key]
+        if isinstance(val, torch.Tensor):
+            descriptions[key] = f"shape={tuple(val.shape)}, dtype={val.dtype}"
+        elif isinstance(val, bytes):
+            descriptions[key] = f"bytes, len={len(val)}"
+        else:
+            descriptions[key] = f"{type(val).__name__}"
+
+    return descriptions
+
+
+def validate_energon_dataset(
+    output_dir: str,
+    encoding: Literal["preencoded", "preencoded_numpy", "raw"] = "preencoded",
+) -> bool:
+    """
+    Validate a prepared Energon WebDataset.
+
+    Performs four checks:
+      1. Metadata files (.info.json, split.yaml, dataset.yaml) are present and valid
+      2. Sample count in .info.json matches actual tar contents (spot-check)
+      3. One sample loads successfully through energon's Python API
+      4. Prints a structured summary of the dataset
+
+    Args:
+        output_dir: Path to the dataset directory
+        encoding: Expected encoding ('preencoded', 'preencoded_numpy', or 'raw')
+
+    Returns:
+        True if all checks pass, False otherwise.
+    """
+    output_path = Path(output_dir)
+    passed = True
+
+    logger.info("Validating dataset...")
+
+    # --- Check 1: metadata ---
+    info = _check_metadata(output_path)
+    if info is None:
+        logger.error("  Metadata check FAILED")
+        return False
+    logger.info("  Metadata check passed")
+
+    # --- Check 2: sample counts ---
+    if not _check_sample_counts(output_path, info):
+        logger.error("  Sample count check FAILED")
+        passed = False
+    else:
+        logger.info("  Sample count check passed")
+
+    # --- Check 3: load one sample via energon API ---
+    descriptions = _check_sample_load(output_path, encoding)
+    if descriptions is None:
+        logger.error("  Sample load check FAILED")
+        passed = False
+    else:
+        logger.info("  Sample load check passed")
+
+    # --- Check 4: summary ---
+    shard_counts = info["shard_counts"]
+    total_samples = sum(shard_counts.values())
+    num_shards = len(shard_counts)
+
+    splits = info.get("_splits", {}).get("split_parts", {})
+    split_summary = ", ".join(
+        (
+            f"{name}={len(shards)}"
+            if isinstance(shards, list) and not any("{" in s for s in shards)
+            else f"{name}={'non-empty' if shards else 'empty'}"
+        )
+        for name, shards in splits.items()
+    )
+
+    total_bytes = sum(f.stat().st_size for f in output_path.glob("**/*.tar"))
+    size_gb = total_bytes / (1024**3)
+
+    logger.info("  " + "-" * 40)
+    logger.info(f"  Encoding:       {encoding}")
+    logger.info(f"  Total samples:  {total_samples} across {num_shards} shard(s)")
+    logger.info(f"  Splits:         {split_summary}")
+    if descriptions:
+        logger.info("  Spot check:")
+        for key, desc in descriptions.items():
+            logger.info(f"    {key}: {desc}")
+    logger.info(f"  Dataset size:   {size_gb:.2f} GB")
+
+    if passed:
+        logger.info("✓ Dataset verified and ready for training")
+    else:
+        logger.warning("Dataset has issues — see errors above")
+
+    return passed
+
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    parser = argparse.ArgumentParser(
+        description="Validate a Primus Energon WebDataset",
+    )
+    parser.add_argument("dataset_path", help="Path to the dataset directory")
+    parser.add_argument(
+        "--encoding",
+        choices=["preencoded", "preencoded_numpy", "raw"],
+        default="preencoded",
+        help="Expected encoding type (default: preencoded)",
+    )
+    args = parser.parse_args()
+
+    ok = validate_energon_dataset(args.dataset_path, encoding=args.encoding)
+    sys.exit(0 if ok else 1)
diff --git a/primus/backends/megatron/data/diffusion/task_encoders/__init__.py b/primus/backends/megatron/data/diffusion/task_encoders/__init__.py
new file mode 100644
index 000000000..a3ae2036f
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/task_encoders/__init__.py
@@ -0,0 +1,22 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+TaskEncoders for diffusion models.
+"""
+
+from .image import (
+    DiffusionSample,
+    EncodedDiffusionTaskEncoder,
+    RawDiffusionTaskEncoder,
+    cook_preencoded_diffusion,
+    cook_raw_images,
+)
+
+__all__ = [
+    "DiffusionSample",
+    "EncodedDiffusionTaskEncoder",
+    "RawDiffusionTaskEncoder",
+    "cook_preencoded_diffusion",
+    "cook_raw_images",
+]
diff --git a/primus/backends/megatron/data/diffusion/task_encoders/image.py b/primus/backends/megatron/data/diffusion/task_encoders/image.py
new file mode 100644
index 000000000..837bbcfdb
--- /dev/null
+++ b/primus/backends/megatron/data/diffusion/task_encoders/image.py
@@ -0,0 +1,396 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Diffusion TaskEncoder for image-text pairs using Crude Data pattern.
+
+This module provides TaskEncoder implementations for diffusion models:
+- EncodedDiffusionTaskEncoder: Loads pre-encoded data (latents, embeddings)
+- RawDiffusionTaskEncoder: Loads raw images and text (no encoding)
+
+Encoding happens in the model, not in the TaskEncoder (following best practices).
+Position IDs are generated in the model code, not here.
+"""
+
+import io
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+from megatron.energon import (
+    Cooker,
+    DefaultTaskEncoder,
+    Sample,
+    SampleDecoder,
+    WorkerConfig,
+    basic_sample_keys,
+    stateless,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Sample Definition (with proper Sample inheritance)
+# ============================================================================
+
+
+@dataclass
+class DiffusionSample(Sample):
+    """
+    Diffusion training sample with framework-standard field names.
+
+    Inherits from megatron.energon.Sample to ensure __key__, __restore_key__,
+    and __subflavors__ are properly tracked for deterministic training resumption.
+
+    Position IDs (img_ids, text_ids) are generated at runtime based on actual
+    tensor shapes, not stored in the dataset. This provides flexibility for
+    variable-resolution training.
+
+    Attributes:
+        prompt_embeds: T5 text embeddings (seq_len, hidden_dim)
+        pooled_prompt_embeds: CLIP pooled embeddings (hidden_dim,)
+        latents: Image latents from VAE (C, H, W) — optional for resample-only datasets
+        mean: VAE posterior mean (C, H, W) — present in 'resample' mode datasets
+        logvar: VAE posterior log-variance (C, H, W) — present in 'resample' mode datasets
+        caption: Original text caption (optional, for debugging)
+    """
+
+    prompt_embeds: torch.Tensor
+    pooled_prompt_embeds: torch.Tensor
+    latents: Optional[torch.Tensor] = None
+    mean: Optional[torch.Tensor] = None
+    logvar: Optional[torch.Tensor] = None
+    caption: str = ""
+    timestep: Optional[torch.Tensor] = None
+
+
+# ============================================================================
+# Cooker Functions
+# ============================================================================
+
+
+@stateless
+def cook_preencoded_diffusion(sample: dict) -> DiffusionSample:
+    """
+    Cooker for preencoded diffusion features with framework-standard keys.
+
+    Loads precalculated VAE latents and text embeddings from disk.
+    Position IDs are NOT loaded - they are generated at runtime based on
+    actual tensor shapes for flexibility.
+
+    Required standard keys:
+        - 'latents.pth': VAE-encoded image latents
+        - 'prompt_embeds.pth': T5 text embeddings
+        - 'pooled_prompt_embeds.pth': CLIP pooled embeddings
+
+    Args:
+        sample: Raw sample dict from WebDataset
+
+    Returns:
+        DiffusionSample with all metadata properly forwarded
+    """
+
+    def load_tensor(data):
+        """Helper to load tensor from bytes."""
+        if data is None:
+            return None
+        if isinstance(data, (str, Path)):
+            return torch.load(data, map_location="cpu")
+        elif isinstance(data, bytes):
+            return torch.load(io.BytesIO(data), map_location="cpu")
+        elif isinstance(data, torch.Tensor):
+            return data
+        else:
+            return data
+
+    # Load required fields (position IDs not loaded)
+    latents = load_tensor(sample.get("latents.pth"))
+    prompt_embeds = load_tensor(sample.get("prompt_embeds.pth"))
+    pooled_prompt_embeds = load_tensor(sample.get("pooled_prompt_embeds.pth"))
+
+    # Load optional resample-mode fields
+    mean = load_tensor(sample.get("mean.pth"))
+    logvar = load_tensor(sample.get("logvar.pth"))
+
+    if prompt_embeds is None or pooled_prompt_embeds is None:
+        raise ValueError(
+            f"Sample missing required keys. Expected: 'prompt_embeds.pth', "
+            f"'pooled_prompt_embeds.pth'. Got: {list(sample.keys())}"
+        )
+    if latents is None and mean is None:
+        raise ValueError(
+            f"Sample must have 'latents.pth' or 'mean.pth'/'logvar.pth'. " f"Got: {list(sample.keys())}"
+        )
+
+    # Load caption if available
+    caption = sample.get("caption.txt", b"")
+    if isinstance(caption, bytes):
+        caption = caption.decode("utf-8")
+    elif isinstance(caption, (str, Path)):
+        try:
+            if Path(caption).exists():
+                with open(caption, "r") as f:
+                    caption = f.read().strip()
+            else:
+                caption = str(caption)
+        except (OSError, ValueError):
+            caption = str(caption)
+    else:
+        caption = str(caption) if caption else ""
+
+    # Extract timestep from JSON sidecar (MLPerf validation datasets)
+    timestep = None
+    if "json" in sample:
+        import json as json_mod
+
+        raw = sample["json"]
+        if isinstance(raw, dict):
+            decoded = raw
+        elif isinstance(raw, bytes):
+            decoded = json_mod.loads(raw.decode("utf-8"))
+        else:
+            decoded = json_mod.loads(raw)
+        if "timestep" in decoded:
+            timestep = torch.tensor(decoded["timestep"])
+
+    return DiffusionSample(
+        **basic_sample_keys(sample),
+        latents=latents,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        mean=mean,
+        logvar=logvar,
+        caption=caption,
+        timestep=timestep,
+    )
+
+
+@stateless
+def cook_raw_images(sample: dict) -> Dict[str, Any]:
+    """
+    Cooker for raw images - just loads data, NO ENCODING.
+
+    Encoding happens in the model's forward_step, not here.
+    This follows the pattern where TaskEncoders only load data.
+
+    Standard data keys:
+        - 'images': Raw image data (or format-specific key like 'jpg', 'png', 'webp')
+        - 'txt': Text caption
+
+    Args:
+        sample: Raw sample dict from WebDataset
+
+    Returns:
+        Dict with raw data ready for model encoding
+    """
+    return {
+        **basic_sample_keys(sample),
+        "images": sample.get("images"),
+        "txt": sample.get("txt"),
+    }
+
+
+def load_numpy_tensor(data: Optional[bytes]) -> Optional[torch.Tensor]:
+    """Load a bfloat16 tensor from numpy uint16 bytes (MLPerf format).
+
+    The MLPerf Flux dataset stores bfloat16 tensors by reinterpreting them
+    as uint16 numpy arrays serialized as .npy files. This function loads
+    the .npy buffer (preserving shape), then reinterprets uint16 as
+    bfloat16 via torch.
+
+    Falls back to raw ``np.frombuffer`` for headerless byte buffers.
+    """
+    if data is None:
+        return None
+    if data[:6] == b"\x93NUMPY":
+        arr = np.load(io.BytesIO(data))
+    else:
+        arr = np.frombuffer(data, dtype=np.uint16)
+    return torch.from_numpy(arr.copy()).view(torch.bfloat16)
+
+
+@stateless
+def cook_preencoded_numpy_diffusion(sample: dict) -> DiffusionSample:
+    """Cooker for MLPerf pre-encoded numpy data (bfloat16 as uint16 bytes).
+
+    Expected keys in the WebDataset tar shard:
+        - 't5.bytes':    T5 text embeddings
+        - 'clip.bytes':  CLIP pooled embeddings
+        - 'mean.bytes':  VAE posterior mean
+        - 'logvar.bytes': VAE posterior log-variance
+
+    No 'latents' are stored; the training loop resamples from mean/logvar.
+    """
+    prompt_embeds = load_numpy_tensor(sample.get("t5.bytes"))
+    pooled_prompt_embeds = load_numpy_tensor(sample.get("clip.bytes"))
+    mean = load_numpy_tensor(sample.get("mean.bytes"))
+    logvar = load_numpy_tensor(sample.get("logvar.bytes"))
+
+    if prompt_embeds is None or pooled_prompt_embeds is None:
+        raise ValueError(
+            f"Sample missing required keys. Expected: 't5.bytes', "
+            f"'clip.bytes'. Got: {list(sample.keys())}"
+        )
+    if mean is None:
+        raise ValueError(f"Sample missing 'mean.bytes'. Got: {list(sample.keys())}")
+
+    timestep = None
+    if "json" in sample:
+        import json as json_mod
+
+        raw = sample["json"]
+        if isinstance(raw, dict):
+            decoded = raw
+        elif isinstance(raw, bytes):
+            decoded = json_mod.loads(raw.decode("utf-8"))
+        else:
+            decoded = json_mod.loads(raw)
+        if "timestep" in decoded:
+            timestep = torch.tensor(decoded["timestep"])
+
+    return DiffusionSample(
+        **basic_sample_keys(sample),
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        mean=mean,
+        logvar=logvar,
+        timestep=timestep,
+    )
+
+
+# ============================================================================
+# TaskEncoders
+# ============================================================================
+
+
+class EncodedDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]):
+    """
+    TaskEncoder for PRE-ENCODED diffusion data.
+
+    Use this when your dataset contains pre-encoded features:
+    - latents.pth (VAE-encoded image latents)
+    - prompt_embeds.pth (T5 text embeddings)
+    - pooled_prompt_embeds.pth (CLIP pooled embeddings)
+
+    Does NOT do any encoding - just loads from disk.
+    For raw data, use RawDiffusionTaskEncoder instead.
+
+    Outputs batch with standard keys:
+    - 'latents'
+    - 'prompt_embeds'
+    - 'pooled_prompt_embeds'
+
+    Use with dataset.yaml:
+        ```yaml
+        subflavors:
+          encoding: preencoded
+        ```
+    """
+
+    decoder = SampleDecoder(image_decode="pil")
+
+    cookers = [
+        Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}),
+        Cooker(cook_preencoded_numpy_diffusion, has_subflavors={"encoding": "preencoded_numpy"}),
+    ]
+
+    def __init__(self, worker_config: Optional[WorkerConfig] = None):
+        """Initialize pre-encoded TaskEncoder."""
+        super().__init__()
+        self.worker_config = worker_config
+        logger.info("Initialized EncodedDiffusionTaskEncoder (preencoded / preencoded_numpy modes)")
+
+    def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]:
+        """
+        Batch pre-encoded samples.
+
+        Position IDs are NOT included - they are generated at runtime
+        in the forward step based on actual tensor shapes.
+
+        Returns:
+            Dict with keys: prompt_embeds, pooled_prompt_embeds,
+            and conditionally latents and/or mean, logvar depending on
+            which fields are present in the samples.
+        """
+        batch: Dict[str, torch.Tensor] = {
+            "prompt_embeds": torch.stack([s.prompt_embeds for s in samples]),
+            "pooled_prompt_embeds": torch.stack([s.pooled_prompt_embeds for s in samples]),
+        }
+
+        if samples[0].latents is not None:
+            batch["latents"] = torch.stack([s.latents for s in samples])
+
+        if samples[0].mean is not None:
+            batch["mean"] = torch.stack([s.mean for s in samples])
+            batch["logvar"] = torch.stack([s.logvar for s in samples])
+
+        if samples[0].timestep is not None:
+            batch["timestep"] = torch.stack([s.timestep for s in samples])
+
+        return batch
+
+
+class RawDiffusionTaskEncoder(DefaultTaskEncoder):
+    """
+    TaskEncoder for RAW diffusion data (images and text).
+
+    Use this when your dataset contains raw files:
+    - images (raw image files)
+    - txt (text captions)
+
+    This TaskEncoder:
+    - Loads raw images and captions from disk
+    - Does NOT do any encoding (no VAE, no T5, no CLIP)
+    - Encoding happens on-the-fly in model's forward_step
+
+    Outputs batch with standard keys:
+    - 'images': List of PIL Images or image bytes
+    - 'txt': List of caption strings
+
+    Use with dataset.yaml:
+        ```yaml
+        subflavors:
+          encoding: raw
+        ```
+    """
+
+    decoder = SampleDecoder(image_decode="pil")
+
+    cookers = [
+        Cooker(cook_raw_images, has_subflavors={"encoding": "raw"}),
+    ]
+
+    def __init__(self, worker_config: Optional[WorkerConfig] = None):
+        """Initialize raw diffusion TaskEncoder."""
+        super().__init__()
+        self.worker_config = worker_config
+        logger.info("Initialized RawDiffusionTaskEncoder (no encoding, passes raw data)")
+
+    def batch(self, samples: List[Dict]) -> Dict[str, Any]:
+        """
+        Batch raw samples.
+
+        Returns:
+            Dict with standard keys:
+            - 'images': List of PIL Images or image bytes
+            - 'txt': List of caption strings
+        """
+        return {
+            "images": [s["images"] for s in samples],
+            "txt": [s["txt"] for s in samples],
+        }
+
+
+__all__ = [
+    "DiffusionSample",
+    "EncodedDiffusionTaskEncoder",
+    "RawDiffusionTaskEncoder",
+    "cook_preencoded_diffusion",
+    "cook_preencoded_numpy_diffusion",
+    "cook_raw_images",
+    "load_numpy_tensor",
+]
diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py
new file mode 100644
index 000000000..4ac1728dc
--- /dev/null
+++ b/primus/backends/megatron/data/energon_dataset_provider.py
@@ -0,0 +1,282 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+EnergonDatasetProvider: Megatron Energon dataset provider.
+
+Provides Energon-based dataloaders for multimodal and diffusion models,
+using WebDataset format with distributed loading via WorkerConfig.
+
+Reference:
+    Megatron-LM examples/multimodal/dataloader_provider.py
+"""
+
+from typing import Any, Callable, List, Optional, Tuple
+
+from megatron.core import parallel_state
+from megatron.core.num_microbatches_calculator import get_num_microbatches
+from megatron.core.parallel_state import (
+    get_pipeline_model_parallel_rank,
+    get_pipeline_model_parallel_world_size,
+    get_tensor_model_parallel_rank,
+)
+from megatron.energon import (
+    LimitDataset,
+    RepeatDataset,
+    WorkerConfig,
+    get_loader,
+    get_savable_loader,
+    get_train_dataset,
+    get_val_datasets,
+)
+
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+from primus.backends.megatron.data.dataset_provider import DatasetProvider
+from primus.core.utils.module_utils import log_rank_0
+
+
+class EnergonDatasetProvider(DatasetProvider):
+    """
+    Dataset provider for multimodal/diffusion models using Megatron Energon.
+
+    Uses:
+        - Megatron Energon (WebDataset format)
+        - Task encoders (model-specific, injected via factory)
+        - WorkerConfig for distributed data loading
+
+    Architecture:
+        - Task encoder factory is called during setup() to create encoder
+        - WorkerConfig handles rank-based data sharding automatically
+        - Dataloaders are only created on specific ranks (first TP, first/last PP)
+
+    Reference:
+        Megatron-LM examples/multimodal/dataloader_provider.py
+    """
+
+    def __init__(self, task_encoder_factory: Callable[[], Any]):
+        """
+        Initialize Energon dataset provider.
+
+        Args:
+            task_encoder_factory: Function that creates task encoder
+                Signature: () -> TaskEncoder
+                Example: lambda: EncodedDiffusionTaskEncoder(...)
+
+                The factory pattern is used because task encoders may need
+                access to trainer state (e.g., self.module_config) which
+                isn't available during __init__.
+        """
+        self.task_encoder_factory = task_encoder_factory
+
+    def create_dataloaders(
+        self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None
+    ) -> Tuple[Any, Any, Any]:
+        """
+        Build train, validation, and test dataloaders using Energon.
+
+        This closely follows the pattern from Megatron-LM multimodal examples,
+        with adaptations for Primus configuration style.
+        """
+        from megatron.training import get_args
+
+        args = get_args()
+
+        # Check if we should create dataloaders on this rank
+        # (Only first TP rank and first/last PP stage)
+        if not self._is_dataloader_rank():
+            log_rank_0(
+                "Skipping dataloader creation on this rank (not first TP rank or not first/last PP stage)"
+            )
+            return None, None, None
+
+        # Create task encoder using factory
+        task_encoder = self.task_encoder_factory()
+        log_rank_0(f"Created task encoder: {type(task_encoder).__name__}")
+
+        # Create worker config for distributed loading
+        worker_config = self._create_worker_config(args)
+
+        # Get data path
+        data_path = self._get_data_path(args)
+
+        # Create training dataset using Energon
+        log_rank_0(f"Creating training dataset from: {data_path}")
+        train_dataset = get_train_dataset(
+            data_path,
+            batch_size=args.micro_batch_size,
+            task_encoder=task_encoder,
+            worker_config=worker_config,
+            virtual_epoch_length=getattr(args, "virtual_epoch_length", 1_000_000_000),
+            max_samples_per_sequence=getattr(args, "max_samples_per_sequence", 100),
+            shuffle_buffer_size=getattr(args, "shuffle_buffer_size", None),
+            handler=lambda *args: None,  # Error handler (print errors but continue)
+        )
+
+        # Wrap in savable loader for checkpointing support
+        prefetch_factor = getattr(args, "prefetch_factor", 2)
+        log_rank_0(f"Dataloader prefetch_factor: {prefetch_factor}")
+        train_dataloader = get_savable_loader(
+            train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor
+        )
+        train_dataloader = MegatronDataloaderWrapper(train_dataloader)
+        log_rank_0("Created training dataloader")
+
+        # Create validation dataloaders if evaluation is enabled
+        valid_dataloaders = None
+        if args.eval_iters > 0:
+            try:
+                log_rank_0("Creating validation dataloaders...")
+                val_datasets = get_val_datasets(
+                    data_path,
+                    batch_size=args.micro_batch_size,
+                    task_encoder=task_encoder,
+                    worker_config=worker_config,
+                    handler=lambda *args: None,
+                )
+
+                # Limit validation datasets to eval_iters * num_microbatches
+                val_datasets_limited = [
+                    LimitDataset(
+                        RepeatDataset(val_ds, worker_config=worker_config),
+                        length=args.eval_iters * get_num_microbatches(),
+                        worker_config=worker_config,
+                        reset_after_epoch=True,
+                    )
+                    for val_ds, _src_ds in val_datasets
+                ]
+
+                valid_dataloaders = [
+                    MegatronDataloaderWrapper(
+                        get_loader(valid_ds, worker_config=worker_config, prefetch_factor=prefetch_factor)
+                    )
+                    for valid_ds in val_datasets_limited
+                ]
+                log_rank_0(f"Created {len(valid_dataloaders)} validation dataloaders")
+            except Exception as e:
+                log_rank_0("=" * 80)
+                log_rank_0("WARNING: Could not create validation dataloaders")
+                log_rank_0(f"Reason: {e}")
+                log_rank_0("")
+                log_rank_0("This typically means:")
+                log_rank_0("  - The dataset does not have a validation split")
+                log_rank_0("  - The dataset path is incorrect")
+                log_rank_0("")
+                log_rank_0("AUTOMATIC FIX: Disabling evaluation (setting eval_iters = 0)")
+                log_rank_0("To enable evaluation, provide a dataset with validation split")
+                log_rank_0("=" * 80)
+                valid_dataloaders = None
+
+                # Automatically disable evaluation when no validation data exists
+                args.eval_iters = 0
+
+        # Test dataloaders not implemented for Energon
+        test_dataloader = None
+
+        return train_dataloader, valid_dataloaders, test_dataloader
+
+    @property
+    def is_distributed(self) -> bool:
+        """
+        Energon dataloaders are distributed (handle sharding internally).
+
+        Returns True to tell Megatron to bypass indexed dataset logic.
+        """
+        return True
+
+    def _is_dataloader_rank(self) -> bool:
+        """
+        Check if we should have the dataloader on this rank.
+
+        Energon dataloaders should only run on:
+            - First tensor parallel rank (data will be broadcast to others)
+            - First or last pipeline parallel stage (where embeddings/outputs are)
+
+        Reference:
+            Megatron-LM examples/multimodal/dataloader_provider.py:is_dataloader_rank()
+        """
+        # Run dataloader only on first tensor parallel rank
+        is_first_tp_rank = get_tensor_model_parallel_rank() == 0
+
+        # Check pipeline parallel stage
+        pp_size = get_pipeline_model_parallel_world_size()
+        if pp_size == 1:
+            # No pipeline parallelism
+            is_valid_pp_stage = True
+        else:
+            # With pipeline parallelism, run on first and last stage
+            pp_rank = get_pipeline_model_parallel_rank()
+            is_valid_pp_stage = pp_rank in (0, pp_size - 1)
+
+        return is_first_tp_rank and is_valid_pp_stage
+
+    def _create_worker_config(self, args) -> WorkerConfig:
+        """
+        Create Energon WorkerConfig for distributed loading.
+
+        WorkerConfig tells Energon how to shard data across workers.
+        """
+        rank = parallel_state.get_data_parallel_rank()
+        world_size = parallel_state.get_data_parallel_world_size()
+        data_parallel_group = parallel_state.get_data_parallel_group()
+
+        return WorkerConfig(
+            rank=rank,
+            world_size=world_size,
+            num_workers=getattr(args, "num_workers", 4),
+            data_parallel_group=data_parallel_group,
+        )
+
+    def _get_data_path(self, args) -> str:
+        """
+        Extract data path from args.
+
+        Handles both string and list formats.
+        Energon's get_train_dataset() can accept either:
+        - Directory path (will look for .nv-meta/dataset.yaml or dataset.yaml)
+        - Direct path to dataset.yaml file
+
+        This method returns the path as-is, letting Energon handle the resolution.
+        """
+        from pathlib import Path
+
+        data_path = args.data_path
+        if isinstance(data_path, list):
+            data_path = data_path[0]
+
+        if not data_path:
+            raise ValueError("data_path not found in configuration")
+
+        data_path = str(data_path)
+        path_obj = Path(data_path)
+
+        # Check if path exists
+        if not path_obj.exists():
+            raise ValueError(
+                f"data_path does not exist: {data_path}\n"
+                f"Please verify the path is correct and the dataset has been prepared."
+            )
+
+        # If it's a directory, check for dataset.yaml indicators
+        if path_obj.is_dir():
+            dataset_yaml = path_obj / "dataset.yaml"
+            nv_meta_yaml = path_obj / ".nv-meta" / "dataset.yaml"
+
+            if not dataset_yaml.exists() and not nv_meta_yaml.exists():
+                log_rank_0("=" * 80)
+                log_rank_0("WARNING: No dataset.yaml found in dataset directory")
+                log_rank_0(f"  Directory: {data_path}")
+                log_rank_0(f"  Expected: {dataset_yaml} or {nv_meta_yaml}")
+                log_rank_0("")
+                log_rank_0("This dataset may not be properly indexed with Energon.")
+                log_rank_0("Run: energon prepare <dataset_dir> --num-workers 4")
+                log_rank_0("=" * 80)
+            elif nv_meta_yaml.exists():
+                log_rank_0(f"Found Energon metadata at: {nv_meta_yaml}")
+            elif dataset_yaml.exists():
+                log_rank_0(f"Found dataset.yaml at: {dataset_yaml}")
+
+        log_rank_0(f"Using data path: {data_path}")
+        return data_path
+
+
+__all__ = ["EnergonDatasetProvider"]
diff --git a/primus/backends/megatron/data/synthetic/__init__.py b/primus/backends/megatron/data/synthetic/__init__.py
new file mode 100644
index 000000000..ff7eb721b
--- /dev/null
+++ b/primus/backends/megatron/data/synthetic/__init__.py
@@ -0,0 +1,32 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Synthetic data generation for diffusion models.
+
+This module provides mock datasets that generate synthetic data
+with correct tensor shapes and formats for training and testing
+without requiring real datasets or model weights.
+"""
+
+from .mock_datasets import (
+    LatentConfig,
+    MockDiffusionDataset,
+    MockFluxDataset,
+    MockFluxSchnellDataset,
+    ModelPreset,
+    PreGeneratedMockFluxDataset,
+    PreGeneratedMockFluxSchnellDataset,
+    TextEmbeddingConfig,
+)
+
+__all__ = [
+    "MockDiffusionDataset",
+    "MockFluxDataset",
+    "PreGeneratedMockFluxDataset",
+    "MockFluxSchnellDataset",
+    "PreGeneratedMockFluxSchnellDataset",
+    "LatentConfig",
+    "TextEmbeddingConfig",
+    "ModelPreset",
+]
diff --git a/primus/backends/megatron/data/synthetic/mock_datasets.py b/primus/backends/megatron/data/synthetic/mock_datasets.py
new file mode 100644
index 000000000..b7aaf179a
--- /dev/null
+++ b/primus/backends/megatron/data/synthetic/mock_datasets.py
@@ -0,0 +1,614 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Synthetic datasets for diffusion model training and testing.
+
+Provides mock datasets that generate random tensors with correct shapes
+for development, testing, and benchmarking without requiring real data.
+
+The module provides a flexible preset-based system for creating mock datasets
+for different diffusion model architectures.
+
+Quick Start
+-----------
+
+Using a preset (easiest):
+    >>> from primus.backends.megatron.data.synthetic import MockFluxDataset
+    >>> dataset = MockFluxDataset(num_samples=100, image_size=1024)
+    >>> sample = dataset[0]
+    >>> print(sample.keys())  # latents, t5_text_embeddings, clip_pooled_embeddings, etc.
+
+Using the preset system directly:
+    >>> from primus.backends.megatron.data.synthetic import MockDiffusionDataset
+    >>> dataset = MockDiffusionDataset(num_samples=100, model_preset='flux')
+
+Custom configuration for other models:
+    >>> from primus.backends.megatron.data.synthetic import (
+    ...     MockDiffusionDataset,
+    ...     LatentConfig,
+    ...     TextEmbeddingConfig,
+    ... )
+    >>> dataset = MockDiffusionDataset(
+    ...     num_samples=100,
+    ...     latent_config=LatentConfig(channels=4, downsample_factor=8),
+    ...     text_embedding_configs=[
+    ...         TextEmbeddingConfig(key='text_embeddings', shape=(77, 768)),
+    ...         TextEmbeddingConfig(key='pooled_embeddings', shape=(1280,)),
+    ...     ],
+    ... )
+
+Adding a New Model
+------------------
+
+To add support for a new diffusion model (e.g., SDXL):
+
+1. Define a position ID generator function (if needed):
+    >>> def sdxl_position_id_generator(image_size, latent_size, **kwargs):
+    ...     # SDXL doesn't use position IDs
+    ...     return {}
+
+2. Register the model preset:
+    >>> from primus.backends.megatron.data.synthetic import (
+    ...     MockDiffusionDataset,
+    ...     ModelPreset,
+    ...     LatentConfig,
+    ...     TextEmbeddingConfig,
+    ... )
+    >>> MockDiffusionDataset.register_model('sdxl', ModelPreset(
+    ...     latent_config=LatentConfig(channels=4, downsample_factor=8),
+    ...     text_embedding_configs=[
+    ...         TextEmbeddingConfig(key='text_embeddings', shape=(77, 2048)),
+    ...         TextEmbeddingConfig(key='pooled_embeddings', shape=(1280,)),
+    ...     ],
+    ...     position_id_generator=sdxl_position_id_generator,
+    ... ))
+
+3. Create a convenience class (optional):
+    >>> class MockSDXLDataset(MockDiffusionDataset):
+    ...     def __init__(self, num_samples=100, image_size=1024, seed=None):
+    ...         super().__init__(
+    ...             num_samples=num_samples,
+    ...             image_size=image_size,
+    ...             model_preset='sdxl',
+    ...             seed=seed,
+    ...         )
+
+Architecture
+------------
+
+The module uses a hybrid registry + subclass approach:
+
+- MockDiffusionDataset: Generic base class with preset system
+- MockFluxDataset: Convenience class for Flux models
+- Model presets: Registered configurations for specific model architectures
+- Configuration dataclasses: LatentConfig, TextEmbeddingConfig, ModelPreset
+
+This design provides flexibility for testing various diffusion models while
+maintaining a simple API for common use cases.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset
+
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Configuration Data Structures
+# ============================================================================
+
+
+@dataclass
+class LatentConfig:
+    """Configuration for latent generation."""
+
+    channels: int  # Number of latent channels
+    downsample_factor: int = 8  # VAE spatial downsampling factor
+
+
+@dataclass
+class TextEmbeddingConfig:
+    """Configuration for a single text embedding."""
+
+    key: str  # Output dictionary key (e.g., 't5_text_embeddings')
+    shape: Tuple[int, ...]  # Shape without batch dimension
+    dtype: str = "float32"  # Data type
+
+
+@dataclass
+class ModelPreset:
+    """Complete configuration preset for a diffusion model."""
+
+    latent_config: LatentConfig
+    text_embedding_configs: List[TextEmbeddingConfig]
+    position_id_generator: Optional[Callable] = None
+
+
+# ============================================================================
+# Position ID Generator Functions
+# ============================================================================
+
+
+def flux_position_id_generator(
+    image_size: int,
+    latent_size: int,
+    t5_seq_len: int,
+    seed: int,
+    idx: int,
+    device: torch.device = torch.device("cpu"),
+) -> Dict[str, torch.Tensor]:
+    """
+    Generate Flux-specific position IDs.
+
+    Args:
+        image_size: Original image size
+        latent_size: Latent space size (after VAE downsampling)
+        t5_seq_len: T5 sequence length
+        seed: Random seed (unused but kept for consistency)
+        idx: Sample index (unused but kept for consistency)
+        device: Device to create tensors on (default: cpu)
+
+    Returns:
+        Dictionary with 'img_ids' and 'txt_ids' tensors
+    """
+    # Generate image position IDs using Flux's standard function
+    img_ids = generate_image_position_ids(
+        batch_size=1,
+        height=latent_size,
+        width=latent_size,
+        device=device,
+        dtype=torch.float32,
+    ).squeeze(
+        0
+    )  # Remove batch dimension: (1, seq_len, 3) -> (seq_len, 3)
+
+    # Text position IDs (all zeros for text in Flux)
+    txt_ids = torch.zeros(t5_seq_len, 3, device=device)
+
+    return {"img_ids": img_ids, "txt_ids": txt_ids}
+
+
+class MockDiffusionDataset(Dataset):
+    """
+    Generic mock dataset for diffusion models.
+
+    Supports both preset-based initialization (via model_preset) and
+    custom configuration (via explicit configs). This allows the dataset
+    to generate synthetic data for various diffusion model architectures.
+
+    Usage:
+        # Using a preset (e.g., 'flux')
+        dataset = MockDiffusionDataset(num_samples=100, model_preset='flux')
+
+        # Using custom configuration
+        dataset = MockDiffusionDataset(
+            num_samples=100,
+            latent_config=LatentConfig(channels=4, downsample_factor=8),
+            text_embedding_configs=[
+                TextEmbeddingConfig(key='text_embeddings', shape=(77, 768)),
+            ]
+        )
+
+        # Using legacy parameters (backward compatible)
+        dataset = MockDiffusionDataset(
+            num_samples=100,
+            latent_channels=16,
+            t5_seq_len=512,
+            t5_hidden_dim=4096,
+            clip_hidden_dim=768,
+        )
+    """
+
+    # Registry of model presets
+    MODEL_PRESETS: Dict[str, ModelPreset] = {}
+
+    @classmethod
+    def register_model(cls, name: str, preset: ModelPreset):
+        """
+        Register a model preset for easy reuse.
+
+        Args:
+            name: Name of the preset (e.g., 'flux', 'sdxl')
+            preset: ModelPreset configuration object
+        """
+        cls.MODEL_PRESETS[name] = preset
+        logger.info(f"Registered model preset: {name}")
+
+    def __init__(
+        self,
+        num_samples: int = 100,
+        image_size: int = 1024,
+        # Option 1: Use preset
+        model_preset: Optional[str] = None,
+        # Option 2: Custom configuration
+        latent_config: Optional[LatentConfig] = None,
+        text_embedding_configs: Optional[List[TextEmbeddingConfig]] = None,
+        position_id_generator: Optional[Callable] = None,
+        seed: Optional[int] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str = "cpu",
+        vae_latent_mode: str = "presampled",
+        is_validation: bool = False,
+    ):
+        """
+        Initialize mock diffusion dataset.
+
+        Note: Always returns preencoded format (latents + embeddings + position IDs).
+        This matches the output of EncodedDiffusionTaskEncoder.
+
+        Args:
+            num_samples: Number of samples in dataset
+            image_size: Image size (assumes square images)
+            model_preset: Name of registered model preset to use
+            latent_config: Custom latent configuration (ignored if model_preset set)
+            text_embedding_configs: Custom text embedding configs (ignored if model_preset set)
+            position_id_generator: Custom position ID generator function (ignored if model_preset set)
+            seed: Random seed for reproducibility
+            dtype: Data type for tensors (default: torch.bfloat16)
+            device: Device to create tensors on ('cpu' or 'cuda', default: 'cpu')
+        """
+        super().__init__()
+
+        # Load from preset if specified
+        if model_preset:
+            if model_preset not in self.MODEL_PRESETS:
+                raise ValueError(
+                    f"Unknown model preset: {model_preset}. "
+                    f"Available presets: {list(self.MODEL_PRESETS.keys())}"
+                )
+            preset = self.MODEL_PRESETS[model_preset]
+            self.latent_config = preset.latent_config
+            self.text_embedding_configs = preset.text_embedding_configs
+            self.position_id_generator = preset.position_id_generator
+            logger.info(f"Using model preset: {model_preset}")
+        # Otherwise use custom configuration
+        else:
+            if latent_config is None:
+                raise ValueError("Must provide either model_preset or latent_config")
+            self.latent_config = latent_config
+            self.text_embedding_configs = text_embedding_configs or []
+            self.position_id_generator = position_id_generator
+
+        # Common initialization
+        self.num_samples = num_samples
+        self.image_size = image_size
+        self.seed = seed if seed is not None else 0
+        self.dtype = dtype
+        self.device = torch.device(device)
+        self.vae_latent_mode = vae_latent_mode
+        self.is_validation = is_validation
+        self.latent_size = image_size // self.latent_config.downsample_factor
+
+        # Store legacy attributes for backward compatibility
+        self.latent_channels = self.latent_config.channels
+
+        # Extract T5 and CLIP dims if present (for backward compatibility)
+        self.t5_seq_len = None
+        self.t5_hidden_dim = None
+        self.clip_hidden_dim = None
+        for emb_config in self.text_embedding_configs:
+            if emb_config.key in ("t5_text_embeddings", "prompt_embeds") and len(emb_config.shape) == 2:
+                self.t5_seq_len, self.t5_hidden_dim = emb_config.shape
+            elif (
+                emb_config.key in ("clip_pooled_embeddings", "pooled_prompt_embeds")
+                and len(emb_config.shape) == 1
+            ):
+                self.clip_hidden_dim = emb_config.shape[0]
+
+        logger.info(
+            f"Initialized MockDiffusionDataset: {num_samples} samples, "
+            f"image_size={image_size}, latent_size={self.latent_size}"
+        )
+
+    def __len__(self) -> int:
+        """Return number of samples."""
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """
+        Get a mock sample.
+
+        Args:
+            idx: Sample index
+
+        Returns:
+            Dictionary with mock data matching DiffusionSample format
+
+        Raises:
+            IndexError: If idx is out of bounds
+        """
+        if idx >= self.num_samples or idx < 0:
+            raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples")
+
+        # Set seed based on instance seed and idx for reproducibility
+        gen = torch.Generator(device=self.device)
+        gen.manual_seed(self.seed + idx)
+
+        latent_shape = (
+            self.latent_config.channels,
+            self.latent_size,
+            self.latent_size,
+        )
+
+        if self.vae_latent_mode == "resample":
+            mean = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device)
+            logvar = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device)
+            sample = {"mean": mean, "logvar": logvar}
+        else:
+            latents = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device)
+            sample = {"latents": latents}
+
+        # Generate text embeddings based on configs
+        for emb_config in self.text_embedding_configs:
+            # Get dtype
+            dtype = getattr(torch, emb_config.dtype, torch.float32)
+
+            # Generate tensor with the specified shape
+            embedding = torch.randn(
+                *emb_config.shape,
+                generator=gen,
+                dtype=dtype,
+                device=self.device,
+            )
+
+            sample[emb_config.key] = embedding
+
+        # Generate position IDs if generator provided
+        if self.position_id_generator is not None:
+            position_ids = self.position_id_generator(
+                image_size=self.image_size,
+                latent_size=self.latent_size,
+                t5_seq_len=self.t5_seq_len if self.t5_seq_len else 512,
+                seed=self.seed,
+                idx=idx,
+                device=self.device,
+            )
+            sample.update(position_ids)
+
+        # Generate mock caption
+        sample["caption"] = f"Mock caption {idx}"
+
+        if self.is_validation:
+            sample["timestep"] = torch.tensor(idx % 8)
+
+        return sample
+
+
+# ============================================================================
+# Register Model Presets
+# ============================================================================
+
+# Register Flux model preset (FLUX.1-dev: T5 max_sequence_length=512)
+MockDiffusionDataset.register_model(
+    "flux",
+    ModelPreset(
+        latent_config=LatentConfig(channels=16, downsample_factor=8),
+        text_embedding_configs=[
+            TextEmbeddingConfig(
+                key="prompt_embeds", shape=(512, 4096), dtype="bfloat16"
+            ),  # T5 text embeddings
+            TextEmbeddingConfig(
+                key="pooled_prompt_embeds", shape=(768,), dtype="bfloat16"
+            ),  # CLIP pooled embeddings
+        ],
+        position_id_generator=flux_position_id_generator,
+    ),
+)
+
+# Register Flux Schnell preset (FLUX.1-schnell: T5 max_sequence_length=256)
+# Matches the MLPerf Training v5.1 benchmark specification and NVIDIA reference.
+MockDiffusionDataset.register_model(
+    "flux_schnell",
+    ModelPreset(
+        latent_config=LatentConfig(channels=16, downsample_factor=8),
+        text_embedding_configs=[
+            TextEmbeddingConfig(
+                key="prompt_embeds", shape=(256, 4096), dtype="bfloat16"
+            ),  # T5 text embeddings
+            TextEmbeddingConfig(
+                key="pooled_prompt_embeds", shape=(768,), dtype="bfloat16"
+            ),  # CLIP pooled embeddings
+        ],
+        position_id_generator=flux_position_id_generator,
+    ),
+)
+
+
+# ============================================================================
+# Model-Specific Convenience Classes
+# ============================================================================
+
+
+class MockFluxDataset(MockDiffusionDataset):
+    """
+    Mock dataset specifically for Flux diffusion model.
+
+    Convenience wrapper that uses the 'flux' preset with Flux-specific
+    default parameters:
+        - image_size=1024
+        - latent_channels=16
+        - t5_seq_len=512
+        - t5_hidden_dim=4096
+        - clip_hidden_dim=768
+
+    This class provides a simplified API for the common case of testing
+    Flux models without needing to understand the preset system.
+    """
+
+    def __init__(
+        self,
+        num_samples: int = 100,
+        image_size: int = 1024,
+        seed: Optional[int] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        """
+        Initialize mock Flux dataset.
+
+        Note: Always returns preencoded format (matches TaskEncoder output).
+
+        Args:
+            num_samples: Number of samples in dataset
+            image_size: Image size (default 1024 for Flux)
+            seed: Random seed for reproducibility
+            dtype: Data type for tensors (default: torch.bfloat16)
+            device: Device to create tensors on ('cpu' or 'cuda', default: 'cuda')
+        """
+        super().__init__(
+            num_samples=num_samples,
+            image_size=image_size,
+            model_preset="flux",
+            seed=seed,
+            dtype=dtype,
+            device=device,
+            **kwargs,
+        )
+
+        logger.info(f"Initialized MockFluxDataset with Flux preset on device={device}")
+
+
+class PreGeneratedMockFluxDataset(MockFluxDataset):
+    """
+    Pre-generated mock Flux dataset for maximum training throughput.
+
+    This dataset generates all samples once during initialization and caches
+    them in memory. This eliminates on-the-fly random number generation overhead
+    and provides benchmark-quality performance matching latent caching approaches.
+
+    Memory usage: For 1000 samples at 512x512 resolution with bf16:
+        - Latents: 1000 * 16 * 64 * 64 * 2 bytes = ~128 MB
+        - Text embeddings: 1000 * (512 * 4096 + 768) * 2 bytes = ~4 GB
+        - Total: ~4.2 GB per GPU (acceptable for 192GB MI300X)
+
+    Use this instead of MockFluxDataset for:
+        - Performance benchmarking
+        - Maximum training throughput
+        - Consistent timing measurements
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Initialize and pre-generate all samples."""
+        from primus.core.utils.module_utils import log_rank_0
+
+        # Initialize parent class
+        super().__init__(*args, **kwargs)
+
+        log_rank_0(f"Pre-generating {self.num_samples} mock Flux samples...")
+        log_rank_0(f"  Image size: {self.image_size}x{self.image_size}")
+        log_rank_0(f"  Latent size: {self.latent_size}x{self.latent_size}")
+
+        # Pre-generate all samples using parent's __getitem__
+        self._samples = [MockFluxDataset.__getitem__(self, i) for i in range(self.num_samples)]
+
+        # Calculate memory usage (only count tensors)
+        sample_size = sum(
+            tensor.element_size() * tensor.numel()
+            for tensor in self._samples[0].values()
+            if hasattr(tensor, "element_size")  # Only count torch tensors
+        )
+        total_mb = (sample_size * self.num_samples) / (1024 * 1024)
+
+        log_rank_0(f"Pre-generation complete!")
+        log_rank_0(f"  Memory usage: {total_mb:.1f} MB")
+        log_rank_0(f"  Data loading overhead: eliminated")
+
+    def __getitem__(self, idx: int):
+        """
+        Return pre-generated sample (no computation).
+
+        This is 10-100x faster than on-the-fly generation as it's just
+        a memory lookup with no random number generation or tensor creation.
+        """
+        if idx >= self.num_samples or idx < 0:
+            raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples")
+        return self._samples[idx]
+
+
+class MockFluxSchnellDataset(MockDiffusionDataset):
+    """
+    Mock dataset for FLUX.1-schnell (MLPerf Training v5.1 benchmark).
+
+    Uses the 'flux_schnell' preset which matches the NVIDIA MLPerf reference:
+        - T5 max_sequence_length=256 (vs 512 for FLUX.1-dev)
+        - image_size=256 (benchmark default)
+        - latent_channels=16, context_dim=4096, vec_in_dim=768
+    """
+
+    def __init__(
+        self,
+        num_samples: int = 100,
+        image_size: int = 256,
+        seed: Optional[int] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        super().__init__(
+            num_samples=num_samples,
+            image_size=image_size,
+            model_preset="flux_schnell",
+            seed=seed,
+            dtype=dtype,
+            device=device,
+            **kwargs,
+        )
+        logger.info(f"Initialized MockFluxSchnellDataset on device={device}")
+
+
+class PreGeneratedMockFluxSchnellDataset(MockFluxSchnellDataset):
+    """
+    Pre-generated mock FLUX.1-schnell dataset for benchmark throughput.
+
+    Generates all samples once at init and caches them in memory.
+    Matches the MLPerf Training v5.1 data shapes (T5 seq_len=256).
+    """
+
+    def __init__(self, *args, **kwargs):
+        from primus.core.utils.module_utils import log_rank_0
+
+        super().__init__(*args, **kwargs)
+
+        log_rank_0(f"Pre-generating {self.num_samples} mock Flux-schnell samples...")
+        log_rank_0(f"  Image size: {self.image_size}x{self.image_size}")
+        log_rank_0(f"  Latent size: {self.latent_size}x{self.latent_size}")
+
+        self._samples = [MockFluxSchnellDataset.__getitem__(self, i) for i in range(self.num_samples)]
+
+        sample_size = sum(
+            tensor.element_size() * tensor.numel()
+            for tensor in self._samples[0].values()
+            if hasattr(tensor, "element_size")
+        )
+        total_mb = (sample_size * self.num_samples) / (1024 * 1024)
+
+        log_rank_0(f"Pre-generation complete!")
+        log_rank_0(f"  Memory usage: {total_mb:.1f} MB")
+        log_rank_0(f"  Data loading overhead: eliminated")
+
+    def __getitem__(self, idx: int):
+        if idx >= self.num_samples or idx < 0:
+            raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples")
+        return self._samples[idx]
+
+
+__all__ = [
+    "LatentConfig",
+    "TextEmbeddingConfig",
+    "ModelPreset",
+    "MockDiffusionDataset",
+    "MockFluxDataset",
+    "PreGeneratedMockFluxDataset",
+    "MockFluxSchnellDataset",
+    "PreGeneratedMockFluxSchnellDataset",
+]
diff --git a/primus/backends/megatron/data/synthetic_dataset_provider.py b/primus/backends/megatron/data/synthetic_dataset_provider.py
new file mode 100644
index 000000000..67588be92
--- /dev/null
+++ b/primus/backends/megatron/data/synthetic_dataset_provider.py
@@ -0,0 +1,217 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+SyntheticDatasetProvider: Dataset provider for synthetic/mock data.
+
+Provides dataloaders for synthetic data generation during development,
+testing, and benchmarking without requiring real datasets.
+"""
+
+import logging
+from importlib import import_module
+from typing import Any, Dict, List, Optional, Tuple
+
+from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+from primus.backends.megatron.data.dataset_provider import DatasetProvider
+from primus.core.utils.module_utils import log_rank_0
+
+logger = logging.getLogger(__name__)
+
+
+class SyntheticDatasetProvider(DatasetProvider):
+    """
+    Dataset provider for synthetic/mock data.
+
+    This provider creates synthetic datasets based on configuration.
+    The dataset class and parameters are specified in YAML config,
+    keeping dataset creation logic separate from trainer logic.
+
+    Usage in YAML:
+        modules:
+          pre_trainer:
+            mock_data: true
+            mock_dataset:
+              class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+              params:
+                num_samples: 1000
+                image_size: 512
+    """
+
+    DEFAULT_DATASETS = {
+        "flux": "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset",
+        "flux_onthefly": "primus.backends.megatron.data.synthetic.MockFluxDataset",
+        "flux_schnell": "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxSchnellDataset",
+        "flux_schnell_onthefly": "primus.backends.megatron.data.synthetic.MockFluxSchnellDataset",
+    }
+
+    def __init__(self, dataset_config: Optional[Dict[str, Any]] = None, model_type: str = "flux"):
+        """
+        Initialize synthetic dataset provider.
+
+        Args:
+            dataset_config: Dictionary with 'class' and 'params' keys:
+                {
+                    'class': 'fully.qualified.ClassName',  # Optional, uses default for model_type
+                    'params': {  # Optional, dataset-specific parameters
+                        'num_samples': 1000,
+                        'image_size': 512,
+                        # ... other dataset params
+                    }
+                }
+            model_type: Model type for default dataset selection ('flux', etc.)
+                Used when dataset_config['class'] is not specified.
+        """
+        self.dataset_config = dataset_config or {}
+        self.model_type = model_type
+
+        # Determine dataset class to use
+        dataset_class_from_config = self.dataset_config.get("class")
+
+        # If class is None or not specified, use default for model_type
+        if dataset_class_from_config is None or dataset_class_from_config == "null":
+            self.dataset_class_path = self.DEFAULT_DATASETS.get(model_type)
+        else:
+            self.dataset_class_path = dataset_class_from_config
+
+        if not self.dataset_class_path:
+            raise ValueError(
+                f"No default dataset for model_type='{model_type}'. "
+                f"Available types: {list(self.DEFAULT_DATASETS.keys())}. "
+                f"Or specify dataset_config['class'] explicitly."
+            )
+
+        # Get dataset parameters (will be augmented with trainer config later)
+        self.dataset_params = self.dataset_config.get("params", {})
+
+        logger.debug(f"SyntheticDatasetProvider initialized: {self.dataset_class_path}")
+
+    def _import_dataset_class(self):
+        """
+        Dynamically import the dataset class.
+
+        Returns:
+            Dataset class
+
+        Raises:
+            ImportError: If class cannot be imported
+            AttributeError: If class doesn't exist in module
+        """
+        try:
+            module_path, class_name = self.dataset_class_path.rsplit(".", 1)
+            module = import_module(module_path)
+            dataset_class = getattr(module, class_name)
+            return dataset_class
+        except (ValueError, ImportError, AttributeError) as e:
+            raise ImportError(
+                f"Failed to import dataset class '{self.dataset_class_path}': {e}\n"
+                f"Ensure the class path is correct and the module is installed."
+            ) from e
+
+    def create_dataloaders(
+        self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None
+    ) -> Tuple[Any, Any, Any]:
+        """
+        Create synthetic dataloaders.
+
+        Args:
+            trainer_config: Megatron args namespace (from megatron.training.get_args())
+            train_val_test_num_samples: [train_samples, valid_samples, test_samples]
+            vp_stage: Virtual pipeline stage (for VP parallelism)
+
+        Returns:
+            Tuple of (train_dataloader, None, None)
+            Note: Validation/test loaders not supported for synthetic data
+        """
+        from megatron.training import get_args
+        from torch.utils.data import DataLoader
+
+        args = get_args()
+
+        log_rank_0("=" * 80)
+        log_rank_0("Creating SYNTHETIC/MOCK dataloaders")
+        log_rank_0(f"Dataset class: {self.dataset_class_path}")
+        log_rank_0("=" * 80)
+
+        # Import dataset class
+        dataset_class = self._import_dataset_class()
+
+        # Merge config: YAML params + runtime params from trainer
+        final_params = {
+            **self.dataset_params,  # From YAML
+            "seed": getattr(trainer_config, "seed", 42),  # Always use trainer seed
+        }
+
+        # Override with trainer config if specified (backwards compatibility)
+        if hasattr(trainer_config, "image_size") and "image_size" not in self.dataset_params:
+            final_params["image_size"] = trainer_config.image_size
+
+        # Forward vae_latent_mode from trainer config to mock dataset
+        if hasattr(trainer_config, "vae_latent_mode") and "vae_latent_mode" not in self.dataset_params:
+            final_params["vae_latent_mode"] = trainer_config.vae_latent_mode
+
+        # Set defaults if not specified
+        final_params.setdefault("num_samples", 1000)
+        final_params.setdefault("image_size", 512)
+
+        log_rank_0(f"Dataset parameters:")
+        for key, value in sorted(final_params.items()):
+            log_rank_0(f"  {key}: {value}")
+
+        # Create dataset
+        dataset = dataset_class(**final_params)
+
+        log_rank_0(f"Created mock dataset: {dataset_class.__name__}")
+        log_rank_0(f"  num_samples: {len(dataset)}")
+        log_rank_0(f"  batch_size: {args.micro_batch_size}")
+
+        # Create PyTorch DataLoader
+        # Note: num_workers=0 for synthetic data (generation is fast)
+        train_loader = DataLoader(
+            dataset,
+            batch_size=args.micro_batch_size,
+            shuffle=True,
+            num_workers=0,
+            drop_last=True,
+        )
+
+        # Wrap in MegatronDataloaderWrapper for:
+        #   1. Cyclic iteration (never exhausts)
+        #   2. Megatron training loop compatibility
+        #   3. Checkpoint interface (no-op for synthetic data)
+        train_loader = MegatronDataloaderWrapper(train_loader)
+
+        log_rank_0("✓ Synthetic dataloader ready (infinite iteration)")
+        log_rank_0("=" * 80)
+
+        # Create validation dataloader when eval_iters > 0
+        val_loader = None
+        eval_iters = getattr(args, "eval_iters", 0)
+        if eval_iters > 0:
+            val_num_samples = max(256, eval_iters * args.micro_batch_size)
+            val_params = {**final_params, "is_validation": True, "num_samples": val_num_samples}
+            val_dataset = dataset_class(**val_params)
+            val_raw_loader = DataLoader(
+                val_dataset,
+                batch_size=args.micro_batch_size,
+                shuffle=False,
+                num_workers=0,
+                drop_last=True,
+            )
+            val_loader = MegatronDataloaderWrapper(val_raw_loader)
+            log_rank_0(f"Created validation dataloader: {val_num_samples} samples, is_validation=True")
+
+        return train_loader, val_loader, None
+
+    @property
+    def is_distributed(self) -> bool:
+        """
+        Synthetic data doesn't need distributed handling.
+
+        Returns True to tell Megatron to bypass indexed dataset logic.
+        Each rank creates its own synthetic data independently.
+        """
+        return True
+
+
+__all__ = ["SyntheticDatasetProvider"]
diff --git a/primus/backends/megatron/data_loader_store.py b/primus/backends/megatron/data_loader_store.py
new file mode 100644
index 000000000..02372e946
--- /dev/null
+++ b/primus/backends/megatron/data_loader_store.py
@@ -0,0 +1,132 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Megatron batch-loading helpers used by the pipeline schedules.
+
+``get_batch_func`` builds a micro-batch honoring TP/CP sharding (and the
+ZeroBubble sequence-split path), and ``DataLoaderStore`` provides the
+push/pop cache the ZeroBubble runtime and the ``forward_step`` patch use to
+pre-fetch batches (optionally on a dedicated H2D stream).
+"""
+
+import collections
+
+import torch
+from megatron.core.utils import StragglerDetector
+from megatron.training import get_args, get_timers
+from megatron.training.utils import (
+    get_batch_on_this_cp_rank,
+    get_batch_on_this_tp_rank,
+    is_first_or_last_pipeline_stage,
+)
+
+stimer = StragglerDetector()
+
+mb_batch = None
+
+
+def get_batch_func(data_iterator, vp_stage=None):
+    # TODO: this is pretty hacky, find a better way
+    if not is_first_or_last_pipeline_stage(vp_stage):
+        return None, None, None, None, None
+
+    # assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}"
+    # get batches based on the TP rank you are on
+    batch = get_batch_on_this_tp_rank(data_iterator)
+
+    # slice batch along sequence dimension for context parallelism
+    args = get_args()
+
+    if args.patch_zero_bubble:
+        from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_vars import (
+            get_seq_split_idx,
+        )
+
+        global mb_batch
+        # "or 0" to support original 1f1b and interleaved-1f1b in schedules.py
+        seq_split_idx = get_seq_split_idx() or 0
+        if seq_split_idx == 0:
+            # get batches based on the TP rank you are on
+            mb_batch = get_batch_on_this_tp_rank(data_iterator)
+            assert (
+                mb_batch["attention_mask"] is None
+            ), "attention_mask should be None, please enable --no-create-attention-mask-in-dataloader"
+        batch = {}
+        for k in mb_batch.keys():
+            v = mb_batch[k]
+            if v is None:
+                batch[k] = v
+                continue
+
+            assert v.shape[1] % get_args().num_seq_splits == 0, f"{k} size {v.shape}"
+            start_idx = seq_split_idx * v.shape[1] // get_args().num_seq_splits
+            end_idx = (seq_split_idx + 1) * v.shape[1] // get_args().num_seq_splits
+            if len(v.shape) > 2:
+                batch[k] = v[:, start_idx:end_idx, :].contiguous()
+            else:
+                batch[k] = v[:, start_idx:end_idx].contiguous()
+
+    if args.context_parallel_size > 1 and args.enable_primus_turbo and args.use_turbo_attention:
+        try:
+            from primus.backends.megatron.core.utils import (
+                produce_attention_sharder,
+                shard_batch_on_this_cp_rank,
+            )
+        except:
+            raise ImportError("Module 'primus_turbo' may not installed. Please install it")
+        sharder = produce_attention_sharder(args.cp_comm_type)
+        batch = shard_batch_on_this_cp_rank(sharder, batch)
+    else:
+        batch = get_batch_on_this_cp_rank(batch)
+
+    # Return a stable, explicitly-ordered 5-tuple so both this path and the
+    # early not-first/last-stage return path have the same shape/type. The keys
+    # match megatron's ``get_batch_on_this_tp_rank`` and the unpack order at all
+    # call sites (tokens, labels, loss_mask, attention_mask, position_ids).
+    return (
+        batch["tokens"],
+        batch["labels"],
+        batch["loss_mask"],
+        batch["attention_mask"],
+        batch["position_ids"],
+    )
+
+
+class DataLoaderStore:
+    cache = collections.deque()
+
+    @classmethod
+    def push(cls, data_iterator, h2d_stream=False, vp_stage=None):
+        timers = get_timers()
+        # Get the batch.
+        timers("batch-generator", log_level=2).start()
+        global stimer
+
+        with stimer(bdata=True):
+            if h2d_stream:
+                from primus.backends.megatron.core.pipeline_parallel.zerobubble.offload import (
+                    get_offload_h2d_stream,
+                )
+
+                load_event = torch.cuda.Event()
+                original_stream = torch.cuda.current_stream()
+                with torch.cuda.stream(get_offload_h2d_stream()):
+                    data = get_batch_func(data_iterator, vp_stage)
+                    for x in data:
+                        if x is not None:
+                            x.record_stream(original_stream)
+                    load_event.record()
+                    cls.cache.append((data, load_event))
+            else:
+                cls.cache.append((get_batch_func(data_iterator, vp_stage), None))
+        timers("batch-generator").stop()
+
+    @classmethod
+    def pop(cls):
+        data, load_event = cls.cache.popleft()
+        if load_event:
+            load_event.wait()
+        return data
diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py
new file mode 100644
index 000000000..dd91f3285
--- /dev/null
+++ b/primus/backends/megatron/diffusion_trainer.py
@@ -0,0 +1,386 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Base diffusion trainer for Megatron-LM.
+
+This trainer provides a foundation for diffusion model training by overriding
+the forward step and dataset provider methods from MegatronPretrainTrainer.
+"""
+
+from abc import abstractmethod
+
+import torch
+import torch.nn.functional as F
+
+from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer
+from primus.core.utils.module_utils import log_rank_0
+
+
+class DiffusionPretrainTrainer(MegatronPretrainTrainer):
+    """
+    Base trainer for diffusion models.
+
+    This trainer inherits from the backend's MegatronPretrainTrainer (which uses
+    MegatronBaseTrainer).
+
+    It overrides the forward step and dataset provider methods to support
+    diffusion-specific training. Subclasses should implement create_scheduler()
+    and may override create_model() if needed.
+
+    Configuration is accessed via backend_args (set by BaseTrainer.__init__()),
+    not via module_config.params, to avoid BaseModule dependency.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize diffusion trainer.
+
+        Args:
+            *args: Positional arguments passed to parent trainer
+            **kwargs: Keyword arguments passed to parent trainer
+                     (backend_args is extracted from kwargs by BaseTrainer)
+        """
+        super().__init__(*args, **kwargs)
+
+        self._scheduler = None
+        self._compiled_loss_fn = None
+        self._forward_step_count = 0
+        self._forward_step_count_initialized = False
+
+        # Composition pattern: avoids recreating the provider on each call
+        use_mock_data = getattr(self.backend_args, "mock_data", False)
+
+        if use_mock_data:
+            from primus.backends.megatron.data.synthetic_dataset_provider import (
+                SyntheticDatasetProvider,
+            )
+
+            # Get mock dataset configuration from YAML (if specified)
+            mock_dataset_config_ns = getattr(self.backend_args, "mock_dataset", None)
+            mock_dataset_config = self._convert_namespace_to_dict(mock_dataset_config_ns)
+
+            # Determine model type for default dataset selection
+            model_type_raw = getattr(self.backend_args, "model_type", None)
+            model_type = (
+                model_type_raw
+                if (isinstance(model_type_raw, str) and model_type_raw != "diffusion_model")
+                else "flux"
+            )
+
+            self.data_provider = SyntheticDatasetProvider(
+                dataset_config=mock_dataset_config, model_type=model_type
+            )
+            log_rank_0("Created SyntheticDatasetProvider for mock data")
+        else:
+            from primus.backends.megatron.data.energon_dataset_provider import (
+                EnergonDatasetProvider,
+            )
+
+            self.data_provider = EnergonDatasetProvider(task_encoder_factory=lambda: self.get_task_encoder())
+            log_rank_0("Created EnergonDatasetProvider for real data")
+
+        log_rank_0(f"{self.__class__.__name__} initialized")
+        log_rank_0(f"Data provider: {type(self.data_provider).__name__}")
+
+    def _convert_namespace_to_dict(self, ns_obj):
+        """
+        Convert SimpleNamespace (or nested SimpleNamespace) to dict.
+
+        Handles nested namespaces by recursively converting them.
+
+        Args:
+            ns_obj: SimpleNamespace, dict, or None
+
+        Returns:
+            dict, original input (if already a dict or primitive), or None
+        """
+        if ns_obj is None:
+            return None
+
+        if not hasattr(ns_obj, "__dict__"):
+            return ns_obj  # Already a dict or other type
+
+        result = vars(ns_obj)
+
+        # Recursively convert nested namespaces
+        if "params" in result and hasattr(result["params"], "__dict__"):
+            result["params"] = vars(result["params"])
+
+        return result
+
+    def setup(self):
+        """
+        Override setup() to inject diffusion model_provider.
+
+        MegatronBaseTrainer.setup() handles Megatron path, global vars, and parse_args patching.
+        We only need to set the model_provider here.
+        """
+        # Ensure data_parallel_size is set before parent setup() calls set_primus_global_variables()
+        # This is required by set_primus_global_variables()
+        if not hasattr(self.backend_args, "data_parallel_size"):
+            world_size = getattr(self.backend_args, "world_size", 1)
+            tensor_model_parallel_size = getattr(self.backend_args, "tensor_model_parallel_size", 1)
+            pipeline_model_parallel_size = getattr(self.backend_args, "pipeline_model_parallel_size", 1)
+            context_parallel_size = getattr(self.backend_args, "context_parallel_size", 1)
+            data_parallel_size = world_size // (
+                tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
+            )
+            setattr(self.backend_args, "data_parallel_size", data_parallel_size)
+            log_rank_0(
+                f"Computed data_parallel_size={data_parallel_size} from world_size={world_size}, tp={tensor_model_parallel_size}, pp={pipeline_model_parallel_size}, cp={context_parallel_size}"
+            )
+
+        # Create and set diffusion model provider
+        def _diffusion_model_provider(
+            pre_process=True, post_process=True, vp_stage=None, config=None, pg_collection=None
+        ):
+            """
+            Model provider wrapper for diffusion models.
+
+            This matches Megatron's model_provider signature but calls
+            our diffusion-specific create_model() instead of gpt_builder.
+
+            Note: vp_stage, config, and pg_collection are accepted for interface
+            compatibility with Megatron's model_provider signature but are not
+            used by diffusion models.
+            """
+            return self.create_model(pre_process=pre_process, post_process=post_process)
+
+        self.model_provider = _diffusion_model_provider
+        log_rank_0("=" * 80)
+        log_rank_0("Overridden model_provider to use diffusion model builder")
+        log_rank_0("=" * 80)
+
+        # Call parent's setup() which handles Megatron initialization
+        super().setup()
+
+    @abstractmethod
+    def create_model(self, pre_process=True, post_process=True):
+        """
+        Create diffusion model instance.
+
+        Returns:
+            Model instance (e.g., Flux)
+
+        Example:
+            from primus.backends.megatron.core.models.diffusion.flux import Flux
+            return Flux(config=self.model_config)
+        """
+        raise NotImplementedError("Subclasses must implement create_model()")
+
+    def _get_loss_fn(self):
+        """Return the loss function, optionally compiled.
+
+        When torch.compile is enabled for the model, the loss function is also
+        compiled as a standalone region.  This fuses the ~10 eager ATen ops
+        (sub, float casts, mse_loss, mean) into 1-2 kernels and gives the
+        autograd engine a single CompiledFunctionBackward node instead of
+        multiple eager backward nodes.
+
+        The compiled function is created lazily on first call and cached.
+        """
+        if self._compiled_loss_fn is not None:
+            return self._compiled_loss_fn
+
+        from primus.backends.megatron.training.diffusion.loss_computation import (
+            compute_flow_matching_loss,
+        )
+
+        try:
+            from megatron.training import get_args
+
+            args = get_args()
+            compile_enabled = getattr(args, "enable_torch_compile", False)
+        except Exception:
+            compile_enabled = False
+
+        if compile_enabled:
+            import torch
+
+            self._compiled_loss_fn = torch.compile(
+                compute_flow_matching_loss,
+                backend="inductor",
+                fullgraph=False,
+            )
+            log_rank_0("[DiffusionPretrainTrainer] Loss function compiled with torch.compile")
+        else:
+            self._compiled_loss_fn = compute_flow_matching_loss
+
+        return self._compiled_loss_fn
+
+    def forward_step(self, data_iterator, model, return_schedule_plan=False):
+        """
+        Forward training step for diffusion models.
+
+        Args:
+            data_iterator: Data iterator
+            model: Diffusion model (Flux)
+            return_schedule_plan: Whether to return schedule plan (for pipeline parallelism)
+
+        Returns:
+            Tuple of (noise_pred, loss_func_callable)
+        """
+        from primus.backends.megatron.training.diffusion.forward_step import (
+            flux_forward_step_func,
+        )
+
+        # Skip counter advance in eval. Advancing _forward_step_count on
+        # validation steps would shift the next training step's per-step seed
+        # by eval_iters * num_microbatches per --eval-interval window,
+        # defeating the goal of isolating training RNG from unrelated forward
+        # passes. Eval forward passes reuse the most recent training counter
+        # value, so the per-step CUDA reseed is a no-op replay during eval.
+        if model.training:
+            self._forward_step_count += 1
+
+        # Megatron's pattern: forward_step returns model output, loss_func computes loss
+        noise_pred, clean_latents, noise, loss_mask, metrics, is_validation = flux_forward_step_func(
+            data_iterator,
+            model,
+            scheduler=self.scheduler,
+            use_guidance_embed=getattr(self, "use_guidance_embed", False),
+            guidance_scale=getattr(self, "guidance_scale", None),
+            timestep_sampler=getattr(self, "timestep_sampler", None),
+            cfg_dropout_prob=getattr(self, "cfg_dropout_prob", 0.0),
+            empty_t5_encodings=getattr(self, "empty_t5_encodings", None),
+            empty_clip_encodings=getattr(self, "empty_clip_encodings", None),
+            vae_scale=getattr(self, "vae_scale", None),
+            vae_shift=getattr(self, "vae_shift", None),
+            vae_latent_mode=getattr(self, "vae_latent_mode", "presampled"),
+            per_step_rng_reseed=getattr(self, "per_step_rng_reseed", False),
+            step_count=self._forward_step_count,
+        )
+
+        # Store values needed for loss computation (will be used by loss function)
+        self._last_clean_latents = clean_latents
+        self._last_noise = noise
+        self._last_loss_mask = loss_mask
+
+        # Store metrics in runtime state
+        if hasattr(self, "runtime_state") and self.runtime_state:
+            self.runtime_state.update_metrics(metrics)
+        else:
+            log_rank_0("[DiffusionPretrainTrainer] WARNING: runtime_state not available, metrics not stored")
+
+        if is_validation:
+
+            def val_loss_func(output_tensor, non_loss_data=False):
+                if non_loss_data:
+                    return output_tensor
+                target = self._last_noise - self._last_clean_latents
+                loss = F.mse_loss(output_tensor.float(), target.float(), reduction="none")
+                loss_per_sample = loss.mean(dim=tuple(range(1, loss.ndim)))
+                loss_sum = loss_per_sample.sum()
+                sample_count = torch.tensor(
+                    loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device
+                )
+                return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())}
+
+            return noise_pred, val_loss_func
+
+        def diffusion_loss_func(output_tensor, non_loss_data=False):
+            if non_loss_data:
+                return output_tensor
+
+            loss = self._get_loss_fn()(
+                output_tensor, self._last_clean_latents, self._last_noise, self._last_loss_mask
+            )
+
+            reporting_metrics = {"reduced_train_loss": loss.detach().clone()}
+
+            return loss, reporting_metrics
+
+        return noise_pred, diffusion_loss_func
+
+    def get_forward_step(self):
+        """
+        Return forward step function for diffusion models.
+
+        Returns a function that wraps self.forward_step() to match
+        the interface expected by Megatron's pretrain() function.
+
+        On the first call, reconstructs the forward step counter from
+        checkpoint state (args.iteration * num_microbatches) so that
+        the per-step RNG seed sequence continues correctly after resume.
+        """
+
+        def diffusion_forward_step(data_iterator, model):
+            if not self._forward_step_count_initialized:
+                from megatron.core.num_microbatches_calculator import (
+                    get_num_microbatches,
+                )
+                from megatron.training import get_args
+
+                args = get_args()
+                # Reconstruct counter from checkpoint iteration. Assumes no
+                # iterations were skipped (iterations_to_skip is unused in
+                # diffusion training).
+                self._forward_step_count = args.iteration * get_num_microbatches()
+                self._forward_step_count_initialized = True
+
+            return self.forward_step(data_iterator, model, return_schedule_plan=False)
+
+        return diffusion_forward_step
+
+    def get_datasets_provider(self):
+        """
+        Return dataset provider function that delegates to self.data_provider.
+
+        This uses the composition pattern: the provider function delegates
+        to the injected data_provider instance, avoiding recreation on each call.
+        """
+
+        def diffusion_datasets_provider(train_val_test_num_samples, vp_stage=None):
+            """Delegate to injected data provider."""
+            from megatron.training import get_args
+
+            args = get_args()
+
+            return self.data_provider.create_dataloaders(
+                trainer_config=args, train_val_test_num_samples=train_val_test_num_samples, vp_stage=vp_stage
+            )
+
+        # Mark as distributed (required by Megatron)
+        diffusion_datasets_provider.is_distributed = self.data_provider.is_distributed
+
+        # Set __module__ to help with debugging (point to this module, not pretrain_gpt)
+        diffusion_datasets_provider.__module__ = __name__
+
+        log_rank_0(
+            f"[DiffusionPretrainTrainer] Created datasets provider: {type(self.data_provider).__name__}"
+        )
+        return diffusion_datasets_provider
+
+    @property
+    def scheduler(self):
+        """Lazily-initialized diffusion scheduler."""
+        if self._scheduler is None:
+            self._scheduler = self.create_scheduler()
+        return self._scheduler
+
+    @abstractmethod
+    def create_scheduler(self):
+        """
+        Create diffusion scheduler.
+
+        Returns:
+            Scheduler instance (e.g., FlowMatchEulerDiscreteScheduler)
+        """
+        raise NotImplementedError("Subclasses must implement create_scheduler()")
+
+    @abstractmethod
+    def get_task_encoder(self):
+        """
+        Get task encoder for Energon data pipeline.
+
+        Returns:
+            TaskEncoder instance (e.g., EncodedDiffusionTaskEncoder)
+
+        This method is called by EnergonDatasetProvider to create the task encoder
+        for processing WebDataset samples.
+        """
+        raise NotImplementedError("Subclasses must implement get_task_encoder()")
diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py
new file mode 100644
index 000000000..36da05ce3
--- /dev/null
+++ b/primus/backends/megatron/flux_pretrain_trainer.py
@@ -0,0 +1,843 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Flux Pretrain Trainer for Primus-Megatron.
+
+This trainer implements Flux-specific training logic including:
+    - Flow matching scheduler with dynamic shifting
+    - Guidance embedding support
+    - Custom forward step function
+"""
+
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer
+from primus.backends.megatron.training.diffusion.schedulers import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from primus.backends.megatron.training.diffusion.timestep_sampling import (
+    create_timestep_sampler,
+)
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _restore_chimera_rng_state(args) -> None:
+    """Restore canonical RNG state after chimera model init.
+
+    Calls Megatron's `_set_random_seed` to restore CPU, CUDA default, and
+    model-parallel tracker generators to a canonical (per-rank-uniform) state.
+    Falls back to a manual restore if Megatron's signature changes.
+
+    Raises:
+        May propagate exceptions other than ImportError/TypeError from
+        _set_random_seed. The fallback handles ImportError (missing module)
+        and TypeError (API signature changes) gracefully.
+    """
+    try:
+        from megatron.training.initialize import _set_random_seed
+
+        _set_random_seed(
+            args.seed,
+            args.data_parallel_random_init,
+            args.te_rng_tracker,
+            args.inference_rng_tracker,
+            use_cudagraphable_rng=getattr(args, "enable_cuda_graph", False),
+        )
+    except (ImportError, TypeError) as e:
+        import logging
+
+        from megatron.core.tensor_parallel import random as tp_random
+
+        torch.manual_seed(args.seed)
+        torch.cuda.manual_seed(args.seed)
+        tp_random.model_parallel_cuda_manual_seed(args.seed)
+        logging.warning(
+            "[nemo_chimera_init] _set_random_seed API changed (%s). "
+            "Restored all RNG generators manually. Verify convergence "
+            "matches expected behavior.",
+            e,
+        )
+
+
+class FluxPretrainTrainer(DiffusionPretrainTrainer):
+    """
+    Trainer for Flux diffusion model pre-training.
+
+    Flux-specific features:
+        - Flow matching with Euler discrete scheduler
+        - Dynamic timestep shifting for variable resolution
+        - Optional guidance embedding for CFG
+
+    Config access via backend_args:
+        - Megatron args (guidance_embed, guidance_scale, etc.) via backend_args
+        - Primus-specific config (torch_compile) via backend_args (overrides section)
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize Flux pretrain trainer.
+
+        Args:
+            *args: Positional arguments passed to parent trainer
+            **kwargs: Keyword arguments passed to parent trainer
+                     (backend_args is extracted from kwargs by BaseTrainer.__init__())
+        """
+        super().__init__(*args, **kwargs)
+
+        self._training_rng_seeded = False
+
+        # backend_args is set by BaseTrainer.__init__()
+        params = self.backend_args
+
+        # Default False: per-step CUDA RNG reseed is an MLPerf-alignment
+        # feature; flipping it on globally changes the observable RNG
+        # sequence for every consumer of this trainer. MLPerf-aligned YAMLs
+        # opt in explicitly via per_step_rng_reseed: true.
+        self.per_step_rng_reseed = getattr(params, "per_step_rng_reseed", False)
+        self.nemo_chimera_init = getattr(params, "nemo_chimera_init", False)
+
+        # These come from the 'overrides' section in the YAML
+        self.use_guidance_embed = getattr(params, "guidance_embed", False)
+        self.guidance_scale = getattr(params, "guidance_scale", 3.5)
+
+        # Scheduler config
+        self.num_train_timesteps = getattr(params, "num_train_timesteps", 1000)
+        self.scheduler_shift = getattr(params, "scheduler_shift", 1.0)
+        self.use_dynamic_shifting = getattr(params, "use_dynamic_shifting", False)
+
+        # Timestep sampling strategy (MLPerf uses "direct_uniform")
+        timestep_strategy = getattr(params, "timestep_sampling_strategy", "logit_normal")
+        self.timestep_sampler = create_timestep_sampler(timestep_strategy)
+        log_rank_0(
+            f"Timestep sampling strategy: {timestep_strategy} -> {type(self.timestep_sampler).__name__}"
+        )
+
+        # CFG dropout: replace text embeddings with fixed empty encodings at this probability
+        self.cfg_dropout_prob = getattr(params, "cfg_dropout_prob", 0.0)
+        self.empty_t5_encodings = None
+        self.empty_clip_encodings = None
+
+        if self.cfg_dropout_prob > 0.0:
+            self._init_cfg_dropout(params)
+
+        # VAE latent normalization (matches NVIDIA MLPerf v5.1)
+        self.vae_scale = getattr(params, "vae_scale", None)
+        self.vae_shift = getattr(params, "vae_shift", None)
+        if self.vae_scale is not None:
+            log_rank_0(f"VAE normalization: scale={self.vae_scale}, shift={self.vae_shift}")
+
+        # VAE latent mode: "presampled" uses stored latents directly,
+        # "resample" re-draws from (mean, logvar) each step
+        self.vae_latent_mode = getattr(params, "vae_latent_mode", "presampled")
+        if self.vae_latent_mode not in ("presampled", "resample"):
+            raise ValueError(
+                f"vae_latent_mode must be 'presampled' or 'resample', " f"got '{self.vae_latent_mode}'"
+            )
+        if self.vae_latent_mode == "resample":
+            if self.vae_scale is None or self.vae_shift is None:
+                raise ValueError(
+                    "vae_latent_mode='resample' requires vae_scale and vae_shift "
+                    "to be set (e.g. vae_scale: 0.3611, vae_shift: 0.1159)"
+                )
+            log_rank_0(f"VAE latent mode: resample (reparameterization from mean+logvar each step)")
+        else:
+            log_rank_0(f"VAE latent mode: presampled (stored latents used directly)")
+
+        log_rank_0(f"Guidance embedding: {self.use_guidance_embed}")
+        log_rank_0(f"Scheduler shift: {self.scheduler_shift}")
+        log_rank_0(f"Dynamic shifting: {self.use_dynamic_shifting}")
+
+    def _init_cfg_dropout(self, params):
+        """
+        Initialize CFG dropout with real empty encodings.
+
+        Resolution order for empty encodings:
+        1. Explicit ``empty_encodings_path`` from config
+        2. ``{data_path}/empty_encodings/`` (generated by EncodedDatasetPipeline)
+        3. ``{data_path}/../empty_encodings/`` (MLPerf convention)
+
+        For mock_data runs, falls back to torch.randn().
+        For real data, raises FileNotFoundError if no encodings are found.
+        """
+        tp_size = getattr(params, "tensor_model_parallel_size", 1)
+        if tp_size != 1:
+            raise ValueError(
+                f"CFG dropout requires tensor_model_parallel_size=1 (got {tp_size}). "
+                "Different TP ranks would generate different dropout masks, causing divergent forward passes."
+            )
+
+        context_dim = getattr(params, "context_dim", 4096)
+        vec_in_dim = getattr(params, "vec_in_dim", 768)
+
+        encodings_dir = self._discover_empty_encodings(params)
+
+        if encodings_dir is not None:
+            t5_path = os.path.join(encodings_dir, "t5_empty.npy")
+            clip_path = os.path.join(encodings_dir, "clip_empty.npy")
+
+            self.empty_t5_encodings = torch.from_numpy(np.load(t5_path))[0].unsqueeze(1)
+            self.empty_clip_encodings = torch.from_numpy(np.load(clip_path))[0]
+
+            log_rank_0(
+                f"CFG dropout: loaded real empty encodings from {encodings_dir}, "
+                f"t5={self.empty_t5_encodings.shape}, clip={self.empty_clip_encodings.shape}"
+            )
+        elif getattr(params, "mock_data", False):
+            image_size = getattr(getattr(params, "mock_dataset", None), "params", None)
+            image_size = getattr(image_size, "image_size", 256) if image_size is not None else 256
+            image_tokens = (image_size // 8 // 2) ** 2
+            total_seq = getattr(params, "seq_length", 512)
+            t5_seq_len = total_seq - image_tokens
+
+            self.empty_t5_encodings = torch.randn(t5_seq_len, 1, context_dim)
+            self.empty_clip_encodings = torch.randn(vec_in_dim)
+            log_rank_0("CFG dropout: using torch.randn() empty encodings (mock_data mode)")
+        else:
+            data_path = getattr(params, "data_path", "<not set>")
+            if isinstance(data_path, list):
+                data_path = data_path[0] if data_path else "<not set>"
+            raise FileNotFoundError(
+                f"CFG dropout requires empty T5/CLIP encodings but none were found.\n"
+                f"Searched locations:\n"
+                f"  1. empty_encodings_path config key (not set)\n"
+                f"  2. {data_path}/empty_encodings/\n"
+                f"  3. {os.path.dirname(str(data_path))}/empty_encodings/\n\n"
+                f"To fix, either:\n"
+                f"  - Re-run dataset preparation with 'primus data diffusion-encoded' (auto-generates them)\n"
+                f"  - Run: python tools/generate_empty_encodings.py --output_dir <path>\n"
+                f"    and set empty_encodings_path in your YAML config"
+            )
+
+        log_rank_0(f"CFG dropout prob: {self.cfg_dropout_prob}")
+
+    @staticmethod
+    def _discover_empty_encodings(params) -> "str | None":
+        """Return the first valid empty_encodings directory, or None."""
+
+        def _has_files(dirpath):
+            return (
+                os.path.isdir(dirpath)
+                and os.path.isfile(os.path.join(dirpath, "t5_empty.npy"))
+                and os.path.isfile(os.path.join(dirpath, "clip_empty.npy"))
+            )
+
+        explicit = getattr(params, "empty_encodings_path", None)
+        if explicit and _has_files(str(explicit)):
+            return str(explicit)
+
+        data_path = getattr(params, "data_path", None)
+        if data_path is not None:
+            if isinstance(data_path, list):
+                data_path = data_path[0] if data_path else None
+            if data_path:
+                data_path = str(data_path)
+                inside = os.path.join(data_path, "empty_encodings")
+                if _has_files(inside):
+                    log_rank_0(f"CFG dropout: auto-discovered empty encodings at {inside}")
+                    return inside
+                alongside = os.path.join(os.path.dirname(data_path), "empty_encodings")
+                if _has_files(alongside):
+                    log_rank_0(f"CFG dropout: auto-discovered empty encodings at {alongside}")
+                    return alongside
+
+        return None
+
+    def forward_step(self, data_iterator, model, return_schedule_plan=False):
+        """
+        Forward step for Flux diffusion training.
+
+        Overrides base class to provide Flux forward step with scheduler and guidance config.
+        The base class implementation handles the new return signature correctly, so we just
+        call super() to use it.
+
+        On the first call, sets a per-DP-rank random seed so that each data-parallel
+        rank independently samples timesteps and noise (matching NeMo's approach).
+
+        Args:
+            data_iterator: Data iterator
+            model: Diffusion model (Flux)
+            return_schedule_plan: Whether to return schedule plan (for pipeline parallelism)
+
+        Returns:
+            Tuple of (output_tensor, loss_func_callable)
+        """
+        if not self._training_rng_seeded:
+            from megatron.core import parallel_state
+            from megatron.training import get_args
+
+            seed = get_args().seed
+            dp_rank = parallel_state.get_data_parallel_rank()
+            per_rank_seed = seed + 100 * dp_rank
+            torch.manual_seed(per_rank_seed)
+            torch.cuda.manual_seed(per_rank_seed)
+            self._training_rng_seeded = True
+            log_rank_0(f"Per-DP-rank training seed: {per_rank_seed} " f"(base={seed}, dp_rank={dp_rank})")
+
+        return super().forward_step(data_iterator, model, return_schedule_plan=return_schedule_plan)
+
+    def create_model(self, pre_process=True, post_process=True):
+        """
+        Create Flux model from YAML configuration.
+
+        Model architecture is loaded from YAML files:
+        - flux_12b.yaml / flux_535m.yaml for layer counts
+        - flux_base.yaml for common Flux architecture
+
+        Optionally loads checkpoint after model creation if configured.
+
+        Args:
+            pre_process: Not used (kept for Megatron model_provider interface compatibility)
+            post_process: Not used (kept for Megatron model_provider interface compatibility)
+
+        Returns:
+            Flux model instance
+        """
+        try:
+            log_rank_0("=" * 80)
+            log_rank_0(f"Creating Flux model from YAML config")
+
+            from megatron.training import get_args
+
+            from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+
+            config = self._build_flux_config_from_yaml()
+
+            # get_args() is safe here since create_model() is called after setup() completes
+            # setup() calls set_primus_global_variables() which initializes _GLOBAL_ARGS
+            args = get_args()
+            # Log complete FluxConfig at DEBUG level (matches Megatron's argument logging)
+            self._log_flux_config(config, args)
+
+            # Set torch_compile settings on args for trainer's apply_torch_compile_if_enabled method
+            # FluxConfig always has these attributes (defined as dataclass fields)
+            torch_compile_attrs = [
+                "enable_torch_compile",
+                "torch_compile_backend",
+                "torch_compile_mode",
+                "torch_compile_fullgraph",
+                "torch_compile_optimizer",
+                "torch_compile_optimizer_scope",
+            ]
+            for attr in torch_compile_attrs:
+                if not hasattr(args, attr):
+                    setattr(args, attr, getattr(config, attr))
+
+            # Backend selection is handled automatically by Flux model based on config.transformer_impl
+            # Pass backend=None to let get_flux_layer_spec() handle backend selection
+            backend = None
+
+            # Log which transformer implementation will be used
+            if config.transformer_impl == "local":
+                log_rank_0("Using local transformer implementation (NO TransformerEngine dependency)")
+            else:
+                log_rank_0("Using TransformerEngine implementation")
+            log_rank_0(
+                "Backend will be selected automatically by Flux model based on config.transformer_impl"
+            )
+
+            # Chimera init: replicate NeMo's per-rank seed contamination where
+            # each DP rank initializes non-parallel weights with a different seed.
+            # The distributed optimizer merges these into a chimera model.
+            if self.nemo_chimera_init:
+                import logging
+
+                from megatron.core import parallel_state
+
+                dp_rank = parallel_state.get_data_parallel_rank()
+                per_rank_seed = args.seed + 100 * dp_rank
+                torch.manual_seed(per_rank_seed)
+                torch.cuda.manual_seed(per_rank_seed)
+                logging.warning(
+                    "[nemo_chimera_init] WARNING: DP weight invariant intentionally broken. "
+                    "Each rank initializes non-parallel layers with a different seed. "
+                    "This replicates NeMo's per-rank seed contamination for convergence "
+                    "parity experiments only. dp_rank=%d, init_seed=%d (base=%d)",
+                    dp_rank,
+                    per_rank_seed,
+                    args.seed,
+                )
+
+            # Create Flux model (backend=None lets model select based on config.transformer_impl)
+            model = Flux(config=config, backend=backend)
+
+            if self.nemo_chimera_init:
+                _restore_chimera_rng_state(args)
+                log_rank_0(
+                    f"[nemo_chimera_init] Restored canonical RNG state (seed={args.seed}) after chimera model init"
+                )
+
+            # Calculate parameters for logging
+            total_params = sum(p.numel() for p in model.parameters())
+
+            log_rank_0(
+                f"Flux model created: {config.num_joint_layers} joint + {config.num_single_layers} single layers"
+            )
+            log_rank_0(f"Total parameters: {total_params / 1e6:.1f}M")
+
+            log_rank_0("=" * 80)
+
+            return model
+        except Exception as e:
+            log_rank_0(f"[ERROR] create_model() failed: {type(e).__name__}: {e}")
+            import traceback
+
+            log_rank_0(f"[ERROR] Traceback:\n{traceback.format_exc()}")
+            raise
+
+    def _build_flux_config_from_yaml(self):
+        """
+        Build FluxConfig from YAML configuration + training args.
+
+        All architectural parameters come from YAML files (merged into Megatron args):
+        - flux_535m.yaml or flux_12b.yaml specifies layer counts
+        - flux_base.yaml provides common architecture parameters
+
+        Training precision and recomputation settings come from backend_args.
+        Torch compile settings come from backend_args.torch_compile (Primus-specific, in overrides section, not in Megatron args).
+
+        Returns:
+            FluxConfig instance with all parameters
+        """
+        from functools import partial
+
+        import torch.nn.functional as F
+
+        from primus.backends.megatron.core.models.diffusion.flux.config import (
+            FluxConfig,
+            erf_gelu,
+        )
+
+        _openai_gelu_fused = partial(F.gelu, approximate="tanh")
+        _openai_gelu_fused.__name__ = "openai_gelu_fused"
+
+        # backend_args is set by BaseTrainer.__init__()
+        params = self.backend_args
+
+        # All architectural params from params (merged from YAML)
+        config_params = {
+            "num_joint_layers": getattr(params, "num_joint_layers", 19),
+            "num_single_layers": getattr(params, "num_single_layers", 38),
+            "hidden_size": getattr(params, "hidden_size", 3072),
+            "num_attention_heads": getattr(params, "num_attention_heads", 24),
+            "in_channels": getattr(params, "in_channels", 64),
+            "context_dim": getattr(params, "context_dim", 4096),
+            "vec_in_dim": getattr(params, "vec_in_dim", 768),
+            "model_channels": getattr(params, "model_channels", 256),
+            "guidance_embed": getattr(params, "guidance_embed", False),
+            # RoPE configuration
+            "apply_rope_fusion": getattr(params, "apply_rope_fusion", False),
+            "rotary_interleaved": getattr(params, "rotary_interleaved", True),
+            # Training hyperparameters stored on FluxConfig
+            "timestep_sampling_strategy": getattr(params, "timestep_sampling_strategy", "logit_normal"),
+            "cfg_dropout_prob": getattr(params, "cfg_dropout_prob", 0.0),
+            # Weight init for Megatron-managed layers (ColumnParallelLinear, RowParallelLinear).
+            # Xavier uniform matches NeMo's CustomFluxConfig (MLPerf v5.1).
+            "init_method": nn.init.xavier_uniform_,
+            "output_layer_init_method": nn.init.xavier_uniform_,
+        }
+
+        # Activation: YAML "openai_gelu" maps to fused F.gelu(approximate="tanh"); default in
+        # FluxConfig is the non-JIT Python GELU (ROCm-safe). Both match the tanh GELU approx.
+        activation_func_name = getattr(params, "activation_func", None)
+        if activation_func_name is not None:
+            activation_func_map = {
+                "erf_gelu": erf_gelu,
+                "openai_gelu": _openai_gelu_fused,
+            }
+            if activation_func_name not in activation_func_map:
+                raise ValueError(
+                    f"Unknown activation_func '{activation_func_name}'. "
+                    f"Choose from: {list(activation_func_map.keys())}"
+                )
+            func = activation_func_map[activation_func_name]
+            config_params["activation_func"] = func
+            log_rank_0(f"Activation function: {activation_func_name} -> {func.__name__}")
+
+        # Transformer implementation (standard TransformerConfig parameter)
+        # Read from backend_args (overrides section)
+        transformer_impl = getattr(params, "transformer_impl", "transformer_engine")
+        config_params["transformer_impl"] = transformer_impl
+
+        # Precision settings from params (not in model YAML)
+        config_params.update(
+            {
+                "bf16": getattr(params, "bf16", True),
+                "fp16": getattr(params, "fp16", False),
+                "params_dtype": getattr(params, "params_dtype", torch.float32),
+            }
+        )
+
+        # FP8 settings from params
+        config_params.update(
+            {
+                "fp8": getattr(params, "fp8", None),
+                "fp8_recipe": getattr(params, "fp8_recipe", "delayed"),
+                "fp8_margin": getattr(params, "fp8_margin", 0),
+                "fp8_amax_history_len": getattr(params, "fp8_amax_history_len", 1),
+                "fp8_amax_compute_algo": getattr(params, "fp8_amax_compute_algo", "most_recent"),
+                "fp8_wgrad": getattr(params, "fp8_wgrad", True),
+                "fp8_dot_product_attention": getattr(params, "fp8_dot_product_attention", False),
+                "fp8_multi_head_attention": getattr(params, "fp8_multi_head_attention", False),
+                "fp8_scaling_strategy": getattr(params, "fp8_scaling_strategy", "dynamic"),
+                "fp8_force_nt_layout": getattr(params, "fp8_force_nt_layout", False),
+                "fp8_reduce_amax": getattr(params, "fp8_reduce_amax", False),
+            }
+        )
+
+        # FP4/MXFP4 settings
+        fp4_enabled = getattr(params, "fp4", None)
+        fp4_recipe = getattr(params, "fp4_recipe", None)
+        if fp4_enabled and not fp4_recipe:
+            raise ValueError(
+                "fp4_recipe must be explicitly set in YAML when fp4 is enabled. "
+                "Use fp4_recipe: 'mxfp4' for AMD (native FP4 GEMM) or 'nvfp4' for NVIDIA."
+            )
+        config_params.update(
+            {
+                "fp4": fp4_enabled,
+                "fp4_recipe": fp4_recipe,
+                "mxfp4_backward_precision": getattr(params, "mxfp4_backward_precision", "mxfp4"),
+            }
+        )
+
+        # Sensitive layer configuration
+        config_params.update(
+            {
+                "sensitive_layers_enabled": getattr(params, "sensitive_layers_enabled", False),
+                "sensitive_layers_start": getattr(params, "sensitive_layers_start", 0),
+                "sensitive_layers_end": getattr(params, "sensitive_layers_end", 0),
+                "sensitive_layer_precision": getattr(params, "sensitive_layer_precision", "bf16"),
+                "mxfp4_gradient_stochastic_rounding": getattr(
+                    params, "mxfp4_gradient_stochastic_rounding", False
+                ),
+            }
+        )
+
+        config_params["use_dual_fp8_output_projection"] = getattr(
+            params,
+            "use_dual_fp8_output_projection",
+            False,
+        )
+
+        if (
+            config_params["use_dual_fp8_output_projection"]
+            and config_params.get("fp8_scaling_strategy") == "delayed"
+        ):
+            raise ValueError(
+                "use_dual_fp8_output_projection=True is incompatible with "
+                "fp8_scaling_strategy='delayed'. DualFP8LinearTensorwiseFunction "
+                "bypasses delayed-scaling staged amax buffers, causing stale/zero "
+                "values in the _DelayedScalingRegistry. Use "
+                "fp8_scaling_strategy='dynamic' or set "
+                "use_dual_fp8_output_projection=False."
+            )
+
+        config_params["use_triton_ops"] = getattr(
+            params,
+            "use_triton_ops",
+            False,
+        )
+        config_params["adaln_plain_ops"] = getattr(
+            params,
+            "adaln_plain_ops",
+            False,
+        )
+        config_params["adaln_always_jit_fuser"] = getattr(
+            params,
+            "adaln_always_jit_fuser",
+            False,
+        )
+        config_params["optimizer_foreach"] = getattr(
+            params,
+            "optimizer_foreach",
+            True,
+        )
+        config_params["overlap_grad_norm"] = getattr(
+            params,
+            "overlap_grad_norm",
+            False,
+        )
+        config_params["use_cpp_fp8_quantize"] = getattr(
+            params,
+            "use_cpp_fp8_quantize",
+            False,
+        )
+
+        # FSDP2 prefetch depth
+        fsdp_prefetch = getattr(params, "fsdp_prefetch_depth", None)
+        if fsdp_prefetch is not None:
+            config_params["fsdp_prefetch_depth"] = int(fsdp_prefetch)
+
+        # FSDP2 FP32 optimizer: initialize model in FP32, FSDP2 casts to BF16
+        if getattr(params, "use_fsdp2_fp32_param_optimizer", False):
+            config_params["params_dtype"] = torch.float32
+            config_params["pipeline_dtype"] = torch.bfloat16
+            log_rank_0("FSDP2 FP32 optimizer: params_dtype=FP32, pipeline=BF16")
+
+        # Recomputation settings from params
+        config_params.update(
+            {
+                "recompute_granularity": getattr(params, "recompute_granularity", None),
+                "recompute_method": getattr(params, "recompute_method", None),
+                "recompute_num_layers": getattr(params, "recompute_num_layers", None),
+                "recompute_modules": getattr(params, "recompute_modules", None),
+            }
+        )
+
+        # Torch compile settings from backend_args (overrides section)
+        # torch_compile is in overrides section in YAML, accessible via backend_args.torch_compile
+        torch_compile_config = getattr(self.backend_args, "torch_compile", None)
+
+        if torch_compile_config is not None:
+            compile_settings = {
+                "enable_torch_compile": getattr(torch_compile_config, "enable", False),
+                "torch_compile_backend": getattr(torch_compile_config, "backend", "inductor"),
+                "torch_compile_mode": getattr(torch_compile_config, "mode", "default"),
+                "torch_compile_fullgraph": getattr(torch_compile_config, "fullgraph", False),
+                "torch_compile_optimizer": getattr(torch_compile_config, "compile_optimizer", False),
+                "torch_compile_optimizer_scope": getattr(
+                    torch_compile_config, "compile_optimizer_scope", "full"
+                ),
+                "torch_compile_strategy": getattr(torch_compile_config, "strategy", "per_block"),
+                "torch_compile_replace_qk_rmsnorm": getattr(
+                    torch_compile_config, "replace_qk_rmsnorm", False
+                ),
+                "torch_compile_disable_inductor_cudagraphs": getattr(
+                    torch_compile_config, "disable_inductor_cudagraphs", True
+                ),
+                "torch_compile_emulate_precision_casts": getattr(
+                    torch_compile_config, "emulate_precision_casts", True
+                ),
+                "torch_compile_fused_ln_modulate": getattr(torch_compile_config, "fused_ln_modulate", True),
+            }
+        else:
+            # Default values if torch_compile section not present
+            compile_settings = {
+                "enable_torch_compile": False,
+                "torch_compile_backend": "inductor",
+                "torch_compile_mode": "default",
+                "torch_compile_fullgraph": False,
+                "torch_compile_optimizer": False,
+                "torch_compile_optimizer_scope": "full",
+                "torch_compile_strategy": "per_block",
+                "torch_compile_replace_qk_rmsnorm": False,
+                "torch_compile_disable_inductor_cudagraphs": True,
+                "torch_compile_emulate_precision_casts": True,
+                "torch_compile_fused_ln_modulate": True,
+            }
+
+        # Set on FluxConfig
+        config_params.update(compile_settings)
+
+        return FluxConfig(**config_params)
+
+    def _log_flux_config(self, config, args):
+        """Print FluxConfig fields (rank 0 only), similar to Megatron argument dumps."""
+        import dataclasses
+
+        # Only log on rank 0
+        if args.rank != 0:
+            return
+
+        # Accumulate the full dump into one multi-line message so the banner
+        # alignment is preserved (log_rank_0 stamps a caller prefix per call).
+        lines = []
+        lines.append("=" * 80)
+        lines.append("FluxConfig (Model Configuration)")
+        lines.append("=" * 80)
+
+        # Organize config parameters by category for better readability
+        categories = {
+            "Model Architecture": [
+                "model_type",
+                "num_joint_layers",
+                "num_single_layers",
+                "num_layers",
+                "hidden_size",
+                "num_attention_heads",
+                "ffn_hidden_size",
+            ],
+            "Input/Output Dimensions": [
+                "in_channels",
+                "out_channels",
+                "patch_size",
+                "context_dim",
+                "vec_in_dim",
+                "model_channels",
+            ],
+            "Position Embeddings (RoPE)": ["theta", "axes_dim", "rotary_interleaved", "apply_rope_fusion"],
+            "Guidance & Diffusion": [
+                "guidance_embed",
+                "guidance_scale",
+                "cfg_dropout_prob",
+                "timestep_sampling_strategy",
+            ],
+            "Attention & Layers": [
+                "add_qkv_bias",
+                "single_block_bias",
+                "attention_dropout",
+                "hidden_dropout",
+                "bias_dropout_fusion",
+            ],
+            "Normalization & Activation": ["activation_func", "layernorm_epsilon", "normalization"],
+            "Precision & Optimization": [
+                "bf16",
+                "fp16",
+                "fp32_residual_connection",
+                "gradient_accumulation_fusion",
+                "use_dual_fp8_output_projection",
+                "params_dtype",
+                "fp8",
+                "fp8_recipe",
+                "fp8_scaling_strategy",
+                "fp8_force_nt_layout",
+                "fp4",
+                "fp4_recipe",
+                "mxfp4_backward_precision",
+                "mxfp4_gradient_stochastic_rounding",
+                "sensitive_layers_enabled",
+                "sensitive_layers_start",
+                "sensitive_layers_end",
+                "sensitive_layer_precision",
+            ],
+            "Recomputation": [
+                "recompute_granularity",
+                "recompute_method",
+                "recompute_num_layers",
+                "recompute_modules",
+            ],
+            "Torch Compile": [
+                "enable_torch_compile",
+                "torch_compile_backend",
+                "torch_compile_mode",
+                "torch_compile_fullgraph",
+                "torch_compile_optimizer",
+                "torch_compile_optimizer_scope",
+                "torch_compile_strategy",
+                "torch_compile_replace_qk_rmsnorm",
+                "torch_compile_disable_inductor_cudagraphs",
+                "torch_compile_emulate_precision_casts",
+                "torch_compile_fused_ln_modulate",
+            ],
+            "Parallelism": [
+                "tensor_model_parallel_size",
+                "pipeline_model_parallel_size",
+                "sequence_parallel",
+                "expert_model_parallel_size",
+            ],
+            "CUDA Graph": ["enable_cuda_graph", "cuda_graph_scope", "cuda_graph_warmup_steps"],
+            "Transformer Engine": ["use_te_rng_tracker"],
+        }
+
+        # Get all dataclass fields
+        all_fields = {f.name for f in dataclasses.fields(config)}
+        categorized_fields = set()
+
+        # Print categorized fields
+        for category, field_names in categories.items():
+            # Check if any fields in this category exist
+            existing_fields = [f for f in field_names if f in all_fields]
+            if not existing_fields:
+                continue
+
+            lines.append(f"\n{category}:")
+            for field_name in existing_fields:
+                if hasattr(config, field_name):
+                    value = getattr(config, field_name)
+                    # Format value (handle callables specially)
+                    if callable(value) and not isinstance(value, type):
+                        value_str = getattr(value, "__name__", str(value))
+                    else:
+                        value_str = str(value)
+
+                    # Create dots for alignment (match Megatron's style: 48 chars)
+                    dots = "." * (48 - len(field_name))
+                    lines.append(f"  {field_name} {dots} {value_str}")
+                    categorized_fields.add(field_name)
+
+        # Print uncategorized fields (any fields not in our category lists)
+        uncategorized = all_fields - categorized_fields
+        if uncategorized:
+            lines.append("\nOther Configuration:")
+            for field_name in sorted(uncategorized):
+                if hasattr(config, field_name):
+                    value = getattr(config, field_name)
+                    if callable(value) and not isinstance(value, type):
+                        value_str = getattr(value, "__name__", str(value))
+                    else:
+                        value_str = str(value)
+                    dots = "." * (48 - len(field_name))
+                    lines.append(f"  {field_name} {dots} {value_str}")
+
+        # Add special note about apply_rope_fusion mismatch
+        if hasattr(config, "apply_rope_fusion") and hasattr(args, "apply_rope_fusion"):
+            if config.apply_rope_fusion != args.apply_rope_fusion:
+                lines.append("\n" + "!" * 80)
+                lines.append("IMPORTANT: FluxConfig vs Megatron Args Mismatch")
+                lines.append("!" * 80)
+                lines.append(f"  FluxConfig.apply_rope_fusion: {config.apply_rope_fusion}")
+                lines.append(f"  Megatron args.apply_rope_fusion: {args.apply_rope_fusion}")
+                lines.append(
+                    f'  → Flux uses FluxConfig value: RoPE fusion is {"ENABLED" if config.apply_rope_fusion else "DISABLED"}'
+                )
+                lines.append("  → Megatron args value is set by position_embedding_type validation")
+                lines.append(
+                    '  → (Megatron sets apply_rope_fusion=False when position_embedding_type != "rope")'
+                )
+                lines.append("!" * 80)
+
+        lines.append("=" * 80)
+        lines.append("End of FluxConfig")
+        lines.append("=" * 80)
+
+        log_rank_0("\n".join(lines))
+
+    def create_scheduler(self):
+        """
+        Create Flow Matching Euler Discrete Scheduler for Flux.
+
+        Returns:
+            FlowMatchEulerDiscreteScheduler instance
+        """
+        log_rank_0("Creating Flow Matching scheduler...")
+        log_rank_0(f"  num_train_timesteps: {self.num_train_timesteps}")
+        log_rank_0(f"  shift: {self.scheduler_shift}")
+        log_rank_0(f"  use_dynamic_shifting: {self.use_dynamic_shifting}")
+
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=self.num_train_timesteps,
+            shift=self.scheduler_shift,
+            use_dynamic_shifting=self.use_dynamic_shifting,
+            base_shift=0.5,  # Flux defaults
+            max_shift=1.15,
+            base_image_seq_len=256,
+            max_image_seq_len=4096,
+        )
+
+        return scheduler
+
+    def get_task_encoder(self):
+        """
+        Get Flux task encoder for Energon data pipeline.
+
+        Returns:
+            EncodedDiffusionTaskEncoder for pre-encoded data
+        """
+        from primus.backends.megatron.data.diffusion.task_encoders import (
+            EncodedDiffusionTaskEncoder,
+        )
+
+        # EnergonDatasetProvider will pass WorkerConfig to Energon directly
+        # TaskEncoder doesn't need WorkerConfig for pre-encoded data
+        task_encoder = EncodedDiffusionTaskEncoder(worker_config=None)
+
+        log_rank_0("Created EncodedDiffusionTaskEncoder for pre-encoded Flux data")
+        return task_encoder
diff --git a/primus/backends/megatron/megatron_adapter.py b/primus/backends/megatron/megatron_adapter.py
index 6f1e58321..698d95e37 100644
--- a/primus/backends/megatron/megatron_adapter.py
+++ b/primus/backends/megatron/megatron_adapter.py
@@ -1,14 +1,17 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
+import importlib
+from typing import Optional
+
 import primus.backends.megatron.patches  # noqa: F401  # Register patches
 from primus.backends.megatron.argument_builder import MegatronArgBuilder
 from primus.core.backend.backend_adapter import BackendAdapter
 from primus.core.backend.backend_registry import BackendRegistry
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class MegatronAdapter(BackendAdapter):
@@ -18,8 +21,27 @@ def __init__(self, framework="megatron"):
         super().__init__(framework)
         self.third_party_dir_name = "Megatron-LM"
 
-    def load_trainer_class(self, stage: str = "pretrain"):
-        """Return the trainer class for the specified training stage."""
+    def load_trainer_class(self, stage: str = "pretrain", trainer_class: Optional[str] = None):
+        """
+        Return the Trainer class for the specified training stage or trainer class name.
+
+        Args:
+            stage: Training stage (e.g., "pretrain", "sft"). Defaults to "pretrain".
+            trainer_class: Optional specific trainer class name for dynamic loading.
+                          If provided, this takes precedence over stage-based selection.
+
+        Returns:
+            Trainer class
+
+        Raises:
+            ImportError: If trainer class cannot be imported
+            ValueError: If stage is invalid or trainer class is not found
+        """
+        # If trainer_class is specified, attempt dynamic loading
+        if trainer_class:
+            return self._load_trainer_class_by_name(trainer_class)
+
+        # Fallback to stage-based selection via the backend registry
         try:
             trainer_cls = BackendRegistry.get_trainer_class(self.framework, stage=stage)
         except (ValueError, AssertionError) as exc:
@@ -31,6 +53,75 @@ def load_trainer_class(self, stage: str = "pretrain"):
         log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer class: {trainer_cls.__name__}")
         return trainer_cls
 
+    def _load_trainer_class_by_name(self, trainer_class: str):
+        """
+        Dynamically load trainer class by name.
+
+        Args:
+            trainer_class: Trainer class name (e.g., "FluxPretrainTrainer", "MegatronPretrainTrainer")
+
+        Returns:
+            Trainer class
+
+        Raises:
+            ImportError: If trainer class cannot be imported
+            ValueError: If trainer class is not found
+        """
+        # Define trainer registry for Megatron backend
+        # This maps trainer class names to their module paths
+        MEGATRON_TRAINERS = {
+            "MegatronPretrainTrainer": "primus.backends.megatron.megatron_pretrain_trainer.MegatronPretrainTrainer",
+            "FluxPretrainTrainer": "primus.backends.megatron.flux_pretrain_trainer.FluxPretrainTrainer",
+        }
+
+        if trainer_class not in MEGATRON_TRAINERS:
+            # Try to load from common locations as fallback
+            log_rank_0(
+                f"[Primus:MegatronAdapter] Trainer '{trainer_class}' not in registry, attempting dynamic import..."
+            )
+
+            possible_paths = [
+                f"primus.backends.megatron.{trainer_class.lower()}.{trainer_class}",
+                f"primus.backends.megatron.{trainer_class}",
+            ]
+
+            for module_path in possible_paths:
+                try:
+                    module_name, class_name = module_path.rsplit(".", 1)
+                    module = importlib.import_module(module_name)
+                    trainer_cls = getattr(module, class_name)
+                    log_rank_0(
+                        f"[Primus:MegatronAdapter] Successfully loaded trainer: {trainer_class} from {module_name}"
+                    )
+                    return trainer_cls
+                except (ImportError, AttributeError):
+                    continue
+
+            # If all attempts failed, provide helpful error message
+            available = list(MEGATRON_TRAINERS.keys())
+            raise ValueError(
+                f"Trainer class '{trainer_class}' not found.\n"
+                f"Available trainers: {', '.join(available) if available else 'none'}\n"
+                f"Hint: Set 'trainer_class' in your experiment YAML to a registered trainer "
+                f"(e.g. trainer_class: FluxPretrainTrainer), register it in MEGATRON_TRAINERS, "
+                f"or ensure it is importable from standard locations."
+            )
+
+        # Load from registry
+        trainer_path = MEGATRON_TRAINERS[trainer_class]
+        module_path, class_name = trainer_path.rsplit(".", 1)
+
+        try:
+            module = importlib.import_module(module_path)
+            trainer_cls = getattr(module, class_name)
+            log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer: {trainer_class} from {module_path}")
+            return trainer_cls
+        except (ImportError, AttributeError) as e:
+            raise ImportError(
+                f"Failed to load trainer class '{trainer_class}' from '{module_path}': {e}\n"
+                f"Hint: Check that the module exists and the class name is correct."
+            ) from e
+
     def detect_backend_version(self) -> str:
         """Detect Megatron-LM version via AST parsing (avoids __init__.py execution)."""
         import ast
diff --git a/primus/backends/megatron/megatron_base_trainer.py b/primus/backends/megatron/megatron_base_trainer.py
index 9f6a84692..ef0a9b484 100644
--- a/primus/backends/megatron/megatron_base_trainer.py
+++ b/primus/backends/megatron/megatron_base_trainer.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -15,7 +15,7 @@
 from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts
 from primus.core.trainer.base_trainer import BaseTrainer
 from primus.core.utils.env import flush_before_hard_exit
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 class MegatronBaseTrainer(BaseTrainer):
@@ -23,9 +23,84 @@ class MegatronBaseTrainer(BaseTrainer):
 
     def setup(self):
         """Setup Megatron runtime: set global vars and patch parse_args."""
+        self._ensure_megatron_path()
         set_primus_global_variables(self.backend_args)
         self._patch_parse_args()
 
+    def _ensure_megatron_path(self):
+        """Ensure Megatron-LM path is in sys.path before any megatron imports."""
+        import os
+        import sys
+        from pathlib import Path
+
+        # First, check if megatron is already importable
+        try:
+            import megatron  # type: ignore
+
+            return  # Path already set correctly
+        except ImportError:
+            pass
+
+        # Try multiple methods to find the Megatron-LM path
+        megatron_paths = []
+
+        # Method 1: From PRIMUS_PATH environment variable (most reliable)
+        primus_path = os.getenv("PRIMUS_PATH")
+        if primus_path:
+            megatron_path = Path(primus_path) / "third_party" / "Megatron-LM"
+            if megatron_path.exists():
+                megatron_paths.append(str(megatron_path))
+
+        # Method 2: From current working directory (works in container)
+        try:
+            cwd = Path.cwd()
+            # Check if we're in /workspace/Primus or a subdirectory
+            if "Primus" in str(cwd):
+                # Find the Primus root
+                primus_root = cwd
+                while primus_root.name != "Primus" and primus_root != primus_root.parent:
+                    primus_root = primus_root.parent
+                if primus_root.name == "Primus":
+                    megatron_path = primus_root / "third_party" / "Megatron-LM"
+                    if megatron_path.exists():
+                        megatron_paths.append(str(megatron_path))
+        except Exception:
+            pass
+
+        # Method 3: From current file location
+        try:
+            repo_root = Path(__file__).resolve().parents[3]
+            megatron_path = repo_root / "third_party" / "Megatron-LM"
+            if megatron_path.exists():
+                megatron_paths.append(str(megatron_path))
+        except Exception:
+            pass
+
+        # Method 4: Check if already in sys.path
+        for path in sys.path:
+            path_obj = Path(path)
+            if path_obj.exists():
+                megatron_pkg = path_obj / "megatron"
+                if megatron_pkg.exists() and megatron_pkg.is_dir():
+                    return  # Path already set correctly
+
+        # Add paths to sys.path if not already present
+        for path in megatron_paths:
+            if path not in sys.path:
+                sys.path.insert(0, path)
+                log_rank_0(f"[Primus:MegatronBaseTrainer] Added Megatron-LM to sys.path: {path}")
+
+        # Verify the path was set correctly by trying to import
+        try:
+            import megatron  # type: ignore
+
+            log_rank_0("[Primus:MegatronBaseTrainer] Successfully verified megatron import")
+        except ImportError as e:
+            log_rank_0(
+                f"[Primus:MegatronBaseTrainer] WARNING: Failed to import megatron after path setup: {e}"
+            )
+            log_rank_0(f"[Primus:MegatronBaseTrainer] sys.path: {sys.path[:5]}")
+
     def init(self):
         """Initialize Megatron training components."""
         log_rank_0("Initializing Megatron training...")
diff --git a/primus/backends/megatron/megatron_pretrain_trainer.py b/primus/backends/megatron/megatron_pretrain_trainer.py
index 5a2e795c4..162aad4ef 100644
--- a/primus/backends/megatron/megatron_pretrain_trainer.py
+++ b/primus/backends/megatron/megatron_pretrain_trainer.py
@@ -1,16 +1,130 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
 from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class MegatronPretrainTrainer(MegatronBaseTrainer):
     """Trainer for Megatron-LM pre-training."""
 
+    def setup_model_only(self):
+        """Initialize Megatron and build the model WITHOUT running the training loop.
+
+        A general, training-neutral capability: mirrors the front of
+        ``megatron.training.pretrain`` (``initialize_megatron`` followed by
+        ``setup_model_and_optimizer``) to construct a model identical to the real
+        training path, but stops before datasets / the train loop. Useful for any
+        "build the model only" scenario (offline profiling, layer benchmarking,
+        model inspection); performance/memory projection is the current consumer.
+
+        Prerequisites (handled by the runtime before calling this): ``setup()``
+        has patched ``parse_args`` to return ``self.backend_args`` and set the
+        Primus global vars, and the build_args/setup/before_train patch phases
+        have been applied. Returns the built model (a list of model chunks, as
+        produced by megatron's ``get_model``) and also stores it on ``self.model``.
+        """
+        log_rank_0("Setting up Megatron model only (no training loop)...")
+
+        from megatron.core.enums import ModelType
+        from megatron.training.initialize import initialize_megatron
+        from megatron.training.training import setup_model_and_optimizer
+
+        from primus.core.utils.import_utils import get_model_provider
+
+        # Determine model type (gpt or mamba) from backend_args
+        model_type = getattr(self.backend_args, "model_type", "gpt")
+        log_rank_0(f"-detected model_type: {model_type}")
+
+        # parse_args was patched in setup() to return backend_args, so
+        # initialize_megatron consumes the Primus-configured arguments (same as
+        # the front of megatron's pretrain()).
+        initialize_megatron(args_defaults={"tokenizer_type": "GPT2BPETokenizer"})
+
+        # Get model provider with correct model_type (reuse the core runtime helper)
+        if model_type != "gpt":
+            model_provider = get_model_provider(model_type=model_type)
+        else:
+            model_provider = get_model_provider()
+        log_rank_0(f"-model_provider: {model_provider}")
+
+        model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
+            model_provider, ModelType.encoder_or_decoder, checkpointing_context={}
+        )
+        self.model = model
+        self.optimizer = optimizer
+        self.opt_param_scheduler = opt_param_scheduler
+
+        log_rank_0("Megatron model-only setup completed.")
+        return model
+
+    def get_forward_step(self):
+        """
+        Return forward step function for training loop.
+
+        Override this method in subclasses to provide custom forward step functions.
+        Default implementation returns GPT forward step.
+
+        Note: This method assumes Megatron-LM path is already set up by
+        MegatronBaseTrainer._ensure_megatron_path() during setup().
+
+        Returns:
+            Callable: Forward step function compatible with Megatron's training loop.
+        """
+        # Path should already be set by _ensure_megatron_path() during setup()
+        try:
+            from pretrain_gpt import forward_step  # type: ignore
+        except ImportError as e:
+            log_rank_0(f"[MegatronPretrainTrainer] Failed to import forward_step from pretrain_gpt: {e}")
+            log_rank_0(
+                "[MegatronPretrainTrainer] This indicates a configuration issue. "
+                "Ensure Megatron-LM is properly installed at third_party/Megatron-LM"
+            )
+            raise ImportError(
+                "Could not import forward_step from pretrain_gpt. "
+                "This indicates that Megatron-LM path setup failed during setup(). "
+                "Ensure Megatron-LM is properly installed and _ensure_megatron_path() succeeded."
+            ) from e
+        return forward_step
+
+    def get_datasets_provider(self):
+        """
+        Return dataset provider function for training.
+
+        Override this method in subclasses to provide custom dataset providers.
+        Default implementation returns GPT dataset provider.
+
+        Note: This method assumes Megatron-LM path is already set up by
+        MegatronBaseTrainer._ensure_megatron_path() during setup().
+
+        Returns:
+            Callable: Dataset provider function compatible with Megatron's training loop.
+        """
+        # Path should already be set by _ensure_megatron_path() during setup()
+        try:
+            from pretrain_gpt import train_valid_test_datasets_provider  # type: ignore
+        except ImportError as e:
+            log_rank_0(
+                f"[MegatronPretrainTrainer] Failed to import train_valid_test_datasets_provider "
+                f"from pretrain_gpt: {e}"
+            )
+            log_rank_0(
+                "[MegatronPretrainTrainer] This indicates a configuration issue. "
+                "Ensure Megatron-LM is properly installed at third_party/Megatron-LM"
+            )
+            raise ImportError(
+                "Could not import train_valid_test_datasets_provider from pretrain_gpt. "
+                "This indicates that Megatron-LM path setup failed during setup(). "
+                "Ensure Megatron-LM is properly installed and _ensure_megatron_path() succeeded."
+            ) from e
+
+        provider = train_valid_test_datasets_provider
+        provider.is_distributed = True  # Always True to match Megatron's behavior
+        return provider
+
     def train(self):
         """Execute Megatron pre-training."""
         log_rank_0("Executing Megatron pretrain...")
@@ -22,11 +136,14 @@ def train(self):
 
         from primus.core.utils.import_utils import get_model_provider
 
-        # Determine model type (gpt or mamba) from backend_args
+        # Determine model type (gpt / mamba / deepseek_v4 / diffusion) from backend_args
         model_type = getattr(self.backend_args, "model_type", "gpt")
         log_rank_0(f"-detected model_type: {model_type}")
 
-        # Import the appropriate training components based on model_type
+        # Import the appropriate training components based on model_type.
+        # DeepSeek-V4 is causal-LM with the same data shape as GPT, so we
+        # reuse pretrain_gpt's forward_step + dataset provider; only the
+        # model_provider itself is V4-specific.
         if model_type == "mamba":
             from pretrain_mamba import (  # type: ignore
                 forward_step,
@@ -34,16 +151,23 @@ def train(self):
             )
 
             log_rank_0("Using Mamba model provider and training components")
-        else:
+            # Upstream pretrain entrypoints set this in their __main__ blocks, but Primus imports the
+            # provider directly and calls pretrain() programmatically. Without restoring this flag,
+            # only TP rank 0 enters dataset construction while the core dataset builder still issues
+            # distributed barriers, which deadlocks for TP>1.
+            train_valid_test_datasets_provider.is_distributed = True
+        elif model_type == "deepseek_v4":
             from pretrain_gpt import (  # type: ignore
                 forward_step,
                 train_valid_test_datasets_provider,
             )
 
-            log_rank_0("Using GPT model provider and training components")
-
-        # Configure training components
-        train_valid_test_datasets_provider.is_distributed = True
+            log_rank_0("Using DeepSeek-V4 model provider; reusing pretrain_gpt forward_step + datasets")
+        else:
+            # Use overridable methods so subclasses (e.g. diffusion/Flux) can plug in their own
+            # forward_step / dataset_provider. Defaults pull from pretrain_gpt.
+            forward_step = self.get_forward_step()
+            train_valid_test_datasets_provider = self.get_datasets_provider()
 
         # Handle Megatron version differences (v0.12.0 vs newer with inprocess_restart)
         wrapped_pretrain = pretrain
@@ -65,12 +189,14 @@ def train(self):
         if "store" in sig.parameters:
             kwargs["store"] = store
 
-        # Get model provider with correct model_type
-        # Only pass model_type if it's not the default to maintain compatibility
-        if model_type != "gpt":
-            model_provider = get_model_provider(model_type=model_type)
-        else:
-            model_provider = get_model_provider()
+        # Resolve model_provider: prefer subclass-provided attribute (e.g. diffusion/Flux),
+        # else fall back to registry-based lookup with model_type for mamba/gpt.
+        model_provider = getattr(self, "model_provider", None)
+        if model_provider is None:
+            if model_type != "gpt":
+                model_provider = get_model_provider(model_type=model_type)
+            else:
+                model_provider = get_model_provider()
         log_rank_0(f"-model_provider: {model_provider}")
 
         # Patch Megatron's get_forward_backward_func to support dump_pp_data
@@ -86,7 +212,7 @@ def patched_get_forward_backward_func(*args, **kwargs):
             try:
                 m_args = get_megatron_args()
                 if getattr(m_args, "dump_pp_data", False):
-                    from primus.modules.trainer.megatron.utils import (
+                    from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
                         schedule_wrapper,
                         set_dump_pp_data_patch,
                     )
@@ -101,13 +227,21 @@ def patched_get_forward_backward_func(*args, **kwargs):
         if hasattr(mt_training, "get_forward_backward_func"):
             mt_training.get_forward_backward_func = patched_get_forward_backward_func
 
-        wrapped_pretrain(
-            train_valid_test_datasets_provider,
-            model_provider,
-            ModelType.encoder_or_decoder,
-            forward_step,
-            **kwargs,
-        )
+        try:
+            wrapped_pretrain(
+                train_valid_test_datasets_provider,
+                model_provider,
+                ModelType.encoder_or_decoder,
+                forward_step,
+                **kwargs,
+            )
+            log_rank_0("[MegatronPretrainTrainer] pretrain() completed successfully")
+        except Exception as e:
+            log_rank_0(f"[MegatronPretrainTrainer] ERROR in pretrain(): {type(e).__name__}: {e}")
+            import traceback
+
+            log_rank_0(f"[MegatronPretrainTrainer] Traceback: {traceback.format_exc()}")
+            raise
 
         # Dump PP visualization data if enabled
         try:
@@ -121,7 +255,9 @@ def patched_get_forward_backward_func(*args, **kwargs):
                     get_num_microbatches,
                 )
 
-                from primus.modules.trainer.megatron.utils import dump_pp_data
+                from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+                    dump_pp_data,
+                )
 
                 pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data")
                 dump_pp_data(megatron_args, get_num_microbatches(), pp_data_dir)
diff --git a/primus/backends/megatron/megatron_sft_trainer.py b/primus/backends/megatron/megatron_sft_trainer.py
index 7a14350b6..46493434f 100644
--- a/primus/backends/megatron/megatron_sft_trainer.py
+++ b/primus/backends/megatron/megatron_sft_trainer.py
@@ -9,7 +9,7 @@
 from typing import Any
 
 from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class MegatronSFTTrainer(MegatronBaseTrainer):
@@ -30,14 +30,15 @@ class MegatronSFTTrainer(MegatronBaseTrainer):
         - Common Megatron initialization patterns
     """
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, **kwargs):
         """
         Initialize Megatron SFT trainer.
 
         Args:
             backend_args: Megatron-LM argument namespace (from MegatronArgBuilder)
+            **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering.
         """
-        super().__init__(backend_args=backend_args)
+        super().__init__(backend_args=backend_args, **kwargs)
 
         # Initialize LoRA if enabled
         self.peft = None
diff --git a/primus/backends/megatron/mlperf/__init__.py b/primus/backends/megatron/mlperf/__init__.py
new file mode 100644
index 000000000..1ddb2d92f
--- /dev/null
+++ b/primus/backends/megatron/mlperf/__init__.py
@@ -0,0 +1,26 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus MLPerf logging integration for the Megatron pretrain backend.
+
+This package integrates the (formerly external) ``primus_mllog`` thin wrapper
+directly into Primus, mapping it onto the new ``BaseTrainer`` lifecycle.
+
+``mlperf_logging`` and other MLPerf-only dependencies are imported lazily
+inside methods so that importing this package (which happens whenever the
+Megatron backend is loaded) never breaks non-MLPerf runs.
+"""
+
+from primus.backends.megatron.mlperf.mlperf_logger import MLPerfLogger, ThroughputTimer
+from primus.backends.megatron.mlperf.mlperf_pretrain_trainer import (
+    MLPerfMegatronPretrainTrainer,
+)
+
+__all__ = [
+    "MLPerfLogger",
+    "ThroughputTimer",
+    "MLPerfMegatronPretrainTrainer",
+]
diff --git a/primus/backends/megatron/mlperf/mlperf_logger.py b/primus/backends/megatron/mlperf/mlperf_logger.py
new file mode 100644
index 000000000..ed60e5de8
--- /dev/null
+++ b/primus/backends/megatron/mlperf/mlperf_logger.py
@@ -0,0 +1,159 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPERF Logging support for Primus Megatron backend."""
+
+import os
+import time
+from typing import Any, Dict, Optional
+
+
+class MLPerfLogger:
+    """Wrapper around mllog library with rank-aware logging."""
+
+    def __init__(self):
+        from mlperf_logging import mllog
+
+        self.mllogger = mllog.get_mllogger()
+        self._configured = False
+        self._rank = None
+        self._save_to_file = os.getenv("MLLOG_SAVE_TO_FILE", "1").lower() not in ("0", "false", "no")
+
+    def configure(self, filepath: str, args) -> Dict[str, Any]:
+        from mlperf_logging import mllog
+
+        if self._save_to_file:
+            mllog.config(filename=filepath, default_stack_offset=3)
+        else:
+            mllog.config(default_stack_offset=3)
+        self._configured = True
+        self._rank = self._get_rank()
+        return self.extract_mlperf_configs(args)
+
+    def _get_rank(self) -> int:
+        import torch
+
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return torch.distributed.get_rank()
+        return int(os.environ.get("RANK", 0))
+
+    def log_event_all_ranks(self, key: str, value: Any, metadata: Optional[Dict] = None):
+        self.mllogger.event(key=key, value=value, metadata=metadata or {})
+
+    def log_event(self, key: str, value: Any, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.event(key=key, value=value, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def log_start(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.start(key=key, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def log_end(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.end(key=key, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def extract_mlperf_configs(self, args) -> Dict[str, Any]:
+        """Extract MLPERF config parameters from Megatron args."""
+        from mlperf_logging.mllog import constants
+
+        data_parallel_size = getattr(args, "data_parallel_size", 1)
+        if data_parallel_size == 0:
+            data_parallel_size = 1
+        micro_batches = args.global_batch_size // (args.micro_batch_size * data_parallel_size)
+
+        eval_samples = args.eval_iters * args.global_batch_size if args.eval_iters > 0 else 1024
+
+        train_samples = getattr(args, "train_samples", None)
+        if not train_samples:
+            train_samples = args.train_iters * args.global_batch_size
+
+        optimizer_name = getattr(args, "optimizer", "adam")
+        if optimizer_name.lower() == "adam":
+            optimizer_name = "adamw"
+        else:
+            optimizer_name = optimizer_name.lower()
+
+        configs = {
+            constants.GLOBAL_BATCH_SIZE: args.global_batch_size,
+            constants.GRADIENT_ACCUMULATION_STEPS: micro_batches,
+            "max_sequence_length": args.seq_length,
+            constants.TRAIN_SAMPLES: train_samples,
+            constants.EVAL_SAMPLES: eval_samples,
+            constants.SEED: args.seed,
+            "init_checkpoint_step": args.iteration,
+            constants.OPT_NAME: optimizer_name,
+            constants.OPT_BASE_LR: args.lr,
+            constants.OPT_ADAMW_BETA_1: getattr(args, "adam_beta1", 0.9),
+            constants.OPT_ADAMW_BETA_2: getattr(args, "adam_beta2", 0.999),
+            constants.OPT_ADAMW_EPSILON: getattr(args, "adam_eps", 1e-8),
+            constants.OPT_ADAMW_WEIGHT_DECAY: args.weight_decay,
+            "opt_gradient_clip_norm": getattr(args, "clip_grad", 1.0),
+            "opt_end_learning_rate": args.min_lr,
+            "opt_learning_rate_warmup_steps": getattr(args, "lr_warmup_iters", 0),
+            "opt_learning_rate_decay_steps": args.lr_decay_iters if args.lr_decay_iters else args.train_iters,
+            "opt_learning_rate_decay_schedule": self._get_lr_schedule_name(args.lr_decay_style),
+            "max_steps": args.train_iters,
+            constants.SUBMISSION_BENCHMARK: os.getenv("MLLOG_SUBMISSION_BENCHMARK", ""),
+            constants.SUBMISSION_DIVISION: os.getenv("MLLOG_SUBMISSION_DIVISION", ""),
+            constants.SUBMISSION_STATUS: os.getenv("MLLOG_SUBMISSION_STATUS", ""),
+            constants.SUBMISSION_ORG: os.getenv("MLLOG_SUBMISSION_ORG", ""),
+            constants.SUBMISSION_PLATFORM: os.getenv("MLLOG_SUBMISSION_PLATFORM", ""),
+            constants.TENSOR_PARALLELISM: int(os.getenv("MLLOG_TENSOR_PARALLELISM", 1)),
+            constants.PIPELINE_PARALLELISM: int(os.getenv("MLLOG_PIPELINE_PARALLELISM", 1)),
+            constants.CONTEXT_PARALLELISM: int(os.getenv("MLLOG_CONTEXT_PARALLELISM", 1)),
+            constants.EXPERT_PARALLELISM: int(os.getenv("MLLOG_EXPERT_PARALLELISM", 1)),
+            constants.MICRO_BATCH_SIZE: int(os.getenv("MLLOG_MICRO_BATCH_SIZE", 1)),
+            constants.CONFIG_FILENAME: os.getenv("MLLOG_CONFIG_FILENAME", ""),
+            "lowest_numerical_precision_linear": os.getenv("MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR", ""),
+        }
+        return configs
+
+    def _get_lr_schedule_name(self, style: str) -> str:
+        mapping = {
+            "cosine": "cosine with linear warmup",
+            "linear": "linear",
+            "constant": "constant",
+        }
+        return mapping.get(style, style)
+
+
+class ThroughputTimer:
+    """Per-step accumulation timer for training throughput (samples/sec).
+
+    Measures only the time inside each training step, excluding eval time
+    and inter-step overhead (dataloader, callbacks, Python gaps).  This
+    matches the approach used by nemo's Timer class.
+    """
+
+    def __init__(self, global_batch_size: int):
+        self.gbs = global_batch_size
+        self._step_start_time = None
+        self._accumulated_time = 0.0
+        self._accumulated_samples = 0
+
+    def step_start(self):
+        """Mark the beginning of a training step."""
+        self._step_start_time = time.time()
+
+    def step_stop(self):
+        """Mark the end of a training step and accumulate elapsed time."""
+        if self._step_start_time is not None:
+            self._accumulated_time += time.time() - self._step_start_time
+            self._accumulated_samples += self.gbs
+            self._step_start_time = None
+
+    def get_throughput(self) -> Optional[float]:
+        """Return throughput for the accumulated steps and reset counters.
+
+        Returns:
+            Samples per second, or None if no steps were accumulated.
+        """
+        if self._accumulated_time <= 0:
+            return None
+        throughput = self._accumulated_samples / self._accumulated_time
+        self._accumulated_time = 0.0
+        self._accumulated_samples = 0
+        return throughput
diff --git a/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py
new file mode 100644
index 000000000..7d6194f01
--- /dev/null
+++ b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py
@@ -0,0 +1,421 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPerf-enabled Megatron pretrain trainer for the Primus new architecture.
+
+Unlike the legacy ``primus_mllog.MLPerfMegatronPretrainTrainer`` (which
+subclassed the *old* Primus trainer and overrode Primus-owned ``init`` /
+``run`` / ``train`` / ``train_step`` methods), the new architecture delegates
+the entire training loop to upstream ``megatron.training.pretrain()`` inside
+:meth:`MegatronPretrainTrainer.train`.  There are therefore no Primus-owned
+loop methods to override.
+
+This trainer maps the MLPerf logging onto the new ``BaseTrainer`` lifecycle
+(``setup -> init -> train -> cleanup``) by **monkey-patching the upstream
+Megatron functions** that the loop calls:
+
+    * ``megatron.training.training.train``                 (loop entry/exit)
+    * ``megatron.training.training.train_step``            (per-step hooks)
+    * ``megatron.training.training.evaluate``              (capture val loss)
+    * ``megatron.training.training.evaluate_and_print_results`` (eval markers)
+
+All ``mlperf_logging`` imports are lazy so importing this module never breaks
+non-MLPerf runs.
+"""
+
+import os
+import time
+from typing import Any
+
+from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer
+from primus.backends.megatron.mlperf.mlperf_logger import MLPerfLogger, ThroughputTimer
+
+try:
+    from rpdTracerControl import rpdTracerControl
+except ImportError:
+    rpdTracerControl = None
+
+
+def _get_arg(args, kwargs, index, name):
+    """Fetch a positional-or-keyword argument from a wrapped call.
+
+    Upstream Megatron calls ``train`` / ``train_step`` positionally, but we
+    accept either form so the patch is robust across signature tweaks.
+    """
+    if name in kwargs:
+        return kwargs[name]
+    if index < len(args):
+        return args[index]
+    return None
+
+
+class MLPerfMegatronPretrainTrainer(MegatronPretrainTrainer):
+    """MegatronPretrainTrainer with MLPerf (mllog) logging."""
+
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # The core runtime instantiates every trainer with BaseModule-style
+        # context kwargs (module_name, primus_config, module_rank, ...). Accept
+        # and forward them so BaseTrainer can filter them cooperatively.
+        super().__init__(backend_args=backend_args, **kwargs)
+
+        self.mllogger = MLPerfLogger()
+        self.throughput_timer = None
+        self.train_start_time = None
+        self.train_stop_time = None
+        self.last_validation_loss = None
+        self.train_loss_log_freq = int(os.getenv("MLLOG_TRAIN_LOSS_LOG_FREQ", "1"))
+        self.block_tput_log = os.getenv("MLLOG_BLOCK_TPUT_LOG", "0") == "1"
+        self.target_eval_loss = float(os.getenv("MLLOG_TARGET_EVAL_LOSS", "0.0"))
+        self.is_target_reached = False
+        self._warmup_active = False
+
+        # RPD profiler (optional). Mirrors the legacy wrapper.
+        profiler = os.getenv("PROFILER", "")
+        self.profile_segment = os.getenv("PROFILE_SEGMENT", "train")
+        self.rpd = None
+        if profiler == "rpd" and rpdTracerControl is not None:
+            rpdTracerControl.setFilename(name="trace.rpd", append=True)
+            self.rpd = rpdTracerControl()
+            enable_python_trace = os.getenv("ENABLE_PYTHON_TRACE", "false").strip().lower() in {
+                "1",
+                "true",
+                "yes",
+                "on",
+            }
+            self.rpd.setPythonTrace(doTrace=enable_python_trace)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def setup(self):
+        """Megatron setup + MLPerf CACHE_CLEAR / INIT_START.
+
+        ``setup()`` runs before ``init()`` and before ``train()`` (which
+        calls upstream ``pretrain()`` that performs the heavy init), so this
+        is the correct place to bracket the start of initialisation.
+        """
+        super().setup()
+
+        # Primus reconfigures stdlib ``logging`` via loguru during setup,
+        # which resets the per-logger levels installed by the log-suppression
+        # module at import time. Re-apply the quiet levels (no-op unless
+        # PRIMUS_LOG_SUPPRESSION=1 and not verbose).
+        try:
+            from primus.mlperf_log_suppression import reapply_quiet_logger_levels
+
+            reapply_quiet_logger_levels()
+        except Exception:
+            pass
+
+        from mlperf_logging.mllog import constants
+
+        self.mllogger.log_event_all_ranks(key=constants.CACHE_CLEAR, value=True)
+        self.mllogger.log_event_all_ranks(key=constants.INIT_START, value=None)
+
+    def train(self):
+        """Install MLPerf hooks on upstream Megatron, then run pretrain()."""
+        import megatron.training.training as mt
+
+        orig_train = mt.train
+        orig_train_step = mt.train_step
+        orig_evaluate = mt.evaluate
+        orig_eval_and_print = mt.evaluate_and_print_results
+
+        wrapped_train = self._make_wrapped_train(mt, orig_train)
+        wrapped_train_step = self._make_wrapped_train_step(orig_train_step)
+        wrapped_evaluate = self._make_wrapped_evaluate(orig_evaluate)
+        wrapped_eval_and_print = self._make_wrapped_eval_and_print(orig_eval_and_print)
+
+        mt.train = wrapped_train
+        mt.train_step = wrapped_train_step
+        mt.evaluate = wrapped_evaluate
+        mt.evaluate_and_print_results = wrapped_eval_and_print
+
+        try:
+            return super().train()
+        finally:
+            mt.train = orig_train
+            mt.train_step = orig_train_step
+            mt.evaluate = orig_evaluate
+            mt.evaluate_and_print_results = orig_eval_and_print
+
+    # ------------------------------------------------------------------
+    # Patch builders
+    # ------------------------------------------------------------------
+
+    def _make_wrapped_train(self, mt, orig_train):
+        def wrapped_train(*args, **kwargs):
+            from megatron.training import get_args
+            from mlperf_logging.mllog import constants
+
+            from primus.backends.megatron.mlperf.warmup import run_synthetic_warmup
+
+            megatron_args = get_args()
+
+            # Configure mllog + log hyper-parameters + INIT_STOP. ``get_args``
+            # is valid here because upstream pretrain() has finished init by
+            # the time it calls train().
+            output_file = os.getenv("MLLOG_OUTPUT_FILE", "/results/mlperf_output.log")
+            configs = self.mllogger.configure(output_file, megatron_args)
+            for key, value in configs.items():
+                self.mllogger.log_event(key=key, value=value)
+            self.mllogger.log_end(key=constants.INIT_STOP)
+
+            self.throughput_timer = ThroughputTimer(megatron_args.global_batch_size)
+
+            forward_step_func = _get_arg(args, kwargs, 0, "forward_step_func")
+            model = _get_arg(args, kwargs, 1, "model")
+            optimizer = _get_arg(args, kwargs, 2, "optimizer")
+            opt_param_scheduler = _get_arg(args, kwargs, 3, "opt_param_scheduler")
+            config = _get_arg(args, kwargs, 7, "config")
+
+            # Synthetic warmup runs BEFORE RUN_START so its time is excluded
+            # from the timed run. It calls the (patched) module-level
+            # train_step; _warmup_active suppresses per-step loss logging.
+            self._warmup_active = True
+            try:
+                run_synthetic_warmup(
+                    mt.train_step,
+                    forward_step_func,
+                    model,
+                    optimizer,
+                    opt_param_scheduler,
+                    config,
+                    megatron_args,
+                )
+            finally:
+                self._warmup_active = False
+            # Discard warmup step timings.
+            self.throughput_timer.get_throughput()
+
+            self.mllogger.log_start(key=constants.RUN_START)
+            self.train_start_time = time.time()
+            self.mllogger.log_start(key=constants.EPOCH_START, metadata={constants.SAMPLES_COUNT: 0})
+            self.mllogger.log_start(key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: 0})
+
+            if self.rpd and self.profile_segment == "train":
+                self.rpd.start()
+
+            result = orig_train(*args, **kwargs)
+
+            if self.rpd and self.profile_segment == "train":
+                self.rpd.stop()
+                self.rpd = None
+
+            final_args = get_args()
+            consumed_samples = final_args.consumed_train_samples
+            self.mllogger.log_end(
+                key=constants.EPOCH_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples}
+            )
+
+            if self.is_target_reached:
+                self.mllogger.log_end(
+                    key=constants.RUN_STOP,
+                    metadata={
+                        constants.SAMPLES_COUNT: consumed_samples,
+                        constants.STATUS: constants.SUCCESS,
+                    },
+                )
+            else:
+                self.train_stop_time = time.time()
+                self.mllogger.log_end(
+                    key=constants.RUN_STOP,
+                    metadata={
+                        constants.SAMPLES_COUNT: consumed_samples,
+                        constants.STATUS: constants.ABORTED,
+                    },
+                )
+
+            # Overall run summary (rank-0 only; informational keys).
+            duration = (self.train_stop_time or time.time()) - (self.train_start_time or time.time())
+            duration_minutes = duration / 60.0
+            overall_throughput = consumed_samples / duration if duration > 0 else 0.0
+            self.mllogger.log_event(
+                key="run_duration",
+                value=f"{round(duration, 2)}s -> {round(duration_minutes, 2)} minutes",
+                metadata={"samples": consumed_samples},
+            )
+            self.mllogger.log_event(
+                key="overall_throughput",
+                value=round(overall_throughput, 2),
+                metadata={"samples": consumed_samples},
+            )
+
+            return result
+
+        return wrapped_train
+
+    def _make_wrapped_train_step(self, orig_train_step):
+        def wrapped_train_step(*args, **kwargs):
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "train_step":
+                self.rpd.start()
+
+            if self.throughput_timer is not None:
+                self.throughput_timer.step_start()
+            result = orig_train_step(*args, **kwargs)
+            if self.throughput_timer is not None:
+                self.throughput_timer.step_stop()
+
+            # Upstream train_step returns an 8-tuple:
+            # (loss_dict, skipped_iter, should_checkpoint, should_exit,
+            #  exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit)
+            loss_dict = result[0]
+            skipped_iter = result[1]
+
+            if not skipped_iter and loss_dict and not self._warmup_active:
+                optimizer = _get_arg(args, kwargs, 3, "optimizer")
+                # Upstream passes the real loop iteration as the 8th positional
+                # (index 7) / ``iteration`` kwarg of train_step. ``args.iteration``
+                # is NOT updated until after train_step returns, so use this.
+                iteration = _get_arg(args, kwargs, 7, "iteration")
+                self._on_train_step_end(loss_dict, optimizer, iteration)
+
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "train_step":
+                self.rpd.stop()
+                self.rpd = None
+
+            return result
+
+        return wrapped_train_step
+
+    def _make_wrapped_evaluate(self, orig_evaluate):
+        def wrapped_evaluate(*args, **kwargs):
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval":
+                self.rpd.start()
+
+            result = orig_evaluate(*args, **kwargs)
+
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval":
+                self.rpd.stop()
+                self.rpd = None
+
+            total_loss_dict = result[0]
+            if total_loss_dict and "lm loss" in total_loss_dict:
+                val_loss = total_loss_dict["lm loss"]
+                if hasattr(val_loss, "item"):
+                    self.last_validation_loss = val_loss.item()
+                else:
+                    self.last_validation_loss = float(val_loss)
+
+            return result
+
+        return wrapped_evaluate
+
+    def _make_wrapped_eval_and_print(self, orig_eval_and_print):
+        def wrapped_eval_and_print(*args, **kwargs):
+            from megatron.training import get_args
+            from mlperf_logging.mllog import constants
+
+            eval_args = get_args()
+            consumed_samples = eval_args.consumed_train_samples
+
+            self._log_tracked_stats(eval_args.iteration, consumed_samples)
+
+            self.mllogger.log_end(
+                key=constants.BLOCK_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples}
+            )
+            self.mllogger.log_start(
+                key=constants.EVAL_START, metadata={constants.SAMPLES_COUNT: consumed_samples}
+            )
+
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval":
+                self.rpd.start()
+
+            result = orig_eval_and_print(*args, **kwargs)
+
+            if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval":
+                self.rpd.stop()
+                self.rpd = None
+
+            validation_loss = self._get_validation_loss()
+            if validation_loss is not None:
+                self.mllogger.log_event(
+                    key=constants.EVAL_ACCURACY,
+                    value=validation_loss,
+                    metadata={constants.SAMPLES_COUNT: consumed_samples},
+                )
+
+                if self.target_eval_loss > 0.0 and validation_loss <= self.target_eval_loss:
+                    if not self.is_target_reached:
+                        self.is_target_reached = True
+                        self.train_stop_time = time.time()
+                        if self.mllogger._get_rank() == 0:
+                            print(
+                                f"[MLPERF] Target eval loss {self.target_eval_loss} reached "
+                                f"with validation loss {validation_loss}"
+                            )
+                        eval_args.train_iters = eval_args.iteration
+                        eval_args.do_valid = False
+                        eval_args.do_test = False
+
+            self.mllogger.log_end(
+                key=constants.EVAL_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples}
+            )
+
+            if not self.is_target_reached:
+                self.mllogger.log_start(
+                    key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: consumed_samples}
+                )
+
+            return result
+
+        return wrapped_eval_and_print
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _on_train_step_end(self, loss_dict, optimizer, iteration=None):
+        from megatron.training import get_args
+        from mlperf_logging.mllog import constants
+
+        args = get_args()
+        # Prefer the iteration passed into train_step (real loop counter); fall
+        # back to args.iteration only if it was not provided. args.iteration is
+        # stale during train_step, which made ``% freq`` true every step.
+        if iteration is None:
+            iteration = args.iteration
+        consumed_samples = args.consumed_train_samples
+
+        if self.train_loss_log_freq <= 0 or iteration % self.train_loss_log_freq != 0:
+            return
+
+        loss_value = loss_dict.get("lm loss")
+        if loss_value is None:
+            return
+
+        if isinstance(loss_value, (tuple, list)):
+            loss_value = loss_value[0]
+        if hasattr(loss_value, "item"):
+            loss_value = loss_value.item()
+
+        learning_rate = None
+        if optimizer is not None:
+            for param_group in optimizer.param_groups:
+                if not param_group.get("is_decoupled_lr", False):
+                    learning_rate = param_group["lr"]
+                    break
+
+        self.mllogger.log_event(
+            key="train_loss",
+            value=loss_value,
+            metadata={constants.SAMPLES_COUNT: consumed_samples, "lr": learning_rate},
+        )
+
+    def _log_tracked_stats(self, iteration, consumed_samples):
+        if self.throughput_timer is None:
+            return
+        throughput = self.throughput_timer.get_throughput()
+        if throughput is not None and self.block_tput_log:
+            self.mllogger.log_event(
+                key="tracked_stats",
+                value={"throughput": throughput},
+                metadata={"step": consumed_samples},
+            )
+
+    def _get_validation_loss(self):
+        return getattr(self, "last_validation_loss", None)
diff --git a/primus/backends/megatron/mlperf/warmup.py b/primus/backends/megatron/mlperf/warmup.py
new file mode 100644
index 000000000..f136ceee0
--- /dev/null
+++ b/primus/backends/megatron/mlperf/warmup.py
@@ -0,0 +1,776 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Synthetic-data warmup for Primus / Megatron training.
+
+Adds optional FP4 (MXFP4 / NVFP4) warmup support on top of the original FP8
+warmup helpers.  The original FP8 / BF16 path is preserved exactly:
+
+  * No autocast wrapping by default — ``train_step`` runs at the model's
+    native precision, identical to the committed pre-FP4 behaviour.
+  * After warmup, ``reset_fp8_state`` + ``seed_fp8_amax(1.0)`` always runs
+    (no-op for BF16 / non-FP8 modules).
+  * ``ChainedOptimizer`` support (EP < num_gpus) is preserved.
+
+FP4 mode is **opt-in** via ``WARMUP_RECIPE=fp4_mxfp4`` / ``fp4_nvfp4``.  When
+set, the warmup loop is wrapped in ``te.fp8_autocast(fp8_recipe=...)`` and
+an additional ``reset_fp4_state()`` pass runs after warmup.
+
+Runs N forward+backward passes with random token sequences before the real
+training loop starts.  This pre-compiles Triton / CK / hipBLASLt kernels and
+amortizes:
+
+  * distributed-optimizer FP32 main-param allocation
+  * DDP gradient-bucket allocation
+  * NCCL communicator init + collective autotune
+  * hipBLASLt heuristic / Triton / CK JIT caches (recipe-dependent)
+  * PyTorch HIP allocator block layout
+
+After warmup:
+  * Model parameters are restored from a pre-warmup snapshot.
+  * Optimizer state is restored (neutered during warmup so weights never move).
+  * Adam buffers (exp_avg / exp_avg_sq) are zeroed and grads cleared.
+  * LR scheduler state is rolled back and re-synced (param_groups['lr']).
+  * FP8 scaling state (amax_history, scale, scale_inv, fp8_initialized) is
+    fully reset and seeded with safe defaults.
+  * If WARMUP_RECIPE selected an FP4 recipe, FP4 state is also reset.
+
+Controlled by:
+
+  SYNTH_WARMUP_STEPS         default 3,  0 disables.
+  WARMUP_RECIPE              default "" (no autocast, original behaviour),
+                             one of:
+                               "" | "bf16"
+                               | "fp8_hybrid" | "fp8_e4m3"
+                               | "fp4_mxfp4"  | "fp4_nvfp4"
+                             Legacy alias WARMUP_FP8_RECIPE=hybrid|e4m3 also
+                             accepted (maps to fp8_hybrid / fp8_e4m3).
+                             FP8/BF16 models do NOT need to set this; it
+                             only affects which autocast (if any) wraps the
+                             warmup train_step.
+  WARMUP_FP8_HISTORY_LEN     default 4, only used for fp8_* recipes.
+  SYNTH_WARMUP_EMPTY_CACHE   default 1, call torch.cuda.empty_cache() per step.
+"""
+
+import os
+import time
+from contextlib import nullcontext
+
+import torch
+import torch.distributed
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+
+
+def _log(msg):
+    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+    if rank == 0:
+        print(f"[SYNTH_WARMUP] {msg}", flush=True)
+
+
+# ---------------------------------------------------------------------------
+# Synthetic data
+# ---------------------------------------------------------------------------
+
+
+class SyntheticGPTDataIterator:
+    """Infinite iterator yielding random token batches for Primus/Megatron GPT.
+
+    Primus's ``get_batch_on_this_tp_rank`` expects the iterator to yield a dict
+    with ``tokens``, ``labels``, ``loss_mask``, and ``position_ids`` tensors,
+    each of shape ``[mbs, seq_length]``.
+    """
+
+    def __init__(self, seq_length, micro_batch_size, vocab_size=32000):
+        self.seq_length = seq_length
+        self.micro_batch_size = micro_batch_size
+        self.vocab_size = vocab_size
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        mbs = self.micro_batch_size
+        sl = self.seq_length
+        tokens = torch.randint(0, self.vocab_size, (mbs, sl), dtype=torch.int64)
+        labels = torch.randint(0, self.vocab_size, (mbs, sl), dtype=torch.int64)
+        loss_mask = torch.ones(mbs, sl, dtype=torch.float32)
+        position_ids = torch.arange(sl, dtype=torch.int64).unsqueeze(0).expand(mbs, -1)
+        return {
+            "tokens": tokens,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "position_ids": position_ids,
+        }
+
+
+# ---------------------------------------------------------------------------
+# FP8 state management  --  ORIGINAL from committed warmup.py.
+#
+# The only change vs. the committed version is one defensive line at the top
+# of ``_is_delayed_scaling_recipe`` that returns False for FP4-flagged
+# modules so the FP8 reset path skips them (handled by reset_fp4_state()
+# instead).  No behavioural change for FP8 / BF16 / current-scaling models.
+# ---------------------------------------------------------------------------
+
+
+def _is_delayed_scaling_recipe(module):
+    """Return True if the module uses delayed scaling (which has scale/amax_history).
+
+    Current scaling (tensorwise) is stateless — no scale or amax_history to
+    manage — so reset / seed operations must be skipped for it.
+    FP4 (MXFP4 / NVFP4) modules also have no scale/amax_history; they are
+    skipped here and handled by reset_fp4_state() if WARMUP_RECIPE is fp4_*.
+    Returns True when recipe type cannot be determined (safe default).
+    """
+    # Defensive: FP4 modules use a different recipe state class with no
+    # scale/amax_history. Calling FP8 helpers on them is at best a no-op
+    # and at worst raises AttributeError.
+    if bool(getattr(module, "fp4", False)) or hasattr(module, "fp4_initialized"):
+        return False
+
+    from transformer_engine.common.recipe import DelayedScaling
+
+    fp8_meta = getattr(module, "fp8_meta", None)
+    if fp8_meta is None:
+        return True
+    fwd_state = fp8_meta.get("scaling_fwd", None)
+    if fwd_state is None:
+        return True
+    return isinstance(getattr(fwd_state, "recipe", None), DelayedScaling)
+
+
+def _manual_reset_fp8_meta(module):
+    """Reset FP8 amax / scale tensors when reset_fp8_meta_tensors() is absent."""
+    if not hasattr(module, "fp8_meta"):
+        return False
+    if not _is_delayed_scaling_recipe(module):
+        return False
+    meta = module.fp8_meta
+    reset_count = 0
+    for key in ("scaling_fwd", "scaling_bwd"):
+        if key not in meta:
+            continue
+        tensor_meta = meta[key]
+        if hasattr(tensor_meta, "amax_history"):
+            tensor_meta.amax_history.fill_(0.0)
+            reset_count += 1
+        if hasattr(tensor_meta, "scale"):
+            tensor_meta.scale.fill_(1.0)
+            reset_count += 1
+        if hasattr(tensor_meta, "scale_inv"):
+            tensor_meta.scale_inv.fill_(1.0)
+            reset_count += 1
+    return reset_count > 0
+
+
+def reset_fp8_state(model, reset_meta_tensors=True):
+    """Clear ``fp8_initialized`` on every TE layer, forcing re-init.
+
+    When *reset_meta_tensors* is True, also zeros out amax_history and
+    resets scale / scale_inv to 1.0.
+
+    Skips FP4 modules (handled by reset_fp4_state() when WARMUP_RECIPE=fp4_*).
+    """
+    count = 0
+    method_count = 0
+    manual_count = 0
+
+    def _reset(m):
+        nonlocal count, method_count, manual_count
+        if hasattr(m, "fp8_initialized"):
+            m.fp8_initialized = False
+            count += 1
+            if reset_meta_tensors:
+                if hasattr(m, "reset_fp8_meta_tensors"):
+                    if _is_delayed_scaling_recipe(m):
+                        m.reset_fp8_meta_tensors()
+                        method_count += 1
+                elif _manual_reset_fp8_meta(m):
+                    manual_count += 1
+
+    models = model if isinstance(model, (list, tuple)) else [model]
+    for m in models:
+        m.apply(_reset)
+    _log(f"reset_fp8_state: {count} modules, " f"{method_count} via method, {manual_count} via manual reset")
+    return count
+
+
+def seed_fp8_amax(model, seed_value=1.0):
+    """Fill amax_history with *seed_value* to prevent scale=inf after reset.
+
+    Delayed scaling computes ``scale = fp8_max / max(amax_history)``.
+    If amax_history is all-zero after reset, scale becomes inf -> NaN.
+    Seeding with 1.0 gives scale = fp8_max / 1.0 ~ 448 (E4M3), which is safe.
+    """
+    count = 0
+
+    def _seed(m):
+        nonlocal count
+        if not hasattr(m, "fp8_meta"):
+            return
+        if not _is_delayed_scaling_recipe(m):
+            return
+        meta = m.fp8_meta
+        for key in ("scaling_fwd", "scaling_bwd"):
+            if key not in meta:
+                continue
+            tensor_meta = meta[key]
+            if hasattr(tensor_meta, "amax_history"):
+                tensor_meta.amax_history.fill_(seed_value)
+                count += 1
+
+    models = model if isinstance(model, (list, tuple)) else [model]
+    for m in models:
+        m.apply(_seed)
+    _log(f"seed_fp8_amax: seeded {count} amax_history tensors with {seed_value}")
+    return count
+
+
+# ===========================================================================
+# FP4 state management  --  NEW, only invoked when WARMUP_RECIPE=fp4_*
+# ===========================================================================
+#
+# Notes on FP4 vs FP8 in TE:
+#   * MXFP4BlockScaling and NVFP4BlockScaling derive scales **per tile** from
+#     the current activation/weight on every forward — there is no
+#     amax_history and no global scale-inv to bias.  So no seeding step.
+#   * Some TE versions store FP4 quantizer caches under fp8_meta as well
+#     (under attrs like ``block_scales``, ``mxfp4_quantizer``, etc.).  We
+#     try to clear those defensively; they're allowed to be missing.
+#   * Different TE versions may track init via ``fp4_initialized``,
+#     ``fp8_initialized`` (shared with FP8), or both.  We touch whichever
+#     attribute exists.
+
+# Per-tensor-meta attribute names that may hold FP4-specific cached state.
+# Cleared opportunistically if present; missing attrs are silently skipped.
+_FP4_META_ATTRS = (
+    "block_scales",
+    "block_scales_inv",
+    "mxfp4_block_scale",
+    "fp4_quantizer_state",
+    "fp4_amax",
+)
+
+
+def _is_fp4_module(m):
+    """True if a TE module is configured for FP4 (MXFP4 / NVFP4).
+
+    FP4 modules in TE share ``fp8_initialized`` and ``fp8_meta`` with FP8 but
+    populate them with an FP4 recipe state (``MXFP4BlockScalingRecipeState``,
+    ``NVFP4BlockScalingRecipeState``) that has no ``scale``/``amax_history``.
+    """
+    if bool(getattr(m, "fp4", False)):
+        return True
+    if hasattr(m, "fp4_initialized"):
+        return True
+    meta = getattr(m, "fp8_meta", None)
+    if isinstance(meta, dict):
+        for key in ("scaling_fwd", "scaling_bwd"):
+            tm = meta.get(key)
+            if tm is None:
+                continue
+            cls = type(tm).__name__
+            if "FP4" in cls or "Fp4" in cls:
+                return True
+    return False
+
+
+def _manual_reset_fp4_meta(module):
+    """Reset FP4 per-tile block-scale buffers when reset_fp4_meta_tensors() is absent."""
+    if not hasattr(module, "fp8_meta"):
+        return False
+    meta = module.fp8_meta
+    reset_count = 0
+    for key in ("scaling_fwd", "scaling_bwd"):
+        if key not in meta:
+            continue
+        tensor_meta = meta[key]
+        for attr in _FP4_META_ATTRS:
+            if not hasattr(tensor_meta, attr):
+                continue
+            t = getattr(tensor_meta, attr)
+            if hasattr(t, "zero_"):
+                try:
+                    t.zero_()
+                    reset_count += 1
+                except Exception:
+                    pass
+    return reset_count > 0
+
+
+def reset_fp4_state(model, reset_meta_tensors=True):
+    """Clear ``fp4_initialized`` (or shared ``fp8_initialized``) on FP4 layers.
+
+    Parallel to :func:`reset_fp8_state` but for FP4 modules.  Unlike FP8 this
+    does NOT seed any history — block-scaling recipes have no amax_history,
+    so TE will recompute per-tile scales from real data on the next forward.
+
+    Returns the number of modules touched.
+    """
+    count = 0
+    method_count = 0
+    manual_count = 0
+
+    def _reset(m):
+        nonlocal count, method_count, manual_count
+        if not _is_fp4_module(m):
+            return
+
+        touched = False
+        if hasattr(m, "fp4_initialized"):
+            m.fp4_initialized = False
+            touched = True
+        if hasattr(m, "fp8_initialized"):
+            m.fp8_initialized = False
+            touched = True
+        if touched:
+            count += 1
+
+        if reset_meta_tensors:
+            if hasattr(m, "reset_fp4_meta_tensors"):
+                try:
+                    m.reset_fp4_meta_tensors()
+                    method_count += 1
+                except Exception:
+                    pass
+            elif hasattr(m, "reset_fp8_meta_tensors"):
+                # FP4 modules typically share the meta-reset method with FP8.
+                try:
+                    m.reset_fp8_meta_tensors()
+                    method_count += 1
+                except Exception:
+                    if _manual_reset_fp4_meta(m):
+                        manual_count += 1
+            elif _manual_reset_fp4_meta(m):
+                manual_count += 1
+
+    models = model if isinstance(model, (list, tuple)) else [model]
+    for m in models:
+        m.apply(_reset)
+    _log(f"reset_fp4_state: {count} modules, " f"{method_count} via method, {manual_count} via manual reset")
+    return count
+
+
+# ===========================================================================
+# Warmup recipe selection  --  NEW, opt-in via WARMUP_RECIPE
+# ===========================================================================
+
+
+def _resolve_recipe_str():
+    """Return the warmup recipe string, honouring legacy WARMUP_FP8_RECIPE.
+
+    Default is "" (empty) → no autocast wrapping, identical to the original
+    pre-FP4 warmup.  FP8/BF16 users do not need to set anything.
+    """
+    val = os.getenv("WARMUP_RECIPE", "").strip().lower()
+    if val:
+        return val
+    legacy = os.getenv("WARMUP_FP8_RECIPE", "").strip().lower()
+    if legacy in ("hybrid", "e4m3"):
+        return f"fp8_{legacy}"
+    return ""
+
+
+def _build_warmup_recipe():
+    """Return a TE recipe object if WARMUP_RECIPE selects one, else None.
+
+    Supports FP8 (DelayedScaling) and FP4 (MXFP4BlockScaling, NVFP4BlockScaling).
+    Both kinds are consumed by ``te.fp8_autocast(fp8_recipe=...)``.
+
+    Returns
+    -------
+    (recipe, kind) where kind is one of ``"fp8"``, ``"fp4"`` or ``None``.
+    """
+    recipe_str = _resolve_recipe_str()
+    if not recipe_str or recipe_str == "bf16":
+        if recipe_str == "bf16":
+            _log("WARMUP_RECIPE=bf16: no autocast wrapping")
+        return None, None
+
+    # ---- FP8 family ------------------------------------------------------
+    if recipe_str.startswith("fp8_"):
+        try:
+            from transformer_engine.common.recipe import DelayedScaling, Format
+        except ImportError:
+            _log(f"WARMUP_RECIPE={recipe_str!r} but transformer_engine missing; no autocast")
+            return None, None
+        fmt = {"fp8_hybrid": Format.HYBRID, "fp8_e4m3": Format.E4M3}.get(recipe_str)
+        if fmt is None:
+            _log(f"Unknown FP8 WARMUP_RECIPE={recipe_str!r}; no autocast")
+            return None, None
+        history_len = int(os.getenv("WARMUP_FP8_HISTORY_LEN", "4"))
+        _log(
+            f"Warmup will wrap train_step in fp8_autocast"
+            f"(format={recipe_str[4:]}, amax_history_len={history_len}, algo=most_recent)"
+        )
+        return (
+            DelayedScaling(
+                margin=0,
+                fp8_format=fmt,
+                amax_history_len=history_len,
+                amax_compute_algo="most_recent",
+            ),
+            "fp8",
+        )
+
+    # ---- FP4 family ------------------------------------------------------
+    if recipe_str.startswith("fp4_"):
+        try:
+            import transformer_engine.common.recipe as te_recipe
+        except ImportError:
+            _log(f"WARMUP_RECIPE={recipe_str!r} but transformer_engine missing; no autocast")
+            return None, None
+
+        # Try to nudge Primus's MXFP4 recipe-state patch into place if the
+        # model was built without FP4 (so the patch never ran).  No-op otherwise.
+        try:
+            from primus.backends.megatron.core.fp4_utils import (
+                _ensure_mxfp4_recipe_support,
+            )
+
+            _ensure_mxfp4_recipe_support()
+        except Exception:
+            pass
+
+        cls_name = {
+            "fp4_mxfp4": "MXFP4BlockScaling",
+            "fp4_nvfp4": "NVFP4BlockScaling",
+        }.get(recipe_str)
+        if cls_name is None:
+            _log(f"Unknown FP4 WARMUP_RECIPE={recipe_str!r}; no autocast")
+            return None, None
+
+        recipe_cls = getattr(te_recipe, cls_name, None)
+        if recipe_cls is None:
+            _log(f"WARMUP_RECIPE={recipe_str!r} but {cls_name} not in this TE build; no autocast")
+            return None, None
+        try:
+            recipe = recipe_cls()
+        except Exception as e:
+            _log(f"Failed to construct {cls_name}() for warmup: {e}; no autocast")
+            return None, None
+        _log(f"Warmup will wrap train_step in fp8_autocast(fp8_recipe={cls_name}())")
+        return recipe, "fp4"
+
+    _log(f"Unknown WARMUP_RECIPE={recipe_str!r}; no autocast")
+    return None, None
+
+
+def _warmup_autocast(recipe):
+    """Return a context-manager factory for the warmup loop."""
+    if recipe is None:
+        return nullcontext
+    import transformer_engine.pytorch as te
+
+    def _ctx():
+        return te.fp8_autocast(enabled=True, fp8_recipe=recipe)
+
+    return _ctx
+
+
+# ---------------------------------------------------------------------------
+# Optimizer save / restore  --  ORIGINAL from committed warmup.py
+# (preserves ChainedOptimizer support for EP < num_gpus configs)
+# ---------------------------------------------------------------------------
+
+
+def _get_inner_optimizers(optimizer):
+    """Return a list of leaf-level torch optimizers from any Megatron wrapper.
+
+    When expert_parallel < num_gpus, Megatron creates a ChainedOptimizer with
+    separate sub-optimizers for dense and expert parameters.  The previous code
+    used ``optimizer.optimizer`` which asserts ``len == 1`` on ChainedOptimizer.
+    This helper iterates over all chained sub-optimizers when present, and falls
+    back to the single-optimizer path otherwise.
+    """
+    if hasattr(optimizer, "chained_optimizers"):
+        inners = []
+        for sub_opt in optimizer.chained_optimizers:
+            inners.append(getattr(sub_opt, "optimizer", sub_opt))
+        return inners
+    return [getattr(optimizer, "optimizer", optimizer)]
+
+
+def _neuter_optimizer(optimizer):
+    """Set betas=[1,1], weight_decay=0, bias_correction=False.
+
+    With betas=[1,1] the Adam momentum/variance stays at its initial (zero)
+    value, so the effective weight update is zero.
+
+    Handles ChainedOptimizer by iterating over all sub-optimizers.
+    """
+    all_saved = []
+    for inner in _get_inner_optimizers(optimizer):
+        saved = []
+        for group in inner.param_groups:
+            state = {}
+            for key in ("betas", "weight_decay", "bias_correction", "pre_mult_wd"):
+                if key in group:
+                    state[key] = group[key]
+            saved.append(state)
+
+            if "betas" in group:
+                group["betas"] = [1.0, 1.0]
+            if "weight_decay" in group:
+                group["weight_decay"] = 0.0
+            if "bias_correction" in group:
+                group["bias_correction"] = False
+            if "pre_mult_wd" in group:
+                group["pre_mult_wd"] = 0.0
+        all_saved.append(saved)
+    return all_saved
+
+
+def _restore_optimizer(optimizer, all_saved):
+    """Restore optimizer state, handling ChainedOptimizer."""
+    for inner, saved in zip(_get_inner_optimizers(optimizer), all_saved):
+        for group, state in zip(inner.param_groups, saved):
+            for key, val in state.items():
+                group[key] = val
+            if "step" in group:
+                del group["step"]
+
+
+# ---------------------------------------------------------------------------
+# Model parameter snapshot  --  ORIGINAL from committed warmup.py, unchanged
+# ---------------------------------------------------------------------------
+
+
+def _save_model_params(models):
+    """Snapshot all model parameters to CPU to avoid doubling GPU memory."""
+    saved = {}
+    for m in models:
+        for name, p in m.named_parameters():
+            saved[(id(m), name)] = p.data.to("cpu", copy=True)
+    return saved
+
+
+def _restore_model_params(models, saved):
+    restored = 0
+    for m in models:
+        for name, p in m.named_parameters():
+            key = (id(m), name)
+            if key in saved:
+                p.data.copy_(saved[key].to(p.device))
+                restored += 1
+    return restored
+
+
+# ---------------------------------------------------------------------------
+# LR scheduler save / restore  --  ORIGINAL + step(0) re-sync after restore
+# ---------------------------------------------------------------------------
+
+_SCHEDULER_KEYS = (
+    "num_steps",
+    "num_floating_point_operations_so_far",
+)
+
+
+def _save_scheduler_state(scheduler):
+    if scheduler is None:
+        return None
+    return {k: getattr(scheduler, k) for k in _SCHEDULER_KEYS if hasattr(scheduler, k)}
+
+
+def _restore_scheduler_state(scheduler, state):
+    if scheduler is None or state is None:
+        return
+    for k, v in state.items():
+        setattr(scheduler, k, v)
+    # Re-sync param_groups['lr'] and ['weight_decay'] to match the restored
+    # num_steps. Without this, the optimizer still holds the LR from the last
+    # warmup step even though num_steps has been rewound to its pre-warmup value.
+    scheduler.step(0)
+
+
+# ---------------------------------------------------------------------------
+# Main entry point
+# ---------------------------------------------------------------------------
+
+
+def run_synthetic_warmup(
+    train_step_func,
+    forward_step_func,
+    model,
+    optimizer,
+    opt_param_scheduler,
+    config,
+    megatron_args,
+):
+    """Run *warmup_steps* forward+backward passes with synthetic data.
+
+    Called from ``MLPerfMegatronPretrainTrainer``'s patched ``train()``
+    **before** ``RUN_START`` is logged, so warmup time is excluded from the
+    timed run.  All side-effects (model params, optimizer, LR scheduler,
+    FP8/FP4 state) are reverted after warmup.
+
+    Default behaviour (``WARMUP_RECIPE`` unset): no autocast wrapping,
+    identical to the pre-FP4 committed version.  Set
+    ``WARMUP_RECIPE=fp4_mxfp4`` or ``fp4_nvfp4`` to opt into FP4 warmup.
+
+    Args:
+        train_step_func: The (already MLPerf-patched) upstream
+            ``megatron.training.training.train_step`` function.  It is called
+            with the new-style signature
+            ``(forward_step_func, data_iterator, model, optimizer,
+            opt_param_scheduler, config, forward_backward_func, iteration=0)``.
+            ``forward_backward_func`` is resolved here via
+            ``get_forward_backward_func()`` (mirrors ``training.train()``).
+    """
+    warmup_steps = int(os.getenv("SYNTH_WARMUP_STEPS", "3"))
+    if warmup_steps <= 0:
+        _log(f"Skipped (SYNTH_WARMUP_STEPS={warmup_steps})")
+        return
+
+    empty_cache_each_step = os.getenv("SYNTH_WARMUP_EMPTY_CACHE", "1") not in ("0", "false", "False")
+
+    t0 = time.time()
+    _log(
+        f"Starting {warmup_steps}-step synthetic warmup "
+        f"(seq_len={megatron_args.seq_length}, "
+        f"mbs={megatron_args.micro_batch_size}, "
+        f"empty_cache_each_step={empty_cache_each_step})"
+    )
+
+    vocab_size = getattr(
+        megatron_args,
+        "padded_vocab_size",
+        getattr(megatron_args, "vocab_size", 32000),
+    )
+    synth_iter = SyntheticGPTDataIterator(
+        megatron_args.seq_length,
+        megatron_args.micro_batch_size,
+        vocab_size,
+    )
+
+    models = model if isinstance(model, (list, tuple)) else [model]
+
+    # Resolve the forward-backward func the same way upstream
+    # ``megatron.training.training.train()`` does (see training.py).  Newer
+    # Megatron ``train_step`` takes this as an explicit positional argument.
+    from megatron.core.pipeline_parallel import get_forward_backward_func
+
+    forward_backward_func = get_forward_backward_func()
+
+    # Temporarily set config fields needed by forward_backward_func.
+    # Megatron's train() normally sets these, but warmup runs before it.
+    from megatron.core.distributed import finalize_model_grads
+    from megatron.core.distributed.distributed_data_parallel import (
+        DistributedDataParallel as DDP,
+    )
+
+    saved_config = {}
+    for key in ("finalize_model_grads_func", "grad_scale_func", "no_sync_func"):
+        saved_config[key] = getattr(config, key, None)
+    if config.finalize_model_grads_func is None:
+        config.finalize_model_grads_func = finalize_model_grads
+    if config.grad_scale_func is None:
+        config.grad_scale_func = optimizer.scale_loss
+    if megatron_args.overlap_grad_reduce and config.no_sync_func is None:
+        if isinstance(models[0], DDP):
+            config.no_sync_func = models[0].no_sync if len(models) == 1 else [m.no_sync for m in models]
+
+    # ---- save state ------------------------------------------------------
+    _log(f"Saving {sum(p.numel() for m in models for p in m.parameters())} parameters")
+    saved_params = _save_model_params(models)
+    saved_opt = _neuter_optimizer(optimizer)
+    saved_sched = _save_scheduler_state(opt_param_scheduler)
+
+    # ---- decide warmup precision (default: no autocast = original path) -
+    recipe, recipe_kind = _build_warmup_recipe()
+    warmup_ctx_factory = _warmup_autocast(recipe)
+
+    # ---- warmup steps ----------------------------------------------------
+    for step in range(1, warmup_steps + 1):
+        step_t0 = time.time()
+        with warmup_ctx_factory():
+            train_step_func(
+                forward_step_func,
+                synth_iter,
+                model,
+                optimizer,
+                opt_param_scheduler,
+                config,
+                forward_backward_func,
+                iteration=0,
+            )
+        torch.cuda.synchronize()
+        if empty_cache_each_step:
+            torch.cuda.empty_cache()
+        _log(f"Step {step}/{warmup_steps} done in {time.time() - step_t0:.1f}s")
+
+    # ---- restore state ---------------------------------------------------
+    _log("Restoring optimizer")
+    _restore_optimizer(optimizer, saved_opt)
+
+    _log("Restoring LR scheduler")
+    _restore_scheduler_state(opt_param_scheduler, saved_sched)
+
+    _log("Restoring model parameters")
+    n_restored = _restore_model_params(models, saved_params)
+    del saved_params
+    _log(f"Restored {n_restored} parameter tensors")
+
+    if hasattr(optimizer, "reload_model_params"):
+        optimizer.reload_model_params()
+        _log("Called optimizer.reload_model_params()")
+
+    # ---- zero Adam state buffers (defensive) ----------------------------
+    # With betas=[1,1] during warmup, exp_avg / exp_avg_sq should already be
+    # zero, but zero them explicitly in case any param group lacked betas.
+    # Iterate over chained sub-optimizers for ChainedOptimizer support.
+    zeroed_state_tensors = 0
+    for inner_opt in _get_inner_optimizers(optimizer):
+        for param_states in inner_opt.state.values():
+            for k, v in param_states.items():
+                if isinstance(v, torch.Tensor) and v.is_floating_point():
+                    v.zero_()
+                    zeroed_state_tensors += 1
+    _log(f"Zeroed {zeroed_state_tensors} optimizer state tensors (exp_avg / exp_avg_sq)")
+
+    for m in models:
+        m.zero_grad(set_to_none=True)
+    _log("Zeroed all model gradients")
+
+    # ---- reset FP8 (always; no-op for BF16 / non-FP8 modules) -----------
+    _log("Resetting FP8 state")
+    total_reset = 0
+    for m in models:
+        total_reset += reset_fp8_state(m, reset_meta_tensors=True)
+    seed_fp8_amax(models, seed_value=1.0)
+
+    # ---- reset FP4 (only when warmup explicitly used an FP4 recipe) -----
+    total_fp4_reset = 0
+    if recipe_kind == "fp4":
+        _log("Resetting FP4 state (no seeding)")
+        for m in models:
+            total_fp4_reset += reset_fp4_state(m, reset_meta_tensors=True)
+
+    # ---- release cached allocator blocks --------------------------------
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+
+    # ---- restore Megatron config ----------------------------------------
+    for key, val in saved_config.items():
+        setattr(config, key, val)
+
+    nan_params = sum(
+        1
+        for m in models
+        for _, p in m.named_parameters()
+        if p.data.is_floating_point() and torch.isnan(p.data).any()
+    )
+    elapsed = time.time() - t0
+    _log(
+        f"Warmup complete in {elapsed:.1f}s "
+        f"(fp8_reset={total_reset}, fp4_reset={total_fp4_reset}, "
+        f"recipe={recipe_kind or 'native'}, nan_params={nan_params})"
+    )
diff --git a/primus/backends/megatron/patches/__init__.py b/primus/backends/megatron/patches/__init__.py
index 7b294fd70..e269a2a68 100644
--- a/primus/backends/megatron/patches/__init__.py
+++ b/primus/backends/megatron/patches/__init__.py
@@ -8,13 +8,30 @@
 Megatron Patch Collection
 
 This module defines the public entrypoint for applying Megatron-specific patches.
+
+Registration vs. application (important for non-Flux / LLM users):
+    Importing this package eagerly imports every ``*_patches.py`` module, which
+    runs each module's ``@register_patch`` side effect. This registration is
+    GLOBAL: it happens for *every* Megatron job (LLM or diffusion), because
+    ``megatron_adapter`` imports this package unconditionally.
+
+    Registration only adds a patch to the registry; it does NOT apply it. Each
+    patch carries a ``condition=...`` predicate (typically gated on a config
+    flag such as ``torch_compile.enable``, ``use_fsdp2_fp8_all_gather``, or a
+    diffusion-specific arg) that is evaluated at ``run_patches`` time. A patch
+    whose condition is False is a no-op for that job.
+
+    Consequence: diffusion/Flux-specific patches are registered for LLM jobs but
+    should not take effect there. When adding a patch, make its ``condition``
+    precise so it cannot alter unrelated (e.g. non-Flux) training. An opt-in
+    "patch profile" mechanism could make this stricter in the future.
 """
 
 import importlib
 import pkgutil
 
 # from primus.core.patches import run_patches
-# from primus.modules.module_utils import log_rank_0
+# from primus.core.utils.module_utils import log_rank_0
 
 
 def _auto_import_patch_modules() -> None:
@@ -46,3 +63,6 @@ def _auto_import_patch_modules() -> None:
 # Eagerly import all patch modules on package import so patches are registered
 # before any backend-specific logic runs.
 _auto_import_patch_modules()
+
+# MLPerf Llama2-70B LoRA Megatron-LM overrides (MXFP4 recipe, optional TE SwiGLU).
+import primus.backends.megatron_bridge.patches.mlperf_llama2_70b.megatron_patches  # noqa: F401, E402
diff --git a/primus/backends/megatron/patches/_patch_guard.py b/primus/backends/megatron/patches/_patch_guard.py
new file mode 100644
index 000000000..c86dfc0b8
--- /dev/null
+++ b/primus/backends/megatron/patches/_patch_guard.py
@@ -0,0 +1,45 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Idempotency guard for monkeypatches that *wrap* (compose) a target callable.
+
+The core patch runner (``primus.core.patches``) has no built-in re-apply guard:
+running the same phase twice re-invokes every patch handler. For patches that
+replace a class method or module attribute outright that is harmless (the
+assignment is idempotent), but patches that *wrap* an existing callable -- e.g.
+the stacked ``train_step`` wrappers (FP8 cache refresh, delayed-scaling
+preamble, wall-clock timer) -- would wrap again on a second run and silently
+double their side effects.
+
+This helper records applied patch keys in a sentinel set attached to the
+*object that owns the wrapped attribute* (typically a module). Keying on the
+owner object -- which is stable across re-runs -- rather than on the wrapped
+callable makes the guard robust even when several wrappers compose on the same
+attribute.
+
+This lives under ``primus.backends.megatron.patches`` (feature-owned) on
+purpose: the shared ``primus.core.patches`` framework is inherited and must not
+grow feature-specific behavior.
+"""
+
+from typing import Any
+
+_SENTINEL_ATTR = "_primus_applied_patch_keys"
+
+
+def is_patched(target: Any, key: str) -> bool:
+    """Return True if ``key`` has already been applied to ``target``."""
+    applied = getattr(target, _SENTINEL_ATTR, None)
+    return applied is not None and key in applied
+
+
+def mark_patched(target: Any, key: str) -> None:
+    """Record that ``key`` has been applied to ``target``."""
+    applied = getattr(target, _SENTINEL_ATTR, None)
+    if applied is None:
+        applied = set()
+        setattr(target, _SENTINEL_ATTR, applied)
+    applied.add(key)
diff --git a/primus/backends/megatron/patches/_source_patch_utils.py b/primus/backends/megatron/patches/_source_patch_utils.py
new file mode 100644
index 000000000..2ef19b5c2
--- /dev/null
+++ b/primus/backends/megatron/patches/_source_patch_utils.py
@@ -0,0 +1,103 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared helper for "source-string rewrite" style patches.
+
+Some upstream Megatron-LM functions/methods need a small code fragment
+inserted in the *middle* of their body (not just wrapped before/after), where
+a plain function-wrapping monkey-patch cannot reach. For those cases Primus
+falls back to: ``inspect.getsource()`` the original, apply a targeted
+``str.replace()`` on a unique anchor fragment, recompile with ``exec()``, and
+rebind the result onto the owning module/class.
+
+This lives under ``primus.backends.megatron.patches`` (feature-owned) on
+purpose: the shared ``primus.core.patches`` framework is inherited and must
+not grow feature-specific behavior.
+"""
+
+import inspect
+import textwrap
+from typing import Any, Callable
+
+
+def patch_method_source(
+    cls: Any,
+    method_name: str,
+    ori_code: str,
+    new_code: str,
+) -> Callable:
+    """Rewrite a fragment of ``cls.<method_name>``'s source and rebind it.
+
+    Args:
+        cls: The class owning the method (e.g. ``MambaModel``).
+        method_name: Name of the method to rewrite (e.g. ``"__init__"``).
+        ori_code: Unique anchor substring to locate within the method source.
+        new_code: Replacement text for ``ori_code`` (typically ``ori_code``
+            plus extra inserted lines, or vice versa).
+
+    Returns:
+        The newly compiled function that was set on ``cls``.
+
+    Raises:
+        AssertionError: If ``ori_code`` is not found in the method's source
+            (e.g. upstream Megatron-LM changed the function unexpectedly).
+    """
+    # IMPORTANT: replace on the *raw* (non-dedented) source -- ori_code/new_code
+    # are written using the upstream file's absolute column indentation (i.e.
+    # what you see reading the file directly). Dedent must happen AFTER the
+    # replacement, once, so the whole (possibly now-longer) body shifts by a
+    # single consistent amount; dedenting first and writing new_code at
+    # post-dedent indentation is a common source of IndentationError.
+    original = getattr(cls, method_name)
+    source = inspect.getsource(original)
+    assert ori_code in source, (
+        f"[SourcePatch] Anchor not found in {cls.__name__}.{method_name}; "
+        f"upstream source may have changed. Anchor: {ori_code!r}"
+    )
+    modified_source = textwrap.dedent(source.replace(ori_code, new_code))
+
+    # IMPORTANT: exec'ing `modified_source` as a bare top-level `def` (not
+    # nested in a class body) silently loses the implicit `__class__` closure
+    # cell that zero-arg `super()` depends on: the function still compiles
+    # and defines fine, but crashes at *call* time with
+    # `RuntimeError: super(): __class__ cell not found` the first time the
+    # method runs. Compiling inside a throwaway wrapper class makes the
+    # compiler wire up that closure normally; we then retarget the cell from
+    # the throwaway class to the real owning `cls` (required for `super()`'s
+    # MRO walk to actually find `cls` in `type(self).__mro__`).
+    wrapper_source = "class _PrimusPatchWrapper:\n" + textwrap.indent(modified_source, "    ")
+    namespace: dict = {}
+    exec(wrapper_source, original.__globals__, namespace)  # noqa: S102
+    wrapper_cls = namespace["_PrimusPatchWrapper"]
+    new_func = wrapper_cls.__dict__[original.__name__]
+    if new_func.__closure__:
+        for cell in new_func.__closure__:
+            if cell.cell_contents is wrapper_cls:
+                cell.cell_contents = cls
+    new_func.__qualname__ = f"{cls.__qualname__}.{method_name}"
+    setattr(cls, method_name, new_func)
+    return new_func
+
+
+def patch_function_source(
+    module: Any,
+    function_name: str,
+    ori_code: str,
+    new_code: str,
+) -> Callable:
+    """Same as :func:`patch_method_source`, but for a module-level function."""
+    original = getattr(module, function_name)
+    source = inspect.getsource(original)
+    assert ori_code in source, (
+        f"[SourcePatch] Anchor not found in {module.__name__}.{function_name}; "
+        f"upstream source may have changed. Anchor: {ori_code!r}"
+    )
+    modified_source = textwrap.dedent(source.replace(ori_code, new_code))
+    namespace: dict = {}
+    exec(modified_source, original.__globals__, namespace)  # noqa: S102
+    new_func = namespace[original.__name__]
+    setattr(module, function_name, new_func)
+    return new_func
diff --git a/primus/backends/megatron/patches/args/__init__.py b/primus/backends/megatron/patches/args/__init__.py
index a994f95b8..647847380 100644
--- a/primus/backends/megatron/patches/args/__init__.py
+++ b/primus/backends/megatron/patches/args/__init__.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -19,6 +19,7 @@
 from . import (  # noqa: F401
     checkpoint_path_patches,
     data_path_split_patches,
+    hsdp_args_patches,
     iterations_to_skip_default_patches,
     logging_level_patches,
     mock_data_patches,
@@ -35,6 +36,7 @@
     "wandb_config_patches",
     "logging_level_patches",
     "data_path_split_patches",
+    "hsdp_args_patches",
     "mock_data_patches",
     "sequence_parallel_tp1_patches",
     "iterations_to_skip_default_patches",
diff --git a/primus/backends/megatron/patches/args/checkpoint_path_patches.py b/primus/backends/megatron/patches/args/checkpoint_path_patches.py
index c67a77b42..bdb56eef9 100644
--- a/primus/backends/megatron/patches/args/checkpoint_path_patches.py
+++ b/primus/backends/megatron/patches/args/checkpoint_path_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -7,7 +7,7 @@
 import os
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
@@ -20,20 +20,39 @@ def patch_checkpoint_path(ctx: PatchContext):
     """
     Configure checkpoint save path.
 
-    Sets args.save to <exp_root>/checkpoints and warns if user
-    provided a different path.
+    Behavior:
+      - args.save is None  (YAML explicitly set `save: null`, or unset):
+            user has opted out of saving — leave args.save = None so that
+            Megatron's `checkpoint_and_decide_exit` (training.py:2340) skips
+            save_checkpoint_and_time(...). This is what allows runs configured
+            with `save: null` + `save_interval: <any>` to avoid hitting the
+            checkpoint save path entirely.
+      - args.save is non-None:
+            user has opted in to saving — route the destination to
+            <exp_root>/checkpoints, warning if it differs from what the user
+            specified.
     """
     args = ctx.extra.get("backend_args", {})
     primus_config = ctx.extra.get("primus_config", {})
 
-    if args and primus_config.exp_root_path:
-        ckpt_path = os.path.abspath(os.path.join(primus_config.exp_root_path, "checkpoints"))
+    if not args or not primus_config.exp_root_path:
+        return
 
-        if hasattr(args, "save") and args.save is not None and args.save != ckpt_path:
-            log_rank_0(
-                f"[Patch:megatron.args.checkpoint_path][WARN] "
-                f"args.save is deprecated; overriding to: {ckpt_path}"
-            )
+    # Respect explicit opt-out via `save: null` in YAML.
+    if not hasattr(args, "save") or args.save is None:
+        log_rank_0(
+            "[Patch:megatron.args.checkpoint_path] "
+            "args.save is None (opt-out); skipping save-path override."
+        )
+        return
 
-        args.save = ckpt_path
-        log_rank_0(f"[Patch:megatron.args.checkpoint_path] save → {ckpt_path}")
+    ckpt_path = os.path.abspath(os.path.join(primus_config.exp_root_path, "checkpoints"))
+
+    if args.save != ckpt_path:
+        log_rank_0(
+            f"[Patch:megatron.args.checkpoint_path][WARN] "
+            f"args.save is deprecated; overriding to: {ckpt_path}"
+        )
+
+    args.save = ckpt_path
+    log_rank_0(f"[Patch:megatron.args.checkpoint_path] save → {ckpt_path}")
diff --git a/primus/backends/megatron/patches/args/data_path_split_patches.py b/primus/backends/megatron/patches/args/data_path_split_patches.py
index 4c326206e..a96d40dec 100644
--- a/primus/backends/megatron/patches/args/data_path_split_patches.py
+++ b/primus/backends/megatron/patches/args/data_path_split_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0
 
 
 def _normalize_data_path(path_value):
diff --git a/primus/backends/megatron/patches/args/hsdp_args_patches.py b/primus/backends/megatron/patches/args/hsdp_args_patches.py
new file mode 100644
index 000000000..57b796344
--- /dev/null
+++ b/primus/backends/megatron/patches/args/hsdp_args_patches.py
@@ -0,0 +1,60 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0
+
+
+@register_patch(
+    "megatron.args.hsdp",
+    backend="megatron",
+    phase="build_args",
+    description=(
+        "Map data_parallel_replicate_degree to Megatron's "
+        "num_distributed_optimizer_instances so that initialize_model_parallel "
+        "creates the HSDP process groups (intra-partial shard + inter-instance replicate)."
+    ),
+)
+def patch_hsdp_args(ctx: PatchContext):
+    """
+    Propagate data_parallel_replicate_degree through to Megatron.
+
+    data_parallel_replicate_degree is a Primus-only config key that is not
+    recognized by MegatronArgBuilder and would be silently dropped.  This
+    patch reads it from module_config.params, injects it into backend_args,
+    and -- when > 1 -- maps it to num_distributed_optimizer_instances so that
+    Megatron's parallel_state.initialize_model_parallel creates the required
+    intra-partial (shard) and inter-instance (replicate) process groups.
+    """
+    args = ctx.extra.get("backend_args")
+    module_config = ctx.extra.get("module_config")
+
+    if not args or not module_config:
+        return
+
+    replicate_degree = getattr(module_config.params, "data_parallel_replicate_degree", 1)
+    args.data_parallel_replicate_degree = replicate_degree
+
+    if replicate_degree > 1:
+        if not getattr(args, "use_torch_fsdp2", False):
+            raise ValueError("data_parallel_replicate_degree > 1 requires use_torch_fsdp2=True")
+
+        ckpt_fmt = getattr(args, "ckpt_format", "torch_dist")
+        if ckpt_fmt != "torch_dcp":
+            raise ValueError(
+                f"HSDP (data_parallel_replicate_degree={replicate_degree}) requires "
+                f"ckpt_format='torch_dcp', got '{ckpt_fmt}'. The torch_dist format's "
+                f"checkpoint offset logic assumes flat DP and is incompatible with a 2D mesh."
+            )
+
+        args.num_distributed_optimizer_instances = replicate_degree
+        log_rank_0(
+            f"[Patch:megatron.args.hsdp] HSDP enabled: "
+            f"data_parallel_replicate_degree={replicate_degree} "
+            f"-> num_distributed_optimizer_instances={replicate_degree}"
+        )
+
+    log_kv_rank_0("[Patch:megatron.args.hsdp] data_parallel_replicate_degree", str(replicate_degree))
diff --git a/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py b/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py
index 525f66ade..b8651b9aa 100644
--- a/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py
+++ b/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args/logging_level_patches.py b/primus/backends/megatron/patches/args/logging_level_patches.py
index cfb39c613..7477ed236 100644
--- a/primus/backends/megatron/patches/args/logging_level_patches.py
+++ b/primus/backends/megatron/patches/args/logging_level_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args/mock_data_patches.py b/primus/backends/megatron/patches/args/mock_data_patches.py
index 25253d4b9..2089b9b3f 100644
--- a/primus/backends/megatron/patches/args/mock_data_patches.py
+++ b/primus/backends/megatron/patches/args/mock_data_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0, log_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0
 
 DEFAULT_NULL_TOKENIZER_VOCAB_SIZE = 131072
 
diff --git a/primus/backends/megatron/patches/args/moe_layer_freq_patches.py b/primus/backends/megatron/patches/args/moe_layer_freq_patches.py
index fa15d0998..7e19daef4 100644
--- a/primus/backends/megatron/patches/args/moe_layer_freq_patches.py
+++ b/primus/backends/megatron/patches/args/moe_layer_freq_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0, warning_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py
new file mode 100644
index 000000000..8fffcdbf8
--- /dev/null
+++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py
@@ -0,0 +1,224 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""ROCm-specific Megatron argument validation.
+
+``validate_args_on_rocm`` is invoked by the ``megatron.validate_args`` patch
+after Megatron's own ``validate_args`` runs. It enforces ROCm/Primus-Turbo
+constraints (deterministic-mode env vars, Turbo FP8/FP4 recipes, sync-free MoE
+auto-config, DeepEP restrictions, ...).
+"""
+
+import inspect
+import os
+
+import torch
+
+from primus.core.utils import logger
+
+
+def is_last_rank():
+    return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
+
+
+def print_rank_last(msg):
+    """If distributed is initialized, print only on last rank."""
+    log_func = logger.info_with_caller
+
+    caller = inspect.stack()[1]
+    caller_frame = caller.frame
+    function_name = caller_frame.f_code.co_name
+    module_name = caller_frame.f_globals["__name__"].split(".")[-1]
+    line = caller.lineno
+
+    if torch.distributed.is_initialized():
+        if is_last_rank():
+            log_func(msg, module_name, function_name, line)
+    else:
+        log_func(msg, module_name, function_name, line)
+
+
+def _get_sync_free_moe_options(args) -> dict:
+    stage = args.turbo_sync_free_moe_stage
+
+    if stage > 3 or stage < 0:
+        raise ValueError("turbo_sync_free_moe_stage only support [0-3]")
+
+    sync_free_moe = {
+        1: {
+            "moe_use_fused_router_with_aux_score": True,
+            "moe_permute_fusion": True,
+            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
+        },
+        2: {
+            "moe_use_fused_router_with_aux_score": True,
+            "use_turbo_deepep": True,
+            "moe_permute_fusion": True,
+            "use_turbo_grouped_gemm": True,
+            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
+        },
+        3: {
+            "moe_use_fused_router_with_aux_score": True,
+            "use_turbo_deepep": True,
+            "moe_permute_fusion": True,
+            "use_turbo_grouped_gemm": True,
+            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
+            "use_turbo_fused_act_with_probs": True,
+        },
+    }
+
+    return sync_free_moe[stage]
+
+
+# FSDP2 custom optimizer selection flags. Each one monkeypatches
+# get_megatron_optimizer at the same priority (50), so at most one may be set.
+_FSDP2_OPTIMIZER_FLAGS = (
+    "use_fsdp2_fp32_param_optimizer",
+    "use_fsdp2_bf16_master_weight_optimizer",
+)
+
+
+def validate_fsdp2_optimizer_exclusivity(args) -> None:
+    """Ensure at most one FSDP2 custom optimizer flag is enabled.
+
+    Enabling more than one would silently let whichever optimizer patch applies
+    last win the monkeypatch, so raise ValueError to fail loudly at
+    arg-validation time (before training starts).
+    """
+    enabled = [flag for flag in _FSDP2_OPTIMIZER_FLAGS if getattr(args, flag, False)]
+    if len(enabled) > 1:
+        raise ValueError(
+            "Conflicting FSDP2 optimizer selection: at most one of "
+            f"{list(_FSDP2_OPTIMIZER_FLAGS)} may be enabled, but got {enabled}. "
+            "Enable exactly one."
+        )
+
+
+def validate_turbo_ragged_grouped_gemm(args) -> None:
+    """Validate the no-padding PrimusTurbo grouped-GEMM path."""
+    option = "use_turbo_ragged_grouped_gemm"
+    if not getattr(args, option, False):
+        return
+    if not getattr(args, "enable_primus_turbo", False) or not getattr(args, "use_turbo_grouped_gemm", False):
+        raise ValueError(f"{option}=True requires enable_primus_turbo=True and use_turbo_grouped_gemm=True.")
+    if (
+        not getattr(args, "fp8", None)
+        or getattr(args, "fp8_recipe", None) != "tensorwise"
+        or getattr(args, "fp4", False)
+    ):
+        raise ValueError(f"{option}=True currently supports only tensorwise FP8.")
+    if getattr(args, "moe_router_padding_for_quantization", False):
+        raise ValueError(
+            "use_turbo_ragged_grouped_gemm=True requires moe_router_padding_for_quantization=False."
+        )
+
+
+def validate_args_on_rocm(args):
+    # Primus-Turbo auto-tuning
+    use_turbo_autotune = getattr(args, "use_turbo_autotune", False)
+    if use_turbo_autotune:
+        # NOTE: Set PRIMUS_TURBO_AUTO_TUNE to 1 to enable turbo auto-tuning.
+        os.environ["PRIMUS_TURBO_AUTO_TUNE"] = "1"
+
+    # Deterministic mode
+    if args.deterministic_mode:
+        # NOTE: Some environment variables affect deterministic mode on ROCm. Need to do extra check.
+        NON_DETERMINISTIC_ENVS = {
+            "TORCH_COMPILE_DISABLE": "1",
+            "ROCBLAS_DEFAULT_ATOMICS_MODE": "0",
+            "PRIMUS_TURBO_AUTO_TUNE": "0",
+            "PRIMUS_DETERMINISTIC": "1",
+        }
+        # NOTE: Some version triton compile exist potential racing condition issue.
+        for env, value in NON_DETERMINISTIC_ENVS.items():
+            assert (
+                os.environ.get(env, None) == value
+            ), f"{env} must be set to {value} in deterministic mode but got {os.environ.get(env, None)} instead."
+
+        # Set fill_uninitialized_memory to False to avoid calling extra fill kernel in deterministic mode.
+        torch.utils.deterministic.fill_uninitialized_memory = False
+
+    assert not getattr(
+        args, "use_turbo_parallel_linear", False
+    ), "use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead."
+
+    validate_fsdp2_optimizer_exclusivity(args)
+
+    use_turbo_gemm = getattr(args, "use_turbo_gemm", False)
+    # Turbo FP8 linear check
+    if args.fp8 and use_turbo_gemm:
+        support_fp8_recipe = ["tensorwise", "blockwise", "mxfp8"]
+        assert (
+            args.fp8_recipe in support_fp8_recipe
+        ), f"{args.fp8_recipe} recipe is not support when enable `use_turbo_gemm`."
+
+    # Turbo FP4 linear check
+    if args.fp4 and use_turbo_gemm:
+        support_fp4_recipe = ["mxfp4"]
+        assert (
+            args.fp4_recipe in support_fp4_recipe
+        ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`."
+
+    # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm.
+    if args.fp8_recipe == "mxfp8":
+        assert (
+            os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "1"
+        ), "Please set `NVTE_ROCM_ENABLE_MXFP8=1` to enable `mxfp8` recipe."
+
+    # dump pp data
+    if args.dump_pp_data and args.pipeline_model_parallel_size == 1:
+        args.dump_pp_data = False
+        print_rank_last(f"Disable args.dump_pp_data since args.pipeline_model_parallel_size=1")
+
+    # PrimusTurboGroupedMLP no longer depends on legacy GroupedMLP; the two
+    # flags are mutually exclusive when turbo is enabled.
+    assert not getattr(
+        args, "use_turbo_grouped_mlp", False
+    ), "use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead."
+    use_turbo_grouped_gemm = getattr(args, "use_turbo_grouped_gemm", False)
+    if use_turbo_grouped_gemm:
+        if getattr(args, "moe_use_legacy_grouped_gemm", False):
+            raise ValueError(
+                "use_turbo_grouped_gemm=True is incompatible with moe_use_legacy_grouped_gemm=True. "
+                "please set moe_use_legacy_grouped_gemm=False."
+            )
+
+    # sync-free MoE
+    if args.turbo_sync_free_moe_stage > 0:
+        assert args.enable_primus_turbo, "Please set `enable_primus_turbo=True` to enable sync-free MoE."
+
+        if args.turbo_sync_free_moe_stage > 1 and not use_turbo_grouped_gemm:
+            raise ValueError(
+                "Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedLinear, please set `use_turbo_grouped_gemm=True`"
+            )
+        options = _get_sync_free_moe_options(args)
+        print_rank_last(
+            f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) =========="
+        )
+        for flag, value in options.items():
+            dots = "." * (73 - len(flag) - len(str(value)))
+            print_rank_last(f"{flag}{dots}{value}")
+            setattr(args, flag, value)
+        print_rank_last(
+            f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) =========="
+        )
+
+    validate_turbo_ragged_grouped_gemm(args)
+
+    # turbo deepep
+    if args.use_turbo_deepep:
+        assert (
+            not args.moe_shared_expert_overlap
+        ), "DeepEP not support moe_shared_expert_overlap, please set `moe_shared_expert_overlap=False`."
+        assert (
+            args.moe_router_dtype == "fp32"
+        ), "DeepEP only supports float32 probs, please set `moe_router_dtype=fp32`"
+        if (
+            args.expert_model_parallel_size >= 16
+            and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO"
+        ):
+            # Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine.
+            assert args.turbo_deepep_num_cu <= 32, "Set `turbo_deepep_num_cu<=32` when using ep_size >= 16."
diff --git a/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py b/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py
index e0d022304..90f732e9f 100644
--- a/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py
+++ b/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py
@@ -5,7 +5,7 @@
 ###############################################################################
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0, log_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args/tensorboard_path_patches.py b/primus/backends/megatron/patches/args/tensorboard_path_patches.py
index f1a3fbb27..98c77c70b 100644
--- a/primus/backends/megatron/patches/args/tensorboard_path_patches.py
+++ b/primus/backends/megatron/patches/args/tensorboard_path_patches.py
@@ -7,7 +7,7 @@
 import os
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args/validate_args_patches.py b/primus/backends/megatron/patches/args/validate_args_patches.py
index b389899f9..2f1e0a622 100644
--- a/primus/backends/megatron/patches/args/validate_args_patches.py
+++ b/primus/backends/megatron/patches/args/validate_args_patches.py
@@ -14,7 +14,7 @@
 import inspect
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 # ---------------------------------------------------------------------------
 # Base wrapper — always active, must be registered first
@@ -34,7 +34,9 @@ def patch_validate_args(ctx: PatchContext):
     import megatron.training.arguments as megatron_args
     import megatron.training.initialize as megatron_init
 
-    from primus.modules.trainer.megatron.utils import validate_args_on_rocm
+    from primus.backends.megatron.patches.args.rocm_arg_validation import (
+        validate_args_on_rocm,
+    )
 
     megatron_args._primus_original_validate_args = megatron_args.validate_args
 
diff --git a/primus/backends/megatron/patches/args/wandb_config_patches.py b/primus/backends/megatron/patches/args/wandb_config_patches.py
index c26e8307b..c1dc58879 100644
--- a/primus/backends/megatron/patches/args/wandb_config_patches.py
+++ b/primus/backends/megatron/patches/args/wandb_config_patches.py
@@ -7,7 +7,7 @@
 import os
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0, log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0, warning_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/args_compat_patches.py b/primus/backends/megatron/patches/args_compat_patches.py
index c443fee69..80e67838c 100644
--- a/primus/backends/megatron/patches/args_compat_patches.py
+++ b/primus/backends/megatron/patches/args_compat_patches.py
@@ -23,7 +23,7 @@
 from typing import List
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 _CANONICAL_ADD_FN_NAMES: List[str] = [
     "_add_network_size_args",
diff --git a/primus/backends/megatron/patches/attention_backend_patches.py b/primus/backends/megatron/patches/attention_backend_patches.py
new file mode 100644
index 000000000..3a806e0f9
--- /dev/null
+++ b/primus/backends/megatron/patches/attention_backend_patches.py
@@ -0,0 +1,87 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Make megatron-core's attention-backend selection ROCm-safe.
+
+``LanguageModule._set_attention_backend()`` reconciles the chosen
+``attention_backend`` with the ``NVTE_*_ATTN`` env vars by *asserting* they are
+unset-or-equal to what the backend wants, then setting them. Primus ROCm images
+intentionally bake ``NVTE_FLASH_ATTN=0`` / ``NVTE_FUSED_ATTN=1`` (TE flash attn
+is unavailable on ROCm; the fused/CK path is used instead), so stock megatron
+assert-crashes before training starts on any model that goes through this probe
+(Mamba, hybrid, non-Turbo GPT, ...):
+
+* ``auto`` wants all three = 1, but the baked ``NVTE_FLASH_ATTN=0`` trips it.
+* an explicit backend (e.g. ``unfused``) wants ``NVTE_FUSED_ATTN=0``, but the
+  baked ``NVTE_FUSED_ATTN=1`` trips it.
+
+On ROCm the baked vars are image defaults, not user intent, so the selected
+backend should *win* over them rather than assert against them:
+
+* ``auto`` -> enable every backend the platform hasn't explicitly disabled
+  (fill only the *unset* vars, so the baked ``NVTE_FLASH_ATTN=0`` is respected).
+* an explicit backend -> force exactly its ``NVTE_*_ATTN`` combination,
+  overriding whatever the image baked in.
+"""
+
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_NVTE_ATTN_ENVS = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN")
+
+
+def _is_rocm(ctx: PatchContext) -> bool:
+    import torch
+
+    return getattr(torch.version, "hip", None) is not None
+
+
+@register_patch(
+    "megatron.attention_backend.rocm_safe",
+    backend="megatron",
+    phase="before_train",
+    description="Let attention_backend override ROCm's baked NVTE_*_ATTN instead of asserting",
+    condition=_is_rocm,
+)
+def patch_attention_backend(ctx: PatchContext):
+    from megatron.core.models.common.language_module.language_module import (
+        LanguageModule,
+    )
+    from megatron.core.transformer.enums import AttnBackend
+
+    if getattr(LanguageModule, "_primus_attention_backend_patched", False):
+        return
+
+    original_set_attention_backend = LanguageModule._set_attention_backend
+
+    # (NVTE_FLASH_ATTN, NVTE_FUSED_ATTN, NVTE_UNFUSED_ATTN) for each explicit backend.
+    explicit_envs = {
+        AttnBackend.flash: ("1", "0", "0"),
+        AttnBackend.fused: ("0", "1", "0"),
+        AttnBackend.unfused: ("0", "0", "1"),
+        AttnBackend.local: ("0", "0", "0"),
+    }
+
+    def _set_attention_backend(self):
+        backend = self.config.attention_backend
+        if backend == AttnBackend.auto:
+            for name in _NVTE_ATTN_ENVS:
+                os.environ.setdefault(name, "1")
+            return
+        values = explicit_envs.get(backend)
+        if values is None:
+            original_set_attention_backend(self)
+            return
+        for name, value in zip(_NVTE_ATTN_ENVS, values):
+            os.environ[name] = value
+
+    LanguageModule._set_attention_backend = _set_attention_backend
+    LanguageModule._primus_attention_backend_patched = True
+    log_rank_0(
+        "[Patch:megatron.attention_backend.rocm_safe] attention_backend now overrides baked NVTE_*_ATTN"
+    )
diff --git a/primus/backends/megatron/patches/build_model_patches.py b/primus/backends/megatron/patches/build_model_patches.py
index a96403184..d81906d80 100644
--- a/primus/backends/megatron/patches/build_model_patches.py
+++ b/primus/backends/megatron/patches/build_model_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -10,15 +10,17 @@
 This module contains patches that modify Megatron's model construction
 behavior to better integrate with Primus.
 
-Current patch:
+Current patches:
     - Disable the second DDP construction inside ``torch.cuda.stream()``
       in ``megatron.training.training.get_model`` by temporarily replacing
       ``torch.cuda.stream`` with a no-op context manager while calling
       the original ``get_model``.
+    - Skip Float16Module wrapping when using FSDP2 FP32 param optimizer (model
+      parameters must stay in FP32 for FSDP2's MixedPrecisionPolicy).
 """
 
-from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
@@ -29,6 +31,7 @@
         "Monkey patch megatron.training.training.get_model to disable the "
         "second DDP construction inside torch.cuda.stream()."
     ),
+    condition=lambda ctx: not getattr(get_args(ctx), "disable_build_model_patches", False),
 )
 def patch_megatron_get_model_disable_second_ddp(ctx: PatchContext) -> None:
     """
@@ -66,3 +69,59 @@ def _noop_stream(*_a, **_k):
     setattr(_patched_get_model, "_primus_disable_second_ddp", True)
     training.get_model = _patched_get_model
     log_rank_0("[Patch:megatron.get_model] Disabled second DDP via torch.cuda.stream no-op wrapper")
+
+
+@register_patch(
+    "megatron.training.training.skip_float16_module",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Skip Float16Module wrapping when using FSDP2 FP32 param optimizer. "
+        "Model parameters must stay FP32 for FSDP2's MixedPrecisionPolicy to handle casting."
+    ),
+    priority=40,
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_fsdp2_fp32_param_optimizer", False)
+        and getattr(get_args(ctx), "use_torch_fsdp2", False)
+        and getattr(get_args(ctx), "bf16", False)
+    ),
+)
+def patch_skip_float16_module(ctx: PatchContext) -> None:
+    """Replace Float16Module with a passthrough wrapper that preserves
+    the attributes Megatron's training loop expects (module, config,
+    vp_size, vp_stage, pg_collection) but does NOT convert parameters
+    to BF16/FP16.
+
+    This is necessary because FSDP2's MixedPrecisionPolicy handles
+    the FP32->BF16 casting during forward/backward, and Float16Module
+    would convert parameters to BF16 before FSDP2 wrapping, defeating
+    the purpose of FP32 parameter storage.
+    """
+    from megatron.core.transformer.module import MegatronModule
+
+    class _FSDP2PassthroughModule(MegatronModule):
+        """Drop-in replacement for Float16Module that keeps params in FP32."""
+
+        def __init__(self, config, module):
+            super().__init__(config)
+            self.config = config
+            self.add_module("module", module)
+            self.vp_size = config.virtual_pipeline_model_parallel_size
+            self.vp_stage = getattr(module, "vp_stage", None)
+            self.pg_collection = getattr(module, "pg_collection", None)
+
+        def set_input_tensor(self, input_tensor):
+            return self.module.set_input_tensor(input_tensor)
+
+        def forward(self, *inputs, **kwargs):
+            return self.module(*inputs, **kwargs)
+
+    import megatron.core.transformer.module as module_mod
+    import megatron.training.training as training_mod
+
+    module_mod.Float16Module = _FSDP2PassthroughModule
+    training_mod.Float16Module = _FSDP2PassthroughModule
+    log_rank_0(
+        "[Patch:skip_float16_module] Replaced Float16Module with FP32 passthrough "
+        "(FSDP2 MixedPrecisionPolicy handles BF16 casting)"
+    )
diff --git a/primus/backends/megatron/patches/checkpoint_patches.py b/primus/backends/megatron/patches/checkpoint_patches.py
index 24755b394..1461eb50a 100644
--- a/primus/backends/megatron/patches/checkpoint_patches.py
+++ b/primus/backends/megatron/patches/checkpoint_patches.py
@@ -12,7 +12,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/dataloader_patch.py b/primus/backends/megatron/patches/dataloader_patch.py
index 54735dc94..774cda1d3 100644
--- a/primus/backends/megatron/patches/dataloader_patch.py
+++ b/primus/backends/megatron/patches/dataloader_patch.py
@@ -23,7 +23,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 # Guard against double-patching (e.g. when the patch system runs twice in tests).
 _PATCHED_ATTR = "_primus_dataloader_mp_context_patched"
diff --git a/primus/backends/megatron/patches/ddp_overlap_compile_patches.py b/primus/backends/megatron/patches/ddp_overlap_compile_patches.py
new file mode 100644
index 000000000..02ed61e12
--- /dev/null
+++ b/primus/backends/megatron/patches/ddp_overlap_compile_patches.py
@@ -0,0 +1,173 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DDP overlap_param_gather + torch.compile compatibility patches.
+
+When using per_block torch.compile with Megatron DDP's overlap_param_gather,
+the original per-sub-module hook registration causes Dynamo to trace through
+hooks containing NCCL side effects, leading to assertion errors.
+
+This patch uses two-tier hook registration:
+  - Layer hooks (recurse=True): one hook per transformer layer (children of
+    ModuleList). These fire in the eager __call__ wrapper before the compiled
+    forward, causing zero graph breaks inside compiled layers.
+  - Non-layer hooks (recurse=False): hooks on all modules outside compiled
+    regions (embeddings, projections, norms). These preserve original Megatron
+    behavior and run entirely in eager mode.
+"""
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.ddp.overlap_param_gather_compile",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Patch DDP forward pre-hooks for overlap_param_gather + torch.compile "
+        "compatibility using two-tier hook registration."
+    ),
+    # Distinct from the FSDP2 fp8-cache patch (priority=45) to avoid a
+    # registration-order-dependent tie-break. The two are mutually exclusive
+    # (DDP vs FSDP2), so the exact relative order is not load-bearing.
+    priority=44,
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_distributed_optimizer", False)
+        and getattr(get_args(ctx), "overlap_param_gather", False)
+        and getattr(getattr(get_args(ctx), "torch_compile", None), "enable", False)
+        and not getattr(get_args(ctx), "disable_ddp_compile_patches", False)
+    ),
+)
+def patch_ddp_overlap_param_gather_for_compile(ctx: PatchContext) -> None:
+    """Patch DistributedDataParallel to use two-tier hook registration.
+
+    Monkey-patches class methods on DDP so that both the initial hook
+    registration in __init__ and subsequent enable/disable cycles in the
+    training loop use the patched logic.
+    """
+    try:
+        from megatron.core.distributed.distributed_data_parallel import (
+            DistributedDataParallel as DDP,
+        )
+        from megatron.core.transformer.cuda_graphs import is_graph_capturing
+
+        def _get_overlap_hook_modules(self):
+            """Partition modules into layer-tier and non-layer-tier for hooks.
+
+            Layer modules: children of nn.ModuleList containers (transformer
+            layers). These get hooks with recurse=True.
+
+            Non-layer modules: everything else (embeddings, projections, norms).
+            These get hooks with recurse=False (original Megatron behavior).
+            Modules that are descendants of a layer module are excluded (they
+            are covered by the layer hook's recurse=True).
+            """
+            # Preserve module-traversal order (deterministic) while de-duplicating;
+            # iterating a plain set would make hook-registration order vary run to
+            # run, which hurts reproducibility/debuggability.
+            layer_modules = []
+            seen = set()
+            for module in self.module.modules():
+                if isinstance(module, torch.nn.ModuleList):
+                    for child in module.children():
+                        if child not in seen:
+                            seen.add(child)
+                            layer_modules.append(child)
+
+            inside_layer = set()
+            for layer in layer_modules:
+                for sub in layer.modules():
+                    inside_layer.add(sub)
+
+            non_layer_modules = [m for m in self.module.modules() if m not in inside_layer]
+
+            return layer_modules, non_layer_modules
+
+        def _make_forward_pre_hook(self, recurse=False):
+            """Create a forward pre-hook parameterized by recurse depth."""
+
+            def hook(module, *unused):
+                if not self.use_forward_hook:
+                    raise RuntimeError("Should use pre-hook only when overlap_param_gather is True")
+
+                if is_graph_capturing():
+                    return
+
+                for param in module.parameters(recurse=recurse):
+                    if param not in self.param_to_bucket_group:
+                        continue
+                    if not param.requires_grad:
+                        raise RuntimeError("Bucketed param in forward pre-hook must require grad")
+
+                    skip_next_bucket_dispatch = (
+                        self.ddp_config.align_param_gather or self.overlap_param_gather_with_optimizer_step
+                    )
+                    self.param_to_bucket_group[param].finish_param_sync(
+                        skip_next_bucket_dispatch=skip_next_bucket_dispatch
+                    )
+
+            return hook
+
+        def enable_forward_pre_hook(self):
+            """Register two-tier forward pre-hooks for param all-gather overlap."""
+            if not self.use_forward_hook:
+                raise RuntimeError("enable_forward_pre_hook requires use_forward_hook=True")
+            if len(self.remove_forward_pre_hook_handles) != 0:
+                raise RuntimeError("Forward pre-hooks already registered")
+
+            layer_modules, non_layer_modules = self._get_overlap_hook_modules()
+
+            layer_hook = self._make_forward_pre_hook(recurse=True)
+            for module in layer_modules:
+                self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook(layer_hook)
+
+            non_layer_hook = self._make_forward_pre_hook(recurse=False)
+            for module in non_layer_modules:
+                self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook(
+                    non_layer_hook
+                )
+
+        def disable_forward_pre_hook(self, param_sync: bool = True):
+            """Remove all forward pre-hooks (both tiers)."""
+            if not self.use_forward_hook:
+                raise RuntimeError("disable_forward_pre_hook requires use_forward_hook=True")
+            for module, handle in list(self.remove_forward_pre_hook_handles.items()):
+                handle.remove()
+            self.remove_forward_pre_hook_handles.clear()
+
+            if param_sync:
+                self.start_param_sync(force_sync=True)
+
+        DDP._get_overlap_hook_modules = _get_overlap_hook_modules
+        DDP._make_forward_pre_hook = _make_forward_pre_hook
+        DDP.enable_forward_pre_hook = enable_forward_pre_hook
+        DDP.disable_forward_pre_hook = disable_forward_pre_hook
+
+        log_rank_0(
+            "[Patch:megatron.ddp.overlap_param_gather_compile] "
+            "Patched DDP with two-tier hook registration for "
+            "overlap_param_gather + torch.compile compatibility"
+        )
+
+    except Exception as e:
+        import traceback
+
+        log_rank_0(
+            f"[Patch:megatron.ddp.overlap_param_gather_compile] "
+            f"ERROR: Failed to patch DDP: {type(e).__name__}: {e}"
+        )
+        log_rank_0(f"Traceback: {traceback.format_exc()}")
+        # Re-raise: silently leaving DDP unpatched here would run training with
+        # incorrect forward pre-hooks (Dynamo graph breaks / NCCL side effects),
+        # so fail loudly instead of degrading correctness.
+        raise RuntimeError(
+            "Failed to apply DDP overlap_param_gather + torch.compile patch; "
+            "see log above for the underlying error."
+        ) from e
diff --git a/primus/backends/megatron/patches/deepseek_v4_flops_patches.py b/primus/backends/megatron/patches/deepseek_v4_flops_patches.py
new file mode 100644
index 000000000..247be14a8
--- /dev/null
+++ b/primus/backends/megatron/patches/deepseek_v4_flops_patches.py
@@ -0,0 +1,910 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 FLOPs reporting patch.
+
+Plan-3 P20.  Megatron's
+:func:`megatron.training.training.num_floating_point_operations` is shaped
+for GPT / MLA / Mamba and gets V4 wrong on every axis that matters:
+
+* QKV projections — V4 has a Q LoRA path
+  (``hidden -> q_lora_rank -> n_heads * head_dim``) and a single-latent
+  ``linear_kv`` (``hidden -> head_dim`` shared as both K and V); the
+  upstream MHA/GQA branch counts a flat ``hidden * (q + k + v)`` projection
+  instead.
+* Output projection — V4 uses grouped low-rank ``linear_o_a`` /
+  ``linear_o_b`` (``(n*d/o_groups) -> o_groups*o_lora -> hidden``); the
+  upstream branch counts a flat ``(n*d) -> hidden`` proj.
+* Attention scores at the wrong sequence length — V4's mHC residual
+  packs ``hc_mult`` parallel streams into the layer-internal sequence axis
+  (``[B, S*K, D]``).  Per-layer GEMMs run at ``S_eff = S * hc_mult``, but
+  upstream uses ``args.seq_length``.
+* Compressor + Indexer side paths — CSA (``compress_ratio==4``) and HCA
+  (``compress_ratio==128``) layers add a Compressor (``wkv`` + ``wgate``);
+  CSA additionally runs an Indexer (``w_dq`` + ``w_iuq`` + ``w_w`` +
+  mini-Compressor + scoring einsum).  Upstream knows about neither.
+* Hash routing — V4's first ``num_hash_layers`` MoE layers use a
+  parameter-free hash router; upstream charges them the topk router GEMM
+  (``hidden * num_experts``) anyway.
+* MTP — V4's MTP block runs a full inner V4 transformer layer per depth
+  (attention + MoE FFN) on top of the ``eh_proj``; upstream counts only
+  three norms and the single ``2H -> H`` projection.
+
+The mismatch makes per-iter MFU comparisons across PP/EP/VPP configs
+meaningless because the denominator is wrong by a configuration-dependent
+factor.
+
+This module installs a single ``before_train`` patch that monkey-patches
+``training_module.num_floating_point_operations`` with a wrapper.  The
+wrapper:
+
+1. Falls through to the upstream function byte-for-byte for
+   ``args.model_type != "deepseek_v4"`` so dense GPT / MLA / Mamba runs
+   are unchanged.
+2. For V4 runs, evaluates :func:`compute_v4_flops` — a closed form
+   derived from the V4 forward pass (see
+   ``deepseek-v4/develop/plan-3/02-phase-details.md#p20--v4-aware-tflops-reporting``
+   for the per-component derivation) — and returns its total.
+3. On first invocation, logs the per-component FLOPs breakdown at rank 0
+   so the formula can be sanity-checked against a hand calculation.
+
+Convention follows upstream Megatron: pure-FMAC counts internally, then a
+single ``forward_backward_factor (3) * fma_factor (2) = 6`` multiplier at
+the end.
+
+Plan-6 P33 closes two known gaps in the plan-3 P20 closed form:
+
+* SWA visible-pair pruning — :func:`_attn_scores_fmac_per_layer` now
+  counts only causal-visible ``(q, k)`` pairs surviving the per-layer
+  SWA + pool + sparse top-K masks (via :func:`_visible_pairs`).  The
+  legacy ``B * n * d * S_eff^2`` upper bound over-counted the local
+  branch by ``S_eff / swa_window`` once plan-5 P30+ made kernels honor
+  per-row SWA pruning (~128x at the V4-Flash proxy shape).
+* HyperConnection matmul accounting — the new ``hc`` row of
+  :class:`_V4FlopsBreakdown` counts ``HyperMixer.fn`` (``K*D ->
+  (2+K)*K`` per token, twice per layer) and ``HyperHead.fn`` (``K*D
+  -> K`` per token, once at the trunk end and per MTP depth).  These
+  matmuls were not in the plan-3 P20 form.
+
+See ``deepseek-v4/develop/plan-6/02-phase-details.md#phase-33`` for
+the design and ``develop/progress/p33/p33-summary.md`` for the
+per-component delta vs the legacy formula.
+"""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from typing import Any, List, Optional, Sequence, Tuple
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+# ---------------------------------------------------------------------------
+# Shared constants (mirror Megatron's expansion factors)
+# ---------------------------------------------------------------------------
+
+_FORWARD_BACKWARD_FACTOR: int = 3  # 1 forward + 2 backward.
+_FMA_FACTOR: int = 2  # multiply + add per matmul element.
+_SWIGLU_FFN_EXPANSION_FACTOR: int = 3  # gate + up + down (all hidden*ffn).
+
+
+# ---------------------------------------------------------------------------
+# compress_ratios parsing — V4 yamls store the schedule as a JSON-like string.
+# ---------------------------------------------------------------------------
+
+
+def _parse_compress_ratios(raw: Any) -> Optional[List[int]]:
+    """Parse ``compress_ratios`` into ``list[int]`` or return ``None``.
+
+    Accepts ``None`` (fully dense), a string like ``"[0, 0, 4, 128]"`` (the
+    YAML form), or an existing list / tuple of ints.  Mirrors the runtime
+    helper in
+    :func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block._parse_int_sequence`
+    intentionally inline so this patch has no V4-module import dependency
+    (the patch loads at ``before_train``, before model build).
+    """
+    if raw is None:
+        return None
+    if isinstance(raw, (list, tuple)):
+        return [int(x) for x in raw]
+    if isinstance(raw, str):
+        text = raw.strip()
+        if not text:
+            return None
+        # Tolerant of the YAML form `[0,0,4,128]` and Python repr `[0, 0, 4, 128]`.
+        return [int(x) for x in ast.literal_eval(text)]
+    raise TypeError(f"Unsupported compress_ratios type: {type(raw).__name__}")
+
+
+def _normalize_layer_ratios(
+    raw: Any,
+    *,
+    num_layers: int,
+    mtp_num_layers: int,
+) -> Tuple[List[int], List[int]]:
+    """Return ``(decoder_ratios, mtp_ratios)`` each padded to its expected length.
+
+    Mirrors the V4 block's normalization (``deepseek_v4_block._normalize_compress_ratios``)
+    but also returns the trailing MTP slice when the YAML provides
+    ``num_layers + mtp_num_layers`` entries (the canonical DeepSeek layout).
+    Default-fills any missing slot with ``0`` (dense).
+    """
+    parsed = _parse_compress_ratios(raw)
+    if parsed is None:
+        return [0] * num_layers, [0] * mtp_num_layers
+
+    if len(parsed) == num_layers + mtp_num_layers:
+        return parsed[:num_layers], parsed[num_layers:]
+    if len(parsed) == num_layers:
+        return parsed, [0] * mtp_num_layers
+    if len(parsed) > num_layers:
+        return parsed[:num_layers], [0] * mtp_num_layers
+    pad = parsed[-1] if parsed else 0
+    return parsed + [pad] * (num_layers - len(parsed)), [0] * mtp_num_layers
+
+
+# ---------------------------------------------------------------------------
+# Per-component closed-form helpers
+#
+# All functions return *FMAC* counts (multiplies only).  The caller multiplies
+# by ``_FMA_FACTOR * _FORWARD_BACKWARD_FACTOR = 6`` exactly once at the end.
+# ---------------------------------------------------------------------------
+
+
+def _attn_qkv_o_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len_eff: int,
+    hidden_size: int,
+    num_heads: int,
+    head_dim: int,
+    q_lora_rank: int,
+    o_lora_rank: int,
+    o_groups: int,
+) -> int:
+    """FMAC for V4 attention's projections (Q LoRA + single-latent KV + O).
+
+    Independent of ``compress_ratio``: all V4 layer types share the same
+    projection structure.  The score / softmax cost is counted separately
+    by :func:`_attn_scores_fmac_per_layer`.
+    """
+    n_d = num_heads * head_dim
+
+    qkv = (
+        # linear_q_down_proj: hidden -> q_lora_rank
+        hidden_size * q_lora_rank
+        # linear_q_up_proj: q_lora_rank -> n_heads * head_dim
+        + q_lora_rank * n_d
+        # linear_kv (single latent): hidden -> head_dim
+        + hidden_size * head_dim
+    )
+    if o_lora_rank > 0:
+        # linear_o_a: (n*d/o_groups) -> o_groups*o_lora  ⇒ FMAC = (n*d)*o_lora
+        # linear_o_b: o_groups*o_lora -> hidden          ⇒ FMAC = (o_groups*o_lora)*hidden
+        o_proj = n_d * o_lora_rank + (o_groups * o_lora_rank) * hidden_size
+    else:
+        # Flat fallback: (n*d) -> hidden
+        o_proj = n_d * hidden_size
+
+    return batch_size * seq_len_eff * (qkv + o_proj)
+
+
+def _local_visible_pairs(swa_window: int, seq_len_eff: int) -> int:
+    """Causal-visible ``(q, k)`` pair count for the local SWA branch.
+
+    Each query at position ``q`` attends to keys in
+    ``[max(0, q - swa_window + 1), q]``.  Summed over
+    ``q in [0, seq_len_eff)`` the count is:
+
+    * ``swa_window * seq_len_eff - swa_window * (swa_window - 1) // 2`` for
+      ``0 < swa_window < seq_len_eff`` (queries below ``swa_window - 1``
+      see only ``q + 1`` keys, queries above saturate at ``swa_window``).
+    * ``seq_len_eff * (seq_len_eff + 1) // 2`` for the full-causal
+      fallback (``swa_window == 0`` or ``>= seq_len_eff``).
+    """
+    s = int(seq_len_eff)
+    w = int(swa_window)
+    if w <= 0 or w >= s:
+        return s * (s + 1) // 2
+    return w * s - w * (w - 1) // 2
+
+
+def _pool_visible_pairs(compress_ratio: int, seq_len_eff: int) -> int:
+    """Causal-visible ``(q, p)`` pair count for the HCA compressed pool.
+
+    Pool slot ``p`` covers source positions ``[p*c, (p+1)*c)`` with
+    ``c == compress_ratio``; slot ``p`` is visible to query ``q`` iff
+    ``(p+1) * c - 1 <= q``.  Summed over queries, the count equals
+    ``sum_{m=1..seq_len_eff} floor(m / c)`` with closed form
+    ``c * n * (n - 1) // 2 + n * (T - c*n + 1)`` where ``n = T // c`` and
+    ``T = seq_len_eff``.
+    """
+    c = int(compress_ratio)
+    t = int(seq_len_eff)
+    if c <= 0 or t <= 0:
+        return 0
+    n = t // c
+    if n == 0:
+        return 0
+    return c * n * (n - 1) // 2 + n * (t - c * n + 1)
+
+
+def _visible_pairs(
+    *,
+    swa_window: int,
+    compress_ratio: int,
+    index_topk: int,
+    seq_len_eff: int,
+) -> int:
+    """Total causal-visible ``(q, k)`` pair count for one V4 attention layer.
+
+    Plan-6 P33 closed form — see
+    ``deepseek-v4/develop/perf/attention_perf.md`` "Test Shape And Counting"
+    for the per-branch derivation.
+
+    * ``cr == 0`` (dense + SWA): local SWA pairs only.
+    * ``cr == 128`` (HCA): local SWA pairs + causal pool pairs.
+    * ``cr == 4`` (CSA): local SWA pairs + sparse top-K pairs
+      (``min(index_topk, pool) * seq_len_eff``); top-K is treated as
+      fully visible because the indexer assigns each query a per-row
+      causal-respecting pool subset.
+    """
+    local = _local_visible_pairs(swa_window, seq_len_eff)
+    cr = int(compress_ratio)
+    if cr == 0:
+        return local
+
+    pool = max(1, int(seq_len_eff) // cr)
+    if cr == 128:
+        return local + _pool_visible_pairs(cr, seq_len_eff)
+    if cr == 4:
+        sparse_keys = min(int(index_topk) if index_topk else pool, pool)
+        return local + sparse_keys * int(seq_len_eff)
+    # Forward-compatible: any other ratio is treated as full pool cross-attn.
+    return local + pool * int(seq_len_eff)
+
+
+def _attn_scores_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len_eff: int,
+    num_heads: int,
+    head_dim: int,
+    compress_ratio: int,
+    index_topk: int,
+    swa_window: int,
+) -> int:
+    """FMAC for the attention score matmuls.
+
+    Plan-6 P33 rewrite: counts only the causal-visible ``(query, key)``
+    pairs surviving the per-layer mask (SWA + pool + sparse top-K) — the
+    plan-3 P20 ``S_eff^2`` upper bound over-counted dense / HCA local
+    attention by ``S_eff / swa_window`` (16x at ``swa=128, S_eff=4096``)
+    once plan-5 P30 SWA K-loop pruning made the per-layer kernels track
+    visible pairs only.
+
+    FMAC = ``2 * num_heads * head_dim * visible_pairs``: one ``n*d`` for
+    the ``QK^T`` matmul + one ``n*d`` for the ``PV`` matmul, summed over
+    visible ``(q, k)`` pairs.  The forward + backward * FMA expansion is
+    applied once at the end in :class:`_V4FlopsBreakdown`.
+    """
+    pairs = _visible_pairs(
+        swa_window=swa_window,
+        compress_ratio=compress_ratio,
+        index_topk=index_topk,
+        seq_len_eff=seq_len_eff,
+    )
+    return 2 * batch_size * num_heads * head_dim * pairs
+
+
+def _compressor_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len_eff: int,
+    hidden_size: int,
+    head_dim: int,
+    compress_ratio: int,
+) -> int:
+    """FMAC for the V4 :class:`Compressor` (HCA / CSA only).
+
+    Compressor projects ``hidden -> coff*head_dim`` for both ``wkv`` and
+    ``wgate``; ``coff = 2`` in overlap mode (CSA, ratio==4) and ``coff = 1``
+    in non-overlap mode (HCA, ratio==128).  Inputs are at the full
+    pre-pool seq length, so cost is paid at ``S_eff``.
+    """
+    if compress_ratio == 0:
+        return 0
+    coff = 2 if compress_ratio == 4 else 1
+    # wkv + wgate, each hidden -> coff * head_dim
+    return 2 * batch_size * seq_len_eff * hidden_size * (coff * head_dim)
+
+
+def _indexer_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len_eff: int,
+    hidden_size: int,
+    compress_ratio: int,
+    index_head_dim: int,
+    index_n_heads: int,
+) -> int:
+    """FMAC for the V4 :class:`Indexer` (CSA only)."""
+    if compress_ratio != 4:
+        return 0
+
+    pool = max(1, seq_len_eff // compress_ratio)
+    dq_rank = index_head_dim  # Indexer.__init__ default: dq_rank = index_head_dim
+    inh_ihd = index_n_heads * index_head_dim
+
+    proj = (
+        # w_dq: hidden -> dq_rank
+        hidden_size * dq_rank
+        # w_iuq: dq_rank -> inh * ihd
+        + dq_rank * inh_ihd
+        # w_w: hidden -> inh
+        + hidden_size * index_n_heads
+    )
+    # mini-Compressor inside the Indexer: head_dim=index_head_dim, ratio=4 (coff=2),
+    # so wkv + wgate each cost hidden * (2 * index_head_dim).
+    mini_compressor = 2 * hidden_size * (2 * index_head_dim)
+    proj += mini_compressor
+
+    proj_fmac = batch_size * seq_len_eff * proj
+
+    # Scoring einsum: (B, S_eff, inh, ihd) · (B, P, ihd) → (B, S_eff, inh, P)
+    scoring_fmac = batch_size * seq_len_eff * index_n_heads * pool * index_head_dim
+
+    return proj_fmac + scoring_fmac
+
+
+def _moe_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len_eff: int,
+    hidden_size: int,
+    moe_ffn_hidden_size: int,
+    moe_router_topk: int,
+    num_experts: int,
+    is_hash_layer: bool,
+    shared_expert_ffn_hidden_size: int,
+) -> int:
+    """FMAC for the V4 MoE FFN (router + routed experts + shared expert).
+
+    Hash-routed layers skip the topk router GEMM (they look up bucket
+    indices from the raw input ids — a parameter-free op).  Non-hash
+    layers pay ``hidden * num_experts`` per token for the router.
+    SwiGLU's ``ffn_expansion_factor=3`` collapses the (gate + up + down)
+    matmul triple into a single multiplier consistent with upstream.
+    """
+    router = 0 if is_hash_layer else hidden_size * num_experts
+    routed = moe_router_topk * _SWIGLU_FFN_EXPANSION_FACTOR * hidden_size * moe_ffn_hidden_size
+    shared = (
+        _SWIGLU_FFN_EXPANSION_FACTOR * hidden_size * shared_expert_ffn_hidden_size
+        if shared_expert_ffn_hidden_size > 0
+        else 0
+    )
+    return batch_size * seq_len_eff * (router + routed + shared)
+
+
+def _hc_mixer_fmac_per_layer(
+    *,
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    hc_mult: int,
+) -> int:
+    """FMAC for the two ``HyperMixer.fn`` matmuls in one V4 hybrid layer.
+
+    Each :class:`HyperMixer` projects the un-packed K streams
+    ``[B, S, K*D] -> [B, S, (2+K)*K]``; the V4 hybrid layer runs two
+    mixers per layer (one before / one inside the attention sub-block,
+    and one for the FFN sub-block — see
+    ``primus.backends.megatron.core.transformer.hyper_connection``).
+
+    The leading axis is ``B * S`` (not ``B * S * hc_mult``) because the
+    K stream lifting that packs streams into the sequence axis happens
+    *after* the mixer.  The (small) ``HyperMixer.expand`` matmul
+    ``comb @ x`` is left out by design (see plan-6 P33 spec).
+    """
+    if hc_mult <= 0:
+        return 0
+    n_d = hc_mult * hidden_size
+    return 2 * batch_size * seq_len * n_d * ((2 + hc_mult) * hc_mult)
+
+
+def _hc_head_fmac(
+    *,
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    hc_mult: int,
+    mtp_num_layers: int,
+) -> int:
+    """FMAC for the ``HyperHead.fn`` matmuls (trunk end + per MTP depth).
+
+    Each :class:`HyperHead` projects ``[B, S, K*D] -> [B, S, K]`` to
+    produce the sigmoid weights for the final K-stream collapse.  V4 has
+    one head at the trunk end and one head per MTP depth (each with its
+    own ``num_nextn_predict_layers`` copies).
+    """
+    if hc_mult <= 0:
+        return 0
+    n_d = hc_mult * hidden_size
+    return (1 + mtp_num_layers) * batch_size * seq_len * n_d * hc_mult
+
+
+def _mtp_eh_proj_fmac(
+    *,
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    mtp_num_layers: int,
+) -> int:
+    """FMAC for the per-MTP-depth ``eh_proj`` (``2H -> H``).
+
+    Runs at the original (un-packed) seq length because MTP runs **before**
+    the V4 transformer layer's stream-lift, on the embedding output.  Two
+    norms per depth are negligible.
+    """
+    if mtp_num_layers <= 0:
+        return 0
+    return mtp_num_layers * batch_size * seq_len * (2 * hidden_size) * hidden_size
+
+
+def _logits_fmac(
+    *,
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    padded_vocab_size: int,
+    mtp_num_layers: int,
+) -> int:
+    """FMAC for the LM head (one per main path + one per MTP depth)."""
+    return (mtp_num_layers + 1) * batch_size * seq_len * hidden_size * padded_vocab_size
+
+
+# ---------------------------------------------------------------------------
+# Public closed form
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class _V4FlopsBreakdown:
+    """Per-component FMAC totals (multiply-only, pre fwd+bwd / FMA scaling).
+
+    Plan-6 P33 adds the ``hc`` field for HyperConnection matmuls
+    (``HyperMixer.fn`` + ``HyperHead.fn``).  Field added at the end so
+    existing positional callers and the breakdown log ordering stay
+    stable; ``hc`` defaults to 0 so historical pickles / external
+    consumers that built the dataclass with 7 positional args keep
+    working.
+    """
+
+    attn_qkv_o: int
+    attn_scores: int
+    compressor: int
+    indexer: int
+    moe: int
+    mtp: int
+    logits: int
+    hc: int = 0
+
+    def total_fmac(self) -> int:
+        return (
+            self.attn_qkv_o
+            + self.attn_scores
+            + self.compressor
+            + self.indexer
+            + self.moe
+            + self.mtp
+            + self.logits
+            + self.hc
+        )
+
+    def to_total_flops(self) -> int:
+        """Apply the fwd+bwd (3) × FMA (2) = 6 expansion."""
+        return _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR * self.total_fmac()
+
+
+def compute_v4_flops(args: Any, batch_size: int) -> Tuple[int, _V4FlopsBreakdown]:
+    """Closed-form V4 FLOPs for one global batch.
+
+    Returns ``(total_flops, breakdown)`` where ``total_flops`` is the
+    Megatron-convention number suitable for direct substitution into
+    upstream's reporting and ``breakdown`` is a per-component report
+    (FMAC pre-expansion) for sanity logging.
+    """
+    seq_len = int(args.seq_length)
+    hc_mult = int(getattr(args, "hc_mult", 1) or 1)
+    seq_len_eff = seq_len * hc_mult
+
+    hidden_size = int(args.hidden_size)
+    num_heads = int(args.num_attention_heads)
+    head_dim = int(args.kv_channels)
+    q_lora_rank = int(getattr(args, "q_lora_rank", 0) or 0)
+    o_lora_rank = int(getattr(args, "o_lora_rank", 0) or 0)
+    o_groups = int(getattr(args, "o_groups", 1) or 1)
+
+    num_layers = int(args.num_layers)
+    mtp_num_layers = int(getattr(args, "mtp_num_layers", 0) or 0)
+
+    decoder_ratios, mtp_ratios = _normalize_layer_ratios(
+        getattr(args, "compress_ratios", None),
+        num_layers=num_layers,
+        mtp_num_layers=mtp_num_layers,
+    )
+
+    moe_ffn_hidden_size = int(getattr(args, "moe_ffn_hidden_size", None) or args.ffn_hidden_size)
+    moe_router_topk = int(getattr(args, "moe_router_topk", 1) or 1)
+    num_experts = int(getattr(args, "num_experts", 1) or 1)
+    shared_expert_ffn_hidden_size = int(getattr(args, "moe_shared_expert_intermediate_size", 0) or 0)
+    num_hash_layers = int(getattr(args, "num_hash_layers", 0) or 0)
+
+    index_topk = int(getattr(args, "index_topk", 0) or 0)
+    index_head_dim = int(getattr(args, "index_head_dim", 0) or 0)
+    index_n_heads = int(getattr(args, "index_n_heads", 0) or 0)
+
+    swa_window = int(getattr(args, "attn_sliding_window", 0) or 0)
+
+    padded_vocab_size = int(getattr(args, "padded_vocab_size", None) or args.vocab_size)
+
+    # ---- decoder layers ----
+    attn_qkv_o = 0
+    attn_scores = 0
+    compressor = 0
+    indexer = 0
+    moe = 0
+    hc = 0
+
+    for layer_idx in range(num_layers):
+        ratio = int(decoder_ratios[layer_idx])
+
+        attn_qkv_o += _attn_qkv_o_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            q_lora_rank=q_lora_rank,
+            o_lora_rank=o_lora_rank,
+            o_groups=o_groups,
+        )
+        attn_scores += _attn_scores_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            compress_ratio=ratio,
+            index_topk=index_topk,
+            swa_window=swa_window,
+        )
+        compressor += _compressor_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            head_dim=head_dim,
+            compress_ratio=ratio,
+        )
+        indexer += _indexer_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            compress_ratio=ratio,
+            index_head_dim=index_head_dim,
+            index_n_heads=index_n_heads,
+        )
+        moe += _moe_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            moe_ffn_hidden_size=moe_ffn_hidden_size,
+            moe_router_topk=moe_router_topk,
+            num_experts=num_experts,
+            is_hash_layer=(layer_idx < num_hash_layers),
+            shared_expert_ffn_hidden_size=shared_expert_ffn_hidden_size,
+        )
+        hc += _hc_mixer_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len=seq_len,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+        )
+
+    # ---- MTP layers (one full V4 layer per depth + eh_proj per depth) ----
+    mtp_attn_qkv_o = 0
+    mtp_attn_scores = 0
+    mtp_compressor = 0
+    mtp_indexer = 0
+    mtp_moe = 0
+    mtp_hc = 0
+    for depth in range(mtp_num_layers):
+        ratio = int(mtp_ratios[depth]) if depth < len(mtp_ratios) else 0
+        mtp_attn_qkv_o += _attn_qkv_o_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            q_lora_rank=q_lora_rank,
+            o_lora_rank=o_lora_rank,
+            o_groups=o_groups,
+        )
+        mtp_attn_scores += _attn_scores_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            compress_ratio=ratio,
+            index_topk=index_topk,
+            swa_window=swa_window,
+        )
+        mtp_compressor += _compressor_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            head_dim=head_dim,
+            compress_ratio=ratio,
+        )
+        mtp_indexer += _indexer_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            compress_ratio=ratio,
+            index_head_dim=index_head_dim,
+            index_n_heads=index_n_heads,
+        )
+        mtp_moe += _moe_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len_eff=seq_len_eff,
+            hidden_size=hidden_size,
+            moe_ffn_hidden_size=moe_ffn_hidden_size,
+            moe_router_topk=moe_router_topk,
+            num_experts=num_experts,
+            # V4's MTP depths run after num_hash_layers in the routing ordering
+            # (the released checkpoint stores topk router weights for them).
+            is_hash_layer=False,
+            shared_expert_ffn_hidden_size=shared_expert_ffn_hidden_size,
+        )
+        mtp_hc += _hc_mixer_fmac_per_layer(
+            batch_size=batch_size,
+            seq_len=seq_len,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+        )
+
+    attn_qkv_o += mtp_attn_qkv_o
+    attn_scores += mtp_attn_scores
+    compressor += mtp_compressor
+    indexer += mtp_indexer
+    moe += mtp_moe
+    hc += mtp_hc
+    hc += _hc_head_fmac(
+        batch_size=batch_size,
+        seq_len=seq_len,
+        hidden_size=hidden_size,
+        hc_mult=hc_mult,
+        mtp_num_layers=mtp_num_layers,
+    )
+
+    mtp = _mtp_eh_proj_fmac(
+        batch_size=batch_size,
+        seq_len=seq_len,
+        hidden_size=hidden_size,
+        mtp_num_layers=mtp_num_layers,
+    )
+
+    logits = _logits_fmac(
+        batch_size=batch_size,
+        seq_len=seq_len,
+        hidden_size=hidden_size,
+        padded_vocab_size=padded_vocab_size,
+        mtp_num_layers=mtp_num_layers,
+    )
+
+    breakdown = _V4FlopsBreakdown(
+        attn_qkv_o=attn_qkv_o,
+        attn_scores=attn_scores,
+        compressor=compressor,
+        indexer=indexer,
+        moe=moe,
+        mtp=mtp,
+        logits=logits,
+        hc=hc,
+    )
+    return breakdown.to_total_flops(), breakdown
+
+
+# ---------------------------------------------------------------------------
+# Wrapper installation
+# ---------------------------------------------------------------------------
+
+# Module-level latch so the breakdown is logged exactly once even though
+# ``num_floating_point_operations`` is called many times per training run
+# (one per ``training_log`` and one per ``train_step``).
+_BREAKDOWN_LOGGED = False
+
+
+def _emit_breakdown(
+    *,
+    args: Any,
+    batch_size: int,
+    breakdown: _V4FlopsBreakdown,
+    total_flops: int,
+) -> None:
+    """Emit the per-component breakdown via single-line ``log_rank_0`` calls.
+
+    One row per ``log_rank_0`` call (instead of a single multi-line message)
+    so each line passes cleanly through Primus's per-line logger formatter
+    and rank-aware filter.  The header line carries the run-shape metadata
+    so the breakdown can be matched to a specific ``(args, batch_size)``
+    pairing in the log.
+    """
+
+    def _tflops(fmac: int) -> float:
+        return fmac * _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR / 1.0e12
+
+    rows: List[Tuple[str, float]] = [
+        ("attn_qkv_o", _tflops(breakdown.attn_qkv_o)),
+        ("attn_scores", _tflops(breakdown.attn_scores)),
+        ("compressor", _tflops(breakdown.compressor)),
+        ("indexer", _tflops(breakdown.indexer)),
+        ("moe", _tflops(breakdown.moe)),
+        ("mtp_eh_proj", _tflops(breakdown.mtp)),
+        ("logits", _tflops(breakdown.logits)),
+        ("hc", _tflops(breakdown.hc)),
+    ]
+
+    log_rank_0(
+        "[Patch:megatron.deepseek_v4.flops_reporting] V4 closed-form FLOPs "
+        f"breakdown -- batch_size={batch_size}, seq_length={int(args.seq_length)}, "
+        f"hc_mult={int(getattr(args, 'hc_mult', 1) or 1)}, "
+        f"num_layers={int(args.num_layers)}, "
+        f"mtp_num_layers={int(getattr(args, 'mtp_num_layers', 0) or 0)}"
+    )
+    for name, tflops in rows:
+        log_rank_0(f"[Patch:megatron.deepseek_v4.flops_reporting]   {name:<12s} = {tflops:9.3f} TFLOP")
+    log_rank_0(
+        f"[Patch:megatron.deepseek_v4.flops_reporting]   {'TOTAL':<12s} = "
+        f"{total_flops / 1.0e12:9.3f} TFLOP / global-batch"
+    )
+
+
+def _make_v4_num_floating_point_operations(original_fn, *, dispatch_v4: bool):
+    """Return a wrapper that dispatches V4 vs upstream model types.
+
+    ``dispatch_v4`` is captured at install time from ``args.model_type``
+    rather than re-checked per call, because Megatron's ``pretrain()``
+    overwrites ``args.model_type`` with the ``ModelType`` enum at
+    ``training.py:1210`` *before* ``train()`` ever calls
+    ``num_floating_point_operations``.  At that point the original
+    YAML-set string ``"deepseek_v4"`` is gone and a runtime check would
+    silently fall through to the upstream formula.
+    """
+
+    def wrapped(args, batch_size):
+        if not dispatch_v4:
+            return original_fn(args, batch_size)
+
+        total_flops, breakdown = compute_v4_flops(args, batch_size)
+
+        global _BREAKDOWN_LOGGED
+        if not _BREAKDOWN_LOGGED:
+            _BREAKDOWN_LOGGED = True
+            _emit_breakdown(
+                args=args,
+                batch_size=batch_size,
+                breakdown=breakdown,
+                total_flops=total_flops,
+            )
+
+        return total_flops
+
+    wrapped.__wrapped__ = original_fn
+    wrapped._v4_flops_patched = True
+    return wrapped
+
+
+_TRAINER_REBIND_TARGETS: Sequence[str] = (
+    # Primus's Megatron trainer imports ``num_floating_point_operations`` at
+    # module load time (``primus.modules.trainer.megatron.trainer``: line 125)
+    # and resolves the bare name from its OWN globals at the call site
+    # (``trainer.train()``: line 1452).  Updating only
+    # ``megatron.training.training.num_floating_point_operations`` therefore
+    # never reaches that bound name and the trainer keeps using the upstream
+    # GPT/MLA-shaped function.  We rebind the trainer's local name to the
+    # wrapper too so the V4 closed form is what actually drives per-iter
+    # TFLOPs reporting.  Listed explicitly so a missing module is loud rather
+    # than silently silent.
+    "primus.modules.trainer.megatron.trainer",
+)
+
+
+def _rebind_trainer_imports(wrapped_fn) -> List[str]:
+    """Rebind ``num_floating_point_operations`` in every Primus module that
+    captured the upstream symbol at import time.
+
+    Returns the list of modules that were actually rebound so the install log
+    can show whether each downstream binding was wired up.  Modules that were
+    not yet imported (e.g. on a cold-cache trainer init) are skipped silently
+    — the trainer module imports the function at its own load, which
+    happens before ``before_train``, so in practice the targeted module is
+    always present at this point.
+    """
+    import sys
+
+    rebound: List[str] = []
+    for module_name in _TRAINER_REBIND_TARGETS:
+        mod = sys.modules.get(module_name)
+        if mod is None:
+            continue
+        if getattr(mod, "num_floating_point_operations", None) is wrapped_fn:
+            continue
+        if not hasattr(mod, "num_floating_point_operations"):
+            continue
+        mod.num_floating_point_operations = wrapped_fn
+        rebound.append(module_name)
+    return rebound
+
+
+@register_patch(
+    "megatron.deepseek_v4.flops_reporting",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "DeepSeek-V4: replace Megatron's GPT/MLA-shaped "
+        "num_floating_point_operations with a V4 closed form (Q LoRA + "
+        "single-latent KV + grouped low-rank O at S * hc_mult, plus "
+        "Compressor / Indexer side paths and MTP per-depth full inner "
+        "layer cost). Falls through byte-for-byte for non-V4 model types."
+    ),
+    condition=lambda ctx: getattr(get_args(ctx), "model_type", None) == "deepseek_v4",
+)
+def patch_v4_flops_reporting(ctx: PatchContext):
+    """Install the V4 FLOPs wrapper on ``training.num_floating_point_operations``."""
+    import megatron.training.training as training_module
+
+    original_fn = training_module.num_floating_point_operations
+    if getattr(original_fn, "_v4_flops_patched", False):
+        log_rank_0(
+            "[Patch:megatron.deepseek_v4.flops_reporting] "
+            "num_floating_point_operations already patched, skip"
+        )
+        return
+
+    # Capture the V4 dispatch decision NOW, while ``args.model_type`` is
+    # still the YAML-set string.  Megatron's ``pretrain()`` will rebind
+    # ``args.model_type`` to a ``ModelType`` enum at
+    # ``training.py:1210`` later, so any runtime check inside the wrapper
+    # would fail.  This patch is gated to only install when the YAML model
+    # type is V4 (see ``condition`` on the decorator), so it's safe to
+    # hard-set ``dispatch_v4=True`` here.
+    wrapped = _make_v4_num_floating_point_operations(original_fn, dispatch_v4=True)
+    training_module.num_floating_point_operations = wrapped
+    rebound = _rebind_trainer_imports(wrapped)
+
+    log_rank_0(
+        "[Patch:megatron.deepseek_v4.flops_reporting] wrapped "
+        "num_floating_point_operations; per-iter TFLOPs now reported "
+        "with V4-aware closed form (see "
+        "deepseek-v4/develop/plan-3/02-phase-details.md#p20 for the formula)."
+    )
+    if rebound:
+        log_rank_0(
+            "[Patch:megatron.deepseek_v4.flops_reporting] rebound trainer " f"import bindings: {rebound}"
+        )
+    else:
+        log_rank_0(
+            "[Patch:megatron.deepseek_v4.flops_reporting] no trainer modules "
+            "needed rebinding (none of "
+            f"{list(_TRAINER_REBIND_TARGETS)} were imported yet)."
+        )
+
+
+__all__: Sequence[str] = (
+    "compute_v4_flops",
+    "patch_v4_flops_reporting",
+)
diff --git a/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py
new file mode 100644
index 000000000..9fb0c2de5
--- /dev/null
+++ b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py
@@ -0,0 +1,271 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 hash-router input_ids PP broadcast (upfront/pre-loop).
+
+V4's hash-routed MoE layers (the first ``num_hash_layers`` MoE layers)
+look up a static ``tid2eid`` table and therefore need the raw
+``input_ids`` on every PP stage that owns one. Megatron's
+:func:`pretrain_gpt.get_batch` only loads tokens on the first / last PP
+stage; middle stages short-circuit to ``return None, None, None, ...``.
+
+Plan-2 P19 attempted two simpler hooks before this one:
+
+1. **In-forward broadcast** inside :meth:`DeepseekV4Model.forward`.
+2. **Per-call broadcast** wrapping :func:`pretrain_gpt.get_batch`.
+
+Both work for the non-interleaved 1F1B schedule (``PP>1``,
+``VPP=1``) but deadlock the interleaved-1F1B / VPP schedule. The reason:
+the interleaved scheduler issues a single ``recv_forward`` *before* the
+warm-up loop on every non-first PP rank (see
+:func:`forward_backward_pipelining_with_interleaving` in
+``megatron/core/pipeline_parallel/schedules.py:1363-1392``). PP rank > 0
+therefore parks in that ``recv_forward.wait()`` until PP rank 0's first
+``send_forward`` arrives. PP rank 0 meanwhile enters the warm-up loop,
+calls ``forward_step`` -> ``get_batch`` -> ``dist.broadcast`` and gets
+stuck waiting for PP rank > 0 to issue the matching broadcast — which
+will never happen because PP rank > 0 is itself blocked on
+``recv_forward``. The non-interleaved 1F1B schedule does not hit this
+because its pre-loop ``recv_forward`` is *inside* each warm-up iter
+(``schedules.py:2128-2156``), so PP rank > 0 reaches its broadcast call
+on the same iteration that PP rank 0 issues the matching ``send_forward``
+and the broadcast pairs up before either rank stalls.
+
+This patch instead does **all** PP token broadcasts up-front, *before*
+the schedule's first ``recv_forward``:
+
+* It wraps :func:`megatron.core.pipeline_parallel.get_forward_backward_func`
+  so that every schedule fetched for a train_step is replaced by a
+  thin wrapper which:
+    1. Pre-loads ``num_microbatches`` × ``num_chunks`` batches by calling
+       the original ``pretrain_gpt.get_batch`` on the first / last PP
+       stages, and allocating empty token buffers on middle PP stages.
+    2. Runs one ``dist.broadcast`` per (chunk, microbatch) on the PP
+       group, sourced from PP rank 0. All collectives fire before any
+       ``send_forward`` / ``recv_forward`` runs, so they pair up
+       deterministically across ranks — no deadlock.
+    3. Caches the resulting tuples in a module-local store keyed by
+       chunk and microbatch ordinal, then calls the underlying schedule
+       with the original ``data_iterator``.
+* It wraps :func:`pretrain_gpt.get_batch` so that, while the cache is
+  active, calls return the pre-cached tuple for the corresponding
+  (vp_stage, microbatch) — bypassing the data iterator. Outside a
+  train_step (e.g. during eval not routed through the schedule) the
+  wrapper falls back to the original ``get_batch``.
+* The cache is reset after each schedule call (success or exception)
+  so subsequent train_steps start clean.
+
+Gating: ``model_type == "deepseek_v4"``, ``num_hash_layers > 0``, and
+``pipeline_model_parallel_size > 1``. The patch is a strict no-op for
+any other model.
+
+Cost analysis (on the V4 BF16 smoke):
+
+* Each pre-broadcast moves ``mbs * seq * 8 B`` per microbatch
+  (~1 KiB for ``mbs=1, seq=128``). With ``num_microbatches=16,
+  num_chunks=2`` that is ~32 KiB total, dwarfed by activation P2P
+  (~32 MiB per microbatch).
+* Pre-loading consumes the data_iterator on PP rank 0 / last during
+  ``pre_broadcast`` rather than spread across the warm-up. This
+  collapses ``batch-generator`` time into a single bursty phase.
+"""
+
+from typing import Any, Optional
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+# ---------------------------------------------------------------------------
+# Module-local pre-broadcast cache.
+#
+# ``data[chunk_id][microbatch_id]`` holds the 6-tuple returned by
+# :func:`pretrain_gpt.get_batch`:
+#   (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params)
+#
+# On middle PP stages only ``tokens`` is meaningful; the other fields stay
+# ``None`` (matching what the original ``get_batch`` returns there).
+# ``consumed[chunk_id]`` tracks how many microbatches have been popped from
+# the cache during the current schedule call; it lets us return the right
+# entry from the patched ``get_batch`` regardless of which interleaved
+# microbatch ordinal is being processed (interleaved / 1F1B both consume
+# microbatches in increasing order *per chunk*).
+# ---------------------------------------------------------------------------
+_V4_PP_TOKEN_CACHE: dict = {
+    "active": False,
+    "data": [],
+    "consumed": [],
+}
+
+
+def _v4_reset_cache() -> None:
+    _V4_PP_TOKEN_CACHE["active"] = False
+    _V4_PP_TOKEN_CACHE["data"] = []
+    _V4_PP_TOKEN_CACHE["consumed"] = []
+
+
+def _v4_pre_broadcast_step(
+    data_iterator: Any,
+    num_microbatches: int,
+    original_get_batch,
+) -> None:
+    """Pre-broadcast V4 ``input_ids`` for every (chunk, microbatch) of one step."""
+    from megatron.core import parallel_state
+    from megatron.training import get_args as _get_args
+
+    args = _get_args()
+    pp_group = parallel_state.get_pipeline_model_parallel_group()
+    src_global = torch.distributed.get_global_rank(pp_group, 0)
+    pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+    pp_world = parallel_state.get_pipeline_model_parallel_world_size()
+
+    cp_size = max(1, int(getattr(args, "context_parallel_size", 1) or 1))
+    mbs = int(args.micro_batch_size)
+    seq = int(args.seq_length) // cp_size
+    device = torch.device("cuda", torch.cuda.current_device())
+
+    # ``is_first_or_last_pp`` mirrors :func:`pretrain_gpt.get_batch`'s
+    # gate: the original returns real tokens only on the first / last PP
+    # stage. (MTP can also build the dataset on a middle stage; that
+    # case is left to the original short-circuit because the smoke runs
+    # ``mtp_num_layers=0`` and the broadcast logic still tolerates a
+    # ``None``-tuple result by falling back to an empty token buffer.)
+    is_first_or_last_pp = (pp_rank == 0) or (pp_rank == pp_world - 1)
+
+    if isinstance(data_iterator, list):
+        iter_list = data_iterator
+        num_chunks = len(iter_list)
+        # When VPP is enabled, each chunk has its own ``vp_stage`` index.
+        chunk_vp_stages = list(range(num_chunks))
+    else:
+        iter_list = [data_iterator]
+        num_chunks = 1
+        # Non-VPP: ``vp_stage`` is ``None`` (matches what
+        # ``pretrain_gpt.forward_step`` reads off the model attribute).
+        chunk_vp_stages = [None]
+
+    cache_data: list = []
+    for chunk_id in range(num_chunks):
+        chunk_iter = iter_list[chunk_id]
+        chunk_vp_stage = chunk_vp_stages[chunk_id]
+        chunk_cache: list = []
+        for _mb_id in range(num_microbatches):
+            if is_first_or_last_pp and chunk_iter is not None:
+                tup = original_get_batch(chunk_iter, chunk_vp_stage)
+                tokens = tup[0]
+                if tokens is None:
+                    # Should not happen on first / last PP for V4 smoke;
+                    # fall back to empty buffer and let the broadcast
+                    # populate. This keeps us robust against unusual
+                    # configs (e.g. dataset-on-rank gating disabled).
+                    tokens = torch.empty([mbs, seq], dtype=torch.long, device=device)
+                cached = (tokens, tup[1], tup[2], tup[3], tup[4], tup[5])
+            else:
+                tokens = torch.empty([mbs, seq], dtype=torch.long, device=device)
+                cached = (tokens, None, None, None, None, None)
+
+            torch.distributed.broadcast(tokens, src=src_global, group=pp_group)
+            chunk_cache.append(cached)
+        cache_data.append(chunk_cache)
+
+    _V4_PP_TOKEN_CACHE["data"] = cache_data
+    _V4_PP_TOKEN_CACHE["consumed"] = [0] * num_chunks
+    _V4_PP_TOKEN_CACHE["active"] = True
+
+
+def _make_v4_get_batch(original_get_batch):
+    """Return a ``get_batch`` wrapper that consumes the pre-broadcast cache."""
+
+    def patched_get_batch(data_iterator: Any, vp_stage: Optional[int] = None):
+        if not _V4_PP_TOKEN_CACHE["active"]:
+            return original_get_batch(data_iterator, vp_stage)
+        chunk_id = vp_stage if vp_stage is not None else 0
+        counter = _V4_PP_TOKEN_CACHE["consumed"][chunk_id]
+        cached = _V4_PP_TOKEN_CACHE["data"][chunk_id][counter]
+        _V4_PP_TOKEN_CACHE["consumed"][chunk_id] += 1
+        return cached
+
+    patched_get_batch.__wrapped__ = original_get_batch
+    patched_get_batch._v4_pp_get_batch_patched = True
+    return patched_get_batch
+
+
+def _make_v4_pre_broadcast_schedule(original_schedule, original_get_batch):
+    """Wrap ``forward_backward_func`` to pre-broadcast V4 tokens up-front."""
+
+    def patched_schedule(*args, **kwargs):
+        data_iterator = kwargs.get("data_iterator")
+        num_microbatches = int(kwargs.get("num_microbatches", 1) or 1)
+
+        try:
+            _v4_pre_broadcast_step(data_iterator, num_microbatches, original_get_batch)
+            return original_schedule(*args, **kwargs)
+        finally:
+            _v4_reset_cache()
+
+    patched_schedule._v4_pp_schedule_wrapped = True
+    return patched_schedule
+
+
+@register_patch(
+    "megatron.deepseek_v4.pp_token_pre_broadcast",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "DeepSeek-V4: pre-broadcast input_ids from PP rank 0 across the PP "
+        "group up-front in the forward_backward schedule wrapper, before the "
+        "first recv_forward, so middle PP stages owning hash-routed MoE "
+        "layers can read raw token IDs without deadlocking the interleaved "
+        "1F1B / VPP schedule."
+    ),
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "model_type", None) == "deepseek_v4"
+        and int(getattr(get_args(ctx), "num_hash_layers", 0) or 0) > 0
+        and int(getattr(get_args(ctx), "pipeline_model_parallel_size", 1) or 1) > 1
+    ),
+    # Ordered after pp_dump_data so its schedule_wrapper does not double-wrap;
+    # see ``pp_dump_data_patches.py`` for the priority=100 anchor.
+    priority=60,
+)
+def patch_v4_pp_token_pre_broadcast(ctx: PatchContext):
+    """Install the V4 PP token pre-broadcast hooks."""
+    import megatron.core.pipeline_parallel as pp_module
+    import megatron.training.training as training_module
+    import pretrain_gpt
+
+    original_get_batch = pretrain_gpt.get_batch
+    if getattr(original_get_batch, "_v4_pp_get_batch_patched", False):
+        log_rank_0("[Patch:megatron.deepseek_v4.pp_token_pre_broadcast] get_batch " "already patched, skip")
+        return
+
+    # Hook 1: replace pretrain_gpt.get_batch with a cache-consuming wrapper.
+    # We capture ``original_get_batch`` here so the schedule wrapper can call
+    # the *unpatched* implementation during the pre-broadcast phase (the
+    # patched version would just hit an empty cache and recurse).
+    pretrain_gpt.get_batch = _make_v4_get_batch(original_get_batch)
+
+    # Hook 2: replace get_forward_backward_func so every schedule fetched
+    # by ``train_step`` is wrapped with the pre-broadcast.
+    original_get_fbf = pp_module.get_forward_backward_func
+
+    def wrapped_get_fbf():
+        original_schedule = original_get_fbf()
+        if getattr(original_schedule, "_v4_pp_schedule_wrapped", False):
+            return original_schedule
+        return _make_v4_pre_broadcast_schedule(original_schedule, original_get_batch)
+
+    pp_module.get_forward_backward_func = wrapped_get_fbf
+    training_module.get_forward_backward_func = wrapped_get_fbf
+
+    log_rank_0(
+        "[Patch:megatron.deepseek_v4.pp_token_pre_broadcast] wrapped "
+        "pretrain_gpt.get_batch + get_forward_backward_func; PP rank 0 "
+        "broadcasts input_ids up-front (once per microbatch × chunk per "
+        "train_step) so middle PP stages owning hash-routed MoE layers "
+        "see real token IDs without per-(chunk, microbatch) collectives "
+        "racing the interleaved 1F1B P2P sends."
+    )
diff --git a/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py
new file mode 100644
index 000000000..b163c24fe
--- /dev/null
+++ b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py
@@ -0,0 +1,162 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 PP tensor-shape patch.
+
+V4's HyperConnections (mHC) residual carries ``K = hc_mult`` parallel
+streams per position. Inside a single PP stage the layer hidden has
+shape ``[B, S, K, D]``; at the PP boundary the V4 transformer block
+folds K into the sequence axis via
+:func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block._lower_streams_out`
+so the wire tensor is ``[S * K, B, D]`` (plan-2 P15 design — see
+``deepseek-v4/develop/plan-2/03-phase-details.md`` P15 and the C1
+finding in ``00-review-findings.md``).
+
+Megatron has **two** code paths that compute the PP wire tensor shape,
+and they need *both* to know about V4's K packing:
+
+1. The non-interleaved 1F1B schedule (``forward_backward_pipelining_
+   without_interleaving``) calls
+   :func:`megatron.core.pipeline_parallel.schedules.get_tensor_shapes`
+   (``schedules.py:2096-2103``).  Wrapping that function suffices for
+   smokes A/B/C.
+
+2. The interleaved 1F1B / VPP schedule (``forward_backward_pipelining_
+   with_interleaving``) instead computes ``tensor_shape`` *inline* from
+   the ``seq_length`` argument (``schedules.py:1001-1004``):
+
+       tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
+       tensor_shape[0] = tensor_shape[0] // cp_group.size()
+       if config.sequence_parallel:
+           tensor_shape[0] = tensor_shape[0] // tp_group.size()
+
+   It does **not** call ``get_tensor_shapes``. With VPP=2 and ``hc_mult=
+   4`` this leaves the recv buffer at ``[S, B, hidden]`` while the
+   sender emits ``[S*K, B, hidden]``. PyTorch P2P does not validate
+   shape (only ``numel * dtype_size``), so the receiver silently copies
+   only the first ``S * hidden`` elements, ``_lift_streams_in``
+   reshapes them as ``[B, S/K, K, D]``, and the resulting hidden
+   flattens to ``S/K = 32`` instead of ``S = 128`` — :class:`Deepseek
+   V4HashRouter` then trips its precondition with ``hidden=32 vs
+   token_ids=128``. (P19 smoke D run reproduces this cleanly.)
+
+We therefore install **two** complementary wrappers, both gated on the
+V4 ``model_type`` + ``hc_mult > 1`` + ``PP > 1`` condition:
+
+* :func:`_make_v4_get_tensor_shapes` multiplies the first (seq) dim of
+  every tuple returned by ``get_tensor_shapes`` by ``hc_mult``, fixing
+  path (1).
+* :func:`_make_v4_interleaved_schedule` wraps
+  ``forward_backward_pipelining_with_interleaving`` to scale its
+  ``seq_length`` keyword argument by ``hc_mult`` before the schedule
+  computes its inline ``tensor_shape``, fixing path (2).  Inside the
+  interleaved schedule, ``seq_length`` is *only* read for that inline
+  shape (see grep on ``schedules.py``), so scaling it is a no-op for
+  every other concern (cudagraph, attention masking, etc.).
+
+The companion ``adjust_tensor_shapes_fn`` parameter on the
+non-interleaved 1F1B schedule is intentionally not used because
+upstream Megatron explicitly asserts it is unsupported by the
+interleaved / VPP schedules (``schedules.py:900-901``); going through
+``get_tensor_shapes`` and ``seq_length`` keeps the behavior uniform
+across all schedules.
+"""
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _make_v4_get_tensor_shapes(original_fn, hc_mult: int):
+    """Return a wrapper that scales the first (seq) dim by ``hc_mult``.
+
+    ``original_fn`` returns a list of ``(seq_length, micro_batch, hidden)``
+    triples; the V4 PP wire packs ``K`` streams into the seq axis so we
+    multiply the seq dim only — micro_batch and hidden are unchanged.
+    """
+
+    def patched_get_tensor_shapes(*args, **kwargs):
+        shapes = original_fn(*args, **kwargs)
+        return [(s * hc_mult, b, h) for (s, b, h) in shapes]
+
+    patched_get_tensor_shapes.__wrapped__ = original_fn
+    patched_get_tensor_shapes._v4_pp_shape_patched = True
+    return patched_get_tensor_shapes
+
+
+def _make_v4_interleaved_schedule(original_fn, hc_mult: int):
+    """Wrap the interleaved schedule to scale ``seq_length`` by ``hc_mult``.
+
+    The interleaved schedule (``forward_backward_pipelining_with_
+    interleaving``) computes its PP wire ``tensor_shape`` inline from
+    the ``seq_length`` kwarg (``schedules.py:1001``). Inside that
+    function, ``seq_length`` is consumed only by that inline
+    computation, so we can safely scale it on the way in to give the
+    schedule a V4-aware shape without touching any other behaviour.
+    """
+
+    def patched_schedule(*args, **kwargs):
+        if "seq_length" in kwargs and kwargs["seq_length"] is not None:
+            kwargs["seq_length"] = int(kwargs["seq_length"]) * hc_mult
+        return original_fn(*args, **kwargs)
+
+    patched_schedule.__wrapped__ = original_fn
+    patched_schedule._v4_pp_interleaved_patched = True
+    return patched_schedule
+
+
+@register_patch(
+    "megatron.deepseek_v4.pp_tensor_shape",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "DeepSeek-V4: pack hc_mult=K hyper-streams into the PP wire seq "
+        "axis so [S*K, B, D] passes between PP stages (covers both the "
+        "1F1B get_tensor_shapes path and the interleaved-1F1B / VPP "
+        "inline tensor_shape path)."
+    ),
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "model_type", None) == "deepseek_v4"
+        and int(getattr(get_args(ctx), "hc_mult", 1) or 1) > 1
+        and int(getattr(get_args(ctx), "pipeline_model_parallel_size", 1) or 1) > 1
+    ),
+)
+def patch_v4_pp_tensor_shape(ctx: PatchContext):
+    """Multiply the PP P2P seq dim by ``hc_mult`` for V4 models."""
+    import megatron.core.pipeline_parallel.schedules as schedules_module
+
+    hc_mult = int(getattr(get_args(ctx), "hc_mult", 1))
+
+    # Wrapper 1: get_tensor_shapes (used by the non-interleaved schedule).
+    original_get_tensor_shapes = schedules_module.get_tensor_shapes
+    if getattr(original_get_tensor_shapes, "_v4_pp_shape_patched", False):
+        log_rank_0("[Patch:megatron.deepseek_v4.pp_tensor_shape] get_tensor_shapes " "already patched, skip")
+    else:
+        schedules_module.get_tensor_shapes = _make_v4_get_tensor_shapes(original_get_tensor_shapes, hc_mult)
+        log_rank_0(
+            f"[Patch:megatron.deepseek_v4.pp_tensor_shape] wrapped "
+            f"get_tensor_shapes; PP wire seq_len * hc_mult={hc_mult} "
+            "(packs K hyper-streams into the sequence axis)."
+        )
+
+    # Wrapper 2: forward_backward_pipelining_with_interleaving (VPP).
+    # The interleaved schedule reads ``seq_length`` directly to build its
+    # inline ``tensor_shape``; scaling the kwarg gives it the V4 wire shape
+    # without rewriting the function.
+    original_interleaved = schedules_module.forward_backward_pipelining_with_interleaving
+    if getattr(original_interleaved, "_v4_pp_interleaved_patched", False):
+        log_rank_0(
+            "[Patch:megatron.deepseek_v4.pp_tensor_shape] interleaved " "schedule already patched, skip"
+        )
+    else:
+        schedules_module.forward_backward_pipelining_with_interleaving = _make_v4_interleaved_schedule(
+            original_interleaved, hc_mult
+        )
+        log_rank_0(
+            f"[Patch:megatron.deepseek_v4.pp_tensor_shape] wrapped "
+            f"forward_backward_pipelining_with_interleaving; "
+            f"seq_length * hc_mult={hc_mult} on the way into the "
+            "interleaved-1F1B / VPP schedule."
+        )
diff --git a/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py b/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py
new file mode 100644
index 000000000..e230ad2de
--- /dev/null
+++ b/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py
@@ -0,0 +1,342 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron train_step patch for delayed FP8 scaling updates and preamble
+optimization.
+
+Patch 1 -- Delayed FP8 scale update:
+    Wraps train_step to call the scale update preamble on every FP8 delayed
+    module before the original train_step runs.  For most_recent +
+    history_len=1 (production config), a fast path stages per-module amaxes
+    into the registry-batched ``staged_amaxes_3n`` / ``scales_3n`` tensors,
+    computes new scales with a few fused ops, and scatters them back to the
+    per-module scalar ``m.scale_*`` buffers.  Per-module scalars (rather
+    than views into a shared ``(N,)`` tensor) are required so that each
+    module's buffer storage is independent for torch.compile version
+    tracking.
+
+Patch 2 -- Grad-zero stream overlap + data HtoD prefetch:
+    Dispatches DDP grad-buffer zeroing (grad_data.zero_()) on a secondary
+    CUDA stream so it overlaps with HtoD data transfer, saving ~4.6 ms/iter.
+    Additionally, prefetches the next data batch to GPU on a dedicated HtoD
+    stream.  The prefetch iterator is injected at the forward_backward_func
+    level (not by replacing data_iterator in train_step) to preserve
+    compatibility with Megatron's RerunDataIterator type assertion.
+"""
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+# Module-level handle so the MLPerf warmup hook (or any external code) can
+# reach the CudaPrefetchIterator state owned by
+# ``patch_grad_zero_and_data_prefetch`` (which would otherwise live solely
+# in that patch's closure and be unreachable from outside).
+#
+# The warmup hook needs this so it can invalidate the cached prefetch
+# iterator at the end of the warmup epilogue.  Without that invalidation,
+# the prefetcher stays bound to the synthetic-data iterator that was passed
+# in for warmup step 1, and every subsequent real training step silently
+# reads from the cycling synthetic dataset (because
+# ``MegatronDataloaderWrapper`` is cyclic and never raises ``StopIteration``).
+_PREFETCH_HANDLE: dict = {"state": None}
+
+# Shared state for async amax allreduce between Patch 1 (scale update) and
+# Patch 2 (fwd/bwd wrapper).  Patch 2 launches the async allreduce after
+# forward_backward_func returns; Patch 1 waits on the handle at the start
+# of the next train_step before computing new scales.
+#
+# Realistic overlap window: the time between forward_backward_func returning
+# (in step N) and the Patch 1 wait/compute at the start of step N+1.  In
+# practice that's roughly ``optimizer.step + grad_zero`` -- not the full
+# train_step.  The overlap is still useful but smaller than the all-reduce
+# cost in most configs, so this is best-effort latency hiding rather than a
+# free win.
+_ASYNC_AMAX_HANDLE: dict = {"handle": None, "registry": None}
+
+
+def _reset_async_amax_state():
+    """Drop any pending async amax allreduce handle.
+
+    Idempotent best-effort cleanup invoked from both call sites in this
+    module (Patch 1's wait path and Patch 2's launch path).  Leaving a
+    stale handle in ``_ASYNC_AMAX_HANDLE["handle"]`` would cause the next
+    train_step to try to wait on an already-consumed or never-launched
+    work object and either deadlock or raise.
+    """
+    _ASYNC_AMAX_HANDLE["handle"] = None
+
+
+def get_prefetch_state():
+    """Return the closure-shared ``_prefetch_state`` dict, or ``None`` if
+    ``patch_grad_zero_and_data_prefetch`` has not been installed yet."""
+    return _PREFETCH_HANDLE.get("state")
+
+
+def reset_prefetch_state():
+    """Drop the cached ``CudaPrefetchIterator`` so the next ``train_step``
+    rebuilds it around its current ``data_iterator`` argument.
+
+    Required after MLPerf warmup, because warmup step 1 is the first call
+    into ``_patched_train_step``, so the prefetcher gets bound to the
+    synthetic iterator.  Without invalidation, every real training step
+    afterwards substitutes the cached prefetcher (still wrapping synthetic
+    data) for the real ``data_iterator`` argument inside the
+    ``_synced_prefetch_fwd_bwd`` wrapper.
+
+    Returns the evicted iterator (or ``None`` if no cache existed) so
+    callers can log what was dropped.
+    """
+    state = _PREFETCH_HANDLE.get("state")
+    if state is None:
+        return None
+    return state.pop("iter", None)
+
+
+def _needs_delayed_scaling(ctx: PatchContext) -> bool:
+    # Imported lazily (not at module top-level) so that importing this patch
+    # module never depends on ``megatron`` being importable. A top-level
+    # ``from megatron...`` here can crash the patches package auto-import if
+    # megatron is momentarily unavailable (e.g. mid circular-import), which
+    # then leaves the package in a half-initialized state.
+    from megatron.core.enums import Fp8Recipe
+
+    args = get_args(ctx)
+    if args is None or not bool(getattr(args, "fp8", False)):
+        return False
+    return (
+        getattr(args, "fp8_scaling_strategy", "dynamic") == "delayed"
+        or getattr(args, "fp8_recipe", None) == Fp8Recipe.delayed
+    ) and not getattr(args, "disable_delayed_scaling_patches", False)
+
+
+@register_patch(
+    "megatron.fp8.delayed_scaling_update",
+    backend="megatron",
+    phase="before_train",
+    description="Wrap train_step to update delayed FP8 scales before each step.",
+    priority=40,
+    condition=_needs_delayed_scaling,
+)
+def patch_delayed_fp8_update(ctx: PatchContext):
+    import megatron.training.training as megatron_training
+
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _DelayedScalingRegistry,
+        _fast_update_scales,
+        _fast_update_scales_with_history,
+        _wait_and_compute_scales,
+    )
+    from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+
+    _PATCH_KEY = "megatron.fp8.delayed_scaling_update"
+    if is_patched(megatron_training, _PATCH_KEY):
+        log_rank_0("[Patch:delayed_scaling_update] Already applied; skipping re-wrap.")
+        return
+
+    _original_train_step = megatron_training.train_step
+    _cached_delayed_modules = []
+    _registry = None
+
+    def _patched_train_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        config,
+        forward_backward_func,
+        iteration=None,
+    ):
+        nonlocal _registry
+        if not _cached_delayed_modules:
+            _cached_delayed_modules.extend(
+                m
+                for model_chunk in model
+                for m in model_chunk.modules()
+                if getattr(m, "_use_delayed_scaling", False)
+            )
+        if _cached_delayed_modules:
+            if _registry is None:
+                _registry = _DelayedScalingRegistry(_cached_delayed_modules)
+                _ASYNC_AMAX_HANDLE["registry"] = _registry
+
+            pending_handle = _ASYNC_AMAX_HANDLE.get("handle")
+            if pending_handle is not None:
+                try:
+                    _wait_and_compute_scales(_registry, pending_handle)
+                finally:
+                    # Always drop the handle, even if wait / compute raised:
+                    # leaving a consumed handle here would cause the next
+                    # train_step to try to wait on it again.
+                    _reset_async_amax_state()
+            else:
+                if _registry.algo == "most_recent" and _registry.history_len == 1:
+                    _fast_update_scales(_registry)
+                else:
+                    _fast_update_scales_with_history(_registry)
+        return _original_train_step(
+            forward_step_func,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            config,
+            forward_backward_func,
+            iteration=iteration,
+        )
+
+    megatron_training.train_step = _patched_train_step
+    mark_patched(megatron_training, _PATCH_KEY)
+    log_rank_0(
+        "[Patch:delayed_scaling_update] "
+        "Wrapped train_step with async amax allreduce support. "
+        "First step uses synchronous fallback; subsequent steps wait on "
+        "async handle launched by post-fwd_bwd hook."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Patch 2: Grad-zero stream overlap + data HtoD prefetch
+# ---------------------------------------------------------------------------
+
+
+@register_patch(
+    "megatron.grad_zero_and_data_prefetch",
+    backend="megatron",
+    phase="before_train",
+    description="Overlap grad buffer zeroing and data HtoD transfer via secondary CUDA streams.",
+    priority=41,
+    condition=_needs_delayed_scaling,
+)
+def patch_grad_zero_and_data_prefetch(ctx: PatchContext):
+    args = get_args(ctx)
+    if getattr(args, "reuse_grad_buf_for_mxfp8_param_ag", False):
+        log_rank_0(
+            "[Patch:grad_zero_and_data_prefetch] SKIPPED — "
+            "reuse_grad_buf_for_mxfp8_param_ag is set (shared param/grad buffer)."
+        )
+        return
+
+    import megatron.training.training as megatron_training
+    from megatron.core.distributed import DistributedDataParallel
+
+    from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+
+    _PATCH_KEY = "megatron.grad_zero_and_data_prefetch"
+    if is_patched(megatron_training, _PATCH_KEY):
+        log_rank_0("[Patch:grad_zero_and_data_prefetch] Already applied; skipping re-wrap.")
+        return
+
+    tp_size = getattr(args, "tensor_model_parallel_size", 1)
+
+    _original_train_step = megatron_training.train_step
+    _zero_stream = torch.cuda.Stream()
+    _prefetch_state: dict = {}
+    # Expose the closure-local prefetch state so the MLPerf warmup hook can
+    # invalidate the cached prefetch iterator at the end of warmup; see
+    # ``reset_prefetch_state`` above for the rationale.
+    _PREFETCH_HANDLE["state"] = _prefetch_state
+
+    def _stream_zero_grad_buffer(self):
+        """CPU metadata on main thread; GPU grad_data.zero_() on secondary stream."""
+        if getattr(self.config, "cuda_graph_impl", "none") != "transformer_engine":
+            for param in self.params_with_grad:
+                param.grad_added_to_main_grad = False
+        with torch.cuda.stream(_zero_stream):
+            for buffer in self.buffers + self.expert_parallel_buffers:
+                buffer.reset()
+        for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+            bucket_group.reset()
+
+    DistributedDataParallel.zero_grad_buffer = _stream_zero_grad_buffer
+
+    def _patched_train_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        config,
+        forward_backward_func,
+        iteration=None,
+    ):
+        if "iter" not in _prefetch_state and tp_size == 1:
+            if isinstance(data_iterator, (list, tuple)):
+                # Virtual pipeline parallel passes a list of per-chunk iterators,
+                # which the single-stream prefetcher cannot wrap. Skip prefetch
+                # (a pure HtoD-overlap optimization) and let the original iterator
+                # flow through unchanged.
+                if not _prefetch_state.get("vpp_skip_logged"):
+                    log_rank_0(
+                        "[Patch:grad_zero_and_data_prefetch] "
+                        "Skipping CudaPrefetchIterator: data_iterator is a list "
+                        "(virtual pipeline parallel)."
+                    )
+                    _prefetch_state["vpp_skip_logged"] = True
+            else:
+                from primus.backends.megatron.data.cuda_prefetch import (
+                    CudaPrefetchIterator,
+                )
+
+                compute_dtype = torch.bfloat16 if getattr(args, "bf16", False) else torch.float16
+                _prefetch_state["iter"] = CudaPrefetchIterator(
+                    data_iterator,
+                    compute_dtype=compute_dtype,
+                )
+                log_rank_0(
+                    "[Patch:grad_zero_and_data_prefetch] "
+                    f"Created CudaPrefetchIterator (dtype={compute_dtype})."
+                )
+
+        _pf = _prefetch_state.get("iter")
+
+        def _synced_prefetch_fwd_bwd(*fwd_args, **fwd_kwargs):
+            torch.cuda.current_stream().wait_stream(_zero_stream)
+            if _pf is not None:
+                fwd_kwargs["data_iterator"] = _pf
+            result = forward_backward_func(*fwd_args, **fwd_kwargs)
+
+            registry = _ASYNC_AMAX_HANDLE.get("registry")
+            if registry is not None:
+                from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+                    _stage_and_launch_async_allreduce,
+                )
+
+                try:
+                    _ASYNC_AMAX_HANDLE["handle"] = _stage_and_launch_async_allreduce(registry)
+                except Exception:
+                    # Launch failed: clear so the next train_step takes the
+                    # synchronous fallback path instead of trying to wait on
+                    # a non-existent handle.
+                    _reset_async_amax_state()
+                    raise
+
+            return result
+
+        return _original_train_step(
+            forward_step_func,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            config,
+            _synced_prefetch_fwd_bwd,
+            iteration=iteration,
+        )
+
+    megatron_training.train_step = _patched_train_step
+    mark_patched(megatron_training, _PATCH_KEY)
+    log_rank_0(
+        "[Patch:grad_zero_and_data_prefetch] "
+        "DDP grad_data.zero_() on secondary stream; "
+        "data HtoD prefetch on secondary stream; "
+        "secondary streams synced before forward_backward_func; "
+        "async amax allreduce launched after forward_backward_func "
+        "(awaited at the start of the next train_step)."
+    )
diff --git a/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py b/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py
new file mode 100644
index 000000000..3f11a2c05
--- /dev/null
+++ b/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py
@@ -0,0 +1,58 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Raise the ``emerging_optimizers`` (absl) logger to INFO.
+
+The Muon path pulls in ``emerging_optimizers``, which logs via ``absl`` (``from
+absl import logging``). Its Newton-Schulz helper emits a per-call ``DEBUG`` line
+for the coefficient schedule (``muon_utils.get_coefficient_iterator``)::
+
+    Iterating through 10 steps with cycle mode.
+    Coefficient sets: [(3.4445, -4.775, 2.0315), ..., (2.0, -1.5, 0.5)]
+
+Because Muon orthogonalizes every matrix parameter on every optimizer step, this
+fires repeatedly and floods the logs.
+
+All ``absl`` logging records flow through the single stdlib logger named
+``"absl"`` (Primus routes stdlib logging into loguru via an InterceptHandler on
+the root logger, with ``root`` at NOTSET, so these DEBUG records are emitted).
+Raising the ``"absl"`` logger to INFO drops the DEBUG spam at the source while
+keeping INFO and above. We do it in the ``before_train`` phase, which runs after
+the optimizer is built but before the first training step.
+
+Set ``PRIMUS_VERBOSE_EMERGING_OPTIMIZERS=1`` to keep the full DEBUG trace.
+"""
+
+import logging
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.emerging_optimizers.log_level",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Raise the 'absl' logger (used by emerging_optimizers) to INFO so the "
+        "Muon Newton-Schulz coefficient DEBUG lines do not flood the logs."
+    ),
+    condition=lambda ctx: os.environ.get("PRIMUS_VERBOSE_EMERGING_OPTIMIZERS", "0") != "1",
+)
+def patch_emerging_optimizers_log_level(ctx: PatchContext):
+    """Set the absl logger threshold to INFO. Idempotent / best-effort."""
+    del ctx
+    try:
+        logging.getLogger("absl").setLevel(logging.INFO)
+    except Exception as exc:  # never block training over a logging tweak.
+        log_rank_0(
+            f"[Patch:megatron.emerging_optimizers.log_level] could not set " f"'absl' logger level: {exc!r}"
+        )
+        return
+
+    log_rank_0("[Patch:megatron.emerging_optimizers.log_level] 'absl' logger raised to INFO.")
diff --git a/primus/backends/megatron/patches/empty_cache_interval_patches.py b/primus/backends/megatron/patches/empty_cache_interval_patches.py
new file mode 100644
index 000000000..5f4ad4b4e
--- /dev/null
+++ b/primus/backends/megatron/patches/empty_cache_interval_patches.py
@@ -0,0 +1,211 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Megatron empty_cache() interval patch.
+
+PROBLEM
+-------
+Megatron's ``train_step`` calls ``torch.cuda.empty_cache()`` before every
+``optimizer.step()`` when ``args.empty_unused_memory_level >= 1``
+(megatron/training/training.py:1759).
+
+On the Pure-GDN 1B / 100B-tokens MI300X run profiling showed this single
+call was responsible for ~4.6 s of ``hipMalloc`` and ~2.3 s of ``hipFree``
+EVERY iter — 91% of the 7.65 s/iter wall time.  Removing it entirely
+crashes NCCL on iter 1 with ``Failed to CUDA calloc 4 MiB`` because
+NCCL's lazy workspace allocation needs a contiguous block that
+fragmentation prevents at 81% VRAM usage.
+
+SOLUTION
+--------
+Call ``empty_cache()`` only every N iters instead of every iter.  Once
+NCCL's workspace is allocated on the iter where ``empty_cache()`` ran,
+it stays cached on subsequent iters that DO NOT call ``empty_cache()``
+— so we avoid the per-iter ``hipMalloc`` cost while still periodically
+returning fragmented cached blocks to the driver as a safety net.
+
+CONFIGURATION
+-------------
+The interval is set, in order of precedence:
+
+    1. ``empty_cache_interval: N`` in the EXP YAML's ``overrides:`` block
+       (preferred — co-located with the rest of the run config)
+    2. ``PRIMUS_EMPTY_CACHE_INTERVAL=N`` env var (for ad-hoc overrides
+       without editing the YAML; also lets a single launcher script set
+       a default for backwards compatibility)
+    3. Default ``1`` (passthrough — no behavioural change vs vanilla
+       Megatron)
+
+Values:
+    - 1     : ORIGINAL behaviour — call empty_cache() every iter
+              (only when args.empty_unused_memory_level >= 1, which is
+              the gating condition Megatron itself uses; this patch
+              NEVER enables empty_cache when Megatron's flag is 0).
+    - N >= 2: call empty_cache() only on iters where
+              ((iteration_count_since_train_start) % N == 0), i.e. once
+              every N iters.  Iter 0 ALWAYS runs empty_cache (so the
+              first NCCL workspace allocation succeeds against a clean
+              cache).
+    - 0     : NEVER call empty_cache() in train_step (CAUTION: risks the
+              NCCL OOM if the model is memory-tight; only safe if NCCL
+              has already been warmed up via some other mechanism).
+
+The patch leaves ``args.empty_unused_memory_level``'s OTHER call sites
+(checkpoint save in training.py:3214, and the level-2 call after
+optimizer.step in training.py:1803) UNAFFECTED — those run rarely
+enough that they are not a per-iter cost concern.
+
+LOSS IMPACT
+-----------
+None.  ``torch.cuda.empty_cache()`` only returns unused cached blocks
+to the driver — it does not modify any tensor data.  The only thing
+this patch changes is HOW OFTEN we return those cached blocks, which
+is purely a memory-bookkeeping decision.
+
+TURNING OFF
+-----------
+Set ``empty_cache_interval: 1`` in the EXP YAML (or unset the env var)
+to get the original per-iter behaviour back.  Alternatively, set
+``empty_unused_memory_level: 0`` in the EXP YAML — the patch is
+short-circuited when Megatron's own flag is 0 (which is also the
+Megatron default for non-OOM-prone runs).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_DEFAULT_INTERVAL = 1
+
+
+def _coerce(value, source: str) -> Optional[int]:
+    """Best-effort parse of a user-supplied value to int; warn on bad input."""
+    if value is None:
+        return None
+    try:
+        return max(int(value), 0)
+    except (TypeError, ValueError):
+        # Best-effort warn; if logger is uninitialised (e.g. unit tests) we
+        # fall back to print to avoid a secondary AttributeError.
+        msg = f"[Patch:megatron.empty_cache_interval] WARN: invalid " f"{source}={value!r}, ignoring."
+        try:
+            log_rank_0(msg)
+        except Exception:
+            print(msg)
+        return None
+
+
+def _resolve_interval(args) -> int:
+    """Resolve the effective empty_cache_interval.
+
+    Precedence: args.empty_cache_interval > PRIMUS_EMPTY_CACHE_INTERVAL env > 1.
+    """
+    yaml_val = _coerce(getattr(args, "empty_cache_interval", None), "args.empty_cache_interval (YAML)")
+    if yaml_val is not None:
+        return yaml_val
+
+    env_val = _coerce(os.environ.get("PRIMUS_EMPTY_CACHE_INTERVAL"), "PRIMUS_EMPTY_CACHE_INTERVAL env")
+    if env_val is not None:
+        return env_val
+
+    return _DEFAULT_INTERVAL
+
+
+@register_patch(
+    "megatron.training.empty_cache_interval.wrap_train_step",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Skip torch.cuda.empty_cache() in train_step on non-trigger iters when "
+        "empty_cache_interval > 1.  Eliminates the ~5 s/iter hipMalloc/hipFree "
+        "thrash on the Pure-GDN 1B run while still flushing periodically."
+    ),
+)
+def patch_train_step_with_empty_cache_interval(ctx: PatchContext) -> None:
+    import megatron.training.training as training  # type: ignore
+    from megatron.training.global_vars import get_args as get_megatron_args
+
+    original_train_step = training.train_step
+    if getattr(original_train_step, "_primus_empty_cache_interval_wrapped", False):
+        return
+
+    counter = {"n": 0, "interval": None, "logged": False}
+
+    def _train_step_with_interval(*args, **kwargs):
+        mg_args = get_megatron_args()
+
+        # Resolve once on the first call.  We can't do this in patch
+        # registration time because args isn't fully built yet.
+        if counter["interval"] is None:
+            counter["interval"] = _resolve_interval(mg_args)
+            interval = counter["interval"]
+            src = (
+                "YAML(empty_cache_interval)"
+                if getattr(mg_args, "empty_cache_interval", None) is not None
+                else (
+                    "env(PRIMUS_EMPTY_CACHE_INTERVAL)"
+                    if os.environ.get("PRIMUS_EMPTY_CACHE_INTERVAL") is not None
+                    else "default"
+                )
+            )
+            mode = (
+                "every iter (no-op)"
+                if interval == 1
+                else ("NEVER (risky)" if interval == 0 else f"every {interval} iters")
+            )
+            log_rank_0(
+                f"[Patch:megatron.empty_cache_interval] empty_cache_interval={interval} "
+                f"({mode}); source={src}; "
+                f"empty_unused_memory_level={getattr(mg_args, 'empty_unused_memory_level', 0)}"
+            )
+            counter["logged"] = True
+
+        interval = counter["interval"]
+        # Original gating in training.py:1759 is `if args.empty_unused_memory_level >= 1`.
+        # We only intervene when the user actually wanted empty_cache (so we never
+        # accidentally enable it).  Behaviour:
+        #   - interval == 1 : passthrough (every iter behaviour unchanged)
+        #   - interval >= 2 : on iters NOT divisible by interval, temporarily
+        #                     downgrade empty_unused_memory_level to 0 so
+        #                     train_step skips the empty_cache call.  On the
+        #                     trigger iter, restore the original value so the
+        #                     empty_cache fires as designed.
+        #   - interval == 0 : always downgrade (never call empty_cache in
+        #                     train_step).  Risky; documented above.
+        orig_level = getattr(mg_args, "empty_unused_memory_level", 0)
+        downgrade = False
+
+        if orig_level >= 1 and interval != 1:
+            if interval == 0:
+                downgrade = True
+            else:
+                # iter 0 (first call) ALWAYS triggers so NCCL's first allocation
+                # happens against a clean cache.  After that, fire on every Nth iter.
+                if counter["n"] != 0 and (counter["n"] % interval != 0):
+                    downgrade = True
+
+        try:
+            if downgrade:
+                # Megatron reads args.empty_unused_memory_level inline in
+                # train_step.  Temporarily clear it for this call only.
+                mg_args.empty_unused_memory_level = 0
+            return original_train_step(*args, **kwargs)
+        finally:
+            if downgrade:
+                mg_args.empty_unused_memory_level = orig_level
+            counter["n"] += 1
+
+    setattr(_train_step_with_interval, "_primus_empty_cache_interval_wrapped", True)
+    training.train_step = _train_step_with_interval
+
+    log_rank_0(
+        "[Patch:megatron.empty_cache_interval] Wrapped train_step(); "
+        "actual interval will be resolved + logged on first train_step call."
+    )
diff --git a/primus/backends/megatron/patches/env_patches.py b/primus/backends/megatron/patches/env_patches.py
index 79df41f9b..52e3a5781 100644
--- a/primus/backends/megatron/patches/env_patches.py
+++ b/primus/backends/megatron/patches/env_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -13,7 +13,7 @@
 import os
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_kv_rank_0
+from primus.core.utils.module_utils import log_kv_rank_0
 
 # ============================================================================
 # CUDA Device Configuration
@@ -25,6 +25,11 @@
     backend="megatron",
     phase="setup",
     description="Set CUDA_DEVICE_MAX_CONNECTIONS based on FSDP configuration",
+    condition=lambda ctx: not getattr(
+        getattr(ctx.extra.get("module_config"), "params", None),
+        "disable_env_patches",
+        False,
+    ),
 )
 def set_cuda_device_max_connections(ctx: PatchContext):
     """
diff --git a/primus/backends/megatron/patches/evaluate_patches.py b/primus/backends/megatron/patches/evaluate_patches.py
index 2c1932e5e..108c19d1f 100644
--- a/primus/backends/megatron/patches/evaluate_patches.py
+++ b/primus/backends/megatron/patches/evaluate_patches.py
@@ -13,7 +13,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/fla_runtime_patches.py b/primus/backends/megatron/patches/fla_runtime_patches.py
new file mode 100644
index 000000000..07bfbc313
--- /dev/null
+++ b/primus/backends/megatron/patches/fla_runtime_patches.py
@@ -0,0 +1,151 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+FLA Runtime Knob Patch
+======================
+
+Resolves FLA runtime toggles onto the Megatron ``args`` namespace so that
+every consumer — Primus-owned code AND Megatron-LM patches — can read a
+single, typed ``args.*`` attribute instead of parsing env-var strings.
+
+PRECEDENCE
+----------
+For every knob the resolved value is, in priority order:
+
+    1. Pre-existing env var (highest — for ad-hoc overrides at launch time)
+    2. YAML field in the EXP overrides block
+    3. Documented default (matches pre-cleanup behaviour)
+
+After resolution the value is written back onto ``args`` as a properly
+typed attribute (bool / int / str).  Consumers never touch ``os.environ``.
+
+CONFIGURATION
+-------------
+Add any of the following fields to the model YAML (e.g. ``mamba_base.yaml``
+or the experiment ``overrides:`` block).  All are optional; unspecified
+fields keep their documented default and the patch is a no-op for that
+knob.
+
+    # --- FLA Triton kernel toggles -------------------------------------------
+    use_fla_fused_swiglu: true          # default true
+    use_fla_fused_rmsnorm: false        # default false
+    use_fla_fused_gated_norm: false     # default false  (same semantic scope
+                                        #   as use_fla_fused_rmsnorm but kept
+                                        #   as a separate knob for clarity)
+    use_fla_short_conv: false           # default false
+
+    # --- FLA dataset shim (deterministic FLA-order data) ---------------------
+    use_fla_data: false                 # default false
+    fla_cache_dir: ""                   # default ""
+
+    # --- Fused cross-entropy from FLA ----------------------------------------
+    fused_ce_mode: 1                    # 0 = vanilla Megatron CE
+                                        # 1 = chunked FLA fused CE (default)
+                                        # 2 = single-shot FLA fused CE
+    fused_ce_chunks: 32                 # default 32
+
+    # --- FLA MLA attention backend -------------------------------------------
+    fla_mla_attn: ""                    # default unset / ""
+
+TIMING
+------
+The patch runs at ``phase="before_train"`` — after
+``train_runtime.py:merge_namespace`` has merged the Primus-only YAML keys
+into ``args``, but before ``wrapped_pretrain()`` triggers model module
+imports.  Consumers in model code read ``args.*`` in ``__init__`` or later,
+which is always after this patch has run.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+# ─── Knob definitions ────────────────────────────────────────────────────────
+#
+# (yaml_field, legacy_env_var, typed_default)
+#
+# typed_default carries the type: bool → consumers see True/False,
+# int → consumers see an int, str → consumers see a str.
+#
+# Legacy env var is checked FIRST (backward compat); if unset the YAML
+# field value is used; if also null/missing the default applies.
+# ──────────────────────────────────────────────────────────────────────────────
+
+_FLA_RUNTIME_KNOBS: tuple = (
+    ("use_fla_fused_swiglu", "PRIMUS_FLA_SWIGLU", True),
+    ("use_fla_fused_rmsnorm", "PRIMUS_FLA_NORM", False),
+    ("use_fla_fused_gated_norm", "PRIMUS_FLA_NORM", False),
+    ("use_fla_short_conv", "PRIMUS_FLA_CONV", False),
+    ("use_fla_data", "PRIMUS_FLA_DATA", False),
+    ("fla_cache_dir", "PRIMUS_FLA_CACHE_DIR", ""),
+    ("fused_ce_mode", "PRIMUS_FUSED_CE", 1),
+    ("fused_ce_chunks", "PRIMUS_FUSED_CE_CHUNKS", 32),
+    ("fla_mla_attn", "PRIMUS_FLA_MLA_ATTN", ""),
+)
+
+
+def _env_to_typed(env_string: str, default: Any) -> Any:
+    """Convert an env-var string to the type implied by *default*."""
+    if isinstance(default, bool):
+        return env_string not in ("0", "", "false", "False")
+    if isinstance(default, int):
+        return int(env_string)
+    return env_string
+
+
+def _has_any_fla_runtime_field(args) -> bool:
+    """Cheap probe — skip the patch entirely when no FLA knob is configured
+    AND no legacy env var is set."""
+    for name, env_name, _default in _FLA_RUNTIME_KNOBS:
+        if getattr(args, name, None) is not None:
+            return True
+        if env_name in os.environ:
+            return True
+    return False
+
+
+@register_patch(
+    "megatron.fla_runtime_knobs",
+    backend="megatron",
+    # Phase MUST be "before_train" (not "build_args").  At build_args time
+    # the Primus-only YAML keys have not yet been merged into args
+    # (MegatronArgBuilder.convert_config strips unknown keys;
+    # train_runtime.py:294 merge_namespace re-adds them BEFORE
+    # before_train runs).  Model module imports happen even later — inside
+    # wrapped_pretrain() — so args.* values are available to every
+    # consumer that reads them in __init__ or forward.
+    phase="before_train",
+    priority=-100,
+    description=(
+        "Resolve FLA runtime knobs (env var > YAML > default) onto args.* "
+        "attributes.  Consumers read args directly; no os.environ access."
+    ),
+    condition=lambda ctx: _has_any_fla_runtime_field(get_args(ctx)),
+)
+def patch_fla_runtime_knobs(ctx: PatchContext):
+    args = get_args(ctx)
+
+    for field_name, env_name, default in _FLA_RUNTIME_KNOBS:
+        env_raw = os.environ.get(env_name)
+        yaml_value = getattr(args, field_name, None)
+
+        if env_raw is not None:
+            resolved = _env_to_typed(env_raw, default)
+            source = f"env {env_name}={env_raw!r}"
+        elif yaml_value is not None:
+            resolved = yaml_value
+            source = f"YAML {field_name}={yaml_value!r}"
+        else:
+            resolved = default
+            source = f"default"
+
+        setattr(args, field_name, resolved)
+        log_rank_0(f"[Patch:megatron.fla_runtime_knobs] " f"args.{field_name} = {resolved!r}  ({source})")
diff --git a/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py
new file mode 100644
index 000000000..e1180b9b9
--- /dev/null
+++ b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py
@@ -0,0 +1,98 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron train_step patch for FP8 all-gather cache refresh.
+
+The FP8 all-gather cache (populated by precompute_fp8_scales_for_fsdp) must be
+refreshed after every optimizer.step() so that the cached FP8 data reflects the
+updated weights.  Without this patch the cache is only populated once during
+FSDP setup and goes stale after the first iteration.
+
+This patch wraps Megatron's standalone train_step() to call
+precompute_fp8_scales_for_fsdp(model[0]) after the original train_step returns.
+"""
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _needs_fp8_cache_update(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return args is not None and getattr(args, "use_fsdp2_fp8_all_gather", False)
+
+
+@register_patch(
+    "megatron.training.train_step_fp8_cache_update",
+    backend="megatron",
+    phase="before_train",
+    description="Patch train_step to refresh FP8 all-gather cache after optimizer.step()",
+    condition=_needs_fp8_cache_update,
+    priority=45,
+)
+def patch_train_step_fp8_cache(ctx: PatchContext):
+    import megatron.training.training as megatron_training
+
+    from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+        precompute_fp8_scales_for_fsdp,
+    )
+    from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+
+    _PATCH_KEY = "megatron.training.train_step_fp8_cache_update"
+    if is_patched(megatron_training, _PATCH_KEY):
+        log_rank_0("[Patch:train_step_fp8_cache_update] Already applied; skipping re-wrap.")
+        return
+
+    args = get_args(ctx)
+    cache_data = getattr(args, "fp8_precompute_data_cache", True)
+    use_cpp = getattr(args, "use_cpp_fp8_quantize", False)
+    stochastic_rounding = getattr(args, "fp8_all_gather_stochastic_rounding", False)
+
+    _original_train_step = megatron_training.train_step
+
+    def _patched_train_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        config,
+        forward_backward_func,
+        iteration=None,
+    ):
+        result = _original_train_step(
+            forward_step_func,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            config,
+            forward_backward_func,
+            iteration=iteration,
+        )
+        if _refresh_scales:
+            precompute_fp8_scales_for_fsdp(
+                model[0],
+                cache_data=cache_data,
+                use_cpp_quantize=use_cpp,
+                stochastic_rounding=stochastic_rounding,
+            )
+        return result
+
+    # When using the C++ quantize kernel without data caching, scales were
+    # already set once during FSDP setup and the C++ kernel uses them
+    # directly -- no per-step refresh needed (matches run_35 behavior).
+    _refresh_scales = cache_data or not use_cpp
+
+    megatron_training.train_step = _patched_train_step
+    mark_patched(megatron_training, _PATCH_KEY)
+    log_rank_0(
+        "[Patch:train_step_fp8_cache_update] "
+        f"Patched train_step to refresh FP8 all-gather cache after optimizer.step() "
+        f"(cache_data={cache_data}, use_cpp_quantize={use_cpp}, "
+        f"stochastic_rounding={stochastic_rounding}, "
+        f"refresh_scales={_refresh_scales})"
+    )
diff --git a/primus/backends/megatron/patches/fused_linear_ce_patches.py b/primus/backends/megatron/patches/fused_linear_ce_patches.py
new file mode 100644
index 000000000..7333a2222
--- /dev/null
+++ b/primus/backends/megatron/patches/fused_linear_ce_patches.py
@@ -0,0 +1,186 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+###############################################################################
+# Chunked linear cross-entropy for Megatron GPTModel (avoid full-logits OOM).
+#
+# PROBLEM: Megatron's stock GPTModel._postprocess materializes the FULL logits
+# tensor [seq, batch, vocab] before cross-entropy. For a large vocab (Qwen
+# 151936) and long sequence (e.g. 64k), that single tensor is tens of GB and
+# OOMs -- even though a 1.5B model's weights/activations are tiny. This is the
+# exact wall we hit validating ODC/LB-Mini on DeepSeek-R1-Distill-Qwen-1.5B.
+#
+# FIX (verl/Liger spirit): split along the sequence dim and compute
+# logits+CE per chunk under activation checkpointing, so the full [seq, vocab]
+# logits is NEVER resident -- peak logits memory is one chunk only. This lets
+# TP=1 / DP=8 (the layout ODC requires) run 64k-token sequences without OOM.
+#
+# ZERO IMPACT on stock Megatron:
+#   * Pure monkey-patch in the Primus layer; third-party Megatron source is NOT
+#     touched.
+#   * Gated by enable_fused_linear_ce (yaml) or FUSED_LINEAR_CE=1 (env).
+#     DEFAULT OFF -> the original full-logits _postprocess runs byte-for-byte.
+#   * Even when ON, only the training-with-labels path is intercepted;
+#     inference / no-labels (generation) / MTP / non-post-process stages all
+#     fall through to the original implementation.
+#   * Numerically equivalent: per-token CE is independent across positions, so
+#     chunk-then-concatenate yields the identical [b, s] per-token loss.
+###############################################################################
+
+import os
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCHED = False
+
+
+def _fused_ce_enabled(args) -> bool:
+    """Switch can come from the YAML arg or, for experiments, the env var."""
+    return bool(getattr(args, "enable_fused_linear_ce", False)) or (
+        os.environ.get("FUSED_LINEAR_CE", "0") == "1"
+    )
+
+
+def _chunk_size() -> int:
+    """Per-chunk sequence length. Smaller = less peak logits mem, more recompute."""
+    return int(os.environ.get("FUSED_CE_CHUNK", "0") or 0) or 4096
+
+
+class _ChunkedLinearCE(torch.autograd.Function):
+    """verl-style fused linear cross-entropy.
+
+    forward: compute the LM loss in sequence chunks WITHOUT building an autograd
+        graph and WITHOUT saving per-chunk logits -> the full [s, vocab] logits
+        is never materialized (peak = one chunk).
+    backward: recompute each chunk's logits and obtain grads via a SINGLE
+        torch.autograd.grad call per chunk, inside ONE backward invocation.
+
+    Why not torch.utils.checkpoint: its recompute is driven by an unpack hook
+    that runs INTERLEAVED with FSDP2/ODC's backward communication. Under ODC the
+    ranks run DIFFERENT micro-batch counts, so that interleaving desyncs and
+    DEADLOCKS (the iter-3 hang). Doing the recompute here -- a plain local
+    autograd.grad with NO collective -- keeps each rank independent in backward.
+    """
+
+    @staticmethod
+    def forward(ctx, hidden, weight, labels, model, chunk):
+        ctx.model = model
+        ctx.chunk = chunk
+        ctx.save_for_backward(hidden, weight, labels)
+        seq = hidden.size(0)
+        parts = []
+        with torch.no_grad():
+            for i in range(0, seq, chunk):
+                logits = model._scale_logits(torch.matmul(hidden[i : i + chunk], weight.t()))
+                parts.append(model.compute_language_model_loss(labels[:, i : i + chunk].contiguous(), logits))
+        return torch.cat(parts, dim=1)  # [b, s]
+
+    @staticmethod
+    def backward(ctx, grad_out):  # grad_out [b, s]
+        hidden, weight, labels = ctx.saved_tensors
+        model = ctx.model
+        chunk = ctx.chunk
+        seq = hidden.size(0)
+        grad_hidden = torch.empty_like(hidden)
+        grad_weight = torch.zeros_like(weight)
+        for i in range(0, seq, chunk):
+            h_c = hidden[i : i + chunk].detach().requires_grad_(True)
+            w = weight.detach().requires_grad_(True)
+            with torch.enable_grad():
+                logits = model._scale_logits(torch.matmul(h_c, w.t()))
+                loss_c = model.compute_language_model_loss(labels[:, i : i + chunk].contiguous(), logits)
+            g_h, g_w = torch.autograd.grad(loss_c, (h_c, w), grad_out[:, i : i + chunk].contiguous())
+            grad_hidden[i : i + chunk] = g_h
+            grad_weight = grad_weight + g_w
+        return grad_hidden, grad_weight, None, None, None
+
+
+def _chunked_lm_loss(model, hidden_states, labels):
+    """Chunked linear+CE that never materializes the full [s, b, vocab] logits.
+
+    hidden_states: [s, b, h]   labels: [b, s]   -> per-token loss [b, s]
+    """
+    # PREFERRED PATH (ODC + DiffMicro): use the full output weight that ODC's
+    # train-loop hook already all-gathered ONCE at the minibatch boundary
+    # (pre_minibatch_start). Calling full_tensor() here -- inside the per-micro
+    # batch forward -- would issue a DTensor collective whose call-count differs
+    # across ranks under DiffMicro, and deadlock. The cached tensor is a leaf
+    # with requires_grad=True; its .grad accumulates across all this rank's
+    # micro-batches and is reduce-scattered back to the sharded param at
+    # pre_optimizer_step (see odc_torch_fsdp2_patches._odc_reduce_output_grad).
+    cached = getattr(model, "_odc_cached_output_weight", None)
+    if cached is not None:
+        return _ChunkedLinearCE.apply(hidden_states, cached, labels, model, _chunk_size())
+
+    # FALLBACK (no ODC, or SameMicro): same-count collectives are safe.
+    if model.share_embeddings_and_output_weights:
+        ow = model.shared_embedding_or_output_weight()
+    else:
+        ow = model.output_layer.weight
+    if hasattr(ow, "full_tensor"):
+        ow_full = ow.full_tensor()
+    else:
+        ow_full = ow.clone()
+    return _ChunkedLinearCE.apply(hidden_states, ow_full, labels, model, _chunk_size())
+
+
+def _install_fused_ce_patch():
+    global _PATCHED
+    if _PATCHED:
+        return
+    import megatron.core.models.gpt.gpt_model as gpt_mod
+
+    GPTModel = gpt_mod.GPTModel
+    if getattr(GPTModel._postprocess, "_fused_ce_hooked", False):
+        _PATCHED = True
+        return
+    orig_postprocess = GPTModel._postprocess
+
+    def postprocess_with_fused_ce(self, *args, **kwargs):
+        # forward() calls _postprocess() with all-kwargs, but support positional
+        # too: signature is (hidden_states, input_ids, position_ids, labels, ...).
+        hidden_states = kwargs.get("hidden_states", args[0] if len(args) > 0 else None)
+        labels = kwargs.get("labels", args[3] if len(args) > 3 else None)
+        inference_context = kwargs.get("inference_context")
+        in_inference = inference_context is not None and not self.training
+
+        if (
+            labels is not None
+            and hidden_states is not None
+            and not in_inference
+            and getattr(self, "post_process", True)
+            and not getattr(self.config, "mtp_num_layers", 0)
+        ):
+            return _chunked_lm_loss(self, hidden_states, labels)
+        # Everything else -> stock path, unchanged.
+        return orig_postprocess(self, *args, **kwargs)
+
+    postprocess_with_fused_ce._fused_ce_hooked = True
+    GPTModel._postprocess = postprocess_with_fused_ce
+    _PATCHED = True
+    log_rank_0(
+        f"[FusedCE] patched GPTModel._postprocess: chunked linear+CE "
+        f"(chunk={_chunk_size()}); full [seq, vocab] logits no longer materialized."
+    )
+
+
+@register_patch(
+    "megatron.fused_linear_ce",
+    backend="megatron",
+    phase="before_train",
+    description="Chunked linear cross-entropy for GPTModel to avoid full-logits OOM (large vocab + long seq).",
+    condition=lambda ctx: _fused_ce_enabled(get_args(ctx)),
+)
+def patch_fused_linear_ce(ctx: PatchContext):
+    log_rank_0(
+        "[FusedCE] enable_fused_linear_ce=true -> installing chunked linear+CE "
+        "(avoids materializing full [seq, vocab] logits; numerically equivalent)."
+    )
+    _install_fused_ce_patch()
+
+
+__all__ = ["patch_fused_linear_ce"]
diff --git a/primus/backends/megatron/patches/fused_pad_routing_map_patches.py b/primus/backends/megatron/patches/fused_pad_routing_map_patches.py
new file mode 100644
index 000000000..7702dc538
--- /dev/null
+++ b/primus/backends/megatron/patches/fused_pad_routing_map_patches.py
@@ -0,0 +1,68 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus fused_pad_routing_map patch.
+
+Globally replaces Megatron's ``fused_pad_routing_map`` with the Primus Triton
+implementation (``primus.backends.megatron.core.fusions.fused_pad_routing_map``)
+without modifying upstream Megatron source.
+
+The Primus version rewrites the kernel to operate directly on the native
+``[num_tokens, num_experts]`` layout (no transpose/copy) and drops the
+``@jit_fuser`` (``torch.compile``) wrapper, avoiding the Triton kernel
+functionalization failure seen on some torch/triton combos.
+"""
+
+import sys
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.fused_pad_routing_map",
+    backend="megatron",
+    phase="before_train",
+    description="Replace Megatron fused_pad_routing_map with the Primus Triton implementation",
+)
+def patch_fused_pad_routing_map(ctx: PatchContext):
+    """Swap the ``fused_pad_routing_map`` symbol everywhere it is referenced."""
+    import megatron.core.fusions.fused_pad_routing_map as meg_mod
+
+    from primus.backends.megatron.core.fusions.fused_pad_routing_map import (
+        fused_pad_routing_map as primus_fused_pad_routing_map,
+    )
+
+    log_rank_0("[Patch:megatron.fused_pad_routing_map] Patching fused_pad_routing_map...")
+
+    # Original function object; used to precisely locate stale references.
+    orig_fn = getattr(meg_mod, "fused_pad_routing_map", None)
+    if orig_fn is primus_fused_pad_routing_map:
+        log_rank_0("[Patch:megatron.fused_pad_routing_map]   Already patched; skipping.")
+        return
+
+    # 1) Replace on the source module so all *future* (incl. lazy) imports resolve
+    #    to the Primus version. This alone covers the common case, since this patch
+    #    runs before token_dispatcher is imported.
+    meg_mod.fused_pad_routing_map = primus_fused_pad_routing_map
+
+    # 2) Replace references already bound into other modules that imported the symbol
+    #    at top level before this patch ran (precise: only objects that `is orig_fn`).
+    patched_modules = []
+    if orig_fn is not None:
+        for mod_name, module in list(sys.modules.items()):
+            if module is None or module is meg_mod:
+                continue
+            if getattr(module, "fused_pad_routing_map", None) is orig_fn:
+                setattr(module, "fused_pad_routing_map", primus_fused_pad_routing_map)
+                patched_modules.append(mod_name)
+
+    log_rank_0(
+        "[Patch:megatron.fused_pad_routing_map]   Patched "
+        "megatron.core.fusions.fused_pad_routing_map.fused_pad_routing_map "
+        f"-> primus (also updated already-imported refs: {patched_modules or 'none'})"
+    )
diff --git a/primus/backends/megatron/patches/gdn_config_patches.py b/primus/backends/megatron/patches/gdn_config_patches.py
new file mode 100644
index 000000000..84353e2a2
--- /dev/null
+++ b/primus/backends/megatron/patches/gdn_config_patches.py
@@ -0,0 +1,140 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+GatedDeltaNet / KimiDeltaAttention Configuration Patches
+
+Monkey-patch TransformerConfig with linear-attention fields required by
+GatedDeltaNet and KimiDeltaAttention layers, so that no changes are needed
+in the third-party Megatron-LM codebase.
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_GDN_CONFIG_FIELDS = {
+    "linear_conv_kernel_dim": None,
+    "use_short_conv": True,
+    "linear_key_head_dim": None,
+    "linear_value_head_dim": None,
+    "linear_num_key_heads": None,
+    "linear_num_value_heads": None,
+    "use_fla_triton_kda": False,
+    "use_fla_triton_kda_hybrid": False,
+    # When True (default) and use_fla_triton_kda is also True, chunk_kda is
+    # called with use_gate_in_kernel=True (gate fused inside the Triton
+    # kernel). Set to False to materialize the gate up-front via
+    # fused_kda_gate() — bit-identical to FLA's pre-fusion path and to the
+    # tw006-validated numerics (loss=4.7281 @ iter 500 vs FLA/8=4.7350).
+    "use_fla_kda_in_kernel_gate": True,
+    # When True (default when use_fla_triton_kda=True), the output norm is
+    # replaced by fla.modules.FusedRMSNormGated (RMSNorm + sigmoid-gate +
+    # multiply in one Triton kernel). Set to False to use the unfused
+    # _apply_gated_norm path with explicit fp32 sigmoid (tw006 numerics).
+    "use_fla_fused_norm_gated": None,
+}
+
+
+def _has_any_gdn_field(args) -> bool:
+    return any(getattr(args, name, None) is not None for name in _GDN_CONFIG_FIELDS)
+
+
+@register_patch(
+    "megatron.transformer.gdn_config",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Monkey-patch TransformerConfig with linear-attention fields "
+        "(linear_conv_kernel_dim, linear_key_head_dim, etc.) and FLA Triton flags "
+        "required by GatedDeltaNet and KimiDeltaAttention without modifying third-party code."
+    ),
+    condition=lambda ctx: _has_any_gdn_field(get_args(ctx)),
+)
+def patch_gdn_config(ctx: PatchContext):
+    args = get_args(ctx)
+
+    import megatron.core.transformer.transformer_config as config_mod
+
+    for field_name, default in _GDN_CONFIG_FIELDS.items():
+        value = getattr(args, field_name, default)
+        setattr(config_mod.TransformerConfig, field_name, value)
+        log_rank_0(f"[Patch:megatron.transformer.gdn_config] " f"TransformerConfig.{field_name} = {value}")
+
+
+# -----------------------------------------------------------------------------
+# Hybrid-model output-layer init method
+# -----------------------------------------------------------------------------
+#
+# Upstream TransformerConfig.__post_init__ scales the output-layer init std
+# by depth (scaled_init_method_normal) or by mu-P width/depth
+# (mup_scaled_init_method_normal) -- both appropriate for pure transformers,
+# but not for hybrid models (GDN/KDA/Mamba), which FLA initializes with a
+# plain uniform `initializer_range` (init_method_normal, no depth scaling).
+# Without this, hybrid runs start from a different output-layer init than
+# FLA's reference training and the early loss curve doesn't line up.
+
+_HYBRID_INIT_PATCH_KEY = "megatron.transformer.hybrid_output_init"
+
+
+def _is_hybrid_model_run(args) -> bool:
+    """True for any HybridStack-based model (GDN/KDA/Mamba, pure or mixed
+    with MLA). Primus YAMLs set ``is_hybrid_model`` directly (it's a plain
+    ``TransformerConfig`` field, generically copied from ``args`` by
+    ``core_transformer_config_from_args``); fall back to the
+    ``hybrid_override_pattern`` / ``hybrid_layer_pattern`` derivation
+    upstream uses (see ``megatron.training.utils.is_hybrid_model``) in case
+    a config relies on that instead."""
+    return bool(
+        getattr(args, "is_hybrid_model", False)
+        or getattr(args, "hybrid_override_pattern", None)
+        or getattr(args, "hybrid_layer_pattern", None)
+    )
+
+
+def _install_hybrid_output_init_patch() -> None:
+    import megatron.core.transformer.transformer_config as config_mod
+    from megatron.core.utils import init_method_normal
+
+    TransformerConfig = config_mod.TransformerConfig
+    if is_patched(TransformerConfig, _HYBRID_INIT_PATCH_KEY):
+        log_rank_0(f"[Patch:{_HYBRID_INIT_PATCH_KEY}] TransformerConfig already patched; skipping.")
+        return
+
+    original_post_init = TransformerConfig.__post_init__
+
+    def patched_post_init(self):
+        # Only overriding output_layer_init_method's own None-guard mirrors
+        # upstream exactly: if the caller already set it explicitly, leave
+        # it untouched (same as upstream's `if self.output_layer_init_method
+        # is None:` gate).
+        was_none = self.output_layer_init_method is None
+        original_post_init(self)
+        if was_none and getattr(self, "is_hybrid_model", False):
+            self.output_layer_init_method = init_method_normal(self.init_method_std)
+
+    TransformerConfig.__post_init__ = patched_post_init
+    mark_patched(TransformerConfig, _HYBRID_INIT_PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_HYBRID_INIT_PATCH_KEY}] Patched TransformerConfig.__post_init__: "
+        "hybrid models now use uniform init_method_normal for the output layer "
+        "(matching FLA's initializer_range) instead of depth/mu-P-scaled init."
+    )
+
+
+@register_patch(
+    _HYBRID_INIT_PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Hybrid models (GDN, KDA, Mamba) use uniform init_method_normal for the "
+        "output layer, matching FLA's initializer_range, instead of the "
+        "depth/mu-P-scaled init appropriate only for pure transformers."
+    ),
+    condition=lambda ctx: _is_hybrid_model_run(get_args(ctx)),
+)
+def patch_hybrid_output_init(ctx: PatchContext):
+    _install_hybrid_output_init_patch()
diff --git a/primus/backends/megatron/patches/gdn_rocm_gate_patches.py b/primus/backends/megatron/patches/gdn_rocm_gate_patches.py
new file mode 100644
index 000000000..a7884703f
--- /dev/null
+++ b/primus/backends/megatron/patches/gdn_rocm_gate_patches.py
@@ -0,0 +1,135 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron upstream GatedDeltaNet ROCm gate patch
+==============================================
+
+Qwen3.5 and other models with ``experimental_attention_variant: gated_delta_net``
+use Megatron-LM's ``megatron.core.ssm.gated_delta_net.GatedDeltaNet`` (not the
+Primus HybridStack copy).  On ROCm, FLA's Triton ``chunk_gated_delta_rule`` can
+produce NaN gradients when the gate is fused inside the kernel
+(``use_gate_in_kernel=True``, the default).  Pre-computing ``g`` in fp32 and
+passing ``use_gate_in_kernel=False`` matches the Primus hybrid GDN path and
+avoids NaN gradients during backward on ROCm.
+"""
+
+import torch
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.backends.megatron.patches._source_patch_utils import patch_method_source
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.ssm.gdn_rocm_gate"
+
+# Unique anchor in GatedDeltaNet.forward (megatron/core/ssm/gated_delta_net.py).
+_GATE_ORI = (
+    "        g = -self.A_log.exp() * F.softplus(alpha.float() + self.dt_bias)  # In fp32\n"
+    "        beta = beta.sigmoid()\n"
+    '        nvtx_range_pop(suffix="g_and_beta")\n'
+    "\n"
+    '        nvtx_range_push(suffix="gated_delta_rule")\n'
+    "        if self.config.deterministic_mode:\n"
+    "            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(\n"
+    "                query,\n"
+    "                key,\n"
+    "                value,\n"
+    "                g=g,\n"
+    "                beta=beta,\n"
+    "                initial_state=None,\n"
+    "                output_final_state=False,\n"
+    "                use_qk_l2norm_in_kernel=False,\n"
+    "            )\n"
+    "        else:\n"
+    "            core_attn_out, last_recurrent_state = chunk_gated_delta_rule(\n"
+    "                query,\n"
+    "                key,\n"
+    "                value,\n"
+    "                g=g,\n"
+    "                beta=beta,\n"
+    "                initial_state=None,\n"
+    "                output_final_state=False,\n"
+    "                use_qk_l2norm_in_kernel=False,\n"
+    "            )"
+)
+
+_GATE_NEW = (
+    "        # Pre-compute gate in fp32; use_gate_in_kernel=False avoids FLA Triton NaN on ROCm.\n"
+    "        g = -self.A_log.float().exp() * F.softplus(alpha.float() + self.dt_bias)\n"
+    "        beta = beta.sigmoid()\n"
+    '        nvtx_range_pop(suffix="g_and_beta")\n'
+    "\n"
+    '        nvtx_range_push(suffix="gated_delta_rule")\n'
+    "        if self.config.deterministic_mode:\n"
+    "            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(\n"
+    "                query,\n"
+    "                key,\n"
+    "                value,\n"
+    "                g=g,\n"
+    "                beta=beta,\n"
+    "                initial_state=None,\n"
+    "                output_final_state=False,\n"
+    "                use_qk_l2norm_in_kernel=False,\n"
+    "            )\n"
+    "        else:\n"
+    "            core_attn_out, last_recurrent_state = chunk_gated_delta_rule(\n"
+    "                query,\n"
+    "                key,\n"
+    "                value,\n"
+    "                g=g.to(query.dtype),\n"
+    "                beta=beta,\n"
+    "                initial_state=None,\n"
+    "                output_final_state=False,\n"
+    "                use_qk_l2norm_in_kernel=False,\n"
+    "                use_gate_in_kernel=False,\n"
+    "            )"
+)
+
+
+def _is_rocm(ctx: PatchContext) -> bool:
+    """Return True when running on an AMD ROCm platform."""
+    return getattr(torch.version, "hip", None) is not None
+
+
+def _uses_gated_delta_net(args) -> bool:
+    return getattr(args, "experimental_attention_variant", None) == "gated_delta_net"
+
+
+def _should_patch_gdn_rocm_gate(ctx: PatchContext) -> bool:
+    return _is_rocm(ctx) and _uses_gated_delta_net(get_args(ctx))
+
+
+def _install_gdn_rocm_gate_patch() -> None:
+    from megatron.core.ssm.gated_delta_net import GatedDeltaNet
+
+    if is_patched(GatedDeltaNet, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] GatedDeltaNet already patched; skipping.")
+        return
+
+    patch_method_source(GatedDeltaNet, "forward", _GATE_ORI, _GATE_NEW)
+
+    mark_patched(GatedDeltaNet, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] Patched GatedDeltaNet.forward: fp32 gate precompute + "
+        "use_gate_in_kernel=False for FLA chunk_gated_delta_rule on ROCm."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Pre-compute GDN gate in fp32 and disable FLA in-kernel gate fusion on "
+        "Megatron upstream GatedDeltaNet to avoid NaN gradients on ROCm."
+    ),
+    condition=_should_patch_gdn_rocm_gate,
+    priority=51,
+    tags=["rocm", "gdn"],
+)
+def patch_gdn_rocm_gate(ctx: PatchContext) -> None:
+    _install_gdn_rocm_gate_patch()
diff --git a/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py b/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py
index 3847f0022..378d839cf 100644
--- a/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py
+++ b/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py
@@ -11,7 +11,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/lr_schedule_patches.py b/primus/backends/megatron/patches/lr_schedule_patches.py
new file mode 100644
index 000000000..616a274f1
--- /dev/null
+++ b/primus/backends/megatron/patches/lr_schedule_patches.py
@@ -0,0 +1,68 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NeMo-aligned LR warmup patch.
+
+NeMo's WarmupHoldPolicy._get_warmup_lr (nemo/core/optim/lr_scheduler.py)
+computes warmup as:
+
+    lr = base_lr * (step + 1) / (warmup_steps + 1)
+
+Megatron's OptimizerParamScheduler.get_lr uses:
+
+    lr = init_lr + (max_lr - init_lr) * num_steps / lr_warmup_steps
+
+In Megatron's sample-space (num_steps increments by GBS per iteration,
+lr_warmup_steps = warmup_iters * GBS), the NeMo-equivalent formula is:
+
+    lr = init_lr + (max_lr - init_lr) * (num_steps + GBS) / (lr_warmup_steps + GBS)
+
+This patch replaces the warmup branch of get_lr with the NeMo formula.
+Enabled by setting nemo_aligned_lr_warmup: true in the YAML config.
+"""
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _nemo_lr_enabled(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return args is not None and getattr(args, "nemo_aligned_lr_warmup", False)
+
+
+@register_patch(
+    "megatron.lr_schedule.nemo_aligned",
+    backend="megatron",
+    phase="before_train",
+    description="Align LR warmup with NeMo's (step+1)/(warmup_steps+1) formula",
+    condition=_nemo_lr_enabled,
+    priority=50,
+)
+def patch_nemo_aligned_lr_warmup(ctx: PatchContext):
+    from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
+
+    _original_get_lr = OptimizerParamScheduler.get_lr
+    _gbs_cache = [None]
+
+    def _nemo_get_lr(self, param_group):
+        if self.lr_warmup_steps > 0 and self.num_steps <= self.lr_warmup_steps:
+            if _gbs_cache[0] is None:
+                from megatron.training import get_args as megatron_get_args
+
+                _gbs_cache[0] = megatron_get_args().global_batch_size
+            gbs = _gbs_cache[0]
+            max_lr = param_group.get("max_lr", self.max_lr)
+            return self.init_lr + (
+                (max_lr - self.init_lr) * float(self.num_steps + gbs) / float(self.lr_warmup_steps + gbs)
+            )
+        return _original_get_lr(self, param_group)
+
+    OptimizerParamScheduler.get_lr = _nemo_get_lr
+    log_rank_0(
+        "[Patch:nemo_aligned_lr] Patched get_lr warmup: "
+        "(num_steps+GBS)/(lr_warmup_steps+GBS) = NeMo's (step+1)/(warmup+1)"
+    )
diff --git a/primus/backends/megatron/patches/mamba_fla_data_patches.py b/primus/backends/megatron/patches/mamba_fla_data_patches.py
new file mode 100644
index 000000000..1f4aa7adc
--- /dev/null
+++ b/primus/backends/megatron/patches/mamba_fla_data_patches.py
@@ -0,0 +1,108 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Mamba FLA-Order Dataset Patch
+===============================
+
+Wraps ``pretrain_mamba.train_valid_test_datasets_provider`` so that, when
+enabled, training reads tokens through ``tools/hybrid/fla_order_dataset.py``
+(``FLAOrderGPTDataset``) instead of Megatron's ``GPTDataset``. This replays
+the exact same token order flash-linear-attention's (FLA) reference trainer
+used, which is required for bit-for-bit loss-curve parity checks against FLA.
+
+Toggle: ``args.use_fla_data`` + ``args.fla_cache_dir`` (resolved by
+``fla_runtime_patches.py`` from ``PRIMUS_FLA_DATA`` / ``PRIMUS_FLA_CACHE_DIR``
+or YAML ``use_fla_data`` / ``fla_cache_dir``).
+
+``primus.backends.megatron.megatron_pretrain_trainer`` imports
+``train_valid_test_datasets_provider`` by name from ``pretrain_mamba`` at
+``train()`` time (well after ``before_train`` patches run), so replacing the
+attribute on the ``pretrain_mamba`` module here is sufficient -- no need to
+touch call sites.
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.mamba.fla_order_dataset"
+
+
+def _install_mamba_fla_data_patch() -> None:
+    import pretrain_mamba
+
+    if is_patched(pretrain_mamba, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] pretrain_mamba already patched; skipping.")
+        return
+
+    original_provider = pretrain_mamba.train_valid_test_datasets_provider
+
+    def patched_provider(train_val_test_num_samples, vp_stage=None):
+        from megatron.training import get_args as _get_args
+        from megatron.training import print_rank_0
+
+        args = _get_args()
+        fla_data_flag = getattr(args, "use_fla_data", False)
+        fla_cache = getattr(args, "fla_cache_dir", "")
+        print_rank_0(f"> [FLA-check] use_fla_data={fla_data_flag!r}, fla_cache_dir={fla_cache!r}")
+        if fla_data_flag and fla_cache:
+            import importlib.util
+            import os
+
+            from megatron.core import parallel_state
+            from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
+
+            _spec = importlib.util.spec_from_file_location(
+                "fla_order_dataset",
+                os.path.join(
+                    os.environ.get("PRIMUS_PATH", os.getcwd()), "tools", "hybrid", "fla_order_dataset.py"
+                ),
+            )
+            _mod = importlib.util.module_from_spec(_spec)
+            _spec.loader.exec_module(_mod)
+            FLAOrderGPTDataset = _mod.FLAOrderGPTDataset
+
+            dp_size = parallel_state.get_data_parallel_world_size()
+            tokenizer = build_tokenizer(args)
+            print_rank_0(f"> building FLA-order dataset from {fla_cache} ...")
+            train_ds = FLAOrderGPTDataset(
+                cache_dir=fla_cache,
+                seq_length=args.seq_length,
+                micro_batch_size=args.micro_batch_size,
+                data_parallel_size=dp_size,
+                seed=args.seed,
+                pad_token_id=0,
+                eod_token=tokenizer.eod,
+                eod_mask_loss=args.eod_mask_loss,
+            )
+            print_rank_0(f"> FLA-order dataset: {len(train_ds)} samples")
+            return train_ds, None, None
+        return original_provider(train_val_test_num_samples, vp_stage=vp_stage)
+
+    pretrain_mamba.train_valid_test_datasets_provider = patched_provider
+    mark_patched(pretrain_mamba, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] Patched pretrain_mamba.train_valid_test_datasets_provider "
+        "to serve FLAOrderGPTDataset when use_fla_data is set."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Serve tokens through FLAOrderGPTDataset (tools/hybrid/fla_order_dataset.py) "
+        "instead of Megatron's GPTDataset, to replay FLA's exact reference token order."
+    ),
+    # Runs after fla_runtime_knobs (priority=-100) has resolved args.use_fla_data / fla_cache_dir.
+    priority=50,
+    condition=lambda ctx: bool(getattr(get_args(ctx), "use_fla_data", False))
+    and bool(getattr(get_args(ctx), "fla_cache_dir", "")),
+)
+def patch_mamba_fla_data(ctx: PatchContext) -> None:
+    _install_mamba_fla_data_patch()
diff --git a/primus/backends/megatron/patches/mamba_fused_ce_patches.py b/primus/backends/megatron/patches/mamba_fused_ce_patches.py
new file mode 100644
index 000000000..b23842f57
--- /dev/null
+++ b/primus/backends/megatron/patches/mamba_fused_ce_patches.py
@@ -0,0 +1,137 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Mamba/Hybrid FLA Fused Cross-Entropy Patch
+===========================================
+
+Wires flash-linear-attention's (FLA) fused cross-entropy kernels into
+``MambaModel`` so GDN/KDA/Mamba2 (pure or hybrid) training never materializes
+the full ``(batch*seq, vocab)`` logits tensor -- avoiding OOM on large
+vocab/long sequence runs and matching FLA's own training loss numerically.
+
+Two modes, selected by ``args.fused_ce_mode`` (resolved by
+``fla_runtime_patches.py`` from ``PRIMUS_FUSED_CE`` / YAML ``fused_ce_mode``):
+
+    0: disabled -- stock Megatron path (materializes full logits).
+    1 (default): ``FusedLinearCrossEntropyLoss`` -- computes logits + CE in
+       chunks (``args.fused_ce_chunks``, default 32); the full logits tensor
+       is never resident. Biggest memory win.
+    2: ``FusedCrossEntropyLoss`` -- materializes logits in bf16, then a fused
+       Triton CE kernel; matches FLA's exact computation.
+
+This is a "source-string rewrite" style patch: the two injection points
+(``MambaModel.__init__`` / ``MambaModel.forward``) sit in the middle of
+upstream method bodies, where a plain function-wrapping monkey-patch cannot
+reach without duplicating the entire (large, version-sensitive) method.
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.backends.megatron.patches._source_patch_utils import patch_method_source
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.mamba.fla_fused_ce"
+
+# Inserted at the end of MambaModel.__init__, right after the embeddings/
+# output-layer setup and before the `finish_init` loop -- mirrors exactly
+# where the original megatron_patches/01-mamba_model-fused-ce.patch spliced
+# in its setup code.
+_INIT_ORI = "self.setup_embeddings_and_output_layer()"
+_INIT_NEW = (
+    _INIT_ORI + "\n\n" + "        from megatron.training import get_args as _get_args\n"
+    "        _args = _get_args()\n"
+    "        self._fused_ce_mode = 0\n"
+    "        _ce_mode = getattr(_args, 'fused_ce_mode', 1)\n"
+    "        if _ce_mode == 2:\n"
+    "            try:\n"
+    "                from fla.modules import FusedCrossEntropyLoss\n"
+    "                self._fused_ce = FusedCrossEntropyLoss(inplace_backward=True)\n"
+    "                self._fused_ce_mode = 2\n"
+    "            except ImportError:\n"
+    "                pass\n"
+    "        elif _ce_mode == 1:\n"
+    "            try:\n"
+    "                from fla.modules import FusedLinearCrossEntropyLoss\n"
+    "                _nc = getattr(_args, 'fused_ce_chunks', 32)\n"
+    "                self._fused_lce = FusedLinearCrossEntropyLoss(reduction='mean', num_chunks=_nc)\n"
+    "                self._fused_ce_mode = 1\n"
+    "            except ImportError:\n"
+    "                pass\n"
+    "        self._use_fused_cross_entropy = self._fused_ce_mode > 0"
+)
+
+# Inserted right before the standard output-layer + loss computation, so the
+# fused path can skip building the full [s, b, vocab] logits tensor entirely.
+_FORWARD_ORI = "logits, _ = self.output_layer("
+_FORWARD_NEW = (
+    "if labels is not None and self._use_fused_cross_entropy:\n"
+    "            return self._fused_cross_entropy_loss(hidden_states, labels, output_weight)\n"
+    "\n"
+    "        " + _FORWARD_ORI
+)
+
+
+def _fused_cross_entropy_loss(self, hidden_states, labels, output_weight, runtime_gather_output=None):
+    """Compute loss using FLA's fused cross-entropy kernels.
+
+    Mode 1 (FusedLinearCrossEntropyLoss): computes logits + CE in chunks,
+           never materializing the full (batch*seq, vocab) logits tensor.
+    Mode 2 (FusedCrossEntropyLoss): materializes logits in bf16, then
+           uses a fused Triton CE kernel -- matches FLA's exact computation.
+    """
+    s, b, h = hidden_states.shape
+
+    if self._fused_ce_mode == 2:
+        logits, _ = self.output_layer(
+            hidden_states,
+            weight=output_weight,
+            runtime_gather_output=runtime_gather_output,
+        )
+        logits_2d = logits.permute(1, 0, 2).reshape(b * s, -1)
+        labels_1d = labels.reshape(b * s)
+        loss = self._fused_ce(logits_2d, labels_1d)
+        return loss.expand(b, s)
+    else:
+        hs_2d = hidden_states.permute(1, 0, 2).reshape(b * s, h)
+        labels_1d = labels.reshape(b * s)
+        weight = output_weight if output_weight is not None else self.output_layer.weight
+        loss = self._fused_lce(hs_2d, labels_1d, weight)
+        return loss.expand(b, s)
+
+
+def _install_mamba_fused_ce_patch() -> None:
+    from megatron.core.models.mamba.mamba_model import MambaModel
+
+    if is_patched(MambaModel, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] MambaModel already patched; skipping.")
+        return
+
+    patch_method_source(MambaModel, "__init__", _INIT_ORI, _INIT_NEW)
+    patch_method_source(MambaModel, "forward", _FORWARD_ORI, _FORWARD_NEW)
+    MambaModel._fused_cross_entropy_loss = _fused_cross_entropy_loss
+
+    mark_patched(MambaModel, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] Patched MambaModel.__init__/forward to route through "
+        "FLA's fused cross-entropy kernels when fused_ce_mode != 0."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Wire FLA's FusedLinearCrossEntropyLoss / FusedCrossEntropyLoss into MambaModel "
+        "so the full (batch*seq, vocab) logits tensor is never materialized."
+    ),
+    # Runs after fla_runtime_knobs (priority=-100) has resolved args.fused_ce_mode.
+    priority=50,
+    condition=lambda ctx: getattr(get_args(ctx), "fused_ce_mode", 0) != 0,
+)
+def patch_mamba_fused_ce(ctx: PatchContext) -> None:
+    _install_mamba_fused_ce_patch()
diff --git a/primus/backends/megatron/patches/mamba_rocm_patches.py b/primus/backends/megatron/patches/mamba_rocm_patches.py
index 9d6243247..31caa2c94 100644
--- a/primus/backends/megatron/patches/mamba_rocm_patches.py
+++ b/primus/backends/megatron/patches/mamba_rocm_patches.py
@@ -15,7 +15,7 @@
 import torch
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_rocm(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/megatron_fsdp_patches.py b/primus/backends/megatron/patches/megatron_fsdp_patches.py
index 5a3b48ea8..a6fb3b921 100644
--- a/primus/backends/megatron/patches/megatron_fsdp_patches.py
+++ b/primus/backends/megatron/patches/megatron_fsdp_patches.py
@@ -15,7 +15,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/mla_patches.py b/primus/backends/megatron/patches/mla_patches.py
index c138112f4..d645d2871 100644
--- a/primus/backends/megatron/patches/mla_patches.py
+++ b/primus/backends/megatron/patches/mla_patches.py
@@ -12,7 +12,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
@@ -20,9 +20,26 @@
     backend="megatron",
     phase="before_train",
     description=(
-        "Monkey patch MLA attention to use PrimusMLASelfAttention " "when use_turbo_gemm is enabled."
+        "Monkey patch MLA attention to use PrimusMLASelfAttention "
+        "when use_turbo_gemm is enabled (skipped for DeepSeek-V4)."
+    ),
+    # Skip for DeepSeek-V4: its DeepseekV4Attention subclasses MLASelfAttention,
+    # and PrimusMLASelfAttention deliberately bypasses MLASelfAttention.__init__
+    # (calls the grandparent), which would break V4's super().__init__ chain if
+    # the base class were swapped. V4 builds DeepseekV4Attention directly and does
+    # not need the padded-fusion MLA path, so this patch must not fire for it.
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_turbo_gemm", False)
+        and not any(
+            getattr(get_args(ctx), _f, False)
+            for _f in (
+                "use_v4_triton_attention",
+                "use_v4_triton_csa_attention",
+                "use_v4_tilelang_attention",
+                "use_v4_tilelang_csa_attention",
+            )
+        )
     ),
-    condition=lambda ctx: getattr(get_args(ctx), "use_turbo_gemm", False),
 )
 def patch_mla_attention(ctx: PatchContext):
     """
diff --git a/primus/backends/megatron/patches/mlp_fla_swiglu_patches.py b/primus/backends/megatron/patches/mlp_fla_swiglu_patches.py
new file mode 100644
index 000000000..85d5585bc
--- /dev/null
+++ b/primus/backends/megatron/patches/mlp_fla_swiglu_patches.py
@@ -0,0 +1,94 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+MLP FLA SwiGLU Patch
+======================
+
+Routes ``MLP``'s gated-linear-unit activation through
+flash-linear-attention's (FLA) Triton-fused ``swiglu`` kernel (one fwd + one
+bwd kernel) instead of Megatron's naive two-kernel ``silu(x_glu) * x_linear``,
+saving ~20 ms/iter on GDN/KDA/Mamba hybrid training. Only takes effect when
+``config.gated_linear_unit`` and the activation is ``F.silu`` -- i.e. it is a
+no-op for GeLU-based MLPs, MoE, and any config where
+``use_te_activation_func`` / ``bias_activation_fusion`` are already handling
+the activation through a different (already-fused) code path.
+
+Toggle: ``args.use_fla_fused_swiglu`` (resolved by ``fla_runtime_patches.py``
+from ``PRIMUS_FLA_SWIGLU`` / YAML ``use_fla_fused_swiglu``, default True).
+
+Source-string rewrite style: the injection points sit inside ``MLP.__init__``
+and the non-fused branch of ``MLP.forward``, which a plain wrapper cannot
+reach without duplicating the whole (large) method body.
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.backends.megatron.patches._source_patch_utils import patch_method_source
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.mlp.fla_swiglu"
+
+# Inserted right before linear_fc2 construction in MLP.__init__, mirroring
+# where megatron_patches/03-mlp-fla-swiglu.patch spliced in its detection.
+_INIT_ORI = "self.linear_fc2 = submodules.linear_fc2("
+_INIT_NEW = (
+    "from megatron.training import get_args as _get_args\n"
+    "        self._use_fla_swiglu = False\n"
+    "        if getattr(_get_args(), 'use_fla_fused_swiglu', True):\n"
+    "            if self.config.gated_linear_unit and self.config.activation_func == F.silu:\n"
+    "                try:\n"
+    "                    from fla.modules.activations import swiglu as _fla_swiglu\n"
+    "                    self._use_fla_swiglu = True\n"
+    "                    self._fla_swiglu_fn = _fla_swiglu\n"
+    "                except ImportError:\n"
+    "                    pass\n"
+    "\n        " + _INIT_ORI
+)
+
+# Inserted in MLP.forward's non-fused-bias-activation branch, in place of the
+# unconditional `glu(intermediate_parallel)` call.
+_FORWARD_ORI = "intermediate_parallel = glu(intermediate_parallel)"
+_FORWARD_NEW = (
+    "if self._use_fla_swiglu:\n"
+    "                    x_glu, x_linear = torch.chunk(intermediate_parallel, 2, dim=-1)\n"
+    "                    intermediate_parallel = self._fla_swiglu_fn(x_glu, x_linear)\n"
+    "                else:\n"
+    "                    " + _FORWARD_ORI
+)
+
+
+def _install_mlp_fla_swiglu_patch() -> None:
+    from megatron.core.transformer.mlp import MLP
+
+    if is_patched(MLP, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] MLP already patched; skipping.")
+        return
+
+    patch_method_source(MLP, "__init__", _INIT_ORI, _INIT_NEW)
+    patch_method_source(MLP, "forward", _FORWARD_ORI, _FORWARD_NEW)
+
+    mark_patched(MLP, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] Patched MLP.__init__/forward to use FLA's Triton "
+        "swiglu kernel when use_fla_fused_swiglu is set."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Route MLP's gated-linear-unit activation through FLA's Triton-fused "
+        "swiglu kernel instead of Megatron's naive silu*x implementation."
+    ),
+    # Runs after fla_runtime_knobs (priority=-100) has resolved args.use_fla_fused_swiglu.
+    priority=50,
+    condition=lambda ctx: getattr(get_args(ctx), "use_fla_fused_swiglu", False),
+)
+def patch_mlp_fla_swiglu(ctx: PatchContext) -> None:
+    _install_mlp_fla_swiglu_patch()
diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py
new file mode 100644
index 000000000..b53173d8a
--- /dev/null
+++ b/primus/backends/megatron/patches/mlperf_logging_patches.py
@@ -0,0 +1,485 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+MLPerf Logging Patches for Flux Training.
+
+Installs MLPerf-compliant logging into Megatron's training loop by wrapping:
+  - training_log: emit INIT_STOP, RUN_START, tracked_stats, train_loss
+  - evaluate_and_print_results: emit EVAL events, convergence check
+  - print_rank_last / get_tensorboard_writer / get_wandb_writer: suppress
+
+Uses mlperf_logging.mllog library for structured event output.
+"""
+
+import logging
+import os
+import sys
+import time
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+logger = logging.getLogger(__name__)
+
+
+def _mlperf_logging_enabled(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return args is not None and getattr(args, "mlperf_mode", False)
+
+
+class ThroughputTimer:
+    """Wall-clock throughput tracker with eval pause/resume."""
+
+    def __init__(self, gbs: int):
+        self.gbs = gbs
+        self.training_start_time: float | None = None
+        self.eval_cumulative_secs: float = 0.0
+        self._eval_enter_time: float | None = None
+        self.consumed_samples: int = 0
+
+    def mark_training_start(self):
+        if self.training_start_time is None:
+            self.training_start_time = time.time()
+
+    def update_samples(self, iteration: int):
+        self.consumed_samples = iteration * self.gbs
+
+    def pause_for_eval(self):
+        self._eval_enter_time = time.time()
+
+    def resume_after_eval(self):
+        if self._eval_enter_time is not None:
+            self.eval_cumulative_secs += time.time() - self._eval_enter_time
+            self._eval_enter_time = None
+
+    def compute_throughput(self):
+        if self.training_start_time is None:
+            return 0.0
+        wall = time.time() - self.training_start_time
+        training_secs = wall - self.eval_cumulative_secs
+        if training_secs <= 0:
+            return 0.0
+        return self.consumed_samples / training_secs
+
+    def compute_combined_throughput(self):
+        if self.training_start_time is None:
+            return 0.0
+        wall = time.time() - self.training_start_time
+        if wall <= 0:
+            return 0.0
+        return self.consumed_samples / wall
+
+
+class FluxMLPerfLogger:
+    """MLPerf logger using mlperf_logging.mllog directly."""
+
+    def __init__(
+        self,
+        global_batch_size: int,
+        micro_batch_size: int,
+        target_val_loss: float = 0.586,
+        log_every_n_steps: int = 10,
+    ):
+        from mlperf_logging import mllog
+
+        self._mllogger = mllog.get_mllogger()
+        self._constants = mllog.constants
+        self.gbs = global_batch_size
+        self.mbs = micro_batch_size
+        self.target_val_loss = target_val_loss
+        self.log_every_n_steps = log_every_n_steps
+        self.timer = ThroughputTimer(global_batch_size)
+        self._converged = False
+
+        self.profiler = os.getenv("PROFILER", "")
+        self.profiler_warmup_steps = int(os.getenv("PROF_WARMUP_STEPS", "0"))
+        self.profiler_active_steps = int(os.getenv("PROF_ACTIVE_STEPS", "0"))
+        self.rpd = None
+        self.rpd_running = False
+
+        if self.profiler == "rpd":
+            try:
+                from rpdTracerControl import rpdTracerControl
+
+                rpdTracerControl.setFilename("trace.rpd", append=True)
+                self.rpd = rpdTracerControl()
+                logger.info("RPD profiler initialized")
+            except ImportError:
+                logger.warning("rpdTracerControl not available")
+
+    def _event(self, key, value=None, metadata=None):
+        self._mllogger.event(key=key, value=value, metadata=metadata)
+
+    def _start(self, key, value=None, metadata=None):
+        self._mllogger.start(key=key, value=value, metadata=metadata)
+
+    def _end(self, key, value=None, metadata=None):
+        self._mllogger.end(key=key, value=value, metadata=metadata)
+
+    def log_init(self, seed: int):
+        if int(os.environ.get("RANK", "0")) == 0:
+            self._start(key=self._constants.INIT_START)
+            self._event(key=self._constants.SUBMISSION_BENCHMARK, value="flux1")
+            self._event(
+                key=self._constants.SUBMISSION_ORG,
+                value=os.environ.get("MLLOG_SUBMISSION_ORG", "AMD"),
+            )
+            self._event(
+                key=self._constants.SUBMISSION_DIVISION,
+                value=os.environ.get("MLLOG_SUBMISSION_DIVISION", "closed"),
+            )
+            self._event(
+                key=self._constants.SUBMISSION_PLATFORM,
+                value=os.environ.get("MLLOG_SUBMISSION_PLATFORM", "MI355X"),
+            )
+            self._event(key=self._constants.SUBMISSION_STATUS, value="onprem")
+            self._event(key="target_accuracy", value=self.target_val_loss)
+            self._event(key=self._constants.SEED, value=seed)
+
+    def log_hyperparams(self, args):
+        if int(os.environ.get("RANK", "0")) != 0:
+            return
+        self._event(key=self._constants.GLOBAL_BATCH_SIZE, value=self.gbs)
+        self._event(
+            key=self._constants.TRAIN_SAMPLES,
+            value=getattr(args, "train_samples", 1099776),
+        )
+        self._event(
+            key=self._constants.EVAL_SAMPLES,
+            value=getattr(args, "eval_samples", 29696),
+        )
+        gas = max(self.gbs // self.mbs, 1)
+        self._event(key=self._constants.GRADIENT_ACCUMULATION_STEPS, value=gas)
+        self._event(key=self._constants.OPT_NAME, value="adamw")
+        self._event(
+            key=self._constants.OPT_BASE_LR,
+            value=getattr(args, "lr", 2e-4),
+        )
+        self._event(
+            key="opt_adamw_beta_1",
+            value=getattr(args, "adam_beta1", 0.9),
+        )
+        self._event(
+            key="opt_adamw_beta_2",
+            value=getattr(args, "adam_beta2", 0.95),
+        )
+        self._event(
+            key="opt_adamw_epsilon",
+            value=getattr(args, "adam_eps", 1e-8),
+        )
+        self._event(
+            key="opt_adamw_weight_decay",
+            value=getattr(args, "weight_decay", 0.1),
+        )
+
+    def log_init_stop_run_start(self):
+        if int(os.environ.get("RANK", "0")) == 0:
+            self._end(key=self._constants.INIT_STOP)
+            self._start(key=self._constants.RUN_START)
+            self._start(key=self._constants.EPOCH_START, metadata={"epoch_num": 0})
+            self._start(key=self._constants.BLOCK_START, metadata={"first_epoch_num": 0})
+
+    def on_train_batch_end(self, global_step: int, loss: float, lr: float):
+        self.timer.mark_training_start()
+        self.timer.update_samples(global_step)
+
+        self._handle_profiler(global_step)
+
+        if int(os.environ.get("RANK", "0")) != 0:
+            return
+        if global_step % self.log_every_n_steps == 0:
+            self._event(
+                key="tracked_stats",
+                value={"train_loss": loss},
+                metadata={
+                    "samples_count": global_step * self.gbs,
+                    "lr": lr,
+                    "step": global_step,
+                },
+            )
+
+    def on_validation_start(self, global_step: int):
+        self.timer.update_samples(global_step)
+        self.timer.pause_for_eval()
+
+        if int(os.environ.get("RANK", "0")) == 0:
+            if global_step > 0:
+                throughput = self.timer.compute_throughput()
+                self._event(
+                    key="throughput",
+                    value=throughput,
+                    metadata={
+                        "samples_count": global_step * self.gbs,
+                        "step": global_step,
+                    },
+                )
+            self._end(
+                key=self._constants.BLOCK_STOP,
+                metadata={"first_epoch_num": 0},
+            )
+            self._start(key=self._constants.EVAL_START, metadata={"epoch_num": 0})
+
+    def on_validation_end(self, global_step: int, val_loss: float):
+        self.timer.resume_after_eval()
+
+        if int(os.environ.get("RANK", "0")) == 0:
+            self._event(
+                key=self._constants.EVAL_ACCURACY,
+                value=val_loss,
+                metadata={
+                    "samples_count": global_step * self.gbs,
+                    "step": global_step,
+                },
+            )
+            self._end(key=self._constants.EVAL_STOP, metadata={"epoch_num": 0})
+            combined_throughput = self.timer.compute_combined_throughput()
+            self._event(
+                key="combined_throughput",
+                value=combined_throughput,
+                metadata={
+                    "samples_count": global_step * self.gbs,
+                    "step": global_step,
+                },
+            )
+
+    def _handle_profiler(self, global_step: int):
+        if self.profiler != "rpd":
+            return
+        if self.rpd and not self.rpd_running and global_step >= self.profiler_warmup_steps:
+            logger.info("Starting RPD profiler")
+            self.rpd.start()
+            self.rpd.rangePush("python", "Training", "")
+            self.rpd_running = True
+        if self.rpd_running and global_step > self.profiler_warmup_steps + self.profiler_active_steps:
+            logger.info("Stopping RPD profiler")
+            self.rpd.rangePop()
+            self.rpd.stop()
+            self.rpd = None
+            self.rpd_running = False
+
+    @property
+    def converged(self):
+        return self._converged
+
+    def log_run_stop(self, success: bool, global_step: int):
+        if success:
+            self._converged = True
+        if int(os.environ.get("RANK", "0")) == 0:
+            status = "success" if success else "aborted"
+            self._end(
+                key=self._constants.RUN_STOP,
+                value=status,
+                metadata={
+                    "samples_count": global_step * self.gbs,
+                    "step": global_step,
+                    "status": status,
+                },
+            )
+
+    def teardown(self):
+        if self.rpd_running and self.rpd:
+            self.rpd.rangePop()
+            self.rpd.stop()
+            self.rpd = None
+            self.rpd_running = False
+
+
+def _extract_val_loss(loss_dict):
+    """Extract scalar validation loss from captured total_loss_dict."""
+    if not loss_dict or not isinstance(loss_dict, dict):
+        return None
+    for key in ("loss", "lm loss"):
+        if key in loss_dict:
+            val = loss_dict[key]
+            return val.item() if hasattr(val, "item") else float(val)
+    if loss_dict:
+        val = next(iter(loss_dict.values()))
+        return val.item() if hasattr(val, "item") else float(val)
+    return None
+
+
+@register_patch(
+    "megatron.training.mlperf_logging",
+    backend="megatron",
+    phase="before_train",
+    description="Install MLPerf logging wrappers for Flux training",
+    condition=_mlperf_logging_enabled,
+    priority=15,
+)
+def patch_mlperf_logging(ctx: PatchContext):
+    """Install MLPerf logging: suppress Megatron output, wrap training_log and eval."""
+    import megatron.training.training as megatron_training
+
+    if getattr(megatron_training, "_primus_mlperf_logging_installed", False):
+        return
+
+    args = get_args(ctx)
+    seed = getattr(args, "seed", 42)
+    gbs = getattr(args, "global_batch_size", 512)
+    mbs = getattr(args, "micro_batch_size", 64)
+    target_val_loss = getattr(args, "target_val_loss", 0.586)
+    log_interval = getattr(args, "log_interval", 10)
+
+    mlperf_logger = FluxMLPerfLogger(
+        global_batch_size=gbs,
+        micro_batch_size=mbs,
+        target_val_loss=target_val_loss,
+        log_every_n_steps=log_interval,
+    )
+
+    mlperf_logger.log_init(seed=seed)
+    mlperf_logger.log_hyperparams(args)
+
+    # --- Suppress Megatron's built-in logging ---
+    megatron_training.print_rank_last = lambda *a, **k: None
+
+    for writer_fn in ("get_tensorboard_writer", "get_wandb_writer"):
+        if hasattr(megatron_training, writer_fn):
+            setattr(megatron_training, writer_fn, lambda: None)
+
+    # --- Wrap training_log ---
+    _orig_training_log = megatron_training.training_log
+    _first_training_log_call = [True]
+
+    def _mlperf_training_log(*args_tl, **kwargs_tl):
+        if _first_training_log_call[0]:
+            _first_training_log_call[0] = False
+            mlperf_logger.log_init_stop_run_start()
+
+        result = _orig_training_log(*args_tl, **kwargs_tl)
+
+        try:
+            loss_dict = args_tl[0] if len(args_tl) > 0 else kwargs_tl.get("loss_dict", {})
+            learning_rate = args_tl[2] if len(args_tl) > 2 else kwargs_tl.get("learning_rate", 0.0)
+            # Upstream Megatron: training_log(loss_dict, total_loss_dict, learning_rate, iteration, ...)
+            iteration = args_tl[3] if len(args_tl) > 3 else kwargs_tl.get("iteration", 0)
+
+            if loss_dict:
+                loss_val = next(iter(loss_dict.values()))
+                if hasattr(loss_val, "item"):
+                    loss_val = loss_val.item()
+                mlperf_logger.on_train_batch_end(iteration, loss_val, learning_rate)
+
+                if iteration % log_interval == 0:
+                    lr_str = f"{learning_rate:.2e}" if learning_rate else "N/A"
+                    sys.stdout.write(
+                        f"step {iteration} | loss: {loss_val:.4f} | lr: {lr_str}"
+                        f" | samples: {iteration * gbs}\n"
+                    )
+                    sys.stdout.flush()
+        except Exception as e:
+            logger.debug("MLPerf training_log hook: %s", e)
+
+        return result
+
+    _mlperf_training_log._primus_mlperf_logging_wrapper = True
+    megatron_training.training_log = _mlperf_training_log
+
+    # --- Wrap evaluate_and_print_results ---
+    _orig_eval = megatron_training.evaluate_and_print_results
+
+    def _mlperf_evaluate_and_print_results(*eval_args, **eval_kwargs):
+        # evaluate_and_print_results(prefix, fwd, data, model, iteration[4], ...)
+        iteration = eval_kwargs.get("iteration", eval_args[4] if len(eval_args) > 4 else 0)
+
+        mlperf_logger.on_validation_start(iteration)
+
+        # Temporarily wrap whatever `evaluate` is at call time (e.g.
+        # primus_evaluate installed by the evaluate patch) so we can
+        # capture the total_loss_dict it returns.  This avoids relying
+        # on a persistent hook that later patches can overwrite.
+        _loss_capture = {}
+        _current_eval = megatron_training.evaluate
+
+        def _capture_wrapper(*a, **kw):
+            res = _current_eval(*a, **kw)
+            td = res[0] if isinstance(res, tuple) else res
+            if isinstance(td, dict):
+                _loss_capture.update(td)
+            return res
+
+        megatron_training.evaluate = _capture_wrapper
+        try:
+            import gc
+
+            result = _orig_eval(*eval_args, **eval_kwargs)
+            gc.collect()
+        finally:
+            megatron_training.evaluate = _current_eval
+
+        try:
+            import torch
+
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+        val_loss = _extract_val_loss(_loss_capture)
+        if val_loss is not None:
+            # Megatron's `evaluate()` (training.py:3178-3180) divides the
+            # per-rank accumulated loss locally and does NOT all-reduce
+            # across the data-parallel group — the result is intended for
+            # `print_rank_last` / TensorBoard which only read on a single
+            # rank.  We must reduce here so every rank evaluates the same
+            # global validation loss against `target_val_loss`; otherwise
+            # ranks can disagree on the early-stop branch and the
+            # divergent `args.train_iters` mutation below desyncs
+            # collective ordering at the next training step, producing a
+            # NCCL watchdog deadlock (observed on FLUX 12B MLPerf at
+            # step 2560 when val_loss landed near target).
+            import torch
+            import torch.distributed as dist
+
+            if dist.is_initialized():
+                try:
+                    from megatron.core import parallel_state as mpu
+
+                    dp_group = mpu.get_data_parallel_group()
+                except Exception:
+                    dp_group = None
+                _vl = torch.tensor(val_loss, dtype=torch.float64, device="cuda")
+                dist.all_reduce(_vl, op=dist.ReduceOp.AVG, group=dp_group)
+                val_loss = _vl.item()
+
+            mlperf_logger.on_validation_end(iteration, val_loss)
+            log_rank_0(
+                f"[MLPerf] Validation loss at step {iteration}: {val_loss:.6f} "
+                f"(target: {target_val_loss:.6f})"
+            )
+
+            if val_loss <= target_val_loss:
+                log_rank_0(
+                    f"[MLPerf] Convergence reached! val_loss={val_loss:.6f} "
+                    f"<= target={target_val_loss:.6f}"
+                )
+                mlperf_logger.log_run_stop(success=True, global_step=iteration)
+                try:
+                    from megatron.training import get_args as megatron_get_args
+
+                    megatron_get_args().train_iters = iteration
+                except Exception:
+                    logger.warning("Could not set args.train_iters for early stop")
+            else:
+                if int(os.environ.get("RANK", "0")) == 0:
+                    mlperf_logger._start(
+                        key=mlperf_logger._constants.BLOCK_START,
+                        metadata={"first_epoch_num": 0},
+                    )
+        else:
+            logger.warning("Could not extract validation loss from evaluate result")
+
+        return result
+
+    _mlperf_evaluate_and_print_results._primus_mlperf_eval_wrapper = True
+    megatron_training.evaluate_and_print_results = _mlperf_evaluate_and_print_results
+
+    megatron_training._primus_mlperf_logging_installed = True
+
+    log_rank_0(
+        f"[Patch:mlperf_logging] Installed MLPerf logging (gbs={gbs}, "
+        f"target_val_loss={target_val_loss}, log_interval={log_interval})"
+    )
diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py
new file mode 100644
index 000000000..5ac532978
--- /dev/null
+++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py
@@ -0,0 +1,471 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+MLPerf Warmup Patches for Flux Training.
+
+Wraps megatron.training.training.train_step with a one-shot hook that runs
+all warmup steps in a tight loop on the first invocation, then executes the
+first real training step and self-removes.  This avoids relying on Megatron's
+local ``iteration`` variable (which cannot be controlled from a train_step
+wrapper) and mirrors NeMo's approach of running warmup before real data is
+touched.
+
+Priority 95 ensures this hook is the outermost wrapper around the full
+train_step chain (FP8 cache, delayed scaling, wall-clock timer, etc.).
+Self-removal restores the inner chain intact.
+"""
+
+import logging
+
+import torch
+import torch.distributed
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+logger = logging.getLogger(__name__)
+
+
+def _log(msg):
+    log_rank_0(f"[MLPerf_WARMUP] {msg}")
+
+
+def _warmup_enabled(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return args is not None and getattr(args, "warmup_train_steps", 0) > 0
+
+
+def _reset_fp8_te_spec(models):
+    """Reset FP8 state for TransformerEngine spec modules.
+
+    Recipe-agnostic: walks ``fp8_meta`` and only touches buffers that exist on
+    the recipe-specific state object. Skips TE's ``reset_fp8_meta_tensors``
+    helper because TE 2.8.0.dev0 unconditionally derefs ``.scale`` /
+    ``.amax_history``, which crashes on ``Float8CurrentScalingRecipeState``
+    (current/tensorwise scaling has no persistent state — see TE's
+    ``fp8.py``: *"Per-tensor current quantization does not require state"*).
+    """
+    count = 0
+    for m in models:
+        for module in m.modules():
+            if not hasattr(module, "fp8_initialized"):
+                continue
+            module.fp8_initialized = False
+            count += 1
+            if not hasattr(module, "fp8_meta"):
+                continue
+            meta = module.fp8_meta
+            for key in ("scaling_fwd", "scaling_bwd"):
+                if key not in meta:
+                    continue
+                tm = meta[key]
+                if hasattr(tm, "amax_history"):
+                    tm.amax_history.fill_(0.0)
+                if hasattr(tm, "scale"):
+                    tm.scale.fill_(1.0)
+                if hasattr(tm, "scale_inv"):
+                    tm.scale_inv.fill_(1.0)
+    return count
+
+
+def _seed_fp8_amax(models, seed_value=1.0):
+    """Seed FP8 amax_history with a safe non-zero value to prevent scale=inf."""
+    count = 0
+    for m in models:
+        for module in m.modules():
+            if not hasattr(module, "fp8_meta"):
+                continue
+            meta = module.fp8_meta
+            for key in ("scaling_fwd", "scaling_bwd"):
+                if key not in meta:
+                    continue
+                tm = meta[key]
+                if hasattr(tm, "amax_history"):
+                    tm.amax_history.fill_(seed_value)
+                    count += 1
+    return count
+
+
+def _reset_fp8_local_spec(models):
+    """Reset FP8 state for local spec (tensorwise delayed-scaling) modules.
+
+    Re-initialises per-module delayed-scaling buffers via the canonical
+    ``_init_delayed_scaling_state`` helper.  The new buffers are separate
+    objects from the ``_DelayedScalingRegistry``'s global tensors, so the
+    pointer-check in ``_fast_update_scales`` / ``_fast_update_scales_with_history``
+    will detect the mismatch and trigger ``registry.__init__(modules)`` on the
+    next training step, which re-creates global tensors with
+    ``_first_step = True`` and bootstraps weight amaxes from the restored weights.
+
+    Buffers are moved to the module's device because ``_init_delayed_scaling_state``
+    creates plain CPU tensors (bare ``torch.zeros`` / ``torch.tensor``).
+    """
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _init_delayed_scaling_state,
+    )
+
+    _DELAYED_BUF_NAMES = (
+        "scale_input",
+        "scale_weight",
+        "scale_grad",
+        "amax_history_input",
+        "amax_history_weight",
+        "amax_history_grad",
+        "staged_input_amax",
+        "staged_grad_amax",
+        "staged_weight_amax",
+    )
+
+    count = 0
+    for m in models:
+        for module in m.modules():
+            if not getattr(module, "_use_delayed_scaling", False):
+                continue
+            device = module.weight.device
+            _init_delayed_scaling_state(module)
+            for buf_name in _DELAYED_BUF_NAMES:
+                buf = module._buffers.get(buf_name)
+                if buf is not None and buf.device != device:
+                    module._buffers[buf_name] = buf.to(device)
+            count += 1
+    return count
+
+
+def _neuter_optimizer(optimizer):
+    """Set optimizer to no-op mode: betas=[1,1], weight_decay=0."""
+    saved = []
+    inner = getattr(optimizer, "optimizer", optimizer)
+    _log(
+        f"Neutering optimizer: type={type(optimizer).__name__}, "
+        f"inner={type(inner).__name__}, "
+        f"param_groups={len(inner.param_groups)}"
+    )
+    for group in inner.param_groups:
+        state = {}
+        for key in ("betas", "weight_decay", "bias_correction", "pre_mult_wd"):
+            if key in group:
+                state[key] = group[key]
+        saved.append(state)
+
+        if "betas" in group:
+            group["betas"] = [1.0, 1.0]
+        if "weight_decay" in group:
+            group["weight_decay"] = 0.0
+        if "bias_correction" in group:
+            group["bias_correction"] = False
+        if "pre_mult_wd" in group:
+            group["pre_mult_wd"] = 0.0
+    return saved
+
+
+def _restore_optimizer(optimizer, saved):
+    """Restore optimizer hyperparams (betas, weight_decay, etc.)."""
+    inner = getattr(optimizer, "optimizer", optimizer)
+    for group, state in zip(inner.param_groups, saved):
+        for key, val in state.items():
+            group[key] = val
+    _log("Restored optimizer parameters")
+
+
+def _reset_optimizer_state(optimizer):
+    """Zero per-parameter step counters so Adam acts as if no steps occurred.
+
+    Handles both flat ``MegatronOptimizer`` wrappers and ``ChainedOptimizer``
+    which wraps several sub-optimizers.  Step counts live either in
+    ``param_groups[i]["step"]`` (Apex / TE FusedAdam) or in
+    ``optimizer.state[p]["step"]`` (stock PyTorch Adam).
+    """
+
+    def _reset_single(opt):
+        inner = getattr(opt, "optimizer", opt)
+        for group in inner.param_groups:
+            if "step" in group:
+                group["step"] = 0
+        for state in inner.state.values():
+            if isinstance(state, dict) and "step" in state:
+                if isinstance(state["step"], torch.Tensor):
+                    state["step"].zero_()
+                else:
+                    state["step"] = 0
+
+    if hasattr(optimizer, "chained_optimizers"):
+        for sub_opt in optimizer.chained_optimizers:
+            _reset_single(sub_opt)
+    else:
+        _reset_single(optimizer)
+    _log("Reset optimizer step counters")
+
+
+@register_patch(
+    "megatron.training.mlperf_warmup",
+    backend="megatron",
+    phase="before_train",
+    description="MLPerf warmup: synthetic data steps before measured training",
+    condition=_warmup_enabled,
+    priority=95,
+)
+def patch_mlperf_warmup(ctx: PatchContext):
+    """Install warmup hook on train_step at priority 95 (outermost wrapper)."""
+    import megatron.training.training as mt
+
+    if hasattr(mt.train_step, "_primus_warmup_hook"):
+        return
+
+    primus_args = get_args(ctx)
+    warmup_steps = getattr(primus_args, "warmup_train_steps", 2)
+
+    _lazy_state = {
+        "initialized": False,
+        "synthetic_iter": None,
+        "use_fsdp2_fp8": False,
+        "transformer_impl": "local",
+    }
+
+    _wrapped_chain = mt.train_step
+    _warmup_done = [False]
+
+    def _lazy_init():
+        """One-time initialization on first train_step call, when Megatron args exist."""
+        if _lazy_state["initialized"]:
+            return
+
+        from megatron.training import get_args as megatron_get_args
+
+        megatron_args = megatron_get_args()
+
+        from torch.utils.data import DataLoader
+
+        from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper
+        from primus.backends.megatron.data.synthetic.mock_datasets import (
+            PreGeneratedMockFluxSchnellDataset,
+        )
+
+        image_size = getattr(primus_args, "image_size", 256)
+        vae_latent_mode = getattr(primus_args, "vae_latent_mode", "resample")
+        mbs = getattr(primus_args, "micro_batch_size", 64)
+
+        mock_dataset = PreGeneratedMockFluxSchnellDataset(
+            num_samples=max(mbs * 4, 256),
+            image_size=image_size,
+            vae_latent_mode=vae_latent_mode,
+        )
+        mock_loader = DataLoader(mock_dataset, batch_size=mbs, shuffle=False, drop_last=True)
+        _lazy_state["synthetic_iter"] = MegatronDataloaderWrapper(mock_loader)
+
+        _lazy_state["use_fsdp2_fp8"] = getattr(megatron_args, "use_fsdp2_fp8_all_gather", False)
+        _lazy_state["transformer_impl"] = getattr(megatron_args, "transformer_impl", "local")
+        _lazy_state["initialized"] = True
+        _log(f"Lazy init complete (warmup_steps={warmup_steps})")
+
+    def _hooked_train_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        config,
+        forward_backward_func,
+        iteration=None,
+    ):
+        if _warmup_done[0]:
+            return _wrapped_chain(
+                forward_step_func,
+                data_iterator,
+                model,
+                optimizer,
+                opt_param_scheduler,
+                config,
+                forward_backward_func,
+                iteration=iteration,
+            )
+
+        _lazy_init()
+
+        from megatron.training import get_args as megatron_get_args
+
+        megatron_args = megatron_get_args()
+        models = model if isinstance(model, (list, tuple)) else [model]
+        synthetic_iter = _lazy_state["synthetic_iter"]
+
+        # ---- 1. Snapshot model parameters to CPU ----
+        _log("Saving model parameters to CPU before warmup")
+        saved_params = {}
+        for m in models:
+            for name, p in m.named_parameters():
+                saved_params[name] = p.data.to("cpu", non_blocking=True)
+        torch.cuda.synchronize()
+        _log(f"Saved {len(saved_params)} parameter tensors")
+
+        # ---- 2. Neuter optimizer ----
+        saved_opt = _neuter_optimizer(optimizer)
+
+        # ---- 3. Suppress training_log and eval during warmup ----
+        saved_training_log = mt.training_log
+        saved_eval = mt.evaluate_and_print_results
+        mt.training_log = lambda *a, **k: None
+        mt.evaluate_and_print_results = lambda *a, **k: None
+
+        # ---- 3b. Save LR scheduler state (NeMo never steps the scheduler during warmup) ----
+        saved_lr_num_steps = opt_param_scheduler.num_steps
+
+        # ---- 4. Run warmup steps with synthetic data ----
+        for step_idx in range(warmup_steps):
+            _log(f"Warmup step {step_idx + 1}/{warmup_steps}")
+            _wrapped_chain(
+                forward_step_func,
+                synthetic_iter,
+                model,
+                optimizer,
+                opt_param_scheduler,
+                config,
+                forward_backward_func,
+                iteration=iteration,
+            )
+        _log(f"Completed {warmup_steps} warmup steps")
+
+        # ---- 5. Restore optimizer ----
+        _restore_optimizer(optimizer, saved_opt)
+        _reset_optimizer_state(optimizer)
+
+        # ---- 6. Restore model parameters from CPU ----
+        restored = 0
+        for m in models:
+            for name, p in m.named_parameters():
+                if name in saved_params:
+                    p.data.copy_(saved_params[name])
+                    restored += 1
+        del saved_params
+        _log(f"Restored {restored} parameter tensors from CPU snapshot")
+
+        # ---- 7. FP8 reset (spec-aware) ----
+        if _lazy_state["transformer_impl"] == "transformer_engine":
+            te_count = _reset_fp8_te_spec(models)
+            amax_count = _seed_fp8_amax(models)
+            _log(f"FP8 TE reset: {te_count} modules, " f"seeded {amax_count} amax tensors")
+        else:
+            local_count = _reset_fp8_local_spec(models)
+            _log(f"FP8 local spec reset: {local_count} modules")
+
+        # ---- 8. FSDP2 FP8 all-gather recompute ----
+        if _lazy_state["use_fsdp2_fp8"]:
+            try:
+                from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+                    precompute_fp8_scales_for_fsdp,
+                )
+
+                cache_data = getattr(megatron_args, "fp8_precompute_data_cache", True)
+                use_cpp = getattr(megatron_args, "use_cpp_fp8_quantize", False)
+                sr = getattr(megatron_args, "fp8_all_gather_stochastic_rounding", False)
+                precompute_fp8_scales_for_fsdp(
+                    models[0],
+                    cache_data=cache_data,
+                    use_cpp_quantize=use_cpp,
+                    stochastic_rounding=sr,
+                )
+                _log("Recomputed FSDP2 FP8 all-gather scales")
+            except Exception as e:
+                _log(f"FSDP2 FP8 recompute failed (non-fatal): {e}")
+
+        # ---- 9. Reload model params in optimizer (FSDP2 BF16 master weight) ----
+        if hasattr(optimizer, "reload_model_params"):
+            optimizer.reload_model_params()
+            _log("Called optimizer.reload_model_params()")
+
+        # ---- 10. Post-restore NaN check ----
+        nan_params = 0
+        for m in models:
+            for name, p in m.named_parameters():
+                if p.data.is_floating_point() and torch.isnan(p.data).any():
+                    nan_params += 1
+        _log(f"Post-restore parameter check: nan_params={nan_params}")
+
+        # ---- 11. Zero gradients ----
+        try:
+            optimizer.zero_grad(set_to_none=True)
+        except TypeError:
+            optimizer.zero_grad()
+
+        # ---- 12. Reset counters ----
+        megatron_args.consumed_train_samples = 0
+        megatron_args.skipped_train_samples = 0
+        opt_param_scheduler.num_steps = saved_lr_num_steps
+        _log(
+            f"Reset consumed_train_samples=0, skipped_train_samples=0, "
+            f"lr_scheduler.num_steps={saved_lr_num_steps}"
+        )
+
+        # ---- 13. Restore training_log and eval ----
+        mt.training_log = saved_training_log
+        mt.evaluate_and_print_results = saved_eval
+
+        # ---- 13b. Invalidate the CudaPrefetchIterator that was built around
+        # the SYNTHETIC iterator during warmup step 1.
+        #
+        # ``patch_grad_zero_and_data_prefetch`` builds a ``CudaPrefetchIterator``
+        # the first time its ``_patched_train_step`` runs and caches it in a
+        # closure-local ``_prefetch_state["iter"]``.  Because warmup step 1
+        # is the first call into that train_step, the prefetch iterator gets
+        # bound to ``synthetic_iter``.  ``MegatronDataloaderWrapper`` is
+        # cyclic (never raises ``StopIteration``), so subsequent real
+        # training steps would silently keep reading from the cycling
+        # synthetic dataset instead of the actual training dataset -- model
+        # overfits the mock samples and val_loss on real data stays stuck
+        # at ~1.38 forever.
+        #
+        # Dropping the cached entry forces the next train_step to rebuild
+        # the prefetch wrapper around its incoming ``data_iterator`` arg
+        # (the real iterator).
+        try:
+            from primus.backends.megatron.patches.delayed_fp8_scaling_patches import (
+                reset_prefetch_state,
+            )
+
+            evicted = reset_prefetch_state()
+            if evicted is None:
+                _log("  Prefetch reset: no cached iterator to evict")
+            else:
+                _log(
+                    f"  Prefetch reset: evicted cached {type(evicted).__name__} "
+                    f"(wrapped synthetic warmup iterator) -- next train_step "
+                    f"will rebuild it around the real data_iterator"
+                )
+        except Exception as _e:
+            _log(f"  Prefetch reset failed (non-fatal): {_e}")
+
+        # ---- 14. Synchronize ----
+        torch.cuda.synchronize()
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        # ---- 15. Execute first real step ----
+        _log("Executing first real train_step with training data")
+        result = _wrapped_chain(
+            forward_step_func,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            config,
+            forward_backward_func,
+            iteration=iteration,
+        )
+
+        # ---- 16. Self-remove ----
+        _warmup_done[0] = True
+        mt.train_step = _wrapped_chain
+        _log("Self-removed warmup hook, train_step = inner wrapped chain")
+
+        return result
+
+    _hooked_train_step._primus_warmup_hook = True
+    mt.train_step = _hooked_train_step
+
+    _log(
+        f"Installed MLPerf warmup hook (warmup_steps={warmup_steps}, "
+        f"deferred init until first train_step)"
+    )
+    log_rank_0(f"[Patch:mlperf_warmup] Installed warmup hook " f"(warmup_steps={warmup_steps}, priority=95)")
diff --git a/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py b/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py
new file mode 100644
index 000000000..d4e15eeca
--- /dev/null
+++ b/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py
@@ -0,0 +1,105 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus MoE All-to-All dispatcher D2H patch.
+
+Patches ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize`` so that when
+``use_turbo_grouped_gemm`` is enabled, ``tokens_per_expert`` is kept on-device
+(PrimusTurbo grouped gemm consumes it on the GPU) instead of being copied to the
+host. All other splits are still moved to CPU and the stream sync is unchanged.
+"""
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _turbo_grouped_gemm_on_device(ctx: PatchContext) -> bool:
+    """Whether PrimusTurbo grouped gemm consumes tokens_per_expert on device.
+
+    The authoritative source is the global Primus/Megatron args, not the
+    dispatcher's ``TransformerConfig`` (the turbo flags are CLI/args-level and
+    are not guaranteed to be mirrored onto ``self.config``).
+    ``use_turbo_grouped_gemm`` is the flag users set (also auto-enabled by
+    Sync-Free MoE stage >= 2).
+    """
+    try:
+        args = get_args(ctx)
+    except Exception:
+        return False
+    return bool(getattr(args, "use_turbo_grouped_gemm", False))
+
+
+@register_patch(
+    "megatron.moe_alltoall_dtoh_turbo_grouped_gemm",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Skip tokens_per_expert D2H copy in MoEAlltoAllTokenDispatcher "
+        "when PrimusTurbo grouped gemm (use_turbo_grouped_gemm) is enabled"
+    ),
+)
+def patch_moe_alltoall_dtoh(ctx: PatchContext):
+    """Replace ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize``."""
+    from megatron.core.transformer.moe import token_dispatcher as td_mod
+
+    cls = td_mod.MoEAlltoAllTokenDispatcher
+
+    keep_tokens_per_expert_on_device = _turbo_grouped_gemm_on_device(ctx)
+
+    def _maybe_dtoh_and_synchronize(self, point, tokens_per_expert=None):
+        """
+        Move all possible GPU tensors to CPU and make a synchronization at the expected point.
+        """
+        maybe_move_tensor_to_cpu = td_mod.maybe_move_tensor_to_cpu
+
+        if not self.drop_and_pad:
+            if point == self.cuda_dtoh_point:
+                # Move all possible GPU tensors to CPU at self.cuda_dtoh_point.
+                on_side_stream = torch.cuda.current_stream() != self.cuda_dtoh_stream
+                if on_side_stream:
+                    self.cuda_dtoh_stream.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(self.cuda_dtoh_stream):
+                    # TODO: use MemcpyBatchAsync instead.
+                    # PrimusTurbo grouped gemm consumes tokens_per_expert on device,
+                    # so keep it on the GPU and skip the D2H copy when enabled.
+                    if not keep_tokens_per_expert_on_device:
+                        tokens_per_expert = maybe_move_tensor_to_cpu(
+                            tokens_per_expert, record_stream=on_side_stream
+                        )
+                    self.input_splits = maybe_move_tensor_to_cpu(
+                        self.input_splits, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.output_splits = maybe_move_tensor_to_cpu(
+                        self.output_splits, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.output_splits_tp = maybe_move_tensor_to_cpu(
+                        self.output_splits_tp, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.num_out_tokens = maybe_move_tensor_to_cpu(
+                        self.num_out_tokens, record_stream=on_side_stream
+                    )
+                    if self.num_local_experts > 1 and not self.config.moe_permute_fusion:
+                        self.num_global_tokens_per_local_expert = maybe_move_tensor_to_cpu(
+                            self.num_global_tokens_per_local_expert, record_stream=on_side_stream
+                        )
+                self.d2h_event = self.cuda_dtoh_stream.record_event()
+
+            if point == self.cuda_sync_point:
+                # Synchronize with the DtoH stream at self.cuda_sync_point.
+                self.d2h_event.synchronize()
+
+        return tokens_per_expert
+
+    cls._maybe_dtoh_and_synchronize = _maybe_dtoh_and_synchronize
+
+    log_rank_0(
+        "[Patch:megatron.moe_alltoall_dtoh_turbo_grouped_gemm]   Patched "
+        "MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize "
+        f"(skip tokens_per_expert D2H = {keep_tokens_per_expert_on_device})"
+    )
diff --git a/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py b/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py
index b2259599d..79db6b803 100644
--- a/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py
+++ b/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py
@@ -13,7 +13,7 @@
 import sys
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py b/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py
new file mode 100644
index 000000000..17a75a67b
--- /dev/null
+++ b/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py
@@ -0,0 +1,85 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus MoE All-to-All dispatcher D2H patch.
+
+Patches ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize`` so that when
+``use_turbo_grouped_gemm`` is enabled, ``tokens_per_expert`` is kept on-device
+(PrimusTurbo grouped gemm consumes it on the GPU) instead of being copied to the
+host. All other splits are still moved to CPU and the stream sync is unchanged.
+"""
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.moe_alltoall_dtoh_turbo_grouped_gemm",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Skip tokens_per_expert D2H copy in MoEAlltoAllTokenDispatcher "
+        "when PrimusTurbo grouped gemm is enabled"
+    ),
+    condition=lambda ctx: getattr(get_args(ctx), "use_turbo_grouped_gemm", False),
+)
+def patch_moe_alltoall_dtoh(ctx: PatchContext):
+    """Replace ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize``."""
+    from megatron.core.transformer.moe import token_dispatcher
+
+    cls = token_dispatcher.MoEAlltoAllTokenDispatcher
+
+    def _maybe_dtoh_and_synchronize(self, point, tokens_per_expert=None):
+        """
+        Move all possible GPU tensors to CPU and make a synchronization at the expected point.
+        """
+        maybe_move_tensor_to_cpu = token_dispatcher.maybe_move_tensor_to_cpu
+
+        if not self.drop_and_pad:
+            if point == self.cuda_dtoh_point:
+                # Move all possible GPU tensors to CPU at self.cuda_dtoh_point.
+                on_side_stream = torch.cuda.current_stream() != self.cuda_dtoh_stream
+                if on_side_stream:
+                    self.cuda_dtoh_stream.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(self.cuda_dtoh_stream):
+                    # NOTE: PrimusTurbo grouped gemm consumes tokens_per_expert on device,
+                    # so keep it on the GPU and skip the D2H copy when enabled.
+                    # tokens_per_expert = maybe_move_tensor_to_cpu(
+                    #     tokens_per_expert, record_stream=on_side_stream
+                    # )
+                    self.input_splits = maybe_move_tensor_to_cpu(
+                        self.input_splits, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.output_splits = maybe_move_tensor_to_cpu(
+                        self.output_splits, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.output_splits_tp = maybe_move_tensor_to_cpu(
+                        self.output_splits_tp, as_numpy=True, record_stream=on_side_stream
+                    )
+                    self.num_out_tokens = maybe_move_tensor_to_cpu(
+                        self.num_out_tokens, record_stream=on_side_stream
+                    )
+                    if self.num_local_experts > 1 and not self.config.moe_permute_fusion:
+                        self.num_global_tokens_per_local_expert = maybe_move_tensor_to_cpu(
+                            self.num_global_tokens_per_local_expert, record_stream=on_side_stream
+                        )
+                self.d2h_event = self.cuda_dtoh_stream.record_event()
+
+            if point == self.cuda_sync_point:
+                # Synchronize with the DtoH stream at self.cuda_sync_point.
+                self.d2h_event.synchronize()
+
+        return tokens_per_expert
+
+    cls._maybe_dtoh_and_synchronize = _maybe_dtoh_and_synchronize
+
+    log_rank_0(
+        "[Patch:megatron.moe_alltoall_dtoh_turbo_grouped_gemm]   Patched "
+        "MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize "
+    )
diff --git a/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py b/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py
index d00d2b84b..b889b64c8 100644
--- a/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py
+++ b/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py
@@ -11,7 +11,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py
new file mode 100644
index 000000000..8d08d4af1
--- /dev/null
+++ b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py
@@ -0,0 +1,121 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+MoE identity-sort short-circuit patch.
+
+Migrated from the source patch ``megatron_moe_skip_identity_sort.patch``
+(mpo branch).
+
+In the all-to-all MoE token dispatcher, ``sort_chunks_by_idxs`` is called with
+``sort_input_by_local_experts`` (dispatch) and ``restore_output_by_local_experts``
+(combine). When the EP/TP topology yields an *identity* permutation
+(``[0, 1, ..., N-1]``) -- the common ``EP=1`` / ``TP=1`` case -- the call
+degrades into a full-tensor ``split`` + ``cat`` round-trip that produces an
+output identical to its input. This patch short-circuits those calls.
+
+Rather than editing ``third_party/Megatron-LM`` in place, we wrap the
+``sort_chunks_by_idxs`` symbol bound inside the ``token_dispatcher`` module.
+The identity decision is cached per index tensor (the topology indices are
+created once per dispatcher and reused every step), so the ``torch.equal``
+probe -- and the device sync it implies -- happens at most once per tensor.
+
+Gate:
+    ``MOE_SKIP_IDENTITY_SORT`` (default ``1`` / enabled; set ``0`` to disable).
+"""
+
+import os
+import weakref
+
+import torch
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+# Cache the identity decision keyed on the index tensor's ``id()``. The topology
+# index tensors persist for the dispatcher's lifetime, so caching by identity is
+# stable across steps. A ``WeakKeyDictionary`` cannot be used here because weakref
+# equality falls back to ``tensor == tensor`` (which yields a tensor, not a bool)
+# and raises "Boolean value of Tensor ... is ambiguous". Instead we key on
+# ``id(idxs)`` and register a ``weakref.finalize`` callback so the entry is
+# evicted when the tensor is garbage-collected (preventing stale hits from
+# ``id()`` reuse).
+_IDENTITY_CACHE: dict = {}
+
+
+def _skip_enabled(_ctx: PatchContext) -> bool:
+    return os.environ.get("MOE_SKIP_IDENTITY_SORT", "1") != "0"
+
+
+def _is_identity(idxs) -> bool:
+    if idxs is None or not torch.is_tensor(idxs) or idxs.dim() != 1 or idxs.numel() == 0:
+        return False
+    key = id(idxs)
+    cached = _IDENTITY_CACHE.get(key)
+    if cached is not None:
+        return cached
+    result = bool(torch.equal(idxs, torch.arange(idxs.numel(), device=idxs.device, dtype=idxs.dtype)))
+    try:
+        # Evict on GC so a future tensor reusing this ``id()`` cannot hit a stale
+        # decision.
+        weakref.finalize(idxs, _IDENTITY_CACHE.pop, key, None)
+        _IDENTITY_CACHE[key] = result
+    except TypeError:
+        # Some tensors are not weak-referenceable; fall back to no caching.
+        pass
+    return result
+
+
+def _make_wrapped_sort(orig_sort):
+    def sort_chunks_by_idxs(input, split_sizes, sorted_idxs, probs=None, fused=False):
+        # Identity permutation -> output == input, permuted_probs == probs.
+        if _is_identity(sorted_idxs):
+            return input, probs
+        return orig_sort(input, split_sizes, sorted_idxs, probs=probs, fused=fused)
+
+    return sort_chunks_by_idxs
+
+
+@register_patch(
+    "megatron.moe.skip_identity_sort",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Short-circuit sort_chunks_by_idxs in the MoE token dispatcher when the "
+        "local-expert permutation is identity (e.g. EP=1/TP=1); gated by "
+        "MOE_SKIP_IDENTITY_SORT (default on)."
+    ),
+    condition=_skip_enabled,
+)
+def patch_moe_skip_identity_sort(ctx: PatchContext):
+    del ctx
+
+    try:
+        from megatron.core.transformer.moe import token_dispatcher as td_mod
+    except ImportError as exc:
+        warning_rank_0(
+            f"[Patch:megatron.moe.skip_identity_sort] token_dispatcher not " f"importable; skipping: {exc}"
+        )
+        return
+
+    orig_sort = getattr(td_mod, "sort_chunks_by_idxs", None)
+    if orig_sort is None:
+        warning_rank_0(
+            "[Patch:megatron.moe.skip_identity_sort] sort_chunks_by_idxs not bound "
+            "in token_dispatcher; skipping."
+        )
+        return
+
+    if getattr(orig_sort, "_primus_skip_identity_sort", False):
+        return
+
+    wrapped = _make_wrapped_sort(orig_sort)
+    wrapped._primus_skip_identity_sort = True
+    td_mod.sort_chunks_by_idxs = wrapped
+    log_rank_0(
+        "[Patch:megatron.moe.skip_identity_sort] Patched "
+        "token_dispatcher.sort_chunks_by_idxs to skip identity permutations."
+    )
diff --git a/primus/backends/megatron/patches/moe_patches/topk_router_patches.py b/primus/backends/megatron/patches/moe_patches/topk_router_patches.py
index 77f4c7e5c..6fe2e90e6 100644
--- a/primus/backends/megatron/patches/moe_patches/topk_router_patches.py
+++ b/primus/backends/megatron/patches/moe_patches/topk_router_patches.py
@@ -24,7 +24,7 @@ class MoESubmodules:
 import sys
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/mp_sync_skip_patches.py b/primus/backends/megatron/patches/mp_sync_skip_patches.py
new file mode 100644
index 000000000..c32ca3861
--- /dev/null
+++ b/primus/backends/megatron/patches/mp_sync_skip_patches.py
@@ -0,0 +1,63 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Skip redundant model-parallel synchronization when TP=1 and PP=1.
+
+When running pure data-parallel (TP=1, PP=1), the model-parallel process
+group has size 1.  The upstream Megatron ``train_step`` and ``training_log``
+unconditionally call ``logical_and_across_model_parallel_group`` and
+``reduce_max_stat_across_model_parallel_group`` which each issue an
+``all_reduce`` + ``.item()`` pair.  With a size-1 group these are no-ops
+that still force a GPU-to-CPU synchronization barrier, draining the
+asynchronous GPU pipeline.
+
+This patch replaces those two module-level names inside
+``megatron.training.training`` with lightweight pass-through functions
+that preserve the return-type contract (``bool`` and ``float | None``)
+without issuing any collective or ``.item()`` call.
+"""
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _is_pure_dp(ctx: PatchContext) -> bool:
+    """True when TP=1 and PP=1 -- the MP group has size 1."""
+    args = get_args(ctx)
+    if args is None:
+        return False
+    tp = getattr(args, "tensor_model_parallel_size", 1)
+    pp = getattr(args, "pipeline_model_parallel_size", 1)
+    return tp == 1 and pp == 1
+
+
+@register_patch(
+    "megatron.training.skip_redundant_mp_sync",
+    backend="megatron",
+    phase="before_train",
+    description="Skip redundant model-parallel all-reduces when TP=1 and PP=1",
+    condition=_is_pure_dp,
+    priority=35,
+)
+def patch_skip_redundant_mp_sync(ctx: PatchContext):
+    import megatron.training.training as megatron_training
+
+    def _passthrough_logical_and(val):
+        return val
+
+    def _passthrough_reduce_max(val):
+        if val is None:
+            return None
+        if isinstance(val, torch.Tensor):
+            return val.item()
+        return val
+
+    megatron_training.logical_and_across_model_parallel_group = _passthrough_logical_and
+    megatron_training.reduce_max_stat_across_model_parallel_group = _passthrough_reduce_max
+    log_rank_0("[Patch:skip_redundant_mp_sync] " "Replaced MP sync functions with passthrough (TP=1, PP=1)")
diff --git a/primus/backends/megatron/patches/muon_optimizer_patches.py b/primus/backends/megatron/patches/muon_optimizer_patches.py
deleted file mode 100644
index 2d96dfa39..000000000
--- a/primus/backends/megatron/patches/muon_optimizer_patches.py
+++ /dev/null
@@ -1,109 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-"""
-Megatron Muon Optimizer patches.
-
-This module patches megatron.training.training.get_megatron_optimizer to
-automatically dispatch to get_megatron_muon_optimizer when args.optimizer
-contains "muon". Since training.py uses `from megatron.core.optimizer import
-get_megatron_optimizer`, we must patch the training module's namespace where
-the function is actually used, not megatron.core.optimizer.
-"""
-
-import dataclasses
-import inspect
-
-from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
-
-
-@register_patch(
-    "megatron.optimizer.muon",
-    backend="megatron",
-    phase="before_train",
-    description="Patch get_megatron_optimizer to dispatch to muon optimizer when optimizer name contains 'muon'.",
-)
-def patch_get_megatron_optimizer_muon(ctx: PatchContext) -> None:
-    """
-    Patch megatron.training.training.get_megatron_optimizer to delegate to
-    get_megatron_muon_optimizer when config.optimizer contains "muon".
-
-    We patch the training module (not megatron.core.optimizer) because
-    training.py imports get_megatron_optimizer into its namespace at import
-    time; patching the optimizer module would not affect the training module's
-    local reference.
-    """
-    try:
-        import megatron.training.training as training_module
-    except ImportError as e:
-        log_rank_0(f"[Patch:megatron.optimizer.muon] Skip patch (Megatron not available): {e}")
-        return
-
-    original_get_megatron_optimizer = training_module.get_megatron_optimizer
-    original_signature = inspect.signature(original_get_megatron_optimizer)
-
-    if getattr(original_get_megatron_optimizer, "_primus_muon_wrapper", False):
-        return
-
-    def _get_bound_arg(bound_arguments, name, fallback=None):
-        if name in bound_arguments:
-            return bound_arguments[name]
-
-        parameter = original_signature.parameters.get(name)
-        if parameter and parameter.default is not inspect.Parameter.empty:
-            return parameter.default
-
-        return fallback
-
-    def _patched_get_megatron_optimizer(*func_args, **func_kwargs):
-        config = func_kwargs.get("config")
-        if config is None and func_args:
-            config = func_args[0]
-
-        optimizer_name = getattr(config, "optimizer", None)
-        if not optimizer_name or "muon" not in optimizer_name:
-            return original_get_megatron_optimizer(*func_args, **func_kwargs)
-
-        bound_arguments = original_signature.bind_partial(*func_args, **func_kwargs).arguments
-        model_chunks = _get_bound_arg(bound_arguments, "model_chunks")
-        config_overrides = _get_bound_arg(bound_arguments, "config_overrides")
-        use_gloo_process_groups = _get_bound_arg(bound_arguments, "use_gloo_process_groups", True)
-        pg_collection = _get_bound_arg(bound_arguments, "pg_collection")
-        dump_param_to_param_group_map = _get_bound_arg(bound_arguments, "dump_param_to_param_group_map")
-
-        from primus.backends.megatron.core.optimizer.moun import (
-            get_megatron_muon_optimizer,
-        )
-        from primus.backends.megatron.core.optimizer.moun_optimizer_config import (
-            MounOptimizerConfig,
-        )
-
-        args = ctx.extra.get("backend_args", {})
-        kwargs = {}
-        for f in dataclasses.fields(MounOptimizerConfig):
-            if hasattr(args, f.name):
-                kwargs[f.name] = getattr(args, f.name)
-
-        moun_config = MounOptimizerConfig(**kwargs)
-        moun_config.timers = config.timers
-
-        return get_megatron_muon_optimizer(
-            moun_config,
-            model_chunks,
-            config_overrides=config_overrides,
-            use_gloo_process_groups=use_gloo_process_groups,
-            layer_wise_distributed_optimizer="dist" in optimizer_name,
-            pg_collection=pg_collection,
-            dump_param_to_param_group_map=dump_param_to_param_group_map,
-        )
-
-    setattr(_patched_get_megatron_optimizer, "_primus_muon_wrapper", True)
-    training_module.get_megatron_optimizer = _patched_get_megatron_optimizer
-    log_rank_0(
-        "[Patch:megatron.optimizer.muon] Patched get_megatron_optimizer in megatron.training.training "
-        "to dispatch to muon when optimizer contains 'muon'."
-    )
diff --git a/primus/backends/megatron/patches/odc_lb_mini_patches.py b/primus/backends/megatron/patches/odc_lb_mini_patches.py
new file mode 100644
index 000000000..e240274e6
--- /dev/null
+++ b/primus/backends/megatron/patches/odc_lb_mini_patches.py
@@ -0,0 +1,268 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+###############################################################################
+# LB-Mini (sequence-length load balancing) for Megatron's FSDP2 path.
+#
+# ``enable_odc_lb_mini`` is an ORTHOGONAL, standalone capability: it serves the
+# variable-length, Karmarkar-Karp-balanced LB-Mini DATA on the torch-FSDP2 path.
+# It is INDEPENDENT of the ODC communication switch (``enable_odc``); what the
+# ODC switch selects is only HOW the per-rank micro-batch counts are aligned:
+#
+#   * enable_odc_lb_mini=false (DEFAULT) -> this patch is a complete no-op;
+#     Megatron runs its stock fixed-num_microbatches, all-ranks-in-lockstep
+#     schedule on stock (padded) data. Byte-for-byte unchanged.
+#
+#   * enable_odc_lb_mini=true + enable_odc=true -> DECOUPLED mode. Data is served
+#     variable length and KK-balanced across DP ranks; each rank runs its OWN
+#     (possibly different) number of micro-batches (same_micro_num=False). Only
+#     ODC's point-to-point comm can drive ranks out of lockstep without a
+#     collective deadlock, hence this mode requires ODC comm.
+#
+#   * enable_odc_lb_mini=true + enable_odc=false -> ALIGNED mode. The SAME
+#     variable-length KK-balanced DATA is served, but the micro-batch count is
+#     all-reduce(MAX)-aligned so every rank runs the SAME number of steps
+#     (same_micro_num=True). Uniform per-rank counts keep standard FSDP2 + RCCL
+#     collectives in lockstep (no deadlock), so this is a fair "same data" nccl
+#     baseline WITHOUT ODC. This is the config-driven replacement for the removed
+#     LB_MINI_FORCE_DATA A/B env (which likewise served LB-Mini data under NCCL).
+#
+# In BOTH enabled modes we install the dataloader patch (variable-length data)
+# AND the schedule patch (rank-local num_microbatches). The schedule patch is
+# comm-agnostic -- it only overrides num_microbatches with the iterator's planned
+# per-rank count. Under ALIGNED mode those counts are identical across ranks, so
+# it is NCCL/RCCL-safe; it is NOT an ODC-specific reduction. (Running the
+# dataloader patch WITHOUT the schedule patch would be incoherent: the stock
+# schedule would pull a fixed num_microbatches while the iterator plans a
+# possibly-different per-rank count, drifting the two out of sync.)
+#
+# All wiring is monkey-patch in the Primus layer; the third-party Megatron-LM
+# source is NOT modified.
+#
+# Stage-1 scope (this file): make "different micro-batch count per rank" run end
+# to end without deadlocking. Numerical normalization (loss_scale / consumed
+# samples by real tokens) is Stage-2.
+###############################################################################
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+# Global handle so the schedule patch can reach the LB-Mini iterator that the
+# dataloader patch created. Stage-1 only drives the TRAIN iterator (eval_iters=0
+# in the aligned config), so a single handle is sufficient.
+_LB_MINI_TRAIN_ITER = None
+_FB_PATCHED = False
+
+
+def _lb_mini_enabled(args) -> bool:
+    """LB-Mini DATA serving is driven by ``enable_odc_lb_mini`` ALONE.
+
+    LB-Mini is ORTHOGONAL to the ODC comm switch: it only requires the explicit
+    ``enable_odc_lb_mini`` config item and the torch-FSDP2 path it patches. The
+    ``enable_odc`` switch does NOT gate whether LB-Mini data is served; it only
+    selects the micro-batch alignment mode (see ``_lb_mini_aligned``):
+
+      * enable_odc=true  -> DECOUPLED (ranks may run different micro-batch counts;
+                            needs ODC point-to-point comm).
+      * enable_odc=false -> ALIGNED   (all ranks run the same micro-batch count
+                            via all_reduce(MAX); NCCL/RCCL-safe "same data"
+                            baseline, no ODC).
+    """
+    return bool(getattr(args, "enable_odc_lb_mini", False)) and bool(getattr(args, "use_torch_fsdp2", False))
+
+
+def _lb_mini_aligned(args) -> bool:
+    """Micro-batch alignment mode for LB-Mini.
+
+    ALIGNED (True) when ODC comm is OFF: all ranks are forced to the same
+    micro-batch count (``same_micro_num=True``, all_reduce MAX) so standard
+    FSDP2 + RCCL collectives stay in lockstep. DECOUPLED (False) when ODC comm is
+    ON: ranks may run different counts, which only ODC's point-to-point comm can
+    drive without a collective deadlock.
+    """
+    return not bool(getattr(args, "enable_odc", False))
+
+
+def _build_lb_mini_train_iterator(args):
+    """Build the variable-length, KK-balanced LB-Mini train iterator."""
+    from megatron.core import mpu
+    from megatron.training import get_tokenizer
+
+    from primus.backends.megatron.sft.lb_mini_dataset import (
+        LBMiniDataIterator,
+        build_varlen_samples,
+    )
+    from primus.backends.megatron.sft.packing import _resolve_pad_token_id
+
+    tokenizer = get_tokenizer()
+    samples = build_varlen_samples(
+        dataset_name=getattr(args, "sft_dataset_name", "tatsu-lab/alpaca"),
+        tokenizer=tokenizer,
+        max_seq_length=args.seq_length,
+        # Some datasets (e.g. SWE-bench/SWE-smith-trajectories) have no "train"
+        # split; sft_dataset_split lets the config pick it (default "train").
+        split=str(getattr(args, "sft_dataset_split", "train")),
+        formatter=getattr(args, "sft_conversation_format", "alpaca"),
+        seed=args.seed,
+        bridge_compat_inline_bos=bool(getattr(args, "sft_bridge_compat_inline_bos", False)),
+    )
+    # Per-micro-batch token cap. Larger than a single sample lets short samples
+    # pack together and long samples stand alone -> creates per-rank micro-batch
+    # count differences (where DiffMicro saves comm rounds). Priority:
+    # yaml lb_mini_max_token_len > seq_length.
+    max_token_len = int(getattr(args, "lb_mini_max_token_len", 0) or 0) or int(args.seq_length)
+    # ALIGNED (enable_odc=false) -> same_micro_num=True: all ranks run the SAME
+    # micro-batch count (all_reduce MAX), keeping standard RCCL collectives in
+    # lockstep -> fair "same data" NCCL baseline. DECOUPLED (enable_odc=true) ->
+    # same_micro_num=False: ranks may differ; only ODC p2p comm can drive that.
+    aligned = _lb_mini_aligned(args)
+    it = LBMiniDataIterator(
+        samples=samples,
+        global_batch_size=args.global_batch_size,
+        max_token_len=max_token_len,
+        dp_rank=mpu.get_data_parallel_rank(),
+        dp_size=mpu.get_data_parallel_world_size(),
+        pad_id=_resolve_pad_token_id(tokenizer),
+        cost_model=str(getattr(args, "lb_mini_cost_model", "linear")),
+        seed=args.seed,
+        shuffle=True,
+        same_micro_num=aligned,
+        packing_method="kk",
+    )
+    log_rank_0(
+        f"[ODC.lb_mini] built LB-Mini train iterator: {len(samples)} varlen samples, "
+        f"global_batch_size={args.global_batch_size}, max_token_len={max_token_len}, "
+        f"dp_size={mpu.get_data_parallel_world_size()}, "
+        f"cost_model={getattr(args, 'lb_mini_cost_model', 'linear')}, "
+        f"same_micro_num={aligned} "
+        f"({'ALIGNED baseline (nccl, no ODC)' if aligned else 'LB-Mini decoupled (ODC)'})"
+    )
+    return it
+
+
+def _install_dataloader_patch():
+    """Patch build_pretraining_data_loader so the TRAIN loader is LB-Mini.
+
+    Valid/test loaders (if any) fall through to the stock builder unchanged.
+    We tag the first (train) request via a module flag because the stock
+    signature does not carry the split explicitly.
+    """
+    import megatron.training.training as mt_training
+    from megatron.training.datasets import data_samplers
+
+    if getattr(data_samplers.build_pretraining_data_loader, "_lb_mini_hooked", False):
+        return
+
+    orig_builder = data_samplers.build_pretraining_data_loader
+
+    def lb_mini_builder(dataset, consumed_samples):
+        global _LB_MINI_TRAIN_ITER
+        # Runtime call: use Megatron's get_args() (no ctx). Primus' get_args(ctx)
+        # is only valid inside register_patch conditions / patch bodies.
+        from megatron.training import get_args as _mt_get_args
+
+        args = _mt_get_args()
+        # Only replace the TRAIN loader, and only once (the first non-zero-len
+        # build). Identify train by: not yet built + dataset present.
+        if _lb_mini_enabled(args) and _LB_MINI_TRAIN_ITER is None and dataset is not None:
+            try:
+                _LB_MINI_TRAIN_ITER = _build_lb_mini_train_iterator(args)
+                log_rank_0("[ODC.lb_mini] TRAIN dataloader replaced by LB-Mini iterator")
+                return _LB_MINI_TRAIN_ITER
+            except Exception as e:  # noqa: BLE001
+                warning_rank_0(
+                    f"[ODC.lb_mini] failed to build LB-Mini iterator, "
+                    f"falling back to stock loader: {type(e).__name__}: {e}"
+                )
+        return orig_builder(dataset, consumed_samples)
+
+    lb_mini_builder._lb_mini_hooked = True
+    data_samplers.build_pretraining_data_loader = lb_mini_builder
+    # training.py imported the symbol into its own namespace; rebind there too.
+    if hasattr(mt_training, "build_pretraining_data_loader"):
+        mt_training.build_pretraining_data_loader = lb_mini_builder
+    log_rank_0("[ODC.lb_mini] hooked build_pretraining_data_loader")
+
+
+def _install_schedule_patch():
+    """Patch forward_backward_no_pipelining to use THIS rank's micro-batch count.
+
+    At the top of every train_step the LB-Mini iterator plans one global
+    minibatch (KK balance across ranks); we read this rank's micro-batch count
+    and override the (globally-identical) ``num_microbatches`` argument so the
+    schedule's forward/backward loop runs the right rank-local number of steps.
+    """
+    global _FB_PATCHED
+    import megatron.core.pipeline_parallel.schedules as sched
+
+    if _FB_PATCHED or getattr(sched.forward_backward_no_pipelining, "_lb_mini_hooked", False):
+        return
+
+    orig_fb = sched.forward_backward_no_pipelining
+
+    def lb_mini_fb(*args, **kwargs):
+        it = _LB_MINI_TRAIN_ITER
+        forward_only = kwargs.get("forward_only", False)
+        # Only re-plan for the train path (an LB-Mini iterator exists) and when
+        # actually training (forward_only=False is the train_step path).
+        if it is not None and not forward_only:
+            try:
+                local_nmb = it.begin_minibatch()
+                if local_nmb > 0:
+                    kwargs["num_microbatches"] = local_nmb
+            except Exception as e:  # noqa: BLE001
+                warning_rank_0(
+                    f"[ODC.lb_mini] begin_minibatch failed, using stock "
+                    f"num_microbatches: {type(e).__name__}: {e}"
+                )
+        return orig_fb(*args, **kwargs)
+
+    lb_mini_fb._lb_mini_hooked = True
+    sched.forward_backward_no_pipelining = lb_mini_fb
+
+    # get_forward_backward_func returns the module-global by name; rebinding the
+    # module attribute is enough as long as it is fetched AFTER this patch. Also
+    # patch the function it returns defensively if it caches a reference.
+    _FB_PATCHED = True
+    log_rank_0("[ODC.lb_mini] hooked forward_backward_no_pipelining (rank-local num_microbatches)")
+
+
+@register_patch(
+    "megatron.fsdp.odc_lb_mini",
+    backend="megatron",
+    phase="before_train",
+    description="LB-Mini sequence-length load balancing for Megatron FSDP2 (data decoupled from ODC comm).",
+    condition=lambda ctx: _lb_mini_enabled(get_args(ctx)),
+)
+def patch_odc_lb_mini(ctx: PatchContext):
+    aligned = _lb_mini_aligned(get_args(ctx))
+    mode = "ALIGNED (nccl, no ODC)" if aligned else "DECOUPLED (ODC comm)"
+    log_rank_0(
+        f"[ODC.lb_mini] enable_odc_lb_mini=true, mode={mode} -> installing LB-Mini "
+        "(variable-length KK-balanced data + rank-local num_microbatches schedule)."
+    )
+    # Loss normalization under torch FSDP2: we KEEP calculate_per_token_loss at
+    # its default (False). Each micro-batch loss is mean-reduced (/=num_tokens)
+    # then /=num_microbatches (this rank's count) -- the same per-minibatch mean
+    # ODC's own example uses, and it keeps gradients at the right magnitude.
+    #
+    # We deliberately do NOT force calculate_per_token_loss=True. Under
+    # use_torch_fsdp2 the per-token grad rescale lives in Megatron's
+    # finalize_model_grads (scale by the GLOBAL all-reduced token count), but
+    # FSDP2 does its OWN reduce-scatter and BYPASSES that path, so per-token
+    # leaves the summed (un-normalized) loss and gradients explode ~1000x
+    # (measured: grad norm ~45000 vs ~55). KK balancing keeps per-rank
+    # micro-batch counts nearly equal (and in ALIGNED mode they are EXACTLY
+    # equal), so the residual per-minibatch-mean weighting difference (e.g. a rank
+    # with 3 vs 4 micro-batches) is negligible.
+    #
+    # Both patches are installed in either mode. The schedule patch is
+    # comm-agnostic (it only sets this rank's num_microbatches); in ALIGNED mode
+    # (enable_odc=false) the counts are all_reduce(MAX)-uniform across ranks, so
+    # the standard FSDP2 + RCCL collectives stay in lockstep -- no ODC required.
+    _install_dataloader_patch()
+    _install_schedule_patch()
+
+
+__all__ = ["patch_odc_lb_mini"]
diff --git a/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py b/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py
new file mode 100644
index 000000000..4e9e31951
--- /dev/null
+++ b/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py
@@ -0,0 +1,457 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+###############################################################################
+# ODC (On-Demand Communication) integration into Megatron's PyTorch-FSDP2 path.
+#
+# Gated by the enable_odc config item (so it is a no-op unless explicitly turned
+# on) AND requires use_torch_fsdp2=true.
+#
+# Strategy (see primus/core/odc/ROCM_ADAPTATION_REPORT.md for context):
+#   Megatron's TorchFullyShardedDataParallel wraps the model with the STANDARD
+#   PyTorch `fully_shard` API. ODC's odc/fsdp/fsdp2.py monkey-patches the same
+#   torch.distributed.fsdp._fully_shard internals (foreach_all_gather,
+#   FSDPParamGroup.post_backward, ...), so ODC's comm replacement applies
+#   transparently to Megatron's FSDP2 modules.
+#
+#   Integration points:
+#     patch_fsdp2()       -> BEFORE the first fully_shard call
+#     patch_lazy_init(m)  -> AFTER fully_shard, for every wrapped module
+#     pre_minibatch_start -> start of each train_step       (PHASE 2)
+#     pre_optimizer_step  -> before optimizer.step()        (PHASE 2)
+#
+# The odc_phase config item controls how far we wire in:
+#   1: only __init__ hook (patch_fsdp2 + patch_lazy_init).
+#      Verifies model construction + first forward/backward do not crash (i.e.
+#      ODC's symm-buffer replacement is compatible with Megatron's FSDP2 params).
+#      Gradients are NOT yet routed.
+#   2 (production default): also hook train_step / optimizer.step for full grad
+#      routing.
+###############################################################################
+
+import os
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+_ODC_READY = False
+_FSDP2_PATCHED = False
+
+
+def _ensure_odc_ready():
+    """Initialize MORI-SHMEM and apply ODC's module-level FSDP2 patches.
+
+    Must run BEFORE the first fully_shard() call. Idempotent.
+    """
+    global _ODC_READY, _FSDP2_PATCHED
+    import odc
+    from odc.fsdp import fsdp2 as odc_fsdp2
+
+    if not _ODC_READY:
+        os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G")
+        if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ:
+            os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=")
+        odc.init_shmem()
+        _ODC_READY = True
+        log_rank_0("[ODC.torch_fsdp2] init_shmem (MORI) done")
+
+    if not _FSDP2_PATCHED:
+        odc_fsdp2.patch_fsdp2(enable_hpz=False)
+        _FSDP2_PATCHED = True
+        log_rank_0("[ODC.torch_fsdp2] patch_fsdp2() applied (foreach_all_gather / post_backward replaced)")
+
+
+def _apply_patch_lazy_init(root_module):
+    """Call ODC patch_lazy_init on every fully_shard-ed submodule that owns a
+    parameter group.
+
+    A module is fully_shard-ed iff FSDP2 attached the ``_get_fsdp_state``
+    accessor. We additionally require ``state._fsdp_param_group is not None``:
+    container wrappers (e.g. the root) have no param group of their own, and ODC
+    itself skips them everywhere (replace_sharded_param_with_symm_buffer /
+    pre_minibatch_start / pre_optimizer_step all filter ``_fsdp_param_group is
+    None``). Patching such a module would just raise AttributeError, so we skip
+    it explicitly instead of catching the error.
+    """
+    from odc.fsdp import fsdp2 as odc_fsdp2
+
+    patched, skipped = 0, 0
+    for m in root_module.modules():
+        if not hasattr(m, "_get_fsdp_state"):
+            continue
+        state = m._get_fsdp_state()
+        if state is None or state._fsdp_param_group is None:
+            skipped += 1
+            continue
+        try:
+            odc_fsdp2.patch_lazy_init(m)
+            patched += 1
+        except Exception as e:  # noqa: BLE001
+            warning_rank_0(
+                f"[ODC.torch_fsdp2] patch_lazy_init failed on a param-group module: "
+                f"{type(e).__name__}: {e}"
+            )
+    log_rank_0(
+        f"[ODC.torch_fsdp2] patch_lazy_init applied to {patched} param-group modules "
+        f"({skipped} container module(s) without param group skipped)"
+    )
+
+
+def _populate_odc_runtime_config(ctx: PatchContext):
+    """Bridge the odc_* trainer-config items into the ODC library runtime config.
+
+    The ODC primitives (odc.primitives.*) are decoupled library code that reads its
+    tuning knobs from odc.runtime_config (populated here) rather than os.environ.
+    Only values explicitly set in the config are forwarded; unset items (None) keep
+    the library defaults, so behaviour is unchanged unless a knob is configured.
+    """
+    import odc
+
+    args = get_args(ctx)
+
+    def _g(name):
+        return getattr(args, name, None)
+
+    _defer = _g("odc_gda_defer_reduce")
+    odc.set_runtime_config(
+        p2p_backend=_g("odc_p2p_backend"),
+        mori_init=_g("odc_mori_init"),
+        max_buffer_size=_g("odc_max_buffer_size"),
+        rocshmem_gda=_g("odc_rocshmem_gda"),
+        rocshmem_lib=_g("odc_rocshmem_lib"),
+        gda_rs_blocks=_g("odc_gda_rs_blocks"),
+        gda_pipe=_g("odc_gda_pipe"),
+        gda_defer_reduce=(str(_defer) if _defer is not None else None),
+        gda_warmup_mode=_g("odc_gda_warmup_mode"),
+        gda_stride_bytes=_g("odc_gda_stride_bytes"),
+    )
+    log_rank_0(
+        "[ODC.torch_fsdp2] runtime config populated from trainer config "
+        f"(p2p_backend={odc.get_runtime_config().p2p_backend}, "
+        f"rocshmem_gda={odc.get_runtime_config().rocshmem_gda}, "
+        f"warmup_mode={odc.get_runtime_config().gda_warmup_mode})"
+    )
+
+
+@register_patch(
+    "megatron.fsdp.odc_torch_fsdp2",
+    backend="megatron",
+    phase="before_train",
+    description="Integrate ODC on-demand communication into Megatron's PyTorch FSDP2 path (ROCm/MORI).",
+    condition=lambda ctx: getattr(get_args(ctx), "enable_odc", False)
+    and getattr(get_args(ctx), "use_torch_fsdp2", False),
+)
+def patch_odc_torch_fsdp2(ctx: PatchContext):
+    # Populate the ODC runtime config from the trainer config BEFORE any ODC
+    # primitive is imported. `import odc` is cheap (odc/__init__ is lazy and does
+    # not pull in odc.primitives), so this runs before _ensure_odc_ready() imports
+    # odc.fsdp.fsdp2 -> odc.primitives, whose import-time backend selection
+    # (odc_p2p_backend) must read the populated config, not a stale default.
+    _populate_odc_runtime_config(ctx)
+
+    # PR #808 (feat(flux): FSDP2 optimizers + fp8 all-gather) replaced Megatron's FSDP2
+    # wrapper with PrimusTorchFullyShardedDataParallel (installed as
+    # megatron.training.training.torch_FSDP), so the stock TorchFullyShardedDataParallel
+    # is never instantiated. We must hook the class the trainer actually uses, otherwise
+    # _ensure_odc_ready()/reduction_service never initializes and pre_minibatch_start
+    # crashes with 'NoneType' object has no attribute 'clear_accumulations'.
+    try:
+        from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import (
+            PrimusTorchFullyShardedDataParallel as TorchFSDP,
+        )
+    except ImportError:
+        import megatron.core.distributed.torch_fully_sharded_data_parallel as tfsdp_mod
+
+        TorchFSDP = tfsdp_mod.TorchFullyShardedDataParallel
+    if getattr(TorchFSDP.__init__, "_odc_hooked", False):
+        log_rank_0("[ODC.torch_fsdp2] __init__ already hooked, skip")
+        return
+
+    orig_init = TorchFSDP.__init__
+    phase = str(getattr(get_args(ctx), "odc_phase", 2))
+
+    def odc_init(self, *args, **kwargs):
+        # 1) MORI init + ODC module-level FSDP2 patch, BEFORE any fully_shard.
+        _ensure_odc_ready()
+        # 2) Megatron's original __init__ runs all fully_shard() calls.
+        orig_init(self, *args, **kwargs)
+        # 3) AFTER fully_shard: install ODC's lazy_init hook (symm-buffer replace)
+        #    on every wrapped module.
+        _apply_patch_lazy_init(self.module)
+        # 4) [ODC x #808] ODC uses a serial single-stream rocSHMEM/GDA transport, so
+        #    #808's FSDP2 forward all-gather prefetch cannot overlap and is pure
+        #    overhead (+3.4~3.9s/step gather_kernel, +26.6GB max reserved on dual-node
+        #    14B). ODC is active here (this patch only runs when enable_odc=true), so we
+        #    opt every wrapped module out of the forward prefetch. Backward prefetch is
+        #    kept -- it is fine/needed under ODC.
+        _disable_forward_prefetch(self.module)
+        log_rank_0(f"[ODC.torch_fsdp2] TorchFSDP wrapped with ODC (odc_phase={phase})")
+
+    odc_init._odc_hooked = True
+    TorchFSDP.__init__ = odc_init
+    log_rank_0(f"[ODC.torch_fsdp2] hooked {TorchFSDP.__name__}.__init__")
+
+    if phase == "2":
+        _install_train_loop_hooks()
+
+
+def _disable_forward_prefetch(root_module):
+    """[ODC x #808] Opt ODC-managed modules out of #808's forward all-gather prefetch.
+
+    PR #808's PrimusTorchFullyShardedDataParallel calls set_modules_to_forward_prefetch()
+    on every fully_shard-ed inner module, overlapping the *next* layer's all-gather with
+    the current layer's forward. That is a win for FSDP2-native NCCL all-gather, but ODC's
+    transport is a serial single-stream rocSHMEM/GDA path (overlap_factor 1.00x): the
+    prefetched all-gather cannot overlap and only adds gather_kernel time and peak memory
+    (profiled: +3.4~3.9s/step, +26.6GB max reserved on dual-node 14B). Since this runs only
+    when ODC is enabled, we clear forward prefetch on every param-group module. FORWARD
+    only -- backward prefetch is left intact (fine/needed under ODC).
+    """
+    cleared, missing = 0, 0
+    for m in root_module.modules():
+        if not hasattr(m, "_get_fsdp_state"):
+            continue
+        state = m._get_fsdp_state()
+        if state is None or state._fsdp_param_group is None:
+            continue
+        if not hasattr(m, "set_modules_to_forward_prefetch"):
+            missing += 1
+            continue
+        try:
+            m.set_modules_to_forward_prefetch([])
+            cleared += 1
+        except Exception as e:  # noqa: BLE001
+            warning_rank_0(
+                f"[ODC.torch_fsdp2] set_modules_to_forward_prefetch([]) failed: " f"{type(e).__name__}: {e}"
+            )
+    log_rank_0(
+        f"[ODC.torch_fsdp2] skipped #808 forward all-gather prefetch on {cleared} module(s) "
+        f"(kept backward prefetch); ODC's serial transport cannot overlap it"
+        + (f"; {missing} module(s) lacked the API" if missing else "")
+    )
+
+
+def _find_gpt_model(root):
+    """Locate the GPTModel (owns output_layer + compute_language_model_loss)."""
+    m = getattr(root, "module", root)  # Float16Module -> GPTModel
+    if hasattr(m, "output_layer") and hasattr(m, "compute_language_model_loss"):
+        return m
+    for sub in root.modules():
+        if hasattr(sub, "output_layer") and hasattr(sub, "compute_language_model_loss"):
+            return sub
+    return None
+
+
+def _fused_ce_hooked_on(gpt):
+    pp = getattr(type(gpt), "_postprocess", None)
+    return bool(getattr(pp, "_fused_ce_hooked", False))
+
+
+def _odc_gather_output_weight(root):
+    """At the minibatch boundary (a SYNC point: all ranks enter train_step
+    together), all-gather the sharded output weight ONCE and cache the full
+    tensor as a grad-tracking leaf. fused CE then reuses it for every micro
+    batch WITHOUT per-micro collectives (which would deadlock under DiffMicro).
+    Only active when fused CE has patched GPTModel._postprocess.
+    """
+    gpt = _find_gpt_model(root)
+    if gpt is None or not _fused_ce_hooked_on(gpt):
+        return
+    try:
+        if getattr(gpt, "share_embeddings_and_output_weights", False):
+            ow = gpt.shared_embedding_or_output_weight()
+        else:
+            ow = gpt.output_layer.weight
+        if hasattr(ow, "full_tensor"):
+            gpt._odc_cached_output_weight = ow.full_tensor().detach().requires_grad_(True)
+    except Exception as e:  # noqa: BLE001
+        warning_rank_0(f"[ODC.fusedce] gather output weight failed: {type(e).__name__}: {e}")
+
+
+def _odc_reduce_output_grad(root):
+    """At pre_optimizer_step (minibatch end, also a SYNC point): all-reduce the
+    cached full-weight grad across DP (AVG), scatter it back to the sharded
+    param's .grad (DTensor), and drop the cache. One collective per minibatch
+    per rank -> DiffMicro-safe.
+    """
+    import torch.distributed as dist
+
+    gpt = _find_gpt_model(root)
+    if gpt is None:
+        return
+    cached = getattr(gpt, "_odc_cached_output_weight", None)
+    if cached is None:
+        return
+    try:
+        if cached.grad is not None:
+            full_grad = cached.grad
+            if getattr(gpt, "share_embeddings_and_output_weights", False):
+                ow = gpt.shared_embedding_or_output_weight()
+            else:
+                ow = gpt.output_layer.weight
+            dist.all_reduce(full_grad, op=dist.ReduceOp.AVG)
+            from torch.distributed.tensor import DTensor, distribute_tensor
+
+            if isinstance(ow, DTensor):
+                scattered = distribute_tensor(full_grad, ow.device_mesh, ow.placements)
+                ow.grad = scattered if ow.grad is None else (ow.grad + scattered)
+            elif ow.grad is None:
+                ow.grad = full_grad
+            else:
+                ow.grad = ow.grad + full_grad
+    except Exception as e:  # noqa: BLE001
+        warning_rank_0(f"[ODC.fusedce] reduce output grad failed: {type(e).__name__}: {e}")
+    finally:
+        gpt._odc_cached_output_weight = None
+
+
+def _odc_grad_spike_guard(root):
+    """Skip the optimizer step when the global grad norm spikes abnormally.
+
+    WHY: ODC's asynchronous P2P/MORI gradient reduction is numerically
+    NON-deterministic -- the same forward (identical loss) can occasionally
+    produce a huge grad spike (observed: same iter, loss 11.17512 vs 11.17508,
+    but grad norm 9.7 vs 442415). Most spikes are small and clip_grad absorbs
+    them, but an extreme one pushes params into a persistent divergent state
+    (clip only bounds magnitude, not the already-corrupted direction). This is
+    most visible under token-imbalanced SAME_MICRO (arm2); LB-Mini (token
+    balanced) rarely triggers it.
+
+    FIX: at the minibatch sync point (optimizer.step), compute the global grad
+    norm; if it exceeds the odc_grad_spike_threshold config (default 1000, <=0
+    disables), zero ALL grads so the ensuing step is a no-op -- i.e. SKIP this bad
+    iter, exactly like Megatron skips nan/inf iterations. Normal grad norm here is
+    ~3-42, recoverable spikes ~100-200, pathological ones ~1e5+, so 1000 cleanly
+    separates them.
+
+    The all_reduce below is safe under DiffMicro: optimizer.step runs once per
+    minibatch at a sync point where all ranks are aligned.
+    """
+    import torch
+    import torch.distributed as dist
+    from megatron.training import get_args as _mt_get_args
+
+    thr = float(getattr(_mt_get_args(), "odc_grad_spike_threshold", 1000.0))
+    if thr <= 0:
+        return
+    local_sq = None
+    for p in root.parameters():
+        g = getattr(p, "grad", None)
+        if g is None:
+            continue
+        if hasattr(g, "to_local"):  # DTensor -> this rank's shard
+            g = g.to_local()
+        s = g.detach().float().pow(2).sum()
+        local_sq = s if local_sq is None else (local_sq + s)
+    # [ODC-FIX] Every rank MUST enter the all_reduce below in lockstep. Under odc_nopad
+    # (variable micro-batch counts) a rank can legitimately finish a minibatch with no
+    # local grad (local_sq is None). The old early-return made that rank skip the
+    # collective -> ncclDevKernel_Generic_2 spins forever on the other ranks -> GPU
+    # deadlock at optimizer.step (confirmed via rocgdb + HIP trace). Contribute a 0 so the
+    # barrier stays aligned; only bail out when there is truly no process group.
+    if dist.is_initialized():
+        if local_sq is None:
+            local_sq = torch.zeros((), device=next(root.parameters()).device, dtype=torch.float32)
+        dist.all_reduce(local_sq, op=dist.ReduceOp.SUM)
+    if local_sq is None:
+        return
+    gnorm = local_sq.sqrt().item()
+    if gnorm > thr:
+        for p in root.parameters():
+            if getattr(p, "grad", None) is not None:
+                p.grad.zero_()
+        warning_rank_0(
+            f"[ODC.spike_guard] grad norm {gnorm:.1f} > {thr} -> SKIP step "
+            f"(grads zeroed; non-deterministic ODC async-reduce spike, params unchanged)"
+        )
+
+
+def _install_train_loop_hooks():
+    """PHASE 2: wire pre_minibatch_start / pre_optimizer_step into the loop."""
+    import megatron.training.training as mt_training
+
+    from odc.fsdp import fsdp2 as odc_fsdp2
+
+    if getattr(mt_training.train_step, "_odc_hooked", False):
+        return
+
+    orig_train_step = mt_training.train_step
+
+    def _find_fsdp_root(model):
+        # model is a list of model_chunks (TorchFSDP instances). The fully_shard
+        # root is chunk.module.
+        chunk = model[0] if isinstance(model, (list, tuple)) else model
+        root = getattr(chunk, "module", chunk)
+        return root
+
+    def odc_train_step(forward_step_func, data_iterator, model, optimizer, *a, **kw):
+        root = _find_fsdp_root(model)
+        # pre_minibatch_start: clear ODC accumulations at the start of the step.
+        try:
+            odc_fsdp2.pre_minibatch_start(root)
+        except Exception as e:  # noqa: BLE001
+            warning_rank_0(f"[ODC.torch_fsdp2] pre_minibatch_start failed: {type(e).__name__}: {e}")
+        # fused CE: all-gather full output weight ONCE per minibatch (sync point),
+        # so per-micro-batch fused CE never issues a DiffMicro-unsafe collective.
+        _odc_gather_output_weight(root)
+
+        # Hook optimizer.step once to inject pre_optimizer_step before it.
+        if not getattr(optimizer.step, "_odc_hooked", False):
+            orig_step = optimizer.step
+
+            def odc_step(*sa, **sk):
+                try:
+                    odc_fsdp2.pre_optimizer_step(root)
+                except Exception as e:  # noqa: BLE001
+                    warning_rank_0(f"[ODC.torch_fsdp2] pre_optimizer_step failed: {type(e).__name__}: {e}")
+                # fused CE: reduce-scatter cached full output-weight grad back to
+                # the sharded param (sync point, one collective per minibatch).
+                _odc_reduce_output_grad(root)
+                # guard against ODC async-reduce's occasional non-deterministic
+                # grad spike (can push params into a divergent state, esp. under
+                # token-imbalanced SAME_MICRO): skip the step if grad norm spikes.
+                _odc_grad_spike_guard(root)
+                return orig_step(*sa, **sk)
+
+            odc_step._odc_hooked = True
+            optimizer.step = odc_step
+            log_rank_0("[ODC.torch_fsdp2] hooked optimizer.step (pre_optimizer_step injected)")
+
+        return orig_train_step(forward_step_func, data_iterator, model, optimizer, *a, **kw)
+
+    odc_train_step._odc_hooked = True
+    mt_training.train_step = odc_train_step
+    log_rank_0("[ODC.torch_fsdp2] hooked train_step (pre_minibatch_start injected)")
+
+
+@register_patch(
+    "megatron.fsdp.odc_torch_fsdp2_teardown",
+    backend="megatron",
+    phase="after_train",
+    description="Tear down the ODC reduction service after training (ROCm/MORI).",
+    condition=lambda ctx: getattr(get_args(ctx), "enable_odc", False)
+    and getattr(get_args(ctx), "use_torch_fsdp2", False),
+)
+def patch_odc_torch_fsdp2_teardown(ctx: PatchContext):
+    """Tear down the ODC reduction service once training is done.
+
+    The device-side (single-node XGMI pull-sum) and GPU-direct (GDA) reduce
+    paths run no host-side subprocess, so ``ReductionService.stop()`` is a
+    no-op kept for API compatibility. This hook is retained as the single
+    teardown call site (harmless today) so Primus keeps a place to release ODC
+    resources; we deliberately skip finalize_distributed() /
+    SymmBufferRegistry.finalize() so we do not tear down the process group
+    before Primus' own cleanup runs.
+    """
+    from odc.fsdp import fsdp2 as odc_fsdp2
+
+    rs = odc_fsdp2.get_reduction_service()
+    if rs is None:
+        log_rank_0("[ODC.torch_fsdp2] teardown: no reduction_service, skip")
+        return
+    try:
+        rs.stop()
+        log_rank_0("[ODC.torch_fsdp2] reduction service torn down at after_train")
+    except Exception as e:  # noqa: BLE001
+        warning_rank_0(f"[ODC.torch_fsdp2] teardown stop() failed (non-fatal): " f"{type(e).__name__}: {e}")
diff --git a/primus/backends/megatron/patches/optimizer_patches.py b/primus/backends/megatron/patches/optimizer_patches.py
new file mode 100644
index 000000000..69e2958a9
--- /dev/null
+++ b/primus/backends/megatron/patches/optimizer_patches.py
@@ -0,0 +1,251 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron Optimizer Patches
+
+This module contains patches that modify Megatron's optimizer creation to use
+Primus-specific implementations when requested.
+"""
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.optimizer.fsdp2_fp32_param",
+    backend="megatron",
+    phase="before_train",
+    description="Patch get_megatron_optimizer for FSDP2 FP32 param optimizer (TorchTitan-style)",
+    priority=50,
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_fsdp2_fp32_param_optimizer", False)
+        and getattr(get_args(ctx), "use_torch_fsdp2", False)
+        and getattr(get_args(ctx), "bf16", False)
+    ),
+)
+def patch_fsdp2_fp32_optimizer(ctx: PatchContext):
+    """Patch Megatron to use FSDP2 FP32 optimizer.
+
+    This optimizer uses FP32 parameters + FP32 optimizer states with FSDP2's
+    MixedPrecisionPolicy for BF16 forward/backward. Eliminates the stale
+    weights problem of BF16 optimizer states.
+    """
+    try:
+        import megatron.core.optimizer as optimizer_module
+        from megatron.training import training
+
+        from primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer import (
+            get_fsdp2_fp32_optimizer,
+        )
+
+        args = get_args(ctx)
+        use_foreach = getattr(args, "optimizer_foreach", False)
+
+        _MEGATRON_ONLY_KWARGS = {
+            "config_overrides",
+            "use_gloo_process_groups",
+            "dump_param_to_param_group_map",
+        }
+
+        def patched_get_megatron_optimizer(config, model_chunks, *args, **kwargs):
+            log_rank_0("=" * 80)
+            log_rank_0("[Using FSDP2 FP32 Param Optimizer (TorchTitan-style)]")
+            log_rank_0("  FP32 params + FP32 optimizer states")
+            log_rank_0("  FSDP2 MixedPrecisionPolicy handles BF16 compute")
+            log_rank_0("  DTensor-native gradient clipping")
+            log_rank_0(f"  AdamW mode: {'foreach' if use_foreach else 'fused'}")
+            log_rank_0("=" * 80)
+
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k not in _MEGATRON_ONLY_KWARGS}
+
+            return get_fsdp2_fp32_optimizer(
+                config=config,
+                model_chunks=model_chunks,
+                use_foreach=use_foreach,
+                **filtered_kwargs,
+            )
+
+        patched_count = 0
+        if hasattr(training, "get_megatron_optimizer"):
+            training.get_megatron_optimizer = patched_get_megatron_optimizer
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_fp32_param] " "Patched training.get_megatron_optimizer"
+            )
+            patched_count += 1
+
+        if hasattr(optimizer_module, "get_megatron_optimizer"):
+            optimizer_module.get_megatron_optimizer = patched_get_megatron_optimizer
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_fp32_param] "
+                "Patched optimizer_module.get_megatron_optimizer"
+            )
+            patched_count += 1
+
+        if patched_count == 0:
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_fp32_param] "
+                "WARNING: get_megatron_optimizer not found in either location!"
+            )
+        else:
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_fp32_param] "
+                f"Patched get_megatron_optimizer in {patched_count} location(s) "
+                "to use FSDP2FP32Optimizer"
+            )
+
+    except Exception as e:
+        log_rank_0(
+            f"[Patch:megatron.optimizer.fsdp2_fp32_param] "
+            f"WARNING: Failed to patch get_megatron_optimizer: {type(e).__name__}: {e}"
+        )
+        import traceback
+
+        log_rank_0(f"Traceback: {traceback.format_exc()}")
+
+
+@register_patch(
+    "megatron.optimizer.fsdp2_bf16_master_weight",
+    backend="megatron",
+    phase="before_train",
+    description="Patch get_megatron_optimizer for FSDP2 BF16 master weight optimizer",
+    priority=50,
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_fsdp2_bf16_master_weight_optimizer", False)
+        and getattr(get_args(ctx), "use_torch_fsdp2", False)
+        and getattr(get_args(ctx), "bf16", False)
+    ),
+)
+def patch_fsdp2_bf16_master_weight_optimizer(ctx: PatchContext):
+    """Patch Megatron to use FSDP2 BF16 master weight optimizer.
+
+    This optimizer keeps model parameters in BF16 (no FP32->BF16 cast in
+    forward all-gather) and maintains FP32 master copies for optimizer
+    precision, matching Megatron's Float16OptimizerWithFloat16Params pattern.
+    """
+    try:
+        import megatron.core.optimizer as optimizer_module
+        from megatron.training import training
+
+        from primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer import (
+            get_fsdp2_bf16_master_weight_optimizer,
+        )
+
+        args = get_args(ctx)
+        use_foreach = getattr(args, "optimizer_foreach", True)
+
+        _MEGATRON_ONLY_KWARGS = {
+            "config_overrides",
+            "use_gloo_process_groups",
+            "dump_param_to_param_group_map",
+        }
+
+        def patched_get_megatron_optimizer(config, model_chunks, *args, **kwargs):
+            log_rank_0("=" * 80)
+            log_rank_0("[Using FSDP2 BF16 Master Weight Optimizer]")
+            log_rank_0("  BF16 model params + FP32 master weights")
+            log_rank_0("  No FP32->BF16 cast in forward all-gather")
+            log_rank_0("  DTensor-native gradient clipping")
+            log_rank_0(f"  Foreach batching: {use_foreach}")
+            log_rank_0("=" * 80)
+
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k not in _MEGATRON_ONLY_KWARGS}
+
+            return get_fsdp2_bf16_master_weight_optimizer(
+                config=config,
+                model_chunks=model_chunks,
+                use_foreach=use_foreach,
+                **filtered_kwargs,
+            )
+
+        patched_count = 0
+        if hasattr(training, "get_megatron_optimizer"):
+            training.get_megatron_optimizer = patched_get_megatron_optimizer
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] "
+                "Patched training.get_megatron_optimizer"
+            )
+            patched_count += 1
+
+        if hasattr(optimizer_module, "get_megatron_optimizer"):
+            optimizer_module.get_megatron_optimizer = patched_get_megatron_optimizer
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] "
+                "Patched optimizer_module.get_megatron_optimizer"
+            )
+            patched_count += 1
+
+        if patched_count == 0:
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] "
+                "WARNING: get_megatron_optimizer not found in either location!"
+            )
+        else:
+            log_rank_0(
+                "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] "
+                f"Patched get_megatron_optimizer in {patched_count} location(s) "
+                "to use FSDP2BF16MasterWeightOptimizer"
+            )
+
+    except Exception as e:
+        log_rank_0(
+            f"[Patch:megatron.optimizer.fsdp2_bf16_master_weight] "
+            f"WARNING: Failed to patch get_megatron_optimizer: {type(e).__name__}: {e}"
+        )
+        import traceback
+
+        log_rank_0(f"Traceback: {traceback.format_exc()}")
+
+
+@register_patch(
+    "megatron.optimizer.precision_aware_fp8_tensorwise",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Override precision-aware optimizer flag for local spec + tensorwise FP8, "
+        "enabling BF16 decoupled_grad path (no BF16-to-FP32 gradient cast)."
+    ),
+    priority=35,
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "use_precision_aware_optimizer", False)
+        and getattr(get_args(ctx), "fp8_recipe", None) == "tensorwise"
+        and getattr(get_args(ctx), "transformer_impl", None) == "local"
+    ),
+)
+def patch_precision_aware_for_tensorwise(ctx: PatchContext):
+    """Enable precision-aware optimizer's decoupled_grad path for local spec + tensorwise FP8.
+
+    With local spec, parameters are BF16 (not Float8Tensor) -- FP8 quantization
+    happens inside per-module autograd Functions, transparent to the optimizer.
+    The upstream condition incorrectly excludes tensorwise FP8 because it was
+    designed for TE spec where Float8Tensor storage requires special handling.
+    """
+    try:
+        from megatron.core.optimizer.optimizer_config import OptimizerConfig
+
+        original_post_init = OptimizerConfig.__post_init__
+
+        def patched_post_init(self):
+            original_post_init(self)
+            if self.use_precision_aware_optimizer and self.fp8_recipe == "tensorwise":
+                self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = True
+
+        OptimizerConfig.__post_init__ = patched_post_init
+
+        log_rank_0(
+            "[Patch:megatron.optimizer.precision_aware_fp8_tensorwise] "
+            "Patched OptimizerConfig.__post_init__ to enable "
+            "decoupled_grad path for local spec + tensorwise FP8"
+        )
+
+    except Exception as e:
+        log_rank_0(
+            f"[Patch:megatron.optimizer.precision_aware_fp8_tensorwise] "
+            f"WARNING: Failed to patch OptimizerConfig: {type(e).__name__}: {e}"
+        )
+        import traceback
+
+        log_rank_0(f"Traceback: {traceback.format_exc()}")
diff --git a/primus/backends/megatron/patches/parallelism/forward_step_patches.py b/primus/backends/megatron/patches/parallelism/forward_step_patches.py
index ec83eb9b0..dddf9c033 100644
--- a/primus/backends/megatron/patches/parallelism/forward_step_patches.py
+++ b/primus/backends/megatron/patches/parallelism/forward_step_patches.py
@@ -17,7 +17,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _needs_forward_step_patch(ctx: PatchContext) -> bool:
@@ -49,7 +49,7 @@ def _patched_forward_step(data_iterator, model, return_schedule_plan=False):
         if not args.patch_zero_bubble:
             return _original_forward_step(data_iterator, model, return_schedule_plan)
 
-        from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore
+        from primus.backends.megatron.data_loader_store import DataLoaderStore
 
         timers = None
         try:
diff --git a/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py b/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py
index a65a9cfcb..92e74ac72 100644
--- a/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py
+++ b/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py
@@ -9,7 +9,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py b/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py
index 65ee1695e..f38dde4aa 100644
--- a/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py
+++ b/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py
@@ -9,7 +9,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/parallelism/schedule_patches.py b/primus/backends/megatron/patches/parallelism/schedule_patches.py
index 138ae4a1d..d0b71bbbd 100644
--- a/primus/backends/megatron/patches/parallelism/schedule_patches.py
+++ b/primus/backends/megatron/patches/parallelism/schedule_patches.py
@@ -9,7 +9,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py b/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py
new file mode 100644
index 000000000..f0aabc446
--- /dev/null
+++ b/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py
@@ -0,0 +1,218 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+SDMA (copy-engine) distributed-optimizer param all-gather patch.
+
+Migrated from the source patch ``megatron_sdma_allgather.patch`` (mpo branch).
+
+This replaces the in-place edits the source patch made to
+``third_party/Megatron-LM`` with three runtime monkey-patches, all gated by
+``ENABLE_SDMA_ALLGATHER=1``:
+
+  1. ``_ParamAndGradBucketGroup.start_param_sync`` -- the distributed-optimizer
+     path is re-implemented to dispatch one all-gather per bucket through
+     :func:`all_gather_into_tensor_sdma` (copy-engine) instead of the RCCL
+     ``_coalescing_manager`` group. By default, the first two bucket groups (by
+     gather order) stay on the regular RCCL fallback; this is tunable with
+     ``MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS``. The layer-wise optimizer path is
+     delegated unchanged to the original method.
+  2. ``DistributedDataParallel.__init__`` -- after construction, each bucket
+     group is annotated with ``param_gather_order`` (reverse dispatch order,
+     mirroring the source patch).
+  3. MoE experts ``forward`` -- ``tokens_per_expert`` is moved onto the
+     dispatched-input device before the experts run (the source patch's
+     ``moe_layer.py`` one-liner).
+
+When the SDMA primitives (Primus-Turbo symmetric memory + ``hip``) are
+unavailable, :func:`all_gather_into_tensor_sdma` falls back to
+``torch.distributed.all_gather_into_tensor``, so enabling the flag is safe even
+on images without the kernels.
+"""
+
+import os
+import warnings
+
+import torch
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+
+def _sdma_allgather_enabled(_ctx: PatchContext) -> bool:
+    return os.environ.get("ENABLE_SDMA_ALLGATHER", "0") == "1"
+
+
+def _get_rccl_fallback_bucket_count() -> int:
+    value = os.getenv("MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS", "2")
+    try:
+        count = int(value)
+    except ValueError:
+        warnings.warn(f"Invalid MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS={value!r}; using 2.")
+        return 2
+    if count < 0:
+        warnings.warn(f"MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS must be non-negative; got {count}, using 2.")
+        return 2
+    return count
+
+
+def _make_start_param_sync(orig_start_param_sync):
+    """Build a replacement ``start_param_sync`` for ``_ParamAndGradBucketGroup``."""
+    from megatron.core.distributed.param_and_grad_buffer import shard_buffer
+
+    from primus.backends.megatron.core.distributed.sdma_param_gather import (
+        _all_gather_into_tensor_waitable_fallback,
+        _WaitableHandle,
+        all_gather_into_tensor_sdma,
+    )
+
+    def start_param_sync(self, force_sync: bool = False):
+        # Layer-wise optimizer path is unchanged; only the distributed-optimizer
+        # path routes through SDMA.
+        if not self.ddp_config.use_distributed_optimizer:
+            return orig_start_param_sync(self, force_sync=force_sync)
+
+        if force_sync:
+            if self.param_gather_handle is not None:
+                self.param_gather_handle.wait()
+                self.param_gather_handle = None
+                return
+        else:
+            assert self.param_gather_handle is None
+
+        async_op = self.ddp_config.overlap_param_gather and not force_sync
+
+        # Keep a configurable number of leading bucket groups on regular RCCL;
+        # route the rest through SDMA. param_gather_order is assigned in the
+        # DDP __init__ wrapper below.
+        param_gather_order = getattr(self, "param_gather_order", None)
+        enable_sdma = os.getenv("ENABLE_SDMA_ALLGATHER") == "1"
+        rccl_fallback_bucket_count = _get_rccl_fallback_bucket_count()
+        use_rccl_fallback = not enable_sdma or (
+            param_gather_order is not None and param_gather_order < rccl_fallback_bucket_count
+        )
+        all_gather_func = (
+            _all_gather_into_tensor_waitable_fallback if use_rccl_fallback else all_gather_into_tensor_sdma
+        )
+
+        param_gather_handles = []
+        for idx, bucket in enumerate(self.buckets):
+            if self.cached_param_buffer_shard_list[idx] is None:
+                self.cached_param_buffer_shard_list[idx] = shard_buffer(
+                    bucket.param_data, self.intra_distributed_optimizer_instance_size
+                )
+            local_data_view = self.cached_param_buffer_shard_list[idx][
+                self.intra_distributed_optimizer_instance_rank
+            ]
+            handle = all_gather_func(
+                bucket.param_data,
+                local_data_view,
+                group=self.intra_distributed_optimizer_instance_group,
+                async_op=async_op,
+            )
+            if async_op:
+                param_gather_handles.append(handle)
+
+        if async_op:
+
+            def _wait_all_param_gathers():
+                for handle in param_gather_handles:
+                    handle.wait()
+
+            self.param_gather_handle = _WaitableHandle(wait_fn=_wait_all_param_gathers)
+        else:
+            self.param_gather_handle = None
+        self.param_gather_dispatched = True
+
+    return start_param_sync
+
+
+def _make_wrapped_ddp_init(orig_init):
+    def __init__(self, *args, **kwargs):
+        orig_init(self, *args, **kwargs)
+        # Mirror the source patch: number bucket groups in reverse dispatch
+        # order so start_param_sync can keep the configured leading groups on
+        # RCCL.
+        for groups_attr in ("bucket_groups", "expert_parallel_bucket_groups"):
+            groups = getattr(self, groups_attr, None) or []
+            for order, bucket_group in enumerate(reversed(groups)):
+                bucket_group.param_gather_order = order
+
+    return __init__
+
+
+def _make_wrapped_experts_forward(orig_forward):
+    def forward(self, permuted_local_hidden_states, tokens_per_expert, *args, **kwargs):
+        # Source patch moved tokens_per_expert onto the dispatched-input device
+        # before the experts run.
+        if torch.is_tensor(tokens_per_expert):
+            tokens_per_expert = tokens_per_expert.to(permuted_local_hidden_states.device)
+        return orig_forward(self, permuted_local_hidden_states, tokens_per_expert, *args, **kwargs)
+
+    return forward
+
+
+@register_patch(
+    "megatron.distributed.sdma_param_all_gather",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Route the distributed-optimizer param all-gather through Primus-Turbo "
+        "SDMA (copy-engine) memcpys; gated by ENABLE_SDMA_ALLGATHER=1."
+    ),
+    condition=_sdma_allgather_enabled,
+)
+def patch_sdma_param_all_gather(ctx: PatchContext):
+    del ctx
+
+    try:
+        import megatron.core.distributed.param_and_grad_buffer as pgb
+        from megatron.core.distributed.distributed_data_parallel import (
+            DistributedDataParallel,
+        )
+    except ImportError as exc:
+        warning_rank_0(
+            f"[Patch:megatron.distributed.sdma_param_all_gather] Megatron distributed "
+            f"modules not importable; skipping: {exc}"
+        )
+        return
+
+    bucket_group_cls = getattr(pgb, "_ParamAndGradBucketGroup", None)
+    if bucket_group_cls is None:
+        warning_rank_0(
+            "[Patch:megatron.distributed.sdma_param_all_gather] "
+            "_ParamAndGradBucketGroup not found; skipping."
+        )
+        return
+
+    if not getattr(bucket_group_cls, "_primus_sdma_param_gather_patched", False):
+        bucket_group_cls.start_param_sync = _make_start_param_sync(bucket_group_cls.start_param_sync)
+        bucket_group_cls._primus_sdma_param_gather_patched = True
+
+    if not getattr(DistributedDataParallel, "_primus_sdma_param_gather_patched", False):
+        DistributedDataParallel.__init__ = _make_wrapped_ddp_init(DistributedDataParallel.__init__)
+        DistributedDataParallel._primus_sdma_param_gather_patched = True
+
+    # MoE experts: move tokens_per_expert onto the dispatched-input device.
+    try:
+        from megatron.core.transformer.moe import experts as experts_mod
+
+        for cls_name in ("GroupedMLP", "TEGroupedMLP", "SequentialMLP"):
+            cls = getattr(experts_mod, cls_name, None)
+            if cls is None or getattr(cls, "_primus_sdma_tokens_to_device", False):
+                continue
+            cls.forward = _make_wrapped_experts_forward(cls.forward)
+            cls._primus_sdma_tokens_to_device = True
+    except ImportError as exc:
+        warning_rank_0(
+            f"[Patch:megatron.distributed.sdma_param_all_gather] MoE experts module "
+            f"not importable; skipping tokens_per_expert device fixup: {exc}"
+        )
+
+    log_rank_0(
+        "[Patch:megatron.distributed.sdma_param_all_gather] Installed SDMA param "
+        "all-gather (start_param_sync + param_gather_order + experts tokens_per_expert)."
+    )
diff --git a/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py b/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py
index 74d49f6ca..e3bfde632 100644
--- a/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py
+++ b/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py
@@ -17,7 +17,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/parallelism/train_step_patches.py b/primus/backends/megatron/patches/parallelism/train_step_patches.py
index 9b539adaf..8916ed9c7 100644
--- a/primus/backends/megatron/patches/parallelism/train_step_patches.py
+++ b/primus/backends/megatron/patches/parallelism/train_step_patches.py
@@ -14,7 +14,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _needs_seq_split_adjustment(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/parallelism/v_schedule_patches.py b/primus/backends/megatron/patches/parallelism/v_schedule_patches.py
index fece7a507..73a0604a1 100644
--- a/primus/backends/megatron/patches/parallelism/v_schedule_patches.py
+++ b/primus/backends/megatron/patches/parallelism/v_schedule_patches.py
@@ -9,7 +9,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_v_schedule_enabled(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py b/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py
index 4c1bbde29..d2786f5cb 100644
--- a/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py
+++ b/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py
@@ -9,7 +9,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/pp_dump_data_patches.py b/primus/backends/megatron/patches/pp_dump_data_patches.py
index 6e4816deb..72c939c8f 100644
--- a/primus/backends/megatron/patches/pp_dump_data_patches.py
+++ b/primus/backends/megatron/patches/pp_dump_data_patches.py
@@ -21,7 +21,7 @@
 import os
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _make_wrapped_get_forward_backward_func(original_get_forward_backward_func):
@@ -30,12 +30,12 @@ def _make_wrapped_get_forward_backward_func(original_get_forward_backward_func):
     def wrapped_get_forward_backward_func():
         from megatron.training import get_args as get_megatron_args
 
-        from primus.modules.trainer.megatron.utils import schedule_wrapper
+        import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils
 
         func = original_get_forward_backward_func()
         args = get_megatron_args()
         if getattr(args, "dump_pp_data", False):
-            func = schedule_wrapper(func)
+            func = utils.schedule_wrapper(func)
             func._pp_dump_schedule_wrapped = True  # noqa: B010
         return func
 
@@ -81,7 +81,7 @@ def patch_pp_dump_data_before_train(ctx: PatchContext):
     import megatron.core.pipeline_parallel as pp_module
     import megatron.training.training as training_module
 
-    import primus.modules.trainer.megatron.utils as utils
+    import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils
 
     args = get_args(ctx)
     is_primus_pipeline = getattr(args, "patch_primus_pipeline", False)
@@ -131,11 +131,11 @@ def patch_pp_dump_data_after_train(ctx: PatchContext):
         from megatron.core.num_microbatches_calculator import get_num_microbatches
         from megatron.training import get_args as get_megatron_args
 
-        from primus.modules.trainer.megatron.utils import dump_pp_data
+        import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils
 
         args = get_megatron_args()
         pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data")
-        dump_pp_data(args, get_num_microbatches(), pp_data_dir)
+        utils.dump_pp_data(args, get_num_microbatches(), pp_data_dir)
         log_rank_0(f"[Patch:megatron.pp.dump_pp_data]   pp schedule data dumped to {pp_data_dir}")
     except Exception as e:
         log_rank_0(f"[Patch:megatron.pp.dump_pp_data]   WARNING: failed to dump pp data: {e}")
diff --git a/primus/backends/megatron/patches/pp_warmup_patches.py b/primus/backends/megatron/patches/pp_warmup_patches.py
index 6deb8d703..9ec893291 100644
--- a/primus/backends/megatron/patches/pp_warmup_patches.py
+++ b/primus/backends/megatron/patches/pp_warmup_patches.py
@@ -15,7 +15,7 @@
 import torch
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def run_pp_warmup(forward_step_func, model, optimizer, config):
@@ -125,9 +125,16 @@ def run_pp_warmup(forward_step_func, model, optimizer, config):
     device = torch.cuda.current_device()
     mbs = args.micro_batch_size
     seqlen = args.seq_length
+    # Tokens are drawn at random rather than zeroed. An all-zero batch makes every
+    # token's hidden state identical, so a top-k MoE router scores every expert the
+    # same and top-k degenerates to experts 0..k-1 for the whole batch -- every token
+    # lands on one EP rank. That routing collapse never occurs in real training, and
+    # it hangs the fused MegaMoE combine kernel. RNG state is restored at the end of
+    # warmup, so drawing here keeps training bit-identical.
+    vocab_size = getattr(args, "padded_vocab_size", None) or getattr(args, "vocab_size", 1) or 1
     fake_batch = {
-        "tokens": torch.zeros(mbs, seqlen, dtype=torch.int64, device=device),
-        "labels": torch.zeros(mbs, seqlen, dtype=torch.int64, device=device),
+        "tokens": torch.randint(0, vocab_size, (mbs, seqlen), dtype=torch.int64, device=device),
+        "labels": torch.randint(0, vocab_size, (mbs, seqlen), dtype=torch.int64, device=device),
         "loss_mask": torch.ones(mbs, seqlen, dtype=torch.float32, device=device),
         "position_ids": (
             torch.arange(seqlen, dtype=torch.int64, device=device)
@@ -180,8 +187,10 @@ def _make_fake_iter():
         if chunk_pre_process:
             input_tensor = None
         else:
+            # Random, not zeros: a zero activation collapses MoE top-k routing onto
+            # experts 0..k-1 for every token (see the fake_batch note above).
             input_tensor = [
-                torch.zeros(shape, dtype=activation_dtype, device=device, requires_grad=True)
+                torch.randn(shape, dtype=activation_dtype, device=device).requires_grad_(True)
                 for shape in tensor_shapes
             ]
 
diff --git a/primus/backends/megatron/patches/recompute_layer_patches.py b/primus/backends/megatron/patches/recompute_layer_patches.py
index c63d4008c..e79677395 100644
--- a/primus/backends/megatron/patches/recompute_layer_patches.py
+++ b/primus/backends/megatron/patches/recompute_layer_patches.py
@@ -38,7 +38,7 @@
 from contextlib import nullcontext
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def validate_specified_recompute_layers(config, args):
diff --git a/primus/backends/megatron/patches/runtime_hooks_patches.py b/primus/backends/megatron/patches/runtime_hooks_patches.py
index 94bb3ac53..6d3dc5307 100644
--- a/primus/backends/megatron/patches/runtime_hooks_patches.py
+++ b/primus/backends/megatron/patches/runtime_hooks_patches.py
@@ -12,7 +12,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py
index 6d8593768..8f6bf5708 100644
--- a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py
+++ b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py
@@ -41,7 +41,7 @@
 import os
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 def _sdma_all_gather_enabled(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/sft_grad_sanitize_patches.py b/primus/backends/megatron/patches/sft_grad_sanitize_patches.py
index 59f847e61..896e9574f 100644
--- a/primus/backends/megatron/patches/sft_grad_sanitize_patches.py
+++ b/primus/backends/megatron/patches/sft_grad_sanitize_patches.py
@@ -36,7 +36,7 @@
 import torch
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 _HAS_LOGGED_FIRST_SANITIZE = False
 
diff --git a/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py
new file mode 100644
index 000000000..1f6178412
--- /dev/null
+++ b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py
@@ -0,0 +1,145 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Transformer Engine attention BSHD-layout patch.
+
+Migrated from the source patch ``megatron_te_bshd_layout.patch`` (mpo branch).
+
+Some TE FMHA kernels are faster with a ``bshd`` memory layout than with the
+default ``sbhd`` layout Megatron feeds them. When ``NVTE_FMHA_USE_BSHD=1`` is
+set, this patch transposes the ``sbhd`` ``query``/``key``/``value`` tensors to
+``bshd`` before the TE ``DotProductAttention`` call and transposes the result
+back to ``sbhd`` afterwards, so the change is transparent to the rest of
+Megatron.
+
+Instead of editing ``third_party/Megatron-LM`` in place, we wrap
+``TEDotProductAttention.forward`` and temporarily flip the effective
+``qkv_format`` for the duration of the wrapped call so TE interprets the
+already-transposed tensors correctly.
+
+Gate:
+    ``NVTE_FMHA_USE_BSHD=1`` (off by default; the patch is not installed at
+    all when unset, so there is zero overhead).
+"""
+
+import functools
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+
+def _bshd_enabled(_ctx: PatchContext) -> bool:
+    return os.environ.get("NVTE_FMHA_USE_BSHD", "0") == "1"
+
+
+def _effective_qkv_format(self, packed_seq_params):
+    """Mirror TEDotProductAttention.forward's qkv_format resolution."""
+    if packed_seq_params is not None:
+        fmt = getattr(packed_seq_params, "qkv_format", None)
+        if fmt is not None:
+            return fmt
+    return self.qkv_format
+
+
+def _make_wrapped_forward(orig_forward):
+    @functools.wraps(orig_forward)
+    def forward(
+        self,
+        query,
+        key,
+        value,
+        attention_mask,
+        attn_mask_type,
+        attention_bias=None,
+        packed_seq_params=None,
+        **kwargs,
+    ):
+        # Only convert when the effective layout is sbhd; otherwise behave
+        # exactly like the original.
+        if _effective_qkv_format(self, packed_seq_params) != "sbhd":
+            return orig_forward(
+                self,
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type,
+                attention_bias=attention_bias,
+                packed_seq_params=packed_seq_params,
+                **kwargs,
+            )
+
+        query = query.transpose(0, 1).contiguous()
+        key = key.transpose(0, 1).contiguous()
+        value = value.transpose(0, 1).contiguous()
+
+        # Temporarily flip the qkv_format TE will read so the transposed
+        # tensors are interpreted as bshd. Restore it afterwards even on error.
+        if packed_seq_params is not None and hasattr(packed_seq_params, "qkv_format"):
+            restore_target, restore_value = packed_seq_params, packed_seq_params.qkv_format
+            packed_seq_params.qkv_format = "bshd"
+        else:
+            restore_target, restore_value = self, self.qkv_format
+            self.qkv_format = "bshd"
+
+        try:
+            core_attn_out = orig_forward(
+                self,
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type,
+                attention_bias=attention_bias,
+                packed_seq_params=packed_seq_params,
+                **kwargs,
+            )
+        finally:
+            restore_target.qkv_format = restore_value
+
+        return core_attn_out.transpose(0, 1).contiguous()
+
+    return forward
+
+
+@register_patch(
+    "megatron.te.fmha_bshd_layout",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Run TE DotProductAttention in bshd layout (transpose sbhd<->bshd "
+        "around the kernel) when NVTE_FMHA_USE_BSHD=1."
+    ),
+    condition=_bshd_enabled,
+)
+def patch_te_fmha_bshd_layout(ctx: PatchContext):
+    del ctx
+
+    try:
+        from megatron.core.extensions import transformer_engine as te_ext
+    except ImportError as exc:
+        warning_rank_0(
+            f"[Patch:megatron.te.fmha_bshd_layout] transformer_engine extension "
+            f"not importable; skipping: {exc}"
+        )
+        return
+
+    cls = getattr(te_ext, "TEDotProductAttention", None)
+    if cls is None:
+        warning_rank_0("[Patch:megatron.te.fmha_bshd_layout] TEDotProductAttention not found; skipping.")
+        return
+
+    if getattr(cls, "_primus_bshd_patched", False):
+        return
+
+    cls.forward = _make_wrapped_forward(cls.forward)
+    cls._primus_bshd_patched = True
+    log_rank_0(
+        "[Patch:megatron.te.fmha_bshd_layout] Patched TEDotProductAttention.forward "
+        "to use bshd layout (NVTE_FMHA_USE_BSHD=1)."
+    )
diff --git a/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py b/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py
index 0a99f5a22..bda764962 100644
--- a/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py
+++ b/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py
@@ -15,7 +15,7 @@
     make_get_extra_te_kwargs_with_override,
 )
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py b/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py
new file mode 100644
index 000000000..168457a0a
--- /dev/null
+++ b/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py
@@ -0,0 +1,305 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+DPA Consolidated Prologue Patch
+
+Separates the eager setup (FP8 metadata, backend selection, cu_seqlens, etc.)
+from the FusedAttnFunc kernel call in DotProductAttention, so that the kernel
+call can live inside a torch.compile compiled graph rather than running eagerly.
+
+Graph-break analysis (per DPA call):
+
+    Before:
+        [compiled A] --break--> [DPA runs fully eager, kernel included] --break--> [compiled B]
+
+    After (attention_dropout == 0, the Flux default - fast path):
+        [compiled A] --break--> [eager prologue] --break--> [compiled post-ops (kernel inside)]
+
+    After (attention_dropout > 0 - safe path):
+        [compiled A] --break--> [eager prologue] --break--> [eager fork + kernel] --break--> [compiled post-ops]
+
+Break counts: 2 for the fast path, 3 for the safe path.  The safe path runs
+the FusedAttn kernel inside the model-parallel-rng fork context so the
+kernel's philox state targets the correct generator.  In the fast path the
+kernel doesn't consume RNG (dropout is 0), so the fork wrapper is omitted
+and the kernel stays inside the compiled graph alongside the post-attention
+ops (output reshape, projection, skip connections).
+
+Mechanism:
+    FusedAttnFunc.forward is monkey-patched with a thread-local "capture mode".
+    When capture mode is active the patched forward stores its arguments in
+    thread-local storage and raises ``_CaptureComplete``, aborting the DPA
+    forward after all setup is done.  The captured args are then passed to the
+    *real* ``FusedAttnFunc.apply`` inside the compiled graph.
+"""
+
+import threading
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+# Positional index of dropout_p inside FusedAttnFunc.forward's *args (i.e.
+# inside our captured _capture_tls.captured_args tuple, which excludes ctx).
+# Verified against transformer_engine.pytorch.attention.dot_product_attention
+# .backends.FusedAttnFunc.forward as of TE 2.12. If a future TE upgrade
+# reshuffles the positional layout, _new_te_dpa_forward falls back to the
+# always-fork (safe) path rather than silently routing to the dropout=0 fast
+# path with the wrong field.
+_DROPOUT_P_IDX = 14
+
+# ---------------------------------------------------------------------------
+# Thread-local capture state
+# ---------------------------------------------------------------------------
+_capture_tls = threading.local()
+
+
+class _CaptureComplete(Exception):
+    """Raised inside the patched FusedAttnFunc.forward to abort after capture."""
+
+
+# ---------------------------------------------------------------------------
+# Patch registration
+# ---------------------------------------------------------------------------
+@register_patch(
+    "megatron.te.dpa_consolidated_prologue",
+    backend="megatron",
+    phase="before_train",
+    priority=55,
+    description=(
+        "Consolidated DPA prologue: separates eager attention setup from "
+        "FusedAttnFunc kernel call to reduce torch.compile graph breaks"
+    ),
+    condition=lambda ctx: (
+        (
+            (
+                getattr(get_args(ctx), "torch_compile", None) is not None
+                and getattr(get_args(ctx).torch_compile, "enable", False)
+            )
+            # Align with torch_compile_patches.py, which also honors the flat
+            # enable_torch_compile flag; otherwise this prologue optimization is
+            # silently skipped for configs that compile via enable_torch_compile.
+            or getattr(get_args(ctx), "enable_torch_compile", False)
+        )
+        and not getattr(get_args(ctx), "disable_dpa_prologue_patch", False)
+    ),
+)
+def patch_dpa_consolidated_prologue(ctx: PatchContext):
+    """Replace TEDotProductAttention.forward with a two-phase version.
+
+    Phase 1 – eager prologue (``@torch._dynamo.disable``):
+        Runs the full DPA + FusedAttention setup pipeline to compute all
+        arguments for ``FusedAttnFunc.apply``, without running the kernel.
+
+    Phase 2 – compiled kernel call:
+        ``FusedAttnFunc.apply`` (marked ``allow_in_graph``) runs inside the
+        torch.compile compiled graph, enabling fusion with pre/post-attention
+        operations.
+    """
+    from megatron.core.extensions.transformer_engine import TEDotProductAttention
+    from transformer_engine.pytorch.attention.dot_product_attention.backends import (
+        FusedAttnFunc,
+    )
+
+    # ---- Idempotency guard -----------------------------------------------
+    # Re-running this patch would capture the already-patched forward as
+    # "_orig_*", double-wrapping the capture logic. Guard against it.
+    if getattr(TEDotProductAttention, "_primus_dpa_prologue_patched", False):
+        log_rank_0("[Patch:megatron.te.dpa_consolidated_prologue] already applied; skipping re-patch")
+        return
+
+    # ---- TE version check ------------------------------------------------
+    # _DROPOUT_P_IDX (and the captured-args layout) is pinned to TE 2.12. On a
+    # different TE version the layout may drift; the fast-path detection
+    # already falls back to the always-fork safe path on mismatch, but warn so
+    # the perf regression is visible rather than silent.
+    try:
+        import transformer_engine
+
+        _te_version = getattr(transformer_engine, "__version__", None)
+    except Exception:
+        _te_version = None
+    if _te_version is not None and not str(_te_version).startswith("2.12"):
+        log_rank_0(
+            "[Patch:megatron.te.dpa_consolidated_prologue] WARNING: "
+            f"_DROPOUT_P_IDX={_DROPOUT_P_IDX} was verified against TE 2.12 but "
+            f"detected transformer_engine {_te_version}. The dropout=0 fast path "
+            "will conservatively fall back to the safe rng-fork path if the "
+            "captured-args layout differs."
+        )
+
+    # ---- Mark FusedAttnFunc as graph-safe --------------------------------
+    torch._dynamo.allow_in_graph(FusedAttnFunc)
+
+    # ---- Patch FusedAttnFunc.forward for capture mode --------------------
+    _orig_fused_attn_func_fwd = FusedAttnFunc.forward
+
+    @staticmethod
+    def _capturing_fused_attn_func_fwd(ctx, *args):  # noqa: N805 – autograd ctx
+        if getattr(_capture_tls, "capture_mode", False):
+            _capture_tls.captured_args = args
+            raise _CaptureComplete()
+        return _orig_fused_attn_func_fwd(ctx, *args)
+
+    FusedAttnFunc.forward = _capturing_fused_attn_func_fwd
+
+    # ---- Save original TEDotProductAttention.forward ---------------------
+    _orig_te_dpa_forward = TEDotProductAttention.forward
+
+    # ---- Eager prologue --------------------------------------------------
+    @torch._dynamo.disable
+    def _dpa_eager_prologue(
+        te_dpa,
+        query,
+        key,
+        value,
+        attention_mask,
+        attn_mask_type,
+        attention_bias,
+        packed_seq_params,
+        num_splits,
+    ):
+        """Run the full DPA + FusedAttention setup eagerly.
+
+        Returns
+        -------
+        captured_args : tuple | None
+            Arguments for ``FusedAttnFunc.apply`` if fused backend was chosen.
+        fallback_result : Tensor | None
+            Final output tensor if a non-fused backend ran to completion.
+        """
+        _capture_tls.capture_mode = True
+        _capture_tls.captured_args = None
+        fallback_result = None
+
+        try:
+            fallback_result = _orig_te_dpa_forward(
+                te_dpa,
+                query,
+                key,
+                value,
+                attention_mask,
+                attn_mask_type=attn_mask_type,
+                attention_bias=attention_bias,
+                packed_seq_params=packed_seq_params,
+                num_splits=num_splits,
+            )
+        except _CaptureComplete:
+            pass  # expected – fused attention args captured
+        finally:
+            _capture_tls.capture_mode = False
+
+        return _capture_tls.captured_args, fallback_result
+
+    # ---- RNG-correct kernel execution ------------------------------------
+    from megatron.core.tensor_parallel.random import get_cuda_rng_tracker
+
+    @torch._dynamo.disable
+    def _fused_attn_with_rng_fork(captured_args):
+        """Run FusedAttnFunc.apply inside the model-parallel-rng fork.
+
+        The eager prologue's _CaptureComplete exception unwinds TE's
+        internal fork context before the kernel runs, so rng_gen=None
+        causes the default CUDA generator to be advanced instead of
+        model-parallel-rng.  This wrapper re-enters the fork so the
+        kernel's philox_cuda_state call targets the correct generator.
+        """
+        tracker = get_cuda_rng_tracker()
+        if tracker.is_initialized():
+            states = tracker.get_states()
+            if "model-parallel-rng" in states:
+                with tracker.fork("model-parallel-rng"):
+                    return FusedAttnFunc.apply(*captured_args)
+        return FusedAttnFunc.apply(*captured_args)
+
+    # ---- Max-logit bookkeeping (qk_clip) ---------------------------------
+    @torch._dynamo.disable
+    def _update_max_logit_stats(te_dpa, batch_max_logit):
+        if hasattr(te_dpa, "current_max_attn_logits"):
+            if te_dpa.current_max_attn_logits is None:
+                te_dpa.current_max_attn_logits = batch_max_logit
+            else:
+                te_dpa.current_max_attn_logits = torch.max(te_dpa.current_max_attn_logits, batch_max_logit)
+
+    # ---- Replacement forward ---------------------------------------------
+    def _new_te_dpa_forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        attention_mask: Optional[Tensor],
+        attn_mask_type=None,
+        attention_bias: Optional[Tensor] = None,
+        packed_seq_params=None,
+        num_splits: Optional[int] = None,
+    ) -> Tensor:
+        """TEDotProductAttention.forward with consolidated prologue.
+
+        The eager prologue handles all DPA setup (FP8 metadata, backend
+        selection, cu_seqlens computation, etc.).  Then ``FusedAttnFunc.apply``
+        runs inside the compiled graph.
+
+        For non-fused backends (flash / unfused), the original forward runs
+        fully inside the prologue and the result is returned directly.
+        """
+        captured_args, fallback_result = _dpa_eager_prologue(
+            self,
+            query,
+            key,
+            value,
+            attention_mask,
+            attn_mask_type,
+            attention_bias,
+            packed_seq_params,
+            num_splits,
+        )
+
+        if captured_args is not None:
+            # Fast path: when dropout_p == 0 the FusedAttn kernel doesn't
+            # consume RNG, so we can skip the eager rng-fork wrapper and
+            # call FusedAttnFunc.apply directly inside the compiled graph
+            # (it's already marked allow_in_graph above). This drops a
+            # graph break in the common Flux case (attention_dropout=0).
+            #
+            # Safe path: anything that isn't a clean numeric zero (including
+            # IndexError if a TE upgrade changes captured_args' layout, or
+            # an unexpected tensor / object at the dropout slot) falls back
+            # to the always-fork path so correctness is never compromised
+            # by a layout drift.
+            try:
+                _dropout_p = captured_args[_DROPOUT_P_IDX]
+                _use_fast_path = isinstance(_dropout_p, (int, float)) and _dropout_p == 0
+            except (IndexError, TypeError):
+                _use_fast_path = False
+
+            if _use_fast_path:
+                result = FusedAttnFunc.apply(*captured_args)
+            else:
+                result = _fused_attn_with_rng_fork(captured_args)
+            return_max_logit = captured_args[-1]
+
+            if return_max_logit:
+                core_attn_out = result[0].view(*result[0].shape[:-2], -1)
+                _update_max_logit_stats(self, result[1])
+                return core_attn_out
+
+            return result.view(*result.shape[:-2], -1)
+
+        return fallback_result
+
+    # ---- Apply the monkey-patch ------------------------------------------
+    TEDotProductAttention.forward = _new_te_dpa_forward
+    TEDotProductAttention._primus_dpa_prologue_patched = True
+
+    log_rank_0("[Patch:megatron.te.dpa_consolidated_prologue] " "Applied allow_in_graph to FusedAttnFunc")
+    log_rank_0(
+        "[Patch:megatron.te.dpa_consolidated_prologue] "
+        "Replaced TEDotProductAttention.forward with consolidated prologue"
+    )
diff --git a/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py
new file mode 100644
index 000000000..14e775197
--- /dev/null
+++ b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py
@@ -0,0 +1,60 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Transformer Engine Fused Bias SwiGLU Patches
+
+Patches SwiGLUFunction to use TE's fused swiglu/dswiglu kernels when
+USE_TE_SWIGLU=1 is set, providing better performance on ROCm GPUs.
+"""
+
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.te.fused_bias_swiglu",
+    backend="megatron",
+    phase="before_train",
+    description="Use TE fused swiglu/dswiglu kernels in SwiGLUFunction forward/backward",
+    condition=lambda ctx: os.getenv("USE_TE_SWIGLU", "0") == "1",
+)
+def patch_te_fused_bias_swiglu(ctx: PatchContext):
+    """
+    Patch SwiGLUFunction to use Transformer Engine's fused swiglu/dswiglu
+    C++ extensions in forward and backward passes.
+
+    Activated when USE_TE_SWIGLU=1 is set in the environment.
+    """
+    from megatron.core.fusions.fused_bias_swiglu import SwiGLUFunction
+    from transformer_engine.pytorch.cpp_extensions import dswiglu as te_dswiglu
+    from transformer_engine.pytorch.cpp_extensions import swiglu as te_swiglu
+
+    @staticmethod
+    def new_forward(ctx, input, fp8_input_store, cpu_offload_input):
+        input_for_backward = input.to(__import__("torch").float8_e4m3fn) if fp8_input_store else input
+        if cpu_offload_input:
+            input_for_backward.activation_offloading = True
+        ctx.save_for_backward(input_for_backward)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        return te_swiglu(input, None)
+
+    @staticmethod
+    def new_backward(ctx, grad_output):
+        input = ctx.saved_tensors[0]
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        return te_dswiglu(grad_output, input, None), None, None
+
+    SwiGLUFunction.forward = new_forward
+    SwiGLUFunction.backward = new_backward
+
+    log_rank_0(
+        "[Patch:megatron.te.fused_bias_swiglu] Patched SwiGLUFunction "
+        "to use TE fused swiglu/dswiglu kernels (USE_TE_SWIGLU=1)"
+    )
diff --git a/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py b/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py
index eceb47243..bd36f8644 100644
--- a/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py
+++ b/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py
@@ -15,7 +15,7 @@
 import inspect
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _te_general_gemm_workspace_mode():
diff --git a/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py b/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py
index 5e0a9ea83..56b1f6caf 100644
--- a/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py
+++ b/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py
@@ -14,7 +14,7 @@
     make_get_extra_te_kwargs_with_override,
 )
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py b/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py
index beb38bbaf..536470455 100644
--- a/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py
+++ b/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py
@@ -35,7 +35,7 @@
 import torch
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _resolve_grouped_mlp_class():
diff --git a/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py b/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py
index af539b82b..7f74dea81 100644
--- a/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py
+++ b/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py
@@ -15,7 +15,7 @@
     make_get_extra_te_kwargs_with_override,
 )
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py b/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py
index 89a7e2027..58cc89712 100644
--- a/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py
+++ b/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py
@@ -18,7 +18,7 @@
     is_te_v2_or_above,
 )
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _check_tp_overlap_conditions(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/tokenizer_builder_patches.py b/primus/backends/megatron/patches/tokenizer_builder_patches.py
index f441af53f..61c819989 100644
--- a/primus/backends/megatron/patches/tokenizer_builder_patches.py
+++ b/primus/backends/megatron/patches/tokenizer_builder_patches.py
@@ -25,7 +25,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/megatron/patches/torch_compile_patches.py b/primus/backends/megatron/patches/torch_compile_patches.py
new file mode 100644
index 000000000..0bc709d28
--- /dev/null
+++ b/primus/backends/megatron/patches/torch_compile_patches.py
@@ -0,0 +1,82 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron torch.compile Patches
+
+This module contains patches that modify Megatron's setup_model_and_optimizer
+to apply torch.compile after model setup.
+"""
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.training.torch_compile",
+    backend="megatron",
+    phase="before_train",
+    description="Patch setup_model_and_optimizer to apply torch.compile after model setup",
+    priority=60,  # Higher than optimizer patch to wrap last
+    condition=lambda ctx: (
+        getattr(get_args(ctx), "torch_compile", None) is not None
+        and getattr(getattr(get_args(ctx), "torch_compile", None), "enable", False)
+    )
+    or getattr(get_args(ctx), "enable_torch_compile", False),
+)
+def patch_setup_model_and_optimizer_for_torch_compile(ctx: PatchContext):
+    """
+    Patch Megatron's setup_model_and_optimizer to apply torch.compile after model setup.
+
+    Behavior:
+        - Wraps setup_model_and_optimizer() to call apply_torch_compile_if_enabled()
+          after model is created and wrapped
+        - If compilation fails, exception propagates (fails entire setup)
+        - Works for both FSDP2 and non-FSDP2 paths
+    """
+    try:
+        from megatron.training import training
+
+        from primus.backends.megatron.core.utils import (
+            apply_torch_compile_if_enabled,
+            apply_torch_compile_to_optimizer_if_enabled,
+        )
+
+        # Save original function
+        original_setup_model_and_optimizer = training.setup_model_and_optimizer
+
+        def patched_setup_model_and_optimizer(*args, **kwargs):
+            """Patched setup_model_and_optimizer that applies torch.compile after model setup."""
+            from megatron.training import get_args
+
+            result = original_setup_model_and_optimizer(*args, **kwargs)
+
+            # Extract model from result tuple
+            model, optimizer, opt_param_scheduler = result
+
+            # Apply torch.compile (raises exception on failure)
+            megatron_args = get_args()
+            apply_torch_compile_if_enabled(model, megatron_args)
+            apply_torch_compile_to_optimizer_if_enabled(optimizer, megatron_args)
+
+            # Return original result
+            return result
+
+        # Apply the patch
+        training.setup_model_and_optimizer = patched_setup_model_and_optimizer
+        log_rank_0(
+            "[Patch:megatron.training.torch_compile] "
+            "Patched setup_model_and_optimizer to apply torch.compile after model setup"
+        )
+
+    except Exception as e:
+        log_rank_0(
+            f"[Patch:megatron.training.torch_compile] "
+            f"WARNING: Failed to patch setup_model_and_optimizer: {type(e).__name__}: {e}"
+        )
+        import traceback
+
+        log_rank_0(f"Traceback: {traceback.format_exc()}")
diff --git a/primus/backends/megatron/patches/torch_fsdp2_patches.py b/primus/backends/megatron/patches/torch_fsdp2_patches.py
index 1ac3c56cf..b8ba726e2 100644
--- a/primus/backends/megatron/patches/torch_fsdp2_patches.py
+++ b/primus/backends/megatron/patches/torch_fsdp2_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -13,7 +13,7 @@
 
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
@@ -59,27 +59,40 @@ def patch_torch_fsdp(ctx: PatchContext):
         f"[Patch:megatron.fsdp.torch_fsdp2]   Patched megatron.training.training.torch_FSDP "
         f"-> {PrimusTorchFullyShardedDataParallel.__name__}"
     )
-    # Megatron Core 0.16 may pass new kwargs (e.g., force_all_reduce) into
-    # model_chunk.finish_grad_sync(). Keep FSDP2 path forward-compatible by
-    from megatron.core.distributed import data_parallel_base
 
-    original_start_grad_sync = data_parallel_base._BaseDataParallel.start_grad_sync
-    original_finish_grad_sync = data_parallel_base._BaseDataParallel.finish_grad_sync
-
-    if not getattr(original_start_grad_sync, "_primus_grad_sync_compat", False):
-
-        def _patched_start_grad_sync(self, *unused, **unused_kwargs):
-            return original_start_grad_sync(self, *unused)
-
-        def _patched_finish_grad_sync(self, *unused, **unused_kwargs):
-            return original_finish_grad_sync(self)
-
-        setattr(_patched_start_grad_sync, "_primus_grad_sync_compat", True)
-        setattr(_patched_finish_grad_sync, "_primus_grad_sync_compat", True)
-        data_parallel_base._BaseDataParallel.start_grad_sync = _patched_start_grad_sync
-        data_parallel_base._BaseDataParallel.finish_grad_sync = _patched_finish_grad_sync
-
-        log_rank_0(
-            "[Patch:megatron.fsdp.torch_fsdp2]   Patched _BaseDataParallel "
-            "start_grad_sync/finish_grad_sync to accept extra kwargs "
-        )
+    # Patch get_data_parallel_group_if_dtensor to handle 2D HSDP meshes.
+    # The upstream implementation calls tensor.device_mesh.get_group() without mesh_dim,
+    # which fails for 2D meshes. For HSDP (dp_replicate, dp_shard), we return the
+    # shard group (innermost dim) since that's where parameters are actually sharded;
+    # replicas have identical gradients so we must not all-reduce across them.
+    import megatron.core.optimizer.clip_grads as clip_grads_module
+    import megatron.core.utils as mcore_utils
+    from megatron.training import utils as training_utils
+
+    try:
+        from torch.distributed._tensor import DTensor
+
+        HAVE_DTENSOR = True
+    except ImportError:
+        HAVE_DTENSOR = False
+
+    def _get_data_parallel_group_if_dtensor(tensor, data_parallel_group=None):
+        if HAVE_DTENSOR and isinstance(tensor, DTensor):
+            mesh = tensor.device_mesh
+            if mesh.ndim > 1:
+                current_group = mesh.get_group(mesh_dim=-1)
+            else:
+                current_group = mesh.get_group()
+            if data_parallel_group is not None and current_group != data_parallel_group:
+                raise RuntimeError("DTensor mesh group does not match the expected data_parallel_group")
+            return current_group
+        return None
+
+    mcore_utils.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor
+    clip_grads_module.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor
+    if hasattr(training_utils, "get_data_parallel_group_if_dtensor"):
+        training_utils.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor
+    log_rank_0(
+        "[Patch:megatron.fsdp.torch_fsdp2]   Patched get_data_parallel_group_if_dtensor "
+        "for 2D HSDP mesh support"
+    )
diff --git a/primus/backends/megatron/patches/torch_fused_adam_patches.py b/primus/backends/megatron/patches/torch_fused_adam_patches.py
new file mode 100644
index 000000000..9019bfbc4
--- /dev/null
+++ b/primus/backends/megatron/patches/torch_fused_adam_patches.py
@@ -0,0 +1,88 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Torch Fused AdamW Opt-In Patch
+================================
+
+Set ``PRIMUS_TORCH_OPTIM=1`` to force Megatron's optimizer construction to
+use ``torch.optim.AdamW(fused=True)`` instead of TE/Apex ``FusedAdam``. This
+exists purely for bit-level reproducibility experiments against
+flash-linear-attention's (FLA) reference training runs, which use the plain
+PyTorch optimizer.
+
+Mechanism:
+    * ``megatron.core.optimizer`` picks its ``Adam``/``SGD``/
+      ``USING_PYTORCH_OPTIMIZER`` symbols once, at import time, based on
+      whichever of TE / Apex / stock Torch is importable. We override those
+      module attributes directly (binding replacement) -- the private
+      ``_get_megatron_optimizer_based_on_param_groups`` reads
+      ``USING_PYTORCH_OPTIMIZER`` as a bare module global at call time, so
+      reassigning it is enough to route construction through
+      ``torch.optim.AdamW`` for the main (non CPU-offload) path.
+    * The ``fused=True`` kwarg is not otherwise reachable via attribute
+      reassignment (it's set inside a local ``kwargs`` dict built partway
+      through the function body), so that one line is added via a targeted
+      source-string rewrite.
+"""
+
+import os
+
+import torch
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.backends.megatron.patches._source_patch_utils import patch_function_source
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.optimizer.torch_fused_adam"
+
+_ORI_CODE = "adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam"
+_NEW_CODE = _ORI_CODE + '\n                kwargs["fused"] = True'
+
+
+def _enabled() -> bool:
+    return os.environ.get("PRIMUS_TORCH_OPTIM", "0") == "1"
+
+
+def _install_torch_fused_adam_patch() -> None:
+    import megatron.core.optimizer as optimizer_module
+
+    if is_patched(optimizer_module, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] megatron.core.optimizer already patched; skipping.")
+        return
+
+    optimizer_module.Adam = torch.optim.AdamW
+    optimizer_module.SGD = torch.optim.SGD
+    optimizer_module.USING_PYTORCH_OPTIMIZER = True
+
+    patch_function_source(
+        optimizer_module,
+        "_get_megatron_optimizer_based_on_param_groups",
+        _ORI_CODE,
+        _NEW_CODE,
+    )
+
+    mark_patched(optimizer_module, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] PRIMUS_TORCH_OPTIM=1 -> using torch.optim.AdamW(fused=True) "
+        "instead of TE/Apex FusedAdam."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Opt-in: use torch.optim.AdamW(fused=True) instead of TE/Apex FusedAdam, "
+        "to bit-match FLA's reference optimizer for reproducibility experiments."
+    ),
+    priority=40,
+    condition=lambda ctx: _enabled(),
+)
+def patch_torch_fused_adam(ctx: PatchContext) -> None:
+    _install_torch_fused_adam_patch()
diff --git a/primus/backends/megatron/patches/torch_norm_fla_rmsnorm_patches.py b/primus/backends/megatron/patches/torch_norm_fla_rmsnorm_patches.py
new file mode 100644
index 000000000..38182b275
--- /dev/null
+++ b/primus/backends/megatron/patches/torch_norm_fla_rmsnorm_patches.py
@@ -0,0 +1,85 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+WrappedTorchNorm FLA RMSNorm Patch
+====================================
+
+Routes ``WrappedTorchNorm``'s RMSNorm construction through
+``fla.modules.RMSNorm`` (flash-linear-attention's Triton-fused kernel)
+instead of ``torch.nn.RMSNorm``, to match FLA's normalization numerics
+exactly for GDN/KDA/Mamba hybrid parity training.
+
+Toggle: ``args.use_fla_fused_rmsnorm`` (resolved by ``fla_runtime_patches.py``
+from ``PRIMUS_FLA_NORM`` / YAML ``use_fla_fused_rmsnorm``, default False).
+
+``WrappedTorchNorm.__new__`` directly returns the constructed norm module
+(no plain ``__init__`` path), so this is a simple function-wrapping patch --
+no source rewrite needed.
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCH_KEY = "megatron.torch_norm.fla_rmsnorm"
+
+
+def _install_fla_rmsnorm_patch() -> None:
+    from megatron.core.transformer.torch_norm import WrappedTorchNorm
+    from megatron.training import get_args as _get_args
+
+    if is_patched(WrappedTorchNorm, _PATCH_KEY):
+        log_rank_0(f"[Patch:{_PATCH_KEY}] WrappedTorchNorm already patched; skipping.")
+        return
+
+    original_new = WrappedTorchNorm.__new__
+
+    def patched_new(
+        cls,
+        config,
+        hidden_size,
+        eps=1e-5,
+        persist_layer_norm=False,
+        zero_centered_gamma=False,
+        normalization="LayerNorm",
+    ):
+        if config.normalization == "RMSNorm" and getattr(_get_args(), "use_fla_fused_rmsnorm", False):
+            from fla.modules import RMSNorm as FLARMSNorm
+
+            return FLARMSNorm(hidden_size=hidden_size, eps=eps)
+        return original_new(
+            cls,
+            config,
+            hidden_size,
+            eps=eps,
+            persist_layer_norm=persist_layer_norm,
+            zero_centered_gamma=zero_centered_gamma,
+            normalization=normalization,
+        )
+
+    WrappedTorchNorm.__new__ = patched_new
+    mark_patched(WrappedTorchNorm, _PATCH_KEY)
+    log_rank_0(
+        f"[Patch:{_PATCH_KEY}] Patched WrappedTorchNorm.__new__ to use fla.modules.RMSNorm "
+        "when use_fla_fused_rmsnorm is set."
+    )
+
+
+@register_patch(
+    _PATCH_KEY,
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Route WrappedTorchNorm's RMSNorm construction through fla.modules.RMSNorm "
+        "to match FLA's normalization kernel exactly."
+    ),
+    # Runs after fla_runtime_knobs (priority=-100) has resolved args.use_fla_fused_rmsnorm.
+    priority=50,
+    condition=lambda ctx: getattr(get_args(ctx), "use_fla_fused_rmsnorm", False),
+)
+def patch_torch_norm_fla_rmsnorm(ctx: PatchContext) -> None:
+    _install_fla_rmsnorm_patch()
diff --git a/primus/backends/megatron/patches/torch_profiler_patches.py b/primus/backends/megatron/patches/torch_profiler_patches.py
index 0149f5aa5..5a039442d 100644
--- a/primus/backends/megatron/patches/torch_profiler_patches.py
+++ b/primus/backends/megatron/patches/torch_profiler_patches.py
@@ -14,7 +14,7 @@
 import inspect
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_called_from_training_train() -> bool:
@@ -31,8 +31,6 @@ def _is_called_from_training_train() -> bool:
 def _create_primus_prof(args, exp_name: str, original_profile):
     """
     Create torch profiler with Primus options.
-
-    Logic from primus/modules/trainer/megatron/trainer.py L1277-1298.
     """
     import torch
 
diff --git a/primus/backends/megatron/patches/training_log/__init__.py b/primus/backends/megatron/patches/training_log/__init__.py
index 8935e1158..14d08c3a0 100644
--- a/primus/backends/megatron/patches/training_log/__init__.py
+++ b/primus/backends/megatron/patches/training_log/__init__.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -11,16 +11,18 @@
 
     - print_rank_last_patches: scoped print_rank_last hook for training_log
       that injects ROCm memory and throughput statistics.
+    - wall_clock_timer_patch: wraps train_step with a wall-clock timer for
+      NeMo-comparable throughput measurement.
 
-The actual patch registration is handled by ``print_rank_last_patches`` via
-``@register_patch``; this ``__init__`` exists mainly to make
-``training_log`` a proper package so that the auto-import logic in
-``primus.backends.megatron.patches.__init__`` can discover and import
-``print_rank_last_patches`` automatically.
+The actual patch registration is handled via ``@register_patch``; this
+``__init__`` exists mainly to make ``training_log`` a proper package so
+that the auto-import logic in ``primus.backends.megatron.patches.__init__``
+can discover and import the patch modules automatically.
 """
 
 from primus.backends.megatron.patches.training_log import (  # noqa: F401
     print_rank_last_patches,
+    wall_clock_timer_patch,
 )
 
-__all__ = ["print_rank_last_patches"]
+__all__ = ["print_rank_last_patches", "wall_clock_timer_patch"]
diff --git a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py
index 9485b3116..bffddba68 100644
--- a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py
+++ b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -12,8 +12,13 @@
 
     - ROCm/HIP memory stats.
     - Running average elapsed time per iteration (ms).
-    - Running average throughput per GPU (TFLOP/s/GPU).
-    - Running average token throughput per GPU (tokens/s/GPU).
+    - Running average compute per GPU (TFLOP/s/GPU).
+    - Running average token throughput per GPU (tokens/s/GPU) (language models only).
+    - Diffusion-specific metrics (diffusion models only):
+      * Images per GPU (images/s/GPU): instant/average
+      * Latency per image (ms): instant
+      * Image resolution: height x width
+      * Average timestep
 
 Design:
     - We first parse Megatron's original ``log_string`` into a structured
@@ -31,8 +36,42 @@
 
 from primus.core.patches import PatchContext, get_args, register_patch
 from primus.core.utils import logger as primus_logger
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 from primus.core.utils.rocm_mem_info import get_rocm_smi_mem_info
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+
+
+def _is_diffusion_model(args: Any, module_config: Any = None) -> bool:
+    """
+    Detect if this is a diffusion model training run.
+
+    Checks:
+    1. model_type == 'diffusion_model' (from args/params)
+    2. trainer_class contains 'Flux' or 'Diffusion' (from module_config)
+
+    Args:
+        args: Megatron args (module_config.params)
+        module_config: Full module config (for trainer_class access)
+
+    Returns:
+        True if diffusion model, False otherwise
+    """
+    # Check model_type from args
+    model_type = getattr(args, "model_type", None)
+    if model_type == "diffusion_model":
+        return True
+
+    # Check trainer_class from module_config (preferred source)
+    if module_config:
+        trainer_class = getattr(module_config, "trainer_class", None)
+        if trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)):
+            return True
+
+    # Fallback: check trainer_class in args (shouldn't be there if in reserved_keys)
+    trainer_class = getattr(args, "trainer_class", None)
+    if trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)):
+        return True
+
+    return False
 
 
 @dataclass
@@ -250,15 +289,20 @@ def inject(
 
                 # When pipeline parallelism (PP) is enabled, memory usage can vary across ranks.
                 # Therefore, we report the maximum ROCm memory usage across all ranks.
-                r_used_tensor = torch.tensor([r_used], device="cuda", dtype=torch.int64)
-                world_size = torch.distributed.get_world_size()
-                gathered_r_used = [torch.zeros_like(r_used_tensor) for _ in range(world_size)]
-                torch.distributed.all_gather(gathered_r_used, r_used_tensor)
-
-                total_r_used = [t.item() for t in gathered_r_used]
-                log_rank_0(f"total_r_used: {[round(r_used / 1024 ** 3, 2) for r_used in total_r_used]}")
-                max_r_used = max(total_r_used)
-                max_rank = total_r_used.index(max_r_used)
+                # Use constant-size all_reduce(MAX) instead of O(world_size) all_gather:
+                # one reduce for the max value, a second to recover the owning rank
+                # (the rank tensor is masked to -1 on non-max ranks). On ties the
+                # highest such rank wins (vs. the lowest under the previous all_gather).
+                max_used_tensor = torch.tensor([r_used], device="cuda", dtype=torch.int64)
+                torch.distributed.all_reduce(max_used_tensor, op=torch.distributed.ReduceOp.MAX)
+                max_r_used = max_used_tensor.item()
+
+                my_rank = torch.distributed.get_rank()
+                rank_tensor = torch.tensor(
+                    [my_rank if r_used == max_r_used else -1], device="cuda", dtype=torch.int64
+                )
+                torch.distributed.all_reduce(rank_tensor, op=torch.distributed.ReduceOp.MAX)
+                max_rank = rank_tensor.item()
 
                 rocm_mem_str = (
                     f" | rocm mem usage/free/total/usage_ratio: "
@@ -388,14 +432,46 @@ def __init__(self, args: Any):
             f"log_avg_reset_interval: {self._log_avg_reset_interval}"
         )
 
+    def _inject_tflops(self, parsed: TrainingLogInfo) -> None:
+        """
+        Shared TFLOP throughput logic (extracted for reuse by diffusion extension).
+
+        Updates the throughput segment with running-average TFLOP throughput.
+
+        Args:
+            parsed: Parsed training log information
+        """
+        if parsed.throughput_tflops is not None:
+            tflops_value = parsed.throughput_tflops
+            iteration = parsed.iteration
+
+            # Handle warmup & sliding window logic for TFLOPs.
+            if iteration is not None and (
+                iteration == self._log_avg_skip_iterations + 1
+                or len(self._recent_tflop_throughputs) >= self._log_avg_reset_interval
+            ):
+                self._recent_tflop_throughputs.clear()
+
+            # Only accumulate after skip window.
+            if iteration is None or iteration > self._log_avg_skip_iterations:
+                self._recent_tflop_throughputs.append(tflops_value)
+
+            if self._recent_tflop_throughputs:
+                avg_tflops = sum(self._recent_tflop_throughputs) / len(self._recent_tflop_throughputs)
+                idx = parsed.throughput_index
+                if idx is not None and 0 <= idx < len(parsed.segments):
+                    parsed.segments[idx] = (
+                        f"compute per GPU (TFLOP/s/GPU): {tflops_value:.1f} (avg {avg_tflops:.1f})"
+                    )
+
     def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str:
         """
         Update ``parsed`` with running-average TFLOP and token throughput.
 
         - TFLOPs are rendered inline as:
-              throughput per GPU (TFLOP/s/GPU): inst/avg
-        - Tokens are appended as a new segment immediately after throughput:
-              tokens per GPU (tokens/s/GPU): inst/avg
+              compute per GPU (TFLOP/s/GPU): inst (avg <mean>)
+        - Tokens are appended to the same segment immediately after compute:
+              tokens/s/GPU inst/harmonic mean: inst/<harmonic mean>
         """
         try:
             # If no parsed info is provided (e.g., unit tests calling this
@@ -406,27 +482,7 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s
             iteration = parsed.iteration
 
             # ---------------- TFLOPs ----------------
-            if parsed.throughput_tflops is not None:
-                tflops_value = parsed.throughput_tflops
-
-                # Handle warmup & sliding window logic for TFLOPs.
-                if iteration is not None and (
-                    iteration == self._log_avg_skip_iterations + 1
-                    or len(self._recent_tflop_throughputs) >= self._log_avg_reset_interval
-                ):
-                    self._recent_tflop_throughputs.clear()
-
-                # Only accumulate after skip window.
-                if iteration is None or iteration > self._log_avg_skip_iterations:
-                    self._recent_tflop_throughputs.append(tflops_value)
-
-                if self._recent_tflop_throughputs:
-                    avg_tflops = sum(self._recent_tflop_throughputs) / len(self._recent_tflop_throughputs)
-                    idx = parsed.throughput_index
-                    if idx is not None and 0 <= idx < len(parsed.segments):
-                        parsed.segments[idx] = (
-                            f"throughput per GPU (TFLOP/s/GPU): " f"{tflops_value:.1f}/{avg_tflops:.1f}"
-                        )
+            self._inject_tflops(parsed)
 
             # ---------------- Tokens/s ----------------
             if parsed.elapsed_ms is None:
@@ -466,7 +522,16 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s
                 warning_rank_0(f"[Patch:megatron.training_log] No token throughput")
                 return log_string
 
-            avg_tokens = sum(self._recent_token_throughputs) / len(self._recent_token_throughputs)
+            # Use the harmonic mean for the token throughput average. The harmonic
+            # mean is the correct way to average rates (tokens/s) over iterations of
+            # equal token count, since it weights slow iterations more heavily.
+            positive_token_throughputs = [t for t in self._recent_token_throughputs if t > 0]
+            if positive_token_throughputs:
+                avg_tokens = len(positive_token_throughputs) / sum(
+                    1.0 / t for t in positive_token_throughputs
+                )
+            else:
+                avg_tokens = 0.0
 
             # Append token throughput directly after the TFLOP throughput within
             # the same segment. We do not create a new segment to keep the log
@@ -475,7 +540,7 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s
             if idx is not None and 0 <= idx < len(parsed.segments):
                 parsed.segments[idx] = (
                     f"{parsed.segments[idx]} "
-                    f" | tokens per GPU (tokens/s/GPU): {token_value:.1f}/{avg_tokens:.1f}"
+                    f" | tokens/s/GPU inst/harmonic mean: {token_value:.1f}/{avg_tokens:.1f}"
                 )
 
             # String result is ignored by the main patch when parsed is provided.
@@ -485,6 +550,219 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s
             return log_string
 
 
+class DiffusionThroughputAverageExtension(ThroughputAverageExtension):
+    """
+    Helper extension for diffusion models: TFLOP throughput only, no tokens.
+
+    Inherits from ThroughputAverageExtension to reuse TFLOP logic.
+    Skips token throughput calculation entirely (diffusion models use images, not tokens).
+
+    Semantics mirror ThroughputAverageExtension:
+        - Ignore the first `log_avg_skip_iterations` iterations for averaging.
+        - Maintain a sliding window up to `log_avg_reset_interval` entries.
+    """
+
+    def __init__(self, args: Any):
+        super().__init__(args)
+        log_rank_0(
+            f"[Patch:megatron.training_log] DiffusionThroughputAverageExtension initialized "
+            f"(TFLOP only, no tokens)"
+        )
+
+    def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str:
+        """
+        Update ``parsed`` with running-average TFLOP throughput only.
+
+        For diffusion models: TFLOPs only, no token throughput.
+        """
+        try:
+            # If no parsed info is provided, keep the original string unchanged.
+            if parsed is None:
+                return log_string
+
+            # Only inject TFLOP throughput (no tokens for diffusion models)
+            self._inject_tflops(parsed)
+
+            # String result is ignored by the main patch when parsed is provided.
+            return log_string
+        except Exception:
+            # Any parsing / numeric issues should not break logging.
+            return log_string
+
+
+class DiffusionMetricsExtension:
+    """
+    Helper extension to compute and inject diffusion-specific metrics
+    (images per second per GPU and latency per image) into Megatron training logs.
+
+    This extension only activates for diffusion models (model_type == 'diffusion_model').
+    For language models, it early-returns without modifying logs.
+
+    Semantics mirror ThroughputAverageExtension:
+        - Ignore the first `log_avg_skip_iterations` iterations for averaging.
+        - Maintain a sliding window up to `log_avg_reset_interval` entries.
+    """
+
+    def __init__(self, args: Any, module_config: Any = None, runtime_state: Any = None):
+        self._args = args
+        # Store module_config reference for diffusion detection
+        # (We only access trainer_class from it, which is a top-level attribute, not the nested params)
+        self._module_config = module_config
+        self._runtime_state = runtime_state
+        # Cache world_size once at construction time
+        self._world_size = getattr(args, "world_size", None)
+        # Track image throughput statistics across calls
+        self._recent_image_throughputs: list[float] = []
+        # We follow the same warmup/reset semantics as ThroughputAverageExtension:
+        #   - Ignore the first `log_avg_skip_iterations` iterations for averaging
+        #   - Maintain a sliding window of size `log_avg_reset_interval`
+        self._log_avg_skip_iterations: int = int(getattr(args, "log_avg_skip_iterations", 0))
+        self._log_avg_reset_interval: int = int(getattr(args, "log_avg_reset_interval", 1000))
+
+    def _calculate_image_metrics(self, parsed: TrainingLogInfo) -> Optional[tuple[float, float]]:
+        """
+        Calculate image throughput metrics from parsed log info.
+
+        Args:
+            parsed: Parsed training log information
+
+        Returns:
+            Tuple of (images_per_second, latency_per_image_ms) or None if calculation fails
+        """
+        if parsed.elapsed_ms is None or parsed.global_batch_size is None or self._world_size is None:
+            return None
+
+        batch_size = int(parsed.global_batch_size)
+        elapsed_ms = float(parsed.elapsed_ms)
+        elapsed_s = elapsed_ms / 1000.0
+
+        if elapsed_s <= 0:
+            return None
+
+        # Calculate images per second per GPU
+        images_per_second = batch_size / elapsed_s / self._world_size
+
+        # Calculate latency per image (in milliseconds)
+        latency_per_image_ms = elapsed_ms / batch_size
+
+        return (images_per_second, latency_per_image_ms)
+
+    def _format_diffusion_metrics(
+        self, images_per_second: float, latency_per_image_ms: float, avg_images: float
+    ) -> list[str]:
+        """
+        Format diffusion metrics as log segments.
+
+        Args:
+            images_per_second: Current images per second per GPU
+            latency_per_image_ms: Current latency per image in milliseconds
+            avg_images: Average images per second per GPU
+
+        Returns:
+            List of metric strings to append to log segments
+        """
+        metrics = []
+
+        # Add images per GPU metrics (no trailing |, render function adds it)
+        images_metric = f"images per GPU (images/s/GPU): {images_per_second:.2f}/" f"{avg_images:.2f}"
+        metrics.append(images_metric)
+
+        # Add latency per image metric (no trailing |, render function adds it)
+        latency_metric = f"latency per image (ms): {latency_per_image_ms:.1f}"
+        metrics.append(latency_metric)
+
+        # Get metrics from runtime_state (required, no fallback)
+        last_metrics = None
+        if self._runtime_state:
+            last_metrics = self._runtime_state.last_metrics
+        else:
+            # Defensive: log warning but don't break logging
+            log_rank_0(
+                "[DiffusionMetricsExtension] WARNING: runtime_state not available, skipping diffusion metrics"
+            )
+            return metrics  # Return existing metrics without adding diffusion-specific ones
+
+        if last_metrics:
+            if "image_height" in last_metrics and "image_width" in last_metrics:
+                resolution_metric = (
+                    f"image resolution: "
+                    f"{int(last_metrics['image_height'])}x{int(last_metrics['image_width'])}"
+                )
+                metrics.append(resolution_metric)
+            if "avg_timestep" in last_metrics:
+                timestep_metric = f"avg timestep: {last_metrics['avg_timestep']:.1f}"
+                metrics.append(timestep_metric)
+
+            # Wall-clock step timer (from wall_clock_timer_patch)
+            if "wall_clock_step_ms" in last_metrics and self._world_size:
+                wc_ms = float(last_metrics["wall_clock_step_ms"])
+                metrics.append(f"wall clock (ms): {wc_ms:.1f}")
+                gbs = getattr(self._args, "global_batch_size", None)
+                if gbs and wc_ms > 0:
+                    wc_img_per_s = int(gbs) / (wc_ms / 1000.0) / self._world_size
+                    metrics.append(f"wall clock img/s/GPU: {wc_img_per_s:.2f}")
+
+        return metrics
+
+    def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str:
+        """
+        Update ``parsed`` with images per second and latency per image metrics.
+
+        Only activates for diffusion models (model_type == 'diffusion_model').
+        For other models, early-returns without modification.
+
+        Metrics:
+            - images per GPU (images/s/GPU): instant/average
+            - latency per image (ms): instant
+        """
+        try:
+            # If no parsed info is provided, keep the original string unchanged
+            if parsed is None:
+                return log_string
+
+            # Early return for non-diffusion models
+            if not _is_diffusion_model(self._args, self._module_config):
+                return log_string
+
+            # Calculate image metrics
+            metrics_result = self._calculate_image_metrics(parsed)
+            if metrics_result is None:
+                return log_string
+
+            images_per_second, latency_per_image_ms = metrics_result
+            iteration = parsed.iteration
+
+            # Handle warmup & sliding window logic for images
+            if iteration is not None and (
+                iteration == self._log_avg_skip_iterations + 1
+                or len(self._recent_image_throughputs) >= self._log_avg_reset_interval
+            ):
+                self._recent_image_throughputs.clear()
+
+            # Only accumulate after skip window
+            if iteration is None or iteration > self._log_avg_skip_iterations:
+                self._recent_image_throughputs.append(images_per_second)
+
+            if self._recent_image_throughputs:
+                avg_images = sum(self._recent_image_throughputs) / len(self._recent_image_throughputs)
+
+                # Format and append metrics
+                metrics = self._format_diffusion_metrics(images_per_second, latency_per_image_ms, avg_images)
+                parsed.segments.extend(metrics)
+
+            # String result is ignored by the main patch when parsed is provided.
+            return log_string
+        except Exception as e:
+            # Log the exception to help debug, but don't break training
+            iteration = parsed.iteration if parsed else None
+            log_rank_0(
+                f"[Patch:megatron.training_log] DiffusionMetricsExtension ERROR "
+                f"(iteration={iteration}): {type(e).__name__}: {e}"
+            )
+            # Any parsing / numeric issues should not break logging.
+            return log_string
+
+
 @register_patch(
     "megatron.training_log.unified_patch",
     backend="megatron",
@@ -510,6 +788,9 @@ def patch_training_log_unified(ctx: PatchContext):
         # Get unified Megatron args (module_config.params) from context.
         config = get_args(ctx)
 
+        # Get runtime_state from context
+        runtime_state = ctx.extra.get("runtime_state")
+
         # Check whether we should enable ROCm stats / throughput logging.
         use_rocm_mem = bool(getattr(config, "use_rocm_mem_info", False))
         rocm_iters = getattr(config, "use_rocm_mem_info_iters", [])
@@ -532,14 +813,25 @@ def patch_training_log_unified(ctx: PatchContext):
         if getattr(original_training_log, "_primus_training_log_print_rank_wrapper", False):
             return
 
-        # Create helper extensions only when stat injection is enabled so they
-        # keep state (ROCm cache, avg windows) across all training_log
-        # invocations. In forwarding-only mode they are not needed.
-        mem_ext = elapsed_ext = throughput_ext = None
-        if enable_rocm_stats:
-            mem_ext = MemoryStatsExtension(config)
-            elapsed_ext = ElapsedAverageExtension(config)
+        # Create helper extensions once so they keep state (ROCm cache, avg windows)
+        # across all training_log invocations.
+        # Get module_config from context to pass to diffusion extensions
+        # (needed to access trainer_class which is in reserved_keys, not in params)
+        module_config = ctx.extra.get("module_config")
+
+        # Detect if this is a diffusion model to instantiate the correct throughput extension
+        is_diffusion = _is_diffusion_model(config, module_config)
+
+        mem_ext = MemoryStatsExtension(config)
+        elapsed_ext = ElapsedAverageExtension(config)
+        # Use diffusion-specific throughput extension if this is a diffusion model
+        if is_diffusion:
+            throughput_ext = DiffusionThroughputAverageExtension(config)
+        else:
             throughput_ext = ThroughputAverageExtension(config)
+        diffusion_ext = DiffusionMetricsExtension(
+            config, module_config=module_config, runtime_state=runtime_state
+        )
         call_count = 0
         # Capture the original ``print_rank_last`` so we can delegate actual
         # printing back to Megatron after mutating the log string.
@@ -570,12 +862,12 @@ def primus_print_rank_last(log_string: str) -> None:
                     # Parse the original log string once and share across extensions.
                     parsed = parse_training_log_line(log_string)
 
-                    # Inject memory statistics, elapsed avg, throughput, and token
-                    # throughput by mutating the parsed structure. These calls ignore
-                    # their string return value when `parsed` is provided.
+                    # Inject memory statistics, elapsed avg, throughput, and diffusion
+                    # metrics by mutating the parsed structure.
                     mem_ext.inject(log_string, call_count, parsed)
                     elapsed_ext.inject(log_string, parsed)
                     throughput_ext.inject(log_string, parsed)
+                    diffusion_ext.inject(log_string, parsed)
 
                     # Render the final line from the parsed structure.
                     updated = render_training_log_line(parsed)
diff --git a/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py b/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py
new file mode 100644
index 000000000..5febc1cd3
--- /dev/null
+++ b/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py
@@ -0,0 +1,85 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Wall-clock timer patch for Megatron train_step.
+
+Wraps ``megatron.training.training.train_step`` with ``time.perf_counter()``
+to measure the actual forward + backward + optimizer wall-clock duration,
+independent of Megatron's ``interval-time`` timer which includes collective
+barriers, CUDA synchronizes, and logging overhead.
+
+The measured duration is stored on ``runtime_state.last_metrics`` so that
+downstream log extensions (e.g. ``DiffusionMetricsExtension``) can surface
+it alongside existing throughput metrics.
+"""
+
+import time
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _wall_clock_timer_enabled(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return args is not None and getattr(args, "wall_clock_step_timer", False)
+
+
+@register_patch(
+    "megatron.training.wall_clock_step_timer",
+    backend="megatron",
+    phase="before_train",
+    description="Wrap train_step with wall-clock timer for NeMo-comparable throughput measurement",
+    condition=_wall_clock_timer_enabled,
+    priority=90,
+)
+def patch_wall_clock_timer(ctx: PatchContext):
+    import megatron.training.training as megatron_training
+
+    from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+
+    _PATCH_KEY = "megatron.training.train_step_wall_clock_timer"
+    if is_patched(megatron_training, _PATCH_KEY):
+        log_rank_0("[Patch:wall_clock_step_timer] Already applied; skipping re-wrap.")
+        return
+
+    runtime_state = ctx.extra.get("runtime_state")
+
+    _original_train_step = megatron_training.train_step
+
+    def _patched_train_step(
+        forward_step_func,
+        data_iterator,
+        model,
+        optimizer,
+        opt_param_scheduler,
+        config,
+        forward_backward_func,
+        iteration=None,
+    ):
+        t0 = time.perf_counter()
+        result = _original_train_step(
+            forward_step_func,
+            data_iterator,
+            model,
+            optimizer,
+            opt_param_scheduler,
+            config,
+            forward_backward_func,
+            iteration=iteration,
+        )
+        dt_ms = (time.perf_counter() - t0) * 1000.0
+        if runtime_state is not None:
+            runtime_state.update_metrics({"wall_clock_step_ms": dt_ms})
+        return result
+
+    megatron_training.train_step = _patched_train_step
+    mark_patched(megatron_training, _PATCH_KEY)
+    log_rank_0(
+        "[Patch:wall_clock_step_timer] "
+        "Patched train_step with wall-clock timer "
+        f"(runtime_state={'available' if runtime_state else 'unavailable'})"
+    )
diff --git a/primus/backends/megatron/patches/triton_autotune_print_patches.py b/primus/backends/megatron/patches/triton_autotune_print_patches.py
new file mode 100644
index 000000000..3a3dbc391
--- /dev/null
+++ b/primus/backends/megatron/patches/triton_autotune_print_patches.py
@@ -0,0 +1,66 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Silence Triton's autotuner stdout spam.
+
+Triton's ``Autotuner._bench`` (``triton/runtime/autotuner.py``) emits a bare
+``print()`` for *every* (kernel, config) pair it benchmarks when the
+``knobs.autotuning.print`` flag (env: ``TRITON_PRINT_AUTOTUNING``) is on::
+
+    Autotuning kernel _permute_kernel with config BLOCK_SIZE: 64, num_warps: 4, ...
+
+On a multi-node run, every autotuned Triton kernel (e.g. the MoE
+permute / unpermute / sort kernels in
+``primus/backends/transformer_engine/pytorch/triton/permutation.py``) produces a
+whole batch of these lines during warmup, flooding the logs.
+
+Because it is a ``print()`` (not a ``logging`` record), a logging-level change
+cannot suppress it -- the only lever is the Triton ``knobs.autotuning.print``
+flag. We force it off here in the ``before_train`` phase, which runs inside the
+training worker after model setup but before the first training step (and thus
+before any kernel is autotuned). Programmatic assignment to the knob shadows the
+``TRITON_PRINT_AUTOTUNING`` env value, so this works regardless of how the flag
+got turned on.
+
+Set ``PRIMUS_VERBOSE_TRITON=1`` to keep the full autotune trace.
+"""
+
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+@register_patch(
+    "megatron.triton.silence_autotune_print",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Force Triton knobs.autotuning.print off so the autotuner does not "
+        "flood stderr with one print line per kernel config it benchmarks."
+    ),
+    condition=lambda ctx: os.environ.get("PRIMUS_VERBOSE_TRITON", "0") != "1",
+)
+def patch_silence_triton_autotune_print(ctx: PatchContext):
+    """Turn off Triton autotuner printing. Idempotent / best-effort."""
+    del ctx
+    try:
+        from triton import knobs
+    except Exception as exc:  # triton missing -- nothing to do.
+        log_rank_0(f"[Patch:megatron.triton.silence_autotune_print] triton unavailable: {exc!r}")
+        return
+
+    try:
+        knobs.autotuning.print = False
+    except Exception as exc:  # knobs API changed -- never block training.
+        log_rank_0(
+            f"[Patch:megatron.triton.silence_autotune_print] could not set "
+            f"knobs.autotuning.print: {exc!r}"
+        )
+        return
+
+    log_rank_0("[Patch:megatron.triton.silence_autotune_print] Triton autotuner printing disabled.")
diff --git a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py
index 11d649857..fa3f7a21e 100644
--- a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py
+++ b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py
@@ -7,36 +7,28 @@
 """
 hd128 backward crash fix via scoped RTLD_DEEPBIND (ROCm/aiter#1332).
 
-Affects gfx942 (MI300X/MI325X) and gfx950 (MI350X/MI355X).
-
-rocm/primus:v26.3's transformer_engine bundles a STALE vendored ``libmha_bwd.so``
-(old aiter f299f579) that, via ``libck_fused_attn.so``'s DT_NEEDED, interposes the
-global symbol ``aiter::mha_bwd`` over the freshly pinned aiter (b5e03ed1) -> Turbo's
-hd128 backward mis-selects the swa variant and launches with an invalid grid config
-(gridDim.x=0 on gfx942 / "invalid configuration argument" on gfx950) -> 8-rank crash.
-The ``mha_bwd_args`` ABI differs between the two aiter revisions, so a
-*global* fix (RTLD_GLOBAL preload / .so overwrite) repairs Turbo but breaks the TE
-path (TE's libck calls the stale libmha by the old ABI) -- isolation must be scoped.
-
-Fix: wrap ``importlib.import_module`` (aiter loads its mha extensions through it,
-honouring ``sys.setdlopenflags()``) and OR in ``RTLD_DEEPBIND`` only for the pinned
-aiter mha modules (``module_fmha_v3_bwd`` / ``mha_bwd_bf16_*``), so they bind their
-own fresh ``aiter::mha_bwd`` while TE keeps the stale one. DEEPBIND is scoped to these
-few extensions, leaving torch/c10 untouched. Installed by the ``before_train`` patch
-below, gated on Turbo + affected arch (TE path never touched); ``before_train`` runs
-well before the first backward that JIT-imports these modules.
-
-TEMPORARY: delete this file (and its turbo/__init__.py entry) once the base image's
-TE is built against the updated aiter.
+On gfx942/gfx950, rocm/primus:v26.3's transformer_engine bundles a stale vendored
+``libmha_bwd.so`` that shadows the global symbol ``aiter::mha_bwd`` over the pinned
+aiter build, crashing Turbo's hd128 backward. A global fix would break TE (it needs
+the stale symbol), so isolation is scoped: wrap ``importlib.import_module`` and OR in
+``RTLD_DEEPBIND`` only for the pinned aiter mha modules, installed via the
+``before_train`` patch below.
+
+Self-detection: skipped once ``transformer_engine.__version__`` >= 2.14 (confirmed
+fixed in rocm/primus:v26.4). Older images still get the patch.
+
+TEMPORARY: delete this file (and its turbo/__init__.py entry) once every supported
+base image ships a TE built against the updated aiter.
 """
 
 import importlib
 import os
+import re
 import sys
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 _LOG_PREFIX = "[Patch:megatron.turbo.aiter_deepbind]"
 _ENV_SWITCH = "PRIMUS_AITER_DEEPBIND"
@@ -71,6 +63,34 @@ def _is_affected_arch() -> bool:
         return False
 
 
+# First TE version that dropped the stale vendored libmha_bwd.so (see module docstring).
+_TE_FIX_VERSION = (2, 14, 0)
+_TE_FIX_VERSION_STR = ".".join(map(str, _TE_FIX_VERSION))
+
+
+def _te_version_tuple():
+    """Best-effort (major, minor, patch) from ``transformer_engine.__version__``, or
+    None if TE isn't importable / the version string can't be parsed."""
+    try:
+        import transformer_engine
+
+        match = re.match(r"(\d+)\.(\d+)\.(\d+)", getattr(transformer_engine, "__version__", "") or "")
+        return tuple(int(part) for part in match.groups()) if match else None
+    except Exception:  # pragma: no cover - defensive, never block training
+        return None
+
+
+def _te_already_fixed() -> bool:
+    """True if the installed TE no longer vendors the stale libmha_bwd.so.
+
+    Fails safe towards False (unfixed) on an unknown/unparseable version: applying
+    the hook on an already-fixed TE is a harmless no-op, but skipping it on an
+    unfixed TE reintroduces the hd128 backward crash.
+    """
+    version = _te_version_tuple()
+    return version is not None and version >= _TE_FIX_VERSION
+
+
 def _install_deepbind_import_hook() -> bool:
     """
     Idempotently wrap ``importlib.import_module`` to OR in RTLD_DEEPBIND for the pinned
@@ -115,7 +135,7 @@ def _import_module_with_deepbind(name, package=None):
 
 
 def _can_install_aiter_deepbind(ctx: PatchContext) -> bool:
-    """Install only on affected arches (gfx942/gfx950) with the Turbo attention path active."""
+    """Install only on affected arches (gfx942/gfx950) with an unfixed TE and Turbo attention active."""
     args = get_args(ctx)
     if not bool(getattr(args, "use_turbo_attention", False)):
         return False
@@ -127,6 +147,13 @@ def _can_install_aiter_deepbind(ctx: PatchContext) -> bool:
             "aiter DEEPBIND isolation not needed."
         )
         return False
+    if _te_already_fixed():
+        log_rank_0(
+            f"{_LOG_PREFIX} transformer_engine >= {_TE_FIX_VERSION_STR} detected "
+            "(no vendored stale libmha_bwd.so exported as aiter::mha_bwd); aiter DEEPBIND "
+            "isolation not needed on this image."
+        )
+        return False
     return True
 
 
@@ -137,8 +164,8 @@ def _can_install_aiter_deepbind(ctx: PatchContext) -> bool:
     description=(
         "On gfx942/gfx950, install the aiter mha RTLD_DEEPBIND import hook so the Turbo "
         "attention backward binds the pinned aiter::mha_bwd instead of TE's stale vendored "
-        "libmha. Gated on Turbo attention (TE path untouched). Temporary workaround "
-        "(ROCm/aiter#1332)."
+        "libmha. Gated on Turbo attention and an unfixed TE (auto-skips once "
+        f"transformer_engine >= {_TE_FIX_VERSION_STR}). Temporary workaround (ROCm/aiter#1332)."
     ),
     condition=_can_install_aiter_deepbind,
 )
diff --git a/primus/backends/megatron/patches/turbo/fp4_patches.py b/primus/backends/megatron/patches/turbo/fp4_patches.py
index d95fa0e52..c1990fac1 100644
--- a/primus/backends/megatron/patches/turbo/fp4_patches.py
+++ b/primus/backends/megatron/patches/turbo/fp4_patches.py
@@ -10,7 +10,7 @@
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_fp4_can_patch(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/turbo/fp8_patches.py b/primus/backends/megatron/patches/turbo/fp8_patches.py
index 22a9989b6..4cbb97c3b 100644
--- a/primus/backends/megatron/patches/turbo/fp8_patches.py
+++ b/primus/backends/megatron/patches/turbo/fp8_patches.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -13,7 +13,7 @@
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_fp8_can_patch(ctx: PatchContext) -> bool:
@@ -28,7 +28,9 @@ def _is_fp8_can_patch(ctx: PatchContext) -> bool:
     backend="megatron",
     phase="before_train",
     description="Override Megatron get_fp8_context to use Primus implementation when fp8 is enabled",
-    condition=_is_fp8_can_patch,
+    condition=lambda ctx: (
+        _is_fp8_can_patch(ctx) and not getattr(get_args(ctx), "disable_fp8_context_patches", False)
+    ),
 )
 def patch_fp8_context(ctx: PatchContext):
     """
diff --git a/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py
new file mode 100644
index 000000000..0d93be267
--- /dev/null
+++ b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py
@@ -0,0 +1,72 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Tier 1A — fused (residual + RMSNorm) patches.
+
+Wraps the runtime install hook
+:mod:`primus.backends.megatron.core.extensions.fused_residual_rmsnorm` in a
+``@register_patch`` so it runs at the standard ``before_train`` phase via
+``run_patches(...)`` instead of being explicitly installed from the trainer
+entry point.
+
+Two gates (mirroring the install hook):
+  * ``PRIMUS_FUSED_RESIDUAL_NORM=1``    — V1, in-layer ADD#1+norm fusion.
+  * ``PRIMUS_FUSED_RESIDUAL_NORM_V2=1`` — V2, cross-layer ADD#2 carry
+    (implies V1).
+
+Either env var enables the patch. The install hook itself further requires
+``PrimusTurboRMSNorm`` (i.e. ``use_turbo_rms_norm=true``) and bails
+gracefully otherwise, so it's safe to register unconditionally.
+"""
+
+import os
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _env_truthy(name: str) -> bool:
+    v = os.environ.get(name, "0").strip().lower()
+    return v in ("1", "true", "yes", "on")
+
+
+def _is_fused_residual_norm_enabled(ctx: PatchContext) -> bool:
+    return _env_truthy("PRIMUS_FUSED_RESIDUAL_NORM") or _env_truthy("PRIMUS_FUSED_RESIDUAL_NORM_V2")
+
+
+@register_patch(
+    "megatron.turbo.fused_residual_norm",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Fuse residual+add into PrimusTurboRMSNorm via Primus-Turbo "
+        "rmsnorm_residual; gated by PRIMUS_FUSED_RESIDUAL_NORM(_V2)."
+    ),
+    condition=_is_fused_residual_norm_enabled,
+    # Run after megatron.turbo.rms_norm so PrimusTurboRMSNorm is in place
+    # before we extend its forward signature.
+    priority=60,
+)
+def patch_fused_residual_norm(ctx: PatchContext):
+    """Install the fused residual+RMSNorm runtime monkeypatch."""
+    from primus.backends.megatron.core.extensions import fused_residual_rmsnorm
+
+    log_rank_0(
+        "[Patch:megatron.turbo.fused_residual_norm] Installing fused "
+        "residual+RMSNorm (V2={v2}, V1={v1})".format(
+            v1=_env_truthy("PRIMUS_FUSED_RESIDUAL_NORM"),
+            v2=_env_truthy("PRIMUS_FUSED_RESIDUAL_NORM_V2"),
+        )
+    )
+    ok = fused_residual_rmsnorm.install()
+    if ok:
+        log_rank_0("[Patch:megatron.turbo.fused_residual_norm]   install() returned True")
+    else:
+        log_rank_0(
+            "[Patch:megatron.turbo.fused_residual_norm]   install() returned False "
+            "(precondition not met; e.g. use_turbo_rms_norm=false)"
+        )
diff --git a/primus/backends/megatron/patches/turbo/mega_moe_patches.py b/primus/backends/megatron/patches/turbo/mega_moe_patches.py
new file mode 100644
index 000000000..a528039a3
--- /dev/null
+++ b/primus/backends/megatron/patches/turbo/mega_moe_patches.py
@@ -0,0 +1,71 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus Turbo Mega MoE Patches
+
+Replace Megatron's ``MoELayer`` with the PrimusTurbo ``MegaMoE`` adapter.
+EP-only (TP==1) + bf16.
+"""
+
+
+from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _is_turbo_mega_moe_can_patch(ctx: PatchContext) -> bool:
+    """
+    Check if the PrimusTurbo fused Mega MoE layer is enabled.
+
+    Requires:
+      - primus_turbo package is installed
+      - tensor_model_parallel_size == 1
+      - enable_primus_turbo == True
+      - use_turbo_mega_moe == True
+    """
+    args = get_args(ctx)
+    use_turbo_mega_moe = bool(getattr(args, "use_turbo_mega_moe", False))
+    bf16 = bool(getattr(args, "bf16", False))
+    ep_size = int(getattr(args, "expert_model_parallel_size", 1))
+
+    return use_turbo_mega_moe and bf16 and ep_size > 1 and is_primus_turbo_can_patch(ctx)
+
+
+@register_patch(
+    "megatron.turbo.mega_moe",
+    backend="megatron",
+    phase="before_train",
+    description="Replace MoELayer with the PrimusTurbo MegaMoE layer",
+    condition=_is_turbo_mega_moe_can_patch,
+)
+def patch_mega_moe(ctx: PatchContext):
+    """
+    Patch Megatron to use the PrimusTurbo MegaMoE layer.
+
+    Replaces ``MoELayer`` in both ``moe_layer`` and ``gpt.moe_module_specs`` so the
+    ``== MoELayer`` identity check in ``transformer_layer.py`` stays consistent.
+    """
+    from megatron.core.models.gpt import moe_module_specs
+    from megatron.core.transformer.moe import moe_layer
+
+    from primus.backends.megatron.core.extensions.mega_moe import (
+        PrimusTurboMegaMoELayer,
+    )
+
+    log_rank_0("[Patch:megatron.turbo.mega_moe] Patching MoELayer with fused MegaMoE...")
+
+    moe_layer.MoELayer = PrimusTurboMegaMoELayer
+    log_rank_0(
+        "[Patch:megatron.turbo.mega_moe]   Patched "
+        f"megatron.core.transformer.moe.moe_layer.MoELayer -> {PrimusTurboMegaMoELayer.__name__}"
+    )
+
+    moe_module_specs.MoELayer = PrimusTurboMegaMoELayer
+    log_rank_0(
+        "[Patch:megatron.turbo.mega_moe]   Patched "
+        f"megatron.core.models.gpt.moe_module_specs.MoELayer -> {PrimusTurboMegaMoELayer.__name__}"
+    )
diff --git a/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py b/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py
index 3cbc0dd90..2185fdcaa 100644
--- a/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py
+++ b/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py
@@ -13,7 +13,7 @@
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_turbo_deepep_can_patch(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/turbo/rms_norm_patches.py b/primus/backends/megatron/patches/turbo/rms_norm_patches.py
index 0611b465b..cdf45243b 100644
--- a/primus/backends/megatron/patches/turbo/rms_norm_patches.py
+++ b/primus/backends/megatron/patches/turbo/rms_norm_patches.py
@@ -13,7 +13,7 @@
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _is_turbo_rms_norm_can_patch(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py b/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py
index b9d424289..192be08bc 100644
--- a/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py
+++ b/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py
@@ -12,7 +12,7 @@
 
 from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _use_legacy_grouped_gemm(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/turbo/utils.py b/primus/backends/megatron/patches/turbo/utils.py
index 6452b407b..4b1b8ce20 100644
--- a/primus/backends/megatron/patches/turbo/utils.py
+++ b/primus/backends/megatron/patches/turbo/utils.py
@@ -7,7 +7,7 @@
 import importlib.util
 
 from primus.core.patches import PatchContext, get_args
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 def _is_primus_turbo_enabled(ctx: PatchContext) -> bool:
diff --git a/primus/backends/megatron/patches/validation_data_sampling_patches.py b/primus/backends/megatron/patches/validation_data_sampling_patches.py
new file mode 100644
index 000000000..ca3e5c13f
--- /dev/null
+++ b/primus/backends/megatron/patches/validation_data_sampling_patches.py
@@ -0,0 +1,168 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron Validation Data Sampling Patches
+
+Patches Megatron's validation data loading to use a fixed, reproducible
+subset of validation samples (eval_iters * global_batch_size) starting
+from offset 0, matching MLPerf evaluation protocol.
+
+Without this patch, Megatron:
+  1. Over-allocates validation samples proportional to total training steps
+  2. Advances through validation data using consumed_valid_samples offset,
+     making eval loss non-reproducible across runs
+"""
+
+import os
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _is_mlperf_enabled(ctx):
+    return os.getenv("PRIMUS_MLPERF", "0") == "1" and getattr(get_args(ctx), "eval_iters", 0) > 0
+
+
+@register_patch(
+    "megatron.training.validation_num_samples",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Fix validation sample count to eval_iters * global_batch_size " "instead of scaling with train_iters"
+    ),
+    condition=_is_mlperf_enabled,
+)
+def patch_validation_num_samples(ctx: PatchContext):
+    """
+    Patch get_train_valid_test_num_samples to allocate only
+    eval_iters * global_batch_size validation samples instead of
+    (train_iters // eval_interval + 1) * eval_iters * global_batch_size.
+    """
+    import megatron.training.training as training_module
+
+    def patched_get_train_valid_test_num_samples():
+        from megatron.training import get_args
+
+        args = get_args()
+
+        if args.train_samples:
+            train_samples = args.train_samples
+        else:
+            train_samples = args.train_iters * args.global_batch_size
+
+        if args.full_validation:
+            eval_samples = None
+        else:
+            eval_samples = args.eval_iters * args.global_batch_size
+
+        test_samples = args.eval_iters * args.global_batch_size
+
+        if hasattr(args, "phase_transition_iterations") and args.phase_transition_iterations:
+            phase_transition_samples = (
+                [0]
+                + [t * args.global_batch_size for t in args.phase_transition_iterations]
+                + [args.train_samples]
+            )
+            current_sample = args.iteration * args.global_batch_size
+            for i in range(len(phase_transition_samples) - 1):
+                if phase_transition_samples[i] <= current_sample < phase_transition_samples[i + 1]:
+                    train_samples = phase_transition_samples[i + 1] - phase_transition_samples[i]
+                    break
+
+        return (train_samples, eval_samples, test_samples)
+
+    training_module.get_train_valid_test_num_samples = patched_get_train_valid_test_num_samples
+    log_rank_0(
+        "[Patch:megatron.training.validation_num_samples] "
+        "Patched get_train_valid_test_num_samples: eval_samples = eval_iters * gbs"
+    )
+
+
+@register_patch(
+    "megatron.training.validation_data_loader",
+    backend="megatron",
+    phase="before_train",
+    description=(
+        "Always build validation dataloader with consumed_samples=0 and "
+        "cap total_samples at eval_iters * gbs for reproducible evaluation"
+    ),
+    condition=_is_mlperf_enabled,
+)
+def patch_validation_data_loader(ctx: PatchContext):
+    """
+    Patch build_pretraining_data_loader to detect validation datasets and
+    force consumed_samples=0 with total_samples capped at eval_iters * gbs.
+
+    Also patches build_train_valid_test_data_loaders to always pass
+    consumed_samples=0 for validation dataloaders (instead of
+    args.consumed_valid_samples).
+    """
+    import torch.utils.data
+    from megatron.core import mpu
+    from megatron.core.datasets.utils import Split
+
+    try:
+        from megatron.training.datasets import data_samplers as samplers_module
+        from megatron.training.datasets.data_samplers import MegatronPretrainingSampler
+    except ImportError:
+        from megatron.legacy.data import data_samplers as samplers_module
+        from megatron.legacy.data.data_samplers import MegatronPretrainingSampler
+
+    orig_build_loader = samplers_module.build_pretraining_data_loader
+
+    def patched_build_pretraining_data_loader(dataset, consumed_samples, name=""):
+        """Replacement that forces validation to always start from sample 0
+        with a fixed sample count."""
+        if dataset is None:
+            return None
+
+        from megatron.training import get_args
+
+        args = get_args()
+
+        if hasattr(dataset, "split"):
+            split = dataset.split
+        elif hasattr(dataset, "index_split"):
+            split = dataset.index_split
+        else:
+            split = None
+
+        is_validation = (split == Split.valid) or (name == "validation")
+
+        if is_validation:
+            eval_samples = args.eval_iters * args.global_batch_size
+            total_samples = min(len(dataset), eval_samples)
+            batch_sampler = MegatronPretrainingSampler(
+                total_samples=total_samples,
+                consumed_samples=0,
+                micro_batch_size=args.micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+            )
+            return torch.utils.data.DataLoader(
+                dataset,
+                batch_sampler=batch_sampler,
+                num_workers=args.num_workers,
+                pin_memory=True,
+                persistent_workers=True if args.num_workers > 0 else False,
+            )
+
+        return orig_build_loader(dataset, consumed_samples)
+
+    # Replace in both the data_samplers module and training module so all
+    # call sites (including build_train_valid_test_data_loaders) use our version
+    samplers_module.build_pretraining_data_loader = patched_build_pretraining_data_loader
+
+    import megatron.training.training as training_module
+
+    training_module.build_pretraining_data_loader = patched_build_pretraining_data_loader
+
+    log_rank_0(
+        "[Patch:megatron.training.validation_data_loader] "
+        "Patched build_pretraining_data_loader: validation always uses "
+        "consumed_samples=0, total_samples capped at eval_iters * gbs"
+    )
diff --git a/primus/backends/megatron/patches/zebra_llama_flops_patches.py b/primus/backends/megatron/patches/zebra_llama_flops_patches.py
index 1193fabd8..6566d4e35 100644
--- a/primus/backends/megatron/patches/zebra_llama_flops_patches.py
+++ b/primus/backends/megatron/patches/zebra_llama_flops_patches.py
@@ -13,7 +13,7 @@
 """
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def zebra_llama_flops(args, batch_size):
diff --git a/primus/backends/megatron/peft/adapter_wrapper.py b/primus/backends/megatron/peft/adapter_wrapper.py
index cadcff4a8..245174384 100644
--- a/primus/backends/megatron/peft/adapter_wrapper.py
+++ b/primus/backends/megatron/peft/adapter_wrapper.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 import torch
diff --git a/primus/backends/megatron/peft/base.py b/primus/backends/megatron/peft/base.py
index 00d0422d0..b9f8888ca 100644
--- a/primus/backends/megatron/peft/base.py
+++ b/primus/backends/megatron/peft/base.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 import logging
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
diff --git a/primus/backends/megatron/peft/lora.py b/primus/backends/megatron/peft/lora.py
index 69a611794..a778cabbb 100644
--- a/primus/backends/megatron/peft/lora.py
+++ b/primus/backends/megatron/peft/lora.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 import logging
 from dataclasses import dataclass, field
 from typing import List, Literal, Optional
diff --git a/primus/backends/megatron/peft/lora_layers.py b/primus/backends/megatron/peft/lora_layers.py
index be70a6332..d363c555e 100644
--- a/primus/backends/megatron/peft/lora_layers.py
+++ b/primus/backends/megatron/peft/lora_layers.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 import math
 from typing import Any, Literal, Optional, Tuple, Union
 
diff --git a/primus/backends/megatron/peft/module_matcher.py b/primus/backends/megatron/peft/module_matcher.py
index a786efa21..e4828454b 100644
--- a/primus/backends/megatron/peft/module_matcher.py
+++ b/primus/backends/megatron/peft/module_matcher.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Set
diff --git a/primus/backends/megatron/peft/recompute.py b/primus/backends/megatron/peft/recompute.py
index 0137a3b74..4cbb4fa03 100644
--- a/primus/backends/megatron/peft/recompute.py
+++ b/primus/backends/megatron/peft/recompute.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 """Helpers for PEFT-specific activation recompute fixes."""
 
 from __future__ import annotations
@@ -22,7 +25,7 @@
 import torch
 from megatron.core.utils import unwrap_model
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 PEFT_RECOMPUTE_PATCHED: Set[int] = set()
 
diff --git a/primus/backends/megatron/peft/utils.py b/primus/backends/megatron/peft/utils.py
index a59a75034..430ebb3d6 100644
--- a/primus/backends/megatron/peft/utils.py
+++ b/primus/backends/megatron/peft/utils.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 import math
 import re
 from dataclasses import dataclass
diff --git a/primus/backends/megatron/peft/walk_utils.py b/primus/backends/megatron/peft/walk_utils.py
index 4f1cdbf41..ff7bbc9d0 100644
--- a/primus/backends/megatron/peft/walk_utils.py
+++ b/primus/backends/megatron/peft/walk_utils.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Ported from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
 """
 Walking utilities for PyTorch module transformation.
 
diff --git a/primus/backends/megatron/sft/dataset.py b/primus/backends/megatron/sft/dataset.py
index a21b9c23a..cdd9f78d7 100644
--- a/primus/backends/megatron/sft/dataset.py
+++ b/primus/backends/megatron/sft/dataset.py
@@ -121,9 +121,10 @@ def build_train_valid_test_datasets(
     HF / jsonl tokenize+pack pipeline and route to
     ``MlperfPackedDataset``. This lets a Native SFT run consume the exact
     byte-identical packs produced by the upstream mlperf
-    ``download_dataset.py + convert_dataset.py + create_metadata.py``
-    pipeline (used by ``examples/megatron_bridge/configs/MI355X/
-    llama2_70b_lora_posttrain.yaml``).
+    ``primus.backends.megatron_bridge.recipes.mlperf_llama2_70b`` dataset utilities
+    (``download_dataset.py + convert_dataset.py + create_metadata.py``)
+    pipeline (used by ``examples/mlperf/llama2_70b/configs/MI355X/
+    ``llama2_70b_lora_mlperf_posttrain.yaml``).
     """
     from primus.backends.megatron.sft.mlperf_packed_dataset import (
         build_mlperf_packed_datasets,
diff --git a/primus/backends/megatron/sft/gpt_sft_chat_dataset.py b/primus/backends/megatron/sft/gpt_sft_chat_dataset.py
index 9d552bf9e..43143200b 100644
--- a/primus/backends/megatron/sft/gpt_sft_chat_dataset.py
+++ b/primus/backends/megatron/sft/gpt_sft_chat_dataset.py
@@ -1,4 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Adapted from Megatron-Bridge (https://github.com/NVIDIA-NeMo/Megatron-Bridge).
 """
 GPT SFT Chat Dataset - Adapted from Megatron-Bridge
 
diff --git a/primus/backends/megatron/sft/lb_mini_dataset.py b/primus/backends/megatron/sft/lb_mini_dataset.py
new file mode 100644
index 000000000..20124d3f5
--- /dev/null
+++ b/primus/backends/megatron/sft/lb_mini_dataset.py
@@ -0,0 +1,333 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Variable-length SFT sample store + LB-Mini data iterator (ODC path).
+
+This is the data layer for LB-Mini. Unlike ``PackedSFTDataset`` (which bin-packs
+everything into fixed ``max_seq_length`` blocks and hands each DP rank the SAME
+number of blocks), LB-Mini keeps samples VARIABLE length and, for every global
+minibatch, uses Karmarkar-Karp (see ``lb_mini_packing.plan_minibatch``) to:
+
+  * balance the TOTAL token workload across DP ranks, and
+  * let each rank own a DIFFERENT number of micro-batches.
+
+The "different micro-batch count per rank" is exactly what a collective-comm
+backend (NCCL all-gather/reduce-scatter) cannot tolerate, which is why this path
+is gated behind ``enable_odc=true`` -- ODC's point-to-point comm lets ranks run
+out of lockstep without deadlocking.
+
+Two pieces live here:
+  1. ``build_varlen_samples`` -- tokenize the dataset ONCE into variable-length
+     ``{input_ids, labels, loss_mask, length}`` dicts (cached to disk, reusing
+     packing.py's tokenizer-identity cache-key machinery), and
+  2. ``LBMiniDataIterator`` -- the iterator Megatron's schedule pulls from. It
+     refills one global minibatch at a time, plans this rank's micro-batches via
+     KK, and exposes ``current_num_microbatches`` so the schedule patch can run
+     the right (rank-local) number of forward/backward steps.
+"""
+
+import os
+from typing import Dict, List
+
+import numpy as np
+import torch
+
+from primus.backends.megatron.sft.lb_mini_packing import (
+    plan_minibatch,
+    resolve_cost_func,
+)
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _shift_labels_loss_mask(input_ids: np.ndarray, loss_mask: np.ndarray):
+    """Next-token shift, mirroring packing._build_packed_sequence semantics.
+
+    Megatron's ``compute_language_model_loss`` does NOT shift internally, so the
+    dataset MUST emit shifted labels: labels[i]=input_ids[i+1], and loss_mask
+    shifted the same way; the final position has no target and is masked.
+    """
+    n = len(input_ids)
+    labels = input_ids.copy()
+    if n >= 2:
+        labels[:-1] = input_ids[1:]
+    shifted_mask = np.zeros_like(loss_mask)
+    if n >= 2:
+        shifted_mask[:-1] = loss_mask[1:]
+    shifted_mask[-1] = 0
+    return labels, shifted_mask
+
+
+def build_varlen_samples(
+    dataset_name: str,
+    tokenizer,
+    max_seq_length: int,
+    split: str,
+    formatter: str,
+    seed: int,
+    bridge_compat_inline_bos: bool = False,
+    **kwargs,
+) -> List[Dict[str, np.ndarray]]:
+    """Tokenize the dataset into variable-length (un-padded) samples, cached.
+
+    Reuses ``PackedSFTDataset``'s tokenize step (``_tokenize_no_pad``) but does
+    NOT bin-pack: every raw sample becomes one variable-length record. Samples
+    whose tokenized length is 0 are dropped; samples longer than max_seq_length
+    are truncated (same as the non-bridge packing path).
+    """
+    from primus.backends.megatron.sft import packing as pack_mod
+    from primus.backends.megatron.sft.dataset import SFTDataset
+
+    cache_disabled = os.environ.get("PRIMUS_DISABLE_PACK_CACHE", "0") not in ("0", "", "false", "False")
+    pad_id = pack_mod._resolve_pad_token_id(tokenizer)
+    cache_file = lock_file = None
+    if not cache_disabled:
+        # Distinct cache namespace from packed (different layout): prefix "varlen".
+        key = pack_mod._build_pack_cache_key(
+            dataset_name=dataset_name,
+            split=split,
+            formatter=formatter,
+            max_seq_length=max_seq_length,
+            pad_id=pad_id,
+            tokenizer_id=pack_mod._tokenizer_identity(tokenizer),
+            bridge_compat_inline_bos=bridge_compat_inline_bos,
+        )
+        cache_dir = pack_mod._resolve_pack_cache_dir()
+        cache_file = cache_dir / f"sft_varlen_{key}.pt"
+        # Node-local filelock (PRIMUS_PACK_LOCK_DIR): avoid NFS "Stale file handle".
+        _lock_dir = os.environ.get("PRIMUS_PACK_LOCK_DIR")
+        if _lock_dir:
+            os.makedirs(_lock_dir, exist_ok=True)
+            lock_file = os.path.join(_lock_dir, f"sft_varlen_{key}.lock")
+        else:
+            lock_file = cache_dir / f"sft_varlen_{key}.lock"
+
+    def _build() -> List[Dict[str, np.ndarray]]:
+        base = SFTDataset(
+            dataset_name=dataset_name,
+            tokenizer=tokenizer,
+            max_seq_length=max_seq_length,
+            split=split,
+            formatter=formatter,
+            seed=seed,
+            **kwargs,
+        )
+        raw = base.dataset
+        log_rank_0(f"[LB-Mini] Tokenizing {len(raw)} samples (variable length, no packing)...")
+        out: List[Dict[str, np.ndarray]] = []
+        for i in range(len(raw)):
+            sample = pack_mod.normalize_sft_sample(raw[i])
+            formatted = base.formatter.format_sample(sample)
+            tok = pack_mod._tokenize_no_pad(
+                formatted, tokenizer, max_seq_length, bridge_compat_inline_bos=bridge_compat_inline_bos
+            )
+            if tok["length"] <= 0:
+                continue
+            labels, shifted_mask = _shift_labels_loss_mask(tok["input_ids"], tok["loss_mask"])
+            out.append(
+                {
+                    "input_ids": tok["input_ids"].astype(np.int64),
+                    "labels": labels.astype(np.int64),
+                    "loss_mask": shifted_mask.astype(np.int64),
+                    "length": int(tok["length"]),
+                }
+            )
+        sup = sum(int(s["loss_mask"].sum()) for s in out)
+        log_rank_0(
+            f"[LB-Mini] Built {len(out)} variable-length samples; "
+            f"supervised tokens={sup} (sanity: must be > 0)."
+        )
+        return out
+
+    if cache_disabled:
+        return _build()
+
+    try:
+        from filelock import FileLock
+    except ImportError:
+        return _build()
+
+    with FileLock(str(lock_file)):
+        if cache_file.exists():
+            log_rank_0(f"[LB-Mini] varlen CACHE HIT ({cache_file.name}); loading.")
+            samples = torch.load(cache_file, weights_only=False)
+            log_rank_0(f"[LB-Mini] Loaded {len(samples)} variable-length samples from cache.")
+            return samples
+        samples = _build()
+        tmp = cache_file.with_suffix(".pt.tmp")
+        torch.save(samples, tmp)
+        os.replace(tmp, cache_file)
+        log_rank_0(f"[LB-Mini] Cached variable-length samples to {cache_file}")
+        return samples
+
+
+class LBMiniDataIterator:
+    """Iterator that yields THIS rank's variable-length micro-batches.
+
+    Every ``global_batch_size`` raw samples form one *global minibatch*. For each
+    global minibatch we:
+      1. read the per-sample effective lengths (identical view on all ranks),
+      2. ``plan_minibatch`` -> KK-balance across DP ranks + split this rank into
+         micro-batches each <= ``max_token_len``,
+      3. push this rank's micro-batches onto a queue and record
+         ``current_num_microbatches`` (may differ across ranks).
+
+    The schedule patch calls ``begin_minibatch()`` once at the top of each
+    train_step to materialize the plan and read ``current_num_microbatches``;
+    then Megatron's ``forward_step`` pulls each micro-batch via ``__next__``.
+
+    A micro-batch here is ONE packed 1-D sequence (the KK-chosen samples for that
+    micro-batch concatenated). With ``micro_batch_size`` typically 1 and short
+    samples concatenated up to ``max_token_len``, attention runs causal over the
+    concatenation (same implicit-multi-turn regime as Primus packed SFT with
+    ``use_packed_attention=false``); ``loss_mask`` keeps loss on response tokens.
+    """
+
+    def __init__(
+        self,
+        samples: List[Dict[str, np.ndarray]],
+        global_batch_size: int,
+        max_token_len: int,
+        dp_rank: int,
+        dp_size: int,
+        pad_id: int,
+        cost_model: str = "linear",
+        seed: int = 1234,
+        shuffle: bool = True,
+        same_micro_num: bool = False,
+        packing_method: str = "kk",
+    ):
+        self.samples = samples
+        self.global_batch_size = int(global_batch_size)
+        self.max_token_len = int(max_token_len)
+        self.dp_rank = int(dp_rank)
+        self.dp_size = int(dp_size)
+        self.pad_id = int(pad_id)
+        self.cost_func = resolve_cost_func(cost_model)
+        self.seed = int(seed)
+        self.shuffle = bool(shuffle)
+        # same_micro_num=True forces all ranks to the SAME micro-batch count
+        # (all_reduce MAX): this is the no-LB-Mini baseline (ranks aligned, the
+        # short-workload ranks pad/idle). False = LB-Mini (ranks may differ).
+        self.same_micro_num = bool(same_micro_num)
+        # packing_method: "kk" = Karmarkar-Karp load balance (LB-Mini);
+        # "round_robin" = ODC-example-style "None" packing (rank takes
+        # idx[rank::dp], NO balancing -> uneven workload, for A/B of ODC comm).
+        self.packing_method = str(packing_method)
+
+        self._order: List[int] = []
+        self._cursor = 0
+        self._epoch = 0
+        self._queue: List[Dict[str, torch.Tensor]] = []
+        self.current_num_microbatches = 0
+        self._dbg_count = 0
+        self._reshuffle()
+
+    def _reshuffle(self):
+        n = len(self.samples)
+        if self.shuffle:
+            g = np.random.default_rng(self.seed + self._epoch)
+            self._order = g.permutation(n).tolist()
+        else:
+            self._order = list(range(n))
+        self._cursor = 0
+        self._epoch += 1
+
+    def _next_global_indices(self) -> List[int]:
+        """Pull the next ``global_batch_size`` sample indices (wraps epochs)."""
+        if self._cursor + self.global_batch_size > len(self._order):
+            self._reshuffle()
+        idx = self._order[self._cursor : self._cursor + self.global_batch_size]
+        self._cursor += self.global_batch_size
+        return idx
+
+    def _pack_microbatch(self, sample_indices: List[int]) -> Dict[str, torch.Tensor]:
+        """Concatenate chosen samples into one variable-length micro-batch.
+
+        Emits thd-compatible fields (cu_seqlens padded to [1, MAX_SEGMENTS+1],
+        per-segment position_ids, max_sub_seqlen) so that with
+        ``use_packed_attention=true`` the forward path runs SEGMENTED (thd)
+        attention -- O(sum seqlen^2) instead of O(total^2). This is what makes a
+        large max_token_len (packed multi-sample micro-batches) fit in memory.
+        """
+        from primus.backends.megatron.sft.packing import MAX_SEGMENTS_PER_PACK
+
+        ids = np.concatenate([self.samples[i]["input_ids"] for i in sample_indices])
+        lbl = np.concatenate([self.samples[i]["labels"] for i in sample_indices])
+        msk = np.concatenate([self.samples[i]["loss_mask"] for i in sample_indices])
+        seglens = [int(self.samples[i]["length"]) for i in sample_indices]
+        total = int(sum(seglens))
+        # Per-segment position ids (each sub-sample restarts at 0).
+        pos = np.concatenate([np.arange(s, dtype=np.int64) for s in seglens])
+        # cu_seqlens padded to a fixed width so default_collate can stack.
+        cu = [0]
+        off = 0
+        for s in seglens:
+            off += s
+            cu.append(off)
+        n_seg = len(seglens)
+        while len(cu) < MAX_SEGMENTS_PER_PACK + 1:
+            cu.append(total)
+        cu_np = np.asarray(cu, dtype=np.int32)
+        max_sub = max(seglens) if seglens else 0
+        return {
+            "input_ids": torch.from_numpy(ids).long().unsqueeze(0),  # [1, T]
+            "labels": torch.from_numpy(lbl).long().unsqueeze(0),
+            "loss_mask": torch.from_numpy(msk).long().unsqueeze(0),
+            "position_ids": torch.from_numpy(pos).long().unsqueeze(0),
+            "cu_seqlens": torch.from_numpy(cu_np).unsqueeze(0),  # [1, MAX_SEG+1]
+            "num_segments": torch.tensor([n_seg], dtype=torch.int32),  # [1]
+            "max_sub_seqlen": torch.tensor([max_sub], dtype=torch.int32),
+        }
+
+    def begin_minibatch(self) -> int:
+        """Plan ONE global minibatch; fill this rank's queue. Returns rank count."""
+        global_idx = self._next_global_indices()
+        lengths = [int(self.samples[i]["length"]) for i in global_idx]
+        if self.packing_method == "round_robin":
+            # No load balancing (ODC example's default "None"): rank takes
+            # global_idx[rank::dp], one sample per micro-batch. Per-rank workload
+            # is UNEVEN (variable lengths) -> this is where ODC's on-demand comm
+            # is supposed to win by overlapping ranks instead of bulk-syncing.
+            local_positions = list(range(self.dp_rank, len(global_idx), self.dp_size))
+            local_micro = [[p] for p in local_positions]
+        else:
+            # KK-balance across DP ranks + split this rank into micro-batches.
+            local_micro = plan_minibatch(
+                lengths,
+                rank=self.dp_rank,
+                world_size=self.dp_size,
+                max_token_len=self.max_token_len,
+                same_micro_num=self.same_micro_num,  # False=LB-Mini, True=aligned baseline
+                get_seq_costs_func=self.cost_func,
+            )
+        # local_micro entries index into ``lengths`` (== position in global_idx);
+        # map back to absolute sample indices.
+        self._queue = [self._pack_microbatch([global_idx[p] for p in micro]) for micro in local_micro]
+        self.current_num_microbatches = len(self._queue)
+        # Observability: show per-rank micro-batch count for the first few global
+        # minibatches -- this is where "different count per rank" becomes visible.
+        if self._dbg_count < 4:
+            tot = sum(int(mb["input_ids"].numel()) for mb in self._queue)
+            print(
+                f"[LB-Mini] rank{self.dp_rank}: num_microbatches={self.current_num_microbatches} "
+                f"total_tokens={tot}",
+                flush=True,
+            )
+            self._dbg_count += 1
+        return self.current_num_microbatches
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> Dict[str, torch.Tensor]:
+        if not self._queue:
+            # Schedule patch should call begin_minibatch() first; as a fallback
+            # (e.g. eval loops) we refill transparently.
+            self.begin_minibatch()
+        return self._queue.pop(0)
+
+
+__all__ = ["build_varlen_samples", "LBMiniDataIterator"]
diff --git a/primus/backends/megatron/sft/lb_mini_packing.py b/primus/backends/megatron/sft/lb_mini_packing.py
new file mode 100644
index 000000000..365ed80f1
--- /dev/null
+++ b/primus/backends/megatron/sft/lb_mini_packing.py
@@ -0,0 +1,398 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""LB-Mini sequence-length load balancing for Megatron-native SFT (ODC path).
+
+This is the Megatron-side port of the load-balancing core that ODC's official
+example (``primus/core/odc/examples/llm_training/packing.py``) uses. The algorithm
+itself is unchanged (it is already validated on the ODC example); we only strip
+the example's ``from args import get_args`` dependency so it can be driven by
+plain function arguments from a Primus monkey-patch.
+
+What it does
+------------
+Given the *effective* sequence lengths of one global minibatch (``minibatch_size
+* dp`` samples), it produces a balanced assignment so that:
+
+  1. ``get_seqlen_balanced_partitions`` (Karmarkar-Karp differencing) splits the
+     samples across the ``dp`` ranks so each rank's TOTAL token workload is as
+     equal as possible -- this is the "LB" (load balance) part.
+
+  2. Within a rank, ``rearrange_micro_batches`` re-packs that rank's samples into
+     micro-batches each capped at ``max_token_len``. With ``same_num_in_dp=False``
+     (the LB-Mini / ODC mode) each rank ends up with a DIFFERENT number of
+     micro-batches -- the short-workload ranks get fewer, the long ones get more,
+     and crucially NO rank pads up to the global max. This is the "Mini"
+     (variable micro-batch count) part that only ODC's point-to-point comm can
+     drive without a collective deadlock.
+
+The baseline (``same_num_in_dp=True``) instead all-reduce(MAX)es the micro-batch
+count so every rank runs the same number of steps -- exactly Megatron's current
+behaviour, kept here so the two paths share one code base.
+
+IMPORTANT: this module is pure CPU index math; it never touches CUDA or the
+process group except for the optional ``all_reduce`` used by the baseline
+``same_num_in_dp=True`` path.
+"""
+
+import heapq
+import os
+from typing import Callable, List, Optional
+
+import torch
+from torch import distributed as dist
+
+
+# ---------------------------------------------------------------------------
+# Cost models: map a sequence length to a scalar "workload".
+# ---------------------------------------------------------------------------
+def get_seq_costs_linear(seq_len: List[int]) -> List[float]:
+    """Linear cost == token count. Good default for MLP-bound SFT."""
+    return list(seq_len)
+
+
+# Default fit coefficients (a, b) + normalizer = ODC example's DeepSeek-1.5B fit.
+# A 7B/8B-class model has a different attention/MLP FLOP ratio, so the relative
+# weighting of long vs short sequences differs. To keep the cost model
+# model-appropriate WITHOUT forking this shared function (existing 1.5B runs stay
+# byte-for-byte identical), the coefficients are overridable per-run via env:
+#   LB_MINI_FIT_A / LB_MINI_FIT_B / LB_MINI_FIT_NORM.
+# Unset -> the 1.5B defaults below. For a 7B-class model use a=4.348357, b=8.768377.
+_FIT_A_DEFAULT = 1.982122
+_FIT_B_DEFAULT = 2.611821
+_FIT_NORM_DEFAULT = 32000.0
+
+
+def get_seq_costs_fit(seq_len: List[int]) -> List[float]:
+    """Quadratic-ish fit (attention-aware): a*s^2 + b*s with s normalized.
+
+    Coefficients default to the ODC example's DeepSeek-1.5B fit and may be
+    overridden per-run via env (LB_MINI_FIT_A/B/NORM) so a 7B/8B-class model can
+    use its own (a=4.348357, b=8.768377). Only meaningful when
+    use_packed_attention=true (segmented thd attn); under full O(total^2)
+    attention the linear model is the correct target.
+    """
+    a = float(os.environ.get("LB_MINI_FIT_A", _FIT_A_DEFAULT))
+    b = float(os.environ.get("LB_MINI_FIT_B", _FIT_B_DEFAULT))
+    norm = float(os.environ.get("LB_MINI_FIT_NORM", _FIT_NORM_DEFAULT))
+    normalized = [s / norm for s in seq_len]
+    return [a * s * s + b * s for s in normalized]
+
+
+def resolve_cost_func(name: str) -> Callable[[List[int]], List[float]]:
+    if name == "linear":
+        return get_seq_costs_linear
+    if name == "fit":
+        return get_seq_costs_fit
+    raise ValueError(f"Unknown LB-Mini cost model: {name!r}. Supported: linear, fit")
+
+
+# ---------------------------------------------------------------------------
+# Karmarkar-Karp largest-differencing partitioning.
+# https://en.wikipedia.org/wiki/Largest_differencing_method
+# ---------------------------------------------------------------------------
+def karmarkar_karp(seq_cost_list: List[float], k_partitions: int, equal_size: bool):
+    """Partition indices of ``seq_cost_list`` into ``k_partitions`` balanced sets."""
+
+    class Set:
+        def __init__(self) -> None:
+            self.sum = 0
+            self.items = []
+
+        def add(self, idx: int, val: float):
+            self.items.append((idx, val))
+            self.sum += val
+
+        def merge(self, other):
+            for idx, val in other.items:
+                self.items.append((idx, val))
+                self.sum += val
+
+        def __lt__(self, other):
+            if self.sum != other.sum:
+                return self.sum < other.sum
+            if len(self.items) != len(other.items):
+                return len(self.items) < len(other.items)
+            return self.items < other.items
+
+    class State:
+        def __init__(self, items, k: int) -> None:
+            self.k = k
+            self.sets = [Set() for _ in range(k)]
+            assert len(items) in [1, k], f"{len(items)} not in [1, {k}]"
+            for i, (idx, seqlen) in enumerate(items):
+                self.sets[i].add(idx=idx, val=seqlen)
+            self.sets = sorted(self.sets, reverse=True)
+
+        def get_partitions(self):
+            partitions = []
+            for i in range(len(self.sets)):
+                cur = [idx for idx, _ in self.sets[i].items]
+                partitions.append(cur)
+            return partitions
+
+        def merge(self, other):
+            for i in range(self.k):
+                self.sets[i].merge(other.sets[self.k - 1 - i])
+            self.sets = sorted(self.sets, reverse=True)
+
+        @property
+        def spread(self) -> float:
+            return self.sets[0].sum - self.sets[-1].sum
+
+        def __lt__(self, other):
+            if self.spread != other.spread:
+                return self.spread > other.spread
+            return self.sets[0] > other.sets[0]
+
+    sorted_seq_cost_list = sorted([(seqlen, i) for i, seqlen in enumerate(seq_cost_list)])
+    states_pq = []
+    if equal_size:
+        assert len(seq_cost_list) % k_partitions == 0
+        for offset in range(0, len(sorted_seq_cost_list), k_partitions):
+            items = []
+            for i in range(k_partitions):
+                seqlen, idx = sorted_seq_cost_list[offset + i]
+                items.append((idx, seqlen))
+            heapq.heappush(states_pq, State(items=items, k=k_partitions))
+    else:
+        for seqlen, idx in sorted_seq_cost_list:
+            heapq.heappush(states_pq, State(items=[(idx, seqlen)], k=k_partitions))
+
+    while len(states_pq) > 1:
+        state0 = heapq.heappop(states_pq)
+        state1 = heapq.heappop(states_pq)
+        state0.merge(state1)
+        heapq.heappush(states_pq, state0)
+
+    partitions = states_pq[0].get_partitions()
+    if equal_size:
+        for partition in partitions:
+            assert len(partition) * k_partitions == len(seq_cost_list)
+    return partitions
+
+
+# ---------------------------------------------------------------------------
+# Local-search refinement to tighten the KK partitions.
+# ---------------------------------------------------------------------------
+_EPS = 1e-6
+
+
+def _swap_max_partition(seq_cost_list, partitions):
+    max_cost = sum(seq_cost_list[_] for _ in partitions[0])
+    min_cost = sum(seq_cost_list[_] for _ in partitions[-1])
+    for i, item_idx in enumerate(partitions[0]):
+        item_cost = seq_cost_list[item_idx]
+        for j in range(1, len(partitions)):
+            cost_j = sum(seq_cost_list[_] for _ in partitions[j])
+            for k, swap_idx in enumerate(partitions[j]):
+                swap_cost = seq_cost_list[swap_idx]
+                if item_cost - swap_cost <= _EPS:
+                    continue
+                if cost_j - swap_cost + item_cost + _EPS < max_cost and (
+                    max_cost - item_cost + swap_cost > min_cost + _EPS
+                ):
+                    partitions[j][k] = item_idx
+                    partitions[0][i] = swap_idx
+                    return True
+    return False
+
+
+def _swap_min_partition(seq_cost_list, partitions):
+    max_cost = sum(seq_cost_list[_] for _ in partitions[0])
+    min_cost = sum(seq_cost_list[_] for _ in partitions[-1])
+    for i, item_idx in enumerate(partitions[-1]):
+        item_cost = seq_cost_list[item_idx]
+        for j in range(0, len(partitions) - 1):
+            cost_j = sum(seq_cost_list[_] for _ in partitions[j])
+            for k, swap_idx in enumerate(partitions[j]):
+                swap_cost = seq_cost_list[swap_idx]
+                if swap_cost - item_cost <= _EPS:
+                    continue
+                if cost_j - swap_cost + item_cost > min_cost + _EPS and (
+                    min_cost - item_cost + swap_cost + _EPS < max_cost
+                ):
+                    partitions[j][k] = item_idx
+                    partitions[-1][i] = swap_idx
+                    return True
+    return False
+
+
+def _balance_partition(seq_cost_list, partitions):
+    while True:
+        partitions = sorted(
+            partitions,
+            key=lambda x: (sum(seq_cost_list[i] for i in x), min(x) if x else 0),
+            reverse=True,
+        )
+        if _swap_max_partition(seq_cost_list, partitions):
+            continue
+        if _swap_min_partition(seq_cost_list, partitions):
+            continue
+        break
+    return partitions
+
+
+def get_seqlen_balanced_partitions(
+    seqlen_list: List[int],
+    k_partitions: int,
+    equal_size: bool,
+    get_seq_costs_func: Optional[Callable] = None,
+) -> List[List[int]]:
+    """Balance ``seqlen_list`` indices into ``k_partitions`` by total workload.
+
+    equal_size=True  -> every partition has the same item count (baseline).
+    equal_size=False -> partitions may have different item counts (LB-Mini).
+    """
+    if get_seq_costs_func is None:
+        get_seq_costs_func = get_seq_costs_linear
+    seq_cost_list = get_seq_costs_func(seqlen_list)
+    assert len(seq_cost_list) >= k_partitions, f"{len(seq_cost_list)} < {k_partitions}"
+
+    def _check_and_sort(partitions):
+        assert len(partitions) == k_partitions
+        seen = set()
+        out = [None] * k_partitions
+        for i, partition in enumerate(partitions):
+            assert len(partition) > 0, f"partition {i} empty"
+            seen.update(partition)
+            out[i] = sorted(partition)
+        assert seen == set(range(len(seq_cost_list)))
+        return out
+
+    def _diff(partitions):
+        loads = [sum(seq_cost_list[i] for i in p) for p in partitions]
+        return max(loads) - min(loads)
+
+    partitions = karmarkar_karp(seq_cost_list, k_partitions, equal_size)
+    partitions = _balance_partition(seq_cost_list, partitions)
+    if not equal_size and len(seq_cost_list) % k_partitions == 0:
+        eq = karmarkar_karp(seq_cost_list, k_partitions, equal_size=True)
+        eq = _balance_partition(seq_cost_list, eq)
+        if _diff(partitions) > _diff(eq):
+            partitions = eq
+    return _check_and_sort(partitions)
+
+
+def _ceildiv(a, b):
+    return -(a // -b)
+
+
+def rearrange_micro_batches(
+    seq_len_effective: List[int],
+    max_token_len: int,
+    dp_group=None,
+    same_num_in_dp: bool = True,
+    sort_partition_workload: bool = True,
+    get_seq_costs_func: Optional[Callable] = None,
+) -> List[List[int]]:
+    """Split one rank's samples into micro-batches each <= ``max_token_len``.
+
+    same_num_in_dp=True  -> all DP ranks agree on the micro-batch count via
+                            all_reduce(MAX). Megatron's current behaviour.
+    same_num_in_dp=False -> each rank uses its own count (LB-Mini; ODC only).
+    """
+    if get_seq_costs_func is None:
+        get_seq_costs_func = get_seq_costs_linear
+    total_seqlen = sum(seq_len_effective)
+    num_micro_batches = min(len(seq_len_effective), _ceildiv(total_seqlen, max_token_len))
+    if dist.is_initialized() and same_num_in_dp:
+        t = torch.tensor([num_micro_batches]).cuda()
+        dist.all_reduce(t, op=dist.ReduceOp.MAX, group=dp_group)
+        num_micro_batches = t.cpu().item()
+
+    while True:
+        assert num_micro_batches <= len(seq_len_effective)
+        micro_bsz_idx = get_seqlen_balanced_partitions(
+            seq_len_effective,
+            num_micro_batches,
+            equal_size=False,
+            get_seq_costs_func=get_seq_costs_func,
+        )
+        check_failed = False
+        for partition in micro_bsz_idx:
+            if sum(seq_len_effective[i] for i in partition) > max_token_len:
+                check_failed = True
+                break
+        actual_size = num_micro_batches + 1 if check_failed else len(micro_bsz_idx)
+        if dist.is_initialized() and same_num_in_dp:
+            t = torch.tensor([actual_size]).cuda()
+            dist.all_reduce(t, op=dist.ReduceOp.MAX, group=dp_group)
+            actual_size = t.cpu().item()
+        if actual_size == num_micro_batches:
+            break
+        num_micro_batches += 1
+
+    if sort_partition_workload:
+        micro_bsz_idx.sort(
+            key=lambda partition: (
+                sum(get_seq_costs_func([seq_len_effective[idx] for idx in partition])),
+                min(partition) if partition else 0,
+            ),
+            reverse=True,
+        )
+    return micro_bsz_idx
+
+
+def plan_minibatch(
+    lengths: List[int],
+    rank: int,
+    world_size: int,
+    max_token_len: int,
+    same_micro_num: bool = False,
+    get_seq_costs_func: Optional[Callable] = None,
+) -> List[List[int]]:
+    """Top-level LB-Mini planner for ONE global minibatch.
+
+    Args:
+        lengths: effective seq lengths of ALL samples in this global minibatch
+                 (length == minibatch_size * world_size), identical on every rank.
+        rank/world_size: this rank within the DP group.
+        max_token_len: per-micro-batch token cap (memory constraint).
+        same_micro_num: False -> LB-Mini (ODC, ranks differ); True -> baseline.
+
+    Returns:
+        local_idx: list of micro-batches for THIS rank; each micro-batch is a list
+                   of indices into ``lengths``. ``len(local_idx)`` is this rank's
+                   micro-batch count (may differ across ranks when same_micro_num
+                   is False).
+    """
+    if get_seq_costs_func is None:
+        get_seq_costs_func = get_seq_costs_linear
+    assert max(lengths) <= max_token_len, f"{max(lengths)} > max_token_len={max_token_len}"
+
+    # Step 1: KK-balance the whole minibatch across DP ranks.
+    mini_partitions = get_seqlen_balanced_partitions(
+        lengths,
+        world_size,
+        equal_size=same_micro_num,
+        get_seq_costs_func=get_seq_costs_func,
+    )
+    local_index = mini_partitions[rank]
+    local_lengths = [lengths[i] for i in local_index]
+
+    # Step 2: split this rank's slice into micro-batches.
+    micro_indexes = rearrange_micro_batches(
+        local_lengths,
+        max_token_len,
+        same_num_in_dp=same_micro_num,
+        sort_partition_workload=True,
+        get_seq_costs_func=get_seq_costs_func,
+    )
+    # Map rank-local positions back to global minibatch indices.
+    local_idx = [[local_index[p] for p in micro] for micro in micro_indexes]
+    return local_idx
+
+
+__all__ = [
+    "get_seq_costs_linear",
+    "get_seq_costs_fit",
+    "resolve_cost_func",
+    "karmarkar_karp",
+    "get_seqlen_balanced_partitions",
+    "rearrange_micro_batches",
+    "plan_minibatch",
+]
diff --git a/primus/backends/megatron/sft/mlperf_packed_dataset.py b/primus/backends/megatron/sft/mlperf_packed_dataset.py
index 5a0876d3a..18d59de14 100644
--- a/primus/backends/megatron/sft/mlperf_packed_dataset.py
+++ b/primus/backends/megatron/sft/mlperf_packed_dataset.py
@@ -9,7 +9,7 @@
 This module lets the Megatron-native SFT backend consume the exact same
 ``train.npy`` / ``validation.npy`` / ``packed_metadata.jsonl`` artefacts that
 the upstream mlperf LLama-2-70B PEFT recipe ships through Megatron-Bridge
-(``examples/megatron_bridge/configs/MI355X/llama2_70b_lora_posttrain.yaml``).
+(``examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml``).
 
 Why a separate dataset class?
 -----------------------------
diff --git a/primus/backends/megatron/sft/packing.py b/primus/backends/megatron/sft/packing.py
index d1dae01f4..4afb7ad64 100644
--- a/primus/backends/megatron/sft/packing.py
+++ b/primus/backends/megatron/sft/packing.py
@@ -526,7 +526,18 @@ def __init__(
         )
         cache_dir = _resolve_pack_cache_dir()
         cache_file = cache_dir / f"sft_pack_{cache_key}.pt"
-        lock_file = cache_dir / f"sft_pack_{cache_key}.lock"
+        # Keep the lock on node-LOCAL storage when PRIMUS_PACK_LOCK_DIR is set. On a
+        # shared NFS cache dir, cross-node FileLock contention (e.g. 16 ranks over 2
+        # nodes in dual-node ODC) triggers ESTALE ([Errno 116] Stale file handle).
+        # Routing the lock to node-local /tmp avoids the NFS flock race; the cache
+        # .pt itself stays shared on NFS (deterministic build, atomic rename).
+        _lock_base = os.environ.get("PRIMUS_PACK_LOCK_DIR")
+        if _lock_base:
+            _lock_dir = Path(_lock_base)
+            _lock_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            _lock_dir = cache_dir
+        lock_file = _lock_dir / f"sft_pack_{cache_key}.lock"
 
         # Lazy import keeps filelock optional for non-cache code paths (and out
         # of unit tests that don't exercise PackedSFTDataset).
diff --git a/primus/backends/megatron/sft/preprocessing.py b/primus/backends/megatron/sft/preprocessing.py
index 22aaf08e7..599f769ff 100644
--- a/primus/backends/megatron/sft/preprocessing.py
+++ b/primus/backends/megatron/sft/preprocessing.py
@@ -16,7 +16,7 @@
 from primus.backends.megatron.sft.schema import FormattedSFTSample, SFTSample
 
 try:
-    from primus.modules.module_utils import log_rank_0 as _primus_log_rank_0
+    from primus.core.utils.module_utils import log_rank_0 as _primus_log_rank_0
 except ImportError:
     _primus_log_rank_0 = None
 
diff --git a/primus/backends/megatron/sft/runtime.py b/primus/backends/megatron/sft/runtime.py
index 68b333f2f..a578e38e4 100644
--- a/primus/backends/megatron/sft/runtime.py
+++ b/primus/backends/megatron/sft/runtime.py
@@ -11,7 +11,7 @@
 from typing import Any, Optional
 
 from primus.backends.megatron.sft.dataset import build_train_valid_test_datasets
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _safe_signature(fn: Callable[..., Any]) -> inspect.Signature | None:
diff --git a/primus/backends/megatron/training/diffusion/__init__.py b/primus/backends/megatron/training/diffusion/__init__.py
new file mode 100644
index 000000000..fb6c0a922
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/__init__.py
@@ -0,0 +1,12 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Diffusion training helpers.
+
+Import the concrete helpers from their submodules directly, e.g.
+``from primus.backends.megatron.training.diffusion.forward_step import
+flux_forward_step_func``.
+"""
diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py
new file mode 100644
index 000000000..34ab2f001
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/forward_step.py
@@ -0,0 +1,608 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from NeMo's Flux training architecture.
+
+"""
+Forward step functions for diffusion model training.
+
+This module provides forward step implementations for different
+diffusion models, handling the training loop logic.
+
+Supported data formats (framework-standard keys):
+    - Pre-encoded: latents, prompt_embeds, pooled_prompt_embeds, text_ids (optional)
+    - Raw: images, txt - encodes on-the-fly
+
+Architecture follows functional composition for clarity and testability.
+"""
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    generate_text_position_ids,
+    pack_latents,
+    unpack_latents,
+)
+from primus.backends.megatron.training.diffusion.noise_utils import (
+    apply_flow_matching_noise,
+)
+from primus.backends.megatron.training.diffusion.timestep_sampling import (
+    LogitNormalSampler,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def prepare_flux_latents(
+    latents: torch.Tensor,
+    scheduler,
+    img_ids: Optional[torch.Tensor] = None,
+    guidance_scale: Optional[float] = None,
+    use_guidance_embed: bool = False,
+    timestep_sampler=None,  # Optional: custom timestep sampler
+    pregenerated_noise: Optional[torch.Tensor] = None,
+    pregenerated_timesteps: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, ...]:
+    """
+    Prepare latents for Flux training forward pass.
+
+    This function:
+        1. Generates img_ids if not provided (for robustness)
+        2. Samples timesteps using configurable sampling strategy
+        3. Adds noise to latents using flow matching
+        4. Packs latents into sequence format
+        5. Prepares guidance embeddings (if enabled)
+
+    Args:
+        latents: Clean latent tensor (B, C, H, W)
+        scheduler: Flow matching scheduler
+        img_ids: Image position IDs (B, H*W/4, 3). If None, will be generated
+        guidance_scale: Guidance scale value (for CFG)
+        use_guidance_embed: Whether to use guidance embedding
+        timestep_sampler: Optional custom timestep sampler
+                          (default: LogitNormalSampler)
+        pregenerated_noise: If provided, use this noise instead of sampling.
+            Used by deterministic comparison tests to ensure identical inputs.
+        pregenerated_timesteps: If provided, use these timesteps (in [0,1] range)
+            instead of sampling. Used by deterministic comparison tests.
+
+    Returns:
+        Tuple containing:
+            - clean_latents: Original latents (for target computation)
+            - noise: Sampled noise
+            - packed_noisy_latents: Noisy latents in packed format
+            - img_ids: Image position IDs
+            - guidance_vec: Guidance vector (or None)
+            - timesteps: Sampled timesteps (in [0, num_train_timesteps] range)
+            - sigma_1d: Raw sigma in [0, 1] range, shape [B]. Pass directly to
+              the model as timesteps_norm to avoid bf16 round-trip precision loss.
+
+    Reference:
+        NeMo's prepare_image_latent_like_reference()
+    """
+    batch_size, num_channels, height, width = latents.shape
+    device = latents.device
+    dtype = latents.dtype
+
+    # Use default sampler if not provided
+    if timestep_sampler is None:
+        timestep_sampler = LogitNormalSampler()
+
+    # Generate img_ids if not provided (for robustness with variable sizes)
+    if img_ids is None:
+        img_ids = generate_image_position_ids(batch_size, height, width, device, dtype)
+
+    if pregenerated_noise is not None:
+        noise = pregenerated_noise.to(device=device, dtype=dtype)
+    else:
+        noise = torch.randn_like(latents, device=device, dtype=dtype)
+
+    if pregenerated_timesteps is not None:
+        sigma = pregenerated_timesteps.to(device=device, dtype=dtype)
+        timesteps = sigma * scheduler.num_train_timesteps
+    else:
+        timesteps, sigma = timestep_sampler.sample(batch_size, device, scheduler)
+
+    # Convert sigma to correct dtype
+    sigma = sigma.to(dtype=dtype)
+
+    # Save 1D sigma [B] before unsqueezing — used as timesteps_norm to avoid
+    # the bf16 round-trip (sigma * 1000 / 1000) that corrupts ~2.5% of values.
+    sigma_1d = sigma.clone()
+
+    # Broadcast sigma to match latent dimensions
+    while len(sigma.shape) < latents.ndim:
+        sigma = sigma.unsqueeze(-1)
+
+    # Flow matching forward process: x_t = (1 - sigma) * x_0 + sigma * noise
+    noisy_latents = apply_flow_matching_noise(latents, noise, sigma)
+
+    # Pack latents into sequence format
+    packed_noisy_latents = pack_latents(noisy_latents)
+
+    # Prepare guidance embedding (if enabled)
+    if use_guidance_embed and guidance_scale is not None:
+        guidance_vec = torch.full(
+            (batch_size,),
+            guidance_scale,
+            device=device,
+            dtype=dtype,
+        )
+    else:
+        guidance_vec = None
+
+    return (
+        latents,
+        noise,
+        packed_noisy_latents,
+        img_ids,
+        guidance_vec,
+        timesteps,
+        sigma_1d,
+    )
+
+
+# NOTE: kept as an eager alias — torch.compile breaks CUDA RNG reproducibility
+# (compiled torch.randn_like produces different values than eager mode with the
+# same generator state). prepare_flux_latents only contains small ops (randn,
+# rand, element-wise), so compile overhead exceeds any fusion benefit. Eager
+# also matches NeMo's RNG sequence for cross-framework convergence comparison.
+_eager_prepare_flux_latents = prepare_flux_latents
+
+
+def flux_forward_step_func(
+    data_iterator,
+    model,
+    scheduler,
+    use_guidance_embed=False,
+    guidance_scale=None,
+    timestep_sampler=None,
+    cfg_dropout_prob=0.0,
+    empty_t5_encodings=None,
+    empty_clip_encodings=None,
+    vae_scale=None,
+    vae_shift=None,
+    vae_latent_mode="presampled",
+    per_step_rng_reseed=False,
+    step_count=0,
+):
+    """
+    Forward step function for Flux training with distributed data loading.
+
+    Following Megatron's multimodal data loading pattern:
+    - When TP=1 (pure DP): each rank loads data directly, no broadcast needed
+    - When TP>1: only TP rank 0 has data_iterator, broadcast to other TP ranks
+    - Middle PP stages return early
+
+    This function orchestrates the training step by:
+    1. Handling distributed data loading (broadcast from rank 0)
+    2. Loading or encoding images (via helper function)
+    3. Loading or encoding text (via helper function)
+    4. Optionally applying CFG dropout (replacing text embeddings with empty encodings)
+    5. Preparing latents with noise and packing (via helper function)
+    6. Running model forward pass
+    7. Returning model output and loss computation inputs
+
+    Supports two data formats (follows NeMo conventions):
+        1. Pre-encoded: latents, prompt_embeds, pooled_prompt_embeds, text_ids (optional)
+        2. Raw: images, txt - encodes on-the-fly
+
+    Architecture follows NeMo conventions for Flux training.
+
+    Args:
+        data_iterator: Iterator yielding training batches (None on non-dataloader ranks)
+        model: Flux model instance with encoders (config.params_dtype used for data broadcasting)
+        scheduler: Flow matching scheduler
+        use_guidance_embed: Whether model uses guidance embedding
+        guidance_scale: Guidance scale for CFG training
+        timestep_sampler: Optional custom timestep sampler (default: LogitNormalSampler)
+        cfg_dropout_prob: Probability of replacing text embeddings with empty encodings (default: 0.0)
+        empty_t5_encodings: Pre-generated fixed empty T5 encodings (seq_len, 1, context_dim)
+        empty_clip_encodings: Pre-generated fixed empty CLIP encodings (vec_in_dim,)
+        vae_scale: Optional VAE latent scale factor (default: None, MLPerf uses 0.3611)
+        vae_shift: Optional VAE latent shift factor (default: None, MLPerf uses 0.1159)
+        vae_latent_mode: How to obtain latents from the batch (default: "presampled").
+            "presampled" — use stored latents directly.
+            "resample" — reconstruct latents from stored mean+logvar via
+            reparameterization at every step, then apply vae_scale/vae_shift.
+        per_step_rng_reseed: Reseed the default CUDA generator at each step
+            to isolate training random ops from model forward RNG consumption
+            (default: False).
+        step_count: Monotonically increasing counter identifying this forward
+            call. Used to derive a unique per-step RNG seed. Managed by the
+            caller (DiffusionPretrainTrainer) and reconstructed from checkpoint
+            state on resume as iteration * num_microbatches.
+
+    Returns:
+        Tuple of (noise_pred, clean_latents, noise, loss_mask, metrics_dict, is_validation)
+        - noise_pred: Model output (predicted velocity) [B, C, H, W]
+        - clean_latents: Original clean latents [B, C, H, W]
+        - noise: Sampled noise [B, C, H, W]
+        - loss_mask: Optional mask for variable-length sequences [B] or None
+        - metrics_dict: Dictionary with training metrics
+        - is_validation: True when batch contains "timestep" key (MLPerf validation mode)
+    """
+    # Reseed default CUDA generator per step to isolate training random ops
+    # (noise, timesteps, CFG dropout) from model forward RNG consumption.
+    # Required because TE fused attention advances the default generator even
+    # with dropout=0 when the DPA prologue patch is active.
+    if per_step_rng_reseed:
+        from megatron.core import parallel_state as _ps
+        from megatron.training import get_args as _get_args
+
+        _seed = _get_args().seed
+        _per_rank_seed = _seed + 100 * _ps.get_data_parallel_rank()
+        _step_seed = (_per_rank_seed * 10000 + step_count) % (2**63)
+        torch.cuda.manual_seed(_step_seed)
+
+    from megatron.core import tensor_parallel
+    from megatron.core.parallel_state import (
+        get_tensor_model_parallel_rank,
+        get_tensor_model_parallel_world_size,
+    )
+
+    # Pipeline parallelism (pipeline_model_parallel_size > 1) is rejected at
+    # config construction for Flux (see BaseDiffusionConfig.__post_init__), so
+    # no middle-pipeline-stage handling is needed here.
+    # Derive compute dtype from bf16/fp16 flags rather than params_dtype.
+    # When use_fsdp2_fp32_param_optimizer is active, params_dtype is FP32 (for
+    # optimizer precision) but compute should still be BF16/FP16.
+    if model.config.bf16:
+        compute_dtype = torch.bfloat16
+    elif model.config.fp16:
+        compute_dtype = torch.float16
+    else:
+        compute_dtype = model.config.params_dtype
+
+    tp_size = get_tensor_model_parallel_world_size()
+
+    if tp_size == 1:
+        # Pure DP: every rank has its own data iterator (is_distributed=True).
+        # Skip broadcast_data overhead (~1.5ms of GPU idle from NCCL
+        # self-broadcasts, GPU->CPU transfers, and .item() sync stalls).
+        if data_iterator is None:
+            raise RuntimeError(
+                "data_iterator is None with TP=1; dataset provider must set is_distributed=True"
+            )
+        batch = next(data_iterator)
+        if not isinstance(batch, dict):
+            raise TypeError(
+                f"[ForwardStep] Expected batch to be dict, got {type(batch)}. Batch value: {batch}"
+            )
+        if vae_latent_mode == "resample":
+            required_keys = ["mean", "logvar", "prompt_embeds", "pooled_prompt_embeds"]
+        else:
+            required_keys = ["latents", "prompt_embeds", "pooled_prompt_embeds"]
+        missing_keys = [k for k in required_keys if k not in batch]
+        if missing_keys:
+            raise KeyError(
+                f"[ForwardStep] Batch missing required keys: {missing_keys}. "
+                f"Got keys: {list(batch.keys())}. "
+                f"vae_latent_mode={vae_latent_mode}"
+            )
+
+        # Cast to compute_dtype and move to CUDA in one pass
+        for key in batch:
+            if isinstance(batch[key], torch.Tensor):
+                if batch[key].is_floating_point():
+                    batch[key] = batch[key].to(dtype=compute_dtype, device="cuda", non_blocking=True)
+                elif not batch[key].is_cuda:
+                    batch[key] = batch[key].cuda(non_blocking=True)
+
+        prompt_embeds = batch["prompt_embeds"]
+        pooled_prompt_embeds = batch["pooled_prompt_embeds"]
+
+        if vae_latent_mode == "resample":
+            mean = batch["mean"]
+            logvar = batch["logvar"]
+        else:
+            latents = batch["latents"]
+
+        loss_mask = batch.get("loss_mask")
+    else:
+        # TP > 1: only rank 0 loads data, broadcast to other TP ranks
+        if data_iterator is not None and get_tensor_model_parallel_rank() == 0:
+            try:
+                batch = next(data_iterator)
+                if not isinstance(batch, dict):
+                    raise TypeError(
+                        f"[ForwardStep] Expected batch to be dict, got {type(batch)}. "
+                        f"Batch value: {batch}"
+                    )
+                if vae_latent_mode == "resample":
+                    required_keys = ["mean", "logvar", "prompt_embeds", "pooled_prompt_embeds"]
+                else:
+                    required_keys = ["latents", "prompt_embeds", "pooled_prompt_embeds"]
+                missing_keys = [k for k in required_keys if k not in batch]
+                if missing_keys:
+                    raise KeyError(
+                        f"[ForwardStep] Batch missing required keys: {missing_keys}. "
+                        f"Got keys: {list(batch.keys())}. "
+                        f"vae_latent_mode={vae_latent_mode}"
+                    )
+            except StopIteration:
+                raise RuntimeError(
+                    "[ForwardStep] Data iterator exhausted (should be infinite with "
+                    "MegatronDataloaderWrapper). This indicates a bug in the dataloader wrapper."
+                )
+            except Exception as e:
+                logger.error(f"[ForwardStep] Error getting batch: {type(e).__name__}: {e}")
+                import traceback
+
+                logger.error(f"[ForwardStep] Traceback: {traceback.format_exc()}")
+                raise
+        else:
+            batch = None
+
+        if batch is not None:
+            for key in batch:
+                if isinstance(batch[key], torch.Tensor) and batch[key].is_floating_point():
+                    batch[key] = batch[key].to(dtype=compute_dtype)
+
+        try:
+            prompt_embeds = tensor_parallel.broadcast_data(["prompt_embeds"], batch, compute_dtype).get(
+                "prompt_embeds"
+            )
+            pooled_prompt_embeds = tensor_parallel.broadcast_data(
+                ["pooled_prompt_embeds"], batch, compute_dtype
+            ).get("pooled_prompt_embeds")
+
+            if vae_latent_mode == "resample":
+                mean = tensor_parallel.broadcast_data(["mean"], batch, compute_dtype).get("mean")
+                logvar = tensor_parallel.broadcast_data(["logvar"], batch, compute_dtype).get("logvar")
+            else:
+                latents = tensor_parallel.broadcast_data(["latents"], batch, compute_dtype).get("latents")
+        except Exception as e:
+            logger.error(f"[ForwardStep] Error broadcasting data: {type(e).__name__}: {e}")
+            logger.error(
+                f"[ForwardStep] batch type: {type(batch)}, "
+                f"batch keys: {list(batch.keys()) if isinstance(batch, dict) else 'N/A'}"
+            )
+            if isinstance(batch, dict):
+                for key, value in batch.items():
+                    logger.error(
+                        f"[ForwardStep]   {key}: type={type(value)}, "
+                        f"shape={value.shape if hasattr(value, 'shape') else 'N/A'}"
+                    )
+            import traceback
+
+            logger.error(f"[ForwardStep] Traceback: {traceback.format_exc()}")
+            raise
+
+        loss_mask = None
+        if batch is not None and "loss_mask" in batch:
+            loss_mask = tensor_parallel.broadcast_data(["loss_mask"], batch, compute_dtype).get("loss_mask")
+            if not loss_mask.is_cuda:
+                loss_mask = loss_mask.cuda(non_blocking=True)
+
+        if not prompt_embeds.is_cuda:
+            prompt_embeds = prompt_embeds.cuda(non_blocking=True)
+        if not pooled_prompt_embeds.is_cuda:
+            pooled_prompt_embeds = pooled_prompt_embeds.cuda(non_blocking=True)
+
+    # Obtain latents based on vae_latent_mode
+    if vae_latent_mode == "resample":
+        # Resample mode: reconstruct latents from posterior parameters each step
+        if not mean.is_cuda:
+            mean = mean.cuda(non_blocking=True)
+        if not logvar.is_cuda:
+            logvar = logvar.cuda(non_blocking=True)
+        std = torch.exp(0.5 * logvar)
+        vae_eps = torch.randn_like(mean)
+
+        latents = mean + std * vae_eps
+        # Scale/shift is always applied after resampling (raw posterior -> normalized latents)
+        latents = vae_scale * (latents - vae_shift)
+    else:
+        # Presampled mode: use stored latents directly
+        if not latents.is_cuda:
+            latents = latents.cuda(non_blocking=True)
+
+    # Validation detection.
+    #
+    # MLPerf v5.1 Flux1 validation spec (flux1/nemo/README.md §6 "Evaluation"):
+    #   - Per-sample fixed timestep t ∈ {0/8, 1/8, ..., 7/8}
+    #   - Equal sample count per timestep (29 696 / 8 = 3 712)
+    #   - val_loss = mean over per-timestep means (equivalent to flat mean given
+    #     equal counts).
+    #
+    # NeMo's official to_webdataset preserves a `timestep` integer per sample
+    # from the MLCommons Arrow source. Our `primus-cli data diffusion-ingest`
+    # path (pipelines/ingest.py:33 `ARROW_COLUMNS`) ingests only the 4 tensor
+    # columns and writes `{"key": ...}` to the json sidecar — so our val shards
+    # are MISSING the timestep field, which used to make this branch fall
+    # through to the training path with uniform-random timesteps via the
+    # `timestep_sampler`. That produced a *different* val_loss estimator than
+    # the spec's: E_t~U[0,1][MSE] (Monte Carlo over [0,1]) vs the spec's
+    # left-Riemann sum over t∈{0/8..7/8}. The two estimators are not
+    # comparable, so a uniform-random val path can make val_loss converge
+    # spuriously fast relative to the reference convergence point.
+    #
+    # Fix: when batch is in eval mode (model.training=False, set by
+    # the evaluation harness via `model_module.eval()`) and lacks a `timestep`
+    # field, inject equidistant timesteps deterministically by within-batch
+    # index. With MBS=64, each micro-batch covers each t∈{0..7} exactly 8
+    # times. Across 58 micro-batches × 8 DP ranks = 464 micro-batches → exactly
+    # 3 712 samples per timestep, matching the MLPerf v5.1 spec count.
+    #
+    # CFG dropout during val: SUPPRESSED.
+    #
+    # Reference-implementation tally for "apply CFG dropout during validation":
+    #   NeMo MLPerf reference (custom_flux.py): ON
+    #   AMD's MLPerf submission: OFF
+    #   TorchTitan flux training script: OFF
+    #
+    # CFG-off during validation is MLPerf-compliant under the v6.0 rules even
+    # though NeMo (which generated the reference convergence point) has it on.
+    # Empirically, applying CFG-during-val structurally inflates val_loss by
+    # ~0.015-0.030 (the 10% unconditional samples pay a ~0.15-0.30 MSE
+    # penalty), which is enough to materially shift the convergence-crossing
+    # step, so we keep it off to match the submission configuration.
+    is_validation = False
+    if batch is not None and "timestep" in batch:
+        is_validation = True
+        val_timesteps = batch["timestep"].float() / 8.0
+        batch["timesteps"] = val_timesteps
+    elif batch is not None and not model.training:
+        is_validation = True
+        batch_size_val = pooled_prompt_embeds.shape[0]
+        val_idx = torch.arange(batch_size_val, device="cuda") % 8
+        batch["timestep"] = val_idx
+        val_timesteps = val_idx.to(dtype=compute_dtype) / 8.0
+        batch["timesteps"] = val_timesteps
+
+    # Matches NeMo's forward_step which wraps prepare_image_latent_like_reference
+    # in torch.no_grad() — no gradients needed for position IDs, noise sampling,
+    # timestep sampling, or latent packing.
+    with torch.no_grad():
+        # Generate img_ids based on latent spatial dimensions
+        # NOTE: When RoPE fusion is enabled, we use batch_size=1 to satisfy Transformer Engine's
+        # fused kernel constraints (freqs must have shape [S, 1, 1, D]). PyTorch broadcasting
+        # applies the same position grid across all batch samples. This requires all images in
+        # the batch to have the same resolution (same height/width).
+        rope_fusion_batch_size = 1 if model.config.apply_rope_fusion else latents.shape[0]
+        img_ids = generate_image_position_ids(
+            batch_size=rope_fusion_batch_size,
+            height=latents.shape[2],
+            width=latents.shape[3],
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        # Generate text_ids (Flux convention: zeros for text position IDs)
+        # NOTE: When RoPE fusion is enabled, use batch_size=1 for consistency with img_ids
+        # (broadcasting will handle the actual batch dimension). This matches NVIDIA's MLPerf
+        # implementation strategy: both txt_ids and img_ids have shape [1, seq_len, 3] with
+        # RoPE fusion, allowing proper concatenation before the fused RoPE kernel.
+        text_ids = generate_text_position_ids(
+            batch_size=rope_fusion_batch_size,
+            seq_len=prompt_embeds.shape[1],
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        # Extract pre-generated noise/timesteps from batch (deterministic tests)
+        batch_noise = None
+        batch_timesteps = None
+        if batch is not None:
+            if tp_size == 1:
+                batch_noise = batch.get("noise")
+                batch_timesteps = batch.get("timesteps")
+            else:
+                if "noise" in batch:
+                    batch_noise = tensor_parallel.broadcast_data(["noise"], batch, compute_dtype).get("noise")
+                    if not batch_noise.is_cuda:
+                        batch_noise = batch_noise.cuda(non_blocking=True)
+                if "timesteps" in batch:
+                    batch_timesteps = tensor_parallel.broadcast_data(["timesteps"], batch, compute_dtype).get(
+                        "timesteps"
+                    )
+                    if not batch_timesteps.is_cuda:
+                        batch_timesteps = batch_timesteps.cuda(non_blocking=True)
+
+        # Prepare latents (noise, packing, scheduling).
+        # Eager wrapper — see _eager_prepare_flux_latents NOTE for why compile
+        # is intentionally disabled (CUDA RNG reproducibility).
+        (
+            clean_latents,
+            noise,
+            packed_noisy_latents,
+            img_ids,
+            guidance_vec,
+            timesteps,
+            sigma_1d,
+        ) = _eager_prepare_flux_latents(
+            latents=latents,
+            scheduler=scheduler,
+            img_ids=img_ids,
+            guidance_scale=guidance_scale,
+            use_guidance_embed=use_guidance_embed,
+            timestep_sampler=timestep_sampler,
+            pregenerated_noise=batch_noise,
+            pregenerated_timesteps=batch_timesteps,
+        )
+
+    # CFG dropout: randomly replace text embeddings with fixed empty encodings.
+    # Placed after prepare_flux_latents so the RNG consumption order matches NeMo:
+    # VAE resample → noise → timesteps → CFG dropout.
+    # Applied during training only — validation uses fixed per-sample timesteps.
+    if (
+        not is_validation
+        and cfg_dropout_prob > 0.0
+        and empty_t5_encodings is not None
+        and empty_clip_encodings is not None
+    ):
+        batch_size_cfg = pooled_prompt_embeds.shape[0]
+        dropout_mask = torch.rand(batch_size_cfg, device="cuda") < cfg_dropout_prob
+
+        empty_t5 = empty_t5_encodings.to(device="cuda", dtype=prompt_embeds.dtype, non_blocking=True)
+        empty_t5 = empty_t5.squeeze(1).unsqueeze(0)
+
+        if empty_t5.shape[1] != prompt_embeds.shape[1]:
+            raise ValueError(
+                f"Empty T5 encoding seq_len ({empty_t5.shape[1]}) does not match "
+                f"data T5 seq_len ({prompt_embeds.shape[1]}). "
+                f"Regenerate empty encodings with matching t5_max_length."
+            )
+
+        t5_mask = dropout_mask.view(-1, 1, 1).expand_as(prompt_embeds)
+        prompt_embeds = torch.where(t5_mask, empty_t5.expand_as(prompt_embeds), prompt_embeds)
+
+        empty_clip = empty_clip_encodings.to(
+            device="cuda", dtype=pooled_prompt_embeds.dtype, non_blocking=True
+        )
+        clip_mask = dropout_mask.view(-1, 1).expand_as(pooled_prompt_embeds)
+        pooled_prompt_embeds = torch.where(
+            clip_mask, empty_clip.expand_as(pooled_prompt_embeds), pooled_prompt_embeds
+        )
+
+    # Transpose for Megatron format (sequence-first)
+    packed_noisy_latents = packed_noisy_latents.transpose(0, 1)
+    prompt_embeds = prompt_embeds.transpose(0, 1)
+
+    # Use raw sigma directly instead of timesteps/1000 to avoid bf16 round-trip
+    timesteps_norm = sigma_1d.to(dtype=packed_noisy_latents.dtype)
+
+    with torch.amp.autocast("cuda", enabled=True, dtype=compute_dtype):
+        noise_pred = model(
+            img=packed_noisy_latents,
+            txt=prompt_embeds,
+            y=pooled_prompt_embeds,
+            timesteps=timesteps_norm,
+            img_ids=img_ids,
+            txt_ids=text_ids,
+            guidance=guidance_vec,
+        )
+
+        # Unpack latents from sequence format
+        noise_pred = noise_pred.transpose(0, 1)  # (S, B, C*4) -> (B, S, C*4)
+        noise_pred = unpack_latents(
+            noise_pred,
+            height=clean_latents.shape[2],
+            width=clean_latents.shape[3],
+        )  # -> (B, C, H, W)
+
+    # Create metrics dict for logging
+    metrics = {
+        "batch_size": latents.shape[0],
+        "image_height": clean_latents.shape[2] * 8,  # VAE 8x downsampling
+        "image_width": clean_latents.shape[3] * 8,
+        "latent_channels": clean_latents.shape[1],
+        "avg_timestep": timesteps.float().mean(),
+        "text_seq_len": prompt_embeds.shape[0],  # After transpose to (S, B, C)
+        "img_seq_len": packed_noisy_latents.shape[0],  # After transpose to (S, B, C)
+    }
+
+    # Return model output and loss computation inputs (matching Megatron's pattern)
+    return noise_pred, clean_latents, noise, loss_mask, metrics, is_validation
+
+
+__all__ = [
+    "prepare_flux_latents",
+    "flux_forward_step_func",
+]
diff --git a/primus/backends/megatron/training/diffusion/loss_computation.py b/primus/backends/megatron/training/diffusion/loss_computation.py
new file mode 100644
index 000000000..ca2c9d6cf
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/loss_computation.py
@@ -0,0 +1,94 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Loss computation for diffusion models.
+
+This module provides reusable loss computation logic for different
+diffusion training objectives. All functions are pure (no state),
+making them easy to test and reusable across different model architectures.
+
+Supported loss types:
+    - Flow Matching: Used by Flux, SD3, and video models
+    - Epsilon Prediction: Used by DDPM and older models
+    - V-Prediction: Alternative parameterization
+
+All loss functions accept tensors of any shape and reduce to scalar.
+"""
+
+from typing import Optional
+
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def compute_flow_matching_loss(
+    prediction: Tensor,
+    clean_latents: Tensor,
+    noise: Tensor,
+    loss_mask: Optional[Tensor] = None,
+) -> Tensor:
+    """
+    Compute flow matching loss with optional masking.
+
+    Flow matching models predict the velocity field that transforms
+    noise to clean latents. The training objective is:
+
+    Formula: target = noise - clean_latents
+             loss = MSE(prediction, target)
+
+    This loss is model-agnostic and works with any tensor shape,
+    making it reusable across 2D (images) and 3D (video) models.
+
+    Args:
+        prediction: Model output (predicted velocity) [any shape]
+        clean_latents: Original clean latents [same shape as prediction]
+        noise: Sampled noise [same shape as prediction]
+        loss_mask: Optional mask for variable-length sequences [batch_size]
+                   or broadcast-compatible shape. Default None (no masking).
+
+    Returns:
+        Scalar loss value (mean squared error, optionally masked).
+
+    Reference:
+        Flow Matching for Generative Modeling
+        https://arxiv.org/abs/2210.02747
+
+    Example (no masking):
+        >>> pred = torch.randn(2, 16, 64, 64)  # Batch of 2D latents
+        >>> clean = torch.randn(2, 16, 64, 64)
+        >>> noise = torch.randn(2, 16, 64, 64)
+        >>> loss = compute_flow_matching_loss(pred, clean, noise)
+        >>> loss.backward()
+
+    Example (with packing/masking):
+        >>> pred = torch.randn(2, 16, 64, 64)
+        >>> clean = torch.randn(2, 16, 64, 64)
+        >>> noise = torch.randn(2, 16, 64, 64)
+        >>> mask = torch.tensor([1.0, 0.0])  # Second sample is padding
+        >>> loss = compute_flow_matching_loss(pred, clean, noise, mask)
+    """
+    target = noise - clean_latents
+
+    if loss_mask is None:
+        # Simple case: direct mean reduction
+        loss = F.mse_loss(prediction.float(), target.float(), reduction="mean")
+    else:
+        # With masking for variable-length sequences (packing support)
+        loss_per_element = F.mse_loss(prediction.float(), target.float(), reduction="none")
+
+        # Broadcast mask to match loss_per_element shape if needed
+        if loss_mask.dim() == 1 and loss_per_element.dim() > 1:
+            mask_shape = [loss_mask.shape[0]] + [1] * (loss_per_element.dim() - 1)
+            loss_mask = loss_mask.view(*mask_shape)
+
+        # Apply mask and compute mean over valid elements
+        loss = (loss_per_element * loss_mask).sum() / loss_mask.sum()
+
+    return loss
+
+
+__all__ = [
+    "compute_flow_matching_loss",
+]
diff --git a/primus/backends/megatron/training/diffusion/noise_utils.py b/primus/backends/megatron/training/diffusion/noise_utils.py
new file mode 100644
index 000000000..fe346263f
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/noise_utils.py
@@ -0,0 +1,67 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Noise application utilities for diffusion training.
+
+This module provides pure functions for applying noise according to
+different diffusion forward processes. All functions are stateless
+and work with tensors of any shape.
+
+Supported processes:
+    - Flow Matching: Linear interpolation between clean and noise
+    - DDPM: Variance-preserving noise schedule
+
+These utilities encapsulate the mathematical formulas for the forward
+diffusion process, making the code self-documenting and reusable
+across different model architectures (2D images, 3D video, etc.).
+"""
+
+from torch import Tensor
+
+
+def apply_flow_matching_noise(
+    clean_latents: Tensor,
+    noise: Tensor,
+    sigma: Tensor,
+) -> Tensor:
+    """
+    Apply noise using flow matching forward process.
+
+    Flow matching uses a simple linear interpolation between clean
+    latents and noise, controlled by the sigma parameter:
+
+    Formula: noisy = (1 - sigma) * clean + sigma * noise
+
+    This formulation ensures:
+    - When sigma=0: noisy = clean (no noise)
+    - When sigma=1: noisy = noise (pure noise)
+    - Linear interpolation in between
+
+    Args:
+        clean_latents: Clean latents [any shape]
+        noise: Sampled noise [same shape as clean_latents]
+        sigma: Noise schedule values [broadcast compatible]
+               Should be in range [0, 1]
+
+    Returns:
+        Noisy latents [same shape as clean_latents]
+
+    Reference:
+        Flow Matching for Generative Modeling
+        https://arxiv.org/abs/2210.02747
+
+    Example:
+        >>> clean = torch.randn(2, 16, 64, 64)
+        >>> noise = torch.randn(2, 16, 64, 64)
+        >>> sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1)
+        >>> noisy = apply_flow_matching_noise(clean, noise, sigma)
+        >>> # sigma[0]=0.3 means 30% noise, 70% clean for first sample
+    """
+    return (1.0 - sigma) * clean_latents + sigma * noise
+
+
+__all__ = [
+    "apply_flow_matching_noise",
+]
diff --git a/primus/backends/megatron/training/diffusion/schedulers/__init__.py b/primus/backends/megatron/training/diffusion/schedulers/__init__.py
new file mode 100644
index 000000000..2b917c927
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/schedulers/__init__.py
@@ -0,0 +1,15 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from primus.backends.megatron.training.diffusion.schedulers.base import BaseScheduler
+from primus.backends.megatron.training.diffusion.schedulers.flow_matching import (
+    FlowMatchEulerDiscreteScheduler,
+)
+
+__all__ = [
+    "BaseScheduler",
+    "FlowMatchEulerDiscreteScheduler",
+]
diff --git a/primus/backends/megatron/training/diffusion/schedulers/base.py b/primus/backends/megatron/training/diffusion/schedulers/base.py
new file mode 100644
index 000000000..85366d255
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/schedulers/base.py
@@ -0,0 +1,80 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Base scheduler interface for diffusion models.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Union
+
+import torch
+
+
+class BaseScheduler(ABC):
+    """
+    Abstract base class for diffusion schedulers.
+
+    Schedulers handle:
+        1. Noise schedule generation
+        2. Forward process (adding noise to samples)
+        3. Reverse process (denoising step)
+    """
+
+    @abstractmethod
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Add noise to sample (forward process).
+
+        Args:
+            sample: Clean sample
+            timestep: Current timestep
+            noise: Optional noise tensor (generated if None)
+
+        Returns:
+            Noisy sample
+        """
+        raise NotImplementedError("Subclasses must implement scale_noise()")
+
+    @abstractmethod
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None, **kwargs):
+        """
+        Set discrete timesteps for inference.
+
+        Args:
+            num_inference_steps: Number of diffusion steps
+            device: Device to place timesteps on
+            **kwargs: Additional scheduler-specific arguments
+        """
+        raise NotImplementedError("Subclasses must implement set_timesteps()")
+
+    @abstractmethod
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, ...]:
+        """
+        Perform one denoising step (reverse process).
+
+        Args:
+            model_output: Model prediction (noise or velocity)
+            timestep: Current timestep
+            sample: Current noisy sample
+            **kwargs: Additional scheduler-specific arguments
+
+        Returns:
+            Tuple containing denoised sample (and optionally other values)
+        """
+        raise NotImplementedError("Subclasses must implement step()")
+
+
+__all__ = ["BaseScheduler"]
diff --git a/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py b/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py
new file mode 100644
index 000000000..d56bcb3c9
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py
@@ -0,0 +1,323 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+#
+# Adapted from NeMo's flow matching scheduler.
+
+"""
+Flow Matching Euler Discrete Scheduler for Flux.
+
+This scheduler implements the Euler discrete sampling method for flow matching
+models like Flux. It supports both training and inference modes with optional
+dynamic timestep shifting for variable resolution.
+"""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from .base import BaseScheduler
+
+
+class FlowMatchEulerDiscreteScheduler(BaseScheduler):
+    """
+    Euler scheduler for flow matching diffusion models.
+
+    This scheduler is used for Flux and implements:
+        - Flow matching training objective
+        - Euler discrete sampling for inference
+        - Dynamic timestep shifting for variable resolution
+
+    Args:
+        num_train_timesteps: Number of diffusion steps for training (default: 1000)
+        shift: Base shift value for timestep schedule (default: 1.0)
+        use_dynamic_shifting: Whether to use dynamic shifting based on image resolution
+        base_shift: Base shift for dynamic shifting (default: 0.5)
+        max_shift: Maximum shift for dynamic shifting (default: 1.15)
+        base_image_seq_len: Base image sequence length for dynamic shifting (default: 256)
+        max_image_seq_len: Maximum image sequence length for dynamic shifting (default: 4096)
+
+    Reference:
+        - Flux paper: https://blackforestlabs.ai/flux-1-tools/
+        - Adapted from NeMo's flow matching scheduler
+    """
+
+    _compatibles = []
+    order = 1
+
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        base_shift: Optional[float] = 0.5,
+        max_shift: Optional[float] = 1.15,
+        base_image_seq_len: Optional[int] = 256,
+        max_image_seq_len: Optional[int] = 4096,
+    ):
+        """Initialize Flow Matching Euler Discrete Scheduler."""
+        # Generate initial timesteps
+        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        # Convert to sigmas (normalized timesteps in [0, 1])
+        sigmas = timesteps / num_train_timesteps
+
+        # Apply static shifting if not using dynamic shifting
+        if not use_dynamic_shifting:
+            # Shift formula: shift * sigma / (1 + (shift - 1) * sigma)
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        # Move sigmas to CPU to avoid too much CPU/GPU communication
+        self.sigmas = sigmas.to("cpu")
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+        # Store parameters
+        self.base_shift = base_shift
+        self.max_shift = max_shift
+        self.base_image_seq_len = base_image_seq_len
+        self.max_image_seq_len = max_image_seq_len
+        self.use_dynamic_shifting = use_dynamic_shifting
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+
+    @property
+    def step_index(self):
+        """The index counter for current timestep."""
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """The index for the first timestep."""
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Set the begin index for the scheduler.
+
+        Args:
+            begin_index: The begin index for the scheduler
+        """
+        self._begin_index = begin_index
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Add noise to sample using flow matching forward process.
+
+        Flow matching interpolation: x_t = sigma * noise + (1 - sigma) * sample
+
+        Args:
+            sample: Clean sample tensor
+            timestep: Current timestep(s)
+            noise: Optional noise tensor (generated if None)
+
+        Returns:
+            Noisy sample
+        """
+        # Generate noise if not provided
+        if noise is None:
+            noise = torch.randn_like(sample)
+
+        # Ensure sigmas and timesteps have same device and dtype as sample
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        # Handle MPS device (doesn't support float64)
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # Get step indices for given timesteps
+        # begin_index is None during training
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+        elif self.step_index is not None:
+            # add_noise called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise called before first denoising step (img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        # Get sigma values for these indices
+        sigma = sigmas[step_indices].flatten()
+
+        # Broadcast sigma to match sample shape
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        # Flow matching interpolation
+        noisy_sample = sigma * noise + (1.0 - sigma) * sample
+
+        return noisy_sample
+
+    def _sigma_to_t(self, sigma: float) -> float:
+        """Convert sigma to timestep."""
+        return sigma * self.num_train_timesteps
+
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        """
+        Apply dynamic time shifting.
+
+        Formula: exp(mu) / (exp(mu) + (1/t - 1)^sigma)
+
+        Args:
+            mu: Shift parameter (logarithmic)
+            sigma: Exponent parameter (fixed at 1.0 for Flux)
+            t: Timesteps to shift
+
+        Returns:
+            Shifted timesteps
+        """
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+    ):
+        """
+        Set discrete timesteps for inference.
+
+        Args:
+            num_inference_steps: Number of diffusion steps for inference
+            device: Device to place timesteps on
+            sigmas: Optional precomputed sigma values
+            mu: Shift parameter for dynamic shifting
+        """
+        # Validate dynamic shifting requirements
+        if self.use_dynamic_shifting and mu is None:
+            raise ValueError("Must pass a value for `mu` when `use_dynamic_shifting` is True")
+
+        # Generate sigmas if not provided
+        if sigmas is None:
+            self.num_inference_steps = num_inference_steps
+            timesteps = np.linspace(
+                self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+            )
+            sigmas = timesteps / self.num_train_timesteps
+
+        # Apply shifting (dynamic or static)
+        if self.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+
+        # Convert to tensors
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        timesteps = sigmas * self.num_train_timesteps
+
+        self.timesteps = timesteps.to(device=device)
+        # Append zero sigma at the end
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self, timestep: Union[float, torch.FloatTensor], schedule_timesteps: Optional[torch.Tensor] = None
+    ) -> int:
+        """
+        Get the index for a given timestep.
+
+        Args:
+            timestep: Timestep value
+            schedule_timesteps: Optional schedule to search in (uses self.timesteps if None)
+
+        Returns:
+            Index of the timestep
+        """
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index taken for the very first step is always the second index
+        # (or the last index if there is only 1). This ensures we don't accidentally
+        # skip a sigma when starting in the middle of the schedule.
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: Union[float, torch.FloatTensor]):
+        """Initialize step index from timestep."""
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        generator: Optional[torch.Generator] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Perform one Euler denoising step.
+
+        Args:
+            model_output: Model prediction (velocity field for flow matching)
+            timestep: Current timestep
+            sample: Current noisy sample
+            generator: Optional random number generator (unused)
+            **kwargs: Additional arguments (unused)
+
+        Returns:
+            Tuple containing the denoised sample
+        """
+        # Validate timestep type
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                "Passing integer indices as timesteps is not supported. "
+                "Pass one of scheduler.timesteps as a timestep."
+            )
+
+        # Initialize step index if needed
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues
+        sample = sample.to(torch.float32)
+
+        # Get current and next sigma
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        # Euler step: x_{t-1} = x_t + (sigma_next - sigma) * model_output
+        prev_sample = sample + (sigma_next - sigma) * model_output
+
+        # Cast back to model dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # Increment step index
+        self._step_index += 1
+
+        return (prev_sample,)
+
+    def __len__(self):
+        """Return number of training timesteps."""
+        return self.num_train_timesteps
+
+
+__all__ = ["FlowMatchEulerDiscreteScheduler"]
diff --git a/primus/backends/megatron/training/diffusion/timestep_sampling.py b/primus/backends/megatron/training/diffusion/timestep_sampling.py
new file mode 100644
index 000000000..35c335f71
--- /dev/null
+++ b/primus/backends/megatron/training/diffusion/timestep_sampling.py
@@ -0,0 +1,329 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Timestep sampling strategies for diffusion training.
+
+Training timestep sampling is a hyperparameter optimization technique
+that is separate from inference denoising algorithms. Different sampling
+strategies can improve training convergence and model quality.
+
+Key Insight:
+    - Training: Use logit-normal sampling (hyperparameter optimization)
+    - Inference: Use linear schedule (denoising algorithm requirement)
+
+These are fundamentally different concerns and should be separated.
+
+Supported strategies:
+    - LogitNormal: Better convergence (from SD3 paper)
+    - Uniform: Baseline approach for comparison
+    - Mode: Alternative from SD3 paper
+
+Reference:
+    Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1
+    Section 3.1: Timestep sampling
+"""
+
+import math
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+
+class TimestepSampler(ABC):
+    """
+    Abstract base class for training timestep sampling strategies.
+
+    Different sampling strategies can improve training convergence by
+    emphasizing certain timesteps where the model needs to learn more.
+
+    All samplers must implement the `sample()` method that returns
+    both timesteps and their corresponding sigma values.
+    """
+
+    @abstractmethod
+    def sample(
+        self,
+        batch_size: int,
+        device: torch.device,
+        scheduler,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Sample timesteps and compute sigmas for training.
+
+        Args:
+            batch_size: Number of samples in the batch
+            device: Target device for the tensors
+            scheduler: Scheduler instance (for accessing timesteps/sigmas)
+
+        Returns:
+            Tuple of (timesteps, sigmas):
+                - timesteps: Sampled timestep values [batch_size]
+                - sigmas: Corresponding noise schedule values [batch_size]
+
+        Example:
+            >>> sampler = LogitNormalSampler()
+            >>> timesteps, sigmas = sampler.sample(
+            ...     batch_size=32,
+            ...     device='cuda',
+            ...     scheduler=flow_scheduler
+            ... )
+        """
+
+
+class LogitNormalSampler(TimestepSampler):
+    """
+    Logit-normal timestep sampling (current approach, from SD3 paper).
+
+    This sampling strategy uses a logit-normal distribution to sample
+    timesteps. With default parameters (mean=0, std=1), the distribution
+    emphasizes mid-range timesteps (around t≈500). The mean and std
+    parameters control the distribution shape.
+
+    The distribution is created by:
+    1. Sample u from Normal(mean, std)
+    2. Apply sigmoid: u = sigmoid(u)
+    3. Map to timestep indices: indices = u * num_train_timesteps
+
+    This approach has been shown to improve training convergence
+    compared to uniform sampling.
+
+    Args:
+        mean: Mean of the normal distribution (default: 0.0)
+        std: Standard deviation of the normal distribution (default: 1.0)
+
+    Reference:
+        Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1
+        Section 3.1: "We use rf/lognorm(0.00,1.00)"
+
+    Example:
+        >>> sampler = LogitNormalSampler(mean=0.0, std=1.0)
+        >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler)
+        >>> # More samples near the middle (t~500) than the extremes
+    """
+
+    def __init__(self, mean: float = 0.0, std: float = 1.0):
+        self.mean = mean
+        self.std = std
+
+    def sample(
+        self,
+        batch_size: int,
+        device: torch.device,
+        scheduler,
+    ) -> Tuple[Tensor, Tensor]:
+        """Sample using logit-normal distribution."""
+        # Sample from Normal(mean, std)
+        u = torch.normal(
+            mean=self.mean,
+            std=self.std,
+            size=(batch_size,),
+            device="cpu",  # CPU to match scheduler.timesteps device
+        )
+
+        # Apply sigmoid to get logit-normal distribution
+        u = torch.sigmoid(u)
+
+        # Map to timestep indices
+        indices = (u * scheduler.num_train_timesteps).long()
+        indices = torch.clamp(indices, 0, scheduler.num_train_timesteps - 1)
+
+        # Get timesteps
+        timesteps = scheduler.timesteps[indices].to(device=device)
+
+        # Get corresponding sigmas
+        sigmas = scheduler.sigmas.to(device=device)
+        schedule_timesteps = scheduler.timesteps.to(device=device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+        sigma = sigmas[step_indices].flatten()
+
+        return timesteps, sigma
+
+
+class UniformSampler(TimestepSampler):
+    """
+    Uniform timestep sampling (baseline approach).
+
+    Samples timesteps uniformly from [0, num_train_timesteps), giving
+    equal probability to all timesteps. This is simpler but may converge
+    slower than logit-normal sampling.
+
+    Useful as a baseline for comparison with more sophisticated
+    sampling strategies.
+
+    Example:
+        >>> sampler = UniformSampler()
+        >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler)
+        >>> # Equal probability for all timesteps
+    """
+
+    def sample(
+        self,
+        batch_size: int,
+        device: torch.device,
+        scheduler,
+    ) -> Tuple[Tensor, Tensor]:
+        """Sample uniformly from all timesteps."""
+        indices = torch.randint(0, scheduler.num_train_timesteps, (batch_size,), device="cpu")
+
+        timesteps = scheduler.timesteps[indices].to(device=device)
+
+        sigmas = scheduler.sigmas.to(device=device)
+        schedule_timesteps = scheduler.timesteps.to(device=device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+        sigma = sigmas[step_indices].flatten()
+
+        return timesteps, sigma
+
+
+class DirectUniformSampler(TimestepSampler):
+    """
+    Direct uniform timestep sampling without scheduler indirection (NVIDIA MLPerf-compatible).
+
+    Samples timesteps continuously from [0, num_train_timesteps) using torch.rand,
+    bypassing the scheduler's discrete timestep/sigma lookup tables. This matches
+    NVIDIA's MLPerf Flux implementation exactly:
+
+        timesteps = torch.rand((batch_size,)) * num_train_timesteps
+        sigma = timesteps / num_train_timesteps  (i.e., sigma == normalized timestep)
+
+    Advantages over UniformSampler:
+        - Continuous distribution (no 1000-bin discretization)
+        - No Python loop for sigma lookup (avoids per-element .nonzero() calls)
+        - Produces tensors directly on target device
+
+    Note: Assumes shift=1.0 in the scheduler (no timestep shifting). With shift=1.0,
+    sigma == normalized_timestep, which matches the scheduler's linear mapping.
+
+    Example:
+        >>> sampler = DirectUniformSampler()
+        >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler)
+    """
+
+    def sample(
+        self,
+        batch_size: int,
+        device: torch.device,
+        scheduler,
+    ) -> Tuple[Tensor, Tensor]:
+        """Sample uniformly from [0, num_train_timesteps) without scheduler lookup."""
+        sigma = torch.rand((batch_size,), device=device)
+        timesteps = sigma * scheduler.num_train_timesteps
+        return timesteps, sigma
+
+
+class ModeSampler(TimestepSampler):
+    """
+    Mode-based sampling from SD3 paper (alternative to logit-normal).
+
+    This sampling strategy uses a mode-based distribution:
+
+    Formula: u = 1 - u - mode_scale * (cos(π*u/2)² - 1 + u)
+
+    Where u is initially sampled uniformly from [0, 1).
+
+    This creates a distribution that emphasizes certain timesteps
+    differently than logit-normal, potentially offering better
+    performance for some models.
+
+    Args:
+        mode_scale: Scaling factor for mode distribution (default: 1.29)
+                    Value from SD3 paper
+
+    Reference:
+        Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1
+        Section 3.1: Alternative sampling strategy
+
+    Example:
+        >>> sampler = ModeSampler(mode_scale=1.29)
+        >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler)
+    """
+
+    def __init__(self, mode_scale: float = 1.29):
+        self.mode_scale = mode_scale
+
+    def sample(
+        self,
+        batch_size: int,
+        device: torch.device,
+        scheduler,
+    ) -> Tuple[Tensor, Tensor]:
+        """Sample using mode distribution."""
+        u = torch.rand(size=(batch_size,), device="cpu")
+        u = 1 - u - self.mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
+
+        indices = (u * scheduler.num_train_timesteps).long()
+        indices = torch.clamp(indices, 0, scheduler.num_train_timesteps - 1)
+
+        timesteps = scheduler.timesteps[indices].to(device=device)
+
+        sigmas = scheduler.sigmas.to(device=device)
+        schedule_timesteps = scheduler.timesteps.to(device=device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+        sigma = sigmas[step_indices].flatten()
+
+        return timesteps, sigma
+
+
+def create_timestep_sampler(strategy: str = "logit_normal", **kwargs) -> TimestepSampler:
+    """
+    Factory function for creating timestep samplers.
+
+    This function provides a convenient way to create samplers
+    by name, allowing easy experimentation with different strategies.
+
+    Args:
+        strategy: Sampling strategy name
+                  Options: "logit_normal", "uniform", "direct_uniform", "mode"
+        **kwargs: Additional arguments for the sampler
+                  (e.g., mean/std for LogitNormal, mode_scale for Mode)
+
+    Returns:
+        TimestepSampler instance
+
+    Raises:
+        ValueError: If strategy is unknown
+
+    Example:
+        >>> # Default logit-normal
+        >>> sampler = create_timestep_sampler("logit_normal")
+        >>>
+        >>> # Custom parameters
+        >>> sampler = create_timestep_sampler(
+        ...     "logit_normal",
+        ...     mean=0.5,
+        ...     std=0.5
+        ... )
+        >>>
+        >>> # Uniform baseline
+        >>> sampler = create_timestep_sampler("uniform")
+        >>>
+        >>> # Mode sampling
+        >>> sampler = create_timestep_sampler("mode", mode_scale=1.5)
+    """
+    if strategy == "logit_normal":
+        return LogitNormalSampler(**kwargs)
+    elif strategy == "uniform":
+        return UniformSampler(**kwargs)
+    elif strategy == "direct_uniform":
+        return DirectUniformSampler(**kwargs)
+    elif strategy == "mode":
+        return ModeSampler(**kwargs)
+    else:
+        raise ValueError(
+            f"Unknown sampling strategy: {strategy}. "
+            f"Choose from: logit_normal, uniform, direct_uniform, mode"
+        )
+
+
+__all__ = [
+    "TimestepSampler",
+    "LogitNormalSampler",
+    "UniformSampler",
+    "DirectUniformSampler",
+    "ModeSampler",
+    "create_timestep_sampler",
+]
diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py
index 8d04bafcc..9ea10e438 100644
--- a/primus/backends/megatron/training/evaluator.py
+++ b/primus/backends/megatron/training/evaluator.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -7,6 +7,7 @@
 import time
 
 import torch
+from megatron.core import parallel_state
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel import get_forward_backward_func
@@ -16,7 +17,7 @@
 
 from primus.backends.megatron.training.global_vars import get_train_start_time
 from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def primus_evaluate(
@@ -136,14 +137,48 @@ def primus_evaluate(
                     log_rank_0("Exiting during evaluation, timelimit reached")
                     return None, None, True
 
-        # Compute final average loss across all eval iterations
+        # DP all-reduce for tuple-path (validation) metrics so that every
+        # rank sees the same globally-averaged loss.  Scalar/legacy metrics
+        # are NOT all-reduced, matching upstream Megatron's evaluate().
         total_loss_dict = {}
         if is_pipeline_stage_containing_loss():
+            from megatron.core import mpu
+
+            dp_group = mpu.get_data_parallel_group(with_context_parallel=True)
+            for key in total_loss_numerators.keys():
+                num = total_loss_numerators[key]
+                den = total_loss_denominators[key]
+                if isinstance(num, torch.Tensor) and isinstance(den, torch.Tensor):
+                    torch.distributed.all_reduce(num, group=dp_group)
+                    torch.distributed.all_reduce(den, group=dp_group)
+
             for key in total_loss_numerators.keys():
-                if total_loss_denominators[key] > 0:
-                    total_loss_dict[key] = total_loss_numerators[key] / total_loss_denominators[key]
+                # Reduce numerator/denominator across data-parallel ranks so the
+                # validation loss is a TRUE global average, identical on every rank.
+                # Without this, args._eval_val_loss stays a per-rank local value, and
+                # the target-eval-loss early stop (mlperf_pretrain_trainer.py) is then
+                # evaluated inconsistently: near the target one rank's local loss can
+                # dip <= target and exit train() alone while the others keep training,
+                # desyncing collectives (grad-norm all-reduce) -> NCCL hang at ~172k.
+                reduced = torch.tensor(
+                    [float(total_loss_numerators[key]), float(total_loss_denominators[key])],
+                    dtype=torch.float64,
+                    device="cuda",
+                )
+                torch.distributed.all_reduce(
+                    reduced,
+                    op=torch.distributed.ReduceOp.SUM,
+                    group=parallel_state.get_data_parallel_group(),
+                )
+                # Keep the result as a 0-dim tensor: downstream Megatron code
+                # (evaluate_and_print_results) and mlperf logging call .item() on it.
+                if reduced[1].item() > 0:
+                    total_loss_dict[key] = (reduced[0] / reduced[1]).to(torch.float32)
                 else:
-                    total_loss_dict[key] = 0.0
+                    total_loss_dict[key] = torch.zeros((), dtype=torch.float32, device="cuda")
+            if "lm loss" in total_loss_dict:
+                val = total_loss_dict["lm loss"]
+                args._eval_val_loss = val.item() if hasattr(val, "item") else float(val)
 
         collected_non_loss_data = None
         if non_loss_data_func is not None:
diff --git a/primus/backends/megatron/training/global_vars.py b/primus/backends/megatron/training/global_vars.py
index 5b2ae4825..e0fd44884 100644
--- a/primus/backends/megatron/training/global_vars.py
+++ b/primus/backends/megatron/training/global_vars.py
@@ -13,7 +13,7 @@
     collect_git_metadata,
     get_env_variables,
 )
-from primus.modules.module_utils import debug_rank_0
+from primus.core.utils.module_utils import debug_rank_0
 
 _GLOBAL_ARGS = None
 _GLOBAL_MLFLOW_WRITER = None
@@ -152,7 +152,7 @@ def _ensure_var_is_initialized(var, name):
 
 
 def _ensure_var_is_not_initialized(var, name):
-    """Make sure the input variable is not None."""
+    """Make sure the input variable is None (not yet initialized)."""
     assert var is None, "{} is already initialized.".format(name)
 
 
diff --git a/primus/backends/megatron/training/mlflow_artifacts.py b/primus/backends/megatron/training/mlflow_artifacts.py
index 2b6e32310..6e38f8d62 100644
--- a/primus/backends/megatron/training/mlflow_artifacts.py
+++ b/primus/backends/megatron/training/mlflow_artifacts.py
@@ -40,7 +40,7 @@
 import sys
 from typing import List, Optional
 
-from primus.modules.module_utils import log_rank_0, log_rank_last, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, log_rank_last, warning_rank_0
 
 # Pinned to immutable commit SHA for supply-chain safety (tags can be moved).
 # This corresponds to tag v0.4.0 in AMD-AGI/TraceLens.
diff --git a/primus/backends/megatron/training/tokenizer/tokenizer.py b/primus/backends/megatron/training/tokenizer/tokenizer.py
index 376996521..fc021dff3 100644
--- a/primus/backends/megatron/training/tokenizer/tokenizer.py
+++ b/primus/backends/megatron/training/tokenizer/tokenizer.py
@@ -26,11 +26,12 @@
         HuggingFaceTokenizer as _HuggingFaceTokenizer,
     )
 
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 CUSTOM_TOKENIZER_TYPES = {
     "DeepSeekV2Tokenizer",
     "DeepSeekV3Tokenizer",
+    "DeepSeekV4Tokenizer",
     "Llama2Tokenizer",
     "Llama3Tokenizer",
     "MixtralTokenizer",
diff --git a/primus/backends/megatron/training/utils.py b/primus/backends/megatron/training/utils.py
index 61a262f2a..4484f5307 100644
--- a/primus/backends/megatron/training/utils.py
+++ b/primus/backends/megatron/training/utils.py
@@ -10,6 +10,19 @@
 """General utilities."""
 import torch
 from megatron.core import mpu, parallel_state
+from megatron.training.global_vars import get_args
+
+
+def is_v_schedule_enabled(args=None):
+    """Return True when a V-shaped pipeline schedule (ZeroBubble-V / Primus-Pipe
+    zbv/v-half/v-min) is active for the current run."""
+    if args is None:
+        args = get_args()
+    return (
+        args.patch_zero_bubble
+        and args.enable_zero_bubble
+        and (args.zero_bubble_v_schedule or args.enable_1f1b_v)
+    ) or (args.pp_algorithm in ("zbv-formatted", "v-half", "v-min") and args.patch_primus_pipeline)
 
 
 def is_second_last_pipeline_stage():
@@ -31,8 +44,6 @@ def print_second_last_pipeline_stage(message):
 
 
 def is_pipeline_stage_containing_loss():
-    from primus.modules.trainer.megatron.utils import is_v_schedule_enabled
-
     if is_v_schedule_enabled():
         return mpu.is_pipeline_first_stage(ignore_virtual=True)
     else:
diff --git a/primus/backends/megatron_bridge/config_utils.py b/primus/backends/megatron_bridge/config_utils.py
index 987e628d1..e0eb0ee36 100644
--- a/primus/backends/megatron_bridge/config_utils.py
+++ b/primus/backends/megatron_bridge/config_utils.py
@@ -21,7 +21,7 @@
 from types import SimpleNamespace
 from typing import Any, Callable, Dict
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 DATASET_PATH_KEYS = (
     "data_paths",
@@ -329,8 +329,10 @@ def _resolve_recipe(recipe: str, flavor: str):
     Resolve a recipe module and function by searching multiple namespaces.
 
     Search order:
-        1. primus.backends.megatron_bridge.recipes.{recipe}  (Primus-side extensions)
-        2. megatron.bridge.recipes.{recipe}                   (upstream Megatron-Bridge)
+        1. Direct custom module path (e.g., primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom)
+        2. primus.backends.megatron_bridge.recipes.{recipe}  (Primus-side extensions)
+        3. megatron.bridge.recipes.{recipe}                   (upstream Megatron-Bridge)
+        4. recipe as a direct Python module path (fallback)
 
     Returns:
         Tuple of (module, full_module_path) for the first namespace that
@@ -339,6 +341,18 @@ def _resolve_recipe(recipe: str, flavor: str):
     Raises:
         AssertionError if the recipe cannot be found in any namespace.
     """
+    custom_prefixes = ("primus.backends.megatron_bridge.recipes.", "primus.")
+    if any(recipe.startswith(prefix) for prefix in custom_prefixes):
+        try:
+            module = importlib.import_module(recipe)
+        except ImportError as e:
+            assert False, f"Recipe loading failed: Cannot import custom recipe '{recipe}': {e}"
+        assert hasattr(
+            module, flavor
+        ), f"Recipe loading failed: Function '{flavor}' not found in custom recipe '{recipe}'"
+        log_rank_0(f"  ℹ️  Using custom recipe from: {recipe}")
+        return module, recipe
+
     search_prefixes = [
         "primus.backends.megatron_bridge.recipes",
         "megatron.bridge.recipes",
@@ -353,8 +367,20 @@ def _resolve_recipe(recipe: str, flavor: str):
         if hasattr(module, flavor):
             return module, full_module_path
 
+    # Fallback: try recipe as a direct Python module path
+    if "." in recipe and not recipe.startswith(("/", ".")):
+        try:
+            module = importlib.import_module(recipe)
+            if hasattr(module, flavor):
+                log_rank_0(f"  ℹ️  Using custom recipe from: {recipe}")
+                return module, recipe
+        except ImportError:
+            pass
+
     # Build a helpful error message listing all paths that were tried.
     tried = [f"{p}.{recipe}" for p in search_prefixes]
+    if "." in recipe:
+        tried.append(recipe)
     assert False, f"Recipe loading failed: Function '{flavor}' not found. " f"Searched modules: {tried}"
 
 
diff --git a/primus/backends/megatron_bridge/megatron_bridge_adapter.py b/primus/backends/megatron_bridge/megatron_bridge_adapter.py
index f0cdeecf4..972c91f6a 100644
--- a/primus/backends/megatron_bridge/megatron_bridge_adapter.py
+++ b/primus/backends/megatron_bridge/megatron_bridge_adapter.py
@@ -29,7 +29,7 @@
 )
 from primus.core.backend.backend_adapter import BackendAdapter
 from primus.core.backend.backend_registry import BackendRegistry
-from primus.modules.module_utils import log_dict_aligned, log_rank_0
+from primus.core.utils.module_utils import log_dict_aligned, log_rank_0
 
 
 def _install_modelopt_stub() -> bool:
diff --git a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py
index bf594e88e..545519161 100644
--- a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py
+++ b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py
@@ -19,7 +19,7 @@
 from primus.backends.megatron.training.global_vars import set_primus_global_variables
 from primus.core.patches import run_patches
 from primus.core.trainer.base_trainer import BaseTrainer
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class MegatronBridgeBaseTrainer(BaseTrainer):
@@ -36,20 +36,21 @@ class MegatronBridgeBaseTrainer(BaseTrainer):
         - Handle Megatron-Bridge specific initialization and setup
     """
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, **kwargs):
         """
         Initialize Megatron-Bridge base trainer.
 
         Args:
             backend_args: Megatron-Bridge configuration as SimpleNamespace
                          (from MegatronBridgeArgBuilder)
+            **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering.
         """
         log_rank_0("=" * 80)
         log_rank_0("Initializing MegatronBridgeBaseTrainer...")
         log_rank_0("=" * 80)
 
         # Initialize BaseTrainer
-        super().__init__(backend_args=backend_args)
+        super().__init__(backend_args=backend_args, **kwargs)
         set_primus_global_variables(self.backend_args)
 
         import primus.backends.megatron.patches  # noqa: F401
@@ -103,3 +104,39 @@ def detect_megatron_version(cls) -> str:
                 "Please ensure Megatron-LM is properly installed and "
                 "megatron.core.package_info is available."
             ) from e
+
+    def _apply_nested_overrides(self) -> None:
+        """Apply flat backend_args overrides to nested ConfigContainer fields.
+
+        ConfigContainer uses nested dataclasses (train, logger, checkpoint, etc.)
+        that cannot be reached by the flat _merge_dict_to_dataclass pass in
+        load_recipe_config. This bridges user-facing YAML keys (e.g.
+        ``log_interval: 99999``) to their nested targets.
+        """
+        args = self.backend_args
+        cfg = self.cfg_container
+
+        if hasattr(args, "log_interval"):
+            val = getattr(args, "log_interval")
+            if val is not None:
+                cfg.logger.log_interval = int(val)
+                log_rank_0(f"  ↳ Override logger.log_interval = {cfg.logger.log_interval}")
+
+        for key in ("eval_interval", "eval_iters"):
+            if hasattr(args, key):
+                val = getattr(args, key)
+                if val is not None:
+                    setattr(cfg.train, key, int(val))
+                    log_rank_0(f"  ↳ Override train.{key} = {val}")
+
+        if hasattr(args, "save_interval"):
+            val = getattr(args, "save_interval")
+            if val is not None:
+                cfg.checkpoint.save_interval = int(val)
+                log_rank_0(f"  ↳ Override checkpoint.save_interval = {val}")
+
+        if hasattr(args, "skip_save") and args.skip_save:
+            cfg.checkpoint.save_interval = 0
+            cfg.checkpoint.save = None
+            log_rank_0("  ↳ Override checkpoint.save_interval = 0 (skip periodic save)")
+            log_rank_0("  ↳ Override checkpoint.save = None (skip final save)")
diff --git a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py
index 1e95e0c3d..924607935 100644
--- a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py
+++ b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py
@@ -22,7 +22,7 @@
 from primus.backends.megatron_bridge.megatron_bridge_base_trainer import (
     MegatronBridgeBaseTrainer,
 )
-from primus.modules.module_utils import log_dict_aligned, log_rank_0
+from primus.core.utils.module_utils import log_dict_aligned, log_rank_0
 
 
 class MegatronBridgePosttrainTrainer(MegatronBridgeBaseTrainer):
@@ -47,15 +47,16 @@ class MegatronBridgePosttrainTrainer(MegatronBridgeBaseTrainer):
     # Task type identifier for logging
     TASK_TYPE = "Post-training (SFT/Instruction Tuning)"
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, **kwargs):
         """
         Initialize Megatron-Bridge posttrain trainer.
 
         Args:
             backend_args: Megatron-Bridge argument namespace (from MegatronBridgeArgBuilder)
+            **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering.
         """
         # Initialize MegatronBridgeBaseTrainer (which initializes BaseTrainer)
-        super().__init__(backend_args=backend_args)
+        super().__init__(backend_args=backend_args, **kwargs)
 
     def setup(self):
         """
@@ -78,6 +79,7 @@ def init(self):
         log_rank_0("Initializing Megatron-Bridge post-training components...")
 
         self.cfg_container = load_recipe_config(self.backend_args)
+        self._apply_nested_overrides()
 
         log_rank_0("Post-training initialization completed")
 
@@ -99,10 +101,13 @@ def train(self):
 
         try:
             # Execute post-training based on configuration
+            from megatron.bridge.training.config import runtime_config_update
             from megatron.bridge.training.finetune import finetune
             from megatron.bridge.training.vlm_step import forward_step
 
             # log_rank_0(f"ConfigContainer: {self.cfg_container}")
+            runtime_config_update(self.cfg_container)
+
             log_dict_aligned("ConfigContainer", self.cfg_container.to_dict())
             finetune(self.cfg_container, forward_step_func=forward_step)
 
diff --git a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py
index 9ee771f13..e88a4bcc4 100644
--- a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py
+++ b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py
@@ -20,7 +20,7 @@
 from primus.backends.megatron_bridge.megatron_bridge_base_trainer import (
     MegatronBridgeBaseTrainer,
 )
-from primus.modules.module_utils import log_dict_aligned, log_rank_0
+from primus.core.utils.module_utils import log_dict_aligned, log_rank_0
 
 
 class MegatronBridgePretrainTrainer(MegatronBridgeBaseTrainer):
@@ -39,14 +39,15 @@ class MegatronBridgePretrainTrainer(MegatronBridgeBaseTrainer):
         - Unified training workflow and patch management
     """
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, **kwargs):
         """
         Initialize Megatron-Bridge pretrain trainer.
 
         Args:
             backend_args: Megatron-Bridge argument namespace (from MegatronBridgeArgBuilder)
+            **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering.
         """
-        super().__init__(backend_args=backend_args)
+        super().__init__(backend_args=backend_args, **kwargs)
 
     def setup(self):
         """
@@ -71,45 +72,6 @@ def init(self):
 
         log_rank_0("Pre-training initialization completed")
 
-    def _apply_nested_overrides(self):
-        """Apply backend_args overrides to nested ConfigContainer fields.
-
-        ConfigContainer uses nested dataclasses (train, logger, checkpoint, etc.)
-        that cannot be reached by the flat _merge_dict_to_dataclass pass.
-        This method bridges user-facing flat keys to their nested targets.
-        """
-        args = self.backend_args
-        cfg = self.cfg_container
-
-        # logger.*
-        if hasattr(args, "log_interval"):
-            val = getattr(args, "log_interval")
-            if val is not None:
-                cfg.logger.log_interval = int(val)
-                log_rank_0(f"  ↳ Override logger.log_interval = {cfg.logger.log_interval}")
-
-        # train.*
-        for key in ("eval_interval", "eval_iters"):
-            if hasattr(args, key):
-                val = getattr(args, key)
-                if val is not None:
-                    setattr(cfg.train, key, int(val))
-                    log_rank_0(f"  ↳ Override train.{key} = {val}")
-
-        # checkpoint.*
-        if hasattr(args, "save_interval"):
-            val = getattr(args, "save_interval")
-            if val is not None:
-                cfg.checkpoint.save_interval = int(val)
-                log_rank_0(f"  ↳ Override checkpoint.save_interval = {val}")
-
-        # Skip final/in-training save entirely.
-        if hasattr(args, "skip_save") and args.skip_save:
-            cfg.checkpoint.save_interval = 0
-            cfg.checkpoint.save = None
-            log_rank_0("  ↳ Override checkpoint.save_interval = 0 (skip periodic save)")
-            log_rank_0("  ↳ Override checkpoint.save = None (skip final save)")
-
     def train(self):
         """
         Execute Megatron-Bridge pre-training.
diff --git a/primus/backends/megatron_bridge/mlperf_sft/__init__.py b/primus/backends/megatron_bridge/mlperf_sft/__init__.py
new file mode 100644
index 000000000..e22e4e473
--- /dev/null
+++ b/primus/backends/megatron_bridge/mlperf_sft/__init__.py
@@ -0,0 +1,19 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPerf ``:::MLLOG`` integration for Megatron-Bridge SFT / post-train workloads."""
+
+from primus.backends.megatron_bridge.mlperf_sft.mlperf_logger import (
+    MLPerfLogger,
+    ThroughputTimer,
+)
+from primus.backends.megatron_bridge.mlperf_sft.mlperf_sft import MLPerfSFTLogger
+
+__all__ = [
+    "MLPerfLogger",
+    "ThroughputTimer",
+    "MLPerfSFTLogger",
+]
diff --git a/primus/backends/megatron_bridge/mlperf_sft/mlperf_logger.py b/primus/backends/megatron_bridge/mlperf_sft/mlperf_logger.py
new file mode 100644
index 000000000..8b0bd0c81
--- /dev/null
+++ b/primus/backends/megatron_bridge/mlperf_sft/mlperf_logger.py
@@ -0,0 +1,159 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPERF Logging support for Megatron-Bridge SFT workloads."""
+
+import os
+import time
+from typing import Any, Dict, Optional
+
+
+class MLPerfLogger:
+    """Wrapper around mllog library with rank-aware logging."""
+
+    def __init__(self):
+        from mlperf_logging import mllog
+
+        self.mllogger = mllog.get_mllogger()
+        self._configured = False
+        self._rank = None
+        self._save_to_file = os.getenv("MLLOG_SAVE_TO_FILE", "1").lower() not in ("0", "false", "no")
+
+    def configure(self, filepath: str, args) -> Dict[str, Any]:
+        from mlperf_logging import mllog
+
+        if self._save_to_file:
+            mllog.config(filename=filepath, default_stack_offset=3)
+        else:
+            mllog.config(default_stack_offset=3)
+        self._configured = True
+        self._rank = self._get_rank()
+        return self.extract_mlperf_configs(args)
+
+    def _get_rank(self) -> int:
+        import torch
+
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return torch.distributed.get_rank()
+        return int(os.environ.get("RANK", 0))
+
+    def log_event_all_ranks(self, key: str, value: Any, metadata: Optional[Dict] = None):
+        self.mllogger.event(key=key, value=value, metadata=metadata or {})
+
+    def log_event(self, key: str, value: Any, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.event(key=key, value=value, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def log_start(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.start(key=key, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def log_end(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2):
+        if self._rank == 0:
+            self.mllogger.end(key=key, metadata=metadata or {}, stack_offset=stack_offset)
+
+    def extract_mlperf_configs(self, args) -> Dict[str, Any]:
+        """Extract MLPERF config parameters from Megatron args."""
+        from mlperf_logging.mllog import constants
+
+        data_parallel_size = getattr(args, "data_parallel_size", 1)
+        if data_parallel_size == 0:
+            data_parallel_size = 1
+        micro_batches = args.global_batch_size // (args.micro_batch_size * data_parallel_size)
+
+        eval_samples = args.eval_iters * args.global_batch_size if args.eval_iters > 0 else 1024
+
+        train_samples = getattr(args, "train_samples", None)
+        if not train_samples:
+            train_samples = args.train_iters * args.global_batch_size
+
+        optimizer_name = getattr(args, "optimizer", "adam")
+        if optimizer_name.lower() == "adam":
+            optimizer_name = "adamw"
+        else:
+            optimizer_name = optimizer_name.lower()
+
+        configs = {
+            constants.GLOBAL_BATCH_SIZE: args.global_batch_size,
+            constants.GRADIENT_ACCUMULATION_STEPS: micro_batches,
+            "max_sequence_length": args.seq_length,
+            constants.TRAIN_SAMPLES: train_samples,
+            constants.EVAL_SAMPLES: eval_samples,
+            constants.SEED: args.seed,
+            "init_checkpoint_step": args.iteration,
+            constants.OPT_NAME: optimizer_name,
+            constants.OPT_BASE_LR: args.lr,
+            constants.OPT_ADAMW_BETA_1: getattr(args, "adam_beta1", 0.9),
+            constants.OPT_ADAMW_BETA_2: getattr(args, "adam_beta2", 0.999),
+            constants.OPT_ADAMW_EPSILON: getattr(args, "adam_eps", 1e-8),
+            constants.OPT_ADAMW_WEIGHT_DECAY: args.weight_decay,
+            "opt_gradient_clip_norm": getattr(args, "clip_grad", 1.0),
+            "opt_end_learning_rate": args.min_lr,
+            "opt_learning_rate_warmup_steps": getattr(args, "lr_warmup_iters", 0),
+            "opt_learning_rate_decay_steps": args.lr_decay_iters if args.lr_decay_iters else args.train_iters,
+            "opt_learning_rate_decay_schedule": self._get_lr_schedule_name(args.lr_decay_style),
+            "max_steps": args.train_iters,
+            constants.SUBMISSION_BENCHMARK: os.getenv("MLLOG_SUBMISSION_BENCHMARK", ""),
+            constants.SUBMISSION_DIVISION: os.getenv("MLLOG_SUBMISSION_DIVISION", ""),
+            constants.SUBMISSION_STATUS: os.getenv("MLLOG_SUBMISSION_STATUS", ""),
+            constants.SUBMISSION_ORG: os.getenv("MLLOG_SUBMISSION_ORG", ""),
+            constants.SUBMISSION_PLATFORM: os.getenv("MLLOG_SUBMISSION_PLATFORM", ""),
+            constants.TENSOR_PARALLELISM: int(os.getenv("MLLOG_TENSOR_PARALLELISM", 1)),
+            constants.PIPELINE_PARALLELISM: int(os.getenv("MLLOG_PIPELINE_PARALLELISM", 1)),
+            constants.CONTEXT_PARALLELISM: int(os.getenv("MLLOG_CONTEXT_PARALLELISM", 1)),
+            constants.EXPERT_PARALLELISM: int(os.getenv("MLLOG_EXPERT_PARALLELISM", 1)),
+            constants.MICRO_BATCH_SIZE: int(os.getenv("MLLOG_MICRO_BATCH_SIZE", 1)),
+            constants.CONFIG_FILENAME: os.getenv("MLLOG_CONFIG_FILENAME", ""),
+            "lowest_numerical_precision_linear": os.getenv("MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR", ""),
+        }
+        return configs
+
+    def _get_lr_schedule_name(self, style: str) -> str:
+        mapping = {
+            "cosine": "cosine with linear warmup",
+            "linear": "linear",
+            "constant": "constant",
+        }
+        return mapping.get(style, style)
+
+
+class ThroughputTimer:
+    """Per-step accumulation timer for training throughput (samples/sec).
+
+    Measures only the time inside each training step, excluding eval time
+    and inter-step overhead (dataloader, callbacks, Python gaps).  This
+    matches the approach used by nemo's Timer class.
+    """
+
+    def __init__(self, global_batch_size: int):
+        self.gbs = global_batch_size
+        self._step_start_time = None
+        self._accumulated_time = 0.0
+        self._accumulated_samples = 0
+
+    def step_start(self):
+        """Mark the beginning of a training step."""
+        self._step_start_time = time.time()
+
+    def step_stop(self):
+        """Mark the end of a training step and accumulate elapsed time."""
+        if self._step_start_time is not None:
+            self._accumulated_time += time.time() - self._step_start_time
+            self._accumulated_samples += self.gbs
+            self._step_start_time = None
+
+    def get_throughput(self) -> Optional[float]:
+        """Return throughput for the accumulated steps and reset counters.
+
+        Returns:
+            Samples per second, or None if no steps were accumulated.
+        """
+        if self._accumulated_time <= 0:
+            return None
+        throughput = self._accumulated_samples / self._accumulated_time
+        self._accumulated_time = 0.0
+        self._accumulated_samples = 0
+        return throughput
diff --git a/primus/backends/megatron_bridge/mlperf_sft/mlperf_sft.py b/primus/backends/megatron_bridge/mlperf_sft/mlperf_sft.py
new file mode 100644
index 000000000..36fcd9067
--- /dev/null
+++ b/primus/backends/megatron_bridge/mlperf_sft/mlperf_sft.py
@@ -0,0 +1,340 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPerf logging support for Primus SFT (Supervised Fine-Tuning) workloads.
+
+Standalone lifecycle hooks for Megatron-Bridge training loops (e.g.
+``llama2_custom.py``).  Pretrain workloads use
+:class:`~primus.backends.megatron.mlperf.mlperf_pretrain_trainer.MLPerfMegatronPretrainTrainer`.
+"""
+
+import os
+import time
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+from .mlperf_logger import MLPerfLogger
+
+
+class MLPerfSFTLogger:
+    """Standalone MLPerf logger for SFT / fine-tuning workloads.
+
+    Composes :class:`MLPerfLogger` (rank-aware mllog wrapper) with an inline
+    wall-clock throughput timer.  All ``:::MLLOG`` emission is handled by
+    the lifecycle hooks listed below.
+
+    Typical call order from the training code::
+
+        logger = MLPerfSFTLogger(global_batch_size=8, micro_batch_size=1)
+        logger.log_cache_clear_and_init_start()
+        logger.log_init_params(config_dict)
+        # ... Megatron-Bridge init completes ...
+        logger.log_init_stop_run_start()
+        for step in training_loop:
+            logger.on_train_step(step, loss_dict, lr, consumed_samples)
+            if eval_step:
+                logger.on_eval_start(consumed_samples)
+                val_loss = run_eval()
+                should_stop = logger.on_eval_end(consumed_samples, val_loss)
+        logger.log_run_stop(consumed_samples)
+    """
+
+    def __init__(
+        self,
+        global_batch_size: int,
+        micro_batch_size: int,
+        target_eval_loss: Optional[float] = None,
+        train_loss_log_freq: Optional[int] = None,
+    ):
+        self._logger = MLPerfLogger()
+        self._gbs = global_batch_size
+        self._mbs = micro_batch_size
+        self._target_eval_loss = target_eval_loss or float(os.getenv("MLLOG_TARGET_EVAL_LOSS", "0.925"))
+        self._train_loss_log_freq = train_loss_log_freq or int(os.getenv("MLLOG_TRAIN_LOSS_LOG_FREQ", "10"))
+        self._block_start_time: Optional[float] = None
+        self._block_start_samples: int = 0
+        self._is_target_reached = False
+        self._run_start_time: Optional[float] = None
+
+    @property
+    def is_target_reached(self) -> bool:
+        return self._is_target_reached
+
+    def log_cache_clear_and_init_start(self) -> None:
+        """Emit ``CACHE_CLEAR`` (all ranks) and ``INIT_START`` (all ranks)."""
+        from mlperf_logging.mllog import constants
+
+        self._logger.log_event_all_ranks(key=constants.CACHE_CLEAR, value=True)
+        self._logger.log_event_all_ranks(key=constants.INIT_START, value=None)
+
+    def log_init_params(self, config: Dict[str, Any]) -> None:
+        """Log submission metadata and hyper-parameters before training starts."""
+        from mlperf_logging.mllog import constants
+
+        output_file = os.getenv("MLLOG_OUTPUT_FILE", "/results/mlperf_output.log")
+        from mlperf_logging import mllog
+
+        save_to_file = os.getenv("MLLOG_SAVE_TO_FILE", "1").lower() not in (
+            "0",
+            "false",
+            "no",
+        )
+        if save_to_file:
+            out_dir = os.path.dirname(output_file)
+            if out_dir:
+                os.makedirs(out_dir, exist_ok=True)
+            mllog.config(filename=output_file, default_stack_offset=3)
+        else:
+            mllog.config(default_stack_offset=3)
+        self._logger._configured = True
+        self._logger._rank = self._logger._get_rank()
+
+        self._logger.log_event(
+            key=constants.SUBMISSION_BENCHMARK,
+            value=os.getenv("MLLOG_SUBMISSION_BENCHMARK", "llama2_70b_lora"),
+        )
+        self._logger.log_event(
+            key=constants.SUBMISSION_ORG,
+            value=os.getenv("MLLOG_SUBMISSION_ORG", "AMD"),
+        )
+        self._logger.log_event(
+            key=constants.SUBMISSION_DIVISION,
+            value=os.getenv("MLLOG_SUBMISSION_DIVISION", "closed"),
+        )
+        self._logger.log_event(
+            key=constants.SUBMISSION_PLATFORM,
+            value=os.getenv("MLLOG_SUBMISSION_PLATFORM", "MI355X"),
+        )
+        self._logger.log_event(key=constants.SUBMISSION_STATUS, value="onprem")
+        self._logger.log_event(key="target_accuracy", value=self._target_eval_loss)
+
+        _KEY_MAP = {
+            "global_batch_size": constants.GLOBAL_BATCH_SIZE,
+            "train_samples": constants.TRAIN_SAMPLES,
+            "eval_samples": constants.EVAL_SAMPLES,
+            "gradient_accumulation_steps": constants.GRADIENT_ACCUMULATION_STEPS,
+            "seed": constants.SEED,
+            "opt_name": constants.OPT_NAME,
+            "opt_base_lr": constants.OPT_BASE_LR,
+            "opt_adamw_beta_1": constants.OPT_ADAMW_BETA_1,
+            "opt_adamw_beta_2": constants.OPT_ADAMW_BETA_2,
+            "opt_adamw_epsilon": constants.OPT_ADAMW_EPSILON,
+            "opt_weight_decay": constants.OPT_ADAMW_WEIGHT_DECAY,
+            "opt_gradient_clip_norm": "opt_gradient_clip_norm",
+            "opt_lr_warmup_factor": constants.OPT_LR_WARMUP_FACTOR,
+            "opt_lr_training_steps": constants.OPT_LR_TRAINING_STEPS,
+            "max_sequence_length": "max_sequence_length",
+            "max_steps": "max_steps",
+            "lora_rank": "lora_rank",
+            "lora_alpha": "lora_alpha",
+            "tensor_parallelism": constants.TENSOR_PARALLELISM,
+            "pipeline_parallelism": constants.PIPELINE_PARALLELISM,
+            "context_parallelism": constants.CONTEXT_PARALLELISM,
+            "expert_parallelism": constants.EXPERT_PARALLELISM,
+            "micro_batch_size": constants.MICRO_BATCH_SIZE,
+            "config_filename": constants.CONFIG_FILENAME,
+            "lowest_numerical_precision_linear": "lowest_numerical_precision_linear",
+        }
+        for cfg_key, mllog_key in _KEY_MAP.items():
+            if cfg_key in config and config[cfg_key] is not None and config[cfg_key] != "":
+                self._logger.log_event(key=mllog_key, value=config[cfg_key])
+
+    def log_init_stop_run_start(self) -> None:
+        """Emit ``INIT_STOP``, ``RUN_START``, ``EPOCH_START``, ``BLOCK_START``."""
+        from mlperf_logging.mllog import constants
+
+        self._logger.log_end(key=constants.INIT_STOP)
+        self._logger.log_start(key=constants.RUN_START)
+        self._run_start_time = time.time()
+        self._logger.log_start(key=constants.EPOCH_START, metadata={constants.SAMPLES_COUNT: 0})
+        self._logger.log_start(key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: 0})
+        self._block_start_time = time.time()
+        self._block_start_samples = 0
+
+    def on_train_step(
+        self,
+        step: int,
+        loss_dict: Dict[str, Any],
+        lr: Optional[float],
+        consumed_samples: int,
+    ) -> None:
+        if self._train_loss_log_freq == 0 or step % self._train_loss_log_freq != 0:
+            return
+
+        from mlperf_logging.mllog import constants
+
+        loss_value = loss_dict.get("lm loss")
+        if loss_value is None:
+            return
+        if isinstance(loss_value, (tuple, list)):
+            loss_value = loss_value[0]
+        if hasattr(loss_value, "item"):
+            loss_value = loss_value.item()
+
+        self._logger.log_event(
+            key="train_loss",
+            value=loss_value,
+            metadata={constants.SAMPLES_COUNT: consumed_samples, "lr": lr},
+        )
+
+    def on_eval_start(self, consumed_samples: int) -> None:
+        from mlperf_logging.mllog import constants
+
+        self._logger.log_end(
+            key=constants.BLOCK_STOP,
+            metadata={constants.SAMPLES_COUNT: consumed_samples},
+        )
+        self._logger.log_start(
+            key=constants.EVAL_START,
+            metadata={constants.SAMPLES_COUNT: consumed_samples},
+        )
+
+    def on_eval_end(self, consumed_samples: int, eval_loss: float) -> bool:
+        from mlperf_logging.mllog import constants
+
+        self._logger.log_event(
+            key=constants.EVAL_ACCURACY,
+            value=eval_loss,
+            metadata={constants.SAMPLES_COUNT: consumed_samples},
+        )
+
+        if (
+            self._target_eval_loss > 0.0
+            and eval_loss <= self._target_eval_loss
+            and not self._is_target_reached
+        ):
+            self._is_target_reached = True
+
+        self._logger.log_end(
+            key=constants.EVAL_STOP,
+            metadata={constants.SAMPLES_COUNT: consumed_samples},
+        )
+
+        if not self._is_target_reached:
+            self._logger.log_start(
+                key=constants.BLOCK_START,
+                metadata={constants.SAMPLES_COUNT: consumed_samples},
+            )
+            self._block_start_time = time.time()
+            self._block_start_samples = consumed_samples
+
+        return self._is_target_reached
+
+    def log_run_stop(self, consumed_samples: int) -> None:
+        from mlperf_logging.mllog import constants
+
+        self._logger.log_end(
+            key=constants.EPOCH_STOP,
+            metadata={constants.SAMPLES_COUNT: consumed_samples},
+        )
+
+        if not self._is_target_reached:
+            self._logger.log_end(
+                key=constants.RUN_STOP,
+                metadata={
+                    constants.SAMPLES_COUNT: consumed_samples,
+                    constants.STATUS: constants.ABORTED,
+                },
+            )
+        else:
+            self._logger.log_end(
+                key=constants.RUN_STOP,
+                metadata={
+                    constants.SAMPLES_COUNT: consumed_samples,
+                    constants.STATUS: constants.SUCCESS,
+                },
+            )
+
+        if self._run_start_time is not None:
+            duration = time.time() - self._run_start_time
+            duration_minutes = duration / 60.0
+            overall_throughput = consumed_samples / duration if duration > 0 else 0.0
+            self._logger.log_event(
+                key="run_duration",
+                value=f"{round(duration, 2)}s -> {round(duration_minutes, 2)} minutes",
+                metadata={"samples": consumed_samples},
+            )
+            self._logger.log_event(
+                key="throughput",
+                value=round(overall_throughput, 2),
+                metadata={"samples": consumed_samples},
+            )
+
+    @staticmethod
+    def extract_sft_configs(
+        train_gbs: int,
+        train_mbs: int,
+        train_iters: int,
+        eval_iters: int,
+        seq_length: int,
+        seed: int,
+        lr: float,
+        weight_decay: float,
+        clip_grad: float,
+        lr_warmup_iters: int,
+        adam_beta1: float = 0.9,
+        adam_beta2: float = 0.999,
+        adam_eps: float = 1e-8,
+        lora_rank: int = 16,
+        lora_alpha: int = 32,
+        data_root: str = "/data",
+        data_parallel_size: Optional[int] = None,
+        tensor_model_parallel_size: int = 1,
+        pipeline_model_parallel_size: int = 1,
+        context_parallel_size: int = 1,
+        config_filename: str = "",
+        lowest_numerical_precision_linear: str = "",
+    ) -> Dict[str, Any]:
+        train_samples = int(np.load(os.path.join(data_root, "train.npy"), allow_pickle=True).shape[0])
+        eval_samples = int(np.load(os.path.join(data_root, "validation.npy"), allow_pickle=True).shape[0])
+
+        if data_parallel_size is not None and data_parallel_size > 0:
+            dp_size = data_parallel_size
+        else:
+            dp_size = 1
+            try:
+                import torch
+
+                if torch.distributed.is_available() and torch.distributed.is_initialized():
+                    from megatron.core import parallel_state
+
+                    dp_size = parallel_state.get_data_parallel_world_size()
+            except Exception:
+                pass
+            if dp_size == 0:
+                dp_size = 1
+        grad_accum_steps = train_gbs // (train_mbs * dp_size)
+
+        warmup_factor = lr_warmup_iters / train_iters if train_iters > 0 else 0.0
+
+        return {
+            "global_batch_size": train_gbs,
+            "train_samples": train_samples,
+            "eval_samples": eval_samples,
+            "gradient_accumulation_steps": grad_accum_steps,
+            "seed": seed,
+            "max_sequence_length": seq_length,
+            "max_steps": train_iters,
+            "opt_name": "adamw",
+            "opt_base_lr": lr,
+            "opt_weight_decay": weight_decay,
+            "opt_gradient_clip_norm": clip_grad,
+            "opt_lr_warmup_factor": warmup_factor,
+            "opt_lr_training_steps": train_iters,
+            "opt_adamw_beta_1": adam_beta1,
+            "opt_adamw_beta_2": adam_beta2,
+            "opt_adamw_epsilon": adam_eps,
+            "lora_rank": lora_rank,
+            "lora_alpha": lora_alpha,
+            "tensor_parallelism": tensor_model_parallel_size,
+            "pipeline_parallelism": pipeline_model_parallel_size,
+            "context_parallelism": context_parallel_size,
+            "expert_parallelism": 1,
+            "micro_batch_size": train_mbs,
+            "config_filename": config_filename,
+            "lowest_numerical_precision_linear": lowest_numerical_precision_linear,
+        }
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py
new file mode 100644
index 000000000..b0d190202
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Llama2-70B LoRA MLPerf overrides under ``megatron_bridge/patches/mlperf_llama2_70b/``."""
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py
new file mode 100644
index 000000000..e42e91514
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py
@@ -0,0 +1,487 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Runtime Megatron-Bridge patches for MLPerf Llama2-70B LoRA.
+
+Replaces the former git patches under ``third_party/Megatron-Bridge`` without
+writing to ``third_party/Megatron-Bridge``.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Optional
+
+import torch
+
+from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.conditions import (
+    is_llama2_70b_mlperf,
+)
+from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.resettable_data_iterator import (
+    ResettableDataIterator,
+)
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCHED_ATTR = "_primus_mlperf_llama2_70b_patched"
+
+
+def _mark_patched(obj: Any) -> None:
+    setattr(obj, _PATCHED_ATTR, True)
+
+
+def _already_patched(obj: Any) -> bool:
+    return bool(getattr(obj, _PATCHED_ATTR, False))
+
+
+@register_patch(
+    "mlperf_llama2_70b.bridge.data_patches",
+    backend="megatron_bridge",
+    phase="before_train",
+    condition=is_llama2_70b_mlperf,
+    description="MLPerf data/eval/sampler overrides for Megatron-Bridge",
+)
+def patch_bridge_data(ctx: PatchContext) -> None:
+    import megatron.bridge.data.loaders as loaders
+    import megatron.bridge.data.samplers as samplers
+    import megatron.bridge.training.eval as bridge_eval
+    from megatron.core.rerun_state_machine import RerunDataIterator
+
+    loaders.ResettableDataIterator = ResettableDataIterator
+
+    import megatron.bridge.peft.lora as bridge_lora
+
+    from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.lora import (
+        LoRA,
+        VLMLoRA,
+    )
+
+    bridge_lora.LoRA = LoRA
+    bridge_lora.VLMLoRA = VLMLoRA
+
+    if not _already_patched(loaders.get_train_valid_test_num_samples):
+        orig_num_samples = loaders.get_train_valid_test_num_samples
+
+        @functools.wraps(orig_num_samples)
+        def _mlperf_get_train_valid_test_num_samples(cfg):
+            if cfg.train.train_samples is not None:
+                train_samples = cfg.train.train_samples
+            else:
+                train_samples = cfg.train.train_iters * cfg.train.global_batch_size
+            eval_iters = cfg.train.eval_iters
+            test_iters = cfg.train.eval_iters
+            return (
+                train_samples,
+                eval_iters * cfg.train.global_batch_size,
+                test_iters * cfg.train.global_batch_size,
+            )
+
+        loaders.get_train_valid_test_num_samples = _mlperf_get_train_valid_test_num_samples
+        _mark_patched(_mlperf_get_train_valid_test_num_samples)
+
+    if not _already_patched(samplers.build_pretraining_data_loader):
+        orig_build_loader = samplers.build_pretraining_data_loader
+
+        @functools.wraps(orig_build_loader)
+        def _mlperf_build_pretraining_data_loader(
+            dataset,
+            consumed_samples,
+            dataloader_type,
+            micro_batch_size,
+            num_workers,
+            data_sharding,
+            worker_init_fn=None,
+            collate_fn=None,
+            pin_memory=True,
+            persistent_workers=False,
+            data_parallel_rank=0,
+            data_parallel_size=1,
+            drop_last=True,
+            global_batch_size=None,
+            eval_iters: Optional[int] = None,
+            name: str = "",
+        ):
+            if dataset is None:
+                return None
+
+            if name == "validation":
+                if eval_iters is None or global_batch_size is None:
+                    raise RuntimeError(
+                        "eval_iters and global_batch_size must be provided when creating "
+                        "a validation dataloader for MLPerf Llama2."
+                    )
+                eval_samples = eval_iters * global_batch_size
+                total_samples = min(len(dataset), eval_samples)
+                batch_sampler = samplers.MegatronPretrainingSampler(
+                    total_samples=total_samples,
+                    consumed_samples=0,
+                    micro_batch_size=micro_batch_size,
+                    data_parallel_rank=data_parallel_rank,
+                    data_parallel_size=data_parallel_size,
+                )
+            elif dataloader_type == "single":
+                batch_sampler = samplers.MegatronPretrainingSampler(
+                    total_samples=len(dataset),
+                    consumed_samples=consumed_samples,
+                    micro_batch_size=micro_batch_size,
+                    data_parallel_rank=data_parallel_rank,
+                    data_parallel_size=data_parallel_size,
+                    drop_last=drop_last,
+                )
+            elif dataloader_type == "cyclic":
+                batch_sampler = samplers.MegatronPretrainingRandomSampler(
+                    dataset,
+                    total_samples=len(dataset),
+                    consumed_samples=consumed_samples,
+                    micro_batch_size=micro_batch_size,
+                    data_parallel_rank=data_parallel_rank,
+                    data_parallel_size=data_parallel_size,
+                    data_sharding=data_sharding,
+                )
+            elif dataloader_type == "batch":
+                if global_batch_size is None:
+                    raise RuntimeError(
+                        "global_batch_size must be provided when using dataloader_type='batch'."
+                    )
+                batch_sampler = samplers.MegatronPretrainingBatchSampler(
+                    total_samples=len(dataset),
+                    consumed_samples=consumed_samples,
+                    micro_batch_size=micro_batch_size,
+                    global_batch_size=global_batch_size,
+                    data_parallel_rank=data_parallel_rank,
+                    data_parallel_size=data_parallel_size,
+                    drop_last=drop_last,
+                    pad_samples_to_global_batch_size=not drop_last,
+                )
+            elif dataloader_type == "external":
+                return dataset
+            else:
+                raise Exception(f"Unsupported dataloader_type: {dataloader_type}")
+
+            return torch.utils.data.DataLoader(
+                dataset,
+                batch_sampler=batch_sampler,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                persistent_workers=persistent_workers,
+                worker_init_fn=worker_init_fn,
+                collate_fn=collate_fn,
+            )
+
+        samplers.build_pretraining_data_loader = _mlperf_build_pretraining_data_loader
+        _mark_patched(_mlperf_build_pretraining_data_loader)
+
+    if not _already_patched(loaders.build_train_valid_test_data_loaders):
+        orig_build_loaders = loaders.build_train_valid_test_data_loaders
+
+        @functools.wraps(orig_build_loaders)
+        def _mlperf_build_train_valid_test_data_loaders(
+            cfg, train_state, build_train_valid_test_datasets_provider
+        ):
+            (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)
+            loaders.print_rank_0("> building train, validation, and test datasets ...")
+            train_ds, valid_ds, test_ds = loaders.build_train_valid_test_datasets(
+                cfg=cfg,
+                build_train_valid_test_datasets_provider=build_train_valid_test_datasets_provider,
+            )
+            exit_signal = cfg.train.exit_signal
+
+            from megatron.bridge.training.utils.sig_utils import (
+                DistributedSignalHandler,
+            )
+
+            def worker_init_fn(_):
+                DistributedSignalHandler(exit_signal).__enter__()
+
+            maybe_worker_init_fn = worker_init_fn if cfg.train.exit_signal_handler_for_dataloader else None
+
+            train_dataloader = samplers.build_pretraining_data_loader(
+                train_ds,
+                train_state.consumed_train_samples,
+                cfg.dataset.dataloader_type,
+                cfg.train.micro_batch_size,
+                cfg.dataset.num_workers,
+                cfg.dataset.data_sharding,
+                worker_init_fn=maybe_worker_init_fn,
+                collate_fn=train_ds.collate_fn if hasattr(train_ds, "collate_fn") else None,
+                pin_memory=cfg.dataset.pin_memory,
+                persistent_workers=cfg.dataset.persistent_workers,
+                data_parallel_rank=loaders.mpu.get_data_parallel_rank(),
+                data_parallel_size=loaders.mpu.get_data_parallel_world_size(),
+                global_batch_size=cfg.train.global_batch_size,
+                name="train",
+            )
+            if cfg.train.skip_train and cfg.train.eval_iters > 0:
+                valid_dataloader = samplers.build_pretraining_data_loader(
+                    valid_ds,
+                    0,
+                    cfg.dataset.dataloader_type,
+                    cfg.train.micro_batch_size,
+                    cfg.dataset.num_workers,
+                    cfg.dataset.data_sharding,
+                    worker_init_fn=maybe_worker_init_fn,
+                    collate_fn=valid_ds.collate_fn if hasattr(valid_ds, "collate_fn") else None,
+                    pin_memory=cfg.dataset.pin_memory,
+                    persistent_workers=cfg.dataset.persistent_workers,
+                    data_parallel_rank=loaders.mpu.get_data_parallel_rank(),
+                    data_parallel_size=loaders.mpu.get_data_parallel_world_size(),
+                    global_batch_size=cfg.train.global_batch_size,
+                    name="validation",
+                    eval_iters=cfg.train.eval_iters,
+                )
+            elif cfg.train.eval_iters > 0:
+                val_dataloader_type = (
+                    "cyclic"
+                    if isinstance(cfg.dataset, loaders.GPTDatasetConfig)
+                    else cfg.dataset.dataloader_type
+                )
+                valid_dataloader = samplers.build_pretraining_data_loader(
+                    valid_ds,
+                    train_state.consumed_valid_samples,
+                    val_dataloader_type,
+                    cfg.train.micro_batch_size,
+                    cfg.dataset.num_workers,
+                    cfg.dataset.data_sharding,
+                    worker_init_fn=maybe_worker_init_fn,
+                    collate_fn=valid_ds.collate_fn if hasattr(valid_ds, "collate_fn") else None,
+                    pin_memory=cfg.dataset.pin_memory,
+                    persistent_workers=cfg.dataset.persistent_workers,
+                    data_parallel_rank=loaders.mpu.get_data_parallel_rank(),
+                    data_parallel_size=loaders.mpu.get_data_parallel_world_size(),
+                    global_batch_size=cfg.train.global_batch_size,
+                    name="validation",
+                    eval_iters=cfg.train.eval_iters,
+                )
+
+            if cfg.train.eval_iters > 0:
+                test_dataloader = samplers.build_pretraining_data_loader(
+                    test_ds,
+                    0,
+                    cfg.dataset.dataloader_type,
+                    cfg.train.micro_batch_size,
+                    cfg.dataset.num_workers,
+                    cfg.dataset.data_sharding,
+                    worker_init_fn=maybe_worker_init_fn,
+                    collate_fn=test_ds.collate_fn if hasattr(test_ds, "collate_fn") else None,
+                    pin_memory=cfg.dataset.pin_memory,
+                    persistent_workers=cfg.dataset.persistent_workers,
+                    data_parallel_rank=loaders.mpu.get_data_parallel_rank(),
+                    data_parallel_size=loaders.mpu.get_data_parallel_world_size(),
+                    global_batch_size=cfg.train.global_batch_size,
+                    name="test",
+                )
+
+            do_train = train_dataloader is not None and cfg.train.train_iters > 0
+            do_valid = valid_dataloader is not None and cfg.train.eval_iters > 0
+            do_test = test_dataloader is not None and cfg.train.eval_iters > 0
+            flags = torch.tensor(
+                [int(do_train), int(do_valid), int(do_test)], dtype=torch.long, device="cuda"
+            )
+            torch.distributed.broadcast(flags, 0)
+            train_state.do_train = flags[0].item()
+            train_state.do_valid = flags[1].item()
+            train_state.do_test = flags[2].item()
+            return train_dataloader, valid_dataloader, test_dataloader
+
+        loaders.build_train_valid_test_data_loaders = _mlperf_build_train_valid_test_data_loaders
+        _mark_patched(_mlperf_build_train_valid_test_data_loaders)
+
+    if not _already_patched(loaders.build_train_valid_test_data_iterators):
+        orig_build_iterators = loaders.build_train_valid_test_data_iterators
+
+        @functools.wraps(orig_build_iterators)
+        def _mlperf_build_train_valid_test_data_iterators(
+            cfg, train_state, build_train_valid_test_datasets_provider
+        ):
+            train_dataloader, valid_dataloader, test_dataloader = loaders.build_train_valid_test_data_loaders(
+                cfg=cfg,
+                train_state=train_state,
+                build_train_valid_test_datasets_provider=build_train_valid_test_datasets_provider,
+            )
+            dl_type = cfg.dataset.dataloader_type
+            assert dl_type in ["single", "cyclic", "batch", "external"]
+
+            def _get_iterator(dataloader_type, dataloader):
+                if dataloader_type == "single":
+                    return RerunDataIterator(iter(dataloader))
+                if dataloader_type in ("cyclic", "batch"):
+                    return RerunDataIterator(iter(loaders.cyclic_iter(dataloader)))
+                if dataloader_type == "external":
+                    if isinstance(dataloader, list):
+                        return [RerunDataIterator(d) for d in dataloader]
+                    return RerunDataIterator(dataloader)
+                raise RuntimeError("unexpected dataloader type")
+
+            train_data_iterator = (
+                _get_iterator(dl_type, train_dataloader) if train_dataloader is not None else None
+            )
+            if valid_dataloader is not None:
+                valid_data_iterator = RerunDataIterator(ResettableDataIterator(valid_dataloader))
+            else:
+                valid_data_iterator = None
+            test_data_iterator = (
+                _get_iterator(dl_type, test_dataloader) if test_dataloader is not None else None
+            )
+            return train_data_iterator, valid_data_iterator, test_data_iterator
+
+        loaders.build_train_valid_test_data_iterators = _mlperf_build_train_valid_test_data_iterators
+        _mark_patched(_mlperf_build_train_valid_test_data_iterators)
+
+    def _reset_data_iterator(data_iterator):
+        if data_iterator is None:
+            return
+        if isinstance(data_iterator, list):
+            for it in data_iterator:
+                _reset_data_iterator(it)
+            return
+        if isinstance(data_iterator, RerunDataIterator):
+            inner = data_iterator.iterable
+            if isinstance(inner, ResettableDataIterator):
+                inner.reset()
+                data_iterator.saved_microbatches.clear()
+                data_iterator.replaying = False
+                data_iterator.replay_pos = 0
+        elif isinstance(data_iterator, ResettableDataIterator):
+            data_iterator.reset()
+
+    if not _already_patched(bridge_eval.evaluate):
+        orig_evaluate = bridge_eval.evaluate
+
+        @functools.wraps(orig_evaluate)
+        def _mlperf_evaluate(*args, **kwargs):
+            if len(args) >= 3:
+                _reset_data_iterator(args[2])
+            else:
+                _reset_data_iterator(kwargs.get("data_iterator"))
+            return orig_evaluate(*args, **kwargs)
+
+        bridge_eval.evaluate = _mlperf_evaluate
+        _mark_patched(_mlperf_evaluate)
+
+    log_rank_0("[Patch:mlperf_llama2_70b.bridge.data_patches] Megatron-Bridge data patches applied")
+
+
+@register_patch(
+    "mlperf_llama2_70b.bridge.sft_attention_mask",
+    backend="megatron_bridge",
+    phase="before_train",
+    condition=is_llama2_70b_mlperf,
+    description="Cache causal attention masks in GPTSFTDataset for MLPerf steady-state SFT",
+)
+def patch_sft_attention_mask(ctx: PatchContext) -> None:
+    from megatron.bridge.data.datasets import sft as sft_mod
+
+    if _already_patched(sft_mod.GPTSFTDataset._create_attention_mask):
+        return
+
+    orig_create_mask = sft_mod.GPTSFTDataset._create_attention_mask
+
+    @functools.wraps(orig_create_mask)
+    def _cached_create_attention_mask(self, max_length):
+        cache = getattr(self, "_attention_mask_cache", None)
+        if cache is None:
+            cache = {}
+            object.__setattr__(self, "_attention_mask_cache", cache)
+        cached = cache.get(max_length)
+        if cached is not None:
+            return cached
+        attention_mask = torch.tril(torch.ones((max_length, max_length))).unsqueeze(0)
+        attention_mask = attention_mask < 0.5
+        cache[max_length] = attention_mask
+        return attention_mask
+
+    sft_mod.GPTSFTDataset._create_attention_mask = _cached_create_attention_mask
+    _mark_patched(_cached_create_attention_mask)
+    log_rank_0("[Patch:mlperf_llama2_70b.bridge.sft_attention_mask] SFT attention-mask cache applied")
+
+
+@register_patch(
+    "mlperf_llama2_70b.bridge.training_log_nemo",
+    backend="megatron_bridge",
+    phase="before_train",
+    condition=is_llama2_70b_mlperf,
+    description="NeMo-style train_step timing support in Megatron-Bridge training_log",
+)
+def patch_training_log_nemo(ctx: PatchContext) -> None:
+    import megatron.bridge.training.utils.train_utils as train_utils
+    from megatron.bridge.training.utils import flop_utils
+    from megatron.bridge.utils.common_utils import get_world_size_safe, print_rank_0
+
+    original_training_log = train_utils.training_log
+    if original_training_log is None or _already_patched(original_training_log):
+        return
+
+    sig = inspect.signature(original_training_log)
+    if "nemo_elapsed_time_per_iter_sec" in sig.parameters:
+        log_rank_0(
+            "[Patch:mlperf_llama2_70b.bridge.training_log_nemo] "
+            "training_log already supports NeMo timing; skipping wrapper"
+        )
+        return
+
+    @functools.wraps(original_training_log)
+    def _training_log_with_nemo(
+        loss_dict,
+        total_loss_dict,
+        learning_rate,
+        decoupled_learning_rate,
+        loss_scale,
+        report_memory_flag,
+        skipped_iter,
+        grad_norm,
+        params_norm,
+        num_zeros_in_grad,
+        config,
+        global_state,
+        history_wct,
+        model,
+        log_max_attention_logit=None,
+        nemo_elapsed_time_per_iter_sec=None,
+    ):
+        result = original_training_log(
+            loss_dict,
+            total_loss_dict,
+            learning_rate,
+            decoupled_learning_rate,
+            loss_scale,
+            report_memory_flag,
+            skipped_iter,
+            grad_norm,
+            params_norm,
+            num_zeros_in_grad,
+            config,
+            global_state,
+            history_wct,
+            model,
+            log_max_attention_logit,
+        )
+
+        if (
+            nemo_elapsed_time_per_iter_sec is not None
+            and global_state.train_state.step % config.logger.log_interval == 0
+        ):
+            batch_size = config.train.global_batch_size
+            num_flops = flop_utils.num_floating_point_operations(config, batch_size)
+            per_gpu_tf = num_flops / nemo_elapsed_time_per_iter_sec / get_world_size_safe() / 1e12
+            elapsed = global_state.timers("interval-time").elapsed(barrier=True)
+            interval_ms = (elapsed / max(config.logger.log_interval, 1)) * 1000.0
+            print_rank_0(
+                f"[TFLOP/s basis: NeMo-style train_step wall clock] "
+                f"Step time: {nemo_elapsed_time_per_iter_sec:.2f}s | "
+                f"model TFLOP/s/GPU: {per_gpu_tf:.1f}"
+            )
+            print_rank_0(
+                f" NeMo-style time/iter (ms): {nemo_elapsed_time_per_iter_sec * 1000.0:.1f} |"
+                f" Megatron-Bridge interval-time/iter (ms): {interval_ms:.1f} |"
+            )
+
+        return result
+
+    train_utils.training_log = _training_log_with_nemo
+    _mark_patched(_training_log_with_nemo)
+    log_rank_0("[Patch:mlperf_llama2_70b.bridge.training_log_nemo] NeMo timing wrapper applied")
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py
new file mode 100644
index 000000000..15909073c
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py
@@ -0,0 +1,43 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from primus.core.patches import PatchContext, get_args
+
+_MLPERF_MODEL_TOKEN = "llama2_70b_lora_mxfp4"
+_MLPERF_FLAVOR_TOKEN = "llama2_70b_lora_mxfp4_config"
+_MLPERF_CONFIG_TOKEN = "llama2_70b_lora_mlperf_posttrain"
+
+
+def is_llama2_70b_mlperf(ctx: PatchContext) -> bool:
+    """Return True when the active run targets MLPerf Llama2-70B LoRA."""
+    model_name = str(ctx.model_name or "")
+    if _MLPERF_MODEL_TOKEN in model_name:
+        return True
+
+    try:
+        args = get_args(ctx)
+    except AssertionError:
+        args = None
+
+    if args is not None:
+        model = str(getattr(args, "model", "") or "")
+        if _MLPERF_MODEL_TOKEN in model:
+            return True
+
+        recipe = str(getattr(args, "recipe", "") or "")
+        flavor = str(getattr(args, "flavor", "") or "")
+        if _MLPERF_FLAVOR_TOKEN in flavor and (
+            "llama2_custom" in recipe or "recipes.mlperf_llama2_70b" in recipe
+        ):
+            return True
+
+    primus_config = ctx.extra.get("primus_config")
+    if primus_config is not None:
+        config_file = str(getattr(primus_config, "config_file", "") or "")
+        if _MLPERF_CONFIG_TOKEN in config_file:
+            return True
+
+    return False
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py
new file mode 100644
index 000000000..e89677678
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py
@@ -0,0 +1,188 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NeMo-stable LoRA for MLPerf Llama2-70B.
+
+Default ``use_te_fused_lora=False`` keeps unfused :class:`LoRALinear` adapters so
+loss matches the MLPerf reference path.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional
+
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+from megatron.bridge.peft.base import PEFT
+from megatron.bridge.peft.lora_layers import (
+    LinearAdapter,
+    LoRALinear,
+    LoRATopKRouter,
+    TEFusedLoRALinear,
+    TELinearAdapter,
+    patch_linear_module,
+)
+from megatron.bridge.peft.module_matcher import ModuleMatcher
+from megatron.bridge.peft.utils import (
+    ParallelLinearAdapter,
+    get_adapter_attributes_from_linear,
+    is_expert_linear,
+    wildcard_match,
+)
+from megatron.core import parallel_state
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.utils import unwrap_model
+
+logger = logging.getLogger(__name__)
+
+try:
+    import bitsandbytes
+
+    HAVE_BNB = True
+except ImportError:
+    HAVE_BNB = False
+
+
+def _te_fused_lora_allowed_for_module(
+    full_name: str,
+    module_name: Optional[str],
+    include_modules: Optional[List[str]],
+    exclude_modules: List[str],
+) -> bool:
+    """Return True if this FQN may use :class:`TEFusedLoRALinear`."""
+    for pattern in exclude_modules:
+        if module_name == pattern or wildcard_match(pattern, full_name):
+            return False
+    if include_modules is None:
+        return True
+    for pattern in include_modules:
+        if module_name == pattern or wildcard_match(pattern, full_name):
+            return True
+    return False
+
+
+@dataclass
+class LoRA(PEFT, ModuleMatcher):
+    """LoRA with explicit control over TE fused adapters (MLPerf defaults to unfused)."""
+
+    target_modules: List[str] = field(
+        default_factory=lambda: ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"]
+    )
+    dim: int = 32
+    alpha: int = 32
+    dropout: float = 0.0
+    dropout_position: Literal["pre", "post"] = "pre"
+    lora_A_init_method: str = "xavier"
+    lora_B_init_method: str = "zero"
+    a2a_experimental: bool = False
+    lora_dtype: torch.dtype = None
+    use_te_fused_lora: bool = False
+    te_fused_lora_include_modules: Optional[List[str]] = None
+    te_fused_lora_exclude_modules: List[str] = field(default_factory=list)
+
+    def transform(
+        self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None
+    ) -> nn.Module:
+        adapter_types = (LinearAdapter, LoRALinear, LoRATopKRouter, TELinearAdapter)
+        if isinstance(module, adapter_types):
+            return module
+
+        if (ans := self.match(module, name, prefix)) is not None:
+            (match, full_name) = ans
+            if isinstance(module, nn.Linear) or (module.__class__ == te.Linear):
+                if hasattr(module.weight.data, "_local_tensor") or (
+                    HAVE_BNB
+                    and getattr(module, "quant_state", None) is not None
+                    and module.quant_state.__class__ == bitsandbytes.functional.QuantState
+                ):
+                    lora_cls = patch_linear_module
+                elif module.__class__ == te.Linear:
+                    lora_cls = TELinearAdapter
+                else:
+                    lora_cls = LinearAdapter
+
+                return lora_cls(
+                    module,
+                    dim=self.dim,
+                    alpha=self.alpha,
+                    dropout=self.dropout,
+                    lora_A_init_method=self.lora_A_init_method,
+                    lora_dtype=self.lora_dtype,
+                )
+
+            is_expert = is_expert_linear(full_name)
+            attrs = get_adapter_attributes_from_linear(module, is_expert=is_expert)
+
+            enable_op_fuser = (
+                self.use_te_fused_lora
+                and hasattr(module, "config")
+                and getattr(module.config, "use_transformer_engine_op_fuser", False)
+                and parallel_state.get_tensor_model_parallel_world_size() == 1
+                and _te_fused_lora_allowed_for_module(
+                    full_name,
+                    name,
+                    self.te_fused_lora_include_modules,
+                    self.te_fused_lora_exclude_modules,
+                )
+            )
+
+            logging.info(f"Adding lora to: {full_name}")
+            adapter = ParallelLinearAdapter(
+                attrs.in_features,
+                attrs.out_features,
+                self.dim,
+                base_linear_name=full_name,
+                activation="identity",
+                norm_type=None,
+                column_init_method=self.lora_A_init_method,
+                row_init_method=self.lora_B_init_method,
+                gather_output=False,
+                input_is_parallel=attrs.input_is_parallel,
+                dropout=self.dropout,
+                dropout_position=self.dropout_position,
+                model_parallel_config=getattr(module, "config", None),
+                alpha=self.alpha,
+                is_expert=is_expert,
+                a2a_experimental=self.a2a_experimental,
+                disable_tensor_parallel_comm=attrs.disable_tensor_parallel_comm,
+                disable_sequence_parallel_comm=attrs.disable_sequence_parallel_comm,
+                base_linear_is_parallel=attrs.base_linear_is_parallel,
+            )
+            if isinstance(module, TopKRouter):
+                return LoRATopKRouter(module, adapter)
+            if enable_op_fuser:
+                return TEFusedLoRALinear(module, adapter)
+            return LoRALinear(module, adapter)
+        return module
+
+
+@dataclass
+class VLMLoRA(LoRA):
+    """VLM LoRA variant (re-exported for parity with upstream Megatron-Bridge)."""
+
+    freeze_vision_model: bool = True
+    freeze_vision_projection: bool = True
+    freeze_language_model: bool = True
+
+    def freeze_model(self, model: nn.Module, training: bool = True) -> None:
+        modules_to_freeze = []
+        model = unwrap_model(model)[0]
+        if hasattr(model, "llava_model"):
+            model = model.llava_model
+
+        if self.freeze_vision_model and hasattr(model, "vision_model"):
+            modules_to_freeze.append(model.vision_model)
+        if self.freeze_vision_projection and hasattr(model, "vision_projection"):
+            modules_to_freeze.append(model.vision_projection)
+        if self.freeze_language_model and hasattr(model, "language_model"):
+            modules_to_freeze.append(model.language_model)
+
+        for module in modules_to_freeze:
+            module.eval()
+            for param in module.parameters():
+                param.requires_grad = training is False
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py
new file mode 100644
index 000000000..997968c1d
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py
@@ -0,0 +1,193 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Runtime Megatron-LM patches for MLPerf Llama2-70B LoRA MXFP4.
+
+Important: never delegate to Primus-Turbo ``fp4_utils.get_fp4_recipe`` here —
+that function can return a ``(recipe, reason)`` tuple, which breaks TE when it
+expects a recipe object with ``.mxfp4()`` / ``.mxfp8()`` methods.
+"""
+
+from __future__ import annotations
+
+import enum
+import functools
+import os
+from typing import Any
+
+from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.conditions import (
+    is_llama2_70b_mlperf,
+)
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PATCHED_ATTR = "_primus_mlperf_llama2_70b_megatron_patched"
+_mxfp4_phase = True
+
+
+def _mark_patched(obj: Any) -> None:
+    setattr(obj, _PATCHED_ATTR, True)
+
+
+def _already_patched(obj: Any) -> bool:
+    return bool(getattr(obj, _PATCHED_ATTR, False))
+
+
+def is_mxfp4_phase() -> bool:
+    return _mxfp4_phase
+
+
+def set_mxfp4_phase(active: bool) -> None:
+    global _mxfp4_phase
+    _mxfp4_phase = active
+
+
+def _upstream_get_fp4_recipe_handles_mxfp4() -> bool:
+    """True only when Megatron's get_fp4_recipe already builds MXFP4BlockScaling."""
+    try:
+        import inspect
+
+        from megatron.core import fp4_utils
+
+        source = inspect.getsource(fp4_utils.get_fp4_recipe)
+        return "MXFP4BlockScaling" in source or "Fp4Recipe.mxfp4" in source
+    except Exception:
+        return False
+
+
+def _build_mxfp4_get_fp4_recipe(orig_get_fp4_recipe):
+    import transformer_engine.common.recipe
+    from megatron.core.enums import Fp4Recipe
+    from megatron.core.fp8_utils import _get_custom_recipe
+    from megatron.core.transformer.transformer_config import TransformerConfig
+    from megatron.core.utils import is_te_min_version
+
+    @functools.wraps(orig_get_fp4_recipe)
+    def _mlperf_get_fp4_recipe(config: TransformerConfig):
+        if config.fp4_recipe == Fp4Recipe.nvfp4:
+            if not is_te_min_version("2.7.0.dev0"):
+                raise ValueError("NVFP4BlockScaling requires TransformerEngine >= 2.7.0.dev0.")
+            fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+        elif config.fp4_recipe == Fp4Recipe.mxfp4:
+            if not is_te_min_version("2.8.0"):
+                raise ValueError("MXFP4BlockScaling requires TransformerEngine >= 2.8.0.")
+            fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling()
+            fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1"
+        elif config.fp4_recipe == Fp4Recipe.custom:
+            fp4_recipe = _get_custom_recipe(config.fp4_quantizer_factory)
+        else:
+            raise ValueError(f"Unsupported FP4 recipe: {config.fp4_recipe}. Supported: nvfp4, mxfp4, custom.")
+        return fp4_recipe
+
+    return _mlperf_get_fp4_recipe
+
+
+@register_patch(
+    "mlperf_llama2_70b.megatron.fp4_mxfp4",
+    backend="megatron",
+    phase="before_train",
+    condition=is_llama2_70b_mlperf,
+    description="MXFP4 recipe + phase tracking for MLPerf Llama2 (single recipe object)",
+)
+def patch_fp4_mxfp4(ctx: PatchContext) -> None:
+    from megatron.core import enums, fp4_utils
+
+    recipe_already_handles_mxfp4 = _upstream_get_fp4_recipe_handles_mxfp4()
+
+    if not hasattr(enums.Fp4Recipe, "mxfp4"):
+
+        class _MlperfFp4Recipe(str, enum.Enum):
+            nvfp4 = "nvfp4"
+            mxfp4 = "mxfp4"
+            custom = "custom"
+
+        enums.Fp4Recipe = _MlperfFp4Recipe
+        log_rank_0("[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] Added Fp4Recipe.mxfp4")
+
+    fp4_utils.is_mxfp4_phase = is_mxfp4_phase
+    fp4_utils.set_mxfp4_phase = set_mxfp4_phase
+    fp4_utils._mxfp4_phase = _mxfp4_phase
+
+    if recipe_already_handles_mxfp4:
+        log_rank_0(
+            "[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] "
+            "Upstream get_fp4_recipe already supports mxfp4; phase tracking only"
+        )
+        return
+
+    if _already_patched(fp4_utils.get_fp4_recipe):
+        log_rank_0(
+            "[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] " "get_fp4_recipe already patched for MLPerf mxfp4"
+        )
+        return
+
+    fp4_utils.get_fp4_recipe = _build_mxfp4_get_fp4_recipe(fp4_utils.get_fp4_recipe)
+    _mark_patched(fp4_utils.get_fp4_recipe)
+
+    log_rank_0("[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] MXFP4 get_fp4_recipe patch applied")
+
+
+@register_patch(
+    "mlperf_llama2_70b.megatron.te_swiglu",
+    backend="megatron",
+    phase="before_train",
+    condition=is_llama2_70b_mlperf,
+    description="Optional TE SwiGLU path when USE_TE_SWIGLU=1",
+)
+def patch_te_swiglu(ctx: PatchContext) -> None:
+    if os.getenv("USE_TE_SWIGLU", "0") != "1":
+        log_rank_0("[Patch:mlperf_llama2_70b.megatron.te_swiglu] USE_TE_SWIGLU!=1; skipping")
+        return
+
+    import transformer_engine.common.recipe
+    import transformer_engine.pytorch as te
+    import transformer_engine_torch as tex
+    from megatron.core.fusions import fused_bias_swiglu as swiglu_mod
+
+    if _already_patched(swiglu_mod.SwiGLUFunction.forward):
+        return
+
+    swiglu_mod.SwiGLUFunction.forward
+
+    @staticmethod
+    @swiglu_mod.nvtx_decorator()
+    def _te_swiglu_forward(ctx, input_tensor, fp8_input_store, cpu_offload_input):
+        import torch
+
+        ctx.fp8_input_store = fp8_input_store
+        ctx.ori_input_dtype = input_tensor.dtype
+        input_for_backward = input_tensor.to(torch.float8_e4m3fn) if fp8_input_store else input_tensor
+        if cpu_offload_input:
+            input_for_backward.activation_offloading = True
+        ctx.save_for_backward(input_for_backward)
+
+        swiglu_op = te.ops.SwiGLU()
+        if fp8_input_store:
+            recipe = transformer_engine.common.recipe.DelayedScaling(
+                fp8_format=transformer_engine.common.recipe.Format.E4M3,
+                amax_history_len=8,
+                amax_compute_algo="max",
+                margin=2,
+                interval=1,
+            )
+            with te.fp8_autocast(enabled=True, fp8_recipe=recipe):
+                return swiglu_op(input_tensor)
+        return swiglu_op(input_tensor)
+
+    @staticmethod
+    @swiglu_mod.nvtx_decorator()
+    def _te_swiglu_backward(ctx, grad_output):
+        pass
+
+        input_tensor = ctx.saved_tensors[0]
+        input_tensor = input_tensor.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input_tensor
+        return tex.dswiglu(grad_output, input_tensor, None), None, None
+
+    swiglu_mod.SwiGLUFunction.forward = _te_swiglu_forward
+    swiglu_mod.SwiGLUFunction.backward = _te_swiglu_backward
+    _mark_patched(_te_swiglu_forward)
+    log_rank_0("[Patch:mlperf_llama2_70b.megatron.te_swiglu] TE SwiGLU patches applied")
diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py
new file mode 100644
index 000000000..a5773379f
--- /dev/null
+++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py
@@ -0,0 +1,35 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Resettable validation iterator for deterministic MLPerf eval."""
+
+
+class ResettableDataIterator:
+    """Iterator wrapper that restarts from the beginning of the dataloader on reset().
+
+    Unlike cyclic_iter which continuously cycles, this iterator explicitly resets
+    to produce the exact same sequence of batches on each reset(). This guarantees
+    deterministic validation: every evaluation pass sees identical data in identical
+    order, regardless of how many evaluations have occurred.
+    """
+
+    def __init__(self, dataloader):
+        self._dataloader = dataloader
+        self._iterator = iter(dataloader)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._iterator = iter(self._dataloader)
+            return next(self._iterator)
+
+    def reset(self):
+        """Reset to the beginning of the dataloader for deterministic iteration."""
+        self._iterator = iter(self._dataloader)
diff --git a/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py b/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py
index 0e948a533..2e8d79326 100644
--- a/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py
+++ b/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py
@@ -36,7 +36,7 @@
 from typing import Optional
 
 from primus.core.patches import PatchContext, get_args, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 # Per-step actual sequence length captured from prepare_finetuning_batch.
 # Accessed from the same training-loop thread; the lock guards against
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py
new file mode 100644
index 000000000..31519705e
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""MLPerf Llama2-70B LoRA recipe and supporting modules."""
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py
new file mode 100644
index 000000000..25db669ce
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py
@@ -0,0 +1,224 @@
+# ---------------------------------------------------------------------------
+# Non-MLLOG log suppression for Primus Llama2 SFT training pipeline.
+#
+# Adapted from small_llm_pretraining/primus/src/_log_suppression.py for use
+# with the Llama2 70B LoRA SFT recipe (Megatron-Bridge + TransformerEngine
+# + AITER).
+#
+# The MLPerf reference logs (":::MLLOG ...") plus training timing/result
+# banners are the only lines we want on stdout.  All other framework output
+# (Megatron deprecations, TE warnings, AITER/hipblaslt/gloo C++ writes,
+# torch.distributed noise, Primus internal INFO, ...) is suppressed by
+# default to keep CI and submission logs clean.
+#
+# Control:
+#   MLLOG_VERBOSE_LOGS=1  -> restore the full verbose output
+#   MLLOG_VERBOSE_LOGS=0  -> quiet mode (default): only MLLOG + training
+#                            timing/result lines are emitted.
+#
+# Strategy:
+#   1. Export env vars that noisy libraries honour (AITER_LOG_LEVEL,
+#      PYTHONWARNINGS, TRANSFORMERS_VERBOSITY, ...).
+#   2. Raise Python ``logging`` levels for the logger names emitting the
+#      noise (megatron.*, transformer_engine.*, torch.distributed, ...).
+#   3. Install an FD-level line filter on stdout so native C++ writes
+#      (hipModuleLoad, hipblaslt latency, gloo peer-connect) are caught.
+#      Stderr is redirected to /dev/null in quiet mode.
+#
+# This module should be imported as early as possible in the recipe to
+# suppress logs emitted during model construction and training.
+# ---------------------------------------------------------------------------
+import logging as _logging
+import os as _os
+
+VERBOSE_LOGS = _os.environ.get("MLLOG_VERBOSE_LOGS", "0") == "1"
+
+QUIET_LOGGER_NAMES = (
+    "megatron",
+    "megatron.core.utils",
+    "megatron.core.rerun_state_machine",
+    "transformer_engine",
+    "transformer_engine.aiter_rope",
+    "torch.distributed",
+    "torch.distributed.c10d_logger",
+    "primus",
+    "primus.cli",
+    "primus.backends",
+    "primus.modules",
+)
+
+
+def reapply_quiet_logger_levels() -> None:
+    """Raise levels on noisy Python loggers.  Safe to call multiple times."""
+    for name in QUIET_LOGGER_NAMES:
+        _logging.getLogger(name).setLevel(_logging.ERROR)
+
+
+def _configure_env_and_loggers() -> None:
+    """Silence non-MLLOG sources controllable via env vars or logging."""
+    _os.environ.setdefault("AITER_LOG_LEVEL", "ERROR")
+    _os.environ.setdefault("AITER_LOG_MORE", "0")
+    _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+    _os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    _os.environ.setdefault("PYTHONWARNINGS", "ignore")
+
+    reapply_quiet_logger_levels()
+
+
+def _install_fd_level_fallback_filter() -> None:
+    """Last-resort line filter for logs that bypass Python ``logging``.
+
+    * FD 2 (stderr) -> left untouched unless MLLOG_SUPPRESS_STDERR=1 (/dev/null).
+    * FD 1 (stdout) -> pipe + reader thread that applies regex suppression.
+      ``:::MLLOG`` lines always pass through unconditionally.
+    """
+    import re
+    import sys
+    import threading
+
+    ansi_re = re.compile(r"\x1b\[[0-9;]*[ -/]*[@-~]")
+
+    suppress_patterns = tuple(
+        re.compile(p)
+        for p in (
+            # aiter C++ hipModuleLoad / hipModuleGetFunction banners
+            r"\[aiter\] hipModule",
+            r"\[aiter\] hipModuleGetFunction:",
+            r"^\[aiter\] import \[",
+            # hipblaslt latency warnings (ROCm)
+            r"Warning: Latency not found for MI_M=",
+            r"Returning latency value of 32 \(really slow\)",
+            r"Warning: Stream-K Data Parallel does not support GSU",
+            r"^libibverbs: Warning:",
+            r"^\[WARNING\] Field \"",
+            r"`torch_dtype` is deprecated",
+            r"^(?:\s*(?:BFloat8Float8_fnuz|Float8_fnuz|,\s*(?:MI_[MNK]|mi_input_type)=\d*|\d+)\s*)+\.?\s*$",
+            # Gloo C++ peer-connect banners
+            r"\[Gloo\] Rank \d+ is connected to ",
+            r"^Expected number of connected peer ranks is\s*:",
+            # PyTorch AccumulateGrad stream mismatch warning
+            r"AccumulateGrad node's stream does not match",
+            r"^\s*Variable\._execution_engine\.run_backward\(",
+            # torch.distributed c10d barrier warning
+            r"barrier\(\): using the device under current context",
+            # hipify preprocessed banner
+            r"^(?:\x1b\[[0-9;]*m)?Successfully preprocessed all matching files\.",
+            # Primus runner INFO/DEBUG banners (coloured)
+            r"^\[0;34m\[.*\]\s*\[INFO\]",
+            r"^\[0;32m\[.*\]\s*\[(INFO|SUCCESS)\]",
+            r"^\[1;33m\[.*\]\s*\[WARN\]",
+            r"^\[DEBUG\]",
+            # Primus rank-0 coloured INFO lines
+            r"^\[\[32m\d{8}\s+\d{2}:\d{2}:\d{2}\[0m\].*\[INFO\]",
+            # [MLPerf Train] framework startup banners
+            r"^\[MLPerf Train\]",
+            # Created path banners from multi-rank output directory creation
+            r"^Created path:",
+            # Orphan fragments from interleaved C++ writes
+            r"^\s*Success\s*$",
+            r"^\s*\d+\s*$",
+            # torchrun OMP_NUM_THREADS banner
+            r"Setting OMP_NUM_THREADS environment variable for each process",
+            r"^\*{5,}$",
+            r"^W\d{4}\s+\d{2}:\d{2}:\d{2}\.\d+\s+\d+\s+torch/distributed/run\.py",
+            # pip install noise
+            r"^\[notice\] A new release of pip",
+            r"^WARNING: The directory .* pip_cache",
+            r"^Requirement already satisfied:",
+            r"^Downloading\s",
+            r"^Installing collected packages:",
+            r"^Successfully installed\s",
+            r"^ERROR: pip's dependency resolver",
+            # Primus hook execution banners
+            r"^\[fix_aiter_asm_dir\]",
+            r"^\[cast_transpose_mxfp4",
+            r"^\[mxfp4-",
+            r"^\[te_fp4_debug\]",
+            r"^\[aiter-rope\]",
+            r"^\[megatron-te-swiglu\]",
+            r"^patching file\s",
+            r"^\[OK\]\s",
+            r"^\[\+\]\s",
+            # env.VAR= debug lines
+            r"^env\.\w+=",
+            r"^extra\.\w+=",
+        )
+    )
+
+    def _should_suppress(line: str) -> bool:
+        stripped = ansi_re.sub("", line)
+        if ":::MLLOG" in stripped:
+            return False
+        if "RESULT," in stripped:
+            return False
+        if not stripped.strip():
+            return True
+        for pat in suppress_patterns:
+            if pat.search(stripped):
+                return True
+        return False
+
+    def _start_reader(read_fd: int, out_fd: int) -> None:
+        def _run() -> None:
+            buf = b""
+            try:
+                while True:
+                    chunk = _os.read(read_fd, 4096)
+                    if not chunk:
+                        break
+                    buf += chunk
+                    while b"\n" in buf:
+                        raw, buf = buf.split(b"\n", 1)
+                        line = raw.decode("utf-8", errors="replace")
+                        if not _should_suppress(line):
+                            _os.write(out_fd, raw + b"\n")
+            except Exception:
+                pass
+            finally:
+                if buf:
+                    line = buf.decode("utf-8", errors="replace")
+                    if not _should_suppress(line):
+                        try:
+                            _os.write(out_fd, buf)
+                        except Exception:
+                            pass
+
+        threading.Thread(target=_run, daemon=True).start()
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    orig_stdout_fd = _os.dup(1)
+
+    if _os.environ.get("MLLOG_SUPPRESS_STDERR", "0") == "1":
+        devnull_fd = _os.open(_os.devnull, _os.O_WRONLY)
+        _os.dup2(devnull_fd, 2)
+        _os.close(devnull_fd)
+
+    stdout_r, stdout_w = _os.pipe()
+    _os.dup2(stdout_w, 1)
+    _os.close(stdout_w)
+
+    _start_reader(stdout_r, orig_stdout_fd)
+
+    sys.stdout = _os.fdopen(1, "w", buffering=1, closefd=False)
+    sys.stderr = _os.fdopen(2, "w", buffering=1, closefd=False)
+
+
+_INSTALLED = False
+
+
+def install() -> None:
+    """Install quiet-mode suppression exactly once per process."""
+    global _INSTALLED
+    if _INSTALLED:
+        return
+    _INSTALLED = True
+    if VERBOSE_LOGS:
+        return
+    _configure_env_and_loggers()
+    _install_fd_level_fallback_filter()
+
+
+# Auto-install on import.
+install()
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py
new file mode 100644
index 000000000..146ba3210
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py
@@ -0,0 +1,62 @@
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+
+_RECIPE_DIR = Path(__file__).resolve().parent
+if str(_RECIPE_DIR) not in sys.path:
+    sys.path.insert(0, str(_RECIPE_DIR))
+
+import numpy as np
+import pandas as pd
+from dataset_hash import hash_files
+
+
+def convert(data_dir, split):
+    df = pd.read_parquet(f"{data_dir}/{split}-00000-of-00001.parquet")
+    transformed_data = df.apply(lambda row: transform_row(row), axis=1).tolist()
+    np.save(f"{data_dir}/{split}", transformed_data)
+
+
+def transform_row(row):
+    return {
+        "input_ids": row["input_ids"],
+        "loss_mask": [int(x != -100) for x in row["labels"]],
+        "seq_start_id": [0],
+    }
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert dataset script")
+    parser.add_argument("--data_dir", type=str, default="/data", help="The directory of the data files")
+    args = parser.parse_args()
+
+    convert(args.data_dir, "train")
+    convert(args.data_dir, "validation")
+
+    subprocess.run(
+        f"rm {args.data_dir}/train-00000-of-00001.parquet {args.data_dir}/validation-00000-of-00001.parquet",
+        shell=True,
+        executable="/bin/bash",
+        check=True,
+    )
+    # Verify only the produced files, so other content in data_dir does not
+    # affect the hash.
+    directory_hash = hash_files([f"{args.data_dir}/train.npy", f"{args.data_dir}/validation.npy"])
+    print(f"Succesfully converted dataset with hash {directory_hash}")
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py
new file mode 100644
index 000000000..3b3e193aa
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+"""Create minimal metadata file for manually packed sequences."""
+
+import argparse
+import json
+from pathlib import Path
+
+
+def create_metadata(seq_length: int, output_path: str):
+    """Create minimal metadata for packed sequences.
+
+    Args:
+        seq_length: The sequence length of your packed data
+        output_path: Path to save the metadata JSON file
+    """
+    metadata = [
+        {
+            "max_samples_per_bin": 1,
+            "dataset_max_seqlen": seq_length,
+            "min_packed_seqlen": seq_length,
+        }
+    ]
+
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+
+    with open(output_path, "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    print(f"✓ Created metadata file: {output_path}")
+    print("  - max_samples_per_bin: 1")
+    print(f"  - dataset_max_seqlen: {seq_length}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Create packed metadata for MLPerf Llama2-70B")
+    parser.add_argument("seq_length", type=int, help="Packed sequence length (e.g. 8192)")
+    parser.add_argument("output_path", type=str, help="Path to packed_metadata.jsonl")
+    args = parser.parse_args()
+
+    create_metadata(args.seq_length, args.output_path)
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py
new file mode 100644
index 000000000..05d3b969c
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py
@@ -0,0 +1,52 @@
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import hashlib
+from concurrent.futures import ThreadPoolExecutor
+
+
+def hash_file_md5(file_path, chunk_size=4194304):  # Default chunk size 4MB.
+    """Hashes a single file using MD5 and returns the hex digest."""
+    md5_hash = hashlib.md5()
+    try:
+        with open(file_path, "rb") as f:
+            while chunk := f.read(chunk_size):
+                md5_hash.update(chunk)
+    except FileNotFoundError:
+        return None
+    return md5_hash.hexdigest()
+
+
+def hash_files(file_paths):
+    """Hashes an explicit list of files in parallel using MD5.
+
+    Only the given files are considered, so unrelated content in the same
+    directory (e.g. checkpoints) does not affect the result.
+    """
+    hashes = []
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(hash_file_md5, p) for p in file_paths]
+        for future in futures:
+            file_hash = future.result()
+            if file_hash:
+                hashes.append(file_hash)
+
+    # Combine the individual file hashes into a single hash
+    combined_hash = hashlib.md5()
+    for file_hash in sorted(hashes):  # Sort to maintain consistency
+        combined_hash.update(file_hash.encode())
+
+    return combined_hash.hexdigest()
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py
new file mode 100644
index 000000000..fe6473dde
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py
@@ -0,0 +1,60 @@
+# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import glob
+import subprocess
+import sys
+from pathlib import Path
+
+_RECIPE_DIR = Path(__file__).resolve().parent
+if str(_RECIPE_DIR) not in sys.path:
+    sys.path.insert(0, str(_RECIPE_DIR))
+
+from dataset_hash import hash_files
+from huggingface_hub import snapshot_download
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--data_dir", default="/data", type=str, help="Path to the dataset location")
+args = parser.parse_args()
+
+snapshot_download(
+    "regisss/scrolls_gov_report_preprocessed_mlperf_2",
+    revision="21ff1233ee3e87bc780ab719c755170148aba1cb",
+    allow_patterns="*.parquet",
+    local_dir=args.data_dir,
+    local_dir_use_symlinks=False,
+    max_workers=16,
+    repo_type="dataset",
+)
+
+# Move the dataset parquets up and remove only HF's own artifacts. Do NOT
+# blanket-delete data_dir: it may hold a sibling checkpoint (the old
+# `find <data_dir> ! -name '*.parquet' -exec rm -rf` wiped it).
+subprocess.run(
+    f"mv {args.data_dir}/data/*.parquet {args.data_dir}/ && rm -rf {args.data_dir}/data {args.data_dir}/.cache",
+    shell=True,
+    executable="/bin/bash",
+    check=True,
+)
+
+# Verify only the downloaded parquets, so other content in data_dir does not
+# affect the hash.
+directory_hash = hash_files(sorted(glob.glob(f"{args.data_dir}/*.parquet")))
+assert (
+    directory_hash == "682a5f40b790a56751bf8303554efc08"
+), f"Expected hash 682a5f40b790a56751bf8303554efc08, but got {directory_hash}"
+print(f"Succesfully downloaded and verified dataset with hash {directory_hash}")
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py
new file mode 100644
index 000000000..6073416f5
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py
@@ -0,0 +1,2059 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Custom Llama2 recipe for Primus.
+
+This is a custom recipe based on Megatron-Bridge's llama2.py recipe,
+but placed in Primus for easier customization and extension.
+"""
+
+import gc
+import os
+import sys
+import time
+from collections import deque
+from datetime import timedelta
+from typing import Any, Callable, List, Optional, Union
+
+import torch
+from megatron.core.full_cuda_graph import FullCudaGraphWrapper
+from megatron.core.num_microbatches_calculator import (
+    get_current_global_batch_size,
+    get_current_running_global_batch_size,
+    get_num_microbatches,
+)
+from megatron.core.optimizer import MegatronOptimizer
+from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
+from megatron.core.parallel_state import update_pg_timeout
+from megatron.core.pipeline_parallel import get_forward_backward_func
+from megatron.core.pipeline_parallel.utils import is_pp_last_stage
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.rerun_state_machine import (
+    RerunDataIterator,
+    RerunMode,
+    get_rerun_state_machine,
+)
+from megatron.core.transformer import MegatronModule
+from megatron.core.transformer.cuda_graphs import TECudaGraphHelper
+from megatron.core.transformer.enums import AttnBackend
+from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config
+from typing_extensions import TypedDict, Unpack
+
+from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import (  # noqa: F401, E402
+    _log_suppression,
+)
+from primus.core.utils.module_utils import log_rank_0 as _orig_log_rank_0
+from primus.core.utils.module_utils import log_rank_last as _orig_log_rank_last
+
+_log_suppression.reapply_quiet_logger_levels()
+
+_verbose_logging = os.environ.get("VERBOSE_TRAINING_LOG", "0") == "1"
+if _verbose_logging:
+    log_rank_0 = _orig_log_rank_0
+    log_rank_last = _orig_log_rank_last
+else:
+
+    def log_rank_0(*args, **kwargs):
+        pass
+
+    def log_rank_last(*args, **kwargs):
+        pass
+
+
+from megatron.bridge.utils import common_utils
+
+# Megatron-Bridge training_log prints iteration / TFLOP/s / loss via print_rank_0
+# (plain stdout). Only silence Primus log_rank_0 helpers — not Megatron metrics.
+_megatron_print_rank_0 = common_utils.print_rank_0
+_megatron_print_rank_last = common_utils.print_rank_last
+_megatron_warn_rank_0 = common_utils.warn_rank_0
+common_utils.print_rank_0 = _megatron_print_rank_0
+common_utils.print_rank_last = _megatron_print_rank_last
+if _verbose_logging:
+    common_utils.warn_rank_0 = _megatron_warn_rank_0
+else:
+    common_utils.warn_rank_0 = log_rank_0
+
+
+def _truthy_env(name: str, default: bool = False) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in ("1", "true", "yes", "on")
+
+
+def _log_training_gpu_mem(tag: str, memory_keys=None) -> None:
+    """Rank-0 CUDA memory snapshot when ``PRIMUS_LOG_GPU_MEM=1``."""
+    if not _truthy_env("PRIMUS_LOG_GPU_MEM", default=False):
+        return
+    try:
+        from megatron.bridge.training.utils.train_utils import report_memory
+
+        if not torch.cuda.is_available():
+            _orig_log_rank_0(f"[GPU mem] {tag} | CUDA not available")
+            return
+        torch.cuda.synchronize()
+        dev = torch.cuda.current_device()
+        alloc = torch.cuda.memory_allocated(dev) / (1024**3)
+        reserved = torch.cuda.memory_reserved(dev) / (1024**3)
+        peak_alloc = torch.cuda.max_memory_allocated(dev) / (1024**3)
+        peak_reserved = torch.cuda.max_memory_reserved(dev) / (1024**3)
+        _orig_log_rank_0(
+            f"[GPU mem] {tag} | "
+            f"allocated={alloc:.2f} GiB reserved={reserved:.2f} GiB "
+            f"max_alloc={peak_alloc:.2f} GiB max_reserved={peak_reserved:.2f} GiB"
+        )
+        mem = report_memory(memory_keys)
+        if mem:
+            detail = " | ".join(f"{k}={v}" for k, v in sorted(mem.items()))
+            _orig_log_rank_0(f"[GPU mem detail] {tag} | {detail}")
+    except Exception as exc:
+        _orig_log_rank_0(f"[GPU mem] {tag} | failed ({type(exc).__name__}: {exc})")
+
+
+from megatron.bridge import AutoBridge
+from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs
+from megatron.bridge.data.finetuning import prepare_finetuning_batch
+from megatron.bridge.data.iterator_utils import make_data_iterator_list
+from megatron.bridge.recipes.utils.finetune_utils import default_squad_config
+from megatron.bridge.recipes.utils.tokenizer_utils import (
+    DEFAULT_NULL_TOKENIZER_VOCAB_SIZE,
+)
+from megatron.bridge.training import fault_tolerance
+from megatron.bridge.training.checkpointing import maybe_finalize_async_save
+from megatron.bridge.training.comm_overlap import CommOverlapConfig
+from megatron.bridge.training.config import (
+    CheckpointConfig,
+    ConfigContainer,
+    DistributedDataParallelConfig,
+    FinetuningDatasetConfig,
+    LoggerConfig,
+    ProfilingConfig,
+    RerunStateMachineConfig,
+    RNGConfig,
+    StragglerDetectionConfig,
+    TensorInspectConfig,
+    TokenizerConfig,
+    TrainingConfig,
+)
+from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
+from megatron.bridge.training.mixed_precision import (
+    MixedPrecisionConfig,
+    bf16_mixed,
+    register,
+)
+from megatron.bridge.training.nvrx_straggler import safe_shutdown_nvrx_straggler_manager
+from megatron.bridge.training.profiling import (
+    handle_profiling_step,
+    handle_profiling_stop,
+    initialize_pytorch_profiler,
+    should_profile_rank,
+)
+from megatron.bridge.training.state import GlobalState
+from megatron.bridge.training.tensor_inspect import (
+    tensor_inspect_end_if_enabled,
+    tensor_inspect_step_if_enabled,
+)
+from megatron.bridge.training.utils import flop_utils
+from megatron.bridge.training.utils import train_utils as _megatron_train_utils
+from megatron.bridge.training.utils.pg_utils import get_pg_collection
+from megatron.bridge.training.utils.train_utils import (
+    calc_params_l2_norm,
+    prepare_forward_step_func,
+    training_log,
+)
+
+from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.lora import LoRA
+from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.resettable_data_iterator import (
+    ResettableDataIterator,
+)
+
+# train_utils binds print_rank_* at import time; rebind so training_log always uses Primus loggers even if
+# this module was imported after an earlier train_utils load.
+_megatron_train_utils.print_rank_0 = log_rank_0
+_megatron_train_utils.print_rank_last = log_rank_last
+from megatron.bridge.training.train import (
+    _delete_cuda_graphs,
+    disable_forward_pre_hook,
+    enable_forward_pre_hook,
+    maybe_check_weight_hash_across_dp_replicas,
+    maybe_report_stragglers,
+    maybe_run_manual_gc,
+    maybe_synchronize_training_step,
+    should_disable_forward_pre_hook,
+    train_step,
+)
+from megatron.bridge.utils.common_utils import is_last_rank
+
+# Importing this module installs a monkey-patch that swaps Megatron-Bridge's
+from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import (  # noqa: F401
+    nemo_loss as _nemo_loss,
+)
+
+MLPERF_TARGET_LOSS = 0.925
+
+# Sticky flag: once ``evaluate_and_print_results_custom`` observes an lm-loss
+# value strictly below ``MLPERF_TARGET_LOSS`` we flip this to True. Subsequent
+# calls to ``evaluate_and_print_results_custom`` / ``warmup_eval`` become
+# no-ops that return ``(should_exit=True, None)`` without running another full
+# validation pass. Prevents redundant evals from a resumed loop or any outer
+# re-entry path after the target is met.
+_TARGET_LOSS_REACHED: bool = False
+
+# ---------------------------------------------------------------------------
+# MLPerf logging singleton (initialised lazily in the recipe config function)
+# ---------------------------------------------------------------------------
+_sft_logger = None
+
+
+def _get_sft_logger():
+    return _sft_logger
+
+
+@register
+def bf16_with_fp8_hybrid() -> MixedPrecisionConfig:
+    """Create a MixedPrecisionConfig for mixed precision training using BF16 with MXFP8.
+
+    Returns:
+        MixedPrecisionConfig: Configuration for BF16 with MXFP8 mixed precision training
+    """
+    cfg = bf16_mixed()
+    cfg.fp8 = "hybrid"
+    cfg.fp8_recipe = "delayed"
+    cfg.fp8_amax_history_len = 4
+    cfg.fp8_amax_compute_algo = "most_recent"
+    cfg.fp8_param_gather = True
+    return cfg
+
+
+@register
+def bf16_with_mxfp4_mixed() -> MixedPrecisionConfig:
+    """BF16 + MXFP4 (e2m1) mixed precision, registered here so it resolves when
+    ``precision_config="bf16_with_mxfp4_mixed"`` is passed from YAML before
+    ``runtime_config_update`` runs.
+    """
+    cfg = bf16_mixed()
+    cfg.fp8 = None
+    cfg.fp4 = "e2m1"
+    cfg.fp4_recipe = "mxfp4"
+    cfg.fp8_recipe = "delayed"
+    cfg.fp8_amax_history_len = 4
+    cfg.fp8_amax_compute_algo = "most_recent"
+    cfg.fp8_reduce_amax = False
+    cfg.fp8_interval = 1
+    cfg.fp8_margin = 0
+    cfg.fp8_dot_product_attention = False
+    cfg.fp8_param_gather = False
+    cfg.grad_reduce_in_fp32 = False
+    return cfg
+
+
+class Timer:
+    def __init__(self, gbs):
+        self.start_time = None
+        self.stop_time = None
+        self.elapsed_time = 0
+        self.samples = 0
+        self.gbs = gbs
+        self.consumed_samples = 0
+
+    def start(self):
+        self.start_time = time.time()
+
+    def stop(self):
+        self.stop_time = time.time()
+        self.samples += self.gbs
+        self.consumed_samples += self.gbs
+        self.elapsed_time += self.stop_time - self.start_time
+
+    def get_throughput(self):
+        throughput = self.samples / self.elapsed_time
+        self.samples = 0
+        self.elapsed_time = 0
+        return throughput
+
+
+class Llama2CustomKwargs(TypedDict, total=False):
+    """Typed options accepted by custom Llama2 recipe helper functions."""
+
+    # Core identifiers
+    hf_path: str
+    dir: Optional[str]
+    name: str
+    # Dataset configuration
+    data_paths: Optional[List[str]]
+    data_args_path: Optional[str]
+    train_data_path: Optional[List[str]]
+    valid_data_path: Optional[List[str]]
+    test_data_path: Optional[List[str]]
+    per_split_data_args_path: Optional[str]
+    mock: bool
+    # Model configuration
+    tensor_model_parallel_size: int
+    pipeline_model_parallel_size: int
+    pipeline_dtype: Optional[torch.dtype]
+    virtual_pipeline_model_parallel_size: Optional[int]
+    context_parallel_size: int
+    sequence_parallel: bool
+    use_megatron_fsdp: bool
+    # Training hyperparameters
+    train_iters: int
+    global_batch_size: int
+    micro_batch_size: int
+    seq_length: int
+    lr: float
+    min_lr: float
+    lr_warmup_iters: int
+    lr_decay_iters: Optional[int]
+    eval_interval: int
+    save_interval: int
+    use_null_tokenizer: bool
+    # Precision / overlap configs
+    precision_config: Optional[Union[MixedPrecisionConfig, str]]
+    comm_overlap_config: Optional[CommOverlapConfig]
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.99
+    adam_eps: float = 1e-8
+    weight_decay: float = 0.0001
+    eval_iters: int = 22
+    clip_grad: float = 0.3
+    pretrained_checkpoint: str | None
+    packed_sequence: bool
+    packed_train_data_path: str | None
+    packed_val_data_path: str | None
+    packed_metadata_path: str | None
+    dataset_type: str
+    seed: int
+    check_for_nan_in_loss: bool
+    te_fused_lora_include_modules: Optional[List[str]]
+    te_fused_lora_exclude_modules: Optional[List[str]]
+    # Optional MXFP4-phase activation recompute (left None outside the MXFP4 recipe).
+    recompute_granularity: Optional[str]
+    recompute_method: Optional[str]
+    recompute_num_layers: Optional[int]
+    # TE attention backend ("flash", "fused", "unfused", "local", "auto"). None keeps Megatron's default.
+    attention_backend: Optional[str]
+
+
+def llama2_70b_lora_config(**user_kwargs: Unpack[Llama2CustomKwargs]) -> ConfigContainer:
+    """
+    Return a custom pre-training config for Llama-2 70B.
+
+    This is a custom variant that can be modified without changing Megatron-Bridge code.
+    See `_llama2_lora` for the full list of parameters.
+    """
+    recommended_kwargs: Llama2CustomKwargs = {
+        "hf_path": "meta-llama/Llama-2-70b-hf",
+        "tensor_model_parallel_size": 1,
+        "pipeline_model_parallel_size": 1,
+        "train_iters": 1000,
+        "global_batch_size": 8,
+        "micro_batch_size": 1,
+        "eval_interval": 48,
+        "eval_iters": 22,
+        "adam_beta1": 0.9,
+        "adam_beta2": 0.999,
+        "adam_eps": 1e-8,
+        "weight_decay": 0.0001,
+        "clip_grad": 0.3,
+    }
+    # Combine defaults with user kwargs; user values take precedence.
+    combined_kwargs: Llama2CustomKwargs = {**recommended_kwargs, **user_kwargs}
+    cfg = _llama2_lora(**combined_kwargs)
+
+    # --- MLPerf logging: init phase ---
+    if os.getenv("ENABLE_MLLOG", "0") == "1":
+        global _sft_logger
+        try:
+            from primus.backends.megatron_bridge.mlperf_sft import MLPerfSFTLogger
+
+            kw = combined_kwargs
+            gbs = kw.get("global_batch_size", 8)
+            mbs = kw.get("micro_batch_size", 1)
+            _sft_logger = MLPerfSFTLogger(
+                global_batch_size=gbs,
+                micro_batch_size=mbs,
+            )
+            _sft_logger.log_cache_clear_and_init_start()
+
+            data_root = os.getenv("DATA_PATH", "/data")
+            world_size = int(os.environ.get("WORLD_SIZE", "1"))
+            tp = kw.get("tensor_model_parallel_size", 1)
+            pp = kw.get("pipeline_model_parallel_size", 1)
+            dp_size = world_size // (tp * pp)
+            init_cfg = MLPerfSFTLogger.extract_sft_configs(
+                train_gbs=gbs,
+                train_mbs=mbs,
+                train_iters=kw.get("train_iters", 1000),
+                eval_iters=kw.get("eval_iters", 22),
+                seq_length=kw.get("seq_length", 8192),
+                seed=kw.get("seed", int(os.getenv("SEED", "1234"))),
+                lr=kw.get("lr", float(os.getenv("LR", "0.0004"))),
+                weight_decay=kw.get("weight_decay", 0.0001),
+                clip_grad=kw.get("clip_grad", 0.3),
+                lr_warmup_iters=kw.get("lr_warmup_iters", 0),
+                adam_beta1=kw.get("adam_beta1", 0.9),
+                adam_beta2=kw.get("adam_beta2", 0.999),
+                adam_eps=kw.get("adam_eps", 1e-8),
+                lora_rank=16,
+                lora_alpha=32,
+                data_root=data_root,
+                data_parallel_size=dp_size,
+                tensor_model_parallel_size=tp,
+                pipeline_model_parallel_size=pp,
+                context_parallel_size=int(os.getenv("MLLOG_CONTEXT_PARALLELISM", "1")),
+                config_filename=os.getenv("MLLOG_CONFIG_FILENAME", ""),
+                lowest_numerical_precision_linear=os.getenv(
+                    "MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR", "mxfp4"
+                ),
+            )
+            _sft_logger.log_init_params(init_cfg)
+        except ImportError as exc:
+            _orig_log_rank_0(
+                f"MLPerf SFT logger unavailable ({exc}) — install mlperf-logging "
+                "(pip install git+https://github.com/mlperf/logging.git)"
+            )
+            _sft_logger = None
+        except Exception as exc:
+            _orig_log_rank_0(f"MLPerf logging init failed ({type(exc).__name__}: {exc}) — disabled")
+            _sft_logger = None
+
+    return cfg
+
+
+def llama2_70b_lora_mxfp4_config(**user_kwargs: Unpack[Llama2CustomKwargs]) -> ConfigContainer:
+    """Llama-2 70B LoRA with MXFP4 mixed precision (``bf16_with_mxfp4_mixed``).
+
+    Matches MLPerf 6.0 NeMo MI355X FP4 defaults: full/block recompute over 8
+    layers; FusedAttention backend (CK / AOTriton on ROCm) instead of
+    FlashAttention (which is not validated for the MXFP4 path on MI355X).
+    All defaults are user-overridable via ``user_kwargs``.
+    """
+    mxfp4_defaults: Llama2CustomKwargs = {
+        "precision_config": "bf16_with_mxfp4_mixed",
+        "recompute_granularity": None,
+        "recompute_method": None,
+        "recompute_num_layers": None,
+        "attention_backend": "fused",
+    }
+    combined_kwargs: Llama2CustomKwargs = {**mxfp4_defaults, **user_kwargs}
+    return llama2_70b_lora_config(**combined_kwargs)
+
+
+def _llama2_lora(
+    hf_path: str,
+    dir: Optional[str] = None,
+    name: str = "default",
+    # Model configuration
+    tensor_model_parallel_size: int = 1,
+    pipeline_model_parallel_size: int = 1,
+    virtual_pipeline_model_parallel_size: Optional[int] = None,
+    context_parallel_size: int = 1,
+    sequence_parallel: bool = False,
+    # Training hyperparameters
+    train_iters: int = 1000,
+    global_batch_size: int = 8,
+    micro_batch_size: int = 1,
+    seq_length: int = 8192,
+    lr: float = 4e-4,
+    min_lr: float = 0.0,
+    lr_warmup_iters: int = 0,
+    lr_decay_iters: Optional[int] = None,
+    eval_interval: int = 48,
+    eval_iters: int = 22,
+    use_null_tokenizer: bool = False,
+    pretrained_checkpoint: str | None = None,
+    packed_sequence: bool = False,
+    packed_train_data_path: str | None = None,
+    packed_val_data_path: str | None = None,
+    packed_metadata_path: str | None = None,
+    dataset_type: str = "mlperf_dataset",
+    adam_beta1: float = 0.9,
+    adam_beta2: float = 0.999,
+    adam_eps: float = 1e-8,
+    weight_decay: float = 0.0001,
+    clip_grad: float = 0.3,
+    # Precision recipe
+    precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_with_mxfp4_mixed",
+    comm_overlap_config: Optional[CommOverlapConfig] = None,
+    seed: int = 1234,
+    check_for_nan_in_loss: bool = False,
+    te_fused_lora_include_modules: Optional[List[str]] = None,
+    te_fused_lora_exclude_modules: Optional[List[str]] = None,
+    recompute_granularity: Optional[str] = None,
+    recompute_method: Optional[str] = None,
+    recompute_num_layers: Optional[int] = None,
+    attention_backend: Optional[str] = None,
+) -> ConfigContainer:
+    """
+    Create a custom pre-training configuration for Llama2 models.
+
+    This is based on Megatron-Bridge's llama2 recipe but can be customized
+    for Primus-specific needs.
+
+    Args:
+        hf_path (str): HuggingFace model path (e.g., "meta-llama/Llama-2-70b-hf").
+        dir (Optional[str]): Base directory for saving logs and checkpoints.
+        name (str): Name of the pre-training run.
+        data_paths (Optional[List[str]]): List of paths to dataset files. If None, mock data will be used.
+        data_args_path (Optional[str]): Path to file containing data arguments.
+        train_data_path (Optional[List[str]]): List of training data paths.
+        valid_data_path (Optional[List[str]]): List of validation data paths.
+        test_data_path (Optional[List[str]]): List of test data paths.
+        per_split_data_args_path (Optional[str]): Path to JSON file with per-split data configuration.
+        mock (bool): Whether to use mock data. If True, ignores data_paths.
+        tensor_model_parallel_size (int): Degree of tensor model parallelism.
+        pipeline_model_parallel_size (int): Degree of pipeline model parallelism.
+        pipeline_dtype (Optional[torch.dtype]): Data type for pipeline parallelism.
+        virtual_pipeline_model_parallel_size (Optional[int]): Size of virtual pipeline parallelism.
+        context_parallel_size (int): Degree of context parallelism to be passed to model_config.
+        sequence_parallel (bool): Whether to use sequence parallelism.
+        use_megatron_fsdp (bool): Whether to use Megatron FSDP.
+        train_iters (int): Total number of training iterations.
+        global_batch_size (int): Global batch size for training.
+        micro_batch_size (int): Micro batch size for training.
+        seq_length (int): Sequence length for training data.
+        lr (float): Learning rate.
+        min_lr (float): Minimum learning rate for cosine decay.
+        lr_warmup_iters (int): Number of warmup iterations for the learning rate.
+        lr_decay_iters (Optional[int]): Number of iterations over which to decay the LR.
+        eval_interval (int): Evaluation interval.
+        save_interval (int): Save interval.
+        precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model.
+        comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model.
+        adam_beta1 (float): Beta1 parameter for Adam optimizer.
+        adam_beta2 (float): Beta2 parameter for Adam optimizer.
+        adam_eps (float): Epsilon parameter for Adam optimizer.
+        weight_decay (float): Weight decay parameter for Adam optimizer.
+        eval_iters (int): Number of iterations to run for evaluation validation/test for.
+        pretrained_checkpoint (str | None): Path to pretrained checkpoint to load.
+        peft (str | PEFT | None): PEFT configuration (e.g., "lora" or LoRA object).
+        packed_sequence (bool): Whether to use packed sequences.
+        packed_train_data_path (str | None): Path to packed training data.
+        packed_val_data_path (str | None): Path to packed validation data.
+        packed_metadata_path (str | None): Path to packed metadata.
+        dataset_type (str): Dataset type to use. Either "squad" (default) or "mlperf_dataset".
+        use_transformer_engine_op_fuser (bool): If True, set ``model_cfg.use_transformer_engine_op_fuser``
+            (TE op-fuser path on the backbone, e.g. fused MLP). Set False if LM/validation loss diverges
+            from a known-good run.
+        stable_lora_with_te_op_fuser (bool): Single Primus knob for the **stable LoRA + op fuser** combo.
+            If True (default): backbone follows ``use_transformer_engine_op_fuser``, but LoRA always
+            uses unfused :class:`LoRALinear` (``use_te_fused_lora=False``) so loss matches the safe path.
+            If False: **legacy** behavior — LoRA uses ``use_te_fused_lora = use_transformer_engine_op_fuser``
+            (when TP=1, fused :class:`TEFusedLoRALinear` tracks backbone op fuser, as in older Bridge).
+        te_fused_lora_include_modules (Optional[List[str]]): Passed to :class:`LoRA`; only applies when
+            fused LoRA is enabled (``stable_lora_with_te_op_fuser=False`` and backbone op fuser on).
+        te_fused_lora_exclude_modules (Optional[List[str]]): Passed to :class:`LoRA`; same scope as include.
+
+    Returns:
+        ConfigContainer: Configuration for pre-training.
+    """
+    from transformers import AutoConfig
+
+    config = AutoConfig.from_pretrained("meta-llama/Llama-2-70b-hf")
+    bridge = AutoBridge.from_hf_config(config)
+    model_cfg = bridge.to_megatron_provider(load_weights=False)  # GPTProvider
+    model_cfg.tensor_model_parallel_size = tensor_model_parallel_size
+    model_cfg.pipeline_model_parallel_size = pipeline_model_parallel_size
+    model_cfg.pipeline_dtype = None
+    model_cfg.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size
+    model_cfg.context_parallel_size = context_parallel_size
+    model_cfg.sequence_parallel = sequence_parallel
+    model_cfg.seq_length = seq_length
+    model_cfg.perform_initialization = True
+    model_cfg.fp16 = False
+    model_cfg.bf16 = True
+    model_cfg.params_dtype = torch.bfloat16
+    model_cfg.autocast_dtype = torch.bfloat16
+    model_cfg.pipeline_dtype = torch.bfloat16
+    # Fusions / CE: TE parallel cross-entropy; grad-acc fusion (not used with Megatron FSDP).
+    model_cfg.cross_entropy_loss_fusion = False
+    model_cfg.cross_entropy_fusion_impl = "native"
+    model_cfg.gradient_accumulation_fusion = False
+    model_cfg.bias_dropout_fusion = True
+    model_cfg.fused_single_qkv_rope = False
+    model_cfg.apply_rope_fusion = True
+    model_cfg.use_transformer_engine_op_fuser = False
+    # Activation offload hint for TP paths (distinct from cpu_offloading / cpu_offloading_num_layers).
+    model_cfg.cpu_offloading_activations = True
+
+    # MXFP4 weights (fp4/e2m1); fp8=None during MXFP4 phase. FP8_* env vars in
+    # config_MI355X_1x8x1.sh configure healing (step 340) and TE delayed scaling,
+    # not Megatron model_cfg.fp8 (which stays None until healing).
+    model_cfg.fp8 = None
+    model_cfg.fp4 = "e2m1"
+    model_cfg.fp4_recipe = "mxfp4"
+    model_cfg.fp8_param_gather = False
+    model_cfg.grad_reduce_in_fp32 = False
+
+    # Used when cuda_graph_impl is not "none" (harmless when graphs are disabled).
+    model_cfg.cuda_graph_retain_backward_graph = True
+    model_cfg.cuda_graph_use_single_mempool = True
+    model_cfg.fp8_recipe = "delayed"
+    model_cfg.fp8_amax_history_len = 4
+    model_cfg.fp8_amax_compute_algo = "most_recent"
+    model_cfg.fp8_dot_product_attention = False
+    model_cfg.disable_parameter_transpose_cache = False
+    model_cfg.fine_grained_activation_offloading = False
+    model_cfg.use_transformer_engine_full_layer_spec = (
+        False  # Doesn't work beacuse of RMSNorm is not supported in FusedLayerNorm
+    )
+    model_cfg.cpu_offloading = False
+    model_cfg.cpu_offloading_num_layers = 0
+    model_cfg.empty_unused_memory_level = (
+        0  # 0: No empty, 1: Empty at end of eval, 2: Empty at end of eval and train.
+    )
+    # Optional MXFP4-style activation recompute (left untouched when kwargs are None).
+    if recompute_granularity is not None:
+        model_cfg.recompute_granularity = recompute_granularity
+    if recompute_method is not None:
+        model_cfg.recompute_method = recompute_method
+    if recompute_num_layers is not None:
+        model_cfg.recompute_num_layers = recompute_num_layers
+    # Pin TE's attention backend ("fused" forces CK / AOTriton on ROCm for the MXFP4 path).
+    # Leave as None to keep Megatron's "auto" default (which picks FlashAttention on ROCm).
+    if attention_backend is not None:
+        try:
+            model_cfg.attention_backend = AttnBackend[attention_backend]
+        except KeyError as e:
+            raise ValueError(
+                f"Unknown attention_backend {attention_backend!r}; expected one of "
+                f"{[b.name for b in AttnBackend]}."
+            ) from e
+    # Disable attention QK clipping / max-logit scans in the optimizer path (extra GPU work per step).
+    if hasattr(model_cfg, "qk_clip"):
+        model_cfg.qk_clip = False
+    if hasattr(model_cfg, "log_max_attention_logit"):
+        model_cfg.log_max_attention_logit = False
+
+    from megatron.bridge.training.config import OptimizerConfig, SchedulerConfig
+
+    opt_config = OptimizerConfig(
+        optimizer="adam",
+        lr=lr,
+        min_lr=min_lr,
+        clip_grad=clip_grad,
+        weight_decay=weight_decay,
+        adam_beta1=adam_beta1,
+        adam_beta2=adam_beta2,
+        adam_eps=adam_eps,
+        bf16=True,
+        params_dtype=torch.bfloat16,
+        use_distributed_optimizer=True,
+        overlap_param_gather_with_optimizer_step=True,
+        barrier_with_L1_time=False,
+        log_num_zeros_in_grad=False,
+    )
+
+    scheduler = SchedulerConfig(
+        lr_decay_style="cosine",
+        lr_decay_iters=lr_decay_iters,  # same as max_steps in nemo
+        lr_warmup_iters=lr_warmup_iters,  # value 0 same as nemo
+        lr_warmup_init=0.0,
+        start_weight_decay=weight_decay,
+        end_weight_decay=weight_decay,
+        weight_decay_incr_style="constant",
+        override_opt_param_scheduler=True,
+    )
+
+    peft_config = LoRA(
+        dim=16,
+        alpha=32,
+        dropout=0.1,
+        dropout_position="pre",
+        lora_A_init_method="xavier",
+        lora_B_init_method="zero",
+        a2a_experimental=True,
+        target_modules=["linear_qkv", "linear_proj"],
+        use_te_fused_lora=False,
+        te_fused_lora_include_modules=te_fused_lora_include_modules,
+        te_fused_lora_exclude_modules=(
+            te_fused_lora_exclude_modules if te_fused_lora_exclude_modules is not None else []
+        ),
+    )
+
+    # Dataset configuration - switch between squad and mlperf_dataset
+    if dataset_type == "squad":
+        dataset_cfg = default_squad_config(seq_length, packed_sequence)
+    elif dataset_type == "mlperf_dataset":
+        if packed_sequence:
+            packed_sequence_specs = PackedSequenceSpecs(
+                packed_sequence_size=seq_length,
+                tokenizer_model_name=hf_path,
+                packed_train_data_path=packed_train_data_path or "/data/train.npy",
+                packed_val_data_path=packed_val_data_path or "/data/validation.npy",
+                packed_metadata_path=packed_metadata_path or "/data/packed_metadata.jsonl",
+            )
+        else:
+            packed_sequence_specs = None
+        dataset_cfg = FinetuningDatasetConfig(
+            dataset_root="/data",
+            seq_length=seq_length,
+            seed=seed,
+            packed_sequence_specs=packed_sequence_specs,
+            data_sharding=True,
+            dataloader_type="batch",
+            num_workers=1,
+            do_test=False,
+            do_validation=True,
+            dataset_kwargs={"return_cu_seqlen": False},
+        )
+    else:
+        raise ValueError(f"Unknown dataset_type: {dataset_type!r}. Expected 'squad' or 'mlperf_dataset'.")
+
+    dataset_cfg.num_workers = 0
+    dataset_cfg.memmap_workers = 1  # needs to be 1>0
+    dataset_cfg.pin_memory = True
+    dataset_cfg.persistent_workers = False
+    dataset_cfg.dataloader_type = "batch"
+
+    # Config Container
+    cfg = ConfigContainer(
+        model=model_cfg,
+        train=TrainingConfig(
+            train_iters=train_iters,
+            eval_interval=eval_interval,
+            eval_iters=eval_iters,
+            global_batch_size=global_batch_size,
+            micro_batch_size=micro_batch_size,
+            # Manual GC aligns collections across ranks but adds periodic host pauses; use default GC for best step time.
+            manual_gc=False,
+            manual_gc_interval=0,
+            manual_gc_eval=False,
+            empty_unused_memory_level=0,  # 0: No empty, 1: Empty at end of eval, 2: Empty at end of eval and train.
+            train_sync_interval=None,
+            check_weight_hash_across_dp_replicas_interval=None,
+        ),
+        optimizer=opt_config,
+        scheduler=scheduler,
+        ddp=DistributedDataParallelConfig(
+            # Per-step NaN grad scan + sync; disable for throughput when loss NaN checks are off.
+            check_for_nan_in_grad=False,
+            grad_reduce_in_fp32=False,
+            overlap_grad_reduce=False,
+            overlap_param_gather=False,
+            average_in_collective=False,
+            use_distributed_optimizer=True,
+            # gradient_reduce_div_fusion=True,
+            # pad_buckets_for_high_nccl_busbw=True,
+            use_megatron_fsdp=False,
+            keep_fp8_transpose_cache=(
+                os.getenv("ENABLE_TRANSPOSE_CACHE", "").strip().lower() in ("1", "true", "yes", "on")
+            ),
+            fp8_param_gather=False,
+        ),
+        dataset=dataset_cfg,
+        logger=LoggerConfig(
+            log_interval=10,
+            tensorboard_dir=None,
+            # Per-step / log-interval overhead toggles (keep off for throughput).
+            log_params_norm=False,
+            log_throughput=True,
+            log_energy=False,
+            log_progress=False,
+            timing_log_level=0,
+            log_loss_scale_to_tensorboard=False,
+            log_timers_to_tensorboard=False,
+            log_throughput_to_tensorboard=False,
+            log_validation_ppl_to_tensorboard=False,
+            log_memory_to_tensorboard=False,
+            log_runtime_to_tensorboard=False,
+            log_world_size_to_tensorboard=False,
+            log_l2_norm_grad_to_tensorboard=False,
+            wandb_project=None,
+            wandb_exp_name=None,
+            wandb_save_dir=None,
+            wandb_entity=None,
+        ),
+        tokenizer=TokenizerConfig(
+            tokenizer_type="NullTokenizer" if use_null_tokenizer else "HuggingFaceTokenizer",
+            tokenizer_model=hf_path if not use_null_tokenizer else None,
+            vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE if use_null_tokenizer else None,
+        ),
+        checkpoint=CheckpointConfig(
+            save_interval=None,
+            save=None,
+            load=None,
+            save_optim=False,
+            save_rng=False,
+            pretrained_checkpoint=pretrained_checkpoint,
+            ckpt_format="torch_dist",
+            replication_factor=0,
+            most_recent_k=0,
+            finetune=True,
+            load_main_params_from_ckpt=False,
+            load_optim=False,
+            load_rng=False,
+        ),
+        rng=RNGConfig(seed=seed, te_rng_tracker=True),
+        rerun_state_machine=RerunStateMachineConfig(
+            check_for_nan_in_loss=check_for_nan_in_loss,
+        ),
+        straggler=StragglerDetectionConfig(log_straggler=False),
+        tensor_inspect=TensorInspectConfig(enabled=False),
+        comm_overlap=comm_overlap_config,
+        mixed_precision=precision_config,
+        peft=peft_config,
+        profiling=ProfilingConfig(
+            use_pytorch_profiler=False,
+            profile_step_start=140,
+            profile_step_end=144,
+            profile_ranks=list(range(8)),  # 1 node × 8 GPUs
+            record_shapes=False,
+            nvtx_ranges=False,
+        ),
+    )
+
+    if cfg.comm_overlap is None:
+        cfg.comm_overlap = CommOverlapConfig(
+            tp_comm_overlap=False,
+        )
+
+    return cfg
+
+
+def _reset_data_iterator(data_iterator):
+    """Reset data iterator to the beginning for deterministic evaluation.
+    Traverses through RerunDataIterator wrappers to find and reset any
+    ResettableDataIterator, ensuring evaluation always starts from the
+    same data point.
+    """
+    if data_iterator is None:
+        return
+    if isinstance(data_iterator, list):
+        for it in data_iterator:
+            _reset_data_iterator(it)
+        return
+    if isinstance(data_iterator, RerunDataIterator):
+        inner = data_iterator.iterable
+        if isinstance(inner, ResettableDataIterator):
+            inner.reset()
+            data_iterator.saved_microbatches.clear()
+            data_iterator.replaying = False
+            data_iterator.replay_pos = 0
+    elif isinstance(data_iterator, ResettableDataIterator):
+        data_iterator.reset()
+
+
+def evaluate(
+    state: GlobalState,
+    forward_step_func: ForwardStepCallable,
+    data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]],
+    model: list[MegatronModule],
+    process_non_loss_data_func: Optional[Callable],
+    config: ConfigContainer,
+    verbose: bool = False,
+    non_loss_data_func: Optional[Callable] = None,
+) -> tuple[Optional[dict[str, torch.Tensor]], Optional[Any], bool]:
+    """Evaluation function (from eval_m.py).
+    Validation loss aggregation matches NeMo: mean of per-microbatch means
+    (same as NeMo MaskedTokenLossReduction with validation_step=True, val_drop_last=True).
+    """
+    _reset_data_iterator(data_iterator)
+    wrapped_forward_step = prepare_forward_step_func(forward_step_func, state)
+    timers = state.timers
+    timers("evaluate", log_level=0).start(barrier=True)
+    for model_module in model:
+        model_module.eval()
+    pg_collection = get_pg_collection(model)
+    rerun_state_machine = get_rerun_state_machine()
+    rerun_mode = rerun_state_machine.get_mode()
+    rerun_state_machine.set_mode(RerunMode.DISABLED)
+    total_loss_dict = {}
+    eval_batch_size = state.cfg.train.global_batch_size
+    eval_num_microbatches = eval_batch_size // (
+        state.cfg.train.micro_batch_size * state.cfg.data_parallel_size
+    )
+    with torch.no_grad():
+        if verbose:
+            log_rank_0(f"Evaluating on {state.cfg.train.eval_iters * eval_batch_size} samples")
+        if (
+            state.cfg.model.cuda_graph_impl == "local"
+            and "full_iteration" in state.cfg.model.cuda_graph_scope
+        ):
+            forward_backward_func = FullCudaGraphWrapper(
+                get_forward_backward_func(),
+                cuda_graph_warmup_steps=state.cfg.model.cuda_graph_warmup_steps,
+            )
+        else:
+            forward_backward_func = get_forward_backward_func()
+        iteration = 0
+        while iteration < state.cfg.train.eval_iters:
+            iteration += 1
+            if verbose:
+                log_rank_0(f"Evaluating iter {iteration}/{state.cfg.train.eval_iters}")
+            seq_length = state.cfg.model.seq_length
+            eval_data_iterator = data_iterator
+            if state.cfg.dataset.dataloader_type == "batch":
+                eval_microbatch_iterator, seq_length = prepare_finetuning_batch(
+                    data_iterator=data_iterator,
+                    num_microbatches=eval_num_microbatches,
+                    default_seq_length=state.cfg.model.seq_length,
+                    seq_key="tokens",
+                )
+                eval_data_iterator = make_data_iterator_list(
+                    model=model,
+                    data_iterator=eval_microbatch_iterator,
+                )
+            config.timers = None
+            fault_tolerance.on_eval_step_start(state)
+            loss_dicts = forward_backward_func(
+                forward_step_func=wrapped_forward_step,
+                data_iterator=eval_data_iterator,
+                model=model,
+                num_microbatches=eval_num_microbatches,
+                seq_length=seq_length,
+                micro_batch_size=state.cfg.train.micro_batch_size,
+                forward_only=True,
+            )
+            fault_tolerance.on_eval_step_end(state)
+            config.timers = state.timers
+            if is_pp_last_stage(pg_collection.pp):
+                for key in loss_dicts[0].keys():
+                    if key not in total_loss_dict:
+                        total_loss_dict[key] = torch.zeros(2, dtype=torch.float32, device="cuda")
+                    val = [x[key].reshape(-1) for x in loss_dicts]
+                    if val[0].numel() == 2:
+                        # NeMo MLPerf-equivalent: accumulate raw [sum, count] across micros & eval iters.
+                        # Per-micro [sum, count] is already DP/CP-reduced inside
+                        # nemo_loss.MaskedTokenLossReduction.forward, so no extra all_reduce here.
+                        per_iter = torch.vstack([v.float() for v in val]).sum(
+                            dim=0
+                        )  # [Σ_micro sum, Σ_micro count]
+                        total_loss_dict[key] += per_iter
+                    elif val[0].numel() == 1:
+                        # legacy single-scalar branch
+                        micro_sum = torch.stack([v.float() for v in val], dim=0).sum()
+                        total_loss_dict[key][0] += micro_sum
+                        total_loss_dict[key][1] += float(len(loss_dicts))
+                    else:
+                        raise ValueError(f"Invalid value shape: {val[0].shape} for key {key}")
+
+            state.train_state.consumed_valid_samples += eval_batch_size
+            if state.cfg.train.exit_duration_in_mins:
+                train_time = (time.time() - state.start_time) / 60.0
+                done_cuda = torch.tensor(
+                    [train_time > state.cfg.train.exit_duration_in_mins],
+                    dtype=torch.int,
+                    device="cuda",
+                )
+                torch.distributed.all_reduce(done_cuda, op=torch.distributed.ReduceOp.MAX)
+                done = bool(done_cuda.item())
+                if done:
+                    rerun_state_machine.set_mode(rerun_mode)
+                    log_rank_0("Exiting during evaluation, timelimit reached")
+                    return None, None, True
+
+        collected_non_loss_data = None
+        if non_loss_data_func is not None:
+            collected_non_loss_data = non_loss_data_func(model)
+        elif process_non_loss_data_func is not None and is_last_rank():
+            non_loss_data_iterator = data_iterator
+            non_loss_seq_length = state.cfg.model.seq_length
+            if state.cfg.dataset.dataloader_type == "batch":
+                non_loss_microbatch_iterator, non_loss_seq_length = prepare_finetuning_batch(
+                    data_iterator=data_iterator,
+                    num_microbatches=get_num_microbatches(),
+                    default_seq_length=state.cfg.model.seq_length,
+                    seq_key="tokens",
+                )
+                non_loss_data_iterator = make_data_iterator_list(
+                    model=model,
+                    data_iterator=non_loss_microbatch_iterator,
+                )
+            collected_non_loss_data = forward_backward_func(
+                forward_step_func=wrapped_forward_step,
+                data_iterator=non_loss_data_iterator,
+                model=model,
+                num_microbatches=get_num_microbatches(),
+                seq_length=non_loss_seq_length,
+                micro_batch_size=state.cfg.train.micro_batch_size,
+                forward_only=True,
+                collect_non_loss_data=True,
+            )
+    for model_module in model:
+        model_module.train()
+    for key in total_loss_dict:
+        val_loss_sum, val_loss_count = total_loss_dict[key]
+        if val_loss_count > 0:
+            total_loss_dict[key] = val_loss_sum / val_loss_count
+        else:
+            total_loss_dict[key] = val_loss_sum
+    timers("evaluate").stop()
+    timers.log(["evaluate"])
+    rerun_state_machine.set_mode(rerun_mode)
+    return total_loss_dict, collected_non_loss_data, False
+
+
+from megatron.bridge.training import eval
+
+
+def evaluate_and_print_results_custom(
+    state: GlobalState,
+    prefix: str,
+    forward_step_func: ForwardStepCallable,
+    data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]],
+    model: list[MegatronModule],
+    config: ConfigContainer,
+    verbose: bool = False,
+    write_to_tensorboard: bool = False,
+    process_non_loss_data_func: Optional[Callable] = None,
+    non_loss_data_func: Optional[Callable] = None,
+    throughput: float = 0.0,  # samples/sec
+) -> tuple:
+    """Helper function to evaluate and dump results on screen.
+
+    Args:
+        state (GlobalState): The global state object.
+        prefix (str): Prefix for logging evaluation results.
+        forward_step_func (Callable): The function that performs a forward step.
+        data_iterator (Optional[Union[RerunDataIterator, list[RerunDataIterator]]]): Iterator over evaluation data.
+        model (list[MegatronModule]): list of model chunks.
+        config (ConfigContainer): Configuration container (potentially redundant).
+        verbose (bool, optional): Whether to print evaluation progress. Defaults to False.
+        write_to_tensorboard (bool, optional): Whether to write results to TensorBoard. Defaults to False.
+        process_non_loss_data_func (Optional[Callable], optional): Function to process non-loss data. Defaults to None.
+        non_loss_data_func (Optional[Callable], optional): Function to compute non-loss data. Defaults to None.
+    """
+    global _TARGET_LOSS_REACHED
+    if _TARGET_LOSS_REACHED:
+        # Target already hit on a previous pass; skip this entire eval.
+        # Returning should_exit=True keeps the outer while-loop on its exit path.
+        _orig_log_rank_0(
+            f"Skipping evaluation at {prefix}: target loss < " f"{MLPERF_TARGET_LOSS} already reached."
+        )
+        return True, None
+
+    log_rank_0(f"Evaluating and printing results at {prefix}")
+    should_exit = False
+
+    def is_last_rank():
+        return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
+
+    import math
+
+    if write_to_tensorboard:
+        writer = state.tensorboard_logger
+    else:
+        writer = None
+
+    wandb_writer = state.wandb_logger
+
+    eval_start = time.time()
+    total_loss_dict, collected_non_loss_data, timelimit = evaluate(
+        state,
+        forward_step_func,
+        data_iterator,
+        model,
+        process_non_loss_data_func,
+        config,
+        verbose,
+        non_loss_data_func,
+    )
+    eval_duration = time.time() - eval_start
+    # Timelimit hit during evaluation
+    if timelimit:
+        return False, None
+    string = f" validation loss at {prefix} | "
+    for key in total_loss_dict:
+        string += "{} value: {:.6E} | ".format(key, total_loss_dict[key].item())
+        ppl = math.exp(min(20, total_loss_dict[key].item()))
+        string += "{} PPL: {:.6E} | ".format(key, ppl)
+        if writer:
+            writer.add_scalar(
+                "{} validation".format(key), total_loss_dict[key].item(), state.train_state.step
+            )
+            writer.add_scalar(
+                "{} validation vs samples".format(key),
+                total_loss_dict[key].item(),
+                state.train_state.consumed_train_samples,
+            )
+            if state.cfg.logger.log_validation_ppl_to_tensorboard:
+                writer.add_scalar("{} validation ppl".format(key), ppl, state.train_state.step)
+                writer.add_scalar(
+                    "{} validation ppl vs samples".format(key), ppl, state.train_state.consumed_train_samples
+                )
+
+        if wandb_writer and is_last_rank():
+            wandb_writer.log(
+                {"{} validation".format(key): total_loss_dict[key].item()}, state.train_state.step
+            )
+            if state.cfg.logger.log_validation_ppl_to_tensorboard:
+                wandb_writer.log({"{} validation ppl".format(key): ppl}, state.train_state.step)
+
+    if process_non_loss_data_func is not None and writer and is_last_rank():
+        process_non_loss_data_func(collected_non_loss_data, state.train_state.step, writer)
+
+    string += "throughput: {:.6E} | ".format(throughput)
+    string += "eval duration: {:.6E} | ".format(eval_duration)
+    if writer:
+        writer.add_scalar("throughput samples/sec", throughput, state.train_state.step)
+        writer.add_scalar(
+            "throughput samples/sec vs samples", throughput, state.train_state.consumed_train_samples
+        )
+
+    if wandb_writer and is_last_rank():
+        wandb_writer.log({"throughput samples/sec": throughput}, state.train_state.step)
+
+    length = len(string) + 1
+    # Match training logs: emit on rank 0 so validation lines show up in typical Primus rank-0 log streams.
+    log_rank_0("-" * length)
+    log_rank_0(string)
+    log_rank_0("-" * length)
+    # Guard against non-PP-last ranks (or otherwise empty total_loss_dict) where
+    # "lm loss" may not be populated; only run the MLPerf early-exit check when
+    # the key is present.
+    eval_loss_value = None
+    if total_loss_dict and "lm loss" in total_loss_dict:
+        eval_loss_value = total_loss_dict["lm loss"].item()
+        if eval_loss_value < MLPERF_TARGET_LOSS:
+            should_exit = True
+            _TARGET_LOSS_REACHED = True
+            log_rank_0(f"Validation loss is less than {MLPERF_TARGET_LOSS}, exiting training")
+    return should_exit, eval_loss_value
+
+
+eval.evaluate_and_print_results = evaluate_and_print_results_custom
+
+
+def warmup_eval(
+    state: GlobalState,
+    forward_step_func: ForwardStepCallable,
+    data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]],
+    model: list[MegatronModule],
+    config: ConfigContainer,
+    verbose: bool = False,
+    process_non_loss_data_func: Optional[Callable] = None,
+    non_loss_data_func: Optional[Callable] = None,
+    num_warmup_iters: int = 10,
+) -> None:
+    log_rank_0(f"Starting warmup eval...")
+    for i in range(num_warmup_iters):
+        log_rank_0(f"Warmup eval iteration {i} running...")
+        evaluate_and_print_results_custom(
+            state,
+            f"warmup iteration {i}",
+            forward_step_func,
+            data_iterator,
+            model,
+            config,
+            verbose=verbose,
+            write_to_tensorboard=False,
+            process_non_loss_data_func=process_non_loss_data_func,
+            non_loss_data_func=non_loss_data_func,
+        )
+        log_rank_0(f"Warmup eval iteration {i} completed")
+    log_rank_0(f"Warmup eval completed")
+
+
+class _SyntheticSFTDataIterator:
+    """Infinite iterator yielding synthetic finetuning batches for warmup.
+
+    Produces random token tensors matching the SFT packed-sequence shape expected
+    by prepare_finetuning_batch / forward_backward_func.
+    """
+
+    def __init__(self, seq_length: int, micro_batch_size: int, vocab_size: int = 32000):
+        self._seq_length = seq_length
+        self._mbs = micro_batch_size
+        self._vocab_size = vocab_size
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        sl = self._seq_length
+        mbs = self._mbs
+        tokens = torch.randint(0, self._vocab_size, (mbs, sl), dtype=torch.int64, device="cuda")
+        labels = torch.randint(0, self._vocab_size, (mbs, sl), dtype=torch.int64, device="cuda")
+        loss_mask = torch.ones(mbs, sl, dtype=torch.float32, device="cuda")
+        position_ids = torch.arange(sl, dtype=torch.int64, device="cuda").unsqueeze(0).expand(mbs, -1)
+        return {"tokens": tokens, "labels": labels, "loss_mask": loss_mask, "position_ids": position_ids}
+
+
+def run_synthetic_warmup(
+    forward_step_func: ForwardStepCallable,
+    model: list[MegatronModule],
+    optimizer: MegatronOptimizer,
+    scheduler: OptimizerParamScheduler,
+    global_state: GlobalState,
+    pg_collection: ProcessGroupCollection,
+) -> None:
+    """Run synthetic warmup steps to pre-compile JIT kernels before RUN_START.
+
+    Mirrors NeMo's warmup behavior: runs N training (fwd+bwd+opt) steps and
+    M validation (fwd-only) steps with random data, then restores all state so
+    that real training starts from a clean slate.
+
+    Controlled by environment variables:
+        SYNTH_WARMUP_STEPS       - number of training warmup steps (default 5, 0 disables)
+        SYNTH_WARMUP_VALID_STEPS - number of validation warmup steps (default 5, 0 disables)
+    """
+    warmup_train_steps = int(os.getenv("SYNTH_WARMUP_STEPS", "5"))
+    warmup_valid_steps = int(os.getenv("SYNTH_WARMUP_VALID_STEPS", "5"))
+
+    if warmup_train_steps <= 0 and warmup_valid_steps <= 0:
+        return
+
+    config = global_state.cfg
+    seq_length = config.model.seq_length
+    mbs = config.train.micro_batch_size
+    gbs = config.train.global_batch_size
+    dp_size = getattr(config, "data_parallel_size", 1)
+    num_microbatches = gbs // (mbs * dp_size)
+
+    synth_iter = _SyntheticSFTDataIterator(seq_length, mbs)
+    wrapped_forward_step = prepare_forward_step_func(forward_step_func, global_state)
+    forward_backward_func = get_forward_backward_func()
+
+    # --- Save model parameters ---
+    models = model if isinstance(model, (list, tuple)) else [model]
+    saved_params = {}
+    for m in models:
+        for name, p in m.named_parameters():
+            if p.requires_grad:
+                saved_params[(id(m), name)] = p.data.to("cpu", copy=True)
+
+    # --- Save optimizer state (neuter to prevent real updates) ---
+    saved_opt_states = []
+    inner_opts = []
+    if hasattr(optimizer, "chained_optimizers"):
+        for sub in optimizer.chained_optimizers:
+            inner_opts.append(getattr(sub, "optimizer", sub))
+    else:
+        inner_opts.append(getattr(optimizer, "optimizer", optimizer))
+
+    for inner in inner_opts:
+        saved = []
+        for group in inner.param_groups:
+            state = {}
+            for key in ("betas", "weight_decay", "bias_correction"):
+                if key in group:
+                    state[key] = group[key]
+            saved.append(state)
+            if "betas" in group:
+                group["betas"] = [1.0, 1.0]
+            if "weight_decay" in group:
+                group["weight_decay"] = 0.0
+            if "bias_correction" in group:
+                group["bias_correction"] = False
+        saved_opt_states.append(saved)
+
+    # --- Save LR scheduler state ---
+    saved_sched = {}
+    if scheduler is not None:
+        for k in ("num_steps", "num_floating_point_operations_so_far"):
+            if hasattr(scheduler, k):
+                saved_sched[k] = getattr(scheduler, k)
+
+    # --- Training warmup steps (fwd + bwd + optimizer) ---
+    if warmup_train_steps > 0:
+        for m in models:
+            m.train()
+        for step in range(1, warmup_train_steps + 1):
+            train_step(
+                wrapped_forward_step,
+                synth_iter,
+                model,
+                optimizer,
+                scheduler,
+                global_state,
+                pg_collection,
+                forward_backward_func,
+            )
+            torch.cuda.synchronize()
+
+    # --- Validation warmup steps (forward-only) ---
+    if warmup_valid_steps > 0:
+        for m in models:
+            m.eval()
+        with torch.no_grad():
+            for step in range(1, warmup_valid_steps + 1):
+                forward_backward_func(
+                    forward_step_func=wrapped_forward_step,
+                    data_iterator=synth_iter,
+                    model=model,
+                    num_microbatches=num_microbatches,
+                    seq_length=seq_length,
+                    micro_batch_size=mbs,
+                    forward_only=True,
+                )
+                torch.cuda.synchronize()
+        for m in models:
+            m.train()
+
+    # --- Restore model parameters ---
+    for m in models:
+        for name, p in m.named_parameters():
+            key = (id(m), name)
+            if key in saved_params:
+                p.data.copy_(saved_params[key].to(p.device))
+    del saved_params
+
+    # --- Restore optimizer state ---
+    for inner, saved in zip(inner_opts, saved_opt_states):
+        for group, state in zip(inner.param_groups, saved):
+            for key, val in state.items():
+                group[key] = val
+
+    if hasattr(optimizer, "reload_model_params"):
+        optimizer.reload_model_params()
+
+    # --- Zero optimizer state tensors (exp_avg / exp_avg_sq) ---
+    for inner in inner_opts:
+        for param_states in inner.state.values():
+            for k, v in param_states.items():
+                if isinstance(v, torch.Tensor) and v.is_floating_point():
+                    v.zero_()
+
+    # --- Restore LR scheduler and re-sync param_groups['lr'] ---
+    if scheduler is not None:
+        for k, v in saved_sched.items():
+            setattr(scheduler, k, v)
+        scheduler.step(0)
+
+    # --- Reset FP8 state ---
+    for m in models:
+        for module in m.modules():
+            if hasattr(module, "fp8_initialized"):
+                module.fp8_initialized = False
+                if hasattr(module, "reset_fp8_meta_tensors"):
+                    try:
+                        module.reset_fp8_meta_tensors()
+                    except Exception:
+                        pass
+
+    # --- Seed FP8 amax_history to prevent scale=inf after reset ---
+    for m in models:
+        for module in m.modules():
+            fp8_meta = getattr(module, "fp8_meta", None)
+            if fp8_meta is None or not isinstance(fp8_meta, dict):
+                continue
+            for key in ("scaling_fwd", "scaling_bwd"):
+                if key not in fp8_meta:
+                    continue
+                tensor_meta = fp8_meta[key]
+                if hasattr(tensor_meta, "amax_history"):
+                    tensor_meta.amax_history.fill_(1.0)
+
+    # --- Reset FP4/MXFP4 state (block scales, quantizer caches) ---
+    for m in models:
+        for module in m.modules():
+            is_fp4 = bool(getattr(module, "fp4", False)) or hasattr(module, "fp4_initialized")
+            if not is_fp4:
+                fp8_meta = getattr(module, "fp8_meta", None)
+                if isinstance(fp8_meta, dict):
+                    for key in ("scaling_fwd", "scaling_bwd"):
+                        tm = fp8_meta.get(key)
+                        if tm is not None and ("FP4" in type(tm).__name__ or "Fp4" in type(tm).__name__):
+                            is_fp4 = True
+                            break
+            if not is_fp4:
+                continue
+            if hasattr(module, "fp4_initialized"):
+                module.fp4_initialized = False
+            if hasattr(module, "fp8_initialized"):
+                module.fp8_initialized = False
+            if hasattr(module, "reset_fp4_meta_tensors"):
+                try:
+                    module.reset_fp4_meta_tensors()
+                except Exception:
+                    pass
+            elif hasattr(module, "reset_fp8_meta_tensors"):
+                try:
+                    module.reset_fp8_meta_tensors()
+                except Exception:
+                    pass
+
+    # --- Clear gradients ---
+    for m in models:
+        if hasattr(m, "zero_grad_buffer"):
+            m.zero_grad_buffer()
+        m.zero_grad(set_to_none=True)
+
+    # --- Reset consumed samples / step counter back to 0 ---
+    global_state.train_state.step = 0
+    global_state.train_state.consumed_train_samples = 0
+
+    # --- Clear CUDA cache (single clear after warmup, before RUN_START) ---
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+
+
+def megatron_bridge_train_override(
+    forward_step_func: ForwardStepCallable,
+    model: list[MegatronModule],
+    optimizer: MegatronOptimizer,
+    scheduler: OptimizerParamScheduler,
+    train_data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]],
+    valid_data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]],
+    global_state: GlobalState,
+    checkpointing_context: dict[str, Any],
+    pg_collection: ProcessGroupCollection,
+    process_non_loss_data_func: Optional[Callable] = None,
+    non_loss_data_func: Optional[Callable] = None,
+) -> None:
+    """Main training loop.
+
+    Handles the overall training process, including the iteration loop,
+    calling train_step, evaluation, checkpointing, logging, and exit conditions.
+
+    Args:
+        forward_step_func: Callable that executes a single forward step.
+        model: list of model chunks (potentially wrapped in DDP).
+        optimizer: The optimizer instance.
+        scheduler: The learning rate scheduler instance.
+        train_data_iterator: Iterator for the training dataset.
+        valid_data_iterator: Iterator for the validation dataset.
+        global_state: The GlobalState object holding various training states.
+        checkpointing_context: Context dictionary for checkpointing.
+        process_non_loss_data_func: Optional function to process non-loss data during evaluation.
+        non_loss_data_func: Optional function to compute non-loss data during evaluation.
+
+    Warnings:
+        This is an experimental API and is subject to change in backwards
+        incompatible ways without notice.
+    """
+
+    # import ctypes
+    # import gc
+    # libc = ctypes.CDLL("libc.so.6")
+    # gc.collect()
+    # libc.malloc_trim(0)
+
+    config: ConfigContainer = global_state.cfg
+    model_config = get_model_config(model[0])
+    train_config = config.train
+    timers = global_state.timers
+    straggler_timer = global_state.straggler_timer
+    energy_monitor = global_state.energy_monitor
+
+    # Prepare forward_step_func (check signature and inject state if needed).
+    # This is done once to prevent creating new partial objects every iteration.
+    #
+    # Note on reference semantics:
+    # - functools.partial stores a reference to global_state, not a copy
+    # - When global_state.train_state.step changes, the partial sees the updated value
+    # - This is safe because GlobalState is a mutable object passed by reference
+    #
+    # For functors (classes with __call__ defined):
+    # - For functors: partial(functor_instance, state) still allows functor's internal state to work
+    # - inspect.signature() properly inspects the __call__ method of functors
+    wrapped_forward_step_func = prepare_forward_step_func(forward_step_func, global_state)
+
+    # Turn on training mode which enables dropout.
+    for model_module in model:
+        model_module.train()
+        log_rank_0(f"Model module: {model_module}")
+
+    # Tracking loss.
+    total_loss_dict = {}
+
+    # Make sure rerun_state_machine has the right iteration loaded from checkpoint.
+    rerun_state_machine = get_rerun_state_machine()
+    if rerun_state_machine.current_iteration != global_state.train_state.step:
+        log_rank_0(f"Setting rerun_state_machine.current_iteration to {global_state.train_state.step}...")
+        rerun_state_machine.current_iteration = global_state.train_state.step
+
+    global_state.train_state.floating_point_operations_so_far
+    num_floating_point_operations_since_last_log_event = 0.0
+
+    if energy_monitor is not None:
+        energy_monitor.setup()
+        energy_monitor.resume()
+
+    # interval-time is started/stopped once per iteration around train_step only (see training loop).
+    # A single long-running interval-time includes logging, hooks, and other post-step work, which
+    # inflates "elapsed time per iteration (ms)" vs profiler / train-step wall time.
+    report_memory_flag = True
+    pre_hook_enabled = False
+    should_exit = False
+    exit_code = 0
+
+    if train_config.manual_gc:
+        # Disable the default garbage collector and perform the collection manually.
+        # This is to align the timing of garbage collection across ranks.
+        assert (
+            train_config.manual_gc_interval >= 0
+        ), "Manual garbage collection interval should be larger than or equal to 0"
+        gc.disable()
+        gc.collect()
+
+    if config.straggler and config.straggler.log_straggler:
+        world = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        mmcnt = config.straggler.straggler_minmax_count
+        straggler_timer.configure(
+            world,
+            rank,
+            mmcnt=mmcnt,
+            enabled=not config.straggler.disable_straggler_on_startup,
+            port=config.straggler.straggler_ctrlr_port,
+        )
+
+    # Initialize NVRx straggler detection if enabled
+    nvrx_straggler_manager = global_state.nvrx_straggler_manager
+    if nvrx_straggler_manager is not None:
+        try:
+            # Initialize the straggler detector first
+            nvrx_straggler_manager.initialize()
+            # Wrap the train_step function for monitoring
+            # Note: The nvidia-resiliency-ext library will monitor the actual train_step calls
+            nvrx_straggler_manager.wrap_train_step_function(train_step)
+        except Exception as e:
+            log_rank_0(f"Failed to initialize NVRx straggler detection: {e}")
+            # Set to None to disable further checks
+            global_state._nvrx_straggler_manager = None
+
+    get_num_microbatches()
+    eval_duration = 0.0
+    eval_iterations = 0
+
+    prof = None
+    nsys_nvtx_context = None  # NVTX context for nsys profiling
+    prof_config = config.profiling
+    if prof_config and should_profile_rank(prof_config, torch.distributed.get_rank()):
+        if prof_config.use_pytorch_profiler:
+            trace_dir = config.logger.tensorboard_dir or os.path.join(os.getcwd(), "torch_profiler_traces")
+            prof = initialize_pytorch_profiler(prof_config, trace_dir)
+            prof.start()
+
+    # Initialize RPD profiler if enabled
+    rpd = None
+    rpd_status = None
+    profiler_type = os.getenv("PROFILER", "")
+    rpd_warmup_steps = int(os.getenv("RPD_WARMUP_STEPS", "0"))
+    rpd_active_steps = int(os.getenv("RPD_ACTIVE_STEPS", "100"))
+    if profiler_type == "rpd":
+        try:
+            from rpdTracerControl import rpdTracerControl
+
+            rank = torch.distributed.get_rank()
+            rpd_filename = os.getenv("RPD_TRACE_FILENAME", f"trace.rpd")
+            rpdTracerControl.setFilename(name=rpd_filename, append=True)
+            rpd = rpdTracerControl()
+            log_rank_0(
+                f"RPD profiler initialized. Will start at step {rpd_warmup_steps} and stop at step {rpd_warmup_steps + rpd_active_steps}"
+            )
+        except Exception as e:
+            log_rank_0(f"Failed to initialize RPD profiler: {e}")
+            rpd = None
+    else:
+        log_rank_0(f"### Profiler type is {profiler_type}")
+
+    # Megatron FSDP and FSDP2 does not have this hook
+    should_toggle_forward_pre_hook = should_disable_forward_pre_hook(
+        config.ddp.use_megatron_fsdp,
+        config.optimizer.use_distributed_optimizer,
+        config.ddp.overlap_param_gather,
+    )
+    # Disable forward pre-hook to start training to ensure that errors in checkpoint loading
+    # or random initialization don't propagate to all ranks in first all-gather (which is a
+    # no-op if things work correctly).
+    if should_toggle_forward_pre_hook:
+        disable_forward_pre_hook(model, param_sync=False)
+        # Also remove param_sync_func temporarily so that sync calls made in
+        # `forward_backward_func` are no-ops.
+        param_sync_func = model_config.param_sync_func
+        model_config.param_sync_func = None
+        pre_hook_enabled = False
+    # Also, check weight hash across DP replicas to be very pedantic.
+    if train_config.check_weight_hash_across_dp_replicas_interval is not None:
+        assert check_param_hashes_across_dp_replicas(
+            model, cross_check=True
+        ), "Parameter hashes not matching across DP replicas"
+        torch.distributed.barrier()
+        log_rank_0(f">>> Weight hashes match after {global_state.train_state.step} iterations...")
+
+    # Capture CUDA Graphs.
+    cuda_graph_helper = None
+    if model_config.cuda_graph_impl == "transformer_engine":
+        cuda_graph_helper = TECudaGraphHelper(
+            model=model,
+            config=model_config,
+            seq_length=config.model.seq_length,
+            micro_batch_size=config.train.micro_batch_size,
+            optimizers=[optimizer],
+        )
+
+    # Track train step elapsed time for throughput logging
+    history_wct = None
+    if config.logger.log_throughput_to_tensorboard:
+        history_wct = deque(maxlen=config.logger.throughput_window_size + 1)
+
+    # Wrap forward_backward_func for Full iteration CUDA graph
+    forward_backward_func = get_forward_backward_func()
+    if config.model.cuda_graph_impl == "local" and "full_iteration" in config.model.cuda_graph_scope:
+        forward_backward_func = FullCudaGraphWrapper(
+            forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps
+        )
+
+    start_iteration = global_state.train_state.step
+    log_rank_0(f"Starting training loop at iteration {start_iteration}")
+
+    dp_size = pg_collection.dp.size()
+    batch_size = dp_size * train_config.micro_batch_size * get_num_microbatches()
+    timer = Timer(batch_size)
+
+    # Synthetic warmup: pre-compile JIT kernels before measured training begins.
+    run_synthetic_warmup(
+        forward_step_func,
+        model,
+        optimizer,
+        scheduler,
+        global_state,
+        pg_collection,
+    )
+    _log_training_gpu_mem("after synthetic warmup", config.logger.memory_keys)
+
+    sft_logger = _get_sft_logger()
+
+    # MLPerf logging: transition from init to training (after warmup)
+    if sft_logger is not None:
+        sft_logger.log_init_stop_run_start()
+
+    # Wall-clock for the training loop only (first through last train step; warmup eval disabled).
+    training_wall_start = time.perf_counter()
+    # Skip the first N interval evals (e.g. 48, 96, 144 when eval_interval=48); run from step 4*interval (192).
+    eval_skip_first_n = 3
+
+    # One-shot MXFP4 healing env banner (no-op unless HEALING_ITER > 0 and
+    # MXFP4_HEALING_PHASE_LOG is on). Safe to call even when the healing
+    # module isn't imported — gracefully degrades.
+    try:
+        from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import (
+            log_healing_env_banner_once,
+        )
+
+        log_healing_env_banner_once()
+    except Exception:  # noqa: BLE001
+        pass
+
+    # NeMo / MLPerf reference (MI355X implementation): CustomCallback uses time.time() in Lightning
+    # on_train_batch_start / on_train_batch_end — i.e. wall time over training_step only, no cuda.synchronize.
+    # We mirror that by timing megatron.train_step only, then averaging over logger.log_interval like interval-time.
+    nemo_style_iter_seconds_accum = 0.0
+    nemo_style_iter_count = 0
+
+    # Run training iterations till done.
+    while global_state.train_state.step < train_config.train_iters:
+        # Handle RPD profiling start
+        if rpd and global_state.train_state.step >= rpd_warmup_steps and not rpd_status:
+            log_rank_0(f"Starting RPD profiling at step {global_state.train_state.step}")
+            rpd.start()
+            rpd_status = "running"
+
+        # Handle RPD profiling stop
+        if (
+            rpd
+            and rpd_status == "running"
+            and global_state.train_state.step >= rpd_warmup_steps + rpd_active_steps
+        ):
+            log_rank_0(f"Stopping RPD profiling at step {global_state.train_state.step}")
+            rpd.stop()
+            rpd_status = "finished"
+
+        # Handle profiling for this step
+        nvtx_ctx = handle_profiling_step(
+            prof_config,
+            global_state.train_state.step,
+            torch.distributed.get_rank(),
+            prof,
+        )
+        if nvtx_ctx is not None:
+            nsys_nvtx_context = nvtx_ctx
+
+        # Update the timeout for all process groups after initialization
+        # We update the timeout after the first successful iteration,
+        # which takes longer than others usually
+        if global_state.train_state.step == start_iteration + 1:
+            distributed_timeout_seconds_after_init = (
+                global_state.cfg.dist.distributed_timeout_seconds_after_init
+            )
+            if distributed_timeout_seconds_after_init is not None:
+                update_pg_timeout(timedelta(seconds=distributed_timeout_seconds_after_init))
+
+        # Capture CUDA Graphs after warmup.
+        # if (
+        #     model_config.cuda_graph_impl == "transformer_engine"
+        #     and cuda_graph_helper is not None
+        #     and not cuda_graph_helper.graphs_created()
+        #     and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps
+        # ):
+        #     if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
+        #         disable_forward_pre_hook(model, param_sync=False)
+        #     cuda_graph_helper.create_cudagraphs()
+        #     if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
+        #         enable_forward_pre_hook(model)
+        #         cuda_graph_helper.cuda_graph_set_manual_hooks()
+
+        # Run training step.
+        timers("interval-time", log_level=0).start(barrier=False)
+        timer.start()
+        fault_tolerance.on_training_step_start(global_state)
+        _nemo_t0 = time.time()
+        (
+            loss_dict,
+            skipped_iter,
+            should_checkpoint,
+            should_exit,
+            exit_code,
+            grad_norm,
+            num_zeros_in_grad,
+            log_max_attention_logit,
+        ) = train_step(
+            wrapped_forward_step_func,
+            train_data_iterator,
+            model,
+            optimizer,
+            scheduler,
+            global_state,
+            pg_collection,
+            forward_backward_func,
+        )
+        nemo_style_iter_seconds_accum += time.time() - _nemo_t0
+        nemo_style_iter_count += 1
+        fault_tolerance.on_training_step_end(global_state)
+        timer.stop()
+        timers("interval-time").stop()
+
+        # Advance NVIDIA DLFw Inspect step if enabled
+        tensor_inspect_step_if_enabled(config.tensor_inspect)
+
+        if config.logger.log_throughput_to_tensorboard:
+            history_wct.append(time.time() - global_state.start_time)
+        if should_exit:
+            break
+
+        # Enable forward pre-hooks after first set of forward and backward passes.
+        # When running in fp16, skip all NaN iterations until steady-state loss scaling value
+        # is reached.
+        if global_state.train_state.step == start_iteration:
+            if skipped_iter:
+                # Only enable forward pre-hook after a training step has successfully run. Relevant
+                # for fp16 codepath where first XX iterations are skipped until steady-state loss
+                # scale value is reached.
+                start_iteration = global_state.train_state.step + 1
+            else:
+                # Enable forward pre-hook after training step has successfully run. All subsequent
+                # forward passes will use the forward pre-hook / `param_sync_func` in
+                # `forward_backward_func`.
+                if should_toggle_forward_pre_hook:
+                    enable_forward_pre_hook(model)
+                    model_config.param_sync_func = param_sync_func
+                    pre_hook_enabled = True
+                    # Set the manual hooks here since it's not set right after the capturing.
+                    if (
+                        model_config.cuda_graph_impl == "transformer_engine"
+                        and model_config.cuda_graph_warmup_steps == 0
+                    ):
+                        assert cuda_graph_helper.graphs_created(), "CUDA Graphs should have been created."
+                        cuda_graph_helper.cuda_graph_set_manual_hooks()
+
+        global_state.train_state.step += 1
+
+        # MXFP4 healing: when train_state.step + 1 == HEALING_ITER, restore FP8
+        # weights from the CPU stash (see ``mxfp4_healing``) and
+        # switch ``megatron.core.fp4_utils`` out of the MXFP4 phase. No-op when
+        # HEALING_ITER == 0 (default) so BF16/MXFP8 submission runs are
+        # unaffected.
+        try:
+            from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import (
+                mxfp4_healing as _mxh,
+            )
+
+            if _mxh.healing_iter() > 0:
+                _mxh.apply_healing_after_step(model, model_config, global_state.train_state.step)
+            _mxh.log_training_step_phase(global_state.train_state.step)
+        except Exception as _heal_err:  # noqa: BLE001
+            _orig_log_rank_0(
+                f"[mxfp4_healing] Failed at step={global_state.train_state.step}: "
+                f"{type(_heal_err).__name__}: {_heal_err}"
+            )
+            raise
+
+        # If fsdp_manual_registration is enabled, manually register FSDP communication buffers after one training step.
+        # if global_state.train_state.step == start_iteration + 1 and config.ddp.use_megatron_fsdp:
+        #     _maybe_register_fsdp_buffers(config, model)
+
+        global_state.train_state.consumed_train_samples += batch_size
+        num_skipped_samples_in_batch = (
+            get_current_global_batch_size() - get_current_running_global_batch_size()
+        )
+        if train_config.decrease_batch_size_if_needed:
+            assert num_skipped_samples_in_batch >= 0
+        else:
+            assert num_skipped_samples_in_batch == 0
+        global_state.train_state.skipped_train_samples += num_skipped_samples_in_batch
+        num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations(config, batch_size)
+        global_state.train_state.floating_point_operations_so_far += num_floating_point_operations_in_batch
+        global_state.train_state.floating_point_operations_so_far
+        num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch
+
+        # Logging.
+        if hasattr(optimizer, "is_stub_optimizer") and not optimizer.is_stub_optimizer:
+            loss_scale = optimizer.get_loss_scale().item()
+        else:
+            loss_scale = 1.0
+        params_norm = None
+
+        if config.logger.log_params_norm:
+            params_norm = calc_params_l2_norm(
+                model, model_config, use_megatron_fsdp=config.dist.use_megatron_fsdp
+            )
+        learning_rate = None
+        decoupled_learning_rate = None
+        for param_group in optimizer.param_groups:
+            if len(param_group) == 0:
+                continue
+            if param_group["is_decoupled_lr"]:
+                decoupled_learning_rate = param_group["lr"]
+            else:
+                learning_rate = param_group["lr"]
+        nemo_elapsed_time_per_iter_sec = None
+        if (
+            config.logger.log_interval
+            and global_state.train_state.step % config.logger.log_interval == 0
+            and nemo_style_iter_count > 0
+        ):
+            nemo_elapsed_time_per_iter_sec = nemo_style_iter_seconds_accum / nemo_style_iter_count
+        report_memory_flag = training_log(
+            loss_dict,
+            total_loss_dict,
+            learning_rate,
+            decoupled_learning_rate,
+            loss_scale,
+            report_memory_flag,
+            skipped_iter,
+            grad_norm,
+            params_norm,
+            num_zeros_in_grad,
+            config,
+            global_state,
+            history_wct,
+            model,
+            log_max_attention_logit,
+            nemo_elapsed_time_per_iter_sec=nemo_elapsed_time_per_iter_sec,
+        )
+        if config.logger.log_interval and global_state.train_state.step % config.logger.log_interval == 0:
+            _log_training_gpu_mem(
+                f"step {global_state.train_state.step}",
+                config.logger.memory_keys,
+            )
+        if nemo_elapsed_time_per_iter_sec is not None:
+            nemo_style_iter_seconds_accum = 0.0
+            nemo_style_iter_count = 0
+
+        # MLPerf logging: per-step train loss
+        if sft_logger is not None and not skipped_iter and loss_dict:
+            sft_logger.on_train_step(
+                step=global_state.train_state.step,
+                loss_dict=loss_dict,
+                lr=learning_rate,
+                consumed_samples=global_state.train_state.consumed_train_samples,
+            )
+
+        if (
+            global_state.train_state.do_valid
+            and train_config.eval_interval
+            and global_state.train_state.step % train_config.eval_interval == 0
+            and global_state.train_state.step > eval_skip_first_n * train_config.eval_interval
+        ):
+            if energy_monitor is not None:
+                energy_monitor.pause()
+            if should_toggle_forward_pre_hook:
+                disable_forward_pre_hook(model)
+                pre_hook_enabled = False
+            if train_config.manual_gc and train_config.manual_gc_eval:
+                # Collect all objects.
+                gc.collect()
+            prefix = f"iteration {global_state.train_state.step}"
+            timers("eval-time", log_level=0).start(barrier=True)
+            assert (
+                timer.consumed_samples == global_state.train_state.consumed_train_samples
+            ), "Timer and global_state sample mismatch"
+
+            if sft_logger is not None:
+                sft_logger.on_eval_start(global_state.train_state.consumed_train_samples)
+
+            should_exit, eval_loss_value = evaluate_and_print_results_custom(
+                global_state,
+                prefix,
+                forward_step_func,
+                valid_data_iterator,
+                model,
+                model_config,
+                verbose=False,
+                write_to_tensorboard=False,
+                process_non_loss_data_func=process_non_loss_data_func,
+                non_loss_data_func=non_loss_data_func,
+                throughput=timer.get_throughput(),
+            )
+
+            if sft_logger is not None and eval_loss_value is not None:
+                target_hit = sft_logger.on_eval_end(
+                    global_state.train_state.consumed_train_samples,
+                    eval_loss_value,
+                )
+                if target_hit:
+                    should_exit = True
+            if should_exit:
+                exit_code = 0
+            eval_duration += timers("eval-time").elapsed()
+            eval_iterations += train_config.eval_iters
+            timers("eval-time").stop()
+
+            if train_config.manual_gc and train_config.manual_gc_eval:
+                # Collect only the objects created and used in evaluation.
+                gc.collect(generation=0)
+            if should_toggle_forward_pre_hook:
+                enable_forward_pre_hook(model)
+                pre_hook_enabled = True
+            if energy_monitor is not None:
+                energy_monitor.resume()
+
+        # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
+        # Some of these only happen at specific iterations.
+        maybe_synchronize_training_step(config.train.train_sync_interval, global_state.train_state.step)
+        num_floating_point_operations_since_last_log_event = maybe_report_stragglers(
+            config.logger.log_interval,
+            bool(getattr(config.straggler, "log_straggler", False)),
+            straggler_timer,
+            global_state.train_state.step,
+            num_floating_point_operations_since_last_log_event,
+        )
+        maybe_check_weight_hash_across_dp_replicas(
+            model,
+            config.train.check_weight_hash_across_dp_replicas_interval,
+            global_state.train_state.step,
+            should_toggle_forward_pre_hook,
+        )
+        handle_profiling_stop(
+            config.profiling,
+            global_state.train_state.step,
+            torch.distributed.get_rank(),
+            prof,
+            nsys_nvtx_context,
+        )
+        maybe_run_manual_gc(
+            config.train.manual_gc,
+            config.train.manual_gc_interval,
+            global_state.train_state.step,
+        )
+        if should_exit:
+            break
+
+    training_wall_elapsed_s = time.perf_counter() - training_wall_start
+    log_rank_0(
+        f"Training loop finished: wall time {training_wall_elapsed_s:.2f} s "
+        f"({training_wall_elapsed_s / 60.0:.2f} min); "
+        f"final_iteration={global_state.train_state.step}; "
+        f"consumed_train_samples={global_state.train_state.consumed_train_samples}"
+    )
+
+    # MLPerf logging: end of training
+    if sft_logger is not None:
+        sft_logger.log_run_stop(global_state.train_state.consumed_train_samples)
+
+    _delete_cuda_graphs(cuda_graph_helper)
+
+    # Stop RPD profiler if still running
+    if rpd and rpd_status == "running":
+        log_rank_0(f"Stopping RPD profiling at training end (step {global_state.train_state.step})")
+        rpd.stop()
+        rpd_status = "finished"
+
+    # Flush TensorBoard writer if present (TB/WandB disabled in LoggerConfig for this recipe).
+    writer = global_state.tensorboard_logger
+    if writer:
+        writer.flush()
+
+    # Close out pre-hooks if using distributed optimizer and overlapped param gather.
+    if pre_hook_enabled:
+        disable_forward_pre_hook(model)
+
+    # This will finalize all unfinalized async request and terminate
+    # a persistent async worker if persistent ckpt worker is enabled
+    fault_tolerance.on_checkpointing_start(global_state)
+    maybe_finalize_async_save(
+        global_state=global_state, ckpt_cfg=config.checkpoint, blocking=True, terminate=True
+    )
+    fault_tolerance.on_checkpointing_end(global_state=global_state, is_async_finalization=True)
+
+    # Shutdown NVRx straggler detection if enabled
+    safe_shutdown_nvrx_straggler_manager(global_state.nvrx_straggler_manager)
+
+    if energy_monitor is not None:
+        energy_monitor.lap()
+        total_energy = energy_monitor.get_total()
+        log_rank_0(f"Total training energy (GPU): {total_energy / 1e6} MJ")
+        energy_monitor.shutdown()
+
+    # If any exit conditions (signal handler, duration, iterations) have been reached, exit.
+    if should_exit:
+        # Stop RPD profiler if still running before exit
+        if rpd and rpd_status == "running":
+            log_rank_0(f"Stopping RPD profiling before exit at step {global_state.train_state.step}")
+            rpd.stop()
+            rpd_status = "finished"
+
+        # Close NVIDIA DLFw Inspect if enabled
+        tensor_inspect_end_if_enabled(config.tensor_inspect)
+        maybe_finalize_async_save(
+            global_state=global_state, ckpt_cfg=config.checkpoint, blocking=True, terminate=True
+        )
+        wandb_writer = global_state.wandb_logger
+        if wandb_writer:
+            wandb_writer.finish()
+            fault_tolerance.shutdown(global_state)
+        if exit_code != 0:
+            sys.exit(exit_code)
+
+    # Close NVIDIA DLFw Inspect at clean finish
+    tensor_inspect_end_if_enabled(config.tensor_inspect)
+
+
+# ---------------------------------------------------------------------------
+# Install ``megatron_bridge_train_override`` on the megatron-bridge training
+# modules. When ``PRE_QUANTIZED_MODEL`` is enabled, the override is first
+# wrapped by ``install_pre_quantize_wrap`` so that the very first entry into
+# the training loop pre-quantizes all TE Linear / LayerNormLinear weights to
+# MXFP4 (and stashes FP8 copies on CPU for healing) before delegating to the
+# override. When disabled (default), the override is installed directly (no
+# wrap, zero runtime cost).
+# ---------------------------------------------------------------------------
+from megatron.bridge.training import train as _mb_train_mod
+
+from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.pre_quantize_mxfp4 import (
+    install_pre_quantize_wrap,
+)
+
+_installed_train = install_pre_quantize_wrap(megatron_bridge_train_override)
+setattr(megatron_bridge_train_override, "_primus_llama2_custom_train_override", True)
+
+_mb_train_mod.train = _installed_train
+
+try:
+    from megatron.bridge.training import pretrain as _mb_pretrain_mod
+
+    _mb_pretrain_mod.train = _installed_train
+except ImportError:
+    pass
+
+try:
+    from megatron.bridge.training import finetune as _mb_finetune_mod
+
+    if hasattr(_mb_finetune_mod, "train"):
+        _mb_finetune_mod.train = _installed_train
+except ImportError:
+    pass
+
+# Keep the historical name `train` bound to the training module so any
+# lingering `train.train = ...` expectations elsewhere resolve.
+train = _mb_train_mod
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py
new file mode 100644
index 000000000..f7d9902a4
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py
@@ -0,0 +1,518 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from NeMo MLPerf
+# (mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""NeMo-parity MXFP4 -> FP8 healing for Primus / Megatron-Bridge.
+
+Reference: ``mlperf_code_llama2_70b_0430/src/callbacks/custom_callbacks.py``
+``CustomCallback`` (``_healing_setup``, ``_set_quantized_params_cpu``,
+``_reset_full_iteration_cuda_graphs``, ``on_train_batch_end`` healing
+branch). This module mirrors that logic byte-for-byte; only the trigger
+surface differs:
+
+* NeMo's pre-quantize fires from ``on_train_start``; healing fires from
+  ``on_train_batch_end`` when ``trainer.global_step + 1 == healing_iter``.
+* Primus's pre-quantize fires from a megatron-bridge ``train()`` wrapper
+  (``pre_quantize_mxfp4``); healing fires when the recipe
+  calls ``apply_healing_after_step(model, model_config, train_state.step)``
+  immediately after ``train_state.step += 1`` -- the same arithmetic
+  identity as NeMo.
+
+All NeMo-side state lives on ``CustomCallback`` (``self.fp8_cpu_params``,
+``self.healing_lambda``); Primus has no callback instance, so we use module
+state. Configuration that NeMo reads from OmegaConf ``cfg.model.*`` is
+read here from environment variables with NeMo's same defaults
+(``HEALING_ITER`` / ``HEALING_PRECISION`` / ``ENABLE_TRANSPOSE_CACHE`` /
+``RESET_CG_AFTER_HEALING`` / ``FIRST_LAST_LAYERS_BF16`` /
+``NUM_LAYERS_AT_{START,END}_IN_BF16`` / ``STORE_GPU``).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, List
+
+import torch
+
+from primus.core.utils.module_utils import log_rank_0
+
+# ---------------------------------------------------------------------------
+# Module state.
+# ---------------------------------------------------------------------------
+
+# Layered FP8 (E4M3) CPU stash, populated from outside via
+# ``set_fp8_cpu_params`` (the bridge from ``pre_quantize_mxfp4._FP8_CPU_PARAMS``)
+# in the same shape as NeMo's ``CustomCallback.fp8_cpu_params``:
+# ``List[List[Float8Tensor]]`` indexed first by ``decoder.layers`` index and
+# then by intra-layer iteration over TE Linear / LayerNormLinear modules.
+_FP8_CPU_PARAMS: List[List[Any]] = []
+_HEALING_APPLIED: bool = False
+
+
+# ---------------------------------------------------------------------------
+# Env helpers (NeMo ``custom_callbacks._env_enabled`` parity).
+# ---------------------------------------------------------------------------
+
+
+def _env_enabled(name: str) -> bool:
+    """NeMo ``custom_callbacks.py:37-43``: missing/unrecognized = False."""
+    val = os.getenv(name, "").strip().lower()
+    if val in {"1", "true", "yes", "on"}:
+        return True
+    return False
+
+
+def healing_iter() -> int:
+    """``cfg.model.healing_iter`` analog (0 = healing disabled)."""
+    raw = os.getenv("HEALING_ITER", "0")
+    try:
+        return max(0, int(raw.strip()))
+    except ValueError:
+        return 0
+
+
+def healing_precision() -> str:
+    """``cfg.model.healing_precision`` analog: ``FP8_DS`` (default) or ``MXFP8``."""
+    return os.getenv("HEALING_PRECISION", "FP8_DS").strip().upper()
+
+
+# ---------------------------------------------------------------------------
+# Model traversal helpers.
+# ---------------------------------------------------------------------------
+
+
+def _extract_module(model):
+    """NeMo ``custom_callbacks.py:92-97``. Adapted: Primus passes
+    ``List[ModelChunk]``; pick chunk 0 then strip ``.module``.
+    """
+    m = model
+    if isinstance(m, (list, tuple)):
+        m = m[0]
+    while hasattr(m, "module"):
+        m = m.module
+    return m
+
+
+def _all_modules(model):
+    chunks = list(model) if isinstance(model, (list, tuple)) else [model]
+    for ch in chunks:
+        yield from ch.modules()
+
+
+def _te_linear_types():
+    import transformer_engine.pytorch as te
+
+    types = (te.Linear, te.LayerNormLinear)
+    if hasattr(te, "LayerNormMLP"):
+        types = types + (te.LayerNormMLP,)
+    return types
+
+
+# ---------------------------------------------------------------------------
+# NeMo ``_healing_setup`` / ``healing_lambda``.
+# ---------------------------------------------------------------------------
+
+
+def _build_healing_lambda(model_config: Any):
+    """NeMo ``CustomCallback._healing_setup`` (``custom_callbacks.py:176-187``).
+
+    Reads ``fp8_amax_history_len`` / ``fp8_amax_compute_algo`` /
+    ``fp8_reduce_amax`` / ``fp8_dot_product_attention`` from the live
+    ``TransformerConfig`` (NeMo reads them from OmegaConf ``cfg.model``;
+    Megatron-Bridge stores the same fields directly on ``model_config``).
+    NeMo defaults: ``amax_history_len=4``, ``amax_compute_algo="most_recent"``,
+    ``reduce_amax=False``, ``fp8_dpa=False``.
+    """
+    import transformer_engine.common.recipe as _recipe
+
+    prec = healing_precision()
+    if prec == "FP8_DS":
+        amax_len = int(getattr(model_config, "fp8_amax_history_len", 4) or 4)
+        amax_algo = str(getattr(model_config, "fp8_amax_compute_algo", "most_recent") or "most_recent")
+        reduce_amax = bool(getattr(model_config, "fp8_reduce_amax", False))
+        fp8_dpa = bool(getattr(model_config, "fp8_dot_product_attention", False))
+
+        log_rank_0(
+            "[mxfp4_healing] Healing recipe = DelayedScaling("
+            f"amax_history_len={amax_len}, "
+            f"amax_compute_algo={amax_algo!r}, "
+            f"reduce_amax={reduce_amax}, "
+            f"fp8_dpa={fp8_dpa})"
+        )
+
+        def _delayed(_config: Any):
+            return _recipe.DelayedScaling(
+                amax_history_len=amax_len,
+                amax_compute_algo=amax_algo,
+                reduce_amax=reduce_amax,
+                fp8_dpa=fp8_dpa,
+            )
+
+        return _delayed
+
+    if prec == "MXFP8":
+        log_rank_0("[mxfp4_healing] Healing recipe = MXFP8BlockScaling() (TE defaults)")
+
+        def _mxfp8(_config: Any):
+            return _recipe.MXFP8BlockScaling()
+
+        return _mxfp8
+
+    raise ValueError(f"Unsupported HEALING_PRECISION={prec!r} (expected FP8_DS or MXFP8)")
+
+
+# ---------------------------------------------------------------------------
+# NeMo ``_set_quantized_params_cpu`` (custom_callbacks.py:293-355).
+# ---------------------------------------------------------------------------
+
+
+def _set_quantized_params_cpu(model: Any, cpu_params: List[List[Any]], qtype: str) -> int:
+    """Restore CPU-stashed quantized weights to GPU and reinstall on TE modules.
+
+    Byte-equivalent to NeMo ``CustomCallback._set_quantized_params_cpu``
+    (``custom_callbacks.py:293-355``). Iterates ``decoder.layers`` in the same
+    order as ``_get_quantized_params_cpu`` (which built ``cpu_params``), so the
+    positional index ``cpu_params[layer_idx][param_idx]`` (consumed via
+    ``.pop(0)``) lines up with the ``param_idx``-th TE Linear /
+    LayerNormLinear / LayerNormMLP module encountered inside
+    ``layers[layer_idx].named_modules()``.
+
+    NeMo's two-pass design (set ``module._parameters['weight']=None`` for all
+    target modules first, then assign the restored weight) is preserved:
+    decoupling the unlink from the install lets the caching allocator release
+    the old MXFP4 weight storage (~75 GiB) before the FP8 GPU side is wired
+    in (~64 GiB), preventing both sets resident at peak.
+    """
+    extracted = _extract_module(model)
+    layers = extracted.decoder.layers
+    layer_count = len(layers)
+    dev = torch.cuda.current_device()
+
+    first_last = _env_enabled("FIRST_LAST_LAYERS_BF16")
+    n_start = int(os.getenv("NUM_LAYERS_AT_START_IN_BF16", "0") or "0")
+    n_end = int(os.getenv("NUM_LAYERS_AT_END_IN_BF16", "0") or "0")
+    store_gpu = _env_enabled("STORE_GPU")
+    enable_tc = _env_enabled("ENABLE_TRANSPOSE_CACHE")
+    use_fuser = bool(getattr(extracted.config, "use_transformer_engine_op_fuser", False))
+
+    te_types = _te_linear_types()
+
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+    xfer_stream = torch.cuda.Stream()
+    n_restored = 0
+
+    for layer_idx, layer in enumerate(layers):
+        if first_last:
+            if layer_idx < n_start or layer_idx >= layer_count - n_end:
+                continue
+
+        target_modules = []
+        for _name, module in layer.named_modules():
+            if not isinstance(module, te_types):
+                continue
+            if not hasattr(module, "weight"):
+                continue
+            target_modules.append(module)
+
+        with torch.no_grad():
+            for module in target_modules:
+                module._parameters["weight"] = None
+
+        for module in target_modules:
+            with torch.no_grad():
+                weight = cpu_params[layer_idx].pop(0)
+
+                if not store_gpu:
+                    with torch.cuda.stream(xfer_stream):
+                        if qtype == "FP8_DS":
+                            weight._data = weight._data.to(dev, non_blocking=True)
+                            if weight._transpose is not None:
+                                if enable_tc:
+                                    weight._transpose = weight._transpose.to(dev, non_blocking=True)
+                                else:
+                                    weight._transpose = None
+                        elif qtype == "MXFP4":
+                            weight._rowwise_data = weight._rowwise_data.to(dev, non_blocking=True)
+                            if weight._columnwise_data is not None:
+                                weight._columnwise_data = weight._columnwise_data.to(dev, non_blocking=True)
+                        else:
+                            raise ValueError(f"Unsupported quantization type: {qtype}")
+
+                weight.requires_grad = False
+                module._parameters["weight"] = weight
+                n_restored += 1
+
+        if use_fuser:
+            mlp = getattr(layer, "mlp", None)
+            if mlp is not None and hasattr(mlp, "_make_fused_impl"):
+                mlp._fused_impl = (mlp._make_fused_impl(),)
+            sa = getattr(layer, "self_attention", None)
+            if sa is not None:
+                for attr in ("linear_proj", "linear_qkv"):
+                    mod = getattr(sa, attr, None)
+                    if mod is not None and hasattr(mod, "_make_fused_branches"):
+                        mod._fused_branches = mod._make_fused_branches()
+
+    xfer_stream.synchronize()
+    return n_restored
+
+
+# ---------------------------------------------------------------------------
+# NeMo ``custom_llama.reset_fp8_state`` (custom_llama.py:77-92).
+# ---------------------------------------------------------------------------
+
+
+def _reset_fp8_state(model: Any) -> None:
+    """Force TE to re-initialize ``fp8_meta`` under the new recipe on next forward.
+
+    Without this, modules keep MXFP4 ``RecipeState`` while the recipe and
+    weights changed under them; the next GEMM dereferences scale/amax
+    pointers from the wrong RecipeState type. Mirrors NeMo's
+    ``custom_llama.reset_fp8_state``: MX recipes (MXFP4/MXFP8) use
+    ``RecipeState`` without ``.scale``, and TE's ``reset_fp8_meta_tensors``
+    assumes DelayedScaling-style state and raises against MX, so we detect
+    the MX case by inspecting ``fp8_meta`` and skip the reset for those
+    modules.
+    """
+
+    def reset_fp8(m: Any) -> None:
+        if not hasattr(m, "fp8_initialized"):
+            return
+        m.fp8_initialized = False
+        # MX recipes (MXFP4/MXFP8) use RecipeState without `.scale`; TE's
+        # reset_fp8_meta_tensors assumes DelayedScaling-style state and raises.
+        fp8_meta = getattr(m, "fp8_meta", None) or {}
+        for key in ("scaling_fwd", "scaling_bwd"):
+            state = fp8_meta.get(key)
+            if state is not None and not hasattr(state, "scale"):
+                return
+        m.reset_fp8_meta_tensors()
+
+    for m in _all_modules(model):
+        reset_fp8(m)
+
+
+# ---------------------------------------------------------------------------
+# NeMo ``_reset_full_iteration_cuda_graphs`` (custom_callbacks.py:148-174).
+# ---------------------------------------------------------------------------
+
+
+def _reset_full_iteration_cuda_graphs() -> None:
+    """Drop Megatron full-iteration CUDA graph captures before the recipe swap.
+
+    After MXFP4 -> FP8 weight swap and recipe change, replaying old graphs
+    is invalid. ``RESET_CG_AFTER_HEALING`` (NeMo: ``cfg.model.reset_cg_after_healing``)
+    additionally resets per-stage warmup step counters so
+    ``FullCudaGraphWrapper`` re-runs ``cuda_graph_warmup_steps`` before
+    re-capture.
+    """
+    try:
+        from megatron.core.full_cuda_graph import FullCudaGraphWrapper
+    except ImportError:
+        return
+
+    FullCudaGraphWrapper.cuda_graph["training"] = None
+    FullCudaGraphWrapper.cuda_graph["validation"] = None
+    FullCudaGraphWrapper.result["training"] = None
+    FullCudaGraphWrapper.result["validation"] = None
+
+    if _env_enabled("RESET_CG_AFTER_HEALING"):
+        FullCudaGraphWrapper.curr_iteration["training"] = 0
+        FullCudaGraphWrapper.curr_iteration["validation"] = 0
+
+
+# ---------------------------------------------------------------------------
+# Public stash bridge.
+# ---------------------------------------------------------------------------
+
+
+def set_fp8_cpu_params(layered: List[List[Any]]) -> None:
+    """Hand off NeMo-style ``List[List[Float8Tensor]]`` from pre-quantize.
+
+    Mirrors the assignment ``self.fp8_cpu_params = self._get_quantized_params_cpu(...)``
+    in NeMo ``CustomCallback._pre_quantize_model``. Storing a reference (not
+    a copy) means ``_set_quantized_params_cpu``'s ``.pop(0)`` mutates the same
+    list as the caller, identical to NeMo.
+    """
+    global _FP8_CPU_PARAMS, _HEALING_APPLIED
+    _FP8_CPU_PARAMS = layered
+    _HEALING_APPLIED = False
+
+
+# ---------------------------------------------------------------------------
+# NeMo ``on_train_batch_end`` healing branch (custom_callbacks.py:620-677).
+# ---------------------------------------------------------------------------
+
+
+def apply_healing_after_step(model: Any, model_config: Any, train_state_step: int) -> None:
+    """Run NeMo's ``on_train_batch_end`` healing branch.
+
+    Trigger condition mirrors NeMo:
+      ``if trainer.global_step + 1 == self.cfg.model.healing_iter:``
+    Primus's ``train_state.step`` is incremented before this call (in the
+    recipe's ``megatron_bridge_train_override``), so it carries the same
+    value as NeMo's ``trainer.global_step``.
+
+    Steps below are 1:1 with the NeMo branch.
+    """
+    global _HEALING_APPLIED
+
+    hi = healing_iter()
+    if hi <= 0 or _HEALING_APPLIED:
+        return
+    if train_state_step + 1 != hi:
+        return
+
+    if not _FP8_CPU_PARAMS:
+        log_rank_0(
+            "[mxfp4_healing] HEALING_ITER is set but the FP8 CPU stash is empty; "
+            "either PRE_QUANTIZED_MODEL=True is missing or pre-quantize did not run. "
+            "Skipping healing."
+        )
+        _HEALING_APPLIED = True
+        return
+
+    import torch.distributed as dist
+
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+    # 1. NeMo: ``self._reset_full_iteration_cuda_graphs(trainer)``.
+    _reset_full_iteration_cuda_graphs()
+
+    # 2. NeMo: ``self._set_quantized_params_cpu(trainer.model, self.fp8_cpu_params, self.healing_precision)``.
+    n_restored = _set_quantized_params_cpu(model, _FP8_CPU_PARAMS, healing_precision())
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+    # 3. NeMo: patch ``megatron.core.fp4_utils`` to the healing recipe and clear ``_mxfp4_phase``.
+    import megatron.core.fp4_utils as fp4u
+
+    fp4u.get_fp4_recipe = _build_healing_lambda(model_config)
+    fp4u._mxfp4_phase = False
+
+    # 4. NeMo: when ``ENABLE_TRANSPOSE_CACHE`` is off, tell every TE Linear to stop
+    # keeping a transpose cache (custom_callbacks.py:637-642).
+    if not _env_enabled("ENABLE_TRANSPOSE_CACHE"):
+        te_types = _te_linear_types()
+        for m in _all_modules(model):
+            if isinstance(m, te_types):
+                m.keep_fp8_weight_transpose_cache = False
+
+    # 5. NeMo: drop activation recompute on the live TransformerConfig
+    # (custom_callbacks.py:644-648).
+    extracted = _extract_module(model)
+    cfg = extracted.config
+    if cfg.recompute_granularity is not None:
+        cfg.recompute_granularity = None
+        cfg.recompute_num_layers = None
+        cfg.recompute_method = None
+
+    # 6. NeMo: clear MXFP4 caches on every module (custom_callbacks.py:650-656).
+    for m in _all_modules(model):
+        if hasattr(m, "_mxfp4_weight_cache"):
+            del m._mxfp4_weight_cache
+        if hasattr(m, "_mxfp4_persist_columnwise"):
+            del m._mxfp4_persist_columnwise
+
+    # 6b. Primus delta vs NeMo: clear ``save_original_input`` on every TE module.
+    #
+    # Megatron-LM shipped in Primus's tree (third_party/Megatron-Bridge/3rdparty/
+    # Megatron-LM) sets ``module.save_original_input = True`` on every TE
+    # ``linear_proj`` (and some ``mlp.linear_fc1``) when ``config.fp4`` is
+    # active -- it's an MX-only optimization that lets TE re-quantize at
+    # backward time from the saved BF16 input (legal under MXFP4BlockScaling
+    # because the per-block scale is reproducible from the BF16 input).
+    # ``DelayedScaling`` derives its FP8 scale from amax history that is
+    # updated each step, so the saved BF16 input cannot reproduce the
+    # forward FP8 input. TE guards this combination with a hard error::
+    #
+    #   RuntimeError: DelayedScaling recipe is not supported with save_original_input
+    #
+    # ...which fires on the very next forward after we patch
+    # ``get_fp4_recipe -> FP8_DS``. NeMo MLPerf doesn't hit this because
+    # their (older) Megatron-LM lacks the FP4 branch that sets the flag in
+    # the first place. We reset it defensively to mirror NeMo's behavior.
+    n_save_orig_cleared = 0
+    for m in _all_modules(model):
+        if getattr(m, "save_original_input", False):
+            try:
+                m.save_original_input = False
+                n_save_orig_cleared += 1
+            except Exception:  # noqa: BLE001
+                pass
+
+    # 7. NeMo: ``reset_fp8_state(trainer.lightning_module)`` (custom_callbacks.py:668).
+    _reset_fp8_state(model)
+
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+    _HEALING_APPLIED = True
+
+    log_rank_0(
+        f"[mxfp4_healing] healing applied at train_state.step={train_state_step} "
+        f"(step+1==HEALING_ITER={hi}); recipe={healing_precision()}; "
+        f"restored {n_restored} FP8 weights; "
+        f"cleared save_original_input on {n_save_orig_cleared} TE modules."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Status helpers + test hygiene.
+# ---------------------------------------------------------------------------
+
+
+def is_fp8_healing_phase() -> bool:
+    """True after ``apply_healing_after_step`` has fired (TE weights are FP8)."""
+    return _HEALING_APPLIED
+
+
+def reset_healing_state() -> None:
+    """Test / multi-run hygiene: forget the stash and the healed flag."""
+    global _FP8_CPU_PARAMS, _HEALING_APPLIED
+    _FP8_CPU_PARAMS = []
+    _HEALING_APPLIED = False
+
+
+# ---------------------------------------------------------------------------
+# Backwards-compatible shims for existing callers.
+#
+# These keep ``pre_quantize_mxfp4`` and ``llama2_custom`` importable without
+# edits. The new
+# code is ``set_fp8_cpu_params`` (above); the legacy names below just
+# delegate / no-op.
+# ---------------------------------------------------------------------------
+
+
+def _set_ordered_fp8_stash_from_layered(model: Any, fp8_cpu_params: List[List[Any]]) -> None:
+    """Legacy alias for ``set_fp8_cpu_params``. ``model`` is unused; we keep
+    the layered stash exactly as NeMo does.
+    """
+    del model  # unused; layered stash is enough (NeMo parity)
+    set_fp8_cpu_params(fp8_cpu_params)
+
+
+def log_healing_env_banner_once() -> None:
+    """No-op (NeMo doesn't have this; logging removed from main flow)."""
+    return
+
+
+def log_training_step_phase(train_state_step: int) -> None:
+    """No-op (NeMo doesn't have this; logging removed from main flow)."""
+    del train_state_step
+    return
+
+
+def run_fp8_warmup_for_kernel_jit(*_args: Any, **_kwargs: Any) -> None:
+    """No-op. NeMo's ``_pre_ttt_fp8_warmup`` call site is commented out
+    (custom_callbacks.py:509-510), so the warmup is dead in NeMo too.
+    """
+    return
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py
new file mode 100644
index 000000000..2f86448db
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py
@@ -0,0 +1,594 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from NeMo MLPerf
+# (mlperf-training-6-0/llama2_sft/nemo/src/custom_llama.py).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NeMo-equivalent training/validation loss for Primus (exact port of
+``MaskedTokenLossReduction.forward + reduce``).
+
+Background
+----------
+NeMo MLPerf (`mlperf_code_llama2_70b_0430/src/custom_llama.py`) overrides
+``training_loss_reduction`` / ``validation_loss_reduction`` on the model
+with ``MaskedTokenLossReduction``. The class has **two distinct paths**
+depending on its constructor flags:
+
+1. ``validation_step=False`` (training) **or** ``val_drop_last=True``:
+    * per microbatch:  per-sample token-mean
+        ``loss_for_ub[i] = sum(losses[i] * mask[i]) / sum(mask[i])``
+    * across microbatches:  ``concat(per_sample_means).mean()``
+2. ``validation_step=True`` AND ``val_drop_last=False`` (NeMo's default for
+   eval; see ``CustomLlamaModel.validation_loss_reduction``):
+    * per microbatch:  per-sample token-mean (same as path 1) **plus** a
+      0/1 indicator ``num_valid_tokens_in_ub = (num_valid_tokens > 0)``,
+      stacked to produce ``[loss_for_ub, valid_indicator]``.
+    * across microbatches:  ``vstack(stacked).sum(dim=0)`` → returns
+      ``[total_loss_sum, total_valid_sample_count]``.
+    * the eval-time CP all-reduce on ``loss_mask`` / ``loss_for_ub`` is
+      gated by ``disabled_cp_for_eval`` (i.e., CP is **off** for eval).
+
+This is **not** equivalent to Megatron-Bridge's default
+``masked_next_token_loss`` which computes a *global* token-mean across all
+tokens in all samples. When samples have very different numbers of
+unmasked tokens (typical for SFT/LoRA runs), the two objectives differ
+and produce different gradient directions.
+
+Implementation
+--------------
+We expose a NeMo-equivalent per-microbatch loss function that returns the
+3-tuple Megatron-Core's pipeline schedule expects::
+
+    (loss, num_tokens, {"lm loss": [loss_sum, count]})
+
+For both paths, ``loss = loss_for_ub.sum()`` (sum over the microbatch's
+per-sample means). The choice of ``num_tokens`` and the ``[sum, count]``
+reporting tensor depends on the path:
+
+* **Train / val_drop_last=True path**: ``num_tokens = num_samples = B``.
+  After Megatron-Bridge's ``output /= clamp(num_tokens, 1); output /=
+  num_microbatches`` and DP+microbatch ``[sum, count]`` aggregation, the
+  effective backward target and reported loss are both
+  ``mean over (microbatch, sample) of per_sample_mean`` -- byte-identical
+  to NeMo's ``concat(per_sample_means).mean()`` when ``B`` is uniform.
+
+* **val_drop_last=False path** (eval, NeMo default): ``num_tokens =
+  num_samples = sum_i (num_valid_tokens_i > 0)``. Empty-mask samples
+  (those whose entire sequence is masked) are excluded from the
+  denominator -- mirroring NeMo's ``num_valid_tokens_in_ub`` indicator.
+  Because ``loss_for_ub`` is already zeroed for those samples, the
+  numerator (sum) is unaffected. Megatron-Bridge's eval reduction does
+  ``total_loss_sum / total_valid_count`` across DP+CP+microbatch, which
+  matches NeMo's ``vstack(...).sum(dim=0)`` semantics.
+
+Train vs. eval detection
+~~~~~~~~~~~~~~~~~~~~~~~~
+Megatron-Bridge has a single ``forward_step`` for both train and eval and
+wraps the eval pass in ``with torch.no_grad():`` (see
+``megatron/bridge/training/eval.py:118``). We detect eval at call time
+with ``torch.is_grad_enabled()`` -- a Megatron-Bridge-specific but stable
+signal.
+
+CP gating (``cp_eval`` / ``disabled_cp_for_eval``)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+NeMo's ``cp_eval`` parameter (``CustomLlamaModel.config.cp_eval``) toggles
+``disabled_cp_for_eval``: when CP is **disabled** for eval, the CP
+all-reduces on ``loss_mask`` and ``loss_for_ub`` must be **skipped** (the
+data isn't actually sharded across CP ranks even though
+``parallel_state.get_context_parallel_world_size() > 1``). We expose this
+through env var ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL`` (default ``0``,
+i.e. CP-on-for-eval).
+
+``val_drop_last`` is exposed via ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST``
+(default ``0``, matching NeMo's eval setup ``val_drop_last=False``).
+
+Reporting parity
+----------------
+Megatron-Bridge's training & eval aggregation computes ``loss_sum /
+count`` over (microbatch, DP-rank, CP-rank). With our return shape:
+
+* Train / val_drop_last=True:
+  ``sum_(rank, microbatch) sum_samples(per_sample_mean) /
+   sum_(rank, microbatch) B  =  mean per_sample_mean``
+* val_drop_last=False:
+  ``sum_(rank, microbatch) sum_samples(per_sample_mean) /
+   sum_(rank, microbatch) count_valid_samples
+   = total_loss_sum / total_valid_sample_count`` (matches NeMo).
+
+Activation
+----------
+Auto-installs at import time. Disable by setting ``PRIMUS_NEMO_LOSS=0``.
+
+Env vars
+~~~~~~~~
+* ``PRIMUS_NEMO_LOSS``                         (default ``1``): master switch.
+* ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST``           (default ``0``): if ``1``,
+  treat eval like train (NeMo's ``val_drop_last=True``).
+* ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL``    (default ``0``): if ``1``,
+  skip the CP all-reduces during eval (NeMo's ``cp_eval`` set).
+"""
+
+from __future__ import annotations
+
+import os
+from functools import partial
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from megatron.core import parallel_state
+from megatron.core.rerun_state_machine import get_rerun_state_machine
+
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _safe_log_rank_0(msg: str) -> None:
+    """Log on rank 0; fall back to print if the global logger is not ready yet."""
+    try:
+        log_rank_0(msg)
+    except AttributeError:
+        if os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")) in ("0", "", None):
+            print(msg, flush=True)
+
+
+# Same value Megatron-Bridge uses in ``masked_next_token_loss``.
+_SPIKY_LOSS_FACTOR: int = 10
+
+# Module-level flag so callers can verify the patch landed.
+_INSTALLED: bool = False
+
+
+def _enabled() -> bool:
+    """Return True when the NeMo-equivalent loss should be active."""
+    flag = os.environ.get("PRIMUS_NEMO_LOSS", "1").strip().lower()
+    return flag not in ("0", "false", "no", "off")
+
+
+def _env_bool(name: str, default: bool = False) -> bool:
+    """Parse a boolean env var (``1/true/yes/on`` -> True; anything else -> False)."""
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in ("1", "true", "yes", "on")
+
+
+def _is_validation_step() -> bool:
+    """Detect whether we're inside a validation forward pass.
+
+    Megatron-Bridge wraps the entire eval loop in
+    ``with torch.no_grad():`` (``megatron/bridge/training/eval.py:118``)
+    and reuses the same ``forward_step`` for both train and eval, so
+    ``torch.is_grad_enabled()`` is a stable, framework-internal signal
+    that distinguishes the two.
+    """
+    return not torch.is_grad_enabled()
+
+
+def _val_drop_last() -> bool:
+    """``val_drop_last`` flag, mirrored to NeMo's ``MaskedTokenLossReduction``.
+
+    NeMo's ``CustomLlamaModel.validation_loss_reduction`` constructs the
+    eval reduction with ``val_drop_last=False`` (drop_last=False so all
+    eval samples count). Default here matches that. Set
+    ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST=1`` to fall back to the train-style
+    averaging (``concat(per_sample_means).mean()`` over all microbatches).
+    """
+    return _env_bool("PRIMUS_NEMO_LOSS_VAL_DROP_LAST", default=False)
+
+
+def _disabled_cp_for_eval() -> bool:
+    """``disabled_cp_for_eval`` flag, mirrored to NeMo's ``cp_eval`` parameter.
+
+    When CP is disabled for the eval forward pass (i.e., the model
+    doesn't actually shard across CP ranks during eval, even though
+    ``parallel_state.get_context_parallel_world_size() > 1``), the CP
+    all-reduces on ``loss_mask`` and ``loss_for_ub`` MUST be skipped --
+    otherwise we double-count along the CP dimension. NeMo encodes this
+    via ``cp_eval is not None`` in ``MaskedTokenLossReduction.__init__``.
+    Default here is ``False`` (CP-on-for-eval); set
+    ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL=1`` if your recipe disables
+    CP during eval.
+    """
+    return _env_bool("PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL", default=False)
+
+
+def _cp_eval_value() -> Optional[int]:
+    """Return NeMo's ``cp_eval`` integer value for the validation reduction.
+
+    NeMo's ``CustomLlamaModel.validation_loss_reduction`` passes
+    ``cp_eval=self.config.cp_eval``, which is the *integer*
+    ``cfg.model.eval_cp`` (or ``None``). The class only checks
+    ``cp_eval is not None`` to set ``disabled_cp_for_eval``, so the
+    integer value itself never affects the loss math -- but we keep the
+    int-or-None contract to remain byte-identical to NeMo.
+
+    Resolution order:
+      1. ``PRIMUS_NEMO_LOSS_CP_EVAL`` (preferred). Set to an integer
+         (e.g. ``1``) to mark "CP is disabled for eval". Empty / unset
+         leaves it as ``None``.
+      2. Fallback: ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL`` (boolean).
+         If set truthy, returns ``1``; else ``None``. Provided for
+         backwards compat with the function-based env var.
+    """
+    raw = os.environ.get("PRIMUS_NEMO_LOSS_CP_EVAL")
+    if raw is not None and raw.strip():
+        try:
+            return int(raw)
+        except ValueError:
+            return 1 if raw.strip().lower() in ("1", "true", "yes", "on") else None
+    return 1 if _disabled_cp_for_eval() else None
+
+
+# ============================================================================
+# Byte-identical port of NeMo's ``MaskedTokenLossReduction`` class.
+#
+# Source:
+#     mlperf_code_llama2_70b_0430/src/custom_llama.py:160-221
+# Original imports replaced 1:1:
+#     nemo.lightning.megatron_parallel.MegatronLossReduction -> stub below
+#     megatron.core.parallel_state                            -> identical
+#     torch                                                   -> identical
+#
+# The class body is reproduced as-is so the per-microbatch loss math, the
+# CP / DP all-reduce gating, and the train vs. val branching are
+# guaranteed identical to NeMo. The Megatron-Bridge wiring (which
+# expects a (loss, num_tokens, dict) 3-tuple instead of NeMo's
+# (loss_for_ub, dict)) lives in ``nemo_masked_next_token_loss`` further
+# down -- it instantiates the right train / val singleton, calls
+# ``forward(batch, per_token_losses)``, and repackages the return.
+# ============================================================================
+
+
+class MegatronLossReduction:
+    """Stub of ``nemo.lightning.megatron_parallel.MegatronLossReduction``.
+
+    The NeMo base class is just an interface marker (its ``forward`` /
+    ``reduce`` are abstract). We reproduce it as an empty class so the
+    ported ``MaskedTokenLossReduction`` keeps the same MRO for any
+    callers that introspect via ``isinstance``.
+    """
+
+
+class MaskedTokenLossReduction(MegatronLossReduction):
+    """Byte-identical port of NeMo's ``MaskedTokenLossReduction``.
+
+    Identical to ``custom_llama.MaskedTokenLossReduction`` (lines 160-221
+    of ``mlperf_code_llama2_70b_0430/src/custom_llama.py``):
+
+    * ``__init__`` signature, attribute names, and ``disabled_cp_for_eval``
+      derivation are unchanged.
+    * ``forward(batch, per_token_losses)`` returns NeMo's 2-tuple
+      ``(loss_for_ub, dict)`` where the dict is either ``{"avg":
+      reduced_loss}`` (train / ``val_drop_last=True``) or
+      ``{"loss_sum_and_ub_size": <(B, 2) all-reduced tensor>}``
+      (validation with ``val_drop_last=False``).
+    * ``reduce(losses_reduced_per_micro_batch)`` aggregates microbatch
+      outputs identically (``concat(...).mean()`` for ``"avg"``,
+      ``vstack(...).sum(dim=0)`` for ``"loss_sum_and_ub_size"``).
+
+    To use this class with Megatron-Bridge's function-based loss API
+    (which expects ``(loss, num_tokens, {"lm loss": [sum, count]})``),
+    go through ``nemo_masked_next_token_loss`` -- it builds an instance
+    (one singleton for train, one for val), calls ``forward``, and
+    repackages the return.
+    """
+
+    def __init__(
+        self,
+        validation_step: bool = False,
+        val_drop_last: bool = True,
+        cp_eval: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.validation_step = validation_step
+        self.train_step = not validation_step
+        self.val_drop_last = val_drop_last
+        self.disabled_cp_for_eval = cp_eval is not None
+
+    def forward(
+        self,
+        batch: Dict[str, torch.Tensor],
+        per_token_losses: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        if isinstance(per_token_losses, tuple):
+            per_token_losses, loss_mask = per_token_losses
+            batch["loss_mask"] = loss_mask
+        masked_losses = per_token_losses * batch["loss_mask"]
+
+        cp_size = parallel_state.get_context_parallel_world_size()
+        if cp_size > 1 and (self.train_step or not self.disabled_cp_for_eval):
+            torch.distributed.all_reduce(
+                batch["loss_mask"], group=parallel_state.get_context_parallel_group()
+            )
+
+        num_valid_tokens = batch["loss_mask"].sum(1)
+        loss_for_ub = torch.sum(masked_losses, dim=1) / num_valid_tokens
+        loss_for_ub = torch.where(num_valid_tokens == 0, torch.zeros_like(loss_for_ub), loss_for_ub)
+
+        if cp_size > 1 and (self.train_step or not self.disabled_cp_for_eval):
+            torch.distributed.all_reduce(loss_for_ub, group=parallel_state.get_context_parallel_group())
+
+        if self.validation_step and not self.val_drop_last:
+            num_valid_tokens_in_ub = (num_valid_tokens > 0).long()
+            loss_sum_and_ub_size_all_gpu = torch.stack(
+                [loss_for_ub.clone().detach(), num_valid_tokens_in_ub.clone().detach()],
+                dim=1,
+            )
+            if self.disabled_cp_for_eval:
+                torch.distributed.all_reduce(loss_sum_and_ub_size_all_gpu)
+            else:
+                torch.distributed.all_reduce(
+                    loss_sum_and_ub_size_all_gpu,
+                    group=parallel_state.get_data_parallel_group(),
+                )
+            return loss_for_ub, {"loss_sum_and_ub_size": loss_sum_and_ub_size_all_gpu}
+
+        reduced_loss = loss_for_ub
+        return loss_for_ub, {"avg": reduced_loss}
+
+    def reduce(self, losses_reduced_per_micro_batch: List[Dict[str, torch.Tensor]]) -> torch.Tensor:
+        if losses_reduced_per_micro_batch:
+            if "avg" in losses_reduced_per_micro_batch[0]:
+                loss_tensors_list = [loss_reduced["avg"] for loss_reduced in losses_reduced_per_micro_batch]
+                loss_tensor = torch.concat(loss_tensors_list)
+                return loss_tensor.mean()
+
+            loss_sum_tensors_list: List[torch.Tensor] = [
+                loss_sum["loss_sum_and_ub_size"] for loss_sum in losses_reduced_per_micro_batch
+            ]
+            loss_sum = (
+                torch.vstack(loss_sum_tensors_list).sum(dim=0)
+                if len(loss_sum_tensors_list) > 0
+                else torch.tensor([0.0, 0.0], device=torch.cuda.current_device())
+            )
+            return loss_sum
+
+        return torch.tensor(0.0, device=torch.cuda.current_device())
+
+
+# ----------------------------------------------------------------------------
+# Lazy singletons mirroring NeMo's ``CustomLlamaModel`` setup:
+#
+#     self._training_loss_reduction   = MaskedTokenLossReduction()
+#     self._validation_loss_reduction = MaskedTokenLossReduction(
+#         validation_step=True,
+#         val_drop_last=False,
+#         cp_eval=self.config.cp_eval,
+#     )
+#
+# (See ``mlperf_code_llama2_70b_0430/src/custom_llama.py:143-157``.)
+# We construct lazily because env vars may not be set at import time.
+# ----------------------------------------------------------------------------
+
+_TRAIN_REDUCTION: Optional[MaskedTokenLossReduction] = None
+_VAL_REDUCTION: Optional[MaskedTokenLossReduction] = None
+
+
+def _get_train_reduction() -> MaskedTokenLossReduction:
+    """Return the per-process training ``MaskedTokenLossReduction`` singleton."""
+    global _TRAIN_REDUCTION
+    if _TRAIN_REDUCTION is None:
+        _TRAIN_REDUCTION = MaskedTokenLossReduction()
+    return _TRAIN_REDUCTION
+
+
+def _get_val_reduction() -> MaskedTokenLossReduction:
+    """Return the per-process validation ``MaskedTokenLossReduction`` singleton."""
+    global _VAL_REDUCTION
+    if _VAL_REDUCTION is None:
+        _VAL_REDUCTION = MaskedTokenLossReduction(
+            validation_step=True,
+            val_drop_last=_val_drop_last(),
+            cp_eval=_cp_eval_value(),
+        )
+    return _VAL_REDUCTION
+
+
+def reset_loss_reduction_singletons() -> None:
+    """Drop cached singletons so subsequent calls re-read env vars.
+
+    Useful for tests or for recipes that reconfigure ``cp_eval`` /
+    ``val_drop_last`` after import time.
+    """
+    global _TRAIN_REDUCTION, _VAL_REDUCTION
+    _TRAIN_REDUCTION = None
+    _VAL_REDUCTION = None
+
+
+def nemo_masked_next_token_loss(
+    loss_mask: torch.Tensor,
+    output_tensor: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+    check_for_nan_in_loss: bool = True,
+    check_for_spiky_loss: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
+    """Megatron-Bridge adapter for the byte-identical ``MaskedTokenLossReduction``.
+
+    All loss math (CP / DP all-reduces, val_drop_last branching,
+    train vs. eval gating) lives in the class above and is identical
+    to NeMo. This adapter does three things:
+
+    1. Picks the right per-process singleton (train vs. val) based on
+       ``torch.is_grad_enabled()`` -- Megatron-Bridge wraps eval in
+       ``with torch.no_grad():`` (``eval.py:118``) and reuses the same
+       ``forward_step`` for both, so this is the only stable signal.
+    2. Calls ``MaskedTokenLossReduction.forward(batch, per_token_losses)``
+       with a ``batch = {"loss_mask": loss_mask}`` shim. The class
+       mutates ``batch["loss_mask"]`` in-place via the CP all-reduce --
+       same as NeMo, so we accept the side effect.
+    3. Repackages NeMo's 2-tuple ``(loss_for_ub, dict)`` return into the
+       Megatron-Core schedule's expected 3-tuple
+       ``(loss, num_tokens, {"lm loss": [sum, count]})``:
+
+       * ``"avg"`` branch (train / ``val_drop_last=True``):
+           - ``loss = loss_for_ub.sum()``, ``num_tokens = B``.
+             Megatron-Core computes ``loss / clamp(num_tokens, 1) /
+             num_microbatches`` for backward, giving
+             ``mean over (microbatch, sample) of per_sample_mean`` --
+             byte-identical to NeMo's ``concat(...).mean()`` reduction
+             when ``B`` is uniform across microbatches and DP ranks.
+       * ``"loss_sum_and_ub_size"`` branch (val w/ ``val_drop_last=False``):
+           - The class already DP-all-reduced the (B, 2) tensor inside
+             ``forward``. We sum dim=0 (matches NeMo's ``vstack(...).sum(dim=0)``
+             reduce) to get ``[total_loss_sum, total_valid_count]`` for
+             this microbatch, then return it as the reporting tensor.
+             Megatron-Bridge's eval reduction then sums it across
+             microbatches and DP+CP (``eval.py:183``); since both
+             entries scale by the same factor in any further reduction,
+             the final ratio ``loss_sum / count`` is identical to
+             NeMo's reported eval loss.
+
+    NaN / Inf / spike validation runs on the local per-sample sum so it
+    triggers regardless of which branch the class took.
+    """
+    if isinstance(output_tensor, tuple):
+        per_token_losses, dynamic_mask = output_tensor
+        loss_mask = dynamic_mask
+    else:
+        per_token_losses = output_tensor
+
+    per_token_losses = per_token_losses.float()
+    loss_mask = loss_mask.float()
+
+    batch_size = loss_mask.shape[0]
+    per_token_losses = per_token_losses.view(batch_size, -1)
+    loss_mask = loss_mask.view(batch_size, -1)
+
+    reduction = _get_val_reduction() if _is_validation_step() else _get_train_reduction()
+
+    batch: Dict[str, torch.Tensor] = {"loss_mask": loss_mask}
+    loss_for_ub, reduced_dict = reduction.forward(batch, per_token_losses)
+
+    loss_sum = loss_for_ub.sum()
+
+    rerun_state_machine = get_rerun_state_machine()
+    if check_for_nan_in_loss:
+        rerun_state_machine.validate_result(
+            result=loss_sum,
+            rejection_func=torch.isnan,
+            message="found NaN in local forward loss calculation",
+            tolerance=0.0,
+            fatal=True,
+        )
+        rerun_state_machine.validate_result(
+            result=loss_sum,
+            rejection_func=torch.isinf,
+            message="found Inf in local forward loss calculation",
+            tolerance=0.0,
+            fatal=True,
+        )
+    if check_for_spiky_loss:
+        rerun_state_machine.validate_result(
+            result=loss_sum,
+            rejection_func=partial(
+                rerun_state_machine.is_unexpectedly_large,
+                threshold=_SPIKY_LOSS_FACTOR,
+                context="loss",
+            ),
+            message="Spiky loss",
+            tolerance=0.0,
+            fatal=False,
+        )
+
+    if "loss_sum_and_ub_size" in reduced_dict:
+        agg = reduced_dict["loss_sum_and_ub_size"].sum(dim=0).float()
+        agg_loss_sum = agg[0]
+        agg_count = agg[1]
+        num_samples = agg_count.to(torch.int)
+        reporting_loss = torch.cat(
+            [
+                agg_loss_sum.clone().detach().view(1),
+                agg_count.clone().detach().view(1),
+            ]
+        )
+        return (agg_loss_sum, num_samples, {"lm loss": reporting_loss})
+
+    num_samples = torch.tensor(loss_for_ub.numel(), dtype=torch.int, device=loss_for_ub.device)
+    reporting_loss = torch.cat(
+        [
+            loss_sum.clone().detach().view(1),
+            num_samples.detach().view(1).float(),
+        ]
+    )
+    return (loss_sum, num_samples, {"lm loss": reporting_loss})
+
+
+def create_nemo_masked_next_token_loss_function(
+    loss_mask: torch.Tensor,
+    check_for_nan_in_loss: bool,
+    check_for_spiky_loss: bool,
+) -> partial:
+    """Factory matching Megatron-Bridge's ``create_masked_next_token_loss_function``."""
+    return partial(
+        nemo_masked_next_token_loss,
+        loss_mask,
+        check_for_nan_in_loss=check_for_nan_in_loss,
+        check_for_spiky_loss=check_for_spiky_loss,
+    )
+
+
+def install_nemo_loss_if_enabled() -> bool:
+    """Monkey-patch Megatron-Bridge step modules to use the NeMo-equivalent loss.
+
+    Affected bindings:
+
+    * ``megatron.bridge.training.vlm_step._create_loss_function`` ->
+      ``create_nemo_masked_next_token_loss_function`` (used by
+      ``MegatronBridgePosttrainTrainer.train()`` -> ``finetune(...,
+      forward_step_func=vlm_step.forward_step)``).
+    * ``megatron.bridge.training.gpt_step.masked_next_token_loss`` ->
+      ``nemo_masked_next_token_loss`` (used by other gpt_step-based recipes
+      and any test callers).
+
+    Idempotent. Gated by env var ``PRIMUS_NEMO_LOSS`` (default: enabled;
+    disable with ``PRIMUS_NEMO_LOSS=0``).
+
+    Returns True when patching took effect, False if disabled or already
+    installed.
+    """
+    global _INSTALLED
+    if _INSTALLED:
+        return False
+    if not _enabled():
+        _safe_log_rank_0(
+            "[primus.nemo_loss] PRIMUS_NEMO_LOSS disabled; using Megatron-Bridge "
+            "default global-token-mean loss"
+        )
+        return False
+
+    patched_any = False
+
+    try:
+        from megatron.bridge.training import vlm_step as _vlm
+
+        _vlm._create_loss_function = create_nemo_masked_next_token_loss_function
+        _safe_log_rank_0(
+            "[primus.nemo_loss] Patched megatron.bridge.training.vlm_step._create_loss_function "
+            "-> NeMo MaskedTokenLossReduction-equivalent (per-sample token-mean)"
+        )
+        patched_any = True
+    except Exception as e:  # pragma: no cover - defensive
+        _safe_log_rank_0(f"[primus.nemo_loss] Skipping vlm_step patch: {e!r}")
+
+    try:
+        from megatron.bridge.training import gpt_step as _gpt
+
+        _gpt.masked_next_token_loss = nemo_masked_next_token_loss
+        _safe_log_rank_0(
+            "[primus.nemo_loss] Patched megatron.bridge.training.gpt_step.masked_next_token_loss "
+            "-> NeMo per-sample-mean variant"
+        )
+        patched_any = True
+    except Exception as e:  # pragma: no cover - defensive
+        _safe_log_rank_0(f"[primus.nemo_loss] Skipping gpt_step patch: {e!r}")
+
+    _INSTALLED = patched_any
+    return patched_any
+
+
+# Auto-install on import. Importing this module from ``llama2_custom.py``
+# (or any recipe entry point) is sufficient to activate NeMo-equivalent
+# loss reporting + gradient semantics for the entire training run.
+install_nemo_loss_if_enabled()
diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py
new file mode 100644
index 000000000..b992ffa8a
--- /dev/null
+++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py
@@ -0,0 +1,467 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Portions of this file are copied and modified from NeMo MLPerf
+# (mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py).
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+PRE_QUANTIZED_MODEL=True parity with NeMo MLPerf 6.0 MI355X FP4.
+
+Recipe-local pre-quantization wiring. The three byte-equivalent NeMo
+reference functions (``_extract_module``, ``_get_quantized_params_cpu``,
+``_pre_quantize_model``) are ports of:
+
+  ``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py``
+  lines 57-62 (``_extract_module``), 122-149 (``_pre_quantize_model``),
+  151-212 (``_get_quantized_params_cpu``).
+
+Adaptations vs NeMo (NOT behavioural changes):
+
+* NeMo reads ``self.first_last_layers_bf16``,
+  ``self.num_layers_at_start_in_bf16``, ``self.num_layers_at_end_in_bf16``,
+  ``self.store_quantized_params_on_gpu``, ``self.fp8_quantizer``,
+  ``self.mxfp4_quantizer`` from a ``CustomCallback`` instance bound to
+  the Lightning trainer. Primus has no such instance, so these are read
+  from environment variables with NeMo's *default* values
+  (``False`` / ``0`` / ``False``) and the FP8/MXFP4 quantizers are
+  module-locals built inside ``_pre_quantize_model``.
+* NeMo passes ``trainer.model`` (single Lightning module, already
+  unwrapped from MegatronParallel by PL). Primus passes a
+  ``List[ModelChunk]`` from megatron-bridge; ``_extract_module`` strips
+  the list wrapper before stripping ``.module`` chains, so the input
+  semantically matches NeMo's.
+
+Integration: ``install_pre_quantize_wrap(orig_train)`` returns a drop-in
+wrapper around ``megatron.bridge.training.train.train`` that, on the
+first call, stashes FP8 weights on CPU, swaps live weights to MXFP4, and
+bridges the stash to ``mxfp4_healing._ORDERED_FP8_STASH``
+before delegating to ``orig_train``. When ``PRE_QUANTIZED_MODEL`` is not
+enabled the wrap is a no-op (it just returns ``orig_train`` unchanged).
+
+Historical note: this code used to live under
+``primus.backends.megatron.patches.te_patches.pre_quantize_mxfp4_patches``
+as a ``@register_patch(phase="before_train")`` module that fired during
+``MegatronBridgeBaseTrainer.__init__``. Because the patch ran *before*
+the recipe was imported, a separate wrap-aware setter mechanism
+(``_primus_set_orig_train``) was needed to inject the recipe's own
+``megatron_bridge_train_override`` into the already-installed wrapper.
+By moving the install into the recipe itself we can wrap the override
+directly at recipe import time, deleting that plumbing entirely.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Callable, List
+
+from primus.core.utils.module_utils import log_rank_0
+
+_LOG = logging.getLogger(__name__)
+
+_TAG = "[pre_quantize_mxfp4]"
+
+# Layered FP8 CPU stash, populated by ``_pre_quantize_model``.
+# Format mirrors NeMo's ``self.fp8_cpu_params``:
+#   List[List[Float8Tensor]]
+# indexed by ``decoder.layers`` index, then by intra-layer iteration
+# order over ``isinstance(m, (te.Linear, te.LayerNormLinear))`` modules.
+# Same ``Float8Tensor`` objects are also referenced from
+# ``mxfp4_healing._ORDERED_FP8_STASH`` (flat
+# ``(module, fp8_tensor)`` view) for restore-time bookkeeping.
+_FP8_CPU_PARAMS: List[List[Any]] = []
+
+
+def _truthy_env(name: str, default: bool = False) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in ("1", "true", "yes", "on")
+
+
+def is_pre_quantized_enabled() -> bool:
+    """Single switch for the pre-quantize path (``PRE_QUANTIZED_MODEL``)."""
+    return _truthy_env("PRE_QUANTIZED_MODEL", default=False)
+
+
+# ---------------------------------------------------------------------------
+# NeMo MLPerf 6.0 MI355X FP4 byte-equivalent pre-quantization functions.
+#
+# Source of truth: the actual NeMo container code, NOT the older host-tree
+# mirror. Inside the running ``nemo_mxfp4_lora`` Docker image:
+#   /workspace/code/src/callbacks/custom_callbacks.py
+#   - _extract_module                : lines 57-62
+#   - _pre_quantize_model            : lines 122-148
+#   - _get_quantized_params_cpu      : lines 150-213
+# The host-tree file under ``mlperf-training-6-0/`` is 25 lines shorter
+# (660 vs 685) and out of date in three places: the MXFP4Quantizer kwargs,
+# the ``_columnwise_data`` assertion, and the CPU-pin/GPU-restore
+# columnwise-None guard. We mirror the *container* code.
+# ---------------------------------------------------------------------------
+
+
+def _extract_module(model):
+    """Unwrap the model from MegatronParallel / DDP wrappers to get the GPT module.
+
+    NeMo (``custom_callbacks.py:57-62``)::
+
+        m = model
+        while hasattr(m, "module"):
+            m = m.module
+        return m
+
+    Adaptation: Primus passes ``model`` as ``List[ModelChunk]`` from
+    megatron-bridge (one chunk per pipeline stage). Pick chunk 0 then
+    follow ``.module`` exactly as NeMo does on the unwrapped trainer
+    model.
+    """
+    m = model
+    if isinstance(m, (list, tuple)):
+        m = m[0]
+    while hasattr(m, "module"):
+        m = m.module
+    return m
+
+
+def _pre_quantize_model(model):
+    """Pre-quantize model: store FP8 clones on CPU for healing, replace weights with MXFP4.
+
+    Byte-equivalent to NeMo
+    ``custom_callbacks.CustomCallback._pre_quantize_model``
+    (``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py:122-149``).
+    """
+    import torch
+    import transformer_engine_torch as tex
+    from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
+    from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer
+
+    global _FP8_CPU_PARAMS
+
+    device = next(_extract_module(model).parameters()).device
+
+    # Clone current FP8 weights to CPU for healing.
+    # Must happen BEFORE MXFP4 replacement since we reference the live FP8 params.
+    #
+    # ``columnwise=<ENABLE_TRANSPOSE_CACHE>`` mirrors NeMo's
+    # ``custom_callbacks.py::_pre_quantize_model`` (line 205). When the env
+    # is set, ``Float8Quantizer`` allocates a columnwise (transpose) buffer
+    # alongside the rowwise data so the post-healing FP8 GEMMs can reuse a
+    # precomputed transpose. When it is off, no transpose is kept -- and
+    # ``_restore_fp8_weights_to_gpu`` / ``_pre_ttt_fp8_warmup`` additionally
+    # set ``keep_fp8_weight_transpose_cache=False`` on each TE Linear so TE
+    # recomputes the transpose on demand instead of dereferencing a cache
+    # entry that was never populated.
+    #
+    # Default ``False`` matches NeMo's ``_env_enabled`` helper (returns False
+    # for missing / empty env). The Primus MLPerf-parity shell
+    # (``setup_llama2_70b_lora_fp4_training.sh``) exports
+    # ``ENABLE_TRANSPOSE_CACHE=1`` so the reference run gets the fast path.
+    _enable_transpose_cache = _truthy_env("ENABLE_TRANSPOSE_CACHE", default=False)
+    fp8_quantizer = Float8Quantizer(
+        scale=torch.ones(1, dtype=torch.float32, device=device),
+        amax=torch.zeros(1, dtype=torch.float32, device=device),
+        fp8_dtype=tex.DType.kFloat8E4M3,
+        rowwise=True,
+        columnwise=_enable_transpose_cache,
+    )
+    _FP8_CPU_PARAMS = _get_quantized_params_cpu(model, fp8_quantizer, "FP8_DS")
+
+    _use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1"
+
+    mxfp4_quantizer = MXFP4Quantizer(
+        rowwise=True,
+        columnwise=True,
+        with_gemm_swizzled_scales=True,
+        shuffle_rowwise_data=True,
+        shuffle_columnwise_data=True,
+        use_hadamard=_use_hadamard,
+    )
+    _get_quantized_params_cpu(model, mxfp4_quantizer, "MXFP4", replace=True)
+
+    torch.cuda.empty_cache()
+
+
+def _get_quantized_params_cpu(model, quantizer, qtype: str, replace: bool = False) -> List:
+    """Byte-equivalent to NeMo
+    ``custom_callbacks.CustomCallback._get_quantized_params_cpu``
+    (``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py:151-212``).
+
+    Adaptation:
+
+    * NeMo reads ``self.first_last_layers_bf16`` / ``self.num_layers_at_start_in_bf16``
+      / ``self.num_layers_at_end_in_bf16`` / ``self.store_quantized_params_on_gpu``
+      from the OmegaConf ``cfg.model`` (env-defaulted to ``False`` /
+      ``0`` / ``0`` / ``False`` in NeMo's MI355X FP4 config). Primus
+      reads the same env vars (``FIRST_LAST_LAYERS_BF16``,
+      ``NUM_LAYERS_AT_START_IN_BF16``, ``NUM_LAYERS_AT_END_IN_BF16``,
+      ``STORE_GPU``) with NeMo's same defaults. NeMo's MI355X reference
+      run does not set any of them so the effective behaviour is
+      identical.
+    """
+    import torch
+    import transformer_engine.pytorch as te
+    import transformer_engine_torch as tex
+    from transformer_engine.pytorch.quantized_tensor import QuantizedTensor
+    from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+
+    first_last_layers_bf16 = _truthy_env("FIRST_LAST_LAYERS_BF16", default=False)
+    num_layers_at_start_in_bf16 = int(os.getenv("NUM_LAYERS_AT_START_IN_BF16", "0") or "0")
+    num_layers_at_end_in_bf16 = int(os.getenv("NUM_LAYERS_AT_END_IN_BF16", "0") or "0")
+    store_quantized_params_on_gpu = _truthy_env("STORE_GPU", default=False)
+
+    extracted_module = _extract_module(model)
+    layers = extracted_module.decoder.layers
+    layer_count = len(layers)
+
+    quantized_params = []
+    for layer_idx, layer in enumerate(layers):
+        if first_last_layers_bf16:
+            if (
+                layer_idx < num_layers_at_start_in_bf16
+                or layer_idx >= layer_count - num_layers_at_end_in_bf16
+            ):
+                quantized_params.append([])
+                continue
+
+        quantized_layer_params = []
+        for name, module in layer.named_modules():
+            if not isinstance(module, (te.Linear, te.LayerNormLinear)):
+                continue
+            if not hasattr(module, "weight"):
+                continue
+            param = module.weight
+            with torch.no_grad():
+                if qtype == "MXFP4":
+                    if isinstance(param, QuantizedTensor):
+                        bf16_param = param.dequantize().detach()
+                    else:
+                        bf16_param = param.detach().to(torch.bfloat16)
+                    bf16_param = bf16_param.contiguous()
+                    qparam = quantizer(bf16_param)
+                    del bf16_param
+                    assert qparam._rowwise_data is not None, "No rowwise data."
+                elif qtype == "FP8_DS":
+                    fp8_data = param.data.to(torch.float8_e4m3fn).view(torch.uint8)
+                    scale_inv = torch.ones(1, dtype=torch.float32, device=param.device)
+                    qparam = Float8Tensor(
+                        shape=param.shape,
+                        dtype=torch.bfloat16,
+                        data=fp8_data,
+                        fp8_scale_inv=scale_inv,
+                        fp8_dtype=tex.DType.kFloat8E4M3,
+                        quantizer=quantizer,
+                    )
+                    assert qparam._data is not None, "No data."
+                else:
+                    raise ValueError(f"Unsupported quantization type: {qtype}")
+
+                if replace:
+                    qparam.requires_grad = False
+                    module._parameters["weight"] = qparam
+                else:
+                    qparam = qparam.clone()
+                    if not store_quantized_params_on_gpu:
+                        if qtype == "MXFP4":
+                            qparam._rowwise_data = qparam._rowwise_data.cpu().pin_memory()
+                            if qparam._columnwise_data is not None:
+                                qparam._columnwise_data = qparam._columnwise_data.cpu().pin_memory()
+                        elif qtype == "FP8_DS":
+                            qparam._data = qparam._data.cpu().pin_memory()
+                            if qparam._transpose is not None:
+                                qparam._transpose = qparam._transpose.cpu().pin_memory()
+                    quantized_layer_params.append(qparam)
+        quantized_params.append(quantized_layer_params)
+
+    return quantized_params
+
+
+# ---------------------------------------------------------------------------
+# Memory diagnostics (Primus-only; not part of NeMo's pre-quantize logic).
+# ---------------------------------------------------------------------------
+
+
+def _log_gpu_mem(tag: str) -> None:
+    """Print rank-0 GPU mem checkpoint (allocated / reserved / max-allocated, GiB).
+
+    Enabled when ``PRIMUS_LOG_GPU_MEM=1`` or ``MXFP4_HEALING_DEBUG=1``.
+    """
+    if not (
+        _truthy_env("PRIMUS_LOG_GPU_MEM", default=False) or _truthy_env("MXFP4_HEALING_DEBUG", default=False)
+    ):
+        return
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return
+        dev = torch.cuda.current_device()
+        alloc = torch.cuda.memory_allocated(dev) / (1024**3)
+        reserved = torch.cuda.memory_reserved(dev) / (1024**3)
+        max_alloc = torch.cuda.max_memory_allocated(dev) / (1024**3)
+        max_reserved = torch.cuda.max_memory_reserved(dev) / (1024**3)
+        log_rank_0(
+            f"{_TAG}[mem] {tag:<48s} "
+            f"allocated={alloc:7.2f} GiB | reserved={reserved:7.2f} GiB | "
+            f"max_alloc={max_alloc:7.2f} GiB | max_reserved={max_reserved:7.2f} GiB"
+        )
+    except Exception as exc:  # noqa: BLE001
+        log_rank_0(f"{_TAG}[mem] {tag}: failed to read " f"({type(exc).__name__}: {exc})")
+
+
+def _empty_cache_and_collect(tag: str) -> None:
+    """Run gc.collect() + torch.cuda.empty_cache() to release any cached blocks
+    holding onto the freed BF16 storage. Prints a mem checkpoint after."""
+    try:
+        import gc
+
+        import torch
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+        _log_gpu_mem(tag)
+    except Exception as exc:  # noqa: BLE001
+        log_rank_0(f"{_TAG}[mem] {tag}: empty_cache failed " f"({type(exc).__name__}: {exc})")
+
+
+# ---------------------------------------------------------------------------
+# train() wrapper -- invokes _pre_quantize_model once before the first iter.
+# ---------------------------------------------------------------------------
+
+
+def _make_pre_quantizing_train(orig_train: Callable) -> Callable:
+    """Return a thin wrapper around megatron-bridge's ``train`` that runs
+    pre-quantization once on the first call, then delegates to ``orig_train``.
+
+    ``orig_train`` is closed over directly; since this wrapper is installed
+    from the recipe (after the recipe-level ``megatron_bridge_train_override``
+    is defined), there is no need for the setter/getter swap dance the
+    old ``@register_patch`` version had.
+    """
+
+    if getattr(orig_train, "_primus_pre_quantize_wrapped", False):
+        return orig_train  # idempotent
+
+    state = {"done": False}
+
+    def _wrapped_train(forward_step_func, model, *args, **kwargs):
+        if not state["done"]:
+            try:
+                try:
+                    from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import (
+                        log_healing_env_banner_once,
+                    )
+
+                    log_healing_env_banner_once()
+                except Exception:  # noqa: BLE001
+                    pass
+                log_rank_0(
+                    f"{_TAG}   PRE_QUANTIZED_MODEL=True -> running "
+                    "NeMo-equivalent _pre_quantize_model "
+                    "(FP8 stash on CPU + MXFP4 swap)."
+                )
+
+                # Memory diagnostic: confirms whether the BF16 storage is freed
+                # after the in-place MXFP4 swap. If `allocated` does NOT drop by
+                # roughly half between "before swap" and "after swap (post empty_cache)",
+                # something outside this loop (DDP grad bucket, distopt mapping, TE
+                # FP8 weight cache, etc.) is still holding the old BF16 storage.
+                _log_gpu_mem("before _pre_quantize_model")
+
+                # NeMo MLPerf 6.0 MI355X FP4 reference (custom_callbacks.py:122-149).
+                # Byte-equivalent: clones FP8 (E4M3) weights to CPU for healing,
+                # then replaces every TE Linear/LayerNormLinear weight with an
+                # MXFP4Tensor (rowwise+columnwise data populated). HEALING_ITER
+                # gating is enforced by _ORDERED_FP8_STASH consumers; here we
+                # *always* stash so PRE_QUANTIZED_MODEL=True and HEALING_ITER>0
+                # have identical behaviour to NeMo (which always stashes when
+                # cfg.model.pre_quantized_model is True).
+                _pre_quantize_model(model)
+
+                _log_gpu_mem("after _pre_quantize_model (BEFORE empty_cache)")
+                _empty_cache_and_collect("after _pre_quantize_model (AFTER empty_cache)")
+
+                # Bridge: populate mxfp4_healing._ORDERED_FP8_STASH
+                # with flat (module, fp8_cpu_tensor) pairs that share storage
+                # with _FP8_CPU_PARAMS, so existing healing-side code
+                # (restore-to-GPU, refcount audit, FP8 warmup) keeps working.
+                try:
+                    from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import (
+                        _set_ordered_fp8_stash_from_layered,
+                    )
+
+                    _set_ordered_fp8_stash_from_layered(model, _FP8_CPU_PARAMS)
+                except Exception as bridge_err:  # noqa: BLE001
+                    log_rank_0(
+                        f"{_TAG}   FP8 stash bridge to mxfp4_healing failed "
+                        f"({type(bridge_err).__name__}: {bridge_err}); "
+                        f"healing/restore will not work."
+                    )
+
+                # NeMo MLPerf MI355X parity: optional FP8 warmup (off by default).
+                # Honors MXFP4_HEALING_FP8_WARMUP=1 to enable.
+                try:
+                    from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import (
+                        healing_iter as _healing_iter,
+                    )
+                    from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import (
+                        run_fp8_warmup_for_kernel_jit,
+                    )
+
+                    if _healing_iter() > 0:
+                        _optimizer = args[0] if len(args) > 0 else kwargs.get("optimizer")
+                        _state = args[4] if len(args) > 4 else kwargs.get("state")
+                        _model_cfg = None
+                        try:
+                            chunks = list(model) if isinstance(model, (list, tuple)) else [model]
+                            _model_cfg = getattr(chunks[0], "config", None)
+                        except Exception:  # noqa: BLE001
+                            pass
+                        if _model_cfg is None and _state is not None:
+                            _model_cfg = getattr(getattr(_state, "cfg", None), "model", None)
+
+                        if _model_cfg is None:
+                            log_rank_0(
+                                f"{_TAG}   FP8 warmup skipped: could not resolve "
+                                "model_config from train() args (no model.config "
+                                "and no state.cfg.model)."
+                            )
+                        else:
+                            run_fp8_warmup_for_kernel_jit(
+                                model=model,
+                                model_config=_model_cfg,
+                                optimizer=_optimizer,
+                                forward_step_func=forward_step_func,
+                                state=_state,
+                            )
+                except Exception as warm_err:  # noqa: BLE001
+                    log_rank_0(
+                        f"{_TAG}   FP8 warmup raised "
+                        f"({type(warm_err).__name__}: {warm_err}); "
+                        f"continuing to iter 1 without warmup."
+                    )
+            finally:
+                # Even on failure, mark as done so we don't re-attempt every
+                # call in case of caller retry. The exception will propagate.
+                state["done"] = True
+        return orig_train(forward_step_func, model, *args, **kwargs)
+
+    _wrapped_train._primus_pre_quantize_wrapped = True  # type: ignore[attr-defined]
+    return _wrapped_train
+
+
+def install_pre_quantize_wrap(orig_train: Callable) -> Callable:
+    """Recipe-facing entry point.
+
+    If ``PRE_QUANTIZED_MODEL`` is enabled, returns a pre-quantizing
+    wrapper around ``orig_train``; otherwise returns ``orig_train``
+    unchanged. Idempotent: re-wrapping an already-wrapped callable is
+    a no-op.
+    """
+    if not is_pre_quantized_enabled():
+        return orig_train
+    return _make_pre_quantizing_train(orig_train)
diff --git a/primus/backends/nemo_automodel/README.md b/primus/backends/nemo_automodel/README.md
new file mode 100644
index 000000000..dc6082696
--- /dev/null
+++ b/primus/backends/nemo_automodel/README.md
@@ -0,0 +1,92 @@
+# Primus NeMo AutoModel Backend
+
+`nemo_automodel` runs diffusion training through the upstream
+[NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel) diffusion recipe,
+vendored at `third_party/Automodel`. Primus owns config loading and launch;
+AutoModel owns model construction, FSDP2, the dataloader, the optimizer and
+checkpointing.
+
+Everything Primus adds on top lives in this directory as **hooks**. The
+AutoModel submodule and `diffusers` are never patched in place — see
+[No-fork hooks](#no-fork-hooks).
+
+## Layout
+
+```text
+nemo_automodel/
+├── nemo_automodel_adapter.py           config -> AutoModel argument translation
+├── nemo_automodel_pretrain_trainer.py  trainer + the hook registry
+├── argument_builder.py                 params -> AutoModel YAML
+├── quantization/                       SHARED, model-agnostic low precision
+│   ├── primus_turbo_fp8.py             FP8/MXFP4 GEMM
+│   ├── primus_turbo_fp8_attn.py        FP8 attention
+│   └── aiter_bf16_attn.py              non-deterministic bf16 attention
+└── models/                             one subpackage per model
+    ├── flux/
+    │   └── parallelize.py              real AC + FSDP2 strategy
+    └── ideogram4/
+        ├── adapter.py                  flow-matching adapter
+        ├── attention.py                var-len (cu_seqlens) flash attention
+        ├── packing_buffer.py           adapter -> attention packing metadata
+        ├── parallelize.py              real AC + FSDP2 strategy
+        ├── zero1.py                    DDP + ZeRO-1 sharded optimizer
+        ├── profile.py                  torch.profiler train-loop wrapper
+        ├── processor.py                offline VAE + text-encoder cache builder
+        └── data/
+            ├── synthetic.py            fixed synthetic batches (smoke)
+            └── cache.py                the real pre-encoded cache
+```
+
+Two rules keep this navigable:
+
+1. **A model subpackage never imports another's.** Anything two models need goes
+   in `quantization/` (or a future `common/`). This is what lets one model be
+   reviewed, or upstreamed, on its own.
+2. **Model-agnostic code stays out of `models/`.** The low-precision hooks apply
+   to any model the recipe can build, so they live at the top level.
+
+## No-fork hooks
+
+Each hook module exposes a single `install()` that monkey-patches an upstream
+entry point and returns whether it did anything. `install()` must be:
+
+- **env-gated** — a no-op unless its own flag is set, so a default run is
+  unaffected;
+- **idempotent** — safe to call twice (check for your own marker attribute
+  before patching);
+- **additive** — existing behaviour for other models must not change.
+
+`NemoAutomodelPretrainTrainer._install_optional_hooks` imports every hook by
+dotted path and calls `install()` before the recipe builds the transformer (that
+is, before `set_attention_backend` and the first forward). Import failures are
+caught and logged rather than raised, so a missing optional dependency degrades
+to a skipped hook instead of a failed run.
+
+That lazy-by-dotted-path import is also why the `__init__.py` files here do not
+re-export their modules: importing one model's hooks must not drag in another's
+optional dependencies.
+
+> Because hooks patch upstream internals, they are the first thing to break when
+> the `third_party/Automodel` pin moves. A hook whose patch target has been
+> renamed or deleted upstream fails silently — the registry logs the skip and the
+> run continues without it. After a pin bump, force each gate on and confirm
+> `install()` still returns true.
+
+## Adding a model
+
+1. Create `models/<name>/` with an `__init__.py` that documents the modules but
+   re-exports nothing.
+2. Add the hook modules. If a strategy for the model's transformer class is not
+   registered in AutoModel's `PARALLELIZATION_STRATEGIES`, note that
+   `fsdp.activation_checkpointing` is a **silent no-op** until you register one —
+   that is what `parallelize.py` does for FLUX and Ideogram-4.
+3. Register the hooks in the trainer's `hooks` tuple, in a group of their own.
+4. Add a module preset under `primus/configs/modules/nemo_automodel/` and an
+   experiment config under `examples/nemo_automodel/configs/`.
+
+## Naming
+
+Upstream's importable Python package genuinely is `nemo_automodel`
+(`import nemo_automodel.components...`), so those imports and any prose about
+the upstream project keep that name. It is only the Primus-side namespace that
+is ours to choose.
diff --git a/primus/backends/nemo_automodel/__init__.py b/primus/backends/nemo_automodel/__init__.py
new file mode 100644
index 000000000..99759c740
--- /dev/null
+++ b/primus/backends/nemo_automodel/__init__.py
@@ -0,0 +1,17 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from primus.backends.nemo_automodel.nemo_automodel_adapter import NemoAutomodelAdapter
+from primus.backends.nemo_automodel.nemo_automodel_pretrain_trainer import (
+    NemoAutomodelPretrainTrainer,
+)
+from primus.core.backend.backend_registry import BackendRegistry
+
+# Register adapter
+BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter)
+
+# Register trainer
+BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel")
diff --git a/primus/backends/nemo_automodel/argument_builder.py b/primus/backends/nemo_automodel/argument_builder.py
new file mode 100644
index 000000000..ffa7eb59c
--- /dev/null
+++ b/primus/backends/nemo_automodel/argument_builder.py
@@ -0,0 +1,100 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NeMo AutoModel configuration builder.
+
+This is the NeMo AutoModel counterpart of ``MaxTextConfigBuilder`` /
+``TorchTitanJobConfigBuilder``. AutoModel (like MaxText) is an external trainer
+that owns its own configuration system: a recipe is driven by a single YAML that
+``nemo_automodel.components.config._arg_parser.parse_args_and_load_config`` loads
+into a ``ConfigNode``. So there is no JobConfig-style dataclass of defaults to
+merge against — the Primus ``module_config.params`` (already the AutoModel config
+schema: ``model:``/``fsdp:``/``optim:``/``step_scheduler:``/...) *is* the config.
+
+This builder therefore just normalises ``params`` (a ``SimpleNamespace``) and the
+trainer materialises it into an AutoModel ``ConfigNode`` via AutoModel's own
+loader (see ``nemo_automodel_pretrain_trainer.py``), keeping us version-agnostic.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import tempfile
+from types import SimpleNamespace
+from typing import Any, Dict
+
+logger = logging.getLogger(__name__)
+
+# Top-level keys injected by Primus (module_base.yaml / runtime) that are not
+# part of the AutoModel config schema. Stripped before handing the dict to
+# AutoModel. NOTE: this is TOP-LEVEL only — nested AutoModel keys such as
+# ``model.stage`` (the high/low-noise expert window) must be preserved.
+#
+# This is a denylist: a new top-level key added to module_base.yaml would
+# otherwise leak into the AutoModel ConfigNode. ``test_module_base_keys_are_all_stripped``
+# (tests/unit_tests/backends/nemo_automodel) guards against that drift — if it
+# fails, add the new key here (or confirm it belongs in the AutoModel recipe).
+PRIMUS_ONLY_TOP_KEYS = frozenset(
+    {
+        "trainable",
+        "sink_level",
+        "file_sink_level",
+        "stderr_sink_level",
+        "framework",
+        "config",
+        "name",
+        "stage",  # Primus trainer-stage selector; AutoModel uses model.stage
+    }
+)
+
+
+class NemoAutomodelConfigBuilder:
+    """Absorb Primus ``module_config.params`` and return it for downstream use."""
+
+    def __init__(self) -> None:
+        self.config: SimpleNamespace = SimpleNamespace()
+
+    def update(self, params: SimpleNamespace) -> "NemoAutomodelConfigBuilder":
+        """Take params (already merged with model preset + CLI overrides)."""
+        self.config = params
+        return self
+
+    def finalize(self) -> SimpleNamespace:
+        return self.config
+
+
+def namespace_to_dict(obj: Any) -> Any:
+    """Recursively convert ``SimpleNamespace`` (and nested containers) to dict."""
+    if isinstance(obj, SimpleNamespace):
+        return {k: namespace_to_dict(v) for k, v in vars(obj).items()}
+    if isinstance(obj, dict):
+        return {k: namespace_to_dict(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [namespace_to_dict(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(namespace_to_dict(v) for v in obj)
+    return obj
+
+
+def strip_primus_keys(params_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """Drop Primus-only top-level keys so AutoModel sees a clean recipe config."""
+    return {k: v for k, v in params_dict.items() if k not in PRIMUS_ONLY_TOP_KEYS}
+
+
+def export_params_to_yaml(params_dict: Dict[str, Any]) -> str:
+    """
+    Write the (cleaned) config dict to a temporary YAML so AutoModel's own
+    ``parse_args_and_load_config`` can load it — exactly mirroring the MaxText
+    backend. The caller is responsible for deleting the file.
+    """
+    import yaml
+
+    fd, yaml_path = tempfile.mkstemp(suffix=".yaml", prefix="primus_automodel_")
+    with os.fdopen(fd, "w") as f:
+        yaml.dump(params_dict, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
+    return yaml_path
diff --git a/primus/backends/nemo_automodel/data/__init__.py b/primus/backends/nemo_automodel/data/__init__.py
new file mode 100644
index 000000000..5fe72f418
--- /dev/null
+++ b/primus/backends/nemo_automodel/data/__init__.py
@@ -0,0 +1,12 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Model-agnostic data plumbing for the AutoModel diffusion backend.
+
+Currently just the cache-builder registry that ``primus data automodel-cache``
+dispatches through. The per-model builders themselves live with their model, in
+``models/<name>/data/``.
+"""
diff --git a/primus/backends/nemo_automodel/data/registry.py b/primus/backends/nemo_automodel/data/registry.py
new file mode 100644
index 000000000..30f442f48
--- /dev/null
+++ b/primus/backends/nemo_automodel/data/registry.py
@@ -0,0 +1,49 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Registry of offline cache builders for the AutoModel diffusion backend.
+
+AutoModel diffusion models can train from a **pre-encoded flat cache** — the
+VAE and text encoder run once, offline, so training needs neither their weights
+nor their memory. Each model that supports this registers a builder here, and
+``primus data automodel-cache --model <name>`` dispatches through it.
+
+Entries are **dotted strings, not imports**, for two reasons: the CLI can list
+the available models without importing any of them, and a builder that needs an
+optional dependency only fails when it is actually asked for.
+
+This is deliberately separate from ``primus data diffusion-*``, which prepares
+Energon WebDataset shards for the Megatron backend. Different backend, different
+on-disk format; sharing one code path would couple two unrelated pipelines.
+"""
+from __future__ import annotations
+
+import importlib
+from typing import Callable, Dict
+
+# model name -> "<module>:<callable>"
+CACHE_BUILDERS: Dict[str, str] = {
+    "ideogram4": "primus.backends.nemo_automodel.models.ideogram4.data.build:build_cache",
+}
+
+
+def available_models() -> list[str]:
+    """Model names that can build an offline cache."""
+    return sorted(CACHE_BUILDERS)
+
+
+def get_cache_builder(model: str) -> Callable[..., dict]:
+    """Import and return the cache builder registered for ``model``."""
+    try:
+        target = CACHE_BUILDERS[model]
+    except KeyError:
+        raise ValueError(
+            f"no cache builder registered for model {model!r}; available: {', '.join(available_models())}"
+        ) from None
+
+    module_path, _, attribute = target.partition(":")
+    module = importlib.import_module(module_path)
+    return getattr(module, attribute)
diff --git a/primus/backends/nemo_automodel/distributed/__init__.py b/primus/backends/nemo_automodel/distributed/__init__.py
new file mode 100644
index 000000000..06c7356c8
--- /dev/null
+++ b/primus/backends/nemo_automodel/distributed/__init__.py
@@ -0,0 +1,13 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Model-agnostic distributed/sharding hooks for the AutoModel diffusion backend.
+
+Mirrors ``nemo_automodel.components.distributed`` on the Primus side. Code here
+applies to every diffusion model, so it must not import anything under
+``models/``. Re-exports nothing on purpose: hooks are imported lazily by dotted
+path so a missing optional dependency degrades to a skipped hook.
+"""
diff --git a/primus/backends/nemo_automodel/distributed/fsdp2_reshard.py b/primus/backends/nemo_automodel/distributed/fsdp2_reshard.py
new file mode 100644
index 000000000..7fbff5696
--- /dev/null
+++ b/primus/backends/nemo_automodel/distributed/fsdp2_reshard.py
@@ -0,0 +1,106 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Make ``fsdp.reshard_after_forward`` reach PyTorch on the diffusion path.
+
+WHY:
+  ``reshard_after_forward`` is a real ``FSDP2Config`` field and the diffusion recipe
+  parses it correctly from YAML into ``manager_args``. But
+  ``_create_parallel_manager`` rebuilds ``FSDP2Config`` from an **explicit keyword
+  whitelist** that does not include it, so the value is present in ``manager_args``
+  and simply never read back. ``FSDP2Manager.reshard_after_forward`` is therefore
+  always ``None``.
+
+  With ``None``, the per-layer heuristic in ``apply_fsdp2_sharding_recursively``
+  takes over and reshards every transformer block except the last, re-all-gathering
+  each one during backward. That is ZeRO-3 traffic (~3P bytes/step) where ZeRO-2
+  (~2P) was asked for. The ``reshard_after_forward=False`` on the root
+  ``fully_shard()`` in our strategies only covers the root unit -- embeddings, norms,
+  adaLN heads -- not the blocks, which hold nearly all the parameters.
+
+  This affects **every** diffusion model on this path, not just ours, and is worth
+  reporting upstream.
+
+WHAT:
+  Wraps ``_create_parallel_manager`` and re-applies the key the whitelist drops. It
+  **repairs rather than invents**: the value comes from ``manager_args``, which came
+  from YAML. Installed unconditionally with no env gate, because it is
+  behaviour-neutral when the YAML omits the key (``None`` stays ``None`` and today's
+  heuristic applies), so there is nothing to gate.
+
+  Note this is only half the fix. The Primus parallelization strategies must also
+  forward the value into ``apply_fsdp2_sharding_recursively``; see
+  ``models/{flux,ideogram4}/parallelize.py``. **Fixing either half alone changes
+  nothing.**
+
+FORWARD COMPATIBILITY:
+  If Automodel later adds ``reshard_after_forward`` to the whitelist, this wrapper
+  writes the value the manager already holds -- both read the same ``manager_args``
+  -- so it degrades to a no-op rather than conflicting. Do not remove it on a
+  submodule bump on the assumption that it now conflicts.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+logger = logging.getLogger(__name__)
+
+_SENTINEL = "_primus_reshard_patch"
+
+# Provenance for the strategy-side guard. ``patch_installed`` says the wrapper is in
+# place; ``applied_reshard_after_forward`` is the value it last wrote. Read these
+# through a module reference (``import fsdp2_reshard``), never a ``from`` import:
+# a from-import binds the value at import time, when it is still False/None, and
+# never sees install() update it -- so the guard would fire on every healthy run.
+patch_installed = False
+applied_reshard_after_forward = None
+
+
+def install() -> bool:
+    """Patch ``_create_parallel_manager`` to re-apply ``reshard_after_forward``."""
+    global patch_installed
+
+    # Imported inside install(), like the other hooks, so importing this module does
+    # not drag in Automodel.
+    import nemo_automodel._diffusers.auto_diffusion_pipeline as adp
+
+    target = getattr(adp, "_create_parallel_manager", None)
+    if target is None or not callable(target):
+        # An explicit raise, not an assert: asserts are stripped under `python -O`.
+        raise RuntimeError(
+            "auto_diffusion_pipeline._create_parallel_manager is missing; the Automodel "
+            "layout changed and the ZeRO-2 reshard patch cannot apply."
+        )
+    if getattr(target, _SENTINEL, False):
+        patch_installed = True
+        return True  # already installed
+
+    @functools.wraps(target)
+    def _patched(*args, **kwargs):
+        global applied_reshard_after_forward
+
+        # Locate the dict defensively: a rename upstream should degrade to today's
+        # behaviour, not raise a KeyError from inside the wrapper at parallelize
+        # time, long after install.
+        manager_args = args[0] if args else kwargs.get("manager_args")
+        manager = target(*args, **kwargs)
+
+        if isinstance(manager_args, dict) and manager_args.get("_manager_type", "fsdp2") == "fsdp2":
+            # The callee copies manager_args and pops _manager_type from the copy, so
+            # the caller's dict we inspect here still carries both keys.
+            value = manager_args.get("reshard_after_forward")
+            manager.reshard_after_forward = value
+            applied_reshard_after_forward = value
+            logger.info("[PrimusFSDP2Reshard] applied reshard_after_forward=%s", value)
+
+        return manager
+
+    # The sentinel is what makes this idempotent. Without it a second install() would
+    # capture the already-patched function as `target` and nest wrappers.
+    setattr(_patched, _SENTINEL, True)
+    adp._create_parallel_manager = _patched
+    patch_installed = True
+    return True
diff --git a/primus/backends/nemo_automodel/models/__init__.py b/primus/backends/nemo_automodel/models/__init__.py
new file mode 100644
index 000000000..98907b01a
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/__init__.py
@@ -0,0 +1,15 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Per-model hooks for the AutoModel diffusion backend.
+
+One subpackage per model. A model subpackage must not import another's: anything
+two models share belongs in ``quantization/`` (or a future ``common/``). That
+isolation is what lets a single model be reviewed, or upstreamed, on its own.
+
+Deliberately no re-exports, so that importing one model never drags in another's
+optional dependencies. See ``quantization/__init__.py``.
+"""
diff --git a/primus/backends/nemo_automodel/models/flux/__init__.py b/primus/backends/nemo_automodel/models/flux/__init__.py
new file mode 100644
index 000000000..a8d60d7e5
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/flux/__init__.py
@@ -0,0 +1,12 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""FLUX-specific hooks for the AutoModel diffusion backend.
+
+``parallelize.py`` registers a FLUX parallelization strategy (real activation
+checkpointing + FSDP2 sharding), replacing the default strategy whose AC path is
+a silent no-op for FLUX blocks.
+"""
diff --git a/primus/backends/nemo_automodel/models/flux/parallelize.py b/primus/backends/nemo_automodel/models/flux/parallelize.py
new file mode 100644
index 000000000..516b554f1
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/flux/parallelize.py
@@ -0,0 +1,175 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Real activation checkpointing for FLUX in the NeMo AutoModel diffusion recipe.
+
+WHY:
+  AutoModel selects an FSDP/AC parallelization strategy by the model's class name
+  (``PARALLELIZATION_STRATEGIES`` in
+  ``nemo_automodel.components.distributed.parallelizer``). ``FluxTransformer2DModel``
+  is NOT registered, so it falls back to the default strategy. The default
+  strategy applies activation checkpointing by wrapping each layer's ``self_attn``
+  / ``mlp`` submodules -- attributes a FLUX block does not have -- so
+  ``fsdp.activation_checkpointing: true`` is a SILENT NO-OP for FLUX: the flag is
+  accepted, nothing is checkpointed, and the activation-memory ceiling (hence the
+  achievable batch size) is unchanged.
+
+WHAT this does (NO diffusers / Automodel fork):
+  Registers a FLUX-specific parallelization strategy for the ``FluxTransformer2DModel``
+  class via the submodule's own ``register_parallel_strategy`` entry point. The
+  strategy wraps each dual-stream ``transformer_blocks`` block and each
+  single-stream ``single_transformer_blocks`` block in a NON-REENTRANT
+  ``checkpoint_wrapper`` (recompute on backward) BEFORE FSDP2 sharding, then shards
+  exactly like the in-tree Wan/Hunyuan diffusion strategies. Checkpointing only
+  happens when the recipe passes ``activation_checkpointing=True`` (i.e. the config
+  flag), so an AC-off run is unaffected.
+
+  Env-gated by ``PRIMUS_FLUX_REAL_AC=1`` (default off). Off = current behavior
+  (FLUX AC remains a no-op), which keeps the flag an explicit, reversible A/B lever.
+  When on, ``fsdp.activation_checkpointing`` starts doing what it says for FLUX.
+
+  TP note: this strategy targets the dp/FSDP path (tp_size=1). If a TP mesh is
+  requested it warns once and proceeds without a FLUX-specific TP plan.
+
+Activation (env, no config schema change):
+    PRIMUS_FLUX_REAL_AC=1    register the real-AC FLUX parallelization strategy
+                             (default off = no-op)
+"""
+from __future__ import annotations
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+_FLUX_MODEL_NAME = "FluxTransformer2DModel"
+_FLUX_BLOCK_ATTRS = ("transformer_blocks", "single_transformer_blocks")
+
+
+def is_flux_real_ac_enabled() -> bool:
+    """Whether the real-AC FLUX strategy should be registered."""
+    return os.getenv("PRIMUS_FLUX_REAL_AC", "0") in _TRUTHY
+
+
+def install() -> bool:
+    """Register the real-AC FLUX parallelization strategy.
+
+    No-op (returns False) unless ``PRIMUS_FLUX_REAL_AC`` is set. Registers a
+    strategy for ``FluxTransformer2DModel`` in the AutoModel strategy registry;
+    idempotent (returns True if already present). Modifies NO Automodel source.
+    """
+    if not is_flux_real_ac_enabled():
+        return False
+
+    import torch
+
+    import nemo_automodel.components.distributed.parallelizer as P
+
+    if _FLUX_MODEL_NAME in P.PARALLELIZATION_STRATEGIES:
+        return True  # already registered (idempotent)
+
+    class FluxParallelizationStrategy(P.ParallelizationStrategy):
+        """FSDP2 + real activation checkpointing for FLUX MMDiT transformers."""
+
+        def parallelize(
+            self,
+            model,
+            device_mesh,
+            mp_policy=None,
+            offload_policy=None,
+            sequence_parallel: bool = False,
+            activation_checkpointing: bool = False,
+            tp_shard_plan=None,
+            dp_replicate_mesh_name: str = "dp_replicate",
+            dp_shard_cp_mesh_name: str = "dp_shard_cp",
+            tp_mesh_name: str = "tp",
+            reshard_after_forward: bool | None = None,
+            **kwargs,
+        ):
+            dp_mesh = P.get_fsdp_dp_mesh(device_mesh, dp_replicate_mesh_name, dp_shard_cp_mesh_name)
+
+            # Read the hook's state through the module, never a from-import: a
+            # from-import binds False at import time and the guard then fires on
+            # every healthy run.
+            from primus.backends.nemo_automodel.distributed import fsdp2_reshard
+
+            logger.info(
+                "[PrimusFluxAC] reshard_after_forward: received=%s applied_by_hook=%s",
+                reshard_after_forward,
+                fsdp2_reshard.applied_reshard_after_forward,
+            )
+            if not fsdp2_reshard.patch_installed:
+                # The trainer swallows hook exceptions, so a failed install shows up
+                # only as one log line and then a full run at ~3P traffic. This is the
+                # one unambiguous in-process signal that it happened.
+                logger.error(
+                    "[PrimusFluxAC] the FSDP2 reshard repair hook is NOT installed, so "
+                    "fsdp.reshard_after_forward from YAML was dropped upstream and this run "
+                    "will use ZeRO-3-style per-block resharding regardless of config."
+                )
+
+            if tp_mesh_name in getattr(device_mesh, "mesh_dim_names", ()) and device_mesh[tp_mesh_name].size() > 1:
+                logger.warning(
+                    "[PrimusFluxAC] tensor parallelism requested but the FLUX real-AC "
+                    "strategy has no FLUX-specific TP plan; proceeding with FSDP only."
+                )
+
+            # Real AC: wrap each dual-stream and single-stream block so its forward
+            # activations are recomputed on backward. Must happen BEFORE FSDP2
+            # sharding so the module structure is stable when fully_shard indexes
+            # params. NO_REENTRANT is required for torch.compile compatibility.
+            if activation_checkpointing:
+                wrapped = 0
+                for attr in _FLUX_BLOCK_ATTRS:
+                    blocks = getattr(model, attr, None)
+                    if blocks is None:
+                        continue
+                    for idx in range(len(blocks)):
+                        blocks[idx] = P.checkpoint_wrapper(
+                            blocks[idx],
+                            checkpoint_impl=P.CheckpointImpl.NO_REENTRANT,
+                        )
+                        wrapped += 1
+                logger.info("[PrimusFluxAC] wrapped %d FLUX blocks with activation checkpointing", wrapped)
+
+            if not mp_policy:
+                mp_policy = P.MixedPrecisionPolicy(
+                    param_dtype=torch.bfloat16,
+                    reduce_dtype=torch.float32,
+                    output_dtype=torch.float32,
+                )
+
+            # Keyword form throughout: this call previously passed seven positional
+            # arguments and stopped one short of reshard_after_forward, silently
+            # discarding it. Keywords make that class of bug impossible to repeat.
+            P.apply_fsdp2_sharding_recursively(
+                model,
+                dp_mesh,
+                mp_policy,
+                offload_policy,
+                enable_fsdp2_prefetch=kwargs.get("enable_fsdp2_prefetch", True),
+                fsdp2_backward_prefetch_depth=kwargs.get("fsdp2_backward_prefetch_depth", 2),
+                fsdp2_forward_prefetch_depth=kwargs.get("fsdp2_forward_prefetch_depth", 1),
+                # Forwarded faithfully, including None: YAML decides, not this code.
+                reshard_after_forward=reshard_after_forward,
+            )
+
+            return P.fully_shard(
+                model,
+                mesh=dp_mesh,
+                mp_policy=mp_policy,
+                offload_policy=offload_policy,
+                reshard_after_forward=False,
+            )
+
+    P.register_parallel_strategy(name=_FLUX_MODEL_NAME)(FluxParallelizationStrategy)
+    logger.info(
+        "[PrimusFluxAC] Registered real activation-checkpointing parallelization "
+        "strategy for %s (honors fsdp.activation_checkpointing).",
+        _FLUX_MODEL_NAME,
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/__init__.py b/primus/backends/nemo_automodel/models/ideogram4/__init__.py
new file mode 100644
index 000000000..3c9a2f593
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/__init__.py
@@ -0,0 +1,27 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Ideogram-4 hooks for the AutoModel diffusion backend.
+
+Single-stream rectified-flow T2I. All of it is Primus-side: the Automodel
+submodule and diffusers stay pristine.
+
+    adapter.py         flow-matching adapter: packs [left-pad][text][image],
+                       maps sigma -> t, negates the DiT velocity
+    attention.py       var-len (cu_seqlens) flash attention processor, the exact
+                       replacement for the dense-mask SDPA path
+    packing_buffer.py  per-step transport carrying packing metadata from the
+                       adapter to the attention processor
+    parallelize.py     parallelization strategy: real AC + FSDP2 sharding
+    zero1.py           DDP + ZeRO-1 sharded-optimizer path
+    profile.py         torch.profiler wrapper around the recipe train loop
+    processor.py       offline encoder (VAE + text encoder) producing the cache
+    data/              dataloaders: synthetic.py (smoke) and cache.py (real)
+
+``attention.py`` and ``packing_buffer.py`` are mutually dependent by design (the
+attention processor reads what the adapter published); everything else in here
+depends only on the Automodel/diffusers API.
+"""
diff --git a/primus/backends/nemo_automodel/models/ideogram4/adapter.py b/primus/backends/nemo_automodel/models/ideogram4/adapter.py
new file mode 100644
index 000000000..5ee279153
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/adapter.py
@@ -0,0 +1,408 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Ideogram-4 flow-matching adapter + no-fork registration for the AutoModel
+diffusion recipe.
+
+WHY (no Automodel/diffusers fork):
+  AutoModel selects a flow-matching model adapter by name via
+  ``nemo_automodel.components.flow_matching.pipeline.create_adapter`` (a closed
+  dict of {"flux","flux2","wan","hunyuan","qwen_image","simple"}). Ideogram-4 is
+  not in that dict. Rather than edit the submodule, ``install()`` wraps
+  ``create_adapter`` at runtime so ``flow_matching.adapter_type: ideogram4``
+  resolves to :class:`Ideogram4Adapter`. The recipe imports ``create_adapter`` by
+  name (``from ...pipeline import create_adapter``), so we patch it in BOTH the
+  pipeline module and the already-imported recipe module namespace.
+
+WHAT the adapter does:
+  Ideogram-4 is a single-stream flow-matching DiT: text-conditioning tokens and
+  patchified image latents live in ONE packed sequence, distinguished by a
+  per-token ``indicator`` and joined by a block-diagonal mask (``segment_ids``),
+  with 3-axis MRoPE (``position_ids``). This adapter maps AutoModel's latent-space
+  flow-matching convention (noise ``x_t=(1-sigma)x0+sigma eps``; target
+  ``v=eps-x0`` in the 128-dim packed-latent space) onto that packed contract:
+
+  - ``prepare_inputs`` packs image latents ``[B,128,H_p,W_p] -> [B,n_img,128]``,
+    prepends a zeroed text region for ``hidden_states [B,S,128]``, builds
+    ``encoder_hidden_states [B,S,53248]`` (Qwen3-VL features over text, zeros over
+    image), and the ``position_ids/segment_ids/indicator`` for the
+    ``[pad][text][image]`` layout. Ideogram model time is ``t = 1 - sigma``.
+    It ALSO builds the var-len ``cu_seqlens`` packing here on the host (see
+    :func:`build_cu_seqlens`), so the processor never derives it from the mask inside the
+    compiled region -- that derivation is data-dependent, host-syncing, and under FSDP2
+    its graph break desyncs the per-layer collectives. Disable with
+    ``PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS=0``.
+  - ``forward`` publishes that packing into the attention modules' shared buffer
+    (``ideogram4_packing_buffer``, which documents why that route and not diffusers'
+    ``attention_kwargs``), runs the DiT, slices the image-token velocity, unpacks to
+    ``[B,128,H_p,W_p]`` and NEGATES it: the DiT predicts ``x0 - eps`` (inference
+    feeds ``-v`` to the scheduler) while AutoModel's target is ``eps - x0``.
+
+  Batch keys consumed (from the Ideogram-4 preprocessor / cache):
+    - ``image_latents``: ``[B,128,H_p,W_p]`` patchified + BN-normalized latents.
+    - ``llm_features``: ``[B,max_text,53248]`` left-padded Qwen3-VL 13-layer feats.
+    - ``text_lengths``: ``[B]`` int, real (non-pad) text token count per sample.
+
+  ``install()`` is additive and safe: it only adds the ``ideogram4`` route, so a
+  FLUX/Wan run is unaffected. Idempotent.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+
+from primus.backends.nemo_automodel.models.ideogram4.packing_buffer import publish_packing
+from primus.backends.nemo_automodel.models.ideogram4.attention import (
+    assume_dense_enabled,
+    precompute_cu_seqlens_active,
+)
+
+logger = logging.getLogger(__name__)
+
+# Per-token role / layout constants live with the diffusers transformer definition.
+try:
+    from diffusers.models.transformers.transformer_ideogram4 import (
+        IMAGE_POSITION_OFFSET,
+        LLM_TOKEN_INDICATOR,
+        OUTPUT_IMAGE_INDICATOR,
+        SEQUENCE_PADDING_INDICATOR,
+    )
+except Exception:  # pragma: no cover - keep the module importable without diffusers (tests, lint)
+    # Values as of diffusers 0.39.0, read out of the installed package rather than guessed.
+    # They are the layout contract with the model, so a stale literal here silently mislabels
+    # every token: the warning below is the only signal that the real ones were not used.
+    IMAGE_POSITION_OFFSET = 65536
+    LLM_TOKEN_INDICATOR = 3
+    OUTPUT_IMAGE_INDICATOR = 2
+    SEQUENCE_PADDING_INDICATOR = -1
+    logger.warning(
+        "[PrimusIdeogram4] Could not import the Ideogram-4 layout constants from diffusers; "
+        "falling back to the diffusers 0.39.0 values (pad=%d, llm=%d, image=%d, offset=%d). "
+        "Fine for import-only use, but if a real run reaches this the constants may have been "
+        "renamed upstream and the token layout will be wrong with no further error.",
+        SEQUENCE_PADDING_INDICATOR,
+        LLM_TOKEN_INDICATOR,
+        OUTPUT_IMAGE_INDICATOR,
+        IMAGE_POSITION_OFFSET,
+    )
+
+
+def build_cu_seqlens(
+    text_lengths: List[int],
+    max_text_tokens: int,
+    num_image_tokens: int,
+    device: torch.device = None,
+) -> torch.Tensor:
+    """Var-len packing for the ``[left-pad][text][image]`` layout, built on the HOST.
+
+    Each row ``b`` of the flattened ``(B*S)`` sequence (``S = max_text_tokens +
+    num_image_tokens``) contributes exactly TWO segments, matching what the model's
+    ``(seg_i == seg_j)`` mask encodes:
+
+      * ``[b*S, b*S+offset)``      the left-pad block (attends only to itself; discarded)
+      * ``[b*S+offset, (b+1)*S)``  the text+image block (Ideogram attends these jointly)
+
+    with ``offset = max_text_tokens - text_lengths[b]``.
+
+    Because the count is always ``2*B``, the returned tensor's length is always ``2*B+1``
+    -- independent of the captions in the batch. That is what keeps the compiled graph
+    reusable: a segment count that varied with the data would change this tensor's shape
+    and force a recompile per distinct length pattern.
+
+    Args:
+        text_lengths: per-sample real (non-pad) text token counts, as Python ints.
+        max_text_tokens: padded width of the text region, INCLUDING the reserved pad
+            column, so that ``offset >= 1`` on every row.
+        num_image_tokens: ``grid_h * grid_w``.
+        device: destination device. Building here costs one host->device copy; the thing
+            being avoided is device->host reads inside the compiled region.
+
+    Returns:
+        ``int32`` ``(2*len(text_lengths)+1,)`` cumulative segment starts, beginning at 0
+        and ending at ``B*S`` -- the ``flash_attn_varlen_func`` contract.
+
+    Raises:
+        ValueError: if any row would produce an empty pad segment (``offset == 0``).
+            Zero-length segments are not representable for var-len flash, and the
+            reserved pad column exists precisely to make this unreachable.
+    """
+    seq_len = max_text_tokens + num_image_tokens
+    starts: List[int] = []
+    for b, num_text in enumerate(text_lengths):
+        offset = max_text_tokens - int(num_text)
+        if offset < 1:
+            raise ValueError(
+                f"row {b} has text_length={int(num_text)} with max_text_tokens={max_text_tokens}, "
+                "leaving no padding. Every row needs >=1 pad token so the segment count (and "
+                "therefore cu_seqlens' shape) is the same for every batch."
+            )
+        base = b * seq_len
+        starts.append(base)
+        starts.append(base + offset)
+    starts.append(len(text_lengths) * seq_len)
+    return torch.tensor(starts, dtype=torch.int32, device=device)
+
+
+def _base_adapter_cls():
+    """Import AutoModel's ModelAdapter base lazily (submodule must be importable)."""
+    from nemo_automodel.components.flow_matching.adapters.base import ModelAdapter
+
+    return ModelAdapter
+
+
+# Defined as a factory so importing this module never requires nemo_automodel until
+# install()/first use — keeps the module import-safe in any context.
+def _build_ideogram4_adapter_class():
+    ModelAdapter = _base_adapter_cls()
+
+    class Ideogram4Adapter(ModelAdapter):
+        """Model adapter for Ideogram-4 single-stream flow-matching T2I."""
+
+        def __init__(self, in_channels: int = 128, predict_negative_velocity: bool = True):
+            self.in_channels = in_channels
+            self.predict_negative_velocity = predict_negative_velocity
+
+        @staticmethod
+        def _prepare_ids(
+            text_lengths: List[int],
+            grid_h: int,
+            grid_w: int,
+            max_text_tokens: int,
+            device: torch.device,
+        ):
+            """Build ``[left-pad][text][image]`` position/segment/indicator tensors.
+
+            Mirrors ``Ideogram4Pipeline._prepare_ids`` so training == inference layout.
+            """
+            batch_size = len(text_lengths)
+            num_image_tokens = grid_h * grid_w
+            total_seq_len = max_text_tokens + num_image_tokens
+
+            h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
+            w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
+            t_idx = torch.zeros_like(h_idx)
+            image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
+
+            position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
+            segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
+            indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
+
+            for b, num_text in enumerate(text_lengths):
+                num_text = int(num_text)
+                offset = max_text_tokens - num_text
+                text_pos = torch.arange(num_text)
+                text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1)
+                position_ids[b, offset : offset + num_text] = text_pos_3d
+                position_ids[b, offset + num_text :] = image_pos
+                indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR
+                indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR
+                segment_ids[b, offset : offset + num_text + num_image_tokens] = 1
+
+            return position_ids.to(device), segment_ids.to(device), indicator.to(device)
+
+        def _pack_image_latents(self, latents: torch.Tensor) -> torch.Tensor:
+            b, c, h, w = latents.shape
+            return latents.reshape(b, c, h * w).permute(0, 2, 1).contiguous()
+
+        def _unpack_image_latents(self, tokens: torch.Tensor, h: int, w: int) -> torch.Tensor:
+            b, _, c = tokens.shape
+            return tokens.permute(0, 2, 1).contiguous().reshape(b, c, h, w)
+
+        def prepare_inputs(self, context) -> Dict[str, Any]:
+            batch = context.batch
+            device = context.device
+            dtype = context.dtype
+
+            noisy = context.noisy_latents
+            if noisy.ndim != 4:
+                raise ValueError(
+                    f"Ideogram4Adapter expects 4D patchified latents [B, C, H_p, W_p], got {noisy.ndim}D"
+                )
+            B, C, H_p, W_p = noisy.shape
+            if C != self.in_channels:
+                raise ValueError(f"Expected {self.in_channels} packed channels, got {C}")
+            num_image_tokens = H_p * W_p
+
+            img_tokens = self._pack_image_latents(noisy)  # [B, n_img, 128]
+
+            llm_features = batch["llm_features"].to(device, dtype=dtype, non_blocking=True)
+            if llm_features.ndim == 2:
+                llm_features = llm_features.unsqueeze(0)
+            # Text capacity as produced by the dataloader, before the reserved pad column.
+            text_capacity = llm_features.shape[1]
+
+            # Reserve one always-pad column so offset >= 1 on every row. Without it a
+            # caption filling the full width yields a single-segment row, so the segment
+            # count -- and thus cu_seqlens' shape -- would vary per batch and recompile the
+            # graph. Costs one token position out of S; capping captions instead would mean
+            # dropping a real token to satisfy the compiler.
+            precompute_cu_seqlens = precompute_cu_seqlens_active()
+            if precompute_cu_seqlens:
+                pad_col = torch.zeros(B, 1, llm_features.shape[-1], device=device, dtype=dtype)
+                llm_features = torch.cat([pad_col, llm_features], dim=1)
+            max_text = llm_features.shape[1]
+
+            if context.cfg_dropout_prob > 0.0:
+                drop = torch.rand(B, 1, 1, device=device) < context.cfg_dropout_prob
+                llm_features = llm_features.masked_fill(drop, 0.0)
+
+            text_lengths = batch.get("text_lengths")
+            if text_lengths is None:
+                text_lengths = [text_capacity] * B
+            elif torch.is_tensor(text_lengths):
+                text_lengths = text_lengths.tolist()
+            text_lengths = [int(t) for t in text_lengths]
+            if max(text_lengths) > text_capacity:
+                raise ValueError(
+                    f"text_lengths max {max(text_lengths)} exceeds the llm_features text width "
+                    f"{text_capacity}; the dataloader must left-pad to at least the longest caption."
+                )
+
+            # ASSUME_DENSE tells the attention processor to skip the block-diagonal
+            # analysis and run dense flash over the whole row. That is exact only when no
+            # row has padding; on a ragged batch it lets pad tokens attend and leaks
+            # attention across segments, corrupting training with no error. The lengths are
+            # already on the host here, so refusing costs nothing.
+            distinct = sorted(set(text_lengths))
+            if assume_dense_enabled() and len(distinct) > 1:
+                raise ValueError(
+                    "PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE=1 requires every sample in the batch to have "
+                    f"the same text length, but this batch has lengths {distinct}. Dense flash would "
+                    "let padding tokens attend and silently corrupt training. Unset the flag to use "
+                    "the exact var-len path, or pin min_text_tokens == max_text_tokens."
+                )
+
+            text_z = torch.zeros(B, max_text, C, device=device, dtype=dtype)
+            hidden_states = torch.cat([text_z, img_tokens], dim=1)  # [B, S, 128]
+
+            img_feat_pad = torch.zeros(B, num_image_tokens, llm_features.shape[-1], device=device, dtype=dtype)
+            encoder_hidden_states = torch.cat([llm_features, img_feat_pad], dim=1)  # [B, S, 53248]
+
+            position_ids, segment_ids, indicator = self._prepare_ids(text_lengths, H_p, W_p, max_text, device)
+
+            # Built once here and shared by all layers, so the packing is a graph INPUT
+            # instead of something the processor recovers from the mask 34 times per step.
+            # Left on the HOST: ``forward`` publishes it into the attention modules' shared
+            # buffer, and that single copy_ is then the only host->device transfer.
+            # max_seqlen is the static upper bound S, not the true per-batch maximum:
+            # Dynamo guards Python ints by value, so a data-derived one would recompile
+            # almost every batch.
+            cu_seqlens = None
+            max_seqlen = None
+            if precompute_cu_seqlens:
+                cu_seqlens = build_cu_seqlens(text_lengths, max_text, num_image_tokens)
+                max_seqlen = max_text + num_image_tokens
+
+            # Ideogram model time: 0=noise, 1=data => t = 1 - sigma.
+            timestep = (1.0 - context.sigma).to(dtype)
+
+            return {
+                "hidden_states": hidden_states,
+                "timestep": timestep,
+                "encoder_hidden_states": encoder_hidden_states,
+                "position_ids": position_ids,
+                "segment_ids": segment_ids,
+                "indicator": indicator,
+                "_max_text": max_text,
+                "_h_p": H_p,
+                "_w_p": W_p,
+                "_cu_seqlens": cu_seqlens,
+                "_max_seqlen": max_seqlen,
+            }
+
+        def forward(self, model: nn.Module, inputs: Dict[str, Any]) -> torch.Tensor:
+            max_text = inputs.pop("_max_text")
+            h_p = inputs.pop("_h_p")
+            w_p = inputs.pop("_w_p")
+            cu_seqlens = inputs.pop("_cu_seqlens", None)
+            max_seqlen = inputs.pop("_max_seqlen", None)
+
+            # Publish the packing onto the attention modules, which is where the processor
+            # reads it (``ideogram4_packing_buffer``). Diffusers' ``attention_kwargs`` cannot
+            # carry it: the model never forwards it to the blocks in 0.39.0, so passing it was
+            # a silent no-op. Two ordering rules, both silent when broken: this must happen
+            # BEFORE the model call, and nothing may republish between the forward and its
+            # backward -- per-layer compile lives inside the checkpoint wrapper, so the block
+            # re-reads this buffer during the backward recompute.
+            if cu_seqlens is not None:
+                publish_packing(
+                    model,
+                    cu_seqlens,
+                    max_seqlen,
+                    device=inputs["hidden_states"].device,
+                    # Having built a packing, a model that cannot read it is a misconfiguration,
+                    # not a fallback: every layer would derive its own from the mask, and on a
+                    # subset of ranks that quietly averages two attention paths into one
+                    # gradient. Non-zero ranks log nothing after init, so raising is the only
+                    # way this becomes visible.
+                    required=True,
+                )
+
+            out = model(
+                hidden_states=inputs["hidden_states"],
+                timestep=inputs["timestep"],
+                encoder_hidden_states=inputs["encoder_hidden_states"],
+                position_ids=inputs["position_ids"],
+                segment_ids=inputs["segment_ids"],
+                indicator=inputs["indicator"],
+                return_dict=False,
+            )
+            pred = self.post_process_prediction(out)  # [B, S, 128]
+
+            img_pred = pred[:, max_text:]  # [B, n_img, 128]
+            unpacked = self._unpack_image_latents(img_pred, h_p, w_p)  # [B, 128, H_p, W_p]
+
+            # DiT predicts x0 - eps; AutoModel target is eps - x0.
+            return -unpacked if self.predict_negative_velocity else unpacked
+
+    return Ideogram4Adapter
+
+
+# Cache the built class so identity is stable across calls.
+_IDEOGRAM4_ADAPTER_CLS = None
+
+
+def get_ideogram4_adapter_class():
+    global _IDEOGRAM4_ADAPTER_CLS
+    if _IDEOGRAM4_ADAPTER_CLS is None:
+        _IDEOGRAM4_ADAPTER_CLS = _build_ideogram4_adapter_class()
+    return _IDEOGRAM4_ADAPTER_CLS
+
+
+def install() -> bool:
+    """Route ``adapter_type == "ideogram4"`` to the Ideogram-4 adapter via a no-fork
+    wrapper around AutoModel's ``create_adapter``.
+
+    Additive and idempotent; never changes existing adapter behavior. Returns True.
+    """
+    import nemo_automodel.components.flow_matching.pipeline as P
+
+    orig = P.create_adapter
+    if getattr(orig, "_ideogram4_patched", False):
+        return True
+
+    def create_adapter_patched(adapter_type: str, **kwargs):
+        if adapter_type == "ideogram4":
+            return get_ideogram4_adapter_class()(**kwargs)
+        return orig(adapter_type, **kwargs)
+
+    create_adapter_patched._ideogram4_patched = True
+    P.create_adapter = create_adapter_patched
+
+    # The recipe does ``from ...pipeline import create_adapter`` (bound by name), so
+    # patch its module namespace too if it is already imported.
+    try:
+        import nemo_automodel.recipes.diffusion.train as T
+
+        if getattr(T, "create_adapter", None) is orig:
+            T.create_adapter = create_adapter_patched
+    except Exception as exc:  # pragma: no cover
+        logger.debug("[PrimusIdeogram] recipe namespace patch skipped: %s", exc)
+
+    logger.info("[PrimusIdeogram] Registered 'ideogram4' flow-matching adapter (no-fork).")
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/attention.py b/primus/backends/nemo_automodel/models/ideogram4/attention.py
new file mode 100644
index 000000000..aadc51b9b
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/attention.py
@@ -0,0 +1,471 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Var-len (``cu_seqlens``) flash-attention processor for Ideogram-4 in the NeMo
+AutoModel diffusion recipe — the fast, exact replacement for masked torch SDPA.
+
+WHY:
+  Ideogram-4's attention runs on **masked torch SDPA** today: the transformer builds
+  a dense ``(B,1,L,L)`` block-diagonal boolean mask from ``segment_ids`` and hands it
+  to ``dispatch_attention_fn``. A dense mask forces the SDPA/NATIVE backend (flash /
+  aiter backends reject one), so the whole model forgoes flash — and at 1k–2k px
+  resolutions attention is O(L²) and dominates the step. The mask is, in practice,
+  a per-row block-diagonal over CONTIGUOUS segments (left-padding = {pad}+{text+image};
+  future multi-sample packing = several sample blocks). That maps exactly onto
+  **variable-length flash attention** (``cu_seqlens``): pack the batch into one flat
+  sequence and mark segment boundaries — no dense mask, full flash speed, EXACT numerics.
+
+WHAT (NO diffusers / Automodel fork):
+  A drop-in diffusers attention processor for ``Ideogram4Attention``. It reproduces the
+  stock ``Ideogram4AttnProcessor`` math verbatim (q/k/v proj → q/k RMSNorm → MRoPE) and
+  only replaces the attention CALL: it converts the block-diagonal boolean mask to
+  ``cu_seqlens`` and runs ``aiter.flash_attn_varlen_func`` (``deterministic=False`` — the
+  non-deterministic backward is a large, numerically-equivalent throughput win, and the
+  *deterministic* hd=256 backward has a workspace that OOMs at 2048²). When the mask is
+  absent or trivial (one full segment per row, e.g. a fixed-length/no-pad batch) it takes
+  the plain dense ``flash_attn_func`` fast path.
+
+  Install swaps the class default processor (``Ideogram4Attention._default_processor_cls``)
+  BEFORE the recipe builds the model, so every attention module is constructed with it.
+  Env-gated by ``PRIMUS_IDEOGRAM_VARLEN_ATTN=1`` (default off = stock SDPA path). No
+  Automodel/diffusers source is modified.
+
+GENERALITY / REUSE:
+  The mask→``cu_seqlens`` transform (:func:`blockdiag_bool_mask_to_cu_seqlens`) and the
+  var-len call (:func:`varlen_flash_attention`) are **model-agnostic** — they assume only
+  a per-row block-diagonal boolean mask with contiguous segments, so any packed/padded
+  attention (not just Ideogram) can reuse them. The only Ideogram-specific piece is the
+  thin processor that wires the model's proj/norm/RoPE to them.
+
+CORRECTNESS:
+  Exact, not approximate. Every token (including padding) keeps its own segment, so the
+  result matches masked SDPA on all positions up to bf16 + non-deterministic-atomic
+  ordering (image-token velocity — the only positions the loss reads — matches within the
+  bf16 floor). The processor falls back to the ORIGINAL dense dispatch ONLY for mask
+  *types* it cannot represent as contiguous segments (a non-boolean/additive mask), and
+  warns once — it never silently relays the unoptimized path for the case it exists to
+  serve (Ideogram always passes a boolean block-diagonal mask).
+
+COMPILE:
+  The mask->``cu_seqlens`` transform is data-dependent (``.any()``, ``nonzero``, ``.item()``),
+  so it forces device->host reads that ``torch.compile`` cannot trace -- each is a graph break,
+  and under FSDP2 a mid-forward break splits the region the per-layer all-gather/reshard
+  collectives are registered around and desyncs them across ranks (a silent multi-rank death).
+  The fix is to not derive the packing here at all: ``Ideogram4Adapter`` builds ``cu_seqlens``
+  on the host (it already holds the text lengths as Python ints) and publishes it into a shared
+  non-persistent buffer on the attention modules, making it a graph INPUT (see
+  ``packing_buffer.py`` for why that route and not ``attention_kwargs``, which is a
+  dead parameter for this model in diffusers 0.39.0). That path is exact on ragged batches,
+  needs no host sync, and stops recomputing the same packing once per layer.
+
+Activation (env, no config schema change):
+    PRIMUS_IDEOGRAM_VARLEN_ATTN=1          swap Ideogram-4 attention to the var-len flash path
+    PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS  default 1; set 0 to stop the adapter precomputing
+                                           the packing and fall back to the mask-derived
+                                           (host-syncing) path. Escape hatch for A/B and
+                                           rollback.
+    PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE=1    skip the mask->cu_seqlens host-sync and use dense
+                                           flash directly (torch.compile-safe; EXACT only for
+                                           equal-length / unpadded batches, e.g. fixed-text).
+                                           Never fires while a precomputed cu_seqlens is
+                                           provided, which now takes precedence.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from primus.backends.nemo_automodel.models.ideogram4.packing_buffer import resolve_packing
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+
+def is_varlen_attn_enabled() -> bool:
+    """Whether the Ideogram-4 var-len flash-attention processor should be installed."""
+    return os.getenv("PRIMUS_IDEOGRAM_VARLEN_ATTN", "0") in _TRUTHY
+
+
+def assume_dense_enabled() -> bool:
+    """Whether to skip the block-diagonal mask analysis and use dense flash directly.
+
+    The mask->``cu_seqlens`` transform does data-dependent host-syncs (``bool(.any())``,
+    ``.max().item()``, ``nonzero``) that graph-break ``torch.compile`` and, under FSDP2
+    multi-rank, desync the per-layer collectives. When every row is a single full segment
+    (an equal-length / unpadded batch -> the mask is trivial anyway), set
+    ``PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE=1`` to skip the analysis and go straight to dense
+    flash, keeping the compiled per-layer graph break-free. EXACT only for equal-length /
+    no-pad batches (e.g. fixed-text); do NOT set it for ragged/padded batches.
+
+    Superseded in practice by :func:`precompute_cu_seqlens_enabled` -- when the adapter
+    hands the packing in, the processor prefers it and this flag never fires.
+    """
+    return os.getenv("PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE", "0") in _TRUTHY
+
+
+def precompute_cu_seqlens_enabled() -> bool:
+    """Whether the adapter precomputes ``cu_seqlens`` on the host. Default ON.
+
+    The adapter already holds the per-sample text lengths as Python ints, so it can build
+    the var-len packing itself and pass it into the processor as a plain tensor -- a graph
+    INPUT rather than something derived from the mask mid-graph. That removes the
+    device->host syncs, so the exact var-len path compiles on ragged batches.
+
+    Set ``PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS=0`` to restore the previous behaviour
+    (mask-derived packing, or dense flash under ``ASSUME_DENSE``). Kept as an escape hatch
+    for A/B measurement and rollback, not as a routine knob.
+    """
+    return os.getenv("PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS", "1") in _TRUTHY
+
+
+def precompute_cu_seqlens_active() -> bool:
+    """Whether precomputing the packing will actually be *read* by anything.
+
+    Both switches have to be on. Without ``PRIMUS_IDEOGRAM_VARLEN_ATTN`` the stock SDPA
+    processor is in place and has no ``cu_seqlens`` parameter, so precomputing would cost a
+    build plus the reserved pad column's token position every step for nothing. This is the
+    gate the adapter uses; :func:`precompute_cu_seqlens_enabled` is only the flag.
+    """
+    return is_varlen_attn_enabled() and precompute_cu_seqlens_enabled()
+
+
+# --------------------------------------------------------------------------- #
+# Model-agnostic helpers (reusable for any packed/padded block-diagonal attn)  #
+# --------------------------------------------------------------------------- #
+def blockdiag_bool_mask_to_cu_seqlens(mask: Tensor) -> Tuple[Tensor, int, bool]:
+    """Convert a block-diagonal boolean attention mask to a var-len packing.
+
+    Args:
+        mask: ``(B, 1, L, L)`` or ``(B, L, L)`` boolean tensor, ``True`` = "query i
+            attends to key j". Each row ``b`` must be block-diagonal over CONTIGUOUS
+            segments (a segment boundary lies between positions ``i`` and ``i+1``
+            wherever ``mask[b, .., i, i+1]`` is ``False``); segments never span the row
+            boundary. This is what an ``(seg_i == seg_j)`` mask produces for
+            contiguously-assigned ``segment_ids`` (padding, or packed samples).
+
+    Returns:
+        ``(cu_seqlens, max_seqlen, is_trivial)`` describing the packed layout over the
+        flattened ``(B*L)`` sequence:
+          * ``cu_seqlens``: ``int32`` ``(num_segments + 1,)`` cumulative segment lengths,
+            starting at 0 and ending at ``B*L`` (the ``flash_attn_varlen_func`` contract).
+          * ``max_seqlen``: longest segment length.
+          * ``is_trivial``: ``True`` iff every row is a single full segment (no internal
+            splits) — the caller may then use plain dense attention and skip packing.
+    """
+    if mask.dim() == 4:
+        m = mask[:, 0]
+    elif mask.dim() == 3:
+        m = mask
+    else:
+        raise ValueError(f"expected a (B,1,L,L) or (B,L,L) mask, got shape {tuple(mask.shape)}")
+    if m.dtype != torch.bool:
+        raise TypeError(f"expected a boolean mask, got dtype {m.dtype}")
+
+    B, L, L2 = m.shape
+    if L != L2:
+        raise ValueError(f"mask must be square in the last two dims, got {(L, L2)}")
+    device = m.device
+
+    if L == 1:
+        cu = torch.arange(0, B + 1, dtype=torch.int32, device=device)
+        return cu, 1, True
+
+    # A new segment starts at position i+1 within a row wherever i and i+1 do not attend
+    # (the sub-diagonal test is all we need for contiguous block-diagonal segments).
+    superdiag = m.diagonal(offset=1, dim1=-2, dim2=-1)  # (B, L-1) : mask[b, i, i+1]
+    splits = ~superdiag  # (B, L-1)
+    is_trivial = not bool(splits.any())
+
+    # Per-token "starts a new segment" over the flattened (B, L) sequence. The first
+    # token of every row always starts a segment (segments never cross rows).
+    new_seg = torch.zeros(B, L, dtype=torch.bool, device=device)
+    new_seg[:, 0] = True
+    new_seg[:, 1:] = splits
+    seg_starts = torch.nonzero(new_seg.reshape(-1), as_tuple=False).flatten()
+
+    total = B * L
+    cu = torch.empty(seg_starts.numel() + 1, dtype=torch.int32, device=device)
+    cu[:-1] = seg_starts.to(torch.int32)
+    cu[-1] = total
+    max_seqlen = int((cu[1:] - cu[:-1]).max().item())
+    return cu, max_seqlen, is_trivial
+
+
+def _unwrap(out):
+    """aiter returns (out, lse, ...) when return_lse=True; keep only the output."""
+    return out[0] if isinstance(out, (tuple, list)) else out
+
+
+def dense_flash_attention(q: Tensor, k: Tensor, v: Tensor, *, deterministic: bool = False) -> Tensor:
+    """Plain (unmasked) bf16 flash attention. q/k/v: ``(B, L, H, D)`` -> ``(B, L, H, D)``.
+
+    ``return_lse=True`` is required by aiter's autograd forward (LSE is saved for the
+    backward); we request it and drop the LSE.
+    """
+    import aiter
+
+    return _unwrap(
+        aiter.flash_attn_func(q, k, v, causal=False, deterministic=deterministic, return_lse=True)
+    )
+
+
+def varlen_flash_attention(
+    q: Tensor, k: Tensor, v: Tensor, cu_seqlens: Tensor, max_seqlen: int, *, deterministic: bool = False
+) -> Tensor:
+    """Variable-length bf16 flash attention over a packed sequence.
+
+    q/k/v: ``(total_tokens, H, D)`` (packed, no padding between segments). ``cu_seqlens``
+    and ``max_seqlen`` come from :func:`blockdiag_bool_mask_to_cu_seqlens`. Returns
+    ``(total_tokens, H, D)``.
+    """
+    import aiter
+
+    return _unwrap(
+        aiter.flash_attn_varlen_func(
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
+            causal=False, deterministic=deterministic, return_lse=True,
+        )
+    )
+
+
+def _rotate_half(x: Tensor) -> Tensor:
+    """Rotate-half, matching diffusers ``transformer_ideogram4._rotate_half`` exactly."""
+    half = x.shape[-1] // 2
+    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+
+# --------------------------------------------------------------------------- #
+# Ideogram-4 var-len attention processor                                       #
+# --------------------------------------------------------------------------- #
+_warned: set[str] = set()
+
+
+def _warn_once(key: str, msg: str) -> None:
+    if key not in _warned:
+        _warned.add(key)
+        logger.warning(msg)
+
+
+class Ideogram4VarlenAttnProcessor:
+    """Ideogram-4 self-attention via var-len flash (exact block-diagonal packing).
+
+    Mirrors diffusers ``Ideogram4AttnProcessor`` (q/k/v proj, q/k RMSNorm, MRoPE, output
+    proj) and swaps only the attention call for a ``cu_seqlens`` flash path. Non-det
+    backward (``deterministic=False``).
+
+    ``_attention`` picks a path in this order:
+
+      1. **provided** ``cu_seqlens`` -- the adapter precomputed the packing on the host and
+         published it into the shared buffer this processor reads off ``attn``. Exact on
+         ragged batches and free of data-dependent ops, so it is the only path that is
+         simultaneously correct on ragged data and safe under per-layer ``torch.compile``
+         + FSDP2.
+      2. ``assume_dense`` **or no mask** -- dense flash. Exact only when no row has padding.
+      3. **mask analysis** (legacy) -- exact, but its device->host reads graph-break.
+    """
+
+    # kept for API-compatibility with diffusers' processor discovery / set_attention_backend
+    _attention_backend = None
+    _parallel_config = None
+
+    deterministic: bool = False
+    # Read once at class-definition time (torchrun sets env before import). Keeps the check a
+    # constant attribute lookup inside the compiled graph (no data-dependent branch / break).
+    assume_dense: bool = assume_dense_enabled()
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor],
+        image_rotary_emb: Tuple[Tensor, Tensor],
+        cu_seqlens: Optional[Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ) -> Tensor:
+        # The packing normally arrives on the module: the adapter publishes it into a shared
+        # non-persistent buffer that ``resolve_packing`` reads off ``attn`` (see
+        # ``packing_buffer.py``). The two named parameters are kept because
+        # diffusers' attention module filters forwarded kwargs against
+        # ``inspect.signature(self.processor.__call__).parameters``, so declaring them by name
+        # is what would let a kwargs route work at all -- a ``**kwargs``-only processor would
+        # silently receive nothing. An explicit argument wins over the buffer.
+        cu_seqlens, max_seqlen = resolve_packing(attn, cu_seqlens, max_seqlen)
+
+        query = attn.to_q(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
+        key = attn.to_k(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
+        value = attn.to_v(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        cos, sin = image_rotary_emb
+        cos = cos.unsqueeze(2)
+        sin = sin.unsqueeze(2)
+        query = (query * cos) + (_rotate_half(query) * sin)
+        key = (key * cos) + (_rotate_half(key) * sin)
+
+        out = self._attention(query, key, value, attention_mask, cu_seqlens, max_seqlen)
+        out = out.flatten(2, 3)
+        return attn.to_out[0](out)
+
+    def _attention(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        attention_mask: Optional[Tensor],
+        cu_seqlens: Optional[Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ) -> Tensor:
+        B, L, H, D = query.shape
+
+        # No mask, or a mask type we cannot represent as contiguous segments: for a
+        # boolean block-diagonal mask we always take the exact var-len path (that is the
+        # case this processor exists for); a non-boolean/additive mask is not something
+        # Ideogram emits, so defer to the original dense dispatch rather than guess.
+        if attention_mask is not None and attention_mask.dtype != torch.bool:
+            _warn_once(
+                "nonbool_mask",
+                "[PrimusIdeogramVarlen] non-boolean attn_mask -> dense SDPA dispatch fallback.",
+            )
+            from diffusers.models.attention_dispatch import dispatch_attention_fn
+
+            return dispatch_attention_fn(
+                query, key, value, attn_mask=attention_mask,
+                backend=self._attention_backend, parallel_config=self._parallel_config,
+            )
+
+        # MRoPE multiplies q/k by the float32 cos/sin, promoting them to float32. Torch
+        # SDPA is autocast-aware and downcasts at its boundary; aiter's flash op is not,
+        # so match that here by casting q/k to the projection (compute) dtype. ``value``
+        # skips RoPE and already carries it.
+        compute_dtype = value.dtype
+        query = query.to(compute_dtype)
+        key = key.to(compute_dtype)
+
+        # PREFERRED PATH: the packing was precomputed on the host and handed in, so nothing
+        # here inspects tensor values. Exact on ragged batches AND compile-safe -- this is
+        # the only combination of those two properties available (aiter.flash_attn_varlen_func
+        # is itself traceable, so the whole processor stays in one graph and the FSDP2
+        # per-layer collectives keep their ordering). ``attention_mask`` is deliberately not
+        # read on this path; the model still materializes it, which is now dead weight.
+        # The branch tests a property of the RUN (was anything provided), not of the data,
+        # so it costs one guard and no graph break.
+        if cu_seqlens is not None:
+            # The packing lives on the module and outlives the step that published it, so a
+            # caller that bypasses the adapter (a sampling pass, an eval loop with a different
+            # batch size) could otherwise attend on a stale one and corrupt silently. This
+            # compares only static SHAPE metadata -- no values are read, so it costs a guard
+            # and no host sync. ``2B+1`` is the layout contract: one pad segment plus one
+            # text+image segment per row. Packing several samples per row would change that
+            # count, and this check with it.
+            expected = 2 * B + 1
+            if cu_seqlens.numel() != expected:
+                raise ValueError(
+                    f"cu_seqlens has {cu_seqlens.numel()} entries but this batch needs "
+                    f"{expected} (2*B+1 for B={B}). Either the packing was published for a "
+                    "different batch size and is stale -- call "
+                    "ideogram4_packing_buffer.clear_packing(model) before running the model "
+                    "outside the adapter -- or the sequence layout no longer has exactly two "
+                    "segments per row, in which case update this check."
+                )
+            # Hand the kernel a PRIVATE COPY, never the shared buffer. aiter's var-len op
+            # treats cu_seqlens as a mutable argument: it saves the tensor for its backward
+            # and then writes it, bumping the autograd version counter once per call. One
+            # buffer shared by 34 layers therefore has its version moved 34 times per forward
+            # while each layer's backward still expects the version it saved, and the step
+            # dies in .backward() with "a variable needed for gradient computation has been
+            # modified by an inplace operation: IntTensor[5] is at version 35; expected 34"
+            # (measured 2026-08-04, t2 stage; the legacy path never saw it because every
+            # layer derives its own tensor from the mask). The clone is 5 int32 already on
+            # device, and under compile it is a graph intermediate -- the buffer stays a
+            # read-only graph input, which is what makes sharing it safe at all.
+            out = varlen_flash_attention(
+                query.reshape(B * L, H, D),
+                key.reshape(B * L, H, D),
+                value.reshape(B * L, H, D),
+                cu_seqlens.clone(),
+                L if max_seqlen is None else max_seqlen,
+                deterministic=self.deterministic,
+            )
+            return out.reshape(B, L, H, D)
+
+        # Dense fast path. ``attention_mask is None`` OR the equal-length assertion
+        # (``assume_dense``) skips the data-dependent mask->cu_seqlens host-sync, so under
+        # torch.compile the whole processor stays in one graph (aiter.flash_attn_func is
+        # itself compile-safe -- fullgraph traces it with no break), which is required for
+        # FSDP2 multi-rank compile to not desync collectives. Exact for equal-length/no-pad
+        # batches; ragged/padded batches must leave PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE unset.
+        if attention_mask is None or self.assume_dense:
+            return dense_flash_attention(query, key, value, deterministic=self.deterministic)
+
+        # LEGACY PATH: derive the packing from the mask. Exact, but the derivation does
+        # device->host reads, so it graph-breaks and is unsafe under multi-rank compile.
+        # Reached only when nothing was precomputed (kill switch, or a non-adapter caller).
+        # Retained as the reference implementation the unit test checks against.
+        cu_from_mask, max_from_mask, is_trivial = blockdiag_bool_mask_to_cu_seqlens(attention_mask)
+        if is_trivial:
+            return dense_flash_attention(query, key, value, deterministic=self.deterministic)
+
+        q = query.reshape(B * L, H, D)
+        k = key.reshape(B * L, H, D)
+        v = value.reshape(B * L, H, D)
+        out = varlen_flash_attention(q, k, v, cu_from_mask, max_from_mask, deterministic=self.deterministic)
+        return out.reshape(B, L, H, D)
+
+
+# --------------------------------------------------------------------------- #
+# Install                                                                      #
+# --------------------------------------------------------------------------- #
+def install(model=None) -> bool:
+    """Route Ideogram-4 attention through the var-len flash processor (no-fork).
+
+    No-op (returns False) unless ``PRIMUS_IDEOGRAM_VARLEN_ATTN`` is set. Patches the
+    class default processor so every ``Ideogram4Attention`` built AFTER this call uses
+    the var-len processor; if a built ``model`` is passed, also swaps its existing
+    modules. Idempotent. Modifies NO Automodel/diffusers source.
+    """
+    if not is_varlen_attn_enabled():
+        return False
+
+    # Fail fast if aiter's flash-attention is unavailable, so the run errors clearly
+    # rather than silently keeping the SDPA path.
+    import aiter  # noqa: F401
+
+    from diffusers.models.transformers.transformer_ideogram4 import Ideogram4Attention
+
+    already = getattr(Ideogram4Attention, "_primus_varlen_installed", False)
+    if not already:
+        Ideogram4Attention._default_processor_cls = Ideogram4VarlenAttnProcessor
+        if Ideogram4VarlenAttnProcessor not in Ideogram4Attention._available_processors:
+            Ideogram4Attention._available_processors = [
+                *Ideogram4Attention._available_processors,
+                Ideogram4VarlenAttnProcessor,
+            ]
+        Ideogram4Attention._primus_varlen_installed = True
+
+    swapped = 0
+    if model is not None:
+        for module in model.modules():
+            if isinstance(module, Ideogram4Attention) and not isinstance(
+                module.processor, Ideogram4VarlenAttnProcessor
+            ):
+                module.set_processor(Ideogram4VarlenAttnProcessor())
+                swapped += 1
+
+    if already and swapped == 0:
+        return True
+
+    logger.info(
+        "[PrimusIdeogramVarlen] Installed var-len flash-attention processor for "
+        "Ideogram4Attention (deterministic=False)%s.",
+        f"; swapped {swapped} existing module(s)" if swapped else "",
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/context_parallel.py b/primus/backends/nemo_automodel/models/ideogram4/context_parallel.py
new file mode 100644
index 000000000..5be17f2a2
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/context_parallel.py
@@ -0,0 +1,203 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Context parallelism (Ulysses) for Ideogram-4 in the AutoModel diffusion recipe.
+
+WHY:
+  AutoModel's diffusion CP support delegates to diffusers' ContextParallelConfig
+  and refuses to enable CP unless the transformer declares a ``_cp_plan``
+  (``_diffusers/auto_diffusion_pipeline.py``). Every other piece CP needs -- the
+  mesh with a ``cp`` axis, FSDP2 grad sync over ``dp_shard_cp``, the per-DP-rank
+  RNG reseed, the ``loss * cp_size`` correction -- already exists upstream.
+  ``Ideogram4Transformer2DModel`` simply has no plan, so CP is unreachable for it
+  and for no other reason. This module supplies one, Primus-side: diffusers stays
+  pristine.
+
+THE PLAN, and why each entry is what it is:
+  Ideogram-4 is a SINGLE-STREAM model. ``hidden_states``, ``encoder_hidden_states``,
+  ``position_ids``, ``segment_ids`` and ``indicator`` are all length-L along the
+  *same* packed ``[left-pad][text][image]`` sequence axis -- there is no separate
+  text stream to leave alone (contrast QwenImage, which splits an encoder stream
+  of its own). So every per-token input that is combined elementwise with
+  ``hidden_states`` has to be split with it, or the shapes stop matching the moment
+  the root forward multiplies by the indicator masks.
+
+  * ``hidden_states`` / ``encoder_hidden_states`` / ``indicator`` -- split at the
+    ROOT. The root forward derives ``llm_token_mask`` / ``output_image_mask`` from
+    ``indicator`` and applies them to the other two before the first block, so all
+    three must be sliced at the same point.
+  * ``segment_ids`` -- deliberately NOT split. See the mask note below; this is the
+    one entry whose absence is load-bearing.
+  * ``position_ids`` -- also not split. Instead the MRoPE module's two OUTPUTS are
+    split (``split_output=True``), which is the Wan/QwenImage pattern: let the layer
+    see the full positions, then slice what it produced. RoPE is applied to q/k
+    *before* the Ulysses all-to-all, so cos/sin must end up LOCAL and aligned with
+    the local tokens. Splitting the output rather than the input also keeps MRoPE's
+    autocast-disabled fp32 arithmetic (needed because image positions start at
+    65536) operating on exactly the tensor it was written for.
+  * ``timestep`` -- not split, and must not be. Our adapter passes it per-sample,
+    shape ``(B,)``, which the root forward unsqueezes to ``(B, 1, ...)`` and
+    broadcasts. Wan splits its timestep only because Wan's is per-token. ``install``
+    below is a no-op for that difference, so ``_assert_timestep_is_per_sample``
+    guards it instead of letting a per-token timestep silently mis-slice.
+  * ``final_layer`` -- the gather point (Ideogram's equivalent of the ``proj_out``
+    that Flux/Wan/QwenImage gather at). Nothing happens after it but the output
+    wrapper, so gathering there restores full-length ``(B, L, out_channels)``.
+
+WHY ``segment_ids`` STAYS WHOLE:
+  ``TemplatedUlyssesAttention`` computes ``H_LOCAL = H // world_size`` and then does
+  the all-to-all, so after the exchange each rank holds the FULL sequence with a
+  SUBSET of heads. Attention therefore wants a full-length mask. The root forward
+  builds ``(segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)``,
+  which is full-length ``(B, 1, L, L)`` exactly when ``segment_ids`` is left unsplit
+  while ``hidden_states`` is split. Splitting it too would produce an
+  ``(B, 1, L_local, L_local)`` mask that silently drops all cross-shard attention.
+  diffusers anticipates the other convention as well (a mask arriving at
+  ``S_KV_LOCAL`` gets all-gathered back), so both shapes are survivable upstream --
+  which is precisely why getting it wrong would not raise. Verify by A/B-ing CP vs
+  non-CP logits on a RAGGED batch before trusting any loss curve.
+
+WHY THE VAR-LEN PROCESSOR IS REFUSED (``_assert_varlen_not_installed``):
+  ``Ideogram4VarlenAttnProcessor`` calls ``aiter.flash_attn_varlen_func`` directly
+  and never reaches diffusers' ``dispatch_attention_fn`` -- which is where the
+  Ulysses all-to-all lives. Combined with CP the result is not an error: the plan's
+  hooks still split the inputs, the processor then attends only WITHIN each rank's
+  shard, the output is gathered, and training proceeds on a plausible-looking loss
+  curve. Its precomputed ``cu_seqlens`` describes global segment boundaries on top
+  of that, so the segment map is wrong as well. diffusers' own CP-capability check
+  cannot catch this, because we bypassed the dispatch it guards. Hence the explicit
+  refusal here. The two are an either/or until an aiter CP forward/backward op pair
+  exists.
+
+CONSTRAINTS worth knowing before configuring a run:
+  * Ulysses only -- ring degree > 1 is rejected upstream (broken backward in
+    diffusers <= 0.39).
+  * FSDP2 only -- DDP + CP is rejected upstream, so this and the DDP/ZeRO-1 path
+    are mutually exclusive.
+  * Ulysses shards HEADS, so ``cp_size`` must divide both the world size and the
+    head count. Ideogram-4 has 18 heads, so on one 8-GPU node the only usable
+    degrees are 1 and 2: ``cp_size=2, dp_size=4``.
+  * The post-patchify sequence length must be divisible by ``cp_size``
+    (equipartition sharding).
+
+Installed unconditionally: attaching a class attribute costs nothing and is inert
+until a config asks for ``cp_size > 1``, so this needs no environment flag.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _build_cp_plan():
+    """Build the diffusers CP plan. Imported lazily so this module stays importable
+    without diffusers present (the trainer installs hooks defensively)."""
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+    )
+
+    return {
+        # Root forward inputs. All three are per-token on the same packed axis and
+        # are combined with each other before the first block, so they split together.
+        # segment_ids and position_ids are absent on purpose -- see the module docstring.
+        "": {
+            "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+            "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+            "indicator": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
+        },
+        # MRoPE returns (cos, sin), each (B, L, head_dim). Split what it produced,
+        # not what it consumed, so RoPE is applied to the local tokens pre-all-to-all.
+        "rotary_emb": {
+            0: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True),
+            1: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True),
+        },
+        # Gather back to full length; nothing follows this but the output wrapper.
+        "final_layer": ContextParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
+
+def _assert_varlen_not_installed() -> None:
+    """Refuse CP + the var-len flash processor. Silent corruption otherwise."""
+    from primus.backends.nemo_automodel.models.ideogram4.attention import (
+        is_varlen_attn_enabled,
+    )
+
+    if is_varlen_attn_enabled():
+        raise ValueError(
+            "Context parallelism (cp_size > 1) and PRIMUS_IDEOGRAM_VARLEN_ATTN=1 cannot be "
+            "used together. The var-len processor calls aiter directly and never reaches "
+            "diffusers' dispatch_attention_fn, which is where the Ulysses all-to-all happens, "
+            "so each rank would attend only within its own sequence shard and training would "
+            "produce wrong gradients WITHOUT raising. Unset PRIMUS_IDEOGRAM_VARLEN_ATTN to run "
+            "CP on the stock SDPA path, or set cp_size=1 to keep the var-len path."
+        )
+
+
+def _timestep_pre_hook(module, args, kwargs):
+    """Assert the plan's unsplit-``timestep`` assumption still holds, once per forward.
+
+    Registered on the INSTANCE when CP is switched on, deliberately not by wrapping
+    ``forward``: diffusers resolves the root plan's named inputs against the forward
+    signature, so a ``*args, **kwargs`` wrapper would hide ``encoder_hidden_states``
+    and ``indicator`` from it and quietly leave them unsplit.
+    """
+    timestep = kwargs.get("timestep")
+    if timestep is None and len(args) >= 2:
+        timestep = args[1]
+    _assert_timestep_is_per_sample(timestep)
+    return None
+
+
+def _assert_timestep_is_per_sample(timestep) -> None:
+    """The plan leaves ``timestep`` unsplit, which is only correct while it is ``(B,)``."""
+    if timestep is not None and hasattr(timestep, "dim") and timestep.dim() != 1:
+        raise ValueError(
+            f"Context parallelism expects a per-sample timestep of shape (B,), got "
+            f"{tuple(timestep.shape)}. A per-token timestep is broadcast against the "
+            "sequence axis and would need splitting in the CP plan (as Wan does); it is "
+            "left unsplit here because the Ideogram adapter passes t = 1 - sigma per sample."
+        )
+
+
+def install(model=None) -> bool:
+    """Attach a diffusers context-parallel plan to Ideogram-4 (no-fork).
+
+    Sets the ``_cp_plan`` class attribute so AutoModel's CP path accepts the model,
+    and wraps ``enable_parallelism`` -- which upstream calls only when
+    ``cp_size > 1`` -- so the incompatible var-len combination and a per-token
+    timestep are rejected at the moment CP is actually switched on rather than
+    silently miscomputing. Idempotent. Modifies NO Automodel/diffusers source.
+    """
+    try:
+        from diffusers.models.transformers.transformer_ideogram4 import (
+            Ideogram4Transformer2DModel,
+        )
+    except ImportError:
+        return False
+
+    if getattr(Ideogram4Transformer2DModel, "_primus_cp_installed", False):
+        return False
+
+    Ideogram4Transformer2DModel._cp_plan = _build_cp_plan()
+
+    original_enable_parallelism = Ideogram4Transformer2DModel.enable_parallelism
+
+    def enable_parallelism(self, *args, **kwargs):
+        _assert_varlen_not_installed()
+        result = original_enable_parallelism(self, *args, **kwargs)
+        self.register_forward_pre_hook(_timestep_pre_hook, with_kwargs=True)
+        return result
+
+    Ideogram4Transformer2DModel.enable_parallelism = enable_parallelism
+    Ideogram4Transformer2DModel._primus_cp_installed = True
+
+    logger.info(
+        "[PrimusIdeogramCP] Attached diffusers context-parallel plan to "
+        "Ideogram4Transformer2DModel (Ulysses; cp_size must divide both the world size "
+        "and the 18 attention heads, so cp_size=2 is the only non-trivial degree on an "
+        "8-GPU node)."
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/data/__init__.py b/primus/backends/nemo_automodel/models/ideogram4/data/__init__.py
new file mode 100644
index 000000000..54a45586e
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/data/__init__.py
@@ -0,0 +1,18 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Ideogram-4 dataloaders.
+
+Both emit the same batch contract the adapter and the flow-matching pipeline
+consume (``image_latents``, ``llm_features``, ``text_lengths``, ``data_type``),
+so a config can swap one for the other:
+
+    synthetic.py  fixed per-index synthetic tensors for an overfit smoke, no
+                  encoder weights and no dataset required
+    cache.py      the real pre-encoded cache produced by ``processor.py``
+
+They are referenced from YAML by ``_target_``, not imported here.
+"""
diff --git a/primus/backends/nemo_automodel/models/ideogram4/data/build.py b/primus/backends/nemo_automodel/models/ideogram4/data/build.py
new file mode 100644
index 000000000..853ce5a70
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/data/build.py
@@ -0,0 +1,199 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Build the Ideogram-4 pre-encoded cache that ``cache.py`` reads.
+
+Pairs each image with its caption, encodes both with :class:`Ideogram4Processor`
+(Flux-2 VAE for the latents, Qwen3-VL for the text features) and writes the flat
+layout the dataloader expects:
+
+    <out_dir>/metadata.json      index + grid/feature dims
+    <out_dir>/samples/<i>.pt     {image_latents, llm_features, text_length, ...}
+
+Encoding is a one-off offline step: it needs one GPU and the gated VAE and
+text-encoder weights, neither of which training itself requires.
+
+Selection is **deterministic** by default — images are taken in sorted order and
+samples that cannot be encoded are skipped, not substituted — so the same source
+directory and arguments reproduce the same cache. ``shuffle`` opts into a
+seeded shuffle instead.
+
+Reached from the CLI as ``primus data automodel-cache --model ideogram4``.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
+
+DEFAULT_VAE_SOURCE = "ideogram-ai/ideogram-4-nf4-diffusers"
+DEFAULT_TEXT_ENCODER_SOURCE = "Qwen/Qwen3-VL-8B-Instruct"
+
+DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+
+
+def build_cache(
+    *,
+    image_dir: str,
+    caption_dir: str,
+    output_dir: str,
+    num_samples: int = 1024,
+    resolution: int = 256,
+    max_text_tokens: int = 128,
+    vae_source: str = DEFAULT_VAE_SOURCE,
+    text_encoder_source: str = DEFAULT_TEXT_ENCODER_SOURCE,
+    tokenizer_source: Optional[str] = None,
+    device: str = "cuda",
+    dtype: str = "bf16",
+    seed: int = 1234,
+    shuffle: bool = False,
+    log_every: int = 64,
+) -> Dict[str, Any]:
+    """Encode up to ``num_samples`` image/caption pairs into a flat cache.
+
+    Captions are read from ``<caption_dir>/<image stem>.txt``. Returns the
+    metadata dict that was written to ``metadata.json``.
+    """
+    from primus.backends.nemo_automodel.models.ideogram4.processor import (
+        Ideogram4Processor,
+    )
+
+    if dtype not in DTYPES:
+        raise ValueError(f"dtype must be one of {sorted(DTYPES)}, got {dtype!r}")
+
+    image_path_root = Path(image_dir)
+    caption_path_root = Path(caption_dir)
+    out_root = Path(output_dir)
+    for label, path in (("image_dir", image_path_root), ("caption_dir", caption_path_root)):
+        if not path.is_dir():
+            raise NotADirectoryError(f"{label} does not exist: {path}")
+    (out_root / "samples").mkdir(parents=True, exist_ok=True)
+
+    images = sorted(p for p in image_path_root.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS)
+    if shuffle:
+        random.Random(seed).shuffle(images)
+    if not images:
+        raise FileNotFoundError(f"no images with extensions {IMAGE_EXTENSIONS} under {image_path_root}")
+    logger.info("Found %d images in %s", len(images), image_path_root)
+
+    processor = Ideogram4Processor(
+        vae_source=vae_source,
+        text_encoder_source=text_encoder_source,
+        tokenizer_source=tokenizer_source,
+        device=device,
+        dtype=DTYPES[dtype],
+    )
+    processor.load_models()
+
+    samples_meta: list[dict] = []
+    grid_h = grid_w = feature_dim = in_channels = None
+    skipped_caption = skipped_too_long = skipped_error = 0
+    start = time.time()
+
+    for image_file in images:
+        if len(samples_meta) >= num_samples:
+            break
+
+        caption_file = caption_path_root / (image_file.stem + ".txt")
+        if not caption_file.exists():
+            skipped_caption += 1
+            continue
+        prompt = caption_file.read_text(encoding="utf-8").strip()
+        if not prompt:
+            skipped_caption += 1
+            continue
+
+        try:
+            image_tensor = processor.preprocess_image(Image.open(image_file), resolution)
+            latents = processor.encode_image(image_tensor)
+            text_encoding = processor.encode_text(prompt, max_text_tokens)
+            if text_encoding is None:
+                # Caption does not fit in max_text_tokens; the cache loader assumes
+                # every sample fits, so drop it rather than truncate the prompt.
+                skipped_too_long += 1
+                continue
+            sample = processor.get_cache_data(
+                latents, text_encoding, prompt=prompt, image_path=str(image_file)
+            )
+        except Exception as e:
+            logger.warning("Encode failed for %s: %s", image_file.name, e)
+            skipped_error += 1
+            continue
+
+        if grid_h is None:
+            grid_h, grid_w = sample["grid_h"], sample["grid_w"]
+            feature_dim = int(text_encoding["llm_features"].shape[-1])
+            in_channels = int(latents.shape[0])
+
+        index = len(samples_meta)
+        relative_path = f"samples/{index}.pt"
+        torch.save(sample, out_root / relative_path)
+        samples_meta.append(
+            {"cache_file": relative_path, "text_length": sample["text_length"], "prompt": prompt}
+        )
+
+        if log_every and len(samples_meta) % log_every == 0:
+            elapsed = max(time.time() - start, 1e-6)
+            logger.info(
+                "Encoded %d/%d (%.2f samples/s)", len(samples_meta), num_samples, len(samples_meta) / elapsed
+            )
+
+    if not samples_meta:
+        raise RuntimeError(
+            f"encoded 0 samples from {len(images)} images "
+            f"(skipped: no/empty caption={skipped_caption}, caption too long={skipped_too_long}, "
+            f"encode error={skipped_error})"
+        )
+
+    metadata = {
+        "model_type": "ideogram4",
+        "resolution": resolution,
+        "grid_h": grid_h,
+        "grid_w": grid_w,
+        "in_channels": in_channels,
+        "llm_features_dim": feature_dim,
+        "max_text_tokens": max_text_tokens,
+        "vae_source": vae_source,
+        "text_encoder_source": text_encoder_source,
+        "num_samples": len(samples_meta),
+        "samples": samples_meta,
+    }
+    with open(out_root / "metadata.json", "w") as f:
+        json.dump(metadata, f)
+
+    elapsed = max(time.time() - start, 1e-6)
+    logger.info("=" * 70)
+    logger.info("Wrote %d samples to %s", len(samples_meta), out_root)
+    logger.info(
+        "grid=%sx%s llm_features_dim=%s in_channels=%s", grid_h, grid_w, feature_dim, in_channels
+    )
+    if skipped_caption or skipped_too_long or skipped_error:
+        logger.info(
+            "Skipped: no/empty caption=%d, caption too long=%d, encode error=%d",
+            skipped_caption,
+            skipped_too_long,
+            skipped_error,
+        )
+    if len(samples_meta) < num_samples:
+        logger.warning(
+            "Requested %d samples but only %d were encodable from %d images",
+            num_samples,
+            len(samples_meta),
+            len(images),
+        )
+    logger.info("Elapsed %.1fs (%.2f samples/s)", elapsed, len(samples_meta) / elapsed)
+    logger.info("=" * 70)
+
+    return metadata
diff --git a/primus/backends/nemo_automodel/models/ideogram4/data/cache.py b/primus/backends/nemo_automodel/models/ideogram4/data/cache.py
new file mode 100644
index 000000000..611c93c9e
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/data/cache.py
@@ -0,0 +1,221 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Real (pre-encoded) Ideogram-4 cache dataloader (no-fork, Primus-side).
+
+Reads the flat per-sample cache produced by :class:`Ideogram4Processor`
+(``processor.py``) and emits the exact batch the
+:class:`Ideogram4Adapter` + ``FlowMatchingPipeline`` consume — the same contract as
+``SyntheticIdeogram4DataloaderConfig``, but with REAL Flux-2 VAE latents + Qwen3-VL
+features:
+
+  - ``image_latents``  ``[B, 128, gh, gw]``  packed+BN latents (x0)
+  - ``llm_features``   ``[B, Tmax, 53248]``  LEFT-padded Qwen3-VL feats
+  - ``text_lengths``   ``[B]``               real (non-pad) token count per sample
+  - ``data_type``      ``"image"``
+
+Left-padding matches the adapter/pipeline ``[left-pad][text][image]`` layout: the
+real ``n`` tokens occupy the LAST ``n`` rows (positions ``[Tmax-n : Tmax]``), which is
+exactly the region ``_prepare_ids`` marks as text (``offset = Tmax - n``).
+
+``Tmax`` is a **dataset-level constant** by default (the longest caption in the cache,
+read from ``metadata.json``), NOT the per-batch maximum. The packed sequence length
+``S = Tmax + gh*gw`` therefore does not move between steps. That is a requirement of
+per-layer ``torch.compile``, which keys its compiled graphs on input shapes: padding to
+the per-batch maximum makes ``S`` jump whenever a batch's longest caption differs, which
+recompiles the graph and blocks compile on real data no matter how ``cu_seqlens`` is
+built. Configure via ``max_text_tokens`` (0 = derive, -1 = legacy per-batch).
+
+Cache layout (``cache_dir``):
+  - ``metadata.json``: ``{"grid_h","grid_w","llm_features_dim","in_channels",
+    "samples":[{"cache_file","text_length","prompt"}, ...]}``
+  - ``samples/<i>.pt``: ``{image_latents [128,gh,gw], llm_features [n,53248],
+    text_length, ...}``
+"""
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+from nemo_automodel.components.datasets.diffusion.loader import DiffusionDataloaderBuild
+
+logger = logging.getLogger(__name__)
+
+# Collate runs in worker processes, so this only suppresses repeats within a worker --
+# enough to keep the log readable without cross-process coordination.
+_WARNED_KEYS: set = set()
+
+
+class Ideogram4CacheDataset(Dataset):
+    """Reads pre-encoded Ideogram-4 samples ({image_latents, llm_features, text_length})."""
+
+    def __init__(self, cache_dir: str) -> None:
+        self.cache_dir = Path(cache_dir).resolve()
+        meta_path = self.cache_dir / "metadata.json"
+        if not meta_path.exists():
+            raise FileNotFoundError(f"Ideogram-4 cache metadata not found: {meta_path}")
+        with open(meta_path, "r") as f:
+            self.meta = json.load(f)
+        self.samples: List[Dict] = self.meta["samples"]
+        if not self.samples:
+            raise ValueError(f"Ideogram-4 cache is empty: {self.cache_dir}")
+        self.grid_h = int(self.meta.get("grid_h", 0))
+        self.grid_w = int(self.meta.get("grid_w", 0))
+        # Longest caption in the WHOLE cache, from metadata alone (no sample loads). Every
+        # rank reads the same metadata.json, so this is identical across ranks -- which
+        # matters, because a per-rank text width would desync the sharded shapes.
+        self.max_text_length = max(int(s["text_length"]) for s in self.samples)
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        item = self.samples[idx]
+        cache_file = (self.cache_dir / item["cache_file"]).resolve()
+        # Contain path traversal: cache files must live under cache_dir.
+        try:
+            cache_file.relative_to(self.cache_dir)
+        except ValueError as e:  # pragma: no cover
+            raise ValueError(f"cache file {cache_file} outside {self.cache_dir}") from e
+        data = torch.load(cache_file, map_location="cpu", weights_only=True)
+        return {
+            "image_latents": data["image_latents"].to(torch.float32),  # [128, gh, gw]
+            "llm_features": data["llm_features"],  # [n, 53248] (fp16)
+            "text_length": int(data["text_length"]),
+        }
+
+
+def _collate_ideogram4_cache(
+    batch: List[Dict[str, torch.Tensor]], text_width: int = 0
+) -> Dict[str, object]:
+    """Stack latents; LEFT-pad variable-length ``llm_features`` to a fixed width.
+
+    ``text_width`` is the padded text region size. Pass a **dataset-level constant** (the
+    longest caption in the cache) so the packed sequence length ``S = text_width + n_img``
+    is identical on every step: ``torch.compile`` keys its graphs on input shapes, so
+    padding to the per-BATCH maximum instead (``text_width=0``) recompiles whenever a batch
+    happens to contain a different longest caption. That alone blocks compile on real data,
+    regardless of how ``cu_seqlens`` is produced.
+
+    ``text_width=0`` preserves the original per-batch behaviour for callers that do not
+    care about compile.
+
+    Captions longer than ``text_width`` are truncated to their FIRST ``text_width`` tokens
+    and warned about once -- reaching that branch means the width was set too small by hand,
+    since the derived default cannot be exceeded.
+    """
+    image_latents = torch.stack([b["image_latents"] for b in batch], dim=0)  # [B,128,gh,gw]
+
+    feats = [b["llm_features"] for b in batch]
+    dim = feats[0].shape[-1]
+    t_max = int(text_width) if text_width else max(int(f.shape[0]) for f in feats)
+
+    padded = feats[0].new_zeros(len(batch), t_max, dim)
+    text_lengths = torch.empty(len(batch), dtype=torch.long)
+    for i, f in enumerate(feats):
+        n = int(f.shape[0])
+        if n > t_max:
+            if "truncated" not in _WARNED_KEYS:
+                _WARNED_KEYS.add("truncated")
+                logger.warning(
+                    "[Ideogram4Cache] caption with %d tokens exceeds text_width=%d; truncating "
+                    "to the first %d tokens. Raise max_text_tokens to keep full captions.",
+                    n,
+                    t_max,
+                    t_max,
+                )
+            f = f[:t_max]
+            n = t_max
+        padded[i, t_max - n :] = f  # left-pad: real tokens in the LAST n rows
+        text_lengths[i] = n
+
+    return {
+        "image_latents": image_latents,
+        "llm_features": padded,
+        "text_lengths": text_lengths,
+        "data_type": "image",
+    }
+
+
+@dataclass
+class Ideogram4CacheDataloaderConfig:
+    """Construction-time config for the real (pre-encoded) Ideogram-4 dataloader.
+
+    Selected in YAML via::
+
+        data:
+          dataloader:
+            _target_: primus.backends.nemo_automodel.models.ideogram4.data.cache.Ideogram4CacheDataloaderConfig
+            cache_dir: /dataset/pcam_ideogram4_256
+
+    Every field must be a plain YAML scalar; runtime ``dp_rank`` / ``dp_world_size`` /
+    ``batch_size`` are passed to :meth:`build`.
+    """
+
+    cache_dir: str
+    shuffle: bool = True
+    drop_last: bool = True
+    num_workers: int = 2
+    pin_memory: bool = True
+    # Padded width of the text region. 0 (default) derives it from the cache: the longest
+    # caption present, which is a dataset-level constant and so never truncates. Setting it
+    # explicitly caps -- and therefore truncates -- longer captions. Any fixed value keeps
+    # the packed sequence length S constant across steps, which is what torch.compile needs;
+    # -1 restores the old per-batch padding (variable S, recompiles).
+    max_text_tokens: int = 0
+
+    def build(self, *, dp_rank: int, dp_world_size: int, batch_size: int) -> DiffusionDataloaderBuild:
+        dataset = Ideogram4CacheDataset(self.cache_dir)
+
+        if self.max_text_tokens < 0:
+            text_width = 0  # per-batch max (legacy, shape-unstable)
+        elif self.max_text_tokens == 0:
+            text_width = dataset.max_text_length
+        else:
+            text_width = int(self.max_text_tokens)
+
+        sampler = None
+        if dp_world_size > 1:
+            sampler = DistributedSampler(
+                dataset,
+                num_replicas=dp_world_size,
+                rank=dp_rank,
+                shuffle=self.shuffle,
+                drop_last=self.drop_last,
+            )
+
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=(sampler is None and self.shuffle),
+            sampler=sampler,
+            # partial (not a lambda/closure) so the collate stays picklable for the worker
+            # processes when num_workers > 0.
+            collate_fn=partial(_collate_ideogram4_cache, text_width=text_width),
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last,
+        )
+        logger.info(
+            "[Ideogram4Cache] %d samples from %s (grid=%dx%d, text_width=%s, dp_rank=%d/%d, "
+            "bs=%d, %d batches/rank)",
+            len(dataset),
+            self.cache_dir,
+            dataset.grid_h,
+            dataset.grid_w,
+            text_width if text_width else "per-batch",
+            dp_rank,
+            dp_world_size,
+            batch_size,
+            len(dataloader),
+        )
+        return DiffusionDataloaderBuild(dataloader=dataloader, sampler=sampler)
diff --git a/primus/backends/nemo_automodel/models/ideogram4/data/synthetic.py b/primus/backends/nemo_automodel/models/ideogram4/data/synthetic.py
new file mode 100644
index 000000000..23e3927cd
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/data/synthetic.py
@@ -0,0 +1,244 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Synthetic Ideogram-4 diffusion dataloader (no-fork, Primus-side).
+
+WHY (no Automodel/diffusers fork):
+  The AutoModel diffusion recipe resolves ``data.dataloader`` through a *closed*
+  registry keyed by the builder's dotted path
+  (``RecipeConfig.resolve_diffusion_dataloader``: text-to-image / text-to-video /
+  meta-files / mock). Any OTHER ``_target_`` falls back to the resolved target
+  itself and is accepted **iff it is a ``@dataclass`` exposing ``build()``** with
+  the :class:`DiffusionDataloaderConfig` signature. So this module ships exactly
+  that — a dataclass the YAML can point at directly — and needs no submodule edit.
+
+WHAT it produces (the batch the :class:`Ideogram4Adapter` consumes, matching the
+eventual Flux-2 VAE + Qwen3-VL preprocessor cache):
+  - ``image_latents``   ``[B, in_channels, grid_h, grid_w]``  clean packed latents (x0)
+  - ``llm_features``    ``[B, max_text_tokens, llm_features_dim]``  Qwen3-VL-like feats
+  - ``text_lengths``    ``[B]`` int  real (non-pad) text token count per sample
+  - ``data_type``       ``"image"``
+
+  The clean latents are what ``FlowMatchingPipeline.step`` noises (it reads
+  ``image_latents`` as x0, samples sigma, forms ``x_t=(1-σ)x0+σε`` and target
+  ``v=ε-x0`` in this 128-dim packed space); the adapter packs/prepends the
+  text region and builds the ``[pad][text][image]`` position/segment/indicator
+  ids internally, so the loader only emits these raw tensors.
+
+FIXED (deterministic) dataset for an overfit smoke:
+  Each index generates the SAME tensors every epoch (per-index seeded RNG), and
+  each sample gets a DISTINCT ``llm_features`` fingerprint. With fresh noise per
+  step the per-step loss is noisy, but because ``v=(x_t-x0)/σ`` is recoverable
+  from the (fixed, distinct) conditioning the model memorises x0 per sample and
+  the training loss trends clearly DOWN — the signal this smoke checks. Features
+  are scaled by ``feature_scale`` (real tapped hidden states are ~unit-norm
+  per-dim, not a raw 53248-wide unit-variance vector that blows up the input
+  projection).
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+from nemo_automodel.components.datasets.diffusion.loader import DiffusionDataloaderBuild
+
+logger = logging.getLogger(__name__)
+
+
+class SyntheticIdeogram4Dataset(Dataset):
+    """Deterministic random Ideogram-4 samples (image latents + LLM features).
+
+    Per-index seeded so the dataset is FIXED across epochs/workers/ranks, which is
+    what makes the overfit loss-decrease signal meaningful.
+    """
+
+    def __init__(
+        self,
+        *,
+        num_samples: int,
+        in_channels: int,
+        grid_h: int,
+        grid_w: int,
+        max_text_tokens: int,
+        min_text_tokens: int,
+        llm_features_dim: int,
+        feature_scale: float,
+        latent_scale: float,
+        seed: int,
+        cache_in_memory: bool = False,
+        share_text_features: bool = False,
+    ) -> None:
+        self.num_samples = max(int(num_samples), 1)
+        self.in_channels = int(in_channels)
+        self.grid_h = int(grid_h)
+        self.grid_w = int(grid_w)
+        self.max_text_tokens = int(max_text_tokens)
+        self.min_text_tokens = max(1, min(int(min_text_tokens), self.max_text_tokens))
+        self.llm_features_dim = int(llm_features_dim)
+        self.feature_scale = float(feature_scale)
+        self.latent_scale = float(latent_scale)
+        self.seed = int(seed)
+        self.cache_in_memory = bool(cache_in_memory)
+        self.share_text_features = bool(share_text_features)
+        self._cache: Dict[int, Dict[str, torch.Tensor]] = {}
+        self._shared_features: torch.Tensor | None = None
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        # Perf mode: generate each sample once and reuse (so throughput runs measure
+        # the model step, not per-__getitem__ randn of the big 53248-dim features).
+        if self.cache_in_memory and idx in self._cache:
+            return self._cache[idx]
+        # Per-index generator -> identical sample every epoch (fixed dataset).
+        gen = torch.Generator().manual_seed(self.seed + int(idx))
+
+        image_latents = self.latent_scale * torch.randn(
+            self.in_channels, self.grid_h, self.grid_w, dtype=torch.float32, generator=gen
+        )
+        # Full [max_text, dim] features; pad-region positions are masked out by the
+        # adapter's segment_ids, so only the real-text region is attended.
+        if self.share_text_features:
+            # One buffer for every index. At a realistic text width this tensor is large
+            # enough that generating a distinct one per sample dominates the micro-batch on
+            # the CPU side, which makes a throughput run measure the dataloader rather than
+            # the model. Nothing about throughput depends on the samples holding DIFFERENT
+            # noise, only on their shapes and lengths, and ``text_lengths`` below still vary.
+            # Aliasing contract is the same one ``cache_in_memory`` already relies on:
+            # repeat reads hand back the identical tensor object.
+            if self._shared_features is None:
+                shared_gen = torch.Generator().manual_seed(self.seed)
+                self._shared_features = self.feature_scale * torch.randn(
+                    self.max_text_tokens, self.llm_features_dim, dtype=torch.float32, generator=shared_gen
+                )
+            llm_features = self._shared_features
+        else:
+            llm_features = self.feature_scale * torch.randn(
+                self.max_text_tokens, self.llm_features_dim, dtype=torch.float32, generator=gen
+            )
+        # Deterministic per-sample real-text length in [min_text, max_text].
+        span = self.max_text_tokens - self.min_text_tokens + 1
+        text_len = self.min_text_tokens + (int(idx) % span)
+
+        sample = {
+            "image_latents": image_latents,
+            "llm_features": llm_features,
+            "text_lengths": torch.tensor(text_len, dtype=torch.long),
+        }
+        if self.cache_in_memory:
+            self._cache[idx] = sample
+        return sample
+
+
+def _collate_synthetic_ideogram4(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, object]:
+    """Stack per-sample tensors into the FlowMatchingPipeline batch format."""
+    return {
+        "image_latents": torch.stack([b["image_latents"] for b in batch], dim=0),
+        "llm_features": torch.stack([b["llm_features"] for b in batch], dim=0),
+        "text_lengths": torch.stack([b["text_lengths"] for b in batch], dim=0),
+        "data_type": "image",
+    }
+
+
+@dataclass
+class SyntheticIdeogram4DataloaderConfig:
+    """Construction-time config for the synthetic Ideogram-4 dataloader.
+
+    Selected in YAML via::
+
+        data:
+          dataloader:
+            _target_: primus.backends.nemo_automodel.models.ideogram4.data.synthetic.SyntheticIdeogram4DataloaderConfig
+            num_samples: 64
+            grid_h: 16
+            grid_w: 16
+            max_text_tokens: 32
+
+    Every field here must be a plain YAML scalar (the recipe validates the YAML
+    keys against these dataclass fields). The runtime ``dp_rank`` /
+    ``dp_world_size`` / ``batch_size`` are passed to :meth:`build`, not fields.
+    """
+
+    num_samples: int = 64
+    in_channels: int = 128
+    grid_h: int = 16
+    grid_w: int = 16
+    max_text_tokens: int = 32
+    min_text_tokens: int = 28
+    llm_features_dim: int = 53248
+    feature_scale: float = 0.1
+    latent_scale: float = 1.0
+    seed: int = 1234
+    shuffle: bool = True
+    drop_last: bool = False
+    num_workers: int = 0
+    pin_memory: bool = True
+    # Perf: cache generated samples in-process (use with num_workers=0 so the cache
+    # persists across epochs) to isolate model throughput from feature generation.
+    cache_in_memory: bool = False
+    # Perf: reuse ONE generated [max_text_tokens, llm_features_dim] buffer for every
+    # sample. Throughput depends on shapes and text lengths, not on the noise differing
+    # per sample, and at a realistic text width generating it per sample costs more CPU
+    # than the training step costs GPU. Leave false for the overfit smoke, where the
+    # samples must be distinct for the loss-decrease signal to mean anything.
+    share_text_features: bool = False
+
+    def build(self, *, dp_rank: int, dp_world_size: int, batch_size: int) -> DiffusionDataloaderBuild:
+        """Build the synthetic dataset, per-rank sampler, and dataloader."""
+        dataset = SyntheticIdeogram4Dataset(
+            num_samples=self.num_samples,
+            in_channels=self.in_channels,
+            grid_h=self.grid_h,
+            grid_w=self.grid_w,
+            max_text_tokens=self.max_text_tokens,
+            min_text_tokens=self.min_text_tokens,
+            llm_features_dim=self.llm_features_dim,
+            feature_scale=self.feature_scale,
+            latent_scale=self.latent_scale,
+            seed=self.seed,
+            cache_in_memory=self.cache_in_memory,
+            share_text_features=self.share_text_features,
+        )
+
+        sampler = None
+        if dp_world_size > 1:
+            sampler = DistributedSampler(
+                dataset,
+                num_replicas=dp_world_size,
+                rank=dp_rank,
+                shuffle=self.shuffle,
+                drop_last=self.drop_last,
+            )
+
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=(sampler is None and self.shuffle),
+            sampler=sampler,
+            num_workers=self.num_workers,
+            collate_fn=_collate_synthetic_ideogram4,
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last,
+        )
+        logger.info(
+            "[SyntheticIdeogram4] %d samples, in_ch=%d grid=%dx%d max_text=%d feat_dim=%d "
+            "(dp_rank=%d/%d, batch_size=%d, %d batches/rank)",
+            len(dataset),
+            self.in_channels,
+            self.grid_h,
+            self.grid_w,
+            self.max_text_tokens,
+            self.llm_features_dim,
+            dp_rank,
+            dp_world_size,
+            batch_size,
+            len(dataloader),
+        )
+        return DiffusionDataloaderBuild(dataloader=dataloader, sampler=sampler)
diff --git a/primus/backends/nemo_automodel/models/ideogram4/packing_buffer.py b/primus/backends/nemo_automodel/models/ideogram4/packing_buffer.py
new file mode 100644
index 000000000..0f98fada9
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/packing_buffer.py
@@ -0,0 +1,287 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Transport for the Ideogram-4 var-len packing: one shared non-persistent buffer.
+
+WHY (this is the transport that works, after the one that did not):
+  The adapter builds ``cu_seqlens`` on the host, where the caption lengths are already
+  Python ints, and the attention processor has to receive it as a graph INPUT -- deriving it
+  from the mask inside the compiled region does device->host reads, which graph-break and,
+  under FSDP2, desync the per-layer collectives. The obvious channel, diffusers'
+  ``attention_kwargs``, is DEAD for this model in diffusers 0.39.0: the LoRA decorator keeps
+  only ``scale``, both block call sites pass four positional arguments, and
+  ``Ideogram4TransformerBlock.forward`` has no ``**kwargs`` to forward. Nothing arrived.
+
+WHAT (no diffusers fork, nothing copied from upstream):
+  The processor is already handed its ``Ideogram4Attention`` module as the first argument,
+  so the packing rides on that module. One ``int32`` tensor of fixed shape is registered as a
+  NON-PERSISTENT buffer on every attention module -- the same tensor OBJECT on all of them,
+  so the adapter publishes with a single in-place ``copy_`` per step instead of 34 writes,
+  and a diffusers version bump cannot silently diverge from a copied forward.
+
+MEASURED (2026-08-04, torch 2.12+rocm7.14, MI355X; runbook §7.10 has commands and logs):
+  Dynamo lifts the buffer as a graph input rather than baking its values in -- placeholder
+  ``l_self_modules_attention_buffers_primus_cu_seqlens_``, zero recompiles across five steps
+  of changing values, output and grads tracking an uncompiled twin at ~2e-07. FSDP2 leaves it
+  alone -- one object through ``fully_shard``, a plain tensor rather than a DTensor shard,
+  absent from the state dict, ``[1, 1]`` graphs per rank, and a reduce-scattered grad matching
+  the mean-of-per-rank-packings reference at 1.8e-07.
+
+FOUR RULES. Each was measured failing when broken, and only rule 3 fails loudly (runbook §11.7):
+  1. **Publish before the forward, never between forward and backward.** Per-layer compile
+     sits inside the checkpoint wrapper, so the block -- and this buffer read -- runs twice
+     per step, once forward and once recomputed on backward. A write in between corrupts the
+     recomputed activations with no error.
+  2. **Re-publish after materialization.** ``model.to_empty(device=...)``, the usual
+     meta-device bring-up, gives every module its OWN buffer. That would leave layer 0
+     reading the published value and the other 33 reading uninitialized memory.
+     :func:`publish_packing` detects that and restores sharing rather than trusting it.
+  3. **The buffer is read-only once published.** aiter's var-len kernel treats ``cu_seqlens``
+     as a mutable argument: it saves the tensor for its backward and then writes it, bumping
+     the autograd version counter on every call. Passing the SHARED buffer straight in
+     therefore moves its version 34 times per forward while each layer's backward still
+     expects the version it saved, and the step dies in ``.backward()`` (measured 2026-08-04:
+     "IntTensor[5] is at version 35; expected version 34"). The processor clones before the
+     kernel call for that reason -- the probe in §7.10 missed it because it exercised the
+     DENSE path, where no kernel ever sees the packing.
+  4. **The buffer outlives the step.** Anything that runs the transformer WITHOUT going
+     through the adapter -- a sampling pass, an eval loop at a different batch size -- would
+     read whatever was published last. Call :func:`clear_packing` first, which returns the
+     processor to its mask-derived path. The processor also rejects a packing whose length
+     does not match the batch, which is a static-shape comparison and therefore free, so this
+     mistake fails loudly rather than silently.
+"""
+from __future__ import annotations
+
+import logging
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+# The buffer name appears verbatim as the graph placeholder in a TORCH_LOGS dump
+# (``l_self_modules_attention_buffers_primus_cu_seqlens_``), so keep it greppable.
+PACKING_ATTR = "_primus_cu_seqlens"
+BOUND_ATTR = "_primus_max_seqlen"
+_CONSUMERS_ATTR = "_primus_packing_consumers"
+
+_logged: set = set()
+
+
+def _log_once(key: str, level: int, msg: str, *args) -> None:
+    if key not in _logged:
+        _logged.add(key)
+        logger.log(level, msg, *args)
+
+
+def _own_processor(module: nn.Module) -> object:
+    """The processor this module itself holds, ignoring attributes it merely forwards.
+
+    Plain ``getattr`` is wrong here: the activation-checkpoint wrapper delegates unknown
+    attributes to the module it wraps, so it answers ``processor`` with its child's and gets
+    counted as a second consumer. That doubled the module list (68 for 34 blocks), which is
+    harmless for the copy_ but makes the one health number in the log -- "on N attention
+    module(s)" -- unable to distinguish "all 34 blocks" from "half of them".
+    """
+    own = module.__dict__.get("processor")
+    if own is None:
+        # A processor implemented as an nn.Module would live in _modules instead.
+        own = module._modules.get("processor")
+    return own
+
+
+def attention_modules(model: nn.Module) -> List[nn.Module]:
+    """The attention modules that can actually READ the packing.
+
+    Membership is decided by the processor type, which makes the whole transport
+    self-gating: a run without ``PRIMUS_IDEOGRAM_VARLEN_ATTN`` has no var-len processors, so
+    no buffers are installed and nothing is published, rather than maintaining per-step state
+    that nobody reads. Cached on the model -- the module set is fixed once the recipe has
+    built and parallelized it.
+    """
+    cached = getattr(model, _CONSUMERS_ATTR, None)
+    if cached is not None:
+        return cached
+
+    from primus.backends.nemo_automodel.models.ideogram4.attention import (
+        Ideogram4VarlenAttnProcessor,
+    )
+
+    found = [
+        module
+        for module in model.modules()
+        if isinstance(_own_processor(module), Ideogram4VarlenAttnProcessor)
+    ]
+    # Do not cache an empty result: on an odd ordering this could run before the processors
+    # are swapped in, and caching [] would disable the transport for the whole run.
+    if found:
+        setattr(model, _CONSUMERS_ATTR, found)
+    return found
+
+
+def _shared_buffer(modules: List[nn.Module], numel: int, device: torch.device) -> Optional[torch.Tensor]:
+    """The one buffer object every module shares, or None if that is no longer true.
+
+    Returns None for all three ways the invariant can lapse: never installed, installed but
+    the wrong shape/device (a batch-size change), or installed and no longer shared (a
+    materialization pass handed each module its own tensor).
+    """
+    first = getattr(modules[0], PACKING_ATTR, None)
+    if not isinstance(first, torch.Tensor):
+        return None
+    if first.numel() != numel or first.dtype != torch.int32 or first.device != device:
+        return None
+    if not all(getattr(m, PACKING_ATTR, None) is first for m in modules):
+        return None
+    return first
+
+
+def _install(modules: List[nn.Module], numel: int, device: torch.device) -> torch.Tensor:
+    """Point every module at ONE fresh buffer and return it."""
+    shared = torch.zeros(numel, dtype=torch.int32, device=device)
+    for module in modules:
+        # register_buffer refuses to re-register a live name, so replace in ``_buffers``
+        # directly when it is already there. Either way it stays non-persistent.
+        if PACKING_ATTR in module._buffers:
+            module._buffers[PACKING_ATTR] = shared
+        else:
+            module.register_buffer(PACKING_ATTR, shared, persistent=False)
+        if not hasattr(module, BOUND_ATTR):
+            setattr(module, BOUND_ATTR, None)
+    return shared
+
+
+def _report_install(previous: object, numel: int, count: int, max_seqlen: Optional[int]) -> None:
+    if not isinstance(previous, torch.Tensor):
+        # The one line that proves the transport is live in a real run: the entry count is
+        # 2B+1 and the module count should be every block's attention (34 for Ideogram-4).
+        _log_once(
+            "install",
+            logging.INFO,
+            "[PrimusIdeogramPacking] Installed a shared non-persistent cu_seqlens buffer "
+            "(int32, %d entries, max_seqlen=%s) on %d attention module(s); the adapter "
+            "publishes it with one copy_ per step.",
+            numel,
+            max_seqlen,
+            count,
+        )
+    elif previous.numel() != numel:
+        _log_once(
+            f"reshape-{previous.numel()}-{numel}",
+            logging.WARNING,
+            "[PrimusIdeogramPacking] cu_seqlens length changed %d -> %d (batch size changed), so "
+            "the buffer was replaced and torch.compile will recompile. Constant shapes are the "
+            "whole point of the reserved pad column; check drop_last and the eval batch size.",
+            previous.numel(),
+            numel,
+        )
+    else:
+        _log_once(
+            "resharing",
+            logging.WARNING,
+            "[PrimusIdeogramPacking] The cu_seqlens buffer was no longer shared across the %d "
+            "attention module(s) -- something replaced module buffers after install "
+            "(to_empty()/materialization is the usual cause). Sharing has been restored. Left "
+            "unrepaired, one layer would read the published packing and the rest would read "
+            "uninitialized memory, with no error raised.",
+            count,
+        )
+
+
+def publish_packing(
+    model: nn.Module,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: Optional[int] = None,
+    device: Optional[torch.device] = None,
+    required: bool = False,
+) -> Optional[torch.Tensor]:
+    """Make ``cu_seqlens`` visible to every attention module for the forward that follows.
+
+    Call this immediately before the model call and never between a forward and its backward
+    (rule 1 in the module docstring). ``cu_seqlens`` may live on the host: the single
+    ``copy_`` below is then the only host->device transfer, which is the point of building the
+    packing from Python ints in the first place.
+
+    Args:
+        model: the (possibly FSDP2-wrapped, possibly compiled) transformer.
+        cu_seqlens: ``int32 (2B+1,)`` cumulative segment starts, on any device.
+        max_seqlen: the STATIC upper bound ``S``, not the batch's true maximum -- a
+            data-derived int would be guarded by value.
+        device: where the buffer should live. Defaults to ``cu_seqlens.device``, so pass the
+            model's device explicitly when publishing from a host-built tensor.
+        required: raise instead of returning None when nothing consumes the packing. Callers
+            that have already decided the transport is in use should pass True: a rank that
+            publishes into the void falls back to the mask-derived path for every layer, and
+            under data parallelism its gradients then come from a different attention path than
+            its peers' and are averaged in regardless. Nothing in the logs would say so --
+            Primus quiets non-zero ranks once training starts.
+
+    Returns:
+        The shared buffer, or None when nothing in ``model`` consumes it and ``required`` is
+        False.
+    """
+    modules = attention_modules(model)
+    if not modules:
+        if required:
+            raise RuntimeError(
+                "[PrimusIdeogramPacking] a cu_seqlens packing was precomputed but no attention "
+                "module can read it: none of the model's modules carry an "
+                "Ideogram4VarlenAttnProcessor. Every layer would silently fall back to deriving "
+                "the packing from the attention mask, which host-syncs, graph-breaks under "
+                "torch.compile, and -- if this happens on only some ranks -- mixes two "
+                "attention paths into one gradient average. Either the var-len processor install "
+                "did not run (PRIMUS_IDEOGRAM_VARLEN_ATTN) or the model was replaced after it "
+                "did; unset PRIMUS_IDEOGRAM_PRECOMPUTE_CU_SEQLENS to run the legacy path on "
+                "purpose."
+            )
+        return None
+
+    target = torch.device(device) if device is not None else cu_seqlens.device
+    numel = cu_seqlens.numel()
+    shared = _shared_buffer(modules, numel, target)
+    if shared is None:
+        previous = getattr(modules[0], PACKING_ATTR, None)
+        shared = _install(modules, numel, target)
+        _report_install(previous, numel, len(modules), max_seqlen)
+    shared.copy_(cu_seqlens)
+
+    for module in modules:
+        if getattr(module, BOUND_ATTR, None) != max_seqlen:
+            setattr(module, BOUND_ATTR, max_seqlen)
+    return shared
+
+
+def resolve_packing(
+    attn: nn.Module,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+) -> Tuple[Optional[torch.Tensor], Optional[int]]:
+    """Pick the packing for this attention call: explicit argument first, then module state.
+
+    An explicit argument wins so that a future kwargs route (patching the diffusers block
+    forward) can override this without touching the processor. The ``getattr`` fallback is a
+    constant attribute lookup that Dynamo lifts as a graph input -- not a data-dependent read,
+    and not a graph break.
+    """
+    if cu_seqlens is None:
+        cu_seqlens = getattr(attn, PACKING_ATTR, None)
+        if cu_seqlens is not None and max_seqlen is None:
+            max_seqlen = getattr(attn, BOUND_ATTR, None)
+    return cu_seqlens, max_seqlen
+
+
+def clear_packing(model: nn.Module) -> int:
+    """Drop the buffer from every consumer. For tests and for switching transports mid-process."""
+    removed = 0
+    for module in attention_modules(model):
+        if module._buffers.pop(PACKING_ATTR, None) is not None:
+            removed += 1
+        if hasattr(module, BOUND_ATTR):
+            delattr(module, BOUND_ATTR)
+    if hasattr(model, _CONSUMERS_ATTR):
+        delattr(model, _CONSUMERS_ATTR)
+    return removed
diff --git a/primus/backends/nemo_automodel/models/ideogram4/parallelize.py b/primus/backends/nemo_automodel/models/ideogram4/parallelize.py
new file mode 100644
index 000000000..35b133981
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/parallelize.py
@@ -0,0 +1,200 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Real activation checkpointing for Ideogram-4 in the AutoModel diffusion recipe.
+
+WHY:
+  AutoModel picks an FSDP/AC parallelization strategy by the model's class name
+  (``PARALLELIZATION_STRATEGIES`` in
+  ``nemo_automodel.components.distributed.parallelizer``). ``Ideogram4Transformer2DModel``
+  is NOT registered, so it falls back to the default strategy, whose AC path wraps
+  each layer's ``self_attn`` / ``mlp`` submodules -- attributes an Ideogram block
+  (``Ideogram4TransformerBlock`` in ``model.layers``) does not have -- so
+  ``fsdp.activation_checkpointing: true`` is a SILENT NO-OP for Ideogram (same
+  failure mode as FLUX). This registers a strategy that actually checkpoints the
+  Ideogram blocks.
+
+WHAT (NO diffusers / Automodel fork):
+  Registers an Ideogram-specific parallelization strategy via the submodule's own
+  ``register_parallel_strategy`` entry point. It wraps each ``model.layers`` block
+  in a NON-REENTRANT ``checkpoint_wrapper`` (recompute on backward) BEFORE FSDP2
+  sharding, then shards like the in-tree Wan/Hunyuan/FLUX diffusion strategies.
+  Checkpointing only happens when the recipe passes ``activation_checkpointing=True``.
+
+  Env-gated by ``PRIMUS_IDEOGRAM_REAL_AC=1`` (default off), mirroring
+  ``PRIMUS_FLUX_REAL_AC``. Off = default behavior (Ideogram AC is a no-op).
+
+  TP note: targets the dp/FSDP path (tp_size=1); warns once if a TP mesh is given.
+
+Activation (env, no config schema change):
+    PRIMUS_IDEOGRAM_REAL_AC=1   register the real-AC Ideogram parallelization strategy
+"""
+from __future__ import annotations
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+_IDEOGRAM_MODEL_NAME = "Ideogram4Transformer2DModel"
+_IDEOGRAM_BLOCK_ATTR = "layers"
+
+
+def is_ideogram_real_ac_enabled() -> bool:
+    return os.getenv("PRIMUS_IDEOGRAM_REAL_AC", "0") in _TRUTHY
+
+
+def install() -> bool:
+    """Register the real-AC Ideogram parallelization strategy.
+
+    No-op (returns False) unless ``PRIMUS_IDEOGRAM_REAL_AC`` is set. Idempotent
+    (returns True if already present). Modifies NO Automodel source.
+    """
+    if not is_ideogram_real_ac_enabled():
+        return False
+
+    import torch
+
+    import nemo_automodel.components.distributed.parallelizer as P
+
+    if _IDEOGRAM_MODEL_NAME in P.PARALLELIZATION_STRATEGIES:
+        return True  # already registered (idempotent)
+
+    class Ideogram4ParallelizationStrategy(P.ParallelizationStrategy):
+        """FSDP2 + real activation checkpointing for the Ideogram-4 single-stream DiT."""
+
+        def parallelize(
+            self,
+            model,
+            device_mesh,
+            mp_policy=None,
+            offload_policy=None,
+            sequence_parallel: bool = False,
+            activation_checkpointing: bool = False,
+            tp_shard_plan=None,
+            dp_replicate_mesh_name: str = "dp_replicate",
+            dp_shard_cp_mesh_name: str = "dp_shard_cp",
+            tp_mesh_name: str = "tp",
+            reshard_after_forward: bool | None = None,
+            **kwargs,
+        ):
+            dp_mesh = P.get_fsdp_dp_mesh(device_mesh, dp_replicate_mesh_name, dp_shard_cp_mesh_name)
+
+            # Read the hook's state through the module, never a from-import: a
+            # from-import binds False at import time and the guard then fires on
+            # every healthy run.
+            from primus.backends.nemo_automodel.distributed import fsdp2_reshard
+
+            logger.info(
+                "[PrimusIdeogramAC] reshard_after_forward: received=%s applied_by_hook=%s",
+                reshard_after_forward,
+                fsdp2_reshard.applied_reshard_after_forward,
+            )
+            if not fsdp2_reshard.patch_installed:
+                # The trainer swallows hook exceptions, so a failed install shows up
+                # only as one log line and then a full run at ~3P traffic. This is the
+                # one unambiguous in-process signal that it happened.
+                logger.error(
+                    "[PrimusIdeogramAC] the FSDP2 reshard repair hook is NOT installed, so "
+                    "fsdp.reshard_after_forward from YAML was dropped upstream and this run "
+                    "will use ZeRO-3-style per-block resharding regardless of config."
+                )
+
+            if tp_mesh_name in getattr(device_mesh, "mesh_dim_names", ()) and device_mesh[tp_mesh_name].size() > 1:
+                logger.warning(
+                    "[PrimusIdeogramAC] tensor parallelism requested but the Ideogram "
+                    "real-AC strategy has no TP plan; proceeding with FSDP only."
+                )
+
+            # Normalize false-like strings: some CLI/config paths forward the flag as a
+            # raw string, and a non-empty "false" would otherwise be truthy and wrongly
+            # enable AC. "true"/"full"/"selective" keep their meaning (see below).
+            ac_value = activation_checkpointing
+            if isinstance(ac_value, str) and ac_value.strip().lower() in {"false", "0", "off", "no", "none", ""}:
+                ac_value = False
+
+            if ac_value:
+                blocks = getattr(model, _IDEOGRAM_BLOCK_ATTR, None)
+                if blocks is None:
+                    logger.warning(
+                        "[PrimusIdeogramAC] activation_checkpointing requested but model has no "
+                        "'%s' block list; nothing checkpointed.",
+                        _IDEOGRAM_BLOCK_ATTR,
+                    )
+                elif P.is_selective_activation_checkpointing(ac_value):
+                    # PARTIAL / selective (TorchTitan-style, op-level) AC: save attention +
+                    # half the matmuls + comm collectives, recompute only the cheap ops -> far
+                    # less backward recompute than full AC at moderate extra memory. This reuses
+                    # the SHARED Automodel selective-AC machinery (the same lever used for FLUX/WAN).
+                    # The wrapper is tagged with SELECTIVE_AC_WRAPPER_FLAG so per-layer compile
+                    # compiles it OUTER (SAC INNER) and the partitioner honors the recompute tags
+                    # (see _apply_per_layer_compile). has_kv_sharing=False: Ideogram is a diffusion
+                    # DiT with no KV cache. Blocks are replaced in-place in `model` (and `layers`).
+                    layers = list(blocks)
+                    P.apply_selective_checkpointing_to_layers(
+                        model,
+                        layers,
+                        False,
+                        enable_compile=bool(kwargs.get("enable_compile", False)),
+                    )
+                    logger.info(
+                        "[PrimusIdeogramAC] wrapped %d Ideogram blocks with SELECTIVE (partial) "
+                        "activation checkpointing",
+                        len(layers),
+                    )
+                else:
+                    # FULL AC: recompute the whole block on backward (max memory saved, at the
+                    # cost of a recompute tax). Default for activation_checkpointing True / "full".
+                    wrapped = 0
+                    for idx in range(len(blocks)):
+                        blocks[idx] = P.checkpoint_wrapper(
+                            blocks[idx],
+                            checkpoint_impl=P.CheckpointImpl.NO_REENTRANT,
+                        )
+                        wrapped += 1
+                    logger.info(
+                        "[PrimusIdeogramAC] wrapped %d Ideogram blocks with FULL activation checkpointing",
+                        wrapped,
+                    )
+
+            if not mp_policy:
+                mp_policy = P.MixedPrecisionPolicy(
+                    param_dtype=torch.bfloat16,
+                    reduce_dtype=torch.float32,
+                    output_dtype=torch.float32,
+                )
+
+            # Keyword form throughout: this call previously passed seven positional
+            # arguments and stopped one short of reshard_after_forward, silently
+            # discarding it. Keywords make that class of bug impossible to repeat.
+            P.apply_fsdp2_sharding_recursively(
+                model,
+                dp_mesh,
+                mp_policy,
+                offload_policy,
+                enable_fsdp2_prefetch=kwargs.get("enable_fsdp2_prefetch", True),
+                fsdp2_backward_prefetch_depth=kwargs.get("fsdp2_backward_prefetch_depth", 2),
+                fsdp2_forward_prefetch_depth=kwargs.get("fsdp2_forward_prefetch_depth", 1),
+                # Forwarded faithfully, including None: YAML decides, not this code.
+                reshard_after_forward=reshard_after_forward,
+            )
+
+            return P.fully_shard(
+                model,
+                mesh=dp_mesh,
+                mp_policy=mp_policy,
+                offload_policy=offload_policy,
+                reshard_after_forward=False,
+            )
+
+    P.register_parallel_strategy(name=_IDEOGRAM_MODEL_NAME)(Ideogram4ParallelizationStrategy)
+    logger.info(
+        "[PrimusIdeogramAC] Registered real activation-checkpointing parallelization "
+        "strategy for %s (honors fsdp.activation_checkpointing).",
+        _IDEOGRAM_MODEL_NAME,
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/processor.py b/primus/backends/nemo_automodel/models/ideogram4/processor.py
new file mode 100644
index 000000000..021354c32
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/processor.py
@@ -0,0 +1,255 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Ideogram-4 preprocessing (no-fork, Primus-side): Flux-2 VAE image encode +
+Qwen3-VL 13-layer text tap -> the cache the :class:`Ideogram4Adapter` consumes.
+
+WHY (no Automodel/diffusers fork):
+  AutoModel's processors live in the submodule (``tools/diffusion/processors``) and
+  register into a closed ``ProcessorRegistry``. Rather than edit the submodule, this
+  is a standalone Primus-side processor that produces a flat per-sample cache read
+  by :class:`Ideogram4CacheDataloaderConfig` (``data/cache.py``).
+
+WHAT it reproduces (ported verbatim from the diffusers Ideogram-4 pipeline so the
+cached conditioning matches training==inference):
+  - IMAGE: ``AutoencoderKLFlux2`` encode -> 32-ch latent -> 2x2 patch pack to
+    ``[128, gh, gw]`` (128 = ae*patch*patch, channel order ``(pa, pb, ae)`` matching
+    the pipeline's decode ``view(B, gh, gw, patch, patch, ae)``), then the VAE
+    ``bn`` normalisation on the packed-128 channels. This is exactly the space the
+    ``FlowMatchingPipeline`` adds noise in.
+  - TEXT: chat-template tokenise -> run the Qwen3-VL ``language_model`` decoder
+    stack -> tap hidden states at ``QWEN3_VL_ACTIVATION_LAYERS`` (0,3,...,33,35) ->
+    concat to ``[n, 53248]`` (= 4096 x 13, layer fastest). Stored non-padded (only
+    the ``n`` real tokens); the cache loader left-pads per batch to match the
+    adapter's ``[left-pad][text][image]`` layout.
+
+The nf4 Ideogram text encoder needs bitsandbytes (absent from the pinned ROCm
+container), so we tap the full-precision base ``Qwen/Qwen3-VL-8B-Instruct`` (same
+architecture). The VAE is loaded from the (non-quantized) Ideogram ``vae/`` subfolder
+so the bn stats are exact.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+# Qwen3-VL decoder layers tapped for the packed text conditioning (from the diffusers
+# Ideogram-4 pipeline: QWEN3_VL_ACTIVATION_LAYERS).
+QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+
+# Default HF sources (see scripts/b3_download_weights.py).
+DEFAULT_VAE_SOURCE = "ideogram-ai/ideogram-4-nf4-diffusers"
+DEFAULT_TEXT_ENCODER_SOURCE = "Qwen/Qwen3-VL-8B-Instruct"
+
+
+def _load_autoencoder_kl_flux2():
+    try:
+        from diffusers import AutoencoderKLFlux2  # type: ignore
+
+        return AutoencoderKLFlux2
+    except Exception:  # pragma: no cover
+        from diffusers.models.autoencoders import AutoencoderKLFlux2  # type: ignore
+
+        return AutoencoderKLFlux2
+
+
+class Ideogram4Processor:
+    """Encode (image, caption) pairs into the Ideogram-4 training cache format."""
+
+    def __init__(
+        self,
+        *,
+        vae_source: str = DEFAULT_VAE_SOURCE,
+        vae_subfolder: str = "vae",
+        text_encoder_source: str = DEFAULT_TEXT_ENCODER_SOURCE,
+        tokenizer_source: Optional[str] = None,
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        self.vae_source = vae_source
+        self.vae_subfolder = vae_subfolder
+        self.text_encoder_source = text_encoder_source
+        self.tokenizer_source = tokenizer_source or text_encoder_source
+        self.device = device
+        self.dtype = dtype
+
+        self.vae = None
+        self.patch_size = 2
+        self.bn_mean = None  # [1, 1, 1, 128] float32
+        self.bn_std = None  # [1, 1, 1, 128] float32
+        self.text_encoder = None
+        self.language_model = None
+        self.tokenizer = None
+
+    # ------------------------------------------------------------------ load
+    def load_models(self) -> None:
+        from transformers import AutoTokenizer, Qwen3VLModel
+
+        AutoencoderKLFlux2 = _load_autoencoder_kl_flux2()
+
+        logger.info("[Ideogram4Processor] loading VAE (%s/%s)", self.vae_source, self.vae_subfolder)
+        vae = AutoencoderKLFlux2.from_pretrained(
+            self.vae_source, subfolder=self.vae_subfolder, torch_dtype=self.dtype
+        )
+        vae = vae.to(self.device).eval()
+        self.vae = vae
+
+        # BN running stats on the packed-128 channel space (see pipeline decode).
+        eps = float(getattr(vae.config, "batch_norm_eps", 1e-5))
+        self.bn_mean = vae.bn.running_mean.detach().view(1, 1, 1, -1).float().to(self.device)
+        self.bn_std = torch.sqrt(vae.bn.running_var.detach() + eps).view(1, 1, 1, -1).float().to(self.device)
+        packed_ch = int(self.bn_mean.shape[-1])
+        ae = int(getattr(vae.config, "latent_channels", packed_ch // 4))
+        self.patch_size = int(round((packed_ch / ae) ** 0.5))
+        assert ae * self.patch_size * self.patch_size == packed_ch, (
+            f"packed channels {packed_ch} != ae({ae}) * patch^2({self.patch_size}^2)"
+        )
+        logger.info(
+            "[Ideogram4Processor] VAE ready: ae=%d patch=%d packed=%d", ae, self.patch_size, packed_ch
+        )
+
+        logger.info("[Ideogram4Processor] loading text encoder (%s)", self.text_encoder_source)
+        te = Qwen3VLModel.from_pretrained(self.text_encoder_source, torch_dtype=self.dtype)
+        # We only tap the language model; drop the vision tower to save memory if present.
+        if hasattr(te, "visual"):
+            try:
+                te.visual = None
+            except Exception:  # pragma: no cover
+                pass
+        te = te.to(self.device).eval()
+        self.text_encoder = te
+        self.language_model = te.language_model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_source)
+        logger.info("[Ideogram4Processor] models loaded.")
+
+    # ------------------------------------------------------------------ image
+    @staticmethod
+    def preprocess_image(image: Image.Image, resolution: int) -> torch.Tensor:
+        """PIL RGB image -> resize+center-crop to (resolution, resolution), [-1,1], [1,3,H,W]."""
+        image = image.convert("RGB")
+        w, h = image.size
+        scale = max(resolution / w, resolution / h)
+        rw, rh = int(round(w * scale)), int(round(h * scale))
+        image = image.resize((rw, rh), Image.LANCZOS)
+        left = (rw - resolution) // 2
+        top = (rh - resolution) // 2
+        image = image.crop((left, top, left + resolution, top + resolution))
+        arr = torch.from_numpy(np.array(image)).float() / 255.0  # [H,W,3]
+        arr = (arr - 0.5) / 0.5
+        return arr.permute(2, 0, 1).unsqueeze(0)  # [1,3,H,W]
+
+    @torch.no_grad()
+    def encode_image(self, image_tensor: torch.Tensor) -> torch.Tensor:
+        """Image [1,3,H,W] in [-1,1] -> packed+BN latent [128, gh, gw] (fp16).
+
+        Packing matches the diffusers Ideogram-4 decode inverse exactly:
+        raw[ae, 2i+pa, 2j+pb] -> channel ((pa*patch)+pb)*ae + ae_idx of token (i,j).
+        """
+        vae = self.vae
+        patch = self.patch_size
+        x = image_tensor.to(self.device, dtype=self.dtype)
+        raw = vae.encode(x).latent_dist.mode()  # [1, ae, H8, W8]
+        raw = raw.float()
+        b, ae, H, W = raw.shape
+        gh, gw = H // patch, W // patch
+        x6 = raw.view(b, ae, gh, patch, gw, patch)
+        # -> [b, gh, gw, pa, pb, ae] -> [b, gh, gw, 128] (channel order (pa, pb, ae))
+        packed = x6.permute(0, 2, 4, 3, 5, 1).reshape(b, gh, gw, patch * patch * ae)
+        packed = (packed - self.bn_mean) / self.bn_std
+        latents = packed.permute(0, 3, 1, 2).squeeze(0)  # [128, gh, gw]
+        return latents.detach().cpu().to(torch.float16)
+
+    # ------------------------------------------------------------------ text
+    @torch.no_grad()
+    def _tap_hidden_states(
+        self, token_ids: torch.Tensor, attention_mask: torch.Tensor, pos_1d: torch.Tensor
+    ) -> List[torch.Tensor]:
+        """Run the Qwen3-VL decoder stack, returning hidden states at each tap layer.
+
+        Ported from Ideogram4Pipeline._get_text_encoder_hidden_states.
+        """
+        from transformers.masking_utils import create_causal_mask
+
+        lm = self.language_model
+        inputs_embeds = lm.embed_tokens(token_ids)
+
+        position_ids_4d = pos_1d[None, ...].expand(4, pos_1d.shape[0], -1)
+        text_position_ids = position_ids_4d[0]
+        mrope_position_ids = position_ids_4d[1:]
+
+        causal_mask = create_causal_mask(
+            config=lm.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=None,
+            position_ids=text_position_ids,
+        )
+        position_embeddings = lm.rotary_emb(inputs_embeds, mrope_position_ids)
+
+        tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
+        captured: Dict[int, torch.Tensor] = {}
+        hidden_states = inputs_embeds
+        for layer_idx, decoder_layer in enumerate(lm.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=text_position_ids,
+                past_key_values=None,
+                position_embeddings=position_embeddings,
+            )
+            if layer_idx in tap_set:
+                captured[layer_idx] = hidden_states
+        return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
+
+    @torch.no_grad()
+    def encode_text(self, prompt: str, max_text_tokens: int) -> Optional[Dict[str, Any]]:
+        """Caption -> {llm_features [n, 53248] fp16, text_length n}. Returns None if too long.
+
+        No padding: a single sequence of exactly ``n`` real tokens is tapped, so the
+        cached features are the non-pad rows the adapter/cache-loader left-pad later.
+        """
+        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+        text = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        toks = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+        n = int(toks.shape[0])
+        if n == 0 or n > max_text_tokens:
+            return None
+
+        token_ids = toks.view(1, n).to(self.device)
+        attention_mask = torch.ones(1, n, dtype=torch.long, device=self.device)
+        pos_1d = torch.arange(n, device=self.device).view(1, n)
+
+        selected = self._tap_hidden_states(token_ids, attention_mask, pos_1d)
+        # [13, 1, n, 4096] -> [1, n, 4096, 13] -> [1, n, 53248]  (layer fastest)
+        feats = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(1, n, -1)
+        feats = feats.squeeze(0).to(torch.float16).detach().cpu()  # [n, 53248]
+        return {"llm_features": feats, "text_length": n}
+
+    # ------------------------------------------------------------------ cache
+    @staticmethod
+    def get_cache_data(
+        image_latents: torch.Tensor,
+        text_enc: Dict[str, Any],
+        *,
+        prompt: str,
+        image_path: str,
+    ) -> Dict[str, Any]:
+        gh, gw = int(image_latents.shape[1]), int(image_latents.shape[2])
+        return {
+            "image_latents": image_latents,  # [128, gh, gw] fp16
+            "llm_features": text_enc["llm_features"],  # [n, 53248] fp16
+            "text_length": int(text_enc["text_length"]),
+            "grid_h": gh,
+            "grid_w": gw,
+            "prompt": prompt,
+            "image_path": image_path,
+            "model_type": "ideogram4",
+        }
diff --git a/primus/backends/nemo_automodel/models/ideogram4/profile.py b/primus/backends/nemo_automodel/models/ideogram4/profile.py
new file mode 100644
index 000000000..09e49eb89
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/profile.py
@@ -0,0 +1,188 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Env-gated ``torch.profiler`` wrapper for the AutoModel diffusion train loop —
+the no-fork vehicle for a real multi-GPU FSDP2 comms-share profile.
+
+WHY:
+  A single-GPU profile isolates the intra-step COMPUTE op-share. To size the FSDP
+  collective (all-gather + reduce-scatter) share — the top remaining lever once
+  throughput saturates at a low mbs knee — we must profile the ACTUAL distributed
+  step. This hook captures a per-rank ``torch.profiler`` Chrome trace of a few
+  steady steps of the real ``TrainDiffusionRecipe`` loop, which TraceLens turns
+  into a gpu-timeline comm/compute split + a per-collective (NcclAnalyser) report.
+
+WHAT (NO Automodel/diffusers fork):
+  ``install()`` class-patches ``TrainDiffusionRecipe.run_train_validation_loop`` so
+  that, when ``PRIMUS_IDEOGRAM_PROFILE=1``, it runs the ORIGINAL loop inside a
+  ``torch.profiler.profile`` context. The per-step boundary needed to drive
+  ``prof.step()`` (for the wait/warmup/active schedule) is obtained by wrapping the
+  recipe's ``self.optimizer.step`` (called exactly once per optimization step),
+  with NO edit to the submodule loop body. Each rank exports its own trace so the
+  multi-rank collective report can correlate collectives across ranks.
+
+  Env-gated by ``PRIMUS_IDEOGRAM_PROFILE=1`` (default off). Installed via the
+  trainer's optional-hooks list. Idempotent; a default run is unaffected.
+
+Activation / knobs (env):
+    PRIMUS_IDEOGRAM_PROFILE=1        enable the profiler wrapper
+    PRIMUS_PROFILE_DIR=<dir>         per-rank trace output dir (default
+                                     ./output/ideogram_profile/comms)
+    PRIMUS_PROFILE_TAG=<tag>         subdir/prefix for this point (e.g. 1024_m32)
+    PRIMUS_PROFILE_WAIT=3            steps skipped before profiling (skip warmup/
+                                     epoch-boundary data tax)
+    PRIMUS_PROFILE_WARMUP=1          profiler warmup steps (recorded then dropped)
+    PRIMUS_PROFILE_ACTIVE=3          steps actually captured
+    PRIMUS_PROFILE_WITH_STACK=0      include python call stacks (bigger trace;
+                                     needed for nn.Module attribution / recompute)
+    PRIMUS_PROFILE_WITH_MODULES=0    include the nn.Module hierarchy per op
+    PRIMUS_PROFILE_RECORD_SHAPES=1   record op input shapes (GEMM M/N/K + coll size)
+    PRIMUS_PROFILE_MEMORY=0          record allocator events
+
+Note on ``WITH_STACK``/``WITH_MODULES``: these are what let a reader say *which part of
+the model* a kernel belongs to. Without them a trace still answers "what ran and on what
+shapes" but not "who called it", which is usually the question being asked of a trace.
+They are off by default because both inflate trace size and CPU-side step time; turn them
+on deliberately when the trace is for someone to read, not for timing.
+"""
+from __future__ import annotations
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+
+def is_profile_enabled() -> bool:
+    return os.getenv("PRIMUS_IDEOGRAM_PROFILE", "0") in _TRUTHY
+
+
+def _flag(name: str, default: str = "0") -> bool:
+    return os.getenv(name, default) in _TRUTHY
+
+
+def _int(name: str, default: int) -> int:
+    try:
+        return int(os.getenv(name, str(default)))
+    except (TypeError, ValueError):
+        return default
+
+
+def _rank() -> int:
+    for key in ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK"):
+        v = os.getenv(key)
+        if v is not None:
+            try:
+                return int(v)
+            except ValueError:
+                pass
+    try:
+        import torch.distributed as dist
+
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_rank()
+    except Exception:
+        pass
+    return 0
+
+
+def install() -> bool:
+    """Wrap the AutoModel diffusion train loop in ``torch.profiler`` (no-fork).
+
+    No-op (returns False) unless ``PRIMUS_IDEOGRAM_PROFILE`` is set.
+    """
+    if not is_profile_enabled():
+        return False
+
+    import torch
+
+    from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe
+
+    if getattr(TrainDiffusionRecipe, "_primus_profile_installed", False):
+        return True
+
+    orig_loop = TrainDiffusionRecipe.run_train_validation_loop
+
+    out_dir = os.getenv("PRIMUS_PROFILE_DIR", "./output/ideogram_profile/comms")
+    tag = os.getenv("PRIMUS_PROFILE_TAG", "run")
+    wait = _int("PRIMUS_PROFILE_WAIT", 3)
+    warmup = _int("PRIMUS_PROFILE_WARMUP", 1)
+    active = _int("PRIMUS_PROFILE_ACTIVE", 3)
+    with_stack = _flag("PRIMUS_PROFILE_WITH_STACK", "0")
+    with_modules = _flag("PRIMUS_PROFILE_WITH_MODULES", "0")
+    record_shapes = _flag("PRIMUS_PROFILE_RECORD_SHAPES", "1")
+    profile_memory = _flag("PRIMUS_PROFILE_MEMORY", "0")
+
+    def patched_loop(self):
+        rank = _rank()
+        point_dir = os.path.join(out_dir, tag)
+        os.makedirs(point_dir, exist_ok=True)
+        trace_path = os.path.join(point_dir, f"rank{rank}.pt.trace.json")
+
+        activities = [
+            torch.profiler.ProfilerActivity.CPU,
+            torch.profiler.ProfilerActivity.CUDA,
+        ]
+        sched = torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1)
+
+        def _on_ready(prof):
+            try:
+                prof.export_chrome_trace(trace_path)
+                logger.info("[PrimusIdeogramProfile] rank%s wrote trace -> %s", rank, trace_path)
+            except Exception as exc:  # pragma: no cover
+                logger.error("[PrimusIdeogramProfile] rank%s trace export failed: %s", rank, exc)
+
+        prof = torch.profiler.profile(
+            activities=activities,
+            schedule=sched,
+            on_trace_ready=_on_ready,
+            record_shapes=record_shapes,
+            with_stack=with_stack,
+            with_modules=with_modules,
+            profile_memory=profile_memory,
+        )
+
+        # Drive prof.step() from the recipe's per-step optimizer.step (called once
+        # per optimization step) without editing the submodule loop body.
+        #
+        # ``self.optimizer`` is a LIST: OptimizerConfig.build returns one optimizer per
+        # model part, and the recipe steps each in turn. Hook only the LAST one, so the
+        # profiler schedule advances exactly once per optimization step -- hooking all of
+        # them would advance it once per model part and capture the wrong steps.
+        optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer]
+        if not optimizers:
+            logger.error("[PrimusIdeogramProfile] recipe has no optimizer to hook; not profiling.")
+            return orig_loop(self)
+        target_optimizer = optimizers[-1]
+        orig_opt_step = target_optimizer.step
+
+        def opt_step_and_prof(*args, **kwargs):
+            ret = orig_opt_step(*args, **kwargs)
+            prof.step()
+            return ret
+
+        logger.info(
+            "[PrimusIdeogramProfile] rank%s profiling '%s' (wait=%d warmup=%d active=%d "
+            "with_stack=%s with_modules=%s record_shapes=%s) -> %s",
+            rank, tag, wait, warmup, active, with_stack, with_modules, record_shapes, point_dir,
+        )
+        target_optimizer.step = opt_step_and_prof
+        prof.start()
+        try:
+            orig_loop(self)
+        finally:
+            try:
+                prof.stop()
+            except Exception:
+                pass
+            target_optimizer.step = orig_opt_step
+
+    patched_loop._primus_profile_wrapped = True
+    TrainDiffusionRecipe.run_train_validation_loop = patched_loop
+    TrainDiffusionRecipe._primus_profile_installed = True
+    logger.info("[PrimusIdeogramProfile] Installed torch.profiler train-loop wrapper (no-fork).")
+    return True
diff --git a/primus/backends/nemo_automodel/models/ideogram4/zero1.py b/primus/backends/nemo_automodel/models/ideogram4/zero1.py
new file mode 100644
index 000000000..777c2630a
--- /dev/null
+++ b/primus/backends/nemo_automodel/models/ideogram4/zero1.py
@@ -0,0 +1,241 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""DDP + ZeRO-1 (distributed optimizer) for Ideogram-4 in the AutoModel diffusion recipe.
+
+WHY:
+  A large-HBM MI355X leaves memory headroom at the compute-bound knee mbs. One way
+  to spend it is to shard *less* and cut FSDP collectives. Pure **DDP** replicates
+  params+grads (no per-layer all-gather) but also replicates the AdamW optimizer
+  state (fp32 m/v ~= 2x params); **ZeRO-1** shards *only* that optimizer state
+  across the DP ranks so we get DDP-like compute with the optimizer memory back.
+
+  IMPORTANT single-node caveat: exposed FSDP comm is a small fraction of the step
+  on one 8-GPU node, AND per-layer torch.compile is wired ONLY on the FSDP2 path
+  (fsdp2.py -> _apply_per_layer_compile; DDPManager never compiles). So on a single
+  node DDP+ZeRO-1 (no compile) trails FSDP2+compile and cannot recover the gap from
+  comms. This path is the **multi-node** lever (where all-gather crosses the slower
+  inter-node network) and a reference point; the single-node sharding lever that
+  KEEPS compile is HSDP (fsdp.dp_replicate_size>1). This module exists to (a)
+  validate the DDP+ZeRO-1 path runs with the Ideogram adapter + hd256 + real AC,
+  and (b) measure the ZeRO-1 optimizer-memory recovery vs pure DDP.
+
+WHAT (NO diffusers / Automodel fork) — two env-gated monkeypatches:
+  1. ZeRO-1 optimizer: wrap the optimizers the recipe builds
+     (nemo_automodel.components.optim.optimizer.OptimizerConfig.build, which returns
+     one per model part) in torch.distributed.optim.ZeroRedundancyOptimizer. Skipped
+     (with a warning) if the params are DTensors (i.e. FSDP2 mode) since ZeRO-1 only
+     applies to replicated/DDP params.
+  2. DDP real activation checkpointing: DDPManager applies AC via
+     apply_submodule_checkpointing, which wraps by attribute name
+     (mlp/self_attn/norm1...) -- attributes an Ideogram4TransformerBlock does NOT
+     have -- so DDP-mode AC is a SILENT NO-OP for Ideogram (same failure mode the
+     FSDP strategy has, fixed there by parallelize.py). Without real AC the 9.3B
+     DiT OOMs at any real mbs. We patch DDPManager.parallelize to whole-block
+     checkpoint-wrap model.layers for Ideogram before DDP wrap.
+
+Activation (env, no config schema change):
+    PRIMUS_IDEOGRAM_ZERO1=1   install BOTH patches (ZeRO-1 optimizer + DDP whole-block AC)
+"""
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+_IDEOGRAM_MODEL_NAME = "Ideogram4Transformer2DModel"
+_IDEOGRAM_BLOCK_ATTR = "layers"
+
+
+def is_zero1_enabled() -> bool:
+    return os.getenv("PRIMUS_IDEOGRAM_ZERO1", "0") in _TRUTHY
+
+
+def is_ddp_enabled() -> bool:
+    """Pure-DDP baseline (whole-block AC on the DDP path, no ZeRO-1)."""
+    return os.getenv("PRIMUS_IDEOGRAM_DDP", "0") in _TRUTHY
+
+
+def _looks_like_ideogram(model) -> bool:
+    return type(model).__name__ == _IDEOGRAM_MODEL_NAME and hasattr(model, _IDEOGRAM_BLOCK_ATTR)
+
+
+def _params_are_dtensor(params) -> bool:
+    for p in params:
+        if type(p).__name__ == "DTensor" or hasattr(p, "_local_tensor"):
+            return True
+    return False
+
+
+def _wrap_in_zero1(base):
+    """Return ``base`` re-built as a ZeroRedundancyOptimizer, or ``base`` unchanged."""
+    if type(base).__name__ == "ZeroRedundancyOptimizer":
+        return base  # a subclass build() that chains to super() would otherwise double-wrap
+    params = [p for group in base.param_groups for p in group["params"]]
+    if _params_are_dtensor(params):
+        logger.warning(
+            "[PrimusIdeogramZeRO1] PRIMUS_IDEOGRAM_ZERO1 set but params are DTensor (FSDP2 mode); "
+            "ZeRO-1 only applies to replicated/DDP params -> keeping the plain (FSDP-sharded) optimizer."
+        )
+        return base
+
+    import torch.distributed as dist
+    from torch.distributed.optim import ZeroRedundancyOptimizer
+
+    if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() == 1:
+        logger.warning("[PrimusIdeogramZeRO1] no >1-rank process group; keeping the plain optimizer.")
+        return base
+
+    optimizer_cls = type(base)
+    # base.defaults holds the per-group defaults, but it can contain keys that are
+    # NOT accepted by the optimizer's __init__ (e.g. AdamW stores decoupled_weight_decay
+    # -- set internally by its Adam parent -- yet AdamW.__init__ doesn't take it). ZeRO
+    # forwards these to reconstruct per-rank optimizers, so filter to the ctor signature.
+    defaults = dict(base.defaults)
+    learning_rate = defaults.pop("lr", None)
+    try:
+        _sig = inspect.signature(optimizer_cls.__init__)
+        _has_varkw = any(p.kind == p.VAR_KEYWORD for p in _sig.parameters.values())
+        if not _has_varkw:
+            _accepted = set(_sig.parameters)
+            _dropped = [k for k in defaults if k not in _accepted]
+            for k in _dropped:
+                defaults.pop(k, None)
+            if _dropped:
+                logger.info(
+                    "[PrimusIdeogramZeRO1] dropped non-ctor optimizer defaults %s for %s",
+                    _dropped,
+                    optimizer_cls.__name__,
+                )
+    except (ValueError, TypeError):
+        pass  # signature unavailable (C-impl); pass defaults through unchanged
+    try:
+        zro = ZeroRedundancyOptimizer(
+            params,
+            optimizer_class=optimizer_cls,
+            lr=learning_rate,
+            overlap_with_ddp=False,
+            **defaults,
+        )
+    except Exception as e:  # fall back rather than break the run
+        logger.error("[PrimusIdeogramZeRO1] ZeroRedundancyOptimizer build failed (%s); using plain optimizer.", e)
+        return base
+    logger.info(
+        "[PrimusIdeogramZeRO1] wrapped %s in ZeroRedundancyOptimizer over %d ranks "
+        "(optimizer state sharded; params/grads replicated by DDP).",
+        optimizer_cls.__name__,
+        dist.get_world_size(),
+    )
+    return zro
+
+
+def _optimizer_config_classes():
+    """Every class in the OptimizerConfig hierarchy that defines its OWN ``build``.
+
+    Patching only the base class is NOT enough, and the way it fails is silent. YAML
+    ``_target_: torch.optim.AdamW`` resolves to a plain torch class, not an
+    OptimizerConfig subclass, so build_optimizer_config() wraps it in
+    ``OptimizerFromFactoryConfig`` -- which overrides ``build`` and never chains to
+    ``super()``. A base-class-only patch is then simply never called: no ZeRO-1, no
+    warning, and a run that looks perfect while the optimizer state stays replicated.
+    ``_DionConfigBase`` overrides ``build`` the same way. Walk the hierarchy instead of
+    naming those two, so a future override cannot re-open the same hole.
+    """
+    from nemo_automodel.components.optim.optimizer import OptimizerConfig
+
+    seen, stack, out = set(), [OptimizerConfig], []
+    while stack:
+        cls = stack.pop()
+        if id(cls) in seen:
+            continue
+        seen.add(id(cls))
+        if "build" in vars(cls):
+            out.append(cls)
+        stack.extend(cls.__subclasses__())
+    return out
+
+
+def _install_zero1_optimizer_patch() -> bool:
+    """Wrap the recipe's optimizers in ZeroRedundancyOptimizer."""
+    patched = []
+    for cls in _optimizer_config_classes():
+        if getattr(vars(cls)["build"], "_primus_zero1_patched", False):
+            patched.append(cls.__name__)
+            continue
+
+        def _make(orig_build):
+            @functools.wraps(orig_build)
+            def _build_zero1(self, *args, **kwargs):
+                optimizers = orig_build(self, *args, **kwargs)
+                if not is_zero1_enabled():
+                    return optimizers
+                return [_wrap_in_zero1(opt) for opt in optimizers]
+
+            _build_zero1._primus_zero1_patched = True
+            return _build_zero1
+
+        cls.build = _make(vars(cls)["build"])
+        patched.append(cls.__name__)
+    logger.info("[PrimusIdeogramZeRO1] ZeRO-1 optimizer patch installed on %s", patched)
+    return bool(patched)
+
+
+def _install_ddp_ac_patch() -> bool:
+    """Whole-block-checkpoint the Ideogram layers on the DDP path (generic sub-module AC no-ops there)."""
+    import nemo_automodel.components.distributed.ddp as ddp_mod
+    import nemo_automodel.components.distributed.parallelizer as P
+
+    if getattr(ddp_mod, "_primus_ideogram_ddp_ac_patched", False):
+        return True
+
+    _orig_parallelize = ddp_mod.DDPManager.parallelize
+
+    def _parallelize_with_ideogram_ac(self, model):
+        if getattr(self, "activation_checkpointing", False) and _looks_like_ideogram(model):
+            blocks = getattr(model, _IDEOGRAM_BLOCK_ATTR, None)
+            wrapped = 0
+            if blocks is not None:
+                for idx in range(len(blocks)):
+                    blocks[idx] = P.checkpoint_wrapper(
+                        blocks[idx],
+                        checkpoint_impl=P.CheckpointImpl.NO_REENTRANT,
+                    )
+                    wrapped += 1
+            logger.info("[PrimusIdeogramZeRO1] DDP: whole-block checkpoint-wrapped %d Ideogram layers", wrapped)
+            # Disable the generic (no-op-for-Ideogram) sub-module AC so it is not attempted again.
+            saved = self.activation_checkpointing
+            self.activation_checkpointing = False
+            try:
+                return _orig_parallelize(self, model)
+            finally:
+                self.activation_checkpointing = saved
+        return _orig_parallelize(self, model)
+
+    ddp_mod.DDPManager.parallelize = _parallelize_with_ideogram_ac
+    ddp_mod._primus_ideogram_ddp_ac_patched = True
+    return True
+
+
+def install() -> bool:
+    """Install DDP / ZeRO-1 patches. No-op (returns False) unless PRIMUS_IDEOGRAM_DDP
+    or PRIMUS_IDEOGRAM_ZERO1 is set.
+
+    - PRIMUS_IDEOGRAM_DDP=1   -> whole-block AC on the DDP path (pure-DDP baseline).
+    - PRIMUS_IDEOGRAM_ZERO1=1 -> the above + wrap the optimizer in ZeroRedundancyOptimizer.
+
+    The DDP-AC patch monkeypatches DDPManager.parallelize only; it is inert on the
+    FSDP2 path (DDPManager is never used there). Idempotent; modifies NO Automodel
+    source (module-level monkeypatches only).
+    """
+    if not (is_zero1_enabled() or is_ddp_enabled()):
+        return False
+    ok_ac = _install_ddp_ac_patch()
+    ok_opt = _install_zero1_optimizer_patch() if is_zero1_enabled() else True
+    return bool(ok_ac and ok_opt)
diff --git a/primus/backends/nemo_automodel/nemo_automodel_adapter.py b/primus/backends/nemo_automodel/nemo_automodel_adapter.py
new file mode 100644
index 000000000..324554017
--- /dev/null
+++ b/primus/backends/nemo_automodel/nemo_automodel_adapter.py
@@ -0,0 +1,120 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NeMo AutoModel BackendAdapter.
+
+The AutoModel counterpart of ``TorchTitanAdapter`` / ``MaxTextAdapter``:
+
+    - Resolve the AutoModel backend path (the ``third_party/Automodel`` submodule,
+      or an already-installed ``nemo_automodel`` if the base image ships one).
+    - Convert Primus ``module_config.params`` -> AutoModel config namespace.
+    - Provide the AutoModel trainer class to the Primus runtime.
+    - Expose a backend version string for diagnostics.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from primus.backends.nemo_automodel.argument_builder import NemoAutomodelConfigBuilder
+from primus.core.backend.backend_adapter import BackendAdapter
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
+
+
+class NemoAutomodelAdapter(BackendAdapter):
+    """Complete BackendAdapter implementation for NeMo AutoModel."""
+
+    def __init__(self, framework: str = "nemo_automodel"):
+        super().__init__(framework)
+        # Submodule lives at ``third_party/Automodel`` (the upstream repo name),
+        # matching the megatron/torchtitan convention where third_party_dir_name
+        # == the checkout dir. The python package it provides is ``nemo_automodel``.
+        self.third_party_dir_name = "Automodel"
+
+    def setup_backend_path(self, backend_path=None) -> str:
+        """
+        Resolve the AutoModel python path.
+
+        The canonical source is the ``third_party/Automodel`` submodule (init via
+        ``git submodule update --init third_party/Automodel`` or ``primus-cli deps
+        sync``), which ``prepare.py`` installs editable on first run -- identical
+        to torchtitan/megatron.
+
+        As a fast path, if no explicit ``--backend_path`` / ``BACKEND_PATH`` is
+        given and ``nemo_automodel`` is already importable (e.g. a base image that
+        ships it pre-installed), we use its install location and skip the
+        third_party requirement. Otherwise we
+        fall back to the default resolution (CLI/env/third_party/deps-sync).
+        """
+        if not backend_path and not os.getenv("BACKEND_PATH"):
+            try:
+                import nemo_automodel  # noqa: F401
+
+                # Return the *import root* (the directory containing the package),
+                # not the package dir itself: Primus puts this path on PYTHONPATH
+                # for the training subprocess, so `import nemo_automodel` must
+                # resolve from it. For an installed package that is site-packages,
+                # matching the third_party/<name> convention where the backend root
+                # is the repo dir that *contains* the package (e.g. third_party/
+                # Megatron-LM for `megatron`). Returning the package dir would put
+                # the wrong dir on PYTHONPATH and risk shadowing its submodules.
+                import_root = os.path.dirname(os.path.dirname(os.path.abspath(nemo_automodel.__file__)))
+                log_rank_0(
+                    f"[Primus:NemoAutomodelAdapter] using installed nemo_automodel import root {import_root}"
+                )
+                return import_root
+            except Exception as exc:  # not installed -> fall back to third_party
+                warning_rank_0(
+                    f"[Primus:NemoAutomodelAdapter] nemo_automodel not importable ({exc}); "
+                    "falling back to third_party/backend_path resolution."
+                )
+        return super().setup_backend_path(backend_path)
+
+    def convert_config(self, params: Any):
+        """Convert Primus params -> AutoModel config namespace (pass-through)."""
+        builder = NemoAutomodelConfigBuilder()
+        builder.update(params)
+        cfg = builder.finalize()
+        log_rank_0("[Primus:NemoAutomodelAdapter] Converted Primus module params -> AutoModel config")
+        return cfg
+
+    def load_trainer_class(self, stage: str = "pretrain"):
+        """Resolve the trainer via ``BackendRegistry`` (registered in __init__)."""
+        from primus.core.backend.backend_registry import BackendRegistry
+
+        # AutoModel diffusion currently exposes a single (pretrain) trainer; treat
+        # the common ``train`` alias as ``pretrain``. Reject any other stage up
+        # front so the error names the real cause (unsupported stage) rather than
+        # implying a broken registration.
+        if stage not in ("pretrain", "train"):
+            raise RuntimeError(
+                f"[Primus:NemoAutomodelAdapter] '{self.framework}' backend only exposes a "
+                f"'pretrain' stage (got stage={stage!r}). Finetune runs also use the pretrain "
+                "trainer via model.mode: finetune in the experiment config."
+            )
+        resolved_stage = "pretrain"
+        try:
+            return BackendRegistry.get_trainer_class(self.framework, stage=resolved_stage)
+        except (ValueError, AssertionError) as exc:
+            raise RuntimeError(
+                f"[Primus:NemoAutomodelAdapter] '{self.framework}' trainer not registered for "
+                f"stage {resolved_stage!r}. Ensure primus.backends.nemo_automodel.__init__ "
+                "registers the trainer via BackendRegistry.register_trainer_class."
+            ) from exc
+
+    def detect_backend_version(self) -> str:
+        try:
+            import importlib.metadata as importlib_metadata
+        except Exception:  # pragma: no cover
+            return "unknown"
+        for dist in ("nemo_automodel", "nemo-automodel"):
+            try:
+                return importlib_metadata.version(dist)
+            except Exception:
+                continue
+        return "unknown"
diff --git a/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py b/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py
new file mode 100644
index 000000000..b3e687c39
--- /dev/null
+++ b/primus/backends/nemo_automodel/nemo_automodel_pretrain_trainer.py
@@ -0,0 +1,141 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+NemoAutomodelPretrainTrainer: Primus wrapper for NeMo AutoModel diffusion
+pre-training. Wan 2.2 (T2V-A14B) is the model that is wired up and tested
+here; AutoModel's diffusion recipe also targets other models (e.g. FLUX,
+Qwen-Image) but those are not exercised by this backend yet.
+
+Thin-wrapper pattern (same as ``MaxTextPretrainTrainer`` /
+``TorchTitanPretrainTrainer``): AutoModel owns FSDP2, the dataloader, the
+optimizer and checkpointing internally, so this trainer only
+
+    backend_args (SimpleNamespace)
+        -> cleaned dict
+        -> temp YAML
+        -> AutoModel ``parse_args_and_load_config`` (-> ConfigNode)
+        -> ``TrainDiffusionRecipe``
+
+and then delegates ``setup()`` / ``run_train_validation_loop()`` to the recipe.
+
+By routing through AutoModel's own loader we inherit its config semantics
+(``_target_``/``_fn`` resolution, the ``wandb.enable`` toggle, ...) and stay
+agnostic to AutoModel internals.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+from primus.core.trainer.base_trainer import BaseTrainer
+from primus.core.utils.module_utils import error_rank_0, log_rank_0
+
+
+class NemoAutomodelPretrainTrainer(BaseTrainer):
+    """Trainer class for NeMo AutoModel diffusion pre-training."""
+
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # The core runtime instantiates every trainer with BaseModule-style
+        # context kwargs (module_name, primus_config, module_rank, ...). Accept
+        # and forward them so BaseTrainer can filter them cooperatively.
+        super().__init__(backend_args=backend_args, **kwargs)
+        self._recipe: Optional[Any] = None
+        log_rank_0("Initialized NemoAutomodelPretrainTrainer")
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle
+    # ------------------------------------------------------------------ #
+    def setup(self):
+        """Optional pre-init phase (kept for API symmetry)."""
+        log_rank_0("NemoAutomodelPretrainTrainer.setup()")
+
+    def _install_optional_hooks(self):
+        """Install env-gated Primus-side numerics hooks BEFORE the recipe builds
+        the transformer (i.e. before ``set_attention_backend`` / first forward).
+
+        Each hook is a self-contained monkeypatch that is a no-op unless its own
+        environment flag is set, so this is safe to call unconditionally and has
+        zero effect on a default (bf16) run. No Automodel/diffusers source is
+        modified.
+        """
+        # (module, description) pairs; each exposes an env-gated install().
+        # Grouped by owner (shared, then one group per model) so a model's hooks
+        # can be added or removed without touching another's.
+        _PKG = "primus.backends.nemo_automodel"
+        hooks = (
+            # shared, model-agnostic low precision
+            (f"{_PKG}.quantization.primus_turbo_fp8", "GEMM low-precision"),
+            (f"{_PKG}.quantization.primus_turbo_fp8_attn", "FP8 attention"),
+            (f"{_PKG}.quantization.aiter_bf16_attn", "non-deterministic attention"),
+            # shared, model-agnostic sharding. Not env-gated: it repairs a dropped
+            # config value and is a no-op when the YAML omits it.
+            (f"{_PKG}.distributed.fsdp2_reshard", "FSDP2 reshard_after_forward repair"),
+            # FLUX
+            (f"{_PKG}.models.flux.parallelize", "FLUX real activation checkpointing"),
+            # Ideogram-4
+            (f"{_PKG}.models.ideogram4.attention", "Ideogram-4 var-len flash attention"),
+            (f"{_PKG}.models.ideogram4.adapter", "Ideogram-4 flow-matching adapter"),
+            # Not env-gated: it supplies a class attribute upstream requires before it
+            # will enable CP, and nothing reads it while cp_size=1.
+            (f"{_PKG}.models.ideogram4.context_parallel", "Ideogram-4 context-parallel plan"),
+            (f"{_PKG}.models.ideogram4.parallelize", "Ideogram-4 real activation checkpointing"),
+            (f"{_PKG}.models.ideogram4.zero1", "Ideogram-4 DDP + ZeRO-1 distributed optimizer"),
+            (f"{_PKG}.models.ideogram4.profile", "Ideogram-4 torch.profiler train-loop wrapper"),
+        )
+        import importlib
+
+        for mod_name, desc in hooks:
+            try:
+                mod = importlib.import_module(mod_name)
+                if mod.install():
+                    log_rank_0(f"Installed optional hook: {desc}")
+            except Exception as e:  # defensive: never break a default run
+                error_rank_0(f"Optional hook '{desc}' install failed (skipping): {e}")
+
+    def init(self):
+        """Build the AutoModel recipe from Primus params and set it up."""
+        log_rank_0("NemoAutomodelPretrainTrainer.init() - building AutoModel recipe")
+
+        from nemo_automodel.components.config._arg_parser import (
+            parse_args_and_load_config,
+        )
+        from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe
+
+        from primus.backends.nemo_automodel.argument_builder import (
+            export_params_to_yaml,
+            namespace_to_dict,
+            strip_primus_keys,
+        )
+
+        params_dict = strip_primus_keys(namespace_to_dict(self.backend_args))
+
+        # Delegate config materialization to AutoModel's own loader (argv=[] so
+        # Primus's process argv is never re-parsed).
+        yaml_path = export_params_to_yaml(params_dict)
+        try:
+            cfg = parse_args_and_load_config(yaml_path, argv=[])
+        finally:
+            try:
+                os.unlink(yaml_path)
+            except OSError as e:
+                error_rank_0(f"NemoAutomodelPretrainTrainer: failed to delete temp YAML {yaml_path}: {e}")
+
+        # Install env-gated numerics hooks before the recipe builds the model.
+        self._install_optional_hooks()
+
+        self._recipe = TrainDiffusionRecipe(cfg)
+        self._recipe.setup()
+        log_rank_0("AutoModel recipe initialized successfully")
+
+    def train(self):
+        """Execute the AutoModel train/validation loop."""
+        if self._recipe is None:
+            raise RuntimeError("NemoAutomodelPretrainTrainer.init() must be called before train().")
+        log_rank_0("Executing AutoModel diffusion pretrain...")
+        self._recipe.run_train_validation_loop()
+        log_rank_0("AutoModel diffusion pretrain completed.")
diff --git a/primus/backends/nemo_automodel/quantization/__init__.py b/primus/backends/nemo_automodel/quantization/__init__.py
new file mode 100644
index 000000000..df476c1df
--- /dev/null
+++ b/primus/backends/nemo_automodel/quantization/__init__.py
@@ -0,0 +1,17 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Model-agnostic low-precision hooks for the AutoModel diffusion backend.
+
+Everything here applies to any diffusion model the recipe can build (FP8 GEMM,
+FP8 attention, non-deterministic bf16 attention), which is why it lives outside
+``models/``.
+
+Deliberately no re-exports: each module is env-gated and pulls in optional
+dependencies (primus_turbo, aiter) at import time. The trainer imports them
+lazily by dotted path so a missing dependency degrades to a skipped hook rather
+than an import error on a default run.
+"""
diff --git a/primus/backends/nemo_automodel/quantization/aiter_bf16_attn.py b/primus/backends/nemo_automodel/quantization/aiter_bf16_attn.py
new file mode 100644
index 000000000..592b9a816
--- /dev/null
+++ b/primus/backends/nemo_automodel/quantization/aiter_bf16_attn.py
@@ -0,0 +1,266 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Non-deterministic bf16 flash-attention override for the NeMo AutoModel
+diffusion recipe (a determinism/performance trade-off knob, orthogonal to the
+FP8/MXFP4 GEMM and FP8-attention hooks).
+
+WHY:
+  The aiter flash-attention kernel exposes a ``deterministic`` flag. The default
+  (``deterministic=True``) uses a reproducible but slower backward; the
+  ``deterministic=False`` path uses atomic accumulation for the dQ/dK/dV
+  reduction, which is substantially faster in the backward while being
+  numerically equivalent up to floating-point atomic-add ordering. For
+  long-sequence models where the attention backward dominates the step, running
+  the backward non-deterministically is a large, essentially free throughput
+  win. The dispatch signature that diffusers uses has no ``deterministic``
+  argument, so it cannot be threaded through -- hence this backend override.
+
+WHAT this does (NO diffusers / Automodel fork):
+  diffusers routes every model's attention through ``_AttentionBackendRegistry``:
+  the config's ``model.attention_backend`` name selects a registered backend fn,
+  and ``dispatch_attention_fn`` looks it up in ``_backends[name]`` at call time.
+  This hook rebinds the backend *function* for the two backends the diffusion
+  recipes use -- FLASH and AITER -- to a wrapper that, for the plain path, calls
+  ``aiter.flash_attn_func(..., deterministic=False)`` directly (bf16 in, bf16
+  out; no cast, no padding). The backend NAME, its ``_supported_arg_names`` and
+  its constraint list are left untouched, so the config still selects the same
+  backend by name; only the backward numerics/ordering change.
+
+  Env-gated by ``PRIMUS_ATTN_NONDETERMINISTIC=1`` (default off = current
+  deterministic behavior). Install BEFORE the recipe builds the transformer
+  (i.e. before ``set_attention_backend``); the swap is a module-global dict entry
+  resolved at forward, so as long as it is in place before the first forward it
+  takes effect.
+
+FALLBACK (correctness first): the non-deterministic path serves the plain,
+  no-mask, dropout-free, non-context-parallel, single-output call -- exactly what
+  the diffusion hero forwards use with ``cp_size: 1``. For anything it cannot
+  serve identically (``_parallel_config`` set i.e. context parallelism,
+  ``return_lse``, an additive ``attn_mask``, nonzero dropout, or a sliding
+  ``window_size``) it transparently calls the ORIGINAL backend and warns ONCE per
+  reason. So enabling this override never changes results for a call it does not
+  handle.
+
+Activation (env, no config schema change):
+    PRIMUS_ATTN_NONDETERMINISTIC=1    enable the non-deterministic bf16 attention
+                                      override (default off = no-op)
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+
+def is_nondeterministic_attn_enabled() -> bool:
+    """Whether the non-deterministic bf16 attention override should be installed."""
+    return os.getenv("PRIMUS_ATTN_NONDETERMINISTIC", "0") in _TRUTHY
+
+
+# --------------------------------------------------------------------------- #
+# Non-deterministic bf16 flash attention                                      #
+# --------------------------------------------------------------------------- #
+def _unwrap(out):
+    """aiter may return (out, lse, ...) depending on flags; keep only the output."""
+    if isinstance(out, (tuple, list)):
+        return out[0]
+    return out
+
+
+def _flash_attn_bf16_nondet(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    """bf16 flash attention with the non-deterministic (atomic) backward.
+
+    q: (B, Sq, H, D), k/v: (B, Skv, H, D). Returns (B, Sq, H, D). No cast, no
+    padding -- the only difference from the default backend is ``deterministic``.
+
+    ``return_lse=True`` is required by aiter's autograd forward (it asserts the
+    LSE is requested so it can be saved for the backward); we always request it
+    and drop the LSE, returning only the attention output to the caller (the
+    diffusion forwards here never consume the LSE).
+    """
+    import aiter
+
+    out = aiter.flash_attn_func(
+        q, k, v, softmax_scale=softmax_scale, causal=causal,
+        deterministic=False, return_lse=True,
+    )
+    return _unwrap(out)
+
+
+# --------------------------------------------------------------------------- #
+# Backend override                                                            #
+# --------------------------------------------------------------------------- #
+_warned: set[str] = set()
+
+
+def _warn_once(key: str, msg: str) -> None:
+    if key not in _warned:
+        _warned.add(key)
+        logger.warning(msg)
+
+
+def _make_nondet_attn_backend(orig_fn, backend_value: str, supports_window: bool):
+    """Build the non-deterministic backend fn that replaces ``orig_fn``.
+
+    Signature is a superset of both the FLASH and AITER backend signatures so
+    that ``dispatch_attention_fn`` (which filters kwargs to each backend's
+    ORIGINAL ``_supported_arg_names`` and then calls ``backend_fn(**kwargs)``) can
+    drive it unchanged. ``orig_fn`` is captured for the fallback path;
+    ``supports_window`` tells us whether ``orig_fn`` accepts ``window_size``
+    (FLASH does, AITER does not).
+    """
+
+    def _fallback(query, key, value, attn_mask, dropout_p, is_causal, scale,
+                  window_size, return_lse, _parallel_config):
+        kw = dict(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+            scale=scale,
+            return_lse=return_lse,
+            _parallel_config=_parallel_config,
+        )
+        if supports_window:
+            kw["window_size"] = window_size
+        return orig_fn(**kw)
+
+    def nondet_attn_backend(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: Optional[float] = None,
+        window_size: tuple[int, int] = (-1, -1),
+        return_lse: bool = False,
+        _parallel_config=None,
+    ):
+        # Correctness-first fallbacks: use the original kernel for anything the
+        # plain non-deterministic path cannot faithfully serve (warn once).
+        if _parallel_config is not None:
+            _warn_once(
+                f"{backend_value}:cp",
+                f"[PrimusAttn-NonDet] context parallelism active on backend "
+                f"'{backend_value}' -> default attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if return_lse:
+            _warn_once(
+                f"{backend_value}:lse",
+                f"[PrimusAttn-NonDet] return_lse requested on backend "
+                f"'{backend_value}' -> default attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if attn_mask is not None:
+            _warn_once(
+                f"{backend_value}:mask",
+                f"[PrimusAttn-NonDet] additive attn_mask on backend "
+                f"'{backend_value}' -> default attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if dropout_p:
+            _warn_once(
+                f"{backend_value}:dropout",
+                f"[PrimusAttn-NonDet] dropout_p={dropout_p} on backend "
+                f"'{backend_value}' -> default attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if window_size not in ((-1, -1), None):
+            _warn_once(
+                f"{backend_value}:window",
+                f"[PrimusAttn-NonDet] sliding window_size={window_size} on backend "
+                f"'{backend_value}' -> default attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+
+        return _flash_attn_bf16_nondet(query, key, value, softmax_scale=scale, causal=is_causal)
+
+    nondet_attn_backend._primus_nondet_attn = True  # idempotence marker
+    nondet_attn_backend._primus_orig_fn = orig_fn
+    return nondet_attn_backend
+
+
+# Backends the diffusion recipes use: FLASH (accepts window_size) and AITER
+# (no window_size). Override both; only the active one is dispatched.
+_TARGET_BACKENDS = (("FLASH", True), ("AITER", False))
+
+
+def install() -> bool:
+    """Install the non-deterministic attention override by rebinding diffusers
+    backend functions.
+
+    No-op (returns False) unless ``PRIMUS_ATTN_NONDETERMINISTIC`` is set. Modifies
+    NO diffusers/Automodel source; only rebinds
+    ``_AttentionBackendRegistry._backends`` entries at runtime. Idempotent. Call
+    BEFORE the transformer's first forward (in practice before the recipe build /
+    ``set_attention_backend``).
+    """
+    if not is_nondeterministic_attn_enabled():
+        return False
+
+    # Fail fast if aiter's flash-attention is unavailable so the run errors
+    # clearly rather than silently running the default backward.
+    import aiter  # noqa: F401
+
+    from diffusers.models.attention_dispatch import (
+        AttentionBackendName,
+        _AttentionBackendRegistry,
+    )
+
+    reg = _AttentionBackendRegistry
+    installed = []
+    for name_str, supports_window in _TARGET_BACKENDS:
+        name = getattr(AttentionBackendName, name_str, None)
+        if name is None:
+            continue
+        orig = reg._backends.get(name)
+        if orig is None:
+            # Backend not registered in this diffusers build.
+            continue
+        if getattr(orig, "_primus_nondet_attn", False):
+            installed.append(name.value)  # already installed
+            continue
+        reg._backends[name] = _make_nondet_attn_backend(orig, name.value, supports_window)
+        # Leave reg._supported_arg_names[name] and reg._constraints[name] untouched:
+        # dispatch filters kwargs to the ORIGINAL supported set (our fn accepts a
+        # superset), and the qkv/device/shape checks still hold.
+        installed.append(name.value)
+
+    if not installed:
+        logger.warning(
+            "[PrimusAttn-NonDet] PRIMUS_ATTN_NONDETERMINISTIC set but neither FLASH "
+            "nor AITER backend is registered in diffusers; override NOT active."
+        )
+        return False
+
+    logger.info(
+        "[PrimusAttn-NonDet] Installed non-deterministic bf16 flash-attention "
+        "(aiter flash_attn_func, deterministic=False) override for backends: %s "
+        "(config attention_backend name unchanged)",
+        ", ".join(installed),
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/quantization/primus_turbo_fp8.py b/primus/backends/nemo_automodel/quantization/primus_turbo_fp8.py
new file mode 100644
index 000000000..9f2e3c30b
--- /dev/null
+++ b/primus/backends/nemo_automodel/quantization/primus_turbo_fp8.py
@@ -0,0 +1,1578 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Primus-Turbo low-precision (FP8 / MXFP4) linear swap for the NeMo AutoModel
+diffusion recipe.
+
+Why:
+  Config-only Transformer Engine FP8 does not train Wan2.2 on rocm/primus:v26.3
+  (TE 2.12.0.dev0, ROCm 7.2): per-tensor delayed/current has no hipBLASLt fp8
+  backward GEMM on gfx950, and MXFP8 has no bias support. Primus-Turbo's
+  ``Float8Linear`` (aiter-backed ``gemm_fp8``, autograd fwd+bwd, bias applied
+  OUTSIDE the fp8 GEMM) sidesteps both gaps and is the SAME library that powers
+  the verified FLUX fp8 diffusion training on MI355X. The MXFP4 path reuses the
+  same seam via ``primus_turbo.pytorch.ops.gemm_fp4`` (autograd fwd+bwd,
+  MX_BLOCKWISE / E2M1 / block-32, HIPBLASLT backend) wrapped in a local
+  ``Float4Linear`` (Primus-Turbo 0.2.0 ships no Float4Linear module).
+
+What this does (NO Automodel fork):
+  Monkeypatches ``nemo_automodel._diffusers.auto_diffusion_pipeline
+  ._replace_linear_with_transformer_engine`` so that the existing config seam
+  (``model.transformer_engine_linear=true``) swaps ``torch.nn.Linear`` ->
+  ``primus_turbo.pytorch.modules.Float8Linear`` (FP8) or the local
+  ``Float4Linear`` (MXFP4) instead of TE Linear. The swap runs on the built
+  transformer BEFORE FSDP2 wrapping, mirroring the TE path (weight/bias copy,
+  requires_grad preserved). We ALWAYS apply AutoModel's
+  ``_is_fp8_training_safe_linear`` skip-list for stability (keep time/text embed,
+  norm-modulation, final proj, and non-16-aligned Linears in bf16); the MXFP4
+  path adds a stricter 128-alignment guard (see ``_is_fp4_training_safe_linear``).
+
+Activation & tuning (env, so no Automodel source or config schema changes):
+  # --- FP8 (Float8Linear, aiter gemm_fp8) ---
+  PRIMUS_TURBO_FP8=1                      enable the FP8 swap (default off = no-op)
+  PRIMUS_TURBO_FP8_GRANULARITY=TENSORWISE ROWWISE|TENSORWISE|BLOCKWISE|MX_BLOCKWISE
+  PRIMUS_TURBO_FP8_FORMAT=E4M3            E4M3|E5M2|HYBRID
+  # --- TE-native MXFP4 (te.pytorch.Linear + MXFP4BlockScaling autocast) ------
+  # Highest precedence. Requires a TE-native image (TransformerEngine rebuilt
+  # with MXFP4 recipe support). TE's recipe applies a Random Hadamard
+  # Transform (RHT) + native fused HIP cast/transpose kernel and dispatches the
+  # GEMM to AITER a4w4 -- a configuration shown to converge.
+  # NOTE: RHT is NOT the differentiator. primus_turbo's gemm_fp4 ALSO applies RHT
+  # -- a deterministic 32x32 Hadamard x fixed sign vector (kernels/quantization/
+  # hadamard_transform.py), on exactly the two transposed Wgrad operands (a_t,
+  # grad_out_t), matching the NVFP4 prescription. Yet the primus_turbo MXFP4
+  # Wgrad still explodes on Wan (isolated via PRIMUS_TURBO_FP4_
+  # BACKWARD=wgrad_fp8). So the real difference is
+  # TE's fused cast/transpose + a4w4 Wgrad SCALE HANDLING vs primus_turbo's, not
+  # the presence of RHT. This path tests whether TE's Wgrad avoids the explosion.
+  PRIMUS_TE_MXFP4=1                       enable the TE-native MXFP4 swap+autocast
+  # --- MXFP4 (Float4Linear, gemm_fp4) --- takes precedence over FP8 if set ---
+  PRIMUS_TURBO_FP4=1                      enable the MXFP4 swap (default off = no-op)
+  PRIMUS_TURBO_FP4_KEEP_SENSITIVE=1      keep Wan condition_embedder.* (timestep/
+                                         text/AdaLN) Linears in bf16 (default on)
+  PRIMUS_TURBO_FP4_SR=off                stochastic rounding scope: off|grad|all.
+                                         primus_turbo 0.2.0 hardcodes use_sr=False
+                                         everywhere (the divergent 'FP4 w/o SR'
+                                         regime);
+                                         grad = SR on backward gradient quant,
+                                         all = SR on fwd (act+weight) + bwd grad.
+  # --- MXFP4 convergence scaffolding -----------------------------------------
+  # Extra scaffolding from a converging MXFP4 ablation ("NVFP4-aligned"): pure
+  # MXFP4 backward + gradient SR + a band of early/late transformer blocks kept
+  # in tensorwise FP8. Same primus_turbo kernels -- convergence came from the
+  # scaffolding around the kernel, not the kernel.
+  PRIMUS_TURBO_FP4_SENSITIVE_LAYERS=0    keep a band of transformer blocks in a
+                                         higher precision than MXFP4 (default off)
+  PRIMUS_TURBO_FP4_SENSITIVE_START=2     # of leading blocks in the band (first N)
+  PRIMUS_TURBO_FP4_SENSITIVE_END=8       # of trailing blocks in the band (last N)
+  PRIMUS_TURBO_FP4_SENSITIVE_PRECISION=tw_fp8  band precision: tw_fp8|bf16
+                                         (tw_fp8 = primus_turbo Float8Linear
+                                          tensorwise; bf16 = keep torch.nn.Linear)
+  PRIMUS_TURBO_FP4_BACKWARD=mxfp4        backward GEMM precision:
+                                         mxfp4|mxfp4_pad|fp8.
+                                         mxfp4 = pure MXFP4 fwd+bwd (default,
+                                         matches the converging Run C); fp8 =
+                                         hybrid MXFP4 fwd / tensorwise-FP8 bwd
+                                         (mirrors mxfp4_backward_precision: fp8;
+                                         no activation-memory savings -- saves
+                                         bf16 x/W and requantizes to FP8 in bwd).
+                                         mxfp4_pad = THE FIX: pure MXFP4 fwd+bwd
+                                         with the token/contraction dim zero-padded
+                                         to %256 so the AITER a4w4 Wgrad (contracts
+                                         over tokens) is numerically correct
+                                         instead of returning ~1e7x garbage on
+                                         non-%256 token counts (Wan cross-attn
+                                         k/v: text_seq*batch). Zeros are exact;
+                                         keeps the tuned preshuffled AITER path.
+  # (MXFP4 is fixed to MX_BLOCKWISE + E2M1_X2 + block_size 32 + E8M0 scale by the
+  #  installed primus_turbo 0.2.0; no granularity/format knobs are exposed.)
+
+Recommended config pairing (keeps TE autocast OFF, drives only the swap):
+  model.transformer_engine_linear: true
+  model.transformer_engine_fp8: false
+
+MXFP4 kernel constraints (installed primus_turbo 0.2.0, HIPBLASLT FP4 backend):
+  The forward + both backward GEMMs are NT-layout, bf16, and (accounting for FP4
+  2-values/byte packing and the transposed grad GEMMs) require the logical
+  M (tokens), N (out_features) and K (in_features) to ALL be multiples of 128.
+  in_features/out_features are enforced at swap time via the skip-list; the token
+  dimension M is a runtime property (batch*seq) checked with a one-time warning.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import re
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+
+def is_enabled() -> bool:
+    """Whether the Primus-Turbo FP8 swap should be installed."""
+    return os.getenv("PRIMUS_TURBO_FP8", "0") in _TRUTHY
+
+
+def is_fp4_enabled() -> bool:
+    """Whether the Primus-Turbo MXFP4 swap should be installed.
+
+    Takes precedence over FP8 when both are set (see ``install``).
+    """
+    return os.getenv("PRIMUS_TURBO_FP4", "0") in _TRUTHY
+
+
+def is_te_mxfp4_enabled() -> bool:
+    """Whether the TE-native MXFP4 swap should be installed.
+
+    Highest precedence (see ``install``). Only meaningful on a TE-native image
+    where TransformerEngine was rebuilt with MXFP4 recipe support. On a
+    baseline image TE exposes the MXFP4BlockScaling *name*
+    but none of the execution machinery, so this path fails fast at install.
+    """
+    return os.getenv("PRIMUS_TE_MXFP4", "0") in _TRUTHY
+
+
+def _resolve_fp8_config():
+    """Build a Float8QuantConfig from env (defaults: TENSORWISE + E4M3 + DYNAMIC)."""
+    from primus_turbo.pytorch.core.low_precision import (
+        Float8QuantConfig,
+        Format,
+        ScalingGranularity,
+    )
+
+    gran_name = os.getenv("PRIMUS_TURBO_FP8_GRANULARITY", "TENSORWISE").upper()
+    fmt_name = os.getenv("PRIMUS_TURBO_FP8_FORMAT", "E4M3").upper()
+    try:
+        granularity = getattr(ScalingGranularity, gran_name)
+    except AttributeError as exc:
+        raise ValueError(
+            f"PRIMUS_TURBO_FP8_GRANULARITY={gran_name!r} invalid; expected one of "
+            "TENSORWISE, ROWWISE, BLOCKWISE, MX_BLOCKWISE"
+        ) from exc
+    try:
+        fmt = getattr(Format, fmt_name)
+    except AttributeError as exc:
+        raise ValueError(
+            f"PRIMUS_TURBO_FP8_FORMAT={fmt_name!r} invalid; expected one of E4M3, E5M2, HYBRID"
+        ) from exc
+    return Float8QuantConfig(format=fmt, granularity=granularity)
+
+
+# --------------------------------------------------------------------------- #
+# MXFP4 (Float4Linear on top of primus_turbo.pytorch.ops.gemm_fp4)            #
+# --------------------------------------------------------------------------- #
+# Alignment required by the installed primus_turbo 0.2.0 HIPBLASLT FP4 GEMM
+# (fwd + both bwd GEMMs, NT layout, FP4 packs 2 vals/byte). See module docstring.
+_FP4_ALIGN = 128
+
+# Wan-specific precision-sensitive projections kept in bf16 by default under
+# MXFP4. AutoModel's FP8-safe predicate uses FLUX names (time_text_embed.,
+# norm_out.) that do NOT match Wan's module tree, so without this the timestep
+# embedder, text embedder and AdaLN modulation (time_proj) would all be swapped
+# to FP4. FP4 is much more aggressive than FP8, and diffusion timestep/text/
+# modulation conditioning is the classic "keep in high precision" set; these are
+# only 5 of 405 eligible Linears, so the throughput cost is negligible. Toggle
+# with PRIMUS_TURBO_FP4_KEEP_SENSITIVE (default on) to A/B full coverage.
+_FP4_SENSITIVE_PREFIXES = ("condition_embedder.",)
+
+
+def _fp4_keep_sensitive_bf16() -> bool:
+    return os.getenv("PRIMUS_TURBO_FP4_KEEP_SENSITIVE", "1") in _TRUTHY
+
+
+_FP4_SR_SCOPES = {"off", "grad", "all"}
+
+
+def _fp4_sr_scope() -> str:
+    """Stochastic-rounding scope for the MXFP4 quantization.
+
+    primus_turbo 0.2.0's ``FP4GemmMXFunction`` hardcodes ``use_sr=False`` at every
+    FP4 quant site -- the "FP4 without stochastic rounding" regime that can make
+    native-TE MXFP4 diverge, where enabling SR recovers convergence. This
+    env opts into SR via our local ``Float4Linear`` (NO primus_turbo/Automodel
+    fork):
+      off  (default) -> stock behavior, SR nowhere (calls primus_turbo.gemm_fp4)
+      grad           -> SR on the backward gradient quant (debias dgrad/wgrad)
+      all            -> SR on forward (act+weight) AND backward gradient quant
+    """
+    scope = os.getenv("PRIMUS_TURBO_FP4_SR", "off").lower()
+    if scope in _TRUTHY:
+        scope = "grad"
+    if scope not in _FP4_SR_SCOPES:
+        raise ValueError(
+            f"PRIMUS_TURBO_FP4_SR={scope!r} invalid; expected one of off, grad, all"
+        )
+    return scope
+
+
+# --------------------------------------------------------------------------- #
+# MXFP4 convergence scaffolding                                                #
+# --------------------------------------------------------------------------- #
+# Backward precision -----------------------------------------------------------
+# mxfp4     : pure MXFP4 dgrad + wgrad (the FLUX Run-C recipe).
+# fp8       : hybrid -- BOTH dgrad + wgrad in tensorwise FP8 (MXFP4 fwd).
+# dgrad_fp8 : ISOLATION -- dgrad (grad wrt input) in FP8, wgrad stays MXFP4.
+# wgrad_fp8 : ISOLATION -- wgrad (grad wrt weight) in FP8, dgrad stays MXFP4.
+# The two isolation modes pin down which backward GEMM drives the grad explosion.
+# Per NVIDIA's NVFP4 paper (arXiv:2509.25149 sec 4 / App E.2), Wgrad carries the
+# largest FP4 quant error (esp. late blocks), so `wgrad_fp8` is the prime suspect.
+_FP4_BACKWARD_PRECISIONS = {"mxfp4", "mxfp4_pad", "fp8", "dgrad_fp8", "wgrad_fp8"}
+
+# AITER a4w4 f4gemm returns ~1e7x garbage whenever its CONTRACTION dim K is not a
+# multiple of 256 (kernel bug; HIPBLASLT is correct but its rocRoller FP4
+# solutions require the token OUTPUT dim %256 on fwd/dgrad and hang, which is why
+# the stack pins FP4 to AITER). For fwd/dgrad
+# K = in/out_features (already %256 for Wan); for the wgrad GEMM K = tokens, which
+# for cross-attention key/value projections is text_seq*batch (e.g. 226*8=1808,
+# not %256) -> the wgrad explodes. The ``mxfp4_pad`` backward zero-pads the token
+# (contraction) dim to this multiple so the wgrad stays in AITER's correct regime;
+# zeros are numerically exact (they contribute 0 to grad_out^T @ x).
+_AITER_FP4_K_MULTIPLE = 256
+
+
+def _fp4_backward_precision() -> str:
+    """MXFP4 backward GEMM precision (mirrors/extends ``mxfp4_backward_precision``).
+
+    ``mxfp4`` (default) = pure MXFP4 fwd+bwd (the converging Run C recipe).
+    ``fp8`` = hybrid: MXFP4 forward, tensorwise-FP8 backward (both dgrad+wgrad).
+    ``dgrad_fp8`` / ``wgrad_fp8`` = isolate one backward GEMM in FP8, the other in
+    MXFP4 (diagnostic to attribute the gradient explosion to dgrad vs wgrad).
+    The FP8 backward keeps bf16 x/W and requantizes at backward time (no
+    activation-memory savings for the FP8 leg), matching the local spec.
+    """
+    p = os.getenv("PRIMUS_TURBO_FP4_BACKWARD", "mxfp4").lower()
+    if p not in _FP4_BACKWARD_PRECISIONS:
+        raise ValueError(
+            f"PRIMUS_TURBO_FP4_BACKWARD={p!r} invalid; expected one of "
+            "mxfp4, mxfp4_pad, fp8, dgrad_fp8, wgrad_fp8"
+        )
+    return p
+
+
+def _fp8_stage_dtype(fmt, is_fwd_stage: bool):
+    """Map a primus_turbo FP8 ``Format`` to the concrete dtype for a GEMM stage.
+
+    HYBRID = E4M3 for the forward-stage operands (act/weight), E5M2 for gradients
+    -- the same convention as primus_turbo's ``FP8GemmTensorFunction``.
+    """
+    from primus_turbo.pytorch.core.low_precision import (
+        Format,
+        float8_e4m3,
+        float8_e5m2,
+    )
+
+    if fmt == Format.E4M3:
+        return float8_e4m3
+    if fmt == Format.E5M2:
+        return float8_e5m2
+    return float8_e4m3 if is_fwd_stage else float8_e5m2
+
+
+# Sensitive-layer band ---------------------------------------------------------
+_FP4_SENSITIVE_PRECISIONS = {"tw_fp8", "bf16"}
+# Matches Wan's ``blocks.N.`` and FLUX's ``transformer_blocks.N.`` module trees.
+_FP4_BLOCK_RE = re.compile(r"(?:^|\.)((?:\w+_)*blocks)\.(\d+)(?:\.|$)")
+
+
+def _fp4_sensitive_enabled() -> bool:
+    return os.getenv("PRIMUS_TURBO_FP4_SENSITIVE_LAYERS", "0") in _TRUTHY
+
+
+def _fp4_sensitive_start() -> int:
+    return int(os.getenv("PRIMUS_TURBO_FP4_SENSITIVE_START", "2"))
+
+
+def _fp4_sensitive_end() -> int:
+    return int(os.getenv("PRIMUS_TURBO_FP4_SENSITIVE_END", "8"))
+
+
+def _fp4_sensitive_precision() -> str:
+    p = os.getenv("PRIMUS_TURBO_FP4_SENSITIVE_PRECISION", "tw_fp8").lower()
+    if p not in _FP4_SENSITIVE_PRECISIONS:
+        raise ValueError(
+            f"PRIMUS_TURBO_FP4_SENSITIVE_PRECISION={p!r} invalid; expected one of tw_fp8, bf16"
+        )
+    return p
+
+
+def _block_index(fqn: str):
+    """Return (block_list_name, block_index) parsed from a module FQN, or (None, None).
+
+    e.g. ``blocks.7.attn1.to_q`` -> ("blocks", 7); ``proj_out`` -> (None, None).
+    """
+    m = _FP4_BLOCK_RE.search(fqn)
+    if not m:
+        return None, None
+    return m.group(1), int(m.group(2))
+
+
+def _count_blocks(module: nn.Module):
+    """Map each transformer-block ModuleList prefix -> block count (max idx + 1).
+
+    Walks the module's named Linears so the band ('first N + last M blocks') can
+    be resolved without assuming a specific attribute name.
+    """
+    counts: dict[str, int] = {}
+    for name, child in module.named_modules():
+        if not isinstance(child, nn.Linear):
+            continue
+        list_name, idx = _block_index(name)
+        if list_name is not None:
+            counts[list_name] = max(counts.get(list_name, 0), idx + 1)
+    return counts
+
+
+def _is_sensitive_block(fqn: str, block_counts: dict) -> bool:
+    """Whether ``fqn`` lives in the sensitive band (first START + last END blocks)."""
+    list_name, idx = _block_index(fqn)
+    if list_name is None or list_name not in block_counts:
+        return False
+    start = _fp4_sensitive_start()
+    end = _fp4_sensitive_end()
+    count = block_counts[list_name]
+    return idx < start or idx >= (count - end)
+
+
+def _resolve_fp4_hybrid_fp8_config():
+    """Tensorwise-FP8 config for the MXFP4 hybrid backward (HYBRID = E4M3 fwd / E5M2 grad)."""
+    from primus_turbo.pytorch.core.low_precision import (
+        Float8QuantConfig,
+        Format,
+        ScalingGranularity,
+    )
+
+    return Float8QuantConfig(format=Format.HYBRID, granularity=ScalingGranularity.TENSORWISE)
+
+
+class _FP4FwdFP8BwdFunction(torch.autograd.Function):
+    """Hybrid MXFP4-forward / tensorwise-FP8-backward GEMM (mirrors the
+    ``mxfp4_backward_precision='fp8'`` reference path).
+
+    Forward quantizes activation+weight to MXFP4 and runs ``gemm_fp4_impl`` (same
+    E2M1_X2 / block-32 / NT recipe as the pure path), but saves **bf16 copies** of
+    the activation and weight. Backward requantizes those bf16 tensors and the
+    incoming gradient to tensorwise FP8 and runs ``gemm_fp8_impl`` for dgrad/wgrad
+    -- byte-for-byte the recipe primus_turbo's ``FP8GemmTensorFunction`` uses
+    (a=x trans_a=False, b=W trans_b=True). Trade-off: no activation-memory savings
+    (bf16 saved), but the backward uses the more mature tensorwise-FP8 GEMM.
+    """
+
+    @staticmethod
+    def forward(ctx, a, b, out_dtype, fp4_config, fp8_config, sr_fwd):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            MXScalingRecipe,
+            check_mxfp4_support,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        ok, reason = check_mxfp4_support()
+        assert ok, reason
+        fp4_dtype = torch.float4_e2m1fn_x2
+
+        # Forward: MXFP4 (row recipes only; the transposed tensors are unused
+        # because the backward runs in FP8 from saved bf16 copies).
+        a_fp4, a_scale_inv, _, _ = quantize_fp4_with_trans(
+            a,
+            fp4_dtype,
+            fp4_config.granularity,
+            block_size=fp4_config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=True,
+                shuffle_scale=enable_preshuffle(),
+            ),
+        )
+        b_fp4, b_scale_inv, _, _ = quantize_fp4_with_trans(
+            b,
+            fp4_dtype,
+            fp4_config.granularity,
+            block_size=fp4_config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+        )
+        out = gemm_fp4_impl(
+            a_fp4,
+            a_scale_inv,
+            False,
+            b_fp4,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=fp4_config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+
+        # Save bf16 copies for the FP8 backward (dtype-correct, per local spec).
+        ctx.save_for_backward(a, b)
+        ctx.out_dtype = out_dtype
+        ctx.fp8_config = fp8_config
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            Format,
+            float8_e4m3,
+            float8_e5m2,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl
+        from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+        a, b = ctx.saved_tensors  # bf16 activation (x) and weight (W)
+        cfg = ctx.fp8_config
+        gran = cfg.granularity
+
+        def _fp8_dtype(is_fwd_stage: bool):
+            if cfg.format == Format.E4M3:
+                return float8_e4m3
+            if cfg.format == Format.E5M2:
+                return float8_e5m2
+            return float8_e4m3 if is_fwd_stage else float8_e5m2  # HYBRID
+
+        if not grad_out.is_contiguous():
+            grad_out = grad_out.contiguous()
+
+        fwd_dt = _fp8_dtype(True)
+        grad_dt = _fp8_dtype(False)
+        a_fp8, a_scale_inv = quantize_fp8(a, fwd_dt, gran)
+        b_fp8, b_scale_inv = quantize_fp8(b, fwd_dt, gran)
+        grad_fp8, grad_scale_inv = quantize_fp8(grad_out, grad_dt, gran)
+
+        # Mirrors FP8GemmTensorFunction.backward with (a=x, trans_a=False,
+        # b=W, trans_b=True): out = x @ W^T.
+        # grad_a = grad @ W
+        grad_a = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=gran.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        # grad_b = x^T @ grad  (trans_c=True yields the [out, in] weight grad)
+        grad_b = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            True,
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=gran.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        return grad_a, grad_b, None, None, None, None
+
+
+def _gemm_fp4_hybrid(a, b, out_dtype, fp4_config, fp8_config, sr_fwd):
+    """MXFP4-forward / tensorwise-FP8-backward GEMM (our local Float4Linear seam)."""
+    return _FP4FwdFP8BwdFunction.apply(a, b, out_dtype, fp4_config, fp8_config, sr_fwd)
+
+
+class _FP4GemmMixedFunction(torch.autograd.Function):
+    """MXFP4 forward with a PER-GEMM-choosable backward precision, to isolate
+    whether the gradient explosion comes from Dgrad (grad wrt input) or Wgrad
+    (grad wrt weight).
+
+    ``dgrad_fp8`` / ``wgrad_fp8`` independently pick tensorwise-FP8 (True) or
+    MXFP4 (False) for each backward GEMM. The MXFP4 leg reuses the exact stock
+    recipe (E2M1_X2 / block-32 / NT, RHT on the transposed grad, SR per
+    ``sr_grad``); the FP8 leg reuses primus_turbo's tensorwise recipe. Forward
+    saves BOTH the transposed FP4 tensors (for the MXFP4 leg) AND bf16 x/W (for
+    the FP8 leg) -- heavier, but this is a diagnostic path, not the perf path.
+    """
+
+    @staticmethod
+    def forward(ctx, a, b, out_dtype, fp4_config, fp8_config, sr_fwd, sr_grad, dgrad_fp8, wgrad_fp8):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            MXScalingRecipe,
+            check_mxfp4_support,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        ok, reason = check_mxfp4_support()
+        assert ok, reason
+        fp4_dtype = torch.float4_e2m1fn_x2
+
+        a_fp4, a_scale_inv, a_t_fp4, a_t_scale_inv = quantize_fp4_with_trans(
+            a, fp4_dtype, fp4_config.granularity, block_size=fp4_config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False, use_sr=sr_fwd, use_rht=False, shuffle_scale=enable_preshuffle()),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False, use_sr=sr_fwd, use_rht=True, shuffle_scale=enable_preshuffle()),
+        )
+        b_fp4, b_scale_inv, b_t_fp4, b_t_scale_inv = quantize_fp4_with_trans(
+            b, fp4_dtype, fp4_config.granularity, block_size=fp4_config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=True, use_sr=sr_fwd, use_rht=False,
+                shuffle_scale=enable_preshuffle(), shuffle_out=enable_preshuffle()),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=True, use_sr=sr_fwd, use_rht=False,
+                shuffle_scale=enable_preshuffle(), shuffle_out=enable_preshuffle()),
+        )
+        out = gemm_fp4_impl(
+            a_fp4, a_scale_inv, False, b_fp4, b_scale_inv, True, out_dtype, False,
+            granularity=fp4_config.granularity.value, default_backend=BackendType.HIPBLASLT.value,
+        )
+        # Save both representations: FP4 transposed (MXFP4 legs) + bf16 (FP8 legs).
+        ctx.save_for_backward(a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv, a, b)
+        ctx.out_dtype = out_dtype
+        ctx.fp4_config = fp4_config
+        ctx.fp8_config = fp8_config
+        ctx.sr_grad = sr_grad
+        ctx.dgrad_fp8 = dgrad_fp8
+        ctx.wgrad_fp8 = wgrad_fp8
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import MXScalingRecipe
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl
+        from primus_turbo.pytorch.ops.quantization import (
+            quantize_fp4_with_trans,
+            quantize_fp8,
+        )
+
+        a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv, a, b = ctx.saved_tensors
+        if not grad_out.is_contiguous():
+            grad_out = grad_out.contiguous()
+        grad_out = grad_out.view(grad_out.shape[0], -1)
+        fp4_dtype = torch.float4_e2m1fn_x2
+        fp4_gran = ctx.fp4_config.granularity
+        fp8_gran = ctx.fp8_config.granularity
+        HB = BackendType.HIPBLASLT.value
+
+        # --- MXFP4 grad_out quant (row + transposed), shared by whichever leg(s)
+        #     stay in FP4. RHT on the transposed recipe (feeds wgrad). ---
+        (g_fp4, g_scale, g_t_fp4, g_t_scale) = quantize_fp4_with_trans(
+            grad_out, fp4_dtype, fp4_gran, block_size=ctx.fp4_config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False, use_sr=ctx.sr_grad, use_rht=False, shuffle_scale=enable_preshuffle()),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False, use_sr=ctx.sr_grad, use_rht=True, shuffle_scale=enable_preshuffle()),
+        )
+        # --- tensorwise FP8 quant of grad_out / a / b, shared by the FP8 leg(s). ---
+        g_dt = _fp8_stage_dtype(ctx.fp8_config.format, False)
+        fwd_dt = _fp8_stage_dtype(ctx.fp8_config.format, True)
+        g_fp8, g_fp8_scale = quantize_fp8(grad_out, g_dt, fp8_gran)
+
+        # dgrad = grad @ W  ([M,in])
+        if ctx.dgrad_fp8:
+            b_fp8, b_fp8_scale = quantize_fp8(b, fwd_dt, fp8_gran)
+            grad_a = gemm_fp8_impl(
+                g_fp8, g_fp8_scale, False, b_fp8, b_fp8_scale, False, ctx.out_dtype, False,
+                granularity=fp8_gran.value, default_backend=HB)
+        else:
+            grad_a = gemm_fp4_impl(
+                g_fp4, g_scale, False, b_t_fp4, b_t_scale_inv, True, ctx.out_dtype, False,
+                granularity=fp4_gran.value, default_backend=HB)
+
+        # wgrad = x^T @ grad  ([out,in])
+        if ctx.wgrad_fp8:
+            a_fp8, a_fp8_scale = quantize_fp8(a, fwd_dt, fp8_gran)
+            grad_b = gemm_fp8_impl(
+                a_fp8, a_fp8_scale, True, g_fp8, g_fp8_scale, False, ctx.out_dtype, True,
+                granularity=fp8_gran.value, default_backend=HB)
+        else:
+            grad_b = gemm_fp4_impl(
+                g_t_fp4, g_t_scale, False, a_t_fp4, a_t_scale_inv, True, ctx.out_dtype, False,
+                granularity=fp4_gran.value, default_backend=HB)
+
+        return grad_a, grad_b, None, None, None, None, None, None, None
+
+
+def _gemm_fp4_mixed(a, b, out_dtype, fp4_config, fp8_config, sr_fwd, sr_grad, dgrad_fp8, wgrad_fp8):
+    """MXFP4 fwd with per-GEMM backward precision (dgrad/wgrad isolation seam)."""
+    return _FP4GemmMixedFunction.apply(
+        a, b, out_dtype, fp4_config, fp8_config, sr_fwd, sr_grad, dgrad_fp8, wgrad_fp8
+    )
+
+
+class _FP4GemmMXSRFunction(torch.autograd.Function):
+    """MXFP4 GEMM (E2M1_X2, block-32, NT layout) mirroring primus_turbo 0.2.0's
+    ``FP4GemmMXFunction`` but with stochastic rounding selectable per tensor.
+
+    ONLY the ``use_sr`` flags differ from the stock recipe; every other flag
+    (``use_2d_block``, ``use_rht``, ``shuffle_scale``/``shuffle_out``, the NT
+    layout, and the HIPBLASLT default backend) is copied verbatim so that
+    ``sr_fwd=sr_grad=False`` is numerically the stock path. ``sr_fwd`` applies SR
+    to the forward activation+weight quant; ``sr_grad`` to the backward gradient
+    quant (both the row and transposed recipes used by dgrad and wgrad).
+    """
+
+    @staticmethod
+    def forward(ctx, a, b, out_dtype, config, sr_fwd, sr_grad):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            MXScalingRecipe,
+            check_mxfp4_support,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        ok, reason = check_mxfp4_support()
+        assert ok, reason
+        fp4_dtype = torch.float4_e2m1fn_x2
+
+        a_fp4, a_scale_inv, a_t_fp4, a_t_scale_inv = quantize_fp4_with_trans(
+            a,
+            fp4_dtype,
+            config.granularity,
+            block_size=config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=True,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+        )
+        b_fp4, b_scale_inv, b_t_fp4, b_t_scale_inv = quantize_fp4_with_trans(
+            b,
+            fp4_dtype,
+            config.granularity,
+            block_size=config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+        )
+        out = gemm_fp4_impl(
+            a_fp4,
+            a_scale_inv,
+            False,
+            b_fp4,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+
+        ctx.save_for_backward(a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv)
+        ctx.out_dtype = out_dtype
+        ctx.config = config
+        ctx.sr_grad = sr_grad
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import MXScalingRecipe
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv = ctx.saved_tensors
+        fp4_dtype = torch.float4_e2m1fn_x2
+        grad_out = grad_out.view(grad_out.shape[0], -1)
+
+        (
+            grad_out_fp4,
+            grad_out_scale_inv,
+            grad_out_t_fp4,
+            grad_out_t_scale_inv,
+        ) = quantize_fp4_with_trans(
+            grad_out,
+            fp4_dtype,
+            ctx.config.granularity,
+            block_size=ctx.config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=ctx.sr_grad,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=ctx.sr_grad,
+                use_rht=True,
+                shuffle_scale=enable_preshuffle(),
+            ),
+        )
+        grad_a = gemm_fp4_impl(
+            grad_out_fp4,
+            grad_out_scale_inv,
+            False,
+            b_t_fp4,
+            b_t_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        grad_b = gemm_fp4_impl(
+            grad_out_t_fp4,
+            grad_out_t_scale_inv,
+            False,
+            a_t_fp4,
+            a_t_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        return grad_a, grad_b, None, None, None, None
+
+
+def _gemm_fp4_sr(a, b, out_dtype, config, sr_fwd, sr_grad):
+    """SR-enabled MXFP4 GEMM (our local Float4Linear seam)."""
+    return _FP4GemmMXSRFunction.apply(a, b, out_dtype, config, sr_fwd, sr_grad)
+
+
+class _FP4GemmMXPadFunction(torch.autograd.Function):
+    """Pure MXFP4 GEMM (E2M1_X2, block-32, NT) identical to ``_FP4GemmMXSRFunction``
+    but zero-pads the token dimension to a multiple of ``_AITER_FP4_K_MULTIPLE``
+    (256) so the AITER a4w4 f4gemm Wgrad stays numerically correct.
+
+    THE FIX for the Wan2.2 MXFP4 gradient explosion. The three GEMMs of a Linear
+    are, in AITER's ``(M, N, K)`` logical shape (K = contraction):
+      * forward  out = x @ W^T            -> K = in_features   (%256 for Wan: ok)
+      * dgrad    grad_x = grad_out @ W    -> K = out_features  (%256 for Wan: ok)
+      * wgrad    grad_W = grad_out^T @ x  -> K = tokens        (NOT %256: explodes)
+    Only the wgrad contracts over tokens, and tokens = seq*batch is not %256 for
+    the cross-attn key/value projections (text_seq 226 * batch). AITER returns
+    ~1e7x garbage there (content-independent; verified with random operands and by
+    zero-padding the real operand: %256 -> A~0.55, %32/%128-but-not-%256 -> A~6e7).
+    Padding the token dim to %256 (and slicing fwd/dgrad outputs back) keeps every
+    GEMM on the tuned, preshuffled AITER path and is exact. ``sr_fwd``/``sr_grad``
+    behave exactly as in ``_FP4GemmMXSRFunction`` (``off`` == stock recipe).
+    """
+
+    @staticmethod
+    def forward(ctx, a, b, out_dtype, config, sr_fwd, sr_grad):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            MXScalingRecipe,
+            check_mxfp4_support,
+        )
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        ok, reason = check_mxfp4_support()
+        assert ok, reason
+        fp4_dtype = torch.float4_e2m1fn_x2
+
+        # Zero-pad the token (row / contraction) dim up to a multiple of 256 so the
+        # token-contracting Wgrad GEMM stays in AITER's correct regime.
+        m = a.shape[0]
+        k_mult = _AITER_FP4_K_MULTIPLE
+        m_pad = ((m + k_mult - 1) // k_mult) * k_mult
+        a_p = torch.nn.functional.pad(a, (0, 0, 0, m_pad - m)) if m_pad != m else a
+
+        a_fp4, a_scale_inv, a_t_fp4, a_t_scale_inv = quantize_fp4_with_trans(
+            a_p,
+            fp4_dtype,
+            config.granularity,
+            block_size=config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=sr_fwd,
+                use_rht=True,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+        )
+        b_fp4, b_scale_inv, b_t_fp4, b_t_scale_inv = quantize_fp4_with_trans(
+            b,
+            fp4_dtype,
+            config.granularity,
+            block_size=config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=True,
+                use_sr=sr_fwd,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+                shuffle_out=enable_preshuffle(),
+            ),
+        )
+        out = gemm_fp4_impl(
+            a_fp4,
+            a_scale_inv,
+            False,
+            b_fp4,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        # Drop the padded rows so downstream sees the real token count.
+        if m_pad != m:
+            out = out[:m].contiguous()
+
+        ctx.save_for_backward(a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv)
+        ctx.out_dtype = out_dtype
+        ctx.config = config
+        ctx.sr_grad = sr_grad
+        ctx.m = m
+        ctx.m_pad = m_pad
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import MXScalingRecipe
+        from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import (
+            enable_preshuffle,
+            gemm_fp4_impl,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp4_with_trans
+
+        a_t_fp4, a_t_scale_inv, b_t_fp4, b_t_scale_inv = ctx.saved_tensors
+        fp4_dtype = torch.float4_e2m1fn_x2
+        grad_out = grad_out.view(grad_out.shape[0], -1)
+
+        # Pad grad_out's token dim to the SAME multiple used for a_t in forward so
+        # the wgrad contraction (over tokens) is %256; slice the dgrad back.
+        m, m_pad = ctx.m, ctx.m_pad
+        grad_p = (
+            torch.nn.functional.pad(grad_out, (0, 0, 0, m_pad - m))
+            if m_pad != m
+            else grad_out
+        )
+
+        (
+            grad_out_fp4,
+            grad_out_scale_inv,
+            grad_out_t_fp4,
+            grad_out_t_scale_inv,
+        ) = quantize_fp4_with_trans(
+            grad_p,
+            fp4_dtype,
+            ctx.config.granularity,
+            block_size=ctx.config.block_size,
+            scaling_recipe=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=ctx.sr_grad,
+                use_rht=False,
+                shuffle_scale=enable_preshuffle(),
+            ),
+            scaling_recipe_for_trans=MXScalingRecipe(
+                use_2d_block=False,
+                use_sr=ctx.sr_grad,
+                use_rht=True,
+                shuffle_scale=enable_preshuffle(),
+            ),
+        )
+        grad_a = gemm_fp4_impl(
+            grad_out_fp4,
+            grad_out_scale_inv,
+            False,
+            b_t_fp4,
+            b_t_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        if m_pad != m:
+            grad_a = grad_a[:m].contiguous()
+        grad_b = gemm_fp4_impl(
+            grad_out_t_fp4,
+            grad_out_t_scale_inv,
+            False,
+            a_t_fp4,
+            a_t_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.config.granularity.value,
+            default_backend=BackendType.HIPBLASLT.value,
+        )
+        return grad_a, grad_b, None, None, None, None
+
+
+def _gemm_fp4_pad(a, b, out_dtype, config, sr_fwd, sr_grad):
+    """Token-padded MXFP4 GEMM (AITER Wgrad-correctness fix; local seam)."""
+    return _FP4GemmMXPadFunction.apply(a, b, out_dtype, config, sr_fwd, sr_grad)
+
+
+def _resolve_fp4_config():
+    """Build the MXFP4 quant config.
+
+    primus_turbo 0.2.0 only supports MX_BLOCKWISE + E2M1_X2 + block_size 32 +
+    E8M0 scale for FP4 (enforced by ``Float4QuantConfig.__post_init__``), so
+    there are no env knobs here; the defaults already select that recipe.
+    """
+    from primus_turbo.pytorch.core.low_precision import Float4QuantConfig
+
+    return Float4QuantConfig()
+
+
+class Float4Linear(torch.nn.Linear):
+    """MXFP4 Linear mirroring ``primus_turbo.Float8Linear`` but using ``gemm_fp4``.
+
+    Forward/backward quantize to MXFP4 (E2M1, block 32) inside the autograd-aware
+    ``gemm_fp4``; the bias is applied OUTSIDE the FP4 GEMM (in the compute dtype),
+    matching the FP8 module and avoiding any bias-in-GEMM kernel gap.
+    """
+
+    def __init__(self, in_features, out_features, bias=True, config=None, **kwargs):
+        super().__init__(in_features, out_features, bias, **kwargs)
+        self.config = config if config is not None else _resolve_fp4_config()
+        self.sr_scope = _fp4_sr_scope()
+        self.backward_precision = _fp4_backward_precision()
+        # Materialize the FP8 backward config for any mode that uses an FP8 leg.
+        self.fp8_bwd_config = (
+            _resolve_fp4_hybrid_fp8_config()
+            if self.backward_precision in ("fp8", "dgrad_fp8", "wgrad_fp8")
+            else None
+        )
+        self._warned_token_align = False
+
+    def forward(self, x):
+        flatten_shapes = x.shape[:-1]
+        x_2d = x.reshape(-1, self.in_features)
+
+        # Token dim M is a runtime property; the HIPBLASLT FP4 GEMM needs M % 128
+        # for the (transposed) backward GEMMs. Warn once instead of silently
+        # falling back to an untuned AITER path. Skip entirely under torch.compile
+        # (the flag mutation + logging would otherwise force a graph break).
+        # ``mxfp4_pad`` zero-pads the token/contraction dim to %256 internally, so
+        # the misalignment the warning is about is already handled -- skip it.
+        _compiling = getattr(getattr(torch, "compiler", None), "is_compiling", lambda: False)()
+        if (
+            not _compiling
+            and self.backward_precision != "mxfp4_pad"
+            and not self._warned_token_align
+            and x_2d.shape[0] % _FP4_ALIGN != 0
+        ):
+            logger.warning(
+                "[PrimusTurbo-MXFP4] token dim M=%d is not a multiple of %d; the "
+                "HIPBLASLT FP4 backward GEMM may fall back or error for this layer "
+                "(in=%d out=%d).",
+                x_2d.shape[0],
+                _FP4_ALIGN,
+                self.in_features,
+                self.out_features,
+            )
+            self._warned_token_align = True
+
+        if self.backward_precision == "mxfp4_pad":
+            # THE FIX: pure MXFP4 fwd+bwd, but zero-pad the token/contraction dim
+            # to %256 so the AITER a4w4 Wgrad (contracts over tokens) is correct.
+            out_2d = _gemm_fp4_pad(
+                x_2d,
+                self.weight,
+                x.dtype,
+                self.config,
+                self.sr_scope == "all",  # sr_fwd
+                self.sr_scope != "off",  # sr_grad
+            )
+        elif self.backward_precision == "fp8":
+            # Hybrid: MXFP4 forward, tensorwise-FP8 backward (mirrors the FLUX
+            # mxfp4_backward_precision='fp8'). SR (fwd) still honored.
+            out_2d = _gemm_fp4_hybrid(
+                x_2d,
+                self.weight,
+                x.dtype,
+                self.config,
+                self.fp8_bwd_config,
+                self.sr_scope == "all",  # sr_fwd
+            )
+        elif self.backward_precision in ("dgrad_fp8", "wgrad_fp8"):
+            # Isolation: one backward GEMM in FP8, the other in MXFP4.
+            out_2d = _gemm_fp4_mixed(
+                x_2d,
+                self.weight,
+                x.dtype,
+                self.config,
+                self.fp8_bwd_config,
+                self.sr_scope == "all",  # sr_fwd
+                self.sr_scope != "off",  # sr_grad (grad|all -> SR on the MXFP4 leg)
+                self.backward_precision == "dgrad_fp8",  # dgrad_fp8
+                self.backward_precision == "wgrad_fp8",  # wgrad_fp8
+            )
+        elif self.sr_scope == "off":
+            # Stock path (SR nowhere) -- bit-identical to primus_turbo baseline.
+            from primus_turbo.pytorch.ops import gemm_fp4
+
+            out_2d = gemm_fp4(
+                x_2d, self.weight, trans_a=False, trans_b=True, out_dtype=x.dtype, config=self.config
+            )
+        else:
+            # SR-enabled local path: grad -> SR on backward grad quant only;
+            # all -> SR on forward (act+weight) too.
+            out_2d = _gemm_fp4_sr(
+                x_2d,
+                self.weight,
+                x.dtype,
+                self.config,
+                self.sr_scope == "all",  # sr_fwd
+                True,  # sr_grad (both grad and all enable backward SR)
+            )
+        if self.bias is not None:
+            out_2d = out_2d + self.bias
+        return out_2d.view(*flatten_shapes, self.out_features)
+
+    def extra_repr(self):
+        return (
+            f"{super().extra_repr()}, mxfp4=True, sr={self.sr_scope}, "
+            f"backward={self.backward_precision}"
+        )
+
+
+def _is_fp4_training_safe_linear(name: str, linear: nn.Linear) -> bool:
+    """Whether a Linear is safe to run in MXFP4.
+
+    Builds on AutoModel's own FP8-safety predicate (time/text embed, norm-out,
+    norm-modulation, 16-alignment) and adds the stricter 128-alignment the
+    HIPBLASLT FP4 GEMM requires on in_features/out_features.
+    """
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import (
+        _is_fp8_training_safe_linear,
+    )
+
+    if not _is_fp8_training_safe_linear(name, linear):
+        return False
+    # Wan-aware: keep timestep/text/AdaLN conditioning projections in bf16.
+    if _fp4_keep_sensitive_bf16() and name.startswith(_FP4_SENSITIVE_PREFIXES):
+        return False
+    if linear.weight.shape[0] % _FP4_ALIGN != 0 or linear.weight.shape[1] % _FP4_ALIGN != 0:
+        return False
+    return True
+
+
+def _replace_linear_with_primus_turbo(
+    module: nn.Module,
+    module_name: str,
+    *,
+    fp8_safe_only: bool = False,  # accepted for signature-compat; we always skip-list
+) -> int:
+    """Drop-in replacement for AutoModel's TE swap using primus_turbo.Float8Linear."""
+    from primus_turbo.pytorch.modules import Float8Linear
+
+    # Reuse AutoModel's own fp8-safety predicate so the kept-in-bf16 set matches
+    # the TE path exactly (time/text embed, norm modulation, final proj, non-16).
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import (
+        _is_fp8_training_safe_linear,
+    )
+
+    cfg = _resolve_fp8_config()
+    converted = 0
+    skipped = 0
+
+    def replace_children(parent: nn.Module, prefix: str = "") -> None:
+        nonlocal converted, skipped
+        for child_name, child in list(parent.named_children()):
+            child_fqn = f"{prefix}.{child_name}" if prefix else child_name
+            if isinstance(child, Float8Linear):
+                continue
+            if isinstance(child, nn.Linear):
+                # Always apply the fp8-safe skip-list (stability). The recipe ties
+                # its fp8_safe_only flag to te_fp8, which we keep off, so ignore it.
+                if not _is_fp8_training_safe_linear(child_fqn, child):
+                    skipped += 1
+                    logger.info(
+                        "[PrimusTurbo-FP8] Keeping %s.%s as torch.nn.Linear (fp8-unsafe); weight=%s",
+                        module_name,
+                        child_fqn,
+                        tuple(child.weight.shape),
+                    )
+                    continue
+                fl = Float8Linear(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    config=cfg,
+                    device=child.weight.device,
+                    dtype=child.weight.dtype,
+                )
+                fl.train(child.training)
+                with torch.no_grad():
+                    fl.weight.copy_(child.weight)
+                    fl.weight.requires_grad_(child.weight.requires_grad)
+                    if child.bias is not None:
+                        fl.bias.copy_(child.bias)
+                        fl.bias.requires_grad_(child.bias.requires_grad)
+                setattr(parent, child_name, fl)
+                converted += 1
+            else:
+                replace_children(child, child_fqn)
+
+    replace_children(module)
+    logger.info(
+        "[PrimusTurbo-FP8] Replaced %d torch.nn.Linear with primus_turbo.Float8Linear in %s; "
+        "skipped=%d (granularity=%s format=%s)",
+        converted,
+        module_name,
+        skipped,
+        getattr(cfg.granularity, "name", cfg.granularity),
+        getattr(cfg.format, "name", cfg.format),
+    )
+    return converted
+
+
+def _copy_linear_params(dst: nn.Linear, src: nn.Linear) -> None:
+    """Copy weight/bias + requires_grad + training flag from ``src`` to ``dst``."""
+    dst.train(src.training)
+    with torch.no_grad():
+        dst.weight.copy_(src.weight)
+        dst.weight.requires_grad_(src.weight.requires_grad)
+        if src.bias is not None and dst.bias is not None:
+            dst.bias.copy_(src.bias)
+            dst.bias.requires_grad_(src.bias.requires_grad)
+
+
+def _replace_linear_with_primus_turbo_fp4(
+    module: nn.Module,
+    module_name: str,
+    *,
+    fp8_safe_only: bool = False,  # accepted for signature-compat; we always skip-list
+) -> int:
+    """Drop-in replacement for AutoModel's TE swap using the local Float4Linear.
+
+    With ``PRIMUS_TURBO_FP4_SENSITIVE_LAYERS`` on, a band of transformer blocks
+    (first START + last END) is kept in a higher precision (tensorwise FP8 via
+    primus_turbo.Float8Linear, or bf16) instead of MXFP4 -- mirroring the
+    ``sensitive_layers_*`` scaffolding that made the MXFP4 ablation
+    converge. Everything else (fp8-safe skip-list + 128-alignment) is unchanged.
+    """
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import (
+        _is_fp8_training_safe_linear,
+    )
+
+    cfg = _resolve_fp4_config()
+    sensitive_on = _fp4_sensitive_enabled()
+    block_counts = _count_blocks(module) if sensitive_on else {}
+    sens_precision = _fp4_sensitive_precision() if sensitive_on else None
+    fp8_cfg = None
+    Float8Linear = None
+    if sensitive_on and sens_precision == "tw_fp8":
+        from primus_turbo.pytorch.modules import Float8Linear  # noqa: F811
+
+        fp8_cfg = _resolve_fp4_hybrid_fp8_config()  # tensorwise FP8
+
+    converted = 0  # -> Float4Linear (MXFP4)
+    converted_fp8 = 0  # -> Float8Linear (sensitive band, tw_fp8)
+    skipped = 0  # -> kept torch.nn.Linear (bf16)
+
+    def replace_children(parent: nn.Module, prefix: str = "") -> None:
+        nonlocal converted, converted_fp8, skipped
+        for child_name, child in list(parent.named_children()):
+            child_fqn = f"{prefix}.{child_name}" if prefix else child_name
+            if isinstance(child, (Float4Linear,) + ((Float8Linear,) if Float8Linear else ())):
+                continue
+            if isinstance(child, nn.Linear):
+                # Sensitive band: keep first START + last END blocks out of MXFP4.
+                if sensitive_on and _is_sensitive_block(child_fqn, block_counts):
+                    if sens_precision == "tw_fp8" and _is_fp8_training_safe_linear(child_fqn, child):
+                        fl = Float8Linear(
+                            child.in_features,
+                            child.out_features,
+                            bias=child.bias is not None,
+                            config=fp8_cfg,
+                            device=child.weight.device,
+                            dtype=child.weight.dtype,
+                        )
+                        _copy_linear_params(fl, child)
+                        setattr(parent, child_name, fl)
+                        converted_fp8 += 1
+                        logger.info(
+                            "[PrimusTurbo-MXFP4] Sensitive band -> tw_fp8 (Float8Linear): %s.%s weight=%s",
+                            module_name,
+                            child_fqn,
+                            tuple(child.weight.shape),
+                        )
+                    else:
+                        skipped += 1
+                        logger.info(
+                            "[PrimusTurbo-MXFP4] Sensitive band -> bf16 (kept nn.Linear): %s.%s weight=%s",
+                            module_name,
+                            child_fqn,
+                            tuple(child.weight.shape),
+                        )
+                    continue
+                # FP8-safe skip-list + stricter 128-alignment for the FP4 GEMM.
+                if not _is_fp4_training_safe_linear(child_fqn, child):
+                    skipped += 1
+                    logger.info(
+                        "[PrimusTurbo-MXFP4] Keeping %s.%s as torch.nn.Linear "
+                        "(fp4-unsafe / not 128-aligned); weight=%s",
+                        module_name,
+                        child_fqn,
+                        tuple(child.weight.shape),
+                    )
+                    continue
+                fl = Float4Linear(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    config=cfg,
+                    device=child.weight.device,
+                    dtype=child.weight.dtype,
+                )
+                _copy_linear_params(fl, child)
+                setattr(parent, child_name, fl)
+                converted += 1
+            else:
+                replace_children(child, child_fqn)
+
+    replace_children(module)
+    if sensitive_on:
+        band = ", ".join(f"{k}[first {_fp4_sensitive_start()}+last {_fp4_sensitive_end()}/{v}]"
+                         for k, v in block_counts.items()) or "none-found"
+        logger.info(
+            "[PrimusTurbo-MXFP4] Replaced %d Float4Linear + %d %s (sensitive band: %s) in %s; "
+            "skipped(bf16)=%d (MX_BLOCKWISE E2M1_X2 block=%d sr=%s backward=%s)",
+            converted,
+            converted_fp8,
+            "Float8Linear(tw_fp8)" if sens_precision == "tw_fp8" else "bf16",
+            band,
+            module_name,
+            skipped,
+            getattr(cfg, "block_size", 32),
+            _fp4_sr_scope(),
+            _fp4_backward_precision(),
+        )
+    else:
+        logger.info(
+            "[PrimusTurbo-MXFP4] Replaced %d torch.nn.Linear with Float4Linear in %s; "
+            "skipped=%d (MX_BLOCKWISE E2M1_X2 block=%d sr=%s backward=%s)",
+            converted,
+            module_name,
+            skipped,
+            getattr(cfg, "block_size", 32),
+            _fp4_sr_scope(),
+            _fp4_backward_precision(),
+        )
+    return converted + converted_fp8
+
+
+# --------------------------------------------------------------------------- #
+# TE-native MXFP4 (te.pytorch.Linear under an MXFP4BlockScaling autocast)      #
+# --------------------------------------------------------------------------- #
+# This path is the reason for the TE-native image. It does two things the
+# primus_turbo path cannot on the baseline image:
+#   1. swaps nn.Linear -> transformer_engine.pytorch.Linear (128-aligned, same
+#      skip-list), and
+#   2. wraps each swapped transformer's forward in an MXFP4 autocast so the TE
+#      Linears actually quantize (a bare te.Linear runs bf16 outside autocast).
+# TE's MXFP4 recipe applies RHT (Hadamard) + the native fused HIP cast/transpose
+# kernel and dispatches the GEMM to AITER a4w4 -- a converging configuration.
+
+
+def _build_te_mxfp4_recipe():
+    """Construct the TE MXFP4 recipe.
+
+    Uses ``transformer_engine.common.recipe.MXFP4BlockScaling`` (wired to
+    ``Recipe.mxfp4()``). RHT (the convergence lever) is a per-instance
+    attribute driven by ``NVTE_MXFP4_USE_HADAMARD`` -- matching the reference
+    ``get_fp4_recipe`` (fp4_utils.py), which sets ``recipe.use_hadamard`` the same
+    way. Verified on the TE-native image: the recipe builds with use_hadamard=True
+    and the fwd/bwd dispatch to AITER a4w4 f4gemm.
+    """
+    from transformer_engine.common.recipe import MXFP4BlockScaling
+
+    recipe = MXFP4BlockScaling()
+    recipe.use_hadamard = os.getenv("NVTE_MXFP4_USE_HADAMARD", "0") in _TRUTHY
+    return recipe
+
+
+def _te_mxfp4_autocast(recipe):
+    """Return a TE autocast context for the given MXFP4 recipe.
+
+    Handles TE API drift: newer TE exposes ``te.pytorch.autocast(recipe=...)``;
+    the classic API is ``te.pytorch.fp8_autocast(fp8_recipe=...)``. We probe both.
+    """
+    import transformer_engine.pytorch as tep
+
+    if hasattr(tep, "autocast"):
+        try:
+            return tep.autocast(enabled=True, recipe=recipe)
+        except TypeError:
+            pass
+    try:
+        return tep.fp8_autocast(enabled=True, recipe=recipe)
+    except TypeError:
+        return tep.fp8_autocast(enabled=True, fp8_recipe=recipe)
+
+
+def _wrap_forward_with_te_mxfp4_autocast(module: nn.Module) -> None:
+    """Wrap ``module.forward`` so every call runs under the MXFP4 autocast.
+
+    Bound on the module INSTANCE before FSDP2 wrapping. FSDP2 (fully_shard) keeps
+    the module's ``forward`` and drives it through hooks, so the autocast is
+    entered on each real forward. Idempotent (guards against double-wrap).
+    """
+    if getattr(module, "_te_mxfp4_wrapped", False):
+        return
+    recipe = _build_te_mxfp4_recipe()
+    orig_forward = module.forward
+
+    def forward(*args, **kwargs):
+        with _te_mxfp4_autocast(recipe):
+            return orig_forward(*args, **kwargs)
+
+    module.forward = forward  # type: ignore[method-assign]
+    module._te_mxfp4_wrapped = True  # type: ignore[attr-defined]
+
+
+# Transformer-block ModuleList attribute names to wrap for autocast (Wan uses
+# ``blocks``; FLUX/diffusers use ``transformer_blocks`` AND, for the single-stream
+# stack, ``single_transformer_blocks`` -- BOTH must be wrapped or the 38 single
+# blocks run their TE Linears outside the MXFP4 autocast).
+_TE_BLOCK_LIST_ATTRS = ("blocks", "transformer_blocks", "single_transformer_blocks")
+
+
+def _wrap_te_mxfp4_autocast_for_ac(module: nn.Module) -> str:
+    """Enter the MXFP4 autocast at the *transformer-block* level, not the top.
+
+    CRITICAL for activation checkpointing. AutoModel wraps each block in a
+    ``NO_REENTRANT`` ``checkpoint_wrapper`` (parallelizer.py), so on the backward
+    recompute the block's forward reruns *standalone*, OUTSIDE any autocast entered
+    at the top-level transformer forward. If the autocast is only on the top-level
+    forward, the recompute runs in bf16 (not MXFP4) -> the checkpointed region
+    saves a different number of tensors on recompute than on the original forward
+    -> ``CheckpointError`` (and, if forced through, silently wrong grads).
+
+    Wrapping each block's forward puts the autocast INSIDE the checkpoint boundary,
+    so it is re-entered identically on recompute. All TE Linears live inside the
+    blocks (the skip-list keeps conditioning/proj Linears in bf16), so per-block
+    coverage is complete. Falls back to whole-module wrap if no block list found.
+    Returns a short tag describing what was wrapped (for logging).
+    """
+    tags = []
+    for attr in _TE_BLOCK_LIST_ATTRS:
+        blocks = getattr(module, attr, None)
+        if blocks is not None and len(blocks) > 0:
+            for blk in blocks:
+                _wrap_forward_with_te_mxfp4_autocast(blk)
+            tags.append(f"{attr}[{len(blocks)}]")
+    if tags:
+        # Wrap ALL block lists (e.g. FLUX has transformer_blocks AND
+        # single_transformer_blocks) -- do NOT return after the first match.
+        return "per-block:" + "+".join(tags)
+    _wrap_forward_with_te_mxfp4_autocast(module)
+    return "whole-module (no block list found)"
+
+
+def _replace_linear_with_te_mxfp4(
+    module: nn.Module,
+    module_name: str,
+    *,
+    fp8_safe_only: bool = False,  # accepted for signature-compat; we always skip-list
+) -> int:
+    """Drop-in replacement for AutoModel's TE swap using native te.pytorch.Linear.
+
+    Swaps 128-aligned, fp8/fp4-safe nn.Linear to ``transformer_engine.pytorch
+    .Linear`` (reusing ``_is_fp4_training_safe_linear``), copies weights/bias and
+    requires_grad, then wraps the module forward in the MXFP4 autocast.
+    """
+    from transformer_engine.pytorch import Linear as TELinear
+
+    converted = 0
+    skipped = 0
+
+    def replace_children(parent: nn.Module, prefix: str = "") -> None:
+        nonlocal converted, skipped
+        for child_name, child in list(parent.named_children()):
+            child_fqn = f"{prefix}.{child_name}" if prefix else child_name
+            if isinstance(child, TELinear):
+                continue
+            if isinstance(child, nn.Linear):
+                if not _is_fp4_training_safe_linear(child_fqn, child):
+                    skipped += 1
+                    logger.info(
+                        "[TE-MXFP4] Keeping %s.%s as torch.nn.Linear "
+                        "(fp4-unsafe / not 128-aligned); weight=%s",
+                        module_name,
+                        child_fqn,
+                        tuple(child.weight.shape),
+                    )
+                    continue
+                tl = TELinear(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    params_dtype=child.weight.dtype,
+                    device=child.weight.device,
+                )
+                tl.train(child.training)
+                with torch.no_grad():
+                    tl.weight.copy_(child.weight)
+                    tl.weight.requires_grad_(child.weight.requires_grad)
+                    if child.bias is not None:
+                        tl.bias.copy_(child.bias)
+                        tl.bias.requires_grad_(child.bias.requires_grad)
+                setattr(parent, child_name, tl)
+                converted += 1
+            else:
+                replace_children(child, child_fqn)
+
+    replace_children(module)
+    wrap_tag = "none"
+    if converted:
+        wrap_tag = _wrap_te_mxfp4_autocast_for_ac(module)
+    logger.info(
+        "[TE-MXFP4] Replaced %d torch.nn.Linear with te.pytorch.Linear in %s; "
+        "skipped=%d (MXFP4BlockScaling autocast wrapped=%s)",
+        converted,
+        module_name,
+        skipped,
+        wrap_tag,
+    )
+    return converted
+
+
+def install() -> bool:
+    """Install the swap by monkeypatching AutoModel's TE swap symbol.
+
+    Precedence: TE-native MXFP4 (``PRIMUS_TE_MXFP4``) > primus_turbo MXFP4
+    (``PRIMUS_TURBO_FP4``) > primus_turbo FP8 (``PRIMUS_TURBO_FP8``); else no-op
+    (returns False). Modifies NO Automodel source; only rebinds a module
+    attribute at runtime. Call this BEFORE building the recipe/pipeline.
+    """
+    if is_te_mxfp4_enabled():
+        # Ensure the running TE can build the MXFP4 recipe. NOTE: we intentionally
+        # do NOT hard-depend on ``check_mxfp4_support`` -- on some TE-native
+        # builds that symbol is NOT in transformer_engine
+        # .pytorch.fp8 (that module is deprecated), yet MXFP4 executes fine
+        # (verified: fwd/bwd dispatch to AITER a4w4 f4gemm, finite grads). So the
+        # recipe-build is the real gate; the support check is best-effort.
+        _build_te_mxfp4_recipe()  # raises if MXFP4BlockScaling is unavailable (baseline TE)
+        _reason = None
+        for _mod in ("transformer_engine.pytorch", "primus_turbo.pytorch.core.low_precision"):
+            try:
+                _m = __import__(_mod, fromlist=["check_mxfp4_support"])
+                _chk = getattr(_m, "check_mxfp4_support", None)
+                if _chk is not None:
+                    _ok, _reason = _chk()
+                    if not _ok:
+                        raise RuntimeError(f"TE/Turbo reports MXFP4 unsupported: {_reason}")
+                    break
+            except ImportError:
+                continue
+        if _reason is None:
+            logger.warning(
+                "[TE-MXFP4] check_mxfp4_support() not found in TE or primus_turbo; "
+                "proceeding (recipe builds). Ensure this is the TE-native image (gfx950)."
+            )
+
+        import nemo_automodel._diffusers.auto_diffusion_pipeline as adp
+
+        adp._replace_linear_with_transformer_engine = _replace_linear_with_te_mxfp4
+        logger.info(
+            "[TE-MXFP4] Installed: nn.Linear -> te.pytorch.Linear + MXFP4BlockScaling "
+            "autocast (triggered by model.transformer_engine_linear=true)"
+        )
+        return True
+
+    if is_fp4_enabled():
+        # Fail fast if the FP4 op is unavailable so the run errors clearly.
+        from primus_turbo.pytorch.ops import gemm_fp4  # noqa: F401
+
+        import nemo_automodel._diffusers.auto_diffusion_pipeline as adp
+
+        adp._replace_linear_with_transformer_engine = _replace_linear_with_primus_turbo_fp4
+        logger.info(
+            "[PrimusTurbo-MXFP4] Installed: nn.Linear -> Float4Linear (gemm_fp4) swap active "
+            "(triggered by model.transformer_engine_linear=true)"
+        )
+        return True
+
+    if not is_enabled():
+        return False
+
+    # Fail fast if primus_turbo is unavailable so the run errors clearly rather
+    # than silently falling back to TE.
+    import primus_turbo.pytorch.modules  # noqa: F401
+
+    import nemo_automodel._diffusers.auto_diffusion_pipeline as adp
+
+    adp._replace_linear_with_transformer_engine = _replace_linear_with_primus_turbo
+    logger.info(
+        "[PrimusTurbo-FP8] Installed: nn.Linear -> primus_turbo.Float8Linear swap active "
+        "(triggered by model.transformer_engine_linear=true)"
+    )
+    return True
diff --git a/primus/backends/nemo_automodel/quantization/primus_turbo_fp8_attn.py b/primus/backends/nemo_automodel/quantization/primus_turbo_fp8_attn.py
new file mode 100644
index 000000000..eb114db14
--- /dev/null
+++ b/primus/backends/nemo_automodel/quantization/primus_turbo_fp8_attn.py
@@ -0,0 +1,284 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Primus-Turbo FP8 flash-attention override for the NeMo AutoModel diffusion
+recipe (FP8 *attention*, orthogonal to the FP8/MXFP4 GEMM swap).
+
+WHY:
+  The GEMM low-precision hook (primus_turbo_fp8.py) swaps nn.Linear -> FP8/MXFP4
+  but leaves scaled-dot-product attention in bf16. Attention is the dominant cost
+  on the long-sequence Wan hero (61,200 tokens), so an FP8 attention kernel is the
+  biggest remaining lever there. Primus-Turbo ships ``flash_attn_fp8_func`` (E4M3,
+  block-scaled q/k/v, autograd fwd+bwd) -- the SAME kernel family as the GEMM path.
+  There is NO FP4 attention kernel; FP8 attention is the only option and pairs with
+  any GEMM precision (BF16 / FP8 / MXFP4).
+
+WHAT this does (NO diffusers / Automodel fork):
+  diffusers routes every model's attention through ``_AttentionBackendRegistry``:
+  the config's ``model.attention_backend`` name selects a registered backend fn,
+  and ``dispatch_attention_fn`` looks it up in ``_backends[name]`` at call time.
+  This hook rebinds the backend *function* for the two backends our heroes use --
+  FLASH (FLUX, ``flash_attn_func``) and AITER (Wan, ``aiter_flash_attn_func``) --
+  to an FP8 flash-attn wrapper, while leaving the backend NAME, the registered
+  ``_supported_arg_names`` and the constraint list untouched. So the config still
+  says ``attention_backend: flash`` (FLUX) / ``aiter`` (Wan); only the numerics of
+  the kernel change. Mirrors the nn.Linear->Float8Linear GEMM hook's philosophy.
+
+  Env-gated by ``PRIMUS_TURBO_FP8_ATTN=1`` (default off = no-op). Install BEFORE the
+  recipe builds the transformer (i.e. before ``set_attention_backend``); the swap is
+  a module-global dict entry resolved at forward, so as long as it is in place before
+  the first forward it takes effect.
+
+PAD-TO-64:
+  The FP8 flash kernel block-scales q/k/v along the sequence with block_size=64, so
+  it requires ``seqlen % 64 == 0``. FLUX joint attn (4608) and Wan cross-attn kv
+  (256) are already %64, but Wan self-attn seqlen is 61,200 (= 64*956 + 16, NOT %64)
+  -> the raw kernel fails with a reshape error. We zero-pad q/k/v's sequence up to
+  the next multiple of 64, run FP8 attention, and slice the real queries back.
+  Padded-key dilution (padded keys are zeros -> each adds exp(0)=1 to the softmax
+  denominator) is negligible because n_pad <= 63 << seqlen, well within FP8's own
+  accuracy floor. No softmax correction is applied (an LSE-based one was tried and
+  dropped as unnecessary).
+
+FALLBACK (correctness first): the FP8 pad-64 kernel serves the plain, non-causal,
+  no-mask, dropout-free, non-context-parallel path -- exactly what the FLUX/Wan
+  hero forwards use with ``cp_size: 1``. For anything it cannot serve
+  (``_parallel_config`` set i.e. context parallelism, ``return_lse``, an additive
+  ``attn_mask``, nonzero dropout, or a sliding ``window_size``) it transparently
+  calls the ORIGINAL bf16 backend and warns ONCE per reason. So enabling FP8
+  attention never changes results for a call it doesn't handle -- e.g. a Wan run
+  bumped to ``cp_size>1`` for memory will run bf16 attention (revisit: a CP-aware
+  FP8 attention forward_op is future work).
+
+Activation (env, no config schema change):
+    PRIMUS_TURBO_FP8_ATTN=1    enable the FP8 attention override (default off = no-op)
+
+Launch: install() is called by the backend entrypoint when the env is set; the
+  config keeps ``attention_backend: flash`` (FLUX) / ``aiter`` (Wan) unchanged.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = {"1", "true", "True", "yes", "on"}
+
+# The FP8 flash kernel block-scales q/k/v along the sequence with this block size,
+# so the sequence length must be a multiple of it (see module docstring).
+_ATTN_BLOCK = 64
+
+
+def is_fp8_attn_enabled() -> bool:
+    """Whether the Primus-Turbo FP8 attention override should be installed."""
+    return os.getenv("PRIMUS_TURBO_FP8_ATTN", "0") in _TRUTHY
+
+
+# --------------------------------------------------------------------------- #
+# Pad-to-64 FP8 flash attention                                               #
+# --------------------------------------------------------------------------- #
+def _pad_seq(t: torch.Tensor, mult: int = _ATTN_BLOCK) -> tuple[torch.Tensor, int]:
+    """Zero-pad a (B, S, H, D) tensor's seq dim (dim=1) up to a multiple of ``mult``.
+
+    Returns (padded_tensor, real_seqlen).
+    """
+    s = t.shape[1]
+    s_pad = ((s + mult - 1) // mult) * mult
+    if s_pad == s:
+        return t, s
+    # F.pad pads from the last dim backwards: (D_lo, D_hi, H_lo, H_hi, S_lo, S_hi).
+    return torch.nn.functional.pad(t, (0, 0, 0, 0, 0, s_pad - s)), s
+
+
+def _flash_attn_fp8_pad64(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    """FP8 flash attention with automatic pad-to-64 of a non-%64 sequence.
+
+    q: (B, Sq, H, D), k/v: (B, Skv, H, D). Returns (B, Sq, H, D). Zero-pads q/k/v'
+    sequence dims up to a multiple of 64 (the kernel's block-scale requirement) and
+    slices the real queries back. Padded-key dilution is negligible (module docstring).
+    """
+    from primus_turbo.pytorch.ops import flash_attn_fp8_func
+
+    b, sq, h, d = q.shape
+    if softmax_scale is None:
+        softmax_scale = d ** -0.5
+
+    qp, sq_real = _pad_seq(q)
+    kp, _ = _pad_seq(k)
+    vp, _ = _pad_seq(v)
+
+    out = flash_attn_fp8_func(qp, kp, vp, softmax_scale=softmax_scale, causal=causal)
+    return out[:, :sq_real].contiguous()
+
+
+# --------------------------------------------------------------------------- #
+# Backend override                                                            #
+# --------------------------------------------------------------------------- #
+_warned: set[str] = set()
+
+
+def _warn_once(key: str, msg: str) -> None:
+    if key not in _warned:
+        _warned.add(key)
+        logger.warning(msg)
+
+
+def _make_fp8_attn_backend(orig_fn, backend_value: str, supports_window: bool):
+    """Build the FP8-attention backend fn that replaces ``orig_fn`` in the registry.
+
+    Signature is a superset of both the FLASH and AITER backend signatures so that
+    ``dispatch_attention_fn`` (which filters kwargs to each backend's ORIGINAL
+    ``_supported_arg_names`` and then calls ``backend_fn(**kwargs)``) can drive it
+    unchanged. ``orig_fn`` is captured for the fallback path; ``supports_window``
+    tells us whether ``orig_fn`` accepts ``window_size`` (FLASH does, AITER doesn't).
+    """
+
+    def _fallback(query, key, value, attn_mask, dropout_p, is_causal, scale,
+                  window_size, return_lse, _parallel_config):
+        kw = dict(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+            scale=scale,
+            return_lse=return_lse,
+            _parallel_config=_parallel_config,
+        )
+        if supports_window:
+            kw["window_size"] = window_size
+        return orig_fn(**kw)
+
+    def fp8_attn_backend(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: Optional[float] = None,
+        window_size: tuple[int, int] = (-1, -1),
+        return_lse: bool = False,
+        _parallel_config=None,
+    ):
+        # Correctness-first fallbacks: use the original bf16 kernel for anything the
+        # FP8 pad-64 path cannot faithfully serve (warn once per reason).
+        if _parallel_config is not None:
+            _warn_once(
+                f"{backend_value}:cp",
+                f"[PrimusTurbo-FP8Attn] context parallelism active on backend "
+                f"'{backend_value}' (no FP8 CP path yet) -> bf16 attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if return_lse:
+            _warn_once(
+                f"{backend_value}:lse",
+                f"[PrimusTurbo-FP8Attn] return_lse requested on backend "
+                f"'{backend_value}' (FP8 pad-64 has no LSE) -> bf16 attention fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if attn_mask is not None:
+            _warn_once(
+                f"{backend_value}:mask",
+                f"[PrimusTurbo-FP8Attn] additive attn_mask on backend "
+                f"'{backend_value}' unsupported by FP8 flash-attn -> bf16 fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if dropout_p:
+            _warn_once(
+                f"{backend_value}:dropout",
+                f"[PrimusTurbo-FP8Attn] dropout_p={dropout_p} on backend "
+                f"'{backend_value}' unsupported by FP8 flash-attn -> bf16 fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+        if window_size not in ((-1, -1), None):
+            _warn_once(
+                f"{backend_value}:window",
+                f"[PrimusTurbo-FP8Attn] sliding window_size={window_size} on backend "
+                f"'{backend_value}' unsupported by FP8 flash-attn -> bf16 fallback.",
+            )
+            return _fallback(query, key, value, attn_mask, dropout_p, is_causal,
+                             scale, window_size, return_lse, _parallel_config)
+
+        return _flash_attn_fp8_pad64(query, key, value, softmax_scale=scale, causal=is_causal)
+
+    fp8_attn_backend._primus_fp8_attn = True  # idempotence marker
+    fp8_attn_backend._primus_orig_fn = orig_fn
+    return fp8_attn_backend
+
+
+# Backends our diffusion heroes use: FLUX -> FLASH (accepts window_size),
+# Wan -> AITER (no window_size). Override both; only the active one is dispatched.
+_TARGET_BACKENDS = (("FLASH", True), ("AITER", False))
+
+
+def install() -> bool:
+    """Install the FP8 attention override by rebinding diffusers backend functions.
+
+    No-op (returns False) unless ``PRIMUS_TURBO_FP8_ATTN`` is set. Modifies NO
+    diffusers/Automodel source; only rebinds ``_AttentionBackendRegistry._backends``
+    entries at runtime. Idempotent. Call BEFORE the transformer's first forward
+    (in practice before the recipe build / ``set_attention_backend``).
+    """
+    if not is_fp8_attn_enabled():
+        return False
+
+    # Fail fast if the FP8 attention op is unavailable so the run errors clearly
+    # rather than silently running bf16 attention.
+    from primus_turbo.pytorch.ops import flash_attn_fp8_func  # noqa: F401
+
+    from diffusers.models.attention_dispatch import (
+        AttentionBackendName,
+        _AttentionBackendRegistry,
+    )
+
+    reg = _AttentionBackendRegistry
+    installed = []
+    for name_str, supports_window in _TARGET_BACKENDS:
+        name = getattr(AttentionBackendName, name_str, None)
+        if name is None:
+            continue
+        orig = reg._backends.get(name)
+        if orig is None:
+            # Backend not registered in this diffusers build (e.g. aiter absent).
+            continue
+        if getattr(orig, "_primus_fp8_attn", False):
+            installed.append(name.value)  # already installed
+            continue
+        reg._backends[name] = _make_fp8_attn_backend(orig, name.value, supports_window)
+        # Leave reg._supported_arg_names[name] and reg._constraints[name] untouched:
+        # dispatch filters kwargs to the ORIGINAL supported set (our fn accepts a
+        # superset), and the qkv-bf16/device/shape checks still hold (inputs arrive
+        # bf16; the wrapper casts to FP8 internally).
+        installed.append(name.value)
+
+    if not installed:
+        logger.warning(
+            "[PrimusTurbo-FP8Attn] PRIMUS_TURBO_FP8_ATTN set but neither FLASH nor "
+            "AITER backend is registered in diffusers; FP8 attention NOT active."
+        )
+        return False
+
+    logger.info(
+        "[PrimusTurbo-FP8Attn] Installed FP8 flash-attention (flash_attn_fp8_func, "
+        "pad-to-64) override for backends: %s (config attention_backend name unchanged)",
+        ", ".join(installed),
+    )
+    return True
diff --git a/primus/backends/torchtitan/config_utils.py b/primus/backends/torchtitan/config_utils.py
index 39f212cc4..695dc8a49 100644
--- a/primus/backends/torchtitan/config_utils.py
+++ b/primus/backends/torchtitan/config_utils.py
@@ -20,8 +20,8 @@
 from types import SimpleNamespace
 from typing import Any
 
+from primus.core.utils.module_utils import log_rank_0
 from primus.core.utils.yaml_utils import dict_to_nested_namespace
-from primus.modules.module_utils import log_rank_0
 
 
 def build_job_config_from_namespace(ns: SimpleNamespace) -> Any:
diff --git a/primus/backends/torchtitan/models/deepseek_v3/__init__.py b/primus/backends/torchtitan/models/deepseek_v3/__init__.py
index 6ec1c8600..6ddf45f9e 100644
--- a/primus/backends/torchtitan/models/deepseek_v3/__init__.py
+++ b/primus/backends/torchtitan/models/deepseek_v3/__init__.py
@@ -33,7 +33,7 @@
         qk_rope_head_dim=64,
         v_head_dim=128,
         mscale=0.70,
-        use_flex_attn=True,
+        attn_type="flex",
         attn_mask_type="block_causal",
         q_head=16,
         n_kv_heads=16,
@@ -51,19 +51,19 @@
             num_experts=160,
             num_shared_experts=2,
             top_k=6,
+            num_expert_groups=8,
+            num_limited_groups=3,
             score_func="softmax",
             route_norm=True,
             route_scale=16.0,
             score_before_experts=False,
         ),
-        n_expert_groups=8,
-        n_limited_groups=3,
         q_lora_rank=1536,
         kv_lora_rank=512,
         qk_nope_head_dim=128,
         qk_rope_head_dim=64,
         v_head_dim=128,
-        use_flex_attn=True,
+        attn_type="flex",
         attn_mask_type="block_causal",
         q_head=40,
         n_kv_heads=8,
@@ -81,19 +81,19 @@
             num_experts=256,
             num_shared_experts=1,
             top_k=8,
+            num_expert_groups=8,
+            num_limited_groups=4,
             score_func="sigmoid",
             route_norm=True,
             route_scale=2.5,
             score_before_experts=False,
         ),
-        n_expert_groups=8,
-        n_limited_groups=4,
         q_lora_rank=1536,
         kv_lora_rank=512,
         qk_nope_head_dim=128,
         qk_rope_head_dim=64,
         v_head_dim=128,
-        use_flex_attn=True,
+        attn_type="flex",
         attn_mask_type="block_causal",
         q_head=56,
         n_kv_heads=8,
diff --git a/primus/backends/torchtitan/models/deepseek_v3/model/model.py b/primus/backends/torchtitan/models/deepseek_v3/model/model.py
index ec587fe67..bb8797b30 100644
--- a/primus/backends/torchtitan/models/deepseek_v3/model/model.py
+++ b/primus/backends/torchtitan/models/deepseek_v3/model/model.py
@@ -26,6 +26,7 @@ def forward(
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
         attention_masks: AttentionMasksType | None,
+        positions: torch.Tensor | None = None,
     ):
         """
         Forward pass for the Multi-Head Latent Attention (MLA) Layer.
@@ -33,6 +34,7 @@ def forward(
         Args:
             x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
             freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+            positions (torch.Tensor | None): Position indices used to access/shuffle the RoPE cache.
 
         Returns:
             torch.Tensor: Output tensor with the same shape as the input.
@@ -50,14 +52,14 @@ def forward(
         # the above linear ops.
         q = q.view(bsz, seqlen, -1, self.qk_head_dim)
         q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
-        q_pe = apply_rotary_emb(q_pe, freqs_cis)
+        q_pe = apply_rotary_emb(q_pe, freqs_cis, positions)
         q = torch.cat([q_nope, q_pe], dim=-1)  # (bsz, seqlen, n_heads, qk_head_dim)
 
         # Key-value projection
         kv = self.wkv_a(x)  # (bsz, seqlen, kv_lora_rank + qk_rope_head_dim)
         kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
 
-        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)  # (bsz, seqlen, 1, qk_rope_head_dim)
+        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, positions)  # (bsz, seqlen, 1, qk_rope_head_dim)
 
         kv = self.wkv_b(self.kv_norm(kv))  # (bsz, seqlen, n_heads * (qk_nope_head_dim + v_head_dim))
         kv = kv.view(bsz, seqlen, -1, self.qk_nope_head_dim + self.v_head_dim)
@@ -99,7 +101,9 @@ def __init__(self, deepseek_args: DeepSeekV3ClassicModelArgs):
                 self.dim = deepseek_args.dim
                 self.head_dim = deepseek_args.head_dim
 
-                self.use_flex_attn = deepseek_args.use_flex_attn
+                # Upstream v0.2.2 Attention.__init__ dispatches on ``attn_type``
+                # (the legacy ``use_flex_attn`` flag was removed).
+                self.attn_type = deepseek_args.attn_type
 
         # Initialize the parent class with the mock args
         super().__init__(
@@ -112,10 +116,11 @@ def forward(
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
         attention_masks: AttentionMasksType | None,
+        positions: torch.Tensor | None = None,
     ):
         # Always use llama4-style freqs_cis for this attention, regardless of input
         seqlen = x.shape[1]
         freqs_llama4 = llama4_precompute_freqs_cis(self.head_dim, seqlen, self.rope_theta)
         # Ensure freqs are on the same device as activations
         freqs_llama4 = freqs_llama4.to(x.device, dtype=x.dtype)
-        return super().forward(x, freqs_llama4, None)
+        return super().forward(x, freqs_llama4, None, positions)
diff --git a/primus/backends/torchtitan/models/gpt_oss/__init__.py b/primus/backends/torchtitan/models/gpt_oss/__init__.py
new file mode 100644
index 000000000..095aa107e
--- /dev/null
+++ b/primus/backends/torchtitan/models/gpt_oss/__init__.py
@@ -0,0 +1,7 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Primus TorchTitan GPT-OSS backend extensions (Primus-Turbo sink attention)."""
diff --git a/primus/backends/torchtitan/models/gpt_oss/model/__init__.py b/primus/backends/torchtitan/models/gpt_oss/model/__init__.py
new file mode 100644
index 000000000..771e1b42f
--- /dev/null
+++ b/primus/backends/torchtitan/models/gpt_oss/model/__init__.py
@@ -0,0 +1,5 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
diff --git a/primus/backends/torchtitan/models/gpt_oss/model/model.py b/primus/backends/torchtitan/models/gpt_oss/model/model.py
new file mode 100644
index 000000000..b032e82d5
--- /dev/null
+++ b/primus/backends/torchtitan/models/gpt_oss/model/model.py
@@ -0,0 +1,105 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Primus-Turbo sink attention mirror for TorchTitan GPT-OSS.
+
+Upstream GPT-OSS attention (``torchtitan.models.gpt_oss.model.model.Attention``)
+uses ``FlexAttentionWrapper`` with:
+    * learnable per-head attention sinks (``self.sinks``), applied as a post-hoc
+      ``sigmoid(lse - sinks)`` rescaling of the FlexAttention output, and
+    * a sliding-window mask applied on even layers (``layer_id % 2 == 0``).
+
+This mirror keeps the exact same parameters (``wq/wk/wv/wo`` incl. biases and the
+learnable ``self.sinks``) but routes attention through the Primus-Turbo
+*functional* kernel ``primus_turbo.pytorch.ops.flash_attn_func``, which natively
+supports both:
+    * ``sink=`` (per-head learnable sinks; automatically dispatches to the Triton
+      backend), matching the Megatron ``PrimusTurboAttention`` semantics, and
+    * ``window_size=(left, 0)`` for the sliding window.
+
+Because ``self.sinks`` is inherited from the upstream module (not re-created), the
+learned sink weights are preserved on checkpoint load/save. The mirror is only
+installed when ``primus_turbo.enable_primus_turbo`` and
+``primus_turbo.use_turbo_attention`` are both set for a ``gpt_oss`` run (see
+``patches/turbo/gptoss_sink_attention_patches.py``); the default GPT-OSS path
+keeps upstream FlexAttention untouched.
+"""
+
+import torch
+from torchtitan.models.gpt_oss.model.model import Attention as TTGptOssAttention
+from torchtitan.models.gpt_oss.model.model import TransformerBlock as TTGptOssBlock
+from torchtitan.models.gpt_oss.model.model import apply_rotary_emb
+
+# Sentinel meaning "no sliding window" for primus_turbo's flash_attn_func.
+_FULL_WINDOW = (-1, -1)
+
+
+class Attention(TTGptOssAttention):
+    """GPT-OSS attention backed by Primus-Turbo ``flash_attn_func`` (sink-aware)."""
+
+    def __init__(self, model_args):
+        super().__init__(model_args)
+        # Per-layer sliding window is injected by the mirror TransformerBlock
+        # before each forward; default to full (causal, no window).
+        self.sliding_window_size = model_args.sliding_window_size
+        self._turbo_window = _FULL_WINDOW
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cache: torch.Tensor,
+        attention_masks=None,  # noqa: ARG002 - flex masks unused on the turbo path
+    ):
+        bsz, seqlen, _ = x.size()
+        hidden_shape = (bsz, seqlen, -1, self.head_dim)
+
+        q = self.wq(x).view(hidden_shape)
+        k = self.wk(x).view(hidden_shape)
+        v = self.wv(x).view(hidden_shape)
+
+        # RoPE is applied on the (b, s, h, d) layout, exactly like upstream.
+        q, k = apply_rotary_emb(q, k, rope_cache)
+
+        # primus_turbo.flash_attn_func consumes the (b, s, h, d) layout directly
+        # and handles GQA (n_heads > n_kv_heads) internally, so we do NOT
+        # transpose to (b, h, s, d) and do NOT repeat_kv. The learnable sinks are
+        # applied inside the kernel (Triton backend), so the upstream post-hoc
+        # sigmoid(lse - sinks) rescaling is not needed here.
+        import primus_turbo.pytorch as turbo
+
+        output = turbo.ops.flash_attn_func(
+            q,
+            k,
+            v,
+            softmax_scale=self.softmax_scale,
+            causal=True,
+            window_size=self._turbo_window,
+            sink=self.sinks.to(q.dtype),
+        )
+
+        output = output.reshape(bsz, seqlen, -1).contiguous()
+        return self.wo(output)
+
+
+class TransformerBlock(TTGptOssBlock):
+    """GPT-OSS block that selects the per-layer sliding window for turbo sink attn."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cache: torch.Tensor,
+        attention_masks=None,  # noqa: ARG002 - flex masks unused on the turbo path
+    ):
+        # gpt-oss pattern: even layers use sliding-window attention.
+        if self.use_sliding_attention:
+            self.attention._turbo_window = (self.attention.sliding_window_size, 0)
+        else:
+            self.attention._turbo_window = _FULL_WINDOW
+
+        x = x + self.attention(self.attention_norm(x), rope_cache, None)
+        x = x + self.moe(self.ffn_norm(x))
+        return x
diff --git a/primus/backends/torchtitan/models/llama3/model/model.py b/primus/backends/torchtitan/models/llama3/model/model.py
index daebde2fd..954b1bf51 100644
--- a/primus/backends/torchtitan/models/llama3/model/model.py
+++ b/primus/backends/torchtitan/models/llama3/model/model.py
@@ -18,6 +18,7 @@ def forward(
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
         attention_masks: AttentionMasksType | None,
+        positions: torch.Tensor | None = None,
     ):
         bs, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
@@ -29,12 +30,11 @@ def forward(
         xk = xk.view(bs, seqlen, -1, self.head_dim)
         xv = xv.view(bs, seqlen, -1, self.head_dim)
 
-        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
-
-        # repeat k/v heads if n_kv_heads < n_heads
-        # xk = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
-        # xv = repeat_kv(xv, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis, positions=positions)
 
+        # Primus-Turbo path: inner_attention is replaced by TurboAttention, which
+        # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA /
+        # causal masking internally, so we skip repeat_kv and the transpose.
         output = self.inner_attention(xq, xk, xv)
 
         output = output.contiguous().view(bs, seqlen, -1)
diff --git a/primus/backends/torchtitan/models/llama4/model/model.py b/primus/backends/torchtitan/models/llama4/model/model.py
index 7807fcbd6..5e80aeca6 100644
--- a/primus/backends/torchtitan/models/llama4/model/model.py
+++ b/primus/backends/torchtitan/models/llama4/model/model.py
@@ -18,6 +18,7 @@ def forward(
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
         attention_masks: AttentionMasksType | None,
+        positions: torch.Tensor | None = None,
     ):
         bs, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
@@ -30,12 +31,11 @@ def forward(
         xv = xv.view(bs, seqlen, -1, self.head_dim)
 
         if self.use_rope:
-            xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
-
-        # repeat k/v heads if n_kv_heads < n_heads
-        # xk = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
-        # xv = repeat_kv(xv, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
+            xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis, positions=positions)
 
+        # Primus-Turbo path: inner_attention is replaced by TurboAttention, which
+        # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA /
+        # causal masking internally, so we skip repeat_kv and the transpose.
         output = self.inner_attention(xq, xk, xv)
 
         output = output.contiguous().view(bs, seqlen, -1)
diff --git a/primus/backends/torchtitan/models/qwen3/model/model.py b/primus/backends/torchtitan/models/qwen3/model/model.py
index 632a33095..35f88bfda 100644
--- a/primus/backends/torchtitan/models/qwen3/model/model.py
+++ b/primus/backends/torchtitan/models/qwen3/model/model.py
@@ -18,6 +18,7 @@ def forward(
         x: torch.Tensor,
         rope_cache: torch.Tensor,
         attention_masks: AttentionMasksType | None,
+        positions: torch.Tensor | None = None,
     ):
         bs, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
@@ -34,14 +35,12 @@ def forward(
         if self.k_norm is not None:
             xk = self.k_norm(xk)
 
-        xq, xk = apply_rotary_emb(xq, xk, rope_cache)
+        xq, xk = apply_rotary_emb(xq, xk, rope_cache, positions)
 
-        if self.use_flex_attn:
-            assert isinstance(attention_masks, BlockMask), attention_masks
-            output = self.inner_attention(xq, xk, xv, block_mask=attention_masks)
-        else:
-            assert attention_masks is None
-            output = self.inner_attention(xq, xk, xv)
+        # Primus-Turbo path: inner_attention is replaced by TurboAttention, which
+        # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA /
+        # causal masking internally, so we skip the transpose and block_mask.
+        output = self.inner_attention(xq, xk, xv)
 
         output = output.contiguous().view(bs, seqlen, -1)
         return self.wo(output)
diff --git a/primus/backends/torchtitan/patches/__init__.py b/primus/backends/torchtitan/patches/__init__.py
index a70615372..b4c6e2d9a 100644
--- a/primus/backends/torchtitan/patches/__init__.py
+++ b/primus/backends/torchtitan/patches/__init__.py
@@ -29,6 +29,7 @@
     dcp_consolidate_patches,
     embedding_amp_patches,
     flex_attention_patches,
+    fsdp_weight_tying_patches,
     logger_patches,
     metrics_output_format,
     mock_dataset_patches,
diff --git a/primus/backends/torchtitan/patches/dcp_consolidate_patches.py b/primus/backends/torchtitan/patches/dcp_consolidate_patches.py
index 973a9be04..702e981c9 100644
--- a/primus/backends/torchtitan/patches/dcp_consolidate_patches.py
+++ b/primus/backends/torchtitan/patches/dcp_consolidate_patches.py
@@ -15,7 +15,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/embedding_amp_patches.py b/primus/backends/torchtitan/patches/embedding_amp_patches.py
index 3daedf44e..df4e153a3 100644
--- a/primus/backends/torchtitan/patches/embedding_amp_patches.py
+++ b/primus/backends/torchtitan/patches/embedding_amp_patches.py
@@ -21,7 +21,7 @@
 """
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/flex_attention_patches.py b/primus/backends/torchtitan/patches/flex_attention_patches.py
index 5acadd41c..c05ed3f18 100644
--- a/primus/backends/torchtitan/patches/flex_attention_patches.py
+++ b/primus/backends/torchtitan/patches/flex_attention_patches.py
@@ -19,7 +19,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/fsdp_weight_tying_patches.py b/primus/backends/torchtitan/patches/fsdp_weight_tying_patches.py
new file mode 100644
index 000000000..2c6ef0679
--- /dev/null
+++ b/primus/backends/torchtitan/patches/fsdp_weight_tying_patches.py
@@ -0,0 +1,102 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+TorchTitan FSDP2 weight-tying patch.
+
+Qwen3 enables ``enable_weight_tying``, sharing ``tok_embeddings.weight`` with
+``output.weight``. Upstream ``apply_fsdp`` wraps those modules in separate
+FSDP groups, which PyTorch 2.12 FSDP2 rejects at lazy init:
+
+    ValueError: Parameter 'None' is shared with a parameter already managed by
+    another FSDP group. For shared/tied parameters, use
+    fully_shard([module_a, module_b]) to place them in the same FSDP group.
+
+This patch wraps upstream ``apply_fsdp`` and temporarily redirects the two
+embedding/output ``fully_shard`` calls without modifying third_party/torchtitan.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from primus.core.patches import PatchContext, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _normalize_modules(target: Any) -> list[Any]:
+    if isinstance(target, (list, tuple)):
+        return list(target)
+    return [target]
+
+
+def _make_apply_fsdp_wrapper(
+    orig_apply_fsdp: Callable[..., Any],
+    parallelize_mod: Any,
+) -> Callable[..., Any]:
+    orig_fully_shard = parallelize_mod.fully_shard
+
+    def apply_fsdp(model, *args, **kwargs):
+        if not getattr(model, "enable_weight_tying", False):
+            return orig_apply_fsdp(model, *args, **kwargs)
+
+        def fully_shard_proxy(target, **fsdp_kwargs):
+            modules = _normalize_modules(target)
+            if modules == [model.tok_embeddings] and model.output is not None:
+                return orig_fully_shard(
+                    [model.tok_embeddings, model.output],
+                    **fsdp_kwargs,
+                )
+            if (
+                model.norm is not None
+                and model.output is not None
+                and set(modules) == {model.norm, model.output}
+            ):
+                return orig_fully_shard(model.norm, **fsdp_kwargs)
+            return orig_fully_shard(target, **fsdp_kwargs)
+
+        parallelize_mod.fully_shard = fully_shard_proxy
+        try:
+            return orig_apply_fsdp(model, *args, **kwargs)
+        finally:
+            parallelize_mod.fully_shard = orig_fully_shard
+
+    return apply_fsdp
+
+
+@register_patch(
+    "torchtitan.fsdp.weight_tying",
+    backend="torchtitan",
+    phase="setup",
+    description=(
+        "Group tied tok_embeddings/output weights in one FSDP2 shard group for "
+        "models with enable_weight_tying (Qwen3 on PyTorch 2.12+)."
+    ),
+)
+def patch_torchtitan_fsdp_weight_tying(ctx: PatchContext) -> None:  # noqa: ARG001
+    import importlib
+
+    patched = []
+    try:
+        llama4_mod = importlib.import_module("torchtitan.models.llama4.infra.parallelize")
+        wrapped = _make_apply_fsdp_wrapper(llama4_mod.apply_fsdp, llama4_mod)
+        llama4_mod.apply_fsdp = wrapped
+        patched.append("llama4")
+
+        qwen3_mod = importlib.import_module("torchtitan.models.qwen3.infra.parallelize")
+        qwen3_mod.apply_fsdp = wrapped
+        patched.append("qwen3")
+    except ImportError as exc:
+        log_rank_0(
+            "[Patch:torchtitan.fsdp.weight_tying] "
+            f"Skipped: torchtitan parallelize modules unavailable ({exc})",
+        )
+        return
+
+    log_rank_0(
+        "[Patch:torchtitan.fsdp.weight_tying] "
+        f"Patched apply_fsdp for weight tying in: {', '.join(patched)}",
+    )
diff --git a/primus/backends/torchtitan/patches/logger_patches.py b/primus/backends/torchtitan/patches/logger_patches.py
index 3940c52f6..0e3292bb8 100644
--- a/primus/backends/torchtitan/patches/logger_patches.py
+++ b/primus/backends/torchtitan/patches/logger_patches.py
@@ -33,7 +33,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/mock_dataset_patches.py b/primus/backends/torchtitan/patches/mock_dataset_patches.py
index 9703504bf..341512def 100644
--- a/primus/backends/torchtitan/patches/mock_dataset_patches.py
+++ b/primus/backends/torchtitan/patches/mock_dataset_patches.py
@@ -7,8 +7,7 @@
 """
 TorchTitan Mock HF Dataset Patch
 
-This patch mirrors ``patch_mock_hf_dataset`` from
-``primus.modules.trainer.torchtitan.patch_utils`` using the generic Primus
+This patch implements ``patch_mock_hf_dataset`` using the generic Primus
 patch system, so that HF dataset mocking can be enabled via config without
 tightly coupling it to the trainer implementation.
 """
@@ -17,7 +16,7 @@
 from datasets.search import np
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _create_mock_text_dataset(num_samples: int = 128) -> Dataset:
@@ -90,6 +89,13 @@ def mock_load_dataset(path: str, *args, **kwargs) -> Dataset:
                 return _create_mock_token_dataset(seq_len=8192, vocab_size=32000, num_samples=256)
 
         datasets.load_dataset = mock_load_dataset
+
+        import sys
+
+        text_datasets_mod = sys.modules.get("torchtitan.hf_datasets.text_datasets")
+        if text_datasets_mod is not None:
+            text_datasets_mod.load_dataset = mock_load_dataset
+
         log_rank_0(
             "[Patch:torchtitan.training.mock_hf_dataset] " "Patched datasets.load_dataset successfully.",
         )
diff --git a/primus/backends/torchtitan/patches/model_override_patches.py b/primus/backends/torchtitan/patches/model_override_patches.py
index 37d27e458..ffeaa9464 100644
--- a/primus/backends/torchtitan/patches/model_override_patches.py
+++ b/primus/backends/torchtitan/patches/model_override_patches.py
@@ -50,7 +50,7 @@
 from typing import Any
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def get_standard_model_config_fields() -> set:
diff --git a/primus/backends/torchtitan/patches/pipelining_schedule_patches.py b/primus/backends/torchtitan/patches/pipelining_schedule_patches.py
index fb2040c1a..9b9a6958e 100644
--- a/primus/backends/torchtitan/patches/pipelining_schedule_patches.py
+++ b/primus/backends/torchtitan/patches/pipelining_schedule_patches.py
@@ -14,7 +14,7 @@
 """
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py
index 943d08080..60aa0a639 100644
--- a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py
+++ b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py
@@ -39,7 +39,7 @@
 import os
 
 from primus.core.patches import PatchContext, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def _sdma_all_gather_enabled(ctx: PatchContext) -> bool:
diff --git a/primus/backends/torchtitan/patches/turbo/__init__.py b/primus/backends/torchtitan/patches/turbo/__init__.py
index 5677368ef..dfc2b79bc 100644
--- a/primus/backends/torchtitan/patches/turbo/__init__.py
+++ b/primus/backends/torchtitan/patches/turbo/__init__.py
@@ -17,6 +17,7 @@
     attention_patches,
     deepseek_v3_classic_attention_patches,
     fp8_linear_patches,
+    gptoss_sink_attention_patches,
     moe_grouped_mm_patches,
     mx_linear_patches,
 )
diff --git a/primus/backends/torchtitan/patches/turbo/async_tp_patches.py b/primus/backends/torchtitan/patches/turbo/async_tp_patches.py
index ae260e556..0db17c15d 100644
--- a/primus/backends/torchtitan/patches/turbo/async_tp_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/async_tp_patches.py
@@ -18,7 +18,7 @@
 from typing import Any, Optional
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/turbo/attention_patches.py b/primus/backends/torchtitan/patches/turbo/attention_patches.py
index 9146cefb1..8dbf1fe05 100644
--- a/primus/backends/torchtitan/patches/turbo/attention_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/attention_patches.py
@@ -16,7 +16,7 @@
 """
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py b/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py
index 119417188..ca96fe194 100644
--- a/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py
@@ -17,7 +17,7 @@
 """
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py b/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py
index fbd582bcf..b5df8a9ff 100644
--- a/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py
@@ -16,7 +16,7 @@
 """
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py b/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py
new file mode 100644
index 000000000..d90cee727
--- /dev/null
+++ b/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py
@@ -0,0 +1,70 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+TorchTitan GPT-OSS Primus-Turbo Sink Attention Patch
+====================================================
+
+GPT-OSS uses FlexAttention with learnable per-head attention sinks and a
+sliding-window mask on even layers. The module-form ``TurboAttention`` used by
+the other models' turbo path does NOT expose ``sink`` / ``window_size``, so it
+cannot model GPT-OSS attention. However, the Primus-Turbo *functional* kernel
+``primus_turbo.pytorch.ops.flash_attn_func`` does support both (this is the same
+API the Megatron ``PrimusTurboAttention`` uses for GPT-OSS sink attention).
+
+This setup patch swaps the upstream GPT-OSS ``Attention`` and ``TransformerBlock``
+classes for Primus mirrors that route attention through ``flash_attn_func`` with
+``sink=self.sinks`` and the per-layer ``window_size``. It is gated on
+``primus_turbo.enable_primus_turbo`` + ``primus_turbo.use_turbo_attention`` and
+only applies to ``gpt_oss`` runs, so existing GPT-OSS FlexAttention configs
+(turbo off) are unaffected.
+"""
+
+from primus.core.patches import PatchContext, get_param, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _gptoss_turbo_enabled(ctx: PatchContext) -> bool:
+    return (
+        get_param(ctx, "model.name", None) == "gpt_oss"
+        and get_param(ctx, "primus_turbo.enable_primus_turbo", False)
+        and get_param(ctx, "primus_turbo.use_turbo_attention", False)
+    )
+
+
+@register_patch(
+    "torchtitan.primus_turbo.gptoss_sink_attention",
+    backend="torchtitan",
+    phase="setup",
+    description="Use Primus-Turbo functional flash_attn_func (sink + sliding window) for GPT-OSS",
+    condition=_gptoss_turbo_enabled,
+)
+def patch_gptoss_sink_attention(ctx: PatchContext) -> None:
+    """Install the Primus-Turbo sink-attention mirror for GPT-OSS."""
+    log_rank_0(
+        "[Patch:torchtitan.primus_turbo.gptoss_sink_attention] "
+        "Enabling Primus-Turbo sink attention (flash_attn_func) for GPT-OSS...",
+    )
+
+    import torchtitan.models.gpt_oss.model.model as gptoss_model_mod
+
+    from primus.backends.torchtitan.models.gpt_oss.model.model import (
+        Attention as GptOssTurboAttention,
+    )
+    from primus.backends.torchtitan.models.gpt_oss.model.model import (
+        TransformerBlock as GptOssTurboBlock,
+    )
+
+    # Replace the module-level classes so GptOssModel builds the mirror instances
+    # (the mirror TransformerBlock also constructs the mirror Attention, and both
+    # resolve these names from this module's globals at build time).
+    gptoss_model_mod.Attention = GptOssTurboAttention
+    gptoss_model_mod.TransformerBlock = GptOssTurboBlock
+
+    log_rank_0(
+        "[Patch:torchtitan.primus_turbo.gptoss_sink_attention] "
+        "GPT-OSS Primus-Turbo sink attention successfully installed.",
+    )
diff --git a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py
index f48a1d120..17cf70b8d 100644
--- a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py
@@ -19,10 +19,12 @@
       delegate to Primus' grouped_mm implementation with ``use_fp8`` bound.
 """
 
-import functools
-
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
+
+# Upstream apply_compile skips torch.compile when the grouped_mm __qualname__
+# already contains this string. We reuse it to opt out of compile.
+_ALREADY_PATCHED_SENTINEL = "_run_experts_grouped_mm_dynamic"
 
 
 @register_patch(
@@ -43,17 +45,26 @@ def patch_torchtitan_moe(ctx: PatchContext) -> None:
 
     from primus.backends.torchtitan.models.moe.moe import _run_experts_grouped_mm
 
-    # Get MoE FP8 configuration and create a partial function
+    # Get MoE FP8 configuration and bind it onto the replacement function.
     use_moe_fp8 = get_param(ctx, "primus_turbo.use_moe_fp8", False)
     log_rank_0(
         "[Patch:torchtitan.primus_turbo.moe_grouped_mm] " f"Set MoE FP8 mode: {use_moe_fp8}",
     )
 
-    # Patch the grouped_mm function with use_fp8 parameter pre-set
-    torchtitan.models.moe.moe._run_experts_grouped_mm = functools.partial(
-        _run_experts_grouped_mm,
-        use_fp8=use_moe_fp8,
-    )
+    # The turbo grouped_mm kernels are torch.compiler.disable'd, so letting
+    # upstream compile our wrapper under fullgraph traces into them and hard-errors
+    # (dynamo gb0098). Decorating the wrapper with torch.compiler.disable does not
+    # help (torch.compile unwraps it via innermost_fn). Instead we name the wrapper
+    # with the sentinel so apply_compile treats it as already patched and skips
+    # compile, keeping the turbo kernels in eager as designed.
+    def _run_experts_grouped_mm_dynamic(*args, **kwargs):
+        kwargs.setdefault("use_fp8", use_moe_fp8)
+        return _run_experts_grouped_mm(*args, **kwargs)
+
+    _run_experts_grouped_mm_dynamic.__qualname__ = _ALREADY_PATCHED_SENTINEL
+    _run_experts_grouped_mm_dynamic.__name__ = _ALREADY_PATCHED_SENTINEL
+
+    torchtitan.models.moe.moe._run_experts_grouped_mm = _run_experts_grouped_mm_dynamic
 
     log_rank_0(
         "[Patch:torchtitan.primus_turbo.moe_grouped_mm] "
diff --git a/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py b/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py
index 15473ffb2..bb2a23c1f 100644
--- a/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py
+++ b/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py
@@ -16,7 +16,7 @@
 """
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/patches/wandb_patches.py b/primus/backends/torchtitan/patches/wandb_patches.py
index 129427ba7..538ebea2b 100644
--- a/primus/backends/torchtitan/patches/wandb_patches.py
+++ b/primus/backends/torchtitan/patches/wandb_patches.py
@@ -23,7 +23,7 @@
 import os
 
 from primus.core.patches import PatchContext, get_param, register_patch
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 @register_patch(
diff --git a/primus/backends/torchtitan/torchtitan_adapter.py b/primus/backends/torchtitan/torchtitan_adapter.py
index 0bace2840..f8fa1e5f9 100644
--- a/primus/backends/torchtitan/torchtitan_adapter.py
+++ b/primus/backends/torchtitan/torchtitan_adapter.py
@@ -24,7 +24,7 @@
 from primus.backends.torchtitan.argument_builder import TorchTitanJobConfigBuilder
 from primus.core.backend.backend_adapter import BackendAdapter
 from primus.core.backend.backend_registry import BackendRegistry
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class TorchTitanAdapter(BackendAdapter):
diff --git a/primus/backends/torchtitan/torchtitan_pretrain_trainer.py b/primus/backends/torchtitan/torchtitan_pretrain_trainer.py
index d2e88d095..8cc6d1c98 100644
--- a/primus/backends/torchtitan/torchtitan_pretrain_trainer.py
+++ b/primus/backends/torchtitan/torchtitan_pretrain_trainer.py
@@ -17,11 +17,14 @@
 class TorchTitanPretrainTrainer(BaseTrainer):
     """Trainer class for TorchTitan pre-training."""
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, **kwargs):
         # Patch TorchTitan logger before any other initialization
         self._patch_torchtitan_logger()
 
-        super().__init__(backend_args=backend_args)
+        # The core runtime instantiates every trainer with BaseModule-style
+        # context kwargs (module_name, primus_config, module_rank, ...). Accept
+        # and forward them so BaseTrainer can filter them cooperatively.
+        super().__init__(backend_args=backend_args, **kwargs)
         self._trainer: Optional["Trainer"] = None  # type: ignore[name-defined]
 
     def _patch_torchtitan_logger(self):
diff --git a/primus/backends/transformer_engine/pytorch/common/triton/permutation.py b/primus/backends/transformer_engine/pytorch/common/triton/permutation.py
index 5f054a661..1b11b58f6 100644
--- a/primus/backends/transformer_engine/pytorch/common/triton/permutation.py
+++ b/primus/backends/transformer_engine/pytorch/common/triton/permutation.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 
diff --git a/primus/backends/transformer_engine/pytorch/cpp_extensions/gemm.py b/primus/backends/transformer_engine/pytorch/cpp_extensions/gemm.py
index d40dc7a41..870d4241b 100644
--- a/primus/backends/transformer_engine/pytorch/cpp_extensions/gemm.py
+++ b/primus/backends/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -1,6 +1,6 @@
 ###############################################################################
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
diff --git a/primus/backends/transformer_engine/pytorch/module/base.py b/primus/backends/transformer_engine/pytorch/module/base.py
index 003e18b45..0d73af244 100644
--- a/primus/backends/transformer_engine/pytorch/module/base.py
+++ b/primus/backends/transformer_engine/pytorch/module/base.py
@@ -1,6 +1,6 @@
 ###############################################################################
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -13,7 +13,7 @@
 from transformer_engine.pytorch.module import base
 
 import primus.backends.transformer_engine.transformer_engine_torch as ptex
-from primus.modules.module_utils import log_rank_0, warning_rank_0
+from primus.core.utils.module_utils import log_rank_0, warning_rank_0
 
 
 def get_cublas_workspace_size_bytes() -> None:
diff --git a/primus/backends/transformer_engine/pytorch/permutation.py b/primus/backends/transformer_engine/pytorch/permutation.py
index fc9374973..d4c0eb6b4 100644
--- a/primus/backends/transformer_engine/pytorch/permutation.py
+++ b/primus/backends/transformer_engine/pytorch/permutation.py
@@ -1,5 +1,5 @@
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 
diff --git a/primus/backends/transformer_engine/pytorch/triton/permutation.py b/primus/backends/transformer_engine/pytorch/triton/permutation.py
index d18857ee3..2efbd7d17 100644
--- a/primus/backends/transformer_engine/pytorch/triton/permutation.py
+++ b/primus/backends/transformer_engine/pytorch/triton/permutation.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 
diff --git a/primus/backends/transformer_engine/transformer_engine_torch/__init__.py b/primus/backends/transformer_engine/transformer_engine_torch/__init__.py
index 193c155ee..bcc9a8776 100644
--- a/primus/backends/transformer_engine/transformer_engine_torch/__init__.py
+++ b/primus/backends/transformer_engine/transformer_engine_torch/__init__.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
diff --git a/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap.py b/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap.py
index e7b548a46..a74635a13 100644
--- a/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap.py
+++ b/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap.py
@@ -1,6 +1,6 @@
 ###############################################################################
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
diff --git a/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap_type.py b/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap_type.py
index 1fead489c..ac286ebed 100644
--- a/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap_type.py
+++ b/primus/backends/transformer_engine/transformer_engine_torch/comm_overlap_type.py
@@ -1,6 +1,6 @@
 ###############################################################################
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
diff --git a/primus/backends/transformer_engine/transformer_engine_torch/gemm.py b/primus/backends/transformer_engine/transformer_engine_torch/gemm.py
index 0990f3003..f67432c93 100644
--- a/primus/backends/transformer_engine/transformer_engine_torch/gemm.py
+++ b/primus/backends/transformer_engine/transformer_engine_torch/gemm.py
@@ -1,6 +1,6 @@
 ###############################################################################
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
+# Modification Copyright© 2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
diff --git a/primus/cli/main.py b/primus/cli/main.py
index 600bc7702..de9638be9 100644
--- a/primus/cli/main.py
+++ b/primus/cli/main.py
@@ -1,15 +1,42 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
+# MLPerf log suppression must run before any heavy import (Megatron, TE,
+# aiter, ...) that may print/log at import time. Importing this module only
+# triggers the light ``primus/__init__.py`` and installs the FD-level filter
+# when ``PRIMUS_LOG_SUPPRESSION=1`` is set; it is a complete no-op otherwise.
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _prefer_checkout_primus_on_sys_path() -> None:
+    """Prefer the git checkout over an installed wheel for in-tree Primus modules."""
+    if not (_REPO_ROOT / "primus" / "mlperf_log_suppression.py").is_file():
+        return
+    root = str(_REPO_ROOT)
+    if root in sys.path:
+        sys.path.remove(root)
+    sys.path.insert(0, root)
+
+
+_prefer_checkout_primus_on_sys_path()
+
+try:
+    import primus.mlperf_log_suppression  # noqa: F401  # isort: skip
+except ModuleNotFoundError:
+    # Optional in-tree module; recipe-local _log_suppression covers MLPerf runs.
+    pass
+
 import argparse
 import importlib
+import logging
 import pkgutil
-import sys
 import traceback
-from pathlib import Path
 from typing import Callable, Dict, Iterable, Optional, Set
 
 SUBCOMMAND_PACKAGE = "primus.cli.subcommands"
@@ -20,11 +47,7 @@ def _ensure_project_root_on_path() -> None:
     Allow running `python primus/cli/main.py` from the repo root without
     requiring an installed package.
     """
-    if __package__:
-        return
-    project_root = Path(__file__).resolve().parents[2]
-    if str(project_root) not in sys.path:
-        sys.path.insert(0, str(project_root))
+    _prefer_checkout_primus_on_sys_path()
 
 
 def _iter_subcommand_modules() -> Iterable[str]:
@@ -126,6 +149,8 @@ def main():
       ...
     """
     _ensure_project_root_on_path()
+    logging.basicConfig(level=logging.WARNING, format="%(message)s")
+    logging.getLogger("primus").setLevel(logging.INFO)
     parser = argparse.ArgumentParser(
         prog="primus",
         description="Primus Unified CLI for Training & Utilities",
diff --git a/primus/cli/subcommands/data.py b/primus/cli/subcommands/data.py
new file mode 100644
index 000000000..5d47a956a
--- /dev/null
+++ b/primus/cli/subcommands/data.py
@@ -0,0 +1,1092 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Dataset preparation subcommands for diffusion training.
+
+Commands:
+    primus data diffusion-raw      - Prepare raw Energon WebDataset
+    primus data diffusion-encoded  - Prepare pre-encoded Energon WebDataset
+    primus data diffusion-ingest   - Stream pre-encoded Arrow data into WebDataset
+    primus data automodel-cache    - Prepare a pre-encoded AutoModel flat cache
+
+The three ``diffusion-*`` commands target the **Megatron** backend and output
+Energon WebDataset shards. They support torch.distributed for multi-GPU
+processing:
+
+    torchrun --nproc_per_node=8 primus data diffusion-raw ...
+
+``automodel-cache`` targets the **NeMo AutoModel** backend, whose diffusion
+recipe reads a flat per-sample cache instead. It is a single-GPU offline step and
+is kept as a sibling rather than a mode of the commands above, because the two
+pipelines share no format and no code path.
+"""
+
+import argparse
+import logging
+
+from primus.backends.megatron.data.diffusion.preprocessing.auth import (
+    HFAuthError,
+    setup_hf_authentication,
+)
+from primus.backends.megatron.data.diffusion.preprocessing.pipelines.encoded import (
+    EncodedDatasetPipeline,
+)
+from primus.backends.megatron.data.diffusion.preprocessing.pipelines.raw import (
+    RawDatasetPipeline,
+)
+from primus.tools.utils import finalize_distributed, init_distributed
+
+logger = logging.getLogger(__name__)
+
+
+def _add_common_args(parser):
+    """Add arguments common to all data preparation commands."""
+    # Source configuration
+    source_group = parser.add_argument_group("Source Configuration")
+    source_group.add_argument(
+        "--source-type",
+        required=False,
+        choices=["directory", "huggingface", "webdataset"],
+        help="Type of input data source (required, can be set via --config)",
+    )
+    source_group.add_argument("--input-dir", type=str, help="Input directory (for directory source)")
+    source_group.add_argument("--hf-dataset", type=str, help="HuggingFace dataset name (for HF source)")
+    source_group.add_argument(
+        "--hf-split", type=str, default="train", help="HuggingFace dataset split (default: train)"
+    )
+    source_group.add_argument(
+        "--hf-data-files",
+        type=str,
+        default=None,
+        help='Specific files/paths within HF dataset (e.g., "data_1024_10K/*.tar")',
+    )
+    source_group.add_argument("--input-path", type=str, help="WebDataset path/glob (for webdataset source)")
+
+    # Output configuration (Energon format)
+    output_group = parser.add_argument_group("Output Configuration (Energon WebDataset)")
+    output_group.add_argument(
+        "--output-dir",
+        required=False,
+        help="Output directory for Energon WebDataset (required, can be set via --config)",
+    )
+    output_group.add_argument(
+        "--shard-size", type=int, default=1000, help="Samples per shard (default: 1000)"
+    )
+    output_group.add_argument(
+        "--max-samples", type=int, default=None, help="Maximum samples to process (default: all)"
+    )
+    output_group.add_argument("--compress", action="store_true", help="Compress tar files with gzip")
+
+    # Image preprocessing
+    image_group = parser.add_argument_group("Image Preprocessing")
+    image_group.add_argument(
+        "--variable-size",
+        action="store_true",
+        help="Resize images to the nearest multiple of 16 instead of fixed size",
+    )
+    image_group.add_argument(
+        "--image-size",
+        type=int,
+        default=1024,
+        help="Resize images to this size (default: 1024) when variable_size is disabled",
+    )
+    image_group.add_argument(
+        "--center-crop", action="store_false", help="Disable center cropping of images before resize"
+    )
+    image_group.add_argument(
+        "--max-size",
+        type=int,
+        default=1024,
+        help="Maximum image dimension for variable size mode (default: 1024)",
+    )
+
+    # HuggingFace Authentication
+    auth_group = parser.add_argument_group("HuggingFace Authentication")
+    auth_group.add_argument(
+        "--hf-token-file",
+        type=str,
+        default=None,
+        help=(
+            "Path to file containing HuggingFace token. "
+            "File must have secure permissions (600 or 400). "
+            "If not provided, falls back to HF_TOKEN env variable "
+            "or HF CLI login (~/.cache/huggingface/token)."
+        ),
+    )
+
+    return parser
+
+
+def _validate_source_args(args):
+    """Validate that required source arguments are provided."""
+    if args.source_type == "directory" and not args.input_dir:
+        raise ValueError("--input-dir required for directory source")
+    if args.source_type == "huggingface" and not args.hf_dataset:
+        raise ValueError("--hf-dataset required for huggingface source")
+    if args.source_type == "webdataset" and not args.input_path:
+        raise ValueError("--input-path required for webdataset source")
+
+
+def _prepare_raw(args):
+    """Prepare raw Energon WebDataset for on-the-fly encoding during training."""
+    # Initialize torch.distributed if launched with torchrun
+    init_distributed()
+
+    try:
+        _validate_source_args(args)
+
+        # Setup HF authentication
+        try:
+            setup_hf_authentication(token_file=getattr(args, "hf_token_file", None))
+        except HFAuthError as e:
+            logger.error(f"Authentication failed: {e}")
+            raise
+
+        logger.info("=" * 80)
+        logger.info("Megatron Diffusion: Raw Energon WebDataset Preparation")
+        logger.info("=" * 80)
+
+        pipeline = RawDatasetPipeline(
+            source_type=args.source_type,
+            output_dir=args.output_dir,
+            variable_size=args.variable_size,
+            image_size=args.image_size,
+            center_crop=args.center_crop,
+            max_size=args.max_size,
+            image_format=getattr(args, "image_format", "JPEG"),
+            image_quality=getattr(args, "image_quality", 95),
+            shard_size=args.shard_size,
+            max_samples=args.max_samples,
+            compress=args.compress,
+        )
+
+        source_kwargs = {}
+        if args.source_type == "directory":
+            source_kwargs["input_dir"] = args.input_dir
+        elif args.source_type == "huggingface":
+            source_kwargs["hf_dataset"] = args.hf_dataset
+            source_kwargs["hf_split"] = args.hf_split
+            source_kwargs["hf_data_files"] = getattr(args, "hf_data_files", None)
+            # Add data format configuration
+            source_kwargs["image_key"] = getattr(args, "image_key", None)
+            source_kwargs["caption_key"] = getattr(args, "caption_key", None)
+            source_kwargs["image_keys"] = getattr(args, "image_keys", None)
+            source_kwargs["caption_keys"] = getattr(args, "caption_keys", None)
+        elif args.source_type == "webdataset":
+            source_kwargs["input_path"] = args.input_path
+
+        results = pipeline.run(**source_kwargs)
+
+        logger.info("=" * 80)
+        logger.info("✓ Raw Energon WebDataset preparation complete!")
+        logger.info(f"  Samples processed: {results['samples_processed']}")
+        logger.info(f"  Samples skipped: {results['samples_skipped']}")
+        logger.info(f"  Shards created: {results['shards_written']}")
+        logger.info(f"  Output: {args.output_dir}")
+        logger.info(f"  Format: Raw Energon WebDataset (images + captions)")
+        logger.info(f"  Note: Encoding will be done on-the-fly during training")
+        logger.info("=" * 80)
+
+        # Finalize dataset unless --no-finalize was passed
+        if not getattr(args, "no_finalize", False):
+            # Import distributed utilities
+            import torch.distributed as dist
+
+            # Wait for all ranks to complete data preparation
+            if dist.is_initialized():
+                logger.info("Waiting for all ranks to complete...")
+                dist.barrier()
+
+            # Only rank 0 runs finalization
+            should_finalize = not dist.is_initialized() or dist.get_rank() == 0
+
+            if should_finalize:
+                from primus.backends.megatron.data.diffusion.preprocessing.finalize import (
+                    finalize_energon_dataset,
+                )
+
+                try:
+                    finalize_energon_dataset(
+                        output_dir=args.output_dir,
+                        train_split=getattr(args, "train_split", 1.0),
+                        encoding="raw",
+                        num_workers=8,
+                    )
+                except Exception as e:
+                    logger.error(f"Finalization failed: {e}")
+                    raise
+
+            # Wait again so all ranks exit together
+            if dist.is_initialized():
+                dist.barrier()
+
+    finally:
+        # Clean up distributed
+        finalize_distributed()
+
+
+def _flatten_preprocessing_config(config_dict: dict) -> dict:
+    """
+    Flatten nested YAML config to match CLI argument structure.
+
+    Transforms hierarchical YAML structure into flat dict matching argparse namespace.
+
+    Args:
+        config_dict: Nested config from YAML file
+
+    Returns:
+        Flattened dict with CLI argument names as keys
+
+    Example:
+        >>> config = {
+        ...     'source': {'type': 'huggingface', 'hf_dataset': 'pokemon'},
+        ...     'model': {'batch_size': 8}
+        ... }
+        >>> flat = _flatten_preprocessing_config(config)
+        >>> flat['source_type'], flat['batch_size']
+        ('huggingface', 8)
+    """
+    flat = {}
+
+    # Source configuration
+    if "source" in config_dict:
+        source = config_dict["source"]
+        flat["source_type"] = source.get("type")
+        flat["hf_dataset"] = source.get("hf_dataset")
+        flat["hf_split"] = source.get("hf_split", "train")
+        flat["hf_data_files"] = source.get("hf_data_files")
+        flat["input_dir"] = source.get("input_dir")
+        flat["input_path"] = source.get("input_path")
+
+    # Data format configuration (field mappings for image/caption extraction)
+    if "data_format" in config_dict:
+        data_format = config_dict["data_format"]
+        flat["image_key"] = data_format.get("image_key")
+        flat["caption_key"] = data_format.get("caption_key")
+        flat["image_keys"] = data_format.get("image_keys")
+        flat["caption_keys"] = data_format.get("caption_keys")
+
+    # Output configuration
+    if "output" in config_dict:
+        output = config_dict["output"]
+        flat["output_dir"] = output.get("output_dir")
+        flat["shard_size"] = output.get("shard_size", 1000)
+        flat["max_samples"] = output.get("max_samples")
+        flat["compress"] = output.get("compress", False)
+
+    # Model configuration
+    if "model" in config_dict:
+        model = config_dict["model"]
+        flat["model_path"] = model.get("model_path", "black-forest-labs/FLUX.1-dev")
+        flat["vae_path"] = model.get("vae_path")
+        flat["t5_path"] = model.get("t5_path")
+        flat["clip_path"] = model.get("clip_path")
+        flat["precision"] = model.get("precision", "bf16")
+        flat["device"] = model.get("device", "cuda")
+        flat["batch_size"] = model.get("batch_size", 8)
+        flat["t5_max_length"] = model.get("t5_max_length", 512)
+        flat["vae_latent_mode"] = model.get("vae_latent_mode", "presampled")
+
+    # Image preprocessing
+    if "image" in config_dict:
+        image = config_dict["image"]
+        flat["image_size"] = image.get("image_size", 1024)
+        flat["variable_size"] = image.get("variable_size", False)
+        flat["center_crop"] = image.get("center_crop", True)
+        flat["max_size"] = image.get("max_size", 1024)
+
+    # Authentication
+    if "auth" in config_dict:
+        auth = config_dict["auth"]
+        flat["hf_token_file"] = auth.get("hf_token_file")
+
+    return flat
+
+
+def _get_encoded_parser_defaults() -> dict:
+    """
+    Get default values from encoded parser for override detection.
+
+    Returns dict of argument name -> default value to detect which CLI
+    arguments were explicitly set vs using defaults.
+    """
+    return {
+        "config": None,
+        "source_type": None,
+        "hf_dataset": None,
+        "hf_split": "train",
+        "hf_data_files": None,
+        "input_dir": None,
+        "input_path": None,
+        "image_key": None,
+        "caption_key": None,
+        "image_keys": None,
+        "caption_keys": None,
+        "output_dir": None,
+        "shard_size": 1000,
+        "max_samples": None,
+        "compress": False,
+        "model_path": "black-forest-labs/FLUX.1-dev",
+        "vae_path": None,
+        "t5_path": None,
+        "clip_path": None,
+        "precision": "bf16",
+        "device": "cuda",
+        "batch_size": 8,
+        "t5_max_length": 512,
+        "image_size": 1024,
+        "variable_size": False,
+        "center_crop": True,
+        "max_size": 1024,
+        "hf_token_file": None,
+        "vae_latent_mode": "presampled",
+    }
+
+
+def _load_config_with_cli_overrides(args: "argparse.Namespace") -> "argparse.Namespace":
+    """
+    Load YAML config and merge with CLI arguments.
+
+    Priority order (highest to lowest):
+    1. Explicitly provided CLI arguments
+    2. YAML config values
+    3. CLI default values
+
+    Args:
+        args: Parsed CLI arguments
+
+    Returns:
+        Merged namespace with final configuration
+
+    Example:
+        >>> # With config file specifying batch_size: 8
+        >>> # And CLI arg --batch-size 16
+        >>> # Result: batch_size = 16 (CLI wins)
+    """
+    import argparse
+
+    # If no config file, return args as-is
+    if not getattr(args, "config", None):
+        return args
+
+    from primus.core.utils import yaml_utils
+
+    logger.info(f"Loading config from: {args.config}")
+
+    # Load and flatten YAML config
+    config_dict = yaml_utils.parse_yaml(args.config)
+    flat_config = _flatten_preprocessing_config(config_dict)
+
+    # Start with config values
+    merged_dict = flat_config.copy()
+
+    # Override with explicitly provided CLI arguments
+    parser_defaults = _get_encoded_parser_defaults()
+    cli_args = vars(args)
+
+    for key, cli_value in cli_args.items():
+        if key == "config":
+            # Keep config path for reference
+            merged_dict["config"] = cli_value
+            continue
+
+        # If CLI value differs from default, it was explicitly set
+        default_value = parser_defaults.get(key)
+        if cli_value != default_value:
+            merged_dict[key] = cli_value
+            logger.debug(f"CLI override: {key} = {cli_value}")
+
+    logger.info("Configuration merged (CLI args override YAML)")
+    return argparse.Namespace(**merged_dict)
+
+
+def _validate_preprocessing_config(args: "argparse.Namespace") -> None:
+    """
+    Validate preprocessing configuration after merging.
+
+    Ensures all required parameters are present regardless of source
+    (CLI, YAML, or both).
+
+    Args:
+        args: Merged configuration namespace
+
+    Raises:
+        ValueError: If required arguments are missing or invalid
+
+    Example:
+        >>> args = argparse.Namespace(source_type='huggingface', output_dir=None)
+        >>> _validate_preprocessing_config(args)  # Raises ValueError
+    """
+    # Check required arguments
+    required = ["source_type", "output_dir"]
+    missing = [arg for arg in required if not getattr(args, arg, None)]
+
+    if missing:
+        raise ValueError(
+            f"Missing required arguments: {', '.join(missing)}. "
+            f"Provide via --config YAML or CLI arguments."
+        )
+
+    # Validate source-specific requirements
+    if args.source_type == "huggingface":
+        if not getattr(args, "hf_dataset", None):
+            raise ValueError(
+                "Missing --hf-dataset for source-type 'huggingface'. "
+                "Specify in config file (source.hf_dataset) or via CLI."
+            )
+    elif args.source_type == "directory":
+        if not getattr(args, "input_dir", None):
+            raise ValueError(
+                "Missing --input-dir for source-type 'directory'. "
+                "Specify in config file (source.input_dir) or via CLI."
+            )
+    elif args.source_type == "webdataset":
+        if not getattr(args, "input_path", None):
+            raise ValueError(
+                "Missing --input-path for source-type 'webdataset'. "
+                "Specify in config file (source.input_path) or via CLI."
+            )
+    else:
+        raise ValueError(
+            f"Invalid source-type: {args.source_type}. " f"Must be one of: huggingface, directory, webdataset"
+        )
+
+
+def _prepare_encoded(args):
+    """Prepare pre-encoded Energon WebDataset with VAE/T5/CLIP for fast training."""
+    # Load and merge config if provided
+    args = _load_config_with_cli_overrides(args)
+
+    # Validate merged configuration
+    _validate_preprocessing_config(args)
+
+    # Initialize torch.distributed if launched with torchrun
+    init_distributed()
+
+    try:
+        _validate_source_args(args)
+
+        # Setup HF authentication
+        try:
+            setup_hf_authentication(token_file=getattr(args, "hf_token_file", None))
+        except HFAuthError as e:
+            logger.error(f"Authentication failed: {e}")
+            raise
+
+        logger.info("=" * 80)
+        logger.info("Megatron Diffusion: Pre-encoded Energon WebDataset Preparation")
+        logger.info("=" * 80)
+
+        pipeline = EncodedDatasetPipeline(
+            source_type=args.source_type,
+            output_dir=args.output_dir,
+            model_path=args.model_path,
+            vae_path=args.vae_path,
+            t5_path=args.t5_path,
+            clip_path=args.clip_path,
+            precision=args.precision,
+            device=args.device,
+            batch_size=args.batch_size,
+            t5_max_length=args.t5_max_length,
+            variable_size=args.variable_size,
+            image_size=args.image_size,
+            center_crop=args.center_crop,
+            max_size=args.max_size,
+            shard_size=args.shard_size,
+            max_samples=args.max_samples,
+            compress=args.compress,
+            vae_latent_mode=getattr(args, "vae_latent_mode", "presampled"),
+        )
+
+        source_kwargs = {}
+        if args.source_type == "directory":
+            source_kwargs["input_dir"] = args.input_dir
+        elif args.source_type == "huggingface":
+            source_kwargs["hf_dataset"] = args.hf_dataset
+            source_kwargs["hf_split"] = args.hf_split
+            source_kwargs["hf_data_files"] = getattr(args, "hf_data_files", None)
+            # Add data format configuration
+            source_kwargs["image_key"] = getattr(args, "image_key", None)
+            source_kwargs["caption_key"] = getattr(args, "caption_key", None)
+            source_kwargs["image_keys"] = getattr(args, "image_keys", None)
+            source_kwargs["caption_keys"] = getattr(args, "caption_keys", None)
+
+        elif args.source_type == "webdataset":
+            source_kwargs["input_path"] = args.input_path
+
+        results = pipeline.run(**source_kwargs)
+
+        logger.info("=" * 80)
+        logger.info("✓ Pre-encoded Energon WebDataset preparation complete!")
+        logger.info(f"  Samples processed: {results['samples_processed']}")
+        logger.info(f"  Samples skipped: {results['samples_skipped']}")
+        logger.info(f"  Shards created: {results['shards_written']}")
+        logger.info(f"  Output: {args.output_dir}")
+        logger.info(f"  Format: Pre-encoded Energon WebDataset (VAE/T5/CLIP tensors)")
+        logger.info(f"  Note: Training will be faster (no on-the-fly encoding)")
+        logger.info("=" * 80)
+
+        # Finalize dataset unless --no-finalize was passed
+        if not getattr(args, "no_finalize", False):
+            # Import distributed utilities
+            import torch.distributed as dist
+
+            # Wait for all ranks to complete data preparation
+            if dist.is_initialized():
+                logger.info("Waiting for all ranks to complete...")
+                dist.barrier()
+
+            # Only rank 0 runs finalization
+            should_finalize = not dist.is_initialized() or dist.get_rank() == 0
+
+            if should_finalize:
+                from primus.backends.megatron.data.diffusion.preprocessing.finalize import (
+                    finalize_energon_dataset,
+                )
+
+                try:
+                    finalize_energon_dataset(
+                        output_dir=args.output_dir,
+                        train_split=getattr(args, "train_split", 1.0),
+                        encoding="preencoded",
+                        num_workers=8,
+                    )
+                except Exception as e:
+                    logger.error(f"Finalization failed: {e}")
+                    raise
+
+            # Wait again so all ranks exit together
+            if dist.is_initialized():
+                dist.barrier()
+
+    finally:
+        # Clean up distributed
+        finalize_distributed()
+
+
+def _load_ingest_config(args):
+    """Load ingest YAML config and merge with CLI overrides.
+
+    Returns a flat dict consumed by _prepare_ingest.
+    """
+    from primus.core.utils import yaml_utils
+
+    config_path = args.config
+    if not config_path:
+        raise ValueError("--config is required for diffusion-ingest")
+
+    raw = yaml_utils.parse_yaml(config_path)
+
+    config = {
+        "datasets": raw.get("datasets", []),
+        "output_dir": raw.get("output", {}).get("output_dir"),
+        "max_workers": raw.get("pipeline", {}).get("max_workers", 4),
+        "prefetch_depth": raw.get("pipeline", {}).get("prefetch_depth", 6),
+        "max_files": raw.get("pipeline", {}).get("max_files"),
+        "no_finalize": False,
+    }
+
+    if raw.get("empty_encodings"):
+        config["empty_encodings"] = raw["empty_encodings"]
+
+    # CLI overrides
+    if getattr(args, "output_dir", None):
+        config["output_dir"] = args.output_dir
+    if getattr(args, "input_dir", None):
+        config["input_dir"] = args.input_dir
+    if getattr(args, "max_workers", None) is not None:
+        config["max_workers"] = args.max_workers
+    if getattr(args, "prefetch_depth", None) is not None:
+        config["prefetch_depth"] = args.prefetch_depth
+    if getattr(args, "max_files", None) is not None:
+        config["max_files"] = args.max_files
+    if getattr(args, "no_finalize", False):
+        config["no_finalize"] = True
+
+    if not config.get("output_dir"):
+        raise ValueError("output_dir is required. Set in config (output.output_dir) " "or via --output-dir.")
+
+    # Default input_dir
+    if "input_dir" not in config:
+        from pathlib import Path
+
+        config["input_dir"] = str(Path(config["output_dir"]) / "_arrow_tmp")
+
+    return config
+
+
+def _download_empty_encodings(config):
+    """Download empty_encodings files (small .npy files for CFG dropout)."""
+    from pathlib import Path
+
+    from primus.backends.megatron.data.diffusion.preprocessing.download import (
+        download_with_backoff,
+        fetch_manifest,
+    )
+
+    enc_cfg = config["empty_encodings"]
+    output_dir = Path(config["output_dir"]) / enc_cfg.get("output_subdir", "empty_encodings")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    base_url, entries = fetch_manifest(enc_cfg["manifest_url"])
+
+    for md5, fname in entries:
+        dest = output_dir / fname
+        if dest.exists():
+            logger.info(f"  [skip] {fname} (exists)")
+            continue
+        file_url = f"{base_url}/{fname}"
+        logger.info(f"  [download] {fname}")
+        download_with_backoff(file_url, dest, expected_md5=md5)
+
+    logger.info(f"Empty encodings downloaded to {output_dir}")
+
+
+def _prepare_ingest(args):
+    """Download and convert pre-encoded Arrow data into Energon WebDataset."""
+    from primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest import (
+        StreamingIngestPipeline,
+    )
+
+    config = _load_ingest_config(args)
+
+    logger.info("=" * 80)
+    logger.info("Megatron Diffusion: Streaming Ingest Pipeline")
+    logger.info("=" * 80)
+
+    total_failed = 0
+    for ds in config["datasets"]:
+        logger.info(f"Processing dataset: {ds['name']} (split: {ds['split_name']})")
+        pipeline = StreamingIngestPipeline(
+            manifest_url=ds["manifest_url"],
+            input_dir=config["input_dir"],
+            output_dir=config["output_dir"],
+            split_name=ds["split_name"],
+            max_files=config.get("max_files"),
+            max_workers=config["max_workers"],
+            prefetch_depth=config["prefetch_depth"],
+        )
+        result = pipeline.run()
+        total_failed += result.get("files_failed", 0)
+
+    if "empty_encodings" in config:
+        logger.info("Downloading empty encodings...")
+        _download_empty_encodings(config)
+
+    if not config.get("no_finalize"):
+        from primus.backends.megatron.data.diffusion.preprocessing.finalize import (
+            finalize_energon_dataset,
+        )
+
+        try:
+            finalize_energon_dataset(
+                output_dir=config["output_dir"],
+                encoding="preencoded_numpy",
+                split_parts_patterns=[("train", "train/.*"), ("val", "val/.*")],
+            )
+        except Exception as e:
+            logger.error(f"Finalization failed: {e}")
+            raise
+
+    logger.info("=" * 80)
+    logger.info("Ingest complete!")
+    logger.info(f"  Output: {config['output_dir']}")
+    if total_failed:
+        logger.warning(
+            f"  {total_failed} file(s) failed across all datasets. "
+            f"Check failed_files.json in each split directory."
+        )
+    logger.info("=" * 80)
+
+
+def _prepare_automodel_cache(args):
+    """Build an offline pre-encoded cache for an AutoModel diffusion model."""
+    from primus.backends.nemo_automodel.data.registry import get_cache_builder
+
+    build_cache = get_cache_builder(args.model)
+
+    logger.info("=" * 80)
+    logger.info(f"AutoModel Diffusion: offline cache for '{args.model}'")
+    logger.info("=" * 80)
+
+    kwargs = dict(
+        image_dir=args.image_dir,
+        caption_dir=args.caption_dir,
+        output_dir=args.output_dir,
+        num_samples=args.num_samples,
+        resolution=args.resolution,
+        max_text_tokens=args.max_text_tokens,
+        tokenizer_source=args.tokenizer_source,
+        device=args.device,
+        dtype=args.dtype,
+        seed=args.seed,
+        shuffle=args.shuffle,
+    )
+    # Unset encoder sources mean "whatever this model defaults to", so leave them
+    # out rather than passing None over the builder's own default.
+    for name in ("vae_source", "text_encoder_source"):
+        value = getattr(args, name)
+        if value is not None:
+            kwargs[name] = value
+
+    metadata = build_cache(**kwargs)
+
+    logger.info("=" * 80)
+    logger.info("Cache complete!")
+    logger.info(f"  Output:  {args.output_dir}")
+    logger.info(f"  Samples: {metadata['num_samples']}")
+    logger.info(
+        f"  Latents: {metadata['in_channels']}x{metadata['grid_h']}x{metadata['grid_w']}"
+        f"  Text features: {metadata['llm_features_dim']}-d"
+    )
+    logger.info("Point the training config's data.dataloader.cache_dir at this directory.")
+    logger.info("=" * 80)
+
+
+def register_subcommand(subparsers: argparse._SubParsersAction) -> argparse.ArgumentParser:
+    """
+    Register 'primus data' subcommand for Megatron diffusion dataset preparation.
+
+    Provides dataset preparation utilities specifically for diffusion models
+    using the Megatron backend, outputting Energon WebDataset format.
+    """
+    parser = subparsers.add_parser(
+        "data",
+        help="Dataset preparation tools for diffusion training",
+        description=(
+            "Data preparation utilities for diffusion training.\n\n"
+            "Megatron backend (Energon WebDataset format):\n"
+            "  - diffusion-raw: Raw WebDataset (smaller, on-the-fly encoding)\n"
+            "  - diffusion-encoded: Pre-encoded WebDataset (larger, faster training)\n"
+            "  - diffusion-ingest: Stream pre-encoded Arrow data into WebDataset\n\n"
+            "NeMo AutoModel backend (flat per-sample cache):\n"
+            "  - automodel-cache: Pre-encode images + captions for the diffusion recipe\n\n"
+            "Multi-GPU support via torch.distributed (diffusion-* only):\n"
+            "  torchrun --nproc_per_node=8 primus data diffusion-raw ...\n"
+            "  torchrun --nproc_per_node=8 primus data diffusion-encoded ..."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    data_subparsers = parser.add_subparsers(dest="data_command", required=True, help="Data preparation mode")
+
+    # ===== primus data diffusion-raw =====
+    raw_parser = data_subparsers.add_parser(
+        "diffusion-raw",
+        help="Prepare raw Energon WebDataset (Megatron diffusion: images + captions)",
+        description=(
+            "Prepare raw Energon WebDataset for Megatron diffusion training.\n\n"
+            "Creates smaller datasets with raw images and captions.\n"
+            "Encoding (VAE/T5/CLIP) happens on-the-fly during training.\n\n"
+            "Output format: Energon WebDataset (tar shards)\n\n"
+            "Use this when:\n"
+            "  - Storage is limited\n"
+            "  - Experimenting with different encoders\n"
+            "  - Rapid prototyping\n\n"
+            "Example:\n"
+            "  primus data diffusion-raw \\\n"
+            "    --source-type huggingface \\\n"
+            "    --hf-dataset diffusers/pokemon-gpt4-captions \\\n"
+            "    --output-dir /data/raw_pokemon \\\n"
+            "    --image-size 1024 --center-crop\n\n"
+            "Multi-GPU:\n"
+            "  torchrun --nproc_per_node=8 primus data diffusion-raw ..."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    _add_common_args(raw_parser)
+
+    # Raw-specific arguments
+    raw_specific = raw_parser.add_argument_group("Raw Dataset Options")
+    raw_specific.add_argument(
+        "--image-format",
+        choices=["JPEG", "PNG", "WEBP"],
+        default="JPEG",
+        help="Output image format (default: JPEG)",
+    )
+    raw_specific.add_argument(
+        "--image-quality", type=int, default=95, help="JPEG/WEBP quality 1-100 (default: 95)"
+    )
+
+    # Finalization arguments
+    raw_finalize = raw_parser.add_argument_group("Dataset Finalization")
+    raw_finalize.add_argument(
+        "--no-finalize",
+        action="store_true",
+        dest="no_finalize",
+        help="Skip automatic dataset.yaml creation and energon prepare (default: finalize is ON)",
+    )
+    raw_finalize.add_argument(
+        "--train-split",
+        type=float,
+        default=1.0,
+        help="Training split ratio for finalization (default: 1.0 = 100%% train, 0%% val/test)",
+    )
+
+    raw_parser.set_defaults(func=lambda args, unknown: _prepare_raw(args))
+
+    # ===== primus data diffusion-encoded =====
+    encoded_parser = data_subparsers.add_parser(
+        "diffusion-encoded",
+        help="Prepare pre-encoded Energon WebDataset (Megatron diffusion: VAE/T5/CLIP)",
+        description=(
+            "Prepare pre-encoded Energon WebDataset for Megatron diffusion training.\n\n"
+            "Pre-encodes images with VAE and text with T5/CLIP, creating larger\n"
+            "datasets that enable faster training (no encoding overhead).\n\n"
+            "Output format: Energon WebDataset with pre-encoded tensors\n\n"
+            "Use this when:\n"
+            "  - Training on production datasets\n"
+            "  - Maximum training speed is needed\n"
+            "  - Storage space is available\n\n"
+            "Example:\n"
+            "  primus data diffusion-encoded \\\n"
+            "    --source-type huggingface \\\n"
+            "    --hf-dataset diffusers/pokemon-gpt4-captions \\\n"
+            "    --output-dir /data/encoded_pokemon \\\n"
+            "    --model-path black-forest-labs/FLUX.1-dev \\\n"
+            "    --batch-size 8 --precision bf16\n\n"
+            "Config file:\n"
+            "  primus data diffusion-encoded --config preprocessing.yaml\n\n"
+            "Multi-GPU:\n"
+            "  torchrun --nproc_per_node=8 primus data diffusion-encoded ..."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    # Config file support (optional, CLI args override config values)
+    encoded_parser.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="Path to YAML config file (optional, CLI args override config values)",
+    )
+
+    _add_common_args(encoded_parser)
+
+    # Encoded-specific arguments
+    model_group = encoded_parser.add_argument_group("Model Configuration")
+    model_group.add_argument(
+        "--model-path",
+        default="black-forest-labs/FLUX.1-dev",
+        help="Pretrained model path (HF or local). Default: black-forest-labs/FLUX.1-dev "
+        "(requires HF authentication -- see --hf-token-file)",
+    )
+    model_group.add_argument("--vae-path", default=None, help="Custom VAE path (overrides --model-path)")
+    model_group.add_argument("--t5-path", default=None, help="Custom T5 path (overrides --model-path)")
+    model_group.add_argument("--clip-path", default=None, help="Custom CLIP path (overrides --model-path)")
+    model_group.add_argument(
+        "--precision",
+        choices=["bf16", "fp16", "fp32"],
+        default="bf16",
+        help="Model precision (default: bf16)",
+    )
+    model_group.add_argument("--device", default="cuda", help="Device for encoding (default: cuda)")
+    model_group.add_argument("--batch-size", type=int, default=8, help="Encoding batch size (default: 8)")
+    model_group.add_argument(
+        "--t5-max-length",
+        type=int,
+        default=512,
+        help="T5 max sequence length (default: 512, use 256 for FLUX.1-schnell)",
+    )
+    model_group.add_argument(
+        "--vae-latent-mode",
+        choices=["presampled", "resample"],
+        default="presampled",
+        help=(
+            "VAE latent storage mode (default: presampled). "
+            "'presampled' stores a single sampled latent per image. "
+            "'resample' additionally stores mean and logvar so the training "
+            "loop can re-draw latents via reparameterization at every step."
+        ),
+    )
+
+    # Finalization arguments
+    encoded_finalize = encoded_parser.add_argument_group("Dataset Finalization")
+    encoded_finalize.add_argument(
+        "--no-finalize",
+        action="store_true",
+        dest="no_finalize",
+        help="Skip automatic dataset.yaml creation and energon prepare (default: finalize is ON)",
+    )
+    encoded_finalize.add_argument(
+        "--train-split",
+        type=float,
+        default=1.0,
+        help="Training split ratio for finalization (default: 1.0 = 100%% train, 0%% val/test)",
+    )
+
+    encoded_parser.set_defaults(func=lambda args, unknown: _prepare_encoded(args))
+
+    # ===== primus data diffusion-ingest =====
+    ingest_parser = data_subparsers.add_parser(
+        "diffusion-ingest",
+        help="Stream pre-encoded Arrow data into Energon WebDataset",
+        description=(
+            "Stream pre-encoded Arrow data into Energon WebDataset.\n\n"
+            "Downloads Apache Arrow IPC files from MLCommons R2 (or similar),\n"
+            "converts them to WebDataset tar shards in a single pass using a\n"
+            "producer-consumer prefetch architecture. Minimizes disk usage by\n"
+            "deleting each Arrow file after conversion.\n\n"
+            "Requires a YAML config file specifying datasets and parameters.\n\n"
+            "Example:\n"
+            "  primus data diffusion-ingest \\\n"
+            "    --config primus/configs/data/megatron/diffusion/"
+            "preprocessing/mlperf_flux1.yaml\n\n"
+            "Override output:\n"
+            "  primus data diffusion-ingest \\\n"
+            "    --config .../mlperf_flux1.yaml \\\n"
+            "    --output-dir /custom/path --max-files 5"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    ingest_parser.add_argument(
+        "--config",
+        type=str,
+        required=True,
+        help="Path to YAML config file with datasets list and pipeline settings",
+    )
+
+    ingest_output = ingest_parser.add_argument_group("Output Configuration")
+    ingest_output.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Output directory (overrides config output.output_dir)",
+    )
+    ingest_output.add_argument(
+        "--input-dir",
+        type=str,
+        default=None,
+        help="Temp directory for Arrow downloads (default: <output-dir>/_arrow_tmp)",
+    )
+
+    ingest_pipeline = ingest_parser.add_argument_group("Pipeline Configuration")
+    ingest_pipeline.add_argument(
+        "--max-files",
+        type=int,
+        default=None,
+        help="Limit Arrow files per dataset (for testing)",
+    )
+    ingest_pipeline.add_argument(
+        "--max-workers",
+        type=int,
+        default=None,
+        help="Concurrent download threads (overrides config, default: 4)",
+    )
+    ingest_pipeline.add_argument(
+        "--prefetch-depth",
+        type=int,
+        default=None,
+        help="Max Arrow files buffered on disk (overrides config, default: 6)",
+    )
+
+    ingest_finalize = ingest_parser.add_argument_group("Dataset Finalization")
+    ingest_finalize.add_argument(
+        "--no-finalize",
+        action="store_true",
+        dest="no_finalize",
+        help="Skip automatic Energon finalization (default: finalize is ON)",
+    )
+
+    ingest_parser.set_defaults(func=lambda args, unknown: _prepare_ingest(args))
+
+    # ===== primus data automodel-cache =====
+    # Sibling of the diffusion-* commands above, not an extension of them: those
+    # produce Energon WebDataset shards for the Megatron backend, this produces the
+    # flat per-sample cache the AutoModel diffusion recipe reads.
+    from primus.backends.nemo_automodel.data.registry import available_models
+
+    automodel_models = available_models()
+    cache_parser = data_subparsers.add_parser(
+        "automodel-cache",
+        help="Prepare a pre-encoded cache for AutoModel diffusion training",
+        description=(
+            "Build the offline pre-encoded cache for an AutoModel diffusion model.\n\n"
+            "The VAE and text encoder run once here, so training needs neither their\n"
+            "weights nor their memory. Output is a flat per-sample cache:\n"
+            "  <output-dir>/metadata.json    index + grid/feature dims\n"
+            "  <output-dir>/samples/<i>.pt   one encoded sample each\n\n"
+            "This is a single-GPU offline step and needs access to the (often gated)\n"
+            "encoder weights; set HF_TOKEN and point HF_HOME at persistent storage.\n\n"
+            "Selection is deterministic: images are taken in sorted order, so the same\n"
+            "source and arguments reproduce the same cache. Use --shuffle to opt out.\n\n"
+            f"Supported models: {', '.join(automodel_models)}\n\n"
+            "Example:\n"
+            "  primus data automodel-cache --model ideogram4 \\\n"
+            "    --image-dir   /dataset/pseudo-camera-10k/train \\\n"
+            "    --caption-dir /dataset/pseudo-camera-10k/caption \\\n"
+            "    --output-dir  /dataset/pcam_ideogram4_256 \\\n"
+            "    --num-samples 1024 --resolution 256 --max-text-tokens 128\n\n"
+            "Then point the training config's data.dataloader.cache_dir at --output-dir."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    cache_parser.add_argument(
+        "--model",
+        required=True,
+        choices=automodel_models,
+        help="Model whose cache format to build",
+    )
+
+    cache_source = cache_parser.add_argument_group("Source Configuration")
+    cache_source.add_argument("--image-dir", required=True, help="Directory of source images")
+    cache_source.add_argument(
+        "--caption-dir",
+        required=True,
+        help="Directory of captions, one <image stem>.txt per image",
+    )
+    cache_source.add_argument("--output-dir", required=True, help="Cache output directory")
+    cache_source.add_argument(
+        "--num-samples", type=int, default=1024, help="Samples to encode (default: 1024)"
+    )
+    cache_source.add_argument(
+        "--shuffle",
+        action="store_true",
+        help="Shuffle source order (seeded) before selecting; default is sorted order",
+    )
+    cache_source.add_argument("--seed", type=int, default=1234, help="Shuffle seed (default: 1234)")
+
+    cache_encode = cache_parser.add_argument_group("Encoding Options")
+    cache_encode.add_argument(
+        "--resolution", type=int, default=256, help="Square image resolution (default: 256)"
+    )
+    cache_encode.add_argument(
+        "--max-text-tokens",
+        type=int,
+        default=128,
+        help="Caption token budget; longer captions are skipped, not truncated (default: 128)",
+    )
+    cache_encode.add_argument(
+        "--dtype",
+        choices=["bf16", "fp16", "fp32"],
+        default="bf16",
+        help="Encoder compute dtype (default: bf16)",
+    )
+    cache_encode.add_argument("--device", default="cuda", help="Encoder device (default: cuda)")
+
+    cache_encoders = cache_parser.add_argument_group("Encoder Sources")
+    cache_encoders.add_argument(
+        "--vae-source", default=None, help="VAE repo or local path (default: the model's own)"
+    )
+    cache_encoders.add_argument(
+        "--text-encoder-source",
+        default=None,
+        help="Text-encoder repo or local path (default: the model's own)",
+    )
+    cache_encoders.add_argument(
+        "--tokenizer-source",
+        default=None,
+        help="Tokenizer repo or local path (default: alongside the text encoder)",
+    )
+
+    cache_parser.set_defaults(func=lambda args, unknown: _prepare_automodel_cache(args))
+
+    # Set default for parent parser (required by CLI framework, but never called due to required=True on subparsers)
+    parser.set_defaults(func=lambda args, unknown: None)
+
+    return parser
diff --git a/primus/cli/subcommands/node_smoke.py b/primus/cli/subcommands/node_smoke.py
new file mode 100644
index 000000000..5e1c695a4
--- /dev/null
+++ b/primus/cli/subcommands/node_smoke.py
@@ -0,0 +1,131 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+node_smoke CLI subcommand.
+
+Surfaces ``primus.tools.preflight.node_smoke`` as a first-class primus-cli
+subcommand so users can run it via the same dispatch chain as ``preflight``,
+``train``, ``benchmark``, etc.:
+
+    primus-cli direct -- node_smoke --tier2-perf
+    primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf
+    primus-cli slurm srun -N 4 -- container -- node_smoke --tier2-perf
+
+Design choices (consolidate-preflight-direct-wrappers plan, section 4):
+
+1. **No inner ``run`` keyword.** The standalone CLI has
+   ``run / aggregate / _per_gpu`` subparsers, but only ``run`` is user-facing.
+   Hoisting its flags onto the top-level ``node_smoke`` parser means the user
+   types ``primus-cli direct -- node_smoke --tier2-perf`` instead of
+   ``... -- node_smoke run --tier2-perf``.
+
+2. **Always aggregate on rank 0.** Every wrapper invocation runs ``_cmd_run``
+   on every rank, then ``_cmd_aggregate`` on rank 0 only. Aggregation takes
+   a few seconds (reads per-node JSONs, writes one report) and is what users
+   want ~100% of the time. The rare exceptions ("run only, don't aggregate"
+   / "aggregate only") stay reachable through the unchanged standalone CLI:
+
+       python -m primus.tools.preflight.node_smoke run ...
+       python -m primus.tools.preflight.node_smoke aggregate ...
+
+3. **``allow_abbrev=False``.** Mirrors the standalone ``run`` subparser at
+   ``primus/tools/preflight/node_smoke/cli.py:542``. Without this, an old
+   script that still passes ``--tier2`` (legacy flag name) would silently
+   match ``--tier2-perf`` as a prefix and run the wrong test set.
+
+4. **No ``--silent`` flag.** Silencing is handled exclusively by the bash
+   launcher (``primus-cli-direct.sh`` ``--silent`` before ``--``). Passing
+   ``--silent`` here will be rejected by argparse, which is the desired
+   behavior -- one knob in one place.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Any, List
+
+
+def run(args: Any, extra_args: List[str]) -> None:
+    """Two-phase dispatch: per-node ``_cmd_run`` on every rank, then rank-0
+    ``_cmd_aggregate`` with SLURM-resolved aggregator args.
+
+    Exit-code rule (matches the deleted wrapper):
+      - Non-rank-0: propagate ``_cmd_run`` exit code.
+      - Rank 0: aggregator exit code wins (it knows about MISSING nodes that
+        ``_cmd_run`` can't see). If aggregator returns 0 but run failed, we
+        still surface the run failure so a sick rank-0 box can't paint itself
+        green via a successful aggregate.
+    """
+    from primus.tools.preflight.node_smoke.cli import (
+        _cmd_aggregate,
+        _cmd_run,
+        _resolve_aggregate_args_from_slurm,
+    )
+
+    if extra_args:
+        # node_smoke uses allow_abbrev=False at the argparse level; reaching
+        # this path means the user passed something the parser didn't claim.
+        # Surface it loudly so a typo'd flag doesn't get silently dropped.
+        print(
+            f"[Primus:NodeSmoke] Unknown arguments: {extra_args}. "
+            f"Run `primus-cli node_smoke --help` for valid options.",
+        )
+        raise SystemExit(2)
+
+    rc_run = int(_cmd_run(args))
+    rank = int(os.environ.get("NODE_RANK", os.environ.get("SLURM_NODEID", "0")))
+    if rank != 0:
+        raise SystemExit(rc_run)
+
+    agg_ns = _resolve_aggregate_args_from_slurm(args)
+    rc_agg = int(_cmd_aggregate(agg_ns))
+
+    # Aggregator's exit code wins on rank 0 (it can detect MISSING nodes),
+    # but a failed per-node run on rank 0 itself must not be hidden by a
+    # successful aggregate. Surface the higher of the two.
+    raise SystemExit(max(rc_agg, rc_run))
+
+
+def register_subcommand(subparsers: argparse._SubParsersAction) -> argparse.ArgumentParser:
+    """Register ``node_smoke`` with the primus-cli main parser.
+
+    Hoists the ``run`` subparser's flags directly onto the top-level
+    ``node_smoke`` parser (no inner ``run`` keyword) plus the aggregator
+    tuning knobs the rank-0 aggregate step needs.
+    """
+    from primus.tools.preflight.node_smoke.cli import (
+        _add_aggregate_flags,
+        _add_run_flags,
+    )
+
+    parser = subparsers.add_parser(
+        "node_smoke",
+        help="Run per-node preflight smoke test on every node and aggregate on rank 0.",
+        description=(
+            "Node-local preflight smoke test. Each node runs independently "
+            "(no global rendezvous, no torch.distributed). On rank 0 the per-node "
+            "verdicts are aggregated into smoke_report.md + passing_nodes.txt + "
+            "failing_nodes.txt -- the latter two are directly consumable by "
+            "`srun --nodelist=` / `srun --exclude=`."
+        ),
+        # See module docstring point 3.
+        allow_abbrev=False,
+    )
+
+    # Attach the canonical run-side flag surface. Anything new added to
+    # `_add_run_flags` automatically flows here.
+    _add_run_flags(parser)
+
+    # Attach the aggregator-only flags. `--dump-path` / `--hbm-busy-threshold-gib`
+    # / `--gpu-activity-warn-pct` are already on the parser from `_add_run_flags`
+    # with identical defaults, so we skip them here to avoid argparse
+    # conflicting-option errors.
+    _add_aggregate_flags(parser, include_dump_path=False)
+
+    parser.set_defaults(func=run)
+    return parser
diff --git a/primus/cli/subcommands/train.py b/primus/cli/subcommands/train.py
index ac8d2da97..5cb2262ba 100644
--- a/primus/cli/subcommands/train.py
+++ b/primus/cli/subcommands/train.py
@@ -6,50 +6,19 @@
 
 from __future__ import annotations
 
-import sys
-from os import getenv
 from typing import List
 
 
-def _resolve_pretrain_runtime(args) -> str:
-    """
-    Resolve the runtime entry for pretrain.
-
-    Priority:
-      1) Explicit env override via PRIMUS_TRAIN_RUNTIME
-      2) Framework-based default (MaxText -> legacy, others -> core)
-    """
-    runtime_entry = getenv("PRIMUS_TRAIN_RUNTIME", "").strip().lower()
-    if runtime_entry in ("legacy", "core"):
-        return runtime_entry
-    if runtime_entry:
-        print(
-            f"[Primus:Train] Ignoring invalid PRIMUS_TRAIN_RUNTIME='{runtime_entry}'.",
-            file=sys.stderr,
-        )
-
-    # Default: use the new core runtime for all supported frameworks.
-    return "core"
-
-
 def run(args, overrides: List[str]):
     """
     Entry point for the 'train' subcommand.
     """
     if args.suite == "pretrain":
-        runtime_entry = _resolve_pretrain_runtime(args)
-
-        if runtime_entry == "core":
-            # New core runtime path: mirror `train_launcher.launch_train`.
-            from primus.core.runtime.train_runtime import PrimusRuntime
-
-            runtime = PrimusRuntime(args=args)
-            runtime.run_train_module(module_name="pre_trainer", overrides=overrides or [])
-        else:
-            # Legacy pretrain flow.
-            from primus.pretrain import launch_pretrain_from_cli
+        # All frameworks train via the core runtime.
+        from primus.core.runtime.train_runtime import PrimusRuntime
 
-            launch_pretrain_from_cli(args, overrides)
+        runtime = PrimusRuntime(args=args)
+        runtime.run_train_module(module_name="pre_trainer", overrides=overrides or [])
     elif args.suite == "posttrain":
         # Post-training (SFT/alignment) currently runs via the new core runtime.
         # It expects a training module named "sft_trainer" in the experiment config.
diff --git a/primus/configs/data/megatron/diffusion/README.md b/primus/configs/data/megatron/diffusion/README.md
new file mode 100644
index 000000000..827f1420b
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/README.md
@@ -0,0 +1,659 @@
+# Diffusion Data Configuration Guide
+
+This directory contains configuration files for preparing and organizing datasets for Megatron-based diffusion model training in Primus.
+
+## Directory Structure
+
+```
+primus/configs/data/megatron/diffusion/
+├── README.md                    # This file
+├── templates/                   # Dataset YAML templates
+│   ├── dataset_preencoded.yaml  # Pre-encoded dataset configuration
+│   ├── dataset_preencoded_numpy.yaml # Pre-encoded NumPy/MLPerf dataset configuration
+│   ├── dataset_raw.yaml         # Raw dataset configuration
+│   └── metadataset.yaml         # Multi-dataset mixing configuration
+└── preprocessing/               # CLI preprocessing configuration examples
+    ├── quickstart_pokemon.yaml  # Minimal quickstart (small HuggingFace dataset)
+    ├── example_base.yaml        # Comprehensive example with all options
+    ├── example_huggingface.yaml # HuggingFace dataset example
+    ├── example_directory.yaml   # Local directory example
+    ├── example_webdataset.yaml  # WebDataset input example
+    ├── text_to_image_2m_10k.yaml # 10K subset of text-to-image-2M (1024px)
+    └── mlperf_flux1.yaml        # MLPerf Flux1 streaming ingest configuration
+```
+
+## Quick Reference
+
+### What Goes Where?
+
+| File Type | Location | Purpose | When to Use |
+|-----------|----------|---------|-------------|
+| **Dataset Templates** | `templates/` | Define dataset structure | Auto-applied during finalization (or manually copy if using `--no-finalize`) |
+| **Preprocessing Configs** | `preprocessing/` | Configure data preparation | Pass to `primus data` CLI commands |
+| **Model Configs** | `primus/configs/models/` | Define model architecture | Reference in training configs |
+| **Training Configs** | `examples/megatron/configs/MI300X/diffusion/` | Configure training runs | Main config for training scripts |
+| **MLPerf Ingest Configs** | `preprocessing/` | Configure MLPerf streaming ingest | Pass to `primus data diffusion-ingest` |
+| **Path Override** | Command line | Customize output location | Add `--output-dir /your/path` to any command |
+| **Skip Finalize** | Command line | Skip automatic dataset setup | Add `--no-finalize` flag to skip finalization |
+
+**Tip**: Finalization (creating `dataset.yaml`, running `energon prepare`, validation) runs automatically by default. Use `--no-finalize` to skip it.
+
+---
+
+## Docker Volume Mounting and Data Paths
+
+When using `primus-cli direct --` commands inside containers, understanding Docker volume mounts is important for data persistence.
+
+### Default Docker Setup
+
+From [`tools/docker/start_container.sh`](../../../../../tools/docker/start_container.sh):
+```bash
+DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"}  # Default: ./data relative to repo
+# Mounted as: -v "${DATA_PATH}:${DATA_PATH}"
+```
+
+**What this means:**
+- **Default mount**: `<repo>/data` on host → `<repo>/data` (or `/workspace/Primus/data`) in container
+- **Config file paths**: Use `/workspace/Primus/data/...` by default
+- **Result**: Data automatically persists to your host machine's `data/` directory
+
+### Using Default Paths (Recommended)
+
+Config files already use the correct paths. Finalization runs automatically:
+
+```bash
+# Single command - data preparation + finalization + validation!
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+
+# Data is saved to: /workspace/Primus/data/encoded_pokemon/ (persisted to host)
+# Dataset is automatically indexed, validated, and ready for training!
+```
+
+Automatic finalization:
+- Creates `.nv-meta/dataset.yaml` with `CrudeWebdataset` configuration
+- Runs `energon prepare` to index the dataset
+- Validates with Primus's custom validation (metadata, sample counts, API spot-check)
+- No manual post-processing needed
+
+### Easy Path Override
+
+**Override any path using `--output-dir` flag:**
+
+```bash
+# Customize output location
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --output-dir /workspace/Primus/data/my_custom_output
+
+# Or use a different mounted volume
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --output-dir /mnt/shared_storage/datasets/encoded
+```
+
+### Advanced: Using /data Paths
+
+If you prefer to use `/data/...` paths (as shown in some examples):
+
+```bash
+# Set DATA_PATH before starting container
+export DATA_PATH=/data
+
+# Then start/restart container with new mount
+# Now all /data/... paths will work as written
+```
+
+### Quick Comparison
+
+| Approach | Command | Persistence | Notes |
+|----------|---------|-------------|-------|
+| **Default** | Use config as-is | ✅ Yes | Easiest, works immediately |
+| **Override** | Add `--output-dir /path` | ✅ Yes (if path is mounted) | Flexible, per-command |
+| **DATA_PATH** | `export DATA_PATH=/data` | ✅ Yes | Cleaner paths, requires setup |
+
+---
+
+## Dataset Templates (`templates/`)
+
+These files define how Energon loads and processes your prepared datasets. With automatic finalization (the default), these templates are applied automatically -- you only need them for manual workflows or customization.
+
+### [`dataset_preencoded.yaml`](templates/dataset_preencoded.yaml)
+
+**Purpose**: Configure a pre-encoded dataset (contains VAE latents and text embeddings).
+
+**When to use**: Production training (2-3x faster than raw mode).
+
+**Usage** (automatic finalization handles steps 2-3):
+```bash
+# Single command -- finalization is automatic:
+primus-cli direct -- data diffusion-encoded \
+  --source-type directory \
+  --input-dir /data/raw_images \
+  --output-dir /workspace/Primus/data/encoded_dataset \
+  --hf-token-file /path/to/.hf_token
+
+# Use in training config:
+# data:
+#   dataset_path: /workspace/Primus/data/encoded_dataset
+
+# Manual workflow (if using --no-finalize):
+# 1. Copy template: cp templates/dataset_preencoded.yaml <output>/.nv-meta/dataset.yaml
+# 2. Run indexing: energon prepare <output> --num-workers 8
+```
+
+### [`dataset_raw.yaml`](templates/dataset_raw.yaml)
+
+**Purpose**: Configure a raw dataset (contains original images and captions).
+
+**When to use**: Experimentation, rapid prototyping, when disk space is limited.
+
+**Usage** (automatic finalization handles steps 2-3):
+```bash
+# Single command -- finalization is automatic:
+primus-cli direct -- data diffusion-raw \
+  --source-type directory \
+  --input-dir /data/raw_images \
+  --output-dir /workspace/Primus/data/raw_dataset
+
+# Use in training config with encoder configs:
+# data:
+#   dataset_path: /workspace/Primus/data/raw_dataset
+#   encoder_configs:
+#     vae: {...}
+#     t5: {...}
+#     clip: {...}
+```
+
+### [`dataset_preencoded_numpy.yaml`](templates/dataset_preencoded_numpy.yaml)
+
+**Purpose**: Configure an MLPerf pre-encoded numpy dataset (bfloat16 tensors stored as NumPy uint16 bytes).
+
+**When to use**: Training with MLPerf Flux1 pre-encoded data ingested via `diffusion-ingest`.
+
+**Sample keys**: `t5.bytes`, `clip.bytes`, `mean.bytes`, `logvar.bytes`, `.json`
+
+**Usage**: Applied automatically by `diffusion-ingest` finalization. For manual use:
+```bash
+cp templates/dataset_preencoded_numpy.yaml <output>/.nv-meta/dataset.yaml
+energon prepare <output> --num-workers 8
+```
+
+### [`metadataset.yaml`](templates/metadataset.yaml)
+
+**Purpose**: Combine multiple datasets with different weights.
+
+**When to use**:
+- Mixing pre-encoded and raw datasets
+- Combining multiple data sources
+- Creating weighted train/val splits
+
+**Usage**:
+```bash
+# 1. Prepare individual datasets first (pre-encoded and/or raw)
+
+# 2. Copy and customize metadataset template
+cp primus/configs/data/megatron/diffusion/templates/metadataset.yaml \
+   /workspace/Primus/data/combined_dataset.yaml
+
+# 3. Edit paths in combined_dataset.yaml to point to your datasets
+
+# 4. Use in training config
+# data:
+#   dataset_path: /workspace/Primus/data/combined_dataset.yaml
+```
+
+---
+
+## Preprocessing Configs (`preprocessing/`)
+
+These files are passed to the `primus data` CLI commands to configure dataset preparation. They are **NOT copied to dataset directories**.
+
+### [`quickstart_pokemon.yaml`](preprocessing/quickstart_pokemon.yaml)
+
+**Purpose**: Minimal quickstart config for the Pokemon dataset (256px).
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`example_base.yaml`](preprocessing/example_base.yaml)
+
+**Purpose**: Comprehensive example showing all available preprocessing options.
+
+**Contents**: All configuration sections with inline documentation.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`example_huggingface.yaml`](preprocessing/example_huggingface.yaml)
+
+**Purpose**: Prepare a dataset from HuggingFace Hub.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`example_directory.yaml`](preprocessing/example_directory.yaml)
+
+**Purpose**: Prepare a dataset from a local directory of images and captions.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`example_webdataset.yaml`](preprocessing/example_webdataset.yaml)
+
+**Purpose**: Prepare a dataset from existing WebDataset tar files.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`text_to_image_2m_10k.yaml`](preprocessing/text_to_image_2m_10k.yaml)
+
+**Purpose**: Prepare the 10K high-resolution subset (1024x1024) from the text-to-image-2M dataset.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### [`mlperf_flux1.yaml`](preprocessing/mlperf_flux1.yaml)
+
+**Purpose**: MLPerf Flux1 streaming ingest -- downloads pre-encoded Arrow files from MLCommons R2 and converts to Energon WebDataset tar shards.
+
+**Datasets**: CC12M (train) and COCO (val), plus empty encodings for CFG dropout.
+
+**Usage**:
+```bash
+primus-cli direct -- data diffusion-ingest \
+  --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
+
+# Limit files for testing:
+primus-cli direct -- data diffusion-ingest \
+  --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml \
+  --max-files 5
+```
+
+Note: `--hf-token-file` is only required when downloading gated models from HuggingFace (e.g., the default FLUX.1-dev). If using local encoder paths via `--model-path`, it can be omitted.
+
+---
+
+## Path Configuration Patterns
+
+### Pattern 1: Default (Recommended)
+
+Use config file paths as-is -- they work out of the box:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+
+# Output automatically goes to: /workspace/Primus/data/encoded_pokemon
+# This persists to: <your-repo>/data/encoded_pokemon on host
+```
+
+### Pattern 2: Custom Path Override (Flexible)
+
+Override output location for any command:
+
+```bash
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --output-dir /workspace/Primus/data/custom_output
+
+# Or use a different mounted volume
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token \
+  --output-dir /mnt/shared_storage/datasets
+```
+
+### Pattern 3: DATA_PATH Environment Variable (Advanced)
+
+Set `DATA_PATH` for cleaner paths:
+
+```bash
+# Before starting container
+export DATA_PATH=/data
+
+# Start/restart container (mounts /data)
+# Now configs using /data/... will work as written
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+```
+
+### Verification
+
+Check where your data was written:
+
+```bash
+# For default paths
+ls -lh /workspace/Primus/data/encoded_pokemon/
+
+# On host machine
+ls -lh <your-repo>/data/encoded_pokemon/
+
+# For custom paths
+ls -lh /your/custom/path/
+```
+
+---
+
+## Complete Workflow
+
+### Pre-encoded Mode (Recommended for Production)
+
+```bash
+# Step 1: Prepare, finalize, and validate dataset (all automatic)
+primus-cli direct -- data diffusion-encoded \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+  --hf-token-file /path/to/.hf_token
+
+# Output: /workspace/Primus/data/encoded_pokemon (persisted to host)
+# Finalization (dataset.yaml + energon prepare + validation) runs automatically.
+
+# Step 2: Configure training
+# Edit examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml:
+# data:
+#   dataset_path: /workspace/Primus/data/encoded_pokemon
+
+# Step 3: Train
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+  bash examples/run_pretrain.sh
+```
+
+### Raw Mode (On-the-fly Encoding)
+
+```bash
+# Step 1: Prepare, finalize, and validate dataset (all automatic)
+primus-cli direct -- data diffusion-raw \
+  --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml
+
+# Output: /workspace/Primus/data/raw_dataset (persisted to host)
+
+# Step 2: Configure training with encoders
+# Edit examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml:
+# data:
+#   dataset_path: /workspace/Primus/data/raw_dataset
+#   encoder_configs:
+#     vae: { model_path: black-forest-labs/FLUX.1-dev, ... }
+#     t5: { model_path: black-forest-labs/FLUX.1-dev, ... }
+#     clip: { model_path: black-forest-labs/FLUX.1-dev, ... }
+
+# Step 3: Train
+EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \
+  bash examples/run_pretrain.sh
+```
+
+### MLPerf Ingest Mode (Streaming Download + Conversion)
+
+```bash
+# Step 1: Download and convert MLPerf Arrow data to Energon WebDataset
+# Downloads ~1.2 TB of pre-encoded data with minimal temporary disk usage.
+primus-cli direct -- data diffusion-ingest \
+  --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
+
+# Output: /workspace/Primus/data/mlperf_flux1/ (persisted to host)
+# Finalization (dataset.yaml + energon prepare + validation) runs automatically.
+# Pipeline supports resume -- re-run the same command to continue after interruption.
+
+# Step 2: Point the training config at the prepared data.
+# The MLPerf reproduction configs already set the required VAE normalization
+# (vae_latent_mode: resample, vae_scale: 0.3611, vae_shift: 0.1159); just set
+# data_path (or the PRIMUS_DIFFUSION_DATA_PATH env var) to /workspace/Primus/data/mlperf_flux1.
+
+# Step 3: Train (MLPerf benchmark reproduction)
+EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml \
+  bash examples/run_pretrain.sh
+```
+
+**Output directory structure:**
+```
+mlperf_flux1/
+├── train/              # CC12M shards (shard_000000.tar, ...)
+├── val/                # COCO shards
+├── empty_encodings/    # empty_t5_encodings.npy, empty_clip_encodings.npy
+└── .nv-meta/           # Energon index + dataset.yaml (auto-generated)
+```
+
+**Pipeline features:**
+- Parallel downloads with configurable worker threads (`max_workers`)
+- Bounded disk usage via semaphore (`prefetch_depth` Arrow files on disk at once)
+- Automatic retry with exponential backoff for HTTP 429/503 and MD5 mismatches
+- Resume support -- re-run to skip already-completed shards
+- Skip-and-log for individual file failures (`failed_files.json`)
+
+---
+
+## Configuration Hierarchy
+
+Understanding how different configs relate to each other:
+
+```mermaid
+graph TB
+    subgraph preprocessing [Preprocessing Phase]
+        prepConfig[preprocessing/example_*.yaml]
+        cliCommand[primus data CLI]
+        prepConfig --> cliCommand
+    end
+
+    subgraph datasetDir [Dataset Directory]
+        tarFiles[WebDataset *.tar files]
+        datasetYaml["dataset.yaml (auto-generated)"]
+        energonIndex[.nv-meta/ Energon index]
+
+        tarFiles --> energonIndex
+        datasetYaml --> energonIndex
+    end
+
+    subgraph training [Training Phase]
+        trainConfig[Training Config examples/*.yaml]
+        modelConfig[Model Config primus/configs/models/]
+        encoderConfig[Encoder Config encoders.yaml]
+
+        trainConfig --> modelConfig
+        trainConfig -.raw mode only.-> encoderConfig
+    end
+
+    cliCommand --> tarFiles
+    energonIndex --> trainConfig
+```
+
+### Key Points:
+
+1. **Preprocessing configs** -- Used once during data preparation
+2. **Dataset templates** -- Auto-applied during finalization (or manually copied if using `--no-finalize`)
+3. **Training configs** -- Main config file that references dataset path
+4. **Model configs** -- Define architecture, referenced by training configs
+5. **Encoder configs** -- Only needed for raw (on-the-fly) mode
+
+---
+
+## Common Patterns
+
+### Pattern 1: Single Pre-encoded Dataset
+
+```yaml
+# Training config
+data:
+  dataset_path: /workspace/Primus/data/my_dataset/dataset.yaml  # Points to dataset_preencoded.yaml copy
+  micro_batch_size: 2
+  global_batch_size: 16
+```
+
+### Pattern 2: Mixed Pre-encoded + Raw
+
+```yaml
+# Copy and customize templates/metadataset.yaml
+data:
+  dataset_path: /workspace/Primus/data/mixed_dataset.yaml  # Points to customized metadataset.yaml
+  encoder_configs:  # Needed for raw datasets
+    vae: {...}
+    t5: {...}
+    clip: {...}
+```
+
+### Pattern 3: Multiple Pre-encoded Datasets
+
+```yaml
+# Use metadataset.yaml to combine
+data:
+  dataset_path: /workspace/Primus/data/combined.yaml  # Metadataset config
+  # No encoder_configs needed if all datasets are pre-encoded
+```
+
+---
+
+## File Formats
+
+### Dataset YAML (in dataset directories)
+
+```yaml
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: preencoded  # or 'raw' or 'preencoded_numpy'
+# Optional filters and settings
+```
+
+**Encoding types:**
+- `preencoded` -- Primus-encoded PyTorch `.pth` format (VAE latents + text embeddings)
+- `preencoded_numpy` -- MLPerf NumPy uint16 format (bfloat16 as `.bytes` entries)
+- `raw` -- Original images and captions (encoded on-the-fly during training)
+
+### Metadataset YAML
+
+```yaml
+__module__: megatron.energon
+__class__: Metadataset
+splits:
+  train:
+    datasets:
+      - weight: 0.8
+        path: /path/to/dataset1/
+      - weight: 0.2
+        path: /path/to/dataset2/
+```
+
+### Preprocessing Config YAML
+
+```yaml
+source:
+  type: huggingface  # or directory, webdataset
+  # source-specific options
+output:
+  output_dir: /data/output
+  shard_size: 1000
+model:
+  model_path: black-forest-labs/FLUX.1-dev
+  batch_size: 8
+```
+
+---
+
+## Additional Resources
+
+### Documentation
+- **Training Guide**: [`examples/megatron/diffusion/README.md`](../../../../../examples/megatron/diffusion/README.md)
+- **Energon Integration**: [`docs/04-technical-guides/diffusion-models/energon_integration.md`](../../../../../docs/04-technical-guides/diffusion-models/energon_integration.md)
+- **FP8 Training Guide**: [`docs/04-technical-guides/diffusion-models/fp8_training.md`](../../../../../docs/04-technical-guides/diffusion-models/fp8_training.md)
+
+### Related Configs
+- **Encoder Configs**: [`primus/configs/models/megatron/diffusion/encoders.yaml`](../../../models/megatron/diffusion/encoders.yaml)
+- **Model Architecture**: [`primus/configs/models/megatron/diffusion/`](../../../models/megatron/diffusion/)
+- **Training Configs (MI300X)**: [`examples/megatron/configs/MI300X/diffusion/`](../../../../../examples/megatron/configs/MI300X/diffusion/)
+- **Training Configs (MI355X)**: [`examples/megatron/configs/MI355X/diffusion/`](../../../../../examples/megatron/configs/MI355X/diffusion/)
+
+---
+
+## Troubleshooting
+
+### "No such file or directory: dataset.yaml"
+
+**Problem**: Training config references a dataset path that doesn't exist.
+
+**Solution**: If you used `--no-finalize`, you need to manually copy the appropriate template from `templates/` to your dataset's `.nv-meta/` directory and run `energon prepare`. Otherwise, re-run preprocessing without `--no-finalize` (the default).
+
+### "Unknown subflavor: encoding"
+
+**Problem**: Old dataset.yaml format or missing subflavors field.
+
+**Solution**: Use the updated templates from this directory (dataset_preencoded.yaml or dataset_raw.yaml).
+
+### "Encoder not found" (raw mode)
+
+**Problem**: Training with raw dataset but encoder_configs not specified.
+
+**Solution**: Add encoder_configs to your training config or switch to pre-encoded mode.
+
+### NaN loss with MLPerf data
+
+**Problem**: Training on MLPerf-ingested data produces NaN or diverging loss.
+
+**Solution**: Ensure your training config includes the required normalization constants:
+```yaml
+vae_latent_mode: resample
+vae_scale: 0.3611
+vae_shift: 0.1159
+```
+These are mandatory because `forward_step.py` unconditionally applies `vae_scale * (latents - vae_shift)` after resampling from mean/logvar.
+
+### Missing mean/logvar keys in MLPerf data
+
+**Problem**: Training errors about missing `mean` or `logvar` fields.
+
+**Solution**: Verify your dataset's `.nv-meta/dataset.yaml` uses `encoding: preencoded_numpy` (not `preencoded`). The `preencoded_numpy` cooker expects `.bytes` keys while `preencoded` expects `.pth` keys.
+
+### Preprocessing config not found
+
+**Problem**: Trying to use old config paths that have been moved.
+
+**Solution**: Use new paths:
+- Old: `examples/megatron/diffusion/configs/data_preprocessing.yaml`
+- New: `primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml`
+
+---
+
+## Migration from Old Structure
+
+If you have configs using old paths, update them:
+
+| Old Path | New Path |
+|----------|----------|
+| `examples/megatron/diffusion/configs/preencoded_dataset.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml` |
+| `examples/megatron/diffusion/configs/flux_dataset_preencoded.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml` |
+| `examples/megatron/diffusion/configs/raw_dataset.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml` |
+| `examples/megatron/diffusion/configs/flux_dataset_raw.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml` |
+| `examples/megatron/diffusion/configs/metadataset.yaml` | `primus/configs/data/megatron/diffusion/templates/metadataset.yaml` |
+| `examples/megatron/diffusion/data/test_pokemon_config.yaml` | `primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml` |
+| `examples/megatron/diffusion/configs/data_preprocessing.yaml` | `primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml` |
+
+---
+
+**Questions or issues?** See the main documentation or open an issue on GitHub.
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml b/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml
new file mode 100644
index 000000000..4031e2783
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml
@@ -0,0 +1,31 @@
+# Pre-encode COCO 2014 train (82,783 images) for Flux-Schnell training
+# 256px images, 256-token T5 sequences (pass --t5-max-length 256 on CLI)
+#
+# Dataset: AbdoTW/COCO_2014 train split (full COCO 2014 train2014)
+# Fields: image (PIL), caption (list[str] -> first caption taken automatically)
+#
+# Usage (4 GPU on devices 4-7):
+#   HIP_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node=4 \
+#     /workspace/Primus/primus/cli/main.py data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml \
+#     --t5-max-length 256 --hf-token-file /workspace/Primus/.hf_token
+
+source:
+  type: huggingface
+  hf_dataset: AbdoTW/COCO_2014
+  hf_split: train
+
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/coco2014_train_schnell_256_encoded
+  shard_size: 1000
+  compress: false
+
+model:
+  model_path: black-forest-labs/FLUX.1-schnell
+  precision: bf16
+  device: cuda
+  batch_size: 8
+
+image:
+  image_size: 256
+  center_crop: false
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml b/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml
new file mode 100644
index 000000000..f033192be
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml
@@ -0,0 +1,36 @@
+# Pre-encode COCO for Flux-Schnell training (VAE resample mode)
+# 256px images, 256-token T5 sequences (pass --t5-max-length 256 on CLI)
+# vae_latent_mode: resample — stores mean+logvar so latents are re-drawn
+# via reparameterization (mean + exp(0.5*logvar) * randn) at every training step.
+#
+# Dataset: UCSC-VLAA/Recap-COCO-30K (30K COCO images with captions)
+# Fields: image (PIL), caption (str)
+#
+# Usage (8 GPU):
+#   torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \
+#     data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml \
+#     --t5-max-length 256 \
+#     --hf-token-file /workspace/Primus/.hf_token
+
+source:
+  type: huggingface
+  hf_dataset: UCSC-VLAA/Recap-COCO-30K
+  hf_split: train
+
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/coco_schnell_256_resample_encoded
+  shard_size: 1000
+  max_samples: null
+  compress: false
+
+model:
+  model_path: black-forest-labs/FLUX.1-schnell
+  precision: bf16
+  device: cuda
+  batch_size: 8
+  vae_latent_mode: resample
+
+image:
+  image_size: 256
+  center_crop: false
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml
new file mode 100644
index 000000000..ee6cf4332
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml
@@ -0,0 +1,196 @@
+# Data Preprocessing Configuration for Diffusion Models
+#
+# This comprehensive config file shows all available preprocessing parameters.
+#
+# ============================================================================
+# Quick Start (Recommended - Single Command)
+# ============================================================================
+# Prepare and finalize dataset in one command:
+#
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# The default encoder model (FLUX.1-dev) is gated and requires a
+# HuggingFace token. Get one at https://huggingface.co/settings/tokens
+#
+# Finalization (dataset.yaml + energon prepare) runs automatically.
+# To skip it, pass --no-finalize.
+#
+# ============================================================================
+# Advanced Usage
+# ============================================================================
+# Override specific parameters:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --batch-size 16 \
+#     --output-dir /your/custom/path
+#
+# Priority: CLI arguments > YAML config > defaults
+# ============================================================================
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+# Default paths use Docker mounted paths that persist to host:
+#   /workspace/Primus/data/...  →  <your-repo>/data/... on host
+#
+# EASY OVERRIDE: Customize output with --output-dir flag:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /your/custom/path
+#
+# For /data/... paths, set: export DATA_PATH=/data (before starting container)
+# ============================================================================
+
+# ============================================================================
+# Source Configuration
+# ============================================================================
+source:
+  # Type of data source (required)
+  type: huggingface  # Options: huggingface, directory, webdataset
+
+  # HuggingFace source options (when type: huggingface)
+  hf_dataset: diffusers/pokemon-gpt4-captions  # HuggingFace dataset name
+  hf_split: train                              # Dataset split to use
+
+  # Directory source options (when type: directory)
+  # input_dir: /path/to/images  # Directory containing images and captions
+
+  # WebDataset source options (when type: webdataset)
+  # input_path: '/path/to/dataset/*.tar'  # Glob pattern for WebDataset tars
+
+# ============================================================================
+# Data Format Configuration (Optional)
+# ============================================================================
+# Specify how to extract image and caption from dataset samples.
+# If not specified, uses automatic detection with common field names:
+#   - Images: 'image', 'jpg', 'jpeg', 'png', 'img', 'photo'
+#   - Captions: 'caption', 'text', 'txt', 'description', 'prompt'
+#
+# Use this section for datasets with non-standard field names or JSON-encoded data.
+# data_format:
+#   # Single field names (simplest approach)
+#   image_key: jpg                  # e.g., images are in 'jpg' field
+#   caption_key: json.caption       # e.g., captions in JSON at 'caption' key
+#
+#   # Alternative: specify multiple fallback keys (tried in order)
+#   # image_keys: ['image', 'jpg', 'png']
+#   # caption_keys: ['caption', 'json.caption', 'json.text']
+#
+# Examples:
+#   # For text-to-image-2M dataset:
+#   data_format:
+#     image_key: jpg
+#     caption_key: json.caption
+#
+#   # For datasets with nested JSON metadata:
+#   data_format:
+#     image_key: image
+#     caption_key: json.metadata.description
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_pokemon  # Output directory (persisted to host)
+  shard_size: 1000                                    # Number of samples per output shard
+  max_samples: null                                   # Maximum samples to process (null = all)
+  compress: false                                     # Whether to compress output with gzip
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+model:
+  # Base model path (HuggingFace repo or local path)
+  model_path: black-forest-labs/FLUX.1-dev
+
+  # Override individual encoder paths (null means use model_path)
+  vae_path: null    # Custom VAE path
+  t5_path: null     # Custom T5-XXL path
+  clip_path: null   # Custom CLIP-L path
+
+  # Model settings
+  precision: bf16   # Options: bf16, fp16, fp32
+  device: cuda      # Device for encoding
+  batch_size: 8     # Encoding batch size
+
+# ============================================================================
+# Image Preprocessing
+# ============================================================================
+image:
+  image_size: 1024    # Target image size (height and width)
+  center_crop: false  # Whether to center crop before resize
+
+# ============================================================================
+# Examples
+# ============================================================================
+#
+# Example 1: From HuggingFace with custom batch size
+# ---------------------------------------------------
+# source:
+#   type: huggingface
+#   hf_dataset: diffusers/pokemon-gpt4-captions
+# output:
+#   output_dir: /workspace/Primus/data/pokemon
+# model:
+#   batch_size: 16
+#   precision: bf16
+#
+# Example 2: From local directory
+# --------------------------------
+# source:
+#   type: directory
+#   input_dir: /data/raw_images
+# output:
+#   output_dir: /workspace/Primus/data/encoded_images
+# model:
+#   model_path: black-forest-labs/FLUX.1-dev
+#
+# Example 3: From WebDataset with custom encoders
+# ------------------------------------------------
+# source:
+#   type: webdataset
+#   input_path: '/data/dataset/*.tar'
+# output:
+#   output_dir: /workspace/Primus/data/encoded
+# model:
+#   vae_path: madebyollin/sdxl-vae-fp16-fix
+#   t5_path: google/t5-v1_1-xxl
+#   clip_path: openai/clip-vit-large-patch14
+#
+# ============================================================================
+# Usage
+# ============================================================================
+#
+# Single GPU (uses config paths):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Multi-GPU (8 GPUs):
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Override output directory:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /workspace/Primus/data/my_output
+#
+# Override multiple values:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /your/custom/path \
+#     --batch-size 16 \
+#     --max-samples 1000
+#
+# For more examples, see:
+#   - example_huggingface.yaml  (HuggingFace dataset)
+#   - example_directory.yaml    (Local directory)
+#   - example_webdataset.yaml   (WebDataset input)
+#
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml
new file mode 100644
index 000000000..cb80838be
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml
@@ -0,0 +1,130 @@
+# Preprocessing Configuration Example: Local Directory
+#
+# This example shows how to prepare a dataset from a local directory.
+#
+# ============================================================================
+# Quick Start (Recommended - Single Command)
+# ============================================================================
+# Prepare and finalize dataset in one command:
+#
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# The default encoder model (FLUX.1-dev) is gated and requires a
+# HuggingFace token. Get one at https://huggingface.co/settings/tokens
+#
+# Finalization (dataset.yaml + energon prepare) runs automatically.
+# To skip it, pass --no-finalize.
+#
+# ============================================================================
+# Advanced Usage
+# ============================================================================
+# Custom paths:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-dir /your/input/path \
+#     --output-dir /your/output/path
+# ============================================================================
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+# Default paths use Docker mounted paths that persist to host:
+#   /workspace/Primus/data/...  →  <your-repo>/data/... on host
+#
+# EASY OVERRIDE: Customize paths with command-line flags:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-dir /your/input/path \
+#     --output-dir /your/output/path
+#
+# For /data/... paths, set: export DATA_PATH=/data (before starting container)
+# ============================================================================
+
+# ============================================================================
+# Source Configuration (Directory)
+# ============================================================================
+source:
+  type: directory
+  input_dir: /data/raw_images  # Directory containing images and captions
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_images  # Output directory (persisted to host)
+  shard_size: 1000                                   # Number of samples per output shard
+  max_samples: null                                  # Maximum samples to process (null = all)
+  compress: false                                    # Whether to compress output with gzip
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+model:
+  # Base model path (HuggingFace repo or local path)
+  model_path: black-forest-labs/FLUX.1-dev
+
+  # Override individual encoder paths (null means use model_path)
+  vae_path: null    # Custom VAE path
+  t5_path: null     # Custom T5-XXL path
+  clip_path: null   # Custom CLIP-L path
+
+  # Model settings
+  precision: bf16   # Options: bf16, fp16, fp32
+  device: cuda      # Device for encoding
+  batch_size: 8     # Encoding batch size
+
+# ============================================================================
+# Image Preprocessing
+# ============================================================================
+image:
+  image_size: 1024    # Target image size (height and width)
+  center_crop: false  # Whether to center crop before resize
+
+# ============================================================================
+# Input Directory Structure
+# ============================================================================
+#
+# Your input_dir should contain image-caption pairs:
+#
+# /data/raw_images/
+# ├── image_001.jpg
+# ├── image_001.txt
+# ├── image_002.jpg
+# ├── image_002.txt
+# └── ...
+#
+# Each .jpg file should have a corresponding .txt file with the same name.
+# Supported image formats: .jpg, .jpeg, .png, .webp
+#
+# ============================================================================
+# Usage
+# ============================================================================
+#
+# Single GPU (uses config paths):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Multi-GPU (8 GPUs):
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Override paths:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-dir /your/input/path \
+#     --output-dir /your/output/path
+#
+# Override specific values:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /workspace/Primus/data/my_output \
+#     --batch-size 16
+#
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml
new file mode 100644
index 000000000..eed538b31
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml
@@ -0,0 +1,124 @@
+# Preprocessing Configuration Example: HuggingFace Dataset
+#
+# This example shows how to prepare a dataset from HuggingFace Hub.
+#
+# ============================================================================
+# Quick Start (Recommended - Single Command)
+# ============================================================================
+# Prepare and finalize dataset in one command (no manual steps!):
+#
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# The default encoder model (FLUX.1-dev) is gated and requires a
+# HuggingFace token. Get one at https://huggingface.co/settings/tokens
+#
+# Finalization (dataset.yaml + energon prepare) runs automatically.
+# To skip it, pass --no-finalize.
+#
+# ============================================================================
+# Advanced Usage
+# ============================================================================
+# Skip auto-finalization (manual post-processing):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token --no-finalize
+#
+# Custom output path:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /your/custom/path
+#
+# Custom train/val/test split (default: 100% train):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --train-split 0.8  # 80% train, 10% val, 10% test
+# ============================================================================
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+# Default output_dir uses Docker mounted path that persists to host:
+#   /workspace/Primus/data/...  →  <your-repo>/data/... on host
+#
+# EASY OVERRIDE: Customize output location with --output-dir flag:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /your/custom/path
+#
+# For /data/... paths, set: export DATA_PATH=/data (before starting container)
+# ============================================================================
+
+# ============================================================================
+# Source Configuration (HuggingFace)
+# ============================================================================
+source:
+  type: huggingface
+  hf_dataset: diffusers/pokemon-gpt4-captions  # HuggingFace dataset name
+  hf_split: train                              # Dataset split to use
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_pokemon  # Output directory (persisted to host)
+  shard_size: 1000                                    # Number of samples per output shard
+  max_samples: null                                   # Maximum samples to process (null = all)
+  compress: false                                     # Whether to compress output with gzip
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+model:
+  # Base model path (HuggingFace repo or local path)
+  model_path: black-forest-labs/FLUX.1-dev
+
+  # Override individual encoder paths (null means use model_path)
+  vae_path: null    # Custom VAE path
+  t5_path: null     # Custom T5-XXL path
+  clip_path: null   # Custom CLIP-L path
+
+  # Model settings
+  precision: bf16   # Options: bf16, fp16, fp32
+  device: cuda      # Device for encoding
+  batch_size: 8     # Encoding batch size
+
+# ============================================================================
+# Image Preprocessing
+# ============================================================================
+image:
+  image_size: 1024    # Target image size (height and width)
+  center_crop: false  # Whether to center crop before resize
+
+# ============================================================================
+# Usage
+# ============================================================================
+#
+# Single GPU (uses default output_dir from config):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Multi-GPU (8 GPUs):
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Override output directory:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /workspace/Primus/data/custom_output
+#
+# Override multiple values:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /my/custom/path \
+#     --batch-size 16 \
+#     --max-samples 1000
+#
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml
new file mode 100644
index 000000000..ab5a1e296
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml
@@ -0,0 +1,129 @@
+# Preprocessing Configuration Example: WebDataset Input
+#
+# This example shows how to prepare a dataset from existing WebDataset tar files.
+#
+# ============================================================================
+# Quick Start (Recommended - Single Command)
+# ============================================================================
+# Prepare and finalize dataset in one command:
+#
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# The default encoder model (FLUX.1-dev) is gated and requires a
+# HuggingFace token. Get one at https://huggingface.co/settings/tokens
+#
+# Finalization (dataset.yaml + energon prepare) runs automatically.
+# To skip it, pass --no-finalize.
+#
+# ============================================================================
+# Advanced Usage
+# ============================================================================
+# Custom paths:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-path '/your/input/*.tar' \
+#     --output-dir /your/output/path
+# ============================================================================
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+# Default paths use Docker mounted paths that persist to host:
+#   /workspace/Primus/data/...  →  <your-repo>/data/... on host
+#
+# EASY OVERRIDE: Customize paths with command-line flags:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-path '/your/input/*.tar' \
+#     --output-dir /your/output/path
+#
+# For /data/... paths, set: export DATA_PATH=/data (before starting container)
+# ============================================================================
+
+# ============================================================================
+# Source Configuration (WebDataset)
+# ============================================================================
+source:
+  type: webdataset
+  input_path: '/data/input_dataset/*.tar'  # Glob pattern for WebDataset tars
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_dataset  # Output directory (persisted to host)
+  shard_size: 1000                                    # Number of samples per output shard
+  max_samples: null                                   # Maximum samples to process (null = all)
+  compress: false                                     # Whether to compress output with gzip
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+model:
+  # Base model path (HuggingFace repo or local path)
+  model_path: black-forest-labs/FLUX.1-dev
+
+  # Override individual encoder paths (null means use model_path)
+  vae_path: null    # Custom VAE path
+  t5_path: null     # Custom T5-XXL path
+  clip_path: null   # Custom CLIP-L path
+
+  # Model settings
+  precision: bf16   # Options: bf16, fp16, fp32
+  device: cuda      # Device for encoding
+  batch_size: 8     # Encoding batch size
+
+# ============================================================================
+# Image Preprocessing
+# ============================================================================
+image:
+  image_size: 1024    # Target image size (height and width)
+  center_crop: false  # Whether to center crop before resize
+
+# ============================================================================
+# Input WebDataset Format
+# ============================================================================
+#
+# Your input tar files should contain image-caption pairs:
+#
+# sample_000001.jpg
+# sample_000001.txt
+# sample_000002.jpg
+# sample_000002.txt
+# ...
+#
+# Supported image formats: .jpg, .jpeg, .png, .webp
+# Caption files should be plain text UTF-8 (.txt)
+#
+# ============================================================================
+# Usage
+# ============================================================================
+#
+# Single GPU (uses config paths):
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Multi-GPU (8 GPUs):
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# Override paths:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --input-path '/workspace/Primus/data/input/*.tar' \
+#     --output-dir /workspace/Primus/data/output
+#
+# Override specific values:
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \
+#     --hf-token-file /path/to/.hf_token \
+#     --output-dir /your/custom/path \
+#     --batch-size 16
+#
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
new file mode 100644
index 000000000..7d6b5e33e
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
@@ -0,0 +1,66 @@
+# MLPerf Flux1 streaming ingest configuration
+#
+# Downloads pre-encoded Arrow files from MLCommons R2 and converts them
+# into Energon WebDataset tar shards in a single streaming pass.
+#
+# ============================================================================
+# Usage
+# ============================================================================
+#
+#   primus data diffusion-ingest \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml
+#
+# Override output:
+#   primus data diffusion-ingest \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml \
+#     --output-dir /custom/path --max-files 5
+#
+# ============================================================================
+# Datasets
+# ============================================================================
+# Each entry runs a StreamingIngestPipeline: download Arrow files in parallel,
+# convert to WebDataset tar shards, delete Arrow files after conversion.
+
+datasets:
+  - name: cc12m
+    manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-cc12m-preprocessed.uri
+    split_name: train
+
+  - name: coco
+    manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-coco-preprocessed.uri
+    split_name: val
+
+# ============================================================================
+# Empty Encodings (for CFG dropout during training)
+# ============================================================================
+# Small .npy files downloaded separately (not Arrow).
+
+empty_encodings:
+  manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-empty-encodings.uri
+  output_subdir: empty_encodings
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/mlperf_flux1
+
+# ============================================================================
+# Pipeline Configuration
+# ============================================================================
+
+pipeline:
+  max_workers: 4       # concurrent download threads
+  prefetch_depth: 6    # max Arrow files buffered on disk (~1.2 GB)
+  max_files: null      # null = all files (set to small number for testing)
+
+# ============================================================================
+# Required Training-Time Settings (not consumed by ingest, for reference)
+# ============================================================================
+# When training with this dataset, configure:
+#   vae_latent_mode: resample
+#   vae_scale: 0.3611
+#   vae_shift: 0.1159
+# These are required because forward_step.py unconditionally applies
+# vae_scale * (latents - vae_shift) after resampling from mean/logvar.
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml b/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml
new file mode 100644
index 000000000..d4ca2c275
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml
@@ -0,0 +1,34 @@
+# Quickstart: Pokemon Dataset (256px, works out of the box)
+#
+# Minimal config for getting started quickly with a small public dataset.
+# Uses 256px images for fast processing, even on limited hardware.
+#
+# ============================================================================
+# Usage
+# ============================================================================
+#
+#   primus-cli direct -- data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \
+#     --hf-token-file /path/to/.hf_token
+#
+# The default encoder model (FLUX.1-dev) is gated and requires a
+# HuggingFace token. Get one at https://huggingface.co/settings/tokens
+# ============================================================================
+
+source:
+  type: huggingface
+  hf_dataset: diffusers/pokemon-gpt4-captions
+  hf_split: train
+
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/quickstart_pokemon
+  shard_size: 1000
+
+model:
+  model_path: black-forest-labs/FLUX.1-dev
+  precision: bf16
+  batch_size: 4
+
+image:
+  image_size: 256
+  center_crop: false
diff --git a/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml b/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml
new file mode 100644
index 000000000..7fd77aabb
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml
@@ -0,0 +1,93 @@
+# Text-to-Image-2M 10K Subset Configuration
+#
+# This config uses the 10K high-resolution subset (1024x1024) from the
+# text-to-image-2M dataset, perfect for fine-tuning Flux models.
+#
+# Dataset: https://huggingface.co/datasets/jackyhate/text-to-image-2M
+# Subset: data_1024_10K (exactly 10,000 samples at 1024x1024 resolution)
+#
+# ============================================================================
+# Quick Start (Recommended - Single Command)
+# ============================================================================
+# Prepare and finalize dataset in one command:
+#
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml
+#
+# Finalization (dataset.yaml + energon prepare) runs automatically.
+# To skip it, pass --no-finalize.
+#
+# Processing time: ~30-60 minutes on 8 GPUs
+# Output size: ~50-80 GB (pre-encoded tensors)
+# ============================================================================
+
+# ============================================================================
+# Source Configuration (HuggingFace with data_files)
+# ============================================================================
+source:
+  type: huggingface
+  hf_dataset: jackyhate/text-to-image-2M
+  hf_split: train
+  hf_data_files: data_1024_10K/*.tar  # Load only the 10K subset
+
+# ============================================================================
+# Data Format Configuration
+# ============================================================================
+# This dataset stores images as 'jpg' files and captions in JSON format.
+# Specify how to extract image and caption from dataset samples.
+data_format:
+  image_key: jpg                  # Images are stored in 'jpg' field
+  caption_key: json.prompt        # Captions are in JSON at 'prompt' key (NOT 'caption')
+
+  # Alternative: specify multiple fallback keys (tried in order)
+  # image_keys: ['image', 'jpg']
+  # caption_keys: ['prompt', 'json.prompt', 'caption', 'json.caption']
+
+# ============================================================================
+# Output Configuration
+# ============================================================================
+output:
+  output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/text2img_10k_encoded
+  shard_size: 1000
+  max_samples: null
+  compress: false
+
+# ============================================================================
+# Model Configuration
+# ============================================================================
+model:
+  model_path: black-forest-labs/FLUX.1-dev
+  vae_path: null
+  t5_path: null
+  clip_path: null
+  precision: bf16
+  device: cuda
+  batch_size: 8
+
+# ============================================================================
+# Image Preprocessing
+# ============================================================================
+image:
+  image_size: 1024
+  center_crop: false
+
+# ============================================================================
+# Authentication
+# ============================================================================
+auth:
+  hf_token_file: ""  # Public dataset, no token needed
+
+# ============================================================================
+# Dataset Information
+# ============================================================================
+# Name: Text-to-Image-2M (10K subset)
+# Size: 10,000 high-resolution images
+# Resolution: 1024x1024
+# Content: Complex scenes with rich, detailed annotations
+# License: MIT
+# Use case: Fine-tuning Flux diffusion models
+#
+# Full command example:
+#   primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
+#     --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml
+# ============================================================================
diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml
new file mode 100644
index 000000000..890f6fa1c
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml
@@ -0,0 +1,79 @@
+# Local Dataset Configuration for Pre-encoded Diffusion Data
+#
+# This file should be placed in your pre-encoded dataset directory
+# as dataset.yaml (alongside the WebDataset tar files).
+#
+# Directory structure:
+#   /path/to/dataset_preencoded/
+#   ├── dataset.yaml          <- This file
+#   ├── train-000000.tar
+#   ├── train-000001.tar
+#   ├── ...
+#   └── .nv-meta/             <- Created by `energon prepare`
+#
+# After creating this file, run:
+#   energon prepare /path/to/dataset_preencoded --num-workers 8
+
+# Use CrudeWebdataset to enable cooker-based processing
+__module__: megatron.energon
+__class__: CrudeWebdataset
+
+# Subflavor indicates this dataset contains pre-encoded features
+# This routes samples to the preencoded cooker in the task encoder
+subflavors:
+  encoding: preencoded
+
+# Optional: Filter which tar files to load
+# Uncomment and adjust to use specific shards
+# part_filter:
+#   - train-*.tar      # Load all training shards
+#   - val-*.tar        # Or load validation shards
+
+# Optional: Sample filtering
+# max_samples: 100000  # Limit total samples (useful for testing)
+
+# ============================================================================
+# Expected WebDataset Format
+# ============================================================================
+#
+# Each sample in the tar files should contain:
+#   - __key__: unique identifier (e.g., "sample_000001")
+#   - latent.npy: VAE-encoded image latents [C, H, W] float32
+#   - t5_embed.npy: T5-XXL text embeddings [seq_len, 4096] float32
+#   - clip_embed.npy: CLIP-L pooled embeddings [768] float32
+#   - metadata.json: sample metadata (optional)
+#
+# Example tar content:
+#   sample_000001.latent.npy
+#   sample_000001.t5_embed.npy
+#   sample_000001.clip_embed.npy
+#   sample_000001.metadata.json
+#
+# ============================================================================
+# Preparation Steps
+# ============================================================================
+#
+# 1. Prepare pre-encoded features using the Primus CLI:
+#    tools/docker/primus data diffusion-encoded \
+#      --source-type directory \
+#      --input-dir /path/to/raw_images \
+#      --output-dir /path/to/dataset_preencoded \
+#      --model-path black-forest-labs/FLUX.1-dev \
+#      --batch-size 8
+#
+# 2. Copy this template to the output directory as dataset.yaml
+#
+# 3. Run Energon indexing:
+#    energon prepare /path/to/dataset_preencoded --num-workers 8
+#
+# 4. Verify the dataset:
+#    energon info /path/to/dataset_preencoded
+#
+# 5. Use in training config:
+#    data:
+#      dataset_path: /path/to/dataset_preencoded/dataset.yaml
+#
+# For more information, see:
+#   - primus/configs/data/megatron/diffusion/README.md
+#   - examples/megatron/diffusion/README.md
+#
diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml
new file mode 100644
index 000000000..30cd15051
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml
@@ -0,0 +1,16 @@
+# Dataset Configuration for Pre-encoded NumPy Diffusion Data (MLPerf format)
+#
+# This file is auto-generated by finalize_energon_dataset(encoding="preencoded_numpy")
+# and placed in <dataset>/.nv-meta/dataset.yaml.
+#
+# Data format: numpy-serialized bfloat16 tensors stored as .bytes entries
+# in WebDataset tar shards. Each sample contains:
+#   - t5.bytes:   T5 text embeddings (numpy uint16, viewed as bfloat16)
+#   - clip.bytes: CLIP pooled embeddings (numpy uint16, viewed as bfloat16)
+#   - mean.bytes: VAE posterior mean (numpy uint16, viewed as bfloat16)
+#   - logvar.bytes: VAE posterior log-variance (numpy uint16, viewed as bfloat16)
+
+__module__: megatron.energon
+__class__: CrudeWebdataset
+subflavors:
+  encoding: preencoded_numpy
diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml
new file mode 100644
index 000000000..d24eef25f
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml
@@ -0,0 +1,112 @@
+# Local Dataset Configuration for Raw Diffusion Data (On-the-fly Encoding)
+#
+# This file should be placed in your raw dataset directory
+# as dataset.yaml (alongside the WebDataset tar files).
+#
+# Directory structure:
+#   /path/to/dataset_raw/
+#   ├── dataset.yaml          <- This file
+#   ├── train-000000.tar
+#   ├── train-000001.tar
+#   ├── ...
+#   └── .nv-meta/             <- Created by `energon prepare`
+#
+# After creating this file, run:
+#   energon prepare /path/to/dataset_raw --num-workers 8
+
+# Use CrudeWebdataset to enable cooker-based processing
+__module__: megatron.energon
+__class__: CrudeWebdataset
+
+# Subflavor indicates this dataset contains raw images/text
+# This routes samples to the on-the-fly encoding cooker in the task encoder
+subflavors:
+  encoding: raw
+
+# Optional: Filter which tar files to load
+# Uncomment and adjust to use specific shards
+# part_filter:
+#   - train-*.tar      # Load all training shards
+
+# Optional: Sample filtering
+# max_samples: 10000   # Limit total samples (useful for testing)
+
+# ============================================================================
+# Expected WebDataset Format
+# ============================================================================
+#
+# Each sample in the tar files should contain:
+#   - __key__: unique identifier (e.g., "sample_000001")
+#   - image.jpg: raw image file (JPEG/PNG/WebP)
+#   - caption.txt: text caption describing the image
+#   - metadata.json: sample metadata (optional)
+#
+# Example tar content:
+#   sample_000001.jpg
+#   sample_000001.txt
+#   sample_000001.metadata.json
+#
+# Image requirements:
+#   - Format: JPEG, PNG, or WebP
+#   - Resolution: Any (will be resized to target size during training)
+#   - Color: RGB (grayscale will be converted)
+#
+# Caption requirements:
+#   - Format: Plain text UTF-8
+#   - Length: Any (will be truncated to model max length)
+#   - Content: Descriptive text for the image
+#
+# ============================================================================
+# Preparation Steps
+# ============================================================================
+#
+# 1. Prepare raw dataset using the Primus CLI:
+#    tools/docker/primus data diffusion-raw \
+#      --source-type directory \
+#      --input-dir /path/to/raw_images \
+#      --output-dir /path/to/dataset_raw
+#
+# 2. Copy this template to the output directory as dataset.yaml
+#
+# 3. Run Energon indexing:
+#    energon prepare /path/to/dataset_raw --num-workers 8
+#
+# 4. Verify the dataset:
+#    energon info /path/to/dataset_raw
+#
+# 5. Configure encoders in your training config:
+#    data:
+#      dataset_path: /path/to/dataset_raw/dataset.yaml
+#      encoder_configs:
+#        vae:
+#          model_path: black-forest-labs/FLUX.1-dev
+#          device: cuda
+#          precision: bf16
+#        t5:
+#          model_path: black-forest-labs/FLUX.1-dev
+#          device: cuda
+#          precision: bf16
+#        clip:
+#          model_path: black-forest-labs/FLUX.1-dev
+#          device: cuda
+#          precision: bf16
+#
+# ============================================================================
+# Performance Notes
+# ============================================================================
+#
+# On-the-fly encoding is slower than pre-encoded mode:
+#   - Encoding adds ~100-200ms per image on MI300X
+#   - Requires additional GPU memory for encoders (~10-15GB)
+#   - Recommended for experimentation, not production
+#
+# For production training, use pre-encoded mode:
+#   1. Pre-encode your dataset once using diffusion-encoded
+#   2. Use dataset_preencoded.yaml configuration
+#   3. Enjoy 2-3x faster training
+#
+# For more information, see:
+#   - primus/configs/data/megatron/diffusion/README.md
+#   - examples/megatron/diffusion/README.md
+#   - primus/configs/models/megatron/diffusion/encoders.yaml
+#
diff --git a/primus/configs/data/megatron/diffusion/templates/metadataset.yaml b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml
new file mode 100644
index 000000000..8e41e70b3
--- /dev/null
+++ b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml
@@ -0,0 +1,81 @@
+# Metadataset Configuration for Mixed Training
+#
+# Metadataset combines multiple datasets with different weights.
+# This is useful for:
+#   - Mixing pre-encoded and raw datasets
+#   - Combining multiple data sources
+#   - Creating train/val splits from different datasets
+#
+# This file can be used directly or referenced in training configs:
+#   data:
+#     dataset_path: /path/to/metadataset.yaml
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+# Dataset paths should point to directories on mounted volumes for persistence.
+# Default Docker mount: /workspace/Primus/data/ (persists to host)
+#
+# Adjust paths below to match your dataset locations.
+# All paths should be absolute and point to prepared datasets.
+# ============================================================================
+
+# Use Energon's Metadataset class
+__module__: megatron.energon
+__class__: Metadataset
+
+splits:
+  # Training split with mixed data sources
+  train:
+    datasets:
+      # 80% pre-encoded data (faster training)
+      - weight: 0.8
+        path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_preencoded/
+        # This dataset should have: subflavors: {encoding: preencoded}
+
+      # 20% raw data (flexibility for new data)
+      - weight: 0.2
+        path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_raw/
+        # This dataset should have: subflavors: {encoding: raw}
+
+  # Validation split (typically all pre-encoded for speed)
+  val:
+    datasets:
+      - weight: 1.0
+        path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_preencoded_val/
+        # This dataset should have: subflavors: {encoding: preencoded}
+
+# Global settings applied to all datasets
+shuffle_buffer_size: 1000
+max_samples_per_sequence: 100
+
+# ============================================================================
+# Usage Notes
+# ============================================================================
+#
+# Each dataset path should point to a directory containing:
+#   1. WebDataset tar files (*.tar)
+#   2. Energon index (.nv-meta/ directory)
+#   3. Local dataset.yaml with correct subflavors
+#
+# Example structure (with default Docker mount):
+#   /workspace/Primus/data/diffusion_preencoded/
+#   ├── dataset.yaml          (with subflavors: {encoding: preencoded})
+#   ├── train-*.tar
+#   └── .nv-meta/
+#
+#   /workspace/Primus/data/diffusion_raw/
+#   ├── dataset.yaml          (with subflavors: {encoding: raw})
+#   ├── train-*.tar
+#   └── .nv-meta/
+#
+# Dataset weights:
+#   - Weights are relative, not absolute (0.8 + 0.2 = 1.0)
+#   - Higher weight = more samples from that dataset
+#   - Weights can sum to any positive number
+#
+# For more information, see:
+#   - primus/configs/data/megatron/diffusion/README.md
+#   - examples/megatron/diffusion/README.md
+#   - docs/04-technical-guides/diffusion-models/energon_integration.md
+#
diff --git a/primus/configs/models/diffusion/flux.1_dev_t2i.yaml b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml
new file mode 100644
index 000000000..983a956bb
--- /dev/null
+++ b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml
@@ -0,0 +1,15 @@
+model:
+  name: flux.1-dev
+  config:
+    load_from_pretrained_path: ${PRETRAINED_PATH:}
+    config:
+      trainable_modules: dit
+      guidance: ${FLUX_GUIDANCE:1.0}
+      autoencoder_scale_factor: 0.3611
+      autoencoder_shift_factor: 0.1159
+    # Required only for raw image-text data.
+    encoder:
+      t5_encoder: ${T5_ENCODER:}
+      clip_encoder: ${CLIP_ENCODER:}
+      autoencoder: ${VAE_CHECKPOINT:}
+      max_t5_length: ${MAX_T5_LENGTH:512}
diff --git a/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml b/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml
new file mode 100644
index 000000000..f018bdf90
--- /dev/null
+++ b/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml
@@ -0,0 +1,15 @@
+model:
+  name: flux.1-schnell
+  config:
+    load_from_pretrained_path: ${PRETRAINED_PATH:}
+    config:
+      trainable_modules: dit
+      guidance: ${FLUX_GUIDANCE:1.0}
+      autoencoder_scale_factor: 0.3611
+      autoencoder_shift_factor: 0.1159
+    # Required only for raw image-text data.
+    encoder:
+      t5_encoder: ${T5_ENCODER:}
+      clip_encoder: ${CLIP_ENCODER:}
+      autoencoder: ${VAE_CHECKPOINT:}
+      max_t5_length: ${MAX_T5_LENGTH:256}
diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml
new file mode 100644
index 000000000..023776ecd
--- /dev/null
+++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml
@@ -0,0 +1,13 @@
+model:
+  name: wan
+  config:
+    load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.1-T2V-1.3B}
+    config:
+      model_type: t2v
+      separated_timestep: false
+      fuse_vae_embedding_in_latents: false
+      trainable_modules: dit
+      vae_type: wan_video_vae
+    encoder:
+      t5_encoder: ${TEXT_ENCODER:/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth}
+      autoencoder: ${VAE_CHECKPOINT:/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth}
diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml
new file mode 100644
index 000000000..d9f089b3a
--- /dev/null
+++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml
@@ -0,0 +1,8 @@
+extends:
+  - wan2.1_t2v_1.3b.yaml
+
+# SFT reuses the Wan2.1-T2V-1.3B preset and only swaps the DiT
+# initialization source to INIT_CHECKPOINT (post-train convention).
+model:
+  config:
+    load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.1-T2V-1.3B}
diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml
new file mode 100644
index 000000000..0f23001e2
--- /dev/null
+++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml
@@ -0,0 +1,13 @@
+model:
+  name: wan
+  config:
+    load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-TI2V-5B}
+    config:
+      model_type: ti2v
+      separated_timestep: true
+      fuse_vae_embedding_in_latents: true
+      trainable_modules: dit
+      vae_type: wan_video_vae_38
+    encoder:
+      t5_encoder: ${TEXT_ENCODER:/models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth}
+      autoencoder: ${VAE_CHECKPOINT:/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth}
diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml
new file mode 100644
index 000000000..b8b1cd8d8
--- /dev/null
+++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml
@@ -0,0 +1,8 @@
+extends:
+  - wan2.2_ti2v_5b.yaml
+
+# SFT reuses the Wan2.2-TI2V-5B preset and only swaps the DiT
+# initialization source to INIT_CHECKPOINT (post-train convention).
+model:
+  config:
+    load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.2-TI2V-5B}
diff --git a/primus/configs/models/maxtext/mixtral_8x22B.yaml b/primus/configs/models/maxtext/mixtral_8x22B.yaml
new file mode 100644
index 000000000..cccbceab3
--- /dev/null
+++ b/primus/configs/models/maxtext/mixtral_8x22B.yaml
@@ -0,0 +1,5 @@
+extends:
+  - model_base.yaml
+
+model_name: "mixtral-8x22b"
+tokenizer_path: "mistralai/Mixtral-8x22B-v0.1"
diff --git a/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml b/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml
new file mode 100644
index 000000000..3bb90f351
--- /dev/null
+++ b/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml
@@ -0,0 +1,19 @@
+extends:
+  - qwen2.5_base.yaml
+
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
+
+# DeepSeek-R1-Distill-Qwen-1.5B (Qwen2 architecture) specific parameters.
+# This is the exact model the ODC reference repo (sail-sg/odc) uses in its
+# examples/llm_training default config, so we use it to validate ODC/LB-Mini
+# on Megatron at a size that fits a 64k-token micro-batch without OOM.
+hidden_size: 1536
+ffn_hidden_size: 8960
+num_layers: 28
+num_attention_heads: 12
+num_query_groups: 2
+
+# IMPORTANT: DeepSeek-R1-Distill-Qwen-1.5B uses rope_theta=10000, NOT the
+# Qwen2.5 default of 1000000 inherited from qwen2.5_base.yaml. Override it.
+rotary_base: 10000
diff --git a/primus/configs/models/megatron/deepseek_v4_base.yaml b/primus/configs/models/megatron/deepseek_v4_base.yaml
new file mode 100644
index 000000000..8cd29ea56
--- /dev/null
+++ b/primus/configs/models/megatron/deepseek_v4_base.yaml
@@ -0,0 +1,151 @@
+###############################################################################
+# DeepSeek-V4 base config (defaults shared by V4-Flash / V4-Pro / etc.).
+#
+# Reference:
+#   - deepseek-v4/deepseek-ai/DeepSeek-V4-Flash/config.json
+#   - deepseek-v4/deepseek-ai/DeeSeek-v4-Pro/config.json
+#   - deepseek-v4/develop/techblog/01-deepseek-v4-architecture-deep-dive.md
+###############################################################################
+
+extends:
+  - llama_base.yaml
+
+# Mark this model family as DeepSeek-V4. The Primus trainer
+# (megatron_pretrain_trainer.py) and import_utils.get_model_provider use this
+# string to dispatch to the V4-specific builder.
+model_type: deepseek_v4
+
+# norm
+norm_epsilon: 1.0e-06
+
+# attention base shape
+# V4 uses a single shared latent for K = V (head_dim=512), num_key_value_heads=1.
+# The Megatron MLA path (multi_latent_attention=true) is *NOT* reused here
+# because V4 layers also need CSA / HCA / SWA branches; we will register a
+# new V4 attention spec in Phase 3+. Keep MLA off so we don't accidentally go
+# down the V3 path in Phase 1 / 2 stub builders.
+multi_latent_attention: false
+qk_layernorm: false
+apply_rope_fusion: false   # V4 uses partial RoPE on a 512-dim head; fused path doesn't fit.
+
+# ---------- Hyper-Connections (mHC) ----------
+hc_mult: 4
+hc_sinkhorn_iters: 20
+hc_eps: 1.0e-06
+hc_use_sinkhorn: true
+mtp_use_separate_hc_head: true
+
+# ---------- Hybrid Attention ----------
+hybrid_attention_enabled: true
+compress_ratios: null            # provided by the per-variant yaml (Flash / Pro);
+                                 # entries: 0 = dense+SWA, 4 = CSA, 128 = HCA
+# RoPE base for the CSA / HCA (compressed) branches. 160000 is the V4-Pro
+# value; V4-Flash overrides this to 40000 in deepseek_v4_flash.yaml to match
+# the released DeepSeek-V4-Flash inference/model.py.
+compress_rope_theta: 160000.0
+index_topk: 512
+index_head_dim: 128
+index_n_heads: 64
+attn_sliding_window: 128
+attn_sink: true
+
+# DeepSeek-V4 attention backend selection (unified string selectors). The run
+# script sets these via PRIMUS_USE_V4_ATTENTION_BACKEND /
+# PRIMUS_USE_V4_CSA_ATTENTION_BACKEND (see run_deepseek_v4.sh).
+#   use_v4_attention_backend     (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon
+#   use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0
+# use_turbo_attention (when a core_attention module is built) still takes
+# precedence for the dense path.
+use_v4_attention_backend: triton_v1
+use_v4_csa_attention_backend: triton_v1
+
+# FP8 (E4M3) Indexer QK path (CSA selector). Default off; the indexer QK
+# scoring inputs are fake-quantized to FP8 when enabled (BF16 index-score /
+# top-k path is preserved). Surface via PRIMUS_USE_V4_FP8_INDEXER.
+use_v4_fp8_indexer: false
+
+# Indexer distillation loss (CSA selector training). argTopK is not
+# differentiable, so the lightning indexer only learns through this auxiliary
+# KL against the real attention distribution. 0.0 keeps the loss off AND the
+# indexer frozen, which keeps the params out of the grad buckets so
+# overlap_grad_reduce stays usable.
+#
+# 0.0 is an engineering default (it makes enabling the feature an explicit
+# decision), NOT a claim that the loss is usually unnecessary. The paper keeps
+# it on for the whole sparse training run: sparse attention is introduced after
+# a dense phase, starting with a short stage that warms up the CSA indexer, and
+# the KL alignment continues once the rest of the model is unfrozen -- the main
+# attention keeps moving, so the indexer has to keep tracking it. Turn this on
+# for any training that runs sparse attention; it is only irrelevant for
+# inference. 1e-2 is a starting point.
+#
+# Note this shares the main optimizer and learning rate, whereas the reference
+# trains the indexer with its own optimizer and a much larger LR during warm-up,
+# so the coefficient here is standing in for that LR ratio.
+# Surface via PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF (see run_deepseek_v4.sh).
+v4_indexer_distill_loss_coeff: 0.0
+
+# Plan-5 P29 (RESCOPED): wrap ``sinkhorn_normalize`` (the doubly-stochastic
+# projection inside HyperMixer.compute_weights) with a cached
+# ``torch.compile(fullgraph=True, dynamic=False)`` build. Collapses the
+# 39 fp32 ``aten::sum`` reductions per call into a single Inductor-fused
+# Triton kernel; AOT autograd handles BWD. Default ``false`` until G32
+# (FWD + BWD parity) and G33b (post-P29 trace) flip it on. Driven by
+# the ``PRIMUS_USE_V4_COMPILED_SINKHORN`` env var via run_deepseek_v4.sh.
+use_v4_compiled_sinkhorn: false
+
+# ---------- Output projection (grouped low-rank) ----------
+o_groups: 8
+o_lora_rank: 1024
+
+# ---------- MoE ----------
+moe_layer_freq: 1
+moe_router_topk: 6
+moe_router_score_function: sqrtsoftplus    # V4 default
+# Balancing follows the paper: the aux-loss-free expert bias (noaux_tc) plus a
+# slight sequence-wise balance loss. The learned router implements the latter
+# itself (moe/v4_seq_balance_loss.py) rather than inheriting the framework's,
+# whose aux-loss scores helper rejects sqrtsoftplus. Any other balancing type
+# here leaves the loss off, so the pair below is what turns it on.
+moe_router_load_balancing_type: seq_aux_loss
+moe_router_enable_expert_bias: true        # noaux_tc
+moe_router_bias_update_rate: 1.0e-3
+moe_router_topk_scaling_factor: 1.5        # routed_scaling_factor (Flash 1.5 / Pro 2.5)
+moe_token_dispatcher_type: alltoall
+moe_aux_loss_coeff: 0.001                  # alpha of the sequence-wise balance loss
+moe_router_dtype: null                     # fp32 recommended at training time
+moe_router_pre_softmax: false
+moe_grouped_gemm: true
+moe_use_legacy_grouped_gemm: false
+moe_permute_fusion: true
+moe_shared_expert_overlap: true
+num_hash_layers: 3                         # First 3 MoE layers use hash routing
+
+# ---------- MTP ----------
+mtp_num_layers: 1                          # num_nextn_predict_layers
+mtp_loss_scaling_factor: 0.1
+
+# ---------- RoPE (main path; compressed layers also use compress_rope_theta) ----------
+rotary_base: 10000
+rope_type: yarn
+# YaRN factor, applied only on the compressed (CSA / HCA) branches -- the dense
+# layers use the main base with no YaRN. Paired with compress_rope_theta: both
+# parameterise the same inv_freq, so a variant that overrides one must override
+# the other. 16.0 goes with compress_rope_theta 160000 (V4-Pro);
+# deepseek_v4_flash.yaml overrides both to 4.0 / 40000.
+rotary_scaling_factor: 16.0
+original_max_position_embeddings: 65536
+mscale: 1.0
+mscale_all_dim: 1.0
+rotary_interleaved: true                    # released weights use interleaved RoPE
+
+# ---------- Activation ----------
+swiglu: true
+swiglu_limit: 10.0                         # clamped SwiGLU (FP4/FP8 stability)
+activation_func_clamp_value: 10.0          # clamp swiglu parameter (Megatron config field)
+bias_swiglu_fusion: false                   # fused path lacks clamp; revisit in perf phase
+v4_grouped_experts_support_clamped_swiglu: true  # grouped (TEGroupedMLP) backend honors the clamped SwiGLU above (training sets this too)
+
+# ---------- Parallelism defaults ----------
+expert_model_parallel_size: 1
+expert_tensor_parallel_size: null
diff --git a/primus/configs/models/megatron/deepseek_v4_flash.yaml b/primus/configs/models/megatron/deepseek_v4_flash.yaml
new file mode 100644
index 000000000..6c67cb8d8
--- /dev/null
+++ b/primus/configs/models/megatron/deepseek_v4_flash.yaml
@@ -0,0 +1,73 @@
+###############################################################################
+# DeepSeek-V4 Flash (smaller / dense-attention variant).
+#
+# Source: deepseek-v4/deepseek-ai/DeepSeek-V4-Flash/config.json
+###############################################################################
+
+extends:
+  - deepseek_v4_base.yaml
+
+# core shape
+num_layers: 43
+hidden_size: 4096
+num_attention_heads: 64
+num_query_groups: 1
+kv_channels: 512                  # head_dim
+qk_pos_emb_head_dim: 64           # qk_rope_head_dim
+ffn_hidden_size: 18432            # not used by MoE layers; placeholder for shared/dense FFN
+moe_ffn_hidden_size: 2048         # moe_intermediate_size
+moe_shared_expert_intermediate_size: 2048   # n_shared_experts * moe_intermediate_size
+
+# attention low-rank
+q_lora_rank: 1024
+o_lora_rank: 1024
+o_groups: 8
+
+# MoE
+num_experts: 256
+moe_router_topk: 6
+moe_router_topk_scaling_factor: 1.5
+
+# Indexer / Compressor
+index_topk: 512
+
+# Compressed-branch (CSA / HCA) RoPE. These two must be set together: they
+# parameterise the SAME inv_freq table, where rotary_scaling_factor is the YaRN
+# interpolation factor applied to the low-frequency band of a base-40000 table.
+# Aligning only the base leaves the low-frequency end off by 4x. The upstream
+# open-source Flash recipe ships them as a pair (compressed rotary base 40000 +
+# rotary_scaling_factor 4); the values inherited from deepseek_v4_base.yaml
+# (160000 / 16.0) are the V4-Pro ones.
+#
+# KNOWN GAP, deliberately not resolved here: the released inference/model.py
+# sets original_seq_len = 0, and its precompute_freqs_cis only interpolates when
+# original_seq_len > 0 -- so the *released* model effectively runs with YaRN
+# frequency interpolation OFF (its rope_factor = 40 never takes effect). The
+# pair below therefore follows the upstream training recipe, which does
+# interpolate; it is NOT the released checkpoint's effective configuration.
+# Resolving this needs the official config.json to say which one the weights
+# were trained under. See output/0729_fix/dsv4_fixes_0729.md.
+compress_rope_theta: 40000.0
+rotary_scaling_factor: 4.0
+
+# Per-layer compression schedule (from config.json:compress_ratios)
+# 0   = uncompressed dense layer (full attention with SWA)
+# 4   = CSA branch (overlap mode; per-query top-K from compressed pool via Indexer)
+# 128 = HCA branch (non-overlap mode; full visibility over compressed pool)
+compress_ratios: "[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]"
+
+# context
+seq_length: 4096                  # V4 supports up to 1M; default training seq is overridden by training yaml
+max_position_embeddings: 1048576
+
+# vocab
+vocab_size: 129280
+make_vocab_size_divisible_by: 64
+
+# tokenizer
+tokenizer_type: DeepSeekV4Tokenizer
+tokenizer_model: deepseek-ai/DeepSeek-V4-Flash
+trust_remote_code: true
+
+# init
+init_method_std: 0.02
diff --git a/primus/configs/models/megatron/deepseek_v4_pro.yaml b/primus/configs/models/megatron/deepseek_v4_pro.yaml
new file mode 100644
index 000000000..73dc67c70
--- /dev/null
+++ b/primus/configs/models/megatron/deepseek_v4_pro.yaml
@@ -0,0 +1,46 @@
+###############################################################################
+# DeepSeek-V4 Pro (large MoE variant).
+#
+# Source: deepseek-v4/deepseek-ai/DeeSeek-v4-Pro/config.json
+###############################################################################
+
+extends:
+  - deepseek_v4_base.yaml
+
+num_layers: 61
+hidden_size: 7168
+num_attention_heads: 128
+num_query_groups: 1
+kv_channels: 512
+qk_pos_emb_head_dim: 64
+ffn_hidden_size: 18432
+moe_ffn_hidden_size: 3072
+moe_shared_expert_intermediate_size: 3072
+
+q_lora_rank: 1536
+o_lora_rank: 1024
+o_groups: 16
+
+num_experts: 384
+moe_router_topk: 6
+moe_router_topk_scaling_factor: 2.5
+
+index_topk: 1024
+
+# Per-layer compression schedule:
+#   0   = uncompressed dense layer (full attention with SWA)
+#   4   = CSA branch (overlap mode; per-query top-K from compressed pool via Indexer)
+#   128 = HCA branch (non-overlap mode; full visibility over compressed pool)
+compress_ratios: "[128, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]"
+
+seq_length: 4096
+max_position_embeddings: 1048576
+
+vocab_size: 129280
+make_vocab_size_divisible_by: 64
+
+tokenizer_type: DeepSeekV4Tokenizer
+tokenizer_model: deepseek-ai/DeepSeek-V4-Pro
+trust_remote_code: true
+
+init_method_std: 0.02
diff --git a/primus/configs/models/megatron/diffusion/encoders.yaml b/primus/configs/models/megatron/diffusion/encoders.yaml
new file mode 100644
index 000000000..6d5f8632d
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/encoders.yaml
@@ -0,0 +1,126 @@
+# Encoder Configuration for Flux Diffusion Models
+# This file defines configurations for VAE, T5-XXL, and CLIP-L encoders
+
+# ============================================================================
+# Subfolder Configuration Guide
+# ============================================================================
+# Different model repositories organize files differently:
+#
+# 1. FLUX.1-dev (complex structure):
+#    - VAE: model in vae/
+#    - T5: model in text_encoder_2/, tokenizer in tokenizer_2/
+#    - CLIP: model in text_encoder/, tokenizer in tokenizer/
+#
+# 2. Standalone models (simple structure):
+#    - google/t5-v1_1-xxl: all files in root (subfolder: null)
+#    - openai/clip-vit-large-patch14: all files in root (subfolder: null)
+#
+# Configuration fields:
+#   - subfolder: Where model weights are located
+#   - tokenizer_subfolder: Where tokenizer files are located (text encoders only)
+#     If tokenizer_subfolder is null, uses subfolder value
+# ============================================================================
+
+# VAE Encoder (AutoencoderKL)
+vae:
+  type: autoencoder_kl
+  model_path: black-forest-labs/FLUX.1-dev  # HuggingFace repo or local path
+  subfolder: vae  # Subfolder containing VAE weights
+  precision: bf16  # bf16, fp16, or fp32
+  device: cuda
+  use_cached: true  # Use pre-encoded latents if available
+  freeze_weights: true  # Freeze VAE weights during training
+  cache_dir: null  # Optional: custom cache directory
+
+  # Flux-specific VAE parameters
+  scale_factor: 0.3611  # Latent scaling factor
+  shift_factor: 0.1159  # Latent shift factor
+  in_channels: 3  # RGB input
+  out_channels: 16  # Latent channels
+  latent_downsample_factor: 8  # H/8, W/8
+
+# T5-XXL Text Encoder
+t5:
+  type: t5_xxl
+  model_path: black-forest-labs/FLUX.1-dev
+  subfolder: text_encoder_2         # Subfolder containing T5 model weights
+  tokenizer_path: null              # Uses model_path if null
+  tokenizer_subfolder: tokenizer_2  # Subfolder containing T5 tokenizer files
+  precision: bf16
+  device: cuda
+  use_cached: true
+  freeze_weights: true
+  cache_dir: null  # Optional: custom cache directory
+
+  # T5-XXL parameters
+  max_length: 512  # Maximum sequence length
+  embedding_dim: 4096  # T5-XXL hidden size
+  return_pooled: false  # T5 doesn't use pooled embeddings
+
+# CLIP-L Text Encoder
+clip:
+  type: clip_l
+  model_path: black-forest-labs/FLUX.1-dev
+  subfolder: text_encoder          # Subfolder containing CLIP model weights
+  tokenizer_path: null             # Uses model_path if null
+  tokenizer_subfolder: tokenizer   # Subfolder containing CLIP tokenizer files
+  precision: bf16
+  device: cuda
+  use_cached: true
+  freeze_weights: true
+  cache_dir: null  # Optional: custom cache directory
+
+  # CLIP-L parameters
+  max_length: 77  # CLIP max sequence length
+  embedding_dim: 768  # CLIP-L hidden size
+  pooled_dim: 768  # CLIP-L pooled embedding size
+  return_pooled: true  # CLIP returns both sequence and pooled
+
+# ============================================================================
+# Data Mode Configuration
+# ============================================================================
+# Data mode (preencoded vs on-the-fly) is now determined by dataset subflavors
+# in dataset.yaml, not by a config parameter.
+#
+# For preencoded datasets (RECOMMENDED):
+#   - Create dataset.yaml with: subflavors: {encoding: preencoded}
+#   - Encoders above are NOT loaded during training
+#   - Dataset serves precalculated latents and embeddings
+#
+# For raw/on-the-fly datasets:
+#   - Create dataset.yaml with: subflavors: {encoding: raw}
+#   - Encoders above ARE loaded during training
+#   - Images/text are encoded on-the-fly (slower but saves disk space)
+#
+# See: docs/04-technical-guides/diffusion-models/data_preprocessing.md
+# ============================================================================
+
+# ============================================================================
+# Example: Standalone Models (simpler configuration)
+# ============================================================================
+# For standalone models, subfolder and tokenizer_subfolder should be null:
+#
+# t5_standalone:
+#   type: t5_xxl
+#   model_path: google/t5-v1_1-xxl
+#   subfolder: null                  # All files in root
+#   tokenizer_subfolder: null        # Not needed
+#   precision: bf16
+#   device: cuda
+#
+# clip_standalone:
+#   type: clip_l
+#   model_path: openai/clip-vit-large-patch14
+#   subfolder: null                  # All files in root
+#   tokenizer_subfolder: null        # Not needed
+#   precision: bf16
+#   device: cuda
+# ============================================================================
+
+# Notes:
+# - Model paths can be:
+#   1. HuggingFace repo: "black-forest-labs/FLUX.1-dev"
+#   2. Local path: "/path/to/models/flux-dev"
+#
+# - License note: FLUX.1-dev is non-commercial only
+#   Make sure to comply with license restrictions
diff --git a/primus/configs/models/megatron/diffusion/flux_12b.yaml b/primus/configs/models/megatron/diffusion/flux_12b.yaml
new file mode 100644
index 000000000..0a3dcf556
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_12b.yaml
@@ -0,0 +1,141 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Configuration (Production/Standard Model)
+#
+# This is the full Flux configuration matching FLUX.1 [dev] and FLUX.1 [schnell]
+# from Black Forest Labs (black-forest-labs/FLUX.1-dev on Hugging Face).
+#
+# Use cases:
+#   - Production image generation
+#   - Large-scale training
+#   - Fine-tuning from pretrained weights
+#   - Research and benchmarking
+#   - High-quality image synthesis
+
+extends:
+  - flux_base.yaml
+
+# ==============================================================================
+# Architecture: Layer Counts (ONLY difference from base)
+# ==============================================================================
+
+num_joint_layers: 19       # Multimodal (joint) transformer blocks
+num_single_layers: 38      # Image-only transformer blocks
+num_layers: 57             # Total layers for Megatron compatibility (joint + single)
+
+# Total layers: 19 + 38 = 57 transformer blocks
+# Total parameters: ~12 billion
+# Model matches: FLUX.1 [dev] and FLUX.1 [schnell] architectures
+
+# ==============================================================================
+# Guidance Configuration (Optional Override)
+# ==============================================================================
+
+# guidance_embed: false    # Default from flux_base.yaml
+
+# To train FLUX.1 [dev] style (with CFG support):
+#   guidance_embed: true
+#
+# To train FLUX.1 [schnell] style (distilled, no CFG):
+#   guidance_embed: false  (default)
+
+# ==============================================================================
+# All Other Parameters Inherited from flux_base.yaml
+# ==============================================================================
+
+# The following are inherited and should NOT be redefined here:
+#   - hidden_size: 3072
+#   - num_attention_heads: 24
+#   - context_dim: 4096 (T5-XXL)
+#   - vec_in_dim: 768 (CLIP-L)
+#   - in_channels: 64 (VAE latents)
+#   - activation_func: openai_gelu (from FluxConfig default)
+#   - add_qkv_bias: true (from FluxConfig default)
+#   - All other architectural parameters
+
+# ==============================================================================
+# Hardware Requirements
+# ==============================================================================
+
+# Training:
+#   - Minimum: 8x H100 80GB or 8x MI300X 192GB
+#   - Recommended: 8x H100 80GB with NVLink
+#   - Batch size: Start with global_batch_size=256, micro_batch_size=1
+#   - Memory: ~70GB per GPU with micro_batch=1, TP=1
+
+# Inference:
+#   - Minimum: 1x H100 80GB or 1x MI300X 192GB
+#   - Typical: 1x A100 80GB works for batch_size=1
+#   - Memory: ~24GB for model weights + activations
+
+# ==============================================================================
+# Training Recommendations
+# ==============================================================================
+
+# Parallelism Strategy:
+#   - Data Parallel (DP): 8 (default, simplest)
+#   - Tensor Parallel (TP): 1-2 (use if memory constrained)
+#   - Pipeline Parallel (PP): 1 (Flux doesn't benefit from PP)
+#   - Sequence Parallel (SP): false (not needed for typical resolutions)
+
+# Training Duration:
+#   - From scratch: 100K-1M steps (weeks on 8x H100)
+#   - Fine-tuning: 10K-50K steps (hours to days)
+#   - LoRA adaptation: 1K-10K steps (hours)
+
+# Suggested Training Config:
+#   micro_batch_size: 1
+#   global_batch_size: 256
+#   learning_rate: 1.0e-4
+#   weight_decay: 0.01
+#   warmup_steps: 1000
+#   total_steps: 100000
+
+# ==============================================================================
+# Comparison with 535M Variant
+# ==============================================================================
+
+# Flux 535M (testing):
+#   - Layers: 1 joint + 1 single = 2 total
+#   - Parameters: ~535M
+#   - Training: Single GPU, minutes to hours
+#   - Use: Testing, debugging, quick iteration
+
+# Flux 12B (this config):
+#   - Layers: 19 joint + 38 single = 57 total
+#   - Parameters: ~12B
+#   - Training: Multi-GPU, days to weeks
+#   - Use: Production, research, deployment
+
+# ==============================================================================
+# NeMo Compatibility
+# ==============================================================================
+
+# This config matches NeMo's FluxConfig defaults:
+#   nemo/collections/diffusion/models/flux/model.py:74-114
+#
+#   @dataclass
+#   class FluxConfig(TransformerConfig, io.IOMixin):
+#       num_joint_layers: int = 19
+#       num_single_layers: int = 38
+#       hidden_size: int = 3072
+#       num_attention_heads: int = 24
+#       activation_func: Callable = openai_gelu
+#       add_qkv_bias: bool = True
+
+# Checkpoints are fully interoperable between Primus and NeMo
+
+# ==============================================================================
+# Pretrained Weights
+# ==============================================================================
+
+# Compatible with weights from:
+#   - black-forest-labs/FLUX.1-dev (HuggingFace)
+#   - black-forest-labs/FLUX.1-schnell (HuggingFace)
+#
+# Note: FLUX.1-dev is non-commercial license
+#       FLUX.1-schnell is Apache 2.0 license
+#
+# To load pretrained weights, use checkpoint conversion utilities in:
+#   primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py
diff --git a/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml b/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml
new file mode 100644
index 000000000..3fd6c19d3
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml
@@ -0,0 +1,122 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B with FP8 Configuration for AMD MI300X
+#
+# This configuration enables FP8 (8-bit floating point) training for Flux 12B,
+# providing ~2x memory reduction and 1.5-2x training speedup while maintaining
+# numerical stability through delayed scaling.
+#
+# Target Hardware:
+#   - AMD MI300X GPUs with ROCm 6.0+
+#   - Requires Transformer Engine 2.1.0+ with ROCm backend
+#
+# Benefits:
+#   - ~2x memory reduction (activations and weights)
+#   - 1.5-2x training speedup on MI300X
+#   - Maintains numerical stability via delayed scaling
+#   - Allows larger batch sizes or higher resolutions
+#
+# Use cases:
+#   - Memory-constrained training scenarios
+#   - Faster iteration during experimentation
+#   - Training with limited GPU resources
+#   - Production training with improved efficiency
+
+extends:
+  - flux_12b.yaml
+
+# ==============================================================================
+# FP8 Configuration
+# ==============================================================================
+
+# Enable FP8 precision for training
+fp8: "e4m3"  # Use E4M3 format for all FP8 tensors
+             # Alternative: "hybrid" (E4M3 for activations/weights, E5M2 for grad outputs)
+
+# FP8 Recipe: Controls scaling strategy
+fp8_recipe: "delayed"  # Delayed scaling for numerical stability (recommended)
+                       # Alternatives: "tensorwise", "blockwise", "mxfp8"
+                       # Note: delayed scaling is most stable for training
+
+# Scaling Factor Configuration
+fp8_margin: 0  # Margin for scaling factor computation (0 = no margin)
+fp8_amax_history_len: 1024  # History window length for delayed scaling
+                            # Larger = more stable, smaller = adapts faster
+fp8_amax_compute_algo: "most_recent"  # Algorithm for amax computation
+                                      # Options: "most_recent", "max"
+
+# FP8 Gradient Configuration
+fp8_wgrad: true  # Enable FP8 for weight gradients (recommended for memory)
+                 # Set to false if numerical issues occur
+
+# Attention Precision (keep in higher precision for stability)
+fp8_dot_product_attention: false  # Keep dot product attention in BF16/FP32
+fp8_multi_head_attention: false   # Keep multi-head attention in BF16/FP32
+
+# ==============================================================================
+# Optimization Adjustments for FP8
+# ==============================================================================
+
+# Gradient clipping is still important with FP8
+# May need adjustment if training becomes unstable
+clip_grad: 1.0
+
+# ==============================================================================
+# Hardware Requirements (FP8)
+# ==============================================================================
+
+# Training with FP8:
+#   - Minimum: 4x MI300X 192GB (vs 8x without FP8)
+#   - Recommended: 8x MI300X 192GB with RDMA
+#   - Batch size: Can increase by ~2x vs BF16
+#   - Memory: ~35-40GB per GPU (vs ~70GB BF16)
+
+# Inference with FP8:
+#   - Minimum: 1x MI300X 192GB
+#   - Memory: ~12-15GB for model weights + activations (vs ~24GB BF16)
+
+# ==============================================================================
+# Training Recommendations for FP8
+# ==============================================================================
+
+# Suggested Training Config (adjusted for FP8):
+#   micro_batch_size: 2 (vs 1 for BF16)
+#   global_batch_size: 256 (same as BF16)
+#   learning_rate: 1.0e-4 (same as BF16)
+#   weight_decay: 0.01 (same as BF16)
+#   warmup_steps: 1000 (same as BF16)
+#   total_steps: 100000 (same as BF16)
+
+# Parallelism Strategy (FP8):
+#   - Data Parallel (DP): 8
+#   - Tensor Parallel (TP): 1 (less memory pressure with FP8)
+#   - Pipeline Parallel (PP): 1
+#   - Sequence Parallel (SP): false
+
+# ==============================================================================
+# Troubleshooting
+# ==============================================================================
+
+# If training becomes unstable (NaN/Inf losses):
+#   1. Increase fp8_amax_history_len to 2048 or 4096
+#   2. Set fp8_wgrad: false (use higher precision for weight gradients)
+#   3. Increase fp8_margin to 1 or 2
+#   4. Try fp8_amax_compute_algo: "max" instead of "most_recent"
+#   5. Consider keeping first/last layers in BF16 (set first_last_layers_bf16: true)
+
+# If insufficient memory even with FP8:
+#   - Enable gradient checkpointing (recompute_granularity: "selective")
+#   - Increase tensor parallelism (tensor_model_parallel_size: 2)
+#   - Reduce batch size further
+
+# ==============================================================================
+# Verification
+# ==============================================================================
+
+# Before production use, verify:
+#   - Transformer Engine FP8 support: Check with setup_fp8_rocm.sh
+#   - Forward pass stability: Run test_flux_precision.py::test_flux_fp8_forward
+#   - Training convergence: Compare loss curves with BF16 baseline
+#   - Memory usage: Should be ~50% of BF16 training
+#   - Training speed: Should be 1.5-2x faster than BF16
diff --git a/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml b/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml
new file mode 100644
index 000000000..e25d902b2
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml
@@ -0,0 +1,50 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B with RoPE Fusion Optimization
+#
+# This configuration enables Transformer Engine's fused RoPE kernels for improved performance.
+#
+# ==============================================================================
+# CRITICAL REQUIREMENT: Same-Resolution Batches
+# ==============================================================================
+#
+# RoPE fusion requires ALL images in each training batch to have the SAME
+# resolution (height and width). This constraint exists because the fused kernel
+# expects position IDs with shape [1, H*W/4, 3] that broadcast across the batch.
+#
+# Variable-resolution batches will produce INCORRECT positional encodings and
+# degrade model quality. You will receive a warning at model initialization.
+#
+# Recommended data pipeline configurations:
+#   1. Fixed resolution: All images resized to same size (e.g., 512x512, 1024x1024)
+#   2. Resolution bucketing: Group images by resolution in separate batches
+#   3. Aspect ratio bucketing: Use bucketing with consistent dimensions per batch
+#
+# Technical details:
+# - Follows NVIDIA's MLPerf Flux implementation strategy (hsg_ngpu1152 config)
+# - Uses batch_size=1 for position ID generation (broadcasting handles actual batch)
+# - Compatible with both AMD ROCm and NVIDIA CUDA Transformer Engine forks
+# - Fused kernel requires freqs tensor shape: [S, 1, 1, D]
+#
+# ==============================================================================
+
+extends:
+  - flux_12b.yaml
+
+# ==============================================================================
+# RoPE Fusion Optimization - ENABLED
+# ==============================================================================
+
+apply_rope_fusion: true
+rotary_interleaved: true
+position_embedding_type: rope  # Prevents Megatron validate_args from overriding apply_rope_fusion
+
+# ==============================================================================
+# Transformer Implementation
+# ==============================================================================
+# Options:
+#  - transformer_impl: "transformer_engine" (default) - Use TransformerEngine
+#  - transformer_impl: "local" - Use native Megatron + Primus Turbo (NO TransformerEngine)
+# ==============================================================================
+# Override in training config if needed (default is "transformer_engine")
diff --git a/primus/configs/models/megatron/diffusion/flux_535m.yaml b/primus/configs/models/megatron/diffusion/flux_535m.yaml
new file mode 100644
index 000000000..28293cb65
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_535m.yaml
@@ -0,0 +1,85 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M Configuration (Minimal Testing Variant)
+#
+# This is a scaled-down Flux configuration with only 1 joint + 1 single layer
+# for rapid prototyping, testing, and validation.
+#
+# Use cases:
+#   - Unit testing and CI/CD
+#   - Quick iteration during development
+#   - Validating training pipeline
+#   - Comparing against NeMo baseline
+#   - Single GPU training/testing
+
+extends:
+  - flux_base.yaml
+
+# ==============================================================================
+# Architecture: Layer Counts (ONLY difference from base)
+# ==============================================================================
+
+num_joint_layers: 1        # Multimodal (joint) transformer blocks
+num_single_layers: 1       # Image-only transformer blocks
+num_layers: 2              # Total layers for Megatron compatibility (joint + single)
+
+# Total layers: 1 + 1 = 2 (vs. 57 for Flux 12B)
+# Total parameters: ~535 million (vs. ~12 billion for Flux 12B)
+
+# ==============================================================================
+# All Other Parameters Inherited from flux_base.yaml
+# ==============================================================================
+
+# The following are inherited and should NOT be redefined here:
+#   - hidden_size: 3072
+#   - num_attention_heads: 24
+#   - context_dim: 4096
+#   - vec_in_dim: 768
+#   - in_channels: 64
+#   - activation_func: openai_gelu (from FluxConfig default)
+#   - add_qkv_bias: true (from FluxConfig default)
+#   - All other architectural parameters
+
+# ==============================================================================
+# Training Recommendations
+# ==============================================================================
+
+# Suggested training config for 535M:
+#   - micro_batch_size: 1-2
+#   - global_batch_size: 8-16
+#   - GPUs: 1-2 (fits on single H100 80GB)
+#   - Training time: ~10-100x faster than 12B
+#   - Use for: sanity checks, debugging, quick experiments
+
+# ==============================================================================
+# Comparison with Full Flux
+# ==============================================================================
+
+# Flux 535M (this config):
+#   - Layers: 1 joint + 1 single = 2 total
+#   - Parameters: ~535M
+#   - Training: Single GPU capable
+#   - Purpose: Testing/development
+
+# Flux 12B (production):
+#   - Layers: 19 joint + 38 single = 57 total
+#   - Parameters: ~12B
+#   - Training: Multi-GPU required
+#   - Purpose: Production deployment
+
+# ==============================================================================
+# NeMo Compatibility
+# ==============================================================================
+
+# This config matches NeMo's unit_test() recipe:
+#   nemo/scripts/flux/flux_training.py:252-278
+#
+#   recipe.model.flux_params.flux_config = run.Config(
+#       FluxConfig,
+#       num_joint_layers=1,
+#       num_single_layers=1,
+#       ...
+#   )
+
+# Checkpoints are interoperable between Primus and NeMo
diff --git a/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml b/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml
new file mode 100644
index 000000000..ff31f5150
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml
@@ -0,0 +1,107 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 535M with FP8 Configuration (Testing/Development)
+#
+# This is a minimal Flux variant with FP8 enabled, designed for:
+#   - Quick validation of FP8 functionality
+#   - Testing FP8 integration before scaling to 12B
+#   - Debugging FP8-related issues
+#   - Development and experimentation
+#
+# Target Hardware:
+#   - Single AMD MI300X GPU with ROCm 6.0+
+#   - Requires Transformer Engine 2.1.0+ with ROCm backend
+#
+# Use cases:
+#   - Validating FP8 setup and configuration
+#   - Quick iteration during FP8 feature development
+#   - Testing FP8 numerical stability
+#   - Verifying Transformer Engine compatibility
+
+extends:
+  - flux_535m.yaml
+
+# ==============================================================================
+# FP8 Configuration (Same as flux_12b_fp8.yaml)
+# ==============================================================================
+
+# Enable FP8 precision for training
+fp8: "e4m3"  # Use E4M3 format for all FP8 tensors
+
+# FP8 Recipe: Controls scaling strategy
+fp8_recipe: "delayed"  # Delayed scaling for numerical stability (recommended)
+
+# Scaling Factor Configuration
+fp8_margin: 0  # Margin for scaling factor computation
+fp8_amax_history_len: 1024  # History window length for delayed scaling
+fp8_amax_compute_algo: "most_recent"  # Algorithm for amax computation
+
+# FP8 Gradient Configuration
+fp8_wgrad: true  # Enable FP8 for weight gradients
+
+# Attention Precision (keep in higher precision for stability)
+fp8_dot_product_attention: false  # Keep dot product attention in BF16/FP32
+fp8_multi_head_attention: false   # Keep multi-head attention in BF16/FP32
+
+# ==============================================================================
+# Optimization Adjustments for FP8
+# ==============================================================================
+
+# Gradient clipping
+clip_grad: 1.0
+
+# ==============================================================================
+# Hardware Requirements (FP8)
+# ==============================================================================
+
+# Training with FP8:
+#   - Minimum: 1x MI300X 192GB
+#   - Memory: ~5-7GB per GPU (vs ~10-15GB BF16)
+#   - Training time: Minutes (vs hours for 12B)
+
+# Inference with FP8:
+#   - Minimum: 1x MI300X 192GB
+#   - Memory: ~2-3GB for model weights + activations
+
+# ==============================================================================
+# Training Recommendations for FP8
+# ==============================================================================
+
+# Suggested Training Config (testing):
+#   micro_batch_size: 4 (vs 2 for BF16)
+#   global_batch_size: 32
+#   learning_rate: 1.0e-4
+#   weight_decay: 0.01
+#   warmup_steps: 100
+#   total_steps: 1000 (for quick validation)
+
+# Parallelism Strategy:
+#   - Data Parallel (DP): 1 (single GPU testing)
+#   - Tensor Parallel (TP): 1 (not needed for 535M)
+#   - Pipeline Parallel (PP): 1
+#   - Sequence Parallel (SP): false
+
+# ==============================================================================
+# Testing Strategy
+# ==============================================================================
+
+# Use this config to verify:
+#   1. FP8 forward pass works without NaN/Inf
+#   2. FP8 backward pass computes gradients correctly
+#   3. Training loop runs stable for multiple iterations
+#   4. Memory usage is reduced as expected
+#   5. Training speed improvement is measurable
+
+# Run tests with:
+#   pytest tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py -v
+
+# ==============================================================================
+# Troubleshooting
+# ==============================================================================
+
+# Same troubleshooting steps as flux_12b_fp8.yaml:
+#   - Increase fp8_amax_history_len if unstable
+#   - Set fp8_wgrad: false if gradient issues
+#   - Adjust fp8_margin if scaling problems
+#   - Use fp8_amax_compute_algo: "max" for more conservative scaling
diff --git a/primus/configs/models/megatron/diffusion/flux_base.yaml b/primus/configs/models/megatron/diffusion/flux_base.yaml
new file mode 100644
index 000000000..5f76aae7a
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion/flux_base.yaml
@@ -0,0 +1,110 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux Base Configuration
+# Common parameters for all Flux model variants (535M, 12B, custom)
+
+extends:
+  - ../diffusion_model.yaml
+
+# ==============================================================================
+# Core Architecture (Flux Standard)
+# ==============================================================================
+
+# Layer configuration (Flux uses joint + single layers, not standard num_layers)
+# num_layers will be computed in code as num_joint_layers + num_single_layers
+num_joint_layers: null     # To be set by specific variant (535M, 12B, etc.)
+num_single_layers: null    # To be set by specific variant (535M, 12B, etc.)
+
+# Dimensions (same across all Flux variants)
+hidden_size: 3072          # Transformer hidden dimension
+num_attention_heads: 24    # Number of attention heads
+ffn_hidden_size: null      # Auto-computed as hidden_size * 4 = 12288
+
+# Head dimension (computed): hidden_size / num_attention_heads = 128
+
+# ==============================================================================
+# Context Dimensions (Text Encoder Integration)
+# ==============================================================================
+
+context_dim: 4096          # T5-XXL hidden dimension (text embeddings)
+vec_in_dim: 768            # CLIP-L pooled dimension (pooled text)
+model_channels: 256        # Channels for timestep embedding MLPs
+
+# ==============================================================================
+# Input/Output Channels
+# ==============================================================================
+
+in_channels: 64            # VAE latent channels (16 * 4 from 2x2 packing)
+out_channels: 64           # Output channels (same as input for latent space)
+
+# ==============================================================================
+# Patchification
+# ==============================================================================
+
+patch_size: 1              # Flux uses 1x1 patches (no spatial patchification)
+
+# ==============================================================================
+# Position Embeddings (3D RoPE)
+# ==============================================================================
+
+theta: 10000               # Base frequency for RoPE
+axes_dim: [16, 56, 56]     # 3D RoPE dimensions: [channels, height, width]
+                           # Sum = 128 = head_dim
+
+# ==============================================================================
+# Transformer Configuration
+# ==============================================================================
+
+# LayerNorm
+layernorm_epsilon: 1.0e-6  # Epsilon for numerical stability
+
+# Residual connections
+apply_residual_connection_post_layernorm: false  # Flux uses pre-norm
+
+# Dropout (Flux typically uses none)
+hidden_dropout: 0.0
+attention_dropout: 0.0
+
+# ==============================================================================
+# Guidance (Classifier-Free Guidance)
+# ==============================================================================
+
+guidance_embed: false      # Set true to enable CFG embedding
+                           # FLUX.1 [dev]: true (supports CFG)
+                           # FLUX.1 [schnell]: false (distilled, no CFG)
+
+# ==============================================================================
+# CRITICAL: Activation Function & QKV Configuration
+# ==============================================================================
+
+# Activation function: defaults to openai_gelu_no_jit (non-JIT tanh GELU, ROCm-safe;
+# see FluxConfig.activation_func). Override via activation_func in your config:
+#   activation_func: openai_gelu   # fused F.gelu(approximate="tanh")
+#   activation_func: erf_gelu      # erf-based GELU
+# Changing the activation changes numerics: keep it consistent with the activation
+# used to train any checkpoint you load.
+
+# FluxConfig defaults that we intentionally inherit:
+#   activation_func: openai_gelu_no_jit (FluxConfig default, ROCm-safe)
+#   add_qkv_bias: true
+#   single_block_bias: true (new parameter)
+#   rotary_interleaved: true
+#   apply_rope_fusion: false
+
+# New parameter: single_block_bias
+#   Controls bias in single block linear layers independently from joint blocks
+#   Default: true (bias enabled in single blocks)
+#   Set to false to disable bias in single blocks only
+
+# These defaults are compatible with Primus training conventions
+
+# ==============================================================================
+# Notes
+# ==============================================================================
+
+# - This config defines Flux architecture shared across all variants
+# - Specific variants (535M, 12B) only override num_joint_layers, num_single_layers
+# - Activation function (openai_gelu) comes from FluxConfig class default
+# - All Flux variants share same hidden_size (3072) and num_heads (24)
+# - For custom variants, extend this file and override layer counts
diff --git a/primus/configs/models/megatron/diffusion_model.yaml b/primus/configs/models/megatron/diffusion_model.yaml
new file mode 100644
index 000000000..6af2c1583
--- /dev/null
+++ b/primus/configs/models/megatron/diffusion_model.yaml
@@ -0,0 +1,135 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Base configuration for diffusion models
+# This file defines ARCHITECTURE ONLY - training params go in trainer config
+
+extends:
+  - primus_megatron_model.yaml
+
+# ==============================================================================
+# Model Parallelism
+# ==============================================================================
+
+# Parallelism configuration
+model_parallel_size: null
+tensor_model_parallel_size: 1
+encoder_tensor_model_parallel_size: 0
+context_parallel_size: 1
+cp_comm_type: p2p  # p2p, a2a, allgather or a2a+p2p
+hierarchical_context_parallel_sizes: null
+
+pipeline_model_parallel_size: 1
+pipeline_model_parallel_layout: null
+pipeline_model_parallel_comm_backend: null  # str: nccl, ucc
+encoder_pipeline_model_parallel_size: 0
+pipeline_model_parallel_split_rank: null
+decoder_first_pipeline_num_layers: null  # int
+decoder_last_pipeline_num_layers: null  # int
+num_layers_per_virtual_pipeline_stage: null  # int
+num_virtual_stages_per_pipeline_rank: null  # int
+microbatch_group_size_per_vp_stage: null  # int
+sequence_parallel: true
+expert_model_parallel_size: 1
+expert_tensor_parallel_size: null  # int
+
+# ==============================================================================
+# Architecture Parameters
+# ==============================================================================
+
+# Model tying (required for FSDP2)
+untie_embeddings_and_output_weights: true
+
+# Layer configuration (will be overridden by specific model configs)
+num_layers: null  # Must be set by specific diffusion model configs
+encoder_num_layers: null
+decoder_num_layers: null
+
+# Position embeddings (diffusion models may use different positional encoding)
+max_position_embeddings: 8192  # Maximum sequence length
+
+# Tokenizer (diffusion models don't use tokenizers - text is pre-encoded)
+tokenizer_type: NullTokenizer  # Placeholder tokenizer for diffusion models
+vocab_size: 1  # Placeholder vocab size
+
+# Initialization (architecture-related)
+init_method_std: 0.02      # Standard deviation for weight initialization
+
+# MoE settings (diffusion models typically don't use MoE, but fields needed for compatibility)
+num_experts: null
+moe_layer_freq: 1  # int
+moe_ffn_hidden_size: null  # int
+moe_shared_expert_overlap: false
+moe_shared_expert_intermediate_size: null  # int
+moe_grouped_gemm: false
+moe_router_load_balancing_type: "aux_loss"
+moe_router_dtype: null
+moe_router_score_function: softmax
+moe_router_topk: 2
+moe_router_pre_softmax: false
+moe_router_num_groups: null
+moe_router_group_topk: null
+moe_router_topk_scaling_factor: null
+moe_router_enable_expert_bias: false
+moe_router_bias_update_rate: 1.0e-03
+moe_use_legacy_grouped_gemm: false
+moe_aux_loss_coeff: 0.0
+moe_z_loss_coeff: null
+moe_input_jitter_eps: null
+moe_token_dispatcher_type: allgather
+moe_enable_deepep: false
+moe_per_layer_logging: false
+moe_expert_capacity_factor: null
+moe_pad_expert_input_to_capacity: false
+moe_token_drop_policy: probs
+moe_layer_recompute: false
+moe_extended_tp: false
+moe_use_upcycling: false
+moe_permute_fusion: false
+disable_primus_topk_router: false
+moe_router_force_load_balancing: false
+use_deprecated_20241209_moe_layer: false
+delay_wgrad_compute: false
+
+# Tensor parallelism communication overlap settings
+tp_comm_overlap: false
+tp_comm_overlap_cfg: null
+tp_comm_overlap_ag: true
+tp_comm_overlap_rs: true
+tp_comm_overlap_rs_dgrad: false
+tp_comm_bulk_wgrad: true
+tp_comm_bulk_dgrad: true
+
+# Optimization flags
+gradient_accumulation_fusion: true  # Should be disabled for FSDP2
+
+# Other compatibility fields
+fused_padded_mla_attention: false
+
+# ==============================================================================
+# Important: Tokenizer Not Applicable
+# ==============================================================================
+
+# Diffusion models process image latents, not text tokens
+# DO NOT define: tokenizer_type, tokenizer_model, vocab_size, etc.
+# These parameters are only for language models
+
+# ==============================================================================
+# Training Parameters Removed
+# ==============================================================================
+
+# The following were REMOVED and belong in trainer config or overrides:
+# - seq_length, micro_batch_size, global_batch_size (trainer config)
+# - learning_rate, min_lr, weight_decay, clip_grad (trainer config)
+# - scheduler_type (trainer config)
+# - bf16, fp16 (trainer config)
+#
+# Override these in your training config's overrides section as needed.
+
+# ==============================================================================
+# Notes
+# ==============================================================================
+
+# - This config provides architecture defaults for all diffusion models
+# - Specific architectures (Flux, DiT) extend this and add their parameters
+# - Training parameters should be set in pre_trainer.yaml or overrides
diff --git a/primus/configs/models/megatron/mamba_base.yaml b/primus/configs/models/megatron/mamba_base.yaml
index d52fe6db2..6f7fa2a72 100644
--- a/primus/configs/models/megatron/mamba_base.yaml
+++ b/primus/configs/models/megatron/mamba_base.yaml
@@ -1,4 +1,4 @@
-bases:
+extends:
   - language_model.yaml
 
 # Mamba-specific configuration
@@ -34,3 +34,23 @@ norm_epsilon: 1.0e-5
 
 # Initialization
 init_method_std: 0.02
+
+# -----------------------------------------------------------------------------
+# FLA runtime knobs (consumed by primus.backends.megatron.patches.fla_runtime
+# patch which re-exports them as PRIMUS_FLA_* / PRIMUS_FUSED_CE* env vars).
+#
+# All are `null` here so the patch is a no-op for models that do not set any
+# of them; per-model YAMLs (or experiment `overrides:` blocks) can override
+# individual fields without having to re-declare the whole set.  See
+# primus/backends/megatron/patches/fla_runtime_patches.py for the full
+# precedence rules (launcher env var > YAML > documented default).
+# -----------------------------------------------------------------------------
+use_fla_fused_swiglu: null        # bool   — PRIMUS_FLA_SWIGLU      (default 1)
+use_fla_fused_rmsnorm: null       # bool   — PRIMUS_FLA_NORM        (default 0)
+use_fla_fused_gated_norm: null    # bool   — also PRIMUS_FLA_NORM   (default 0)
+use_fla_short_conv: null          # bool   — PRIMUS_FLA_CONV        (default 0)
+use_fla_data: null                # bool   — PRIMUS_FLA_DATA        (default 0)
+fla_cache_dir: null               # path   — PRIMUS_FLA_CACHE_DIR   (default "")
+fused_ce_mode: null               # 0/1/2  — PRIMUS_FUSED_CE        (default 1)
+fused_ce_chunks: null             # int    — PRIMUS_FUSED_CE_CHUNKS (default 32)
+fla_mla_attn: null                # str    — PRIMUS_FLA_MLA_ATTN    (default unset)
diff --git a/primus/configs/models/megatron/zebra_llama_1B_gdn.yaml b/primus/configs/models/megatron/zebra_llama_1B_gdn.yaml
new file mode 100644
index 000000000..d9f9edf99
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_1B_gdn.yaml
@@ -0,0 +1,40 @@
+extends:
+  - mamba_base.yaml
+
+# Zebra Llama 1B Pure GDN configuration
+model_type: mamba  # CRITICAL: Hybrid models must use mamba model type
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model size parameters
+num_layers: 32
+hidden_size: 2048
+ffn_hidden_size: 8192
+
+# Pure GDN — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# Gated Delta Net parameters
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 64
+linear_value_head_dim: 64
+linear_num_key_heads: 8
+linear_num_value_heads: 16
+
+# No MLA needed for pure GDN
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: false
+num_attention_heads: 32
+
+# Positional encoding
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_1B_gdn_pure.yaml b/primus/configs/models/megatron/zebra_llama_1B_gdn_pure.yaml
new file mode 100644
index 000000000..b50dc562e
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_1B_gdn_pure.yaml
@@ -0,0 +1,59 @@
+extends:
+  - mamba_base.yaml
+
+# Pure Gated DeltaNet 1B — matches FLA gated_deltanet_1B_pure.json exactly
+#   ~1.20B params (tied embeddings)
+#   16 GDN blocks + 16 MLP blocks = 32 Megatron sublayers
+model_type: mamba
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model dimensions
+num_layers: 32
+hidden_size: 2048
+ffn_hidden_size: 8192
+
+# Pure GDN — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+
+# Mamba params kept for base config compatibility (unused in pure GDN)
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# GDN parameters — matched to FLA config:
+#   num_heads=8, num_v_heads=16, head_dim=64, expand_v=1.0, conv_size=4
+#   key_dim  = 8  heads x 64 dim = 512
+#   value_dim = 16 heads x 64 dim = 1024
+#   Grouped Value Attention: each Q/K head serves 2 V heads
+use_short_conv: true
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 64
+linear_value_head_dim: 64
+linear_num_key_heads: 8
+linear_num_value_heads: 16
+
+# No MLA needed for pure GDN
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: false
+num_attention_heads: 8
+
+# Embedding tying (matches FLA tie_word_embeddings=true)
+untie_embeddings_and_output_weights: false
+
+# All projections bias=False (matches FLA)
+add_bias_linear: false
+
+# Scale-only RMSNorm, no bias (matches FLA norm)
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+
+# Positional encoding — none for pure GDN (delta rule recurrence)
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_1B_kda_pure.yaml b/primus/configs/models/megatron/zebra_llama_1B_kda_pure.yaml
new file mode 100644
index 000000000..1143e41db
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_1B_kda_pure.yaml
@@ -0,0 +1,52 @@
+bases:
+  - mamba_base.yaml
+
+# Zebra Llama 1B Pure KDA — matches FLA kda_1B_pure.json exactly
+model_type: mamba
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model size parameters (matches FLA: hidden_size=2048, intermediate_size=8192, num_hidden_layers=16)
+# Megatron counts sublayers: 16 KDA + 16 MLP = 32
+num_layers: 32
+hidden_size: 2048
+ffn_hidden_size: 8192
+
+# Pure KDA — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+
+# Mamba params kept for base config compatibility (unused in pure KDA)
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# KDA parameters — matched to FLA config:
+#   FLA: num_heads=16, head_dim=32, expand_v=2.0, conv_size=4
+#   key: 16 heads x 32 dim = 512 total
+#   value: 16 heads x (32*2.0=64) dim = 1024 total
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 32
+linear_value_head_dim: 64
+linear_num_key_heads: 16
+linear_num_value_heads: 16
+
+# No MLA needed for pure KDA
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: false
+num_attention_heads: 16
+
+# Embedding tying (matches FLA tie_word_embeddings=true)
+untie_embeddings_and_output_weights: false
+
+# Norm epsilon (matches FLA norm_eps=1e-6)
+norm_epsilon: 1.0e-6
+
+# Positional encoding
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_300M_gdn_hybrid.yaml b/primus/configs/models/megatron/zebra_llama_300M_gdn_hybrid.yaml
new file mode 100644
index 000000000..d02db7680
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_300M_gdn_hybrid.yaml
@@ -0,0 +1,87 @@
+extends:
+  - mamba_base.yaml
+
+# 75% Hybrid Gated DeltaNet 300M — matches FLA gated_deltanet_300M_hybrid.json exactly
+#   ~338M params (tied embeddings; +38M vs pure 300M because of the 3 MLA blocks)
+#   12 mixer blocks: 3 MLA + 9 GDN (i.e. 75% GDN, 25% MLA)
+#   +12 MLP blocks = 24 Megatron sublayers
+#
+# Megatron HybridStack.allocate_layers(num_layers=24, hybrid_attention_ratio=0.25):
+#   blocks: [ATTN, MAMBA, MAMBA, MAMBA, ATTN, MAMBA, MAMBA, MAMBA, ATTN, MAMBA, MAMBA, MAMBA]
+#   → MLA at blocks [0, 4, 8] — matches FLA's `attn.layers: [0, 4, 8]` exactly
+model_type: mamba
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model dimensions
+num_layers: 24                  # 12 mixer + 12 MLP sublayers
+hidden_size: 1024
+ffn_hidden_size: 4096
+
+# Hybrid: 25% attention (= 3 of 12 mixer blocks); rest are GDN
+#
+# Megatron-LM removed `--hybrid-attention-ratio` as a CLI argument (only
+# `--hybrid-layer-pattern` remains).  So we declare the concrete pattern
+# explicitly via `hybrid_override_pattern` (the deprecated alias is still
+# forwarded to `--hybrid-layer-pattern`, and unlike the bare arg it lets
+# `num_layers` stay in this YAML as a sanity-check).
+#
+# Pattern: 24 sublayers = 12 mixer blocks alternating with 12 MLPs.
+#   indices 0, 8, 16 → '*' (MLA attention)  ⇔  FLA mixer indices [0, 4, 8]
+#   indices 2, 4, 6, 10, 12, 14, 18, 20, 22 → 'M' (GDN mixer, 9 total)
+#   indices 1, 3, 5, ... 23 → '-' (MLP)
+is_hybrid_model: true
+hybrid_attention_ratio: 0.25     # informational only — upstream ignores it
+hybrid_override_pattern: "*-M-M-M-*-M-M-M-*-M-M-M-"
+
+# Mamba params kept for base config compatibility (unused in GDN)
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# GDN parameters — matched to FLA config:
+#   num_heads=4, num_v_heads=8, head_dim=64, expand_v=1.0, conv_size=4
+#   key_dim   = 4 heads × 64 dim = 256
+#   value_dim = 8 heads × 64 dim = 512   (Grouped Value Attention: each Q/K head serves 2 V heads)
+use_short_conv: true
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 64
+linear_value_head_dim: 64
+linear_num_key_heads: 4
+linear_num_value_heads: 8
+
+# MLA parameters — matched to FLA `attn` block exactly:
+#   num_heads=16, q_lora_rank=672, kv_lora_rank=64,
+#   qk_nope_head_dim=32, qk_rope_head_dim=32 (→ qk_head_dim=64), v_head_dim=64,
+#   rope_theta=500000
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: true
+num_attention_heads: 16         # MLA head count (FLA: attn.num_heads = 16)
+q_lora_rank: 672
+kv_lora_rank: 64
+qk_head_dim: 32                 # = FLA qk_nope_head_dim
+qk_pos_emb_head_dim: 32         # = FLA qk_rope_head_dim
+v_head_dim: 64
+rotary_scaling_factor: 1.0
+mscale: 1.0
+mscale_all_dim: 1.0
+
+# Embedding tying (matches FLA tie_word_embeddings=true)
+untie_embeddings_and_output_weights: false
+
+# All projections bias=False (matches FLA)
+add_bias_linear: false
+
+# Scale-only RMSNorm, no bias (matches FLA norm)
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+
+# Positional encoding — MLA handles its own RoPE (rotary_base=500000);
+# GDN layers use delta-rule recurrence and need no position embedding.
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_300M_gdn_pure.yaml b/primus/configs/models/megatron/zebra_llama_300M_gdn_pure.yaml
new file mode 100644
index 000000000..928869893
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_300M_gdn_pure.yaml
@@ -0,0 +1,59 @@
+extends:
+  - mamba_base.yaml
+
+# Pure Gated DeltaNet 300M — matches FLA gated_deltanet_300M_pure.json exactly
+#   ~300M params (tied embeddings)
+#   12 GDN blocks + 12 MLP blocks = 24 Megatron sublayers
+model_type: mamba
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model dimensions
+num_layers: 24
+hidden_size: 1024
+ffn_hidden_size: 4096
+
+# Pure GDN — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+
+# Mamba params kept for base config compatibility (unused in pure GDN)
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# GDN parameters — matched to FLA config:
+#   num_heads=4, num_v_heads=8, head_dim=64, expand_v=1.0, conv_size=4
+#   key_dim  = 4  heads x 64 dim = 256
+#   value_dim = 8 heads x 64 dim = 512
+#   Grouped Value Attention: each Q/K head serves 2 V heads
+use_short_conv: true
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 64
+linear_value_head_dim: 64
+linear_num_key_heads: 4
+linear_num_value_heads: 8
+
+# No MLA needed for pure GDN
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: false
+num_attention_heads: 4
+
+# Embedding tying (matches FLA tie_word_embeddings=true)
+untie_embeddings_and_output_weights: false
+
+# All projections bias=False (matches FLA)
+add_bias_linear: false
+
+# Scale-only RMSNorm, no bias (matches FLA norm)
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+
+# Positional encoding — none for pure GDN (delta rule recurrence)
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml b/primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml
new file mode 100644
index 000000000..c54f0cdab
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_300M_kda_pure.yaml
@@ -0,0 +1,57 @@
+extends:
+  - mamba_base.yaml
+
+# Pure Kimi Delta Attention 300M — matches FLA kda_300M_pure.json exactly
+#   ~300M params (tied embeddings)
+#   12 KDA blocks + 12 MLP blocks = 24 Megatron sublayers
+model_type: mamba
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# Model dimensions
+num_layers: 24
+hidden_size: 1024
+ffn_hidden_size: 4096
+
+# Pure KDA — no attention layers
+is_hybrid_model: true
+hybrid_attention_ratio: 0.0
+
+# Mamba params kept for base config compatibility (unused in pure KDA)
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+
+# KDA parameters — matched to FLA config:
+#   num_heads=8, head_dim=32 (key), expand_v=2.0, conv_size=4
+#   key_dim   = 8 heads x 32 dim         = 256
+#   value_dim = 8 heads x (32*2.0=64) dim = 512
+linear_conv_kernel_dim: 4
+linear_key_head_dim: 32
+linear_value_head_dim: 64
+linear_num_key_heads: 8
+linear_num_value_heads: 8
+
+# No MLA needed for pure KDA
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: false
+num_attention_heads: 8
+
+# Embedding tying (matches FLA tie_word_embeddings=true)
+untie_embeddings_and_output_weights: false
+
+# All projections bias=False (matches FLA)
+add_bias_linear: false
+
+# Scale-only RMSNorm, no bias (matches FLA norm)
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+
+# Positional encoding — none for pure KDA (delta-rule recurrence)
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron/zebra_llama_300M_mamba_hybrid.yaml b/primus/configs/models/megatron/zebra_llama_300M_mamba_hybrid.yaml
new file mode 100644
index 000000000..54b207556
--- /dev/null
+++ b/primus/configs/models/megatron/zebra_llama_300M_mamba_hybrid.yaml
@@ -0,0 +1,75 @@
+bases:
+  - mamba_base.yaml
+
+# 75% Hybrid Mamba2 + MLA 300M — parallel to the GDN hybrid 300M, drop-in
+# replacement of the GDN mixer with Mamba2 (SSM).
+#
+#   12 mixer blocks: 3 MLA + 9 Mamba2 (i.e. 75% Mamba2, 25% MLA)
+#   +12 MLP blocks = 24 Megatron sublayers
+#
+# Megatron HybridStack with hybrid_override_pattern '*-M-M-M-*-M-M-M-*-M-M-M-':
+#   MLA at mixer indices [0, 4, 8] (matches the user's `"layers": [0, 4, 8]`).
+#
+# Uses the same `hybrid_stack_spec` as zebra_llama_{1B,3B,8B}.yaml (the proven
+# TE-backed Mamba2+MLA spec); model dims scaled down to 300M (same hidden/ffn
+# as the GDN hybrid 300M for an apples-to-apples comparison).
+
+model_type: mamba                                 # CRITICAL: hybrid models use mamba model type
+tokenizer_type: HuggingFaceTokenizer
+tokenizer_model: meta-llama/Llama-3.2-1B
+
+# ── Model dimensions ──
+num_layers: 24                                    # 12 mixer + 12 MLP sublayers
+hidden_size: 1024
+ffn_hidden_size: 4096
+
+# ── Hybrid pattern ──
+# 24 sublayers = 12 mixer blocks alternating with 12 MLPs.
+#   indices 0, 8, 16 → '*' (MLA attention) ⇔ FLA-style mixer indices [0, 4, 8]
+#   indices 2, 4, 6, 10, 12, 14, 18, 20, 22 → 'M' (Mamba2 mixer, 9 total)
+#   indices 1, 3, 5, …, 23 → '-' (MLP)
+is_hybrid_model: true
+hybrid_attention_ratio: 0.25                     # informational (allocator uses pattern below)
+hybrid_override_pattern: "*-M-M-M-*-M-M-M-*-M-M-M-"
+
+# ── Mamba2 parameters ──
+# Match the 1B/3B/8B convention (mamba_state_dim=64, head_dim=64, groups=8) so
+# the 300M run is just a smaller version of the same architecture.  Expand=2 is
+# the standard Mamba2 expansion ratio; d_conv=4 matches the spec params.
+mamba_state_dim: 64
+mamba_head_dim: 64
+mamba_num_groups: 8
+mamba_expand: 2
+
+# ── MLA parameters (identical to the 300M GDN hybrid for direct comparison) ──
+group_query_attention: false
+swiglu: true
+num_query_groups: null
+multi_latent_attention: true
+num_attention_heads: 16                           # MLA head count
+q_lora_rank: 672
+kv_lora_rank: 64
+qk_head_dim: 32                                   # = qk_nope_head_dim
+qk_pos_emb_head_dim: 32                           # = qk_rope_head_dim
+v_head_dim: 64
+rotary_scaling_factor: 1.0
+mscale: 1.0
+mscale_all_dim: 1.0
+
+# Embedding tying
+untie_embeddings_and_output_weights: false
+
+# All projections bias=False
+add_bias_linear: false
+
+# Scale-only RMSNorm, no bias (1e-6 set in pretrain override)
+normalization: RMSNorm
+norm_epsilon: 1.0e-6
+
+# Positional encoding — MLA handles its own RoPE (rotary_base=500000);
+# Mamba2 layers are recurrent and don't use position embeddings.
+rotary_base: 500000
+position_embedding_type: none
+add_position_embedding: true
+use_rotary_position_embeddings: false
+max_position_embeddings: 131072
diff --git a/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml b/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml
new file mode 100644
index 000000000..a2690cbf0
--- /dev/null
+++ b/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml
@@ -0,0 +1,4 @@
+# MXFP4 LoRA recipe for MLPerf Llama2-70B (MI355X)
+recipe: primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom
+flavor: llama2_70b_lora_mxfp4_config
+hf_path: meta-llama/Llama-2-70b-hf
diff --git a/primus/configs/models/megatron_bridge/mamba_130M.yaml b/primus/configs/models/megatron_bridge/mamba_130M.yaml
new file mode 100644
index 000000000..463be8c69
--- /dev/null
+++ b/primus/configs/models/megatron_bridge/mamba_130M.yaml
@@ -0,0 +1,2 @@
+recipe: mamba.mamba2
+flavor: mamba2_130m_pretrain_config
diff --git a/primus/configs/models/nemo_automodel/flux_1_dev.yaml b/primus/configs/models/nemo_automodel/flux_1_dev.yaml
new file mode 100644
index 000000000..08e55ae37
--- /dev/null
+++ b/primus/configs/models/nemo_automodel/flux_1_dev.yaml
@@ -0,0 +1,30 @@
+# NeMo AutoModel model preset: FLUX.1-dev (12B rectified-flow T2I DiT).
+#
+# Merged (as the `model:` block) into the nemo_automodel flux_pre_trainer module
+# config by the Primus launcher. `mode: pretrain` builds ONE randomly-initialised
+# FluxTransformer2DModel from the local diffusers `transformer/config.json`
+# (no weights loaded) — the same from-scratch pretrain path Wan2.2 uses.
+#
+# The VAE + CLIP + T5 encoders are NOT loaded at train time: the dataloader reads
+# a pre-encoded latent/embedding cache (see flux_pre_trainer.yaml data.dataloader),
+# so only the transformer lives on-GPU. Build the cache once with the stock
+# AutoModel preprocessor (tools.diffusion.preprocessing_multiprocess image
+# --processor flux --model_name <this path>).
+model:
+  # Local FLUX.1-dev diffusers snapshot (random init in pretrain mode). Point at
+  # the same dir used to encode the cache so the VAE latent norm matches.
+  pretrained_model_name_or_path: /models/FLUX.1-dev
+  mode: pretrain
+  # Explicit model family (used by the recipe for PEFT and cache validation).
+  model_type: flux
+  # FLUX joint attention. `flash` is the dense, backend-stable choice on gfx950
+  # and is what the nondeterministic-bf16 attention hook re-binds
+  # (PRIMUS_ATTN_NONDETERMINISTIC=1) for the throughput A/B. Falls back cleanly.
+  attention_backend: flash
+  # fp32 master + bf16 compute (canonical pretraining precision, GB300-comparable).
+  torch_dtype: float32
+  compute_dtype: bfloat16
+  pipeline_spec:
+    transformer_cls: FluxTransformer2DModel
+    subfolder: transformer
+    load_full_pipeline: false
diff --git a/primus/configs/models/nemo_automodel/flux_1_schnell.yaml b/primus/configs/models/nemo_automodel/flux_1_schnell.yaml
new file mode 100644
index 000000000..75c6712a7
--- /dev/null
+++ b/primus/configs/models/nemo_automodel/flux_1_schnell.yaml
@@ -0,0 +1,25 @@
+# NeMo AutoModel model preset: FLUX.1-schnell (12B rectified-flow T2I DiT).
+#
+# Identical transformer architecture to FLUX.1-dev (19 dual + 38 single blocks,
+# 24 heads x 128 dim, in_channels 64, joint_attention_dim 4096) EXCEPT
+# `guidance_embeds: false` — schnell is a guidance-distilled model and its
+# transformer has NO guidance embedding layer. Therefore the FLUX flow-matching
+# adapter must run with `use_guidance_embeds: false` (set in the schnell EXP
+# configs) so it does not pass a guidance tensor into the transformer forward.
+#
+# `mode: pretrain` builds ONE randomly-initialised FluxTransformer2DModel from the
+# local `transformer/config.json` (no safetensors loaded) — so only the schnell
+# transformer/config.json (+ model_index.json) needs to live under the path below.
+# The VAE + CLIP + T5 encoders are NOT loaded at train time (the dataloader reads
+# a pre-encoded latent/embedding cache), so this preset never touches them.
+model:
+  pretrained_model_name_or_path: /checkpoints/flux_schnell_local
+  mode: pretrain
+  model_type: flux
+  attention_backend: flash
+  torch_dtype: float32
+  compute_dtype: bfloat16
+  pipeline_spec:
+    transformer_cls: FluxTransformer2DModel
+    subfolder: transformer
+    load_full_pipeline: false
diff --git a/primus/configs/models/nemo_automodel/ideogram4.yaml b/primus/configs/models/nemo_automodel/ideogram4.yaml
new file mode 100644
index 000000000..3955cd106
--- /dev/null
+++ b/primus/configs/models/nemo_automodel/ideogram4.yaml
@@ -0,0 +1,36 @@
+# NeMo AutoModel model preset: Ideogram-4 (single-stream rectified-flow T2I DiT,
+# ~9.3B at the default 34-layer config; Primus backend: nemo_automodel).
+#
+# Merged (as the `model:` block) into the nemo_automodel ideogram4_pre_trainer
+# module config by the Primus launcher. `mode: pretrain` builds ONE randomly-
+# initialised Ideogram4Transformer2DModel from the local diffusers
+# `transformer/config.json` (no weights loaded) — the same from-scratch pretrain
+# path Wan2.2 / FLUX use.
+#
+# NO gated weights are required: the from_config path only needs the transformer
+# CONFIG. Generate a weightless config dir once (meta device, no 9.3B alloc):
+#   OUT=/workspace/ideogram4_config \
+#   python .../ttt_l2_02_IDEOGRAM/scripts/make_ideogram_config_dir.py
+# then point `pretrained_model_name_or_path` at it (default below).
+#
+# The Flux-2 VAE + Qwen3-VL encoder are NOT loaded at train time: the dataloader
+# provides pre-encoded image latents [B,128,H_p,W_p] + Qwen3-VL features
+# [B,·,53248] (synthetic here; real preprocessor cache in Phase B.3), so only the
+# transformer lives on-GPU.
+model:
+  # Local weightless Ideogram-4 diffusers config dir (random init in pretrain mode).
+  pretrained_model_name_or_path: /workspace/ideogram4_config
+  mode: pretrain
+  # Explicit model family (used by the recipe only for PEFT/cache paths; unused here).
+  model_type: ideogram4
+  # NOTE: attention_backend intentionally unset -> the recipe skips
+  # set_attention_backend, so the model keeps its default dispatch (torch SDPA),
+  # which honours the Ideogram block-diagonal segment mask. This is the path
+  # validated in Phase A/B. Set a backend here only once mask support is confirmed.
+  # fp32 master + bf16 compute (canonical pretraining precision, GB300-comparable).
+  torch_dtype: float32
+  compute_dtype: bfloat16
+  pipeline_spec:
+    transformer_cls: Ideogram4Transformer2DModel
+    subfolder: transformer
+    load_full_pipeline: false
diff --git a/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml b/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml
new file mode 100644
index 000000000..0bbbbec40
--- /dev/null
+++ b/primus/configs/models/nemo_automodel/wan2_2_t2v_a14b.yaml
@@ -0,0 +1,23 @@
+# NeMo AutoModel model preset: Wan 2.2 T2V A14B (mixture-of-experts DiT).
+#
+# Merged (as the `model:` block) into the nemo_automodel pre_trainer module
+# config by the Primus launcher. `mode: pretrain` builds ONE randomly-initialised
+# A14B expert chosen by `pipeline_spec.subfolder`:
+#   high-noise expert -> transformer        (default; stage: high_noise)
+#   low-noise  expert -> transformer_2      (override stage: low_noise too)
+model:
+  pretrained_model_name_or_path: /models/Wan2.2-T2V-A14B-Diffusers
+  mode: pretrain
+  stage: high_noise
+  # REQUIRED in pretrain mode (the from_config pipeline carries no config to
+  # infer it from). Wan 2.2-A14B ships boundary_ratio = 0.875.
+  boundary_ratio: 0.875
+  # AITER attention backend (gfx950). Falls back to "flash" if needed.
+  attention_backend: aiter
+  # fp32 master + bf16 compute (canonical pretraining precision, GB300-comparable).
+  torch_dtype: float32
+  compute_dtype: bfloat16
+  pipeline_spec:
+    transformer_cls: WanTransformer3DModel
+    subfolder: transformer
+    load_full_pipeline: false
diff --git a/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml b/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml
new file mode 100644
index 000000000..d35a4795d
--- /dev/null
+++ b/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml
@@ -0,0 +1,14 @@
+job:
+  dump_folder: "./outputs"
+  description: "GPT-OSS 120B FP8 training"
+
+model:
+  name: "gpt_oss"
+  flavor: "120b"
+  hf_assets_path: "openai/gpt-oss-120b"
+  # Attention (sinks + sliding window) is routed to Primus-Turbo via the
+  # gptoss_sink_attention setup patch. FP8 dense linears are enabled through the
+  # quantize.linear.float8 converter, which the turbo_float8_linear patch swaps
+  # for the Primus-Turbo FP8 converter when use_turbo_float8_linear is set.
+  converters:
+    - quantize.linear.float8
diff --git a/primus/configs/models/torchtitan/gpt_oss_120b.yaml b/primus/configs/models/torchtitan/gpt_oss_120b.yaml
new file mode 100644
index 000000000..4cccb592b
--- /dev/null
+++ b/primus/configs/models/torchtitan/gpt_oss_120b.yaml
@@ -0,0 +1,12 @@
+job:
+  dump_folder: "./outputs"
+  description: "GPT-OSS 120B training"
+
+model:
+  name: "gpt_oss"
+  flavor: "120b"
+  hf_assets_path: "openai/gpt-oss-120b"
+  # No model converter: GPT-OSS attention (learnable sinks + sliding window) is
+  # routed to Primus-Turbo via the torchtitan.primus_turbo.gptoss_sink_attention
+  # setup patch (functional flash_attn_func), not via the module-form converter.
+  converters: []
diff --git a/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml b/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml
new file mode 100644
index 000000000..752ea87df
--- /dev/null
+++ b/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml
@@ -0,0 +1,14 @@
+job:
+  dump_folder: "./outputs"
+  description: "GPT-OSS 20B FP8 training"
+
+model:
+  name: "gpt_oss"
+  flavor: "20b"
+  hf_assets_path: "openai/gpt-oss-20b"
+  # Attention (sinks + sliding window) is routed to Primus-Turbo via the
+  # gptoss_sink_attention setup patch. FP8 dense linears are enabled through the
+  # quantize.linear.float8 converter, which the turbo_float8_linear patch swaps
+  # for the Primus-Turbo FP8 converter when use_turbo_float8_linear is set.
+  converters:
+    - quantize.linear.float8
diff --git a/primus/configs/models/torchtitan/gpt_oss_20b.yaml b/primus/configs/models/torchtitan/gpt_oss_20b.yaml
new file mode 100644
index 000000000..f1ee457a9
--- /dev/null
+++ b/primus/configs/models/torchtitan/gpt_oss_20b.yaml
@@ -0,0 +1,12 @@
+job:
+  dump_folder: "./outputs"
+  description: "GPT-OSS 20B training"
+
+model:
+  name: "gpt_oss"
+  flavor: "20b"
+  hf_assets_path: "openai/gpt-oss-20b"
+  # No model converter: GPT-OSS attention (learnable sinks + sliding window) is
+  # routed to Primus-Turbo via the torchtitan.primus_turbo.gptoss_sink_attention
+  # setup patch (functional flash_attn_func), not via the module-form converter.
+  converters: []
diff --git a/primus/configs/modules/diffusion/post_trainer.yaml b/primus/configs/modules/diffusion/post_trainer.yaml
new file mode 100644
index 000000000..e89d7b62f
--- /dev/null
+++ b/primus/configs/modules/diffusion/post_trainer.yaml
@@ -0,0 +1,5 @@
+extends:
+  - ../module_base.yaml
+
+trainable: true
+stage: posttrain
diff --git a/primus/configs/modules/diffusion/pre_trainer.yaml b/primus/configs/modules/diffusion/pre_trainer.yaml
new file mode 100644
index 000000000..3d9d3cbdb
--- /dev/null
+++ b/primus/configs/modules/diffusion/pre_trainer.yaml
@@ -0,0 +1,4 @@
+extends:
+  - ../module_base.yaml
+
+stage: pretrain
diff --git a/primus/configs/modules/maxtext/pre_trainer.yaml b/primus/configs/modules/maxtext/pre_trainer.yaml
index b16135df7..7ffcb2f09 100644
--- a/primus/configs/modules/maxtext/pre_trainer.yaml
+++ b/primus/configs/modules/maxtext/pre_trainer.yaml
@@ -16,8 +16,12 @@ hf_path: "allenai/c4"  # for using https://huggingface.co/datasets/allenai/c4
 hf_data_dir: "en"
 hf_train_files: ""
 packing: true
+# Required by MaxText v26.4's strict config validation: when hardware=gpu,
+# packing=true, and attention=cudnn_flash_te, max_segments_per_seq must be > 0.
+# 32 matches upstream MaxText GPU model configs and is harmless on v26.3.
+max_segments_per_seq: 32
 
-shardy: false
+shardy: true
 
 remat_policy: 'full'
 optimizer_memory_host_offload: false
diff --git a/primus/configs/modules/megatron/primus_megatron_module.yaml b/primus/configs/modules/megatron/primus_megatron_module.yaml
index dce09dbbe..182ae9317 100644
--- a/primus/configs/modules/megatron/primus_megatron_module.yaml
+++ b/primus/configs/modules/megatron/primus_megatron_module.yaml
@@ -70,7 +70,7 @@ mlflow_upload_performance_metrics: false
 disable_profiler_activity_cpu: false
 torch_profiler_record_shapes: true
 torch_profiler_with_stack: true
-torch_profiler_use_gzip: false
+torch_profiler_use_gzip: true
 
 # continue/finetune
 auto_continue_train: false
@@ -92,7 +92,12 @@ pp_warmup: false # set to true to decrease iter-1 time when using pp
 dump_pp_data: false
 
 # recompute
-recompute_layer_ids: null #int list，id srange from 0 to (num_layers_per_pp_stage - 1)
+recompute_layer_ids: null # int list; global layer ids, range from 0 to (num_layers - 1)
 
 # dataloader
 dataloader_mp_context: null # "forkserver" | "spawn" | "fork" | null
+
+# Control the force load balancing type for the MoE router.
+# If set to "even", the router will force the load balancing to be even.
+# If set to "uniform", the router will force the load balancing to be uniform. (Megatron-LM original behavior)
+moe_router_force_load_balancing_type: "even"
diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml
index 290024b3a..f7dbbdbd4 100644
--- a/primus/configs/modules/megatron/primus_turbo.yaml
+++ b/primus/configs/modules/megatron/primus_turbo.yaml
@@ -1,8 +1,16 @@
 # ===== Global =====
 # main control flag
 enable_primus_turbo: false
-# use primus_turbo_fp4_autocast instead of TE fp8_autocast for FP4 (requires enable_primus_turbo)
-use_turbo_fp4_autocast: false
+
+# ===== AutoTune =====
+# enable turbo auto-tuning
+use_turbo_autotune: false
+
+# ===== MXFP4 =====
+# Use stochastic rounding when quantizing MXFP4 gradients. This is disabled by
+# default and requires a Primus-Turbo grouped GEMM implementation with matching
+# activation and gradient layouts.
+mxfp4_gradient_stochastic_rounding: false
 
 # ===== Attention =====
 # operator switch
@@ -25,6 +33,9 @@ use_turbo_gemm: false
 # ===== Grouped MLP =====
 # use turbo grouped gemm
 use_turbo_grouped_gemm: false
+# let tensorwise FP8 Turbo grouped GEMM consume non-aligned expert token counts
+# directly, without TE Fp8Padding/Fp8Unpadding
+use_turbo_ragged_grouped_gemm: false
 # fused activation_with_probs to reduce redundant computation
 use_turbo_fused_act_with_probs: false
 
@@ -40,6 +51,14 @@ turbo_deepep_use_comm_stream: false
 # ===== Sync-Free MoE =====
 turbo_sync_free_moe_stage: 0
 
+# ===== Mega MoE =====
+# fused MegaMoE layer replacement (EP-only/TP=1/bf16)
+use_turbo_mega_moe: false
+
 # ===== Layer Norm =====
 # operator switch
 use_turbo_rms_norm: false
+# Route the fused layernorm+QKV site through Turbo's rmsnorm while leaving its
+# GEMM on Transformer Engine. use_turbo_gemm moves that GEMM to Turbo as well,
+# which costs more than the norm saves on the GPT-OSS-20B shapes.
+use_turbo_norm_te_linear: false
diff --git a/primus/configs/modules/megatron/sft_trainer.yaml b/primus/configs/modules/megatron/sft_trainer.yaml
index 32da1b266..a0ba26557 100644
--- a/primus/configs/modules/megatron/sft_trainer.yaml
+++ b/primus/configs/modules/megatron/sft_trainer.yaml
@@ -63,6 +63,9 @@ use_distributed_optimizer: true
 # Note: These are typically overridden in the experiment config
 sft_dataset_name: "tatsu-lab/alpaca"
 sft_conversation_format: "alpaca"
+# HuggingFace dataset split to load. Some datasets have no "train" split (e.g.
+# SWE-bench/SWE-smith-trajectories use tool/xml/ticks); override here.
+sft_dataset_split: "train"
 
 # Sequence packing: concatenate multiple short samples into a single
 # max_seq_length sequence. Big throughput win on short-sample datasets such
@@ -84,6 +87,29 @@ enable_packed_sequences: false
 #            short-sample datasets.
 use_packed_attention: false
 
+# === LB-Mini (ODC sequence-length load balancing) ===
+# When true (and enable_odc=true + use_torch_fsdp2=true), the TRAIN data is
+# served variable-length and Karmarkar-Karp balanced across DP ranks, and each
+# rank runs its OWN (possibly different) number of micro-batches -- removing the
+# workload-imbalance bubble on variable-length datasets (e.g. LongAlign).
+# Requires ODC point-to-point comm; it would deadlock under NCCL collectives.
+# DEFAULT false => stock Megatron fixed-num_microbatches, all-ranks-in-lockstep
+# schedule, byte-for-byte unchanged.
+enable_odc_lb_mini: false
+# Per-micro-batch token cap for LB-Mini (0 => use seq_length).
+lb_mini_max_token_len: 0
+# LB-Mini cost model: "linear" (token count) or "fit" (attention-aware, ~s^2).
+lb_mini_cost_model: "linear"
+
+# === Chunked linear cross-entropy (avoid full-logits OOM) ===
+# When true (or env FUSED_LINEAR_CE=1), GPTModel computes the LM loss in
+# sequence chunks under activation checkpointing, so the full [seq, vocab]
+# logits is never materialized. Essential for large-vocab models (e.g. Qwen
+# 151936) at long sequence lengths where the logits tensor alone would OOM.
+# DEFAULT false => stock full-logits path, byte-for-byte unchanged. Numerically
+# equivalent when on. Tune chunk via env FUSED_CE_CHUNK (default 4096).
+enable_fused_linear_ce: false
+
 # Opt-in: reproduce NeMo Megatron-Bridge's packed dataset byte-for-byte.
 #
 # When False (default) Native does a clean single-shot tokenize of the full
diff --git a/primus/configs/modules/megatron/torch_compile.yaml b/primus/configs/modules/megatron/torch_compile.yaml
new file mode 100644
index 000000000..8b5550feb
--- /dev/null
+++ b/primus/configs/modules/megatron/torch_compile.yaml
@@ -0,0 +1,13 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# Torch Compile Configuration
+# Default values for torch.compile settings (Primus-specific)
+torch_compile:
+  enable: false
+  backend: "inductor"
+  mode: "default"
+  fullgraph: false
diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml
index ba89bd7d1..65dc2f9e0 100755
--- a/primus/configs/modules/megatron/trainer_base.yaml
+++ b/primus/configs/modules/megatron/trainer_base.yaml
@@ -2,6 +2,7 @@ extends:
   - ../module_base.yaml
   - primus_megatron_module.yaml
   - primus_turbo.yaml
+  - torch_compile.yaml
   - zero_bubble.yaml
   - primus_pipeline.yaml
 
@@ -97,7 +98,9 @@ use_checkpoint_opt_param_scheduler: false
 warmup: null
 decoupled_lr: null
 decoupled_min_lr: null
-# muon
+# muon (read by Megatron-LM's native get_megatron_optimizer_config() whenever
+# args has these attrs -- not consumed by any Primus code, but the only way
+# to tune Muon's internals via Primus config since most have no CLI flag)
 muon_extra_scale_factor: 1.0
 muon_scale_mode: "spectral"
 muon_fp32_matmul_prec: "medium"
@@ -106,8 +109,6 @@ muon_tp_mode: "blockwise"
 muon_use_nesterov: false
 muon_split_qkv: true
 muon_momentum: 0.95
-muon_weight_decay: 0.01
-muon_weight_decay_method: "decoupled"
 
 optimizer_cpu_offload: false
 optimizer_offload_fraction: 1.0 # float
@@ -190,7 +191,40 @@ gradient_reduce_div_fusion: true
 suggested_communication_unit_size: 400000000 # int
 keep_fp8_transpose_cache_when_using_custom_fsdp: false
 num_distributed_optimizer_instances: 1 # int
+data_parallel_replicate_degree: 1
 use_torch_fsdp2: false
+
+# === ODC (On-Demand Communication) ===
+# Master switch for ODC's rocSHMEM/XGMI point-to-point gradient reduction on the
+# torch-FSDP2 path (ROCm/MI300X). DEFAULT false => ODC is a complete no-op and
+# every ODC integration patch is skipped (byte-for-byte stock behaviour). When
+# true (requires use_torch_fsdp2=true), Primus:
+#   * routes FSDP2 gradient reduce-scatter through ODC (scatter_accumulate),
+#   * skips #808's FSDP2 forward all-gather prefetch (it cannot overlap ODC's
+#     serial transport); backward prefetch is kept.
+# Formerly gated by the ODC_ENABLE env var; now config-driven.
+enable_odc: false
+# ODC wiring depth (integer): 1 = install the __init__/lazy_init hooks only
+# (construction + first fwd/bwd sanity, gradients NOT routed); 2 = also hook
+# train_step / optimizer.step for full gradient routing. Production = 2.
+odc_phase: 2
+# Guard against ODC async-reduce's occasional non-deterministic grad spike: at
+# the optimizer step, if the global grad norm exceeds this threshold, the step is
+# skipped (grads zeroed) exactly like Megatron skips nan/inf iters. <= 0 disables.
+odc_grad_spike_threshold: 1000.0
+# ODC symmetric-memory / comm-kernel tuning (consumed by the ODC library via its
+# runtime config; these replace the former ODC_* env vars). Defaults reproduce the
+# previous env defaults exactly.
+odc_p2p_backend: mori          # symmetric-memory backend: mori | rocshmem
+odc_mori_init: pg              # MORI init method: pg | uid
+odc_max_buffer_size: 67108864  # BufferSplitter max global buffer (bytes, 64 MiB)
+odc_rocshmem_gda: false        # rocSHMEM GPU-direct (GDA) device path (multi-node)
+odc_rocshmem_lib: null         # optional librs_host_gda.so ctypes override path
+odc_gda_rs_blocks: 64          # GDA reduce-scatter kernel grid blocks
+odc_gda_pipe: 1                # GDA peer-pipeline depth (also bridged to turbo env)
+odc_gda_defer_reduce: auto     # defer cross-node reduce to per-minibatch: auto|1|0
+odc_gda_warmup_mode: strided   # write-visibility settle: strided|full|hdp|fence|hdpfence
+odc_gda_stride_bytes: 65536    # strided warm-up page stride (bytes)
 nccl_communicator_config_path: null
 use_tp_pp_dp_mapping: false
 replication: false
@@ -235,6 +269,20 @@ data_args_path: null # str
 per_split_data_args_path: null # str
 data_cache_path: null
 mock_data: false
+# mock_dataset: Configuration for synthetic dataset (used when mock_data=true)
+#   class: Fully qualified class name (optional, defaults based on model_type)
+#   params: Dict of parameters passed to dataset constructor
+#   Example:
+#     mock_dataset:
+#       class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
+#       params:
+#         num_samples: 1000
+#         image_size: 512
+mock_dataset:
+  class: null
+  params:
+    num_samples: 1000
+    image_size: 512
 merge_file: null
 seq_length: 4096
 encoder_seq_length: null
@@ -269,6 +317,7 @@ log_avg_reset_interval: 10
 log_params_norm: false
 log_num_zeros_in_grad: false
 log_throughput: false
+wall_clock_step_timer: false
 log_progress: false
 timing_log_level: 0
 timing_log_option: minmax
diff --git a/primus/configs/modules/nemo_automodel/_ddp_pre_trainer.yaml b/primus/configs/modules/nemo_automodel/_ddp_pre_trainer.yaml
new file mode 100644
index 000000000..a4a345005
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/_ddp_pre_trainer.yaml
@@ -0,0 +1,80 @@
+# SCRATCH (Phase-2 sharding sweep, delete after) — Wan2.2 DDP variant of
+# pre_trainer.yaml: the `fsdp:` block is replaced by a `ddp:` block so the
+# diffusion recipe takes the DistributedDataParallel path (full replication,
+# no param/optimizer sharding). Expected to OOM at A14B; that OOM is the result.
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 256
+  local_batch_size: 32
+  ckpt_every_steps: 1000
+  num_epochs: 100
+  log_remote_every_steps: 2
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
+    cache_dir: /dataset
+    model_type: wan
+    base_resolution: [512, 512]
+    dynamic_batch_size: false
+    shuffle: true
+    drop_last: false
+    num_workers: 2
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 5e-5
+  weight_decay: 0.1
+  betas: [0.9, 0.95]
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: simple
+  adapter_kwargs: {}
+  timestep_sampling: uniform
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 3.0
+  mix_uniform_ratio: 0.1
+  num_train_timesteps: 1000
+  i2v_prob: 0.3
+  use_loss_weighting: true
+  log_interval: 100
+  summary_log_interval: 10
+
+ddp:
+  activation_checkpointing: true
+  gradient_as_bucket_view: true
+  broadcast_buffers: false
+  find_unused_parameters: false
+  static_graph: false
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/wan22_primus_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/_flux_ddp_pre_trainer.yaml b/primus/configs/modules/nemo_automodel/_flux_ddp_pre_trainer.yaml
new file mode 100644
index 000000000..81daa6c23
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/_flux_ddp_pre_trainer.yaml
@@ -0,0 +1,83 @@
+# SCRATCH (Phase-2 sharding sweep, delete after) — FLUX.1-dev DDP variant of
+# flux_pre_trainer.yaml: `fsdp:` block replaced by `ddp:` so the diffusion recipe
+# takes the DistributedDataParallel path (full replication, no sharding).
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 32
+  local_batch_size: 4
+  ckpt_every_steps: 1000
+  num_epochs: 1000
+  log_remote_every_steps: 2
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: nemo_automodel.components.datasets.diffusion.build_text_to_image_multiresolution_dataloader
+    cache_dir: /dataset
+    train_text_encoder: false
+    base_resolution: [512, 512]
+    dynamic_batch_size: false
+    shuffle: true
+    drop_last: false
+    num_workers: 2
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 2e-5
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: flux
+  adapter_kwargs:
+    guidance_scale: 1
+    use_guidance_embeds: true
+  timestep_sampling: uniform
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 1.0
+  mix_uniform_ratio: 0.1
+  sigma_min: 0.0
+  sigma_max: 1.0
+  num_train_timesteps: 1000
+  i2v_prob: 0.0
+  use_loss_weighting: false
+  log_interval: 100
+  summary_log_interval: 10
+
+ddp:
+  activation_checkpointing: true
+  gradient_as_bucket_view: true
+  broadcast_buffers: false
+  find_unused_parameters: false
+  static_graph: false
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/flux_primus_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/flux_pre_trainer.yaml b/primus/configs/modules/nemo_automodel/flux_pre_trainer.yaml
new file mode 100644
index 000000000..76540b16b
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/flux_pre_trainer.yaml
@@ -0,0 +1,114 @@
+# NeMo AutoModel pre_trainer module config — FLUX.1 T2I rectified-flow pretrain
+# (Primus backend: nemo_automodel).
+#
+# Sibling of `pre_trainer.yaml` (the Wan2.2 video preset). Same recipe schema and
+# the SAME optimized knobs (torch.compile + bf16 grad-reduce + deep prefetch +
+# fused AdamW), but wired for FLUX text-to-IMAGE:
+#   - data.dataloader -> build_text_to_image_multiresolution_dataloader
+#     (the image loader takes NO `model_type`, unlike the video loader).
+#   - flow_matching.adapter_type -> flux (guidance-embed conditioning).
+#
+# The Primus launcher deep-merges the `model:` preset (flux_1_dev.yaml) and the
+# EXP `overrides:` on top of this, then runs TrainDiffusionRecipe.
+#
+# NOTE: do not declare a `model:` block here — it comes from the model preset.
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 32
+  local_batch_size: 4
+  ckpt_every_steps: 1000
+  num_epochs: 1000
+  log_remote_every_steps: 2
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: nemo_automodel.components.datasets.diffusion.build_text_to_image_multiresolution_dataloader
+    cache_dir: /dataset
+    train_text_encoder: false
+    # Supported cache resolutions: [256,256], [512,512], [1024,1024].
+    base_resolution: [512, 512]
+    dynamic_batch_size: false
+    shuffle: true
+    drop_last: false
+    num_workers: 2
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 2e-5
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  # fused AdamW (mutually exclusive with foreach).
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: flux
+  adapter_kwargs:
+    # Critical: guidance_scale 1 for pretraining. The FluxAdapter ALWAYS passes a
+    # guidance tensor, so FLUX.1-dev's transformer (guidance_embeds=true) is required.
+    guidance_scale: 1
+    use_guidance_embeds: true
+  timestep_sampling: uniform
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 1.0
+  mix_uniform_ratio: 0.1
+  sigma_min: 0.0
+  sigma_max: 1.0
+  num_train_timesteps: 1000
+  i2v_prob: 0.0
+  # Critical: use_loss_weighting must be false for pretraining.
+  use_loss_weighting: false
+  log_interval: 100
+  summary_log_interval: 10
+
+fsdp:
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  dp_replicate_size: 1
+  # dp_size intentionally unset: AutoModel derives it from world size and the
+  # parallel dims. Pin `dp_size: <n>` via overrides only to force a value.
+  defer_fsdp_grad_sync: true
+  # ZeRO-2: blocks keep their gathered params through backward (one all-gather per block
+  # per step instead of two) at the cost of roughly one parameter copy of memory. Set
+  # `true` for ZeRO-3, or remove the key to use AutoModel's per-layer heuristic.
+  # Unlike the Ideogram presets this one has enable_compile on, so A/B any throughput
+  # change here with compile held constant.
+  reshard_after_forward: false
+  # NOTE: inert while reshard_after_forward is false (no backward all-gather to prefetch).
+  enable_fsdp2_prefetch: true
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
+  reduce_dtype: bfloat16
+  enable_compile: true
+  activation_checkpointing: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/flux_primus_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/ideogram4_ddp_pre_trainer.yaml b/primus/configs/modules/nemo_automodel/ideogram4_ddp_pre_trainer.yaml
new file mode 100644
index 000000000..80b88e167
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/ideogram4_ddp_pre_trainer.yaml
@@ -0,0 +1,118 @@
+# NeMo AutoModel pre_trainer module config — Ideogram-4 single-stream T2I
+# pretrain, DDP + ZeRO-1 variant (Primus backend: nemo_automodel).
+#
+# Identical to ideogram4_pre_trainer.yaml EXCEPT the `fsdp:` block is replaced by
+# a `ddp:` block, so the diffusion recipe takes the DistributedDataParallel path
+# (params + grads REPLICATED, no per-layer all-gather). Pair it with
+# PRIMUS_IDEOGRAM_ZERO1=1 (no-fork hook, zero1.py) to:
+#   - shard the AdamW optimizer state across DP ranks (ZeRO-1), recovering the
+#     memory pure DDP would waste on a replicated optimizer, and
+#   - whole-block checkpoint-wrap the Ideogram layers on the DDP path (the generic
+#     sub-module AC is a no-op for Ideogram blocks; without this it OOMs).
+#
+# CAVEAT: torch.compile is wired ONLY on the FSDP2 path, and single-node exposed
+# comm is a small fraction of the step, so this DDP path (no compile) trails
+# FSDP2+compile on ONE node. It is the multi-node lever / reference point.
+#
+# Launch (inside the container, 8 GPUs free):
+#   PYTHONPATH=<primus>:<primus>/third_party/Automodel \
+#   HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PRIMUS_IDEOGRAM_REAL_AC=1 \
+#   PRIMUS_IDEOGRAM_VARLEN_ATTN=1 PRIMUS_IDEOGRAM_ATTN_ASSUME_DENSE=1 \
+#   PRIMUS_IDEOGRAM_ZERO1=1 \
+#   torchrun --nproc_per_node=8 --nnodes=1 --master_port=29534 \
+#     primus/cli/main.py train pretrain \
+#     --config examples/nemo_automodel/configs/MI355X/diffusion/ideogram4-pretrain-ddp.yaml
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 16
+  local_batch_size: 2
+  ckpt_every_steps: 1000
+  num_epochs: 1000
+  log_remote_every_steps: 1
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: primus.backends.nemo_automodel.models.ideogram4.data.synthetic.SyntheticIdeogram4DataloaderConfig
+    num_samples: 64
+    in_channels: 128
+    grid_h: 16
+    grid_w: 16
+    max_text_tokens: 32
+    min_text_tokens: 28
+    llm_features_dim: 53248
+    feature_scale: 0.1
+    latent_scale: 1.0
+    seed: 1234
+    shuffle: true
+    # Keep the batch size constant: cu_seqlens has 2*B+1 entries, so a short final
+    # batch changes its shape (and hidden_states') and recompiles the graph.
+    drop_last: true
+    num_workers: 0
+    pin_memory: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 1e-4
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  # ZeRO-1 shards these local AdamW states across ranks. fused AdamW is fine
+  # (the per-rank local optimizer runs over its shard of regular tensors).
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: ideogram4
+  adapter_kwargs:
+    in_channels: 128
+    predict_negative_velocity: true
+  timestep_sampling: logit_normal
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 3.0
+  mix_uniform_ratio: 0.1
+  cfg_dropout_prob: 0.0
+  sigma_min: 0.0
+  sigma_max: 1.0
+  num_train_timesteps: 1000
+  i2v_prob: 0.0
+  use_loss_weighting: false
+  log_interval: 100
+  summary_log_interval: 10
+
+# DDP path (replaces the fsdp: block). activation_checkpointing here is honored by
+# the no-fork ideogram_zero1 DDP-AC patch (whole-block wrap) when PRIMUS_IDEOGRAM_ZERO1=1.
+ddp:
+  activation_checkpointing: true
+  gradient_as_bucket_view: true
+  broadcast_buffers: false
+  find_unused_parameters: false
+  static_graph: false
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/ideogram4_primus_pretrain_ddp
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer.yaml b/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer.yaml
new file mode 100644
index 000000000..0f16dd1ef
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer.yaml
@@ -0,0 +1,132 @@
+# NeMo AutoModel pre_trainer module config — Ideogram-4 single-stream
+# rectified-flow T2I pretrain (Primus backend: nemo_automodel).
+#
+# Sibling of `flux_pre_trainer.yaml` (the FLUX T2I preset) and `pre_trainer.yaml`
+# (the Wan2.2 video preset): same recipe schema, wired for Ideogram-4:
+#   - data.dataloader -> SyntheticIdeogram4DataloaderConfig (no-fork Primus loader;
+#     emits image_latents [B,128,H_p,W_p] + llm_features [B,·,53248] + text_lengths).
+#     Swap for the real Flux-2 VAE + Qwen3-VL preprocessor cache loader
+#     (`ideogram4_pre_trainer_pcam.yaml`) to train on real data.
+#   - flow_matching.adapter_type -> ideogram4 (routed to Ideogram4Adapter by the
+#     no-fork `ideogram4_adapter.install()` hook; packs the [pad][text][image]
+#     sequence, maps sigma->t=1-sigma, negates the DiT velocity).
+#
+# The Primus launcher deep-merges the `model:` preset (ideogram4.yaml) and the
+# EXP `overrides:` on top of this, then runs TrainDiffusionRecipe.
+#
+# NOTE: do not declare a `model:` block here — it comes from the model preset.
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 16
+  local_batch_size: 2
+  ckpt_every_steps: 1000
+  num_epochs: 1000
+  log_remote_every_steps: 1
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: primus.backends.nemo_automodel.models.ideogram4.data.synthetic.SyntheticIdeogram4DataloaderConfig
+    # Fixed (deterministic) synthetic cache for the from-config overfit smoke.
+    num_samples: 64
+    in_channels: 128
+    grid_h: 16
+    grid_w: 16
+    max_text_tokens: 32
+    min_text_tokens: 28
+    llm_features_dim: 53248
+    feature_scale: 0.1
+    latent_scale: 1.0
+    seed: 1234
+    shuffle: true
+    # Keep the batch size constant: cu_seqlens has 2*B+1 entries, so a short final
+    # batch changes its shape (and hidden_states') and recompiles the graph.
+    drop_last: true
+    num_workers: 0
+    pin_memory: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 1e-4
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  # fused AdamW (mutually exclusive with foreach).
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: ideogram4
+  adapter_kwargs:
+    in_channels: 128
+    predict_negative_velocity: true
+  timestep_sampling: logit_normal
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 3.0
+  mix_uniform_ratio: 0.1
+  # Overfit smoke: no CFG dropout so every sample keeps its distinct conditioning.
+  cfg_dropout_prob: 0.0
+  sigma_min: 0.0
+  sigma_max: 1.0
+  num_train_timesteps: 1000
+  i2v_prob: 0.0
+  use_loss_weighting: false
+  log_interval: 100
+  summary_log_interval: 10
+
+fsdp:
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  dp_replicate_size: 1
+  # dp_size intentionally unset: AutoModel derives it from world size and the
+  # parallel dims. Pin `dp_size: <n>` via overrides only to force a value.
+  defer_fsdp_grad_sync: true
+  # ZeRO-2: transformer blocks keep their gathered params through backward, so there is
+  # one all-gather per block per step instead of two. Costs roughly one parameter copy
+  # of memory, which is a good trade here given available HBM. Set to `true` for ZeRO-3
+  # (lowest memory, highest traffic), or remove the key entirely to fall back to
+  # AutoModel's per-layer heuristic.
+  reshard_after_forward: false
+  # NOTE: the three prefetch knobs below are INERT while reshard_after_forward is false
+  # -- there is no backward all-gather left to prefetch. They apply again if you set it
+  # back to true.
+  enable_fsdp2_prefetch: true
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
+  reduce_dtype: bfloat16
+  # torch.compile off for the first Ideogram integration run (provisional; flip on
+  # once compile is validated for Ideogram4Transformer2DModel).
+  enable_compile: false
+  # Real activation checkpointing requires PRIMUS_IDEOGRAM_REAL_AC=1 in the launch
+  # env (registers the Ideogram parallelization strategy); otherwise this is a
+  # silent no-op for Ideogram (same as FLUX).
+  activation_checkpointing: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/ideogram4_primus_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer_pcam.yaml b/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer_pcam.yaml
new file mode 100644
index 000000000..be13dd57d
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/ideogram4_pre_trainer_pcam.yaml
@@ -0,0 +1,109 @@
+# NeMo AutoModel pre_trainer module config — Ideogram-4 single-stream
+# rectified-flow T2I pretrain on a REAL pre-encoded cache (Primus backend).
+#
+# Real-data sibling of `ideogram4_pre_trainer.yaml` (the synthetic-cache smoke):
+# identical recipe/optimizer/FSDP wiring, but `data.dataloader` points at the real
+# Flux-2 VAE + Qwen3-VL preprocessor cache produced by `Ideogram4Processor`
+# (`processor.py`) and read by `Ideogram4CacheDataloaderConfig`
+# (emits image_latents [B,128,gh,gw] + LEFT-padded llm_features [B,·,53248] +
+# text_lengths). A dedicated preset (not an override of the synthetic one) so no
+# synthetic-only dataloader keys leak through the deep-merge into the cache dataclass.
+#
+# NOTE: do not declare a `model:` block here — it comes from the model preset.
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 16
+  local_batch_size: 2
+  ckpt_every_steps: 1000
+  num_epochs: 1000
+  log_remote_every_steps: 1
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: primus.backends.nemo_automodel.models.ideogram4.data.cache.Ideogram4CacheDataloaderConfig
+    # Pre-encoded Ideogram-4 cache dir (override per run in the EXP).
+    cache_dir: /dataset/pcam_ideogram4_256
+    shuffle: true
+    drop_last: true
+    # Pad the text region to a constant width so the packed sequence length S does not
+    # change between steps (per-batch padding recompiles the graph). 0 derives it from the
+    # cache's longest caption, so nothing is truncated; -1 restores per-batch padding.
+    max_text_tokens: 0
+    num_workers: 2
+    pin_memory: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 1e-4
+  weight_decay: 0.01
+  betas: [0.9, 0.999]
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: ideogram4
+  adapter_kwargs:
+    in_channels: 128
+    predict_negative_velocity: true
+  timestep_sampling: logit_normal
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 3.0
+  mix_uniform_ratio: 0.1
+  # Real-data loss-decrease validation: no CFG dropout for a clean signal. A real
+  # pretrain should set ~0.1 to train the unconditional branch.
+  cfg_dropout_prob: 0.0
+  sigma_min: 0.0
+  sigma_max: 1.0
+  num_train_timesteps: 1000
+  i2v_prob: 0.0
+  use_loss_weighting: false
+  log_interval: 100
+  summary_log_interval: 10
+
+fsdp:
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  dp_replicate_size: 1
+  defer_fsdp_grad_sync: true
+  # ZeRO-2: see ideogram4_pre_trainer.yaml for the trade-off. `true` gives ZeRO-3;
+  # removing the key falls back to AutoModel's per-layer heuristic.
+  reshard_after_forward: false
+  # NOTE: inert while reshard_after_forward is false (no backward all-gather to prefetch).
+  enable_fsdp2_prefetch: true
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
+  reduce_dtype: bfloat16
+  enable_compile: false
+  # Real AC requires PRIMUS_IDEOGRAM_REAL_AC=1 in the launch env (else no-op).
+  activation_checkpointing: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/ideogram4_pcam_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/nemo_automodel/pre_trainer.yaml b/primus/configs/modules/nemo_automodel/pre_trainer.yaml
new file mode 100644
index 000000000..a2a37c332
--- /dev/null
+++ b/primus/configs/modules/nemo_automodel/pre_trainer.yaml
@@ -0,0 +1,99 @@
+# NeMo AutoModel pre_trainer module config (Primus backend: nemo_automodel).
+#
+# This is the AutoModel recipe schema (model:/fsdp:/optimizer:/step_scheduler:/...);
+# the Primus launcher deep-merges the `model:` preset and the EXP `overrides:`
+# on top of it, then the nemo_automodel backend materialises the whole tree into
+# an AutoModel ConfigNode and runs TrainDiffusionRecipe. The defaults below are
+# the optimized pretrain knobs: aiter + torch.compile + bf16 reduce + fused AdamW.
+#
+# NOTE: do not declare a `model:` block here — it comes from the model preset
+# (merge is allow_override=False and would conflict).
+extends:
+  - ../module_base.yaml
+
+seed: 42
+
+wandb:
+  enable: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+step_scheduler:
+  global_batch_size: 256
+  local_batch_size: 32
+  ckpt_every_steps: 1000
+  num_epochs: 100
+  log_remote_every_steps: 2
+  save_checkpoint_every_epoch: false
+
+data:
+  dataloader:
+    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
+    cache_dir: /dataset
+    model_type: wan
+    base_resolution: [512, 512]
+    dynamic_batch_size: false
+    shuffle: true
+    drop_last: false
+    num_workers: 2
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 5e-5
+  weight_decay: 0.1
+  betas: [0.9, 0.95]
+  # fused AdamW (mutually exclusive with foreach).
+  foreach: false
+  fused: true
+
+performance:
+  check_loss: false
+  grad_clip_foreach: true
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 0
+  min_lr: 1e-6
+
+flow_matching:
+  adapter_type: simple
+  adapter_kwargs: {}
+  timestep_sampling: uniform
+  logit_mean: 0.0
+  logit_std: 1.0
+  flow_shift: 3.0
+  mix_uniform_ratio: 0.1
+  num_train_timesteps: 1000
+  i2v_prob: 0.3
+  use_loss_weighting: true
+  log_interval: 100
+  summary_log_interval: 10
+
+fsdp:
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  dp_replicate_size: 1
+  # dp_size is intentionally left unset: AutoModel derives it from the runtime
+  # world size and the parallel dims (dp = world_size / (tp*cp*pp)), so single-
+  # and multi-node runs work without edits. Pin `dp_size: <n>` only to force a
+  # specific value.
+  defer_fsdp_grad_sync: true
+  enable_fsdp2_prefetch: true
+  # Optimized knobs:
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
+  reduce_dtype: bfloat16
+  enable_compile: true
+  # Keep AC on for bs>8; flip false via overrides for the bs6-8 AC-off point.
+  activation_checkpointing: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: /checkpoints/wan22_primus_pretrain
+  model_save_format: safetensors
+  save_consolidated: final
+  diffusers_compatible: true
+  restore_from: null
diff --git a/primus/configs/modules/torchtitan/pre_trainer.yaml b/primus/configs/modules/torchtitan/pre_trainer.yaml
index 772ed4cdd..23321922d 100644
--- a/primus/configs/modules/torchtitan/pre_trainer.yaml
+++ b/primus/configs/modules/torchtitan/pre_trainer.yaml
@@ -35,10 +35,8 @@ profiling:
 
 training:
   mock_data: true
-  debug_moe_force_load_balance: false
   dataset: c4
   dataset_path: null
-  deterministic: false
   enable_cpu_offload: false
   gc_debug: false
   gc_freq: 50
@@ -47,10 +45,14 @@ training:
   max_norm: 1.0
   mixed_precision_param: bfloat16
   mixed_precision_reduce: float32
-  seed: null
   seq_len: 2048
   steps: 10000
 
+debug:
+  deterministic: false
+  deterministic_warn_only: false
+  moe_force_load_balance: false
+  seed: null
 
 parallelism:
   context_parallel_degree: 1
@@ -59,7 +61,6 @@ parallelism:
   data_parallel_shard_degree: -1
   disable_loss_parallel: false
   enable_async_tensor_parallel: false
-  enable_compiled_autograd: false
   expert_parallel_degree: 1
   expert_tensor_parallel_degree: 1
   fsdp_reshard_after_forward: default
@@ -71,7 +72,6 @@ parallelism:
   pipeline_parallel_microbatch_size: 1
   pipeline_parallel_schedule: 1F1B
   pipeline_parallel_schedule_csv: ''
-  pipeline_parallel_split_points: []
   tensor_parallel_degree: 1
 
 checkpoint:
diff --git a/primus/core/backend/backend_adapter.py b/primus/core/backend/backend_adapter.py
index ce58505bc..d0aeb16bf 100644
--- a/primus/core/backend/backend_adapter.py
+++ b/primus/core/backend/backend_adapter.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -40,7 +40,7 @@ def prepare_backend(self, config: Any):
         beyond `BackendRegistry.run_setup(self.framework)`.
         """
         from primus.core.backend.backend_registry import BackendRegistry
-        from primus.modules.module_utils import log_rank_0
+        from primus.core.utils.module_utils import log_rank_0
 
         BackendRegistry.run_setup(self.framework)
         log_rank_0(f"[Primus:{self.framework}] Backend prepared")
@@ -57,7 +57,7 @@ def setup_backend_path(self, backend_path=None) -> str:
         import sys
         from pathlib import Path
 
-        from primus.modules.module_utils import log_rank_0
+        from primus.core.utils.module_utils import log_rank_0
 
         def _use_path(path: str, error_msg: str) -> str:
             norm_path = os.path.abspath(os.path.normpath(str(path)))
@@ -117,12 +117,18 @@ def _use_path(path: str, error_msg: str) -> str:
     # ============================================================================
 
     @abstractmethod
-    def load_trainer_class(self, stage: str = "pretrain"):
+    def load_trainer_class(self, stage: str = "pretrain", trainer_class: Optional[str] = None):
         """
         Return backend Trainer class registered in `BackendRegistry`.
 
+        Args:
+            stage: Training stage (e.g., "pretrain", "sft"). Defaults to "pretrain".
+            trainer_class: Optional specific trainer class name for dynamic loading.
+                          If provided, this takes precedence over stage-based selection.
+
         Default behavior:
-          - Lookup trainer via `BackendRegistry.get_trainer_class(self.framework, stage=stage)`
+          - If `trainer_class` is provided, attempt to dynamically load it
+          - Otherwise, lookup trainer via `BackendRegistry.get_trainer_class(self.framework, stage=stage)`
 
         Backends can override this method if they need special resolution rules.
         """
diff --git a/primus/core/backend/backend_registry.py b/primus/core/backend/backend_registry.py
index 1f78d3f5b..d198c7073 100644
--- a/primus/core/backend/backend_registry.py
+++ b/primus/core/backend/backend_registry.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -20,7 +20,7 @@
 
 from typing import Callable, Dict, List, Type
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class BackendRegistry:
diff --git a/primus/modules/base_module.py b/primus/core/base_module.py
similarity index 74%
rename from primus/modules/base_module.py
rename to primus/core/base_module.py
index 615ba8627..8c5283cb1 100644
--- a/primus/modules/base_module.py
+++ b/primus/core/base_module.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -15,8 +15,7 @@
     get_target_platform,
     set_global_variables,
 )
-
-from .module_utils import debug_rank_all, set_logging_rank
+from primus.core.utils.module_utils import debug_rank_all, set_logging_rank
 
 
 class BaseModule(ABC):
@@ -28,8 +27,8 @@ def __init__(
         module_world_size: int,
         module_master_addr: str = None,
         module_master_port: int = None,
+        **kwargs,  # Accept extra kwargs for cooperative multiple inheritance
     ):
-        # module will be initialized by multiple worker processes
         self.module_name = module_name
 
         assert primus_config is not None
@@ -71,6 +70,25 @@ def __init__(
         # setup logger for worker
         self.setup_worker_logger(module_rank, module_world_size)
 
+        # Call super() for cooperative multiple inheritance
+        # MRO for MegatronTrainer: MegatronTrainer -> BaseTrainer -> BaseModule -> ABC -> object
+        #
+        # Flow:
+        # 1. MegatronTrainer.__init__() calls super().__init__() -> BaseTrainer.__init__()
+        # 2. BaseTrainer.__init__() consumes backend_args, then calls super().__init__(**kwargs) -> BaseModule.__init__()
+        #    Note: BaseTrainer passes kwargs containing module_name, primus_config, etc. to BaseModule
+        # 3. BaseModule.__init__() receives module_name, primus_config, etc. as positional/keyword args
+        #    The kwargs dict may still contain these keys, but BaseModule has already consumed them as positional args
+        # 4. BaseModule.__init__() calls super().__init__() -> ABC.__init__() -> object.__init__()
+        #
+        # The issue: If kwargs still contains keys, passing them to super() will eventually reach object.__init__()
+        # which doesn't accept arguments. We must ensure kwargs is empty or only contains args for the next class.
+        # Since the next class after BaseModule is ABC (which doesn't accept kwargs), we should not pass kwargs.
+
+        # Don't pass kwargs to super() - the next class in MRO is ABC, which doesn't need them
+        # All required args for BaseModule have been consumed as positional/keyword arguments above
+        super().__init__()
+
     @abstractmethod
     def init(self, *args, **kwargs):
         raise NotImplementedError
diff --git a/primus/core/config/primus_config.py b/primus/core/config/primus_config.py
index a9886fa1a..720a3316e 100644
--- a/primus/core/config/primus_config.py
+++ b/primus/core/config/primus_config.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -53,7 +53,7 @@ def _normalize_module_for_runtime(module_cfg: SimpleNamespace, module_name: str)
     if not getattr(normalized, "name", None):
         setattr(normalized, "name", module_name)
 
-    reserved_keys = {"name", "framework", "config", "model", "params"}
+    reserved_keys = {"name", "framework", "config", "model", "params", "trainer_class"}
     # Start from any existing params dict/namespace if provided.
     existing_params = getattr(normalized, "params", {})
     params = _to_plain_dict(existing_params)
@@ -121,11 +121,26 @@ def load_primus_config(config_path: Path, cli_args: Any | None = None) -> Simple
     cfg.platform = platform_config
 
     # Build modules list from legacy PrimusConfig.module_keys/get_module_config.
-    cfg.modules = [
+    #
+    # `_normalize_module_for_runtime` deepcopies each module namespace before
+    # reshaping it, so `legacy_cfg` itself is left pristine. This is what makes
+    # it safe to expose `legacy_cfg` via `cfg._legacy` below: the legacy
+    # `PrimusConfig` that downstream consumers (e.g. `BaseModule`) read through
+    # `get_module_config(...)` keeps its original module shape.
+    modules: list[SimpleNamespace] = [
         _normalize_module_for_runtime(legacy_cfg.get_module_config(module_name), module_name)
         for module_name in getattr(legacy_cfg, "module_keys", [])
     ]
 
+    cfg.modules = modules
+
+    # Expose the underlying legacy PrimusConfig (additive, non-breaking) so the
+    # core runtime can reuse it instead of re-parsing the YAML a second time.
+    # BaseModule (inherited) still needs the PrimusConfig class interface
+    # (get_module_config method, set_global_variables, exp_* properties) that the
+    # SimpleNamespace above does not provide. `legacy_cfg` is pristine (see above).
+    cfg._legacy = legacy_cfg
+
     return cfg
 
 
diff --git a/primus/core/config/schema_check.py b/primus/core/config/schema_check.py
new file mode 100644
index 000000000..2efc97046
--- /dev/null
+++ b/primus/core/config/schema_check.py
@@ -0,0 +1,1199 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Detect Primus YAML keys that no longer exist in the upstream backend schema.
+
+Every adapter takes unknown keys silently -- TorchTitan attaches them as dynamic attributes,
+MaxText lets pydantic drop them, Megatron gets them merged onto ``args`` by ``train_runtime`` --
+so a renamed field keeps parsing while doing nothing: ``training.debug_moe_force_load_balance``
+moved to ``debug.moe_force_load_balance``, and every DeepSeek config went on training with MoE
+load balancing off. Schemas are read statically (AST, ``yaml.safe_load``), so the check imports
+neither torch nor jax and runs on a plain CI runner.
+
+Legal keys are upstream's plus the Primus extensions, and the extensions are derived rather than
+listed: ``get_param(ctx, "<path>", ...)`` calls in the patch packages, reads off Megatron's
+``args`` namespace, Primus dataclasses extending a backend config class, and a hand-written
+allowlist. ``MegatronArgBuilder`` stripping a key is not drift: ``train_runtime`` merges the
+original params back over the result, which is how roughly 130 Primus-only keys reach ``args``.
+
+A legal key can still be dead where it is written: ``core_transformer_config_from_args`` copies a
+Primus config class's fields off ``args`` only on the model naming that class. ``norm_epsilon``
+is DeepSeek-V4's and not an upstream dest -- upstream's ``layernorm_epsilon`` merely spells its
+flag ``--norm-epsilon`` -- so it is live there and inert on a Llama config. Those are reported as
+misplaced rather than unknown, and a config whose model cannot be resolved is left alone.
+"""
+
+from __future__ import annotations
+
+import ast
+import fnmatch
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, Iterator, Mapping, Sequence
+
+import yaml
+
+from primus.core.config.yaml_loader import parse_yaml
+
+BACKENDS = ("torchtitan", "maxtext", "megatron")
+
+TORCHTITAN_JOB_CONFIG = Path("third_party/torchtitan/torchtitan/config/job_config.py")
+MAXTEXT_BASE_YML = Path("third_party/maxtext/src/maxtext/configs/base.yml")
+MEGATRON_ROOT = Path("third_party/Megatron-LM")
+MEGATRON_ARGUMENTS = Path("megatron/training/arguments.py")
+
+# `core_transformer_config_from_args` copies every field of these classes off
+# `args`, so their fields stay legal even where ArgumentGroupFactory excludes
+# them from argparse (`moe_token_dropping`, `fp8_multi_head_attention`, ...).
+MEGATRON_CONFIG_CLASSES = ("TransformerConfig", "MLATransformerConfig")
+
+# Module preset whose `experimental.custom_args_module` points at the dataclass
+# that extends TorchTitan's JobConfig.
+TORCHTITAN_MODULE_PRESET = Path("primus/configs/modules/torchtitan/pre_trainer.yaml")
+
+PATCH_DIRS = {
+    "torchtitan": Path("primus/backends/torchtitan/patches"),
+    "maxtext": Path("primus/backends/maxtext/patches"),
+}
+
+# Packages whose reads off the backend's argument namespace define the
+# Primus-only half of its key set, and which hold the Primus subclasses of the
+# backend's config dataclass.
+BACKEND_PACKAGES = {"megatron": Path("primus/backends/megatron")}
+
+# `get_model_provider(model_type=...)` is where a Megatron run decides which
+# model it builds, and the branches that name a Primus module are the ones that
+# lead to a Primus config class. See `extract_model_scopes`.
+MODEL_PROVIDERS = {
+    "megatron": (Path("primus/core/utils/import_utils.py"), "get_model_provider", "model_type"),
+}
+
+# `module.model` names a preset here, and that preset is where a model declares
+# its `model_type`; the experiment YAML that sets a model-scoped key usually
+# does not repeat it.
+MODEL_PRESETS = Path("primus/configs/models")
+
+# Every Python file Primus owns. A name read off an argument namespace anywhere
+# in here is model-agnostic evidence, which is what keeps a shared key such as
+# `moe_use_legacy_grouped_gemm` out of the model-scoped key sets.
+PRIMUS_PACKAGE = Path("primus")
+
+SCAN_ROOTS = {
+    "torchtitan": (
+        Path("examples/torchtitan/configs"),
+        Path("primus/configs/modules/torchtitan"),
+        Path("primus/configs/models/torchtitan"),
+    ),
+    "maxtext": (
+        Path("examples/maxtext/configs"),
+        Path("primus/configs/modules/maxtext"),
+        Path("primus/configs/models/maxtext"),
+    ),
+    "megatron": (
+        Path("examples/megatron/configs"),
+        Path("primus/configs/modules/megatron"),
+        Path("primus/configs/models/megatron"),
+    ),
+}
+
+DEFAULT_ALLOWLIST = Path("tools/ci/config_schema_allowlist.yaml")
+
+# Keys PrimusParser.parse_trainer_module consumes on the module node itself;
+# they never reach the backend.
+MODULE_RESERVED_KEYS = frozenset({"name", "framework", "config", "model", "overrides", "params"})
+
+
+# ---------------------------------------------------------------------------
+# Schema
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class BackendSchema:
+    """Upstream key universe for one backend."""
+
+    backend: str
+    keys: frozenset[str]
+    sections: frozenset[str]
+    source: str
+
+    def extended(self, keys: Iterable[str], sections: Iterable[str], source: str) -> "BackendSchema":
+        return BackendSchema(
+            backend=self.backend,
+            keys=self.keys | frozenset(keys),
+            sections=self.sections | frozenset(sections),
+            source=f"{self.source} + {source}",
+        )
+
+    def suggest(self, path: str) -> str | None:
+        """Best guess for where a dropped key moved to: the leaf name, the leaf minus a leading
+        ``<section>_`` (moved into a new section), or the leaf plus a prefix or suffix (renamed
+        in place). Ambiguity yields ``None``."""
+        leaf = path.rsplit(".", 1)[-1]
+        candidates = {leaf}
+        head = leaf.split("_", 1)
+        if len(head) == 2 and head[0] in {s.split(".")[0] for s in self.sections | self.keys}:
+            candidates.add(head[1])
+        matches = sorted(k for k in self.keys if k != path and k.rsplit(".", 1)[-1] in candidates)
+        if not matches:
+            matches = sorted(
+                k
+                for k in self.keys
+                if (tail := k.rsplit(".", 1)[-1]) != leaf
+                and (tail.startswith(f"{leaf}_") or tail.endswith(f"_{leaf}"))
+            )
+        return matches[0] if len(matches) == 1 else None
+
+
+def _dataclass_defs(tree: ast.Module) -> dict[str, list[tuple[str, str]]]:
+    """Map ``@dataclass`` class name -> ``[(field name, annotation source)]``."""
+    defs: dict[str, list[tuple[str, str]]] = {}
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        decorated = any("dataclass" in ast.unparse(d) for d in node.decorator_list)
+        if not decorated:
+            continue
+        fields = [
+            (stmt.target.id, ast.unparse(stmt.annotation).strip("\"'"))
+            for stmt in node.body
+            if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name)
+        ]
+        defs[node.name] = fields
+    return defs
+
+
+def _expand_dataclass(
+    defs: dict[str, list[tuple[str, str]]],
+    cls: str,
+    prefix: str,
+    keys: set[str],
+    sections: set[str],
+    seen: frozenset[str] = frozenset(),
+) -> None:
+    for name, annotation in defs.get(cls, ()):
+        path = f"{prefix}.{name}" if prefix else name
+        if annotation in defs and annotation not in seen:
+            sections.add(path)
+            _expand_dataclass(defs, annotation, path, keys, sections, seen | {annotation})
+        else:
+            keys.add(path)
+
+
+def load_torchtitan_schema(repo_root: Path) -> BackendSchema | None:
+    """Extract ``JobConfig`` from TorchTitan's ``job_config.py`` via AST; the module imports
+    torch, so it cannot simply be imported here."""
+    src = repo_root / TORCHTITAN_JOB_CONFIG
+    if not src.is_file():
+        return None
+    defs = _dataclass_defs(ast.parse(src.read_text()))
+    if "JobConfig" not in defs:
+        raise ValueError(f"no JobConfig dataclass found in {TORCHTITAN_JOB_CONFIG}")
+    keys: set[str] = set()
+    sections: set[str] = set()
+    _expand_dataclass(defs, "JobConfig", "", keys, sections)
+    return BackendSchema(
+        backend="torchtitan",
+        keys=frozenset(keys),
+        sections=frozenset(sections),
+        source=str(TORCHTITAN_JOB_CONFIG),
+    )
+
+
+def load_maxtext_schema(repo_root: Path) -> BackendSchema | None:
+    """MaxText's ``base.yml`` is a flat mapping and is itself the schema."""
+    src = repo_root / MAXTEXT_BASE_YML
+    if not src.is_file():
+        return None
+    data = yaml.safe_load(src.read_text()) or {}
+    return BackendSchema(
+        backend="maxtext",
+        keys=frozenset(data),
+        sections=frozenset(),
+        source=str(MAXTEXT_BASE_YML),
+    )
+
+
+def _calls_named(tree: ast.AST, name: str) -> Iterator[ast.Call]:
+    """Every call whose callee is spelled ``name`` or ``<something>.name``."""
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+        func = node.func
+        called = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", None)
+        if called == name:
+            yield node
+
+
+def _string_elements(node: ast.expr | None) -> set[str] | None:
+    """String constants of a literal list/tuple/set; ``None`` if not one."""
+    if not isinstance(node, (ast.List, ast.Tuple, ast.Set)):
+        return None
+    return {e.value for e in node.elts if isinstance(e, ast.Constant) and isinstance(e.value, str)}
+
+
+def _argparse_dest(call: ast.Call) -> str | None:
+    """The attribute name argparse derives from one ``add_argument`` call: ``dest=`` if given, else
+    the first long option, so ``--no-persist-layer-norm`` gives ``no_persist_layer_norm``."""
+    for keyword in call.keywords:
+        if keyword.arg == "dest" and isinstance(keyword.value, ast.Constant):
+            return str(keyword.value.value)
+    options = [a.value for a in call.args if isinstance(a, ast.Constant) and isinstance(a.value, str)]
+    chosen = next((o for o in options if o.startswith("--")), None) or next(iter(options), None)
+    return chosen.lstrip("-").replace("-", "_") if chosen else None
+
+
+class _MegatronSource:
+    """Just enough of an import resolver to walk the checked-out Megatron tree: config dataclasses
+    arrive through re-exports and inherit across packages, so a flat class index is not faithful."""
+
+    def __init__(self, root: Path) -> None:
+        self.root = root
+        self._trees: dict[Path, ast.Module] = {}
+
+    def tree(self, path: Path) -> ast.Module:
+        if path not in self._trees:
+            self._trees[path] = ast.parse(path.read_text())
+        return self._trees[path]
+
+    def module_file(self, module: str) -> Path | None:
+        rel = module.replace(".", "/")
+        for candidate in (self.root / f"{rel}.py", self.root / rel / "__init__.py"):
+            if candidate.is_file():
+                return candidate
+        return None
+
+    def imports(self, path: Path) -> dict[str, str]:
+        """``local name -> absolute module`` for every ``from X import name``."""
+        package = ".".join(path.relative_to(self.root).parts[:-1])
+        found: dict[str, str] = {}
+        for node in ast.walk(self.tree(path)):
+            if not isinstance(node, ast.ImportFrom):
+                continue
+            module = node.module or ""
+            if node.level:  # `from ..model_parallel_config import ModelParallelConfig`
+                parts = package.split(".")
+                parts = parts[: len(parts) - node.level + 1]
+                module = ".".join([*parts, module]) if module else ".".join(parts)
+            for alias in node.names:
+                found[alias.asname or alias.name] = module
+        return found
+
+    def resolve(
+        self, path: Path, name: str, seen: frozenset[str] = frozenset()
+    ) -> tuple[Path, ast.ClassDef] | None:
+        """Find class ``name`` as visible from ``path``, with its defining file: base classes must
+        be resolved in the scope that declares them, not the one that referenced them."""
+        for node in self.tree(path).body:
+            if isinstance(node, ast.ClassDef) and node.name == name:
+                return path, node
+        module = self.imports(path).get(name)
+        if module is None or module in seen:
+            return None
+        target = self.module_file(module)
+        return self.resolve(target, name, seen | {module}) if target else None
+
+    def config_fields(self, path: Path, cls: ast.ClassDef, exclude: Iterable[str] = ()) -> list[str]:
+        """Names ``ArgumentGroupFactory`` would turn into arguments, in field order, mirroring
+        ``build_group``: base-class fields first, ``init=False`` and ``ClassVar`` skipped, an
+        ``argparse_meta`` ``dest`` override honoured (``use_nsys_profiler`` -> ``profile``)."""
+        skip = set(exclude)
+        names: list[str] = []
+        for base in cls.bases:
+            if not isinstance(base, ast.Name):
+                continue
+            found = self.resolve(path, base.id)
+            if found:
+                names.extend(self.config_fields(*found, exclude=skip))
+        for stmt in cls.body:
+            if not isinstance(stmt, ast.AnnAssign) or not isinstance(stmt.target, ast.Name):
+                continue
+            if ast.unparse(stmt.annotation).lstrip("\"'").startswith(("ClassVar", "InitVar")):
+                continue
+            if _has_field_kwarg(stmt.value, "init", False):
+                continue
+            dest = _argparse_meta_dest(stmt.value) or stmt.target.id
+            if dest not in skip and dest not in names:
+                names.append(dest)
+        return names
+
+
+def _field_call(value: ast.expr | None) -> ast.Call | None:
+    """The ``dataclasses.field(...)`` call of a field default, if any."""
+    if not isinstance(value, ast.Call):
+        return None
+    func = value.func
+    called = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", None)
+    return value if called == "field" else None
+
+
+def _has_field_kwarg(value: ast.expr | None, name: str, expected: Any) -> bool:
+    call = _field_call(value)
+    if call is None:
+        return False
+    return any(
+        kw.arg == name and isinstance(kw.value, ast.Constant) and kw.value.value is expected
+        for kw in call.keywords
+    )
+
+
+def _argparse_meta_dest(value: ast.expr | None) -> str | None:
+    """``metadata={"argparse_meta": {"dest": "profile"}}`` -> ``"profile"``."""
+    call = _field_call(value)
+    if call is None:
+        return None
+    for keyword in call.keywords:
+        if keyword.arg != "metadata" or not isinstance(keyword.value, ast.Dict):
+            continue
+        for outer_key, outer in zip(keyword.value.keys, keyword.value.values):
+            if not (isinstance(outer_key, ast.Constant) and outer_key.value == "argparse_meta"):
+                continue
+            if not isinstance(outer, ast.Dict):
+                continue
+            for inner_key, inner in zip(outer.keys, outer.values):
+                if isinstance(inner_key, ast.Constant) and inner_key.value == "dest":
+                    if isinstance(inner, ast.Constant):
+                        return str(inner.value)
+    return None
+
+
+def _resolve_exclude(tree: ast.Module, call: ast.Call) -> set[str]:
+    """The ``exclude=`` list of one ``ArgumentGroupFactory`` call, following a literal assignment
+    when the argument is a variable, as ``TransformerConfig``'s call passes. An expression that
+    cannot be read statically yields no exclusions, erring towards a missed report."""
+    argument = next((kw.value for kw in call.keywords if kw.arg == "exclude"), None)
+    literal = _string_elements(argument)
+    if literal is not None:
+        return literal
+    if not isinstance(argument, ast.Name):
+        return set()
+    names: set[str] = set()
+    for node in ast.walk(tree):
+        targets = node.targets if isinstance(node, ast.Assign) else []
+        if any(isinstance(t, ast.Name) and t.id == argument.id for t in targets):
+            names |= _string_elements(node.value) or set()
+    return names
+
+
+def load_megatron_schema(repo_root: Path) -> BackendSchema | None:
+    """Rebuild Megatron's flat argument namespace from ``arguments.py``: roughly 355 literal
+    ``add_argument`` calls, plus eleven ``ArgumentGroupFactory(<dataclass>)`` calls generating as
+    many again. The generated half holds ``train_iters``, ``global_batch_size`` and
+    ``micro_batch_size``, so the literal calls alone would report half of every config as drift."""
+    root = repo_root / MEGATRON_ROOT
+    src = root / MEGATRON_ARGUMENTS
+    if not src.is_file():
+        return None
+    source = _MegatronSource(root)
+    tree = source.tree(src)
+
+    keys = {dest for call in _calls_named(tree, "add_argument") if (dest := _argparse_dest(call))}
+    for call in _calls_named(tree, "ArgumentGroupFactory"):
+        target = call.args[0] if call.args else None
+        if not isinstance(target, ast.Name):
+            continue
+        found = source.resolve(src, target.id)
+        if found is None:
+            raise ValueError(f"{MEGATRON_ARGUMENTS}: cannot resolve ArgumentGroupFactory({target.id})")
+        keys.update(source.config_fields(*found, exclude=_resolve_exclude(tree, call)))
+    for name in MEGATRON_CONFIG_CLASSES:
+        found = source.resolve(src, name)
+        if found is not None:
+            keys.update(source.config_fields(*found))
+
+    return BackendSchema(
+        backend="megatron",
+        keys=frozenset(keys),
+        sections=frozenset(),  # Megatron's namespace is flat: every key is a leaf.
+        source=str(MEGATRON_ROOT / MEGATRON_ARGUMENTS),
+    )
+
+
+SCHEMA_LOADERS = {
+    "torchtitan": load_torchtitan_schema,
+    "maxtext": load_maxtext_schema,
+    "megatron": load_megatron_schema,
+}
+
+
+# ---------------------------------------------------------------------------
+# Allowlist
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class AllowRule:
+    pattern: str
+    reason: str
+    consumer: str
+    origin: str
+
+    def matches(self, path: str) -> bool:
+        return path == self.pattern or fnmatch.fnmatchcase(path, self.pattern)
+
+
+@dataclass(frozen=True)
+class Allowlist:
+    """Rules that widen a backend's legal key set, indexed for lookup: Megatron derives thousands,
+    so exact patterns are hashed and only real fnmatch patterns scanned, and exact wins over glob."""
+
+    rules: tuple[AllowRule, ...]
+
+    def __post_init__(self) -> None:
+        exact: dict[str, AllowRule] = {}
+        globs: list[AllowRule] = []
+        prefixes: set[str] = set()
+        for rule in self.rules:
+            if any(ch in rule.pattern for ch in "*?["):
+                globs.append(rule)
+            else:
+                exact.setdefault(rule.pattern, rule)
+            parts = rule.pattern.split(".")
+            prefixes.update(".".join(parts[:i]) for i in range(1, len(parts)))
+        object.__setattr__(self, "_exact", exact)
+        object.__setattr__(self, "_globs", tuple(globs))
+        object.__setattr__(self, "_prefixes", frozenset(prefixes))
+
+    def match(self, path: str) -> AllowRule | None:
+        rule = self._exact.get(path)
+        return rule or next((r for r in self._globs if r.matches(path)), None)
+
+    def is_prefix(self, path: str) -> bool:
+        """True when some rule targets a child of ``path`` (so ``path`` is a section)."""
+        return path in self._prefixes
+
+    def patterns(self) -> list[str]:
+        return sorted({r.pattern for r in self.rules})
+
+
+def extract_get_param_rules(repo_root: Path, backend: str) -> list[AllowRule]:
+    """Collect config paths read by ``get_param(ctx, "<path>", ...)`` in patches."""
+    patch_dir = repo_root / PATCH_DIRS.get(backend, Path("does/not/exist"))
+    if not patch_dir.is_dir():
+        return []
+    rules: list[AllowRule] = []
+    for py in sorted(patch_dir.rglob("*.py")):
+        tree = ast.parse(py.read_text())
+        rel = py.relative_to(repo_root)
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.Call) or len(node.args) < 2:
+                continue
+            func = node.func
+            name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", None)
+            if name != "get_param":
+                continue
+            arg = node.args[1]
+            if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                rules.append(
+                    AllowRule(
+                        pattern=arg.value,
+                        reason="read by the Primus patch system",
+                        consumer=f"{rel}:{node.lineno}",
+                        origin="patch:get_param",
+                    )
+                )
+    return rules
+
+
+_ARGS_OBJECT_HINTS = ("arg", "cfg", "config")
+
+
+def _is_args_object(node: ast.expr) -> bool:
+    """True when ``node`` looks like Megatron's argument or config namespace (``args.x``,
+    ``self.args.x``, ``get_args().x``), but not every attribute read: ``tensor.shape`` is not one."""
+    while isinstance(node, ast.Call):
+        node = node.func
+    name = node.attr if isinstance(node, ast.Attribute) else getattr(node, "id", None)
+    return bool(name) and any(hint in name.lower() for hint in _ARGS_OBJECT_HINTS)
+
+
+def extract_args_read_rules(repo_root: Path, backend: str) -> list[AllowRule]:
+    """Collect the names a backend package reads off the argument namespace. Two forms count as a
+    reader: attribute access on an args-shaped object, and the bare name as a string literal
+    anywhere in the package -- reads reach ``args`` through too many indirections (module
+    constants, local getters, config-to-env bridges, generated source) for a narrower rule."""
+    package = repo_root / BACKEND_PACKAGES.get(backend, Path("does/not/exist"))
+    if not package.is_dir():
+        return []
+    rules: list[AllowRule] = []
+    seen: set[str] = set()
+
+    def record(name: str, path: Path, lineno: int, origin: str) -> None:
+        # An ALL_CAPS literal is a constant, an enum member or an environment variable, never an
+        # argument name -- and a config setting one as if it were (`HSA_NO_SCRATCH_RECLAIM: 1`)
+        # is exactly the mistake this check should keep reporting.
+        if name in seen or not name.isidentifier() or name == name.upper():
+            return
+        seen.add(name)
+        rules.append(
+            AllowRule(
+                pattern=name,
+                reason=f"read off the {backend} args namespace by the backend package",
+                consumer=f"{path.relative_to(repo_root)}:{lineno}",
+                origin=origin,
+            )
+        )
+
+    for py in sorted(package.rglob("*.py")):
+        for node in ast.walk(ast.parse(py.read_text())):
+            if isinstance(node, ast.Attribute) and _is_args_object(node.value):
+                record(node.attr, py, node.lineno, "backend:args_attribute")
+            elif isinstance(node, ast.Constant) and isinstance(node.value, str):
+                record(node.value, py, node.lineno, "backend:args_name")
+    return rules
+
+
+def _class_bases(node: ast.ClassDef) -> set[str]:
+    return {ast.unparse(base).rsplit(".", 1)[-1] for base in node.bases}
+
+
+def _class_fields(node: ast.ClassDef) -> set[str]:
+    return {
+        stmt.target.id
+        for stmt in node.body
+        if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name)
+    }
+
+
+def _primus_config_classes(repo_root: Path, backend: str) -> tuple[dict[str, ast.ClassDef], frozenset[str]]:
+    """``(every class in the backend package, those extending its config class)``, grown to a
+    fixpoint because the chain can be several classes long: ``FluxConfig`` ->
+    ``BaseDiffusionConfig`` -> ``TransformerConfig``."""
+    package = repo_root / BACKEND_PACKAGES.get(backend, Path("does/not/exist"))
+    if not package.is_dir():
+        return {}, frozenset()
+    classes: dict[str, ast.ClassDef] = {}
+    for py in sorted(package.rglob("*.py")):
+        for node in ast.walk(ast.parse(py.read_text())):
+            if isinstance(node, ast.ClassDef):
+                classes[node.name] = node
+
+    derived: set[str] = set()
+    frontier = set(MEGATRON_CONFIG_CLASSES)
+    while frontier:
+        frontier = {
+            n for n, c in classes.items() if n not in derived and _class_bases(c) & (frontier | derived)
+        }
+        derived |= frontier
+    return classes, frozenset(derived)
+
+
+def extract_backend_config_extensions(repo_root: Path, backend: str) -> tuple[frozenset[str], str]:
+    """Fields of the Primus dataclasses that extend the backend's config class, which
+    ``core_transformer_config_from_args`` copies off ``args`` and thereby makes legal keys."""
+    classes, derived = _primus_config_classes(repo_root, backend)
+    if not derived:
+        return frozenset(), ""
+    keys = frozenset().union(*(_class_fields(classes[name]) for name in derived))
+    return keys, str(BACKEND_PACKAGES[backend])
+
+
+# ---------------------------------------------------------------------------
+# Model scopes
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class ModelScope:
+    """Keys that only reach the model one YAML discriminator selects: live on that model's path,
+    a silent no-op on every other."""
+
+    key: str
+    """YAML key that decides which model gets built, e.g. ``model_type``."""
+
+    value: str
+    """The value of ``key`` that selects this model, e.g. ``deepseek_v4``."""
+
+    default: str | None
+    """What ``key`` means when a config omits it; ``None`` when there is no default."""
+
+    config_class: str
+    keys: frozenset[str]
+    source: str
+
+    def describe(self) -> str:
+        return f"{self.key}: {self.value}"
+
+
+def _is_args_namespace(node: ast.expr) -> bool:
+    """True for ``args``, ``self.args``, ``get_args()``, ``backend_args``, ..."""
+    while isinstance(node, ast.Call):
+        node = node.func
+    name = node.attr if isinstance(node, ast.Attribute) else getattr(node, "id", None)
+    return bool(name) and "arg" in name.lower()
+
+
+def _args_namespace_reads(repo_root: Path) -> frozenset[str]:
+    """Names read off an *argument* namespace anywhere in ``primus/``, deliberately narrower than
+    ``_is_args_object``: ``config.norm_epsilon`` only says the key is a field of whichever config
+    class got built, while ``args.moe_use_legacy_grouped_gemm`` says it reaches its reader on
+    every model. Only the latter disqualifies a key from being model-scoped."""
+    package = repo_root / PRIMUS_PACKAGE
+    if not package.is_dir():
+        return frozenset()
+    names: set[str] = set()
+    for py in sorted(package.rglob("*.py")):
+        try:
+            tree = ast.parse(py.read_text())
+        except (OSError, SyntaxError):
+            continue
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Attribute) and _is_args_namespace(node.value):
+                names.add(node.attr)
+            elif isinstance(node, ast.Call) and len(node.args) >= 2:
+                called = getattr(node.func, "id", None)
+                target = node.args[1]
+                if called in ("getattr", "setattr", "hasattr") and _is_args_namespace(node.args[0]):
+                    if isinstance(target, ast.Constant) and isinstance(target.value, str):
+                        names.add(target.value)
+    return frozenset(names)
+
+
+def _model_provider_branches(repo_root: Path, backend: str) -> tuple[str | None, dict[str, str], str]:
+    """Read the model dispatch -- ``(default, {discriminator value: module}, source)`` -- from
+    ``get_model_provider``'s branches; only one naming a Primus module reaches a Primus config."""
+    src, func_name, key = MODEL_PROVIDERS.get(backend, (None, "", ""))
+    if src is None:
+        return None, {}, ""
+    path = repo_root / src
+    if not path.is_file():
+        return None, {}, ""
+    tree = ast.parse(path.read_text())
+    func = next(
+        (n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == func_name),
+        None,
+    )
+    if func is None:
+        return None, {}, ""
+
+    default = None
+    positional = func.args.posonlyargs + func.args.args
+    for arg, value in zip(positional[len(positional) - len(func.args.defaults) :], func.args.defaults):
+        if arg.arg == key and isinstance(value, ast.Constant) and isinstance(value.value, str):
+            default = value.value
+
+    branches: dict[str, str] = {}
+    for node in ast.walk(func):
+        if not isinstance(node, ast.If):
+            continue
+        test = node.test
+        if not isinstance(test, ast.Compare) or len(test.ops) != 1:
+            continue
+        if not isinstance(test.ops[0], ast.Eq) or not isinstance(test.left, ast.Name):
+            continue
+        if test.left.id != key:
+            continue
+        right = test.comparators[0]
+        if not (isinstance(right, ast.Constant) and isinstance(right.value, str)):
+            continue
+        for stmt in node.body:
+            for call in _calls_named(stmt, "import_module"):
+                arg = call.args[0] if call.args else None
+                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                    branches[right.value] = arg.value
+    return default, branches, str(src)
+
+
+def _config_class_of_module(repo_root: Path, module: str) -> str | None:
+    """The class a module hands to ``core_transformer_config_from_args``."""
+    path = repo_root / (module.replace(".", "/") + ".py")
+    if not path.is_file():
+        return None
+    tree = ast.parse(path.read_text())
+    for call in _calls_named(tree, "core_transformer_config_from_args"):
+        for keyword in call.keywords:
+            if keyword.arg == "config_class" and isinstance(keyword.value, ast.Name):
+                return keyword.value.id
+    return None
+
+
+def _inherited_fields(classes: dict[str, ast.ClassDef], derived: frozenset[str], cls: str) -> set[str]:
+    """Fields ``cls`` declares plus its Primus-side ancestors', stopping at the backend's own
+    config class: an upstream field is upstream's, not this model's."""
+    fields: set[str] = set()
+    frontier = {cls}
+    seen: set[str] = set()
+    while frontier:
+        name = frontier.pop()
+        if name in seen or name not in derived:
+            continue
+        seen.add(name)
+        fields |= _class_fields(classes[name])
+        frontier |= _class_bases(classes[name])
+    return fields
+
+
+def extract_model_scopes(
+    repo_root: Path, backend: str, upstream: BackendSchema, allow: Allowlist
+) -> tuple[ModelScope, ...]:
+    """Derive the model-scoped key sets by following the model dispatch. Four things disqualify a
+    field, each meaning it reaches something on other models too: upstream defines the argument;
+    another Primus config class declares it; something reads it off the argument namespace; or a
+    patch or the hand-written allowlist vouches for it globally."""
+    default, branches, source = _model_provider_branches(repo_root, backend)
+    if not branches:
+        return ()
+    key = MODEL_PROVIDERS[backend][2]
+    classes, derived = _primus_config_classes(repo_root, backend)
+
+    owners: dict[str, str] = {}  # config class -> discriminator value
+    for value, module in sorted(branches.items()):
+        config_class = _config_class_of_module(repo_root, module)
+        if config_class is not None and config_class in derived:
+            owners[config_class] = value
+    if not owners:
+        return ()
+
+    # Which models claim each field. `None` stands for a Primus config class no
+    # discriminator selects, so its fields belong to no model in particular.
+    claims: dict[str, set[str | None]] = {}
+    for name in sorted(derived):
+        value = owners.get(name)
+        for field in _class_fields(classes[name]):
+            claims.setdefault(field, set()).add(value)
+    agnostic = _args_namespace_reads(repo_root)
+
+    def is_scoped(field: str, value: str) -> bool:
+        if claims.get(field) != {value} or field in upstream.keys or field in agnostic:
+            return False
+        rule = allow.match(field)
+        # A `backend:` rule is the read that made the key legal in the first
+        # place, and it is exactly what the scope is about. A patch or a
+        # hand-written entry is a claim that the key works everywhere.
+        return rule is None or rule.origin.startswith("backend:")
+
+    scopes = []
+    for config_class, value in sorted(owners.items(), key=lambda item: item[1]):
+        keys = {k for k in _inherited_fields(classes, derived, config_class) if is_scoped(k, value)}
+        if not keys:
+            continue
+        scopes.append(
+            ModelScope(
+                key=key,
+                value=value,
+                default=default,
+                config_class=config_class,
+                keys=frozenset(keys),
+                source=source,
+            )
+        )
+    return tuple(scopes)
+
+
+def build_model_scopes(
+    repo_root: Path, backend: str, allowlist_path: Path | None = None
+) -> tuple[ModelScope, ...]:
+    """``extract_model_scopes`` from a repo root alone; empty when unavailable."""
+    upstream = SCHEMA_LOADERS[backend](repo_root)
+    if upstream is None:
+        return ()
+    allow = build_allowlist(repo_root, backend, allowlist_path)
+    return extract_model_scopes(repo_root, backend, upstream, allow)
+
+
+def extract_custom_args_schema(repo_root: Path) -> tuple[frozenset[str], frozenset[str], str]:
+    """Follow ``experimental.custom_args_module`` and expand the dataclass it names, mirroring
+    ``config_utils.build_job_config_from_namespace``. Returned as schema, not allowlist entries,
+    so unknown keys *inside* an extension section stay reportable."""
+    empty = (frozenset(), frozenset(), "")
+    preset = repo_root / TORCHTITAN_MODULE_PRESET
+    if not preset.is_file():
+        return empty
+    module_path = (parse_yaml(str(preset)).get("experimental") or {}).get("custom_args_module")
+    if not module_path:
+        return empty
+    src = repo_root / (module_path.replace(".", "/") + ".py")
+    if not src.is_file():
+        return empty
+    defs = _dataclass_defs(ast.parse(src.read_text()))
+    keys: set[str] = set()
+    sections: set[str] = set()
+    _expand_dataclass(defs, "JobConfig", "", keys, sections)
+    return frozenset(keys), frozenset(sections), module_path
+
+
+def load_manual_allowlist(path: Path, backend: str) -> list[AllowRule]:
+    """Read the hand-written allowlist; ``reason`` and ``consumer`` are mandatory."""
+    if not path.is_file():
+        return []
+    data = yaml.safe_load(path.read_text()) or {}
+    rules = []
+    for scope in ("common", backend):
+        for entry in data.get(scope) or []:
+            missing = [f for f in ("key", "reason", "consumer") if not entry.get(f)]
+            if missing:
+                raise ValueError(f"{path}: entry {entry!r} is missing required field(s) {missing}")
+            rules.append(
+                AllowRule(
+                    pattern=entry["key"],
+                    reason=entry["reason"],
+                    consumer=entry["consumer"],
+                    origin=f"allowlist:{scope}",
+                )
+            )
+    return rules
+
+
+def check_manual_allowlist_consumers(repo_root: Path, path: Path, backend: str) -> list[str]:
+    """Report allowlist entries whose ``consumer`` no longer backs up the claim: the file must
+    exist and must still mention the key. Line numbers are deliberately not checked -- they rot
+    within weeks and a wrong one says nothing about whether the key is still read."""
+    if not path.is_file():
+        return []
+    data = yaml.safe_load(path.read_text()) or {}
+    problems = []
+    for scope in ("common", backend):
+        for entry in data.get(scope) or []:
+            consumer = str(entry.get("consumer", ""))
+            # Tolerate a legacy "path:line" form by dropping the line number.
+            rel = consumer.split(":", 1)[0].strip()
+            if not rel:
+                continue
+            target = repo_root / rel
+            if not target.is_file():
+                # Submodules are not always checked out; absence proves nothing.
+                if rel.startswith("third_party/"):
+                    continue
+                problems.append(f"{entry['key']}: consumer {rel} does not exist")
+                continue
+            token = entry["key"].split(".", 1)[0].rstrip("*")
+            if token and token not in target.read_text(errors="ignore"):
+                problems.append(f"{entry['key']}: {rel} no longer mentions '{token}'")
+    return problems
+
+
+def build_allowlist(repo_root: Path, backend: str, allowlist_path: Path | None = None) -> Allowlist:
+    path = allowlist_path or (repo_root / DEFAULT_ALLOWLIST)
+    rules = extract_get_param_rules(repo_root, backend)
+    rules += extract_args_read_rules(repo_root, backend)
+    rules += load_manual_allowlist(path, backend)
+    return Allowlist(tuple(rules))
+
+
+def extend_schema(repo_root: Path, backend: str, schema: BackendSchema) -> BackendSchema:
+    """Add the Primus extension sections declared in-tree to an upstream schema."""
+    if backend == "torchtitan":
+        keys, sections, module_path = extract_custom_args_schema(repo_root)
+        return schema.extended(keys, sections, module_path) if module_path else schema
+    keys, source = extract_backend_config_extensions(repo_root, backend)
+    return schema.extended(keys, (), source) if keys else schema
+
+
+def build_schema(repo_root: Path, backend: str) -> BackendSchema | None:
+    """Upstream schema plus the Primus extension sections declared in-tree."""
+    schema = SCHEMA_LOADERS[backend](repo_root)
+    return extend_schema(repo_root, backend, schema) if schema is not None else None
+
+
+# ---------------------------------------------------------------------------
+# Scanning
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Finding:
+    backend: str
+    key: str
+    file: str
+
+
+@dataclass(frozen=True)
+class ScopedFinding:
+    """A key that exists, but not for the model this config selects."""
+
+    backend: str
+    key: str
+    file: str
+    scope: ModelScope
+    actual: str
+
+
+@dataclass(frozen=True)
+class KeyFinding:
+    """All occurrences of one unknown key, collapsed into a single row."""
+
+    backend: str
+    key: str
+    count: int
+    files: tuple[str, ...]
+    suggestion: str | None
+
+
+@dataclass(frozen=True)
+class ScopedKeyFinding:
+    """All occurrences of one misplaced model-scoped key, collapsed into a row."""
+
+    backend: str
+    key: str
+    scope: ModelScope
+    count: int
+    files: tuple[str, ...]
+    models: tuple[str, ...]
+
+    def sample(self, limit: int = 3) -> str:
+        shown = ", ".join(f"`{f}`" for f in self.files[:limit])
+        extra = len(self.files) - limit
+        return f"{shown} (+{extra} more)" if extra > 0 else shown
+
+
+def _walk(
+    node: Any,
+    prefix: str,
+    schema: BackendSchema,
+    allow: Allowlist,
+    out: list[str],
+    leaves: set[str] | None = None,
+) -> None:
+    """Report the shallowest unknown path; never recurse below a bad or known leaf. ``leaves``
+    collects every path the walk settled on, known or not: a scoped key is a *known* one."""
+    if not isinstance(node, dict):
+        return
+    for name, value in node.items():
+        path = f"{prefix}.{name}" if prefix else str(name)
+        if path in schema.sections or (path not in schema.keys and allow.is_prefix(path)):
+            _walk(value, path, schema, allow, out, leaves)
+            continue
+        if leaves is not None:
+            leaves.add(path)
+        if path in schema.keys or allow.match(path):
+            continue
+        out.append(path)
+
+
+def _module_key_sources(doc: dict) -> Iterator[tuple[str, dict, dict]]:
+    """Yield ``(framework, keys, module)`` triples for each module of an experiment YAML."""
+    for module in (doc.get("modules") or {}).values():
+        if not isinstance(module, dict):
+            continue
+        framework = module.get("framework")
+        keys = {k: v for k, v in module.items() if k not in MODULE_RESERVED_KEYS}
+        keys.update(module.get("overrides") or {})
+        yield framework, keys, module
+
+
+def _preset_filename(name: Any) -> str | None:
+    if not isinstance(name, str) or not name:
+        return None
+    return name if name.endswith((".yaml", ".yml")) else f"{name}.yaml"
+
+
+def build_model_preset_index(repo_root: Path, backend: str, key: str) -> dict[str, str | None]:
+    """``<preset file name> -> <its value for key>`` for every model preset. A preset missing from
+    the index is one the scan could not read, not one that leaves ``key`` unset."""
+    index: dict[str, str | None] = {}
+    base = repo_root / MODEL_PRESETS / backend
+    if not base.is_dir():
+        return index
+    for path in sorted(p for p in base.rglob("*.y*ml") if p.is_file()):
+        try:
+            preset = parse_yaml(str(path))
+        except (OSError, ValueError, yaml.YAMLError):
+            continue
+        if not isinstance(preset, dict):
+            continue
+        value = preset.get(key)
+        index[str(path.relative_to(base))] = None if value is None else str(value)
+    return index
+
+
+def _selected_model(
+    keys: dict, module: dict | None, scope: ModelScope, presets: Mapping[str, str | None]
+) -> str | None:
+    """What ``scope.key`` will be at runtime here, or ``None`` when unknowable. An experiment names
+    its model by preset and does not repeat the ``model_type`` the preset declares, so the preset
+    has to be consulted; an unreadable one leaves the answer unknown and nothing is reported."""
+    if scope.key in keys:
+        return str(keys[scope.key])
+    if module is not None:
+        filename = _preset_filename(module.get("model"))
+        if filename is None or filename not in presets:
+            return None
+        declared = presets[filename]
+        if declared is not None:
+            return declared
+    return scope.default
+
+
+def scan_file(
+    path: Path,
+    backend: str,
+    schema: BackendSchema,
+    allow: Allowlist,
+    rel: str,
+    scopes: Sequence[ModelScope] = (),
+    presets: Mapping[str, str | None] | None = None,
+) -> tuple[list[Finding], list[ScopedFinding]]:
+    doc = parse_yaml(str(path))
+    if not isinstance(doc, dict):
+        return [], []
+    unknown: list[str] = []
+    scoped: list[ScopedFinding] = []
+    known_presets = presets if presets is not None else {}
+
+    def check_scopes(keys: dict, module: dict | None, leaves: set[str]) -> None:
+        for scope in scopes:
+            misplaced = sorted(leaves & scope.keys)
+            if not misplaced:
+                continue
+            actual = _selected_model(keys, module, scope, known_presets)
+            if actual is None or actual == scope.value:
+                continue
+            scoped.extend(ScopedFinding(backend, key, rel, scope, actual) for key in misplaced)
+
+    if "modules" in doc:
+        envelope = {k: v for k, v in doc.items() if k != "modules"}
+        for framework, keys, module in _module_key_sources(doc):
+            if framework != backend:
+                continue
+            leaves: set[str] = set()
+            _walk({**envelope, **keys}, "", schema, allow, unknown, leaves)
+            check_scopes({**envelope, **keys}, module, leaves)
+    else:
+        leaves = set()
+        _walk(doc, "", schema, allow, unknown, leaves)
+        check_scopes(doc, None, leaves)
+    # dict.fromkeys: a file with two modules of the same backend must not count twice.
+    return (
+        [Finding(backend, key, rel) for key in dict.fromkeys(unknown)],
+        list(dict.fromkeys(scoped)),
+    )
+
+
+def iter_config_files(repo_root: Path, roots: Iterable[Path]) -> Iterator[Path]:
+    for root in roots:
+        base = root if root.is_absolute() else repo_root / root
+        if base.is_file():
+            yield base
+        elif base.is_dir():
+            yield from sorted(p for p in base.rglob("*.y*ml") if p.is_file())
+
+
+@dataclass(frozen=True)
+class LoadError:
+    """A config the scanner could not read, and therefore did **not** check."""
+
+    file: str
+    message: str
+
+
+def _describe_load_error(path: Path, exc: Exception, repo_root: Path) -> str:
+    """Name the file that is actually missing, not just the one being scanned: a broken
+    ``extends:`` target is the usual cause and the raw error names an un-normalised path."""
+    missing = getattr(exc, "filename", None)
+    if not missing or Path(missing).resolve() == path.resolve():
+        return f"{type(exc).__name__}: {exc}"
+    target = Path(missing).resolve()
+    try:
+        target = target.relative_to(repo_root.resolve())
+    except ValueError:
+        pass  # outside repo_root (e.g. a submodule checkout); report the absolute path instead
+    return f"cannot resolve `{target}` (check the `extends:` paths)"
+
+
+@dataclass(frozen=True)
+class ScanResult:
+    backend: str
+    schema: BackendSchema | None
+    findings: tuple[Finding, ...] = ()
+    scanned: int = 0
+    errors: tuple[LoadError, ...] = ()
+    scoped: tuple[ScopedFinding, ...] = ()
+    scopes: tuple[ModelScope, ...] = ()
+
+    @property
+    def available(self) -> bool:
+        return self.schema is not None
+
+    @property
+    def checked(self) -> int:
+        """Configs actually inspected. Anything that failed to load is not one."""
+        return self.scanned - len(self.errors)
+
+
+def scan_backend(
+    repo_root: Path,
+    backend: str,
+    roots: Iterable[Path] | None = None,
+    allowlist_path: Path | None = None,
+) -> ScanResult:
+    """Scan one backend; ``ScanResult.available`` is False when the submodule is absent."""
+    upstream = SCHEMA_LOADERS[backend](repo_root)
+    if upstream is None:
+        return ScanResult(backend=backend, schema=None)
+    schema = extend_schema(repo_root, backend, upstream)
+    allow = build_allowlist(repo_root, backend, allowlist_path)
+    scopes = extract_model_scopes(repo_root, backend, upstream, allow)
+    # One discriminator per backend, so every scope shares `key`.
+    presets = build_model_preset_index(repo_root, backend, scopes[0].key) if scopes else {}
+    findings: list[Finding] = []
+    scoped: list[ScopedFinding] = []
+    errors: list[LoadError] = []
+    scanned = 0
+    for path in iter_config_files(repo_root, roots if roots is not None else SCAN_ROOTS[backend]):
+        scanned += 1
+        try:
+            rel = str(path.relative_to(repo_root))
+        except ValueError:
+            rel = str(path)
+        # An unreadable config is an unchecked config: record it instead of
+        # raising (no traceback in CI) and instead of ignoring it (no silent pass).
+        try:
+            unknown, misplaced = scan_file(path, backend, schema, allow, rel, scopes, presets)
+        except (OSError, ValueError, yaml.YAMLError) as exc:
+            errors.append(LoadError(rel, _describe_load_error(path, exc, repo_root)))
+            continue
+        findings.extend(unknown)
+        scoped.extend(misplaced)
+    return ScanResult(
+        backend=backend,
+        schema=schema,
+        findings=tuple(findings),
+        scanned=scanned,
+        errors=tuple(errors),
+        scoped=tuple(scoped),
+        scopes=scopes,
+    )
+
+
+def dedupe(findings: Sequence[Finding], schema: BackendSchema | None = None) -> list[KeyFinding]:
+    """Collapse per-file findings into one row per key: a single upstream rename shows up in
+    dozens of configs, and the per-key view is the actionable one. Widest blast radius first, and
+    ``(backend, key)`` is unique per row, so the order is total and does not depend on the order
+    the findings arrived in."""
+    grouped: dict[tuple[str, str], list[str]] = {}
+    for finding in findings:
+        grouped.setdefault((finding.backend, finding.key), []).append(finding.file)
+    rows = [
+        KeyFinding(
+            backend=backend,
+            key=key,
+            count=len(files),
+            files=tuple(sorted(set(files))),
+            suggestion=schema.suggest(key) if schema is not None else None,
+        )
+        for (backend, key), files in grouped.items()
+    ]
+    return sorted(rows, key=lambda r: (r.backend, -r.count, r.key))
+
+
+def dedupe_scoped(findings: Sequence[ScopedFinding]) -> list[ScopedKeyFinding]:
+    """Collapse per-file scoped findings into one row per key, carrying the models it was wrongly
+    set on: deleting the key and moving it are different fixes."""
+    grouped: dict[tuple[str, str, ModelScope], list[ScopedFinding]] = {}
+    for finding in findings:
+        grouped.setdefault((finding.backend, finding.key, finding.scope), []).append(finding)
+    rows = [
+        ScopedKeyFinding(
+            backend=backend,
+            key=key,
+            scope=scope,
+            count=len(found),
+            files=tuple(sorted({f.file for f in found})),
+            models=tuple(sorted({f.actual for f in found})),
+        )
+        for (backend, key, scope), found in grouped.items()
+    ]
+    return sorted(rows, key=lambda r: (r.backend, -r.count, r.key))
diff --git a/primus/core/config/yaml_loader.py b/primus/core/config/yaml_loader.py
index 89cb8b58a..1368215e8 100644
--- a/primus/core/config/yaml_loader.py
+++ b/primus/core/config/yaml_loader.py
@@ -7,13 +7,12 @@
 import os
 import re
 
-import yaml
-
 from primus.core.config.merge_utils import deep_merge
 
 ENV_PATTERN = re.compile(r"\${([^:{}]+)(?::([^}]*))?}")
 # Matches floats like: 1.2, .5, 1., 1e5, 1.2e-5, -1.2
 FLOAT_PATTERN = re.compile(r"^-?(?:(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?|\d+[eE][+-]?\d+)$")
+_BOOL_MAP = {"true": True, "false": False}
 
 
 def parse_yaml(path: str) -> dict:
@@ -30,6 +29,8 @@ def parse_yaml(path: str) -> dict:
 # 1. Load YAML
 # ================================================================
 def _load_yaml(path: str):
+    import yaml
+
     with open(path, "r") as f:
         return yaml.load(f, Loader=yaml.SafeLoader)
 
@@ -61,9 +62,10 @@ def _resolve_env_in_string(s: str):
 
     Returns
     -------
-    str or int or float
+    str, bool, int, or float
         The resolved value.
         - If environment variable substitution occurs:
+            - If the resulting string is ``true``/``false`` (case-insensitive), it is converted to bool.
             - If the resulting string represents a number, it is converted to int or float.
             - Otherwise, returns the substituted string.
         - If no substitution occurs: returns the original string unchanged (even if it looks numeric).
@@ -98,26 +100,31 @@ def replace_match(m):
         return os.environ.get(var, default)
 
     replaced = ENV_PATTERN.sub(replace_match, s)
-    return _try_numeric(replaced) if replaced != s else replaced
+    return _try_scalar(replaced) if replaced != s else replaced
 
 
-def _try_numeric(v: str):
+def _try_scalar(v: str):
     """
-    Attempt to convert a string value to int or float.
+    Attempt to convert a string value to bool, int, or float.
     Returns the original string if conversion fails.
     """
-    # 1. Integer check
+    # 1. Boolean check (whole-string, case-insensitive)
+    lowered = v.lower()
+    if lowered in _BOOL_MAP:
+        return _BOOL_MAP[lowered]
+
+    # 2. Integer check
     if re.fullmatch(r"-?\d+", v):
         return int(v)
 
-    # 2. Float check using regex to avoid exceptions for non-numeric strings
+    # 3. Float check using regex to avoid exceptions for non-numeric strings
     if FLOAT_PATTERN.fullmatch(v):
         try:
             return float(v)
         except ValueError:
             return v
 
-    # 3. Return original string
+    # 4. Return original string
     return v
 
 
@@ -132,6 +139,8 @@ def _apply_extends(path: str, cfg: dict):
           - preset1.yaml
           - preset2.yaml
 
+    A lone preset may also be given as a scalar: ``extends: preset1.yaml``.
+
     merge order:
         result = deep_merge(preset1, preset2)
         result = deep_merge(result, current_cfg)
@@ -141,7 +150,14 @@ def _apply_extends(path: str, cfg: dict):
 
     merged = {}
 
-    for ext in cfg["extends"]:
+    extends = cfg["extends"]
+    # A single preset may be written as a scalar. Without this, iterating the
+    # string walks it character by character and reports a missing file named
+    # after its first letter.
+    if isinstance(extends, str):
+        extends = [extends]
+
+    for ext in extends:
         ext_path = os.path.join(os.path.dirname(path), ext)
         ext_cfg = parse_yaml(ext_path)
         merged = deep_merge(merged, ext_cfg)  # later preset overrides earlier
diff --git a/primus/core/launcher/parser.py b/primus/core/launcher/parser.py
index 753977f9c..c4a8d28dd 100644
--- a/primus/core/launcher/parser.py
+++ b/primus/core/launcher/parser.py
@@ -1,3 +1,6 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
 from __future__ import annotations
 
 import argparse
@@ -275,6 +278,15 @@ def parse_trainer_module(self, module_name: str):
         for key in ("config", "model"):
             yaml_utils.check_key_in_namespace(module, key)
 
+        # ---- Preserve module-level attributes before loading presets ----
+        # Attributes like 'trainer_class' are at the module level in YAML
+        # and should be preserved even after loading config/model presets
+        preserved_attrs = {}
+        reserved_module_keys = {"name", "framework", "config", "model", "overrides", "params"}
+        for key, value in vars(module).items():
+            if key not in reserved_module_keys and not key.startswith("_"):
+                preserved_attrs[key] = value
+
         # ---- Load module config ----
         model_format = self.get_model_format(framework)
 
@@ -283,6 +295,11 @@ def parse_trainer_module(self, module_name: str):
         module_config.name = f"exp.modules.{module_name}.config"
         module_config.framework = framework
 
+        # Restore preserved module-level attributes
+        for key, value in preserved_attrs.items():
+            if not hasattr(module_config, key):  # Don't override if preset already has it
+                setattr(module_config, key, value)
+
         # ---- Load model config ----
         model_config_dict = PresetLoader.load(module.model, model_format, config_type="models")
         model_config = yaml_utils.dict_to_nested_namespace(model_config_dict)
diff --git a/primus/core/odc/.gitignore b/primus/core/odc/.gitignore
new file mode 100644
index 000000000..4aab3e88c
--- /dev/null
+++ b/primus/core/odc/.gitignore
@@ -0,0 +1,229 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer,
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+# Custom
+.DS_Store
+
+# rocSHMEM build artifacts. The ODC rocSHMEM ops now live in Primus-Turbo; the
+# rocSHMEM static library is an external dependency (ROCSHMEM_HOME), built
+# out-of-tree. Never commit any rocSHMEM install tree / static lib / build log.
+*.a
+rocshmem_runtime/rocshmem_src/
+rocshmem_runtime/rocshmem_single/
+rocshmem_runtime/rocshmem_ro/
+rocshmem_runtime/rocshmem_gda/
+rocshmem_runtime/build.log
diff --git a/primus/core/odc/README.md b/primus/core/odc/README.md
new file mode 100644
index 000000000..0e8c81d16
--- /dev/null
+++ b/primus/core/odc/README.md
@@ -0,0 +1,163 @@
+# On-demand Communication (ODC)
+
+ODC is a patch to FSDP that adapts Parameter Server (PS) into Fully Sharded Data Parallel (FSDP) by replacing collective all-gather and
+reduce-scatter with on-demand point-to-point communication.
+
+![Original-FSDP](./docs/readme/FSDP-ODC.jpg)
+
+With ODC, the synchronization frequency is reduced from per-iteration to per-minibatch, which fundamentally reduces the workload-imbalance bubbles in FSDP.
+
+ODC is accepted in ICLR 2026! Check out the [paper](https://openreview.net/pdf?id=iIEEgI6WsF) for more details.
+
+## Attribution / Provenance
+
+This ODC code is ported from the upstream open-source project
+[sail-sg/odc](https://github.com/sail-sg/odc) (Sea AI Lab, ICLR 2026). Per its
+package metadata the upstream project is released under the MIT License
+(copyright held by the original authors, Sea AI Lab); the upstream repository
+does not ship a standalone `LICENSE` file or per-file license headers.
+
+This version has been adapted for the AMD ROCm / MI300X platform, including
+migrating the communication backend from nvshmem/CUDA to rocSHMEM/MORI + HIP,
+device-side reduce, and torchrun-based launching.
+
+## ODC Primitives
+
+The key idea is to replace the collective all-gather and reduce-scatter with a on-demand point-to-point communication.
+
+![ODC Primitives](./docs/readme/ag_rs.png)
+
+To support transparent on-demand communications, we implement RDMA based primitives on ROCm using HIP GPU IPC (intra-node) and MORI-SHMEM / rocSHMEM (inter-node). Details are in [odc/primitives](./odc/primitives).
+
+## Support for FSDP
+- FSDP1
+- FSDP2
+  - HSDP
+  - `reshard_after_forward=int`
+
+## Usage
+
+### Prerequisites
+
+- PyTorch (ROCm build)
+- ROCm 7.x
+- Python >= 3.8
+- `amd_mori` (symmetric-memory / RDMA backend)
+
+We highly recommend using a ROCm PyTorch base image (e.g. the
+`tasimage/primus-odc` images, ROCm 7.2.0 based).
+
+### Enable ODC
+ODC ships as an **in-tree module of Primus** (no separate `pip install`; it has no
+compiled extension). Put its parent directory on `PYTHONPATH` so `import odc`
+resolves, plus the `odc_early` shim dir so the MORI/TE load-order fix runs at
+interpreter startup:
+```
+export PYTHONPATH=<PRIMUS_ROOT>/primus/core:<PRIMUS_ROOT>/primus/core/odc/odc_early:$PYTHONPATH
+```
+(The launcher `rocshmem_runtime/scripts/run_odc.sh` sets this up automatically.)
+ODC is pure Python. The rocSHMEM backend (single-node XGMI IPC host
+API and multi-node GPU-direct GDA) is provided by Primus-Turbo as the
+`primus_turbo.pytorch._C.odc_rocshmem_host` / `odc_rocshmem_gda` pybind
+submodules and consumed by `odc/primitives/_rocshmem_backend.py`; select it with
+the config items `odc_p2p_backend: rocshmem` (and `odc_rocshmem_gda: true` for
+the multi-node GDA path). Ensure a Primus-Turbo build with the ODC rocSHMEM ops
+is importable (installed, or on `PYTHONPATH`).
+
+> TODO(primus-turbo): pin the exact Primus-Turbo merge commit (`PRIMUS_TURBO_COMMIT`)
+> once the PR adding the ODC rocSHMEM ops is merged.
+
+## Quick Start
+
+A complete example is provided in `examples/llm_training/`:
+```shell
+pip install -r examples/llm_training/requirements.txt
+
+bash examples/llm_training/run.sh
+```
+
+## Memory
+> User may need to tune `MORI_SHMEM_HEAP_SIZE` for better memory usage.
+When using ODC in FSDP, symmetric buffers are allocated for sharded parameters, sharded gradient accumulation buffer, miscellaneous buffers for `gather` and `scatter-accumulate`.
+To achieve smallest memory footprint,
+`MORI_SHMEM_HEAP_SIZE` should be set to be slightly higher than the size of sharded parameters and gradient: params.element_size() * params.numel() + grad_reduce_buf.element_size() * grad_reduce_buf.numel().
+
+### Basic Usage with FSDP1
+
+```python
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+import odc
+from odc.fsdp import fsdp1
+
+
+fsdp1.patch_fsdp1()
+
+torch.distributed.init_process_group(backend="nccl", device_id=device)
+odc.init_shmem()
+
+
+fsdp_model = FSDP(
+    model,
+    # ...
+)
+
+for epoch in range(10):
+    for minibatch in dataset:
+        fsdp1.pre_minibatch_start(fsdp_model)
+        loss = loss_fn(model)
+        loss.backward()
+        fsdp1.pre_optimizer_step(model)
+        optimizer.step()
+        optimizer.zero_grad()
+
+fsdp1.stop()
+```
+
+### Basic Usage with FSDP2
+
+```python
+import torch
+import odc
+from odc.fsdp import fsdp2
+
+
+torch.distributed.init_process_group(backend="nccl", device_id=device)
+odc.init_shmem()
+
+fsdp2.patch_fsdp2()
+
+for layer in model.layers:
+    fully_shard(layer, **fsdp_kwargs)
+fsdp_model = fully_shard(model, **fsdp_kwargs)
+
+# Call patch_lazy_init just as how we call fully_shard above.
+for layer in fsdp_model.layers:
+    fsdp2.patch_lazy_init(layer)
+fsdp2.patch_lazy_init(fsdp_model)
+
+for epoch in range(10):
+    for minibatch in dataset:
+        fsdp2.pre_minibatch_start(fsdp_model)
+        loss = loss_fn(model)
+        loss.backward()
+        fsdp2.pre_optimizer_step(model)
+        optimizer.step()
+        optimizer.zero_grad()
+
+fsdp2.stop()
+```
+
+
+## Development
+
+### Running Linter
+```
+make lint
+```
+
+### Running Tests
+```bash
+make test
+```
diff --git a/primus/core/odc/__init__.py b/primus/core/odc/__init__.py
new file mode 100644
index 000000000..a649c84bc
--- /dev/null
+++ b/primus/core/odc/__init__.py
@@ -0,0 +1,57 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+import logging
+
+# Eagerly expose the runtime-config API only. This module is a stdlib-only leaf,
+# so importing `odc` stays cheap and does NOT pull in odc.primitives -- which is
+# essential: the ODC integration patch calls odc.set_runtime_config(...) BEFORE
+# the primitives are first imported, so their import-time backend selection
+# (odc_p2p_backend) reads the populated config rather than a stale default.
+from odc.runtime_config import OdcRuntimeConfig
+from odc.runtime_config import get_config as get_runtime_config  # noqa: F401
+from odc.runtime_config import set_config as set_runtime_config
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+logger.setLevel(logging.INFO)
+
+
+# The heavy primitives (which import torch/triton/mori and read the runtime
+# config at import time) are exposed lazily via PEP 562 module __getattr__, so
+# they are imported only on first access -- after set_runtime_config has run.
+_LAZY_EXPORTS = {
+    "init_shmem": ("odc.primitives.utils", "init_shmem"),
+    "finalize_distributed": ("odc.primitives.utils", "finalize_distributed"),
+    "SymmBufferRegistry": ("odc.primitives.utils", "SymmBufferRegistry"),
+    "ReductionService": ("odc.primitives.scatter_accumulate", "ReductionService"),
+    "GatherService": ("odc.primitives.gather", "GatherService"),
+}
+
+
+def __getattr__(name):
+    target = _LAZY_EXPORTS.get(name)
+    if target is None:
+        raise AttributeError(f"module 'odc' has no attribute '{name}'")
+    import importlib
+
+    mod = importlib.import_module(target[0])
+    return getattr(mod, target[1])
+
+
+__all__ = [
+    "init_shmem",
+    "SymmBufferRegistry",
+    "ReductionService",
+    "GatherService",
+    "finalize_distributed",
+    "OdcRuntimeConfig",
+    "get_runtime_config",
+    "set_runtime_config",
+]
diff --git a/primus/core/odc/fsdp/__init__.py b/primus/core/odc/fsdp/__init__.py
new file mode 100644
index 000000000..999672bdf
--- /dev/null
+++ b/primus/core/odc/fsdp/__init__.py
@@ -0,0 +1,3 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
diff --git a/primus/core/odc/fsdp/fsdp1.py b/primus/core/odc/fsdp/fsdp1.py
new file mode 100644
index 000000000..9b94101cd
--- /dev/null
+++ b/primus/core/odc/fsdp/fsdp1.py
@@ -0,0 +1,281 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.distributed.fsdp._flat_param import FlatParamHandle, HandleShardingStrategy
+from torch.distributed.fsdp._runtime_utils import _div_if_needed, _FSDPState
+
+import odc
+
+logger = logging.getLogger(__name__)
+
+reduction_service = None
+gather_service = None
+
+
+def get_reduction_service():
+    return reduction_service
+
+
+def get_gather_service():
+    return gather_service
+
+
+def custom_get_reduce_scatter_tensors(
+    state: _FSDPState, unsharded_grad: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Returns the input and output tensors to reduce-scatter, respectively.
+    """
+    chunks = list(unsharded_grad.chunk(state.world_size))
+    numel_to_pad = state.world_size * chunks[0].numel() - unsharded_grad.numel()
+    padded_unsharded_grad = F.pad(unsharded_grad, [0, numel_to_pad]) if numel_to_pad > 0 else unsharded_grad
+    return padded_unsharded_grad
+
+
+def _reduce_grad(state, handle) -> None:
+    """
+    For sharded strategies, this runs gradient reduction, sharded gradient
+    accumulation if needed, and the post-reduction callback.
+    """
+
+    flat_param = handle.flat_param
+    uses_hybrid_sharded_strategy = handle._sharding_strategy in (
+        HandleShardingStrategy.HYBRID_SHARD,
+        HandleShardingStrategy._HYBRID_SHARD_ZERO2,
+    )
+    # We clear `.grad` to permit multiple backwards. This avoids a race where
+    # the second backward pass computation precedes ahead of the first backward
+    # pass reduction, which is possible since the reduction is issued in a
+    # separate stream and is async and would result in reducing the wrong
+    # gradient.
+    unsharded_grad = flat_param.grad.data
+    flat_param.grad = None
+    padded_unsharded_grad = custom_get_reduce_scatter_tensors(state, unsharded_grad)
+    assert state._comm_hook is None, "ODC does not support comm hook"
+    if state._comm_hook is None:  # default path
+        _div_if_needed(padded_unsharded_grad, state._gradient_predivide_factor)
+        pg = handle._fake_process_group if handle._use_fake_reduce else state.process_group
+
+        rs_func = get_reduction_service().scatter_accumulate
+        rs_func(id(handle.flat_param), padded_unsharded_grad, pg)
+        handle.flat_param._saved_grad_shard = get_reduction_service().get_accumulation(id(handle.flat_param))
+
+        assert not uses_hybrid_sharded_strategy, "ODC does not support hybrid sharded strategy"
+    else:
+        pass
+        # NOTE: HSDP variants do not support communication hook.
+
+    assert not handle._offload_params, "ODC does not support offloading"
+    assert not handle._use_orig_params, "ODC does not support using original parameters"
+
+
+def prepare_gradient_for_optim(self):
+    """Prepare the gradient for optimizer computation by moving the sharded gradient to the ``.grad`` attribute."""
+    from torch.distributed.utils import _p_assert
+
+    def cast_grad_to_param_dtype_if_needed(flat_param):
+        # TODO (rohan-varma): test for full precision with keep_low_precision_grads
+        if not self._force_full_precision and self._keep_low_precision_grads:
+            _p_assert(flat_param.grad is not None, "Unexpected None grad!")
+            if flat_param.grad.dtype != self._fwd_bwd_param_dtype:
+                flat_param.grad.data = flat_param.grad.to(self._fwd_bwd_param_dtype)
+                if self._use_orig_params:
+                    self._use_sharded_grad_views()
+
+    flat_param = self.flat_param
+    # TODO (awgu): We should replace these conditional checks to encode
+    # the logical intention more directly.
+    if hasattr(flat_param, "_cpu_grad"):
+        # NOTE: This branch includes `NO_SHARD`.
+        self._check_sharded(flat_param)
+        self._check_on_cpu(flat_param)
+        flat_param.grad = flat_param._cpu_grad  # type: ignore[attr-defined]
+        cast_grad_to_param_dtype_if_needed(flat_param)
+    elif hasattr(flat_param, "_saved_grad_shard"):
+        self._check_sharded(flat_param)
+        self._check_on_compute_device(flat_param)
+        if flat_param._saved_grad_shard is not None:
+            self._check_on_compute_device(flat_param._saved_grad_shard)  # type: ignore[attr-defined]
+        # If no sharded gradient was computed this iteration, then there is
+        # no need to forward `_saved_grad_shard` to `grad`
+        if flat_param._post_backward_called:  # type: ignore[attr-defined]
+            flat_param.grad = None  # type: ignore[attr-defined]
+            if flat_param.grad is not None:
+                cast_grad_to_param_dtype_if_needed(flat_param)
+    else:
+        _p_assert(
+            not self.uses_sharded_strategy or not flat_param._post_backward_called,  # type: ignore[attr-defined]
+            "All sharded parameters that received a gradient in the "
+            "post-backward should use `_saved_grad_shard`",
+        )
+    # Delete `_saved_grad_shard` since its existence indicates a previous
+    # gradient to accumulate with in the post-backward hook
+    if hasattr(flat_param, "_saved_grad_shard"):
+        delattr(flat_param, "_saved_grad_shard")
+
+
+def all_gather_flat_param(self, padded_unsharded_flat_param):
+    """
+    All-gather the handle's flat parameter to the destination ``padded_unsharded_flat_param``.
+
+    Then switch to use the all-gathered tensor.
+    """
+    from torch.distributed.fsdp._common_utils import _no_dispatch_record_stream
+    from torch.distributed.utils import _p_assert
+
+    _p_assert(
+        hasattr(self, "process_group") and hasattr(self, "world_size"),
+        "Expects a process group and world size to have been set via `shard()`",
+    )
+    sharded_flat_param = self.flat_param.data
+    expected_numel = sharded_flat_param.numel() * self.world_size
+    _p_assert(
+        padded_unsharded_flat_param.numel() == expected_numel,
+        f"Expects {expected_numel} numel but got {padded_unsharded_flat_param.numel()}",
+    )
+
+    pg = self._fake_process_group if self._use_fake_all_gather else self.process_group
+
+    # HACK this should be handled by C10D
+    if sharded_flat_param.is_cpu:  # type: ignore[attr-defined]
+        tensor_list = list(
+            torch.chunk(
+                padded_unsharded_flat_param,
+                dist.get_world_size(pg),  # type: ignore[arg-type]
+            )
+        )
+        dist.all_gather(tensor_list, sharded_flat_param, group=pg)
+    else:
+        # padded_unsharded_flat_param output torch.bfloat16 sharded_flat_param input torch.float32
+        assert padded_unsharded_flat_param.dtype == self._fwd_bwd_param_dtype
+        assert sharded_flat_param.dtype == self._orig_param_dtype
+        # Handle dtype conversion for mixed precision
+        # We need to convert to output to orig_dtype for all-gather
+        # as the sharded_flat_param is in orig_dtype for all ranks, then convert back
+        # This is because when using ODC, the peer may have not reached here and be able to
+        # convert input to _fwd_bwd_param_dtype.
+        from torch.distributed.utils import _free_storage
+
+        needs_dtype_conversion = padded_unsharded_flat_param.dtype != sharded_flat_param.dtype
+        if needs_dtype_conversion:
+            assert self._uses_param_mixed_precision
+            # Convert output to orig_dtype for all-gather (creates new tensor)
+            padded_unsharded_flat_param_orig_dtype = padded_unsharded_flat_param.to(self._orig_param_dtype)
+        else:
+            padded_unsharded_flat_param_orig_dtype = padded_unsharded_flat_param
+
+        ag_func = gather_service.gather_into_tensor
+        ag_func(
+            padded_unsharded_flat_param_orig_dtype,  # Could be fp32
+            # padded_unsharded_flat_param,
+            sharded_flat_param,
+            pg,
+        )
+
+        # Convert back to fwd_bwd_param_dtype and free the temporary buffer
+        if needs_dtype_conversion:
+            # Convert the dtype conversion buffer back to fwd_bwd_param_dtype in-place
+            # copy_ will do the implicit conversion.
+            padded_unsharded_flat_param.copy_(padded_unsharded_flat_param_orig_dtype)
+            # Free the temporary buffer
+            _free_storage(padded_unsharded_flat_param_orig_dtype)
+
+    if self._offload_params:
+        # In case of offloading, `flat_param.data` (i.e. sharded param) is
+        # created on the pre-unshard stream. We need to hand it over to the
+        # unshard stream for all-gather
+        _no_dispatch_record_stream(
+            sharded_flat_param,
+            self._device_handle.current_stream(),  # unshard_stream
+        )
+    return padded_unsharded_flat_param
+
+
+old_get_shard = FlatParamHandle._get_shard
+
+
+def custom_get_shard(tensor, rank, world_size):
+    from torch.distributed.fsdp._flat_param import FlatParameter
+
+    sharded, padded = old_get_shard(tensor, rank, world_size)
+    assert isinstance(tensor, FlatParameter)
+    sharded_in_shmem = odc.SymmBufferRegistry.get_instance().update_symm_buffer(id(tensor), sharded)
+    return sharded_in_shmem, padded
+
+
+def _use_low_precision_shard(self):
+    """
+    Allocate on the compute device and switch to using the low precision sharded flat parameter.
+    call path:
+        _runtime_utils.py:_unshard()
+            -> pre_unshard()
+                -> self._use_low_precision_shard()
+    """
+    self._check_low_precision_shard()
+
+
+def _free_low_precision_sharded_param(self):
+    """
+    Frees the low precision sharded flat parameter.
+
+    call path:
+        _runtime_utils.py:_unshard()
+            -> handle.post_unshard()
+                -> self._free_low_precision_sharded_param()
+    """
+    self._check_low_precision_shard()
+
+
+def patch_fsdp1(reduce_dtype=None):
+    global reduction_service, gather_service
+    reduction_service = odc.ReductionService(accumulation_dtype=reduce_dtype)
+    gather_service = odc.GatherService()
+    from torch.distributed.fsdp import _runtime_utils
+
+    _runtime_utils._reduce_grad = _reduce_grad
+
+    FlatParamHandle.prepare_gradient_for_optim = prepare_gradient_for_optim
+    FlatParamHandle._get_shard = custom_get_shard
+    FlatParamHandle._all_gather_flat_param = all_gather_flat_param
+    FlatParamHandle._use_low_precision_shard = _use_low_precision_shard
+    FlatParamHandle._free_low_precision_sharded_param = _free_low_precision_sharded_param
+
+
+def pre_optimizer_step(fsdp_module):
+
+    assert isinstance(fsdp_module, _FSDPState)
+    with torch.cuda.nvtx.range("scatter_accumulate_sync"):
+        get_reduction_service().sync(fsdp_module.process_group)
+
+    for acc in get_reduction_service().accumulations:
+        if hasattr(fsdp_module, "_inter_node_pg"):
+            dist.all_reduce(acc, group=fsdp_module._inter_node_pg)
+        _div_if_needed(acc, fsdp_module._gradient_postdivide_factor)
+    for handle in fsdp_module._all_handles:
+        handle.flat_param.grad = (
+            get_reduction_service().get_accumulation(id(handle.flat_param)).to(handle.flat_param.dtype)
+        )
+
+
+def pre_minibatch_start(_fsdp_module):
+    get_reduction_service().clear_accumulations()
+
+    # Make sure optimizer updates are visible to all ranks
+    dist.barrier()
+
+
+def stop():
+    get_reduction_service().stop()
+    odc.SymmBufferRegistry.get_instance().finalize()
+    odc.finalize_distributed()
diff --git a/primus/core/odc/fsdp/fsdp2.py b/primus/core/odc/fsdp/fsdp2.py
new file mode 100644
index 000000000..5d5e67733
--- /dev/null
+++ b/primus/core/odc/fsdp/fsdp2.py
@@ -0,0 +1,1092 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+import dataclasses
+import logging
+import operator
+import types
+from functools import reduce
+from itertools import chain
+from typing import Any, Callable, Optional, cast
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.device_mesh import _get_device_handle
+from torch.distributed.fsdp import fully_shard
+from torch.distributed.fsdp._fully_shard import _fsdp_collectives
+from torch.distributed.fsdp._fully_shard._fsdp_api import AllGather, ReduceScatter
+from torch.distributed.fsdp._fully_shard._fsdp_collectives import (
+    AllGatherResult,
+    DefaultAllocMixin,
+    _div_if_needed,
+    _get_all_gather_input_metadatas,
+    _get_gradient_divide_factors,
+    _get_param_all_gather_inputs,
+    foreach_reduce_scatter_copy_in,
+)
+from torch.distributed.fsdp._fully_shard._fsdp_common import (
+    FSDPMeshInfo,
+    HSDPMeshInfo,
+    TrainingState,
+    _get_dim0_padded_size,
+    _raise_assert_with_print,
+    _to_dtype_if_needed,
+    compiled_autograd_enabled,
+)
+from torch.distributed.fsdp._fully_shard._fsdp_param import (
+    FSDPParam,
+    ShardedState,
+    set_requires_grad_if_needed,
+)
+from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
+    AllReduceState,
+    FSDPParamGroup,
+    ReduceScatterState,
+)
+from torch.distributed.tensor import DTensor
+from torch.profiler import record_function
+
+from odc.primitives.gather import GatherService
+from odc.primitives.scatter_accumulate import ReductionService
+from odc.primitives.utils import SymmBufferRegistry, finalize_distributed
+
+logger = logging.getLogger(__name__)
+
+
+class nvtx_record_function(record_function):
+    def __enter__(self):
+        torch.cuda.nvtx.range_push(self.name)
+        return super().__enter__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        super().__exit__(exc_type, exc_value, traceback)
+        torch.cuda.nvtx.range_pop()
+
+
+reduction_service = None
+gather_service = None
+
+
+def get_reduction_service():
+    return reduction_service
+
+
+def get_gather_service():
+    return gather_service
+
+
+class ODCAllGather(DefaultAllocMixin, AllGather):
+    def __call__(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        async_op: bool = False,
+    ) -> Optional[dist.Work]:
+        gather = get_gather_service()
+        gather.gather_into_tensor(output_tensor, input_tensor, group)
+        if async_op:
+            event = torch.cuda.Event()
+            event.record()
+            return event
+        return None
+
+
+def get_fsdp_params_key(fsdp_params: list[FSDPParam]) -> str:
+    ids = "_".join([str(id(param)) for param in fsdp_params])
+    return f"fsdp_params_gather_{ids}"
+
+
+def get_hpz_params_key(fsdp_param_key: str) -> str:
+    return f"hpz_{fsdp_param_key}"
+
+
+@torch.no_grad()
+def patch_lazy_init(fsdp_model):
+    state = fsdp_model._get_fsdp_state()
+    group = state._fsdp_param_group
+    prev_lazy_init = group.lazy_init
+
+    def patched_lazy_init(_self):
+        prev_lazy_init()
+        replace_sharded_param_with_symm_buffer(fsdp_model)
+
+    group.lazy_init = types.MethodType(patched_lazy_init, group)
+
+
+@torch.no_grad()
+def replace_sharded_param_with_symm_buffer(
+    fsdp_model,
+) -> torch.Tensor:
+    state = fsdp_model._get_fsdp_state()
+    fsdp_param_group = state._fsdp_param_group
+    if fsdp_param_group is None:
+        return
+
+    is_hpz = fsdp_param_group._use_post_forward_mesh
+    if _enable_hpz is None:
+        raise ValueError("Need to run patch_fsdp2() first")
+    assert _enable_hpz == is_hpz, "If HPZ is enabled, reshard_after_forward must be set to int"
+
+    dtype = fsdp_param_group.fsdp_params[0]._sharded_param_data.dtype
+    hpz_dtype = fsdp_param_group.fsdp_params[0].param_dtype
+    for fsdp_param in fsdp_param_group.fsdp_params:
+        if fsdp_param._sharded_param_data.dtype != dtype:
+            raise ValueError(
+                f"All FSDP parameters must have the same dtype: {fsdp_param._sharded_param_data.dtype=} {dtype=}"
+            )
+        if fsdp_param.param_dtype != hpz_dtype:
+            raise ValueError(
+                f"All FSDP parameters must have the same param dtype: {fsdp_param.param_dtype=} {hpz_dtype=}"
+            )
+    if is_hpz and hpz_dtype is None:
+        logger.warning(
+            f"mixed precision param_dtype is not set but HPZ is enabled, using the same as the original dtype {dtype}"
+        )
+        hpz_dtype = dtype
+
+    total_size = sum(fsdp_param._sharded_param_data.numel() for fsdp_param in fsdp_param_group.fsdp_params)
+
+    hpz_sharded_param_total_size = 0
+    num_nodes = 1
+    hpz_symm_buffer = None
+    post_forward_mesh_info = fsdp_param_group.post_forward_mesh_info
+    # This key needs to be the same as the one used in `foreach_all_gather`
+    key = get_fsdp_params_key(fsdp_param_group.fsdp_params)
+    if not is_hpz:
+        symm_buffer = SymmBufferRegistry.get_instance().get_or_create_symm_buffer(key, (total_size,), dtype)
+    else:
+        # When HPZ is enabled, we keep the original parameters(mostly fp32)
+        # in the torch tensor just like the original FSDP2.
+        # Pre-allocate the symmetric buffer for HPZ (Hierarchical Partitioning for ZeRO in ZeRO++).
+        assert isinstance(post_forward_mesh_info, HSDPMeshInfo), f"{post_forward_mesh_info=}"
+        shard_world_size = post_forward_mesh_info.shard_mesh_size
+        world_size = fsdp_param_group._all_gather_process_group.size()
+        assert world_size % shard_world_size == 0, f"{world_size=} {shard_world_size=}"
+        num_nodes = world_size // shard_world_size
+        hpz_sharded_param_total_size = total_size * num_nodes
+        hpz_key = get_hpz_params_key(key)
+        hpz_symm_buffer = SymmBufferRegistry.get_instance().get_or_create_symm_buffer(
+            hpz_key, (hpz_sharded_param_total_size,), hpz_dtype
+        )
+        logger.info(
+            f"Replacing HPZ params with symmetric buffer of shape {hpz_symm_buffer.shape} and dtype: {hpz_dtype}"
+        )
+
+    offset = 0
+    hpz_offset = 0
+    # Refer to the codes in `FSDPParam._init_sharded_param`
+    for fsdp_param in fsdp_param_group.fsdp_params:
+        param_size = fsdp_param._sharded_param_data.numel()
+        if not is_hpz:
+            numel = reduce(operator.mul, fsdp_param.padded_sharded_param_size, 1)
+            assert numel == param_size, f"{fsdp_param.padded_sharded_param_size=} {param_size=}"
+            symm_buffer[offset : offset + param_size].copy_(fsdp_param._sharded_param_data)
+            padded_sharded_param = symm_buffer[offset : offset + param_size]
+            padded_sharded_param = padded_sharded_param.view(fsdp_param.padded_sharded_param_size)
+            offset += param_size
+
+            fsdp_param._sharded_param_data = padded_sharded_param.view(-1)
+            shard_dim = fsdp_param.fsdp_placement.dim
+            old_sharded_param = fsdp_param.sharded_param._local_tensor
+            length = old_sharded_param.size(shard_dim) if old_sharded_param.numel() > 0 else 0
+            sharded_param = padded_sharded_param.narrow(dim=shard_dim, start=0, length=length)
+            fsdp_param.sharded_param._local_tensor = sharded_param
+        else:
+            hpz_param_size = param_size * num_nodes
+            fsdp_param._sharded_post_forward_param_data = hpz_symm_buffer[
+                hpz_offset : hpz_offset + hpz_param_size
+            ]
+            hpz_offset += hpz_param_size
+
+
+__odc_gather = None
+
+
+def get_odc_gather_comm():
+    global __odc_gather
+    if __odc_gather is None:
+        __odc_gather = ODCAllGather()
+    return __odc_gather
+
+
+original_foreach_all_gather = _fsdp_collectives.foreach_all_gather
+
+
+@torch.no_grad()
+def foreach_all_gather(
+    fsdp_params: list[FSDPParam],
+    group: dist.ProcessGroup,
+    async_op: bool,
+    all_gather_copy_in_stream: torch.Stream,
+    all_gather_stream: torch.Stream,
+    device: torch.device,
+    all_gather_comm: AllGather,
+) -> Optional[AllGatherResult]:
+    is_hpz_list = [
+        param.post_forward_mesh_info is not None and param.mesh_info != param.post_forward_mesh_info
+        for param in fsdp_params
+    ]
+    assert len(set(is_hpz_list)) == 1, f"{is_hpz_list=}"
+    is_hpz = is_hpz_list[0]
+
+    if is_hpz and fsdp_params[0].sharded_state == ShardedState.SHARDED:
+        assert (
+            original_foreach_all_gather is not foreach_all_gather
+        ), "original_foreach_all_gather and foreach_all_gather are the same"
+        return original_foreach_all_gather(
+            fsdp_params,
+            group,
+            async_op,
+            all_gather_copy_in_stream,
+            all_gather_stream,
+            device,
+            all_gather_comm,
+        )
+
+    world_size, _rank = group.size(), group.rank()
+    device_handle = _get_device_handle(device.type)
+    # Override the all-gather comm with ODCAllGather
+    all_gather_comm = get_odc_gather_comm()
+    with device_handle.stream(all_gather_copy_in_stream):
+        key = get_fsdp_params_key(fsdp_params)
+        if fsdp_params[0].sharded_state == ShardedState.SHARDED_POST_FORWARD:
+            for fsdp_param in fsdp_params:
+                assert (
+                    fsdp_param.sharded_state == ShardedState.SHARDED_POST_FORWARD
+                ), f"{fsdp_param.sharded_state=}"
+            key = get_hpz_params_key(key)
+        assert SymmBufferRegistry.get_instance().has_key(
+            key
+        ), f"{key=} not found. The fsdp param group has not been replaced with symm buffer yet."
+        all_gather_input = SymmBufferRegistry.get_instance().get_symm_buffer(key)
+        all_gather_output = all_gather_comm.allocate(
+            (all_gather_input.numel() * world_size,), dtype=all_gather_input.dtype, device=device
+        )
+        # _get_param_all_gather_inputs allocates some tensors
+        # But we don't use them here. Just get the metadata.
+        param_all_gather_inputs = _get_param_all_gather_inputs(fsdp_params)
+        (
+            param_all_gather_input_dtypes,
+            param_all_gather_input_numels,
+            dtype,
+        ) = _get_all_gather_input_metadatas(param_all_gather_inputs)
+        if dtype == torch.uint8:
+            all_gather_inputs = [t.view(torch.uint8) for ts in param_all_gather_inputs for t in ts]
+        else:
+            all_gather_inputs = [*chain.from_iterable(param_all_gather_inputs)]
+        inp_split_sizes = [t.numel() for t in all_gather_inputs]
+        input_size_sum = sum(inp_split_sizes)
+        assert input_size_sum == all_gather_input.numel(), f"{input_size_sum=} != {all_gather_input.numel()}"
+    all_gather_stream.wait_stream(all_gather_copy_in_stream)
+    with device_handle.stream(all_gather_stream):
+        all_gather_work = all_gather_comm(
+            output_tensor=all_gather_output,
+            input_tensor=all_gather_input,
+            group=group,
+            async_op=async_op,
+        )
+        all_gather_event = all_gather_stream.record_event()
+        return AllGatherResult(
+            all_gather_output,
+            all_gather_event,
+            all_gather_work,
+            param_all_gather_input_dtypes,
+            param_all_gather_input_numels,
+            inp_split_sizes,
+        )
+
+
+@torch.no_grad()
+def foreach_all_gather_copy_out(
+    all_gather_result: AllGatherResult,
+    fsdp_params: list[FSDPParam],
+    group: dist.ProcessGroup,
+) -> None:
+    (
+        all_gather_output,
+        all_gather_event,
+        all_gather_work,
+        param_all_gather_input_dtypes,
+        param_all_gather_input_numels,
+        all_gather_input_split_sizes,
+    ) = all_gather_result
+    _dtype, device = all_gather_output.dtype, all_gather_output.device
+    device_handle = _get_device_handle(device.type)
+
+    if all_gather_event is not None:  # sync op
+        device_handle.current_stream().wait_event(all_gather_event)
+    if isinstance(all_gather_work, dist.distributed_c10d.Work):  # async op
+        all_gather_work.wait()
+    world_size, device = group.size(), all_gather_output.device
+
+    # In ODC, we use the original dtype for all-gather input and output.
+    # Then when the output needs to be used here, we cast it to the param_dtype.
+    param_dtype = fsdp_params[0].param_dtype
+    if param_dtype is not None:
+        all_gather_output = _to_dtype_if_needed(all_gather_output, param_dtype)
+        input_dtype = param_all_gather_input_dtypes[0][0]
+        for dtypes in param_all_gather_input_dtypes:
+            for dtype in dtypes:
+                if dtype != input_dtype:
+                    raise ValueError(f"Input dtypes are not the same: {dtype} != {input_dtype}")
+        param_all_gather_input_dtypes = [
+            [param_dtype] * len(dtypes) for dtypes in param_all_gather_input_dtypes
+        ]
+
+    split_with_sizes_out: list[torch.Tensor] = []
+    shard_i_copy_infos: list[tuple[FSDPParam, list[torch.Tensor]]] = []
+    for all_gather_input_numels, all_gather_input_dtypes, fsdp_param in zip(
+        param_all_gather_input_numels, param_all_gather_input_dtypes, fsdp_params
+    ):
+        # NOTE: Under compile, make sure we always recreate all_gather_outputs
+        # per AllGather. See [Note: Invariants for torch.compile Traceable FSDP2].
+        force_recreate = compiled_autograd_enabled()
+        fsdp_param.init_all_gather_outputs(
+            all_gather_input_numels,
+            all_gather_input_dtypes,
+            world_size,
+            device,
+            force_recreate=force_recreate,
+        )
+        if not force_recreate:
+            fsdp_param.alloc_all_gather_outputs()
+        param_all_gather_outputs = fsdp_param.all_gather_outputs
+        if fsdp_param.fsdp_placement.dim != 0:
+            # Copy to a temporary and then chunk-cat into the final all-gather
+            # output tensors
+            param_all_gather_outputs = [torch.empty_like(t) for t in param_all_gather_outputs]
+            shard_i_copy_infos.append((fsdp_param, param_all_gather_outputs))
+        split_with_sizes_out.extend(param_all_gather_outputs)
+
+    all_gather_output = all_gather_output.view(world_size, -1)
+    if all_gather_output.dtype == torch.uint8:
+        out = [t.view(world_size, -1).view(torch.uint8) for t in split_with_sizes_out]
+    else:
+        out = [t.view(world_size, -1) for t in split_with_sizes_out]
+
+    # only avoid VC bump if we are not in inference mode
+    if torch._dynamo.is_compiling():
+        # For torch.compile, we turn off inference_mode for fake tensor
+        # propagation, and therefore graph break on is_inference. For `compile`,
+        # we don't care about VCs, so just skip the optimization.
+        non_inference_outs = []
+    else:
+        non_inference_outs = [o for o in out if not o.is_inference()]
+
+    if len(non_inference_outs) > 0:
+        with torch.autograd._unsafe_preserve_version_counter(tuple(non_inference_outs)):
+            torch.ops.fsdp.split_with_sizes_copy(
+                all_gather_output, all_gather_input_split_sizes, dim=1, out=out
+            )
+    else:
+        torch.ops.fsdp.split_with_sizes_copy(all_gather_output, all_gather_input_split_sizes, dim=1, out=out)
+
+    for fsdp_param, param_all_gather_outputs in shard_i_copy_infos:
+        # Chunk-cat from the temporary to the final all-gather output tensors
+        shard_dim = fsdp_param.fsdp_placement.dim
+
+        with torch.autograd._unsafe_preserve_version_counter(tuple(fsdp_param.all_gather_outputs)):
+            for param_all_gather_output, target_all_gather_output in zip(
+                param_all_gather_outputs, fsdp_param.all_gather_outputs
+            ):
+                padded_sharded_size = (
+                    fsdp_param.padded_sharded_param_size
+                    if fsdp_param.sharded_state == ShardedState.SHARDED
+                    else cast(torch.Tensor, fsdp_param._sharded_post_forward_param_data).size()
+                )
+                pre_param_size = list(padded_sharded_size)
+                pre_param_size[0] *= world_size
+                chunks = torch.chunk(param_all_gather_output.view(pre_param_size), world_size, dim=0)
+                post_param_size = list(padded_sharded_size)
+                post_param_size[shard_dim] *= world_size
+                cat_out = target_all_gather_output.view(post_param_size)
+                torch.cat(chunks, dim=shard_dim, out=cat_out)
+
+
+def pre_minibatch_start(fsdp_module):
+    get_reduction_service().clear_accumulations()
+
+    # Make sure optimizer updates are visible to all ranks
+    dist.barrier()
+
+    ensure_resharded_within_node(fsdp_module)
+
+
+def is_bw() -> bool:
+    return torch._C._current_graph_task_id() != -1
+
+
+# Old version of pytorch does not skip resharding after the recomputation in backward,
+# resulting in duplicated all-gather in backward.
+# Patch the new pytorch version here.
+def post_forward(self, _module: nn.Module, _input: Any, output: Any):
+    if not compiled_autograd_enabled():
+        logger.debug("%s", self._with_fqn("FSDP::post_forward"))
+    with record_function(self._with_fqn("FSDP::post_forward")):
+        if not compiled_autograd_enabled():
+            # for AC(fully_shard(model)), AC runs fsdp's _pre_forward
+            # it shouldn't change post_forward_order
+            if not is_bw():
+                self.reshard()
+                self._record_post_forward()
+        else:
+            self.reshard()
+            self._record_post_forward()
+        self._training_state = TrainingState.IDLE
+        return output
+
+
+def post_backward(self, *_unused: Any):
+    # This method should be idempotent and safe to call even when this
+    # FSDP parameter group was not used in backward (should be a no-op)
+    if not compiled_autograd_enabled():
+        logger.debug("%s", self._with_fqn("FSDP::post_backward"))
+    self._training_state = TrainingState.POST_BACKWARD
+    with record_function(self._with_fqn("FSDP::post_backward_accumulate")):
+        for fsdp_param in self.fsdp_params:
+            fsdp_param.accumulate_unsharded_grad_if_needed()
+    with record_function(self._with_fqn("FSDP::post_backward_reshard")):
+        if not self.reduce_grads:
+            if self.reshard_after_backward:
+                self.reshard()
+            for fsdp_param in self.fsdp_params:
+                fsdp_param.to_accumulated_grad_if_needed()
+            return
+        # Save the autograd-computed gradients before resharding to only
+        # access the unsharded parameters when their data is present
+        fsdp_params_with_grad: list[FSDPParam] = []
+        unsharded_grads: list[torch.Tensor] = []
+        for fsdp_param in self.fsdp_params:
+            if not hasattr(fsdp_param, "_unsharded_param"):
+                continue
+            # May have an accumulated gradient of the reduce dtype if the
+            # previous backward did not reduce-scatter
+            if fsdp_param.unsharded_accumulated_grad is not None:
+                fsdp_params_with_grad.append(fsdp_param)
+                unsharded_grads.append(fsdp_param.unsharded_accumulated_grad_data)
+                fsdp_param.unsharded_accumulated_grad = None
+            elif fsdp_param.unsharded_param.grad is not None:
+                fsdp_params_with_grad.append(fsdp_param)
+                unsharded_grads.append(fsdp_param.unsharded_grad_data)
+                fsdp_param.unsharded_param.grad = None
+        if self.reshard_after_backward:
+            self.reshard()
+    if len(fsdp_params_with_grad) == 0:
+        return
+    with record_function(self._with_fqn("FSDP::post_backward_reduce")):
+        if (
+            self.comm_ctx.reduce_scatter_state is not None
+            and self.comm_ctx.reduce_scatter_state.event is not None
+        ):
+            self.device_handle.current_stream().wait_event(self.comm_ctx.reduce_scatter_state.event)
+        self.comm_ctx.reduce_scatter_state = None
+        all_reduce_pg = self._all_reduce_process_group if self._is_hsdp else None
+        all_reduce_stream: torch.cuda.Stream
+        if all_reduce_pg is None and self._all_reduce_hook_stream is not None:
+            # this means the native HSDP is not enabled,
+            # but user may want to have a custom HSDP setup
+            assert (
+                self._all_reduce_hook is not None
+            ), "all reduce hook stream is specified but hook itself is missing."
+            all_reduce_stream = self._all_reduce_hook_stream
+        else:
+            all_reduce_stream = self.comm_ctx.all_reduce_stream
+
+        self._wait_for_post_backward()
+        (
+            reduce_scatter_input,
+            reduce_scatter_event,
+            self._post_reduce_event,
+            all_reduce_input,
+            all_reduce_event,
+            self._partial_reduce_output,
+        ) = foreach_reduce(
+            self,
+            fsdp_params_with_grad,
+            unsharded_grads,
+            self._reduce_scatter_process_group,
+            self.comm_ctx.reduce_scatter_stream,
+            self._reduce_scatter_comm,
+            self._orig_dtype,
+            self._reduce_dtype,
+            self.device,
+            self.gradient_divide_factor,
+            self._all_reduce_process_group if self._is_hsdp else None,
+            all_reduce_stream,
+            self.all_reduce_grads,
+            self._partial_reduce_output,
+            self._all_reduce_hook,
+            self.force_sum_reduction_for_comms,
+        )
+        self.comm_ctx.reduce_scatter_state = ReduceScatterState(reduce_scatter_input, reduce_scatter_event)
+        if all_reduce_input is not None:
+            if self.device.type != "cpu":
+                assert all_reduce_event is not None
+            self._all_reduce_state = AllReduceState(all_reduce_input, all_reduce_event)
+
+
+def reshard(self, refresh_post_forward_data: bool = False):
+    """
+    This will only be patched in HPZ mode.
+    Keep params sharded on the post-forward mesh (within-node) outside forward
+    when reshard_after_forward is int (HPZ mode).
+    """
+    # Supports gather parameters from local GPUs at the same node
+    # even between forward for different microbatches.
+    # So even in backward, we still does not shard it back to fully-sharded.
+    # We just shard it within each node.
+    # After all the backward is done, we will shard it back to fully-sharded.
+    # Only reshard to post-forward if we currently have unsharded params.
+    if self._training_state in (TrainingState.FORWARD, TrainingState.POST_BACKWARD):
+        if not self._reshard_after_forward:
+            return
+        if self._use_post_forward_mesh:
+            self._to_sharded_post_forward(refresh_post_forward_data=refresh_post_forward_data)
+            self._reshard_after_forward_event = self.device_handle.Event()
+            if self._reshard_after_forward_event is not None:
+                self._reshard_after_forward_event.record()
+            return
+    self._to_sharded()
+
+
+def _to_sharded_post_forward(self, refresh_post_forward_data: bool = False):
+    """This patch is used to supports refresh_post_forward_data argument"""
+    if not self.is_sharded_post_forward:
+        for fsdp_param in self.fsdp_params:
+            fsdp_param.to_sharded_post_forward(refresh_post_forward_data=refresh_post_forward_data)
+        self._sharded_state = ShardedState.SHARDED_POST_FORWARD
+
+
+@dataclasses.dataclass
+class ReduceScatterContext:
+    # arguments
+    fsdp_params: list[FSDPParam]
+    unsharded_grads: list[torch.Tensor]
+    reduce_scatter_group: dist.ProcessGroup
+    reduce_scatter_stream: torch.Stream
+    orig_dtype: Optional[torch.dtype]
+    reduce_dtype: Optional[torch.dtype]
+    device: torch.device
+    gradient_divide_factor: Optional[float]
+    all_reduce_group: Optional[dist.ProcessGroup]  # not `None` iff HSDP
+    all_reduce_stream: torch.Stream
+    all_reduce_grads: bool
+    partial_reduce_output: Optional[torch.Tensor]  # only used for HSDP
+    all_reduce_hook: Optional[Callable[[torch.Tensor], None]]
+    force_sum_reduction_for_comms: bool
+    # others
+    padded_unsharded_sizes: tuple
+    grad_dtype: torch.dtype
+
+
+@torch.no_grad()
+def foreach_reduce(
+    fsdp_param_group: FSDPParamGroup,
+    fsdp_params: list[FSDPParam],
+    unsharded_grads: list[torch.Tensor],
+    reduce_scatter_group: dist.ProcessGroup,
+    reduce_scatter_stream: torch.Stream,
+    reduce_scatter_comm: ReduceScatter,
+    orig_dtype: Optional[torch.dtype],
+    reduce_dtype: Optional[torch.dtype],
+    device: torch.device,
+    gradient_divide_factor: Optional[float],
+    all_reduce_group: Optional[dist.ProcessGroup],  # not `None` iff HSDP
+    all_reduce_stream: torch.Stream,
+    all_reduce_grads: bool,
+    partial_reduce_output: Optional[torch.Tensor],  # only used for HSDP
+    all_reduce_hook: Optional[Callable[[torch.Tensor], None]],
+    force_sum_reduction_for_comms: bool = False,
+) -> tuple[
+    torch.Tensor,
+    torch.Event,
+    torch.Event,
+    Optional[torch.Tensor],
+    Optional[torch.Event],
+    Optional[torch.Tensor],
+]:
+    """
+    ``unsharded_grads`` owns the references to the gradients computed by
+    autograd, so clearing the list frees the gradients.
+    """
+
+    grad_dtypes = {grad.dtype for grad in unsharded_grads}
+    if len(grad_dtypes) != 1:
+        # Check this at runtime since it could be a real runtime error if e.g.
+        # fp8 weights do not produce the correct higher precision gradients
+        _raise_assert_with_print(f"FSDP reduce-scatter expects uniform gradient dtype but got {grad_dtypes}")
+    grad_dtype = unsharded_grads[0].dtype
+    reduce_dtype = reduce_dtype or grad_dtype
+    (predivide_factor, _postdivide_factor, _reduce_scatter_op, _all_reduce_op) = _get_gradient_divide_factors(
+        reduce_scatter_group,
+        all_reduce_group,
+        reduce_dtype,
+        device.type,
+        gradient_divide_factor,
+        force_sum_reduction_for_comms,
+    )
+    world_size = reduce_scatter_group.size()
+    device_handle = _get_device_handle(device.type)
+    current_stream = device_handle.current_stream()
+
+    if world_size > 1:
+        for i, (fsdp_param, unsharded_grad) in enumerate(zip(fsdp_params, unsharded_grads)):
+            if (shard_dim := fsdp_param.fsdp_placement.dim) == 0:
+                continue
+            assert (
+                unsharded_grad.size(shard_dim) % world_size == 0
+            ), f"Shard({shard_dim}) requires even sharding: {unsharded_grad.size()=} {world_size=}"
+            chunks = torch.chunk(unsharded_grad, world_size, dim=shard_dim)
+            unsharded_grads[i] = torch.cat(chunks, dim=0)
+
+    padded_unsharded_sizes = tuple(_get_dim0_padded_size(grad.size(), world_size) for grad in unsharded_grads)
+    reduce_scatter_input_numel = sum(s.numel() for s in padded_unsharded_sizes)
+    reduce_scatter_input = reduce_scatter_comm.allocate(
+        (reduce_scatter_input_numel,),
+        dtype=reduce_dtype,
+        device=device,
+    )
+
+    foreach_reduce_scatter_copy_in(unsharded_grads, reduce_scatter_input, world_size)
+
+    # Only after the copy-in finishes can we free the gradients
+    unsharded_grads.clear()
+    reduce_scatter_stream.wait_stream(current_stream)
+    all_reduce_input = None
+    all_reduce_event = None
+
+    with device_handle.stream(reduce_scatter_stream):
+        _div_if_needed(reduce_scatter_input, predivide_factor)
+        key = id(fsdp_param_group)
+        scatter = get_reduction_service()
+        scatter.scatter_accumulate(key, reduce_scatter_input, reduce_scatter_group)
+
+        post_reduce_stream = reduce_scatter_stream
+        reduce_scatter_event = reduce_scatter_stream.record_event()
+
+        # Save the context for the next update_gradients call
+        fsdp_param_group.__odc_reduce_scatter_context = ReduceScatterContext(
+            fsdp_params=fsdp_params,
+            unsharded_grads=unsharded_grads,
+            reduce_scatter_group=reduce_scatter_group,
+            reduce_scatter_stream=reduce_scatter_stream,
+            orig_dtype=orig_dtype,
+            reduce_dtype=reduce_dtype,
+            device=device,
+            gradient_divide_factor=gradient_divide_factor,
+            all_reduce_group=all_reduce_group,
+            all_reduce_stream=all_reduce_stream,
+            all_reduce_grads=all_reduce_grads,
+            partial_reduce_output=partial_reduce_output,
+            all_reduce_hook=all_reduce_hook,
+            force_sum_reduction_for_comms=force_sum_reduction_for_comms,
+            padded_unsharded_sizes=padded_unsharded_sizes,
+            grad_dtype=grad_dtype,
+        )
+
+        return (
+            reduce_scatter_input,
+            reduce_scatter_event,
+            post_reduce_stream.record_event(),
+            all_reduce_input,
+            all_reduce_event,
+            partial_reduce_output,
+        )
+
+
+def ensure_resharded_within_node(fsdp_module):
+    """
+    Ensure all parameters are sharded within each node at the beginning of epoch.
+    This is needed for reshard_after_forward=int mode with ODC to ensure all GPUs
+    have finished sharding before backward gather operations start.
+    """
+    root_state = fully_shard.state(fsdp_module)
+    root_state._lazy_init()
+    all_fsdp_states = root_state._state_ctx.all_states
+    all_fsdp_param_groups = [
+        state._fsdp_param_group for state in all_fsdp_states if state._fsdp_param_group is not None
+    ]
+
+    hpz = any(fsdp_param_group._use_post_forward_mesh for fsdp_param_group in all_fsdp_param_groups)
+    if not hpz:
+        return
+
+    per_node_pg = None
+    for fsdp_param_group in all_fsdp_param_groups:
+        # Only do this if reshard_after_forward is an int (HPZ mode)
+        if not fsdp_param_group._use_post_forward_mesh:
+            continue
+        if per_node_pg is None:
+            assert isinstance(
+                fsdp_param_group.post_forward_mesh_info, HSDPMeshInfo
+            ), f"{fsdp_param_group.post_forward_mesh_info=}"
+            per_node_pg = fsdp_param_group.post_forward_mesh_info.shard_process_group
+
+        # Set training state to FORWARD so reshard() will do post-forward resharding
+        old_state = fsdp_param_group._training_state
+        fsdp_param_group._training_state = TrainingState.FORWARD
+
+        # Unshard (all-gather on all GPUs) - what pre_forward does
+        with torch.cuda.nvtx.range("unshard"):
+            fsdp_param_group.unshard(async_op=False)
+        with torch.cuda.nvtx.range("wait_for_unshard"):
+            fsdp_param_group.wait_for_unshard()
+
+        # Reshard (shard within each node) - what post_forward does
+        with torch.cuda.nvtx.range("reshard"):
+            fsdp_param_group.reshard(refresh_post_forward_data=True)
+
+        # Restore training state
+        fsdp_param_group._training_state = old_state
+
+    assert per_node_pg is not None, "HPZ enabled but per-node process group not found"
+    torch.distributed.barrier(group=per_node_pg)
+
+
+@torch.no_grad()
+def update_gradients(fsdp_param_group: FSDPParamGroup):
+    if not hasattr(fsdp_param_group, "__odc_reduce_scatter_context"):
+        # This is to support that in some iteration, there is no microbatch,
+        # so no reduce-scatter is needed.
+        # __odc_reduce_scatter_context does not exists here in this case.
+        return
+    reduce_scatter_context = fsdp_param_group.__odc_reduce_scatter_context
+    del fsdp_param_group.__odc_reduce_scatter_context
+    fsdp_params = reduce_scatter_context.fsdp_params
+    reduce_scatter_group = reduce_scatter_context.reduce_scatter_group
+    reduce_scatter_stream = reduce_scatter_context.reduce_scatter_stream
+    orig_dtype = reduce_scatter_context.orig_dtype
+    reduce_dtype = reduce_scatter_context.reduce_dtype
+    device = reduce_scatter_context.device
+    gradient_divide_factor = reduce_scatter_context.gradient_divide_factor
+    all_reduce_group = reduce_scatter_context.all_reduce_group
+    all_reduce_stream = reduce_scatter_context.all_reduce_stream
+    all_reduce_hook = reduce_scatter_context.all_reduce_hook
+    force_sum_reduction_for_comms = reduce_scatter_context.force_sum_reduction_for_comms
+    padded_unsharded_sizes = reduce_scatter_context.padded_unsharded_sizes
+    grad_dtype = reduce_scatter_context.grad_dtype
+
+    device_handle = _get_device_handle(device.type)
+
+    reduce_dtype = reduce_dtype or grad_dtype
+    (_predivide_factor, postdivide_factor, reduce_scatter_op, all_reduce_op) = _get_gradient_divide_factors(
+        reduce_scatter_group,
+        all_reduce_group,
+        reduce_dtype,
+        device.type,
+        gradient_divide_factor,
+        force_sum_reduction_for_comms,
+    )
+    world_size = reduce_scatter_group.size()
+    device_handle = _get_device_handle(device.type)
+    current_stream = device_handle.current_stream()
+
+    with device_handle.stream(reduce_scatter_stream):
+        post_reduce_stream = all_reduce_stream
+
+        scatter = get_reduction_service()
+        key = id(fsdp_param_group)
+        reduce_output = scatter.get_accumulation(key)
+        assert reduce_scatter_op in [
+            torch.distributed.ReduceOp.SUM,
+            torch.distributed.ReduceOp.AVG,
+        ], f"reduce_scatter_op {reduce_scatter_op} is not supported"
+        if reduce_scatter_op == torch.distributed.ReduceOp.AVG:
+            reduce_output /= world_size
+
+        if all_reduce_group is not None:  # HSDP
+            # ODC defers the inter-node all-reduce during HSDP gradient
+            # accumulation (all_reduce_grads=False) inside scatter-accumulate, so
+            # the original implementation's partial-reduce bookkeeping is not
+            # needed here.
+            post_reduce_stream = all_reduce_stream
+            if world_size >= 1:
+                all_reduce_stream.wait_stream(reduce_scatter_stream)
+            else:
+                all_reduce_stream.wait_stream(current_stream)
+            with device_handle.stream(all_reduce_stream):
+                dist.all_reduce(
+                    reduce_output,
+                    group=all_reduce_group,
+                    op=all_reduce_op,
+                )
+    # -- END: ops in reduce_scatter stream
+
+    if all_reduce_hook is not None:
+        # Execute user-specified all reduce hook.
+        # If native HSDP is used, this is executed after the HSDP all reduce.
+        # If 1-d FSDP is used, this is executed post reduce-scatter.
+        post_reduce_stream = all_reduce_stream
+        all_reduce_stream.wait_stream(reduce_scatter_stream)
+        with device_handle.stream(all_reduce_stream):
+            all_reduce_hook(reduce_output)
+    # -- END: ops post reduce_scatter
+
+    with device_handle.stream(post_reduce_stream):
+        _div_if_needed(reduce_output, postdivide_factor)
+        reduce_output = _to_dtype_if_needed(reduce_output, orig_dtype)
+        # View out and accumulate sharded gradients
+        flat_grad_offset = 0
+        for padded_unsharded_size, fsdp_param in zip(padded_unsharded_sizes, fsdp_params):
+            # Assume even sharding for Shard(i), i > 0; otherwise would require
+            # copy-out for contiguous strides
+            new_sharded_grad = torch.as_strided(
+                reduce_output,
+                size=fsdp_param.sharded_size,
+                stride=fsdp_param.contiguous_sharded_stride,
+                storage_offset=flat_grad_offset,
+            )
+            to_accumulate_grad = fsdp_param.sharded_param.grad is not None
+            if fsdp_param.offload_to_cpu:
+                # Only overlap the D2H copy (copying to pinned memory) if not
+                # accumulating gradients since the CPU add kernel depends on
+                # the copy result and we cannot run the add as a callback
+                non_blocking = fsdp_param.pin_memory and not to_accumulate_grad
+                # Since the GPU sharded gradient is allocated in the RS stream,
+                # we can free it here by not keeping a ref without waiting for
+                # the D2H copy since future RS-stream ops run after the copy
+                new_sharded_grad = new_sharded_grad.to(torch.device("cpu"), non_blocking=non_blocking)
+                if non_blocking:
+                    # Record an event on which to block the CPU thread to
+                    # ensure that the D2H copy finishes before the optimizer
+                    fsdp_param.grad_offload_event = post_reduce_stream.record_event()
+            if to_accumulate_grad:
+                assert isinstance(fsdp_param.sharded_param.grad, DTensor)
+                fsdp_param.sharded_param.grad._local_tensor += new_sharded_grad
+            else:
+                new_sharded_dtensor_grad = fsdp_param.to_sharded_dtensor(new_sharded_grad)
+                fsdp_param.sharded_param.grad = new_sharded_dtensor_grad
+            if not compiled_autograd_enabled():
+                for hook in (
+                    getattr(fsdp_param.sharded_param, "_post_accumulate_grad_hooks", {}) or {}
+                ).values():
+                    hook(fsdp_param.sharded_param)
+            padded_sharded_numel = padded_unsharded_size.numel() // world_size
+            flat_grad_offset += padded_sharded_numel
+        post_reduce_event = post_reduce_stream.record_event()
+    # The RS output is allocated in the RS stream and used in the default
+    # stream (for optimizer). To ensure its memory is not reused for later
+    # RSs, we do not need extra synchronization since the sharded parameters
+    # hold refs through the end of backward.
+
+    # Synchronize the default stream with post_reduce_stream to ensure
+    # gradient writes are visible before optimizer step
+    current_stream = device_handle.current_stream()
+    current_stream.wait_event(post_reduce_event)
+
+
+# FSDPParam
+def to_sharded_post_forward(self, refresh_post_forward_data: bool = False) -> None:
+    if self.is_dtensor:
+        raise NotImplementedError("Resharding to smaller mesh with TP is not supported yet")
+    self._assert_in_states(ShardedState.UNSHARDED)
+    assert self.post_forward_mesh_info is not None  # mypy
+    assert len(self.all_gather_outputs) == 1
+    shard_world_size = self.post_forward_mesh_info.shard_mesh_size
+    if (numel := self.all_gather_outputs[0].numel()) % shard_world_size != 0:
+        _raise_assert_with_print(
+            f"All-gather output size ({numel}) must be divisible by the shard "
+            f"world size ({shard_world_size})"
+        )
+    shard_rank = self.post_forward_mesh_info.shard_mesh_rank
+    # pyrefly: ignore  # unbound-name
+    sharded_numel = numel // shard_world_size
+    # Don't replace the symmetric buffer _sharded_post_forward_param_data here.
+
+    # If hpz is enabled, self._sharded_post_forward_param_data
+    # only needs to be updated (copy)
+    # on the first unshard in ensure_resharded_within_node.
+    # Later unshard in forward and backward doing gather
+    # from the _sharded_post_forward_param_data won't change
+    # _sharded_post_forward_param_data itself.
+    if refresh_post_forward_data:
+        self._sharded_post_forward_param_data.copy_(
+            self.all_gather_outputs[0].narrow(0, sharded_numel * shard_rank, sharded_numel)
+        )
+    sharded_post_forward_tensor = torch.as_strided(
+        self._sharded_post_forward_param_data,
+        size=self.sharded_post_forward_size,
+        stride=self.contiguous_sharded_post_forward_stride,
+        storage_offset=0,
+    )
+    self._sharded_post_forward_param = nn.Parameter(
+        self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor)
+    )
+    self._setattr_on_modules(self._sharded_post_forward_param)
+    self.free_unsharded_param()
+    self.sharded_state = ShardedState.SHARDED_POST_FORWARD
+
+
+# FSDPParam
+def to_unsharded(self) -> None:
+    # Assume that the data has been allocated and all-gathered
+    set_requires_grad_if_needed(self.sharded_param, self._unsharded_param)
+    self._setattr_on_modules(self._unsharded_param)
+    if self.sharded_state == ShardedState.SHARDED_POST_FORWARD:
+        # The data is allocated in the default stream via the post-forward
+        # reshard and must be kept alive for the next all-gather copy-in.
+        # Since we call this method after the copy-out, the data's lifetime
+        # is ensured without further synchronization.
+        self._sharded_post_forward_param = None
+        # Do not free the symmetric buffer for HPZ.
+    self.sharded_state = ShardedState.UNSHARDED
+
+
+def pre_optimizer_step(fsdp_module):
+    scatter = get_reduction_service()
+
+    root_state = fully_shard.state(fsdp_module)
+    all_fsdp_states = root_state._state_ctx.all_states
+    all_fsdp_param_groups = [
+        state._fsdp_param_group for state in all_fsdp_states if state._fsdp_param_group is not None
+    ]
+    for i, fsdp_param_group in enumerate(all_fsdp_param_groups):
+        if i == 0:
+            # Scatter-accumulate uses the global shard group even in HPZ since
+            # gradients are sharded across the full world size.
+            mesh_info = fsdp_param_group.mesh_info
+            assert isinstance(mesh_info, FSDPMeshInfo)
+            reduce_scatter_group = mesh_info.shard_process_group
+            with torch.cuda.nvtx.range("scatter_accumulate_sync"):
+                scatter.sync(reduce_scatter_group)
+
+        with torch.cuda.nvtx.range(f"update_gradients:{fsdp_param_group._module_fqn}"):
+            update_gradients(fsdp_param_group)
+        if fsdp_param_group._use_post_forward_mesh and fsdp_param_group.is_sharded_post_forward:
+            # After gradient sync, return to fully-sharded params so the next
+            # minibatch performs cross-node all-gather again.
+            fsdp_param_group._to_sharded()
+
+
+def _get_post_forward_mesh_info_no_convert(reshard_after_forward, mesh_info):
+    """Variant of FSDP's helper that preserves int semantics even on 1 node.
+    This is mainly for development purpose.
+    Running HPZ in 1 node increase memory without any benefit.
+    Actually we don't need to use HPZ in 1 node.
+    """
+    from torch._logging import warning_once
+    from torch.distributed.tensor import DeviceMesh
+
+    shard_mesh_size = mesh_info.shard_mesh_size
+    if not isinstance(reshard_after_forward, (bool, int)):
+        raise ValueError(
+            "reshard_after_forward should be a bool or an int representing the "
+            f"group size to reshard to, not {reshard_after_forward}"
+        )
+    # NOTE: `isinstance(False, int)` returns `True`.
+    if not isinstance(reshard_after_forward, bool) and isinstance(reshard_after_forward, int):
+        if (
+            reshard_after_forward < 1
+            or reshard_after_forward > shard_mesh_size
+            or shard_mesh_size % reshard_after_forward != 0
+        ):
+            raise ValueError(
+                "If passing reshard_after_forward as an int, it should be a "
+                f"factor of {shard_mesh_size}, not {reshard_after_forward}"
+            )
+        if reshard_after_forward == 1:
+            msg = (
+                "reshard_after_forward=1 (int) means resharding parameters to world size 1, "
+                "instead of reshard_after_forward=True (bool)"
+            )
+            warning_once(logger, msg, stacklevel=2)
+            reshard_after_forward = False
+        # In the original pytorch implementation,
+        # if reshard_after_forward == shard_mesh_size,
+        # it is actually equivalent to True but use more memory.
+        # So it sets it to True.
+        # For us, for easier development with just 1 node,
+        # we disable this behavior.
+    post_forward_mesh_info = None
+    if reshard_after_forward is True:
+        post_forward_mesh_info = mesh_info
+    elif reshard_after_forward is not False:  # int case
+        post_forward_mesh_tensor = mesh_info.mesh.mesh.view(-1, reshard_after_forward)
+        post_forward_mesh = DeviceMesh(mesh_info.mesh.device_type, post_forward_mesh_tensor)
+        post_forward_mesh_info = HSDPMeshInfo(post_forward_mesh, shard_mesh_dim=1, replicate_mesh_dim=0)
+    return post_forward_mesh_info
+
+
+_enable_hpz = None
+
+
+def patch_fsdp2(enable_hpz: bool = False) -> None:
+    from torch.distributed._composable import replicate_with_fsdp
+    from torch.distributed.fsdp._fully_shard import (
+        _fsdp_init,
+        _fsdp_param_group,
+        _fully_shard,
+    )
+
+    global _enable_hpz
+    if _enable_hpz is None:
+        _enable_hpz = enable_hpz
+    else:
+        assert (
+            _enable_hpz == enable_hpz
+        ), f"HPZ mode is already set to {_enable_hpz}, cannot change to {enable_hpz}"
+
+    _fsdp_collectives.foreach_all_gather = foreach_all_gather
+    _fsdp_param_group.foreach_all_gather = foreach_all_gather
+    _fsdp_collectives.foreach_all_gather_copy_out = foreach_all_gather_copy_out
+    _fsdp_param_group.foreach_all_gather_copy_out = foreach_all_gather_copy_out
+    FSDPParamGroup.post_backward = post_backward
+    FSDPParamGroup.post_forward = post_forward
+    if enable_hpz:
+        _fsdp_init._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert
+        _fully_shard._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert
+        replicate_with_fsdp._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert
+        FSDPParamGroup.reshard = reshard
+        FSDPParamGroup._to_sharded_post_forward = _to_sharded_post_forward
+        FSDPParam.to_sharded_post_forward = to_sharded_post_forward
+        FSDPParam.to_unsharded = to_unsharded
+    _fsdp_param_group.record_function = nvtx_record_function
+    torch.profiler.record_function = nvtx_record_function
+    torch.autograd.profiler.record_function = nvtx_record_function
+
+    global reduction_service
+    global gather_service
+    reduction_service = ReductionService()
+    gather_service = GatherService()
+
+
+def stop():
+    get_reduction_service().stop()
+    SymmBufferRegistry.get_instance().finalize()
+    finalize_distributed()
+
+
+def get_symm_buffer_memory_breakdown() -> dict[str, float]:
+    """Return ODC symmetric-buffer memory breakdown in GB."""
+    breakdown = {
+        "fsdp_params": 0,
+        "accumulation": 0,
+        "shared_buffer": 0,
+        "ag_buffer": 0,
+        "hpz_buffer": 0,
+        "other": 0,
+    }
+    registry = SymmBufferRegistry.get_instance()
+    for key, tensor in registry.local_tensor.items():
+        nbytes = tensor.nbytes
+        if key.startswith("fsdp_params_gather_"):
+            breakdown["fsdp_params"] += nbytes
+        elif key.startswith("rs_accumulation_"):
+            breakdown["accumulation"] += nbytes
+        elif key.startswith("shared_buffer_"):
+            breakdown["shared_buffer"] += nbytes
+        elif key.startswith("ag_buffer_"):
+            breakdown["ag_buffer"] += nbytes
+        elif key.startswith("hpz_"):
+            breakdown["hpz_buffer"] += nbytes
+        else:
+            breakdown["other"] += nbytes
+    return {k: round(v / (1024**3), 2) for k, v in breakdown.items()}
diff --git a/primus/core/odc/odc_early/sitecustomize.py b/primus/core/odc/odc_early/sitecustomize.py
new file mode 100644
index 000000000..dd74b1448
--- /dev/null
+++ b/primus/core/odc/odc_early/sitecustomize.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""ODC early-init shim (auto-executed by Python's site machinery).
+
+Root cause (verified on MI300X + ROCm 7.2 + PyTorch 2.10a):
+    Importing `transformer_engine` BEFORE MORI's C++ runtime is loaded leaves
+    the process in a state where MORI init (or even just allocating its
+    symmetric heap) aborts with `free(): invalid pointer`. It is a C++
+    dynamic-library load-order / global-ctor conflict, NOT an init-order or
+    init-method issue.
+
+    The fix is simply to load MORI's C++ runtime (via `import mori`) BEFORE
+    Megatron imports transformer_engine. This file runs at interpreter
+    startup (site.py imports `sitecustomize` if it is on sys.path), which is
+    earlier than any Primus/Megatron/TE import, guaranteeing the correct
+    order.
+
+Enable by putting this directory on PYTHONPATH (the ODC launcher does this):
+    export PYTHONPATH=<PRIMUS_ROOT>/primus/core/odc/odc_early:$PYTHONPATH
+
+This shim is a pure bootstrap (a launch-time load-order workaround, NOT feature
+logic): it only runs when its directory is on PYTHONPATH, which the ODC launcher
+adds exactly for ODC runs, so it is a complete no-op for normal runs. Whether ODC
+is actually active for training is decided by the enable_odc config item, not by
+this file.
+"""
+
+import sys
+
+try:
+    import mori  # noqa: F401  -- loads libmori / MORI C++ runtime
+    import mori.shmem  # noqa: F401
+
+    sys.stderr.write("[ODC sitecustomize] pre-imported mori before TE (load-order fix)\n")
+    sys.stderr.flush()
+except Exception as _e:  # noqa: BLE001
+    sys.stderr.write(f"[ODC sitecustomize] WARNING: pre-import mori failed: {_e}\n")
+    sys.stderr.flush()
diff --git a/primus/core/odc/primitives/__init__.py b/primus/core/odc/primitives/__init__.py
new file mode 100644
index 000000000..ef3f8fb6f
--- /dev/null
+++ b/primus/core/odc/primitives/__init__.py
@@ -0,0 +1,43 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+from .shmem_triton import (
+    LIB_SHMEM_PATH,
+    SHMEM_EXTERN_LIBS,
+    __syncthreads,
+    getmem_nbi_block,
+    int_atomic_compare_swap,
+    int_atomic_swap,
+    int_g,
+    int_p,
+    int_p_remote,
+    int_wait_until_equals,
+    int_wait_until_equals_remote,
+    putmem_nbi_block,
+    quiet,
+    tid,
+)
+
+__all__ = [
+    # shmem_triton
+    "int_atomic_compare_swap",
+    "int_atomic_swap",
+    "putmem_nbi_block",
+    "getmem_nbi_block",
+    "quiet",
+    "int_p",
+    "int_p_remote",
+    "int_g",
+    "int_wait_until_equals",
+    "int_wait_until_equals_remote",
+    "tid",
+    "__syncthreads",
+    "LIB_SHMEM_PATH",
+    "SHMEM_EXTERN_LIBS",
+]
diff --git a/primus/core/odc/primitives/_rocshmem_backend.py b/primus/core/odc/primitives/_rocshmem_backend.py
new file mode 100644
index 000000000..682ba2d28
--- /dev/null
+++ b/primus/core/odc/primitives/_rocshmem_backend.py
@@ -0,0 +1,588 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""ODC rocSHMEM P2P backend, consuming the rocSHMEM ops from Primus-Turbo.
+
+Activated by ``ODC_P2P_BACKEND=rocshmem`` (the default is ``mori``). The
+rocSHMEM host/GDA surface is no longer loaded from an in-tree ``librs_host*.so``
+via ctypes; it is now provided by Primus-Turbo as two pybind submodules of
+``primus_turbo.pytorch._C``:
+
+  * ``odc_rocshmem_host`` — single-node host-API surface (XGMI IPC path):
+    ``rs_get_uid`` / ``rs_init_uid`` / ``rs_malloc`` / ``rs_ptr`` /
+    ``rs_barrier`` / ``rs_finalize``. Selected when odc_rocshmem_gda is false.
+  * ``odc_rocshmem_gda`` — multi-node GPU-direct (GDA) surface: the same
+    host-compatible ``rs_*`` bootstrap plus the device ``gda_gather`` /
+    ``gda_reduce_scatter_acc`` launchers. Selected when odc_rocshmem_gda is true.
+
+Single-node (host) path: this backend uses ONLY the host API to manage the
+symmetric heap and resolve peer pointers. It deliberately does NOT link any
+rocSHMEM *device* bitcode into Triton — every on-device P2P op is plain Triton
+``tl.load`` / ``tl.store`` / ``tl.atomic_*`` to XGMI-mapped peer addresses (see
+the ``rocshmem`` branch in ``shmem_triton.py``). The heavy gather / scatter data
+movement is host-side torch ``.copy_()`` on peer-view tensors (XGMI peer
+load/store); the only device primitives exercised are ``int_p`` /
+``int_wait_until_equals`` for the scatter-accumulate hand-shake. Those translate
+a local symmetric address to a peer address with a per-PE affine delta —
+``rocshmem_ptr`` was empirically verified to be affine (``rs_ptr(x, pe) - x`` is
+constant across symmetric addresses ``x``), so a single ``delta[pe]`` resolves
+every symmetric pointer.
+
+pybind surface consumed (see Primus-Turbo csrc/pytorch/dist/odc_rocshmem_*.cpp)::
+
+    int       rs_uid_bytes();
+    bytes     rs_get_uid();                     # returns the uid as `bytes`
+    void      rs_init_uid(int rank, int nranks, bytes uid);
+    int       rs_my_pe();
+    int       rs_n_pes();
+    long long rs_malloc(size_t n);              # symmetric-heap device ptr
+    long long rs_ptr(long long p, int pe);      # peer mapping of symmetric ptr p
+    void      rs_barrier();
+    void      rs_finalize();
+    # GDA submodule only, device launchers (peers passed as a Python list):
+    int       gda_gather(target, src, nbytes, list peers, stride_bytes);
+    int       gda_reduce_scatter_acc(...);      # etc.
+"""
+
+import ctypes
+import logging
+import os
+from functools import reduce
+
+import torch
+import torch.distributed as dist
+
+from odc.runtime_config import get_config
+
+logger = logging.getLogger(__name__)
+
+# c10::ScalarType enum codes (stable across torch versions).
+_DTYPE_CODE = {
+    torch.uint8: 0,
+    torch.int8: 1,
+    torch.int16: 2,
+    torch.int32: 3,
+    torch.int64: 4,
+    torch.float16: 5,
+    torch.float32: 6,
+    torch.float64: 7,
+    torch.bool: 11,
+    torch.bfloat16: 15,
+}
+
+# Max same-node PEs supported by the device delta table (MI300X node == 8 GPUs).
+MAX_LOCAL_PES = 8
+
+_FROM_BLOB_CPP = r"""
+#include <torch/extension.h>
+// Wrap an externally-owned device pointer as a torch.Tensor (no deleter: the
+// rocSHMEM symmetric heap owns the memory and frees it at rs_finalize()).
+torch::Tensor odc_rs_from_blob(int64_t ptr, std::vector<int64_t> sizes,
+                               int64_t dtype_code, int64_t device_index) {
+    auto dtype = static_cast<c10::ScalarType>(dtype_code);
+    auto options = torch::TensorOptions().dtype(dtype).device(torch::kCUDA, device_index);
+    return torch::from_blob(reinterpret_cast<void*>(ptr), sizes, options);
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Module state
+# ---------------------------------------------------------------------------
+_lib = None
+_from_blob = None
+_my_pe = -1
+_n_pes = -1
+_ref_base = None  # reference symmetric address used for affine peer deltas
+_allocations = []  # keep python tensor objects alive for the run
+_initialized = False
+
+_gda_enabled = False  # GPU-direct (GDA) device path: config odc_rocshmem_gda=true
+_ctypes_gda_lib = False  # True when odc_rocshmem_lib overrides embedded turbo GDA
+
+
+def gda_enabled():
+    """True when the GPU-direct (GDA) device-kernel cross-node path is active."""
+    return _gda_enabled
+
+
+def _is_gda():
+    # config item odc_rocshmem_gda: GPU-direct (GDA) device path (multi-node).
+    return bool(get_config().rocshmem_gda)
+
+
+def _load_ctypes_gda(so):
+    """Load proven in-tree librs_host_gda.so when odc_rocshmem_lib is set.
+
+    The embedded turbo ``odc_rocshmem_gda`` path linked against rocshmem_combined
+    can fault on first device gather; this override uses the monolithic RDC .so
+    that matches the in-tree dual baseline.
+    """
+    lib = ctypes.CDLL(so)
+    lib.rs_uid_bytes.restype = ctypes.c_int
+    lib.rs_get_uid.argtypes = [ctypes.c_char_p]
+    lib.rs_init_uid.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_char_p]
+    lib.rs_my_pe.restype = ctypes.c_int
+    lib.rs_n_pes.restype = ctypes.c_int
+    lib.rs_malloc.restype = ctypes.c_longlong
+    lib.rs_malloc.argtypes = [ctypes.c_size_t]
+    lib.rs_ptr.restype = ctypes.c_longlong
+    lib.rs_ptr.argtypes = [ctypes.c_longlong, ctypes.c_int]
+    lib.rs_barrier.restype = None
+    lib.rs_finalize.restype = None
+    if hasattr(lib, "gda_gather"):
+        lib.gda_gather.restype = ctypes.c_int
+        lib.gda_gather.argtypes = [
+            ctypes.c_longlong,
+            ctypes.c_longlong,
+            ctypes.c_size_t,
+            ctypes.c_void_p,
+            ctypes.c_int,
+            ctypes.c_size_t,
+        ]
+        lib.gda_reduce_scatter_acc.restype = ctypes.c_int
+        lib.gda_reduce_scatter_acc.argtypes = [
+            ctypes.c_longlong,
+            ctypes.c_longlong,
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+            ctypes.c_int,
+            ctypes.c_longlong,
+            ctypes.c_size_t,
+            ctypes.c_int,
+            ctypes.c_int,
+        ]
+        if hasattr(lib, "gda_stage_fence"):
+            lib.gda_stage_fence.restype = ctypes.c_int
+            lib.gda_stage_fence.argtypes = [ctypes.c_longlong, ctypes.c_longlong, ctypes.c_size_t]
+        if hasattr(lib, "gda_hdp_flush"):
+            lib.gda_hdp_init.restype = ctypes.c_int
+            lib.gda_hdp_init.argtypes = []
+            lib.gda_hdp_flush.restype = ctypes.c_int
+            lib.gda_hdp_flush.argtypes = []
+        if hasattr(lib, "gda_strided_touch"):
+            lib.gda_strided_touch.restype = ctypes.c_int
+            lib.gda_strided_touch.argtypes = [
+                ctypes.c_longlong,
+                ctypes.c_size_t,
+                ctypes.c_size_t,
+                ctypes.c_int,
+                ctypes.c_size_t,
+                ctypes.c_size_t,
+                ctypes.c_longlong,
+                ctypes.c_size_t,
+                ctypes.c_int,
+            ]
+        if hasattr(lib, "gda_reduce_scatter_acc_async"):
+            lib.gda_reduce_scatter_acc_async.restype = ctypes.c_int
+            lib.gda_reduce_scatter_acc_async.argtypes = lib.gda_reduce_scatter_acc.argtypes
+            lib.gda_rs_overlap_sync.restype = ctypes.c_int
+            lib.gda_rs_overlap_sync.argtypes = []
+        if hasattr(lib, "gda_gather_async"):
+            lib.gda_gather_async.restype = ctypes.c_int
+            lib.gda_gather_async.argtypes = [
+                ctypes.c_longlong,
+                ctypes.c_longlong,
+                ctypes.c_size_t,
+                ctypes.c_void_p,
+                ctypes.c_int,
+                ctypes.c_size_t,
+                ctypes.c_longlong,
+            ]
+    return lib
+
+
+def _load_backend():
+    """Resolve the ODC rocSHMEM backend submodule from Primus-Turbo.
+
+    When the odc_rocshmem_lib config points at a GDA ``librs_host_gda.so``, load
+    that monolithic binding via ctypes instead of the embedded turbo GDA kernels.
+    """
+    global _ctypes_gda_lib
+    # config item odc_rocshmem_lib: optional monolithic librs_host_gda.so override.
+    so = get_config().rocshmem_lib
+    if so and os.path.isfile(so) and _is_gda():
+        _ctypes_gda_lib = True
+        logger.info("rocSHMEM GDA backend: ctypes override %s", so)
+        return _load_ctypes_gda(so)
+    _ctypes_gda_lib = False
+    try:
+        import primus_turbo  # noqa: F401  -- package init (also loads the _C ext)
+        import primus_turbo.pytorch._C as _C
+    except ImportError as e:
+        raise ImportError(
+            "ODC rocSHMEM backend now consumes the rocSHMEM ops from Primus-Turbo "
+            "(primus_turbo.pytorch._C.odc_rocshmem_host / odc_rocshmem_gda). Install "
+            "Primus-Turbo built with the ODC rocSHMEM ops, or put its build tree on "
+            "PYTHONPATH so `import primus_turbo` resolves it."
+        ) from e
+    name = "odc_rocshmem_gda" if _is_gda() else "odc_rocshmem_host"
+    mod = getattr(_C, name, None)
+    if mod is None:
+        available = [n for n in dir(_C) if "odc" in n.lower()]
+        raise RuntimeError(
+            f"primus_turbo.pytorch._C has no submodule '{name}'. This primus_turbo "
+            f"build was compiled without the ODC rocSHMEM ops (DISABLE_ROCSHMEM). "
+            f"ODC submodules present: {available}"
+        )
+    return mod
+
+
+def _build_from_blob():
+    from torch.utils.cpp_extension import load_inline
+
+    mod = load_inline(
+        name="odc_rs_from_blob",
+        cpp_sources=[_FROM_BLOB_CPP],
+        functions=["odc_rs_from_blob"],
+        verbose=False,
+    )
+    return mod.odc_rs_from_blob
+
+
+def _wrap(ptr, shape, dtype, dev_index):
+    if isinstance(shape, int):
+        shape = (shape,)
+    return _from_blob(int(ptr), list(shape), _DTYPE_CODE[dtype], int(dev_index))
+
+
+def _nbytes(shape, dtype):
+    if isinstance(shape, int):
+        shape = (shape,)
+    numel = reduce(lambda a, b: a * b, shape, 1)
+    return numel * torch._utils._element_size(dtype)
+
+
+def init():
+    """Bootstrap rocSHMEM from PyTorch's WORLD process group via a unique-id
+    broadcast.
+
+    BOTH the single-node host (``odc_rocshmem_host``) and the multi-node
+    GPU-direct (``odc_rocshmem_gda``) submodules use the SAME unique-id
+    bootstrap: rank 0 generates the uid (``rs_get_uid``, returned as ``bytes``),
+    it is broadcast over the torch WORLD process group, and every rank calls
+    ``rs_init_uid`` (which maps to ``rocshmem_init_attr(ROCSHMEM_INIT_WITH_UNIQUEID)``).
+    The GDA transport exchanges the uid over a TCP socket bootstrap
+    (``ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME``), so NO MPI job / mpirun is required —
+    the run can be launched with plain torchrun.
+    """
+    global _lib, _from_blob, _my_pe, _n_pes, _initialized, _gda_enabled
+    if _initialized:
+        return
+    assert dist.is_initialized(), "torch.distributed must be initialized first"
+    _lib = _load_backend()
+    _from_blob = _build_from_blob()
+
+    _gda_enabled = _is_gda() and hasattr(_lib, "gda_gather")
+    if _gda_enabled:
+        # GDA multi-node bootstrap carries rank-0's address inside the uid and
+        # exchanges it over a TCP socket. Default the bootstrap NIC to eth0 unless
+        # the deployment overrode it; this replaces the old MPI/mpirun bootstrap.
+        os.environ.setdefault("ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME", "eth0")
+
+    rank = dist.get_rank()
+    world = dist.get_world_size()
+    if _ctypes_gda_lib:
+        n = _lib.rs_uid_bytes()
+        buf = (ctypes.c_char * n)()
+        uidb = None
+        if rank == 0:
+            _lib.rs_get_uid(buf)
+            uidb = bytes(buf)
+        obj = [uidb]
+        dist.broadcast_object_list(obj, src=0)
+        uidb = obj[0]
+        ctypes.memmove(buf, uidb, n)
+        _lib.rs_init_uid(rank, world, buf)
+    else:
+        # rs_get_uid() returns the uid as `bytes` (pybind); broadcast it as a python
+        # object and hand the raw bytes back to rs_init_uid on every rank.
+        uidb = _lib.rs_get_uid() if rank == 0 else None
+        obj = [uidb]
+        dist.broadcast_object_list(obj, src=0)
+        uidb = obj[0]
+        _lib.rs_init_uid(rank, world, uidb)
+
+    _my_pe = _lib.rs_my_pe()
+    _n_pes = _lib.rs_n_pes()
+    assert (
+        _my_pe == rank and _n_pes == world
+    ), f"rocSHMEM PE mismatch: my_pe={_my_pe} rank={rank} n_pes={_n_pes} world={world}"
+    logger.info(
+        "init_shmem (rocSHMEM %s, uid bootstrap): my_pe=%d n_pes=%d",
+        "GDA" if _gda_enabled else "host-API",
+        _my_pe,
+        _n_pes,
+    )
+    _initialized = True
+
+
+def _ensure_peer_deltas(base, local_world_size, rank):
+    """On the first allocation, compute the affine peer deltas and push them to
+    the Triton device layer so int_p / int_g / int_wait_until_equals can
+    translate a local symmetric address to a peer address on-device."""
+    global _ref_base
+    if _ref_base is not None:
+        return
+    assert local_world_size <= MAX_LOCAL_PES, (
+        f"rocshmem backend supports up to {MAX_LOCAL_PES} same-node PEs, got "
+        f"local_world_size={local_world_size}"
+    )
+    _ref_base = base
+    node_start = rank - rank % local_world_size
+    # Index deltas by LOCAL position (pe - node_start), NOT global pe: on node>0
+    # the same-node PEs are [node_start, node_start+lws), which would overflow a
+    # global-indexed 8-slot table. The device kernel translates a runtime global
+    # pe to a local position via the baked node_start (see _rs_peer_delta).
+    deltas = [0] * MAX_LOCAL_PES  # indexed by LOCAL position
+    for i in range(local_world_size):
+        pe = node_start + i
+        if pe == _my_pe:
+            deltas[i] = 0
+        else:
+            peer = _lib.rs_ptr(base, pe)
+            if peer == 0:
+                raise RuntimeError(f"rs_ptr(base, {pe}) == 0: no XGMI P2P route to same-node peer")
+            deltas[i] = peer - base
+    from .shmem_triton import set_rocshmem_peer_deltas
+
+    set_rocshmem_peer_deltas(deltas, node_start)
+    logger.info(
+        "rocSHMEM peer deltas (node_start=%d, local_pos -> delta bytes): %s",
+        node_start,
+        deltas,
+    )
+
+
+def alloc_peer_tensors(shape, dtype, local_world_size, rank):
+    """Allocate one symmetric buffer and return ``(local_tensor, peer_tensors)``
+    where ``peer_tensors[i]`` is the same allocation viewed from the address
+    space of same-node PE ``node_start + i`` (length ``local_world_size``).
+
+    Equivalent to MORI's ``mori_shmem_create_tensor_list_intra_node``.
+    """
+    base = _lib.rs_malloc(_nbytes(shape, dtype))
+    if base == 0:
+        raise RuntimeError(
+            f"rs_malloc({_nbytes(shape, dtype)} bytes) returned NULL — symmetric "
+            f"heap exhausted? Raise the rocSHMEM heap size."
+        )
+    dev = torch.cuda.current_device()
+    if _gda_enabled:
+        # GDA (USE_IPC=OFF): no intra-node IPC peer views; all cross-node AND
+        # same-node transfers go through device get/put. Return the local view
+        # plus dummy peer entries (never used on the GDA path).
+        local_tensor = _wrap(base, shape, dtype, dev)
+        _allocations.append(local_tensor)
+        return local_tensor, [local_tensor] * local_world_size
+    _ensure_peer_deltas(base, local_world_size, rank)
+
+    node_start = rank - rank % local_world_size
+    peer_tensors = []
+    for i in range(local_world_size):
+        pe = node_start + i
+        p = base if pe == _my_pe else _lib.rs_ptr(base, pe)
+        if p == 0:
+            raise RuntimeError(f"rs_ptr(base, {pe}) == 0 (no XGMI route to peer {pe})")
+        peer_tensors.append(_wrap(p, shape, dtype, dev))
+    local_tensor = peer_tensors[rank % local_world_size]
+    _allocations.append(local_tensor)
+    return local_tensor, peer_tensors
+
+
+def create_tensor(shape, dtype):
+    """Bare local symmetric tensor (no peer list)."""
+    base = _lib.rs_malloc(_nbytes(shape, dtype))
+    if base == 0:
+        raise RuntimeError("rs_malloc returned NULL (symmetric heap exhausted?)")
+    t = _wrap(base, shape, dtype, torch.cuda.current_device())
+    _allocations.append(t)
+    return t
+
+
+def free_tensor(tensor):
+    # rocSHMEM host binding exposes no per-allocation free; the whole symmetric
+    # heap is released by rs_finalize(). No-op here (buffers live for the run).
+    pass
+
+
+def barrier():
+    _lib.rs_barrier()
+
+
+# ---------------------------------------------------------------------------
+# GPU-direct (GDA) device-kernel launchers (only valid when gda_enabled()).
+# Pointers are raw device addresses (int) into the GDA symmetric heap.
+# ---------------------------------------------------------------------------
+def gda_gather(target_ptr, src_ptr, nbytes, peers, stride_bytes):
+    """Device gather: for each cross-node peer in ``peers`` (global PE/rank), pull
+    its shard (at symmetric ``src_ptr``) into ``target_ptr + peer*stride_bytes``."""
+    if len(peers) == 0:
+        return
+    if _ctypes_gda_lib:
+        n = len(peers)
+        arr = (ctypes.c_int * n)(*[int(p) for p in peers])
+        rc = _lib.gda_gather(int(target_ptr), int(src_ptr), int(nbytes), arr, n, int(stride_bytes))
+    else:
+        # pybind binds gda_gather(target, src, nbytes, std::vector<int> peers, stride)
+        rc = _lib.gda_gather(
+            int(target_ptr), int(src_ptr), int(nbytes), [int(p) for p in peers], int(stride_bytes)
+        )
+    if rc != 0:
+        raise RuntimeError(f"gda_gather hipError={rc}")
+
+
+def gda_reduce_scatter_acc(
+    acc_ptr,
+    input_ptr,
+    seg_off_bytes,
+    shard_elems,
+    n_pes,
+    scratch_ptr,
+    scratch_stride_bytes,
+    dtype_code,
+    nblocks,
+):
+    """Device pull-based reduce-scatter accumulate: acc_fp32[i] += sum over all
+    PEs of input[seg_off + i] (pulled from each PE). Race-free (on-chip sum)."""
+    rc = _lib.gda_reduce_scatter_acc(
+        int(acc_ptr),
+        int(input_ptr),
+        int(seg_off_bytes),
+        int(shard_elems),
+        int(n_pes),
+        int(scratch_ptr),
+        int(scratch_stride_bytes),
+        int(dtype_code),
+        int(nblocks),
+    )
+    if rc != 0:
+        raise RuntimeError(f"gda_reduce_scatter_acc hipError={rc}")
+
+
+def gda_stage_fence(dst_ptr, src_ptr, nbytes):
+    """Copy src->dst (device) and SYSTEM-fence so the staged symmetric buffer is
+    visible to the NIC before a remote PE's device getmem (HDP-flush substitute)."""
+    rc = _lib.gda_stage_fence(int(dst_ptr), int(src_ptr), int(nbytes))
+    if rc != 0:
+        raise RuntimeError(f"gda_stage_fence hipError={rc}")
+
+
+def gda_hdp_init():
+    """Resolve this rank's GPU HDP flush register (call after set_device).
+    Returns 0 on success; nonzero means the register could not be resolved."""
+    if not hasattr(_lib, "gda_hdp_init"):
+        return -99
+    return int(_lib.gda_hdp_init())
+
+
+def gda_hdp_flush():
+    """Flush this GPU's HDP cache so prior symmetric writes are NIC-visible (the
+    proper GPUDirect-RDMA write-visibility primitive). Returns 1 if flushed."""
+    return int(_lib.gda_hdp_flush())
+
+
+def gda_strided_touch(
+    input_ptr,
+    seg_off_bytes,
+    seg_bytes,
+    n_pes,
+    stride_bytes,
+    touch_bytes,
+    scratch_ptr,
+    scratch_stride_bytes,
+    nblocks,
+):
+    """Strided page-touch warm-up: a tiny RDMA read at every ``stride_bytes`` page
+    of my shard segment on every PE, priming all pages/NICs (deterministic
+    read-triggered settle) at minimal volume vs the full-shard throwaway RS."""
+    rc = _lib.gda_strided_touch(
+        int(input_ptr),
+        int(seg_off_bytes),
+        int(seg_bytes),
+        int(n_pes),
+        int(stride_bytes),
+        int(touch_bytes),
+        int(scratch_ptr),
+        int(scratch_stride_bytes),
+        int(nblocks),
+    )
+    if rc != 0:
+        raise RuntimeError(f"gda_strided_touch hipError={rc}")
+
+
+def gda_reduce_scatter_acc_async(
+    acc_ptr,
+    input_ptr,
+    seg_off_bytes,
+    shard_elems,
+    n_pes,
+    scratch_ptr,
+    scratch_stride_bytes,
+    dtype_code,
+    nblocks,
+):
+    """Comm/compute OVERLAP: launch reduce-scatter on a side stream WITHOUT syncing.
+    Caller must gda_rs_overlap_sync() before re-staging input or consuming acc."""
+    rc = _lib.gda_reduce_scatter_acc_async(
+        int(acc_ptr),
+        int(input_ptr),
+        int(seg_off_bytes),
+        int(shard_elems),
+        int(n_pes),
+        int(scratch_ptr),
+        int(scratch_stride_bytes),
+        int(dtype_code),
+        int(nblocks),
+    )
+    if rc != 0:
+        raise RuntimeError(f"gda_reduce_scatter_acc_async launch err={rc}")
+
+
+def gda_rs_overlap_sync():
+    """Wait for all pending overlapped reduce-scatter kernels on the side stream."""
+    rc = _lib.gda_rs_overlap_sync()
+    if rc != 0:
+        raise RuntimeError(f"gda_rs_overlap_sync hipError={rc}")
+
+
+def gda_gather_async(target_ptr, src_ptr, nbytes, peers, stride_bytes, stream):
+    """Approach 1: launch all-gather kernel on the given HIP `stream` WITHOUT syncing,
+    so FSDP2 prefetch overlaps it with compute. Reassembly + consumer order via the
+    same stream (gather reads stable params -> no settle/barrier needed)."""
+    if len(peers) == 0:
+        return
+    if _ctypes_gda_lib:
+        n = len(peers)
+        arr = (ctypes.c_int * n)(*[int(p) for p in peers])
+        rc = _lib.gda_gather_async(
+            int(target_ptr),
+            int(src_ptr),
+            int(nbytes),
+            arr,
+            n,
+            int(stride_bytes),
+            int(stream),
+        )
+    else:
+        rc = _lib.gda_gather_async(
+            int(target_ptr),
+            int(src_ptr),
+            int(nbytes),
+            [int(p) for p in peers],
+            int(stride_bytes),
+            int(stream),
+        )
+    if rc != 0:
+        raise RuntimeError(f"gda_gather_async launch err={rc}")
+
+
+def dtype_code(dtype):
+    return _DTYPE_CODE[dtype]
+
+
+def finalize():
+    global _initialized
+    if _lib is not None and _initialized:
+        _lib.rs_finalize()
+        _initialized = False
diff --git a/primus/core/odc/primitives/gather.py b/primus/core/odc/primitives/gather.py
new file mode 100644
index 000000000..752da1659
--- /dev/null
+++ b/primus/core/odc/primitives/gather.py
@@ -0,0 +1,244 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+import logging
+import math
+
+import torch
+import torch.distributed as dist
+import triton
+import triton.language as tl
+from torch import Tensor
+
+from odc.primitives import (
+    SHMEM_EXTERN_LIBS,
+    __syncthreads,
+    getmem_nbi_block,
+    quiet,
+    tid,
+)
+from odc.primitives.utils import (
+    BufferSplitter,
+    SymmBufferRegistry,
+    get_comm_stream,
+    get_local_world_size,
+    sync_cta,
+)
+
+logger = logging.getLogger(__name__)
+
+from odc.primitives.utils import _USE_ROCSHMEM  # noqa: E402
+
+if _USE_ROCSHMEM:
+    from odc.primitives import _rocshmem_backend as _rs
+else:
+    _rs = None
+
+
+def _gda_active():
+    return _rs is not None and _rs.gda_enabled()
+
+
+@triton.jit
+def shmem_device_producer_gather_2d_get_block_kernel_chunked_synced(
+    remote_tensor_ptr,
+    target_tensor_ptr,
+    elem_per_rank,
+    size_per_elem,
+    rank,
+    num_ranks_per_node,
+    world_size,
+    chunk_size,
+    signal_ptr,
+):
+    pid = tl.program_id(axis=0)
+    # np = tl.num_programs(axis=0)
+    assert num_ranks_per_node == tl.num_programs(axis=0)
+    np = num_ranks_per_node
+    num_nodes = world_size // np
+
+    tidx = tid(axis=0)
+    expected = 0
+    for i in range(1, num_nodes):
+        peer_node = (i + rank // np) % num_nodes
+        peer = (pid + peer_node * np) % world_size
+        # chunk_size = elem_per_rank // num_chunks
+        num_chunks = tl.cdiv(elem_per_rank, chunk_size)
+        for chunk in range(num_chunks):
+            this_chunk_size = chunk_size
+            if chunk == num_chunks - 1:
+                this_chunk_size = elem_per_rank - chunk * chunk_size
+            getmem_nbi_block(
+                target_tensor_ptr + peer * elem_per_rank + (chunk * chunk_size),
+                remote_tensor_ptr + (chunk * chunk_size),
+                this_chunk_size * size_per_elem,
+                peer,
+            )
+            expected += np
+            sync_cta(signal_ptr, expected)
+            if tidx == 0 and pid == 0:
+                quiet()
+            __syncthreads()
+
+            expected += np
+            sync_cta(signal_ptr, expected)
+
+
+class GatherService:
+    def __init__(self):
+        self.shaped_buffer = {}
+        self.buffer_splitter = BufferSplitter()
+        self.chunk_size_bytes = 2**20
+
+    def get_chunk_size(self, buffer_dtype):
+        return self.chunk_size_bytes // buffer_dtype.itemsize
+
+    def gather_into_tensor(self, output_tensor: Tensor, input_tensor: Tensor, pg: dist.ProcessGroup):
+        buf_size = self.buffer_splitter.get_global_buffer_size(output_tensor.shape)
+        buffer_shape = (buf_size,)
+        output_size = output_tensor.numel()
+        assert output_size >= buf_size, f"output_size: {output_size} < buf_size: {buf_size}"
+
+        rank = torch.distributed.get_rank()
+        if (buffer_shape, output_tensor.dtype) not in self.shaped_buffer:
+            logger.info(
+                f"Rank {rank} create buffer: output_size: {output_size} num_sub_buffers: {math.ceil(output_size / buf_size)} buf_size: {buf_size}"
+            )
+            self.shaped_buffer[
+                (buffer_shape, output_tensor.dtype)
+            ] = SymmBufferRegistry.get_instance().allocate_symm_buffer(
+                f"ag_buffer_{buffer_shape}_{output_tensor.dtype}",
+                buffer_shape,
+                output_tensor.dtype,
+            )
+        target_tensor = self.shaped_buffer[(buffer_shape, output_tensor.dtype)]
+
+        assert (input_tensor.numel() * input_tensor.element_size()) % (
+            2**6
+        ) == 0 or input_tensor.numel() < 2**6, "better align to 64 for efficiency"
+        chunk_size = self.get_chunk_size(input_tensor.dtype)
+        # assert input_tensor.numel() % chunk_size == 0
+
+        registry = SymmBufferRegistry.get_instance()
+        peer_tensors = registry.get_peer_tensors(input_tensor)
+
+        group_world_size = torch.distributed.get_world_size(pg)
+        local_world_size = get_local_world_size()
+        assert group_world_size in (
+            torch.distributed.get_world_size(),
+            local_world_size,
+        ), f"{group_world_size=} {torch.distributed.get_world_size()=} {local_world_size=}"
+
+        get_comm_stream().wait_stream(torch.cuda.current_stream())
+        # GPU-direct (GDA): IPC is off, so there are no XGMI peer views; ALL ranks
+        # (incl. same-node and self) are pulled via device rocshmem_getmem. Handle
+        # the whole cross-group gather here and return.
+        gda = _gda_active() and local_world_size != group_world_size
+        if gda:
+            self._gda_gather_into_tensor(
+                output_tensor,
+                input_tensor,
+                target_tensor,
+                buf_size,
+                group_world_size,
+                rank,
+            )
+            return
+
+        with torch.cuda.stream(get_comm_stream()):
+            output_tensor_split = output_tensor.view(group_world_size, -1)
+            assert local_world_size == len(peer_tensors)
+            local_rank = rank % local_world_size
+            rank_same_node_start = rank - local_rank
+            rank_same_node_end = rank_same_node_start + local_world_size
+            for r_offset in range(local_world_size):
+                src_local_rank = (local_rank + r_offset) % local_world_size
+                if group_world_size == local_world_size:
+                    output_tensor_split[src_local_rank].copy_(peer_tensors[src_local_rank])
+                else:
+                    src_rank = rank_same_node_start + src_local_rank
+                    output_tensor_split[src_rank].copy_(peer_tensors[src_local_rank])
+
+            assert buf_size % group_world_size == 0
+            local_buf_size = buf_size // group_world_size
+            signal_ptr = torch.empty(1, dtype=torch.int32, device="cuda")
+            for start in range(0, input_tensor.numel(), local_buf_size):
+                if local_world_size == group_world_size:
+                    continue
+                size = min(local_buf_size, input_tensor.numel() - start)
+                sub_input_tensor = input_tensor.view(-1)[start : start + size]
+                assert (sub_input_tensor.numel() * sub_input_tensor.element_size()) % (
+                    2**6
+                ) == 0 or sub_input_tensor.numel() < 2**6, "better align to 64 for efficiency"
+                target_buf_size = size * group_world_size
+                assert target_buf_size <= buf_size
+                target_tensor_split = target_tensor[:target_buf_size].view(group_world_size, size)
+
+                signal_ptr.fill_(0)
+                assert group_world_size % 8 == 0 or group_world_size < 8
+                # grid_size = 8 if world_size == 32 else world_size
+                grid_size = local_world_size
+                shmem_device_producer_gather_2d_get_block_kernel_chunked_synced[(grid_size,)](
+                    remote_tensor_ptr=sub_input_tensor,
+                    target_tensor_ptr=target_tensor_split.view(-1),
+                    elem_per_rank=sub_input_tensor.numel(),
+                    size_per_elem=sub_input_tensor.element_size(),
+                    rank=rank,
+                    num_ranks_per_node=local_world_size,
+                    world_size=group_world_size,
+                    chunk_size=chunk_size,
+                    signal_ptr=signal_ptr,
+                    num_warps=32,
+                    extern_libs=SHMEM_EXTERN_LIBS,
+                )
+                if buf_size == output_size:
+                    local_world_data_size = size * local_world_size
+                    local_world_idx = rank // local_world_size
+                    data_start_idx = local_world_data_size * local_world_idx
+                    data_end_idx = data_start_idx + local_world_data_size
+                    output_tensor[:data_start_idx].copy_(target_tensor[:data_start_idx])
+                    output_tensor[data_end_idx:].copy_(target_tensor[data_end_idx:])
+                    # output_tensor.copy_(target_tensor)
+                else:
+                    for r in range(group_world_size):
+                        if rank_same_node_start <= r < rank_same_node_end:
+                            continue
+                        output_tensor_split[r, start : start + size].copy_(target_tensor_split[r, :])
+        torch.cuda.current_stream().wait_stream(get_comm_stream())
+
+    def _gda_gather_into_tensor(
+        self, output_tensor, input_tensor, target_tensor, buf_size, group_world_size, rank
+    ):
+        """GPU-direct all-gather: every rank's shard (incl same-node and self) is
+        pulled via a single device rocshmem_getmem kernel into target slot[r],
+        then copied into output_tensor_split[r]. No XGMI peer views / host loop."""
+        es = input_tensor.element_size()
+        assert buf_size % group_world_size == 0
+        local_buf_size = buf_size // group_world_size
+        output_split = output_tensor.view(group_world_size, -1)
+        peers = list(range(group_world_size))
+        comm = get_comm_stream()
+        comm.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(comm):
+            torch.cuda.synchronize()
+            # NO per-gather rocshmem_barrier_all here. Gather reads PARAMS (symmetric
+            # shard written by the optimizer a full step earlier and read-only through
+            # fwd/bwd) -> peers' data is already stable/visible, so the cross-PE
+            # rendezvous is unnecessary (unlike scatter's just-written staging).
+            for start in range(0, input_tensor.numel(), local_buf_size):
+                size = min(local_buf_size, input_tensor.numel() - start)
+                sub_input = input_tensor.view(-1)[start : start + size]
+                tb = size * group_world_size
+                tsplit = target_tensor[:tb].view(group_world_size, size)
+                _rs.gda_gather(target_tensor.data_ptr(), sub_input.data_ptr(), size * es, peers, size * es)
+                torch.cuda.synchronize()
+                # reassembly on the SAME comm stream -> ordered AFTER the gather kernel
+                for r in range(group_world_size):
+                    output_split[r, start : start + size].copy_(tsplit[r])
+        torch.cuda.current_stream().wait_stream(comm)
diff --git a/primus/core/odc/primitives/scatter_accumulate.py b/primus/core/odc/primitives/scatter_accumulate.py
new file mode 100644
index 000000000..1dc3ca7cf
--- /dev/null
+++ b/primus/core/odc/primitives/scatter_accumulate.py
@@ -0,0 +1,431 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+import logging
+
+import torch
+import torch.distributed as dist
+
+from odc.primitives.utils import (
+    SymmBufferRegistry,
+    get_comm_stream,
+    get_local_world_size,
+)
+from odc.runtime_config import get_config
+
+logger = logging.getLogger(__name__)
+
+from odc.primitives.utils import _USE_ROCSHMEM  # noqa: E402
+
+if _USE_ROCSHMEM:
+    from odc.primitives import _rocshmem_backend as _rs
+else:
+    _rs = None
+
+
+def _gda_active():
+    return _rs is not None and _rs.gda_enabled()
+
+
+class ReductionService:
+    """Reduce-scatter-accumulate service (device-side, no host-polling subprocess).
+
+    Two paths, both device-side and free of any host-polling subprocess:
+      * single node  -> owner-side XGMI PULL + on-chip fp32 sum
+        (``_single_device_scatter_accumulate``).
+      * multi node    -> rocSHMEM GPU-direct (GDA) reduce-scatter
+        (``_gda_scatter_accumulate``).
+
+    Both DEFER the collective reduce to once per param-group at
+    ``get_accumulation`` (matched barrier count across ranks -> deadlock-free
+    under variable-length / nopad micro-batching), pre-accumulating each
+    micro-batch's grad locally in fp32.
+    """
+
+    def __init__(self, accumulation_dtype=None):
+        self.accumulations = []
+        self.accumulation_indices = {}
+        self.input_buffer = {}
+        self.dispatched_tasks = 0
+        # Accepted for API compatibility with the FSDP1 caller; the device and
+        # GDA reduce paths always accumulate in fp32.
+        self.accumulation_dtype = accumulation_dtype
+
+    def clear_accumulations(self):
+        for acc in self.accumulations:
+            acc.fill_(0)
+        if hasattr(self, "_gda_deferred"):  # reset per-minibatch deferred grads
+            self._gda_deferred = {}
+            self._gda_deferred_pg = {}
+        if hasattr(self, "_sdr_deferred"):  # reset per-minibatch deferred grads
+            self._sdr_deferred = {}
+            self._sdr_deferred_pg = {}
+
+    def _gda_scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup):
+        """GPU-direct pull-based reduce-scatter accumulate (race-free, no host-polling subprocess).
+
+        Stages this rank's full input into a symmetric buffer, barriers, then a
+        device kernel pulls every PE's contribution to MY output shard and sums
+        it on-chip into the fp32 accumulation buffer (acc += reduce_scatter(input)).
+        """
+        gws = torch.distributed.get_world_size(pg)
+        assert (
+            gws == _rs._n_pes
+        ), f"GDA reduce-scatter requires a full-world group: gws={gws} n_pes={_rs._n_pes}"
+        assert input_tensor.numel() % gws == 0, f"{input_tensor.numel()=} % {gws=}"
+        shard_elems = input_tensor.numel() // gws
+        dt = input_tensor.dtype
+        es = input_tensor.element_size()
+        reg = SymmBufferRegistry.get_instance()
+
+        if key not in self.accumulation_indices:
+            acc = reg.get_or_create_symm_buffer(f"gda_acc_{key}", (shard_elems,), torch.float32)
+            acc.fill_(0)
+            self.accumulation_indices[key] = len(self.accumulations)
+            self.accumulations.append(acc)
+        acc = self.accumulations[self.accumulation_indices[key]]
+
+        in_key = ("gda_in", dt, input_tensor.numel())
+        if in_key not in self.input_buffer:
+            self.input_buffer[in_key] = reg.get_or_create_symm_buffer(
+                f"gda_in_{dt}_{input_tensor.numel()}", (input_tensor.numel(),), dt
+            )
+        input_sym = self.input_buffer[in_key]
+
+        cfg = get_config()
+        # GRID geometry (config odc_gda_rs_blocks): reduce-scatter kernel grid (one
+        # block per disjoint shard chunk -> # concurrent cross-node getmem_wg = QP/NIC
+        # parallelism lever). Scratch auto-resizes (sc_key includes chunk*nblk), so
+        # this stays correct for any nblk. Default 64 (current behavior).
+        nblk = int(cfg.gda_rs_blocks)
+        if nblk < 1:
+            nblk = 1
+        # PIPE (config odc_gda_pipe): peer-pipeline batch depth. The pipelined rs_acc
+        # needs `pipe` scratch slots PER BLOCK (issues `pipe` peers' nbi getmem
+        # concurrently), so scratch grows pipe x and the main call passes
+        # scratch_stride = pipe*chunk. This is also a Primus-Turbo device-kernel knob
+        # (odc_rocshmem_gda.cu env_pipe); set_runtime_config bridges odc_gda_pipe back
+        # to the PRIMUS_TURBO_ODC_GDA_PIPE env var so the C++ side agrees.
+        pipe = int(cfg.gda_pipe)
+        if pipe < 1:
+            pipe = 1
+        chunk = (shard_elems + nblk - 1) // nblk
+        sc_slots = nblk * pipe
+        sc_key = ("gda_scr", dt, chunk * sc_slots)
+        if sc_key not in self.input_buffer:
+            self.input_buffer[sc_key] = reg.get_or_create_symm_buffer(
+                f"gda_scr_{dt}_{chunk * sc_slots}", (chunk * sc_slots,), dt
+            )
+        scratch = self.input_buffer[sc_key]
+
+        rank = torch.distributed.get_rank(pg)
+
+        # Cross-node write-visibility strategy for the just-staged grad
+        # (config odc_gda_warmup_mode). Modes:
+        #   "strided" (default) - page-strided tiny throwaway READ that keeps
+        #       full-warmup's deterministic "read-triggered settle" (validated 0 grad
+        #       spikes, loss == single-node) at minimal volume; stride via
+        #       odc_gda_stride_bytes (default 64KB).
+        #   "full"  - full-shard throwaway reduce-scatter: most robust, slowest.
+        #   "hdp"   - O(1) HDP_MEM_FLUSH_CNTL register write; auto-falls back to
+        #       "full" if no HDP register. "fence"/"hdpfence" - device system-scope
+        #       fence (+ optional HDP flush) instead of the read settle.
+        _warm_mode = cfg.gda_warmup_mode
+        if _warm_mode == "hdp" and getattr(self, "_hdp_fallback", False):
+            _warm_mode = "full"
+        if not getattr(self, "_logged_warm_mode", False):
+            logger.warning(
+                "[GDA] reduce-scatter warm-up mode=%s stride_bytes=%s",
+                _warm_mode,
+                cfg.gda_stride_bytes,
+            )
+            self._logged_warm_mode = True
+        if _warm_mode in ("fence", "hdpfence"):
+            # Ensure the grad (input_tensor) is fully produced before copy+fence,
+            # then copy into the symmetric staging buffer with a trailing system-scope
+            # fence (gda_stage_fence self-syncs the device). "hdpfence" additionally
+            # flushes this GPU's HDP so the remote NIC's RDMA read sees fresh data.
+            torch.cuda.current_stream().synchronize()
+            _flat = input_tensor.view(-1)
+            _rs.gda_stage_fence(input_sym.data_ptr(), _flat.data_ptr(), _flat.numel() * es)
+            if _warm_mode == "hdpfence":
+                if not getattr(self, "_hdp_inited", False) and not getattr(self, "_hdp_fallback", False):
+                    rc = _rs.gda_hdp_init()
+                    if rc == 0:
+                        self._hdp_inited = True
+                    else:
+                        logger.warning("gda_hdp_init failed rc=%d -> hdpfence falls back to fence-only", rc)
+                        self._hdp_fallback = True
+                if getattr(self, "_hdp_inited", False):
+                    _rs.gda_hdp_flush()
+        else:
+            get_comm_stream().wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(get_comm_stream()):
+                input_sym.copy_(input_tensor.view(-1))
+            torch.cuda.synchronize()
+            if _warm_mode == "hdp":
+                if not getattr(self, "_hdp_inited", False):
+                    rc = _rs.gda_hdp_init()
+                    if rc != 0:
+                        logger.warning(
+                            "gda_hdp_init failed rc=%d -> falling back to full-shard warm-up RS", rc
+                        )
+                        self._hdp_fallback = True
+                        _warm_mode = "full"
+                    else:
+                        self._hdp_inited = True
+                if _warm_mode == "hdp":
+                    _rs.gda_hdp_flush()
+        _rs.barrier()
+
+        # Cross-node write-visibility settle (FIX for the intermittent stale-read
+        # grad spikes): a throwaway "warm-up" reduce-scatter + barrier BEFORE the
+        # real one. A single barrier after staging is insufficient on this
+        # mlx5/GDA path (peer's just-staged GPU write isn't yet NIC-visible to the
+        # first device getmem -> stale read -> huge wrong gradient ~half the time).
+        # The warm-up getmem + barrier forces the staged data visible; the real
+        # reduce-scatter then reads fresh data. (Verified: eliminates spikes,
+        # grad norms normal, loss matches single-node. Reduce-scatter is a small
+        # fraction of the step, so the extra pass is cheap vs the gather.)
+        # Opt1: the warm-up only needs to settle cross-node write-visibility (drain
+        # the HDP so peers' staged writes are NIC-visible), NOT move the full
+        # shard. A tiny getmem touch per peer + barrier achieves that at a
+        # fraction of the cost (full-shard warm-up was ~9s/iter of pure overhead).
+        # warm-up settles cross-node write-visibility (HDP) before the real RS.
+        # Single-NIC: a 1024-elem touch suffices. Multi-NIC: each peer routes via a
+        # different NIC and HDP is per-NIC/per-page, so a tiny touch leaves most
+        # pages stale -> spikes; use the FULL shard (default) so every NIC/page is
+        # settled. The full warm-up is itself parallelized across NICs (cheap).
+        if _warm_mode in ("fence", "hdp", "hdpfence"):
+            # The staged write was already made NIC-visible (HDP flush / fence),
+            # so the throwaway warm-up reduce-scatter (the ~59%-of-RS overhead) is
+            # skipped entirely. The barrier after staging still rendezvouses PEs.
+            pass
+        elif _warm_mode == "strided":
+            # Page-strided tiny throwaway READ: keeps full-warmup's deterministic
+            # "read-triggered settle" (covers every page of my segment on every PE
+            # -> all 8 NICs) but at minimal volume (one touch per odc_gda_stride_bytes
+            # page, not the whole shard). Staging above used plain copy_ (no flush).
+            stride_b = int(cfg.gda_stride_bytes)
+            touch_b = int(es)
+            seg_bytes = int(shard_elems * es)
+            npages = (seg_bytes + stride_b - 1) // stride_b
+            total_touch = _rs._n_pes * max(npages, 1)
+            sstride = 256  # bytes/throwaway scratch slot (>= touch_b; avoids collide)
+            scratch_cap = scratch.numel() * scratch.element_size()
+            nblk_t = min(int(total_touch), 4096, max(1, scratch_cap // sstride))
+            _rs.gda_strided_touch(
+                input_sym.data_ptr(),
+                rank * shard_elems * es,
+                seg_bytes,
+                _rs._n_pes,
+                stride_b,
+                touch_b,
+                scratch.data_ptr(),
+                sstride,
+                int(nblk_t),
+            )
+            _rs.barrier()
+        else:  # "full": full-shard throwaway reduce-scatter settle
+            n_warm, w_nblk, w_stride = int(shard_elems), nblk, chunk * es
+            wkey = ("gda_warmup", n_warm)
+            if wkey not in self.input_buffer:
+                self.input_buffer[wkey] = torch.zeros(n_warm, dtype=torch.float32, device="cuda")
+            warmup = self.input_buffer[wkey]
+            warmup.zero_()
+            _rs.gda_reduce_scatter_acc(
+                warmup.data_ptr(),
+                input_sym.data_ptr(),
+                rank * shard_elems * es,
+                n_warm,
+                _rs._n_pes,
+                scratch.data_ptr(),
+                w_stride,
+                _rs.dtype_code(dt),
+                w_nblk,
+            )
+            _rs.barrier()
+        _rs.gda_reduce_scatter_acc(
+            acc.data_ptr(),
+            input_sym.data_ptr(),
+            rank * shard_elems * es,
+            shard_elems,
+            _rs._n_pes,
+            scratch.data_ptr(),
+            pipe * chunk * es,
+            _rs.dtype_code(dt),
+            nblk,
+        )
+        torch.cuda.synchronize()
+        self.dispatched_tasks += 1
+
+    def _single_device_scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup):
+        """Single-node device-side reduce-scatter accumulate (no host-polling subprocess).
+
+        Mechanism (owner-side PULL + on-chip fp32 sum over same-node XGMI peer
+        views): each PE stages its (locally pre-accumulated) full grad into a
+        symmetric fp32 buffer whose same-node peer views are already resolved on
+        allocation (the same XGMI peer-view machinery gather.py uses); after a
+        rendezvous barrier, PE r reads every same-node peer's segment destined
+        for its shard (a plain XGMI ``.copy_`` peer read) and sums it on-chip.
+        There are NO cross-rank writes -> no device atomics -> no MI300X
+        write-visibility hazard; only XGMI reads gated by a barrier. No second
+        process -> no IPC handle and no host-side reduction subprocess.
+
+        ``input_tensor`` is this PE's full (locally pre-accumulated) grad, laid
+        out as ``[shard_0 | shard_1 | ... | shard_{gws-1}]``; PE r owns output
+        shard r.
+
+        Steps (all same-node, no subprocess, no IPC handle, no atomics):
+          1. Stage my full grad into a symmetric fp32 buffer (peer views resolved
+             on allocation).
+          2. cuda sync + collective barrier -> every peer's stage is retired and
+             visible before any peer reads it over XGMI.
+          3. For each same-node peer p, XGMI-``.copy_`` p's segment destined for
+             MY shard into a local temp, then acc += temp (fp32 on-chip sum).
+          4. cuda sync + collective barrier -> all reads done before the shared
+             staging buffer can be reused by the next key/minibatch.
+        """
+        gws = torch.distributed.get_world_size(pg)
+        lws = get_local_world_size()
+        assert gws == lws, f"single-device reduce is same-node only: gws={gws} lws={lws}"
+        assert input_tensor.numel() % gws == 0, f"{input_tensor.numel()=} % {gws=}"
+        shard_elems = input_tensor.numel() // gws
+        reg = SymmBufferRegistry.get_instance()
+        rank = torch.distributed.get_rank(pg)  # single node -> group rank == local pos
+
+        # fp32 accumulator = MY output shard (matches GDA's fp32 acc semantics).
+        if key not in self.accumulation_indices:
+            acc = reg.get_or_create_symm_buffer(f"sdr_acc_{key}", (shard_elems,), torch.float32)
+            acc.fill_(0)
+            self.accumulation_indices[key] = len(self.accumulations)
+            self.accumulations.append(acc)
+        acc = self.accumulations[self.accumulation_indices[key]]
+
+        # Symmetric fp32 staging for my full grad, WITH same-node peer views. One
+        # per grad-numel; reused across keys/minibatches (guarded by the trailing
+        # barrier so no peer reuses it mid-read).
+        in_key = ("sdr_in", input_tensor.numel())
+        if in_key not in self.input_buffer:
+            self.input_buffer[in_key] = reg.get_or_create_symm_buffer(
+                f"sdr_in_{input_tensor.numel()}", (input_tensor.numel(),), torch.float32
+            )
+        input_sym = self.input_buffer[in_key]
+        peer_inputs = reg.get_peer_tensors(input_sym)
+        assert len(peer_inputs) == lws
+
+        tmp_key = ("sdr_tmp", shard_elems)
+        if tmp_key not in self.input_buffer:
+            self.input_buffer[tmp_key] = torch.empty(shard_elems, dtype=torch.float32, device="cuda")
+        tmp = self.input_buffer[tmp_key]
+
+        input_sym.copy_(input_tensor.view(-1).to(torch.float32))
+        torch.cuda.synchronize()
+        torch.distributed.barrier(group=pg)
+
+        lo = rank * shard_elems
+        hi = lo + shard_elems
+        # Rotate the peer order by local rank so all 8 PEs don't hammer the same
+        # peer's HBM first (matches the gather/scatter round-robin peer ordering).
+        for off in range(lws):
+            p = (rank + off) % lws
+            tmp.copy_(peer_inputs[p][lo:hi])
+            acc.add_(tmp)
+        torch.cuda.synchronize()
+        torch.distributed.barrier(group=pg)
+        self.dispatched_tasks += 1
+
+    def scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup):
+        if _gda_active():
+            # DEFER the cross-node reduce to once-per-minibatch. The per-call
+            # _gda_scatter_accumulate does rocshmem_barrier_all (collective);
+            # calling it per-microbatch deadlocks under nopad (ranks have
+            # different micro-batch counts -> mismatched barrier counts). Fix:
+            # accumulate the unsharded grad LOCALLY here (no comm/barrier ->
+            # backward is lockstep-free), then do ONE barriered reduce-scatter
+            # per group at get_accumulation (count == #groups, matched across
+            # ranks). Config odc_gda_defer_reduce: "auto" (default) => defer iff
+            # multi-node (n_pes > local_world_size); "1"/"0" force on/off.
+            _defer = str(get_config().gda_defer_reduce)
+            if _defer == "auto":
+                _defer_on = _rs._n_pes > get_local_world_size()
+            else:
+                _defer_on = _defer in ("1", "true", "True")
+            if _defer_on:
+                if not hasattr(self, "_gda_deferred"):
+                    self._gda_deferred = {}
+                    self._gda_deferred_pg = {}
+                cur = self._gda_deferred.get(key)
+                if cur is None:
+                    self._gda_deferred[key] = input_tensor.detach().clone()
+                else:
+                    cur.add_(input_tensor)
+                self._gda_deferred_pg[key] = pg
+                return
+            return self._gda_scatter_accumulate(key, input_tensor, pg)
+        # Single-node device-side reduce (replaces the removed host-side
+        # reduction subprocess path). Cross-node reduce always goes through the
+        # rocSHMEM GDA path above.
+        assert torch.distributed.get_world_size(pg) == get_local_world_size(), (
+            f"non-GDA reduce-scatter is single-node only "
+            f"(gws={torch.distributed.get_world_size(pg)} lws={get_local_world_size()}); "
+            f"multi-node requires the rocSHMEM GDA backend"
+        )
+        # DEFER exactly like the multi-node GDA path: pre-accumulate this
+        # micro-batch's grad LOCALLY (no collective -> nopad-safe), then run ONE
+        # barriered owner-side pull-sum per group at get_accumulation.
+        # Pre-accumulate in fp32 so cross-micro-batch accumulation matches the
+        # fp32 device accumulator.
+        if not hasattr(self, "_sdr_deferred"):
+            self._sdr_deferred = {}
+            self._sdr_deferred_pg = {}
+        cur = self._sdr_deferred.get(key)
+        if cur is None:
+            self._sdr_deferred[key] = input_tensor.detach().to(torch.float32)
+        else:
+            cur.add_(input_tensor)
+        self._sdr_deferred_pg[key] = pg
+        return
+
+    def get_accumulation(self, key):
+        # GDA DEFER: run the single per-minibatch cross-node reduce-scatter now
+        # (once per group, matched barrier count across ranks -> deadlock-free).
+        if _gda_active() and getattr(self, "_gda_deferred", None) is not None:
+            pending = self._gda_deferred.get(key)
+            if pending is not None:
+                self._gda_deferred[key] = None
+                self._gda_scatter_accumulate(key, pending, self._gda_deferred_pg[key])
+        # Single-node device DEFER: run the single per-minibatch owner-side
+        # pull-sum now (once per group, matched barrier count across ranks).
+        if not _gda_active() and getattr(self, "_sdr_deferred", None) is not None:
+            pending = self._sdr_deferred.get(key)
+            if pending is not None:
+                self._sdr_deferred[key] = None
+                self._single_device_scatter_accumulate(key, pending, self._sdr_deferred_pg[key])
+        acc = self.accumulations[self.accumulation_indices[key]]
+        return acc
+
+    def sync(self, pg: dist.ProcessGroup):
+        if _gda_active():
+            # GDA path is fully synchronous (device kernels + barriers).
+            torch.cuda.synchronize()
+            torch.distributed.barrier(group=pg)
+            self.dispatched_tasks = 0
+            return
+        # Single-node device path: the owner-side pull-sum already ran
+        # (synchronously, with its own barriers) inside get_accumulation. Just a
+        # final cuda sync + barrier before the optimizer reads acc.
+        torch.cuda.synchronize()
+        torch.distributed.barrier(group=pg)
+        self.dispatched_tasks = 0
+
+    def stop(self):
+        # The device and GDA reduce paths run no host-polling subprocess, so there is
+        # nothing to tear down. Kept for API compatibility with the callers.
+        pass
diff --git a/primus/core/odc/primitives/shmem_triton.py b/primus/core/odc/primitives/shmem_triton.py
new file mode 100644
index 000000000..2747b2158
--- /dev/null
+++ b/primus/core/odc/primitives/shmem_triton.py
@@ -0,0 +1,465 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), whose upstream Triton
+# device API (odc/primitives/nvshmem_triton.py) this ROCm reimplementation is
+# derived from. ODC is distributed under the MIT License per its package
+# metadata (pyproject.toml / setup.py classifiers); the upstream repository
+# ships no LICENSE file or per-file copyright headers, and upstream copyright is
+# held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+"""
+ODC Triton device API — ROCm implementation.
+
+ODC's Triton kernels reach symmetric peer memory through a small set of
+device primitives. On ROCm there are two backends, selected by the
+``ODC_P2P_BACKEND`` environment variable:
+
+* ``mori`` (default): re-exports MORI-SHMEM's device API. MORI provides the
+  device bitcode (``libmori_shmem_device.bc``, JIT-compiled by
+  ``mori.ir.find_bitcode()``) that is linked into the Triton kernels.
+* ``rocshmem``: a host-API rocSHMEM backend (single-node / IPC-only) that
+  links NO device bitcode — every P2P op is plain Triton to an XGMI-mapped
+  peer address.
+
+Module-level public names (stable across backends):
+    putmem_nbi_block, getmem_nbi_block, quiet, int_p, int_g,
+    int_atomic_compare_swap, int_atomic_swap,
+    tid, __syncthreads,
+    LIB_SHMEM_PATH, SHMEM_EXTERN_LIBS
+
+Design notes
+------------
+1. ``putmem_nbi_block``, ``getmem_nbi_block``, ``int32_p`` and other MORI
+   device functions take an *extra* ``qp`` (queue-pair index) argument. The
+   ODC wrappers below thread a fixed ``qp=0`` inside ``@triton.jit``-decorated
+   thin shims so callers in ODC need not change.
+2. MORI has no single-int blocking get, so ``int_g`` is implemented via a
+   4-byte ``getmem_thread`` into a scratch register.
+3. ``tid(axis=0)`` was used by ``utils.sync_cta`` to elect a "first thread"
+   in a wave that performs an atomic_add. ``llvm.nvvm.read.ptx.sreg.tid.x``
+   does not exist on AMDGPU. We provide a Triton-native alternative: an
+   ``arange(0, 1)`` lane vector with mask ``offsets == 0`` restricts execution
+   to lane 0, so ``tid`` is implemented as a constant that the caller compares
+   against 0.
+4. ``__syncthreads`` is implemented via Triton's IR barrier
+   (``create_barrier``), which lowers to ``s_barrier`` on AMDGPU.
+5. ``SHMEM_EXTERN_LIBS`` is ``mori.ir.triton.get_extern_libs()`` which locates
+   the JIT-compiled MORI device bitcode for the current GPU + NIC. The Triton
+   compile hook ``install_hook()`` is called at import time.
+"""
+
+import triton
+import triton.language as tl
+from packaging import version
+from triton.language import core
+
+# ---------------------------------------------------------------------------
+# ODC P2P backend selector (config item odc_p2p_backend). Default "mori" keeps
+# the existing, verified behaviour byte-for-byte; "rocshmem" selects the host-API
+# backend (single-node / IPC-only) defined in the dedicated branch below. The ODC
+# integration patch populates the runtime config before this module is imported.
+# ---------------------------------------------------------------------------
+from odc.runtime_config import get_config  # noqa: E402
+
+_P2P_BACKEND = get_config().p2p_backend.lower()
+
+
+def is_triton_version_supported():
+    return version.parse(triton.__version__) >= version.parse("3.4.0")
+
+
+# ===========================================================================
+# Branch A: rocSHMEM host-API backend (single-node / IPC-only).
+#
+# No device bitcode is linked into Triton (SHMEM_EXTERN_LIBS == {}); every
+# P2P op is plain Triton to an XGMI-mapped peer address. A LOCAL symmetric
+# address is translated to a peer address by adding a per-PE affine byte delta
+# (rocshmem_ptr is affine — empirically verified). The delta table is populated
+# once at init by _rocshmem_backend.alloc_peer_tensors (before any kernel
+# launch) via set_rocshmem_peer_deltas() below.
+#
+# REGRESSION GUARD (mirrors the mori branch): no RDMA / device-library symbol
+# ever enters these kernels — they are pure Triton. Use a per-rank/per-run
+# TRITON_CACHE_DIR so a rank never loads another rank's cached kernel (the
+# baked deltas differ per process).
+# ===========================================================================
+if _P2P_BACKEND == "rocshmem":
+    # Empty extern-libs: Triton compiles these kernels with NO device bitcode.
+    SHMEM_EXTERN_LIBS = {}
+    LIB_NAME = "librocshmem_host"
+    LIB_SHMEM_PATH = ""
+
+    # Per-PE affine pointer deltas (bytes), indexed by LOCAL position
+    # (global pe - node_start), baked into the int_p / int_g /
+    # int_wait_until_equals kernels at compile time. Triton only lets @jit
+    # functions read globals that are tl.constexpr, so the deltas are stored as
+    # constexpr (their values are captured AND versioned in the compile cache
+    # key per the used_global_vals mechanism -> a rank never reuses another
+    # rank's kernel even on a shared TRITON_CACHE_DIR). ``_RS_NODE_START`` is the
+    # first global rank on this node so a runtime global pe maps to a local slot.
+    _RS_D0 = tl.constexpr(0)
+    _RS_D1 = tl.constexpr(0)
+    _RS_D2 = tl.constexpr(0)
+    _RS_D3 = tl.constexpr(0)
+    _RS_D4 = tl.constexpr(0)
+    _RS_D5 = tl.constexpr(0)
+    _RS_D6 = tl.constexpr(0)
+    _RS_D7 = tl.constexpr(0)
+    _RS_NODE_START = tl.constexpr(0)
+
+    def set_rocshmem_peer_deltas(deltas, node_start=0):
+        """Populate the device peer-delta table. Called once at init, BEFORE any
+        kernel launch. ``deltas[i]`` == peer_addr(x, node_start+i) - x for
+        symmetric x (indexed by LOCAL position i). ``node_start`` is this node's
+        first global rank, used to map a runtime global pe to a local slot."""
+        global _RS_D0, _RS_D1, _RS_D2, _RS_D3, _RS_D4, _RS_D5, _RS_D6, _RS_D7
+        global _RS_NODE_START
+        d = list(deltas) + [0] * (8 - len(deltas))
+        _RS_D0 = tl.constexpr(d[0])
+        _RS_D1 = tl.constexpr(d[1])
+        _RS_D2 = tl.constexpr(d[2])
+        _RS_D3 = tl.constexpr(d[3])
+        _RS_D4 = tl.constexpr(d[4])
+        _RS_D5 = tl.constexpr(d[5])
+        _RS_D6 = tl.constexpr(d[6])
+        _RS_D7 = tl.constexpr(d[7])
+        _RS_NODE_START = tl.constexpr(node_start)
+
+    @triton.jit
+    def _rs_peer_delta(pe):
+        # Select the baked delta for runtime global ``pe`` as int64 via a masked
+        # sum over LOCAL positions (exactly one term is non-zero for a same-node
+        # pe; all terms are int64).
+        p = pe.to(tl.int64) - _RS_NODE_START
+        d = (p == 0).to(tl.int64) * _RS_D0
+        d += (p == 1).to(tl.int64) * _RS_D1
+        d += (p == 2).to(tl.int64) * _RS_D2
+        d += (p == 3).to(tl.int64) * _RS_D3
+        d += (p == 4).to(tl.int64) * _RS_D4
+        d += (p == 5).to(tl.int64) * _RS_D5
+        d += (p == 6).to(tl.int64) * _RS_D6
+        d += (p == 7).to(tl.int64) * _RS_D7
+        return d
+
+    @triton.jit
+    def _rs_peer_addr(local_ptr, pe):
+        # Translate a LOCAL symmetric pointer to the peer's XGMI-mapped address.
+        delta = _rs_peer_delta(pe).to(tl.uint64, bitcast=True)
+        return local_ptr.to(tl.uint64, bitcast=True) + delta
+
+    @triton.jit
+    def int_p(dest, value, pe):
+        """Same-node single-int put: SYSTEM-scope release atomic store to the
+        peer's XGMI-mapped fine-grained slot (visible to a peer CPU-side read).
+        ``pe == my_pe`` -> delta 0 -> local store."""
+        peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True)
+        tl.atomic_xchg(peer, value, sem="release", scope="sys")
+
+    @triton.jit
+    def int_g(src, pe):
+        """Single int read from a peer slot. A volatile load is fine for a
+        one-shot read; spin loops must use int_wait_until_equals."""
+        peer = _rs_peer_addr(src, pe).to(tl.pointer_type(tl.int32), bitcast=True)
+        return tl.load(peer, volatile=True)
+
+    @triton.jit
+    def int_wait_until_equals(ptr, expected, pe):
+        """Spin until the peer slot == ``expected`` using a SYSTEM-scope acquire
+        atomic load (atomic_add of 0). This bypasses the MI300X L2 staleness that
+        makes a naive volatile-load spin never observe a peer's ack
+        (verified: ~sub-ms wait instead of a hang). The pure-Triton equivalent
+        of mori's int32_wait_until_equals."""
+        peer = _rs_peer_addr(ptr, pe).to(tl.pointer_type(tl.int32), bitcast=True)
+        got = expected - 1
+        while got != expected:
+            got = tl.atomic_add(peer, 0, sem="acquire", scope="sys")
+
+    @triton.jit
+    def quiet():
+        # Single-node: bulk transfer is host-side copy_ and the same-node
+        # signalling kernels never call quiet(). No-op shim so the module
+        # imports; never on the single-node hot path.
+        pass
+
+    @triton.jit
+    def putmem_nbi_block(dest, source, nbytes, pe):
+        # Single-node bulk moves are host-side copy_; this shim exists only so
+        # the (never-launched-on-single-node) cross-node kernels compile.
+        pass
+
+    @triton.jit
+    def getmem_nbi_block(dest, source, nbytes, pe):
+        pass
+
+    @triton.jit
+    def int_p_remote(dest, value, pe):
+        # Cross-node is unsupported in the single-node host-API backend.
+        pass
+
+    @triton.jit
+    def int_wait_until_equals_remote(scratch_ptr, src_ptr, expected, pe):
+        pass
+
+    @triton.jit
+    def int_atomic_compare_swap(dest, cond, value, pe):
+        peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True)
+        return tl.atomic_cas(peer, cond, value, sem="acq_rel", scope="sys")
+
+    @triton.jit
+    def int_atomic_swap(dest, value, pe):
+        peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True)
+        return tl.atomic_xchg(peer, value, sem="acq_rel", scope="sys")
+
+    @core.extern
+    def tid(axis: core.constexpr, _semantic=None):
+        # Platform-portable: return constant 0 (see mori branch rationale).
+        del axis
+        return core.full((), 0, dtype=tl.int32, _semantic=_semantic)
+
+    @core.extern
+    def __syncthreads(_semantic=None):
+        return tl.tensor(_semantic.builder.create_barrier(), tl.void)
+
+
+# ===========================================================================
+# Branch B: MORI backend (default) — re-export MORI device API.
+# ===========================================================================
+else:
+    import mori.shmem as _ms  # noqa: F401  (used by users importing this module)
+    from mori.ir import triton as _mt
+    from mori.ir.triton import get_extern_libs as _mori_get_extern_libs
+    from mori.ir.triton import install_hook as _mori_install_hook
+
+    # Install MORI's Triton compile hook. After install_hook(), Triton
+    # compilation of any kernel that uses MORI device functions will be linked
+    # against libmori_shmem_device.bc.
+    _mori_install_hook()
+
+    # SHMEM_EXTERN_LIBS is a dict {LIB_NAME: LIB_SHMEM_PATH}; MORI returns the
+    # same shape from ``get_extern_libs()``.
+    SHMEM_EXTERN_LIBS = _mori_get_extern_libs()
+    LIB_NAME = "libmori_shmem"
+    # LIB_SHMEM_PATH kept for backward compatibility with code that probes the
+    # dict; on ROCm this is the path to the JIT-compiled MORI bitcode.
+    if isinstance(SHMEM_EXTERN_LIBS, dict) and SHMEM_EXTERN_LIBS:
+        LIB_SHMEM_PATH = next(iter(SHMEM_EXTERN_LIBS.values()))
+    else:
+        LIB_SHMEM_PATH = ""
+
+    # -----------------------------------------------------------------------
+    # Device API: thin @triton.jit shims that fold the extra ``qp`` argument.
+    # -----------------------------------------------------------------------
+    # NOTE: MORI device functions take an extra qp index. ODC's signatures pass
+    # (dest, src, nbytes, pe). We hard-code qp=0 for now (single-QP is the
+    # default); if/when ODC wants multi-QP, change here.
+
+    @triton.jit
+    def putmem_nbi_block(dest, source, nbytes, pe):
+        return _mt.putmem_nbi_block(dest, source, nbytes, pe, 0)
+
+    @triton.jit
+    def getmem_nbi_block(dest, source, nbytes, pe):
+        return _mt.getmem_nbi_block(dest, source, nbytes, pe, 0)
+
+    @triton.jit
+    def quiet():
+        return _mt.quiet_thread()
+
+    @triton.jit
+    def int_p(dest, value, pe):
+        """Write a 32-bit int to ``dest`` on rank ``pe`` from this rank.
+
+        SAME-NODE ONLY.
+        ------------------------------------------------------------------
+        Resolve the peer pointer via ``ptr_p2p`` (XGMI direct mapping)
+        and write through it with ``tl.store``. This is the recommended
+        same-node path in MORI's own examples.
+
+        IMPORTANT (regression guard): this primitive is compiled into the
+        SAME-NODE request kernel that runs on every single-node iteration.
+        It must NOT reference any RDMA device function (``int32_p`` etc.).
+        Linking an RDMA primitive into this kernel activates MORI's IBGDA
+        path, which on a host without a usable RoCE/IB NIC silently breaks
+        the same-node store (empirically: ODC grad reduce produced zero
+        gradients on ~half the iters). Cross-node puts live in the separate
+        ``int_p_remote`` below, which is only ever compiled into the
+        cross-node-only kernels (never launched on a single node).
+
+        Known limitation (investigated empirically on MI300X + ROCm 7.2)
+        ---------------------------------------------------------------
+        ROCm's XGMI cache coherence between the writer GPU and the
+        reader GPU's CPU-side D2H copy is not strict: a P2P-written
+        int32 can take O(1 second) before the destination GPU's
+        ``cudaMemcpy(D2H)`` returns the new value.
+        """
+        my_pe = _mt.my_pe()
+        if pe == my_pe:
+            tl.store(dest.to(tl.pointer_type(tl.int32)), value)
+        else:
+            peer_raw = _mt.ptr_p2p(dest.to(tl.uint64, bitcast=True), my_pe, pe)
+            peer_ptr = peer_raw.to(tl.pointer_type(tl.int32), bitcast=True)
+            tl.store(peer_ptr, value)
+
+    @triton.jit
+    def int_p_remote(dest, value, pe):
+        """Cross-node single-int put: write ``value`` to the symmetric address
+        ``dest`` on the remote PE ``pe`` over RDMA.
+
+        Uses MORI's ``int32_p`` which the device API dispatches onto the RDMA
+        (IBGDA) transport for a cross-node PE (and P2P for a same-node PE, so
+        it is also safe if called with a same-node peer). ``dest`` is the
+        LOCAL symmetric address; MORI resolves the remote PE's matching
+        symmetric offset internally (qp=0).
+
+        This primitive is intentionally SEPARATE from ``int_p`` so that the
+        RDMA device symbol is only ever linked into the cross-node-only
+        kernels (``shmem_*_remote_node_kernel`` / ``shmem_cross_node_*``),
+        which are never compiled/launched on a single node. This keeps the
+        single-node code path unchanged.
+        """
+        _mt.int32_p(dest.to(tl.uint64, bitcast=True), value, pe, 0)
+
+    @triton.jit
+    def int_g(src, pe):
+        """Read a 32-bit int from ``src`` on rank ``pe`` to this rank (single read).
+
+        SAME-NODE ONLY (same regression-guard rationale as ``int_p``).
+        ------------------------------------------------------------------
+        Resolve the peer pointer via ``ptr_p2p`` (XGMI direct mapping) and
+        read with ``tl.load``. NOTE: a *single* read is fine, but do NOT
+        use this inside a hot in-kernel spin loop — on ROCm the GPU L2
+        caches the peer address and a repeated volatile load never
+        observes a cross-process update. Use ``int_wait_until_equals``
+        (same node) or ``int_wait_until_equals_remote`` (cross node) for
+        spin-waiting instead.
+        """
+        my_pe = _mt.my_pe()
+        if pe == my_pe:
+            return tl.load(src.to(tl.pointer_type(tl.int32)), volatile=True)
+        peer_raw = _mt.ptr_p2p(src.to(tl.uint64, bitcast=True), my_pe, pe)
+        peer_ptr = peer_raw.to(tl.pointer_type(tl.int32), bitcast=True)
+        return tl.load(peer_ptr, volatile=True)
+
+    @triton.jit
+    def int_wait_until_equals(ptr, expected, pe):
+        """Block until the int32 at ``ptr`` on rank ``pe`` equals ``expected``.
+
+        This is the correct replacement for a naive
+            ``while r != expected: quiet(); r = int_g(...)``
+        spin loop on ROCm.
+
+        Root cause it fixes
+        -------------------
+        A naive spin ``while ...: tl.load(volatile=True)`` does not observe
+        cross-process updates on MI300X: the GPU L2 caches the peer address
+        and a long-lived in-kernel spin never sees a cross-process peer's
+        ack write, causing the scatter_accumulate "60s hang".
+
+        MORI's ``int32_wait_until_equals`` is implemented as
+            ``while (AtomicLoadRelaxedSystem(addr) != val) {}``
+        i.e. a SYSTEM-SCOPE atomic load that bypasses/refreshes the L2, so
+        the cross-process / cross-GPU write becomes visible. Verified on
+        MI300X: replacing the volatile spin with this turns a hang into a
+        ~0.19 ms wait. We pass the PEER's P2P-resolved address so the
+        same-node read path is exercised.
+        """
+        my_pe = _mt.my_pe()
+        if pe == my_pe:
+            target = ptr.to(tl.uint64, bitcast=True)
+        else:
+            target = _mt.ptr_p2p(ptr.to(tl.uint64, bitcast=True), my_pe, pe)
+        _mt.int32_wait_until_equals(target, expected)
+
+    @triton.jit
+    def int_wait_until_equals_remote(scratch_ptr, src_ptr, expected, pe):
+        """Cross-node spin-wait: block until the int32 at ``src_ptr`` on rank
+        ``pe`` equals ``expected``.
+
+        Why a separate primitive (vs ``int_wait_until_equals``)
+        ------------------------------------------------------
+        MORI's ``int32_wait_until_equals(addr, val)`` polls a *local*
+        address (``while AtomicLoadRelaxedSystem(addr) != val``). The
+        same-node ``int_wait_until_equals`` makes that work by passing the
+        peer's XGMI-mapped address (the peer's memory is visible in this
+        GPU's address space). A cross-node peer has **no** such mapping —
+        ``ptr_p2p`` returns 0 and there is no local address that aliases
+        the remote slot — so we must actively pull the value over RDMA.
+
+        Implementation
+        --------------
+        Repeatedly RDMA-``get`` the 4-byte remote slot into a local
+        symmetric scratch and compare. ``getmem`` carries no atomic-alignment
+        constraint (unlike a remote atomic), so any int32 slot works.
+
+        Args:
+            scratch_ptr: a LOCAL symmetric/registered int32 scratch slot
+                (RDMA needs the get destination to be a registered MR;
+                MORI symmetric-heap tensors qualify).
+            src_ptr: the LOCAL symmetric address of the slot to watch; the
+                value is fetched from the same symmetric offset on ``pe``.
+            expected: the int32 value to wait for.
+            pe: the remote PE that owns the authoritative slot.
+        """
+        got = expected - 1
+        while got != expected:
+            _mt.getmem_nbi_block(
+                scratch_ptr.to(tl.uint64, bitcast=True),
+                src_ptr.to(tl.uint64, bitcast=True),
+                4,
+                pe,
+                0,
+            )
+            _mt.quiet_thread()
+            got = tl.load(scratch_ptr.to(tl.pointer_type(tl.int32)), volatile=True)
+
+    @triton.jit
+    def int_atomic_compare_swap(dest, cond, value, pe):
+        # MORI: atomic_uint32_fetch_thread(dest, val, cmp, op, pe, qp)
+        # CAS op code = 1 (per mori/src/shmem/atomic_ops.hpp; if different,
+        # adjust). Only used by ODC tests, not by the production gather/
+        # scatter path, so a slightly-incorrect op code would not affect
+        # the FSDP integration.
+        # We expose the same call shape; the op code is hard-coded as the
+        # MORI CAS-fetch primitive.
+        # TODO: cross-check op code with mori headers when wiring tests.
+        return _mt.atomic_uint32_fetch_thread(dest, value, cond, 1, pe, 0)
+
+    @triton.jit
+    def int_atomic_swap(dest, value, pe):
+        # MORI atomic SWAP op code = 0 (placeholder; not used in production).
+        return _mt.atomic_uint32_fetch_thread(dest, value, 0, 0, pe, 0)
+
+    # -----------------------------------------------------------------------
+    # tid / __syncthreads — platform-portable shims
+    # -----------------------------------------------------------------------
+    # ODC uses tid(axis=0) only to elect "first thread does the atomic_add"
+    # in sync_cta. The Triton-native way to do this is a mask pattern:
+    #     offsets = tl.arange(0, 1)
+    #     tl.atomic_add(signal_ptr + offsets, 1, mask=offsets == 0)
+    # rather than ``if tidx == 0: tl.atomic_add(signal_ptr, 1)``.
+    # We patch sync_cta in utils.py to use the mask form; here, ``tid``
+    # returns a constant 0 so that ``if tidx == 0`` branches taken in
+    # ODC's existing kernels (request_accumulation_same_node_kernel and
+    # wait_accumulation_same_node_kernel) still execute as before. In a
+    # ``@triton.jit`` kernel, all lanes follow the same control-flow path
+    # at the IR level, so ``if tid == 0`` becomes a conditional that
+    # produces the same result on every lane — semantically equivalent to
+    # "everybody does it" plus relying on the operation being idempotent
+    # under repetition. For ``int_p`` (memory store) this is fine because
+    # we issue the same write from every lane and MORI's underlying
+    # ``mori_shmem_int32_p`` is a wave-collective store (single transaction
+    # per wave).
+    @core.extern
+    def tid(axis: core.constexpr, _semantic=None):
+        # axis is a Triton constexpr (0/1/2). Return constant 0 — all lanes
+        # are treated identically. See note above for correctness rationale.
+        # We return a scalar int32 value.
+        del axis  # unused on ROCm path
+        return core.full((), 0, dtype=tl.int32, _semantic=_semantic)
+
+    @core.extern
+    def __syncthreads(_semantic=None):
+        # Platform-portable: Triton IR barrier. Lowers to s_barrier on AMDGPU.
+        return tl.tensor(_semantic.builder.create_barrier(), tl.void)
diff --git a/primus/core/odc/primitives/utils.py b/primus/core/odc/primitives/utils.py
new file mode 100644
index 000000000..fed1ebee2
--- /dev/null
+++ b/primus/core/odc/primitives/utils.py
@@ -0,0 +1,488 @@
+# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under
+# the MIT License per its package metadata (pyproject.toml / setup.py
+# classifiers). The upstream repository ships no LICENSE file or per-file
+# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab).
+#
+# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+
+"""
+ODC primitives utilities — ROCm symmetric-memory management.
+
+ODC allocates symmetric memory and enumerates same-node peer views. On ROCm
+there are two backends, selected by ``ODC_P2P_BACKEND``:
+
+* ``mori`` (default):
+    - ``init_shmem`` initializes MORI-SHMEM from the PyTorch process
+      group via ``mori.shmem.shmem_torch_process_group_init("default")``.
+    - ``shmem_create_tensor`` calls ``mori.shmem.mori_shmem_create_tensor``.
+    - Same-node peer views are obtained in a single MORI call:
+      ``mori_shmem_create_tensor_list_intra_node`` which returns the
+      list of symmetric views across all same-node ranks. Because that
+      API allocates *and* returns the peer-view list together, we
+      structure ``SymmBufferRegistry.allocate_symm_buffer`` to use it,
+      keeping the same external contract (a single tensor handed back
+      to the caller, plus a separately stored peer list).
+    - ``shmem_free_tensor_sync`` and ``finalize_distributed`` map to
+      ``mori.shmem.mori_shmem_free_tensor`` / ``shmem_finalize``.
+* ``rocshmem``: a host-API rocSHMEM backend (single-node / IPC-only)
+  implemented in ``_rocshmem_backend.py``.
+
+NOTE on environment variables (mori backend):
+    Set ``MORI_SHMEM_HEAP_SIZE`` (e.g. ``"4G"``) *before* the first MORI
+    call. Also set ``MORI_SOCKET_IFNAME`` (defaults to the value of
+    ``NCCL_SOCKET_IFNAME`` if available) for the MORI bootstrap.
+"""
+
+import logging
+import os
+from functools import reduce
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+
+from odc.primitives import __syncthreads
+from odc.runtime_config import get_config
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# ODC P2P backend selector (config item odc_p2p_backend). Default "mori"
+# preserves the existing, verified behaviour byte-for-byte; "rocshmem" selects
+# the host-API backend (single-node / IPC-only) implemented in
+# _rocshmem_backend.py. Read HERE at import time -- the ODC integration patch
+# populates the runtime config before odc.primitives is first imported.
+# ---------------------------------------------------------------------------
+_P2P_BACKEND = get_config().p2p_backend.lower()
+_USE_ROCSHMEM = _P2P_BACKEND == "rocshmem"
+
+
+# ---------------------------------------------------------------------------
+# Backend-specific imports
+# ---------------------------------------------------------------------------
+if _USE_ROCSHMEM:
+    from . import _rocshmem_backend as _rs
+else:
+    import mori.shmem as _mori_shmem
+
+
+# ---------------------------------------------------------------------------
+# init_shmem — initialize the symmetric-memory runtime.
+# ---------------------------------------------------------------------------
+def init_shmem():
+    """Initialize the symmetric-memory runtime on the global process group.
+
+    mori backend:
+        Bootstrap MORI-SHMEM from PyTorch's WORLD process group by name.
+        Requires that PyTorch distributed is already initialized.
+        ``MORI_SHMEM_HEAP_SIZE`` and ``MORI_SOCKET_IFNAME`` should be
+        exported in the environment before this call.
+
+    rocshmem backend:
+        Bootstrap the rocSHMEM host-API runtime via a unique-id broadcast.
+    """
+    assert torch.distributed.is_initialized()
+
+    if _USE_ROCSHMEM:
+        # rocSHMEM host-API backend: bootstrap via unique-id broadcast.
+        _rs.init()
+        return
+
+    # MORI requires MORI_SHMEM_HEAP_SIZE to be set. Inherit a sensible
+    # default if the user did not export one — this is purely a guard;
+    # production paths should always set it explicitly.
+    os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "4G")
+
+    # MORI uses its own TCP bootstrap socket. If the user set
+    # NCCL_SOCKET_IFNAME but not MORI_SOCKET_IFNAME, mirror it so they
+    # stay aligned.
+    if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ:
+        os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=")
+
+    # Two init methods (select via the odc_mori_init config item):
+    #   "pg"  (default): shmem_torch_process_group_init("default")
+    #   "uid": shmem_init_attr(WITH_UNIQUEID, rank, world, uid) — does NOT
+    #          touch PyTorch process-group registration. Use this when the
+    #          host framework (e.g. Megatron) has already created many
+    #          named process groups and the PG-based init crashes
+    #          (observed: `free(): invalid pointer` during
+    #          shmem_torch_process_group_init under Megatron).
+    init_method = get_config().mori_init
+    logger.info(
+        "init_shmem (MORI): heap=%s, sock_ifname=%s, method=%s",
+        os.environ.get("MORI_SHMEM_HEAP_SIZE"),
+        os.environ.get("MORI_SOCKET_IFNAME", "<auto>"),
+        init_method,
+    )
+
+    if init_method == "uid":
+        rank_id = torch.distributed.get_rank()
+        num_ranks = torch.distributed.get_world_size()
+        # rank 0 generates the 128-byte unique id; broadcast to all ranks.
+        uid_holder = [_mori_shmem.shmem_get_unique_id() if rank_id == 0 else None]
+        torch.distributed.broadcast_object_list(uid_holder, src=0)
+        torch.distributed.barrier()
+        unique_id = uid_holder[0]
+        _mori_shmem.shmem_init_attr(
+            _mori_shmem.MORI_SHMEM_INIT_WITH_UNIQUEID,
+            rank_id,
+            num_ranks,
+            unique_id,
+        )
+        return
+
+    # method == "pg": register WORLD as "default", then PG-based init.
+    world_group = torch.distributed.group.WORLD
+    try:
+        torch._C._distributed_c10d._register_process_group("default", world_group)
+    except RuntimeError:
+        # Already registered — ignore.
+        pass
+    _mori_shmem.shmem_torch_process_group_init("default")
+
+
+# ---------------------------------------------------------------------------
+# Symmetric tensor creation / peer-view enumeration / free
+# ---------------------------------------------------------------------------
+def shmem_create_tensor(shape, dtype) -> torch.Tensor:
+    """Allocate a symmetric tensor on the symmetric heap.
+
+    The full allocate + same-node peer enumeration is done in one call inside
+    ``SymmBufferRegistry.allocate_symm_buffer`` (see below). This top-level
+    helper is kept for callers that want a bare local tensor without the
+    peer list (e.g. when only the local view is needed). On the mori backend
+    it calls ``mori_shmem_create_tensor``, which allocates from MORI's
+    symmetric heap.
+    """
+    torch.cuda.synchronize()
+    if _USE_ROCSHMEM:
+        tensor = _rs.create_tensor(shape, dtype)
+    else:
+        tensor = _mori_shmem.mori_shmem_create_tensor(shape, dtype)
+    torch.cuda.synchronize()
+    return tensor
+
+
+def get_same_node_tensors(tensor, rank, local_world_size) -> List[torch.Tensor]:
+    """Return the list of peer-view tensors for the given symmetric ``tensor``
+    across the same-node ranks (``rank // local_world_size``'s peers).
+
+    This API is only safe when ``tensor`` was created by the paired
+    ``mori_shmem_create_tensor_list_intra_node`` (see
+    ``SymmBufferRegistry.allocate_symm_buffer``); the list is retrieved
+    from the registry. As a stand-alone fall-back, we synthesize it via
+    ``shmem_ptr_p2p`` and wrap raw device pointers with
+    ``torch.frombuffer``-style construction.
+    """
+    # When the tensor was created via the registry path, the registry
+    # already stores the peer list — callers should prefer
+    # ``SymmBufferRegistry.get_peer_tensors``.
+    # As a stand-alone fallback for ad-hoc symmetric tensors, build
+    # peer views via shmem_ptr_p2p + torch.Tensor reconstruction.
+    peer_tensors: List[torch.Tensor] = []
+    local_rank = rank % local_world_size
+    rank_on_same_node_start = rank - local_rank
+    for peer in range(rank_on_same_node_start, rank_on_same_node_start + local_world_size):
+        if peer == rank:
+            peer_tensors.append(tensor)
+            continue
+        peer_ptr = _mori_shmem.shmem_ptr_p2p(tensor.data_ptr(), rank, peer)
+        if peer_ptr == 0:
+            raise RuntimeError(
+                f"shmem_ptr_p2p returned 0 for peer {peer}: no XGMI P2P route "
+                "from PE {rank}. This usually means peer is on a different node "
+                "(requires RDMA, not P2P)."
+            )
+        # Build a torch.Tensor view at peer_ptr.
+        peer_tensors.append(_tensor_from_raw_ptr(peer_ptr, tensor.shape, tensor.dtype, tensor.device))
+    return peer_tensors
+
+
+def _tensor_from_raw_ptr(raw_ptr: int, shape, dtype, device) -> torch.Tensor:
+    """Wrap a raw device pointer as a torch.Tensor without taking ownership.
+
+    Used on ROCm to materialize peer-view tensors from
+    ``mori.shmem.shmem_ptr_p2p``. The returned tensor shares storage with
+    the original symmetric allocation (which is owned by the registry);
+    no deleter is attached, so the caller must guarantee that the symm
+    allocation outlives all peer views.
+    """
+    import ctypes
+
+    # Compute total bytes
+    if isinstance(shape, int):
+        shape = (shape,)
+    nbytes = reduce(lambda a, b: a * b, shape) * torch._utils._element_size(dtype)
+
+    # Build a ctypes array view at raw_ptr (zero-copy) — only used to
+    # pass through PyTorch's from_blob. PyTorch will create a tensor that
+    # references this storage. The CObject is intentionally kept alive
+    # via attribute attachment so GC does not free the wrapper.
+    array_t = (ctypes.c_uint8 * nbytes).from_address(raw_ptr)
+    t = torch.frombuffer(array_t, dtype=torch.uint8, count=nbytes).view(dtype).view(*shape).to(device)
+    # NOTE: ``torch.frombuffer`` materializes a CPU tensor; ``.to(device)`` would
+    # COPY. We can't use that path for a real peer-view (would defeat the purpose).
+    # Instead use the explicit cuda path:
+    raise NotImplementedError(
+        "Direct raw-ptr → torch.Tensor wrapping is not implemented in this "
+        "iteration. Use SymmBufferRegistry.allocate_symm_buffer (which uses "
+        "mori_shmem_create_tensor_list_intra_node) — that path is the one "
+        "exercised by ODC's gather/scatter primitives. The bare "
+        "get_same_node_tensors fall-back is provided only as a placeholder."
+    )
+
+
+def shmem_free_tensor_sync(tensor):
+    torch.cuda.synchronize()
+    if _USE_ROCSHMEM:
+        _rs.free_tensor(tensor)
+    else:
+        _mori_shmem.mori_shmem_free_tensor(tensor)
+    torch.cuda.synchronize()
+
+
+def finalize_distributed():
+    if _USE_ROCSHMEM:
+        _rs.finalize()
+    else:
+        _mori_shmem.shmem_finalize()
+
+
+# ---------------------------------------------------------------------------
+# Symmetric buffer registry (the heart of ODC memory management)
+# ---------------------------------------------------------------------------
+class SymmBufferRegistry:
+    def __init__(self):
+        self.local_tensor = {}
+        self.local_tensor_to_keys = {}
+        self.updated = set()
+        self.peer_tensors = {}
+        self.allocations = []
+
+    @classmethod
+    def get_instance(cls):
+        if not hasattr(cls, "_instance"):
+            cls._instance = SymmBufferRegistry()
+        return cls._instance
+
+    # we'll mark all symm buffer as dirty, and next update_symm_buffer will copy the data to the symm buffer
+    def flush(self):
+        self.updated.clear()
+
+    def update_symm_buffer(self, buffer_key, values):
+        values = values.contiguous()
+        if buffer_key not in self.local_tensor:
+            self.allocate_symm_buffer(buffer_key, values.shape, values.dtype)
+
+        if buffer_key not in self.updated:
+            self.updated.add(buffer_key)
+            self.local_tensor[buffer_key].copy_(values)
+            # Make sure updated buffer is visible to all ranks
+            torch.distributed.barrier()
+        return self.local_tensor[buffer_key]
+
+    @classmethod
+    def set_shmem_flag(cls, tensor):
+        tensor._odc_is_shmem = True
+
+    @classmethod
+    def is_shmem_tensor(cls, tensor):
+        return hasattr(tensor, "_odc_is_shmem") and tensor._odc_is_shmem
+
+    def allocate_symm_buffer(self, key, shape, dtype):
+        """Allocate a symmetric tensor and record same-node peer views.
+
+        mori backend: one-shot via ``mori_shmem_create_tensor_list_intra_node``,
+        which returns ``list[torch.Tensor]`` of length ``local_world_size``,
+        each entry being the same symmetric allocation viewed from the
+        corresponding peer's address space. ``tensors[my_local_rank]`` is
+        the local view; the entire list is recorded as the peer-view list.
+
+        rocshmem backend: one-shot ``_rs.alloc_peer_tensors`` (equivalent).
+        """
+        assert key not in self.local_tensor
+        local_world_size = get_local_world_size()
+        rank = torch.distributed.get_rank()
+        local_rank = rank % local_world_size
+
+        if _USE_ROCSHMEM:
+            # rocSHMEM host-API: one-shot symmetric alloc + same-node peer-view
+            # list (equivalent to MORI's mori_shmem_create_tensor_list_intra_node).
+            tensor, peer_tensors = _rs.alloc_peer_tensors(shape, dtype, local_world_size, rank)
+            self.allocations.append(tensor)
+        else:
+            # node_start == 0 covers BOTH the single-node case and node 0 of a
+            # multi-node run. In that case the same-node PEs are exactly
+            # [0, local_world_size), so we call MORI's one-shot helper VERBATIM
+            # — byte-for-byte the original single-node behaviour (single-node
+            # protected, zero change).
+            #
+            # node_start > 0 (only reachable on node>0 of a multi-node run):
+            # MORI's helper hardcodes GLOBAL PEs range(lws) == [0, lws), which
+            # are cross-node from this node (ptr_p2p -> 0 -> bogus tensor at
+            # address 0, also corrupting the "local" view). We instead allocate
+            # the symmetric tensor once and enumerate THIS node's real PE range
+            # [node_start, node_start + lws).
+            node_start = rank - local_rank
+            if node_start == 0:
+                peer_tensors = _mori_shmem.mori_shmem_create_tensor_list_intra_node(
+                    shape, dtype, local_world_size
+                )
+                assert len(peer_tensors) == local_world_size, (
+                    f"mori_shmem_create_tensor_list_intra_node returned "
+                    f"{len(peer_tensors)} tensors, expected {local_world_size}"
+                )
+                tensor = peer_tensors[local_rank]
+            else:
+                base_tensor = _mori_shmem.mori_shmem_create_tensor(shape, dtype)
+                peer_tensors = [
+                    _mori_shmem.symm_mori_shmem_tensor(base_tensor, node_start + i)
+                    for i in range(local_world_size)
+                ]
+                assert len(peer_tensors) == local_world_size, (
+                    f"expected {local_world_size} same-node peer views, " f"got {len(peer_tensors)}"
+                )
+                # peer_tensors[local_rank] resolves to peer == my_pe, i.e. the
+                # true local allocation (base_tensor) — right object to free.
+                tensor = peer_tensors[local_rank]
+            # Keep the same allocation list to track for finalize.
+            self.allocations.append(tensor)
+
+        self.set_shmem_flag(tensor)
+        assert len(peer_tensors) == local_world_size
+
+        # ranks inside the same node must be contiguous.
+        self.local_tensor[key] = tensor
+        self.peer_tensors[key] = peer_tensors
+        self.local_tensor_to_keys[self.local_tensor[key].data_ptr()] = key
+        logger.info(
+            f"Rank {rank} create tensor {key} with shape {shape} and dtype {dtype} "
+            f"and ptr {self.local_tensor[key].data_ptr():#x}"
+        )
+        return tensor
+
+    def has_key(self, key):
+        return key in self.local_tensor
+
+    def get_or_create_symm_buffer(self, key, shape, dtype):
+        if self.has_key(key):
+            return self.local_tensor[key]
+        return self.allocate_symm_buffer(key, shape, dtype)
+
+    def get_symm_buffer(self, key):
+        if self.has_key(key):
+            return self.local_tensor[key]
+        raise ValueError(f"Symm buffer {key} not found")
+
+    def get_peer_tensors(self, local_tensor):
+        # Returns tensors in the same node.
+        buffer_key = self.local_tensor_to_keys[local_tensor.data_ptr()]
+        return self.peer_tensors[buffer_key]
+
+    def finalize(self):
+        for t in self.allocations:
+            shmem_free_tensor_sync(t)
+        self.local_tensor.clear()
+        self.local_tensor_to_keys.clear()
+        self.updated.clear()
+        self.peer_tensors.clear()
+
+    def memory_allocated(self):
+        return sum(t.nbytes for t in self.allocations)
+
+
+# ---------------------------------------------------------------------------
+# Process group helpers — unchanged from upstream
+# ---------------------------------------------------------------------------
+local_world_pg = None
+
+
+def get_local_world_pg(pg: torch.distributed.ProcessGroup):
+    local_world_size = get_local_world_size()
+    assert torch.distributed.get_world_size() == torch.distributed.get_world_size(
+        group=pg
+    ), "Cached AG only supports pure data parallelism"
+    rank = torch.distributed.get_rank()
+    global local_world_pg
+    if local_world_pg is None:
+        for i in range(0, torch.distributed.get_world_size(), local_world_size):
+            ranks = list(range(i, i + local_world_size))
+            new_gp = torch.distributed.new_group(ranks=ranks, backend="nccl")
+            if rank in ranks:
+                local_world_pg = new_gp
+    assert local_world_pg is not None
+    return local_world_pg
+
+
+def get_local_world_size():
+    if "RAY_LOCAL_WORLD_SIZE" in os.environ:
+        return int(os.environ["RAY_LOCAL_WORLD_SIZE"])
+    else:
+        return int(os.environ["LOCAL_WORLD_SIZE"])
+
+
+stream = None
+
+
+def get_comm_stream():
+    global stream
+    if stream is None:
+        stream = torch.cuda.Stream()
+    return stream
+
+
+class BufferSplitter:
+    def __init__(self):
+        self.round_data_size = 2**6
+
+    def get_max_global_buffer_size(self):
+        # config item odc_max_buffer_size (default 64 MiB).
+        max_buffer_size = int(get_config().max_buffer_size)
+        return max_buffer_size
+
+    def get_global_buffer_size(self, original_buffer_shape):
+        original_size = reduce(lambda x, y: x * y, original_buffer_shape)
+        max_buffer_size = self.get_max_global_buffer_size()
+        if max_buffer_size <= 0:
+            return original_size
+        buf_size = min(max_buffer_size, original_size)
+        return buf_size
+
+    def round_to_data_size(self, size):
+        return (size + self.round_data_size - 1) // self.round_data_size * self.round_data_size
+
+    def get_local_buffer_size(self, original_buffer_shape, world_size):
+        original_size = reduce(lambda x, y: x * y, original_buffer_shape)
+        max_buffer_size = self.get_max_global_buffer_size()
+        if max_buffer_size <= 0:
+            return self.round_to_data_size(original_size)
+        assert (
+            max_buffer_size % world_size == 0
+        ), f"ODC_MAX_BUFFER_SIZE: {max_buffer_size} % world_size: {world_size} != 0"
+        local_max_buffer_size = max_buffer_size // world_size
+        buf_size = min(local_max_buffer_size, original_size)
+        return self.round_to_data_size(buf_size)
+
+
+# ---------------------------------------------------------------------------
+# sync_cta — CTA-internal synchronization helper
+# ---------------------------------------------------------------------------
+# Elect a single thread (lane 0) to do the atomic_add using the Triton vector
+# mask pattern (no NVVM tid intrinsic exists on AMDGPU).
+@triton.jit
+def sync_cta(signal_ptr, expected):
+    # Use a length-1 lane vector with mask to ensure only lane 0 of
+    # the wave performs the atomic_add. This is the Triton-native way
+    # to elect a "first thread".
+    offsets = tl.arange(0, 1)
+    mask = offsets == 0
+    tl.atomic_add(signal_ptr + offsets, 1, mask=mask)
+    __syncthreads()
+    r = 0
+    while r < expected:
+        signals = tl.load(signal_ptr + offsets, mask=mask, volatile=True)
+        r = tl.max(signals)
diff --git a/primus/core/odc/rocshmem_runtime/README.md b/primus/core/odc/rocshmem_runtime/README.md
new file mode 100644
index 000000000..98e65f248
--- /dev/null
+++ b/primus/core/odc/rocshmem_runtime/README.md
@@ -0,0 +1,58 @@
+# rocshmem_runtime — ODC runtime helpers (rocSHMEM ops now in Primus-Turbo)
+
+The ODC rocSHMEM host/GDA ops were migrated to **Primus-Turbo**
+(`primus_turbo.pytorch._C.odc_rocshmem_host` / `odc_rocshmem_gda`) and are
+consumed by `odc/primitives/_rocshmem_backend.py`. The former in-tree ctypes
+bindings (`host_bindings/rs_host.cpp`, `gda_backend/rs_host_gda.cpp`) and the
+`build_rocshmem_backend.sh` build script have been **removed**.
+
+The rocSHMEM **static library** (`librocshmem.a` + headers) that Primus-Turbo
+links against is an **external dependency**: point the Turbo build's
+`ROCSHMEM_HOME` at a rocSHMEM install. Any local rocSHMEM build tree
+(`rocshmem_src/`, `rocshmem_{single,gda}/`, `*.a`, `*.so`) is `.gitignore`d and
+must never be committed.
+
+> TODO(primus-turbo): pin the exact Primus-Turbo merge commit
+> (`PRIMUS_TURBO_COMMIT`) once the ODC-ops PR merges.
+
+This directory now only holds runtime helpers:
+
+- `scripts/run_odc.sh` — single-node ODC launcher (plain torchrun). Default P2P
+  backend is `mori`; pass `rocshmem` to use the Turbo-provided ops. Set
+  `PRIMUS_TURBO_PATH` to a Primus-Turbo build tree if it is not already importable.
+- `scripts/cleanup_rs.sh` — leftover-process cleanup.
+
+## Selecting the rocSHMEM backend
+
+These are Primus config items (read by the ODC library via its runtime config):
+
+- `odc_p2p_backend: rocshmem` — use the rocSHMEM ops (default is `mori`).
+- `odc_rocshmem_gda: true` — multi-node GPU-direct (GDA) path; otherwise
+  single-node XGMI IPC host path.
+
+The ops resolve from `primus_turbo.pytorch._C`. As an escape hatch,
+`odc_rocshmem_lib: <path>/librs_host_gda.so` loads an external monolithic ctypes
+binding instead of the Turbo submodule (see `_rocshmem_backend.py`).
+
+## Multi-node (GDA) correctness & deployment env
+
+**Correctness (automatic, no env needed).** Multi-node (`n_pes > GPUs-per-node`)
+defaults to `ODC_GDA_DEFER_REDUCE=1`: each micro-batch's unsharded grad is
+accumulated locally, then ONE barriered reduce-scatter runs per minibatch. This
+keeps the collective-barrier count equal across ranks, so variable-length
+(`nopad`) packing does not deadlock. Single-node defaults to `0`.
+
+**Cluster deployment env** (set per your cluster; not hardcoded in the repo):
+
+| env | example | notes |
+|-----|---------|-------|
+| `ROCSHMEM_HCA_LIST` | `mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9` | this cluster's compute NICs (else all traffic funnels through `mlx5_0`) |
+| `ROCSHMEM_GDA_PROVIDER` | `mlx5` | RDMA provider |
+| `ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME` | `eth0` | uid-over-socket bootstrap NIC |
+| `ROCSHMEM_HEAP_SIZE` | `8589934592` | symmetric heap RAW bytes (decimal only, no K/M/G) |
+| `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` | `eth0` | torch.distributed control plane |
+| `NCCL_IB_GID_INDEX` | `3` | RoCE v2 GID (fabric-dependent) |
+| `ODC_GDA_WARMUP_MODE` / `ODC_GDA_STRIDE_BYTES` | `strided` / `65536` | GDA connection warmup |
+
+The GDA backend bootstraps rocSHMEM with a unique-id over a TCP socket, so the
+job launches with **plain torchrun** — no MPI / mpirun.
diff --git a/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh b/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh
new file mode 100755
index 000000000..8206cc2b7
--- /dev/null
+++ b/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+# shellcheck disable=SC2009  # need full args column; pgrep can't match python -m patterns
+# Terminate all residual training/compile processes from the abandoned rocSHMEM
+# smoke in THIS container. Patterns live in this file (not in the invoking
+# shell's argv), so `pkill -f` cannot match its own parent shell.
+echo "=== BEFORE: residual training/compile procs ==="
+ps -eo pid,etimes,comm,args 2>/dev/null \
+  | grep -iE 'torchrun|primus/cli|torch.distributed|hipcc|clang|cmake|ninja|llvm|multiprocessing' \
+  | grep -vE 'grep|cleanup_rs.sh' || echo "(none before)"
+
+for pat in 'torchrun' 'torch.distributed.run' 'primus/cli' 'run_pretrain.sh' \
+           'run_odc.sh' 'multiprocessing.spawn' 'multiprocessing.resource_tracker' \
+           'hipcc' 'cmake' 'ninja'; do
+  pkill -9 -f "$pat" 2>/dev/null
+done
+pkill -9 -x python3 2>/dev/null
+pkill -9 -x pt_main_thread 2>/dev/null
+pkill -9 clang 2>/dev/null
+pkill -9 'clang++' 2>/dev/null
+pkill -9 llvm-as 2>/dev/null
+pkill -9 lld 2>/dev/null
+sleep 4
+
+echo "=== AFTER: residual training/compile procs (want NONE) ==="
+LEFT=$(ps -eo pid,etimes,comm,args 2>/dev/null \
+  | grep -iE 'torchrun|primus/cli|torch.distributed|hipcc|clang|cmake|ninja|llvm|multiprocessing|python3' \
+  | grep -vE 'grep|cleanup_rs.sh')
+if [ -z "$LEFT" ]; then echo "(none)"; else echo "$LEFT"; fi
+
+echo "=== rocm-smi GPU use% ==="
+rocm-smi --showuse 2>/dev/null | grep -E 'GPU use \(%\)'
+echo "=== rocm-smi compute pids ==="
+rocm-smi --showpids 2>/dev/null | sed -n '1,12p'
diff --git a/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh b/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh
new file mode 100755
index 000000000..5d0bc76a9
--- /dev/null
+++ b/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh
@@ -0,0 +1,97 @@
+#!/bin/bash
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+# =============================================================================
+# run_odc.sh — portable ODC LB-mini launcher (single-node)
+#
+# Runs an ODC LB-mini training job. Default P2P backend is `mori`; passing
+# `rocshmem` switches to the rocSHMEM backend. The rocSHMEM ops are consumed
+# from Primus-Turbo (primus_turbo.pytorch._C.odc_rocshmem_host / _gda); set
+# PRIMUS_TURBO_PATH to a Primus-Turbo build tree if it is not already importable
+# in the environment (it is prepended to PYTHONPATH).
+#
+# usage: run_odc.sh <mori|rocshmem> <pad|nopad> <exp_yaml_relpath> <exp_name> [KEY=VAL ...]
+#   pad|nopad is retained for backwards-compatible invocation but is now a no-op
+#     (the aligned-vs-decoupled A/B study knob was removed; LB-Mini always runs
+#     decoupled when enabled). ODC feature switches live in the EXP yaml config.
+#   extra KEY=VAL args are exported verbatim.
+#
+# Overridable env (all have portable defaults):
+#   PRIMUS_ROOT            project root (auto-derived from this script's path)
+#   PRIMUS_TURBO_PATH      Primus-Turbo build tree to prepend to PYTHONPATH so
+#                          `import primus_turbo` (with the ODC rocSHMEM ops)
+#                          resolves; leave unset if already installed.
+#   HF_HOME                HF cache dir (default: /workspace/hf_cache)
+#   PRIMUS_PACK_CACHE_DIR  packed-sequence cache (default: $HOME/primus_packed)
+#   TRITON_CACHE_DIR       triton cache (rocshmem: fresh per-run unless pinned)
+#   TRAIN_LOG_DIR          where to write runlog_*.log (default: $HOME/odc_logs)
+#   MASTER_PORT            default 29600
+#   ROCSHMEM_HEAP_SIZE     symmetric heap RAW BYTES (default 8 GiB)
+# =============================================================================
+set -u
+BACKEND=$1; PAD=$2; EXP_REL=$3; EXPNAME=$4; shift 4
+
+# --- derive project root from this script's location (portable) -------------
+# scripts/ -> rocshmem_runtime/ -> odc/ -> core/ -> primus/ -> <PRIMUS_ROOT>
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RUNTIME_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"            # rocshmem_runtime/
+ODC_ROOT="$(cd "$RUNTIME_DIR/.." && pwd)"              # primus/core/odc/
+PRIMUS_ROOT="${PRIMUS_ROOT:-$(cd "$ODC_ROOT/../../.." && pwd)}"
+cd "$PRIMUS_ROOT" || exit 1
+
+export EXP=$EXP_REL
+# ODC arm env. Prepend a Primus-Turbo build tree if provided so the rocSHMEM
+# backend (primus_turbo.pytorch._C.odc_rocshmem_*) is importable.
+# The `odc` package now lives directly at $ODC_ROOT (primus/core/odc/), so put
+# its PARENT (primus/core/) on PYTHONPATH for `import odc`; odc_early holds the
+# sitecustomize load-order shim.
+export PYTHONPATH="${PRIMUS_TURBO_PATH:+$PRIMUS_TURBO_PATH:}$ODC_ROOT/odc_early:${ODC_ROOT%/*}"
+# ODC feature switches (enable_odc, odc_phase, enable_odc_lb_mini, ...) are now
+# CONFIG items set in the EXP yaml, NOT env vars. Only genuine infra env is set
+# here. MORI_SHMEM_HEAP_SIZE is MORI runtime infra (symmetric heap size).
+export MORI_SHMEM_HEAP_SIZE=8G
+# NOTE: the $PAD positional arg is retained for backwards-compatible invocation
+# but no longer toggles an aligned-vs-decoupled A/B baseline (that was a study
+# knob and has been removed); LB-Mini always runs decoupled when enabled.
+# public env
+export HF_HOME=${HF_HOME:-/workspace/hf_cache} DATA_PATH=${DATA_PATH:-/workspace}
+export GLOO_SOCKET_IFNAME=lo NCCL_SOCKET_IFNAME=lo NCCL_IB_DISABLE=1
+export FUSED_LINEAR_CE=1
+export PRIMUS_PACK_CACHE_DIR=${PRIMUS_PACK_CACHE_DIR:-$HOME/primus_packed}
+export PRIMUS_EXP_NAME=$EXPNAME
+export MASTER_PORT=${MASTER_PORT:-29600}
+mkdir -p "$PRIMUS_PACK_CACHE_DIR"
+# backend selection
+if [ "$BACKEND" = "rocshmem" ]; then
+  export ODC_P2P_BACKEND=rocshmem
+  # The rocSHMEM ops are consumed from Primus-Turbo (see PRIMUS_TURBO_PATH above);
+  # no in-tree librs_host*.so is loaded anymore.
+  export ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME=lo
+  # rocSHMEM symmetric heap size, RAW BYTES (the env parser is decimal-only and
+  # does NOT accept K/M/G suffixes). 8 GiB matches MORI_SHMEM_HEAP_SIZE.
+  export ROCSHMEM_HEAP_SIZE=${ROCSHMEM_HEAP_SIZE:-8589934592}
+  # IMPORTANT: the rocshmem device kernels bake per-PE peer deltas as Triton
+  # constexpr. Reusing a Triton cache built by a *different* toolchain (or a
+  # different launch) was observed to silently load mismatched kernels ->
+  # garbage int_p/wait signalling -> NaN grads from iter 1. Always start from a
+  # fresh per-run cache for rocshmem unless the caller pins TRITON_CACHE_DIR.
+  export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/tmp/tcache_rocshmem_$(date +%Y%m%d_%H%M%S)_$$}
+else
+  export ODC_P2P_BACKEND=mori
+  export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/tmp/tcache_mori}
+fi
+# extra KEY=VAL env
+# shellcheck disable=SC2163  # kv is a literal KEY=VAL token, so export works
+for kv in "$@"; do export "$kv"; done
+
+# Unique per-run timestamp so reruns never clobber earlier logs. Override
+# TRAIN_LOG_TS to pin a specific stamp (e.g. shared across multinode ranks).
+TRAIN_LOG_TS=${TRAIN_LOG_TS:-$(date +%Y%m%d_%H%M%S)}
+TRAIN_LOG_DIR=${TRAIN_LOG_DIR:-$HOME/odc_logs}; mkdir -p "$TRAIN_LOG_DIR"
+export TRAIN_LOG="$TRAIN_LOG_DIR/runlog_${EXPNAME}_${TRAIN_LOG_TS}.log"
+echo "[run_odc] ROOT=$PRIMUS_ROOT BACKEND=$BACKEND PAD=$PAD P2P=$ODC_P2P_BACKEND EXP=$EXP NAME=$EXPNAME TS=$TRAIN_LOG_TS"
+echo "[run_odc] TURBO_PATH=${PRIMUS_TURBO_PATH:-<installed>} TRITON_CACHE_DIR=$TRITON_CACHE_DIR LOG=$TRAIN_LOG"
+bash examples/run_pretrain.sh
+echo "[run_odc] DONE exit=$? log=$TRAIN_LOG"
diff --git a/primus/core/odc/runtime_config.py b/primus/core/odc/runtime_config.py
new file mode 100644
index 000000000..adf9970e4
--- /dev/null
+++ b/primus/core/odc/runtime_config.py
@@ -0,0 +1,86 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""ODC runtime configuration.
+
+The ODC primitives (``odc.primitives.*``) are decoupled library code with no
+knowledge of Primus' config system. Historically they read their tuning knobs
+straight from ``os.environ`` (ODC_P2P_BACKEND, ODC_GDA_*, ...), which violates
+the Primus rule that formal feature logic must be config-driven, not env-driven.
+
+This leaf module (imports only stdlib, so it can be populated BEFORE the heavy
+``odc.primitives`` modules are first imported) holds a single ``OdcRuntimeConfig``
+instance. The Primus ODC integration patch reads the ``odc_*`` config items from
+the trainer config and calls :func:`set_config` at ``before_train`` -- i.e. after
+the Primus config is available but before ``odc.primitives`` (and their
+import-time backend selection) run. The primitives then read values from here
+instead of the environment.
+
+Defaults below are byte-for-byte the previous env defaults, so an unset config
+reproduces the prior behaviour exactly.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class OdcRuntimeConfig:
+    # --- symmetric-memory backend selection (was ODC_P2P_BACKEND) ---
+    # "mori" (default) | "rocshmem"
+    p2p_backend: str = "mori"
+    # MORI init method (was ODC_MORI_INIT): "pg" (default) | "uid"
+    mori_init: str = "pg"
+    # BufferSplitter max global buffer size in bytes (was ODC_MAX_BUFFER_SIZE).
+    max_buffer_size: int = 64 * 1024 * 1024
+
+    # --- rocSHMEM backend (was ODC_ROCSHMEM_GDA / ODC_ROCSHMEM_LIB) ---
+    # GPU-direct (GDA) device-kernel cross-node path; required for multi-node.
+    rocshmem_gda: bool = False
+    # Optional path to a monolithic librs_host_gda.so ctypes override; None ->
+    # consume the rocSHMEM ops from Primus-Turbo.
+    rocshmem_lib: Optional[str] = None
+
+    # --- GDA reduce-scatter tuning knobs ---
+    # reduce-scatter kernel grid blocks (was ODC_GDA_RS_BLOCKS).
+    gda_rs_blocks: int = 64
+    # peer-pipeline batch depth (was PRIMUS_TURBO_ODC_GDA_PIPE). ALSO consumed by
+    # the Primus-Turbo device kernel via getenv, so set_config bridges it back to
+    # the PRIMUS_TURBO_ODC_GDA_PIPE env var for the C++ side.
+    gda_pipe: int = 1
+    # defer the cross-node reduce to once-per-minibatch (was ODC_GDA_DEFER_REDUCE):
+    # "auto" (default: on iff multi-node, i.e. n_pes > local_world_size) | "1" | "0".
+    gda_defer_reduce: str = "auto"
+    # cross-node write-visibility warm-up mode (was ODC_GDA_WARMUP_MODE):
+    # "strided" (default) | "full" | "hdp" | "fence" | "hdpfence".
+    gda_warmup_mode: str = "strided"
+    # strided warm-up page stride in bytes (was ODC_GDA_STRIDE_BYTES).
+    gda_stride_bytes: int = 65536
+
+
+_CONFIG = OdcRuntimeConfig()
+
+
+def get_config() -> OdcRuntimeConfig:
+    """Return the process-wide ODC runtime config (populated by set_config)."""
+    return _CONFIG
+
+
+def set_config(**kwargs) -> OdcRuntimeConfig:
+    """Update the ODC runtime config in place from keyword overrides.
+
+    Only known fields are applied; None values are ignored (keep the default).
+    Returns the config instance. Also bridges gda_pipe back to the
+    PRIMUS_TURBO_ODC_GDA_PIPE env var, which the Primus-Turbo device kernel reads
+    via getenv (the one knob that must stay visible to the C++ side).
+    """
+    import os
+
+    valid = set(OdcRuntimeConfig.__dataclass_fields__)
+    for k, v in kwargs.items():
+        if k in valid and v is not None:
+            setattr(_CONFIG, k, v)
+    # Bridge the pipe depth to the env var the turbo C++ kernel reads.
+    os.environ["PRIMUS_TURBO_ODC_GDA_PIPE"] = str(int(_CONFIG.gda_pipe))
+    return _CONFIG
diff --git a/primus/core/patches/patch.py b/primus/core/patches/patch.py
index 50522f163..466bfddfc 100644
--- a/primus/core/patches/patch.py
+++ b/primus/core/patches/patch.py
@@ -15,7 +15,7 @@
 
 from primus.core.patches.context import PatchContext
 from primus.core.patches.utils import version_in_range
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 # -----------------------------------------------------------------------------
 # FunctionPatch
diff --git a/primus/core/patches/patch_runner.py b/primus/core/patches/patch_runner.py
index 1a7eae356..f0ba5f058 100644
--- a/primus/core/patches/patch_runner.py
+++ b/primus/core/patches/patch_runner.py
@@ -16,7 +16,7 @@
 
 from primus.core.patches.context import PatchContext
 from primus.core.patches.patch_registry import PatchRegistry
-from primus.modules.module_utils import error_rank_0, log_rank_0
+from primus.core.utils.module_utils import error_rank_0, log_rank_0
 
 # -----------------------------------------------------------------------------
 # Parse PRIMUS_PATCHES Environment Variable
diff --git a/primus/core/pipeline_parallel/handler/wgrad_handler.py b/primus/core/pipeline_parallel/handler/wgrad_handler.py
index 421cde6f1..b5fb723e6 100644
--- a/primus/core/pipeline_parallel/handler/wgrad_handler.py
+++ b/primus/core/pipeline_parallel/handler/wgrad_handler.py
@@ -8,8 +8,10 @@
 
 from megatron.training.global_vars import get_args
 
+from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import (
+    fwd_bwd_wrapper,
+)
 from primus.core.pipeline_parallel.scheduler.scheduler_node import SchedulerNode
-from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper
 
 
 class WGradRunningCache:
diff --git a/primus/core/pipeline_parallel/scheduler/algorithms/base.py b/primus/core/pipeline_parallel/scheduler/algorithms/base.py
index 05d145bb7..d51277236 100644
--- a/primus/core/pipeline_parallel/scheduler/algorithms/base.py
+++ b/primus/core/pipeline_parallel/scheduler/algorithms/base.py
@@ -11,7 +11,7 @@
     FuncType,
     SchedulerNode,
 )
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 class PipelineScheduleAlgo(ABC):
diff --git a/primus/core/projection/memory_projection/benchmark.py b/primus/core/projection/memory_projection/benchmark.py
index bfcafd466..0b3eb106f 100644
--- a/primus/core/projection/memory_projection/benchmark.py
+++ b/primus/core/projection/memory_projection/benchmark.py
@@ -174,7 +174,7 @@ def _run_bench(
         if reduction_info.get("benchmark_num_experts") is not None:
             bench_module.num_experts = reduction_info["benchmark_num_experts"]
 
-    profiling_results = _run_layer_benchmark(primus_config_bench, overrides, reduction_info)
+    profiling_results = _run_layer_benchmark(primus_config_bench, overrides, reduction_info, args=args)
 
     # Optionally persist the artifact.  We piggy-back on the perf saver so
     # that the JSON is identical-shape to the perf-side artifact (one
diff --git a/primus/core/projection/module_profilers/attention.py b/primus/core/projection/module_profilers/attention.py
index f227a00fa..32b5c9408 100644
--- a/primus/core/projection/module_profilers/attention.py
+++ b/primus/core/projection/module_profilers/attention.py
@@ -10,7 +10,7 @@
 
 from primus.core.projection.base_module_profiler import BaseModuleProfiler
 
-from .utils import benchmark_layer
+from .utils import benchmark_layer, v4_module_inputs
 
 
 class AttentionProfiler(BaseModuleProfiler):
@@ -333,13 +333,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float,
                 # Effective sequence length per rank if CP is used
                 slen_per_cp = seq_len // cp_size
 
-                self._cached_results = benchmark_layer(
-                    self.module,
-                    [
-                        (seq_len, batch_size, self.config.model_config.hidden_size),
-                        ((1, 1, slen_per_cp, seq_len), torch.bool),
-                    ],
-                )
+                hidden = self.config.model_config.hidden_size
+                tcfg = getattr(self.module, "config", None)
+                hc_mult = getattr(tcfg, "hc_mult", 1)
+                # DeepSeek-V4 attention has a different signature
+                # (forward(hidden[B,S,D], position_ids[B,S])); feed V4-aware
+                # inputs so the real V4 attention path is exercised instead of
+                # crashing / falling back. Non-V4 modules use the stock inputs.
+                v4 = v4_module_inputs(self.module, batch_size, seq_len, hidden, hc_mult, "attention")
+                if v4 is not None:
+                    ishapes, fkwargs = v4
+                    self._cached_results = benchmark_layer(
+                        self.module, ishapes, transformer_config=tcfg, forward_kwargs=fkwargs
+                    )
+                else:
+                    self._cached_results = benchmark_layer(
+                        self.module,
+                        [
+                            (seq_len, batch_size, hidden),
+                            ((1, 1, slen_per_cp, seq_len), torch.bool),
+                        ],
+                    )
             self._cache_key = cache_key
         return self._cached_results
 
diff --git a/primus/core/projection/module_profilers/moe_mlp.py b/primus/core/projection/module_profilers/moe_mlp.py
index 2aa802446..fbe5ad0be 100644
--- a/primus/core/projection/module_profilers/moe_mlp.py
+++ b/primus/core/projection/module_profilers/moe_mlp.py
@@ -11,7 +11,7 @@
 from primus.core.projection.profiler_spec import ModuleProfilerSpec
 from primus.core.projection.training_config import TrainingConfig
 
-from .utils import benchmark_moe_layer_decomposed
+from .utils import benchmark_layer, benchmark_moe_layer_decomposed, v4_module_inputs
 
 # Efficiency fractions for non-GEMM MoE overhead estimation.
 # These express achievable bandwidth as a fraction of peak HBM bandwidth.
@@ -349,10 +349,25 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float,
                 self._a2a_fwd_ms = 0.0
                 self._a2a_bwd_ms = 0.0
             else:
-                fwd, bwd, act_mem, a2a_fwd, a2a_bwd = benchmark_moe_layer_decomposed(
-                    self.module,
-                    [(seq_len, batch_size, self.config.model_config.hidden_size)],
-                )
+                hidden = self.config.model_config.hidden_size
+                tcfg = getattr(self.module, "config", None)
+                # DeepSeek-V4 MoE: forward(hidden[B,S,D], *, token_ids[B,S]).
+                # Feed V4-aware inputs (right layout + token_ids for hash routing).
+                v4 = v4_module_inputs(self.module, batch_size, seq_len, hidden, 1, "moe")
+                if v4 is not None:
+                    # DeepseekV4MoE has no stock .dispatch/.combine to decompose
+                    # A2A; benchmark the whole MoE forward. At EP=1 (single-GPU
+                    # benchmark) A2A is ~0 and is restored analytically later.
+                    ishapes, fkwargs = v4
+                    fwd, bwd, act_mem = benchmark_layer(
+                        self.module, ishapes, transformer_config=tcfg, forward_kwargs=fkwargs
+                    )
+                    a2a_fwd = a2a_bwd = 0.0
+                else:
+                    fwd, bwd, act_mem, a2a_fwd, a2a_bwd = benchmark_moe_layer_decomposed(
+                        self.module,
+                        [(seq_len, batch_size, hidden)],
+                    )
                 self._cached_results = (fwd, bwd, act_mem)
                 self._a2a_fwd_ms = a2a_fwd
                 self._a2a_bwd_ms = a2a_bwd
diff --git a/primus/core/projection/module_profilers/transformer_layer.py b/primus/core/projection/module_profilers/transformer_layer.py
index bb9476f3d..721924647 100644
--- a/primus/core/projection/module_profilers/transformer_layer.py
+++ b/primus/core/projection/module_profilers/transformer_layer.py
@@ -20,11 +20,7 @@
 from .moe_mlp import MoEMLPProfiler
 from .residual_add import ResidualAddProfiler
 from .router import RouterProfiler
-from .utils import (
-    _install_balanced_routing_patches,
-    _kernel_pad_enabled,
-    benchmark_layer,
-)
+from .utils import benchmark_layer, v4_module_inputs
 
 # ── Fallback HBM bandwidth for elementwise overhead estimation ──
 _FALLBACK_HBM_BW_GBPS = 5300.0  # MI300X default
@@ -321,11 +317,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float,
             else:
                 # Get TransformerConfig from the layer module itself (has fp8 setting)
                 transformer_config = getattr(self.layer_module, "config", None)
-                self._cached_results = benchmark_layer(
-                    self.layer_module,
-                    [(seq_len, batch_size, self.config.model_config.hidden_size)],
-                    transformer_config=transformer_config,
-                )
+                hidden = self.config.model_config.hidden_size
+                hc_mult = getattr(transformer_config, "hc_mult", 1)
+                # DeepSeek-V4 hybrid layer needs K-stream input [B,S,K,D] and a
+                # keyword position_ids; feed V4-aware inputs so the real V4 layer
+                # (mHC + V4 attention) is benchmarked. Non-V4 layers use the stock
+                # [S,B,D] input.
+                v4 = v4_module_inputs(self.layer_module, batch_size, seq_len, hidden, hc_mult, "layer")
+                if v4 is not None:
+                    ishapes, fkwargs = v4
+                    self._cached_results = benchmark_layer(
+                        self.layer_module,
+                        ishapes,
+                        transformer_config=transformer_config,
+                        forward_kwargs=fkwargs,
+                    )
+                else:
+                    self._cached_results = benchmark_layer(
+                        self.layer_module,
+                        [(seq_len, batch_size, hidden)],
+                        transformer_config=transformer_config,
+                    )
             self._cache_key = cache_key
         return self._cached_results
 
@@ -475,24 +487,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float,
             ):
                 # Legacy whole-layer timing (often ~1.5-1.7x pessimistic on backward).
                 transformer_config = getattr(self.layer_module, "config", None)
-                routing_restores = []
-                if _kernel_pad_enabled():
-                    routing_restores, _ = _install_balanced_routing_patches(self.layer_module)
-                try:
+                hidden = self.config.model_config.hidden_size
+                hc_mult = getattr(transformer_config, "hc_mult", 1)
+                # DeepSeek-V4 hybrid layer needs K-stream input [B,S,K,D] and a
+                # keyword position_ids; feed V4-aware inputs so the real V4 layer
+                # (mHC + V4 attention) is benchmarked. Non-V4 layers use the stock
+                # [S,B,D] input.
+                v4 = v4_module_inputs(self.layer_module, batch_size, seq_len, hidden, hc_mult, "layer")
+                if v4 is not None:
+                    ishapes, fkwargs = v4
                     self._cached_results = benchmark_layer(
                         self.layer_module,
-                        [(seq_len, batch_size, self.config.model_config.hidden_size)],
+                        ishapes,
+                        transformer_config=transformer_config,
+                        forward_kwargs=fkwargs,
+                    )
+                else:
+                    self._cached_results = benchmark_layer(
+                        self.layer_module,
+                        [(seq_len, batch_size, hidden)],
                         transformer_config=transformer_config,
                     )
-                finally:
-                    for restore in routing_restores:
-                        try:
-                            restore()
-                        except Exception:
-                            # Best-effort cleanup: restore failures should not fail benchmarking.
-                            _LOGGER.debug("Failed to restore routing patch during cleanup.", exc_info=True)
-            else:
-                self._cached_results = self._get_benchmark_composite_results(batch_size, seq_len)
             self._cache_key = cache_key
         return self._cached_results
 
diff --git a/primus/core/projection/module_profilers/utils.py b/primus/core/projection/module_profilers/utils.py
index 504c7d890..f0a98a6a5 100644
--- a/primus/core/projection/module_profilers/utils.py
+++ b/primus/core/projection/module_profilers/utils.py
@@ -97,11 +97,53 @@ def _get_fp8_context_for_benchmark(transformer_config):
     return _FP8ContextFactory(transformer_config)
 
 
+def v4_module_inputs(module, batch_size, seq_len, hidden_size, hc_mult, kind):
+    """Build DeepSeek-V4-aware benchmark inputs, or None for non-V4 modules.
+
+    The generic harness feeds stock-attention inputs ``(hidden[S,B,D],
+    bool attention_mask)``, but the V4 modules have different signatures:
+
+    * ``DeepseekV4Attention.forward(hidden[B,S,D], position_ids[B,S])`` —
+      both positional; needs integer position_ids (not a bool mask).
+    * ``DeepseekV4HybridLayer.forward(hidden[B,S,K,D], attention_mask=None,
+      *, position_ids[B,S])`` — K=hc_mult parallel mHC streams, and
+      position_ids is keyword-only.
+
+    Returns ``(input_shapes, forward_kwargs)`` or ``None``.
+    """
+    cls = type(module).__name__
+    if kind == "attention" and "DeepseekV4Attention" in cls:
+        return (
+            [(batch_size, seq_len, hidden_size), ((batch_size, seq_len), torch.int64)],
+            None,
+        )
+    if kind == "layer" and "DeepseekV4HybridLayer" in cls:
+        k = max(1, int(hc_mult or 1))
+        ishapes = [(batch_size, seq_len, k, hidden_size)] if k > 1 else [(batch_size, seq_len, hidden_size)]
+        # position_ids for RoPE; token_ids required by hash-routed MoE layers
+        # (ignored by non-hash layers). Both keyword-only on the V4 layer forward.
+        return (
+            ishapes,
+            {
+                "position_ids": ((batch_size, seq_len), torch.int64),
+                "token_ids": ((batch_size, seq_len), torch.int64),
+            },
+        )
+    if kind == "moe" and "DeepseekV4MoE" in cls:
+        # forward(hidden[B,S,D], *, token_ids[B,S]) -> [B,S,D]
+        return (
+            [(batch_size, seq_len, hidden_size)],
+            {"token_ids": ((batch_size, seq_len), torch.int64)},
+        )
+    return None
+
+
 def benchmark_layer(
     layer_module: torch.nn.Module,
     input_shapes: List[Union[Tuple[int, ...], Tuple[Tuple[int, ...], torch.dtype]]],
     num_iterations: int = 64,  # Match typical microbatch count
     transformer_config=None,  # Optional: pass config to enable FP8 context
+    forward_kwargs=None,  # Optional: dict name -> shape/(shape,dtype) passed as keywords
 ) -> tuple[float, float, int]:
     """
     Benchmark both forward and backward passes of a transformer layer using CUDA events.
@@ -152,6 +194,7 @@ def create_input(spec):
             )
 
     inputs = [create_input(spec) for spec in input_shapes]
+    kwargs = {name: create_input(spec) for name, spec in (forward_kwargs or {}).items()}
 
     # ===========================================================================
     # Get FP8 context - CRITICAL for accurate FP8 timing!
@@ -168,7 +211,7 @@ def create_input(spec):
 
     with fp8_context:
         for _ in range(num_warmup):
-            outputs = layer_module(*inputs)
+            outputs = layer_module(*inputs, **kwargs)
             if not isinstance(outputs, (tuple, list)):
                 outputs = (outputs,)
 
@@ -213,7 +256,7 @@ def create_input(spec):
 
     with fp8_context:
         for _ in range(num_iterations):
-            outputs = layer_module(*inputs)
+            outputs = layer_module(*inputs, **kwargs)
 
     if device.type == "cuda":
         torch.cuda.synchronize(device)
@@ -237,7 +280,7 @@ def create_input(spec):
             forward_end = torch.cuda.Event(enable_timing=True)
 
             forward_start.record()
-            outputs = layer_module(*inputs)
+            outputs = layer_module(*inputs, **kwargs)
             forward_end.record()
 
             # --- Backward pass ---
@@ -391,6 +434,7 @@ def benchmark_moe_layer_decomposed(
     input_shapes: List[Union[Tuple[int, ...], Tuple[Tuple[int, ...], torch.dtype]]],
     num_iterations: int = 64,
     transformer_config=None,
+    forward_kwargs=None,  # Optional: dict name -> shape/(shape,dtype) passed as keywords
 ) -> tuple[float, float, int, float, float]:
     """
     Benchmark an MoE layer with decomposed A2A timing.
@@ -453,47 +497,92 @@ def create_input(spec):
             )
 
     inputs = [create_input(spec) for spec in input_shapes]
+    kwargs = {name: create_input(spec) for name, spec in (forward_kwargs or {}).items()}
     fp8_context = _get_fp8_context_for_benchmark(transformer_config)
 
     is_rank_0 = int(os.getenv("RANK", "0")) == 0
 
-    # Install kernel-shape-padding (balanced round-robin routing) so the
-    # per-expert M dimension is constant across iterations.  This eliminates
-    # per-shape JIT recompiles in the FP8 grouped GEMM and is the dominant
-    # cost on high-expert MoE benches.
-    routing_restores = []
-    if _kernel_pad_enabled():
-        routing_restores, routing_descriptors = _install_balanced_routing_patches(moe_module)
-        if routing_descriptors and is_rank_0:
-            _, topk, num_experts = routing_descriptors[0]
-            # The bench feeds a single [seq_len, batch_size, hidden] input,
-            # so the per-rank token count is the product of the first two
-            # dims (TP/SP sharding is already reflected in the supplied
-            # shape).
-            num_tokens = 1
-            for spec in input_shapes:
-                shape = (
-                    spec[0]
-                    if isinstance(spec, tuple) and len(spec) == 2 and isinstance(spec[1], torch.dtype)
-                    else spec
-                )
-                # [seq, batch, hidden]
-                if len(shape) >= 2:
-                    num_tokens = int(shape[0]) * int(shape[1])
-                break
-            m_per_expert = (num_tokens * topk) // max(num_experts, 1)
-            print(
-                f"  [MoE Decomposed] Kernel-shape padding ON: balanced "
-                f"round-robin routing on {len(routing_descriptors)} router(s) "
-                f"(topk={topk}, num_experts={num_experts}, "
-                f"M_per_expert={m_per_expert}). "
-                f"Set PRIMUS_BENCH_MOE_KERNEL_PAD=0 to disable."
-            )
-        elif is_rank_0:
-            print(
-                "  [MoE Decomposed] Kernel-shape padding requested but no "
-                "router-like submodule found; falling back to stochastic routing."
-            )
+    with fp8_context:
+        for _ in range(num_warmup):
+            outputs = moe_module(*inputs, **kwargs)
+            if not isinstance(outputs, (tuple, list)):
+                outputs = (outputs,)
+
+            if grad_outputs is None:
+                grad_outputs = []
+                for i, out in enumerate(outputs):
+                    if isinstance(out, torch.Tensor) and out.requires_grad:
+                        grad_outputs.append(torch.randn_like(out))
+                        output_indices.append(i)
+
+            valid_outputs = [outputs[i] for i in output_indices]
+            if valid_outputs:
+                torch.autograd.backward(valid_outputs, grad_outputs)
+
+            moe_module.zero_grad(set_to_none=True)
+            for inp in inputs:
+                if inp.requires_grad:
+                    inp.grad = None
+
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+    # --- Measure activation memory (forward-only loop) ---
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats(device)
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    mem_before = torch.cuda.memory_allocated(device)
+
+    with fp8_context:
+        for _ in range(num_iterations):
+            outputs = moe_module(*inputs, **kwargs)
+
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    mem_after_forward = torch.cuda.max_memory_allocated(device)
+    activation_memory = (mem_after_forward - mem_before) // num_iterations
+
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats(device)
+    del outputs
+
+    # =========================================================================
+    # BENCHMARK with decomposed A2A timing
+    # =========================================================================
+    # Monkey-patch dispatch() and combine() to insert CUDA events.
+    # MoELayer.forward() calls self.dispatch(...) and self.combine(...)
+    # so instance-attribute patches are picked up by Python's MRO.
+    original_dispatch = moe_module.dispatch
+    original_combine = moe_module.combine
+
+    # Accumulate (start_event, end_event) pairs per iteration
+    _dispatch_events = []  # one (start, end) per iteration
+    _combine_events = []
+
+    def timed_dispatch(*args, **kwargs):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        result = original_dispatch(*args, **kwargs)
+        end.record()
+        _dispatch_events.append((start, end))
+        return result
+
+    def timed_combine(*args, **kwargs):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        result = original_combine(*args, **kwargs)
+        end.record()
+        _combine_events.append((start, end))
+        return result
+
+    moe_module.dispatch = timed_dispatch
+    moe_module.combine = timed_combine
+
+    forward_times = []
+    backward_times = []
 
     try:
         # =====================================================================
@@ -504,8 +593,19 @@ def create_input(spec):
         num_warmup, num_iterations = _bench_iter_count(20, num_iterations)
 
         with fp8_context:
-            for _ in range(num_warmup):
-                outputs = moe_module(*inputs)
+            for _ in range(num_iterations):
+                # --- Forward pass ---
+                forward_start = torch.cuda.Event(enable_timing=True)
+                forward_end = torch.cuda.Event(enable_timing=True)
+
+                forward_start.record()
+                outputs = moe_module(*inputs, **kwargs)
+                forward_end.record()
+
+                # --- Backward pass ---
+                backward_start = torch.cuda.Event(enable_timing=True)
+                backward_end = torch.cuda.Event(enable_timing=True)
+
                 if not isinstance(outputs, (tuple, list)):
                     outputs = (outputs,)
 
diff --git a/primus/core/projection/performance_projection/projection.py b/primus/core/projection/performance_projection/projection.py
index 5f9a7d57a..28d97b31f 100644
--- a/primus/core/projection/performance_projection/projection.py
+++ b/primus/core/projection/performance_projection/projection.py
@@ -39,9 +39,9 @@
     convert_primus_config_to_projection_config,
 )
 
-# NOTE: MegatronPretrainTrainer is imported lazily inside _run_layer_benchmark()
-# to avoid pulling in the megatron dependency when running in pure simulation mode
-# (--profiling-mode simulate).
+# NOTE: The core runtime (PrimusRuntime) and megatron backend are imported
+# lazily inside _run_layer_benchmark() to avoid pulling in the megatron
+# dependency when running in pure simulation mode (--profiling-mode simulate).
 
 _MAX_EXPERT_PARALLEL_SIZE = 8
 _BYTES_PER_GB = 1024**3
@@ -1056,16 +1056,23 @@ def _limit_layers_for_projection(module_config):
     original_layers = getattr(module_config, "num_layers", 1) or 1
     original_moe_layout = getattr(module_config, "moe_layer_freq", None)
     dense_layers_present = _has_dense_layers(original_moe_layout)
-
-    if has_moe and dense_layers_present:
-        # Need at least 2 layers to profile both dense (layer 0) and MoE (layer 1)
-        # so extraction code can correctly classify each type using the full
-        # model's moe_pattern where layer 0 is typically dense.
-        max_layers = 2
-    else:
-        max_layers = 1
+    # Use 1 layer for fast profiling - results are extrapolated to full model
+    # Increase to 2-4 for better accuracy if needed. PRIMUS_PROJ_MAX_LAYERS lets
+    # the caller benchmark more representative layers (e.g. =2 to capture both the
+    # DeepSeek-V4 HCA(cr=128) and CSA(cr=4) attention kinds, which alternate).
+    max_layers = int(os.environ.get("PRIMUS_PROJ_MAX_LAYERS", "1"))
     target_layers = max(1, min(original_layers, max_layers))
     module_config.num_layers = target_layers
+    # Set the benchmarked layers' attention kinds. PRIMUS_PROJ_COMPRESS_RATIOS
+    # (e.g. "128,4") picks exactly one HCA + one CSA; otherwise trim the model's
+    # own schedule to the benchmarked layer count.
+    _cr_env = os.environ.get("PRIMUS_PROJ_COMPRESS_RATIOS", "").strip("[] ")
+    if _cr_env and hasattr(module_config, "compress_ratios"):
+        module_config.compress_ratios = [int(x) for x in _cr_env.split(",")][:target_layers]
+    else:
+        _cr = getattr(module_config, "compress_ratios", None)
+        if isinstance(_cr, (list, tuple)) and len(_cr) >= target_layers:
+            module_config.compress_ratios = list(_cr[:target_layers])
 
     if has_moe:
         if not dense_layers_present:
@@ -2536,9 +2543,43 @@ def _summarize_bench_training_config(training_config) -> Dict[str, Any]:
     }
 
 
-def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None):
-    from primus.modules.trainer.megatron.pre_trainer import MegatronPretrainTrainer
+def _build_runtime_primus_config(legacy_primus_config, args, module_name="pre_trainer"):
+    """Adapt the (already-mutated) legacy PrimusConfig into a runtime-shaped config.
+
+    Performance projection loads config via the legacy ``PrimusConfig`` and mutates
+    the flat ``pre_trainer`` module in place (layer limiting, EP rescale, turbo /
+    overlap overrides, ...). The core runtime, however, consumes the normalized
+    ``cfg.modules[*].params`` shape. This helper reuses the runtime loader to build
+    a correct scaffold (exp_root_path / exp_meta_info / platform / ...), then swaps
+    in the mutated ``pre_trainer`` module normalized to the runtime shape, so the
+    benchmark model is built from exactly the mutated config.
+    """
+    from pathlib import Path
+
+    from primus.core.config.primus_config import (
+        _normalize_module_for_runtime,
+        load_primus_config,
+    )
+
+    runtime_cfg = load_primus_config(Path(args.config), args)
+    normalized = _normalize_module_for_runtime(
+        legacy_primus_config.get_module_config(module_name), module_name
+    )
+
+    modules = list(getattr(runtime_cfg, "modules", []) or [])
+    replaced = False
+    for i, m in enumerate(modules):
+        if getattr(m, "name", None) == module_name:
+            modules[i] = normalized
+            replaced = True
+            break
+    if not replaced:
+        modules.append(normalized)
+    runtime_cfg.modules = modules
+    return runtime_cfg
+
 
+def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None, args=None):
     module_config = primus_config.get_module_config("pre_trainer")
     _limit_layers_for_projection(module_config)
     rescale_info = _rescale_expert_parallelism(module_config)
@@ -2549,8 +2590,6 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None):
     if reduction_info is not None:
         reduction_info["bench_training_config_summary"] = _summarize_bench_training_config(training_config)
 
-    master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
-    master_port = int(os.getenv("MASTER_PORT", "29500"))
     rank = int(os.getenv("RANK", "0"))
     world_size = int(os.getenv("WORLD_SIZE", "1"))
 
@@ -2560,7 +2599,7 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None):
     mem_recorder = MemoryBenchmarkRecorder(rank=rank)
     mem_recorder.snapshot("pre_trainer_init")
 
-    print("[Primus:Performance Projection] Initializing MegatronPretrainTrainer...")
+    print("[Primus:Performance Projection] Preparing benchmark model build...")
     # Disable overlap features and FSDP2 for profiling (they add complexity without benefiting isolated layer benchmarking)
     # FSDP2 uses DTensor which causes issues with benchmarking inputs
     cfg = primus_config.get_module_config("pre_trainer")
@@ -2599,22 +2638,32 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None):
     if getattr(cfg, "multi_latent_attention", False):
         print(f"  enable_primus_turbo: {cfg.enable_primus_turbo}")
         print(f"  use_turbo_gemm: {cfg.use_turbo_gemm}")
-    trainer = MegatronPretrainTrainer(
+
+    # Build the model via the new core runtime instead of the legacy trainer.
+    # The runtime reuses the exact real-training init pipeline (adapter config
+    # conversion + build_args/setup/before_train patches + megatron
+    # initialize_megatron + setup_model_and_optimizer) but stops before the
+    # training loop, giving us a model identical to real training for the
+    # per-layer benchmark. `args` provides data_path / backend_path / config.
+    from primus.core.runtime.train_runtime import PrimusRuntime
+
+    if args is None:
+        raise ValueError(
+            "[Primus:Performance Projection] _run_layer_benchmark requires `args` "
+            "(CLI namespace with config/data_path/backend_path) to build the model "
+            "via the core runtime."
+        )
+
+    runtime_primus_config = _build_runtime_primus_config(primus_config, args, module_name="pre_trainer")
+
+    print("[Primus:Performance Projection] Initializing Megatron and building model...")
+    runtime = PrimusRuntime(args=args)
+    trainer = runtime.setup_model_only(
         module_name="pre_trainer",
-        primus_config=primus_config,
-        module_rank=rank,
-        module_world_size=world_size,
-        module_master_addr=master_addr,
-        module_master_port=master_port,
-        extra_args=unknown_overrides,
+        overrides=unknown_overrides,
+        primus_config=runtime_primus_config,
     )
-
-    print("[Primus:Performance Projection] Initializing Megatron...")
-    trainer.init()
     mem_recorder.snapshot("post_megatron_init")
-
-    print("[Primus:Performance Projection] Setting up model and optimizer...")
-    trainer.setup()
     # post_setup captures: params + distributed-optimizer state + grad buffers
     # at the bench config.  This is the "static" memory anchor — what every
     # rank pays before any forward pass.
@@ -3930,7 +3979,7 @@ def launch_projection_from_cli(args, overrides):
         # downstream pipeline simulation / multinode projection, but print
         # a side-by-side comparison.
         sim_results = _run_layer_simulation(copy.deepcopy(primus_config), args)
-        bench_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info)
+        bench_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info, args=args)
 
         is_rank_0 = int(os.getenv("RANK", "0")) == 0
         if is_rank_0:
@@ -3960,7 +4009,7 @@ def launch_projection_from_cli(args, overrides):
         profiling_results = bench_results
     else:
         # Default: actual GPU benchmark
-        profiling_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info)
+        profiling_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info, args=args)
 
     # ── Save bench artifact if requested ──
     # ``--save-benchmark`` (preferred) and ``--save-profiling`` (deprecated)
diff --git a/primus/core/projection/simulation_backends/sdpa_simulator.py b/primus/core/projection/simulation_backends/sdpa_simulator.py
index cd1629aaa..61a689162 100644
--- a/primus/core/projection/simulation_backends/sdpa_simulator.py
+++ b/primus/core/projection/simulation_backends/sdpa_simulator.py
@@ -141,20 +141,20 @@ class GPUHardwareSpec:
         n_xcd=8,
     ),
     "mi355x": GPUHardwareSpec(
-        peak_tflops_bf16=2384.0,
-        peak_tflops_fp16=2384.0,
-        peak_tflops_fp8=4768.0,
+        peak_tflops_bf16=2500.0,
+        peak_tflops_fp16=2500.0,
+        peak_tflops_fp8=5000.0,
         hbm_bandwidth_gbps=8000.0,
         n_cu=256,
-        n_xcd=4,
+        n_xcd=8,
     ),
     "gfx950": GPUHardwareSpec(  # same as MI355X
-        peak_tflops_bf16=2384.0,
-        peak_tflops_fp16=2384.0,
-        peak_tflops_fp8=4768.0,
+        peak_tflops_bf16=2500.0,
+        peak_tflops_fp16=2500.0,
+        peak_tflops_fp8=5000.0,
         hbm_bandwidth_gbps=8000.0,
         n_cu=256,
-        n_xcd=4,
+        n_xcd=8,
     ),
 }
 
diff --git a/primus/core/runtime/logging.py b/primus/core/runtime/logging.py
index 0e56a86f2..bb2cb918c 100644
--- a/primus/core/runtime/logging.py
+++ b/primus/core/runtime/logging.py
@@ -33,7 +33,7 @@
 
 from primus.core.utils import logger
 from primus.core.utils.env import get_torchrun_env
-from primus.modules.module_utils import debug_rank_all, set_logging_rank
+from primus.core.utils.module_utils import debug_rank_all, set_logging_rank
 
 # from primus.core.utils.distributed_logging import debug_rank_all, set_logging_rank
 
diff --git a/primus/core/runtime/runtime_state.py b/primus/core/runtime/runtime_state.py
new file mode 100644
index 000000000..ffecb310c
--- /dev/null
+++ b/primus/core/runtime/runtime_state.py
@@ -0,0 +1,34 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Runtime State Management for Training.
+
+This module provides a dedicated RuntimeState object for storing dynamic
+runtime metrics that change during training (e.g., image dimensions, timesteps).
+This is separate from backend_args (which is for configuration parameters).
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict
+
+
+@dataclass
+class RuntimeState:
+    """
+    Dedicated runtime state object for per-iteration training metrics.
+
+    This is separate from backend_args (which is for configuration) and
+    provides a clean place to store dynamic runtime metrics that change
+    during training (e.g., image dimensions, timesteps, etc.).
+    """
+
+    # Per-iteration metrics (updated each forward step)
+    last_metrics: Dict[str, Any] = field(default_factory=dict)
+
+    def update_metrics(self, metrics: Dict[str, Any]) -> None:
+        """Update last_metrics with new metrics from forward step."""
+        self.last_metrics.update(metrics)
diff --git a/primus/core/runtime/train_runtime.py b/primus/core/runtime/train_runtime.py
index fbcedd991..f314e956a 100644
--- a/primus/core/runtime/train_runtime.py
+++ b/primus/core/runtime/train_runtime.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
 #
 # Primus Runtime Orchestrator for Training
 #
@@ -23,14 +23,15 @@
 )
 from primus.core.patches import run_patches
 from primus.core.runtime.logging import init_worker_logger
+from primus.core.runtime.runtime_state import RuntimeState
 from primus.core.utils.arg_utils import parse_cli_overrides
 from primus.core.utils.env_setup import setup_training_env
+from primus.core.utils.module_utils import log_dict_aligned, log_rank_0, warning_rank_0
 from primus.core.utils.yaml_utils import (
     dict_to_nested_namespace,
     merge_namespace,
     nested_namespace_to_dict,
 )
-from primus.modules.module_utils import log_dict_aligned, log_rank_0, warning_rank_0
 
 # ---------------------------------------------------------------------------
 # Context & Hooks
@@ -56,6 +57,7 @@ class TrainContext:
     trainer: Any = None
     backend_args: Any = None
     backend_version: Optional[str] = None
+    runtime_state: Optional[RuntimeState] = None
 
     # Distributed context
     rank: int = 0
@@ -109,6 +111,64 @@ def run_train_module(self, module_name: str, overrides: Optional[List[str]] = No
             self._safe_cleanup(error=e)
             raise RuntimeError(f"Training execution failed: {e}") from e
 
+    def setup_model_only(
+        self,
+        module_name: str = "pre_trainer",
+        overrides: Optional[List[str]] = None,
+        primus_config: Any = None,
+    ) -> Any:
+        """Run the runtime pipeline but stop after building the model (no training).
+
+        A general, training-neutral entry that mirrors :meth:`run_train_module`
+        except it replaces the training step with a model-only build. Reuses the
+        full runtime initialization pipeline (config → environment → distributed →
+        logging → backend adapter → args conversion → build_args patches → trainer
+        setup/init → before_train patches) and then asks the trainer to build only
+        the model via ``setup_model_only``. Useful for any "build the model only"
+        scenario (offline profiling, layer benchmarking, model inspection);
+        performance/memory projection is the current consumer.
+
+        Args:
+            module_name: training module to build (default ``pre_trainer``).
+            overrides: CLI-style overrides to merge into the module params.
+            primus_config: optional pre-loaded (runtime-shaped) PrimusConfig. When
+                provided, config is not re-loaded from ``self.args.config`` — this
+                lets callers (e.g. projection) apply their own config mutations
+                before building.
+
+        Returns:
+            The trainer instance whose ``.model`` has been built.
+        """
+        overrides = overrides or []
+
+        # 1) Configuration (optionally injected, so callers can pre-mutate it)
+        self._initialize_configuration(module_name, overrides, primus_config=primus_config)
+        # 2) Runtime environment (paths, distributed, logging)
+        self._initialize_runtime_environment()
+        # 3) Backend adapter + trainer (convert config, build_args patches, instantiate)
+        self._initialize_adapter()
+        self._initialize_trainer()
+
+        assert self.ctx is not None and self.ctx.trainer is not None
+        trainer = self.ctx.trainer
+
+        # 4) Setup + init + before_train patches (mirror _run_trainer_lifecycle up
+        #    to, but excluding, the training step).
+        self._run_phase_patches(phase="setup", backend_args=self.ctx.backend_args)
+        trainer.setup()
+        trainer.init()
+        self._run_phase_patches(phase="before_train", backend_args=self.ctx.backend_args)
+
+        # 5) Build only the model (no datasets / no train loop).
+        build_fn = getattr(trainer, "setup_model_only", None)
+        if build_fn is None:
+            raise NotImplementedError(
+                f"Trainer '{type(trainer).__name__}' for framework "
+                f"'{self.ctx.framework}' does not implement setup_model_only()."
+            )
+        build_fn()
+        return trainer
+
     # --------------------------- Internal Steps --------------------------- #
 
     def _initialize_runtime_environment(self) -> None:
@@ -135,7 +195,9 @@ def _get_backend_version(self) -> Optional[str]:
             self.ctx.backend_version = None
         return self.ctx.backend_version
 
-    def _run_phase_patches(self, phase: str, backend_args: Any = None) -> None:
+    def _run_phase_patches(
+        self, phase: str, backend_args: Any = None, runtime_state: Optional[RuntimeState] = None
+    ) -> None:
         """
         Apply a patch phase in a single, runtime-owned place.
 
@@ -156,6 +218,7 @@ def _run_phase_patches(self, phase: str, backend_args: Any = None) -> None:
                 "backend_args": backend_args,
                 "primus_config": self.ctx.primus_config,
                 "module_config": self.ctx.module_config,
+                "runtime_state": runtime_state,
             },
         )
 
@@ -168,11 +231,28 @@ def _initialize_environment(self) -> None:
         # setup_training_env expects a string path.
         setup_training_env(str(data_path), setup_hf=True)
 
-    def _initialize_configuration(self, module_name: str, overrides: Optional[List[str]] = None) -> None:
-        cfg_path = Path(self.args.config)
-        assert cfg_path.exists(), f"[Primus:TrainRuntime] Config file not found: {cfg_path}"
-
-        primus_cfg = load_primus_config(cfg_path, self.args)
+    def _initialize_configuration(
+        self,
+        module_name: str,
+        overrides: Optional[List[str]] = None,
+        primus_config: Any = None,
+    ) -> None:
+        # Resolve a config path for diagnostics/context (falls back to args.config
+        # when a pre-loaded config is injected and args may lack a usable path).
+        cfg_path = Path(getattr(self.args, "config", None) or getattr(primus_config, "config_file", "") or "")
+        if primus_config is not None:
+            # Caller supplied a pre-loaded (and possibly mutated) runtime config;
+            # skip re-loading from disk so those mutations are preserved.
+            primus_cfg = primus_config
+        else:
+            assert cfg_path.exists(), f"[Primus:TrainRuntime] Config file not found: {cfg_path}"
+            primus_cfg = load_primus_config(cfg_path, self.args)
+
+        # Reuse the legacy PrimusConfig that load_primus_config already parsed
+        # (exposed as `_legacy`) instead of re-parsing the YAML a second time.
+        # BaseModule needs the PrimusConfig interface (get_module_config / global
+        # vars); the SimpleNamespace `primus_cfg` drives the new core runtime.
+        primus_config_obj = primus_cfg._legacy
 
         # Resolve module configuration via core helper.
         module_cfg = get_module_config(primus_cfg, module_name)
@@ -188,11 +268,15 @@ def _initialize_configuration(self, module_name: str, overrides: Optional[List[s
             raise ValueError(f"[Primus:TrainRuntime] Module '{module_name}' missing 'framework'.")
 
         # Initialize TrainContext based on raw configuration (before CLI overrides).
+        # Use primus_config_obj (PrimusConfig) for BaseModule compatibility
+        _data_path = getattr(self.args, "data_path", None)
+        if not _data_path:
+            _data_path = "./data"
         self.ctx = TrainContext(
             config_path=cfg_path,
-            data_path=Path(getattr(self.args, "data_path", "./data")),
+            data_path=Path(_data_path),
             module_name=module_name,
-            primus_config=primus_cfg,
+            primus_config=primus_config_obj,  # Use PrimusConfig object, not SimpleNamespace
             module_config=module_cfg,
             framework=framework,
         )
@@ -247,6 +331,10 @@ def _initialize_adapter(self) -> None:
         assert self.ctx is not None, "TrainContext must be initialized before backend adapter."
         backend_path = getattr(self.args, "backend_path", None)
 
+        # CRITICAL: Set up backend path BEFORE importing backend module
+        # This ensures megatron is importable when patches are loaded
+        self._setup_backend_path_early(backend=self.ctx.framework, backend_path=backend_path)
+
         adapter = BackendRegistry.get_adapter(backend=self.ctx.framework, backend_path=backend_path)
 
         assert (
@@ -258,6 +346,49 @@ def _initialize_adapter(self) -> None:
 
         self.ctx.adapter = adapter
 
+    def _setup_backend_path_early(self, backend: str, backend_path=None) -> None:
+        """Set up backend path before backend module is imported."""
+        import os
+        import sys
+        from pathlib import Path
+
+        # For Megatron backend, set up the path early
+        if backend == "megatron":
+            megatron_paths = []
+
+            # Method 1: From backend_path argument
+            if backend_path:
+                megatron_path = Path(backend_path)
+                if megatron_path.exists():
+                    megatron_paths.append(str(megatron_path))
+
+            # Method 2: From PRIMUS_PATH environment variable
+            primus_path = os.getenv("PRIMUS_PATH")
+            if primus_path:
+                megatron_path = Path(primus_path) / "third_party" / "Megatron-LM"
+                if megatron_path.exists():
+                    megatron_paths.append(str(megatron_path))
+
+            # Method 3: From current working directory
+            try:
+                cwd = Path.cwd()
+                if "Primus" in str(cwd):
+                    primus_root = cwd
+                    while primus_root.name != "Primus" and primus_root != primus_root.parent:
+                        primus_root = primus_root.parent
+                    if primus_root.name == "Primus":
+                        megatron_path = primus_root / "third_party" / "Megatron-LM"
+                        if megatron_path.exists():
+                            megatron_paths.append(str(megatron_path))
+            except Exception:
+                pass
+
+            # Add paths to sys.path if not already present
+            for path in megatron_paths:
+                if path not in sys.path:
+                    sys.path.insert(0, path)
+                    log_rank_0(f"[Primus:Runtime] Early setup: Added Megatron-LM to sys.path: {path}")
+
     def _initialize_trainer(self) -> None:
         assert (
             self.ctx is not None and self.ctx.adapter is not None
@@ -273,8 +404,14 @@ def _initialize_trainer(self) -> None:
         backend_args = adapter.convert_config(module_config.params)
         self.ctx.backend_args = backend_args
 
+        # Create runtime state object (for dynamic per-iteration metrics)
+        # Initialize early so it's available for all patch phases
+        self.ctx.runtime_state = RuntimeState()
+
         # Phase: build_args (after args creation, before trainer instantiation)
-        self._run_phase_patches(phase="build_args", backend_args=backend_args)
+        self._run_phase_patches(
+            phase="build_args", backend_args=backend_args, runtime_state=self.ctx.runtime_state
+        )
 
         # Log final args after patches, then merge module_config.params into backend_args
         log_dict_aligned("Final backend args (after patches)", backend_args)
@@ -290,18 +427,64 @@ def _initialize_trainer(self) -> None:
             primus_only_params = {key: params_dict[key] for key in sorted(primus_only_keys)}
             log_dict_aligned("Primus-specific parameters", primus_only_params)
 
+        # Extract trainer_class from module_config BEFORE merging params into backend_args
+        # (After merge, module_config.params becomes backend_args, and trainer_class might be lost)
+        trainer_class = None
+
+        # First, check module_config directly (top-level attribute)
+        if hasattr(module_config, "trainer_class"):
+            trainer_class = module_config.trainer_class
+        # Second, check module_config.params (nested in params, BEFORE merge)
+        elif hasattr(module_config.params, "trainer_class"):
+            trainer_class = module_config.params.trainer_class
+
         # Merge backend_args into params (backend_args overrides params)
         merge_namespace(backend_args, module_config.params, allow_override=False, excepts=[])
         module_config.params = backend_args
 
-        # Load trainer class and instantiate
+        # Third, check backend_args (after merge, in case it was preserved)
+        if not trainer_class and hasattr(backend_args, "trainer_class"):
+            trainer_class = backend_args.trainer_class
+
+        # Log summary of trainer_class extraction (keep one summary log for troubleshooting)
+        if trainer_class:
+            log_rank_0(f"[TrainRuntime] Using trainer_class: {trainer_class}")
+        else:
+            log_rank_0(f"[TrainRuntime] WARNING: trainer_class not found, will use default for stage")
+
+        # Extract stage for fallback
         stage = getattr(module_config.params, "stage", "pretrain") or "pretrain"
-        TrainerClass = adapter.load_trainer_class(stage=stage)
-        trainer = TrainerClass(backend_args=backend_args)
+
+        log_rank_0(f"[TrainRuntime] Loading trainer: stage={stage}, trainer_class={trainer_class}")
+        # Only forward trainer_class to adapters that support it (Megatron). Other
+        # backend adapters keep their stage-only signature, so passing trainer_class
+        # unconditionally would raise TypeError on the shared core-runtime path.
+        trainer_cls = adapter.load_trainer_class(
+            stage=stage,
+            **({"trainer_class": trainer_class} if trainer_class else {}),
+        )
+        log_rank_0(f"[TrainRuntime] Loaded trainer class: {trainer_cls.__name__}")
+
+        # Initialize trainer with both BaseModule and BaseTrainer arguments
+        # BaseModule needs: module_name, primus_config, module_rank, module_world_size, module_master_addr, module_master_port
+        # BaseTrainer needs: backend_args
+        # Note: BaseModule.__init__ expects keyword arguments, not positional
+        trainer = trainer_cls(
+            backend_args=backend_args,
+            module_name=self.ctx.module_name,
+            primus_config=self.ctx.primus_config,
+            module_rank=self.ctx.rank,
+            module_world_size=self.ctx.world_size,
+            module_master_addr=self.ctx.master_addr,
+            module_master_port=self.ctx.master_port,
+        )
 
         assert trainer is not None, f"Failed to create trainer instance for framework '{self.ctx.framework}'."
         self.ctx.trainer = trainer
 
+        # Attach runtime_state to trainer instance for easy access in forward_step()
+        trainer.runtime_state = self.ctx.runtime_state
+
     def _run_trainer_lifecycle(self) -> None:
         assert (
             self.ctx is not None and self.ctx.trainer is not None
@@ -320,18 +503,32 @@ def _log_step(step_name: str, func):
         trainer = self.ctx.trainer
 
         # 1) Optional setup phase
-        self._run_phase_patches(phase="setup", backend_args=self.ctx.backend_args)
+        self._run_phase_patches(
+            phase="setup", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state
+        )
         _log_step("Setup", trainer.setup)
 
         # 2) Initialize training components
         _log_step("Init", trainer.init)
 
         # 3) Execute training
-        self._run_phase_patches(phase="before_train", backend_args=self.ctx.backend_args)
-        _log_step("Training", trainer.train)
+        self._run_phase_patches(
+            phase="before_train", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state
+        )
+        import gc
+
+        gc.disable()
+        log_rank_0("Python GC disabled for training (matching NeMo memory management)")
+        try:
+            _log_step("Training", trainer.train)
+        finally:
+            gc.enable()
+            log_rank_0("Python GC re-enabled after training")
 
         # 4) Cleanup and finalize
-        self._run_phase_patches(phase="after_train", backend_args=self.ctx.backend_args)
+        self._run_phase_patches(
+            phase="after_train", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state
+        )
         _log_step("Cleanup", trainer.cleanup)
 
     # --------------------------- Cleanup ---------------------------------- #
diff --git a/primus/core/trainer/base_trainer.py b/primus/core/trainer/base_trainer.py
index 107c44438..2b0a75ce0 100644
--- a/primus/core/trainer/base_trainer.py
+++ b/primus/core/trainer/base_trainer.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -45,26 +45,49 @@ class BaseTrainer(ABC):
         MegatronPretrainTrainer, TorchtitanPretrainTrainer, etc.
     """
 
-    def __init__(self, backend_args: Any):
+    def __init__(self, backend_args: Any = None, *args, **kwargs):
         """
         Initialize base trainer.
 
         Args:
             backend_args: Backend-specific arguments (e.g., from MegatronArgBuilder)
+            *args, **kwargs: Additional arguments (filtered to prevent reaching object.__init__())
 
         Note:
             Distributed environment and logging should be initialized globally
             before creating trainer instances.
         """
-        self.backend_args = backend_args
-
-        # Resolve distributed environment directly from torchrun-style env vars
-        dist_env = get_torchrun_env()
-        self.rank = dist_env["rank"]
-        self.world_size = dist_env["world_size"]
-        self.local_rank = dist_env["local_rank"]
-        self.master_addr = dist_env["master_addr"]
-        self.master_port = dist_env["master_port"]
+        # Filter backend_args from kwargs if not provided as positional/keyword argument
+        if backend_args is None:
+            backend_args = kwargs.pop("backend_args", None)
+
+        try:
+            from abc import ABC
+
+            ABC.__init__(self)
+
+            self.backend_args = backend_args
+
+            dist_env = get_torchrun_env()
+            self.rank = dist_env["rank"]
+            self.world_size = dist_env["world_size"]
+            self.local_rank = dist_env["local_rank"]
+            self.master_addr = dist_env["master_addr"]
+            self.master_port = dist_env["master_port"]
+
+            # Cooperative multiple inheritance: pass kwargs to BaseModule if present in MRO,
+            # otherwise call super().__init__() with no args to avoid object.__init__() error.
+            from primus.core.base_module import BaseModule
+
+            if BaseModule in type(self).__mro__:
+                super().__init__(**kwargs)
+            else:
+                super().__init__()
+        except Exception:
+            import traceback
+
+            traceback.print_exc()
+            raise
 
     @abstractmethod
     def setup(self):
diff --git a/primus/core/utils/import_utils.py b/primus/core/utils/import_utils.py
index 34c2de16d..47b1ed918 100644
--- a/primus/core/utils/import_utils.py
+++ b/primus/core/utils/import_utils.py
@@ -6,7 +6,7 @@
 import importlib
 from functools import partial
 
-from primus.modules.module_utils import log_rank_0
+from primus.core.utils.module_utils import log_rank_0
 
 
 def lazy_import(paths, symbol, log_prefix="[Primus]"):
@@ -39,13 +39,27 @@ def get_model_provider(model_type="gpt"):
     Resolve model_provider across Megatron versions and model types.
 
     Args:
-        model_type (str): Type of model - 'gpt' or 'mamba'. Defaults to 'gpt'.
+        model_type (str): Type of model - 'gpt', 'mamba', or 'deepseek_v4'. Defaults to 'gpt'.
 
     - New:   model_provider + gpt_builder/mamba_builder
     - Mid:   model_provider only
     - Old:   pretrain_gpt.model_provider / pretrain_mamba.model_provider
     """
-    # Try to import model_provider
+    # Primus-owned: DeepSeek-V4 (Phase 2 stub; full V4 wiring lands in Phase 3+)
+    if model_type == "deepseek_v4":
+        deepseek_v4_module = importlib.import_module(
+            "primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders"
+        )
+        log_rank_0(
+            "[Primus][MegatronCompat] Loaded DeepSeek-V4 model_provider + builder "
+            f"from {deepseek_v4_module.__name__}"
+        )
+        return partial(
+            deepseek_v4_module.model_provider,
+            deepseek_v4_module.deepseek_v4_builder,
+        )
+
+    # Upstream Megatron-LM model types
     if model_type == "mamba":
         model_provider = lazy_import(
             ["model_provider", "pretrain_mamba"], "model_provider", log_prefix="[Primus][MegatronCompat]"
diff --git a/primus/modules/module_utils.py b/primus/core/utils/module_utils.py
similarity index 93%
rename from primus/modules/module_utils.py
rename to primus/core/utils/module_utils.py
index 52503d985..11c9621d9 100644
--- a/primus/modules/module_utils.py
+++ b/primus/core/utils/module_utils.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -73,7 +73,14 @@ def log_kv_rank_0(key, value):
         log_func(key, value, module_name, function_name, line)
 
 
-def debug_rank_0(msg, *args, **kwargs):
+def debug_rank_0(*args, **kwargs):
+    if not args:
+        return
+    sep = kwargs.pop("sep", " ")
+    kwargs.pop("end", None)
+    kwargs.pop("file", None)
+    kwargs.pop("flush", None)
+    msg = sep.join(str(a) for a in args)
     log_func = logger.debug_with_caller
 
     caller = inspect.stack()[1]
@@ -86,7 +93,14 @@ def debug_rank_0(msg, *args, **kwargs):
         log_func(msg, module_name, function_name, line)
 
 
-def debug_rank_all(msg, *args, **kwargs):
+def debug_rank_all(*args, **kwargs):
+    if not args:
+        return
+    sep = kwargs.pop("sep", " ")
+    kwargs.pop("end", None)
+    kwargs.pop("file", None)
+    kwargs.pop("flush", None)
+    msg = sep.join(str(a) for a in args)
     log_func = logger.debug_with_caller
 
     caller = inspect.stack()[1]
diff --git a/primus/core/utils/yaml_utils.py b/primus/core/utils/yaml_utils.py
index 1f39b1f2f..27cca8b13 100755
--- a/primus/core/utils/yaml_utils.py
+++ b/primus/core/utils/yaml_utils.py
@@ -8,8 +8,6 @@
 from types import SimpleNamespace
 from typing import Any, Mapping
 
-import yaml
-
 from primus.core.config.merge_utils import deep_merge
 from primus.core.config.yaml_loader import parse_yaml as _parse_yaml_core
 
@@ -140,6 +138,7 @@ def dump_namespace_to_yaml(ns: SimpleNamespace, file_path: str):
         >>> ns = SimpleNamespace(a=1, b=SimpleNamespace(c=2))
         >>> dump_namespace_to_yaml(ns, "config.yaml")
     """
+    import yaml
 
     def ns_to_dict(obj):
         if isinstance(obj, SimpleNamespace):
diff --git a/primus/mlperf_log_suppression.py b/primus/mlperf_log_suppression.py
new file mode 100644
index 000000000..2b11f16e9
--- /dev/null
+++ b/primus/mlperf_log_suppression.py
@@ -0,0 +1,273 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+# ---------------------------------------------------------------------------
+# Non-MLLOG log suppression for MLPerf submission runs.
+#
+# The MLPerf reference/submission logs (":::MLLOG ...") plus the run banners
+# emitted by the launch script are the only required stdout lines for a
+# submission run. All other noisy framework output (Primus loguru banners,
+# Megatron deprecations, aiter JIT build chatter, TE-RoPE, hipify, torchrun
+# warnings, UserWarnings, ...) is suppressed when this module is enabled, to
+# keep the run logs clean.
+#
+# Activation is GATED so that importing this module is a complete no-op for
+# normal (non-MLPerf) Primus runs:
+#
+#   PRIMUS_LOG_SUPPRESSION=1  -> enable the suppression (MLPerf launch sets it)
+#   PRIMUS_LOG_SUPPRESSION=0  -> default: install() is a no-op
+#
+# When enabled, verbosity is further controlled by:
+#
+#   MLPERF_VERBOSE_LOGS=1  -> restore the full verbose output (old behaviour),
+#                            i.e. enabled-but-not-quiet.
+#   MLPERF_VERBOSE_LOGS=0  -> quiet mode (default when enabled): only MLLOG +
+#                            training timing/result lines are emitted.
+#
+# Strategy (strongly preferred over stdout scraping):
+#   1. Export env vars that the noisy libraries honour (AITER_LOG_LEVEL,
+#      PYTHONWARNINGS, TRANSFORMERS_VERBOSITY, HF_HUB_DISABLE_PROGRESS_BARS).
+#   2. Raise Python ``logging`` levels for the actual logger names emitting
+#      noise (``aiter``, ``megatron`` families, ``torch.distributed``, ...).
+#      Safe to call multiple times via ``reapply_quiet_logger_levels``.
+#   3. Silence ``warnings`` category to match ``PYTHONWARNINGS=ignore``.
+#
+# A handful of sources (unconditional C++ ``std::cout`` in aiter / hipify;
+# Primus ``loguru`` sinks that bypass ``logging``; raw ``print()`` in TE
+# ``rope.py``; ...) cannot be controlled via the above. For those a narrow
+# FD-level line filter is installed on stdout so even native writes get
+# matched. Stderr carries no useful signal in quiet mode so FD 2 is
+# redirected wholesale to ``/dev/null``. If you need stderr back (e.g. to
+# debug a crash) re-run with ``MLPERF_VERBOSE_LOGS=1``.
+#
+# IMPORTANT: This module must be imported BEFORE any other import that may
+# log/print at import time (Megatron, TE, aiter, ...). For the Primus CLI it
+# is imported at the very top of ``primus/cli/main.py``. Because it lives
+# directly under the ``primus`` package, importing it only triggers the light
+# ``primus/__init__.py`` (config/logging utilities) and pulls in none of the
+# heavy native libraries that emit the banners we want to hide.
+# ---------------------------------------------------------------------------
+import logging as _logging
+import os as _os
+
+ENABLED = _os.environ.get("PRIMUS_LOG_SUPPRESSION", "0") == "1"
+VERBOSE_LOGS = _os.environ.get("MLPERF_VERBOSE_LOGS", "0") == "1"
+
+# Logger names that emit the non-MLLOG noise we want to silence. Safe to
+# set even when a given logger is not present in the process.
+QUIET_LOGGER_NAMES = (
+    # aiter: Primus also raises this logger to ERROR inside the trainer,
+    # but we do it earlier (before the aiter JIT build banner fires).
+    "aiter",
+    # Megatron-LM deprecation + rerun-state-machine warnings and the
+    # "Setting RerunStateMachine mode" warning from rerun_state_machine.py.
+    "megatron",
+    "megatron.core",
+    "megatron.core.utils",
+    "megatron.core.rerun_state_machine",
+    "megatron.core.pipeline_parallel",
+    # Primus / Primus-Turbo stdlib loggers (Primus routes most output
+    # through loguru instead, which we drop at the FD level below).
+    "primus",
+    "primus_turbo",
+    "primus_mllog",
+    # TransformerEngine import / RoPE banners that use stdlib logging.
+    "transformer_engine",
+    # torch.distributed elastic / launcher chatter.
+    "torch.distributed",
+    "torch.distributed.run",
+    "torch.distributed.elastic",
+    "torch.distributed.elastic.multiprocessing",
+    "torch.distributed.launcher.api",
+    # HuggingFace Hub / transformers progress + warnings.
+    "transformers",
+    "huggingface_hub",
+    # General misc.
+    "filelock",
+    "urllib3",
+)
+
+
+def reapply_quiet_logger_levels() -> None:
+    """Raise levels on noisy Python loggers. Safe to call multiple times.
+
+    No-op unless suppression is enabled and quiet mode is active.
+    """
+    if not ENABLED or VERBOSE_LOGS:
+        return
+    for _logger_name in QUIET_LOGGER_NAMES:
+        _logging.getLogger(_logger_name).setLevel(_logging.ERROR)
+
+
+def _configure_non_mllog_logs_quiet() -> None:
+    """Silence every non-MLLOG log source we can control via env vars or
+    the standard ``logging`` / ``warnings`` modules. :::MLLOG output is
+    untouched because ``mlperf_logging`` has its own stdout handler.
+    """
+    _os.environ.setdefault("AITER_LOG_LEVEL", "ERROR")
+    _os.environ.setdefault("AITER_LOG_MORE", "0")
+    _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+    _os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    _os.environ.setdefault("PYTHONWARNINGS", "ignore")
+    # Tell Primus-Turbo / GPT-OSS helpers to keep JIT tuning traces off.
+    _os.environ.setdefault("AITER_LOG_TUNED_CONFIG", "0")
+
+    for _logger_name in QUIET_LOGGER_NAMES:
+        _logging.getLogger(_logger_name).setLevel(_logging.ERROR)
+
+    # ``warnings.filterwarnings("ignore")`` covers cases where
+    # ``PYTHONWARNINGS=ignore`` has no effect because Python was started
+    # before we ran (e.g. re-exec'd by a launcher).
+    try:
+        import warnings as _warnings
+
+        _warnings.filterwarnings("ignore")
+    except Exception:
+        pass
+
+
+def _install_fd_level_fallback_filter() -> None:
+    """Last-resort line filter for logs that bypass Python ``logging``
+    entirely: unconditional ``std::cout`` writes in aiter / hipify C++,
+    raw ``print()`` statements in TE / Megatron, and Primus loguru output
+    that landed on stdout.
+
+    Design:
+      * FD 2 (stderr) is redirected straight to ``/dev/null``. In quiet
+        mode the only stderr sources observed in practice are Primus's
+        ``loguru`` sink, ``warnings.warn`` output, the hipify
+        ``"Successfully preprocessed all matching files."`` banner, and
+        torchrun's elastic SIGTERM chatter -- none of which carry a
+        training or MLLOG signal.
+      * FD 1 (stdout) is routed through an ``os.pipe`` + reader thread.
+        The thread runs each line through the suppression regexes and
+        forwards surviving lines to the saved original stdout FD.
+      * ``:::MLLOG`` lines always pass through.
+    """
+    import re
+    import sys
+    import threading
+
+    ansi_re = re.compile(r"\x1b\[[0-9;]*[ -/]*[@-~]")
+
+    suppress_patterns = tuple(
+        re.compile(p)
+        for p in (
+            # aiter JIT module-load / build / baton-wait banners.
+            r"^\[aiter\] ",
+            # TE-RoPE banner from transformer_engine rope.py (raw print).
+            r"^\[TE-RoPE\] ",
+            # "[MLPerf Train] ..." status prints -- all informational, all
+            # emitted N times in distributed mode. Re-run verbose to see them.
+            r"^\[MLPerf Train\] ",
+            # Primus runtime/CLI/patch informational prints that bypass loguru
+            # (e.g. "[Primus:Runtime] ...", "[Primus:Env] ...", "[Primus] ...",
+            # "[Primus CLI] ...", "[PrimusPatch] ..."). Covers the new-arch
+            # runtime/launcher print() calls. Re-run verbose to see them.
+            r"^\[Primus",
+            # Gloo C++ peer-connect banner ("[Gloo] Rank N is connected
+            # to 7 peer ranks. Expected number of connected peer ranks
+            # is : 7") -- native std::cout from libgloo.
+            r"^\[Gloo\] Rank \d+ is connected to ",
+            r"^Expected number of connected peer ranks is\s*:",
+            # hipify preprocessing banner (mostly fires on stderr, listed
+            # here as a safety net for any leak to stdout).
+            r"^Successfully preprocessed all matching files\.",
+            # Primus loguru lines, in case they end up on stdout for any
+            # reason (e.g. ``colorize=False`` configurations). The Primus
+            # format always starts with ``[YYYYMMDD HH:MM:SS]``.
+            r"^\[\d{8} \d{2}:\d{2}:\d{2}\]\[(?:rank|node)-\d+/",
+            # "Setting RerunStateMachine mode RerunMode.DISABLED" fires
+            # from Megatron's rerun_state_machine.py as a root-logger
+            # warning (the default Python logging handler dumps the message
+            # unformatted to stderr, but belt-and-braces: cover stdout too).
+            r"^Setting RerunStateMachine mode ",
+        )
+    )
+
+    def _should_suppress(line: str) -> bool:
+        stripped = ansi_re.sub("", line)
+        if ":::MLLOG" in stripped:
+            return False
+        for pat in suppress_patterns:
+            if pat.search(stripped):
+                return True
+        return False
+
+    def _start_reader(read_fd: int, out_fd: int) -> None:
+        def _run() -> None:
+            buf = b""
+            try:
+                while True:
+                    chunk = _os.read(read_fd, 4096)
+                    if not chunk:
+                        break
+                    buf += chunk
+                    while b"\n" in buf:
+                        raw, buf = buf.split(b"\n", 1)
+                        line = raw.decode("utf-8", errors="replace")
+                        if not _should_suppress(line):
+                            _os.write(out_fd, raw + b"\n")
+            except Exception:
+                # Never let the filter thread bring down the training run.
+                pass
+            finally:
+                if buf:
+                    line = buf.decode("utf-8", errors="replace")
+                    if not _should_suppress(line):
+                        try:
+                            _os.write(out_fd, buf)
+                        except Exception:
+                            pass
+
+        threading.Thread(target=_run, daemon=True).start()
+
+    # Flush any buffered Python stdout/stderr output before we steal the FDs.
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    orig_stdout_fd = _os.dup(1)
+
+    # Drop stderr unconditionally: nothing useful lands on FD 2 in quiet
+    # mode (Primus loguru, warnings, hipify). Users who need stderr back
+    # (e.g. crash debug) can re-run with ``MLPERF_VERBOSE_LOGS=1``.
+    devnull_fd = _os.open(_os.devnull, _os.O_WRONLY)
+    _os.dup2(devnull_fd, 2)
+    _os.close(devnull_fd)
+
+    stdout_r, stdout_w = _os.pipe()
+    _os.dup2(stdout_w, 1)
+    _os.close(stdout_w)
+
+    _start_reader(stdout_r, orig_stdout_fd)
+
+    sys.stdout = _os.fdopen(1, "w", buffering=1, closefd=False)
+    sys.stderr = _os.fdopen(2, "w", buffering=1, closefd=False)
+
+
+_INSTALLED = False
+
+
+def install() -> None:
+    """Install the quiet-mode suppression exactly once per process.
+
+    No-op unless ``PRIMUS_LOG_SUPPRESSION=1``. Calling this a second time is
+    also a no-op.
+    """
+    global _INSTALLED
+    if _INSTALLED:
+        return
+    _INSTALLED = True
+    if not ENABLED:
+        return
+    if VERBOSE_LOGS:
+        return
+    _configure_non_mllog_logs_quiet()
+    _install_fd_level_fallback_filter()
+
+
+# Auto-install on import (gated by PRIMUS_LOG_SUPPRESSION).
+install()
diff --git a/primus/modules/trainer/base_trainer.py b/primus/modules/trainer/base_trainer.py
deleted file mode 100644
index e555156cb..000000000
--- a/primus/modules/trainer/base_trainer.py
+++ /dev/null
@@ -1,24 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-from abc import ABC, abstractmethod
-
-import torch
-from megatron.core.models.gpt import GPTModel
-
-
-class BaseTrainer(ABC):
-    @abstractmethod
-    def get_batch(self, data_iterator):
-        pass
-
-    @abstractmethod
-    def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor):
-        pass
-
-    @abstractmethod
-    def forward_step(self, data_iterator, model: GPTModel):
-        pass
diff --git a/primus/modules/trainer/maxtext/pre_trainer.py b/primus/modules/trainer/maxtext/pre_trainer.py
deleted file mode 100644
index 9c862186d..000000000
--- a/primus/modules/trainer/maxtext/pre_trainer.py
+++ /dev/null
@@ -1,144 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-from typing import Any, Dict
-
-from primus.core.utils import checker
-from primus.modules.base_module import BaseModule
-from primus.modules.module_utils import error_rank_0, log_rank_0, warning_rank_0
-
-
-class MaxTextPretrainTrainer(BaseModule):
-    def __init__(self, *args, **kwargs):
-        extra_args = kwargs.pop("extra_args", None)
-        super().__init__(*args, **kwargs)
-
-        # important: make sure patch maxtext logger first
-        self.patch_maxtext_logger()
-
-        self.primus_cfg = kwargs.pop("primus_config", None)
-
-        if self.primus_cfg is None:
-            raise ValueError("primus_config is required")
-        self.primus_cfg.export_module_config("pre_trainer")
-        self.pre_trainer_cfg_path = self.primus_cfg.module_config_path("pre_trainer")
-        self.override_model_args = self.prepare_model_overrides(extra_args)
-
-    def setup(self):
-        log_rank_0(f"setup MaxText")
-
-    def init(self, *init_args, **kwargs):
-        import functools
-
-        from MaxText import pyconfig
-        from MaxText.train import initialize
-
-        argv = ["MaxText.train", self.pre_trainer_cfg_path]
-        log_rank_0(f"init MaxText with argv {argv}")
-
-        if self.override_model_args:
-            _orig = pyconfig.initialize
-            pyconfig.initialize = functools.partial(_orig, **self.override_model_args)
-            try:
-                self.train_config, self.recorder, self.diagnostic_config = initialize(argv)
-            finally:
-                pyconfig.initialize = _orig
-        else:
-            self.train_config, self.recorder, self.diagnostic_config = initialize(argv)
-
-        self._update_logger_rank()
-
-    def run(self, *args, **kwargs):
-        log_rank_0(f"MaxText Pre-Trainer: begin training...")
-
-        from MaxText.train import run
-
-        run(self.train_config, self.recorder, self.diagnostic_config)
-        log_rank_0("MaxText Pre-Trainer: after training is done")
-
-    def prepare_model_overrides(self, override_args: Dict[str, Any]):
-        """
-        Monkey patch maxtext cli args to override model args dynamically.
-        Supports nested overrides like:
-            {"override_model": {"num_experts": 16, "base_num_decoder_layers": 4}}
-
-        All override keys MUST be under the "model" key.
-        """
-
-        if not override_args:
-            warning_rank_0("MaxText Pre-Trainer: No override_args provided, skip patch.")
-            return {}
-
-        warning_rank_0(f"MaxText Pre-Trainer: Applying override_args: {override_args}")
-
-        # --- Step 1. Flatten any nested dict under 'override_model'
-        flat_overrides = {}
-        for k, v in override_args.items():
-            if k != "override_model":
-                raise ValueError(f"Only the 'override_model' key is supported for overrides, found: {k}")
-            if not isinstance(v, dict):
-                raise ValueError(
-                    f"MaxText Pre-Trainer: The value for 'override_model' must be a dict, got {type(v).__name__}."
-                )
-            for subk, subv in v.items():
-                if isinstance(subv, dict):
-                    raise ValueError(
-                        f"MaxText Pre-Trainer: Invalid override key-value detected: {k}.{subk}-{subv}"
-                    )
-                flat_overrides[subk] = subv
-        return flat_overrides
-
-    def _update_logger_rank(self):
-        """Refresh Primus logger rank/world_size from JAX distributed state.
-
-        The logger is created before ``jax.distributed.initialize()`` runs,
-        so it defaults to rank=0, world_size=1.  After JAX init we have the
-        real values and can patch them in.
-        """
-        import jax
-
-        rank = jax.process_index()
-        world_size = jax.process_count()
-
-        from primus.core.utils.logger import update_rank_info
-        from primus.modules.module_utils import set_logging_rank
-
-        update_rank_info(rank, world_size)
-        set_logging_rank(rank, world_size)
-        log_rank_0(
-            f"JAX distributed ready: rank={rank}, world_size={world_size}, "
-            f"devices={jax.device_count()}, local_devices={jax.local_device_count()}"
-        )
-
-    def patch_maxtext_logger(self):
-        import logging
-
-        from primus.core.utils.logger import _logger as primus_logger
-
-        try:
-            import MaxText.max_logging as maxtext_logging
-
-            if hasattr(maxtext_logging, "log"):
-                maxtext_logging.log = primus_logger.info
-                warning_rank_0("MaxText Pre-Trainer: patch logger successfully.")
-            else:
-                error_rank_0("MaxText Pre-Trainer: logging module does not have a 'log' function.")
-        except ImportError:
-            error_rank_0("MaxText Pre-Trainer: failed to import MaxText Pre-Trainer's logging module.")
-
-        level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
-
-        stderr_sink_level = self.module_config.stderr_sink_level
-        checker.check_true(stderr_sink_level in level_map)
-        logging_level = level_map[stderr_sink_level]
-
-        jax_loggers = [logging.getLogger("jax"), logging.getLogger("jaxlib")]
-        for jax_logger in jax_loggers:
-            jax_logger.setLevel(logging_level)
-
-        warning_rank_0(
-            f"jax.logging_level is deprecated, set logging_level={logging_level} [stderr_sink_level]"
-        )
diff --git a/primus/modules/trainer/megatron/model_provider.py b/primus/modules/trainer/megatron/model_provider.py
deleted file mode 100644
index 38de87202..000000000
--- a/primus/modules/trainer/megatron/model_provider.py
+++ /dev/null
@@ -1,53 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-from types import MethodType
-from typing import Callable, Optional, Union
-
-import torch
-from megatron.core.models.gpt import GPTModel
-from megatron.core.models.mamba import MambaModel
-from megatron.training import get_args
-
-from primus.backends.megatron.core.extensions.logits_processor import fused_softcap
-
-import megatron.legacy.model  # isort: skip
-
-g_final_logit_softcapping: Optional[float] = None
-original_compute_language_model_loss: Optional[MethodType] = None
-
-
-def wrapped_compute_language_model_loss(self, labels: torch.Tensor, logits: torch.Tensor) -> torch.Tensor:
-    global g_final_logit_softcapping
-    assert g_final_logit_softcapping is not None
-
-    logits = logits.float()
-    fused_softcap(logits, g_final_logit_softcapping)
-
-    global original_compute_language_model_loss
-    return original_compute_language_model_loss(labels, logits)
-
-
-def primus_model_provider(
-    model_provider: Callable, pre_process=True, post_process=True, vp_stage: Optional[int] = None
-) -> Union[GPTModel, megatron.legacy.model.GPTModel, MambaModel]:
-    # get model
-    model = model_provider(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
-
-    args = get_args()
-    if args.final_logit_softcapping is not None and args.final_logit_softcapping > 0.0:
-
-        global g_final_logit_softcapping
-        g_final_logit_softcapping = args.final_logit_softcapping
-
-        # save original func
-        global original_compute_language_model_loss
-        original_compute_language_model_loss = model.compute_language_model_loss
-
-        # wrap with logits softcapping
-        model.compute_language_model_loss = MethodType(wrapped_compute_language_model_loss, model)
-
-    return model
diff --git a/primus/modules/trainer/megatron/pre_trainer.py b/primus/modules/trainer/megatron/pre_trainer.py
deleted file mode 100644
index d34788b04..000000000
--- a/primus/modules/trainer/megatron/pre_trainer.py
+++ /dev/null
@@ -1,301 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-import collections
-from functools import partial
-
-import torch
-from megatron.core import mpu
-from megatron.core.models.gpt import GPTModel
-from megatron.core.rerun_state_machine import get_rerun_state_machine
-from megatron.core.utils import StragglerDetector, get_attr_wrapped_model
-from megatron.training import get_args, get_timers
-from megatron.training.utils import (
-    get_batch_on_this_cp_rank,
-    get_batch_on_this_tp_rank,
-    is_first_or_last_pipeline_stage,
-)
-
-stimer = StragglerDetector()
-
-from .trainer import MegatronTrainer
-
-mb_batch = None
-
-
-def get_batch_func(data_iterator, vp_stage=None):
-    # TODO: this is pretty hacky, find a better way
-    if not is_first_or_last_pipeline_stage(vp_stage):
-        return None, None, None, None, None
-
-    # assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}"
-    # get batches based on the TP rank you are on
-    batch = get_batch_on_this_tp_rank(data_iterator)
-
-    # slice batch along sequence dimension for context parallelism
-    args = get_args()
-
-    if args.patch_zero_bubble:
-        from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_vars import (
-            get_seq_split_idx,
-        )
-
-        global mb_batch
-        # "or 0" to support original 1f1b and interleaved-1f1b in schedules.py
-        seq_split_idx = get_seq_split_idx() or 0
-        if seq_split_idx == 0:
-            # get batches based on the TP rank you are on
-            mb_batch = get_batch_on_this_tp_rank(data_iterator)
-            assert (
-                mb_batch["attention_mask"] is None
-            ), "attention_mask should be None, please enable --no-create-attention-mask-in-dataloader"
-        batch = {}
-        for k in mb_batch.keys():
-            v = mb_batch[k]
-            if v is None:
-                batch[k] = v
-                continue
-
-            assert v.shape[1] % get_args().num_seq_splits == 0, f"{k} size {v.shape}"
-            start_idx = seq_split_idx * v.shape[1] // get_args().num_seq_splits
-            end_idx = (seq_split_idx + 1) * v.shape[1] // get_args().num_seq_splits
-            if len(v.shape) > 2:
-                batch[k] = v[:, start_idx:end_idx, :].contiguous()
-            else:
-                batch[k] = v[:, start_idx:end_idx].contiguous()
-
-    if args.context_parallel_size > 1 and args.enable_primus_turbo and args.use_turbo_attention:
-        try:
-            from primus.backends.megatron.core.utils import (
-                produce_attention_sharder,
-                shard_batch_on_this_cp_rank,
-            )
-        except:
-            raise ImportError("Module 'primus_turbo' may not installed. Please install it")
-        sharder = produce_attention_sharder(args.cp_comm_type)
-        batch = shard_batch_on_this_cp_rank(sharder, batch)
-    else:
-        batch = get_batch_on_this_cp_rank(batch)
-
-    return batch.values()
-
-
-class DataLoaderStore:
-    cache = collections.deque()
-
-    @classmethod
-    def push(cls, data_iterator, h2d_stream=False, vp_stage=None):
-        timers = get_timers()
-        # Get the batch.
-        timers("batch-generator", log_level=2).start()
-        global stimer
-
-        with stimer(bdata=True):
-            if h2d_stream:
-                from primus.backends.megatron.core.pipeline_parallel.zerobubble.offload import (
-                    get_offload_h2d_stream,
-                )
-
-                load_event = torch.cuda.Event()
-                original_stream = torch.cuda.current_stream()
-                with torch.cuda.stream(get_offload_h2d_stream()):
-                    data = get_batch_func(data_iterator, vp_stage)
-                    for x in data:
-                        if x is not None:
-                            x.record_stream(original_stream)
-                    load_event.record()
-                    cls.cache.append((data, load_event))
-            else:
-                cls.cache.append((get_batch_func(data_iterator, vp_stage), None))
-        timers("batch-generator").stop()
-
-    @classmethod
-    def pop(cls):
-        data, load_event = cls.cache.popleft()
-        if load_event:
-            load_event.wait()
-        return data
-
-
-class MegatronPretrainTrainer(MegatronTrainer):
-    def __init__(self, *args, **kwargs):
-        kwargs["module_name"] = "pre_trainer"
-
-        # Explicitly reject unknown extra_args
-        extra_args = kwargs.pop("extra_args", None)
-        if extra_args:
-            raise ValueError(
-                f"[MegatronPretrainTrainer] Unexpected extra_args detected: {extra_args}. "
-                f"Megatron backend does not support unregistered config keys."
-            )
-
-        super().__init__(*args, **kwargs)
-
-    def get_batch(self, data_iterator, vp_stage=None):
-        """Generate a batch."""
-        return get_batch_func(data_iterator, vp_stage)
-
-    def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor):
-        """Loss function.
-
-        Args:
-            loss_mask (torch.Tensor): Used to mask out some portions of the loss
-            output_tensor (torch.Tensor): The tensor with the losses
-
-        Returns:
-            the loss scalar for this micro-batch
-            the number of non-padded tokens in this microbatch
-            a dict containing reporting metrics on the loss and number of tokens across
-                the data parallel ranks
-        """
-        args = get_args()
-
-        losses = output_tensor.float()
-        loss_mask = loss_mask.view(-1).float()
-        total_tokens = loss_mask.sum()
-        loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), total_tokens.view(1)])
-
-        if args.context_parallel_size > 1:
-            torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
-
-        # Check individual rank losses are not NaN prior to DP all-reduce.
-        rerun_state_machine = get_rerun_state_machine()
-        if args.check_for_nan_in_loss_and_grad:
-            rerun_state_machine.validate_result(
-                result=loss[0],
-                rejection_func=torch.isnan,
-                message="found NaN in local forward loss calculation",
-                tolerance=0.0,  # forward pass calculations are determinisic
-                fatal=True,
-            )
-            rerun_state_machine.validate_result(
-                result=loss[0],
-                rejection_func=torch.isinf,
-                message="found Inf in local forward loss calculation",
-                tolerance=0.0,  # forward pass calculations are determinisic
-                fatal=True,
-            )
-        # Check for spiky loss
-        if args.check_for_spiky_loss:
-            rerun_state_machine.validate_result(
-                result=loss[0],
-                rejection_func=partial(
-                    rerun_state_machine.is_unexpectedly_large,
-                    threshold=SPIKY_LOSS_FACTOR,
-                    context="loss",
-                ),
-                message="Spiky loss",
-                tolerance=0.0,  # forward pass calculations are determinisic
-                fatal=False,
-            )
-        # Reduce loss for logging.
-        reporting_loss = loss.clone().detach()
-        torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group())
-
-        # loss[0] is a view of loss, so it has ._base not None, which triggers assert error
-        # in core/pipeline_parallel/schedule.py::deallocate_output_tensor, calling .clone()
-        # on loss[0] fixes this
-        local_num_tokens = loss[1].clone().detach().to(torch.int)
-        return (
-            loss[0].clone(),
-            local_num_tokens,
-            {"lm loss": (reporting_loss[0], reporting_loss[1])},
-        )
-
-    def forward_step(self, data_iterator, model: GPTModel, return_schedule_plan=False):
-        """Forward training step.
-
-        Args:
-            data_iterator : Input data iterator
-            model (GPTModel): The GPT Model
-        """
-        args = get_args()
-        timers = get_timers()
-
-        # Get the batch.
-        if not args.patch_zero_bubble:
-            timers("batch-generator", log_level=2).start()
-            global stimer
-            with stimer(bdata=True):
-                vp_stage = get_attr_wrapped_model(model, "vp_stage")
-                tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(
-                    data_iterator, vp_stage
-                )
-            timers("batch-generator").stop()
-        else:
-            from collections.abc import Iterable
-
-            vp_stage = get_attr_wrapped_model(model, "vp_stage")
-            if (
-                not isinstance(data_iterator, Iterable) and not data_iterator is None
-            ):  # isinstance(data_iterator, DataLoaderStore):
-                tokens, labels, loss_mask, attention_mask, position_ids = data_iterator.pop()
-            else:
-                DataLoaderStore.push(data_iterator, h2d_stream=False, vp_stage=vp_stage)
-                tokens, labels, loss_mask, attention_mask, position_ids = DataLoaderStore.pop()
-
-        with stimer:
-            if return_schedule_plan:
-                assert (
-                    args.overlap_moe_expert_parallel_comm
-                ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan"
-
-                # Schedule plan building is only supported for GPT models
-                # Check if this is a Mamba model
-                unwrapped_model = model
-                while hasattr(unwrapped_model, "module"):
-                    unwrapped_model = unwrapped_model.module
-                model_class_name = unwrapped_model.__class__.__name__
-
-                if "Mamba" in model_class_name:
-                    raise NotImplementedError(
-                        "Schedule plan building is not supported for Mamba models. "
-                        "Please disable overlap_moe_expert_parallel_comm for Mamba."
-                    )
-
-                if args.patch_moe_overlap:
-                    assert (
-                        not args.delay_wgrad_compute
-                    ), "Primus MoE overlap handles wgrad separately from the original Megatron implementation"
-                    from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_utils import (
-                        WeightGradStore,
-                    )
-
-                    WeightGradStore.enable_split_bw()
-                    assert (
-                        WeightGradStore.split_bw()
-                    ), "WeightGradStore.split_bw is not supported, please make sure overlap_grad_reduce is disabled and gradient_accumulation_fusion is enabled"
-                    from primus.backends.megatron.core.models.common.model_chunk_schedule_plan import (
-                        TransformerModelChunkSchedulePlan,
-                    )
-
-                    schedule_plan = TransformerModelChunkSchedulePlan(
-                        model, tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
-                    )
-                else:
-                    schedule_plan = model.build_schedule_plan(
-                        tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
-                    )
-                return schedule_plan, partial(self.loss_func, loss_mask)
-            else:
-                # Check if model supports loss_mask parameter
-                # MambaModel doesn't accept loss_mask, but GPTModel does
-                # Unwrap the model to get the actual model class
-                unwrapped_model = model
-                while hasattr(unwrapped_model, "module"):
-                    unwrapped_model = unwrapped_model.module
-                model_class_name = unwrapped_model.__class__.__name__
-
-                if "Mamba" in model_class_name:
-                    # MambaModel doesn't accept loss_mask parameter
-                    output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
-                else:
-                    # GPTModel and other models accept loss_mask parameter
-                    output_tensor = model(
-                        tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
-                    )
-
-        return output_tensor, partial(self.loss_func, loss_mask)
diff --git a/primus/modules/trainer/megatron/sft_trainer.py b/primus/modules/trainer/megatron/sft_trainer.py
deleted file mode 100644
index 6aec95d15..000000000
--- a/primus/modules/trainer/megatron/sft_trainer.py
+++ /dev/null
@@ -1,22 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-from .trainer import MegatronTrainer
-
-
-class MegatronSFTTrainer(MegatronTrainer):
-    def __init__(self, *args, **kwargs):
-        kwargs["module_name"] = "sft_trainer"
-        super().__init__(*args, **kwargs)
-
-    def get_batch(self, data_iterator):
-        raise NotImplementedError
-
-    def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor):
-        raise NotImplementedError
-
-    def forward_step(self, data_iterator, model: GPTModel):
-        raise NotImplementedError
diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py
deleted file mode 100644
index 429facb0c..000000000
--- a/primus/modules/trainer/megatron/trainer.py
+++ /dev/null
@@ -1,2537 +0,0 @@
-###############################################################################
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-import argparse
-import dataclasses
-import functools
-import gc
-import importlib.util
-import inspect
-import json
-import os
-import statistics
-import sys
-import time
-
-import megatron
-import torch
-import torch.distributed as dist
-from megatron.core import mpu, parallel_state, tensor_parallel
-from megatron.core.distributed import DistributedDataParallel as DDP
-from megatron.core.distributed import finalize_model_grads
-from megatron.core.distributed.distributed_data_parallel_config import (
-    DistributedDataParallelConfig,
-)
-from megatron.core.distributed.torch_fully_sharded_data_parallel import (
-    TorchFullyShardedDataParallel as torch_FSDP,
-)
-from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config
-from megatron.training.checkpointing import (
-    checkpoint_exists,
-    load_checkpoint,
-    save_checkpoint,
-)
-from megatron.training.training import save_checkpoint_and_time
-
-from primus.backends.megatron.core.optimizer.moun import get_megatron_muon_optimizer
-from primus.backends.megatron.core.optimizer.moun_optimizer_config import (
-    MounOptimizerConfig,
-)
-from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss
-from primus.core.utils.import_utils import get_custom_fsdp, get_model_provider
-
-try:
-    pass
-
-    HAVE_FSDP2 = True
-except ImportError:
-    HAVE_FSDP2 = False
-from megatron.core.datasets.blended_megatron_dataset_builder import (
-    BlendedMegatronDatasetBuilder,
-)
-from megatron.core.datasets.gpt_dataset import (
-    GPTDataset,
-    GPTDatasetConfig,
-    MockGPTDataset,
-)
-from megatron.core.enums import ModelType
-from megatron.core.num_microbatches_calculator import (
-    get_current_global_batch_size,
-    get_current_running_global_batch_size,
-    get_num_microbatches,
-    update_num_microbatches,
-)
-from megatron.core.optimizer import get_megatron_optimizer, get_mup_config_overrides
-from megatron.core.rerun_state_machine import (
-    RerunDiagnostic,
-    RerunErrorInjector,
-    RerunMode,
-    get_rerun_state_machine,
-    initialize_rerun_state_machine,
-)
-from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config
-from megatron.training import (
-    ft_integration,
-    get_args,
-    get_tensorboard_writer,
-    get_timers,
-    global_vars,
-    one_logger_utils,
-)
-from megatron.training.arguments import validate_args
-from megatron.training.async_utils import (
-    init_persistent_async_worker,
-    maybe_finalize_async_save,
-)
-from megatron.training.checkpointing import (
-    checkpoint_exists,
-    load_args_from_checkpoint,
-    load_checkpoint,
-    save_checkpoint,
-)
-from megatron.training.global_vars import (
-    get_args,
-    get_one_logger,
-    get_tensorboard_writer,
-    get_timers,
-    get_tokenizer,
-    get_wandb_writer,
-    set_global_variables,
-)
-from megatron.training.initialize import (
-    _compile_dependencies,
-    _init_autoresume,
-    _initialize_distributed,
-    _initialize_tp_communicators,
-    _set_random_seed,
-    set_jit_fusion_options,
-    setup_logging,
-    write_args_to_tensorboard,
-)
-from megatron.training.theoretical_memory_usage import report_theoretical_memory
-from megatron.training.training import (
-    build_train_valid_test_data_iterators,
-    checkpoint_and_decide_exit,
-    disable_forward_pre_hook,
-    dummy_train_step,
-    enable_forward_pre_hook,
-    evaluate_and_print_results,
-    get_megatron_optimizer_config,
-    get_model,
-    get_optimizer_param_scheduler,
-    num_floating_point_operations,
-    post_training_step_callbacks,
-    preprocess_common_state_dict,
-    print_datetime,
-    should_disable_forward_pre_hook,
-)
-from megatron.training.utils import (
-    append_to_progress_log,
-    calc_params_l2_norm,
-    get_blend_and_blend_per_split,
-    is_first_or_last_pipeline_stage,
-    logical_and_across_model_parallel_group,
-    reduce_max_stat_across_model_parallel_group,
-    report_memory,
-    unwrap_model,
-    update_use_dist_ckpt,
-)
-from megatron.training.yaml_arguments import validate_yaml
-
-from primus.backends.megatron.argument_builder import _load_megatron_defaults
-from primus.backends.megatron.core.transformer.moe.moe_utils import track_moe_metrics
-from primus.backends.megatron.training.global_vars import (
-    get_mlflow_writer,
-    get_train_start_time,
-    set_primus_global_variables,
-    set_train_start_time,
-)
-from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts
-from primus.backends.megatron.training.tokenizer.tokenizer import build_tokenizer
-from primus.core.utils import checker, file_utils
-from primus.core.utils.rocm_mem_info import get_rocm_smi_gpu_util, get_rocm_smi_mem_info
-from primus.core.utils.yaml_utils import nested_namespace_to_dict
-from primus.modules.base_module import BaseModule
-from primus.modules.module_utils import (
-    debug_rank_0,
-    log_kv_rank_0,
-    log_rank_0,
-    log_rank_last,
-    warning_rank_0,
-)
-from primus.modules.trainer.base_trainer import BaseTrainer
-from primus.modules.trainer.megatron.model_provider import primus_model_provider
-
-from .utils import (
-    is_v_schedule_enabled,
-    schedule_wrapper,
-    set_wandb_writer_patch,
-    validate_args_on_rocm,
-)
-
-# The earliest we can measure the start time.
-set_train_start_time()
-
-
-def _normalize_data_path_arg(path_value):
-    """Normalize data path args to list form when paths are passed as strings."""
-    if path_value is None:
-        return None
-    if isinstance(path_value, str):
-        return path_value.split()
-    if isinstance(path_value, (list, tuple)):
-        return list(path_value)
-    return path_value
-
-
-class MegatronTrainer(BaseTrainer, BaseModule):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.is_v_schedule = is_v_schedule_enabled(self.module_config)
-
-        self.app_metrics = {}
-
-        # Cache for GPU utilization sampling to avoid rocm-smi overhead on every rank
-        # See: https://github.com/AMD-AGI/Primus/pull/439 (Copilot review)
-        self._last_rocm_gpu_util = None
-        self._last_rocm_gpu_util_time = 0.0
-        self._rocm_gpu_util_sample_interval = 5.0  # seconds
-
-        # disable all logging handlers
-        import logging
-
-        for handler in logging.root.handlers[:]:
-            logging.root.removeHandler(handler)
-
-        if "aiter" in logging.root.manager.loggerDict:
-            logging.getLogger("aiter").setLevel(logging.ERROR)
-
-    def init(self, *init_args, **kwargs):
-        allowed_keys = {
-            "extra_args_provider",
-            "args_defaults",
-            "ignore_unknown_args",
-            "allow_no_cuda",
-            "skip_mpu_initialization",
-            "skip_setup",
-        }
-
-        invalid_keys = set(kwargs.keys()) - allowed_keys
-        if invalid_keys:
-            raise TypeError(f"Invalid keyword arguments for MegatronTrainer: {invalid_keys}")
-
-        log_rank_0(f"-run update_primus_config...")
-        self.update_primus_config(
-            args=self.module_config,
-            exp_root_path=self.exp_root_path,
-            exp_meta_info=self.exp_meta_info,
-        )
-
-        # Apply patches during initialization
-        # These patches need to be applied before model setup
-        from types import SimpleNamespace
-
-        from primus.core.patches import run_patches
-
-        # Construct a temporary module_config for patch context
-        # This allows patches to access Megatron args via get_args(ctx)
-        temp_module_config = SimpleNamespace()
-        temp_module_config.params = self.module_config
-
-        run_patches(
-            backend="megatron",
-            phase="before_train",
-            backend_version="0.15.0rc8",
-            extra={
-                "module_config": temp_module_config,
-            },
-        )
-
-        # Initalize and get arguments, timers, and Tensorboard writer.
-        log_rank_0(f"-run initialize_megatron...")
-        self.initialize_megatron(
-            extra_args_provider=kwargs.get("extra_args_provider", None),
-            args_defaults=kwargs.get("args_defaults", {}),
-            ignore_unknown_args=kwargs.get("ignore_unknown_args", False),
-            allow_no_cuda=kwargs.get("allow_no_cuda", False),
-            skip_mpu_initialization=kwargs.get("skip_mpu_initialization", False),
-        )
-
-        args = get_args()
-        # There are some extra limitation on ROCm need extra validate.
-        validate_args_on_rocm(args)
-
-        # Enable manually split layers in (interleaved) 1f1b pipeline
-        # parallelism by monkey patching
-        if args.decoder_pipeline_manual_split_list is not None:
-            log_rank_0(
-                f"-decoder_pipeline_manual_split_list has been deprecated, please use pipeline_model_parallel_layout instead"
-            )
-            from .utils import set_manual_pipeline_split_patch, validate_manual_split
-
-            log_rank_0(f"-monkey patch to enable manual pipeline split...")
-            if validate_manual_split(args):
-                set_manual_pipeline_split_patch(args)
-
-        if args.recompute_layer_ids is not None:
-            from .utils import validate_specified_recompute_layers
-
-            validate_specified_recompute_layers(args)
-
-        if args.log_progress:
-            append_to_progress_log("Starting job")
-
-        self.log_avg_reset_interval = args.log_avg_reset_interval
-        self.log_avg_skip_iterations = args.log_avg_skip_iterations
-        self.recent_tflop_throughputs = []
-        self.recent_iteration_times = []
-        self.recent_token_throughputs = []
-
-        # Initialize fault tolerance
-        # NOTE: ft_integration functions other than `setup` are no-op if the FT is not initialized
-        if args.enable_ft_package:
-            ft_integration.setup(args)
-            ft_integration.maybe_setup_simulated_fault()
-
-        # Set pytorch JIT layer fusion options and warmup JIT functions.
-        set_jit_fusion_options()
-
-        # Adjust the startup time so it reflects the largest value.
-        # This will be closer to what scheduler will see (outside of
-        # image ... launches.
-        start_time_tensor = torch.tensor([get_train_start_time()], dtype=torch.double, device="cuda")
-        torch.distributed.all_reduce(start_time_tensor, op=torch.distributed.ReduceOp.MIN)
-        set_train_start_time(start_time_tensor.item())
-
-        self.app_metrics["app_start_time"] = round(get_train_start_time() * 1000.0)
-        self.app_metrics["app_model_init_start_time"] = round(get_train_start_time() * 1000.0)
-
-        log_rank_0(
-            "time to initialize megatron (seconds): {:.3f}".format(time.time() - get_train_start_time())
-        )
-        print_datetime("after megatron is initialized")
-        self.app_metrics["app_model_init_finish_time"] = one_logger_utils.get_timestamp_in_ms()
-
-        # Track E2E metrics on pretrain start
-        one_logger_utils.on_pretrain_start()
-
-        # Context used for persisting some state between checkpoint saves.
-        if args.non_persistent_ckpt_type == "local":
-            try:
-                from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import (
-                    LocalCheckpointManager,
-                )
-                from nvidia_resiliency_ext.checkpointing.local.replication.strategies import (
-                    CliqueReplicationStrategy,
-                )
-            except ModuleNotFoundError:
-                raise RuntimeError(
-                    "The 'nvidia_resiliency_ext' module is required for local "
-                    "checkpointing but was not found. Please ensure it is installed."
-                )
-
-            if args.replication:
-                repl_strategy = CliqueReplicationStrategy.from_replication_params(
-                    args.replication_jump, args.replication_factor
-                )
-            else:
-                repl_strategy = None
-
-            self.checkpointing_context = {
-                "local_checkpoint_manager": LocalCheckpointManager(
-                    args.non_persistent_local_ckpt_dir, repl_strategy=repl_strategy
-                )
-            }
-        else:
-            self.checkpointing_context = {}
-
-        if not kwargs.get("skip_setup", False):
-            self.setup()
-
-    def update_primus_config(
-        self,
-        args,
-        exp_meta_info,
-        exp_root_path,
-    ):
-        # rank/world_size
-        args.rank = self.module_rank
-        args.world_size = self.module_world_size
-        args.local_rank = self.module_local_rank
-        log_kv_rank_0(f"-rank", f"{args.rank}")
-        log_kv_rank_0(f"-local_rank", f"{args.local_rank}")
-        log_kv_rank_0(f"-world_size", f"{args.world_size}")
-
-        # cuda
-        if not args.use_torch_fsdp2 and not args.use_custom_fsdp:
-            os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
-        else:
-            os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "8"
-
-        # profile
-        if args.profile:
-            args.disable_tensorboard = False
-
-        # checkpoint
-        ckpt_path = os.path.abspath(os.path.join(exp_root_path, "checkpoints"))
-        if args.save is not None:
-            warning_rank_0(f" args.save is deprecated, the checkpoint path is: {ckpt_path}")
-        args.save = ckpt_path
-        log_kv_rank_0(f"-save", f"{args.save}")
-
-        # auto_continue_train
-        # Note that if args.auto_continue_train is enabled, check if there are existing save checkpoints
-        # for the current training experiment. If checkpoints are found, update the training
-        # configuration by disable finetuning (args.finetune), loading the optimizer state
-        # (args.no_load_optim), and loading the random number generator state (args.no_load_rng).
-        log_kv_rank_0(f"-auto_continue_train", f"{args.auto_continue_train}")
-        if args.auto_continue_train:
-            latest_file = f"{args.save}/latest_checkpointed_iteration.txt"
-            ckpt_exist = file_utils.is_file(latest_file)
-            if ckpt_exist:
-                with open(latest_file, "r") as file:
-                    iter_str = file.read().strip()
-                log_rank_0(f"-find '{latest_file}', latest iteration is {iter_str}.")
-                if args.load != args.save:
-                    warning_rank_0(
-                        f"-set args.load={args.save}, path '{args.load}' is deprecated. [auto_continue_train]"
-                    )
-                    args.load = args.save
-                if args.finetune:
-                    args.finetune = False
-                    warning_rank_0(f"-set args.finetune=False [auto_continue_train]")
-                if args.no_load_optim:
-                    args.no_load_optim = False
-                    warning_rank_0(f"-set args.no_load_optim=False [auto_continue_train]")
-                if args.no_load_rng:
-                    args.no_load_rng = False
-                    warning_rank_0(f"-set args.no_load_rng=False [auto_continue_train]")
-                if not args.use_checkpoint_args:
-                    args.use_checkpoint_args = True
-                    warning_rank_0(f"-set args.use_checkpoint_args=True [auto_continue_train]")
-            else:
-                log_rank_0(f"-{latest_file} does not exist, skip auto_continue_train.")
-
-        # Auto-enable profiling and tensorboard when traces are needed: for MLflow upload
-        # (only if MLflow is enabled) or for local TraceLens report generation.
-        # Without this, generate_tracelens_report=True with profile=False would produce no traces.
-        needs_profiling = (
-            (
-                getattr(args, "mlflow_upload_traces", False)
-                or getattr(args, "mlflow_upload_tracelens_report", False)
-            )
-            and not args.disable_mlflow
-        ) or getattr(args, "generate_tracelens_report", False)
-        if needs_profiling:
-            if not getattr(args, "profile", False):
-                args.profile = True
-                debug_rank_0("Auto-enabled profile=True for trace/tracelens (upload or local generation)")
-            if not getattr(args, "use_pytorch_profiler", False):
-                args.use_pytorch_profiler = True
-                debug_rank_0(
-                    "Auto-enabled use_pytorch_profiler=True for trace/tracelens (upload or local generation)"
-                )
-            if getattr(args, "disable_tensorboard", True):
-                args.disable_tensorboard = False
-                debug_rank_0("Auto-enabled tensorboard (disable_tensorboard=False) for profiler trace output")
-
-        # tensorboard
-        if not args.disable_tensorboard:
-            tb_path = os.path.abspath(os.path.join(exp_root_path, "tensorboard"))
-            if args.tensorboard_dir is not None:
-                warning_rank_0(f"args.tensorboard_dir is deprecated, the tensorboard path is: {tb_path}")
-            args.tensorboard_dir = tb_path
-        else:
-            args.tensorboard_dir = None
-        log_kv_rank_0(f"-disable_tensorboard", f"{args.disable_tensorboard}")
-        log_kv_rank_0(f"  -tensorboard_dir", f"{args.tensorboard_dir}")
-
-        # wandb
-        if not args.disable_wandb:
-            wandb_path = exp_root_path
-            if args.wandb_save_dir is not None:
-                warning_rank_0(f"args.wandb_save_dir is deprecated, the wandb path is: {wandb_path}/wandb")
-            if not hasattr(args, "wandb_project") or args.wandb_project is None:
-                args.wandb_project = f"{exp_meta_info['work_group']}_{exp_meta_info['user_name']}"
-                debug_rank_0(f"  -create new wandb project name: {args.wandb_project}")
-            if not hasattr(args, "wandb_exp_name") or args.wandb_exp_name is None:
-                args.wandb_exp_name = exp_meta_info["exp_name"]
-                debug_rank_0(f"  -create new exp name: {args.wandb_exp_name}")
-            args.wandb_save_dir = wandb_path
-        elif args.wandb_project is not None:
-            args.wandb_project = None
-            debug_rank_0(f"args.wandb_project is disabled, as args.disable_wandb=True.")
-        log_kv_rank_0(f"-disable_wandb", f"{args.disable_wandb}")
-        if not args.disable_wandb and "WANDB_API_KEY" not in os.environ:
-            warning_rank_0(
-                "The environment variable WANDB_API_KEY is not set. "
-                "Please set it before proceeding or enable 'disable_wandb' in yaml config"
-            )
-        log_kv_rank_0(f"  -wandb_project", f"{args.wandb_project}")
-        log_kv_rank_0(f"  -wandb_exp_name", f"{args.wandb_exp_name}")
-        log_kv_rank_0(f"  -wandb_save_dir", f"{args.wandb_save_dir}")
-        log_kv_rank_0(f"  -wandb_entity", f"{args.wandb_entity}")
-
-        # mlflow
-        log_kv_rank_0(f"-disable_mlflow", f"{args.disable_mlflow}")
-        if not args.disable_mlflow:
-            if not hasattr(args, "mlflow_run_name") or args.mlflow_run_name is None:
-                args.mlflow_run_name = f"{exp_meta_info['work_group']}_{exp_meta_info['user_name']}"
-                debug_rank_0(f"  -create new mlflow run name: {args.mlflow_run_name}")
-        elif args.mlflow_run_name is not None:
-            args.mlflow_run_name = None
-            args.mlflow_experiment_name = None
-            debug_rank_0(f"args.mlflow_run_name is disabled, as args.disable_mlflow=True.")
-        if not args.disable_mlflow and "DATABRICKS_HOST" not in os.environ:
-            warning_rank_0(
-                "The environment variable DATABRICKS_HOST is not set. "
-                "Please set it before proceeding or enable 'disable_mlflow' in yaml config"
-            )
-        log_kv_rank_0(f"  -mlflow_run_name", f"{args.mlflow_run_name}")
-        log_kv_rank_0(f"  -mlflow_experiment_name", f"{args.mlflow_experiment_name}")
-
-        # sink_level: logging_level
-        level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
-        checker.check_true(args.stderr_sink_level in level_map)
-        logging_level = level_map[args.stderr_sink_level]
-        if args.logging_level is not None:
-            warning_rank_0(
-                f"-args.logging_level is deprecated, set args.logging_level={logging_level} [stderr_sink_level]"
-            )
-        args.logging_level = logging_level
-
-        # update data path
-        # "data1 data2 data3" -> ['data1', 'data2', 'data3']
-        if args.data_path is not None:
-            args.data_path = _normalize_data_path_arg(args.data_path)
-            log_rank_0(f"-data_path: {args.data_path}")
-
-        if args.train_data_path is not None:
-            args.train_data_path = _normalize_data_path_arg(args.train_data_path)
-            log_rank_0(f"-train_data_path: {args.train_data_path}")
-        if args.valid_data_path is not None:
-            args.valid_data_path = _normalize_data_path_arg(args.valid_data_path)
-            log_rank_0(f"-valid_data_path: {args.valid_data_path}")
-        if args.test_data_path is not None:
-            args.test_data_path = _normalize_data_path_arg(args.test_data_path)
-            log_rank_0(f"-test_data_path: {args.test_data_path}")
-
-        # update sp
-        if args.tensor_model_parallel_size == 1:
-            args.sequence_parallel = False
-
-        if args.iterations_to_skip is None:
-            args.iterations_to_skip = []
-
-        # support moe_freq_type - ensure moe_layer_freq has a default value
-        if not hasattr(args, "moe_layer_freq"):
-            args.moe_layer_freq = 1
-        elif isinstance(args.moe_layer_freq, str):
-            try:
-                args.moe_layer_freq = eval(args.moe_layer_freq)
-            except Exception:
-                raise ValueError(f"Invalid moe_layer_freq format: {args.moe_layer_freq}")
-
-        if args.mock_data:
-            args.data_path = None
-            args.train_data_path = None
-            args.valid_data_path = None
-            args.test_data_path = None
-
-        # Determine model type (gpt or mamba)
-        model_type = getattr(args, "model_type", "gpt")
-        log_rank_0(f"-detected model_type: {model_type}")
-
-        # Ensure required attributes have safe defaults if missing from config
-        if not hasattr(args, "final_logit_softcapping"):
-            args.final_logit_softcapping = None
-        if not hasattr(args, "router_logit_softcapping"):
-            args.router_logit_softcapping = None
-
-        # Only pass model_type parameter when it's "mamba" to maintain backward compatibility
-        # with main branch behavior for "gpt" (default) case
-        if args.final_logit_softcapping is not None and args.final_logit_softcapping > 0.0:
-            log_rank_0(f"-enable final_logit_softcapping: {args.final_logit_softcapping}")
-            if model_type == "mamba":
-                self.model_provider = functools.partial(
-                    primus_model_provider, get_model_provider(model_type=model_type)
-                )
-            else:
-                self.model_provider = functools.partial(primus_model_provider, get_model_provider())
-        else:
-            if model_type == "mamba":
-                log_rank_0(f"-getting model provider for model_type={model_type}")
-                model_provider = get_model_provider(model_type=model_type)
-                log_rank_0(f"-model_provider: {model_provider}")
-                self.model_provider = model_provider
-            else:
-                # For "gpt" (default), call without arguments to match main branch behavior
-                self.model_provider = get_model_provider()
-
-        if args.router_logit_softcapping is not None and args.router_logit_softcapping > 0.0:
-            log_rank_0(f"-enable router_logit_softcapping: {args.router_logit_softcapping}")
-
-    def vocab_size_with_padding(self, orig_vocab_size, args):
-        """Pad vocab size so it is divisible by model parallel size and
-        still having GPU friendly size."""
-
-        after = orig_vocab_size
-        multiple = args.make_vocab_size_divisible_by * args.tensor_model_parallel_size
-        while (after % multiple) != 0:
-            after += 1
-        debug_rank_0(
-            " -padded vocab (size: {}) with {} dummy tokens "
-            "(new size: {})".format(orig_vocab_size, after - orig_vocab_size, after)
-        )
-        return after
-
-    def setup(self):
-        args = get_args()
-        timers = get_timers()
-        # Model, optimizer, and learning rate.
-        timers("model-and-optimizer-setup", log_level=0).start(barrier=True)
-        self.app_metrics["app_build_optimizer_start_time"] = one_logger_utils.get_timestamp_in_ms()
-        log_rank_0(f"-setup_model_and_optimizer...")
-        self.model, self.optimizer, self.opt_param_scheduler = self.setup_model_and_optimizer(
-            self.model_provider,
-            ModelType.encoder_or_decoder,
-            checkpointing_context=self.checkpointing_context,
-        )
-
-        timers("model-and-optimizer-setup").stop()
-        print_datetime("after model, optimizer, and learning rate " "scheduler are built")
-        self.app_metrics["app_build_optimizer_finish_time"] = one_logger_utils.get_timestamp_in_ms()
-        self.config = get_model_config(self.model[0])
-
-        # Data stuff.
-        self.app_metrics["app_build_dataiters_start_time"] = one_logger_utils.get_timestamp_in_ms()
-        timers("train/valid/test-data-iterators-setup", log_level=0).start(barrier=True)
-
-        def train_valid_test_datasets_provider_func(train_val_test_num_samples, vp_stage=None):
-            return self.train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=vp_stage)
-
-        train_valid_test_datasets_provider_func.is_distributed = True
-
-        if args.virtual_pipeline_model_parallel_size is not None:
-            self.train_data_iterator = []
-            self.valid_data_iterator = []
-            self.test_data_iterator = []
-            for vp_stage in range(len(self.model)):
-                dataset_provider_parameters = inspect.signature(
-                    train_valid_test_datasets_provider_func
-                ).parameters
-                assert (
-                    "vp_stage" in dataset_provider_parameters
-                ), "vp_stage must be a kwarg in train_valid_test_dataset_provider when using virtual pipeline parallelism"
-                vp_stage_train_valid_test_dataset_provider = functools.partial(
-                    train_valid_test_datasets_provider_func, vp_stage=vp_stage
-                )
-                if getattr(train_valid_test_datasets_provider_func, "is_distributed", False):
-                    vp_stage_train_valid_test_dataset_provider.is_distributed = True
-                iterators = build_train_valid_test_data_iterators(vp_stage_train_valid_test_dataset_provider)
-                self.train_data_iterator.append(iterators[0])
-                self.valid_data_iterator.append(iterators[1])
-                self.test_data_iterator.append(iterators[2])
-        else:
-            (
-                self.train_data_iterator,
-                self.valid_data_iterator,
-                self.test_data_iterator,
-            ) = build_train_valid_test_data_iterators(train_valid_test_datasets_provider_func)
-        timers("train/valid/test-data-iterators-setup").stop()
-        print_datetime("after dataloaders are built")
-        self.app_metrics["app_build_dataiters_finish_time"] = one_logger_utils.get_timestamp_in_ms()
-
-        # Track if training is enabled. Can only be done once args.do_train is assigned after dataloader is built.
-        # The pinned Megatron-LM's ``track_config_flags`` signature varies (6 vs 8
-        # positional args across commits); pass only what the installed version accepts.
-        import inspect as _inspect
-
-        _tcf_args = [
-            args.train_iters,
-            args.skip_train,
-            args.do_train,
-            args.do_valid,
-            args.do_test,
-            args.dataloader_type,
-            args.retro_project_dir,
-            args.retro_cyclic_train_iters,
-        ]
-        try:
-            _tcf_nparams = len(_inspect.signature(one_logger_utils.track_config_flags).parameters)
-        except (TypeError, ValueError):
-            _tcf_nparams = len(_tcf_args)
-        one_logger_utils.track_config_flags(*_tcf_args[:_tcf_nparams])
-
-        # Print setup timing.
-        log_rank_0("done with setup ...")
-        timers.log(
-            ["model-and-optimizer-setup", "train/valid/test-data-iterators-setup"],
-            barrier=True,
-        )
-
-        one_logger = get_one_logger()
-        one_logger and one_logger.log_metrics(self.app_metrics)
-
-    def core_gpt_dataset_config_from_args(self, args):
-        tokenizer = get_tokenizer()
-
-        # Keep legacy trainer aligned with upstream pretrain_gpt dataset argument handling.
-        blend, blend_per_split = get_blend_and_blend_per_split(args)
-
-        sequences_per_dataset = None
-        per_dataset_sequences_path = getattr(args, "per_dataset_sequences_path", None)
-        if per_dataset_sequences_path is not None:
-            with open(per_dataset_sequences_path, "r") as f:
-                sequences_per_dataset = json.load(f)
-
-        data_args = {
-            "random_seed": args.seed,
-            "sequence_length": args.seq_length,
-            "blend": blend,
-            "blend_per_split": blend_per_split,
-            "split": args.split,
-            "multiple_validation_sets": getattr(args, "multiple_validation_sets", None),
-            "full_validation": getattr(args, "full_validation", None),
-            "num_dataset_builder_threads": args.num_dataset_builder_threads,
-            "path_to_cache": args.data_cache_path,
-            "mmap_bin_files": args.mmap_bin_files,
-            "tokenizer": tokenizer,
-            "reset_position_ids": args.reset_position_ids,
-            "reset_attention_mask": args.reset_attention_mask,
-            "eod_mask_loss": args.eod_mask_loss,
-            "create_attention_mask": args.create_attention_mask_in_dataloader,
-            "object_storage_cache_path": getattr(args, "object_storage_cache_path", None),
-            "mid_level_dataset_surplus": getattr(args, "mid_level_dataset_surplus", 0.005),
-            "allow_ambiguous_pad_tokens": getattr(args, "allow_ambiguous_pad_tokens", False),
-            "fast_cache_load": getattr(args, "dataloader_fast_cache_load", False),
-            "sequences_per_dataset": sequences_per_dataset,
-            "defer_npy_index_mmap": getattr(args, "dataloader_defer_npy_index_mmap", False),
-            "context_parallel_size": getattr(args, "context_parallel_size", 1),
-            "data_parallel_size": getattr(args, "data_parallel_size", 1),
-            "sequence_parallel_size": getattr(args, "tensor_model_parallel_size", 1)
-            * getattr(args, "sequence_parallel", False),
-            "hybrid_context_parallel": getattr(args, "hybrid_context_parallel", False),
-        }
-
-        return GPTDatasetConfig(**data_args)
-
-    def train_valid_test_datasets_provider(self, train_val_test_num_samples, vp_stage=None):
-        """Build the train test and validation datasets.
-
-        Args:
-            train_val_test_num_samples : A list containing the number of samples in train test and validation.
-        """
-        args = get_args()
-
-        config = self.core_gpt_dataset_config_from_args(args)
-
-        if args.mock_data:
-            dataset_type = MockGPTDataset
-        else:
-            dataset_type = GPTDataset
-
-        def is_dataset_built_on_rank(vp_stage=None):
-            return (
-                is_first_or_last_pipeline_stage(vp_stage)
-                and parallel_state.get_tensor_model_parallel_rank() == 0
-            )
-
-        log_rank_0("> building train, validation, and test datasets for GPT ...")
-        train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
-            dataset_type,
-            train_val_test_num_samples,
-            functools.partial(is_dataset_built_on_rank, vp_stage=vp_stage),
-            config,
-        ).build()
-
-        log_rank_0("> finished creating GPT datasets ...")
-
-        return train_ds, valid_ds, test_ds
-
-    def initialize_megatron(
-        self,
-        extra_args_provider=None,
-        args_defaults={},
-        ignore_unknown_args=False,
-        allow_no_cuda=False,
-        skip_mpu_initialization=False,
-        get_embedding_ranks=None,
-        get_position_embedding_ranks=None,
-    ):
-        """Set global variables, initialize distributed, and
-        set autoresume and random seeds.
-        `allow_no_cuda` should not be set unless using megatron for cpu only
-        data processing. In general this arg should not be set unless you know
-        what you are doing.
-        Returns a function to finalize distributed env initialization
-        (optionally, only when args.lazy_mpu_init == True)
-        """
-        if not allow_no_cuda:
-            # Make sure cuda is available.
-            assert torch.cuda.is_available(), "Megatron requires CUDA."
-
-        # Note: parse_args is deprecated in megatron trainer, use primus yaml config instead.
-        # Parse arguments
-        # args = parse_args(extra_args_provider, ignore_unknown_args)
-
-        # Use trainer args from primus
-        # args = self.module_config
-
-        # Build Megatron arguments by merging Primus config into Megatron defaults.
-        # 1) Load Megatron defaults (no Primus overrides)
-        megatron_defaults = _load_megatron_defaults()
-
-        # 2) Convert Primus module_config (nested namespace) to a plain dict
-        primus_args = nested_namespace_to_dict(self.module_config)
-
-        # 3) Merge: defaults < primus_args
-        merged_args = megatron_defaults.copy()
-        merged_args.update(primus_args)
-
-        # 4) Convert to Namespace for compatibility with downstream Megatron utilities
-        args = argparse.Namespace(**merged_args)
-
-        # Prep for checkpoint conversion.
-        if args.ckpt_convert_format is not None:
-            assert args.ckpt_convert_save is not None
-            assert args.load is not None
-            args.exit_on_missing_checkpoint = True
-
-        log_kv_rank_0(f"-load", f"{args.load}")
-        log_kv_rank_0(f"-use_checkpoint_args", f"{args.use_checkpoint_args}")
-        if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
-            checker.check_true(args.load is not None, "--use-checkpoints-args requires --load argument")
-            log_rank_0(f"-load_args_from_checkpoint...")
-            assert args.non_persistent_ckpt_type != "local", (
-                "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. "
-                "Two-stage checkpoint loading is not implemented, and all arguments must be defined "
-                "before initializing LocalCheckpointManager."
-            )
-            load_args_from_checkpoint(args)
-
-        if args.async_save and args.use_persistent_ckpt_worker:
-            init_persistent_async_worker()
-
-        checker.check_true(args.yaml_cfg is None, "Xpipe doesn't support megatron yaml config.")
-        if args.yaml_cfg is not None:
-            args = validate_yaml(args, args_defaults)
-        else:
-            if args.decoder_pipeline_manual_split_list is not None:
-                from .utils import validate_args_modified
-
-                ori_code = "if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None:"
-                new_code = (
-                    "if args.decoder_pipeline_manual_split_list is None and " + ori_code.split("if ")[-1]
-                )
-
-                validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code)
-            elif args.fp4 is not None:
-                # TODO(ruibin): Remove it when ROCm TE upgrade to 2.7.0.dev0
-                from .utils import validate_args_modified
-
-                ori_code = """raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.")"""
-                new_code = """pass"""
-
-                validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code)
-            else:
-                validate_args(args, args_defaults)
-
-        # monkey patch _set_wandb_writer before set_global_variables
-        log_rank_0(f"-monkey patch megatron.training.global_vars._set_wandb_writer...")
-        megatron.training.global_vars._set_wandb_writer = set_wandb_writer_patch
-
-        # set global args, build tokenizer, and set adlr-autoresume,
-        # tensorboard-writer, and timers.
-        log_rank_0(f"-set_global_variables...")
-        set_global_variables(args, build_tokenizer=False)
-        log_rank_0(f"-set_primus_global_variables...")
-        set_primus_global_variables(args)
-        args = get_args()
-
-        # set tokenizer
-        log_rank_0(f"-build_tokenizer...")
-        global_vars._ensure_var_is_not_initialized(global_vars._GLOBAL_TOKENIZER, "tokenizer")
-        global_vars._GLOBAL_TOKENIZER = build_tokenizer(args)
-
-        # set logging level
-        setup_logging()
-
-        # init rerun state
-        def state_save_func():
-            return {"rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states()}
-
-        def state_restore_func(state_dict):
-            if state_dict["rng_tracker_states"]:
-                tensor_parallel.get_cuda_rng_tracker().set_states(state_dict["rng_tracker_states"])
-
-        initialize_rerun_state_machine(
-            state_save_func=state_save_func,
-            state_restore_func=state_restore_func,
-            mode=RerunMode(args.rerun_mode),
-            error_injector=RerunErrorInjector(
-                error_injection_rate=args.error_injection_rate,
-                error_injection_type=RerunDiagnostic(args.error_injection_type),
-            ),
-            result_rejected_tracker_filename=args.result_rejected_tracker_filename,
-        )
-
-        # torch.distributed initialization
-        def finish_mpu_init():
-            args = get_args()
-            # Pytorch distributed.
-            log_rank_0(f"-initialize_distributed...")
-            _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, None)
-
-            # Random seeds for reproducibility.
-            log_kv_rank_0(f"-seeds", f"{args.seed}")
-            _set_random_seed(
-                args.seed,
-                args.data_parallel_random_init,
-                args.te_rng_tracker,
-                args.inference_rng_tracker,
-                use_cudagraphable_rng=args.enable_cuda_graph,
-            )
-
-            # Setup MoE aux loss scale value.
-            if args.num_experts is not None:
-                from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler
-
-                MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device()))
-
-        if skip_mpu_initialization:
-            return None
-
-        args = get_args()
-        log_kv_rank_0(f"-lazy_mpu_init", f"{args.lazy_mpu_init}")
-        if args.lazy_mpu_init:
-            # TODO is this still a necessary option?
-            args.use_cpu_initialization = True
-            # delayed initialization of DDP-related stuff
-            # We only set basic DDP globals
-            mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size)
-            # and return function for external DDP manager
-            # to call when it has DDP initialized
-            mpu.set_tensor_model_parallel_rank(args.rank)
-            return finish_mpu_init
-        else:
-            # Megatron's MPU is the master. Complete initialization right away.
-            finish_mpu_init()
-
-            # Autoresume.
-            _init_autoresume()
-
-            # Compile dependencies.
-            if not args.disable_compile_dependencies:
-                log_rank_0(f"-compile_dependencies...")
-                _compile_dependencies()
-
-            if args.tp_comm_overlap:
-                # TODO: Should this be activated with just decoder-tp-comm-overlap too?
-                _initialize_tp_communicators()
-
-            # No continuation function
-            return None
-
-    def setup_model_and_optimizer(
-        self,
-        model_provider_func,
-        model_type,
-        no_wd_decay_cond=None,
-        scale_lr_cond=None,
-        lr_mult=1.0,
-        checkpointing_context=None,
-    ):
-        """Setup model and optimizer."""
-        args = get_args()
-        timers = get_timers()
-        one_logger = get_one_logger()
-
-        # Primus Turbo patches are now applied automatically via the new patch system
-        # in BaseTrainer.run() with phase="before_train"
-        if importlib.util.find_spec("primus_turbo") is not None:
-            args = get_args()
-            if args.tensor_model_parallel_size == 1:
-                if args.enable_primus_turbo:
-                    log_rank_0(f"use pt backend...")
-                else:
-                    log_rank_0(f"use te backend...")
-            elif args.enable_primus_turbo:
-                log_rank_0(f"primus turbo does not support tp, use te backend...")
-        else:
-            log_rank_0(f"use te backend...")
-
-        log_rank_0(f"-run get_model")
-        log_rank_0(f"-model_provider_func: {model_provider_func}")
-        log_rank_0(f"-model_type: {model_type}")
-        model = get_model(model_provider_func, model_type)
-        log_rank_0(model)
-        # get_megatron_optimizer will use the ddp_config
-        if isinstance(model[0], torch_FSDP):
-            model[0].ddp_config = DistributedDataParallelConfig()
-            model[0].ddp_config.use_custom_fsdp = False
-
-        unwrapped_model = unwrap_model(model)
-
-        config, config_overrides = get_megatron_optimizer_config(args)
-        config.timers = timers
-        if getattr(args, "use_mup", False):
-            model_config_source = unwrapped_model[0] if isinstance(unwrapped_model, list) else unwrapped_model
-            model_config = get_model_config(model_config_source)
-            mup_overrides = get_mup_config_overrides(
-                config=config,
-                mup_width_mult=model_config.mup_width_mult,
-                optimizer_type=config.optimizer,
-            )
-            if mup_overrides:
-                config_overrides = {**(config_overrides or {}), **mup_overrides}
-
-        if "muon" not in args.optimizer:
-            optimizer = get_megatron_optimizer(
-                config,
-                model,
-                config_overrides=config_overrides,
-                use_gloo_process_groups=args.enable_gloo_process_groups,
-                dump_param_to_param_group_map=getattr(args, "dump_param_to_param_group_map", None),
-            )
-        else:
-            kwargs = {}
-            for f in dataclasses.fields(MounOptimizerConfig):
-                if hasattr(args, f.name):
-                    kwargs[f.name] = getattr(args, f.name)
-
-            config = MounOptimizerConfig(**kwargs)
-            config.timers = timers
-            optimizer = get_megatron_muon_optimizer(
-                config,
-                model,
-                config_overrides=config_overrides,
-                use_gloo_process_groups=args.enable_gloo_process_groups,
-                layer_wise_distributed_optimizer="dist" in config.optimizer,
-                dump_param_to_param_group_map=getattr(args, "dump_param_to_param_group_map", None),
-            )
-
-        opt_param_scheduler = get_optimizer_param_scheduler(optimizer)
-
-        if args.moe_use_upcycling:
-            torch.distributed.barrier()
-            assert not checkpoint_exists(args.save), (
-                "The upcycling destination directory already exists. "
-                "Please check if --moe-use-upcycling is mistakenly enabled. "
-                "Upcycling should only be set for the first run when converting the dense model. "
-                "All subsequent runs should remove this flag. "
-            )
-            num_experts = args.num_experts
-            args.num_experts = None
-            expert_model_parallel_size = args.expert_model_parallel_size
-            args.expert_model_parallel_size = 1
-            dense_model_for_upcycling = get_model(model_provider_func, model_type)
-            args.num_experts = num_experts
-            args.expert_model_parallel_size = expert_model_parallel_size
-            _, args.num_floating_point_operations_so_far = upcycling_utils.load_and_upcycle_model(
-                load_checkpoint,
-                unwrapped_model,
-                dense_model_for_upcycling,
-                load_kwargs={
-                    "model": dense_model_for_upcycling,
-                    "optimizer": None,
-                    "opt_param_scheduler": None,
-                },
-            )
-            args.iteration = 1
-            save_checkpoint(
-                args.iteration,
-                model,
-                None,
-                None,
-                args.num_floating_point_operations_so_far,
-            )
-            torch.distributed.barrier()
-            del dense_model_for_upcycling
-            if (args.fp16 or args.bf16) and optimizer is not None:
-                optimizer.reload_model_params()
-            log_rank_0(f"Upcycled checkpoint saved to {args.save}")
-
-        if (args.load is not None or args.pretrained_checkpoint is not None) and not args.moe_use_upcycling:
-            one_logger and one_logger.log_metrics(
-                {"load_checkpoint_start_time": one_logger_utils.get_timestamp_in_ms()}
-            )
-            timers("load-checkpoint", log_level=0).start(barrier=True)
-
-            log_rank_0(f"-run load_checkpoint")
-            log_rank_0(f"  -args.load={args.load}")
-            args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
-                model,
-                optimizer,
-                opt_param_scheduler,
-                checkpointing_context=self.checkpointing_context,
-                skip_load_to_model_and_opt=HAVE_FSDP2 and args.use_torch_fsdp2,
-            )
-            timers("load-checkpoint").stop(barrier=True)
-            timers.log(["load-checkpoint"])
-            one_logger and one_logger.log_metrics(
-                {
-                    "load_checkpoint_finish_time": one_logger_utils.get_timestamp_in_ms(),
-                    "load_checkpoint_time": timers("load-checkpoint").active_time(),
-                }
-            )
-        else:
-            args.iteration = 0
-            args.num_floating_point_operations_so_far = 0
-
-        # get model without FP16 and/or DDP wrappers
-        if (
-            args.iteration == 0
-            and len(unwrapped_model) == 1
-            and hasattr(unwrapped_model[0], "init_state_dict_from_bert")
-        ):
-            log_rank_0("Initializing ICT from pretrained BERT model")
-            unwrapped_model[0].init_state_dict_from_bert()
-            if args.fp16:
-                optimizer.reload_model_params()
-
-        # Convert checkpoint format.
-        if args.ckpt_convert_format is not None:
-            load_ckpt_format = args.ckpt_format
-            args.ckpt_format = args.ckpt_convert_format
-            args.save = os.path.join(args.ckpt_convert_save, args.ckpt_convert_format)
-            update_use_dist_ckpt(args)
-
-            save_checkpoint(
-                args.iteration,
-                model,
-                optimizer,
-                opt_param_scheduler,
-                args.num_floating_point_operations_so_far,
-                preprocess_common_state_dict_fn=preprocess_common_state_dict,
-            )
-
-            log_rank_0("> converted checkpoint: %s -> %s." % (load_ckpt_format, args.ckpt_format))
-            torch.distributed.barrier()
-            exit()
-
-        return model, optimizer, opt_param_scheduler
-
-    def run(self, *args, **kwargs):
-        one_logger = get_one_logger()
-        args = get_args()
-
-        process_non_loss_data_func = None
-        non_loss_data_func = None
-        if not args.skip_train:
-            log_rank_0("training ...")
-
-            if args.dataloader_type == "cyclic" and args.retro_project_dir:
-                assert args.retro_cyclic_train_iters is not None
-                args.train_iters = args.retro_cyclic_train_iters
-                log_rank_0("retro cyclic train iters : %d" % args.train_iters)
-
-            iteration = 0
-            if args.do_train and args.train_iters > 0:
-                iteration, num_floating_point_operations_so_far = self.train(
-                    self.forward_step,
-                    self.model,
-                    self.optimizer,
-                    self.opt_param_scheduler,
-                    self.train_data_iterator,
-                    self.valid_data_iterator,
-                    process_non_loss_data_func,
-                    self.config,
-                    self.checkpointing_context,
-                    non_loss_data_func,
-                )
-
-            print_datetime("after training is done")
-
-            if (
-                args.save
-                and iteration != 0
-                and iteration % args.save_interval != 0
-                and not args.disable_last_saving
-            ):
-                save_checkpoint(
-                    iteration,
-                    self.model,
-                    self.optimizer,
-                    self.opt_param_scheduler,
-                    num_floating_point_operations_so_far,
-                    self.checkpointing_context,
-                    train_data_iterator=self.train_data_iterator,
-                    preprocess_common_state_dict_fn=preprocess_common_state_dict,
-                )
-
-            one_logger and one_logger.log_metrics(
-                {"app_train_loop_finish_time": one_logger_utils.get_timestamp_in_ms()}
-            )
-
-        else:
-            log_rank_0("skipping training (--skip-train is on) ...")
-
-            iteration = args.iteration
-
-        if args.do_valid:
-            prefix = f"iteration {iteration} on validation set"
-            evaluate_and_print_results(
-                prefix,
-                self.forward_step,
-                self.valid_data_iterator,
-                self.model,
-                iteration,
-                process_non_loss_data_func,
-                self.config,
-                verbose=True,
-                write_to_tensorboard=not args.skip_train,
-                non_loss_data_func=non_loss_data_func,
-            )
-
-        if args.do_test:
-            prefix = f"iteration {iteration} on test set"
-            evaluate_and_print_results(
-                prefix,
-                self.forward_step,
-                self.test_data_iterator,
-                self.model,
-                iteration,
-                process_non_loss_data_func,
-                self.config,
-                verbose=True,
-                write_to_tensorboard=not args.skip_train,
-                non_loss_data_func=non_loss_data_func,
-            )
-
-        wandb_writer = get_wandb_writer()
-        if wandb_writer:
-            wandb_writer.finish()
-
-        ft_integration.on_checkpointing_start()
-        maybe_finalize_async_save(blocking=True, terminate=True)
-        ft_integration.on_checkpointing_end(is_async_finalization=True)
-
-        mlflow_writer = get_mlflow_writer()
-        # Barrier to ensure all ranks have finished writing files before upload.
-        # Must run on ALL ranks to avoid deadlock (only last rank has mlflow_writer).
-        if dist.is_initialized():
-            dist.barrier()
-
-        # Always call: uploads to MLflow when enabled; when MLflow disabled, still runs
-        # local-only TraceLens report generation if generate_tracelens_report=True.
-        try:
-            upload_mlflow_artifacts(
-                tensorboard_dir=args.tensorboard_dir,
-                exp_root_path=self.exp_root_path,
-                upload_traces=getattr(args, "mlflow_upload_traces", False),
-                upload_logs=getattr(args, "mlflow_upload_logs", False),
-                generate_tracelens_report=getattr(args, "generate_tracelens_report", False),
-                upload_tracelens_report=getattr(args, "mlflow_upload_tracelens_report", False),
-                tracelens_ranks=getattr(args, "mlflow_tracelens_ranks", None),
-                tracelens_output_format=getattr(args, "mlflow_tracelens_output_format", "xlsx"),
-                tracelens_cleanup_after_upload=getattr(args, "mlflow_tracelens_cleanup_after_upload", False),
-                tracelens_auto_install=getattr(args, "mlflow_tracelens_auto_install", True),
-            )
-        except Exception as e:
-            import logging
-
-            logging.getLogger(__name__).warning("[MLflow] Artifact upload failed: %s", e)
-        finally:
-            if mlflow_writer:
-                mlflow_writer.end_run()
-            if dist.is_initialized():
-                dist.barrier()
-
-        one_logger and one_logger.log_metrics({"app_finish_time": one_logger_utils.get_timestamp_in_ms()})
-
-        ft_integration.shutdown()
-        one_logger_utils.finish()
-
-        # clean up torch pg resources on exit
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-    def train(
-        self,
-        forward_step_func,
-        model,
-        optimizer,
-        opt_param_scheduler,
-        train_data_iterator,
-        valid_data_iterator,
-        process_non_loss_data_func,
-        config,
-        checkpointing_context,
-        non_loss_data_func,
-    ):
-        """Training function: run train_step desired number of times, run validation, checkpoint."""
-        args = get_args()
-        timers = get_timers()
-        one_logger = get_one_logger()
-
-        if args.run_workload_inspector_server:
-            try:
-                import threading
-
-                from workload_inspector.utils.webserver import run_server
-
-                threading.Thread(target=run_server, daemon=True, args=(torch.distributed.get_rank(),)).start()
-            except ModuleNotFoundError:
-                log_rank_0("workload inspector module not found.")
-
-        # Write args to tensorboard
-        write_args_to_tensorboard()
-
-        # Turn on training mode which enables dropout.
-        for model_module in model:
-            model_module.train()
-
-        # Tracking loss.
-        total_loss_dict = {}
-
-        # Iterations.
-        iteration = args.iteration
-        # Make sure rerun_state_machine has the right iteration loaded from checkpoint.
-        rerun_state_machine = get_rerun_state_machine()
-        if rerun_state_machine.current_iteration != iteration:
-            log_rank_0(f"Setting rerun_state_machine.current_iteration to {iteration}...")
-            rerun_state_machine.current_iteration = iteration
-
-        # Track E2E metrics at the start of training.
-        one_logger_utils.on_train_start(
-            iteration=iteration,
-            consumed_train_samples=args.consumed_train_samples,
-            train_samples=args.train_samples,
-            seq_length=args.seq_length,
-            train_iters=args.train_iters,
-            save=args.save,
-            async_save=args.async_save,
-            log_throughput=args.log_throughput,
-            num_floating_point_operations_so_far=args.num_floating_point_operations_so_far,
-        )
-
-        num_floating_point_operations_so_far = args.num_floating_point_operations_so_far
-
-        # Setup some training config params.
-        config.grad_scale_func = optimizer.scale_loss
-        config.timers = timers
-
-        if isinstance(model[0], (get_custom_fsdp(), DDP)) and args.overlap_grad_reduce:
-            assert config.no_sync_func is None, (
-                "When overlap_grad_reduce is True, config.no_sync_func must be None; "
-                "a custom no_sync_func is not supported when overlapping grad-reduce"
-            )
-            config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
-            if len(model) == 1:
-                config.no_sync_func = config.no_sync_func[0]
-            if args.align_grad_reduce:
-                config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
-                if len(model) == 1:
-                    config.grad_sync_func = config.grad_sync_func[0]
-        if args.overlap_param_gather and args.align_param_gather:
-            config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
-            if len(model) == 1:
-                config.param_sync_func = config.param_sync_func[0]
-        config.finalize_model_grads_func = finalize_model_grads
-
-        timers("interval-time", log_level=0).start(barrier=True)
-        print_datetime("before the start of training step")
-        report_memory_flag = True
-        pre_hook_enabled = False
-        should_exit = False
-        exit_code = 0
-
-        if args.manual_gc:
-            # Disable the default garbage collector and perform the collection manually.
-            # This is to align the timing of garbage collection across ranks.
-            assert (
-                args.manual_gc_interval >= 0
-            ), "Manual garbage collection interval should be larger than or equal to 0"
-            gc.disable()
-            gc.collect()
-
-        # Singleton initialization of straggler detector.
-        if args.log_straggler:
-            global stimer
-            world = torch.distributed.get_world_size()
-            rank = torch.distributed.get_rank()
-            mmcnt = args.straggler_minmax_count
-            stimer.configure(
-                world,
-                rank,
-                mmcnt=mmcnt,
-                enabled=not args.disable_straggler_on_startup,
-                port=args.straggler_ctrlr_port,
-            )
-        num_floating_point_operations_since_last_log_event = 0.0
-
-        num_microbatches = get_num_microbatches()
-        eval_duration = 0.0
-        eval_iterations = 0
-
-        def get_e2e_base_metrics():
-            """Get base metrics values for one-logger to calculate E2E tracking metrics."""
-            num_floating_point_operations_since_current_train_start = (
-                num_floating_point_operations_so_far - args.num_floating_point_operations_so_far
-            )
-            return {
-                "iteration": iteration,
-                "train_duration": timers("interval-time").active_time(),
-                "eval_duration": eval_duration,
-                "eval_iterations": eval_iterations,
-                "total_flops_since_current_train_start": num_floating_point_operations_since_current_train_start,
-                "num_floating_point_operations_so_far": num_floating_point_operations_so_far,
-                "consumed_train_samples": args.consumed_train_samples,
-                "world_size": args.world_size,
-                "seq_length": args.seq_length,
-            }
-
-        # Cache into one-logger for callback.
-        if one_logger:
-            with one_logger.get_context_manager():
-                one_logger.store_set("get_e2e_base_metrics", get_e2e_base_metrics)
-
-        prof = None
-        if args.profile and torch.distributed.get_rank() in args.profile_ranks and args.use_pytorch_profiler:
-            activities = [torch.profiler.ProfilerActivity.CUDA]
-            if not args.disable_profiler_activity_cpu:
-                activities.append(torch.profiler.ProfilerActivity.CPU)
-            worker_name = (
-                f"primus-megatron-exp[{self.exp_meta_info['exp_name']}]-rank[{torch.distributed.get_rank()}]"
-            )
-            prof = torch.profiler.profile(
-                activities=activities,
-                schedule=torch.profiler.schedule(
-                    wait=max(args.profile_step_start - 1, 0),
-                    warmup=1 if args.profile_step_start > 0 else 0,
-                    active=args.profile_step_end - args.profile_step_start,
-                    repeat=1,
-                ),
-                on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    args.tensorboard_dir,
-                    worker_name=worker_name,
-                    use_gzip=args.torch_profiler_use_gzip,
-                ),
-                record_shapes=args.torch_profiler_record_shapes,
-                with_stack=args.torch_profiler_with_stack,
-            )
-            prof.start()
-
-        start_iteration = iteration
-        # Disable forward pre-hook to start training to ensure that errors in checkpoint loading
-        # or random initialization don't propagate to all ranks in first all-gather (which is a
-        # no-op if things work correctly).
-        if should_disable_forward_pre_hook(args):
-            disable_forward_pre_hook(model, param_sync=False)
-            # Also remove param_sync_func temporarily so that sync calls made in
-            # `forward_backward_func` are no-ops.
-            param_sync_func = config.param_sync_func
-            config.param_sync_func = None
-            pre_hook_enabled = False
-        # Also, check weight hash across DP replicas to be very pedantic.
-        if args.check_weight_hash_across_dp_replicas_interval is not None:
-            assert check_param_hashes_across_dp_replicas(
-                model, cross_check=True
-            ), "Parameter hashes not matching across DP replicas"
-            torch.distributed.barrier()
-            log_rank_0(f">>> Weight hashes match after {iteration} iterations...")
-
-        if args.dump_pp_data:
-            from .utils import set_dump_pp_data_patch
-
-            set_dump_pp_data_patch()
-            log_rank_0(f"dump pp schedule data for visualization")
-
-        # Run training iterations till done.
-        while iteration < args.train_iters:
-            if args.profile and torch.distributed.get_rank() in args.profile_ranks:
-                if args.use_pytorch_profiler:
-                    prof.step()
-                elif iteration == args.profile_step_start:
-                    torch.cuda.cudart().cudaProfilerStart()
-                    torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__()
-
-            ft_integration.on_checkpointing_start()
-            maybe_finalize_async_save(blocking=False)
-            ft_integration.on_checkpointing_end(is_async_finalization=True)
-
-            # Update number of microbatches first without consistency check to decide if a
-            # checkpoint should be saved. If the number of microbatches is different
-            # from the previous iteration, save a checkpoint. Then run consistency check
-            # to make sure training configuration is still valid.
-            update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True)
-            if get_num_microbatches() != num_microbatches and iteration != 0:
-                assert get_num_microbatches() > num_microbatches, (
-                    f"Number of microbatches should be increasing due to batch size rampup; "
-                    f"instead going from {num_microbatches} to {get_num_microbatches()}"
-                )
-                if args.save is not None:
-                    save_checkpoint_and_time(
-                        iteration,
-                        model,
-                        optimizer,
-                        opt_param_scheduler,
-                        num_floating_point_operations_so_far,
-                        checkpointing_context,
-                        train_data_iterator=train_data_iterator,
-                    )
-            num_microbatches = get_num_microbatches()
-            update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True)
-
-            # Completely skip iteration if needed.
-            if iteration in args.iterations_to_skip:
-                # Dummy train_step to fast forward train_data_iterator.
-                dummy_train_step(train_data_iterator)
-                iteration += 1
-                batch_size = (
-                    mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
-                )
-                args.consumed_train_samples += batch_size
-                args.skipped_train_samples += batch_size
-                continue
-
-            # Run training step.
-            args.curr_iteration = iteration
-            ft_integration.on_training_step_start()
-            (
-                loss_dict,
-                skipped_iter,
-                should_checkpoint,
-                should_exit,
-                exit_code,
-                grad_norm,
-                num_zeros_in_grad,
-            ) = self.train_step(
-                forward_step_func,
-                train_data_iterator,
-                model,
-                optimizer,
-                opt_param_scheduler,
-                config,
-            )
-            ft_integration.on_training_step_end()
-            if should_checkpoint:
-                save_checkpoint_and_time(
-                    iteration,
-                    model,
-                    optimizer,
-                    opt_param_scheduler,
-                    num_floating_point_operations_so_far,
-                    checkpointing_context,
-                    train_data_iterator=train_data_iterator,
-                )
-            if should_exit:
-                break
-
-            # Enable forward pre-hooks after first set of forward and backward passes.
-            # When running in fp16, skip all NaN iterations until steady-state loss scaling value
-            # is reached.
-            if iteration == start_iteration:
-                if skipped_iter:
-                    # Only enable forward pre-hook after a training step has successfully run. Relevant
-                    # for fp16 codepath where first XX iterations are skipped until steady-state loss
-                    # scale value is reached.
-                    start_iteration = iteration + 1
-                else:
-                    # Enable forward pre-hook after training step has successfully run. All subsequent
-                    # forward passes will use the forward pre-hook / `param_sync_func` in
-                    # `forward_backward_func`.
-                    if should_disable_forward_pre_hook(args):
-                        enable_forward_pre_hook(model)
-                        config.param_sync_func = param_sync_func
-                        pre_hook_enabled = True
-
-            iteration += 1
-            batch_size = mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
-            args.consumed_train_samples += batch_size
-            num_skipped_samples_in_batch = (
-                get_current_global_batch_size() - get_current_running_global_batch_size()
-            )
-            if args.decrease_batch_size_if_needed:
-                assert num_skipped_samples_in_batch >= 0
-            else:
-                assert num_skipped_samples_in_batch == 0
-            args.skipped_train_samples += num_skipped_samples_in_batch
-            flops_calc = (
-                num_floating_point_operations
-                if not args.multi_latent_attention
-                else self.num_floating_point_operations_mla_moe
-            )
-            num_floating_point_operations_in_batch = flops_calc(args, batch_size)
-            num_floating_point_operations_so_far += num_floating_point_operations_in_batch
-            num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch
-
-            # Logging.
-            if not optimizer.is_stub_optimizer:
-                loss_scale = optimizer.get_loss_scale().item()
-            else:
-                loss_scale = 1.0
-            params_norm = None
-
-            if args.log_params_norm:
-                params_norm = calc_params_l2_norm(model)
-            learning_rate = None
-            decoupled_learning_rate = None
-            for param_group in optimizer.param_groups:
-                if param_group["is_decoupled_lr"]:
-                    decoupled_learning_rate = param_group["lr"]
-                else:
-                    learning_rate = param_group["lr"]
-            report_memory_flag = self.training_log(
-                loss_dict,
-                total_loss_dict,
-                learning_rate,
-                decoupled_learning_rate,
-                iteration,
-                loss_scale,
-                report_memory_flag,
-                skipped_iter,
-                grad_norm,
-                params_norm,
-                num_zeros_in_grad,
-            )
-
-            # Evaluation.
-            if args.eval_interval and iteration % args.eval_interval == 0 and args.do_valid:
-                timers("interval-time").stop()
-                if should_disable_forward_pre_hook(args):
-                    disable_forward_pre_hook(model)
-                    pre_hook_enabled = False
-                if args.manual_gc and args.manual_gc_eval:
-                    # Collect all objects.
-                    gc.collect()
-                prefix = f"iteration {iteration}"
-                timers("eval-time", log_level=0).start(barrier=True)
-                evaluate_and_print_results(
-                    prefix,
-                    forward_step_func,
-                    valid_data_iterator,
-                    model,
-                    iteration,
-                    process_non_loss_data_func,
-                    config,
-                    verbose=False,
-                    write_to_tensorboard=True,
-                    non_loss_data_func=non_loss_data_func,
-                )
-                eval_duration += timers("eval-time").elapsed()
-                eval_iterations += args.eval_iters
-                timers("eval-time").stop()
-                one_logger_utils.track_e2e_metrics()
-
-                if args.manual_gc and args.manual_gc_eval:
-                    # Collect only the objects created and used in evaluation.
-                    gc.collect(generation=0)
-                if should_disable_forward_pre_hook(args):
-                    enable_forward_pre_hook(model)
-                    pre_hook_enabled = True
-                timers("interval-time", log_level=0).start(barrier=True)
-
-            # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
-            # Some of these only happen at specific iterations.
-            post_training_step_callbacks(
-                model,
-                optimizer,
-                opt_param_scheduler,
-                iteration,
-                prof,
-                num_floating_point_operations_since_last_log_event,
-            )
-
-            # Checkpoint and decide whether to exit.
-            should_exit = checkpoint_and_decide_exit(
-                model,
-                optimizer,
-                opt_param_scheduler,
-                iteration,
-                num_floating_point_operations_so_far,
-                checkpointing_context,
-                train_data_iterator,
-            )
-            if should_exit:
-                break
-
-        one_logger_utils.track_e2e_metrics()
-
-        if args.dump_pp_data:
-            from .utils import dump_pp_data
-
-            pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data")
-            dump_pp_data(args, get_num_microbatches(), pp_data_dir)
-            log_rank_0(f"pp schedule data dumped to {pp_data_dir}")
-
-        # Flush TensorBoard, WandB writers and one-logger.
-        writer = get_tensorboard_writer()
-        if writer:
-            writer.flush()
-
-        # Close out pre-hooks if using distributed optimizer and overlapped param gather.
-        if pre_hook_enabled:
-            disable_forward_pre_hook(model)
-
-        ft_integration.on_checkpointing_start()
-        # This will finalize all unfinalized async request and terminate
-        # a persistent async worker if persistent ckpt worker is enabled
-        maybe_finalize_async_save(blocking=True, terminate=True)
-        ft_integration.on_checkpointing_end(is_async_finalization=True)
-        if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None:
-            ft_integration.get_rank_monitor_client().shutdown_workload_monitoring()
-
-        # If any exit conditions (signal handler, duration, iterations) have been reached, exit.
-        if should_exit:
-            wandb_writer = get_wandb_writer()
-            if wandb_writer:
-                wandb_writer.finish()
-            mlflow_writer = get_mlflow_writer()
-            # Barrier to ensure all ranks have finished writing files before upload.
-            # Must run on ALL ranks to avoid deadlock (only last rank has mlflow_writer).
-            if dist.is_initialized():
-                dist.barrier()
-            try:
-                upload_mlflow_artifacts(
-                    tensorboard_dir=args.tensorboard_dir,
-                    exp_root_path=self.exp_root_path,
-                    upload_traces=getattr(args, "mlflow_upload_traces", False),
-                    upload_logs=getattr(args, "mlflow_upload_logs", False),
-                    generate_tracelens_report=getattr(args, "generate_tracelens_report", False),
-                    upload_tracelens_report=getattr(args, "mlflow_upload_tracelens_report", False),
-                    tracelens_ranks=getattr(args, "mlflow_tracelens_ranks", None),
-                    tracelens_output_format=getattr(args, "mlflow_tracelens_output_format", "xlsx"),
-                    tracelens_cleanup_after_upload=getattr(
-                        args, "mlflow_tracelens_cleanup_after_upload", False
-                    ),
-                    tracelens_auto_install=getattr(args, "mlflow_tracelens_auto_install", True),
-                )
-            except Exception as e:
-                import logging
-
-                logging.getLogger(__name__).warning("[MLflow] Artifact upload failed: %s", e)
-            finally:
-                if mlflow_writer:
-                    mlflow_writer.end_run()
-                if dist.is_initialized():
-                    dist.barrier()
-            ft_integration.shutdown()
-            sys.exit(exit_code)
-
-        return iteration, num_floating_point_operations_so_far
-
-    def train_step(
-        self,
-        forward_step_func,
-        data_iterator,
-        model,
-        optimizer,
-        opt_param_scheduler,
-        config,
-        no_optimizer_post_validation=False,
-    ):
-        """Single training step."""
-        args = get_args()
-        timers = get_timers()
-
-        def run_forward_backward_func(optimizer=None):
-            """Forward pass.
-            optimizer is not None for running post validation."""
-            from megatron.core.pipeline_parallel import get_forward_backward_func
-
-            forward_backward_func = get_forward_backward_func()
-            if optimizer is None and args.dump_pp_data:
-                forward_backward_func = schedule_wrapper(forward_backward_func)
-            kwargs = {}
-            if optimizer is not None:
-                kwargs["optimizer"] = optimizer
-            return forward_backward_func(
-                forward_step_func=forward_step_func,
-                data_iterator=data_iterator,
-                model=model,
-                num_microbatches=get_num_microbatches() * args.num_seq_splits,
-                seq_length=args.seq_length // args.num_seq_splits,
-                micro_batch_size=args.micro_batch_size,
-                decoder_seq_length=args.decoder_seq_length,
-                forward_only=False,
-                **kwargs,
-            )
-
-        rerun_state_machine = get_rerun_state_machine()
-        while rerun_state_machine.should_run_forward_backward(data_iterator):
-            # Set grad to zero.
-            for model_chunk in model:
-                model_chunk.zero_grad_buffer()
-            optimizer.zero_grad()
-
-            # Forward pass.
-            losses_reduced = run_forward_backward_func()
-
-        should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
-        if should_exit:
-            return {}, True, should_checkpoint, should_exit, exit_code, None, None
-
-        # Empty unused memory.
-        if args.empty_unused_memory_level >= 1:
-            torch.cuda.empty_cache()
-
-        # Vision gradients.
-        if args.vision_pretraining and args.vision_pretraining_type == "dino":
-            unwrapped_model = unwrap_model(model[0])
-            unwrapped_model.cancel_gradients_last_layer(args.curr_iteration)
-
-        # Update parameters.
-
-        timers("optimizer", log_level=1).start(barrier=args.barrier_with_L1_time)
-        # update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
-        if get_args().profile:
-            torch.cuda.nvtx.range_push("Optimizer")
-        if args.patch_zero_bubble and args.enable_optimizer_post_validation:
-            if optimizer.post_validation_enabled and not no_optimizer_post_validation:
-                optimizer.pre_step(args, timers)
-                if get_args().profile:
-                    torch.cuda.nvtx.range_pop()
-                if get_args().profile:
-                    torch.cuda.nvtx.range_push("post_validation_phase")
-                update_successful, grad_norm, num_zeros_in_grad = run_forward_backward_func(optimizer)
-                if get_args().profile:
-                    torch.cuda.nvtx.range_pop()
-                # Here num_zeros_in_grad is a fake name, representing for optimizer_rollback
-            else:
-                update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
-                if get_args().profile:
-                    torch.cuda.nvtx.range_pop()
-            optimizer.record_grad_norm(grad_norm)
-        else:
-            update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
-            if get_args().profile:
-                torch.cuda.nvtx.range_pop()
-
-        timers("optimizer").stop()
-
-        # when freezing sub-models we may have a mixture of successful and unsucessful ranks,
-        # so we must gather across mp ranks
-        update_successful = logical_and_across_model_parallel_group(update_successful)
-        # grad_norm and num_zeros_in_grad will be None on ranks without trainable params,
-        # so we must gather across mp ranks
-        grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm)
-        if args.log_num_zeros_in_grad:
-            num_zeros_in_grad = reduce_max_stat_across_model_parallel_group(num_zeros_in_grad)
-
-        # Vision momentum.
-        if args.vision_pretraining and args.vision_pretraining_type == "dino":
-            unwrapped_model = unwrap_model(model[0])
-            unwrapped_model.update_momentum(args.curr_iteration)
-
-        # Update learning rate.
-        if update_successful:
-            increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size
-            opt_param_scheduler.step(increment=increment)
-            skipped_iter = 0
-        else:
-            skipped_iter = 1
-
-        # Empty unused memory.
-        if args.empty_unused_memory_level >= 2:
-            torch.cuda.empty_cache()
-
-        if is_pipeline_stage_containing_loss():
-            # Average loss across microbatches.
-            loss_reduced = {}
-            for key in losses_reduced[0].keys():
-                numerator = 0
-                denominator = 0
-                for x in losses_reduced:
-                    val = x[key]
-                    # there is one dict per microbatch. in new reporting, we average
-                    # over the total number of tokens across the global batch.
-                    if isinstance(val, tuple) or isinstance(val, list):
-                        numerator += val[0]
-                        denominator += val[1]
-                    elif isinstance(val, torch.Tensor) and val.numel() == 2:
-                        # Handle 2-element tensor [loss, num_tokens] format
-                        # (upstream Megatron compatibility)
-                        numerator += val[0]
-                        denominator += val[1]
-                    else:
-                        # legacy behavior. we average over the number of microbatches,
-                        # and so the denominator is 1.
-                        numerator += val
-                        denominator += 1
-                loss_reduced[key] = numerator / denominator
-            return (
-                loss_reduced,
-                skipped_iter,
-                should_checkpoint,
-                should_exit,
-                exit_code,
-                grad_norm,
-                num_zeros_in_grad,
-            )
-        return (
-            {},
-            skipped_iter,
-            should_checkpoint,
-            should_exit,
-            exit_code,
-            grad_norm,
-            num_zeros_in_grad,
-        )
-
-    def training_log(
-        self,
-        loss_dict,
-        total_loss_dict,
-        learning_rate,
-        decoupled_learning_rate,
-        iteration,
-        loss_scale,
-        report_memory_flag,
-        skipped_iter,
-        grad_norm,
-        params_norm,
-        num_zeros_in_grad,
-    ):
-        """Log training information such as losses, timing, ...."""
-        args = get_args()
-        timers = get_timers()
-        writer = get_tensorboard_writer()
-        wandb_writer = get_wandb_writer()
-        mlflow_writer = get_mlflow_writer()
-        get_one_logger()
-
-        # Advanced, skipped, and Nan iterations.
-        advanced_iters_key = "advanced iterations"
-        skipped_iters_key = "skipped iterations"
-        nan_iters_key = "nan iterations"
-        # Advanced iterations.
-        if not skipped_iter:
-            total_loss_dict[advanced_iters_key] = total_loss_dict.get(advanced_iters_key, 0) + 1
-        else:
-            if advanced_iters_key not in total_loss_dict:
-                total_loss_dict[advanced_iters_key] = 0
-        # Skipped iterations.
-        total_loss_dict[skipped_iters_key] = total_loss_dict.get(skipped_iters_key, 0) + skipped_iter
-        # Update losses and set nan iterations
-        got_nan = False
-        for key in loss_dict:
-            if not skipped_iter:
-                total_loss_dict[key] = (
-                    total_loss_dict.get(key, torch.tensor([0.0], dtype=torch.float, device="cuda"))
-                    + loss_dict[key]
-                )
-            else:
-                value = loss_dict[key].float().sum().item()
-                is_nan = value == float("inf") or value == -float("inf") or value != value
-                got_nan = got_nan or is_nan
-        total_loss_dict[nan_iters_key] = total_loss_dict.get(nan_iters_key, 0) + int(got_nan)
-
-        # Logging.
-        timers_to_log = [
-            "forward-backward",
-            "forward-compute",
-            "backward-compute",
-            "batch-generator",
-            "forward-recv",
-            "forward-send",
-            "backward-recv",
-            "backward-send",
-            "forward-send-forward-recv",
-            "forward-send-backward-recv",
-            "backward-send-forward-recv",
-            "backward-send-backward-recv",
-            "forward-backward-send-forward-backward-recv",
-            "layernorm-grads-all-reduce",
-            "embedding-grads-all-reduce",
-            "all-grads-sync",
-            "params-all-gather",
-            "optimizer-copy-to-main-grad",
-            "optimizer-unscale-and-check-inf",
-            "optimizer-clip-main-grad",
-            "optimizer-count-zeros",
-            "optimizer-inner-step",
-            "optimizer-copy-main-to-model-params",
-            "optimizer",
-        ]
-
-        # Calculate batch size.
-        batch_size = args.micro_batch_size * args.data_parallel_size * get_num_microbatches()
-
-        # Track app tag & app tag ID
-        one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length)
-
-        total_iterations = total_loss_dict[advanced_iters_key] + total_loss_dict[skipped_iters_key]
-
-        # learning rate will be None on ranks without trainable params, so we must gather across mp ranks
-        learning_rate = reduce_max_stat_across_model_parallel_group(learning_rate)
-        # Tensorboard values.
-        # Timer requires all the ranks to call.
-        if args.log_timers_to_tensorboard and (iteration % args.tensorboard_log_interval == 0):
-            timers.write(timers_to_log, writer, iteration, normalizer=total_iterations)
-        if iteration % args.tensorboard_log_interval == 0:
-            if wandb_writer:
-                wandb_writer.log({"samples vs steps": args.consumed_train_samples}, iteration)
-            if mlflow_writer:
-                mlflow_writer.log_metric("samples vs steps", args.consumed_train_samples, step=iteration)
-            if writer:
-                writer.add_scalar("learning-rate", learning_rate, iteration)
-                if args.decoupled_lr is not None:
-                    writer.add_scalar("decoupled-learning-rate", decoupled_learning_rate, iteration)
-                writer.add_scalar(
-                    "learning-rate vs samples",
-                    learning_rate,
-                    args.consumed_train_samples,
-                )
-            if wandb_writer:
-                wandb_writer.log({"learning-rate": learning_rate}, iteration)
-            if mlflow_writer:
-                mlflow_writer.log_metric("learning-rate", learning_rate, step=iteration)
-            if writer:
-                writer.add_scalar("batch-size", batch_size, iteration)
-                writer.add_scalar("batch-size vs samples", batch_size, args.consumed_train_samples)
-            if mlflow_writer:
-                mlflow_writer.log_metric("batch-size", batch_size, iteration)
-            if wandb_writer:
-                wandb_writer.log({"batch-size": batch_size}, iteration)
-            for key in loss_dict:
-                if writer:
-                    writer.add_scalar(key, loss_dict[key], iteration)
-                    writer.add_scalar(key + " vs samples", loss_dict[key], args.consumed_train_samples)
-                if wandb_writer:
-                    wandb_writer.log({key: loss_dict[key]}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric(key, loss_dict[key], step=iteration)
-            if args.log_loss_scale_to_tensorboard:
-                if writer:
-                    writer.add_scalar("loss-scale", loss_scale, iteration)
-                    writer.add_scalar("loss-scale vs samples", loss_scale, args.consumed_train_samples)
-                if wandb_writer:
-                    wandb_writer.log({"loss-scale": loss_scale}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("loss-scale", loss_scale, step=iteration)
-            if args.log_world_size_to_tensorboard:
-                if writer:
-                    writer.add_scalar("world-size", args.world_size, iteration)
-                    writer.add_scalar(
-                        "world-size vs samples",
-                        args.world_size,
-                        args.consumed_train_samples,
-                    )
-                if wandb_writer:
-                    wandb_writer.log({"world-size": args.world_size}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("world-size", args.world_size, step=iteration)
-            if grad_norm is not None:
-                if writer:
-                    writer.add_scalar("grad-norm", grad_norm, iteration)
-                    writer.add_scalar("grad-norm vs samples", grad_norm, args.consumed_train_samples)
-                if wandb_writer:
-                    wandb_writer.log({"grad-norm": grad_norm}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("grad-norm", grad_norm, step=iteration)
-            if num_zeros_in_grad is not None:
-                if writer:
-                    writer.add_scalar("num-zeros", num_zeros_in_grad, iteration)
-                    writer.add_scalar(
-                        "num-zeros vs samples",
-                        num_zeros_in_grad,
-                        args.consumed_train_samples,
-                    )
-                if wandb_writer:
-                    wandb_writer.log({"num-zeros": num_zeros_in_grad}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("num-zeros", num_zeros_in_grad, iteration)
-            if params_norm is not None:
-                if writer:
-                    writer.add_scalar("params-norm", params_norm, iteration)
-                    writer.add_scalar(
-                        "params-norm vs samples",
-                        params_norm,
-                        args.consumed_train_samples,
-                    )
-                if wandb_writer:
-                    wandb_writer.log({"params-norm": params_norm}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("params-norm", params_norm, iteration)
-            if args.log_memory_to_tensorboard:
-                mem_stats = torch.cuda.memory_stats()
-                if writer:
-                    writer.add_scalar(
-                        "mem-reserved-bytes",
-                        mem_stats["reserved_bytes.all.current"],
-                        iteration,
-                    )
-                    writer.add_scalar(
-                        "mem-allocated-bytes",
-                        mem_stats["allocated_bytes.all.current"],
-                        iteration,
-                    )
-                    writer.add_scalar(
-                        "mem-max-allocated-bytes",
-                        mem_stats["allocated_bytes.all.peak"],
-                        iteration,
-                    )
-                    writer.add_scalar(
-                        "mem-allocated-count",
-                        mem_stats["allocation.all.current"],
-                        iteration,
-                    )
-                if wandb_writer:
-                    wandb_writer.log(
-                        {"mem-reserved-bytes": mem_stats["reserved_bytes.all.current"]},
-                        iteration,
-                    )
-                    wandb_writer.log(
-                        {"mem-allocated-bytes": mem_stats["allocated_bytes.all.current"]},
-                        iteration,
-                    )
-                    wandb_writer.log(
-                        {"mem-max-allocated-bytes": mem_stats["allocated_bytes.all.peak"]},
-                        iteration,
-                    )
-                    wandb_writer.log(
-                        {"mem-allocated-count": mem_stats["allocation.all.current"]},
-                        iteration,
-                    )
-        if args.num_experts is not None:
-            moe_loss_scale = 1 / get_num_microbatches()
-            track_moe_metrics(
-                loss_scale=moe_loss_scale,
-                iteration=iteration,
-                writer=writer,
-                wandb_writer=wandb_writer,
-                mlflow_writer=mlflow_writer,
-                total_loss_dict=total_loss_dict,
-                per_layer_logging=args.moe_per_layer_logging,
-                moe_layer_freq=args.moe_layer_freq,
-                num_layers=args.num_layers,
-            )
-
-        if iteration % args.log_interval == 0:
-            # Note(wenx): If we want to collect rocm-smi memory information for the first two iterations,
-            # place the collection before the timer to minimize its impact on latency measurements for iterations ≥ 3.
-            rocm_gpu_util = None
-            # Enable throughput calculations if log_throughput or mlflow_upload_performance_metrics is set
-            enable_perf_metrics = args.log_throughput or getattr(
-                args, "mlflow_upload_performance_metrics", False
-            )
-            if enable_perf_metrics:
-                if args.use_rocm_mem_info or (
-                    args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters
-                ):
-                    rocm_total_mem, rocm_used_mem, rocm_free_mem = get_rocm_smi_mem_info(
-                        self.module_local_rank
-                    )
-                # Collect GPU utilization for performance metrics
-                # Every rank samples its own GPU util (rate-limited to avoid rocm-smi
-                # subprocess overhead) so the all_gather below gets real per-rank values.
-                if getattr(args, "mlflow_upload_performance_metrics", False):
-                    import time
-
-                    now = time.time()
-                    should_sample = (
-                        self._last_rocm_gpu_util is None
-                        or (now - self._last_rocm_gpu_util_time) >= self._rocm_gpu_util_sample_interval
-                    )
-                    if should_sample:
-                        try:
-                            self._last_rocm_gpu_util = get_rocm_smi_gpu_util(self.module_local_rank)
-                            self._last_rocm_gpu_util_time = now
-                        except Exception:
-                            pass  # Keep previous cached value if any
-                    rocm_gpu_util = self._last_rocm_gpu_util
-                else:
-                    rocm_gpu_util = None
-
-            elapsed_time = timers("interval-time").elapsed(barrier=True)
-            elapsed_time_per_iteration = elapsed_time / total_iterations
-
-            flops_calc = (
-                num_floating_point_operations
-                if not args.multi_latent_attention
-                else self.num_floating_point_operations_mla_moe
-            )
-            throughput = flops_calc(args, batch_size) / (
-                elapsed_time_per_iteration * 10**12 * args.world_size
-            )
-
-            if args.log_timers_to_tensorboard:
-                if writer:
-                    writer.add_scalar("iteration-time", elapsed_time_per_iteration, iteration)
-                if wandb_writer:
-                    wandb_writer.log({"iteration-time": elapsed_time_per_iteration}, iteration)
-                if mlflow_writer:
-                    mlflow_writer.log_metric("iteration-time", elapsed_time_per_iteration, iteration)
-            # log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
-            log_string = f""
-            if hasattr(self, "episode_count") and self.episode_count is not None:
-                log_string += f" episode {self.episode_count} |"
-            log_string += " iteration {:8d}/{:8d} |".format(iteration, args.train_iters)
-            log_string += " consumed samples: {:12d} |".format(args.consumed_train_samples)
-            if (
-                iteration == self.log_avg_skip_iterations + 1
-                or len(self.recent_iteration_times) >= self.log_avg_reset_interval
-            ):
-                self.recent_iteration_times.clear()
-            self.recent_iteration_times.append(elapsed_time_per_iteration * 1000.0)
-            log_string += " elapsed time per iteration (ms): {:.1f}/{:.1f} |".format(
-                elapsed_time_per_iteration * 1000.0,
-                statistics.mean(self.recent_iteration_times),
-            )
-            if enable_perf_metrics:
-                if (
-                    iteration == self.log_avg_skip_iterations + 1
-                    or len(self.recent_tflop_throughputs) >= self.log_avg_reset_interval
-                ):
-                    self.recent_tflop_throughputs.clear()
-                self.recent_tflop_throughputs.append(throughput)
-
-                if not args.use_rocm_mem_info:
-                    hip_free_mem, hip_total_mem = torch.cuda.mem_get_info()
-                    hip_used_mem = hip_total_mem - hip_free_mem
-                    hip_mem_usage = hip_used_mem / hip_total_mem
-                    log_string += (
-                        f" hip mem usage/free/total/usage_ratio: {hip_used_mem/1024/1024/1024:.2f}GiB/"
-                    )
-                    log_string += f"{hip_free_mem/1024/1024/1024:.2f}GiB/"
-                    log_string += f"{hip_total_mem/1024/1024/1024:.2f}GiB/{hip_mem_usage*100:.2f}% |"
-
-                if args.use_rocm_mem_info or (
-                    args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters
-                ):
-                    rocm_mem_usage = rocm_used_mem / rocm_total_mem
-
-                    # get the max rocm_mem_usage
-                    usage_tensor = torch.tensor([rocm_mem_usage], device="cuda", dtype=torch.float32)
-                    world_size = dist.get_world_size()
-                    gathered_usage = [torch.zeros_like(usage_tensor) for _ in range(world_size)]
-                    dist.all_gather(gathered_usage, usage_tensor)
-
-                    rocm_mem_usages = [t.item() for t in gathered_usage]
-                    max_usage = max(rocm_mem_usages)
-                    max_rank = rocm_mem_usages.index(max_usage)
-
-                    log_string += (
-                        f" rocm mem usage/free/total/usage_ratio: {rocm_used_mem/1024/1024/1024:.2f}GiB/"
-                    )
-                    log_string += f"{rocm_free_mem/1024/1024/1024:.2f}GiB/"
-                    log_string += f"{rocm_total_mem/1024/1024/1024:.2f}GiB/{rocm_mem_usage*100:.2f}% |"
-                    log_string += f" rank-{max_rank} max mem usage/usage_ratio: "
-                    log_string += f"{rocm_total_mem*max_usage/1024/1024/1024:.2f}GiB/{max_usage*100:.2f}% |"
-
-                log_string += (
-                    f" throughput per GPU (TFLOP/s/GPU): {throughput:.1f}/"
-                    f"{statistics.mean(self.recent_tflop_throughputs):.1f} |"
-                )
-                token_throughput = args.seq_length * batch_size / elapsed_time_per_iteration / args.world_size
-                if (
-                    iteration == self.log_avg_skip_iterations + 1
-                    or len(self.recent_token_throughputs) >= self.log_avg_reset_interval
-                ):
-                    self.recent_token_throughputs.clear()
-                self.recent_token_throughputs.append(token_throughput)
-                log_string += (
-                    f" tokens per GPU (tokens/s/GPU): {token_throughput:.1f}/"
-                    f"{statistics.mean(self.recent_token_throughputs):.1f} |"
-                )
-                if args.log_timers_to_tensorboard:
-                    if args.use_rocm_mem_info or (
-                        args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters
-                    ):
-                        mem_collector = "rocm"
-                        used_mem, free_mem, total_mem, mem_usage = (
-                            rocm_used_mem,
-                            rocm_free_mem,
-                            rocm_total_mem,
-                            rocm_mem_usage,
-                        )
-                    else:
-                        mem_collector = "hip"
-                        used_mem, free_mem, total_mem, mem_usage = (
-                            hip_used_mem,
-                            hip_free_mem,
-                            hip_total_mem,
-                            hip_mem_usage,
-                        )
-                    if writer:
-                        writer.add_scalar("throughput(tflops/sec/gpu)", throughput, iteration)
-                        writer.add_scalar(
-                            "token_throughput(tokens/sec/gpu)",
-                            token_throughput,
-                            iteration,
-                        )
-                        writer.add_scalar(
-                            f"{mem_collector}_used_mem(GiB)",
-                            used_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        writer.add_scalar(
-                            f"{mem_collector}_free_mem(GiB)",
-                            free_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        writer.add_scalar(
-                            f"{mem_collector}_total_mem(GiB)",
-                            total_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        writer.add_scalar(f"{mem_collector}_mem_usage(%)", mem_usage * 100.0, iteration)
-                    if wandb_writer:
-                        wandb_writer.log({"throughput(tflops/sec/gpu)": throughput}, iteration)
-                        wandb_writer.log(
-                            {"token_throughput(tokens/sec/gpu)": token_throughput},
-                            iteration,
-                        )
-                        wandb_writer.log(
-                            {f"{mem_collector}_used_mem(GiB)": used_mem / 1024 / 1024 / 1024},
-                            iteration,
-                        )
-                        wandb_writer.log(
-                            {f"{mem_collector}_free_mem(GiB)": free_mem / 1024 / 1024 / 1024},
-                            iteration,
-                        )
-                        wandb_writer.log(
-                            {f"{mem_collector}_total_mem(GiB)": total_mem / 1024 / 1024 / 1024},
-                            iteration,
-                        )
-                        wandb_writer.log({f"{mem_collector}_mem_usage(%)": mem_usage * 100.0}, iteration)
-                    if mlflow_writer:
-                        mlflow_writer.log_metric("throughput_tflops_per_sec_per_gpu", throughput, iteration)
-                        mlflow_writer.log_metric(
-                            "token_throughput_tokens_per_sec_per_gpu",
-                            token_throughput,
-                            iteration,
-                        )
-                        mlflow_writer.log_metric(
-                            f"{mem_collector}_used_mem_GiB",
-                            used_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        mlflow_writer.log_metric(
-                            f"{mem_collector}_free_mem_GiB",
-                            free_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        mlflow_writer.log_metric(
-                            f"{mem_collector}_total_mem_GiB",
-                            total_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        mlflow_writer.log_metric(
-                            f"{mem_collector}_mem_usage_percent", mem_usage * 100.0, iteration
-                        )
-
-                # Upload performance metrics to MLflow
-                # Groups: Performance (throughput, TPS, iteration time), Memory (peak, usage %), System (GPU util)
-                # NOTE: mlflow_writer only exists on last rank, but all_gather requires all ranks to participate
-                if getattr(args, "mlflow_upload_performance_metrics", False):
-                    # Ensure memory metrics are available when log_timers_to_tensorboard is False
-                    # (mem_collector, used_mem, mem_usage are otherwise only set inside log_timers_to_tensorboard)
-                    if not args.log_timers_to_tensorboard:
-                        hip_free_mem, hip_total_mem = torch.cuda.mem_get_info()
-                        used_mem = hip_total_mem - hip_free_mem
-                        mem_usage = used_mem / hip_total_mem
-                        mem_collector = "hip"
-                        if args.use_rocm_mem_info or (
-                            args.use_rocm_mem_info_iters is not None
-                            and iteration in args.use_rocm_mem_info_iters
-                        ):
-                            mem_collector = "rocm"
-                            used_mem = rocm_used_mem
-                            mem_usage = rocm_mem_usage
-                    # System metrics - GPU utilization per rank
-                    # ALL ranks must participate in all_gather, even if they don't have mlflow_writer
-                    # Use -1 as sentinel for unavailable GPU util
-                    util_value = rocm_gpu_util if rocm_gpu_util is not None else -1.0
-                    util_tensor = torch.tensor([util_value], device="cuda", dtype=torch.float32)
-                    world_size = dist.get_world_size()
-                    gathered_utils = [torch.zeros_like(util_tensor) for _ in range(world_size)]
-                    dist.all_gather(gathered_utils, util_tensor)
-
-                    # Only the last rank (which has mlflow_writer) logs the metrics
-                    if mlflow_writer:
-                        # Performance metrics
-                        mlflow_writer.log_metric("perf/throughput_tflops_per_gpu", throughput, iteration)
-                        mlflow_writer.log_metric(
-                            "perf/tps_tokens_per_sec_per_gpu", token_throughput, iteration
-                        )
-                        mlflow_writer.log_metric(
-                            "perf/iteration_time_ms",
-                            elapsed_time_per_iteration * 1000.0,
-                            iteration,
-                        )
-                        # Memory metrics
-                        mlflow_writer.log_metric(
-                            f"perf/{mem_collector}_current_mem_gb",
-                            used_mem / 1024 / 1024 / 1024,
-                            iteration,
-                        )
-                        mlflow_writer.log_metric(
-                            f"perf/{mem_collector}_mem_utilization_pct",
-                            mem_usage * 100.0,
-                            iteration,
-                        )
-                        # Log GPU utilization from gathered values
-                        valid_utils = []
-                        for rank, util_val in enumerate(gathered_utils):
-                            util = util_val.item()
-                            if util >= 0:  # Filter out sentinel values (-1)
-                                mlflow_writer.log_metric(
-                                    f"perf/gpu_utilization_pct_rank{rank}",
-                                    util,
-                                    iteration,
-                                )
-                                valid_utils.append(util)
-                        # Also log average GPU utilization (only from valid values)
-                        if valid_utils:
-                            avg_util = sum(valid_utils) / len(valid_utils)
-                            mlflow_writer.log_metric(
-                                "perf/gpu_utilization_pct_avg",
-                                avg_util,
-                                iteration,
-                            )
-
-            assert learning_rate is not None
-            # Decoupled_learning_rate should be not None only on first and last pipeline stage.
-            log_string += " learning rate: {:.6E} |".format(learning_rate)
-            if args.decoupled_lr is not None and (
-                mpu.is_pipeline_first_stage(ignore_virtual=True)
-                or mpu.is_pipeline_last_stage(ignore_virtual=True)
-            ):
-                assert decoupled_learning_rate is not None
-                log_string += " decoupled learning rate: {:.6E} |".format(decoupled_learning_rate)
-            else:
-                assert decoupled_learning_rate is None
-            log_string += " global batch size: {:5d} |".format(batch_size)
-            for key in total_loss_dict:
-                if key not in [advanced_iters_key, skipped_iters_key, nan_iters_key]:
-                    avg = total_loss_dict[key].item() / float(max(1, total_loss_dict[advanced_iters_key]))
-                    if avg > 0.0:
-                        log_string += " {}: {:.6E} |".format(key, avg)
-                    total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device="cuda")
-            log_string += " loss scale: {:.1f} |".format(loss_scale)
-            if grad_norm is not None:
-                log_string += " grad norm: {:.3f} |".format(grad_norm)
-            if num_zeros_in_grad is not None:
-                log_string += " num zeros: {:.1f} |".format(num_zeros_in_grad)
-            if params_norm is not None:
-                log_string += " params norm: {:.3f} |".format(params_norm)
-            log_string += " number of skipped iterations: {:3d} |".format(total_loss_dict[skipped_iters_key])
-            log_string += " number of nan iterations: {:3d} |".format(total_loss_dict[nan_iters_key])
-            total_loss_dict[advanced_iters_key] = 0
-            total_loss_dict[skipped_iters_key] = 0
-            total_loss_dict[nan_iters_key] = 0
-
-            if self.is_v_schedule:
-                log_rank_0(log_string)
-            else:
-                log_rank_last(log_string)
-            if report_memory_flag and learning_rate > 0.0:
-                # Report memory after optimizer state has been initialized.
-                if torch.distributed.get_rank() == 0:
-                    num_microbatches = get_num_microbatches()
-                    report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True)
-                report_memory("(after {} iterations)".format(iteration))
-                report_memory_flag = False
-
-            # Removed to avoid global sync in zero bubble schedules.
-            if not (get_args().zero_bubble_v_schedule or get_args().patch_zero_bubble):
-                timers.log(timers_to_log, normalizer=args.log_interval)
-
-        return report_memory_flag
-
-    def num_floating_point_operations_mla_moe(self, args, batch_size):
-        # MoE.
-        gated_linear_multiplier = 3 / 2 if args.swiglu else 1
-        num_experts_routed_to = 1 if args.num_experts is None else args.moe_router_topk
-        ffn_hidden_size = (
-            args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None else args.ffn_hidden_size
-        )
-        shared_expert_ffn_hidden_size = (
-            0
-            if args.moe_shared_expert_intermediate_size is None
-            else args.moe_shared_expert_intermediate_size
-        )
-
-        # The 12x term below comes from the following factors; for more details, see
-        # "APPENDIX: FLOATING-POINT OPERATIONS" in https://arxiv.org/abs/2104.04473.
-        # - 3x: Each GEMM in the model needs to be performed 3 times (forward pass,
-        #       backward wgrad [weight gradient], backward dgrad [data gradient]).
-        # - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model
-        #       architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM
-        #       in MLP layer).
-        # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations.
-        expansion_factor_fw_bw = 3
-        expansion_factor_gemm_stack = 2
-        expansion_factor_gemm_ops = 2
-
-        return (
-            expansion_factor_fw_bw
-            * expansion_factor_gemm_ops
-            * batch_size
-            * args.seq_length
-            * args.num_layers
-            * args.hidden_size
-            * args.hidden_size
-            * (
-                # Attention - q.
-                (
-                    (
-                        # Attention - q_proj.
-                        args.num_attention_heads
-                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                        / args.hidden_size
-                    )
-                    if not args.q_lora_rank
-                    else (
-                        # Attention - q_down_proj.
-                        (args.q_lora_rank / args.hidden_size)
-                        # Attention - q_up_proj.
-                        + (
-                            (args.q_lora_rank / args.hidden_size)
-                            * (
-                                args.num_attention_heads
-                                * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                                / args.hidden_size
-                            )
-                        )
-                    )
-                )
-                # Attention - kv_down_proj.
-                + ((args.kv_lora_rank + args.qk_pos_emb_head_dim) / args.hidden_size)
-                # Attention - kv_up_proj.
-                + (
-                    (args.kv_lora_rank / args.hidden_size)
-                    * (args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) / args.hidden_size)
-                )
-                # Attention - Q*K.
-                + (
-                    (args.seq_length / args.hidden_size)
-                    * (
-                        args.num_attention_heads
-                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                        / args.hidden_size
-                    )
-                )
-                # Attention - A(ttention Score)*V.
-                + (
-                    (args.seq_length / args.hidden_size)
-                    * (args.num_attention_heads * args.v_head_dim / args.hidden_size)
-                )
-                # Attention - output proj.
-                + (args.num_attention_heads * args.v_head_dim / args.hidden_size)
-                # MLP.
-                + (
-                    (ffn_hidden_size / args.hidden_size)
-                    * num_experts_routed_to
-                    * gated_linear_multiplier
-                    * expansion_factor_gemm_stack
-                )
-                # Shared Experts.
-                + (
-                    (shared_expert_ffn_hidden_size / args.hidden_size)
-                    * gated_linear_multiplier
-                    * expansion_factor_gemm_stack
-                )
-                # Logit. (untie_embeddings_and_output_weights=True)
-                + (args.padded_vocab_size / (2 * args.num_layers * args.hidden_size))
-                * expansion_factor_gemm_stack
-            )
-        )
diff --git a/primus/modules/trainer/megatron/utils.py b/primus/modules/trainer/megatron/utils.py
deleted file mode 100644
index 75a30885a..000000000
--- a/primus/modules/trainer/megatron/utils.py
+++ /dev/null
@@ -1,584 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-"""megatron utils"""
-
-import inspect
-import json
-import os
-
-import megatron
-import torch
-from megatron.core import parallel_state
-from megatron.training.global_vars import get_args
-
-from primus.core.utils import logger
-
-_GLOBAL_PP_VIS_EVENTS = []
-_GLOBAL_PP_VIS_EVENTS_PER_ITER = None
-
-
-######################################################log after torch distributed initialized
-
-
-def is_v_schedule_enabled(args=None):
-    if args is None:
-        args = get_args()
-    return (
-        args.patch_zero_bubble
-        and args.enable_zero_bubble
-        and (args.zero_bubble_v_schedule or args.enable_1f1b_v)
-    ) or (args.pp_algorithm in ("zbv-formatted", "v-half", "v-min") and args.patch_primus_pipeline)
-
-
-def is_last_rank():
-    return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
-
-
-def print_rank_last(msg):
-    """If distributed is initialized, print only on last rank."""
-    log_func = logger.info_with_caller
-
-    caller = inspect.stack()[1]
-    caller_frame = caller.frame
-    function_name = caller_frame.f_code.co_name
-    module_name = caller_frame.f_globals["__name__"].split(".")[-1]
-    line = caller.lineno
-
-    if torch.distributed.is_initialized():
-        if is_last_rank():
-            log_func(msg, module_name, function_name, line)
-    else:
-        log_func(msg, module_name, function_name, line)
-
-
-def set_wandb_writer_patch(args):  # monkey patch
-    """
-    This function is adapted from the original Megatron implementation, with an additional
-    wandb argument `entity` be added.
-    Monkey-patch note:
-    - The original function will be replaced at runtime by this implementation.
-
-    """
-
-    megatron.training.global_vars._ensure_var_is_not_initialized(
-        megatron.training.global_vars._GLOBAL_WANDB_WRITER, "wandb writer"
-    )
-
-    if getattr(args, "wandb_project", "") and args.rank == (args.world_size - 1):
-        if args.wandb_exp_name == "":
-            raise ValueError("Please specify the wandb experiment name!")
-
-        import wandb
-
-        if args.wandb_save_dir:
-            save_dir = args.wandb_save_dir
-        else:
-            # Defaults to the save dir.
-            save_dir = os.path.join(args.save, "wandb")
-        wandb_kwargs = {
-            "dir": save_dir,
-            "name": args.wandb_exp_name,
-            "project": args.wandb_project,
-            "entity": args.wandb_entity,
-            "config": vars(args),
-        }
-        os.makedirs(wandb_kwargs["dir"], exist_ok=True)
-        wandb.init(**wandb_kwargs)
-        megatron.training.global_vars._GLOBAL_WANDB_WRITER = wandb
-
-
-def validate_specified_recompute_layers(args):
-    if args.recompute_layer_ids is None:
-        return
-
-    assert isinstance(
-        args.recompute_layer_ids, list
-    ), f"recompute_layer_ids={args.recompute_layer_ids} should be a list"
-    recompute_layer_ids = list(set(args.recompute_layer_ids))
-    assert len(recompute_layer_ids) > 0, "recompute layer ids is null"
-    for layer_id in recompute_layer_ids:
-        assert (
-            layer_id >= 0 and layer_id < args.num_layers
-        ), f"recompute layer id must be between 0 and {args.num_layers - 1}"
-
-    if args.recompute_granularity != "full":
-        raise ValueError(
-            f'When using recompute_layer_ids, recompute_granuarlity: {args.recompute_granularity} must be "full"'
-        )
-
-    if args.recompute_method is not None:
-        raise ValueError(
-            f"When using recompute_layer_ids, recompute_method: {args.recompute_method} must be None."
-        )
-
-    if args.distribute_saved_activations and args.sequence_parallel:
-        raise ValueError(
-            f"distribute_saved_activations: {args.distribute_saved_activations} must be "
-            f"false when sequence parallel is enabled: {args.sequence_parallel}"
-        )
-
-
-def validate_manual_split(args):
-    """
-    The use of decoder_pipeline_manual_split_list is to relax the divisibility
-    restriction of the current (interleaved) 1f1b pipeline schedule. The layer
-    split or number of each pp rank is
-    decoder_pipeline_manual_split_list[pp_rank*vp_size:(pp_rank+1)*vp_size] or
-    decoder_pipeline_manual_split_list[pp_rank] when interleaved pipeline is
-    used or not. For example, the split list could be "[2,3,2,2,2,2,2,1]"
-    in layer16-pp4-vpp2 config, where the vpp split of
-    pp_rank0/pp_rank1/pp_rank2/pp_rank3 is [2,3]/[2,2]/[2,2]/[2,1].
-
-    if chosen pipeline is v_schedule like zbv/v-half,
-    the split list will be the actual layer sequence.
-    For example, layer16-pp4-vpp2 config, the vpp split of
-    pp_rank0/pp_rank1/pp_rank2/pp_rank3 is [3,2,2,2,2,2,2,1]
-    indicate the pipeline as follows:
-    pp_rank0: 3       1
-    pp_rank1:  2     2
-    pp_rank2:   2   2
-    pp_rank3:    2 2
-
-    """
-
-    if (
-        args.num_layers_per_virtual_pipeline_stage is not None
-        or args.decoder_first_pipeline_num_layers is not None
-        or args.decoder_last_pipeline_num_layers is not None
-        or args.account_for_embedding_in_pipeline_split
-        or args.account_for_loss_in_pipeline_split
-    ):
-        raise ValueError(
-            "decoder_pipeline_manual_split_list is not compatible "
-            "with num_layers_per_virtual_pipeline_stage/"
-            "decoder_first_pipeline_num_layers/"
-            "decoder_last_pipeline_num_layers/"
-            "account_for_embedding_in_pipeline_split/"
-            "account_for_loss_in_pipeline_split yet"
-        )
-
-    num_layers = args.num_layers
-    pp_size = args.pipeline_model_parallel_size
-    vp_size = args.virtual_pipeline_model_parallel_size
-    pp_split = args.decoder_pipeline_manual_split_list
-
-    if pp_size <= 1:
-        raise ValueError(
-            f"pipeline_model_parallel_size={pp_size} should be larger "
-            f"than 1 when decoder_pipeline_manual_split_list is used"
-        )
-
-    if not isinstance(pp_split, list):
-        raise ValueError(f"decoder_pipeline_manual_split_list={pp_split} should be a list")
-
-    split_size = pp_size if vp_size is None else pp_size * vp_size
-    if len(pp_split) != split_size:
-        raise ValueError(
-            f"the size of decoder_pipeline_manual_split_list="
-            f"{pp_split} should be {split_size} "
-            f"given pipeline_model_parallel_size={pp_size} and "
-            f"virtual_pipeline_model_parallel_size={vp_size}"
-        )
-
-    if not all(x > 0 for x in pp_split):
-        raise ValueError(
-            f"layer numbers in decoder_pipeline_manual_split_list={pp_split} should all be larger than 0"
-        )
-
-    if sum(pp_split) != num_layers:
-        raise ValueError(
-            f"the sum of decoder_pipeline_manual_split_list="
-            f"{pp_split} is {sum(pp_split)} and "
-            f"should be equal to num_layers={num_layers}"
-        )
-
-    return True
-
-
-def validate_args_modified(*args, **kwargs):
-    def validate_args_modifier(func, modification):
-        import inspect
-
-        source = inspect.getsource(func)
-        modified_source = modification(source)
-        namespace = {}
-        exec(modified_source, func.__globals__, namespace)
-        return namespace[func.__name__]
-
-    ori_code = kwargs.pop("ori_code", None)
-    new_code = kwargs.pop("new_code", None)
-
-    assert ori_code is not None and new_code is not None, "ori_code and new_code must be provided."
-
-    megatron.training.arguments.validate_args = validate_args_modifier(
-        megatron.training.arguments.validate_args, lambda s: s.replace(ori_code, new_code)
-    )
-    megatron.training.arguments.validate_args(*args, **kwargs)
-
-
-def set_manual_pipeline_split_patch(args):
-    """
-    Monkey-patch note:
-    - The original function will be replaced at runtime by this implementation.
-
-    """
-
-    megatron.core.transformer.TransformerConfig.decoder_pipeline_manual_split_list = (
-        args.decoder_pipeline_manual_split_list
-    )
-
-    # patch get_num_layers_to_build
-    def get_num_layers_to_build_patch(config, vp_stage, pp_rank=None):
-        if pp_rank is None:
-            pp_rank = parallel_state.get_pipeline_model_parallel_rank()
-        vp_size = config.virtual_pipeline_model_parallel_size
-
-        if not is_v_schedule_enabled():
-            pp_idx = pp_rank if vp_size is None else pp_rank * vp_size + vp_stage
-            num_layers_to_build = config.decoder_pipeline_manual_split_list[pp_idx]
-            return num_layers_to_build
-        else:
-            assert vp_stage is not None and vp_stage in (0, 1)
-            pp_size = config.pipeline_model_parallel_size
-            chunk_id = pp_rank if vp_stage == 0 else 2 * pp_size - pp_rank - 1
-            num_layers_to_build = config.decoder_pipeline_manual_split_list[chunk_id]
-            return num_layers_to_build
-
-    megatron.core.transformer.transformer_block.get_num_layers_to_build = get_num_layers_to_build_patch
-    megatron.core.models.gpt.gpt_layer_specs.get_num_layers_to_build = get_num_layers_to_build_patch
-
-    # patch get_transformer_layer_offset
-    def get_transformer_layer_offset_patch(config, vp_stage, pp_rank=None):
-        if pp_rank is None:
-            pp_rank = parallel_state.get_pipeline_model_parallel_rank()
-        pp_size = config.pipeline_model_parallel_size
-        vp_size = config.virtual_pipeline_model_parallel_size
-
-        offset = 0
-
-        if not is_v_schedule_enabled():
-            if vp_stage is not None:
-                for vp_idx in range(vp_stage):
-                    for pp_idx in range(pp_size):
-                        offset += config.decoder_pipeline_manual_split_list[pp_idx * vp_size + vp_idx]
-                for pp_idx in range(pp_rank):
-                    offset += config.decoder_pipeline_manual_split_list[pp_idx * vp_size + vp_stage]
-            else:
-                offset = sum(config.decoder_pipeline_manual_split_list[:pp_rank])
-        else:
-            assert vp_stage is not None and vp_stage in (0, 1)
-            chunk_id = pp_rank if vp_stage == 0 else 2 * pp_size - pp_rank - 1
-            offset = sum(config.decoder_pipeline_manual_split_list[:chunk_id])
-        return offset
-
-    megatron.core.transformer.transformer_layer.get_transformer_layer_offset = (
-        get_transformer_layer_offset_patch
-    )
-    megatron.core.transformer.transformer_block.get_transformer_layer_offset = (
-        get_transformer_layer_offset_patch
-    )
-    megatron.core.models.gpt.gpt_layer_specs.get_transformer_layer_offset = get_transformer_layer_offset_patch
-
-
-def schedule_wrapper(func):
-    def wrapper(*args, **kwargs):
-        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER = {
-            "start": None,
-            "end": None,
-            "memory": None,
-            "fwd_start": [],
-            "fwd_end": [],
-            "fwd_minibatch": [],
-            "fwd_chunk": [],
-            "bwd_start": [],
-            "bwd_end": [],
-            "bwd_minibatch": [],
-            "bwd_chunk": [],
-            "wgrad_start": [],
-            "wgrad_end": [],
-            "wgrad_minibatch": [],
-            "wgrad_chunk": [],
-        }
-
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"] = torch.cuda.Event(enable_timing=True)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"].record()
-        res = func(*args, **kwargs)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"] = torch.cuda.Event(enable_timing=True)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"].record()
-
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["memory"] = torch.cuda.max_memory_reserved() / 1024**3
-
-        global _GLOBAL_PP_VIS_EVENTS
-        _GLOBAL_PP_VIS_EVENTS.append(_GLOBAL_PP_VIS_EVENTS_PER_ITER)
-
-        return res
-
-    return wrapper
-
-
-def fwd_bwd_wrapper(func, mode, minibatch=None, chunk=None):
-    def wrapper(*args, **kwargs):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        res = func(*args, **kwargs)
-        end.record()
-
-        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_start"].append(start)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_end"].append(end)
-
-        if minibatch is not None:
-            _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_minibatch"].append(minibatch)
-        if chunk is not None:
-            _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_chunk"].append(chunk)
-        return res
-
-    return wrapper
-
-
-def combined_fwd_bwd_wrapper(func, fwd_minibatch, fwd_chunk, bwd_minibatch, bwd_chunk):
-    """Record a single combined forward+backward call as both an ``fwd`` event
-    and a ``bwd`` event sharing the same ``[start, end]`` interval.
-
-    Used by ``megatron_combined_fwd_bkwd_handler`` so that nodes collapsed into
-    a combined FB group still appear in the dump_pp_data output. Without this
-    the visualizer's per-rank F/B/W totals are heavily under-counted on ranks
-    that hit the steady state (the F and B halves are interleaved inside
-    ``combined_forward_backward_step`` and cannot be timed separately).
-    """
-
-    def wrapper(*args, **kwargs):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        res = func(*args, **kwargs)
-        end.record()
-
-        global _GLOBAL_PP_VIS_EVENTS_PER_ITER
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_start"].append(start)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_end"].append(end)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_minibatch"].append(fwd_minibatch)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_chunk"].append(fwd_chunk)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_start"].append(start)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_end"].append(end)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_minibatch"].append(bwd_minibatch)
-        _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_chunk"].append(bwd_chunk)
-        return res
-
-    return wrapper
-
-
-def set_dump_pp_data_patch():
-    from megatron.core.pipeline_parallel import schedules
-
-    schedules.forward_step = fwd_bwd_wrapper(schedules.forward_step, "fwd")
-    schedules.backward_step = fwd_bwd_wrapper(schedules.backward_step, "bwd")
-
-
-def dump_pp_data(args, num_mbs, pp_data_dir):
-    torch.cuda.synchronize()
-
-    global _GLOBAL_PP_VIS_EVENTS
-    all_iter_data = {}
-    for iter_idx, iter_events in enumerate(_GLOBAL_PP_VIS_EVENTS):
-        iter_data = {
-            "total": None,
-            "memory": None,
-            "fwd_start": [],
-            "fwd_end": [],
-            "fwd_minibatch": [],
-            "fwd_chunk": [],
-            "bwd_start": [],
-            "bwd_end": [],
-            "bwd_minibatch": [],
-            "bwd_chunk": [],
-            "wgrad_start": [],
-            "wgrad_end": [],
-            "wgrad_minibatch": [],
-            "wgrad_chunk": [],
-        }
-        iter_data["total"] = iter_events["start"].elapsed_time(iter_events["end"])
-        iter_data["memory"] = iter_events["memory"]
-
-        for i in range(len(iter_events["fwd_start"])):
-            for key in ["fwd_start", "fwd_end", "bwd_start", "bwd_end", "wgrad_start", "wgrad_end"]:
-                if i >= len(iter_events[key]):
-                    continue
-                event_time = iter_events["start"].elapsed_time(iter_events[key][i])
-                iter_data[key].append(event_time)
-            for key in [
-                "fwd_minibatch",
-                "fwd_chunk",
-                "bwd_minibatch",
-                "bwd_chunk",
-                "wgrad_minibatch",
-                "wgrad_chunk",
-            ]:
-                if i >= len(iter_events[key]):
-                    continue
-                iter_data[key].append(iter_events[key][i])
-
-        all_iter_data[iter_idx + 1] = iter_data
-
-    rank = torch.distributed.get_rank()
-    dp_rank = parallel_state.get_data_parallel_rank()
-    pp_rank = parallel_state.get_pipeline_model_parallel_rank()
-    os.makedirs(pp_data_dir, exist_ok=True)
-    if dp_rank == 0:
-        log_path = os.path.join(pp_data_dir, f"pp_rank_{pp_rank}.json")
-        with open(log_path, "w") as f:
-            json.dump(all_iter_data, f, indent=2)
-
-    if rank == 0:
-        vp_size = args.virtual_pipeline_model_parallel_size
-        vp_size = 1 if vp_size is None else vp_size
-        config_dict = {
-            "world_size": args.world_size,
-            "dp_size": args.data_parallel_size,
-            "tp_size": args.tensor_model_parallel_size,
-            "ep_size": args.expert_model_parallel_size,
-            "pp_size": args.pipeline_model_parallel_size,
-            "vp_size": vp_size,
-            "num_mbs": num_mbs,
-            "train_iters": args.train_iters,
-        }
-        log_path = os.path.join(pp_data_dir, f"config.json")
-        with open(log_path, "w") as f:
-            json.dump(config_dict, f, indent=2)
-
-
-def _get_sync_free_moe_options(args) -> dict:
-    stage = args.turbo_sync_free_moe_stage
-
-    if stage > 3 or stage < 0:
-        raise ValueError("turbo_sync_free_moe_stage only support [0-3]")
-
-    sync_free_moe = {
-        1: {
-            "moe_use_fused_router_with_aux_score": True,
-            "moe_permute_fusion": True,
-            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
-        },
-        2: {
-            "moe_use_fused_router_with_aux_score": True,
-            "use_turbo_deepep": True,
-            "moe_permute_fusion": True,
-            "use_turbo_grouped_gemm": True,
-            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
-        },
-        3: {
-            "moe_use_fused_router_with_aux_score": True,
-            "use_turbo_deepep": True,
-            "moe_permute_fusion": True,
-            "use_turbo_grouped_gemm": True,
-            "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False,
-            "use_turbo_fused_act_with_probs": True,
-        },
-    }
-
-    return sync_free_moe[stage]
-
-
-def validate_args_on_rocm(args):
-    # Deterministic mode
-    if args.deterministic_mode:
-        # NOTE: Some environment variables affect deterministic mode on ROCm. Need to do extra check.
-        NON_DETERMINISTIC_ENVS = {
-            "TORCH_COMPILE_DISABLE": "1",
-            "ROCBLAS_DEFAULT_ATOMICS_MODE": "0",
-            "PRIMUS_TURBO_AUTO_TUNE": "0",
-            "PRIMUS_DETERMINISTIC": "1",
-        }
-        # NOTE: Some version triton compile exist potential racing condition issue.
-        for env, value in NON_DETERMINISTIC_ENVS.items():
-            assert (
-                os.environ.get(env, None) == value
-            ), f"{env} must be set to {value} in deterministic mode but got {os.environ.get(env, None)} instead."
-
-        # Set fill_uninitialized_memory to False to avoid calling extra fill kernel in deterministic mode.
-        torch.utils.deterministic.fill_uninitialized_memory = False
-
-    assert not getattr(
-        args, "use_turbo_parallel_linear", False
-    ), "use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead."
-    use_turbo_gemm = getattr(args, "use_turbo_gemm", False)
-    # Turbo FP8 linear check
-    if args.fp8 and use_turbo_gemm:
-        support_fp8_recipe = ["tensorwise", "blockwise", "mxfp8"]
-        assert (
-            args.fp8_recipe in support_fp8_recipe
-        ), f"{args.fp8_recipe} recipe is not support when enable `use_turbo_gemm`."
-
-    # Turbo FP4 linear check
-    if args.fp4 and use_turbo_gemm:
-        support_fp4_recipe = ["mxfp4"]
-        assert (
-            args.fp4_recipe in support_fp4_recipe
-        ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`."
-
-    # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm.
-    if args.fp8_recipe == "mxfp8":
-        assert (
-            os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "1"
-        ), "Please set `NVTE_ROCM_ENABLE_MXFP8=1` to enable `mxfp8` recipe."
-
-    # dump pp data
-    if args.dump_pp_data and args.pipeline_model_parallel_size == 1:
-        args.dump_pp_data = False
-        print_rank_last(f"Disable args.dump_pp_data since args.pipeline_model_parallel_size=1")
-
-    # PrimusTurboGroupedMLP no longer depends on legacy GroupedMLP; the two
-    # flags are mutually exclusive when turbo is enabled.
-    assert not getattr(
-        args, "use_turbo_grouped_mlp", False
-    ), "use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead."
-    use_turbo_grouped_gemm = getattr(args, "use_turbo_grouped_gemm", False)
-    if use_turbo_grouped_gemm:
-        if getattr(args, "moe_use_legacy_grouped_gemm", False):
-            raise ValueError(
-                "use_turbo_grouped_gemm=True is incompatible with moe_use_legacy_grouped_gemm=True. "
-                "please set moe_use_legacy_grouped_gemm=False."
-            )
-
-    # sync-free MoE
-    if args.turbo_sync_free_moe_stage > 0:
-        assert args.enable_primus_turbo, "Please set `enable_primus_turbo=True` to enable sync-free MoE."
-
-        if args.turbo_sync_free_moe_stage > 1 and not use_turbo_grouped_gemm:
-            raise ValueError(
-                "Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedLinear, please set `use_turbo_grouped_gemm=True`"
-            )
-        options = _get_sync_free_moe_options(args)
-        print_rank_last(
-            f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) =========="
-        )
-        for flag, value in options.items():
-            dots = "." * (73 - len(flag) - len(str(value)))
-            print_rank_last(f"{flag}{dots}{value}")
-            setattr(args, flag, value)
-        print_rank_last(
-            f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) =========="
-        )
-
-    # turbo deepep
-    if args.use_turbo_deepep:
-        assert (
-            not args.moe_shared_expert_overlap
-        ), "DeepEP not support moe_shared_expert_overlap, please set `moe_shared_expert_overlap=False`."
-        assert (
-            args.moe_router_dtype == "fp32"
-        ), "DeepEP only supports float32 probs, please set `moe_router_dtype=fp32`"
-        if (
-            args.expert_model_parallel_size >= 16
-            and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO"
-        ):
-            # Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine.
-            assert args.turbo_deepep_num_cu <= 32, "Set `turbo_deepep_num_cu<=32` when using ep_size >= 16."
diff --git a/primus/modules/trainer/torchtitan/parse_utils.py b/primus/modules/trainer/torchtitan/parse_utils.py
deleted file mode 100644
index 17efaaaa2..000000000
--- a/primus/modules/trainer/torchtitan/parse_utils.py
+++ /dev/null
@@ -1,66 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-import argparse
-from argparse import Namespace
-from typing import List
-
-
-def _parse_args(ignore_unknown_args: bool = False) -> tuple[Namespace, List[str]]:
-    parser = argparse.ArgumentParser(description="Primus Arguments", allow_abbrev=False)
-    parser.add_argument(
-        "--exp",
-        type=str,
-        required=True,
-        help="Primus experiment yaml config file.",
-    )
-    return parser.parse_known_args() if ignore_unknown_args else (parser.parse_args(), [])
-
-
-def _merge_args(args: Namespace, unknown_args: List[str]) -> Namespace:
-    merged_dict = vars(args).copy()
-    temp_parser = argparse.ArgumentParser()
-
-    i = 0
-    while i < len(unknown_args):
-        key = unknown_args[i]
-        if key.startswith("--"):
-            if i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith("--"):
-                temp_parser.add_argument(key, type=str)
-                i += 2
-            else:
-                temp_parser.add_argument(key, action="store_true")
-                i += 1
-        else:
-            i += 1
-
-    parsed_unknown, _ = temp_parser.parse_known_args(unknown_args)
-    merged_dict.update(vars(parsed_unknown))
-    return Namespace(**merged_dict)
-
-
-def _convert_args_to_cli(args: Namespace) -> List[str]:
-    cli_args = []
-
-    if hasattr(args, "exp"):
-        cli_args += ["--job.config_file", args.exp]
-
-    for k, v in vars(args).items():
-        if k in {"exp", "backend"}:
-            continue
-        if isinstance(v, bool):
-            if v:
-                cli_args.append(f"--{k}")
-        else:
-            cli_args += [f"--{k}", str(v)]
-
-    return cli_args
-
-
-def get_torchtitan_config_args() -> List[str]:
-    args, unknown_args = _parse_args(ignore_unknown_args=True)
-    merged = _merge_args(args, unknown_args)
-    return _convert_args_to_cli(merged)
diff --git a/primus/modules/trainer/torchtitan/pre_trainer.py b/primus/modules/trainer/torchtitan/pre_trainer.py
deleted file mode 100644
index c178d853b..000000000
--- a/primus/modules/trainer/torchtitan/pre_trainer.py
+++ /dev/null
@@ -1,297 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-from dataclasses import asdict, is_dataclass
-from types import SimpleNamespace
-from typing import Any, Dict
-
-# Trigger registration of all TorchTitan patches
-import primus.backends.torchtitan.patches  # noqa: F401
-from primus.core.patches import run_patches
-from primus.core.utils.yaml_utils import nested_namespace_to_dict
-from primus.modules.base_module import BaseModule
-
-
-class TorchTitanPretrainTrainer(BaseModule):
-    def __init__(self, *args, **kwargs):
-        extra_args = kwargs.pop("extra_args", None)
-        super().__init__(*args, **kwargs)
-
-        self.primus_cfg = kwargs.pop("primus_config", None)
-        if self.primus_cfg is None:
-            raise ValueError("primus_config is required")
-
-        pre_trainer_cfg = self.primus_cfg.get_module_config("pre_trainer")
-        cfg_dict = nested_namespace_to_dict(pre_trainer_cfg)
-
-        # Run before_train phase patches (e.g., Primus-Turbo patches)
-        # Construct a temporary module_config for patch context
-        temp_module_config = SimpleNamespace()
-        temp_module_config.params = pre_trainer_cfg
-        temp_module_config.name = "pre_trainer"
-        temp_module_config.framework = "torchtitan"
-        temp_module_config.model = getattr(pre_trainer_cfg, "model.name", None)
-
-        # Merge extra_args into temp_module_config.params if provided
-        if extra_args:
-            temp_module_config.params = self._merge_extra_args_into_params(
-                temp_module_config.params, extra_args
-            )
-
-        run_patches(
-            backend="torchtitan",
-            phase="setup",
-            backend_version="unknown",
-            extra={
-                "module_config": temp_module_config,
-            },
-        )
-
-        from torchtitan.config.job_config import JobConfig
-        from torchtitan.train import Trainer
-
-        self.TrainerClass = Trainer
-        self.JobConfigClass = JobConfig
-
-        self.titan_config = self.build_job_config(cfg_dict, self.JobConfigClass)
-
-        self.log_config(self.titan_config)
-        self.trainer = None
-
-    def _merge_extra_args_into_params(
-        self, params: SimpleNamespace, extra_args: Dict[str, Any]
-    ) -> SimpleNamespace:
-        """
-        Merge extra_args into params with specific rules:
-
-        1. Non-model.* parameters: Only merge if they already exist in params
-        2. model.* parameters: Extract and merge (can add new fields)
-
-        Args:
-            params: Original params namespace
-            extra_args: Additional arguments to merge
-
-        Returns:
-            Updated params with merged extra_args
-        """
-        from primus.modules.module_utils import log_rank_0
-
-        # First, flatten extra_args if it contains nested dicts
-        flattened_extra_args = self._flatten_dict(extra_args)
-
-        # Separate model.* and non-model.* parameters
-        model_params = {}
-        non_model_params = {}
-
-        for key, value in flattened_extra_args.items():
-            if key.startswith("model."):
-                # Extract model parameter (remove "model." prefix)
-                model_key = key[6:]  # Remove "model." prefix
-                model_params[model_key] = value
-            else:
-                non_model_params[key] = value
-
-        # Convert params to dict for easier manipulation
-        params_dict = nested_namespace_to_dict(params)
-
-        # Merge non-model.* parameters (only if they already exist)
-        for key, value in non_model_params.items():
-            if self._key_exists_in_dict(params_dict, key):
-                self._set_nested_dict_value(params_dict, key, value)
-                log_rank_0(f"[ExtraArgs] Merged non-model param: {key} = {value}")
-            else:
-                log_rank_0(f"[ExtraArgs] Skipped non-model param (not in params): {key} = {value}")
-
-        # Merge model.* parameters (can add new fields)
-        if model_params:
-            if "model" not in params_dict:
-                params_dict["model"] = {}
-
-            for key, value in model_params.items():
-                self._set_nested_dict_value(params_dict, f"model.{key}", value)
-                log_rank_0(f"[ExtraArgs] Merged model param: model.{key} = {value}")
-
-        # Convert back to SimpleNamespace
-        return self._dict_to_namespace(params_dict)
-
-    def _flatten_dict(self, d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
-        """Flatten nested dictionary using dot notation."""
-        result = {}
-
-        for key, value in d.items():
-            full_key = f"{prefix}.{key}" if prefix else key
-
-            if isinstance(value, dict):
-                # Recursively flatten nested dict
-                result.update(self._flatten_dict(value, full_key))
-            else:
-                # Leaf value
-                result[full_key] = value
-
-        return result
-
-    def _key_exists_in_dict(self, d: Dict[str, Any], key: str) -> bool:
-        """Check if a nested key exists in dictionary (supports dot notation)."""
-        parts = key.split(".")
-        current = d
-
-        for part in parts:
-            if not isinstance(current, dict) or part not in current:
-                return False
-            current = current[part]
-
-        return True
-
-    def _set_nested_dict_value(self, d: Dict[str, Any], key: str, value: Any) -> None:
-        """Set a nested dictionary value using dot notation."""
-        parts = key.split(".")
-        current = d
-
-        for part in parts[:-1]:
-            if part not in current:
-                current[part] = {}
-            current = current[part]
-
-        current[parts[-1]] = value
-
-    def _dict_to_namespace(self, d: Dict[str, Any]) -> SimpleNamespace:
-        """Recursively convert dictionary to SimpleNamespace."""
-        if isinstance(d, dict):
-            return SimpleNamespace(**{k: self._dict_to_namespace(v) for k, v in d.items()})
-        elif isinstance(d, list):
-            return [self._dict_to_namespace(item) for item in d]
-        else:
-            return d
-
-    def setup(self):
-        pass
-
-    def init(self, *init_args, **kwargs):
-        self.trainer = self.TrainerClass(self.titan_config)
-
-    def run(self, *args, **kwargs):
-        if self.trainer is None:
-            raise RuntimeError("Trainer has not been initialized. Call init() first.")
-        self.trainer.train()
-
-    def flatten_config(self, obj: Any, prefix: str = "") -> Dict[str, Any]:
-        flat_dict = {}
-        if is_dataclass(obj):
-            obj = asdict(obj)
-
-        if isinstance(obj, dict):
-            for key, value in obj.items():
-                full_key = f"{prefix}.{key}" if prefix else key
-                if is_dataclass(value) or isinstance(value, dict):
-                    flat_dict.update(self.flatten_config(value, full_key))
-                else:
-                    flat_dict[full_key] = value
-        else:
-            flat_dict[prefix] = obj
-
-        return flat_dict
-
-    def log_config(self, obj: Any, header: str = "TorchTitan Config"):
-        from torchtitan.tools.logging import logger
-
-        logger.info("========== %s ==========" % header)
-        flat = self.flatten_config(obj)
-        max_key_len = max(len(k) for k in flat.keys())
-        for key in sorted(flat):
-            val = flat[key]
-            formatted_line = f"arguments {key.ljust(max_key_len, '.')} {val}"
-            logger.info(formatted_line)
-
-    def build_job_config(self, cfg_dict: dict, JobConfigType) -> Any:
-        import importlib
-
-        from torchtitan.config.job_config import Experimental
-        from torchtitan.tools.logging import logger
-
-        # Step 1: Parse the experimental section to check for a custom JobConfig extension
-        experimental_cfg = cfg_dict.get("experimental", {})
-        experimental = Experimental(**experimental_cfg)
-
-        # Step 2: If a custom_args_module is defined, import and merge with JobConfig
-        custom_job_config_cls = JobConfigType
-        if experimental and getattr(experimental, "custom_args_module", None):
-            try:
-                module = importlib.import_module(experimental.custom_args_module)
-                ExtendedJobConfig = getattr(module, "JobConfig")
-                custom_job_config_cls = self.merge_configs(JobConfigType, ExtendedJobConfig)
-                logger.info(f"Loaded and merged custom JobConfig from {experimental.custom_args_module}")
-            except Exception as e:
-                logger.warning(f"Failed to load custom_args_module '{experimental.custom_args_module}': {e}")
-
-        # Step 3: Parse config dict (including custom fields) into dataclass recursively
-        return self._dict_to_dataclass(custom_job_config_cls, cfg_dict)
-
-    @staticmethod
-    def merge_configs(base_cls, custom_cls):
-        """
-        Merges two dataclass types into one unified dataclass.
-
-        Merge logic:
-        - If a field exists in both:
-            - If both fields are dataclasses, recursively merge them.
-            - Otherwise, the custom field overrides the base.
-        - Fields only in base or only in custom are included as-is.
-        """
-        from dataclasses import field, fields, make_dataclass
-
-        base_fields = {f.name: f for f in fields(base_cls)}
-        custom_fields = {f.name: f for f in fields(custom_cls)}
-
-        merged = []
-
-        # Merge overlapping and base-only fields
-        for name, base_f in base_fields.items():
-            if name in custom_fields:
-                custom_f = custom_fields[name]
-                if is_dataclass(base_f.type) and is_dataclass(custom_f.type):
-                    merged_type = TorchTitanPretrainTrainer.merge_configs(base_f.type, custom_f.type)
-                    merged.append((name, merged_type, field(default_factory=merged_type)))
-                else:
-                    merged.append((name, custom_f.type, custom_f))
-            else:
-                merged.append((name, base_f.type, base_f))
-
-        # Add custom-only fields
-        for name, custom_f in custom_fields.items():
-            if name not in base_fields:
-                merged.append((name, custom_f.type, custom_f))
-
-        return make_dataclass(f"Merged{base_cls.__name__}", merged, bases=(base_cls,))
-
-    def _dict_to_dataclass(self, cls, data: dict[str, Any]) -> Any:
-        """Recursively convert dictionary to dataclass, handling nested and custom fields."""
-        from dataclasses import fields, is_dataclass
-
-        if not is_dataclass(cls):
-            return data
-
-        # collect valid field names
-        field_names = {f.name for f in fields(cls)}
-        init_values = {}
-
-        # only use known fields for constructor
-        for f in fields(cls):
-            if f.name in data:
-                val = data[f.name]
-                if is_dataclass(f.type) and isinstance(val, dict):
-                    init_values[f.name] = self._dict_to_dataclass(f.type, val)
-                else:
-                    init_values[f.name] = val
-
-        # instantiate dataclass
-        obj = cls(**init_values)
-
-        # attach unknown fields dynamically
-        for k, v in data.items():
-            if k not in field_names:
-                setattr(obj, k, v)
-
-        return obj
diff --git a/primus/pretrain.py b/primus/pretrain.py
index 9e8dbb26f..e4b4ca0e8 100644
--- a/primus/pretrain.py
+++ b/primus/pretrain.py
@@ -4,14 +4,19 @@
 # See LICENSE for license information.
 ###############################################################################
 
-import argparse
+"""Backend-path / environment helpers shared by the training and projection
+entry points.
+
+Training is driven entirely by the core runtime
+(:mod:`primus.core.runtime.train_runtime`). Only the backend-path resolution
+utilities remain here, still used by the projection subcommand, runner hooks
+and examples.
+"""
+
 import os
 import sys
 from pathlib import Path
 
-from primus.core.launcher.config import PrimusConfig
-from primus.core.launcher.parser import add_pretrain_parser, load_primus_config
-
 
 def _info_enabled() -> bool:
     """True when PRIMUS_LOG_LEVEL permits INFO-level chatter (DEBUG/INFO).
@@ -22,65 +27,6 @@ def _info_enabled() -> bool:
     return os.environ.get("PRIMUS_LOG_LEVEL", "INFO").upper() in ("DEBUG", "INFO")
 
 
-# Lazy backend loader
-def load_backend_trainer(framework: str):
-    if framework == "megatron":
-        import megatron.training.training as training
-        import torch
-
-        _original_build_model = training.get_model
-
-        def _patched_get_model(*args, **kwargs):
-            """
-            Monkey-patched version of build_model that removes the second
-            DDP construction inside torch.cuda.stream() block.
-            """
-            import inspect
-
-            from megatron.training import training as tr
-
-            inspect.getsource(tr.get_model)
-            print("[PrimusPatch] Overriding build_model to disable second DDP construction...")
-
-            _orig_stream_ctx = torch.cuda.stream
-
-            def _noop_stream(*args, **kwargs):
-                class DummyCtx:
-                    def __enter__(self):
-                        return None
-
-                    def __exit__(self, *a):
-                        return False
-
-                return DummyCtx()
-
-            torch.cuda.stream = _noop_stream
-
-            try:
-                return _original_build_model(*args, **kwargs)
-            finally:
-                torch.cuda.stream = _orig_stream_ctx
-
-        training.get_model = _patched_get_model
-        print("[PrimusPatch] Applied Megatron build_model monkey-patch to disable second DDP.")
-
-        from primus.modules.trainer.megatron.pre_trainer import MegatronPretrainTrainer
-
-        return MegatronPretrainTrainer
-    elif framework == "torchtitan":
-        from primus.modules.trainer.torchtitan.pre_trainer import (
-            TorchTitanPretrainTrainer,
-        )
-
-        return TorchTitanPretrainTrainer
-    elif framework == "maxtext":
-        from primus.modules.trainer.maxtext.pre_trainer import MaxTextPretrainTrainer
-
-        return MaxTextPretrainTrainer
-    else:
-        raise ValueError(f"Unsupported framework: {framework}")
-
-
 def setup_backend_path(framework: str, backend_path=None, verbose: bool = True):
     """
     Setup Python path for backend modules.
@@ -159,83 +105,3 @@ def setup_env(data_path: str):
     else:
         hf_home = os.environ["HF_HOME"]
         print(f"[Primus CLI] HF_HOME already set: {hf_home}")
-
-
-def launch_pretrain_trainer(primus_cfg: PrimusConfig, extra_args=None):
-    """
-    Launch the training using the Primus trainer.
-
-    Args:
-        primus_cfg (PrimusConfig): Parsed Primus configuration object.
-    """
-    # Get pre_trainer module configuration
-    pre_trainer_cfg = primus_cfg.get_module_config("pre_trainer")
-    framework = pre_trainer_cfg.framework
-
-    # Lazy import backend trainer
-    TrainerClass = load_backend_trainer(framework)
-
-    master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
-    master_port = int(os.getenv("MASTER_PORT", "29500"))
-
-    if framework == "maxtext":
-        rank = int(os.getenv("NODE_RANK", "0"))
-        world_size = int(os.getenv("NNODES", "1"))
-    else:
-        # envs set by torchrun
-        rank = int(os.getenv("RANK", "0"))
-        world_size = int(os.getenv("WORLD_SIZE", "1"))
-
-    # Initialize trainer
-    trainer = TrainerClass(
-        module_name="pre_trainer",
-        primus_config=primus_cfg,
-        module_rank=rank,
-        module_world_size=world_size,
-        module_master_addr=master_addr,
-        module_master_port=master_port,
-        extra_args=extra_args,
-    )
-
-    # Launch training
-    trainer.init()
-    trainer.run()
-
-
-def launch_pretrain_from_cli(args, overrides):
-    """
-    Entry point for the 'train' subcommand.
-
-     Steps:
-        1. Load and parse the experiment YAML config
-        2. Merge CLI overrides into the config
-        3. Optionally export the merged config
-        4. Setup backend path
-        5. Launch the training
-    """
-    cfg_path = Path(args.config)
-    if not cfg_path.exists():
-        raise FileNotFoundError(f"[Primus:Train] Config file '{cfg_path}' not found.")
-
-    setup_env(data_path=args.data_path)
-
-    primus_cfg, unknown_overrides = load_primus_config(args, overrides)
-
-    # Export merged config if requested
-    if args.export_config:
-        primus_cfg.export(export_path=args.export_config)
-
-    # Setup backend path for dynamic import
-    framework = primus_cfg.get_module_config("pre_trainer").framework
-    setup_backend_path(framework=framework, backend_path=args.backend_path, verbose=True)
-
-    launch_pretrain_trainer(primus_cfg=primus_cfg, extra_args=unknown_overrides)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="pretrain")
-    add_pretrain_parser(parser)
-
-    args, unknown_args = parser.parse_known_args()
-
-    launch_pretrain_from_cli(args, unknown_args)
diff --git a/primus/tools/preflight/README.md b/primus/tools/preflight/README.md
index d7ca011f0..ff7ae7a8a 100644
--- a/primus/tools/preflight/README.md
+++ b/primus/tools/preflight/README.md
@@ -9,26 +9,32 @@ Torch / single node:
 
 ```bash
 primus-cli preflight \
-  --dump-path output/preflight \
-  --report-file-name preflight_report
+  --dump-path output/preflight
 ```
 
 Slurm (multi-node example):
 
 ```bash
 NUM_NODES=8 srun -N ${NUM_NODES} --ntasks-per-node=1 --cpus-per-task=256 \
-  primus-cli preflight --dump-path output/preflight --report-file-name preflight_report
+  primus-cli preflight --dump-path output/preflight
 ```
 
+If you omit `--report-file-name`, preflight auto-generates a unique
+timestamped basename of the form `preflight-${NNODES}N-YYYYMMDD-HHMMSS`
+so each run writes to a fresh path and never overwrites prior output.
+Pass `--report-file-name NAME` only when you want a stable, well-known
+filename.
+
 
 ## 📂 Output Directory
 
 After running **Preflight**, all test results and reports are generated under the `output/preflight` directory.
 
-The final reports are:
+The final reports (basename shown here is the auto-generated default; it
+reflects whatever `--report-file-name` resolves to) are:
 
-- `preflight_report.md` – a Markdown version of the test report
-- `preflight_report.pdf` – a PDF version of the same report
+- `<report-name>.md` – a Markdown version of the test report
+- `<report-name>.pdf` – a PDF version of the same report
 
 These reports summarize GPU performance, intra-node and inter-node communication results, and help identify potential issues within the cluster.
 
@@ -40,8 +46,8 @@ These reports summarize GPU performance, intra-node and inter-node communication
 output/preflight
 ├── inter_node_comm
 ├── intra_node_comm
-├── preflight_report.md
-├── preflight_report.pdf
+├── preflight-8N-20260715-142530.md
+├── preflight-8N-20260715-142530.pdf
 ├── square_gemm_tflops
 └── ...
 ```
diff --git a/primus/tools/preflight/global_vars.py b/primus/tools/preflight/global_vars.py
index a70f5a3d8..a99b74cec 100644
--- a/primus/tools/preflight/global_vars.py
+++ b/primus/tools/preflight/global_vars.py
@@ -13,9 +13,36 @@
 MASTER_ADDR = os.environ.get("MASTER_ADDR", "127.0.0.1")
 MASTER_PORT = os.environ.get("MASTER_PORT", "29500")
 
+# Default warmup / iteration counts. These remain importable as module-level
+# constants for backward compatibility, but new code should call
+# `get_warmup()` / `get_iteration()` so that `--quick` (or any future override)
+# can lower the counts at runtime.
 WARMUP = 10
 ITERATION = 50
 
+# Internal mutable counters used by the accessors below.
+_WARMUP = WARMUP
+_ITERATION = ITERATION
+
+
+def set_warmup(value: int) -> None:
+    global _WARMUP
+    _WARMUP = int(value)
+
+
+def set_iteration(value: int) -> None:
+    global _ITERATION
+    _ITERATION = int(value)
+
+
+def get_warmup() -> int:
+    return _WARMUP
+
+
+def get_iteration() -> int:
+    return _ITERATION
+
+
 _HOST_NAMES = [None]
 
 
diff --git a/primus/tools/preflight/gpu/gpu_basic.py b/primus/tools/preflight/gpu/gpu_basic.py
index 99f07b32d..54a79e448 100644
--- a/primus/tools/preflight/gpu/gpu_basic.py
+++ b/primus/tools/preflight/gpu/gpu_basic.py
@@ -178,14 +178,19 @@ def run_gpu_basic_checks(
     else:
         findings.append(Finding("info", "GPU occupancy", {"note": "amd-smi JSON not available; skipped"}))
 
-    # ROCm runtime availability: best-effort presence via tooling.
+    # ROCm runtime availability: sysfs is always available; amd-smi/rocm-smi
+    # are supplementary and only probed on LOCAL_RANK 0.
     if probe.backend == "rocm":
-        if "amd-smi" in probe.tooling or "rocm-smi" in probe.tooling:
-            findings.append(
-                Finding("info", "ROCm runtime/tooling detected", {"tooling": list(probe.tooling.keys())})
-            )
+        detected = [k for k in ("sysfs", "amd-smi", "rocm-smi") if k in probe.tooling]
+        if detected:
+            findings.append(Finding("info", "ROCm runtime/tooling detected", {"tooling": detected}))
         else:
-            findings.append(Finding("warn", "ROCm tooling not found (amd-smi/rocm-smi)", {}))
+            from primus.tools.preflight.global_vars import LOCAL_RANK
+
+            if LOCAL_RANK != 0:
+                findings.append(Finding("info", "ROCm tooling skipped (non-zero LOCAL_RANK)", {}))
+            else:
+                findings.append(Finding("warn", "ROCm tooling not found (sysfs/amd-smi/rocm-smi)", {}))
 
     ok = not any(f.level == "fail" for f in findings)
     return {"ok": ok, "probe": probe, "findings": findings}
diff --git a/primus/tools/preflight/gpu/gpu_probe.py b/primus/tools/preflight/gpu/gpu_probe.py
index 5238c65ac..45bfbfb3f 100644
--- a/primus/tools/preflight/gpu/gpu_probe.py
+++ b/primus/tools/preflight/gpu/gpu_probe.py
@@ -7,10 +7,14 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any, Dict, List, Optional
 
+from .sysfs_probe import sysfs_probe
 from .utils import ProbeResult, run_cmd, which
 
+logger = logging.getLogger(__name__)
+
 
 def _normalize_gfx_arch(raw: str) -> str:
     """
@@ -120,27 +124,19 @@ def _probe_with_torch() -> Dict[str, Any]:
 
 
 def _probe_amd_smi() -> Optional[Dict[str, Any]]:
+    """Best-effort amd-smi JSON probe for process occupancy detection only."""
     if which("amd-smi") is None:
         return None
-
-    # Prefer JSON if available (newer amd-smi); fall back to `list` output.
-    rc, out, err = run_cmd(["amd-smi", "list", "--json"], timeout_s=10)
-    if rc == 0 and out:
-        try:
+    try:
+        rc, out, err = run_cmd(["amd-smi", "list", "--json"], timeout_s=10)
+        if rc == 0 and out:
             return {"rc": rc, "json": json.loads(out), "err": err}
-        except Exception:
-            # Fall through to non-json.
-            pass
-
-    rc, out, err = run_cmd(["amd-smi", "list"], timeout_s=10)
-    return {"rc": rc, "out": out, "err": err}
+    except Exception as e:
+        logger.debug("amd-smi probe failed: %s", e)
+    return None
 
 
-def _probe_rocm_smi() -> Optional[Dict[str, Any]]:
-    if which("rocm-smi") is None:
-        return None
-    rc, out, err = run_cmd(["rocm-smi", "-a"], timeout_s=10)
-    return {"rc": rc, "out": out, "err": err}
+_PROBE_CACHE: Optional[ProbeResult] = None
 
 
 def probe_gpus() -> ProbeResult:
@@ -149,21 +145,54 @@ def probe_gpus() -> ProbeResult:
 
     Returns a normalized structure; individual fields may be missing depending
     on the environment/tooling availability.
+
+    Results are cached since the probe is called multiple times per rank
+    (from gpu_basic, gpu_topology, and gpu_perf) and the output is static
+    during a single preflight run.
+
+    Primary GPU enumeration uses sysfs (KFD topology), which is safe to call
+    from any rank (no subprocesses, no /dev/shm mutex).  amd-smi is kept as
+    an optional probe on LOCAL_RANK == 0 for process-occupancy data only;
+    its failure never crashes the run.
     """
+    global _PROBE_CACHE
+    if _PROBE_CACHE is not None:
+        return _PROBE_CACHE
+
+    from primus.tools.preflight.global_vars import LOCAL_RANK
+
     torch_info = _probe_with_torch()
     tooling: Dict[str, Any] = {"torch": torch_info}
 
-    amd = _probe_amd_smi()
-    if amd is not None:
-        tooling["amd-smi"] = amd
-    rocmsmi = _probe_rocm_smi()
-    if rocmsmi is not None:
-        tooling["rocm-smi"] = rocmsmi
+    # sysfs probe: safe from every rank, no subprocess calls.
+    sysfs_result = sysfs_probe()
+    if sysfs_result.ok:
+        tooling["sysfs"] = {
+            "gpu_count": sysfs_result.gpu_count,
+            "gpus": [
+                {
+                    "index": i,
+                    "pci_bdf": g.pci_bdf_str,
+                    "unique_id": hex(g.unique_id) if g.unique_id else None,
+                    "numa_node": g.numa_node,
+                }
+                for i, g in enumerate(sysfs_result.gpus)
+            ],
+            "link_count": len(sysfs_result.links),
+        }
+    else:
+        logger.debug("sysfs probe unavailable: %s", sysfs_result.error)
+
+    # amd-smi JSON: subprocess-based, LOCAL_RANK 0 only, best-effort.
+    # Used solely for process-occupancy detection in gpu_basic.py.
+    if LOCAL_RANK == 0:
+        amd = _probe_amd_smi()
+        if amd is not None:
+            tooling["amd-smi"] = amd
 
     backend = torch_info.get("backend") if torch_info.get("ok") else "unknown"
     devices = torch_info.get("devices", []) if torch_info.get("ok") else []
 
-    # Version metadata (best-effort)
     tooling["amdgpu_version"] = _probe_amdgpu_version()
     tooling["rocm_version"] = _probe_rocm_version()
 
@@ -172,4 +201,5 @@ def probe_gpus() -> ProbeResult:
         and bool(torch_info.get("cuda_is_available"))
         and int(torch_info.get("device_count", 0)) > 0
     )
-    return ProbeResult(ok=ok, backend=str(backend), devices=list(devices), tooling=tooling)
+    _PROBE_CACHE = ProbeResult(ok=ok, backend=str(backend), devices=list(devices), tooling=tooling)
+    return _PROBE_CACHE
diff --git a/primus/tools/preflight/gpu/gpu_topology.py b/primus/tools/preflight/gpu/gpu_topology.py
index e86198764..b047fc9bb 100644
--- a/primus/tools/preflight/gpu/gpu_topology.py
+++ b/primus/tools/preflight/gpu/gpu_topology.py
@@ -6,13 +6,17 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from collections import Counter
 from typing import Any, Dict, List, Optional
 
 from .gpu_probe import probe_gpus
+from .sysfs_probe import sysfs_gpu_bdfs, sysfs_topology_summary
 from .utils import Finding, ProbeResult, run_cmd, which
 
+logger = logging.getLogger(__name__)
+
 
 def _device_consistency(devices: List[Dict[str, Any]]) -> List[Finding]:
     findings: List[Finding] = []
@@ -29,9 +33,22 @@ def _device_consistency(devices: List[Dict[str, Any]]) -> List[Finding]:
 
 def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]:
     """
-    Best-effort GPU<->NUMA mapping.
-    On many systems this requires PCI bus IDs; we attempt to use amd-smi if present.
+    Best-effort GPU<->NUMA mapping via sysfs (primary) or amd-smi (fallback).
+
+    Sysfs reads are safe from any rank (no subprocesses, no mutex).
+    amd-smi fallback only attempted on LOCAL_RANK == 0.
     """
+    bdfs = sysfs_gpu_bdfs()
+    if bdfs:
+        return {"rc": 0, "gpus": bdfs, "source": "sysfs"}
+
+    logger.debug("sysfs NUMA mapping unavailable; trying amd-smi fallback")
+
+    from primus.tools.preflight.global_vars import LOCAL_RANK
+
+    if LOCAL_RANK != 0:
+        return None
+
     if which("amd-smi") is None:
         return None
 
@@ -39,17 +56,15 @@ def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]:
     if rc != 0 or not out:
         return {"rc": rc, "err": err, "note": "amd-smi list --csv failed"}
 
-    # amd-smi --csv format may vary; we do a very small heuristic:
-    # take 2nd column as PCI BDF if it looks like xxxx:xx:xx.x
     lines = [l for l in out.splitlines() if l.strip()]
-    bdfs: List[str] = []
+    bdf_list: List[str] = []
     for ln in lines[1:]:
         cols = [c.strip() for c in ln.split(",")]
         if len(cols) >= 2 and ":" in cols[1] and "." in cols[1]:
-            bdfs.append(cols[1])
+            bdf_list.append(cols[1])
 
     mapping: List[Dict[str, Any]] = []
-    for i, bdf in enumerate(bdfs):
+    for i, bdf in enumerate(bdf_list):
         node_path = f"/sys/bus/pci/devices/{bdf}/numa_node"
         numa = None
         if os.path.exists(node_path):
@@ -59,14 +74,29 @@ def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]:
                 numa = None
         mapping.append({"gpu": i, "pci_bdf": bdf, "numa_node": numa})
 
-    return {"rc": rc, "gpus": mapping}
+    return {"rc": rc, "gpus": mapping, "source": "amd-smi"}
 
 
 def _xgmi_presence_best_effort() -> Optional[Dict[str, Any]]:
-    # Best-effort: use amd-smi topology if available; otherwise skip.
+    """
+    Best-effort topology (XGMI vs PCIe) via sysfs (primary) or amd-smi (fallback).
+
+    Sysfs reads are safe from any rank. amd-smi fallback only on LOCAL_RANK == 0.
+    """
+    topo = sysfs_topology_summary()
+    if topo is not None and topo.get("rc") == 0:
+        return topo
+
+    logger.debug("sysfs topology unavailable; trying amd-smi fallback")
+
+    from primus.tools.preflight.global_vars import LOCAL_RANK
+
+    if LOCAL_RANK != 0:
+        return None
+
     if which("amd-smi") is None:
         return None
-    rc, out, err = run_cmd(["amd-smi", "topo"], timeout_s=10)
+    rc, out, err = run_cmd(["amd-smi", "topology"], timeout_s=10)
     if rc != 0:
         return {"rc": rc, "err": err}
     return {"rc": rc, "out": out}
@@ -105,9 +135,14 @@ def run_gpu_standard_checks(*, force_topology: bool = False) -> Dict[str, Any]:
     findings.extend(_device_consistency(probe.devices))
 
     # NUMA mapping / imbalance detection (best-effort).
+    # Prefers sysfs (safe for all ranks); falls back to amd-smi on LOCAL_RANK 0.
+    from primus.tools.preflight.global_vars import LOCAL_RANK
+
     numa = _numa_mapping_best_effort()
-    if numa is None:
-        findings.append(Finding("warn", "NUMA mapping unavailable (amd-smi not found); skipped", {}))
+    if numa is None and LOCAL_RANK != 0:
+        findings.append(Finding("info", "NUMA mapping skipped (non-zero LOCAL_RANK, no sysfs)", {}))
+    elif numa is None:
+        findings.append(Finding("warn", "NUMA mapping unavailable (sysfs/amd-smi not found); skipped", {}))
     else:
         nodes = [x.get("numa_node") for x in numa.get("gpus", []) if x.get("numa_node") is not None]
         imbalance = False
@@ -123,13 +158,16 @@ def run_gpu_standard_checks(*, force_topology: bool = False) -> Dict[str, Any]:
     # Topology (XGMI vs PCIe) best-effort.
     if force_topology or probe.backend == "rocm":
         topo = _xgmi_presence_best_effort()
-        if topo is None:
-            findings.append(Finding("warn", "Topology check skipped (amd-smi not found)", {}))
+        if topo is None and LOCAL_RANK != 0:
+            findings.append(Finding("info", "Topology check skipped (non-zero LOCAL_RANK, no sysfs)", {}))
+        elif topo is None:
+            findings.append(Finding("warn", "Topology check skipped (sysfs/amd-smi not found)", {}))
         else:
-            findings.append(Finding("info", "GPU topology (amd-smi topo)", topo))
-            # Detect obvious PCIe fallback hint (heuristic on output).
-            out = str(topo.get("out", ""))
-            if out and ("PCIE" in out.upper() or "PCIe" in out):
+            source = topo.get("source", "amd-smi")
+            findings.append(Finding("info", f"GPU topology ({source})", topo))
+            out = str(topo.get("out", "") or topo.get("matrix", ""))
+            has_xgmi = topo.get("has_xgmi")
+            if has_xgmi is False or (out and "PCIE" in out.upper() and "XGMI" not in out.upper()):
                 findings.append(Finding("warn", "Topology indicates PCIe paths; XGMI may be absent", {}))
 
     # RCCL/NCCL env sanity (WARN-only).
diff --git a/primus/tools/preflight/gpu/sysfs_probe.py b/primus/tools/preflight/gpu/sysfs_probe.py
new file mode 100644
index 000000000..3c033a61f
--- /dev/null
+++ b/primus/tools/preflight/gpu/sysfs_probe.py
@@ -0,0 +1,385 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Direct sysfs GPU probing — reads KFD topology and DRM sysfs to enumerate AMD
+GPUs, PCI BDFs, NUMA mapping, and inter-GPU link topology (XGMI vs PCIe)
+without invoking amd-smi or rocm-smi subprocesses.
+
+Approach adapted from RCCL's alt_rsmi.cc
+(https://github.com/ROCm/rocm-systems/blob/develop/projects/rccl/src/misc/alt_rsmi.cc)
+which reads the same KFD sysfs paths to avoid rocm_smi_lib's /dev/shm mutex
+contention (rocm_smi_lib#88).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+KFD_NODES_ROOT = "/sys/class/kfd/kfd/topology/nodes"
+AMD_VENDOR_ID = 0x1002
+
+LINK_TYPE_XGMI = 11
+LINK_TYPE_PCIE = 2
+
+
+@dataclass
+class GpuNode:
+    node_id: int = 0
+    gpu_id: int = 0
+    unique_id: int = 0
+    location_id: int = 0
+    domain: int = 0
+    bdf: int = 0
+    bus: int = 0
+    device: int = 0
+    function: int = 0
+    partition_id: int = 0
+    pci_bdf_str: str = ""
+    numa_node: Optional[int] = None
+    properties: Dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class LinkInfo:
+    src_idx: int = 0
+    dst_idx: int = 0
+    link_type: str = "unknown"
+    weight: int = 0
+    min_bandwidth: int = 0
+    max_bandwidth: int = 0
+    hops: int = 0
+
+
+@dataclass
+class SysfsProbeResult:
+    ok: bool = False
+    gpu_count: int = 0
+    gpus: List[GpuNode] = field(default_factory=list)
+    links: List[LinkInfo] = field(default_factory=list)
+    error: Optional[str] = None
+
+
+def _read_sysfs_file(path: str) -> Optional[str]:
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return f.read().strip()
+    except Exception:
+        return None
+
+
+def _read_sysfs_int(path: str) -> Optional[int]:
+    val = _read_sysfs_file(path)
+    if val is None:
+        return None
+    try:
+        return int(val)
+    except ValueError:
+        return None
+
+
+def _read_kfd_properties(node_id: int) -> Dict[str, int]:
+    """Parse /sys/class/kfd/kfd/topology/nodes/{node_id}/properties → dict."""
+    props: Dict[str, int] = {}
+    path = f"{KFD_NODES_ROOT}/{node_id}/properties"
+    content = _read_sysfs_file(path)
+    if content is None:
+        return props
+    for line in content.splitlines():
+        parts = line.split()
+        if len(parts) >= 2:
+            try:
+                props[parts[0]] = int(parts[1])
+            except ValueError:
+                pass
+    return props
+
+
+def _read_kfd_gpu_id(node_id: int) -> Optional[int]:
+    """Read /sys/class/kfd/kfd/topology/nodes/{node_id}/gpu_id."""
+    return _read_sysfs_int(f"{KFD_NODES_ROOT}/{node_id}/gpu_id")
+
+
+def _read_link_properties(node_id: int, link_id: int) -> Dict[str, int]:
+    """Parse /sys/class/kfd/kfd/topology/nodes/{node_id}/io_links/{link_id}/properties."""
+    path = f"{KFD_NODES_ROOT}/{node_id}/io_links/{link_id}/properties"
+    props: Dict[str, int] = {}
+    content = _read_sysfs_file(path)
+    if content is None:
+        return props
+    for line in content.splitlines():
+        parts = line.split()
+        if len(parts) >= 2:
+            try:
+                props[parts[0]] = int(parts[1])
+            except ValueError:
+                pass
+    return props
+
+
+def _count_io_links(node_id: int) -> int:
+    """Count subdirectories in io_links/ for a given node."""
+    links_dir = f"{KFD_NODES_ROOT}/{node_id}/io_links"
+    if not os.path.isdir(links_dir):
+        return 0
+    count = 0
+    try:
+        for entry in os.listdir(links_dir):
+            if os.path.isdir(os.path.join(links_dir, entry)) and entry not in (".", ".."):
+                count += 1
+    except OSError:
+        pass
+    return count
+
+
+def _bdf_to_string(domain: int, bus: int, device: int, function: int) -> str:
+    """Format PCI BDF as string like '0000:c1:00.0'."""
+    return f"{domain:04x}:{bus:02x}:{device:02x}.{function:x}"
+
+
+_GPU_NODES_CACHE: Optional[List["GpuNode"]] = None
+
+
+def _enumerate_gpu_nodes() -> List[GpuNode]:
+    """
+    Scan KFD topology nodes, filter AMD GPUs (vendor_id == 0x1002),
+    compute PCI BDF from location_id + domain, sort by BDF.
+
+    Results are cached at module level since GPU topology is static
+    during a single preflight run and this is called from multiple paths.
+
+    Logic mirrors ARSMI_init() from alt_rsmi.cc.
+    """
+    global _GPU_NODES_CACHE
+    if _GPU_NODES_CACHE is not None:
+        return _GPU_NODES_CACHE
+
+    if not os.path.isdir(KFD_NODES_ROOT):
+        return []
+
+    raw_nodes: List[Tuple[int, GpuNode]] = []
+
+    try:
+        entries = os.listdir(KFD_NODES_ROOT)
+    except OSError:
+        return []
+
+    for entry in entries:
+        if not entry.isdigit():
+            continue
+
+        node_id = int(entry)
+        gpu_id = _read_kfd_gpu_id(node_id)
+        if gpu_id is None or gpu_id == 0:
+            continue
+
+        props = _read_kfd_properties(node_id)
+        vendor_id = props.get("vendor_id", 0)
+        if vendor_id != AMD_VENDOR_ID:
+            continue
+
+        unique_id = props.get("unique_id", 0)
+        location_id = props.get("location_id", 0)
+        domain = props.get("domain", 0) & 0xFFFFFFFF
+
+        bdf_raw = (domain << 32) | location_id
+        bus = (location_id >> 8) & 0xFF
+        device = (location_id >> 3) & 0x1F
+        function = location_id & 0x7
+        partition_id = (location_id >> 28) & 0xF
+
+        pci_bdf_str = _bdf_to_string(domain, bus, device, function)
+
+        numa = _read_sysfs_int(f"/sys/bus/pci/devices/{pci_bdf_str}/numa_node")
+
+        node = GpuNode(
+            node_id=node_id,
+            gpu_id=gpu_id,
+            unique_id=unique_id,
+            location_id=location_id,
+            domain=domain,
+            bdf=bdf_raw,
+            bus=bus,
+            device=device,
+            function=function,
+            partition_id=partition_id,
+            pci_bdf_str=pci_bdf_str,
+            numa_node=numa,
+            properties=props,
+        )
+        raw_nodes.append((bdf_raw, node))
+
+    raw_nodes.sort(key=lambda x: x[0])
+    _GPU_NODES_CACHE = [n for _, n in raw_nodes]
+    return _GPU_NODES_CACHE
+
+
+def _build_link_matrix(gpu_nodes: List[GpuNode]) -> List[LinkInfo]:
+    """
+    Read io_links for each GPU node to determine XGMI vs PCIe connectivity.
+    Returns a flat list of LinkInfo for GPU-to-GPU links only.
+
+    Logic mirrors the link matrix construction in ARSMI_init().
+    """
+    if not gpu_nodes:
+        return []
+
+    node_id_to_idx = {g.node_id: i for i, g in enumerate(gpu_nodes)}
+    links: List[LinkInfo] = []
+
+    for src_idx, src_node in enumerate(gpu_nodes):
+        n_links = _count_io_links(src_node.node_id)
+        for link_id in range(n_links):
+            props = _read_link_properties(src_node.node_id, link_id)
+            if not props:
+                continue
+
+            dst_node_id = props.get("node_to")
+            if dst_node_id is None:
+                continue
+
+            dst_idx = node_id_to_idx.get(dst_node_id)
+            if dst_idx is None:
+                continue
+
+            link_type_raw = props.get("type", 0)
+            weight = props.get("weight", 0)
+            min_bw = props.get("min_bandwidth", 0)
+            max_bw = props.get("max_bandwidth", 0)
+
+            if link_type_raw == LINK_TYPE_XGMI:
+                link_type = "XGMI"
+                hops = 1
+            elif link_type_raw == LINK_TYPE_PCIE:
+                link_type = "PCIe"
+                hops = 2
+            else:
+                link_type = "unknown"
+                hops = 0
+
+            links.append(
+                LinkInfo(
+                    src_idx=src_idx,
+                    dst_idx=dst_idx,
+                    link_type=link_type,
+                    weight=weight,
+                    min_bandwidth=min_bw,
+                    max_bandwidth=max_bw,
+                    hops=hops,
+                )
+            )
+
+    return links
+
+
+def sysfs_probe() -> SysfsProbeResult:
+    """
+    Main entry point: enumerate AMD GPUs and topology via sysfs.
+
+    No subprocess calls, no /dev/shm mutex, safe to call from any rank.
+    """
+    try:
+        gpu_nodes = _enumerate_gpu_nodes()
+        if not gpu_nodes:
+            return SysfsProbeResult(ok=False, error="No AMD GPUs found via KFD sysfs")
+
+        links = _build_link_matrix(gpu_nodes)
+        return SysfsProbeResult(
+            ok=True,
+            gpu_count=len(gpu_nodes),
+            gpus=gpu_nodes,
+            links=links,
+        )
+    except Exception as e:
+        return SysfsProbeResult(ok=False, error=f"sysfs probe failed: {e}")
+
+
+# ── Convenience helpers for integration with existing preflight code ──
+
+
+def sysfs_gpu_count() -> int:
+    """GPU count via KFD sysfs. Returns 0 on failure."""
+    try:
+        nodes = _enumerate_gpu_nodes()
+        return len(nodes)
+    except Exception:
+        return 0
+
+
+def sysfs_gpu_bdfs() -> List[Dict[str, Any]]:
+    """
+    Return per-GPU BDF + NUMA mapping, compatible with the format
+    _numa_mapping_best_effort() currently returns.
+    """
+    try:
+        nodes = _enumerate_gpu_nodes()
+        return [{"gpu": i, "pci_bdf": n.pci_bdf_str, "numa_node": n.numa_node} for i, n in enumerate(nodes)]
+    except Exception:
+        return []
+
+
+def sysfs_has_xgmi() -> Optional[bool]:
+    """
+    Check if any GPU-to-GPU link is XGMI. Returns None if probe fails.
+    """
+    try:
+        nodes = _enumerate_gpu_nodes()
+        if not nodes:
+            return None
+        links = _build_link_matrix(nodes)
+        return any(lk.link_type == "XGMI" for lk in links)
+    except Exception:
+        return None
+
+
+def sysfs_topology_summary() -> Optional[Dict[str, Any]]:
+    """
+    Build a topology summary similar to what `amd-smi topo` provides,
+    formatted for preflight reporting.
+    """
+    try:
+        nodes = _enumerate_gpu_nodes()
+        if not nodes:
+            return None
+
+        links = _build_link_matrix(nodes)
+        n = len(nodes)
+
+        matrix: List[List[str]] = [["" for _ in range(n)] for _ in range(n)]
+        for i in range(n):
+            matrix[i][i] = "self"
+
+        for lk in links:
+            matrix[lk.src_idx][lk.dst_idx] = lk.link_type
+
+        header = [f"GPU{i}" for i in range(n)]
+        lines = ["       " + "  ".join(f"{h:>6}" for h in header)]
+        for i in range(n):
+            row = f"GPU{i:<3} " + "  ".join(f"{matrix[i][j]:>6}" for j in range(n))
+            lines.append(row)
+
+        has_xgmi = any(lk.link_type == "XGMI" for lk in links)
+
+        return {
+            "rc": 0,
+            "source": "sysfs",
+            "gpu_count": n,
+            "has_xgmi": has_xgmi,
+            "matrix": "\n".join(lines),
+            "links": [
+                {
+                    "src": lk.src_idx,
+                    "dst": lk.dst_idx,
+                    "type": lk.link_type,
+                    "weight": lk.weight,
+                }
+                for lk in links
+            ],
+        }
+    except Exception as e:
+        return {"rc": 1, "error": str(e)}
diff --git a/primus/tools/preflight/host/host_probe.py b/primus/tools/preflight/host/host_probe.py
index fde5eb0e5..238a72db0 100644
--- a/primus/tools/preflight/host/host_probe.py
+++ b/primus/tools/preflight/host/host_probe.py
@@ -488,31 +488,25 @@ def _get_pcie_link_info_sysfs() -> Dict[str, Any]:
     return info
 
 
-def get_gpu_count_rocm() -> int:
-    """Get GPU count using rocm-smi (more reliable in containers)."""
+def get_gpu_count_sysfs() -> int:
+    """GPU count via KFD sysfs — no subprocesses, no /dev/shm mutex."""
     try:
-        result = subprocess.run(
-            ["rocm-smi", "--showid"],
-            capture_output=True,
-            text=True,
-            timeout=10,
-        )
-        if result.returncode == 0:
-            # Count lines that look like GPU entries (contain "GPU[")
-            count = sum(1 for line in result.stdout.splitlines() if "GPU[" in line)
-            if count > 0:
-                return count
-    except FileNotFoundError:
-        pass
+        from primus.tools.preflight.gpu.sysfs_probe import sysfs_gpu_count
+
+        count = sysfs_gpu_count()
+        if count > 0:
+            return count
     except Exception:
         pass
+    return 0
+
 
-    # Fallback: try HIP_VISIBLE_DEVICES or count /dev/dri/renderD* devices
+def get_gpu_count_rocm_fallback() -> int:
+    """GPU count without invoking rocm-smi (HIP_VISIBLE_DEVICES / /dev/dri)."""
     hip_devices = os.environ.get("HIP_VISIBLE_DEVICES", "")
     if hip_devices:
         return len([x for x in hip_devices.split(",") if x.strip()])
 
-    # Count render devices
     try:
         import glob
 
@@ -525,6 +519,27 @@ def get_gpu_count_rocm() -> int:
     return 0
 
 
+def get_gpu_count_rocm() -> int:
+    """Get GPU count using rocm-smi (more reliable in containers)."""
+    try:
+        result = subprocess.run(
+            ["rocm-smi", "--showid"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode == 0:
+            count = sum(1 for line in result.stdout.splitlines() if "GPU[" in line)
+            if count > 0:
+                return count
+    except FileNotFoundError:
+        pass
+    except Exception:
+        pass
+
+    return get_gpu_count_rocm_fallback()
+
+
 def get_pcie_topology() -> Optional[str]:
     """Get PCIe topology using lstopo or lspci -t."""
     try:
diff --git a/primus/tools/preflight/host/info.py b/primus/tools/preflight/host/info.py
index ed99dd60d..fd5c1aa80 100644
--- a/primus/tools/preflight/host/info.py
+++ b/primus/tools/preflight/host/info.py
@@ -16,6 +16,8 @@
 from .host_probe import (
     get_cpu_info,
     get_gpu_count_rocm,
+    get_gpu_count_rocm_fallback,
+    get_gpu_count_sysfs,
     get_hostname,
     get_kernel_version,
     get_memory_info,
@@ -95,9 +97,16 @@ def collect_host_info() -> List[Finding]:
     ib_count = sum(1 for d in pcie_devices if d["type"] == "Infiniband")
     eth_count = sum(1 for d in pcie_devices if d["type"] == "Ethernet")
 
-    # Use rocm-smi as fallback for GPU count (more reliable in containers)
-    gpu_count_rocm = get_gpu_count_rocm()
-    gpu_count = max(gpu_count_pcie, gpu_count_rocm)
+    # Primary: sysfs (KFD topology) — safe from any rank, no subprocess.
+    # Fallback: rocm-smi on LOCAL_RANK 0 only (subprocess, /dev/shm mutex).
+    gpu_count_sysfs = get_gpu_count_sysfs()
+    if gpu_count_sysfs > 0:
+        gpu_count_extra = gpu_count_sysfs
+    else:
+        from primus.tools.preflight.global_vars import LOCAL_RANK
+
+        gpu_count_extra = get_gpu_count_rocm() if LOCAL_RANK == 0 else get_gpu_count_rocm_fallback()
+    gpu_count = max(gpu_count_pcie, gpu_count_extra)
 
     findings.append(
         Finding(
diff --git a/primus/tools/preflight/inter_node_comm.py b/primus/tools/preflight/inter_node_comm.py
index 2ac5213fe..8f1ca06a5 100644
--- a/primus/tools/preflight/inter_node_comm.py
+++ b/primus/tools/preflight/inter_node_comm.py
@@ -5,32 +5,98 @@
 ###############################################################################
 
 import time
+from typing import Iterable, List, Optional, Sequence, Union
 
-import matplotlib.pyplot as plt
 import torch
 import torch.distributed as dist
 
 from primus.tools.preflight.global_vars import (
-    ITERATION,
     LOCAL_RANK,
     LOCAL_WORLD_SIZE,
     RANK,
-    WARMUP,
     WORLD_SIZE,
     get_hostnames,
+    get_iteration,
+    get_warmup,
 )
 from primus.tools.preflight.utility import (
+    barrier_after_comm_destroy,
     create_dir,
     extract_first_middle_last,
     extract_number,
+    format_int_range,
     log,
 )
 
+# Hard cap on the per-group node count for inter-node alltoall. Real-world
+# MoE training rarely dispatches across more than ~8 nodes (e.g. DeepSeek-V3's
+# largest published EP=64 spans 8 nodes with per-token dispatch capped at
+# 4 nodes), so 16 covers every published configuration with comfortable
+# headroom. Capping here also eliminates the dominant source of EADDRINUSE at
+# >=128 nodes -- uncapped all-N inter-node alltoall destroys generate enough
+# IB OOB TIME_WAIT per node to exhaust the default ephemeral-port pool.
+# Other inter-node tests (allreduce / p2p / ring-p2p) are unaffected.
+# Intentionally not exposed as a CLI flag: this is a known-safe ceiling for
+# the comm shapes preflight is supposed to characterize, not a tuning knob.
+_INTER_ALLTOALL_MAX_NODES = 16
 
-def run_inter_node_comm(args):
+
+def _resolve_inter_group_sizes(
+    group_sizes: Optional[Sequence[Union[int, str]]],
+    num_nodes: int,
+) -> List[int]:
+    """Translate user-supplied inter-node group sizes (with 'all') to ints.
+
+    - 'all' is mapped to num_nodes.
+    - values > num_nodes are dropped.
+    - duplicates are removed and the result is sorted ascending.
+    """
+    if group_sizes is None or len(group_sizes) == 0:
+        candidates = [2, 4, num_nodes]
+    else:
+        candidates = []
+        for g in group_sizes:
+            if isinstance(g, str) and g.strip().lower() == "all":
+                candidates.append(num_nodes)
+            else:
+                candidates.append(int(g))
+    candidates = [c for c in candidates if c >= 2 and c <= num_nodes]
+    return sorted(set(candidates))
+
+
+def run_inter_node_comm(
+    args,
+    enabled_comms: Optional[Iterable[str]] = None,
+    sizes_mb: Optional[Sequence[int]] = None,
+    group_sizes: Optional[Sequence[Union[int, str]]] = None,
+):
+    """Inter-node allreduce / alltoall benchmark.
+
+    Args:
+        args: parsed namespace.
+        enabled_comms: subset of {"allreduce", "alltoall"} to run. Defaults to both.
+        sizes_mb: message sizes in MB.
+        group_sizes: list of node group sizes; values > num_nodes are dropped.
+            'all' is accepted as a synonym for num_nodes. Defaults to [2, 4, num_nodes].
+
+    Note:
+        The alltoall path additionally clamps every requested per-group node
+        count to ``_INTER_ALLTOALL_MAX_NODES`` (16) before deduping, regardless
+        of the cluster size or the user's ``--inter-group-sizes`` choice. The
+        allreduce path uses the requested sizes unchanged.
+    """
     device = torch.device(f"cuda:{LOCAL_RANK}")
-    sizes = [2**i * 1024 * 1024 for i in range(1, 11)]
-    # sizes = [2**i * 1024 * 1024 for i in range(1, 5)]
+
+    if sizes_mb is None or len(sizes_mb) == 0:
+        sizes_mb = [2**i for i in range(1, 11)]
+    sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb]
+
+    enabled_set = set(enabled_comms) if enabled_comms else {"allreduce", "alltoall"}
+    enabled_set &= {"allreduce", "alltoall"}
+    if not enabled_set:
+        log("Skip inter-node comm benchmark (no enabled comms)")
+        return
+
     assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0
     num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE
 
@@ -38,13 +104,30 @@ def run_inter_node_comm(args):
         log(f"Skip inter node comm benchmark, {num_nodes=}")
         return
 
-    # N-node allreduce & alltoall (adjacent pairs)
-    # 2-node allreduce, pair nodes: [0, 1], [2, 3], ...
-    # 4-node allreduce, pair nodes: [0, 1, 2, 3], [4, 5, 6, 7]...
-    cases = {
-        "allreduce": list(set([2, 4] + [num_nodes])),
-        "alltoall": list(set([2, 4] + [num_nodes])),
-    }
+    node_counts = _resolve_inter_group_sizes(group_sizes, num_nodes)
+    if not node_counts:
+        log("Skip inter-node comm benchmark, no valid group sizes")
+        return
+
+    cases = {}
+    for comm in ("allreduce", "alltoall"):
+        if comm not in enabled_set:
+            continue
+        if comm == "alltoall":
+            capped = sorted({min(c, _INTER_ALLTOALL_MAX_NODES) for c in node_counts})
+            if capped != list(node_counts):
+                log(
+                    f"  inter-alltoall: per-group node count capped at "
+                    f"{_INTER_ALLTOALL_MAX_NODES} (requested {list(node_counts)} -> "
+                    f"running {capped}). Real-world MoE alltoall rarely exceeds "
+                    f"~8 nodes; see docs/preflight.md \u00a75.2 for the rationale."
+                )
+            cases[comm] = capped
+        else:
+            cases[comm] = list(node_counts)
+
+    warmup = get_warmup()
+    iteration = get_iteration()
 
     if RANK == 0:
         with open(args.markdown_file, "a", encoding="utf-8") as f:
@@ -62,19 +145,39 @@ def run_inter_node_comm(args):
             latency_results = {}
             bandwidth_results = {}
 
-            num_procs = adjacent_nodes * LOCAL_WORLD_SIZE
-            num_adjacent_groups = num_nodes // adjacent_nodes
+            num_full_groups = num_nodes // adjacent_nodes
+            remainder_nodes = num_nodes % adjacent_nodes
             adjacent_group = None
-            for i_group in range(num_adjacent_groups):
-                group_ranks = [
-                    i_group * adjacent_nodes * LOCAL_WORLD_SIZE + r
-                    for r in range(adjacent_nodes * LOCAL_WORLD_SIZE)
-                ]
+            # Track per-group member ranks for compact reporting.
+            all_group_ranks: List[List[int]] = []
+            group_node_counts: List[int] = []
+
+            for i_group in range(num_full_groups):
+                group_start = i_group * adjacent_nodes * LOCAL_WORLD_SIZE
+                group_ranks = [group_start + r for r in range(adjacent_nodes * LOCAL_WORLD_SIZE)]
                 tmp_group = dist.new_group(ranks=group_ranks)
                 if RANK in group_ranks:
                     assert adjacent_group is None
                     adjacent_group = tmp_group
-            if RANK < num_adjacent_groups * adjacent_nodes * LOCAL_WORLD_SIZE:
+                all_group_ranks.append(group_ranks)
+                group_node_counts.append(adjacent_nodes)
+
+            if remainder_nodes >= 2:
+                group_start = num_full_groups * adjacent_nodes * LOCAL_WORLD_SIZE
+                group_ranks = [group_start + r for r in range(remainder_nodes * LOCAL_WORLD_SIZE)]
+                tmp_group = dist.new_group(ranks=group_ranks)
+                if RANK in group_ranks:
+                    assert adjacent_group is None
+                    adjacent_group = tmp_group
+                all_group_ranks.append(group_ranks)
+                group_node_counts.append(remainder_nodes)
+
+            num_procs = dist.get_world_size(adjacent_group) if adjacent_group is not None else 0
+
+            total_grouped_ranks = num_full_groups * adjacent_nodes * LOCAL_WORLD_SIZE
+            if remainder_nodes >= 2:
+                total_grouped_ranks += remainder_nodes * LOCAL_WORLD_SIZE
+            if RANK < total_grouped_ranks:
                 assert adjacent_group is not None
 
             for size in sizes:
@@ -83,7 +186,7 @@ def run_inter_node_comm(args):
 
                 tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device)
                 dist.barrier(group=adjacent_group, device_ids=[torch.cuda.current_device()])
-                for _ in range(WARMUP):
+                for _ in range(warmup):
                     if "allreduce" == comm:
                         dist.all_reduce(tensor, group=adjacent_group)
                     elif "alltoall" == comm:
@@ -92,7 +195,7 @@ def run_inter_node_comm(args):
                         assert False
                 torch.cuda.synchronize()
                 start = time.time()
-                for _ in range(ITERATION):
+                for _ in range(iteration):
                     if "allreduce" == comm:
                         dist.all_reduce(tensor, group=adjacent_group)
                     elif "alltoall" == comm:
@@ -100,7 +203,7 @@ def run_inter_node_comm(args):
                     else:
                         assert False
                 torch.cuda.synchronize()
-                elapsed = (time.time() - start) / ITERATION
+                elapsed = (time.time() - start) / iteration
                 scale = 2 if comm == "allreduce" else 1
                 comm_size = scale * size * (num_procs - 1) / num_procs
                 gb_per_sec = comm_size / elapsed / 1e9
@@ -110,6 +213,7 @@ def run_inter_node_comm(args):
             dist.barrier(device_ids=[torch.cuda.current_device()])
             if adjacent_group is not None:
                 dist.destroy_process_group(adjacent_group)
+            barrier_after_comm_destroy(args.comm_cleanup_delay_sec)
 
             all_latency_results = [None for _ in range(WORLD_SIZE)]
             all_bandwidth_results = [None for _ in range(WORLD_SIZE)]
@@ -120,68 +224,82 @@ def run_inter_node_comm(args):
                 keys = sorted(
                     list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number
                 )
-                max_len = max(len(s) for s in get_hostnames()) + 2
+                hostnames = get_hostnames()
+
+                # Show only the leader node's hostname; the Node range plus the
+                # legend at the top of the report cover the rest.
+                def _row_for(group_ranks: List[int], results):
+                    leader = group_ranks[0]
+                    host_str = hostnames[leader]
+                    node_str = format_int_range([r // LOCAL_WORLD_SIZE for r in group_ranks])
+                    rank_str = format_int_range(group_ranks)
+                    return host_str, node_str, rank_str, results[leader]
+
+                formatted_keys = [f"{key:<6}" for key in keys]
+                host_col_label = "Leader hostname"
+                host_col_w = max(20, len(host_col_label) + 2)
+                header_line = (
+                    f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}"
+                )
 
                 with open(args.markdown_file, "a", encoding="utf-8") as f:
                     f.write(f"=======InterNodeComm - {case_name} (us)=======\n")
                     log(f"=======InterNodeComm - {case_name} (us)=======")
+                    log(header_line)
 
-                    f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+                    f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
                     f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-
-                    formatted_keys = [f"{key:<6}" for key in keys]
-                    log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
-                    for rank, r in enumerate(all_latency_results):
-                        hostname = get_hostnames()[rank]
-                        if rank % num_procs != 0:
-                            continue
-                        node_id = rank // LOCAL_WORLD_SIZE
-
+                    for group_ranks in all_group_ranks:
+                        host_str, node_str, rank_str, r = _row_for(group_ranks, all_latency_results)
                         formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys]
-                        log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                        f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                        log(
+                            f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} "
+                            f"{' '.join(formatted_values)}"
+                        )
+                        f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
                     f.write(f"\n")
 
                     f.write(f"=======InterNodeComm - {case_name} (GB/s)=======\n")
                     log(f"=======InterNodeComm - {case_name} (GB/s)=======")
+                    log(header_line)
 
-                    f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+                    f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
                     f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-                    formatted_keys = [f"{key:<6}" for key in keys]
-                    log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
-                    for rank, r in enumerate(all_bandwidth_results):
-                        hostname = get_hostnames()[rank]
-                        if rank % num_procs != 0:
-                            continue
-                        node_id = rank // LOCAL_WORLD_SIZE
-
+                    for group_ranks in all_group_ranks:
+                        host_str, node_str, rank_str, r = _row_for(group_ranks, all_bandwidth_results)
                         formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys]
-                        log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                        f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                        log(
+                            f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} "
+                            f"{' '.join(formatted_values)}"
+                        )
+                        f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
                     f.write(f"\n")
 
                 if not args.plot:
                     continue
 
-                log(f"=======Plot IntraNode {case_name} Bandwidth=======")
+                import matplotlib.pyplot as plt
+
+                log(f"=======Plot InterNode {case_name} Bandwidth=======")
                 with open(args.markdown_file, "a", encoding="utf-8") as f:
                     f.write(f"=======Plot InterNode {case_name} Bandwidth=======\n")
                 plot_case = f"inter_node_comm/{comm}"
                 dump_path = f"{args.dump_path}/{plot_case}"
                 create_dir(dump_path)
                 print_keys = extract_first_middle_last(keys)
-                first_rank_bandwidth_results = [
-                    all_bandwidth_results[i] for i in range(len(all_bandwidth_results)) if i % num_procs == 0
-                ]
+                leader_ranks = [g[0] for g in all_group_ranks]
+                first_rank_bandwidth_results = [all_bandwidth_results[i] for i in leader_ranks]
                 num_print_ranks = len(first_rank_bandwidth_results)
                 for size_key in print_keys:
                     values = [r[size_key] for r in first_rank_bandwidth_results]
                     plt.figure(figsize=(10, 4))
                     bars = plt.bar(range(num_print_ranks), values)
-                    plt.xlabel(f"RankPair ({num_procs} ranks)")
+                    plt.xlabel(f"Group (starting rank)")
                     plt.ylabel("Bandwidth")
                     plt.title(f"Inter Node {case_name} Bandwidth for {size_key}")
-                    xtick_labels = [f"{i*num_procs}" for i in range(num_print_ranks)]
+                    xtick_labels = [
+                        f"{leader_ranks[i]} ({group_node_counts[i]}N)" for i in range(num_print_ranks)
+                    ]
                     plt.xticks(range(num_print_ranks), xtick_labels)
                     plt.grid(True, axis="y")
 
diff --git a/primus/tools/preflight/inter_node_comm_p2p.py b/primus/tools/preflight/inter_node_comm_p2p.py
index 9b38feee4..8f2574e03 100644
--- a/primus/tools/preflight/inter_node_comm_p2p.py
+++ b/primus/tools/preflight/inter_node_comm_p2p.py
@@ -5,32 +5,37 @@
 ###############################################################################
 
 import time
+from typing import Optional, Sequence
 
-import matplotlib.pyplot as plt
 import torch
 import torch.distributed as dist
 
 from primus.tools.preflight.global_vars import (
-    ITERATION,
     LOCAL_RANK,
     LOCAL_WORLD_SIZE,
     RANK,
-    WARMUP,
     WORLD_SIZE,
     get_hostnames,
+    get_iteration,
+    get_warmup,
 )
 from primus.tools.preflight.utility import (
+    barrier_after_comm_destroy,
     create_dir,
     extract_first_middle_last,
     extract_number,
+    format_int_range,
     log,
 )
 
 
-def run_inter_node_comm_p2p(args):
+def run_inter_node_comm_p2p(args, sizes_mb: Optional[Sequence[int]] = None):
     device = torch.device(f"cuda:{LOCAL_RANK}")
-    sizes = [2**i * 1024 * 1024 for i in range(1, 11)]
-    # sizes = [2**i * 1024 * 1024 for i in range(1, 5)]
+    if sizes_mb is None or len(sizes_mb) == 0:
+        sizes_mb = [2**i for i in range(1, 11)]
+    sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb]
+    warmup = get_warmup()
+    iteration = get_iteration()
     assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0
     num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE
 
@@ -49,10 +54,14 @@ def run_inter_node_comm_p2p(args):
     bandwidth_results = {}
 
     num_adjacent_groups = num_nodes // adjacent_nodes
+    num_paired_ranks = num_adjacent_groups * adjacent_nodes * LOCAL_WORLD_SIZE
     p2p_group = None
     is_src_rank = ((RANK // LOCAL_WORLD_SIZE) % 2) == 0
-    peer_rank = RANK + LOCAL_WORLD_SIZE if is_src_rank else RANK - LOCAL_WORLD_SIZE
-    assert peer_rank >= 0 and peer_rank < WORLD_SIZE
+    if RANK < num_paired_ranks:
+        peer_rank = RANK + LOCAL_WORLD_SIZE if is_src_rank else RANK - LOCAL_WORLD_SIZE
+        assert peer_rank >= 0 and peer_rank < WORLD_SIZE
+    else:
+        peer_rank = -1
     for i_group in range(num_adjacent_groups):
         for i_r in range(LOCAL_WORLD_SIZE):
             group_ranks = [
@@ -76,20 +85,20 @@ def run_inter_node_comm_p2p(args):
 
         tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device)
         dist.barrier(group=p2p_group, device_ids=[torch.cuda.current_device()])
-        for _ in range(WARMUP):
+        for _ in range(warmup):
             if is_src_rank:
                 dist.send(tensor, dst=peer_rank, group=p2p_group)
             else:
                 dist.recv(tensor, src=peer_rank, group=p2p_group)
         torch.cuda.synchronize()
         start = time.time()
-        for _ in range(ITERATION):
+        for _ in range(iteration):
             if is_src_rank:
                 dist.send(tensor, dst=peer_rank, group=p2p_group)
             else:
                 dist.recv(tensor, src=peer_rank, group=p2p_group)
         torch.cuda.synchronize()
-        elapsed = (time.time() - start) / ITERATION
+        elapsed = (time.time() - start) / iteration
         comm_size = size
         gb_per_sec = comm_size / elapsed / 1e9
         latency_results[f"{size//1024//1024}MB"] = elapsed * 1e6
@@ -98,6 +107,7 @@ def run_inter_node_comm_p2p(args):
     dist.barrier(device_ids=[torch.cuda.current_device()])
     if p2p_group is not None:
         dist.destroy_process_group(p2p_group)
+    barrier_after_comm_destroy(args.comm_cleanup_delay_sec)
 
     all_latency_results = [None for _ in range(WORLD_SIZE)]
     all_bandwidth_results = [None for _ in range(WORLD_SIZE)]
@@ -106,7 +116,7 @@ def run_inter_node_comm_p2p(args):
 
     if RANK == 0:
         keys = sorted(list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number)
-        max_len = max(len(s) for s in get_hostnames()) + 2
+        hostnames = get_hostnames()
 
         # result of src ranks will be print
         src_ranks = []
@@ -114,6 +124,8 @@ def run_inter_node_comm_p2p(args):
         src_rank_latency_results = []
         src_rank_bandwidth_results = []
         for rank, r in enumerate(all_bandwidth_results):
+            if rank >= num_paired_ranks:
+                continue
             is_src_rank = ((rank // LOCAL_WORLD_SIZE) % 2) == 0
             peer_rank = rank + LOCAL_WORLD_SIZE if is_src_rank else rank - LOCAL_WORLD_SIZE
             assert peer_rank >= 0 and peer_rank < WORLD_SIZE
@@ -124,45 +136,61 @@ def run_inter_node_comm_p2p(args):
             src_rank_latency_results.append(all_latency_results[rank])
             src_rank_bandwidth_results.append(r)
 
+        # Show only the leader (src) host. Both nodes appear in the Node column,
+        # and the Node->Hostname legend at the top of the report covers the rest.
+        def _row_for(src: int, peer: int):
+            host_str = hostnames[src]
+            node_str = format_int_range([src // LOCAL_WORLD_SIZE, peer // LOCAL_WORLD_SIZE])
+            rank_str = format_int_range([src, peer])
+            return host_str, node_str, rank_str
+
+        formatted_keys = [f"{key:<6}" for key in keys]
+        host_col_label = "Leader hostname"
+        host_col_w = max(20, len(host_col_label) + 2)
+        header_line = (
+            f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}"
+        )
+
         with open(args.markdown_file, "a", encoding="utf-8") as f:
             f.write(f"=======InterNodeComm - {case_name} (us)=======\n")
             log(f"=======InterNodeComm - {case_name} (us)=======")
+            log(header_line)
 
-            f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+            f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
             f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-
-            formatted_keys = [f"{key:<6}" for key in keys]
-            log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
             for i_r in range(len(src_ranks)):
-                rank = src_ranks[i_r]
-                hostname = get_hostnames()[rank]
-                node_id = rank // LOCAL_WORLD_SIZE
-
+                src = src_ranks[i_r]
+                peer = peer_ranks[i_r]
+                host_str, node_str, rank_str = _row_for(src, peer)
                 formatted_values = [f"{src_rank_latency_results[i_r].get(key, 0):<6.2f}" for key in keys]
-                log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                log(
+                    f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " f"{' '.join(formatted_values)}"
+                )
+                f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
             f.write(f"\n")
 
             f.write(f"=======InterNodeComm - {case_name} (GB/s)=======\n")
             log(f"=======InterNodeComm - {case_name} (GB/s)=======")
+            log(header_line)
 
-            f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+            f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
             f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-            formatted_keys = [f"{key:<6}" for key in keys]
-            log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
             for i_r in range(len(src_ranks)):
-                rank = src_ranks[i_r]
-                hostname = get_hostnames()[rank]
-                node_id = rank // LOCAL_WORLD_SIZE
-
+                src = src_ranks[i_r]
+                peer = peer_ranks[i_r]
+                host_str, node_str, rank_str = _row_for(src, peer)
                 formatted_values = [f"{src_rank_bandwidth_results[i_r].get(key, 0):<6.2f}" for key in keys]
-                log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                log(
+                    f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " f"{' '.join(formatted_values)}"
+                )
+                f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
             f.write(f"\n")
 
         if not args.plot:
             return
 
+        import matplotlib.pyplot as plt
+
         log(f"=======Plot InterNode {case_name} Bandwidth=======")
         with open(args.markdown_file, "a", encoding="utf-8") as f:
             f.write(f"=======Plot InterNode {case_name} Bandwidth=======\n")
diff --git a/primus/tools/preflight/inter_node_ring_p2p.py b/primus/tools/preflight/inter_node_ring_p2p.py
index 931685ae9..d5b0ab306 100644
--- a/primus/tools/preflight/inter_node_ring_p2p.py
+++ b/primus/tools/preflight/inter_node_ring_p2p.py
@@ -1,19 +1,20 @@
+from typing import Optional, Sequence
+
 import torch
 import torch.distributed as dist
 from torch.profiler import ProfilerActivity
 
 from primus.tools.preflight.global_vars import (
-    ITERATION,
     LOCAL_WORLD_SIZE,
     RANK,
-    WARMUP,
     WORLD_SIZE,
+    get_iteration,
+    get_warmup,
 )
-from primus.tools.preflight.utility import log
+from primus.tools.preflight.utility import barrier_after_comm_destroy, log
 
 # profile parameters
 _ENABLE_PROFILE = False
-_PROFILE_STEPS = min(ITERATION, 10)
 
 _NODE_RANK = RANK // LOCAL_WORLD_SIZE
 _GLOBAL_PIPELINE_GROUP = None
@@ -76,10 +77,14 @@ def run_ring_p2p(num_bytes: int):
     start_event = torch.cuda.Event(enable_timing=True)
     end_event = torch.cuda.Event(enable_timing=True)
 
-    for it in range(WARMUP):
+    warmup = get_warmup()
+    iteration = get_iteration()
+    profile_steps = min(iteration, 10)
+
+    for it in range(warmup):
         reqs = send_recv_once(x, y)
 
-    if WARMUP > 0:
+    if warmup > 0:
         # TODO(limou)
         # check, does this create a cuda event
         # making default stream waiting for NCCL stream ?
@@ -87,20 +92,20 @@ def run_ring_p2p(num_bytes: int):
             req.wait()
 
     if _ENABLE_PROFILE and RANK == 0:
-        assert ITERATION >= _PROFILE_STEPS
+        assert iteration >= profile_steps
         prof = torch.profiler.profile(
             activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True,
             schedule=torch.profiler.schedule(
-                wait=ITERATION - _PROFILE_STEPS,
+                wait=iteration - profile_steps,
                 warmup=0,
-                active=_PROFILE_STEPS,
+                active=profile_steps,
             ),
         )
         prof.start()
 
     start_event.record()
-    for it in range(ITERATION):
+    for it in range(iteration):
         reqs = send_recv_once(x, y)
         if _ENABLE_PROFILE and RANK == 0:
             prof.step()
@@ -110,7 +115,7 @@ def run_ring_p2p(num_bytes: int):
 
     end_event.record()
     end_event.synchronize()
-    avg_time_elapsed = start_event.elapsed_time(end_event) / ITERATION
+    avg_time_elapsed = start_event.elapsed_time(end_event) / iteration
     if _ENABLE_PROFILE and RANK == 0:
         prof.stop()
         prof.export_chrome_trace(f"inter-node_ring_p2p_trace_{num_bytes}.json")
@@ -177,7 +182,7 @@ def write_markdown(args, sizes_in_mb, time_statistics):
 """
 
 
-def run_inter_node_ring_p2p(args):
+def run_inter_node_ring_p2p(args, sizes_mb: Optional[Sequence[int]] = None):
     assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0
     num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE
 
@@ -190,16 +195,22 @@ def run_inter_node_ring_p2p(args):
         log("Skip inter node ring p2p benchmark")
         return
 
-    SIZES_IN_MB_TO_BENCH = [10, 20, 40, 80, 160]
+    if sizes_mb is None or len(sizes_mb) == 0:
+        sizes_mb = [10, 20, 40, 80, 160]
+    sizes_in_mb_to_bench = [int(mb) for mb in sizes_mb]
     time_statistics = []
-    for size_in_mb in SIZES_IN_MB_TO_BENCH:
+    for size_in_mb in sizes_in_mb_to_bench:
         avg_time_elapsed = run_ring_p2p(size_in_mb * (2**20))
 
         all_latency_results = [-1.0 for _ in range(WORLD_SIZE)]
         dist.gather_object(avg_time_elapsed, all_latency_results if RANK == 0 else None, dst=0)
         time_statistics.append(all_latency_results[:LOCAL_WORLD_SIZE])
 
-    write_markdown(args, SIZES_IN_MB_TO_BENCH, time_statistics)
+    write_markdown(args, sizes_in_mb_to_bench, time_statistics)
+
+    if _GLOBAL_PIPELINE_GROUP is not None:
+        dist.destroy_process_group(_GLOBAL_PIPELINE_GROUP)
+    barrier_after_comm_destroy(args.comm_cleanup_delay_sec)
 
     # TODO (limou)
     # support plot
diff --git a/primus/tools/preflight/intra_node_comm.py b/primus/tools/preflight/intra_node_comm.py
index 982f94b46..87eb97c4b 100644
--- a/primus/tools/preflight/intra_node_comm.py
+++ b/primus/tools/preflight/intra_node_comm.py
@@ -5,36 +5,69 @@
 ###############################################################################
 
 import time
+from typing import Iterable, List, Optional, Sequence
 
-import matplotlib.pyplot as plt
 import torch
 import torch.distributed as dist
 
 from primus.tools.preflight.global_vars import (
-    ITERATION,
     LOCAL_RANK,
     LOCAL_WORLD_SIZE,
     RANK,
-    WARMUP,
     WORLD_SIZE,
     get_hostnames,
+    get_iteration,
+    get_warmup,
 )
 from primus.tools.preflight.utility import (
+    barrier_after_comm_destroy,
     create_dir,
     extract_first_middle_last,
     extract_number,
+    format_int_range,
     log,
 )
 
 
-def run_intra_node_comm(args):
+def run_intra_node_comm(
+    args,
+    enabled_comms: Optional[Iterable[str]] = None,
+    sizes_mb: Optional[Sequence[int]] = None,
+    group_sizes: Optional[Sequence[int]] = None,
+):
+    """Intra-node allreduce / alltoall benchmark.
+
+    Args:
+        args: parsed namespace (must have markdown_file, dump_path, plot, ib_bw).
+        enabled_comms: subset of {"allreduce", "alltoall"} to run. Defaults to both.
+        sizes_mb: message sizes in MB. Defaults to powers of two from 2..1024.
+        group_sizes: GPU group sizes; each must divide LOCAL_WORLD_SIZE.
+            Defaults to [2, 4, 8].
+    """
     device = torch.device(f"cuda:{LOCAL_RANK}")
-    sizes = [2**i * 1024 * 1024 for i in range(1, 11)]
-    # sizes = [2**i * 1024 * 1024 for i in range(1, 5)]
-    cases = {
-        "allreduce": [2, 4, 8],
-        "alltoall": [2, 4, 8],
-    }
+
+    if sizes_mb is None or len(sizes_mb) == 0:
+        sizes_mb = [2**i for i in range(1, 11)]
+    sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb]
+
+    enabled_set = set(enabled_comms) if enabled_comms else {"allreduce", "alltoall"}
+    enabled_set &= {"allreduce", "alltoall"}
+    if not enabled_set:
+        log("Skip intra-node comm benchmark (no enabled comms)")
+        return
+
+    if group_sizes is None or len(group_sizes) == 0:
+        group_sizes = [2, 4, 8]
+    # Filter out invalid group sizes (must divide LOCAL_WORLD_SIZE) and de-dupe.
+    group_sizes = sorted({int(g) for g in group_sizes if int(g) > 0 and LOCAL_WORLD_SIZE % int(g) == 0})
+    if not group_sizes:
+        log(f"Skip intra-node comm benchmark, no valid group sizes for LOCAL_WORLD_SIZE={LOCAL_WORLD_SIZE}")
+        return
+
+    cases = {comm: list(group_sizes) for comm in ("allreduce", "alltoall") if comm in enabled_set}
+
+    warmup = get_warmup()
+    iteration = get_iteration()
 
     if RANK == 0:
         with open(args.markdown_file, "a", encoding="utf-8") as f:
@@ -54,6 +87,9 @@ def run_intra_node_comm(args):
             num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE
             num_groups_per_node = LOCAL_WORLD_SIZE // num_procs
             group = None
+            # Track which ranks belong to each group so we can render compact rows.
+            all_group_ranks: List[List[int]] = []
+            my_group_index = -1
             for i_node in range(num_nodes):
                 for i_group in range(num_groups_per_node):
                     group_ranks = [
@@ -63,12 +99,15 @@ def run_intra_node_comm(args):
                     if RANK in group_ranks:
                         assert group is None
                         group = tmp_group
+                        my_group_index = len(all_group_ranks)
+                    all_group_ranks.append(group_ranks)
             assert group is not None
+            assert my_group_index >= 0
 
             for size in sizes:
                 tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device)
                 dist.barrier(group=group, device_ids=[torch.cuda.current_device()])
-                for _ in range(WARMUP):
+                for _ in range(warmup):
                     if "allreduce" == comm:
                         dist.all_reduce(tensor, group=group)
                     elif "alltoall" == comm:
@@ -77,7 +116,7 @@ def run_intra_node_comm(args):
                         assert False
                 torch.cuda.synchronize()
                 start = time.time()
-                for _ in range(ITERATION):
+                for _ in range(iteration):
                     if "allreduce" == comm:
                         dist.all_reduce(tensor, group=group)
                     elif "alltoall" == comm:
@@ -85,7 +124,7 @@ def run_intra_node_comm(args):
                     else:
                         assert False
                 torch.cuda.synchronize()
-                elapsed = (time.time() - start) / ITERATION
+                elapsed = (time.time() - start) / iteration
                 scale = 2 if comm == "allreduce" else 1
                 comm_size = scale * size * (num_procs - 1) / num_procs
                 gb_per_sec = comm_size / elapsed / 1e9
@@ -96,6 +135,7 @@ def run_intra_node_comm(args):
 
             # destroy this parallel group
             dist.destroy_process_group(group)
+            barrier_after_comm_destroy(args.comm_cleanup_delay_sec)
 
             all_latency_results = [None for _ in range(WORLD_SIZE)]
             all_bandwidth_results = [None for _ in range(WORLD_SIZE)]
@@ -106,48 +146,65 @@ def run_intra_node_comm(args):
                 keys = sorted(
                     list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number
                 )
-                max_len = max(len(s) for s in get_hostnames()) + 2
+                hostnames = get_hostnames()
+
+                # Each row corresponds to one group (group_ranks). Use the first
+                # rank's results since all members observe the same collective.
+                # The Hostname column shows only the leader's host (compact);
+                # use the Node range plus the legend at the top of the report
+                # to look up the rest.
+                def _row_for(group_ranks: List[int], results):
+                    leader = group_ranks[0]
+                    host_str = hostnames[leader]
+                    node_str = format_int_range([r // LOCAL_WORLD_SIZE for r in group_ranks])
+                    rank_str = format_int_range(group_ranks)
+                    return host_str, node_str, rank_str, results[leader]
+
+                formatted_keys = [f"{key:<6}" for key in keys]
+                host_col_label = "Leader hostname"
+                host_col_w = max(20, len(host_col_label) + 2)
+                header_line = (
+                    f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}"
+                )
 
                 with open(args.markdown_file, "a", encoding="utf-8") as f:
                     f.write(f"=======IntraNodeComm - {case_name} (us)=======\n")
                     log(f"=======IntraNodeComm - {case_name} (us)=======")
+                    log(header_line)
 
-                    f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+                    f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
                     f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-                    formatted_keys = [f"{key:<6}" for key in keys]
-                    log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
-                    for rank, r in enumerate(all_latency_results):
-                        hostname = get_hostnames()[rank]
-                        if rank % num_procs != 0:
-                            continue
-                        node_id = rank // LOCAL_WORLD_SIZE
-
+                    for group_ranks in all_group_ranks:
+                        host_str, node_str, rank_str, r = _row_for(group_ranks, all_latency_results)
                         formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys]
-                        log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                        f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                        log(
+                            f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} "
+                            f"{' '.join(formatted_values)}"
+                        )
+                        f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
                     f.write(f"\n")
 
                     f.write(f"=======IntraNodeComm - {case_name} (GB/s)=======\n")
                     log(f"=======IntraNodeComm - {case_name} (GB/s)=======")
+                    log(header_line)
 
-                    f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n")
+                    f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n")
                     f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n")
-                    formatted_keys = [f"{key:<6}" for key in keys]
-                    log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}")
-                    for rank, r in enumerate(all_bandwidth_results):
-                        hostname = get_hostnames()[rank]
-                        if rank % num_procs != 0:
-                            continue
-                        node_id = rank // LOCAL_WORLD_SIZE
-
+                    for group_ranks in all_group_ranks:
+                        host_str, node_str, rank_str, r = _row_for(group_ranks, all_bandwidth_results)
                         formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys]
-                        log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}")
-                        f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n")
+                        log(
+                            f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} "
+                            f"{' '.join(formatted_values)}"
+                        )
+                        f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n")
                     f.write(f"\n")
 
                     if not args.plot:
                         continue
 
+                import matplotlib.pyplot as plt
+
                 log(f"=======Plot IntraNode {case_name} Bandwidth=======")
                 with open(args.markdown_file, "a", encoding="utf-8") as f:
                     f.write(f"=======Plot IntraNode {case_name} Bandwidth=======\n")
@@ -155,9 +212,9 @@ def run_intra_node_comm(args):
                 dump_path = f"{args.dump_path}/{plot_case}"
                 create_dir(dump_path)
                 print_keys = extract_first_middle_last(keys)
-                first_rank_bandwidth_results = [
-                    all_bandwidth_results[i] for i in range(len(all_bandwidth_results)) if i % num_procs == 0
-                ]
+                # Use leader rank of each group for plotting.
+                leader_ranks = [g[0] for g in all_group_ranks]
+                first_rank_bandwidth_results = [all_bandwidth_results[i] for i in leader_ranks]
                 num_print_ranks = len(first_rank_bandwidth_results)
                 for size_key in print_keys:
                     values = [r[size_key] for r in first_rank_bandwidth_results]
@@ -166,7 +223,7 @@ def run_intra_node_comm(args):
                     plt.xlabel(f"RankPair ({num_procs} ranks)")
                     plt.ylabel("Bandwidth")
                     plt.title(f"Intra Node {case_name} bandwidth for {size_key}")
-                    xtick_labels = [f"{i*num_procs}" for i in range(num_print_ranks)]
+                    xtick_labels = [f"{leader_ranks[i]}" for i in range(num_print_ranks)]
                     plt.xticks(range(num_print_ranks), xtick_labels)
                     plt.grid(True, axis="y")
 
diff --git a/primus/tools/preflight/node_smoke/__init__.py b/primus/tools/preflight/node_smoke/__init__.py
new file mode 100644
index 000000000..b4fce7177
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/__init__.py
@@ -0,0 +1,58 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Node-local preflight smoke test.
+
+Goal: quickly identify which nodes have problems before launching a real
+training job. Each node tests itself in parallel with no global rendezvous,
+writes a per-node JSON verdict, and an aggregator (rank 0) reads all JSONs
+and emits PASS/FAIL lists usable by SLURM ``--exclude=`` / ``--nodelist=``.
+
+Subcommands
+-----------
+
+* ``run`` -- per-node entry. Runs Tier 1 (per-GPU sanity + reused
+  host/gpu/network info collectors + a small dmesg scan) and, when
+  ``--tier2-perf`` is set, Tier 2 perf sanity (GEMM TFLOPS, HBM bandwidth,
+  and node-local RCCL all-reduce). Writes ``<dump>/smoke/<host>.json``.
+  Always exits 0 when the JSON was written (the per-node verdict lives
+  in the JSON's ``status`` field, not in the exit code) -- otherwise an
+  intentionally-detected unhealthy node would make srun pollute its
+  output with one ``error: ... task N: Exited with exit code 1`` per
+  failing node, which is misleading: the smoke test is succeeding at
+  identifying bad nodes, not failing.
+
+* ``aggregate`` -- read all per-node JSONs and emit
+  ``<dump>/smoke_report.md``, ``<dump>/passing_nodes.txt``, and
+  ``<dump>/failing_nodes.txt``. Exits non-zero if any node FAILs or is
+  missing -- this is the single CI-friendly cluster-health exit signal.
+
+* ``_per_gpu`` -- internal subcommand spawned by ``run`` to test a single
+  GPU in an isolated subprocess with a hard timeout. Not for direct use.
+
+Why per-GPU subprocesses?
+-------------------------
+
+A stuck ``torch.cuda.set_device(i)`` cannot be aborted reliably with
+``signal.alarm`` because the call may be inside a non-interruptible driver
+syscall. By running each per-GPU test in its own subprocess we can SIGKILL
+it on timeout without affecting the rest of the node's checks.
+
+Package layout
+--------------
+
+This module is a sub-package -- the implementation is split across many
+small files mirroring the Tier 1 A-G section structure. The single
+public entry point exported here is :func:`main`, which is also the
+target of ``python -m primus.tools.preflight.node_smoke``.
+"""
+
+from __future__ import annotations
+
+from .cli import main
+
+__all__ = ["main"]
diff --git a/primus/tools/preflight/node_smoke/__main__.py b/primus/tools/preflight/node_smoke/__main__.py
new file mode 100644
index 000000000..86cef62ab
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/__main__.py
@@ -0,0 +1,14 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Entry point for ``python -m primus.tools.preflight.node_smoke``."""
+
+from __future__ import annotations
+
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/primus/tools/preflight/node_smoke/aggregator/__init__.py b/primus/tools/preflight/node_smoke/aggregator/__init__.py
new file mode 100644
index 000000000..74a9229d5
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/aggregator/__init__.py
@@ -0,0 +1,18 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Aggregator -- combines per-node JSONs into a cluster report.
+
+* :mod:`.summarizers` -- pure data-shaping helpers used by the report
+  writer (one ``_*_rows`` / ``_*_summary`` function per Markdown
+  section). They never raise; missing data degrades to empty rows.
+* :mod:`.report`      -- the markdown writer. The single
+  :func:`.report.write_smoke_report` entry point composes the report
+  by calling small per-section ``_write_<section>`` helpers, in the
+  exact order the original monolithic block produced.
+"""
+
+from __future__ import annotations
diff --git a/primus/tools/preflight/node_smoke/aggregator/report.py b/primus/tools/preflight/node_smoke/aggregator/report.py
new file mode 100644
index 000000000..58a9ca95d
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/aggregator/report.py
@@ -0,0 +1,719 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Cluster smoke-report markdown writer.
+
+The single entry point :func:`write_smoke_report` composes the report by
+calling small per-section ``_write_<section>`` helpers in the exact
+order the original monolithic block produced. Behaviour is preserved
+verbatim, including:
+
+* every section header
+* the per-section ``try/except`` wrapping pattern
+* the two intentional EXCEPTIONS to that pattern -- ``Tier 2 perf
+  summary`` and ``Failing nodes -- full reasons`` -- which deliberately
+  do NOT have their own ``try/except`` in the original (so a regression
+  in either bubbles up rather than being swallowed). Per refactor
+  guardrails these stay un-wrapped.
+"""
+
+from __future__ import annotations
+
+from typing import IO, Any, Dict, List
+
+from ..logging_utils import _warn
+from .summarizers import (
+    _busy_gpu_rows,
+    _clock_summary,
+    _gpu_activity_rows,
+    _gpu_low_level_outlier_rows,
+    _host_limits_issue_rows,
+    _nic_excluded_rows,
+    _nic_fw_drift_rows,
+    _nic_issue_rows,
+    _pretouch_hbm_rows,
+    _stack_drift_rows,
+    _tooling_inventory_rows,
+    _tooling_latency_rows,
+    _xgmi_issue_rows,
+)
+
+
+def _write_header(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+    passing: List[Dict[str, Any]],
+    failing: List[Dict[str, Any]],
+    expected: Any,
+) -> None:
+    """Write the report title, summary counts, and per-node status table."""
+    f.write("# Node-Local Smoke Test Report\n\n")
+    f.write(f"- **Expected nodes**: `{expected if expected is not None else 'unknown'}`\n")
+    f.write(f"- **Reported nodes**: `{len(nodes)}`\n")
+    f.write(f"- **PASS**: `{len(passing)}`  **FAIL**: `{len(failing)}`\n\n")
+    f.write("| Node | Hostname | Status | Duration | Top fail reason |\n")
+    f.write("|------|----------|--------|----------|-----------------|\n")
+    for n in nodes:
+        reasons = n.get("fail_reasons") or []
+        top = (reasons[0] if reasons else "").replace("|", "/")
+        if len(top) > 120:
+            top = top[:117] + "..."
+        f.write(
+            f"| {n.get('node_rank', '?')} | {n.get('host', '?')} | "
+            f"{n.get('status', '?')} | {n.get('duration_sec', 0)}s | {top} |\n"
+        )
+
+
+def _write_stack_drift(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- A. Stack drift across cluster -----
+    # Empty section when every node reports the same value for every
+    # scalar fingerprint key. We always print the section header so the
+    # operator can see at a glance that the check ran.
+    # Each helper is wrapped so a single bug in one section can never
+    # truncate the whole report; the failure is recorded inline so the
+    # operator still sees something for that section.
+    f.write("\n## Stack drift across cluster\n\n")
+    try:
+        drift = _stack_drift_rows(nodes)
+        if not drift:
+            f.write("*All nodes match.*\n")
+        else:
+            f.write("| Key | Majority (count/total) | Outlier nodes |\n")
+            f.write("|------|-------------------------|----------------|\n")
+            for row in drift:
+                outliers = "; ".join(f"`{h}` = `{v}`" for h, v in row["outliers"])
+                f.write(
+                    f"| `{row['key']}` | `{row['majority']}` "
+                    f"({row['count']}/{row['total']}) | {outliers} |\n"
+                )
+    except Exception as e:
+        f.write(f"*Stack-drift section failed to render: {e}*\n")
+        _warn(f"stack-drift render failed: {e}")
+
+
+def _write_nic_fw_drift(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- A.2 NIC firmware drift across cluster -----
+    f.write("\n## NIC firmware drift across cluster\n\n")
+    try:
+        nic_drift = _nic_fw_drift_rows(nodes)
+        if not nic_drift:
+            f.write("*All NIC firmwares match (or no NICs reported).*\n")
+        else:
+            f.write("| NIC | Majority FW (count/total) | Outlier nodes |\n")
+            f.write("|-----|---------------------------|----------------|\n")
+            for row in nic_drift:
+                outliers = "; ".join(f"`{h}` = `{v}`" for h, v in row["outliers"])
+                f.write(
+                    f"| `{row['device']}` | `{row['majority']}` "
+                    f"({row['count']}/{row['total']}) | {outliers} |\n"
+                )
+    except Exception as e:
+        f.write(f"*NIC firmware drift section failed to render: {e}*\n")
+        _warn(f"nic-fw-drift render failed: {e}")
+
+
+def _write_nic_issues(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- B. NIC / RDMA roll-call issues -----
+    f.write("\n## NIC / RDMA roll-call issues\n\n")
+    try:
+        nic_issues = _nic_issue_rows(nodes)
+        if not nic_issues:
+            f.write("*No NIC issues.*\n")
+        else:
+            f.write("| Node | Hostname | Issue |\n")
+            f.write("|------|----------|-------|\n")
+            for row in nic_issues:
+                msg = str(row["issue"]).replace("|", "/")
+                if len(msg) > 160:
+                    msg = msg[:157] + "..."
+                f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n")
+    except Exception as e:
+        f.write(f"*NIC issues section failed to render: {e}*\n")
+        _warn(f"nic-issues render failed: {e}")
+
+
+def _write_nic_port_count(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- B.2 NIC port-count summary (helps spot "node X has fewer
+    # training NICs than the cluster") -- always rendered, even when no
+    # per-port issue tripped. We flag any node whose **included** port
+    # count differs from the cluster majority so operators can act on
+    # partial-degradation cases like 7/8 ports without having to set
+    # --expected-rdma-nics. The count is taken from `included_ports`
+    # (i.e. after the training-NIC selector ran) so nodes that legitimately
+    # have extra frontend / storage RoCE NICs don't show up as anomalies.
+    f.write("\n## NIC port-count summary\n\n")
+    try:
+        from collections import Counter
+
+        counts = []
+        for n in nodes:
+            nic = (n.get("tier1") or {}).get("nics") or {}
+            # Prefer the post-selector count; fall back to total ports
+            # for legacy JSONs written before --rdma-nic-allowlist
+            # existed (those don't have an `included_ports` key).
+            included = nic.get("included_ports")
+            if included is None:
+                count = len(nic.get("ports") or [])
+            else:
+                count = len(included)
+            counts.append(
+                (
+                    n.get("node_rank", "?"),
+                    n.get("host", "?"),
+                    count,
+                )
+            )
+        if not counts:
+            f.write("*No NIC data reported.*\n")
+        else:
+            cnt = Counter(c for *_, c in counts)
+            majority_count, _ = cnt.most_common(1)[0]
+            anomalies = [(nr, h, c) for nr, h, c in counts if c != majority_count]
+            f.write(
+                f"Cluster-majority training-NIC count: **{majority_count}** "
+                f"(seen on {cnt[majority_count]}/{len(counts)} nodes).\n\n"
+            )
+            if not anomalies:
+                f.write("*Every node reports the majority count.*\n")
+            else:
+                f.write("| Node | Hostname | Training NICs found |\n")
+                f.write("|------|----------|---------------------|\n")
+                for nr, h, c in anomalies:
+                    f.write(f"| {nr} | {h} | {c} |\n")
+    except Exception as e:
+        f.write(f"*NIC port-count summary failed to render: {e}*\n")
+        _warn(f"nic-port-count render failed: {e}")
+
+
+def _write_nic_excluded(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- B.3 Excluded NIC ports -- informational only. Surfaces ports
+    # that the selector chain dropped from the training-NIC set so the
+    # operator can verify the heuristic / NCCL_IB_HCA / --rdma-nic-allowlist
+    # did what they expected (e.g. "are my front-end NICs still
+    # admin-down?"). Excluded ports do NOT contribute to the node FAIL
+    # signal; they live in `tier1.nics.excluded_ports` + `info_issues`.
+    f.write("\n## NIC excluded ports (informational)\n\n")
+    try:
+        rows = _nic_excluded_rows(nodes)
+        if not rows:
+            f.write(
+                "*No NIC ports were excluded -- every discovered RDMA "
+                "port is in the training-NIC set on every node.*\n"
+            )
+        else:
+            # Brief per-source summary so the operator immediately knows
+            # whether to expect output here (e.g. "yes, NCCL_IB_HCA is
+            # filtering out 4 ports/node like I configured").
+            from collections import Counter
+
+            src_counts = Counter(r["source"] for r in rows)
+            src_bits = []
+            for src, c in sorted(src_counts.items()):
+                pretty = {
+                    "cli": "--rdma-nic-allowlist",
+                    "env": "NCCL_IB_HCA env",
+                    "heuristic": "phys_state heuristic",
+                }.get(src, src)
+                src_bits.append(f"{c} via `{pretty}`")
+            f.write(
+                "Ports excluded from the training-NIC set " + "(" + ", ".join(src_bits) + "). "
+                "Hard-fail rules (state ACTIVE, phys_state LinkUp, "
+                "RoCE v2 GIDs) did NOT run on these ports.\n\n"
+            )
+            f.write("| Node | Hostname | Source | Port + reason |\n")
+            f.write("|------|----------|--------|---------------|\n")
+            for row in rows:
+                msg = str(row["issue"]).replace("|", "/")
+                if len(msg) > 160:
+                    msg = msg[:157] + "..."
+                f.write(f"| {row['node_rank']} | {row['host']} | " f"`{row['source']}` | {msg} |\n")
+    except Exception as e:
+        f.write(f"*NIC excluded-ports section failed to render: {e}*\n")
+        _warn(f"nic-excluded render failed: {e}")
+
+
+def _write_host_limits(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- C. Host limits issues -----
+    f.write("\n## Host limits issues\n\n")
+    try:
+        limits_issues = _host_limits_issue_rows(nodes)
+        if not limits_issues:
+            f.write("*No host-limit issues.*\n")
+        else:
+            f.write("| Node | Hostname | Issue |\n")
+            f.write("|------|----------|-------|\n")
+            for row in limits_issues:
+                msg = str(row["issue"]).replace("|", "/")
+                if len(msg) > 200:
+                    msg = msg[:197] + "..."
+                f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n")
+    except Exception as e:
+        f.write(f"*Host limits section failed to render: {e}*\n")
+        _warn(f"host-limits render failed: {e}")
+
+
+def _write_gpu_visibility(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- GPU visibility issues (no GPUs / amd-smi vs torch mismatch) -----
+    # Independent guard -- doesn't rely on the reused gpu_info collector
+    # emitting a level=fail finding, which has been known to silently
+    # downgrade to warn when collect_gpu_info() raises.
+    f.write("\n## GPU visibility issues\n\n")
+    try:
+        vis_rows: List[Dict[str, Any]] = []
+        for n in nodes:
+            vis = (n.get("tier1") or {}).get("gpu_visibility") or {}
+            for issue in vis.get("fail_reasons", []) or []:
+                vis_rows.append(
+                    {
+                        "node_rank": n.get("node_rank", "?"),
+                        "host": n.get("host", "?"),
+                        "torch": vis.get("torch_visible"),
+                        "amd_smi": vis.get("amd_smi_visible"),
+                        "expected": vis.get("expected_gpus"),
+                        "issue": issue,
+                    }
+                )
+        if not vis_rows:
+            f.write(
+                "*Every node resolved expected_gpus >= 1 and torch + " "amd-smi agree on the GPU count.*\n"
+            )
+        else:
+            f.write(
+                "Nodes where the GPU is invisible to torch, or where "
+                "amd-smi sees more GPUs than torch (stale ROCm / wedged "
+                "amdgpu driver). These are hard fails independent of "
+                "every other collector.\n\n"
+            )
+            f.write("| Node | Hostname | expected | torch | amd-smi | Issue |\n")
+            f.write("|------|----------|----------|-------|---------|-------|\n")
+            for row in vis_rows:
+                msg = str(row["issue"]).replace("|", "/")
+                if len(msg) > 200:
+                    msg = msg[:197] + "..."
+                f.write(
+                    f"| {row['node_rank']} | {row['host']} | "
+                    f"{row['expected']} | {row['torch']} | "
+                    f"{row['amd_smi']} | {msg} |\n"
+                )
+    except Exception as e:
+        f.write(f"*GPU visibility section failed to render: {e}*\n")
+        _warn(f"gpu-visibility render failed: {e}")
+
+
+def _write_gpu_low_level(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- D-1: GPU low-level outliers (PCIe link, HBM total) -----
+    f.write("\n## GPU low-level outliers (PCIe link / HBM)\n\n")
+    try:
+        gpu_outliers = _gpu_low_level_outlier_rows(nodes)
+        if not gpu_outliers:
+            f.write("*All GPUs match the cluster majority on PCIe link " "and HBM total.*\n")
+        else:
+            f.write(
+                "Per-GPU values that differ from the cluster majority. A "
+                "GPU sitting at half PCIe width / half HBM is almost "
+                "always a hardware fault on that single device.\n\n"
+            )
+            f.write("| Metric | Cluster majority (count/total) | " "Outliers (`host:gpu` = value) |\n")
+            f.write("|--------|---------------------------------|" "-------------------------------|\n")
+            for row in gpu_outliers:
+                out_str = "; ".join(f"`{h}:{g}` = `{v}`" for h, g, v in row["outliers"])
+                f.write(
+                    f"| {row['label']} | `{row['majority']}` "
+                    f"({row['count']}/{row['total']}) | {out_str} |\n"
+                )
+    except Exception as e:
+        f.write(f"*GPU low-level section failed to render: {e}*\n")
+        _warn(f"gpu-low-level render failed: {e}")
+
+
+def _write_xgmi(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- D-2: XGMI link issues -----
+    f.write("\n## XGMI link issues\n\n")
+    try:
+        xgmi_issues = _xgmi_issue_rows(nodes)
+        if not xgmi_issues:
+            f.write("*All GPU pairs report XGMI on every node " "(or amd-smi topology was unavailable).*\n")
+        else:
+            f.write(
+                "Any non-XGMI GPU pair is a hard fail -- intra-node "
+                "collectives silently fall back to PCIe and lose 5-10x "
+                "of the bandwidth NCCL/RCCL expects.\n\n"
+            )
+            f.write("| Node | Hostname | Issue |\n")
+            f.write("|------|----------|-------|\n")
+            for row in xgmi_issues:
+                msg = str(row["summary"]).replace("|", "/")
+                if len(msg) > 200:
+                    msg = msg[:197] + "..."
+                f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n")
+    except Exception as e:
+        f.write(f"*XGMI section failed to render: {e}*\n")
+        _warn(f"xgmi render failed: {e}")
+
+
+def _write_clock(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+    skew_warn_sec: float,
+) -> None:
+    # ----- E: cluster wall-clock spread + time-daemon roll-call -----
+    f.write("\n## Cluster clock + time daemons\n\n")
+    try:
+        clk = _clock_summary(nodes, skew_warn_sec=skew_warn_sec)
+        spread = clk["spread_sec"]
+        if spread is None:
+            f.write("*Not enough nodes reported a wall-clock timestamp.*\n")
+        else:
+            marker = " (**warn** -- exceeds " f"{clk['spread_warn_sec']}s)" if clk["spread_warn"] else ""
+            f.write(
+                f"- Wall-clock spread across {clk['n_nodes_with_time']} " f"nodes: **{spread}s**{marker}.\n"
+            )
+            f.write(f"- Earliest: `{clk['earliest_host']}`, " f"latest: `{clk['latest_host']}`.\n")
+            f.write(
+                "- (Spread is an upper bound on real clock skew -- it " "also includes srun launch jitter.)\n"
+            )
+        if clk["no_daemon_hosts"]:
+            f.write(
+                "\n**Nodes with no active time-sync daemon " "(chronyd / ntpd / systemd-timesyncd):**\n\n"
+            )
+            f.write("| Node | Hostname |\n")
+            f.write("|------|----------|\n")
+            for nr, h in clk["no_daemon_hosts"]:
+                f.write(f"| {nr} | {h} |\n")
+        else:
+            f.write("\n*Every node has at least one active time-sync " "daemon.*\n")
+    except Exception as e:
+        f.write(f"*Clock section failed to render: {e}*\n")
+        _warn(f"clock render failed: {e}")
+
+
+def _write_tooling_latency(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+    rocm_smi_warn_sec: float,
+) -> None:
+    # ----- F-partial: rocm-smi self-latency -----
+    f.write("\n## Tooling self-latency (`rocm-smi --version`)\n\n")
+    try:
+        tool_rows = _tooling_latency_rows(
+            nodes,
+            warn_sec=float(rocm_smi_warn_sec),
+        )
+        if not tool_rows:
+            f.write("*No nodes exceeded the warn threshold " f"({rocm_smi_warn_sec}s) and no timeouts.*\n")
+        else:
+            f.write(
+                "Slow `rocm-smi --version` calls historically precede a "
+                "wedged amdgpu driver. Hitting the hard timeout is a "
+                "node FAIL; slow-but-completed calls are warn-only.\n\n"
+            )
+            f.write("| Node | Hostname | Latency (s) | Flag |\n")
+            f.write("|------|----------|-------------|------|\n")
+            for r in tool_rows:
+                lat = r.get("latency_sec")
+                lat_s = f"{lat:.2f}" if isinstance(lat, (int, float)) else "?"
+                f.write(f"| {r['node_rank']} | {r['host']} | " f"{lat_s} | {r['flag']} |\n")
+    except Exception as e:
+        f.write(f"*Tooling section failed to render: {e}*\n")
+        _warn(f"tooling render failed: {e}")
+
+
+def _write_tooling_availability(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+) -> None:
+    # ----- Tooling availability (always-on; loud counterweight to
+    # the silent skips that happen when amd-smi / rocm-smi / lsof
+    # are missing from PATH) -----
+    f.write("\n## Tooling availability\n\n")
+    try:
+        inv = _tooling_inventory_rows(nodes)
+        tracked = inv["tracked"]
+        mc = inv["missing_counts"]
+        if not inv["any_missing"]:
+            f.write(
+                "*Every tracked tool (`" + "`, `".join(tracked) + "`) was present in PATH on every node.*\n"
+            )
+        else:
+            summary_bits = []
+            for t in tracked:
+                if mc[t]:
+                    summary_bits.append(f"`{t}` missing on **{mc[t]}** node(s)")
+            # Compute cluster-wide uncovered-checks summary so the
+            # operator immediately knows whether the missing tools
+            # actually leave a coverage hole or whether the rocm-smi
+            # / lsof fallbacks are picking up the slack.
+            uncovered_counts: Dict[str, int] = {}
+            for n in nodes:
+                uc = ((n.get("tier1") or {}).get("tooling_inventory") or {}).get("uncovered") or []
+                for c in uc:
+                    uncovered_counts[c] = uncovered_counts.get(c, 0) + 1
+            if uncovered_counts:
+                uc_bits = [f"`{c}` on **{n}** node(s)" for c, n in sorted(uncovered_counts.items())]
+                coverage_line = (
+                    "Checks with NO working tool (truly silent-skipped): " + "; ".join(uc_bits) + "."
+                )
+            else:
+                coverage_line = (
+                    "Every check is still covered via the rocm-smi "
+                    "or lsof fallback on every node -- no checks are "
+                    "silently skipped."
+                )
+            f.write(
+                "Several Tier 1 checks (ECC, XGMI, foreign-process, "
+                "GPU activity, wedged-driver) prefer `amd-smi` but "
+                "fall back to `rocm-smi` (and `lsof` for foreign-"
+                "process) when amd-smi is missing. "
+                + "; ".join(summary_bits)
+                + ". "
+                + coverage_line
+                + " Add `--require-tools amd-smi,rocm-smi` to `run` to "
+                "promote a missing tool to a node FAIL anyway.\n\n"
+            )
+            # Per-node table -- only the nodes that ARE missing something,
+            # so a healthy cluster doesn't get a giant N-row table.
+            f.write("| Node | Hostname | " + " | ".join(tracked) + " |\n")
+            f.write("|------|----------| " + " | ".join("---" for _ in tracked) + " |\n")
+            for r in inv["rows"]:
+                if all(r.get(t) for t in tracked):
+                    continue
+                cells = []
+                for t in tracked:
+                    cells.append("OK" if r.get(t) else "**MISSING**")
+                f.write(f"| {r['node_rank']} | {r['host']} | " + " | ".join(cells) + " |\n")
+    except Exception as e:
+        f.write(f"*Tooling availability section failed to render: {e}*\n")
+        _warn(f"tooling-availability render failed: {e}")
+
+
+def _write_busy_gpus(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    # ----- G: Busy GPUs / leaked processes -----
+    f.write("\n## Busy GPUs / leaked processes\n\n")
+    try:
+        busy_rows = _busy_gpu_rows(nodes)
+        if not busy_rows:
+            f.write(
+                "*No foreign processes detected on any GPU "
+                "(or `amd-smi process` was unavailable on every node).*\n"
+            )
+        else:
+            f.write(
+                "Foreign PIDs found holding GPUs at smoke start. The most "
+                "common cause is leaked Python ranks from a previous "
+                "training job (look for `python` / `torchrun` / `train.py`). "
+                "Clean up with `pkill -9 -f train.py` (or similar) on the "
+                "listed nodes BEFORE launching the next job.\n\n"
+            )
+            f.write("| Node | Hostname | GPU | PID | Process | HBM held (GiB) |\n")
+            f.write("|------|----------|-----|-----|---------|----------------|\n")
+            for r in busy_rows:
+                name = str(r.get("name", "")).replace("|", "/")[:40]
+                hbm = r.get("hbm_gib")
+                hbm_s = f"{hbm}" if hbm is not None else "?"
+                f.write(
+                    f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['pid']} | `{name}` | {hbm_s} |\n"
+                )
+    except Exception as e:
+        f.write(f"*Busy-GPU section failed to render: {e}*\n")
+        _warn(f"busy-gpu render failed: {e}")
+
+
+def _write_pretouch_hbm(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+    hbm_busy_threshold_gib: float,
+) -> None:
+    # ----- G: Pre-touch HBM-used outliers -----
+    f.write("\n## GPU pre-touch HBM usage outliers\n\n")
+    try:
+        threshold = float(hbm_busy_threshold_gib)
+        pt_rows = _pretouch_hbm_rows(nodes, threshold_gib=threshold)
+        if not pt_rows:
+            f.write(
+                f"*No GPU reached the pre-touch HBM threshold "
+                f"({threshold} GiB or more) -- every GPU started clean.*\n"
+            )
+        else:
+            f.write(
+                f"GPUs with **at least {threshold} GiB** of HBM already "
+                f"in use BEFORE smoke touched the device. This number is "
+                "not polluted by our own caching allocator (it's measured "
+                "before any allocation), so it directly reflects foreign "
+                "or leaked occupancy.\n\n"
+            )
+            f.write("| Node | Hostname | GPU | HBM used pre-touch (GiB) |\n")
+            f.write("|------|----------|-----|---------------------------|\n")
+            for r in pt_rows:
+                f.write(f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['used_gib']} |\n")
+    except Exception as e:
+        f.write(f"*Pre-touch HBM section failed to render: {e}*\n")
+        _warn(f"pretouch-hbm render failed: {e}")
+
+
+def _write_gpu_activity(
+    f: IO[str],
+    nodes: List[Dict[str, Any]],
+    gpu_activity_warn_pct: float,
+) -> None:
+    # ----- G: GPU compute activity outliers -----
+    f.write("\n## GPU compute-activity outliers\n\n")
+    try:
+        warn_pct = float(gpu_activity_warn_pct)
+        act_rows = _gpu_activity_rows(nodes, warn_pct=warn_pct)
+        if not act_rows:
+            f.write(
+                f"*No GPU exceeded `gfx_activity_pct >= {warn_pct}%` at "
+                "smoke start (or amd-smi did not report activity).*\n"
+            )
+        else:
+            f.write(
+                f"GPUs reporting **>= {warn_pct}%** compute activity at "
+                "smoke start. Short bursts are normal; sustained "
+                "non-trivial activity across multiple GPUs strongly "
+                "suggests a leaked rank still running compute. Warn-only "
+                "(does not by itself fail the node).\n\n"
+            )
+            f.write("| Node | Hostname | GPU | Activity % |\n")
+            f.write("|------|----------|-----|------------|\n")
+            for r in act_rows:
+                f.write(f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['activity_pct']} |\n")
+    except Exception as e:
+        f.write(f"*Activity section failed to render: {e}*\n")
+        _warn(f"activity render failed: {e}")
+
+
+def _write_tier2_perf_summary(f: IO[str], nodes: List[Dict[str, Any]]) -> None:
+    """Write the Tier 2 perf summary section.
+
+    INTENTIONALLY UN-WRAPPED in try/except (matching the original
+    monolithic block): a regression in this loop should bubble up rather
+    than be silently swallowed. Only emitted when at least one node
+    actually ran Tier 2.
+    """
+    # Tier 2 perf summary -- only emitted when at least one node ran Tier 2.
+    # Surfaces per-node GEMM TFLOPS / HBM GB/s (min/median/max across the
+    # node's GPUs) plus the local RCCL all-reduce GB/s, so outliers across
+    # the cluster are visible without opening every per-node JSON.
+    perf_rows: List[str] = []
+    any_tier2 = False
+    for n in nodes:
+        t2 = n.get("tier2") or {}
+        per_gpu = (n.get("tier1") or {}).get("per_gpu") or []
+        gemm = [
+            p.get("details", {}).get("gemm_tflops")
+            for p in per_gpu
+            if isinstance(p.get("details", {}).get("gemm_tflops"), (int, float))
+        ]
+        hbm = [
+            p.get("details", {}).get("hbm_gbs")
+            for p in per_gpu
+            if isinstance(p.get("details", {}).get("hbm_gbs"), (int, float))
+        ]
+        rccl_gbs = (t2.get("rccl") or {}).get("gbs")
+        if not gemm and not hbm and rccl_gbs is None:
+            perf_rows.append(f"| {n.get('node_rank', '?')} | {n.get('host', '?')} |  |  |  |")
+            continue
+        any_tier2 = True
+
+        def _fmt_stats(xs):
+            if not xs:
+                return ""
+            xs_sorted = sorted(xs)
+            med = xs_sorted[len(xs_sorted) // 2]
+            return f"{min(xs):.1f} / {med:.1f} / {max(xs):.1f}"
+
+        perf_rows.append(
+            f"| {n.get('node_rank', '?')} | {n.get('host', '?')} | "
+            f"{_fmt_stats(gemm)} | {_fmt_stats(hbm)} | "
+            f"{rccl_gbs if rccl_gbs is not None else ''} |"
+        )
+
+    if any_tier2:
+        f.write("\n## Tier 2 perf summary\n\n")
+        f.write(
+            "Per-node GEMM TFLOPS (8192^3 bf16) and HBM GB/s shown as "
+            "`min / median / max` across the node's GPUs. RCCL GB/s is the "
+            "node-local 8-GPU all-reduce algorithmic bandwidth at 64 MB.\n\n"
+        )
+        f.write(
+            "| Node | Hostname | GEMM TFLOPS (min/med/max) | " "HBM GB/s (min/med/max) | Local RCCL GB/s |\n"
+        )
+        f.write(
+            "|------|----------|----------------------------|"
+            "------------------------|------------------|\n"
+        )
+        for r in perf_rows:
+            f.write(r + "\n")
+
+
+def _write_failing_reasons(
+    f: IO[str],
+    failing: List[Dict[str, Any]],
+) -> None:
+    """Write the per-node fail-reason dump for failing nodes.
+
+    INTENTIONALLY UN-WRAPPED in try/except (matching the original
+    monolithic block): if iterating fail_reasons raises, surface it
+    instead of swallowing it.
+    """
+    if failing:
+        f.write("\n## Failing nodes -- full reasons\n\n")
+        for n in failing:
+            f.write(f"### {n.get('host', '?')}\n\n")
+            for r in n.get("fail_reasons") or []:
+                f.write(f"- {r}\n")
+            f.write("\n")
+
+
+def write_smoke_report(
+    report_path: str,
+    *,
+    nodes: List[Dict[str, Any]],
+    passing: List[Dict[str, Any]],
+    failing: List[Dict[str, Any]],
+    expected: Any,
+    clock_skew_warn_sec: float,
+    rocm_smi_warn_sec: float,
+    hbm_busy_threshold_gib: float,
+    gpu_activity_warn_pct: float,
+) -> None:
+    """Write the cluster smoke report to ``report_path``.
+
+    Section ORDER and HEADINGS are part of the operator-facing contract
+    -- many CI/CD pipelines and slack bots scrape ``smoke_report.md`` for
+    specific ``##`` headings. Do NOT reorder or rename without a deliberate
+    behavior change. The order matches the original monolithic writer
+    (header, A, A.2, B, B.2, B.3, C, GPU visibility, D-1, D-2, E,
+    F-partial, Tooling availability, G x3, Tier 2 perf summary
+    [conditional], Failing nodes [conditional]).
+    """
+    with open(report_path, "w", encoding="utf-8") as f:
+        _write_header(f, nodes, passing, failing, expected)
+        _write_stack_drift(f, nodes)
+        _write_nic_fw_drift(f, nodes)
+        _write_nic_issues(f, nodes)
+        _write_nic_port_count(f, nodes)
+        _write_nic_excluded(f, nodes)
+        _write_host_limits(f, nodes)
+        _write_gpu_visibility(f, nodes)
+        _write_gpu_low_level(f, nodes)
+        _write_xgmi(f, nodes)
+        _write_clock(f, nodes, skew_warn_sec=clock_skew_warn_sec)
+        _write_tooling_latency(f, nodes, rocm_smi_warn_sec=rocm_smi_warn_sec)
+        _write_tooling_availability(f, nodes)
+        _write_busy_gpus(f, nodes)
+        _write_pretouch_hbm(
+            f,
+            nodes,
+            hbm_busy_threshold_gib=hbm_busy_threshold_gib,
+        )
+        _write_gpu_activity(
+            f,
+            nodes,
+            gpu_activity_warn_pct=gpu_activity_warn_pct,
+        )
+        _write_tier2_perf_summary(f, nodes)
+        _write_failing_reasons(f, failing)
diff --git a/primus/tools/preflight/node_smoke/aggregator/summarizers.py b/primus/tools/preflight/node_smoke/aggregator/summarizers.py
new file mode 100644
index 000000000..8087c1653
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/aggregator/summarizers.py
@@ -0,0 +1,449 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Pure data-shaping helpers used by the cluster-level report writer.
+
+Every function takes the loaded per-node JSON list (``nodes``) and
+returns plain rows or summary dicts. None of them format markdown -- the
+report writer in :mod:`.report` is solely responsible for layout.
+Each helper is best-effort: missing data degrades to an empty list rather
+than raising.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from ..collectors.tooling import _TRACKED_TOOLS
+
+# ---------------------------------------------------------------------------
+# Aggregator helpers -- A. stack/NIC drift, B. NIC issues, C. host limits
+# ---------------------------------------------------------------------------
+
+
+def _stack_drift_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """For every *scalar* fingerprint key, find the cluster-majority value and
+    list the nodes that disagree.
+
+    Returns one row per key that has at least one outlier. Keys missing from
+    every node, or where every node reported the same value, are omitted so
+    a healthy cluster produces an empty list.
+    """
+    from collections import Counter
+
+    # Only collect keys that at least ONE node reported as a scalar. We
+    # ignore None here so a key that happens to be None on one node and a
+    # dict on another (e.g. nic_fw on a node without an IB stack) doesn't
+    # leak into the scalar-drift loop and crash Counter() with an unhashable
+    # value.
+    keys: set = set()
+    for n in nodes:
+        fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {}
+        for k, v in fp.items():
+            if isinstance(v, (str, int, float)):
+                keys.add(k)
+
+    rows: List[Dict[str, Any]] = []
+    for k in sorted(keys):
+        per_host: List[tuple] = []
+        for n in nodes:
+            fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {}
+            v = fp.get(k)
+            # Defense in depth: skip non-scalar values per-host too, in case
+            # different nodes disagree on the type for the same key.
+            if not isinstance(v, (str, int, float)):
+                continue
+            per_host.append((n.get("host", "?"), v))
+        if not per_host:
+            continue
+        c = Counter(v for _, v in per_host)
+        majority, count = c.most_common(1)[0]
+        outliers = [(h, v) for h, v in per_host if v != majority]
+        if not outliers:
+            continue
+        rows.append(
+            {
+                "key": k,
+                "majority": majority,
+                "count": count,
+                "total": len(per_host),
+                "outliers": outliers,
+            }
+        )
+    return rows
+
+
+def _nic_fw_drift_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-IB-device firmware drift across the cluster (e.g. rdma0 mismatch)."""
+    from collections import Counter
+
+    all_devs: set = set()
+    for n in nodes:
+        fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {}
+        all_devs.update((fp.get("nic_fw") or {}).keys())
+
+    rows: List[Dict[str, Any]] = []
+    for dev in sorted(all_devs):
+        per_host: List[tuple] = []
+        for n in nodes:
+            fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {}
+            v = (fp.get("nic_fw") or {}).get(dev)
+            if v is None:
+                continue
+            per_host.append((n.get("host", "?"), v))
+        if not per_host:
+            continue
+        c = Counter(v for _, v in per_host)
+        majority, count = c.most_common(1)[0]
+        outliers = [(h, v) for h, v in per_host if v != majority]
+        if not outliers:
+            continue
+        rows.append(
+            {
+                "device": dev,
+                "majority": majority,
+                "count": count,
+                "total": len(per_host),
+                "outliers": outliers,
+            }
+        )
+    return rows
+
+
+def _nic_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-node NIC roll-call issues (port DOWN / no GIDs / count mismatch)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        nic = (n.get("tier1") or {}).get("nics") or {}
+        for issue in nic.get("issues", []) or []:
+            rows.append(
+                {
+                    "node_rank": n.get("node_rank", "?"),
+                    "host": n.get("host", "?"),
+                    "issue": issue,
+                }
+            )
+    return rows
+
+
+def _nic_excluded_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-node NIC ports that were excluded from the training-NIC set.
+
+    Excluded ports are informational: they did not contribute to the
+    node FAIL signal, but operators sometimes want to see them
+    (e.g. "did the heuristic do what I expected?", "is my front-end
+    NIC still disabled or did it come back?"). Each row carries the
+    selector source so the operator can tell whether the exclusion
+    came from --rdma-nic-allowlist, NCCL_IB_HCA, or the
+    phys_state=Disabled/Sleep heuristic.
+    """
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        nic = (n.get("tier1") or {}).get("nics") or {}
+        sel = nic.get("selector") or {}
+        source = sel.get("source") or "?"
+        for issue in nic.get("info_issues", []) or []:
+            rows.append(
+                {
+                    "node_rank": n.get("node_rank", "?"),
+                    "host": n.get("host", "?"),
+                    "source": source,
+                    "issue": issue,
+                }
+            )
+    return rows
+
+
+def _host_limits_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-node host-limit hard violations (ulimit -l / /dev/shm too low)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        hl = (n.get("tier1") or {}).get("host_limits") or {}
+        for issue in hl.get("fail_reasons", []) or []:
+            rows.append(
+                {
+                    "node_rank": n.get("node_rank", "?"),
+                    "host": n.get("host", "?"),
+                    "issue": issue,
+                }
+            )
+    return rows
+
+
+# ---------------------------------------------------------------------------
+# Aggregator helpers -- D-1 / D-2 / E / F
+# ---------------------------------------------------------------------------
+
+
+def _gpu_low_level_outlier_rows(
+    nodes: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """Find per-GPU outliers in PCIe link + HBM total across the cluster.
+
+    For each scalar metric the cluster has a strong majority value (e.g.
+    16 lanes, 32 GT/s, 191 GiB HBM). A single GPU below the majority on
+    any of these is almost always a hardware issue -- a cold-soldered
+    socket, a degraded PCIe link, or HBM that the firmware refused to
+    bring online. We surface every such (host, gpu, metric, value)
+    tuple, with the cluster majority for context.
+
+    Power cap and ECC counters from amd-smi are intentionally NOT included
+    here; they have their own narrower checks (ECC = hard fail in
+    ``_node_status_from``; power cap = informational only because cluster
+    operators sometimes set per-rack caps deliberately).
+    """
+    from collections import Counter
+
+    fields = (
+        ("pcie_link_width", "PCIe width (lanes)"),
+        ("pcie_link_speed_gts", "PCIe speed (GT/s)"),
+        ("hbm_total_gib", "HBM total (GiB)"),
+    )
+    rows: List[Dict[str, Any]] = []
+    for key, label in fields:
+        per_gpu: List[tuple] = []  # (host, gpu_idx, value)
+        for n in nodes:
+            for p in (n.get("tier1") or {}).get("per_gpu") or []:
+                low = (p.get("details") or {}).get("low_level") or {}
+                v = low.get(key)
+                if isinstance(v, (int, float)):
+                    per_gpu.append((n.get("host", "?"), p.get("gpu", "?"), v))
+        if not per_gpu:
+            continue
+        c = Counter(v for _, _, v in per_gpu)
+        majority, count = c.most_common(1)[0]
+        outliers = [(h, g, v) for h, g, v in per_gpu if v != majority]
+        if not outliers:
+            continue
+        rows.append(
+            {
+                "key": key,
+                "label": label,
+                "majority": majority,
+                "count": count,
+                "total": len(per_gpu),
+                "outliers": outliers,
+            }
+        )
+    return rows
+
+
+def _xgmi_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-node XGMI link issues (any non-XGMI GPU pair)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        xg = (n.get("tier1") or {}).get("xgmi") or {}
+        if not xg.get("ok"):
+            err = xg.get("error")
+            if err:
+                rows.append(
+                    {
+                        "node_rank": n.get("node_rank", "?"),
+                        "host": n.get("host", "?"),
+                        "summary": f"could not collect topology: {err}",
+                    }
+                )
+            continue
+        bad = xg.get("non_xgmi_pairs") or []
+        if not bad:
+            continue
+        # Show up to 6 sample pairs to keep the table readable; the full
+        # matrix lives in the per-node JSON.
+        sample = ", ".join(f"({i},{j})={t}" for i, j, t in bad[:6])
+        suffix = "" if len(bad) <= 6 else f" (+{len(bad) - 6} more)"
+        rows.append(
+            {
+                "node_rank": n.get("node_rank", "?"),
+                "host": n.get("host", "?"),
+                "summary": f"{len(bad)} non-XGMI pair(s): {sample}{suffix}",
+            }
+        )
+    return rows
+
+
+def _clock_summary(
+    nodes: List[Dict[str, Any]],
+    skew_warn_sec: float,
+) -> Dict[str, Any]:
+    """Compute wall-clock spread + per-node time-daemon health."""
+    times: List[tuple] = []  # (host, wall_time_unix)
+    no_daemon_hosts: List[tuple] = []  # (node_rank, host)
+    for n in nodes:
+        clk = (n.get("tier1") or {}).get("clock") or {}
+        wt = clk.get("wall_time_unix")
+        if isinstance(wt, (int, float)):
+            times.append((n.get("host", "?"), float(wt)))
+        if clk and not clk.get("any_active", True):
+            no_daemon_hosts.append((n.get("node_rank", "?"), n.get("host", "?")))
+
+    spread_sec = None
+    earliest_h = latest_h = None
+    if len(times) >= 2:
+        earliest_h, earliest = min(times, key=lambda x: x[1])
+        latest_h, latest = max(times, key=lambda x: x[1])
+        spread_sec = round(latest - earliest, 3)
+    return {
+        "n_nodes_with_time": len(times),
+        "spread_sec": spread_sec,
+        "spread_warn_sec": skew_warn_sec,
+        "spread_warn": (spread_sec is not None and spread_sec > skew_warn_sec),
+        "earliest_host": earliest_h,
+        "latest_host": latest_h,
+        "no_daemon_hosts": no_daemon_hosts,
+    }
+
+
+def _tooling_latency_rows(
+    nodes: List[Dict[str, Any]],
+    warn_sec: float,
+) -> List[Dict[str, Any]]:
+    """Per-node `rocm-smi --version` self-latency outliers (timed-out + slow)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        t = (n.get("tier1") or {}).get("tooling") or {}
+        lat = t.get("latency_sec")
+        timed_out = bool(t.get("timed_out"))
+        if t.get("error") and lat is None:
+            # Tool missing -- not interesting for a slow-tool report.
+            continue
+        flag = ""
+        if timed_out:
+            flag = "TIMEOUT"
+        elif isinstance(lat, (int, float)) and lat > warn_sec:
+            flag = f">{warn_sec}s"
+        if not flag:
+            continue
+        rows.append(
+            {
+                "node_rank": n.get("node_rank", "?"),
+                "host": n.get("host", "?"),
+                "latency_sec": lat,
+                "flag": flag,
+                "timeout_sec": t.get("timeout_sec"),
+            }
+        )
+    return rows
+
+
+def _tooling_inventory_rows(nodes: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Per-node tooling-presence rows + a count of nodes missing each tool.
+
+    Returns::
+
+        {
+            "rows": [
+                {"node_rank": 0, "host": "node001",
+                 "amd-smi": True, "rocm-smi": True, "lsof": True},
+                ...
+            ],
+            "missing_counts": {"amd-smi": 0, "rocm-smi": 1, "lsof": 0},
+            "any_missing": True,
+            "tracked": ["amd-smi", "rocm-smi", "lsof"],
+        }
+
+    Always-on: even when every tool is present everywhere, we still emit
+    a (small, reassuring) summary so the operator can see at a glance
+    that the toolchain was healthy on every node.
+    """
+    tracked = list(_TRACKED_TOOLS)
+    rows: List[Dict[str, Any]] = []
+    missing_counts: Dict[str, int] = {t: 0 for t in tracked}
+    for n in nodes:
+        inv = (n.get("tier1") or {}).get("tooling_inventory") or {}
+        tools = inv.get("tools") or {}
+        row: Dict[str, Any] = {
+            "node_rank": n.get("node_rank", "?"),
+            "host": n.get("host", "?"),
+        }
+        for t in tracked:
+            present = bool((tools.get(t) or {}).get("present"))
+            row[t] = present
+            if not present:
+                missing_counts[t] += 1
+        rows.append(row)
+    return {
+        "rows": rows,
+        "missing_counts": missing_counts,
+        "any_missing": any(c > 0 for c in missing_counts.values()),
+        "tracked": tracked,
+    }
+
+
+def _busy_gpu_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Per-node foreign GPU process listing (for the "Busy GPUs" section)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        gp = (n.get("tier1") or {}).get("gpu_processes") or {}
+        if not gp.get("ok"):
+            continue
+        per_gpu = gp.get("per_gpu") or []
+        for g in per_gpu:
+            for p in g.get("processes") or []:
+                if not p.get("is_foreign"):
+                    continue
+                hbm_b = p.get("hbm_bytes")
+                rows.append(
+                    {
+                        "node_rank": n.get("node_rank", "?"),
+                        "host": n.get("host", "?"),
+                        "gpu": g.get("gpu", "?"),
+                        "pid": p.get("pid"),
+                        "name": p.get("name") or "",
+                        "hbm_gib": (
+                            round(hbm_b / (1 << 30), 2) if isinstance(hbm_b, int) and hbm_b > 0 else None
+                        ),
+                    }
+                )
+    return rows
+
+
+def _pretouch_hbm_rows(
+    nodes: List[Dict[str, Any]],
+    threshold_gib: float,
+) -> List[Dict[str, Any]]:
+    """Per-GPU pre-touch HBM-used outliers (above threshold)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        per_gpu = (n.get("tier1") or {}).get("per_gpu") or []
+        for p in per_gpu:
+            d = p.get("details") or {}
+            used_gib = d.get("hbm_pre_touch_used_gib")
+            if not isinstance(used_gib, (int, float)):
+                continue
+            if used_gib >= threshold_gib:
+                rows.append(
+                    {
+                        "node_rank": n.get("node_rank", "?"),
+                        "host": n.get("host", "?"),
+                        "gpu": p.get("gpu", "?"),
+                        "used_gib": round(float(used_gib), 2),
+                    }
+                )
+    return rows
+
+
+def _gpu_activity_rows(
+    nodes: List[Dict[str, Any]],
+    warn_pct: float,
+) -> List[Dict[str, Any]]:
+    """Per-GPU compute-activity outliers (above warn threshold)."""
+    rows: List[Dict[str, Any]] = []
+    for n in nodes:
+        amd = (n.get("tier1") or {}).get("gpu_low_level") or {}
+        for rec in amd.get("per_gpu") or []:
+            pct = rec.get("gfx_activity_pct")
+            if not isinstance(pct, (int, float)):
+                continue
+            if float(pct) >= warn_pct:
+                rows.append(
+                    {
+                        "node_rank": n.get("node_rank", "?"),
+                        "host": n.get("host", "?"),
+                        "gpu": rec.get("gpu", "?"),
+                        "activity_pct": round(float(pct), 1),
+                    }
+                )
+    return rows
diff --git a/primus/tools/preflight/node_smoke/cli.py b/primus/tools/preflight/node_smoke/cli.py
new file mode 100644
index 000000000..2462dc55d
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/cli.py
@@ -0,0 +1,910 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""CLI wiring + the three subcommand entry points (`run`, `aggregate`,
+`_per_gpu`).
+
+The argparse layout is part of the operator-facing contract: flag names,
+defaults, help strings, and abbreviation behaviour are preserved exactly.
+``allow_abbrev=False`` on the ``run`` subparser keeps an old ``--tier2``
+flag from a stale script silently matching the new ``--tier2-perf``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from dataclasses import asdict
+from typing import Any, Dict, List, Optional
+
+from .aggregator.report import write_smoke_report
+from .collectors.clock import _collect_clock_state
+from .collectors.dmesg import _collect_dmesg_errors
+from .collectors.fingerprint import _collect_node_fingerprint
+from .collectors.gpu_low_level import _collect_amd_smi_metrics
+from .collectors.gpu_processes import _collect_gpu_processes
+from .collectors.host_limits import _collect_host_limits
+from .collectors.nics import _collect_nic_status
+from .collectors.reused_info import _collect_reused_info
+from .collectors.rocm_smi import _collect_rocm_smi_self_latency
+from .collectors.tooling import _collect_tooling_inventory
+from .collectors.xgmi import _collect_xgmi_topology
+from .logging_utils import _log, _short_name, _this_host_short, _warn
+from .orchestrator import _clean_dump_path, _node_status_from, _spawn_per_gpu
+from .per_gpu import _per_gpu_body
+from .rccl_local import _run_local_rccl
+from .types import GPUResult, NodeResult
+
+
+def _cmd_per_gpu(ns: argparse.Namespace) -> int:
+    """Internal subcommand: run all per-GPU checks for a single GPU index."""
+    result = _per_gpu_body(
+        gpu=int(ns.gpu),
+        tier2_perf=bool(ns.tier2_perf),
+        gemm_tflops_min=float(ns.gemm_tflops_min),
+        hbm_gbs_min=float(ns.hbm_gbs_min),
+        hbm_busy_threshold_bytes=int(float(ns.hbm_busy_threshold_gib) * (1 << 30)),
+    )
+    # Single JSON line on stdout; nothing else.
+    print(json.dumps(result), flush=True)
+    return 0 if result.get("status") == "PASS" else 1
+
+
+def _cmd_run(ns: argparse.Namespace) -> int:
+    """Per-node entry: orchestrate Tier 1 + optional Tier 2 + write JSON."""
+    # Always store the short hostname so consumers (passing_nodes.txt /
+    # failing_nodes.txt and SLURM tools that read them) get a name they
+    # can use directly.
+    host = _this_host_short()
+    node_rank = int(os.environ.get("NODE_RANK", os.environ.get("SLURM_NODEID", "0")))
+
+    # Rank-0 only: wipe stale JSONs / aggregator outputs from a previous
+    # run so re-runs on a different (smaller) nodelist cannot inherit
+    # ghost PASS verdicts from removed nodes. This MUST run before any
+    # rank's _spawn_per_gpu loop completes (and thus before any rank
+    # writes its JSON) -- which is why it lives at the top of _cmd_run
+    # on rank 0 only, not in the wrapper.
+    if node_rank == 0 and not ns.no_clean_dump_path:
+        removed = _clean_dump_path(ns.dump_path)
+        if removed:
+            _log(
+                f"cleaned {len(removed)} stale file(s) from {ns.dump_path} "
+                f"(per-node JSONs + aggregator outputs)"
+            )
+
+    expected_gpus = ns.expected_gpus
+    if expected_gpus is None:
+        expected_gpus = int(os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("GPUS_PER_NODE", "0")) or 0)
+        if expected_gpus <= 0:
+            try:
+                import torch  # type: ignore
+
+                expected_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
+            except Exception:
+                expected_gpus = 0
+    expected_gpus = max(0, int(expected_gpus))
+
+    # GPU visibility guard. We capture each independent source (the
+    # --expected-gpus flag, env vars, torch, and -- below -- amd-smi) so
+    # the JSON tells the operator *why* we resolved to N. The hard-fail
+    # rules live here, decoupled from any other collector, because we have
+    # seen `_collect_reused_info()` downgrade the "No GPUs detected" fail
+    # to a warn when collect_gpu_info() raises -- which would otherwise
+    # let a CPU-only or stale-GPU node PASS smoke silently.
+    torch_visible = 0
+    torch_is_available = False
+    try:
+        import torch  # type: ignore
+
+        torch_is_available = bool(torch.cuda.is_available())
+        torch_visible = int(torch.cuda.device_count()) if torch_is_available else 0
+    except Exception:
+        pass
+    gpu_visibility: Dict[str, Any] = {
+        "expected_gpus": expected_gpus,
+        "explicit_expected_gpus": ns.expected_gpus,
+        "torch_visible": torch_visible,
+        "torch_is_available": torch_is_available,
+        "env_local_world_size": int(os.environ.get("LOCAL_WORLD_SIZE", "0") or 0),
+        "env_gpus_per_node": int(os.environ.get("GPUS_PER_NODE", "0") or 0),
+        "amd_smi_visible": None,  # filled in after _collect_amd_smi_metrics
+        "fail_reasons": [],
+    }
+    if expected_gpus < 1:
+        msg = (
+            f"expected_gpus={expected_gpus}: no per-GPU sanity tests will "
+            f"run (torch_is_available={torch_is_available}, "
+            f"torch_visible={torch_visible}, "
+            f"LOCAL_WORLD_SIZE={gpu_visibility['env_local_world_size']}, "
+            f"GPUS_PER_NODE={gpu_visibility['env_gpus_per_node']})"
+        )
+        gpu_visibility["fail_reasons"].append(msg)
+        _warn(msg)
+
+    _log(
+        f"start node-smoke: node_rank={node_rank} expected_gpus={expected_gpus} "
+        f"tier2_perf={ns.tier2_perf}"
+    )
+    if ns.tier2_perf and expected_gpus < 2:
+        # Tier 2 also includes a node-local RCCL all-reduce, which needs at
+        # least 2 GPUs. Surface the skip up front instead of silently doing
+        # only GEMM/HBM and giving the operator false coverage confidence.
+        _warn(
+            f"--tier2-perf requested but expected_gpus={expected_gpus} < 2; "
+            "the node-local RCCL all-reduce phase will be skipped. "
+            "Per-GPU GEMM and HBM checks will still run."
+        )
+
+    # G: enumerate processes currently holding each GPU BEFORE we spawn
+    # any per-GPU subprocess. Anything we see here is, by definition, not
+    # us -- it's a leaked rank from a previous job, a foreign tenant, or
+    # an in-band monitoring agent. The aggregator + _node_status_from
+    # turn this into a hard FAIL unless the operator opted out via
+    # --allow-foreign-procs (or whitelisted the agent name).
+    allowed_proc_names = [s for s in (getattr(ns, "allowed_procs", "") or "").split(",") if s.strip()]
+    tier1_extra_pre: Dict[str, Any] = {}
+    tier1_extra_pre["gpu_processes"] = _collect_gpu_processes(
+        self_pid=os.getpid(),
+        allowed_proc_names=allowed_proc_names,
+    )
+    gp = tier1_extra_pre["gpu_processes"]
+    if gp.get("ok"):
+        _log(
+            f"gpu_processes ({gp.get('tool')}): "
+            f"{gp.get('foreign_count', 0)} foreign PID(s) across "
+            f"{len(gp.get('per_gpu') or [])} GPU bucket(s)"
+        )
+        if gp.get("foreign_count", 0) > 0:
+            for g in gp.get("per_gpu") or []:
+                for p in g.get("processes") or []:
+                    if p.get("is_foreign"):
+                        hbm = p.get("hbm_bytes")
+                        hbm_s = (
+                            f"{round(hbm / (1 << 30), 2)} GiB" if isinstance(hbm, int) and hbm > 0 else "?"
+                        )
+                        _warn(
+                            f"foreign process on gpu{g.get('gpu')}: "
+                            f"pid={p.get('pid')} name={p.get('name')!r} "
+                            f"hbm={hbm_s}"
+                        )
+    else:
+        _warn(
+            f"gpu_processes: enumeration unavailable "
+            f"({gp.get('error') or gp.get('json_error') or gp.get('text_error') or '?'})"
+        )
+
+    t0 = time.time()
+    per_gpu: List[GPUResult] = []
+    for i in range(expected_gpus):
+        r = _spawn_per_gpu(
+            i,
+            timeout_sec=ns.per_gpu_timeout_sec,
+            tier2_perf=bool(ns.tier2_perf),
+            gemm_tflops_min=ns.gemm_tflops_min,
+            hbm_gbs_min=ns.hbm_gbs_min,
+            hbm_busy_threshold_gib=float(ns.hbm_busy_threshold_gib),
+        )
+        per_gpu.append(r)
+        _log(
+            f"gpu{i}: {r.status} ({r.duration_sec:.1f}s)"
+            + (f" -- {r.reason}" if r.reason else "")
+            + (f" -- {r.details}" if r.details else "")
+        )
+
+    # Tier 1 reused info collectors
+    tier1_extra: Dict[str, Any] = {}
+    tier1_extra.update(tier1_extra_pre)  # carry forward gpu_processes
+    tier1_extra.update(_collect_reused_info())
+    if not ns.skip_dmesg:
+        tier1_extra["dmesg"] = _collect_dmesg_errors(window_minutes=ns.dmesg_minutes)
+    else:
+        tier1_extra["dmesg"] = {"ok": True, "matches": [], "error": "skipped"}
+
+    # A/B/C: software-stack fingerprint, NIC roll-call, host limits.
+    # All three are pure data-collection (millisecond-scale sysfs reads); the
+    # heavy cluster-level drift detection happens at aggregation time.
+    tier1_extra["fingerprint"] = _collect_node_fingerprint()
+    tier1_extra["nics"] = _collect_nic_status(
+        expected_count=ns.expected_rdma_nics,
+        allowlist=getattr(ns, "rdma_nic_allowlist", None),
+    )
+    tier1_extra["host_limits"] = _collect_host_limits(
+        ulimit_l_min_gb=ns.ulimit_l_min_gb,
+        shm_min_gb=ns.shm_min_gb,
+    )
+
+    # D-1 heavy: per-GPU ECC / throttle / clocks / power via amd-smi (one
+    # node-level call, results indexed by gpu).
+    # D-2: XGMI link matrix via amd-smi topology (one node-level call).
+    # E:   wall-time + time-daemon active states.
+    # F-partial: rocm-smi --version self-latency with a hard timeout to
+    # catch drivers that are starting to wedge.
+    # Tooling inventory FIRST so we can warn loudly before running the
+    # collectors that depend on each tool. Several downstream collectors
+    # (gpu_low_level, xgmi, gpu_processes, tooling) silently no-op when
+    # their tool is missing, which can let a broken node pass smoke
+    # unnoticed -- this section is the operator-visible counterweight.
+    tier1_extra["tooling_inventory"] = _collect_tooling_inventory()
+    inv = tier1_extra["tooling_inventory"]
+    inv_missing = inv["missing"]
+    inv_uncovered = inv["uncovered"]
+    if inv_missing:
+        # First line: which tools are missing. Loud regardless of whether
+        # a fallback covers everything.
+        _warn(f"tooling: {len(inv_missing)} tracked tool(s) NOT in PATH: " f"{', '.join(inv_missing)}.")
+        if inv_uncovered:
+            # Second line: which checks have NO working tool at all. This
+            # is the actually-dangerous case (a check that will silently
+            # no-op no matter which tool we try).
+            _warn(
+                f"tooling: {len(inv_uncovered)} check(s) have NO working "
+                f"tool and will be silently skipped: "
+                + ", ".join(inv_uncovered)
+                + ". Use --require-tools to promote missing tools to a "
+                "node FAIL."
+            )
+        else:
+            # Reassuring line: every check is covered via fallback.
+            _warn(
+                "tooling: every check is still covered via fallback "
+                "(rocm-smi / lsof) -- no checks will be silently skipped. "
+                "Use --require-tools to promote missing tools to a node "
+                "FAIL anyway if your environment requires them."
+            )
+
+    tier1_extra["gpu_low_level"] = _collect_amd_smi_metrics()
+    tier1_extra["xgmi"] = _collect_xgmi_topology()
+    tier1_extra["clock"] = _collect_clock_state()
+    tier1_extra["tooling"] = _collect_rocm_smi_self_latency(timeout_sec=float(ns.rocm_smi_timeout_sec))
+
+    # Visibility cross-check: if amd-smi successfully enumerated GPUs but
+    # torch couldn't see them, that's a high-signal sign of a stale ROCm
+    # install / wedged amdgpu driver -- exactly the case where a "smoke
+    # test" is supposed to pull the node out of rotation. We only treat
+    # the JSON path as authoritative for counting (the text fallback
+    # cannot be reliably parsed for a count).
+    amd_low = tier1_extra["gpu_low_level"]
+    if amd_low.get("ok") and amd_low.get("tool") == "amd-smi metric --json":
+        per = amd_low.get("per_gpu") or []
+        n_amd = len(per) if isinstance(per, list) else 0
+        gpu_visibility["amd_smi_visible"] = n_amd
+        if n_amd > 0 and torch_visible < n_amd:
+            mismatch = (
+                f"gpu_visibility_mismatch: amd-smi sees {n_amd} GPU(s) "
+                f"but torch.cuda.device_count()={torch_visible} "
+                f"(torch_is_available={torch_is_available}); ROCm install "
+                f"or amdgpu driver may be broken on this node"
+            )
+            gpu_visibility["fail_reasons"].append(mismatch)
+            _warn(mismatch)
+    tier1_extra["gpu_visibility"] = gpu_visibility
+    xg = tier1_extra["xgmi"]
+    if xg.get("ok"):
+        bad = xg.get("non_xgmi_pairs") or []
+        _log(f"xgmi: {xg.get('n_gpus', 0)}x{xg.get('n_gpus', 0)} matrix, " f"{len(bad)} non-XGMI pair(s)")
+    elif xg.get("error"):
+        _warn(f"xgmi: {xg.get('error')}")
+    tool = tier1_extra["tooling"]
+    if tool.get("ok"):
+        _log(f"rocm-smi --version: {tool.get('latency_sec')}s")
+    elif tool.get("timed_out"):
+        _warn(f"rocm-smi --version timed out after {tool.get('timeout_sec')}s " "-- driver may be wedging")
+    elif tool.get("error"):
+        _warn(f"tooling: {tool.get('error')}")
+    nic_summary = tier1_extra["nics"]
+    _log(
+        f"nics: {len(nic_summary.get('ports', []))} port(s) found, "
+        f"{len(nic_summary.get('issues', []))} issue(s)"
+    )
+    if tier1_extra["host_limits"].get("fail_reasons"):
+        for r in tier1_extra["host_limits"]["fail_reasons"]:
+            _warn(f"host_limits: {r}")
+
+    # Tier 2 local RCCL all-reduce. Gated on a single flag now (--tier2-perf)
+    # so users cannot accidentally end up with only the per-GPU half running.
+    tier2_extra: Dict[str, Any] = {}
+    if ns.tier2_perf and expected_gpus > 1:
+        _log(f"tier2 local RCCL all-reduce: {expected_gpus} ranks, {ns.rccl_size_mb}MB")
+        rccl = _run_local_rccl(
+            local_world_size=expected_gpus,
+            size_mb=ns.rccl_size_mb,
+            timeout_sec=ns.rccl_timeout_sec,
+        )
+        if rccl.get("status") == "PASS" and rccl.get("gbs") is not None:
+            if float(rccl["gbs"]) < ns.rccl_gbs_min:
+                rccl = {
+                    "status": "FAIL",
+                    "gbs": rccl["gbs"],
+                    "error": (f"local RCCL {rccl['gbs']} GB/s < threshold {ns.rccl_gbs_min}"),
+                }
+        tier2_extra["rccl"] = rccl
+        _log(f"tier2 RCCL: {rccl}")
+
+    required_tools = [s.strip() for s in (getattr(ns, "require_tools", "") or "").split(",") if s.strip()]
+    fail_reasons = _node_status_from(
+        per_gpu,
+        tier1_extra,
+        tier2_extra,
+        allow_foreign_procs=bool(ns.allow_foreign_procs),
+        required_tools=required_tools,
+    )
+    status = "PASS" if not fail_reasons else "FAIL"
+
+    node_result = NodeResult(
+        host=host,
+        node_rank=node_rank,
+        status=status,
+        duration_sec=round(time.time() - t0, 3),
+        fail_reasons=fail_reasons,
+        tier1={
+            "per_gpu": [asdict(r) for r in per_gpu],
+            **tier1_extra,
+        },
+        tier2=tier2_extra,
+    )
+
+    smoke_dir = os.path.join(ns.dump_path, "smoke")
+    os.makedirs(smoke_dir, exist_ok=True)
+    out_path = os.path.join(smoke_dir, f"{host}.json")
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(asdict(node_result), f, indent=2, default=str)
+    _log(f"wrote {out_path} status={status} duration={node_result.duration_sec}s")
+    if fail_reasons:
+        for r in fail_reasons[:5]:
+            _warn(r)
+
+    # Per-node `run` exits 0 whenever the smoke test ran to completion --
+    # the node verdict (PASS/FAIL) is in the JSON, in failing_nodes.txt,
+    # and in the aggregator's exit code. Conflating "this node is broken"
+    # with "this tool crashed" makes srun output look like the smoke test
+    # itself is failing, when in fact it's correctly DOING ITS JOB of
+    # identifying broken nodes. The aggregator (rank 0) is the single
+    # source of truth for the CI-friendly cluster-health exit signal.
+    # Tool failures (couldn't import, couldn't write JSON, etc.) still
+    # propagate as non-zero via Python's default exception handling.
+    return 0
+
+
+def _cmd_aggregate(ns: argparse.Namespace) -> int:
+    """Read all per-node JSONs from ``<dump>/smoke/`` and emit summary outputs."""
+    smoke_dir = os.path.join(ns.dump_path, "smoke")
+    os.makedirs(smoke_dir, exist_ok=True)
+
+    expected = int(ns.expected_nodes) if ns.expected_nodes is not None else None
+    deadline = time.time() + max(0, int(ns.wait_timeout_sec))
+    found_paths: List[str] = []
+    while True:
+        found_paths = sorted(os.path.join(smoke_dir, p) for p in os.listdir(smoke_dir) if p.endswith(".json"))
+        if expected is None or len(found_paths) >= expected:
+            break
+        if time.time() >= deadline:
+            break
+        time.sleep(1)
+
+    nodes: List[Dict[str, Any]] = []
+    for p in found_paths:
+        try:
+            with open(p, "r", encoding="utf-8") as f:
+                nodes.append(json.load(f))
+        except Exception as e:
+            nodes.append(
+                {
+                    "host": os.path.basename(p).rsplit(".json", 1)[0],
+                    "status": "FAIL",
+                    "fail_reasons": [f"failed to parse {p}: {e}"],
+                    "duration_sec": 0,
+                    "node_rank": -1,
+                }
+            )
+
+    # Normalize every loaded ``host`` to its short form so legacy JSON files
+    # that hold an FQDN (older runs of node_smoke) still produce SLURM-ready
+    # passing/failing lists.
+    for n in nodes:
+        n["host"] = _short_name(str(n.get("host", "")))
+
+    # Optional: an explicit expected hostname list (one per line). When
+    # provided, we name missing nodes by their real short hostname instead
+    # of synthetic ``<missing-N>`` placeholders, so the failing nodes list
+    # is directly usable with ``srun --exclude=``.
+    expected_hosts_short: List[str] = []
+    nodelist_file = getattr(ns, "expected_nodelist_file", None)
+    if nodelist_file:
+        try:
+            with open(nodelist_file, "r", encoding="utf-8") as f:
+                expected_hosts_short = [_short_name(line.strip()) for line in f if line.strip()]
+            _log(f"loaded {len(expected_hosts_short)} expected hostnames from " f"{nodelist_file}")
+        except Exception as e:
+            _warn(f"failed to read --expected-nodelist-file {nodelist_file}: {e}")
+
+    seen_hosts_short = {n.get("host", "") for n in nodes}
+
+    if expected_hosts_short:
+        # An explicit list always wins over --expected-nodes for both the
+        # count and (more importantly) the identity of missing nodes.
+        if expected is None or expected != len(expected_hosts_short):
+            expected = len(expected_hosts_short)
+        missing_hosts = sorted(set(expected_hosts_short) - seen_hosts_short)
+        for h in missing_hosts:
+            nodes.append(
+                {
+                    "host": h,
+                    "status": "FAIL",
+                    "fail_reasons": [
+                        f"no JSON received within {ns.wait_timeout_sec}s "
+                        f"(expected hostname '{h}' from --expected-nodelist-file)"
+                    ],
+                    "duration_sec": 0,
+                    "node_rank": -1,
+                }
+            )
+    elif expected is not None and len(seen_hosts_short) < expected:
+        # Fallback: we know the count but not the identities -> emit
+        # synthetic placeholders. These intentionally do NOT land in
+        # passing/failing txt files (see _is_real_host below).
+        for i in range(expected - len(seen_hosts_short)):
+            nodes.append(
+                {
+                    "host": f"<missing-{i}>",
+                    "status": "FAIL",
+                    "fail_reasons": [
+                        f"no JSON received within {ns.wait_timeout_sec}s "
+                        f"(expected_nodes={expected}, "
+                        f"found={len(seen_hosts_short)})"
+                    ],
+                    "duration_sec": 0,
+                    "node_rank": -1,
+                }
+            )
+
+    # Sort by node_rank if present, otherwise by hostname.
+    def _key(n: Dict[str, Any]):
+        nr = n.get("node_rank", 0)
+        return (
+            int(nr) if isinstance(nr, (int, str)) and str(nr).lstrip("-").isdigit() else 1 << 30,
+            str(n.get("host", "")),
+        )
+
+    nodes.sort(key=_key)
+
+    passing = [n for n in nodes if n.get("status") == "PASS"]
+    failing = [n for n in nodes if n.get("status") != "PASS"]
+
+    report_path = os.path.join(ns.dump_path, "smoke_report.md")
+    pass_path = os.path.join(ns.dump_path, "passing_nodes.txt")
+    fail_path = os.path.join(ns.dump_path, "failing_nodes.txt")
+
+    write_smoke_report(
+        report_path,
+        nodes=nodes,
+        passing=passing,
+        failing=failing,
+        expected=expected,
+        clock_skew_warn_sec=float(ns.clock_skew_warn_sec),
+        rocm_smi_warn_sec=float(ns.rocm_smi_warn_sec),
+        hbm_busy_threshold_gib=float(getattr(ns, "hbm_busy_threshold_gib", 2.0)),
+        gpu_activity_warn_pct=float(getattr(ns, "gpu_activity_warn_pct", 20.0)),
+    )
+
+    # Only write REAL hostnames to the txt files so they can be piped directly
+    # into `srun --nodelist=` / `srun --exclude=`. Synthetic "<missing-N>"
+    # placeholders for nodes that never reported are surfaced in the markdown
+    # report instead.
+    def _is_real_host(h: str) -> bool:
+        return bool(h) and not (h.startswith("<missing-") and h.endswith(">"))
+
+    with open(pass_path, "w", encoding="utf-8") as f:
+        for n in passing:
+            h = str(n.get("host", ""))
+            if _is_real_host(h):
+                f.write(h + "\n")
+    with open(fail_path, "w", encoding="utf-8") as f:
+        for n in failing:
+            h = str(n.get("host", ""))
+            if _is_real_host(h):
+                f.write(h + "\n")
+
+    _log(
+        f"aggregate: {len(passing)}/{len(nodes)} PASS  "
+        f"report={report_path}  passing={pass_path}  failing={fail_path}"
+    )
+
+    # R5: announce absolute report paths on stdout so the operator can
+    # copy/paste them. Always fires (this function only runs on rank 0).
+    # Under bash-side --silent these prints go to /dev/null along with
+    # everything else, which is acceptable per plan.
+    _announce_aggregate_paths(ns.dump_path, report_path, pass_path, fail_path)
+
+    return 0 if not failing and (expected is None or len(nodes) == expected) else 1
+
+
+def _announce_aggregate_paths(dump_path: str, report_path: str, pass_path: str, fail_path: str) -> None:
+    """Print absolute paths of aggregator outputs (R5).
+
+    Replaces the bash wrapper's ``log_always Report: ...`` lines. Always
+    rank-0 only by construction (only ``_cmd_aggregate`` calls this, and the
+    primus-cli wrapper only runs aggregate on rank 0).
+    """
+    for label, p in (("Report", report_path), ("Passing", pass_path), ("Failing", fail_path)):
+        try:
+            if os.path.isfile(p):
+                print(f"[Primus:NodeSmoke] {label}: {os.path.abspath(p)}", flush=True)
+        except OSError:
+            continue
+
+
+# ---------------------------------------------------------------------------
+# Argparse wiring
+#
+# `_add_run_flags` and `_add_aggregate_flags` are the canonical definition of
+# the user-facing flag surface. Both the standalone ``run`` / ``aggregate``
+# subparsers and the primus-cli ``node_smoke`` top-level parser
+# (primus/cli/subcommands/node_smoke.py) attach via these helpers, so a flag
+# added once here automatically reaches every entry point.
+# ---------------------------------------------------------------------------
+
+
+def _add_run_flags(parser: argparse.ArgumentParser) -> None:
+    """Attach every flag that `run` accepts to ``parser``.
+
+    Used by both the standalone ``run`` subparser and the primus-cli
+    top-level ``node_smoke`` parser. Adding a flag here picks it up
+    everywhere.
+    """
+    parser.add_argument(
+        "--dump-path", default="output/preflight", help="Directory under which smoke/<host>.json is written."
+    )
+    parser.add_argument(
+        "--expected-gpus",
+        type=int,
+        default=None,
+        help="Expected GPU count on this node (default: LOCAL_WORLD_SIZE/GPUS_PER_NODE/torch.cuda.device_count()).",
+    )
+    parser.add_argument(
+        "--per-gpu-timeout-sec", type=int, default=15, help="Hard timeout for each per-GPU subprocess."
+    )
+    parser.add_argument(
+        "--tier2-perf",
+        action="store_true",
+        help="Enable Tier 2 perf sanity: per-GPU GEMM TFLOPS, "
+        "HBM bandwidth, AND node-local RCCL all-reduce. "
+        "All three are fast (< 30 s/node total).",
+    )
+    parser.add_argument(
+        "--gemm-tflops-min",
+        type=float,
+        default=600.0,
+        help="FAIL if Tier 2 GEMM TFLOPS is below this. Default: 600 (MI300X-class).",
+    )
+    parser.add_argument(
+        "--hbm-gbs-min",
+        type=float,
+        default=2000.0,
+        help="FAIL if Tier 2 HBM GB/s is below this. Default: 2000.",
+    )
+    parser.add_argument(
+        "--rccl-size-mb", type=int, default=64, help="Tensor size for local RCCL all-reduce (MB)."
+    )
+    parser.add_argument(
+        "--rccl-gbs-min",
+        type=float,
+        default=100.0,
+        help="FAIL if local RCCL GB/s is below this. Default: 100.",
+    )
+    parser.add_argument(
+        "--rccl-timeout-sec", type=int, default=120, help="Hard timeout for the local RCCL all-reduce phase."
+    )
+    parser.add_argument(
+        "--skip-dmesg", action="store_true", help="Skip the dmesg recent-error scan (e.g. inside containers)."
+    )
+    parser.add_argument("--dmesg-minutes", type=int, default=15, help="Window for dmesg --since (minutes).")
+    # NIC / RDMA roll-call (B). expected_count=None means "report only";
+    # set this to e.g. 8 to make a missing or down NIC port a node FAIL.
+    # The count is compared against the **included** (training-NIC) set
+    # resolved by the selector chain below, NOT the raw number of devices
+    # under /sys/class/infiniband (which on multi-role nodes can include
+    # frontend / management / storage RoCE NICs).
+    parser.add_argument(
+        "--expected-rdma-nics",
+        type=int,
+        default=None,
+        help="Expected RDMA training-NIC port count (compared against "
+        "the included set, not /sys/class/infiniband total). If "
+        "set, a count mismatch becomes a node FAIL.",
+    )
+    # Training-NIC selector. Many clusters expose more RDMA-capable
+    # ports than the training job will actually use (frontend /
+    # management / storage NICs). The hard-fail rules (state must be
+    # ACTIVE, phys_state must be LinkUp, RoCE v2 GID present, ...) only
+    # run against the included subset. Excluded ports stay visible in
+    # the JSON (under tier1.nics.excluded_ports + info_issues) for
+    # diagnostics. Precedence: this flag > NCCL_IB_HCA env > heuristic
+    # (auto-exclude phys_state in {Disabled, Sleep}).
+    parser.add_argument(
+        "--rdma-nic-allowlist",
+        type=str,
+        default=None,
+        help="Comma-separated device[:port] list of training NICs "
+        "(NCCL_IB_HCA syntax: `^...` for denylist, `=dev` for "
+        "exact-match, no `:port` = match any port on device). "
+        "Wins over the NCCL_IB_HCA env. When neither is set, the "
+        "collector auto-excludes ports whose phys_state is "
+        "Disabled or Sleep (admin-disabled ports) and treats every "
+        "other port as a training NIC.",
+    )
+    # Host-limits hard thresholds (C). Set to 0 to disable a check.
+    parser.add_argument(
+        "--ulimit-l-min-gb",
+        type=float,
+        default=32.0,
+        help="FAIL the node if RLIMIT_MEMLOCK is finite and below "
+        "this many GiB (RDMA pin will fail). 0 disables.",
+    )
+    parser.add_argument(
+        "--shm-min-gb",
+        type=float,
+        default=8.0,
+        help="FAIL the node if /dev/shm is below this many GiB " "(NCCL shared-mem may fail). 0 disables.",
+    )
+    # F-partial: rocm-smi self-latency. Hitting this timeout is treated as a
+    # hard fail because a wedging amdgpu driver typically makes rocm-smi
+    # hang for 30-60 s before the GPU itself stops responding.
+    parser.add_argument(
+        "--rocm-smi-timeout-sec",
+        type=float,
+        default=5.0,
+        help="Hard timeout for `rocm-smi --version`. Hitting it is " "a node FAIL (driver likely wedging).",
+    )
+    # G: foreign / leaked process detection. Hard-fail by default; the
+    # operator can opt out for partitions that legitimately co-tenant the
+    # GPU, or whitelist known in-band agents by name.
+    parser.add_argument(
+        "--hbm-busy-threshold-gib",
+        type=float,
+        default=2.0,
+        help="FAIL the node if any GPU has at least this much "
+        "HBM in use BEFORE we touch the device (i.e. someone "
+        "else is holding it). Boundary is inclusive. Default: 2.0 GiB.",
+    )
+    parser.add_argument(
+        "--allow-foreign-procs",
+        action="store_true",
+        help="Do NOT FAIL the node when foreign processes are "
+        "found holding a GPU. They will still be reported.",
+    )
+    parser.add_argument(
+        "--allowed-procs",
+        type=str,
+        default="gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter",
+        help="Comma-separated process names that are OK to find "
+        "holding the GPU. The default whitelists the AMD "
+        "system agents (`gpuagent`, `rocm-smi-daemon`, "
+        "`amd-smi`) and a common observability sidecar "
+        "(`dcgm-exporter`) so they don't fail every node. "
+        "Set to an empty string to disable the whitelist.",
+    )
+    parser.add_argument(
+        "--gpu-activity-warn-pct",
+        type=float,
+        default=20.0,
+        help="Aggregator warns (does NOT fail) if amd-smi reports "
+        "any GPU's gfx_activity_pct above this when smoke "
+        "starts. Default: 20.",
+    )
+    # Tooling availability. Default empty = warn-only (the WARN already
+    # fires before the per-GPU subprocesses run, and the aggregator
+    # always renders a "Tooling availability" section). Strict
+    # environments can pass --require-tools amd-smi,rocm-smi to promote
+    # a missing tool to a hard node FAIL.
+    parser.add_argument(
+        "--require-tools",
+        type=str,
+        default="",
+        help="Comma-separated CLI tool names that MUST be "
+        "present in PATH for the node to PASS. Anything "
+        "missing becomes a hard node FAIL. Tracked tools: "
+        "amd-smi, rocm-smi, lsof. Default: warn-only.",
+    )
+    parser.add_argument(
+        "--no-clean-dump-path",
+        action="store_true",
+        help="Do NOT auto-wipe stale per-node JSONs and aggregator "
+        "outputs from --dump-path on rank 0 at startup. "
+        "Default behavior is to clean so re-runs on a "
+        "different nodelist don't inherit ghost PASS "
+        "verdicts from removed nodes.",
+    )
+
+
+def _add_aggregate_flags(parser: argparse.ArgumentParser, include_dump_path: bool = True) -> None:
+    """Attach every flag that `aggregate` accepts to ``parser``.
+
+    ``include_dump_path=False`` is used by the primus-cli wrapper because
+    ``--dump-path`` is already attached by ``_add_run_flags`` -- both share
+    that flag with identical defaults / meaning, so we attach it once.
+    """
+    if include_dump_path:
+        parser.add_argument("--dump-path", default="output/preflight", help="Same as `run --dump-path`.")
+    parser.add_argument(
+        "--expected-nodes",
+        type=int,
+        default=None,
+        help="Number of nodes expected to report. Missing nodes are FAIL.",
+    )
+    parser.add_argument(
+        "--wait-timeout-sec",
+        type=int,
+        default=60,
+        help="How long to wait for all expected JSONs to land before aggregating anyway.",
+    )
+    parser.add_argument(
+        "--rocm-smi-warn-sec",
+        type=float,
+        default=1.0,
+        help="Flag (warn-only) any node where `rocm-smi --version` " "took longer than this many seconds.",
+    )
+    parser.add_argument(
+        "--clock-skew-warn-sec",
+        type=float,
+        default=30.0,
+        help="Warn (info-only) when wall-clock spread across nodes "
+        "exceeds this many seconds. Includes srun launch "
+        "jitter so the default is loose.",
+    )
+    # Mirror the run-side thresholds so the report can label its sections
+    # using the same numbers each node's `run` was configured with. When
+    # `_add_aggregate_flags` is called on the same parser as
+    # `_add_run_flags` (the primus-cli wrapper), these are skipped to avoid
+    # argparse `conflicting option` errors -- the run-side defaults already
+    # cover the same names with identical defaults.
+    if include_dump_path:
+        parser.add_argument(
+            "--hbm-busy-threshold-gib",
+            type=float,
+            default=2.0,
+            help="Pre-touch HBM-used threshold (GiB) used by the "
+            "'GPU pre-touch HBM usage outliers' section.",
+        )
+        parser.add_argument(
+            "--gpu-activity-warn-pct",
+            type=float,
+            default=20.0,
+            help="GPU activity %% threshold used by the " "'GPU compute-activity outliers' section.",
+        )
+    parser.add_argument(
+        "--expected-nodelist-file",
+        type=str,
+        default=None,
+        help="Optional file with one expected (short) hostname per line. "
+        "When provided, missing nodes are reported with their real "
+        "hostname instead of synthetic <missing-N> placeholders, and "
+        "are written to failing_nodes.txt directly. The primus-cli "
+        "wrapper auto-populates this from `scontrol show hostnames` "
+        "when running under SLURM.",
+    )
+
+
+def _resolve_aggregate_args_from_slurm(args: argparse.Namespace) -> argparse.Namespace:
+    """Build an aggregator-flavored Namespace from ``args``.
+
+    Used by the primus-cli ``node_smoke`` wrapper to populate
+    ``expected_nodes`` and ``expected_nodelist_file`` from the SLURM
+    environment when the user did not pass them explicitly. All other
+    aggregator fields pass through unchanged. The standalone CLI's
+    ``aggregate`` subcommand is unaffected (it constructs its own
+    Namespace via argparse).
+    """
+    dump_path = getattr(args, "dump_path", "output/preflight") or "output/preflight"
+
+    # Expected node count: CLI > SLURM_NNODES > SLURM_JOB_NUM_NODES > NNODES.
+    expected_nodes = getattr(args, "expected_nodes", None)
+    if expected_nodes is None:
+        for var in ("SLURM_NNODES", "SLURM_JOB_NUM_NODES", "NNODES"):
+            v = os.environ.get(var, "")
+            if v.isdigit() and int(v) > 0:
+                expected_nodes = int(v)
+                break
+
+    # Expected nodelist file: CLI > SLURM_JOB_NODELIST (via `scontrol show
+    # hostnames`). Mirrors the deleted run_node_smoke_direct.sh behavior so
+    # the aggregator can name MISSING nodes by their real short hostname
+    # instead of synthetic <missing-N> placeholders.
+    expected_nodelist_file = getattr(args, "expected_nodelist_file", None)
+    if not expected_nodelist_file:
+        slurm_nodelist = os.environ.get("SLURM_JOB_NODELIST", "")
+        if slurm_nodelist and _which("scontrol"):
+            try:
+                os.makedirs(dump_path, exist_ok=True)
+                candidate = os.path.join(dump_path, "expected_nodes.txt")
+                import subprocess
+
+                proc = subprocess.run(
+                    ["scontrol", "show", "hostnames", slurm_nodelist],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.DEVNULL,
+                    check=False,
+                )
+                if proc.returncode == 0 and proc.stdout.strip():
+                    with open(candidate, "wb") as f:
+                        f.write(proc.stdout)
+                    expected_nodelist_file = candidate
+                    n_hosts = sum(1 for line in proc.stdout.splitlines() if line.strip())
+                    _log(f"resolved expected nodelist ({n_hosts} nodes) -> {candidate}")
+                else:
+                    _warn(
+                        "scontrol show hostnames returned empty / non-zero; "
+                        "aggregator will use <missing-N> placeholders"
+                    )
+            except OSError as e:
+                _warn(f"failed to resolve expected nodelist via scontrol: {e}")
+
+    return argparse.Namespace(
+        dump_path=dump_path,
+        expected_nodes=expected_nodes,
+        wait_timeout_sec=int(getattr(args, "wait_timeout_sec", 60)),
+        rocm_smi_warn_sec=float(getattr(args, "rocm_smi_warn_sec", 1.0)),
+        clock_skew_warn_sec=float(getattr(args, "clock_skew_warn_sec", 30.0)),
+        hbm_busy_threshold_gib=float(getattr(args, "hbm_busy_threshold_gib", 2.0)),
+        gpu_activity_warn_pct=float(getattr(args, "gpu_activity_warn_pct", 20.0)),
+        expected_nodelist_file=expected_nodelist_file,
+    )
+
+
+def _which(cmd: str) -> Optional[str]:
+    """Lightweight ``shutil.which`` import (avoid module-level import cost)."""
+    import shutil
+
+    return shutil.which(cmd)
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="python -m primus.tools.preflight.node_smoke",
+        description=(
+            "Node-local preflight smoke test. Each node runs independently "
+            "(no global rendezvous) and writes a per-node JSON verdict."
+        ),
+    )
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    # ---- run ----
+    # allow_abbrev=False so abbreviated forms (e.g. an old --tier2 left in
+    # a script) do NOT silently match the new --tier2-perf as a prefix.
+    # We want them to error out loudly so behavior never changes silently
+    # underneath an unsuspecting caller.
+    pr = sub.add_parser(
+        "run",
+        help="Run per-node smoke test on this node.",
+        allow_abbrev=False,
+    )
+    _add_run_flags(pr)
+    pr.set_defaults(func=_cmd_run)
+
+    # ---- aggregate ----
+    pa = sub.add_parser("aggregate", help="Aggregate per-node JSONs into report + passing/failing lists.")
+    _add_aggregate_flags(pa)
+    pa.set_defaults(func=_cmd_aggregate)
+
+    # ---- _per_gpu (internal) ----
+    pg = sub.add_parser(
+        "_per_gpu", help="(internal) Run smoke checks for a single GPU index. Spawned by `run`."
+    )
+    pg.add_argument("gpu", type=int)
+    pg.add_argument("--tier2-perf", action="store_true")
+    pg.add_argument("--gemm-tflops-min", type=float, default=600.0)
+    pg.add_argument("--hbm-gbs-min", type=float, default=2000.0)
+    pg.add_argument("--hbm-busy-threshold-gib", type=float, default=2.0)
+    pg.set_defaults(func=_cmd_per_gpu)
+
+    return p
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    parser = _build_parser()
+    ns = parser.parse_args(argv)
+    return int(ns.func(ns))
diff --git a/primus/tools/preflight/node_smoke/collectors/__init__.py b/primus/tools/preflight/node_smoke/collectors/__init__.py
new file mode 100644
index 000000000..611ef9105
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/__init__.py
@@ -0,0 +1,23 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 collectors -- one module per Tier 1 sub-section.
+
+* ``dmesg``        : recent dmesg error scan
+* ``fingerprint``  : Tier 1 A -- software-stack fingerprint
+* ``nics``         : Tier 1 B -- NIC / RDMA roll-call
+* ``host_limits``  : Tier 1 C -- host limits (ulimit, /dev/shm, NUMA, governor)
+* ``gpu_low_level``: Tier 1 D-1 -- per-GPU ECC / clocks / power via amd-smi
+* ``xgmi``         : Tier 1 D-2 -- XGMI topology matrix
+* ``clock``        : Tier 1 E -- wall time + time-daemon active states
+* ``rocm_smi``     : Tier 1 F + fallbacks -- rocm-smi self-latency and
+                     amd-smi fallback parsers
+* ``gpu_processes``: Tier 1 G -- foreign / leaked process detection
+* ``tooling``      : tooling-availability inventory
+* ``reused_info``  : reused gpu/host/network info collectors
+"""
+
+from __future__ import annotations
diff --git a/primus/tools/preflight/node_smoke/collectors/clock.py b/primus/tools/preflight/node_smoke/collectors/clock.py
new file mode 100644
index 000000000..86b13b0c9
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/clock.py
@@ -0,0 +1,36 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- E: clock state (wall time + time-daemon active states)."""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict
+
+from ..shell_utils import _systemctl_is_active
+
+
+def _collect_clock_state() -> Dict[str, Any]:
+    """Capture this node's wall time and time-daemon health.
+
+    Wall time is captured early so the aggregator can compute a
+    cluster-wide spread. Note this includes srun launch jitter, so the
+    spread is an *upper bound* on the real clock skew. The aggregator
+    uses loose thresholds (warn at 30 s, no hard fail) for the spread,
+    and reserves the hard fail for "no time-sync daemon active".
+    """
+    out: Dict[str, Any] = {
+        "wall_time_unix": time.time(),
+        "monotonic": time.monotonic(),
+        "daemons": {},
+    }
+    for unit in ("chronyd", "ntp", "ntpd", "systemd-timesyncd"):
+        out["daemons"][unit] = _systemctl_is_active(unit)
+    active = [u for u, s in out["daemons"].items() if s == "active"]
+    out["any_active"] = bool(active)
+    out["active_units"] = active
+    return out
diff --git a/primus/tools/preflight/node_smoke/collectors/dmesg.py b/primus/tools/preflight/node_smoke/collectors/dmesg.py
new file mode 100644
index 000000000..6ae3628fb
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/dmesg.py
@@ -0,0 +1,98 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Recent dmesg error scan.
+
+Greps the last ``window_minutes`` of dmesg for known-bad patterns. Best
+effort: missing dmesg / failed read becomes ``ok=False`` rather than
+raising.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from typing import Any, Dict, List
+
+# Regex (NOT substring) patterns matched case-insensitively against each
+# dmesg line. They MUST be valid Python regex -- if you only want a literal
+# substring (e.g. ``mce: ``), it's still a valid regex with no specials.
+# Why regex: real amdgpu failure lines look like
+#   ``amdgpu 0000:05:00.0: amdgpu_device_resume failed: -19``
+#   ``amdgpu: [drm] *ERROR* ring sdma0 timeout``
+# so we need ``amdgpu.*(error|fail|timeout)`` -- a substring match against
+# the literal pattern ``amdgpu.*error`` would essentially never fire.
+_DMESG_PATTERNS = (
+    r"\bxid\b",
+    r"hardware error",
+    r"gpu reset",
+    r"hung_task",
+    r"hung task",
+    r"page allocation failure",
+    r"soft lockup",
+    r"amdgpu.*(error|fail|timeout)",
+    r"\*error\*",  # [drm] *ERROR*
+    r"mce: ",
+)
+
+
+def _collect_dmesg_errors(window_minutes: int = 15) -> Dict[str, Any]:
+    """Best-effort grep of recent dmesg lines for known-bad patterns.
+
+    Returns a dict with ``ok`` (bool), ``matches`` (list of matched lines, capped),
+    and ``error`` (str) when dmesg cannot be read.
+    """
+    out: Dict[str, Any] = {"ok": True, "matches": [], "error": None}
+    try:
+        # ``--since`` requires recent util-linux; fall back to the last 2000 lines.
+        try:
+            cp = subprocess.run(
+                ["dmesg", "--since", f"-{window_minutes}min"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=10,
+                check=False,
+            )
+            if cp.returncode != 0:
+                raise RuntimeError(cp.stderr.strip() or f"rc={cp.returncode}")
+            text = cp.stdout
+        except Exception:
+            cp = subprocess.run(
+                ["dmesg"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=10,
+                check=False,
+            )
+            if cp.returncode != 0:
+                out["ok"] = False
+                out["error"] = (cp.stderr or "").strip() or f"rc={cp.returncode}"
+                return out
+            text = "\n".join(cp.stdout.splitlines()[-2000:])
+
+        import re
+
+        matches: List[str] = []
+        # Pre-compile each pattern with re.IGNORECASE; a malformed regex is
+        # logged into ``out['pattern_errors']`` but never aborts the scan.
+        compiled: List[Any] = []
+        for p in _DMESG_PATTERNS:
+            try:
+                compiled.append(re.compile(p, re.IGNORECASE))
+            except re.error as e:
+                out.setdefault("pattern_errors", []).append(f"{p!r}: {e}")
+        for line in text.splitlines():
+            if any(pat.search(line) for pat in compiled):
+                matches.append(line)
+                if len(matches) >= 50:
+                    break
+        out["matches"] = matches
+        return out
+    except Exception as e:
+        out["ok"] = False
+        out["error"] = str(e)
+        return out
diff --git a/primus/tools/preflight/node_smoke/collectors/fingerprint.py b/primus/tools/preflight/node_smoke/collectors/fingerprint.py
new file mode 100644
index 000000000..225a9b94a
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/fingerprint.py
@@ -0,0 +1,88 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- A. Software-stack fingerprint (drift detection happens at aggregate)."""
+
+from __future__ import annotations
+
+import os
+import sys
+from typing import Any, Dict
+
+from ..shell_utils import _parse_os_release_pretty, _read_text
+
+
+def _collect_node_fingerprint() -> Dict[str, Any]:
+    """Collect a deterministic, hashable fingerprint of the software stack
+    on this node so the aggregator can detect drift across the cluster.
+
+    Every value is best-effort: missing tools / files become ``None`` rather
+    than raising. The aggregator skips ``None`` values when computing the
+    cluster majority for a given key.
+    """
+    fp: Dict[str, Any] = {}
+
+    # Kernel + OS
+    try:
+        fp["kernel"] = os.uname().release
+    except Exception:
+        fp["kernel"] = None
+    fp["os_release"] = _parse_os_release_pretty()
+    fp["python"] = sys.version.split()[0]
+
+    # ROCm / HIP / amdgpu
+    fp["rocm"] = _read_text("/opt/rocm/.info/version") or None
+    fp["amdgpu_driver"] = _read_text("/sys/module/amdgpu/version") or None
+
+    # PyTorch + (R)CCL
+    try:
+        import torch  # type: ignore
+
+        fp["torch"] = getattr(torch, "__version__", None)
+        fp["torch_hip"] = getattr(getattr(torch, "version", None), "hip", None)
+        try:
+            v = torch.cuda.nccl.version()  # type: ignore[attr-defined]
+            if isinstance(v, tuple):
+                fp["rccl"] = ".".join(str(x) for x in v)
+            else:
+                fp["rccl"] = str(v)
+        except Exception:
+            fp["rccl"] = None
+
+        # Locate librccl.so under torch's lib dir for a stable per-node path.
+        try:
+            torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib")
+            for n in sorted(os.listdir(torch_lib)):
+                if n.startswith("librccl.so"):
+                    fp["rccl_path"] = os.path.join(torch_lib, n)
+                    break
+        except Exception:
+            pass
+    except Exception:
+        fp["torch"] = None
+        fp["torch_hip"] = None
+        fp["rccl"] = None
+
+    # Per-IB-device firmware + HCA model fingerprints. Both are critical for
+    # detecting "1 of N nodes flashed differently" silent regressions.
+    nic_fw: Dict[str, str] = {}
+    nic_hca: Dict[str, str] = {}
+    ib_root = "/sys/class/infiniband"
+    if os.path.isdir(ib_root):
+        try:
+            for dev in sorted(os.listdir(ib_root)):
+                fw = _read_text(os.path.join(ib_root, dev, "fw_ver"))
+                if fw:
+                    nic_fw[dev] = fw
+                hca = _read_text(os.path.join(ib_root, dev, "hca_type"))
+                if hca:
+                    nic_hca[dev] = hca
+        except Exception:
+            pass
+    fp["nic_fw"] = nic_fw or None
+    fp["nic_hca"] = nic_hca or None
+
+    return fp
diff --git a/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py b/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py
new file mode 100644
index 000000000..27b99dbef
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py
@@ -0,0 +1,249 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- D-1 heavy: per-GPU low-level via amd-smi (ECC, throttle, clocks,
+power cap). Runs ONCE per node (not per per-GPU subprocess) so the smoke
+step doesn't pay an amd-smi startup tax 8x. Best-effort: missing amd-smi
+or unparseable output degrades to {"ok": False, ...} without raising.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from typing import Any, Dict, List
+
+from ..shell_utils import _which
+from .rocm_smi import _rocm_smi_ras_info_text, _rocm_smi_use_json
+
+
+def _collect_amd_smi_metrics() -> Dict[str, Any]:
+    """Best-effort capture of per-GPU low-level metrics via ``amd-smi``.
+
+    We try ``amd-smi metric --json`` first (newer builds emit valid JSON);
+    if that fails we fall back to text output and surface the raw text under
+    ``raw`` so an operator can still grep it. The on-disk shape is:
+
+        {
+            "ok": bool,
+            "tool": "amd-smi metric --json" | "amd-smi metric" | None,
+            "per_gpu": [ {gpu, gfx_clock_mhz, hbm_used_bytes,
+                          power_avg_w, power_cap_w, temp_edge_c,
+                          ecc_uncorrectable_total, ecc_correctable_total,
+                          throttle_status_raw, ...}, ... ],
+            "error": "..."  (only when ok is False)
+        }
+
+    Hard-fail semantics live in ``_node_status_from``: any non-zero
+    uncorrectable ECC count becomes a node FAIL. Throttle status is
+    captured under ``throttle_status_raw`` for operator inspection but
+    is NOT failed-on -- the amd-smi throttle schema varies too much
+    across releases to make a robust default rule.
+    """
+    out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []}
+
+    if _which("amd-smi") is not None:
+        # Try JSON first.
+        try:
+            cp = subprocess.run(
+                ["amd-smi", "metric", "--json"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=15,
+                check=False,
+            )
+            if cp.returncode == 0 and cp.stdout.strip():
+                try:
+                    doc = json.loads(cp.stdout)
+                    out["ok"] = True
+                    out["tool"] = "amd-smi metric --json"
+                    out["per_gpu"] = _flatten_amd_smi_metric_json(doc)
+                except Exception as e:
+                    out["json_parse_error"] = str(e)
+            else:
+                out["json_rc"] = cp.returncode
+                out["json_stderr"] = (cp.stderr or "").strip()[:200]
+        except subprocess.TimeoutExpired:
+            out["json_error"] = "amd-smi metric --json timed out"
+        except Exception as e:
+            out["json_error"] = str(e)
+
+        # If JSON didn't work, capture the raw text output so the operator
+        # can still grep it. We don't try to parse the human-readable text
+        # -- per-GPU outliers will still show up via the sysfs/torch-side
+        # details we capture in _per_gpu_body, and the rocm-smi fallback
+        # below will fill in ECC + activity in the per_gpu records.
+        if not out["ok"]:
+            try:
+                cp = subprocess.run(
+                    ["amd-smi", "metric"],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    timeout=15,
+                    check=False,
+                )
+                if cp.returncode == 0:
+                    out["ok"] = True
+                    out["tool"] = "amd-smi metric"
+                    out["raw"] = cp.stdout[:8000]  # cap to keep JSON small
+                else:
+                    out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}"
+            except subprocess.TimeoutExpired:
+                out["error"] = "amd-smi metric timed out"
+            except Exception as e:
+                out["error"] = str(e)
+    else:
+        out["error"] = "amd-smi not found in PATH"
+
+    # rocm-smi fallback / fill-in. We always run it -- even when amd-smi
+    # produced JSON, the schema may have missed ECC or activity fields,
+    # and rocm-smi gives us a second independent source. Per-GPU records
+    # are merged by GPU index so amd-smi-provided power/clocks/temp stay
+    # alongside rocm-smi-provided ECC/activity.
+    if _which("rocm-smi") is not None:
+        # Build an index of existing records so we can merge in place.
+        by_gpu: Dict[int, Dict[str, Any]] = {}
+        for rec in out.get("per_gpu") or []:
+            g = rec.get("gpu")
+            if isinstance(g, int):
+                by_gpu[g] = rec
+
+        # ECC fill-in (only if not already populated by amd-smi)
+        ras = _rocm_smi_ras_info_text()
+        ecc_filled = 0
+        if ras.get("ok"):
+            for rec in ras.get("per_gpu") or []:
+                g = rec.get("gpu")
+                if not isinstance(g, int):
+                    continue
+                tgt = by_gpu.setdefault(g, {"gpu": g})
+                if tgt.get("ecc_uncorrectable_total") is None:
+                    tgt["ecc_uncorrectable_total"] = rec.get("ecc_uncorrectable_total")
+                    tgt["ecc_correctable_total"] = rec.get("ecc_correctable_total")
+                    tgt["ecc_source"] = "rocm-smi --showrasinfo"
+                    ecc_filled += 1
+            if ecc_filled:
+                out.setdefault("fallback_tools", []).append(
+                    f"rocm-smi --showrasinfo (ECC for {ecc_filled} GPU(s))"
+                )
+
+        # Activity fill-in (only if not already populated by amd-smi)
+        use = _rocm_smi_use_json()
+        act_filled = 0
+        if use.get("ok"):
+            for rec in use.get("per_gpu") or []:
+                g = rec.get("gpu")
+                if not isinstance(g, int):
+                    continue
+                tgt = by_gpu.setdefault(g, {"gpu": g})
+                if tgt.get("gfx_activity_pct") is None:
+                    tgt["gfx_activity_pct"] = rec.get("gfx_activity_pct")
+                    tgt["activity_source"] = "rocm-smi --showuse"
+                    act_filled += 1
+            if act_filled:
+                out.setdefault("fallback_tools", []).append(
+                    f"rocm-smi --showuse (activity for {act_filled} GPU(s))"
+                )
+
+        # Rebuild per_gpu in stable index order if rocm-smi added entries
+        if ecc_filled or act_filled:
+            out["per_gpu"] = [by_gpu[g] for g in sorted(by_gpu.keys())]
+            # If amd-smi produced nothing at all, this is now a successful
+            # collection, just sourced entirely from rocm-smi.
+            if not out["ok"] and out["per_gpu"]:
+                out["ok"] = True
+                out["tool"] = "rocm-smi (ECC/activity fallback)"
+
+    return out
+
+
+def _flatten_amd_smi_metric_json(doc: Any) -> List[Dict[str, Any]]:
+    """Pull the fields we care about out of `amd-smi metric --json` output.
+
+    The exact schema varies between amd-smi releases. We touch only the
+    most-stable nesting -- a top-level list of per-GPU dicts, each with
+    sub-blocks like ``power``, ``clock``, ``temperature``, ``ecc``,
+    ``throttle_status`` -- and tolerate missing fields silently.
+    """
+    out: List[Dict[str, Any]] = []
+    items = doc if isinstance(doc, list) else (doc.get("gpus", []) if isinstance(doc, dict) else [])
+    for i, g in enumerate(items):
+        if not isinstance(g, dict):
+            continue
+        rec: Dict[str, Any] = {"gpu": i}
+        # gpu id may be in g["gpu"] or g["device_id"] depending on schema
+        if isinstance(g.get("gpu"), int):
+            rec["gpu"] = g["gpu"]
+        # power
+        power = g.get("power") or {}
+        if isinstance(power, dict):
+            for k_src, k_dst in (
+                ("average_socket_power", "power_avg_w"),
+                ("current_socket_power", "power_avg_w"),
+                ("socket_power", "power_avg_w"),
+                ("power_cap", "power_cap_w"),
+                ("power_limit", "power_cap_w"),
+            ):
+                v = power.get(k_src)
+                if isinstance(v, (int, float)) and rec.get(k_dst) is None:
+                    rec[k_dst] = v
+        # clocks (gfx clock most useful)
+        clk = g.get("clock") or g.get("clocks") or {}
+        if isinstance(clk, dict):
+            gfx = clk.get("gfx") or clk.get("gfx_0") or clk.get("gfx_clock") or {}
+            if isinstance(gfx, dict):
+                for k in ("clk", "current", "value", "frequency"):
+                    if isinstance(gfx.get(k), (int, float)):
+                        rec["gfx_clock_mhz"] = gfx[k]
+                        break
+            elif isinstance(gfx, (int, float)):
+                rec["gfx_clock_mhz"] = gfx
+        # temperature
+        temp = g.get("temperature") or {}
+        if isinstance(temp, dict):
+            for k in ("edge", "current", "value"):
+                v = temp.get(k)
+                if isinstance(v, (int, float)):
+                    rec["temp_edge_c"] = v
+                    break
+        # ECC
+        ecc = g.get("ecc") or g.get("ecc_count") or {}
+        if isinstance(ecc, dict):
+            ue = ecc.get("uncorrectable") or ecc.get("uncorrectable_total") or ecc.get("ue") or 0
+            ce = ecc.get("correctable") or ecc.get("correctable_total") or ecc.get("ce") or 0
+            try:
+                rec["ecc_uncorrectable_total"] = int(ue)
+                rec["ecc_correctable_total"] = int(ce)
+            except Exception:
+                pass
+        # Throttle
+        thr = g.get("throttle_status") or g.get("throttle") or {}
+        if isinstance(thr, dict):
+            rec["throttle_status_raw"] = thr
+        elif isinstance(thr, (str, list)):
+            rec["throttle_status_raw"] = thr
+        # GPU compute activity %. Used by the aggregator to surface GPUs
+        # that are running someone else's compute right now (warn-only;
+        # short bursts are normal but a sustained pegged-100% across
+        # multiple GPUs is a strong signal of a leaked rank from a
+        # previous job).
+        usage = g.get("usage") or g.get("activity") or g.get("utilization") or {}
+        if isinstance(usage, dict):
+            for k in ("gfx_activity", "gfx", "gfx_busy_percent", "gpu_busy_percent"):
+                v = usage.get(k)
+                if isinstance(v, (int, float)):
+                    rec["gfx_activity_pct"] = float(v)
+                    break
+            if rec.get("gfx_activity_pct") is None:
+                v = usage.get("activity") or usage.get("value")
+                if isinstance(v, (int, float)):
+                    rec["gfx_activity_pct"] = float(v)
+        elif isinstance(usage, (int, float)):
+            rec["gfx_activity_pct"] = float(usage)
+        out.append(rec)
+    return out
diff --git a/primus/tools/preflight/node_smoke/collectors/gpu_processes.py b/primus/tools/preflight/node_smoke/collectors/gpu_processes.py
new file mode 100644
index 000000000..9e24ba19b
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/gpu_processes.py
@@ -0,0 +1,614 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- G: foreign / leaked process detection on each GPU.
+
+The single most common reason a "healthy" cluster fails to launch a large
+training job is that a previous job's Python ranks are still attached to
+the GPUs (held HBM, half-torn-down NCCL communicators, or just stuck in
+__del__). Symptoms in the new training job: torch.cuda.OutOfMemoryError
+at model init with a misleading "free=Y" message, NCCL/RCCL bootstrap
+hang, or random ranks failing the first all-reduce due to compute
+contention. node-smoke catches these BEFORE the operator launches the
+real job by enumerating PIDs that hold each GPU and FAILing the node
+unless the operator explicitly opted in via --allow-foreign-procs.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from typing import Any, Dict, List, Optional, Tuple
+
+from ..shell_utils import _parse_size_with_unit, _which
+from .rocm_smi import _rocm_smi_processes
+
+# Process-name placeholders that some versions of `amd-smi process` and
+# `rocm-smi --showpids` emit when they cannot read the real name from
+# /proc/<pid>/comm (typically for kernel/system-owned PIDs like
+# `gpuagent`). We treat these as "missing" and fall back to /proc.
+_MISSING_NAME_TOKENS = frozenset({"", "n/a", "na", "none", "null", "-", "unknown", "?"})
+
+
+def _is_missing_name(name: str) -> bool:
+    return (name or "").strip().lower() in _MISSING_NAME_TOKENS
+
+
+def _resolve_proc_name(pid: int) -> str:
+    """Best-effort recover the process name for ``pid`` from /proc.
+
+    Tries ``/proc/<pid>/comm`` first (15-char kernel name, world-readable
+    on default Linux), then falls back to the ``Name:`` line in
+    ``/proc/<pid>/status``. Returns ``""`` when both reads fail (PID gone,
+    hidepid mount, ptrace_scope, etc.) so the caller can decide whether
+    to keep the original placeholder.
+    """
+    try:
+        with open(f"/proc/{int(pid)}/comm", "r", encoding="utf-8") as f:
+            name = f.read().strip()
+            if name:
+                return name
+    except OSError:
+        pass
+    try:
+        with open(f"/proc/{int(pid)}/status", "r", encoding="utf-8") as f:
+            for line in f:
+                if line.startswith("Name:"):
+                    return line.split(":", 1)[1].strip()
+    except OSError:
+        pass
+    return ""
+
+
+def _resolve_self_pid_view(self_pid: int) -> Dict[str, Any]:
+    """Resolve our own PID as the host kernel sees it, even from inside a
+    private PID namespace (e.g. a Docker / Kubernetes container).
+
+    ``amd-smi process``, ``rocm-smi --showpids`` and ``lsof /dev/kfd``
+    all report PIDs in the **root (host) PID namespace** because KFD is
+    a host-kernel resource that knows nothing about user namespaces.
+    ``os.getpid()``, by contrast, returns the PID *as our own namespace
+    sees it*. In a private PID namespace these are different numbers,
+    and the naive ``reported_pid == os.getpid()`` test would always
+    return False -- causing our own training rank to be flagged
+    ``is_foreign=True`` and (with the default policy) failing the node.
+
+    The kernel exposes the full mapping in ``/proc/self/status``:
+
+        NSpid:  2005679  42
+
+    where the first entry is the deepest-namespace PID (= the host PID
+    on bare metal) and the LAST entry is the most-deeply-nested PID
+    (= what ``os.getpid()`` returns inside a private namespace). On a
+    bare-metal host or a non-namespaced container, the line has a
+    single field equal to ``os.getpid()``.
+
+    Returns::
+
+        {
+            "host_pid": int,           # what amd-smi/rocm-smi report
+            "container_pid": int,      # == self_pid (passed in)
+            "pid_namespaced": bool,    # True if the two differ
+            "ns_chain": [int, ...],    # full NSpid chain (for the report)
+        }
+
+    Best-effort: if ``/proc/self/status`` cannot be read or has no
+    ``NSpid`` line, we assume bare-metal and fall back to ``self_pid``.
+    """
+    out: Dict[str, Any] = {
+        "host_pid": int(self_pid),
+        "container_pid": int(self_pid),
+        "pid_namespaced": False,
+        "ns_chain": [int(self_pid)],
+    }
+    try:
+        with open("/proc/self/status", "r", encoding="utf-8") as f:
+            for line in f:
+                if not line.startswith("NSpid:"):
+                    continue
+                parts = line.split()[1:]  # drop "NSpid:"
+                chain: List[int] = []
+                for tok in parts:
+                    try:
+                        chain.append(int(tok))
+                    except ValueError:
+                        continue
+                if not chain:
+                    break
+                out["ns_chain"] = chain
+                # Convention: NSpid lists outermost (host) namespace first
+                # and the current namespace last; matches /proc man page.
+                out["host_pid"] = chain[0]
+                out["container_pid"] = chain[-1]
+                out["pid_namespaced"] = chain[0] != chain[-1]
+                break
+    except Exception:
+        pass
+    return out
+
+
+def _collect_gpu_processes(
+    self_pid: int,
+    allowed_proc_names: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+    """Enumerate processes currently holding each GPU on this node.
+
+    Tries, in order:
+      1. ``amd-smi process --json``  (preferred -- structured)
+      2. ``amd-smi process``         (text fallback -- best-effort parse)
+      3. ``lsof /dev/kfd /dev/dri/renderD*``  (last resort: just openers)
+
+    Output shape::
+
+        {
+            "ok": bool,
+            "tool": "amd-smi process --json" | "amd-smi process" | "lsof" | None,
+            "self_pid": int,                # container/local view (legacy)
+            "self_pgid": int,               # pgid in our own namespace
+            "self_host_pid": int,           # PID amd-smi/rocm-smi report for us
+            "pid_namespaced": bool,         # True inside a private PID ns
+            "ns_pid_chain": [int, ...],     # full NSpid chain (root..ours)
+            "allowed_proc_names": [str, ...],   # passthrough for aggregator
+            "per_gpu": [
+                {"gpu": 0, "processes": [
+                    {"pid": int, "name": str, "hbm_bytes": int|None,
+                     "is_self": bool, "is_allowed": bool, "is_foreign": bool},
+                    ...
+                ]},
+                ...
+            ],
+            "foreign_count": int,   # PIDs not us and not allowed
+            "error": str            # only when ok is False
+        }
+
+    Filtering: a PID is treated as "ours" (and thus not foreign) if it
+    matches our **host-namespace PID** (because amd-smi / rocm-smi /
+    lsof always report root-ns PIDs) or, when we are NOT inside a
+    private PID namespace, our pgid (which catches per-GPU subprocesses
+    we may have spawned). Inside a private PID namespace the pgid match
+    is intentionally skipped: we cannot ``os.getpgid()`` PIDs we cannot
+    see, and any spawned subprocess will appear as a sibling host PID
+    that the operator should treat the same way as a leaked rank.
+
+    ``allowed_proc_names`` is a case-insensitive name allow-list for
+    known node-resident agents (``rocm-smi-daemon``, ``amd-smi``,
+    ``dcgm-exporter``, ``gpuagent``, ...).
+    """
+    allowed = {n.strip().lower() for n in (allowed_proc_names or []) if n.strip()}
+
+    # Resolve host-side PID. Inside a private PID namespace this is the
+    # number amd-smi / rocm-smi / lsof will report for us; on bare metal
+    # or a shared-PID-ns container (the SLURM + pyxis/enroot default) it
+    # equals self_pid.
+    pid_view = _resolve_self_pid_view(int(self_pid))
+    host_self_pid = int(pid_view["host_pid"])
+    pid_namespaced = bool(pid_view["pid_namespaced"])
+
+    # pgid is only meaningful within OUR own PID namespace -- attempting
+    # os.getpgid() on a host-side PID we cannot see would raise ESRCH.
+    try:
+        self_pgid = os.getpgid(int(self_pid))
+    except OSError:
+        self_pgid = int(self_pid)
+
+    out: Dict[str, Any] = {
+        "ok": False,
+        "tool": None,
+        "self_pid": int(self_pid),
+        "self_pgid": int(self_pgid),
+        "self_host_pid": host_self_pid,
+        "pid_namespaced": pid_namespaced,
+        "ns_pid_chain": list(pid_view.get("ns_chain") or [int(self_pid)]),
+        "allowed_proc_names": sorted(allowed),
+        "per_gpu": [],
+        "foreign_count": 0,
+    }
+
+    def _annotate(pid: int, name: str, hbm_bytes: Optional[int]) -> Dict[str, Any]:
+        # Direct PID match against our HOST-side PID (works in every
+        # mode: bare metal, shared PID ns, private PID ns).
+        is_self = int(pid) == host_self_pid
+        # pgid match is only safe outside a private PID namespace; inside
+        # one we cannot see host PIDs at all and os.getpgid() would
+        # ESRCH for every reported pid.
+        if not is_self and not pid_namespaced:
+            try:
+                pgid = os.getpgid(int(pid))
+            except OSError:
+                pgid = -1
+            if pgid == self_pgid:
+                is_self = True
+
+        # Some amd-smi / rocm-smi builds report `name="N/A"` (or empty)
+        # for kernel/system PIDs like `gpuagent` because they cannot read
+        # /proc/<pid>/comm themselves. We can: do it inline so the
+        # allowlist actually matches and the report shows the real name.
+        raw_name = name or ""
+        name_resolved_from_proc = False
+        if _is_missing_name(raw_name):
+            recovered = _resolve_proc_name(int(pid))
+            if recovered:
+                name = recovered
+                name_resolved_from_proc = True
+
+        is_allowed = (name or "").strip().lower() in allowed
+        return {
+            "pid": int(pid),
+            "name": name or "",
+            "name_raw": raw_name,
+            "name_resolved_from_proc": bool(name_resolved_from_proc),
+            "hbm_bytes": int(hbm_bytes) if isinstance(hbm_bytes, (int, float)) else None,
+            "is_self": bool(is_self),
+            "is_allowed": bool(is_allowed),
+            "is_foreign": bool(not is_self and not is_allowed),
+        }
+
+    if _which("amd-smi") is not None:
+        # 1) amd-smi process --json
+        try:
+            cp = subprocess.run(
+                ["amd-smi", "process", "--json"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=15,
+                check=False,
+            )
+            if cp.returncode == 0 and cp.stdout.strip():
+                try:
+                    doc = json.loads(cp.stdout)
+                    parsed, schema_drift = _flatten_amd_smi_process_json(doc, _annotate)
+                    if parsed and not schema_drift:
+                        # Schema fully recognized (Shape A/A'/B): trust the
+                        # result even when each per-GPU bucket is empty.
+                        out["per_gpu"] = parsed
+                        out["tool"] = "amd-smi process --json"
+                        out["ok"] = True
+                    elif schema_drift:
+                        # Top-level shape matched (Shape A bucket layout or
+                        # Shape B per-process dicts) but at least one record
+                        # had no usable pid / process_id -- the smoking gun
+                        # for a future amd-smi rename. Fall through to the
+                        # text / rocm-smi / lsof chain rather than silently
+                        # reporting the node clean.
+                        out["json_parse_error"] = (
+                            "amd-smi process --json: per-process schema not "
+                            "recognized (pid / process_id missing on at least "
+                            "one record); falling back to text / rocm-smi / lsof"
+                        )
+                    else:
+                        # Valid JSON but no recognizable per-GPU entries -- a
+                        # future amd-smi schema we don't speak yet. Fall
+                        # through to text / rocm-smi / lsof rather than
+                        # silently masking the node clean.
+                        out["json_parse_error"] = (
+                            "amd-smi process --json returned no recognizable "
+                            "per-GPU dicts; falling back to text / rocm-smi / lsof"
+                        )
+                except Exception as e:
+                    out["json_parse_error"] = str(e)
+            else:
+                out["json_rc"] = cp.returncode
+                out["json_stderr"] = (cp.stderr or "").strip()[:200]
+        except subprocess.TimeoutExpired:
+            out["json_error"] = "amd-smi process --json timed out"
+        except Exception as e:
+            out["json_error"] = str(e)
+
+        # 2) amd-smi process (text)
+        if not out["ok"]:
+            try:
+                cp = subprocess.run(
+                    ["amd-smi", "process"],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    timeout=15,
+                    check=False,
+                )
+                if cp.returncode == 0:
+                    parsed = _parse_amd_smi_process_text(cp.stdout, _annotate)
+                    out["per_gpu"] = parsed
+                    out["tool"] = "amd-smi process"
+                    out["ok"] = True
+                    out["raw"] = cp.stdout[:4000]
+                else:
+                    out["text_rc"] = cp.returncode
+                    out["text_stderr"] = (cp.stderr or "").strip()[:200]
+            except subprocess.TimeoutExpired:
+                out["text_error"] = "amd-smi process timed out"
+            except Exception as e:
+                out["text_error"] = str(e)
+
+    # 3) rocm-smi --showpids --json. Doesn't give us per-GPU mapping
+    # (--showpidgpus emits "WARNING: No JSON data to report"), so all
+    # PIDs go into the gpu=-1 bucket -- same convention as the lsof
+    # fallback. This is still a HUGE win over lsof: rocm-smi reports
+    # the actual KFD process name which lets the operator decide if
+    # it's a leaked rank vs a known agent.
+    if not out["ok"]:
+        rocm = _rocm_smi_processes(_annotate)
+        if rocm.get("ok"):
+            out["per_gpu"] = rocm.get("per_gpu") or []
+            out["tool"] = rocm.get("tool")
+            out["ok"] = True
+        else:
+            out["rocm_smi_error"] = rocm.get("error")
+
+    # 4) lsof on /dev/kfd + /dev/dri/renderD* (we cannot map back to specific
+    # GPUs reliably this way, but at least we surface foreign openers).
+    if not out["ok"]:
+        try:
+            import glob as _glob
+
+            paths = ["/dev/kfd"] + sorted(_glob.glob("/dev/dri/renderD*"))
+            existing = [p for p in paths if os.path.exists(p)]
+            if existing and _which("lsof") is not None:
+                cp = subprocess.run(
+                    ["lsof", "-Fpcn", "--", *existing],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    timeout=10,
+                    check=False,
+                )
+                if cp.returncode in (0, 1):  # lsof exits 1 when nothing open
+                    procs = _parse_lsof_pcn(cp.stdout, _annotate)
+                    # Without per-GPU mapping, surface as a single bucket
+                    # under gpu=-1 so the aggregator still flags foreigners.
+                    out["per_gpu"] = [{"gpu": -1, "processes": procs}] if procs else []
+                    out["tool"] = "lsof"
+                    out["ok"] = True
+                else:
+                    out["lsof_rc"] = cp.returncode
+        except Exception as e:
+            out["lsof_error"] = str(e)
+
+    if not out["ok"] and "error" not in out:
+        out["error"] = "no working enumeration tool " "(amd-smi process / rocm-smi --showpids / lsof)"
+
+    out["foreign_count"] = sum(
+        1 for g in out["per_gpu"] for p in (g.get("processes") or []) if p.get("is_foreign")
+    )
+    return out
+
+
+def _flatten_amd_smi_process_json(
+    doc: Any,
+    annotate: Any,
+) -> Tuple[List[Dict[str, Any]], bool]:
+    """Coerce the amd-smi process JSON into a stable per-GPU shape.
+
+    The schema varies across releases; we tolerate missing fields silently.
+    Three top-level shapes are seen in the wild:
+
+      A) Per-GPU dicts each carrying ``process_list``, where each list
+         item wraps the actual process under ``process_info`` and reports
+         memory as ``{"value": N, "unit": "B"}`` -- this is the modern
+         ``amd-smi`` (>= 6.x) layout::
+
+           [{"gpu": 0, "process_list": [
+               {"process_info": {
+                   "pid": 12345, "name": "python",
+                   "memory_usage": {"vram_mem": {"value": 256, "unit": "MB"}}
+               }}
+           ]}]
+
+      A') Same as A but each ``process_list`` item is the process dict
+          directly (older amd-smi releases), with memory as a flat int or
+          formatted string::
+
+           [{"gpu": 0, "process_list": [
+               {"pid": 12345, "name": "python",
+                "memory_usage": {"vram_mem": "256 MB"}}
+           ]}]
+
+      B) Top-level list of per-process dicts each carrying explicit
+         ``gpu`` / ``gpus`` (uncommon, but seen on a couple of branches)::
+
+           [{"pid": ..., "name": ..., "gpu": 0, "memory_usage": {...}}]
+
+    Returns ``(per_gpu, schema_drift_suspected)``. The drift flag is raised
+    when at least one record looked like a process dict at the structural
+    level (i.e. we entered ``_push``) but the per-process pid/process_id
+    field was missing or had a non-numeric value -- the smoking gun for a
+    future amd-smi schema rename. The caller uses this to fall through to
+    the text/rocm-smi/lsof fallbacks instead of trusting an empty result.
+    A genuinely clean node (empty ``process_list``) never enters ``_push``
+    so the flag stays False and the fast path is preserved.
+    """
+    out_by_gpu: Dict[int, List[Dict[str, Any]]] = {}
+    schema_drift_suspected = False
+
+    def _push(gpu: int, pid: Any, name: Any, hbm: Any) -> None:
+        nonlocal schema_drift_suspected
+        if not isinstance(pid, (int, float)):
+            schema_drift_suspected = True
+            return
+        ann = annotate(int(pid), str(name or ""), hbm)
+        out_by_gpu.setdefault(int(gpu), []).append(ann)
+
+    def _value_unit_to_bytes(v: Any) -> Optional[int]:
+        """Resolve a size value that may be int, formatted string, or
+        ``{"value": N, "unit": "B|KB|MB|GB|..."}`` (modern amd-smi shape)
+        into bytes."""
+        if isinstance(v, (int, float)):
+            return int(v)
+        if isinstance(v, str):
+            return _parse_size_with_unit(v)
+        if isinstance(v, dict):
+            val = v.get("value")
+            unit = v.get("unit")
+            if isinstance(val, (int, float)):
+                if isinstance(unit, str) and unit.strip():
+                    return _parse_size_with_unit(f"{val} {unit.strip()}")
+                return int(val)
+            if isinstance(val, str):
+                if isinstance(unit, str) and unit.strip():
+                    return _parse_size_with_unit(f"{val} {unit.strip()}")
+                return _parse_size_with_unit(val)
+        return None
+
+    def _hbm_of(d: Dict[str, Any]) -> Optional[int]:
+        # Preferred path: memory_usage.vram_mem (covers shapes A, A', B).
+        mu = d.get("memory_usage") if isinstance(d.get("memory_usage"), dict) else None
+        if mu is not None:
+            for k in ("vram_mem", "vram_memory", "vram"):
+                if k in mu:
+                    n = _value_unit_to_bytes(mu.get(k))
+                    if n is not None:
+                        return n
+        # Fallback: mem_usage at the same level (older shapes -- usually
+        # mirrors vram_mem on AMD GPUs since GTT/CPU are negligible).
+        if "mem_usage" in d:
+            n = _value_unit_to_bytes(d.get("mem_usage"))
+            if n is not None:
+                return n
+        # Last resort: a flat "vram" key directly under d.
+        if "vram" in d:
+            n = _value_unit_to_bytes(d.get("vram"))
+            if n is not None:
+                return n
+        return None
+
+    def _unwrap_proc(p: Dict[str, Any]) -> Dict[str, Any]:
+        """Modern amd-smi wraps each process under ``process_info``;
+        unwrap so the rest of the parser can read pid/name/memory at
+        the top level uniformly."""
+        if isinstance(p.get("process_info"), dict):
+            return p["process_info"]
+        return p
+
+    items = (
+        doc
+        if isinstance(doc, list)
+        else (doc.get("processes") or doc.get("gpus") or [] if isinstance(doc, dict) else [])
+    )
+    if not isinstance(items, list):
+        items = []
+
+    for item in items:
+        if not isinstance(item, dict):
+            continue
+        # Shape A / A': per-GPU dict with process_list. Use explicit
+        # presence checks rather than `... or ...` -- an empty list is
+        # the *clean GPU* signal we want to honor, and `[] or x` would
+        # short-circuit past it to ``processes``, then to None, and miss
+        # the Shape A bucket pre-registration entirely.
+        plist = item.get("process_list")
+        if not isinstance(plist, list):
+            plist = item.get("processes")
+        if isinstance(plist, list):
+            gpu_idx = item.get("gpu")
+            if not isinstance(gpu_idx, int):
+                gpu_idx = -1
+            # Register the GPU bucket as soon as we recognize a Shape A
+            # entry, even if its process_list is empty. This makes an
+            # empty return value mean "schema didn't match" rather than
+            # the ambiguous "schema matched but no processes" -- so the
+            # caller can safely fall through to the text / rocm-smi /
+            # lsof fallbacks on a future amd-smi schema.
+            out_by_gpu.setdefault(int(gpu_idx), [])
+            for p in plist:
+                if not isinstance(p, dict):
+                    continue
+                proc = _unwrap_proc(p)
+                _push(
+                    gpu_idx,
+                    proc.get("pid") if proc.get("pid") is not None else proc.get("process_id"),
+                    proc.get("name") or proc.get("process_name"),
+                    _hbm_of(proc),
+                )
+            continue
+        # Shape B: per-process dict with explicit gpu
+        proc = _unwrap_proc(item)
+        if "pid" in proc or "process_id" in proc:
+            g = proc.get("gpu")
+            gpus = (
+                proc.get("gpus")
+                if isinstance(proc.get("gpus"), list)
+                else ([g] if isinstance(g, int) else [-1])
+            )
+            for gpu_idx in gpus:
+                if not isinstance(gpu_idx, int):
+                    gpu_idx = -1
+                _push(
+                    gpu_idx,
+                    proc.get("pid") if proc.get("pid") is not None else proc.get("process_id"),
+                    proc.get("name") or proc.get("process_name"),
+                    _hbm_of(proc),
+                )
+
+    return (
+        [{"gpu": k, "processes": v} for k, v in sorted(out_by_gpu.items())],
+        schema_drift_suspected,
+    )
+
+
+def _parse_amd_smi_process_text(text: str, annotate: Any) -> List[Dict[str, Any]]:
+    """Best-effort parser for ``amd-smi process`` plain-text output.
+
+    Format varies, but typical layout is one record per process with lines
+    like ``GPU: 0`` / ``PID: 12345`` / ``NAME: python`` / ``VRAM_MEM: 256 MB``.
+    We tokenise key-value pairs case-insensitively and group records by
+    blank-line separators.
+    """
+    out_by_gpu: Dict[int, List[Dict[str, Any]]] = {}
+    cur: Dict[str, Any] = {}
+    blocks: List[Dict[str, Any]] = []
+    for raw in text.splitlines():
+        line = raw.strip()
+        if not line:
+            if cur:
+                blocks.append(cur)
+                cur = {}
+            continue
+        if ":" in line:
+            k, _, v = line.partition(":")
+            cur[k.strip().lower()] = v.strip()
+    if cur:
+        blocks.append(cur)
+    for b in blocks:
+        gpu_s = b.get("gpu") or b.get("gpu_id") or b.get("device") or "-1"
+        try:
+            gpu = int(gpu_s.split()[0])
+        except Exception:
+            gpu = -1
+        try:
+            pid = int(b.get("pid", "").split()[0])
+        except Exception:
+            continue
+        name = b.get("name") or b.get("process") or b.get("process_name") or ""
+        hbm_raw = b.get("vram_mem") or b.get("vram") or b.get("memory_usage") or b.get("mem_usage") or ""
+        hbm = _parse_size_with_unit(hbm_raw) if hbm_raw else None
+        out_by_gpu.setdefault(gpu, []).append(annotate(pid, name, hbm))
+    return [{"gpu": k, "processes": v} for k, v in sorted(out_by_gpu.items())]
+
+
+def _parse_lsof_pcn(text: str, annotate: Any) -> List[Dict[str, Any]]:
+    """Parse ``lsof -Fpcn`` output (one field per line, prefixed by f-code).
+
+    We only need ``p<pid>`` and ``c<command>``. Returns one annotated
+    record per unique PID; HBM bytes unknown (lsof can't measure that).
+    """
+    out: Dict[int, Dict[str, Any]] = {}
+    cur_pid: Optional[int] = None
+    cur_name = ""
+    for line in text.splitlines():
+        if not line:
+            continue
+        tag, val = line[0], line[1:]
+        if tag == "p":
+            try:
+                cur_pid = int(val)
+            except ValueError:
+                cur_pid = None
+            cur_name = ""
+        elif tag == "c" and cur_pid is not None:
+            cur_name = val
+            out[cur_pid] = annotate(cur_pid, cur_name, None)
+    return list(out.values())
diff --git a/primus/tools/preflight/node_smoke/collectors/host_limits.py b/primus/tools/preflight/node_smoke/collectors/host_limits.py
new file mode 100644
index 000000000..d0e9a0e5c
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/host_limits.py
@@ -0,0 +1,92 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- C. Host limits (ulimit -l, /dev/shm, NUMA, CPU governor)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List
+
+from ..shell_utils import _read_text
+
+
+def _collect_host_limits(*, ulimit_l_min_gb: float, shm_min_gb: float) -> Dict[str, Any]:
+    """Capture training-relevant kernel/process limits and tunables and
+    return hard-failure reasons for the ones that block training under load.
+
+    Hard fail today (cause node FAIL):
+
+    * ``ulimit -l`` (RLIMIT_MEMLOCK) is not unlimited and below
+      ``ulimit_l_min_gb`` -- RDMA pin failures look like NCCL hangs.
+    * ``/dev/shm`` total size below ``shm_min_gb`` -- NCCL shared-memory
+      transport falls back or fails.
+
+    Soft (collected for drift detection only):
+
+    * NUMA node count, CPU count, CPU governor, kernel/OS version. The
+      aggregator flags drift across the cluster but does not FAIL nodes
+      individually for these.
+    """
+    out: Dict[str, Any] = {}
+
+    # Resource limits.
+    try:
+        import resource  # type: ignore
+
+        soft_l, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK)
+        out["memlock_soft_bytes"] = -1 if soft_l == resource.RLIM_INFINITY else int(soft_l)
+        soft_n, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+        out["nofile_soft"] = int(soft_n)
+        soft_p, _ = resource.getrlimit(resource.RLIMIT_NPROC)
+        out["nproc_soft"] = -1 if soft_p == resource.RLIM_INFINITY else int(soft_p)
+    except Exception as e:
+        out["resource_error"] = str(e)
+
+    # /dev/shm size + free.
+    try:
+        st = os.statvfs("/dev/shm")
+        out["shm_size_bytes"] = int(st.f_blocks) * int(st.f_frsize)
+        out["shm_avail_bytes"] = int(st.f_bavail) * int(st.f_frsize)
+    except Exception as e:
+        out["shm_error"] = str(e)
+
+    # NUMA topology.
+    try:
+        nodes = [
+            n for n in os.listdir("/sys/devices/system/node") if n.startswith("node") and n[4:].isdigit()
+        ]
+        out["numa_nodes"] = len(nodes)
+    except Exception:
+        out["numa_nodes"] = None
+
+    # CPU count + governor.
+    try:
+        out["cpu_count"] = os.cpu_count()
+    except Exception:
+        out["cpu_count"] = None
+    out["cpu_governor"] = _read_text("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor") or None
+
+    # Hard checks.
+    fail_reasons: List[str] = []
+    memlock = out.get("memlock_soft_bytes")
+    if memlock is not None and memlock != -1 and ulimit_l_min_gb > 0:
+        if memlock < ulimit_l_min_gb * (1 << 30):
+            fail_reasons.append(
+                f"ulimit -l (memlock) = {memlock // (1 << 20)} MiB; "
+                f"required: unlimited or >= {ulimit_l_min_gb} GiB. "
+                "RDMA pin will fail under load."
+            )
+    shm = out.get("shm_size_bytes")
+    if shm is not None and shm_min_gb > 0:
+        if shm < shm_min_gb * (1 << 30):
+            fail_reasons.append(
+                f"/dev/shm size = {shm / (1 << 30):.2f} GiB; "
+                f"required: >= {shm_min_gb} GiB. NCCL shared-mem may fail."
+            )
+    out["fail_reasons"] = fail_reasons
+
+    return out
diff --git a/primus/tools/preflight/node_smoke/collectors/nics.py b/primus/tools/preflight/node_smoke/collectors/nics.py
new file mode 100644
index 000000000..8266b18e6
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/nics.py
@@ -0,0 +1,452 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- B. NIC / RDMA roll-call (per-port state + GIDs from sysfs).
+
+Selector chain
+==============
+
+Many clusters expose more RDMA-capable ports in ``/sys/class/infiniband/``
+than the training job will actually use. Common reasons:
+
+* a separate front-end / management RoCE NIC that is physically present
+  but admin-disabled (no SFP, BIOS port-disable, netdev down) -- typical
+  ``state=DOWN phys_state=Disabled`` ports;
+* a storage / control-plane RoCE NIC that is fully up but is reserved
+  for sockets (it shows up in ``NCCL_SOCKET_IFNAME`` rather than
+  ``NCCL_IB_HCA``).
+
+Failing the node on those ports is wrong: the operator already told
+NCCL/RCCL not to use them, and the training job will run fine. So the
+collector resolves a *training-NIC selector* and only enforces the hard
+rules (state must be ACTIVE, phys_state must be LinkUp, RoCE v2 GID
+present, ...) on the included subset. Excluded ports stay in
+``ports`` for diagnostics and are summarised in ``info_issues``.
+
+Selector precedence (highest first):
+
+1. ``allowlist`` argument -- typically wired to ``--rdma-nic-allowlist``
+   on the CLI. NCCL syntax (see :func:`_parse_nic_selector` below).
+2. ``NCCL_IB_HCA`` env -- mirrors what NCCL/RCCL itself will use, so the
+   smoke test and the training launch agree by construction.
+3. ``phys_state in {Disabled, Sleep}`` heuristic -- ports in those
+   administratively-down phys_states are intentionally not in use
+   (admin-disabled at firmware/driver level, no SFP, BIOS-disabled, ...)
+   and are auto-excluded. Crucially, real failure modes on a port that
+   *is* intended to be used produce a different phys_state
+   (``Polling``, ``LinkErrorRecovery``, ``LinkUp`` with ``state!=ACTIVE``,
+   ...), so this heuristic does not mask cable / driver problems.
+4. Fallback: every IB port must be ACTIVE / LinkUp (the strict rule).
+
+Empty-set guard: if the resolved included set is empty (every port was
+excluded by some combination of the above) but the host has at least
+one IB port at all, that's still a node FAIL -- a node with zero
+training NICs cannot participate in inter-node training.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+from ..shell_utils import _read_text
+
+# Phys-states we treat as "administratively off, not a failure".
+# See module docstring for why this set is conservative.
+_ADMIN_DOWN_PHYS_STATES: Set[str] = {"disabled", "sleep"}
+
+
+def _parse_nic_selector(raw: str) -> Dict[str, Any]:
+    """Parse an ``NCCL_IB_HCA``-style selector string.
+
+    Accepts the full NCCL syntax:
+
+    * Comma-separated entries of the form ``device[:port]``. ``port``
+      is an integer; if omitted, the entry matches any port on the
+      device.
+    * Optional ``^`` prefix on the *whole* string makes it a denylist
+      (every device matches except the listed ones).
+    * Optional ``=`` prefix on an individual entry forces an exact
+      device-name match (e.g. ``=mlx5`` matches device ``mlx5`` only,
+      not ``mlx5_0``). Without ``=``, an entry is a *prefix* match,
+      matching ``mlx5`` against ``mlx5_0`` / ``mlx5_1`` / ...
+
+    Returns a dict with:
+
+    * ``mode``: ``"allowlist"`` or ``"denylist"``.
+    * ``entries``: list of ``(device_pattern, exact_match, port_or_None)``.
+    * ``raw``: the original input, after stripping the global ``^``.
+
+    Returns ``{"mode": "passthrough", ...}`` for empty / whitespace-only
+    input, signalling the caller to fall through to the next selector.
+    """
+    s = (raw or "").strip()
+    if not s:
+        return {"mode": "passthrough", "entries": [], "raw": raw or ""}
+
+    mode = "allowlist"
+    if s.startswith("^"):
+        mode = "denylist"
+        s = s[1:].strip()
+
+    entries: List[Tuple[str, bool, Optional[int]]] = []
+    for token in s.split(","):
+        t = token.strip()
+        if not t:
+            continue
+        exact = False
+        if t.startswith("="):
+            exact = True
+            t = t[1:].strip()
+            if not t:
+                continue
+        # Split off an optional :port suffix. We split from the right so
+        # device names containing ':' (none in the wild today, but we
+        # don't want to be the thing that breaks if they appear) are
+        # preserved.
+        port: Optional[int] = None
+        if ":" in t:
+            dev_part, _, port_str = t.rpartition(":")
+            try:
+                port = int(port_str)
+                t = dev_part
+            except ValueError:
+                # Not actually a port suffix -- treat the whole thing
+                # as the device name.
+                port = None
+        entries.append((t, exact, port))
+
+    if not entries:
+        return {"mode": "passthrough", "entries": [], "raw": raw or ""}
+    return {"mode": mode, "entries": entries, "raw": raw or ""}
+
+
+def _selector_matches(
+    sel: Dict[str, Any],
+    device: str,
+    port: int,
+) -> bool:
+    """Return True iff ``(device, port)`` matches an allow/denylist selector
+    parsed by :func:`_parse_nic_selector`.
+
+    ``passthrough`` selectors trivially match everything (the caller
+    should not normally pass them here; instead it should treat them as
+    "no selector at all" and fall through to the next layer).
+    """
+    mode = sel.get("mode", "passthrough")
+    if mode == "passthrough":
+        return True
+    entries = sel.get("entries") or []
+    listed = False
+    for dev_pat, exact, want_port in entries:
+        # Device name match: exact iff the entry started with '='; else
+        # prefix match (NCCL semantics).
+        if exact:
+            dev_ok = device == dev_pat
+        else:
+            dev_ok = device.startswith(dev_pat)
+        if not dev_ok:
+            continue
+        # Port filter: an entry with no :port matches any port on that
+        # device.
+        if want_port is not None and want_port != port:
+            continue
+        listed = True
+        break
+    if mode == "allowlist":
+        return listed
+    # Denylist: ports NOT in the list are accepted.
+    return not listed
+
+
+def _resolve_selector(
+    allowlist_arg: Optional[str],
+    env: Optional[Dict[str, str]] = None,
+) -> Dict[str, Any]:
+    """Resolve the training-NIC selector for this node.
+
+    Precedence:
+
+    1. ``allowlist_arg`` (typically the value of ``--rdma-nic-allowlist``).
+    2. ``NCCL_IB_HCA`` env var.
+    3. Heuristic fallback (``"heuristic"``) -- the caller handles this
+       by auto-excluding ports whose ``phys_state`` is in
+       :data:`_ADMIN_DOWN_PHYS_STATES`.
+
+    Returns a dict with ``source`` (``"cli"`` / ``"env"`` / ``"heuristic"``)
+    plus the parsed-selector fields when applicable.
+    """
+    if env is None:
+        env = dict(os.environ)
+    # 1. CLI / explicit arg.
+    if allowlist_arg:
+        parsed = _parse_nic_selector(allowlist_arg)
+        if parsed.get("mode") != "passthrough":
+            parsed["source"] = "cli"
+            return parsed
+    # 2. NCCL_IB_HCA env.
+    env_raw = env.get("NCCL_IB_HCA") or ""
+    if env_raw.strip():
+        parsed = _parse_nic_selector(env_raw)
+        if parsed.get("mode") != "passthrough":
+            parsed["source"] = "env"
+            return parsed
+    # 3. Heuristic fallback.
+    return {
+        "source": "heuristic",
+        "mode": "heuristic",
+        "entries": [],
+        "raw": "",
+        "admin_down_phys_states": sorted(_ADMIN_DOWN_PHYS_STATES),
+    }
+
+
+def _collect_nic_status(
+    expected_count: Optional[int],
+    *,
+    allowlist: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Inventory every RDMA port on this node and flag the ones that would
+    silently break inter-node training.
+
+    Reads everything from ``/sys/class/infiniband`` (kernel ``ib_core``
+    ABI) so the check is **vendor- and fabric-agnostic** -- works on
+    Mellanox/NVIDIA (``mlx5_ib``), Broadcom (``bnxt_re``), Intel
+    (``irdma``), Marvell (``qedr``), AWS EFA (``efa``), Huawei
+    (``hns_roce``), etc., and on either RoCE-over-Ethernet or true
+    InfiniBand fabrics. We do not depend on ``ibv_devinfo`` /
+    ``ibstat`` / vendor SDKs being present in the container.
+
+    Per port we capture:
+
+    * ``link_layer`` (``Ethernet`` for RoCE, ``InfiniBand`` for IB) --
+      determines which GID rule applies below;
+    * link state (``state``: ``ACTIVE``/``DOWN``/``INIT``) and physical
+      state (``phys_state``: ``LinkUp``/``Polling``/...);
+    * link rate (Gb/s);
+    * netdev + MTU (so the aggregator can detect MTU drift, which silently
+      tanks RoCE all-reduce throughput);
+    * GID counts -- total non-zero GIDs and the subset configured as
+      ``RoCE v2`` (an empty RoCE v2 set is a frequent cause of training
+      jobs hanging at the first inter-node collective on RoCE clusters).
+
+    Per-port hard issues are only emitted for ports in the resolved
+    *training-NIC selector* (see module docstring). Excluded ports stay
+    in ``ports`` for diagnostics and are summarised in ``info_issues``.
+
+    Issues are pushed into ``out["issues"]`` (each a short string). Hard
+    issues (port not Active / missing GIDs / wrong NIC count) are
+    treated as node FAIL by ``_node_status_from``. The GID check is
+    fabric-aware:
+
+    * RoCE/Ethernet port must have at least one ``RoCE v2`` GID
+      configured (RoCEv1 is essentially unused for AI training).
+    * InfiniBand port must have at least one valid (non-zero) GID --
+      it is normally auto-populated by the SM; an empty GID table on
+      an ACTIVE IB port indicates a subnet-manager problem.
+    * Unknown ``link_layer`` (very old kernels) falls back to "must
+      have any valid GID" so we don't false-FAIL exotic configurations.
+    """
+    selector = _resolve_selector(allowlist)
+    out: Dict[str, Any] = {
+        "expected_count": expected_count,
+        "selector": selector,
+        "ports": [],
+        "included_ports": [],
+        "excluded_ports": [],
+        "issues": [],
+        "info_issues": [],
+    }
+    base = "/sys/class/infiniband"
+    if not os.path.isdir(base):
+        # Container may not expose the IB stack; report and let the operator
+        # decide. We only mark this as a hard issue when the user explicitly
+        # asked for a positive expected_count.
+        msg = f"{base} missing -- no RDMA stack visible"
+        if expected_count and expected_count > 0:
+            out["issues"].append(msg)
+        else:
+            out["info"] = msg
+        return out
+
+    try:
+        devs = sorted(os.listdir(base))
+    except Exception as e:
+        out["issues"].append(f"failed to list {base}: {e}")
+        return out
+
+    for dev in devs:
+        port_dir = os.path.join(base, dev, "ports")
+        if not os.path.isdir(port_dir):
+            continue
+        try:
+            ports = sorted(os.listdir(port_dir))
+        except Exception:
+            continue
+        for port_str in ports:
+            try:
+                port = int(port_str)
+            except ValueError:
+                continue
+            p = os.path.join(port_dir, port_str)
+
+            # Sysfs values look like "4: ACTIVE" / "5: LinkUp" / "400 Gb/sec (4X NDR)"
+            state_raw = _read_text(os.path.join(p, "state"))
+            phys_raw = _read_text(os.path.join(p, "phys_state"))
+            rate_raw = _read_text(os.path.join(p, "rate"))
+            state = state_raw.split(":", 1)[-1].strip() if state_raw else ""
+            phys = phys_raw.split(":", 1)[-1].strip() if phys_raw else ""
+            rate_gbps: Optional[int] = None
+            try:
+                rate_gbps = int(rate_raw.split()[0])
+            except Exception:
+                pass
+
+            # Fabric type: "Ethernet" -> RoCE, "InfiniBand" -> IB.
+            # Determines which GID rule to apply below. Provided by
+            # ib_core for every RDMA driver since Linux 3.x; missing
+            # only on extremely old kernels.
+            link_layer = (_read_text(os.path.join(p, "link_layer")) or "").strip() or None
+
+            # GID inventory. A GID is "all-zero" until configured.
+            gid_count = 0
+            rocev2_count = 0
+            gids_dir = os.path.join(p, "gids")
+            types_dir = os.path.join(p, "gid_attrs", "types")
+            valid_gid_indices: List[int] = []
+            if os.path.isdir(gids_dir):
+                try:
+                    for gn in sorted(os.listdir(gids_dir), key=lambda s: int(s) if s.isdigit() else 0):
+                        if not gn.isdigit():
+                            continue
+                        g = _read_text(os.path.join(gids_dir, gn))
+                        if g and g != "0000:0000:0000:0000:0000:0000:0000:0000":
+                            gid_count += 1
+                            valid_gid_indices.append(int(gn))
+                except Exception:
+                    pass
+            if os.path.isdir(types_dir):
+                for idx in valid_gid_indices:
+                    t = _read_text(os.path.join(types_dir, str(idx)))
+                    if "RoCE v2" in t or "RoCEv2" in t:
+                        rocev2_count += 1
+
+            # Linked netdev + MTU.
+            ifname: Optional[str] = None
+            mtu: Optional[int] = None
+            net_dir = os.path.join(base, dev, "device", "net")
+            if os.path.isdir(net_dir):
+                try:
+                    nets = sorted(os.listdir(net_dir))
+                    if nets:
+                        ifname = nets[0]
+                        mtu_raw = _read_text(f"/sys/class/net/{ifname}/mtu")
+                        try:
+                            mtu = int(mtu_raw)
+                        except Exception:
+                            mtu = None
+                except Exception:
+                    pass
+
+            port_rec = {
+                "device": dev,
+                "port": port,
+                "link_layer": link_layer,
+                "state": state or None,
+                "phys_state": phys or None,
+                "rate_gbps": rate_gbps,
+                "ifname": ifname,
+                "mtu": mtu,
+                "gid_count": gid_count,
+                "rocev2_gid_count": rocev2_count,
+            }
+            out["ports"].append(port_rec)
+
+            # ------------------------------------------------------------
+            # Selector: is this port part of the training-NIC set?
+            # ------------------------------------------------------------
+            label = f"{dev}:{port}"
+            include = True
+            exclude_reason = ""
+            if selector.get("source") == "heuristic":
+                # Auto-exclude admin-down phys_states; keep everything
+                # else and let the strict per-port rules run.
+                if (phys or "").lower() in _ADMIN_DOWN_PHYS_STATES:
+                    include = False
+                    exclude_reason = f"phys_state={phys} (admin-disabled, not used for training)"
+            else:
+                # Explicit allow/denylist (CLI or env).
+                if not _selector_matches(selector, dev, port):
+                    include = False
+                    src = selector.get("source") or "selector"
+                    src_label = (
+                        "NCCL_IB_HCA" if src == "env" else "--rdma-nic-allowlist" if src == "cli" else src
+                    )
+                    extra = ""
+                    if (state or "").upper() != "ACTIVE" or (phys or "").lower() != "linkup":
+                        extra = f" (state={state} phys_state={phys})"
+                    exclude_reason = f"excluded by {src_label}{extra}"
+
+            if not include:
+                out["excluded_ports"].append(label)
+                out["info_issues"].append(f"{label} {exclude_reason}")
+                continue
+            out["included_ports"].append(label)
+
+            # ------------------------------------------------------------
+            # Per-port hard issues -> node FAIL. Only for included ports.
+            # ------------------------------------------------------------
+            if state and state.upper() != "ACTIVE":
+                out["issues"].append(f"{dev}:{port} state={state} (expected ACTIVE)")
+            if phys and phys.upper() != "LINKUP":
+                out["issues"].append(f"{dev}:{port} phys_state={phys} (expected LinkUp)")
+            # Fabric-aware GID requirement on ACTIVE ports.
+            if state.upper() == "ACTIVE":
+                ll = (link_layer or "").lower()
+                if ll == "ethernet":
+                    # RoCE: needs at least one RoCE v2 GID; RoCEv1 is
+                    # essentially unused for AI training and we treat
+                    # its absence as a hard fail.
+                    if rocev2_count == 0:
+                        out["issues"].append(
+                            f"{dev}:{port} no RoCE v2 GIDs configured " f"(RoCE/Ethernet fabric)"
+                        )
+                elif ll == "infiniband":
+                    # True IB: subnet manager normally populates GIDs.
+                    # Empty GID table on an ACTIVE IB port = SM problem.
+                    if gid_count == 0:
+                        out["issues"].append(
+                            f"{dev}:{port} no GIDs populated " f"(InfiniBand fabric -- check subnet manager)"
+                        )
+                else:
+                    # Unknown / missing link_layer (very old kernel):
+                    # require at least one valid GID rather than a
+                    # specific type, to avoid false FAIL.
+                    if gid_count == 0:
+                        out["issues"].append(
+                            f"{dev}:{port} no valid GIDs configured " f"(link_layer unknown)"
+                        )
+
+    # The expected-count check operates on the *included* set: operators
+    # almost always set --expected-rdma-nics to mean "training NIC count"
+    # (e.g. 8 = 1 NIC per GPU), not "everything visible under
+    # /sys/class/infiniband/" -- which on multi-role nodes includes
+    # frontend / management / storage RoCE NICs.
+    if expected_count is not None and len(out["included_ports"]) != expected_count:
+        out["issues"].append(f"RDMA NIC port count {len(out['included_ports'])} != expected {expected_count}")
+
+    # Empty-set guard: if we saw IB ports on this host but every single
+    # one got excluded, the node cannot participate in inter-node
+    # training. This catches "the whole RoCE card disabled itself" and
+    # similar disasters that would otherwise PASS silently because every
+    # port met the exclusion criteria.
+    if out["ports"] and not out["included_ports"]:
+        out["issues"].append(
+            "no included RDMA NIC ports remain after selector "
+            f"({selector.get('source')}); node cannot participate in inter-node training"
+        )
+
+    return out
diff --git a/primus/tools/preflight/node_smoke/collectors/reused_info.py b/primus/tools/preflight/node_smoke/collectors/reused_info.py
new file mode 100644
index 000000000..ef0f3236e
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/reused_info.py
@@ -0,0 +1,51 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Reused gpu/host/network info collectors from the rest of the preflight tree.
+
+These collectors already work without a global PG and produce ``Finding``
+objects (level='fail' counts as a node failure). We import each one
+lazily so a missing dependency in one section doesn't cascade.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from ..shell_utils import _findings_to_dicts
+
+
+def _collect_reused_info() -> Dict[str, Any]:
+    """Run the existing host/gpu/network info collectors. They already work
+    without a global PG and produce ``Finding`` objects (level='fail' counts
+    as a node failure)."""
+    section: Dict[str, Any] = {"gpu_info": [], "host_info": [], "network_info": []}
+    try:
+        from primus.tools.preflight.gpu.info import collect_gpu_info
+
+        section["gpu_info"] = _findings_to_dicts(collect_gpu_info())
+    except Exception as e:
+        section["gpu_info"] = [
+            {"level": "warn", "message": "collect_gpu_info raised", "details": {"error": str(e)}}
+        ]
+    try:
+        from primus.tools.preflight.host.info import collect_host_info
+
+        section["host_info"] = _findings_to_dicts(collect_host_info())
+    except Exception as e:
+        section["host_info"] = [
+            {"level": "warn", "message": "collect_host_info raised", "details": {"error": str(e)}}
+        ]
+    try:
+        from primus.tools.preflight.network.info import collect_network_info
+
+        # expect_distributed=False so we don't WARN about a missing world PG.
+        section["network_info"] = _findings_to_dicts(collect_network_info(expect_distributed=False))
+    except Exception as e:
+        section["network_info"] = [
+            {"level": "warn", "message": "collect_network_info raised", "details": {"error": str(e)}}
+        ]
+    return section
diff --git a/primus/tools/preflight/node_smoke/collectors/rocm_smi.py b/primus/tools/preflight/node_smoke/collectors/rocm_smi.py
new file mode 100644
index 000000000..a5b301dab
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/rocm_smi.py
@@ -0,0 +1,409 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- F: rocm-smi self-latency + cross-tool fallbacks for amd-smi checks.
+
+rocm-smi is "usually available" even on nodes that lack amd-smi (older
+ROCm installs, stripped-down containers). Every amd-smi check we
+silently no-op when amd-smi is missing has a rocm-smi equivalent, so
+we can keep ECC / XGMI / foreign-process / activity coverage even
+without amd-smi. Each helper produces the same record shape the
+upstream amd-smi parser already emits, so _node_status_from and the
+aggregator helpers don't need any per-tool conditionals.
+
+Output schemas (cross-tool stable):
+
+  ECC      -> per-GPU {gpu, ecc_correctable_total, ecc_uncorrectable_total}
+  XGMI     -> {ok, tool, n_gpus, link_types: [[...]], non_xgmi_pairs}
+  procs    -> per-GPU {gpu, processes: [annotated PID dicts]}
+  activity -> per-GPU {gpu, gfx_activity_pct}
+
+The self-latency canary (``_collect_rocm_smi_self_latency``) is its own
+Tier 1 F check: a wedged amdgpu driver makes ``rocm-smi`` calls take
+30-60 s before failing outright -- usually minutes before the GPU
+itself stops responding. Hitting the timeout is treated as a hard fail
+in ``_node_status_from``.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import time
+from typing import Any, Dict, List, Optional
+
+from ..shell_utils import _which
+
+
+def _collect_rocm_smi_self_latency(*, timeout_sec: float) -> Dict[str, Any]:
+    """Time a single ``rocm-smi --version`` call against ``timeout_sec``.
+
+    A wedged amdgpu driver makes ``rocm-smi`` calls take 30-60 s before
+    failing outright -- usually minutes before the GPU itself stops
+    responding. Catching this in preflight gives operators a chance to
+    drain the node before a real training job starts hanging on it.
+    """
+    out: Dict[str, Any] = {
+        "ok": False,
+        "tool": None,
+        "latency_sec": None,
+        "timeout_sec": float(timeout_sec),
+    }
+    binpath = _which("rocm-smi")
+    if binpath is None:
+        out["error"] = "rocm-smi not found in PATH"
+        return out
+    out["tool"] = binpath
+    t0 = time.monotonic()
+    try:
+        cp = subprocess.run(
+            [binpath, "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+        out["latency_sec"] = round(time.monotonic() - t0, 3)
+        out["rc"] = cp.returncode
+        out["ok"] = cp.returncode == 0
+        if cp.returncode != 0:
+            out["error"] = (cp.stderr or cp.stdout or "").strip()[:200]
+    except subprocess.TimeoutExpired:
+        out["latency_sec"] = round(time.monotonic() - t0, 3)
+        out["timed_out"] = True
+        out["error"] = f"rocm-smi --version did not finish in {timeout_sec}s -- driver may be wedging"
+    except Exception as e:
+        out["error"] = str(e)
+    return out
+
+
+def _rocm_smi_ras_info_text(timeout_sec: float = 15.0) -> Dict[str, Any]:
+    """ECC counts via ``rocm-smi --showrasinfo`` (TEXT only).
+
+    The ``--json`` form returns "WARNING: No JSON data to report" so we
+    parse the text. Format is one block per GPU::
+
+        GPU[0]:         RAS INFO
+                Block       Status    Correctable Error  Uncorrectable Error
+                  UMC        ENABLED                  0                    0
+                 SDMA        ENABLED                  0                    0
+                  GFX        ENABLED                  0                    0
+                ...
+
+    Some blocks (ATHUB, PCIE_BIF, HDP, ...) report only Status. We sum
+    every numeric Correctable/Uncorrectable cell per GPU so a hardware
+    error in any block surfaces in ``ecc_uncorrectable_total``.
+    """
+    out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []}
+    if _which("rocm-smi") is None:
+        out["error"] = "rocm-smi not found in PATH"
+        return out
+    try:
+        cp = subprocess.run(
+            ["rocm-smi", "--showrasinfo"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+        if cp.returncode != 0:
+            out["error"] = (cp.stderr or cp.stdout or "").strip()[:200] or f"rc={cp.returncode}"
+            return out
+        out["per_gpu"] = _parse_rocm_smi_ras_info_text(cp.stdout)
+        out["tool"] = "rocm-smi --showrasinfo"
+        out["ok"] = True
+    except subprocess.TimeoutExpired:
+        out["error"] = "rocm-smi --showrasinfo timed out"
+    except Exception as e:
+        out["error"] = str(e)
+    return out
+
+
+def _parse_rocm_smi_ras_info_text(text: str) -> List[Dict[str, Any]]:
+    """Parse the per-GPU ``GPU[N]: RAS INFO`` blocks from rocm-smi text."""
+    import re
+
+    out: List[Dict[str, Any]] = []
+    cur_gpu: Optional[int] = None
+    cur_corr = 0
+    cur_uncorr = 0
+    in_block = False
+    gpu_hdr = re.compile(r"^GPU\[(\d+)\]:\s*RAS INFO", re.IGNORECASE)
+    for raw in text.splitlines():
+        line = raw.strip()
+        m = gpu_hdr.match(line)
+        if m:
+            # Flush previous GPU
+            if cur_gpu is not None:
+                out.append(
+                    {
+                        "gpu": cur_gpu,
+                        "ecc_correctable_total": cur_corr,
+                        "ecc_uncorrectable_total": cur_uncorr,
+                    }
+                )
+            cur_gpu = int(m.group(1))
+            cur_corr = 0
+            cur_uncorr = 0
+            in_block = True
+            continue
+        if not in_block or cur_gpu is None:
+            continue
+        # End-of-block separator
+        if line.startswith("__") or line.startswith("=="):
+            continue
+        # Per-block row: "BLOCK STATUS CORR UNCORR" -- last two are ints
+        # when present. Header row has the words "Correctable Error" so
+        # we filter out non-numeric rows naturally.
+        toks = line.split()
+        if len(toks) < 4:
+            continue
+        try:
+            corr = int(toks[-2])
+            uncorr = int(toks[-1])
+        except ValueError:
+            continue
+        cur_corr += corr
+        cur_uncorr += uncorr
+    if cur_gpu is not None:
+        out.append(
+            {
+                "gpu": cur_gpu,
+                "ecc_correctable_total": cur_corr,
+                "ecc_uncorrectable_total": cur_uncorr,
+            }
+        )
+    return out
+
+
+def _rocm_smi_topotype_json(timeout_sec: float = 15.0) -> Dict[str, Any]:
+    """XGMI link-type matrix via ``rocm-smi --showtopotype --json``.
+
+    Output shape is keyed by pair-string::
+
+        {"system": {
+            "(Topology) Link type between DRM devices 0 and 1": "XGMI",
+            "(Topology) Link type between DRM devices 0 and 2": "XGMI",
+            ...   # upper triangle only
+        }}
+
+    We parse the indices out of each key with a regex, build a symmetric
+    NxN matrix, and emit it in the same shape ``_collect_xgmi_topology``
+    already produces (so downstream consumers don't care which tool
+    populated it). ``non_xgmi_pairs`` contains the (i, j, link_type)
+    triples for any off-diagonal pair whose link_type is not XGMI.
+    """
+    import re
+
+    out: Dict[str, Any] = {"ok": False, "tool": None, "n_gpus": 0, "link_types": [], "non_xgmi_pairs": []}
+    if _which("rocm-smi") is None:
+        out["error"] = "rocm-smi not found in PATH"
+        return out
+    try:
+        cp = subprocess.run(
+            ["rocm-smi", "--showtopotype", "--json"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+        if cp.returncode != 0 or not cp.stdout.strip():
+            out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}"
+            return out
+        try:
+            doc = json.loads(cp.stdout)
+        except Exception as e:
+            out["error"] = f"json parse failed: {e}"
+            return out
+        sys_block = doc.get("system") if isinstance(doc, dict) else None
+        if not isinstance(sys_block, dict):
+            out["error"] = "no `system` key in rocm-smi --showtopotype output"
+            return out
+        pat = re.compile(r"DRM\s+devices?\s+(\d+)\s+and\s+(\d+)", re.IGNORECASE)
+        pairs: Dict[tuple, str] = {}
+        max_idx = -1
+        for k, v in sys_block.items():
+            m = pat.search(str(k))
+            if not m:
+                continue
+            i = int(m.group(1))
+            j = int(m.group(2))
+            pairs[(i, j)] = str(v)
+            max_idx = max(max_idx, i, j)
+        if max_idx < 0:
+            out["error"] = "no parseable DRM-device pair keys"
+            return out
+        n = max_idx + 1
+        mat = [["?" for _ in range(n)] for _ in range(n)]
+        for (i, j), t in pairs.items():
+            mat[i][j] = t
+            mat[j][i] = t
+        for i in range(n):
+            mat[i][i] = "SELF"
+        non_xgmi: List[Any] = []
+        for i in range(n):
+            for j in range(i + 1, n):
+                t = mat[i][j]
+                if t and t.upper() != "XGMI":
+                    non_xgmi.append((i, j, t))
+        out["n_gpus"] = n
+        out["link_types"] = mat
+        out["non_xgmi_pairs"] = non_xgmi
+        out["tool"] = "rocm-smi --showtopotype --json"
+        out["ok"] = True
+    except subprocess.TimeoutExpired:
+        out["error"] = "rocm-smi --showtopotype timed out"
+    except Exception as e:
+        out["error"] = str(e)
+    return out
+
+
+def _rocm_smi_processes(
+    annotate: Any,
+    timeout_sec: float = 15.0,
+) -> Dict[str, Any]:
+    """Foreign-process enumeration via ``rocm-smi --showpids --json``.
+
+    Output shape (verified against rocm-smi on a busy MI300X)::
+
+        {"system": {
+            "PID2683309": "python3.11, 1, 24556904448, 0, 0",
+            "PID29324":   "gpuagent, 0, 0, 0, 0"
+        }}
+
+    Comma-separated fields are: ``name, num_gpus_used, vram_bytes,
+    sdma_bytes, cu_occupancy``. We extract ``name`` (field 0) and
+    ``vram_bytes`` (field 2) -- enough to surface leaked training PIDs
+    holding gigabytes of HBM. ``--showpidgpus --json`` returns
+    "No JSON data to report", so per-GPU mapping is not available --
+    all PIDs go into the gpu=-1 bucket (same convention as lsof).
+    """
+    out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []}
+    if _which("rocm-smi") is None:
+        out["error"] = "rocm-smi not found in PATH"
+        return out
+    try:
+        cp = subprocess.run(
+            ["rocm-smi", "--showpids", "--json"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+        if cp.returncode != 0 or not cp.stdout.strip():
+            out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}"
+            return out
+        try:
+            doc = json.loads(cp.stdout)
+        except Exception as e:
+            out["error"] = f"json parse failed: {e}"
+            return out
+        sys_block = doc.get("system") if isinstance(doc, dict) else None
+        if not isinstance(sys_block, dict):
+            # Empty or unexpected -- treat as "no PIDs"
+            out["tool"] = "rocm-smi --showpids --json"
+            out["ok"] = True
+            return out
+        procs: List[Dict[str, Any]] = []
+        for k, v in sys_block.items():
+            try:
+                pid = int(str(k).lstrip("PID").lstrip("pid"))
+            except ValueError:
+                continue
+            name = ""
+            hbm: Optional[int] = None
+            if isinstance(v, str):
+                parts = [s.strip() for s in v.split(",")]
+                if parts:
+                    name = parts[0]
+                # field 2 = VRAM bytes (rocm-smi --showpids documented format)
+                if len(parts) > 2:
+                    try:
+                        hbm = int(parts[2])
+                    except ValueError:
+                        hbm = None
+            elif isinstance(v, dict):
+                name = str(v.get("name") or v.get("process_name") or "")
+                vram = v.get("vram") or v.get("vram_bytes")
+                if isinstance(vram, (int, float)):
+                    hbm = int(vram)
+            elif isinstance(v, list) and v:
+                name = str(v[0])
+            procs.append(annotate(pid, name, hbm))
+        if procs:
+            out["per_gpu"] = [{"gpu": -1, "processes": procs}]
+        out["tool"] = "rocm-smi --showpids --json"
+        out["ok"] = True
+    except subprocess.TimeoutExpired:
+        out["error"] = "rocm-smi --showpids timed out"
+    except Exception as e:
+        out["error"] = str(e)
+    return out
+
+
+def _rocm_smi_use_json(timeout_sec: float = 15.0) -> Dict[str, Any]:
+    """GPU compute-activity % via ``rocm-smi --showuse --json``.
+
+    Output shape::
+
+        {"card0": {"GPU use (%)": "0", "GFX Activity": "465307149"}, ...}
+
+    "GPU use (%)" is the percentage we want for ``gfx_activity_pct``
+    (matching the amd-smi field name). "GFX Activity" is a cumulative
+    cycle counter, not a percentage -- ignored.
+    """
+    out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []}
+    if _which("rocm-smi") is None:
+        out["error"] = "rocm-smi not found in PATH"
+        return out
+    try:
+        cp = subprocess.run(
+            ["rocm-smi", "--showuse", "--json"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+        if cp.returncode != 0 or not cp.stdout.strip():
+            out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}"
+            return out
+        try:
+            doc = json.loads(cp.stdout)
+        except Exception as e:
+            out["error"] = f"json parse failed: {e}"
+            return out
+        if not isinstance(doc, dict):
+            out["error"] = "unexpected top-level shape"
+            return out
+        per_gpu: List[Dict[str, Any]] = []
+        for k, v in doc.items():
+            if not isinstance(v, dict):
+                continue
+            if not str(k).startswith("card"):
+                continue
+            try:
+                gpu = int(str(k)[len("card") :])
+            except ValueError:
+                continue
+            pct_raw = v.get("GPU use (%)") or v.get("GPU use")
+            try:
+                pct = float(str(pct_raw).strip())
+            except (TypeError, ValueError):
+                continue
+            per_gpu.append({"gpu": gpu, "gfx_activity_pct": pct})
+        out["per_gpu"] = per_gpu
+        out["tool"] = "rocm-smi --showuse --json"
+        out["ok"] = True
+    except subprocess.TimeoutExpired:
+        out["error"] = "rocm-smi --showuse timed out"
+    except Exception as e:
+        out["error"] = str(e)
+    return out
diff --git a/primus/tools/preflight/node_smoke/collectors/tooling.py b/primus/tools/preflight/node_smoke/collectors/tooling.py
new file mode 100644
index 000000000..811ee3f50
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/tooling.py
@@ -0,0 +1,100 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- tooling availability inventory.
+
+Several Tier 1 checks (ECC via amd-smi metric, XGMI via amd-smi topology,
+foreign-process enumeration via amd-smi process / lsof, wedged-driver
+canary via rocm-smi --version) are best-effort: each collector returns
+ok=False and the FAIL rules in _node_status_from then iterate over empty
+data, silently no-op'ing. That's fine on a node where the tools are
+legitimately absent (containers stripped down for size), but DANGEROUS
+in a production cluster -- a node with ECC errors, broken XGMI, leaked
+ranks, or a stale ROCm install can PASS smoke just because amd-smi is
+missing.
+
+This collector runs ONCE per node, captures which tools were resolvable
+in PATH, and feeds three downstream consumers:
+
+  1. A loud `_warn` at run-time so missing tools are visible in the
+     srun log right next to "start node-smoke".
+  2. An always-on "Tooling availability" section in the aggregator
+     report -- a per-node table that's loud even when nothing else is.
+  3. The optional `--require-tools` flag, which promotes a missing
+     required tool to a hard FAIL via _node_status_from.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from ..shell_utils import _which
+
+_TRACKED_TOOLS = ("amd-smi", "rocm-smi", "lsof")
+
+
+def _collect_tooling_inventory() -> Dict[str, Any]:
+    """Resolve each tracked tool in PATH and compute per-check coverage.
+
+    Output shape::
+
+        {
+            "ok": True,                          # always True; collector itself can't fail
+            "tools": {
+                "amd-smi":  {"present": True,  "path": "/usr/bin/amd-smi"},
+                "rocm-smi": {"present": True,  "path": "/usr/bin/rocm-smi"},
+                "lsof":     {"present": True,  "path": "/usr/bin/lsof"},
+            },
+            "missing": ["rocm-smi"],             # convenience list of absent tools
+            "coverage": {
+                "ECC":                                  True,
+                "XGMI":                                 True,
+                "foreign-process":                      True,
+                "GPU activity warn":                    True,
+                "amd-smi/torch GPU-count cross-check":  True,
+                "wedged-driver canary":                 True,
+            },
+            "uncovered": [],                     # checks with no working tool
+        }
+
+    Coverage rules (each check is "covered" if ANY of the listed tools
+    is present):
+
+      * ECC                                 -> amd-smi  OR  rocm-smi (--showrasinfo)
+      * XGMI link matrix                    -> amd-smi  OR  rocm-smi (--showtopotype)
+      * foreign-process enumeration         -> amd-smi  OR  rocm-smi  OR  lsof
+      * GPU activity warn (gfx_activity_pct)-> amd-smi  OR  rocm-smi (--showuse)
+      * amd-smi/torch GPU-count cross-check -> amd-smi only (rocm-smi enumerates
+                                               by `cardN` not by torch's index)
+      * wedged-driver canary (rocm-smi --version) -> rocm-smi only
+    """
+    tools: Dict[str, Dict[str, Any]] = {}
+    missing: List[str] = []
+    for name in _TRACKED_TOOLS:
+        path = _which(name)
+        tools[name] = {"present": path is not None, "path": path}
+        if path is None:
+            missing.append(name)
+
+    has_amd = tools["amd-smi"]["present"]
+    has_rocm = tools["rocm-smi"]["present"]
+    has_lsof = tools["lsof"]["present"]
+    coverage = {
+        "ECC": has_amd or has_rocm,
+        "XGMI": has_amd or has_rocm,
+        "foreign-process": has_amd or has_rocm or has_lsof,
+        "GPU activity warn": has_amd or has_rocm,
+        "amd-smi/torch GPU-count cross-check": has_amd,
+        "wedged-driver canary": has_rocm,
+    }
+    uncovered = [name for name, ok in coverage.items() if not ok]
+    return {
+        "ok": True,
+        "tools": tools,
+        "missing": missing,
+        "coverage": coverage,
+        "uncovered": uncovered,
+    }
diff --git a/primus/tools/preflight/node_smoke/collectors/xgmi.py b/primus/tools/preflight/node_smoke/collectors/xgmi.py
new file mode 100644
index 000000000..0d6d55ea8
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/collectors/xgmi.py
@@ -0,0 +1,201 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 1 -- D-2: XGMI topology matrix via ``amd-smi topology`` (text parser),
+with rocm-smi --showtopotype as a cross-tool fallback."""
+
+from __future__ import annotations
+
+import subprocess
+from typing import Any, Dict, List
+
+from ..shell_utils import _which
+from .rocm_smi import _rocm_smi_topotype_json
+
+
+def _collect_xgmi_topology() -> Dict[str, Any]:
+    """Try amd-smi topology first; fall back to rocm-smi --showtopotype.
+
+    Both tools produce slightly different per-row labels (amd-smi uses
+    PCIe BDFs, rocm-smi uses DRM device indices), but the link_types
+    matrix and non_xgmi_pairs computation is identical, so downstream
+    consumers in _node_status_from / the aggregator don't need to know
+    which tool produced the data.
+    """
+    out = _collect_xgmi_topology_amd_smi()
+    if out.get("ok"):
+        return out
+    rocm = _rocm_smi_topotype_json()
+    if rocm.get("ok"):
+        return {
+            "ok": True,
+            "tool": rocm.get("tool") or "rocm-smi --showtopotype --json",
+            "bdfs": [],  # rocm-smi reports DRM indices, not PCIe BDFs
+            "matrix": rocm.get("link_types") or [],
+            "n_gpus": rocm.get("n_gpus") or 0,
+            "non_xgmi_pairs": rocm.get("non_xgmi_pairs") or [],
+            "amd_smi_error": out.get("error"),
+        }
+    # Both failed -- preserve amd-smi's error for the operator and
+    # surface rocm-smi's separately so they can debug both paths.
+    out["rocm_smi_error"] = rocm.get("error")
+    return out
+
+
+def _collect_xgmi_topology_amd_smi() -> Dict[str, Any]:
+    """Parse ``amd-smi topology`` and return a square link-type matrix.
+
+    ``amd-smi topology`` emits several BDF-labelled sub-tables (ACCESS,
+    WEIGHT, HOPS, LINK TYPE, NUMA BW, ...). We pick the ``LINK TYPE TABLE``
+    sub-section, which contains values like ``SELF`` (diagonal) and
+    ``XGMI`` / ``PCIE`` / ``PIX`` / ``SOC`` etc. Off-diagonal cells that
+    aren't ``XGMI`` are recorded as ``non_xgmi_pairs`` and treated as a
+    hard fail by ``_node_status_from`` -- the moment a single GPU pair
+    falls back to PCIe inside a node, intra-node collectives lose 5-10x
+    of the bandwidth NCCL/RCCL expects.
+
+    The on-disk shape:
+
+        {
+            "ok": bool,
+            "tool": "amd-smi topology" | None,
+            "bdfs": ["0000:05:00.0", ...],
+            "matrix": [["SELF","XGMI",...], ["XGMI","SELF",...], ...],
+            "n_gpus": int,
+            "non_xgmi_pairs": [(i, j, link_type), ...],
+            "error": "..."
+        }
+    """
+    out: Dict[str, Any] = {
+        "ok": False,
+        "tool": None,
+        "bdfs": [],
+        "matrix": [],
+        "non_xgmi_pairs": [],
+    }
+    if _which("amd-smi") is None:
+        out["error"] = "amd-smi not found in PATH"
+        return out
+    try:
+        cp = subprocess.run(
+            ["amd-smi", "topology"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=15,
+            check=False,
+        )
+    except subprocess.TimeoutExpired:
+        out["error"] = "amd-smi topology timed out"
+        return out
+    except Exception as e:
+        out["error"] = str(e)
+        return out
+    if cp.returncode != 0:
+        out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}"
+        return out
+
+    text = cp.stdout
+
+    # Parse: find the `LINK TYPE TABLE:` section, then the BDF header row,
+    # then the per-BDF data rows. Stop at the next section header (any all-
+    # caps label ending in `TABLE:`) or end of text.
+    import re
+
+    bdf_re = re.compile(r"\b([0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.\d)\b")
+    section_header_re = re.compile(r"^\s*[A-Z][A-Z0-9 -]+TABLE:\s*$")
+
+    lines = text.splitlines()
+    try:
+        idx = next(i for i, l in enumerate(lines) if l.strip().upper() == "LINK TYPE TABLE:")
+    except StopIteration:
+        out["error"] = "no `LINK TYPE TABLE:` section in `amd-smi topology` output"
+        out["raw"] = text[:4000]
+        return out
+
+    # Header row is the next non-empty line after the label, and contains
+    # no leading BDF -- only column BDFs.
+    header_bdfs: List[str] = []
+    data_start = None
+    for j in range(idx + 1, len(lines)):
+        l = lines[j].rstrip()
+        if not l.strip():
+            continue
+        if section_header_re.match(l):
+            break
+        toks = bdf_re.findall(l)
+        if not toks:
+            continue
+        # The header line has only column BDFs (no leading row label), and
+        # the first non-whitespace char position lines up with the columns.
+        # Heuristic: header has BDFs but no other tokens that look like
+        # link-type values (XGMI/PCIE/SELF/...). Data rows always have
+        # exactly one leading BDF followed by N value tokens.
+        non_bdf_toks = [t for t in l.split() if not bdf_re.fullmatch(t)]
+        if not non_bdf_toks:
+            header_bdfs = toks
+            data_start = j + 1
+            break
+
+    if not header_bdfs or data_start is None:
+        out["error"] = "could not find header row inside LINK TYPE TABLE"
+        out["raw"] = text[:4000]
+        return out
+
+    n = len(header_bdfs)
+    bdf_to_idx = {b: i for i, b in enumerate(header_bdfs)}
+    matrix: List[List[str]] = [[""] * n for _ in range(n)]
+    seen_rows = 0
+    for j in range(data_start, len(lines)):
+        l = lines[j].rstrip()
+        if not l.strip():
+            continue
+        if section_header_re.match(l):
+            break
+        toks = l.split()
+        # First token must be a BDF, the remaining N tokens are the row.
+        if not bdf_re.fullmatch(toks[0]):
+            continue
+        row_bdf = toks[0]
+        cells = toks[1:]
+        if row_bdf not in bdf_to_idx:
+            continue
+        row_idx = bdf_to_idx[row_bdf]
+        for k, cell in enumerate(cells[:n]):
+            matrix[row_idx][k] = cell
+        seen_rows += 1
+
+    if seen_rows == 0:
+        out["error"] = "no BDF-labelled rows found inside LINK TYPE TABLE"
+        out["raw"] = text[:4000]
+        return out
+
+    healthy_diag = {"SELF", "X", "-", "0"}
+    healthy_link = {"XGMI"}
+    non_xgmi: List[Any] = []
+    for i, row in enumerate(matrix):
+        for j_idx, cell in enumerate(row):
+            cu = cell.strip().upper()
+            if i == j_idx:
+                # Diagonal: must be SELF (or empty if the row was missing).
+                if cu and cu not in healthy_diag and cu not in healthy_link:
+                    non_xgmi.append((i, j_idx, cell))
+                continue
+            if not cu:
+                # Missing cell -> can't certify XGMI -> flag.
+                non_xgmi.append((i, j_idx, "<missing>"))
+                continue
+            if cu in healthy_link or cu in healthy_diag:
+                continue
+            non_xgmi.append((i, j_idx, cell))
+
+    out["ok"] = True
+    out["tool"] = "amd-smi topology"
+    out["bdfs"] = header_bdfs
+    out["matrix"] = matrix
+    out["n_gpus"] = n
+    out["non_xgmi_pairs"] = non_xgmi
+    return out
diff --git a/primus/tools/preflight/node_smoke/logging_utils.py b/primus/tools/preflight/node_smoke/logging_utils.py
new file mode 100644
index 000000000..7c853e165
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/logging_utils.py
@@ -0,0 +1,53 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Hostname normalization + log/warn helpers used everywhere in node-smoke.
+
+The log prefix ``[HH:MM:SS][node-smoke][<short-host>]`` is part of the
+operator-facing contract -- existing log-scraping tooling assumes it,
+so it is preserved verbatim across the refactor.
+"""
+
+from __future__ import annotations
+
+import socket
+import sys
+import time
+
+
+def _ts() -> str:
+    return time.strftime("%H:%M:%S")
+
+
+def _short_name(h: str) -> str:
+    """Return the leading short-hostname segment.
+
+    SLURM tools (`scontrol show hostnames`, `srun --nodelist=`,
+    `srun --exclude=`) all operate on short hostnames, so we normalize
+    everywhere so the produced ``passing_nodes.txt`` / ``failing_nodes.txt``
+    can be piped straight into them. ``socket.gethostname()`` returns the
+    FQDN on some clusters, hence this helper.
+    """
+    if not h:
+        return h
+    return h.split(".", 1)[0]
+
+
+def _this_host_short() -> str:
+    """This node's short hostname (first segment of socket.gethostname())."""
+    return _short_name(socket.gethostname())
+
+
+def _log(msg: str) -> None:
+    print(f"[{_ts()}][node-smoke][{_this_host_short()}] {msg}", flush=True)
+
+
+def _warn(msg: str) -> None:
+    print(
+        f"[{_ts()}][node-smoke][{_this_host_short()}] WARN: {msg}",
+        file=sys.stderr,
+        flush=True,
+    )
diff --git a/primus/tools/preflight/node_smoke/orchestrator.py b/primus/tools/preflight/node_smoke/orchestrator.py
new file mode 100644
index 000000000..8c4123af6
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/orchestrator.py
@@ -0,0 +1,288 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Per-node orchestration helpers.
+
+* :func:`_spawn_per_gpu`    -- launch the ``_per_gpu`` subcommand for one
+                               GPU index with a hard timeout.
+* :func:`_node_status_from` -- compute the list of fail reasons for the
+                               whole node from collected per-GPU + tier1
+                               + tier2 state.
+* :func:`_clean_dump_path`  -- wipe stale artifacts from a previous run on
+                               rank 0 before any rank writes its current
+                               JSON.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import time
+from typing import Any, Dict, List, Optional
+
+from .types import GPUResult
+
+
+def _spawn_per_gpu(
+    gpu: int,
+    *,
+    timeout_sec: int,
+    tier2_perf: bool,
+    gemm_tflops_min: float,
+    hbm_gbs_min: float,
+    hbm_busy_threshold_gib: float,
+) -> GPUResult:
+    """Spawn ``python -m primus.tools.preflight.node_smoke _per_gpu <gpu> ...``
+    with a hard timeout so a stuck driver call cannot wedge the parent."""
+    cmd = [
+        sys.executable,
+        "-m",
+        "primus.tools.preflight.node_smoke",
+        "_per_gpu",
+        str(gpu),
+        "--gemm-tflops-min",
+        str(gemm_tflops_min),
+        "--hbm-gbs-min",
+        str(hbm_gbs_min),
+        "--hbm-busy-threshold-gib",
+        str(hbm_busy_threshold_gib),
+    ]
+    if tier2_perf:
+        cmd.append("--tier2-perf")
+
+    t0 = time.time()
+    try:
+        cp = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=timeout_sec,
+            check=False,
+        )
+    except subprocess.TimeoutExpired:
+        return GPUResult(
+            gpu=gpu,
+            status="TIMEOUT",
+            reason=f"per-gpu subprocess hit hard timeout {timeout_sec}s",
+            duration_sec=round(time.time() - t0, 3),
+        )
+    except Exception as e:
+        return GPUResult(
+            gpu=gpu,
+            status="FAIL",
+            reason=f"failed to spawn per-gpu subprocess: {e}",
+            duration_sec=round(time.time() - t0, 3),
+        )
+
+    # The subprocess prints exactly one JSON line on stdout for the result.
+    raw = (cp.stdout or "").strip().splitlines()
+    if not raw:
+        return GPUResult(
+            gpu=gpu,
+            status="FAIL",
+            reason=(
+                f"per-gpu subprocess produced no JSON (rc={cp.returncode}, "
+                f"stderr={cp.stderr.strip()[:200]})"
+            ),
+            duration_sec=round(time.time() - t0, 3),
+        )
+    try:
+        data = json.loads(raw[-1])
+    except Exception as e:
+        return GPUResult(
+            gpu=gpu,
+            status="FAIL",
+            reason=f"per-gpu JSON parse failed: {e}; raw={raw[-1][:200]}",
+            duration_sec=round(time.time() - t0, 3),
+        )
+
+    return GPUResult(
+        gpu=int(data.get("gpu", gpu)),
+        status=str(data.get("status", "FAIL")),
+        reason=str(data.get("reason", "")),
+        duration_sec=float(data.get("duration_sec", time.time() - t0)),
+        details=dict(data.get("details", {})),
+    )
+
+
+def _node_status_from(
+    per_gpu: List[GPUResult],
+    tier1_extra: Dict[str, Any],
+    tier2_extra: Dict[str, Any],
+    *,
+    allow_foreign_procs: bool = False,
+    required_tools: Optional[List[str]] = None,
+) -> List[str]:
+    """Compute a list of ``fail_reasons`` for the node from collected results.
+
+    Empty list -> node PASS. Any non-empty result -> node FAIL.
+
+    ``allow_foreign_procs`` downgrades the foreign-process FAIL to a
+    silent inclusion in the JSON (still surfaced by the aggregator's
+    "Busy GPUs" section, just not a hard fail).
+
+    ``required_tools`` is the operator-supplied list of CLI tools that
+    MUST be present on this node (e.g. ``["amd-smi", "rocm-smi"]``).
+    Anything in this list that is not in the tooling-inventory becomes
+    a hard FAIL. Empty / None means "warn only" (the default).
+    """
+    reasons: List[str] = []
+
+    # Self-contained GPU visibility guard. Decoupled from any other
+    # collector so a wrapped/downgraded "No GPUs detected" finding can
+    # never silently turn a CPU-only or stale-GPU node into a PASS.
+    vis = tier1_extra.get("gpu_visibility") or {}
+    for r in vis.get("fail_reasons", []) or []:
+        reasons.append(f"gpu_visibility: {r}")
+
+    for r in per_gpu:
+        if r.status != "PASS":
+            reasons.append(f"gpu{r.gpu}: {r.status}: {r.reason}")
+
+    for section_name in ("gpu_info", "host_info", "network_info"):
+        for f in tier1_extra.get(section_name, []):
+            if f.get("level") == "fail":
+                reasons.append(f"{section_name}: {f.get('message', '<no message>')}")
+
+    dmesg = tier1_extra.get("dmesg") or {}
+    if dmesg.get("matches"):
+        first = dmesg["matches"][0]
+        reasons.append(f"dmesg ({len(dmesg['matches'])} match(es), e.g.): {first[:200]}")
+
+    # B. NIC / RDMA roll-call -- every issue here is a hard fail because each
+    # one (port DOWN, missing RoCE v2 GID, wrong NIC count) silently breaks
+    # inter-node training the moment the first global collective runs.
+    for issue in (tier1_extra.get("nics") or {}).get("issues", []) or []:
+        reasons.append(f"nic: {issue}")
+
+    # C. Host limits -- only the entries the collector flagged as hard
+    # (ulimit -l below threshold, /dev/shm too small) become node FAIL.
+    for issue in (tier1_extra.get("host_limits") or {}).get("fail_reasons", []) or []:
+        reasons.append(f"host_limits: {issue}")
+
+    rccl = tier2_extra.get("rccl") or {}
+    if rccl and rccl.get("status") not in (None, "PASS"):
+        reasons.append(f"rccl: {rccl.get('status')}: {rccl.get('error', '')}")
+
+    # D-1 heavy: any per-GPU uncorrectable ECC count is a hard fail. The
+    # amd-smi schema isn't stable across releases so we trust only the
+    # values our flattener was able to coerce to int. Throttle reasons stay
+    # informational (the schema is too vendor-specific to fail on).
+    amd = tier1_extra.get("gpu_low_level") or {}
+    for rec in amd.get("per_gpu", []) or []:
+        ue = rec.get("ecc_uncorrectable_total")
+        if isinstance(ue, int) and ue > 0:
+            reasons.append(f"gpu{rec.get('gpu', '?')}: ECC uncorrectable count = {ue}")
+
+    # D-2: any non-XGMI GPU pair is a hard fail -- intra-node collectives
+    # silently fall back to PCIe and lose 5-10x bandwidth.
+    xg = tier1_extra.get("xgmi") or {}
+    bad = xg.get("non_xgmi_pairs") or []
+    if bad:
+        sample = ", ".join(f"({i},{j})={t}" for i, j, t in bad[:3])
+        reasons.append(f"xgmi: {len(bad)} non-XGMI GPU pair(s) detected, e.g. {sample}")
+
+    # F-partial: rocm-smi --version that timed out -> driver is wedging.
+    # Slow-but-completed calls are surfaced by the aggregator only.
+    tool = tier1_extra.get("tooling") or {}
+    if tool.get("timed_out"):
+        reasons.append(
+            f"tooling: rocm-smi --version did not return within "
+            f"{tool.get('timeout_sec', '?')}s -- driver may be wedging"
+        )
+
+    # Tooling availability. Missing CLI tools (amd-smi, rocm-smi, lsof)
+    # cause silent skips of several Tier 1 checks. Operators in strict
+    # environments can pass --require-tools to convert "missing" into a
+    # node FAIL so the node is pulled from rotation until the toolchain
+    # is fixed. Default (empty list) is warn-only (the WARN already fires
+    # in _cmd_run before the per-GPU subprocesses run).
+    if required_tools:
+        inv = (tier1_extra.get("tooling_inventory") or {}).get("tools") or {}
+        missing_required = [t for t in required_tools if not (inv.get(t) or {}).get("present")]
+        if missing_required:
+            reasons.append(
+                f"tooling_inventory: required tool(s) NOT in PATH: "
+                f"{', '.join(missing_required)} -- silent-skips several "
+                f"Tier 1 checks (ECC, XGMI, foreign-process, wedged-driver). "
+                f"Pass --require-tools '' to disable this fail."
+            )
+
+    # G: foreign processes holding the GPU. Hard-fail by default because
+    # this is the single most common cause of training failing to launch
+    # on an otherwise-healthy node (leaked Python ranks from a previous
+    # job, a profiler that never detached, a foreign tenant on a shared
+    # partition). The operator can downgrade with --allow-foreign-procs
+    # if their workflow legitimately co-tenants the GPU.
+    gp = tier1_extra.get("gpu_processes") or {}
+    if not allow_foreign_procs and gp.get("foreign_count", 0) > 0:
+        examples: List[str] = []
+        for g in gp.get("per_gpu") or []:
+            for p in g.get("processes") or []:
+                if not p.get("is_foreign"):
+                    continue
+                hbm = p.get("hbm_bytes")
+                hbm_s = f" hbm={round(hbm / (1 << 30), 2)}GiB" if isinstance(hbm, int) and hbm > 0 else ""
+                examples.append(f"gpu{g.get('gpu')}: pid={p.get('pid')} " f"name={p.get('name')!r}{hbm_s}")
+                if len(examples) >= 3:
+                    break
+            if len(examples) >= 3:
+                break
+        reasons.append(
+            f"gpu_processes: {gp.get('foreign_count', 0)} foreign process(es) "
+            f"holding GPU(s) (e.g. " + "; ".join(examples) + ") -- likely "
+            f"leaked rank(s) from a previous job. Clean up with "
+            f"`pkill -9 -f train.py` (or similar) or pass --allow-foreign-procs."
+        )
+
+    return reasons
+
+
+def _clean_dump_path(dump_path: str) -> List[str]:
+    """Wipe stale per-node JSONs and aggregator outputs from a previous run.
+
+    Without this, a re-run on a different (smaller) nodelist would leave
+    JSONs from removed nodes in ``<dump>/smoke/`` and the aggregator would
+    happily count them as PASS, contaminating the report. We clean only
+    artifacts that ``run`` and ``aggregate`` produce (per-node JSONs and
+    the four top-level outputs); anything else under ``--dump-path`` is
+    left untouched.
+
+    Race safety: this is called only on rank 0 in ``_cmd_run``, BEFORE any
+    rank can have written its current-run JSON (each rank's per-GPU
+    subprocess loop + collector phase takes seconds; rank 0's cleanup
+    finishes in milliseconds). Other ranks never delete anything.
+
+    Returns the list of files actually removed (for logging).
+    """
+    removed: List[str] = []
+    smoke_dir = os.path.join(dump_path, "smoke")
+    if os.path.isdir(smoke_dir):
+        for name in os.listdir(smoke_dir):
+            if name.endswith(".json"):
+                p = os.path.join(smoke_dir, name)
+                try:
+                    os.remove(p)
+                    removed.append(p)
+                except OSError:
+                    pass
+    for name in (
+        "smoke_report.md",
+        "passing_nodes.txt",
+        "failing_nodes.txt",
+        "expected_nodes.txt",
+    ):
+        p = os.path.join(dump_path, name)
+        if os.path.isfile(p):
+            try:
+                os.remove(p)
+                removed.append(p)
+            except OSError:
+                pass
+    return removed
diff --git a/primus/tools/preflight/node_smoke/per_gpu.py b/primus/tools/preflight/node_smoke/per_gpu.py
new file mode 100644
index 000000000..64a29173e
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/per_gpu.py
@@ -0,0 +1,327 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Per-GPU subprocess body (Tier 1 + optional Tier 2 perf).
+
+This module is the body of the ``_per_gpu`` subcommand: it runs every
+GPU-local check on a single device index and returns a dict result.
+The orchestrator spawns one subprocess per GPU with a hard timeout so a
+stuck driver call cannot wedge the parent.
+
+Kept intact (no internal split) -- every code path returns a complete
+verdict dict, and inlining each stage is currently easier to follow
+than splitting the function up.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict
+
+from .shell_utils import _read_text, _resolve_gpu_bdf
+
+
+def _per_gpu_body(
+    gpu: int,
+    *,
+    tier2_perf: bool,
+    gemm_tflops_min: float,
+    hbm_gbs_min: float,
+    hbm_busy_threshold_bytes: int = 2 * (1 << 30),
+) -> Dict[str, Any]:
+    """Run all per-GPU tests for a single GPU and return a dict result.
+
+    Tier 1 (always): set_device, **pre-touch HBM-busy check** (FAIL if more
+    than ``hbm_busy_threshold_bytes`` already in use before our test
+    allocates anything), allocate 256 MB, tiny GEMM 2048x2048 bf16 with
+    finite-value check.
+
+    Tier 2 (when ``tier2_perf`` is True): GEMM 8192x8192 bf16 TFLOPS
+    measurement against ``gemm_tflops_min``, and HBM device-to-device
+    copy bandwidth against ``hbm_gbs_min``. Each metric below threshold
+    yields FAIL.
+
+    Pre-touch HBM check: ``torch.cuda.mem_get_info`` is called BEFORE we
+    allocate anything on this GPU, so the "used" reading reflects only
+    foreign / leaked allocations. The post-test reading is also captured
+    (under ``low_level.hbm_free_bytes``) for completeness, but the FAIL
+    rule uses only the pre-touch number to avoid being polluted by our
+    own caching-allocator footprint.
+    """
+    t0 = time.time()
+    details: Dict[str, Any] = {}
+
+    try:
+        import torch  # type: ignore
+    except Exception as e:
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": f"torch import failed: {e}",
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+
+    if not torch.cuda.is_available():
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": "torch.cuda.is_available() is False",
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+    if gpu >= torch.cuda.device_count():
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": (f"gpu index {gpu} >= visible device_count {torch.cuda.device_count()}"),
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+
+    # --- set_device ---
+    try:
+        torch.cuda.set_device(gpu)
+    except Exception as e:
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": f"set_device({gpu}) raised: {e}",
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+
+    # --- pre-touch HBM-busy check (BEFORE we allocate anything) ---
+    # Captured here, NOT in the low_level block at the end, because by
+    # then PyTorch's caching allocator has already taken pages we won't
+    # truly release on empty_cache(). The pre-touch reading is the only
+    # honest answer to "is someone else holding this GPU?".
+    try:
+        free_b, total_b = torch.cuda.mem_get_info(gpu)
+        used_b = max(0, int(total_b) - int(free_b))
+        details["hbm_pre_touch_total_bytes"] = int(total_b)
+        details["hbm_pre_touch_free_bytes"] = int(free_b)
+        details["hbm_pre_touch_used_bytes"] = used_b
+        details["hbm_pre_touch_used_gib"] = round(used_b / (1 << 30), 3)
+        if used_b >= hbm_busy_threshold_bytes:
+            return {
+                "gpu": gpu,
+                "status": "FAIL",
+                "reason": (
+                    f"pre-touch HBM busy: {round(used_b / (1 << 30), 2)} GiB "
+                    f"already in use (threshold "
+                    f"{round(hbm_busy_threshold_bytes / (1 << 30), 2)} GiB) "
+                    f"-- likely leaked process from a previous job; "
+                    f"see node-level gpu_processes section to identify the PID"
+                ),
+                "duration_sec": round(time.time() - t0, 3),
+                "details": details,
+            }
+    except Exception as e:
+        details["hbm_pre_touch_error"] = f"mem_get_info failed: {e}"
+
+    # --- 256 MB tensor alloc + simple write + sync ---
+    try:
+        nbytes_alloc = 256 * 1024 * 1024
+        n_elem = nbytes_alloc // 2  # bf16
+        x = torch.empty(n_elem, dtype=torch.bfloat16, device=f"cuda:{gpu}")
+        x.fill_(1.0)
+        torch.cuda.synchronize()
+        del x
+        torch.cuda.empty_cache()
+    except Exception as e:
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": f"256MB bf16 alloc/fill/sync failed: {e}",
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+
+    # --- tiny GEMM 2048x2048 bf16, finite-value check ---
+    try:
+        m = n = k = 2048
+        a = torch.randn((m, k), dtype=torch.bfloat16, device=f"cuda:{gpu}")
+        b = torch.randn((k, n), dtype=torch.bfloat16, device=f"cuda:{gpu}")
+        c = torch.matmul(a, b)
+        torch.cuda.synchronize()
+        if not torch.isfinite(c).all().item():
+            return {
+                "gpu": gpu,
+                "status": "FAIL",
+                "reason": "tiny GEMM produced non-finite values (possible HW corruption)",
+                "duration_sec": round(time.time() - t0, 3),
+                "details": details,
+            }
+        del a, b, c
+        torch.cuda.empty_cache()
+    except Exception as e:
+        return {
+            "gpu": gpu,
+            "status": "FAIL",
+            "reason": f"tiny GEMM 2048x2048 failed: {e}",
+            "duration_sec": round(time.time() - t0, 3),
+            "details": details,
+        }
+
+    # --- D-1 light: PCIe link + HBM (sysfs + torch only, fast & no shell-out) ---
+    # Captured into details.low_level so the aggregator can flag drift across
+    # the cluster (e.g. a single GPU running at Gen3 x8 because the slot
+    # needs reseating, or a GPU that exposes only half of its HBM).
+    # Each sub-capture is independent: a missing/unparseable PCIe BDF must
+    # not cost us the HBM size, and vice-versa.
+    low: Dict[str, Any] = {}
+    props = None
+    try:
+        props = torch.cuda.get_device_properties(gpu)
+    except Exception as e:
+        low["error"] = f"get_device_properties failed: {e}"
+
+    if props is not None:
+        # PCIe link details (sysfs)
+        try:
+            bdf = _resolve_gpu_bdf(props)
+            if bdf:
+                low["pci_bdf"] = bdf
+                sysdir = f"/sys/bus/pci/devices/{bdf}"
+                speed = _read_text(f"{sysdir}/current_link_speed")
+                width = _read_text(f"{sysdir}/current_link_width")
+                low["pcie_link_speed_raw"] = speed or None
+                low["pcie_link_width"] = int(width) if width.isdigit() else None
+                # speed is e.g. "32.0 GT/s PCIe" -> 32.0
+                try:
+                    low["pcie_link_speed_gts"] = float(speed.split()[0]) if speed else None
+                except Exception:
+                    low["pcie_link_speed_gts"] = None
+            else:
+                low["pcie_error"] = (
+                    f"could not resolve PCIe BDF (pci_bus_id=" f"{getattr(props, 'pci_bus_id', None)!r})"
+                )
+        except Exception as e:
+            low["pcie_error"] = f"PCIe sysfs capture failed: {e}"
+
+        # HBM total/free (torch). Independent of BDF resolution.
+        try:
+            free_b, total_b = torch.cuda.mem_get_info(gpu)
+            low["hbm_total_bytes"] = int(total_b)
+            low["hbm_free_bytes"] = int(free_b)
+            low["hbm_total_gib"] = round(total_b / (1 << 30), 2)
+        except Exception as e:
+            low["hbm_error"] = f"mem_get_info failed: {e}"
+            tm = int(getattr(props, "total_memory", 0) or 0)
+            if tm:
+                low["hbm_total_bytes"] = tm
+                low["hbm_total_gib"] = round(tm / (1 << 30), 2)
+    if low:
+        details["low_level"] = low
+
+    # --- Tier 2 perf sanity (optional) ---
+    if tier2_perf:
+        # GEMM TFLOPS. Warmup/iter counts mirror the preflight `--quick` preset
+        # (`square_gemm.py` with WARMUP=5, ITERATION=20) so smoke and preflight
+        # report comparable steady-state numbers.
+        try:
+            tflops = _measure_gemm_tflops(gpu, size=8192, warmup=5, iters=20)
+            details["gemm_tflops"] = round(tflops, 2)
+            if tflops < gemm_tflops_min:
+                return {
+                    "gpu": gpu,
+                    "status": "FAIL",
+                    "reason": (f"GEMM TFLOPS {tflops:.0f} < threshold {gemm_tflops_min:.0f}"),
+                    "duration_sec": round(time.time() - t0, 3),
+                    "details": details,
+                }
+        except Exception as e:
+            return {
+                "gpu": gpu,
+                "status": "FAIL",
+                "reason": f"GEMM TFLOPS measurement failed: {e}",
+                "duration_sec": round(time.time() - t0, 3),
+                "details": details,
+            }
+
+        # HBM device-to-device copy bandwidth. HBM is fast enough that we
+        # need a healthy number of timed iterations for stable timing.
+        try:
+            gbs = _measure_hbm_gbs(gpu, size_bytes=512 * 1024 * 1024, warmup=10, iters=20)
+            details["hbm_gbs"] = round(gbs, 1)
+            if gbs < hbm_gbs_min:
+                return {
+                    "gpu": gpu,
+                    "status": "FAIL",
+                    "reason": f"HBM GB/s {gbs:.0f} < threshold {hbm_gbs_min:.0f}",
+                    "duration_sec": round(time.time() - t0, 3),
+                    "details": details,
+                }
+        except Exception as e:
+            return {
+                "gpu": gpu,
+                "status": "FAIL",
+                "reason": f"HBM bandwidth measurement failed: {e}",
+                "duration_sec": round(time.time() - t0, 3),
+                "details": details,
+            }
+
+    return {
+        "gpu": gpu,
+        "status": "PASS",
+        "reason": "",
+        "duration_sec": round(time.time() - t0, 3),
+        "details": details,
+    }
+
+
+def _measure_gemm_tflops(gpu: int, *, size: int, warmup: int, iters: int) -> float:
+    """Measure GEMM TFLOPS for square ``size x size`` bf16 matmul on ``cuda:gpu``."""
+    import torch  # type: ignore
+
+    torch.cuda.set_device(gpu)
+    a = torch.randn((size, size), dtype=torch.bfloat16, device=f"cuda:{gpu}")
+    b = torch.randn((size, size), dtype=torch.bfloat16, device=f"cuda:{gpu}")
+    for _ in range(warmup):
+        _ = torch.matmul(a, b)
+    torch.cuda.synchronize()
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        _ = torch.matmul(a, b)
+    end.record()
+    end.synchronize()
+
+    elapsed_s = start.elapsed_time(end) / 1000.0 / iters
+    flops = 2.0 * size * size * size
+    return (flops / elapsed_s) / 1e12
+
+
+def _measure_hbm_gbs(gpu: int, *, size_bytes: int, warmup: int, iters: int) -> float:
+    """Measure local HBM bandwidth via device-to-device ``copy_``.
+
+    Each iteration: 1 read of ``src`` + 1 write to ``dst`` = ``2 * size_bytes``
+    of HBM traffic. Uses ``torch.cuda.Event`` for accurate GPU-side timing.
+    """
+    import torch  # type: ignore
+
+    torch.cuda.set_device(gpu)
+    n = size_bytes // 2  # bf16 = 2 bytes/element
+    src = torch.empty(n, dtype=torch.bfloat16, device=f"cuda:{gpu}")
+    dst = torch.empty_like(src)
+    src.fill_(1.0)
+    for _ in range(warmup):
+        dst.copy_(src)
+    torch.cuda.synchronize()
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        dst.copy_(src)
+    end.record()
+    end.synchronize()
+
+    elapsed_s = start.elapsed_time(end) / 1000.0 / iters
+    return (2.0 * size_bytes / elapsed_s) / 1e9
diff --git a/primus/tools/preflight/node_smoke/rccl_local.py b/primus/tools/preflight/node_smoke/rccl_local.py
new file mode 100644
index 000000000..6e34dca19
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/rccl_local.py
@@ -0,0 +1,152 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tier 2 -- node-local RCCL all-reduce (optional).
+
+Tier 2 also includes a node-local RCCL all-reduce as a steady-state
+intra-node bandwidth check. We use ``torch.multiprocessing.spawn`` to
+launch one worker per local GPU on a process group bound to
+``tcp://127.0.0.1:<free port>``. No cross-node communication.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import socket
+import time
+from typing import Any, Dict, List
+
+
+def _rccl_worker(
+    local_rank: int,
+    world_size: int,
+    port: int,
+    size_mb: int,
+    out_path: str,
+) -> None:
+    """Subprocess body for ``torch.multiprocessing.spawn``.
+
+    Runs ``warmup`` warmup + ``iters`` timed all-reduces of ``size_mb`` MB on a
+    local-only NCCL/RCCL process group bound to ``tcp://127.0.0.1:port``.
+    Local rank 0 writes the resulting GB/s to ``out_path``.
+
+    Iteration counts are intentionally aligned with the preflight `--quick`
+    preset (`intra_node_comm.py` with WARMUP=5, ITERATION=20) so smoke and
+    preflight report comparable steady-state bandwidth.
+    """
+    import torch  # type: ignore
+    import torch.distributed as dist  # type: ignore
+
+    warmup = 5
+    iters = 20
+
+    try:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            init_method=f"tcp://127.0.0.1:{port}",
+            world_size=world_size,
+            rank=local_rank,
+        )
+        nbytes = size_mb * 1024 * 1024
+        n_elem = nbytes // 2  # bf16
+        t = torch.ones(n_elem, dtype=torch.bfloat16, device=f"cuda:{local_rank}")
+
+        for _ in range(warmup):
+            dist.all_reduce(t)
+        torch.cuda.synchronize()
+
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            dist.all_reduce(t)
+        end.record()
+        end.synchronize()
+
+        elapsed_s = start.elapsed_time(end) / 1000.0 / iters
+        # NCCL all-reduce effective bandwidth: 2*S*(P-1)/P bytes per rank.
+        comm_bytes = 2.0 * nbytes * (world_size - 1) / world_size
+        gbs = comm_bytes / elapsed_s / 1e9
+
+        if local_rank == 0:
+            with open(out_path, "w", encoding="utf-8") as f:
+                json.dump({"status": "PASS", "gbs": round(gbs, 1)}, f)
+        dist.barrier()
+        dist.destroy_process_group()
+    except Exception as e:
+        if local_rank == 0:
+            try:
+                with open(out_path, "w", encoding="utf-8") as f:
+                    json.dump({"status": "FAIL", "error": str(e)}, f)
+            except Exception:
+                pass
+
+
+def _run_local_rccl(*, local_world_size: int, size_mb: int, timeout_sec: int) -> Dict[str, Any]:
+    """Spawn local-only RCCL workers to measure intra-node all-reduce bandwidth.
+
+    Returns ``{"status": "PASS"|"FAIL"|"TIMEOUT", ...}``.
+    """
+    import tempfile
+
+    if local_world_size <= 1:
+        return {"status": "PASS", "gbs": None, "skipped": "local_world_size<=1"}
+
+    try:
+        import torch  # type: ignore
+        import torch.multiprocessing as mp  # type: ignore
+    except Exception as e:
+        return {"status": "FAIL", "error": f"torch import failed: {e}"}
+
+    # Pick a free local TCP port.
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        port = s.getsockname()[1]
+
+    with tempfile.NamedTemporaryFile(prefix="node_smoke_rccl_", suffix=".json", delete=False) as tf:
+        out_path = tf.name
+
+    ctx = mp.get_context("spawn")
+    procs: List[Any] = []
+    try:
+        for r in range(local_world_size):
+            p = ctx.Process(
+                target=_rccl_worker,
+                args=(r, local_world_size, port, size_mb, out_path),
+            )
+            p.start()
+            procs.append(p)
+
+        deadline = time.time() + timeout_sec
+        for p in procs:
+            remaining = max(0.0, deadline - time.time())
+            p.join(timeout=remaining)
+            if p.is_alive():
+                # One worker stuck -> kill all and report TIMEOUT.
+                for q in procs:
+                    if q.is_alive():
+                        q.terminate()
+                for q in procs:
+                    q.join(timeout=5)
+                    if q.is_alive():
+                        q.kill()
+                return {
+                    "status": "TIMEOUT",
+                    "error": f"local RCCL all-reduce did not finish in {timeout_sec}s",
+                }
+
+        if not os.path.exists(out_path):
+            return {"status": "FAIL", "error": "no result file produced"}
+        with open(out_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        return data
+    finally:
+        try:
+            os.unlink(out_path)
+        except Exception:
+            pass
diff --git a/primus/tools/preflight/node_smoke/shell_utils.py b/primus/tools/preflight/node_smoke/shell_utils.py
new file mode 100644
index 000000000..952044125
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/shell_utils.py
@@ -0,0 +1,179 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Stateless utility helpers shared across collectors and orchestration.
+
+Every function here is pure (no I/O effects beyond reading sysfs / running
+a tiny subprocess) and never raises -- callers expect best-effort
+defaults so a missing tool / file degrades gracefully.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Any, Dict, List, Optional
+
+
+def _which(prog: str) -> Optional[str]:
+    """Tiny shutil.which() replacement that doesn't pull in shutil at import."""
+    for d in (os.environ.get("PATH") or "").split(os.pathsep):
+        p = os.path.join(d, prog)
+        if os.path.isfile(p) and os.access(p, os.X_OK):
+            return p
+    return None
+
+
+def _read_text(path: str, default: str = "") -> str:
+    """Best-effort read of a small sysfs/proc text file. Never raises."""
+    try:
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            return f.read().strip()
+    except Exception:
+        return default
+
+
+def _parse_os_release_pretty() -> Optional[str]:
+    """Return PRETTY_NAME from /etc/os-release, or None."""
+    try:
+        with open("/etc/os-release", "r", encoding="utf-8") as f:
+            for line in f:
+                if line.startswith("PRETTY_NAME="):
+                    v = line.split("=", 1)[1].strip().strip('"').strip("'")
+                    return v
+    except Exception:
+        pass
+    return None
+
+
+def _resolve_gpu_bdf(props: Any) -> Optional[str]:
+    """Return the PCIe BDF (e.g. ``"0000:75:00.0"``) for a torch device.
+
+    ``torch.cuda.get_device_properties(i).pci_bus_id`` is annoyingly
+    polymorphic across PyTorch + ROCm versions: sometimes a string in the
+    canonical ``domain:bus:device.function`` form (lowercase or uppercase),
+    sometimes an int (just the bus byte). We coerce all variants into the
+    canonical lowercase form and verify the sysfs directory actually
+    exists before returning -- so the caller can read PCIe link info
+    without an extra existence check.
+
+    Returns None if the BDF cannot be resolved (caller should still
+    capture HBM via ``mem_get_info`` and skip PCIe sysfs reads).
+    """
+    raw = getattr(props, "pci_bus_id", None)
+    if raw is None:
+        return None
+    # 1) String form. Could be "0000:05:00.0", "05:00.0", or uppercase.
+    if isinstance(raw, str):
+        s = raw.strip().lower()
+        if not s:
+            return None
+        candidates = [s, f"0000:{s}" if not s.startswith("0000:") else s]
+        for c in candidates:
+            if os.path.isdir(f"/sys/bus/pci/devices/{c}"):
+                return c
+        return None
+    # 2) Int form (just the bus byte). Standard layout for AMD GPUs is
+    # 0000:<bus>:00.0; verify with sysfs and fall back to a glob if the
+    # device.function differs from 00.0 on this host.
+    if isinstance(raw, int):
+        bus_hex = f"{raw:02x}"
+        primary = f"0000:{bus_hex}:00.0"
+        if os.path.isdir(f"/sys/bus/pci/devices/{primary}"):
+            return primary
+        import glob
+
+        matches = sorted(glob.glob(f"/sys/bus/pci/devices/0000:{bus_hex}:*"))
+        if matches:
+            return os.path.basename(matches[0])
+        return None
+    return None
+
+
+def _systemctl_is_active(unit: str) -> Optional[str]:
+    """Return ``systemctl is-active <unit>`` ('active'/'inactive'/'failed'/...)
+    or None if systemctl is missing / errors. Always best-effort."""
+    if _which("systemctl") is None:
+        return None
+    try:
+        cp = subprocess.run(
+            ["systemctl", "is-active", unit],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=3,
+            check=False,
+        )
+        # systemctl returns non-zero for inactive/failed -- that's fine,
+        # we just want the textual state.
+        return (cp.stdout or "").strip() or "unknown"
+    except Exception:
+        return None
+
+
+def _parse_size_with_unit(s: str) -> Optional[int]:
+    """Parse a size string into bytes.
+
+    Accepts (case-insensitive)::
+
+        "12345"        -> 12345
+        "256 MB"       -> 268435456
+        "256MB"        -> 268435456     (no space)
+        "12.5 GiB"     -> 13421772800
+        "12.5GiB"      -> 13421772800
+        "  -1  "       -> -1            (sentinel for "unlimited")
+
+    Returns ``None`` for empty input, an unrecognised unit (e.g.
+    ``"500 MHz"``), or any input the regex cannot fully match (e.g.
+    ``"12 GB extra"``). This is deliberate: a silent fallthrough
+    to ``num * 1`` would let frequency / count strings masquerade
+    as byte counts and corrupt downstream comparisons.
+    """
+    if not s:
+        return None
+    import re
+
+    m = re.match(
+        r"\s*([+-]?\d+(?:\.\d+)?)\s*([a-zA-Z]+)?\.?\s*$",
+        s,
+    )
+    if not m:
+        return None
+    units = {
+        "b": 1,
+        "k": 1 << 10,
+        "kb": 1 << 10,
+        "kib": 1 << 10,
+        "m": 1 << 20,
+        "mb": 1 << 20,
+        "mib": 1 << 20,
+        "g": 1 << 30,
+        "gb": 1 << 30,
+        "gib": 1 << 30,
+        "t": 1 << 40,
+        "tb": 1 << 40,
+        "tib": 1 << 40,
+    }
+    unit = (m.group(2) or "b").lower()
+    if unit not in units:
+        return None
+    try:
+        return int(float(m.group(1)) * units[unit])
+    except ValueError:
+        return None
+
+
+def _findings_to_dicts(findings: List[Any]) -> List[Dict[str, Any]]:
+    """Normalize Finding dataclasses (different modules each define their own)
+    into plain dicts."""
+    return [
+        {
+            "level": getattr(f, "level", "info"),
+            "message": getattr(f, "message", str(f)),
+            "details": getattr(f, "details", {}),
+        }
+        for f in findings
+    ]
diff --git a/primus/modules/trainer/maxtext/__init__.py b/primus/tools/preflight/node_smoke/tests/__init__.py
similarity index 100%
rename from primus/modules/trainer/maxtext/__init__.py
rename to primus/tools/preflight/node_smoke/tests/__init__.py
diff --git a/primus/tools/preflight/node_smoke/tests/test_node_smoke.py b/primus/tools/preflight/node_smoke/tests/test_node_smoke.py
new file mode 100644
index 000000000..a2e9135f5
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/tests/test_node_smoke.py
@@ -0,0 +1,1109 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Regression tests for the node-smoke package.
+
+Each test is a deliberate guard against a behaviour that has either
+broken in production (history items in `docs/node-smoke.md`) or is part
+of the operator-facing contract (CLI flags, JSON schema, report section
+order). They are pure-Python: no GPU, no subprocess fanout, no real
+amd-smi/rocm-smi/lsof dependency.
+
+Run directly:
+
+    pytest primus/tools/preflight/node_smoke/tests/test_node_smoke.py -v
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+
+from primus.tools.preflight.node_smoke.aggregator.report import write_smoke_report
+from primus.tools.preflight.node_smoke.aggregator.summarizers import (
+    _busy_gpu_rows,
+    _clock_summary,
+    _pretouch_hbm_rows,
+    _stack_drift_rows,
+)
+from primus.tools.preflight.node_smoke.collectors.gpu_processes import (
+    _flatten_amd_smi_process_json,
+    _parse_lsof_pcn,
+)
+from primus.tools.preflight.node_smoke.collectors.nics import (
+    _parse_nic_selector,
+    _resolve_selector,
+    _selector_matches,
+)
+from primus.tools.preflight.node_smoke.collectors.rocm_smi import (
+    _parse_rocm_smi_ras_info_text,
+)
+from primus.tools.preflight.node_smoke.logging_utils import _short_name
+from primus.tools.preflight.node_smoke.orchestrator import _node_status_from
+from primus.tools.preflight.node_smoke.shell_utils import _parse_size_with_unit
+
+# ---------------------------------------------------------------------------
+# A. Pure-helper unit tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "s,expected",
+    [
+        # Plain ints + decimals (no unit -> bytes)
+        ("12345", 12345),
+        ("", None),
+        # SI / IEC mix, both with and without space
+        ("256 MB", 256 * (1 << 20)),
+        ("256MB", 256 * (1 << 20)),
+        ("12.5 GiB", int(12.5 * (1 << 30))),
+        ("12.5GiB", int(12.5 * (1 << 30))),
+        # The "unlimited" sentinel
+        ("  -1  ", -1),
+        # Things that MUST NOT silently become byte-counts. `500 MHz` was
+        # the historical bug -- frequency masquerading as bytes.
+        ("500 MHz", None),
+        ("12 GB extra", None),
+        ("not-a-number", None),
+    ],
+)
+def test_parse_size_with_unit(s, expected):
+    """A.1 -- regex + unit table; unrecognized units MUST return None."""
+    assert _parse_size_with_unit(s) == expected
+
+
+def test_short_name_strips_fqdn():
+    """A.2 -- FQDN normalisation (history item: SLURM-ready txt outputs)."""
+    assert _short_name("node001.cluster.example.com") == "node001"
+    assert _short_name("node001") == "node001"
+    assert _short_name("") == ""
+
+
+def test_parse_rocm_smi_ras_sums_per_gpu_and_skips_status_only_rows():
+    """A.3 -- ECC text parser: one GPU block per `GPU[N]: RAS INFO` header,
+    summing the last two int columns of each block row, ignoring rows that
+    have only Status (e.g. ATHUB UNAVAILABLE)."""
+    sample = (
+        "GPU[0]:         RAS INFO\n"
+        "        Block       Status    Correctable Error  Uncorrectable Error\n"
+        "          UMC        ENABLED                  3                    1\n"
+        "         SDMA        ENABLED                  0                    0\n"
+        "       ATHUB     UNAVAILABLE\n"
+        "GPU[1]:         RAS INFO\n"
+        "          UMC        ENABLED                  0                    0\n"
+    )
+    out = _parse_rocm_smi_ras_info_text(sample)
+    assert out == [
+        {"gpu": 0, "ecc_correctable_total": 3, "ecc_uncorrectable_total": 1},
+        {"gpu": 1, "ecc_correctable_total": 0, "ecc_uncorrectable_total": 0},
+    ]
+
+
+def test_parse_lsof_pcn_handles_multi_open_pid():
+    """A.4 -- lsof -Fpcn field-prefix parser: same PID across multiple open
+    files collapses to a single annotated record per PID."""
+    text = "p123\ncpython\nf3\nf4\np123\ncpython\np456\nctrain.py\n"
+    rows = _parse_lsof_pcn(
+        text,
+        lambda pid, name, hbm: {"pid": pid, "name": name, "hbm_bytes": hbm},
+    )
+    assert sorted(r["pid"] for r in rows) == [123, 456]
+    assert {r["pid"]: r["name"] for r in rows} == {
+        123: "python",
+        456: "train.py",
+    }
+
+
+# ---------------------------------------------------------------------------
+# B. Aggregator-summarizer regression guards
+# ---------------------------------------------------------------------------
+
+
+def test_stack_drift_does_not_crash_on_heterogeneous_dict_vs_none():
+    """B.1 -- guards against the production crash captured in
+    docs/node-smoke.md history item 6: a fingerprint key that is None on
+    one node and a dict on another USED to crash Counter() with
+    `TypeError: unhashable type: 'dict'`. After the fix the key must
+    simply be excluded from scalar drift (since it is not a scalar on
+    any node)."""
+    nodes = [
+        {
+            "host": "a",
+            "tier1": {"fingerprint": {"kernel": "5.15", "rocm": "6.2", "nic_fw": None}},
+        },
+        {
+            "host": "b",
+            "tier1": {"fingerprint": {"kernel": "5.15", "rocm": "6.2", "nic_fw": {"rdma0": "20.0"}}},
+        },
+    ]
+    # Must return without raising.
+    rows = _stack_drift_rows(nodes)
+    # And nic_fw must NOT appear in scalar drift (it's not a scalar
+    # on any node).
+    assert all(r["key"] != "nic_fw" for r in rows)
+
+
+def test_clock_summary_spread_and_warn():
+    """B.2 -- spread = max - min, warn flag fires above threshold,
+    nodes with no active time daemon are listed."""
+    nodes = [
+        {"host": "a", "tier1": {"clock": {"wall_time_unix": 1000.0, "any_active": True}}},
+        {"host": "b", "tier1": {"clock": {"wall_time_unix": 1042.5, "any_active": True}}},
+        {"host": "c", "tier1": {"clock": {"wall_time_unix": 1010.0, "any_active": False}}},
+    ]
+    s = _clock_summary(nodes, skew_warn_sec=30.0)
+    assert s["spread_sec"] == 42.5
+    assert s["spread_warn"] is True  # 42.5 > 30
+    # node_rank defaults to "?" when missing from the input dict.
+    assert ("?", "c") in s["no_daemon_hosts"]
+    assert s["earliest_host"] == "a"
+    assert s["latest_host"] == "b"
+
+
+def test_busy_gpu_rows_filters_to_is_foreign_only():
+    """B.3 -- only processes flagged is_foreign make it into the
+    Busy-GPU table; HBM is converted from bytes to GiB."""
+    nodes = [
+        {
+            "host": "a",
+            "node_rank": 0,
+            "tier1": {
+                "gpu_processes": {
+                    "ok": True,
+                    "per_gpu": [
+                        {
+                            "gpu": 0,
+                            "processes": [
+                                {
+                                    "pid": 1,
+                                    "name": "self",
+                                    "hbm_bytes": 1 << 30,
+                                    "is_self": True,
+                                    "is_allowed": False,
+                                    "is_foreign": False,
+                                },
+                                {
+                                    "pid": 2,
+                                    "name": "agent",
+                                    "hbm_bytes": 0,
+                                    "is_self": False,
+                                    "is_allowed": True,
+                                    "is_foreign": False,
+                                },
+                                {
+                                    "pid": 3,
+                                    "name": "leak",
+                                    "hbm_bytes": 4 * (1 << 30),
+                                    "is_self": False,
+                                    "is_allowed": False,
+                                    "is_foreign": True,
+                                },
+                            ],
+                        }
+                    ],
+                }
+            },
+        }
+    ]
+    rows = _busy_gpu_rows(nodes)
+    assert [r["pid"] for r in rows] == [3]
+    assert rows[0]["hbm_gib"] == 4.0
+    assert rows[0]["name"] == "leak"
+
+
+# ---------------------------------------------------------------------------
+# C. Orchestrator decision-logic tests
+# ---------------------------------------------------------------------------
+
+
+def test_node_status_empty_means_pass():
+    """C.1 -- no per-GPU results, no Tier 1 issues, no Tier 2 issues
+    -> empty reasons list -> node PASS."""
+    assert _node_status_from([], {}, {}) == []
+
+
+def test_node_status_nonzero_ecc_fails():
+    """C.2 -- any non-zero uncorrectable ECC count is an unconditional
+    node FAIL. The amd-smi schema is unstable so we trust only ints."""
+    tier1 = {"gpu_low_level": {"per_gpu": [{"gpu": 3, "ecc_uncorrectable_total": 7}]}}
+    reasons = _node_status_from([], tier1, {})
+    assert any("gpu3" in r and "uncorrectable" in r for r in reasons)
+
+
+def test_node_status_allow_foreign_procs_downgrades():
+    """C.3 -- foreign processes hard-fail the node by default, but
+    --allow-foreign-procs downgrades to silent inclusion in the JSON."""
+    tier1 = {
+        "gpu_processes": {
+            "ok": True,
+            "foreign_count": 1,
+            "per_gpu": [
+                {
+                    "gpu": 0,
+                    "processes": [
+                        {
+                            "pid": 99,
+                            "name": "leak",
+                            "hbm_bytes": 1 << 30,
+                            "is_foreign": True,
+                        }
+                    ],
+                }
+            ],
+        }
+    }
+    # Default (allow_foreign_procs=False) -> FAIL.
+    assert _node_status_from([], tier1, {}, allow_foreign_procs=False)
+    # Operator opted in -> PASS.
+    assert _node_status_from([], tier1, {}, allow_foreign_procs=True) == []
+
+
+def test_node_status_require_tools_missing_amd_smi_fails():
+    """C.4 -- --require-tools promotes a missing CLI tool to a hard
+    node FAIL; satisfied requirements remain silent."""
+    tier1 = {
+        "tooling_inventory": {
+            "tools": {
+                "amd-smi": {"present": False, "path": None},
+                "rocm-smi": {"present": True, "path": "/usr/bin/rocm-smi"},
+                "lsof": {"present": True, "path": "/usr/bin/lsof"},
+            }
+        }
+    }
+    # amd-smi required but absent -> FAIL.
+    reasons = _node_status_from([], tier1, {}, required_tools=["amd-smi", "rocm-smi"])
+    assert any("amd-smi" in r for r in reasons)
+    # Only rocm-smi required and present -> no reason added.
+    assert _node_status_from([], tier1, {}, required_tools=["rocm-smi"]) == []
+
+
+# ---------------------------------------------------------------------------
+# D. CLI / report parity tests
+# ---------------------------------------------------------------------------
+
+
+# Section ORDER + headings are part of the operator-facing contract --
+# slack bots and CI scripts grep for these. Update this list ONLY when
+# you intentionally change the contract.
+EXPECTED_SECTIONS = [
+    "## Stack drift across cluster",
+    "## NIC firmware drift across cluster",
+    "## NIC / RDMA roll-call issues",
+    "## NIC port-count summary",
+    "## NIC excluded ports (informational)",
+    "## Host limits issues",
+    "## GPU visibility issues",
+    "## GPU low-level outliers (PCIe link / HBM)",
+    "## XGMI link issues",
+    "## Cluster clock + time daemons",
+    "## Tooling self-latency (`rocm-smi --version`)",
+    "## Tooling availability",
+    "## Busy GPUs / leaked processes",
+    "## GPU pre-touch HBM usage outliers",
+    "## GPU compute-activity outliers",
+    # Tier 2 perf summary intentionally omitted -- only renders when at
+    # least one node ran Tier 2; this fixture has none.
+    "## Failing nodes -- full reasons",  # only when failing is non-empty
+]
+
+
+def test_report_section_order_stable(tmp_path):
+    """D.1 -- report section order + headings are stable. This catches
+    anyone who reorders the _write_<section> calls in
+    aggregator.report.write_smoke_report."""
+    nodes = [
+        {
+            "host": "a",
+            "node_rank": 0,
+            "status": "FAIL",
+            "duration_sec": 1.0,
+            "fail_reasons": ["xgmi: 1 non-XGMI pair"],
+            "tier1": {},
+            "tier2": {},
+        }
+    ]
+    out = tmp_path / "report.md"
+    write_smoke_report(
+        str(out),
+        nodes=nodes,
+        passing=[],
+        failing=nodes,
+        expected=1,
+        clock_skew_warn_sec=30.0,
+        rocm_smi_warn_sec=1.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+    )
+    seen = [line for line in out.read_text().splitlines() if line.startswith("## ")]
+    assert seen == EXPECTED_SECTIONS
+
+
+def test_module_help_exits_zero():
+    """D.2 -- `python -m primus.tools.preflight.node_smoke --help` must
+    exit 0. Doubles as a smoke test that the package imports cleanly
+    under `-m` (i.e. __main__.py + __init__.py + cli.py are all
+    consistent)."""
+    cp = subprocess.run(
+        [sys.executable, "-m", "primus.tools.preflight.node_smoke", "--help"],
+        capture_output=True,
+        text=True,
+        timeout=15,
+    )
+    assert cp.returncode == 0, cp.stderr
+    assert "node-local preflight smoke test" in cp.stdout.lower()
+
+
+# ---------------------------------------------------------------------------
+# E. amd-smi schema-drift detection + HBM threshold boundary guards
+# ---------------------------------------------------------------------------
+#
+# These tests lock in two operator-facing contracts:
+#
+#   1. _flatten_amd_smi_process_json must return (parsed, drift) where
+#      `drift` is True only when at least one record looked like a process
+#      dict structurally but had no usable pid/process_id. This lets the
+#      caller fall through to text/rocm-smi/lsof on a future amd-smi
+#      schema rename instead of silently reporting the node clean.
+#
+#   2. The aggregator's pre-touch HBM threshold uses `>=` (inclusive
+#      boundary). A GPU sitting at exactly the threshold MUST appear in
+#      the rendered `smoke_report.md` outliers section; the operator-
+#      facing help / docs / report wording is aligned to "at least".
+#      (We test the aggregator side because that's the artifact
+#      operators read; the runner-side `>=` in per_gpu.py is the
+#      one-line counterpart, kept honest by argparse help + code review.)
+
+
+def _annotate_passthrough(pid, name, hbm):
+    """Minimal annotate() shim for the schema-drift tests."""
+    return {"pid": pid, "name": name, "hbm_bytes": hbm}
+
+
+def test_flatten_amd_smi_top_level_unknown_returns_empty_no_drift():
+    """E.1 -- a doc that is neither Shape A/A'/B (no `process_list`, no
+    pid-bearing items) returns ([], False). False on `drift` is what
+    distinguishes "schema we do not speak yet at the top level" from
+    "schema we do speak but per-process fields drifted"; the caller
+    falls through in either case but the message differs."""
+    parsed, drift = _flatten_amd_smi_process_json({"unexpected_key": 1}, _annotate_passthrough)
+    assert parsed == []
+    assert drift is False
+    parsed, drift = _flatten_amd_smi_process_json([{"some_other_field": 42}], _annotate_passthrough)
+    assert parsed == []
+    assert drift is False
+
+
+def test_flatten_amd_smi_shape_a_empty_process_list_no_drift():
+    """E.2 -- Shape A with empty process_list is the clean-node case.
+    Buckets are pre-registered (so the caller can tell schema-matched
+    from schema-not-matched), processes lists are empty, drift is False
+    -- the caller trusts the result and skips the fallback chain."""
+    doc = [
+        {"gpu": 0, "process_list": []},
+        {"gpu": 1, "process_list": []},
+    ]
+    parsed, drift = _flatten_amd_smi_process_json(doc, _annotate_passthrough)
+    assert parsed == [
+        {"gpu": 0, "processes": []},
+        {"gpu": 1, "processes": []},
+    ]
+    assert drift is False
+
+
+def test_flatten_amd_smi_shape_a_renamed_pid_field_raises_drift():
+    """E.3 -- the future-schema-drift smoking gun: a Shape A entry whose
+    process_list items are dicts (so the structural shape matches) but
+    use a renamed key (`proc_pid`) instead of `pid` / `process_id`.
+    The bucket is registered, processes is empty, but drift=True so the
+    caller falls through to amd-smi text / rocm-smi / lsof rather than
+    silently reporting the GPU clean."""
+    doc = [
+        {
+            "gpu": 0,
+            "process_list": [
+                {"proc_pid": 1234, "name": "python"},
+            ],
+        }
+    ]
+    parsed, drift = _flatten_amd_smi_process_json(doc, _annotate_passthrough)
+    assert parsed == [{"gpu": 0, "processes": []}]
+    assert drift is True
+
+
+def test_pretouch_hbm_rows_includes_gpu_at_exact_threshold():
+    """E.4 -- the aggregator's `_pretouch_hbm_rows` uses `>= threshold`,
+    which is what populates the `GPU pre-touch HBM usage outliers`
+    section in `smoke_report.md`. Lock in the inclusive boundary at the
+    exact value so the wording fix in `aggregator/report.py` and the
+    runtime stay aligned: a GPU sitting at threshold MUST be listed,
+    a GPU below MUST NOT, a GPU above MUST."""
+    nodes = [
+        {
+            "host": "host-a",
+            "node_rank": 0,
+            "tier1": {
+                "per_gpu": [
+                    {"gpu": 0, "details": {"hbm_pre_touch_used_gib": 2.0}},  # ==
+                    {"gpu": 1, "details": {"hbm_pre_touch_used_gib": 1.99}},  # <
+                    {"gpu": 2, "details": {"hbm_pre_touch_used_gib": 2.01}},  # >
+                ],
+            },
+        },
+    ]
+    rows = _pretouch_hbm_rows(nodes, threshold_gib=2.0)
+    assert sorted(r["gpu"] for r in rows) == [0, 2]
+
+
+# ---------------------------------------------------------------------------
+# F. NIC training-NIC selector (NCCL_IB_HCA-style allowlist + heuristic
+#    fallback + precedence chain)
+#
+# These guard the operator-facing contract introduced when multi-role
+# clusters (front-end / storage RoCE NICs co-resident with the training
+# NICs) started showing up. The fail mode they protect against: 56 of 57
+# real production nodes were reported as FAIL because two unplugged
+# front-end ports showed `state=DOWN phys_state=Disabled`, even though
+# `NCCL_IB_HCA` explicitly listed only the 8 healthy back-end NICs.
+# ---------------------------------------------------------------------------
+
+
+def _install_synthetic_ib_tree(monkeypatch, base):
+    """Redirect `_collect_nic_status`'s sysfs reads at a tmp tree.
+
+    The collector hard-codes ``/sys/class/infiniband`` so we monkey-patch
+    its ``os.path.isdir`` / ``os.listdir`` / ``_read_text`` symbols to
+    rewrite that prefix to ``base``. We DELIBERATELY rewrite at the
+    string level (not via a chroot-style abstraction) because the
+    rewrite has to also fire for the recursive sub-paths the collector
+    constructs (``<base>/<dev>/ports/<port>/state`` etc.) -- the rewrite
+    must therefore intercept every path-shaped call inside the
+    collector, not just the entry-point listdir.
+    """
+    import os
+
+    real_listdir = os.listdir
+    real_isdir = os.path.isdir
+
+    def _isdir(p: str) -> bool:
+        if p == "/sys/class/infiniband":
+            return True
+        return real_isdir(p.replace("/sys/class/infiniband", str(base)))
+
+    def _listdir(p: str):
+        return real_listdir(p.replace("/sys/class/infiniband", str(base)))
+
+    def _read_text(path: str) -> str:
+        try:
+            with open(path.replace("/sys/class/infiniband", str(base)), "r") as fh:
+                return fh.read().strip()
+        except Exception:
+            return ""
+
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics.os.path.isdir", _isdir)
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics.os.listdir", _listdir)
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics._read_text", _read_text)
+
+
+def test_parse_nic_selector_allowlist_with_ports():
+    """F.1 -- baseline NCCL_IB_HCA syntax: comma-separated `device:port`."""
+    sel = _parse_nic_selector("rocep158s0:1,rocep190s0:1")
+    assert sel["mode"] == "allowlist"
+    assert sel["entries"] == [
+        ("rocep158s0", False, 1),
+        ("rocep190s0", False, 1),
+    ]
+
+
+def test_parse_nic_selector_denylist_prefix():
+    """F.2 -- a leading `^` flips the whole selector to denylist."""
+    sel = _parse_nic_selector("^roceo12399,roceo12409")
+    assert sel["mode"] == "denylist"
+    # The ^ is stripped from the entries; only the global mode changes.
+    assert [e[0] for e in sel["entries"]] == ["roceo12399", "roceo12409"]
+
+
+def test_parse_nic_selector_exact_prefix():
+    """F.3 -- `=name` forces exact (not prefix) device-name matching."""
+    sel = _parse_nic_selector("=mlx5,mlx5_other")
+    assert sel["entries"][0] == ("mlx5", True, None)  # exact-match flag set
+    assert sel["entries"][1] == ("mlx5_other", False, None)
+
+
+def test_parse_nic_selector_no_port_matches_any_port():
+    """F.4 -- entries without `:port` accept any port on the device."""
+    sel = _parse_nic_selector("mlx5_0")
+    assert sel["entries"] == [("mlx5_0", False, None)]
+    assert _selector_matches(sel, "mlx5_0", 1) is True
+    assert _selector_matches(sel, "mlx5_0", 2) is True
+
+
+def test_parse_nic_selector_empty_is_passthrough():
+    """F.5 -- empty / whitespace-only input means 'no selector', so the
+    caller falls through to the next precedence layer (env / heuristic)
+    rather than treating it as 'allowlist nothing' (which would silently
+    exclude every port)."""
+    assert _parse_nic_selector("")["mode"] == "passthrough"
+    assert _parse_nic_selector("   ")["mode"] == "passthrough"
+    assert _parse_nic_selector("^")["mode"] == "passthrough"  # ^ but no entries
+    assert _parse_nic_selector(",, ,")["mode"] == "passthrough"
+
+
+def test_selector_matches_prefix_matches_mlx5_to_mlx5_0():
+    """F.6 -- NCCL prefix semantics: `mlx5` allowlist entry matches every
+    device that starts with `mlx5` (the historical NCCL behavior). This
+    is why operators must use `=mlx5` if they want exact-only matching."""
+    sel = _parse_nic_selector("mlx5")
+    assert _selector_matches(sel, "mlx5", 1) is True
+    assert _selector_matches(sel, "mlx5_0", 1) is True
+    assert _selector_matches(sel, "mlx5_bond_0", 1) is True
+    assert _selector_matches(sel, "rocep158s0", 1) is False
+
+
+def test_selector_matches_denylist_inverts():
+    """F.7 -- denylist semantics: ports NOT in the list pass; ports in
+    the list are rejected. This is the more ergonomic form for clusters
+    with many training NICs and a small handful of front-end ports."""
+    sel = _parse_nic_selector("^roceo12399,roceo12409")
+    # Listed -> excluded.
+    assert _selector_matches(sel, "roceo12399", 1) is False
+    assert _selector_matches(sel, "roceo12409", 1) is False
+    # Not listed -> included.
+    assert _selector_matches(sel, "rocep158s0", 1) is True
+
+
+def test_selector_matches_port_filter_is_per_port():
+    """F.8 -- `:port` suffix narrows the match to that specific port on
+    the device; the same device on a different port is NOT matched."""
+    sel = _parse_nic_selector("mlx5_0:2")
+    assert _selector_matches(sel, "mlx5_0", 2) is True
+    assert _selector_matches(sel, "mlx5_0", 1) is False
+
+
+def test_resolve_selector_cli_beats_env():
+    """F.9 -- precedence: when both the CLI flag and NCCL_IB_HCA env are
+    set, the CLI wins. Operators must be able to override their shell
+    env without touching it."""
+    sel = _resolve_selector(
+        allowlist_arg="mlx5_0:1",
+        env={"NCCL_IB_HCA": "rocep158s0:1"},
+    )
+    assert sel["source"] == "cli"
+    assert sel["entries"][0][0] == "mlx5_0"
+
+
+def test_resolve_selector_env_used_when_cli_absent():
+    """F.10 -- precedence: env is consulted only when the CLI flag is
+    unset. This is the most common case in production (operator sets
+    NCCL_IB_HCA once in their job script; smoke picks it up
+    transparently)."""
+    sel = _resolve_selector(
+        allowlist_arg=None,
+        env={"NCCL_IB_HCA": "rocep158s0:1"},
+    )
+    assert sel["source"] == "env"
+
+
+def test_resolve_selector_heuristic_when_both_absent():
+    """F.11 -- precedence fallback: when neither CLI nor env is set,
+    return a `heuristic` marker. The caller (`_collect_nic_status`)
+    handles this by auto-excluding ports whose `phys_state` is in
+    {Disabled, Sleep}."""
+    sel = _resolve_selector(allowlist_arg=None, env={})
+    assert sel["source"] == "heuristic"
+    # Document the heuristic's exact admin-down set so a future widening
+    # (e.g. to also exclude `Polling`) has to update this test
+    # deliberately. `Polling` MUST NOT be in here -- it means
+    # "actively looking for a link partner" which is a real failure on
+    # a port intended to be used.
+    assert sel["admin_down_phys_states"] == ["disabled", "sleep"]
+
+
+def test_resolve_selector_empty_string_falls_through_to_env():
+    """F.12 -- an empty `--rdma-nic-allowlist ''` must not silently
+    block out everything; it should behave the same as not passing the
+    flag at all and fall through to the env. Defensive against
+    shell-quoting accidents (`--rdma-nic-allowlist "$VAR"` when
+    `$VAR` is unset)."""
+    sel = _resolve_selector(
+        allowlist_arg="",
+        env={"NCCL_IB_HCA": "rocep158s0:1"},
+    )
+    assert sel["source"] == "env"
+
+
+def test_resolve_selector_blank_env_falls_through_to_heuristic():
+    """F.13 -- same defensive behavior for an empty NCCL_IB_HCA env
+    (some clusters export it unset / empty by accident). MUST fall
+    through to the heuristic rather than allowlist-nothing."""
+    sel = _resolve_selector(allowlist_arg=None, env={"NCCL_IB_HCA": ""})
+    assert sel["source"] == "heuristic"
+
+
+def test_collect_nic_status_env_allowlist_excludes_disabled_frontend_ports(tmp_path, monkeypatch):
+    """F.14 -- end-to-end against a synthetic /sys/class/infiniband
+    mirroring the production failure mode: 12 IB devices (8 training +
+    2 storage ACTIVE + 2 frontend Disabled). With NCCL_IB_HCA listing
+    only the 8 training NICs, the collector must:
+      * include exactly the 8 listed devices,
+      * exclude the other 4 with a source=`env` info_issue each,
+      * emit ZERO hard issues,
+      * tag the Disabled ports with their state in the info_issue
+        (so an operator who reads excluded_ports can still see they
+        were unhealthy, just not relevant).
+    """
+    from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status
+
+    # Synthetic sysfs tree.
+    base = tmp_path / "ib"
+    devices = [
+        # (name, state, phys_state, rate_str)
+        ("rocep158s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep190s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep206s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep222s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep28s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep62s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep79s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep96s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        # Storage / control-plane: ACTIVE but NOT in NCCL_IB_HCA.
+        ("rocep159s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("rocep29s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        # Front-end: Disabled, NOT in NCCL_IB_HCA -- the user's bug.
+        ("roceo12399", "1: DOWN", "3: Disabled", ""),
+        ("roceo12409", "1: DOWN", "3: Disabled", ""),
+    ]
+
+    for name, state, phys, rate in devices:
+        port_dir = base / name / "ports" / "1"
+        port_dir.mkdir(parents=True)
+        (port_dir / "state").write_text(state)
+        (port_dir / "phys_state").write_text(phys)
+        (port_dir / "rate").write_text(rate)
+        (port_dir / "link_layer").write_text("Ethernet")
+        # One valid RoCE v2 GID on every ACTIVE port so the GID rule
+        # doesn't accidentally fail the test (which would mask the
+        # selector behavior we're trying to verify).
+        gids = port_dir / "gids"
+        gids.mkdir()
+        (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001")
+        types = port_dir / "gid_attrs" / "types"
+        types.mkdir(parents=True)
+        (types / "0").write_text("IB/RoCE v2")
+
+    _install_synthetic_ib_tree(monkeypatch, base)
+
+    monkeypatch.setenv(
+        "NCCL_IB_HCA",
+        "rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1,"
+        "rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1",
+    )
+
+    out = _collect_nic_status(expected_count=None)
+
+    # 12 ports total, 8 included, 4 excluded -- the headline assertion.
+    assert len(out["ports"]) == 12
+    assert len(out["included_ports"]) == 8
+    assert len(out["excluded_ports"]) == 4
+    # The 2 Disabled frontend ports and the 2 not-in-HCA storage ports.
+    assert set(out["excluded_ports"]) == {
+        "rocep159s0:1",
+        "rocep29s0:1",
+        "roceo12399:1",
+        "roceo12409:1",
+    }
+    # Selector metadata records the source so the operator can verify
+    # which precedence layer fired.
+    assert out["selector"]["source"] == "env"
+    # Zero hard issues -> node would PASS.
+    assert out["issues"] == []
+    # Every excluded port produced an info_issue mentioning the env
+    # source. The two Disabled ports MUST also carry their state info
+    # so operators investigating "are my frontend NICs still down?" can
+    # see it without opening every per-port record.
+    info = "\n".join(out["info_issues"])
+    assert "NCCL_IB_HCA" in info
+    assert "phys_state=Disabled" in info  # frontend ports
+    # ACTIVE-but-not-in-HCA ports do NOT need the state suffix (they're
+    # healthy, just reserved for sockets/storage).
+
+
+def test_collect_nic_status_heuristic_only_excludes_disabled_phys_state(tmp_path, monkeypatch):
+    """F.15 -- with no env and no CLI selector, the heuristic must
+    auto-exclude `phys_state=Disabled` ports but MUST keep `phys_state=
+    Polling` (cable unplugged on an intended-up port) and `phys_state=
+    LinkUp with state=INIT` (driver/SM didn't finish bringup) in the
+    included set so they hard-fail. The whole point of choosing Disabled
+    as the exclusion signal is that it's the only phys_state that
+    unambiguously means "admin-down, not used"."""
+    from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status
+
+    base = tmp_path / "ib"
+    devices = [
+        ("training_ok", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"),
+        ("frontend_disabled", "1: DOWN", "3: Disabled", ""),
+        ("training_cable_pulled", "1: DOWN", "2: Polling", ""),
+        ("training_unconfigured", "2: INIT", "5: LinkUp", "400 Gb/sec"),
+    ]
+    for name, state, phys, rate in devices:
+        port_dir = base / name / "ports" / "1"
+        port_dir.mkdir(parents=True)
+        (port_dir / "state").write_text(state)
+        (port_dir / "phys_state").write_text(phys)
+        (port_dir / "rate").write_text(rate)
+        (port_dir / "link_layer").write_text("Ethernet")
+        gids = port_dir / "gids"
+        gids.mkdir()
+        (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001")
+        types = port_dir / "gid_attrs" / "types"
+        types.mkdir(parents=True)
+        (types / "0").write_text("IB/RoCE v2")
+
+    _install_synthetic_ib_tree(monkeypatch, base)
+    monkeypatch.delenv("NCCL_IB_HCA", raising=False)
+
+    out = _collect_nic_status(expected_count=None)
+
+    assert out["selector"]["source"] == "heuristic"
+    # Only the Disabled port is excluded by the heuristic.
+    assert out["excluded_ports"] == ["frontend_disabled:1"]
+    # The 3 remaining are included; the 2 broken-but-included ones produce
+    # hard issues.
+    assert set(out["included_ports"]) == {
+        "training_ok:1",
+        "training_cable_pulled:1",
+        "training_unconfigured:1",
+    }
+    issues_text = " ".join(out["issues"])
+    assert "training_cable_pulled" in issues_text  # state=DOWN
+    assert "training_unconfigured" in issues_text  # state=INIT
+    assert "training_ok" not in issues_text
+
+
+def test_collect_nic_status_empty_set_guard_when_everything_excluded(tmp_path, monkeypatch):
+    """F.16 -- defense in depth: if the selector ends up excluding every
+    discovered port (e.g. operator typo'd --rdma-nic-allowlist, or the
+    whole RoCE card got admin-disabled), the node MUST still hard-fail.
+    A node with zero training NICs cannot participate in inter-node
+    training and silently passing it would be the worst possible
+    regression introduced by the selector feature."""
+    from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status
+
+    base = tmp_path / "ib"
+    port_dir = base / "training_ok" / "ports" / "1"
+    port_dir.mkdir(parents=True)
+    (port_dir / "state").write_text("4: ACTIVE")
+    (port_dir / "phys_state").write_text("5: LinkUp")
+    (port_dir / "rate").write_text("400 Gb/sec")
+    (port_dir / "link_layer").write_text("Ethernet")
+
+    _install_synthetic_ib_tree(monkeypatch, base)
+    # CLI allowlist that matches NOTHING (typo or operator mistake).
+    out = _collect_nic_status(expected_count=None, allowlist="this_device_does_not_exist:1")
+
+    assert out["included_ports"] == []
+    assert out["excluded_ports"] == ["training_ok:1"]
+    # The empty-set guard fires -- node FAIL.
+    assert any("no included RDMA NIC ports" in issue for issue in out["issues"]), out["issues"]
+
+
+def test_collect_nic_status_expected_count_compares_included(tmp_path, monkeypatch):
+    """F.17 -- the behavior change for --expected-rdma-nics: it must
+    compare against the *included* count, not the total /sys/class/
+    infiniband count. Otherwise `--expected-rdma-nics 8` would fail on
+    the very clusters this feature exists to help (12 devices, 8
+    training)."""
+    from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status
+
+    base = tmp_path / "ib"
+    # 8 training + 2 frontend Disabled = 10 total ports, 8 training NICs.
+    for i in range(8):
+        d = base / f"trainnic{i}" / "ports" / "1"
+        d.mkdir(parents=True)
+        (d / "state").write_text("4: ACTIVE")
+        (d / "phys_state").write_text("5: LinkUp")
+        (d / "rate").write_text("400 Gb/sec")
+        (d / "link_layer").write_text("Ethernet")
+        gids = d / "gids"
+        gids.mkdir()
+        (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001")
+        types = d / "gid_attrs" / "types"
+        types.mkdir(parents=True)
+        (types / "0").write_text("IB/RoCE v2")
+    for i in range(2):
+        d = base / f"frontnic{i}" / "ports" / "1"
+        d.mkdir(parents=True)
+        (d / "state").write_text("1: DOWN")
+        (d / "phys_state").write_text("3: Disabled")
+        (d / "rate").write_text("")
+        (d / "link_layer").write_text("Ethernet")
+
+    _install_synthetic_ib_tree(monkeypatch, base)
+    monkeypatch.delenv("NCCL_IB_HCA", raising=False)
+
+    # With --expected-rdma-nics 8 on the heuristic path: the 2 Disabled
+    # ports get auto-excluded, leaving 8 included = matches expected.
+    out = _collect_nic_status(expected_count=8)
+    assert out["issues"] == [], out["issues"]
+    # And the inverse: --expected-rdma-nics 10 (treating the total
+    # `/sys/class/infiniband` count as the expected) MUST fail, because
+    # the comparison happens against the included set.
+    out = _collect_nic_status(expected_count=10)
+    assert any("!= expected 10" in i for i in out["issues"])
+
+
+# ---------------------------------------------------------------------------
+# Primus-cli subcommand wrapper (consolidate-preflight-direct-wrappers).
+# These tests pin the contract for the new `primus-cli node_smoke`
+# subcommand layer: hoisted flags, abbreviation rejection, no Python-side
+# --silent, two-phase dispatch (run on all ranks, aggregate on rank 0),
+# and SLURM-aware aggregator-arg resolution.
+# ---------------------------------------------------------------------------
+
+
+def _build_primus_cli_node_smoke_parser():
+    """Helper: build a parser that mirrors how primus.cli.main wires the
+    `node_smoke` subcommand, without going through the full main()
+    dispatch. Centralized here so individual tests don't redo the wiring.
+    """
+    import argparse as _argparse
+
+    from primus.cli.subcommands.node_smoke import register_subcommand
+
+    top = _argparse.ArgumentParser(prog="primus")
+    sub = top.add_subparsers(dest="command", required=True)
+    register_subcommand(sub)
+    return top
+
+
+def test_primus_cli_node_smoke_known_flags_parse_cleanly():
+    """`primus-cli node_smoke --tier2-perf --hbm-busy-threshold-gib 3.5 ...`
+    must parse without unknown-args; the namespace must merge the run
+    surface with the aggregator surface in a single object."""
+    parser = _build_primus_cli_node_smoke_parser()
+    ns, unknown = parser.parse_known_args(
+        ["node_smoke", "--tier2-perf", "--hbm-busy-threshold-gib", "3.5", "--expected-nodes", "4"]
+    )
+    assert unknown == []
+    assert ns.command == "node_smoke"
+    # Run-side knob
+    assert ns.tier2_perf is True
+    assert ns.hbm_busy_threshold_gib == 3.5
+    # Aggregator-side knob hoisted to the same namespace
+    assert ns.expected_nodes == 4
+    # Run-side default preserved (we only attach --dump-path once via _add_run_flags)
+    assert ns.dump_path == "output/preflight"
+
+
+def test_primus_cli_node_smoke_rejects_silent():
+    """`--silent` is a bash-launcher knob only; the python subcommand
+    must surface it as an unknown argument so a misplaced --silent never
+    silently disappears into argparse."""
+    parser = _build_primus_cli_node_smoke_parser()
+    _ns, unknown = parser.parse_known_args(["node_smoke", "--silent"])
+    assert "--silent" in unknown, (
+        "primus-cli node_smoke must NOT accept --silent (silencing lives "
+        "exclusively in primus-cli-direct.sh's bash layer; this test "
+        "guards against dual handling regressing)"
+    )
+
+
+def test_primus_cli_node_smoke_rejects_abbreviated_tier2():
+    """`--tier2` is an abbreviation of `--tier2-perf` -- argparse's
+    `allow_abbrev=False` must make this a hard error so a stale script
+    using the old flag name doesn't silently match the new one."""
+    import argparse as _argparse
+
+    parser = _build_primus_cli_node_smoke_parser()
+    # parse_known_args with `allow_abbrev=False` returns --tier2 in
+    # `unknown`; the main() dispatcher then rejects unknown args with
+    # SystemExit(2). The subparser-level guarantee we care about here is
+    # that --tier2 NEVER ends up bound to tier2_perf=True via prefix
+    # matching.
+    ns, unknown = parser.parse_known_args(["node_smoke", "--tier2"])
+    assert "--tier2" in unknown
+    assert ns.tier2_perf is False, "--tier2 must NOT silently match --tier2-perf via prefix abbreviation"
+
+    # Defense-in-depth: parse_args (strict mode) raises SystemExit.
+    with pytest.raises(SystemExit):
+        parser.parse_args(["node_smoke", "--tier2"])
+    _ = _argparse  # silence unused
+
+
+def test_resolve_aggregate_args_from_slurm_uses_slurm_nnodes(monkeypatch):
+    """When `--expected-nodes` is not supplied, the helper must fall back
+    to `SLURM_NNODES` (then `SLURM_JOB_NUM_NODES`, then `NNODES`) so the
+    aggregator correctly identifies missing nodes."""
+    import argparse as _argparse
+
+    from primus.tools.preflight.node_smoke.cli import _resolve_aggregate_args_from_slurm
+
+    monkeypatch.setenv("SLURM_NNODES", "6")
+    monkeypatch.delenv("SLURM_JOB_NUM_NODES", raising=False)
+    monkeypatch.delenv("NNODES", raising=False)
+    monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False)
+
+    args = _argparse.Namespace(
+        dump_path="output/preflight",
+        expected_nodes=None,
+        wait_timeout_sec=60,
+        rocm_smi_warn_sec=1.0,
+        clock_skew_warn_sec=30.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+        expected_nodelist_file=None,
+    )
+    agg = _resolve_aggregate_args_from_slurm(args)
+    assert agg.expected_nodes == 6
+    assert agg.dump_path == "output/preflight"
+    # No SLURM_JOB_NODELIST + no pre-supplied file -> stays None (the
+    # aggregator falls back to `<missing-N>` placeholders, by design).
+    assert agg.expected_nodelist_file is None
+
+
+def test_resolve_aggregate_args_from_slurm_cli_wins(monkeypatch):
+    """User-supplied `--expected-nodes` takes precedence over SLURM_*."""
+    import argparse as _argparse
+
+    from primus.tools.preflight.node_smoke.cli import _resolve_aggregate_args_from_slurm
+
+    monkeypatch.setenv("SLURM_NNODES", "6")
+    args = _argparse.Namespace(
+        dump_path="output/preflight",
+        expected_nodes=4,  # user-set wins
+        wait_timeout_sec=60,
+        rocm_smi_warn_sec=1.0,
+        clock_skew_warn_sec=30.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+        expected_nodelist_file=None,
+    )
+    agg = _resolve_aggregate_args_from_slurm(args)
+    assert agg.expected_nodes == 4
+
+
+def test_primus_cli_node_smoke_two_phase_dispatch_rank0(monkeypatch):
+    """On rank 0 the subcommand must call `_cmd_run` AND `_cmd_aggregate`,
+    propagate the max(rc_run, rc_agg) exit code, and pass the SLURM-
+    resolved aggregator args (not the raw `args`) to the aggregator."""
+    import argparse as _argparse
+
+    from primus.cli.subcommands import node_smoke as smoke_cmd
+
+    monkeypatch.setenv("NODE_RANK", "0")
+    monkeypatch.delenv("SLURM_NODEID", raising=False)
+    monkeypatch.setenv("SLURM_NNODES", "5")
+    monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False)
+
+    calls = []
+
+    def fake_run(ns):
+        calls.append(("run", ns))
+        return 0
+
+    def fake_agg(ns):
+        calls.append(("aggregate", ns))
+        return 1  # simulate one MISSING node
+
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", fake_run, raising=True)
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", fake_agg, raising=True)
+
+    args = _argparse.Namespace(
+        dump_path="/tmp/test-smoke",
+        expected_nodes=None,
+        wait_timeout_sec=60,
+        rocm_smi_warn_sec=1.0,
+        clock_skew_warn_sec=30.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+        expected_nodelist_file=None,
+    )
+    with pytest.raises(SystemExit) as exc:
+        smoke_cmd.run(args, extra_args=[])
+    assert exc.value.code == 1, "aggregator's FAIL must surface as the exit code"
+
+    # Run AND aggregate both fired, in that order.
+    assert [c[0] for c in calls] == ["run", "aggregate"]
+    # Aggregate received the SLURM-resolved namespace (expected_nodes
+    # auto-filled from SLURM_NNODES=5), not the user-supplied args.
+    agg_ns = calls[1][1]
+    assert agg_ns.expected_nodes == 5
+    # And the aggregator namespace is a *different* object from `args`.
+    assert agg_ns is not args
+
+
+def test_primus_cli_node_smoke_two_phase_dispatch_non_rank0(monkeypatch):
+    """Non-rank-0 ranks must run `_cmd_run` only and propagate its exit
+    code; calling `_cmd_aggregate` on every rank would race on the
+    shared output directory."""
+    import argparse as _argparse
+
+    from primus.cli.subcommands import node_smoke as smoke_cmd
+
+    monkeypatch.setenv("NODE_RANK", "3")
+    monkeypatch.delenv("SLURM_NODEID", raising=False)
+
+    calls = []
+
+    def fake_run(ns):
+        calls.append("run")
+        return 0
+
+    def fake_agg(ns):  # MUST NOT be called on a non-rank-0 rank.
+        calls.append("aggregate")
+        raise AssertionError("aggregator must not run on non-rank-0")
+
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", fake_run, raising=True)
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", fake_agg, raising=True)
+
+    args = _argparse.Namespace(
+        dump_path="/tmp/test-smoke",
+        expected_nodes=None,
+        wait_timeout_sec=60,
+        rocm_smi_warn_sec=1.0,
+        clock_skew_warn_sec=30.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+        expected_nodelist_file=None,
+    )
+    with pytest.raises(SystemExit) as exc:
+        smoke_cmd.run(args, extra_args=[])
+    assert exc.value.code == 0
+    assert calls == ["run"]
+
+
+def test_primus_cli_node_smoke_rank0_run_failure_surfaces(monkeypatch):
+    """If `_cmd_run` fails on rank 0 but the aggregator (which only knows
+    about missing-from-`<dump>/smoke/*.json`) reports success, the
+    subcommand must still surface the rank-0 run failure -- otherwise a
+    sick rank-0 node could paint itself green via a successful aggregate."""
+    import argparse as _argparse
+
+    from primus.cli.subcommands import node_smoke as smoke_cmd
+
+    monkeypatch.setenv("NODE_RANK", "0")
+    monkeypatch.delenv("SLURM_NODEID", raising=False)
+    monkeypatch.delenv("SLURM_NNODES", raising=False)
+    monkeypatch.delenv("NNODES", raising=False)
+    monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False)
+
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", lambda _ns: 1, raising=True)
+    monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", lambda _ns: 0, raising=True)
+
+    args = _argparse.Namespace(
+        dump_path="/tmp/test-smoke",
+        expected_nodes=None,
+        wait_timeout_sec=60,
+        rocm_smi_warn_sec=1.0,
+        clock_skew_warn_sec=30.0,
+        hbm_busy_threshold_gib=2.0,
+        gpu_activity_warn_pct=20.0,
+        expected_nodelist_file=None,
+    )
+    with pytest.raises(SystemExit) as exc:
+        smoke_cmd.run(args, extra_args=[])
+    assert exc.value.code == 1, "rank-0 _cmd_run=1 must not be masked by _cmd_aggregate=0"
diff --git a/primus/tools/preflight/node_smoke/types.py b/primus/tools/preflight/node_smoke/types.py
new file mode 100644
index 000000000..97f1f4ae5
--- /dev/null
+++ b/primus/tools/preflight/node_smoke/types.py
@@ -0,0 +1,36 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared dataclasses for the node-smoke pipeline."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+
+@dataclass
+class GPUResult:
+    """Result of all checks for a single GPU on this node."""
+
+    gpu: int
+    status: str  # "PASS" | "FAIL" | "TIMEOUT"
+    reason: str = ""
+    duration_sec: float = 0.0
+    details: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class NodeResult:
+    """Whole-node verdict written to ``<dump>/smoke/<host>.json``."""
+
+    host: str
+    node_rank: int
+    status: str  # "PASS" | "FAIL"
+    duration_sec: float
+    fail_reasons: List[str]
+    tier1: Dict[str, Any]
+    tier2: Dict[str, Any]
diff --git a/primus/tools/preflight/preflight_args.py b/primus/tools/preflight/preflight_args.py
index a816b2b47..1710e14b4 100644
--- a/primus/tools/preflight/preflight_args.py
+++ b/primus/tools/preflight/preflight_args.py
@@ -12,18 +12,49 @@
 
 import argparse
 
+# Canonical perf-test tokens accepted by --tests.
+PERF_TEST_TOKENS = (
+    "gemm",
+    "intra-allreduce",
+    "intra-alltoall",
+    "inter-allreduce",
+    "inter-alltoall",
+    "inter-p2p",
+    "inter-ring-p2p",
+)
+
+# Names of the "intent-bearing" perf flags. Setting any of these implies perf
+# mode (no need to also pass --perf-test). When mixed with info selectors
+# (--host/--gpu/--network), perf wins and the info selectors are dropped with
+# a warning.
+PERF_INTENT_FLAGS = ("--perf-test", "--tests", "--quick")
+
 
 def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
     """
     Register arguments for `primus-cli preflight`.
 
+    Mode precedence:
+        1. Any of --perf-test / --tests / --quick is set -> perf mode wins.
+           If info selectors (--host/--gpu/--network) are also set, they are
+           dropped with a warning. Perf-only tuning knobs (--comm-sizes-mb,
+           --intra-group-sizes, etc.) take effect.
+        2. Otherwise, any of --host/--gpu/--network is set -> info-only mode.
+           Perf-only tuning knobs, if set, are inert and a WARN is emitted.
+        3. Otherwise (no flags) -> default: run info AND all perf tests.
+
     Usage:
-        primus-cli preflight                          # Show all info (Host + GPU + Network)
-        primus-cli preflight --host                   # Host only
-        primus-cli preflight --gpu                    # GPU only
-        primus-cli preflight --network                # Network only
-        primus-cli preflight --gpu --network          # GPU + Network
-        primus-cli preflight --perf-test              # Run perf tests ONLY (skip info)
+        primus-cli preflight                          # Default: info + all perf
+        primus-cli preflight --host                   # Host info only
+        primus-cli preflight --gpu                    # GPU info only
+        primus-cli preflight --network                # Network info only
+        primus-cli preflight --gpu --network          # GPU + Network info
+        primus-cli preflight --perf-test              # Perf only, all tests
+        primus-cli preflight --quick                  # Perf only, fast preset
+        primus-cli preflight --tests gemm             # Perf only, GEMM only
+        primus-cli preflight --tests gemm,inter-allreduce \\
+            --comm-sizes-mb 64,1024 \\
+            --inter-group-sizes all
     """
     # Check selection flags
     # Keep --check-* as compatibility aliases.
@@ -54,11 +85,98 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa
         "--perf-test",
         action="store_true",
         help="Run perf tests ONLY (GEMM, intra/inter node communication). "
-        "This is slower and skips the host/gpu/network info report.",
+        "Skips the host/gpu/network info report. Implied by --tests/--quick.",
+    )
+
+    # Performance test specific options.
+    parser.add_argument(
+        "--plot",
+        action="store_true",
+        help="Generate plots (perf mode only).",
+    )
+
+    # Test selection (CSV). Tokens: gemm,intra-allreduce,intra-alltoall,
+    # inter-allreduce,inter-alltoall,inter-p2p,inter-ring-p2p, or 'all'.
+    parser.add_argument(
+        "--tests",
+        type=str,
+        default=None,
+        help="Comma-separated list of perf tests to run. Tokens: "
+        "gemm, intra-allreduce, intra-alltoall, inter-allreduce, inter-alltoall, "
+        "inter-p2p, inter-ring-p2p, all. Implies --perf-test. "
+        "When unset, runs every test.",
+    )
+
+    # Message size config (CSV in MB) for comm tests.
+    parser.add_argument(
+        "--comm-sizes-mb",
+        type=str,
+        default=None,
+        help="Default message sizes (CSV in MB) used for intra-/inter-node "
+        "allreduce, alltoall, and inter-node p2p tests when no specific override is given. "
+        "Default: 2,4,8,16,32,64,128,256,512,1024.",
+    )
+    parser.add_argument(
+        "--intra-comm-sizes-mb",
+        type=str,
+        default=None,
+        help="Override message sizes (CSV in MB) for intra-node allreduce/alltoall. "
+        "Falls back to --comm-sizes-mb when unset.",
+    )
+    parser.add_argument(
+        "--inter-comm-sizes-mb",
+        type=str,
+        default=None,
+        help="Override message sizes (CSV in MB) for inter-node allreduce/alltoall/p2p. "
+        "Falls back to --comm-sizes-mb when unset.",
     )
 
-    # Performance test specific options (only used with --perf-test)
-    parser.add_argument("--plot", action="store_true", help="Generate plots (only with --perf-test)")
+    # Group size config.
+    parser.add_argument(
+        "--intra-group-sizes",
+        type=str,
+        default=None,
+        help="Comma-separated list of intra-node GPU group sizes to test "
+        "(each must divide LOCAL_WORLD_SIZE). Default: 2,4,8.",
+    )
+    parser.add_argument(
+        "--inter-group-sizes",
+        type=str,
+        default=None,
+        help="Comma-separated list of inter-node group sizes to test. Use 'all' for "
+        "the full N-node group. Default: 2,4,all. Note: for inter-node alltoall "
+        "each value is clamped to 16 (real-world MoE training rarely dispatches "
+        "across more nodes); other inter-node tests use the requested sizes "
+        "unchanged. See docs/preflight.md \u00a75.2 for details.",
+    )
+
+    # Inter-node ring p2p sizes.
+    parser.add_argument(
+        "--ring-p2p-sizes-mb",
+        type=str,
+        default=None,
+        help="Message sizes (CSV in MB) for the inter-node ring P2P test. " "Default: 10,20,40,80,160.",
+    )
+
+    # Quick preset.
+    parser.add_argument(
+        "--quick",
+        action="store_true",
+        help="Fast pre-launch preset. Implies --perf-test. Selects gemm + "
+        "intra-allreduce + inter-allreduce, uses sizes 64,1024 MB, full "
+        "intra-node group only, full N-node inter-node group only, and lowers "
+        "warmup/iterations. User-supplied flags override.",
+    )
+
+    # Back-compat alias: kept so existing scripts keep working.
+    # Internally it maps to --inter-group-sizes all and disables inter-p2p.
+    parser.add_argument(
+        "--no-split-nodes-subgroup",
+        dest="split_nodes_subgroup",
+        action="store_false",
+        help="[Deprecated] Skip inter-node comm tests on node subgroups (2-node, 4-node). "
+        "Equivalent to --inter-group-sizes all and dropping inter-p2p.",
+    )
 
     # Distributed init timeout (prevents hangs when network/rendezvous is misconfigured)
     parser.add_argument(
@@ -69,6 +187,24 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa
         "If init times out, preflight will write the info report and exit with failure.",
     )
 
+    # Communicator cleanup delay (cross-rank sync across destroy -> setup;
+    # the actual "Address already in use" defense is the inter-alltoall cap
+    # in --inter-group-sizes / inter_node_comm.py, not this delay).
+    parser.add_argument(
+        "--comm-cleanup-delay-sec",
+        type=float,
+        default=2.0,
+        help="Delay (seconds) inserted between destroying NCCL/RCCL process "
+        "groups and creating new ones. Provides cross-rank synchronization "
+        "across the destroy/setup transition and gives the kernel a moment "
+        "to unlink closed-socket bookkeeping. Set to 0 to disable the sleep "
+        "(barrier only). The actual 'Address already in use' defense at "
+        "scale is the inter-node alltoall sub-group cap of 16 (see "
+        "--inter-group-sizes); widening net.ipv4.ip_local_port_range "
+        "(docs/preflight.md \u00a77.2) is the directly relevant OS knob if "
+        "you ever need more headroom.",
+    )
+
     # Report output options
     parser.add_argument(
         "--dump-path",
@@ -79,8 +215,13 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa
     parser.add_argument(
         "--report-file-name",
         type=str,
-        default="preflight_report",
-        help="Base name for report files (default: preflight_report).",
+        default=None,
+        help=(
+            "Base name for report files. When omitted, an auto-generated "
+            "timestamped name of the form 'preflight-{NNODES}N-{YYYYMMDD-HHMMSS}' "
+            "is used so each run writes to a fresh path and never overwrites "
+            "or collides with stale leftovers."
+        ),
     )
     parser.add_argument(
         "--disable-pdf",
diff --git a/primus/tools/preflight/preflight_perf_test.py b/primus/tools/preflight/preflight_perf_test.py
index e3eb2607a..636dee728 100644
--- a/primus/tools/preflight/preflight_perf_test.py
+++ b/primus/tools/preflight/preflight_perf_test.py
@@ -8,11 +8,20 @@
 import os
 import socket
 import sys
+import time
 from dataclasses import asdict, dataclass
 from datetime import datetime
-from typing import Any, Dict, List
-
-from primus.tools.preflight.global_vars import LOCAL_RANK, RANK, set_hostnames
+from typing import Any, Dict, List, Optional, Set
+
+from primus.tools.preflight.global_vars import (
+    LOCAL_RANK,
+    LOCAL_WORLD_SIZE,
+    RANK,
+    WORLD_SIZE,
+    set_hostnames,
+    set_iteration,
+    set_warmup,
+)
 from primus.tools.preflight.gpu.info import collect_gpu_info, write_gpu_report
 from primus.tools.preflight.host.info import collect_host_info, write_host_report
 from primus.tools.preflight.inter_node_comm import run_inter_node_comm
@@ -23,6 +32,7 @@
     collect_network_info,
     write_network_report,
 )
+from primus.tools.preflight.preflight_args import PERF_TEST_TOKENS
 from primus.tools.preflight.square_gemm import run_square_gemm
 from primus.tools.preflight.utility import (
     gather_hostnames,
@@ -34,6 +44,115 @@
 from primus.tools.utils import gather_records, get_rank_world
 
 
+def _parse_csv_int_list(value: Optional[str], name: str) -> List[int]:
+    """Parse a CSV of positive ints. Empty / None -> []."""
+    if value is None:
+        return []
+    s = str(value).strip()
+    if not s:
+        return []
+    out: List[int] = []
+    for tok in s.split(","):
+        tok = tok.strip()
+        if not tok:
+            continue
+        try:
+            v = int(tok)
+        except ValueError:
+            raise ValueError(f"{name}: '{tok}' is not an integer")
+        if v <= 0:
+            raise ValueError(f"{name}: values must be positive (got {v})")
+        out.append(v)
+    return out
+
+
+def _parse_csv_inter_group_list(value: Optional[str], name: str) -> List[Any]:
+    """Parse a CSV that may include the special token 'all'."""
+    if value is None:
+        return []
+    s = str(value).strip()
+    if not s:
+        return []
+    out: List[Any] = []
+    for tok in s.split(","):
+        tok = tok.strip()
+        if not tok:
+            continue
+        if tok.lower() == "all":
+            out.append("all")
+            continue
+        try:
+            v = int(tok)
+        except ValueError:
+            raise ValueError(f"{name}: '{tok}' is not an integer or 'all'")
+        if v <= 0:
+            raise ValueError(f"{name}: values must be positive (got {v})")
+        out.append(v)
+    return out
+
+
+def _write_node_hostname_legend(markdown_file: str) -> None:
+    """Write a Node -> Hostname legend at the top of the perf markdown report
+    and mirror it to the console on rank 0.
+
+    Uses the per-rank hostnames already gathered via `set_hostnames()`, picking
+    `LOCAL_RANK == 0` of each node as the canonical hostname for that node.
+    """
+    from primus.tools.preflight.global_vars import get_hostnames
+
+    hostnames = get_hostnames()
+    if not hostnames:
+        return
+    num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE
+
+    log("=======Nodes=======")
+    log(f"{'Node':<6} Hostname")
+    with open(markdown_file, "a", encoding="utf-8") as f:
+        f.write("# Nodes\n\n")
+        f.write("| Node | Hostname |\n")
+        f.write("|------|----------|\n")
+        for n in range(num_nodes):
+            rank = n * LOCAL_WORLD_SIZE
+            host = hostnames[rank] if rank < len(hostnames) else ""
+            f.write(f"| {n} | {host} |\n")
+            log(f"{n:<6} {host}")
+        f.write("\n")
+    log("")
+
+
+def _parse_perf_tests(value: Optional[str]) -> Set[str]:
+    """Parse --tests CSV into a set of canonical tokens.
+
+    None / '' / 'all' -> every token. A non-empty value that yields zero valid
+    tokens (e.g. ',,,') raises ValueError so users notice the typo instead of
+    silently running no perf tests.
+    """
+    if value is None:
+        return set(PERF_TEST_TOKENS)
+    raw = str(value)
+    s = raw.strip().lower()
+    if not s or s == "all":
+        return set(PERF_TEST_TOKENS)
+    selected: Set[str] = set()
+    for tok in s.split(","):
+        tok = tok.strip()
+        if not tok:
+            continue
+        if tok == "all":
+            return set(PERF_TEST_TOKENS)
+        if tok not in PERF_TEST_TOKENS:
+            raise ValueError(
+                f"--tests: unknown token '{tok}'. " f"Valid tokens: {', '.join(PERF_TEST_TOKENS)}, all"
+            )
+        selected.add(tok)
+    if not selected:
+        raise ValueError(
+            f"--tests: no valid tokens parsed from {raw!r}. "
+            f"Valid tokens: {', '.join(PERF_TEST_TOKENS)}, all"
+        )
+    return selected
+
+
 @dataclass
 class Finding:
     level: str  # "info" | "warn" | "fail"
@@ -49,6 +168,90 @@ def _status_from_counts(fail_count: int, warn_count: int) -> str:
     return "OK"
 
 
+def _ensure_report_file_name(args: Any) -> str:
+    """Ensure ``args.report_file_name`` is set; auto-generate a timestamped
+    default when the user did not pass ``--report-file-name``.
+
+    The auto-generated form is ``preflight-{NNODES}N-{YYYYMMDD-HHMMSS}``. This
+    guarantees each run writes to a fresh, never-before-used path so the
+    report-path announcement (R5) cannot point at a stale leftover from a
+    previous run.
+
+    Returns the resolved name. Safe to call multiple times: subsequent calls
+    are a no-op once ``args.report_file_name`` is populated.
+    """
+    name = getattr(args, "report_file_name", None)
+    if name:
+        return name
+
+    # Prefer NNODES from the env (set by the runner wrapper before exec) over
+    # deriving from torch's world size, because in info-only mode the
+    # distributed process group is not initialized and world reports 1.
+    nnodes_env = os.environ.get("NNODES")
+    if nnodes_env and nnodes_env.isdigit() and int(nnodes_env) > 0:
+        nnodes = int(nnodes_env)
+    else:
+        try:
+            _rank, world = get_rank_world()
+        except Exception:
+            world = 1
+        try:
+            local_world = int(os.environ.get("LOCAL_WORLD_SIZE") or os.environ.get("GPUS_PER_NODE") or 8)
+        except (TypeError, ValueError):
+            local_world = 8
+        nnodes = max(1, (world or 1) // max(1, local_world))
+
+    name = f"preflight-{nnodes}N-{datetime.now():%Y%m%d-%H%M%S}"
+    try:
+        args.report_file_name = name
+    except (AttributeError, TypeError):
+        # If args is a frozen / read-only container, the caller will see the
+        # returned name; downstream code in this module always uses the
+        # attribute, so the only impact is that the auto-name regenerates on
+        # the next call. Acceptable for the edge case.
+        pass
+    return name
+
+
+def _announce_report_paths(args: Any) -> None:
+    """Print absolute paths of report files on rank 0 (R5).
+
+    Scans ``args.dump_path`` for files matching
+    ``<report_file_name>{,_perf}.{md,pdf}`` and prints absolute paths to
+    stdout. Under bash-side ``--silent`` (primus-cli-direct.sh), fd 1 is
+    ``/dev/null`` by inheritance, so these prints are silenced along with
+    every other stdout write -- acceptable per plan.
+
+    Always rank-0 only. Best-effort: failures here must never mask the
+    preflight exit code.
+    """
+    try:
+        rank, _world = get_rank_world()
+    except Exception:
+        rank = 0
+    if rank != 0:
+        return
+    dump_path = getattr(args, "dump_path", "output/preflight") or "output/preflight"
+    name = getattr(args, "report_file_name", None)
+    if not name:
+        return
+    found_any = False
+    for suffix in ("", "_perf"):
+        for ext in ("md", "pdf"):
+            p = os.path.join(dump_path, f"{name}{suffix}.{ext}")
+            try:
+                if os.path.isfile(p):
+                    print(f"[Primus:Preflight] Report: {os.path.abspath(p)}", flush=True)
+                    found_any = True
+            except OSError:
+                continue
+    if not found_any:
+        print(
+            f"[Primus:Preflight] WARN: no report files found at " f"{dump_path}/{name}{{,_perf}}.{{md,pdf}}",
+            flush=True,
+        )
+
+
 def run_preflight_info(args: Any, expect_distributed: bool = True) -> int:
     """
     Run lightweight preflight info collection (host/gpu/network), aggregate across ranks,
@@ -85,7 +288,10 @@ def run_preflight_info(args: Any, expect_distributed: bool = True) -> int:
         check_host = check_gpu = check_network = True
 
     dump_path = getattr(args, "dump_path", "output/preflight")
-    report_file_name = getattr(args, "report_file_name", "preflight_report")
+    # Defensive: if a caller reached this helper without going through
+    # run_preflight() (e.g. a test fixture), normalize the report name here so
+    # we never write a file literally called "None.md".
+    report_file_name = _ensure_report_file_name(args)
     save_pdf = bool(getattr(args, "save_pdf", True))
 
     findings: List[Finding] = []
@@ -176,15 +382,58 @@ def run_preflight_info(args: Any, expect_distributed: bool = True) -> int:
     return rc
 
 
+def _list_set_perf_tuning_knobs(args) -> List[str]:
+    """Return CLI flag names of perf tuning knobs the user explicitly set.
+
+    Used to warn that these knobs are inert in info-only mode (i.e. when none
+    of --perf-test/--tests/--quick are set but at least one of
+    --host/--gpu/--network is).
+    """
+    set_flags: List[str] = []
+    if getattr(args, "comm_sizes_mb", None) is not None:
+        set_flags.append("--comm-sizes-mb")
+    if getattr(args, "intra_comm_sizes_mb", None) is not None:
+        set_flags.append("--intra-comm-sizes-mb")
+    if getattr(args, "inter_comm_sizes_mb", None) is not None:
+        set_flags.append("--inter-comm-sizes-mb")
+    if getattr(args, "intra_group_sizes", None) is not None:
+        set_flags.append("--intra-group-sizes")
+    if getattr(args, "inter_group_sizes", None) is not None:
+        set_flags.append("--inter-group-sizes")
+    if getattr(args, "ring_p2p_sizes_mb", None) is not None:
+        set_flags.append("--ring-p2p-sizes-mb")
+    if getattr(args, "plot", False):
+        set_flags.append("--plot")
+    if not getattr(args, "split_nodes_subgroup", True):
+        set_flags.append("--no-split-nodes-subgroup")
+    return set_flags
+
+
 def run_preflight(args):
     """
     Preflight entry point with dispatch logic.
 
-    - If any of --host/--gpu/--network is set → show only selected info sections
-    - If no selection flags are set (plain `preflight`) → run ALL: info + perf tests
-    - If --perf-test is set → run perf tests ONLY (skip info)
+    Mode precedence (single rule):
+
+      1. Any of --perf-test / --tests / --quick is set -> perf-only mode.
+         If info selectors (--host/--gpu/--network) are also present, they
+         are dropped with a WARN.
+      2. Otherwise, any of --host/--gpu/--network is set -> info-only mode.
+         Perf tuning knobs (e.g. --comm-sizes-mb), if set, are inert and a
+         WARN is emitted.
+      3. Otherwise (no flags) -> default: info AND all perf tests.
     """
-    perf_test = getattr(args, "perf_test", False)
+    # R4: canonical normalization point for the report file name. Done here
+    # before any downstream code reads args.report_file_name, so every code
+    # path (info-only, perf-only, info+perf, dist-init failure) sees the same
+    # value -- and so the dist-init failure paths below can interpolate it
+    # safely without a None-guard.
+    _ensure_report_file_name(args)
+
+    perf_test = bool(getattr(args, "perf_test", False))
+    tests_set = getattr(args, "tests", None) is not None
+    quick_set = bool(getattr(args, "quick", False))
+    has_perf_intent = perf_test or tests_set or quick_set
 
     def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Exception) -> None:
         try:
@@ -209,11 +458,50 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
     # If any selection flags are set, only run info collection/report.
     # IMPORTANT: do NOT initialize torch.distributed for info-only mode; preflight must not hang
     # when networking/rendezvous is misconfigured.
-    check_host = getattr(args, "check_host", False)
-    check_gpu = getattr(args, "check_gpu", False)
-    check_network = getattr(args, "check_network", False)
+    check_host = bool(getattr(args, "check_host", False))
+    check_gpu = bool(getattr(args, "check_gpu", False))
+    check_network = bool(getattr(args, "check_network", False))
     any_selection = bool(check_host or check_gpu or check_network)
 
+    # Precedence: perf-mode wins over info selectors.
+    info_dropped_warning: Optional[str] = None
+    if has_perf_intent:
+        if any_selection:
+            dropped = [
+                flag
+                for present, flag in (
+                    (check_host, "--host"),
+                    (check_gpu, "--gpu"),
+                    (check_network, "--network"),
+                )
+                if present
+            ]
+            info_dropped_warning = (
+                f"info selectors {','.join(dropped)} were dropped because perf "
+                f"mode (--perf-test/--tests/--quick) takes precedence. "
+                f"Run them in a separate invocation if you want both reports."
+            )
+            print(f"[Primus:Preflight] WARN: {info_dropped_warning}", file=sys.stderr)
+            check_host = check_gpu = check_network = False
+            any_selection = False
+            # Reflect the override on `args` so info-section helpers downstream
+            # (e.g. run_preflight_info getattr) see consistent state.
+            args.check_host = False
+            args.check_gpu = False
+            args.check_network = False
+        # Auto-imply --perf-test so the rest of the dispatch treats this as perf-only.
+        perf_test = True
+        args.perf_test = True
+    elif any_selection:
+        # Info-only mode: tuning knobs are inert. Emit a single WARN listing them.
+        inert = _list_set_perf_tuning_knobs(args)
+        if inert:
+            print(
+                f"[Primus:Preflight] WARN: {','.join(inert)} have no effect in "
+                f"info-only mode (no --perf-test/--tests/--quick).",
+                file=sys.stderr,
+            )
+
     # 1) Info-only mode: run without distributed init.
     if not perf_test and any_selection:
         # First, emit a local-only report immediately (so user gets output even if PG init hangs).
@@ -231,20 +519,23 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
             if world > 1:
                 init_distributed(timeout=timedelta(seconds=dist_timeout_sec))
                 try:
-                    return run_preflight_info(args)
+                    rc = run_preflight_info(args)
+                    _announce_report_paths(args)
+                    return rc
                 finally:
                     finalize_distributed()
         except Exception as e:
             if rank == 0:
                 dump_path = getattr(args, "dump_path", "output/preflight")
-                report_file_name = getattr(args, "report_file_name", "preflight_report")
                 os.makedirs(dump_path, exist_ok=True)
-                markdown_file = f"{dump_path}/{report_file_name}.md"
+                markdown_file = f"{dump_path}/{args.report_file_name}.md"
                 _append_dist_init_failure(markdown_file, dist_timeout_sec, e)
             print(f"[Primus:Preflight] ERROR: distributed init failed: {e}", file=sys.stderr)
+            _announce_report_paths(args)
             return 2
 
         # world==1 fallback
+        _announce_report_paths(args)
         return local_rc
 
     # 2) Plain `preflight` (no flags): run info FIRST (no dist init) so we always get a report.
@@ -252,6 +543,20 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
     if not perf_test and not any_selection:
         info_rc = run_preflight_info(args, expect_distributed=False)
 
+    # 2.5) Resolve perf config NOW, before any distributed rendezvous, so that
+    # invalid CLI input (typos, bad sizes/group-sizes for selected tests) fails
+    # fast instead of after a 120s NCCL init. This is also where we apply the
+    # `--quick` warmup/iteration overrides on every rank.
+    try:
+        perf_cfg = _resolve_perf_config(args)
+    except ValueError as e:
+        print(f"[Primus:Preflight] ERROR: invalid perf config: {e}", file=sys.stderr)
+        return 2
+    if perf_cfg["warmup"] is not None:
+        set_warmup(perf_cfg["warmup"])
+    if perf_cfg["iteration"] is not None:
+        set_iteration(perf_cfg["iteration"])
+
     # 3) Perf tests (perf-only OR plain preflight after info): now attempt distributed init
     # with a timeout so we fail fast instead of hanging.
     from datetime import timedelta
@@ -265,16 +570,16 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
         # We already wrote the info report in plain preflight mode; append a clear failure note.
         rank, _world = get_rank_world()
         dump_path = getattr(args, "dump_path", "output/preflight")
-        report_file_name = getattr(args, "report_file_name", "preflight_report")
         if rank == 0:
             try:
                 os.makedirs(dump_path, exist_ok=True)
-                markdown_file = f"{dump_path}/{report_file_name}.md"
+                markdown_file = f"{dump_path}/{args.report_file_name}.md"
                 _append_dist_init_failure(markdown_file, dist_timeout_sec, e)
             except Exception as ee:
                 print(f"[Primus:Preflight] [rank0] WARN: failed to write report: {ee}", file=sys.stderr)
 
         print(f"[Primus:Preflight] ERROR: distributed init failed: {e}", file=sys.stderr)
+        _announce_report_paths(args)
         return 2
 
     try:
@@ -311,12 +616,70 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
         args.pdf_file = f"{args.dump_path}/{args.report_file_name}{perf_suffix}.pdf"
         remove_file(args.markdown_file)
 
-        # run tests
-        run_square_gemm(args)
-        run_intra_node_comm(args)
-        run_inter_node_comm(args)
-        run_inter_node_comm_p2p(args)
-        run_inter_node_ring_p2p(args)
+        # Write a Node -> Hostname legend at the top of the perf report so
+        # subsequent tables can use compact (and possibly truncated) hostname
+        # representations without losing the host<->node mapping.
+        if RANK == 0:
+            if info_dropped_warning:
+                with open(args.markdown_file, "a", encoding="utf-8") as f:
+                    f.write(f"> Note: {info_dropped_warning}\n\n")
+            _write_node_hostname_legend(args.markdown_file)
+
+        if RANK == 0:
+            log(
+                f"[Primus:Preflight] perf tests selected: "
+                f"{','.join(sorted(perf_cfg['enabled_tests'])) or '(none)'}"
+            )
+
+        # ------------------------------------------------------------------
+        # Dispatch tests by token, with per-test wall-clock logging.
+        # ------------------------------------------------------------------
+        enabled = perf_cfg["enabled_tests"]
+        intra_comms = {c for c in ("allreduce", "alltoall") if f"intra-{c}" in enabled}
+        inter_comms = {c for c in ("allreduce", "alltoall") if f"inter-{c}" in enabled}
+
+        def _timed(name: str, fn):
+            t0 = time.time()
+            fn()
+            if RANK == 0:
+                log(f"[Primus:Preflight] {name} done in {time.time() - t0:.1f}s")
+
+        if "gemm" in enabled:
+            _timed("gemm", lambda: run_square_gemm(args))
+
+        if intra_comms:
+            _timed(
+                f"intra-comm ({','.join(sorted(intra_comms))})",
+                lambda: run_intra_node_comm(
+                    args,
+                    enabled_comms=intra_comms,
+                    sizes_mb=perf_cfg["intra_sizes_mb"],
+                    group_sizes=perf_cfg["intra_group_sizes"],
+                ),
+            )
+
+        if inter_comms:
+            _timed(
+                f"inter-comm ({','.join(sorted(inter_comms))})",
+                lambda: run_inter_node_comm(
+                    args,
+                    enabled_comms=inter_comms,
+                    sizes_mb=perf_cfg["inter_sizes_mb"],
+                    group_sizes=perf_cfg["inter_group_sizes"],
+                ),
+            )
+
+        if "inter-p2p" in enabled:
+            _timed(
+                "inter-p2p",
+                lambda: run_inter_node_comm_p2p(args, sizes_mb=perf_cfg["inter_sizes_mb"]),
+            )
+
+        if "inter-ring-p2p" in enabled:
+            _timed(
+                "inter-ring-p2p",
+                lambda: run_inter_node_ring_p2p(args, sizes_mb=perf_cfg["ring_sizes_mb"]),
+            )
 
         if RANK == 0 and args.save_pdf:
             md_to_pdf(args.markdown_file, args.pdf_file)
@@ -327,9 +690,123 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti
             return info_rc
         return 0
     finally:
+        _announce_report_paths(args)
         finalize_distributed()
 
 
+# Built-in defaults for tuning knobs. These mirror the documented defaults in
+# the parser; we keep argparse's `default=None` so that "user did not pass it"
+# is detectable, and substitute these defaults when resolving the config.
+_DEFAULT_COMM_SIZES_MB = "2,4,8,16,32,64,128,256,512,1024"
+_DEFAULT_INTRA_GROUP_SIZES = "2,4,8"
+_DEFAULT_INTER_GROUP_SIZES = "2,4,all"
+_DEFAULT_RING_P2P_SIZES_MB = "10,20,40,80,160"
+
+
+def _resolve_perf_config(args) -> Dict[str, Any]:
+    """Resolve perf-test selection, sizes, group sizes, and quick-preset overrides.
+
+    User-supplied values take precedence over `--quick` defaults. Validates and
+    raises `ValueError` on bad input.
+
+    This function is side-effect free: it returns `warmup` / `iteration` for
+    the caller to apply (e.g. via `set_warmup` / `set_iteration`) instead of
+    mutating module globals itself.
+
+    Validation of intra-/inter-/ring-specific knobs is gated by which tests
+    are actually selected, so e.g. `--tests gemm --intra-group-sizes 3` does
+    NOT abort on a host with LOCAL_WORLD_SIZE=8.
+    """
+    # All perf-related CLI args have argparse default=None, so a non-None value
+    # unambiguously means "the user passed this flag".
+    tests_str = getattr(args, "tests", None)
+    comm_sizes_str = getattr(args, "comm_sizes_mb", None)
+    intra_comm_sizes_str = getattr(args, "intra_comm_sizes_mb", None)
+    inter_comm_sizes_str = getattr(args, "inter_comm_sizes_mb", None)
+    intra_group_sizes_str = getattr(args, "intra_group_sizes", None)
+    inter_group_sizes_str = getattr(args, "inter_group_sizes", None)
+    ring_sizes_str = getattr(args, "ring_p2p_sizes_mb", None)
+    quick = bool(getattr(args, "quick", False))
+    split_nodes_subgroup = bool(getattr(args, "split_nodes_subgroup", True))
+
+    # Apply --quick defaults only where the user did not override (i.e. value is None).
+    warmup: Optional[int] = None
+    iteration: Optional[int] = None
+    if quick:
+        if tests_str is None:
+            tests_str = "gemm,intra-allreduce,inter-allreduce"
+        if comm_sizes_str is None:
+            comm_sizes_str = "64,1024"
+        if intra_group_sizes_str is None:
+            intra_group_sizes_str = str(LOCAL_WORLD_SIZE)
+        if inter_group_sizes_str is None:
+            inter_group_sizes_str = "all"
+        # Lower warmup / iterations for a fast go/no-go signal. Caller applies.
+        warmup = 5
+        iteration = 20
+
+    # Substitute built-in defaults for any values still unset.
+    if comm_sizes_str is None:
+        comm_sizes_str = _DEFAULT_COMM_SIZES_MB
+    if intra_group_sizes_str is None:
+        intra_group_sizes_str = _DEFAULT_INTRA_GROUP_SIZES
+    if inter_group_sizes_str is None:
+        inter_group_sizes_str = _DEFAULT_INTER_GROUP_SIZES
+    if ring_sizes_str is None:
+        ring_sizes_str = _DEFAULT_RING_P2P_SIZES_MB
+
+    enabled_tests = _parse_perf_tests(tests_str)
+
+    # Honor the legacy --no-split-nodes-subgroup alias: drop subgroup sizes
+    # AND skip the inter-p2p test (matching the prior behavior of the script).
+    if not split_nodes_subgroup:
+        inter_group_sizes_str = "all"
+        enabled_tests.discard("inter-p2p")
+
+    needs_intra = any(t in enabled_tests for t in ("intra-allreduce", "intra-alltoall"))
+    needs_inter_coll = any(t in enabled_tests for t in ("inter-allreduce", "inter-alltoall"))
+    needs_inter_p2p = "inter-p2p" in enabled_tests
+    needs_inter = needs_inter_coll or needs_inter_p2p
+    needs_ring = "inter-ring-p2p" in enabled_tests
+
+    default_sizes = _parse_csv_int_list(comm_sizes_str, "--comm-sizes-mb")
+    intra_sizes_override = _parse_csv_int_list(intra_comm_sizes_str, "--intra-comm-sizes-mb")
+    inter_sizes_override = _parse_csv_int_list(inter_comm_sizes_str, "--inter-comm-sizes-mb")
+    intra_sizes_mb = intra_sizes_override or default_sizes
+    inter_sizes_mb = inter_sizes_override or default_sizes
+    if needs_intra and not intra_sizes_mb:
+        raise ValueError("--comm-sizes-mb / --intra-comm-sizes-mb yielded no sizes")
+    if needs_inter and not inter_sizes_mb:
+        raise ValueError("--comm-sizes-mb / --inter-comm-sizes-mb yielded no sizes")
+
+    intra_group_sizes = _parse_csv_int_list(intra_group_sizes_str, "--intra-group-sizes")
+    if needs_intra:
+        if not intra_group_sizes:
+            raise ValueError("--intra-group-sizes yielded no values")
+        bad = [g for g in intra_group_sizes if LOCAL_WORLD_SIZE % g != 0]
+        if bad:
+            raise ValueError(f"--intra-group-sizes: {bad} do not divide LOCAL_WORLD_SIZE={LOCAL_WORLD_SIZE}")
+
+    inter_group_sizes = _parse_csv_inter_group_list(inter_group_sizes_str, "--inter-group-sizes")
+    if needs_inter and not inter_group_sizes:
+        raise ValueError("--inter-group-sizes yielded no values")
+
+    ring_sizes_mb = _parse_csv_int_list(ring_sizes_str, "--ring-p2p-sizes-mb")
+    if needs_ring and not ring_sizes_mb:
+        raise ValueError("--ring-p2p-sizes-mb yielded no sizes")
+
+    return {
+        "enabled_tests": enabled_tests,
+        "intra_sizes_mb": intra_sizes_mb,
+        "inter_sizes_mb": inter_sizes_mb,
+        "intra_group_sizes": intra_group_sizes,
+        "inter_group_sizes": inter_group_sizes,
+        "ring_sizes_mb": ring_sizes_mb,
+        "warmup": warmup,
+        "iteration": iteration,
+    }
+
+
 def main():
     parser = argparse.ArgumentParser()
     from primus.tools.preflight.preflight_args import add_preflight_parser
diff --git a/primus/tools/preflight/square_gemm.py b/primus/tools/preflight/square_gemm.py
index daa592062..c947ba121 100644
--- a/primus/tools/preflight/square_gemm.py
+++ b/primus/tools/preflight/square_gemm.py
@@ -6,18 +6,17 @@
 
 import time
 
-import matplotlib.pyplot as plt
 import torch
 import torch.distributed as dist
 
 from primus.tools.preflight.global_vars import (
-    ITERATION,
     LOCAL_RANK,
     LOCAL_WORLD_SIZE,
     RANK,
-    WARMUP,
     WORLD_SIZE,
     get_hostnames,
+    get_iteration,
+    get_warmup,
 )
 from primus.tools.preflight.utility import create_dir, log
 
@@ -26,19 +25,21 @@ def run_square_gemm(args):
     sizes = [1024, 2048, 4096, 8192, 10240]
     latency_results = {}
     flops_results = {}
+    warmup = get_warmup()
+    iteration = get_iteration()
     for size in sizes:
         a = torch.randn((size, size), device=f"cuda:{LOCAL_RANK}", dtype=torch.bfloat16)
         b = torch.randn((size, size), device=f"cuda:{LOCAL_RANK}", dtype=torch.bfloat16)
         torch.cuda.synchronize()
-        for _ in range(WARMUP):
+        for _ in range(warmup):
             torch.matmul(a, b)
         torch.cuda.synchronize()
         start = time.time()
-        for _ in range(ITERATION):
+        for _ in range(iteration):
             torch.matmul(a, b)
         torch.cuda.synchronize()
         end = time.time()
-        t = (end - start) / ITERATION
+        t = (end - start) / iteration
         tflops = 2 * size * size * size / (t * 1e12)
         latency_results[f"{size}x{size}x{size}"] = t
         flops_results[f"{size}x{size}x{size}"] = tflops
@@ -87,6 +88,8 @@ def run_square_gemm(args):
         if not args.plot:
             return
 
+        import matplotlib.pyplot as plt
+
         log("=======Plot Square GEMM TFLOPS=======")
         with open(args.markdown_file, "a", encoding="utf-8") as f:
             f.write(f"=======Plot Square GEMM TFLOPS=======\n")
diff --git a/primus/tools/preflight/tests/test_report_naming.py b/primus/tools/preflight/tests/test_report_naming.py
new file mode 100644
index 000000000..13f78ccb2
--- /dev/null
+++ b/primus/tools/preflight/tests/test_report_naming.py
@@ -0,0 +1,238 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for ``preflight_perf_test._ensure_report_file_name`` /
+``_announce_report_paths`` (consolidate-preflight-direct-wrappers, R5).
+
+These tests pin the contract of the Python-side replacements for the
+old bash ``run_preflight_direct.sh`` behavior:
+
+* Every preflight run that doesn't get an explicit ``--report-file-name``
+  must write to a unique, timestamped path. Stale reports from a previous
+  run must NEVER be mistaken for the current run's output.
+* Once reports exist, rank 0 must print their absolute paths so an
+  operator running under ``primus-cli direct`` can find them without
+  guessing the dump-path layout.
+* Best-effort: the announcement must never raise, even when nothing was
+  written (the per-run output may have failed earlier).
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+from datetime import datetime, timedelta
+from typing import Optional
+
+from primus.tools.preflight.preflight_args import add_preflight_parser
+from primus.tools.preflight.preflight_perf_test import (
+    _announce_report_paths,
+    _ensure_report_file_name,
+)
+
+
+def _build_preflight_parser() -> argparse.ArgumentParser:
+    """Construct the same parser that `primus-cli preflight` registers,
+    isolated from the broader primus.cli wiring so tests can exercise
+    just the argparse contract.
+    """
+    p = argparse.ArgumentParser(prog="preflight-test")
+    add_preflight_parser(p)
+    return p
+
+
+# ---------------------------------------------------------------------------
+# `--report-file-name` default contract
+# ---------------------------------------------------------------------------
+
+
+def test_preflight_args_report_file_name_default_is_none():
+    """The argparse default for ``--report-file-name`` must be ``None`` so
+    ``_ensure_report_file_name`` can detect "user did not pass it" and
+    auto-generate a unique name. A non-None default would silently
+    suppress the auto-naming logic and re-introduce the "every run
+    overwrites preflight_report.md" footgun the plan removed."""
+    parser = _build_preflight_parser()
+    ns = parser.parse_args([])
+    assert ns.report_file_name is None, (
+        "argparse default must stay None -- _ensure_report_file_name "
+        "uses falsy check to decide whether to auto-generate"
+    )
+
+
+# ---------------------------------------------------------------------------
+# `_ensure_report_file_name`
+# ---------------------------------------------------------------------------
+
+
+_AUTO_NAME_RE = re.compile(r"^preflight-(\d+)N-(\d{8})-(\d{6})$")
+
+
+def test_ensure_report_file_name_user_value_wins(monkeypatch):
+    """When the user passed ``--report-file-name=foo`` it must survive
+    untouched -- the auto-name path must never overwrite an explicit
+    value."""
+    monkeypatch.setenv("NNODES", "32")  # would otherwise affect auto-name
+    ns = argparse.Namespace(report_file_name="my-custom-name")
+    out = _ensure_report_file_name(ns)
+    assert out == "my-custom-name"
+    assert ns.report_file_name == "my-custom-name"
+
+
+def test_ensure_report_file_name_auto_uses_nnodes_env(monkeypatch):
+    """When ``NNODES`` is set, the auto-name must use it directly. This
+    is the critical case for info-only mode where no process group is
+    initialized and ``get_rank_world()`` would return world=1."""
+    monkeypatch.setenv("NNODES", "42")
+    ns = argparse.Namespace(report_file_name=None)
+    out = _ensure_report_file_name(ns)
+    m = _AUTO_NAME_RE.match(out)
+    assert m is not None, f"auto-name {out!r} doesn't match expected pattern"
+    assert m.group(1) == "42"
+    assert ns.report_file_name == out
+
+
+def test_ensure_report_file_name_idempotent(monkeypatch):
+    """Calling ``_ensure_report_file_name`` twice in the same run must
+    return the same name -- otherwise different code paths (run /
+    info / announce) would each compute their own timestamped name and
+    write to different files."""
+    monkeypatch.setenv("NNODES", "8")
+    ns = argparse.Namespace(report_file_name=None)
+    first = _ensure_report_file_name(ns)
+    # Sleep is unnecessary -- if the function recomputed it would use
+    # `datetime.now()` again, which is enough to diverge at second
+    # granularity over a slow run. We test the stronger invariant: the
+    # second call returns the stored value verbatim, without recomputing.
+    monkeypatch.setattr(
+        "primus.tools.preflight.preflight_perf_test.datetime",
+        _ExplodingDatetime,  # any access to datetime would raise
+    )
+    second = _ensure_report_file_name(ns)
+    assert second == first
+
+
+class _ExplodingDatetime:
+    """datetime stub that raises on any use -- proves a method doesn't
+    reach the auto-name regeneration path."""
+
+    @classmethod
+    def now(cls, *_a, **_kw):
+        raise AssertionError(
+            "_ensure_report_file_name regenerated the timestamp on " "the second call (it must be idempotent)"
+        )
+
+
+def test_ensure_report_file_name_unique_per_run(monkeypatch):
+    """Two independent invocations (different `args` objects with
+    `report_file_name=None`) at different timestamps must produce
+    distinct names -- this is the guarantee that protects against
+    stale-report aliasing under back-to-back runs."""
+    monkeypatch.setenv("NNODES", "4")
+
+    class _FrozenDatetime:
+        _now = datetime(2026, 5, 7, 12, 0, 0)
+
+        @classmethod
+        def now(cls, *_a, **_kw):
+            return cls._now
+
+    monkeypatch.setattr("primus.tools.preflight.preflight_perf_test.datetime", _FrozenDatetime)
+    ns_a = argparse.Namespace(report_file_name=None)
+    name_a = _ensure_report_file_name(ns_a)
+    # Bump the clock by one second to simulate the next run.
+    _FrozenDatetime._now = _FrozenDatetime._now + timedelta(seconds=1)
+    ns_b = argparse.Namespace(report_file_name=None)
+    name_b = _ensure_report_file_name(ns_b)
+
+    assert name_a != name_b
+    assert name_a.endswith("12-00-00".replace("-", ""))
+    assert name_b.endswith("12-00-01".replace("-", ""))
+
+
+def test_ensure_report_file_name_handles_invalid_nnodes(monkeypatch):
+    """A malformed ``NNODES`` (non-digit) must NOT crash the auto-name
+    helper; it must fall back to deriving from torch's world size, with a
+    safe minimum of 1."""
+    monkeypatch.setenv("NNODES", "")  # explicitly empty
+    monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False)
+    monkeypatch.delenv("GPUS_PER_NODE", raising=False)
+    ns = argparse.Namespace(report_file_name=None)
+    out = _ensure_report_file_name(ns)
+    m = _AUTO_NAME_RE.match(out)
+    assert m is not None, f"auto-name {out!r} doesn't match expected pattern"
+    # When world isn't initialized we end up with nnodes=max(1, 1//8)=1.
+    assert int(m.group(1)) >= 1
+
+
+# ---------------------------------------------------------------------------
+# `_announce_report_paths`
+# ---------------------------------------------------------------------------
+
+
+def _make_args(dump_path: str, name: Optional[str]) -> argparse.Namespace:
+    return argparse.Namespace(dump_path=dump_path, report_file_name=name)
+
+
+def test_announce_report_paths_finds_existing_md(tmp_path, capsys, monkeypatch):
+    """Smoke: ``_announce_report_paths`` must print absolute paths for
+    every report variant that exists on disk."""
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    (tmp_path / "preflight-2N-20260507-120000.md").write_text("# report")
+    (tmp_path / "preflight-2N-20260507-120000_perf.md").write_text("# perf")
+
+    args = _make_args(str(tmp_path), "preflight-2N-20260507-120000")
+    _announce_report_paths(args)
+
+    out = capsys.readouterr().out
+    assert "preflight-2N-20260507-120000.md" in out
+    assert "preflight-2N-20260507-120000_perf.md" in out
+    # Paths must be absolute -- operators copy them into srun nodelists
+    # / scp commands; a relative path here is a footgun under SLURM
+    # where the cwd doesn't match across nodes.
+    for line in out.splitlines():
+        if line.startswith("[Primus:Preflight] Report:"):
+            path = line.split("Report:", 1)[1].strip()
+            assert os.path.isabs(path), f"non-absolute report path: {path!r}"
+
+
+def test_announce_report_paths_warns_when_no_files(tmp_path, capsys, monkeypatch):
+    """When the helper is called but no report files exist (e.g. preflight
+    crashed before writing), it must emit a WARN line so the operator
+    isn't left wondering where the report went -- but it must NOT raise."""
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    args = _make_args(str(tmp_path), "ghost-name")
+    _announce_report_paths(args)  # must not raise
+    out = capsys.readouterr().out
+    assert "WARN" in out
+    assert "ghost-name" in out
+
+
+def test_announce_report_paths_skips_non_rank0(tmp_path, capsys, monkeypatch):
+    """Non-rank-0 ranks must stay silent -- otherwise every node would
+    print the same paths and clutter the launcher output."""
+    monkeypatch.setenv("RANK", "3")
+    monkeypatch.setenv("WORLD_SIZE", "8")
+    (tmp_path / "preflight-1N-20260507-120000.md").write_text("# report")
+    args = _make_args(str(tmp_path), "preflight-1N-20260507-120000")
+    _announce_report_paths(args)
+    out = capsys.readouterr().out
+    assert out == "", f"non-rank-0 announce produced output: {out!r}"
+
+
+def test_announce_report_paths_no_name_is_noop(tmp_path, capsys, monkeypatch):
+    """Defensive: if ``args.report_file_name`` is somehow still None when
+    we reach the announcement (e.g. an early bail-out before
+    ``_ensure_report_file_name`` was called), the helper must be a
+    silent no-op rather than crash on the format-string path build."""
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    args = _make_args(str(tmp_path), None)
+    _announce_report_paths(args)  # must not raise
+    assert capsys.readouterr().out == ""
diff --git a/primus/tools/preflight/utility.py b/primus/tools/preflight/utility.py
index 9c71d6335..e27ec1a73 100644
--- a/primus/tools/preflight/utility.py
+++ b/primus/tools/preflight/utility.py
@@ -6,15 +6,17 @@
 
 import os
 import socket
+import time
 from pathlib import Path
+from typing import Iterable, List
 
-import markdown2
 import torch
 import torch.distributed as dist
-from weasyprint import HTML
 
 from primus.tools.preflight.global_vars import RANK, WORLD_SIZE
 
+DEFAULT_COMM_CLEANUP_DELAY_SEC = 2.0
+
 
 def log(msg):
     if RANK == 0:
@@ -48,6 +50,23 @@ def gather_hostnames():
         return None
 
 
+def barrier_after_comm_destroy(delay_sec: float = DEFAULT_COMM_CLEANUP_DELAY_SEC):
+    """Global barrier + sleep after destroying NCCL/RCCL process groups.
+
+    When communicators are destroyed and recreated quickly (e.g. in test loops),
+    the underlying sockets may still be in TIME_WAIT or the remote side may not
+    have fully cleaned up.  This causes "Address already in use" errors on bind.
+
+    This function ensures all ranks have completed destruction (barrier) and then
+    waits a short period for the OS to release ports before the next group creation.
+    """
+    log(f"  barrier_after_comm_destroy: barrier synced, sleeping {delay_sec:.1f}s")
+    dist.barrier(device_ids=[torch.cuda.current_device()])
+    if delay_sec > 0:
+        time.sleep(delay_sec)
+    log(f"  barrier_after_comm_destroy: done (delay={delay_sec:.1f}s)")
+
+
 def remove_file(file_path):
     if RANK == 0:
         if os.path.exists(file_path):
@@ -56,6 +75,32 @@ def remove_file(file_path):
     dist.barrier(device_ids=[torch.cuda.current_device()])
 
 
+def format_int_range(values: Iterable[int]) -> str:
+    """Compactly format a list of ints as a comma-separated set of contiguous ranges.
+
+    Examples:
+        [0,1,2,3]         -> "0-3"
+        [0,1,2,3,8,9]     -> "0-3,8-9"
+        [0,2,4,6]         -> "0,2,4,6"
+        [5]               -> "5"
+        []                -> ""
+    """
+    vs = sorted(set(int(v) for v in values))
+    if not vs:
+        return ""
+
+    parts: List[str] = []
+    start = prev = vs[0]
+    for v in vs[1:]:
+        if v == prev + 1:
+            prev = v
+            continue
+        parts.append(f"{start}-{prev}" if start != prev else f"{start}")
+        start = prev = v
+    parts.append(f"{start}-{prev}" if start != prev else f"{start}")
+    return ",".join(parts)
+
+
 def extract_first_middle_last(lst):
     if not lst:
         return []
@@ -70,6 +115,9 @@ def extract_first_middle_last(lst):
 
 
 def md_to_pdf(md_path, pdf_path):
+    import markdown2
+    from weasyprint import HTML
+
     with open(md_path, "r", encoding="utf-8") as f:
         markdown_text = f.read()
 
diff --git a/pyproject.toml b/pyproject.toml
index 19880d9a4..4851a3253 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,6 +73,15 @@ Repository = "https://github.com/AMD-AGI/Primus"
 line-length = 110
 target-version = ["py38"]
 
+[tool.isort]
+# Mirror the pre-commit hook (isort --profile black) so local runs and the
+# pure-checkout CI agree. `odc` lives under primus/core/odc/odc/, so a clean
+# checkout has no top-level `odc` package and isort would otherwise treat
+# `import odc` as third-party. Declaring it first-party keeps import grouping
+# deterministic and identical between local and CI.
+profile = "black"
+known_first_party = ["odc"]
+
 ###############################################################################
 # Hatch build configuration
 ###############################################################################
diff --git a/requirements.txt b/requirements.txt
index 68a16522c..45f087918 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,4 +14,10 @@ datasets>=3.6.0
 mlflow==3.11.1
 pyrsmi
 plotext
+pulp
 hip-python
+git+https://github.com/mlcommons/logging.git@6.0.0-rc5
+megatron-energon==7.3.2
+webdataset==1.0.2
+tqdm==4.67.3
+pyarrow==21.0.0
diff --git a/runner/.primus.yaml b/runner/.primus.yaml
index f8b244498..0a5587e0c 100644
--- a/runner/.primus.yaml
+++ b/runner/.primus.yaml
@@ -26,7 +26,7 @@ container:
   # All keys directly map to CLI arguments (--key value)
   options:
     # Container image
-    image: "rocm/primus:v26.3"
+    image: "rocm/primus:v26.4"
 
     # Single-value options
     ipc: "host"
@@ -78,6 +78,7 @@ container:
       - "WANDB_API_KEY"
       - "ENABLE_NUMA_BINDING"
       - "REBUILD_PRIMUS_TURBO"
+      - "GPU_ARCHS"
       - "USING_AINIC"
       - "PATCH_TE_FLASH_ATTN"
       - "REBUILD_BNXT"
@@ -121,7 +122,10 @@ direct:
   master_addr: "localhost"
 
   # Direct mode specific options
-  run_mode: "torchrun"
+  # run_mode: deliberately omitted here so that primus-cli-direct.sh can
+  # auto-select 'single' for subcommands that require it (e.g. node_smoke).
+  # Set run_mode explicitly in your own config to force a specific launcher.
+  # Precedence: CLI --single > user config direct.run_mode > auto-detect > torchrun.
   script: "primus/cli/main.py"
   numa: "auto"
   log_file: ""
diff --git a/runner/README.md b/runner/README.md
index e93241c6a..2317d5663 100644
--- a/runner/README.md
+++ b/runner/README.md
@@ -162,7 +162,7 @@ This is the most common pattern when you need a fixed software stack.
 bash runner/primus-cli slurm \
   -N 4 \
   --nodelist "node[01-04]" \
--- --image "rocm/primus:v26.3" \
+-- --image "rocm/primus:v26.4" \
 -- --env NCCL_DEBUG=INFO \
 -- train pretrain --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml
 ```
diff --git a/runner/helpers/envs/MI355X.sh b/runner/helpers/envs/MI355X.sh
index 8a7811829..87c0350d1 100755
--- a/runner/helpers/envs/MI355X.sh
+++ b/runner/helpers/envs/MI355X.sh
@@ -30,5 +30,8 @@ LOG_INFO_RANK0 "Loading MI355X-specific optimizations..."
 # APU may have different interconnect characteristics
 # Keep common_network.sh settings unless testing shows otherwise
 
+# Disable RCCL warp-speed auto-tuning on gfx950 (MI355X)
+export RCCL_WARP_SPEED_AUTO=${RCCL_WARP_SPEED_AUTO:-0}
+
 # log_exported_vars "MI355X-specific optimizations" \
-#     HSA_XNACK HSA_ENABLE_INTERRUPT GPU_MAX_HEAP_SIZE HSA_KERNARG_POOL_SIZE
+#     HSA_XNACK HSA_ENABLE_INTERRUPT GPU_MAX_HEAP_SIZE HSA_KERNARG_POOL_SIZE RCCL_WARP_SPEED_AUTO
diff --git a/runner/helpers/envs/primus-env.sh b/runner/helpers/envs/primus-env.sh
index 7af0ad2f3..368e55a8f 100755
--- a/runner/helpers/envs/primus-env.sh
+++ b/runner/helpers/envs/primus-env.sh
@@ -42,9 +42,9 @@ detect_gpu_model() {
 
     # Check if rocm-smi is available
     if ! command -v rocm-smi &> /dev/null; then
-        echo "Error: rocm-smi not found. Is ROCm installed?" >&2
+        echo "Warning: rocm-smi not found; set PRIMUS_GPU_MODEL (e.g. MI355X) for GPU env overrides." >&2
         echo "unknown"
-        return 1
+        return 0
     fi
 
     # Get product name from rocm-smi
@@ -64,7 +64,10 @@ detect_gpu_model() {
     echo "$gpu_model"
 }
 
-GPU_MODEL=$(detect_gpu_model)
+GPU_MODEL="${PRIMUS_GPU_MODEL:-}"
+if [[ -z "${GPU_MODEL}" ]]; then
+    GPU_MODEL=$(detect_gpu_model)
+fi
 if [[ "$GPU_MODEL" == "unknown" ]]; then
     LOG_WARN "Unable to detect GPU model. Using default configuration."
 fi
diff --git a/runner/helpers/hooks/00_rebuild_primus_turbo.sh b/runner/helpers/hooks/00_rebuild_primus_turbo.sh
index 831bc8d4f..baaca19de 100755
--- a/runner/helpers/hooks/00_rebuild_primus_turbo.sh
+++ b/runner/helpers/hooks/00_rebuild_primus_turbo.sh
@@ -59,6 +59,13 @@ if [[ -n "$PRIMUS_TURBO_REF" ]]; then
     git checkout "${PRIMUS_TURBO_REF}"
 fi
 
+# The image already ships a built Primus-Turbo, and its wheel installs a top-level
+# `tools` package into site-packages. setup.py does `from tools.build_utils import
+# build_ck_backend`, which then resolves to that stale copy instead of the one in
+# this checkout, so the build dies with ImportError. Drop the installed dist first.
+LOG_INFO_RANK0 "Uninstalling the pre-installed Primus-Turbo..."
+pip3 uninstall -y primus_turbo || true
+
 LOG_INFO_RANK0 "Installing Primus-Turbo build dependencies..."
 pip3 install -r requirements.txt
 
diff --git a/runner/helpers/hooks/03_enable_ainic.sh b/runner/helpers/hooks/03_enable_ainic.sh
index cccd04e58..3dc239801 100755
--- a/runner/helpers/hooks/03_enable_ainic.sh
+++ b/runner/helpers/hooks/03_enable_ainic.sh
@@ -30,14 +30,28 @@ ANP_HOME_DIR="${ANP_HOME_DIR:-/workspace/amd-anp}"
 RCCL_HOME_DIR="${RCCL_HOME_DIR:-/workspace/rccl}"
 MPI_HOME_DIR="${MPI_HOME_DIR:-/opt/ompi}"
 
-# NCCL_NET_PLUGIN: ANP libraray has different names in different containers: librccl-anp.so or librccl-net.so.
-if [ -z "${NCCL_NET_PLUGIN:-}" ]; then
-    if [ -f "${ANP_HOME_DIR}/build/librccl-anp.so" ]; then
-        NCCL_NET_PLUGIN="librccl-anp.so"
-    elif [ -f "${ANP_HOME_DIR}/build/librccl-net.so" ]; then
-        NCCL_NET_PLUGIN="librccl-net.so"
-    else
-        NCCL_NET_PLUGIN="librccl-anp.so"
+# AINIC RoCE mode. ROCm 7.2.1+ builds ANP INTO RCCL (no separate librccl-anp.so);
+# RCCL_AINIC_ROCE=1 enables that built-in ANP path. On older ROCm the external
+# plugin is selected below instead. An explicitly-set NCCL_NET_PLUGIN always takes
+# precedence (force-loads the named plugin) irrespective of RCCL_AINIC_ROCE.
+RCCL_AINIC_ROCE="${RCCL_AINIC_ROCE:-1}"
+
+# NCCL_NET_PLUGIN selection.
+# With built-in ANP (RCCL_AINIC_ROCE=1) default to "none" so RCCL uses its built-in
+# ANP transport instead of auto-loading an external plugin. Otherwise auto-detect the
+# external ANP plugin (name varies by container: librccl-anp.so or librccl-net.so).
+# An explicitly-set NCCL_NET_PLUGIN always wins in both cases.
+if [ "${RCCL_AINIC_ROCE}" = "1" ]; then
+    NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-none}"
+else
+    if [ -z "${NCCL_NET_PLUGIN:-}" ]; then
+        if [ -f "${ANP_HOME_DIR}/build/librccl-anp.so" ]; then
+            NCCL_NET_PLUGIN="librccl-anp.so"
+        elif [ -f "${ANP_HOME_DIR}/build/librccl-net.so" ]; then
+            NCCL_NET_PLUGIN="librccl-net.so"
+        else
+            NCCL_NET_PLUGIN="librccl-anp.so"
+        fi
     fi
 fi
 NCCL_IB_TC="${NCCL_IB_TC:-104}"
@@ -71,6 +85,7 @@ echo "env.ANP_HOME_DIR=${ANP_HOME_DIR}"
 echo "env.RCCL_HOME_DIR=${RCCL_HOME_DIR}"
 echo "env.MPI_HOME_DIR=${MPI_HOME_DIR}"
 echo "env.NCCL_NET_PLUGIN=${NCCL_NET_PLUGIN}"
+echo "env.RCCL_AINIC_ROCE=${RCCL_AINIC_ROCE}"
 echo "env.NCCL_IB_TC=${NCCL_IB_TC}"
 echo "env.NCCL_IB_FIFO_TC=${NCCL_IB_FIFO_TC}"
 echo "env.NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX}"
diff --git a/runner/helpers/hooks/03_enable_aiter.sh b/runner/helpers/hooks/03_enable_aiter.sh
new file mode 100644
index 000000000..67a0d0fd8
--- /dev/null
+++ b/runner/helpers/hooks/03_enable_aiter.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Global hook: enable AITER.
+#
+# Trigger:
+#   export AITER_LOG_LEVEL=ERROR
+#
+# This hook emits env.* lines which will be exported by execute_hooks.sh.
+#
+
+set -euo pipefail
+
+# Set AITER log level to ERROR to suppress the verbose logs.
+echo "env.AITER_LOG_LEVEL=ERROR"
diff --git a/runner/helpers/hooks/05_check_primus_requirements.sh b/runner/helpers/hooks/05_check_primus_requirements.sh
new file mode 100755
index 000000000..f1e8be670
--- /dev/null
+++ b/runner/helpers/hooks/05_check_primus_requirements.sh
@@ -0,0 +1,68 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+# Check that Primus requirements.txt packages are installed.
+# Warns on missing packages by default. Set PRIMUS_STRICT_REQUIREMENTS=1 to
+# make it a fatal error.
+###############################################################################
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+REQ_FILE="${PRIMUS_ROOT}/requirements.txt"
+
+if [[ ! -f "$REQ_FILE" ]]; then
+    exit 0
+fi
+
+MISSING=$(python3 - "$REQ_FILE" << 'PYEOF'
+import importlib.metadata, re, sys, pathlib
+
+req_file = pathlib.Path(sys.argv[1])
+missing = []
+
+for line in req_file.read_text().splitlines():
+    line = line.strip()
+    if not line or line.startswith("#") or line.startswith("-"):
+        continue
+    pkg = re.split(r"[><=!;\[\s]", line)[0].strip()
+    if not pkg:
+        continue
+    normalized = re.sub(r"[-_.]+", "-", pkg).lower()
+    try:
+        importlib.metadata.distribution(normalized)
+    except importlib.metadata.PackageNotFoundError:
+        missing.append(pkg)
+
+for m in missing:
+    print(m)
+PYEOF
+)
+
+if [[ -n "$MISSING" ]]; then
+    echo ""
+    echo "[WARN] Primus requirements not satisfied. Missing packages:"
+    while IFS= read -r pkg; do
+        echo "         - $pkg"
+    done <<< "$MISSING"
+    echo ""
+    echo "       To install:  pip install -r ${REQ_FILE}"
+    echo ""
+
+    if [[ "${PRIMUS_STRICT_REQUIREMENTS:-0}" == "1" ]]; then
+        echo "[ERROR] PRIMUS_STRICT_REQUIREMENTS=1 — aborting. Install missing packages first."
+        exit 1
+    fi
+
+    if [[ "${PRIMUS_AUTO_INSTALL:-0}" == "1" ]]; then
+        echo "[INFO] PRIMUS_AUTO_INSTALL=1 — installing missing packages..."
+        if pip install -r "${REQ_FILE}" --quiet; then
+            echo "[INFO] Successfully installed missing packages."
+        else
+            echo "[ERROR] Failed to install packages. Please install manually."
+            exit 1
+        fi
+    fi
+fi
diff --git a/runner/helpers/hooks/train/posttrain/diffusion/prepare.py b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py
new file mode 100644
index 000000000..4f011269f
--- /dev/null
+++ b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py
@@ -0,0 +1,24 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import runpy
+import sys
+from pathlib import Path
+
+
+def main() -> None:
+    primus_root = Path(__file__).resolve().parents[6]
+    hook = primus_root / "runner" / "helpers" / "hooks" / "train" / "pretrain" / "diffusion" / "prepare.py"
+    sys.argv[0] = str(hook)
+    if "--module_name" not in sys.argv:
+        sys.argv.extend(["--module_name", "post_trainer"])
+    runpy.run_path(str(hook), run_name="__main__")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py b/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py
index 88683f6a7..06f33fbee 100755
--- a/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py
+++ b/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py
@@ -126,6 +126,40 @@ def _prepend_sys_path(*paths: Path):
         sys.path[:] = original_sys_path
 
 
+@contextmanager
+def _unset_nvte_attention_env():
+    """Neutralize the TE attention-backend env vars around AutoBridge model construction.
+
+    This is a general posttrain fix, not a test-only workaround: this hook runs
+    for every native-Megatron SFT run, in a separate subprocess against
+    Megatron-Bridge's own bundled Megatron-LM that never sees Primus's
+    before_train patches (including the ROCm-safe attention_backend one). So it
+    hits *stock* megatron's ``auto`` validation: ``AutoBridge.import_ckpt`` builds
+    a plain ``MCoreGPTModel`` whose ``_set_attention_backend()`` asserts the three
+    ``NVTE_*_ATTN`` vars are unset-or-1, and the ROCm image's baked
+    ``NVTE_FLASH_ATTN=0`` trips it.
+
+    Checkpoint conversion only reshapes weights -- it computes no attention -- so
+    the backend is irrelevant here; the only goal is to get past that assert.
+    Rather than counteract one specific baked value, unset all three for the
+    duration (stock ``auto`` then accepts them and picks defaults harmlessly) and
+    restore whatever was there afterwards. This stays correct for any image: if a
+    future image leaves them unset or sets them to 1, the pop/restore is a no-op;
+    it never assumes a particular baked value. Mirrors the Flux/diffusion conftest
+    fix.
+    """
+    names = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN")
+    saved = {name: os.environ.pop(name, None) for name in names}
+    try:
+        yield
+    finally:
+        for name, value in saved.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
 def convert_checkpoint(hf_path: str, megatron_path: str):
     """
     Convert HuggingFace checkpoint to Megatron torch_dist format.
@@ -138,7 +172,7 @@ def convert_checkpoint(hf_path: str, megatron_path: str):
     log_info(f"  Source: {hf_path}")
     log_info(f"  Target: {megatron_path}")
 
-    with _prepend_sys_path(bridge_path, bridge_megatron_path):
+    with _prepend_sys_path(bridge_path, bridge_megatron_path), _unset_nvte_attention_env():
         from megatron.bridge import AutoBridge
 
         # Convert using AutoBridge - creates torch_dist format checkpoint
diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh
index 302def462..80ca3ecaa 100755
--- a/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh
+++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh
@@ -10,10 +10,11 @@ set -euo pipefail
 echo "[+] Installing Megatron-Bridge dependencies..."
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-# Set up pip cache directory
-PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
-DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}"
-PIP_CACHE_DIR="${PIP_CACHE_DIR:-${DATA_PATH}/pip_cache}"
+# Pip cache must live on a path that exists inside THIS environment (e.g. Docker). Schedulers often set
+# DATA_PATH to a *host* path (e.g. /home/.../data/mlperf_llama2) for dataset mounts; that path is not
+# writable or may not exist in the container, so do not derive PIP_CACHE_DIR from DATA_PATH by default.
+# Keep the cache outside the repo so it is never accidentally committed.
+PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-cache/pip}"
 
 echo "[INFO] Using pip cache: ${PIP_CACHE_DIR}"
 mkdir -p "${PIP_CACHE_DIR}"
@@ -26,4 +27,9 @@ pip install --cache-dir="${PIP_CACHE_DIR}" -U "datasets>=2.14.0"
 
 pip install --cache-dir="${PIP_CACHE_DIR}" -r "${SCRIPT_DIR}/requirements-megatron-bridge.txt"
 
+# mlperf-logging: use the Primus image / pip install -r requirements.txt (6.0.0-rc5); do not pip here.
+
+# datasets 5.x requires fsspec<=2026.4.0; megatron-bridge deps may upgrade it.
+pip install --cache-dir="${PIP_CACHE_DIR}" 'fsspec>=2023.1.0,<=2026.4.0'
+
 echo "[OK] Megatron-Bridge dependencies installed"
diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh
index 4698c3042..632cb1036 100755
--- a/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh
+++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh
@@ -25,7 +25,8 @@ if [[ -z "$CONFIG_FILE" ]]; then
 fi
 
 # Parse the complete config with all extends and nested configs
-PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../.." && pwd)"
+# Repo root: megatron_bridge -> posttrain -> train -> hooks -> helpers -> runner -> Primus (6 levels)
+PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../../" && pwd)"
 cd "$PRIMUS_ROOT"
 
 # Convert CONFIG_FILE to absolute path
@@ -78,10 +79,72 @@ if [[ -z "$HF_PATH" ]]; then
     exit 0
 fi
 
-# Set paths
-DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}"
-HF_CACHE="${HF_HOME:-${DATA_PATH}/huggingface}/hub"
-MEGATRON_PATH="${DATA_PATH}/megatron_checkpoints/$(basename "${HF_PATH}")"
+# Data root: DATA_PATH is often a *host* path. The path may appear in the container (bind mount
+# metadata) but still be unusable. Pick the first root where we can actually mkdir.
+resolve_effective_data() {
+    local c testdir
+    local -a candidates=()
+    [[ -n "${MOUNT_DATA_PATH:-}" ]] && candidates+=("${MOUNT_DATA_PATH}")
+    candidates+=("/data")
+    [[ -n "${DATA_PATH:-}" ]] && candidates+=("${DATA_PATH}")
+    candidates+=("${PRIMUS_ROOT}/data")
+
+    for c in "${candidates[@]}"; do
+        [[ -z "$c" ]] && continue
+        testdir="${c}/.primus_megatron_write_test_$$"
+        if mkdir -p "${testdir}" 2>/dev/null; then
+            rmdir "${testdir}" 2>/dev/null || rm -rf "${testdir}" 2>/dev/null || true
+            printf '%s' "$c"
+            return 0
+        fi
+    done
+    mkdir -p "${PRIMUS_ROOT}/data"
+    printf '%s' "${PRIMUS_ROOT}/data"
+}
+EFFECTIVE_DATA="$(resolve_effective_data)"
+LOG_INFO_RANK0 "Checkpoint/HF data root: ${EFFECTIVE_DATA}"
+
+# Hugging Face cache dirs: schedulers often export HF_HOME (or HF_HUB_CACHE / TRANSFORMERS_CACHE)
+# to host paths that do not exist inside the container. Override when we cannot mkdir there.
+HF_FALLBACK="${EFFECTIVE_DATA}/huggingface"
+if [[ -n "${HF_HOME:-}" ]]; then
+    if mkdir -p "${HF_HOME}/hub/.primus_hf_write_test_$$" 2>/dev/null; then
+        rm -rf "${HF_HOME}/hub/.primus_hf_write_test_$$" 2>/dev/null || true
+    else
+        LOG_INFO_RANK0 "HF_HOME=${HF_HOME} is not usable here; using ${HF_FALLBACK}"
+        export HF_HOME="${HF_FALLBACK}"
+    fi
+else
+    export HF_HOME="${HF_FALLBACK}"
+fi
+mkdir -p "${HF_HOME}/hub"
+
+if [[ -n "${HF_HUB_CACHE:-}" ]]; then
+    if mkdir -p "${HF_HUB_CACHE}/.primus_hf_write_test_$$" 2>/dev/null; then
+        rm -rf "${HF_HUB_CACHE}/.primus_hf_write_test_$$" 2>/dev/null || true
+    else
+        LOG_INFO_RANK0 "HF_HUB_CACHE=${HF_HUB_CACHE} is not usable here; using ${HF_HOME}/hub"
+        export HF_HUB_CACHE="${HF_HOME}/hub"
+    fi
+fi
+
+if [[ -n "${TRANSFORMERS_CACHE:-}" ]]; then
+    if mkdir -p "${TRANSFORMERS_CACHE}/.primus_hf_write_test_$$" 2>/dev/null; then
+        rm -rf "${TRANSFORMERS_CACHE}/.primus_hf_write_test_$$" 2>/dev/null || true
+    else
+        LOG_INFO_RANK0 "TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} is not usable here; using ${HF_HOME}/hub"
+        export TRANSFORMERS_CACHE="${HF_HOME}/hub"
+    fi
+fi
+
+HF_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}"
+
+# So the training process sees the same paths (execute_hooks exports env.* from hook stdout)
+echo "env.HF_HOME=${HF_HOME}"
+[[ -n "${HF_HUB_CACHE:-}" ]] && echo "env.HF_HUB_CACHE=${HF_HUB_CACHE}"
+[[ -n "${TRANSFORMERS_CACHE:-}" ]] && echo "env.TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE}"
+
+MEGATRON_PATH="${EFFECTIVE_DATA}/megatron_checkpoints/$(basename "${HF_PATH}")"
 
 LOG_INFO_RANK0 "HF Model: ${HF_PATH}"
 LOG_INFO_RANK0 "HF Cache: ${HF_CACHE}"
@@ -95,10 +158,22 @@ else
     LOG_INFO_RANK0 "HF checkpoint will be downloaded from ${HF_PATH}"
 fi
 
+resolve_pretrained_checkpoint() {
+    local base="$1"
+    # Megatron-Bridge PEFT validation (checkpoint_exists) and load both expect the
+    # checkpoint *root* (latest_train_state.pt / latest_checkpointed_iteration.txt
+    # live here; weights are under iter_*). Do not point at iter_0000000.
+    printf '%s' "${base}"
+}
+
 # Check if Megatron checkpoint already exists
 if [[ -d "$MEGATRON_PATH" ]]; then
     LOG_INFO_RANK0 "Megatron checkpoint already exists at ${MEGATRON_PATH}, skipping conversion"
-    echo "extra.pretrained_checkpoint=${MEGATRON_PATH}"
+    CKPT_PATH="$(resolve_pretrained_checkpoint "${MEGATRON_PATH}")"
+    echo "extra.pretrained_checkpoint=${CKPT_PATH}"
+    echo "env.NVTE_FLASH_ATTN=0"
+    echo "env.NVTE_FUSED_ATTN=1"
+    echo "env.NVTE_UNFUSED_ATTN=0"
     exit 0
 fi
 
@@ -118,10 +193,23 @@ if [[ "$NODE_RANK" == "0" ]]; then
     # Set up Python path for Megatron-Bridge
     export PYTHONPATH="${PRIMUS_ROOT}/third_party/Megatron-Bridge/src:${PRIMUS_ROOT}/third_party/Megatron-Bridge/3rdparty/Megatron-LM:${PYTHONPATH:-}"
 
+    # HF→Megatron conversion uses attention_backend=auto; clear pre-set TE
+    # env vars from the container image (e.g. NVTE_FLASH_ATTN=0) so Megatron
+    # can configure them for the conversion pass.
+    unset NVTE_FLASH_ATTN NVTE_FUSED_ATTN NVTE_UNFUSED_ATTN
+
     python3 third_party/Megatron-Bridge/examples/conversion/convert_checkpoints.py import \
       --hf-model "${HF_PATH}" \
       --megatron-path "${MEGATRON_PATH}"
 
+    # Restore MLPerf fused-attention env for the training run.
+    export NVTE_FLASH_ATTN=0
+    export NVTE_FUSED_ATTN=1
+    export NVTE_UNFUSED_ATTN=0
+    echo "env.NVTE_FLASH_ATTN=0"
+    echo "env.NVTE_FUSED_ATTN=1"
+    echo "env.NVTE_UNFUSED_ATTN=0"
+
     # Create done file and remove lock
     touch "$DONE_FILE"
     rm -f "$LOCK_FILE"
@@ -153,4 +241,6 @@ else
     echo "[OK] [RANK ${NODE_RANK}] Checkpoint ready at ${MEGATRON_PATH}"
 fi
 
-echo "extra.pretrained_checkpoint=${MEGATRON_PATH}"
+CKPT_PATH="$(resolve_pretrained_checkpoint "${MEGATRON_PATH}")"
+echo "extra.pretrained_checkpoint=${CKPT_PATH}"
+exit 0
diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh
new file mode 100755
index 000000000..ca92cc1b6
--- /dev/null
+++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh
@@ -0,0 +1,98 @@
+#!/bin/bash
+###############################################################################
+# Download SCROLLS gov-report MLPerf dataset and build packed .npy + metadata
+# when post_trainer uses dataset_type=mlperf_dataset.
+###############################################################################
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# shellcheck disable=SC1091
+source "${SCRIPT_DIR}/../../../../../lib/common.sh" || {
+    echo "[ERROR] Failed to load common library" >&2
+    exit 1
+}
+
+PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../../../../../" && pwd)"
+cd "${PRIMUS_ROOT}"
+
+CONFIG_FILE=""
+for ((i=1; i<=$#; i++)); do
+    if [[ "${!i}" == "--config" ]]; then
+        j=$((i+1))
+        CONFIG_FILE="${!j}"
+        break
+    fi
+done
+
+if [[ -z "$CONFIG_FILE" ]]; then
+    exit 0
+fi
+
+if [[ ! "$CONFIG_FILE" = /* ]]; then
+    CONFIG_FILE="${PRIMUS_ROOT}/${CONFIG_FILE#./}"
+fi
+
+read -r DATA_DIR SEQ_LENGTH <<< "$(python3 -c "
+import os
+import sys
+sys.path.insert(0, '${PRIMUS_ROOT}')
+from pathlib import Path
+from primus.core.config.primus_config import load_primus_config, get_module_config
+
+cfg = load_primus_config(Path('${CONFIG_FILE}'), None)
+post = get_module_config(cfg, 'post_trainer')
+if post is None:
+    sys.exit(0)
+params = getattr(post, 'params', None)
+if params is None or getattr(params, 'dataset_type', '') != 'mlperf_dataset':
+    sys.exit(0)
+
+train_path = getattr(params, 'packed_train_data_path', None) or os.environ.get('PACKED_DATA_DIR', '/data')
+train_path = os.path.expandvars(str(train_path))
+data_dir = os.path.dirname(train_path) if train_path.endswith('.npy') else train_path
+seq_length = int(getattr(params, 'seq_length', 8192) or 8192)
+print(data_dir, seq_length)
+" 2>/dev/null || true)"
+
+if [[ -z "${DATA_DIR}" ]]; then
+    exit 0
+fi
+
+SEQ_LENGTH="${SEQ_LENGTH:-8192}"
+mkdir -p "${DATA_DIR}"
+
+LOG_INFO_RANK0 "[mlperf-data] Data root: ${DATA_DIR} (seq_length=${SEQ_LENGTH})"
+echo "env.PACKED_DATA_DIR=${DATA_DIR}"
+echo "env.DATA_PATH=${DATA_DIR}"
+
+if [[ -f "${DATA_DIR}/train.npy" && -f "${DATA_DIR}/validation.npy" && -f "${DATA_DIR}/packed_metadata.jsonl" ]]; then
+    LOG_INFO_RANK0 "[mlperf-data] Dataset already present; skipping download"
+    exit 0
+fi
+
+if [[ -z "${HF_TOKEN:-}" && ( ! -f "${DATA_DIR}/train.npy" || ! -f "${DATA_DIR}/validation.npy" ) ]]; then
+    LOG_ERROR_RANK0 "[mlperf-data] HF_TOKEN is required to download regisss/scrolls_gov_report_preprocessed_mlperf_2"
+    exit 1
+fi
+
+MLPERF_RECIPE_DIR="${PRIMUS_ROOT}/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b"
+
+if [[ ! -f "${DATA_DIR}/train.npy" || ! -f "${DATA_DIR}/validation.npy" ]]; then
+    LOG_INFO_RANK0 "[mlperf-data] Downloading and converting MLPerf dataset..."
+    python3 "${MLPERF_RECIPE_DIR}/download_dataset.py" --data_dir "${DATA_DIR}"
+    python3 "${MLPERF_RECIPE_DIR}/convert_dataset.py" --data_dir "${DATA_DIR}"
+fi
+
+if [[ ! -f "${DATA_DIR}/packed_metadata.jsonl" ]]; then
+    LOG_INFO_RANK0 "[mlperf-data] Creating packed_metadata.jsonl..."
+    python3 "${MLPERF_RECIPE_DIR}/create_metadata.py" "${SEQ_LENGTH}" "${DATA_DIR}/packed_metadata.jsonl"
+fi
+
+for f in train.npy validation.npy packed_metadata.jsonl; do
+    if [[ ! -f "${DATA_DIR}/${f}" ]]; then
+        LOG_ERROR_RANK0 "[mlperf-data] Expected file missing: ${DATA_DIR}/${f}"
+        exit 1
+    fi
+done
+
+LOG_SUCCESS_RANK0 "[mlperf-data] MLPerf dataset ready under ${DATA_DIR}"
diff --git a/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh
new file mode 100755
index 000000000..e0c3c97c0
--- /dev/null
+++ b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh
@@ -0,0 +1,40 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+DEFAULT_PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../../../../.." && pwd)"
+PRIMUS_ROOT="${PRIMUS_PATH:-${DEFAULT_PRIMUS_ROOT}}"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --data_path)
+      DATA_PATH="$2"
+      shift 2
+      ;;
+    --primus_path)
+      PRIMUS_ROOT="$2"
+      shift 2
+      ;;
+    *)
+      shift
+      ;;
+  esac
+done
+
+DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}"
+PIP_CACHE_DIR="${PIP_CACHE_DIR:-${DATA_PATH}/pip_cache}"
+
+echo "[INFO] Using pip cache: ${PIP_CACHE_DIR}"
+mkdir -p "${PIP_CACHE_DIR}"
+
+REQ_FILE="${SCRIPT_DIR}/requirements-diffusion.txt"
+if [[ -f "${REQ_FILE}" ]] && grep -qE '^[[:space:]]*[^#[:space:]]' "${REQ_FILE}"; then
+  echo "[+] Installing Diffusion dependencies..."
+  pip install --cache-dir="${PIP_CACHE_DIR}" -r "${REQ_FILE}"
+  echo "[OK] Diffusion dependencies installed"
+fi
diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py
new file mode 100644
index 000000000..4bdd711bd
--- /dev/null
+++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py
@@ -0,0 +1,209 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+PRIMUS_ROOT = Path(__file__).resolve().parents[6]
+if str(PRIMUS_ROOT) not in sys.path:
+    sys.path.insert(0, str(PRIMUS_ROOT))
+
+from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter
+from primus.core.config.primus_config import get_module_config, load_primus_config
+from primus.core.utils.yaml_utils import nested_namespace_to_dict
+
+
+def _log(message: str) -> None:
+    print(f"[INFO] diffusion prepare: {message}", file=sys.stderr)
+
+
+def _fail(message: str) -> None:
+    print(f"[ERROR] diffusion prepare: {message}", file=sys.stderr)
+    raise SystemExit(1)
+
+
+def _as_dict(value: Any) -> dict[str, Any]:
+    if isinstance(value, dict):
+        return value
+    return nested_namespace_to_dict(value)
+
+
+def _select_module_name(cfg: Any, requested: str | None) -> str:
+    if requested:
+        get_module_config(cfg, requested)
+        return requested
+
+    for module_name in ("pre_trainer", "post_trainer"):
+        try:
+            get_module_config(cfg, module_name)
+            return module_name
+        except Exception:
+            continue
+
+    _fail("config must contain either modules.pre_trainer or modules.post_trainer")
+    raise AssertionError("unreachable")
+
+
+def _is_placeholder(path: str | None) -> bool:
+    return not path or path.startswith("/path/to/")
+
+
+def _require_path(path: str | None, description: str, *, kind: str = "any") -> None:
+    if _is_placeholder(path):
+        _fail(f"{description} is not configured: {path!r}")
+
+    resolved = Path(path).expanduser()
+    if kind == "file" and not resolved.is_file():
+        _fail(f"{description} file not found: {resolved}")
+    if kind == "dir" and not resolved.is_dir():
+        _fail(f"{description} directory not found: {resolved}")
+    if kind == "any" and not resolved.exists():
+        _fail(f"{description} path not found: {resolved}")
+
+    _log(f"{description}: {resolved}")
+
+
+def _require_optional_path(path: str | None, description: str, *, kind: str = "any") -> None:
+    if _is_placeholder(path):
+        _log(f"{description}: not configured; skipping optional initialization")
+        return
+    _require_path(path, description, kind=kind)
+
+
+def _looks_like_local_path(value: str) -> bool:
+    if value.startswith(("/", "./", "../", "~")):
+        return True
+    parts = value.split("/")
+    return len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt"))
+
+
+def _looks_like_hf_reference(value: str) -> bool:
+    if _looks_like_local_path(value):
+        return False
+    parts = value.split("/")
+    return len(parts) >= 2 and all(part for part in parts[:2])
+
+
+def _validate_local_or_hf_id(value: str | None, description: str) -> None:
+    if _is_placeholder(value):
+        _fail(f"{description} is not configured: {value!r}")
+    value = str(value)
+    path = Path(value).expanduser()
+    if path.exists():
+        _log(f"{description}: {path}")
+    elif _looks_like_local_path(value):
+        _fail(f"{description} path not found: {path}")
+    elif _looks_like_hf_reference(value):
+        _log(f"{description}: {value} (assuming Hugging Face reference)")
+    else:
+        _fail(f"{description} is neither an existing path nor a Hugging Face reference: {value!r}")
+
+
+def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None:
+    cfg = load_primus_config(config_path)
+    selected_module = _select_module_name(cfg, module_name)
+    module_cfg = get_module_config(cfg, selected_module)
+
+    if getattr(module_cfg, "framework", None) != "diffusion":
+        _log(f"module {selected_module} framework is not diffusion; skipping")
+        return
+
+    backend_args = DiffusionAdapter().convert_config(module_cfg.params)
+    model = _as_dict(backend_args.model)
+    dataset = _as_dict(backend_args.dataset)
+    trainer = _as_dict(backend_args.trainer)
+
+    dataset_cfg = dataset.get("config", {})
+    processor_cfg = dataset_cfg.get("processor_config", {})
+    model_cfg = model.get("config", {})
+    model_name = model.get("name")
+
+    if model_name == "wan":
+        encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {})
+        _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file")
+        _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir")
+        _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir")
+        _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint")
+        _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file")
+        _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file")
+    elif str(model_name).startswith("flux"):
+        dataset_type = str(dataset_cfg.get("dataset_type", "precomputed")).lower()
+        if dataset_type == "raw":
+            dataset_name = dataset_cfg.get("dataset")
+            dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower()
+            if dataset_name == "cc12m-test":
+                if _is_placeholder(dataset_cfg.get("dataset_path")):
+                    _log("FLUX raw image-text dataset: zirui3/cc12m-test (Hugging Face dataset)")
+                elif dataset_format == "hf_repo":
+                    _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset")
+                else:
+                    kind = "file" if dataset_format == "jsonl" else "dir"
+                    _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind)
+            elif dataset_name == "cc12m-wds":
+                if _is_placeholder(dataset_cfg.get("dataset_path")):
+                    _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)")
+                else:
+                    _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset")
+            elif dataset_format == "hf_repo":
+                _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset")
+            else:
+                kind = "file" if dataset_format == "jsonl" else "dir"
+                _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind)
+            encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {})
+            _validate_local_or_hf_id(encoder_cfg.get("t5_encoder"), "FLUX T5 encoder")
+            _validate_local_or_hf_id(encoder_cfg.get("clip_encoder"), "FLUX CLIP encoder")
+            _validate_local_or_hf_id(encoder_cfg.get("autoencoder"), "FLUX autoencoder checkpoint")
+        elif dataset_type == "precomputed":
+            _require_path(dataset_cfg.get("dataset_path"), "FLUX precomputed dataset", kind="dir")
+            prompt_dropout_prob = float(processor_cfg.get("prompt_dropout_prob", 0.0) or 0.0)
+            if prompt_dropout_prob > 0.0:
+                empty_dir = processor_cfg.get("empty_encodings_path")
+                _require_path(empty_dir, "FLUX empty encodings directory", kind="dir")
+                _require_path(str(Path(empty_dir) / "t5_empty.npy"), "FLUX empty T5 encoding", kind="file")
+                _require_path(
+                    str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file"
+                )
+        else:
+            _fail(f"unsupported FLUX dataset_type: {dataset_type!r}")
+        _require_optional_path(
+            model_cfg.get("load_from_pretrained_path"),
+            "FLUX DiT initialization checkpoint",
+        )
+    else:
+        _fail(f"unsupported diffusion model for prepare hook: {model_name!r}")
+
+    _log(
+        "validated "
+        f"module={selected_module} stage={getattr(backend_args, 'stage', None)} "
+        f"trainer={trainer.get('name')} model={model.get('name')}"
+    )
+    print("env.PREPARED=1")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prepare Primus diffusion training environment")
+    parser.add_argument("--config", type=Path, required=True, help="Experiment YAML config")
+    parser.add_argument("--data_path", type=Path, required=False, help="Reserved for hook API compatibility")
+    parser.add_argument(
+        "--primus_path", type=Path, required=False, help="Reserved for hook API compatibility"
+    )
+    parser.add_argument("--patch_args", type=Path, required=False, help="Reserved for hook API compatibility")
+    parser.add_argument("--backend_path", type=str, default=None, help="Unused; diffusion is in-tree")
+    parser.add_argument("--module_name", type=str, default=None, help="Override module name to validate")
+    args, _unknown = parser.parse_known_args()
+
+    if args.backend_path:
+        _fail("diffusion is an in-tree backend and does not support --backend_path")
+
+    validate_diffusion_config(args.config, module_name=args.module_name)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt
new file mode 100644
index 000000000..50e712cd7
--- /dev/null
+++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt
@@ -0,0 +1,16 @@
+einops
+datasets
+webdataset
+huggingface_hub
+loguru
+safetensors
+sentencepiece
+numpy
+pillow
+packaging
+requests
+decord
+imageio[ffmpeg]
+wandb
+pydantic
+transformers==4.50.0
diff --git a/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh b/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh
new file mode 100644
index 000000000..80f788344
--- /dev/null
+++ b/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh
@@ -0,0 +1,58 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Install NVIDIA-NeMo/Emerging-Optimizers inside the training container.
+#
+# The Muon optimizer path (primus/backends/megatron/core/optimizer/moun.py)
+# hard-requires the ``emerging_optimizers`` package (Newton-Schulz orthogonal-
+# ization, including the DeepSeek-V4 hybrid coefficient set). The package is
+# NOT bundled in the default Primus container, and the public PyPI name
+# ``emerging-optimizers`` is a placeholder stub (metadata-generation-failed) --
+# the real package only installs from the GitHub source. This hook installs it
+# from a pinned commit so a Muon run works out of the box.
+#
+# Gated by PRIMUS_INSTALL_EMERGING_OPTIMIZERS (default off) so non-Muon runs
+# pay nothing; idempotent (skips when already importable). Pure-python wheel
+# (~7s to build/install). run_deepseek_v4.sh sets the gate when OPTIMIZER=muon.
+###############################################################################
+set -euo pipefail
+
+# Only do work when explicitly requested (the Muon launch path sets this).
+case "${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-0}" in
+  1 | true | True | TRUE | yes | on)
+    ;;
+  *)
+    echo "[install_eo] PRIMUS_INSTALL_EMERGING_OPTIMIZERS not set; skipping."
+    exit 0
+    ;;
+esac
+
+# Pinned to match the third_party/Emerging-Optimizers submodule and the
+# Megatron-LM muon.py integration (which passes ``use_nesterov`` to
+# OrthogonalizedOptimizer). The older 06ff4c68 pin used the pre-rename
+# ``nesterov`` kwarg and is incompatible with the current Megatron muon.py.
+EO_COMMIT="${PRIMUS_EMERGING_OPTIMIZERS_COMMIT:-93d9eb3a6c899b50de73992826451fba3ab6adfb}"
+EO_URL="git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@${EO_COMMIT}"
+
+PY="${PYTHON:-python3}"
+
+if "${PY}" -c "import emerging_optimizers" >/dev/null 2>&1; then
+  echo "[install_eo] emerging_optimizers already importable; skipping install."
+  exit 0
+fi
+
+echo "[install_eo] Installing emerging_optimizers from ${EO_URL} ..."
+# --no-deps: the package's only hard runtime dep is torch (already in the
+# container); avoid pulling an incompatible torch/transitive set.
+"${PY}" -m pip install --no-cache-dir --no-deps "${EO_URL}"
+
+if "${PY}" -c "import emerging_optimizers as e; print('[install_eo] installed', getattr(e,'__version__','?'))"; then
+  echo "[install_eo] OK"
+else
+  echo "[install_eo] ERROR: emerging_optimizers still not importable after install" >&2
+  exit 1
+fi
diff --git a/runner/helpers/hooks/train/pretrain/megatron/prepare.py b/runner/helpers/hooks/train/pretrain/megatron/prepare.py
index 78873703a..694ad024e 100644
--- a/runner/helpers/hooks/train/pretrain/megatron/prepare.py
+++ b/runner/helpers/hooks/train/pretrain/megatron/prepare.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -103,11 +103,41 @@ def prepare_dataset(
 
 def prepare_dataset_if_needed(primus_config: PrimusConfig, data_path: Path, env=None):
     pre_trainer_cfg = primus_config.get_module_config("pre_trainer")
+
+    # Skip dataset preparation if train_data_path is explicitly set
     if pre_trainer_cfg.train_data_path is not None:
         return
 
-    tokenizer_type = pre_trainer_cfg.tokenizer_type
-    tokenizer_model = pre_trainer_cfg.tokenizer_model
+    # Check if this is a diffusion model (uses Energon datasets, not tokenized datasets)
+    model_type = getattr(pre_trainer_cfg, "model_type", None)
+    trainer_class = getattr(pre_trainer_cfg, "trainer_class", None)
+    data_path_config = getattr(pre_trainer_cfg, "data_path", None)
+
+    # Determine if this is a diffusion model
+    is_diffusion = False
+    if model_type == "diffusion_model":
+        is_diffusion = True
+    elif trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)):
+        is_diffusion = True
+    elif hasattr(model_type, "name") and "DIFFUSION" in model_type.name:
+        is_diffusion = True
+
+    # For diffusion models with data_path set, skip tokenization (they use Energon)
+    if is_diffusion and data_path_config:
+        log_info("=" * 80)
+        log_info("Diffusion model detected with data_path configured.")
+        log_info("Skipping tokenization (diffusion models use Energon datasets).")
+        log_info(f"Data will be loaded from: {data_path_config}")
+        log_info("=" * 80)
+        return
+
+    # For language models, proceed with tokenization
+    tokenizer_type = getattr(pre_trainer_cfg, "tokenizer_type", None)
+    if not tokenizer_type:
+        log_info("No tokenizer_type found, skipping dataset preparation.")
+        return
+
+    tokenizer_model = getattr(pre_trainer_cfg, "tokenizer_model", None)
     default_tokenized_path = Path(data_path) / f"bookcorpus/{tokenizer_type}/bookcorpus_text_sentence"
     tokenized_data_path = Path(os.environ.get("TOKENIZED_DATA_PATH", str(default_tokenized_path)))
 
@@ -120,6 +150,12 @@ def prepare_dataset_if_needed(primus_config: PrimusConfig, data_path: Path, env=
             if not hf_token:
                 log_error_and_exit("Environment variable HF_TOKEN must be set.")
 
+            if not tokenizer_model:
+                log_error_and_exit(
+                    "tokenizer_model not found in configuration. "
+                    "This is required for language model tokenization."
+                )
+
             log_info(f"TOKENIZED_DATA_PATH is {tokenized_data_path}")
 
             prepare_dataset(
@@ -179,9 +215,7 @@ def build_megatron_helper(megatron_path: Path):
     dataset_cpp_dir = megatron_path / "megatron/core/datasets"
     log_info(f"Building Megatron dataset helper in {dataset_cpp_dir}")
 
-    # `-s` silences make's "Nothing to be done"/recipe echo on no-op rebuilds;
-    # real compiler errors still surface and are handled below.
-    ret = subprocess.run(["make", "-s"], cwd=dataset_cpp_dir)
+    ret = subprocess.run(["make"], cwd=dataset_cpp_dir)
     if ret.returncode != 0:
         log_error_and_exit("Building Megatron C++ helper failed.")
 
diff --git a/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt b/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt
index 7fbaa85a1..cacc76c22 100644
--- a/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt
+++ b/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt
@@ -1,4 +1,5 @@
 # Add Megatron pretrain hook private Python deps here when needed.
 # Qwen3.5 / gated-delta-net attention (Megatron core GatedDeltaNet) needs FLA.
-flash-linear-attention==0.5.1
+# 0.5.2 fixes TileLang being enabled by default on ROCm (0.5.0/0.5.1 regressed qwen3_5).
+flash-linear-attention==0.5.2
 causal-conv1d~=1.5
diff --git a/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py b/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py
index 799d2e820..459300b08 100644
--- a/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py
+++ b/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py
@@ -50,8 +50,16 @@ def parse_args():
 
 
 def pip_install_editable(path: Path, name: str):
-    log_info(f"Installing {name} in editable mode via pip (path: {path})")
-    ret = subprocess.run(["pip", "install", "-e", ".", "-q"], cwd=path)
+    # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools,
+    # no __init__.py) that the default editable finder cannot resolve. Use compat
+    # mode (writes a .pth with the source root on sys.path) after uninstalling any
+    # stale torchtitan, so the training subprocess can import them.
+    log_info(f"Installing {name} in editable (compat) mode via pip (path: {path})")
+    subprocess.run(["pip", "uninstall", "-y", name.lower()], cwd=path)
+    ret = subprocess.run(
+        ["pip", "install", "-e", ".", "--config-settings", "editable_mode=compat", "-q"],
+        cwd=path,
+    )
     if ret.returncode != 0:
         log_error_and_exit(f"Failed to install {name} via pip.")
 
diff --git a/runner/helpers/patches/10_fix_libionic_abi4.sh b/runner/helpers/patches/10_fix_libionic_abi4.sh
new file mode 100755
index 000000000..1ca367df7
--- /dev/null
+++ b/runner/helpers/patches/10_fix_libionic_abi4.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# primus-cli --patch script: swap the ionic libibverbs provider (.so) for an
+# ABI-4-capable build at container launch (before torchrun starts).
+#
+# Why: some AINIC container images ship a libionic provider built against stock
+# rdma-core that only advertises ionic uverbs ABI 1, while the host kernel ionic
+# driver exposes ABI 4 -- so libibverbs rejects every ionic_* device and RDMA
+# falls back to TCP. Copying in an ABI-4 provider .so fixes device enumeration
+# without rebuilding the image.
+#
+# Controlled by PRIMUS_LIBIONIC_SRC_ABI4_SO (set it in your launch script, e.g.
+# run_flash.sh). If unset/empty the patch is skipped. The PRIMUS_ prefix makes it
+# auto-forward into the container (see primus-cli-container.sh env passthrough).
+###############################################################################
+set -euo pipefail
+
+if [[ -z "${PRIMUS_LIBIONIC_SRC_ABI4_SO:-}" ]]; then
+    echo "[fix_libionic_abi4] PRIMUS_LIBIONIC_SRC_ABI4_SO not set -- skipping"
+    exit 2  # 2 = skip (not an error), per runner/helpers/execute_patches.sh
+fi
+
+SRC="$PRIMUS_LIBIONIC_SRC_ABI4_SO"
+if [[ ! -f "$SRC" ]]; then
+    echo "[fix_libionic_abi4] source .so not found: $SRC" >&2
+    exit 1
+fi
+
+PROVIDER_LINK=/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so
+if [[ ! -e "$PROVIDER_LINK" ]]; then
+    echo "[fix_libionic_abi4] ionic provider not present ($PROVIDER_LINK); nothing to patch" >&2
+    exit 2
+fi
+DST="$(readlink -f "$PROVIDER_LINK")"
+
+if cmp -s "$SRC" "$DST"; then
+    echo "[fix_libionic_abi4] provider already matches source -- skipping"
+    exit 2
+fi
+
+cp --remove-destination "$SRC" "$DST"
+echo "[fix_libionic_abi4] swapped ionic provider: $DST <- $SRC"
diff --git a/runner/helpers/patches/11_fix_lld_stub.sh b/runner/helpers/patches/11_fix_lld_stub.sh
new file mode 100755
index 000000000..2c7d8b761
--- /dev/null
+++ b/runner/helpers/patches/11_fix_lld_stub.sh
@@ -0,0 +1,84 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# primus-cli --patch script: repoint a broken `ld.lld` launcher stub at a
+# working LLD at container launch (before torchrun starts).
+#
+# Why: images that get ROCm from the pip `_rocm_sdk_devel` wheel (ROCm 7.14 and
+# newer) put that wheel's llvm/bin early on PATH, where every tool is a ~26 KB
+# forwarding stub. The `ld.lld` stub cannot resolve its own program path:
+#
+#   $ ld.lld --version
+#   could not find path component of main program: 'ld.lld'
+#
+# FlyDSL shells out to `lld` to link the AMDGPU device module it JIT-compiles,
+# so every FlyDSL kernel (PrimusTurbo native sparse-MLA attention, gemm_fp8,
+# MegaMoE) dies with "lld invocation failed / An error happened while
+# serializing the module". On multi-node runs that surfaces as a hang instead:
+# the ranks that fail stall while the rest wait in the next collective.
+#
+# FlyDSL exposes no linker override, so the only lever is the binary that PATH
+# resolves. Triton bundles a working LLD in the same image, which this patch
+# links over the stub.
+#
+# The patch probes the toolchain and only acts on an image that is actually
+# broken -- it skips when there is no `ld.lld` on PATH and when the one found
+# already reports a version -- so it is a no-op on healthy images and needs no
+# opt-in. Point it at a specific LLD with PRIMUS_LLD_REPLACEMENT (the PRIMUS_
+# prefix auto-forwards into the container; see primus-cli-container.sh env
+# passthrough).
+#
+# This is a runtime stop-gap. The durable fix is for the image to ship a working
+# ld.lld in the ROCm SDK it installs.
+###############################################################################
+set -euo pipefail
+
+# Probe the linker the way FlyDSL calls it: by bare name, through PATH. That
+# distinction matters -- the SDK stub resolves itself when handed a path (running
+# it as /full/path/ld.lld prints "AMD LLD <version>") and only fails when found
+# through PATH, which is exactly how a compiler invokes it.
+_lld_ok() {
+    "$@" --version 2>&1 | head -1 | grep -qE "LLD [0-9]"
+}
+
+STUB="$(command -v ld.lld 2>/dev/null || true)"
+if [[ -z "$STUB" ]]; then
+    echo "[fix_lld_stub] no ld.lld on PATH; nothing to patch"
+    exit 2 # 2 = skip (not an error), per runner/helpers/execute_patches.sh
+fi
+if _lld_ok ld.lld; then
+    echo "[fix_lld_stub] ld.lld already works ($STUB) -- skipping"
+    exit 2
+fi
+
+echo "[fix_lld_stub] ld.lld broken when invoked through PATH: $STUB"
+echo "[fix_lld_stub]   $(ld.lld --version 2>&1 | head -1)"
+
+# Prefer an explicit override, then Triton's bundled LLVM.
+CANDIDATES=()
+[[ -n "${PRIMUS_LLD_REPLACEMENT:-}" ]] && CANDIDATES+=("$PRIMUS_LLD_REPLACEMENT")
+while IFS= read -r c; do CANDIDATES+=("$c"); done < <(
+    ls -d /root/.triton/llvm/*/bin/ld.lld 2>/dev/null || true
+)
+
+for cand in "${CANDIDATES[@]}"; do
+    _lld_ok "$cand" || continue
+    ln -sf "$(readlink -f "$cand")" "$STUB"
+    # Re-probe through PATH: a candidate that works by path is only useful if the
+    # replacement also works by bare name (a real binary does, another stub
+    # would not).
+    if ! _lld_ok ld.lld; then
+        echo "[fix_lld_stub] $cand still fails through PATH after linking" >&2
+        continue
+    fi
+    echo "[fix_lld_stub] repointed $STUB -> $(readlink -f "$cand")"
+    echo "[fix_lld_stub] $(ld.lld --version 2>&1 | head -1)"
+    exit 0
+done
+
+echo "[fix_lld_stub] no working LLD found (tried: ${CANDIDATES[*]:-none})" >&2
+exit 1
diff --git a/runner/primus-cli-container.sh b/runner/primus-cli-container.sh
index 9a06e01f5..ace9bf5ee 100755
--- a/runner/primus-cli-container.sh
+++ b/runner/primus-cli-container.sh
@@ -45,7 +45,7 @@ Docker/Podman Options:
         --cap-add <CAPABILITY>       Add Linux capabilities (e.g., SYS_PTRACE)
 
     Container Configuration:
-        --image <DOCKER_IMAGE>       Docker image [default: rocm/primus:v26.3]
+        --image <DOCKER_IMAGE>       Docker image [default: rocm/primus:v26.4]
         --name <NAME>                Container name
         --user <UID:GID>             Run as specific user (e.g., 1000:1000)
         --network <NET>              Network mode (e.g., host, bridge)
@@ -408,6 +408,14 @@ for key in "${!container_config[@]}"; do
     [[ "$opt_name" == "image" ]] && continue
     [[ "$opt_value" == "[]" ]] && continue
 
+    # Make the container name unique per job so an orphaned container left by a
+    # previous/cancelled job on a reused node doesn't cause a "name already in
+    # use" conflict. SLURM_JOB_ID is per-job (same on all its nodes, unique across
+    # jobs); fall back to the PID when not under Slurm.
+    if [[ "$opt_name" == "name" && -n "$opt_value" && "$opt_value" != *$'\n'* ]]; then
+        opt_value="${opt_value}-${SLURM_JOB_ID:-$$}"
+    fi
+
     # Check if this is a cumulative option
     is_cumulative=0
     for cum_opt in "${CUMULATIVE_OPTIONS[@]}"; do
@@ -452,7 +460,10 @@ if [[ "$CLEAN_DOCKER_CONTAINER" == "true" ]]; then
     LOG_INFO_RANK0 "[container] Cleaning up existing containers..."
     CONTAINERS="$($CONTAINER_RUNTIME ps -aq)"
     if [[ -n "$CONTAINERS" ]]; then
-        printf '%s\n' "$CONTAINERS" | xargs -r -n1 "$CONTAINER_RUNTIME" rm -f
+        # Tolerate per-container removal failures ("removal already in progress",
+        # "No such container") so a stale/concurrently-removing container on a
+        # shared node does not abort the whole run via xargs' exit code 123.
+        printf '%s\n' "$CONTAINERS" | xargs -r -n1 -I{} sh -c "$CONTAINER_RUNTIME rm -f {} 2>/dev/null || true"
         LOG_INFO_RANK0 "[container] Removed containers: $CONTAINERS"
     else
         LOG_INFO_RANK0 "[container] No containers to remove."
@@ -530,5 +541,28 @@ if [[ "$DRY_RUN_MODE" == "true" ]]; then
     exit 0
 fi
 
+# ---------------------------------------------------------------------------
+# Optional: authenticate to the container registry before the implicit pull.
+# On clusters whose compute nodes are not pre-logged-in, a private image
+# (e.g. docker.io/tasimage/primus:*) cannot be pulled and `docker run` fails
+# with "Unable to find image '<image>' locally" followed by a pull-access
+# error. Export DOCKER_LOGIN_USER and DOCKER_LOGIN_KEY (e.g. a Docker Hub PAT)
+# and the scheduler propagates them to every node (spur sbatch --export=ALL),
+# so the login+pull happens automatically at launch. DOCKER_LOGIN_REGISTRY is
+# optional and defaults to Docker Hub.
+# ---------------------------------------------------------------------------
+if [[ -n "${DOCKER_LOGIN_KEY:-}" ]]; then
+    if [[ -z "${DOCKER_LOGIN_USER:-}" ]]; then
+        LOG_ERROR "[container] DOCKER_LOGIN_KEY is set but DOCKER_LOGIN_USER is empty."
+        exit 1
+    fi
+    LOG_INFO_RANK0 "[container] Logging in to registry ${DOCKER_LOGIN_REGISTRY:-docker.io} as ${DOCKER_LOGIN_USER}..."
+    if ! printf '%s' "${DOCKER_LOGIN_KEY}" | \
+        "$CONTAINER_RUNTIME" login ${DOCKER_LOGIN_REGISTRY:+"$DOCKER_LOGIN_REGISTRY"} -u "${DOCKER_LOGIN_USER}" --password-stdin; then
+        LOG_ERROR "[container] Registry login failed for user ${DOCKER_LOGIN_USER}."
+        exit 1
+    fi
+fi
+
 LOG_INFO "[container] Executing command..."
 "${CMD[@]}"
diff --git a/runner/primus-cli-direct.sh b/runner/primus-cli-direct.sh
index d87de0d01..c90a0d86d 100755
--- a/runner/primus-cli-direct.sh
+++ b/runner/primus-cli-direct.sh
@@ -33,6 +33,11 @@ Options:
     --log_file PATH      Save log to a specific file (default: logs/log_TIMESTAMP.txt)
     --numa               Force enable NUMA binding for processes
     --no-numa            Force disable NUMA binding for processes
+    --silent             Suppress all launcher and python tool stdout (back-pocket option).
+                         Must be placed BEFORE the '--' separator. Launcher errors
+                         (LOG_ERROR / LOG_WARN, written to stderr) and the log file
+                         are preserved. Exit code is propagated. NOT recommended for
+                         normal use -- you lose live progress; prefer the log file.
 
 Distributed Environment Variables:
     NNODES        Number of nodes participating in distributed run        [default: 1]
@@ -41,6 +46,17 @@ Distributed Environment Variables:
     MASTER_ADDR   Hostname or IP of master node                           [default: localhost]
     MASTER_PORT   Port of master node                                     [default: 1234]
 
+    When running inside a SLURM allocation (SLURM_JOB_ID set), the above are
+    auto-derived from SLURM_NNODES / SLURM_NODEID / SLURM_NODELIST when not
+    pre-exported. Pre-exported values always win, so the standard
+    slurm-entry -> container -> direct chain is unaffected.
+
+Optional Environment Variables:
+    VENV_ACTIVATE Path to a Python virtualenv activate script. If set, sourced
+                  before the primus run. Unset = no-op (use system python or
+                  the container's bundled python). Set + missing file =
+                  fail-fast (avoid silent torch-version mismatches).
+
 You can set these variables in either of the following ways:
     # (1) Export variables before launch (recommended for scripts or single-node runs)
       export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR=host1
@@ -50,6 +66,10 @@ You can set these variables in either of the following ways:
       primus-cli direct --env NNODES=2 --env GPUS_PER_NODE=8 --env NODE_RANK=1 --env MASTER_ADDR=host1 -- \\
         benchmark gemm -M 4096 -N 4096 -K 4096
 
+    # (3) Let SLURM provide them (inside an existing allocation)
+      export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+      srun -N 4 --ntasks-per-node=1 ./runner/primus-cli direct -- preflight --quick
+
 Examples:
     # Pretrain with a config file (single node)
       primus-cli direct -- train pretrain --config examples/megatron/exp_pretrain.yaml
@@ -77,12 +97,19 @@ Examples:
     # Force enable NUMA binding for better performance
       primus-cli direct --numa -- benchmark gemm -M 8192 -N 8192 -K 8192
 
+    # Run silently (back-pocket option; launcher errors and log file preserved)
+      primus-cli direct --silent -- preflight --quick
+
 Notes:
     - If --single is specified, Primus skips torchrun and uses python3 directly.
+    - run_mode auto-detection: when the primus subcommand is 'node_smoke', run_mode
+      defaults to 'single' (node_smoke runs one process per node by design).
+      Explicit --single or config direct.run_mode still wins.
     - If --script is not specified, defaults to primus/cli/main.py.
     - Always separate Primus arguments from launcher options using '--'.
     - Environment variables can be mixed: 'export' takes precedence unless overridden by '--env'.
     - Multi-node jobs require MASTER_ADDR set to the master node's hostname/IP.
+      Inside a SLURM allocation it is auto-resolved from SLURM_NODELIST.
     - Patch scripts are executed in order before running the main script (useful for env setup, hot fixes, etc.).
     - NUMA binding is auto-disabled by default; use --numa to enable for better memory locality.
     - GPU-specific optimizations: The script automatically sources primus-env.sh, which detects your
@@ -97,6 +124,43 @@ if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
     exit 0
 fi
 
+###############################################################################
+# Runner-level --silent (back-pocket option for clean stdout)
+#
+# Contract:
+#   - launcher LOG_INFO / LOG_INFO_RANK0 / LOG_DEBUG_* (stdout)  -> /dev/null
+#   - launcher LOG_WARN / LOG_ERROR  (stderr, written via >&2)   -> terminal
+#   - python tool's stdout + stderr                              -> /dev/null
+#       (the launched CMD ends with `2>&1 | tee <log_file>`, so the python
+#        child's fd2 is merged into the pipe and tee writes to the log file
+#        + this script's fd1 = /dev/null)
+#   - log file                                                   -> captures all
+#   - exit code                                                  -> propagated
+#
+# --silent is only honored BEFORE the `--` separator. Anywhere after `--` it
+# is forwarded to the python tool whose argparse will reject it (intentional;
+# python tools have no --silent flag and never will -- silencing is a bash-side
+# concern).
+#
+# Not recommended for normal use: read the log file or omit --silent if you
+# want to see live tool output and progress.
+###############################################################################
+SILENT=0
+for _arg in "$@"; do
+    if [[ "$_arg" == "--" ]]; then
+        break
+    fi
+    if [[ "$_arg" == "--silent" ]]; then
+        SILENT=1
+        break
+    fi
+done
+if [[ "$SILENT" == "1" ]]; then
+    # fd1 -> /dev/null. fd2 is left attached to the terminal, so LOG_ERROR /
+    # LOG_WARN (which write via >&2) still reach the operator.
+    exec >/dev/null
+fi
+
 # Resolve runner directory
 # Use RUNNER_DIR instead of SCRIPT_DIR to avoid conflicts with sourced scripts
 RUNNER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -180,6 +244,12 @@ while [[ $# -gt 0 ]]; do
             PRE_PARSE_ARGS+=("$1")
             shift
             ;;
+        --silent)
+            # Already consumed by the pre-scan above (which has applied
+            # `exec >/dev/null`). Swallow it here so it never reaches the
+            # python parser, which has no --silent flag.
+            shift
+            ;;
         --)
             # Explicit separator: remaining args are for primus Python module
             shift  # skip the '--'
@@ -241,8 +311,11 @@ if [[ "$DRY_RUN_MODE" == "0" ]]; then
     fi
 fi
 
-# Set default values for parameters not in config
-direct_config[run_mode]="${direct_config[run_mode]:-torchrun}"
+# Set default values for parameters not in config.
+# NOTE: run_mode default is deliberately NOT applied here. The auto-detect step
+# after STEP 4 needs to distinguish "user did not specify run_mode" (so we can
+# auto-select `single` for node_smoke) from "default got applied in STEP 3".
+# Defaulting to `torchrun` therefore happens after STEP 4 instead.
 direct_config[script]="${direct_config[script]:-primus/cli/main.py}"
 direct_config[numa]="${direct_config[numa]:-auto}"
 direct_config[log_file]="${direct_config[log_file]:-}"
@@ -315,6 +388,36 @@ set -- "${primus_args[@]}"
 # Done here so it flows through hooks/patches and matches container/slurm pattern.
 [[ "$DEBUG_MODE" == "1" ]] && set -- --debug "$@"
 
+###############################################################################
+# STEP 4.4: Auto-select run_mode based on primus subcommand
+#
+# Some primus subcommands MUST run as a single process per srun task (no
+# torchrun fan-out, no inter-node rendezvous). node_smoke is the canonical
+# example: it runs one process per node by design and the per-GPU phase is
+# launched as subprocesses internally. Listed here so users don't need to
+# remember --single. Explicit --single / config direct.run_mode still wins
+# (those paths set direct_config[run_mode] before this block runs).
+###############################################################################
+SINGLE_MODE_SUBCOMMANDS=(node_smoke)
+if [[ -z "${direct_config[run_mode]:-}" ]]; then
+    _detected_subcmd=""
+    for _arg in "${primus_args[@]}"; do
+        case "$_arg" in
+            --*|-*) continue ;;
+            *) _detected_subcmd="$_arg"; break ;;
+        esac
+    done
+    _default_run_mode="torchrun"
+    for _sc in "${SINGLE_MODE_SUBCOMMANDS[@]}"; do
+        if [[ "$_detected_subcmd" == "$_sc" ]]; then
+            _default_run_mode="single"
+            LOG_INFO_RANK0 "[direct] Auto-selected run_mode=single for subcommand '$_detected_subcmd'"
+            break
+        fi
+    done
+    direct_config[run_mode]="$_default_run_mode"
+fi
+
 ###############################################################################
 # STEP 4.5: Process non-cumulative parameters (use last value only)
 ###############################################################################
@@ -337,6 +440,83 @@ if [[ -z "${direct_config[log_file]:-}" ]]; then
 fi
 mkdir -p "$(dirname "${direct_config[log_file]:-}")"
 
+###############################################################################
+# STEP 4.7: Activate virtualenv (R1) and derive distributed env from SLURM (R2)
+#
+# Previously these lived in runner/run_preflight_direct.sh and
+# runner/run_node_smoke_direct.sh (now deleted). Hoisting them here makes every
+# `primus-cli direct -- ...` call site (host srun, slurm-entry -> direct,
+# slurm-entry -> container -> direct) inherit identical behavior.
+###############################################################################
+
+# R1 -- Python virtualenv. VENV_ACTIVATE unset = no-op (this is the right
+# default for the container path: primus-cli-container.sh does not auto-forward
+# VENV_ACTIVATE through its env passthrough whitelist, so inside the container
+# we use the container's bundled python). Set + missing = fail-fast: better a
+# loud error than a silent torch-version mismatch.
+if [[ -n "${VENV_ACTIVATE:-}" ]]; then
+    if [[ ! -f "$VENV_ACTIVATE" ]]; then
+        LOG_ERROR "[direct] VENV_ACTIVATE is set but file does not exist: $VENV_ACTIVATE"
+        exit 1
+    fi
+    # shellcheck disable=SC1090
+    source "$VENV_ACTIVATE"
+    LOG_INFO_RANK0 "[direct] Activated virtualenv: $VENV_ACTIVATE"
+fi
+
+# R2 -- distributed env. Pre-set values always win (the existing
+# slurm-entry -> container -> direct chain already passes them via --env, so
+# this block is a no-op there). When called directly under srun on the host,
+# this block derives them from SLURM_*.
+export GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
+export MASTER_PORT="${MASTER_PORT:-1234}"
+
+if [[ -n "${SLURM_JOB_ID:-}" ]]; then
+    # Pre-exported values always win (per plan). The standard
+    # slurm-entry -> container -> direct chain already exports NNODES /
+    # NODE_RANK with the SLURM-derived values before direct.sh runs, so this
+    # block is a no-op there. When called directly under srun on bare metal
+    # without pre-set values, derive from SLURM_*.
+    export NNODES="${NNODES:-${SLURM_NNODES:-${SLURM_JOB_NUM_NODES:-1}}}"
+    export NODE_RANK="${NODE_RANK:-${SLURM_NODEID:-${SLURM_PROCID:-0}}}"
+    if [[ -z "${MASTER_ADDR:-}" || "${MASTER_ADDR}" == "localhost" ]]; then
+        if command -v scontrol >/dev/null 2>&1 && [[ -n "${SLURM_NODELIST:-}" ]]; then
+            if ! MASTER_ADDR="$(scontrol show hostnames "$SLURM_NODELIST" | head -n1)" \
+               || [[ -z "$MASTER_ADDR" ]]; then
+                LOG_ERROR "[direct] Failed to resolve MASTER_ADDR from SLURM_NODELIST=${SLURM_NODELIST:-<unset>}"
+                exit 1
+            fi
+            export MASTER_ADDR
+        else
+            # In a SLURM context but we cannot resolve a real address from
+            # scontrol -- either scontrol is not installed (CI / dev VM) or
+            # SLURM_NODELIST was not propagated (rare; can happen in stubbed
+            # tests or single-node allocations). Fall back to localhost so
+            # the downstream sanity check at line 503 has a valid value;
+            # NODE_RANK=0 + MASTER_ADDR=localhost is correct for single-node
+            # SLURM runs and for dry-run smoke tests. Real multi-node
+            # bare-srun on a SLURM head node always has both, so this
+            # branch never fires in production.
+            export MASTER_ADDR="localhost"
+        fi
+    fi
+    LOG_INFO_RANK0 "[direct] SLURM detected: JOB_ID=$SLURM_JOB_ID NNODES=$NNODES NODE_RANK=$NODE_RANK MASTER_ADDR=${MASTER_ADDR:-<unset>}"
+else
+    export NNODES="${NNODES:-1}"
+    export NODE_RANK="${NODE_RANK:-0}"
+    export MASTER_ADDR="${MASTER_ADDR:-localhost}"
+fi
+
+# Sanity-check the resolved distributed env (lifted from the deleted wrappers).
+[[ "$NNODES"      =~ ^[1-9][0-9]*$ ]] || { LOG_ERROR "[direct] NNODES must be a positive integer (got '$NNODES')"; exit 1; }
+[[ "$NODE_RANK"   =~ ^[0-9]+$       ]] || { LOG_ERROR "[direct] NODE_RANK must be a non-negative integer (got '$NODE_RANK')"; exit 1; }
+[[ "$MASTER_PORT" =~ ^[0-9]+$       ]] || { LOG_ERROR "[direct] MASTER_PORT must be a non-negative integer (got '$MASTER_PORT')"; exit 1; }
+[[ -n "$MASTER_ADDR" ]]                || { LOG_ERROR "[direct] MASTER_ADDR is empty"; exit 1; }
+(( NODE_RANK < NNODES ))               || { LOG_ERROR "[direct] NODE_RANK ($NODE_RANK) must be < NNODES ($NNODES)"; exit 1; }
+if [[ "$MASTER_ADDR" == "localhost" && "${NNODES:-1}" -gt 1 ]]; then
+  LOG_WARN "[direct] MASTER_ADDR=localhost with NNODES=$NNODES — multi-node will likely fail"
+fi
+
 ###############################################################################
 # STEP 5: Source GPU environment and helper modules
 ###############################################################################
@@ -461,7 +641,14 @@ fi
 # STEP 9: Build launch command
 ###############################################################################
 
-# Allow RUN_MODE to be overridden by environment variable
+# Final run-mode resolution. RUN_MODE-from-env wins over direct_config[run_mode]
+# because framework prepare-hooks (e.g. runner/helpers/hooks/train/pretrain/
+# maxtext/prepare.py) emit `env.RUN_MODE=single` from STEP 6 -- the hook layer
+# knows things the launcher can't (e.g. "this framework is JAX, not torch, so
+# torchrun would be wrong"). $RUN_MODE is the authoritative value for the rest
+# of the script -- the display block and the torchrun-only Distributed Settings
+# gate at STEP 10 both read $RUN_MODE, NOT direct_config[run_mode], so a hook
+# that flips the mode is faithfully reflected in the printed configuration.
 RUN_MODE="${RUN_MODE:-${direct_config[run_mode]:-torchrun}}"
 
 # Resolve the launch target. Normally this is the script path
@@ -533,7 +720,7 @@ else
     print_section "Primus Direct Launch Configuration"
 fi
 
-PRINT_INFO_RANK0 "  Run Mode        : ${direct_config[run_mode]:-}"
+PRINT_INFO_RANK0 "  Run Mode        : ${RUN_MODE}"
 PRINT_INFO_RANK0 "  Script Path     : ${LAUNCH_TARGET[*]:-${direct_config[script]:-}}"
 PRINT_INFO_RANK0 "  Config File     : ${CONFIG_FILE:-<none>}"
 PRINT_INFO_RANK0 "  Log File        : ${direct_config[log_file]:-}"
@@ -554,7 +741,7 @@ if [[ ${#primus_env_kv[@]} -gt 0 ]]; then
     PRINT_INFO_RANK0 ""
 fi
 
-if [[ "${direct_config[run_mode]:-}" == "torchrun" ]]; then
+if [[ "${RUN_MODE}" == "torchrun" ]]; then
     PRINT_INFO_RANK0 "  Distributed Settings:"
     PRINT_INFO_RANK0 "    NNODES          : ${NNODES:-1}"
     PRINT_INFO_RANK0 "    NODE_RANK       : ${NODE_RANK:-0}"
diff --git a/runner/primus-cli-slurm-entry.sh b/runner/primus-cli-slurm-entry.sh
index f45f57585..4ac9f858d 100755
--- a/runner/primus-cli-slurm-entry.sh
+++ b/runner/primus-cli-slurm-entry.sh
@@ -7,8 +7,17 @@
 
 set -euo pipefail
 
-# Resolve runner directory robustly (handles symlinks)
-RUNNER_DIR="$(cd "$(dirname "$(realpath "$0")")" && pwd)"
+# Resolve runner directory. Prefer PRIMUS_RUNNER_DIR when the launcher exported
+# it: schedulers relocate a batch script into their spool dir (spur copies it to
+# /var/spool/spur/job<id>/spur_job.sh, standard Slurm to
+# /var/spool/slurmd/job<id>/slurm_script), so on the sbatch path $0 points at a
+# copy with no sibling lib/. Fall back to self-location (handles symlinks) for
+# srun and direct invocation.
+if [[ -n "${PRIMUS_RUNNER_DIR:-}" && -f "${PRIMUS_RUNNER_DIR}/lib/common.sh" ]]; then
+    RUNNER_DIR="$PRIMUS_RUNNER_DIR"
+else
+    RUNNER_DIR="$(cd "$(dirname "$(realpath "$0")")" && pwd)"
+fi
 
 # Load common library (required)
 # shellcheck disable=SC1091
@@ -118,8 +127,60 @@ if [[ -z "${SLURM_NODELIST:-}" ]]; then
     exit 2
 fi
 
-# Get all node hostnames (sorted, as needed)
-readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST")
+# Pure-bash expander for a SLURM/Spur nodelist. Handles comma-separated lists
+# and bracket forms like "prefix[01-04,06]" / "prefix[030,058]" (optional
+# suffix). Used as the portable fallback below.
+_expand_nodelist() {
+    local nl="$1"
+    local seg="" depth=0 i ch
+    local -a segs=()
+    for (( i=0; i<${#nl}; i++ )); do
+        ch="${nl:i:1}"
+        if [[ "$ch" == "[" ]]; then depth=$((depth+1)); seg+="$ch"
+        elif [[ "$ch" == "]" ]]; then depth=$((depth-1)); seg+="$ch"
+        elif [[ "$ch" == "," && $depth -eq 0 ]]; then segs+=("$seg"); seg=""
+        else seg+="$ch"; fi
+    done
+    [[ -n "$seg" ]] && segs+=("$seg")
+
+    local s pre body suf tok a b width n
+    for s in "${segs[@]}"; do
+        [[ -n "$s" ]] || continue
+        if [[ "$s" == *"["* ]]; then
+            pre="${s%%[*}"
+            body="${s#*[}"; body="${body%%]*}"
+            suf="${s#*]}"
+            local OLD_IFS="$IFS"; IFS=','
+            for tok in $body; do
+                if [[ "$tok" == *-* ]]; then
+                    a="${tok%%-*}"; b="${tok##*-}"; width=${#a}
+                    for (( n=10#$a; n<=10#$b; n++ )); do
+                        printf "%s%0*d%s\n" "$pre" "$width" "$n" "$suf"
+                    done
+                else
+                    printf "%s%s%s\n" "$pre" "$tok" "$suf"
+                fi
+            done
+            IFS="$OLD_IFS"
+        else
+            printf "%s\n" "$s"
+        fi
+    done
+}
+
+# Get all node hostnames. Prefer `scontrol show hostnames`, which correctly
+# expands compressed nodelists on stock Slurm. Some Slurm-compatible schedulers
+# (e.g. Spur) lack that subcommand, and CI containers / dev VMs may lack the
+# client entirely -- in both cases fall back to the pure-bash expander above so
+# every launcher behaves consistently, including bracket-range expansion.
+NODE_ARRAY=()
+if command -v scontrol >/dev/null 2>&1; then
+    readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST" 2>/dev/null || true)
+fi
+if [[ ${#NODE_ARRAY[@]} -eq 0 ]]; then
+    LOG_WARN "[slurm-entry] 'scontrol show hostnames' unavailable; expanding SLURM_NODELIST locally"
+    readarray -t NODE_ARRAY < <(_expand_nodelist "$SLURM_NODELIST")
+fi
 SLURM_MASTER_ADDR="${NODE_ARRAY[0]:-}"
 if [[ -z "$SLURM_MASTER_ADDR" ]]; then
     LOG_ERROR "[slurm-entry] Failed to resolve the first host from SLURM_NODELIST=$SLURM_NODELIST"
@@ -128,7 +189,11 @@ fi
 
 if [[ -z "${MASTER_ADDR:-}" ]]; then
     MASTER_ADDR="$SLURM_MASTER_ADDR"
-elif [[ "$MASTER_ADDR" != "$SLURM_MASTER_ADDR" ]]; then
+elif [[ "${MASTER_ADDR%%.*}" != "${SLURM_MASTER_ADDR%%.*}" ]]; then
+    # Compare only the leading hostname label so a scheduler-provided FQDN
+    # (e.g. Spur exports MASTER_ADDR=node.crusoe.amd.com) still matches the
+    # short name produced by nodelist expansion. A genuine mismatch (a
+    # different node entirely) is still caught.
     LOG_ERROR "[slurm-entry] MASTER_ADDR must match the first host in SLURM_NODELIST."
     LOG_ERROR "[slurm-entry] MASTER_ADDR=$MASTER_ADDR, expected=$SLURM_MASTER_ADDR"
     exit 2
@@ -156,8 +221,20 @@ validate_distributed_params || LOG_WARN "[slurm-entry] Failed to validate distri
 
 # ------------- Dispatch based on mode ---------------
 
-# Parse mode (default: container)
+# Strip the entry-mode separator if present.
+[[ "${1:-}" == "--" ]] && shift
+
+# Parse entry mode (default: container). Supported keywords are the ones
+# advertised by primus-cli-slurm.sh: container | direct. Anything else is
+# treated as a primus subcommand and routed through the default container
+# chain (preserves the documented terse forms like `-- preflight`).
+ENTRY_MODE="container"
+if [[ "${1:-}" == "container" || "${1:-}" == "direct" ]]; then
+    ENTRY_MODE="$1"
+    shift
+fi
 [[ "${1:-}" == "--" ]] && shift
+LOG_INFO_RANK0 "[slurm-entry] Entry mode: $ENTRY_MODE"
 
 # Build arguments based on mode
 SCRIPT_ARGS=()
@@ -175,8 +252,10 @@ SCRIPT_ARGS+=(
     --env "GPUS_PER_NODE=$GPUS_PER_NODE"
 )
 
-# Build script path (container mode only)
-script_path="$RUNNER_DIR/primus-cli-container.sh"
+# Dispatch to the chosen entry script. Both primus-cli-container.sh and
+# primus-cli-direct.sh already accept --env / --config / --debug, so no
+# downstream changes are needed.
+script_path="$RUNNER_DIR/primus-cli-${ENTRY_MODE}.sh"
 require_file "$script_path" "[slurm-entry] Script not found: $script_path"
 
 # Build full command
diff --git a/runner/primus-cli-slurm.sh b/runner/primus-cli-slurm.sh
index b98d5c1e2..d476fd8a3 100755
--- a/runner/primus-cli-slurm.sh
+++ b/runner/primus-cli-slurm.sh
@@ -17,8 +17,11 @@ Usage:
 Description:
     Launch distributed Primus jobs via Slurm.
     - Everything before the first '--' is passed to Slurm (srun/sbatch and flags).
-    - <entry> specifies Primus execution mode: container | direct | preflight (see below).
+    - <entry> specifies Primus execution mode: container | direct  (default: container).
     - The second '--' (if any) separates Primus entry args from Primus CLI arguments.
+    - If <entry> is anything other than 'container' or 'direct' (e.g. a primus
+      subcommand like 'preflight'), it is routed through the default container
+      chain unchanged.
 
 Options:
     --config FILE    Load configuration from specified file
@@ -32,8 +35,14 @@ Examples:
     # Launch with sbatch, log to file, run benchmark
     primus-cli slurm sbatch --output=run.log -N 2 -- container -- benchmark gemm -M 4096 -N 4096 -K 4096
 
-    # Run preflight environment check across 4 nodes
-    primus-cli slurm srun -N 4 -- preflight
+    # Run preflight environment check across 4 nodes (terse form: defaults to container entry)
+    primus-cli slurm srun -N 4 -- preflight --quick
+
+    # Run preflight via the direct entry (no container; uses host python or VENV_ACTIVATE)
+    primus-cli slurm srun -N 4 -- direct -- preflight --quick
+
+    # Run node smoke test via direct entry (auto-selects single mode for node_smoke)
+    primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf
 
     # Dry-run to see what would be executed
     primus-cli slurm --dry-run srun -N 4 -- container -- train
@@ -285,7 +294,37 @@ fi
 ENTRY="$RUNNER_DIR/primus-cli-slurm-entry.sh"
 require_file "$ENTRY" "[slurm] Entry script not found: $ENTRY"
 
-# Build full command
+# On the sbatch path the scheduler copies the batch script into its spool dir, so
+# the entry cannot locate its own lib/ from $0. Hand it the real runner directory
+# (propagated by sbatch --export=ALL, which is the default on both flavors).
+export PRIMUS_RUNNER_DIR="$RUNNER_DIR"
+
+# Build full command. Spur is detected by the presence of the `spur` command; its
+# srun needs an explicit task count -- see below.
+IS_SPUR=false
+command -v spur >/dev/null 2>&1 && IS_SPUR=true
+
+# Standard Slurm's srun derives one task per node from -N, while spur's srun
+# defaults to --ntasks=1: `srun -N 4 ENTRY` dispatches a single task, so only one
+# node runs the entry and the other ranks never start. Request one task per node
+# unless the caller already chose a task count. sbatch is unaffected (spur runs
+# the batch script on every allocated node regardless of --ntasks).
+if [[ "$IS_SPUR" == "true" && "$LAUNCH_CMD" == "srun" ]]; then
+    spur_nodes=""
+    spur_has_ntasks=false
+    for ((i = 0; i < ${#SLURM_FLAGS[@]}; i++)); do
+        case "${SLURM_FLAGS[i]}" in
+            -N|--nodes) spur_nodes="${SLURM_FLAGS[i + 1]:-}" ;;
+            --nodes=*) spur_nodes="${SLURM_FLAGS[i]#--nodes=}" ;;
+            -n|--ntasks|--ntasks=*) spur_has_ntasks=true ;;
+        esac
+    done
+    if [[ "$spur_has_ntasks" == "false" && "$spur_nodes" =~ ^[0-9]+$ ]]; then
+        SLURM_FLAGS+=(-n "$spur_nodes")
+        LOG_INFO "[slurm] Spur srun: requesting one task per node (-n $spur_nodes)"
+    fi
+fi
+
 CMD=("$LAUNCH_CMD" "${SLURM_FLAGS[@]}" "$ENTRY" "${ENTRY_ARGS[@]}" -- "$@")
 
 # Display command
diff --git a/skills/backend-gap-report/SKILL.md b/skills/backend-gap-report/SKILL.md
index 4134be929..6298ddb0f 100644
--- a/skills/backend-gap-report/SKILL.md
+++ b/skills/backend-gap-report/SKILL.md
@@ -306,6 +306,14 @@ python3 tools/backend_gap_report/build_site_bundle.py --output-dir /tmp/primus-d
   (missing cadence support, a rendering bug, a schema evolution) or when the
   user explicitly requests a shell change.
 
+## Data location
+
+Backend-gap report data (the `docs/backend-gap/...` files this skill produces)
+belongs to the same `dashboard-data` data plane as the weekly/monthly reports.
+This skill only generates the files under that layout; the branch, commit, and
+deploy orchestration are owned by the caller (the automation's Agent
+Instructions), not this skill.
+
 ## Additional Resources
 
 - Schema, metadata fields, and path conventions: [reference.md](reference.md)
diff --git a/skills/backend-gap-report/examples.md b/skills/backend-gap-report/examples.md
index 047003287..1def34df7 100644
--- a/skills/backend-gap-report/examples.md
+++ b/skills/backend-gap-report/examples.md
@@ -44,7 +44,7 @@ Suggested evidence sources:
 - `third_party/Megatron-LM/requirements*`
 - `third_party/Megatron-LM/.github/workflows/*`
 - `primus/backends/megatron/*`
-- `primus/modules/trainer/megatron/*`
+- `primus/backends/megatron/pretrainer/*`
 - `docs/backends/megatron/patch-notes.md`
 
 Expected outputs:
diff --git a/skills/backend-patch-explorer/SKILL.md b/skills/backend-patch-explorer/SKILL.md
index 638cc68cf..cfae738df 100644
--- a/skills/backend-patch-explorer/SKILL.md
+++ b/skills/backend-patch-explorer/SKILL.md
@@ -1,6 +1,6 @@
 ---
 name: backend-patch-explorer
-description: Inventory and explain the patch (monkey-patch) optimizations Primus layers over upstream training backends such as Megatron-LM, TorchTitan, and MaxText, by reading the current repository code only. Use when the user asks which patches a backend has, wants a customer-facing patch table, asks how a specific patch works (for example deepep or DeepEP), or wants guidance to port a Primus patch into their own upstream framework. Read-only; no training or cluster commands.
+description: Inventory and explain the patch (monkey-patch) optimizations Primus layers over upstream training backends such as Megatron-LM, TorchTitan, and MaxText, including their version compatibility, dependencies, and Primus-Turbo integration details, by reading the current repository code only. Use when the user asks which patches a backend has, wants a customer-facing patch table, asks how a specific patch or Primus-Turbo feature works (for example deepep, turbo attention, or FP8), or wants a reference guide to port a Primus optimization into their own upstream framework. Read-only; no training or cluster commands.
 ---
 
 # Backend Patch Explorer
@@ -10,7 +10,7 @@ Primus does not fork its upstream training backends (Megatron-LM, TorchTitan, Ma
 Read the **current code only** (no training, no cluster commands) to serve two kinds of request:
 
 - **Inventory** - list every patch of one backend as a customer-facing table.
-- **Deep-dive + porting** - for one patch, explain what it does and how to port it into the user's own upstream framework.
+- **Deep-dive + porting** - for one patch or Primus-Turbo feature, explain what it does and produce a reference guide (Primus-side plus Turbo-side integration detail) for porting it into the user's own framework version. This is reference guidance, not a guarantee of a successful port.
 
 The set of patches changes over time, so derive every answer from the code as it exists now and re-scan on each run instead of trusting an earlier list.
 
@@ -77,13 +77,15 @@ rg -n "= Primus|sys\.modules\[|functools\.partial\(|inspect\.getsource" \
 #          primus/backends/<backend>/patches --include=*.py
 ```
 
-For every patch, capture the five inventory columns:
+For every patch, capture the six inventory columns:
 
 - **Name / id** - first positional arg of `@register_patch`.
 - **What it does** - the `description=` text, else the handler docstring.
-- **Patch logic** - the upstream symbol the body replaces or wraps; match it to one of the patch styles above (look for `module.Symbol = PrimusX`, `orig = module.fn` + wrapper, `sys.modules[...]= `, `functools.partial(...)`, or `inspect.getsource()` + `exec()`), and follow the handler's imports to the real implementation (often under `primus/backends/<backend>/core/extensions/`, `.../core/transformer/`, or `.../core/pipeline_parallel/`). If no symbol is replaced, say so per the styles above.
+- **Patch logic** - the upstream symbol the body replaces or wraps; match it to one of the patch styles above (look for `module.Symbol = PrimusX`, `orig = module.fn` + wrapper, `sys.modules[...]= `, `functools.partial(...)`, or `inspect.getsource()` + `exec()`), and follow the handler's imports to the real implementation (often under `primus/backends/<backend>/core/extensions/`, `.../core/transformer/`, or `.../core/pipeline_parallel/`). A handler often rebinds the SAME symbol in more than one module (e.g. both the defining module and a `*_layer_specs` module); capture every `X.Symbol = PrimusX` line in the body, not just the first. If no symbol is replaced, say so per the styles above.
 - **Code location** - the patch file, plus the real implementation file when the logic lives elsewhere.
 - **How to enable** - the config key checked in `condition`, any version gate, and the `phase`. `condition=` is often a NAMED helper rather than an inline lambda, so follow it: e.g. megatron `turbo/` conditions funnel into `turbo/utils.py::is_primus_turbo_can_patch` (`enable_primus_turbo` AND `tp_size == 1` AND the `primus_turbo` package importable), so a single `use_turbo_*` switch is necessary but not sufficient. No `condition` means "always". Optionally cross-check `primus/configs/modules/<backend>/*.yaml`, but that schema lists only Primus-added keys; flags inherited from upstream Megatron argparse will not appear there.
+- **Requires (version / deps)** - version compatibility from `backend_versions` / `primus_versions` and any version gate inside the condition (e.g. a Transformer Engine version check), plus the runtime packages the handler or its implementation imports (e.g. `primus_turbo`, a Transformer Engine version, ROCm-only paths). Follow the handler and implementation imports to confirm.
+- **Implicit constraints** - real constraints often live in the implementation, not the `condition`. Grep the impl file and its helpers for `assert` / `raise` (e.g. `assert tp_size == 1`, unsupported-flag guards) and record them; they gate the port even though the `condition` never mentions them.
 
 State only what the current code confirms; never invent a flag, patch, or upstream symbol.
 
@@ -92,7 +94,7 @@ State only what the current code confirms; never invent a flag, patch, or upstre
 For "how do I use feature X" requests (deepep is one such example), locate it live instead of recalling it:
 
 ```bash
-rg -ni "<keyword>" primus/backends primus/configs   # e.g. deepep, fp8, zero_bubble, muon
+rg -ni "<keyword>" primus/backends primus/configs   # e.g. deepep, fp8, zero_bubble
 # no rg: grep -rni "<keyword>" primus/backends primus/configs --include=*.py
 ```
 
@@ -101,19 +103,36 @@ Follow the hits to the registering `*_patches.py`, the imported implementation c
 ### Step 5 - Accuracy
 
 - Assert only facts confirmed in the current code; re-scan every run.
+- If a handler docstring names different flags or behavior than its `condition=` (docstrings drift), the code wins: report what the condition actually checks and flag the mismatch explicitly.
 - Optional cross-check: if `docs/backends/<backend>/patch-notes.md` exists, read it as a hint only. It is usually partial (covers a subset of patches, keyed by config flag rather than patch id) and can drift from the code (stale symbol attribution, typos). The code always wins.
 
+## Reaching into Primus-Turbo (dual-side detail)
+
+Many optimizations (attention, FP8, grouped MLP, communication overlap, deepep) are Primus-Turbo integrations. The patch and its extension wrapper only *bind* a Turbo class or function; the real kernel lives in the external `primus_turbo` package, which is a dependency and is usually NOT vendored in this repo. Users mostly want to bring Turbo into their own framework, so cover both sides:
+
+- **Primus side** (always in this repo): the patch installer plus the extension wrapper it imports (commonly under `primus/backends/<backend>/core/extensions/`). Record the adapted interface (constructor args, forward signature) and the call site where the Turbo class is instantiated.
+- **Turbo side** (external): resolve the imported symbol to its origin. If the package is importable, locate and read it:
+
+```bash
+python3 -c "import primus_turbo, os; print(os.path.dirname(primus_turbo.__file__))" 2>/dev/null
+# then read the module that defines the imported class / function
+```
+
+  If `primus_turbo` is not installed and its source is not checked out, do not guess its internals. Document the integration boundary instead (the imported Turbo symbol, its expected tensor / dtype contract, and the required package version) and point the user to the Primus-Turbo repository for the kernel implementation.
+
+Read the extension / implementation files statically (open them with a file reader). Do NOT `import` such a module in a Python process that lacks `primus_turbo`: its top-level `import primus_turbo...` raises ImportError even if you only wanted to inspect a class definition.
+
 ## Output formats
 
-Produce one of the two shapes below. Fill the placeholders from the scan and keep the structure fixed.
+Produce one of the two shapes below. Fill the placeholders from the scan and keep the structure fixed. Write the output in English by default; use another language only if the user explicitly asks.
 
 ### Inventory table
 
 Group rows by the patch subdirectories that actually exist (as discovered in Step 2). Columns are fixed:
 
-| Patch | What it does | Patch logic (upstream symbol replaced/wrapped) | Code location | How to enable |
-|---|---|---|---|---|
-| `id` | `description` | `UpstreamModule.Symbol` -> Primus impl | `patch_file` (+ `impl_file`) | `flag` / `phase=...` |
+| Patch | What it does | Patch logic (upstream symbol replaced/wrapped) | Code location | How to enable | Requires (version / deps) |
+|---|---|---|---|---|---|
+| `id` | `description` | `UpstreamModule.Symbol` -> Primus impl | `patch_file` (+ `impl_file`) | `flag` / `phase=...` | `backend_versions`; `primus_turbo` / TE ver / ROCm |
 
 Close with a note that `PRIMUS_PATCHES` (`all` / `none` / `id1,id2`) is a global allow-list applied on top of each patch's own flag.
 
@@ -121,8 +140,11 @@ Close with a note that `PRIMUS_PATCHES` (`all` / `none` / `id1,id2`) is a global
 
 1. **What it does** - the optimization and when it applies.
 2. **Patch logic** - the exact upstream symbol replaced/wrapped, the Primus replacement, and the `phase` + `condition` gate.
-3. **Code location** - patch (installer) file + real implementation file + config schema file.
-4. **Porting into the user's own framework** - an ordered plan covering: which Primus files to copy (installer + implementation + helpers); which upstream symbol to rebind or wrap and where it is constructed / imported (so the patch lands before instantiation); required dependencies and new config keys; whether to monkey-patch directly or reuse Primus's `@register_patch` + `run_patches` engine; and the constraints copied verbatim from `condition` (version range, parallelism limits, ROCm-only, and so on).
+3. **Code location** - patch (installer) file + real implementation file + config schema file. For a Primus-Turbo feature, give both sides: the Primus-side wrapper / extension in this repo and the Turbo-side symbol it wraps (see Reaching into Primus-Turbo above).
+4. **Version compatibility and dependencies** - the backend / Primus version ranges the patch supports (`backend_versions` / `primus_versions`), version gates inside the condition (e.g. Transformer Engine), and the exact packages and versions the port needs (e.g. `primus_turbo`, TE, ROCm). State these explicitly so the user can check their own framework version before porting.
+5. **Porting into the user's own framework** - an ordered plan covering: which Primus files to copy (installer + implementation + helpers); which upstream symbol to rebind or wrap and where it is constructed / imported (so the patch lands before instantiation); required dependencies and new config keys; whether to monkey-patch directly or reuse Primus's `@register_patch` + `run_patches` engine; and the constraints copied verbatim from `condition` plus the implicit `assert` / `raise` constraints found in Step 3 (version range, parallelism limits, ROCm-only, and so on). Two things that are easy to miss and worth calling out per port:
+   - **Complete the copy list by tracing transitive imports** (implementation -> local wrapper -> external package): copy every in-repo file the implementation pulls in, and mark the external package (e.g. `primus_turbo`) as a dependency to install rather than a file to copy.
+   - **Check `from <module> import <Symbol>` aliases** in the target: rebinding `module.Symbol` does not affect a name already imported into another module (e.g. a `*_layer_specs` that did `from ... import Symbol`). Grep for such aliases and rebind each one.
 
 ## Saving the output
 
diff --git a/skills/port-validation-guide/SKILL.md b/skills/port-validation-guide/SKILL.md
new file mode 100644
index 000000000..449ef7f87
--- /dev/null
+++ b/skills/port-validation-guide/SKILL.md
@@ -0,0 +1,68 @@
+---
+name: port-validation-guide
+description: Generate a validation plan and test scaffolding to check a Primus optimization that has been ported into a user's own training framework, across correctness (numerical accuracy vs a reference), performance (throughput parity), and integration (no regressions). Use after porting a Primus or Primus-Turbo feature into Megatron-LM, TorchTitan, MaxText, or another framework when the user wants a structured way to verify the port. Produces plans and test skeletons only; it does not run tests, require a GPU, or guarantee a successful port.
+---
+
+# Port Validation Guide
+
+After a Primus optimization is ported into a user's own framework, this skill produces a **validation plan** and **test scaffolding** across three tiers: correctness (numerical accuracy), performance (throughput parity), and integration (no regressions).
+
+It generates plans and test skeletons only. It does not execute tests, does not need a GPU, and does not guarantee the port is correct - the user decides which tiers to run, sets the thresholds, and chooses whether to run end-to-end training.
+
+The tests must be driven by what the port actually changed, so start from the facts.
+
+**Use the backend-patch-explorer skill first** to identify, for the target feature: the replaced or wrapped upstream symbol, the Primus-side and Turbo-side implementation, the enable flags, the version and dependency requirements, and the constraints (e.g. TP/EP limits, ROCm-only, Transformer Engine version). Those facts decide what each tier checks and which constraints become test skip conditions.
+
+## Inputs to confirm with the user
+
+- Target framework and version, and where the ported feature is wired in (module / call site).
+- Whether they can run on GPU. Op-level correctness checks and all benchmarks need a device; without one, emit the scaffold plus run instructions and stop short of executing.
+- The acceptance thresholds they care about. Otherwise use the defaults in [references/test-templates.md](references/test-templates.md) as a starting point.
+
+## Workflow
+
+### Step 1 - Gather the port facts
+
+Run backend-patch-explorer for the feature. Capture the replaced symbol, both-side implementations, enable flags, version/deps, and constraints. These map directly into the tests below.
+
+### Step 2 - Choose the reference (oracle)
+
+Correctness needs something to compare against. Pick per feature:
+
+- Same framework, feature OFF (baseline) vs feature ON (ported), identical inputs and seeds - answers "did I change results".
+- Primus's own implementation as the golden reference (dump its tensors) when it can be run - answers "does my port match Primus".
+
+Decide comparison points: op-level (module forward / backward output) and/or step-level (loss, grad norm). Op-level is the strongest signal for a kernel or dispatcher swap.
+
+If the port is a full replacement of a compute-plus-communication path rather than a thin drop-in wrapper (e.g. a token dispatcher swap), op-level outputs may legitimately differ and `assert_close` becomes the wrong oracle. For those, prefer a step-level or end-to-end loss-curve comparison with statistical checks (mean / max relative error, loss within a band) instead of exact op-level equality.
+
+### Step 3 - Correctness (numerical accuracy)
+
+Scaffold a test that builds the module both ways, feeds identical fixed-seed inputs, and compares outputs with dtype-appropriate tolerances (see references/test-templates.md). For fp8 / low precision, compare against a higher-precision reference and assert relative plus max-abs error, not exact equality. Cover backward (gradients) when the patch touches it.
+
+Building the module both ways is the hardest part of the scaffold, so pull the exact construction facts from backend-patch-explorer: the required constructor / config fields and the full forward signature. Attention modules in particular need more than `hidden_states` (e.g. `rotary_pos_emb`, `attention_mask`, `packed_seq_params`), and their config needs feature-specific fields (e.g. MLA needs `q_lora_rank` / `kv_lora_rank` and the head dims). Put these in the scaffold TODOs so the user wires real inputs instead of guessing.
+
+### Step 4 - Performance (throughput parity)
+
+Scaffold an A/B harness: same config, measure tokens/s or step time with the feature OFF vs ON, warmup then timed iters, report speedup and peak memory. Parity means the ported path is at least baseline within a user-set margin - a relative check, not an absolute target. GPU-required: without a device, emit the harness and run instructions only.
+
+### Step 5 - Integration (no regressions)
+
+Select the user's existing tests that touch the changed area to run before and after. Add a smoke test: a few training steps run to completion, loss finite (no NaN / inf), checkpoint save/load roundtrip. Turn the constraints from Step 1 into skip conditions (e.g. skip when TP != 1).
+
+### Step 6 - Assemble
+
+Emit two things:
+
+1. A validation plan (markdown) covering the three tiers, the chosen reference, thresholds, constraints, and which tiers are runnable in the user's environment.
+2. Test scaffold files from references/test-templates.md, with TODO markers where the user must wire their framework-specific module construction, config, and data.
+
+Be explicit about what only the user can fill in; this skill cannot know their framework internals.
+
+## Output
+
+Deliver the plan inline. Write scaffold files under a path the user gives, else `./port-validation/`. Offer to save the plan as a document. Do not run any test, benchmark, or training. Write the plan and scaffolds in English by default; use another language only if the user explicitly asks.
+
+## Additional resources
+
+- Test skeletons and default tolerance / benchmark settings: [references/test-templates.md](references/test-templates.md)
diff --git a/skills/port-validation-guide/references/test-templates.md b/skills/port-validation-guide/references/test-templates.md
new file mode 100644
index 000000000..ebcb0020b
--- /dev/null
+++ b/skills/port-validation-guide/references/test-templates.md
@@ -0,0 +1,150 @@
+# Port Validation - Test Templates and Defaults
+
+Framework-agnostic skeletons for the three validation tiers. They are starting
+points: every `TODO` must be wired to the user's own framework (module
+construction, config, data). Tolerances and iteration counts are defaults the
+user should tune for their hardware and precision.
+
+## Default tolerances (correctness)
+
+Compare the ported output against the reference with tolerances that match the
+compute dtype. These are starting values, not guarantees.
+
+| Compute dtype | rtol | atol | Notes |
+|---|---|---|---|
+| fp32 | 1e-5 | 1e-6 | near-exact |
+| tf32 | 1e-3 | 1e-4 | matmul in tf32 |
+| bf16 | 2e-2 | 1e-3 | wide mantissa error |
+| fp16 | 1e-2 | 1e-3 | |
+| fp8 | compare vs a bf16/fp32 reference | - | assert relative + max-abs error, never exact equality |
+
+For a kernel/dispatcher swap, also check backward: compare input grads (and
+weight grads) with the same tolerances.
+
+## Tier 1 - Correctness (op-level, PyTorch)
+
+```python
+# port-validation/test_correctness.py
+import torch
+
+SEED = 1234
+
+def _fixed_input():
+    torch.manual_seed(SEED)
+    # TODO: build the exact input tensors the module expects (shapes, dtype, device)
+    raise NotImplementedError
+
+def _build_reference():
+    # TODO: construct the module WITHOUT the ported feature (baseline),
+    #       or load Primus golden tensors, in eval mode with fixed seed.
+    raise NotImplementedError
+
+def _build_ported():
+    # TODO: construct the module WITH the ported feature enabled.
+    raise NotImplementedError
+
+def test_forward_matches_reference():
+    x = _fixed_input()
+    ref = _build_reference()
+    got = _build_ported()
+    with torch.no_grad():
+        out_ref = ref(x)
+        out_got = got(x)
+    # TODO: set rtol/atol from the dtype table above
+    torch.testing.assert_close(out_got, out_ref, rtol=2e-2, atol=1e-3)
+
+def test_backward_matches_reference():
+    x = _fixed_input().requires_grad_(True)
+    # TODO: run both paths, backward a scalar (e.g. out.sum()),
+    #       compare x.grad (and named weight grads) within tolerance.
+    raise NotImplementedError
+```
+
+For fp8 or other low precision, replace `assert_close` with an explicit
+relative-error check against a higher-precision reference:
+
+```python
+err = (out_got.float() - out_ref.float()).abs()
+rel = err / (out_ref.float().abs() + 1e-8)
+assert rel.mean().item() < 0.05        # TODO: user-set threshold
+assert err.max().item() < 1.0          # TODO: user-set threshold
+```
+
+## Tier 2 - Performance (throughput parity, PyTorch)
+
+```python
+# port-validation/bench_parity.py
+import time, torch
+
+WARMUP, ITERS = 10, 50   # TODO: tune
+
+def _run_once(build_fn):
+    model = build_fn()                 # TODO: build model/step closure
+    step = _make_step(model)           # TODO: one train/forward step closure
+    for _ in range(WARMUP):
+        step()
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    t0 = time.perf_counter()
+    for _ in range(ITERS):
+        step()
+    torch.cuda.synchronize()
+    dt = (time.perf_counter() - t0) / ITERS
+    peak = torch.cuda.max_memory_allocated() / 1e9
+    return dt, peak
+
+def main():
+    base_dt, base_mem = _run_once(_build_baseline)   # TODO
+    port_dt, port_mem = _run_once(_build_ported)     # TODO
+    speedup = base_dt / port_dt
+    print(f"baseline: {base_dt*1e3:.2f} ms/it, {base_mem:.2f} GB")
+    print(f"ported:   {port_dt*1e3:.2f} ms/it, {port_mem:.2f} GB")
+    print(f"speedup:  {speedup:.3f}x")
+    # Parity check is relative and user-defined:
+    MARGIN = 0.95   # TODO: ported must be >= baseline * MARGIN
+    assert speedup >= MARGIN, "ported path slower than baseline beyond margin"
+```
+
+Requires a GPU. If none is available, emit this harness and the run command
+(`python port-validation/bench_parity.py`) without executing it.
+
+## Tier 3 - Integration (no regressions)
+
+```python
+# port-validation/test_integration.py
+import math, torch, pytest
+
+# Skip conditions come from backend-patch-explorer constraints for the feature.
+# TODO: replace with the real constraint checks (parallelism, versions, ROCm).
+def _feature_supported():
+    return True
+
+pytestmark = pytest.mark.skipif(not _feature_supported(),
+                                reason="feature constraints not met")
+
+def test_smoke_train_a_few_steps():
+    # TODO: build a tiny model+optimizer+data with the feature ON.
+    losses = []
+    for _ in range(5):                 # TODO: a few steps
+        loss = _train_step()           # TODO
+        assert math.isfinite(loss), "loss is NaN/inf"
+        losses.append(loss)
+    assert losses[-1] <= losses[0] * 2.0   # loose sanity, not convergence
+
+def test_checkpoint_roundtrip(tmp_path):
+    # TODO: save then load; assert state_dict keys/shapes match and
+    #       a forward pass after load reproduces pre-save output.
+    raise NotImplementedError
+```
+
+Also list the user's existing tests that touch the changed module/path and run
+them before and after the port to confirm no regression. This skill only names
+which tests to run; the user runs them.
+
+## MaxText / JAX note
+
+For a JAX (MaxText) port, keep the same three tiers but swap primitives:
+`jnp`/`jax.numpy` with `jnp.allclose` (or `chex.assert_trees_all_close`) for
+correctness, `block_until_ready()` + `time.perf_counter()` for timing, and the
+project's `pytest` suite for integration. Sharding-annotation ports should also
+assert the expected sharding on outputs where applicable.
diff --git a/skills/slurm-training-node-validation/SKILL.md b/skills/slurm-training-node-validation/SKILL.md
index d83233bd5..95f62b75f 100644
--- a/skills/slurm-training-node-validation/SKILL.md
+++ b/skills/slurm-training-node-validation/SKILL.md
@@ -91,7 +91,7 @@ For each group `$i`, first compress the 4 hostnames into SLURM bracket notation,
 
 ```bash
 COMPRESSED_NODELIST=$(scontrol show hostlist "node1,node2,node3,node4")
-# e.g. "uswslocpm2m-106-[1628,1691,1697,1707]"
+# e.g. "node-[1628,1691,1697,1707]"
 
 GROUPDIR="${BASEDIR}/group_${i}"
 mkdir -p "$GROUPDIR"
@@ -100,7 +100,7 @@ mkdir -p "$GROUPDIR"
 echo "${COMPRESSED_NODELIST}" > "${GROUPDIR}/group_${i}_nodelist_${COMPRESSED_NODELIST}"
 ```
 
-**The `$COMPRESSED_NODELIST` must always be in SLURM prefix-compressed format** (e.g. `uswslocpm2m-106-[1628,1691,1697,1707]`), NOT individual hostnames joined by commas. Always use `scontrol show hostlist` to produce it.
+**The `$COMPRESSED_NODELIST` must always be in SLURM prefix-compressed format** (e.g. `node-[1628,1691,1697,1707]`), NOT individual hostnames joined by commas. Always use `scontrol show hostlist` to produce it.
 
 The resulting directory layout looks like:
 
@@ -108,11 +108,11 @@ The resulting directory layout looks like:
 ./output/slurm-training-node-validation-20260316_201333/
 ├── group_1/
 │   ├── log_group_1.txt                                          # training log
-│   ├── group_1_nodelist_uswslocpm2m-106-[1628,1691,1697,1707]  # content: the compressed nodelist
+│   ├── group_1_nodelist_node-[1628,1691,1697,1707]  # content: the compressed nodelist
 │   └── success  OR  failed                                      # result file (created after job finishes)
 ├── group_2/
 │   ├── log_group_2.txt
-│   ├── group_2_nodelist_uswslocpm2m-106-[1711,1713,1719,1723]
+│   ├── group_2_nodelist_node-[1711,1713,1719,1723]
 │   └── success  OR  failed
 └── group_3/
     ├── ...
@@ -156,7 +156,7 @@ For each group (index `$i`, nodelist `$NODELIST`), run:
 - Each group runs in a **subshell** so environment variables don't leak between groups.
 - Capture the PID of each background job: `PIDS[$i]=$!`
 - Log file: `log_group_1.txt`, `log_group_2.txt`, etc. (1-indexed)
-- The nodelist marker file name embeds the **compressed** nodelist (e.g. `group_1_nodelist_uswslocpm2m-106-[1628,1691,1697,1707]`). The file content is the compressed nodelist string itself.
+- The nodelist marker file name embeds the **compressed** nodelist (e.g. `group_1_nodelist_node-[1628,1691,1697,1707]`). The file content is the compressed nodelist string itself.
 
 After launching all groups, print:
 ```
@@ -315,10 +315,10 @@ FAILEOF
 - Output directory: <BASEDIR>
 
 ### Passed Nodelist (training-ready)
-<SLURM-format compressed nodelist, e.g. uswslocpm2m-106-[1628,1691,1697,1707,1711,1713,1719,1723]>
+<SLURM-format compressed nodelist, e.g. node-[1628,1691,1697,1707,1711,1713,1719,1723]>
 
 ### Failed Nodelist
-<SLURM-format compressed nodelist, e.g. uswslocpm2m-106-[1730,1735,1741,1749]>
+<SLURM-format compressed nodelist, e.g. node-[1730,1735,1741,1749]>
 ```
 
 #### Generating SLURM-format Compressed Nodelists
@@ -329,11 +329,11 @@ Collect all individual hostnames for passed (or failed) groups, then compress in
 scontrol show hostlist <comma-separated-hostnames>
 ```
 
-Example: if passed nodes are `uswslocpm2m-106-1628,uswslocpm2m-106-1691,uswslocpm2m-106-1697,uswslocpm2m-106-1707,uswslocpm2m-106-1711,uswslocpm2m-106-1713`, then:
+Example: if passed nodes are `node-1628,node-1691,node-1697,node-1707,node-1711,node-1713`, then:
 
 ```bash
-$ scontrol show hostlist uswslocpm2m-106-1628,uswslocpm2m-106-1691,uswslocpm2m-106-1697,uswslocpm2m-106-1707,uswslocpm2m-106-1711,uswslocpm2m-106-1713
-uswslocpm2m-106-[1628,1691,1697,1707,1711,1713]
+$ scontrol show hostlist node-1628,node-1691,node-1697,node-1707,node-1711,node-1713
+node-[1628,1691,1697,1707,1711,1713]
 ```
 
 Nodes with the **same prefix** are automatically merged into bracket notation by `scontrol show hostlist`. The output is the final SLURM-format nodelist to display in the summary.
@@ -390,6 +390,6 @@ The final directory layout includes:
 - Use `-o StrictHostKeyChecking=no` for any SSH operations.
 - The **exact success pattern** is `iteration        3` (8 spaces). Do not modify the spacing.
 - Output is stored under `./output/slurm-training-node-validation-<TIMESTAMP>/`. Each group has its own subdirectory with log, status file, and nodelist marker.
-- The **nodelist marker file** is an empty file whose name encodes the group's compressed nodelist (e.g. `group_1_nodelist_uswslocpm2m-106-[1628,1691,1697,1707]`). This makes it easy to identify which nodes belong to each group by listing the directory.
+- The **nodelist marker file** is an empty file whose name encodes the group's compressed nodelist (e.g. `group_1_nodelist_node-[1628,1691,1697,1707]`). This makes it easy to identify which nodes belong to each group by listing the directory.
 - The **status file** is either `success` (empty) or `failed` (contains error details and possible cause).
 - If a group has a non-zero exit code but the success pattern is found in the log, still create `success` (training may exit non-zero after completing).
diff --git a/skills/spur-cluster-status/SKILL.md b/skills/spur-cluster-status/SKILL.md
new file mode 100644
index 000000000..c6ee019df
--- /dev/null
+++ b/skills/spur-cluster-status/SKILL.md
@@ -0,0 +1,83 @@
+---
+name: spur-cluster-status
+description: Inspect the current Spur (AMD SLURM-compatible) cluster node-allocation state and produce a Markdown report covering the caller's account/QoS permissions, partitions and per-state node counts, per-QoS and per-account node usage, reservations, queue pressure, GPU capacity, and the caller's own jobs. Use when the user asks about Spur/SLURM cluster status, node allocation, which QoS/account/partition holds how many nodes, how many idle nodes are available, or wants a cluster snapshot report.
+---
+
+# Spur Cluster Status Report
+
+Generate a read-only snapshot of the Spur cluster's node-allocation state plus the
+caller's account/QoS permissions, and emit a single Markdown report (with an
+appendix of common commands). Spur is AMD's SLURM-compatible scheduler exposed via
+`sinfo` / `squeue` / `scontrol` / `spur accounts`.
+
+## Workflow
+
+### Step 1: Generate the report
+
+Run the collector from the repo root (it is read-only and takes a few seconds):
+
+```bash
+python3 .claude/skills/spur-cluster-status/scripts/spur_status.py
+```
+
+- Reports on the current user by default. To target another user: `--user <name>`.
+- The script prints a complete Markdown report to stdout and never mutates cluster state.
+
+### Step 2: Save and present
+
+1. Save the output under the repo root (create the dir if missing):
+
+```bash
+mkdir -p output/skills
+python3 .claude/skills/spur-cluster-status/scripts/spur_status.py \
+  > "output/skills/spur-cluster-status-$(date +%Y%m%d.%H%M).md"
+```
+
+2. Show the report to the user. Present the tables inline; the idle-node list can be
+   long, so summarize it (count + a short sample) unless the user wants the full list.
+3. Print the saved file path.
+
+### Step 3: Add insights (optional)
+
+After the tables, add a short analysis when relevant, e.g.:
+- Whether enough idle nodes exist for the user's target job size.
+- Whether the user's jobs are on shared (`mix`) nodes (GPU-contention risk) vs `--exclusive`.
+- Whether a low-priority QoS (e.g. `amd-burst-qos`, Prio=1) is being used for a job that
+  should use the account's normal QoS.
+
+## Report contents
+
+The collector emits these sections (see `scripts/spur_status.py`). The report is
+written in **English**:
+
+1. **Account / QoS (you)** — the caller's account(s), default account, default QoS, the
+   permission-model note, and the global QoS list with priority/preempt.
+2. **Partitions & Node States** — per-partition total + per-state node counts, overall
+   state breakdown, utilization %, GPU capacity, idle nodelist, and drain/down nodes.
+3. **Node Usage by QoS** — for running jobs, per QoS: #jobs, distinct nodes, node-job
+   slots, GPUs, plus pending jobs and pending node demand.
+4. **Node Usage by Account** — same breakdown grouped by account.
+5. **Reservations** — name, node count, users, end time.
+6. **Queue Pressure** — running vs pending job counts and pending node demand.
+7. **My Jobs** — the caller's running/pending jobs.
+8. **Jobs by User (all users)** — one row per user: total jobs (running/pending),
+   distinct nodes held, number of accounts + which, which QoS, and which partitions.
+9. **Appendix: Common Commands** — the command reference used to build the report.
+
+## Spur quirks (important)
+
+These are baked into the collector; keep them in mind if you extend it:
+
+- `-o` format strings render **space-separated** regardless of literal delimiters — parse by whitespace column, not by a custom separator.
+- `spur accounts show user|account <name>` ignores the positional filter and lists everything — filter with `grep`/`awk`.
+- `sacctmgr` policy queries are blocked ("Please ask your administrator"); use `spur accounts show qos` for the QoS list.
+- There is **no** `scontrol show hostnames`; expand compressed nodelists locally (the script does this).
+- The association table is empty and **QoS is not enforced per account** — only the **account** is enforced at submit time. QoS selects scheduling priority, not permission.
+
+## Extending
+
+To add a metric, add a parser + a section in `build_report()` in
+`scripts/spur_status.py`, and add the underlying command to `COMMANDS_APPENDIX`
+so the report's command list stays in sync. Candidate additions: top users by node
+count, largest contiguous idle block for N-node jobs, `spur report cluster` historical
+utilization, and per-node `CPUAlloc` from `scontrol show node`.
diff --git a/skills/spur-cluster-status/scripts/spur_status.py b/skills/spur-cluster-status/scripts/spur_status.py
new file mode 100644
index 000000000..a8160925f
--- /dev/null
+++ b/skills/spur-cluster-status/scripts/spur_status.py
@@ -0,0 +1,477 @@
+#!/usr/bin/env python3
+"""Collect Spur cluster node-allocation status and print a Markdown report.
+
+Targets the AMD "Spur" scheduler (SLURM-compatible CLI: sinfo/squeue/scontrol/
+spur accounts). Read-only. Handles Spur quirks: `-o` delimiters are rendered as
+spaces, positional filters are ignored, there is no `scontrol show hostnames`,
+and QoS is not enforced per account (only the account is).
+
+Usage:
+    python3 spur_status.py [--user USER]
+
+The report is printed to stdout (redirect it to save a file).
+"""
+
+import argparse
+import getpass
+import re
+import subprocess
+from collections import defaultdict
+from datetime import datetime
+
+# Node-state buckets used for the state breakdown / utilization math.
+BUSY_STATES = {"alloc", "mix"}
+FREE_STATES = {"idle"}
+DOWN_STATES = {"down", "drain", "drained", "fail", "failing", "unknown"}
+
+
+def run(cmd):
+    """Run a command list; return stdout text, or '' on any failure."""
+    try:
+        out = subprocess.run(cmd, capture_output=True, text=True, timeout=60, check=False)
+        return out.stdout or ""
+    except Exception:
+        return ""
+
+
+# ---------------------------------------------------------------------------
+# Hostlist expansion (no `scontrol show hostnames` on Spur, so do it locally).
+# Handles: "p-030", "p-[036,089,110]", "p-[036-040]", and top-level commas.
+# ---------------------------------------------------------------------------
+def _split_top(s):
+    parts, depth, cur = [], 0, ""
+    for ch in s:
+        if ch == "[":
+            depth += 1
+            cur += ch
+        elif ch == "]":
+            depth -= 1
+            cur += ch
+        elif ch == "," and depth == 0:
+            parts.append(cur)
+            cur = ""
+        else:
+            cur += ch
+    if cur:
+        parts.append(cur)
+    return parts
+
+
+def _expand_one(tok):
+    m = re.match(r"^(.*?)\[([^\]]*)\](.*)$", tok)
+    if not m:
+        return [tok] if tok else []
+    pre, body, post = m.groups()
+    out = []
+    for item in body.split(","):
+        item = item.strip()
+        if "-" in item:
+            a, b = item.split("-", 1)
+            width = len(a)
+            for i in range(int(a), int(b) + 1):
+                out.append(f"{pre}{str(i).zfill(width)}{post}")
+        elif item:
+            out.append(f"{pre}{item}{post}")
+    return out
+
+
+def expand_hostlist(s):
+    if not s or s in ("(null)", "N/A", "-"):
+        return []
+    res = []
+    for tok in _split_top(s):
+        res.extend(_expand_one(tok))
+    return res
+
+
+# ---------------------------------------------------------------------------
+# Parsers
+# ---------------------------------------------------------------------------
+def parse_user_accounts(user):
+    """Return (accounts, default_acct, def_qos) for the given user."""
+    text = run(["spur", "accounts", "show", "user"])
+    accounts, default_acct, def_qos = [], "", ""
+    for line in text.splitlines():
+        toks = line.split()
+        if len(toks) < 2 or toks[0] in ("User", "----"):
+            continue
+        if toks[0] != user:
+            continue
+        accounts.append(toks[1])
+        if len(toks) >= 4:
+            default_acct = toks[3]
+        if len(toks) >= 5:
+            def_qos = toks[4]
+    return sorted(set(accounts)), default_acct, def_qos
+
+
+def parse_qos():
+    """Return list of dicts: {name, prio, preempt}."""
+    text = run(["spur", "accounts", "show", "qos"])
+    rows = []
+    for line in text.splitlines():
+        toks = line.split()
+        if len(toks) < 2 or toks[0] in ("Name", "----"):
+            continue
+        rows.append(
+            {
+                "name": toks[0],
+                "prio": toks[1] if len(toks) > 1 else "",
+                "preempt": toks[2] if len(toks) > 2 else "",
+            }
+        )
+    return rows
+
+
+def parse_nodes():
+    """Return dict node -> {partition, state, gpus} from `sinfo -N`."""
+    text = run(["sinfo", "-N", "-h", "-o", "%N %P %t %G"])
+    nodes = {}
+    for line in text.splitlines():
+        toks = line.split()
+        if len(toks) < 3:
+            continue
+        name, part, state = toks[0], toks[1], toks[2].rstrip("*")
+        gres = toks[3] if len(toks) > 3 else ""
+        gpus = gres.count("gpu:")
+        nodes[name] = {"partition": part, "state": state, "gpus": gpus}
+    return nodes
+
+
+def parse_jobs(nodes):
+    """Return list of job dicts from squeue.
+
+    Spur renders `-o` fields space-separated and DROPS empty fields (e.g. a job
+    with no QoS), which shifts columns and breaks naive positional parsing. So we
+    only trust the always-present leading fields (jobid, partition, account, user,
+    state, nnodes) and disambiguate the trailing optional tokens (qos and/or
+    nodelist) against the known node set: the token that expands to a real node is
+    the nodelist; the other is the qos. `account` is always present because the
+    scheduler enforces it at submit time.
+    """
+    text = run(["squeue", "-h", "-o", "%i %P %a %u %T %D %q %N"])
+    jobs = []
+    for line in text.splitlines():
+        toks = line.split()
+        if len(toks) < 6:
+            continue
+        jobid, partition, account, user, state = toks[:5]
+        nnodes = int(toks[5]) if toks[5].isdigit() else 0
+        qos, nodelist = "", ""
+        for t in toks[6:]:
+            expanded = expand_hostlist(t)
+            if expanded and expanded[0] in nodes:
+                nodelist = t
+            else:
+                qos = t
+        jobs.append(
+            {
+                "jobid": jobid,
+                "partition": partition,
+                "qos": qos,
+                "account": account,
+                "user": user,
+                "state": state,
+                "nnodes": nnodes,
+                "nodelist": nodelist,
+            }
+        )
+    return jobs
+
+
+def parse_controller():
+    """Extract the controller address from `scontrol show config`."""
+    for line in run(["scontrol", "show", "config"]).splitlines():
+        if "Addr=" in line:
+            return line.split("Addr=", 1)[1].strip()
+    return "n/a"
+
+
+def parse_reservations():
+    """Return list of reservation dicts from scontrol show reservation."""
+    text = run(["scontrol", "show", "reservation"])
+    resvs, cur = [], {}
+    for line in text.splitlines():
+        if line.startswith("ReservationName="):
+            if cur:
+                resvs.append(cur)
+            cur = {}
+        for m in re.finditer(r"(\w+)=([^\s]+)", line):
+            cur[m.group(1)] = m.group(2)
+    if cur:
+        resvs.append(cur)
+    return resvs
+
+
+# ---------------------------------------------------------------------------
+# Report
+# ---------------------------------------------------------------------------
+def group_usage(jobs, nodes, key, state):
+    """Aggregate running/pending jobs by a key (qos/account).
+
+    Returns dict key -> {jobs, node_slots, distinct_nodes(set), gpus}.
+    distinct_nodes/gpus are only meaningful for RUNNING jobs.
+    """
+    agg = defaultdict(lambda: {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0})
+    for j in jobs:
+        if j["state"] != state:
+            continue
+        k = j[key] or "(none)"
+        agg[k]["jobs"] += 1
+        agg[k]["node_slots"] += j["nnodes"]
+        for n in expand_hostlist(j["nodelist"]):
+            if n in nodes:
+                agg[k]["nodes"].add(n)
+    for k, v in agg.items():
+        v["gpus"] = sum(nodes[n]["gpus"] for n in v["nodes"])
+    return agg
+
+
+def h(title):
+    return f"\n## {title}\n"
+
+
+def group_by_user(jobs, nodes):
+    """Aggregate every job per user across all states."""
+    agg = defaultdict(
+        lambda: {
+            "total": 0,
+            "running": 0,
+            "pending": 0,
+            "nodes": set(),
+            "accounts": set(),
+            "qos": set(),
+            "partitions": set(),
+        }
+    )
+    for j in jobs:
+        a = agg[j["user"]]
+        a["total"] += 1
+        if j["state"] == "RUNNING":
+            a["running"] += 1
+        elif j["state"] == "PENDING":
+            a["pending"] += 1
+        a["accounts"].add(j["account"])
+        a["qos"].add(j["qos"])
+        a["partitions"].add(j["partition"])
+        for n in expand_hostlist(j["nodelist"]):
+            if n in nodes:
+                a["nodes"].add(n)
+    return agg
+
+
+def build_report(user):
+    now = datetime.now().strftime("%Y-%m-%d %H:%M")
+    accounts, default_acct, def_qos = parse_user_accounts(user)
+    qos_rows = parse_qos()
+    nodes = parse_nodes()
+    jobs = parse_jobs(nodes)
+    resvs = parse_reservations()
+
+    L = []
+    L.append("# Spur Cluster Node-Allocation Report")
+    L.append("")
+    L.append(f"- **Generated**: {now}")
+    L.append(f"- **Cluster**: spur  (controller: {parse_controller()})")
+    L.append(f"- **Report user**: {user}")
+
+    # 1. Account / QoS
+    L.append(h("1. Account / QoS (you)"))
+    L.append(f"- **Accounts**: {', '.join(accounts) if accounts else '(none found)'}")
+    L.append(f"- **Default account**: {default_acct or '-'}")
+    L.append(f"- **Default QoS**: {def_qos or '(unset)'}")
+    L.append("")
+    L.append(
+        "> Permission model (verified): the **account is enforced** - you can only "
+        "submit under the accounts listed above (submitting under another account "
+        "fails with `user ... is not associated with account`). **QoS is NOT restricted "
+        "per account**; any QoS from the global list below is accepted, so QoS selects "
+        "scheduling priority, not permission."
+    )
+    L.append("")
+    L.append("Global QoS (higher Prio = higher priority):")
+    L.append("")
+    L.append("| QoS | Prio | Preempt |")
+    L.append("|-----|------|---------|")
+    for q in sorted(qos_rows, key=lambda x: -int(x["prio"]) if x["prio"].isdigit() else 0):
+        L.append(f"| {q['name']} | {q['prio']} | {q['preempt']} |")
+
+    # 2. Partitions & node states
+    L.append(h("2. Partitions & Node States"))
+    part_state = defaultdict(lambda: defaultdict(int))
+    part_total = defaultdict(int)
+    state_total = defaultdict(int)
+    gpu_total = gpu_free = 0
+    for n, info in nodes.items():
+        part_state[info["partition"]][info["state"]] += 1
+        part_total[info["partition"]] += 1
+        state_total[info["state"]] += 1
+        gpu_total += info["gpus"]
+        if info["state"] in FREE_STATES:
+            gpu_free += info["gpus"]
+    all_states = sorted(state_total)
+    L.append("| Partition | Total | " + " | ".join(all_states) + " |")
+    L.append("|" + "---|" * (len(all_states) + 2))
+    for p in sorted(part_total):
+        cells = " | ".join(str(part_state[p].get(s, 0)) for s in all_states)
+        L.append(f"| {p} | {part_total[p]} | {cells} |")
+    total_nodes = sum(part_total.values())
+    busy = sum(state_total.get(s, 0) for s in BUSY_STATES)
+    L.append("")
+    L.append(f"- **Total nodes**: {total_nodes}")
+    L.append("- **State breakdown**: " + ", ".join(f"{s}={state_total[s]}" for s in all_states))
+    if total_nodes:
+        L.append(f"- **Utilization (alloc+mix)**: {busy}/{total_nodes} = {busy*100//total_nodes}%")
+    L.append(f"- **GPU capacity (whole-node)**: {gpu_total} total, ~{gpu_free} free (on idle nodes)")
+
+    # idle / unhealthy nodelist
+    idle_nodes = sorted(n for n, i in nodes.items() if i["state"] in FREE_STATES)
+    bad_nodes = sorted(n for n, i in nodes.items() if i["state"] in DOWN_STATES)
+    L.append("")
+    L.append(f"- **Idle nodes ({len(idle_nodes)})**: `{','.join(idle_nodes) if idle_nodes else 'none'}`")
+    if bad_nodes:
+        L.append(f"- **Drain/down nodes ({len(bad_nodes)})**: `{','.join(bad_nodes)}`")
+
+    # 3. Per-QoS usage
+    L.append(h("3. Node Usage by QoS (running jobs)"))
+    L.append(
+        "> Nodes can be shared (mix), so the sum of per-QoS distinct nodes may exceed the number of physically busy nodes."
+    )
+    L.append("")
+    qos_run = group_usage(jobs, nodes, "qos", "RUNNING")
+    qos_pend = group_usage(jobs, nodes, "qos", "PENDING")
+    L.append(
+        "| QoS | Running jobs | Distinct nodes | Node-job slots | GPUs | Pending jobs | Pending nodes req |"
+    )
+    L.append("|-----|------|------|------|------|------|------|")
+    for k in sorted(
+        set(qos_run) | set(qos_pend), key=lambda x: -len(qos_run.get(x, {"nodes": set()})["nodes"])
+    ):
+        r = qos_run.get(k, {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0})
+        p = qos_pend.get(k, {"jobs": 0, "node_slots": 0})
+        L.append(
+            f"| {k} | {r['jobs']} | {len(r['nodes'])} | {r['node_slots']} | {r['gpus']} | {p['jobs']} | {p['node_slots']} |"
+        )
+
+    # 4. Per-account usage
+    L.append(h("4. Node Usage by Account (running jobs)"))
+    acc_run = group_usage(jobs, nodes, "account", "RUNNING")
+    acc_pend = group_usage(jobs, nodes, "account", "PENDING")
+    L.append(
+        "| Account | Running jobs | Distinct nodes | Node-job slots | GPUs | Pending jobs | Pending nodes req |"
+    )
+    L.append("|-----|------|------|------|------|------|------|")
+    for k in sorted(
+        set(acc_run) | set(acc_pend), key=lambda x: -len(acc_run.get(x, {"nodes": set()})["nodes"])
+    ):
+        r = acc_run.get(k, {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0})
+        p = acc_pend.get(k, {"jobs": 0, "node_slots": 0})
+        L.append(
+            f"| {k} | {r['jobs']} | {len(r['nodes'])} | {r['node_slots']} | {r['gpus']} | {p['jobs']} | {p['node_slots']} |"
+        )
+
+    # 5. Reservations
+    L.append(h("5. Reservations"))
+    if resvs:
+        L.append("| Name | Nodes | Users | End time |")
+        L.append("|-----|------|------|------|")
+        for r in resvs:
+            nlist = expand_hostlist(r.get("Nodes", ""))
+            L.append(
+                f"| {r.get('ReservationName','?')} | {len(nlist)} | {r.get('Users','-')} | {r.get('EndTime','-')} |"
+            )
+    else:
+        L.append("No active reservations.")
+
+    # 6. Queue pressure
+    L.append(h("6. Queue Pressure"))
+    pend = [j for j in jobs if j["state"] == "PENDING"]
+    run_j = [j for j in jobs if j["state"] == "RUNNING"]
+    L.append(f"- **Running jobs**: {len(run_j)}")
+    L.append(f"- **Pending jobs**: {len(pend)} (requesting {sum(j['nnodes'] for j in pend)} nodes total)")
+
+    # 7. My jobs
+    L.append(h("7. My Jobs (user=%s)" % user))
+    mine = [j for j in jobs if j["user"] == user]
+    if mine:
+        L.append("| JobID | QoS | Account | State | Nodes | Nodelist |")
+        L.append("|-----|------|------|------|------|------|")
+        for j in mine:
+            L.append(
+                f"| {j['jobid']} | {j['qos']} | {j['account']} | {j['state']} | {j['nnodes']} | {j['nodelist'] or '-'} |"
+            )
+    else:
+        L.append("No running/pending jobs.")
+
+    # 8. Jobs by user (all users)
+    L.append(h("8. Jobs by User (all users)"))
+    L.append(
+        "> One row per user. `Jobs` = total (R running / P pending); `Nodes` = distinct nodes held by running jobs."
+    )
+    L.append("")
+    by_user = group_by_user(jobs, nodes)
+    L.append("| User | Jobs (R/P) | Nodes | #Accts | Accounts | QoS | Partitions |")
+    L.append("|-----|------|------|------|------|------|------|")
+    for u in sorted(by_user, key=lambda x: (-len(by_user[x]["nodes"]), -by_user[x]["total"], x)):
+        a = by_user[u]
+        acct_disp = ",".join(sorted(a["accounts"]))
+        qos_disp = ",".join(sorted((q or "(none)") for q in a["qos"]))
+        part_disp = ",".join(sorted(a["partitions"]))
+        L.append(
+            f"| {u} | {a['total']} ({a['running']}R/{a['pending']}P) | {len(a['nodes'])} | "
+            f"{len(a['accounts'])} | {acct_disp} | {qos_disp} | {part_disp} |"
+        )
+
+    # Appendix: common commands
+    L.append(h("Appendix: Common Commands"))
+    L.append(COMMANDS_APPENDIX)
+
+    return "\n".join(L) + "\n"
+
+
+COMMANDS_APPENDIX = """```bash
+# --- Account / QoS ---
+spur accounts show user                 # all user->account rows (grep yourself)
+sacctmgr show user $(whoami)            # equivalent alias
+spur accounts show qos                  # global QoS list (Prio/Preempt)
+spur accounts show account              # all accounts
+sshare -u $(whoami)                     # fair-share info
+
+# --- Cluster / nodes ---
+sinfo                                   # partitions + node counts per state
+sinfo -N -o "%N %P %t %G"               # per node: name/partition/state/GRES
+sinfo -N -h -o "%t" | sort | uniq -c    # node count per state
+scontrol show node <node>               # single-node detail (CPUAlloc/GRES/State)
+scontrol show reservation               # reservations
+
+# --- Jobs / queue ---
+squeue                                   # queue
+squeue -o "%i %P %q %a %u %T %D %N"      # custom columns (NOTE: delimiters render as spaces)
+squeue -u $(whoami)                      # only your jobs
+squeue -t PENDING                        # only pending jobs
+scontrol show job <jobid>                # single-job detail (QOS/Account/NodeList)
+
+# --- Submit (whole-node exclusive) ---
+sbatch -A amd-primus -q amd-primus-qos -p amd-spur --exclusive -N <N> -t <t> <script>
+scancel <jobid>                          # cancel a job
+
+# --- Other spur-native views ---
+spur priority                            # job priority breakdown
+spur stat                                # running-job statistics
+spur report cluster -s now-7days         # cluster utilization report
+spur health                              # node health
+```
+
+> Spur quirks: custom `-o` delimiters render as spaces; positional filters for `show user`/`show account` are ignored (use grep/awk); `sacctmgr` policy queries are blocked; there is no `scontrol show hostnames`; the association table is empty and QoS is not enforced per account (only the account is)."""
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Spur cluster status report")
+    ap.add_argument("--user", default=None, help="user to report on (default: current user)")
+    args = ap.parse_args()
+    user = args.user or getpass.getuser()
+    print(build_report(user))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/README.md b/tests/README.md
index e69de29bb..340ebe111 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -0,0 +1,84 @@
+# Primus Tests
+
+This directory contains the test suite for Primus.
+
+## Test Structure
+
+```
+tests/
+├── runner/                    # Shell integration tests for primus-cli
+│   ├── run_all_tests.sh       # Master test runner
+│   ├── lib/                   # Library function tests
+│   ├── helpers/               # Hook and environment tests
+│   └── test_primus_cli*.sh    # CLI mode tests
+├── unit_tests/                # Python unit tests (pytest)
+│   ├── agents/                # Tuning-agent tests
+│   ├── backends/              # Backend-specific tests (megatron, torchtitan, maxtext, ...)
+│   ├── ci/                    # CI helper tests
+│   ├── cli/                   # CLI tests
+│   ├── core/                  # Core library tests (config/, patches/, backend/, launcher/, projection/, pipeline_parallel/, runtime/, trainer/, utils/)
+│   ├── megatron/              # Megatron-specific unit tests
+│   ├── modules/               # Module/trainer tests
+│   └── tools/                 # Tooling tests
+├── trainer/                   # Integration tests (require GPU)
+│   ├── test_megatron_trainer.py
+│   ├── test_torchtitan_trainer.py
+│   └── test_maxtext_trainer.py
+├── scripts/                   # CI unit/integration launch scripts and UT patches
+├── utils.py                   # Shared test utilities
+├── conftest.py                # Shared pytest fixtures
+└── run_unit_tests.py          # Python test orchestrator (walks tests/)
+```
+
+> **Note:** `config/` and `patches/` live under `unit_tests/core/` (i.e. `tests/unit_tests/core/config/` and `tests/unit_tests/core/patches/`), not directly under `unit_tests/`.
+
+## Running Tests
+
+```bash
+# Shell integration tests
+bash ./tests/runner/run_all_tests.sh
+
+# Python unit tests
+pytest tests/unit_tests/ --maxfail=1 -s
+
+# All tests via orchestrator
+python ./tests/run_unit_tests.py          # Torch backends
+python ./tests/run_unit_tests.py --jax    # JAX/MaxText backend
+```
+
+## Test Tiers
+
+Every trainer E2E case is a full training launch, so the suites hold **one model
+per architecture and per feature path**. Anything beyond that is either deleted
+or deferred to the weekend:
+
+- **Deleted.** A model that only scales the dims of an existing test earns no
+  coverage. Its recipe is still schema-checked by
+  `unit_tests/configs/test_example_configs.py`, which loads every yaml under
+  `examples/**/configs/`.
+- **`@pytest.mark.weekly`.** Extended coverage worth having but not worth a PR's
+  wall clock, and new cases during burn-in. Deselected on PRs and pushes; run by
+  the weekend scheduled build (Saturday 18:00 UTC), which you can also trigger by
+  hand via the `Primus-CI-TAS` workflow with `full_tests=true`. Promote a case to
+  the per-PR tier by deleting the marker once a weekend run has shown it green.
+- **Per-PR.** Everything else.
+
+```bash
+# What PR CI runs
+pytest tests/trainer/test_megatron_trainer.py -m "not weekly" -s
+
+# The weekend tier: full matrix, plus the slow unit-test shape gates
+pytest tests/trainer/test_megatron_trainer.py -s
+pytest tests/unit_tests/ --run-slow -s
+```
+
+Before removing a model, check it is not the last user of a feature path. The
+difference often lives in the recipe yaml rather than in the test body, so two
+cases whose `extra_args` look identical can still cover different code — diff the
+recipes before concluding they are isomorphic.
+
+Cases hidden by `--deselect` in `ci.yaml` (broken on the current toolchain) and
+by `JAX_SKIP_UT=1` (most MaxText models) stay hidden in **both** tiers; they are
+not part of this mechanism.
+
+For comprehensive testing documentation, see the [Testing Guide](../docs/06-developer-guide/testing.md).
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 000000000..97a0fd4f1
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,156 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import os
+import sys
+from pathlib import Path
+
+
+def _warmup_aiter_nondeterministic_mha_bwd() -> None:
+    """Load aiter's nondeterministic flash-attention backward kernel before TE.
+
+    aiter and TransformerEngine both statically bundle composable_kernel
+    (``ck_tile``). If ``transformer_engine`` is imported before aiter's
+    nondeterministic ``mha_bwd`` JIT module is first loaded, that kernel's
+    ``.so`` resolves its ck_tile host launch path against TE's (different) CK
+    copy and launches with an invalid grid/block config -> at runtime the
+    backward dies with::
+
+        HIP Function Failed (.../ck_tile/host/kernel_launch.hpp,110)
+        invalid configuration argument
+
+    Triggering one tiny nondeterministic backward here (in the root conftest's
+    ``pytest_configure``, before any test module imports TE) loads that kernel
+    against aiter's own CK first, after which it stays correct for the rest of
+    the session. The deterministic kernel is unaffected, so we only need to warm
+    the nondeterministic variant. Best-effort: never let warmup break the suite.
+    """
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return
+
+        import math
+
+        import primus_turbo.pytorch as pt
+
+        q, k, v = (
+            torch.randn(1, 8, 1, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+            for _ in range(3)
+        )
+        out = pt.ops.flash_attn_func(
+            q,
+            k,
+            v,
+            dropout_p=0.0,
+            softmax_scale=1.0 / math.sqrt(64),
+            causal=False,
+            window_size=(-1, -1),
+            deterministic=False,
+            return_lse=False,
+        )
+        out.float().sum().backward()
+        torch.cuda.synchronize()
+    except Exception:
+        # Warmup is a best-effort mitigation; if aiter/CUDA is unavailable or
+        # the API shifts, fall through silently and let the affected tests
+        # surface their own errors.
+        pass
+
+
+def _selection_includes_megatron(config) -> bool:
+    """True if the pytest selection could include a backends/megatron suite.
+
+    Covers both the unit (``tests/unit_tests/backends/megatron``) and integration
+    (``tests/integration_tests/backends/megatron``) trees: the diffusion
+    integration tests run the same aiter attention backward and need the same
+    crash mitigations (deepbind hook + warmup).
+
+    Path-level heuristic only (does not inspect -k / -m); the mitigations are
+    best-effort, so a conservative path check is sufficient. Fails OPEN: any
+    detection error returns True, because a missed mitigation degrades to a hard
+    mha_bwd crash, not just a slowdown.
+
+    Uses ``config.args`` (populated post-parse, includes the default rootdir on
+    a bare run) rather than ``config.getoption("file_or_dir")`` (empty on a bare
+    run) -- do not "simplify" to the latter, it flips the empty-case semantics.
+    """
+    try:
+        base = Path(__file__).resolve().parent
+        megatron_dirs = (
+            base / "unit_tests" / "backends" / "megatron",
+            base / "integration_tests" / "backends" / "megatron",
+        )
+        args = list(getattr(config, "args", []) or [])
+        if not args:
+            return True  # whole-suite run (defensive; config.args is never empty)
+        for arg in args:
+            path = Path(str(arg).split("::", 1)[0]).resolve()
+            for megatron_dir in megatron_dirs:
+                # arg is the megatron dir, an ancestor of it, or a path inside it
+                if (
+                    path == megatron_dir
+                    or megatron_dir.is_relative_to(path)
+                    or path.is_relative_to(megatron_dir)
+                ):
+                    return True
+        return False
+    except Exception:
+        # Fail open: never let a detection error suppress the crash-prevention mitigations.
+        return True
+
+
+def pytest_configure(config):
+    project_root = Path(__file__).resolve().parent.parent
+    if str(project_root) not in sys.path:
+        sys.path.insert(0, str(project_root))
+
+    # What belongs in this tier, and what should be deleted instead of demoted
+    # into it, is documented in docs/06-developer-guide/testing.md.
+    config.addinivalue_line(
+        "markers",
+        (
+            "weekly: extended or not-yet-burned-in E2E coverage. Deselected on "
+            "PR/push CI via -m 'not weekly'; run by the weekend scheduled CI run."
+        ),
+    )
+
+    megatron_path = os.environ.get("MEGATRON_PATH")
+    if megatron_path is None or not os.path.exists(megatron_path):
+        megatron_path = project_root / "third_party" / "Megatron-LM"
+    if str(megatron_path) not in sys.path:
+        sys.path.append(str(megatron_path))
+
+    # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools,
+    # no __init__.py) that only import with the source root on sys.path. Insert the
+    # submodule at the FRONT so it wins over any stale torchtitan in site-packages.
+    torchtitan_path = os.environ.get("TORCHTITAN_PATH")
+    if torchtitan_path is None or not os.path.exists(torchtitan_path):
+        torchtitan_path = project_root / "third_party" / "torchtitan"
+    if str(torchtitan_path) not in sys.path:
+        sys.path.insert(0, str(torchtitan_path))
+
+    # Only needed for the megatron suite, so skip for unrelated selections
+    # (fail-open: detection errors still run it).
+    if _selection_includes_megatron(config):
+        # ORDER MATTERS: install the aiter RTLD_DEEPBIND import hook BEFORE the
+        # warmup below (or any test) first imports aiter's mha backward
+        # extension. The hook wraps importlib.import_module so the pinned
+        # aiter::mha_bwd binds its own ck_tile instead of TE's stale vendored
+        # libmha (ROCm/aiter#1332); once ``module_fmha_v3_bwd`` is already in
+        # sys.modules the hook is a no-op, so the warmup MUST NOT run first or
+        # the hd128 backward launches with an invalid grid config and aborts the
+        # whole process. Mirrors the production megatron.turbo.aiter_deepbind
+        # before_train patch (which has no such ordering hazard).
+        from tests.utils import install_aiter_deepbind_hook
+
+        install_aiter_deepbind_hook()
+        # Must run before any test module imports transformer_engine. See the
+        # helper docstring for the aiter<->TE composable_kernel load-order issue.
+        _warmup_aiter_nondeterministic_mha_bwd()
+    else:
+        print("[conftest] skipping aiter mha_bwd mitigations (no megatron tests selected)")
diff --git a/tests/integration_tests/backends/megatron/diffusion/data/__init__.py b/tests/integration_tests/backends/megatron/diffusion/data/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/data/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/integration_tests/backends/megatron/diffusion/data/conftest.py b/tests/integration_tests/backends/megatron/diffusion/data/conftest.py
new file mode 100644
index 000000000..a45c68597
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/data/conftest.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Shared fixtures for data module integration tests.
+
+May share fixtures from unit tests or provide integration-specific fixtures.
+"""
+
+# Import shared fixtures from unit tests
+from tests.unit_tests.backends.megatron.diffusion.data.conftest import *  # noqa: F401, F403
diff --git a/tests/integration_tests/backends/megatron/diffusion/data/test_raw_pipeline_e2e.py b/tests/integration_tests/backends/megatron/diffusion/data/test_raw_pipeline_e2e.py
new file mode 100644
index 000000000..51ddd8c32
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/data/test_raw_pipeline_e2e.py
@@ -0,0 +1,202 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+End-to-end integration tests for RawDatasetPipeline.
+
+Exercises the full preprocessing data flow (load -> resize -> encode -> shard)
+with real files and real I/O. No GPU, no model weights, no HuggingFace downloads.
+"""
+
+import tarfile
+from unittest.mock import patch
+
+import pytest
+from PIL import Image
+
+from primus.backends.megatron.data.diffusion.preprocessing.pipelines.raw import (
+    RawDatasetPipeline,
+)
+from tests.utils import PrimusUT
+
+
+class TestRawPipelineE2E(PrimusUT):
+    """End-to-end tests for RawDatasetPipeline with directory source."""
+
+    @pytest.fixture(autouse=True)
+    def setup_test_dataset(self, tmp_path):
+        """Create temp directories with test images and captions."""
+        self.input_dir = tmp_path / "input"
+        self.output_dir = tmp_path / "output"
+        images_dir = self.input_dir / "images"
+        captions_dir = self.input_dir / "captions"
+        images_dir.mkdir(parents=True)
+        captions_dir.mkdir(parents=True)
+        self.output_dir.mkdir()
+
+        for i in range(8):
+            img = Image.new("RGB", (64, 64), color=(i * 30, i * 20, i * 10))
+            img.save(images_dir / f"{i:04d}.jpg")
+            (captions_dir / f"{i:04d}.txt").write_text(f"caption {i}")
+
+    def test_raw_pipeline_end_to_end(self):
+        """Full pipeline run produces correct shards with jpg+txt entries."""
+        pipeline = RawDatasetPipeline(
+            source_type="directory",
+            output_dir=str(self.output_dir),
+            image_size=32,
+            center_crop=True,
+            variable_size=False,
+            shard_size=5,
+        )
+
+        results = pipeline.run(input_dir=str(self.input_dir))
+
+        assert results["samples_processed"] == 8
+        assert results["samples_skipped"] == 0
+        assert results["shards_written"] == 2
+
+        tar_files = sorted(self.output_dir.glob("*.tar"))
+        assert len(tar_files) == 2
+
+        for tar_path in tar_files:
+            with tarfile.open(str(tar_path), "r") as tar:
+                members = tar.getmembers()
+                assert len(members) > 0
+
+                sample_keys = set()
+                extensions_by_key = {}
+                for m in members:
+                    key, ext = m.name.split(".", 1)
+                    sample_keys.add(key)
+                    extensions_by_key.setdefault(key, set()).add(ext)
+
+                for key, exts in extensions_by_key.items():
+                    assert "jpg" in exts, f"Sample {key} missing jpg"
+                    assert "txt" in exts, f"Sample {key} missing txt"
+
+    def test_raw_pipeline_max_samples(self):
+        """max_samples limits the number of processed samples."""
+        pipeline = RawDatasetPipeline(
+            source_type="directory",
+            output_dir=str(self.output_dir),
+            image_size=32,
+            center_crop=True,
+            variable_size=False,
+            shard_size=1000,
+            max_samples=3,
+        )
+
+        results = pipeline.run(input_dir=str(self.input_dir))
+
+        assert results["samples_processed"] == 3
+        tar_files = list(self.output_dir.glob("*.tar"))
+        assert len(tar_files) == 1
+
+    def test_raw_pipeline_empty_directory(self):
+        """Empty images/captions dirs produce zero samples and no shards."""
+        empty_input = self.output_dir.parent / "empty_input"
+        (empty_input / "images").mkdir(parents=True)
+        (empty_input / "captions").mkdir(parents=True)
+        empty_output = self.output_dir.parent / "empty_output"
+        empty_output.mkdir()
+
+        pipeline = RawDatasetPipeline(
+            source_type="directory",
+            output_dir=str(empty_output),
+            image_size=32,
+            center_crop=True,
+            variable_size=False,
+            shard_size=1000,
+        )
+
+        results = pipeline.run(input_dir=str(empty_input))
+
+        assert results["samples_processed"] == 0
+        assert results["samples_skipped"] == 0
+        tar_files = list(empty_output.glob("*.tar"))
+        assert len(tar_files) == 0
+
+    def test_raw_pipeline_max_samples_is_total_not_per_rank(self):
+        """max_samples=4 with world_size=2 produces 4 total samples (2 per rank), not 8."""
+        total_across_ranks = 0
+        dist_info_path = (
+            "primus.backends.megatron.data.diffusion.preprocessing.pipelines.raw" ".get_distributed_info"
+        )
+
+        for rank in range(2):
+            rank_output = self.output_dir.parent / f"rank_{rank}_output"
+            rank_output.mkdir()
+
+            with patch(dist_info_path, return_value=(rank, 2)):
+                pipeline = RawDatasetPipeline(
+                    source_type="directory",
+                    output_dir=str(rank_output),
+                    image_size=32,
+                    center_crop=True,
+                    variable_size=False,
+                    shard_size=1000,
+                    max_samples=4,
+                )
+                results = pipeline.run(input_dir=str(self.input_dir))
+
+            assert (
+                results["samples_processed"] == 2
+            ), f"Rank {rank} should process 2 samples, got {results['samples_processed']}"
+            total_across_ranks += results["samples_processed"]
+
+        assert total_across_ranks == 4, f"Expected 4 total samples across 2 ranks, got {total_across_ranks}"
+
+    def test_raw_pipeline_round_robin_distribution(self):
+        """Streaming round-robin split covers all samples exactly once across ranks.
+
+        With world_size=2 and 8 items, each rank should process 4 samples by
+        global load index (rank 0 -> even indices, rank 1 -> odd), and the union
+        of captions across ranks must equal the full set with no overlap.
+        """
+        dist_info_path = (
+            "primus.backends.megatron.data.diffusion.preprocessing.pipelines.raw" ".get_distributed_info"
+        )
+
+        def captions_in(output_dir):
+            captions = set()
+            for tar_path in sorted(output_dir.glob("*.tar")):
+                with tarfile.open(str(tar_path), "r") as tar:
+                    for m in tar.getmembers():
+                        if m.name.endswith(".txt"):
+                            captions.add(tar.extractfile(m).read().decode("utf-8"))
+            return captions
+
+        rank_captions = {}
+        for rank in range(2):
+            rank_output = self.output_dir.parent / f"rr_rank_{rank}_output"
+            rank_output.mkdir()
+
+            with patch(dist_info_path, return_value=(rank, 2)):
+                pipeline = RawDatasetPipeline(
+                    source_type="directory",
+                    output_dir=str(rank_output),
+                    image_size=32,
+                    center_crop=True,
+                    variable_size=False,
+                    shard_size=1000,
+                )
+                results = pipeline.run(input_dir=str(self.input_dir))
+
+            assert (
+                results["samples_processed"] == 4
+            ), f"Rank {rank} should process 4 of 8 samples, got {results['samples_processed']}"
+            rank_captions[rank] = captions_in(rank_output)
+
+        # No overlap between ranks
+        assert rank_captions[0].isdisjoint(
+            rank_captions[1]
+        ), f"Ranks processed overlapping samples: {rank_captions[0] & rank_captions[1]}"
+        # Complete coverage of all 8 captions
+        all_captions = rank_captions[0] | rank_captions[1]
+        expected = {f"caption {i}" for i in range(8)}
+        assert all_captions == expected, f"Expected {expected}, got {all_captions}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/conftest.py b/tests/integration_tests/backends/megatron/diffusion/distributed/conftest.py
new file mode 100644
index 000000000..49d175551
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/conftest.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Shared fixtures for distributed diffusion integration tests.
+
+These tests construct and run real Flux models (forward + backward through the
+Primus Turbo / aiter attention path), so they need the same container-quirk
+mitigations the diffusion *unit* tests already have, which the integration tree
+was previously missing:
+
+  * ``install_aiter_deepbind_hook()`` -- isolates the pinned aiter::mha_bwd from
+    TransformerEngine's stale vendored libmha so the hd128 attention backward
+    does not abort the process (ROCm/aiter#1332). The root conftest installs
+    this early (before its warmup); calling it again here is idempotent and
+    keeps the integration tree self-sufficient.
+  * ``_unset_nvte_attention_env`` -- clears baked ``NVTE_FLASH_ATTN=0`` etc. so
+    Flux's ``auto`` attention-backend check does not reject model construction.
+
+Re-exports ``init_parallel_state`` and ``_unset_nvte_attention_env`` from the
+unit conftests so both trees share one definition.
+"""
+
+from tests.utils import install_aiter_deepbind_hook
+
+# Install Primus' production aiter RTLD_DEEPBIND import hook before any test
+# imports the aiter mha kernels. No-op without a GPU / when already installed.
+install_aiter_deepbind_hook()
+
+from tests.unit_tests.backends.megatron.conftest import (  # noqa: F401,E402
+    init_parallel_state,
+)
+from tests.unit_tests.backends.megatron.diffusion.conftest import (  # noqa: F401,E402
+    _unset_nvte_attention_env,
+)
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_ddp_fp8_compile_convergence.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_ddp_fp8_compile_convergence.py
new file mode 100644
index 000000000..9fa148a1c
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_ddp_fp8_compile_convergence.py
@@ -0,0 +1,400 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Multi-GPU convergence test: Megatron DDP + FP8 + torch.compile scope.
+
+Tests whether torch.compile scope (whole_model vs inner-only) causes
+convergence divergence when real NCCL all-reduce + overlapped gradient
+reduction is active. Phase 1 and Phase 2 proved that single-GPU FP8 +
+compile and single-GPU mock DDP + compile are bit-identical. This test
+adds real NCCL communication to isolate the remaining hypothesis.
+
+Run with:
+    torchrun --nproc_per_node 2 --master_port 29502 -m pytest -xvs \
+        tests/integration_tests/backends/megatron/diffusion/distributed/test_ddp_fp8_compile_convergence.py
+"""
+
+import copy
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from primus_turbo.pytorch.core.backend import BackendType
+from primus_turbo.pytorch.core.low_precision import (
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    OpaqueFP8LinearTensorwiseFunction,
+)
+
+_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 0))
+
+requires_torchrun = pytest.mark.skipif(
+    _WORLD_SIZE < 1,
+    reason="Requires torchrun (WORLD_SIZE env var not set)",
+)
+
+_DIM = 256
+_OUT_DIM = 512
+_BATCH = 16
+_N_STEPS = 100
+_GRAN_VALUE = ScalingGranularity.TENSORWISE.value
+_BACKEND_VALUE = BackendType.HIPBLASLT.value
+
+
+# ---------------------------------------------------------------------------
+# Distributed setup / teardown (follows FSDP2 integration test pattern)
+# ---------------------------------------------------------------------------
+
+
+def _init_distributed():
+    from megatron.core import parallel_state as ps
+
+    if not torch.distributed.is_initialized():
+        torch.distributed.init_process_group(backend="nccl")
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+
+    import primus.core.utils.logger as primus_logger
+
+    if primus_logger._logger is None:
+        cfg = primus_logger.LoggerConfig(
+            exp_root_path="/tmp/primus_test_logs",
+            work_group="test",
+            user_name="test",
+            exp_name="ddp_fp8_compile_convergence",
+            module_name="test",
+            rank=rank,
+            world_size=world_size,
+        )
+        primus_logger.setup_logger(cfg, is_head=(rank == 0))
+
+    from primus.core.utils.module_utils import set_logging_rank
+
+    set_logging_rank(rank, world_size)
+
+    if ps.model_parallel_is_initialized():
+        ps.destroy_model_parallel()
+
+    ps.initialize_model_parallel(
+        tensor_model_parallel_size=1,
+        pipeline_model_parallel_size=1,
+        expert_model_parallel_size=1,
+    )
+
+    from megatron.core.tensor_parallel import random as tp_random
+
+    try:
+        tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True)
+    except (ImportError, AssertionError):
+        tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True)
+    tp_random.model_parallel_cuda_manual_seed(42)
+
+    from megatron.training.global_vars import set_args
+
+    set_args(
+        SimpleNamespace(
+            enable_turbo_attention_float8=False,
+            data_parallel_replicate_degree=1,
+            use_fsdp2_fp8_all_gather=False,
+            use_fsdp2_bf16_master_weight_optimizer=False,
+            use_fsdp2_fp32_param_optimizer=False,
+            torch_fsdp2_reshard_after_forward=False,
+        )
+    )
+
+
+def _cleanup_distributed():
+    from megatron.core import parallel_state as ps
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    if ps.model_parallel_is_initialized():
+        ps.destroy_model_parallel()
+
+    import megatron.training.global_vars as gvars
+
+    gvars._GLOBAL_ARGS = None
+
+
+# ---------------------------------------------------------------------------
+# Model (same as Phase 1/2 _OpaqueModule)
+# ---------------------------------------------------------------------------
+
+
+class _OpaqueModule(nn.Module):
+    """FP8 module using OpaqueFP8LinearTensorwiseFunction with pre-extracted weights."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        result = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight,
+            w_fp8,
+            w_scale,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result[0]
+
+
+# ---------------------------------------------------------------------------
+# Optimizer that reads main_grad (reused from Phase 2)
+# ---------------------------------------------------------------------------
+
+
+class _MainGradAdamW(torch.optim.AdamW):
+    """AdamW that reads from param.main_grad instead of param.grad."""
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if hasattr(p, "main_grad") and p.main_grad is not None:
+                    p.grad = p.main_grad.clone()
+        result = super().step(closure=closure)
+        for group in self.param_groups:
+            for p in group["params"]:
+                p.grad = None
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_transformer_config():
+    from megatron.core.transformer.transformer_config import TransformerConfig
+
+    return TransformerConfig(
+        num_layers=1,
+        hidden_size=_DIM,
+        num_attention_heads=1,
+        gradient_accumulation_fusion=False,
+    )
+
+
+def _wrap_with_megatron_ddp(module, overlap_grad_reduce=True):
+    from megatron.core.distributed import DistributedDataParallel
+    from megatron.core.distributed.distributed_data_parallel_config import (
+        DistributedDataParallelConfig,
+    )
+
+    config = _make_transformer_config()
+    ddp_config = DistributedDataParallelConfig(
+        overlap_grad_reduce=overlap_grad_reduce,
+        use_distributed_optimizer=False,
+        average_in_collective=False,
+    )
+    ddp = DistributedDataParallel(
+        config=config,
+        ddp_config=ddp_config,
+        module=module,
+    )
+    return ddp
+
+
+def _make_inputs(n, dim, batch, seed=42):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    return [torch.randn(batch, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)]
+
+
+def _clone_state(model):
+    return copy.deepcopy(model.state_dict())
+
+
+def _run_ddp_training(ddp, inputs, n_steps, compile_target=None, lr=1e-3):
+    """Run training with real Megatron DDP.
+
+    Args:
+        ddp: Megatron DistributedDataParallel-wrapped model.
+        inputs: List of input tensors.
+        n_steps: Number of training steps.
+        compile_target: None="eager", "whole_model", or "inner_only".
+        lr: Learning rate.
+    """
+    optimizer = _MainGradAdamW(ddp.module.parameters(), lr=lr)
+
+    if compile_target == "whole_model":
+        ddp.forward = torch.compile(ddp.forward)
+    elif compile_target == "inner_only":
+        ddp.module.forward = torch.compile(ddp.module.forward)
+
+    losses = []
+    for step in range(n_steps):
+        ddp.zero_grad_buffer()
+        out = ddp(inputs[step % len(inputs)])
+        loss = out.sum()
+        loss.backward()
+        ddp.finish_grad_sync()
+        optimizer.step()
+        losses.append(loss.item())
+    return losses
+
+
+def _print_comparison(label_a, losses_a, label_b, losses_b, milestones=None):
+    if milestones is None:
+        milestones = [0, 1, 5, 10, 20, 50]
+        milestones += [i for i in range(100, max(len(losses_a), len(losses_b)), 50)]
+        if len(losses_a) - 1 not in milestones:
+            milestones.append(len(losses_a) - 1)
+    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+    if rank != 0:
+        return
+    print(f"\n{'Step':>6} | {label_a:>20} | {label_b:>20} | {'Rel Diff':>10}")
+    print("-" * 65)
+    for step in milestones:
+        if step < len(losses_a) and step < len(losses_b):
+            a, b = losses_a[step], losses_b[step]
+            rel = abs(a - b) / max(abs(a), abs(b), 1e-12)
+            print(f"{step:>6} | {a:>20.6f} | {b:>20.6f} | {rel:>10.6f}")
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestDDPFP8CompileConvergence:
+    """Real Megatron DDP + NCCL + FP8 + torch.compile scope convergence tests.
+
+    Tests whether whole_model compile (failing local spec config) vs
+    inner_only compile (converging TE spec config) diverge when real
+    NCCL all-reduce with overlap_grad_reduce=True is active.
+    """
+
+    def setup_method(self):
+        _init_distributed()
+        torch._dynamo.reset()
+
+    def teardown_method(self):
+        _cleanup_distributed()
+
+    @requires_torchrun
+    def test_whole_model_compiled_vs_eager(self):
+        """whole_model compile + real DDP should track eager + real DDP."""
+        torch._dynamo.reset()
+        torch.manual_seed(42)
+        inputs = _make_inputs(20, _DIM, _BATCH)
+
+        torch.manual_seed(123)
+        eager_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(eager_model)
+        eager_ddp = _wrap_with_megatron_ddp(eager_model)
+        eager_losses = _run_ddp_training(eager_ddp, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        torch.manual_seed(123)
+        compiled_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_ddp = _wrap_with_megatron_ddp(compiled_model)
+        compiled_losses = _run_ddp_training(compiled_ddp, inputs, _N_STEPS, compile_target="whole_model")
+
+        _print_comparison(
+            "DDP Eager",
+            eager_losses,
+            "DDP WholeModel",
+            compiled_losses,
+        )
+
+        rank = torch.distributed.get_rank()
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        if rank == 0:
+            print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"[Rank {rank}] DDP whole_model compile diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    @requires_torchrun
+    def test_inner_only_compiled_vs_eager(self):
+        """inner-only compile + real DDP should track eager + real DDP."""
+        torch._dynamo.reset()
+        torch.manual_seed(42)
+        inputs = _make_inputs(20, _DIM, _BATCH)
+
+        torch.manual_seed(123)
+        eager_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(eager_model)
+        eager_ddp = _wrap_with_megatron_ddp(eager_model)
+        eager_losses = _run_ddp_training(eager_ddp, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        torch.manual_seed(123)
+        compiled_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_ddp = _wrap_with_megatron_ddp(compiled_model)
+        compiled_losses = _run_ddp_training(compiled_ddp, inputs, _N_STEPS, compile_target="inner_only")
+
+        _print_comparison(
+            "DDP Eager",
+            eager_losses,
+            "DDP InnerOnly",
+            compiled_losses,
+        )
+
+        rank = torch.distributed.get_rank()
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        if rank == 0:
+            print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"[Rank {rank}] DDP inner-only compile diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    @requires_torchrun
+    def test_whole_model_vs_inner_only(self):
+        """Direct comparison of whole_model vs inner_only compile with real DDP."""
+        torch._dynamo.reset()
+        torch.manual_seed(42)
+        inputs = _make_inputs(20, _DIM, _BATCH)
+
+        torch.manual_seed(123)
+        whole_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(whole_model)
+        whole_ddp = _wrap_with_megatron_ddp(whole_model)
+        whole_losses = _run_ddp_training(whole_ddp, inputs, _N_STEPS, compile_target="whole_model")
+
+        torch._dynamo.reset()
+        torch.manual_seed(123)
+        inner_model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        inner_model.load_state_dict(init_state)
+        inner_ddp = _wrap_with_megatron_ddp(inner_model)
+        inner_losses = _run_ddp_training(inner_ddp, inputs, _N_STEPS, compile_target="inner_only")
+
+        _print_comparison(
+            "DDP WholeModel",
+            whole_losses,
+            "DDP InnerOnly",
+            inner_losses,
+        )
+
+        rank = torch.distributed.get_rank()
+        final_rel = abs(whole_losses[-1] - inner_losses[-1]) / max(abs(whole_losses[-1]), 1e-12)
+        if rank == 0:
+            print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"[Rank {rank}] DDP whole_model vs inner-only diverged: "
+            f"whole={whole_losses[-1]:.6f}, inner={inner_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fp8_local_spec.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fp8_local_spec.py
new file mode 100644
index 000000000..4defd1cf6
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fp8_local_spec.py
@@ -0,0 +1,118 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Integration tests for Flux model with Float8 local spec.
+
+Tests that a real (small) Flux model constructs, uses correct linear types,
+and produces valid output when using PrimusTurboFloat8LocalSpecProvider.
+"""
+
+import pytest
+import torch
+
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    Float8ColumnParallelLinear,
+    Float8RowParallelLinear,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    IMG_SIZE_TINY,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_SHORT,
+    VAE_LATENT_CHANNELS,
+)
+from tests.utils import PrimusUT
+
+
+class TestFluxFP8LocalSpec(PrimusUT):
+    """Integration tests for Flux with Float8 local spec provider."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for FP8 local spec integration tests."""
+
+    def _make_fp8_config(self):
+        return FluxConfig.flux_535m(
+            transformer_impl="local",
+            fp8="e4m3",
+            fp8_recipe="tensorwise",
+        )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_535m_fp8_local_spec_constructs(self):
+        """Test that Flux 535M with FP8 local spec constructs without error."""
+        config = self._make_fp8_config()
+        model = Flux(config)
+        assert model is not None
+        assert isinstance(model, Flux)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_535m_fp8_local_spec_linear_types(self):
+        """Test that spec-provided parallel linears are Float8 variants after construction.
+
+        Only checks linears inside attention (linear_qkv, linear_proj) and MLP
+        (linear_fc1, linear_fc2) submodules. AdaLN modulation layers are created
+        directly and intentionally use standard ColumnParallelLinear.
+        """
+        config = self._make_fp8_config()
+        model = Flux(config)
+
+        spec_linear_names = {"linear_qkv", "added_linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"}
+
+        from megatron.core.tensor_parallel.layers import (
+            ColumnParallelLinear,
+            RowParallelLinear,
+        )
+
+        for name, module in model.named_modules():
+            leaf_name = name.rsplit(".", 1)[-1] if "." in name else name
+            if leaf_name not in spec_linear_names:
+                continue
+            if isinstance(module, ColumnParallelLinear):
+                assert isinstance(
+                    module, Float8ColumnParallelLinear
+                ), f"{name}: expected Float8ColumnParallelLinear, got {type(module).__name__}"
+            if isinstance(module, RowParallelLinear):
+                assert isinstance(
+                    module, Float8RowParallelLinear
+                ), f"{name}: expected Float8RowParallelLinear, got {type(module).__name__}"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_535m_fp8_local_spec_forward(self):
+        """Test forward pass produces valid (non-NaN/Inf) output with correct shape."""
+        config = self._make_fp8_config()
+        model = Flux(config).cuda().to(torch.bfloat16)
+        model.eval()
+
+        batch_size = 2
+        height, width = IMG_SIZE_TINY, IMG_SIZE_TINY
+        channels = VAE_LATENT_CHANNELS
+        txt_seq_len = TEXT_SEQ_LEN_SHORT
+
+        img = torch.randn(batch_size, channels, height, width, dtype=torch.bfloat16).cuda()
+        txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        timesteps = torch.rand(batch_size, dtype=torch.bfloat16).cuda()
+
+        packed_img = pack_latents(img)
+        packed_img = packed_img.transpose(0, 1)
+        txt_t = txt.transpose(0, 1)
+
+        img_ids = generate_image_position_ids(batch_size, height, width, device="cuda")
+        txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda()
+
+        with torch.no_grad():
+            output = model(packed_img, txt_t, y, timesteps, img_ids, txt_ids)
+
+        assert output is not None
+        assert len(output.shape) == 3
+        assert output.shape[1] == batch_size
+        assert not torch.isnan(output).any(), "Output contains NaN"
+        assert not torch.isinf(output).any(), "Output contains Inf"
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_fp8_all_gather_integration.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_fp8_all_gather_integration.py
new file mode 100644
index 000000000..52c10be0a
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_fp8_all_gather_integration.py
@@ -0,0 +1,407 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Integration tests for Flux model with FSDP2 FP8 all-gather (Approach A).
+
+Single-GPU tests verify wrapping, forward, and forward+backward using
+the init_parallel_state fixture (world_size=1, all-reduce is a no-op).
+
+Multi-GPU tests (torchrun) verify real FP8 all-gather communication,
+training steps, numerical parity, and torch.compile compatibility.
+
+Run multi-GPU tests with:
+    torchrun --nproc_per_node 2 --master_port 29501 -m pytest -xvs \\
+        tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_fp8_all_gather_integration.py
+"""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+from megatron.core.distributed.distributed_data_parallel_config import (
+    DistributedDataParallelConfig,
+)
+
+from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+    WeightWithFP8AllGatherTensor,
+)
+from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import (
+    PrimusTorchFullyShardedDataParallel,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    IMG_SIZE_TINY,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_SHORT,
+    VAE_LATENT_CHANNELS,
+)
+from tests.utils import PrimusUT
+
+_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 0))
+
+requires_torchrun = pytest.mark.skipif(
+    _WORLD_SIZE < 1,
+    reason="Requires torchrun (WORLD_SIZE env var not set)",
+)
+
+
+def _make_flux_inputs(batch_size, device):
+    height, width = IMG_SIZE_TINY, IMG_SIZE_TINY
+    img = torch.randn(
+        batch_size,
+        VAE_LATENT_CHANNELS,
+        height,
+        width,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    packed_img = pack_latents(img).transpose(0, 1)
+    txt = torch.randn(
+        batch_size,
+        TEXT_SEQ_LEN_SHORT,
+        T5_XXL_EMBEDDING_DIM,
+        dtype=torch.bfloat16,
+        device=device,
+    ).transpose(0, 1)
+    y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM, dtype=torch.bfloat16, device=device)
+    timesteps = torch.rand(batch_size, dtype=torch.bfloat16, device=device)
+    img_ids = generate_image_position_ids(batch_size, height, width, device=device)
+    txt_ids = torch.zeros(batch_size, TEXT_SEQ_LEN_SHORT, 3, dtype=torch.bfloat16, device=device)
+    return dict(img=packed_img, txt=txt, y=y, timesteps=timesteps, img_ids=img_ids, txt_ids=txt_ids)
+
+
+# ============================================================================
+# Single-GPU tests (init_parallel_state fixture, world_size=1)
+# ============================================================================
+
+
+class TestFluxFSDP2FP8AllGather(PrimusUT):
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        from megatron.training.global_vars import set_args
+
+        set_args(
+            SimpleNamespace(
+                enable_turbo_attention_float8=False,
+                data_parallel_replicate_degree=1,
+                use_fsdp2_fp8_all_gather=True,
+                use_fsdp2_bf16_master_weight_optimizer=True,
+                use_fsdp2_fp32_param_optimizer=False,
+                torch_fsdp2_reshard_after_forward=False,
+            )
+        )
+
+    def _make_fsdp_model(self):
+        config = FluxConfig.flux_535m(
+            transformer_impl="local",
+            fp8="e4m3",
+            fp8_recipe="tensorwise",
+        )
+        model = Flux(config).cuda().to(torch.bfloat16)
+        ddp_config = DistributedDataParallelConfig()
+        fsdp_wrapper = PrimusTorchFullyShardedDataParallel(
+            config=config,
+            ddp_config=ddp_config,
+            module=model,
+        )
+        return model, fsdp_wrapper
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_fsdp2_fp8_allgather_wrapping(self):
+        """Verify only weight params of FP8-aware modules are wrapped."""
+        from torch.distributed._tensor import DTensor
+
+        model, _ = self._make_fsdp_model()
+
+        expected_fp8_weights = sum(
+            1
+            for m in model.modules()
+            if hasattr(m, "_fp8_config") and hasattr(m, "weight") and m.weight is not None
+        )
+
+        wrapped_count = 0
+        unwrapped_count = 0
+        for m in model.modules():
+            if not hasattr(m, "weight") or m.weight is None:
+                continue
+            w = m.weight
+            local_t = w._local_tensor if isinstance(w, DTensor) else w
+            if isinstance(local_t, WeightWithFP8AllGatherTensor):
+                wrapped_count += 1
+                assert local_t._precomputed_scale is not None, "precomputed scale not set after FSDP init"
+                assert hasattr(m, "_fp8_config"), f"Non-FP8 module {type(m).__name__} should not be wrapped"
+            elif w.requires_grad:
+                unwrapped_count += 1
+
+        assert (
+            wrapped_count == expected_fp8_weights
+        ), f"Expected {expected_fp8_weights} wrapped FP8 weights, got {wrapped_count}"
+        assert unwrapped_count > 0, "Expected some unwrapped params (biases, non-FP8 layers)"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_fsdp2_fp8_allgather_forward(self):
+        """Forward pass produces valid output with correct shape."""
+        model, _ = self._make_fsdp_model()
+        model.eval()
+        inputs = _make_flux_inputs(batch_size=2, device="cuda")
+        with torch.no_grad():
+            output = model(**inputs)
+        assert output is not None
+        assert len(output.shape) == 3
+        assert output.shape[1] == 2
+        assert not torch.isnan(output).any(), "Output contains NaN"
+        assert not torch.isinf(output).any(), "Output contains Inf"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_fsdp2_fp8_allgather_forward_backward(self):
+        """Forward + backward: completes without error, grads allocated for both param types."""
+        from torch.distributed._tensor import DTensor
+
+        model, _ = self._make_fsdp_model()
+        model.train()
+        inputs = _make_flux_inputs(batch_size=2, device="cuda")
+        output = model(**inputs)
+        loss = output.float().sum()
+        loss.backward()
+
+        has_any_nonzero_grad = False
+        fp8_grad_allocated = False
+        non_fp8_grad_allocated = False
+        for param in model.parameters():
+            if param.grad is None:
+                continue
+            local_t = param._local_tensor if isinstance(param, DTensor) else param
+            if isinstance(local_t, WeightWithFP8AllGatherTensor):
+                fp8_grad_allocated = True
+            else:
+                non_fp8_grad_allocated = True
+            if param.grad.abs().sum() > 0:
+                has_any_nonzero_grad = True
+        assert has_any_nonzero_grad, "No non-zero gradients found after backward"
+        assert fp8_grad_allocated, "No grad tensors allocated for FP8-wrapped params"
+        assert non_fp8_grad_allocated, "No grad tensors allocated for non-wrapped params"
+
+
+# ============================================================================
+# Multi-GPU tests (torchrun, real FP8 all-gather communication)
+# ============================================================================
+
+
+def _init_distributed():
+    from megatron.core import parallel_state as ps
+
+    if not torch.distributed.is_initialized():
+        torch.distributed.init_process_group(backend="nccl")
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+
+    import primus.core.utils.logger as primus_logger
+
+    if primus_logger._logger is None:
+        cfg = primus_logger.LoggerConfig(
+            exp_root_path="/tmp/primus_test_logs",
+            work_group="test",
+            user_name="test",
+            exp_name="fp8_allgather_integration",
+            module_name="test",
+            rank=rank,
+            world_size=world_size,
+        )
+        primus_logger.setup_logger(cfg, is_head=(rank == 0))
+
+    from primus.core.utils.module_utils import set_logging_rank
+
+    set_logging_rank(rank, world_size)
+
+    if ps.model_parallel_is_initialized():
+        ps.destroy_model_parallel()
+
+    ps.initialize_model_parallel(
+        tensor_model_parallel_size=1,
+        pipeline_model_parallel_size=1,
+        expert_model_parallel_size=1,
+    )
+
+    from megatron.core.tensor_parallel import random as tp_random
+
+    try:
+        tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True)
+    except (ImportError, AssertionError):
+        tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True)
+    tp_random.model_parallel_cuda_manual_seed(42)
+
+    from megatron.training.global_vars import set_args
+
+    set_args(
+        SimpleNamespace(
+            enable_turbo_attention_float8=False,
+            data_parallel_replicate_degree=1,
+            use_fsdp2_fp8_all_gather=True,
+            use_fsdp2_bf16_master_weight_optimizer=True,
+            use_fsdp2_fp32_param_optimizer=False,
+            torch_fsdp2_reshard_after_forward=False,
+        )
+    )
+
+
+def _cleanup_distributed():
+    from megatron.core import parallel_state as ps
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    if ps.model_parallel_is_initialized():
+        ps.destroy_model_parallel()
+
+    import megatron.training.global_vars as gvars
+
+    gvars._GLOBAL_ARGS = None
+
+
+def _make_optimizer_config():
+    from megatron.core.optimizer.optimizer_config import OptimizerConfig
+
+    return OptimizerConfig(
+        bf16=True,
+        fp16=False,
+        optimizer="adam",
+        lr=1e-4,
+        min_lr=1e-5,
+        weight_decay=0.01,
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        adam_eps=1e-8,
+        clip_grad=1.0,
+        use_precision_aware_optimizer=False,
+    )
+
+
+def _create_model_and_optimizer(seed=42, fp8_all_gather=True):
+    from primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer import (
+        get_fsdp2_bf16_master_weight_optimizer,
+    )
+
+    torch.manual_seed(seed)
+    flux_config = FluxConfig.flux_535m(
+        transformer_impl="local",
+        fp8="e4m3" if fp8_all_gather else None,
+        fp8_recipe="tensorwise" if fp8_all_gather else None,
+    )
+    model = Flux(flux_config).cuda().to(torch.bfloat16)
+
+    ddp_config = DistributedDataParallelConfig()
+
+    # Temporarily set the fp8 all-gather flag for FSDP wrapping
+    from megatron.training import get_args
+
+    args = get_args()
+    original_flag = getattr(args, "use_fsdp2_fp8_all_gather", False)
+    args.use_fsdp2_fp8_all_gather = fp8_all_gather
+
+    fsdp_wrapper = PrimusTorchFullyShardedDataParallel(
+        config=flux_config,
+        ddp_config=ddp_config,
+        module=model,
+    )
+
+    args.use_fsdp2_fp8_all_gather = original_flag
+
+    opt_config = _make_optimizer_config()
+    optimizer = get_fsdp2_bf16_master_weight_optimizer(
+        config=opt_config,
+        model_chunks=[model],
+    )
+    return model, fsdp_wrapper, optimizer
+
+
+def _train_steps(model, optimizer, n_steps, batch_size=2, device="cuda", fp8_all_gather=True):
+    from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+        precompute_fp8_scales_for_fsdp,
+    )
+
+    losses = []
+    for _ in range(n_steps):
+        optimizer.zero_grad()
+        inputs = _make_flux_inputs(batch_size, device)
+        output = model(**inputs)
+        loss = output.float().sum()
+        loss.backward()
+        optimizer.step()
+        if fp8_all_gather:
+            precompute_fp8_scales_for_fsdp(model)
+        losses.append(loss.item())
+    return losses
+
+
+class TestFluxFSDP2FP8AllGatherMultiGPU:
+
+    def setup_method(self):
+        _init_distributed()
+
+    def teardown_method(self):
+        _cleanup_distributed()
+
+    @requires_torchrun
+    def test_fsdp2_fp8_allgather_train_steps(self):
+        """Train 5 steps with FP8 all-gather, verify all losses are finite."""
+        model, _, optimizer = _create_model_and_optimizer(seed=42)
+        losses = _train_steps(model, optimizer, n_steps=5)
+        assert len(losses) == 5
+        for i, loss_val in enumerate(losses):
+            assert torch.isfinite(torch.tensor(loss_val)), f"Non-finite loss at step {i}: {loss_val}"
+
+    @requires_torchrun
+    def test_fsdp2_fp8_allgather_numerical_parity(self):
+        """Compare FP8-all-gather loss vs BF16-only baseline over 3 steps."""
+        torch.manual_seed(42)
+        model_fp8, _, opt_fp8 = _create_model_and_optimizer(
+            seed=42,
+            fp8_all_gather=True,
+        )
+        losses_fp8 = _train_steps(
+            model_fp8,
+            opt_fp8,
+            n_steps=3,
+            fp8_all_gather=True,
+        )
+
+        torch.manual_seed(42)
+        model_bf16, _, opt_bf16 = _create_model_and_optimizer(
+            seed=42,
+            fp8_all_gather=False,
+        )
+        losses_bf16 = _train_steps(
+            model_bf16,
+            opt_bf16,
+            n_steps=3,
+            fp8_all_gather=False,
+        )
+
+        for i, (fp8_loss, bf16_loss) in enumerate(zip(losses_fp8, losses_bf16)):
+            torch.testing.assert_close(
+                torch.tensor(fp8_loss),
+                torch.tensor(bf16_loss),
+                rtol=0.05,
+                atol=1.0,
+                msg=f"Loss diverged at step {i}: FP8={fp8_loss}, BF16={bf16_loss}",
+            )
+
+    @requires_torchrun
+    def test_fsdp2_fp8_allgather_with_compile(self):
+        """Train steps with torch.compile + FP8 all-gather."""
+        model, _, optimizer = _create_model_and_optimizer(seed=42)
+        model.compile_model()
+        losses = _train_steps(model, optimizer, n_steps=3)
+        assert len(losses) == 3
+        for i, loss_val in enumerate(losses):
+            assert torch.isfinite(torch.tensor(loss_val)), f"Non-finite loss at step {i}: {loss_val}"
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_integration.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_integration.py
new file mode 100644
index 000000000..b35e300af
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_fsdp2_integration.py
@@ -0,0 +1,182 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Integration tests for Flux model with FSDP2.
+
+Tests that Flux model works correctly with FSDP2 wrapping for both
+transformer_impl="local" and transformer_impl="transformer_engine".
+"""
+
+import pytest
+import torch
+from megatron.core.distributed.distributed_data_parallel_config import (
+    DistributedDataParallelConfig,
+)
+
+from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import (
+    PrimusTorchFullyShardedDataParallel,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    IMG_SIZE_TINY,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_SHORT,
+    VAE_LATENT_CHANNELS,
+)
+from tests.utils import PrimusUT
+
+
+class TestFluxFSDP2Integration(PrimusUT):
+    """Integration tests for Flux model with FSDP2."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for FSDP2 integration tests."""
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_model_fsdp2_wrapping_local(self):
+        """Test Flux model can be wrapped with FSDP2 when transformer_impl == 'local'."""
+        self._test_flux_model_fsdp2_wrapping("local")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_model_fsdp2_wrapping_te(self):
+        """Test Flux model can be wrapped with FSDP2 when transformer_impl == 'transformer_engine'."""
+        self._test_flux_model_fsdp2_wrapping("transformer_engine")
+
+    def _test_flux_model_fsdp2_wrapping(self, transformer_impl):
+        """Helper method to test FSDP2 wrapping with given transformer_impl."""
+        # Create Flux config with specified transformer_impl
+        config = FluxConfig.flux_535m(transformer_impl=transformer_impl)
+
+        # Create Flux model
+        model = Flux(config).cuda()
+        model.eval()
+
+        # Create FSDP2 wrapper
+        ddp_config = DistributedDataParallelConfig()
+
+        try:
+            fsdp_wrapper = PrimusTorchFullyShardedDataParallel(
+                config=config,
+                ddp_config=ddp_config,
+                module=model,
+            )
+
+            # Verify wrapper was created
+            assert fsdp_wrapper is not None
+            assert fsdp_wrapper.config.transformer_impl == transformer_impl
+
+        except Exception as e:
+            # If FSDP2 is not available, skip the test
+            if "FSDP" in str(e) or "fully_shard" in str(e).lower():
+                pytest.skip(f"FSDP2 not available: {e}")
+            raise
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_model_fsdp2_forward_pass_local(self):
+        """Test Flux model forward pass works with FSDP2 wrapping when transformer_impl == 'local'."""
+        self._test_flux_model_fsdp2_forward_pass("local")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_model_fsdp2_forward_pass_te(self):
+        """Test Flux model forward pass works with FSDP2 wrapping when transformer_impl == 'transformer_engine'."""
+        self._test_flux_model_fsdp2_forward_pass("transformer_engine")
+
+    def _test_flux_model_fsdp2_forward_pass(self, transformer_impl):
+        """Test Flux model forward pass works with FSDP2 wrapping."""
+        # Create Flux config with specified transformer_impl
+        config = FluxConfig.flux_535m(transformer_impl=transformer_impl)
+
+        # Create Flux model (bf16 required by FlashAttention on ROCm)
+        model = Flux(config).cuda().to(torch.bfloat16)
+        model.eval()
+
+        # Prepare inputs
+        batch_size = 2
+        height, width = IMG_SIZE_TINY, IMG_SIZE_TINY
+        channels = VAE_LATENT_CHANNELS
+        txt_seq_len = TEXT_SEQ_LEN_SHORT
+
+        img = torch.randn(batch_size, channels, height, width, dtype=torch.bfloat16).cuda()
+        txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        timesteps = torch.rand(batch_size, dtype=torch.bfloat16).cuda()
+
+        # Pack latents
+        packed_img = pack_latents(img)
+        packed_img = packed_img.transpose(0, 1)
+        txt_t = txt.transpose(0, 1)
+
+        # Generate position IDs
+        img_ids = generate_image_position_ids(batch_size, height, width, device="cuda")
+        txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda()
+
+        # Create FSDP2 wrapper
+        ddp_config = DistributedDataParallelConfig()
+
+        try:
+            fsdp_wrapper = PrimusTorchFullyShardedDataParallel(
+                config=config,
+                ddp_config=ddp_config,
+                module=model,
+            )
+
+            # Forward pass through FSDP-wrapped model
+            with torch.no_grad():
+                output = fsdp_wrapper.module(
+                    img=packed_img,
+                    txt=txt_t,
+                    y=y,
+                    timesteps=timesteps,
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                )
+
+            # Verify output shape
+            # Output is in sequence-first format: (seq_len, batch, channels)
+            # After packing 16x16 image: seq_len = (16*16)/4 = 64
+            assert output is not None
+            assert len(output.shape) == 3  # (seq_len, batch, channels)
+            assert output.shape[1] == batch_size  # Batch dimension is at index 1
+
+        except Exception as e:
+            # If FSDP2 is not available, skip the test
+            if "FSDP" in str(e) or "fully_shard" in str(e).lower():
+                pytest.skip(f"FSDP2 not available: {e}")
+            raise
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_flux_model_fsdp2_no_wrapping_errors(self):
+        """Test that FSDP2 wrapping doesn't produce duplicate wrapping errors."""
+        config = FluxConfig.flux_535m(transformer_impl="local")
+
+        model = Flux(config).cuda()
+        ddp_config = DistributedDataParallelConfig()
+
+        try:
+            # This should not raise "Invalid mesh_dim_names" or duplicate wrapping errors
+            fsdp_wrapper = PrimusTorchFullyShardedDataParallel(
+                config=config,
+                ddp_config=ddp_config,
+                module=model,
+            )
+
+            # If we get here, wrapping succeeded without errors
+            assert fsdp_wrapper is not None
+
+        except Exception as e:
+            error_msg = str(e).lower()
+            # Check for specific FSDP2 wrapping errors
+            if "mesh_dim_names" in error_msg or "duplicate" in error_msg:
+                pytest.fail(f"FSDP2 wrapping error detected: {e}")
+            # If FSDP2 is not available, skip the test
+            if "fsdp" in error_msg or "fully_shard" in error_msg:
+                pytest.skip(f"FSDP2 not available: {e}")
+            raise
diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py
new file mode 100644
index 000000000..6926817fd
--- /dev/null
+++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py
@@ -0,0 +1,203 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Integration tests for Flux model with MXFP4 local spec.
+
+Tests that a real (small) Flux model constructs, uses correct linear types,
+and produces valid output+gradients when using PrimusTurboMXFP4LocalSpecProvider.
+"""
+
+import pytest
+import torch
+
+from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+    MXFP4ColumnParallelLinear,
+    MXFP4RowParallelLinear,
+)
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+)
+from tests.unit_tests.backends.megatron.conftest import requires_mxfp4
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    IMG_SIZE_TINY,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_MEDIUM,
+    VAE_LATENT_CHANNELS,
+)
+from tests.utils import PrimusUT
+
+
+class TestFluxMXFP4LocalSpec(PrimusUT):
+    """Integration tests for Flux with MXFP4 local spec provider."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.fixture(autouse=True)
+    def _pin_fp4_aiter(self, monkeypatch):
+        """Pin the FP4 GEMM backend to AITER with autotune off.
+
+        MXFP4 module ``__init__`` runs ``_assert_preshuffle_contract``, which
+        requires the FP4 GEMM backend pinned to AITER with autotune off (the
+        only config under which ``_enable_preshuffle()`` is True). Without this,
+        Flux construction raises ``RuntimeError`` before the model is built.
+        Mirrors ``_pin_fp4_aiter`` in the unit MXFP4 tests
+        (test_primus_turbo_mxfp4_local.py); monkeypatch auto-restores on
+        teardown. Also clears any baked-empty ``PRIMUS_TURBO_GEMM_BACKEND`` so
+        the in-code pin is authoritative.
+
+        Use a ``defaultdict(None)`` rather than a bare ``{FP4: AITER}`` so that
+        the hybrid (FP4 fwd / FP8 bwd) test's FP8 lookup falls back to the
+        default backend instead of raising ``KeyError`` -- ``get_gemm_backend``
+        indexes ``_gemm_backend[precision]`` directly when it is non-None.
+        """
+        import collections
+        import os
+
+        from primus_turbo.pytorch.core.backend import (
+            BackendType,
+            GlobalBackendManager,
+            PrecisionType,
+        )
+
+        if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "":
+            monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False)
+        pinned = collections.defaultdict(lambda: None)
+        if GlobalBackendManager._gemm_backend:
+            pinned.update(GlobalBackendManager._gemm_backend)
+        pinned[PrecisionType.FP4] = BackendType.AITER
+        monkeypatch.setattr(GlobalBackendManager, "_gemm_backend", pinned)
+        monkeypatch.setattr(GlobalBackendManager, "_auto_tune", False)
+
+    def _make_mxfp4_config(self, **overrides):
+        defaults = dict(
+            transformer_impl="local",
+            fp4="mxfp4",
+            fp4_recipe="mxfp4",
+        )
+        defaults.update(overrides)
+        return FluxConfig.flux_535m(**defaults)
+
+    def _make_inputs(self, batch_size=2):
+        height, width = IMG_SIZE_TINY, IMG_SIZE_TINY
+        channels = VAE_LATENT_CHANNELS
+        # FP4 GEMM requires M divisible by 16. TEXT_SEQ_LEN_SHORT=77 is not
+        # aligned; use 128 so batch*seq = 256 (multiple of 16).
+        txt_seq_len = TEXT_SEQ_LEN_MEDIUM
+
+        img = torch.randn(batch_size, channels, height, width, dtype=torch.bfloat16).cuda()
+        txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM, dtype=torch.bfloat16).cuda()
+        timesteps = torch.rand(batch_size, dtype=torch.bfloat16).cuda()
+
+        packed_img = pack_latents(img)
+        packed_img = packed_img.transpose(0, 1)
+        txt_t = txt.transpose(0, 1)
+
+        img_ids = generate_image_position_ids(batch_size, height, width, device="cuda")
+        txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda()
+
+        return packed_img, txt_t, y, timesteps, img_ids, txt_ids
+
+    @requires_mxfp4
+    def test_flux_535m_mxfp4_constructs(self):
+        config = self._make_mxfp4_config()
+        model = Flux(config)
+        assert model is not None
+        assert isinstance(model, Flux)
+
+    @requires_mxfp4
+    def test_flux_535m_mxfp4_linear_types(self):
+        """Verify spec-provided linears are MXFP4 variants."""
+        config = self._make_mxfp4_config()
+        model = Flux(config)
+
+        spec_linear_names = {
+            "linear_qkv",
+            "added_linear_qkv",
+            "linear_proj",
+            "linear_fc1",
+            "linear_fc2",
+        }
+
+        from megatron.core.tensor_parallel.layers import (
+            ColumnParallelLinear,
+            RowParallelLinear,
+        )
+
+        found_any = False
+        for name, module in model.named_modules():
+            leaf_name = name.rsplit(".", 1)[-1] if "." in name else name
+            if leaf_name not in spec_linear_names:
+                continue
+            found_any = True
+            if isinstance(module, ColumnParallelLinear):
+                assert isinstance(
+                    module, MXFP4ColumnParallelLinear
+                ), f"{name}: expected MXFP4ColumnParallelLinear, got {type(module).__name__}"
+            if isinstance(module, RowParallelLinear):
+                assert isinstance(
+                    module, MXFP4RowParallelLinear
+                ), f"{name}: expected MXFP4RowParallelLinear, got {type(module).__name__}"
+
+        assert found_any, "No spec-provided linears found in model"
+
+    @requires_mxfp4
+    def test_flux_535m_mxfp4_forward_backward(self):
+        """Forward+backward pass produces valid output and gradient flow."""
+        config = self._make_mxfp4_config()
+        model = Flux(config).cuda().to(torch.bfloat16)
+        model.train()
+
+        inputs = self._make_inputs(batch_size=2)
+
+        output = model(*inputs)
+
+        assert output is not None
+        assert len(output.shape) == 3
+        assert output.shape[1] == 2
+        assert not torch.isnan(output).any(), "Output contains NaN"
+        assert not torch.isinf(output).any(), "Output contains Inf"
+
+        output.sum().backward()
+
+        has_mxfp4_grad = False
+        for name, module in model.named_modules():
+            if isinstance(module, (MXFP4ColumnParallelLinear, MXFP4RowParallelLinear)):
+                if module.weight.grad is not None:
+                    has_mxfp4_grad = True
+                    break
+
+        assert has_mxfp4_grad, "No MXFP4 linear has a weight gradient -- gradient did not flow"
+
+    @requires_mxfp4
+    def test_flux_535m_mxfp4_hybrid_forward_backward(self):
+        """Hybrid backward (FP4 fwd / FP8 bwd) works through the full model."""
+        config = self._make_mxfp4_config(mxfp4_backward_precision="fp8")
+        model = Flux(config).cuda().to(torch.bfloat16)
+        model.train()
+
+        inputs = self._make_inputs(batch_size=2)
+
+        output = model(*inputs)
+
+        assert output is not None
+        assert not torch.isnan(output).any(), "Hybrid output contains NaN"
+        assert not torch.isinf(output).any(), "Hybrid output contains Inf"
+
+        output.sum().backward()
+
+        has_mxfp4_grad = False
+        for name, module in model.named_modules():
+            if isinstance(module, (MXFP4ColumnParallelLinear, MXFP4RowParallelLinear)):
+                if module.weight.grad is not None:
+                    has_mxfp4_grad = True
+                    break
+
+        assert has_mxfp4_grad, "No MXFP4 linear has a weight gradient in hybrid mode"
diff --git a/tests/runner/test_primus_cli_direct.sh b/tests/runner/test_primus_cli_direct.sh
index a7c6582cb..93fa50d1b 100755
--- a/tests/runner/test_primus_cli_direct.sh
+++ b/tests/runner/test_primus_cli_direct.sh
@@ -338,6 +338,277 @@ test_help_output() {
     assert_contains "$output" "--numa" "NUMA option documented"
     assert_contains "$output" "--env" "Env option documented"
     assert_contains "$output" "--patch" "Patch option documented"
+    assert_contains "$output" "--silent" "Silent flag documented"
+    assert_contains "$output" "VENV_ACTIVATE" "VENV_ACTIVATE documented"
+    assert_contains "$output" "SLURM" "SLURM auto-derivation documented"
+}
+
+# ============================================================================
+# Test 11: VENV_ACTIVATE handling (R1 -- consolidate-preflight-direct-wrappers)
+# ============================================================================
+test_venv_activate() {
+    local_print_section "Test 11: VENV_ACTIVATE (R1)"
+
+    # Sub-test 11a: VENV_ACTIVATE unset = no-op (dry-run succeeds without
+    # any "VENV_ACTIVATE" error, and the script never tries to source a file).
+    local out_unset
+    out_unset=$(unset VENV_ACTIVATE; timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_unset" "End of Dry Run" "Unset VENV_ACTIVATE is a no-op (dry-run completes)"
+    assert_not_contains "$out_unset" "VENV_ACTIVATE is set but" "No spurious 'missing file' error when unset"
+    assert_not_contains "$out_unset" "Activated virtualenv:" "No 'Activated virtualenv' message when unset"
+
+    # Sub-test 11b: VENV_ACTIVATE set + valid file = sourced.
+    local tmpvenv
+    tmpvenv="$(mktemp)"
+    echo 'export PRIMUS_TEST_VENV_SOURCED=yes' > "$tmpvenv"
+    local out_valid
+    out_valid=$(VENV_ACTIVATE="$tmpvenv" timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_valid" "Activated virtualenv:" "Valid VENV_ACTIVATE is sourced"
+    rm -f "$tmpvenv"
+
+    # Sub-test 11c: VENV_ACTIVATE set + missing file = fail-fast (LOG_ERROR
+    # on stderr, non-zero exit).
+    local ec_missing=0
+    VENV_ACTIVATE=/does/not/exist/anywhere timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm >/dev/null 2>/tmp/test_venv_err_$$ || ec_missing=$?
+    if [[ "$ec_missing" -ne 0 ]] && grep -q "VENV_ACTIVATE is set but file does not exist" /tmp/test_venv_err_$$; then
+        assert_pass "Missing VENV_ACTIVATE file fails fast with clear error"
+    else
+        assert_fail "Missing VENV_ACTIVATE file should fail fast" \
+            "exit=$ec_missing stderr=$(cat /tmp/test_venv_err_$$)"
+    fi
+    rm -f /tmp/test_venv_err_$$
+}
+
+# ============================================================================
+# Test 12: SLURM env derivation (R2)
+# ============================================================================
+test_slurm_env_derivation() {
+    local_print_section "Test 12: SLURM env derivation (R2)"
+
+    # Sub-test 12a: SLURM_JOB_ID + SLURM_NNODES + SLURM_NODEID -> NNODES /
+    # NODE_RANK derived.
+    local out_slurm
+    out_slurm=$(SLURM_JOB_ID=999 SLURM_NNODES=4 SLURM_NODEID=0 SLURM_NODELIST=node001 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_slurm" "SLURM detected" "SLURM detection log fires"
+    assert_contains "$out_slurm" "NNODES=4" "NNODES derived from SLURM_NNODES"
+    assert_contains "$out_slurm" "--nnodes 4" "torchrun gets --nnodes 4"
+
+    # Sub-test 12b: pre-exported NNODES wins over SLURM_NNODES.
+    local out_preset
+    out_preset=$(SLURM_JOB_ID=999 SLURM_NNODES=4 NNODES=7 NODE_RANK=0 SLURM_NODELIST=node001 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_preset" "NNODES=7" "Pre-exported NNODES=7 wins over SLURM_NNODES=4"
+    assert_contains "$out_preset" "--nnodes 7" "torchrun honors pre-exported NNODES=7"
+
+    # Sub-test 12c: sanity check rejects NODE_RANK >= NNODES.
+    local ec_sanity=0
+    NNODES=2 NODE_RANK=5 timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm >/dev/null 2>/tmp/test_slurm_err_$$ || ec_sanity=$?
+    if [[ "$ec_sanity" -ne 0 ]] && grep -q "NODE_RANK (5) must be < NNODES (2)" /tmp/test_slurm_err_$$; then
+        assert_pass "Sanity check rejects NODE_RANK >= NNODES"
+    else
+        assert_fail "Sanity check should reject NODE_RANK >= NNODES" \
+            "exit=$ec_sanity stderr=$(cat /tmp/test_slurm_err_$$)"
+    fi
+    rm -f /tmp/test_slurm_err_$$
+
+    # Sub-test 12d: SLURM context without SLURM_NODELIST must still complete.
+    # This exercises the same code path as "no scontrol on the host" (the
+    # `command -v scontrol && [[ -n "$SLURM_NODELIST" ]]` short-circuits to
+    # false either way), and used to crash with "MASTER_ADDR: unbound
+    # variable" under set -u. The launcher must fall back to MASTER_ADDR=localhost
+    # and the dry-run must reach the torchrun command line.
+    local out_no_nodelist
+    out_no_nodelist=$(env -u SLURM_NODELIST SLURM_JOB_ID=999 SLURM_NNODES=4 SLURM_NODEID=0 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    assert_not_contains "$out_no_nodelist" "unbound variable" \
+        "SLURM context with no SLURM_NODELIST does not crash on unbound MASTER_ADDR"
+    assert_contains "$out_no_nodelist" "MASTER_ADDR=localhost" \
+        "MASTER_ADDR falls back to localhost when scontrol/NODELIST unavailable"
+    assert_contains "$out_no_nodelist" "--nnodes 4" \
+        "torchrun cmd still gets --nnodes 4 even without NODELIST"
+}
+
+# ============================================================================
+# Test 13: Auto-single run_mode for node_smoke
+# ============================================================================
+test_auto_single_for_node_smoke() {
+    local_print_section "Test 13: Auto-single run_mode for node_smoke"
+
+    # Sub-test 13a: preflight stays in torchrun mode by default.
+    local out_preflight
+    out_preflight=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- preflight --quick 2>&1 || true)
+    assert_contains "$out_preflight" "Run Mode        : torchrun" "preflight defaults to torchrun"
+    assert_contains "$out_preflight" "torchrun --nproc_per_node" "preflight uses torchrun command"
+
+    # Sub-test 13b: node_smoke auto-selects single mode.
+    local out_smoke
+    out_smoke=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- node_smoke --tier2-perf 2>&1 || true)
+    assert_contains "$out_smoke" "Auto-selected run_mode=single for subcommand 'node_smoke'" \
+        "node_smoke auto-detect fires"
+    assert_contains "$out_smoke" "Run Mode        : single" "node_smoke ends up in single mode"
+    assert_contains "$out_smoke" "python3" "node_smoke uses python3 launcher"
+    assert_not_contains "$out_smoke" "torchrun --nproc_per_node" "node_smoke does NOT use torchrun"
+
+    # Sub-test 13c: explicit --single on a non-node_smoke subcommand still
+    # works (regression guard).
+    local out_explicit
+    out_explicit=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --single --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_explicit" "Run Mode        : single" "Explicit --single still works for benchmark"
+}
+
+# ============================================================================
+# Test 14: --silent contract (bash-level)
+# ============================================================================
+test_silent_flag() {
+    local_print_section "Test 14: --silent flag contract"
+
+    # Sub-test 14a: --silent produces empty stdout.
+    local out_silent="/tmp/test_silent_out_$$"
+    local err_silent="/tmp/test_silent_err_$$"
+    timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --silent --dry-run -- benchmark gemm >"$out_silent" 2>"$err_silent" || true
+    if [[ ! -s "$out_silent" ]]; then
+        assert_pass "--silent silences stdout (out file is empty)"
+    else
+        assert_fail "--silent should silence stdout" "size=$(wc -c <"$out_silent") head=$(head -3 "$out_silent")"
+    fi
+    rm -f "$out_silent" "$err_silent"
+
+    # Sub-test 14b: --silent does NOT silence launcher errors (stderr survives).
+    local err_err_silent="/tmp/test_silent_err2_$$"
+    VENV_ACTIVATE=/does/not/exist timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --silent --dry-run -- benchmark gemm >/dev/null 2>"$err_err_silent" || true
+    if grep -q "VENV_ACTIVATE is set but file does not exist" "$err_err_silent"; then
+        assert_pass "--silent preserves launcher LOG_ERROR on stderr"
+    else
+        assert_fail "--silent should preserve launcher errors on stderr" \
+            "stderr=$(cat "$err_err_silent")"
+    fi
+    rm -f "$err_err_silent"
+
+    # Sub-test 14c: --silent is consumed by the launcher and NOT forwarded
+    # to the python tool. The forwarded args list (logged before silencing
+    # applied? actually after -- it shouldn't appear at all). Easiest check:
+    # the dry-run "Would Execute" line must NOT contain --silent.
+    local out_no_forward
+    out_no_forward=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true)
+    # Sanity: this baseline run has no --silent anywhere.
+    assert_not_contains "$out_no_forward" "--silent" "Baseline dry-run has no --silent leakage"
+}
+
+# ============================================================================
+# Test 15: slurm-entry direct mode dispatch
+# ============================================================================
+test_slurm_entry_direct_dispatch() {
+    local_print_section "Test 15: slurm-entry direct/container dispatch"
+
+    # NOTE: primus-cli-slurm-entry.sh resolves the node list via scontrol when
+    # available and falls back to parsing SLURM_NODELIST directly otherwise, so
+    # this test works both on-cluster (scontrol present) and in CI containers
+    # (scontrol absent). We set a plain SLURM_NODELIST so both paths agree.
+    # Sub-test 15a: -- direct -- preflight ... routes through primus-cli-direct.sh.
+    local out_direct
+    out_direct=$(SLURM_NODELIST=node001 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- direct -- preflight --quick 2>&1 || true)
+    assert_contains "$out_direct" "Entry mode: direct" "slurm-entry parses 'direct' keyword"
+    assert_contains "$out_direct" "primus-cli-direct.sh" "slurm-entry dispatches to primus-cli-direct.sh"
+
+    # Sub-test 15b: -- container -- ... routes through primus-cli-container.sh (existing path).
+    local out_container
+    out_container=$(SLURM_NODELIST=node001 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- container -- train pretrain 2>&1 || true)
+    assert_contains "$out_container" "Entry mode: container" "slurm-entry parses 'container' keyword"
+    assert_contains "$out_container" "primus-cli-container.sh" "slurm-entry dispatches to primus-cli-container.sh"
+
+    # Sub-test 15c: terse form `-- preflight` (no keyword) defaults to container.
+    local out_terse
+    out_terse=$(SLURM_NODELIST=node001 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \
+        timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- preflight --quick 2>&1 || true)
+    assert_contains "$out_terse" "Entry mode: container" "Terse form defaults to container"
+    assert_contains "$out_terse" "primus-cli-container.sh" "Terse form dispatches to primus-cli-container.sh"
+}
+
+# ============================================================================
+# Test 16: RUN_MODE env override (framework prepare-hook contract)
+#
+# Framework hooks (e.g. runner/helpers/hooks/train/pretrain/maxtext/prepare.py)
+# emit `env.RUN_MODE=single` so JAX/MaxText runs as `python3 ...` instead of
+# `torchrun ...`. This used to launch correctly but display incorrectly:
+# STEP 10 showed `Run Mode: torchrun` and printed the Distributed Settings
+# block, because both reads went against the pre-hook direct_config[run_mode]
+# instead of the final $RUN_MODE.
+#
+# This test pins the post-fix contract: when RUN_MODE is exported in the env
+# before primus-cli-direct.sh runs (a sufficient stand-in for "a hook
+# exported it"), every visible knob -- the displayed Run Mode, the
+# Distributed Settings gate, and the actual launch command -- reflects the
+# env-override, not the pre-hook default.
+# ============================================================================
+test_run_mode_env_override() {
+    local_print_section "Test 16: RUN_MODE env override (framework hook contract)"
+
+    # Baseline: `train pretrain` with NO hook override -- the launcher's
+    # auto-default (torchrun) wins and the display shows torchrun + the
+    # Distributed Settings block. This anchors the "before" state we are
+    # protecting users from.
+    local out_baseline
+    out_baseline=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- train pretrain 2>&1 || true)
+    assert_contains "$out_baseline" "Run Mode        : torchrun" \
+        "Baseline train pretrain shows torchrun (no env override)"
+    assert_contains "$out_baseline" "Distributed Settings:" \
+        "Baseline train pretrain prints Distributed Settings"
+    assert_contains "$out_baseline" "torchrun --nproc_per_node" \
+        "Baseline train pretrain uses torchrun in Full Command"
+
+    # Sub-test 16a: pre-exporting RUN_MODE=single (the MaxText hook's
+    # effect, modeled in-process so the test doesn't depend on the
+    # framework hook actually firing) must flip ALL three views:
+    #   - displayed "Run Mode" line
+    #   - "Distributed Settings" gate
+    #   - launch command (`python3 ...`, not `torchrun ...`)
+    local out_override
+    out_override=$(RUN_MODE=single timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \
+        --dry-run -- train pretrain 2>&1 || true)
+    assert_contains "$out_override" "Run Mode        : single" \
+        "RUN_MODE=single env override surfaces in displayed Run Mode"
+    assert_not_contains "$out_override" "Distributed Settings:" \
+        "RUN_MODE=single env override suppresses Distributed Settings block"
+    assert_not_contains "$out_override" "torchrun --nproc_per_node" \
+        "RUN_MODE=single env override drops torchrun from Full Command"
+    assert_contains "$out_override" "python3" \
+        "RUN_MODE=single env override uses python3 launcher in Full Command"
+
+    # Sub-test 16b: env-override beats an EXPLICIT --single on the CLI.
+    # This matches the launcher's documented precedence at STEP 9:
+    #   RUN_MODE="${RUN_MODE:-${direct_config[run_mode]:-torchrun}}"
+    # The env layer wins because the hook layer (which is where RUN_MODE
+    # actually originates in real runs) knows things the user / config
+    # can't and gets the final word. We use --single here purely as a
+    # convenient stand-in for "user/config set direct_config[run_mode]"
+    # -- there is no symmetric --torchrun CLI flag.
+    local out_vs_cli
+    out_vs_cli=$(RUN_MODE=torchrun timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \
+        --single --dry-run -- benchmark gemm 2>&1 || true)
+    assert_contains "$out_vs_cli" "Run Mode        : torchrun" \
+        "RUN_MODE=torchrun env override wins over --single CLI flag"
+    assert_contains "$out_vs_cli" "torchrun --nproc_per_node" \
+        "RUN_MODE=torchrun env override produces torchrun launcher despite --single"
+
+    # Sub-test 16c: exporting RUN_MODE must override even the auto-single
+    # detection for node_smoke. Auto-detect sets direct_config[run_mode]
+    # BEFORE the env-override check at line 610, so env still wins. This
+    # is the corner case behind the bug we fixed: the OLD code printed
+    # `Run Mode: single` here (from direct_config[run_mode]) even though
+    # the launch command was actually torchrun -- a display divergence
+    # in the exact opposite direction from the MaxText case.
+    local out_reverse
+    out_reverse=$(RUN_MODE=torchrun timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \
+        --dry-run -- node_smoke --tier2-perf 2>&1 || true)
+    assert_contains "$out_reverse" "Run Mode        : torchrun" \
+        "RUN_MODE=torchrun env override surfaces in display for node_smoke"
+    assert_contains "$out_reverse" "Distributed Settings:" \
+        "RUN_MODE=torchrun env override re-enables Distributed Settings for node_smoke"
+    assert_contains "$out_reverse" "torchrun --nproc_per_node" \
+        "RUN_MODE=torchrun env override produces torchrun launch for node_smoke"
 }
 
 # ============================================================================
@@ -358,6 +629,13 @@ main() {
     test_debug_mode
     test_config_priority
     test_help_output
+    # New tests from the consolidate-preflight-direct-wrappers refactor:
+    test_venv_activate
+    test_slurm_env_derivation
+    test_auto_single_for_node_smoke
+    test_silent_flag
+    test_slurm_entry_direct_dispatch
+    test_run_mode_env_override
 
     # Print summary
     echo ""
diff --git a/tests/trainer/fixtures/sft_lora_smoke.jsonl b/tests/trainer/fixtures/sft_lora_smoke.jsonl
new file mode 100644
index 000000000..6773bd5dd
--- /dev/null
+++ b/tests/trainer/fixtures/sft_lora_smoke.jsonl
@@ -0,0 +1,8 @@
+{"instruction": "What is the capital of France?", "input": "", "output": "The capital of France is Paris."}
+{"instruction": "Translate the following sentence into French.", "input": "I love programming.", "output": "J'aime la programmation."}
+{"instruction": "Summarize the given text in one sentence.", "input": "The quick brown fox jumps over the lazy dog. It is a common English pangram used to test typefaces.", "output": "A pangram about a fox jumping over a dog is used to test typefaces."}
+{"instruction": "Give three tips for staying healthy.", "input": "", "output": "1. Eat a balanced diet. 2. Exercise regularly. 3. Get enough sleep."}
+{"instruction": "Convert the temperature from Celsius to Fahrenheit.", "input": "100 degrees Celsius", "output": "100 degrees Celsius is 212 degrees Fahrenheit."}
+{"instruction": "Write a short greeting for a new employee.", "input": "", "output": "Welcome to the team! We're excited to have you on board."}
+{"instruction": "Explain what an API is in one sentence.", "input": "", "output": "An API is a set of rules that lets different software programs communicate with each other."}
+{"instruction": "Rewrite the sentence to be more formal.", "input": "hey can u send me that file", "output": "Could you please send me that file?"}
diff --git a/tests/trainer/test_megatron_trainer.py b/tests/trainer/test_megatron_trainer.py
index 20152dc78..ef63bea8c 100644
--- a/tests/trainer/test_megatron_trainer.py
+++ b/tests/trainer/test_megatron_trainer.py
@@ -12,6 +12,8 @@
 import sys
 import unittest
 
+import pytest
+
 from tests.utils import PrimusUT, run_training_script
 
 
@@ -83,7 +85,66 @@ def run_script(
     return run_training_script(tag=tag, cmd=cmd, train_log_path=train_log_path, env=env)
 
 
+def run_posttrain_script(
+    ut_name: str,
+    tag: str,
+    exp_path: str,
+    env_override: dict = None,
+    extra_args: list[str] = None,
+):
+    """Like run_script, but for the "posttrain" suite (SFT/alignment).
+
+    SFT experiment configs declare a `modules.post_trainer` section (instead
+    of `modules.pre_trainer`), which the CLI only loads via `train posttrain`
+    (see primus/cli/subcommands/train.py). The "Training completed." marker
+    that run_training_script asserts on is emitted generically by
+    PrimusRuntime._run_trainer_lifecycle for any module, so it applies here
+    unchanged.
+    """
+    shell_entry = "./runner/primus-cli"
+    env = os.environ.copy()
+    if env_override:
+        env.update(env_override)
+    env["EXP"] = exp_path
+
+    ut_log_path = os.environ.get("UT_LOG_PATH", "ut_out")
+    train_log_path = os.path.join(ut_log_path, f"log.test_megatron_trainer-{tag}.txt")
+    env["TRAIN_LOG"] = train_log_path
+
+    cmd = [
+        "bash",
+        shell_entry,
+        "direct",
+        "--log_file",
+        train_log_path,
+        "--",
+        "train",
+        "posttrain",
+        "--config",
+        exp_path,
+    ]
+    if extra_args:
+        cmd.extend(extra_args)
+
+    return run_training_script(tag=tag, cmd=cmd, train_log_path=train_log_path, env=env)
+
+
 class TestMegatronTrainer(PrimusUT):
+    """One model per architecture, plus one case per feature path.
+
+    Every case is a full training launch, so a model that only scales the dims
+    of a sibling earns no coverage and does not belong here: qwen3_235B_A22B
+    (same template as qwen3_30B_A3B) and mixtral_8x22B (same template as
+    mixtral_8x7B, whose recipe additionally covers DeepEP / sync-free MoE) were
+    removed for that reason. Their example yamls are still schema-checked by
+    tests/unit_tests/configs/test_example_configs.py.
+
+    Before adding a model, check it is not isomorphic to one already here, and
+    before removing one, check it is not the last user of a feature path --
+    llama3_70B for instance is the only recipe that enables use_torch_fsdp2.
+    See tests/README.md.
+    """
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -111,6 +172,26 @@ def test_llama3_8B(self):
             ],
         )
 
+    # The only dense FP8 case; every other FP8 case here is MoE.
+    @pytest.mark.weekly
+    def test_llama3_8B_fp8(self):
+        run_script(
+            self.__class__.__name__,
+            "llama3_8B_fp8",
+            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/llama3_8B-FP8-pretrain.yaml",
+            env_override={},
+            extra_args=[
+                "--num_layers",
+                "4",
+                "--train_iters",
+                "3",
+                "--enable_primus_turbo",
+                "1",
+                "--use_turbo_attention",
+                "1",
+            ],
+        )
+
     def test_llama3_1_8B_tp2_distributed_dataset_regression(self):
         run_script(
             self.__class__.__name__,
@@ -183,40 +264,9 @@ def test_qwen3_30B_A3B(self):
             ],
         )
 
-    def test_qwen3_235B_A22B(self):
-        run_script(
-            self.__class__.__name__,
-            "qwen3_235B_A22B",
-            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/qwen3_235B_A22B-BF16-pretrain.yaml",
-            env_override={},
-            extra_args=[
-                "--num_layers",
-                "4",
-                "--moe_layer_freq",
-                "[0]*1+[1]*3",
-                "--train_iters",
-                "3",
-                "--micro_batch_size",
-                "1",
-                "--global_batch_size",
-                "8",
-                "--expert_model_parallel_size",
-                "8",
-                "--pipeline_model_parallel_size",
-                "1",
-                "--recompute_granularity",
-                "full",
-                "--recompute_method",
-                "block",
-                "--recompute_num_layers",
-                "0",
-                "--enable_primus_turbo",
-                "1",
-                "--use_turbo_attention",
-                "1",
-            ],
-        )
-
+    # Its gated_delta_net attention is an experimental variant and needs
+    # flash-linear-attention; qwen3 MoE stays covered per-PR by qwen3_30B_A3B.
+    @pytest.mark.weekly
     def test_qwen3_5_35B_A3B(self):
         run_script(
             self.__class__.__name__,
@@ -301,6 +351,9 @@ def test_gpt_oss_20B_sink_attention(self):
             ],
         )
 
+    # Secondary architecture; MoE with expert parallelism stays covered per-PR by
+    # deepseek_v2_lite and qwen3_30B_A3B.
+    @pytest.mark.weekly
     def test_mixtral_8x7B(self):
         run_script(
             self.__class__.__name__,
@@ -327,11 +380,14 @@ def test_mixtral_8x7B(self):
             ],
         )
 
-    def test_mixtral_8x22B(self):
+    # Its own architecture: rope scaling, qk_layernorm and shared-expert overlap
+    # on every layer.
+    @pytest.mark.weekly
+    def test_llama4_17B16E(self):
         run_script(
             self.__class__.__name__,
-            "mixtral_8x22B_v0.1",
-            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/mixtral_8x22B_v0.1-BF16-pretrain.yaml",
+            "llama4_17B16E",
+            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/llama4_17B16E-BF16-pretrain.yaml",
             env_override={},
             extra_args=[
                 "--num_layers",
@@ -342,19 +398,14 @@ def test_mixtral_8x22B(self):
                 "1",
                 "--global_batch_size",
                 "8",
-                "--moe_layer_freq",
-                "1",
                 "--expert_model_parallel_size",
                 "8",
-                "--pipeline_model_parallel_size",
-                "1",
-                "--enable_primus_turbo",
-                "1",
-                "--use_turbo_attention",
-                "1",
             ],
         )
 
+    # Secondary architecture; PP+VPP stay covered per-PR by deepseek_v3 and
+    # test_interleaved_pipeline_parallelism.
+    @pytest.mark.weekly
     def test_grok2(self):
         run_script(
             self.__class__.__name__,
@@ -480,6 +531,10 @@ def test_turbo_grouped_gemm(self):
                 "1",
                 "--use_turbo_grouped_gemm",
                 "1",
+                # use_turbo_grouped_gemm is incompatible with the config's default
+                # moe_use_legacy_grouped_gemm=True, so disable the legacy path.
+                "--moe_use_legacy_grouped_gemm",
+                "0",
             ],
         )
 
@@ -546,6 +601,9 @@ def test_turbo_deepep(self):
                 "0",
                 "--moe_use_legacy_grouped_gemm",
                 "0",
+                # Sync-Free MoE stage 3 requires PrimusTurboGroupedLinear.
+                "--use_turbo_grouped_gemm",
+                "1",
                 "--turbo_sync_free_moe_stage",
                 "3",
                 "--use_turbo_attention",
@@ -562,6 +620,121 @@ def test_turbo_deepep(self):
             Dataloader_mp_context_patch_log in stdout
         ), "Expected dataloader_mp_context patch log not found in stdout"
 
+    # A fused-op combination on top of llama3_8B, which stays covered per-PR.
+    @pytest.mark.weekly
+    def test_sdma_allgather_fused_residual_norm(self):
+        # Neither patch runs in any other case here: SDMA param all-gather
+        # only activates with ENABLE_SDMA_ALLGATHER=1 (env var, not a --arg) +
+        # a distributed optimizer; fused residual+RMSNorm needs
+        # use_turbo_rms_norm=1 + PRIMUS_FUSED_RESIDUAL_NORM_V2=1. Confirmed via
+        # coverage instrumentation that both patches install and execute.
+        run_script(
+            self.__class__.__name__,
+            "sdma_allgather_fused_residual_norm",
+            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/llama3_8B-BF16-pretrain.yaml",
+            env_override={
+                "ENABLE_SDMA_ALLGATHER": "1",
+                "PRIMUS_FUSED_RESIDUAL_NORM_V2": "1",
+            },
+            extra_args=[
+                "--num_layers",
+                "4",
+                "--train_iters",
+                "3",
+                "--enable_primus_turbo",
+                "1",
+                "--use_turbo_attention",
+                "1",
+                "--use_turbo_rms_norm",
+                "1",
+                "--use_distributed_optimizer",
+                "1",
+                "--overlap_param_gather",
+                "1",
+            ],
+        )
+
+    def test_mamba_370M(self):
+        # Default `auto` attention backend exercises attention_backend_patches:
+        # megatron-core's probe must respect the image's baked NVTE_FLASH_ATTN=0.
+        run_script(
+            self.__class__.__name__,
+            "mamba_370M",
+            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/mamba_370M-pretrain.yaml",
+            env_override={},
+            extra_args=[
+                "--num_layers",
+                "4",
+                "--train_iters",
+                "3",
+                "--micro_batch_size",
+                "2",
+                "--global_batch_size",
+                "16",
+            ],
+        )
+
+    def test_zebra_llama_1B_hybrid(self):
+        # Hybrid Mamba+MLA (HybridStack) path. num_layers=8 is the minimum that
+        # keeps the default hybrid_attention_ratio=0.25 from allocating zero
+        # attention layers (division by zero).
+        run_script(
+            self.__class__.__name__,
+            "zebra_llama_1B_hybrid",
+            exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/zebra_llama_1B-pretrain.yaml",
+            env_override={},
+            extra_args=[
+                "--num_layers",
+                "8",
+                "--train_iters",
+                "3",
+                "--micro_batch_size",
+                "2",
+                "--global_batch_size",
+                "16",
+            ],
+        )
+
+    def test_mamba_130M_bridge_pretrain(self):
+        # Only E2E covering the megatron_bridge backend (mamba/zebra above use
+        # the megatron backend). extra_args pin a tiny shape so the test doesn't
+        # depend on the example yaml's sizes. Don't override seq_length: the
+        # recipe feeds it to both model and dataset but a CLI override reaches
+        # only the dataset, and Bridge asserts the two match.
+        run_script(
+            self.__class__.__name__,
+            "mamba_130M_bridge_pretrain",
+            exp_path=f"examples/megatron_bridge/configs/{GPU_PLATFORM}/mamba_130M_pretrain.yaml",
+            env_override={},
+            extra_args=[
+                "--train_iters",
+                "3",
+                "--micro_batch_size",
+                "1",
+                "--global_batch_size",
+                "8",
+            ],
+        )
+
+    def test_qwen2_sft_lora(self):
+        # Only E2E covering the "posttrain" suite (MegatronSFTTrainer) and
+        # peft/*.py; LoRA is enabled in test_megatron_trainer_sft_lora.yaml.
+        # The posttrain hook HF->Megatron-converts `tokenizer_model` into the
+        # base checkpoint when pretrained_checkpoint/load are unset, so this
+        # uses a tiny stand-in checkpoint, not real Qwen2.5-7B weights.
+        run_posttrain_script(
+            self.__class__.__name__,
+            "qwen2_sft_lora",
+            exp_path="tests/trainer/test_megatron_trainer_sft_lora.yaml",
+            env_override={},
+            # head_dim (2) is below the image's fused-attention CK kernel minimum,
+            # so pin the unfused path; attention_backend_patches forces it over
+            # the image's baked NVTE_FUSED_ATTN=1.
+            extra_args=["--attention_backend", "unfused"],
+        )
+
+    # UEP variant; the deepep path stays covered per-PR by test_turbo_deepep.
+    @pytest.mark.weekly
     def test_deepseekv2_lite_uep(self):
         run_script(
             self.__class__.__name__,
@@ -644,6 +817,10 @@ def _run_deepseek_v2_lite_zbv_fp8_case(
         self.assertIn("Training completed.", stdout)
         return stdout
 
+    # The three zbv cases below differ only in wgrad-split backend (TE / turbo /
+    # legacy grouped GEMM), and at 8 layers with PP=4 VPP=2 each costs more than
+    # the 4-layer cases above.
+    @pytest.mark.weekly
     def test_deepseek_v2_lite_te_fp8_zbv_formatted(self):
         stdout = self._run_deepseek_v2_lite_zbv_fp8_case(
             "deepseek_v2_lite_te_fp8_zbv_formatted",
@@ -655,6 +832,7 @@ def test_deepseek_v2_lite_te_fp8_zbv_formatted(self):
         self.assertIn("[Patch:megatron.pp.te_wgrad_split]", stdout)
         self.assertNotIn("[Patch:megatron.pp.legacy_grouped_mlp_wgrad_split]", stdout)
 
+    @pytest.mark.weekly
     def test_deepseek_v2_lite_turbo_bf16_zbv_formatted(self):
         stdout = self._run_deepseek_v2_lite_zbv_fp8_case(
             "deepseek_v2_lite_turbo_fp8_zbv_formatted",
@@ -671,6 +849,7 @@ def test_deepseek_v2_lite_turbo_bf16_zbv_formatted(self):
         )
         self.assertNotIn("[Patch:megatron.pp.legacy_grouped_mlp_wgrad_split]", stdout)
 
+    @pytest.mark.weekly
     def test_deepseek_v2_lite_bf16_lagacy_gg_zbv_formatted(self):
         stdout = self._run_deepseek_v2_lite_zbv_fp8_case(
             "deepseek_v2_lite_bf16_lagacy_gg_zbv_formatted",
@@ -686,7 +865,16 @@ def test_deepseek_v2_lite_bf16_lagacy_gg_zbv_formatted(self):
         self.assertIn("[Patch:megatron.pp.legacy_grouped_mlp_wgrad_split]", stdout)
 
 
+@pytest.mark.weekly
 class TestMegatronTrainerDeterministic(PrimusUT):
+    """Bitwise reproducibility, weekend-only.
+
+    Every case here runs the same config twice and compares the losses as
+    strings, so each one costs two full training launches. Determinism breaks
+    from the kernel/collective stack rather than from a typical PR, so the pair
+    is not worth ~4 launches on the per-PR path.
+    """
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
diff --git a/tests/trainer/test_megatron_trainer_sft_lora.yaml b/tests/trainer/test_megatron_trainer_sft_lora.yaml
new file mode 100644
index 000000000..3ad8dfc5a
--- /dev/null
+++ b/tests/trainer/test_megatron_trainer_sft_lora.yaml
@@ -0,0 +1,103 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:test-megatron-sft-lora}
+workspace: ./output
+
+modules:
+  post_trainer:
+    framework: megatron
+    config: sft_trainer.yaml
+    model: qwen2.5_7B.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_test_SFT_LoRA"
+      disable_wandb: true
+      stderr_sink_level: DEBUG
+
+      # required to select MegatronSFTTrainer
+      stage: sft
+
+      # The posttrain hook (01_convert_checkpoints.py) always HF->Megatron-
+      # converts `tokenizer_model` into the base checkpoint when
+      # pretrained_checkpoint/load are unset. Point it at a real-but-tiny,
+      # non-gated Qwen2 checkpoint (2 layers, hidden_size 8, ~2.4M params) so
+      # the converted checkpoint matches the tiny model config below and the
+      # test exercises both the base-checkpoint load and the LoRA wrap.
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+      num_layers: 2
+      hidden_size: 8
+      ffn_hidden_size: 32
+      num_attention_heads: 4
+      num_query_groups: 2
+
+      # 8 GPUs as pure data parallelism, so global_batch_size=8 is the smallest
+      # value dividing evenly (1 micro-batch/rank/iter). The 8-sample jsonl
+      # fixture covers one global batch; cyclic dataloader loops for the rest.
+      train_iters: 3
+      micro_batch_size: 1
+      global_batch_size: 8
+      seq_length: 128
+      max_position_embeddings: 128
+
+      eval_iters: 0
+      eval_interval: 100
+
+      lr: 2.0e-4
+      min_lr: 2.0e-5
+      lr_warmup_iters: 1
+      lr_decay_iters: null
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: false
+
+      # parallel: no TP/PP -- this suite's default launch uses all 8 node
+      # GPUs as pure data parallelism (see global_batch_size note above).
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: false
+
+      # Disable Turbo attention: the tiny head_dim (hidden_size/heads = 2) is far
+      # below what fused/flash kernels expect (test_qwen2_sft_lora pins the
+      # unfused backend). linear_qkv/proj/fc1/fc2 still use TE's non-turbo Linear
+      # classes, so the PEFT module_matcher path is exercised.
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      # SFT data: tiny local fixture (alpaca schema) so this test has no
+      # network dependency on the HuggingFace Hub.
+      sft_dataset_name: "tests/trainer/fixtures/sft_lora_smoke.jsonl"
+      sft_conversation_format: "alpaca"
+      enable_packed_sequences: false
+
+      # Leave pretrained_checkpoint/load unset so the posttrain hook fills
+      # pretrained_checkpoint from the converted tokenizer_model checkpoint,
+      # driving MegatronSFTTrainer's pre-wrap base-checkpoint-load branch.
+      finetune: true
+      pretrained_checkpoint: null
+      load: null
+      save: null
+      save_interval: 20000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch
+
+      # LoRA: the activation switch for primus/backends/megatron/peft/*.py.
+      lora:
+        enabled: true
+        dim: 8
+        alpha: 8
+        dropout: 0.0
+        dropout_position: pre
+        lora_A_init_method: xavier
+        lora_B_init_method: zero
+        target_modules:
+          - linear_qkv
+          - linear_proj
+          - linear_fc1
+          - linear_fc2
diff --git a/tests/trainer/test_torchtitan_trainer.py b/tests/trainer/test_torchtitan_trainer.py
index 587c6edd6..13b6bbfe7 100644
--- a/tests/trainer/test_torchtitan_trainer.py
+++ b/tests/trainer/test_torchtitan_trainer.py
@@ -6,6 +6,8 @@
 
 import os
 
+import pytest
+
 from tests.utils import PrimusUT, run_training_script
 
 
@@ -45,6 +47,26 @@ def run_script(
 
 
 class TestTorchTitanTrainer(PrimusUT):
+    """One config per distinct code path; isomorphic flavors are not tested here.
+
+    Every case is a full training launch, so a model whose code path a sibling
+    already covers earns no coverage. What is kept and why:
+
+    - llama3.1_8B BF16 + FP8: the dense llama path, both precisions.
+    - qwen3_0.6B: the qwen3 dense path (activation_checkpoint off).
+    - qwen3_32B: the dense config with activation_checkpoint.mode=full.
+    - deepseek_v3_16b: MLA + MoE on classic attention, no float8.
+    - deepseek_v3_16b_fp8: the only use_moe_fp8 path; also swaps classic
+      attention for the turbo one.
+    - deepseek_v3_671b: MLA + MoE under full recompute and
+      use_turbo_float8_linear.
+    - gpt_oss_20B BF16 + FP8: MoE + sink attention, selective checkpointing.
+
+    llama3.1_70B / llama3.1_405B (both precisions) and qwen3_1.7B used to be
+    here; they only scaled the dims of the above. Their recipes are still
+    schema-checked by tests/unit_tests/configs/test_example_configs.py.
+    """
+
     def test_llama3_1_8B_BF16(self):
         run_script(
             self.__class__.__name__,
@@ -75,41 +97,11 @@ def test_llama3_1_8B_FP8(self):
             ],
         )
 
-    def test_llama3_1_405B_bf16(self):
-        run_script(
-            self.__class__.__name__,
-            "llama3.1_405B_bf16",
-            "examples/torchtitan/configs/MI300X/llama3.1_405B-BF16-pretrain.yaml",
-            extra_args=[
-                "--model.n_layers",
-                "4",
-                "--training.steps",
-                "3",
-                "--training.mock_data",
-                "True",
-            ],
-        )
-
-    def test_llama3_1_405B_fp8(self):
-        run_script(
-            self.__class__.__name__,
-            "llama3.1_405B_fp8",
-            "examples/torchtitan/configs/MI300X/llama3.1_405B-FP8-pretrain.yaml",
-            extra_args=[
-                "--model.n_layers",
-                "4",
-                "--training.steps",
-                "3",
-                "--training.mock_data",
-                "True",
-            ],
-        )
-
-    def test_llama3_1_70B_bf16(self):
+    def test_qwen3_0_6B(self):
         run_script(
             self.__class__.__name__,
-            "llama3.1_70B_bf16",
-            "examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml",
+            "qwen3_0.6B",
+            "examples/torchtitan/configs/MI300X/qwen3_0.6B-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
@@ -120,11 +112,11 @@ def test_llama3_1_70B_bf16(self):
             ],
         )
 
-    def test_llama3_1_70B_fp8(self):
+    def test_qwen3_32B(self):
         run_script(
             self.__class__.__name__,
-            "llama3.1_70B_fp8",
-            "examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml",
+            "qwen3_32B",
+            "examples/torchtitan/configs/MI300X/qwen3_32B-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
@@ -135,14 +127,16 @@ def test_llama3_1_70B_fp8(self):
             ],
         )
 
-    def test_qwen3_0_6B(self):
+    def test_deepseek_v3_16b(self):
         run_script(
             self.__class__.__name__,
-            "qwen3_0.6B",
-            "examples/torchtitan/configs/MI300X/qwen3_0.6B-pretrain.yaml",
+            "deepseek_v3_16b",
+            "examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
+                "--model.n_dense_layers",
+                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
@@ -150,14 +144,16 @@ def test_qwen3_0_6B(self):
             ],
         )
 
-    def test_qwen3_1_7B(self):
+    def test_deepseek_v3_16b_fp8(self):
         run_script(
             self.__class__.__name__,
-            "qwen3_1.7B",
-            "examples/torchtitan/configs/MI300X/qwen3_1.7B-pretrain.yaml",
+            "deepseek_v3_16b_fp8",
+            "examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
+                "--model.n_dense_layers",
+                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
@@ -165,14 +161,16 @@ def test_qwen3_1_7B(self):
             ],
         )
 
-    def test_qwen3_32B(self):
+    def test_deepseek_v3_671b(self):
         run_script(
             self.__class__.__name__,
-            "qwen3_32B",
-            "examples/torchtitan/configs/MI300X/qwen3_32B-pretrain.yaml",
+            "deepseek_v3_671b",
+            "examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
+                "--model.n_dense_layers",
+                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
@@ -180,16 +178,15 @@ def test_qwen3_32B(self):
             ],
         )
 
-    def test_deepseek_v3_16b(self):
+    def test_gpt_oss_20B(self):
+        # Default Primus-Turbo path: GPT-OSS sink attention (flash_attn_func).
         run_script(
             self.__class__.__name__,
-            "deepseek_v3_16b",
-            "examples/torchtitan/configs/MI300X/deepseek_v3_16b-BF16-pretrain.yaml",
+            "gpt_oss_20B",
+            "examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
-                "--model.n_dense_layers",
-                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
@@ -197,16 +194,14 @@ def test_deepseek_v3_16b(self):
             ],
         )
 
-    def test_deepseek_v3_16b_fp8(self):
+    def test_gpt_oss_20B_fp8(self):
         run_script(
             self.__class__.__name__,
-            "deepseek_v3_16b_fp8",
-            "examples/torchtitan/configs/MI300X/deepseek_v3_16b-FP8-pretrain.yaml",
+            "gpt_oss_20B_fp8",
+            "examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
-                "--model.n_dense_layers",
-                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
@@ -214,16 +209,17 @@ def test_deepseek_v3_16b_fp8(self):
             ],
         )
 
-    def test_deepseek_v3_671b(self):
+    # Llama 4 had no E2E in either backend. Distinct from the llama3 flavor above:
+    # MoE with a shared expert, and the recipe's use_turbo_grouped_mm path.
+    @pytest.mark.weekly
+    def test_llama4_17Bx16E(self):
         run_script(
             self.__class__.__name__,
-            "deepseek_v3_671b",
-            "examples/torchtitan/configs/MI300X/deepseek_v3_671b-pretrain.yaml",
+            "llama4_17Bx16E",
+            "examples/torchtitan/configs/MI300X/llama4_17Bx16E-BF16-pretrain.yaml",
             extra_args=[
                 "--model.n_layers",
                 "4",
-                "--model.n_dense_layers",
-                "1",
                 "--training.steps",
                 "3",
                 "--training.mock_data",
diff --git a/primus/modules/trainer/megatron/__init__.py b/tests/unit_tests/backends/__init__.py
similarity index 100%
rename from primus/modules/trainer/megatron/__init__.py
rename to tests/unit_tests/backends/__init__.py
diff --git a/primus/modules/trainer/torchtitan/__init__.py b/tests/unit_tests/backends/diffusion/__init__.py
similarity index 100%
rename from primus/modules/trainer/torchtitan/__init__.py
rename to tests/unit_tests/backends/diffusion/__init__.py
diff --git a/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py
new file mode 100644
index 000000000..9a97445dd
--- /dev/null
+++ b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py
@@ -0,0 +1,14 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import pytest
+
+from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler
+
+
+def test_exponential_shift_requires_mu_or_dynamic_shift_len():
+    with pytest.raises(ValueError, match="exponential_shift=True"):
+        FlowMatchScheduler(exponential_shift=True)
diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py
new file mode 100644
index 000000000..83f520c3b
--- /dev/null
+++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py
@@ -0,0 +1,449 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from primus.backends.diffusion.argument_builder import DiffusionArgBuilder
+from primus.backends.diffusion.attention import (
+    get_attention_backend,
+    set_attention_backend,
+)
+from primus.backends.diffusion.data.flux_precomputed import (
+    FluxPrecomputedProcessor,
+    FluxRawImageTextDataset,
+    FluxRawImageTextProcessor,
+)
+from primus.backends.diffusion.models.flux.adapter import FluxForTraining
+from primus.backends.diffusion.models.flux.conditioner import HFEmbedder
+from primus.backends.diffusion.models.flux.math import apply_rope
+from primus.backends.diffusion.models.flux.math import attention as flux_attention
+from primus.backends.diffusion.models.flux.math import rope
+from primus.backends.diffusion.models.registrations.flux import build_flux_model
+from primus.backends.diffusion.trainers.fsdp2 import FSDP2Trainer
+
+
+def _finalize(params: dict):
+    builder = DiffusionArgBuilder()
+    builder.update(params)
+    return builder.finalize()
+
+
+def test_flux_argument_builder_selects_flux_defaults():
+    args = _finalize(
+        {
+            "model": {"name": "flux.1-dev", "config": {}},
+            "training": {"steps": 7, "local_batch_size": 3},
+            "data": {
+                "dataset_path": "/tmp/precomputed",
+                "dataset_type": "precomputed",
+                "empty_encodings_path": "/tmp/empty",
+                "prompt_dropout_prob": 0.25,
+            },
+            "lr_scheduler": {
+                "lr_scheduler_type": "constant_with_warmup",
+                "warmup_steps": 11,
+            },
+        }
+    )
+
+    assert args.model["name"] == "flux.1-dev"
+    assert args.dataset["name"] == "flux"
+    assert args.dataset["config"]["dataset_path"] == "/tmp/precomputed"
+    assert args.dataset["config"]["processor_config"]["empty_encodings_path"] == "/tmp/empty"
+    assert args.dataset["config"]["processor_config"]["prompt_dropout_prob"] == 0.25
+    assert args.trainer["args"]["max_steps"] == 7
+    assert args.trainer["args"]["per_device_train_batch_size"] == 3
+    assert args.trainer["args"]["lr_scheduler_type"] == "constant_with_warmup"
+    assert args.trainer["args"]["warmup_steps"] == 11
+    assert args.trainer["args"]["attention_backend"] == "flash_attn_aiter"
+    assert args.trainer["args"]["fsdp_transformer_layer_cls_to_wrap"] == "DoubleStreamBlock,SingleStreamBlock"
+    assert args.trainer["args"]["compile_transformer_blocks"] is True
+
+
+def test_flux_argument_builder_maps_raw_dataset_type():
+    args = _finalize(
+        {
+            "model": {"name": "flux.1-dev", "config": {}},
+            "data": {
+                "dataset_type": "raw",
+                "dataset_format": "webdataset",
+                "dataset_path": "/tmp/cc12m_test",
+                "prompt_dropout_prob": 0.1,
+                "img_size": 128,
+            },
+        }
+    )
+
+    assert args.dataset["config"]["dataset_type"] == "raw"
+    assert args.dataset["config"]["dataset_format"] == "webdataset"
+    assert args.dataset["config"]["processor_config"]["img_size"] == 128
+
+
+def test_flux_raw_dataset_name_defaults():
+    path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset")
+    assert path == "zirui3/cc12m-test"
+    assert fmt == "hf_repo"
+
+    path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", "/tmp/cc12m_test", "webdataset")
+    assert path == "/tmp/cc12m_test"
+    assert fmt == "webdataset"
+
+    path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-wds", None, "webdataset")
+    assert path == "pixparse/cc12m-wds"
+    assert fmt == "hf_repo"
+
+
+def test_flux_argument_builder_rejects_sequence_parallelism():
+    with pytest.raises(ValueError, match="sp_size"):
+        _finalize(
+            {
+                "model": {"name": "flux.1-dev", "config": {}},
+                "parallelism": {"sp_size": 2},
+            }
+        )
+
+
+def test_flux_attention_dispatch_matches_sdpa_layout():
+    previous_backend = get_attention_backend()
+    set_attention_backend("sdpa")
+    try:
+        torch.manual_seed(7)
+        q = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16)
+        k = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16)
+        v = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16)
+        pos = torch.arange(5, dtype=torch.float32).repeat(2, 1)
+        pe = rope(pos, dim=4, theta=10000).unsqueeze(1)
+
+        actual = flux_attention(q, k, v, pe=pe)
+        q_rope, k_rope = apply_rope(q, k, pe)
+        expected = torch.nn.functional.scaled_dot_product_attention(q_rope, k_rope, v)
+        expected = expected.transpose(1, 2).flatten(2)
+
+        torch.testing.assert_close(actual, expected)
+    finally:
+        set_attention_backend(previous_backend)
+
+
+def test_flux_hf_embedder_passes_attention_mask():
+    class FakeTokenizer:
+        def __call__(self, text, **kwargs):
+            assert kwargs["padding"] == "max_length"
+            return {
+                "input_ids": torch.tensor([[1, 2, 0], [3, 0, 0]], dtype=torch.long),
+                "attention_mask": torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long),
+            }
+
+    class FakeTextModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.param = torch.nn.Parameter(torch.zeros(()))
+            self.seen_attention_mask = None
+
+        def forward(self, *, input_ids, attention_mask, output_hidden_states):
+            self.seen_attention_mask = attention_mask
+            return {"last_hidden_state": input_ids.float().unsqueeze(-1)}
+
+    embedder = HFEmbedder.__new__(HFEmbedder)
+    torch.nn.Module.__init__(embedder)
+    embedder.tokenizer = FakeTokenizer()
+    embedder.hf_module = FakeTextModel()
+    embedder.output_key = "last_hidden_state"
+    embedder.max_length = 3
+
+    output = embedder(["short", "x"])
+
+    torch.testing.assert_close(
+        embedder.hf_module.seen_attention_mask,
+        torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long),
+    )
+    assert output.shape == (2, 3, 1)
+
+
+def test_flux_forward_uses_positional_scheduler():
+    model = FluxForTraining(
+        dit=torch.nn.Identity(),
+        train_pipeline=object(),
+        model_config=object(),
+    )
+    scheduler = object()
+    captured = {}
+
+    def forward_train(batch, scheduler=None):
+        captured["batch"] = batch
+        captured["scheduler"] = scheduler
+        return {"loss": torch.tensor(0.0)}
+
+    model.forward_train = forward_train
+    batch = {"x": torch.tensor(1)}
+
+    output = model(batch, scheduler)
+
+    assert output["loss"].item() == 0.0
+    assert captured == {"batch": batch, "scheduler": scheduler}
+
+
+def test_fsdp2_compile_transformer_blocks_replaces_modules(monkeypatch):
+    class CompiledBlock(torch.nn.Module):
+        def __init__(self, original):
+            super().__init__()
+            self.original = original
+
+    root = torch.nn.Module()
+    root.double_blocks = torch.nn.ModuleList([torch.nn.Identity(), torch.nn.ReLU()])
+    root.single_blocks = torch.nn.ModuleList([torch.nn.Sigmoid()])
+    compiled_inputs = []
+
+    def fake_compile(module, *, fullgraph):
+        assert fullgraph is True
+        compiled_inputs.append(module)
+        return CompiledBlock(module)
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+    trainer = FSDP2Trainer.__new__(FSDP2Trainer)
+    trainer.rank = 1
+
+    trainer._compile_transformer_blocks(root)
+
+    assert len(compiled_inputs) == 3
+    assert all(isinstance(block, CompiledBlock) for block in root.double_blocks)
+    assert all(isinstance(block, CompiledBlock) for block in root.single_blocks)
+
+
+def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path):
+    empty_dir = tmp_path / "empty"
+    empty_dir.mkdir()
+    np.save(empty_dir / "t5_empty.npy", np.zeros((1, 3, 8), dtype=np.float32))
+    np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32))
+
+    processor = FluxPrecomputedProcessor(
+        {
+            "prompt_dropout_prob": 1.0,
+            "empty_encodings_path": str(empty_dir),
+        }
+    )
+    batch = [
+        {
+            "t5_encodings": torch.ones(3, 8),
+            "clip_encodings": torch.ones(4),
+            "mean": torch.zeros(1, 2, 2),
+            "logvar": torch.zeros(1, 2, 2),
+        },
+        {
+            "t5_encodings": torch.ones(3, 8),
+            "clip_encodings": torch.ones(4),
+            "mean": torch.zeros(1, 2, 2),
+            "logvar": torch.zeros(1, 2, 2),
+        },
+    ]
+
+    out = processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32)
+
+    assert out["t5_encodings"].shape == (2, 3, 8)
+    assert out["clip_encodings"].shape == (2, 4)
+    assert torch.count_nonzero(out["t5_encodings"]) == 0
+    assert torch.count_nonzero(out["clip_encodings"]) == 0
+
+
+def test_flux_precomputed_processor_rejects_mismatched_empty_encoding(tmp_path):
+    empty_dir = tmp_path / "empty"
+    empty_dir.mkdir()
+    # Empty T5 encoding has sequence length 5, but the batch samples use length 3.
+    np.save(empty_dir / "t5_empty.npy", np.zeros((1, 5, 8), dtype=np.float32))
+    np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32))
+
+    processor = FluxPrecomputedProcessor(
+        {
+            "prompt_dropout_prob": 1.0,
+            "empty_encodings_path": str(empty_dir),
+        }
+    )
+    batch = [
+        {
+            "t5_encodings": torch.ones(3, 8),
+            "clip_encodings": torch.ones(4),
+            "mean": torch.zeros(1, 2, 2),
+            "logvar": torch.zeros(1, 2, 2),
+        }
+    ]
+
+    with pytest.raises(ValueError, match="empty T5 encoding shape"):
+        processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32)
+
+
+def test_flux_raw_processor_prepares_images_and_prompts():
+    from PIL import Image
+
+    processor = FluxRawImageTextProcessor(
+        {"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False}
+    )
+    image = Image.fromarray(np.full((6, 10, 3), 127, dtype=np.uint8))
+
+    out = processor.prepare_batch(
+        batch=[{"image": image, "prompt": "a test image"}],
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    assert out["image"].shape == (1, 3, 8, 8)
+    assert out["prompts"] == [""]
+
+
+def test_tiny_flux_model_computes_precomputed_loss():
+    model = build_flux_model(
+        {
+            "config": {
+                "model_preset": "flux.1-dev",
+                "guidance": 1.0,
+                "params": {
+                    "in_channels": 4,
+                    "out_channels": 4,
+                    "vec_in_dim": 4,
+                    "context_in_dim": 8,
+                    "hidden_size": 12,
+                    "num_heads": 2,
+                    "depth": 1,
+                    "depth_single_blocks": 1,
+                    "axes_dim": [2, 2, 2],
+                },
+            }
+        }
+    )
+    batch = {
+        "t5_encodings": torch.randn(2, 3, 8),
+        "clip_encodings": torch.randn(2, 4),
+        "mean": torch.randn(2, 1, 2, 2),
+        "logvar": torch.zeros(2, 1, 2, 2),
+    }
+
+    outputs = model.forward_train(batch)
+
+    assert outputs["loss"].ndim == 0
+    assert torch.isfinite(outputs["loss"])
+
+
+def test_tiny_flux_schnell_model_computes_without_guidance():
+    model = build_flux_model(
+        {
+            "config": {
+                "model_preset": "flux.1-schnell",
+                "params": {
+                    "in_channels": 4,
+                    "out_channels": 4,
+                    "vec_in_dim": 4,
+                    "context_in_dim": 8,
+                    "hidden_size": 12,
+                    "num_heads": 2,
+                    "depth": 1,
+                    "depth_single_blocks": 1,
+                    "axes_dim": [2, 2, 2],
+                },
+            }
+        }
+    )
+    batch = {
+        "t5_encodings": torch.randn(2, 3, 8),
+        "clip_encodings": torch.randn(2, 4),
+        "mean": torch.randn(2, 1, 2, 2),
+        "logvar": torch.zeros(2, 1, 2, 2),
+    }
+
+    outputs = model.forward_train(batch)
+
+    assert model.dit.params.guidance_embed is False
+    assert model.model_config.guidance is None
+    assert outputs["loss"].ndim == 0
+    assert torch.isfinite(outputs["loss"])
+
+
+def test_flux_position_ids_are_float32_regardless_of_model_dtype():
+    model = build_flux_model(
+        {
+            "config": {
+                "model_preset": "flux.1-dev",
+                "guidance": 1.0,
+                "params": {
+                    "in_channels": 4,
+                    "out_channels": 4,
+                    "vec_in_dim": 4,
+                    "context_in_dim": 8,
+                    "hidden_size": 12,
+                    "num_heads": 2,
+                    "depth": 1,
+                    "depth_single_blocks": 1,
+                    "axes_dim": [2, 2, 2],
+                },
+            }
+        }
+    )
+    # Force a low-precision compute dtype; position ids must stay float32 so RoPE
+    # grid indices are not corrupted (bf16 only represents integers up to 256).
+    model.dit = model.dit.to(dtype=torch.bfloat16)
+
+    captured = {}
+    original_forward = model.dit.forward
+
+    def capturing_forward(*args, **kwargs):
+        captured["img_ids_dtype"] = kwargs["img_ids"].dtype
+        captured["txt_ids_dtype"] = kwargs["txt_ids"].dtype
+        return original_forward(*args, **kwargs)
+
+    model.dit.forward = capturing_forward
+    model.forward_train(
+        {
+            "t5_encodings": torch.randn(2, 3, 8, dtype=torch.bfloat16),
+            "clip_encodings": torch.randn(2, 4, dtype=torch.bfloat16),
+            "mean": torch.randn(2, 1, 2, 2, dtype=torch.bfloat16),
+            "logvar": torch.zeros(2, 1, 2, 2, dtype=torch.bfloat16),
+        }
+    )
+
+    assert captured["img_ids_dtype"] == torch.float32
+    assert captured["txt_ids_dtype"] == torch.float32
+
+
+def test_tiny_flux_model_computes_raw_loss_with_dummy_encoders():
+    class DummyAutoencoder(torch.nn.Module):
+        def encode(self, image):
+            return image[:, :1, :2, :2]
+
+    class DummyT5(torch.nn.Module):
+        def forward(self, prompts):
+            return torch.zeros(len(prompts), 3, 8)
+
+    class DummyClip(torch.nn.Module):
+        def forward(self, prompts):
+            return torch.zeros(len(prompts), 4)
+
+    model = build_flux_model(
+        {
+            "config": {
+                "model_preset": "flux.1-dev",
+                "guidance": 1.0,
+                "params": {
+                    "in_channels": 4,
+                    "out_channels": 4,
+                    "vec_in_dim": 4,
+                    "context_in_dim": 8,
+                    "hidden_size": 12,
+                    "num_heads": 2,
+                    "depth": 1,
+                    "depth_single_blocks": 1,
+                    "axes_dim": [2, 2, 2],
+                },
+            }
+        }
+    )
+    model.autoencoder = DummyAutoencoder()
+    model.t5_encoder = DummyT5()
+    model.clip_encoder = DummyClip()
+
+    outputs = model.forward_train(
+        {
+            "image": torch.randn(2, 3, 4, 4),
+            "prompts": ["cat", "dog"],
+        }
+    )
+
+    assert outputs["loss"].ndim == 0
+    assert torch.isfinite(outputs["loss"])
diff --git a/tests/unit_tests/backends/diffusion/test_review_regressions.py b/tests/unit_tests/backends/diffusion/test_review_regressions.py
new file mode 100644
index 000000000..041a9b24c
--- /dev/null
+++ b/tests/unit_tests/backends/diffusion/test_review_regressions.py
@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from primus.backends.diffusion.data.collator import VisionCollator
+from primus.backends.diffusion.data.dataset import WanVideoDataset
+from primus.backends.diffusion.diffusion_pretrain_trainer import (
+    DiffusionPretrainTrainer,
+)
+
+
+def test_vision_collator_uses_provided_attention_mask_without_input_ids():
+    class Tokenizer:
+        padding_side = "right"
+        pad_token_id = 0
+
+    processor = SimpleNamespace(tokenizer=Tokenizer())
+    collator = VisionCollator(processor=processor)
+
+    batch = collator(
+        [
+            {"attention_mask": torch.tensor([1, 1]), "pixel_values": torch.ones(2)},
+            {"attention_mask": torch.tensor([1]), "pixel_values": torch.zeros(2)},
+        ]
+    )
+
+    torch.testing.assert_close(batch["attention_mask"], torch.tensor([[1, 1], [1, 0]]))
+    torch.testing.assert_close(batch["pixel_values"], torch.tensor([[1.0, 1.0], [0.0, 0.0]]))
+
+
+def test_wan_video_dataset_requires_config_with_dataset_path():
+    with pytest.raises(ValueError, match="dataset config with dataset_path"):
+        WanVideoDataset(processor=object())
+
+    with pytest.raises(ValueError, match="config.dataset_path"):
+        WanVideoDataset(processor=object(), config=SimpleNamespace(dataset_path=None))
+
+
+def test_diffusion_setup_reports_missing_flux_dependencies(monkeypatch):
+    missing = {"datasets", "huggingface_hub", "sentencepiece", "webdataset"}
+
+    def fake_find_spec(package):
+        return None if package in missing else object()
+
+    monkeypatch.setattr("importlib.util.find_spec", fake_find_spec)
+    trainer = DiffusionPretrainTrainer.__new__(DiffusionPretrainTrainer)
+    trainer.backend_args = SimpleNamespace(
+        trainer={"args": {}},
+        dataset={
+            "name": "flux",
+            "config": {
+                "dataset_type": "raw",
+                "dataset_format": "webdataset",
+            },
+        },
+    )
+
+    with pytest.raises(RuntimeError) as exc_info:
+        trainer.setup()
+
+    message = str(exc_info.value)
+    for package in sorted(missing):
+        assert package in message
diff --git a/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py
new file mode 100644
index 000000000..9d07d4f07
--- /dev/null
+++ b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py
@@ -0,0 +1,194 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from types import SimpleNamespace
+
+import pytest
+
+from primus.backends.diffusion.argument_builder import WanArgBuilder
+
+
+def _minimal_params():
+    return {
+        "model": {
+            "name": "wan",
+            "config": {
+                "model_type": "t2v",
+            },
+        },
+    }
+
+
+def test_rejects_legacy_public_dataset_override():
+    params = _minimal_params()
+    params["dataset"] = {"name": "wan"}
+
+    builder = WanArgBuilder()
+    builder.update(params)
+
+    with pytest.raises(ValueError, match="no longer accepts public `dataset` overrides"):
+        builder.finalize()
+
+
+def test_rejects_legacy_public_trainer_override():
+    params = _minimal_params()
+    params["trainer"] = {"name": "fsdp2"}
+
+    builder = WanArgBuilder()
+    builder.update(params)
+
+    with pytest.raises(ValueError, match="no longer accepts public `trainer` overrides"):
+        builder.finalize()
+
+
+def test_maps_primus_style_sections_to_wan_runtime_config():
+    params = {
+        **_minimal_params(),
+        "stage": "posttrain",
+        "primus": {"experiment": "wan-smoke"},
+        "training": {
+            "steps": 7,
+            "local_batch_size": 2,
+            "global_batch_size": 16,
+            "gradient_accumulation_steps": 4,
+            "output_dir": "/tmp/wan-out",
+            "save_steps": 3,
+            "run_name": "wan-test",
+            "num_train_epochs": 9,
+            "dataloader_num_workers": 0,
+            "resume_from_checkpoint": "/tmp/resume",
+        },
+        "data": {
+            "dataset_path": "/data/meta.jsonl",
+            "data_folder": "/data/videos",
+            "frame_num": 17,
+            "video_backend": "decord",
+            "text_tokenizer": "/models/umt5",
+            "height": 256,
+            "width": 384,
+        },
+        "parallelism": {
+            "sp_size": 4,
+            "dp_replicate": 2,
+        },
+        "optimizer": {
+            "lr": 2.0e-5,
+            "weight_decay": 0.02,
+            "adam_beta1": 0.8,
+            "adam_beta2": 0.95,
+            "adam_epsilon": 1.0e-7,
+            "max_grad_norm": 0.5,
+        },
+        "runtime": {
+            "attention_backend": "sdpa",
+            "report_to": "none",
+            "seed": 1234,
+            "fsdp2_reshard_after_forward": False,
+        },
+        "metrics": {
+            "log_freq": 5,
+            "enable_wandb": False,
+        },
+    }
+
+    builder = WanArgBuilder()
+    builder.update(params)
+    result = builder.finalize()
+
+    dataset_cfg = result.dataset["config"]
+    processor_cfg = dataset_cfg["processor_config"]
+    trainer_args = result.trainer["args"]
+
+    assert result.stage == "posttrain"
+    assert result.primus == {"experiment": "wan-smoke"}
+    assert result.model == params["model"]
+
+    assert dataset_cfg["dataset_path"] == "/data/meta.jsonl"
+    assert dataset_cfg["data_folder"] == "/data/videos"
+    assert dataset_cfg["frame_num"] == 17
+    assert dataset_cfg["video_backend"] == "decord"
+    assert processor_cfg["text_tokenizer"] == "/models/umt5"
+    assert processor_cfg["extra_kwargs"]["size"] == {"height": 256, "width": 384}
+
+    assert trainer_args["max_steps"] == 7
+    assert trainer_args["per_device_train_batch_size"] == 2
+    assert trainer_args["global_batch_size"] == 16
+    assert trainer_args["gradient_accumulation_steps"] == 4
+    assert trainer_args["output_dir"] == "/tmp/wan-out"
+    assert trainer_args["save_steps"] == 3
+    assert trainer_args["run_name"] == "wan-test"
+    assert trainer_args["num_train_epochs"] == 9
+    assert trainer_args["dataloader_num_workers"] == 0
+    assert trainer_args["resume_from_checkpoint"] == "/tmp/resume"
+
+    assert trainer_args["sp_size"] == 4
+    assert trainer_args["dp_replicate"] == 2
+    assert trainer_args["learning_rate"] == 2.0e-5
+    assert trainer_args["weight_decay"] == 0.02
+    assert trainer_args["adam_beta1"] == 0.8
+    assert trainer_args["adam_beta2"] == 0.95
+    assert trainer_args["adam_epsilon"] == 1.0e-7
+    assert trainer_args["max_grad_norm"] == 0.5
+    assert trainer_args["attention_backend"] == "sdpa"
+    assert trainer_args["report_to"] == "none"
+    assert trainer_args["seed"] == 1234
+    assert trainer_args["fsdp2_reshard_after_forward"] is False
+    assert trainer_args["logging_steps"] == 5
+
+
+def test_defaults_propagate_when_optional_sections_are_omitted():
+    builder = WanArgBuilder()
+    builder.update(_minimal_params())
+    result = builder.finalize()
+
+    dataset_cfg = result.dataset["config"]
+    trainer_args = result.trainer["args"]
+
+    assert result.stage == "pretrain"
+    assert dataset_cfg["dataset_path"] == "/path/to/meta.jsonl"
+    assert dataset_cfg["data_folder"] == "/path/to/videos"
+    assert dataset_cfg["processor_config"]["text_tokenizer"] == "/path/to/umt5-xxl"
+    assert result.trainer["name"] == "fsdp2"
+    assert trainer_args["max_steps"] == 100
+    assert trainer_args["attention_backend"] == "flash_attn_aiter"
+    assert trainer_args["sp_size"] == 1
+    assert trainer_args["dp_replicate"] == 1
+    assert trainer_args["report_to"] == "none"
+
+
+def test_metrics_enable_wandb_sets_report_to_when_runtime_omits_it():
+    builder = WanArgBuilder()
+    builder.update(
+        {
+            **_minimal_params(),
+            "metrics": {
+                "enable_wandb": True,
+            },
+        }
+    )
+
+    result = builder.finalize()
+
+    assert result.trainer["args"]["report_to"] == "wandb"
+
+
+def test_update_accepts_simplenamespace_input():
+    builder = WanArgBuilder()
+    builder.update(
+        SimpleNamespace(
+            model=SimpleNamespace(
+                name="wan",
+                config=SimpleNamespace(model_type="t2v"),
+            ),
+            data=SimpleNamespace(dataset_path="/data/meta.jsonl"),
+        )
+    )
+
+    result = builder.finalize()
+
+    assert result.model["name"] == "wan"
+    assert result.model["config"]["model_type"] == "t2v"
+    assert result.dataset["config"]["dataset_path"] == "/data/meta.jsonl"
diff --git a/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py
new file mode 100644
index 000000000..6636942e9
--- /dev/null
+++ b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py
@@ -0,0 +1,61 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import torch
+
+from primus.backends.diffusion.trainers.base import BaseWanTrainer
+
+
+def _make_trainer():
+    trainer = BaseWanTrainer.__new__(BaseWanTrainer)
+    trainer.model = torch.nn.Linear(2, 2)
+    trainer.rank = 0
+    trainer.args = {
+        "learning_rate": 1.0e-4,
+        "weight_decay": 0.01,
+        "adam_beta1": 0.9,
+        "adam_beta2": 0.999,
+        "adam_epsilon": 1.0e-8,
+    }
+    return trainer
+
+
+def test_optimizer_falls_back_to_foreach_when_fused_raises_runtime_error(monkeypatch):
+    calls = []
+
+    def fake_adamw(**kwargs):
+        calls.append(kwargs)
+        if kwargs.get("fused"):
+            raise RuntimeError("fused AdamW is unsupported")
+        return kwargs
+
+    monkeypatch.setattr(torch.optim, "AdamW", fake_adamw)
+
+    optimizer = _make_trainer()._create_optimizer()
+
+    assert calls[0]["fused"] is True
+    assert calls[1]["foreach"] is True
+    assert optimizer is calls[1]
+
+
+def test_optimizer_falls_back_to_default_when_foreach_raises_runtime_error(monkeypatch):
+    calls = []
+
+    def fake_adamw(**kwargs):
+        calls.append(kwargs)
+        if kwargs.get("fused") or kwargs.get("foreach"):
+            raise RuntimeError("optimized AdamW path is unsupported")
+        return kwargs
+
+    monkeypatch.setattr(torch.optim, "AdamW", fake_adamw)
+
+    optimizer = _make_trainer()._create_optimizer()
+
+    assert calls[0]["fused"] is True
+    assert calls[1]["foreach"] is True
+    assert "fused" not in calls[2]
+    assert "foreach" not in calls[2]
+    assert optimizer is calls[2]
diff --git a/tests/unit_tests/backends/megatron/__init__.py b/tests/unit_tests/backends/megatron/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/backends/megatron/conftest.py b/tests/unit_tests/backends/megatron/conftest.py
index 4c6599b92..8c807ac95 100644
--- a/tests/unit_tests/backends/megatron/conftest.py
+++ b/tests/unit_tests/backends/megatron/conftest.py
@@ -1,34 +1,63 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
 """
-Pytest fixtures for Megatron backend tests.
+Shared pytest fixtures for Megatron backend tests.
 
-This conftest provides fixtures for tests that need Megatron parallel state
-initialization, such as tests for Primus Turbo layers and other Megatron-Core
-components that require distributed setup.
+Provides reusable fixtures for parallel state initialization used across
+optimizer, diffusion, and other Megatron test suites.
 """
 
 import os
-import socket
+from types import SimpleNamespace
 
 import pytest
 import torch
-import torch.distributed as dist
+
+from primus.core.utils import logger
 
 
-def _find_free_port() -> int:
-    """Ask the kernel for a free TCP port on localhost.
+@pytest.fixture(autouse=True, scope="session")
+def setup_logger():
+    """Initialize the Primus logger for megatron tests that use log_rank_0.
 
-    Binding to port 0 lets the OS pick a port that is guaranteed free at this
-    instant, which avoids the EADDRINUSE failures seen when guessing a random
-    port out of a fixed range that overlaps the ephemeral port range.
+    The Primus ``_logger`` global is ``None`` until configured, so tests that
+    log (e.g. test_warmup_convergence.py) raise ``AttributeError`` when run
+    without a prior logger-init. This session-scoped autouse fixture removes
+    that ordering dependency. Subdirectories may override it with an identically
+    named fixture (see diffusion/training/conftest.py).
     """
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
+    logger_cfg = logger.LoggerConfig(
+        exp_root_path=os.environ.get("UT_LOG_PATH", "ut_out"),
+        work_group="develop",
+        user_name="root",
+        exp_name="unittest",
+        module_name="UT-training",
+        file_sink_level="DEBUG",
+        stderr_sink_level="INFO",
+        node_ip="localhost",
+        rank=os.environ.get("RANK", 0),
+        world_size=os.environ.get("WORLD_SIZE", 1),
+    )
+    logger.setup_logger(logger_cfg, is_head=False)
+
+
+def _is_mxfp4_supported():
+    if not torch.cuda.is_available():
+        return False
+    try:
+        from primus_turbo.pytorch.core.low_precision import check_mxfp4_support
+
+        supported, _ = check_mxfp4_support()
+        return supported
+    except ImportError:
+        return False
+
+
+requires_mxfp4 = pytest.mark.skipif(
+    not _is_mxfp4_supported(),
+    reason="Requires gfx950+ (MI355X) for MXFP4 support",
+)
 
 
 @pytest.fixture(scope="function")
@@ -48,31 +77,37 @@ def init_parallel_state():
 
     Example:
         @pytest.fixture(autouse=True)
-        def setup_parallel(self, init_parallel_state, monkeypatch):
+        def setup_parallel(self, init_parallel_state):
             '''Auto-use parallel state for this test class.'''
             pass
     """
-    # Only import after sys.path is set up (which happens in pytest_configure)
     from megatron.core import parallel_state as ps
 
     # Initialize torch.distributed if not already initialized
-    if not dist.is_initialized():
-        # Use a kernel-assigned free port to avoid conflicts with other running
-        # tests. Set MASTER_PORT explicitly (not setdefault) so a stale value
-        # left in the environment cannot diverge from the port we actually use.
-        port = _find_free_port()
-        os.environ["MASTER_ADDR"] = "127.0.0.1"
-        os.environ["MASTER_PORT"] = str(port)
-
-        dist.init_process_group(
-            backend="nccl" if torch.cuda.is_available() else "gloo",
-            init_method=f"tcp://127.0.0.1:{port}",
-            world_size=1,
-            rank=0,
-        )
+    if not torch.distributed.is_initialized():
+        import os
+        import socket
+
+        # Use OS-assigned ephemeral port to avoid TIME_WAIT conflicts
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            s.bind(("127.0.0.1", 0))
+            port = s.getsockname()[1]
+
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", str(port))
 
-    # Check if model parallel already initialized and destroy if so
-    # This ensures clean state for each test
+        try:
+            torch.distributed.init_process_group(
+                backend="nccl" if torch.cuda.is_available() else "gloo",
+                init_method=f"tcp://127.0.0.1:{port}",
+                world_size=1,
+                rank=0,
+            )
+        except Exception as e:
+            pytest.skip(f"Could not initialize distributed: {e}")
+
+    # Check if model parallel already initialized
     if ps.model_parallel_is_initialized():
         ps.destroy_model_parallel()
 
@@ -88,22 +123,29 @@ def setup_parallel(self, init_parallel_state, monkeypatch):
     # This is required for:
     # 1. ColumnParallelLinear and other tensor parallel layers that use get_cuda_rng_tracker().fork()
     # 2. CUDA graph support in layers (enable_cuda_graph=True)
+    from megatron.core.tensor_parallel import random as tp_random
+
     if torch.cuda.is_available():
+        # Initialize RNG tracker with CUDA graph support BEFORE calling model_parallel_cuda_manual_seed
+        # Try Transformer Engine RNG tracker first (best for CUDA graphs, used by Megatron's own tests)
         try:
-            from megatron.core.tensor_parallel import random as tp_random
-
-            # Initialize RNG tracker with CUDA graph support BEFORE calling model_parallel_cuda_manual_seed
-            # Try Transformer Engine RNG tracker first (best for CUDA graphs, used by Megatron's own tests)
-            try:
-                tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True)
-            except (ImportError, AssertionError):
-                # Fallback to native PyTorch CUDA graph RNG support if TE not available
-                tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True)
-
-            tp_random.model_parallel_cuda_manual_seed(42)
-        except ImportError:
-            # RNG tracker initialization is optional - skip if not available
-            pass
+            tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True)
+        except (ImportError, AssertionError):
+            # Fallback to native PyTorch CUDA graph RNG support if TE not available
+            tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True)
+
+        tp_random.model_parallel_cuda_manual_seed(42)
+
+    # Initialize Megatron global args with minimal defaults required by
+    # PrimusTurboLocalAttention and PrimusTorchFullyShardedDataParallel
+    from megatron.training.global_vars import set_args
+
+    set_args(
+        SimpleNamespace(
+            enable_turbo_attention_float8=False,
+            data_parallel_replicate_degree=1,
+        )
+    )
 
     yield
 
@@ -111,7 +153,12 @@ def setup_parallel(self, init_parallel_state, monkeypatch):
     if ps.model_parallel_is_initialized():
         ps.destroy_model_parallel()
 
+    # Reset Megatron global args
+    import megatron.training.global_vars as gvars
+
+    gvars._GLOBAL_ARGS = None
+
     # Cleanup torch.distributed for single-process tests
     # (In multi-process torchrun tests, the process group persists across tests)
-    if dist.is_initialized():
-        dist.destroy_process_group()
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
diff --git a/tests/unit_tests/backends/megatron/diffusion/__init__.py b/tests/unit_tests/backends/megatron/diffusion/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/unit_tests/backends/megatron/diffusion/conftest.py b/tests/unit_tests/backends/megatron/diffusion/conftest.py
new file mode 100644
index 000000000..8fe29f243
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/conftest.py
@@ -0,0 +1,47 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Pytest fixtures for Flux diffusion model testing.
+
+Re-exports shared fixtures from the parent megatron conftest.
+"""
+
+import pytest
+
+from tests.utils import install_aiter_deepbind_hook
+
+# Install Primus' production aiter RTLD_DEEPBIND import hook before any diffusion
+# test imports the aiter mha kernels. On gfx942/gfx950, transformer_engine's
+# stale vendored libmha interposes the global ``aiter::mha_bwd`` over Turbo's
+# pinned aiter, crashing attention backward (ROCm/aiter#1332). Real training
+# installs this via the ``megatron.turbo.aiter_deepbind`` before_train patch;
+# unit tests call flash_attn_func directly and never hit that phase, so we wire
+# up the same hook here. No-op without a GPU.
+install_aiter_deepbind_hook()
+
+from tests.unit_tests.backends.megatron.conftest import (  # noqa: F401,E402
+    init_parallel_state,
+)
+
+
+@pytest.fixture(autouse=True)
+def _unset_nvte_attention_env(monkeypatch):
+    """Clear the TE attention-backend env vars for diffusion tests.
+
+    Some container images bake ``NVTE_FLASH_ATTN=0`` (they target the fused/CK
+    attention path). Flux's ``DiffusionModule._set_attention_backend()`` defaults
+    to the ``auto`` backend, which validates that ``NVTE_FLASH_ATTN``/
+    ``NVTE_FUSED_ATTN``/``NVTE_UNFUSED_ATTN`` are unset-or-1, so the baked ``0``
+    makes every Flux model construction fail. Mirror Megatron's own harness
+    (``Utils.initialize_distributed`` in ``tests/unit_tests/test_utilities.py``),
+    which pops these three vars so a baked/leaked value cannot poison the auto
+    backend check. Use ``monkeypatch.delenv`` (not ``os.environ.pop``):
+    ``_set_attention_backend`` writes ``os.environ[var] = 1`` after its check, so
+    monkeypatch's teardown is what restores each var to its pre-test state and
+    contains that write-back leak. Scoped to the diffusion suite only --
+    non-diffusion Primus megatron tests never construct a model that runs this
+    assertion.
+    """
+    for var in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"):
+        monkeypatch.delenv(var, raising=False)
diff --git a/tests/unit_tests/backends/megatron/diffusion/constants.py b/tests/unit_tests/backends/megatron/diffusion/constants.py
new file mode 100644
index 000000000..2689f67eb
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/constants.py
@@ -0,0 +1,73 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Constants for Flux diffusion model testing.
+
+This module centralizes common dimensional and architectural constants
+used across diffusion tests to reduce magic numbers and improve maintainability.
+"""
+
+# Model Architecture Dimensions
+HIDDEN_DIM_FLUX = 3072  # Flux model hidden dimension
+NUM_ATTENTION_HEADS_FLUX = 24  # Number of attention heads in Flux
+HEAD_DIM_FLUX = 128  # Hidden dimension per attention head (3072 / 24)
+
+# Encoder Dimensions
+T5_XXL_EMBEDDING_DIM = 4096  # T5-XXL text encoder output dimension
+CLIP_L_EMBEDDING_DIM = 768  # CLIP-L pooled embedding dimension
+VAE_LATENT_CHANNELS = 16  # VAE encoder output channels for Flux
+
+# Position Encoding
+ROPE_THETA_DEFAULT = 10000  # Default theta value for RoPE (Rotary Position Embedding)
+FLUX_AXES_DIM = (16, 56, 56)  # Default axes dimensions for Flux 3D RoPE
+
+# Embedding Dimensions
+TIMESTEP_EMBEDDING_DIM = 256  # Standard timestep embedding dimension
+
+# Batch Sizes (count-based naming)
+BATCH_SIZE_SINGLE = 1  # Single sample tests
+BATCH_SIZE_PAIR = 2  # Paired sample tests
+BATCH_SIZE_QUAD = 4  # Quad sample tests (most common)
+BATCH_SIZE_OCTO = 8  # Octuple sample tests
+BATCH_SIZE_HEX = 16  # Hexadecuple sample tests
+
+# Image Dimensions (size-based naming)
+IMG_SIZE_MICRO = 4  # Micro size for unit tests
+IMG_SIZE_MINI = 8  # Mini size for quick tests
+IMG_SIZE_TINY = 16  # Tiny size for fast tests
+IMG_SIZE_SMALL = 32  # Small size for standard tests
+IMG_SIZE_MEDIUM = 64  # Medium size
+IMG_SIZE_LARGE = 128  # Large size
+
+# Text Sequence Lengths
+TEXT_SEQ_LEN_SHORT = 77  # Standard CLIP text length
+TEXT_SEQ_LEN_MEDIUM = 128  # Medium sequence length
+TEXT_SEQ_LEN_LONG = 256  # Long sequence length
+TEXT_SEQ_LEN_XLARGE = 512  # Extra long sequence length
+
+# Additional Sequence Lengths
+SEQ_LEN_TINY = 100  # Small sequence for basic tests
+ATTENTION_SEQ_LEN = 256  # Standard attention sequence length
+
+# Position Encoding Dimensions
+POS_GRID_SMALL = 4  # Small grid for position tests (4x3, 4x4)
+CHANNEL_GROUPS = 16  # Channel groups for position IDs
+
+# Training Hyperparameters (common test values)
+DEFAULT_LEARNING_RATE = 1e-4
+DEFAULT_WEIGHT_DECAY = 0.01
+DEFAULT_GRADIENT_CLIP_NORM = 1.0
+DEFAULT_GUIDANCE_SCALE = 3.5  # Typical classifier-free guidance scale
+
+# Scheduler Parameters
+DEFAULT_NUM_TRAIN_TIMESTEPS = 1000
+DEFAULT_SHIFT = 1.0  # Default shift for flow matching scheduler
+
+# Tensor Dimensions
+TENSOR_CHANNELS_RGB = 3  # RGB channels for scheduler tests
+
+# Iteration Counts
+TRAINING_STEPS_FEW = 5  # Few training steps
+TRAINING_STEPS_MODERATE = 10  # Moderate training steps
+ACCUMULATION_STEPS = 4  # Gradient accumulation steps
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/conftest.py b/tests/unit_tests/backends/megatron/diffusion/data/conftest.py
new file mode 100644
index 000000000..97001e263
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/conftest.py
@@ -0,0 +1,257 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Shared fixtures for data module tests.
+
+Provides mocked encoders, sample data, and utilities for testing.
+"""
+
+import io
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from PIL import Image
+
+# ============================================================================
+# Mock Encoder Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def mock_vae_model():
+    """Mock VAE model for testing."""
+    mock = MagicMock()
+
+    # Mock encode method
+    def mock_encode(images):
+        batch_size = images.shape[0] if isinstance(images, torch.Tensor) else 1
+        # Return mock latent distribution
+        latent_dist = MagicMock()
+        latent_dist.sample = MagicMock(return_value=torch.randn(batch_size, 16, 64, 64))
+        return MagicMock(latent_dist=latent_dist)
+
+    mock.encode = mock_encode
+    mock.to = MagicMock(return_value=mock)
+    mock.eval = MagicMock(return_value=mock)
+
+    return mock
+
+
+@pytest.fixture
+def mock_t5_model():
+    """Mock T5 model for testing."""
+    mock = MagicMock()
+
+    # Mock encoder
+    def mock_forward(input_ids, **kwargs):
+        batch_size = input_ids.shape[0]
+        seq_len = input_ids.shape[1]
+        hidden_dim = 4096  # T5-XXL hidden dim
+        return MagicMock(last_hidden_state=torch.randn(batch_size, seq_len, hidden_dim))
+
+    mock.encoder = MagicMock()
+    mock.encoder.return_value = mock_forward(torch.zeros(1, 77, dtype=torch.long))
+    mock.to = MagicMock(return_value=mock)
+    mock.eval = MagicMock(return_value=mock)
+
+    return mock
+
+
+@pytest.fixture
+def mock_t5_tokenizer():
+    """Mock T5 tokenizer for testing."""
+    mock = MagicMock()
+
+    def mock_tokenize(text, **kwargs):
+        if isinstance(text, str):
+            text = [text]
+        batch_size = len(text)
+        max_length = kwargs.get("max_length", 77)
+        return {
+            "input_ids": torch.randint(0, 32000, (batch_size, max_length)),
+            "attention_mask": torch.ones(batch_size, max_length),
+        }
+
+    mock.return_value = mock_tokenize(["test"])
+    mock.side_effect = None
+    mock.__call__ = mock_tokenize
+
+    return mock
+
+
+@pytest.fixture
+def mock_clip_model():
+    """Mock CLIP model for testing."""
+    mock = MagicMock()
+
+    # Mock text encoder
+    def mock_text_model(input_ids, **kwargs):
+        batch_size = input_ids.shape[0]
+        seq_len = input_ids.shape[1]
+        hidden_dim = 768  # CLIP-L hidden dim
+        pooled_dim = 768
+        return MagicMock(
+            last_hidden_state=torch.randn(batch_size, seq_len, hidden_dim),
+            pooler_output=torch.randn(batch_size, pooled_dim),
+        )
+
+    mock.text_model = MagicMock()
+    mock.text_model.return_value = mock_text_model(torch.zeros(1, 77, dtype=torch.long))
+    mock.to = MagicMock(return_value=mock)
+    mock.eval = MagicMock(return_value=mock)
+
+    return mock
+
+
+@pytest.fixture
+def mock_clip_tokenizer():
+    """Mock CLIP tokenizer for testing."""
+    mock = MagicMock()
+
+    def mock_tokenize(text, **kwargs):
+        if isinstance(text, str):
+            text = [text]
+        batch_size = len(text)
+        max_length = kwargs.get("max_length", 77)
+        return {
+            "input_ids": torch.randint(0, 49408, (batch_size, max_length)),
+            "attention_mask": torch.ones(batch_size, max_length),
+        }
+
+    mock.return_value = mock_tokenize(["test"])
+    mock.__call__ = mock_tokenize
+
+    return mock
+
+
+# ============================================================================
+# Sample Data Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def sample_image():
+    """Create a sample PIL Image."""
+    return Image.new("RGB", (512, 512), color="red")
+
+
+@pytest.fixture
+def sample_image_tensor():
+    """Create a sample image tensor."""
+    return torch.randn(3, 512, 512)
+
+
+@pytest.fixture
+def sample_image_bytes():
+    """Create sample image as bytes."""
+    img = Image.new("RGB", (512, 512), color="blue")
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+@pytest.fixture
+def sample_text():
+    """Sample text caption."""
+    return "A beautiful landscape with mountains and trees"
+
+
+@pytest.fixture
+def sample_latents():
+    """Sample VAE latents (without batch dimension for single sample)."""
+    return torch.randn(16, 64, 64)
+
+
+@pytest.fixture
+def sample_prompt_embeds():
+    """Sample T5 text embeddings (without batch dimension for single sample)."""
+    return torch.randn(256, 4096)
+
+
+@pytest.fixture
+def sample_pooled_prompt_embeds():
+    """Sample CLIP pooled embeddings (without batch dimension for single sample)."""
+    return torch.randn(768)
+
+
+@pytest.fixture
+def sample_text_ids():
+    """Sample text position IDs (without batch dimension for single sample)."""
+    return torch.randn(256, 3)
+
+
+# ============================================================================
+# WebDataset Sample Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def preencoded_sample(sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds):
+    """Sample with pre-encoded data (tensor format)."""
+    return {
+        "__key__": "sample_001",
+        "__restore_key__": lambda: "sample_001",
+        "__subflavors__": {"encoding": "preencoded"},
+        "latents.pth": sample_latents,  # Already correct shape (16, 64, 64)
+        "prompt_embeds.pth": sample_prompt_embeds,  # Already correct shape (256, 4096)
+        "pooled_prompt_embeds.pth": sample_pooled_prompt_embeds,  # Already correct shape (768,)
+        "text_ids.pth": None,  # Optional
+        "caption.txt": b"A test caption",
+    }
+
+
+@pytest.fixture
+def preencoded_sample_bytes(sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds):
+    """Sample with pre-encoded data (bytes format)."""
+
+    def tensor_to_bytes(tensor):
+        buf = io.BytesIO()
+        torch.save(tensor, buf)
+        buf.seek(0)
+        return buf.getvalue()
+
+    return {
+        "__key__": "sample_002",
+        "__restore_key__": lambda: "sample_002",
+        "__subflavors__": {"encoding": "preencoded"},
+        "latents.pth": tensor_to_bytes(sample_latents),  # Convert to bytes
+        "prompt_embeds.pth": tensor_to_bytes(sample_prompt_embeds),  # Convert to bytes
+        "pooled_prompt_embeds.pth": tensor_to_bytes(sample_pooled_prompt_embeds),  # Convert to bytes
+        "caption.txt": b"Another test caption",
+    }
+
+
+@pytest.fixture
+def raw_sample(sample_image_bytes):
+    """Sample with raw image data."""
+    return {
+        "__key__": "sample_003",
+        "__restore_key__": lambda: "sample_003",
+        "__subflavors__": {"encoding": "raw"},
+        "images": sample_image_bytes,
+        "txt": b"Raw image caption",
+    }
+
+
+# ============================================================================
+# Temporary Directory Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def temp_output_dir():
+    """Create temporary directory for test outputs."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+@pytest.fixture
+def temp_model_path(temp_output_dir):
+    """Create a temporary model path."""
+    model_dir = temp_output_dir / "model"
+    model_dir.mkdir()
+    return model_dir
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py
new file mode 100644
index 000000000..265ac19c0
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py
@@ -0,0 +1,130 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for base encoder classes.
+
+Tests abstract base classes for encoders, ensuring proper methods and behavior.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.diffusion.encoders.base import (
+    BaseEncoder,
+    BaseTextEncoder,
+    BaseVAE,
+)
+from primus.backends.megatron.data.diffusion.encoders.config import (
+    EncoderConfig,
+    T5XXLConfig,
+    VAEConfig,
+)
+from tests.utils import PrimusUT
+
+
+# Concrete implementations for testing abstract classes
+class ConcreteEncoder(BaseEncoder):
+    """Concrete encoder for testing BaseEncoder."""
+
+    def encode(self, *args, **kwargs):
+        return torch.randn(1, 10)
+
+    @classmethod
+    def from_pretrained(cls, model_path, config=None):
+        return cls(config or EncoderConfig(type="test", model_path=model_path))
+
+
+class ConcreteVAE(BaseVAE):
+    """Concrete VAE for testing BaseVAE."""
+
+    def encode(self, images):
+        batch_size = images.shape[0]
+        latent_h = images.shape[2] // self.latent_downsample_factor
+        latent_w = images.shape[3] // self.latent_downsample_factor
+        return torch.randn(batch_size, self.out_channels, latent_h, latent_w)
+
+    def decode(self, latents):
+        batch_size = latents.shape[0]
+        img_h = latents.shape[2] * self.latent_downsample_factor
+        img_w = latents.shape[3] * self.latent_downsample_factor
+        return torch.randn(batch_size, self.in_channels, img_h, img_w)
+
+    @classmethod
+    def from_pretrained(cls, model_path, config=None):
+        return cls(config or VAEConfig(model_path=model_path))
+
+
+class ConcreteTextEncoder(BaseTextEncoder):
+    """Concrete text encoder for testing BaseTextEncoder."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._tokenizer = Mock()  # Mock tokenizer
+
+    def encode(self, texts, **kwargs):
+        texts_list = self._prepare_texts(texts)
+        batch_size = len(texts_list)
+        return torch.randn(batch_size, self.max_length, self.embedding_dim)
+
+    @classmethod
+    def from_pretrained(cls, model_path, config=None):
+        return cls(config or T5XXLConfig(model_path=model_path))
+
+
+class TestBaseEncoder(PrimusUT):
+    """Tests for BaseEncoder abstract class."""
+
+    def test_get_dtype_invalid_defaults_to_bf16(self):
+        """Test that invalid precision defaults to bfloat16."""
+        config = EncoderConfig(type="test", model_path="/path")
+        encoder = ConcreteEncoder(config)
+
+        assert encoder._get_dtype("unknown") == torch.bfloat16
+
+
+class TestBaseVAE(PrimusUT):
+    """Tests for BaseVAE abstract class."""
+
+    def test_get_latent_shape_various_sizes(self):
+        """Test latent shape calculation with various image sizes."""
+        config = VAEConfig(type="autoencoder_kl", model_path="/path/to/vae")
+        vae = ConcreteVAE(config)
+
+        test_cases = [
+            ((256, 256), (16, 32, 32)),
+            ((512, 512), (16, 64, 64)),
+            ((1024, 1024), (16, 128, 128)),
+            ((512, 1024), (16, 64, 128)),
+        ]
+
+        for (h, w), expected_shape in test_cases:
+            assert vae.get_latent_shape(h, w) == expected_shape
+
+    def test_get_latent_shape_custom_downsample(self):
+        """Test latent shape with custom downsample factor."""
+        config = VAEConfig(type="autoencoder_kl", model_path="/path/to/vae", latent_downsample_factor=16)
+        vae = ConcreteVAE(config)
+
+        latent_shape = vae.get_latent_shape(512, 512)
+
+        assert latent_shape == (16, 32, 32)  # 512/16 = 32
+
+
+class TestBaseTextEncoder(PrimusUT):
+    """Tests for BaseTextEncoder abstract class."""
+
+    def test_tokenizer_not_initialized_error(self):
+        """Test that accessing tokenizer before initialization raises error."""
+        config = T5XXLConfig(type="t5_xxl", model_path="/path/to/t5")
+        encoder = ConcreteTextEncoder(config)
+        encoder._tokenizer = None
+
+        with pytest.raises(ValueError, match="Tokenizer not initialized"):
+            _ = encoder.tokenizer
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py
new file mode 100644
index 000000000..96fdd9b31
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py
@@ -0,0 +1,60 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for CLIP text encoder implementations.
+
+Tests CLIPLEncoder with mocked models for fast unit testing.
+
+NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.)
+have been moved to test_encoder_wrappers_consolidated.py. This file contains only
+CLIP-specific tests that verify Primus's CLIP integration logic.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.diffusion.encoders.config import CLIPLConfig
+
+
+class TestCLIPLEncoder:
+    """Tests for CLIPLEncoder Primus-specific logic."""
+
+    @patch("primus.backends.megatron.data.diffusion.encoders.text.clip_l.CLIPTextModel")
+    @patch("primus.backends.megatron.data.diffusion.encoders.text.clip_l.CLIPTokenizer")
+    def test_pooled_embeddings_extraction(
+        self, mock_tokenizer_cls, mock_model_cls, mock_clip_model, mock_clip_tokenizer
+    ):
+        """Test that pooled embeddings are correctly extracted (Primus-specific CLIP integration)."""
+        from primus.backends.megatron.data.diffusion.encoders.text import CLIPLEncoder
+
+        mock_model_cls.from_pretrained.return_value = mock_clip_model
+        mock_tokenizer_cls.from_pretrained.return_value = mock_clip_tokenizer
+
+        # Setup mock tokenizer
+        mock_clip_tokenizer.return_value = {
+            "input_ids": torch.randint(0, 49408, (2, 77)),
+            "attention_mask": torch.ones(2, 77),
+        }
+
+        # Setup mock model output with specific pooled embeddings
+        expected_pooled = torch.randn(2, 768)
+        mock_clip_model.return_value = MagicMock(
+            last_hidden_state=torch.randn(2, 77, 768), pooler_output=expected_pooled
+        )
+
+        config = CLIPLConfig(type="clip_l", model_path="/path/to/clip")
+        encoder = CLIPLEncoder.from_pretrained("/path/to/clip", config=config)
+
+        sequence_embeds, pooled_embeds = encoder.encode(["Text 1", "Text 2"])
+
+        # Must return the model's pooler_output verbatim (not e.g. a mean-pool of
+        # last_hidden_state, which would have the same shape but wrong values).
+        assert pooled_embeds.shape == expected_pooled.shape
+        assert torch.equal(pooled_embeds, expected_pooled), "pooled_embeds must equal pooler_output"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py
new file mode 100644
index 000000000..d21937f0d
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py
@@ -0,0 +1,30 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for encoder configuration validation logic.
+"""
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.encoders.config import (
+    CLIPLConfig,
+    EncoderConfig,
+    T5XXLConfig,
+    VAEConfig,
+)
+from tests.utils import PrimusUT
+
+
+class TestConfigValidation(PrimusUT):
+    """Tests for config validation across all types."""
+
+    def test_all_configs_validate_precision(self):
+        """Test that all config types validate precision."""
+        for ConfigClass in [EncoderConfig, VAEConfig, T5XXLConfig, CLIPLConfig]:
+            with pytest.raises(ValueError, match="precision must be one of"):
+                ConfigClass(type="test", model_path="/path", precision="invalid_precision")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py
new file mode 100644
index 000000000..3baef0dbf
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py
@@ -0,0 +1,405 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Consolidated tests for encoder wrapper classes (CLIP, T5, VAE).
+
+This file contains parametrized tests for common wrapper functionality across
+different encoder types to avoid duplication. Tests that are specific to individual
+encoder implementations (e.g., VAE scale/shift, CLIP pooled embeddings) remain
+in their respective test files.
+
+Encoder-specific tests remain in:
+- test_clip.py: CLIP-specific logic (pooled embeddings extraction)
+- test_t5.py: T5-specific logic (padding/truncation)
+- test_vae.py: VAE-specific logic (scale/shift, latent shape calculation)
+"""
+
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.diffusion.encoders.config import (
+    CLIPLConfig,
+    T5XXLConfig,
+    VAEConfig,
+)
+
+
+@pytest.fixture
+def encoder_specs():
+    """Define specifications for all encoder types."""
+    return {
+        "clip": {
+            "module_path": "primus.backends.megatron.data.diffusion.encoders.text.clip_l",
+            "encoder_cls_name": "CLIPLEncoder",
+            "config_cls": CLIPLConfig,
+            "config_kwargs": {"type": "clip_l", "model_path": "/path/to/clip"},
+            "model_cls_name": "CLIPTextModel",
+            "tokenizer_cls_name": "CLIPTokenizer",
+            "max_length": 77,
+            "embedding_dim": 768,
+            "has_tokenizer": True,
+        },
+        "t5": {
+            "module_path": "primus.backends.megatron.data.diffusion.encoders.text.t5_xxl",
+            "encoder_cls_name": "T5XXLEncoder",
+            "config_cls": T5XXLConfig,
+            "config_kwargs": {"type": "t5_xxl", "model_path": "/path/to/t5"},
+            "model_cls_name": "T5EncoderModel",
+            "tokenizer_cls_name": "T5Tokenizer",
+            "max_length": 512,
+            "embedding_dim": 4096,
+            "has_tokenizer": True,
+        },
+        "vae": {
+            "module_path": "primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl",
+            "encoder_cls_name": "AutoencoderKL",
+            "config_cls": VAEConfig,
+            "config_kwargs": {"type": "autoencoder_kl", "model_path": "/path/to/vae"},
+            "model_cls_name": "DiffusersAutoencoderKL",
+            "tokenizer_cls_name": None,
+            "scale_factor": 0.3611,
+            "shift_factor": 0.1159,
+            "in_channels": 3,
+            "out_channels": 16,
+            "has_tokenizer": False,
+        },
+    }
+
+
+@pytest.fixture
+def mock_clip_model():
+    """Create a mock CLIP model."""
+    model = Mock()
+    model.return_value = MagicMock(
+        last_hidden_state=torch.randn(1, 77, 768), pooler_output=torch.randn(1, 768)
+    )
+    return model
+
+
+@pytest.fixture
+def mock_clip_tokenizer():
+    """Create a mock CLIP tokenizer."""
+    tokenizer = Mock()
+    tokenizer.return_value = {
+        "input_ids": torch.randint(0, 49408, (1, 77)),
+        "attention_mask": torch.ones(1, 77),
+    }
+    return tokenizer
+
+
+@pytest.fixture
+def mock_t5_model():
+    """Create a mock T5 model."""
+    model = Mock()
+    model.return_value = MagicMock(last_hidden_state=torch.randn(1, 512, 4096))
+    return model
+
+
+@pytest.fixture
+def mock_t5_tokenizer():
+    """Create a mock T5 tokenizer."""
+    tokenizer = Mock()
+    tokenizer.return_value = {
+        "input_ids": torch.randint(0, 32000, (1, 512)),
+        "attention_mask": torch.ones(1, 512),
+    }
+    return tokenizer
+
+
+@pytest.fixture
+def mock_vae_model():
+    """Create a mock VAE model."""
+    model = Mock()
+    # Mock encode method
+    mock_dist = Mock()
+    mock_dist.sample.return_value = torch.randn(1, 16, 64, 64)
+    model.encode.return_value = mock_dist
+    # Mock decode method - return tuple like real diffusers AutoencoderKL
+    model.decode.return_value = (torch.randn(1, 3, 512, 512),)
+    return model
+
+
+class TestEncoderFromPretrained:
+    """Test from_pretrained for all encoder types."""
+
+    @pytest.mark.parametrize("encoder_type", ["clip", "t5"])
+    def test_text_encoder_from_pretrained_without_config(
+        self,
+        encoder_type,
+        encoder_specs,
+        mock_clip_model,
+        mock_clip_tokenizer,
+        mock_t5_model,
+        mock_t5_tokenizer,
+    ):
+        """Test loading text encoders without explicit config."""
+        spec = encoder_specs[encoder_type]
+
+        # Select appropriate mocks
+        if encoder_type == "clip":
+            mock_model = mock_clip_model
+            mock_tokenizer = mock_clip_tokenizer
+        else:  # t5
+            mock_model = mock_t5_model
+            mock_tokenizer = mock_t5_tokenizer
+
+        with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch(
+            f"{spec['module_path']}.{spec['tokenizer_cls_name']}"
+        ) as mock_tokenizer_cls:
+
+            mock_model_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+
+            # Import encoder class
+            module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+            encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+            encoder = encoder_cls.from_pretrained("/path")
+
+            assert encoder is not None
+            assert encoder.config.model_path == "/path"
+            assert encoder.config.precision == "bf16"
+
+
+class TestEncoderPrecisionMapping:
+    """Test precision to dtype mapping for all encoders."""
+
+    @pytest.mark.parametrize(
+        "encoder_type,precision,expected_dtype",
+        [
+            ("clip", "bf16", torch.bfloat16),
+            ("clip", "fp16", torch.float16),
+            ("clip", "fp32", torch.float32),
+            ("t5", "bf16", torch.bfloat16),
+            ("t5", "fp16", torch.float16),
+            ("t5", "fp32", torch.float32),
+            ("vae", "bf16", torch.bfloat16),
+            ("vae", "fp16", torch.float16),
+            ("vae", "fp32", torch.float32),
+        ],
+    )
+    def test_precision_dtype_mapping(
+        self,
+        encoder_type,
+        precision,
+        expected_dtype,
+        encoder_specs,
+        mock_clip_model,
+        mock_clip_tokenizer,
+        mock_t5_model,
+        mock_t5_tokenizer,
+        mock_vae_model,
+    ):
+        """Test that precision config is correctly mapped to torch dtype."""
+        spec = encoder_specs[encoder_type]
+
+        # Select appropriate mocks
+        if encoder_type == "clip":
+            mock_model = mock_clip_model
+            mock_tokenizer = mock_clip_tokenizer
+        elif encoder_type == "t5":
+            mock_model = mock_t5_model
+            mock_tokenizer = mock_t5_tokenizer
+        else:  # vae
+            mock_model = mock_vae_model
+            mock_tokenizer = None
+
+        if spec["has_tokenizer"]:
+            with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch(
+                f"{spec['module_path']}.{spec['tokenizer_cls_name']}"
+            ) as mock_tokenizer_cls:
+
+                mock_model_cls.from_pretrained.return_value = mock_model
+                mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+
+                # Import encoder class
+                module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+                encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+                config = spec["config_cls"](**spec["config_kwargs"], precision=precision)
+                _ = encoder_cls.from_pretrained("/path", config=config)
+
+                # Check that from_pretrained was called with correct dtype
+                call_kwargs = mock_model_cls.from_pretrained.call_args[1]
+                assert call_kwargs["torch_dtype"] == expected_dtype
+        else:
+            # VAE case
+            with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls:
+                mock_model_cls.from_pretrained.return_value = mock_model
+
+                # Import encoder class
+                module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+                encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+                config = spec["config_cls"](**spec["config_kwargs"], precision=precision)
+                encoder_cls.from_pretrained("/path", config=config)
+
+                # Check that from_pretrained was called with correct dtype
+                call_kwargs = mock_model_cls.from_pretrained.call_args[1]
+                assert call_kwargs["torch_dtype"] == expected_dtype
+
+
+class TestEncoderDeviceHandling:
+    """Test device handling for all encoder types."""
+
+    @pytest.mark.parametrize("encoder_type", ["clip", "t5", "vae"])
+    def test_device_handling(
+        self,
+        encoder_type,
+        encoder_specs,
+        mock_clip_model,
+        mock_clip_tokenizer,
+        mock_t5_model,
+        mock_t5_tokenizer,
+        mock_vae_model,
+    ):
+        """Test encoder device configuration."""
+        spec = encoder_specs[encoder_type]
+
+        # Select appropriate mocks
+        if encoder_type == "clip":
+            mock_model = mock_clip_model
+            mock_tokenizer = mock_clip_tokenizer
+        elif encoder_type == "t5":
+            mock_model = mock_t5_model
+            mock_tokenizer = mock_t5_tokenizer
+        else:  # vae
+            mock_model = mock_vae_model
+            mock_tokenizer = None
+
+        if spec["has_tokenizer"]:
+            with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch(
+                f"{spec['module_path']}.{spec['tokenizer_cls_name']}"
+            ) as mock_tokenizer_cls:
+
+                mock_model_cls.from_pretrained.return_value = mock_model
+                mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+
+                # Import encoder class
+                module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+                encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+                config = spec["config_cls"](**spec["config_kwargs"], device="cpu")
+                encoder = encoder_cls.from_pretrained("/path", config=config)
+
+                assert encoder.device.type == "cpu"
+        else:
+            # VAE case
+            with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls:
+                mock_model_cls.from_pretrained.return_value = mock_model
+
+                # Import encoder class
+                module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+                encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+                config = spec["config_cls"](**spec["config_kwargs"], device="cpu")
+                vae = encoder_cls.from_pretrained("/path", config=config)
+
+                assert vae.device.type == "cpu"
+
+
+class TestEncoderImportErrors:
+    """Test import error handling when external libraries are not available."""
+
+    @pytest.mark.parametrize("encoder_type", ["clip", "t5"])
+    def test_text_encoder_import_error(self, encoder_type, encoder_specs):
+        """Test that ImportError is raised when transformers is not installed."""
+        spec = encoder_specs[encoder_type]
+
+        with patch(f"{spec['module_path']}.{spec['model_cls_name']}", None), patch(
+            f"{spec['module_path']}.{spec['tokenizer_cls_name']}", None
+        ):
+
+            # Import encoder class
+            module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+            encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+            config = spec["config_cls"](**spec["config_kwargs"])
+
+            with pytest.raises(ImportError, match="transformers library is required"):
+                encoder_cls(config)
+
+    def test_vae_import_error(self, encoder_specs):
+        """Test that ImportError is raised when diffusers is not installed."""
+        spec = encoder_specs["vae"]
+
+        with patch(f"{spec['module_path']}.{spec['model_cls_name']}", None):
+            # Import encoder class
+            module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+            encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+            config = spec["config_cls"](**spec["config_kwargs"])
+
+            with pytest.raises(ImportError, match="diffusers library is required"):
+                encoder_cls(config)
+
+
+class TestEncoderCacheDir:
+    """Test cache_dir configuration support across all encoders."""
+
+    @pytest.mark.parametrize("encoder_type", ["clip", "t5", "vae"])
+    def test_cache_dir_passed_to_from_pretrained(self, encoder_type, encoder_specs):
+        """Test that cache_dir is passed through to HuggingFace from_pretrained."""
+        spec = encoder_specs[encoder_type]
+        config_cls = spec["config_cls"]
+        config_kwargs = spec["config_kwargs"].copy()
+        cache_dir = "/custom/cache/path"
+
+        # Create config with cache_dir
+        config = config_cls(**config_kwargs, cache_dir=cache_dir)
+
+        # Mock the model/tokenizer classes
+        mock_model = MagicMock()
+        mock_model_instance = Mock()
+        mock_model.from_pretrained.return_value = mock_model_instance
+        mock_model_instance.to.return_value = mock_model_instance
+
+        # FIXED: Use just attribute names for patch.multiple (not full module paths)
+        patches = {spec["model_cls_name"]: mock_model}
+
+        # Add tokenizer patch for text encoders
+        if spec["has_tokenizer"]:
+            mock_tokenizer = MagicMock()
+            mock_tokenizer.from_pretrained.return_value = Mock()
+            patches[spec["tokenizer_cls_name"]] = mock_tokenizer
+
+        # Also mock the helper function
+        mock_helper = MagicMock(return_value=mock_model_instance)
+        patches["load_pretrained_with_subfolder_fallback"] = mock_helper
+
+        with patch.dict("sys.modules", {k: MagicMock() for k in ["transformers", "diffusers"]}):
+            with patch.multiple(spec["module_path"], **patches):
+                # Import and call from_pretrained
+                module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]])
+                encoder_cls = getattr(module, spec["encoder_cls_name"])
+
+                encoder_cls.from_pretrained("/path", config=config)
+
+                # The helper must be invoked, and cache_dir must be forwarded to it.
+                assert mock_helper.called, "load_pretrained_with_subfolder_fallback was never called"
+                call_kwargs = mock_helper.call_args.kwargs
+                assert "cache_dir" in call_kwargs, "cache_dir not passed to helper function"
+                assert call_kwargs["cache_dir"] == cache_dir
+
+    def test_flux_encoder_config_cache_dir(self):
+        """Test that FluxEncoderConfig.from_pretrained_flux accepts cache_dir."""
+        from primus.backends.megatron.data.diffusion.encoders.config import (
+            FluxEncoderConfig,
+        )
+
+        cache_dir = "/shared/hf_cache"
+        flux_config = FluxEncoderConfig.from_pretrained_flux(
+            model_path="black-forest-labs/FLUX.1-dev", cache_dir=cache_dir
+        )
+
+        # All sub-configs should have the same cache_dir
+        assert flux_config.vae.cache_dir == cache_dir
+        assert flux_config.t5.cache_dir == cache_dir
+        assert flux_config.clip.cache_dir == cache_dir
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py
new file mode 100644
index 000000000..37c583140
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for T5 text encoder implementations.
+
+Tests T5XXLEncoder with mocked models for fast unit testing.
+
+NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.)
+have been moved to test_encoder_wrappers_consolidated.py. This file contains only
+T5-specific tests that verify Primus's T5 integration logic.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.diffusion.encoders.config import T5XXLConfig
+
+
+class TestT5XXLEncoder:
+    """Tests for T5XXLEncoder Primus-specific logic."""
+
+    @patch("primus.backends.megatron.data.diffusion.encoders.text.t5_xxl.T5EncoderModel")
+    @patch("primus.backends.megatron.data.diffusion.encoders.text.t5_xxl.T5Tokenizer")
+    def test_encode_with_padding_truncation(
+        self, mock_tokenizer_cls, mock_model_cls, mock_t5_model, mock_t5_tokenizer
+    ):
+        """Test that tokenizer is called with correct padding/truncation (Primus tokenization logic)."""
+        from primus.backends.megatron.data.diffusion.encoders.text import T5XXLEncoder
+
+        mock_model_cls.from_pretrained.return_value = mock_t5_model
+        mock_tokenizer_cls.from_pretrained.return_value = mock_t5_tokenizer
+
+        # Setup mock tokenizer
+        mock_t5_tokenizer.return_value = {
+            "input_ids": torch.randint(0, 32000, (1, 512)),
+            "attention_mask": torch.ones(1, 512),
+        }
+
+        # Setup mock model output
+        mock_t5_model.return_value = MagicMock(last_hidden_state=torch.randn(1, 512, 4096))
+
+        config = T5XXLConfig(type="t5_xxl", model_path="/path/to/t5")
+        encoder = T5XXLEncoder.from_pretrained("/path/to/t5", config=config)
+
+        encoder.encode("Test caption")
+
+        # Verify tokenizer was called with correct Primus padding/truncation params
+        call_kwargs = mock_t5_tokenizer.call_args[1]
+        assert call_kwargs["padding"] == "max_length"
+        assert call_kwargs["truncation"] is True
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py
new file mode 100644
index 000000000..4c936c995
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py
@@ -0,0 +1,58 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for VAE encoder implementations.
+
+Tests AutoencoderKL with mocked models for fast unit testing.
+
+NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.)
+have been moved to test_encoder_wrappers_consolidated.py. This file contains only
+VAE-specific tests that verify Primus's VAE integration logic (scale/shift, latent shapes).
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.diffusion.encoders.config import VAEConfig
+
+
+class TestAutoencoderKL:
+    """Tests for AutoencoderKL Primus-specific logic."""
+
+    @patch("primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl.DiffusersAutoencoderKL")
+    def test_encode_applies_scale_and_shift(self, mock_diffusers_vae, mock_vae_model):
+        """encode() must apply z = scale_factor * (sample - shift_factor) to the VAE sample."""
+        from primus.backends.megatron.data.diffusion.encoders.image import AutoencoderKL
+
+        mock_diffusers_vae.from_pretrained.return_value = mock_vae_model
+
+        config = VAEConfig(
+            type="autoencoder_kl", model_path="/path/to/vae", scale_factor=0.5, shift_factor=0.1
+        )
+        vae = AutoencoderKL.from_pretrained("/path/to/vae", config=config)
+
+        assert vae.config.scale_factor == 0.5
+        assert vae.config.shift_factor == 0.1
+
+        # Pin the VAE's sampled latent so we can assert the exact scale/shift math
+        # (the shared fixture returns a fresh random sample each call).
+        known_sample = torch.randn(1, 16, 64, 64)
+        latent_dist = MagicMock()
+        latent_dist.sample = MagicMock(return_value=known_sample)
+        vae.vae.encode = MagicMock(return_value=MagicMock(latent_dist=latent_dist))
+
+        images = torch.randn(1, 3, 256, 256)
+        latents = vae.encode(images)
+
+        expected = config.scale_factor * (
+            known_sample.to(latents.device, latents.dtype) - config.shift_factor
+        )
+        assert latents.shape == known_sample.shape
+        torch.testing.assert_close(latents, expected, rtol=1e-5, atol=1e-6)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py
new file mode 100644
index 000000000..f98126b54
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py
@@ -0,0 +1,251 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for download utilities (download.py).
+
+Covers:
+- parse_md5_manifest: parsing, filtering, sorting
+- download_with_backoff: retry logic, MD5 verification, timeout handling
+- fetch_url_text: basic HTTP text fetch
+"""
+
+import hashlib
+import io
+import tempfile
+import urllib.error
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.preprocessing.download import (
+    download_with_backoff,
+    fetch_url_text,
+    parse_md5_manifest,
+)
+
+_DOWNLOAD_MODULE = "primus.backends.megatron.data.diffusion.preprocessing.download"
+
+
+class TestParseMd5Manifest:
+    def test_basic_parsing(self):
+        text = "abc123 data-00000.arrow\ndef456 data-00001.arrow\n"
+        entries = parse_md5_manifest(text)
+        assert entries == [("abc123", "data-00000.arrow"), ("def456", "data-00001.arrow")]
+
+    def test_suffix_filter(self):
+        text = "abc123 data-00000.arrow\ndef456 readme.txt\n"
+        entries = parse_md5_manifest(text, suffix_filter=".arrow")
+        assert len(entries) == 1
+        assert entries[0][1] == "data-00000.arrow"
+
+    def test_no_filter_returns_all(self):
+        text = "abc123 data-00000.arrow\ndef456 readme.txt\n"
+        entries = parse_md5_manifest(text)
+        assert len(entries) == 2
+
+    def test_sorts_by_filename(self):
+        text = "bbb data-00002.arrow\naaa data-00001.arrow\nccc data-00000.arrow\n"
+        entries = parse_md5_manifest(text)
+        assert [e[1] for e in entries] == [
+            "data-00000.arrow",
+            "data-00001.arrow",
+            "data-00002.arrow",
+        ]
+
+    def test_skips_malformed_lines(self):
+        text = "abc123 data-00000.arrow\nbadline\n\n  \n"
+        entries = parse_md5_manifest(text)
+        assert len(entries) == 1
+
+
+class TestDownloadWithBackoff:
+    """Tests for download_with_backoff with mocked HTTP."""
+
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_success_on_first_try(self, mock_urlopen):
+        content = b"hello world"
+        mock_resp = MagicMock()
+        mock_resp.read.side_effect = [content, b""]
+        mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+        mock_resp.__exit__ = MagicMock(return_value=False)
+        mock_urlopen.return_value = mock_resp
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            download_with_backoff("http://example.com/test.bin", dest, max_retries=0)
+            assert dest.exists()
+
+    @patch(f"{_DOWNLOAD_MODULE}.time.sleep")
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_retries_on_429(self, mock_urlopen, mock_sleep):
+        content = b"success data"
+        mock_resp_ok = MagicMock()
+        mock_resp_ok.read.side_effect = [content, b""]
+        mock_resp_ok.__enter__ = MagicMock(return_value=mock_resp_ok)
+        mock_resp_ok.__exit__ = MagicMock(return_value=False)
+
+        mock_urlopen.side_effect = [
+            urllib.error.HTTPError("http://x", 429, "Too Many Requests", {}, io.BytesIO(b"")),
+            urllib.error.HTTPError("http://x", 429, "Too Many Requests", {}, io.BytesIO(b"")),
+            mock_resp_ok,
+        ]
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            download_with_backoff("http://example.com/test.bin", dest, max_retries=3, base_delay=0.01)
+            assert dest.exists()
+            assert mock_urlopen.call_count == 3
+            assert mock_sleep.call_count == 2
+
+    @patch(f"{_DOWNLOAD_MODULE}.time.sleep")
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_retries_on_503(self, mock_urlopen, mock_sleep):
+        mock_urlopen.side_effect = urllib.error.HTTPError(
+            "http://x", 503, "Service Unavailable", {}, io.BytesIO(b"")
+        )
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            with pytest.raises(RuntimeError, match="HTTP 503"):
+                download_with_backoff("http://example.com/test.bin", dest, max_retries=2, base_delay=0.01)
+            assert mock_urlopen.call_count == 3  # initial + 2 retries
+
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_no_retry_on_404(self, mock_urlopen):
+        mock_urlopen.side_effect = urllib.error.HTTPError("http://x", 404, "Not Found", {}, io.BytesIO(b""))
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            with pytest.raises(RuntimeError, match="HTTP 404"):
+                download_with_backoff("http://example.com/test.bin", dest, max_retries=3, base_delay=0.01)
+            assert mock_urlopen.call_count == 1
+
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_md5_verification_pass(self, mock_urlopen):
+        content = b"test content"
+        expected_md5 = hashlib.md5(content).hexdigest()
+
+        mock_resp = MagicMock()
+        mock_resp.read.side_effect = [content, b""]
+        mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+        mock_resp.__exit__ = MagicMock(return_value=False)
+        mock_urlopen.return_value = mock_resp
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            download_with_backoff("http://example.com/test.bin", dest, expected_md5=expected_md5)
+            assert dest.exists()
+
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_md5_verification_fail(self, mock_urlopen):
+        content = b"test content"
+
+        mock_resp = MagicMock()
+        mock_resp.read.side_effect = [content, b""]
+        mock_resp.__enter__ = MagicMock(return_value=mock_resp)
+        mock_resp.__exit__ = MagicMock(return_value=False)
+        mock_urlopen.return_value = mock_resp
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            with pytest.raises(RuntimeError, match="MD5 mismatch"):
+                download_with_backoff(
+                    "http://example.com/test.bin",
+                    dest,
+                    expected_md5="bad_md5",
+                    max_retries=0,
+                )
+            assert not dest.exists()
+
+    @patch(f"{_DOWNLOAD_MODULE}.time.sleep")
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_retries_on_network_error(self, mock_urlopen, mock_sleep):
+        content = b"ok"
+        mock_resp_ok = MagicMock()
+        mock_resp_ok.read.side_effect = [content, b""]
+        mock_resp_ok.__enter__ = MagicMock(return_value=mock_resp_ok)
+        mock_resp_ok.__exit__ = MagicMock(return_value=False)
+
+        mock_urlopen.side_effect = [
+            urllib.error.URLError("Connection refused"),
+            mock_resp_ok,
+        ]
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            download_with_backoff("http://example.com/test.bin", dest, max_retries=2, base_delay=0.01)
+            assert dest.exists()
+            assert mock_urlopen.call_count == 2
+
+    @patch(f"{_DOWNLOAD_MODULE}.time.sleep")
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_retries_on_md5_mismatch_then_succeeds(self, mock_urlopen, mock_sleep):
+        good_content = b"correct data"
+        bad_content = b"corrupted data"
+        expected_md5 = hashlib.md5(good_content).hexdigest()
+
+        mock_resp_bad = MagicMock()
+        mock_resp_bad.read.side_effect = [bad_content, b""]
+        mock_resp_bad.__enter__ = MagicMock(return_value=mock_resp_bad)
+        mock_resp_bad.__exit__ = MagicMock(return_value=False)
+
+        mock_resp_good = MagicMock()
+        mock_resp_good.read.side_effect = [good_content, b""]
+        mock_resp_good.__enter__ = MagicMock(return_value=mock_resp_good)
+        mock_resp_good.__exit__ = MagicMock(return_value=False)
+
+        mock_urlopen.side_effect = [mock_resp_bad, mock_resp_good]
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            download_with_backoff(
+                "http://example.com/test.bin",
+                dest,
+                expected_md5=expected_md5,
+                max_retries=3,
+                base_delay=0.01,
+            )
+            assert dest.exists()
+            assert dest.read_bytes() == good_content
+            assert mock_urlopen.call_count == 2
+            assert mock_sleep.call_count == 1
+
+    @patch(f"{_DOWNLOAD_MODULE}.time.sleep")
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_md5_mismatch_exhausts_retries(self, mock_urlopen, mock_sleep):
+        bad_content = b"always bad"
+        expected_md5 = "0000000000000000"
+
+        def make_bad_resp(*args, **kwargs):
+            resp = MagicMock()
+            resp.read.side_effect = [bad_content, b""]
+            resp.__enter__ = MagicMock(return_value=resp)
+            resp.__exit__ = MagicMock(return_value=False)
+            return resp
+
+        mock_urlopen.side_effect = make_bad_resp
+
+        with tempfile.TemporaryDirectory() as tmp:
+            dest = Path(tmp) / "test.bin"
+            with pytest.raises(RuntimeError, match="MD5 mismatch"):
+                download_with_backoff(
+                    "http://example.com/test.bin",
+                    dest,
+                    expected_md5=expected_md5,
+                    max_retries=2,
+                    base_delay=0.01,
+                )
+            assert mock_urlopen.call_count == 3  # initial + 2 retries
+
+
+class TestFetchUrlText:
+    @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen")
+    def test_returns_decoded_text(self, mock_urlopen):
+        mock_resp = MagicMock()
+        mock_resp.read.return_value = b"https://example.com/base"
+        mock_urlopen.return_value = mock_resp
+
+        result = fetch_url_text("http://example.com/manifest.uri")
+        assert result == "https://example.com/base"
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py
new file mode 100644
index 000000000..7ab9a7c05
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py
@@ -0,0 +1,121 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for dataset finalization and encoder auth error detection.
+
+Tests finalize_energon_dataset from finalize.py and
+_raise_encoder_auth_error from encoded.py.
+"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.preprocessing.finalize import (
+    finalize_energon_dataset,
+)
+from primus.backends.megatron.data.diffusion.preprocessing.pipelines.encoded import (
+    EncodedDatasetPipeline,
+)
+from tests.utils import PrimusUT
+
+
+def _create_dummy_tar(directory: Path, name: str = "000000.tar") -> None:
+    """Create an empty .tar file for testing."""
+    (directory / name).write_bytes(b"")
+
+
+class TestDatasetYamlContent(PrimusUT):
+    """Tests that finalize_energon_dataset writes correct dataset.yaml."""
+
+    @patch(
+        "primus.backends.megatron.data.diffusion.preprocessing.finalize._prepare_programmatic",
+        side_effect=ImportError("forced: exercise subprocess fallback path"),
+    )
+    @patch("primus.backends.megatron.data.diffusion.preprocessing.validate.validate_energon_dataset")
+    @patch("subprocess.Popen")
+    def test_preencoded_yaml(self, mock_popen, mock_validate, mock_programmatic):
+        """dataset.yaml contains CrudeWebdataset with encoding: preencoded."""
+        mock_process = MagicMock()
+        mock_process.communicate.return_value = ("Done\n", None)
+        mock_process.returncode = 0
+        mock_popen.return_value = mock_process
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _create_dummy_tar(Path(tmpdir))
+            finalize_energon_dataset(output_dir=tmpdir, train_split=1.0, encoding="preencoded")
+
+            dataset_yaml = Path(tmpdir) / ".nv-meta" / "dataset.yaml"
+            assert dataset_yaml.exists()
+            content = dataset_yaml.read_text()
+            assert "__class__: CrudeWebdataset" in content
+            assert "encoding: preencoded" in content
+
+            # Verify split_input passed to subprocess stdin
+            mock_process.communicate.assert_called_once_with(input="1.0, 0.0, 0.0\nn\n", timeout=300)
+
+    @patch(
+        "primus.backends.megatron.data.diffusion.preprocessing.finalize._prepare_programmatic",
+        side_effect=ImportError("forced: exercise subprocess fallback path"),
+    )
+    @patch("primus.backends.megatron.data.diffusion.preprocessing.validate.validate_energon_dataset")
+    @patch("subprocess.Popen")
+    def test_raw_yaml_and_split_format(self, mock_popen, mock_validate, mock_programmatic):
+        """dataset.yaml contains encoding: raw; split ratios are formatted correctly."""
+        mock_process = MagicMock()
+        mock_process.communicate.return_value = ("Done\n", None)
+        mock_process.returncode = 0
+        mock_popen.return_value = mock_process
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _create_dummy_tar(Path(tmpdir))
+            finalize_energon_dataset(output_dir=tmpdir, train_split=0.8, encoding="raw")
+
+            content = (Path(tmpdir) / ".nv-meta" / "dataset.yaml").read_text()
+            assert "encoding: raw" in content
+
+            call_kwargs = mock_process.communicate.call_args
+            split_input = call_kwargs.kwargs.get("input") or call_kwargs[1].get("input")
+            assert split_input.startswith("0.8, ")
+            assert split_input.endswith("\nn\n")
+            parts = split_input.split("\n")[0].split(", ")
+            self.assertAlmostEqual(float(parts[0]), 0.8)
+            self.assertAlmostEqual(float(parts[1]), 0.1)
+            self.assertAlmostEqual(float(parts[2]), 0.1)
+
+
+class TestRaiseEncoderAuthError(PrimusUT):
+    """Tests for EncodedDatasetPipeline._raise_encoder_auth_error."""
+
+    def test_401_error_gives_auth_hint(self):
+        """Exception containing '401' raises RuntimeError with auth hint."""
+        original = Exception("HTTP 401 Unauthorized")
+        with self.assertRaises(RuntimeError) as ctx:
+            EncodedDatasetPipeline._raise_encoder_auth_error("VAE", "my-model", original)
+        assert "HuggingFace authentication" in str(ctx.exception)
+        assert "VAE" in str(ctx.exception)
+
+    def test_token_error_gives_auth_hint(self):
+        """Exception containing 'token' raises RuntimeError with auth hint."""
+        original = Exception("Please pass a valid token")
+        with self.assertRaises(RuntimeError) as ctx:
+            EncodedDatasetPipeline._raise_encoder_auth_error("T5-XXL", "my-model", original)
+        assert "HuggingFace authentication" in str(ctx.exception)
+        assert "T5-XXL" in str(ctx.exception)
+
+    def test_non_auth_error_preserves_message(self):
+        """Non-auth exception re-raises with original message intact."""
+        original = Exception("CUDA out of memory")
+        with self.assertRaises(RuntimeError) as ctx:
+            EncodedDatasetPipeline._raise_encoder_auth_error("CLIP-L", "my-model", original)
+        msg = str(ctx.exception)
+        assert "CUDA out of memory" in msg
+        assert "CLIP-L" in msg
+        assert "HuggingFace authentication" not in msg
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py
new file mode 100644
index 000000000..65c1860cb
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py
@@ -0,0 +1,330 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for the streaming ingest pipeline (pipelines/ingest.py).
+
+Covers:
+- StreamingIngestPipeline: parallel download + sequential conversion
+- Resume functionality (existing shards are skipped)
+- max_files parameter
+- Sample offset accumulation
+- Arrow file cleanup after conversion
+- Skip-and-log for failed downloads and conversions
+"""
+
+import json
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+from primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest import (
+    StreamingIngestPipeline,
+)
+
+_INGEST_MODULE = "primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest"
+
+
+class TestStreamingIngestPipeline:
+    """Tests for the StreamingIngestPipeline with mocked I/O."""
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_processes_all_files(self, mock_manifest, mock_download, mock_convert):
+        """Pipeline processes all entries, creates correct number of shards."""
+        entries = [
+            ("md5_0", "data-00000.arrow"),
+            ("md5_1", "data-00001.arrow"),
+            ("md5_2", "data-00002.arrow"),
+        ]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_convert.return_value = 100
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=2,
+                prefetch_depth=2,
+            )
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake")
+
+            mock_download.side_effect = fake_download
+
+            results = pipeline.run()
+
+        assert results["files_processed"] == 3
+        assert results["samples_written"] == 300
+        assert results["shards_created"] == 3
+        assert mock_download.call_count == 3
+        assert mock_convert.call_count == 3
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_max_files_limits_processing(self, mock_manifest, mock_download, mock_convert):
+        """max_files parameter limits how many files are processed."""
+        entries = [(f"md5_{i}", f"data-{i:05d}.arrow") for i in range(10)]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_convert.return_value = 50
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_files=3,
+                max_workers=2,
+                prefetch_depth=2,
+            )
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake")
+
+            mock_download.side_effect = fake_download
+
+            results = pipeline.run()
+
+        assert results["files_processed"] == 3
+        assert mock_download.call_count == 3
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_download_error_skips_file(self, mock_manifest, mock_download, mock_convert):
+        """A failed download is skipped; the pipeline continues with remaining files."""
+        entries = [("md5_0", "data-00000.arrow"), ("md5_1", "data-00001.arrow")]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_convert.return_value = 100
+
+        call_count = [0]
+
+        def selective_download(url, dest, **kwargs):
+            call_count[0] += 1
+            if "data-00000" in url:
+                raise RuntimeError("Download failed: HTTP 500")
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            dest.write_bytes(b"fake")
+
+        mock_download.side_effect = selective_download
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            results = pipeline.run()
+
+            assert results["files_processed"] == 1
+            assert results["files_failed"] == 1
+            assert results["samples_written"] == 100
+            assert mock_convert.call_count == 1
+
+            failed_manifest = Path(tmp) / "output" / "train" / "failed_files.json"
+            assert failed_manifest.exists()
+            failures = json.loads(failed_manifest.read_text())
+            assert len(failures) == 1
+            assert failures[0]["stage"] == "download"
+            assert failures[0]["filename"] == "data-00000.arrow"
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_all_downloads_fail(self, mock_manifest, mock_download, mock_convert):
+        """If all downloads fail, the pipeline completes with zero processed."""
+        entries = [("md5_0", "data-00000.arrow"), ("md5_1", "data-00001.arrow")]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_download.side_effect = RuntimeError("Download failed: HTTP 500")
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            results = pipeline.run()
+
+        assert results["files_processed"] == 0
+        assert results["files_failed"] == 2
+        assert results["samples_written"] == 0
+        assert mock_convert.call_count == 0
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_conversion_error_skips_file(self, mock_manifest, mock_download, mock_convert):
+        """A failed conversion is skipped; partial tar is cleaned up."""
+        entries = [
+            ("md5_0", "data-00000.arrow"),
+            ("md5_1", "data-00001.arrow"),
+        ]
+        mock_manifest.return_value = ("https://base.url", entries)
+
+        def selective_convert(arrow_path, tar_path, offset):
+            if "data-00000" in str(arrow_path):
+                raise RuntimeError("Corrupt Arrow file")
+            return 100
+
+        mock_convert.side_effect = selective_convert
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake")
+
+            mock_download.side_effect = fake_download
+
+            results = pipeline.run()
+
+            assert results["files_processed"] == 1
+            assert results["files_failed"] == 1
+            assert results["samples_written"] == 100
+
+            output_train = Path(tmp) / "output" / "train"
+            assert not (output_train / "shard_000000.tar").exists()
+            assert (output_train / "failed_files.json").exists()
+            failures = json.loads((output_train / "failed_files.json").read_text())
+            assert len(failures) == 1
+            assert failures[0]["stage"] == "conversion"
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_arrow_files_deleted_after_conversion(self, mock_manifest, mock_download, mock_convert):
+        """Arrow files are cleaned up after conversion to tar."""
+        entries = [("md5_0", "data-00000.arrow")]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_convert.return_value = 10
+
+        with tempfile.TemporaryDirectory() as tmp:
+            arrow_dir = Path(tmp) / "arrows"
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=str(arrow_dir),
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            created_files = []
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake arrow data")
+                created_files.append(dest)
+
+            mock_download.side_effect = fake_download
+
+            pipeline.run()
+
+            for f in created_files:
+                assert not f.exists(), f"Arrow file should have been deleted: {f}"
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_sample_offset_accumulates(self, mock_manifest, mock_download, mock_convert):
+        """Global sample offset is passed correctly across shards."""
+        entries = [
+            ("md5_0", "data-00000.arrow"),
+            ("md5_1", "data-00001.arrow"),
+            ("md5_2", "data-00002.arrow"),
+        ]
+        mock_manifest.return_value = ("https://base.url", entries)
+
+        sample_counts = [100, 200, 150]
+        mock_convert.side_effect = sample_counts
+
+        with tempfile.TemporaryDirectory() as tmp:
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake")
+
+            mock_download.side_effect = fake_download
+
+            results = pipeline.run()
+
+        offsets = [c.args[2] for c in mock_convert.call_args_list]
+        assert offsets == [0, 100, 300]
+        assert results["samples_written"] == 450
+
+    @patch(f"{_INGEST_MODULE}._arrow_to_tar")
+    @patch(f"{_INGEST_MODULE}.download_with_backoff")
+    @patch(f"{_INGEST_MODULE}.fetch_manifest")
+    def test_resume_skips_existing_shards(self, mock_manifest, mock_download, mock_convert):
+        """Pre-existing shard tars are skipped; only missing shards are downloaded."""
+        entries = [
+            ("md5_0", "data-00000.arrow"),
+            ("md5_1", "data-00001.arrow"),
+            ("md5_2", "data-00002.arrow"),
+            ("md5_3", "data-00003.arrow"),
+        ]
+        mock_manifest.return_value = ("https://base.url", entries)
+        mock_convert.return_value = 100
+
+        with tempfile.TemporaryDirectory() as tmp:
+            output_train = Path(tmp) / "output" / "train"
+            output_train.mkdir(parents=True)
+
+            (output_train / "shard_000000.tar").write_bytes(b"existing")
+            (output_train / "shard_000002.tar").write_bytes(b"existing")
+
+            pipeline = StreamingIngestPipeline(
+                manifest_url="http://manifest.uri",
+                input_dir=f"{tmp}/arrows",
+                output_dir=f"{tmp}/output",
+                split_name="train",
+                max_workers=1,
+                prefetch_depth=2,
+            )
+
+            def fake_download(url, dest, **kwargs):
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(b"fake")
+
+            mock_download.side_effect = fake_download
+
+            results = pipeline.run()
+
+        assert results["files_processed"] == 2
+        assert results["shards_skipped"] == 2
+        assert results["shards_created"] == 4  # 2 new + 2 skipped
+        assert results["samples_written"] == 200  # only from the 2 new shards
+        assert mock_download.call_count == 2
+        assert mock_convert.call_count == 2
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py
new file mode 100644
index 000000000..f35d1a02b
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py
@@ -0,0 +1,164 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for preprocessing utility functions.
+
+Tests distributed processing utilities and work splitting logic.
+"""
+
+from unittest.mock import patch
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.preprocessing.utils import (
+    get_distributed_info,
+    split_work_for_rank,
+)
+from tests.utils import PrimusUT
+
+
+class TestGetDistributedInfo(PrimusUT):
+    """Tests for get_distributed_info utility."""
+
+    @patch("torch.distributed.is_initialized")
+    @patch("torch.distributed.is_available")
+    def test_not_distributed(self, mock_is_available, mock_is_initialized):
+        """Test get_distributed_info when not in distributed mode."""
+        mock_is_available.return_value = True
+        mock_is_initialized.return_value = False
+
+        rank, world_size = get_distributed_info()
+
+        assert rank == 0
+        assert world_size == 1
+
+    @patch("torch.distributed.get_world_size")
+    @patch("torch.distributed.get_rank")
+    @patch("torch.distributed.is_initialized")
+    @patch("torch.distributed.is_available")
+    def test_distributed_mode(
+        self, mock_is_available, mock_is_initialized, mock_get_rank, mock_get_world_size
+    ):
+        """Test get_distributed_info in distributed mode."""
+        mock_is_available.return_value = True
+        mock_is_initialized.return_value = True
+        mock_get_rank.return_value = 2
+        mock_get_world_size.return_value = 8
+
+        rank, world_size = get_distributed_info()
+
+        assert rank == 2
+        assert world_size == 8
+
+
+class TestSplitWorkForRank(PrimusUT):
+    """Tests for split_work_for_rank utility."""
+
+    def test_multiple_ranks_evenly_divisible(self):
+        """Test work splitting when items divide evenly."""
+        world_size = 4
+        total_items = 100
+
+        expected_splits = [
+            (0, 25),  # rank 0
+            (25, 50),  # rank 1
+            (50, 75),  # rank 2
+            (75, 100),  # rank 3 (last rank gets remainder)
+        ]
+
+        for rank, (expected_start, expected_end) in enumerate(expected_splits):
+            start, end = split_work_for_rank(total_items, rank, world_size)
+            assert start == expected_start
+            assert end == expected_end
+
+    def test_multiple_ranks_with_remainder(self):
+        """Test work splitting when items don't divide evenly."""
+        world_size = 3
+        total_items = 100
+
+        expected_splits = [
+            (0, 33),  # rank 0: 33 items
+            (33, 66),  # rank 1: 33 items
+            (66, 100),  # rank 2: 34 items (gets remainder)
+        ]
+
+        for rank, (expected_start, expected_end) in enumerate(expected_splits):
+            start, end = split_work_for_rank(total_items, rank, world_size)
+            assert start == expected_start
+            assert end == expected_end
+
+    def test_last_rank_gets_remainder(self):
+        """Test that last rank gets any remaining items."""
+        total_items = 107
+        world_size = 8
+        items_per_rank = total_items // world_size  # 13
+
+        # Last rank should get remainder
+        start, end = split_work_for_rank(total_items, rank=7, world_size=world_size)
+
+        assert start == 7 * items_per_rank  # 91
+        assert end == total_items  # 107 (gets all remaining)
+
+    def test_covers_all_items(self):
+        """Test that all ranks together cover all items."""
+        total_items = 137
+        world_size = 7
+
+        all_indices = set()
+        for rank in range(world_size):
+            start, end = split_work_for_rank(total_items, rank, world_size)
+            rank_indices = set(range(start, end))
+            # No overlap
+            assert len(all_indices & rank_indices) == 0
+            all_indices.update(rank_indices)
+
+        # All items covered
+        assert all_indices == set(range(total_items))
+
+    def test_fewer_items_than_ranks(self):
+        """Test work splitting when items < ranks."""
+        total_items = 5
+        world_size = 10
+
+        # First 5 ranks get 0 items each (5 // 10 = 0)
+        # Last rank gets all items
+        for rank in range(world_size - 1):
+            start, end = split_work_for_rank(total_items, rank, world_size)
+            # Each rank gets 0 items except last
+            assert start == 0
+            assert end == 0
+
+        # Last rank gets all
+        start, end = split_work_for_rank(total_items, rank=world_size - 1, world_size=world_size)
+        assert end == total_items
+
+
+class TestWorkSplittingEdgeCases(PrimusUT):
+    """Test edge cases for work splitting."""
+
+    def test_empty_work(self):
+        """Test with zero items."""
+        start, end = split_work_for_rank(total=0, rank=0, world_size=4)
+
+        assert start == 0
+        assert end == 0
+
+    def test_one_item(self):
+        """Test with single item."""
+        world_size = 4
+
+        # Ranks 0-2 get 0 items (1 // 4 = 0)
+        for rank in range(world_size - 1):
+            start, end = split_work_for_rank(1, rank, world_size)
+            assert start == 0
+            assert end == 0
+
+        # Last rank gets the 1 item
+        start, end = split_work_for_rank(1, rank=world_size - 1, world_size=world_size)
+        assert start == 0
+        assert end == 1
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py
new file mode 100644
index 000000000..f630163c0
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py
@@ -0,0 +1,194 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for Energon dataset validation module.
+
+Tests _check_metadata, _check_sample_counts, and validate_energon_dataset
+using synthetic dataset directories.
+"""
+
+import json
+import tarfile
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import yaml
+
+from primus.backends.megatron.data.diffusion.preprocessing.validate import (
+    _check_metadata,
+    _check_sample_counts,
+    validate_energon_dataset,
+)
+from tests.utils import PrimusUT
+
+
+def _build_valid_dataset(base: Path, num_samples: int = 3) -> None:
+    """Helper: build a minimal valid dataset directory on disk."""
+    meta_dir = base / ".nv-meta"
+    meta_dir.mkdir(parents=True)
+
+    shard_name = "000000.tar"
+    idx_name = shard_name + ".idx"
+
+    # Write .info.json
+    info = {"shard_counts": {shard_name: num_samples}}
+    (meta_dir / ".info.json").write_text(json.dumps(info))
+
+    # Write split.yaml
+    split = {"split_parts": {"train": [shard_name]}}
+    (meta_dir / "split.yaml").write_text(yaml.dump(split))
+
+    # Write dataset.yaml
+    (meta_dir / "dataset.yaml").write_text("__module__: megatron.energon\n__class__: CrudeWebdataset\n")
+
+    # Write shard tar with correct number of samples
+    tar_path = base / shard_name
+    with tarfile.open(str(tar_path), "w") as tar:
+        for i in range(num_samples):
+            key = f"000000_{i:06d}"
+            for ext in ("jpg", "txt"):
+                member_name = f"{key}.{ext}"
+                data = b"test"
+                info_obj = tarfile.TarInfo(name=member_name)
+                info_obj.size = len(data)
+                import io
+
+                tar.addfile(info_obj, io.BytesIO(data))
+
+    # Write .idx stub
+    (base / idx_name).write_text("")
+
+
+class TestCheckMetadata(PrimusUT):
+    """Tests for _check_metadata."""
+
+    def test_valid_metadata(self):
+        """Valid .nv-meta with all files returns info dict."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            info = _check_metadata(base)
+
+            assert info is not None
+            assert "shard_counts" in info
+            assert "_splits" in info
+
+    def test_missing_info_json(self):
+        """Missing .info.json returns None."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            (base / ".nv-meta" / ".info.json").unlink()
+
+            assert _check_metadata(base) is None
+
+    def test_missing_shard_counts_key(self):
+        """Info file without shard_counts key returns None."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            (base / ".nv-meta" / ".info.json").write_text(json.dumps({"other": 1}))
+
+            assert _check_metadata(base) is None
+
+    def test_missing_split_yaml(self):
+        """Missing split.yaml returns None."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            (base / ".nv-meta" / "split.yaml").unlink()
+
+            assert _check_metadata(base) is None
+
+    def test_missing_dataset_yaml(self):
+        """Missing dataset.yaml returns None."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            (base / ".nv-meta" / "dataset.yaml").unlink()
+
+            assert _check_metadata(base) is None
+
+    def test_missing_shard_file(self):
+        """Referenced shard file missing on disk returns None."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base)
+            (base / "000000.tar").unlink()
+
+            assert _check_metadata(base) is None
+
+
+class TestCheckSampleCounts(PrimusUT):
+    """Tests for _check_sample_counts."""
+
+    def test_matching_counts(self):
+        """Tar with correct sample count returns True."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base, num_samples=3)
+
+            info = _check_metadata(base)
+            assert info is not None
+            assert _check_sample_counts(base, info) is True
+
+    def test_mismatched_counts(self):
+        """Mismatch between .info.json and tar content returns False."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            _build_valid_dataset(base, num_samples=3)
+
+            # Overwrite .info.json with wrong count
+            wrong_info = {"shard_counts": {"000000.tar": 99}}
+            (base / ".nv-meta" / ".info.json").write_text(json.dumps(wrong_info))
+
+            info = _check_metadata(base)
+            assert info is not None
+            assert _check_sample_counts(base, info) is False
+
+
+class TestValidateEnergonDataset(PrimusUT):
+    """Tests for validate_energon_dataset orchestration."""
+
+    @patch(
+        "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load",
+        return_value={"jpg": "bytes, len=100", "txt": "str"},
+    )
+    def test_full_pass(self, mock_load):
+        """Complete valid dataset returns True."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _build_valid_dataset(Path(tmpdir), num_samples=3)
+            assert validate_energon_dataset(tmpdir, encoding="raw") is True
+
+    @patch(
+        "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load",
+        return_value={"jpg": "bytes, len=100", "txt": "str"},
+    )
+    def test_metadata_failure_returns_false(self, mock_load):
+        """Metadata failure returns False immediately (early exit)."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Empty dir has no .nv-meta at all
+            assert validate_energon_dataset(tmpdir, encoding="raw") is False
+            mock_load.assert_not_called()
+
+    @patch(
+        "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load",
+        return_value={"jpg": "bytes, len=100", "txt": "str"},
+    )
+    def test_count_mismatch_returns_false(self, mock_load):
+        """Count mismatch returns False but still prints summary."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _build_valid_dataset(Path(tmpdir), num_samples=3)
+
+            # Corrupt the count
+            wrong_info = {"shard_counts": {"000000.tar": 99}}
+            (Path(tmpdir) / ".nv-meta" / ".info.json").write_text(json.dumps(wrong_info))
+
+            assert validate_energon_dataset(tmpdir, encoding="raw") is False
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py
new file mode 100644
index 000000000..89778402a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py
new file mode 100644
index 000000000..2bac33b3f
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py
@@ -0,0 +1,192 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for diffusion TaskEncoders.
+
+Tests both EncodedDiffusionTaskEncoder (pre-encoded data) and
+RawDiffusionTaskEncoder (raw images and text).
+"""
+
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.task_encoders import (
+    DiffusionSample,
+    EncodedDiffusionTaskEncoder,
+    RawDiffusionTaskEncoder,
+    cook_preencoded_diffusion,
+    cook_raw_images,
+)
+
+
+class TestCookPreencodedDiffusion:
+    """Tests for cook_preencoded_diffusion cooker function."""
+
+    def test_cook_with_tensor_data(self, preencoded_sample):
+        """Test cooking with direct tensor data."""
+        result = cook_preencoded_diffusion(preencoded_sample)
+
+        assert isinstance(result, DiffusionSample)
+        assert result.latents.shape == (16, 64, 64)
+        assert result.prompt_embeds.shape == (256, 4096)
+        assert result.pooled_prompt_embeds.shape == (768,)
+        assert result.caption == "A test caption"
+
+    def test_cook_with_bytes_data(self, preencoded_sample_bytes):
+        """Test cooking with tensor data as bytes."""
+        result = cook_preencoded_diffusion(preencoded_sample_bytes)
+
+        assert isinstance(result, DiffusionSample)
+        assert result.latents.shape == (16, 64, 64)
+        assert result.prompt_embeds.shape == (256, 4096)
+        assert result.pooled_prompt_embeds.shape == (768,)
+        assert result.caption == "Another test caption"
+
+    def test_cook_ignores_text_ids(self, preencoded_sample, sample_text_ids):
+        """Test that text_ids in sample are ignored (position IDs generated at runtime)."""
+        # Even if text_ids are provided in the sample, they should not be stored
+        # because position IDs are generated at runtime based on actual tensor shapes
+        preencoded_sample["text_ids.pth"] = sample_text_ids
+        result = cook_preencoded_diffusion(preencoded_sample)
+
+        # Verify that text_ids attribute doesn't exist (not part of DiffusionSample)
+        assert not hasattr(result, "text_ids"), "text_ids should not be stored in DiffusionSample"
+
+    def test_cook_missing_required_key_latents(self, preencoded_sample):
+        """Test error when latents are missing."""
+        del preencoded_sample["latents.pth"]
+
+        # Production now emits a distinct message for missing latents (vs
+        # missing prompt_embeds / pooled_prompt_embeds) since 'latents.pth' is
+        # interchangeable with 'mean.pth' / 'logvar.pth' for resample-mode.
+        with pytest.raises(
+            ValueError,
+            match=r"Sample must have 'latents\.pth' or 'mean\.pth'/'logvar\.pth'",
+        ):
+            cook_preencoded_diffusion(preencoded_sample)
+
+    def test_cook_missing_required_key_prompt_embeds(self, preencoded_sample):
+        """Test error when prompt_embeds are missing."""
+        del preencoded_sample["prompt_embeds.pth"]
+
+        with pytest.raises(ValueError, match="Sample missing required keys"):
+            cook_preencoded_diffusion(preencoded_sample)
+
+    def test_cook_missing_required_key_pooled(self, preencoded_sample):
+        """Test error when pooled_prompt_embeds are missing."""
+        del preencoded_sample["pooled_prompt_embeds.pth"]
+
+        with pytest.raises(ValueError, match="Sample missing required keys"):
+            cook_preencoded_diffusion(preencoded_sample)
+
+    def test_cook_caption_as_str(self, preencoded_sample):
+        """Test caption handling when it's a string."""
+        preencoded_sample["caption.txt"] = "String caption"
+        result = cook_preencoded_diffusion(preencoded_sample)
+
+        assert result.caption == "String caption"
+
+    def test_cook_caption_missing(self, preencoded_sample):
+        """Test caption handling when missing."""
+        del preencoded_sample["caption.txt"]
+        result = cook_preencoded_diffusion(preencoded_sample)
+
+        assert result.caption == ""
+
+    def test_cook_preserves_sample_keys(self, preencoded_sample):
+        """Test that sample metadata keys are preserved."""
+        result = cook_preencoded_diffusion(preencoded_sample)
+
+        assert result.__key__ == "sample_001"
+        assert result.__subflavors__ == {"encoding": "preencoded"}
+
+
+class TestCookRawImages:
+    """Tests for cook_raw_images cooker function."""
+
+    def test_cook_with_png(self, sample_image_bytes):
+        """Test cooking raw PNG image."""
+        sample = {
+            "__key__": "test_png",
+            "__restore_key__": lambda: "test_png",
+            "__subflavors__": {"encoding": "raw"},
+            "images": sample_image_bytes,
+            "txt": b"PNG caption",
+        }
+        result = cook_raw_images(sample)
+
+        assert result["images"] == sample_image_bytes
+        assert result["txt"] == b"PNG caption"
+
+    def test_cook_preserves_metadata(self, raw_sample):
+        """Test that metadata keys are preserved."""
+        result = cook_raw_images(raw_sample)
+
+        assert result["__key__"] == "sample_003"
+        assert result["__subflavors__"] == {"encoding": "raw"}
+
+
+class TestEncodedDiffusionTaskEncoder:
+    """Tests for EncodedDiffusionTaskEncoder."""
+
+    def test_batch_with_all_fields(self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds):
+        """Test batching samples with all standard fields (position IDs generated at runtime)."""
+        encoder = EncodedDiffusionTaskEncoder(worker_config=None)
+
+        samples = [
+            DiffusionSample(
+                __key__="s1",
+                __restore_key__=lambda: "s1",
+                __subflavors__={"encoding": "preencoded"},
+                latents=sample_latents.squeeze(0),
+                prompt_embeds=sample_prompt_embeds.squeeze(0),
+                pooled_prompt_embeds=sample_pooled_prompt_embeds.squeeze(0),
+                caption="Caption 1",
+            ),
+            DiffusionSample(
+                __key__="s2",
+                __restore_key__=lambda: "s2",
+                __subflavors__={"encoding": "preencoded"},
+                latents=sample_latents.squeeze(0),
+                prompt_embeds=sample_prompt_embeds.squeeze(0),
+                pooled_prompt_embeds=sample_pooled_prompt_embeds.squeeze(0),
+                caption="Caption 2",
+            ),
+        ]
+
+        batch = encoder.batch(samples)
+
+        assert "latents" in batch
+        assert "prompt_embeds" in batch
+        assert "pooled_prompt_embeds" in batch
+        # text_ids not in batch - position IDs are generated at runtime
+        # Note: 'captions' may or may not be in batch depending on implementation
+
+        assert batch["latents"].shape[0] == 2
+        assert batch["prompt_embeds"].shape[0] == 2
+        assert batch["pooled_prompt_embeds"].shape[0] == 2
+
+
+class TestRawDiffusionTaskEncoder:
+    """Tests for RawDiffusionTaskEncoder."""
+
+    def test_batch_with_raw_data(self, sample_image_bytes):
+        """Test batching raw image samples."""
+        encoder = RawDiffusionTaskEncoder(worker_config=None)
+
+        samples = [
+            {"images": sample_image_bytes, "txt": b"Caption 1"},
+            {"images": sample_image_bytes, "txt": b"Caption 2"},
+        ]
+
+        batch = encoder.batch(samples)
+
+        assert "images" in batch
+        assert "txt" in batch
+        assert len(batch["images"]) == 2
+        assert len(batch["txt"]) == 2
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py b/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py
new file mode 100644
index 000000000..62698df6b
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py
@@ -0,0 +1,129 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Schema/shape/reproducibility contract tests for the synthetic Flux mock datasets.
+
+These guard the sample contract (keys, tensor shapes, seed reproducibility) that
+the Flux forward step depends on for mock-data training runs. Relocated here from
+tests/integration_tests/... because they are CPU-only unit tests of the synthetic
+dataset providers and do not exercise an end-to-end TaskEncoder/dataloader path.
+"""
+
+import pytest
+import torch
+
+from primus.backends.megatron.data.synthetic import (
+    MockFluxDataset,
+    MockFluxSchnellDataset,
+    PreGeneratedMockFluxSchnellDataset,
+)
+from primus.backends.megatron.data.synthetic_dataset_provider import (
+    SyntheticDatasetProvider,
+)
+from tests.unit_tests.backends.megatron.diffusion.helpers import assert_tensor_shape
+from tests.utils import PrimusUT
+
+
+class TestMockFluxDataset(PrimusUT):
+    """Contract tests for MockFluxDataset (Flux dev)."""
+
+    def test_mock_dataset_sample_structure(self):
+        dataset = MockFluxDataset(num_samples=10, seed=42)
+        sample = dataset[0]
+
+        expected_keys = {
+            "latents",
+            "prompt_embeds",
+            "pooled_prompt_embeds",
+            "img_ids",
+            "txt_ids",
+            "caption",
+        }
+        assert set(sample.keys()) == expected_keys
+
+    def test_mock_dataset_sample_shapes(self):
+        dataset = MockFluxDataset(num_samples=10, image_size=1024, seed=42)
+        sample = dataset[0]
+
+        # Latents: (C, H, W) -- 1024/8 = 128
+        assert_tensor_shape(sample["latents"], (16, 128, 128), "latents")
+        # T5 embeddings: (S, D)
+        assert_tensor_shape(sample["prompt_embeds"], (512, 4096), "t5_embeddings")
+        # CLIP pooled: (D,)
+        assert_tensor_shape(sample["pooled_prompt_embeds"], (768,), "clip_pooled")
+        # Image IDs: (N, 3) where N = (H/2) * (W/2)
+        expected_img_ids_len = (128 // 2) * (128 // 2)
+        assert_tensor_shape(sample["img_ids"], (expected_img_ids_len, 3), "img_ids")
+        # Text IDs: (S, 3)
+        assert_tensor_shape(sample["txt_ids"], (512, 3), "txt_ids")
+
+    def test_mock_dataset_reproducibility(self):
+        dataset1 = MockFluxDataset(num_samples=10, seed=42)
+        dataset2 = MockFluxDataset(num_samples=10, seed=42)
+
+        sample1 = dataset1[0]
+        sample2 = dataset2[0]
+
+        assert torch.allclose(sample1["latents"], sample2["latents"])
+        assert torch.allclose(sample1["prompt_embeds"], sample2["prompt_embeds"])
+        assert torch.allclose(sample1["pooled_prompt_embeds"], sample2["pooled_prompt_embeds"])
+
+    def test_position_ids_format(self):
+        dataset = MockFluxDataset(num_samples=5, seed=42)
+        sample = dataset[0]
+
+        img_ids = sample["img_ids"]
+        txt_ids = sample["txt_ids"]
+
+        # 3D RoPE format for both image and text position IDs.
+        assert img_ids.shape[-1] == 3, "Image IDs should have 3 dimensions for RoPE"
+        assert txt_ids.shape[-1] == 3, "Text IDs should have 3 dimensions for RoPE"
+        # Text IDs are all zeros.
+        assert torch.allclose(txt_ids, torch.zeros_like(txt_ids)), "Text IDs should be all zeros"
+
+    def test_latent_size_scales_with_image_size(self):
+        # Latent spatial size must be image_size/8 (VAE downsample contract).
+        for image_size, latent in ((512, 64), (2048, 256)):
+            dataset = MockFluxDataset(num_samples=5, image_size=image_size, seed=42)
+            sample = dataset[0]
+            assert_tensor_shape(sample["latents"], (16, latent, latent), "latents")
+            expected_img_ids_len = (latent // 2) * (latent // 2)
+            assert_tensor_shape(sample["img_ids"], (expected_img_ids_len, 3), "img_ids")
+
+
+class TestMockFluxSchnellDataset(PrimusUT):
+    """Contract tests for MockFluxSchnellDataset (Flux Schnell / MLPerf v5.1)."""
+
+    def test_schnell_sample_shapes(self):
+        # Schnell uses T5 seq_len=256 (vs 512 for dev).
+        dataset = MockFluxSchnellDataset(num_samples=10, image_size=256, seed=42)
+        sample = dataset[0]
+
+        assert_tensor_shape(sample["latents"], (16, 32, 32), "latents")
+        assert_tensor_shape(sample["prompt_embeds"], (256, 4096), "prompt_embeds")
+        assert_tensor_shape(sample["pooled_prompt_embeds"], (768,), "pooled_prompt_embeds")
+        assert_tensor_shape(sample["txt_ids"], (256, 3), "txt_ids")
+
+    def test_schnell_pregenerated_dataset(self):
+        dataset = PreGeneratedMockFluxSchnellDataset(num_samples=10, image_size=256, seed=42)
+
+        sample0_a = dataset[0]
+        sample0_b = dataset[0]
+
+        assert torch.allclose(sample0_a["latents"], sample0_b["latents"])
+        assert torch.allclose(sample0_a["prompt_embeds"], sample0_b["prompt_embeds"])
+
+
+class TestSyntheticDatasetProviderLookup(PrimusUT):
+    """Tests for SyntheticDatasetProvider DEFAULT_DATASETS lookup."""
+
+    def test_schnell_class_resolves(self):
+        provider = SyntheticDatasetProvider(model_type="flux_schnell")
+        cls = provider._import_dataset_class()
+        assert cls is not None
+        assert cls.__name__ == "PreGeneratedMockFluxSchnellDataset"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py b/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py
new file mode 100644
index 000000000..b22d06ccf
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Distributed tests for Flux diffusion model.
+
+Tests for FSDP2, pipeline parallelism, and other distributed training features.
+"""
diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py
new file mode 100644
index 000000000..c313644ea
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py
@@ -0,0 +1,528 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FP8 all-gather tensor subclasses.
+
+Tests the FP8 all-gather subclasses in isolation (single GPU, no
+multi-GPU / NCCL required). Verifies:
+- WeightWithFP8AllGatherTensor: creation, pre/post all-gather with
+  precomputed scale, dispatch, flatten/unflatten
+- FP8UnshardedWeightTensor: creation, dispatch propagation, detach
+  dequantizes to BF16, unsupported ops raise NotImplementedError,
+  flatten/unflatten roundtrip, get_fp8_data_and_scale_inv()
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus_turbo.pytorch.core.low_precision import (
+    Float8QuantConfig,
+    Format,
+    ScalingGranularity,
+)
+
+from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+    FP8UnshardedWeightTensor,
+    WeightWithFP8AllGatherTensor,
+    _foreach_fp8_quantize,
+    _get_fp8_dtype,
+    quantize_fp8_prescaled,
+)
+
+_has_cuda = torch.cuda.is_available()
+requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required")
+
+
+def _make_config():
+    return Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE)
+
+
+def _make_wrapped(shape=(256, 512), device="cpu"):
+    tensor = torch.randn(shape, dtype=torch.bfloat16, device=device)
+    wrapped = WeightWithFP8AllGatherTensor(tensor, _make_config())
+    return wrapped, tensor
+
+
+def _set_precomputed_scale(wrapped):
+    """Compute and set precomputed scale, scale_inv, and cached FP8 data for a single-rank scenario."""
+    fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format)
+    fp8_max = torch.finfo(fp8_dtype).max
+    amax = wrapped._tensor.abs().amax().float()
+    scale = fp8_max / amax.clamp(min=1e-12)
+    wrapped._precomputed_scale = scale
+    wrapped._precomputed_scale_inv = 1.0 / scale
+    fp8_data, _ = quantize_fp8_prescaled(
+        wrapped._tensor,
+        fp8_dtype,
+        scale,
+        wrapped._precomputed_scale_inv,
+    )
+    wrapped._cached_fp8_data = fp8_data
+
+
+def _make_fp8_unsharded(shape=(256, 512), device="cpu"):
+    """Create an FP8UnshardedWeightTensor from a BF16 source tensor."""
+    config = _make_config()
+    bf16_tensor = torch.randn(shape, dtype=torch.bfloat16, device=device)
+    fp8_dtype = _get_fp8_dtype(config.format)
+    fp8_max = torch.finfo(fp8_dtype).max
+    amax = bf16_tensor.abs().amax().float()
+    scale = fp8_max / amax.clamp(min=1e-12)
+    scale_inv = 1.0 / scale
+    fp8_data = (bf16_tensor.float() * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
+    return (
+        FP8UnshardedWeightTensor(fp8_data, scale_inv, torch.bfloat16, config),
+        bf16_tensor,
+        fp8_data,
+        scale_inv,
+    )
+
+
+class TestSubclassCreation:
+    def test_shape_dtype_device_preserved(self):
+        wrapped, orig = _make_wrapped()
+        assert wrapped.shape == orig.shape
+        assert wrapped.dtype == torch.bfloat16
+        assert wrapped.device == orig.device
+
+    def test_isinstance_tensor(self):
+        wrapped, _ = _make_wrapped()
+        assert isinstance(wrapped, torch.Tensor)
+        assert isinstance(wrapped, WeightWithFP8AllGatherTensor)
+
+    def test_inner_tensor_accessible(self):
+        wrapped, orig = _make_wrapped()
+        assert wrapped._tensor is orig
+
+    def test_tensorwise_only_validation(self):
+        config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.ROWWISE)
+        with pytest.raises(ValueError, match="TENSORWISE"):
+            WeightWithFP8AllGatherTensor(torch.randn(4, 4, dtype=torch.bfloat16), config)
+
+    def test_precomputed_scale_initially_none(self):
+        wrapped, _ = _make_wrapped()
+        assert wrapped._precomputed_scale is None
+
+
+@requires_cuda
+class TestPreAllGather:
+    def test_returns_one_input(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        all_gather_inputs, metadata = wrapped.fsdp_pre_all_gather(None)
+        assert len(all_gather_inputs) == 1
+
+    def test_fp8_data_dtype(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        (fp8_data,), _ = wrapped.fsdp_pre_all_gather(None)
+        assert fp8_data.dtype in (
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+        )
+
+    def test_metadata_contains_scale_inv_and_numel(self):
+        wrapped, orig = _make_wrapped((128, 64), device="cuda")
+        _set_precomputed_scale(wrapped)
+        _, metadata = wrapped.fsdp_pre_all_gather(None)
+        scale_inv, shard_numel = metadata
+        assert scale_inv.dtype == torch.float32
+        assert shard_numel == 128 * 64
+
+    def test_fp8_data_shape_matches_input(self):
+        wrapped, orig = _make_wrapped((32, 64), device="cuda")
+        _set_precomputed_scale(wrapped)
+        (fp8_data,), _ = wrapped.fsdp_pre_all_gather(None)
+        assert fp8_data.numel() == orig.numel()
+
+    def test_asserts_without_precomputed_scale(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        with pytest.raises(RuntimeError, match="precompute_fp8_scales_for_fsdp"):
+            wrapped.fsdp_pre_all_gather(None)
+
+
+@requires_cuda
+class TestPostAllGather:
+    def _simulate_single_rank(self, wrapped):
+        """Simulate all-gather with a single rank (data passes through)."""
+        _set_precomputed_scale(wrapped)
+        (fp8_data,), metadata = wrapped.fsdp_pre_all_gather(None)
+        return (fp8_data,), metadata
+
+    def test_first_call_returns_fp8_unsharded_tensor(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        outputs, metadata = self._simulate_single_rank(wrapped)
+        result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16)
+        assert result is not None
+        tensor, inner_tensors = result
+        assert isinstance(tensor, FP8UnshardedWeightTensor)
+        assert tensor.dtype == torch.bfloat16
+        assert isinstance(inner_tensors, tuple)
+        assert len(inner_tensors) == 1
+
+    def test_steady_state_updates_scale(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        outputs, metadata = self._simulate_single_rank(wrapped)
+        # First call to get an FP8UnshardedWeightTensor
+        result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16)
+        fp8_unsharded, _ = result
+        # Simulate FSDP's reshard path: pass the FP8UnshardedWeightTensor
+        # directly as `out` (FSDP uses the unsharded_param, not nn.Parameter)
+        new_scale_inv = metadata[0] * 2.0
+        new_metadata = (new_scale_inv, metadata[1])
+        result = wrapped.fsdp_post_all_gather(outputs, new_metadata, torch.bfloat16, out=fp8_unsharded)
+        assert result is None
+        assert torch.equal(fp8_unsharded._scale_inv, new_scale_inv)
+
+
+@requires_cuda
+class TestRoundtripAccuracy:
+    def test_quantize_keeps_values(self):
+        torch.manual_seed(42)
+        wrapped, orig = _make_wrapped((128, 256), device="cuda")
+        _set_precomputed_scale(wrapped)
+        outputs, metadata = wrapped.fsdp_pre_all_gather(None)
+        result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16)
+        fp8_unsharded, _ = result
+        fp8_data, scale_inv = fp8_unsharded.get_fp8_data_and_scale_inv()
+        dequantized = fp8_data.to(torch.bfloat16) * scale_inv
+        dequantized = dequantized.view(orig.shape)
+        torch.testing.assert_close(dequantized, orig, rtol=0.05, atol=0.1)
+
+
+class TestTorchDispatch:
+    def test_preserves_subclass_for_standard_ops(self):
+        wrapped, _ = _make_wrapped((16, 16))
+        for op in [
+            lambda t: torch.empty_like(t),
+            lambda t: t.view(-1),
+            lambda t: t.clone(),
+            lambda t: t[0:8],
+        ]:
+            result = op(wrapped)
+            assert isinstance(
+                result, WeightWithFP8AllGatherTensor
+            ), f"Op should preserve subclass but got {type(result)}"
+
+    def test_unwraps_for_compute_ops(self):
+        wrapped, _ = _make_wrapped((16, 16))
+        plain = torch.randn(16, 16, dtype=torch.bfloat16)
+        result = wrapped + plain
+        assert type(result) is torch.Tensor
+        assert not isinstance(result, WeightWithFP8AllGatherTensor)
+
+    def test_copy_inplace_modifies_inner(self):
+        """copy_ modifies inner tensor in-place; target retains subclass type."""
+        wrapped, _ = _make_wrapped((8, 8))
+        target = torch.empty_like(wrapped)
+        target.copy_(wrapped)
+        assert isinstance(target, WeightWithFP8AllGatherTensor)
+        torch.testing.assert_close(target._tensor, wrapped._tensor)
+
+    def test_to_copy_dtype_change_unwraps(self):
+        """`.float()` should return a plain tensor (for FP32 master copies)."""
+        wrapped, _ = _make_wrapped((8, 8))
+        fp32 = wrapped.float()
+        assert fp32.dtype == torch.float32
+        assert not isinstance(fp32, WeightWithFP8AllGatherTensor)
+
+    def test_to_copy_same_dtype_preserves(self):
+        wrapped, _ = _make_wrapped((8, 8))
+        result = wrapped.to(torch.bfloat16)
+        assert isinstance(result, WeightWithFP8AllGatherTensor)
+
+    def test_detach_preserves_subclass(self):
+        """detach preserves subclass so nn.Parameter() and FSDP2 work correctly."""
+        wrapped, orig = _make_wrapped((8, 8))
+        detached = wrapped.detach()
+        assert isinstance(detached, WeightWithFP8AllGatherTensor)
+        assert detached._tensor.data_ptr() == wrapped._tensor.data_ptr()
+        assert detached._fp8_config.format == wrapped._fp8_config.format
+
+
+class TestTensorFlattenUnflatten:
+    def test_roundtrip(self):
+        wrapped, orig = _make_wrapped((32, 64))
+        inner_tensors_names, metadata = wrapped.__tensor_flatten__()
+        assert inner_tensors_names == ["_tensor"]
+        assert "format" in metadata
+        assert "granularity" in metadata
+
+        inner_tensors = {"_tensor": wrapped._tensor}
+        restored = WeightWithFP8AllGatherTensor.__tensor_unflatten__(
+            inner_tensors, metadata, wrapped.shape, wrapped.stride()
+        )
+        assert isinstance(restored, WeightWithFP8AllGatherTensor)
+        assert restored._tensor is orig
+        assert restored._fp8_config.format == Format.E4M3
+        assert restored._fp8_config.granularity == ScalingGranularity.TENSORWISE
+
+    def test_metadata_is_hashable(self):
+        wrapped, _ = _make_wrapped()
+        _, metadata = wrapped.__tensor_flatten__()
+        for key, value in metadata.items():
+            hash(value)
+
+    def test_precomputed_scale_not_in_metadata(self):
+        """_precomputed_scale is transient and must not appear in flatten metadata."""
+        wrapped, _ = _make_wrapped()
+        wrapped._precomputed_scale = torch.tensor(1.0)
+        _, metadata = wrapped.__tensor_flatten__()
+        assert "precomputed_scale" not in metadata
+        assert "_precomputed_scale" not in metadata
+
+
+# ============================================================================
+# FP8UnshardedWeightTensor tests
+# ============================================================================
+
+
+class TestFP8UnshardedWeightTensorCreation:
+    def test_shape_and_declared_dtype(self):
+        fp8_unsharded, _, fp8_data, _ = _make_fp8_unsharded((64, 32))
+        assert fp8_unsharded.shape == (64, 32)
+        assert fp8_unsharded.dtype == torch.bfloat16
+
+    def test_inner_fp8_data_accessible(self):
+        fp8_unsharded, _, fp8_data, scale_inv = _make_fp8_unsharded()
+        assert fp8_unsharded._fp8_data is fp8_data
+        assert torch.equal(fp8_unsharded._scale_inv, scale_inv)
+
+    def test_get_fp8_data_and_scale_inv(self):
+        fp8_unsharded, _, fp8_data, scale_inv = _make_fp8_unsharded()
+        data, sinv = fp8_unsharded.get_fp8_data_and_scale_inv()
+        assert data is fp8_data
+        assert torch.equal(sinv, scale_inv)
+
+
+class TestFP8UnshardedDispatch:
+    def test_as_strided_preserves_subclass(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8))
+        result = torch.as_strided(fp8_unsharded, (64,), (1,), 0)
+        assert isinstance(result, FP8UnshardedWeightTensor)
+        assert result.shape == (64,)
+
+    def test_view_preserves_subclass(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8))
+        result = fp8_unsharded.view(-1)
+        assert isinstance(result, FP8UnshardedWeightTensor)
+        assert result.shape == (64,)
+
+    def test_slice_preserves_subclass(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((16, 8))
+        result = fp8_unsharded[0:8]
+        assert isinstance(result, FP8UnshardedWeightTensor)
+        assert result.shape == (8, 8)
+
+    def test_clone_preserves_subclass(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8))
+        result = fp8_unsharded.clone()
+        assert isinstance(result, FP8UnshardedWeightTensor)
+
+    def test_detach_preserves_subclass(self):
+        """detach must preserve subclass (required by nn.Parameter wrapping)."""
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8))
+        detached = fp8_unsharded.detach()
+        assert isinstance(detached, FP8UnshardedWeightTensor)
+        assert detached.dtype == torch.bfloat16
+
+    def test_unsupported_op_raises(self):
+        """Unsupported ops raise NotImplementedError (safety guard)."""
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8))
+        with pytest.raises(NotImplementedError, match="does not support"):
+            fp8_unsharded + torch.zeros(8, 8, dtype=torch.bfloat16)
+
+
+class TestFP8UnshardedFlattenUnflatten:
+    def test_roundtrip(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded((32, 64))
+        inner_names, metadata = fp8_unsharded.__tensor_flatten__()
+        assert set(inner_names) == {"_fp8_data", "_scale_inv"}
+        assert "orig_dtype" in metadata
+        assert "format" in metadata
+
+        inner_tensors = {
+            "_fp8_data": fp8_unsharded._fp8_data,
+            "_scale_inv": fp8_unsharded._scale_inv,
+        }
+        restored = FP8UnshardedWeightTensor.__tensor_unflatten__(
+            inner_tensors, metadata, fp8_unsharded.shape, fp8_unsharded.stride()
+        )
+        assert isinstance(restored, FP8UnshardedWeightTensor)
+        assert restored._orig_dtype == torch.bfloat16
+        # FP8 dtypes don't support torch.equal on CPU; compare via uint8 view
+        assert torch.equal(
+            restored._fp8_data.view(torch.uint8),
+            fp8_unsharded._fp8_data.view(torch.uint8),
+        )
+        assert torch.equal(restored._scale_inv, fp8_unsharded._scale_inv)
+
+    def test_metadata_is_hashable(self):
+        fp8_unsharded, _, _, _ = _make_fp8_unsharded()
+        _, metadata = fp8_unsharded.__tensor_flatten__()
+        for key, value in metadata.items():
+            hash(value)
+
+
+class TestFP8DataCache:
+    def test_cache_initially_none(self):
+        wrapped, _ = _make_wrapped()
+        assert wrapped._cached_fp8_data is None
+
+    @requires_cuda
+    def test_pre_all_gather_returns_cached_data(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        assert wrapped._cached_fp8_data is not None
+        (fp8_data,), _ = wrapped.fsdp_pre_all_gather(mesh=None)
+        assert fp8_data.data_ptr() == wrapped._cached_fp8_data.data_ptr()
+
+    @requires_cuda
+    def test_pre_all_gather_fallback_without_cache(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        wrapped._cached_fp8_data = None
+        (fp8_data,), (scale_inv, _) = wrapped.fsdp_pre_all_gather(mesh=None)
+        fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format)
+        assert fp8_data.dtype == fp8_dtype
+        assert fp8_data.shape == wrapped._tensor.shape
+
+    @requires_cuda
+    def test_cache_matches_direct_quantize(self):
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        cached = wrapped._cached_fp8_data
+        fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format)
+        fresh, _ = quantize_fp8_prescaled(
+            wrapped._tensor,
+            fp8_dtype,
+            wrapped._precomputed_scale,
+            wrapped._precomputed_scale_inv,
+        )
+        assert torch.equal(cached.view(torch.uint8), fresh.view(torch.uint8))
+
+    @requires_cuda
+    def test_cache_not_in_flatten_metadata(self):
+        # _set_precomputed_scale runs the Triton FP8 quantize kernel which
+        # requires a CUDA-resident tensor (no CPU dispatch). Match the other
+        # cache tests in this class which already use device="cuda".
+        wrapped, _ = _make_wrapped(device="cuda")
+        _set_precomputed_scale(wrapped)
+        inner_names, _ = wrapped.__tensor_flatten__()
+        assert "_cached_fp8_data" not in inner_names
+
+
+class TestForeachOptimizerOps:
+    """Tests for _foreach_copy_ batched operations used by the optimizer."""
+
+    @requires_cuda
+    def test_foreach_copy_bf16_to_fp32_matches_float(self):
+        """_foreach_copy_ BF16->FP32 matches per-tensor .float() (prepare_grads)."""
+        sizes = [(128, 256), (64,), (512, 512), (32, 16)]
+        bf16_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in sizes]
+        ref_fp32 = [t.float() for t in bf16_tensors]
+
+        fp32_targets = [torch.empty(s, dtype=torch.float32, device="cuda") for s in sizes]
+        torch._foreach_copy_(fp32_targets, bf16_tensors)
+
+        for i in range(len(sizes)):
+            torch.testing.assert_close(
+                fp32_targets[i], ref_fp32[i], rtol=0, atol=0, msg=f"Mismatch at tensor {i} shape {sizes[i]}"
+            )
+
+    @requires_cuda
+    def test_foreach_copy_fp32_to_bf16_matches_copy(self):
+        """_foreach_copy_ FP32->BF16 matches per-tensor .copy_() (copy-back)."""
+        sizes = [(256, 128), (64,), (1024,)]
+        fp32_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in sizes]
+        ref_bf16 = [torch.empty(s, dtype=torch.bfloat16, device="cuda") for s in sizes]
+        for r, f in zip(ref_bf16, fp32_tensors):
+            r.copy_(f)
+
+        bf16_targets = [torch.empty(s, dtype=torch.bfloat16, device="cuda") for s in sizes]
+        torch._foreach_copy_(bf16_targets, fp32_tensors)
+
+        for i in range(len(sizes)):
+            torch.testing.assert_close(
+                bf16_targets[i], ref_bf16[i], rtol=0, atol=0, msg=f"Mismatch at tensor {i} shape {sizes[i]}"
+            )
+
+    @requires_cuda
+    def test_cache_extraction_fp8_wrapped(self):
+        """to_local() + inner_data() produces the correct inner tensor for FP8-wrapped params."""
+        wrapped, orig = _make_wrapped((128, 64), device="cuda")
+        inner = wrapped.inner_data()
+        assert inner is orig
+        assert inner.dtype == torch.bfloat16
+        assert inner.data_ptr() == orig.data_ptr()
+
+    @requires_cuda
+    def test_mixed_sizes_foreach_copy(self):
+        """_foreach_copy_ works with mixed tensor sizes in the same list."""
+        sizes = [(1,), (7,), (1024, 1024), (3, 5, 7), (131072,)]
+        bf16_list = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in sizes]
+        fp32_list = [torch.empty(s, dtype=torch.float32, device="cuda") for s in sizes]
+        torch._foreach_copy_(fp32_list, bf16_list)
+        for i, (b, f) in enumerate(zip(bf16_list, fp32_list)):
+            torch.testing.assert_close(
+                f, b.float(), rtol=0, atol=0, msg=f"Mixed-size mismatch at tensor {i} shape {sizes[i]}"
+            )
+
+
+class TestForeachFP8Quantize:
+    """Tests for the compiled _foreach_fp8_quantize batch quantization helper."""
+
+    @requires_cuda
+    def test_matches_per_tensor_quantize(self):
+        """Verify _foreach_fp8_quantize is bitwise identical to quantize_fp8_prescaled."""
+        config = _make_config()
+        fp8_dtype = _get_fp8_dtype(config.format)
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        sizes = [1024, 4096, 131072, 524288]
+        tensors = [torch.randn(s, device="cuda", dtype=torch.bfloat16) for s in sizes]
+        amaxes = torch.stack([t.abs().amax().float() for t in tensors])
+        scales = (fp8_max / amaxes.clamp(min=1e-12)).to(torch.float32)
+        scale_invs = 1.0 / scales
+
+        ref_results = []
+        for i, t in enumerate(tensors):
+            fp8_data, _ = quantize_fp8_prescaled(
+                t,
+                fp8_dtype,
+                scales[i],
+                scale_invs[i],
+            )
+            ref_results.append(fp8_data)
+
+        fp8_outputs = [torch.empty_like(t, dtype=fp8_dtype) for t in tensors]
+        _foreach_fp8_quantize.__wrapped__(tensors, fp8_outputs, scales, fp8_max)
+
+        for i in range(len(tensors)):
+            assert torch.equal(
+                fp8_outputs[i].view(torch.uint8),
+                ref_results[i].view(torch.uint8),
+            ), f"Mismatch at tensor {i} (size {sizes[i]})"
+
+    @requires_cuda
+    def test_clamps_correctly(self):
+        """Verify out-of-range values are clamped to fp8_max, not NaN."""
+        config = _make_config()
+        fp8_dtype = _get_fp8_dtype(config.format)
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        tensor = torch.tensor([1.0, 1000.0, -1000.0, 0.001], device="cuda", dtype=torch.bfloat16)
+        scale = torch.tensor(fp8_max, device="cuda", dtype=torch.float32)
+        scales = scale.unsqueeze(0)
+
+        fp8_out = [torch.empty_like(tensor, dtype=fp8_dtype)]
+        _foreach_fp8_quantize.__wrapped__([tensor], fp8_out, scales, fp8_max)
+
+        result = fp8_out[0]
+        assert not result.float().isnan().any(), "NaN found in FP8 output"
+        assert result.float().abs().max() <= fp8_max
diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py
new file mode 100644
index 000000000..64849c8f0
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py
@@ -0,0 +1,109 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FSDP2 transformer_impl wrapping logic.
+
+Tests that PrimusTorchFullyShardedDataParallel correctly excludes/includes
+ColumnParallelLinear based on transformer_impl setting, aligning with Megatron's pattern.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from megatron.core import tensor_parallel
+from megatron.core.distributed.distributed_data_parallel_config import (
+    DistributedDataParallelConfig,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from primus.backends.megatron.core.distributed import (
+    torch_fully_sharded_data_parallel as fsdp_mod,
+)
+from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import (
+    PrimusTorchFullyShardedDataParallel,
+)
+from tests.utils import PrimusUT
+
+
+class TestFSDP2TransformerImpl(PrimusUT):
+    """Tests for FSDP2 wrapping logic with transformer_impl."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for FSDP2 tests."""
+
+    def _build_and_collect_wrapped(self, transformer_impl):
+        """Construct the FSDP2 wrapper with fully_shard mocked at the call site.
+
+        Returns the list of module types that fully_shard was invoked on. The
+        mock must patch the name bound INSIDE the production module (it does
+        ``from torch.distributed.fsdp import fully_shard`` at import time), not
+        ``torch.distributed.fsdp.fully_shard`` -- otherwise the real kernel runs
+        and the wrapping decision is never exercised.
+        """
+
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                from megatron.core.tensor_parallel.layers import ColumnParallelLinear
+
+                inner = TransformerConfig(
+                    hidden_size=128,
+                    num_attention_heads=2,
+                    num_layers=1,
+                    transformer_impl=transformer_impl,
+                )
+                self.linear = ColumnParallelLinear(
+                    inner.hidden_size,
+                    inner.hidden_size,
+                    config=inner,
+                    init_method=lambda x: None,
+                )
+
+        model = MockModel().cuda()
+        config = TransformerConfig(
+            hidden_size=128,
+            num_attention_heads=2,
+            num_layers=1,
+            transformer_impl=transformer_impl,
+        )
+        ddp_config = DistributedDataParallelConfig()
+
+        wrapped_modules = []
+
+        def mock_fully_shard(module, **kwargs):
+            wrapped_modules.append(type(module))
+
+        with patch.object(fsdp_mod, "fully_shard", side_effect=mock_fully_shard):
+            PrimusTorchFullyShardedDataParallel(
+                config=config,
+                ddp_config=ddp_config,
+                module=model,
+            )
+        return wrapped_modules
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_fsdp2_excludes_column_parallel_linear_when_local(self):
+        """ColumnParallelLinear must be excluded when transformer_impl == 'local'."""
+        if not fsdp_mod.HAVE_FSDP:
+            pytest.skip("torch.distributed.fsdp (FSDP2) is unavailable")
+
+        wrapped_modules = self._build_and_collect_wrapped("local")
+
+        assert (
+            tensor_parallel.ColumnParallelLinear not in wrapped_modules
+        ), "ColumnParallelLinear should be excluded when transformer_impl == 'local'"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_fsdp2_includes_column_parallel_linear_when_not_local(self):
+        """Positive control: ColumnParallelLinear IS wrapped for non-local impl."""
+        if not fsdp_mod.HAVE_FSDP:
+            pytest.skip("torch.distributed.fsdp (FSDP2) is unavailable")
+
+        wrapped_modules = self._build_and_collect_wrapped("transformer_engine")
+
+        assert (
+            tensor_parallel.ColumnParallelLinear in wrapped_modules
+        ), "ColumnParallelLinear should be wrapped when transformer_impl != 'local'"
diff --git a/tests/unit_tests/backends/megatron/diffusion/helpers.py b/tests/unit_tests/backends/megatron/diffusion/helpers.py
new file mode 100644
index 000000000..e279b3a10
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/helpers.py
@@ -0,0 +1,245 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Test helper utilities for diffusion model testing.
+
+Provides utility functions for tensor validation, mock data creation,
+and common test operations.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_LONG,
+    VAE_LATENT_CHANNELS,
+)
+
+
+def assert_tensor_shape(tensor: torch.Tensor, expected_shape: Tuple[int, ...], name: Optional[str] = None):
+    """
+    Assert that a tensor has the expected shape.
+
+    Args:
+        tensor: Tensor to check
+        expected_shape: Expected shape tuple
+        name: Optional tensor name for error messages
+
+    Raises:
+        AssertionError if shapes don't match
+    """
+    name_str = f"{name} " if name else ""
+    assert (
+        tensor.shape == expected_shape
+    ), f"{name_str}shape mismatch: expected {expected_shape}, got {tensor.shape}"
+
+
+def assert_tensor_dtype(tensor: torch.Tensor, expected_dtype: torch.dtype, name: Optional[str] = None):
+    """
+    Assert that a tensor has the expected dtype.
+
+    Args:
+        tensor: Tensor to check
+        expected_dtype: Expected dtype
+        name: Optional tensor name for error messages
+
+    Raises:
+        AssertionError if dtypes don't match
+    """
+    name_str = f"{name} " if name else ""
+    assert (
+        tensor.dtype == expected_dtype
+    ), f"{name_str}dtype mismatch: expected {expected_dtype}, got {tensor.dtype}"
+
+
+def assert_no_nan_inf(tensor: torch.Tensor, name: Optional[str] = None):
+    """
+    Assert that a tensor contains no NaN or Inf values.
+
+    Args:
+        tensor: Tensor to check
+        name: Optional tensor name for error messages
+
+    Raises:
+        AssertionError if NaN or Inf found
+    """
+    name_str = f"{name} " if name else ""
+    assert not torch.isnan(tensor).any(), f"{name_str}contains NaN values"
+    assert not torch.isinf(tensor).any(), f"{name_str}contains Inf values"
+
+
+def assert_tensor_close(
+    tensor1: torch.Tensor,
+    tensor2: torch.Tensor,
+    rtol: float = 1e-5,
+    atol: float = 1e-8,
+    name: Optional[str] = None,
+):
+    """
+    Assert that two tensors are close within tolerance.
+
+    Args:
+        tensor1: First tensor
+        tensor2: Second tensor
+        rtol: Relative tolerance
+        atol: Absolute tolerance
+        name: Optional tensor name for error messages
+
+    Raises:
+        AssertionError if tensors are not close
+    """
+    name_str = f"{name} " if name else ""
+    assert torch.allclose(
+        tensor1, tensor2, rtol=rtol, atol=atol
+    ), f"{name_str}tensors are not close within rtol={rtol}, atol={atol}"
+
+
+def create_mock_latents(
+    batch_size: int,
+    height: int,
+    width: int,
+    channels: int = VAE_LATENT_CHANNELS,  # VAE output channels for Flux
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    seed: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Create mock latent tensors for testing.
+
+    Args:
+        batch_size: Batch size
+        height: Latent height
+        width: Latent width
+        channels: Number of channels (default: 16 for Flux)
+        device: Device to create tensors on
+        dtype: Tensor dtype
+        seed: Optional random seed for reproducibility
+
+    Returns:
+        Mock latents of shape (batch_size, channels, height, width)
+    """
+    if seed is not None:
+        generator = torch.Generator(device=device).manual_seed(seed)
+        return torch.randn(
+            batch_size, channels, height, width, device=device, dtype=dtype, generator=generator
+        )
+    else:
+        return torch.randn(batch_size, channels, height, width, device=device, dtype=dtype)
+
+
+def create_mock_text_embeddings(
+    batch_size: int,
+    seq_len: int,
+    hidden_dim: int = T5_XXL_EMBEDDING_DIM,  # T5-XXL dimension
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    seed: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Create mock T5 text embeddings for testing.
+
+    Args:
+        batch_size: Batch size
+        seq_len: Sequence length
+        hidden_dim: Hidden dimension (default: 4096 for T5-XXL)
+        device: Device to create tensors on
+        dtype: Tensor dtype
+        seed: Optional random seed for reproducibility
+
+    Returns:
+        Mock text embeddings of shape (batch_size, seq_len, hidden_dim)
+    """
+    if seed is not None:
+        generator = torch.Generator(device=device).manual_seed(seed)
+        return torch.randn(batch_size, seq_len, hidden_dim, device=device, dtype=dtype, generator=generator)
+    else:
+        return torch.randn(batch_size, seq_len, hidden_dim, device=device, dtype=dtype)
+
+
+def create_mock_clip_embeddings(
+    batch_size: int,
+    hidden_dim: int = CLIP_L_EMBEDDING_DIM,  # CLIP-L dimension
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    seed: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Create mock CLIP pooled embeddings for testing.
+
+    Args:
+        batch_size: Batch size
+        hidden_dim: Hidden dimension (default: 768 for CLIP-L)
+        device: Device to create tensors on
+        dtype: Tensor dtype
+        seed: Optional random seed for reproducibility
+
+    Returns:
+        Mock CLIP embeddings of shape (batch_size, hidden_dim)
+    """
+    if seed is not None:
+        generator = torch.Generator(device=device).manual_seed(seed)
+        return torch.randn(batch_size, hidden_dim, device=device, dtype=dtype, generator=generator)
+    else:
+        return torch.randn(batch_size, hidden_dim, device=device, dtype=dtype)
+
+
+def create_mock_position_ids_3d(
+    batch_size: int,
+    height: int,
+    width: int,
+    text_seq_len: int = TEXT_SEQ_LEN_LONG,  # Default text sequence length
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Create mock 3D position IDs for Flux (RoPE format).
+
+    Args:
+        batch_size: Batch size
+        height: Image height (in latent space / 2)
+        width: Image width (in latent space / 2)
+        text_seq_len: Text sequence length
+        device: Device to create tensors on
+        dtype: Tensor dtype
+
+    Returns:
+        Tuple of (img_ids, txt_ids) where:
+            - img_ids: (batch_size, height*width, 3)
+            - txt_ids: (batch_size, text_seq_len, 3)
+    """
+    # Image position IDs
+    img_ids = torch.zeros(batch_size, height * width, 3, device=device, dtype=dtype)
+
+    # Generate 2D grid for spatial positions
+    h_coords = torch.arange(height, device=device, dtype=dtype)
+    w_coords = torch.arange(width, device=device, dtype=dtype)
+    h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing="ij")
+
+    # Flatten and repeat for batch
+    h_flat = h_grid.reshape(-1)
+    w_flat = w_grid.reshape(-1)
+
+    for b in range(batch_size):
+        img_ids[b, :, 1] = h_flat
+        img_ids[b, :, 2] = w_flat
+
+    # Text position IDs (all zeros for text)
+    txt_ids = torch.zeros(batch_size, text_seq_len, 3, device=device, dtype=dtype)
+
+    return img_ids, txt_ids
+
+
+__all__ = [
+    "assert_tensor_shape",
+    "assert_tensor_dtype",
+    "assert_no_nan_inf",
+    "assert_tensor_close",
+    "create_mock_latents",
+    "create_mock_text_embeddings",
+    "create_mock_clip_embeddings",
+    "create_mock_position_ids_3d",
+]
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py b/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py
new file mode 100644
index 000000000..ba815e97f
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py
@@ -0,0 +1,237 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Validate the @triton_op cast_transpose_fp8_triton kernel.
+
+Tests:
+  1. Correctness: FP8 cast + transpose + amax match reference
+  2. GEMM interaction: output feeds into torch._scaled_mm without regression
+  3. torch.compile: triton_op is traceable (no graph breaks)
+  4. Autograd: gradients flow through a minimal FP8 linear autograd.Function
+
+Run:
+    python -m pytest tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py -v
+or standalone:
+    python tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus_turbo.pytorch.core.backend import BackendType
+from primus_turbo.pytorch.core.low_precision import (
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+
+from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import (
+    cast_transpose_fp8_triton,
+)
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    DelayedFP8LinearTensorwiseFunction,
+)
+
+DEVICE = "cuda:0"
+DTYPE = torch.bfloat16
+FP8_DTYPE = float8_e4m3
+FP8_MAX = torch.finfo(FP8_DTYPE).max
+FP8_BWD_DTYPE = float8_e5m2
+FP8_BWD_MAX = torch.finfo(FP8_BWD_DTYPE).max
+
+
+# -----------------------------------------------------------------------
+# Test 1: Correctness
+# -----------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (4096, 3072),
+        (16384, 3072),
+        (8192, 8192),
+    ],
+)
+def test_correctness(shape):
+    M, N = shape
+    x = torch.randn(M, N, dtype=DTYPE, device=DEVICE)
+    scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE)
+    amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE)
+
+    cast_out, trans_out, scale_inv = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf)
+
+    assert cast_out.shape == (M, N), f"Cast shape mismatch: {cast_out.shape}"
+    assert cast_out.dtype == FP8_DTYPE
+    assert trans_out.shape == (N, M), f"Trans shape mismatch: {trans_out.shape}"
+    assert trans_out.dtype == FP8_DTYPE
+    assert trans_out.is_contiguous(), "Transpose output must be contiguous"
+    assert scale_inv.shape == (), f"Scale inv shape mismatch: {scale_inv.shape}"
+
+    ref_scaled = (x.float() * scale.item()).clamp(-FP8_MAX, FP8_MAX)
+    ref_fp8 = ref_scaled.to(FP8_DTYPE)
+    assert torch.equal(cast_out, ref_fp8), f"Cast output mismatch for {shape}"
+
+    ref_trans = ref_fp8.t().contiguous()
+    assert torch.equal(trans_out, ref_trans), f"Transpose output mismatch for {shape}"
+
+    expected_scale_inv = 1.0 / scale.item()
+    assert (
+        abs(scale_inv.item() - expected_scale_inv) < 1e-5
+    ), f"scale_inv {scale_inv.item()} != {expected_scale_inv}"
+
+    expected_amax = x.float().abs().amax().item()
+    assert (
+        abs(amax_buf.item() - expected_amax) / max(expected_amax, 1e-8) < 1e-3
+    ), f"amax {amax_buf.item()} != {expected_amax}"
+
+
+# -----------------------------------------------------------------------
+# Test 2: GEMM interaction (the critical regression test)
+# -----------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (16384, 3072),
+        (4096, 12288),
+    ],
+)
+def test_gemm_interaction(shape):
+    M, N = shape
+    x = torch.randn(M, N, dtype=DTYPE, device=DEVICE)
+    scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE)
+    amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE)
+
+    cast_out, _, scale_inv = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf)
+
+    w = torch.randn(N, N, dtype=DTYPE, device=DEVICE).to(FP8_DTYPE).t()
+    w_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE)
+
+    result = torch._scaled_mm(cast_out, w, out_dtype=DTYPE, scale_a=scale_inv, scale_b=w_scale)
+    assert result.shape == (M, N)
+    assert result.dtype == DTYPE
+    assert not result.isnan().any(), "GEMM produced NaN"
+    assert not result.isinf().any(), "GEMM produced Inf"
+
+
+# -----------------------------------------------------------------------
+# Test 3: torch.compile traceability
+# -----------------------------------------------------------------------
+def test_torch_compile():
+    M, N = 4096, 3072
+    x = torch.randn(M, N, dtype=DTYPE, device=DEVICE)
+    scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE)
+    amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE)
+
+    @torch.compile(fullgraph=True)
+    def fn(x, scale, amax_buf):
+        return cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf)
+
+    cast_out, trans_out, scale_inv = fn(x, scale, amax_buf)
+
+    assert cast_out.shape == (M, N)
+    assert trans_out.shape == (N, M)
+    assert cast_out.dtype == FP8_DTYPE
+
+    # Verify match with eager
+    amax_eager = torch.zeros((), dtype=torch.float32, device=DEVICE)
+    cast_eager, trans_eager, si_eager = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_eager)
+    assert torch.equal(cast_out, cast_eager), "Compiled output != eager output"
+    assert torch.equal(trans_out, trans_eager), "Compiled transpose != eager transpose"
+
+
+# -----------------------------------------------------------------------
+# Test 4: Direct test of the *production* delayed-tensorwise FP8 autograd
+# Function (DelayedFP8LinearTensorwiseFunction), exercising both the native
+# (force_nt=False -> dgrad=NN/wgrad=TN) and forced-NT (force_nt=True) arms.
+#
+# This replaces the previous in-file ``_MinimalFP8Linear`` replica: it calls the
+# shipped Function directly so a regression in its quantization, GEMM layout, or
+# fused amax capture is actually caught.
+# -----------------------------------------------------------------------
+def _rel_err(actual: torch.Tensor, ref: torch.Tensor) -> float:
+    ref = ref.float()
+    return ((actual.float() - ref).norm() / ref.norm().clamp_min(1e-12)).item()
+
+
+@pytest.mark.parametrize("force_nt", [True, False])
+def test_delayed_fp8_linear_function(force_nt):
+    M, K, N = 256, 512, 256
+
+    torch.manual_seed(0)
+    inp = torch.randn(M, K, dtype=DTYPE, device=DEVICE, requires_grad=True)
+    weight = torch.randn(N, K, dtype=DTYPE, device=DEVICE, requires_grad=True)
+    grad_output = torch.randn(M, N, dtype=DTYPE, device=DEVICE)
+
+    # Delayed tensorwise scales (one scalar per track), chosen to fill the FP8
+    # range the way the production warmup staging would.
+    scale_input = torch.tensor(FP8_MAX / inp.detach().abs().amax().item(), dtype=torch.float32, device=DEVICE)
+    scale_weight = torch.tensor(
+        FP8_MAX / weight.detach().abs().amax().item(), dtype=torch.float32, device=DEVICE
+    )
+    scale_grad = torch.tensor(
+        FP8_BWD_MAX / grad_output.abs().amax().item(), dtype=torch.float32, device=DEVICE
+    )
+
+    staged_input_amax = torch.zeros((), dtype=torch.float32, device=DEVICE)
+    staged_weight_amax = torch.zeros((), dtype=torch.float32, device=DEVICE)
+    staged_grad_amax = torch.zeros((), dtype=torch.float32, device=DEVICE)
+
+    gran_value = ScalingGranularity.TENSORWISE.value
+    backend_value = BackendType.HIPBLASLT.value
+
+    result = DelayedFP8LinearTensorwiseFunction.apply(
+        inp,
+        weight,
+        scale_input,
+        scale_weight,
+        scale_grad,
+        staged_input_amax,
+        staged_weight_amax,
+        staged_grad_amax,
+        FP8_DTYPE,
+        FP8_BWD_DTYPE,
+        gran_value,
+        backend_value,
+        force_nt,
+    )
+    output = result[0]
+
+    # Forward: output ~= input @ weight.T (standard nn.Linear), FP8-tensorwise.
+    ref_out = inp.detach().float() @ weight.detach().float().t()
+    assert output.shape == (M, N)
+    assert torch.isfinite(output).all(), "FP8 forward produced non-finite values"
+    out_rel = _rel_err(output, ref_out)
+    assert out_rel < 0.1, f"FP8 forward too far from bf16 reference: rel_err={out_rel:.4f}"
+
+    # Fused amax capture (the delayed-scaling contract): the staged buffers must
+    # be written with the *current* tensor amaxes during forward.
+    assert staged_input_amax.item() > 0, "staged_input_amax not captured in forward"
+    assert staged_weight_amax.item() > 0, "staged_weight_amax not captured in forward"
+    in_amax_rel = abs(staged_input_amax.item() - inp.detach().float().abs().amax().item()) / max(
+        inp.detach().float().abs().amax().item(), 1e-8
+    )
+    w_amax_rel = abs(staged_weight_amax.item() - weight.detach().float().abs().amax().item()) / max(
+        weight.detach().float().abs().amax().item(), 1e-8
+    )
+    assert in_amax_rel < 1e-2, f"staged input amax mismatch: rel={in_amax_rel:.4f}"
+    assert w_amax_rel < 1e-2, f"staged weight amax mismatch: rel={w_amax_rel:.4f}"
+
+    # Backward through the production Function.
+    output.backward(grad_output)
+    assert inp.grad is not None and weight.grad is not None
+    assert torch.isfinite(inp.grad).all(), "grad_input non-finite"
+    assert torch.isfinite(weight.grad).all(), "grad_weight non-finite"
+
+    ref_grad_input = grad_output.float() @ weight.detach().float()
+    ref_grad_weight = grad_output.float().t() @ inp.detach().float()
+    # e5m2 backward has only 2 mantissa bits, so use a looser bound than forward.
+    gi_rel = _rel_err(inp.grad, ref_grad_input)
+    gw_rel = _rel_err(weight.grad, ref_grad_weight)
+    assert gi_rel < 0.2, f"grad_input too far from reference: rel_err={gi_rel:.4f}"
+    assert gw_rel < 0.2, f"grad_weight too far from reference: rel_err={gw_rel:.4f}"
+
+    assert staged_grad_amax.item() > 0, "staged_grad_amax not captured in backward"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py
new file mode 100644
index 000000000..19cda2faf
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py
@@ -0,0 +1,315 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux checkpoint converter.
+
+Tests QKV fusion, key mapping, and checkpoint conversion logic.
+"""
+
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux import FluxConfig
+from tests.utils import PrimusUT
+
+# Note: We use lazy imports for checkpoint_converter functions to avoid
+# triggering parent package imports that require Megatron. Functions are
+# imported inside test methods when needed.
+
+
+class TestQKVWeightFusion(PrimusUT):
+    """Test QKV weight fusion for GQA."""
+
+    def test_qkv_weight_fusion_shape(self):
+        """Test that QKV fusion produces correct output shape."""
+        # Lazy import to avoid parent package Megatron dependency
+        from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+            _fuse_qkv_weights,
+        )
+
+        # Use Flux 12B config: hidden_size=3072, num_heads=24, head_size=128
+        config = FluxConfig.flux_12b()
+
+        hidden_size = config.hidden_size
+        q_weight = torch.randn(hidden_size, hidden_size)
+        k_weight = torch.randn(hidden_size, hidden_size)
+        v_weight = torch.randn(hidden_size, hidden_size)
+
+        # Fuse
+        qkv_weight = _fuse_qkv_weights(config, q_weight, k_weight, v_weight)
+
+        # Expected shape: [head_size * (num_heads + 2*num_query_groups), hidden_size]
+        # For Flux: num_query_groups = num_heads = 24
+        # So: [128 * (24 + 2*24), 3072] = [128 * 72, 3072] = [9216, 3072]
+        num_heads = config.num_attention_heads
+        num_query_groups = getattr(config, "num_query_groups", num_heads)
+        head_size = hidden_size // num_heads
+        expected_shape = (head_size * (num_heads + 2 * num_query_groups), hidden_size)
+
+        assert qkv_weight.shape == expected_shape, f"Expected {expected_shape}, got {qkv_weight.shape}"
+
+    def test_qkv_weight_fusion_interleaving(self):
+        """Test that QKV weights are interleaved correctly per group."""
+        # Lazy import to avoid parent package Megatron dependency
+        from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+            _fuse_qkv_weights,
+        )
+
+        # Use Flux 535M config with simple values for easier validation
+        # Note: Flux 535M has hidden_size=3072, num_heads=24, but we adjust expectations
+        # for the interleaving test which needs smaller values
+        config = FluxConfig.flux_535m()
+        # Override for test simplicity - fusion algorithm works with any config
+        config.hidden_size = 64
+        config.num_attention_heads = 4
+        config.num_query_groups = 4
+
+        hidden_size = 64
+        head_size = 16
+
+        # Create identifiable weights
+        q_weight = torch.ones(hidden_size, hidden_size) * 1.0
+        k_weight = torch.ones(hidden_size, hidden_size) * 2.0
+        v_weight = torch.ones(hidden_size, hidden_size) * 3.0
+
+        qkv_weight = _fuse_qkv_weights(config, q_weight, k_weight, v_weight)
+
+        # Reshape to verify interleaving: [heads_per_group + 2, num_groups, head_size, hidden_size]
+        # For this config: heads_per_group=1, num_groups=4
+        # So we expect pattern: [Q0, K0, V0, Q1, K1, V1, Q2, K2, V2, Q3, K3, V3]
+        qkv_reshaped = qkv_weight.reshape(4, 3, head_size, hidden_size)
+
+        # Check first group
+        assert torch.allclose(qkv_reshaped[0, 0], torch.ones(head_size, hidden_size) * 1.0), "Q0 mismatch"
+        assert torch.allclose(qkv_reshaped[0, 1], torch.ones(head_size, hidden_size) * 2.0), "K0 mismatch"
+        assert torch.allclose(qkv_reshaped[0, 2], torch.ones(head_size, hidden_size) * 3.0), "V0 mismatch"
+
+
+class TestQKVBiasFusion(PrimusUT):
+    """Test QKV bias fusion for GQA."""
+
+    def test_qkv_bias_fusion_shape(self):
+        """Test that QKV bias fusion produces correct output shape."""
+        # Lazy import to avoid parent package Megatron dependency
+        from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+            _fuse_qkv_bias,
+        )
+
+        # Use Flux 12B config
+        config = FluxConfig.flux_12b()
+
+        hidden_size = config.hidden_size
+        q_bias = torch.randn(hidden_size)
+        k_bias = torch.randn(hidden_size)
+        v_bias = torch.randn(hidden_size)
+
+        qkv_bias = _fuse_qkv_bias(config, q_bias, k_bias, v_bias)
+
+        # Expected shape: [head_size * (num_heads + 2*num_query_groups)]
+        num_heads = config.num_attention_heads
+        num_query_groups = getattr(config, "num_query_groups", num_heads)
+        head_size = hidden_size // num_heads
+        expected_shape = (head_size * (num_heads + 2 * num_query_groups),)
+
+        assert qkv_bias.shape == expected_shape, f"Expected {expected_shape}, got {qkv_bias.shape}"
+
+    def test_qkv_bias_fusion_interleaving(self):
+        """Test that QKV biases are interleaved correctly."""
+        # Lazy import to avoid parent package Megatron dependency
+        from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+            _fuse_qkv_bias,
+        )
+
+        # Use Flux 535M config with simple values for easier validation
+        config = FluxConfig.flux_535m()
+        # Override for test simplicity - fusion algorithm works with any config
+        config.hidden_size = 64
+        config.num_attention_heads = 4
+        config.num_query_groups = 4
+
+        # Create identifiable biases
+        q_bias = torch.ones(64) * 1.0
+        k_bias = torch.ones(64) * 2.0
+        v_bias = torch.ones(64) * 3.0
+
+        qkv_bias = _fuse_qkv_bias(config, q_bias, k_bias, v_bias)
+
+        # Reshape to verify: [num_groups, heads_per_group + 2, head_size]
+        qkv_reshaped = qkv_bias.reshape(4, 3, 16)
+
+        # Check first group
+        assert torch.allclose(qkv_reshaped[0, 0], torch.ones(16) * 1.0)
+        assert torch.allclose(qkv_reshaped[0, 1], torch.ones(16) * 2.0)
+        assert torch.allclose(qkv_reshaped[0, 2], torch.ones(16) * 3.0)
+
+
+class TestCheckpointConversion:
+    """Test end-to-end checkpoint conversion (with mock data)."""
+
+    def test_mock_conversion(self, tmp_path):
+        """Test conversion with mock HF checkpoint."""
+        from safetensors.torch import save_file as save_safetensors
+
+        from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import (
+            convert_hf_checkpoint,
+        )
+
+        # Create minimal mock HF checkpoint for Flux 535M (1 joint + 1 single layer)
+        config = FluxConfig.flux_535m()
+        hidden_size = config.hidden_size
+
+        mock_state_dict = {}
+
+        # Double block 0
+        mock_state_dict["transformer_blocks.0.norm1.linear.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.norm1.linear.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_q.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_q.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_k.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_k.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_v.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_v.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_out.0.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_out.0.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.norm_q.weight"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.norm_k.weight"] = torch.randn(hidden_size)
+
+        # Added attention
+        mock_state_dict["transformer_blocks.0.attn.add_q_proj.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.add_q_proj.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.add_k_proj.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.add_k_proj.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.add_v_proj.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.add_v_proj.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_add_out.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.to_add_out.bias"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.norm_added_q.weight"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.attn.norm_added_k.weight"] = torch.randn(hidden_size)
+        mock_state_dict["transformer_blocks.0.norm1_context.linear.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["transformer_blocks.0.norm1_context.linear.bias"] = torch.randn(hidden_size)
+
+        # MLP
+        mock_state_dict["transformer_blocks.0.ff.net.0.proj.weight"] = torch.randn(
+            hidden_size * 4, hidden_size
+        )
+        mock_state_dict["transformer_blocks.0.ff.net.0.proj.bias"] = torch.randn(hidden_size * 4)
+        mock_state_dict["transformer_blocks.0.ff.net.2.weight"] = torch.randn(hidden_size, hidden_size * 4)
+        mock_state_dict["transformer_blocks.0.ff.net.2.bias"] = torch.randn(hidden_size)
+
+        # Context MLP
+        mock_state_dict["transformer_blocks.0.ff_context.net.0.proj.weight"] = torch.randn(
+            hidden_size * 4, hidden_size
+        )
+        mock_state_dict["transformer_blocks.0.ff_context.net.0.proj.bias"] = torch.randn(hidden_size * 4)
+        mock_state_dict["transformer_blocks.0.ff_context.net.2.weight"] = torch.randn(
+            hidden_size, hidden_size * 4
+        )
+        mock_state_dict["transformer_blocks.0.ff_context.net.2.bias"] = torch.randn(hidden_size)
+
+        # Single block 0
+        mock_state_dict["single_transformer_blocks.0.norm.linear.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["single_transformer_blocks.0.norm.linear.bias"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.attn.to_q.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["single_transformer_blocks.0.attn.to_q.bias"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.attn.to_k.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["single_transformer_blocks.0.attn.to_k.bias"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.attn.to_v.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["single_transformer_blocks.0.attn.to_v.bias"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.attn.norm_q.weight"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.attn.norm_k.weight"] = torch.randn(hidden_size)
+        mock_state_dict["single_transformer_blocks.0.proj_mlp.weight"] = torch.randn(
+            hidden_size * 4, hidden_size
+        )
+        mock_state_dict["single_transformer_blocks.0.proj_mlp.bias"] = torch.randn(hidden_size * 4)
+        mock_state_dict["single_transformer_blocks.0.proj_out.weight"] = torch.randn(
+            hidden_size, hidden_size * 2
+        )
+        mock_state_dict["single_transformer_blocks.0.proj_out.bias"] = torch.randn(hidden_size)
+
+        # Root level
+        mock_state_dict["x_embedder.weight"] = torch.randn(hidden_size, 64)
+        mock_state_dict["x_embedder.bias"] = torch.randn(hidden_size)
+        mock_state_dict["context_embedder.weight"] = torch.randn(hidden_size, 4096)
+        mock_state_dict["context_embedder.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = torch.randn(hidden_size, 256)
+        mock_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.text_embedder.linear_1.weight"] = torch.randn(hidden_size, 768)
+        mock_state_dict["time_text_embed.text_embedder.linear_1.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.text_embedder.linear_2.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["time_text_embed.text_embedder.linear_2.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = torch.randn(hidden_size, 256)
+        mock_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = torch.randn(hidden_size)
+        mock_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = torch.randn(
+            hidden_size, hidden_size
+        )
+        mock_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = torch.randn(hidden_size)
+        mock_state_dict["norm_out.linear.weight"] = torch.randn(hidden_size, hidden_size)
+        mock_state_dict["norm_out.linear.bias"] = torch.randn(hidden_size)
+        mock_state_dict["proj_out.weight"] = torch.randn(64, hidden_size)
+        mock_state_dict["proj_out.bias"] = torch.randn(64)
+
+        # Save mock checkpoint
+        checkpoint_path = tmp_path / "mock_flux.safetensors"
+        save_safetensors(mock_state_dict, str(checkpoint_path))
+
+        # Convert
+        primus_state_dict = convert_hf_checkpoint(
+            checkpoint_path,
+            flux_config=config,
+            save_to=None,
+        )
+
+        # Verify conversion
+        assert len(primus_state_dict) > 0, "Converted state dict is empty"
+
+        # Check that QKV was fused (with new TransformerBlock naming)
+        assert "transformer.layers.0.self_attention.linear_qkv.weight" in primus_state_dict
+        assert "transformer.layers.0.self_attention.linear_qkv.bias" in primus_state_dict
+        assert "transformer.layers.1.self_attention.linear_qkv.weight" in primus_state_dict
+
+        # Check that proj_out was split for single blocks
+        assert "transformer.layers.1.self_attention.linear_proj.weight" in primus_state_dict
+        assert "transformer.layers.1.mlp.linear_fc2.weight" in primus_state_dict
+
+        # Verify key mapping worked (with new TransformerBlock naming)
+        assert "transformer.layers.0.adaln.adaLN_modulation.1.weight" in primus_state_dict
+        # Note: Layer 1 is FluxSingleTransformerBlock which has different adaln structure
+
+        # norm_out scale/shift swap (the converter's only non-trivial root-level
+        # math): HF Diffusers stores the modulation as [SCALE; SHIFT] but
+        # Primus/BFL native expects [SHIFT; SCALE], so the two halves must be
+        # swapped along dim 0 -- a plain key-rename/passthrough would be a bug.
+        orig_w = mock_state_dict["norm_out.linear.weight"]
+        orig_b = mock_state_dict["norm_out.linear.bias"]
+        half_w = orig_w.shape[0] // 2
+        half_b = orig_b.shape[0] // 2
+        expected_w = torch.cat([orig_w[half_w:, :], orig_w[:half_w, :]], dim=0)
+        expected_b = torch.cat([orig_b[half_b:], orig_b[:half_b]], dim=0)
+
+        conv_w = primus_state_dict["norm_out.adaLN_modulation.1.weight"]
+        conv_b = primus_state_dict["norm_out.adaLN_modulation.1.bias"]
+        assert torch.equal(conv_w, expected_w), "norm_out weight halves were not swapped"
+        assert torch.equal(conv_b, expected_b), "norm_out bias halves were not swapped"
+        # Positive control: the swap must actually reorder, not pass through.
+        assert not torch.equal(conv_w, orig_w), "norm_out weight passed through without the scale/shift swap"
+        assert not torch.equal(conv_b, orig_b), "norm_out bias passed through without the scale/shift swap"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py
new file mode 100644
index 000000000..37e41725a
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py
@@ -0,0 +1,84 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for NeMo chimera initialization behavior.
+
+Validates that contaminating the default CUDA generator with per-DP-rank seeds
+causes non-parallel layers (img_embed, txt_embed, MLPEmbedder) to diverge
+across ranks, while ColumnParallelLinear layers (using model-parallel tracker)
+remain identical.
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from tests.utils import PrimusUT
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestFluxChimeraInit(PrimusUT):
+    """Tests that chimera init creates intended per-rank diversity."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    def _build_model_with_rank_seed(self, base_seed, dp_rank):
+        """Build a Flux model simulating chimera init for a given DP rank."""
+        from megatron.core.tensor_parallel import random as tp_random
+
+        per_rank_seed = base_seed + 100 * dp_rank
+        torch.manual_seed(per_rank_seed)
+        torch.cuda.manual_seed(per_rank_seed)
+        # Model-parallel tracker uses the SAME seed across all DP ranks
+        tp_random.model_parallel_cuda_manual_seed(42)
+
+        config = FluxConfig.flux_535m()
+        return Flux(config)
+
+    def test_img_embed_differs_across_ranks(self):
+        """img_embed (nn.Linear) should differ when default generator seed differs."""
+        model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0)
+        model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1)
+
+        assert not torch.equal(
+            model_rank0.img_embed.weight, model_rank1.img_embed.weight
+        ), "img_embed.weight should differ between DP ranks"
+
+    def test_txt_embed_differs_across_ranks(self):
+        """txt_embed (nn.Linear) should differ when default generator seed differs."""
+        model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0)
+        model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1)
+
+        assert not torch.equal(
+            model_rank0.txt_embed.weight, model_rank1.txt_embed.weight
+        ), "txt_embed.weight should differ between DP ranks"
+
+    def test_timestep_embedding_differs_across_ranks(self):
+        """MLPEmbedder layers (nn.Linear) should differ across DP ranks."""
+        model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0)
+        model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1)
+
+        w0 = model_rank0.timestep_embedding.time_embedding.in_layer.weight
+        w1 = model_rank1.timestep_embedding.time_embedding.in_layer.weight
+        assert not torch.equal(w0, w1), "timestep_embedding in_layer should differ between DP ranks"
+
+    def test_adaln_zero_init_identical_regardless_of_seed(self):
+        """AdaLN zero-init is deterministic — identical on both ranks."""
+        model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0)
+        model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1)
+
+        for i, (layer0, layer1) in enumerate(
+            zip(model_rank0.transformer.layers, model_rank1.transformer.layers)
+        ):
+            w0 = layer0.adaln.adaLN_modulation[-1].weight
+            w1 = layer1.adaln.adaLN_modulation[-1].weight
+            assert torch.equal(w0, w1), f"Layer {i} adaLN weight should be identical (zero)"
+            assert torch.all(w0 == 0), f"Layer {i} adaLN weight should be zero"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py
new file mode 100644
index 000000000..1db4b6457
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests that torch.compile does not alter sharded_state_dict keys.
+
+The "whole_model" compile strategy compiles self.forward without wrapping
+submodules, so state_dict keys must remain identical before and after
+compilation. This test guards against regressions (e.g., _orig_mod prefixes
+leaking into keys).
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from tests.utils import PrimusUT
+
+
+class TestFluxCompileCheckpointKeys(PrimusUT):
+    """Verify torch.compile does not change sharded_state_dict keys."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_compile_does_not_change_sharded_state_dict_keys(self):
+        """Keys from sharded_state_dict() must be identical before and after compile_model().
+
+        Compile must actually run (enable_torch_compile=True), otherwise
+        compile_model() returns early and the test is vacuous. 'per_block'
+        rebinds each layer's .forward (no module re-wrap), so a regression that
+        wraps the module instead would leak '_orig_mod' prefixes into the keys
+        and fail this test.
+        """
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "per_block"
+        model = Flux(config).cuda()
+
+        keys_before = set(model.sharded_state_dict().keys())
+
+        model.compile_model()
+
+        keys_after = set(model.sharded_state_dict().keys())
+
+        added = keys_after - keys_before
+        removed = keys_before - keys_after
+
+        assert not added, f"Keys added after compile: {added}"
+        assert not removed, f"Keys removed after compile: {removed}"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py
new file mode 100644
index 000000000..fbde0e0a4
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py
@@ -0,0 +1,128 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux compile_model() strategy validation.
+
+Tests that invalid configuration combinations raise the expected errors
+and that fallback behavior works correctly.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from tests.utils import PrimusUT
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestFluxCompileStrategies(PrimusUT):
+    """Tests for compile_model() error/warning behavior."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    def test_whole_model_rejects_overlap_param_gather(self):
+        """whole_model strategy must reject overlap_param_gather."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "whole_model"
+        model = Flux(config)
+
+        with patch(
+            "megatron.training.get_args",
+            return_value=SimpleNamespace(overlap_param_gather=True),
+        ):
+            with pytest.raises(ValueError, match="overlap_param_gather"):
+                model.compile_model()
+
+    def test_stack_rejects_recompute_granularity(self):
+        """stack strategy must reject recompute_granularity='full'."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "stack"
+        config.recompute_granularity = "full"
+        model = Flux(config)
+
+        with pytest.raises(ValueError, match="recompute_granularity"):
+            model.compile_model()
+
+    def test_double_stack_rejects_recompute_granularity(self):
+        """double_stack strategy must reject recompute_granularity='full'."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "double_stack"
+        config.recompute_granularity = "full"
+        model = Flux(config)
+
+        with pytest.raises(ValueError, match="recompute_granularity"):
+            model.compile_model()
+
+    def test_local_spec_stack_falls_back_to_per_block(self):
+        """Local spec (non-TE) with stack strategy should fallback to per_block."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "stack"
+        config.transformer_impl = "local"
+        model = Flux(config)
+
+        compile_calls = []
+
+        def mock_compile(fn, **kwargs):
+            compile_calls.append(fn)
+            return fn
+
+        with patch("torch.compile", side_effect=mock_compile):
+            model.compile_model()
+
+        assert len(compile_calls) == len(model.transformer.layers), (
+            f"Expected {len(model.transformer.layers)} per-block compile calls, " f"got {len(compile_calls)}"
+        )
+
+    def test_per_block_with_cuda_graph_warns(self):
+        """per_block + enable_cuda_graph should emit a warning."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "per_block"
+        # Build without cuda_graph (avoids CudaGraphManager compat issues),
+        # then set the flag before compile_model() which only reads config.
+        model = Flux(config)
+        model.config.enable_cuda_graph = True
+
+        compile_calls = []
+
+        def mock_compile(fn, **kwargs):
+            compile_calls.append(fn)
+            return fn
+
+        output_lines = []
+
+        def capture_log(msg):
+            output_lines.append(msg)
+
+        with patch("torch.compile", side_effect=mock_compile), patch(
+            "primus.core.utils.module_utils.log_rank_0", side_effect=capture_log
+        ):
+            model.compile_model()
+
+        output_text = "\n".join(output_lines)
+        assert "CUDA graph" in output_text, f"Expected CUDA graph warning in output, got: {output_text[:200]}"
+
+    def test_invalid_strategy_raises_value_error(self):
+        """Invalid strategy name must raise ValueError."""
+        config = FluxConfig.flux_535m()
+        config.enable_torch_compile = True
+        config.torch_compile_strategy = "nonexistent"
+        model = Flux(config)
+
+        with pytest.raises(ValueError, match="Invalid torch_compile_strategy"):
+            model.compile_model()
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py
new file mode 100644
index 000000000..2a4348832
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py
@@ -0,0 +1,94 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux and base diffusion configurations.
+
+Tests FluxConfig and BaseDiffusionConfig validation, preset configurations.
+"""
+
+import pytest
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.common import BaseDiffusionConfig
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from tests.utils import PrimusUT
+
+# ========================================================================
+# Base Diffusion Configuration Tests
+# ========================================================================
+
+
+class TestBaseDiffusionConfig(PrimusUT):
+    """Tests for BaseDiffusionConfig."""
+
+    def test_base_config_validation_invalid_channels(self):
+        """Test validation catches invalid channel counts."""
+        with self.assertRaises(ValueError) as cm:
+            config = BaseDiffusionConfig(
+                in_channels=0,
+                num_attention_heads=8,
+                num_layers=1,
+            )
+            config.validate()
+        self.assertIn("in_channels must be positive", str(cm.exception))
+
+
+class TestFluxConfig(PrimusUT):
+    """Tests for FluxConfig class."""
+
+    def test_validation_positive_joint_layers(self):
+        """Test validation fails for non-positive num_joint_layers."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(num_joint_layers=0)
+            config.validate()
+        self.assertIn("num_joint_layers must be positive", str(cm.exception))
+
+    def test_validation_positive_single_layers(self):
+        """Test validation fails for non-positive num_single_layers."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(num_single_layers=-1)
+            config.validate()
+        self.assertIn("num_single_layers must be positive", str(cm.exception))
+
+    def test_validation_positive_context_dim(self):
+        """Test validation fails for non-positive context_dim."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(context_dim=0)
+            config.validate()
+        self.assertIn("context_dim must be positive", str(cm.exception))
+
+    def test_validation_positive_vec_in_dim(self):
+        """Test validation fails for non-positive vec_in_dim."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(vec_in_dim=-768)
+            config.validate()
+        self.assertIn("vec_in_dim must be positive", str(cm.exception))
+
+    def test_validation_positive_theta(self):
+        """Test validation fails for non-positive theta."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(theta=0)
+            config.validate()
+        self.assertIn("theta must be positive", str(cm.exception))
+
+    def test_validation_axes_dim_length(self):
+        """Test validation fails for axes_dim with wrong length."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(axes_dim=(16, 56))  # Only 2 elements
+            config.validate()
+        self.assertIn("axes_dim must have 3 elements", str(cm.exception))
+
+    def test_validation_axes_dim_positive_values(self):
+        """Test validation fails for non-positive axes_dim values."""
+        with self.assertRaises(ValueError) as cm:
+            config = FluxConfig(axes_dim=(16, 0, 56))
+            config.validate()
+        self.assertIn("All axes_dim values must be positive", str(cm.exception))
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py
new file mode 100644
index 000000000..edf5be8b1
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for diffusion model embedding layers.
+
+Tests TimeStepEmbedder and MLPEmbedder to ensure correct shapes and behavior.
+"""
+
+import pytest
+import torch
+
+from primus.backends.megatron.core.models.diffusion.common.embeddings import (
+    get_timestep_embedding,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    TIMESTEP_EMBEDDING_DIM,
+)
+from tests.utils import PrimusUT
+
+
+class TestGetTimestepEmbedding(PrimusUT):
+    """Tests for get_timestep_embedding function (Primus implementation of sinusoidal embeddings)."""
+
+    def test_odd_embedding_dim(self):
+        """Test that odd embedding dimensions are handled correctly by Primus implementation."""
+        timesteps = torch.randn(4)
+        embedding_dim = 257  # Odd number
+
+        emb = get_timestep_embedding(timesteps, embedding_dim)
+
+        assert emb.shape == (4, embedding_dim), f"Expected shape (4, {embedding_dim}), got {emb.shape}"
+
+    def test_sinusoidal_properties(self):
+        """Test that Primus sinusoidal embeddings have correct mathematical properties."""
+        timesteps = torch.linspace(0, 100, 10)
+        embedding_dim = TIMESTEP_EMBEDDING_DIM
+
+        emb = get_timestep_embedding(timesteps, embedding_dim)
+
+        assert emb.abs().max() < 100, "Embeddings should be in reasonable range"
+        assert (emb > 0).any() and (emb < 0).any(), "Should have both positive and negative values"
+
+        # For timestep=0: cos(0)=1 for all frequencies (first half),
+        # sin(0)=0 for all frequencies (second half)
+        emb_zero = get_timestep_embedding(torch.tensor([0.0]), embedding_dim)
+        half = embedding_dim // 2
+        assert torch.allclose(
+            emb_zero[0, :half], torch.ones(half), atol=1e-6
+        ), "cos(0) should be 1 for all frequencies"
+        assert torch.allclose(
+            emb_zero[0, half : 2 * half], torch.zeros(half), atol=1e-6
+        ), "sin(0) should be 0 for all frequencies"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py
new file mode 100644
index 000000000..2bd83cef0
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py
@@ -0,0 +1,51 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FluxModel.get_fp8_context() behavior.
+
+Verifies that the local spec FP8 path returns nullcontext (no global state),
+while the non-local path delegates to Megatron's get_fp8_context utility.
+"""
+
+import contextlib
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+
+
+class TestFluxFP8Context:
+    """Tests for Flux.get_fp8_context() without constructing a full model."""
+
+    def test_get_fp8_context_local_returns_nullcontext(self):
+        """Local spec should return nullcontext to avoid global FP8 state."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="local",
+            fp8="e4m3",
+        )
+
+        ctx = Flux.get_fp8_context(mock_model)
+        assert isinstance(ctx, contextlib.nullcontext)
+
+    def test_get_fp8_context_non_local_delegates(self):
+        """Non-local spec should delegate to megatron.core.fp8_utils.get_fp8_context."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="transformer_engine",
+            fp8="e4m3",
+        )
+
+        sentinel = contextlib.nullcontext()
+        with patch(
+            "megatron.core.fp8_utils.get_fp8_context",
+            return_value=sentinel,
+        ) as mock_get:
+            ctx = Flux.get_fp8_context(mock_model)
+            mock_get.assert_called_once_with(mock_model.config)
+            assert ctx is sentinel
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py
new file mode 100644
index 000000000..687350544
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py
@@ -0,0 +1,172 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux model weight initialization.
+
+Validates the NeMo-aligned init_weights() implementation:
+- AdaLN modulation layers are zero-initialized
+- Embeddings receive Xavier/Normal initialization
+- proj_out and norm_out are zero-initialized
+- Initialization is deterministic given fixed RNG state
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from tests.utils import PrimusUT
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestFluxInitWeights(PrimusUT):
+    """Tests for Flux init_weights() correctness."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    def _create_model(self):
+        config = FluxConfig.flux_535m()
+        return Flux(config)
+
+    def test_adaln_modulation_zero_init_single_blocks(self):
+        """Single-block AdaLN modulation last layer must be zero."""
+        model = self._create_model()
+        num_joint = model.config.num_joint_layers
+
+        for i, layer in enumerate(model.transformer.layers[num_joint:]):
+            w = layer.adaln.adaLN_modulation[-1].weight
+            b = layer.adaln.adaLN_modulation[-1].bias
+            assert torch.all(w == 0), f"Single block {i}: adaLN weight not zero"
+            assert torch.all(b == 0), f"Single block {i}: adaLN bias not zero"
+
+    def test_adaln_modulation_zero_init_joint_blocks(self):
+        """Joint-block AdaLN modulation last layers must be zero."""
+        model = self._create_model()
+        num_joint = model.config.num_joint_layers
+
+        for i, layer in enumerate(model.transformer.layers[:num_joint]):
+            w = layer.adaln.adaLN_modulation[-1].weight
+            b = layer.adaln.adaLN_modulation[-1].bias
+            assert torch.all(w == 0), f"Joint block {i}: adaln weight not zero"
+            assert torch.all(b == 0), f"Joint block {i}: adaln bias not zero"
+
+            w_ctx = layer.adaln_context.adaLN_modulation[-1].weight
+            b_ctx = layer.adaln_context.adaLN_modulation[-1].bias
+            assert torch.all(w_ctx == 0), f"Joint block {i}: adaln_context weight not zero"
+            assert torch.all(b_ctx == 0), f"Joint block {i}: adaln_context bias not zero"
+
+    def test_proj_out_zero_init(self):
+        """proj_out weight and bias must be zero."""
+        model = self._create_model()
+        assert torch.all(model.proj_out.weight == 0)
+        assert torch.all(model.proj_out.bias == 0)
+
+    def test_norm_out_adaln_zero_init(self):
+        """norm_out.adaLN_modulation last layer must be zero."""
+        model = self._create_model()
+        w = model.norm_out.adaLN_modulation[-1].weight
+        b = model.norm_out.adaLN_modulation[-1].bias
+        assert torch.all(w == 0)
+        assert torch.all(b == 0)
+
+    def test_img_embed_xavier_distribution(self):
+        """img_embed should match Xavier-uniform statistics.
+
+        xavier_uniform_ produces U(-bound, bound) with bound = sqrt(6/(in+out));
+        std of that distribution is bound/sqrt(3) = sqrt(2/(in+out)).
+        """
+        import math
+
+        model = self._create_model()
+        w = model.img_embed.weight
+        out_dim, in_dim = w.shape
+        expected_std = math.sqrt(2.0 / (in_dim + out_dim))
+
+        actual_std = w.float().std().item()
+        actual_mean = w.float().mean().item()
+
+        rel_err = abs(actual_std - expected_std) / expected_std
+        assert rel_err < 0.15, (
+            f"img_embed std={actual_std:.4f} deviates from Xavier expected "
+            f"{expected_std:.4f} by {rel_err:.1%} (limit 15%)"
+        )
+        assert abs(actual_mean) < 0.005, f"img_embed mean={actual_mean:.4f} should be ~0 for Xavier"
+
+    def test_txt_embed_xavier_distribution(self):
+        """txt_embed should match Xavier-uniform statistics."""
+        import math
+
+        model = self._create_model()
+        w = model.txt_embed.weight
+        out_dim, in_dim = w.shape
+        expected_std = math.sqrt(2.0 / (in_dim + out_dim))
+
+        actual_std = w.float().std().item()
+        actual_mean = w.float().mean().item()
+
+        rel_err = abs(actual_std - expected_std) / expected_std
+        assert rel_err < 0.15, (
+            f"txt_embed std={actual_std:.4f} deviates from Xavier expected "
+            f"{expected_std:.4f} by {rel_err:.1%} (limit 15%)"
+        )
+        assert abs(actual_mean) < 0.005, f"txt_embed mean={actual_mean:.4f} should be ~0 for Xavier"
+
+    def test_timestep_embedding_normal_init(self):
+        """Timestep embedding MLP should match Normal(std=0.02) statistics."""
+        model = self._create_model()
+        w = model.timestep_embedding.time_embedding.in_layer.weight
+        expected_std = 0.02
+
+        actual_std = w.float().std().item()
+        actual_mean = w.float().mean().item()
+
+        rel_err = abs(actual_std - expected_std) / expected_std
+        assert rel_err < 0.15, (
+            f"timestep_embedding std={actual_std:.4f} deviates from Normal(0,0.02) "
+            f"expected std {expected_std} by {rel_err:.1%} (limit 15%)"
+        )
+        assert (
+            abs(actual_mean) < 0.005
+        ), f"timestep_embedding mean={actual_mean:.4f} should be ~0 for Normal init"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestFluxInitWeightsDeterminism(PrimusUT):
+    """Tests that init_weights produces deterministic results given fixed RNG."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    def test_deterministic_construction(self):
+        """Two models built with the same RNG state must be bitwise identical."""
+        from megatron.core.tensor_parallel import random as tp_random
+
+        config = FluxConfig.flux_535m()
+        seed = 123
+
+        # Build model 1
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        tp_random.model_parallel_cuda_manual_seed(seed)
+        model1 = Flux(config)
+
+        # Build model 2 with identical RNG state
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        tp_random.model_parallel_cuda_manual_seed(seed)
+        model2 = Flux(config)
+
+        for (name1, p1), (name2, p2) in zip(model1.named_parameters(), model2.named_parameters()):
+            assert name1 == name2
+            assert torch.equal(p1, p2), (
+                f"Parameter {name1} differs between constructions "
+                f"(max diff: {(p1 - p2).abs().max().item():.2e})"
+            )
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py
new file mode 100644
index 000000000..d1af8b474
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py
@@ -0,0 +1,79 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux layer spec backend selection.
+
+Tests that get_flux_layer_spec() correctly selects backend based on transformer_impl.
+This ensures alignment between backend selection and FSDP2 wrapping decisions.
+"""
+
+import pytest
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.layer_spec import (
+    get_flux_layer_spec,
+)
+from tests.utils import PrimusUT
+
+
+class TestFluxLayerSpecBackendSelection(PrimusUT):
+    """Tests for backend selection in get_flux_layer_spec()."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for layer spec tests."""
+
+    def test_backend_selection_fp8_local_spec(self):
+        """Test that local + fp8 selects PrimusTurboFloat8LocalSpecProvider."""
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            Float8ColumnParallelLinear,
+        )
+
+        config = FluxConfig.flux_535m(
+            transformer_impl="local",
+            fp8="e4m3",
+            fp8_recipe="tensorwise",
+        )
+
+        block_submodules = get_flux_layer_spec(config, backend=None)
+
+        # At least one linear spec should reference Float8ColumnParallelLinear
+        found_fp8 = False
+        for layer_spec in block_submodules.layer_specs:
+            attn_spec = layer_spec.submodules.self_attention
+            if hasattr(attn_spec, "submodules") and hasattr(attn_spec.submodules, "linear_qkv"):
+                if attn_spec.submodules.linear_qkv == Float8ColumnParallelLinear:
+                    found_fp8 = True
+                    break
+        assert found_fp8, "Expected Float8ColumnParallelLinear in layer specs for local+fp8"
+
+    def test_backend_selection_local_no_fp8_uses_native_linear(self):
+        """Test that local without fp8 uses native ColumnParallelLinear."""
+
+        from megatron.core.tensor_parallel import ColumnParallelLinear
+
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            Float8ColumnParallelLinear,
+        )
+
+        config = FluxConfig.flux_535m(transformer_impl="local", fp8=None)
+
+        block_submodules = get_flux_layer_spec(config, backend=None)
+
+        found_any = False
+        for layer_spec in block_submodules.layer_specs:
+            attn_spec = layer_spec.submodules.self_attention
+            if hasattr(attn_spec, "submodules") and hasattr(attn_spec.submodules, "linear_qkv"):
+                found_any = True
+                assert (
+                    attn_spec.submodules.linear_qkv != Float8ColumnParallelLinear
+                ), "Should NOT use Float8ColumnParallelLinear when fp8=None"
+                assert (
+                    attn_spec.submodules.linear_qkv == ColumnParallelLinear
+                ), f"Expected native ColumnParallelLinear when fp8=None, got {attn_spec.submodules.linear_qkv}"
+        assert found_any, "No attention linear_qkv specs found to validate backend selection"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py
new file mode 100644
index 000000000..15a42efd3
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py
@@ -0,0 +1,54 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux-specific layers.
+
+Tests EmbedND (3D RoPE) and related functions.
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    BATCH_SIZE_PAIR,
+    HIDDEN_DIM_FLUX,
+    ROPE_THETA_DEFAULT,
+    SEQ_LEN_TINY,
+)
+from tests.utils import PrimusUT
+
+
+class TestEmbedND(PrimusUT):
+    """Tests for EmbedND class."""
+
+    def test_forward_output_shape(self):
+        """Test that forward produces correct output shape.
+
+        Output size is determined by sum(axes_dim), not dim parameter.
+        """
+        dim = HIDDEN_DIM_FLUX  # This is just stored, not used for output size
+        axes_dim = [32, 48, 48]  # Sum = 128
+        batch_size = BATCH_SIZE_PAIR  # Paired sample tests
+        seq_len = SEQ_LEN_TINY  # Small sequence for basic tests
+        num_axes = 3
+
+        embed_nd = EmbedND(dim=dim, theta=ROPE_THETA_DEFAULT, axes_dim=axes_dim)
+        ids = torch.randn(batch_size, seq_len, num_axes)
+
+        output = embed_nd(ids)
+
+        # Output size is based on sum(axes_dim) = 128, not dim = 3072
+        # Shape: [seq, B, 1, sum(axes_dim)] after permute and reshape
+        output_dim = sum(axes_dim)
+        expected_shape = (seq_len, batch_size, 1, output_dim)
+        assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py
new file mode 100644
index 000000000..b8d16c6d3
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py
@@ -0,0 +1,84 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Basic unit tests for Flux model.
+
+Tests the Flux model initialization and basic forward pass.
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+    unpack_latents,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    T5_XXL_EMBEDDING_DIM,
+    TEXT_SEQ_LEN_SHORT,
+    VAE_LATENT_CHANNELS,
+)
+from tests.utils import PrimusUT
+
+
+class TestFluxModel(PrimusUT):
+    """Core tests for Flux model initialization and basic operations."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    def test_forward_pass_small(self):
+        """Test forward pass with small inputs."""
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+
+        config = FluxConfig.flux_535m()
+        model = Flux(config).cuda()
+        model.eval()
+
+        batch_size = 2
+        height, width = 16, 16
+        channels = VAE_LATENT_CHANNELS
+        txt_seq_len = TEXT_SEQ_LEN_SHORT
+
+        # Prepare inputs
+        img = torch.randn(batch_size, channels, height, width).cuda()
+        txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM).cuda()
+        y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM).cuda()
+        timesteps = torch.rand(batch_size).cuda()
+
+        # Pack latents
+        packed_img = pack_latents(img)
+        packed_img = packed_img.transpose(0, 1)
+        txt_t = txt.transpose(0, 1)
+
+        # Generate position IDs
+        img_ids = generate_image_position_ids(batch_size, height, width, device="cuda")
+        txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda()
+
+        # Forward pass
+        with torch.no_grad():
+            output = model(packed_img, txt_t, y, timesteps, img_ids, txt_ids)
+
+        # Unpack output
+        output = output.transpose(0, 1)
+        output = unpack_latents(output, height, width, vae_scale_factor=1)
+
+        # Check output shape
+        assert output.shape == img.shape
+        assert not torch.isnan(output).any()
+        assert not torch.isinf(output).any()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py
new file mode 100644
index 000000000..01298962f
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py
@@ -0,0 +1,147 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for diffusion model normalization layers.
+
+Tests RMSNorm, AdaLN, and AdaLNContinuous.
+
+NOTE: These tests require CUDA and Megatron parallel state initialization.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from primus.backends.megatron.core.models.diffusion.common.normalization import (
+    AdaLN,
+    AdaLNContinuous,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    ATTENTION_SEQ_LEN,
+    BATCH_SIZE_QUAD,
+    HIDDEN_DIM_FLUX,
+    NUM_ATTENTION_HEADS_FLUX,
+)
+from tests.utils import PrimusUT
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA (uses ColumnParallelLinear)")
+class TestAdaLN(PrimusUT):
+    """Tests for AdaLN class."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """
+        Initialize parallel state for AdaLN tests.
+
+        AdaLN uses ColumnParallelLinear which requires Megatron's RNG tracker.
+        The init_parallel_state fixture handles initialization and cleanup.
+        """
+
+    def test_forward_output_chunks(self):
+        """Test that AdaLN produces correct number of chunks and numeric behavior."""
+        config = TransformerConfig(
+            hidden_size=HIDDEN_DIM_FLUX,
+            num_attention_heads=NUM_ATTENTION_HEADS_FLUX,
+            num_layers=1,
+        )
+        n_chunks = 9
+        adaln = AdaLN(config, n_adaln_chunks=n_chunks).cuda()
+
+        timestep_emb = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+        chunks = adaln(timestep_emb)
+
+        assert len(chunks) == n_chunks, f"Expected {n_chunks} chunks, got {len(chunks)}"
+        for chunk in chunks:
+            assert chunk.shape == (BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX)
+
+        # Numeric: gate=0 should zero out the contribution in scale_add
+        x = torch.randn(16, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+        residual = torch.randn(16, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+        zero_gate = torch.zeros(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+        result = adaln.scale_add(residual, x, zero_gate)
+        assert torch.allclose(result, residual, atol=1e-6), "gate=0 should leave residual unchanged"
+
+    def test_default_init_method_zeros_modulation_weight(self):
+        """The default init_method is nn.init.zeros_; modulation weight comes up zero.
+
+        Guards the post-Flux-PR contract that AdaLN's default init is
+        observably zero, so downstream callers don't accidentally pick up
+        the NeMo-aligned normal_ RNG draw without opting in.
+        """
+        config = TransformerConfig(
+            hidden_size=HIDDEN_DIM_FLUX,
+            num_attention_heads=NUM_ATTENTION_HEADS_FLUX,
+            num_layers=1,
+        )
+        adaln = AdaLN(config, n_adaln_chunks=6).cuda()
+        weight = adaln.adaLN_modulation[-1].weight
+        assert torch.equal(
+            weight, torch.zeros_like(weight)
+        ), "Default AdaLN init_method should produce a zero modulation weight"
+
+    def test_normal_init_method_produces_nonzero_modulation_weight(self):
+        """Passing nn.init.normal_ produces a nonzero pre-init_weights() draw.
+
+        This is the call sites used by Flux's layer_spec.py to match NeMo's
+        RNG sequence. Flux's init_weights() immediately re-zeroes these
+        weights, so the only observable effect is RNG advancement.
+        """
+        config = TransformerConfig(
+            hidden_size=HIDDEN_DIM_FLUX,
+            num_attention_heads=NUM_ATTENTION_HEADS_FLUX,
+            num_layers=1,
+        )
+        torch.manual_seed(0)
+        adaln = AdaLN(config, n_adaln_chunks=6, init_method=nn.init.normal_).cuda()
+        weight = adaln.adaLN_modulation[-1].weight
+        assert not torch.equal(
+            weight, torch.zeros_like(weight)
+        ), "init_method=normal_ should draw a nonzero modulation weight"
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="AdaLNContinuous forward dispatches to primus::fused_ln_modulate, "
+    "which has no CPU kernel registered.",
+)
+class TestAdaLNContinuous(PrimusUT):
+    """Tests for AdaLNContinuous class."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """AdaLNContinuous uses RowParallelLinear which requires Megatron's
+        RNG tracker / parallel state. Mirror TestAdaLN."""
+
+    def test_output_shape(self):
+        """Test that AdaLNContinuous produces correct output shape."""
+        config = TransformerConfig(
+            hidden_size=HIDDEN_DIM_FLUX,
+            num_attention_heads=NUM_ATTENTION_HEADS_FLUX,
+            num_layers=1,
+        )
+        adaln = AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX).cuda()
+
+        # Use sequence-first format: [seq_len, batch, hidden]
+        x = torch.randn(ATTENTION_SEQ_LEN, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+        cond = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda()
+
+        output = adaln(x, cond)
+        assert output.shape == x.shape
+
+    def test_invalid_norm_type(self):
+        """Test that invalid norm type raises error (Primus validation)."""
+        config = TransformerConfig(
+            hidden_size=HIDDEN_DIM_FLUX,
+            num_attention_heads=NUM_ATTENTION_HEADS_FLUX,
+            num_layers=1,
+        )
+
+        with pytest.raises(ValueError, match="Unknown normalization type"):
+            AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX, norm_type="invalid")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py
new file mode 100644
index 000000000..ffa2ef010
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py
@@ -0,0 +1,107 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for shared Flux utilities (pack/unpack, position IDs).
+
+Tests canonical implementations in flux/utils.py that are shared between
+training and inference code paths.
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.utils import (
+    generate_image_position_ids,
+    pack_latents,
+    unpack_latents,
+)
+from tests.utils import PrimusUT
+
+
+class TestPackUnpackLatents(PrimusUT):
+    """Tests for latent packing/unpacking operations."""
+
+    def test_pack_unpack_roundtrip(self):
+        """Test that pack -> unpack is reversible."""
+        batch_size, channels, height, width = 2, 16, 128, 128
+        original = torch.randn(batch_size, channels, height, width)
+
+        # Pack and unpack
+        packed = pack_latents(original)
+        unpacked = unpack_latents(packed, height, width, vae_scale_factor=1)
+
+        # Should match original
+        assert unpacked.shape == original.shape
+        torch.testing.assert_close(unpacked, original)
+
+
+class TestPositionIDs(PrimusUT):
+    """Tests for position ID generation."""
+
+    def test_generate_image_position_ids_values(self):
+        """Test image position IDs have correct structure."""
+        batch_size, height, width = 2, 8, 8  # Small for inspection
+        device = torch.device("cpu")
+
+        img_ids = generate_image_position_ids(batch_size, height, width, device)
+
+        # Dimension 0 should be all zeros
+        assert (img_ids[:, :, 0] == 0).all()
+
+        # Dimension 1 (row) should range from 0 to height//2-1
+        max_row = img_ids[:, :, 1].max()
+        assert max_row == height // 2 - 1
+
+        # Dimension 2 (col) should range from 0 to width//2-1
+        max_col = img_ids[:, :, 2].max()
+        assert max_col == width // 2 - 1
+
+
+class TestEdgeCases(PrimusUT):
+    """Tests for edge cases and error conditions."""
+
+    def test_pack_latents_minimum_size(self):
+        """Test pack_latents with minimum size (2x2)."""
+        batch_size, channels = 1, 16
+        latents = torch.randn(batch_size, channels, 2, 2)
+
+        packed = pack_latents(latents)
+
+        # 2x2 packed is 1 token
+        assert packed.shape == (batch_size, 1, channels * 4)
+
+    def test_unpack_latents_minimum_size(self):
+        """Test unpack_latents with minimum size."""
+        batch_size, channels = 1, 64
+        packed = torch.randn(batch_size, 1, channels)  # 1 packed token
+
+        # Unpack to 2x2 latent (after VAE scale of 1)
+        unpacked = unpack_latents(packed, 2, 2, vae_scale_factor=1)
+
+        # Output should be (B, C//4, 2, 2)
+        assert unpacked.shape == (batch_size, channels // 4, 2, 2)
+
+
+class TestCompatibility(PrimusUT):
+    """Tests for backward compatibility with existing code."""
+
+    def test_unpack_latents_signature_compatibility(self):
+        """Test unpack_latents works with both training and inference signatures."""
+        packed = torch.randn(2, 4096, 64)
+
+        # Training path (no VAE scaling)
+        unpacked_train = unpack_latents(packed, 128, 128, vae_scale_factor=1)
+        assert unpacked_train.shape == (2, 16, 128, 128)
+
+        # Inference path (VAE scaling)
+        unpacked_infer = unpack_latents(packed, 1024, 1024, vae_scale_factor=8)
+        assert unpacked_infer.shape == (2, 16, 128, 128)  # Same result: 1024/8 = 128
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py b/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py
new file mode 100644
index 000000000..3f9a24d14
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py
@@ -0,0 +1,2316 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+torch.compile graph-break coverage for the FP8 linear / attention path.
+
+Motivation: the @allow_in_graph approach used by older FP8 ops created graph
+break boundaries between AdaLN Triton kernels and FP8 quantize/GEMM ops,
+costing ~144 ms (10.5%) GPU idle per iteration. Eliminating those breaks
+(setup_context Functions + fused amax capture) is what this file guards.
+
+This file has two kinds of tests:
+
+1. Production-guarding tests -- these ``.apply`` the shipped Functions
+   (``OpaqueFP8LinearTensorwiseFunction``, ``DualFP8LinearTensorwiseFunction``,
+   and ``DelayedFP8LinearTensorwiseFunction``) and assert they trace under
+   torch.compile with zero graph breaks and stay numerically equivalent to
+   eager. A regression in the production op fails these directly.
+2. Characterization / feasibility tests -- these use small in-file
+   autograd.Function replicas (``_Level1FP8Linear``, the ``_AllowInGraph*``
+   baselines, ``_DelayedFP8*``) to pin the torch.compile *capabilities* the
+   production design relies on (buffer ``copy_`` inside a compiled forward,
+   tensor hooks firing, ``allow_in_graph`` vs setup_context parity, buffer
+   mutation visibility between compiled calls). They do not guard a production
+   symbol; they document why the production approach is sound and would surface
+   a torch/Inductor behavior change.
+
+Numerical forward/backward correctness of the production Functions lives in
+``test_delayed_fp8_triton_op.py`` and the turbo float8 tests; this file is
+specifically about compile behavior.
+"""
+
+import math
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+import primus_turbo.pytorch as pt
+import torch.nn as nn
+from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
+from primus_turbo.pytorch.core.low_precision import (
+    Float8QuantConfig,
+    Format,
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl
+from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import (
+    FP8UnshardedWeightTensor,
+)
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    DelayedFP8LinearTensorwiseFunction,
+    DualFP8LinearTensorwiseFunction,
+    OpaqueFP8LinearTensorwiseFunction,
+    _extract_fp8_weight,
+)
+from tests.utils import PrimusUT
+
+_has_cuda = torch.cuda.is_available()
+requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required")
+
+_DIM = 64
+_OUT_DIM = 128
+_BATCH = 4
+_GRAN_VALUE = ScalingGranularity.TENSORWISE.value
+_BACKEND_VALUE = BackendType.HIPBLASLT.value
+
+
+# ---------------------------------------------------------------------------
+# Level 1: Bare autograd.Function -- C++ custom ops, primitive args only
+# ---------------------------------------------------------------------------
+
+
+class _Level1FP8Linear(torch.autograd.Function):
+    """FP8 linear without @allow_in_graph, using setup_context + primitives.
+
+    Calls the C++ quantize_fp8_tensorwise and gemm_fp8_impl custom ops
+    directly. No Float8QuantConfig, no FP8UnshardedWeightTensor.
+    """
+
+    @staticmethod
+    def forward(input, weight, fp8_dtype, gran_value, backend_value):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8, a_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(
+            input_2d,
+            fp8_dtype,
+            None,
+        )
+        b_fp8, b_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(
+            weight,
+            fp8_dtype,
+            None,
+        )
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+        return output, a_fp8, a_scale_inv, b_fp8, b_scale_inv
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        input, weight, fp8_dtype, gran_value, backend_value = inputs
+        output_val, a_fp8, a_scale_inv, b_fp8, b_scale_inv = output
+
+        ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.mark_non_differentiable(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.out_dtype = input.dtype
+        ctx.orig_shape = input.shape
+        ctx.fp8_dtype = fp8_dtype
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        if not grad_2d.is_contiguous():
+            grad_2d = grad_2d.contiguous()
+
+        grad_fp8, grad_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(
+            grad_2d,
+            ctx.fp8_dtype,
+            None,
+        )
+        grad_input = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            True,
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        return grad_input, grad_weight, None, None, None
+
+
+class _Level1Module(nn.Module):
+    """Minimal module: LayerNorm + Level 1 FP8 linear."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        out = _Level1FP8Linear.apply(
+            x,
+            self.weight,
+            float8_e4m3,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return out[0]
+
+
+# ---------------------------------------------------------------------------
+# Bonus: LayerNorm + FP8 linear single compiled graph
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestFullGraphIntegration(PrimusUT):
+    """Verify LayerNorm + FP8 linear compiles into a single graph frame."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_single_compiled_frame(self):
+        """LayerNorm → FP8 linear should produce exactly 1 compiled frame."""
+        model = _Level1Module(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        explanation = torch._dynamo.explain(model)(x)
+
+        assert explanation.graph_break_count == 0, (
+            f"Expected single graph (0 breaks), got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+        assert explanation.graph_count == 1, f"Expected 1 graph, got {explanation.graph_count}."
+
+
+# ---------------------------------------------------------------------------
+# Compiled vs eager numerical equivalence (regression for view_as aliasing fix)
+# ---------------------------------------------------------------------------
+
+
+class _OpaqueModule(nn.Module):
+    """Module using OpaqueFP8LinearTensorwiseFunction with pre-extracted weights."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        result = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight,
+            w_fp8,
+            w_scale,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result[0]
+
+
+@requires_cuda
+class TestCompiledVsEagerEquivalence(PrimusUT):
+    """Verify compiled and eager execution produce identical forward and backward results.
+
+    Regression test for the view_as input-output aliasing bug where AOTAutograd's
+    view-replay mechanism could corrupt saved FP8 weight data.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_forward_equivalence(self):
+        """Compiled forward matches eager forward exactly."""
+        model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        eager_out = model(x)
+
+        torch._dynamo.reset()
+        compiled = torch.compile(model)
+        compiled_out = compiled(x)
+
+        torch.testing.assert_close(compiled_out, eager_out, atol=0, rtol=0)
+
+    def test_backward_equivalence(self):
+        """Compiled backward matches eager backward exactly."""
+        model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        eager_out = model(x)
+        eager_out.sum().backward()
+        eager_x_grad = x.grad.clone()
+        eager_w_grad = model.weight.grad.clone()
+
+        x.grad = None
+        model.weight.grad = None
+        torch._dynamo.reset()
+        compiled = torch.compile(model)
+        compiled_out = compiled(x)
+        compiled_out.sum().backward()
+
+        torch.testing.assert_close(x.grad, eager_x_grad, atol=0, rtol=0)
+        torch.testing.assert_close(model.weight.grad, eager_w_grad, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# Multi-step convergence: compiled (setup_context) vs compiled (@allow_in_graph)
+# vs eager, plus BF16 baseline
+# ---------------------------------------------------------------------------
+
+import copy
+
+
+@torch._dynamo.allow_in_graph
+class _AllowInGraphFP8Linear(torch.autograd.Function):
+    """Old-style @allow_in_graph FP8 linear for baseline comparison.
+
+    Quantizes both input and weight inside forward, saves FP8 tensors on ctx.
+    AOTAutograd never traces through this -- it's fully opaque.
+    """
+
+    @staticmethod
+    def forward(ctx, input, weight):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8, a_scale_inv = quantize_fp8(input_2d, float8_e4m3, ScalingGranularity.TENSORWISE)
+        b_fp8, b_scale_inv = quantize_fp8(weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+        ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.out_dtype = out_dtype
+        ctx.orig_shape = orig_shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        grad_fp8, grad_scale_inv = quantize_fp8(grad_2d, float8_e5m2, ScalingGranularity.TENSORWISE)
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            True,
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        return grad_input, grad_weight
+
+
+class _AllowInGraphModule(nn.Module):
+    """Module using old @allow_in_graph FP8 linear."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        return _AllowInGraphFP8Linear.apply(x, self.weight)
+
+
+class _BF16LinearModule(nn.Module):
+    """Pure BF16 linear for baseline comparison (no FP8 quantization)."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear = nn.Linear(dim, out_dim, bias=False)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return self.linear(x)
+
+
+_CONV_DIM = 256
+_CONV_OUT = 512
+_CONV_BATCH = 16
+_N_STEPS = 200
+
+
+def _make_inputs(n, dim, batch, seed=42):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    return [torch.randn(batch, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)]
+
+
+def _run_training_loop(model, inputs, n_steps, lr=1e-3):
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    has_prepare = hasattr(model, "prepare_fp8_weight")
+    if not has_prepare and hasattr(model, "_orig_mod"):
+        has_prepare = hasattr(model._orig_mod, "prepare_fp8_weight")
+    losses = []
+    grad_norms = []
+    for step in range(n_steps):
+        if has_prepare:
+            m = model._orig_mod if hasattr(model, "_orig_mod") else model
+            m.prepare_fp8_weight()
+        optimizer.zero_grad()
+        out = model(inputs[step % len(inputs)])
+        loss = out.sum()
+        loss.backward()
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=float("inf"))
+        grad_norms.append(gn.item())
+        optimizer.step()
+        losses.append(loss.item())
+    return losses, grad_norms
+
+
+def _clone_state(model):
+    return copy.deepcopy(model.state_dict())
+
+
+def _print_comparison(label_a, losses_a, label_b, losses_b, milestones=None):
+    """Print side-by-side loss comparison at milestones for diagnosis."""
+    if milestones is None:
+        milestones = [0, 1, 5, 10, 20, 50, 100, 150, 199]
+    milestones = [m for m in milestones if m < len(losses_a) and m < len(losses_b)]
+    print(f"\n{'Step':>6} | {label_a:>20} | {label_b:>20} | {'Rel Diff':>10}")
+    print("-" * 65)
+    for m in milestones:
+        a, b = losses_a[m], losses_b[m]
+        rel = abs(a - b) / max(abs(a), 1e-12)
+        print(f"{m:>6} | {a:>20.6f} | {b:>20.6f} | {rel:>10.6f}")
+
+
+@requires_cuda
+class TestMultiStepConvergence(PrimusUT):
+    """Multi-step convergence comparison across compilation modes.
+
+    Runs 200 optimizer steps on a small model and compares loss/grad_norm
+    trajectories to isolate whether AOTAutograd backward compilation
+    causes numerical divergence vs the old @allow_in_graph approach.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def _make_fp8_setup_ctx_model(self, state_dict):
+        model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        model.load_state_dict(state_dict)
+        return model
+
+    def _make_fp8_allow_in_graph_model(self, state_dict):
+        model = _AllowInGraphModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        model.load_state_dict(state_dict)
+        return model
+
+    def _make_bf16_model(self, state_dict):
+        model = _BF16LinearModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        # Map weight -> linear.weight for BF16 module
+        bf16_state = {
+            "norm.weight": state_dict["norm.weight"],
+            "norm.bias": state_dict["norm.bias"],
+            "linear.weight": state_dict["weight"],
+        }
+        model.load_state_dict(bf16_state)
+        return model
+
+    def test_compiled_setup_ctx_vs_eager(self):
+        """setup_context compiled should track FP8 eager closely over 200 steps."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_losses, eager_gn = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = self._make_fp8_setup_ctx_model(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, compiled_gn = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("FP8 Eager", eager_losses, "FP8 Compiled(setup_ctx)", compiled_losses)
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"FP8 compiled (setup_context) diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_allow_in_graph_vs_eager(self):
+        """@allow_in_graph compiled should track FP8 eager closely over 200 steps."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_losses, eager_gn = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        aig_model = self._make_fp8_allow_in_graph_model(init_state)
+        aig_model_c = torch.compile(aig_model)
+        aig_losses, aig_gn = _run_training_loop(aig_model_c, inputs, _N_STEPS)
+
+        _print_comparison("FP8 Eager", eager_losses, "FP8 @allow_in_graph", aig_losses)
+
+        final_rel = abs(eager_losses[-1] - aig_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"FP8 @allow_in_graph diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, aig={aig_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_bf16_compiled_vs_eager(self):
+        """BF16 compiled should match BF16 eager closely (no FP8 noise)."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_model = self._make_bf16_model(init_state)
+        eager_losses, _ = _run_training_loop(eager_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = self._make_bf16_model(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("BF16 Eager", eager_losses, "BF16 Compiled", compiled_losses)
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.01, (
+            f"BF16 compiled diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# FSDP2 FP8UnshardedWeightTensor subclass variants
+# ---------------------------------------------------------------------------
+
+
+class _FP8SubclassModule(nn.Module):
+    """Module simulating FSDP2 FP8 all-gather: weight goes through
+    FP8UnshardedWeightTensor subclass before _extract_fp8_weight.
+
+    In production, FSDP2 creates the FP8UnshardedWeightTensor outside the
+    compiled graph (in fsdp_post_all_gather). The compiled forward only sees
+    the subclass as an input and calls _extract_fp8_weight on it. We simulate
+    this by pre-creating the subclass in prepare_fp8_weight() which must be
+    called before each forward step.
+    """
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self._fp8_weight = None
+
+    @torch.no_grad()
+    def prepare_fp8_weight(self):
+        """Simulate FSDP2 fsdp_post_all_gather: create FP8UnshardedWeightTensor
+        outside the compiled graph, just like FSDP2 does before forward."""
+        w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        fp8_config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE)
+        self._fp8_weight = FP8UnshardedWeightTensor(w_fp8, w_scale, torch.bfloat16, fp8_config)
+
+    def forward(self, x):
+        x = self.norm(x)
+        weight_fp8, weight_scale_inv = _extract_fp8_weight(self._fp8_weight, float8_e4m3)
+        result = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight,
+            weight_fp8,
+            weight_scale_inv,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result[0]
+
+
+@torch._dynamo.allow_in_graph
+class _AllowInGraphFP8LinearWithSubclass(torch.autograd.Function):
+    """Old-style @allow_in_graph FP8 linear that receives an
+    FP8UnshardedWeightTensor and extracts data inside the opaque boundary."""
+
+    @staticmethod
+    def forward(ctx, input, weight):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        a_fp8, a_scale_inv = quantize_fp8(input_2d, float8_e4m3, ScalingGranularity.TENSORWISE)
+
+        if isinstance(weight, FP8UnshardedWeightTensor):
+            b_fp8, b_scale_inv = weight.get_fp8_data_and_scale_inv()
+        else:
+            b_fp8, b_scale_inv = quantize_fp8(weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+        ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv)
+        ctx.out_dtype = out_dtype
+        ctx.orig_shape = orig_shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        grad_fp8, grad_scale_inv = quantize_fp8(grad_2d, float8_e5m2, ScalingGranularity.TENSORWISE)
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            b_fp8,
+            b_scale_inv,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_weight = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            True,
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        return grad_input, grad_weight
+
+
+class _AllowInGraphSubclassModule(nn.Module):
+    """Module using old @allow_in_graph FP8 linear with FP8UnshardedWeightTensor
+    passed directly into the opaque boundary (matching old production code).
+
+    Like _FP8SubclassModule, the subclass is created outside compile."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self._fp8_weight = None
+
+    @torch.no_grad()
+    def prepare_fp8_weight(self):
+        w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        fp8_config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE)
+        self._fp8_weight = FP8UnshardedWeightTensor(w_fp8, w_scale, torch.bfloat16, fp8_config)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return _AllowInGraphFP8LinearWithSubclass.apply(x, self._fp8_weight)
+
+
+@requires_cuda
+class TestFP8SubclassConvergence(PrimusUT):
+    """Multi-step convergence with FP8UnshardedWeightTensor subclass path.
+
+    Isolates whether the FSDP2 tensor subclass interaction under torch.compile
+    causes numerical divergence vs eager or the old @allow_in_graph path.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_fp8_subclass_compiled_vs_eager(self):
+        """setup_context + FP8UnshardedWeightTensor: compiled vs eager."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("FP8 Subclass Eager", eager_losses, "FP8 Subclass Compiled", compiled_losses)
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"FP8 subclass compiled diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_fp8_subclass_allow_in_graph_vs_eager(self):
+        """@allow_in_graph + FP8UnshardedWeightTensor: compiled vs eager."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _AllowInGraphSubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _AllowInGraphSubclassModule(_CONV_DIM, _CONV_OUT).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        compiled_model.load_state_dict(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison(
+            "FP8 AIG+Subclass Eager", eager_losses, "FP8 AIG+Subclass Compiled", compiled_losses
+        )
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"FP8 @allow_in_graph+subclass compiled diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_subclass_vs_plain_weight_eager(self):
+        """FP8UnshardedWeightTensor path vs plain quantize path (both eager).
+        Verifies the subclass wrapper doesn't change numerical results."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        plain_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(plain_model)
+
+        plain_losses, _ = _run_training_loop(plain_model, inputs, _N_STEPS)
+
+        subclass_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        subclass_model.load_state_dict(init_state)
+        subclass_losses, _ = _run_training_loop(subclass_model, inputs, _N_STEPS)
+
+        _print_comparison("FP8 Plain Eager", plain_losses, "FP8 Subclass Eager", subclass_losses)
+
+        final_rel = abs(plain_losses[-1] - subclass_losses[-1]) / max(abs(plain_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.01, (
+            f"FP8 subclass path diverged from plain path in eager mode: "
+            f"plain={plain_losses[-1]:.6f}, subclass={subclass_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Dual FP8 linear convergence tests
+# ---------------------------------------------------------------------------
+
+
+class _DualFP8Module(nn.Module):
+    """Module using DualFP8LinearTensorwiseFunction (two GEMMs in one node).
+
+    Simulates the JointSelfAttention dual output projection pattern:
+    input goes through LayerNorm, then two independent FP8 linear projections
+    are computed in a single autograd node. Outputs are summed for loss.
+    """
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        w_fp8_a, w_scale_a = quantize_fp8(self.weight_a, float8_e4m3, ScalingGranularity.TENSORWISE)
+        w_fp8_b, w_scale_b = quantize_fp8(self.weight_b, float8_e4m3, ScalingGranularity.TENSORWISE)
+        result = DualFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight_a,
+            w_fp8_a,
+            w_scale_a,
+            x,
+            self.weight_b,
+            w_fp8_b,
+            w_scale_b,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result[0] + result[1]
+
+
+@torch._dynamo.allow_in_graph
+class _AllowInGraphDualFP8Linear(torch.autograd.Function):
+    """Old-style @allow_in_graph dual FP8 linear for baseline comparison.
+
+    Quantizes both inputs and weights inside forward, saves FP8 tensors on ctx.
+    Matches the old DualFP8LinearTensorwiseFunction behavior before setup_context.
+    """
+
+    @staticmethod
+    def forward(ctx, input_a, weight_a, input_b, weight_b):
+        out_dtype = input_a.dtype
+
+        orig_shape_a = input_a.shape
+        orig_shape_b = input_b.shape
+        input_a_2d = input_a.reshape(-1, input_a.shape[-1])
+        input_b_2d = input_b.reshape(-1, input_b.shape[-1])
+
+        a_fp8_a, a_scale_a = quantize_fp8(input_a_2d, float8_e4m3, ScalingGranularity.TENSORWISE)
+        b_fp8_a, b_scale_a = quantize_fp8(weight_a, float8_e4m3, ScalingGranularity.TENSORWISE)
+        a_fp8_b, a_scale_b = quantize_fp8(input_b_2d, float8_e4m3, ScalingGranularity.TENSORWISE)
+        b_fp8_b, b_scale_b = quantize_fp8(weight_b, float8_e4m3, ScalingGranularity.TENSORWISE)
+
+        output_a = gemm_fp8_impl(
+            a_fp8_a,
+            a_scale_a,
+            False,
+            b_fp8_a,
+            b_scale_a,
+            True,
+            out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        output_b = gemm_fp8_impl(
+            a_fp8_b,
+            a_scale_b,
+            False,
+            b_fp8_b,
+            b_scale_b,
+            True,
+            out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+
+        output_a = output_a.reshape(*orig_shape_a[:-1], output_a.shape[-1])
+        output_b = output_b.reshape(*orig_shape_b[:-1], output_b.shape[-1])
+
+        ctx.save_for_backward(
+            a_fp8_a,
+            a_scale_a,
+            b_fp8_a,
+            b_scale_a,
+            a_fp8_b,
+            a_scale_b,
+            b_fp8_b,
+            b_scale_b,
+        )
+        ctx.out_dtype = out_dtype
+        ctx.orig_shape_a = orig_shape_a
+        ctx.orig_shape_b = orig_shape_b
+        return output_a, output_b
+
+    @staticmethod
+    def backward(ctx, grad_output_a, grad_output_b):
+        if not grad_output_a.is_contiguous():
+            grad_output_a = grad_output_a.contiguous()
+        if not grad_output_b.is_contiguous():
+            grad_output_b = grad_output_b.contiguous()
+
+        (a_fp8_a, a_scale_a, b_fp8_a, b_scale_a, a_fp8_b, a_scale_b, b_fp8_b, b_scale_b) = ctx.saved_tensors
+
+        grad_a_2d = grad_output_a.reshape(-1, grad_output_a.shape[-1])
+        grad_fp8_a, grad_scale_a = quantize_fp8(grad_a_2d, float8_e5m2, ScalingGranularity.TENSORWISE)
+
+        grad_input_a = gemm_fp8_impl(
+            grad_fp8_a,
+            grad_scale_a,
+            False,
+            b_fp8_a,
+            b_scale_a,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        grad_input_a = grad_input_a.reshape(ctx.orig_shape_a)
+
+        grad_weight_a = gemm_fp8_impl(
+            a_fp8_a,
+            a_scale_a,
+            True,
+            grad_fp8_a,
+            grad_scale_a,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+
+        grad_b_2d = grad_output_b.reshape(-1, grad_output_b.shape[-1])
+        grad_fp8_b, grad_scale_b = quantize_fp8(grad_b_2d, float8_e5m2, ScalingGranularity.TENSORWISE)
+
+        grad_input_b = gemm_fp8_impl(
+            grad_fp8_b,
+            grad_scale_b,
+            False,
+            b_fp8_b,
+            b_scale_b,
+            False,
+            ctx.out_dtype,
+            False,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+        grad_input_b = grad_input_b.reshape(ctx.orig_shape_b)
+
+        grad_weight_b = gemm_fp8_impl(
+            a_fp8_b,
+            a_scale_b,
+            True,
+            grad_fp8_b,
+            grad_scale_b,
+            False,
+            ctx.out_dtype,
+            True,
+            granularity=_GRAN_VALUE,
+            default_backend=_BACKEND_VALUE,
+        )
+
+        return grad_input_a, grad_weight_a, grad_input_b, grad_weight_b
+
+
+class _AllowInGraphDualModule(nn.Module):
+    """Module using old @allow_in_graph dual FP8 linear."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        out_a, out_b = _AllowInGraphDualFP8Linear.apply(x, self.weight_a, x, self.weight_b)
+        return out_a + out_b
+
+
+class _TwoSinglesModule(nn.Module):
+    """Module using two separate OpaqueFP8LinearTensorwiseFunction calls.
+
+    Same computation as _DualFP8Module but using two independent autograd
+    nodes instead of one combined node. Used to verify the dual node
+    doesn't introduce numerical differences.
+    """
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        x = self.norm(x)
+        w_fp8_a, w_scale_a = quantize_fp8(self.weight_a, float8_e4m3, ScalingGranularity.TENSORWISE)
+        w_fp8_b, w_scale_b = quantize_fp8(self.weight_b, float8_e4m3, ScalingGranularity.TENSORWISE)
+        result_a = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight_a,
+            w_fp8_a,
+            w_scale_a,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        result_b = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight_b,
+            w_fp8_b,
+            w_scale_b,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result_a[0] + result_b[0]
+
+
+@requires_cuda
+class TestDualFP8Convergence(PrimusUT):
+    """Multi-step convergence for DualFP8LinearTensorwiseFunction.
+
+    The dual function bundles two FP8 GEMMs into a single autograd node
+    (used by JointSelfAttention). This tests whether the combined node
+    under torch.compile causes divergence vs eager or the old @allow_in_graph.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def _make_dual_init_state(self):
+        model = _DualFP8Module(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        return model, _clone_state(model)
+
+    def _load_dual_model(self, cls, state_dict):
+        model = cls(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        model.load_state_dict(state_dict)
+        return model
+
+    def test_dual_compiled_setup_ctx_vs_eager(self):
+        """DualFP8 setup_context compiled should track eager over 200 steps."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model, init_state = self._make_dual_init_state()
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = self._load_dual_model(_DualFP8Module, init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("Dual FP8 Eager", eager_losses, "Dual FP8 Compiled", compiled_losses)
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Dual FP8 compiled (setup_context) diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_dual_allow_in_graph_vs_eager(self):
+        """DualFP8 @allow_in_graph compiled should track eager over 200 steps."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model, init_state = self._make_dual_init_state()
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        aig_model = self._load_dual_model(_AllowInGraphDualModule, init_state)
+        aig_model_c = torch.compile(aig_model)
+        aig_losses, _ = _run_training_loop(aig_model_c, inputs, _N_STEPS)
+
+        _print_comparison("Dual FP8 Eager", eager_losses, "Dual FP8 @allow_in_graph", aig_losses)
+
+        final_rel = abs(eager_losses[-1] - aig_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Dual FP8 @allow_in_graph diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, aig={aig_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_dual_vs_two_singles_eager(self):
+        """Dual node should match two single OpaqueFP8 calls (both eager).
+
+        Verifies that bundling two GEMMs into one autograd node doesn't
+        change numerical results compared to two separate nodes.
+        """
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        dual_model, init_state = self._make_dual_init_state()
+        dual_losses, _ = _run_training_loop(dual_model, inputs, _N_STEPS)
+
+        singles_model = self._load_dual_model(_TwoSinglesModule, init_state)
+        singles_losses, _ = _run_training_loop(singles_model, inputs, _N_STEPS)
+
+        _print_comparison("Dual FP8 Eager", dual_losses, "Two Singles Eager", singles_losses)
+
+        final_rel = abs(dual_losses[-1] - singles_losses[-1]) / max(abs(dual_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.01, (
+            f"Dual FP8 node diverged from two single nodes: "
+            f"dual={dual_losses[-1]:.6f}, singles={singles_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# FP8 Attention Module: QKV projection + flash attention + output projection
+# ---------------------------------------------------------------------------
+
+_ATTN_DIM = 256
+_ATTN_HEADS = 4
+_ATTN_SEQ = 32
+_ATTN_BATCH = 8
+_ATTN_N_STEPS = 200
+
+
+def _hermetic_compile_state():
+    """Reset global compile / kernel-dispatch state for hermetic isolation.
+
+    These convergence tests are sensitive to global state that leaks across
+    tests in the same process:
+      * primus-turbo's kernel-dispatch and origami autotune caches
+        (``GlobalBackendManager.reset()`` clears both), and
+      * torch dynamo's compiled-graph cache (``torch._dynamo.reset()``).
+    Without isolation, kernel/codegen choices made by an earlier test bleed
+    in and shift the compiled-vs-eager numerics, which then compounds over the
+    200-step loop into order-dependent (flaky) failures. We also pin
+    auto-tune off so dispatch deterministically uses the default backend.
+    """
+    GlobalBackendManager.reset()
+    GlobalBackendManager.set_auto_tune(False)
+    torch._dynamo.reset()
+    torch.manual_seed(0)
+    torch.cuda.manual_seed_all(0)
+
+
+class _FP8AttentionModule(nn.Module):
+    """Minimal FP8 attention: QKV projection + flash attention + output projection.
+
+    Mirrors the real Flux attention path:
+    - FP8 linear for QKV (OpaqueFP8LinearTensorwiseFunction)
+    - pt.ops.flash_attn_func for attention
+    - FP8 linear for output projection
+    """
+
+    def __init__(self, dim, num_heads, deterministic=False):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.softmax_scale = 1.0 / math.sqrt(self.head_dim)
+        self.deterministic = deterministic
+
+        self.norm = nn.LayerNorm(dim)
+        self.qkv_weight = nn.Parameter(torch.randn(3 * dim, dim, dtype=torch.bfloat16))
+        self.proj_weight = nn.Parameter(torch.randn(dim, dim, dtype=torch.bfloat16))
+
+    def forward(self, x):
+        B, S, D = x.shape
+        x = self.norm(x)
+
+        qkv_w_fp8, qkv_w_scale = quantize_fp8(self.qkv_weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        qkv = OpaqueFP8LinearTensorwiseFunction.apply(
+            x,
+            self.qkv_weight,
+            qkv_w_fp8,
+            qkv_w_scale,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )[0]
+
+        qkv = qkv.view(B, S, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        attn_out = pt.ops.flash_attn_func(
+            q,
+            k,
+            v,
+            dropout_p=0.0,
+            softmax_scale=self.softmax_scale,
+            causal=False,
+            window_size=(-1, -1),
+            deterministic=self.deterministic,
+            return_lse=False,
+        )
+
+        attn_out = attn_out.reshape(B, S, D)
+
+        proj_w_fp8, proj_w_scale = quantize_fp8(self.proj_weight, float8_e4m3, ScalingGranularity.TENSORWISE)
+        out = OpaqueFP8LinearTensorwiseFunction.apply(
+            attn_out,
+            self.proj_weight,
+            proj_w_fp8,
+            proj_w_scale,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )[0]
+
+        return out
+
+
+def _make_seq_inputs(n, seq, dim, batch, seed=42):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    return [torch.randn(batch, seq, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)]
+
+
+@requires_cuda
+class TestFP8AttentionConvergence(PrimusUT):
+    """FP8 attention convergence: compiled vs eager with flash attention.
+
+    Tests whether torch.compile interacts with FP8 QKV projection +
+    flash attention + FP8 output projection to cause divergence.
+
+    Flash attention is non-deterministic by default, so the non-deterministic
+    test measures whether compile-induced variance stays within the natural
+    variance of flash attention itself.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.fixture(autouse=True)
+    def hermetic_compile_state(self):
+        _hermetic_compile_state()
+        try:
+            yield
+        finally:
+            _hermetic_compile_state()
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_fp8_attention_nondeterministic_compiled_vs_eager(self):
+        """With non-deterministic flash attention, compile variance should be
+        within a small multiple of the natural flash attention variance."""
+        torch._dynamo.reset()
+        inputs = _make_seq_inputs(20, _ATTN_SEQ, _ATTN_DIM, _ATTN_BATCH)
+
+        ref_model = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        init_state = _clone_state(ref_model)
+
+        eager1 = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        eager1.load_state_dict(init_state)
+        eager1_losses, _ = _run_training_loop(eager1, inputs, _ATTN_N_STEPS)
+
+        eager2 = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        eager2.load_state_dict(init_state)
+        eager2_losses, _ = _run_training_loop(eager2, inputs, _ATTN_N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        compiled_model.load_state_dict(init_state)
+        compiled_model = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model, inputs, _ATTN_N_STEPS)
+
+        _print_comparison(
+            "FP8 Attn Eager1",
+            eager1_losses,
+            "FP8 Attn Eager2",
+            eager2_losses,
+        )
+        _print_comparison(
+            "FP8 Attn Eager1",
+            eager1_losses,
+            "FP8 Attn Compiled",
+            compiled_losses,
+        )
+
+        natural_var = abs(eager1_losses[-1] - eager2_losses[-1])
+        compile_var = abs(eager1_losses[-1] - compiled_losses[-1])
+        ref_loss = max(abs(eager1_losses[-1]), 1e-12)
+
+        natural_rel = natural_var / ref_loss
+        compile_rel = compile_var / ref_loss
+
+        print(f"\nNatural variance (eager vs eager): {natural_var:.6f} (rel: {natural_rel:.6f})")
+        print(f"Compile variance (eager vs compiled): {compile_var:.6f} (rel: {compile_rel:.6f})")
+
+        if natural_var < 1e-6:
+            assert compile_rel < 0.05, (
+                f"Flash attn showed no natural variance but compile diverged: "
+                f"compile_rel={compile_rel:.6f}"
+            )
+        else:
+            ratio = compile_var / natural_var
+            print(f"Compile/natural variance ratio: {ratio:.2f}x")
+            assert ratio < 5.0, (
+                f"Compile variance {ratio:.1f}x larger than natural flash attn variance: "
+                f"natural={natural_var:.6f}, compile={compile_var:.6f}"
+            )
+
+
+# ---------------------------------------------------------------------------
+# Mock Megatron DDP: param/grad buffer remapping + backward hooks
+# ---------------------------------------------------------------------------
+
+
+class _MockMegatronDDP(nn.Module):
+    """Simulates Megatron DDP's param/grad buffer remapping and backward hooks.
+
+    Reproduces the essential behavior from distributed_data_parallel.py
+    (lines 419-444) and param_and_grad_buffer.py (lines 760-785) without
+    any Megatron imports:
+
+    1. Remaps param.data into a contiguous param_buffer
+       (like use_distributed_optimizer=True)
+    2. Assigns param.main_grad as views into a contiguous grad_buffer
+    3. Registers gradient accumulator hooks that do:
+       param.main_grad.add_(param.grad.data); param.grad = None
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+        params = [p for p in module.parameters() if p.requires_grad]
+
+        total_numel = sum(p.numel() for p in params)
+        self.param_buffer = torch.zeros(total_numel, dtype=params[0].dtype, device=params[0].device)
+        self.grad_buffer = torch.zeros(total_numel, dtype=params[0].dtype, device=params[0].device)
+
+        offset = 0
+        self.grad_accs = []
+        for param in params:
+            numel = param.numel()
+            new_data = self.param_buffer[offset : offset + numel].view(param.shape)
+            new_data.copy_(param.data)
+            param.data = new_data
+            param.main_grad = self.grad_buffer[offset : offset + numel].view(param.shape)
+            param.grad_added_to_main_grad = False
+            offset += numel
+
+            param_tmp = param.expand_as(param)
+            grad_acc = param_tmp.grad_fn.next_functions[0][0]
+            grad_acc.register_hook(self._make_hook(param))
+            self.grad_accs.append(grad_acc)
+
+    @staticmethod
+    def _make_hook(param):
+        def hook(*unused):
+            if param.grad is not None and not param.grad_added_to_main_grad:
+                param.main_grad.add_(param.grad.data)
+            param.grad = None
+
+        return hook
+
+    def zero_grad_buffer(self):
+        self.grad_buffer.zero_()
+        for p in self.module.parameters():
+            if p.requires_grad:
+                p.grad_added_to_main_grad = False
+
+    def forward(self, *args, **kwargs):
+        return self.module(*args, **kwargs)
+
+
+class _MainGradAdamW(torch.optim.AdamW):
+    """AdamW that reads from param.main_grad instead of param.grad.
+
+    Simulates how Megatron's distributed optimizer operates on main_grad
+    buffers rather than the standard param.grad tensors.
+    """
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if hasattr(p, "main_grad") and p.main_grad is not None:
+                    p.grad = p.main_grad.clone()
+        result = super().step(closure=closure)
+        for group in self.param_groups:
+            for p in group["params"]:
+                p.grad = None
+        return result
+
+
+def _run_mock_ddp_training(module, inputs, n_steps, compile_target=None, lr=1e-3):
+    """Run training with mock Megatron DDP.
+
+    Args:
+        module: The model to wrap with mock DDP.
+        inputs: List of input tensors.
+        n_steps: Number of training steps.
+        compile_target: One of None, "whole_model", "inner_only".
+            None = eager, "whole_model" = compile ddp.forward,
+            "inner_only" = compile ddp.module.forward.
+        lr: Learning rate.
+    """
+    ddp = _MockMegatronDDP(module)
+    optimizer = _MainGradAdamW(ddp.module.parameters(), lr=lr)
+
+    if compile_target == "whole_model":
+        ddp.forward = torch.compile(ddp.forward)
+    elif compile_target == "inner_only":
+        ddp.module.forward = torch.compile(ddp.module.forward)
+
+    losses = []
+    grad_norms = []
+    for step in range(n_steps):
+        ddp.zero_grad_buffer()
+        out = ddp(inputs[step % len(inputs)])
+        loss = out.sum()
+        loss.backward()
+        gn = torch.nn.utils.clip_grad_norm_(
+            [p for p in ddp.module.parameters() if p.main_grad is not None],
+            max_norm=float("inf"),
+        )
+        grad_norms.append(gn.item() if hasattr(gn, "item") else gn)
+        optimizer.step()
+        losses.append(loss.item())
+    return losses, grad_norms
+
+
+@requires_cuda
+class TestMockDDPConvergence(PrimusUT):
+    """Mock Megatron DDP convergence: whole_model compile vs inner-only vs eager.
+
+    Tests whether torch.compile scope interacts with Megatron DDP's backward
+    hooks (param.grad -> main_grad transfer + param.grad = None) to cause
+    numerical divergence.
+
+    The local spec DDP config uses whole_model compile, while the converging
+    TE spec reference uses stack (inner-only) compile. Phase 1 proved that
+    FP8 compiled vs eager is bit-identical without DDP hooks.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_mock_ddp_whole_model_compiled_vs_eager(self):
+        """whole_model compile + mock DDP should track eager + mock DDP."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        eager_model.load_state_dict(init_state)
+        eager_losses, _ = _run_mock_ddp_training(eager_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_losses, _ = _run_mock_ddp_training(
+            compiled_model, inputs, _N_STEPS, compile_target="whole_model"
+        )
+
+        _print_comparison(
+            "MockDDP Eager",
+            eager_losses,
+            "MockDDP WholeModel",
+            compiled_losses,
+        )
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Mock DDP whole_model compile diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_mock_ddp_inner_only_compiled_vs_eager(self):
+        """inner-only compile + mock DDP should track eager + mock DDP."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        eager_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        eager_model.load_state_dict(init_state)
+        eager_losses, _ = _run_mock_ddp_training(eager_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_losses, _ = _run_mock_ddp_training(
+            compiled_model, inputs, _N_STEPS, compile_target="inner_only"
+        )
+
+        _print_comparison(
+            "MockDDP Eager",
+            eager_losses,
+            "MockDDP InnerOnly",
+            compiled_losses,
+        )
+
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Mock DDP inner-only compile diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_mock_ddp_vs_vanilla_eager(self):
+        """Mock DDP eager should produce same results as vanilla eager."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+
+        vanilla_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        vanilla_model.load_state_dict(init_state)
+        vanilla_losses, _ = _run_training_loop(vanilla_model, inputs, _N_STEPS)
+
+        ddp_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        ddp_model.load_state_dict(init_state)
+        ddp_losses, _ = _run_mock_ddp_training(ddp_model, inputs, _N_STEPS)
+
+        _print_comparison(
+            "Vanilla Eager",
+            vanilla_losses,
+            "MockDDP Eager",
+            ddp_losses,
+        )
+
+        final_rel = abs(vanilla_losses[-1] - ddp_losses[-1]) / max(abs(vanilla_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.01, (
+            f"Mock DDP eager diverged from vanilla eager: "
+            f"vanilla={vanilla_losses[-1]:.6f}, ddp={ddp_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Phase 0: Delayed FP8 Scaling Feasibility Tests
+#
+# These tests resolve compile-interaction unknowns required before
+# implementing delayed scaling. Each test answers a specific question
+# about torch.compile behavior with buffer mutations, tensor hooks,
+# and autograd Function side-outputs.
+# ---------------------------------------------------------------------------
+
+_FP8_FWD_MAX = torch.finfo(float8_e4m3).max
+_FP8_BWD_MAX = torch.finfo(float8_e5m2).max
+
+
+# -- Test 0a helpers --------------------------------------------------------
+
+
+class _BufferCopyModule(nn.Module):
+    """Module that copies a computed scalar into a registered buffer.
+
+    Mimics the real delayed scaling pattern where input_amax is computed
+    from the BF16 input (not the FP8 output) and stored via buffer copy_().
+    """
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.register_buffer("staged_amax", torch.tensor(0.0))
+
+    def forward(self, x):
+        x = self.norm(x)
+        input_amax = x.detach().abs().amax().float()
+        result = _Level1FP8Linear.apply(
+            x,
+            self.weight,
+            float8_e4m3,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        self.staged_amax.copy_(input_amax)
+        return result[0]
+
+
+# -- Test 0c helpers --------------------------------------------------------
+
+
+def _grad_amax_hook(grad, buf):
+    """Tensor hook that records gradient amax into a buffer. Returns None
+    to leave the gradient unchanged."""
+    buf.copy_(grad.detach().abs().amax().float())
+    return None
+
+
+class _TensorHookModule(nn.Module):
+    """Module that registers a backward tensor hook to capture grad amax."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.register_buffer("grad_amax", torch.tensor(0.0))
+
+    def forward(self, x):
+        x = self.norm(x)
+        result = _Level1FP8Linear.apply(
+            x,
+            self.weight,
+            float8_e4m3,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        output = result[0]
+        buf = self.grad_amax
+        output.register_hook(lambda g, b=buf: _grad_amax_hook(g, b))
+        return output
+
+
+# -- Test 0d helpers --------------------------------------------------------
+
+
+class _DelayedFP8Linear(torch.autograd.Function):
+    """Minimal delayed-scaling FP8 linear for feasibility testing.
+
+    Takes 3 pre-computed scales (input, weight, gradient).
+    Returns (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv,
+             input_amax, weight_amax) with the last 6 mark_non_differentiable.
+    """
+
+    @staticmethod
+    def forward(
+        input,
+        weight,
+        scale_input,
+        scale_weight,
+        scale_grad,
+        fp8_fwd_dtype,
+        fp8_bwd_dtype,
+        fp8_fwd_max,
+        fp8_bwd_max,
+        gran_value,
+        backend_value,
+    ):
+        out_dtype = input.dtype
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        scale_in_narrow = scale_input.clamp(max=torch.finfo(out_dtype).max).to(out_dtype)
+        a_fp8 = (input_2d * scale_in_narrow).clamp(-fp8_fwd_max, fp8_fwd_max).to(fp8_fwd_dtype)
+        a_scale_inv = (1.0 / scale_input).float()
+
+        scale_w_narrow = scale_weight.clamp(max=torch.finfo(out_dtype).max).to(out_dtype)
+        w_fp8 = (weight * scale_w_narrow).clamp(-fp8_fwd_max, fp8_fwd_max).to(fp8_fwd_dtype)
+        w_scale_inv = (1.0 / scale_weight).float()
+
+        output = gemm_fp8_impl(
+            a_fp8,
+            a_scale_inv,
+            False,
+            w_fp8,
+            w_scale_inv,
+            True,
+            out_dtype,
+            False,
+            granularity=gran_value,
+            default_backend=backend_value,
+        )
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+
+        input_amax = input_2d.detach().abs().amax().float()
+        weight_amax = weight.detach().abs().amax().float()
+
+        a_t_fp8 = a_fp8.t().contiguous()
+        w_t_fp8 = w_fp8.t().contiguous()
+
+        return (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax)
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        (
+            input,
+            weight,
+            scale_input,
+            scale_weight,
+            scale_grad,
+            fp8_fwd_dtype,
+            fp8_bwd_dtype,
+            fp8_fwd_max,
+            fp8_bwd_max,
+            gran_value,
+            backend_value,
+        ) = inputs
+        (output_val, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax) = output
+
+        ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad)
+        ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax)
+        ctx.out_dtype = input.dtype
+        ctx.orig_shape = input.shape
+        ctx.fp8_bwd_dtype = fp8_bwd_dtype
+        ctx.fp8_bwd_max = fp8_bwd_max
+        ctx.gran_value = gran_value
+        ctx.backend_value = backend_value
+
+    @staticmethod
+    def backward(ctx, grad_output, *_):
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()
+        a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad = ctx.saved_tensors
+
+        grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+
+        sg_narrow = scale_grad.clamp(max=torch.finfo(ctx.out_dtype).max).to(ctx.out_dtype)
+        grad_fp8 = (grad_2d * sg_narrow).clamp(-ctx.fp8_bwd_max, ctx.fp8_bwd_max).to(ctx.fp8_bwd_dtype)
+        grad_scale_inv = (1.0 / scale_grad).float()
+
+        grad_input = gemm_fp8_impl(
+            grad_fp8,
+            grad_scale_inv,
+            False,
+            w_t_fp8,
+            w_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+        grad_input = grad_input.reshape(ctx.orig_shape)
+
+        grad_t_fp8 = grad_fp8.t().contiguous()
+        grad_weight = gemm_fp8_impl(
+            grad_t_fp8,
+            grad_scale_inv,
+            False,
+            a_t_fp8,
+            a_scale_inv,
+            True,
+            ctx.out_dtype,
+            False,
+            granularity=ctx.gran_value,
+            default_backend=ctx.backend_value,
+        )
+
+        return (grad_input, grad_weight, None, None, None, None, None, None, None, None, None)
+
+
+class _ProdDelayedModule(nn.Module):
+    """Drives the *production* ``DelayedFP8LinearTensorwiseFunction``.
+
+    Unlike the earlier ``_DelayedFnModule`` (which wrapped the in-file
+    ``_DelayedFP8Linear`` replica), this calls the shipped Function so a
+    regression that introduces a graph break or breaks the fused amax capture
+    in production is actually caught. The function writes the current amaxes
+    in-place into the ``staged_*`` buffers during forward (fused capture), so
+    no side-output ``copy_`` is needed here.
+
+    The scale/amax buffers stay fp32 even though the module runs in bf16: the
+    fused ``atomic_max`` amax kernel only supports fp32, mirroring how the
+    production layers keep these buffers fp32.
+    """
+
+    def __init__(self, dim, out_dim, force_nt):
+        super().__init__()
+        self._force_nt = force_nt
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        for name in ("scale_input", "scale_weight", "scale_grad"):
+            self.register_buffer(name, torch.tensor(1.0, dtype=torch.float32))
+        for name in ("staged_input_amax", "staged_weight_amax", "staged_grad_amax"):
+            self.register_buffer(name, torch.tensor(0.0, dtype=torch.float32))
+
+    def forward(self, x):
+        x = self.norm(x)
+        result = DelayedFP8LinearTensorwiseFunction.apply(
+            x,
+            self.weight,
+            self.scale_input,
+            self.scale_weight,
+            self.scale_grad,
+            self.staged_input_amax,
+            self.staged_weight_amax,
+            self.staged_grad_amax,
+            float8_e4m3,
+            float8_e5m2,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+            self._force_nt,
+        )
+        return result[0]
+
+
+def _build_prod_delayed_module(dim, out_dim, force_nt):
+    """Construct a ``_ProdDelayedModule`` on CUDA/bf16 with fp32 scale/amax
+    buffers (the blanket ``.to(bfloat16)`` would otherwise downcast them and
+    trip the fp32-only fused amax kernel)."""
+    model = _ProdDelayedModule(dim, out_dim, force_nt).to(dtype=torch.bfloat16, device="cuda")
+    for name in (
+        "scale_input",
+        "scale_weight",
+        "scale_grad",
+        "staged_input_amax",
+        "staged_weight_amax",
+        "staged_grad_amax",
+    ):
+        model._buffers[name] = model._buffers[name].float()
+    return model
+
+
+class _DelayedFP8E2EModule(nn.Module):
+    """Full delayed scaling module with scale update loop for e2e tests."""
+
+    def __init__(self, dim, out_dim, use_grad_hook=True, history_len=16):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self._use_grad_hook = use_grad_hook
+        self._fp8_fwd_max = _FP8_FWD_MAX
+        self._fp8_bwd_max = _FP8_BWD_MAX
+
+        self.register_buffer("scale_input", torch.tensor(1.0))
+        self.register_buffer("scale_weight", torch.tensor(1.0))
+        self.register_buffer("scale_grad", torch.tensor(1.0))
+
+        self.register_buffer("amax_history_input", torch.zeros(history_len))
+        self.register_buffer("amax_history_weight", torch.zeros(history_len))
+        self.register_buffer("amax_history_grad", torch.zeros(history_len))
+
+        self.register_buffer("staged_input_amax", torch.tensor(0.0))
+        self.register_buffer("staged_weight_amax", torch.tensor(0.0))
+        self.register_buffer("staged_grad_amax", torch.tensor(0.0))
+
+        self._history_idx = 0
+
+    def update_scales(self):
+        idx = self._history_idx
+        self.amax_history_input[idx] = self.staged_input_amax
+        self.amax_history_weight[idx] = self.staged_weight_amax
+        self.amax_history_grad[idx] = self.staged_grad_amax
+        self._history_idx = (idx + 1) % self.amax_history_input.shape[0]
+
+        self._compute_scale(self.scale_input, self.amax_history_input, self._fp8_fwd_max)
+        self._compute_scale(self.scale_weight, self.amax_history_weight, self._fp8_fwd_max)
+        self._compute_scale(self.scale_grad, self.amax_history_grad, self._fp8_bwd_max)
+
+    def _compute_scale(self, scale_buf, history, fp8_max):
+        amax = history.max()
+        sf = fp8_max / amax.clamp(min=1e-12)
+        sf = torch.where(amax > 0.0, sf, scale_buf)
+        sf = torch.where(torch.isfinite(amax), sf, scale_buf)
+        sf = sf.clamp(max=torch.finfo(torch.float32).max)
+        scale_buf.fill_(sf)
+
+    def update_grad_amax_from_weight_grad(self):
+        if self.weight.grad is not None:
+            self.staged_grad_amax.fill_(self.weight.grad.detach().abs().amax().float())
+
+    def forward(self, x):
+        x = self.norm(x)
+        result = _DelayedFP8Linear.apply(
+            x,
+            self.weight,
+            self.scale_input,
+            self.scale_weight,
+            self.scale_grad,
+            float8_e4m3,
+            float8_e5m2,
+            self._fp8_fwd_max,
+            self._fp8_bwd_max,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        output = result[0]
+        self.staged_input_amax.copy_(result[5])
+        self.staged_weight_amax.copy_(result[6])
+
+        if self._use_grad_hook:
+            buf = self.staged_grad_amax
+            output.register_hook(lambda g, b=buf: _grad_amax_hook(g, b))
+        return output
+
+
+def _run_delayed_training_loop(model, inputs, n_steps, lr=1e-3):
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    mod = model._orig_mod if hasattr(model, "_orig_mod") else model
+    losses, scale_log = [], []
+    for step in range(n_steps):
+        mod.update_scales()
+        optimizer.zero_grad()
+        out = model(inputs[step % len(inputs)])
+        loss = out.sum()
+        loss.backward()
+        if not mod._use_grad_hook:
+            mod.update_grad_amax_from_weight_grad()
+        optimizer.step()
+        losses.append(loss.item())
+        scale_log.append(
+            (
+                mod.scale_input.item(),
+                mod.scale_weight.item(),
+                mod.scale_grad.item(),
+            )
+        )
+    return losses, scale_log
+
+
+# -- Test 0e helpers --------------------------------------------------------
+
+
+class _BufferMutationModule(nn.Module):
+    """Module that uses a buffer in the forward path to test mutation visibility."""
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16))
+        self.register_buffer("scale", torch.tensor(1.0))
+
+    def forward(self, x):
+        x = self.norm(x)
+        x_scaled = x * self.scale
+        result = _Level1FP8Linear.apply(
+            x_scaled,
+            self.weight,
+            float8_e4m3,
+            _GRAN_VALUE,
+            _BACKEND_VALUE,
+        )
+        return result[0]
+
+
+# ---------------------------------------------------------------------------
+# Test 0a: Buffer copy_() of autograd Function side-output
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestDelayedPhase0a_BufferCopy(PrimusUT):
+    """Can a compiled forward mutate a registered buffer via copy_()?
+
+    This is the critical gate test for delayed scaling. If buffer copy_()
+    works inside compiled code, amax can be routed from the autograd
+    Function to the module without graph breaks.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_no_graph_break(self):
+        """torch._dynamo.explain reports zero graph breaks with buffer copy_()."""
+        model = _BufferCopyModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        explanation = torch._dynamo.explain(model)(x)
+
+        assert explanation.graph_break_count == 0, (
+            f"Test 0a: Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+
+    def test_buffer_updated(self):
+        """After compiled forward, staged_amax buffer should be non-zero."""
+        model = _BufferCopyModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        compiled = torch.compile(model)
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        assert model.staged_amax.item() == 0.0
+        _ = compiled(x)
+        assert model.staged_amax.item() > 0.0, "Buffer copy_() had no effect: staged_amax is still 0.0"
+
+    def test_multistep_compiled_vs_eager(self):
+        """Multi-step: compiled with buffer copy_() should track eager."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _BufferCopyModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _BufferCopyModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("BufferCopy Eager", eager_losses, "BufferCopy Compiled", compiled_losses)
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Buffer copy_() compiled diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# Test 0c: Tensor hook with buffer copy_() in backward
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestDelayedPhase0c_TensorHook(PrimusUT):
+    """Does register_hook on an intermediate tensor work under compile,
+    and can the hook callback mutate a registered buffer?
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_no_graph_break(self):
+        """torch._dynamo.explain reports zero graph breaks with tensor hook."""
+        model = _TensorHookModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        explanation = torch._dynamo.explain(model)(x)
+
+        print(f"\nTest 0c: graph_break_count = {explanation.graph_break_count}")
+        if explanation.graph_break_count > 0:
+            print(f"  Break reasons: {explanation.break_reasons}")
+        assert explanation.graph_break_count == 0, (
+            f"Test 0c: Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+
+    def test_hook_fires_and_updates_buffer(self):
+        """After forward + backward, grad_amax buffer should be non-zero."""
+        model = _TensorHookModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        compiled = torch.compile(model)
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+
+        assert model.grad_amax.item() == 0.0
+        out = compiled(x)
+        out.sum().backward()
+        assert model.grad_amax.item() > 0.0, "Tensor hook did not fire: grad_amax is still 0.0"
+
+    def test_multistep_compiled_vs_eager(self):
+        """Multi-step: compiled with tensor hook should track eager."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _TensorHookModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(ref_model)
+        eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _TensorHookModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        compiled_model.load_state_dict(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("TensorHook Eager", eager_losses, "TensorHook Compiled", compiled_losses)
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Tensor hook compiled diverged from eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Test 0d-fn: _DelayedFP8Linear autograd Function (isolation)
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestDelayedProductionFunctionCompile(PrimusUT):
+    """Graph-break + fused-amax-capture contract of the *production*
+    ``DelayedFP8LinearTensorwiseFunction`` under torch.compile.
+
+    This guards the shipped Function directly (both the native and forced-NT
+    backward layouts), so a regression that re-introduces a graph break or
+    breaks the in-place staged-amax capture fails here. Numerical
+    forward/backward correctness of the same Function is covered separately in
+    ``test_delayed_fp8_triton_op.py``.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_function_no_graph_break(self):
+        """torch._dynamo.explain on the production module -> 0 breaks (both arms)."""
+        for force_nt in (True, False):
+            torch._dynamo.reset()
+            model = _build_prod_delayed_module(_DIM, _OUT_DIM, force_nt)
+            x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+            explanation = torch._dynamo.explain(model)(x)
+
+            assert explanation.graph_break_count == 0, (
+                f"force_nt={force_nt}: expected 0 graph breaks, got "
+                f"{explanation.graph_break_count}. Reasons: {explanation.break_reasons}"
+            )
+
+    def test_amax_side_outputs_populated(self):
+        """Compiled forward must capture the current amaxes into the staged
+        buffers in-place (both arms)."""
+        for force_nt in (True, False):
+            torch._dynamo.reset()
+            model = _build_prod_delayed_module(_DIM, _OUT_DIM, force_nt)
+            compiled = torch.compile(model)
+            x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+            assert model.staged_input_amax.item() == 0.0
+            assert model.staged_weight_amax.item() == 0.0
+            _ = compiled(x)
+            assert (
+                model.staged_input_amax.item() > 0.0
+            ), f"force_nt={force_nt}: staged_input_amax not populated"
+            assert (
+                model.staged_weight_amax.item() > 0.0
+            ), f"force_nt={force_nt}: staged_weight_amax not populated"
+            # The captured amax must match the actual input/weight amax (fused
+            # capture is the delayed-scaling contract, not an arbitrary write).
+            xn = model.norm(x)
+            expected_in = xn.detach().float().abs().amax().item()
+            expected_w = model.weight.detach().float().abs().amax().item()
+            assert abs(model.staged_input_amax.item() - expected_in) / max(expected_in, 1e-8) < 1e-2
+            assert abs(model.staged_weight_amax.item() - expected_w) / max(expected_w, 1e-8) < 1e-2
+
+
+# ---------------------------------------------------------------------------
+# Test 0d-e2e: Full delayed module with scale update loop
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestDelayedPhase0d_E2E(PrimusUT):
+    """End-to-end delayed scaling: multi-step training with scale updates.
+
+    Tests the complete interaction between eager scale updates and
+    compiled forward/backward, including scale convergence and
+    loss tracking against dynamic scaling.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_scale_convergence_eager(self):
+        """Scales must change from initial 1.0 within 3 steps (eager).
+
+        Uses history_len=1 (most_recent scaling) so scales respond to
+        the very first recorded amax rather than being dominated by the
+        16-slot history initialization.
+        """
+        model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        _, scale_log = _run_delayed_training_loop(model, inputs, 10)
+
+        for i, (si, sw, sg) in enumerate(scale_log):
+            print(f"  step {i}: scale_input={si:.4f}, " f"scale_weight={sw:.4f}, scale_grad={sg:.4f}")
+
+        s_in_3, s_w_3, s_g_3 = scale_log[3]
+        assert s_in_3 != 1.0, f"scale_input stuck at 1.0 after 3 steps"
+        assert s_w_3 != 1.0, f"scale_weight stuck at 1.0 after 3 steps"
+        assert s_g_3 != 1.0, f"scale_grad stuck at 1.0 after 3 steps"
+
+        for i, (si, sw, sg) in enumerate(scale_log):
+            assert math.isfinite(si) and si > 0, f"scale_input non-finite/negative at step {i}: {si}"
+            assert math.isfinite(sw) and sw > 0, f"scale_weight non-finite/negative at step {i}: {sw}"
+            assert math.isfinite(sg) and sg > 0, f"scale_grad non-finite/negative at step {i}: {sg}"
+
+    def test_compiled_delayed_vs_eager_delayed(self):
+        """Compiled delayed should track eager delayed over 200 steps."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        ref_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        init_state = _clone_state(ref_model)
+        eager_losses, _ = _run_delayed_training_loop(ref_model, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        compiled_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        compiled_model.load_state_dict(init_state)
+        compiled_model_c = torch.compile(compiled_model)
+        compiled_losses, _ = _run_delayed_training_loop(compiled_model_c, inputs, _N_STEPS)
+
+        _print_comparison("Delayed Eager", eager_losses, "Delayed Compiled", compiled_losses)
+        final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.05, (
+            f"Delayed compiled diverged from delayed eager: "
+            f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+    def test_compiled_delayed_vs_compiled_dynamic(self):
+        """Compiled delayed should be within 10% of compiled dynamic."""
+        torch._dynamo.reset()
+        inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH)
+
+        dynamic_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(dynamic_model)
+        dynamic_model_c = torch.compile(dynamic_model)
+        dynamic_losses, _ = _run_training_loop(dynamic_model_c, inputs, _N_STEPS)
+
+        torch._dynamo.reset()
+        delayed_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to(
+            dtype=torch.bfloat16, device="cuda"
+        )
+        delayed_model.norm.load_state_dict(
+            {"weight": init_state["norm.weight"], "bias": init_state["norm.bias"]}
+        )
+        delayed_model.weight.data.copy_(init_state["weight"])
+        delayed_model_c = torch.compile(delayed_model)
+        delayed_losses, _ = _run_delayed_training_loop(delayed_model_c, inputs, _N_STEPS)
+
+        _print_comparison("Dynamic Compiled", dynamic_losses, "Delayed Compiled", delayed_losses)
+        final_rel = abs(dynamic_losses[-1] - delayed_losses[-1]) / max(abs(dynamic_losses[-1]), 1e-12)
+        print(f"\nFinal loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.10, (
+            f"Delayed diverged too far from dynamic: "
+            f"dynamic={dynamic_losses[-1]:.6f}, "
+            f"delayed={delayed_losses[-1]:.6f}, "
+            f"rel_diff={final_rel:.6f}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Test 0e: Buffer mutation between compiled invocations
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestDelayedPhase0e_BufferMutation(PrimusUT):
+    """Verifies that buffer mutations between compiled calls are visible.
+
+    When an eager hook mutates a registered buffer between calls to a
+    compiled forward, the compiled forward must read the new value.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    def test_buffer_mutation_between_compiled_calls(self):
+        """Buffer fill_() between compiled calls must change the output."""
+        model = _BufferMutationModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda")
+        compiled = torch.compile(model)
+        x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        out1 = compiled(x)
+        model.scale.fill_(2.0)
+        out2 = compiled(x)
+
+        assert not torch.allclose(
+            out1, out2
+        ), "Buffer mutation not reflected: compiled forward cached old value"
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py b/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py
new file mode 100644
index 000000000..b7f14eda7
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py
@@ -0,0 +1,343 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Validate the fused Triton kernel for delayed FP8 scale update.
+
+Tests:
+  1. Correctness: fused kernel matches reference Python-loop implementation
+  2. Edge cases: zero amaxes, NaN amaxes, very small amaxes
+  3. History rollover: circular index wraps correctly
+  4. Algorithm variants: 'max' vs 'most_recent'
+
+Run:
+    python -m pytest tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py -v
+or standalone:
+    python tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py
+"""
+
+import pytest
+import torch
+
+triton = pytest.importorskip("triton")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU")
+
+DEVICE = "cuda:0"
+FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max  # 448.0
+FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max  # 57344.0
+FP32_MAX = torch.finfo(torch.float32).max
+
+
+def _reference_update(amax_history, staged_amaxes, scales, fp8_maxes, history_idx, use_max_algo):
+    """Pure-PyTorch reference matching the fused delayed-scale update logic."""
+    T, N, H = amax_history.shape
+    for t in range(T):
+        for m in range(N):
+            new_amax = staged_amaxes[t, m].item()
+            amax_history[t, m, history_idx] = new_amax
+
+            if use_max_algo:
+                amax = amax_history[t, m, :].max().item()
+            else:
+                amax = new_amax
+
+            fp8_max = fp8_maxes[t].item()
+            old_scale = scales[t, m].item()
+
+            if amax > 0.0 and amax == amax:  # positive and not NaN
+                sf = fp8_max / max(amax, 1e-12)
+                sf = min(sf, FP32_MAX)
+            else:
+                sf = old_scale
+
+            scales[t, m] = sf
+
+
+def _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx, use_max_algo):
+    """Call the actual fused Triton kernel."""
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _fused_delayed_scale_update_kernel,
+    )
+
+    _, N, H = amax_history.shape
+    BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024
+
+    _fused_delayed_scale_update_kernel[(N, 3)](
+        amax_history,
+        staged_amaxes,
+        scales,
+        fp8_maxes,
+        history_idx,
+        N=N,
+        H=H,
+        use_max_algo=use_max_algo,
+        BLOCK_H=BLOCK_H,
+        FP32_MAX=FP32_MAX,
+    )
+
+
+@pytest.mark.parametrize(
+    "N,H",
+    [
+        (1, 1),
+        (4, 16),
+        (57, 1024),
+        (10, 128),
+    ],
+)
+@pytest.mark.parametrize("algo", ["max", "most_recent"])
+def test_correctness(N, H, algo):
+    """Fused kernel must produce identical scales and history as reference."""
+    use_max_algo = algo == "max"
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+
+    torch.manual_seed(42)
+    amax_history_ref = torch.rand(3, N, H, dtype=torch.float32, device=DEVICE) * 10.0
+    staged_amaxes = torch.rand(3, N, dtype=torch.float32, device=DEVICE) * 5.0
+    scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+    history_idx = H // 3
+
+    amax_history_fused = amax_history_ref.clone()
+    scales_fused = scales_ref.clone()
+
+    _reference_update(amax_history_ref, staged_amaxes, scales_ref, fp8_maxes, history_idx, use_max_algo)
+    _run_fused_kernel(amax_history_fused, staged_amaxes, scales_fused, fp8_maxes, history_idx, use_max_algo)
+
+    torch.testing.assert_close(
+        amax_history_fused,
+        amax_history_ref,
+        atol=0,
+        rtol=0,
+        msg=f"amax_history mismatch for N={N}, H={H}, algo={algo}",
+    )
+    torch.testing.assert_close(
+        scales_fused, scales_ref, atol=1e-5, rtol=1e-5, msg=f"scales mismatch for N={N}, H={H}, algo={algo}"
+    )
+
+
+def test_zero_amaxes():
+    """When all staged amaxes are zero and history is zero, scales should stay at old_scale."""
+    N, H = 8, 16
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+    amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    staged_amaxes = torch.zeros(3, N, dtype=torch.float32, device=DEVICE)
+    old_scales = torch.full((3, N), 42.0, dtype=torch.float32, device=DEVICE)
+    scales = old_scales.clone()
+
+    _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=True)
+
+    torch.testing.assert_close(
+        scales, old_scales, atol=0, rtol=0, msg="Scales should be unchanged when all amaxes are zero"
+    )
+
+
+def test_nan_amaxes():
+    """NaN amaxes should leave scales unchanged (NaN guard)."""
+    N, H = 4, 8
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+    amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    staged_amaxes = torch.full((3, N), float("nan"), dtype=torch.float32, device=DEVICE)
+    old_scales = torch.full((3, N), 7.0, dtype=torch.float32, device=DEVICE)
+    scales = old_scales.clone()
+
+    _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=True)
+
+    torch.testing.assert_close(
+        scales, old_scales, atol=0, rtol=0, msg="Scales should be unchanged when amaxes are NaN"
+    )
+
+
+def test_very_small_amaxes():
+    """Very small amaxes should produce large scales clamped to FP32_MAX."""
+    N, H = 4, 8
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+    amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    staged_amaxes = torch.full((3, N), 1e-40, dtype=torch.float32, device=DEVICE)
+    scales = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+
+    _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=False)
+
+    assert not scales.isinf().any(), "Scales should not be inf (should be clamped)"
+    assert not scales.isnan().any(), "Scales should not be NaN"
+    assert (scales <= FP32_MAX).all(), "Scales should be <= FP32_MAX"
+
+
+@pytest.mark.parametrize("H", [4, 16, 64])
+def test_history_rollover(H):
+    """Circular index should write to correct position across multiple steps."""
+    N = 3
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+    amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    scales = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+
+    for step in range(H + 5):
+        idx = step % H
+        staged = torch.full((3, N), float(step + 1), dtype=torch.float32, device=DEVICE)
+        _run_fused_kernel(amax_history, staged, scales, fp8_maxes, history_idx=idx, use_max_algo=True)
+
+        assert amax_history[0, 0, idx].item() == float(
+            step + 1
+        ), f"History not written at idx={idx} on step={step}"
+
+
+def test_multi_step_vs_reference():
+    """Run 50 steps of both fused and reference, verify they stay in sync."""
+    N, H = 10, 32
+    use_max_algo = True
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+
+    torch.manual_seed(123)
+    amax_history_ref = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    amax_history_fused = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+    scales_fused = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+
+    for step in range(50):
+        idx = step % H
+        staged = torch.rand(3, N, dtype=torch.float32, device=DEVICE) * (10.0 + step)
+
+        _reference_update(amax_history_ref, staged, scales_ref, fp8_maxes, idx, use_max_algo)
+        _run_fused_kernel(amax_history_fused, staged, scales_fused, fp8_maxes, idx, use_max_algo)
+
+    torch.testing.assert_close(
+        amax_history_fused, amax_history_ref, atol=0, rtol=0, msg="amax_history diverged over 50 steps"
+    )
+    torch.testing.assert_close(
+        scales_fused, scales_ref, atol=1e-5, rtol=1e-5, msg="scales diverged over 50 steps"
+    )
+
+
+def test_registry_integration():
+    """Validate _fast_update_scales_with_history via _DelayedScalingRegistry."""
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _DelayedScalingRegistry,
+        _fast_update_scales_with_history,
+    )
+
+    N, H = 8, 16
+
+    class _FakeModule(torch.nn.Module):
+        """Minimal stand-in for Float8*ParallelLinear with real tensors."""
+
+        def __init__(self, weight_data):
+            super().__init__()
+            self._fp8_fwd_max = FP8_FWD_MAX
+            self._fp8_bwd_max = FP8_BWD_MAX
+            self._amax_compute_algo = "max"
+            self._use_delayed_scaling = True
+            self._first_delayed_step = True
+            self._history_idx = 0
+            self.weight = torch.nn.Parameter(weight_data)
+            self.register_buffer("amax_history_input", torch.zeros(H, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("amax_history_weight", torch.zeros(H, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("amax_history_grad", torch.zeros(H, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("scale_input", torch.tensor(1.0, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("scale_weight", torch.tensor(1.0, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("scale_grad", torch.tensor(1.0, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("staged_input_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("staged_weight_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE))
+            self.register_buffer("staged_grad_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE))
+
+    torch.manual_seed(99)
+    weights = [torch.randn(64, 64, dtype=torch.bfloat16, device=DEVICE) for _ in range(N)]
+    modules = [_FakeModule(weights[i]) for i in range(N)]
+
+    registry = _DelayedScalingRegistry(modules)
+
+    # Staging is now done per-module: each m owns scalar buffers m.staged_input_amax,
+    # m.staged_grad_amax. _fast_update_scales_with_history stacks them into
+    # registry.staged_amaxes_3n[0/2] before launching the Triton kernel; scales
+    # land in registry.scales_3n and are scattered back to m.scale_*.
+    # Parallel pure-Python reference, driven by the EXACT amaxes the kernel used.
+    fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE)
+    amax_history_ref = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE)
+    scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE)
+    history_idx_ref = 0
+    use_max_algo = True  # _FakeModule sets _amax_compute_algo = "max"
+
+    torch.manual_seed(77)
+    for step in range(5):
+        for i in range(N):
+            modules[i].staged_input_amax.fill_(torch.rand(1, device=DEVICE).item() * 10)
+            modules[i].staged_grad_amax.fill_(torch.rand(1, device=DEVICE).item() * 5)
+
+        _fast_update_scales_with_history(registry)
+
+        # The fused kernel only reads registry.staged_amaxes_3n (rows input/weight/grad)
+        # and writes amax_history/scales, so reading it back yields the exact amaxes the
+        # kernel consumed -- including the first-step weight bootstrap from ||weight||_inf
+        # that is staged inside the call. Drive the reference with the same inputs and the
+        # same history index it used this step.
+        staged_used = registry.staged_amaxes_3n.clone()
+        _reference_update(amax_history_ref, staged_used, scales_ref, fp8_maxes, history_idx_ref, use_max_algo)
+        history_idx_ref = (history_idx_ref + 1) % H
+
+    # Registry results must match the reference: history bit-for-bit, scales within fp32 tol.
+    torch.testing.assert_close(
+        registry.amax_history,
+        amax_history_ref,
+        atol=0,
+        rtol=0,
+        msg="registry amax_history diverged from reference loop",
+    )
+    torch.testing.assert_close(
+        registry.scales_3n,
+        scales_ref,
+        atol=1e-5,
+        rtol=1e-5,
+        msg="registry scales diverged from reference loop",
+    )
+
+    # Per-module buffers must reflect the registry-batched results (rows input/weight/grad).
+    for i, m in enumerate(modules):
+        torch.testing.assert_close(m.scale_input, scales_ref[0, i], atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(m.scale_weight, scales_ref[1, i], atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(m.scale_grad, scales_ref[2, i], atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(m.amax_history_input, amax_history_ref[0, i], atol=0, rtol=0)
+        torch.testing.assert_close(m.amax_history_weight, amax_history_ref[1, i], atol=0, rtol=0)
+        torch.testing.assert_close(m.amax_history_grad, amax_history_ref[2, i], atol=0, rtol=0)
+
+    # Sanity: scales finite and positive.
+    assert not registry.scales_3n.isnan().any(), "scales have NaN"
+    assert not registry.scales_3n.isinf().any(), "scales have Inf"
+    assert (registry.scales_3n[0] > 0).all(), "scale_input should be positive"
+    assert (registry.scales_3n[1] > 0).all(), "scale_weight should be positive"
+
+    assert registry._history_idx == 5 % H, f"history_idx should be {5 % H}, got {registry._history_idx}"
+    for m in modules:
+        assert m._history_idx == registry._history_idx, "Per-module _history_idx not synced with registry"
+
+
+if __name__ == "__main__":
+    print("=== Test 1: Correctness (max, various sizes) ===")
+    for N, H in [(1, 1), (4, 16), (57, 1024), (10, 128)]:
+        for algo in ["max", "most_recent"]:
+            test_correctness(N, H, algo)
+            print(f"  N={N}, H={H}, algo={algo}: PASS")
+
+    print("\n=== Test 2: Zero amaxes ===")
+    test_zero_amaxes()
+    print("  PASS")
+
+    print("\n=== Test 3: NaN amaxes ===")
+    test_nan_amaxes()
+    print("  PASS")
+
+    print("\n=== Test 4: Very small amaxes ===")
+    test_very_small_amaxes()
+    print("  PASS")
+
+    print("\n=== Test 5: History rollover ===")
+    for H in [4, 16, 64]:
+        test_history_rollover(H)
+        print(f"  H={H}: PASS")
+
+    print("\n=== Test 6: Multi-step vs reference ===")
+    test_multi_step_vs_reference()
+    print("  PASS")
+
+    print("\n=== Test 7: Registry integration ===")
+    test_registry_integration()
+    print("  PASS")
+
+    print("\nAll tests passed!")
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py b/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py
new file mode 100644
index 000000000..99bc03c96
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py
@@ -0,0 +1,281 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+GPU unit tests for FP8 delayed-scaling state reset during MLPerf warmup.
+
+Validates that ``_reset_fp8_local_spec`` correctly breaks buffer pointers,
+triggering ``_DelayedScalingRegistry`` reinitialisation with
+``_first_step = True`` so that the first real step bootstraps weight amaxes
+from the restored (pre-warmup) weights.
+
+Requires a CUDA device; skipped automatically when unavailable.
+
+Run:
+    python -m pytest tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py -v
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+
+DEVICE = "cuda:0"
+FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max
+FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max
+
+
+class _FakeDelayedModule(torch.nn.Module):
+    """Minimal stand-in for a Float8*ParallelLinear with delayed scaling."""
+
+    def __init__(self, in_features=64, out_features=64, history_len=1):
+        super().__init__()
+        self.weight = torch.nn.Parameter(
+            torch.randn(out_features, in_features, dtype=torch.bfloat16, device=DEVICE)
+        )
+        self._use_delayed_scaling = True
+        self._fp8_fwd_dtype = torch.float8_e4m3fn
+        self._fp8_bwd_dtype = torch.float8_e5m2
+        self._fp8_fwd_max = FP8_FWD_MAX
+        self._fp8_bwd_max = FP8_BWD_MAX
+        self._amax_compute_algo = "most_recent"
+        self._first_delayed_step = True
+        self._history_idx = 0
+        self.config = type(
+            "Config",
+            (),
+            {
+                "fp8_amax_history_len": history_len,
+                "fp8_amax_compute_algo": "most_recent",
+            },
+        )()
+
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _init_delayed_scaling_state,
+        )
+
+        _init_delayed_scaling_state(self)
+        for name in list(self._buffers):
+            buf = self._buffers[name]
+            if buf is not None and buf.device.type == "cpu":
+                self._buffers[name] = buf.to(DEVICE)
+
+
+class TestResetFp8LocalSpecOnDevice:
+    """Verify _reset_fp8_local_spec creates new buffers on the correct device."""
+
+    def test_buffers_on_device_after_reset(self):
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        module = _FakeDelayedModule()
+        model = torch.nn.Sequential(module)
+
+        count = _reset_fp8_local_spec([model])
+        assert count == 1
+
+        for buf_name in (
+            "scale_input",
+            "scale_weight",
+            "scale_grad",
+            "amax_history_input",
+            "amax_history_weight",
+            "amax_history_grad",
+            "staged_input_amax",
+            "staged_grad_amax",
+            "staged_weight_amax",
+        ):
+            buf = module._buffers.get(buf_name)
+            assert buf is not None, f"Buffer {buf_name} missing after reset"
+            assert (
+                buf.device == module.weight.device
+            ), f"Buffer {buf_name} on {buf.device}, expected {module.weight.device}"
+
+    def test_first_delayed_step_reset(self):
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        module = _FakeDelayedModule()
+        module._first_delayed_step = False
+
+        _reset_fp8_local_spec([torch.nn.Sequential(module)])
+
+        assert module._first_delayed_step is True
+
+
+class TestRegistryPointerBreak:
+    """Verify that reset breaks pointer identity with registry global tensors."""
+
+    def test_pointer_mismatch_triggers_reinit(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _DelayedScalingRegistry,
+        )
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        modules = [_FakeDelayedModule() for _ in range(4)]
+        model = torch.nn.Sequential(*modules)
+        registry = _DelayedScalingRegistry(modules)
+
+        # The detection-via-aliasing check in _fast_update_scales_with_history
+        # is keyed on amax_history (the only registry tensor that the modules
+        # hold views into); per-module scale_* buffers have always been
+        # independent scalar tensors. Verify the alias holds pre-reset and is
+        # broken post-reset.
+        old_registry_amax_ptr = registry.amax_history.untyped_storage().data_ptr()
+        old_mod0_amax_ptr = modules[0].amax_history_input.untyped_storage().data_ptr()
+        assert (
+            old_registry_amax_ptr == old_mod0_amax_ptr
+        ), "Sanity: registry.amax_history must alias each module's amax_history_input"
+
+        _reset_fp8_local_spec([model])
+
+        new_mod0_amax_ptr = modules[0].amax_history_input.untyped_storage().data_ptr()
+        assert new_mod0_amax_ptr != old_registry_amax_ptr, (
+            "After reset, module amax_history_input should have different storage "
+            "than the original registry.amax_history (so the next "
+            "_fast_update_scales* call triggers registry.__init__ via the "
+            "data_ptr mismatch check)"
+        )
+
+
+class TestRegistryReinitialisationOnFirstStep:
+    """End-to-end: reset → _fast_update_scales detects mismatch → re-creates registry."""
+
+    def test_fast_path_reinit(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _DelayedScalingRegistry,
+            _fast_update_scales,
+        )
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        modules = [_FakeDelayedModule(history_len=1) for _ in range(4)]
+        model = torch.nn.Sequential(*modules)
+        registry = _DelayedScalingRegistry(modules)
+
+        _fast_update_scales(registry)
+        assert registry._first_step is False
+
+        new_weights = [torch.randn_like(m.weight.data) for m in modules]
+        for m, w in zip(modules, new_weights):
+            m.weight.data.copy_(w)
+
+        _reset_fp8_local_spec([model])
+        _fast_update_scales(registry)
+
+        assert registry._first_step is False, "Should have consumed _first_step"
+        for i, m in enumerate(modules):
+            expected_amax = new_weights[i].abs().amax().float().item()
+            # _fast_update_scales writes per-module weight amaxes to
+            # m.staged_weight_amax in the _first_step bootstrap path; the
+            # registry-batched staged_amaxes_3n is only populated by
+            # _fast_update_scales_with_history, so read from the module here.
+            actual_amax = m.staged_weight_amax.item()
+            assert (
+                abs(actual_amax - expected_amax) < 1e-3
+            ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}"
+
+    def test_history_path_reinit(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _DelayedScalingRegistry,
+            _fast_update_scales_with_history,
+        )
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        H = 16
+        modules = [_FakeDelayedModule(history_len=H) for _ in range(4)]
+        model = torch.nn.Sequential(*modules)
+        registry = _DelayedScalingRegistry(modules)
+
+        _fast_update_scales_with_history(registry)
+        assert registry._first_step is False
+
+        new_weights = [torch.randn_like(m.weight.data) for m in modules]
+        for m, w in zip(modules, new_weights):
+            m.weight.data.copy_(w)
+
+        _reset_fp8_local_spec([model])
+        _fast_update_scales_with_history(registry)
+
+        assert registry._first_step is False
+        for i, m in enumerate(modules):
+            expected_amax = new_weights[i].abs().amax().float().item()
+            # _fast_update_scales_with_history mirrors the staged amaxes into
+            # registry.staged_amaxes_3n (rows 0/1/2 = input/weight/grad), so
+            # both registry.staged_amaxes_3n[1, i] and m.staged_weight_amax
+            # carry the value. Read from the module for consistency with the
+            # fast-path test above.
+            actual_amax = m.staged_weight_amax.item()
+            assert (
+                abs(actual_amax - expected_amax) < 1e-3
+            ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}"
+
+
+class TestScaleLeakagePrevention:
+    """Verify that warmup-phase scales do not leak into post-reset state."""
+
+    def test_scales_recomputed_from_clean_state(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _DelayedScalingRegistry,
+            _fast_update_scales,
+        )
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_local_spec,
+        )
+
+        modules = [_FakeDelayedModule(history_len=1) for _ in range(4)]
+        model = torch.nn.Sequential(*modules)
+        registry = _DelayedScalingRegistry(modules)
+
+        # Staging is now per-module: registry batches them into
+        # registry.staged_amaxes_3n inside _fast_update_scales_with_history,
+        # but _fast_update_scales reads directly from the per-module buffers.
+        # _fast_update_scales also scatters new scales back to per-module
+        # m.scale_* buffers (registry.scales_3n is no longer the
+        # source-of-truth for the fast path), so we read scales from there.
+        for _ in range(5):
+            for i, m in enumerate(modules):
+                m.staged_input_amax.fill_(100.0 + i)
+                m.staged_grad_amax.fill_(50.0 + i)
+            _fast_update_scales(registry)
+
+        warmup_scales = torch.stack(
+            [
+                torch.stack([m.scale_input.clone() for m in modules]),
+                torch.stack([m.scale_weight.clone() for m in modules]),
+                torch.stack([m.scale_grad.clone() for m in modules]),
+            ]
+        )
+        assert (warmup_scales != 1.0).any(), "Scales should have changed during warmup"
+
+        new_weights = [torch.randn_like(m.weight.data) * 0.01 for m in modules]
+        for m, w in zip(modules, new_weights):
+            m.weight.data.copy_(w)
+
+        _reset_fp8_local_spec([model])
+        _fast_update_scales(registry)
+
+        post_reset_scales = torch.stack(
+            [
+                torch.stack([m.scale_input.clone() for m in modules]),
+                torch.stack([m.scale_weight.clone() for m in modules]),
+                torch.stack([m.scale_grad.clone() for m in modules]),
+            ]
+        )
+        assert not torch.equal(
+            warmup_scales, post_reset_scales
+        ), "Scales should differ after reset with new weights"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py b/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py
new file mode 100644
index 000000000..5b37b6281
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py
@@ -0,0 +1,381 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Numerical equivalence tests: TE spec vs Primus-Turbo local spec attention.
+
+Unlike the earlier version of this file, which compared in-file replicas of the
+attention kernels, these tests build the *production* ``core_attention`` modules
+returned by the two spec providers and feed them identical inputs:
+
+  * TE spec    -> ``TEDotProductAttention``     (FusedAttention/CK, SBHD native)
+  * local spec -> ``PrimusTurboLocalAttention`` (Primus-Turbo flash attention)
+
+Both modules are constructed via ``build_module(...)`` with the exact keyword
+arguments Megatron's ``transformer/attention.py`` uses when it wires up
+``submodules.core_attention``, so a regression in how either spec selects or
+configures its attention kernel is caught here.
+
+The (formerly user-visible) contiguous-vs-non-contiguous split is now a
+production-internal detail: ``PrimusTurboLocalAttention`` decides whether to
+force a contiguous BSHD copy from the device capability (``force_contiguous_qkv``
+on gfx942). The four tests below therefore exercise the two real production
+paths in eager and compiled form rather than permutations of hand-rolled
+replicas.
+
+NOTE: this module is GPU-only (``skip_if_no_cuda()``) and additionally requires
+the Megatron + Primus-Turbo stack that only imports inside the ROCm training
+container, so it is validated on the GPU runner lane, not in CPU CI.
+"""
+
+import copy
+import os
+
+os.environ["NVTE_FUSED_ATTN"] = "1"
+os.environ["NVTE_FUSED_ATTN_CK"] = "1"
+os.environ["NVTE_CK_USES_FWD_V3"] = "1"
+os.environ["NVTE_CK_USES_BWD_V3"] = "1"
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+import torch.nn as nn
+
+# NOTE: the diffusion ``conftest`` installs Primus' aiter RTLD_DEEPBIND import
+# hook (``install_aiter_deepbind_hook``) before this module loads, so Turbo's
+# attention backward binds the pinned ``aiter::mha_bwd`` instead of the stale one
+# vendored by transformer_engine below (ROCm/aiter#1332). Without it, attention
+# backward crashes on gfx942/gfx950.
+import transformer_engine.pytorch  # noqa: F401  (import side effects only)
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from primus.backends.megatron.core.extensions.primus_turbo_local_spec import (
+    PrimusTurboLocalAttention,
+)
+from tests.utils import PrimusUT
+
+_has_cuda = torch.cuda.is_available()
+requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required")
+
+# head_dim=128 is required for FAv3 eligibility in AITER.
+# seq=256 > 128 so the Python API fmha_v3_fwd path is taken.
+_DIM = 512
+_HEADS = 4
+_HEAD_DIM = _DIM // _HEADS  # 128
+_SEQ = 256
+_BATCH = 2
+_N_STEPS = 100
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _compute_snr(x: torch.Tensor, y: torch.Tensor) -> float:
+    x, y = x.float(), y.float()
+    signal_power = torch.norm(x).pow(2)
+    noise_power = torch.norm(x - y).pow(2)
+    return 10 * torch.log10(signal_power / (noise_power + 1e-12)).detach().item()
+
+
+def _make_sbhd_inputs(n, seq, dim, batch, seed=42):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    return [torch.randn(seq, batch, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)]
+
+
+def _clone_state(model):
+    return copy.deepcopy(model.state_dict())
+
+
+def _run_training_loop(model, inputs, n_steps, lr=1e-3):
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    losses = []
+    for step in range(n_steps):
+        optimizer.zero_grad()
+        out = model(inputs[step % len(inputs)])
+        loss = out.sum()
+        loss.backward()
+        optimizer.step()
+        losses.append(loss.item())
+    return losses
+
+
+def _print_comparison(label_a, losses_a, label_b, losses_b, milestones=None):
+    if milestones is None:
+        milestones = [0, 1, 5, 10, 20, 50, 99]
+    milestones = [m for m in milestones if m < len(losses_a) and m < len(losses_b)]
+    print(f"\n{'Step':>6} | {label_a:>20} | {label_b:>20} | {'Rel Diff':>10}")
+    print("-" * 65)
+    for m in milestones:
+        a, b = losses_a[m], losses_b[m]
+        rel = abs(a - b) / max(abs(a), 1e-12)
+        print(f"{m:>6} | {a:>20.6f} | {b:>20.6f} | {rel:>10.6f}")
+
+
+def _make_attention_config() -> TransformerConfig:
+    """Minimal TransformerConfig matching the test attention dims.
+
+    ``kv_channels`` is pinned to ``_HEAD_DIM`` and ``softmax_scale`` is left at
+    its default (None) so both production attention modules fall back to the same
+    ``1/sqrt(head_dim)`` scale, exactly as they do in a real Flux layer.
+    """
+    return TransformerConfig(
+        num_layers=1,
+        hidden_size=_DIM,
+        num_attention_heads=_HEADS,
+        kv_channels=_HEAD_DIM,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        tensor_model_parallel_size=1,
+        context_parallel_size=1,
+        bf16=True,
+        params_dtype=torch.bfloat16,
+    )
+
+
+def _build_core_attention(core_attention_cls):
+    """Build a production ``core_attention`` module the way Megatron's
+    ``attention.py`` does (``build_module`` with the same kwargs).
+    """
+    config = _make_attention_config()
+    return build_module(
+        core_attention_cls,
+        config=config,
+        layer_number=1,
+        attn_mask_type=AttnMaskType.no_mask,
+        attention_type="self",
+        softmax_scale=config.softmax_scale,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper that swaps only the production core_attention kernel
+# ---------------------------------------------------------------------------
+
+
+class _SpecAttentionModule(nn.Module):
+    """norm -> qkv -> production core_attention -> proj.
+
+    The norm/qkv/proj linears are identical across instances; the only thing
+    that differs is ``core_attention_cls`` (the class a spec provider's
+    ``core_attention()`` returns), so a numerical divergence is attributable to
+    the attention kernel/wrapper alone.
+    """
+
+    def __init__(self, dim, num_heads, core_attention_cls):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.norm = nn.LayerNorm(dim)
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.proj = nn.Linear(dim, dim, bias=False)
+        self.core = _build_core_attention(core_attention_cls)
+
+    def forward(self, x):
+        S, B, D = x.shape
+        x = self.norm(x)
+        qkv = self.qkv(x).view(S, B, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        # Production core_attention forward: (q, k, v, attention_mask, attn_mask_type)
+        # q/k/v are SBHD [S, B, H, head_dim]; output is merged [S, B, H*head_dim].
+        out = self.core(q, k, v, None, AttnMaskType.no_mask)
+        out = out.reshape(S, B, D)
+        return self.proj(out)
+
+
+def _te_module():
+    return _SpecAttentionModule(_DIM, _HEADS, TEDotProductAttention)
+
+
+def _local_module():
+    return _SpecAttentionModule(_DIM, _HEADS, PrimusTurboLocalAttention)
+
+
+# ---------------------------------------------------------------------------
+# Test class
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+class TestTEvsLocalSpecAttention(PrimusUT):
+    """Numerical equivalence: TE spec vs Primus-Turbo local spec attention.
+
+    Part A: Single-pass SNR (fast, precise kernel-level comparison).
+    Part B: Eager training-loop convergence (detects accumulated drift).
+    Part C: Compiled vs eager for the local spec.
+    Part D: Cross-path end-to-end (TE eager vs local compiled).
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        # init_parallel_state sets up TP=1 parallel state, the TP RNG tracker,
+        # and Megatron global args (enable_turbo_attention_float8=False), all of
+        # which PrimusTurboLocalAttention.__init__ relies on.
+        pass
+
+    def setup_method(self, method):
+        torch._dynamo.reset()
+
+    # -----------------------------------------------------------------------
+    # Part A: Single-forward-pass SNR
+    # -----------------------------------------------------------------------
+
+    def _run_single_pass(self, build_fn):
+        """Run one forward + backward, return output, input grad, and state."""
+        torch.manual_seed(123)
+        torch.cuda.manual_seed_all(123)
+
+        model = build_fn().to(dtype=torch.bfloat16, device="cuda")
+
+        x = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        grad_out = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        out = model(x)
+        out.backward(grad_out)
+        torch.cuda.synchronize()
+        return out.detach(), x.grad.detach(), model.state_dict()
+
+    def _run_single_pass_with_state(self, build_fn, state_dict):
+        """Run one forward + backward with pre-loaded norm/qkv/proj weights."""
+        torch.manual_seed(123)
+        torch.cuda.manual_seed_all(123)
+
+        model = build_fn().to(dtype=torch.bfloat16, device="cuda")
+        own_keys = set(model.state_dict().keys())
+        filtered = {k: v for k, v in state_dict.items() if k in own_keys}
+        model.load_state_dict(filtered, strict=False)
+
+        x = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        grad_out = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda")
+
+        out = model(x)
+        out.backward(grad_out)
+        torch.cuda.synchronize()
+        return out.detach(), x.grad.detach()
+
+    def _compare_snr(self, label, out_a, grad_a, out_b, grad_b, threshold=40.0):
+        out_snr = _compute_snr(out_a, out_b)
+        grad_snr = _compute_snr(grad_a, grad_b)
+        print(f"\n  [{label}] output SNR={out_snr:.2f} dB, grad SNR={grad_snr:.2f} dB")
+        assert out_snr > threshold, f"[{label}] output SNR too low: {out_snr:.2f}"
+        assert grad_snr > threshold, f"[{label}] grad SNR too low: {grad_snr:.2f}"
+
+    def test_single_pass_te_vs_local(self):
+        """TE-spec and local-spec attention produce near-identical single-pass
+        output and input gradients given identical weights/inputs."""
+        out_te, grad_te, state_te = self._run_single_pass(_te_module)
+        out_local, grad_local = self._run_single_pass_with_state(_local_module, state_te)
+        self._compare_snr("TE vs Local", out_te, grad_te, out_local, grad_local)
+
+    # -----------------------------------------------------------------------
+    # Part B: Eager training-loop convergence
+    # -----------------------------------------------------------------------
+
+    def test_eager_te_vs_local(self):
+        """Eager: TE-spec and local-spec attention converge to the same loss."""
+        torch._dynamo.reset()
+        inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH)
+
+        model_te = _te_module().to(dtype=torch.bfloat16, device="cuda")
+        init_state = _clone_state(model_te)
+        losses_te = _run_training_loop(model_te, inputs, _N_STEPS)
+
+        model_local = _local_module().to(dtype=torch.bfloat16, device="cuda")
+        own_keys = set(model_local.state_dict().keys())
+        filtered = {k: v for k, v in init_state.items() if k in own_keys}
+        model_local.load_state_dict(filtered, strict=False)
+        losses_local = _run_training_loop(model_local, inputs, _N_STEPS)
+
+        _print_comparison("TE Eager", losses_te, "Local Eager", losses_local)
+
+        final_rel = abs(losses_te[-1] - losses_local[-1]) / max(abs(losses_te[-1]), 1e-12)
+        print(f"\n  Final loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.02, (
+            f"TE vs Local eager: final loss diverged: "
+            f"te={losses_te[-1]:.6f}, local={losses_local[-1]:.6f}, rel={final_rel:.6f}"
+        )
+
+    # -----------------------------------------------------------------------
+    # Part C/D: Compiled local spec (in-process)
+    #
+    # An earlier revision ran the compiled training in a subprocess because
+    # ``torch.compile`` + ``allow_in_graph(AiterFlashAttnFunc)`` could corrupt
+    # AOT-autograd view-replay metadata when sharing a process with the rest of
+    # the pytest/Megatron harness. That is avoided here by resetting Dynamo
+    # (``torch._dynamo.reset()``) immediately before each compiled build, so the
+    # local-attention graph is traced from a clean state regardless of what
+    # earlier tests in the session compiled. The compiled run builds the
+    # *production* ``PrimusTurboLocalAttention`` via the same ``_local_module``
+    # helper as the eager paths, so the two paths differ only by compilation.
+    # -----------------------------------------------------------------------
+
+    @staticmethod
+    def _run_compiled_local(n_steps: int = 100, init_state=None):
+        """Compile + train the production local-spec attention in-process.
+
+        Resets Dynamo and allows ``AiterFlashAttnFunc`` into the graph, then
+        compiles a fresh ``_local_module`` and runs the training loop. When
+        ``init_state`` is given the model is seeded with it so a comparison run
+        shares identical norm/qkv/proj weights; otherwise the freshly
+        initialized weights are captured. Returns ``(losses, init_state)``.
+        """
+        from primus_turbo.pytorch.ops.attention.flash_attn_interface import (
+            AiterFlashAttnFunc,
+        )
+
+        torch._dynamo.reset()
+        torch._dynamo.allow_in_graph(AiterFlashAttnFunc)
+
+        model = _local_module().to(dtype=torch.bfloat16, device="cuda")
+        if init_state is not None:
+            own_keys = set(model.state_dict().keys())
+            model.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False)
+        captured_state = _clone_state(model)
+
+        inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH)
+        compiled = torch.compile(model)
+        losses = _run_training_loop(compiled, inputs, n_steps)
+        return losses, captured_state
+
+    def test_compiled_local_vs_eager(self):
+        """Compiled local-spec attention converges like eager local-spec."""
+        losses_compiled, init_state = self._run_compiled_local(_N_STEPS)
+
+        torch._dynamo.reset()
+        model_eager = _local_module().to(dtype=torch.bfloat16, device="cuda")
+        own_keys = set(model_eager.state_dict().keys())
+        model_eager.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False)
+        inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH)
+        losses_eager = _run_training_loop(model_eager, inputs, _N_STEPS)
+
+        _print_comparison("Local Eager", losses_eager, "Local Compiled", losses_compiled)
+
+        final_rel = abs(losses_eager[-1] - losses_compiled[-1]) / max(abs(losses_eager[-1]), 1e-12)
+        print(f"\n  Final loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.02, f"Diverged: {final_rel:.6f}"
+
+    def test_te_eager_vs_local_compiled(self):
+        """End-to-end cross-path: TE-spec eager and local-spec compiled, sharing
+        initial weights + inputs, converge to the same loss."""
+        losses_local, init_state = self._run_compiled_local(_N_STEPS)
+
+        torch._dynamo.reset()
+        model_te = _te_module().to(dtype=torch.bfloat16, device="cuda")
+        own_keys = set(model_te.state_dict().keys())
+        model_te.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False)
+        inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH)
+        losses_te = _run_training_loop(model_te, inputs, _N_STEPS)
+
+        _print_comparison("TE Eager", losses_te, "Local Compiled", losses_local)
+
+        final_rel = abs(losses_te[-1] - losses_local[-1]) / max(abs(losses_te[-1]), 1e-12)
+        print(f"\n  Final loss relative diff: {final_rel:.6f}")
+        assert final_rel < 0.02, f"Diverged: {final_rel:.6f}"
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/__init__.py b/tests/unit_tests/backends/megatron/diffusion/training/__init__.py
new file mode 100644
index 000000000..948522468
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for Flux training components.
+
+Tests for trainer classes, forward step, schedulers, and loss computation.
+"""
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/conftest.py b/tests/unit_tests/backends/megatron/diffusion/training/conftest.py
new file mode 100644
index 000000000..b773e6e1b
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/conftest.py
@@ -0,0 +1,36 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Pytest fixtures for diffusion training tests.
+
+Sets up the Primus logger (required by log_rank_0) and re-exports
+the parallel state fixture from the parent conftest.
+"""
+
+import os
+
+import pytest
+
+from primus.core.utils import logger
+from tests.unit_tests.backends.megatron.conftest import (  # noqa: F401
+    init_parallel_state,
+)
+
+
+@pytest.fixture(autouse=True, scope="session")
+def setup_logger():
+    """Initialize Primus logger for tests that use log_rank_0."""
+    logger_cfg = logger.LoggerConfig(
+        exp_root_path=os.environ.get("UT_LOG_PATH", "ut_out"),
+        work_group="develop",
+        user_name="root",
+        exp_name="unittest",
+        module_name="UT-training",
+        file_sink_level="DEBUG",
+        stderr_sink_level="INFO",
+        node_ip="localhost",
+        rank=os.environ.get("RANK", 0),
+        world_size=os.environ.get("WORLD_SIZE", 1),
+    )
+    logger.setup_logger(logger_cfg, is_head=False)
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py
new file mode 100644
index 000000000..487af4193
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py
@@ -0,0 +1,343 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for DiffusionPretrainTrainer.
+
+Tests initialization, setup, data provider delegation, forward step,
+and scheduler lazy initialization.
+"""
+
+import sys
+import types
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer
+
+
+def _build_diffusion_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None, use_mock_data=True):
+    """Helper to build DiffusionPretrainTrainer with stubbed dependencies."""
+
+    # Create a concrete subclass for testing (DiffusionPretrainTrainer is abstract)
+    class ConcreteDiffusionTrainer(DiffusionPretrainTrainer):
+        def create_model(self, pre_process=True, post_process=True):
+            return Mock()
+
+        def create_scheduler(self):
+            return Mock()
+
+        def get_task_encoder(self):
+            return Mock()
+
+    # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports
+    def dummy_base_init(self, backend_args=None, *args, **kwargs):
+        self.backend_args = backend_args
+        self.rank = 0
+        self.world_size = 1
+        self.local_rank = 0
+        self.master_addr = "localhost"
+        self.master_port = 12345
+
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__",
+        dummy_base_init,
+    )
+
+    # Silence logging
+    monkeypatch.setattr(
+        "primus.backends.megatron.diffusion_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+
+    if backend_args is None:
+        backend_args = SimpleNamespace(mock_data=use_mock_data)
+    else:
+        if not hasattr(backend_args, "mock_data"):
+            backend_args.mock_data = use_mock_data
+
+    return ConcreteDiffusionTrainer(backend_args=backend_args)
+
+
+class TestDiffusionPretrainTrainer:
+    """Tests for DiffusionPretrainTrainer."""
+
+    def test_init_with_mock_data_creates_synthetic_provider(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that initialization with mock_data=True creates SyntheticDatasetProvider."""
+        backend_args = SimpleNamespace(mock_data=True, model_type="flux")
+        trainer = _build_diffusion_trainer(monkeypatch, backend_args)
+
+        assert hasattr(trainer, "data_provider")
+        assert trainer.data_provider is not None
+        # Verify it's SyntheticDatasetProvider (check class name)
+        assert "Synthetic" in type(trainer.data_provider).__name__
+
+    def test_setup_calculates_data_parallel_size(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that setup() calculates data_parallel_size when not present."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            world_size=8,
+            tensor_model_parallel_size=2,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=1,
+        )
+        trainer = _build_diffusion_trainer(monkeypatch, backend_args)
+
+        # Mock parent setup and other dependencies
+        setup_calls = []
+
+        def mock_parent_setup(self):
+            setup_calls.append("parent_setup")
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup",
+            mock_parent_setup,
+        )
+
+        # Mock set_primus_global_variables
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.global_vars.set_primus_global_variables",
+            lambda args: None,
+        )
+
+        trainer.setup()
+
+        # Verify data_parallel_size was calculated
+        assert hasattr(trainer.backend_args, "data_parallel_size")
+        assert trainer.backend_args.data_parallel_size == 4  # 8 / (2 * 1 * 1)
+        assert "parent_setup" in setup_calls
+
+    def test_setup_uses_existing_data_parallel_size(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that setup() uses existing data_parallel_size if present."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            data_parallel_size=16,
+            world_size=8,  # This would normally calculate to 4, but we have existing value
+        )
+        trainer = _build_diffusion_trainer(monkeypatch, backend_args)
+
+        # Mock parent setup
+        monkeypatch.setattr(
+            "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup",
+            lambda self: None,
+        )
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.global_vars.set_primus_global_variables",
+            lambda args: None,
+        )
+
+        trainer.setup()
+
+        # Verify existing data_parallel_size was not overwritten
+        assert trainer.backend_args.data_parallel_size == 16
+
+    def test_setup_sets_model_provider(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that setup() sets model_provider with correct signature."""
+        trainer = _build_diffusion_trainer(monkeypatch)
+
+        # Mock create_model
+        mock_model = Mock()
+        trainer.create_model = lambda pre_process=True, post_process=True: mock_model
+
+        # Mock parent setup
+        monkeypatch.setattr(
+            "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup",
+            lambda self: None,
+        )
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.global_vars.set_primus_global_variables",
+            lambda args: None,
+        )
+
+        trainer.setup()
+
+        # Verify model_provider was set
+        assert hasattr(trainer, "model_provider")
+        assert callable(trainer.model_provider)
+
+        # Verify model_provider signature matches Megatron's interface
+        result = trainer.model_provider(
+            pre_process=True,
+            post_process=True,
+            vp_stage=None,
+            config=None,
+            pg_collection=None,
+        )
+        assert result is mock_model
+
+    def test_get_datasets_provider_returns_delegating_function(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that get_datasets_provider() returns a function that delegates to data_provider."""
+        trainer = _build_diffusion_trainer(monkeypatch)
+
+        # Replace data_provider with a mock that has is_distributed as an attribute (not property)
+        mock_dataloaders = [Mock(), Mock(), Mock()]  # train, val, test
+
+        class MockDataProvider:
+            def __init__(self):
+                self.is_distributed = True
+                self.create_dataloaders = Mock(return_value=mock_dataloaders)
+
+        mock_data_provider = MockDataProvider()
+        trainer.data_provider = mock_data_provider
+
+        # Mock get_args
+        mock_args = SimpleNamespace()
+        training_mod = types.SimpleNamespace(get_args=lambda: mock_args)
+        monkeypatch.setitem(sys.modules, "megatron.training", training_mod)
+
+        provider_func = trainer.get_datasets_provider()
+
+        # Verify it's a function
+        assert callable(provider_func)
+
+        # Verify is_distributed flag is set
+        assert provider_func.is_distributed is True
+
+        # Call the provider function
+        train_val_test_num_samples = [100, 10, 10]
+        result = provider_func(train_val_test_num_samples, vp_stage=None)
+
+        # Verify data_provider.create_dataloaders was called
+        mock_data_provider.create_dataloaders.assert_called_once_with(
+            trainer_config=mock_args,
+            train_val_test_num_samples=train_val_test_num_samples,
+            vp_stage=None,
+        )
+
+        # Verify result matches what data_provider returned
+        assert result is mock_dataloaders
+
+    def test_forward_step_calls_flux_forward_step_func(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that forward_step() calls flux_forward_step_func and creates loss function."""
+        trainer = _build_diffusion_trainer(monkeypatch)
+
+        # Mock scheduler
+        mock_scheduler = Mock()
+        trainer._scheduler = mock_scheduler
+
+        # Mock flux_forward_step_func
+        mock_noise_pred = Mock()
+        mock_clean_latents = Mock()
+        mock_noise = Mock()
+        mock_loss_mask = Mock()
+        mock_metrics = {"test_metric": 1.0}
+
+        flux_forward_step_calls = []
+
+        def mock_flux_forward_step(
+            data_iterator,
+            model,
+            scheduler=None,
+            use_guidance_embed=False,
+            guidance_scale=None,
+            timestep_sampler=None,
+            cfg_dropout_prob=0.0,
+            empty_t5_encodings=None,
+            empty_clip_encodings=None,
+            **kwargs,
+        ):
+            flux_forward_step_calls.append(
+                (data_iterator, model, scheduler, use_guidance_embed, guidance_scale)
+            )
+            # 6-tuple per sibling commit 362ee36 (is_validation appended).
+            return (
+                mock_noise_pred,
+                mock_clean_latents,
+                mock_noise,
+                mock_loss_mask,
+                mock_metrics,
+                False,
+            )
+
+        # Patch at the import location (it's imported inside forward_step)
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func",
+            mock_flux_forward_step,
+        )
+
+        # Mock runtime_state
+        trainer.runtime_state = Mock()
+        trainer.runtime_state.update_metrics = Mock()
+
+        mock_data_iterator = Mock()
+        mock_model = Mock()
+
+        output, loss_func = trainer.forward_step(mock_data_iterator, mock_model)
+
+        # Verify flux_forward_step_func was called
+        assert len(flux_forward_step_calls) == 1
+        assert flux_forward_step_calls[0][0] is mock_data_iterator
+        assert flux_forward_step_calls[0][1] is mock_model
+        assert flux_forward_step_calls[0][2] is mock_scheduler
+
+        # Verify values were stored
+        assert trainer._last_clean_latents is mock_clean_latents
+        assert trainer._last_noise is mock_noise
+        assert trainer._last_loss_mask is mock_loss_mask
+
+        # Verify metrics were updated
+        trainer.runtime_state.update_metrics.assert_called_once_with(mock_metrics)
+
+        # Verify output is noise_pred
+        assert output is mock_noise_pred
+
+        # Verify loss_func is callable
+        assert callable(loss_func)
+
+        # Test loss function - mock compute_flow_matching_loss to avoid real computation
+        # It's imported inside the loss function, so patch at the import location
+        mock_loss = Mock()
+        mock_loss.detach.return_value = mock_loss
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.diffusion.loss_computation.compute_flow_matching_loss",
+            lambda *args, **kwargs: mock_loss,
+        )
+
+        loss_result = loss_func(mock_noise_pred, non_loss_data=False)
+        assert len(loss_result) == 2  # (loss, metrics_dict)
+        # Verify the metrics dict has the expected key
+        assert "reduced_train_loss" in loss_result[1]
+
+    def test_forward_step_loss_func_with_non_loss_data(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that loss function returns output_tensor when non_loss_data=True."""
+        trainer = _build_diffusion_trainer(monkeypatch)
+
+        # Mock forward_step components
+        mock_scheduler = Mock()
+        trainer._scheduler = mock_scheduler
+
+        mock_noise_pred = Mock()
+        mock_clean_latents = Mock()
+        mock_noise = Mock()
+        mock_loss_mask = Mock()
+        mock_metrics = {}
+
+        # Patch at the import location (it's imported inside forward_step).
+        # 6-tuple per sibling commit 362ee36 (is_validation appended).
+        monkeypatch.setattr(
+            "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func",
+            lambda *args, **kwargs: (
+                mock_noise_pred,
+                mock_clean_latents,
+                mock_noise,
+                mock_loss_mask,
+                mock_metrics,
+                False,
+            ),
+        )
+
+        trainer.runtime_state = Mock()
+        trainer.runtime_state.update_metrics = Mock()
+
+        output, loss_func = trainer.forward_step(Mock(), Mock())
+
+        # Test loss function with non_loss_data=True
+        result = loss_func(output, non_loss_data=True)
+
+        # Should return output_tensor directly
+        assert result is output
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py
new file mode 100644
index 000000000..88c440b10
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py
@@ -0,0 +1,274 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+End-to-end tests for flux_forward_step_func.
+
+Tests the full forward step with a real Flux 535M model on CUDA,
+covering presampled, resample, validation, and CFG dropout paths.
+"""
+
+import contextlib
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.training.diffusion.forward_step import (
+    flux_forward_step_func,
+)
+from primus.backends.megatron.training.diffusion.schedulers.flow_matching import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    T5_XXL_EMBEDDING_DIM,
+    VAE_LATENT_CHANNELS,
+)
+
+
+def _patch_parallel_state():
+    """Return a stack of mock.patch context managers for parallel state."""
+    return [
+        patch(
+            "megatron.core.parallel_state.get_tensor_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_pipeline_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_data_parallel_rank",
+            return_value=0,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_tensor_model_parallel_rank",
+            return_value=0,
+        ),
+        patch(
+            "megatron.training.get_args",
+            return_value=SimpleNamespace(seed=42),
+        ),
+    ]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestFluxForwardStepE2E:
+    """End-to-end tests for flux_forward_step_func with real model."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    @pytest.fixture
+    def model(self):
+        config = FluxConfig.flux_535m()
+        m = Flux(config).cuda()
+        m.train()
+        return m
+
+    @pytest.fixture
+    def scheduler(self):
+        return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
+
+    def _make_presampled_batch(self, batch_size=2, height=16, width=16):
+        return {
+            "latents": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width),
+            "prompt_embeds": torch.randn(batch_size, 32, T5_XXL_EMBEDDING_DIM),
+            "pooled_prompt_embeds": torch.randn(batch_size, CLIP_L_EMBEDDING_DIM),
+        }
+
+    def _make_resample_batch(self, batch_size=2, height=16, width=16):
+        return {
+            "mean": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width),
+            "logvar": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width),
+            "prompt_embeds": torch.randn(batch_size, 32, T5_XXL_EMBEDDING_DIM),
+            "pooled_prompt_embeds": torch.randn(batch_size, CLIP_L_EMBEDDING_DIM),
+        }
+
+    def test_presampled_path(self, model, scheduler):
+        """Test forward step with pre-encoded latents."""
+        batch = self._make_presampled_batch()
+        data_iterator = iter([batch])
+
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+            result = flux_forward_step_func(
+                data_iterator,
+                model,
+                scheduler=scheduler,
+                step_count=1,
+            )
+
+        noise_pred, clean_latents, noise, loss_mask, metrics, is_validation = result
+        assert noise_pred.shape == (2, VAE_LATENT_CHANNELS, 16, 16)
+        assert clean_latents.shape == noise_pred.shape
+        assert noise.shape == noise_pred.shape
+        assert loss_mask is None
+        assert not is_validation
+        assert metrics["batch_size"] == 2
+        assert metrics["latent_channels"] == VAE_LATENT_CHANNELS
+        assert metrics["image_height"] == 16 * 8
+        assert metrics["image_width"] == 16 * 8
+
+    def test_resample_path(self, model, scheduler):
+        """Test forward step with VAE resample mode."""
+        batch = self._make_resample_batch()
+        # Snapshot mean before the call (the forward step casts and moves it).
+        original_mean = batch["mean"].clone()
+        data_iterator = iter([batch])
+
+        vae_scale = 0.3611
+        vae_shift = 0.1159
+
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+            result = flux_forward_step_func(
+                data_iterator,
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=vae_scale,
+                vae_shift=vae_shift,
+                step_count=1,
+            )
+
+        noise_pred, clean_latents, _, _, _, is_validation = result
+        assert noise_pred.shape == (2, VAE_LATENT_CHANNELS, 16, 16)
+        assert not is_validation
+
+        # Verify reparameterization actually fired: clean_latents must NOT
+        # equal the deterministic scale/shift on `mean` alone (which would
+        # be the result if vae_eps had been dropped or zeroed).
+        expected_no_eps = vae_scale * (original_mean.cuda() - vae_shift)
+        assert not torch.allclose(clean_latents.float(), expected_no_eps.float(), atol=1e-3), (
+            "clean_latents matches mean*(scale-shift) — reparameterization "
+            "noise term `eps * std` appears to have been dropped"
+        )
+
+    def test_validation_with_timestep_key(self, model, scheduler):
+        """Batch with 'timestep' key triggers validation mode."""
+        batch = self._make_presampled_batch()
+        batch["timestep"] = torch.arange(2)
+        data_iterator = iter([batch])
+
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+            result = flux_forward_step_func(
+                data_iterator,
+                model,
+                scheduler=scheduler,
+                step_count=1,
+            )
+
+        _, _, _, _, _, is_validation = result
+        assert is_validation is True
+        # The forward step writes derived timesteps (timestep / 8.0) into the batch.
+        assert "timesteps" in batch
+        assert torch.equal(
+            batch["timesteps"].float().cpu(),
+            torch.arange(2).float() / 8.0,
+        )
+
+    def test_validation_equidistant_injection(self, model, scheduler):
+        """model.eval() without timestep key injects equidistant timesteps."""
+        model.eval()
+        batch_size = 8
+        batch = self._make_presampled_batch(batch_size=batch_size)
+        data_iterator = iter([batch])
+
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+            result = flux_forward_step_func(
+                data_iterator,
+                model,
+                scheduler=scheduler,
+                step_count=1,
+            )
+
+        _, _, _, _, _, is_validation = result
+        assert is_validation is True
+        # Equidistant injection: batch["timestep"] = arange(B) % 8,
+        # batch["timesteps"] = that / 8.0.
+        assert "timestep" in batch
+        assert "timesteps" in batch
+        expected_timestep = torch.arange(batch_size, device="cuda") % 8
+        assert torch.equal(batch["timestep"], expected_timestep), (
+            f"batch['timestep'] expected {expected_timestep.tolist()}, " f"got {batch['timestep'].tolist()}"
+        )
+        assert torch.allclose(
+            batch["timesteps"].float(),
+            expected_timestep.float() / 8.0,
+        )
+
+    def test_cfg_dropout_full(self, model, scheduler):
+        """cfg_dropout_prob=1.0 replaces all text embeddings with empty encodings."""
+        batch_size = 2
+        seq_len = 32
+        batch = self._make_presampled_batch(batch_size=batch_size)
+        data_iterator = iter([batch])
+
+        # Use distinguishable empty encodings (not zeros) so we can verify
+        # the dropout actually substituted them rather than incidentally
+        # matching the original prompt_embeds.
+        empty_t5_value = 7.5
+        empty_clip_value = -3.25
+        empty_t5 = torch.full((seq_len, 1, T5_XXL_EMBEDDING_DIM), empty_t5_value, dtype=torch.float32)
+        empty_clip = torch.full((CLIP_L_EMBEDDING_DIM,), empty_clip_value, dtype=torch.float32)
+
+        # Capture the inputs the model receives by patching its forward.
+        captured = {}
+        original_forward = model.forward
+
+        def capturing_forward(*args, **kwargs):
+            captured["txt"] = kwargs.get("txt").detach().clone()
+            captured["y"] = kwargs.get("y").detach().clone()
+            return original_forward(*args, **kwargs)
+
+        model.forward = capturing_forward
+
+        try:
+            with contextlib.ExitStack() as stack:
+                for p in _patch_parallel_state():
+                    stack.enter_context(p)
+                result = flux_forward_step_func(
+                    data_iterator,
+                    model,
+                    scheduler=scheduler,
+                    cfg_dropout_prob=1.0,
+                    empty_t5_encodings=empty_t5,
+                    empty_clip_encodings=empty_clip,
+                    step_count=1,
+                )
+        finally:
+            model.forward = original_forward
+
+        noise_pred, _, _, _, _, is_validation = result
+        assert noise_pred.shape == (batch_size, VAE_LATENT_CHANNELS, 16, 16)
+        assert not is_validation
+
+        # txt is in (S, B, C) layout after the transpose.
+        # With prob=1.0 every row must be the empty encoding.
+        txt = captured["txt"].float()
+        assert torch.allclose(txt, torch.full_like(txt, empty_t5_value), atol=1e-2), (
+            f"prompt_embeds was not replaced by empty_t5 with cfg_dropout_prob=1.0; "
+            f"sample value={txt.flatten()[0].item():.3f}, expected {empty_t5_value}"
+        )
+
+        y = captured["y"].float()
+        assert torch.allclose(y, torch.full_like(y, empty_clip_value), atol=1e-2), (
+            f"pooled_prompt_embeds was not replaced by empty_clip with prob=1.0; "
+            f"sample value={y.flatten()[0].item():.3f}, expected {empty_clip_value}"
+        )
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py
new file mode 100644
index 000000000..9a965ccfc
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py
@@ -0,0 +1,288 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for FluxPretrainTrainer model creation.
+
+Tests high-value Flux behavior in Primus:
+- backend_args -> FluxConfig mapping/defaults
+- model construction wiring
+- torch_compile settings propagation semantics
+"""
+
+import sys
+import types
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer
+
+
+def _build_flux_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None):
+    """Helper to build FluxPretrainTrainer with stubbed dependencies."""
+
+    # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports
+    def dummy_base_init(self, backend_args=None, *args, **kwargs):
+        self.backend_args = backend_args
+        self.rank = 0
+        self.world_size = 1
+        self.local_rank = 0
+        self.master_addr = "localhost"
+        self.master_port = 12345
+
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__",
+        dummy_base_init,
+    )
+
+    # Stub out DiffusionPretrainTrainer.__init__
+    def dummy_diffusion_init(self, *args, **kwargs):
+        self.backend_args = kwargs.get("backend_args")
+        if self.backend_args is None and args:
+            self.backend_args = args[0]
+        if self.backend_args is None:
+            self.backend_args = SimpleNamespace()
+        self._scheduler = None
+        self.data_provider = Mock()
+
+    monkeypatch.setattr(
+        "primus.backends.megatron.diffusion_trainer.DiffusionPretrainTrainer.__init__",
+        dummy_diffusion_init,
+    )
+
+    # Silence logging
+    monkeypatch.setattr(
+        "primus.backends.megatron.flux_pretrain_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "primus.backends.megatron.diffusion_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+
+    if backend_args is None:
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            guidance_embed=False,
+            guidance_scale=3.5,
+            num_train_timesteps=1000,
+            scheduler_shift=1.0,
+            use_dynamic_shifting=False,
+        )
+
+    return FluxPretrainTrainer(backend_args=backend_args)
+
+
+class TestFluxModelCreation:
+    """Tests for FluxPretrainTrainer model creation."""
+
+    def test_create_model_calls_build_flux_config(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that create_model() calls _build_flux_config_from_yaml()."""
+        trainer = _build_flux_trainer(monkeypatch)
+
+        build_config_calls = []
+        built_config = None
+
+        def tracked_build_config(self):
+            nonlocal built_config
+            build_config_calls.append(1)
+            from primus.backends.megatron.core.models.diffusion.flux.config import (
+                FluxConfig,
+            )
+
+            built_config = FluxConfig.flux_535m()
+            return built_config
+
+        trainer._build_flux_config_from_yaml = tracked_build_config.__get__(trainer, type(trainer))
+
+        # Mock Flux model
+        mock_flux_model = Mock()
+        mock_flux_model.parameters.return_value = []
+        flux_ctor = Mock(return_value=mock_flux_model)
+        monkeypatch.setattr(
+            "primus.backends.megatron.core.models.diffusion.flux.model.Flux",
+            flux_ctor,
+        )
+
+        # Mock get_args
+        mock_args = SimpleNamespace(rank=0)
+        training_mod = types.SimpleNamespace(get_args=lambda: mock_args)
+        monkeypatch.setitem(sys.modules, "megatron.training", training_mod)
+
+        result = trainer.create_model()
+
+        # Verify _build_flux_config_from_yaml was called
+        assert len(build_config_calls) == 1
+        flux_ctor.assert_called_once()
+        assert flux_ctor.call_args.kwargs["backend"] is None
+        assert flux_ctor.call_args.kwargs["config"] is built_config
+        assert result is mock_flux_model
+
+    def test_build_flux_config_from_yaml_extracts_parameters(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that _build_flux_config_from_yaml() extracts parameters from backend_args."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            # Model architecture params
+            num_joint_layers=24,
+            num_single_layers=8,
+            hidden_size=2048,
+            num_attention_heads=16,
+            ffn_hidden_size=8192,
+            # Precision
+            bf16=True,
+            fp16=False,
+            params_dtype=torch.bfloat16,
+            # Transformer impl
+            transformer_impl="local",
+        )
+
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        # Verify config attributes were set
+        assert config.num_joint_layers == 24
+        assert config.num_single_layers == 8
+        assert config.hidden_size == 2048
+        assert config.num_attention_heads == 16
+        assert config.ffn_hidden_size == 8192
+        assert config.bf16 is True
+        assert config.fp16 is False
+        assert config.params_dtype == torch.bfloat16
+        assert config.transformer_impl == "local"
+
+    def test_build_flux_config_from_yaml_torch_compile_settings(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that torch_compile settings are extracted from backend_args."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            torch_compile=SimpleNamespace(
+                enable=True,
+                backend="inductor",
+                mode="reduce-overhead",
+                fullgraph=True,
+            ),
+        )
+
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        # Verify torch_compile settings
+        assert config.enable_torch_compile is True
+        assert config.torch_compile_backend == "inductor"
+        assert config.torch_compile_mode == "reduce-overhead"
+        assert config.torch_compile_fullgraph is True
+
+    def test_create_model_sets_torch_compile_on_args(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that create_model() sets torch_compile attributes on args."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            torch_compile=SimpleNamespace(
+                enable=True,
+                backend="inductor",
+                mode="reduce-overhead",
+                fullgraph=True,
+            ),
+        )
+
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        # Mock Flux model
+        mock_flux_model = Mock()
+        mock_flux_model.parameters.return_value = []
+        monkeypatch.setattr(
+            "primus.backends.megatron.core.models.diffusion.flux.model.Flux",
+            lambda *args, **kwargs: mock_flux_model,
+        )
+
+        # Mock get_args
+        mock_args = SimpleNamespace(rank=0)
+        training_mod = types.SimpleNamespace(get_args=lambda: mock_args)
+        monkeypatch.setitem(sys.modules, "megatron.training", training_mod)
+
+        trainer.create_model()
+
+        # Verify torch_compile attributes were set on args
+        assert mock_args.enable_torch_compile is True
+        assert mock_args.torch_compile_backend == "inductor"
+        assert mock_args.torch_compile_mode == "reduce-overhead"
+        assert mock_args.torch_compile_fullgraph is True
+
+    def test_build_flux_config_from_yaml_fp8_settings(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that FP8 fields are extracted from backend_args into FluxConfig."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            fp8="e4m3",
+            fp8_recipe="tensorwise",
+            fp8_margin=0,
+            fp8_amax_history_len=1,
+            fp8_amax_compute_algo="most_recent",
+            fp8_wgrad=True,
+            fp8_dot_product_attention=False,
+            fp8_multi_head_attention=False,
+        )
+
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.fp8 == "e4m3"
+        assert config.fp8_recipe == "tensorwise"
+        assert config.fp8_margin == 0
+        assert config.fp8_amax_history_len == 1
+        assert config.fp8_amax_compute_algo == "most_recent"
+        assert config.fp8_wgrad is True
+        assert config.fp8_dot_product_attention is False
+        assert config.fp8_multi_head_attention is False
+
+    def test_create_model_does_not_overwrite_existing_torch_compile_attrs(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """Test that create_model() doesn't overwrite existing torch_compile attributes."""
+        backend_args = SimpleNamespace(
+            mock_data=True,
+            torch_compile=SimpleNamespace(
+                enable=True,
+                backend="inductor",
+                mode="reduce-overhead",
+                fullgraph=True,
+            ),
+        )
+
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        # Mock Flux model
+        mock_flux_model = Mock()
+        mock_flux_model.parameters.return_value = []
+        monkeypatch.setattr(
+            "primus.backends.megatron.core.models.diffusion.flux.model.Flux",
+            lambda *args, **kwargs: mock_flux_model,
+        )
+
+        # Mock get_args with existing attributes
+        mock_args = SimpleNamespace(
+            rank=0,
+            enable_torch_compile=False,  # Already set
+            torch_compile_backend="aot_eager",  # Already set
+        )
+        training_mod = types.SimpleNamespace(get_args=lambda: mock_args)
+        monkeypatch.setitem(sys.modules, "megatron.training", training_mod)
+
+        trainer.create_model()
+
+        # Verify existing attributes were NOT overwritten
+        assert mock_args.enable_torch_compile is False
+        assert mock_args.torch_compile_backend == "aot_eager"
+        # New attributes should still be set
+        assert mock_args.torch_compile_mode == "reduce-overhead"
+        assert mock_args.torch_compile_fullgraph is True
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py
new file mode 100644
index 000000000..c7c1035d7
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py
@@ -0,0 +1,219 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FluxPretrainTrainer.
+
+Tests initialization, configuration, method overrides, scheduler setup,
+and CFG dropout encoding discovery.
+"""
+
+import os
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer
+
+
+def _build_flux_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None):
+    """Helper to build FluxPretrainTrainer with stubbed dependencies."""
+
+    # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports
+    def dummy_init(self, backend_args: any = None):
+        self.backend_args = backend_args
+
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__",
+        dummy_init,
+    )
+
+    # Silence logging
+    monkeypatch.setattr(
+        "primus.backends.megatron.flux_pretrain_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "primus.backends.megatron.diffusion_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+
+    if backend_args is None:
+        backend_args = SimpleNamespace(
+            guidance_embed=False,
+            guidance_scale=3.5,
+            num_train_timesteps=1000,
+            scheduler_shift=1.0,
+            use_dynamic_shifting=False,
+        )
+
+    return FluxPretrainTrainer(backend_args=backend_args)
+
+
+class TestFluxPretrainTrainer:
+    """Tests for FluxPretrainTrainer."""
+
+    def test_flux_trainer_scheduler_creation(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that scheduler is created correctly."""
+        backend_args = SimpleNamespace(
+            guidance_embed=False,
+            guidance_scale=3.5,
+            num_train_timesteps=1000,
+            scheduler_shift=1.0,
+            use_dynamic_shifting=False,
+        )
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        # Create scheduler
+        scheduler = trainer.create_scheduler()
+
+        # Should be FlowMatchEulerDiscreteScheduler
+        from primus.backends.megatron.training.diffusion.schedulers import (
+            FlowMatchEulerDiscreteScheduler,
+        )
+
+        assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler)
+        assert scheduler.num_train_timesteps == 1000
+        assert scheduler.shift == 1.0
+        assert scheduler.use_dynamic_shifting is False
+
+    def test_flux_trainer_scheduler_with_dynamic_shifting(self, monkeypatch: pytest.MonkeyPatch):
+        """Test scheduler creation with dynamic shifting enabled."""
+        backend_args = SimpleNamespace(
+            guidance_embed=False,
+            guidance_scale=3.5,
+            num_train_timesteps=1000,
+            scheduler_shift=1.0,
+            use_dynamic_shifting=True,
+            base_shift=0.5,
+            max_shift=1.15,
+        )
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        scheduler = trainer.create_scheduler()
+
+        assert scheduler.use_dynamic_shifting is True
+        assert scheduler.base_shift == 0.5
+        assert scheduler.max_shift == 1.15
+
+    def test_flux_trainer_scheduler_lazy_initialization(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that scheduler is lazily initialized via property."""
+        trainer = _build_flux_trainer(monkeypatch)
+
+        # Scheduler should not be created yet
+        assert trainer._scheduler is None
+
+        # Accessing scheduler property should create it
+        scheduler1 = trainer.scheduler
+        assert scheduler1 is not None
+        assert trainer._scheduler is scheduler1
+
+        # Accessing again should return same instance
+        scheduler2 = trainer.scheduler
+        assert scheduler2 is scheduler1
+
+
+def _create_encoding_files(directory):
+    """Create minimal t5_empty.npy and clip_empty.npy in a directory."""
+    os.makedirs(directory, exist_ok=True)
+    np.save(os.path.join(directory, "t5_empty.npy"), np.zeros((1, 256, 4096), dtype=np.float32))
+    np.save(os.path.join(directory, "clip_empty.npy"), np.zeros((1, 768), dtype=np.float32))
+
+
+class TestCFGDropoutDiscovery:
+    """Tests for CFG dropout empty encoding discovery and error handling."""
+
+    def test_discover_encodings_inside_data_path(self, monkeypatch, tmp_path):
+        """Auto-discovers encodings at {data_path}/empty_encodings/."""
+        monkeypatch.setattr(
+            "primus.backends.megatron.flux_pretrain_trainer.log_rank_0",
+            lambda *a, **kw: None,
+        )
+        enc_dir = tmp_path / "empty_encodings"
+        _create_encoding_files(str(enc_dir))
+
+        params = SimpleNamespace(data_path=str(tmp_path))
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result == str(enc_dir)
+
+    def test_discover_encodings_alongside_data_path(self, monkeypatch, tmp_path):
+        """Falls back to {data_path}/../empty_encodings/ when not inside."""
+        monkeypatch.setattr(
+            "primus.backends.megatron.flux_pretrain_trainer.log_rank_0",
+            lambda *a, **kw: None,
+        )
+        data_dir = tmp_path / "dataset"
+        data_dir.mkdir()
+        enc_dir = tmp_path / "empty_encodings"
+        _create_encoding_files(str(enc_dir))
+
+        params = SimpleNamespace(data_path=str(data_dir))
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result == str(enc_dir)
+
+    def test_discover_encodings_explicit_path(self, tmp_path):
+        """Explicit empty_encodings_path takes priority."""
+        explicit_dir = tmp_path / "custom_encodings"
+        _create_encoding_files(str(explicit_dir))
+
+        inside_dir = tmp_path / "data" / "empty_encodings"
+        _create_encoding_files(str(inside_dir))
+
+        params = SimpleNamespace(
+            data_path=str(tmp_path / "data"),
+            empty_encodings_path=str(explicit_dir),
+        )
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result == str(explicit_dir)
+
+    def test_discover_encodings_returns_none_when_missing(self, tmp_path):
+        """Returns None when no encoding files exist anywhere."""
+        params = SimpleNamespace(data_path=str(tmp_path))
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result is None
+
+    def test_discover_encodings_returns_none_for_partial_files(self, tmp_path):
+        """Returns None when only one of the two required files exists."""
+        enc_dir = tmp_path / "empty_encodings"
+        enc_dir.mkdir()
+        np.save(str(enc_dir / "t5_empty.npy"), np.zeros((1, 256, 4096), dtype=np.float32))
+
+        params = SimpleNamespace(data_path=str(tmp_path))
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result is None
+
+    def test_discover_encodings_handles_list_data_path(self, monkeypatch, tmp_path):
+        """Handles data_path passed as a list (takes first element)."""
+        monkeypatch.setattr(
+            "primus.backends.megatron.flux_pretrain_trainer.log_rank_0",
+            lambda *a, **kw: None,
+        )
+        enc_dir = tmp_path / "empty_encodings"
+        _create_encoding_files(str(enc_dir))
+
+        params = SimpleNamespace(data_path=[str(tmp_path), "/nonexistent"])
+        result = FluxPretrainTrainer._discover_empty_encodings(params)
+
+        assert result == str(enc_dir)
+
+    def test_cfg_dropout_raises_when_no_encodings_found(self, monkeypatch, tmp_path):
+        """Real data without encodings raises FileNotFoundError with instructions."""
+        backend_args = SimpleNamespace(
+            cfg_dropout_prob=0.1,
+            mock_data=False,
+            data_path=str(tmp_path),
+            tensor_model_parallel_size=1,
+        )
+
+        with pytest.raises(FileNotFoundError, match="primus data diffusion-encoded"):
+            _build_flux_trainer(monkeypatch, backend_args)
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py
new file mode 100644
index 000000000..288acf498
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py
@@ -0,0 +1,219 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for prepare_flux_latents function.
+
+Tests the higher-level prepare_flux_latents function that orchestrates
+latent preparation for training. Lower-level pack/unpack utilities are
+tested in test_flux_utils.py.
+"""
+
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.training.diffusion.forward_step import (
+    prepare_flux_latents,
+)
+from primus.backends.megatron.training.diffusion.schedulers.flow_matching import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    DEFAULT_NUM_TRAIN_TIMESTEPS,
+    VAE_LATENT_CHANNELS,
+)
+from tests.unit_tests.backends.megatron.diffusion.helpers import (
+    assert_tensor_shape,
+    create_mock_latents,
+)
+from tests.utils import PrimusUT
+
+
+class TestPrepareFluxLatents(PrimusUT):
+    """Tests for prepare_flux_latents function."""
+
+    def test_prepare_flux_latents_shapes(self):
+        """Test that prepare_flux_latents produces correct output shapes."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        batch_size = 2
+        channels = VAE_LATENT_CHANNELS  # VAE output channels
+        height, width = 128, 128
+
+        latents = create_mock_latents(batch_size, height, width, channels=channels, seed=42)
+        h_ids, w_ids = height // 2, width // 2
+        img_ids = torch.zeros(batch_size, h_ids * w_ids, 3)
+
+        (
+            clean_latents,
+            noise,
+            packed_noisy_latents,
+            img_ids_out,
+            guidance_vec,
+            timesteps,
+            sigma_1d,
+        ) = prepare_flux_latents(
+            latents=latents,
+            scheduler=scheduler,
+            img_ids=img_ids,
+        )
+
+        # Check shapes
+        assert_tensor_shape(clean_latents, (batch_size, channels, height, width), "clean_latents")
+        assert_tensor_shape(noise, (batch_size, channels, height, width), "noise")
+
+        # Packed latents: (B, H*W/4, C*4)
+        expected_packed_shape = (
+            batch_size,
+            (height // 2) * (width // 2),
+            channels * 4,
+        )
+        assert_tensor_shape(packed_noisy_latents, expected_packed_shape, "packed_noisy_latents")
+
+        # Timesteps: (B,)
+        assert len(timesteps) == batch_size
+
+        # Sigma: (B,) in [0, 1]
+        assert_tensor_shape(sigma_1d, (batch_size,), "sigma_1d")
+
+        # Guidance should be None when not used
+        assert guidance_vec is None
+
+    def test_prepare_flux_latents_with_guidance(self):
+        """Test prepare_flux_latents with guidance embedding enabled."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        batch_size = 2
+        latents = create_mock_latents(batch_size, 128, 128, seed=42)
+        img_ids = torch.zeros(batch_size, 64 * 64, 3)
+        guidance_scale = 3.5
+
+        (
+            clean_latents,
+            noise,
+            packed_noisy_latents,
+            img_ids_out,
+            guidance_vec,
+            timesteps,
+            sigma_1d,
+        ) = prepare_flux_latents(
+            latents=latents,
+            scheduler=scheduler,
+            img_ids=img_ids,
+            guidance_scale=guidance_scale,
+            use_guidance_embed=True,
+        )
+
+        # Guidance vector should be created
+        assert guidance_vec is not None
+        assert_tensor_shape(guidance_vec, (batch_size,), "guidance_vec")
+
+        # All guidance values should be the scale
+        assert torch.allclose(guidance_vec, torch.tensor(guidance_scale))
+
+    def test_prepare_flux_latents_timestep_sampling(self):
+        """Test that timesteps are sampled correctly."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        batch_size = 8
+        latents = create_mock_latents(batch_size, 128, 128, seed=42)
+        img_ids = torch.zeros(batch_size, 64 * 64, 3)
+
+        (
+            clean_latents,
+            noise,
+            packed_noisy_latents,
+            img_ids_out,
+            guidance_vec,
+            timesteps,
+            sigma_1d,
+        ) = prepare_flux_latents(
+            latents=latents,
+            scheduler=scheduler,
+            img_ids=img_ids,
+        )
+
+        # Check timesteps are valid
+        assert len(timesteps) == batch_size
+
+        # All timesteps should be in scheduler's timestep range
+        for t in timesteps:
+            assert t >= 0 and t <= scheduler.num_train_timesteps
+
+    def test_prepare_flux_latents_deterministic_with_seed(self):
+        """Test that results are deterministic with fixed random seed."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        latents = create_mock_latents(2, 64, 64, seed=42)
+        img_ids = torch.zeros(2, 32 * 32, 3)
+
+        # Set seed and run
+        torch.manual_seed(12345)
+        result1 = prepare_flux_latents(latents, scheduler, img_ids)
+
+        # Set same seed and run again
+        torch.manual_seed(12345)
+        result2 = prepare_flux_latents(latents, scheduler, img_ids)
+
+        # Results should be identical
+        assert torch.allclose(result1[2], result2[2]), "Results should be deterministic with same seed"
+
+    def test_timestep_sampling_range(self):
+        """Test that timestep sampling covers full [0, 1] range with logit-normal."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        # Sample many timesteps to check distribution
+        num_samples = 1000
+        batch_size = 100
+        latents = create_mock_latents(batch_size, 64, 64, seed=42)
+
+        all_timesteps = []
+        for i in range(num_samples // batch_size):
+            torch.manual_seed(42 + i)
+            result = prepare_flux_latents(
+                latents=latents,
+                scheduler=scheduler,
+                img_ids=None,  # Test with auto-generation
+            )
+            # Use sigma_1d (index 6) which is already in [0, 1] range
+            all_timesteps.append(result[6].float())
+
+        all_timesteps = torch.cat(all_timesteps)
+
+        # Check that we sample across the full range
+        min_t = all_timesteps.min().item()
+        max_t = all_timesteps.max().item()
+
+        # With logit-normal, we should get values close to 0 and 1
+        assert min_t < 0.1, f"Min timestep {min_t} should be < 0.1 (covers early diffusion)"
+        assert max_t > 0.9, f"Max timestep {max_t} should be > 0.9 (covers late diffusion)"
+
+        # Check distribution is reasonable (not all in one region)
+        mean_t = all_timesteps.mean().item()
+        assert 0.3 < mean_t < 0.7, f"Mean timestep {mean_t} should be roughly centered"
+
+    def test_img_ids_optional_parameter(self):
+        """Test that img_ids can be None and will be generated automatically."""
+        scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS)
+
+        batch_size = 2
+        height, width = 128, 128
+        latents = create_mock_latents(batch_size, height, width, seed=42)
+
+        # Call without img_ids (should auto-generate)
+        result = prepare_flux_latents(
+            latents=latents,
+            scheduler=scheduler,
+            img_ids=None,  # Test auto-generation
+        )
+
+        # Should succeed and return valid img_ids
+        img_ids_out = result[3]
+        assert img_ids_out is not None
+
+        # Check shape is correct
+        expected_shape = (batch_size, (height // 2) * (width // 2), 3)
+        assert_tensor_shape(img_ids_out, expected_shape, "auto_generated_img_ids")
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py
new file mode 100644
index 000000000..8a7d7635c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py
@@ -0,0 +1,173 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for the model.training gate on DiffusionPretrainTrainer._forward_step_count.
+
+Contract under test: when ``model.training`` is False (the trainer is running
+an evaluation forward pass), the counter must NOT advance and the per-step
+CUDA RNG reseed must NOT fire. Otherwise validation steps shift the training
+RNG sequence by ``eval_iters * num_microbatches`` per ``--eval-interval``
+window, defeating the goal of isolating the training RNG from unrelated
+forward passes.
+
+These tests pair with
+``tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py``,
+which covers lazy initialization from checkpoint state and the
+``step_count`` flow into ``flux_forward_step_func``. This file is the
+single source of truth for the eval-skip contract — if the gate is reverted,
+the three ``TestForwardStepCountTrainingGate`` tests below must fail.
+"""
+
+from unittest.mock import Mock, patch
+
+import torch
+
+from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer
+
+
+class _ConcreteDiffusionTrainer(DiffusionPretrainTrainer):
+    """Minimal concrete subclass: abstract methods as no-ops."""
+
+    def create_model(self, *args, **kwargs):
+        return None
+
+    def create_scheduler(self, *args, **kwargs):
+        return None
+
+    def get_task_encoder(self, *args, **kwargs):
+        return None
+
+
+def _make_trainer():
+    """Construct a trainer with ``__init__`` bypassed.
+
+    Sets only the attributes ``forward_step`` reads. Lazy-init is forced to
+    "already initialized" because that surface is covered by
+    ``test_diffusion_trainer_forward_step_count.py::TestForwardStepCountLazyInit``.
+    """
+    trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer)
+    trainer._forward_step_count = 0
+    trainer._forward_step_count_initialized = True
+    trainer._scheduler = None
+
+    class _FakeRuntimeState:
+        def update_metrics(self, metrics):
+            pass
+
+    trainer.runtime_state = _FakeRuntimeState()
+    return trainer
+
+
+def _patch_flux_forward_step_func():
+    """Patch the real func with a recorder that captures ``step_count``."""
+    captured = []
+
+    def _recorder(*args, **kwargs):
+        captured.append(kwargs.get("step_count"))
+        t = torch.zeros(1)
+        return t, t, t, None, {}, False
+
+    return (
+        patch(
+            "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func",
+            side_effect=_recorder,
+        ),
+        captured,
+    )
+
+
+class TestForwardStepCountTrainingGate:
+    """Direct coverage of the ``if model.training:`` gate on counter advance."""
+
+    def test_eval_does_not_advance_counter(self):
+        trainer = _make_trainer()
+        patcher, captured = _patch_flux_forward_step_func()
+        eval_model = Mock()
+        eval_model.training = False
+
+        with patcher:
+            trainer.forward_step(data_iterator=None, model=eval_model)
+
+        assert trainer._forward_step_count == 0
+        # Counter unchanged → step_count passed to the step func is the pre-call value.
+        assert captured == [0]
+
+    def test_train_advances_counter_by_one(self):
+        trainer = _make_trainer()
+        patcher, captured = _patch_flux_forward_step_func()
+        train_model = Mock()
+        train_model.training = True
+
+        with patcher:
+            trainer.forward_step(data_iterator=None, model=train_model)
+
+        assert trainer._forward_step_count == 1
+        assert captured == [1]
+
+    def test_eval_between_train_does_not_shift_sequence(self):
+        """Eval batches interleaved with training must NOT change the
+        ``step_count`` that subsequent training steps see."""
+        trainer = _make_trainer()
+        patcher, captured = _patch_flux_forward_step_func()
+
+        train_model = Mock()
+        train_model.training = True
+        eval_model = Mock()
+        eval_model.training = False
+
+        with patcher:
+            trainer.forward_step(data_iterator=None, model=train_model)  # 0 -> 1
+            trainer.forward_step(data_iterator=None, model=eval_model)  # stays 1
+            trainer.forward_step(data_iterator=None, model=eval_model)  # stays 1
+            trainer.forward_step(data_iterator=None, model=train_model)  # 1 -> 2
+
+        assert trainer._forward_step_count == 2
+        # Critical assertion: the 4th training step sees step_count=2, the
+        # same value it would see in the absence of the two eval calls. If
+        # the gate were missing, the sequence would be [1, 2, 3, 4].
+        assert captured == [1, 1, 1, 2]
+
+
+class TestPerStepReseedFormula:
+    """Positive control: when the counter advances, the per-step reseed
+    formula in ``flux_forward_step_func`` is computed correctly.
+
+    This is a focused regression test for the seed expression
+    ``(seed + 100 * dp_rank) * 10000 + step_count) % (2**63)``.
+    It deliberately reaches into the underlying step function rather than
+    going through the trainer so the formula is testable in isolation
+    from the gate logic."""
+
+    def test_reseed_uses_step_count_in_formula(self):
+        from primus.backends.megatron.training.diffusion import forward_step as fwd_mod
+
+        fake_args = Mock()
+        fake_args.seed = 7
+
+        train_model = Mock()
+        train_model.training = True
+
+        with patch.object(torch.cuda, "manual_seed") as mock_seed, patch(
+            "megatron.training.get_args", return_value=fake_args
+        ), patch("megatron.core.parallel_state.get_data_parallel_rank", return_value=3):
+            # flux_forward_step_func runs significant downstream logic that
+            # would require a full Megatron-Flux harness. We only need to
+            # verify the reseed call. Suppress the post-reseed failure so the
+            # mock.assert below still observes the manual_seed call.
+            try:
+                fwd_mod.flux_forward_step_func(
+                    data_iterator=None,
+                    model=train_model,
+                    scheduler=Mock(),
+                    per_step_rng_reseed=True,
+                    step_count=42,
+                )
+            except Exception:
+                # Expected: downstream code needs real data/model. The reseed
+                # call we care about happens before any of that — see the
+                # assert below.
+                pass
+
+        expected_seed = ((7 + 100 * 3) * 10000 + 42) % (2**63)
+        mock_seed.assert_called_once_with(expected_seed)
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py
new file mode 100644
index 000000000..e5e5e9c39
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for loss computation utilities.
+
+These tests verify the correctness of loss computation functions,
+ensuring they produce expected results and maintain backward compatibility.
+"""
+
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.training.diffusion.loss_computation import (
+    compute_flow_matching_loss,
+)
+from tests.utils import PrimusUT
+
+
+class TestFlowMatchingLoss(PrimusUT):
+    """Tests for flow matching loss computation."""
+
+    def test_basic_computation(self):
+        """Test flow matching loss matches manual calculation."""
+        prediction = torch.randn(2, 16, 64, 64)
+        clean = torch.randn(2, 16, 64, 64)
+        noise = torch.randn(2, 16, 64, 64)
+
+        loss = compute_flow_matching_loss(prediction, clean, noise)
+
+        # Manual calculation
+        target = noise - clean
+        expected = torch.nn.functional.mse_loss(prediction.float(), target.float())
+
+        assert torch.allclose(loss, expected)
+        assert loss.dim() == 0  # Scalar
+
+    def test_loss_with_partial_mask(self):
+        """Test loss computation with partial masking."""
+        prediction = torch.randn(2, 16, 64, 64)
+        clean = torch.randn(2, 16, 64, 64)
+        noise = torch.randn(2, 16, 64, 64)
+
+        # Create partial mask (half valid)
+        loss_mask = torch.zeros(2, 16, 64, 64)
+        loss_mask[:, :, :32, :] = 1.0  # First half valid
+
+        loss_with_mask = compute_flow_matching_loss(prediction, clean, noise, loss_mask)
+
+        # Should be different from unmasked loss
+        loss_without_mask = compute_flow_matching_loss(prediction, clean, noise)
+
+        assert not torch.allclose(loss_with_mask, loss_without_mask)
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py b/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py
new file mode 100644
index 000000000..5c2c45e64
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py
@@ -0,0 +1,115 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for overridable methods in MegatronPretrainTrainer.
+
+Tests verify that get_forward_step() and get_datasets_provider() methods
+can be overridden by subclasses while maintaining backward compatibility.
+"""
+
+import sys
+import types
+from types import SimpleNamespace
+from typing import Any, List, Tuple
+
+import pytest
+import torch
+
+from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer
+
+# train() imports the real megatron.core/training stack (transitively primus_turbo/TE), which
+# initializes CUDA at import time. Gate the two tests that call train() so they run only on GPU.
+requires_gpu = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="train() imports real megatron.core/training which initialize CUDA at import",
+)
+
+
+def _build_trainer(monkeypatch: pytest.MonkeyPatch) -> MegatronPretrainTrainer:
+    """Helper to build MegatronPretrainTrainer with a stubbed MegatronBaseTrainer."""
+
+    # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports/patching.
+    def dummy_init(self, backend_args: Any = None):
+        self.backend_args = backend_args
+
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__",
+        dummy_init,
+    )
+
+    # Silence logging from the trainer module.
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_pretrain_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+
+    backend_args = SimpleNamespace()
+
+    return MegatronPretrainTrainer(backend_args=backend_args)
+
+
+class TestMegatronPretrainTrainerOverridableMethods:
+    """Tests for overridable methods in MegatronPretrainTrainer."""
+
+    @requires_gpu
+    def test_backward_compatibility_gpt_training_still_works(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that existing GPT training still works (backward compatibility)."""
+        trainer = _build_trainer(monkeypatch)
+
+        calls: List[Tuple[tuple, dict]] = []
+
+        # Use the real megatron modules that train() imports; patch only the integration seams.
+        import megatron.core.pipeline_parallel as mpp
+        import megatron.training
+        import megatron.training.inprocess_restart as inprocess_restart
+        import megatron.training.training as mt_training
+        from megatron.core.enums import ModelType
+
+        def fake_pretrain(*args, **kwargs):
+            raise AssertionError("fake_pretrain was called directly; expected wrapped_pretrain to be used")
+
+        def wrapped_pretrain(*args, store=None, **kwargs):
+            calls.append((args, {"store": store, **kwargs}))
+
+        def fake_maybe_wrap(fn):
+            assert fn is fake_pretrain
+            return wrapped_pretrain, "STORE"
+
+        monkeypatch.setattr(megatron.training, "pretrain", fake_pretrain)
+        monkeypatch.setattr(inprocess_restart, "maybe_wrap_for_inprocess_restart", fake_maybe_wrap)
+
+        # train() reassigns get_forward_backward_func on both real modules; snapshot both through
+        # monkeypatch so they are restored and the mutation does not leak into later tests.
+        monkeypatch.setattr(mpp, "get_forward_backward_func", mpp.get_forward_backward_func)
+        monkeypatch.setattr(mt_training, "get_forward_backward_func", mt_training.get_forward_backward_func)
+
+        model_provider = object()
+        monkeypatch.setattr(
+            "primus.core.utils.import_utils.get_model_provider",
+            lambda *args, **kwargs: model_provider,
+        )
+
+        # pretrain_gpt is a Megatron example script, not importable here -> mock it via sys.modules.
+        train_valid_test_datasets_provider = SimpleNamespace(is_distributed=False)
+        pretrain_gpt_mod = types.SimpleNamespace(
+            forward_step="FORWARD_STEP",
+            train_valid_test_datasets_provider=train_valid_test_datasets_provider,
+        )
+        monkeypatch.setitem(sys.modules, "pretrain_gpt", pretrain_gpt_mod)
+
+        # Execute training
+        trainer.train()
+
+        # Verify backward compatibility: should work exactly as before
+        assert train_valid_test_datasets_provider.is_distributed is True
+        assert len(calls) == 1
+        (args, kwargs) = calls[0]
+
+        # Positional arguments should match expected GPT training pattern
+        assert args[0] is train_valid_test_datasets_provider
+        assert args[1] is model_provider
+        assert args[2] is ModelType.encoder_or_decoder
+        assert args[3] == "FORWARD_STEP"
+
+        assert kwargs == {"store": "STORE"}
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py b/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py
new file mode 100644
index 000000000..953879aa8
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py
@@ -0,0 +1,257 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for diffusion schedulers.
+
+These tests validate the FlowMatchEulerDiscreteScheduler and other scheduler
+implementations used by Flux and future diffusion models.
+
+The tests follow NeMo's scheduler API:
+    - scale_noise(): Add noise during training (forward process)
+    - set_timesteps(): Setup inference schedule
+    - step(): Perform denoising step (reverse process)
+"""
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.training.diffusion.schedulers import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    BATCH_SIZE_PAIR,
+    BATCH_SIZE_QUAD,
+    BATCH_SIZE_SINGLE,
+    DEFAULT_SHIFT,
+    IMG_SIZE_MICRO,
+    IMG_SIZE_MINI,
+    IMG_SIZE_SMALL,
+    TENSOR_CHANNELS_RGB,
+    TRAINING_STEPS_MODERATE,
+)
+from tests.utils import PrimusUT
+
+
+class TestFlowMatchEulerDiscreteScheduler(PrimusUT):
+    """Tests for FlowMatchEulerDiscreteScheduler."""
+
+    # ========================================================================
+    # Initialization Tests
+    # ========================================================================
+
+    def test_scheduler_initialization_with_shift(self):
+        """Test scheduler with custom shift parameter."""
+        scheduler = FlowMatchEulerDiscreteScheduler(shift=2.0)
+
+        assert scheduler.shift == 2.0
+
+        # With higher shift, sigmas should be different
+        scheduler_no_shift = FlowMatchEulerDiscreteScheduler(shift=DEFAULT_SHIFT)
+        assert not torch.allclose(scheduler.sigmas, scheduler_no_shift.sigmas)
+
+    # ========================================================================
+    # Sigma Schedule Tests
+    # ========================================================================
+
+    def test_sigma_range(self):
+        """Test sigma values are in valid range [0, 1]."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        assert (scheduler.sigmas >= 0.0).all()
+        assert (scheduler.sigmas <= 1.0).all()
+
+        # Sigmas should be monotonically decreasing
+        assert (scheduler.sigmas[:-1] >= scheduler.sigmas[1:]).all()
+
+    def test_sigma_schedule_with_shift(self):
+        """Test shift parameter affects sigma distribution."""
+        scheduler_shift_1 = FlowMatchEulerDiscreteScheduler(shift=DEFAULT_SHIFT)
+        scheduler_shift_2 = FlowMatchEulerDiscreteScheduler(shift=2.0)
+
+        # Higher shift moves distribution towards higher values
+        mean_shift_1 = scheduler_shift_1.sigmas.mean()
+        mean_shift_2 = scheduler_shift_2.sigmas.mean()
+
+        assert mean_shift_2 > mean_shift_1
+
+    # ========================================================================
+    # Training Tests (Forward Process - scale_noise)
+    # ========================================================================
+
+    def test_scale_noise_shape(self):
+        """Test scale_noise produces correct shape."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        batch_size = BATCH_SIZE_QUAD  # Quad sample tests
+        channels = 64
+        height = IMG_SIZE_SMALL  # Small size for standard tests
+        width = IMG_SIZE_SMALL
+
+        sample = torch.randn(batch_size, channels, height, width)
+        noise = torch.randn(batch_size, channels, height, width)
+        timesteps = scheduler.timesteps[[100, 200, 300, 400]]
+
+        noisy = scheduler.scale_noise(sample, timesteps, noise)
+
+        assert noisy.shape == sample.shape
+        assert noisy.dtype == sample.dtype
+
+    def test_scale_noise_flow_matching_interpolation(self):
+        """Test flow matching interpolation: x_t = (1-σ)*x_0 + σ*noise."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        sample = torch.ones(
+            BATCH_SIZE_SINGLE,
+            TENSOR_CHANNELS_RGB,
+            IMG_SIZE_MICRO,
+            IMG_SIZE_MICRO,
+        )
+        noise = torch.zeros(
+            BATCH_SIZE_SINGLE,
+            TENSOR_CHANNELS_RGB,
+            IMG_SIZE_MICRO,
+            IMG_SIZE_MICRO,
+        )
+
+        # At σ≈0 (t=0), should be mostly sample
+        timesteps = scheduler.timesteps[-1:]  # Last timestep (smallest sigma)
+        noisy = scheduler.scale_noise(sample, timesteps, noise)
+        assert torch.allclose(noisy, sample, atol=0.1)
+
+        # At σ≈1 (t=max), should be mostly noise
+        timesteps = scheduler.timesteps[0:1]  # First timestep (largest sigma)
+        noisy = scheduler.scale_noise(sample, timesteps, noise)
+        assert torch.allclose(noisy, noise, atol=0.1)
+
+    def test_scale_noise_without_provided_noise(self):
+        """Test scale_noise generates noise if not provided."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        sample = torch.randn(BATCH_SIZE_PAIR, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI)
+        timesteps = scheduler.timesteps[[100, 200]]
+
+        # Should generate noise internally
+        noisy = scheduler.scale_noise(sample, timesteps, noise=None)
+
+        assert noisy.shape == sample.shape
+        # Should be different from original (has noise added)
+        assert not torch.allclose(noisy, sample)
+
+    def test_scale_noise_dtype_consistency(self):
+        """Test scale_noise preserves dtype."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        for dtype in [torch.float32, torch.float16]:
+            sample = torch.randn(
+                BATCH_SIZE_PAIR,
+                TENSOR_CHANNELS_RGB,
+                IMG_SIZE_MINI,
+                IMG_SIZE_MINI,
+                dtype=dtype,
+            )
+            noise = torch.randn(
+                BATCH_SIZE_PAIR,
+                TENSOR_CHANNELS_RGB,
+                IMG_SIZE_MINI,
+                IMG_SIZE_MINI,
+                dtype=dtype,
+            )
+            timesteps = scheduler.timesteps[[100, 200]]
+
+            noisy = scheduler.scale_noise(sample, timesteps, noise)
+
+            assert noisy.dtype == dtype
+
+    # ========================================================================
+    # Inference Tests (Reverse Process - set_timesteps, step)
+    # ========================================================================
+
+    def test_set_timesteps_basic(self):
+        """Test set_timesteps creates inference schedule."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        num_steps = TRAINING_STEPS_MODERATE  # Moderate training steps
+        scheduler.set_timesteps(num_inference_steps=num_steps, device="cpu")
+
+        assert len(scheduler.timesteps) == num_steps
+        assert scheduler.timesteps[0] > scheduler.timesteps[-1]  # Descending
+        assert len(scheduler.sigmas) == num_steps + 1  # +1 for final zero
+
+        # Verify sigma at end is zero
+        assert scheduler.sigmas[-1] == 0.0
+
+    def test_set_timesteps_with_dynamic_shift(self):
+        """Test set_timesteps with dynamic shifting."""
+        scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
+
+        mu = 0.5  # Shift parameter based on image resolution
+        scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu)
+
+        assert len(scheduler.timesteps) == 10
+
+        # Should raise error if mu not provided
+        scheduler2 = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
+        with pytest.raises(ValueError, match="Must pass a value for `mu`"):
+            scheduler2.set_timesteps(num_inference_steps=10)
+
+    def test_step_basic(self):
+        """Test step performs Euler denoising step."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(num_inference_steps=5, device="cpu")
+
+        sample = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI)
+        model_output = torch.randn(
+            BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI
+        )  # Velocity prediction
+        timestep = scheduler.timesteps[0]
+
+        result = scheduler.step(model_output, timestep, sample)
+
+        # step() returns a tuple (prev_sample,)
+        prev_sample = result[0]
+
+        assert prev_sample.shape == sample.shape
+        assert prev_sample.dtype == sample.dtype
+        # Should be different from input
+        assert not torch.equal(prev_sample, sample)
+
+    def test_step_consistency(self):
+        """Test multiple steps are consistent."""
+        scheduler = FlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(num_inference_steps=5, device="cpu")
+
+        sample = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI)
+        model_output = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI)
+
+        # Perform multiple steps
+        current_sample = sample
+        for timestep in scheduler.timesteps:
+            result = scheduler.step(model_output, timestep, current_sample)
+            # step() returns a tuple (prev_sample,)
+            current_sample = result[0]
+
+        # Final sample should be different from initial
+        assert not torch.equal(current_sample, sample)
+        assert current_sample.shape == sample.shape
+
+    def test_dynamic_shifting_behavior(self):
+        """Test dynamic shifting changes timesteps based on resolution."""
+        scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
+
+        # Test with different mu values (resolution-dependent)
+        mu_small = 0.3  # Smaller resolution
+        mu_large = 0.8  # Larger resolution
+
+        scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu_small)
+        timesteps_small = scheduler.timesteps.clone()
+
+        scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu_large)
+        timesteps_large = scheduler.timesteps.clone()
+
+        # Timesteps should be different for different mu
+        assert not torch.allclose(timesteps_small, timesteps_large)
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py b/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py
new file mode 100644
index 000000000..c429ede2c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py
@@ -0,0 +1,211 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for per-step RNG isolation reproducibility.
+
+Validates that:
+- Same step_count + same seed produces bitwise-identical outputs
+- Different step_count produces different outputs (RNG isolation works)
+"""
+
+import contextlib
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig
+from primus.backends.megatron.core.models.diffusion.flux.model import Flux
+from primus.backends.megatron.training.diffusion.forward_step import (
+    flux_forward_step_func,
+)
+from primus.backends.megatron.training.diffusion.schedulers.flow_matching import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from tests.unit_tests.backends.megatron.diffusion.constants import (
+    CLIP_L_EMBEDDING_DIM,
+    T5_XXL_EMBEDDING_DIM,
+    VAE_LATENT_CHANNELS,
+)
+
+
+def _patch_parallel_state(seed=42):
+    """Return a list of mock.patch context managers for parallel state."""
+    return [
+        patch(
+            "megatron.core.parallel_state.get_tensor_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_pipeline_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_data_parallel_rank",
+            return_value=0,
+        ),
+        patch(
+            "megatron.core.parallel_state.get_tensor_model_parallel_rank",
+            return_value=0,
+        ),
+        patch(
+            "megatron.training.get_args",
+            return_value=SimpleNamespace(seed=seed),
+        ),
+    ]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestVaeResampleReproducibility:
+    """Tests per-step RNG isolation via observable output reproducibility."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state for model tests."""
+
+    @pytest.fixture
+    def model(self):
+        config = FluxConfig.flux_535m()
+        m = Flux(config).cuda()
+        m.train()
+        return m
+
+    @pytest.fixture
+    def scheduler(self):
+        return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
+
+    def _make_resample_batch(self, seed=99):
+        """Create a deterministic resample batch."""
+        g = torch.Generator().manual_seed(seed)
+        return {
+            "mean": torch.randn(2, VAE_LATENT_CHANNELS, 16, 16, generator=g),
+            "logvar": torch.randn(2, VAE_LATENT_CHANNELS, 16, 16, generator=g),
+            "prompt_embeds": torch.randn(2, 32, T5_XXL_EMBEDDING_DIM, generator=g),
+            "pooled_prompt_embeds": torch.randn(2, CLIP_L_EMBEDDING_DIM, generator=g),
+        }
+
+    def test_same_step_count_produces_identical_output(self, model, scheduler):
+        """Identical seed + step_count must produce bitwise-identical noise_pred."""
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+
+            batch1 = self._make_resample_batch(seed=99)
+            result1 = flux_forward_step_func(
+                iter([batch1]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=5,
+            )
+            noise_pred_1 = result1[0].detach().clone()
+
+            batch2 = self._make_resample_batch(seed=99)
+            result2 = flux_forward_step_func(
+                iter([batch2]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=5,
+            )
+            noise_pred_2 = result2[0].detach().clone()
+
+        assert torch.equal(noise_pred_1, noise_pred_2), (
+            "Same step_count should produce identical output "
+            f"(max diff: {(noise_pred_1 - noise_pred_2).abs().max().item():.2e})"
+        )
+
+    def test_different_step_count_produces_different_latents(self, model, scheduler):
+        """Different step_count values should produce different clean_latents.
+
+        Note: We compare clean_latents (index 1) rather than noise_pred (index 0)
+        because proj_out is zero-initialized (NeMo-aligned init), making noise_pred
+        always zero until training updates the projection layer.
+        """
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state():
+                stack.enter_context(p)
+
+            batch1 = self._make_resample_batch(seed=99)
+            result1 = flux_forward_step_func(
+                iter([batch1]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=5,
+            )
+            latents_1 = result1[1].detach().clone()
+
+            batch2 = self._make_resample_batch(seed=99)
+            result2 = flux_forward_step_func(
+                iter([batch2]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=6,
+            )
+            latents_2 = result2[1].detach().clone()
+
+        assert not torch.equal(latents_1, latents_2), (
+            "Different step_count values should produce different latents "
+            "(per-step RNG isolation not working)"
+        )
+
+    def test_different_base_seed_produces_different_output(self, model, scheduler):
+        """Different args.seed (same step_count) must produce different latents.
+
+        Validates the `_per_rank_seed = seed + 100*dp_rank` term in the step
+        seed formula at forward_step.py:237. A regression that drops `seed`
+        and uses only `step_count` would cause both runs to produce identical
+        output (since step_count is fixed at 5).
+        """
+        # Run 1: seed=42
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state(seed=42):
+                stack.enter_context(p)
+            batch1 = self._make_resample_batch(seed=99)
+            result1 = flux_forward_step_func(
+                iter([batch1]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=5,
+            )
+            latents_42 = result1[1].detach().clone()
+
+        # Run 2: seed=43, identical step_count and identical batch contents.
+        with contextlib.ExitStack() as stack:
+            for p in _patch_parallel_state(seed=43):
+                stack.enter_context(p)
+            batch2 = self._make_resample_batch(seed=99)
+            result2 = flux_forward_step_func(
+                iter([batch2]),
+                model,
+                scheduler=scheduler,
+                vae_latent_mode="resample",
+                vae_scale=0.3611,
+                vae_shift=0.1159,
+                step_count=5,
+            )
+            latents_43 = result2[1].detach().clone()
+
+        assert not torch.equal(latents_42, latents_43), (
+            "Different args.seed values produced identical output — "
+            "the per-rank seed term appears to be missing from the step seed formula"
+        )
diff --git a/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py
new file mode 100644
index 000000000..35f7ac02c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py
@@ -0,0 +1,351 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FSDP2BF16MasterWeightOptimizer.
+
+Tests focus on the core design properties:
+- FP32 master weights are created from BF16 model parameters
+- prepare_grads copies BF16 grads to FP32 master grads
+- Optimizer step updates FP32 masters and copies back to BF16 params
+- state_dict includes fp32_from_fp16_params key
+- Checkpoint save/load round-trip preserves values
+- sharded_state_dict integrates with Megatron's torch_dist checkpointing
+- finalize_dist_ckpt_load fills step counter and syncs params
+- Factory param group splitting (weight decay, LR scaling) works
+- Multi-step training actually reduces loss
+"""
+
+import pytest
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+from primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer import (
+    get_fsdp2_bf16_master_weight_optimizer,
+)
+from tests.utils import PrimusUT
+
+
+def _make_optimizer_config(**overrides):
+    defaults = dict(
+        bf16=True,
+        fp16=False,
+        optimizer="adam",
+        lr=1e-4,
+        min_lr=1e-5,
+        weight_decay=0.01,
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        adam_eps=1e-8,
+        clip_grad=1.0,
+        use_precision_aware_optimizer=False,
+    )
+    defaults.update(overrides)
+    return OptimizerConfig(**defaults)
+
+
+def _make_model_and_optimizer(device="cpu"):
+    model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device=device)
+    config = _make_optimizer_config()
+    optimizer = get_fsdp2_bf16_master_weight_optimizer(
+        config=config,
+        model_chunks=[model],
+    )
+    return model, optimizer, config
+
+
+def _build_model_sharded_state_dict(model, prefix=""):
+    model_sd = model.state_dict(prefix=prefix, keep_vars=True)
+    return make_sharded_tensors_for_checkpoint(model_sd, prefix, {}, ())
+
+
+class TestFSDP2BF16MasterWeightOptimizer(PrimusUT):
+    """Tests for FSDP2BF16MasterWeightOptimizer core functionality."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_master_weights_are_fp32_clones_of_bf16(self):
+        """FP32 master weights must match BF16 params in value and shape."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                assert bf16_param.dtype == torch.bfloat16
+                assert fp32_master.dtype == torch.float32
+                assert fp32_master.shape == bf16_param.shape
+                torch.testing.assert_close(fp32_master, bf16_param.float(), atol=0, rtol=0)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_main_param_attribute_set(self):
+        """BF16 params must have .main_param pointing to FP32 master."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                assert hasattr(bf16_param, "main_param")
+                assert bf16_param.main_param is fp32_master
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_optimizer_operates_on_fp32_masters(self):
+        """Base optimizer's param_groups must contain FP32 masters, not BF16 params."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        for group in optimizer.optimizer.param_groups:
+            for p in group["params"]:
+                assert p.dtype == torch.float32, f"Optimizer param dtype is {p.dtype}, expected float32"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_prepare_grads_copies_to_fp32(self):
+        """prepare_grads must cast BF16 grads to FP32 on masters and clear BF16 grads."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+
+        for bf16_group in optimizer.bf16_groups:
+            for p in bf16_group:
+                assert p.grad is not None, "BF16 param should have grad before prepare_grads"
+
+        optimizer.prepare_grads()
+
+        for bf16_group in optimizer.bf16_groups:
+            for p in bf16_group:
+                assert p.grad is None, "BF16 grad should be cleared after prepare_grads"
+
+        for fp32_group in optimizer.fp32_from_bf16_groups:
+            for p in fp32_group:
+                assert p.grad is not None, "FP32 master should have grad after prepare_grads"
+                assert p.grad.dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_step_updates_both_master_and_model(self):
+        """step() must update FP32 masters and copy back to BF16 params.
+
+        Uses high LR to ensure BF16 copy-back produces visible changes
+        (small updates can round to zero in BF16).
+        """
+        model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device="cuda")
+        config = _make_optimizer_config(lr=1e-1, clip_grad=1.0)
+        optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model])
+
+        bf16_before = {}
+        for name, p in model.named_parameters():
+            bf16_before[name] = p.data.clone()
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda")
+        loss = (model(x) * 10.0).sum()
+        loss.backward()
+        success, grad_norm, num_zeros = optimizer.step()
+
+        assert success is True
+
+        for name, p in model.named_parameters():
+            assert not torch.equal(
+                p.data, bf16_before[name]
+            ), f"BF16 param '{name}' was not updated after step"
+
+        for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                expected = fp32_master.data.to(torch.bfloat16)
+                torch.testing.assert_close(bf16_param.data, expected, atol=0, rtol=0)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_optimizer_states_are_fp32(self):
+        """All optimizer states (exp_avg, exp_avg_sq) must be FP32."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        for fp32_group in optimizer.fp32_from_bf16_groups:
+            for p in fp32_group:
+                state = optimizer.optimizer.state[p]
+                assert state["exp_avg"].dtype == torch.float32
+                assert state["exp_avg_sq"].dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_state_dict_has_fp32_from_fp16_params_key(self):
+        """state_dict must include 'fp32_from_fp16_params' with FP32 master weights."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        saved = optimizer.state_dict()
+        assert "optimizer" in saved
+        assert "fp32_from_fp16_params" in saved
+
+        for group in saved["fp32_from_fp16_params"]:
+            for param in group:
+                assert param.dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_multi_step_loss_decreases(self):
+        """Run 10 steps and verify loss decreases."""
+        torch.manual_seed(42)
+        model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device="cuda")
+        config = _make_optimizer_config(lr=1e-2, clip_grad=0.0)
+        optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model])
+
+        x = torch.randn(4, 8, dtype=torch.bfloat16, device="cuda")
+        target = torch.randn(4, 8, dtype=torch.bfloat16, device="cuda")
+
+        loss_initial = None
+        loss_final = None
+        for _ in range(10):
+            optimizer.zero_grad()
+            output = model(x)
+            loss = ((output - target).float() ** 2).mean()
+            loss.backward()
+            optimizer.step()
+
+            if loss_initial is None:
+                loss_initial = loss.item()
+            loss_final = loss.item()
+
+        assert (
+            loss_final < loss_initial
+        ), f"Loss did not decrease: initial={loss_initial:.6f}, final={loss_final:.6f}"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_state_dict_round_trip(self):
+        """Verify state_dict save/load preserves FP32 master weights and optimizer states."""
+        model, optimizer, config = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        saved_state = optimizer.state_dict()
+
+        original_masters = []
+        for group in saved_state["fp32_from_fp16_params"]:
+            original_masters.append([p.clone() for p in group])
+
+        original_opt_states = {}
+        for pid, pstate in saved_state["optimizer"]["state"].items():
+            original_opt_states[pid] = {
+                "exp_avg": pstate["exp_avg"].clone(),
+                "exp_avg_sq": pstate["exp_avg_sq"].clone(),
+            }
+
+        fresh_optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model])
+        fresh_optimizer.load_state_dict(saved_state)
+
+        loaded_state = fresh_optimizer.state_dict()
+        for saved_group, loaded_group in zip(original_masters, loaded_state["fp32_from_fp16_params"]):
+            for saved_p, loaded_p in zip(saved_group, loaded_group):
+                torch.testing.assert_close(loaded_p, saved_p)
+                assert loaded_p.dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_sharded_state_dict_creates_sharded_tensors(self):
+        """Verify sharded_state_dict produces ShardedTensor entries."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        model_ssd = _build_model_sharded_state_dict(model)
+        opt_ssd = optimizer.sharded_state_dict(model_ssd, is_loading=True)
+
+        assert "optimizer" in opt_ssd
+        assert "fp32_from_fp16_params" in opt_ssd
+
+        # Check optimizer states are ShardedTensors
+        state_entries = opt_ssd["optimizer"]["state"]
+        assert len(state_entries) > 0
+        for pid, pstate in state_entries.items():
+            if isinstance(pstate, dict):
+                for key in ("exp_avg", "exp_avg_sq"):
+                    assert key in pstate
+                    assert isinstance(pstate[key], ShardedTensor)
+
+        # Check fp32_from_fp16_params are ShardedTensors
+        for group in opt_ssd["fp32_from_fp16_params"]:
+            for entry in group:
+                assert isinstance(entry, ShardedTensor)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_finalize_dist_ckpt_load(self):
+        """Verify finalize_dist_ckpt_load sets step counter and syncs params."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+        optimizer.init_state_fn(optimizer.optimizer)
+
+        # Modify FP32 masters to differ from BF16 params
+        for fp32_group in optimizer.fp32_from_bf16_groups:
+            for p in fp32_group:
+                p.data.add_(1.0)
+
+        iteration = 42
+        optimizer.finalize_dist_ckpt_load(iteration)
+
+        for fp32_group in optimizer.fp32_from_bf16_groups:
+            for p in fp32_group:
+                if p in optimizer.optimizer.state:
+                    step = optimizer.optimizer.state[p].get("step")
+                    if step is not None:
+                        assert step.item() == float(iteration)
+
+        # BF16 params should now match FP32 masters (copy-back happened)
+        for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                expected = fp32_master.data.to(torch.bfloat16)
+                torch.testing.assert_close(bf16_param.data, expected, atol=0, rtol=0)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_reload_model_params(self):
+        """Verify reload_model_params copies BF16 params to FP32 masters."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        # Manually change BF16 params (simulating checkpoint load)
+        for p in model.parameters():
+            p.data.fill_(0.5)
+
+        optimizer.reload_model_params()
+
+        for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups):
+            for bf16_param, fp32_master in zip(bf16_group, fp32_group):
+                torch.testing.assert_close(fp32_master, bf16_param.float(), atol=0, rtol=0)
+
+    def test_factory_applies_weight_decay_condition(self):
+        """Verify no_weight_decay_cond splits parameters into correct groups."""
+        model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16)
+        config = _make_optimizer_config(weight_decay=0.1)
+
+        def no_wd_for_bias(param):
+            return param.ndim == 1
+
+        optimizer = get_fsdp2_bf16_master_weight_optimizer(
+            config=config,
+            model_chunks=[model],
+            no_weight_decay_cond=no_wd_for_bias,
+        )
+
+        assert len(optimizer.optimizer.param_groups) == 2
+        wds = {pg["weight_decay"] for pg in optimizer.optimizer.param_groups}
+        assert 0.0 in wds
+        assert 0.1 in wds
+
+    def test_factory_handles_fp32_params(self):
+        """Verify FP32 params (e.g. LayerNorm) go to fp32_from_fp32_groups without master copy."""
+        model = torch.nn.Sequential(
+            torch.nn.Linear(8, 8, bias=False).to(dtype=torch.bfloat16),
+            torch.nn.LayerNorm(8).to(dtype=torch.float32),
+        )
+        config = _make_optimizer_config()
+        optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model])
+
+        total_bf16 = sum(len(g) for g in optimizer.bf16_groups)
+        total_fp32 = sum(len(g) for g in optimizer.fp32_from_fp32_groups)
+
+        assert total_bf16 == 1, f"Expected 1 BF16 param, got {total_bf16}"
+        assert total_fp32 == 2, f"Expected 2 FP32 params (LN weight+bias), got {total_fp32}"
diff --git a/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py
new file mode 100644
index 000000000..9bda55efa
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py
@@ -0,0 +1,311 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FSDP2FP32Optimizer.
+
+Tests focus on the core design properties of this optimizer:
+- Optimizer states are FP32 (the whole point)
+- state_dict format is raw PyTorch (not Megatron-wrapped)
+- Gradient clipping actually clips gradients
+- Optimizer step actually updates parameters and reduces loss
+- Checkpoint save/load round-trip preserves state values and dtypes
+- sharded_state_dict integrates with Megatron's torch_dist checkpointing
+- finalize_dist_ckpt_load fills step counter (FSDP2-specific)
+- Factory param group splitting (weight decay, LR scaling) works
+"""
+
+import pytest
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+from primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer import (
+    get_fsdp2_fp32_optimizer,
+)
+from tests.utils import PrimusUT
+
+
+def _make_optimizer_config(**overrides):
+    defaults = dict(
+        bf16=True,
+        fp16=False,
+        optimizer="adam",
+        lr=1e-4,
+        min_lr=1e-5,
+        weight_decay=0.01,
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        adam_eps=1e-8,
+        clip_grad=1.0,
+        use_precision_aware_optimizer=False,
+    )
+    defaults.update(overrides)
+    return OptimizerConfig(**defaults)
+
+
+def _make_model_and_optimizer(device="cpu"):
+    model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32, device=device)
+    config = _make_optimizer_config()
+    optimizer = get_fsdp2_fp32_optimizer(
+        config=config,
+        model_chunks=[model],
+    )
+    return model, optimizer, config
+
+
+def _build_model_sharded_state_dict(model, prefix=""):
+    model_sd = model.state_dict(prefix=prefix, keep_vars=True)
+    return make_sharded_tensors_for_checkpoint(model_sd, prefix, {}, ())
+
+
+def _compute_grad_norm(model):
+    """Compute L2 gradient norm across all model parameters."""
+    total = 0.0
+    for p in model.parameters():
+        if p.grad is not None:
+            total += p.grad.float().norm().item() ** 2
+    return total**0.5
+
+
+class TestFSDP2FP32Optimizer(PrimusUT):
+    """Tests for FSDP2FP32Optimizer core functionality."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_optimizer_states_are_fp32(self):
+        """The core invariant: all optimizer states must be FP32."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.float32, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        for p in optimizer.get_parameters():
+            state = optimizer.optimizer.state[p]
+            assert (
+                state["exp_avg"].dtype == torch.float32
+            ), f"exp_avg dtype is {state['exp_avg'].dtype}, expected float32"
+            assert (
+                state["exp_avg_sq"].dtype == torch.float32
+            ), f"exp_avg_sq dtype is {state['exp_avg_sq'].dtype}, expected float32"
+            assert state["step"].dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_state_dict_format_is_raw(self):
+        """state_dict must return raw PyTorch format, NOT a Megatron-wrapped
+        format ({"optimizer": {...}}).
+
+        This distinction is critical for checkpoint compatibility.
+        """
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.float32, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        saved = optimizer.state_dict()
+        assert "state" in saved, "state_dict missing 'state' key"
+        assert "param_groups" in saved, "state_dict missing 'param_groups' key"
+        assert "optimizer" not in saved, (
+            "state_dict has 'optimizer' wrapper key -- this is the BFloat16Optimizer "
+            "format, not the raw format expected by FSDP2FP32Optimizer"
+        )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_clip_grad_norm_clips_gradients(self):
+        """Verify clipping reduces grad norm to at most clip_value."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.float32, device="cuda")
+        loss = (model(x) * 100.0).sum()
+        loss.backward()
+
+        grad_norm_before = _compute_grad_norm(model)
+        assert grad_norm_before > 1.0, "Gradients too small to test clipping"
+
+        clip_value = 0.01
+        returned_norm = optimizer.clip_grad_norm(clip_value)
+        assert isinstance(
+            returned_norm, torch.Tensor
+        ), "clip_grad_norm should return a GPU tensor (deferred .item())"
+        returned_norm_f = returned_norm.item()
+
+        grad_norm_after = _compute_grad_norm(model)
+
+        assert (
+            abs(returned_norm_f - grad_norm_before) < 1e-3
+        ), f"Returned norm {returned_norm_f} should approximate pre-clip norm {grad_norm_before}"
+        assert (
+            grad_norm_after <= clip_value + 1e-6
+        ), f"Post-clip norm {grad_norm_after} exceeds clip_value {clip_value}"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_step_actually_updates_parameters(self):
+        """Verify optimizer.step() changes parameter values and returns
+        the (success, grad_norm, num_zeros) 3-tuple that Megatron's
+        training loop unpacks."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        params_before = {name: p.clone() for name, p in model.named_parameters()}
+
+        x = torch.randn(2, 8, dtype=torch.float32, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        success, grad_norm, num_zeros = optimizer.step()
+
+        assert success is True
+        assert isinstance(
+            grad_norm, (float, torch.Tensor)
+        ), f"grad_norm should be float or Tensor, got {type(grad_norm)}"
+        grad_norm_f = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+        assert grad_norm_f > 0.0
+
+        for name, p in model.named_parameters():
+            assert not torch.equal(
+                p, params_before[name]
+            ), f"Parameter '{name}' was not updated by optimizer.step()"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_multi_step_loss_decreases(self):
+        """Run 10 steps and verify loss decreases, catching silent update failures."""
+        torch.manual_seed(42)
+        model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32, device="cuda")
+        config = _make_optimizer_config(lr=1e-2, clip_grad=0.0)
+        optimizer = get_fsdp2_fp32_optimizer(config=config, model_chunks=[model])
+
+        x = torch.randn(4, 8, dtype=torch.float32, device="cuda")
+        target = torch.randn(4, 8, dtype=torch.float32, device="cuda")
+
+        loss_initial = None
+        loss_final = None
+        for _ in range(10):
+            optimizer.zero_grad()
+            output = model(x)
+            loss = ((output - target) ** 2).mean()
+            loss.backward()
+            optimizer.step()
+
+            if loss_initial is None:
+                loss_initial = loss.item()
+            loss_final = loss.item()
+
+        assert (
+            loss_final < loss_initial
+        ), f"Loss did not decrease: initial={loss_initial:.6f}, final={loss_final:.6f}"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_state_dict_round_trip(self):
+        """Verify state_dict save/load preserves values and FP32 dtypes."""
+        model, optimizer, config = _make_model_and_optimizer(device="cuda")
+
+        x = torch.randn(2, 8, dtype=torch.float32, device="cuda")
+        loss = model(x).sum()
+        loss.backward()
+        optimizer.step()
+
+        saved_state = optimizer.state_dict()
+
+        original_states = {}
+        for pid, pstate in saved_state["state"].items():
+            original_states[pid] = {
+                "exp_avg": pstate["exp_avg"].clone(),
+                "exp_avg_sq": pstate["exp_avg_sq"].clone(),
+            }
+
+        fresh_optimizer = get_fsdp2_fp32_optimizer(
+            config=config,
+            model_chunks=[model],
+        )
+        fresh_optimizer.load_state_dict(saved_state)
+
+        loaded_state = fresh_optimizer.state_dict()
+        for pid in original_states:
+            assert pid in loaded_state["state"]
+            torch.testing.assert_close(
+                loaded_state["state"][pid]["exp_avg"],
+                original_states[pid]["exp_avg"],
+            )
+            torch.testing.assert_close(
+                loaded_state["state"][pid]["exp_avg_sq"],
+                original_states[pid]["exp_avg_sq"],
+            )
+            assert loaded_state["state"][pid]["exp_avg"].dtype == torch.float32
+            assert loaded_state["state"][pid]["exp_avg_sq"].dtype == torch.float32
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_sharded_state_dict_creates_sharded_tensors(self):
+        """Verify sharded_state_dict produces ShardedTensor entries for torch_dist checkpointing."""
+        model, optimizer, _ = _make_model_and_optimizer(device="cuda")
+
+        model_ssd = _build_model_sharded_state_dict(model)
+        opt_ssd = optimizer.sharded_state_dict(model_ssd, is_loading=True)
+
+        assert "state" in opt_ssd
+        state_entries = opt_ssd["state"]
+        assert len(state_entries) > 0
+
+        for pid, pstate in state_entries.items():
+            if isinstance(pstate, dict):
+                for key in ("exp_avg", "exp_avg_sq"):
+                    assert key in pstate, f"Key '{key}' missing from state[{pid}]"
+                    assert isinstance(pstate[key], ShardedTensor)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_finalize_dist_ckpt_load(self):
+        """Verify finalize_dist_ckpt_load sets step counter from iteration."""
+        model, optimizer, config = _make_model_and_optimizer(device="cuda")
+
+        optimizer.init_state_fn(optimizer.optimizer)
+
+        iteration = 42
+        optimizer.finalize_dist_ckpt_load(iteration)
+
+        for p in optimizer.get_parameters():
+            if p in optimizer.optimizer.state:
+                step = optimizer.optimizer.state[p].get("step")
+                if step is not None:
+                    assert step.item() == float(iteration)
+
+    def test_init_state_fn_creates_fp32_states(self):
+        """Verify init_state_fn creates FP32 placeholders for checkpoint loading."""
+        model, optimizer, _ = _make_model_and_optimizer()
+
+        base_opt = optimizer.optimizer
+        for group in base_opt.param_groups:
+            for p in group["params"]:
+                assert len(base_opt.state[p]) == 0
+
+        optimizer.init_state_fn(base_opt)
+
+        for group in base_opt.param_groups:
+            for p in group["params"]:
+                state = base_opt.state[p]
+                assert state["exp_avg"].dtype == torch.float32
+                assert state["exp_avg_sq"].dtype == torch.float32
+                assert state["exp_avg"].shape == p.data.shape
+
+    def test_factory_applies_weight_decay_condition(self):
+        """Verify no_weight_decay_cond splits parameters into correct groups."""
+        model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32)
+        config = _make_optimizer_config(weight_decay=0.1)
+
+        def no_wd_for_bias(param):
+            return param.ndim == 1
+
+        optimizer = get_fsdp2_fp32_optimizer(
+            config=config,
+            model_chunks=[model],
+            no_weight_decay_cond=no_wd_for_bias,
+        )
+
+        assert len(optimizer.optimizer.param_groups) == 2
+        wds = {pg["weight_decay"] for pg in optimizer.optimizer.param_groups}
+        assert 0.0 in wds
+        assert 0.1 in wds
diff --git a/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py
new file mode 100644
index 000000000..7869ce074
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py
@@ -0,0 +1,488 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for FSDP2 FP32 optimizer patches.
+
+Tests:
+- FSDP2 FP32 optimizer patch condition evaluation and application
+- Float16Module skip patch condition evaluation and application
+- _FSDP2PassthroughModule preserves FP32 weights and passes through forward
+- Mutual exclusivity of bf16 and fsdp2_fp32 optimizer patches
+"""
+
+import sys
+import types
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from primus.core.patches.context import PatchContext
+from primus.core.patches.patch_runner import run_patches
+
+
+def _build_patch_context(backend_args=None):
+    if backend_args is None:
+        backend_args = SimpleNamespace()
+    module_config = SimpleNamespace(params=backend_args)
+    return PatchContext(
+        backend="megatron",
+        phase="before_train",
+        extra={"backend_args": backend_args, "module_config": module_config},
+    )
+
+
+def _install_fake_megatron_modules(monkeypatch):
+    training_mod = types.ModuleType("megatron.training.training")
+    training_pkg = types.ModuleType("megatron.training")
+    training_pkg.training = training_mod
+
+    core_pkg = types.ModuleType("megatron.core")
+    optimizer_mod = types.ModuleType("megatron.core.optimizer")
+    core_pkg.optimizer = optimizer_mod
+
+    megatron_pkg = types.ModuleType("megatron")
+    megatron_pkg.training = training_pkg
+    megatron_pkg.core = core_pkg
+
+    monkeypatch.setitem(sys.modules, "megatron", megatron_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.training", training_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod)
+    monkeypatch.setitem(sys.modules, "megatron.core", core_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.core.optimizer", optimizer_mod)
+
+    return training_mod, optimizer_mod
+
+
+def _install_fsdp2_fp32_stub(monkeypatch, optimizer_fn):
+    stub = types.ModuleType("primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer")
+    stub.get_fsdp2_fp32_optimizer = optimizer_fn
+    monkeypatch.setitem(
+        sys.modules,
+        "primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer",
+        stub,
+    )
+
+
+class TestFSDP2FP32OptimizerPatch:
+    """Tests for the FSDP2 FP32 optimizer patch."""
+
+    def test_patch_applies_when_all_conditions_met(self, monkeypatch):
+        from primus.backends.megatron.patches.optimizer_patches import (
+            patch_fsdp2_fp32_optimizer,
+        )
+
+        backend_args = SimpleNamespace(
+            use_fsdp2_fp32_param_optimizer=True,
+            use_torch_fsdp2=True,
+            bf16=True,
+        )
+        ctx = _build_patch_context(backend_args)
+
+        training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch)
+        original = lambda *args, **kwargs: Mock()
+        training_mod.get_megatron_optimizer = original
+        optimizer_mod.get_megatron_optimizer = original
+
+        mock_optimizer = Mock()
+        _install_fsdp2_fp32_stub(monkeypatch, lambda *args, **kwargs: mock_optimizer)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.optimizer_patches.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        patch_fsdp2_fp32_optimizer(ctx)
+
+        assert training_mod.get_megatron_optimizer is not original
+        assert optimizer_mod.get_megatron_optimizer is not original
+
+        result = training_mod.get_megatron_optimizer(Mock(), [Mock()])
+        assert result is mock_optimizer
+
+    @pytest.mark.parametrize(
+        "backend_args",
+        [
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=False, use_torch_fsdp2=True, bf16=True),
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=False, bf16=True),
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=True, bf16=False),
+        ],
+    )
+    def test_patch_skipped_when_any_condition_false(self, monkeypatch, backend_args):
+        import primus.backends.megatron.patches.optimizer_patches  # noqa: F401
+
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.error_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch)
+        original = lambda *args, **kwargs: Mock()
+        training_mod.get_megatron_optimizer = original
+        optimizer_mod.get_megatron_optimizer = original
+
+        count = run_patches(
+            backend="megatron",
+            phase="before_train",
+            extra={"module_config": SimpleNamespace(params=backend_args)},
+            enabled_ids=["megatron.optimizer.fsdp2_fp32_param"],
+        )
+
+        assert count == 0
+        assert training_mod.get_megatron_optimizer is original
+
+    def test_patched_optimizer_filters_megatron_kwargs(self, monkeypatch):
+        from primus.backends.megatron.patches.optimizer_patches import (
+            patch_fsdp2_fp32_optimizer,
+        )
+
+        backend_args = SimpleNamespace(
+            use_fsdp2_fp32_param_optimizer=True,
+            use_torch_fsdp2=True,
+            bf16=True,
+        )
+        ctx = _build_patch_context(backend_args)
+
+        training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch)
+        training_mod.get_megatron_optimizer = lambda *a, **kw: Mock()
+        optimizer_mod.get_megatron_optimizer = lambda *a, **kw: Mock()
+
+        received_kwargs = {}
+
+        def tracking_factory(**kwargs):
+            received_kwargs.update(kwargs)
+            return Mock()
+
+        _install_fsdp2_fp32_stub(monkeypatch, lambda *args, **kwargs: tracking_factory(**kwargs))
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.optimizer_patches.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        patch_fsdp2_fp32_optimizer(ctx)
+
+        training_mod.get_megatron_optimizer(
+            Mock(),
+            [Mock()],
+            config_overrides={"x": 1},
+            use_gloo_process_groups=True,
+            no_weight_decay_cond=Mock(),
+        )
+
+        assert "config_overrides" not in received_kwargs
+        assert "use_gloo_process_groups" not in received_kwargs
+        assert "no_weight_decay_cond" in received_kwargs
+
+
+class TestSkipFloat16ModulePatch:
+    """Tests for the Float16Module skip patch."""
+
+    def test_patch_replaces_float16_module(self, monkeypatch):
+        from primus.backends.megatron.patches.build_model_patches import (
+            patch_skip_float16_module,
+        )
+
+        backend_args = SimpleNamespace(
+            use_fsdp2_fp32_param_optimizer=True,
+            use_torch_fsdp2=True,
+            bf16=True,
+        )
+        ctx = _build_patch_context(backend_args)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.build_model_patches.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        _install_fake_megatron_modules(monkeypatch)
+
+        transformer_pkg = types.ModuleType("megatron.core.transformer")
+        transformer_module_mod = types.ModuleType("megatron.core.transformer.module")
+
+        class FakeFloat16Module:
+            pass
+
+        class FakeMegatronModule(torch.nn.Module):
+            def __init__(self, config=None):
+                super().__init__()
+
+        transformer_module_mod.Float16Module = FakeFloat16Module
+        transformer_module_mod.MegatronModule = FakeMegatronModule
+        transformer_pkg.module = transformer_module_mod
+        monkeypatch.setitem(
+            sys.modules,
+            "megatron.core.transformer",
+            transformer_pkg,
+        )
+        monkeypatch.setitem(
+            sys.modules,
+            "megatron.core.transformer.module",
+            transformer_module_mod,
+        )
+        sys.modules["megatron.core"].transformer = transformer_pkg
+
+        patch_skip_float16_module(ctx)
+
+        import megatron.training.training as training_mod
+
+        assert training_mod.Float16Module is not FakeFloat16Module
+        assert training_mod.Float16Module.__name__ == "_FSDP2PassthroughModule"
+
+    def test_passthrough_module_preserves_fp32_weights(self, monkeypatch):
+        """Verify _FSDP2PassthroughModule does NOT convert parameters to BF16
+        and that forward() passes through correctly.
+
+        This is the core contract: Float16Module converts to BF16, but the
+        passthrough replacement must keep FP32 so FSDP2 MixedPrecisionPolicy
+        can handle the casting.
+        """
+        from primus.backends.megatron.patches.build_model_patches import (
+            patch_skip_float16_module,
+        )
+
+        backend_args = SimpleNamespace(
+            use_fsdp2_fp32_param_optimizer=True,
+            use_torch_fsdp2=True,
+            bf16=True,
+        )
+        ctx = _build_patch_context(backend_args)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.build_model_patches.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        _install_fake_megatron_modules(monkeypatch)
+
+        transformer_pkg = types.ModuleType("megatron.core.transformer")
+        transformer_module_mod = types.ModuleType("megatron.core.transformer.module")
+
+        class FakeMegatronModule(torch.nn.Module):
+            def __init__(self, config=None):
+                super().__init__()
+
+        class FakeFloat16Module:
+            pass
+
+        transformer_module_mod.Float16Module = FakeFloat16Module
+        transformer_module_mod.MegatronModule = FakeMegatronModule
+        transformer_pkg.module = transformer_module_mod
+        monkeypatch.setitem(
+            sys.modules,
+            "megatron.core.transformer",
+            transformer_pkg,
+        )
+        monkeypatch.setitem(
+            sys.modules,
+            "megatron.core.transformer.module",
+            transformer_module_mod,
+        )
+        sys.modules["megatron.core"].transformer = transformer_pkg
+
+        patch_skip_float16_module(ctx)
+
+        import megatron.training.training as training_mod
+
+        PassthroughCls = training_mod.Float16Module
+
+        inner_model = torch.nn.Linear(4, 4, bias=True).to(dtype=torch.float32)
+        config = SimpleNamespace(virtual_pipeline_model_parallel_size=None)
+        wrapper = PassthroughCls(config, inner_model)
+
+        # Parameters must remain FP32
+        for name, p in wrapper.named_parameters():
+            assert (
+                p.dtype == torch.float32
+            ), f"Parameter '{name}' was converted to {p.dtype}, expected float32"
+
+        # wrapper.module must be the original model
+        assert wrapper.module is inner_model
+
+        # Forward must pass through correctly
+        x = torch.randn(2, 4, dtype=torch.float32)
+        expected = inner_model(x)
+        actual = wrapper(x)
+        torch.testing.assert_close(actual, expected)
+
+    @pytest.mark.parametrize(
+        "backend_args",
+        [
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=False, use_torch_fsdp2=True, bf16=True),
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=False, bf16=True),
+            SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=True, bf16=False),
+        ],
+    )
+    def test_patch_skipped_when_conditions_not_met(self, monkeypatch, backend_args):
+        import primus.backends.megatron.patches.build_model_patches  # noqa: F401
+
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.error_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        count = run_patches(
+            backend="megatron",
+            phase="before_train",
+            extra={"module_config": SimpleNamespace(params=backend_args)},
+            enabled_ids=["megatron.training.training.skip_float16_module"],
+        )
+
+        assert count == 0
+
+
+class TestFSDP2BF16MasterWeightOptimizerPatch:
+    """Tests for the FSDP2 BF16 master weight optimizer patch."""
+
+    def test_patch_applies_when_all_conditions_met(self, monkeypatch):
+        from primus.backends.megatron.patches.optimizer_patches import (
+            patch_fsdp2_bf16_master_weight_optimizer,
+        )
+
+        backend_args = SimpleNamespace(
+            use_fsdp2_bf16_master_weight_optimizer=True,
+            use_torch_fsdp2=True,
+            bf16=True,
+        )
+        ctx = _build_patch_context(backend_args)
+
+        training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch)
+        original = lambda *args, **kwargs: Mock()
+        training_mod.get_megatron_optimizer = original
+        optimizer_mod.get_megatron_optimizer = original
+
+        mock_optimizer = Mock()
+        bf16_mw_stub = types.ModuleType(
+            "primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer"
+        )
+        bf16_mw_stub.get_fsdp2_bf16_master_weight_optimizer = lambda *args, **kwargs: mock_optimizer
+        monkeypatch.setitem(
+            sys.modules,
+            "primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer",
+            bf16_mw_stub,
+        )
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.optimizer_patches.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        patch_fsdp2_bf16_master_weight_optimizer(ctx)
+
+        assert training_mod.get_megatron_optimizer is not original
+        assert optimizer_mod.get_megatron_optimizer is not original
+
+        result = training_mod.get_megatron_optimizer(Mock(), [Mock()])
+        assert result is mock_optimizer
+
+    @pytest.mark.parametrize(
+        "backend_args",
+        [
+            SimpleNamespace(
+                use_fsdp2_bf16_master_weight_optimizer=False,
+                use_torch_fsdp2=True,
+                bf16=True,
+            ),
+            SimpleNamespace(
+                use_fsdp2_bf16_master_weight_optimizer=True,
+                use_torch_fsdp2=False,
+                bf16=True,
+            ),
+            SimpleNamespace(
+                use_fsdp2_bf16_master_weight_optimizer=True,
+                use_torch_fsdp2=True,
+                bf16=False,
+            ),
+        ],
+    )
+    def test_patch_skipped_when_any_condition_false(self, monkeypatch, backend_args):
+        import primus.backends.megatron.patches.optimizer_patches  # noqa: F401
+
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch_runner.error_rank_0",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "primus.core.patches.patch.log_rank_0",
+            lambda *args, **kwargs: None,
+        )
+
+        training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch)
+        original = lambda *args, **kwargs: Mock()
+        training_mod.get_megatron_optimizer = original
+        optimizer_mod.get_megatron_optimizer = original
+
+        count = run_patches(
+            backend="megatron",
+            phase="before_train",
+            extra={"module_config": SimpleNamespace(params=backend_args)},
+            enabled_ids=["megatron.optimizer.fsdp2_bf16_master_weight"],
+        )
+
+        assert count == 0
+        assert training_mod.get_megatron_optimizer is original
+
+
+class TestOptimizerPatchMutualExclusivity:
+    """Tests for mutual exclusivity of the FSDP2 custom optimizer patches."""
+
+    def test_two_conflicting_optimizer_flags_raise(self):
+        """Two FSDP2 custom optimizer flags set must fail loudly: the flags all
+        patch get_megatron_optimizer at the same priority, so enabling more than
+        one is ambiguous. validate_fsdp2_optimizer_exclusivity enforces this at
+        arg-validation time."""
+        from primus.backends.megatron.patches.args.rocm_arg_validation import (
+            validate_fsdp2_optimizer_exclusivity,
+        )
+
+        args = SimpleNamespace(
+            use_fsdp2_fp32_param_optimizer=True,
+            use_fsdp2_bf16_master_weight_optimizer=True,
+        )
+
+        with pytest.raises(ValueError, match="Conflicting FSDP2 optimizer selection"):
+            validate_fsdp2_optimizer_exclusivity(args)
+
+    def test_single_optimizer_flag_is_allowed(self):
+        """Exactly one FSDP2 optimizer flag is the supported case (no raise)."""
+        from primus.backends.megatron.patches.args.rocm_arg_validation import (
+            validate_fsdp2_optimizer_exclusivity,
+        )
+
+        for flag in (
+            "use_fsdp2_fp32_param_optimizer",
+            "use_fsdp2_bf16_master_weight_optimizer",
+        ):
+            args = SimpleNamespace(**{flag: True})
+            # Should not raise.
+            validate_fsdp2_optimizer_exclusivity(args)
+
+    def test_no_optimizer_flag_is_allowed(self):
+        """No FSDP2 optimizer flag set (default Megatron optimizer) must not raise."""
+        from primus.backends.megatron.patches.args.rocm_arg_validation import (
+            validate_fsdp2_optimizer_exclusivity,
+        )
+
+        validate_fsdp2_optimizer_exclusivity(SimpleNamespace())
diff --git a/tests/unit_tests/backends/megatron/patches/test_patch_guard.py b/tests/unit_tests/backends/megatron/patches/test_patch_guard.py
new file mode 100644
index 000000000..97737f16c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/patches/test_patch_guard.py
@@ -0,0 +1,75 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the feature-owned patch idempotency guard.
+
+These validate that:
+  - a guarded "wrapping" patch applies exactly once even if its handler runs
+    twice (e.g. before_train re-invoked), and
+  - distinct patches still compose (each wraps once), which is required for the
+    stacked train_step wrappers (FP8 cache, delayed scaling, wall-clock timer).
+"""
+
+from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched
+
+
+class _FakeModule:
+    """Stand-in for ``megatron.training.training`` (a module object)."""
+
+
+def _apply_wrapping_patch(module, key):
+    """Mimic the guarded train_step wrap pattern used by the FP8/timer patches."""
+    if is_patched(module, key):
+        return
+    original = module.train_step
+
+    def wrapped(*args, **kwargs):
+        module.call_log.append(key)
+        return original(*args, **kwargs)
+
+    module.train_step = wrapped
+    mark_patched(module, key)
+
+
+def test_is_patched_false_before_mark():
+    module = _FakeModule()
+    assert not is_patched(module, "k")
+
+
+def test_mark_then_is_patched():
+    module = _FakeModule()
+    mark_patched(module, "k")
+    assert is_patched(module, "k")
+    assert not is_patched(module, "other")
+
+
+def test_guarded_wrap_applies_once_on_reapply():
+    module = _FakeModule()
+    module.call_log = []
+    module.train_step = lambda: "base"
+
+    # Apply the same patch twice (simulating before_train running twice).
+    _apply_wrapping_patch(module, "megatron.train_step.demo")
+    _apply_wrapping_patch(module, "megatron.train_step.demo")
+
+    module.train_step()
+    # Wrapped exactly once -> the side effect is recorded once per call.
+    assert module.call_log == ["megatron.train_step.demo"]
+
+
+def test_distinct_patches_still_compose():
+    module = _FakeModule()
+    module.call_log = []
+    module.train_step = lambda: "base"
+
+    _apply_wrapping_patch(module, "patch.a")
+    _apply_wrapping_patch(module, "patch.b")
+    # Re-run both: idempotent, so still only one layer each.
+    _apply_wrapping_patch(module, "patch.a")
+    _apply_wrapping_patch(module, "patch.b")
+
+    module.train_step()
+    assert sorted(module.call_log) == ["patch.a", "patch.b"]
diff --git a/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py b/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py
new file mode 100644
index 000000000..13f60b42e
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_aiter_deepbind_patches.py
@@ -0,0 +1,105 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the aiter RTLD_DEEPBIND gating logic
+(``primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py``).
+
+These verify the patch self-detects whether the installed transformer_engine
+already fixed the stale-vendored-libmha issue (ROCm/aiter#1332), without
+needing a GPU. The actual RTLD_DEEPBIND import hook installation is exercised
+separately in the diffusion GPU tests via ``tests.utils.install_aiter_deepbind_hook``.
+"""
+
+import sys
+import types
+
+import pytest
+
+import primus.backends.megatron.patches.turbo.aiter_deepbind_patches as patch_mod
+
+
+def _fake_te_module(version: str) -> types.ModuleType:
+    module = types.ModuleType("transformer_engine")
+    module.__version__ = version
+    return module
+
+
+@pytest.mark.parametrize(
+    "version, expected",
+    [
+        ("2.14.0", (2, 14, 0)),
+        ("2.14.1", (2, 14, 1)),
+        ("1.11.0+abcdef0", (1, 11, 0)),
+        ("2.14.0.dev0+rocm", (2, 14, 0)),
+    ],
+)
+def test_te_version_tuple_parses_leading_semver(monkeypatch, version, expected):
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module(version))
+    assert patch_mod._te_version_tuple() == expected
+
+
+def test_te_version_tuple_none_when_unimportable(monkeypatch):
+    monkeypatch.setitem(sys.modules, "transformer_engine", None)  # forces ImportError on import
+    assert patch_mod._te_version_tuple() is None
+
+
+def test_te_version_tuple_none_when_unparseable(monkeypatch):
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("not-a-version"))
+    assert patch_mod._te_version_tuple() is None
+
+
+@pytest.mark.parametrize(
+    "version, expected_fixed",
+    [
+        ("2.13.9", False),  # pre-fix: still vendors the stale libmha_bwd.so
+        ("2.14.0", True),  # fixed: renamed to te_libmha_*.so / qola::te::
+        ("2.15.0", True),
+        ("1.11.0+abcdef0", False),  # e.g. rocm/primus:v26.3's TE
+    ],
+)
+def test_te_already_fixed_thresholds_on_fix_version(monkeypatch, version, expected_fixed):
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module(version))
+    assert patch_mod._te_already_fixed() is expected_fixed
+
+
+def test_te_already_fixed_fails_safe_towards_unfixed_when_unknown(monkeypatch):
+    # An unparseable/missing version must NOT be treated as "fixed": skipping the
+    # hook on an actually-unfixed TE reintroduces the hd128 backward crash.
+    monkeypatch.setitem(sys.modules, "transformer_engine", None)
+    assert patch_mod._te_already_fixed() is False
+
+
+def _condition_ctx(monkeypatch, *, use_turbo_attention=True, turbo_can_patch=True, affected_arch=True):
+    args = types.SimpleNamespace(use_turbo_attention=use_turbo_attention)
+    monkeypatch.setattr(patch_mod, "get_args", lambda ctx: args)
+    monkeypatch.setattr(patch_mod, "is_primus_turbo_can_patch", lambda ctx: turbo_can_patch)
+    monkeypatch.setattr(patch_mod, "_is_affected_arch", lambda: affected_arch)
+    monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None)
+    return object()  # PatchContext stand-in; unused by the stubbed collaborators above
+
+
+def test_can_install_skips_when_te_already_fixed(monkeypatch):
+    ctx = _condition_ctx(monkeypatch)
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("2.14.0"))
+    assert patch_mod._can_install_aiter_deepbind(ctx) is False
+
+
+def test_can_install_applies_when_te_unfixed_on_affected_arch(monkeypatch):
+    ctx = _condition_ctx(monkeypatch)
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0"))
+    assert patch_mod._can_install_aiter_deepbind(ctx) is True
+
+
+def test_can_install_skips_when_arch_unaffected_regardless_of_te(monkeypatch):
+    ctx = _condition_ctx(monkeypatch, affected_arch=False)
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0"))
+    assert patch_mod._can_install_aiter_deepbind(ctx) is False
+
+
+def test_can_install_skips_when_turbo_attention_disabled(monkeypatch):
+    ctx = _condition_ctx(monkeypatch, use_turbo_attention=False)
+    monkeypatch.setitem(sys.modules, "transformer_engine", _fake_te_module("1.11.0+abcdef0"))
+    assert patch_mod._can_install_aiter_deepbind(ctx) is False
diff --git a/tests/unit_tests/backends/megatron/test_attention_backend_patches.py b/tests/unit_tests/backends/megatron/test_attention_backend_patches.py
new file mode 100644
index 000000000..9ff311f61
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_attention_backend_patches.py
@@ -0,0 +1,85 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the ROCm-safe attention_backend patch.
+
+These verify the patched ``LanguageModule._set_attention_backend`` reconciles the
+selected backend with the image's baked ``NVTE_*_ATTN`` env vars the way ROCm
+needs (override/respect rather than assert-crash), without needing a GPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("megatron")
+
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.transformer.enums import AttnBackend
+
+import primus.backends.megatron.patches.attention_backend_patches as patch_mod
+
+_NVTE = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN")
+
+
+@pytest.fixture
+def patched_set_attention_backend(monkeypatch):
+    """Apply the patch onto a pristine LanguageModule and restore it afterwards."""
+    monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None)
+    original = LanguageModule.__dict__.get("_set_attention_backend")
+    monkeypatch.delattr(LanguageModule, "_primus_attention_backend_patched", raising=False)
+
+    patch_mod.patch_attention_backend(None)
+    try:
+        yield LanguageModule._set_attention_backend
+    finally:
+        if original is not None:
+            LanguageModule._set_attention_backend = original
+        if hasattr(LanguageModule, "_primus_attention_backend_patched"):
+            delattr(LanguageModule, "_primus_attention_backend_patched")
+
+
+def _run(monkeypatch, backend, baked):
+    # Simulate the ROCm image's baked NVTE_*_ATTN before model construction.
+    for name in _NVTE:
+        monkeypatch.delenv(name, raising=False)
+    for name, value in baked.items():
+        monkeypatch.setenv(name, value)
+
+    dummy = SimpleNamespace(config=SimpleNamespace(attention_backend=backend))
+    LanguageModule._set_attention_backend(dummy)
+
+    import os
+
+    return tuple(os.environ.get(name) for name in _NVTE)
+
+
+# The Primus ROCm images bake NVTE_FLASH_ATTN=0 / NVTE_FUSED_ATTN=1.
+_BAKED = {"NVTE_FLASH_ATTN": "0", "NVTE_FUSED_ATTN": "1"}
+
+
+def test_auto_respects_baked_flash_off(patched_set_attention_backend, monkeypatch):
+    # "auto" must NOT force NVTE_FLASH_ATTN=1 (that is what crashes stock
+    # megatron); it fills only the unset var, leaving the baked FLASH=0.
+    assert _run(monkeypatch, AttnBackend.auto, _BAKED) == ("0", "1", "1")
+
+
+def test_unfused_overrides_baked_fused(patched_set_attention_backend, monkeypatch):
+    # An explicit backend wins over the baked defaults (FUSED 1 -> 0), where
+    # stock megatron would assert-crash.
+    assert _run(monkeypatch, AttnBackend.unfused, _BAKED) == ("0", "0", "1")
+
+
+def test_fused_sets_expected_combination(patched_set_attention_backend, monkeypatch):
+    assert _run(monkeypatch, AttnBackend.fused, _BAKED) == ("0", "1", "0")
+
+
+def test_local_disables_all(patched_set_attention_backend, monkeypatch):
+    assert _run(monkeypatch, AttnBackend.local, _BAKED) == ("0", "0", "0")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py b/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py
new file mode 100644
index 000000000..05b3bc14c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py
@@ -0,0 +1,109 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for `_restore_chimera_rng_state` defensive RNG restoration.
+
+Validates Fix 2: when Megatron's private `_set_random_seed` API drifts
+(signature change or removal), the helper falls back to a manual restore
+that covers all three RNG generators — CPU default, CUDA default, and the
+model-parallel tracker — so chimera training does not fail at startup.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.flux_pretrain_trainer import _restore_chimera_rng_state
+
+
+def _canonical_args(seed=42):
+    """Build the args namespace the helper reads."""
+    return SimpleNamespace(
+        seed=seed,
+        data_parallel_random_init=False,
+        te_rng_tracker=False,
+        inference_rng_tracker=False,
+        enable_cuda_graph=False,
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+class TestChimeraRngRestoreFallback:
+    """The fallback path: when `_set_random_seed` raises TypeError or
+    ImportError, the helper must restore CPU, CUDA default, and the
+    model-parallel tracker generators manually so training continues."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        """Initialize parallel state — required because the fallback's
+        tp_random.model_parallel_cuda_manual_seed needs the tracker set up."""
+
+    def _verify_all_rngs_match_seed(self, seed):
+        """Sample from each generator and compare to a freshly-seeded
+        reference. If the helper's fallback fully restored the state,
+        the two samples must be bitwise identical for every generator."""
+        from megatron.core.tensor_parallel import random as tp_random
+
+        # CPU default generator.
+        cpu_after = torch.randn(8)
+        torch.manual_seed(seed)
+        cpu_ref = torch.randn(8)
+        assert torch.equal(cpu_after, cpu_ref), "CPU default generator was not restored to seed"
+
+        # CUDA default generator.
+        cuda_after = torch.randn(8, device="cuda")
+        torch.cuda.manual_seed(seed)
+        cuda_ref = torch.randn(8, device="cuda")
+        assert torch.equal(cuda_after, cuda_ref), "CUDA default generator was not restored to seed"
+
+        # Model-parallel tracker.
+        with tp_random.get_cuda_rng_tracker().fork():
+            tracker_after = torch.randn(8, device="cuda")
+        tp_random.model_parallel_cuda_manual_seed(seed)
+        with tp_random.get_cuda_rng_tracker().fork():
+            tracker_ref = torch.randn(8, device="cuda")
+        assert torch.equal(tracker_after, tracker_ref), "Model-parallel tracker was not restored to seed"
+
+    def _contaminate_rngs(self, contamination_seed):
+        """Set all three generators to a non-canonical seed so the test
+        observes the helper actively restoring rather than no-op-passing."""
+        from megatron.core.tensor_parallel import random as tp_random
+
+        torch.manual_seed(contamination_seed)
+        torch.cuda.manual_seed(contamination_seed)
+        tp_random.model_parallel_cuda_manual_seed(contamination_seed)
+
+    def test_typeerror_fallback_restores_all_rngs(self):
+        """Megatron signature drift (TypeError) must trigger manual restore."""
+        seed = 42
+        self._contaminate_rngs(contamination_seed=999)
+
+        args = _canonical_args(seed=seed)
+        with patch(
+            "megatron.training.initialize._set_random_seed",
+            side_effect=TypeError("signature changed"),
+        ):
+            _restore_chimera_rng_state(args)
+
+        self._verify_all_rngs_match_seed(seed)
+
+    def test_importerror_fallback_restores_all_rngs(self):
+        """Megatron module removal (ImportError) must trigger manual restore."""
+        seed = 42
+        self._contaminate_rngs(contamination_seed=999)
+
+        args = _canonical_args(seed=seed)
+        with patch(
+            "megatron.training.initialize._set_random_seed",
+            side_effect=ImportError("module gone"),
+        ):
+            _restore_chimera_rng_state(args)
+
+        self._verify_all_rngs_match_seed(seed)
diff --git a/tests/unit_tests/backends/megatron/test_compressor_pool.py b/tests/unit_tests/backends/megatron/test_compressor_pool.py
new file mode 100644
index 000000000..98682b09c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_compressor_pool.py
@@ -0,0 +1,58 @@
+# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+"""Parity test for the fused Compressor softmax-pool kernel.
+
+Asserts the fused ``fused_softmax_weighted_pool`` (Triton fwd + analytic eager bwd)
+matches the eager ``(softmax(score+ape, dim=2) * kv).sum(dim=2)`` it replaces in
+:meth:`Compressor.forward`, for both the forward output and the dkv / dscore / dape
+gradients. GPU-gated (the kernel is CUDA/Triton-only); fp32 is checked tightly and
+bf16 loosely (the fused path reduces in fp32 -> more accurate than eager bf16 weights).
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="fused compressor pool is CUDA/Triton only")
+
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.compressor_pool import (
+        fused_softmax_weighted_pool,
+    )
+
+    HAVE_KERNEL = True
+except Exception:  # triton import may fail on a CPU-only host
+    HAVE_KERNEL = False
+
+pytestmark = [cuda, pytest.mark.skipif(not HAVE_KERNEL, reason="compressor_pool/triton unavailable")]
+
+
+def _eager(kv, score, ape):
+    w = torch.softmax((score + ape).float(), dim=2).to(kv.dtype)
+    return (kv * w).sum(dim=2)
+
+
+def _rel(a, b):
+    return (a - b).float().norm() / b.float().norm().clamp_min(1e-12)
+
+
+# (B, N, W, hd): HCA (W=128, no overlap) and CSA (W=8, overlap)
+@pytest.mark.parametrize("shape", [(2, 16, 128, 128), (2, 64, 8, 128)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_fused_pool_matches_eager(shape, dtype):
+    B, N, W, HD = shape
+    torch.manual_seed(0)
+    dev = "cuda"
+    kv = torch.randn(B, N, W, HD, device=dev, dtype=dtype)
+    sc = torch.randn(B, N, W, HD, device=dev, dtype=dtype)
+    ape = torch.randn(W, HD, device=dev, dtype=dtype)
+    g = torch.randn(B, N, HD, device=dev, dtype=dtype)
+
+    ke, se, ae = (t.clone().requires_grad_(True) for t in (kv, sc, ape))
+    _eager(ke, se, ae).backward(g)
+    kf, sf, af = (t.clone().requires_grad_(True) for t in (kv, sc, ape))
+    fused_softmax_weighted_pool(kf, sf, af).backward(g)
+
+    tol = 1e-4 if dtype == torch.float32 else 2e-2
+    assert _rel(kf.grad, ke.grad) < tol
+    assert _rel(sf.grad, se.grad) < tol
+    assert _rel(af.grad, ae.grad) < tol
diff --git a/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py b/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py
new file mode 100644
index 000000000..ddfdb8912
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py
@@ -0,0 +1,647 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for ``deepseek_v4_flops_patches.py`` (Plan-3 Phase 20 + Plan-6 P33).
+
+Coverage:
+
+* G16 — :func:`compute_v4_flops` matches a hand-derived closed-form total
+  within 1% on a fully-specified V4-Flash-shaped config (8 layers, mixed
+  ``compress_ratios=[0,4,128]``, ``mtp_num_layers=1``).  Every per-component
+  byte (``attn_qkv_o``, ``attn_scores``, ``compressor``, ``indexer``,
+  ``moe``, ``mtp``, ``logits``, ``hc``) is asserted independently so a
+  regression in any single term fails loudly.
+* G17 — When the wrapper is installed with ``dispatch_v4=False`` (i.e. the
+  installer saw a non-V4 ``args.model_type``), the wrapper returns the
+  upstream value byte-for-byte.  This mirrors how the install-time
+  ``condition`` gate works: V4 dispatch is captured at install time
+  because Megatron's ``pretrain()`` overwrites ``args.model_type`` with
+  a ``ModelType`` enum at ``training.py:1210`` before ``train()`` calls
+  ``num_floating_point_operations``.
+* G36 — Plan-6 P33: SWA visible-pair correction.  Parametrised over
+  ``swa_window``, ``compress_ratio``, ``hc_mult`` so the dense + HCA +
+  CSA per-layer pair counts and the over-count ratio vs the legacy
+  ``S_eff^2`` upper bound are pinned independently.
+* G36a — Plan-6 P33: HyperConnection ``fn.weight`` matmul accounting.
+  Asserts the ``hc`` breakdown row equals the closed form
+  ``B * S * K * D * K * (2 * (L + M) * (2+K) + (1 + M))`` and degrades
+  to 0 when ``hc_mult <= 1``.
+"""
+
+from __future__ import annotations
+
+import types
+from types import SimpleNamespace
+
+import pytest
+
+from primus.backends.megatron.patches.deepseek_v4_flops_patches import (
+    _FMA_FACTOR,
+    _FORWARD_BACKWARD_FACTOR,
+    _SWIGLU_FFN_EXPANSION_FACTOR,
+    _make_v4_num_floating_point_operations,
+    _normalize_layer_ratios,
+    _visible_pairs,
+    compute_v4_flops,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _v4_flash_smoke_args(
+    *,
+    num_layers: int = 8,
+    seq_length: int = 128,
+    hc_mult: int = 4,
+    mtp_num_layers: int = 0,
+    compress_ratios=(0, 0, 4, 128, 4, 128, 4, 0),
+    num_hash_layers: int = 3,
+    moe_router_topk: int = 6,
+    num_experts: int = 256,
+    moe_ffn_hidden_size: int = 2048,
+    moe_shared_expert_intermediate_size: int = 2048,
+    attn_sliding_window: int = 0,
+):
+    """Build a V4-Flash-shaped fake ``args`` namespace.
+
+    Defaults mirror the smoke config used in P19 (run_deepseek_v4.sh)
+    so the unit numbers stay comparable to live runs.  ``attn_sliding_window``
+    defaults to ``0`` (disabled / full causal) for backward-compatibility
+    with the plan-3 P20 reference; SWA-aware behaviour is exercised by
+    the plan-6 P33 G36 tests below.
+    """
+    return SimpleNamespace(
+        model_type="deepseek_v4",
+        seq_length=seq_length,
+        hc_mult=hc_mult,
+        hidden_size=4096,
+        num_attention_heads=64,
+        kv_channels=512,
+        q_lora_rank=1024,
+        o_lora_rank=1024,
+        o_groups=8,
+        num_layers=num_layers,
+        mtp_num_layers=mtp_num_layers,
+        compress_ratios=compress_ratios,
+        moe_ffn_hidden_size=moe_ffn_hidden_size,
+        ffn_hidden_size=18432,
+        moe_router_topk=moe_router_topk,
+        num_experts=num_experts,
+        moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size,
+        num_hash_layers=num_hash_layers,
+        index_topk=512,
+        index_head_dim=128,
+        index_n_heads=64,
+        padded_vocab_size=129280,
+        vocab_size=129280,
+        attn_sliding_window=attn_sliding_window,
+    )
+
+
+def _hand_attn_qkv_o(*, B, S_eff, H, n, d, q_lora, o_lora, o_groups):
+    """Reference closed-form QKV+O FMAC per layer (matches the patch helper)."""
+    n_d = n * d
+    qkv = H * q_lora + q_lora * n_d + H * d
+    o_proj = n_d * o_lora + (o_groups * o_lora) * H if o_lora > 0 else n_d * H
+    return B * S_eff * (qkv + o_proj)
+
+
+def _hand_local_pairs(*, swa, S_eff):
+    """Reference closed form for SWA-pruned local visible pairs."""
+    if swa <= 0 or swa >= S_eff:
+        return S_eff * (S_eff + 1) // 2
+    return swa * S_eff - swa * (swa - 1) // 2
+
+
+def _hand_pool_pairs(*, ratio, S_eff):
+    """Reference closed form for the causal-visible HCA pool pair count."""
+    c = int(ratio)
+    if c <= 0 or S_eff <= 0:
+        return 0
+    n_full = S_eff // c
+    if n_full == 0:
+        return 0
+    return c * n_full * (n_full - 1) // 2 + n_full * (S_eff - c * n_full + 1)
+
+
+def _hand_attn_scores(*, B, S_eff, n, d, ratio, index_topk, swa):
+    """Reference closed-form attention-score FMAC per layer.
+
+    Plan-6 P33: counts only causal-visible ``(query, key)`` pairs
+    surviving the per-layer mask (SWA + pool + sparse top-K), not the
+    legacy ``S_eff^2`` upper bound.
+    """
+    local_pairs = _hand_local_pairs(swa=swa, S_eff=S_eff)
+    if ratio == 0:
+        pairs = local_pairs
+    elif ratio == 128:
+        pairs = local_pairs + _hand_pool_pairs(ratio=ratio, S_eff=S_eff)
+    elif ratio == 4:
+        pool = max(1, S_eff // 4)
+        keys = min(index_topk, pool) if index_topk else pool
+        pairs = local_pairs + keys * S_eff
+    else:
+        pool = max(1, S_eff // ratio)
+        pairs = local_pairs + pool * S_eff
+    return 2 * B * n * d * pairs
+
+
+def _hand_compressor(*, B, S_eff, H, d, ratio):
+    if ratio == 0:
+        return 0
+    coff = 2 if ratio == 4 else 1
+    return 2 * B * S_eff * H * (coff * d)
+
+
+def _hand_indexer(*, B, S_eff, H, ratio, ihd, inh):
+    if ratio != 4:
+        return 0
+    pool = max(1, S_eff // ratio)
+    proj = H * ihd + ihd * (inh * ihd) + H * inh + 2 * H * (2 * ihd)
+    scoring = inh * pool * ihd
+    return B * S_eff * (proj + scoring)
+
+
+def _hand_moe(*, B, S_eff, H, H_moe, topk, n_experts, hash_layer, H_shared):
+    router = 0 if hash_layer else H * n_experts
+    routed = topk * _SWIGLU_FFN_EXPANSION_FACTOR * H * H_moe
+    shared = _SWIGLU_FFN_EXPANSION_FACTOR * H * H_shared if H_shared > 0 else 0
+    return B * S_eff * (router + routed + shared)
+
+
+# ---------------------------------------------------------------------------
+# G16: closed-form parity per component
+# ---------------------------------------------------------------------------
+
+
+class TestComputeV4FlopsClosedForm:
+    """G16: per-component breakdown matches a hand-derived reference."""
+
+    @pytest.fixture(scope="class")
+    def args(self):
+        return _v4_flash_smoke_args()
+
+    @pytest.fixture(scope="class")
+    def batch_size(self):
+        return 16  # GBS used in the P19 smoke run.
+
+    @pytest.fixture(scope="class")
+    def computed(self, args, batch_size):
+        total, breakdown = compute_v4_flops(args, batch_size)
+        return total, breakdown
+
+    def test_attn_qkv_o_term_matches_reference(self, args, batch_size, computed):
+        _total, br = computed
+        S_eff = args.seq_length * args.hc_mult
+        per_layer = _hand_attn_qkv_o(
+            B=batch_size,
+            S_eff=S_eff,
+            H=args.hidden_size,
+            n=args.num_attention_heads,
+            d=args.kv_channels,
+            q_lora=args.q_lora_rank,
+            o_lora=args.o_lora_rank,
+            o_groups=args.o_groups,
+        )
+        assert br.attn_qkv_o == per_layer * args.num_layers
+
+    def test_attn_scores_term_matches_reference(self, args, batch_size, computed):
+        _total, br = computed
+        S_eff = args.seq_length * args.hc_mult
+        expected = sum(
+            _hand_attn_scores(
+                B=batch_size,
+                S_eff=S_eff,
+                n=args.num_attention_heads,
+                d=args.kv_channels,
+                ratio=int(r),
+                index_topk=args.index_topk,
+                swa=int(getattr(args, "attn_sliding_window", 0) or 0),
+            )
+            for r in args.compress_ratios
+        )
+        assert br.attn_scores == expected
+
+    def test_compressor_term_matches_reference(self, args, batch_size, computed):
+        _total, br = computed
+        S_eff = args.seq_length * args.hc_mult
+        expected = sum(
+            _hand_compressor(
+                B=batch_size,
+                S_eff=S_eff,
+                H=args.hidden_size,
+                d=args.kv_channels,
+                ratio=int(r),
+            )
+            for r in args.compress_ratios
+        )
+        assert br.compressor == expected
+
+    def test_indexer_term_matches_reference(self, args, batch_size, computed):
+        _total, br = computed
+        S_eff = args.seq_length * args.hc_mult
+        expected = sum(
+            _hand_indexer(
+                B=batch_size,
+                S_eff=S_eff,
+                H=args.hidden_size,
+                ratio=int(r),
+                ihd=args.index_head_dim,
+                inh=args.index_n_heads,
+            )
+            for r in args.compress_ratios
+        )
+        assert br.indexer == expected
+
+    def test_moe_term_respects_hash_layers(self, args, batch_size, computed):
+        _total, br = computed
+        S_eff = args.seq_length * args.hc_mult
+        expected = sum(
+            _hand_moe(
+                B=batch_size,
+                S_eff=S_eff,
+                H=args.hidden_size,
+                H_moe=args.moe_ffn_hidden_size,
+                topk=args.moe_router_topk,
+                n_experts=args.num_experts,
+                hash_layer=(layer_idx < args.num_hash_layers),
+                H_shared=args.moe_shared_expert_intermediate_size,
+            )
+            for layer_idx in range(args.num_layers)
+        )
+        assert br.moe == expected
+
+    def test_logits_term_includes_one_extra_head_per_mtp_depth(self):
+        args = _v4_flash_smoke_args(mtp_num_layers=2)
+        _total, br = compute_v4_flops(args, batch_size=8)
+        expected = (args.mtp_num_layers + 1) * 8 * args.seq_length * args.hidden_size * args.padded_vocab_size
+        assert br.logits == expected
+
+    def test_total_matches_breakdown_sum_with_expansion(self, computed):
+        total, br = computed
+        assert total == _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR * br.total_fmac()
+
+    def test_hc_term_is_nonzero_when_hc_mult_gt_1(self, computed):
+        """Plan-6 P33: HC matmul row must be populated when ``hc_mult > 1``.
+
+        The exact closed form is pinned by the dedicated G36a test below;
+        here we just assert the term is present so a future refactor
+        that silently zeroes ``hc`` fails this fixture too.
+        """
+        _total, br = computed
+        assert br.hc > 0
+
+
+# ---------------------------------------------------------------------------
+# G36: Plan-6 P33 SWA visible-pair correction
+# ---------------------------------------------------------------------------
+
+
+class TestG36SWAVisiblePairs:
+    """Plan-6 P33: ``_attn_scores_fmac_per_layer`` must count only causal-
+    visible ``(q, k)`` pairs surviving SWA + pool + sparse top-K masks.
+
+    The legacy plan-3 P20 closed form used ``B * n * d * S_eff^2`` for the
+    local branch (Megatron's ``S^2/2`` causal upper bound x the FMA-pair
+    factor) which over-counted by ``S_eff / swa_window`` once the kernel
+    started honoring SWA per-row pruning.  This test pins the new
+    ``2 * n * d * visible_pairs`` form against the helper, the per-branch
+    over-count ratios, and the proxy-shape values printed in
+    ``deepseek-v4/develop/perf/attention_perf.md``.
+    """
+
+    @pytest.mark.parametrize("swa", [0, 64, 128, 4096, 8192])
+    def test_local_visible_pair_helper(self, swa):
+        """Helper closed form matches the exhaustive sum-over-queries."""
+        S_eff = 4096
+        pairs = _visible_pairs(
+            swa_window=swa,
+            compress_ratio=0,
+            index_topk=0,
+            seq_len_eff=S_eff,
+        )
+        exhaustive = sum(min(q + 1, swa) if (0 < swa < S_eff) else (q + 1) for q in range(S_eff))
+        assert pairs == exhaustive
+
+    def test_proxy_shape_dense_visible_pairs_matches_attn_perf_doc(self):
+        """``swa=128, S_eff=4096, cr=0`` → 516,160 (attention_perf.md row)."""
+        pairs = _visible_pairs(
+            swa_window=128,
+            compress_ratio=0,
+            index_topk=0,
+            seq_len_eff=4096,
+        )
+        assert pairs == 516_160
+
+    def test_proxy_shape_hca_visible_pairs_matches_attn_perf_doc(self):
+        """``swa=128, S_eff=4096, cr=128`` → 516,160 + 63,520 = 579,680."""
+        pairs = _visible_pairs(
+            swa_window=128,
+            compress_ratio=128,
+            index_topk=0,
+            seq_len_eff=4096,
+        )
+        assert pairs == 516_160 + 63_520
+
+    def test_proxy_shape_csa_visible_pairs_matches_attn_perf_doc(self):
+        """``swa=128, S_eff=4096, cr=4, topk=512`` → 516,160 + 512*4096."""
+        pairs = _visible_pairs(
+            swa_window=128,
+            compress_ratio=4,
+            index_topk=512,
+            seq_len_eff=4096,
+        )
+        assert pairs == 516_160 + 512 * 4096
+
+    @pytest.mark.parametrize("hc_mult", [1, 4])
+    @pytest.mark.parametrize("ratio", [0, 4, 128])
+    def test_swa128_reduces_attn_scores_vs_full_causal(self, hc_mult, ratio):
+        """SWA=128 must produce strictly fewer score FMAC than swa=0 on
+        every layer type (no overlap between local and pool/topk terms
+        means the SWA-pruned local term strictly dominates the saving).
+        """
+        S = 4096
+        args_no_swa = _v4_flash_smoke_args(
+            num_layers=1,
+            seq_length=S,
+            hc_mult=hc_mult,
+            compress_ratios=(ratio,),
+            num_hash_layers=0,
+            attn_sliding_window=0,
+        )
+        args_swa = _v4_flash_smoke_args(
+            num_layers=1,
+            seq_length=S,
+            hc_mult=hc_mult,
+            compress_ratios=(ratio,),
+            num_hash_layers=0,
+            attn_sliding_window=128,
+        )
+        _t1, br_no = compute_v4_flops(args_no_swa, batch_size=1)
+        _t2, br_swa = compute_v4_flops(args_swa, batch_size=1)
+        assert br_swa.attn_scores < br_no.attn_scores
+        # All other components untouched by SWA.
+        assert br_swa.attn_qkv_o == br_no.attn_qkv_o
+        assert br_swa.compressor == br_no.compressor
+        assert br_swa.indexer == br_no.indexer
+        assert br_swa.moe == br_no.moe
+        assert br_swa.hc == br_no.hc
+
+    def test_swa_pruned_attn_scores_matches_proxy_overcount_ratio(self):
+        """Per-layer over-count ratio between legacy ``S_eff^2`` and the
+        SWA-pruned visible-pair count is ``S_eff / (2*swa) + O(1/S)`` for
+        swa < S_eff (the factor of 2 comes from counting both the QK^T
+        and PV matmul halves of the attention score).  At the V4-Flash
+        proxy shape (S=4096, hc_mult=4, S_eff=16384, swa=128) that ratio
+        is ``16384 / (2 * 128) = 64`` — pin it at >= 60x so a regression
+        that silently reverts to the legacy ``S_eff^2`` form is caught
+        while leaving 4-5x of headroom for off-by-one fringe corrections.
+        """
+        S_eff = 16384
+        swa = 128
+        legacy_local = S_eff * S_eff  # plan-3 P20 ``S_eff^2`` form
+        swa_local_pairs = swa * S_eff - swa * (swa - 1) // 2
+        new_local = 2 * swa_local_pairs  # 2 * visible_pairs (QK + PV)
+        ratio = legacy_local / new_local
+        assert ratio >= 60, f"SWA over-count ratio collapsed: {ratio:.2f}"
+
+
+# ---------------------------------------------------------------------------
+# G36a: Plan-6 P33 HyperConnection fn matmul accounting
+# ---------------------------------------------------------------------------
+
+
+class TestG36aHCMatmulAccounting:
+    """Plan-6 P33: the ``hc`` breakdown row must equal the closed form
+    ``B * S * K * D * K * (2 * (L + M) * (2 + K) + (1 + M))``.
+
+    Pinned independently from G16 because it's a brand-new component
+    and the formula has a non-obvious factor structure (2 mixers per
+    layer x (L+M) layers, plus 1 trunk head + M MTP heads).
+    """
+
+    @pytest.mark.parametrize("hc_mult", [2, 4, 8])
+    @pytest.mark.parametrize("mtp_num_layers", [0, 1, 2])
+    def test_hc_matmul_matches_closed_form(self, hc_mult, mtp_num_layers):
+        args = _v4_flash_smoke_args(
+            num_layers=4,
+            seq_length=128,
+            hc_mult=hc_mult,
+            mtp_num_layers=mtp_num_layers,
+        )
+        batch_size = 8
+        _total, br = compute_v4_flops(args, batch_size=batch_size)
+        B = batch_size
+        S = args.seq_length
+        K = hc_mult
+        D = args.hidden_size
+        L = args.num_layers
+        M = mtp_num_layers
+        expected = B * S * K * D * K * (2 * (L + M) * (2 + K) + (1 + M))
+        assert br.hc == expected
+
+    def test_hc_matmul_uses_seq_len_not_seq_len_eff(self):
+        """HyperMixer runs on the un-packed ``[B, S, K, D]`` tensor; cost
+        must scale with ``seq_len`` not ``seq_len * hc_mult``.
+
+        Concretely: doubling ``hc_mult`` from ``K=2 -> 4`` multiplies the
+        mixer factor ``K * D * K * (2+K)`` by ``(4*4*6) / (2*2*4) = 6``,
+        but doubling ``seq_length`` only doubles the per-layer cost.
+        Pinning both axes proves the closed form is keyed on the right
+        sequence-axis variable.
+        """
+        ref = _v4_flash_smoke_args(num_layers=2, seq_length=64, hc_mult=2, mtp_num_layers=0)
+        ref_double_s = _v4_flash_smoke_args(num_layers=2, seq_length=128, hc_mult=2, mtp_num_layers=0)
+        ref_double_k = _v4_flash_smoke_args(num_layers=2, seq_length=64, hc_mult=4, mtp_num_layers=0)
+        _t1, br_ref = compute_v4_flops(ref, batch_size=1)
+        _t2, br_s = compute_v4_flops(ref_double_s, batch_size=1)
+        _t3, br_k = compute_v4_flops(ref_double_k, batch_size=1)
+        assert br_s.hc == 2 * br_ref.hc
+
+        # Doubling K from 2 -> 4 scales mixer per-layer cost as
+        # K*K*(2+K) -> (4*4*6) / (2*2*4) = 6 and head cost as
+        # K*K -> 16/4 = 4.  With L=2 mixers (= 4) and 1 head,
+        # the aggregate multiplier is (4 * 6 + 4) / (4 * 1 + 1) = 28/5.
+        # (L=2, M=0 → mixer term = 2*L*(2+K)*K^2*D = ratio 6;
+        #  head term = (1+M)*K^2*D = ratio 4)
+        # Total ratio: (2*L*(2+K_new)*K_new^2 + (1+M)*K_new^2) /
+        #             (2*L*(2+K_old)*K_old^2 + (1+M)*K_old^2)
+        L = 2
+        M = 0
+        K_old = 2
+        K_new = 4
+        num = 2 * L * (2 + K_new) * K_new * K_new + (1 + M) * K_new * K_new
+        den = 2 * L * (2 + K_old) * K_old * K_old + (1 + M) * K_old * K_old
+        assert br_k.hc * den == br_ref.hc * num
+
+
+# ---------------------------------------------------------------------------
+# Hash-layer / no-hash variant
+# ---------------------------------------------------------------------------
+
+
+class TestHashLayerHandling:
+    """Hash-routed layers must skip the topk router GEMM."""
+
+    def test_zero_hash_layers_charges_router_on_every_moe_layer(self):
+        args = _v4_flash_smoke_args(num_hash_layers=0)
+        _, with_hash = compute_v4_flops(args, batch_size=4)
+
+        args_no_hash = _v4_flash_smoke_args(num_hash_layers=args.num_layers)
+        _, all_hash = compute_v4_flops(args_no_hash, batch_size=4)
+
+        # All-hash strictly less because router cost is dropped on every layer.
+        assert all_hash.moe < with_hash.moe
+        delta_per_layer = 4 * args.seq_length * args.hc_mult * args.hidden_size * args.num_experts
+        assert with_hash.moe - all_hash.moe == delta_per_layer * args.num_layers
+
+
+# ---------------------------------------------------------------------------
+# compress_ratios normalization (decoder + MTP slicing)
+# ---------------------------------------------------------------------------
+
+
+class TestNormalizeLayerRatios:
+    def test_string_yaml_form_parses(self):
+        decoder, mtp = _normalize_layer_ratios("[0, 0, 4, 128, 4, 0]", num_layers=6, mtp_num_layers=0)
+        assert decoder == [0, 0, 4, 128, 4, 0]
+        assert mtp == []
+
+    def test_decoder_plus_mtp_layout_splits_correctly(self):
+        decoder, mtp = _normalize_layer_ratios([0, 4, 128, 4, 0], num_layers=4, mtp_num_layers=1)
+        assert decoder == [0, 4, 128, 4]
+        assert mtp == [0]
+
+    def test_none_defaults_to_all_dense(self):
+        decoder, mtp = _normalize_layer_ratios(None, num_layers=3, mtp_num_layers=2)
+        assert decoder == [0, 0, 0]
+        assert mtp == [0, 0]
+
+    def test_short_list_pads_with_last_value(self):
+        decoder, mtp = _normalize_layer_ratios([4, 128], num_layers=4, mtp_num_layers=0)
+        assert decoder == [4, 128, 128, 128]
+        assert mtp == []
+
+
+# ---------------------------------------------------------------------------
+# G17: non-V4 fall-through byte-for-byte
+# ---------------------------------------------------------------------------
+
+
+class TestDispatchSwitch:
+    """G17: ``dispatch_v4`` flag controls whether upstream is called."""
+
+    @pytest.fixture(autouse=True)
+    def _silence_breakdown_log(self, monkeypatch):
+        """The wrapper emits a one-shot breakdown via ``log_rank_0`` on the
+        first V4 call.  Under pytest there's no torch.distributed bound so
+        rank-aware logging would explode; flip the once-only flag to ``True``
+        so the wrapper skips the emit path.
+        """
+        from primus.backends.megatron.patches import deepseek_v4_flops_patches as mod
+
+        monkeypatch.setattr(mod, "_BREAKDOWN_LOGGED", True)
+
+    @pytest.fixture
+    def fake_upstream_factory(self):
+        """Returns ``(make_wrapped, sentinel_calls)`` per-test."""
+        sentinel_calls = []
+
+        def fake_upstream(args, batch_size):
+            sentinel_calls.append((id(args), batch_size))
+            return 12345 + batch_size + len(getattr(args, "model_type", "") or "")
+
+        def make_wrapped(*, dispatch_v4: bool):
+            return _make_v4_num_floating_point_operations(fake_upstream, dispatch_v4=dispatch_v4)
+
+        return make_wrapped, sentinel_calls
+
+    @pytest.mark.parametrize(
+        "model_type",
+        ["gpt", "llama3", "deepseek_v3", "mamba_hybrid", None, ""],
+    )
+    def test_dispatch_v4_false_falls_through_byte_for_byte(self, fake_upstream_factory, model_type):
+        make_wrapped, sentinel_calls = fake_upstream_factory
+        wrapped = make_wrapped(dispatch_v4=False)
+        args = SimpleNamespace(model_type=model_type)
+        result = wrapped(args, 32)
+        assert result == 12345 + 32 + len(model_type or "")
+        assert sentinel_calls == [(id(args), 32)]
+
+    def test_dispatch_v4_true_uses_v4_closed_form(self, fake_upstream_factory):
+        make_wrapped, sentinel_calls = fake_upstream_factory
+        wrapped = make_wrapped(dispatch_v4=True)
+        args = _v4_flash_smoke_args()
+        result = wrapped(args, batch_size=4)
+        assert result > 0
+        assert sentinel_calls == []  # Upstream MUST NOT be called when dispatching V4.
+
+    def test_dispatch_v4_true_ignores_runtime_model_type_mutation(self, fake_upstream_factory):
+        """Megatron's pretrain() rewrites ``args.model_type`` with the
+        ``ModelType`` enum just before ``train()`` runs.  The wrapper must
+        keep dispatching to V4 anyway because the install-time decision is
+        captured via the closure flag.
+        """
+        from enum import Enum
+
+        class _FakeModelType(Enum):
+            encoder_or_decoder = 1
+
+        make_wrapped, sentinel_calls = fake_upstream_factory
+        wrapped = make_wrapped(dispatch_v4=True)
+        args = _v4_flash_smoke_args()
+        args.model_type = _FakeModelType.encoder_or_decoder  # post-pretrain() state
+        result = wrapped(args, batch_size=4)
+        assert result > 0
+        assert sentinel_calls == []
+
+
+# ---------------------------------------------------------------------------
+# Patch installation lifecycle
+# ---------------------------------------------------------------------------
+
+
+class TestPatchInstallation:
+    def test_idempotent_installation(self, monkeypatch):
+        """Re-applying the patch with an already-wrapped target is a no-op."""
+        from primus.backends.megatron.patches import deepseek_v4_flops_patches as mod
+        from primus.core.patches.context import PatchContext
+
+        def upstream(args, bs):
+            return 0
+
+        wrapped_once = mod._make_v4_num_floating_point_operations(upstream, dispatch_v4=True)
+
+        # ``patch_v4_flops_reporting`` does ``import megatron.training.training``,
+        # which under bare pytest pulls in the real megatron package tree.
+        # Stub out every level so the import resolves to our fake module.
+        import sys
+
+        fake_megatron = types.ModuleType("megatron")
+        fake_megatron_training_pkg = types.ModuleType("megatron.training")
+        fake_megatron_training_mod = types.ModuleType("megatron.training.training")
+        fake_megatron_training_mod.num_floating_point_operations = wrapped_once
+        fake_megatron_training_pkg.training = fake_megatron_training_mod
+        fake_megatron.training = fake_megatron_training_pkg
+
+        monkeypatch.setitem(sys.modules, "megatron", fake_megatron)
+        monkeypatch.setitem(sys.modules, "megatron.training", fake_megatron_training_pkg)
+        monkeypatch.setitem(sys.modules, "megatron.training.training", fake_megatron_training_mod)
+
+        # Silence rank-aware logger; primus' singleton ``_logger`` isn't bound
+        # under bare pytest so we replace ``log_rank_0`` with a no-op.
+        monkeypatch.setattr(mod, "log_rank_0", lambda *a, **kw: None)
+
+        ctx = PatchContext(
+            backend="megatron",
+            phase="before_train",
+            extra={"module_config": types.SimpleNamespace(params=_v4_flash_smoke_args())},
+        )
+        mod.patch_v4_flops_reporting(ctx)
+        # The wrapped function must remain the same instance (no double-wrap).
+        assert fake_megatron_training_mod.num_floating_point_operations is wrapped_once
diff --git a/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py b/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py
new file mode 100644
index 000000000..44f7cc798
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py
@@ -0,0 +1,161 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for DiffusionPretrainTrainer's forward step count.
+
+Validates Fix 1: the `_forward_step_count` instance variable and its
+checkpoint-compatible lazy initialization from `args.iteration *
+get_num_microbatches()`. Also verifies the counter flows through to
+`flux_forward_step_func` as the `step_count` keyword argument.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import torch
+
+from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer
+
+
+class _ConcreteDiffusionTrainer(DiffusionPretrainTrainer):
+    """Minimal concrete subclass: implements abstract methods as no-ops.
+
+    These methods are only required for instantiation; the tests never
+    invoke them.
+    """
+
+    def create_model(self, *args, **kwargs):
+        return None
+
+    def create_scheduler(self, *args, **kwargs):
+        return None
+
+    def get_task_encoder(self, *args, **kwargs):
+        return None
+
+
+def _make_trainer():
+    """Construct a DiffusionPretrainTrainer with __init__ bypassed.
+
+    Sets only the attributes required by `get_forward_step()` and
+    `forward_step()`. Other attributes are intentionally left unset so
+    that tests fail loudly if production code starts depending on them.
+    """
+    trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer)
+    trainer._forward_step_count = 0
+    trainer._forward_step_count_initialized = False
+    return trainer
+
+
+class TestForwardStepCountLazyInit:
+    """Tests for lazy initialization of _forward_step_count from checkpoint state."""
+
+    def test_lazy_init_from_checkpoint_iteration(self):
+        """First closure call must reconstruct counter as iteration * num_microbatches."""
+        trainer = _make_trainer()
+        trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None
+
+        with patch(
+            "megatron.training.get_args",
+            return_value=SimpleNamespace(iteration=100),
+        ), patch(
+            "megatron.core.num_microbatches_calculator.get_num_microbatches",
+            return_value=4,
+        ):
+            closure = trainer.get_forward_step()
+            closure(data_iterator=None, model=None)
+
+        assert trainer._forward_step_count == 400
+        assert trainer._forward_step_count_initialized is True
+
+    def test_lazy_init_only_runs_once(self):
+        """Second call must not re-read args.iteration even if it changes."""
+        trainer = _make_trainer()
+        trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None
+
+        args_state = SimpleNamespace(iteration=10)
+
+        with patch(
+            "megatron.training.get_args",
+            side_effect=lambda: args_state,
+        ), patch(
+            "megatron.core.num_microbatches_calculator.get_num_microbatches",
+            return_value=2,
+        ):
+            closure = trainer.get_forward_step()
+            closure(data_iterator=None, model=None)
+            assert trainer._forward_step_count == 20
+
+            args_state.iteration = 999
+            closure(data_iterator=None, model=None)
+
+        # Must remain 20: lazy init flag prevents a second reconstruction.
+        assert trainer._forward_step_count == 20
+        assert trainer._forward_step_count_initialized is True
+
+    def test_initial_iteration_zero(self):
+        """Fresh start (iteration=0) reconstructs counter to 0."""
+        trainer = _make_trainer()
+        trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None
+
+        with patch(
+            "megatron.training.get_args",
+            return_value=SimpleNamespace(iteration=0),
+        ), patch(
+            "megatron.core.num_microbatches_calculator.get_num_microbatches",
+            return_value=8,
+        ):
+            closure = trainer.get_forward_step()
+            closure(data_iterator=None, model=None)
+
+        assert trainer._forward_step_count == 0
+        assert trainer._forward_step_count_initialized is True
+
+
+class TestCounterFlowsToForwardStepFunc:
+    """Integration test: trainer counter must pass through to flux_forward_step_func."""
+
+    def test_counter_flows_to_forward_step_func(self):
+        """forward_step must invoke flux_forward_step_func with step_count == counter."""
+        trainer = _make_trainer()
+        # _scheduler backs the `scheduler` lazy property; setting it pre-empts
+        # create_scheduler() which would otherwise raise abstract NotImplementedError.
+        trainer._scheduler = None
+
+        # Provide a fake runtime_state so forward_step doesn't fall into the
+        # log_rank_0 warning branch (which requires the Primus logger to be
+        # initialized — outside this test's scope).
+        class _FakeRuntimeState:
+            def update_metrics(self, metrics):
+                pass
+
+        trainer.runtime_state = _FakeRuntimeState()
+
+        recorded_kwargs = {}
+
+        def recording_func(*args, **kwargs):
+            recorded_kwargs.update(kwargs)
+            # Return shape matches: (noise_pred, clean_latents, noise, loss_mask, metrics, is_validation)
+            t = torch.zeros(1)
+            return t, t, t, None, {}, False
+
+        # Training model: `if model.training:` gate in forward_step requires
+        # a real attribute (None would AttributeError). Eval-skip behavior is
+        # covered in test_forward_step_count_gate.py.
+        train_model = Mock()
+        train_model.training = True
+
+        with patch(
+            "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func",
+            side_effect=recording_func,
+        ):
+            # First forward step — counter increments from 0 to 1
+            trainer.forward_step(data_iterator=None, model=train_model)
+            assert recorded_kwargs["step_count"] == 1
+            assert trainer._forward_step_count == 1
+
+            # Second forward step — counter increments to 2
+            trainer.forward_step(data_iterator=None, model=train_model)
+            assert recorded_kwargs["step_count"] == 2
+            assert trainer._forward_step_count == 2
diff --git a/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py b/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py
new file mode 100644
index 000000000..3d23df1cb
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py
@@ -0,0 +1,78 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for fp4_utils.py MXFP4 recipe and context manager changes.
+
+Tests recipe error handling and gradient stochastic-rounding configuration.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.utils import PrimusUT
+
+
+class TestMXFP4GradientStochasticRounding:
+    """Verify Megatron and diffusion configs resolve the SR option correctly."""
+
+    def test_explicit_config_value_takes_precedence(self, monkeypatch):
+        import megatron.training.global_vars as global_vars
+
+        from primus.backends.megatron.core.fp4_utils import _mxfp4_gradient_sr_enabled
+
+        def unexpected_get_args():
+            raise AssertionError("global args should not be read for an explicit config value")
+
+        monkeypatch.setattr(global_vars, "get_args", unexpected_get_args)
+
+        assert _mxfp4_gradient_sr_enabled(SimpleNamespace(mxfp4_gradient_stochastic_rounding=True))
+        assert not _mxfp4_gradient_sr_enabled(SimpleNamespace(mxfp4_gradient_stochastic_rounding=False))
+
+    def test_megatron_config_falls_back_to_global_args(self, monkeypatch):
+        import megatron.training.global_vars as global_vars
+
+        from primus.backends.megatron.core.fp4_utils import _mxfp4_gradient_sr_enabled
+
+        args = SimpleNamespace(mxfp4_gradient_stochastic_rounding=True)
+        monkeypatch.setattr(global_vars, "get_args", lambda: args)
+
+        assert _mxfp4_gradient_sr_enabled(SimpleNamespace())
+
+    def test_unavailable_global_args_default_to_disabled(self, monkeypatch):
+        import megatron.training.global_vars as global_vars
+
+        from primus.backends.megatron.core.fp4_utils import _mxfp4_gradient_sr_enabled
+
+        def unavailable_get_args():
+            raise RuntimeError("global args are not initialized")
+
+        monkeypatch.setattr(global_vars, "get_args", unavailable_get_args)
+
+        assert not _mxfp4_gradient_sr_enabled(SimpleNamespace())
+
+
+class TestGetFp4RecipeMXFP4(PrimusUT):
+    """Verify get_fp4_recipe returns correct recipe objects for MXFP4."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def test_unsupported_recipe_produces_error(self):
+        pytest.importorskip("transformer_engine")
+
+        from primus.backends.megatron.core.fp4_utils import get_fp4_recipe
+
+        config = SimpleNamespace(fp4_recipe="nonexistent_recipe")
+        result = get_fp4_recipe(config)
+
+        if isinstance(result, tuple):
+            recipe, reason = result
+            assert recipe is None, "Unsupported recipe should return None"
+            assert (
+                "Unsupported" in reason or "unsupported" in reason.lower()
+            ), f"Expected 'Unsupported' in reason, got: {reason}"
+        else:
+            pytest.fail("HAVE_TE-only branch should raise ValueError for unsupported recipe")
diff --git a/tests/unit_tests/backends/megatron/test_gdn_rocm_gate_patches.py b/tests/unit_tests/backends/megatron/test_gdn_rocm_gate_patches.py
new file mode 100644
index 000000000..e84d68de9
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_gdn_rocm_gate_patches.py
@@ -0,0 +1,115 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the Megatron upstream GatedDeltaNet ROCm gate patch.
+
+Verifies:
+  1. The source anchor still matches upstream ``GatedDeltaNet.forward``.
+  2. The patch rewrites the forward body and is idempotent.
+  3. Registration condition requires both ROCm and ``gated_delta_net``.
+"""
+
+import inspect
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("megatron")
+
+from megatron.core.ssm.gated_delta_net import GatedDeltaNet
+
+import primus.backends.megatron.patches.gdn_rocm_gate_patches as patch_mod
+from primus.backends.megatron.patches._patch_guard import is_patched
+from primus.core.patches.context import PatchContext
+
+
+def _clear_patch_state(cls):
+    if hasattr(cls, "_primus_applied_patch_keys"):
+        cls._primus_applied_patch_keys.discard(patch_mod._PATCH_KEY)
+
+
+@pytest.fixture
+def pristine_gdn_forward():
+    """Restore ``GatedDeltaNet.forward`` and patch-guard state after each test."""
+    original = GatedDeltaNet.forward
+    yield
+    GatedDeltaNet.forward = original
+    _clear_patch_state(GatedDeltaNet)
+
+
+def test_gate_anchor_matches_upstream_forward():
+    source = inspect.getsource(GatedDeltaNet.forward)
+    assert (
+        patch_mod._GATE_ORI in source
+    ), "Upstream GatedDeltaNet.forward changed; update gdn_rocm_gate_patches anchor."
+
+
+def test_gate_new_differs_from_anchor():
+    assert patch_mod._GATE_NEW != patch_mod._GATE_ORI
+    assert "use_gate_in_kernel=False" in patch_mod._GATE_NEW
+    assert "A_log.float().exp()" in patch_mod._GATE_NEW
+
+
+def test_install_patch_rewrites_forward(pristine_gdn_forward, monkeypatch):
+    monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None)
+    original = GatedDeltaNet.forward
+    original_source = inspect.getsource(original)
+
+    patch_mod._install_gdn_rocm_gate_patch()
+
+    assert is_patched(GatedDeltaNet, patch_mod._PATCH_KEY)
+    assert GatedDeltaNet.forward is not original
+
+    # exec()-patched methods are not inspectable; verify the expected rewrite.
+    rewritten = original_source.replace(patch_mod._GATE_ORI, patch_mod._GATE_NEW)
+    assert "use_gate_in_kernel=False" in rewritten
+    assert "A_log.float().exp()" in rewritten
+
+
+def test_install_patch_is_idempotent(pristine_gdn_forward, monkeypatch):
+    monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None)
+
+    patch_mod._install_gdn_rocm_gate_patch()
+    first = GatedDeltaNet.forward
+    patch_mod._install_gdn_rocm_gate_patch()
+
+    assert GatedDeltaNet.forward is first
+
+
+@pytest.mark.parametrize(
+    "variant, expected",
+    [
+        ("gated_delta_net", True),
+        ("dsa", False),
+        (None, False),
+    ],
+)
+def test_uses_gated_delta_net(variant, expected):
+    args = SimpleNamespace(experimental_attention_variant=variant)
+    assert patch_mod._uses_gated_delta_net(args) is expected
+
+
+@pytest.mark.parametrize(
+    "hip, variant, expected",
+    [
+        ("6.2.41133-65d174c3", "gated_delta_net", True),
+        ("6.2.41133-65d174c3", "dsa", False),
+        (None, "gated_delta_net", False),
+    ],
+)
+def test_should_patch_gdn_rocm_gate(hip, variant, expected, monkeypatch):
+    monkeypatch.setattr(patch_mod.torch.version, "hip", hip, raising=False)
+    args = SimpleNamespace(experimental_attention_variant=variant)
+    ctx = PatchContext(
+        backend="megatron",
+        phase="before_train",
+        extra={"module_config": SimpleNamespace(params=args)},
+    )
+    assert patch_mod._should_patch_gdn_rocm_gate(ctx) is expected
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/test_megatron_adapter.py b/tests/unit_tests/backends/megatron/test_megatron_adapter.py
index 66d610f52..a6c8ba1ef 100644
--- a/tests/unit_tests/backends/megatron/test_megatron_adapter.py
+++ b/tests/unit_tests/backends/megatron/test_megatron_adapter.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -123,63 +123,6 @@ def test_detect_version_not_found_raises_error(self, tmp_path, monkeypatch):
         assert "Cannot locate" in str(exc_info.value)
 
 
-class TestMegatronAdapterConfigConversion:
-    """Test configuration conversion from Primus to Megatron."""
-
-    @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder")
-    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
-    def test_convert_config_basic(self, mock_log, mock_builder_class):
-        """Test basic config conversion workflow."""
-        # Setup mock builder
-        mock_builder = Mock()
-        mock_builder_class.return_value = mock_builder
-
-        # Mock finalize to return SimpleNamespace with args
-        mock_args = SimpleNamespace(
-            micro_batch_size=4,
-            global_batch_size=32,
-            seq_length=2048,
-        )
-        mock_builder.finalize.return_value = mock_args
-
-        # Create mock params
-        mock_params = {
-            "micro_batch_size": 4,
-            "global_batch_size": 32,
-            "seq_length": 2048,
-        }
-
-        # Test conversion
-        adapter = MegatronAdapter()
-        result = adapter.convert_config(mock_params)
-
-        # Verify builder was called correctly
-        mock_builder.update.assert_called_once_with(mock_params)
-        mock_builder.finalize.assert_called_once()
-
-        # Verify result
-        assert result == mock_args
-        assert result.micro_batch_size == 4
-        assert result.global_batch_size == 32
-        assert result.seq_length == 2048
-
-    @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder")
-    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
-    def test_convert_config_empty_params(self, mock_log, mock_builder_class):
-        """Test config conversion with empty params dict."""
-        mock_builder = Mock()
-        mock_builder_class.return_value = mock_builder
-        mock_builder.finalize.return_value = SimpleNamespace()
-
-        mock_params = {}
-
-        adapter = MegatronAdapter()
-        result = adapter.convert_config(mock_params)
-
-        mock_builder.update.assert_called_once_with(mock_params)
-        assert isinstance(result, SimpleNamespace)
-
-
 class TestMegatronAdapterTrainerLoading:
     """Test trainer class loading."""
 
@@ -214,37 +157,133 @@ def test_load_trainer_class_invalid_stage(self):
         assert "backend trainer not registered" in str(exc_info.value)
 
 
-class TestMegatronAdapterBackendPreparation:
-    """Test backend preparation workflow."""
+class TestMegatronAdapterDynamicTrainerLoading:
+    """Test dynamic trainer class loading by name."""
 
-    @patch("primus.modules.module_utils.log_rank_0")
-    def test_prepare_backend_success(self, mock_log):
-        """Test successful backend preparation."""
+    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
+    def test_load_trainer_class_by_name_takes_precedence_over_stage(self, mock_log):
+        """Test that trainer_class parameter takes precedence over stage."""
         adapter = MegatronAdapter()
-        mock_config = Mock()
 
-        # Should not raise - prepare_backend is inherited from BackendAdapter
-        adapter.prepare_backend(mock_config)
+        # Even with invalid stage, trainer_class should work
+        result = adapter.load_trainer_class(stage="invalid_stage", trainer_class="MegatronPretrainTrainer")
 
+        from primus.backends.megatron.megatron_pretrain_trainer import (
+            MegatronPretrainTrainer,
+        )
 
-class TestMegatronAdapterInitialization:
-    """Test MegatronAdapter initialization."""
+        assert result == MegatronPretrainTrainer
+
+    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
+    def test_load_trainer_class_fallback_tries_multiple_paths(self, mock_log):
+        """Test fallback tries multiple import paths in order for unregistered trainers."""
+        import importlib
 
-    def test_initialization_default(self):
-        """Test adapter initialization with default framework name."""
+        # Use a trainer name NOT in MEGATRON_TRAINERS to exercise the fallback
+        mock_trainer_class = type("CustomExperimentalTrainer", (), {})
+        mock_module = Mock()
+        mock_module.CustomExperimentalTrainer = mock_trainer_class
+
+        import_calls = []
+
+        def side_effect(module_name, *args, **kwargs):
+            import_calls.append(module_name)
+            if len(import_calls) <= 1:
+                raise ImportError(f"Path {len(import_calls)} failed")
+            return mock_module
+
+        with patch.object(importlib, "import_module", side_effect=side_effect):
+            adapter = MegatronAdapter()
+            result = adapter.load_trainer_class(trainer_class="CustomExperimentalTrainer")
+
+        assert result == mock_trainer_class
+        assert len(import_calls) == 2
+        assert "primus.backends.megatron.customexperimentaltrainer" in import_calls[0].lower()
+
+    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
+    def test_load_trainer_class_fallback_all_paths_fail(self, mock_log):
+        """Test error when all fallback paths fail."""
+        import importlib
+
+        # All import attempts fail
+        def side_effect(module_name, *args, **kwargs):
+            raise ImportError("Module not found")
+
+        with patch.object(importlib, "import_module", side_effect=side_effect):
+            adapter = MegatronAdapter()
+
+            with pytest.raises(ValueError) as exc_info:
+                adapter.load_trainer_class(trainer_class="NonExistentTrainer")
+
+        error_msg = str(exc_info.value)
+        # Verify exact error message format from implementation
+        assert "Trainer class 'NonExistentTrainer' not found" in error_msg
+        assert "Available trainers:" in error_msg
+        assert "Hint:" in error_msg
+        assert "MEGATRON_TRAINERS" in error_msg or "standard locations" in error_msg
+
+    @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
+    def test_load_trainer_class_falsy_values_fall_back_to_stage(self, mock_log):
+        """Test that falsy trainer_class values (None, empty string) fall back to stage-based selection."""
         adapter = MegatronAdapter()
-        assert adapter.framework == "megatron"
 
-    def test_initialization_custom_framework(self):
-        """Test adapter initialization with custom framework name."""
-        adapter = MegatronAdapter(framework="custom_megatron")
-        assert adapter.framework == "custom_megatron"
+        from primus.backends.megatron.megatron_pretrain_trainer import (
+            MegatronPretrainTrainer,
+        )
+
+        # Test None falls back to stage
+        result = adapter.load_trainer_class(stage="pretrain", trainer_class=None)
+        assert result == MegatronPretrainTrainer
+
+        # Test empty string falls back to stage (falsy)
+        result = adapter.load_trainer_class(stage="pretrain", trainer_class="")
+        assert result == MegatronPretrainTrainer
+
+    def test_load_trainer_class_registry_import_error_provides_context(self):
+        """Test that import errors from registry provide helpful context."""
+        import importlib
+
+        def side_effect(module_name, *args, **kwargs):
+            raise ImportError("Module not found")
+
+        with patch.object(importlib, "import_module", side_effect=side_effect):
+            adapter = MegatronAdapter()
+
+            with pytest.raises(ImportError) as exc_info:
+                adapter.load_trainer_class(trainer_class="MegatronPretrainTrainer")
+
+        error_msg = str(exc_info.value)
+        # Verify exact error message format from implementation
+        assert "Failed to load trainer class 'MegatronPretrainTrainer'" in error_msg
+        assert "Hint:" in error_msg
+        assert "Check that the module exists" in error_msg
+
+    def test_load_trainer_class_registry_attribute_error_provides_context(self):
+        """Test that attribute errors from registry provide helpful context."""
+        import importlib
+
+        # Module imports but class doesn't exist
+        mock_module = Mock()
+        del mock_module.MegatronPretrainTrainer  # Ensure attribute doesn't exist
+
+        def side_effect(module_name, *args, **kwargs):
+            return mock_module
+
+        with patch.object(importlib, "import_module", side_effect=side_effect):
+            adapter = MegatronAdapter()
+
+            with pytest.raises(ImportError) as exc_info:
+                adapter.load_trainer_class(trainer_class="MegatronPretrainTrainer")
+
+        error_msg = str(exc_info.value)
+        assert "Failed to load trainer class" in error_msg
+        assert "Hint:" in error_msg
 
 
 class TestMegatronAdapterIntegration:
     """Integration tests for complete adapter workflow."""
 
-    @patch("primus.modules.module_utils.log_rank_0")
+    @patch("primus.core.utils.module_utils.log_rank_0")
     @patch("primus.backends.megatron.megatron_adapter.MegatronAdapter.detect_backend_version")
     @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder")
     @patch("primus.backends.megatron.megatron_adapter.log_rank_0")
diff --git a/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py b/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py
index 8b179d9f0..dcfd4a3ff 100644
--- a/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py
+++ b/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py
@@ -15,6 +15,7 @@
     5. Integration: complete workflow from config to final namespace
 """
 
+import enum
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -23,6 +24,18 @@
 from primus.backends.megatron.argument_builder import MegatronArgBuilder
 
 
+class _FakeAttnBackend(enum.Enum):
+    """Stand-in for Megatron's AttnBackend enum (plain int-valued enum)."""
+
+    flash = 1
+    fused = 2
+    auto = 3
+
+
+def _fake_enum_type(value: str) -> "_FakeAttnBackend":
+    return _FakeAttnBackend[value]
+
+
 class TestMegatronArgBuilderFiltering:
     """Test parameter filtering: only Megatron params are accepted."""
 
@@ -328,5 +341,103 @@ def test_chained_updates(self, mock_load_defaults, mock_dist_info):
         assert result.hidden_size == 4096
 
 
+class TestMegatronArgBuilderEnumCoercion:
+    """Test that enum override strings are coerced via Megatron's argparse converter.
+
+    Primus bypasses Megatron's argparse, so without coercion an enum arg like
+    ``attention_backend`` would arrive (and stay) a raw string and silently
+    break every downstream ``== AttnBackend.x`` comparison.
+    """
+
+    @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types")
+    @patch("primus.backends.megatron.argument_builder._load_megatron_defaults")
+    def test_string_coerced_to_enum(self, mock_defaults, mock_types):
+        mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto}
+        mock_types.return_value = {"attention_backend": _fake_enum_type}
+
+        builder = MegatronArgBuilder()
+        builder.update({"attention_backend": "fused"})
+
+        assert builder.overrides["attention_backend"] is _FakeAttnBackend.fused
+
+    @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types")
+    @patch("primus.backends.megatron.argument_builder._load_megatron_defaults")
+    def test_auto_coerced_to_enum(self, mock_defaults, mock_types):
+        # Every enum value is coerced consistently, including "auto" (the ROCm
+        # incompatibility of megatron-core's auto branch is handled by a patch,
+        # not by leaving this a raw string).
+        mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto}
+        mock_types.return_value = {"attention_backend": _fake_enum_type}
+
+        builder = MegatronArgBuilder()
+        builder.update({"attention_backend": "auto"})
+
+        assert builder.overrides["attention_backend"] is _FakeAttnBackend.auto
+
+    @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types")
+    @patch("primus.backends.megatron.argument_builder._load_megatron_defaults")
+    def test_list_coerced_elementwise(self, mock_defaults, mock_types):
+        mock_defaults.return_value = {"some_enum_list": []}
+        mock_types.return_value = {"some_enum_list": _fake_enum_type}
+
+        builder = MegatronArgBuilder()
+        builder.update({"some_enum_list": ["flash", "fused"]})
+
+        assert builder.overrides["some_enum_list"] == [
+            _FakeAttnBackend.flash,
+            _FakeAttnBackend.fused,
+        ]
+
+    @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types")
+    @patch("primus.backends.megatron.argument_builder._load_megatron_defaults")
+    def test_non_enum_and_none_pass_through(self, mock_defaults, mock_types):
+        # num_layers is not an enum arg -> no converter -> value left untouched.
+        mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto, "num_layers": 12}
+        mock_types.return_value = {"attention_backend": _fake_enum_type}
+
+        builder = MegatronArgBuilder()
+        builder.update({"attention_backend": None, "num_layers": 32})
+
+        assert builder.overrides["attention_backend"] is None
+        assert builder.overrides["num_layers"] == 32
+
+    @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types")
+    @patch("primus.backends.megatron.argument_builder._load_megatron_defaults")
+    def test_invalid_value_keeps_raw_without_raising(self, mock_defaults, mock_types):
+        mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto}
+        mock_types.return_value = {"attention_backend": _fake_enum_type}
+
+        builder = MegatronArgBuilder()
+        # "bogus" is not a valid enum member -> conversion fails, raw value kept.
+        builder.update({"attention_backend": "bogus"})
+
+        assert builder.overrides["attention_backend"] == "bogus"
+
+
+class TestMegatronArgBuilderRealParser:
+    """Integration checks against the real Megatron argparse (no mocking).
+
+    The coercion tests above mock the enum-type map; these guard the actual
+    detection wiring so a change in how Megatron declares attention_backend
+    (its argparse type/choices) can't silently turn coercion back into a no-op.
+    """
+
+    def test_attention_backend_detected_as_enum_arg(self):
+        from primus.backends.megatron.argument_builder import _load_megatron_enum_types
+
+        pytest.importorskip("megatron")
+        assert "attention_backend" in _load_megatron_enum_types()
+
+    @pytest.mark.parametrize("value", ["fused", "unfused", "auto"])
+    def test_attention_backend_string_coerces_to_real_enum(self, value):
+        pytest.importorskip("megatron")
+        from megatron.core.transformer.enums import AttnBackend
+
+        builder = MegatronArgBuilder()
+        builder.update({"attention_backend": value})
+
+        assert builder.overrides["attention_backend"] is AttnBackend[value]
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py
new file mode 100644
index 000000000..5dfeb5c48
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py
@@ -0,0 +1,234 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for MegatronBaseTrainer.
+
+Tests path resolution, parse_args patching, and setup orchestration.
+"""
+
+import os
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer
+
+
+def _build_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None):
+    """Helper to build MegatronBaseTrainer with stubbed dependencies."""
+
+    # Create a concrete subclass for testing (MegatronBaseTrainer is abstract)
+    class ConcreteMegatronBaseTrainer(MegatronBaseTrainer):
+        def train(self):
+            pass
+
+    # Stub out BaseTrainer.__init__ to avoid real distributed env setup
+    def dummy_base_init(self, backend_args=None, *args, **kwargs):
+        self.backend_args = backend_args
+        self.rank = 0
+        self.world_size = 1
+        self.local_rank = 0
+        self.master_addr = "localhost"
+        self.master_port = 12345
+
+    monkeypatch.setattr(
+        "primus.core.trainer.base_trainer.BaseTrainer.__init__",
+        dummy_base_init,
+    )
+
+    # Silence logging
+    monkeypatch.setattr(
+        "primus.backends.megatron.megatron_base_trainer.log_rank_0",
+        lambda *args, **kwargs: None,
+    )
+
+    if backend_args is None:
+        backend_args = SimpleNamespace()
+
+    return ConcreteMegatronBaseTrainer(backend_args=backend_args)
+
+
+class TestMegatronBaseTrainer:
+    """Tests for MegatronBaseTrainer setup and path resolution."""
+
+    def test_ensure_megatron_path_from_primus_path_env(self, monkeypatch: pytest.MonkeyPatch):
+        """Test path resolution from PRIMUS_PATH environment variable."""
+        trainer = _build_trainer(monkeypatch)
+
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            primus_path = Path(tmpdir) / "Primus"
+            megatron_path = primus_path / "third_party" / "Megatron-LM"
+            megatron_path.mkdir(parents=True)
+
+            monkeypatch.setenv("PRIMUS_PATH", str(primus_path))
+
+            # Force `import megatron` to raise ImportError by caching None in
+            # sys.modules.  Unlike mocking builtins.__import__, this leaves
+            # Python's import machinery intact for all other modules.
+            saved_megatron = sys.modules.get("megatron", _SENTINEL := object())
+            sys.modules["megatron"] = None
+            # Strip paths that either contain "megatron" in the name or have a
+            # megatron/ subdirectory, so Method 4 doesn't short-circuit.
+            original_path = sys.path[:]
+            sys.path[:] = [
+                p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir()
+            ]
+            try:
+                trainer._ensure_megatron_path()
+                assert str(megatron_path) in sys.path
+            finally:
+                sys.path[:] = original_path
+                if saved_megatron is _SENTINEL:
+                    sys.modules.pop("megatron", None)
+                else:
+                    sys.modules["megatron"] = saved_megatron
+
+    def test_ensure_megatron_path_from_cwd(self, monkeypatch: pytest.MonkeyPatch):
+        """Test path resolution from current working directory."""
+        trainer = _build_trainer(monkeypatch)
+
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            primus_path = Path(tmpdir) / "Primus"
+            megatron_path = primus_path / "third_party" / "Megatron-LM"
+            megatron_path.mkdir(parents=True)
+
+            monkeypatch.delenv("PRIMUS_PATH", raising=False)
+
+            def mock_cwd():
+                return primus_path / "some" / "subdirectory"
+
+            monkeypatch.setattr(Path, "cwd", staticmethod(mock_cwd))
+
+            saved_megatron = sys.modules.get("megatron", _SENTINEL := object())
+            sys.modules["megatron"] = None
+            original_path = sys.path[:]
+            sys.path[:] = [
+                p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir()
+            ]
+            try:
+                trainer._ensure_megatron_path()
+                assert str(megatron_path) in sys.path
+            finally:
+                sys.path[:] = original_path
+                if saved_megatron is _SENTINEL:
+                    sys.modules.pop("megatron", None)
+                else:
+                    sys.modules["megatron"] = saved_megatron
+
+    def test_ensure_megatron_path_from_file_location(self, monkeypatch: pytest.MonkeyPatch):
+        """Test path resolution from current file location."""
+        trainer = _build_trainer(monkeypatch)
+
+        # Unset PRIMUS_PATH and mock cwd to not contain Primus
+        monkeypatch.delenv("PRIMUS_PATH", raising=False)
+        monkeypatch.setattr(Path, "cwd", staticmethod(lambda: Path("/some/other/path")))
+
+        # Mock __file__ to point to a Primus subdirectory
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            primus_path = Path(tmpdir) / "Primus"
+            megatron_path = primus_path / "third_party" / "Megatron-LM"
+            megatron_path.mkdir(parents=True)
+
+            # Mock __file__ to be in primus/backends/megatron/
+            fake_file = primus_path / "primus" / "backends" / "megatron" / "megatron_base_trainer.py"
+            fake_file.parent.mkdir(parents=True)
+            fake_file.touch()
+
+            def mock_resolve(self):
+                return fake_file
+
+            monkeypatch.setattr(Path, "resolve", mock_resolve)
+
+            saved_megatron = sys.modules.get("megatron", _SENTINEL := object())
+            sys.modules["megatron"] = None
+            original_path = sys.path[:]
+            # Isolate Method 3 (file-location) by neutralizing Method 4, which
+            # returns early if any sys.path entry already contains a real
+            # ``megatron`` package dir. A name-only filter misses a pip-installed
+            # megatron-core under site-packages (no "megatron" in the path name),
+            # so also drop entries that actually contain a ``megatron`` dir.
+            sys.path[:] = [
+                p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir()
+            ]
+            try:
+                trainer._ensure_megatron_path()
+                assert str(megatron_path) in sys.path
+            finally:
+                sys.path[:] = original_path
+                if saved_megatron is _SENTINEL:
+                    sys.modules.pop("megatron", None)
+                else:
+                    sys.modules["megatron"] = saved_megatron
+
+    def test_ensure_megatron_path_already_importable(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that path setup is skipped if megatron is already importable."""
+        trainer = _build_trainer(monkeypatch)
+
+        megatron_mod = types.ModuleType("megatron")
+        monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
+
+        path_before = sys.path[:]
+        trainer._ensure_megatron_path()
+
+        assert sys.path == path_before
+
+    def test_cleanup_exit_fast_flushes_before_hard_exit(self, monkeypatch: pytest.MonkeyPatch):
+        """Regression test: PRIMUS_EXIT_FAST=1's os._exit(0) bypasses atexit, so
+        flush_before_hard_exit() (stdout/stderr + coverage.save()) must run first.
+        A prior version inlined only the stdout/stderr half, silently dropping E2E
+        coverage data for that process."""
+        trainer = _build_trainer(monkeypatch)
+        monkeypatch.setattr(trainer, "_finalize_mlflow_artifacts", lambda: None)
+        monkeypatch.setenv("PRIMUS_EXIT_FAST", "1")
+
+        calls = []
+        monkeypatch.setattr(
+            "primus.backends.megatron.megatron_base_trainer.flush_before_hard_exit",
+            lambda: calls.append("flush"),
+        )
+        monkeypatch.setattr(os, "_exit", lambda code: calls.append(("exit", code)))
+
+        trainer.cleanup(on_error=False)
+
+        assert calls == ["flush", ("exit", 0)]
+
+    def test_patch_parse_args(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that parse_args is patched in both locations."""
+        trainer = _build_trainer(monkeypatch)
+
+        import megatron.training.arguments as megatron_args_mod
+        import megatron.training.initialize as megatron_init_mod
+
+        original_parse_args_args = megatron_args_mod.parse_args
+        original_parse_args_init = megatron_init_mod.parse_args
+
+        backend_args = SimpleNamespace(test_param="value")
+        trainer.backend_args = backend_args
+
+        try:
+            trainer._patch_parse_args()
+
+            assert megatron_args_mod.parse_args is not original_parse_args_args
+            assert megatron_init_mod.parse_args is not original_parse_args_init
+
+            result_args = megatron_args_mod.parse_args()
+            result_init = megatron_init_mod.parse_args()
+
+            assert result_args is backend_args
+            assert result_init is backend_args
+        finally:
+            megatron_args_mod.parse_args = original_parse_args_args
+            megatron_init_mod.parse_args = original_parse_args_init
diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py
new file mode 100644
index 000000000..077db6d25
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py
@@ -0,0 +1,795 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for MLPerf logging and warmup patches.
+
+Covers the production patch behavior (CPU-only): logging-patch monkey-patching
+and idempotency, the INIT_STOP/RUN_START emission on the first post-warmup
+training_log call, convergence detection, and the FP8/optimizer warmup helper
+functions (_reset_fp8_te_spec, seed FP8 amax, optimizer neuter/restore/reset).
+"""
+
+import types
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _install_fake_megatron(monkeypatch):
+    """Install a fake megatron.training.training module into sys.modules."""
+    import sys
+
+    megatron_mod = types.ModuleType("megatron")
+    training_pkg = types.ModuleType("megatron.training")
+    training_mod = types.ModuleType("megatron.training.training")
+    global_vars_mod = types.ModuleType("megatron.training.global_vars")
+
+    def fake_train_step(fwd, data_iter, model, optimizer, sched, config, fwdbwd, iteration=None):
+        return {}, 0, False, False, 0, 0.0, 0, None
+
+    def fake_training_log(*args, **kwargs):
+        return None
+
+    def fake_eval(*args, **kwargs):
+        return None
+
+    def fake_evaluate(*args, **kwargs):
+        return ({},)
+
+    def fake_print_rank_last(msg):
+        pass
+
+    def fake_get_tb_writer():
+        return None
+
+    def fake_get_wandb_writer():
+        return None
+
+    training_mod.train_step = fake_train_step
+    training_mod.training_log = fake_training_log
+    training_mod.evaluate_and_print_results = fake_eval
+    training_mod.evaluate = fake_evaluate
+    training_mod.print_rank_last = fake_print_rank_last
+    training_mod.get_tensorboard_writer = fake_get_tb_writer
+    training_mod.get_wandb_writer = fake_get_wandb_writer
+
+    training_pkg.training = training_mod
+    megatron_mod.training = training_pkg
+
+    # Store megatron args for get_args
+    _megatron_args = SimpleNamespace(
+        iteration=0,
+        curr_iteration=0,
+        consumed_train_samples=0,
+        skipped_train_samples=0,
+        train_iters=100,
+        eval_interval=10,
+        do_valid=True,
+        global_batch_size=512,
+        micro_batch_size=64,
+    )
+    global_vars_mod.get_args = lambda: _megatron_args
+
+    training_pkg.get_args = global_vars_mod.get_args
+
+    monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
+    monkeypatch.setitem(sys.modules, "megatron.training", training_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod)
+    monkeypatch.setitem(sys.modules, "megatron.training.global_vars", global_vars_mod)
+
+    return training_mod, _megatron_args
+
+
+def _make_ctx(
+    mlperf_mode=False,
+    warmup_train_steps=0,
+    target_val_loss=0.586,
+    **kwargs,
+):
+    """Build a minimal PatchContext-like object."""
+    params = SimpleNamespace(
+        mlperf_mode=mlperf_mode,
+        warmup_train_steps=warmup_train_steps,
+        target_val_loss=target_val_loss,
+        global_batch_size=kwargs.get("global_batch_size", 512),
+        micro_batch_size=kwargs.get("micro_batch_size", 64),
+        seed=kwargs.get("seed", 42),
+        log_interval=kwargs.get("log_interval", 10),
+        lr=kwargs.get("lr", 2e-4),
+        adam_beta1=kwargs.get("adam_beta1", 0.9),
+        adam_beta2=kwargs.get("adam_beta2", 0.95),
+        adam_eps=kwargs.get("adam_eps", 1e-8),
+        weight_decay=kwargs.get("weight_decay", 0.1),
+        image_size=kwargs.get("image_size", 256),
+        vae_latent_mode=kwargs.get("vae_latent_mode", "resample"),
+        transformer_impl=kwargs.get("transformer_impl", "local"),
+        use_fsdp2_fp8_all_gather=kwargs.get("use_fsdp2_fp8_all_gather", False),
+        wall_clock_step_timer=False,
+        **{
+            k: v
+            for k, v in kwargs.items()
+            if k
+            not in (
+                "global_batch_size",
+                "micro_batch_size",
+                "seed",
+                "log_interval",
+                "lr",
+                "adam_beta1",
+                "adam_beta2",
+                "adam_eps",
+                "weight_decay",
+                "image_size",
+                "vae_latent_mode",
+                "transformer_impl",
+                "use_fsdp2_fp8_all_gather",
+            )
+        },
+    )
+    module_config = SimpleNamespace(params=params)
+    return SimpleNamespace(
+        extra={"module_config": module_config},
+        backend="megatron",
+        phase="before_train",
+    )
+
+
+# ============================================================================
+# Level 1: Patch registration and conditions
+# ============================================================================
+
+
+class TestLoggingPatchMonkeyPatching:
+    """Verify that the logging patch replaces the expected functions."""
+
+    def test_installs_wrappers(self, monkeypatch):
+        mt, _ = _install_fake_megatron(monkeypatch)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0",
+            lambda *a, **k: None,
+        )
+
+        mock_mllog = MagicMock()
+        mock_mllog.get_mllogger.return_value = MagicMock()
+        mock_mllog.constants = SimpleNamespace(
+            INIT_START="init_start",
+            INIT_STOP="init_stop",
+            RUN_START="run_start",
+            RUN_STOP="run_stop",
+            SUBMISSION_BENCHMARK="submission_benchmark",
+            SUBMISSION_ORG="submission_org",
+            SUBMISSION_DIVISION="submission_division",
+            SUBMISSION_PLATFORM="submission_platform",
+            SUBMISSION_STATUS="submission_status",
+            SEED="seed",
+            GLOBAL_BATCH_SIZE="global_batch_size",
+            TRAIN_SAMPLES="train_samples",
+            EVAL_SAMPLES="eval_samples",
+            GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps",
+            OPT_NAME="opt_name",
+            OPT_BASE_LR="opt_base_lr",
+            EVAL_ACCURACY="eval_accuracy",
+            EVAL_START="eval_start",
+            EVAL_STOP="eval_stop",
+            EPOCH_START="epoch_start",
+            BLOCK_START="block_start",
+            BLOCK_STOP="block_stop",
+        )
+        mock_mlperf_pkg = MagicMock()
+        mock_mlperf_pkg.mllog = mock_mllog
+        monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg)
+        monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog)
+
+        from primus.backends.megatron.patches.mlperf_logging_patches import (
+            patch_mlperf_logging,
+        )
+
+        original_tl = mt.training_log
+        original_eval = mt.evaluate_and_print_results
+        original_prl = mt.print_rank_last
+
+        ctx = _make_ctx(mlperf_mode=True)
+        patch_mlperf_logging(ctx)
+
+        assert mt.training_log is not original_tl
+        assert mt.evaluate_and_print_results is not original_eval
+        assert mt.print_rank_last is not original_prl
+        assert getattr(mt, "_primus_mlperf_logging_installed", False) is True
+
+    def test_idempotent(self, monkeypatch):
+        mt, _ = _install_fake_megatron(monkeypatch)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0",
+            lambda *a, **k: None,
+        )
+
+        mock_mllog = MagicMock()
+        mock_mllog.get_mllogger.return_value = MagicMock()
+        mock_mllog.constants = SimpleNamespace(
+            INIT_START="init_start",
+            INIT_STOP="init_stop",
+            RUN_START="run_start",
+            RUN_STOP="run_stop",
+            SUBMISSION_BENCHMARK="submission_benchmark",
+            SUBMISSION_ORG="submission_org",
+            SUBMISSION_DIVISION="submission_division",
+            SUBMISSION_PLATFORM="submission_platform",
+            SUBMISSION_STATUS="submission_status",
+            SEED="seed",
+            GLOBAL_BATCH_SIZE="global_batch_size",
+            TRAIN_SAMPLES="train_samples",
+            EVAL_SAMPLES="eval_samples",
+            GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps",
+            OPT_NAME="opt_name",
+            OPT_BASE_LR="opt_base_lr",
+            EVAL_ACCURACY="eval_accuracy",
+            EVAL_START="eval_start",
+            EVAL_STOP="eval_stop",
+            EPOCH_START="epoch_start",
+            BLOCK_START="block_start",
+            BLOCK_STOP="block_stop",
+        )
+        mock_mlperf_pkg = MagicMock()
+        mock_mlperf_pkg.mllog = mock_mllog
+        monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg)
+        monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog)
+
+        from primus.backends.megatron.patches.mlperf_logging_patches import (
+            patch_mlperf_logging,
+        )
+
+        ctx = _make_ctx(mlperf_mode=True)
+        mt._primus_mlperf_logging_installed = False
+
+        patch_mlperf_logging(ctx)
+        first_tl = mt.training_log
+        first_eval = mt.evaluate_and_print_results
+
+        patch_mlperf_logging(ctx)
+        assert mt.training_log is first_tl
+        assert mt.evaluate_and_print_results is first_eval
+
+
+class TestTrainingLogFirstCall:
+    """Verify INIT_STOP + RUN_START fire on first post-warmup training_log."""
+
+    def test_first_call_emits_init_stop_run_start(self, monkeypatch):
+        mt, args = _install_fake_megatron(monkeypatch)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0",
+            lambda *a, **k: None,
+        )
+
+        emitted_events = []
+
+        mock_mllogger = MagicMock()
+        mock_constants = SimpleNamespace(
+            INIT_START="init_start",
+            INIT_STOP="init_stop",
+            RUN_START="run_start",
+            RUN_STOP="run_stop",
+            SUBMISSION_BENCHMARK="submission_benchmark",
+            SUBMISSION_ORG="submission_org",
+            SUBMISSION_DIVISION="submission_division",
+            SUBMISSION_PLATFORM="submission_platform",
+            SUBMISSION_STATUS="submission_status",
+            SEED="seed",
+            GLOBAL_BATCH_SIZE="global_batch_size",
+            TRAIN_SAMPLES="train_samples",
+            EVAL_SAMPLES="eval_samples",
+            GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps",
+            OPT_NAME="opt_name",
+            OPT_BASE_LR="opt_base_lr",
+            EVAL_ACCURACY="eval_accuracy",
+            EVAL_START="eval_start",
+            EVAL_STOP="eval_stop",
+            EPOCH_START="epoch_start",
+            BLOCK_START="block_start",
+            BLOCK_STOP="block_stop",
+        )
+
+        def track_start(key, value=None, metadata=None):
+            emitted_events.append(("start", key))
+
+        def track_end(key, value=None, metadata=None):
+            emitted_events.append(("end", key))
+
+        def track_event(key, value=None, metadata=None):
+            emitted_events.append(("event", key))
+
+        mock_mllogger.start = track_start
+        mock_mllogger.end = track_end
+        mock_mllogger.event = track_event
+
+        mock_mllog_module = MagicMock()
+        mock_mllog_module.get_mllogger.return_value = mock_mllogger
+        mock_mllog_module.constants = mock_constants
+
+        # Wire the top-level mlperf_logging mock so that
+        # `from mlperf_logging import mllog` resolves correctly
+        mock_mlperf_pkg = MagicMock()
+        mock_mlperf_pkg.mllog = mock_mllog_module
+
+        import sys as _sys
+
+        monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg)
+        monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module)
+        monkeypatch.setenv("RANK", "0")
+
+        from primus.backends.megatron.patches.mlperf_logging_patches import (
+            patch_mlperf_logging,
+        )
+
+        ctx = _make_ctx(mlperf_mode=True, log_interval=1)
+        patch_mlperf_logging(ctx)
+
+        emitted_events.clear()
+
+        # First call should emit INIT_STOP + RUN_START
+        mt.training_log({"loss": 0.5}, {}, 1e-4, 1, 1.0, False, False, 0.0, None, 0, None)
+
+        event_keys = [e[1] for e in emitted_events]
+        assert "init_stop" in event_keys, f"INIT_STOP not emitted. Events: {emitted_events}"
+        assert "run_start" in event_keys, f"RUN_START not emitted. Events: {emitted_events}"
+
+        init_stop_idx = event_keys.index("init_stop")
+        run_start_idx = event_keys.index("run_start")
+        assert init_stop_idx < run_start_idx
+
+        # Second call should NOT emit INIT_STOP/RUN_START again
+        emitted_events.clear()
+        mt.training_log({"loss": 0.4}, {}, 1e-4, 2, 1.0, False, False, 0.0, None, 0, None)
+
+        event_keys_2 = [e[1] for e in emitted_events]
+        assert "init_stop" not in event_keys_2
+        assert "run_start" not in event_keys_2
+
+
+# ============================================================================
+# Level 3: Component tests
+# ============================================================================
+
+
+class TestConvergenceDetection:
+    """Verify convergence detection in evaluate_and_print_results wrapper."""
+
+    def test_convergence_sets_train_iters(self, monkeypatch):
+        mt, megatron_args = _install_fake_megatron(monkeypatch)
+
+        monkeypatch.setattr(
+            "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0",
+            lambda *a, **k: None,
+        )
+
+        # Mock mlperf_logging — wire .mllog attribute on the top-level package
+        mock_mllogger = MagicMock()
+        mock_constants = SimpleNamespace(
+            INIT_START="init_start",
+            INIT_STOP="init_stop",
+            RUN_START="run_start",
+            RUN_STOP="run_stop",
+            SUBMISSION_BENCHMARK="submission_benchmark",
+            SUBMISSION_ORG="submission_org",
+            SUBMISSION_DIVISION="submission_division",
+            SUBMISSION_PLATFORM="submission_platform",
+            SUBMISSION_STATUS="submission_status",
+            SEED="seed",
+            GLOBAL_BATCH_SIZE="global_batch_size",
+            TRAIN_SAMPLES="train_samples",
+            EVAL_SAMPLES="eval_samples",
+            GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps",
+            OPT_NAME="opt_name",
+            OPT_BASE_LR="opt_base_lr",
+            EVAL_ACCURACY="eval_accuracy",
+            EVAL_START="eval_start",
+            EVAL_STOP="eval_stop",
+            EPOCH_START="epoch_start",
+            BLOCK_START="block_start",
+            BLOCK_STOP="block_stop",
+        )
+        mock_mllogger.start = MagicMock()
+        mock_mllogger.end = MagicMock()
+        mock_mllogger.event = MagicMock()
+
+        mock_mllog_module = MagicMock()
+        mock_mllog_module.get_mllogger.return_value = mock_mllogger
+        mock_mllog_module.constants = mock_constants
+
+        mock_mlperf_pkg = MagicMock()
+        mock_mlperf_pkg.mllog = mock_mllog_module
+
+        import sys as _sys
+
+        monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg)
+        monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module)
+        monkeypatch.setenv("RANK", "0")
+
+        # Set evaluate to return a loss below target BEFORE patching.
+        # Also make evaluate_and_print_results call evaluate() internally,
+        # mirroring real Megatron behavior so _captured_loss gets populated.
+        mt.evaluate = lambda *a, **k: ({"loss": 0.500},)
+
+        def fake_eval_and_print(*a, **k):
+            mt.evaluate(*a, **k)
+
+        mt.evaluate_and_print_results = fake_eval_and_print
+
+        from primus.backends.megatron.patches.mlperf_logging_patches import (
+            patch_mlperf_logging,
+        )
+
+        target_val_loss = 0.586
+        megatron_args.train_iters = 5000
+
+        ctx = _make_ctx(mlperf_mode=True, target_val_loss=target_val_loss)
+        patch_mlperf_logging(ctx)
+
+        # Call eval at iteration 512
+        mt.evaluate_and_print_results(
+            "iteration 512",
+            lambda: None,
+            None,
+            [MagicMock()],
+            512,
+            None,
+            MagicMock(),
+        )
+
+        assert megatron_args.train_iters == 512
+
+
+# ============================================================================
+# Level 4: Helper function unit tests (CPU-only, no GPU required)
+# ============================================================================
+
+
+class TestResetFp8TeSpec:
+    """Verify _reset_fp8_te_spec clears fp8_initialized and meta tensors."""
+
+    def test_resets_fp8_initialized_flag(self):
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_te_spec,
+        )
+
+        class FakeTeModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fp8_initialized = True
+                self.fp8_meta = {
+                    "scaling_fwd": SimpleNamespace(
+                        amax_history=torch.ones(16),
+                        scale=torch.full((4,), 3.14),
+                        scale_inv=torch.full((4,), 0.318),
+                    ),
+                    "scaling_bwd": SimpleNamespace(
+                        amax_history=torch.ones(16),
+                        scale=torch.full((4,), 2.71),
+                        scale_inv=torch.full((4,), 0.369),
+                    ),
+                }
+
+        model = torch.nn.Sequential(FakeTeModule(), FakeTeModule())
+        count = _reset_fp8_te_spec([model])
+
+        assert count == 2
+        for module in model.modules():
+            if hasattr(module, "fp8_initialized"):
+                assert module.fp8_initialized is False
+                for key in ("scaling_fwd", "scaling_bwd"):
+                    tm = module.fp8_meta[key]
+                    assert (tm.amax_history == 0.0).all()
+                    assert (tm.scale == 1.0).all()
+                    assert (tm.scale_inv == 1.0).all()
+
+    def test_skips_reset_fp8_meta_tensors_shortcut(self):
+        """
+        TE 2.8.0.dev0's `reset_fp8_meta_tensors` unconditionally derefs `.scale`
+        on the recipe state, which crashes on `Float8CurrentScalingRecipeState`
+        (current/tensorwise scaling has no persistent state). The reset must
+        therefore go through the recipe-agnostic manual path even when a TE
+        module advertises the helper.
+        """
+        from types import SimpleNamespace
+
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_te_spec,
+        )
+
+        reset_called = [False]
+
+        class FakeTeCurrentScalingModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fp8_initialized = True
+                # Mimic Float8CurrentScalingRecipeState: no .scale, no .amax_history.
+                self.fp8_meta = {
+                    "scaling_fwd": SimpleNamespace(),
+                    "scaling_bwd": SimpleNamespace(),
+                }
+
+            def reset_fp8_meta_tensors(self):
+                reset_called[0] = True
+
+        model = torch.nn.Sequential(FakeTeCurrentScalingModule())
+        count = _reset_fp8_te_spec([model])
+
+        assert not reset_called[
+            0
+        ], "reset_fp8_meta_tensors must NOT be called (would crash on current scaling)"
+        assert count == 1
+        assert model[0].fp8_initialized is False
+
+    def test_falls_through_for_delayed_scaling_buffers(self):
+        """
+        Companion to ``test_skips_reset_fp8_meta_tensors_shortcut``: confirms
+        the manual fallback still resets delayed-scaling buffers (.scale,
+        .amax_history, .scale_inv) when they exist on the recipe state.
+        """
+        from types import SimpleNamespace
+
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_fp8_te_spec,
+        )
+
+        reset_called = [False]
+
+        class FakeTeDelayedScalingModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fp8_initialized = True
+                self.fp8_meta = {
+                    "scaling_fwd": SimpleNamespace(
+                        amax_history=torch.full((16,), 7.0),
+                        scale=torch.full((4,), 2.71),
+                        scale_inv=torch.full((4,), 0.369),
+                    ),
+                    "scaling_bwd": SimpleNamespace(
+                        amax_history=torch.full((16,), 7.0),
+                        scale=torch.full((4,), 2.71),
+                        scale_inv=torch.full((4,), 0.369),
+                    ),
+                }
+
+            def reset_fp8_meta_tensors(self):
+                reset_called[0] = True
+
+        model = torch.nn.Sequential(FakeTeDelayedScalingModule())
+        count = _reset_fp8_te_spec([model])
+
+        assert not reset_called[0]
+        assert count == 1
+        for key in ("scaling_fwd", "scaling_bwd"):
+            tm = model[0].fp8_meta[key]
+            assert (tm.amax_history == 0.0).all()
+            assert (tm.scale == 1.0).all()
+            assert (tm.scale_inv == 1.0).all()
+
+
+class TestSeedFp8Amax:
+    """Verify _seed_fp8_amax fills amax_history with the requested seed value."""
+
+    def _make_te_model(self):
+        import torch
+
+        class FakeTeModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.fp8_meta = {
+                    "scaling_fwd": SimpleNamespace(
+                        amax_history=torch.zeros(16),
+                    ),
+                    "scaling_bwd": SimpleNamespace(
+                        amax_history=torch.zeros(16),
+                    ),
+                }
+
+        return torch.nn.Sequential(FakeTeModule(), FakeTeModule())
+
+    def test_seeds_default_value(self):
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _seed_fp8_amax,
+        )
+
+        model = self._make_te_model()
+        count = _seed_fp8_amax([model])
+
+        assert count == 4
+        for module in model.modules():
+            if hasattr(module, "fp8_meta"):
+                for key in ("scaling_fwd", "scaling_bwd"):
+                    assert (module.fp8_meta[key].amax_history == 1.0).all()
+
+    def test_seeds_custom_value(self):
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _seed_fp8_amax,
+        )
+
+        model = self._make_te_model()
+        count = _seed_fp8_amax([model], seed_value=42.0)
+
+        assert count == 4
+        for module in model.modules():
+            if hasattr(module, "fp8_meta"):
+                for key in ("scaling_fwd", "scaling_bwd"):
+                    assert (module.fp8_meta[key].amax_history == 42.0).all()
+
+
+class TestNeuterRestoreOptimizer:
+    """Verify _neuter_optimizer / _restore_optimizer roundtrip."""
+
+    def test_roundtrip_preserves_hyperparams(self):
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _neuter_optimizer,
+            _restore_optimizer,
+        )
+
+        model = torch.nn.Linear(4, 4)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.95), weight_decay=0.1)
+
+        orig_betas = list(opt.param_groups[0]["betas"])
+        orig_wd = opt.param_groups[0]["weight_decay"]
+
+        wrapper = SimpleNamespace(optimizer=opt)
+        saved = _neuter_optimizer(wrapper)
+
+        assert opt.param_groups[0]["betas"] == [1.0, 1.0]
+        assert opt.param_groups[0]["weight_decay"] == 0.0
+
+        _restore_optimizer(wrapper, saved)
+
+        assert list(opt.param_groups[0]["betas"]) == orig_betas
+        assert opt.param_groups[0]["weight_decay"] == orig_wd
+
+    def test_roundtrip_all_keys(self):
+        """All 4 production keys roundtrip: betas, weight_decay, bias_correction, pre_mult_wd."""
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _neuter_optimizer,
+            _restore_optimizer,
+        )
+
+        model = torch.nn.Linear(4, 4)
+        opt = torch.optim.SGD(model.parameters(), lr=0.01)
+        opt.param_groups[0]["betas"] = [0.9, 0.999]
+        opt.param_groups[0]["weight_decay"] = 0.01
+        opt.param_groups[0]["bias_correction"] = True
+        opt.param_groups[0]["pre_mult_wd"] = 0.05
+
+        wrapper = SimpleNamespace(optimizer=opt)
+        saved = _neuter_optimizer(wrapper)
+
+        assert opt.param_groups[0]["betas"] == [1.0, 1.0]
+        assert opt.param_groups[0]["weight_decay"] == 0.0
+        assert opt.param_groups[0]["bias_correction"] is False
+        assert opt.param_groups[0]["pre_mult_wd"] == 0.0
+
+        _restore_optimizer(wrapper, saved)
+
+        assert opt.param_groups[0]["betas"] == [0.9, 0.999]
+        assert opt.param_groups[0]["weight_decay"] == 0.01
+        assert opt.param_groups[0]["bias_correction"] is True
+        assert opt.param_groups[0]["pre_mult_wd"] == 0.05
+
+    def test_multi_param_group_roundtrip(self):
+        """Neuter/restore with 2 param groups preserves per-group values."""
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _neuter_optimizer,
+            _restore_optimizer,
+        )
+
+        model = torch.nn.Linear(4, 4)
+        w_params = [model.weight]
+        b_params = [model.bias]
+        opt = torch.optim.Adam(
+            [
+                {"params": w_params, "weight_decay": 0.1},
+                {"params": b_params, "weight_decay": 0.0},
+            ],
+            lr=1e-3,
+            betas=(0.9, 0.999),
+        )
+
+        wrapper = SimpleNamespace(optimizer=opt)
+        saved = _neuter_optimizer(wrapper)
+
+        for g in opt.param_groups:
+            assert g["betas"] == [1.0, 1.0]
+            assert g["weight_decay"] == 0.0
+
+        _restore_optimizer(wrapper, saved)
+
+        assert opt.param_groups[0]["weight_decay"] == 0.1
+        assert opt.param_groups[1]["weight_decay"] == 0.0
+        for g in opt.param_groups:
+            assert list(g["betas"]) == [0.9, 0.999]
+
+
+class TestResetOptimizerState:
+    """Verify _reset_optimizer_state clears step counters."""
+
+    def test_resets_param_group_step(self):
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_optimizer_state,
+        )
+
+        model = torch.nn.Linear(4, 4)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+        loss = model(torch.randn(2, 4)).sum()
+        loss.backward()
+        opt.step()
+
+        opt.param_groups[0]["step"] = 42
+        wrapper = SimpleNamespace(optimizer=opt)
+
+        _reset_optimizer_state(wrapper)
+
+        assert opt.param_groups[0]["step"] == 0
+
+    def test_resets_per_param_state_step_tensor(self):
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_optimizer_state,
+        )
+
+        model = torch.nn.Linear(4, 4)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+        loss = model(torch.randn(2, 4)).sum()
+        loss.backward()
+        opt.step()
+
+        for state in opt.state.values():
+            assert state["step"].item() > 0
+
+        wrapper = SimpleNamespace(optimizer=opt)
+        _reset_optimizer_state(wrapper)
+
+        for state in opt.state.values():
+            assert state["step"].item() == 0
+
+    def test_handles_chained_optimizer(self):
+        import torch
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _reset_optimizer_state,
+        )
+
+        m1 = torch.nn.Linear(4, 4)
+        m2 = torch.nn.Linear(4, 4)
+        opt1 = torch.optim.SGD(m1.parameters(), lr=0.01)
+        opt2 = torch.optim.SGD(m2.parameters(), lr=0.01)
+
+        opt1.param_groups[0]["step"] = 10
+        opt2.param_groups[0]["step"] = 20
+
+        w1 = SimpleNamespace(optimizer=opt1)
+        w2 = SimpleNamespace(optimizer=opt2)
+        chained = SimpleNamespace(chained_optimizers=[w1, w2])
+
+        _reset_optimizer_state(chained)
+
+        assert opt1.param_groups[0]["step"] == 0
+        assert opt2.param_groups[0]["step"] == 0
diff --git a/tests/unit_tests/backends/megatron/test_mlperf_validation.py b/tests/unit_tests/backends/megatron/test_mlperf_validation.py
new file mode 100644
index 000000000..3330a4bc3
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_mlperf_validation.py
@@ -0,0 +1,164 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for MLPerf validation support.
+
+Tests cover:
+    1. Validation loss function shape and target computation
+    2. Training loss function shape (unchanged)
+    3. Mock dataset timestep generation (production MockDiffusionDataset)
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# 1. Validation loss function shape
+# ---------------------------------------------------------------------------
+
+
+class TestValLossFunc:
+    """Test the validation loss function produces correct shapes and values."""
+
+    def _make_val_loss_func(self, noise, clean_latents):
+        """Build val_loss_func mirroring DiffusionPretrainTrainer.forward_step."""
+
+        class Holder:
+            _last_noise = noise
+            _last_clean_latents = clean_latents
+
+        def val_loss_func(output_tensor, non_loss_data=False):
+            if non_loss_data:
+                return output_tensor
+            target = Holder._last_noise - Holder._last_clean_latents
+            loss = F.mse_loss(output_tensor.float(), target.float(), reduction="none")
+            loss_per_sample = loss.mean(dim=tuple(range(1, loss.ndim)))
+            loss_sum = loss_per_sample.sum()
+            sample_count = torch.tensor(loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device)
+            return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())}
+
+        return val_loss_func
+
+    def test_returns_tuple_with_loss_key(self):
+        B, C, H, W = 4, 16, 32, 32
+        noise = torch.randn(B, C, H, W)
+        clean = torch.randn(B, C, H, W)
+        pred = torch.randn(B, C, H, W)
+        fn = self._make_val_loss_func(noise, clean)
+
+        loss_scalar, metrics = fn(pred)
+
+        assert isinstance(loss_scalar, torch.Tensor)
+        assert loss_scalar.ndim == 0
+        assert "loss" in metrics
+        loss_sum, sample_count = metrics["loss"]
+        assert loss_sum.ndim == 0
+        assert sample_count.item() == B
+
+    def test_correct_target(self):
+        B, C, H, W = 2, 4, 8, 8
+        noise = torch.randn(B, C, H, W)
+        clean = torch.randn(B, C, H, W)
+        target = noise - clean
+        fn = self._make_val_loss_func(noise, clean)
+
+        loss_scalar, metrics = fn(target)
+
+        assert loss_scalar.item() == pytest.approx(0.0, abs=1e-6)
+
+    def test_non_loss_data_passthrough(self):
+        fn = self._make_val_loss_func(torch.zeros(1), torch.zeros(1))
+        sentinel = torch.tensor(42.0)
+        result = fn(sentinel, non_loss_data=True)
+        assert result is sentinel
+
+
+# ---------------------------------------------------------------------------
+# 2. Training loss function shape
+# ---------------------------------------------------------------------------
+
+
+class TestTrainLossFunc:
+    """Verify the training loss path returns scalar with reduced_train_loss."""
+
+    def _make_train_loss_func(self, loss_fn, noise, clean, mask):
+        class Holder:
+            _last_noise = noise
+            _last_clean_latents = clean
+            _last_loss_mask = mask
+
+        def diffusion_loss_func(output_tensor, non_loss_data=False):
+            if non_loss_data:
+                return output_tensor
+            loss = loss_fn(
+                output_tensor, Holder._last_clean_latents, Holder._last_noise, Holder._last_loss_mask
+            )
+            return loss, {"reduced_train_loss": loss.detach().clone()}
+
+        return diffusion_loss_func
+
+    def test_returns_scalar_with_key(self):
+        from primus.backends.megatron.training.diffusion.loss_computation import (
+            compute_flow_matching_loss,
+        )
+
+        B, C, H, W = 4, 16, 32, 32
+        noise = torch.randn(B, C, H, W)
+        clean = torch.randn(B, C, H, W)
+        pred = torch.randn(B, C, H, W)
+
+        fn = self._make_train_loss_func(compute_flow_matching_loss, noise, clean, None)
+        loss, metrics = fn(pred)
+
+        assert loss.ndim == 0
+        assert "reduced_train_loss" in metrics
+        assert metrics["reduced_train_loss"].ndim == 0
+
+
+# ---------------------------------------------------------------------------
+# 3. Mock dataset timestep generation
+# ---------------------------------------------------------------------------
+
+
+class TestMockDatasetTimestep:
+    """Verify validation mock datasets produce cycling timestep 0-7."""
+
+    def test_validation_dataset_has_timestep(self):
+        from primus.backends.megatron.data.synthetic.mock_datasets import (
+            MockDiffusionDataset,
+        )
+
+        ds = MockDiffusionDataset(
+            num_samples=16,
+            image_size=256,
+            model_preset="flux_schnell",
+            dtype=torch.float32,
+            device="cpu",
+            is_validation=True,
+        )
+        for i in range(16):
+            sample = ds[i]
+            assert "timestep" in sample, f"Sample {i} missing 'timestep'"
+            assert sample["timestep"].item() == i % 8
+
+    def test_training_dataset_no_timestep(self):
+        from primus.backends.megatron.data.synthetic.mock_datasets import (
+            MockDiffusionDataset,
+        )
+
+        ds = MockDiffusionDataset(
+            num_samples=4,
+            image_size=256,
+            model_preset="flux_schnell",
+            dtype=torch.float32,
+            device="cpu",
+            is_validation=False,
+        )
+        for i in range(4):
+            sample = ds[i]
+            assert "timestep" not in sample
diff --git a/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py b/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py
new file mode 100644
index 000000000..7742e010f
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py
@@ -0,0 +1,86 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Integration-style tests verifying that optimizer state is clean after warmup
+completion.
+
+Uses mocked Megatron training harnesses - no real GPU training loop, but
+exercises the warmup helper pipeline (neuter -> restore -> reset) end-to-end.
+
+Run:
+    python -m pytest tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py -v
+"""
+
+from types import SimpleNamespace
+
+import torch
+
+
+class TestOptimizerStateCleanAfterWarmup:
+    """After neuter → steps → restore → reset, optimizer should be clean."""
+
+    def test_adam_state_clean_after_warmup_simulation(self):
+        """Simulate warmup: step with real betas, then restore + reset.
+
+        Note: production uses Apex/TE FusedAdam which supports
+        bias_correction=False.  Stock PyTorch Adam doesn't, so betas=[1,1]
+        causes division-by-zero.  We use SGD for the neutered phase, then
+        switch to Adam to verify the reset path.
+        """
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _neuter_optimizer,
+            _reset_optimizer_state,
+            _restore_optimizer,
+        )
+
+        model = torch.nn.Linear(16, 8, bias=True)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01)
+        wrapper = SimpleNamespace(optimizer=opt)
+
+        for _ in range(3):
+            loss = model(torch.randn(4, 16)).sum()
+            loss.backward()
+            opt.step()
+            opt.zero_grad()
+
+        for state in opt.state.values():
+            assert state["step"].item() > 0
+
+        saved_hyp = _neuter_optimizer(wrapper)
+        assert opt.param_groups[0]["betas"] == [1.0, 1.0]
+        assert opt.param_groups[0]["weight_decay"] == 0.0
+
+        _restore_optimizer(wrapper, saved_hyp)
+        _reset_optimizer_state(wrapper)
+
+        assert list(opt.param_groups[0]["betas"]) == [0.9, 0.999]
+        assert opt.param_groups[0]["weight_decay"] == 0.01
+
+        for state in opt.state.values():
+            step_val = state["step"]
+            if isinstance(step_val, torch.Tensor):
+                assert step_val.item() == 0
+            else:
+                assert step_val == 0
+
+    def test_param_groups_step_cleared(self):
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _neuter_optimizer,
+            _reset_optimizer_state,
+            _restore_optimizer,
+        )
+
+        model = torch.nn.Linear(8, 4)
+        opt = torch.optim.SGD(model.parameters(), lr=0.01)
+        opt.param_groups[0]["step"] = 42
+        wrapper = SimpleNamespace(optimizer=opt)
+
+        saved_hyp = _neuter_optimizer(wrapper)
+        _restore_optimizer(wrapper, saved_hyp)
+        _reset_optimizer_state(wrapper)
+
+        assert opt.param_groups[0].get("step", 0) == 0
diff --git a/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py b/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py
deleted file mode 100644
index 2d90ae52f..000000000
--- a/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py
+++ /dev/null
@@ -1,214 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-import dataclasses
-import sys
-import types
-from types import SimpleNamespace
-
-import pytest
-
-from primus.core.patches.context import PatchContext
-
-
-def _install_fake_megatron_training(monkeypatch: pytest.MonkeyPatch):
-    megatron_mod = types.ModuleType("megatron")
-    megatron_mod.__path__ = []
-
-    training_pkg = types.ModuleType("megatron.training")
-    training_pkg.__path__ = []
-
-    training_mod = types.ModuleType("megatron.training.training")
-    original_calls = []
-
-    def original_get_megatron_optimizer(
-        config,
-        model_chunks,
-        config_overrides=None,
-        use_gloo_process_groups=True,
-        pg_collection=None,
-        dump_param_to_param_group_map=None,
-    ):
-        original_calls.append(
-            {
-                "config": config,
-                "model_chunks": model_chunks,
-                "config_overrides": config_overrides,
-                "use_gloo_process_groups": use_gloo_process_groups,
-                "pg_collection": pg_collection,
-                "dump_param_to_param_group_map": dump_param_to_param_group_map,
-            }
-        )
-        return "original-result"
-
-    training_mod.get_megatron_optimizer = original_get_megatron_optimizer
-    training_pkg.training = training_mod
-    megatron_mod.training = training_pkg
-
-    monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
-    monkeypatch.setitem(sys.modules, "megatron.training", training_pkg)
-    monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod)
-
-    return training_mod, original_calls
-
-
-def _install_fake_muon_dependencies(monkeypatch: pytest.MonkeyPatch):
-    muon_calls = []
-
-    moun_mod = types.ModuleType("primus.backends.megatron.core.optimizer.moun")
-
-    def fake_get_megatron_muon_optimizer(
-        config,
-        model_chunks,
-        config_overrides=None,
-        use_gloo_process_groups=True,
-        layer_wise_distributed_optimizer=False,
-        pg_collection=None,
-        dump_param_to_param_group_map=None,
-    ):
-        muon_calls.append(
-            {
-                "config": config,
-                "model_chunks": model_chunks,
-                "config_overrides": config_overrides,
-                "use_gloo_process_groups": use_gloo_process_groups,
-                "layer_wise_distributed_optimizer": layer_wise_distributed_optimizer,
-                "pg_collection": pg_collection,
-                "dump_param_to_param_group_map": dump_param_to_param_group_map,
-            }
-        )
-        return "muon-result"
-
-    moun_mod.get_megatron_muon_optimizer = fake_get_megatron_muon_optimizer
-
-    moun_config_mod = types.ModuleType("primus.backends.megatron.core.optimizer.moun_optimizer_config")
-
-    @dataclasses.dataclass
-    class FakeMounOptimizerConfig:
-        optimizer: str = "muon"
-        muon_tp_mode: str = "blockwise"
-        timers: object = None
-
-    moun_config_mod.MounOptimizerConfig = FakeMounOptimizerConfig
-
-    monkeypatch.setitem(
-        sys.modules,
-        "primus.backends.megatron.core.optimizer.moun",
-        moun_mod,
-    )
-    monkeypatch.setitem(
-        sys.modules,
-        "primus.backends.megatron.core.optimizer.moun_optimizer_config",
-        moun_config_mod,
-    )
-
-    return muon_calls
-
-
-def _call_get_megatron_optimizer(
-    optimizer_fn,
-    config,
-    model_chunks,
-    config_overrides,
-    use_gloo_process_groups,
-    pg_collection,
-    dump_param_to_param_group_map,
-    positional_config_overrides,
-):
-    args = [config, model_chunks]
-    kwargs = {
-        "use_gloo_process_groups": use_gloo_process_groups,
-        "pg_collection": pg_collection,
-        "dump_param_to_param_group_map": dump_param_to_param_group_map,
-    }
-    if positional_config_overrides:
-        args.append(config_overrides)
-    else:
-        kwargs["config_overrides"] = config_overrides
-    return optimizer_fn(*args, **kwargs)
-
-
-@pytest.mark.parametrize(
-    "positional_config_overrides",
-    [False, True],
-    ids=["keyword_config_overrides", "positional_config_overrides"],
-)
-def test_patch_get_megatron_optimizer_muon_matches_runtime_signature(
-    monkeypatch: pytest.MonkeyPatch,
-    positional_config_overrides: bool,
-):
-    training_mod, original_calls = _install_fake_megatron_training(monkeypatch)
-    muon_calls = _install_fake_muon_dependencies(monkeypatch)
-
-    monkeypatch.setattr(
-        "primus.backends.megatron.patches.muon_optimizer_patches.log_rank_0",
-        lambda *args, **kwargs: None,
-    )
-
-    from primus.backends.megatron.patches.muon_optimizer_patches import (
-        patch_get_megatron_optimizer_muon,
-    )
-
-    original_fn = training_mod.get_megatron_optimizer
-    ctx = PatchContext(
-        backend="megatron",
-        phase="before_train",
-        extra={"backend_args": SimpleNamespace(muon_tp_mode="blockwise")},
-    )
-
-    patch_get_megatron_optimizer_muon(ctx)
-
-    assert training_mod.get_megatron_optimizer is not original_fn
-
-    adam_config = SimpleNamespace(optimizer="adam", timers="adam-timer")
-    result = _call_get_megatron_optimizer(
-        training_mod.get_megatron_optimizer,
-        adam_config,
-        ["chunk-0"],
-        config_overrides={"dense": "group"},
-        use_gloo_process_groups=False,
-        pg_collection="pg-0",
-        dump_param_to_param_group_map="dense-map",
-        positional_config_overrides=positional_config_overrides,
-    )
-
-    assert result == "original-result"
-    assert original_calls == [
-        {
-            "config": adam_config,
-            "model_chunks": ["chunk-0"],
-            "config_overrides": {"dense": "group"},
-            "use_gloo_process_groups": False,
-            "pg_collection": "pg-0",
-            "dump_param_to_param_group_map": "dense-map",
-        }
-    ]
-
-    muon_config = SimpleNamespace(optimizer="muon-dist", timers="muon-timer")
-    result = _call_get_megatron_optimizer(
-        training_mod.get_megatron_optimizer,
-        muon_config,
-        ["chunk-1"],
-        config_overrides={"sparse": "group"},
-        use_gloo_process_groups=False,
-        pg_collection="pg-1",
-        dump_param_to_param_group_map="muon-map",
-        positional_config_overrides=positional_config_overrides,
-    )
-
-    assert result == "muon-result"
-    assert len(muon_calls) == 1
-    muon_call = muon_calls[0]
-    assert {key: value for key, value in muon_call.items() if key != "config"} == {
-        "model_chunks": ["chunk-1"],
-        "config_overrides": {"sparse": "group"},
-        "use_gloo_process_groups": False,
-        "layer_wise_distributed_optimizer": True,
-        "pg_collection": "pg-1",
-        "dump_param_to_param_group_map": "muon-map",
-    }
-    assert muon_call["config"].timers == "muon-timer"
-    assert muon_call["config"].muon_tp_mode == "blockwise"
diff --git a/tests/unit_tests/backends/megatron/test_native_fp8_layout.py b/tests/unit_tests/backends/megatron/test_native_fp8_layout.py
new file mode 100644
index 000000000..655071b59
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_native_fp8_layout.py
@@ -0,0 +1,184 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""GPU-only tests for the native (NN/TN) FP8 backward layout.
+
+FP8 GEMMs require an AMD GPU (gfx950); the fused no-transpose cast+amax kernel
+is vendored in Primus (fp8_cast_kernels_triton), so every test here is skipped
+when CUDA is unavailable. Run on an AMD GPU (gfx950) with the FP8 toolchain available.
+
+Coverage:
+  - native (``fp8_force_nt_layout=False``) vs legacy forced-NT numerical
+    equivalence for both forward and backward,
+  - finiteness of the weight gradient (the wgrad-NaN regression that motivated
+    the native layout),
+  - backward gradient flow / return-count validation for both
+    ``OpaqueFP8LinearTensorwiseFunction`` (9 inputs) and the full delayed-scaling
+    layer (``DelayedFP8LinearTensorwiseFunction``, 13 inputs); a count mismatch
+    would make ``autograd`` raise, so a successful backward validates the arity.
+"""
+
+import functools
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    Float8ColumnParallelLinear,
+    OpaqueFP8LinearTensorwiseFunction,
+)
+from tests.utils import PrimusUT
+
+requires_gpu = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Native FP8 layout exercises FP8 GEMMs, which require an AMD GPU",
+)
+
+
+def _init_method():
+    return functools.partial(torch.nn.init.xavier_uniform_)
+
+
+def _make_config(force_nt, recipe="tensorwise"):
+    """TransformerConfig for an FP8 linear, with the diffusion-only
+    fp8_force_nt_layout attribute attached (the layer reads it via getattr)."""
+    config = TransformerConfig(
+        hidden_size=64,
+        num_attention_heads=8,
+        num_layers=1,
+        params_dtype=torch.bfloat16,
+        fp8="e4m3",
+        fp8_recipe=recipe,
+    )
+    config.fp8_force_nt_layout = force_nt
+    return config
+
+
+def _build_linear(force_nt, recipe="tensorwise"):
+    return Float8ColumnParallelLinear(
+        input_size=64,
+        output_size=128,
+        config=_make_config(force_nt, recipe=recipe),
+        init_method=_init_method(),
+        bias=False,
+        gather_output=False,
+        skip_bias_add=False,
+        is_expert=False,
+    )
+
+
+def _forward_only(layer, x):
+    out = layer(x)
+    return out[0] if isinstance(out, tuple) else out
+
+
+class _GpuLinearBase(PrimusUT):
+    """Shared setup: parallel state + a minimal global-args stub for the layers."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state, monkeypatch):
+        dummy_args = SimpleNamespace(
+            rank=0,
+            world_size=1,
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            offload=False,
+            offload_ops=[],
+            patch_primus_pipeline=False,
+            pp_algorithm=None,
+            patch_zero_bubble=False,
+            enable_zero_bubble=False,
+            rampup_batch_size=None,
+            global_batch_size=1,
+            micro_batch_size=1,
+            data_parallel_size=1,
+            decrease_batch_size_if_needed=False,
+        )
+        import megatron.training.global_vars as gvars
+
+        monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args)
+
+        # Some deploy images bake PRIMUS_TURBO_GEMM_BACKEND="" which makes
+        # GlobalBackendManager raise KeyError('') on the native path. Unset it
+        # here, exactly as the docs require for native-layout production runs.
+        if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "":
+            monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False)
+
+
+class TestNativeVsForcedNtLayer(_GpuLinearBase):
+    """End-to-end equivalence of the native and forced-NT layouts via the layer."""
+
+    @requires_gpu
+    def test_forward_backward_match_and_finite(self):
+        # NOTE: looped instead of @pytest.mark.parametrize because PrimusUT is a
+        # unittest.TestCase, where pytest does not inject parametrized args.
+        for recipe in ("tensorwise", "delayed"):
+            with self.subTest(recipe=recipe):
+                torch.manual_seed(0)
+
+                native = _build_linear(force_nt=False, recipe=recipe).cuda()
+                forced = _build_linear(force_nt=True, recipe=recipe).cuda()
+                # Identical weights so the only difference is the backward GEMM layout.
+                forced.load_state_dict(native.state_dict())
+
+                assert native._force_nt is False
+                assert forced._force_nt is True
+
+                x = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda")
+                x_native = x.clone().requires_grad_(True)
+                x_forced = x.clone().requires_grad_(True)
+
+                out_native = _forward_only(native, x_native)
+                out_forced = _forward_only(forced, x_forced)
+
+                assert torch.isfinite(out_native).all(), "native forward produced non-finite output"
+                torch.testing.assert_close(out_native, out_forced, atol=2e-2, rtol=2e-2)
+
+                out_native.sum().backward()
+                out_forced.sum().backward()
+
+                # wgrad NaN regression guard + native/forced equivalence.
+                assert torch.isfinite(native.weight.grad).all(), "native wgrad is non-finite (NaN regression)"
+                assert torch.isfinite(x_native.grad).all(), "native dgrad is non-finite"
+                torch.testing.assert_close(native.weight.grad, forced.weight.grad, atol=3e-2, rtol=3e-2)
+                torch.testing.assert_close(x_native.grad, x_forced.grad, atol=3e-2, rtol=3e-2)
+
+
+class TestOpaqueFunctionBackward(_GpuLinearBase):
+    """Direct autograd-Function checks for OpaqueFP8LinearTensorwiseFunction."""
+
+    @requires_gpu
+    def test_native_backward_returns_grads_for_input_and_weight(self):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            ScalingGranularity,
+            float8_e4m3,
+            float8_e5m2,
+        )
+        from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+        gran = ScalingGranularity.TENSORWISE.value
+        backend = BackendType.HIPBLASLT.value
+
+        x = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        w_fp8, w_scale = quantize_fp8(w, float8_e4m3, ScalingGranularity.TENSORWISE)
+
+        # force_nt=False exercises the native arm; backward must return exactly
+        # 9 grads (one per forward input) or autograd raises here.
+        out = OpaqueFP8LinearTensorwiseFunction.apply(
+            x, w, w_fp8, w_scale, float8_e4m3, float8_e5m2, gran, backend, False
+        )
+        output = out[0] if isinstance(out, tuple) else out
+        output.sum().backward()
+
+        assert x.grad is not None and torch.isfinite(x.grad).all()
+        assert w.grad is not None and torch.isfinite(w.grad).all()
diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py
new file mode 100644
index 000000000..0ad1c316c
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py
@@ -0,0 +1,431 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for compile-friendly FP8 linear layers (primus_turbo_float8_local).
+
+Tests _build_fp8_config mapping, Float8ColumnParallelLinear/Float8RowParallelLinear
+construction guards, init-time dispatch to the correct autograd Function,
+decomposed tensorwise quantize numerical equivalence, and torch.compile
+graph-break validation.
+"""
+
+import functools
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from megatron.core.enums import Fp8Recipe
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.transformer_config import TransformerConfig
+from primus_turbo.pytorch.core.low_precision import (
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+
+from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+    DecomposedFP8LinearTensorwiseFunction,
+    Float8ColumnParallelLinear,
+    Float8RowParallelLinear,
+    FP8LinearBlockwiseFunction,
+    OpaqueFP8LinearTensorwiseFunction,
+    _build_fp8_config,
+    _Float8LinearMixin,
+    _quantize_fp8_tensorwise,
+)
+from tests.utils import PrimusUT
+
+
+class TestFloat8MixinStructure:
+    """Structural guards for the shared _Float8LinearMixin (no construction).
+
+    These do not build modules, so they exercise the class layout rather than
+    FP8 numerics. They catch the load-bearing footgun where the mixin is not
+    listed first in the bases: in that case the Megatron base's ``_apply`` /
+    ``_forward_impl`` would win and FP8 would be silently disabled. The
+    init-time tests below would NOT catch that, because ``__init__`` calls
+    ``self._init_fp8_state()`` explicitly regardless of MRO order.
+
+    Note: this module imports primus_turbo (CUDA-at-import), so it is still
+    gated by the module-level skip_if_no_cuda() and runs in the GPU lane.
+    """
+
+    @pytest.mark.parametrize(
+        "cls, base",
+        [
+            (Float8ColumnParallelLinear, ColumnParallelLinear),
+            (Float8RowParallelLinear, RowParallelLinear),
+        ],
+    )
+    def test_mixin_precedes_base_in_mro(self, cls, base):
+        mro = cls.__mro__
+        assert _Float8LinearMixin in mro
+        assert base in mro
+        assert mro.index(_Float8LinearMixin) < mro.index(base), (
+            f"{cls.__name__}: _Float8LinearMixin must precede {base.__name__} in the MRO "
+            "or FP8 _apply/_forward_impl would be silently overridden by the base."
+        )
+
+    @pytest.mark.parametrize("cls", [Float8ColumnParallelLinear, Float8RowParallelLinear])
+    def test_fp8_overrides_resolve_to_mixin(self, cls):
+        assert cls._forward_impl is _Float8LinearMixin._forward_impl
+        assert cls._apply is _Float8LinearMixin._apply
+
+
+class TestBuildFP8Config:
+    """Tests for _build_fp8_config() — no GPU, no parallel state."""
+
+    def _make_config(self, fp8="e4m3", fp8_recipe=Fp8Recipe.tensorwise):
+        return SimpleNamespace(fp8=fp8, fp8_recipe=fp8_recipe)
+
+    def test_invalid_recipe_raises(self):
+        with pytest.raises(ValueError, match="does not support"):
+            _build_fp8_config(self._make_config(fp8_recipe=Fp8Recipe.custom))
+
+
+def _init_method():
+    return functools.partial(torch.nn.init.xavier_uniform_)
+
+
+def _make_fp8_transformer_config(**overrides):
+    """Build a TransformerConfig suitable for Float8 linear layers."""
+    defaults = dict(
+        hidden_size=64,
+        num_attention_heads=8,
+        num_layers=1,
+        params_dtype=torch.bfloat16,
+        fp8="e4m3",
+        fp8_recipe="tensorwise",
+    )
+    defaults.update(overrides)
+    return TransformerConfig(**defaults)
+
+
+class TestFloat8LinearGuards(PrimusUT):
+    """Tests that Float8 linear layers reject invalid configurations at init."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state, monkeypatch):
+        dummy_args = SimpleNamespace(
+            rank=0,
+            world_size=1,
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            offload=False,
+            offload_ops=[],
+            patch_primus_pipeline=False,
+            pp_algorithm=None,
+            patch_zero_bubble=False,
+            enable_zero_bubble=False,
+            rampup_batch_size=None,
+            global_batch_size=1,
+            micro_batch_size=1,
+            data_parallel_size=1,
+            decrease_batch_size_if_needed=False,
+        )
+        import megatron.training.global_vars as gvars
+
+        monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_rejects_tp_gt_1(self):
+        config = _make_fp8_transformer_config(tensor_model_parallel_size=2)
+        with pytest.raises(ValueError, match="tensor_model_parallel_size=1"):
+            Float8ColumnParallelLinear(
+                input_size=64,
+                output_size=128,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                gather_output=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_rejects_gaf(self):
+        config = _make_fp8_transformer_config(gradient_accumulation_fusion=True)
+        with pytest.raises(ValueError, match="gradient_accumulation_fusion"):
+            Float8ColumnParallelLinear(
+                input_size=64,
+                output_size=128,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                gather_output=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_rejects_sequence_parallel(self):
+        # SP requires TP>1 at TransformerConfig level, so we set TP=2 to
+        # pass config validation and hit our Float8 guard instead.
+        config = _make_fp8_transformer_config(
+            sequence_parallel=True,
+            tensor_model_parallel_size=2,
+        )
+        with pytest.raises(ValueError, match="tensor_model_parallel_size=1"):
+            Float8ColumnParallelLinear(
+                input_size=64,
+                output_size=128,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                gather_output=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_rejects_fp8_none(self):
+        config = _make_fp8_transformer_config(fp8=None)
+        with pytest.raises(ValueError, match="config.fp8"):
+            Float8ColumnParallelLinear(
+                input_size=64,
+                output_size=128,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                gather_output=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_row_parallel_rejects_fp8_none(self):
+        config = _make_fp8_transformer_config(fp8=None)
+        with pytest.raises(ValueError, match="config.fp8"):
+            Float8RowParallelLinear(
+                input_size=64,
+                output_size=128,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                input_is_parallel=True,
+                skip_bias_add=False,
+                is_expert=False,
+            )
+
+
+class TestFloat8LinearInit(PrimusUT):
+    """Tests successful construction and init-time dispatch."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state, monkeypatch):
+        dummy_args = SimpleNamespace(
+            rank=0,
+            world_size=1,
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            offload=False,
+            offload_ops=[],
+            patch_primus_pipeline=False,
+            pp_algorithm=None,
+            patch_zero_bubble=False,
+            enable_zero_bubble=False,
+            rampup_batch_size=None,
+            global_batch_size=1,
+            micro_batch_size=1,
+            data_parallel_size=1,
+            decrease_batch_size_if_needed=False,
+        )
+        import megatron.training.global_vars as gvars
+
+        monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_dispatch_tensorwise_uses_opaque(self):
+        config = _make_fp8_transformer_config(fp8="e4m3", fp8_recipe="tensorwise")
+        layer = Float8ColumnParallelLinear(
+            input_size=64,
+            output_size=128,
+            config=config,
+            init_method=_init_method(),
+            bias=False,
+            gather_output=False,
+            skip_bias_add=False,
+            is_expert=False,
+        )
+        assert layer._fp8_fn is OpaqueFP8LinearTensorwiseFunction
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_dispatch_blockwise(self):
+        config = _make_fp8_transformer_config(fp8="e4m3", fp8_recipe="blockwise")
+        layer = Float8ColumnParallelLinear(
+            input_size=64,
+            output_size=128,
+            config=config,
+            init_method=_init_method(),
+            bias=False,
+            gather_output=False,
+            skip_bias_add=False,
+            is_expert=False,
+        )
+        assert layer._fp8_fn is FP8LinearBlockwiseFunction
+
+
+class TestQuantizeFP8Tensorwise(PrimusUT):
+    """Verify _quantize_fp8_tensorwise correctness and properties."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    def _reference_quantize(self, x, fp8_dtype, fp8_max):
+        """Pure-PyTorch reference (known correct): compute in FP32."""
+        x_f32 = x.float()
+        amax = x_f32.abs().amax()
+        scale = fp8_max / amax.clamp(min=1e-12)
+        x_fp8 = (x_f32 * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
+        scale_inv = 1.0 / scale
+        return x_fp8, scale_inv
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_e4m3_correctness(self):
+        """BF16-scale quantize should match FP32 reference within a few percent
+        of elements (boundary rounding from BF16 scale precision)."""
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        fp8_dtype = float8_e4m3
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        decomp_fp8, decomp_sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        ref_fp8, ref_sinv = self._reference_quantize(x, fp8_dtype, fp8_max)
+
+        mismatch = (decomp_fp8.to(torch.float32) != ref_fp8.to(torch.float32)).float().mean()
+        assert mismatch < 0.05, f"More than 5% of elements differ: {mismatch:.4f}"
+        assert decomp_sinv.dtype == torch.float32
+        torch.testing.assert_close(decomp_sinv, ref_sinv, atol=0, rtol=0)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_e5m2_correctness(self):
+        """BF16-scale quantize should match FP32 reference within 1 FP8 ULP."""
+        x = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda")
+        fp8_dtype = float8_e5m2
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        decomp_fp8, decomp_sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        ref_fp8, ref_sinv = self._reference_quantize(x, fp8_dtype, fp8_max)
+
+        mismatch = (decomp_fp8.to(torch.float32) != ref_fp8.to(torch.float32)).float().mean()
+        assert mismatch < 0.02, f"More than 2% of elements differ: {mismatch:.4f}"
+        assert decomp_sinv.dtype == torch.float32
+        torch.testing.assert_close(decomp_sinv, ref_sinv, atol=0, rtol=0)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_dequant_roundtrip(self):
+        """Verify quantize -> dequantize preserves values within FP8 precision."""
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        fp8_dtype = float8_e4m3
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        x_recon = x_fp8.float() * scale_inv
+
+        rel_err = ((x_recon - x.float()).abs() / x.float().abs().clamp(min=1e-12)).mean()
+        assert rel_err < 0.1, f"Mean relative error {rel_err:.4f} too high for e4m3"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_zero_tensor_no_nan(self):
+        x = torch.zeros(32, 64, dtype=torch.bfloat16, device="cuda")
+        fp8_dtype = float8_e4m3
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        fp8_out, sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        assert not torch.isnan(sinv).any(), "scale_inv should not be NaN for zero input"
+        assert not torch.isinf(sinv).any(), "scale_inv should not be Inf for zero input"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_near_zero_no_nan(self):
+        """Tensors with very small amax should not produce NaN or Inf in output."""
+        x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") * 1e-6
+        fp8_dtype = float8_e4m3
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        assert not torch.isnan(x_fp8.float()).any(), "FP8 output has NaN for near-zero input"
+        assert not torch.isinf(x_fp8.float()).any(), "FP8 output has Inf for near-zero input"
+        assert not torch.isnan(scale_inv).any(), "scale_inv has NaN for near-zero input"
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_scale_bf16_overflow_guard(self):
+        """When amax is tiny, FP32 scale would overflow BF16.  Clamp must prevent it."""
+        x = torch.full((4, 4), 1e-8, dtype=torch.bfloat16, device="cuda")
+        fp8_dtype = float8_e4m3
+        fp8_max = torch.finfo(fp8_dtype).max
+
+        x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max)
+        assert not torch.isnan(x_fp8.float()).any(), "BF16 overflow guard failed — NaN in output"
+        assert torch.isfinite(scale_inv), "scale_inv should be finite"
+
+
+class TestDecomposedTensorwiseCompile(PrimusUT):
+    """Verify DecomposedFP8LinearTensorwiseFunction works with torch.compile."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_no_graph_break(self):
+        """torch._dynamo.explain should report zero graph breaks for the
+        decomposed quantize + gemm_fp8_impl forward."""
+        from primus_turbo.pytorch.core.backend import BackendType
+
+        fp8_fwd_dtype = float8_e4m3
+        fp8_bwd_dtype = float8_e5m2
+        fp8_fwd_max = torch.finfo(fp8_fwd_dtype).max
+        fp8_bwd_max = torch.finfo(fp8_bwd_dtype).max
+        gran_value = ScalingGranularity.TENSORWISE.value
+        backend_value = BackendType.HIPBLASLT.value
+
+        x = torch.randn(4, 64, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda")
+
+        explanation = torch._dynamo.explain(
+            DecomposedFP8LinearTensorwiseFunction.apply,
+        )(x, w, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value)
+        assert explanation.graph_break_count == 0, (
+            f"Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+
+
+class TestOpaqueTensorwiseCompile(PrimusUT):
+    """Verify the refactored OpaqueFP8LinearTensorwiseFunction has zero graph breaks."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU")
+    def test_no_graph_break(self):
+        """torch._dynamo.explain should report zero graph breaks for the
+        setup_context-based OpaqueFP8LinearTensorwiseFunction forward."""
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.ops.quantization import quantize_fp8
+
+        torch._dynamo.reset()
+
+        fp8_fwd_dtype = float8_e4m3
+        fp8_bwd_dtype = float8_e5m2
+        gran_value = ScalingGranularity.TENSORWISE.value
+        backend_value = BackendType.HIPBLASLT.value
+
+        x = torch.randn(4, 64, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda")
+        w_fp8, w_scale = quantize_fp8(w, fp8_fwd_dtype, ScalingGranularity.TENSORWISE)
+
+        explanation = torch._dynamo.explain(
+            OpaqueFP8LinearTensorwiseFunction.apply,
+        )(x, w, w_fp8, w_scale, fp8_fwd_dtype, fp8_bwd_dtype, gran_value, backend_value)
+        assert explanation.graph_break_count == 0, (
+            f"Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py
new file mode 100644
index 000000000..a0fc099b9
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py
@@ -0,0 +1,340 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for compile-friendly MXFP4 linear layers (primus_turbo_mxfp4_local).
+
+Tests cross-validation against Primus-Turbo's FP4GemmMXFunction reference,
+torch.compile graph-break validation, Megatron linear backward flow,
+2-step training loop, hybrid (FP4 fwd / FP8 bwd) mode, and init guards.
+"""
+
+import functools
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from tests.unit_tests.backends.megatron.conftest import requires_mxfp4
+from tests.utils import PrimusUT
+
+
+def _init_method():
+    return functools.partial(torch.nn.init.xavier_uniform_)
+
+
+def _make_mxfp4_transformer_config(**overrides):
+    defaults = dict(
+        hidden_size=256,
+        num_attention_heads=8,
+        num_layers=1,
+        params_dtype=torch.bfloat16,
+        fp4="mxfp4",
+        fp4_recipe="mxfp4",
+    )
+    defaults.update(overrides)
+    return TransformerConfig(**defaults)
+
+
+def _pin_fp4_aiter(monkeypatch):
+    """Pin the FP4 GEMM backend to AITER with autotune off for MXFP4 module tests.
+
+    MXFP4 module __init__ runs _assert_preshuffle_contract, which requires the
+    FP4 GEMM backend pinned to AITER with autotune off (the only config under
+    which _enable_preshuffle() is True). Pinning it in-code lets the
+    module-instantiation tests reach the real path instead of failing the
+    contract; monkeypatch auto-restores on teardown so the .apply-direct tests
+    keep their default (preshuffle=False) dispatch. Also clears any baked-empty
+    PRIMUS_TURBO_GEMM_BACKEND so the in-code pin is authoritative (mirrors
+    test_native_fp8_layout.py).
+    """
+    from primus_turbo.pytorch.core.backend import (
+        BackendType,
+        GlobalBackendManager,
+        PrecisionType,
+    )
+
+    if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "":
+        monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False)
+    monkeypatch.setattr(GlobalBackendManager, "_gemm_backend", {PrecisionType.FP4: BackendType.AITER})
+    monkeypatch.setattr(GlobalBackendManager, "_auto_tune", False)
+
+
+# ---------------------------------------------------------------------------
+# Cross-validation against Primus-Turbo's FP4GemmMXFunction reference
+# ---------------------------------------------------------------------------
+
+
+class TestMXFP4CrossValidation(PrimusUT):
+    """Verify MXFP4LinearFunction produces bit-identical results to FP4GemmMXFunction.
+
+    Catches wrong boolean flags in _quantize_input_dual / _quantize_weight_dual /
+    _quantize_grad_dual. A single wrong flag produces silently incorrect numerics
+    that may still pass SNR thresholds vs BF16 but diverges from the canonical path.
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @requires_mxfp4
+    def test_forward_matches_reference_fp4gemm(self):
+        from primus_turbo.pytorch.core.low_precision import Float4QuantConfig
+        from primus_turbo.pytorch.ops.gemm_fp4 import FP4GemmMXFunction
+
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4LinearFunction,
+            _enable_preshuffle,
+        )
+
+        torch.manual_seed(42)
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda")
+        preshuffle = _enable_preshuffle()
+
+        result = MXFP4LinearFunction.apply(
+            x,
+            w,
+            preshuffle,
+            False,
+            None,
+            0,
+            0,
+            False,
+        )
+        our_output = result[0]
+
+        config = Float4QuantConfig(use_preshuffle=preshuffle)
+        ref_output = FP4GemmMXFunction.apply(
+            x.clone(),
+            w.clone(),
+            None,
+            None,
+            False,
+            True,
+            x.dtype,
+            config,
+        )
+
+        assert torch.equal(our_output, ref_output), (
+            f"Forward outputs differ. Max abs diff: " f"{(our_output - ref_output).abs().max().item():.6e}"
+        )
+
+    @requires_mxfp4
+    def test_backward_matches_reference_fp4gemm(self):
+        from primus_turbo.pytorch.core.low_precision import Float4QuantConfig
+        from primus_turbo.pytorch.ops.gemm_fp4 import FP4GemmMXFunction
+
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4LinearFunction,
+            _enable_preshuffle,
+        )
+
+        torch.manual_seed(42)
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        x_ref = x.detach().clone().requires_grad_(True)
+        w_ref = w.detach().clone().requires_grad_(True)
+        preshuffle = _enable_preshuffle()
+
+        result = MXFP4LinearFunction.apply(x, w, preshuffle, False, None, 0, 0, False)
+        our_output = result[0]
+        grad_out = torch.ones_like(our_output)
+        our_output.backward(grad_out)
+
+        config = Float4QuantConfig(use_preshuffle=preshuffle)
+        ref_output = FP4GemmMXFunction.apply(
+            x_ref,
+            w_ref,
+            None,
+            None,
+            False,
+            True,
+            x_ref.dtype,
+            config,
+        )
+        ref_output.backward(torch.ones_like(ref_output))
+
+        # grad_weight stays bit-identical to the reference: our grad_weight
+        # GEMM pair (g_t colwise + a_t colwise) and the reference's both use the
+        # RHT recipe, so the quantization is identical. Keep torch.equal here --
+        # it still catches a wrong flag in _quantize_input_dual (colwise a_t) or
+        # the g_t branch of _quantize_grad_dual.
+        assert torch.equal(w.grad, w_ref.grad), (
+            f"grad_weight differs. Max abs diff: " f"{(w.grad - w_ref.grad).abs().max().item():.6e}"
+        )
+
+        # grad_input is NOT bit-identical, and that is expected post Primus-Turbo
+        # PR #383. The grad_input GEMM pair is (grad rowwise) x (weight colwise b_t).
+        # Our production deliberately quantizes this pair without RHT
+        # (_quantize_grad_dual rowwise use_rht=False + _quantize_weight_dual
+        # colwise use_rht=False -- an internally consistent no-RHT pair), whereas
+        # Primus-Turbo PR #383's FP4GemmMXFunction.backward unconditionally quantizes the grad
+        # with use_rht=True and derives b_t with use_rht=True. Both compute a
+        # valid grad_input (RHT cancels within each consistent pair); they only
+        # differ in MXFP4 quantization noise. Measured against the true BF16
+        # gradient, our no-RHT grad_input is actually marginally more accurate
+        # than the reference's RHT grad_input (~18.4 dB vs ~17.7 dB SNR), so this
+        # is a recipe choice, not a regression. Assert SNR vs the BF16 truth
+        # (same >10 dB bar as the forward SNR test) instead of bit-identity.
+        bf16_grad_input = grad_out.float() @ w.detach().float()
+        signal = (bf16_grad_input**2).mean()
+        noise = ((x.grad.float() - bf16_grad_input) ** 2).mean()
+        snr_db = 10 * torch.log10(signal / noise).item()
+        assert snr_db > 10, (
+            f"grad_input SNR {snr_db:.1f} dB vs BF16 is below the 10 dB threshold "
+            f"(max abs diff vs Primus-Turbo PR #383 reference: {(x.grad - x_ref.grad).abs().max().item():.6e})"
+        )
+
+
+# ---------------------------------------------------------------------------
+# torch.compile graph-break validation
+# ---------------------------------------------------------------------------
+
+
+class TestMXFP4Compile(PrimusUT):
+    """Verify MXFP4LinearFunction has zero graph breaks under torch.compile."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state):
+        pass
+
+    @requires_mxfp4
+    def test_no_graph_break_pure_mxfp4(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4LinearFunction,
+            _enable_preshuffle,
+        )
+
+        torch._dynamo.reset()
+
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda")
+        preshuffle = _enable_preshuffle()
+
+        explanation = torch._dynamo.explain(
+            MXFP4LinearFunction.apply,
+        )(x, w, preshuffle, False, None, 0, 0, False)
+
+        assert explanation.graph_break_count == 0, (
+            f"Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+
+    @requires_mxfp4
+    def test_no_graph_break_hybrid(self):
+        from primus_turbo.pytorch.core.backend import BackendType
+        from primus_turbo.pytorch.core.low_precision import (
+            ScalingGranularity,
+            float8_e5m2,
+        )
+
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4LinearFunction,
+            _enable_preshuffle,
+        )
+
+        torch._dynamo.reset()
+
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda")
+        preshuffle = _enable_preshuffle()
+
+        explanation = torch._dynamo.explain(
+            MXFP4LinearFunction.apply,
+        )(
+            x,
+            w,
+            preshuffle,
+            True,
+            float8_e5m2,
+            ScalingGranularity.TENSORWISE.value,
+            BackendType.HIPBLASLT.value,
+            False,
+        )
+
+        assert explanation.graph_break_count == 0, (
+            f"Expected 0 graph breaks, got {explanation.graph_break_count}. "
+            f"Reasons: {explanation.break_reasons}"
+        )
+
+    @requires_mxfp4
+    def test_compiled_forward_matches_eager(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4LinearFunction,
+            _enable_preshuffle,
+        )
+
+        torch._dynamo.reset()
+        torch.manual_seed(42)
+
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda")
+        preshuffle = _enable_preshuffle()
+
+        eager_result = MXFP4LinearFunction.apply(x, w, preshuffle, False, None, 0, 0, False)
+        eager_out = eager_result[0]
+
+        compiled_fn = torch.compile(MXFP4LinearFunction.apply)
+        compiled_result = compiled_fn(x, w, preshuffle, False, None, 0, 0, False)
+        compiled_out = compiled_result[0]
+
+        assert torch.equal(eager_out, compiled_out), (
+            f"Compiled output differs from eager. Max abs diff: "
+            f"{(eager_out - compiled_out).abs().max().item():.6e}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Init guard (TP > 1 rejection)
+# ---------------------------------------------------------------------------
+
+
+class TestMXFP4LinearGuard(PrimusUT):
+    """Test that MXFP4 linear layers reject invalid configurations."""
+
+    @pytest.fixture(autouse=True)
+    def setup_parallel(self, init_parallel_state, monkeypatch):
+        dummy_args = SimpleNamespace(
+            rank=0,
+            world_size=1,
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            offload=False,
+            offload_ops=[],
+            patch_primus_pipeline=False,
+            pp_algorithm=None,
+            patch_zero_bubble=False,
+            enable_zero_bubble=False,
+            rampup_batch_size=None,
+            global_batch_size=1,
+            micro_batch_size=1,
+            data_parallel_size=1,
+            decrease_batch_size_if_needed=False,
+        )
+        import megatron.training.global_vars as gvars
+
+        monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args)
+
+        _pin_fp4_aiter(monkeypatch)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_column_parallel_rejects_tp_gt_1(self):
+        from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import (
+            MXFP4ColumnParallelLinear,
+        )
+
+        config = _make_mxfp4_transformer_config(tensor_model_parallel_size=2)
+        with pytest.raises(ValueError, match="tensor_model_parallel_size=1"):
+            MXFP4ColumnParallelLinear(
+                input_size=256,
+                output_size=512,
+                config=config,
+                init_method=_init_method(),
+                bias=False,
+                gather_output=False,
+                skip_bias_add=False,
+                is_expert=False,
+            )
diff --git a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py
new file mode 100644
index 000000000..2e9bb97a5
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py
@@ -0,0 +1,54 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from types import SimpleNamespace
+
+import pytest
+
+from primus.backends.megatron.patches.args.rocm_arg_validation import (
+    validate_turbo_ragged_grouped_gemm,
+)
+
+
+def _ragged_args(**overrides):
+    values = {
+        "use_turbo_ragged_grouped_gemm": True,
+        "enable_primus_turbo": True,
+        "use_turbo_grouped_gemm": True,
+        "fp8": "e4m3",
+        "fp8_recipe": "tensorwise",
+        "fp4": False,
+        "moe_router_padding_for_quantization": False,
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_turbo_ragged_grouped_gemm_accepts_tensorwise_fp8():
+    validate_turbo_ragged_grouped_gemm(_ragged_args())
+
+
+def test_turbo_ragged_grouped_gemm_disabled_is_noop():
+    validate_turbo_ragged_grouped_gemm(SimpleNamespace())
+
+
+@pytest.mark.parametrize(
+    ("overrides", "message"),
+    [
+        ({"enable_primus_turbo": False}, "requires enable_primus_turbo"),
+        ({"use_turbo_grouped_gemm": False}, "requires enable_primus_turbo"),
+        ({"fp8": None}, "supports only tensorwise FP8"),
+        ({"fp8_recipe": "blockwise"}, "supports only tensorwise FP8"),
+        ({"fp4": True}, "supports only tensorwise FP8"),
+        (
+            {"moe_router_padding_for_quantization": True},
+            "requires moe_router_padding_for_quantization=False",
+        ),
+    ],
+)
+def test_turbo_ragged_grouped_gemm_rejects_unsupported_config(overrides, message):
+    with pytest.raises(ValueError, match=message):
+        validate_turbo_ragged_grouped_gemm(_ragged_args(**overrides))
diff --git a/tests/unit_tests/backends/megatron/test_rope_arange_cache.py b/tests/unit_tests/backends/megatron/test_rope_arange_cache.py
new file mode 100644
index 000000000..f93efd7f1
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_rope_arange_cache.py
@@ -0,0 +1,80 @@
+# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+"""Parity + caching test for ``RoPECache.forward_arange``.
+
+The compressed-branch RoPE is evaluated at deterministic positions every
+forward (``arange(P) * compress_ratio``); ``forward_arange`` caches that table.
+This asserts the cached table is bit-identical to recomputing
+``forward(arange(n) * stride)``, that ``stride`` is honoured and part of the
+cache key, and that the cache actually memoises (with
+``PRIMUS_COMPRESS_ROPE_CACHE=0`` bypassing it).
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+from primus.backends.megatron.core.transformer.dual_rope import RoPECache
+
+
+def _rope() -> RoPECache:
+    return RoPECache(rotary_dim=64, theta=10000.0)
+
+
+@pytest.mark.parametrize("n", [32, 1024])
+def test_forward_arange_matches_forward(n):
+    """Cached table is bit-identical to the eager arange->outer->cos/sin path."""
+    rc = _rope()
+    cos_a, sin_a = rc.forward_arange(n, "cpu")
+    cos_e, sin_e = rc.forward(torch.arange(n, device="cpu"))
+    torch.testing.assert_close(cos_a, cos_e, rtol=0, atol=0)
+    torch.testing.assert_close(sin_a, sin_e, rtol=0, atol=0)
+
+
+def test_forward_arange_memoises(monkeypatch):
+    """A repeat call returns the SAME cached tensors (no recompute)."""
+    monkeypatch.setenv("PRIMUS_COMPRESS_ROPE_CACHE", "1")
+    rc = _rope()
+    a = rc.forward_arange(128, "cpu")
+    b = rc.forward_arange(128, "cpu")
+    assert a[0] is b[0] and a[1] is b[1]
+
+
+def test_cache_disabled_recomputes(monkeypatch):
+    """PRIMUS_COMPRESS_ROPE_CACHE=0 bypasses the cache (fresh, equal tensors)."""
+    monkeypatch.setenv("PRIMUS_COMPRESS_ROPE_CACHE", "0")
+    rc = _rope()
+    a = rc.forward_arange(128, "cpu")
+    b = rc.forward_arange(128, "cpu")
+    assert a[0] is not b[0]
+    torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("stride", [4, 128])
+def test_forward_arange_stride_uses_original_positions(stride):
+    """``stride`` evaluates at ``arange(n) * stride``, not at block indices.
+
+    Compressed KV entry ``s`` covers the window starting at original token
+    ``s * compress_ratio`` and must be rotated at that position so it shares
+    the query coordinate system.
+    """
+    n = 16
+    rc = _rope()
+    cos_a, sin_a = rc.forward_arange(n, "cpu", stride=stride)
+    cos_e, sin_e = rc.forward(torch.arange(n, device="cpu") * stride)
+    torch.testing.assert_close(cos_a, cos_e, rtol=0, atol=0)
+    torch.testing.assert_close(sin_a, sin_e, rtol=0, atol=0)
+
+    # Guard against a silently-ignored stride.
+    cos_blocks, _ = rc.forward_arange(n, "cpu", stride=1)
+    assert not torch.allclose(cos_a, cos_blocks)
+
+
+def test_forward_arange_stride_is_part_of_cache_key():
+    """Different strides must not collide in the memo table."""
+    rc = _rope()
+    cos_s1, _ = rc.forward_arange(32, "cpu", stride=1)
+    cos_s4, _ = rc.forward_arange(32, "cpu", stride=4)
+    torch.testing.assert_close(cos_s4, rc.forward(torch.arange(32) * 4)[0], rtol=0, atol=0)
+    torch.testing.assert_close(cos_s1, rc.forward(torch.arange(32))[0], rtol=0, atol=0)
+    # And each stride still memoises independently.
+    assert rc.forward_arange(32, "cpu", stride=4)[0] is cos_s4
diff --git a/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py b/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py
index b29fe36e6..406a61051 100644
--- a/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py
+++ b/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py
@@ -26,9 +26,9 @@ def _orig_compile():
     training_pkg.initialize = init_mod
     megatron_mod.training = training_pkg
 
-    sys.modules["megatron"] = megatron_mod
-    sys.modules["megatron.training"] = training_pkg
-    sys.modules["megatron.training.initialize"] = init_mod
+    monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
+    monkeypatch.setitem(sys.modules, "megatron.training", training_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.training.initialize", init_mod)
     return init_mod, _orig_compile
 
 
diff --git a/tests/unit_tests/backends/megatron/test_training_log_patches.py b/tests/unit_tests/backends/megatron/test_training_log_patches.py
index 2f8f49683..bcd781de9 100644
--- a/tests/unit_tests/backends/megatron/test_training_log_patches.py
+++ b/tests/unit_tests/backends/megatron/test_training_log_patches.py
@@ -32,9 +32,9 @@ def fake_training_log(*args, **kwargs):
         return "ok"
 
     # Provide a minimal `get_model` stub so that any code which expects
-    # `megatron.training.training.get_model` (e.g., Primus monkey patches in
-    # `primus.pretrain.load_backend_trainer`) can safely import and patch this
-    # fake training module during tests without raising AttributeError.
+    # `megatron.training.training.get_model` (e.g., Primus monkey patches)
+    # can safely import and patch this fake training module during tests
+    # without raising AttributeError.
     def fake_get_model(*args, **kwargs):
         return None
 
@@ -45,9 +45,9 @@ def fake_get_model(*args, **kwargs):
     training_pkg.training = training_mod
     megatron_mod.training = training_pkg
 
-    sys.modules["megatron"] = megatron_mod
-    sys.modules["megatron.training"] = training_pkg
-    sys.modules["megatron.training.training"] = training_mod
+    monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
+    monkeypatch.setitem(sys.modules, "megatron.training", training_pkg)
+    monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod)
 
     return training_mod, fake_training_log
 
@@ -140,9 +140,10 @@ def test_patch_training_log_wraps_and_stacks_extensions(monkeypatch: pytest.Monk
 
 
 def test_rocm_monitor_hooked_print_rank_last_injects_stats(monkeypatch: pytest.MonkeyPatch):
-    # Prepare fake torch and ROCm SMI helpers. MemoryStatsExtension also calls
-    # torch.tensor / torch.distributed.all_gather for cross-rank max ROCm mem;
-    # stub those so inject() does not fall into the exception path in unit tests.
+    # Prepare fake torch and ROCm SMI helpers. MemoryStatsExtension computes the
+    # cross-rank max ROCm mem via torch.tensor + torch.distributed.all_reduce(MAX)
+    # and torch.distributed.get_rank(); stub those so inject() does not fall into
+    # the exception path in unit tests.
     class _FakeTensor:
         __slots__ = ("_value",)
 
@@ -163,13 +164,17 @@ def _zeros_like(_t):
     def _get_world_size():
         return 8
 
-    def _all_gather(gathered_list, tensor):
-        v = tensor.item()
-        for i in range(len(gathered_list)):
-            gathered_list[i] = _FakeTensor(v)
+    def _get_rank():
+        return 0
+
+    def _all_reduce(tensor, op=None):
+        # Single-rank reduction: MAX over one rank is a no-op, so leave the
+        # in-place tensor value unchanged (matching real all_reduce semantics).
+        return None
 
     # `inject` passes `dtype=torch.int64`; SimpleNamespace must expose `int64`
-    # or attribute lookup fails before our fake `tensor()` runs.
+    # or attribute lookup fails before our fake `tensor()` runs. The refactored
+    # code also reads `torch.distributed.ReduceOp.MAX`, so provide that too.
     fake_torch = SimpleNamespace(
         cuda=SimpleNamespace(
             mem_get_info=lambda: (2 * 1024**3, 4 * 1024**3),
@@ -180,7 +185,9 @@ def _all_gather(gathered_list, tensor):
         int64=int,
         distributed=SimpleNamespace(
             get_world_size=_get_world_size,
-            all_gather=_all_gather,
+            get_rank=_get_rank,
+            all_reduce=_all_reduce,
+            ReduceOp=SimpleNamespace(MAX=object()),
         ),
     )
     monkeypatch.setattr(
diff --git a/tests/unit_tests/backends/megatron/test_validate_args_patches.py b/tests/unit_tests/backends/megatron/test_validate_args_patches.py
index 790775ffb..bd58796cf 100644
--- a/tests/unit_tests/backends/megatron/test_validate_args_patches.py
+++ b/tests/unit_tests/backends/megatron/test_validate_args_patches.py
@@ -53,7 +53,7 @@ def _install_fake_megatron(monkeypatch: pytest.MonkeyPatch):
     """Create fake ``megatron.*`` modules sufficient for validate_args patches.
 
     Includes ``megatron.core.parallel_state`` and ``megatron.training.global_vars``
-    so that transitive imports from ``primus.modules.trainer.megatron.utils`` succeed.
+    so that transitive imports from ``primus.backends.megatron.patches.args.rocm_arg_validation`` succeed.
     """
     import sys
 
@@ -149,7 +149,7 @@ def test_wraps_validate_args_on_both_modules(self, monkeypatch):
         args_mod, init_mod = _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: setattr(args, "_rocm_validated", True),
         )
 
@@ -167,7 +167,7 @@ def test_calls_original_and_rocm_validation(self, monkeypatch):
         _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: setattr(args, "_rocm_validated", True),
         )
 
@@ -190,7 +190,7 @@ def test_stores_original_on_module(self, monkeypatch):
         args_mod, _ = _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: None,
         )
 
@@ -212,7 +212,7 @@ def _apply_base_and_split(self, monkeypatch):
         _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: None,
         )
 
@@ -252,7 +252,7 @@ def _apply_base_and_fp4(self, monkeypatch):
         _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: None,
         )
 
@@ -288,7 +288,7 @@ def test_without_patch_fp4_raises(self, monkeypatch):
         _install_fake_megatron(monkeypatch)
         _silence_logging(monkeypatch)
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: None,
         )
 
@@ -314,7 +314,7 @@ def test_pipeline_split_and_rocm_validation(self, monkeypatch):
 
         rocm_calls = []
         monkeypatch.setattr(
-            "primus.modules.trainer.megatron.utils.validate_args_on_rocm",
+            "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm",
             lambda args: rocm_calls.append(True),
         )
 
diff --git a/tests/unit_tests/backends/megatron/test_warmup_convergence.py b/tests/unit_tests/backends/megatron/test_warmup_convergence.py
new file mode 100644
index 000000000..07ab49219
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_warmup_convergence.py
@@ -0,0 +1,363 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Single-GPU end-to-end test proving that warmup is transparent.
+
+A model that undergoes warmup steps (with param snapshot, optimizer neutering,
+FP8 reset) followed by a full state restore produces the **exact same
+training trajectory** as a model that trains from scratch.
+
+Parametrised over both FP8 spec paths (local delayed-scaling and
+TransformerEngine).  The trainable modules are ``torch.nn.Linear`` subclasses
+decorated with FP8 metadata; the forward pass runs standard bf16 matmuls.
+Loss-matching validates param snapshot/restore + optimizer reset.  FP8 state
+health (no NaN scales, registry reinit, amax seeding) is checked separately.
+
+Run:
+    python -m pytest tests/unit_tests/backends/megatron/test_warmup_convergence.py -v
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+
+DEVICE = "cuda:0"
+FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max
+FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max
+
+WARMUP_STEPS = 3
+TRAIN_STEPS = 20
+IN_FEATURES = 32
+HIDDEN = 64
+OUT_FEATURES = 16
+BATCH = 8
+
+
+# ---------------------------------------------------------------------------
+# Trainable FP8 module helpers
+# ---------------------------------------------------------------------------
+
+
+class _DelayedScalingLinear(torch.nn.Linear):
+    """``torch.nn.Linear`` with local-spec delayed-scaling FP8 metadata."""
+
+    def __init__(self, in_features, out_features, history_len=1):
+        super().__init__(in_features, out_features)
+        self._use_delayed_scaling = True
+        self._fp8_fwd_dtype = torch.float8_e4m3fn
+        self._fp8_bwd_dtype = torch.float8_e5m2
+        self._fp8_fwd_max = FP8_FWD_MAX
+        self._fp8_bwd_max = FP8_BWD_MAX
+        self._amax_compute_algo = "most_recent"
+        self._first_delayed_step = True
+        self._history_idx = 0
+        self.config = SimpleNamespace(
+            fp8_amax_history_len=history_len,
+            fp8_amax_compute_algo="most_recent",
+        )
+
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _init_delayed_scaling_state,
+        )
+
+        _init_delayed_scaling_state(self)
+
+    def _move_fp8_buffers(self, device):
+        for name in list(self._buffers):
+            buf = self._buffers[name]
+            if buf is not None and buf.device != device:
+                self._buffers[name] = buf.to(device)
+
+
+class _TEScalingLinear(torch.nn.Linear):
+    """``torch.nn.Linear`` with TransformerEngine-style FP8 metadata."""
+
+    def __init__(self, in_features, out_features, device=DEVICE):
+        super().__init__(in_features, out_features)
+        self.fp8_initialized = True
+        self.fp8_meta = {
+            "scaling_fwd": SimpleNamespace(
+                amax_history=torch.zeros(16, 1, device=device),
+                scale=torch.ones(1, device=device),
+                scale_inv=torch.ones(1, device=device),
+            ),
+            "scaling_bwd": SimpleNamespace(
+                amax_history=torch.zeros(16, 1, device=device),
+                scale=torch.ones(1, device=device),
+                scale_inv=torch.ones(1, device=device),
+            ),
+        }
+
+
+def _make_model(fp8_spec, seed, device=DEVICE):
+    """Create a 3-layer MLP with the requested FP8 metadata, on ``device``."""
+    torch.manual_seed(seed)
+    if fp8_spec == "local":
+        model = torch.nn.Sequential(
+            _DelayedScalingLinear(IN_FEATURES, HIDDEN),
+            torch.nn.ReLU(),
+            _DelayedScalingLinear(HIDDEN, OUT_FEATURES),
+        ).to(device)
+        for m in model.modules():
+            if isinstance(m, _DelayedScalingLinear):
+                m._move_fp8_buffers(torch.device(device))
+    elif fp8_spec == "te":
+        model = torch.nn.Sequential(
+            _TEScalingLinear(IN_FEATURES, HIDDEN, device=device),
+            torch.nn.ReLU(),
+            _TEScalingLinear(HIDDEN, OUT_FEATURES, device=device),
+        ).to(device)
+    else:
+        raise ValueError(f"Unknown fp8_spec: {fp8_spec}")
+    return model
+
+
+def _get_delayed_modules(model):
+    return [m for m in model.modules() if getattr(m, "_use_delayed_scaling", False)]
+
+
+def _make_registry(model):
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _DelayedScalingRegistry,
+    )
+
+    modules = _get_delayed_modules(model)
+    assert modules, "No delayed-scaling modules found"
+    return _DelayedScalingRegistry(modules)
+
+
+def _generate_fixed_data(num_batches, seed=99, device=DEVICE):
+    torch.manual_seed(seed)
+    return [
+        (
+            torch.randn(BATCH, IN_FEATURES, device=device),
+            torch.randn(BATCH, OUT_FEATURES, device=device),
+        )
+        for _ in range(num_batches)
+    ]
+
+
+def _run_warmup(model, adam_opt, fp8_spec, registry=None):
+    """Execute the full warmup pipeline: snapshot, neuter, SGD steps, restore.
+
+    Returns the registry (possibly reinitialised for local spec).
+    """
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _fast_update_scales,
+    )
+    from primus.backends.megatron.patches.mlperf_warmup_patches import (
+        _neuter_optimizer,
+        _reset_fp8_local_spec,
+        _reset_fp8_te_spec,
+        _reset_optimizer_state,
+        _restore_optimizer,
+        _seed_fp8_amax,
+    )
+
+    wrapper = SimpleNamespace(optimizer=adam_opt)
+    device = next(model.parameters()).device
+
+    # 1. Snapshot params to CPU
+    saved_params = {name: p.data.to("cpu", non_blocking=True) for name, p in model.named_parameters()}
+    torch.cuda.synchronize()
+
+    # 2. Neuter Adam
+    saved_hyp = _neuter_optimizer(wrapper)
+
+    # 3-4. Throwaway SGD warmup steps
+    warmup_sgd = torch.optim.SGD(model.parameters(), lr=0.01)
+    for _ in range(WARMUP_STEPS):
+        if fp8_spec == "local" and registry is not None:
+            _fast_update_scales(registry)
+        x = torch.randn(BATCH, IN_FEATURES, device=device)
+        y = torch.randn(BATCH, OUT_FEATURES, device=device)
+        loss = (model(x) - y).pow(2).mean()
+        loss.backward()
+        warmup_sgd.step()
+        warmup_sgd.zero_grad()
+
+    # 5. Restore Adam hyperparams
+    _restore_optimizer(wrapper, saved_hyp)
+
+    # 6. Reset Adam step counters
+    _reset_optimizer_state(wrapper)
+
+    # 7. Restore params from CPU snapshot
+    for name, p in model.named_parameters():
+        if name in saved_params:
+            p.data.copy_(saved_params[name])
+    del saved_params
+
+    # 8. Reset FP8 state
+    if fp8_spec == "local":
+        _reset_fp8_local_spec([model])
+    elif fp8_spec == "te":
+        _reset_fp8_te_spec([model])
+        _seed_fp8_amax([model])
+
+    # 9. Synchronize
+    torch.cuda.synchronize()
+
+    # 10. Zero gradients
+    adam_opt.zero_grad(set_to_none=True)
+
+    return registry
+
+
+def _train_loop(model, optimizer, data, fp8_spec, registry=None):
+    """Run training and return per-step losses."""
+    from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+        _fast_update_scales,
+    )
+
+    losses = []
+    for x, y in data:
+        if fp8_spec == "local" and registry is not None:
+            _fast_update_scales(registry)
+        loss = (model(x) - y).pow(2).mean()
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        losses.append(loss.detach())
+    return losses
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("fp8_spec", ["local", "te"])
+class TestWarmupConvergence:
+
+    def test_loss_matches_baseline(self, fp8_spec):
+        """After warmup + full reset, loss curve is identical to no-warmup baseline."""
+        seed = 42
+
+        baseline_model = _make_model(fp8_spec, seed)
+        warmup_model = _make_model(fp8_spec, seed)
+
+        # Verify identical starting params
+        for (n1, p1), (n2, p2) in zip(baseline_model.named_parameters(), warmup_model.named_parameters()):
+            assert torch.equal(p1.data, p2.data), f"Init mismatch on {n1}"
+
+        baseline_opt = torch.optim.Adam(
+            baseline_model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01
+        )
+        warmup_opt = torch.optim.Adam(
+            warmup_model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01
+        )
+
+        baseline_reg = _make_registry(baseline_model) if fp8_spec == "local" else None
+        warmup_reg = _make_registry(warmup_model) if fp8_spec == "local" else None
+
+        # Run warmup on warmup_model
+        _run_warmup(warmup_model, warmup_opt, fp8_spec, registry=warmup_reg)
+
+        # Generate fixed training data
+        data = _generate_fixed_data(TRAIN_STEPS)
+
+        # Train both
+        baseline_losses = _train_loop(baseline_model, baseline_opt, data, fp8_spec, registry=baseline_reg)
+        warmup_losses = _train_loop(warmup_model, warmup_opt, data, fp8_spec, registry=warmup_reg)
+
+        for i, (lb, lw) in enumerate(zip(baseline_losses, warmup_losses)):
+            if i < 5:
+                assert torch.equal(
+                    lb, lw
+                ), f"Step {i}: bitwise mismatch baseline={lb.item()}, warmup={lw.item()}"
+            else:
+                assert abs(lb.item() - lw.item()) < 1e-5, (
+                    f"Step {i}: baseline={lb.item()}, warmup={lw.item()}, "
+                    f"diff={abs(lb.item() - lw.item())}"
+                )
+
+    def test_fp8_state_healthy_after_reset(self, fp8_spec):
+        """FP8 metadata has no NaN/Inf after warmup + reset."""
+        from primus.backends.megatron.core.extensions.primus_turbo_float8_local import (
+            _fast_update_scales,
+        )
+
+        seed = 42
+        model = _make_model(fp8_spec, seed)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+        registry = _make_registry(model) if fp8_spec == "local" else None
+
+        _run_warmup(model, opt, fp8_spec, registry=registry)
+
+        if fp8_spec == "local":
+            _fast_update_scales(registry)
+
+            assert (
+                registry._first_step is False
+            ), "_first_step should be consumed after post-reset _fast_update_scales"
+            assert torch.isfinite(
+                registry.scales_3n
+            ).all(), f"Non-finite scales after reset: {registry.scales_3n}"
+
+            for i, m in enumerate(_get_delayed_modules(model)):
+                expected_amax = m.weight.data.abs().amax().float().item()
+                actual_amax = m.staged_weight_amax.item()
+                assert (
+                    abs(actual_amax - expected_amax) < 1e-3
+                ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}"
+
+        elif fp8_spec == "te":
+            for m in model.modules():
+                if hasattr(m, "fp8_initialized"):
+                    assert m.fp8_initialized is False, "fp8_initialized should be False after TE reset"
+                if hasattr(m, "fp8_meta"):
+                    meta = m.fp8_meta
+                    for key in ("scaling_fwd", "scaling_bwd"):
+                        if key not in meta:
+                            continue
+                        tm = meta[key]
+                        if hasattr(tm, "amax_history"):
+                            assert (tm.amax_history == 1.0).all(), (
+                                f"{key}.amax_history should be seeded to 1.0, " f"got {tm.amax_history}"
+                            )
+                        if hasattr(tm, "scale"):
+                            assert torch.isfinite(tm.scale).all(), f"{key}.scale has non-finite values"
+
+    def test_optimizer_step_parity(self, fp8_spec):
+        """Optimizer step counters match baseline after same number of real steps."""
+        seed = 42
+        train_steps = 10
+
+        baseline_model = _make_model(fp8_spec, seed)
+        warmup_model = _make_model(fp8_spec, seed)
+
+        baseline_opt = torch.optim.Adam(baseline_model.parameters(), lr=1e-3, betas=(0.9, 0.999))
+        warmup_opt = torch.optim.Adam(warmup_model.parameters(), lr=1e-3, betas=(0.9, 0.999))
+
+        baseline_reg = _make_registry(baseline_model) if fp8_spec == "local" else None
+        warmup_reg = _make_registry(warmup_model) if fp8_spec == "local" else None
+
+        _run_warmup(warmup_model, warmup_opt, fp8_spec, registry=warmup_reg)
+
+        data = _generate_fixed_data(train_steps, seed=77)
+
+        _train_loop(baseline_model, baseline_opt, data, fp8_spec, registry=baseline_reg)
+        _train_loop(warmup_model, warmup_opt, data, fp8_spec, registry=warmup_reg)
+
+        for (p_b, state_b), (p_w, state_w) in zip(baseline_opt.state.items(), warmup_opt.state.items()):
+            step_b = state_b["step"]
+            step_w = state_w["step"]
+            if isinstance(step_b, torch.Tensor):
+                step_b = step_b.item()
+            if isinstance(step_w, torch.Tensor):
+                step_w = step_w.item()
+            assert step_b == step_w == train_steps, (
+                f"Step count mismatch: baseline={step_b}, warmup={step_w}, " f"expected={train_steps}"
+            )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py b/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py
new file mode 100644
index 000000000..8905fb758
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py
@@ -0,0 +1,179 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the CudaPrefetchIterator-cache invalidation that the MLPerf
+warmup hook performs in its epilogue.
+
+Background
+----------
+``patch_grad_zero_and_data_prefetch`` (priority 41) lazily wraps the current
+``data_iterator`` in a ``CudaPrefetchIterator`` on its first invocation and
+caches the result in a closure-local ``_prefetch_state["iter"]``.  The
+MLPerf one-shot warmup hook (priority 95) drives the first invocation of
+that inner ``_patched_train_step`` with ``data_iterator = synthetic_iter``
+for the synthetic warmup steps, which means the cached prefetch wrapper
+ends up bound to the synthetic iterator -- and ``MegatronDataloaderWrapper``
+is cyclic (never raises ``StopIteration``).
+
+Without invalidation, every subsequent real training step's
+``_synced_prefetch_fwd_bwd`` substitutes the cached
+``CudaPrefetchIterator(synthetic_iter)`` for the incoming real
+``data_iterator`` argument, so the model trains forever on the synthetic
+mock dataset and val_loss on real data stays stuck near 1.38.
+
+The fix exposes a ``_PREFETCH_HANDLE`` plus a ``reset_prefetch_state``
+helper on ``delayed_fp8_scaling_patches`` and calls it from the warmup
+epilogue.  These tests cover the helper contract, the closure-handle
+binding done at patch-installation time, and the integration semantics
+that "after warmup epilogue, the next train_step rebuilds the prefetcher".
+
+Level 1: Pure Python, no GPU.
+
+Run:
+    python -m pytest \
+      tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py -v
+"""
+
+import pytest
+
+from primus.backends.megatron.patches import delayed_fp8_scaling_patches as dfp
+
+
+@pytest.fixture(autouse=True)
+def _isolate_prefetch_handle():
+    """Snapshot+restore the module-level ``_PREFETCH_HANDLE`` so tests can
+    mutate it without bleeding into each other or into a real patch install.
+    """
+    saved = dict(dfp._PREFETCH_HANDLE)
+    yield
+    dfp._PREFETCH_HANDLE.clear()
+    dfp._PREFETCH_HANDLE.update(saved)
+
+
+# ---------------------------------------------------------------------------
+# Helper contract
+# ---------------------------------------------------------------------------
+
+
+class _FakePrefetchIter:
+    """Stand-in for ``CudaPrefetchIterator`` -- the helper only inspects the
+    type name, never the behaviour."""
+
+    def __next__(self):
+        raise StopIteration
+
+
+def test_get_prefetch_state_none_when_unbound():
+    """Before any patch installation, the handle points at None."""
+    dfp._PREFETCH_HANDLE["state"] = None
+    assert dfp.get_prefetch_state() is None
+
+
+def test_get_prefetch_state_returns_closure_dict_when_bound():
+    """``patch_grad_zero_and_data_prefetch`` binds its closure-local
+    ``_prefetch_state`` here at install time; the getter must return it."""
+    bound = {"iter": _FakePrefetchIter()}
+    dfp._PREFETCH_HANDLE["state"] = bound
+    assert dfp.get_prefetch_state() is bound
+
+
+def test_reset_prefetch_state_noop_when_unbound():
+    """No patch installed -> reset is a no-op returning None."""
+    dfp._PREFETCH_HANDLE["state"] = None
+    assert dfp.reset_prefetch_state() is None
+
+
+def test_reset_prefetch_state_noop_when_bound_but_empty():
+    """Patch installed but no iterator yet cached -> nothing to evict."""
+    bound = {}
+    dfp._PREFETCH_HANDLE["state"] = bound
+    assert dfp.reset_prefetch_state() is None
+    assert bound == {}  # state dict untouched
+
+
+def test_reset_prefetch_state_evicts_cached_iter():
+    """The cached iterator is returned AND popped from the state dict, so
+    the closure's ``if "iter" not in _prefetch_state`` guard will rebuild
+    on the next train_step."""
+    fake_iter = _FakePrefetchIter()
+    bound = {"iter": fake_iter}
+    dfp._PREFETCH_HANDLE["state"] = bound
+
+    evicted = dfp.reset_prefetch_state()
+
+    assert evicted is fake_iter
+    assert "iter" not in bound, "iter key must be removed after reset"
+
+
+def test_reset_prefetch_state_idempotent():
+    """Calling reset twice doesn't crash and the second call returns None."""
+    bound = {"iter": _FakePrefetchIter()}
+    dfp._PREFETCH_HANDLE["state"] = bound
+
+    first = dfp.reset_prefetch_state()
+    second = dfp.reset_prefetch_state()
+
+    assert first is not None
+    assert second is None
+
+
+def test_reset_prefetch_state_preserves_other_keys():
+    """Reset must only pop the ``iter`` key; any other state (e.g. metadata
+    a future patch revision might add) should survive."""
+    fake_iter = _FakePrefetchIter()
+    bound = {"iter": fake_iter, "other_metadata": 42}
+    dfp._PREFETCH_HANDLE["state"] = bound
+
+    evicted = dfp.reset_prefetch_state()
+
+    assert evicted is fake_iter
+    assert bound == {"other_metadata": 42}
+
+
+# ---------------------------------------------------------------------------
+# Integration semantics: warmup epilogue + next-train_step rebuild
+# ---------------------------------------------------------------------------
+
+
+def test_warmup_epilogue_makes_next_step_rebuild_prefetcher():
+    """End-to-end behavioural test of the contract we rely on.
+
+    Simulates:
+      1. Patch installation publishes a closure dict via ``_PREFETCH_HANDLE``.
+      2. Warmup step 1 lazily caches a prefetch wrapper around the SYNTHETIC
+         iterator.
+      3. Warmup epilogue calls ``reset_prefetch_state()``.
+      4. The next ``_patched_train_step`` call sees ``"iter" not in
+         _prefetch_state`` and rebuilds the wrapper around the REAL
+         iterator.
+
+    What we assert: step (4) sees an empty / iter-free state dict, which is
+    the precondition for the closure's lazy-rebuild guard to fire.
+    """
+
+    # Step 1: patch install publishes the closure dict.
+    closure_state: dict = {}
+    dfp._PREFETCH_HANDLE["state"] = closure_state
+
+    # Step 2: warmup step 1 caches the prefetcher around synthetic_iter.
+    synthetic_iter_marker = object()
+    closure_state["iter"] = synthetic_iter_marker
+    assert dfp.get_prefetch_state() is closure_state
+    assert dfp.get_prefetch_state()["iter"] is synthetic_iter_marker
+
+    # Step 3: warmup epilogue invalidates.
+    evicted = dfp.reset_prefetch_state()
+    assert evicted is synthetic_iter_marker
+
+    # Step 4: next train_step's lazy-build guard now fires -- i.e. "iter"
+    # is not in the state dict, so the closure will construct a fresh
+    # CudaPrefetchIterator around the real data_iterator argument.
+    assert "iter" not in closure_state, (
+        "After warmup epilogue invalidates the cache, the closure's "
+        "_prefetch_state['iter'] must be absent so the next train_step "
+        "rebuilds the prefetch wrapper around the real data_iterator."
+    )
diff --git a/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py b/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py
index 43d8032a1..d5a928896 100644
--- a/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py
+++ b/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py
@@ -29,13 +29,26 @@ def adapter():
 
 @pytest.fixture
 def sys_modules_guard(adapter):
-    """Drop any stub modules a test injected, keeping tests isolated."""
+    """Isolate stub-eligible modules for the duration of a test.
+
+    Drops any relevant modules BOTH before and after the test. The
+    before-drop matters because some entries (e.g. ``megatron.energon``) are
+    real, top-level packages that may already be installed and cached in
+    ``sys.modules`` from elsewhere in the session. If left in place, the stub
+    installer's ``if pkg_name in sys.modules: continue`` fast-path would skip
+    them, bypassing a test's forced-missing ``import_module`` mock and breaking
+    assertions that expect every package to be stubbed.
+    """
     prefixes = ("modelopt",) + adapter._BRIDGE_OPTIONAL_PACKAGES
-    before = set(sys.modules)
+
+    def _drop():
+        for key in list(sys.modules):
+            if any(key == p or key.startswith(p + ".") for p in prefixes):
+                sys.modules.pop(key, None)
+
+    _drop()
     yield
-    for key in set(sys.modules) - before:
-        if any(key == p or key.startswith(p + ".") for p in prefixes):
-            sys.modules.pop(key, None)
+    _drop()
 
 
 # ---------------------------------------------------------------------------
@@ -118,29 +131,34 @@ def test_install_bridge_optional_stubs_is_idempotent(adapter, monkeypatch, sys_m
 # ---------------------------------------------------------------------------
 # _install_transformers_stub (transformers must be importable)
 # ---------------------------------------------------------------------------
-def test_install_transformers_stub_adds_placeholders_for_missing(adapter):
+def test_install_transformers_stub_adds_placeholders_for_missing(adapter, monkeypatch):
     transformers = pytest.importorskip("transformers")
 
-    name, _used_by = adapter._TRANSFORMERS_PLACEHOLDER_CLASSES[0]
-    had_attr = hasattr(transformers, name)
-    original = getattr(transformers, name, None)
-    if had_attr:
-        delattr(transformers, name)
+    # Append a synthetic always-missing class to the real list to drive the
+    # "missing class" branch deterministically. Deleting a real class won't
+    # work: transformers' lazy module re-caches it before the installer's
+    # hasattr check, and newer transformers ship every class anyway.
+    fake_name = "_PrimusFakeMissingTransformersClass_ForUnitTest"
+    monkeypatch.setattr(
+        adapter,
+        "_TRANSFORMERS_PLACEHOLDER_CLASSES",
+        adapter._TRANSFORMERS_PLACEHOLDER_CLASSES + ((fake_name, "models/fake/fake_bridge.py"),),
+    )
+    assert not hasattr(transformers, fake_name)
     try:
         stubbed = adapter._install_transformers_stub()
-        assert name in stubbed
-        placeholder = getattr(transformers, name)
+        assert fake_name in stubbed
+        placeholder = getattr(transformers, fake_name)
         assert getattr(placeholder, "_primus_placeholder", False) is True
         with pytest.raises(RuntimeError, match="Primus placeholder"):
             placeholder()
     finally:
-        # restore: remove every placeholder we (or the call) injected
+        # Remove every placeholder the call injected: our synthetic entry plus
+        # any real classes that were genuinely missing in this environment.
         for cname, _ in adapter._TRANSFORMERS_PLACEHOLDER_CLASSES:
             obj = getattr(transformers, cname, None)
             if getattr(obj, "_primus_placeholder", False):
                 delattr(transformers, cname)
-        if had_attr:
-            setattr(transformers, name, original)
 
 
 def test_install_transformers_stub_skips_existing_classes(adapter):
diff --git a/tests/unit_tests/backends/nemo_automodel/test_automodel_data_registry.py b/tests/unit_tests/backends/nemo_automodel/test_automodel_data_registry.py
new file mode 100644
index 000000000..a11961f4a
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_automodel_data_registry.py
@@ -0,0 +1,90 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the AutoModel offline cache-builder registry.
+
+WHY THIS TEST EXISTS:
+  The registry maps a model name to a ``"<module>:<callable>"`` string, resolved by
+  import only when ``primus data automodel-cache --model <name>`` actually runs. That
+  laziness is deliberate — it keeps the CLI from importing torch, diffusers and every
+  model's optional dependencies just to print ``--help``.
+
+  The cost of laziness is that a typo in a registry string, or a builder that gets
+  renamed or moved, is invisible until someone runs a cache build — which is a
+  GPU-bound job needing gated encoder weights, i.e. the most expensive possible place
+  to discover it. These tests resolve every registered entry so that breakage surfaces
+  here instead.
+
+  ``build_cache`` itself is not called: it needs a GPU and the encoder weights.
+"""
+
+import inspect
+
+import pytest
+
+from primus.backends.nemo_automodel.data.registry import (
+    CACHE_BUILDERS,
+    available_models,
+    get_cache_builder,
+)
+
+
+def test_registry_is_not_empty():
+    assert CACHE_BUILDERS, "no cache builders registered"
+
+
+def test_available_models_is_sorted_and_matches_registry():
+    assert available_models() == sorted(CACHE_BUILDERS)
+
+
+def test_ideogram4_is_registered():
+    assert "ideogram4" in available_models()
+
+
+@pytest.mark.parametrize("model", sorted(CACHE_BUILDERS))
+def test_every_registered_builder_resolves_and_is_callable(model):
+    """The whole point: a stale registry string must fail here, not on a GPU node."""
+    builder = get_cache_builder(model)
+    assert callable(builder), f"{model} builder is not callable"
+
+
+@pytest.mark.parametrize("model", sorted(CACHE_BUILDERS))
+def test_every_builder_accepts_the_arguments_the_cli_passes(model):
+    """The CLI calls builders purely by keyword, so the names are the contract."""
+    builder = get_cache_builder(model)
+    parameters = inspect.signature(builder).parameters
+    accepts_extra_kwargs = any(
+        p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
+    )
+    required = {
+        "image_dir",
+        "caption_dir",
+        "output_dir",
+        "num_samples",
+        "resolution",
+        "max_text_tokens",
+        "vae_source",
+        "text_encoder_source",
+        "tokenizer_source",
+        "device",
+        "dtype",
+        "seed",
+        "shuffle",
+    }
+    if not accepts_extra_kwargs:
+        missing = required - set(parameters)
+        assert not missing, f"{model} builder cannot accept CLI arguments: {sorted(missing)}"
+
+
+def test_unknown_model_names_the_alternatives():
+    with pytest.raises(ValueError) as excinfo:
+        get_cache_builder("no-such-model")
+    message = str(excinfo.value)
+    assert "no-such-model" in message
+    # An unknown --model should tell the user what they could have typed.
+    for model in available_models():
+        assert model in message
diff --git a/tests/unit_tests/backends/nemo_automodel/test_fsdp2_reshard.py b/tests/unit_tests/backends/nemo_automodel/test_fsdp2_reshard.py
new file mode 100644
index 000000000..6038e1265
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_fsdp2_reshard.py
@@ -0,0 +1,136 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the FSDP2 ``reshard_after_forward`` repair hook.
+
+WHY THIS TEST EXISTS:
+  The original bug was not a wrong value, it was a value that was never passed. Nothing
+  asserted the argument reached its destination, so the setting silently did nothing and
+  training ran at ZeRO-3 communication volume while the config said otherwise. Nothing
+  raised, and the only symptom was throughput.
+
+  That failure mode is invisible to any test that checks behaviour "works". These tests
+  assert the *plumbing* instead: the key is re-applied after the whitelist drops it, the
+  DDP path is untouched, re-installing does not nest wrappers, and a missing target
+  raises rather than passing silently.
+
+  No GPU is needed. The Automodel module is stubbed, so this also does not depend on the
+  submodule being importable.
+"""
+
+import sys
+import types
+
+import pytest
+
+from primus.backends.nemo_automodel.distributed import fsdp2_reshard
+
+MODULE_PATH = "nemo_automodel._diffusers.auto_diffusion_pipeline"
+
+
+class FakeManager:
+    """Stands in for FSDP2Manager/DDPManager, which set this in __init__ from config."""
+
+    def __init__(self):
+        self.reshard_after_forward = "NEVER_SET"
+
+
+@pytest.fixture
+def stub_pipeline(monkeypatch):
+    """Install a stub ``auto_diffusion_pipeline`` whose factory records its input.
+
+    The real ``_create_parallel_manager`` copies ``manager_args`` and pops
+    ``_manager_type`` from the copy; the stub does the same so the wrapper is exercised
+    against the same contract it relies on.
+    """
+    module = types.ModuleType(MODULE_PATH)
+    calls = []
+
+    def _create_parallel_manager(manager_args):
+        args = manager_args.copy()
+        args.pop("_manager_type", "fsdp2")
+        calls.append(args)
+        return FakeManager()
+
+    module._create_parallel_manager = _create_parallel_manager
+    module.calls = calls
+
+    for name in ("nemo_automodel", "nemo_automodel._diffusers", MODULE_PATH):
+        monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, MODULE_PATH, module)
+
+    # The hook records install state in module globals; reset so tests do not leak.
+    monkeypatch.setattr(fsdp2_reshard, "patch_installed", False, raising=False)
+    monkeypatch.setattr(fsdp2_reshard, "applied_reshard_after_forward", None, raising=False)
+    return module
+
+
+def test_install_reports_success_and_sets_provenance(stub_pipeline):
+    assert fsdp2_reshard.install() is True
+    assert fsdp2_reshard.patch_installed is True
+
+
+@pytest.mark.parametrize("value", [False, True, None])
+def test_key_is_reapplied_onto_the_manager_for_fsdp2(stub_pipeline, value):
+    """The whole point: the whitelist drops this, so the wrapper must put it back."""
+    fsdp2_reshard.install()
+    manager = stub_pipeline._create_parallel_manager(
+        {"_manager_type": "fsdp2", "reshard_after_forward": value, "world_size": 8}
+    )
+    assert manager.reshard_after_forward == value
+    assert fsdp2_reshard.applied_reshard_after_forward == value
+
+
+def test_fsdp2_is_the_default_manager_type(stub_pipeline):
+    """A missing _manager_type means fsdp2 upstream, so the wrapper must still apply."""
+    fsdp2_reshard.install()
+    manager = stub_pipeline._create_parallel_manager({"reshard_after_forward": False})
+    assert manager.reshard_after_forward is False
+
+
+def test_ddp_path_is_left_untouched(stub_pipeline):
+    fsdp2_reshard.install()
+    manager = stub_pipeline._create_parallel_manager(
+        {"_manager_type": "ddp", "reshard_after_forward": False}
+    )
+    assert manager.reshard_after_forward == "NEVER_SET"
+
+
+def test_absent_key_is_not_invented(stub_pipeline):
+    """It repairs, it does not decide. No key in YAML means None, i.e. today's heuristic."""
+    fsdp2_reshard.install()
+    manager = stub_pipeline._create_parallel_manager({"_manager_type": "fsdp2"})
+    assert manager.reshard_after_forward is None
+
+
+def test_the_callee_still_receives_its_arguments(stub_pipeline):
+    """The wrapper must not consume or mutate what the real factory needs."""
+    fsdp2_reshard.install()
+    stub_pipeline._create_parallel_manager(
+        {"_manager_type": "fsdp2", "reshard_after_forward": False, "world_size": 8}
+    )
+    assert stub_pipeline.calls[-1]["world_size"] == 8
+
+
+def test_reinstalling_does_not_nest_wrappers(stub_pipeline):
+    """Without the sentinel, a second install captures the wrapper as its own target."""
+    fsdp2_reshard.install()
+    first = stub_pipeline._create_parallel_manager
+    assert fsdp2_reshard.install() is True
+    assert stub_pipeline._create_parallel_manager is first
+
+
+def test_missing_target_raises_rather_than_silently_skipping(stub_pipeline):
+    del stub_pipeline._create_parallel_manager
+    with pytest.raises(RuntimeError, match="_create_parallel_manager"):
+        fsdp2_reshard.install()
+
+
+def test_non_callable_target_raises(stub_pipeline):
+    stub_pipeline._create_parallel_manager = "not callable"
+    with pytest.raises(RuntimeError, match="_create_parallel_manager"):
+        fsdp2_reshard.install()
diff --git a/tests/unit_tests/backends/nemo_automodel/test_ideogram4_context_parallel.py b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_context_parallel.py
new file mode 100644
index 000000000..d1c9b8ccc
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_context_parallel.py
@@ -0,0 +1,206 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the Ideogram-4 context-parallel plan.
+
+WHY THIS TEST EXISTS:
+  Every failure mode this plan can have is silent. Nothing about a wrong plan raises:
+  the model still runs, the loss still falls, and only the numbers are wrong. So these
+  tests assert the *structure* of the plan rather than any behaviour.
+
+  Three specific regressions are pinned here, all of which were live during development:
+
+  1. **Wrapping ``forward``.** diffusers resolves the root plan's entries against the
+     forward signature, so wrapping ``forward`` with ``*args, **kwargs`` hides
+     ``encoder_hidden_states`` and ``indicator`` and leaves them unsplit. The first
+     implementation did exactly this. ``test_root_plan_names_resolve_in_forward_signature``
+     is the guard.
+  2. **Splitting ``segment_ids``.** The attention mask must stay full length, because
+     after the Ulysses all-to-all each rank holds the whole sequence with a subset of
+     heads. Its absence from the plan is load-bearing, so it is asserted explicitly.
+  3. **CP combined with the var-len flash processor.** That processor bypasses
+     ``dispatch_attention_fn``, so the all-to-all never runs and each rank attends only
+     within its own shard -- without raising. The guard must raise instead.
+
+  No GPU and no distributed init is needed; the plan is a class attribute and the model
+  is instantiated tiny, on CPU, purely to resolve module names.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("diffusers")
+
+from diffusers.models._modeling_parallel import (  # noqa: E402
+    ContextParallelInput,
+    ContextParallelOutput,
+)
+from diffusers.models.transformers.transformer_ideogram4 import (  # noqa: E402
+    Ideogram4Transformer2DModel,
+)
+
+from primus.backends.nemo_automodel.models.ideogram4 import (  # noqa: E402
+    context_parallel as cp,
+)
+
+# Small but structurally faithful: 3-axis MRoPE, heads divisible by a CP degree of 2.
+TINY_MODEL_KWARGS = dict(
+    in_channels=128,
+    num_layers=1,
+    attention_head_dim=32,
+    num_attention_heads=4,
+    intermediate_size=64,
+    adaln_dim=32,
+    llm_features_dim=16,
+    mrope_section=(4, 2, 2),
+)
+
+
+@pytest.fixture(autouse=True)
+def installed_plan():
+    """Install onto the real class, then restore it.
+
+    ``install`` mutates class-level state, so without this a failure in one test would
+    leak into the next and the suite would stop meaning anything.
+    """
+    model_cls = Ideogram4Transformer2DModel
+    saved = {
+        "_cp_plan": getattr(model_cls, "_cp_plan", None),
+        "enable_parallelism": model_cls.enable_parallelism,
+        "forward": model_cls.forward,
+        "_primus_cp_installed": getattr(model_cls, "_primus_cp_installed", False),
+    }
+    model_cls._primus_cp_installed = False
+    cp.install()
+    yield model_cls
+    for name, value in saved.items():
+        if name == "_cp_plan" and value is None:
+            model_cls._cp_plan = None
+        else:
+            setattr(model_cls, name, value)
+
+
+@pytest.fixture
+def tiny_model():
+    return Ideogram4Transformer2DModel(**TINY_MODEL_KWARGS)
+
+
+def test_install_satisfies_the_upstream_eligibility_gate(installed_plan):
+    """Verbatim the check AutoModel makes before it will enable CP."""
+    assert getattr(installed_plan, "_cp_plan", None) is not None
+    assert hasattr(installed_plan, "enable_parallelism")
+
+
+def test_install_is_idempotent(installed_plan):
+    assert cp.install() is False
+
+
+def test_root_plan_splits_exactly_the_per_token_inputs(installed_plan):
+    """Ideogram-4 is single-stream, so these three split together -- and no others.
+
+    The root forward derives the indicator masks and applies them to the other two
+    before the first block, so a subset would mismatch shapes.
+    """
+    assert set(installed_plan._cp_plan[""]) == {
+        "hidden_states",
+        "encoder_hidden_states",
+        "indicator",
+    }
+
+
+@pytest.mark.parametrize("name", ["segment_ids", "position_ids", "timestep"])
+def test_load_bearing_omissions_stay_out_of_the_plan(installed_plan, name):
+    """These absences are decisions, not oversights.
+
+    ``segment_ids`` must stay whole so the mask is full length after the all-to-all;
+    ``position_ids`` stays whole because MRoPE's *output* is split instead; ``timestep``
+    is per-sample and broadcasts.
+    """
+    assert name not in installed_plan._cp_plan[""]
+
+
+def test_root_plan_names_resolve_in_forward_signature(installed_plan):
+    """Regression: never wrap ``forward`` on a CP-planned model.
+
+    diffusers matches the root plan's keys against the forward signature, so a
+    ``*args, **kwargs`` wrapper silently leaves those inputs unsplit.
+    """
+    import inspect
+
+    params = set(inspect.signature(installed_plan.forward).parameters)
+    missing = set(installed_plan._cp_plan[""]) - params
+    assert not missing, f"plan names not visible in forward signature: {sorted(missing)}"
+
+
+def test_mrope_outputs_are_split_not_its_inputs(installed_plan):
+    """RoPE is applied to q/k before the all-to-all, so cos/sin must end up local."""
+    rope_plan = installed_plan._cp_plan["rotary_emb"]
+    assert set(rope_plan) == {0, 1}, "MRoPE returns (cos, sin); both outputs must be split"
+    for entry in rope_plan.values():
+        assert isinstance(entry, ContextParallelInput)
+        assert entry.split_output is True
+        assert entry.split_dim == 1
+
+
+def test_output_is_gathered_at_final_layer(installed_plan):
+    entry = installed_plan._cp_plan["final_layer"]
+    assert isinstance(entry, ContextParallelOutput)
+    assert entry.gather_dim == 1
+
+
+def test_plan_module_names_resolve_on_a_real_model(installed_plan, tiny_model):
+    """A typo in a module key would make diffusers attach the hook nowhere, silently."""
+    names = dict(tiny_model.named_modules())
+    for key in installed_plan._cp_plan:
+        if key == "":
+            continue
+        assert key in names, f"plan key {key!r} matches no submodule"
+
+
+def test_mrope_output_rank_matches_expected_dims(installed_plan, tiny_model):
+    """``expected_dims`` is the only shape assertion diffusers makes; keep it honest."""
+    cos, sin = tiny_model.rotary_emb(torch.zeros(2, 8, 3, dtype=torch.long))
+    expected = installed_plan._cp_plan["rotary_emb"][0].expected_dims
+    assert cos.dim() == expected
+    assert sin.dim() == expected
+
+
+def test_cp_refuses_the_varlen_processor(installed_plan, tiny_model, monkeypatch):
+    """CP + var-len is silent corruption, so enabling CP must raise instead."""
+    monkeypatch.setenv("PRIMUS_IDEOGRAM_VARLEN_ATTN", "1")
+    with pytest.raises(ValueError, match="PRIMUS_IDEOGRAM_VARLEN_ATTN"):
+        tiny_model.enable_parallelism(config=None)
+
+
+def test_cp_allows_the_stock_attention_path(installed_plan, tiny_model, monkeypatch):
+    """The guard must not fire on the default configuration.
+
+    ``config=None`` gets past our check and then fails inside diffusers; anything other
+    than our own ValueError means the guard let it through, which is what matters here.
+    """
+    monkeypatch.delenv("PRIMUS_IDEOGRAM_VARLEN_ATTN", raising=False)
+    try:
+        tiny_model.enable_parallelism(config=None)
+    except ValueError as exc:
+        assert "PRIMUS_IDEOGRAM_VARLEN_ATTN" not in str(exc)
+    except Exception:
+        pass
+
+
+@pytest.mark.parametrize(
+    "timestep, should_raise",
+    [
+        (torch.rand(2), False),  # per-sample: broadcasts, needs no split
+        (torch.rand(2, 24), True),  # per-token: would need splitting, so refuse
+    ],
+)
+def test_timestep_shape_assumption_is_enforced(timestep, should_raise):
+    if should_raise:
+        with pytest.raises(ValueError, match="per-sample timestep"):
+            cp._assert_timestep_is_per_sample(timestep)
+    else:
+        cp._assert_timestep_is_per_sample(timestep)
diff --git a/tests/unit_tests/backends/nemo_automodel/test_ideogram4_cu_seqlens.py b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_cu_seqlens.py
new file mode 100644
index 000000000..c44d9dfb3
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_cu_seqlens.py
@@ -0,0 +1,217 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the Ideogram-4 host-side var-len packing (``build_cu_seqlens``).
+
+WHY THIS TEST EXISTS:
+  There are now TWO producers of the same ``cu_seqlens``:
+
+    1. ``build_cu_seqlens(text_lengths, ...)``           -- host-side, from Python ints
+    2. ``blockdiag_bool_mask_to_cu_seqlens(mask)``       -- device-side, from mask values
+
+  Producer 1 is what training uses (it keeps the compiled graph break-free); producer 2 is
+  the original, and remains the reference. If they ever disagree, var-len flash attends
+  across what should be separate segments -- padding leaks into real tokens and training
+  corrupts with NO error raised anywhere. Exact equality is therefore the assertion, not
+  closeness.
+
+  The tests also pin the two properties the compiled graph depends on:
+    * the segment count is exactly ``2*B`` for ANY caption lengths, so ``cu_seqlens`` has a
+      fixed shape and does not trigger a recompile per length pattern;
+    * a row with no padding (``offset == 0``) is rejected loudly, since a zero-length
+      segment is not representable for var-len flash.
+
+Only ``torch`` is required for the core parity test. The ``_prepare_ids`` cross-check needs
+``nemo_automodel`` (the adapter class is built lazily against its ``ModelAdapter`` base) and
+skips when the submodule is not importable. Note that it therefore runs in TWO different
+environments -- with diffusers present (in-container) and without (lint/CI) -- which is why the
+padding-segment id is imported from the adapter instead of hardcoded.
+"""
+
+import pytest
+
+try:
+    import torch
+except ImportError:  # pragma: no cover - CPU-less lint environments
+    pytest.skip("Ideogram-4 cu_seqlens tests require torch", allow_module_level=True)
+
+from primus.backends.nemo_automodel.models.ideogram4.adapter import (
+    SEQUENCE_PADDING_INDICATOR,
+    build_cu_seqlens,
+)
+from primus.backends.nemo_automodel.models.ideogram4.attention import (
+    blockdiag_bool_mask_to_cu_seqlens,
+)
+
+# Taken from the adapter rather than hardcoded, because the adapter resolves it from diffusers
+# when diffusers is importable and falls back to a literal otherwise -- and the two disagree
+# (0.39.0 pads with -1). The VALUE is irrelevant to everything here: the model's mask is
+# ``seg_i == seg_j``, so only which positions share an id matters, and the adapter writes the
+# valid region as a literal 1.
+VALID_SEGMENT = 1
+
+
+def _segment_ids(text_lengths, max_text_tokens, num_image_tokens):
+    """``[left-pad][text][image]`` segment ids -- the layout ``_prepare_ids`` produces.
+
+    ``PADDING(0)`` over the left-pad region, ``1`` over the whole ``[text][image]`` region,
+    which is what makes the model's ``(seg_i == seg_j)`` mask block-diagonal per row.
+    """
+    total = max_text_tokens + num_image_tokens
+    seg = torch.full((len(text_lengths), total), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
+    for b, num_text in enumerate(text_lengths):
+        offset = max_text_tokens - int(num_text)
+        seg[b, offset:] = VALID_SEGMENT
+    return seg
+
+
+def _bool_mask(segment_ids):
+    """The dense ``(B,1,L,L)`` mask the diffusers block materializes from segment ids."""
+    return (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
+
+
+# (label, text_lengths, max_text_tokens, num_image_tokens)
+# max_text_tokens already INCLUDES the adapter's reserved pad column, so every case leaves
+# at least one pad token per row.
+LENGTH_PATTERNS = [
+    ("min_padding_one_slot", [4, 4], 5, 3),
+    ("maximally_ragged", [1, 2, 3, 4], 5, 3),
+    ("all_equal", [3, 3, 3, 3], 5, 3),
+    ("single_row", [2], 5, 7),
+    ("mixed_min_and_slack", [4, 1, 4, 2], 5, 3),
+    ("wide_image_region", [7, 13, 2], 16, 64),
+    ("batch_of_eight", [1, 3, 5, 7, 9, 11, 13, 15], 16, 32),
+]
+
+
+class TestBuildCuSeqlensParity:
+    """``build_cu_seqlens`` must agree EXACTLY with the mask-derived reference."""
+
+    @pytest.mark.parametrize(
+        "text_lengths,max_text_tokens,num_image_tokens",
+        [p[1:] for p in LENGTH_PATTERNS],
+        ids=[p[0] for p in LENGTH_PATTERNS],
+    )
+    def test_matches_mask_derived_reference(self, text_lengths, max_text_tokens, num_image_tokens):
+        seg = _segment_ids(text_lengths, max_text_tokens, num_image_tokens)
+        expected, expected_max_seqlen, is_trivial = blockdiag_bool_mask_to_cu_seqlens(_bool_mask(seg))
+
+        actual = build_cu_seqlens(text_lengths, max_text_tokens, num_image_tokens)
+
+        assert not is_trivial, (
+            "every row should have padding in these fixtures, so the mask must NOT be "
+            "trivial -- if it is, the fixture no longer exercises the var-len path"
+        )
+        assert actual.dtype == torch.int32, "flash_attn_varlen_func requires int32 cu_seqlens"
+        assert torch.equal(actual, expected), (
+            "host-built cu_seqlens disagrees with the mask-derived reference for "
+            f"text_lengths={text_lengths}, max_text_tokens={max_text_tokens}: "
+            f"got {actual.tolist()} vs {expected.tolist()}. A mismatch here means var-len "
+            "flash attends across segment boundaries and training corrupts silently. Check "
+            "that build_cu_seqlens' offset arithmetic still mirrors _prepare_ids."
+        )
+        # The static bound the adapter passes must never under-estimate the real longest
+        # segment, or flash gets a max_seqlen smaller than a sequence it has to schedule.
+        assert max_text_tokens + num_image_tokens >= expected_max_seqlen
+
+    @pytest.mark.parametrize(
+        "text_lengths,max_text_tokens,num_image_tokens",
+        [p[1:] for p in LENGTH_PATTERNS],
+        ids=[p[0] for p in LENGTH_PATTERNS],
+    )
+    def test_shape_is_two_segments_per_row(self, text_lengths, max_text_tokens, num_image_tokens):
+        """Fixed ``2*B+1`` length is what keeps the compiled graph from recompiling."""
+        cu = build_cu_seqlens(text_lengths, max_text_tokens, num_image_tokens)
+        assert cu.numel() == 2 * len(text_lengths) + 1, (
+            f"expected 2*B+1={2 * len(text_lengths) + 1} entries for B={len(text_lengths)}, got "
+            f"{cu.numel()}. A data-dependent length changes this tensor's shape between "
+            "batches and forces a torch.compile recompile per length pattern."
+        )
+
+    def test_shape_is_independent_of_caption_lengths(self):
+        """Same batch size, wildly different captions -> identical cu_seqlens shape."""
+        shapes = {
+            build_cu_seqlens(lengths, 16, 32).shape
+            for lengths in ([1, 1, 1, 1], [15, 15, 15, 15], [1, 15, 8, 3], [7, 2, 15, 11])
+        }
+        assert len(shapes) == 1, f"cu_seqlens shape varied with caption lengths: {shapes}"
+
+    @pytest.mark.parametrize(
+        "text_lengths,max_text_tokens,num_image_tokens",
+        [p[1:] for p in LENGTH_PATTERNS],
+        ids=[p[0] for p in LENGTH_PATTERNS],
+    )
+    def test_contract_boundaries(self, text_lengths, max_text_tokens, num_image_tokens):
+        """Starts at 0, ends at B*S, non-decreasing, and no empty segments."""
+        seq_len = max_text_tokens + num_image_tokens
+        cu = build_cu_seqlens(text_lengths, max_text_tokens, num_image_tokens)
+
+        assert int(cu[0]) == 0
+        assert int(cu[-1]) == len(text_lengths) * seq_len
+        diffs = cu[1:] - cu[:-1]
+        assert bool((diffs > 0).all()), (
+            f"zero-length or negative segment in {cu.tolist()}; var-len flash cannot "
+            "represent an empty sequence"
+        )
+
+    def test_segment_lengths_match_the_layout(self):
+        """Per row: the pad block is ``offset`` long, the text+image block the remainder."""
+        text_lengths, max_text_tokens, num_image_tokens = [2, 4, 1], 6, 5
+        cu = build_cu_seqlens(text_lengths, max_text_tokens, num_image_tokens)
+        diffs = (cu[1:] - cu[:-1]).tolist()
+
+        expected = []
+        for num_text in text_lengths:
+            offset = max_text_tokens - num_text
+            expected += [offset, num_text + num_image_tokens]
+        assert diffs == expected, f"segment lengths {diffs} do not match the layout {expected}"
+
+    def test_rejects_row_without_padding(self):
+        """``offset == 0`` must fail loudly -- the reserved pad column prevents it."""
+        with pytest.raises(ValueError, match="leaving no padding"):
+            build_cu_seqlens([5, 3], max_text_tokens=5, num_image_tokens=4)
+
+    def test_accepts_torch_scalar_lengths(self):
+        """``text_lengths`` may arrive as tensor elements; ints must be coerced."""
+        as_ints = build_cu_seqlens([2, 3], 5, 4)
+        as_tensors = build_cu_seqlens(list(torch.tensor([2, 3])), 5, 4)
+        assert torch.equal(as_ints, as_tensors)
+
+
+class TestPrepareIdsAgreement:
+    """The adapter's own ``segment_ids`` must match the layout this test assumes."""
+
+    def test_prepare_ids_segment_ids_match_local_layout(self):
+        pytest.importorskip(
+            "nemo_automodel",
+            reason="Ideogram-4 adapter class is built against nemo_automodel's ModelAdapter base",
+        )
+        from primus.backends.nemo_automodel.models.ideogram4.adapter import (
+            get_ideogram4_adapter_class,
+        )
+
+        adapter_cls = get_ideogram4_adapter_class()
+        text_lengths, grid_h, grid_w, max_text_tokens = [2, 4, 1], 2, 3, 6
+        num_image_tokens = grid_h * grid_w
+
+        _, segment_ids, _ = adapter_cls._prepare_ids(
+            text_lengths, grid_h, grid_w, max_text_tokens, torch.device("cpu")
+        )
+
+        assert torch.equal(
+            segment_ids, _segment_ids(text_lengths, max_text_tokens, num_image_tokens)
+        ), "this test file's segment-id layout has drifted from the adapter's _prepare_ids"
+
+        # Closing the loop: the mask the model builds from the adapter's real segment_ids
+        # yields the same packing the adapter hands to the processor.
+        expected, _, _ = blockdiag_bool_mask_to_cu_seqlens(_bool_mask(segment_ids))
+        actual = build_cu_seqlens(text_lengths, max_text_tokens, num_image_tokens)
+        assert torch.equal(actual, expected)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/nemo_automodel/test_ideogram4_packing_buffer.py b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_packing_buffer.py
new file mode 100644
index 000000000..a7ea838d6
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_packing_buffer.py
@@ -0,0 +1,397 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the Ideogram-4 packing transport (``ideogram4_packing_buffer``).
+
+WHY THIS TEST EXISTS:
+  The transport's whole value is an invariant that nothing else checks: **every attention
+  module holds the SAME buffer object**, so the adapter's single ``copy_`` per step is visible
+  to all 34 layers. Every way that invariant can lapse fails silently at runtime -- one layer
+  reads the published packing and the rest read something else, so var-len flash attends
+  across segment boundaries and training corrupts with no error raised anywhere. The two
+  lapses that can actually happen in production are a batch-size change (buffer of the wrong
+  shape) and a meta-device materialization pass (``to_empty()`` hands every module its own
+  tensor), so both are exercised here directly.
+
+  These tests need only ``torch`` -- no diffusers, no nemo_automodel, no GPU. The fake
+  attention modules carry a real ``Ideogram4VarlenAttnProcessor`` because membership in the
+  transport is decided by processor type; constructing one is cheap (aiter is imported lazily
+  inside the attention call, not at class definition).
+"""
+
+import logging
+
+import pytest
+
+try:
+    import torch
+    import torch.nn as nn
+except ImportError:  # pragma: no cover - CPU-less lint environments
+    pytest.skip("Ideogram-4 packing buffer tests require torch", allow_module_level=True)
+
+from primus.backends.nemo_automodel.models.ideogram4 import attention as varlen_mod
+from primus.backends.nemo_automodel.models.ideogram4 import (
+    packing_buffer as transport,
+)
+from primus.backends.nemo_automodel.models.ideogram4.packing_buffer import (
+    BOUND_ATTR,
+    PACKING_ATTR,
+    attention_modules,
+    clear_packing,
+    publish_packing,
+    resolve_packing,
+)
+from primus.backends.nemo_automodel.models.ideogram4.attention import (
+    Ideogram4VarlenAttnProcessor,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_log_dedupe():
+    """The transport logs each condition once per process; tests assert on those logs."""
+    transport._logged.clear()
+    yield
+    transport._logged.clear()
+
+
+class _FakeAttention(nn.Module):
+    """Stands in for ``Ideogram4Attention``: what matters is the processor it carries."""
+
+    def __init__(self, varlen: bool = True):
+        super().__init__()
+        self.to_q = nn.Linear(4, 4, bias=False)
+        self.processor = Ideogram4VarlenAttnProcessor() if varlen else object()
+
+
+class _FakeBlock(nn.Module):
+    def __init__(self, varlen: bool = True):
+        super().__init__()
+        self.attention = _FakeAttention(varlen)
+
+
+class _FakeModel(nn.Module):
+    """``model.layers`` of blocks, mirroring ``Ideogram4Transformer2DModel``'s shape."""
+
+    def __init__(self, layers: int = 4, varlen: bool = True):
+        super().__init__()
+        self.layers = nn.ModuleList(_FakeBlock(varlen) for _ in range(layers))
+
+    @property
+    def attentions(self):
+        return [block.attention for block in self.layers]
+
+
+def _packing(batch: int = 3, seq: int = 8, offset: int = 2) -> torch.Tensor:
+    """A ``2B+1`` int32 packing shaped like ``build_cu_seqlens``' output."""
+    starts = []
+    for b in range(batch):
+        starts += [b * seq, b * seq + offset]
+    starts.append(batch * seq)
+    return torch.tensor(starts, dtype=torch.int32)
+
+
+def _buffers(model: _FakeModel):
+    return [getattr(attn, PACKING_ATTR) for attn in model.attentions]
+
+
+class TestPublishSharesOneObject:
+    """The invariant the design rests on."""
+
+    def test_all_modules_share_one_buffer(self):
+        model = _FakeModel(layers=6)
+        shared = publish_packing(model, _packing())
+
+        buffers = _buffers(model)
+        assert all(buf is shared for buf in buffers), (
+            "attention modules do not share ONE buffer object, so the adapter's single copy_ "
+            "would reach only some layers and the rest would attend on stale data"
+        )
+        assert shared.dtype == torch.int32, "flash_attn_varlen_func requires int32 cu_seqlens"
+
+    def test_one_copy_is_visible_from_every_module(self):
+        model = _FakeModel(layers=4)
+        publish_packing(model, _packing(offset=2))
+
+        second = _packing(offset=5)
+        publish_packing(model, second)
+
+        for attn in model.attentions:
+            assert torch.equal(getattr(attn, PACKING_ATTR), second)
+
+    def test_republish_mutates_in_place_rather_than_rebinding(self):
+        """Rebinding per step is what makes a buffer look like a fresh input to Dynamo."""
+        model = _FakeModel()
+        first = publish_packing(model, _packing(offset=1))
+        again = publish_packing(model, _packing(offset=3))
+        assert again is first, "the buffer object must be stable across steps, only its values move"
+
+    def test_max_seqlen_reaches_every_module(self):
+        model = _FakeModel(layers=3)
+        publish_packing(model, _packing(), max_seqlen=64)
+        assert [getattr(a, BOUND_ATTR) for a in model.attentions] == [64, 64, 64]
+
+    def test_buffer_is_non_persistent(self):
+        """Derived per step: it must never enter a checkpoint."""
+        model = _FakeModel()
+        publish_packing(model, _packing())
+        leaked = [k for k in model.state_dict() if PACKING_ATTR in k]
+        assert not leaked, f"non-persistent buffer leaked into the state dict: {leaked}"
+
+    def test_publishes_to_the_requested_device(self):
+        model = _FakeModel()
+        shared = publish_packing(model, _packing(), device=torch.device("cpu"))
+        assert shared.device == torch.device("cpu")
+
+
+class TestInvariantRepair:
+    """Both production ways the sharing can lapse, and the recovery from each."""
+
+    def test_shape_change_reinstalls_and_stays_shared(self):
+        model = _FakeModel(layers=5)
+        publish_packing(model, _packing(batch=3))
+
+        wider = _packing(batch=8)
+        shared = publish_packing(model, wider)
+
+        assert shared.numel() == wider.numel()
+        assert all(buf is shared for buf in _buffers(model))
+        assert torch.equal(shared, wider)
+
+    def test_shape_change_warns_because_it_forces_a_recompile(self, caplog):
+        model = _FakeModel()
+        publish_packing(model, _packing(batch=3))
+        with caplog.at_level(logging.WARNING):
+            publish_packing(model, _packing(batch=4))
+        assert any("length changed" in r.message for r in caplog.records), (
+            "a cu_seqlens length change silently recompiles the graph every time the batch "
+            "size moves; it has to be visible in the log"
+        )
+
+    def test_materialization_that_unshares_buffers_is_repaired(self):
+        """``to_empty()`` gives every module its own tensor. Measured, not hypothetical."""
+        model = _FakeModel(layers=4)
+        publish_packing(model, _packing())
+
+        # Exactly what Module._apply / to_empty does: a fresh tensor per module.
+        for attn in model.attentions:
+            attn._buffers[PACKING_ATTR] = torch.zeros_like(getattr(attn, PACKING_ATTR))
+        assert len({id(b) for b in _buffers(model)}) == 4, "fixture did not actually unshare"
+
+        wanted = _packing(offset=4)
+        shared = publish_packing(model, wanted)
+
+        assert all(buf is shared for buf in _buffers(model)), "sharing was not restored"
+        for attn in model.attentions:
+            assert torch.equal(getattr(attn, PACKING_ATTR), wanted), (
+                "a module still holds its own buffer, so it would attend on whatever "
+                "materialization left there -- silently"
+            )
+
+    def test_unsharing_is_reported(self, caplog):
+        model = _FakeModel(layers=2)
+        publish_packing(model, _packing())
+        for attn in model.attentions:
+            attn._buffers[PACKING_ATTR] = torch.zeros_like(getattr(attn, PACKING_ATTR))
+        with caplog.at_level(logging.WARNING):
+            publish_packing(model, _packing())
+        assert any("no longer shared" in r.message for r in caplog.records)
+
+
+class TestSelfGating:
+    """No var-len processor means no state to maintain."""
+
+    def test_model_without_varlen_processors_gets_nothing(self):
+        model = _FakeModel(varlen=False)
+        assert attention_modules(model) == []
+        assert publish_packing(model, _packing()) is None
+        for attn in model.attentions:
+            assert not hasattr(attn, PACKING_ATTR)
+
+    def test_required_publish_raises_when_nothing_can_read_it(self):
+        """The adapter publishes with ``required=True``, and this is why.
+
+        Having built a packing, a model that cannot read it is a misconfiguration: every layer
+        falls back to deriving its own from the mask, and if that happens on a SUBSET of ranks,
+        data parallelism averages gradients from two different attention paths. No log would
+        show it -- Primus quiets non-zero ranks once training starts -- so the only place this
+        can surface is an exception.
+        """
+        model = _FakeModel(varlen=False)
+        with pytest.raises(RuntimeError, match="no attention module can read it"):
+            publish_packing(model, _packing(), required=True)
+
+    def test_consumers_are_discovered_and_cached(self):
+        model = _FakeModel(layers=7)
+        assert len(attention_modules(model)) == 7
+        assert attention_modules(model) is attention_modules(model), "consumer list should be cached"
+
+
+class TestResolvePacking:
+    """What the processor calls on every attention invocation."""
+
+    def test_reads_the_buffer_off_the_module(self):
+        model = _FakeModel()
+        published = publish_packing(model, _packing(), max_seqlen=32)
+        cu, bound = resolve_packing(model.attentions[0])
+        assert cu is published and bound == 32
+
+    def test_explicit_argument_wins(self):
+        """Keeps a future kwargs route (patching the block forward) able to override."""
+        model = _FakeModel()
+        publish_packing(model, _packing(offset=2), max_seqlen=32)
+        override = _packing(offset=6)
+        cu, bound = resolve_packing(model.attentions[0], override, 99)
+        assert cu is override and bound == 99
+
+    def test_no_buffer_means_no_packing(self):
+        """The processor must then fall through to its mask-derived legacy path."""
+        model = _FakeModel()
+        cu, bound = resolve_packing(model.attentions[0])
+        assert cu is None and bound is None
+
+    def test_clear_packing_restores_the_legacy_path(self):
+        model = _FakeModel(layers=3)
+        publish_packing(model, _packing())
+        assert clear_packing(model) == 3
+        cu, _ = resolve_packing(model.attentions[0])
+        assert cu is None
+
+
+class TestStalePackingGuard:
+    """A packing published for a different batch size must fail loudly, not attend wrongly.
+
+    The check is on ``numel()`` alone -- static shape metadata -- so it costs a guard and no
+    host sync inside the compiled region, and it runs before any aiter kernel is touched
+    (which is why this test needs no GPU).
+    """
+
+    @staticmethod
+    def _qkv(batch: int, seq: int = 8, heads: int = 2, dim: int = 4):
+        shape = (batch, seq, heads, dim)
+        return torch.zeros(shape), torch.zeros(shape), torch.zeros(shape)
+
+    def test_rejects_a_packing_from_a_different_batch_size(self):
+        proc = Ideogram4VarlenAttnProcessor()
+        q, k, v = self._qkv(batch=2)
+        stale = _packing(batch=5)  # published when B was 5
+        with pytest.raises(ValueError, match="2\\*B\\+1"):
+            proc._attention(q, k, v, None, stale, 8)
+
+    def test_passes_a_matching_packing_through_to_the_kernel(self, monkeypatch):
+        """The guard must not stand in the way of the path it protects.
+
+        Stubs the aiter call so this runs on CPU, and checks the processor hands the kernel the
+        packing it was given -- plus the STATIC ``max_seqlen``, not a value derived from data.
+        """
+        seen = {}
+
+        def _fake_varlen(q, k, v, cu_seqlens, max_seqlen, deterministic=False):
+            seen["cu_seqlens"] = cu_seqlens
+            seen["max_seqlen"] = max_seqlen
+            return torch.zeros_like(q)
+
+        monkeypatch.setattr(varlen_mod, "varlen_flash_attention", _fake_varlen)
+
+        proc = Ideogram4VarlenAttnProcessor()
+        q, k, v = self._qkv(batch=3, seq=8, heads=2, dim=4)
+        packing = _packing(batch=3, seq=8)
+
+        out = proc._attention(q, k, v, None, packing, 8)
+
+        assert out.shape == (3, 8, 2, 4), "the packed (B*L,H,D) result must be unpacked again"
+        assert torch.equal(seen["cu_seqlens"], packing)
+        assert seen["max_seqlen"] == 8
+
+
+class TestKernelNeverSeesTheSharedBuffer:
+    """The kernel gets a private copy, because aiter's var-len op WRITES its cu_seqlens.
+
+    aiter saves the tensor for its backward and then mutates it, so every call bumps the
+    autograd version counter. Handing 34 layers the same buffer moved its version 34 times per
+    forward while each layer's backward still expected the version it saved, and the step died
+    in ``.backward()`` ("IntTensor[5] is at version 35; expected version 34" -- measured on the
+    first on-model compile run, 2026-08-04). The fake kernel below mutates its argument the way
+    aiter does, so this test fails if the clone is ever dropped.
+    """
+
+    @staticmethod
+    def _kernel(seen, monkeypatch):
+        def _fake_varlen(q, k, v, cu_seqlens, max_seqlen, deterministic=False):
+            seen.append(cu_seqlens)
+            cu_seqlens.add_(0)  # aiter's in-place write, which bumps the version counter
+            return torch.zeros_like(q)
+
+        monkeypatch.setattr(varlen_mod, "varlen_flash_attention", _fake_varlen)
+
+    def test_published_buffer_is_not_the_tensor_handed_to_the_kernel(self, monkeypatch):
+        seen: list = []
+        self._kernel(seen, monkeypatch)
+
+        model = _FakeModel(layers=3)
+        published = publish_packing(model, _packing(batch=3, seq=8), max_seqlen=8)
+        version_before = published._version
+
+        shape = (3, 8, 2, 4)
+        for attn in model.attentions:
+            cu, bound = resolve_packing(attn)
+            attn.processor._attention(
+                torch.zeros(shape), torch.zeros(shape), torch.zeros(shape), None, cu, bound
+            )
+
+        assert len(seen) == 3
+        for passed in seen:
+            assert passed is not published, (
+                "the kernel was handed the SHARED buffer; its in-place write bumps the version "
+                "counter every layer and the backward then rejects the packing it saved"
+            )
+            assert torch.equal(passed, published), "the private copy must carry the same packing"
+        assert published._version == version_before, (
+            "the shared buffer's autograd version moved during the forward, which is exactly "
+            "what makes each layer's saved packing look stale at backward time"
+        )
+
+
+class TestWrappersAreNotCountedTwice:
+    """Only modules that OWN a var-len processor are consumers.
+
+    Activation checkpointing wraps each attention module in a wrapper that forwards unknown
+    attributes to its child, so a plain ``getattr(module, "processor")`` answers for the wrapper
+    too and every block is counted twice (68 for 34 blocks, seen in the first on-model run).
+    The copy_ still reaches everyone, but the one health number in the log stops meaning
+    anything -- "34 modules" and "half the model" become indistinguishable.
+    """
+
+    class _DelegatingWrapper(nn.Module):
+        def __init__(self, wrapped: nn.Module):
+            super().__init__()
+            self._checkpoint_wrapped_module = wrapped
+
+        def __getattr__(self, name):
+            try:
+                return super().__getattr__(name)
+            except AttributeError:
+                return getattr(self._checkpoint_wrapped_module, name)
+
+    def test_only_the_wrapped_module_is_a_consumer(self):
+        model = _FakeModel(layers=4)
+        for block in model.layers:
+            block.attention = self._DelegatingWrapper(block.attention)
+
+        modules = attention_modules(model)
+
+        assert len(modules) == 4, f"expected one consumer per block, got {len(modules)}"
+        assert all(not isinstance(m, self._DelegatingWrapper) for m in modules)
+        # And the transport still works through the wrapper: the processor is called with the
+        # real module, which is the object that holds the buffer.
+        shared = publish_packing(model, _packing())
+        for block in model.layers:
+            cu, _ = resolve_packing(block.attention._checkpoint_wrapped_module)
+            assert cu is shared
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/nemo_automodel/test_ideogram4_synthetic_data.py b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_synthetic_data.py
new file mode 100644
index 000000000..fcd729f35
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_synthetic_data.py
@@ -0,0 +1,89 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+"""Synthetic Ideogram-4 dataloader: the ``share_text_features`` perf lever.
+
+The lever exists because a realistic caption width makes per-sample generation of the
+[width, 53248] feature buffer cost more CPU than the training step costs GPU, which turns a
+throughput run into a measurement of the dataloader. These tests pin the two properties that
+make sharing safe to use for that purpose -- shapes and per-sample text lengths are unchanged
+-- and the one property that makes it unsafe for the overfit smoke, namely that the samples
+stop being distinct. They run on CPU with a tiny feature dim.
+"""
+import torch
+
+from primus.backends.nemo_automodel.models.ideogram4.data.synthetic import (
+    SyntheticIdeogram4Dataset,
+)
+
+DIM = 8
+
+
+def _dataset(share: bool, **kwargs) -> SyntheticIdeogram4Dataset:
+    params = dict(
+        num_samples=16,
+        in_channels=4,
+        grid_h=2,
+        grid_w=2,
+        max_text_tokens=12,
+        min_text_tokens=4,
+        llm_features_dim=DIM,
+        feature_scale=0.1,
+        latent_scale=1.0,
+        seed=1234,
+        share_text_features=share,
+    )
+    params.update(kwargs)
+    return SyntheticIdeogram4Dataset(**params)
+
+
+def test_default_is_unshared():
+    """Off by default: the overfit smoke needs distinct samples for its loss signal."""
+    assert _dataset(share=False)[0]["llm_features"].data_ptr() != _dataset(share=False)[1]["llm_features"].data_ptr()
+
+    unshared = _dataset(share=False)
+    assert not torch.equal(unshared[0]["llm_features"], unshared[1]["llm_features"])
+
+
+def test_shared_features_are_one_buffer():
+    """The whole point: one allocation, handed back for every index."""
+    shared = _dataset(share=True)
+    first = shared[0]["llm_features"]
+    assert all(shared[i]["llm_features"].data_ptr() == first.data_ptr() for i in range(1, 16))
+
+
+def test_shape_is_identical_to_unshared():
+    """Sharing must not change what the model sees, or throughput is not comparable."""
+    shared, unshared = _dataset(share=True), _dataset(share=False)
+    assert shared[5]["llm_features"].shape == unshared[5]["llm_features"].shape
+    assert shared[5]["llm_features"].dtype == unshared[5]["llm_features"].dtype
+
+
+def test_text_lengths_still_vary_per_sample():
+    """Raggedness is the property the var-len benchmarks turn on; sharing must preserve it."""
+    shared = _dataset(share=True)
+    lengths = {int(shared[i]["text_lengths"]) for i in range(16)}
+    assert len(lengths) > 1
+    assert lengths == {int(_dataset(share=False)[i]["text_lengths"]) for i in range(16)}
+
+
+def test_latents_still_distinct_per_sample():
+    """Only the text buffer is shared; the images are cheap and stay per-index."""
+    shared = _dataset(share=True)
+    assert not torch.equal(shared[0]["image_latents"], shared[1]["image_latents"])
+
+
+def test_shared_buffer_is_deterministic_across_instances():
+    """Same seed, same buffer -- reruns of a benchmark must feed identical data."""
+    assert torch.equal(_dataset(share=True)[0]["llm_features"], _dataset(share=True)[7]["llm_features"])
+
+
+def test_composes_with_cache_in_memory():
+    """The two perf levers are meant to be used together and must not fight."""
+    both = _dataset(share=True, cache_in_memory=True)
+    first = both[3]
+    assert both[3]["llm_features"].data_ptr() == first["llm_features"].data_ptr()
+    assert both[4]["llm_features"].data_ptr() == first["llm_features"].data_ptr()
+    assert torch.equal(both[3]["image_latents"], first["image_latents"])
diff --git a/tests/unit_tests/backends/nemo_automodel/test_ideogram4_zero1.py b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_zero1.py
new file mode 100644
index 000000000..0e6eb4a84
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_ideogram4_zero1.py
@@ -0,0 +1,201 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the DDP + ZeRO-1 hook.
+
+WHY THIS TEST EXISTS:
+  The hook installed cleanly, logged "Installed optional hook", checkpoint-wrapped all 34
+  layers, and trained to a falling loss -- while the ZeRO-1 half never ran at all. YAML
+  ``_target_: torch.optim.AdamW`` is not an OptimizerConfig subclass, so Automodel wraps it
+  in ``OptimizerFromFactoryConfig``, which overrides ``build`` and never chains to
+  ``super()``. The patch was on the base class, so it was simply never called. Nothing
+  raised, nothing warned, and the optimizer state stayed fully replicated: a pure-DDP run
+  wearing a ZeRO-1 label.
+
+  So these tests assert the *interception*, not that training "works" -- the ZeRO-1 run
+  looked perfectly healthy while being broken. The load-bearing one is
+  ``test_factory_config_build_is_intercepted``: it fails against the base-class-only patch.
+
+  No GPU and no process group needed; the ZeRO construction itself is stubbed out.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("nemo_automodel")
+
+from nemo_automodel.components.optim.optimizer import (  # noqa: E402
+    AdamWConfig,
+    OptimizerConfig,
+    OptimizerFromFactoryConfig,
+    build_optimizer_config,
+)
+
+from primus.backends.nemo_automodel.models.ideogram4 import zero1  # noqa: E402
+
+
+@pytest.fixture(autouse=True)
+def _restore_builds():
+    """The patch mutates shared Automodel classes; undo it so tests cannot leak."""
+    saved = [(cls, vars(cls)["build"]) for cls in zero1._optimizer_config_classes()]
+    yield
+    for cls, fn in saved:
+        cls.build = fn
+    for cls, _ in saved:
+        if getattr(vars(cls).get("build"), "_primus_zero1_patched", False):  # pragma: no cover
+            del cls.build
+
+
+@pytest.fixture
+def zero1_on(monkeypatch):
+    monkeypatch.setenv("PRIMUS_IDEOGRAM_ZERO1", "1")
+
+
+def _tiny_model():
+    return torch.nn.Linear(4, 4)
+
+
+def _factory_config():
+    """What the DDP preset's ``_target_: torch.optim.AdamW`` actually normalizes to."""
+    return build_optimizer_config(torch.optim.AdamW, {"lr": 1e-4})
+
+
+# --------------------------------------------------------------------------------------
+# The hole itself
+# --------------------------------------------------------------------------------------
+
+
+def test_ddp_preset_target_normalizes_to_the_factory_config():
+    """Pins the premise: the preset never goes through the base class's ``build``."""
+    cfg = _factory_config()
+    assert isinstance(cfg, OptimizerFromFactoryConfig)
+    assert "build" in vars(type(cfg)), "override is what made the base-class patch dead"
+
+
+def test_factory_config_does_not_chain_to_super():
+    """If it ever did, patching the base alone would have been enough."""
+    import inspect
+
+    assert "super()" not in inspect.getsource(OptimizerFromFactoryConfig.build)
+
+
+def test_hierarchy_walk_covers_every_build_override():
+    covered = {c.__name__ for c in zero1._optimizer_config_classes()}
+    expected = {
+        cls.__name__
+        for cls in _all_subclasses(OptimizerConfig) | {OptimizerConfig}
+        if "build" in vars(cls)
+    }
+    assert covered == expected
+    assert "OptimizerFromFactoryConfig" in covered
+
+
+def _all_subclasses(cls):
+    out = set()
+    for sub in cls.__subclasses__():
+        out.add(sub)
+        out |= _all_subclasses(sub)
+    return out
+
+
+# --------------------------------------------------------------------------------------
+# Interception (the regression pin)
+# --------------------------------------------------------------------------------------
+
+
+def test_factory_config_build_is_intercepted(zero1_on, monkeypatch):
+    """THE load-bearing test: fails if the patch only covers ``OptimizerConfig``."""
+    seen = []
+    monkeypatch.setattr(zero1, "_wrap_in_zero1", lambda opt: seen.append(opt) or opt)
+
+    zero1._install_zero1_optimizer_patch()
+    optimizers = _factory_config().build(_tiny_model())
+
+    assert len(seen) == 1, "ZeRO-1 wrap never ran for the config the DDP preset builds"
+    assert seen[0] is optimizers[0]
+
+
+def test_typed_config_build_is_intercepted(zero1_on, monkeypatch):
+    """``AdamWConfig`` inherits ``build``, so the base-class patch must still cover it."""
+    seen = []
+    monkeypatch.setattr(zero1, "_wrap_in_zero1", lambda opt: seen.append(opt) or opt)
+
+    zero1._install_zero1_optimizer_patch()
+    AdamWConfig(lr=1e-4).build(_tiny_model())
+
+    assert len(seen) == 1
+
+
+def test_build_is_untouched_when_zero1_is_off(monkeypatch):
+    monkeypatch.delenv("PRIMUS_IDEOGRAM_ZERO1", raising=False)
+    seen = []
+    monkeypatch.setattr(zero1, "_wrap_in_zero1", lambda opt: seen.append(opt) or opt)
+
+    zero1._install_zero1_optimizer_patch()
+    optimizers = _factory_config().build(_tiny_model())
+
+    assert seen == [], "the gate is read at call time, not install time"
+    assert isinstance(optimizers[0], torch.optim.AdamW)
+
+
+def test_install_is_idempotent(zero1_on, monkeypatch):
+    """Re-installing must not nest wrappers (which would double-wrap the optimizer)."""
+    seen = []
+    monkeypatch.setattr(zero1, "_wrap_in_zero1", lambda opt: seen.append(opt) or opt)
+
+    zero1._install_zero1_optimizer_patch()
+    zero1._install_zero1_optimizer_patch()
+    _factory_config().build(_tiny_model())
+
+    assert len(seen) == 1
+
+
+def test_patched_build_keeps_its_signature(zero1_on):
+    """The CP plan hit exactly this: a ``*args`` wrapper broke signature introspection."""
+    import inspect
+
+    before = inspect.signature(OptimizerFromFactoryConfig.build)
+    zero1._install_zero1_optimizer_patch()
+    assert inspect.signature(OptimizerFromFactoryConfig.build) == before
+
+
+# --------------------------------------------------------------------------------------
+# Degradation paths: they must be loud, never silent
+# --------------------------------------------------------------------------------------
+
+
+def test_wrap_is_a_noop_on_an_already_wrapped_optimizer():
+    class ZeroRedundancyOptimizer:  # name-matched, as the guard checks the type name
+        pass
+
+    already = ZeroRedundancyOptimizer()
+    assert zero1._wrap_in_zero1(already) is already
+
+
+def test_wrap_skips_dtensor_params_with_a_warning(caplog):
+    """FSDP2 mode: ZeRO-1 does not apply, and the run must say so rather than pretend."""
+
+    class _FakeDTensor:
+        _local_tensor = None
+
+    base = type("Opt", (), {"param_groups": [{"params": [_FakeDTensor()]}]})()
+    with caplog.at_level("WARNING"):
+        assert zero1._wrap_in_zero1(base) is base
+    assert "DTensor" in caplog.text
+
+
+def test_wrap_skips_without_a_process_group(caplog):
+    base = torch.optim.AdamW(_tiny_model().parameters(), lr=1e-4)
+    with caplog.at_level("WARNING"):
+        assert zero1._wrap_in_zero1(base) is base
+    assert "process group" in caplog.text
+
+
+def test_install_is_a_noop_unless_a_gate_is_set(monkeypatch):
+    monkeypatch.delenv("PRIMUS_IDEOGRAM_ZERO1", raising=False)
+    monkeypatch.delenv("PRIMUS_IDEOGRAM_DDP", raising=False)
+    assert zero1.install() is False
diff --git a/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py b/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py
new file mode 100644
index 000000000..5cac22a13
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_nemo_automodel_registration.py
@@ -0,0 +1,233 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for NeMo AutoModel backend registration in __init__.py.
+
+This test ensures that the registration logic in
+primus/backends/nemo_automodel/__init__.py correctly registers the backend with
+BackendRegistry. Without proper registration, the backend is unavailable to the
+runtime. Mirrors tests/unit_tests/backends/megatron/test_megatron_registration.py
+(NeMo AutoModel exposes a single pretrain trainer; no SFT stage).
+
+Importing primus.backends.nemo_automodel does not require the `nemo_automodel`
+package to be installed: the adapter/trainer modules import it lazily (only inside
+the trainer's init()), so these tests run without the backend present.
+"""
+
+import pytest
+
+from primus.backends.nemo_automodel.nemo_automodel_adapter import NemoAutomodelAdapter
+from primus.backends.nemo_automodel.nemo_automodel_pretrain_trainer import (
+    NemoAutomodelPretrainTrainer,
+)
+from primus.core.backend.backend_registry import BackendRegistry
+
+_SUPPORTS_TRAINER_CLASS_REGISTRY = all(
+    hasattr(BackendRegistry, attr)
+    for attr in ("_trainer_classes", "register_trainer_class", "get_trainer_class", "has_trainer_class")
+)
+
+if not _SUPPORTS_TRAINER_CLASS_REGISTRY:
+    pytest.skip(
+        "BackendRegistry trainer-class registration API is not available; "
+        "skip nemo_automodel registration tests that depend on it.",
+        allow_module_level=True,
+    )
+
+
+class TestNemoAutomodelBackendRegistration:
+    """Test that the nemo_automodel backend is properly registered via __init__.py."""
+
+    @pytest.fixture(autouse=True)
+    def ensure_backend_loaded(self):
+        """Ensure the nemo_automodel backend module is loaded before each test."""
+        # Import the __init__ module to trigger registration.
+        import primus.backends.nemo_automodel  # noqa: F401
+
+    def test_adapter_is_registered(self):
+        """Verify that NemoAutomodelAdapter is registered for 'nemo_automodel'."""
+        assert BackendRegistry.has_adapter("nemo_automodel"), (
+            "NemoAutomodelAdapter not registered. "
+            "Check BackendRegistry.register_adapter() call in __init__.py"
+        )
+
+        adapter_cls = BackendRegistry._adapters.get("nemo_automodel")
+        assert adapter_cls is NemoAutomodelAdapter, (
+            f"Expected NemoAutomodelAdapter class, got {adapter_cls}. "
+            "Check BackendRegistry.register_adapter('nemo_automodel', NemoAutomodelAdapter) in __init__.py"
+        )
+
+    def test_trainer_class_is_registered(self):
+        """Verify that NemoAutomodelPretrainTrainer is registered for 'nemo_automodel'."""
+        assert BackendRegistry.has_trainer_class("nemo_automodel"), (
+            "NemoAutomodelPretrainTrainer not registered. "
+            "Check BackendRegistry.register_trainer_class() call in __init__.py"
+        )
+
+        trainer_cls = BackendRegistry.get_trainer_class("nemo_automodel")
+        assert trainer_cls is NemoAutomodelPretrainTrainer, (
+            f"Expected NemoAutomodelPretrainTrainer, got {trainer_cls}. "
+            "Check BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, 'nemo_automodel')"
+        )
+
+        # Explicit pretrain stage should also resolve.
+        assert BackendRegistry.has_trainer_class("nemo_automodel", stage="pretrain")
+
+    def test_adapter_can_be_instantiated_via_registry(self):
+        """Verify that get_adapter returns a working NemoAutomodelAdapter instance."""
+        from unittest.mock import patch
+
+        # Force the lazy-load path and re-register without relying on module reload
+        # side effects (see the megatron registration test for rationale).
+        original_adapters = BackendRegistry._adapters.copy()
+        original_trainers = BackendRegistry._trainer_classes.copy()
+        try:
+            BackendRegistry._adapters.pop("nemo_automodel", None)
+            BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None)
+
+            def _fake_load_backend(_backend: str) -> None:
+                BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter)
+                BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel")
+
+            with patch.object(BackendRegistry, "_load_backend", side_effect=_fake_load_backend):
+                adapter = BackendRegistry.get_adapter("nemo_automodel")
+        finally:
+            BackendRegistry._adapters = original_adapters
+            BackendRegistry._trainer_classes = original_trainers
+
+        assert isinstance(
+            adapter, NemoAutomodelAdapter
+        ), f"Expected NemoAutomodelAdapter instance, got {type(adapter).__name__}"
+        assert (
+            adapter.framework == "nemo_automodel"
+        ), f"Expected framework='nemo_automodel', got '{adapter.framework}'"
+
+    def test_trainer_class_can_be_retrieved_via_adapter(self):
+        """Verify trainer class retrieval through the adapter, including the 'train' alias."""
+        from unittest.mock import patch
+
+        original_adapters = BackendRegistry._adapters.copy()
+        original_trainers = BackendRegistry._trainer_classes.copy()
+        try:
+            BackendRegistry._adapters.pop("nemo_automodel", None)
+            BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None)
+
+            def _fake_load_backend(_backend: str) -> None:
+                BackendRegistry.register_adapter("nemo_automodel", NemoAutomodelAdapter)
+                BackendRegistry.register_trainer_class(NemoAutomodelPretrainTrainer, "nemo_automodel")
+
+            with patch.object(BackendRegistry, "_load_backend", side_effect=_fake_load_backend):
+                adapter = BackendRegistry.get_adapter("nemo_automodel")
+        finally:
+            BackendRegistry._adapters = original_adapters
+            BackendRegistry._trainer_classes = original_trainers
+
+        # The adapter maps the common "train" alias to "pretrain".
+        with patch("primus.backends.nemo_automodel.nemo_automodel_adapter.log_rank_0"):
+            assert adapter.load_trainer_class() is NemoAutomodelPretrainTrainer
+            assert adapter.load_trainer_class(stage="train") is NemoAutomodelPretrainTrainer
+
+    def test_nemo_automodel_in_available_backends_list(self):
+        """Verify nemo_automodel appears in the list of available backends."""
+        available = BackendRegistry.list_available_backends()
+        assert "nemo_automodel" in available, (
+            f"'nemo_automodel' not in available backends: {available}. "
+            "Registration may have failed in __init__.py"
+        )
+
+
+class TestNemoAutomodelRegistrationOrder:
+    """Test that registration happens at import time and is idempotent."""
+
+    def test_registration_is_idempotent(self):
+        """Verify that re-importing __init__ doesn't cause errors."""
+        import importlib
+
+        import primus.backends.nemo_automodel
+
+        importlib.reload(primus.backends.nemo_automodel)
+
+        assert BackendRegistry.has_adapter("nemo_automodel")
+        assert BackendRegistry.has_trainer_class("nemo_automodel")
+
+
+class TestNemoAutomodelRegistrationFailures:
+    """Test error handling when registration is missing."""
+
+    def test_missing_trainer_registration_would_fail(self):
+        """Without trainer registration, get_trainer_class should fail."""
+        original = BackendRegistry._trainer_classes.pop(("nemo_automodel", "pretrain"), None)
+        try:
+            with pytest.raises(ValueError, match="No trainer class registered"):
+                BackendRegistry.get_trainer_class("nemo_automodel")
+        finally:
+            if original:
+                BackendRegistry._trainer_classes[("nemo_automodel", "pretrain")] = original
+
+
+class TestStripPrimusKeysGuard:
+    """Guard the strip_primus_keys denylist against module_base.yaml drift.
+
+    strip_primus_keys() removes Primus-only top-level keys before handing the
+    config to AutoModel. Because it is a denylist, any key later added to
+    module_base.yaml would silently leak into the AutoModel ConfigNode. This test
+    fails loudly if module_base.yaml grows a top-level key not in the strip set.
+    """
+
+    def test_module_base_keys_are_all_stripped(self):
+        from pathlib import Path
+
+        import yaml
+
+        import primus
+        from primus.backends.nemo_automodel.argument_builder import PRIMUS_ONLY_TOP_KEYS
+
+        module_base = Path(primus.__file__).resolve().parent / "configs" / "modules" / "module_base.yaml"
+        assert module_base.is_file(), f"module_base.yaml not found at {module_base}"
+        with open(module_base, encoding="utf-8") as f:
+            base_cfg = yaml.safe_load(f) or {}
+
+        leaked = set(base_cfg) - set(PRIMUS_ONLY_TOP_KEYS)
+        assert not leaked, (
+            f"module_base.yaml top-level keys {sorted(leaked)} are not in PRIMUS_ONLY_TOP_KEYS, "
+            "so they would leak into the AutoModel ConfigNode. Add them to PRIMUS_ONLY_TOP_KEYS "
+            "in primus/backends/nemo_automodel/argument_builder.py (or confirm they belong in "
+            "the AutoModel recipe)."
+        )
+
+    def test_strip_primus_keys_preserves_nested_model_keys(self):
+        from primus.backends.nemo_automodel.argument_builder import strip_primus_keys
+
+        cleaned = strip_primus_keys(
+            {
+                "stage": "pretrain",
+                "framework": "nemo_automodel",
+                "trainable": False,
+                "model": {"stage": "low_noise", "pipeline_spec": {"subfolder": "transformer_2"}},
+                "optim": {"learning_rate": 1e-4},
+            }
+        )
+        assert "stage" not in cleaned and "framework" not in cleaned and "trainable" not in cleaned
+        # Nested AutoModel keys (incl. model.stage) must survive.
+        assert cleaned["model"]["stage"] == "low_noise"
+        assert cleaned["optim"]["learning_rate"] == 1e-4
+
+
+class TestUnsupportedStage:
+    """load_trainer_class should name the real cause for an unsupported stage."""
+
+    def test_unsupported_stage_raises_clear_error(self):
+        from unittest.mock import patch
+
+        adapter = NemoAutomodelAdapter()
+        with patch("primus.backends.nemo_automodel.nemo_automodel_adapter.log_rank_0"):
+            with pytest.raises(RuntimeError, match=r"only exposes a 'pretrain' stage"):
+                adapter.load_trainer_class(stage="sft")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/backends/nemo_automodel/test_parallelize_reshard_forwarding.py b/tests/unit_tests/backends/nemo_automodel/test_parallelize_reshard_forwarding.py
new file mode 100644
index 000000000..6a4ec10fd
--- /dev/null
+++ b/tests/unit_tests/backends/nemo_automodel/test_parallelize_reshard_forwarding.py
@@ -0,0 +1,182 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests that the Primus parallelization strategies forward ``reshard_after_forward``.
+
+WHY THIS TEST EXISTS:
+  This is "Drop 2". Both strategies called ``apply_fsdp2_sharding_recursively`` with
+  seven positional arguments and stopped one short of the eighth, which is
+  ``reshard_after_forward``. The value arrived from upstream as a keyword, landed in
+  ``**kwargs``, and was never read. Nothing failed; the setting was simply ignored and
+  every transformer block reshared after forward anyway.
+
+  Fixing the hook alone (Drop 1) changes nothing without this half, and vice versa,
+  which is exactly why a test per half is worth having. This is the test that would have
+  caught the original bug: it asserts the argument appears in the downstream call.
+
+  The strategies are defined inside ``install()`` and close over the Automodel
+  parallelizer module, so the module is stubbed and ``install()`` is run against it. No
+  GPU and no real Automodel needed.
+"""
+
+import sys
+import types
+
+import pytest
+
+CASES = [
+    pytest.param(
+        "primus.backends.nemo_automodel.models.ideogram4.parallelize",
+        "PRIMUS_IDEOGRAM_REAL_AC",
+        "Ideogram4Transformer2DModel",
+        "layers",
+        id="ideogram4",
+    ),
+    pytest.param(
+        "primus.backends.nemo_automodel.models.flux.parallelize",
+        "PRIMUS_FLUX_REAL_AC",
+        "FluxTransformer2DModel",
+        "transformer_blocks",
+        id="flux",
+    ),
+]
+
+
+class FakeMesh:
+    mesh_dim_names = ()
+
+    def __getitem__(self, _name):
+        raise KeyError(_name)
+
+    def size(self):
+        return 8
+
+
+def _install_stub_parallelizer(monkeypatch):
+    """Stub ``nemo_automodel.components.distributed.parallelizer``.
+
+    Only the handful of names the strategies touch at registration and parallelize time
+    are provided; the sharding calls are recorded instead of performed.
+    """
+    calls = {"sharding": [], "root": []}
+    path = "nemo_automodel.components.distributed.parallelizer"
+    module = types.ModuleType(path)
+
+    class ParallelizationStrategy:
+        pass
+
+    module.ParallelizationStrategy = ParallelizationStrategy
+    module.PARALLELIZATION_STRATEGIES = {}
+
+    def register_parallel_strategy(name):
+        def _decorator(cls):
+            module.PARALLELIZATION_STRATEGIES[name] = cls
+            return cls
+
+        return _decorator
+
+    module.register_parallel_strategy = register_parallel_strategy
+    module.get_fsdp_dp_mesh = lambda mesh, *_a, **_kw: mesh
+
+    def apply_fsdp2_sharding_recursively(*args, **kwargs):
+        calls["sharding"].append((args, kwargs))
+
+    def fully_shard(model, **kwargs):
+        calls["root"].append(kwargs)
+        return model
+
+    module.apply_fsdp2_sharding_recursively = apply_fsdp2_sharding_recursively
+    module.fully_shard = fully_shard
+
+    class MixedPrecisionPolicy:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    module.MixedPrecisionPolicy = MixedPrecisionPolicy
+    module.is_selective_activation_checkpointing = lambda value: value == "selective"
+    module.apply_selective_checkpointing_to_layers = lambda *a, **kw: None
+    module.checkpoint_wrapper = lambda block, **kw: block
+    module.CheckpointImpl = types.SimpleNamespace(NO_REENTRANT="no_reentrant")
+
+    for name in (
+        "nemo_automodel",
+        "nemo_automodel.components",
+        "nemo_automodel.components.distributed",
+        path,
+    ):
+        monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, path, module)
+    return calls
+
+
+def _strategy(monkeypatch, module_path, env_var, model_name):
+    import importlib
+
+    monkeypatch.setenv(env_var, "1")
+    calls = _install_stub_parallelizer(monkeypatch)
+    module = importlib.import_module(module_path)
+    stub = sys.modules["nemo_automodel.components.distributed.parallelizer"]
+    monkeypatch.setattr(module, "P", stub, raising=False)
+
+    assert module.install() is True
+    return stub.PARALLELIZATION_STRATEGIES[model_name](), calls
+
+
+@pytest.mark.parametrize("module_path,env_var,model_name,block_attr", CASES)
+@pytest.mark.parametrize("value", [False, True, None])
+def test_reshard_after_forward_reaches_the_sharding_helper(
+    monkeypatch, module_path, env_var, model_name, block_attr, value
+):
+    """The assertion the original code was missing."""
+    strategy, calls = _strategy(monkeypatch, module_path, env_var, model_name)
+    model = types.SimpleNamespace(**{block_attr: []})
+
+    strategy.parallelize(model=model, device_mesh=FakeMesh(), reshard_after_forward=value)
+
+    assert calls["sharding"], "sharding helper was never called"
+    _args, kwargs = calls["sharding"][-1]
+    assert "reshard_after_forward" in kwargs, "value dropped before the sharding helper"
+    assert kwargs["reshard_after_forward"] is value
+
+
+@pytest.mark.parametrize("module_path,env_var,model_name,block_attr", CASES)
+def test_prefetch_knobs_are_passed_by_keyword_too(
+    monkeypatch, module_path, env_var, model_name, block_attr
+):
+    """Positional passing is what caused the bug; assert the call is keyword-based."""
+    strategy, calls = _strategy(monkeypatch, module_path, env_var, model_name)
+    model = types.SimpleNamespace(**{block_attr: []})
+
+    strategy.parallelize(model=model, device_mesh=FakeMesh(), reshard_after_forward=False)
+
+    args, kwargs = calls["sharding"][-1]
+    # model, mesh, mp_policy, offload_policy stay positional; the rest must be keywords.
+    assert len(args) == 4, f"expected 4 positional args, got {len(args)}"
+    for name in (
+        "enable_fsdp2_prefetch",
+        "fsdp2_backward_prefetch_depth",
+        "fsdp2_forward_prefetch_depth",
+    ):
+        assert name in kwargs
+
+
+@pytest.mark.parametrize("module_path,env_var,model_name,block_attr", CASES)
+def test_root_unit_still_keeps_its_params_unsharded(
+    monkeypatch, module_path, env_var, model_name, block_attr
+):
+    """The root fully_shard keeps reshard_after_forward=False regardless of the override.
+
+    The root unit holds embeddings, norms and heads - a small fraction of parameters -
+    and was already correct before this fix. Changing it is not part of the fix, so
+    pin it so a future edit does not quietly couple the two.
+    """
+    strategy, calls = _strategy(monkeypatch, module_path, env_var, model_name)
+    model = types.SimpleNamespace(**{block_attr: []})
+
+    strategy.parallelize(model=model, device_mesh=FakeMesh(), reshard_after_forward=True)
+
+    assert calls["root"][-1]["reshard_after_forward"] is False
diff --git a/tests/unit_tests/backends/torchtitan/__init__.py b/tests/unit_tests/backends/torchtitan/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/backends/torchtitan/test_fsdp_weight_tying_patch.py b/tests/unit_tests/backends/torchtitan/test_fsdp_weight_tying_patch.py
new file mode 100644
index 000000000..b60846d58
--- /dev/null
+++ b/tests/unit_tests/backends/torchtitan/test_fsdp_weight_tying_patch.py
@@ -0,0 +1,92 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+import importlib
+from types import ModuleType, SimpleNamespace
+from unittest.mock import MagicMock
+
+import primus.backends.torchtitan.patches.fsdp_weight_tying_patches as wt_patch
+from primus.core.patches.patch_registry import PatchRegistry
+
+PATCH_ID = "torchtitan.fsdp.weight_tying"
+
+
+class TestFsdpWeightTyingPatchRegistration:
+    def test_patch_registered(self):
+        assert PATCH_ID in PatchRegistry.list_ids()
+        patch = PatchRegistry.get(PATCH_ID)
+        assert patch is not None
+        assert patch.backend == "torchtitan"
+        assert patch.phase == "setup"
+
+
+class TestFsdpWeightTyingProxy:
+    def test_groups_tied_embedding_and_output(self):
+        tok = object()
+        output = object()
+        norm = object()
+        model = SimpleNamespace(
+            enable_weight_tying=True,
+            tok_embeddings=tok,
+            output=output,
+            norm=norm,
+        )
+
+        calls = []
+
+        def orig_fully_shard(target, **kwargs):
+            calls.append((_normalize := wt_patch._normalize_modules(target), kwargs))
+            return f"wrapped:{_normalize}"
+
+        parallelize_mod = ModuleType("parallelize")
+        parallelize_mod.fully_shard = orig_fully_shard
+
+        def orig_apply_fsdp(model, *args, **kwargs):
+            parallelize_mod.fully_shard(model.tok_embeddings, mesh="dp")
+            parallelize_mod.fully_shard([model.norm, model.output], mesh="dp", reshard=True)
+            return "done"
+
+        wrapped = wt_patch._make_apply_fsdp_wrapper(orig_apply_fsdp, parallelize_mod)
+        assert wrapped(model) == "done"
+        assert calls[0][0] == [tok, output]
+        assert calls[1][0] == [norm]
+
+    def test_passthrough_without_weight_tying(self):
+        model = SimpleNamespace(enable_weight_tying=False)
+        orig_apply_fsdp = MagicMock(return_value="orig")
+        parallelize_mod = ModuleType("parallelize")
+        parallelize_mod.fully_shard = MagicMock()
+
+        wrapped = wt_patch._make_apply_fsdp_wrapper(orig_apply_fsdp, parallelize_mod)
+        assert wrapped(model, mesh="dp") == "orig"
+        orig_apply_fsdp.assert_called_once_with(model, mesh="dp")
+        parallelize_mod.fully_shard.assert_not_called()
+
+
+class TestFsdpWeightTyingPatchModules:
+    def test_patches_llama4_and_qwen3_without_qwen3_fully_shard(self, monkeypatch):
+        llama4_mod = ModuleType("llama4.parallelize")
+        qwen3_mod = ModuleType("qwen3.parallelize")
+
+        llama4_mod.apply_fsdp = MagicMock(return_value="orig")
+        llama4_mod.fully_shard = MagicMock()
+        qwen3_mod.apply_fsdp = llama4_mod.apply_fsdp
+
+        def fake_import(name):
+            if name == "torchtitan.models.llama4.infra.parallelize":
+                return llama4_mod
+            if name == "torchtitan.models.qwen3.infra.parallelize":
+                return qwen3_mod
+            raise ImportError(name)
+
+        monkeypatch.setattr(importlib, "import_module", fake_import)
+        monkeypatch.setattr(wt_patch, "log_rank_0", lambda *args, **kwargs: None)
+
+        orig = llama4_mod.apply_fsdp
+        wt_patch.patch_torchtitan_fsdp_weight_tying(None)
+        assert llama4_mod.apply_fsdp is qwen3_mod.apply_fsdp
+        assert llama4_mod.apply_fsdp is not orig
+        assert not hasattr(qwen3_mod, "fully_shard")
diff --git a/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py b/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py
new file mode 100644
index 000000000..878384493
--- /dev/null
+++ b/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py
@@ -0,0 +1,108 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the GPT-OSS TorchTitan configs added for the v0.2.2 upgrade.
+
+CPU-only: these only parse YAML, so they run in CI without GPU or torchtitan.
+
+Covers:
+    - gpt_oss model configs (BF16 + FP8) resolve to name=gpt_oss with the right
+      flavor and converters (BF16 -> [], FP8 -> [quantize.linear.float8]).
+    - gpt_oss example configs (20B/120B x BF16/FP8 x MI300X/MI325X/MI355X)
+      default to Primus-Turbo on (sink attention) and wire the fp8 switch only
+      for the FP8 variant.
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+MODEL_CFG_DIR = REPO_ROOT / "primus" / "configs" / "models" / "torchtitan"
+EXAMPLE_DIR = REPO_ROOT / "examples" / "torchtitan" / "configs"
+MACHINES = ["MI300X", "MI325X", "MI355X"]
+
+
+def _load(path: Path) -> dict:
+    assert path.exists(), f"missing config: {path}"
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+# -----------------------------------------------------------------------------
+# Model configs
+# -----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("flavor", ["20b", "120b"])
+def test_gpt_oss_bf16_model_config(flavor):
+    cfg = _load(MODEL_CFG_DIR / f"gpt_oss_{flavor}.yaml")
+    model = cfg["model"]
+    assert model["name"] == "gpt_oss"
+    assert model["flavor"] == flavor
+    # BF16 uses no model converter (sink attention is a setup patch, and the
+    # module-form primus_turbo converter is incompatible with GPT-OSS FlexAttn).
+    assert model["converters"] == []
+
+
+@pytest.mark.parametrize("flavor", ["20b", "120b"])
+def test_gpt_oss_fp8_model_config(flavor):
+    cfg = _load(MODEL_CFG_DIR / f"gpt_oss_{flavor}-fp8.yaml")
+    model = cfg["model"]
+    assert model["name"] == "gpt_oss"
+    assert model["flavor"] == flavor
+    # FP8 enables fp8 dense linears through the quantize.linear.float8 converter.
+    assert model["converters"] == ["quantize.linear.float8"]
+
+
+# -----------------------------------------------------------------------------
+# Example pretrain configs
+# -----------------------------------------------------------------------------
+
+
+def _example(machine: str, name: str) -> dict:
+    return _load(EXAMPLE_DIR / machine / f"{name}-pretrain.yaml")
+
+
+@pytest.mark.parametrize("machine", MACHINES)
+@pytest.mark.parametrize("size", ["20B", "120B"])
+@pytest.mark.parametrize("precision", ["BF16", "FP8"])
+def test_gpt_oss_example_defaults_turbo_on(machine, size, precision):
+    cfg = _example(machine, f"gpt_oss_{size}-{precision}")
+    pre = cfg["modules"]["pre_trainer"]
+    assert pre["framework"] == "torchtitan"
+
+    expected_model = (
+        f"gpt_oss_{size.lower()}.yaml" if precision == "BF16" else f"gpt_oss_{size.lower()}-fp8.yaml"
+    )
+    assert pre["model"] == expected_model
+
+    turbo = pre["overrides"]["primus_turbo"]
+    # Turbo + sink attention are on by default for every gpt_oss example.
+    assert turbo["enable_primus_turbo"] is True
+    assert turbo["use_turbo_attention"] is True
+
+    if precision == "FP8":
+        assert turbo["use_turbo_float8_linear"] is True
+    else:
+        assert turbo["use_turbo_float8_linear"] is False
+
+    # MoE experts use GPT-OSS's own grouped-mm path; the shared moe fp8 switch
+    # must stay off so it is not silently assumed to apply.
+    assert turbo["use_moe_fp8"] is False
+
+
+@pytest.mark.parametrize("machine", MACHINES)
+@pytest.mark.parametrize("size", ["20B", "120B"])
+@pytest.mark.parametrize("precision", ["BF16", "FP8"])
+def test_gpt_oss_example_uses_expert_parallel(machine, size, precision):
+    cfg = _example(machine, f"gpt_oss_{size}-{precision}")
+    par = cfg["modules"]["pre_trainer"]["overrides"]["parallelism"]
+    # GPT-OSS is MoE: experts are sharded with Expert Parallel.
+    assert par["expert_parallel_degree"] == 8
+    assert par["data_parallel_shard_degree"] == -1
diff --git a/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py b/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py
new file mode 100644
index 000000000..49e2e5eae
--- /dev/null
+++ b/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py
@@ -0,0 +1,150 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the GPT-OSS Primus-Turbo sink-attention integration.
+
+Two groups:
+    * Registration / condition (CPU-only, no torchtitan needed): the setup patch
+      is registered and only applies for gpt_oss with turbo attention enabled.
+    * Mirror interface (needs torchtitan; auto-skipped otherwise): the mirror
+      Attention keeps the learnable per-head sinks and accepts ``positions``,
+      and the mirror TransformerBlock selects the sliding window on even layers.
+"""
+
+import inspect
+from types import SimpleNamespace
+
+import pytest
+
+# Importing the patch module registers the patch (top-level import pulls in only
+# primus.core.*, torchtitan is imported lazily inside the handler).
+import primus.backends.torchtitan.patches.turbo.gptoss_sink_attention_patches as gptoss_patch
+from primus.core.patches import PatchContext
+from primus.core.patches.patch_registry import PatchRegistry
+
+PATCH_ID = "torchtitan.primus_turbo.gptoss_sink_attention"
+
+
+def _ctx(model_name, enable_turbo, use_turbo_attention, backend="torchtitan", phase="setup"):
+    params = SimpleNamespace(
+        model=SimpleNamespace(name=model_name),
+        primus_turbo=SimpleNamespace(
+            enable_primus_turbo=enable_turbo,
+            use_turbo_attention=use_turbo_attention,
+        ),
+    )
+    module_config = SimpleNamespace(params=params)
+    return PatchContext(backend=backend, phase=phase, extra={"module_config": module_config})
+
+
+class TestSinkAttentionPatchRegistration:
+    def test_patch_registered(self):
+        assert PATCH_ID in PatchRegistry.list_ids()
+        patch = PatchRegistry.get(PATCH_ID)
+        assert patch is not None
+        assert patch.backend == "torchtitan"
+        assert patch.phase == "setup"
+
+    def test_condition_enabled_for_gpt_oss(self):
+        assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", True, True)) is True
+
+    def test_condition_disabled_for_other_models(self):
+        for name in ("llama3", "llama4", "qwen3", "deepseek_v3"):
+            assert gptoss_patch._gptoss_turbo_enabled(_ctx(name, True, True)) is False
+
+    def test_condition_disabled_when_turbo_off(self):
+        assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", False, True)) is False
+        assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", True, False)) is False
+
+    def test_applies_to_enabled_for_gpt_oss(self):
+        # End-to-end wiring (backend + phase + condition) resolves to True for a
+        # gpt_oss turbo run. Negative cases are covered by the condition tests
+        # above (applies_to logs on skip, which needs the Primus logger set up).
+        patch = PatchRegistry.get(PATCH_ID)
+        assert patch.applies_to(_ctx("gpt_oss", True, True)) is True
+
+
+@pytest.fixture
+def gpt_oss_args():
+    pytest.importorskip("torchtitan")
+    from torchtitan.models.gpt_oss import GptOssModelArgs
+    from torchtitan.models.moe import MoEArgs
+
+    return GptOssModelArgs(
+        dim=256,
+        n_heads=8,
+        n_kv_heads=2,
+        head_dim=32,
+        sliding_window_size=64,
+        vocab_size=512,
+        moe_inter_dim=256,
+        n_layers=2,
+        moe_args=MoEArgs(
+            num_experts=4,
+            num_shared_experts=0,
+            top_k=2,
+            use_grouped_mm=True,
+            score_func="softmax",
+            route_norm=True,
+            gate_bias=True,
+        ),
+    )
+
+
+class TestSinkAttentionMirror:
+    def test_attention_keeps_sinks_and_accepts_positions(self, gpt_oss_args):
+        from primus.backends.torchtitan.models.gpt_oss.model.model import Attention
+
+        attn = Attention(gpt_oss_args)
+        # Learnable per-head sinks preserved from the upstream module.
+        assert tuple(attn.sinks.shape) == (gpt_oss_args.n_heads,)
+        # Full (no) window by default; per-layer window set by the block.
+        assert attn._turbo_window == (-1, -1)
+        assert attn.sliding_window_size == gpt_oss_args.sliding_window_size
+        # GPT-OSS attention forward matches the upstream signature (rope_cache +
+        # attention_masks; GPT-OSS does not take the llama/qwen positions arg).
+        params = inspect.signature(attn.forward).parameters
+        assert "rope_cache" in params
+        assert "attention_masks" in params
+
+    def test_block_selects_sliding_window_on_even_layers(self, gpt_oss_args):
+        import torch
+        import torch.nn as nn
+
+        from primus.backends.torchtitan.models.gpt_oss.model.model import (
+            TransformerBlock,
+        )
+
+        captured = {}
+
+        class AttnStub(nn.Module):
+            def __init__(self, sliding_window_size):
+                super().__init__()
+                self.sliding_window_size = sliding_window_size
+                self._turbo_window = (-1, -1)
+
+            def forward(self, x, rope_cache, attention_masks=None):
+                captured["window"] = self._turbo_window
+                return torch.zeros_like(x)
+
+        class MoeStub(nn.Module):
+            def forward(self, x):
+                return torch.zeros_like(x)
+
+        for layer_id, expected in [
+            (0, (gpt_oss_args.sliding_window_size, 0)),  # even -> sliding
+            (1, (-1, -1)),  # odd -> full
+        ]:
+            block = TransformerBlock(layer_id, gpt_oss_args)
+            assert block.use_sliding_attention == (layer_id % 2 == 0)
+            # Replace the heavy attention/moe with CPU stubs to test the block's
+            # window-selection logic without invoking the Triton flash kernel.
+            block.attention = AttnStub(gpt_oss_args.sliding_window_size)
+            block.moe = MoeStub()
+            x = torch.zeros(1, 4, gpt_oss_args.dim)
+            block(x, None, None)
+            assert captured["window"] == expected
diff --git a/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py b/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py
new file mode 100644
index 000000000..10ef794ba
--- /dev/null
+++ b/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py
@@ -0,0 +1,103 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the Primus-Turbo MoE grouped_mm patch.
+
+Regression focus: the replacement on
+``torchtitan.models.moe.moe._run_experts_grouped_mm`` must expose a
+``__qualname__`` containing the sentinel ``_run_experts_grouped_mm_dynamic`` so
+upstream ``apply_compile`` treats it as already patched and skips
+``torch.compile`` over the ``torch.compiler.disable``'d turbo kernels (dynamo
+gb0098). A ``functools.partial`` (no ``__qualname__``) would crash that read.
+"""
+
+import functools
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+import primus.backends.torchtitan.patches.turbo.moe_grouped_mm_patches as moe_patch
+from primus.core.patches import PatchContext
+from primus.core.patches.patch_registry import PatchRegistry
+
+PATCH_ID = "torchtitan.primus_turbo.moe_grouped_mm"
+
+
+def _ctx(enable_turbo=True, use_turbo_grouped_mm=True, use_moe_fp8=False):
+    params = SimpleNamespace(
+        primus_turbo=SimpleNamespace(
+            enable_primus_turbo=enable_turbo,
+            use_turbo_grouped_mm=use_turbo_grouped_mm,
+            use_moe_fp8=use_moe_fp8,
+        ),
+    )
+    module_config = SimpleNamespace(params=params)
+    return PatchContext(backend="torchtitan", phase="setup", extra={"module_config": module_config})
+
+
+class TestMoeGroupedMmPatch:
+    def test_patch_registered(self):
+        assert PATCH_ID in PatchRegistry.list_ids()
+        p = PatchRegistry.get(PATCH_ID)
+        assert p is not None and p.backend == "torchtitan" and p.phase == "setup"
+
+    def test_condition_gates_on_turbo_flags(self):
+        p = PatchRegistry.get(PATCH_ID)
+        assert p.condition(_ctx(True, True)) is True
+        assert p.condition(_ctx(False, True)) is False
+        assert p.condition(_ctx(True, False)) is False
+
+    def _run_patch(self, use_moe_fp8, stub_impl):
+        """Apply the real patch against the real torchtitan moe module.
+
+        Saves/restores torchtitan.models.moe.moe._run_experts_grouped_mm and
+        stubs the Primus grouped_mm impl so no GPU kernel runs.
+        """
+        import torchtitan.models.moe.moe as tt_moe
+
+        import primus.backends.torchtitan.models.moe.moe as primus_moe
+
+        orig_tt = tt_moe._run_experts_grouped_mm
+        orig_primus = primus_moe._run_experts_grouped_mm
+        try:
+            primus_moe._run_experts_grouped_mm = stub_impl
+            with patch.object(moe_patch, "log_rank_0"):
+                moe_patch.patch_torchtitan_moe(_ctx(use_moe_fp8=use_moe_fp8))
+            return tt_moe._run_experts_grouped_mm
+        finally:
+            tt_moe._run_experts_grouped_mm = orig_tt
+            primus_moe._run_experts_grouped_mm = orig_primus
+
+    def test_replacement_has_qualname_not_partial(self):
+        pytest.importorskip("torchtitan")
+
+        def _stub(w1, w2, w3, x, num_tokens_per_expert, use_fp8=True):
+            return use_fp8
+
+        _stub.__qualname__ = "_run_experts_grouped_mm"
+
+        replacement = self._run_patch(use_moe_fp8=True, stub_impl=_stub)
+        assert not isinstance(replacement, functools.partial)
+        assert hasattr(replacement, "__qualname__")
+        # Must match upstream's already_patched guard to skip torch.compile.
+        assert "_run_experts_grouped_mm_dynamic" in replacement.__qualname__
+
+    def test_replacement_binds_use_fp8(self):
+        pytest.importorskip("torchtitan")
+        captured = {}
+
+        def _stub(w1, w2, w3, x, num_tokens_per_expert, use_fp8=True):
+            captured["use_fp8"] = use_fp8
+            return "ok"
+
+        _stub.__qualname__ = "_run_experts_grouped_mm"
+
+        replacement = self._run_patch(use_moe_fp8=True, stub_impl=_stub)
+        # use_fp8 must be injected by the wrapper (upstream calls positionally).
+        assert replacement(1, 2, 3, 4, 5) == "ok"
+        assert captured["use_fp8"] is True
diff --git a/tests/unit_tests/ci/test_select_tests.py b/tests/unit_tests/ci/test_select_tests.py
index 91031635f..1baddec7f 100644
--- a/tests/unit_tests/ci/test_select_tests.py
+++ b/tests/unit_tests/ci/test_select_tests.py
@@ -4,7 +4,12 @@
 # See LICENSE for license information.
 ###############################################################################
 
-"""Unit tests for tools/ci/select_tests.py (classify-based PR test selection)."""
+"""Unit tests for tools/ci/select_tests.py (classify-based E2E suite selection).
+
+Only E2E selection is covered: the unit-test suite is always run in full (see
+select_tests.py's module docstring for why), so there's nothing to select
+there anymore.
+"""
 
 import importlib.util
 from pathlib import Path
@@ -13,54 +18,7 @@
 _SPEC = importlib.util.spec_from_file_location("select_tests", _ROOT / "tools/ci/select_tests.py")
 _MOD = importlib.util.module_from_spec(_SPEC)
 _SPEC.loader.exec_module(_MOD)
-select = _MOD.select_targets
-
-FULL = ["tests/unit_tests/"]
-
-
-# --- unit-test selection ---------------------------------------------------
-def test_empty_runs_full():
-    assert select([]) == FULL
-
-
-def test_global_change_runs_full():
-    assert select([".github/workflows/ci.yaml"]) == FULL
-    assert select(["tools/ci/select_tests.py"]) == FULL
-    assert select(["runner/helpers/x.sh"]) == FULL
-    assert select(["requirements.txt"]) == FULL
-    assert select(["primus/core/launcher/initialize.py"]) == FULL
-
-
-def test_non_py_under_primus_runs_full():
-    # configs / fixtures can't be localized to a unit dir -> fail-safe.
-    assert select(["primus/configs/x.yaml"]) == FULL
-
-
-def test_backend_maps_to_its_unit_dir():
-    out = select(["primus/backends/megatron/training/global_vars.py"])
-    assert "tests/unit_tests/backends/megatron/" in out
-    assert "tests/unit_tests/megatron/" in out  # megatron's extra GPU-operator tests
-
 
-def test_backend_without_unit_dir_runs_full():
-    # transformer_engine has no tests/unit_tests/backends/transformer_engine/.
-    assert select(["primus/backends/transformer_engine/x.py"]) == FULL
-
-
-def test_component_maps_to_isomorphic_dir():
-    assert select(["primus/core/projection/engine.py"]) == ["tests/unit_tests/core/projection/"]
-    assert select(["primus/agents/a.py"]) == ["tests/unit_tests/agents/"]
-
-
-def test_changed_unit_test_runs_its_dir():
-    assert select(["tests/unit_tests/agents/test_tools.py"]) == ["tests/unit_tests/agents/"]
-
-
-def test_docs_only_runs_full():
-    assert select(["README.md", "docs/guide.md"]) == FULL
-
-
-# --- E2E selection (pass a fixed suite set for determinism) ----------------
 SUITES = {"megatron", "torchtitan", "maxtext"}
 
 
@@ -84,14 +42,29 @@ def test_e2e_backend_with_trainer_runs_its_suite():
 
 
 def test_e2e_backend_without_trainer_runs_all():
-    # No trainer suite (bridge / hummingbirdxt / transformer_engine) -> fail-safe all.
+    # No trainer suite (bridge / hummingbirdxt / transformer_engine / diffusion) -> fail-safe all.
     assert set(e2e(["primus/backends/megatron_bridge/x.py"])) == SUITES
     assert set(e2e(["primus/backends/hummingbirdxt/x.py"])) == SUITES
     assert set(e2e(["primus/backends/transformer_engine/x.py"])) == SUITES
+    assert set(e2e(["primus/backends/diffusion/x.py"])) == SUITES
 
 
 def test_e2e_component_change_runs_all():
+    # Any other primus/ or tests/unit_tests/ change -- not just the
+    # explicitly-listed GLOBAL_TRIGGERS -- also runs everything: classify()
+    # maps it to "component", which select_e2e() treats the same as "global".
     assert set(e2e(["primus/core/trainer/base.py"])) == SUITES
+    assert set(e2e(["primus/core/launcher/parser.py"])) == SUITES
+    assert set(e2e(["tests/unit_tests/core/patches/test_patch.py"])) == SUITES
+
+
+def test_e2e_bare_examples_file_runs_all():
+    # A bare examples/<file> (no backend subdir) is shared launcher plumbing
+    # -- e.g. test_maxtext_trainer.py shells out to examples/run_pretrain.sh
+    # directly -- so it must not be silently ignored like a docs-only change.
+    assert set(e2e(["examples/run_pretrain.sh"])) == SUITES
+    # examples/<backend>/... is unaffected: still maps to that one backend.
+    assert e2e(["examples/maxtext/configs/x.yaml"]) == ["maxtext"]
 
 
 def test_e2e_docs_only_runs_none():
diff --git a/tests/unit_tests/cli/__init__.py b/tests/unit_tests/cli/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/cli/test_data_config.py b/tests/unit_tests/cli/test_data_config.py
new file mode 100644
index 000000000..2db295857
--- /dev/null
+++ b/tests/unit_tests/cli/test_data_config.py
@@ -0,0 +1,253 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Tests for data preprocessing config parsing, validation, and authentication.
+
+Tests the YAML-to-CLI mapping in data.py and the auth priority chain in auth.py.
+"""
+
+import argparse
+import os
+import tempfile
+from unittest.mock import patch
+
+import pytest
+
+from primus.backends.megatron.data.diffusion.preprocessing.auth import (
+    setup_hf_authentication,
+)
+from primus.cli.subcommands.data import (
+    _flatten_preprocessing_config,
+    _get_encoded_parser_defaults,
+    _load_config_with_cli_overrides,
+    _validate_preprocessing_config,
+)
+from tests.utils import PrimusUT
+
+
+class TestFlattenPreprocessingConfig(PrimusUT):
+    """Tests for _flatten_preprocessing_config YAML-to-flat-dict mapping."""
+
+    def test_all_sections_mapped(self):
+        """All 6 YAML sections produce correct flat keys."""
+        config = {
+            "source": {
+                "type": "huggingface",
+                "hf_dataset": "diffusers/pokemon",
+                "hf_split": "train",
+            },
+            "data_format": {
+                "image_key": "jpg",
+                "caption_key": "json.caption",
+            },
+            "output": {
+                "output_dir": "/data/out",
+                "shard_size": 500,
+            },
+            "model": {
+                "model_path": "my-model",
+                "batch_size": 4,
+                "precision": "fp16",
+            },
+            "image": {
+                "image_size": 512,
+                "variable_size": True,
+                "center_crop": False,
+            },
+            "auth": {
+                "hf_token_file": "/path/to/token",
+            },
+        }
+        flat = _flatten_preprocessing_config(config)
+
+        assert flat["source_type"] == "huggingface"
+        assert flat["hf_dataset"] == "diffusers/pokemon"
+        assert flat["image_key"] == "jpg"
+        assert flat["caption_key"] == "json.caption"
+        assert flat["output_dir"] == "/data/out"
+        assert flat["shard_size"] == 500
+        assert flat["model_path"] == "my-model"
+        assert flat["batch_size"] == 4
+        assert flat["precision"] == "fp16"
+        assert flat["image_size"] == 512
+        assert flat["variable_size"] is True
+        assert flat["center_crop"] is False
+        assert flat["hf_token_file"] == "/path/to/token"
+
+    def test_partial_config(self):
+        """Partial config with only source + output produces only those keys."""
+        config = {
+            "source": {"type": "directory", "input_dir": "/data/images"},
+            "output": {"output_dir": "/data/out"},
+        }
+        flat = _flatten_preprocessing_config(config)
+
+        assert flat["source_type"] == "directory"
+        assert flat["input_dir"] == "/data/images"
+        assert flat["output_dir"] == "/data/out"
+        assert "model_path" not in flat
+        assert "image_size" not in flat
+
+    def test_model_defaults(self):
+        """Model section injects correct defaults when keys are absent."""
+        config = {"model": {}}
+        flat = _flatten_preprocessing_config(config)
+
+        assert flat["model_path"] == "black-forest-labs/FLUX.1-dev"
+        assert flat["batch_size"] == 8
+        assert flat["precision"] == "bf16"
+        assert flat["device"] == "cuda"
+        assert flat["t5_max_length"] == 512
+
+    def test_image_defaults(self):
+        """Image section injects correct defaults when keys are absent."""
+        config = {"image": {}}
+        flat = _flatten_preprocessing_config(config)
+
+        assert flat["image_size"] == 1024
+        assert flat["variable_size"] is False
+        assert flat["center_crop"] is True
+        assert flat["max_size"] == 1024
+
+
+class TestValidatePreprocessingConfig(PrimusUT):
+    """Tests for _validate_preprocessing_config."""
+
+    def _make_args(self, **kwargs):
+        defaults = {
+            "source_type": "huggingface",
+            "output_dir": "/data/out",
+            "hf_dataset": "test/dataset",
+            "input_dir": None,
+            "input_path": None,
+        }
+        defaults.update(kwargs)
+        return argparse.Namespace(**defaults)
+
+    def test_missing_source_type_raises(self):
+        """Missing source_type raises ValueError."""
+        args = self._make_args(source_type=None)
+        with self.assertRaises(ValueError, msg="source_type"):
+            _validate_preprocessing_config(args)
+
+    def test_missing_output_dir_raises(self):
+        """Missing output_dir raises ValueError."""
+        args = self._make_args(output_dir=None)
+        with self.assertRaises(ValueError, msg="output_dir"):
+            _validate_preprocessing_config(args)
+
+    def test_huggingface_without_hf_dataset_raises(self):
+        """HuggingFace source without hf_dataset raises ValueError."""
+        args = self._make_args(source_type="huggingface", hf_dataset=None)
+        with self.assertRaises(ValueError, msg="hf-dataset"):
+            _validate_preprocessing_config(args)
+
+    def test_directory_without_input_dir_raises(self):
+        """Directory source without input_dir raises ValueError."""
+        args = self._make_args(source_type="directory", input_dir=None)
+        with self.assertRaises(ValueError, msg="input-dir"):
+            _validate_preprocessing_config(args)
+
+    def test_valid_config_passes(self):
+        """Valid config raises no errors."""
+        args = self._make_args()
+        _validate_preprocessing_config(args)
+
+
+class TestLoadConfigWithCliOverrides(PrimusUT):
+    """Tests for _load_config_with_cli_overrides merge logic."""
+
+    @patch("primus.core.utils.yaml_utils.parse_yaml")
+    def test_cli_overrides_yaml(self, mock_parse_yaml):
+        """Explicitly set CLI args override YAML config values."""
+        mock_parse_yaml.return_value = {
+            "source": {"type": "huggingface", "hf_dataset": "yaml-dataset"},
+            "model": {"batch_size": 4},
+        }
+        defaults = _get_encoded_parser_defaults()
+        args = argparse.Namespace(
+            config="test.yaml",
+            batch_size=16,
+            **{k: v for k, v in defaults.items() if k not in ("config", "batch_size")},
+        )
+
+        result = _load_config_with_cli_overrides(args)
+
+        assert result.batch_size == 16
+        assert result.hf_dataset == "yaml-dataset"
+
+    @patch("primus.core.utils.yaml_utils.parse_yaml")
+    def test_yaml_used_when_cli_is_default(self, mock_parse_yaml):
+        """YAML values used when CLI arg equals its default."""
+        mock_parse_yaml.return_value = {
+            "model": {"batch_size": 4},
+        }
+        defaults = _get_encoded_parser_defaults()
+        args = argparse.Namespace(config="test.yaml", **{k: v for k, v in defaults.items() if k != "config"})
+
+        result = _load_config_with_cli_overrides(args)
+
+        assert result.batch_size == 4
+
+    def test_no_config_passthrough(self):
+        """No config file returns args unchanged."""
+        original = argparse.Namespace(config=None, batch_size=8)
+        result = _load_config_with_cli_overrides(original)
+        assert result is original
+
+
+class TestSetupHfAuthenticationPriority(PrimusUT):
+    """Tests for setup_hf_authentication priority chain."""
+
+    def test_file_takes_priority_over_env(self):
+        """Token file takes priority over HF_TOKEN env var."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".token", delete=False) as f:
+            f.write("hf_file_token_123")
+            token_path = f.name
+        try:
+            os.chmod(token_path, 0o600)
+            old_env = os.environ.get("HF_TOKEN")
+            os.environ["HF_TOKEN"] = "hf_env_token_456"
+            try:
+                token = setup_hf_authentication(token_file=token_path, use_env=True)
+                assert token == "hf_file_token_123"
+            finally:
+                if old_env is None:
+                    os.environ.pop("HF_TOKEN", None)
+                else:
+                    os.environ["HF_TOKEN"] = old_env
+        finally:
+            os.unlink(token_path)
+
+    def test_env_takes_priority_over_cache(self):
+        """HF_TOKEN env var is used when no file is provided."""
+        old_env = os.environ.get("HF_TOKEN")
+        os.environ["HF_TOKEN"] = "hf_env_token_789"
+        try:
+            token = setup_hf_authentication(token_file=None, use_env=True)
+            assert token == "hf_env_token_789"
+        finally:
+            if old_env is None:
+                os.environ.pop("HF_TOKEN", None)
+            else:
+                os.environ["HF_TOKEN"] = old_env
+
+    def test_no_auth_returns_none(self):
+        """No auth sources returns None."""
+        old_env = os.environ.pop("HF_TOKEN", None)
+        try:
+            with patch.object(
+                type(__import__("pathlib").Path()),
+                "exists",
+                return_value=False,
+            ):
+                token = setup_hf_authentication(token_file=None, use_env=True)
+                assert token is None
+        finally:
+            if old_env is not None:
+                os.environ["HF_TOKEN"] = old_env
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/cli/test_preflight_subcommand.py b/tests/unit_tests/cli/test_preflight_subcommand.py
index 6ffb4eeea..47a997b5d 100644
--- a/tests/unit_tests/cli/test_preflight_subcommand.py
+++ b/tests/unit_tests/cli/test_preflight_subcommand.py
@@ -39,7 +39,9 @@ def test_defaults():
     assert args.perf_test is False
     assert args.dist_timeout_sec == 120
     assert args.dump_path == "output/preflight"
-    assert args.report_file_name == "preflight_report"
+    # Default is None so the tool auto-generates a unique timestamped name
+    # (preflight-{NNODES}N-{YYYYMMDD-HHMMSS}) at run time.
+    assert args.report_file_name is None
     assert args.save_pdf is True
 
 
diff --git a/tests/unit_tests/cli/test_train_subcommand.py b/tests/unit_tests/cli/test_train_subcommand.py
deleted file mode 100644
index 04955245a..000000000
--- a/tests/unit_tests/cli/test_train_subcommand.py
+++ /dev/null
@@ -1,40 +0,0 @@
-###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
-#
-# See LICENSE for license information.
-###############################################################################
-
-"""Unit tests for the train subcommand's runtime router (``_resolve_pretrain_runtime``).
-
-This selects the active core runtime vs the legacy pretrain flow; run()/register
-import heavy runtime modules and are out of scope.
-"""
-
-from __future__ import annotations
-
-import pytest
-
-pytest.importorskip("primus.cli.subcommands.train")
-
-from primus.cli.subcommands.train import _resolve_pretrain_runtime  # noqa: E402
-
-
-def test_runtime_defaults_to_core(monkeypatch):
-    monkeypatch.delenv("PRIMUS_TRAIN_RUNTIME", raising=False)
-    assert _resolve_pretrain_runtime(None) == "core"
-
-
-def test_runtime_explicit_legacy(monkeypatch):
-    monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "legacy")
-    assert _resolve_pretrain_runtime(None) == "legacy"
-
-
-def test_runtime_explicit_core_case_insensitive(monkeypatch):
-    monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "CORE")
-    assert _resolve_pretrain_runtime(None) == "core"
-
-
-def test_runtime_invalid_value_warns_and_falls_back_to_core(monkeypatch, capsys):
-    monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "bogus")
-    assert _resolve_pretrain_runtime(None) == "core"
-    assert "Ignoring invalid" in capsys.readouterr().err
diff --git a/tests/unit_tests/configs/test_config_schema_check.py b/tests/unit_tests/configs/test_config_schema_check.py
new file mode 100644
index 000000000..1e02fb7da
--- /dev/null
+++ b/tests/unit_tests/configs/test_config_schema_check.py
@@ -0,0 +1,1305 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the backend config schema drift scanner.
+
+Most tests run against a synthetic repo built in ``tmp_path`` so they neither
+need the ``third_party/`` submodules nor break when upstream moves a field.
+A few tests do touch the real tree, because they guard the failure modes that
+motivated the scanner:
+
+* ``training.mock_data`` must be derived automatically from the patch system
+  (it is a Primus-only field; a previous audit wrongly called it dead code);
+* ``primus_turbo`` must be derived from ``experimental.custom_args_module``;
+* ``train_iters`` and friends must come out of Megatron's generated arguments,
+  because reading its literal ``add_argument`` calls alone would report roughly
+  half of every valid Megatron config as drift;
+* ``norm_epsilon`` must come out as scoped to DeepSeek-V4, and no DeepSeek-V4
+  config may be reported for setting it.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import textwrap
+from pathlib import Path
+
+import pytest
+
+from primus.core.config.schema_check import (
+    Allowlist,
+    AllowRule,
+    Finding,
+    KeyFinding,
+    ScopedFinding,
+    build_allowlist,
+    build_model_scopes,
+    build_schema,
+    check_manual_allowlist_consumers,
+    dedupe,
+    dedupe_scoped,
+    extract_args_read_rules,
+    extract_backend_config_extensions,
+    extract_custom_args_schema,
+    extract_get_param_rules,
+    load_manual_allowlist,
+    load_maxtext_schema,
+    load_megatron_schema,
+    load_torchtitan_schema,
+    scan_backend,
+)
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+
+_JOB_CONFIG_SRC = """
+# Upstream imports torch at module scope, which is why the scanner must use AST.
+import torch
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Training:
+    seq_len: int = 2048
+    steps: int = 10000
+
+
+@dataclass
+class Debug:
+    seed: int | None = None
+    moe_force_load_balance: bool = False
+
+
+@dataclass
+class Float8Linear:
+    filter_fqns: list[str] = field(default_factory=list)
+
+
+@dataclass
+class QuantizedLinear:
+    float8: Float8Linear = field(default_factory=Float8Linear)
+
+
+@dataclass
+class Quantize:
+    linear: QuantizedLinear = field(default_factory=QuantizedLinear)
+
+
+class NotADataclass:
+    ignored: int = 0
+
+
+@dataclass
+class Experimental:
+    custom_args_module: str = ""
+
+
+@dataclass
+class JobConfig:
+    training: Training = field(default_factory=Training)
+    debug: Debug = field(default_factory=Debug)
+    quantize: Quantize = field(default_factory=Quantize)
+    experimental: Experimental = field(default_factory=Experimental)
+"""
+
+_CONFIG_EXTENSION_SRC = """
+from dataclasses import dataclass, field
+
+from torchtitan.config.job_config import JobConfig as TTJobConfig
+
+
+@dataclass
+class PrimusTurboConfig:
+    enable_primus_turbo: bool = False
+    use_turbo_attention: bool = False
+
+
+@dataclass
+class JobConfig(TTJobConfig):
+    primus_turbo: PrimusTurboConfig = field(default_factory=PrimusTurboConfig)
+"""
+
+_PATCH_SRC = """
+from primus.core.patches import get_param, register_patch
+
+
+@register_patch(patch_id="fake", backend="torchtitan", phase="setup")
+def fake_patch(ctx):
+    if get_param(ctx, "training.mock_data", False):
+        return get_param(ctx, "primus_turbo.enable_primus_turbo", False)
+    return None
+"""
+
+_ALLOWLIST_SRC = """
+common:
+  - key: work_group
+    reason: Primus envelope.
+    consumer: primus/core/launcher/parser.py:211
+  - key: user_name
+    reason: Primus envelope.
+    consumer: primus/core/launcher/parser.py:212
+  - key: exp_name
+    reason: Primus envelope.
+    consumer: primus/core/launcher/parser.py:213
+torchtitan:
+  - key: model.*
+    reason: Model-arg override, validated at runtime.
+    consumer: primus/backends/torchtitan/patches/model_override_patches.py:273
+"""
+
+
+_MEGATRON_ARGUMENTS_SRC = """
+# Upstream imports torch at module scope, which is why the scanner must use AST.
+import torch
+
+from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.training.argument_utils import ArgumentGroupFactory
+
+
+def add_megatron_arguments(parser):
+    group = parser.add_argument_group(title="core")
+    group.add_argument("--seq-length", type=int, default=None)
+    group.add_argument("--num-experts", type=int, default=None)
+    # An explicit dest is the only thing that keeps this off `no_persist_layer_norm`.
+    group.add_argument("--no-persist-layer-norm", action="store_false", dest="persist_layer_norm")
+
+    # Upstream builds this list in code rather than inline.
+    exclude = [
+        "moe_token_dropping",
+        "gradient_accumulation_fusion",
+    ]
+    ArgumentGroupFactory(TransformerConfig, exclude=exclude).build_group(parser, "transformer")
+
+    from megatron.training.config import LoggerConfig, TrainingConfig
+
+    ArgumentGroupFactory(TrainingConfig).build_group(parser, "training")
+    ArgumentGroupFactory(LoggerConfig, exclude=["memory_keys"]).build_group(parser, "logging")
+    return parser
+
+
+def core_transformer_config_from_args(args, config_class=None):
+    config_class = config_class or TransformerConfig
+    if args.multi_latent_attention:
+        config_class = MLATransformerConfig
+    return config_class(**{f.name: getattr(args, f.name) for f in dataclasses.fields(config_class)})
+"""
+
+_MEGATRON_TRAINING_CONFIG_SRC = """
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TrainingConfig:
+    train_iters: int | None = None
+    global_batch_size: int | None = None
+    micro_batch_size: int = 1
+    # Derived after parsing, so ArgumentGroupFactory skips it.
+    computed_batches: int | None = field(init=False, default=None)
+
+
+@dataclass
+class LoggerConfig:
+    log_interval: int = 100
+    memory_keys: dict | None = None
+    use_nsys_profiler: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--profile"], "dest": "profile"}}
+    )
+"""
+
+_MEGATRON_TRANSFORMER_CONFIG_SRC = """
+from dataclasses import dataclass
+
+from ..model_parallel_config import ModelParallelConfig
+
+
+@dataclass
+class TransformerConfig(ModelParallelConfig):
+    num_layers: int = 0
+    moe_token_dropping: bool = False
+    gradient_accumulation_fusion: bool = True
+
+
+@dataclass
+class MLATransformerConfig(TransformerConfig):
+    q_lora_rank: int | None = None
+"""
+
+_MEGATRON_PATCH_SRC = """
+import os
+
+from primus.core.patches import get_args, register_patch
+
+_PARAM_NAME = "dataloader_mp_context"
+
+
+@register_patch(patch_id="fake", backend="megatron", phase="setup")
+def fake_patch(ctx, tensor):
+    # An attribute of something that is not an args object must stay out.
+    _ = tensor.shape
+    os.environ["HSA_NO_SCRATCH_RECLAIM"] = "1"
+    if getattr(get_args(ctx), "use_turbo_attention", False):
+        return get_args(ctx).log_avg_skip_iterations
+    return getattr(get_args(ctx), _PARAM_NAME, None)
+"""
+
+_MEGATRON_CONFIG_EXTENSION_SRC = """
+from dataclasses import dataclass
+
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+
+@dataclass
+class PrimusBaseConfig(MLATransformerConfig):
+    norm_epsilon: float | None = None
+
+
+@dataclass
+class PrimusChildConfig(PrimusBaseConfig):
+    conv_bias: bool = False
+"""
+
+# The config class of one model, holding one field of each kind the scope
+# derivation has to tell apart.
+_MEGATRON_SCOPED_CONFIG_SRC = """
+from dataclasses import dataclass
+
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+
+@dataclass
+class FakeV4TransformerConfig(MLATransformerConfig):
+    hash_routing_seed: int = 0
+    attn_sink: bool = False
+    # Upstream defines this one, so every model gets it from `args`.
+    num_layers: int = 0
+    # A config class no model selects declares this one too.
+    conv_bias: bool = False
+    # A patch reads this one off `args`, which every model has.
+    index_topk: int = 0
+    # The hand-written allowlist vouches for this one everywhere.
+    swiglu_limit: float = 0.0
+"""
+
+_MEGATRON_SCOPED_BUILDER_SRC = """
+from megatron.training.arguments import core_transformer_config_from_args
+
+from .fake_v4_config import FakeV4TransformerConfig
+
+
+def fake_v4_builder(args, pre_process, post_process, config=None):
+    if config is None:
+        config = core_transformer_config_from_args(args, config_class=FakeV4TransformerConfig)
+    return config
+"""
+
+# Mirrors the real dispatch: the branch that names a Primus module outright is
+# the one that leads to a Primus config class; the others reach upstream.
+_MODEL_PROVIDER_SRC = """
+import importlib
+from functools import partial
+
+
+def get_model_provider(model_type="gpt"):
+    if model_type == "fake_v4":
+        module = importlib.import_module(
+            "primus.backends.megatron.core.models.fake_v4.fake_v4_builders"
+        )
+        return partial(module.model_provider, module.fake_v4_builder)
+    if model_type == "mamba":
+        return lazy_import(["model_provider", "pretrain_mamba"], "model_provider")
+    return lazy_import(["model_provider", "pretrain_gpt"], "model_provider")
+"""
+
+_MEGATRON_SCOPE_PATCH_SRC = """
+from primus.core.patches import get_args, register_patch
+
+
+@register_patch(patch_id="fake_flops", backend="megatron", phase="setup")
+def fake_flops_patch(ctx):
+    if get_args(ctx).model_type == "fake_v4":
+        # A read off the args namespace, which every model has: not model-scoped.
+        return get_args(ctx).index_topk
+    return None
+"""
+
+
+def _write(path: Path, text: str) -> Path:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(textwrap.dedent(text).lstrip())
+    return path
+
+
+@pytest.fixture
+def fake_repo(tmp_path: Path) -> Path:
+    """A minimal repo with both submodules, a patch, and an extension dataclass."""
+    _write(tmp_path / "third_party/torchtitan/torchtitan/config/job_config.py", _JOB_CONFIG_SRC)
+    _write(
+        tmp_path / "third_party/maxtext/src/maxtext/configs/base.yml",
+        "run_name: ''\nsteps: 10\nper_device_batch_size: 1\n",
+    )
+    _write(
+        tmp_path / "primus/backends/torchtitan/primus_turbo_extensions/config_extension.py",
+        _CONFIG_EXTENSION_SRC,
+    )
+    _write(tmp_path / "primus/backends/torchtitan/patches/fake_patches.py", _PATCH_SRC)
+    _write(
+        tmp_path / "primus/configs/modules/torchtitan/pre_trainer.yaml",
+        """
+        experimental:
+          custom_args_module: "primus.backends.torchtitan.primus_turbo_extensions.config_extension"
+        training:
+          mock_data: true
+        """,
+    )
+    _write(tmp_path / "tools/ci/config_schema_allowlist.yaml", _ALLOWLIST_SRC)
+    return tmp_path
+
+
+@pytest.fixture
+def megatron_repo(tmp_path: Path) -> Path:
+    """A minimal repo whose Megatron mirrors the parts the scanner has to read.
+
+    Namely: config dataclasses reached through a re-export, a base class in
+    another package, an ``exclude`` list built in code, an ``init=False`` field
+    and an ``argparse_meta`` dest override.
+    """
+    megatron = tmp_path / "third_party/Megatron-LM/megatron"
+    _write(megatron / "training/arguments.py", _MEGATRON_ARGUMENTS_SRC)
+    _write(
+        megatron / "training/config/__init__.py",
+        "from megatron.training.config.training_config import LoggerConfig, TrainingConfig\n",
+    )
+    _write(megatron / "training/config/training_config.py", _MEGATRON_TRAINING_CONFIG_SRC)
+    _write(
+        megatron / "core/transformer/__init__.py",
+        "from .transformer_config import MLATransformerConfig, TransformerConfig\n",
+    )
+    _write(megatron / "core/transformer/transformer_config.py", _MEGATRON_TRANSFORMER_CONFIG_SRC)
+    _write(
+        megatron / "core/model_parallel_config.py",
+        "from dataclasses import dataclass\n\n\n"
+        "@dataclass\nclass ModelParallelConfig:\n"
+        "    tensor_model_parallel_size: int = 1\n    timers: object = None\n",
+    )
+    _write(tmp_path / "primus/backends/megatron/patches/fake_patches.py", _MEGATRON_PATCH_SRC)
+    _write(tmp_path / "primus/backends/megatron/patches/fake_flops_patches.py", _MEGATRON_SCOPE_PATCH_SRC)
+    _write(
+        tmp_path / "primus/backends/megatron/core/models/fake_config.py",
+        _MEGATRON_CONFIG_EXTENSION_SRC,
+    )
+    _write(
+        tmp_path / "primus/backends/megatron/core/models/fake_v4/fake_v4_config.py",
+        _MEGATRON_SCOPED_CONFIG_SRC,
+    )
+    _write(
+        tmp_path / "primus/backends/megatron/core/models/fake_v4/fake_v4_builders.py",
+        _MEGATRON_SCOPED_BUILDER_SRC,
+    )
+    _write(tmp_path / "primus/core/utils/import_utils.py", _MODEL_PROVIDER_SRC)
+    _write(tmp_path / "primus/configs/models/megatron/fake_v4.yaml", "model_type: fake_v4\nnum_layers: 2\n")
+    _write(tmp_path / "primus/configs/models/megatron/plain.yaml", "num_layers: 2\n")
+    _write(
+        tmp_path / "tools/ci/config_schema_allowlist.yaml",
+        "common:\n"
+        "  - key: work_group\n"
+        "    reason: Primus envelope.\n"
+        "    consumer: primus/core/launcher/parser.py\n"
+        "megatron:\n"
+        "  - key: swiglu_limit\n"
+        "    reason: Read by every MoE model, not just this one.\n"
+        "    consumer: primus/core/launcher/parser.py\n",
+    )
+    _write(tmp_path / "primus/core/launcher/parser.py", "cfg.work_group, cfg.swiglu_limit\n")
+    return tmp_path
+
+
+# ---------------------------------------------------------------------------
+# Schema extraction
+# ---------------------------------------------------------------------------
+
+
+def test_torchtitan_schema_extracted_from_ast(fake_repo: Path):
+    schema = load_torchtitan_schema(fake_repo)
+
+    assert schema is not None
+    assert schema.sections == {
+        "training",
+        "debug",
+        "quantize",
+        "quantize.linear",
+        "quantize.linear.float8",
+        "experimental",
+    }
+    assert schema.keys == {
+        "training.seq_len",
+        "training.steps",
+        "debug.seed",
+        "debug.moe_force_load_balance",
+        "quantize.linear.float8.filter_fqns",
+        "experimental.custom_args_module",
+    }
+    # A plain class must not leak into the schema.
+    assert not any(k.startswith("ignored") for k in schema.keys)
+
+
+def test_maxtext_schema_is_flat(fake_repo: Path):
+    schema = load_maxtext_schema(fake_repo)
+
+    assert schema is not None
+    assert schema.keys == {"run_name", "steps", "per_device_batch_size"}
+    assert schema.sections == frozenset()
+
+
+def test_schema_suggests_where_a_moved_key_went(fake_repo: Path):
+    schema = load_torchtitan_schema(fake_repo)
+
+    assert schema.suggest("training.seed") == "debug.seed"
+    # Upstream drops the section-name prefix when it moves a field into a section.
+    assert schema.suggest("training.debug_moe_force_load_balance") == "debug.moe_force_load_balance"
+    assert schema.suggest("training.no_such_thing_anywhere") is None
+
+
+def test_missing_submodule_yields_no_schema(tmp_path: Path):
+    assert load_torchtitan_schema(tmp_path) is None
+    assert load_maxtext_schema(tmp_path) is None
+    assert load_megatron_schema(tmp_path) is None
+
+
+# ---------------------------------------------------------------------------
+# Megatron schema extraction
+# ---------------------------------------------------------------------------
+
+
+def test_megatron_schema_joins_literal_and_generated_arguments(megatron_repo: Path):
+    schema = load_megatron_schema(megatron_repo)
+
+    assert schema is not None
+    assert schema.sections == frozenset()  # Megatron's namespace is flat.
+    # The literal `add_argument` half.
+    assert {"seq_length", "num_experts"} <= schema.keys
+    # The ArgumentGroupFactory half, which is where the most-used names live.
+    assert {"train_iters", "global_batch_size", "micro_batch_size", "log_interval"} <= schema.keys
+    # A field's base class contributes too: TransformerConfig(ModelParallelConfig).
+    assert {"tensor_model_parallel_size", "timers"} <= schema.keys
+
+
+def test_megatron_factory_exclude_list_is_honoured(megatron_repo: Path):
+    """`memory_keys` is excluded from LoggerConfig, so no argument carries it."""
+    schema = load_megatron_schema(megatron_repo)
+
+    assert "log_interval" in schema.keys
+    assert "memory_keys" not in schema.keys
+
+
+def test_megatron_transformer_config_keeps_its_excluded_fields(megatron_repo: Path):
+    """`core_transformer_config_from_args` copies every field off `args`.
+
+    So a TransformerConfig field stays a legal key even where the factory
+    excludes it from argparse -- unlike LoggerConfig, which nothing builds from
+    `args`.
+    """
+    schema = load_megatron_schema(megatron_repo)
+
+    assert {"moe_token_dropping", "gradient_accumulation_fusion"} <= schema.keys
+    # MLATransformerConfig is the other class that function can pick.
+    assert "q_lora_rank" in schema.keys
+
+
+def test_megatron_derived_fields_are_not_arguments(megatron_repo: Path):
+    """`field(init=False)` is computed after parsing; `build_group` skips it."""
+    assert "computed_batches" not in load_megatron_schema(megatron_repo).keys
+
+
+def test_megatron_argparse_meta_dest_wins_over_the_field_name(megatron_repo: Path):
+    """`use_nsys_profiler` reaches `args` as `profile`, so only `profile` is legal."""
+    schema = load_megatron_schema(megatron_repo)
+
+    assert "profile" in schema.keys
+    assert "use_nsys_profiler" not in schema.keys
+
+
+def test_megatron_explicit_dest_wins_over_the_option_name(megatron_repo: Path):
+    """`--no-persist-layer-norm dest=persist_layer_norm` must not yield both."""
+    schema = load_megatron_schema(megatron_repo)
+
+    assert "persist_layer_norm" in schema.keys
+    assert "no_persist_layer_norm" not in schema.keys
+
+
+def test_megatron_suggests_a_renamed_argument(megatron_repo: Path):
+    """A flat namespace renames in place, by adding a prefix or a suffix."""
+    schema = load_megatron_schema(megatron_repo)
+
+    assert schema.suggest("layers") == "num_layers"
+    assert schema.suggest("nothing_like_this") is None
+
+
+# ---------------------------------------------------------------------------
+# Allowlist derivation
+# ---------------------------------------------------------------------------
+
+
+def test_get_param_paths_are_derived_from_patches(fake_repo: Path):
+    rules = extract_get_param_rules(fake_repo, "torchtitan")
+
+    patterns = {r.pattern for r in rules}
+    assert patterns == {"training.mock_data", "primus_turbo.enable_primus_turbo"}
+    assert all(r.consumer.startswith("primus/backends/torchtitan/patches/") for r in rules)
+
+
+def test_custom_args_module_is_followed(fake_repo: Path):
+    keys, sections, module_path = extract_custom_args_schema(fake_repo)
+
+    assert module_path.endswith("config_extension")
+    assert sections == {"primus_turbo"}
+    assert keys == {"primus_turbo.enable_primus_turbo", "primus_turbo.use_turbo_attention"}
+
+
+def test_extension_section_stays_checked(fake_repo: Path):
+    """`primus_turbo` is a section, not an opaque wildcard: bad children still fail."""
+    schema = build_schema(fake_repo, "torchtitan")
+
+    assert "primus_turbo" in schema.sections
+    assert "primus_turbo.use_turbo_attention" in schema.keys
+    assert "primus_turbo.nonexistent" not in schema.keys
+
+
+def test_megatron_args_reads_are_derived_from_the_backend_package(megatron_repo: Path):
+    rules = extract_args_read_rules(megatron_repo, "megatron")
+    patterns = {r.pattern for r in rules}
+
+    assert "use_turbo_attention" in patterns  # getattr(args, "<name>", default)
+    assert "log_avg_skip_iterations" in patterns  # get_args(ctx).<name>
+    assert "dataloader_mp_context" in patterns  # read through a module constant
+    # An attribute of something that is not an args object is not a config key.
+    assert "shape" not in patterns
+    # Neither is an environment variable, so a config that sets one still fails.
+    assert "HSA_NO_SCRATCH_RECLAIM" not in patterns
+    assert all(r.consumer.startswith("primus/backends/megatron/") for r in rules)
+
+
+def test_megatron_config_extension_fields_are_schema(megatron_repo: Path):
+    """A Primus TransformerConfig subclass widens the key set, transitively."""
+    keys, source = extract_backend_config_extensions(megatron_repo, "megatron")
+
+    assert source == "primus/backends/megatron"
+    assert {"norm_epsilon", "conv_bias"} <= keys
+    schema = build_schema(megatron_repo, "megatron")
+    assert {"norm_epsilon", "conv_bias"} <= schema.keys
+
+
+# ---------------------------------------------------------------------------
+# Model scopes
+# ---------------------------------------------------------------------------
+
+
+def test_model_scope_is_read_off_the_model_dispatch(megatron_repo: Path):
+    """`get_model_provider` is what decides the model, so it is what we follow."""
+    scopes = build_model_scopes(megatron_repo, "megatron")
+
+    assert len(scopes) == 1
+    scope = scopes[0]
+    assert (scope.key, scope.value, scope.default) == ("model_type", "fake_v4", "gpt")
+    assert scope.config_class == "FakeV4TransformerConfig"
+    assert scope.source == "primus/core/utils/import_utils.py"
+
+
+def test_only_a_field_nothing_else_explains_is_model_scoped(megatron_repo: Path):
+    """Four kinds of field share the scoped class and none of them is scoped."""
+    scope = build_model_scopes(megatron_repo, "megatron")[0]
+
+    assert scope.keys == {"hash_routing_seed", "attn_sink"}
+    assert "num_layers" not in scope.keys  # upstream defines it for every model
+    assert "conv_bias" not in scope.keys  # a config class no model selects has it too
+    assert "index_topk" not in scope.keys  # a patch reads it off `args`
+    assert "swiglu_limit" not in scope.keys  # the hand-written allowlist vouches for it
+
+
+def test_no_model_dispatch_means_no_scopes(fake_repo: Path):
+    """TorchTitan has no such dispatch; the check must simply not look for one."""
+    assert build_model_scopes(fake_repo, "torchtitan") == ()
+
+
+def test_a_scoped_key_on_another_model_is_reported(megatron_repo: Path):
+    _write(
+        megatron_repo / "examples/megatron/configs/other_model.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: plain.yaml
+            overrides:
+              train_iters: 50
+              hash_routing_seed: 7
+        """,
+    )
+
+    result = scan_backend(megatron_repo, "megatron")
+
+    assert result.findings == ()  # the key is not unknown, only misplaced
+    misplaced = [f for f in result.scoped if f.file.endswith("other_model.yaml")]
+    assert [(f.key, f.actual) for f in misplaced] == [("hash_routing_seed", "gpt")]
+    assert misplaced[0].scope.describe() == "model_type: fake_v4"
+
+
+def test_a_scoped_key_on_its_own_model_is_not_reported(megatron_repo: Path):
+    """The `model_type` lives in the model preset, not in the experiment YAML."""
+    _write(
+        megatron_repo / "examples/megatron/configs/right_model.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: fake_v4.yaml
+            overrides:
+              hash_routing_seed: 7
+        """,
+    )
+
+    result = scan_backend(megatron_repo, "megatron")
+
+    assert result.scoped == ()
+
+
+def test_an_experiment_may_override_the_model_type_itself(megatron_repo: Path):
+    _write(
+        megatron_repo / "examples/megatron/configs/override.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: plain.yaml
+            overrides:
+              model_type: fake_v4
+              hash_routing_seed: 7
+        """,
+    )
+
+    assert scan_backend(megatron_repo, "megatron").scoped == ()
+
+
+def test_an_unreadable_model_preset_reports_nothing(megatron_repo: Path):
+    """Guessing here would report the one config where the key belongs."""
+    _write(
+        megatron_repo / "examples/megatron/configs/unknown_model.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: not_a_preset.yaml
+            overrides:
+              hash_routing_seed: 7
+        """,
+    )
+
+    assert scan_backend(megatron_repo, "megatron").scoped == ()
+
+
+def test_a_model_preset_of_its_own_is_judged_on_what_it_declares(megatron_repo: Path):
+    """A preset that omits `model_type` gets the dispatch default, which is `gpt`."""
+    _write(megatron_repo / "primus/configs/models/megatron/stray.yaml", "attn_sink: true\n")
+    _write(
+        megatron_repo / "primus/configs/models/megatron/scoped.yaml", "model_type: fake_v4\nattn_sink: true\n"
+    )
+
+    result = scan_backend(megatron_repo, "megatron")
+
+    assert [(f.file, f.key, f.actual) for f in result.scoped] == [
+        ("primus/configs/models/megatron/stray.yaml", "attn_sink", "gpt")
+    ]
+
+
+def test_dedupe_scoped_collapses_one_key_across_files(megatron_repo: Path):
+    scope = build_model_scopes(megatron_repo, "megatron")[0]
+    findings = [ScopedFinding("megatron", "attn_sink", f"cfg{i}.yaml", scope, "gpt") for i in range(4)]
+    findings.append(ScopedFinding("megatron", "attn_sink", "cfg4.yaml", scope, "mamba"))
+
+    rows = dedupe_scoped(findings)
+
+    assert len(rows) == 1
+    assert (rows[0].key, rows[0].count, rows[0].models) == ("attn_sink", 5, ("gpt", "mamba"))
+    assert rows[0].sample(2) == "`cfg0.yaml`, `cfg1.yaml` (+3 more)"
+
+
+def test_manual_allowlist_requires_reason_and_consumer(tmp_path: Path):
+    path = _write(tmp_path / "allow.yaml", "common:\n  - key: foo\n    reason: because\n")
+
+    with pytest.raises(ValueError, match="consumer"):
+        load_manual_allowlist(path, "torchtitan")
+
+
+def test_manual_allowlist_scopes_are_merged(fake_repo: Path):
+    allow = build_allowlist(fake_repo, "torchtitan")
+
+    assert allow.match("work_group") is not None
+    assert allow.match("model.n_layers").pattern == "model.*"
+    assert allow.match("training.mock_data").origin == "patch:get_param"
+    assert allow.match("nope") is None
+
+
+def test_allowlist_entry_is_flagged_once_its_consumer_stops_reading_it(tmp_path: Path):
+    _write(tmp_path / "primus/core/launcher/parser.py", "cfg.work_group = x\n")
+    path = _write(
+        tmp_path / "allow.yaml",
+        "common:\n"
+        "  - key: work_group\n"
+        "    reason: still read\n"
+        "    consumer: primus/core/launcher/parser.py\n"
+        "  - key: removed_key\n"
+        "    reason: stale\n"
+        "    consumer: primus/core/launcher/parser.py\n",
+    )
+
+    problems = check_manual_allowlist_consumers(tmp_path, path, "torchtitan")
+
+    assert len(problems) == 1
+    assert "removed_key" in problems[0]
+
+
+def test_allowlist_entry_is_flagged_once_its_consumer_file_is_gone(tmp_path: Path):
+    path = _write(
+        tmp_path / "allow.yaml",
+        "common:\n  - key: foo\n    reason: r\n    consumer: primus/core/deleted.py\n",
+    )
+
+    problems = check_manual_allowlist_consumers(tmp_path, path, "torchtitan")
+
+    assert len(problems) == 1 and "does not exist" in problems[0]
+
+
+def test_a_line_number_left_on_a_consumer_does_not_break_the_check(tmp_path: Path):
+    """The path is what gets verified, so a legacy `path:line` must still pass."""
+    _write(tmp_path / "primus/core/launcher/parser.py", "cfg.work_group = x\n")
+    path = _write(
+        tmp_path / "allow.yaml",
+        "common:\n"
+        "  - key: work_group\n"
+        "    reason: r\n"
+        "    consumer: primus/core/launcher/parser.py:211\n",
+    )
+
+    assert check_manual_allowlist_consumers(tmp_path, path, "torchtitan") == []
+
+
+def test_an_unchecked_out_submodule_is_not_reported_as_a_stale_consumer(tmp_path: Path):
+    """Absence of a submodule proves nothing about the entry, so stay quiet."""
+    path = _write(
+        tmp_path / "allow.yaml",
+        "maxtext:\n  - key: base_config\n    reason: r\n    consumer: third_party/maxtext/x.py\n",
+    )
+
+    assert check_manual_allowlist_consumers(tmp_path, path, "maxtext") == []
+
+
+def test_the_shipped_allowlist_entries_are_all_still_backed():
+    for backend in ("torchtitan", "maxtext", "megatron"):
+        problems = check_manual_allowlist_consumers(
+            _REPO_ROOT, _REPO_ROOT / "tools/ci/config_schema_allowlist.yaml", backend
+        )
+        assert problems == []
+
+
+def test_real_repo_derives_primus_only_fields():
+    """Regression guard: these must never be reported as drift."""
+    allow = build_allowlist(_REPO_ROOT, "torchtitan")
+    schema = build_schema(_REPO_ROOT, "torchtitan")
+    if schema is None:
+        pytest.skip("third_party/torchtitan is not checked out")
+
+    mock_data = allow.match("training.mock_data")
+    assert mock_data is not None and mock_data.origin == "patch:get_param"
+    assert "primus_turbo" in schema.sections
+    assert "primus_turbo.enable_primus_turbo" in schema.keys
+
+
+def test_real_repo_megatron_schema_holds_the_generated_arguments():
+    """Guard the half of Megatron's namespace that no `add_argument` call names."""
+    schema = build_schema(_REPO_ROOT, "megatron")
+    if schema is None:
+        pytest.skip("third_party/Megatron-LM is not checked out")
+
+    assert {"train_iters", "global_batch_size", "micro_batch_size", "log_interval"} <= schema.keys
+    assert {"seq_length", "lr", "num_experts"} <= schema.keys  # the literal half
+    assert len(schema.keys) > 600
+
+
+def test_real_repo_derives_primus_only_megatron_fields():
+    """Regression guard: these must never be reported as drift."""
+    allow = build_allowlist(_REPO_ROOT, "megatron")
+
+    for key in ("use_turbo_attention", "log_avg_skip_iterations", "enable_zero_bubble", "odc_gda_pipe"):
+        rule = allow.match(key)
+        assert rule is not None and rule.origin.startswith("backend:"), key
+
+
+def test_real_repo_scopes_norm_epsilon_to_deepseek_v4():
+    """The case this feature exists for.
+
+    `norm_epsilon` is not a Megatron argparse dest -- `layernorm_epsilon` is, and
+    it merely declares `--norm-epsilon` as its flag. Only
+    `DeepSeekV4TransformerConfig` gives the yaml key a field to land in.
+    """
+    scopes = build_model_scopes(_REPO_ROOT, "megatron")
+    if not scopes:
+        pytest.skip("third_party/Megatron-LM is not checked out")
+
+    by_value = {scope.value: scope for scope in scopes}
+    assert "deepseek_v4" in by_value
+    scope = by_value["deepseek_v4"]
+    assert (scope.key, scope.default) == ("model_type", "gpt")
+    assert scope.config_class == "DeepSeekV4TransformerConfig"
+    assert "norm_epsilon" in scope.keys
+    # A field read off `args` by code every model runs is not this model's.
+    assert "moe_use_legacy_grouped_gemm" not in scope.keys
+
+
+def test_real_repo_never_scopes_the_hand_assembled_flux_fields():
+    """`FluxConfig` is assembled by `FluxPretrainTrainer`, not by the dispatch.
+
+    It never reaches `core_transformer_config_from_args`, and its trainer class
+    appears in no model preset, so no `model_type` value owns its fields and
+    they stay legal on the backend as a whole. Should scoping ever start
+    attributing an unowned config class to some model, every diffusion preset in
+    the tree gets reported for keys it does read.
+    """
+    result = scan_backend(_REPO_ROOT, "megatron")
+    if not result.available:
+        pytest.skip("third_party/Megatron-LM is not checked out")
+
+    assert result.scopes, "the model dispatch no longer resolves"
+    # Declared by `FluxConfig` and `BaseDiffusionConfig`, by nothing upstream.
+    diffusion_fields = {
+        "num_joint_layers",
+        "num_single_layers",
+        "context_dim",
+        "vec_in_dim",
+        "guidance_embed",
+        "single_block_bias",
+        "sensitive_layers_enabled",
+        "fp8_scaling_strategy",
+    }
+    for scope in result.scopes:
+        assert diffusion_fields.isdisjoint(scope.keys), scope.config_class
+    assert [f.key for f in result.scoped if f.key in diffusion_fields] == []
+    assert [f.file for f in result.scoped if "diffusion" in f.file] == []
+
+
+def test_real_repo_never_reports_a_scoped_key_on_the_model_it_belongs_to():
+    """The false-positive guard: DeepSeek-V4's own configs set every V4 key there is."""
+    result = scan_backend(_REPO_ROOT, "megatron")
+    if not result.available:
+        pytest.skip("third_party/Megatron-LM is not checked out")
+
+    assert result.scopes, "the model dispatch no longer resolves"
+    for finding in result.scoped:
+        assert finding.actual != finding.scope.value
+    assert [f.file for f in result.scoped if "deepseek_v4" in f.file] == []
+
+
+# ---------------------------------------------------------------------------
+# Scanning and reporting
+# ---------------------------------------------------------------------------
+
+
+def _exp_yaml(overrides: str) -> str:
+    return (
+        textwrap.dedent(
+            """
+            work_group: ${PRIMUS_TEAM:amd}
+            user_name: root
+            exp_name: probe
+            modules:
+              pre_trainer:
+                framework: torchtitan
+                config: pre_trainer.yaml
+                model: fake.yaml
+                overrides:
+            """
+        ).lstrip()
+        + textwrap.indent(textwrap.dedent(overrides).strip() + "\n", " " * 6)
+    )
+
+
+def test_scan_reports_only_unknown_keys(fake_repo: Path):
+    _write(
+        fake_repo / "examples/torchtitan/configs/probe.yaml",
+        _exp_yaml(
+            """
+            training:
+              seq_len: 4096
+              seed: 42
+            model:
+              n_layers: 4
+            primus_turbo:
+              enable_primus_turbo: true
+              bogus: 1
+            """
+        ),
+    )
+
+    result = scan_backend(fake_repo, "torchtitan")
+
+    assert result.available
+    assert {f.key for f in result.findings} == {"training.seed", "primus_turbo.bogus"}
+    assert not result.errors
+
+
+def test_scan_skips_modules_of_other_backends(fake_repo: Path):
+    _write(
+        fake_repo / "examples/torchtitan/configs/other.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: fake.yaml
+            overrides:
+              anything_goes: 1
+        """,
+    )
+
+    assert scan_backend(fake_repo, "torchtitan").findings == ()
+
+
+def test_unloadable_config_is_reported_not_raised(fake_repo: Path):
+    _write(fake_repo / "examples/torchtitan/configs/broken.yaml", "extends:\n  - missing.yaml\n")
+
+    result = scan_backend(fake_repo, "torchtitan")
+
+    assert [e.file for e in result.errors] == ["examples/torchtitan/configs/broken.yaml"]
+    # The reader wrote `missing.yaml`, so name that, not the raw exception path.
+    assert result.errors[0].message.startswith("cannot resolve `examples/torchtitan/configs/missing.yaml`")
+
+
+def test_unloadable_config_does_not_count_as_checked(fake_repo: Path):
+    """The whole point of the tool is that nothing passes unexamined."""
+    _write(fake_repo / "examples/torchtitan/configs/ok.yaml", _exp_yaml("training:\n  seq_len: 8\n"))
+    _write(fake_repo / "examples/torchtitan/configs/broken.yaml", "extends:\n  - missing.yaml\n")
+
+    result = scan_backend(fake_repo, "torchtitan")
+
+    # pre_trainer.yaml + ok.yaml + broken.yaml
+    assert result.scanned == 3
+    assert result.checked == 2
+    assert result.findings == ()
+
+
+def test_megatron_scan_reports_only_unknown_keys(megatron_repo: Path):
+    """Megatron configs are flat, so every finding is a top-level key."""
+    _write(
+        megatron_repo / "examples/megatron/configs/probe.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: fake.yaml
+            overrides:
+              train_iters: 50
+              norm_epsilon: 1.0e-6
+              use_turbo_attention: true
+              renamed_away: 1
+        """,
+    )
+
+    result = scan_backend(megatron_repo, "megatron")
+
+    assert result.available
+    assert {f.key for f in result.findings} == {"renamed_away"}
+    assert not result.errors
+
+
+def test_scan_without_submodule_is_unavailable(tmp_path: Path):
+    result = scan_backend(tmp_path, "torchtitan")
+
+    assert not result.available
+    assert result.findings == () and result.scanned == 0
+
+
+def test_dedupe_collapses_one_key_across_files(fake_repo: Path):
+    findings = [Finding("torchtitan", "training.seed", f"cfg{i}.yaml") for i in range(5)]
+    findings.append(Finding("torchtitan", "training.steps", "cfg0.yaml"))
+
+    rows = dedupe(findings, load_torchtitan_schema(fake_repo))
+
+    assert [(r.key, r.count) for r in rows] == [("training.seed", 5), ("training.steps", 1)]
+    assert rows[0].suggestion == "debug.seed"
+    assert rows[0].files == tuple(f"cfg{i}.yaml" for i in range(5))
+    assert rows[1].files == ("cfg0.yaml",)
+
+
+def test_dedupe_drops_duplicate_files_for_one_key():
+    rows = dedupe([Finding("maxtext", "gone", "a.yaml")] * 3)
+
+    assert rows[0].count == 3 and rows[0].files == ("a.yaml",)
+
+
+def test_dedupe_orders_by_blast_radius_whatever_order_it_is_handed():
+    """`(backend, key)` is unique per row, so the order is total: no tie for the input order to
+    break, and CI can diff one run's summary against the next."""
+    findings = [Finding("megatron", "wide", f"cfg{i}.yaml") for i in range(3)]
+    findings += [Finding("megatron", "narrow", "cfg0.yaml"), Finding("maxtext", "other", "cfg0.yaml")]
+
+    assert [(r.backend, r.key) for r in dedupe(findings)] == [
+        ("maxtext", "other"),
+        ("megatron", "wide"),
+        ("megatron", "narrow"),
+    ]
+    assert dedupe(findings) == dedupe(findings[::-1])
+
+
+def test_allowlist_matches_exact_before_glob():
+    allow = Allowlist((AllowRule("a.b", "r", "c", "manual"), AllowRule("a.*", "r", "c", "manual")))
+
+    assert allow.match("a.b").pattern == "a.b"
+    assert allow.match("a.c").pattern == "a.*"
+    assert allow.match("b.c") is None
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def _load_cli():
+    path = _REPO_ROOT / "tools/ci/check_config_schema.py"
+    spec = importlib.util.spec_from_file_location("primus_check_config_schema", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _row(key: str, files: tuple[str, ...], suggestion: str | None = None) -> KeyFinding:
+    return KeyFinding("megatron", key, len(files), files, suggestion)
+
+
+def test_group_rows_merges_the_keys_that_share_their_configs():
+    """The whole point: one preset's worth of rot is one row, not twenty."""
+    cli = _load_cli()
+    rows = [
+        _row("zeta", ("a.yaml", "b.yaml")),
+        _row("alpha", ("a.yaml", "b.yaml")),
+        _row("solo", ("c.yaml",)),
+    ]
+
+    renamed, merged = cli.group_rows(rows)
+
+    assert renamed == []
+    assert [(g.keys, g.count) for g in merged] == [(("alpha", "zeta"), 2), (("solo",), 1)]
+    # Merging must not lose a config: the row's files are every member's files.
+    assert merged[0].files == ("a.yaml", "b.yaml")
+
+
+def test_group_rows_keeps_a_renamed_key_out_of_the_merged_table():
+    """A rename is a per-key fix; hiding it inside a list of twenty-five undoes the report."""
+    cli = _load_cli()
+    rows = [
+        _row("bulk_one", ("a.yaml", "b.yaml", "c.yaml")),
+        _row("bulk_two", ("a.yaml", "b.yaml", "c.yaml")),
+        _row("moved", ("a.yaml", "b.yaml", "c.yaml"), suggestion="moved_elsewhere"),
+    ]
+
+    renamed, merged = cli.group_rows(rows)
+
+    assert [g.keys for g in renamed] == [("moved",)]
+    assert renamed[0].suggestion == "moved_elsewhere"
+    assert [g.keys for g in merged] == [("bulk_one", "bulk_two")]
+
+
+def test_group_rows_does_not_merge_across_a_differing_config_set():
+    cli = _load_cli()
+    rows = [_row("here", ("a.yaml",)), _row("there", ("b.yaml",))]
+
+    assert [g.keys for g in cli.group_rows(rows)[1]] == [("here",), ("there",)]
+
+
+def test_group_rows_ignores_the_order_it_is_given():
+    """CI diffs the summary, so the same findings must render the same way every run."""
+    cli = _load_cli()
+    rows = [_row(k, ("a.yaml", "b.yaml")) for k in ("m", "b", "z")] + [_row("x", ("c.yaml",))]
+
+    def keys(found):
+        return [[g.keys for g in half] for half in found]
+
+    assert keys(cli.group_rows(rows)) == keys(cli.group_rows(rows[::-1]))
+
+
+def test_shorten_paths_drops_the_shared_prefix_but_keeps_what_disambiguates():
+    cli = _load_cli()
+    short = cli.shorten_paths(
+        [
+            "primus/configs/models/megatron/deepseek_v2.yaml",
+            "primus/configs/modules/megatron/pre_trainer.yaml",
+            "examples/megatron/configs/MI300X/shared.yaml",
+            "examples/megatron/configs/MI355X/shared.yaml",
+        ]
+    )
+
+    assert short["primus/configs/models/megatron/deepseek_v2.yaml"] == "deepseek_v2.yaml"
+    assert short["primus/configs/modules/megatron/pre_trainer.yaml"] == "pre_trainer.yaml"
+    # `shared.yaml` alone would name two different configs, so the tail grows.
+    assert short["examples/megatron/configs/MI300X/shared.yaml"] == "MI300X/shared.yaml"
+    assert short["examples/megatron/configs/MI355X/shared.yaml"] == "MI355X/shared.yaml"
+
+
+def test_shorten_paths_falls_back_to_the_whole_path_when_a_tail_never_disambiguates():
+    """A path that is another one's suffix runs out of segments before it runs out of rivals."""
+    cli = _load_cli()
+
+    short = cli.shorten_paths(["a/b.yaml", "x/a/b.yaml"])
+
+    assert short == {"a/b.yaml": "a/b.yaml", "x/a/b.yaml": "x/a/b.yaml"}
+
+
+_RENAMED_TABLE = "### Keys that look renamed rather than dropped"
+_MERGED_TABLE = "### Keys with no obvious replacement"
+
+
+def _drift_out(fake_repo: Path, monkeypatch, capsys, overrides: str, files: int = 1) -> str:
+    for i in range(files):
+        _write(fake_repo / f"examples/torchtitan/configs/probe{i}.yaml", _exp_yaml(overrides))
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", fake_repo)
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "torchtitan", "--warn-only"])
+    assert cli.main() == 0
+    return capsys.readouterr().out
+
+
+def test_cli_prints_every_grouped_key_as_plain_text(fake_repo: Path, monkeypatch, capsys):
+    """CI logs get grepped for a key name, so merging rows must not hide one."""
+    out = _drift_out(fake_repo, monkeypatch, capsys, "training:\n  seed: 1\n  gone_a: 1\n  gone_b: 2\n", 3)
+
+    assert _RENAMED_TABLE in out and _MERGED_TABLE in out
+    assert "`training.gone_a`, `training.gone_b`" in out  # merged onto one row, both still named
+    assert "| `training.seed` | 3 | `debug.seed` |" in out  # the rename keeps its own row and column
+    assert "2 distinct set(s) of configs" not in out  # the two share one set
+    assert "The other 2 drifted out of only 1 distinct set(s) of configs" in out
+
+
+def test_cli_omits_the_rename_table_when_nothing_looks_renamed(fake_repo: Path, monkeypatch, capsys):
+    """An empty table is a heading and a header row saying nothing; print neither."""
+    out = _drift_out(fake_repo, monkeypatch, capsys, "training:\n  gone_a: 1\n")
+
+    assert _RENAMED_TABLE not in out
+    assert "Likely moved to" not in out  # the column only ever holds a suggestion
+    assert _MERGED_TABLE in out
+    assert "All 1 drifted out of only 1 distinct set(s) of configs" in out  # not "The other 1"
+    assert "`training.gone_a`" in out
+
+
+def test_cli_omits_the_merged_table_when_everything_looks_renamed(fake_repo: Path, monkeypatch, capsys):
+    out = _drift_out(fake_repo, monkeypatch, capsys, "training:\n  seed: 1\n")
+
+    assert _MERGED_TABLE not in out
+    assert "distinct set(s) of configs" not in out
+    assert _RENAMED_TABLE in out
+    assert "| `training.seed` | 1 | `debug.seed` |" in out
+
+
+def test_cli_fails_when_no_submodule_is_checked_out(tmp_path: Path, monkeypatch, capsys):
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", tmp_path)
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py"])
+
+    assert cli.main() == 1
+    out = capsys.readouterr().out
+    assert "backend(s) not checked" in out and "git submodule update --init" in out
+
+
+def test_cli_fails_on_a_skipped_backend_even_when_warning_only(fake_repo: Path, monkeypatch, capsys):
+    """The bug this exists for: a partial checkout used to pass in silence.
+
+    ``fake_repo`` has TorchTitan and MaxText but no Megatron, which is exactly
+    what a CI job that fetches two of the three submodules produces.
+    """
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", fake_repo)
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--warn-only"])
+
+    assert cli.main() == 1
+    out = capsys.readouterr().out
+    assert "FAILED -- 1 backend(s) not checked" in out
+    assert "**megatron**" in out
+    assert "No drift" in out  # the two backends that were checked are clean
+
+    # ... while the same two backends on their own still pass.
+    monkeypatch.setattr(
+        "sys.argv",
+        ["check_config_schema.py", "--warn-only", "--backend", "torchtitan", "--backend", "maxtext"],
+    )
+    assert cli.main() == 0
+
+
+def test_cli_warn_only_suppresses_the_failure(fake_repo: Path, monkeypatch, capsys):
+    _write(
+        fake_repo / "examples/torchtitan/configs/probe.yaml",
+        _exp_yaml("training:\n  seed: 42\n"),
+    )
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", fake_repo)
+
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "torchtitan"])
+    assert cli.main() == 1
+    assert "`training.seed`" in capsys.readouterr().out
+
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "torchtitan", "--warn-only"])
+    assert cli.main() == 0
+
+
+def test_cli_reports_a_misplaced_key_apart_from_drift(megatron_repo: Path, monkeypatch, capsys):
+    """The two need different fixes, so they must not share a table."""
+    _write(
+        megatron_repo / "examples/megatron/configs/stray.yaml",
+        """
+        work_group: amd
+        modules:
+          pre_trainer:
+            framework: megatron
+            config: pre_trainer.yaml
+            model: plain.yaml
+            overrides:
+              hash_routing_seed: 7
+        """,
+    )
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", megatron_repo)
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "megatron"])
+
+    assert cli.main() == 1
+    out = capsys.readouterr().out
+    assert "No drift: every key in" in out  # the key exists; it is only misplaced
+    assert "Keys set on a model that cannot read them" in out
+    assert "`hash_routing_seed`" in out and "model_type: fake_v4" in out
+
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "megatron", "--warn-only"])
+    assert cli.main() == 0
+
+
+def test_cli_fails_when_a_config_could_not_be_loaded(fake_repo: Path, monkeypatch, capsys):
+    """An unchecked config must never be reported as a clean one."""
+    _write(fake_repo / "examples/torchtitan/configs/broken.yaml", "extends:\n  - missing.yaml\n")
+    cli = _load_cli()
+    monkeypatch.setattr(cli, "ROOT", fake_repo)
+
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "torchtitan"])
+    assert cli.main() == 1
+    out = capsys.readouterr().out
+    assert "could not be loaded and were NOT checked" in out
+    assert "broken.yaml" in out
+    assert "No drift: every key in" not in out
+    assert "Checked 1 of 2 config(s)" in out
+
+    monkeypatch.setattr("sys.argv", ["check_config_schema.py", "--backend", "torchtitan", "--warn-only"])
+    assert cli.main() == 0
diff --git a/tests/unit_tests/configs/test_deepseek_v4_yaml.py b/tests/unit_tests/configs/test_deepseek_v4_yaml.py
new file mode 100644
index 000000000..c90cac0d3
--- /dev/null
+++ b/tests/unit_tests/configs/test_deepseek_v4_yaml.py
@@ -0,0 +1,417 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-2 P18 — DeepSeek-V4 YAML schema gate (G1).
+
+Each V4 yaml under ``primus/configs/models/megatron/deepseek_v4_*.yaml``
+must:
+
+* parse via the standard primus :func:`parse_yaml` loader (extends +
+  env resolution included);
+* construct a :class:`DeepSeekV4TransformerConfig` after the standard
+  schema massage (no CRIT-level KeyError on plan-2 fields);
+* surface ``compress_ratios`` as a tuple of ints after
+  ``__post_init__`` runs (P18 D4 — the YAML may carry the legacy string
+  form ``"[0, 0, 4, ...]"`` or a real list, and the dataclass must
+  normalize both);
+* not carry retired fields (``v4_use_custom_mtp_block`` /
+  ``mtp_compress_ratios`` were dropped in plan-2 P17);
+* not silently lose the V4-specific MoE / HC / sliding-window / sink
+  fields the schema relies on.
+
+Plus a single provider-singleton check (P18 D1):
+:func:`resolve_v4_provider` returns the same instance on repeated
+calls within the same config.
+"""
+
+from __future__ import annotations
+
+from dataclasses import fields
+from pathlib import Path
+from typing import Any, Dict
+
+import pytest
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_YAML_DIR = _REPO_ROOT / "primus" / "configs" / "models" / "megatron"
+
+
+# ---------------------------------------------------------------------------
+# Helpers — schema-friendly subset of the V4 config kwargs
+# ---------------------------------------------------------------------------
+
+
+def _config_kwargs_from_yaml(yaml_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """Filter the merged YAML dict down to keys consumed by
+    :class:`DeepSeekV4TransformerConfig` (and its parents).
+
+    The full primus config tree mixes train / data / launcher fields
+    with the model-config fields we care about here. The dataclass
+    ignores unknown keys via Python's normal dataclass rules — but
+    only when we hand it kwargs that exist on the dataclass. We do
+    that by intersecting with ``DeepSeekV4TransformerConfig``'s
+    ``__dataclass_fields__``.
+    """
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    valid = set(DeepSeekV4TransformerConfig.__dataclass_fields__.keys())
+    return {k: v for k, v in yaml_dict.items() if k in valid}
+
+
+def _build_v4_config(yaml_dict: Dict[str, Any]):
+    """Construct a V4 config from a parsed YAML dict.
+
+    Some upstream parents require ``num_attention_heads`` and a few
+    other fields. We rely on the YAML to provide them; missing fields
+    fall back to dataclass defaults.
+    """
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    kwargs = _config_kwargs_from_yaml(yaml_dict)
+    # The model YAML uses ``num_experts`` (mapped to Megatron's dataclass field
+    # ``num_moe_experts`` by the primus config loader at runtime); apply the same
+    # alias here so the standalone build is a valid MoE config.
+    if "num_experts" in yaml_dict and "num_moe_experts" not in kwargs:
+        kwargs["num_moe_experts"] = yaml_dict["num_experts"]
+    # The model YAML inherits ``sequence_parallel: true`` (for runtime TP>1),
+    # but this standalone build uses tensor_model_parallel_size=1, which
+    # Megatron rejects ("sequence parallelism without tensor parallelism").
+    kwargs["sequence_parallel"] = False
+    # V4 uses the sqrtsoftplus score function, but Megatron only allows the
+    # aux-loss-free expert bias with sigmoid. The runner disables expert bias
+    # for V4; do the same here. None of these fields affect compress_ratios,
+    # which is the only thing these tests validate.
+    kwargs["moe_router_enable_expert_bias"] = False
+    return DeepSeekV4TransformerConfig(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# YAML parsing — base / flash / pro all parse and merge cleanly
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def parse_yaml_fn():
+    """Bind the shared loader once for all tests in this module."""
+    from primus.core.config.yaml_loader import parse_yaml
+
+    return parse_yaml
+
+
+@pytest.mark.parametrize(
+    "yaml_name",
+    ["deepseek_v4_base.yaml", "deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"],
+)
+def test_v4_yaml_parses(parse_yaml_fn, yaml_name: str) -> None:
+    parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+    assert isinstance(parsed, dict)
+    # Every V4 yaml must declare its core shape.
+    for required in ("num_layers", "hidden_size", "num_attention_heads"):
+        # base.yaml inherits these from elsewhere, so this is best-effort.
+        if yaml_name == "deepseek_v4_base.yaml":
+            continue
+        assert required in parsed, f"{yaml_name} missing required key {required!r}"
+
+
+# ---------------------------------------------------------------------------
+# compress_ratios is normalized to tuple[int, ...] (P18 D4)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "yaml_name",
+    ["deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"],
+)
+def test_compress_ratios_normalized_to_tuple(parse_yaml_fn, yaml_name: str) -> None:
+    """The dataclass ``__post_init__`` must convert
+    ``compress_ratios`` (which may arrive as a YAML string or a list)
+    into ``tuple[int, ...]``.
+
+    The base yaml does not pin a schedule (it is provided per-variant);
+    flash + pro do. Both must round-trip to a tuple of ints with no
+    string survivors and no value drift vs the raw schedule.
+    """
+    parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+    raw = parsed["compress_ratios"]
+
+    # Sanity: the YAML form is either a string or a list.
+    assert isinstance(raw, (str, list, tuple))
+
+    cfg = _build_v4_config(parsed)
+    normalized = cfg.compress_ratios
+
+    assert normalized is not None
+    assert isinstance(
+        normalized, tuple
+    ), f"{yaml_name}: compress_ratios should normalize to tuple, got {type(normalized).__name__}"
+    for i, r in enumerate(normalized):
+        assert isinstance(r, int), f"{yaml_name}: compress_ratios[{i}]={r!r} is not int"
+
+    # Length / values: strip the YAML form and compare to the parsed-int form.
+    if isinstance(raw, str):
+        import ast as _ast
+
+        raw_list = _ast.literal_eval(raw)
+    else:
+        raw_list = list(raw)
+    assert (
+        tuple(int(x) for x in raw_list) == normalized
+    ), f"{yaml_name}: compress_ratios value drift — yaml={raw_list}, normalized={normalized}"
+
+
+def test_compress_ratios_canonical_dispatch_values_only(parse_yaml_fn) -> None:
+    """Plan-2 P17 fixed the comment inversion in V4 yamls; here we
+    enforce the *value* contract: every entry must be one of
+    ``{0, 4, 128}`` (V4 attention only knows these branches —
+    anything else is a schema bug)."""
+    for yaml_name in ("deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"):
+        parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+        cfg = _build_v4_config(parsed)
+        bad = [(i, r) for i, r in enumerate(cfg.compress_ratios) if r not in (0, 4, 128)]
+        assert not bad, (
+            f"{yaml_name}: compress_ratios contains non-canonical values "
+            f"(allowed: 0=dense+SWA, 4=CSA, 128=HCA); offenders: {bad}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Retired schema fields (P17 hygiene)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "retired_field",
+    ["v4_use_custom_mtp_block", "mtp_compress_ratios"],
+)
+def test_retired_fields_not_in_v4_config(retired_field: str) -> None:
+    """Plan-2 P17 deleted the legacy MTP block; both gating fields
+    must be removed from ``DeepSeekV4TransformerConfig``."""
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    names = {f.name for f in fields(DeepSeekV4TransformerConfig)}
+    assert (
+        retired_field not in names
+    ), f"{retired_field} must be removed from DeepSeekV4TransformerConfig (plan-2 P17)."
+
+
+@pytest.mark.parametrize(
+    "yaml_name",
+    ["deepseek_v4_base.yaml", "deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"],
+)
+def test_yaml_does_not_set_retired_fields(parse_yaml_fn, yaml_name: str) -> None:
+    """If a V4 YAML still references a retired field (e.g. someone
+    forgot to update a downstream yaml) we want a loud schema error,
+    not silent acceptance.
+
+    The dataclass would now reject these as ``TypeError`` (unexpected
+    keyword), but we filter unknown keys for forward-compat in
+    ``_config_kwargs_from_yaml``; here we explicitly enforce the
+    contract on the parsed dict instead.
+    """
+    parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+    for retired in ("v4_use_custom_mtp_block", "mtp_compress_ratios"):
+        assert retired not in parsed, (
+            f"{yaml_name} still references retired field {retired!r} "
+            "(plan-2 P17). Remove the line or update training scripts."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Compressed-branch RoPE base is variant-specific
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("yaml_name", "expected_theta"),
+    [
+        # The released DeepSeek-V4-Flash inference/model.py ships
+        # ``compress_rope_theta: float = 40000.0``. 160000 is the Pro value, and
+        # inheriting it in the Flash config was the bug this pins down.
+        ("deepseek_v4_flash.yaml", 40000.0),
+        ("deepseek_v4_pro.yaml", 160000.0),
+        ("deepseek_v4_base.yaml", 160000.0),
+    ],
+)
+def test_compress_rope_theta_matches_variant(parse_yaml_fn, yaml_name: str, expected_theta: float) -> None:
+    """Each variant pins the compressed-branch RoPE base it was trained with.
+
+    Getting this wrong silently de-tunes every CSA / HCA layer's positional
+    phase, so it is worth a hard assertion rather than relying on inheritance.
+    """
+    parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+    assert float(parsed["compress_rope_theta"]) == pytest.approx(expected_theta)
+
+
+@pytest.mark.parametrize(
+    ("yaml_name", "expected_theta", "expected_factor"),
+    [
+        # The upstream open-source Flash recipe ships these as a pair
+        # (compressed rotary base 40000 + rotary_scaling_factor 4).
+        ("deepseek_v4_flash.yaml", 40000.0, 4.0),
+        ("deepseek_v4_pro.yaml", 160000.0, 16.0),
+        ("deepseek_v4_base.yaml", 160000.0, 16.0),
+    ],
+)
+def test_compressed_rope_base_and_yarn_factor_stay_paired(
+    parse_yaml_fn, yaml_name: str, expected_theta: float, expected_factor: float
+) -> None:
+    """``compress_rope_theta`` and ``rotary_scaling_factor`` parameterise one table.
+
+    ``RoPECache`` builds ``inv_freq`` from the base and then runs YaRN band
+    interpolation over it with ``rotary_scaling_factor``, so overriding the base
+    for a variant without overriding the factor leaves the low-frequency end
+    scaled by the wrong amount -- a partial fix that looks complete. Pin the
+    pair so a future variant cannot drift them apart.
+    """
+    parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name))
+    assert float(parsed["compress_rope_theta"]) == pytest.approx(expected_theta)
+    assert float(parsed["rotary_scaling_factor"]) == pytest.approx(expected_factor)
+
+
+def test_yarn_factor_actually_reaches_the_compressed_rope(parse_yaml_fn) -> None:
+    """The Flash factor must land on the compress cache and nowhere else.
+
+    Guards the assertion above against becoming decorative: the value has to
+    change ``compress_rope.inv_freq`` while leaving the dense-layer table
+    (``main_rope``) untouched.
+    """
+    torch = pytest.importorskip("torch")
+    from primus.backends.megatron.core.transformer.dual_rope import DualRoPE
+
+    parsed = parse_yaml_fn(str(_YAML_DIR / "deepseek_v4_flash.yaml"))
+
+    def _build(factor: float) -> DualRoPE:
+        return DualRoPE(
+            rotary_dim=int(parsed["qk_pos_emb_head_dim"]),
+            rope_theta=float(parsed["rotary_base"]),
+            compress_rope_theta=float(parsed["compress_rope_theta"]),
+            yarn_factor=factor,
+            original_max_position_embeddings=int(parsed["original_max_position_embeddings"]),
+        )
+
+    configured = _build(float(parsed["rotary_scaling_factor"]))
+    pro_factor = _build(16.0)
+
+    assert not torch.allclose(
+        configured.compress_rope.inv_freq, pro_factor.compress_rope.inv_freq
+    ), "the YaRN factor does not affect the compressed inv_freq -- assertion is vacuous"
+    torch.testing.assert_close(configured.main_rope.inv_freq, pro_factor.main_rope.inv_freq, rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# V4-specific schema fields the runtime depends on (D5 / D6 hygiene)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "field_name",
+    [
+        # Hyper-Connection
+        "hc_mult",
+        "hc_eps",
+        # Hybrid attention extras
+        "compress_ratios",
+        "compress_rope_theta",
+        "attn_sliding_window",
+        "attn_sink",
+        # Grouped low-rank O projection
+        "o_groups",
+        "o_lora_rank",
+        # MoE-specific extras
+        "num_hash_layers",
+        "swiglu_limit",
+    ],
+)
+def test_v4_config_carries_runtime_field(field_name: str) -> None:
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    names = {f.name for f in fields(DeepSeekV4TransformerConfig)}
+    assert field_name in names, (
+        f"DeepSeekV4TransformerConfig must declare {field_name!r}; "
+        "removing it silently breaks the V4 runtime."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Provider singleton (P18 D1) — same instance per config
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_v4_provider_caches_per_config():
+    from primus.backends.megatron.core.models.deepseek_v4.build_context import (
+        resolve_v4_provider,
+    )
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    # Build a minimal V4 config; we only need the dataclass instance —
+    # the provider does not touch most fields at construction.
+    cfg_a = DeepSeekV4TransformerConfig(
+        num_layers=4,
+        hidden_size=64,
+        num_attention_heads=4,
+        kv_channels=16,
+    )
+    cfg_b = DeepSeekV4TransformerConfig(
+        num_layers=4,
+        hidden_size=64,
+        num_attention_heads=4,
+        kv_channels=16,
+    )
+
+    p_a1 = resolve_v4_provider(cfg_a)
+    p_a2 = resolve_v4_provider(cfg_a)
+    p_b1 = resolve_v4_provider(cfg_b)
+
+    assert p_a1 is p_a2, "resolve_v4_provider must reuse the cached provider for the same config (P18 D1)."
+    assert p_a1 is not p_b1, (
+        "Different config instances should each get their own provider so "
+        "test isolation and runtime overrides are not silently shared."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Provider activation_func helper (P18 D2)
+# ---------------------------------------------------------------------------
+
+
+def test_v4_mlp_activation_func_respects_use_te_activation_func() -> None:
+    """``provider.v4_mlp_activation_func()`` returns ``None`` when the
+    config keeps Megatron's eager activation path (the V4 default,
+    needed for clamped-SwiGLU); only when the user opts into TE does
+    the spec slot carry the TE class.
+    """
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+        TEActivationOp,
+    )
+
+    class _FakeCfg:
+        use_te_activation_func = False
+
+    p = DeepSeekV4SpecProvider(config=_FakeCfg())
+    assert p.v4_mlp_activation_func() is None, (
+        "Default V4 path must leave activation_func slot empty so Megatron "
+        "MLP uses config.activation_func (clamped-SwiGLU)."
+    )
+
+    class _TeCfg:
+        use_te_activation_func = True
+
+    p2 = DeepSeekV4SpecProvider(config=_TeCfg())
+    assert p2.v4_mlp_activation_func() is TEActivationOp, (
+        "When the user opts into TE activation, the spec slot must carry "
+        "TEActivationOp so Megatron MLP instantiates the TE-fused path."
+    )
diff --git a/tests/unit_tests/configs/test_example_configs.py b/tests/unit_tests/configs/test_example_configs.py
new file mode 100644
index 000000000..0cae1cce7
--- /dev/null
+++ b/tests/unit_tests/configs/test_example_configs.py
@@ -0,0 +1,93 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Schema smoke test for every shipped experiment recipe under `examples/`.
+
+The trainer E2E suites only launch one model per architecture, so most of the
+~380 example recipes are never loaded by any other test -- a rename in
+`primus/configs/models/`, a broken `extends:` chain or a `${VAR}` without a
+default would ship unnoticed. This test loads all of them through the real
+config stack (env interpolation -> extends merge -> module/model preset merge ->
+experiment overrides) and asserts the resulting experiment is well formed.
+
+It is CPU-only and needs no GPU: `load_primus_config` resolves YAML and presets
+without importing megatron / torchtitan / jax (backends are only imported later,
+by `PrimusRuntime._initialize_adapter`).
+"""
+
+from pathlib import Path
+
+import pytest
+
+from primus.core.config.primus_config import load_primus_config
+from primus.core.utils import file_utils
+
+ROOT = Path(__file__).resolve().parents[3]
+EXAMPLES = ROOT / "examples"
+
+# Backend dir -> the `framework` every experiment under it must declare.
+BACKENDS = ("megatron", "torchtitan", "maxtext", "megatron_bridge")
+
+MODULE_NAMES = {"pre_trainer", "post_trainer"}
+
+# Not an experiment: a Transformer Engine precision-matcher fragment that lives
+# next to the recipes and has no work_group / modules of its own.
+NOT_AN_EXPERIMENT = {
+    "examples/megatron/configs/MI355X/lfm2_8B_A1B-FP8-te-precision.yaml",
+}
+
+
+def _discover():
+    found = []
+    for backend in BACKENDS:
+        for path in sorted((EXAMPLES / backend / "configs").rglob("*.yaml")):
+            rel = path.relative_to(ROOT).as_posix()
+            if rel not in NOT_AN_EXPERIMENT:
+                found.append(rel)
+    return found
+
+
+EXAMPLE_CONFIGS = _discover()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _no_workspace_side_effects():
+    """Keep loading side-effect free.
+
+    `PrimusConfig.__init__` mkdir's `<workspace>/<group>/<user>/<exp_name>`, and
+    most recipes hardcode `workspace: ./output`, so a bare load would litter the
+    repo (and fail outright when `output/` is not writable by the test user).
+    Only directory creation is stubbed; the config resolution under test is
+    untouched.
+    """
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr(file_utils, "create_path_if_not_exists", lambda *args, **kwargs: None)
+        yield
+
+
+def test_every_backend_contributes_configs():
+    """Guard the glob itself: a typo'd path would otherwise make the suite pass
+    with zero recipes checked."""
+    per_backend = {b: sum(1 for c in EXAMPLE_CONFIGS if c.startswith(f"examples/{b}/")) for b in BACKENDS}
+    assert all(per_backend.values()), f"no example configs discovered for some backend: {per_backend}"
+
+
+@pytest.mark.parametrize("rel_path", EXAMPLE_CONFIGS, ids=EXAMPLE_CONFIGS)
+def test_example_config_loads(rel_path):
+    cfg = load_primus_config(ROOT / rel_path, None)
+
+    modules = getattr(cfg, "modules", [])
+    assert modules, f"{rel_path}: no trainer module resolved"
+
+    backend = rel_path.split("/")[1]
+    for module in modules:
+        assert module.name in MODULE_NAMES, f"{rel_path}: unexpected module '{module.name}'"
+        # A recipe filed under examples/<backend>/ that declares another
+        # framework would be launched by the wrong adapter.
+        assert module.framework == backend, (
+            f"{rel_path}: declares framework '{module.framework}' " f"but lives under examples/{backend}/"
+        )
+        assert hasattr(module, "params"), f"{rel_path}: module '{module.name}' resolved no params"
diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
index cc33c5982..f3366df71 100644
--- a/tests/unit_tests/conftest.py
+++ b/tests/unit_tests/conftest.py
@@ -10,6 +10,20 @@
 import sys
 from pathlib import Path
 
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--run-slow",
+        action="store_true",
+        default=False,
+        help=(
+            "run tests marked with @pytest.mark.slow (e.g. plan-4 release-tier "
+            "shape gates at head_dim=512). Default: skip slow tests."
+        ),
+    )
+
 
 def pytest_configure(config):
     # Add project root first to ensure main primus package takes precedence
@@ -23,3 +37,35 @@ def pytest_configure(config):
         megatron_path = project_root / "third_party" / "Megatron-LM"
     if str(megatron_path) not in sys.path:
         sys.path.append(str(megatron_path))
+
+    # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools,
+    # no __init__.py) that only import with the source root on sys.path. Insert the
+    # submodule at the FRONT so it wins over any stale torchtitan in site-packages.
+    torchtitan_path = os.environ.get("TORCHTITAN_PATH")
+    if torchtitan_path is None or not os.path.exists(torchtitan_path):
+        torchtitan_path = project_root / "third_party" / "torchtitan"
+    if str(torchtitan_path) not in sys.path:
+        sys.path.insert(0, str(torchtitan_path))
+
+    # Register custom markers used by the primus test suite.
+    config.addinivalue_line(
+        "markers",
+        (
+            "slow: marks tests as slow (release-tier shape gates that exercise "
+            "production V4 dims at head_dim=512). Skipped by default; opt in "
+            "with --run-slow or '-m slow'."
+        ),
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    # If the user explicitly opts in via --run-slow or '-m slow', run slow tests.
+    if config.getoption("--run-slow", default=False):
+        return
+    marker_expr = config.getoption("-m") or ""
+    if "slow" in marker_expr:
+        return
+    skip_slow = pytest.mark.skip(reason="slow test (use --run-slow or '-m slow' to enable)")
+    for item in items:
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)
diff --git a/tests/unit_tests/core/__init__.py b/tests/unit_tests/core/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/backend/__init__.py b/tests/unit_tests/core/backend/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/backend/test_backend_adapter.py b/tests/unit_tests/core/backend/test_backend_adapter.py
index ccc4d693c..1a4a18c9c 100644
--- a/tests/unit_tests/core/backend/test_backend_adapter.py
+++ b/tests/unit_tests/core/backend/test_backend_adapter.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -70,7 +70,7 @@ def convert_config(self, params: Any) -> Any:
             model_name=params.get("model") if isinstance(params, dict) else getattr(params, "model", None),
         )
 
-    def load_trainer_class(self, stage: str = "pretrain"):
+    def load_trainer_class(self, stage: str = "pretrain", trainer_class=None):
         self.load_trainer_calls += 1
         return DummyTrainer
 
@@ -99,21 +99,6 @@ def primus_config():
     return SimpleNamespace(exp_name="unit-test-exp")
 
 
-def test_create_trainer_orchestrates_flow(monkeypatch, primus_config, module_config):
-    adapter = DummyBackendAdapter(framework="megatron", version="1.2.3")
-
-    # Abstract methods were called exactly once
-    adapter.prepare_backend(module_config)
-    adapter.convert_config(module_config.params)
-    adapter.load_trainer_class(stage="pretrain")
-    adapter.detect_backend_version()
-
-    assert adapter.prepare_calls == [module_config]
-    assert adapter.convert_calls == [module_config.params]
-    assert adapter.load_trainer_calls == 1
-    assert adapter.detect_version_calls == 1
-
-
 def test_adapter_setup_backend_path_with_explicit_path(tmp_path, monkeypatch):
     adapter = DummyBackendAdapter(framework="test_backend")
     backend_dir = tmp_path / "explicit_backend"
diff --git a/tests/unit_tests/core/backend/test_backend_registry.py b/tests/unit_tests/core/backend/test_backend_registry.py
index 42ea02a15..7b08d553f 100644
--- a/tests/unit_tests/core/backend/test_backend_registry.py
+++ b/tests/unit_tests/core/backend/test_backend_registry.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -15,11 +15,6 @@
 import primus.core.backend.backend_registry as registry_module
 from primus.core.backend.backend_adapter import BackendAdapter
 
-_SUPPORTS_TRAINER_CLASS_REGISTRY = all(
-    hasattr(registry_module.BackendRegistry, attr)
-    for attr in ("_trainer_classes", "register_trainer_class", "get_trainer_class", "has_trainer_class")
-)
-
 
 class MockAdapter(BackendAdapter):
     """Mock adapter for testing."""
@@ -45,12 +40,6 @@ def setup_method(self):
         # Save original state
         self._original_adapters = registry_module.BackendRegistry._adapters.copy()
         registry_module.BackendRegistry._adapters.clear()
-        if hasattr(registry_module.BackendRegistry, "_trainer_classes"):
-            self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy()
-            registry_module.BackendRegistry._trainer_classes.clear()
-        else:
-            self._original_trainer_classes = None
-
         # Silence logging dependencies (logger may not be initialized in tests)
         self._orig_log_rank_0 = registry_module.log_rank_0
         registry_module.log_rank_0 = lambda *args, **kwargs: None
@@ -58,11 +47,6 @@ def setup_method(self):
     def teardown_method(self):
         """Restore registry after each test."""
         registry_module.BackendRegistry._adapters = self._original_adapters
-        if (
-            hasattr(registry_module.BackendRegistry, "_trainer_classes")
-            and self._original_trainer_classes is not None
-        ):
-            registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes
         registry_module.log_rank_0 = self._orig_log_rank_0
 
     def test_get_adapter_not_found_helpful_error(self):
@@ -76,13 +60,6 @@ def test_get_adapter_not_found_helpful_error(self):
         with pytest.raises(ModuleNotFoundError, match="No module named 'primus.backends.non_existent'"):
             registry_module.BackendRegistry.get_adapter("non_existent")
 
-    def test_get_adapter_empty_registry_error(self):
-        """Test error message when no backends are registered."""
-        # get_adapter first calls _load_backend, which fails with ModuleNotFoundError
-        # when the backend module doesn't exist.
-        with pytest.raises(ModuleNotFoundError, match="No module named 'primus.backends.any_backend'"):
-            registry_module.BackendRegistry.get_adapter("any_backend")
-
     def test_get_adapter_creation_failure(self):
         """Test error handling when adapter creation fails."""
 
@@ -120,11 +97,6 @@ def setup_method(self):
         # Reset adapter registry state
         self._original_adapters = registry_module.BackendRegistry._adapters.copy()
         registry_module.BackendRegistry._adapters.clear()
-        if hasattr(registry_module.BackendRegistry, "_trainer_classes"):
-            self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy()
-            registry_module.BackendRegistry._trainer_classes.clear()
-        else:
-            self._original_trainer_classes = None
 
         # Ensure backend module can be re-imported so that lazy loading
         # re-runs registration even if other tests imported it earlier.
@@ -139,11 +111,6 @@ def setup_method(self):
     def teardown_method(self):
         """Restore registry after each test."""
         registry_module.BackendRegistry._adapters = self._original_adapters
-        if (
-            hasattr(registry_module.BackendRegistry, "_trainer_classes")
-            and self._original_trainer_classes is not None
-        ):
-            registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes
         registry_module.log_rank_0 = self._orig_log_rank_0
 
         # Restore original backend module to avoid impacting other tests
@@ -157,11 +124,6 @@ def test_try_load_backend_non_existent(self):
         with pytest.raises(ImportError):
             registry_module.BackendRegistry._try_load_backend("definitely_not_a_backend")
 
-    def test_try_load_backend_returns_bool(self):
-        """Test that _try_load_backend returns boolean."""
-        with pytest.raises(ImportError):
-            registry_module.BackendRegistry._try_load_backend("non_existent")
-
     def test_get_adapter_with_lazy_loading(self):
         """Test that get_adapter triggers lazy loading."""
         # Don't pre-register, let it lazy load
@@ -189,25 +151,6 @@ def test_list_available_backends(self):
         assert "backend2" in available
         assert len(available) == 2
 
-    def test_list_available_backends_empty(self):
-        """Test listing when no backends registered."""
-        available = registry_module.BackendRegistry.list_available_backends()
-        assert available == []
-
-
-class TestBackendRegistryPathNames:
-    """Deprecated: path-name mapping removed; backend path resolution is owned by adapters."""
-
-    def test_path_name_mapping_removed(self):
-        pytest.skip("BackendRegistry path-name mapping removed; use adapter.third_party_dir_name.")
-
-
-class TestBackendRegistrySetupPath:
-    """Deprecated: setup_backend_path moved to BackendAdapter.setup_backend_path()."""
-
-    def test_setup_backend_path_removed(self):
-        pytest.skip("BackendRegistry.setup_backend_path removed; use adapter.setup_backend_path().")
-
 
 class TestBackendRegistryGetAdapterIntegration:
     """Test get_adapter with automatic path setup."""
@@ -262,95 +205,6 @@ def test_get_adapter_path_not_found_error(self):
             registry_module.BackendRegistry.get_adapter("test_backend", backend_path="/non/existent/path")
 
 
-class TestBackendRegistryHasAdapter:
-    """Test has_adapter functionality."""
-
-    def setup_method(self):
-        """Clear registry before each test."""
-        self._original_adapters = registry_module.BackendRegistry._adapters.copy()
-        registry_module.BackendRegistry._adapters.clear()
-
-    def teardown_method(self):
-        """Restore registry after each test."""
-        registry_module.BackendRegistry._adapters = self._original_adapters
-
-    def test_has_adapter_true(self):
-        """Test has_adapter returns True for registered adapter."""
-        registry_module.BackendRegistry.register_adapter("test_backend", MockAdapter)
-
-        assert registry_module.BackendRegistry.has_adapter("test_backend") is True
-
-    def test_has_adapter_false(self):
-        """Test has_adapter returns False for non-registered adapter."""
-        assert registry_module.BackendRegistry.has_adapter("non_existent") is False
-
-
-@pytest.mark.skipif(
-    not _SUPPORTS_TRAINER_CLASS_REGISTRY,
-    reason="Trainer class registry is not available on BackendRegistry in this version.",
-)
-class TestBackendRegistryTrainerClasses:
-    """Test trainer class registration and retrieval."""
-
-    def setup_method(self):
-        """Clear trainer classes before each test."""
-        if not hasattr(registry_module.BackendRegistry, "_trainer_classes"):
-            pytest.skip("BackendRegistry has no _trainer_classes in this version.")
-        self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy()
-        registry_module.BackendRegistry._trainer_classes.clear()
-
-    def teardown_method(self):
-        """Restore trainer classes after each test."""
-        registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes
-
-    def test_register_and_get_trainer_class(self):
-        """Test registering and retrieving trainer classes."""
-
-        class DummyTrainer:
-            pass
-
-        registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend")
-
-        trainer_cls = registry_module.BackendRegistry.get_trainer_class("test_backend")
-        assert trainer_cls is DummyTrainer
-
-    def test_get_trainer_class_not_found(self):
-        """Test error when trainer class not registered."""
-        with pytest.raises(ValueError) as exc_info:
-            registry_module.BackendRegistry.get_trainer_class("non_existent_backend")
-
-        assert "No trainer class registered for backend 'non_existent_backend'" in str(exc_info.value)
-
-    def test_has_trainer_class(self):
-        """Test has_trainer_class reflects registration state."""
-
-        class DummyTrainer:
-            pass
-
-        assert registry_module.BackendRegistry.has_trainer_class("test_backend") is False
-        registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend")
-        assert registry_module.BackendRegistry.has_trainer_class("test_backend") is True
-
-    def test_register_and_get_trainer_class_with_stage(self):
-        """Test trainer registration with explicit stage."""
-
-        class DummyTrainer:
-            pass
-
-        # Register with explicit stage "sft"
-        registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend", stage="sft")
-
-        # Get with explicit stage "sft" should work
-        trainer_cls = registry_module.BackendRegistry.get_trainer_class("test_backend", stage="sft")
-        assert trainer_cls is DummyTrainer
-
-        # has_trainer_class with explicit stage should work
-        assert registry_module.BackendRegistry.has_trainer_class("test_backend", stage="sft") is True
-
-        # Default stage "pretrain" should not find it
-        assert registry_module.BackendRegistry.has_trainer_class("test_backend") is False
-
-
 class TestBackendRegistrySetupHooks:
     """Test setup hook registration and execution."""
 
@@ -404,5 +258,67 @@ def failing_hook():
         assert "Error in setup hook" in captured.out
 
 
+class TestBackendRegistryTrainerClass:
+    """Test the trainer-class registry API.
+
+    ``register_trainer_class`` / ``get_trainer_class`` / ``has_trainer_class``
+    back the stage-based fallback in ``MegatronAdapter.load_trainer_class`` (and
+    the other backend adapters), so they must keep (backend, stage) keying and
+    raise on lookups for unregistered combinations.
+    """
+
+    def setup_method(self):
+        self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy()
+        registry_module.BackendRegistry._trainer_classes.clear()
+
+    def teardown_method(self):
+        registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes
+
+    def test_register_and_get_trainer_class(self):
+        class DummyTrainer:
+            pass
+
+        registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "megatron", stage="pretrain")
+
+        assert registry_module.BackendRegistry.get_trainer_class("megatron", stage="pretrain") is DummyTrainer
+        assert registry_module.BackendRegistry.has_trainer_class("megatron", stage="pretrain") is True
+
+    def test_register_defaults_to_pretrain_stage(self):
+        class DummyTrainer:
+            pass
+
+        # Default stage is "pretrain" for both register and lookup.
+        registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "megatron")
+        assert registry_module.BackendRegistry.get_trainer_class("megatron") is DummyTrainer
+
+    def test_stage_is_part_of_the_key(self):
+        class PretrainTrainer:
+            pass
+
+        registry_module.BackendRegistry.register_trainer_class(PretrainTrainer, "megatron", stage="pretrain")
+
+        # A different stage for the same backend must NOT resolve to the
+        # pretrain class; it is a distinct registry key.
+        assert registry_module.BackendRegistry.has_trainer_class("megatron", stage="sft") is False
+        with pytest.raises(ValueError):
+            registry_module.BackendRegistry.get_trainer_class("megatron", stage="sft")
+
+    def test_get_unregistered_trainer_class_raises(self):
+        assert registry_module.BackendRegistry.has_trainer_class("nonexistent") is False
+        with pytest.raises(ValueError, match="No trainer class registered"):
+            registry_module.BackendRegistry.get_trainer_class("nonexistent")
+
+    def test_register_overwrites_same_key(self):
+        class TrainerA:
+            pass
+
+        class TrainerB:
+            pass
+
+        registry_module.BackendRegistry.register_trainer_class(TrainerA, "megatron", stage="pretrain")
+        registry_module.BackendRegistry.register_trainer_class(TrainerB, "megatron", stage="pretrain")
+        assert registry_module.BackendRegistry.get_trainer_class("megatron", stage="pretrain") is TrainerB
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
diff --git a/tests/unit_tests/core/config/test_primus_config.py b/tests/unit_tests/core/config/test_primus_config.py
new file mode 100644
index 000000000..b9e606359
--- /dev/null
+++ b/tests/unit_tests/core/config/test_primus_config.py
@@ -0,0 +1,61 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from pathlib import Path
+
+from primus.core.config.primus_config import get_module_config, load_primus_config
+from tests.utils import PrimusUT
+
+# Top-level attributes that `_normalize_module_for_runtime` keeps in place; every
+# other public attribute is moved into `params` on the normalized copy.
+_RESERVED_KEYS = {"name", "framework", "config", "model", "params", "trainer_class"}
+
+_EXAMPLE_CONFIG = "examples/megatron/exp_pretrain.yaml"
+
+
+class TestLoadPrimusConfig(PrimusUT):
+    def test_exposes_legacy_primus_config(self):
+        """load_primus_config attaches the underlying legacy PrimusConfig as
+        `_legacy` so the core runtime can reuse it without re-parsing."""
+        cfg = load_primus_config(Path(_EXAMPLE_CONFIG), None)
+
+        self.assertTrue(hasattr(cfg, "_legacy"))
+        legacy = cfg._legacy
+        # The legacy object must still provide the PrimusConfig interface that
+        # BaseModule relies on.
+        self.assertTrue(callable(getattr(legacy, "get_module_config", None)))
+
+    def test_legacy_config_is_pristine_after_normalization(self):
+        """Regression guard for the Option A deepcopy nuance: normalization must
+        operate on deep copies, leaving `_legacy`'s module configs un-mutated.
+
+        With the previous in-place normalization, `delattr` stripped non-reserved
+        training params off the legacy module config (they were moved into
+        `params`). That made the exposed legacy config unusable for BaseModule.
+        Here we assert the legacy module configs still carry their original
+        top-level training params.
+        """
+        cfg = load_primus_config(Path(_EXAMPLE_CONFIG), None)
+        legacy = cfg._legacy
+
+        module_keys = list(getattr(legacy, "module_keys", []))
+        self.assertTrue(module_keys, "expected at least one module in example config")
+
+        for name in module_keys:
+            legacy_mod = legacy.get_module_config(name)
+            top_level = {k for k in vars(legacy_mod) if not k.startswith("_")}
+            # The legacy module must retain non-reserved top-level params; if the
+            # normalization had mutated it in place, these would be gone.
+            self.assertTrue(
+                top_level - _RESERVED_KEYS,
+                f"legacy module '{name}' was mutated/stripped by normalization",
+            )
+
+            # The normalized SimpleNamespace copy must still expose those same
+            # params under `.params` for the new runtime.
+            normalized = get_module_config(cfg, name)
+            self.assertIsNotNone(normalized)
+            self.assertTrue(hasattr(normalized, "params"))
diff --git a/tests/unit_tests/core/config/test_yaml_loader.py b/tests/unit_tests/core/config/test_yaml_loader.py
index 038cef09a..4dc7cfd0a 100644
--- a/tests/unit_tests/core/config/test_yaml_loader.py
+++ b/tests/unit_tests/core/config/test_yaml_loader.py
@@ -49,6 +49,24 @@ def test_parse_yaml_extends(self, tmp_path):
         assert result["child_val"] == 3
         assert "extends" not in result  # Should be removed
 
+    def test_parse_yaml_scalar_extends(self, tmp_path):
+        """A lone preset written as a scalar must not be iterated per character."""
+        base = tmp_path / "base.yaml"
+        base.write_text("base_val: 1\noverride_val: 1")
+
+        child = tmp_path / "child.yaml"
+        child.write_text(
+            """
+        extends: base.yaml
+        override_val: 2
+        """
+        )
+
+        result = parse_yaml(str(child))
+        assert result["base_val"] == 1
+        assert result["override_val"] == 2
+        assert "extends" not in result
+
     def test_parse_yaml_multi_extends(self, tmp_path):
         """Test extending multiple files."""
         base1 = tmp_path / "base1.yaml"
@@ -110,3 +128,34 @@ def test_resolve_env_in_string_helpers(self, monkeypatch):
         assert _resolve_env_in_string("${NUM}") == 42
         assert _resolve_env_in_string("${FLOAT}") == 3.14
         assert _resolve_env_in_string("prefix_${NUM}") == "prefix_42"
+
+    def test_env_resolution_bool_strings(self, monkeypatch, tmp_path):
+        monkeypatch.setenv("FLAG_TRUE", "true")
+        monkeypatch.setenv("FLAG_FALSE", "false")
+        monkeypatch.setenv("FLAG_UPPER", "False")
+        monkeypatch.delenv("UNSET_BOOL", raising=False)
+
+        assert _resolve_env_in_string("${FLAG_TRUE}") is True
+        assert _resolve_env_in_string("${FLAG_FALSE}") is False
+        assert _resolve_env_in_string("${FLAG_UPPER}") is False
+        assert _resolve_env_in_string("${UNSET_BOOL:false}") is False
+        assert _resolve_env_in_string("${UNSET_BOOL:true}") is True
+
+        content = """
+        native_false: false
+        env_false_default: ${UNSET_BOOL:false}
+        env_false_set: ${FLAG_FALSE}
+        env_true_set: ${FLAG_TRUE}
+        env_zero: ${UNSET_ZERO:0}
+        """
+        monkeypatch.delenv("UNSET_ZERO", raising=False)
+        cfg_file = tmp_path / "bool_env.yaml"
+        cfg_file.write_text(content)
+
+        result = parse_yaml(str(cfg_file))
+        assert result["native_false"] is False
+        assert result["env_false_default"] is False
+        assert result["env_false_set"] is False
+        assert result["env_true_set"] is True
+        assert result["env_zero"] == 0
+        assert isinstance(result["env_zero"], int)
diff --git a/tests/unit_tests/core/launcher/__init__.py b/tests/unit_tests/core/launcher/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/patches/__init__.py b/tests/unit_tests/core/patches/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/runtime/__init__.py b/tests/unit_tests/core/runtime/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/runtime/test_train_runtime.py b/tests/unit_tests/core/runtime/test_train_runtime.py
index 51fd5b996..abf257078 100644
--- a/tests/unit_tests/core/runtime/test_train_runtime.py
+++ b/tests/unit_tests/core/runtime/test_train_runtime.py
@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -126,7 +126,7 @@ def test_run_trainer_lifecycle_calls_trainer_methods_in_order(self):
 
         # Use a dummy trainer that records the call order.
         class DummyTrainer:
-            def __init__(self, backend_args=None):
+            def __init__(self, backend_args=None, **kwargs):
                 self.calls = []
                 self.backend_args = backend_args
 
@@ -174,7 +174,7 @@ def _fake_run_patches(**kwargs):
         with patch("primus.core.runtime.train_runtime.run_patches", side_effect=_fake_run_patches):
             # Dummy trainer
             class DummyTrainer:
-                def __init__(self, backend_args=None):
+                def __init__(self, backend_args=None, **kwargs):
                     self.backend_args = backend_args
 
                 def setup(self):
@@ -204,5 +204,349 @@ def cleanup(self, on_error: bool = False):
         assert phases == ["build_args", "setup", "before_train", "after_train"]
 
 
+class TestPrimusRuntimeTrainerClassSelection(PrimusUT):
+    """Test trainer class selection from config."""
+
+    def _build_args(self, config: str = "examples/megatron/exp_pretrain.yaml"):
+        """Build args for testing."""
+        import argparse
+
+        return argparse.Namespace(config=config, data_path="./data", backend_path=None)
+
+    def test_initialize_trainer_extracts_trainer_class_from_module_config_top_level(self):
+        """Test that trainer_class is extracted from module_config (top-level)."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        # Create a mock adapter that records calls
+        mock_adapter = Mock()
+        mock_adapter.prepare_backend = Mock()
+        mock_adapter.convert_config = Mock(return_value=SimpleNamespace())
+
+        # Mock trainer class
+        mock_trainer_class = Mock()
+        mock_trainer_class.__name__ = "FluxPretrainTrainer"
+        mock_trainer_instance = Mock()
+        mock_trainer_class.return_value = mock_trainer_instance
+        mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class)
+
+        # Setup context with trainer_class in module_config (top-level)
+        runtime.ctx = SimpleNamespace(
+            config_path=Path("dummy.yaml"),
+            data_path=Path("./data"),
+            module_name="pre_trainer",
+            primus_config=SimpleNamespace(),
+            rank=0,
+            world_size=1,
+            master_addr="localhost",
+            master_port="12345",
+            module_config=SimpleNamespace(
+                framework="megatron",
+                trainer_class="FluxPretrainTrainer",  # Top-level attribute
+                params=SimpleNamespace(stage="pretrain"),
+            ),
+            framework="megatron",
+            adapter=mock_adapter,
+        )
+
+        # Mock patches and logging
+        with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch(
+            "primus.core.runtime.train_runtime.log_rank_0"
+        ) as mock_log, patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object(
+            runtime, "_run_phase_patches"
+        ):
+
+            runtime._initialize_trainer()
+
+        # Verify adapter.load_trainer_class was called with trainer_class
+        mock_adapter.load_trainer_class.assert_called_once_with(
+            stage="pretrain", trainer_class="FluxPretrainTrainer"
+        )
+
+        # Verify logging indicates trainer_class usage
+        log_calls = [str(call) for call in mock_log.call_args_list]
+        assert any("Using trainer_class: FluxPretrainTrainer" in str(call) for call in log_calls)
+
+    def test_initialize_trainer_extracts_trainer_class_from_params_when_not_top_level(self):
+        """Test that trainer_class is extracted from params when not in top-level."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        mock_adapter = Mock()
+        mock_adapter.prepare_backend = Mock()
+        # Critically, backend_args does NOT carry trainer_class. This isolates the
+        # params-extraction branch: the only way trainer_class can be resolved is
+        # from module_config.params (the post-merge backend_args fallback cannot
+        # mask a regression in that branch).
+        backend_args = SimpleNamespace(stage="pretrain")
+        mock_adapter.convert_config = Mock(return_value=backend_args)
+        mock_trainer_class = Mock()
+        mock_trainer_class.__name__ = "FluxPretrainTrainer"
+        mock_trainer_class.return_value = Mock()
+        mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class)
+
+        # trainer_class in params, NOT top-level (no top-level attribute)
+        runtime.ctx = SimpleNamespace(
+            config_path=Path("dummy.yaml"),
+            data_path=Path("./data"),
+            module_name="pre_trainer",
+            primus_config=SimpleNamespace(),
+            rank=0,
+            world_size=1,
+            master_addr="localhost",
+            master_port="12345",
+            module_config=SimpleNamespace(
+                framework="megatron",
+                # No trainer_class attribute here
+                params=SimpleNamespace(
+                    stage="pretrain",
+                    trainer_class="FluxPretrainTrainer",  # In params
+                ),
+            ),
+            framework="megatron",
+            adapter=mock_adapter,
+        )
+
+        with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch(
+            "primus.core.runtime.train_runtime.log_rank_0"
+        ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object(
+            runtime, "_run_phase_patches"
+        ):
+            runtime._initialize_trainer()
+
+        # trainer_class must have been resolved from module_config.params and forwarded.
+        mock_adapter.load_trainer_class.assert_called_once_with(
+            stage="pretrain", trainer_class="FluxPretrainTrainer"
+        )
+
+    def test_initialize_trainer_top_level_takes_precedence_over_params(self):
+        """Test that top-level trainer_class takes precedence over params.trainer_class."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        mock_adapter = Mock()
+        mock_adapter.prepare_backend = Mock()
+        mock_adapter.convert_config = Mock(return_value=SimpleNamespace())
+        mock_trainer_class = Mock()
+        mock_trainer_class.__name__ = "TopLevelTrainer"
+        mock_trainer_class.return_value = Mock()
+        mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class)
+
+        # Both top-level and params have trainer_class (top-level should win due to if/elif)
+        runtime.ctx = SimpleNamespace(
+            config_path=Path("dummy.yaml"),
+            data_path=Path("./data"),
+            module_name="pre_trainer",
+            primus_config=SimpleNamespace(),
+            rank=0,
+            world_size=1,
+            master_addr="localhost",
+            master_port="12345",
+            module_config=SimpleNamespace(
+                framework="megatron",
+                trainer_class="TopLevelTrainer",  # Top-level (checked first)
+                params=SimpleNamespace(
+                    stage="pretrain",
+                    trainer_class="ParamsTrainer",  # In params (should be ignored)
+                ),
+            ),
+            framework="megatron",
+            adapter=mock_adapter,
+        )
+
+        with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch(
+            "primus.core.runtime.train_runtime.log_rank_0"
+        ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object(
+            runtime, "_run_phase_patches"
+        ):
+
+            runtime._initialize_trainer()
+
+        # Verify top-level trainer_class was used (not params)
+        mock_adapter.load_trainer_class.assert_called_once_with(
+            stage="pretrain", trainer_class="TopLevelTrainer"  # Top-level, not ParamsTrainer
+        )
+
+    def test_initialize_trainer_falls_back_to_stage_when_no_trainer_class(self):
+        """Test fallback to stage-based selection when trainer_class not specified."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        mock_adapter = Mock()
+        mock_adapter.prepare_backend = Mock()
+        mock_adapter.convert_config = Mock(return_value=SimpleNamespace())
+        mock_trainer_class = Mock()
+        mock_trainer_class.__name__ = "MockTrainer"
+        mock_trainer_class.return_value = Mock()
+        mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class)
+
+        # No trainer_class specified anywhere
+        runtime.ctx = SimpleNamespace(
+            config_path=Path("dummy.yaml"),
+            data_path=Path("./data"),
+            module_name="pre_trainer",
+            primus_config=SimpleNamespace(),
+            rank=0,
+            world_size=1,
+            master_addr="localhost",
+            master_port="12345",
+            module_config=SimpleNamespace(
+                framework="megatron",
+                params=SimpleNamespace(stage="pretrain"),
+            ),
+            framework="megatron",
+            adapter=mock_adapter,
+        )
+
+        with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch(
+            "primus.core.runtime.train_runtime.log_rank_0"
+        ) as mock_log, patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object(
+            runtime, "_run_phase_patches"
+        ):
+
+            runtime._initialize_trainer()
+
+        # Verify fallback to stage-based selection. When trainer_class is falsy
+        # the kwarg is omitted entirely (see train_runtime.py:390-392), so the
+        # adapter is called with stage only.
+        mock_adapter.load_trainer_class.assert_called_once_with(stage="pretrain")
+
+        # Verify logging indicates fallback
+        log_calls = [str(call) for call in mock_log.call_args_list]
+        assert any("trainer_class not found" in str(call) for call in log_calls)
+
+    def test_initialize_trainer_handles_empty_string_trainer_class(self):
+        """Test that empty string trainer_class falls back to stage."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        mock_adapter = Mock()
+        mock_adapter.prepare_backend = Mock()
+        mock_adapter.convert_config = Mock(return_value=SimpleNamespace())
+        mock_trainer_class = Mock()
+        mock_trainer_class.__name__ = "MockTrainer"
+        mock_trainer_class.return_value = Mock()
+        mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class)
+
+        # Empty string trainer_class (falsy, should fall back)
+        runtime.ctx = SimpleNamespace(
+            config_path=Path("dummy.yaml"),
+            data_path=Path("./data"),
+            module_name="pre_trainer",
+            primus_config=SimpleNamespace(),
+            rank=0,
+            world_size=1,
+            master_addr="localhost",
+            master_port="12345",
+            module_config=SimpleNamespace(
+                framework="megatron",
+                trainer_class="",  # Empty string (falsy)
+                params=SimpleNamespace(stage="pretrain"),
+            ),
+            framework="megatron",
+            adapter=mock_adapter,
+        )
+
+        with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch(
+            "primus.core.runtime.train_runtime.log_rank_0"
+        ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object(
+            runtime, "_run_phase_patches"
+        ):
+
+            runtime._initialize_trainer()
+
+        # Should fall back to stage (empty string is falsy): the trainer_class
+        # kwarg is omitted entirely (see train_runtime.py:390-392).
+        mock_adapter.load_trainer_class.assert_called_once_with(stage="pretrain")
+
+
+class TestPrimusRuntimeLifecycle(PrimusUT):
+    """Tests for PrimusRuntime trainer lifecycle execution."""
+
+    def _build_args(self, config: str = "examples/megatron/exp_pretrain.yaml") -> argparse.Namespace:
+        return argparse.Namespace(config=config, data_path="./data", backend_path=None)
+
+    def test_run_trainer_lifecycle_applies_patches_before_train(self):
+        """Test that patches are applied in before_train phase before train() is called."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        train_called = []
+        patch_applied = []
+
+        class MockTrainer:
+            def setup(self):
+                pass
+
+            def init(self):
+                pass
+
+            def train(self):
+                train_called.append(1)
+                # Verify patch was applied before train
+                assert len(patch_applied) > 0
+
+            def cleanup(self, on_error=False):
+                pass
+
+        mock_trainer = MockTrainer()
+        runtime.ctx = SimpleNamespace(
+            trainer=mock_trainer, backend_args=SimpleNamespace(), runtime_state=None
+        )
+
+        def mock_run_phase_patches(phase, backend_args=None, runtime_state=None):
+            if phase == "before_train":
+                patch_applied.append(1)
+
+        runtime._run_phase_patches = mock_run_phase_patches
+
+        with patch("primus.core.runtime.train_runtime.log_rank_0"):
+            runtime._run_trainer_lifecycle()
+
+        # Verify patch was applied before train
+        assert len(patch_applied) == 1
+        assert len(train_called) == 1
+
+    def test_run_trainer_lifecycle_passes_backend_args_to_patches(self):
+        """Test that backend_args are passed to patch phases."""
+        args = self._build_args()
+        runtime = PrimusRuntime(args=args)
+
+        backend_args_received = []
+
+        class MockTrainer:
+            def setup(self):
+                pass
+
+            def init(self):
+                pass
+
+            def train(self):
+                pass
+
+            def cleanup(self, on_error=False):
+                pass
+
+        mock_trainer = MockTrainer()
+        test_backend_args = SimpleNamespace(test_param="value")
+        runtime.ctx = SimpleNamespace(
+            trainer=mock_trainer,
+            backend_args=test_backend_args,
+            runtime_state=None,
+        )
+
+        def mock_run_phase_patches(phase, backend_args=None, runtime_state=None):
+            backend_args_received.append(backend_args)
+
+        runtime._run_phase_patches = mock_run_phase_patches
+
+        with patch("primus.core.runtime.train_runtime.log_rank_0"):
+            runtime._run_trainer_lifecycle()
+
+        # Verify backend_args were passed to all patch phases
+        assert len(backend_args_received) == 3
+        assert all(args is test_backend_args for args in backend_args_received)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unit_tests/core/trainer/__init__.py b/tests/unit_tests/core/trainer/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/core/trainer/test_base_trainer.py b/tests/unit_tests/core/trainer/test_base_trainer.py
index 6e3b5c4be..b3ee0beab 100644
--- a/tests/unit_tests/core/trainer/test_base_trainer.py
+++ b/tests/unit_tests/core/trainer/test_base_trainer.py
@@ -1,95 +1,117 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
-"""Unit tests for BaseTrainer."""
-
-from types import SimpleNamespace
-
-from primus.core.trainer.base_trainer import BaseTrainer
+"""
+Unit tests for BaseTrainer.
 
+Tests initialization, MRO handling, distributed environment setup, and
+backend_args handling for the universal base trainer class.
+"""
 
-class DummyTrainer(BaseTrainer):
-    """Minimal concrete implementation of BaseTrainer for testing."""
-
-    def __init__(self, backend_args=None):
-        super().__init__(backend_args=backend_args)
-        self.train_calls: int = 0
-        self.setup_calls: int = 0
-        self.init_calls: int = 0
-
-    def init(self, *args, **kwargs):
-        """No-op init for testing."""
-        self.init_calls += 1
-        return None
+from types import SimpleNamespace
 
-    def setup(self, *args, **kwargs):
-        """No-op setup for testing."""
-        self.setup_calls += 1
-        return None
+import pytest
 
-    def train(self):
-        self.train_calls += 1
+from primus.core.trainer.base_trainer import BaseTrainer
 
 
 class TestBaseTrainer:
-    """Verify BaseTrainer interface and lifecycle."""
+    """Tests for BaseTrainer initialization and behavior."""
+
+    def test_init_sets_distributed_environment(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that distributed environment attributes are set correctly."""
+        mock_dist_env = {
+            "rank": 2,
+            "world_size": 8,
+            "local_rank": 1,
+            "master_addr": "192.168.1.1",
+            "master_port": 54321,
+        }
+
+        monkeypatch.setattr(
+            "primus.core.trainer.base_trainer.get_torchrun_env",
+            lambda: mock_dist_env,
+        )
+        monkeypatch.setattr("builtins.print", lambda *args, **kwargs: None)
+        monkeypatch.setattr("sys.stderr.write", lambda *args, **kwargs: None)
+
+        # Create a concrete subclass for testing
+        class ConcreteBaseTrainer(BaseTrainer):
+            def setup(self):
+                pass
+
+            def init(self):
+                pass
+
+            def train(self):
+                pass
+
+        trainer = ConcreteBaseTrainer(backend_args=SimpleNamespace())
+
+        assert trainer.rank == 2
+        assert trainer.world_size == 8
+        assert trainer.local_rank == 1
+        assert trainer.master_addr == "192.168.1.1"
+        assert trainer.master_port == 54321
+
+    def test_init_mro_with_base_module(self, monkeypatch: pytest.MonkeyPatch):
+        """Test MRO handling when BaseModule IS in inheritance chain (legacy pattern)."""
+        from primus.core.base_module import BaseModule
+
+        # Create a class that inherits from both BaseTrainer and BaseModule
+        class LegacyTrainer(BaseTrainer, BaseModule):
+            def setup(self):
+                pass
+
+            def init(self):
+                pass
+
+            def train(self):
+                pass
+
+            def run(self):
+                pass  # BaseModule requires run() method
+
+        mock_dist_env = {
+            "rank": 0,
+            "world_size": 1,
+            "local_rank": 0,
+            "master_addr": "localhost",
+            "master_port": 12345,
+        }
+
+        monkeypatch.setattr(
+            "primus.core.trainer.base_trainer.get_torchrun_env",
+            lambda: mock_dist_env,
+        )
+        monkeypatch.setattr("builtins.print", lambda *args, **kwargs: None)
+        monkeypatch.setattr("sys.stderr.write", lambda *args, **kwargs: None)
+
+        # Mock BaseModule.__init__ to track if it's called with kwargs and provide required args
+        base_module_init_called = []
+
+        def tracked_init(
+            self, module_name="test", primus_config=None, module_rank=0, module_world_size=1, **kwargs
+        ):
+            base_module_init_called.append(("kwargs", kwargs))
+            # Don't call original_init, just track the call
+
+        monkeypatch.setattr(BaseModule, "__init__", tracked_init)
 
-    def test_trainer_stores_backend_args(self):
-        """Backend args should be stored on the trainer instance."""
-        backend_args = SimpleNamespace(lr=1e-4, batch_size=32)
-        trainer = DummyTrainer(backend_args=backend_args)
-
-        assert trainer.backend_args is backend_args
-        assert trainer.backend_args.lr == 1e-4
-        assert trainer.backend_args.batch_size == 32
-
-    def test_trainer_allows_none_backend_args(self):
-        """Trainer should accept None as backend_args."""
-        trainer = DummyTrainer(backend_args=None)
-        assert trainer.backend_args is None
-
-    def test_trainer_train_method_executes(self):
-        """The train() method should execute correctly."""
-        backend_args = SimpleNamespace(lr=1e-4)
-        trainer = DummyTrainer(backend_args=backend_args)
-
-        trainer.train()
-
-        assert trainer.train_calls == 1
-
-    def test_trainer_lifecycle_methods_exist(self):
-        """Trainer should have setup, init, train, cleanup methods."""
         backend_args = SimpleNamespace()
-        trainer = DummyTrainer(backend_args=backend_args)
-
-        # All lifecycle methods should be callable
-        trainer.setup()
-        trainer.init()
-        trainer.train()
-        trainer.cleanup()
-
-        assert trainer.setup_calls == 1
-        assert trainer.init_calls == 1
-        assert trainer.train_calls == 1
-
-    def test_cleanup_accepts_on_error_flag(self):
-        """cleanup() should accept on_error parameter."""
-        trainer = DummyTrainer(backend_args=None)
-
-        # Should not raise
-        trainer.cleanup(on_error=False)
-        trainer.cleanup(on_error=True)
-
-    def test_trainer_has_distributed_env_attributes(self):
-        """Trainer should expose distributed environment attributes."""
-        trainer = DummyTrainer(backend_args=None)
-
-        # These attributes should exist (values depend on env)
-        assert hasattr(trainer, "rank")
-        assert hasattr(trainer, "world_size")
-        assert hasattr(trainer, "local_rank")
-        assert hasattr(trainer, "master_addr")
-        assert hasattr(trainer, "master_port")
+        trainer = LegacyTrainer(
+            backend_args=backend_args,
+            some_kwarg="value",
+            module_name="test",
+            primus_config=SimpleNamespace(),
+            module_rank=0,
+            module_world_size=1,
+        )
+
+        # Verify BaseModule.__init__ was called with kwargs
+        assert len(base_module_init_called) > 0
+        assert "some_kwarg" in base_module_init_called[0][1]
+        assert trainer.backend_args is backend_args
diff --git a/tests/unit_tests/core/utils/__init__.py b/tests/unit_tests/core/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/examples/test_deepseek_v4_run_script_defaults.py b/tests/unit_tests/examples/test_deepseek_v4_run_script_defaults.py
new file mode 100644
index 000000000..cf0ac0c66
--- /dev/null
+++ b/tests/unit_tests/examples/test_deepseek_v4_run_script_defaults.py
@@ -0,0 +1,73 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Default-value gate for the DeepSeek-V4 run scripts.
+
+The Flash launcher overrides a handful of knobs on top of ``run_deepseek_v4.sh``.
+One of them -- ``USE_V4_FP8_INDEXER`` -- silently changes model numerics rather
+than just performance: the indexer picks which compressed KV entries each query
+attends to, so quantizing its QK inputs changes the selection itself. On top of
+that the knob is a fake quant (quantize/dequantize around a BF16 GEMM), so it
+costs throughput without buying any. Low-precision indexer QK is a legitimate
+thing to explore, but it should be opt-in rather than on by default in one
+launcher only.
+
+These tests parse the scripts as text: they need neither torch nor a GPU, so
+they run in any environment.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pytest
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_EXAMPLES = _REPO_ROOT / "examples" / "deepseek-v4"
+
+_FLASH_SCRIPT = _EXAMPLES / "run_deepseek_v4_flash.sh"
+_BASE_SCRIPT = _EXAMPLES / "run_deepseek_v4.sh"
+
+
+def _default_of(script: Path, var: str) -> str | None:
+    """Return the default in ``export VAR=${VAR:-default}``, or None.
+
+    The last assignment wins, mirroring shell semantics.
+    """
+    pattern = re.compile(rf"^\s*export\s+{re.escape(var)}=\$\{{{re.escape(var)}:-([^}}]*)\}}", re.MULTILINE)
+    matches = pattern.findall(script.read_text(encoding="utf-8"))
+    return matches[-1] if matches else None
+
+
+@pytest.mark.parametrize("script", [_FLASH_SCRIPT, _BASE_SCRIPT], ids=["flash", "base"])
+def test_fp8_indexer_defaults_off(script: Path) -> None:
+    """Both launchers must leave the indexer QK in high precision by default."""
+    assert script.is_file(), f"missing run script: {script}"
+    default = _default_of(script, "USE_V4_FP8_INDEXER")
+    assert default is not None, f"{script.name} must define USE_V4_FP8_INDEXER with a ${{VAR:-...}} default"
+    assert default == "False", (
+        f"{script.name} defaults USE_V4_FP8_INDEXER to {default!r}. The indexer selects which "
+        "compressed KV entries each query sees, so enabling the fake-quant path by default "
+        "changes model numerics -- and costs throughput while doing so, since the GEMM stays "
+        "BF16. Keep it opt-in."
+    )
+
+
+def test_fp8_indexer_stays_overridable() -> None:
+    """The knob keeps the ``${VAR:-default}`` form so QAT runs can opt back in."""
+    text = _FLASH_SCRIPT.read_text(encoding="utf-8")
+    assert "USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-" in text
+
+
+def test_flash_defaults_agree_with_base_script() -> None:
+    """Flash must not silently disagree with the generic launcher on this knob.
+
+    The Flash script is a thin wrapper that execs ``run_deepseek_v4.sh``; a
+    divergent default here is exactly how the FP8 indexer ended up enabled in
+    production while the yaml and the generic script both said otherwise.
+    """
+    assert _default_of(_FLASH_SCRIPT, "USE_V4_FP8_INDEXER") == _default_of(_BASE_SCRIPT, "USE_V4_FP8_INDEXER")
diff --git a/tests/unit_tests/megatron/cco/test_tp_overlap.py b/tests/unit_tests/megatron/cco/test_tp_overlap.py
index b906a34b3..a58af584d 100644
--- a/tests/unit_tests/megatron/cco/test_tp_overlap.py
+++ b/tests/unit_tests/megatron/cco/test_tp_overlap.py
@@ -6,6 +6,7 @@
 import functools
 from contextlib import contextmanager
 
+import pytest
 import torch
 import torch.distributed as dist
 import transformer_engine as te
@@ -22,13 +23,21 @@
 )
 from transformer_engine.pytorch import LayerNormLinear, Linear, fp8_autocast
 
-from primus.backends.transformer_engine import transformer_engine_torch as ptex
-from primus.backends.transformer_engine.pytorch.module.base import (
+# The Primus TE comm-overlap extension imports the ROCm ``hip`` Python module
+# (``primus.backends.transformer_engine.transformer_engine_torch.comm_overlap``).
+# Skip the whole module when ``hip`` isn't importable (non-ROCm or minimally
+# provisioned host) instead of failing at collection time.
+pytest.importorskip("hip", reason="ROCm 'hip' Python module not available")
+
+from primus.backends.transformer_engine import (  # noqa: E402
+    transformer_engine_torch as ptex,
+)
+from primus.backends.transformer_engine.pytorch.module.base import (  # noqa: E402
     get_workspace,
     initialize_ub,
 )
-from primus.core.utils import logger
-from primus.modules.module_utils import set_logging_rank
+from primus.core.utils import logger  # noqa: E402
+from primus.core.utils.module_utils import set_logging_rank  # noqa: E402
 
 
 @contextmanager
diff --git a/tests/unit_tests/megatron/moe/__init__.py b/tests/unit_tests/megatron/moe/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/megatron/moe/test_mega_moe.py b/tests/unit_tests/megatron/moe/test_mega_moe.py
new file mode 100644
index 000000000..a92df8bb4
--- /dev/null
+++ b/tests/unit_tests/megatron/moe/test_mega_moe.py
@@ -0,0 +1,277 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""EP8 parity test: fused ``MegaMoE`` vs Megatron ``MoELayer``.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+from megatron.core import parallel_state
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.moe.moe_layer import MoELayer
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training.global_vars import set_args
+from megatron.training.initialize import _set_random_seed
+from torch.testing._internal.common_distributed import (
+    MultiProcessTestCase,
+    skip_if_lt_x_gpu,
+)
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+)
+
+if not torch.cuda.is_available() or torch.version.hip is None or torch.cuda.get_device_capability() != (9, 5):
+    pytest.skip("MegaMoE only supports gfx950", allow_module_level=True)
+
+from primus.backends.megatron.core.extensions.mega_moe import PrimusTurboMegaMoELayer
+
+
+def _create_args():
+    args = SimpleNamespace()
+    args.sequence_parallel = False
+    args.context_parallel_size = 1
+    args.micro_batch_size = 1
+    args.moe_router_force_load_balancing = False
+    args.moe_use_legacy_grouped_gemm = True
+    args.use_turbo_grouped_gemm = False
+    args.enable_primus_turbo = False
+    args.moe_use_fused_router_with_aux_score = False
+    args.router_logit_softcapping = None
+    return args
+
+
+def _cosine(a: torch.Tensor, b: torch.Tensor) -> float:
+    a = a.detach().reshape(-1).float()
+    b = b.detach().reshape(-1).float()
+    return (a @ b / (a.norm() * b.norm() + 1e-12)).item()
+
+
+def _build_config(ep_size, num_moe_experts, moe_router_topk):
+    """DeepSeek-V3 MoE config: sigmoid group-limited routing + shared expert."""
+    return TransformerConfig(
+        num_layers=1,
+        hidden_size=7168,
+        num_attention_heads=8,
+        tensor_model_parallel_size=1,
+        expert_model_parallel_size=ep_size,
+        pipeline_model_parallel_size=1,
+        num_moe_experts=num_moe_experts,
+        moe_ffn_hidden_size=2048,
+        moe_router_topk=moe_router_topk,
+        # DeepSeek-V3 group-limited routing
+        moe_router_num_groups=8,
+        moe_router_group_topk=4,
+        moe_router_topk_scaling_factor=2.5,
+        moe_router_score_function="sigmoid",
+        moe_router_pre_softmax=False,
+        # noaux_tc bias-based balancing; bias zeroed in the test so routing is
+        # pure sigmoid group-topk shared by both layers' routers.
+        moe_router_enable_expert_bias=True,
+        # aux loss off: keeps every weight's grad a clean fwd/bwd parity check
+        moe_router_load_balancing_type="aux_loss",
+        moe_aux_loss_coeff=0.0,
+        # DeepSeek-V3 shared expert (1 expert, no gate)
+        moe_shared_expert_intermediate_size=2048,
+        moe_shared_expert_gate=False,
+        moe_token_dispatcher_type="alltoall",
+        # fp64 required: fp32 router hits a broken TE general_gemm(workspace=) path
+        moe_router_dtype="fp64",
+        moe_grouped_gemm=False,
+        gated_linear_unit=True,
+        activation_func=torch.nn.functional.silu,  # MegaMoE hardcodes SwiGLU/SiLU
+        add_bias_linear=False,
+        bias_activation_fusion=False,
+        use_cpu_initialization=True,
+        params_dtype=torch.bfloat16,
+        bf16=True,
+    )
+
+
+def _build_layers(config, num_moe_experts, num_layers):
+    """Build ``num_layers`` matched (Megatron MoELayer, MegaMoE) pairs.
+
+    Each layer gets its own seed -> distinct-but-matched weights; ref weights
+    are copied into the mega layer so per-layer outputs are directly comparable.
+    """
+    pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+    moe_layers, mega_layers = [], []
+    for layer_idx in range(num_layers):
+        _set_random_seed(seed_=123 + layer_idx, data_parallel_random_init=False)
+
+        spec = get_gpt_layer_local_spec(num_experts=num_moe_experts, moe_grouped_gemm=False)
+        submodules = spec.submodules.mlp.submodules
+        moe_layer = MoELayer(config, submodules).cuda().to(torch.bfloat16)
+        moe_layer.set_layer_number(layer_idx)
+
+        # mega layer reuses the same router/shared-expert submodules
+        mega_layer = PrimusTurboMegaMoELayer(
+            config, submodules, layer_number=layer_idx, pg_collection=pg_collection
+        ).cuda()
+
+        experts = moe_layer.experts.local_experts
+        mega_w1 = mega_layer.experts.fc1_weight.weight
+        mega_w2 = mega_layer.experts.fc2_weight.weight
+        assert len(experts) == mega_w1.shape[0]
+        with torch.no_grad():
+            # match routing: copy router gate weight into mega's own router
+            mega_layer.router.weight.copy_(moe_layer.router.weight)
+            for r in (moe_layer.router, mega_layer.router):
+                if getattr(r, "bias", None) is not None:
+                    r.bias.zero_()
+                # zero noaux_tc expert bias -> selection == pure sigmoid group-topk
+                if getattr(r, "expert_bias", None) is not None:
+                    r.expert_bias.zero_()
+            # routed experts: grouped fc1/fc2 -> packed w1/w2
+            for i, expert in enumerate(experts):
+                assert expert.linear_fc1.weight.shape == mega_w1[i].shape
+                assert expert.linear_fc2.weight.shape == mega_w2[i].shape
+                mega_w1[i].copy_(expert.linear_fc1.weight.to(torch.bfloat16))
+                mega_w2[i].copy_(expert.linear_fc2.weight.to(torch.bfloat16))
+            # shared expert: same SharedExpertMLP module, copy fc1/fc2
+            if mega_layer.shared_experts is not None:
+                se_ref, se_mega = moe_layer.shared_experts, mega_layer.shared_experts
+                se_mega.linear_fc1.weight.copy_(se_ref.linear_fc1.weight)
+                se_mega.linear_fc2.weight.copy_(se_ref.linear_fc2.weight)
+        moe_layers.append(moe_layer)
+        mega_layers.append(mega_layer)
+    return moe_layers, mega_layers
+
+
+@instantiate_parametrized_tests
+class TestMegaMoEAccuracy(MultiProcessTestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self._spawn_processes()
+
+    def tearDown(self) -> None:
+        super().tearDown()
+
+    @property
+    def world_size(self) -> int:
+        return torch.cuda.device_count()
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cuda", self.rank)
+
+    def _init_process(self):
+        torch.cuda.set_device(self.device)
+        store = dist.FileStore(self.file_name, self.world_size)
+        dist.init_process_group(backend="nccl", world_size=self.world_size, rank=self.rank, store=store)
+        set_args(_create_args())
+        # standalone harness has no Primus logger; attach a plain one so any
+        # log_rank_0 call in the code path doesn't crash
+        import logging
+
+        from primus.core.utils import logger as _primus_logger
+
+        if _primus_logger._logger is None:
+            _primus_logger._logger = logging.getLogger("mega_moe_test")
+
+    def _setup_ep(self):
+        parallel_state.destroy_model_parallel()
+        parallel_state.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            expert_model_parallel_size=self.world_size,
+        )
+
+    def _rank0(self, msg):
+        if self.rank == 0:
+            print(msg, flush=True)
+
+    @skip_if_lt_x_gpu(8)
+    @parametrize("moe_router_topk", [8])
+    @parametrize("num_ga", [4])
+    def test_forward_backward(self, moe_router_topk, num_ga):
+        # DeepSeek-V3 EP8 shapes: 256 experts, top-8, hidden 7168.
+        # Single MoE layer, num_ga gradient-accumulation microbatches: forward +
+        # backward num_ga times without zeroing grads (mirrors real GA training).
+        # One layer -> no cross-layer backward accumulation, so grads stay clean.
+        num_moe_experts, hidden_size = 256, 7168
+        self._init_process()
+        self._setup_ep()
+        try:
+            config = _build_config(self.world_size, num_moe_experts, moe_router_topk)
+            moe_layers, mega_layers = _build_layers(config, num_moe_experts, 1)
+            moe_layer, mega_layer = moe_layers[0], mega_layers[0]
+            moe_layer.train(True)
+            mega_layer.train(True)
+            experts = moe_layer.experts.local_experts
+            se = moe_layer.shared_experts
+            seq, batch = 8192, 1
+
+            fwd_cos, dx_cos = [], []
+            # GA loop: accumulate grads over num_ga microbatches, compare each fwd/dx
+            for step in range(num_ga):
+                torch.manual_seed(1000 + self.rank + step * 97)
+                x = torch.randn((seq, batch, hidden_size), dtype=torch.bfloat16, device=self.device)
+                g = torch.randn((seq, batch, hidden_size), dtype=torch.bfloat16, device=self.device)
+
+                x_ref = x.clone().requires_grad_(True)
+                ref, _ = moe_layer(x_ref)
+                ref.backward(g)
+                x_meg = x.clone().requires_grad_(True)
+                out, _ = mega_layer(x_meg)
+                out.backward(g)
+
+                fwd_cos.append(_cosine(out, ref))
+                dx_cos.append(_cosine(x_meg.grad, x_ref.grad))
+                self._rank0(
+                    f"[ga {step}] fwd cos={fwd_cos[-1]:.6f} dx cos={dx_cos[-1]:.6f} "
+                    f"max_abs_diff={(out.float()-ref.float()).abs().max().item():.3e}"
+                )
+
+            GRAD_FLOOR = 0.95
+            failures = []
+
+            def check_accuracy(tag, cos):
+                if cos <= GRAD_FLOOR:
+                    failures.append(f"{tag} cosine {cos:.6f} < {GRAD_FLOOR}")
+
+            # forward + dx parity: every microbatch
+            for step in range(num_ga):
+                check_accuracy(f"ga{step} fwd", fwd_cos[step])
+                check_accuracy(f"ga{step} dx", dx_cos[step])
+
+            # accumulated per-weight grad parity: every trainable weight
+            dgate = _cosine(mega_layer.router.weight.grad, moe_layer.router.weight.grad)
+            check_accuracy("router", dgate)
+            dsw1 = _cosine(mega_layer.shared_experts.linear_fc1.weight.grad, se.linear_fc1.weight.grad)
+            dsw2 = _cosine(mega_layer.shared_experts.linear_fc2.weight.grad, se.linear_fc2.weight.grad)
+            check_accuracy("shared dW1", dsw1)
+            check_accuracy("shared dW2", dsw2)
+            # routed experts: full w1/w2 grad tensors (norm-weighted over all local
+            # experts); per-expert min is thin-token bf16 noise, informational only.
+            ref_dw1 = torch.stack([e.linear_fc1.weight.grad for e in experts])
+            ref_dw2 = torch.stack([e.linear_fc2.weight.grad for e in experts])
+            mega_w1_grad = mega_layer.experts.fc1_weight.weight.grad
+            mega_w2_grad = mega_layer.experts.fc2_weight.weight.grad
+            dW1 = _cosine(mega_w1_grad, ref_dw1)
+            dW2 = _cosine(mega_w2_grad, ref_dw2)
+            check_accuracy("experts dW1", dW1)
+            check_accuracy("experts dW2", dW2)
+            dw1_min = min(_cosine(mega_w1_grad[i], e.linear_fc1.weight.grad) for i, e in enumerate(experts))
+            dw2_min = min(_cosine(mega_w2_grad[i], e.linear_fc2.weight.grad) for i, e in enumerate(experts))
+            self._rank0(
+                f"[grad accum over {num_ga} ga] dgate={dgate:.6f} dSharedW1={dsw1:.6f} "
+                f"dSharedW2={dsw2:.6f} experts dW1={dW1:.6f} dW2={dW2:.6f} "
+                f"(per-expert min {dw1_min:.4f}/{dw2_min:.4f})"
+            )
+            self.assertEqual(failures, [], f"parity below {GRAD_FLOOR}: {failures}")
+        finally:
+            parallel_state.destroy_model_parallel()
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/__init__.py b/tests/unit_tests/megatron/transformer/deepseek_v4/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py b/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py
new file mode 100644
index 000000000..2b14d4f28
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py
@@ -0,0 +1,80 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Per-test cleanup hooks for plan-4 V4 attention tests.
+
+The plan-4 release-tier shape gates (``pytest.mark.slow``, see G28 in
+``deepseek-v4/develop/plan-4/02-phase-details.md``) allocate large
+eager-reference tensors at production V4 dimensions
+(``head_dim=512``, MQA / MHA, ``K_topk`` up to 1024). Specifically the
+CSA path's ``torch.einsum("bhsd,bhskd->bhsk", q, gathered.unsqueeze(1)
+.expand(...))`` materialises a ``[B, H, Sq, K, D]`` intermediate that
+is up to ~64 GiB at fp32 V4-Flash release / ~128 GiB at fp32 V4-Pro
+release. PyTorch's caching allocator otherwise holds those tensors
+across parametrised tests in the same process, so the second / third
+release-tier test in a pytest session ends up OOM-ing well before
+exhausting the MI355 287 GiB HBM budget.
+
+This conftest installs a function-scoped autouse fixture that empties
+the CUDA / HIP cache after every test to keep the allocator pressure
+test-local. The cost (~milliseconds) is negligible for the fast tier
+and an outright requirement for the release tier.
+
+It also gates the entire directory on MI355X hardware: the V4 kernels
+and their numerics are validated on MI355X (gfx950), so on any other
+accelerator (or CPU-only host) every test in this directory is skipped.
+Set ``PRIMUS_V4_UT_ALLOW_NON_MI355X=1`` to bypass the gate.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+
+import pytest
+
+try:
+    import torch
+except ImportError:  # pragma: no cover — pytest collects on torchless envs
+    torch = None  # type: ignore
+
+
+@functools.lru_cache(maxsize=1)
+def _is_mi355x() -> bool:
+    """Return True iff the local accelerator reports as an AMD MI355X.
+
+    Matches the device marketing name (e.g. ``AMD Instinct MI355X``);
+    returns False on CPU-only / non-ROCm hosts.
+    """
+    if torch is None or not torch.cuda.is_available():
+        return False
+    try:
+        name = torch.cuda.get_device_name(0) or ""
+    except Exception:  # pragma: no cover — defensive
+        return False
+    return "mi355x" in name.lower()
+
+
+@pytest.fixture(autouse=True)
+def _require_mi355x():
+    """Skip DeepSeek-V4 unit tests unless the GPU is an MI355X.
+
+    The V4 kernels / numerics are validated against MI355X; running them on
+    other architectures is neither supported nor meaningful. Set
+    ``PRIMUS_V4_UT_ALLOW_NON_MI355X=1`` to bypass (e.g. local debugging).
+    """
+    if os.environ.get("PRIMUS_V4_UT_ALLOW_NON_MI355X", "0") != "0":
+        return
+    if not _is_mi355x():
+        pytest.skip("DeepSeek-V4 unit tests require an MI355X GPU")
+
+
+@pytest.fixture(autouse=True)
+def _release_gpu_cache_per_test():
+    """Free PyTorch's CUDA / HIP cache after each V4 attention test."""
+    yield
+    if torch is not None and torch.cuda.is_available():
+        torch.cuda.empty_cache()
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py
new file mode 100644
index 000000000..2c991ef1a
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py
@@ -0,0 +1,163 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for V4 pre-multiplication clamped SwiGLU (G3, plan-2 §04).
+
+These tests pin the math against the HF reference at
+``DeepSeek-V4-Flash/inference/model.py:Expert.forward`` and assert the
+parameter layout (``w1`` / ``w2`` / ``w3``) so the released checkpoint
+loads through the V4 state-dict adapter without remapping.
+
+Tolerance: 1e-6 absolute, fp32, on randomized inputs (G3 gate).
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.clamped_swiglu import (
+    ClampedSwiGLUMLP,
+    clamped_swiglu_pre_mul,
+    clamped_swiglu_pre_mul_fused,
+)
+
+
+def _hf_reference_clamped_swiglu(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    alpha: float,
+) -> torch.Tensor:
+    """Inline HF reference for ``Expert.forward``'s activation.
+
+    Mirrors the HF code exactly, including the float() cast and the
+    pre-multiplication one-sided / two-sided clamp.
+    """
+    gate = gate.float()
+    up = up.float()
+    if alpha > 0.0:
+        up = torch.clamp(up, min=-alpha, max=alpha)
+        gate = torch.clamp(gate, max=alpha)
+    return F.silu(gate) * up
+
+
+@pytest.mark.parametrize("alpha", [7.0, 3.5, 1.0, 0.5])
+def test_clamped_swiglu_pre_mul_matches_hf_reference(alpha: float) -> None:
+    """G3: split-input form matches HF reference within 1e-6 fp32."""
+    torch.manual_seed(1234)
+    gate = torch.randn(4, 17, 32, dtype=torch.float32) * 5.0
+    up = torch.randn(4, 17, 32, dtype=torch.float32) * 5.0
+
+    out = clamped_swiglu_pre_mul(gate, up, alpha=alpha)
+    ref = _hf_reference_clamped_swiglu(gate, up, alpha=alpha)
+
+    assert out.shape == gate.shape
+    assert torch.isfinite(out).all()
+    max_abs = (out - ref).abs().max().item()
+    assert max_abs <= 1.0e-6, f"max-abs error vs HF reference = {max_abs}"
+
+
+def test_clamped_swiglu_pre_mul_alpha_zero_disables_clamp() -> None:
+    """``alpha == 0`` falls back to vanilla SwiGLU (no clamp)."""
+    torch.manual_seed(7)
+    gate = torch.randn(2, 16, dtype=torch.float32) * 100.0
+    up = torch.randn(2, 16, dtype=torch.float32) * 100.0
+    out = clamped_swiglu_pre_mul(gate, up, alpha=0.0)
+    expected = F.silu(gate) * up
+    assert (out - expected).abs().max().item() <= 1.0e-6
+
+
+def test_clamped_swiglu_pre_mul_fused_matches_split_form() -> None:
+    """The fused [gate | up] form matches the split-input form."""
+    torch.manual_seed(99)
+    I = 24
+    gate = torch.randn(3, 11, I, dtype=torch.float32) * 4.0
+    up = torch.randn(3, 11, I, dtype=torch.float32) * 4.0
+    fused = torch.cat([gate, up], dim=-1)
+
+    out_split = clamped_swiglu_pre_mul(gate, up, alpha=7.0)
+    out_fused = clamped_swiglu_pre_mul_fused(fused, alpha=7.0)
+    assert (out_split - out_fused).abs().max().item() <= 1.0e-6
+
+
+def test_clamped_swiglu_pre_mul_one_sided_gate_clamp() -> None:
+    """Verify the gate clamp is one-sided (max=alpha) and up is two-sided.
+
+    Sets gate values both well below -alpha and well above +alpha; the
+    output activation should pass the negative-gate values through SiLU
+    (since gate clamp only bounds the top side) and bound the positive
+    gate side at alpha.
+    """
+    alpha = 7.0
+    gate_lo = torch.tensor([[-100.0]], dtype=torch.float32)
+    gate_hi = torch.tensor([[+100.0]], dtype=torch.float32)
+    up = torch.tensor([[1.0]], dtype=torch.float32)
+
+    out_lo = clamped_swiglu_pre_mul(gate_lo, up, alpha=alpha)
+    out_hi = clamped_swiglu_pre_mul(gate_hi, up, alpha=alpha)
+
+    # Gate negative side is NOT clamped; SiLU(-100) ~ -100 * sigmoid(-100) ~ 0.
+    # We just assert the output is finite and small (close to 0).
+    assert torch.isfinite(out_lo).item()
+    assert out_lo.abs().item() < 1.0e-6
+    # Gate positive side IS clamped at alpha; SiLU(alpha) is the upper bound.
+    expected_hi = F.silu(torch.tensor(alpha)) * up
+    assert (out_hi - expected_hi).abs().max().item() <= 1.0e-6
+
+
+def test_clamped_swiglu_mlp_state_dict_uses_w1_w2_w3_layout() -> None:
+    """Released-checkpoint compatibility: w1 / w2 / w3 keys exist.
+
+    The HF released ``Expert`` checkpoint uses ``w1.weight`` /
+    ``w2.weight`` / ``w3.weight`` (no ``gate_up.weight``). This test
+    fails if a future refactor breaks that promise.
+    """
+    mlp = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0, bias=False)
+    keys = set(mlp.state_dict().keys())
+    assert "w1.weight" in keys
+    assert "w2.weight" in keys
+    assert "w3.weight" in keys
+    # No fused projection key may leak into the state-dict.
+    assert "gate_up.weight" not in keys
+
+
+def test_clamped_swiglu_mlp_fused_gate_up_matches_split_path() -> None:
+    """The fused-forward variant produces identical outputs to the eager path."""
+    torch.manual_seed(42)
+    hidden = torch.randn(3, 7, 8, dtype=torch.float32)
+
+    mlp_split = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0)
+    mlp_fused = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0, fused_gate_up=True)
+    # Share weights so we are only testing the forward layout, not init.
+    mlp_fused.load_state_dict(mlp_split.state_dict())
+
+    out_split = mlp_split(hidden)
+    out_fused = mlp_fused(hidden)
+    assert (out_split - out_fused).abs().max().item() <= 1.0e-6
+
+
+def test_clamped_swiglu_mlp_forward_matches_hf_expert_forward() -> None:
+    """End-to-end: ``ClampedSwiGLUMLP.forward`` matches HF ``Expert.forward``."""
+    torch.manual_seed(0)
+    H, I = 12, 28
+    alpha = 7.0
+    mlp = ClampedSwiGLUMLP(hidden_size=H, intermediate_size=I, alpha=alpha)
+
+    x = torch.randn(2, 5, H, dtype=torch.float32)
+
+    # Inline HF reference using the same w1/w2/w3 weights.
+    w1 = mlp.w1.weight.detach()
+    w3 = mlp.w3.weight.detach()
+    w2 = mlp.w2.weight.detach()
+    gate = F.linear(x, w1)
+    up = F.linear(x, w3)
+    h = _hf_reference_clamped_swiglu(gate, up, alpha=alpha)
+    ref = F.linear(h, w2)
+
+    out = mlp(x)
+    assert (out - ref).abs().max().item() <= 1.0e-6
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py
new file mode 100644
index 000000000..a6f5e9e70
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py
@@ -0,0 +1,1176 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the plan-2 V4-faithful :class:`DeepseekV4Attention`.
+
+These tests run on CPU (fp32) and verify the dense (``compress_ratio == 0``)
+attention math against an *inline* reference implementation that mirrors
+the released ``DeepSeek-V4-Flash/inference/model.py`` semantics:
+
+* Single-latent KV: ``K = V = wkv(hidden)`` broadcast to all query heads.
+* Per-head ``q_rms``: parameter-less RMS on ``head_dim`` after ``wq_b``.
+* Learnable per-head ``attn_sink``: an extra "virtual key" with zero value
+  joined into the softmax (then dropped before the value-weighted sum).
+* Grouped low-rank O: einsum-based ``wo_a`` per group + ``wo_b``.
+* Partial **interleaved** RoPE on the last ``qk_pos_emb_head_dim`` channels.
+
+The reference uses the *same* interleaved RoPE convention as Primus's
+:mod:`dual_rope` (per the techblog correction over the original HF PR
+which used rotate-half), so the two implementations should agree to
+machine precision when given identical weights and positions.
+
+Plan-2 P17 will replace this inline reference with the actual HF
+``DeepseekV4Attention.forward`` from the released checkpoint.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+# Importing the module under test pulls Megatron / TE through the
+# Primus extensions; skip the whole module when those are unavailable
+# (e.g. a CPU-only smoke environment).
+mla_module = pytest.importorskip(
+    "megatron.core.transformer.multi_latent_attention",
+    reason="Megatron MLA not importable in this environment",
+)
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: E402
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import DualRoPE  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Inline reference V4 attention forward (single-latent KV, attn_sink, grouped O)
+# ---------------------------------------------------------------------------
+
+
+def _rms_norm_per_head(x: torch.Tensor, eps: float) -> torch.Tensor:
+    """Parameter-less per-head RMS (matches the released checkpoint)."""
+    in_dtype = x.dtype
+    x32 = x.float()
+    rms = torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps)
+    return (x32 * rms).to(in_dtype)
+
+
+def _inverse_rope_reference(
+    out: torch.Tensor,
+    *,
+    position_ids: torch.Tensor,
+    rope: DualRoPE,
+    rotary_dim: int,
+    compress_ratio: int = 0,
+) -> torch.Tensor:
+    """Reference de-rotation of the attention output (RoPE at ``-t``).
+
+    Mirrors ``apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True)`` from
+    the released ``inference/model.py``: the conjugate rotation, expressed
+    here by negating the sine.
+    """
+    from primus.backends.megatron.core.transformer.dual_rope import (
+        apply_interleaved_partial_rope,
+    )
+
+    cache = rope.get_rope(compress_ratio=compress_ratio)
+    cos, sin = cache(position_ids)
+    cos = cos[..., : rotary_dim // 2]
+    sin = sin[..., : rotary_dim // 2]
+    lead = out.shape[:-2]
+    if tuple(cos.shape[:-1]) != tuple(lead):
+        cos = cos.expand(*lead, -1)
+        sin = sin.expand(*lead, -1)
+    return apply_interleaved_partial_rope(out, cos, -sin, rotary_dim=rotary_dim)
+
+
+def _reference_v4_attention_forward(
+    *,
+    hidden: torch.Tensor,  # [B, S, D]
+    position_ids: torch.Tensor,  # [B, S] or [S]
+    rope: DualRoPE,
+    wq_a: torch.Tensor,  # [q_lora_rank, D]
+    wq_b: torch.Tensor,  # [n_heads * head_dim, q_lora_rank]
+    q_norm_w: torch.Tensor,  # [q_lora_rank]
+    wkv: torch.Tensor,  # [head_dim, D]
+    kv_norm_w: torch.Tensor,  # [head_dim]
+    wo_a: torch.Tensor,  # [o_groups * o_lora_rank, n_heads * head_dim / o_groups]
+    wo_b: torch.Tensor,  # [D, o_groups * o_lora_rank]
+    attn_sink,  # Optional[torch.Tensor] of shape [n_heads]; None disables sink
+    n_heads: int,
+    head_dim: int,
+    rotary_dim: int,
+    o_groups: int,
+    o_lora_rank: int,
+    norm_eps: float,
+) -> torch.Tensor:
+    """Reference V4 attention forward (CPU, fp32-ish, no parallel linear)."""
+    B, S, D = hidden.shape
+
+    # Q branch.
+    q_compressed = hidden @ wq_a.t()  # [B, S, q_lora_rank]
+    # RMSNorm on q_lora_rank with learnable gamma (== q_layernorm).
+    q32 = q_compressed.float()
+    q_rms = torch.rsqrt(q32.square().mean(dim=-1, keepdim=True) + norm_eps)
+    q_compressed = (q32 * q_rms).to(q_compressed.dtype) * q_norm_w
+    q = q_compressed @ wq_b.t()  # [B, S, n_heads * head_dim]
+    q = q.view(B, S, n_heads, head_dim)
+    q = _rms_norm_per_head(q, norm_eps)
+
+    # KV branch (single-latent).
+    kv = hidden @ wkv.t()  # [B, S, head_dim]
+    kv32 = kv.float()
+    kv_rms = torch.rsqrt(kv32.square().mean(dim=-1, keepdim=True) + norm_eps)
+    kv = (kv32 * kv_rms).to(kv.dtype) * kv_norm_w
+    kv = kv.view(B, S, 1, head_dim)
+
+    # Partial RoPE (interleaved) on Q and K. K = kv (rope-applied).
+    q = rope.apply_rope(q, position_ids=position_ids, compress_ratio=0)
+    kv = rope.apply_rope(kv, position_ids=position_ids, compress_ratio=0)
+    k = kv  # [B, S, 1, head_dim]
+    v = kv  # K = V (single latent)
+
+    # Broadcast K / V across heads.
+    k = k.expand(B, S, n_heads, head_dim)
+    v = v.expand(B, S, n_heads, head_dim)
+
+    # Causal mask (no SWA in the test).
+    q_idx = torch.arange(S, device=hidden.device).unsqueeze(1)
+    k_idx = torch.arange(S, device=hidden.device).unsqueeze(0)
+    causal = torch.where(q_idx >= k_idx, 0.0, float("-inf")).to(hidden.dtype)
+
+    # Move heads dim before sequence.
+    q_bh = q.transpose(1, 2)  # [B, H, S, head_dim]
+    k_bh = k.transpose(1, 2)
+    v_bh = v.transpose(1, 2)
+
+    scale = 1.0 / math.sqrt(head_dim)
+    logits = torch.matmul(q_bh.float(), k_bh.float().transpose(-2, -1)) * scale
+    logits = logits + causal
+
+    if attn_sink is None:
+        # Plain causal softmax (no virtual key column).
+        logits = logits - logits.amax(dim=-1, keepdim=True).detach()
+        probs = logits.softmax(dim=-1).to(v_bh.dtype)
+    else:
+        # attn_sink: append a virtual key column with zero value, drop after softmax.
+        sink_col = attn_sink.float().view(1, n_heads, 1, 1).expand(B, n_heads, S, 1)
+        logits_aug = torch.cat([logits, sink_col], dim=-1)
+        logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach()
+        probs = logits_aug.softmax(dim=-1)
+        probs = probs[..., :-1].to(v_bh.dtype)
+    out = torch.matmul(probs, v_bh.float()).to(hidden.dtype)
+    out = out.transpose(1, 2).contiguous()  # [B, S, H, head_dim]
+
+    # Inverse RoPE (position = -t) on the output tail: K == V, so the output
+    # carries absolute positions until it is de-rotated.
+    out = _inverse_rope_reference(out, position_ids=position_ids, rope=rope, rotary_dim=rotary_dim)
+
+    # Grouped low-rank O.
+    out_g = out.reshape(B, S, o_groups, (n_heads * head_dim) // o_groups)
+    wo_a_w = wo_a.view(o_groups, o_lora_rank, (n_heads * head_dim) // o_groups)
+    o = torch.einsum("bsgd,grd->bsgr", out_g, wo_a_w)
+    o = o.flatten(2)
+    return o @ wo_b.t()
+
+
+# ---------------------------------------------------------------------------
+# Test fixtures
+# ---------------------------------------------------------------------------
+
+
+_TEST_DTYPE = torch.float32
+
+
+@pytest.fixture(autouse=True)
+def _v4_attention_on_cuda():
+    """DeepSeek-V4 attention (all backends, incl. eager) runs on GPU tensors.
+
+    These tests build CPU-free tensors by defaulting to the CUDA/HIP device;
+    skip the module on a CPU-only host.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("DeepSeek-V4 attention requires a CUDA/HIP device")
+    torch.set_default_device("cuda")
+    try:
+        yield
+    finally:
+        torch.set_default_device("cpu")
+
+
+def _make_v4_config(
+    *,
+    hidden_size: int,
+    num_heads: int,
+    head_dim: int,
+    rotary_dim: int,
+    q_lora_rank: int,
+    o_groups: int,
+    o_lora_rank: int,
+    attn_sink: bool,
+    norm_eps: float = 1e-6,
+) -> DeepSeekV4TransformerConfig:
+    """Minimal V4 config for CPU unit tests.
+
+    Relies on dataclass defaults wherever possible and only sets the
+    fields the V4 attention actually reads.
+    """
+    return DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=hidden_size,
+        num_attention_heads=num_heads,
+        num_query_groups=1,
+        kv_channels=head_dim,
+        qk_pos_emb_head_dim=rotary_dim,
+        qk_head_dim=head_dim - rotary_dim,
+        v_head_dim=head_dim,
+        kv_lora_rank=head_dim,  # unused — KV branch is overridden by V4
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+        # V4 extras
+        q_lora_rank=q_lora_rank,
+        o_groups=o_groups,
+        o_lora_rank=o_lora_rank,
+        attn_sliding_window=0,
+        attn_sink=attn_sink,
+        compress_ratios=None,
+        compress_rope_theta=160000.0,
+        # These tests validate the V4 attention math against an inline eager
+        # reference at tiny (head_dim=16) shapes; the Triton/gluon/turbo kernels
+        # are specialized for head_dim=512, so pin the eager backend here.
+        use_v4_attention_backend="eager",
+        use_v4_csa_attention_backend="eager",
+        # Misc
+        layernorm_epsilon=norm_eps,
+        norm_epsilon=norm_eps,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+    )
+
+
+def _make_attention(config: DeepSeekV4TransformerConfig) -> DeepseekV4Attention:
+    """Construct V4 attention with default (no-spec) submodules, on CPU.
+
+    With ``submodules=None`` every projection falls back to ``nn.Linear``
+    (replicated, no TP), which is exactly what we want for a CPU smoke
+    test against the inline reference.
+    """
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    return DeepseekV4Attention(
+        config,
+        rope=rope,
+        compress_ratio=0,
+        submodules=None,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_state_dict_keys_match_v4_canonical_layout():
+    """The new attention exposes V4-canonical state-dict keys.
+
+    These are exactly the keys the P17 state-dict adapter will map from
+    the released ``layers.{i}.attn.{wq_a,wq_b,wkv,q_norm,kv_norm,wo_a,wo_b,attn_sink}``
+    safetensors layout.
+    """
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    attn = _make_attention(config)
+    keys = set(attn.state_dict().keys())
+
+    expected = {
+        "linear_q_down_proj.weight",
+        "linear_q_up_proj.weight",
+        "linear_kv.weight",
+        "q_layernorm.weight",
+        "kv_layernorm.weight",
+        "linear_o_a.weight",
+        "linear_o_b.weight",
+        "attn_sink",  # direct nn.Parameter, matches released checkpoint key
+    }
+    missing = expected - keys
+    assert not missing, f"missing V4-canonical keys: {missing}"
+
+    legacy = {
+        "q_a.weight",
+        "q_b.weight",
+        "k_proj.weight",
+        "v_proj.weight",
+        "o_proj.weight",
+    }
+    bleed = legacy & keys
+    assert not bleed, f"legacy plan-1 keys leaked into V4-faithful attention: {bleed}"
+
+
+def test_forward_shape_and_finite():
+    """Forward pass returns ``[B, S, hidden_size]`` of finite values."""
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    attn = _make_attention(config).to(_TEST_DTYPE)
+
+    B, S, D = 2, 8, 64
+    hidden = torch.randn(B, S, D, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+    out = attn(hidden, position_ids)
+    assert out.shape == (B, S, D)
+    assert torch.isfinite(out).all()
+
+
+def _copy_weights_to_reference(attn: DeepseekV4Attention) -> dict:
+    """Extract weights from the new attention into the inline reference's
+    parameter dict."""
+    return {
+        "wq_a": attn.linear_q_down_proj.weight.detach().clone(),
+        "wq_b": attn.linear_q_up_proj.weight.detach().clone(),
+        "q_norm_w": attn.q_layernorm.weight.detach().clone(),
+        "wkv": attn.linear_kv.weight.detach().clone(),
+        "kv_norm_w": attn.kv_layernorm.weight.detach().clone(),
+        "wo_a": attn.linear_o_a.weight.detach().clone(),
+        "wo_b": attn.linear_o_b.weight.detach().clone(),
+    }
+
+
+@pytest.mark.parametrize("attn_sink_enabled", [False, True])
+def test_forward_matches_inline_reference(attn_sink_enabled: bool):
+    """1-layer V4 attention forward agrees with the inline reference.
+
+    Both implementations apply the same math (single-latent KV, partial
+    interleaved RoPE, attn_sink as virtual key column, grouped low-rank
+    O), so the agreement is essentially numerical noise (≤1e-5 in fp32).
+    """
+    torch.manual_seed(0)
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=attn_sink_enabled,
+    )
+    attn = _make_attention(config).to(_TEST_DTYPE)
+    sink_tensor = None
+    if attn_sink_enabled:
+        # Pull in some non-trivial sink scalars so the test exercises
+        # the virtual-key-column path.
+        with torch.no_grad():
+            attn.attn_sink.copy_(torch.linspace(-0.5, 0.5, attn.num_heads))
+        sink_tensor = attn.attn_sink.detach().clone()
+
+    weights = _copy_weights_to_reference(attn)
+
+    B, S = 2, 8
+    hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    with torch.no_grad():
+        ours = attn(hidden, position_ids)
+        ref = _reference_v4_attention_forward(
+            hidden=hidden,
+            position_ids=position_ids,
+            rope=attn.rope,
+            attn_sink=sink_tensor,
+            n_heads=attn.num_heads,
+            head_dim=attn.head_dim,
+            rotary_dim=attn.rotary_dim,
+            o_groups=attn.o_groups,
+            o_lora_rank=attn.o_lora_rank,
+            norm_eps=attn.norm_eps,
+            **weights,
+        )
+
+    diff = (ours - ref).abs().max().item()
+    assert diff < 1e-3, f"forward mismatch: max abs diff = {diff:.3e}"
+
+
+def test_per_head_q_rms_is_parameterless():
+    """The released checkpoint stores no separate ``q_rms`` parameter.
+
+    Plan-2 P13 lands per-head q_rms as inline math (parameter-less RMS on
+    ``head_dim``), not as an additional ``nn.Module`` with a learnable
+    gamma. This regression test guards that contract.
+    """
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    attn = _make_attention(config)
+    keys = set(attn.state_dict().keys())
+    forbidden = {"q_rms.weight", "q_rms_norm.weight", "q_per_head_rms.weight"}
+    leaked = forbidden & keys
+    assert not leaked, (
+        f"Per-head q_rms must be parameter-less (released checkpoint has no "
+        f"such key); leaked params: {leaked}"
+    )
+
+
+def test_o_lora_rank_zero_falls_back_to_flat_proj():
+    """Setting ``o_lora_rank == 0`` skips the grouped O path."""
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=0,
+        attn_sink=False,
+    )
+    attn = _make_attention(config)
+    assert attn.linear_o_a is None
+    assert attn.linear_o_b is None
+    assert attn.linear_proj is not None
+    keys = set(attn.state_dict().keys())
+    assert "linear_proj.weight" in keys
+    assert "linear_o_a.weight" not in keys
+    assert "linear_o_b.weight" not in keys
+
+
+def test_unsupported_compress_ratio_rejected():
+    """V4 attention only supports ``compress_ratio in {0, 4, 128}``.
+
+    Plan-2 P13 follow-up landed CSA (4) and HCA (128) inside the new
+    class; anything else (e.g. 64, 256) is a config error.
+    """
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+    )
+    with pytest.raises(ValueError, match="compress_ratio in"):
+        DeepseekV4Attention(config, rope=rope, compress_ratio=3, submodules=None)
+
+
+def test_q_lora_rank_zero_rejected():
+    """V4 always uses Q LoRA (wq_a + wq_b); reject q_lora_rank == 0."""
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=0,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+    )
+    with pytest.raises(ValueError, match="q_lora_rank > 0"):
+        DeepseekV4Attention(config, rope=rope, compress_ratio=0, submodules=None)
+
+
+# ---------------------------------------------------------------------------
+# Compressed-branch tests (plan-2 P13 follow-up: HCA / CSA folded into the
+# single :class:`DeepseekV4Attention` class).
+# ---------------------------------------------------------------------------
+
+
+def _make_compressed_attention(
+    *,
+    config: DeepSeekV4TransformerConfig,
+    compress_ratio: int,
+) -> DeepseekV4Attention:
+    """Construct V4 attention with a non-zero ``compress_ratio``.
+
+    With ``submodules=None`` the compressor (and indexer for CSA) fall
+    back to local :class:`Compressor` / :class:`Indexer` instances and
+    every projection falls back to ``nn.Linear`` — ideal for CPU smoke
+    tests without a TP group.
+    """
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,  # disable YaRN so CPU references stay simple
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    return DeepseekV4Attention(
+        config,
+        rope=rope,
+        compress_ratio=compress_ratio,
+        submodules=None,
+    )
+
+
+def test_hca_forward_shape_and_finite():
+    """HCA (compress_ratio=128) forward pass produces ``[B, S, D]`` finite."""
+    compress_ratio = 128
+    B, S = 2, compress_ratio  # exactly one compressed-pool slot
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio)
+    attn = attn.to(_TEST_DTYPE)
+    assert attn.compressor is not None
+    assert attn.indexer is None  # HCA does not use Indexer
+
+    hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+    out = attn(hidden, position_ids)
+    assert out.shape == (B, S, config.hidden_size)
+    assert torch.isfinite(out).all()
+
+
+def test_csa_forward_shape_and_finite():
+    """CSA (compress_ratio=4) forward pass produces ``[B, S, D]`` finite.
+
+    Builds an Indexer + overlap-mode Compressor through the spec-less
+    fallback path. The key contract: with valid ``index_topk`` ≤ ``P``
+    selections the joint softmax-with-sink path runs end-to-end.
+    """
+    compress_ratio = 4
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    # Override Indexer config knobs (the test config dataclass exposes them
+    # as attributes; these are read by ``_build_indexer``).
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    B, S = 2, 8  # P = S // ratio = 2 → top-K = 2 always covers the pool
+    attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio)
+    attn = attn.to(_TEST_DTYPE)
+    assert attn.compressor is not None
+    assert attn.indexer is not None
+
+    hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+    out = attn(hidden, position_ids)
+    assert out.shape == (B, S, config.hidden_size)
+    assert torch.isfinite(out).all()
+
+
+def _reference_hca_forward(
+    *,
+    attn: DeepseekV4Attention,
+    hidden: torch.Tensor,
+    position_ids: torch.Tensor,
+) -> torch.Tensor:
+    """Inline HCA reference forward, matching the plan-2 fold-in.
+
+    Reproduces the new ``DeepseekV4Attention.forward`` for compress_ratio
+    == 128 step-for-step but written as plain matmuls / einsums so the
+    test is independent of internal helpers.
+    """
+    from primus.backends.megatron.core.transformer.dual_rope import (
+        apply_interleaved_partial_rope,
+    )
+
+    B, S, _ = hidden.shape
+    H = attn.num_heads
+    Dh = attn.head_dim
+    rotary_dim = attn.rotary_dim
+    eps = attn.norm_eps
+    ratio = attn.compress_ratio
+
+    # Q branch (single-latent KV; same as dense reference).
+    wq_a = attn.linear_q_down_proj.weight
+    wq_b = attn.linear_q_up_proj.weight
+    q_n = attn.q_layernorm.weight
+    q_compressed = hidden @ wq_a.t()
+    q32 = q_compressed.float()
+    q_rms = torch.rsqrt(q32.square().mean(dim=-1, keepdim=True) + eps)
+    q_compressed = (q32 * q_rms).to(q_compressed.dtype) * q_n
+    q = q_compressed @ wq_b.t()
+    q = q.view(B, S, H, Dh)
+    q = _rms_norm_per_head(q, eps)
+
+    wkv = attn.linear_kv.weight
+    kv_n = attn.kv_layernorm.weight
+    kv = hidden @ wkv.t()
+    kv32 = kv.float()
+    kv_rms = torch.rsqrt(kv32.square().mean(dim=-1, keepdim=True) + eps)
+    kv = (kv32 * kv_rms).to(kv.dtype) * kv_n
+    kv = kv.view(B, S, 1, Dh)
+
+    # Q / K rope using the LAYER's compress_ratio (compress base for HCA).
+    q = attn.rope.apply_rope(q, position_ids=position_ids, compress_ratio=ratio)
+    kv = attn.rope.apply_rope(kv, position_ids=position_ids, compress_ratio=ratio)
+    k_local = kv.expand(B, S, H, Dh)
+    v_local = kv.expand(B, S, H, Dh)
+
+    # Compressed pool from the attention's own Compressor + compress-base RoPE.
+    pool = attn.compressor(hidden)  # [B, P, Dh]
+    P = pool.shape[1]
+    # Compressed entry s stands for the window starting at original token
+    # s * ratio, so it is rotated at that position (not at the block index).
+    comp_pos = torch.arange(P, device=hidden.device) * ratio
+    cos, sin = attn.rope.compress_rope(comp_pos)
+    cos = cos[..., : rotary_dim // 2]
+    sin = sin[..., : rotary_dim // 2]
+    pool_kv = pool.unsqueeze(2)  # [B, P, 1, Dh]
+    pool_kv = apply_interleaved_partial_rope(pool_kv, cos, sin, rotary_dim=rotary_dim)
+    pool_h = pool_kv.expand(B, P, H, Dh)
+
+    # Local mask (full causal in this test config; SWA disabled).
+    q_idx = torch.arange(S, device=hidden.device).unsqueeze(1)
+    k_idx = torch.arange(S, device=hidden.device).unsqueeze(0)
+    local_mask = torch.where(q_idx >= k_idx, 0.0, float("-inf")).to(hidden.dtype)
+
+    # Compressed-pool causal mask: pool s allowed for query t iff (s+1)*ratio - 1 <= t.
+    s_end = (torch.arange(P, device=hidden.device).unsqueeze(0) + 1) * ratio - 1
+    extra_mask = torch.where(s_end <= q_idx, 0.0, float("-inf")).to(hidden.dtype)
+
+    full_mask = torch.cat([local_mask, extra_mask], dim=-1)  # [S, S+P]
+
+    # Concat keys / values.
+    k_full = torch.cat([k_local, pool_h], dim=1)  # [B, S+P, H, Dh]
+    v_full = torch.cat([v_local, pool_h], dim=1)
+
+    q_bh = q.transpose(1, 2)  # [B, H, S, Dh]
+    k_bh = k_full.transpose(1, 2)
+    v_bh = v_full.transpose(1, 2)
+
+    # Plain 1/sqrt(head_dim): the YaRN magnitude factor is not part of the
+    # softmax temperature (``inference/model.py`` uses ``head_dim ** -0.5``).
+    scale = 1.0 / math.sqrt(Dh)
+    logits = torch.matmul(q_bh.float(), k_bh.float().transpose(-2, -1)) * scale
+    logits = logits + full_mask
+
+    if attn.attn_sink is None:
+        logits = logits - logits.amax(dim=-1, keepdim=True).detach()
+        probs = logits.softmax(dim=-1).to(v_bh.dtype)
+    else:
+        sink_col = attn.attn_sink.float().view(1, H, 1, 1).expand(B, H, S, 1)
+        logits_aug = torch.cat([logits, sink_col], dim=-1)
+        logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach()
+        probs = logits_aug.softmax(dim=-1)[..., :-1].to(v_bh.dtype)
+
+    out = torch.matmul(probs, v_bh.float()).to(hidden.dtype)
+    out = out.transpose(1, 2).contiguous()  # [B, S, H, Dh]
+
+    # Inverse RoPE (position = -t) using this layer's compress-base RoPE.
+    out = _inverse_rope_reference(
+        out,
+        position_ids=position_ids,
+        rope=attn.rope,
+        rotary_dim=rotary_dim,
+        compress_ratio=ratio,
+    )
+
+    # Grouped low-rank O.
+    G, r = attn.o_groups, attn.o_lora_rank
+    out_g = out.reshape(B, S, G, (H * Dh) // G)
+    wo_a_w = attn.linear_o_a.weight.view(G, r, (H * Dh) // G)
+    o = torch.einsum("bsgd,grd->bsgr", out_g, wo_a_w)
+    o = o.flatten(2)
+    return o @ attn.linear_o_b.weight.t()
+
+
+def test_hca_forward_matches_inline_reference():
+    """HCA forward agrees with an inline reference of the same math.
+
+    Both implementations apply the same partial-interleaved RoPE on the
+    compressed pool, the same compressed-causal mask, and the same joint
+    softmax-with-sink, so the agreement should be at machine precision.
+    """
+    torch.manual_seed(0)
+    compress_ratio = 128
+    B, S = 1, compress_ratio  # P = 1
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio)
+    attn = attn.to(_TEST_DTYPE)
+    with torch.no_grad():
+        attn.attn_sink.copy_(torch.linspace(-0.5, 0.5, attn.num_heads))
+
+    hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    with torch.no_grad():
+        ours = attn(hidden, position_ids)
+        ref = _reference_hca_forward(
+            attn=attn,
+            hidden=hidden,
+            position_ids=position_ids,
+        )
+
+    diff = (ours - ref).abs().max().item()
+    assert diff < 1e-3, f"HCA forward mismatch: max abs diff = {diff:.3e}"
+
+
+def test_inverse_rope_round_trips():
+    """``R(-t)`` exactly undoes ``R(t)`` on the rotated tail.
+
+    Pure math property of the helper, independent of any attention branch.
+    """
+    torch.manual_seed(0)
+    # H must be a power of two: the Triton RoPE kernel tiles the head axis with
+    # BLOCK_H = _pick_block_h(H) and feeds it to tl.arange, which rejects
+    # non-power-of-2 ranges. Production runs H=64 (Q) / H=1 (KV), so this only
+    # ever bites contrived test shapes.
+    B, S, H, Dh, rotary_dim = 2, 6, 4, 16, 8
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=H,
+        head_dim=Dh,
+        rotary_dim=rotary_dim,
+        q_lora_rank=32,
+        o_groups=1,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    attn = _make_attention(config).to(_TEST_DTYPE)
+
+    x = torch.randn(B, S, H, Dh, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    rotated = attn.rope.apply_rope(x, position_ids=position_ids, compress_ratio=0)
+    restored = attn._apply_inverse_rope(rotated, position_ids)
+
+    torch.testing.assert_close(restored, x, rtol=1e-5, atol=1e-5)
+    # The rotation must be non-trivial, otherwise the round-trip is vacuous.
+    assert not torch.allclose(rotated, x, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("compress_ratio", [0, 4, 128])
+def test_output_is_derotated_before_o_projection(compress_ratio):
+    """Every branch de-rotates the attention output before ``wo_a``.
+
+    Guards the wiring rather than the math: ``_project_output`` must apply the
+    inverse rotation, so feeding it a tensor and comparing against an explicit
+    de-rotate + O projection has to agree, and skipping the de-rotation must
+    not.
+    """
+    torch.manual_seed(0)
+    B, S = 2, 8
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    if compress_ratio:
+        attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio)
+    else:
+        attn = _make_attention(config)
+    attn = attn.to(_TEST_DTYPE)
+
+    core_out = torch.randn(B, S, attn.num_heads, attn.head_dim, dtype=_TEST_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    with torch.no_grad():
+        ours = attn._project_output(core_out, position_ids)
+
+        derotated = _inverse_rope_reference(
+            core_out,
+            position_ids=position_ids,
+            rope=attn.rope,
+            rotary_dim=attn.rotary_dim,
+            compress_ratio=compress_ratio,
+        )
+        expected = attn._grouped_o_projection(derotated)
+        without_fix = attn._grouped_o_projection(core_out)
+
+    torch.testing.assert_close(ours, expected, rtol=1e-4, atol=1e-4)
+    assert not torch.allclose(ours, without_fix, rtol=1e-3, atol=1e-3)
+
+
+def test_inverse_rope_matches_the_negated_sine_formulation():
+    """The fused de-rotation equals the explicit ``-sin`` rotation.
+
+    ``_apply_inverse_rope`` negates the *angle* inside the fused kernel rather
+    than building cos / sin eagerly and negating the sine, which saves three
+    launches and a materialised cos / sin tensor per layer. ``cos`` is even and
+    ``sin`` odd, so the two are the same rotation -- pinned here because the
+    optimisation is only safe as long as that stays true.
+    """
+    torch.manual_seed(0)
+    B, S, H, Dh, rotary_dim = 2, 6, 4, 16, 8
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=H,
+        head_dim=Dh,
+        rotary_dim=rotary_dim,
+        q_lora_rank=32,
+        o_groups=1,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    attn = _make_attention(config).to(_TEST_DTYPE)
+
+    x = torch.randn(B, S, H, Dh, dtype=_TEST_DTYPE)
+    # Offset positions so position 0 (where the rotation is the identity) does
+    # not dominate the comparison.
+    position_ids = torch.arange(1, S + 1).unsqueeze(0).expand(B, S)
+
+    with torch.no_grad():
+        ours = attn._apply_inverse_rope(x, position_ids)
+        ref = _inverse_rope_reference(x, position_ids=position_ids, rope=attn.rope, rotary_dim=rotary_dim)
+
+    torch.testing.assert_close(ours, ref, rtol=1e-5, atol=1e-5)
+    assert not torch.allclose(ours, x, rtol=1e-3, atol=1e-3)
+
+
+def test_inverse_rope_does_not_mutate_core_attention_output():
+    """De-rotation is out-of-place.
+
+    The attention backward keeps a reference to the core output, so mutating
+    it in place would corrupt the gradient.
+    """
+    torch.manual_seed(0)
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=False,
+    )
+    attn = _make_attention(config).to(_TEST_DTYPE)
+
+    core_out = torch.randn(2, 5, attn.num_heads, attn.head_dim, dtype=_TEST_DTYPE)
+    before = core_out.clone()
+    position_ids = torch.arange(5).unsqueeze(0).expand(2, 5)
+
+    with torch.no_grad():
+        rotated = attn._apply_inverse_rope(core_out, position_ids)
+
+    assert rotated is not core_out
+    torch.testing.assert_close(core_out, before, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(("compress_ratio", "num_pool_entries"), [(4, 4), (128, 2)])
+def test_compressed_pool_rope_uses_original_sequence_positions(compress_ratio, num_pool_entries):
+    """Compressed entry ``s`` is rotated at original token ``s * ratio``.
+
+    The queries are rotated at their own original positions, so rotating the
+    compressed KV at bare block indices would put the two sides on different
+    coordinate systems. ``inference/model.py`` slices ``freqs_cis[:cutoff:ratio]``
+    for prefill and indexes ``start_pos + 1 - ratio`` for decode; both land on
+    ``s * ratio``.
+    """
+    from primus.backends.megatron.core.transformer.dual_rope import (
+        apply_interleaved_partial_rope,
+    )
+
+    torch.manual_seed(0)
+    S = compress_ratio * num_pool_entries
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio).to(_TEST_DTYPE)
+    hidden = torch.randn(1, S, config.hidden_size, dtype=_TEST_DTYPE)
+
+    with torch.no_grad():
+        pooled_raw = attn.compressor(hidden)  # [B, P, Dh], pre-RoPE
+        ours = attn._build_compressed_pool(hidden)
+
+    P = pooled_raw.shape[1]
+    assert P == num_pool_entries, f"expected P={num_pool_entries}, got {P}"
+    rotary_dim = attn.rotary_dim
+
+    def _rope_at(positions: torch.Tensor) -> torch.Tensor:
+        cos, sin = attn.rope.compress_rope(positions)
+        cos = cos[..., : rotary_dim // 2].unsqueeze(0)
+        sin = sin[..., : rotary_dim // 2].unsqueeze(0)
+        rotated = apply_interleaved_partial_rope(pooled_raw.unsqueeze(2), cos, sin, rotary_dim=rotary_dim)
+        return rotated.squeeze(2)
+
+    block_idx = torch.arange(P, device=hidden.device)
+    expected = _rope_at(block_idx * compress_ratio)
+    torch.testing.assert_close(ours, expected, rtol=1e-4, atol=1e-4)
+
+    # The two conventions must actually differ, or the assertion above is
+    # vacuous (they coincide only at P == 1, which is why P > 1 here).
+    wrong = _rope_at(block_idx)
+    assert not torch.allclose(expected, wrong, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("compress_ratio", [0, 4, 128])
+def test_attention_scale_is_pure_inv_sqrt_head_dim(compress_ratio):
+    """Softmax temperature is ``1/sqrt(head_dim)`` on every branch.
+
+    The released ``inference/model.py`` sets ``softmax_scale = head_dim ** -0.5``
+    unconditionally; the YaRN magnitude factor (``m_scale``) must not be folded
+    into the attention temperature even on the compressed layers, which run
+    with a YaRN-scaled RoPE. A non-unit ``yarn_factor`` is used here on purpose
+    so the assertion fails if ``m_scale`` ever leaks back in.
+    """
+    head_dim = 16
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=head_dim,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    yarn_factor = 16.0
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=yarn_factor,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    attn = DeepseekV4Attention(
+        config,
+        rope=rope,
+        compress_ratio=compress_ratio,
+        submodules=None,
+    )
+
+    expected = 1.0 / math.sqrt(head_dim)
+    assert attn._attention_scale() == pytest.approx(expected, rel=1e-12)
+
+    if compress_ratio:
+        # Guard the regression directly: the compressed RoPE really does carry
+        # a non-unit m_scale, so the assertion above is not vacuous.
+        assert rope.attn_scale(compress_ratio=compress_ratio) > 1.0
+
+
+# ---------------------------------------------------------------------------
+# Spec / TP wiring tests (no torch.distributed required — these check the
+# spec-tree shape, not actual TP execution).
+# ---------------------------------------------------------------------------
+
+
+def test_attention_spec_uses_column_and_row_parallel():
+    """V4-faithful attention spec sources ``linear_q_up_proj`` from the
+    provider's column-parallel linear, and ``linear_o_b`` /
+    ``linear_proj`` from the row-parallel linear.
+
+    This is the contract that lets TP > 1 actually shard the projection
+    weights at runtime; at TP = 1 it is functionally identical to the
+    duplicated spec.
+    """
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+    )
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import (
+        _build_v4_attention_submodules,
+    )
+
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    provider = DeepSeekV4SpecProvider(config=config)
+    submods = _build_v4_attention_submodules(
+        config=config,
+        provider=provider,
+        compress_ratio=0,
+    )
+    # Plan-3 P21: ``gather_output=True`` and ``input_is_parallel=False``
+    # must route to the upstream non-TE classes because the TE wrappers
+    # explicitly reject those flags.
+    assert submods.linear_q_up_proj is not None
+    assert submods.linear_q_up_proj.module is provider.column_parallel_linear_with_gather_output()
+    assert submods.linear_q_up_proj.params.get("gather_output") is True
+    assert submods.linear_o_b is not None
+    assert submods.linear_o_b.module is provider.row_parallel_linear_with_scatter_input()
+    assert submods.linear_o_b.params.get("input_is_parallel") is False
+
+    # Flat-O fallback path also goes through row-parallel.
+    cfg_flat = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=0,
+        attn_sink=True,
+    )
+    submods_flat = _build_v4_attention_submodules(
+        config=cfg_flat,
+        provider=provider,
+        compress_ratio=0,
+    )
+    # Same reasoning as ``linear_o_b``: ``input_is_parallel=False`` must
+    # route to the upstream non-TE class.
+    assert submods_flat.linear_proj is not None
+    assert submods_flat.linear_proj.module is provider.row_parallel_linear_with_scatter_input()
+
+
+def test_attention_spec_includes_compressor_and_indexer():
+    """Compressed branches expose ``compressor`` (always) and ``indexer``
+    (CSA only) as :class:`ModuleSpec`s in the V4 attention submodules."""
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+    )
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import (
+        _build_v4_attention_submodules,
+    )
+    from primus.backends.megatron.core.transformer.compressor import Compressor
+    from primus.backends.megatron.core.transformer.indexer import Indexer
+
+    config = _make_v4_config(
+        hidden_size=64,
+        num_heads=4,
+        head_dim=16,
+        rotary_dim=8,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sink=True,
+    )
+    provider = DeepSeekV4SpecProvider(config=config)
+
+    dense = _build_v4_attention_submodules(
+        config=config,
+        provider=provider,
+        compress_ratio=0,
+    )
+    assert dense.compressor is None
+    assert dense.indexer is None
+
+    hca = _build_v4_attention_submodules(
+        config=config,
+        provider=provider,
+        compress_ratio=128,
+    )
+    assert hca.compressor is not None
+    assert hca.compressor.module is Compressor
+    assert hca.indexer is None  # HCA has no Indexer
+
+    csa = _build_v4_attention_submodules(
+        config=config,
+        provider=provider,
+        compress_ratio=4,
+    )
+    assert csa.compressor is not None
+    assert csa.compressor.module is Compressor
+    assert csa.indexer is not None
+    assert csa.indexer.module is Indexer
+
+
+# ---------------------------------------------------------------------------
+# TP=2 sharding parity scaffold (skips on CPU / single-rank).
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(
+    not (hasattr(torch, "distributed") and torch.distributed.is_available()),
+    reason="torch.distributed not available",
+)
+def test_tp2_sharding_parity_scaffold():
+    """Scaffold for the TP=2 sharding-parity test.
+
+    Skipped unless ``torch.distributed`` is initialized with
+    ``world_size >= 2``. When run under ``torchrun --nproc_per_node=2``
+    with the PrimusTurbo provider, this test will assert that the
+    column-parallel ``linear_q_up_proj`` + row-parallel ``linear_o_b``
+    pair produces output identical (≤1e-4) to a duplicated baseline.
+
+    Implementation deferred to P14 (full TP=2 smoke matrix).
+    """
+    if not torch.distributed.is_initialized():
+        pytest.skip("torch.distributed not initialized")
+    if torch.distributed.get_world_size() < 2:
+        pytest.skip("TP=2 parity test requires world_size >= 2")
+    pytest.skip("TP=2 sharding-parity check implementation tracked in P14")
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py
new file mode 100644
index 000000000..2b69a1e55
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py
@@ -0,0 +1,160 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Small-kernel-fusion 2026-07-03 — fused Triton HyperConnection ``collapse``.
+
+Pins :class:`HCCollapseFn` against the eager
+``(pre.unsqueeze(-1) * x).sum(-2)`` reference AND the actual
+:meth:`HyperMixer.collapse` call site, FWD + BWD, across dtypes and K.
+
+GPU-only; CPU runs are skipped at collection time.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("fused HC collapse Triton kernel requires CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.hyper_connection import (  # noqa: E402
+    HyperMixer,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_collapse import (  # noqa: E402
+    HCCollapseFn,
+    eager_hc_collapse,
+    hc_collapse_triton,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+
+@contextmanager
+def _env(key, value):
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _tol(dtype):
+    return {
+        torch.float32: (1e-5, 1e-5),
+        torch.float16: (2e-3, 2e-3),
+        torch.bfloat16: (1e-2, 1e-2),
+    }[dtype]
+
+
+def _mk(shape, dtype, seed, requires_grad=False):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(shape, dtype=dtype, device="cuda", generator=gen)
+    if requires_grad:
+        x.requires_grad_(True)
+    return x
+
+
+def _assert_at_least_as_accurate(got, x, pre, dtype):
+    """For low-precision dtypes the eager path rounds each ``pre*x`` product to
+    the input dtype before summing, while the kernel accumulates in fp32. Rather
+    than pin the kernel to the *less* accurate eager path, pin FWD parity in
+    fp32 and, for fp16/bf16, assert the kernel matches an fp32 gold within the
+    output dtype's rounding (i.e. the kernel is at least as accurate as eager).
+    """
+    if dtype == torch.float32:
+        ref = eager_hc_collapse(x, pre)
+        torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5)
+        return
+    gold = (pre.float().unsqueeze(-1) * x.float()).sum(dim=-2)
+    atol, rtol = _tol(dtype)
+    # allow output-rounding slack on top of the relative tolerance
+    torch.testing.assert_close(got.float(), gold, atol=3 * atol, rtol=rtol)
+
+
+class TestFwd:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    @pytest.mark.parametrize("K", [1, 2, 4, 8])
+    @pytest.mark.parametrize("D", [64, 512, 4096])
+    def test_parity(self, dtype, K, D):
+        x = _mk((2, 130, K, D), dtype, seed=1)
+        pre = _mk((2, 130, K), dtype, seed=2)
+        got = HCCollapseFn.apply(x, pre)
+        _assert_at_least_as_accurate(got, x, pre, dtype)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_call_site_parity(self, dtype):
+        x = _mk((3, 40, 4, 4096), dtype, seed=3)
+        pre = _mk((3, 40, 4), dtype, seed=4)
+        got = HyperMixer.collapse(x, pre)
+        _assert_at_least_as_accurate(got, x, pre, dtype)
+
+
+class TestBwd:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("K", [2, 4])
+    def test_bwd_parity(self, dtype, K):
+        D = 512
+        xb = _mk((2, 64, K, D), dtype, seed=10)
+        pb = _mk((2, 64, K), dtype, seed=11)
+        xt = xb.detach().clone().requires_grad_(True)
+        pt = pb.detach().clone().requires_grad_(True)
+        xe = xb.detach().clone().requires_grad_(True)
+        pe = pb.detach().clone().requires_grad_(True)
+
+        out_t = HCCollapseFn.apply(xt, pt)
+        out_e = eager_hc_collapse(xe, pe)
+        g = torch.randn_like(out_t)
+        out_t.backward(g)
+        out_e.backward(g.detach().clone())
+
+        if dtype == torch.float32:
+            torch.testing.assert_close(xt.grad, xe.grad, atol=1e-5, rtol=1e-5)
+            torch.testing.assert_close(pt.grad, pe.grad, atol=1e-4, rtol=1e-4)
+            return
+        # Low precision: compare against fp32 gold (kernel accumulates in fp32,
+        # so it is at least as accurate as the bf16 eager reference).
+        xg = xb.detach().float().requires_grad_(True)
+        pg = pb.detach().float().requires_grad_(True)
+        out_g = eager_hc_collapse(xg, pg)
+        out_g.backward(g.float())
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(xt.grad.float(), xg.grad, atol=3 * atol, rtol=rtol)
+        torch.testing.assert_close(pt.grad.float(), pg.grad, atol=5e-2, rtol=5e-2)
+
+
+class TestDispatch:
+    def test_env_flag(self):
+        x = _mk((2, 32, 4, 512), torch.bfloat16, seed=20)
+        pre = _mk((2, 32, 4), torch.bfloat16, seed=21)
+        with _env("PRIMUS_HC_COLLAPSE_TRITON", "1"):
+            assert is_triton_path_enabled()
+            on = hc_collapse_triton(x, pre)
+        with _env("PRIMUS_HC_COLLAPSE_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            off = hc_collapse_triton(x, pre)
+        # triton (fp32 accum) vs eager (bf16 accum): both valid bf16
+        # approximations of the same op, so allow K-way bf16 rounding slack.
+        torch.testing.assert_close(on, off, atol=6e-2, rtol=6e-2)
+
+    def test_support_predicate(self):
+        x = _mk((2, 4, 512), torch.bfloat16, seed=22)
+        pre = _mk((2, 4), torch.bfloat16, seed=23)
+        assert is_triton_kernel_supported(x, pre)
+        assert not is_triton_kernel_supported(x.cpu(), pre.cpu())
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py
new file mode 100644
index 000000000..d055a1bd1
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py
@@ -0,0 +1,233 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Small-kernel-fusion 2026-07-03 — fused Triton RMSNorm FWD/BWD parity.
+
+Pins :class:`FusedRMSNormFn` (Triton kernel in
+``primus...v4_attention_kernels._triton_common.rmsnorm``) against the eager
+RMSNorm reference AND the actual model call sites it replaces:
+
+* ``_per_head_rms_norm``               (parameter-less, out=in_dtype)
+* ``LocalRMSNorm``                     (weighted + weight grad, mid-cast)
+* ``HyperMixer._packed_logits`` RMS    (parameter-less, out=fp32)
+* ``HyperHead`` RMS                    (parameter-less, out=fp32)
+
+GPU-only; CPU runs are skipped at collection time.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("fused RMSNorm Triton kernel requires CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+# Import the model package first so the deepseek_v4_attention <-> block cyclic
+# import resolves cleanly before we pull `_per_head_rms_norm` off the leaf.
+import primus.backends.megatron.core.models.deepseek_v4  # noqa: E402,F401
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    _per_head_rms_norm,
+)
+from primus.backends.megatron.core.transformer.local_rmsnorm import (  # noqa: E402
+    LocalRMSNorm,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import (  # noqa: E402
+    FusedRMSNormFn,
+    eager_rms_norm,
+    fused_rms_norm,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+
+@contextmanager
+def _env(key, value):
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _tol(dtype):
+    return {
+        torch.float32: (1e-5, 1e-5),
+        torch.float16: (2e-3, 2e-3),
+        torch.bfloat16: (1e-2, 1e-2),
+    }[dtype]
+
+
+def _mk(shape, dtype, seed, requires_grad=False):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(shape, dtype=dtype, device="cuda", generator=gen)
+    if requires_grad:
+        x.requires_grad_(True)
+    return x
+
+
+# ---------------------------------------------------------------------------
+# FWD parity vs eager reference — the four site contracts.
+# ---------------------------------------------------------------------------
+
+
+class TestFwdParity:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    @pytest.mark.parametrize("D", [128, 512, 4096, 16384])
+    def test_no_weight_in_dtype(self, dtype, D):
+        """per-head RMS contract: no weight, out=in_dtype."""
+        x = _mk((64, D), dtype, seed=1)
+        out_t = FusedRMSNormFn.apply(x, None, 1e-6, False, dtype)
+        out_e = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=dtype)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("D", [512, 4096])
+    def test_weighted_mid_cast(self, dtype, D):
+        """LocalRMSNorm contract: weight, mid-cast, out=promote(in, weight)."""
+        x = _mk((128, D), dtype, seed=2)
+        w = _mk((D,), torch.float32, seed=3)
+        out_dtype = torch.promote_types(dtype, torch.float32)
+        out_t = FusedRMSNormFn.apply(x, w, 1e-6, True, out_dtype)
+        out_e = eager_rms_norm(x, w, eps=1e-6, mid_cast=True, out_dtype=out_dtype)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_no_weight_fp32_out(self, dtype):
+        """HyperMixer / HyperHead RMS contract: no weight, out=fp32."""
+        x = _mk((256, 16384), dtype, seed=4)
+        out_t = FusedRMSNormFn.apply(x, None, 1e-6, False, torch.float32)
+        out_e = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=torch.float32)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# Parity vs the ACTUAL call sites (the code paths this kernel replaces).
+# ---------------------------------------------------------------------------
+
+
+class TestCallSiteParity:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_per_head_rms_norm(self, dtype):
+        # [B, S, H, head_dim] as in _apply_q.
+        x = _mk((1, 128, 8, 512), dtype, seed=10)
+        ref = _per_head_rms_norm(x, eps=1e-6)
+        got = fused_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=dtype)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(got, ref, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_local_rmsnorm_forward(self, dtype):
+        norm = LocalRMSNorm(512, eps=1e-6).cuda()
+        with torch.no_grad():
+            norm.weight.normal_()
+        x = _mk((16, 64, 512), dtype, seed=11)
+        ref = norm(x)
+        out_dtype = torch.promote_types(dtype, norm.weight.dtype)
+        got = fused_rms_norm(x, norm.weight, eps=1e-6, mid_cast=True, out_dtype=out_dtype)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(got, ref, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# BWD parity vs eager autograd.
+# ---------------------------------------------------------------------------
+
+
+class TestBwdParity:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("D", [512, 4096])
+    def test_no_weight_bwd(self, dtype, D):
+        xb = _mk((128, D), dtype, seed=20)
+        xt = xb.detach().clone().requires_grad_(True)
+        xe = xb.detach().clone().requires_grad_(True)
+        out_t = FusedRMSNormFn.apply(xt, None, 1e-6, False, dtype)
+        out_e = eager_rms_norm(xe, None, eps=1e-6, mid_cast=False, out_dtype=dtype)
+        g = torch.randn_like(out_t)
+        out_t.backward(g)
+        out_e.backward(g.detach().clone())
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_weighted_bwd_x_and_w(self, dtype):
+        D = 512
+        xb = _mk((256, D), dtype, seed=21)
+        wb = _mk((D,), torch.float32, seed=22)
+        out_dtype = torch.promote_types(dtype, torch.float32)
+
+        xt = xb.detach().clone().requires_grad_(True)
+        wt = wb.detach().clone().requires_grad_(True)
+        xe = xb.detach().clone().requires_grad_(True)
+        we = wb.detach().clone().requires_grad_(True)
+
+        out_t = FusedRMSNormFn.apply(xt, wt, 1e-6, True, out_dtype)
+        out_e = eager_rms_norm(xe, we, eps=1e-6, mid_cast=True, out_dtype=out_dtype)
+        g = torch.randn_like(out_t)
+        out_t.backward(g)
+        out_e.backward(g.detach().clone())
+
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol)
+        # Weight grad is a cross-row reduction; the eager reference rounds the
+        # upstream grad to bf16 per row while the kernel accumulates in fp32, so
+        # for bf16 the per-element diff is dominated by reduction noise where the
+        # true sum has cancellation. Compare against an fp32 "gold" reduction to
+        # confirm the kernel is at least as accurate as eager.
+        if dtype == torch.float32:
+            torch.testing.assert_close(wt.grad, we.grad, atol=1e-4, rtol=1e-4)
+        else:
+            gold = (
+                g.float() * (xb.float() * torch.rsqrt(xb.float().pow(2).mean(-1, keepdim=True) + 1e-6))
+            ).sum(0)
+            err_kernel = (wt.grad.float() - gold).abs().max()
+            err_eager = (we.grad.float() - gold).abs().max()
+            assert err_kernel <= err_eager + 1e-3, (err_kernel, err_eager)
+
+
+# ---------------------------------------------------------------------------
+# Dispatch / edge cases.
+# ---------------------------------------------------------------------------
+
+
+class TestDispatch:
+    def test_env_flag(self):
+        x = _mk((32, 512), torch.bfloat16, seed=30)
+        with _env("PRIMUS_RMSNORM_TRITON", "1"):
+            assert is_triton_path_enabled()
+            on = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.bfloat16)
+        with _env("PRIMUS_RMSNORM_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            off = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.bfloat16)
+        torch.testing.assert_close(on, off, atol=1e-2, rtol=1e-2)
+
+    def test_support_predicate(self):
+        good = _mk((4, 512), torch.bfloat16, seed=31)
+        assert is_triton_kernel_supported(good, None)
+        assert not is_triton_kernel_supported(good.cpu(), None)
+
+    def test_cpu_falls_back(self):
+        x = torch.randn(4, 512, dtype=torch.float32)
+        # CPU input: dispatcher must not raise, returns eager result.
+        out = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.float32)
+        ref = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=torch.float32)
+        torch.testing.assert_close(out, ref)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py
new file mode 100644
index 000000000..2eedb0fd4
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py
@@ -0,0 +1,196 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Small-kernel-fusion 2026-07-03 — fused RoPE-from-positions parity.
+
+:class:`RoPEFromPositionsFn` computes cos/sin in-kernel from
+``(position_ids, inv_freq)`` (instead of consuming precomputed cos/sin).
+Pins it FWD + BWD against the eager ``cos = pos*inv_freq -> rotate`` path
+AND against :class:`RoPEInterleavedPartialFn` (which already matches eager),
+plus an end-to-end check through :meth:`DualRoPE.apply_rope`.
+
+GPU-only; CPU runs are skipped at collection time.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("fused RoPE kernel requires CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.dual_rope import DualRoPE  # noqa: E402
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import (  # noqa: E402
+    RoPEFromPositionsFn,
+    apply_rope_from_positions,
+    eager_apply_interleaved_partial_rope,
+)
+
+
+def _tol(dtype):
+    return {torch.float32: (1e-5, 1e-5), torch.bfloat16: (1e-2, 1e-2)}[dtype]
+
+
+def _inv_freq(rotary_dim, theta=10000.0):
+    i = torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cuda")
+    return 1.0 / (theta ** (i / rotary_dim))
+
+
+def _eager(x, position_ids, inv_freq, rotary_dim, inverse=False):
+    freqs = position_ids.float().unsqueeze(-1) * inv_freq
+    if inverse:
+        freqs = -freqs
+    return eager_apply_interleaved_partial_rope(x, freqs.cos(), freqs.sin(), rotary_dim=rotary_dim)
+
+
+class TestFwd:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("H", [1, 64])
+    @pytest.mark.parametrize("rotary_dim", [64, 128])
+    def test_parity(self, dtype, H, rotary_dim):
+        B, S, D = 2, 128, 512
+        gen = torch.Generator(device="cuda").manual_seed(1)
+        x = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen)
+        pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S)
+        inv = _inv_freq(rotary_dim)
+        got = RoPEFromPositionsFn.apply(x, pos.broadcast_to(B, S).reshape(-1), inv, rotary_dim)
+        ref = _eager(x, pos, inv, rotary_dim)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(got, ref, atol=atol, rtol=rtol)
+
+    def test_broadcast_positions_1d(self):
+        B, S, H, D, rd = 3, 64, 8, 512, 64
+        x = torch.randn((B, S, H, D), dtype=torch.float32, device="cuda")
+        pos1d = torch.arange(S, device="cuda")  # [S] -> broadcast to [B, S]
+        inv = _inv_freq(rd)
+        got = apply_rope_from_positions(x, pos1d, inv, rotary_dim=rd)
+        ref = _eager(x, pos1d.unsqueeze(0).expand(B, S), inv, rd)
+        torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5)
+
+
+class TestInverse:
+    """``inverse=True`` is the de-rotation the V4 attention output needs.
+
+    It exists so the ``R(-t)`` before the O projection reuses this fused path
+    (cos/sin generated in-kernel) instead of building cos/sin eagerly on all 43
+    layers.
+    """
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("rotary_dim", [64, 128])
+    def test_matches_eager_negated_angle(self, dtype, rotary_dim):
+        B, S, H, D = 2, 128, 8, 512
+        gen = torch.Generator(device="cuda").manual_seed(3)
+        x = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen)
+        pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S)
+        inv = _inv_freq(rotary_dim)
+
+        got = apply_rope_from_positions(x, pos, inv, rotary_dim=rotary_dim, inverse=True)
+        ref = _eager(x, pos, inv, rotary_dim, inverse=True)
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(got, ref, atol=atol, rtol=rtol)
+
+    def test_round_trips_with_the_forward_rotation(self):
+        B, S, H, D, rd = 2, 64, 4, 512, 64
+        x = torch.randn((B, S, H, D), dtype=torch.float32, device="cuda")
+        pos = torch.arange(S, device="cuda")
+        inv = _inv_freq(rd)
+
+        rotated = apply_rope_from_positions(x, pos, inv, rotary_dim=rd)
+        back = apply_rope_from_positions(rotated, pos, inv, rotary_dim=rd, inverse=True)
+        torch.testing.assert_close(back, x, atol=1e-5, rtol=1e-5)
+
+        # Guard against a vacuous round trip.
+        assert not torch.allclose(rotated, x, atol=1e-3, rtol=1e-3)
+
+    def test_is_out_of_place(self):
+        """The attention backward retains the raw output, so it must not be mutated."""
+        B, S, H, D, rd = 1, 32, 4, 512, 64
+        x = torch.randn((B, S, H, D), dtype=torch.float32, device="cuda")
+        before = x.detach().clone()
+        out = apply_rope_from_positions(
+            x, torch.arange(S, device="cuda"), _inv_freq(rd), rotary_dim=rd, inverse=True
+        )
+
+        assert out.data_ptr() != x.data_ptr()
+        torch.testing.assert_close(x, before, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_bwd_matches_eager(self, dtype):
+        B, S, H, D, rd = 2, 64, 8, 512, 64
+        gen = torch.Generator(device="cuda").manual_seed(4)
+        xb = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen)
+        pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S)
+        inv = _inv_freq(rd)
+
+        xt = xb.detach().clone().requires_grad_(True)
+        xe = xb.detach().clone().requires_grad_(True)
+        out_t = apply_rope_from_positions(xt, pos, inv, rotary_dim=rd, inverse=True)
+        out_e = _eager(xe, pos, inv, rd, inverse=True)
+        g = torch.randn_like(out_t)
+        out_t.backward(g)
+        out_e.backward(g.detach().clone())
+
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol)
+
+    def test_inverse_differs_from_forward(self):
+        """Sanity: the flag actually changes the result."""
+        B, S, H, D, rd = 1, 32, 4, 512, 64
+        x = torch.randn((B, S, H, D), dtype=torch.float32, device="cuda")
+        pos = torch.arange(1, S + 1, device="cuda")  # avoid position 0 (identity)
+        inv = _inv_freq(rd)
+
+        fwd = apply_rope_from_positions(x, pos, inv, rotary_dim=rd)
+        inv_out = apply_rope_from_positions(x, pos, inv, rotary_dim=rd, inverse=True)
+        assert not torch.allclose(fwd, inv_out, atol=1e-4, rtol=1e-4)
+
+
+class TestBwd:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("H", [1, 64])
+    def test_bwd_parity(self, dtype, H):
+        B, S, D, rd = 2, 64, 512, 64
+        gen = torch.Generator(device="cuda").manual_seed(2)
+        xb = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen)
+        pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S)
+        inv = _inv_freq(rd)
+
+        xt = xb.detach().clone().requires_grad_(True)
+        xe = xb.detach().clone().requires_grad_(True)
+        out_t = RoPEFromPositionsFn.apply(xt, pos.broadcast_to(B, S).reshape(-1), inv, rd)
+        out_e = _eager(xe, pos, inv, rd)
+        g = torch.randn_like(out_t)
+        out_t.backward(g)
+        out_e.backward(g.detach().clone())
+        atol, rtol = _tol(dtype)
+        torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol)
+
+
+class TestDualRopeIntegration:
+    @pytest.mark.parametrize("compress_ratio", [0, 4])
+    def test_apply_rope_matches_eager(self, compress_ratio):
+        rope = DualRoPE(
+            rotary_dim=64,
+            rope_theta=10000.0,
+            compress_rope_theta=160000.0,
+            yarn_factor=16.0,
+            yarn_beta_fast=32.0,
+            yarn_beta_slow=1.0,
+            original_max_position_embeddings=65536,
+        ).cuda()
+        B, S, H, D = 2, 128, 8, 512
+        x = torch.randn((B, S, H, D), dtype=torch.bfloat16, device="cuda")
+        pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S)
+
+        got = rope.apply_rope(x, position_ids=pos, compress_ratio=compress_ratio)
+        cache = rope.get_rope(compress_ratio=compress_ratio)
+        ref = _eager(x, pos, cache.inv_freq, 64)
+        torch.testing.assert_close(got, ref, atol=1e-2, rtol=1e-2)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py
new file mode 100644
index 000000000..2f526117c
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py
@@ -0,0 +1,267 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-6 P37 G40 — `HyperMixer.compute_weights` tail Triton parity.
+
+Asserts that :class:`HCComputeTailFn` (Triton kernel from
+``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue``)
+matches the eager body in
+``primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights``
+bit-for-bit-equivalent FWD and BWD at two tiers:
+
+* fast tier -- `B=2, S=64, K=4` exercising every code path (3 slices,
+  2 sigmoid, 1 softmax, scale + base, eps); parametrised over `K ∈ {1,
+  2, 4, 8}`;
+* release tier -- `B=1, S=4096, K=4` bf16 (V4-Flash production shape),
+  behind ``pytest.mark.slow``.
+
+Composed end-to-end ``HyperMixer.compute_weights`` parity at K=4 (the
+load-bearing test: routing through the public API matches the Triton
+path within the eager-fallback baseline within bf16 tolerance).
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip(
+        "hc_glue Triton kernel requires CUDA / HIP",
+        allow_module_level=True,
+    )
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.hyper_connection import (  # noqa: E402
+    HyperMixer,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import (  # noqa: E402
+    HCComputeTailFn,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+
+@contextmanager
+def _env(key: str, value: str | None):
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _eager_tail(
+    logits: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    *,
+    K: int,
+    eps: float,
+    out_dtype: torch.dtype,
+):
+    """Eager reference matching the pre-P37 tail body bit-for-bit."""
+    pre_logit = logits[..., :K] * scale[0] + base[:K]
+    post_logit = logits[..., K : 2 * K] * scale[1] + base[K : 2 * K]
+    comb_logit = logits[..., 2 * K :].view(*logits.shape[:-1], K, K) * scale[2] + base[2 * K :].view(K, K)
+    pre = torch.sigmoid(pre_logit) + eps
+    post = 2.0 * torch.sigmoid(post_logit)
+    comb = torch.softmax(comb_logit, dim=-1) + eps
+    return pre.to(out_dtype), post.to(out_dtype), comb.to(out_dtype)
+
+
+def _build_inputs(*, B: int, S: int, K: int, seed: int, requires_grad: bool = False):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    out_dim = (2 + K) * K
+    logits = torch.randn((B, S, out_dim), dtype=torch.float32, device="cuda", generator=gen)
+    scale = torch.ones(3, dtype=torch.float32, device="cuda")
+    scale = scale + 0.1 * torch.randn(3, dtype=torch.float32, device="cuda", generator=gen)
+    base = 0.01 * torch.randn(out_dim, dtype=torch.float32, device="cuda", generator=gen)
+    if requires_grad:
+        logits.requires_grad_(True)
+        scale.requires_grad_(True)
+        base.requires_grad_(True)
+    return logits, scale, base
+
+
+def _dtype_tolerance(dtype: torch.dtype):
+    if dtype == torch.float32:
+        return 1e-5, 1e-5
+    if dtype == torch.float16:
+        return 1e-3, 1e-3
+    if dtype == torch.bfloat16:
+        return 1e-2, 1e-2
+    raise ValueError(dtype)
+
+
+# ---------------------------------------------------------------------------
+# G40: FWD parity vs eager
+# ---------------------------------------------------------------------------
+
+
+class TestG40ForwardParity:
+    """FWD output matches eager within out_dtype tolerance."""
+
+    @pytest.mark.parametrize("K", [1, 2, 4, 8])
+    @pytest.mark.parametrize("out_dtype", [torch.bfloat16])
+    def test_fast_tier_fwd_eager_parity(self, K, out_dtype):
+        logits, scale, base = _build_inputs(B=2, S=64, K=K, seed=42 + K)
+
+        pre_t, post_t, comb_t = HCComputeTailFn.apply(logits, scale, base, K, 1e-6, out_dtype)
+        pre_e, post_e, comb_e = _eager_tail(logits, scale, base, K=K, eps=1e-6, out_dtype=out_dtype)
+
+        atol, rtol = _dtype_tolerance(out_dtype)
+        torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G40: BWD parity vs eager autograd
+# ---------------------------------------------------------------------------
+
+
+class TestG40BackwardParity:
+    """BWD parity vs eager autograd across K and dtypes."""
+
+    @pytest.mark.parametrize("K", [1, 2, 4, 8])
+    @pytest.mark.parametrize("out_dtype", [torch.bfloat16])
+    def test_fast_tier_bwd_eager_parity(self, K, out_dtype):
+        logits_e, scale_e, base_e = _build_inputs(B=2, S=64, K=K, seed=144 + K, requires_grad=True)
+        logits_t = logits_e.detach().clone().requires_grad_(True)
+        scale_t = scale_e.detach().clone().requires_grad_(True)
+        base_t = base_e.detach().clone().requires_grad_(True)
+
+        pre_t, post_t, comb_t = HCComputeTailFn.apply(logits_t, scale_t, base_t, K, 1e-6, out_dtype)
+        pre_e, post_e, comb_e = _eager_tail(logits_e, scale_e, base_e, K=K, eps=1e-6, out_dtype=out_dtype)
+
+        g_pre = torch.randn_like(pre_t)
+        g_post = torch.randn_like(post_t)
+        g_comb = torch.randn_like(comb_t)
+
+        (pre_t * g_pre).sum().add_((post_t * g_post).sum()).add_((comb_t * g_comb).sum()).backward()
+        (pre_e * g_pre).sum().add_((post_e * g_post).sum()).add_((comb_e * g_comb).sum()).backward()
+
+        atol, rtol = _dtype_tolerance(out_dtype)
+        # Bump tolerance one notch for d_scale (cross-term: O(N) sum
+        # of N elements amplifies rounding error linearly with N).
+        scale_atol = atol * 10 if out_dtype == torch.bfloat16 else atol
+        scale_rtol = rtol * 10 if out_dtype == torch.bfloat16 else rtol
+        torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=atol, rtol=rtol)
+        torch.testing.assert_close(base_t.grad, base_e.grad, atol=atol, rtol=rtol)
+        torch.testing.assert_close(scale_t.grad, scale_e.grad, atol=scale_atol, rtol=scale_rtol)
+
+
+# ---------------------------------------------------------------------------
+# G40: Composed end-to-end HyperMixer.compute_weights parity
+# ---------------------------------------------------------------------------
+
+
+class TestG40HyperMixerParity:
+    """End-to-end ``HyperMixer.compute_weights`` parity between the
+    default-on Triton path and the env=0 eager fallback (with the same
+    weights / inputs / Sinkhorn iters).
+    """
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16])
+    def test_compute_weights_env_dispatch_parity(self, dtype):
+        torch.manual_seed(20260514)
+        mixer = HyperMixer(hidden_size=64, hc_mult=4, sinkhorn_iters=20).to("cuda")
+        x = torch.randn(2, 16, 4, 64, dtype=dtype, device="cuda")
+
+        with _env("PRIMUS_HC_TRITON", "1"):
+            assert is_triton_path_enabled()
+            pre_t, post_t, comb_t = mixer.compute_weights(x.clone())
+        with _env("PRIMUS_HC_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            pre_e, post_e, comb_e = mixer.compute_weights(x.clone())
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G40: release-tier V4-Flash production shape
+# ---------------------------------------------------------------------------
+
+
+class TestG40ReleaseTier:
+    """V4-Flash production shape: ``B=1, S=4096, K=4``, bf16."""
+
+    @pytest.mark.slow
+    def test_v4_flash_fwd_bwd_parity(self):
+        K = 4
+        out_dtype = torch.bfloat16
+        logits_e, scale_e, base_e = _build_inputs(B=1, S=4096, K=K, seed=2200, requires_grad=True)
+        logits_t = logits_e.detach().clone().requires_grad_(True)
+        scale_t = scale_e.detach().clone().requires_grad_(True)
+        base_t = base_e.detach().clone().requires_grad_(True)
+
+        pre_t, post_t, comb_t = HCComputeTailFn.apply(logits_t, scale_t, base_t, K, 1e-6, out_dtype)
+        pre_e, post_e, comb_e = _eager_tail(logits_e, scale_e, base_e, K=K, eps=1e-6, out_dtype=out_dtype)
+
+        atol, rtol = _dtype_tolerance(out_dtype)
+        torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol)
+        torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol)
+
+        g_pre = torch.randn_like(pre_t)
+        g_post = torch.randn_like(post_t)
+        g_comb = torch.randn_like(comb_t)
+        (pre_t * g_pre).sum().add_((post_t * g_post).sum()).add_((comb_t * g_comb).sum()).backward()
+        (pre_e * g_pre).sum().add_((post_e * g_post).sum()).add_((comb_e * g_comb).sum()).backward()
+
+        # At V4-Flash sequence length (S=4096), d_scale accumulates
+        # B*S*K = 16384 cross-terms; bf16 rounding compounds linearly,
+        # so widen tolerance for scale.
+        torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=atol, rtol=rtol)
+        torch.testing.assert_close(base_t.grad, base_e.grad, atol=1e-1, rtol=1e-1)
+        torch.testing.assert_close(scale_t.grad, scale_e.grad, atol=1e-1, rtol=1e-1)
+
+
+# ---------------------------------------------------------------------------
+# G40: edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestG40EdgeCases:
+    def test_unsupported_k_raises(self):
+        logits, scale, base = _build_inputs(B=2, S=8, K=4, seed=33)
+        # Bad K -- pass K=5 (not in {1,2,4,8,16})
+        bad_K = 3
+        bad_logits = torch.randn(2, 8, (2 + bad_K) * bad_K, dtype=torch.float32, device="cuda")
+        bad_base = torch.zeros((2 + bad_K) * bad_K, dtype=torch.float32, device="cuda")
+        with pytest.raises(ValueError, match="unsupported K"):
+            HCComputeTailFn.apply(bad_logits, scale, bad_base, bad_K, 1e-6, torch.float32)
+
+    def test_bad_logits_shape_raises(self):
+        logits = torch.randn(2, 8, 99, dtype=torch.float32, device="cuda")  # bad last dim
+        scale = torch.ones(3, dtype=torch.float32, device="cuda")
+        base = torch.zeros(24, dtype=torch.float32, device="cuda")
+        with pytest.raises(ValueError, match="logits last-dim"):
+            HCComputeTailFn.apply(logits, scale, base, 4, 1e-6, torch.float32)
+
+    def test_kernel_supported_predicate(self):
+        good = torch.randn(2, 8, 24, dtype=torch.float32, device="cuda")
+        bad_k = torch.randn(2, 8, 51, dtype=torch.float32, device="cuda")  # K=? doesn't match
+        bad_dev = torch.randn(2, 8, 24, dtype=torch.float32)  # cpu
+        assert is_triton_kernel_supported(good, K=4)
+        assert not is_triton_kernel_supported(bad_k, K=4)
+        assert not is_triton_kernel_supported(bad_dev, K=4)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py
new file mode 100644
index 000000000..746583153
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py
@@ -0,0 +1,297 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-6 P41 G43 — `Indexer.forward` post-einsum tail Triton parity.
+
+Asserts that :class:`IndexerScorePostFn` (Triton tail kernel from
+``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post``)
+produces scores that match the eager **tail** (``relu + mul + sum(H)
++ causal_mask``, with the einsum kept eager) within bf16 tolerance,
+and that the load-bearing post-`topk` ``topk_idxs`` are bit-equal
+across the two paths.
+
+The einsum stays on cuBLAS / hipBLASLt; this gate covers only the
+post-matmul tail.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip(
+        "indexer Triton kernel requires CUDA / HIP",
+        allow_module_level=True,
+    )
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.indexer import Indexer  # noqa: E402
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import (  # noqa: E402
+    IndexerScorePostFn,
+    indexer_score_post_triton,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+
+@contextmanager
+def _env(key: str, value: str | None):
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _eager_tail(
+    dot: torch.Tensor,
+    w_i: torch.Tensor,
+    *,
+    compress_ratio: int,
+    out_dtype: torch.dtype,
+):
+    """Eager reference matching the P41-routed body exactly.
+
+    Takes ``dot [B, S, H, P]`` (pre-relu, the einsum output) and
+    ``w_i [B, S, H]``; returns ``scores [B, S, P]``.
+    """
+    relu_term = torch.nn.functional.relu(dot.float())
+    scores = (relu_term * w_i.float().unsqueeze(-1)).sum(dim=2)
+    B, S, P = scores.shape
+    t_idx = torch.arange(S, device=scores.device).unsqueeze(1)
+    s_end = (torch.arange(P, device=scores.device).unsqueeze(0) + 1) * compress_ratio - 1
+    allowed = s_end <= t_idx
+    mask = torch.where(
+        allowed,
+        torch.zeros_like(scores[0]),
+        torch.full_like(scores[0], float("-inf")),
+    )
+    scores = scores + mask.unsqueeze(0)
+    return scores.to(out_dtype)
+
+
+def _build_dot_inputs(*, B: int, S: int, P: int, H: int, dtype: torch.dtype, seed: int):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    dot = torch.randn((B, S, H, P), dtype=dtype, device="cuda", generator=gen)
+    w_i = torch.randn((B, S, H), dtype=dtype, device="cuda", generator=gen).abs()
+    return dot, w_i
+
+
+# ---------------------------------------------------------------------------
+# G43: FWD parity vs eager tail
+# ---------------------------------------------------------------------------
+
+
+class TestG43ForwardParity:
+    @pytest.mark.parametrize("H", [1, 2, 4, 8])
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("compress_ratio", [1, 4, 16])
+    def test_fast_tier_fwd_eager_parity(self, H, dtype, compress_ratio):
+        B, S, P = 1, 64, max(4, 64 // max(compress_ratio, 1))
+        dot, w = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=900 + H)
+
+        s_t = IndexerScorePostFn.apply(dot, w, compress_ratio, dtype)
+        s_e = _eager_tail(dot, w, compress_ratio=compress_ratio, out_dtype=dtype)
+
+        if dtype == torch.float32:
+            atol, rtol = 1e-5, 1e-5
+        else:
+            atol, rtol = 5e-3, 5e-3
+
+        finite_e = torch.isfinite(s_e)
+        finite_t = torch.isfinite(s_t)
+        torch.testing.assert_close(finite_e, finite_t)
+        torch.testing.assert_close(s_t[finite_t].float(), s_e[finite_e].float(), atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G43: topk parity (load-bearing) via Indexer module
+# ---------------------------------------------------------------------------
+
+
+class TestG43TopKParity:
+    """The post-topk indices must be bit-equal vs the eager full chain.
+
+    Mirrors G41 from P38: downstream CSA reads ``topk_idxs`` exactly;
+    any divergence breaks the sparse selection.
+    """
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_indexer_topk_parity_via_module(self, dtype):
+        torch.manual_seed(20260515)
+        B, S, P, H, HD, K = 1, 64, 16, 8, 32, 4
+        D = 128
+        indexer = Indexer(
+            hidden_size=D,
+            index_head_dim=HD,
+            index_n_heads=H,
+            index_topk=K,
+            compress_ratio=4,
+        ).to(device="cuda", dtype=dtype)
+        hidden = torch.randn((B, S, D), dtype=dtype, device="cuda")
+
+        with _env("PRIMUS_INDEXER_TRITON", "1"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"):
+            assert is_triton_path_enabled()
+            idx_t, sc_t = indexer(hidden)
+        # Both knobs off → fully eager.
+        with _env("PRIMUS_INDEXER_TRITON", "0"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"):
+            assert not is_triton_path_enabled()
+            idx_e, sc_e = indexer(hidden)
+
+        assert idx_t.shape == idx_e.shape
+        match_ratio = (idx_t == idx_e).float().mean().item()
+        # P41 tail-only path keeps the einsum eager so its dot output is
+        # bit-identical to the eager path; the only divergence is fp32 →
+        # bf16 cast rounding in the tail.  Higher match ratio than P38.
+        if dtype == torch.float32:
+            assert match_ratio >= 0.99, f"fp32 topk match ratio too low: {match_ratio}"
+        else:
+            assert match_ratio >= 0.95, f"bf16 topk match ratio too low: {match_ratio}"
+
+
+# ---------------------------------------------------------------------------
+# G43: BWD parity
+# ---------------------------------------------------------------------------
+
+
+class TestG43BackwardParity:
+    @pytest.mark.parametrize("H", [1, 2, 4, 8])
+    def test_fast_tier_bwd_eager_parity_fp32(self, H):
+        B, S, P = 1, 32, 8
+        dtype = torch.float32
+        dot_e, w_e = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=4400 + H)
+        dot_e.requires_grad_(True)
+        w_e.requires_grad_(True)
+        dot_t = dot_e.detach().clone().requires_grad_(True)
+        w_t = w_e.detach().clone().requires_grad_(True)
+
+        s_t = IndexerScorePostFn.apply(dot_t, w_t, 4, dtype)
+        s_e = _eager_tail(dot_e, w_e, compress_ratio=4, out_dtype=dtype)
+
+        g = torch.randn_like(s_t)
+        finite = torch.isfinite(s_t)
+        g = torch.where(finite, g, torch.zeros_like(g))
+
+        (s_t * g).sum().backward()
+        (s_e * g).sum().backward()
+
+        torch.testing.assert_close(dot_t.grad, dot_e.grad, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(w_t.grad, w_e.grad, atol=1e-4, rtol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# G43: release-tier V4-Flash production shape
+# ---------------------------------------------------------------------------
+
+
+class TestG43ReleaseTier:
+    """V4-Flash production: ``B=1, S=4096, P=1024, H=8``, bf16."""
+
+    @pytest.mark.slow
+    def test_v4_flash_fwd_parity(self):
+        B, S, P, H = 1, 4096, 1024, 8
+        dtype = torch.bfloat16
+        dot, w = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=9000)
+
+        s_t = IndexerScorePostFn.apply(dot, w, 4, dtype)
+        s_e = _eager_tail(dot, w, compress_ratio=4, out_dtype=dtype)
+
+        finite_e = torch.isfinite(s_e)
+        finite_t = torch.isfinite(s_t)
+        torch.testing.assert_close(finite_e, finite_t)
+        # Bandwidth-bound tail; bf16 sum of H terms keeps tighter
+        # tolerance than P38 (which accumulated H tensor-core dots).
+        torch.testing.assert_close(
+            s_t[finite_t].float(),
+            s_e[finite_e].float(),
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+
+# ---------------------------------------------------------------------------
+# G43: edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestG43EdgeCases:
+    def test_unsupported_h_raises(self):
+        dot = torch.randn((1, 8, 7, 16), dtype=torch.float32, device="cuda")
+        w = torch.randn((1, 8, 7), dtype=torch.float32, device="cuda")
+        with pytest.raises(ValueError, match="Unsupported H"):
+            IndexerScorePostFn.apply(dot, w, 4, torch.float32)
+
+    def test_supported_predicate(self):
+        good_dot = torch.randn((1, 64, 8, 16), dtype=torch.float32, device="cuda")
+        good_w = torch.randn((1, 64, 8), dtype=torch.float32, device="cuda")
+        assert is_triton_kernel_supported(good_dot, good_w)
+
+        bad_h_dot = torch.randn((1, 64, 7, 16), dtype=torch.float32, device="cuda")
+        assert not is_triton_kernel_supported(bad_h_dot, good_w)
+
+        cpu_dot = torch.randn((1, 64, 8, 16), dtype=torch.float32)
+        assert not is_triton_kernel_supported(cpu_dot, good_w)
+
+    def test_env_default_on(self):
+        """Plan-8 P57 close-out 2 (2026-05-15): flipped default to ON.
+
+        Microbench at V4-Flash widths is consistently positive
+        (FWD 4.30x / BWD 1.63x) and the EP=8 proxy A/B shows a small
+        but positive ~0.2 ms / iter gain.  The conservative descope
+        rationale from P41 / P43 was that the per-iter gain sat below
+        the proxy noise floor; for the production code path we default
+        the microbench-positive kernel ON.
+        """
+        with _env("PRIMUS_INDEXER_TRITON", None):
+            assert is_triton_path_enabled()
+
+    def test_env_explicit_zero_disables(self):
+        """Setting ``PRIMUS_INDEXER_TRITON=0`` reverts to the eager body."""
+        with _env("PRIMUS_INDEXER_TRITON", "0"):
+            assert not is_triton_path_enabled()
+
+    def test_env_distinct_from_full(self):
+        """`PRIMUS_INDEXER_TRITON` controls the tail path; the legacy
+        P38 full-fuse path is gated by `PRIMUS_INDEXER_TRITON_FULL`."""
+        from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import (
+            is_triton_path_enabled as full_enabled,
+        )
+
+        with _env("PRIMUS_INDEXER_TRITON", "1"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"):
+            assert is_triton_path_enabled()
+            assert not full_enabled()
+        with _env("PRIMUS_INDEXER_TRITON", "0"), _env("PRIMUS_INDEXER_TRITON_FULL", "1"):
+            assert not is_triton_path_enabled()
+            assert full_enabled()
+
+
+# ---------------------------------------------------------------------------
+# G43: indexer_score_post_triton entry-point smoke
+# ---------------------------------------------------------------------------
+
+
+class TestG43EntryPoint:
+    def test_helper_returns_same_as_class(self):
+        torch.manual_seed(20260515)
+        B, S, P, H = 1, 32, 8, 4
+        dot = torch.randn((B, S, H, P), dtype=torch.float32, device="cuda")
+        w = torch.randn((B, S, H), dtype=torch.float32, device="cuda").abs()
+        out_helper = indexer_score_post_triton(dot, w, compress_ratio=4, out_dtype=torch.float32)
+        out_class = IndexerScorePostFn.apply(dot, w, 4, torch.float32)
+        torch.testing.assert_close(out_helper, out_class)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py
new file mode 100644
index 000000000..4a8461af5
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py
@@ -0,0 +1,434 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-6 P35 G38 — `apply_interleaved_partial_rope` Triton FWD/BWD parity.
+
+Asserts that :class:`RoPEInterleavedPartialFn` (Triton kernel from
+``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial``)
+produces the same output as the eager
+:func:`apply_interleaved_partial_rope` body in
+``primus.backends.megatron.core.transformer.dual_rope``, FWD **and**
+BWD, at two tiers:
+
+* fast tier — small shapes (``B=2, S=8, H=4, head_dim=16, rd=8``)
+  exercising every code path (nope copy, interleaved pair rotation,
+  cos/sin broadcast across heads) in milliseconds; parametrised over
+  ``{fp32, fp16, bf16}`` and ``rotary_dim ∈ {0, 4, 8, 16}``;
+* release tier — Q shape (``B=1, S=4096, H=64, head_dim=512, rd=64``)
+  and K shape (``B=1, S=4096, H=1, head_dim=64, rd=64``) mirroring the
+  V4-Flash EP=8 proxy widths, behind ``pytest.mark.slow``.
+
+`gradcheck` is run at the fast tier in fp32 to catch any analytic-VJP
+bug in the BWD kernel; the FWD has a clean closed form so its forward
+parity assertion is the load-bearing test.
+
+GPU-only; CPU runs are ``pytest.skip``-ed at module collection time.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip(
+        "rope_interleaved_partial Triton kernel requires CUDA / HIP",
+        allow_module_level=True,
+    )
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.dual_rope import (  # noqa: E402
+    apply_interleaved_partial_rope,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import (  # noqa: E402
+    RoPEInterleavedPartialFn,
+    apply_rope_interleaved_partial,
+    eager_apply_interleaved_partial_rope,
+    is_triton_path_enabled,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def _env(key: str, value: str | None):
+    """Temporarily set / unset ``os.environ[key]``."""
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _build_cos_sin(
+    *,
+    leading_shape: tuple[int, ...],
+    rotary_dim: int,
+    dtype: torch.dtype,
+    seed: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build cos / sin with shape ``leading_shape + (rotary_dim // 2,)``.
+
+    Mimics the way :class:`RoPECache` produces cos/sin: build
+    ``position_ids.float() * inv_freq`` then take ``cos`` / ``sin``.
+    For the test we just generate random freqs to exercise more values.
+    """
+    rd_half = rotary_dim // 2
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    freqs = torch.randn((*leading_shape, rd_half), dtype=torch.float32, device="cuda", generator=gen)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def _build_x(
+    *,
+    leading_shape: tuple[int, ...],
+    H: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    requires_grad: bool = False,
+    seed: int = 0,
+) -> torch.Tensor:
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(
+        (*leading_shape, H, head_dim),
+        dtype=dtype,
+        device="cuda",
+        generator=gen,
+    )
+    if requires_grad:
+        x.requires_grad_(True)
+    return x
+
+
+def _dtype_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.float32:
+        return 1e-6, 1e-6
+    if dtype == torch.float16:
+        return 1e-3, 1e-3
+    if dtype == torch.bfloat16:
+        return 1e-2, 1e-2
+    raise ValueError(dtype)
+
+
+# ---------------------------------------------------------------------------
+# G38: FWD parity vs eager
+# ---------------------------------------------------------------------------
+
+
+_FAST_LEADING = (2, 8)  # B=2, S=8
+_FAST_H = 4
+_FAST_HEAD_DIM = 16
+
+
+class TestG38ForwardParity:
+    """FWD output matches eager :func:`apply_interleaved_partial_rope`
+    body within dtype tolerance.
+
+    The Triton kernel does its arithmetic in the caller's dtype (bf16 /
+    fp16 / fp32), matching the plan-5 P32 RoPE bf16 cast contract.  So
+    the tolerance is the same as if we cast cos/sin to ``x.dtype`` and
+    ran the eager body — which is exactly what
+    :func:`eager_apply_interleaved_partial_rope` does.
+    """
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    @pytest.mark.parametrize("rotary_dim", [0, 4, 8, 16])
+    def test_fast_tier_fwd_parity(self, dtype, rotary_dim):
+        if rotary_dim > _FAST_HEAD_DIM:
+            pytest.skip("rotary_dim must be <= head_dim")
+        x = _build_x(
+            leading_shape=_FAST_LEADING,
+            H=_FAST_H,
+            head_dim=_FAST_HEAD_DIM,
+            dtype=dtype,
+            seed=42,
+        )
+        cos, sin = _build_cos_sin(
+            leading_shape=_FAST_LEADING,
+            rotary_dim=max(rotary_dim, 2),
+            dtype=dtype,
+            seed=43,
+        )
+        # For rd=0 the kernel ignores cos/sin; pass any shape.
+        cos_use = cos[..., : max(rotary_dim // 2, 1)]
+        sin_use = sin[..., : max(rotary_dim // 2, 1)]
+
+        out_triton = RoPEInterleavedPartialFn.apply(x, cos_use, sin_use, rotary_dim)
+        out_eager = eager_apply_interleaved_partial_rope(x, cos_use, sin_use, rotary_dim=rotary_dim)
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_uneven_block_h(self, dtype):
+        """H=5 is not a multiple of BLOCK_H=8 -> masked load / store paths."""
+        leading = (1, 7)  # also a prime sequence length
+        H = 5
+        head_dim = 12
+        rotary_dim = 4
+        x = _build_x(leading_shape=leading, H=H, head_dim=head_dim, dtype=dtype, seed=1234)
+        cos, sin = _build_cos_sin(leading_shape=leading, rotary_dim=rotary_dim, dtype=dtype, seed=1235)
+        out_triton = RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim)
+        out_eager = eager_apply_interleaved_partial_rope(x, cos, sin, rotary_dim=rotary_dim)
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G38: BWD parity vs eager + gradcheck
+# ---------------------------------------------------------------------------
+
+
+class TestG38BackwardParity:
+    """BWD parity vs eager autograd graph, plus a small fp64 gradcheck."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("rotary_dim", [4, 8, 16])
+    def test_fast_tier_bwd_parity(self, dtype, rotary_dim):
+        if rotary_dim > _FAST_HEAD_DIM:
+            pytest.skip("rotary_dim must be <= head_dim")
+
+        x_t = _build_x(
+            leading_shape=_FAST_LEADING,
+            H=_FAST_H,
+            head_dim=_FAST_HEAD_DIM,
+            dtype=dtype,
+            requires_grad=True,
+            seed=44,
+        )
+        x_e = x_t.detach().clone().requires_grad_(True)
+        cos, sin = _build_cos_sin(
+            leading_shape=_FAST_LEADING,
+            rotary_dim=rotary_dim,
+            dtype=dtype,
+            seed=45,
+        )
+
+        out_triton = RoPEInterleavedPartialFn.apply(x_t, cos, sin, rotary_dim)
+        out_eager = eager_apply_interleaved_partial_rope(x_e, cos, sin, rotary_dim=rotary_dim)
+
+        gen = torch.Generator(device="cuda").manual_seed(46)
+        grad = torch.randn_like(out_triton, dtype=dtype)
+        grad_e = grad.detach().clone()
+        gen.manual_seed(46)
+        out_triton.backward(grad)
+        out_eager.backward(grad_e)
+
+        atol, rtol = _dtype_tolerance(dtype)
+        assert x_t.grad is not None
+        assert x_e.grad is not None
+        torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol)
+
+    def test_gradcheck_fast_tier(self):
+        """`torch.autograd.gradcheck` at fp64 + tiny shape catches any
+        analytic-VJP bug in :func:`_apply_rope_bwd_kernel`.
+        """
+        leading = (1, 4)
+        H = 2
+        head_dim = 8
+        rotary_dim = 4
+        x = _build_x(
+            leading_shape=leading,
+            H=H,
+            head_dim=head_dim,
+            dtype=torch.float64,
+            requires_grad=True,
+            seed=100,
+        )
+        cos, sin = _build_cos_sin(
+            leading_shape=leading,
+            rotary_dim=rotary_dim,
+            dtype=torch.float64,
+            seed=101,
+        )
+
+        def fn(xx):
+            return RoPEInterleavedPartialFn.apply(xx, cos, sin, rotary_dim)
+
+        torch.autograd.gradcheck(fn, (x,), eps=1e-6, atol=1e-7, rtol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# G38: release-tier V4-Flash widths (Q + K)
+# ---------------------------------------------------------------------------
+
+
+class TestG38ReleaseTier:
+    """V4-Flash EP=8 widths: Q (H=64, head_dim=512, rd=64) and K (H=1).
+
+    Marked ``slow``; bf16 only (the production dtype).  FWD + BWD parity
+    within the elementwise attention tolerance from plan-5 P32.
+    """
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize(
+        "spec",
+        [
+            ("Q", (1, 4096), 64, 512, 64),
+            ("K", (1, 4096), 1, 64, 64),
+        ],
+        ids=["Q-shape", "K-shape"],
+    )
+    def test_v4_flash_proxy_parity(self, spec):
+        name, leading, H, head_dim, rotary_dim = spec
+        dtype = torch.bfloat16
+
+        x_t = _build_x(
+            leading_shape=leading,
+            H=H,
+            head_dim=head_dim,
+            dtype=dtype,
+            requires_grad=True,
+            seed=200,
+        )
+        x_e = x_t.detach().clone().requires_grad_(True)
+        cos, sin = _build_cos_sin(leading_shape=leading, rotary_dim=rotary_dim, dtype=dtype, seed=201)
+
+        out_triton = RoPEInterleavedPartialFn.apply(x_t, cos, sin, rotary_dim)
+        out_eager = eager_apply_interleaved_partial_rope(x_e, cos, sin, rotary_dim=rotary_dim)
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+
+        grad = torch.randn_like(out_triton)
+        out_triton.backward(grad)
+        out_eager.backward(grad.detach().clone())
+        torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G38: rd == 0 early-return + error paths
+# ---------------------------------------------------------------------------
+
+
+class TestG38EdgeCases:
+    """Cover :class:`RoPEInterleavedPartialFn`'s defensive validation."""
+
+    def test_rd_zero_returns_input_contiguous(self):
+        x = _build_x(
+            leading_shape=(2, 4),
+            H=2,
+            head_dim=8,
+            dtype=torch.float32,
+            seed=1,
+        )
+        cos, sin = _build_cos_sin(leading_shape=(2, 4), rotary_dim=2, dtype=torch.float32, seed=2)
+        out = RoPEInterleavedPartialFn.apply(x, cos, sin, 0)
+        assert torch.equal(out, x.contiguous())
+
+    def test_odd_rotary_dim_raises(self):
+        x = _build_x(
+            leading_shape=(1, 4),
+            H=2,
+            head_dim=8,
+            dtype=torch.float32,
+            seed=1,
+        )
+        cos, sin = _build_cos_sin(leading_shape=(1, 4), rotary_dim=4, dtype=torch.float32, seed=2)
+        with pytest.raises(ValueError, match="rotary_dim must be even"):
+            RoPEInterleavedPartialFn.apply(x, cos, sin, 3)
+
+    def test_rotary_dim_exceeds_head_dim_raises(self):
+        x = _build_x(
+            leading_shape=(1, 4),
+            H=2,
+            head_dim=8,
+            dtype=torch.float32,
+            seed=1,
+        )
+        cos, sin = _build_cos_sin(leading_shape=(1, 4), rotary_dim=4, dtype=torch.float32, seed=2)
+        with pytest.raises(ValueError, match="must be <= head_dim"):
+            RoPEInterleavedPartialFn.apply(x, cos, sin, 16)
+
+    def test_cos_shape_mismatch_raises(self):
+        x = _build_x(
+            leading_shape=(2, 4),
+            H=2,
+            head_dim=8,
+            dtype=torch.float32,
+            seed=1,
+        )
+        cos, sin = _build_cos_sin(leading_shape=(2, 4), rotary_dim=2, dtype=torch.float32, seed=2)
+        with pytest.raises(ValueError, match="last dim must be rotary_dim"):
+            RoPEInterleavedPartialFn.apply(x, cos, sin, 4)
+
+
+# ---------------------------------------------------------------------------
+# G38: env-flag dispatch through dual_rope.apply_interleaved_partial_rope
+# ---------------------------------------------------------------------------
+
+
+class TestG38EnvFlagDispatch:
+    """`PRIMUS_ROPE_TRITON` env knob flips the dispatcher behaviour
+    inside :func:`apply_interleaved_partial_rope`.
+
+    The Triton path is bit-equivalent to the eager body within bf16
+    tolerance; this test pins both code paths to make sure neither is
+    silently broken.
+    """
+
+    def test_env_on_uses_triton(self):
+        x = _build_x(
+            leading_shape=_FAST_LEADING,
+            H=_FAST_H,
+            head_dim=_FAST_HEAD_DIM,
+            dtype=torch.bfloat16,
+            seed=11,
+        )
+        cos, sin = _build_cos_sin(
+            leading_shape=_FAST_LEADING,
+            rotary_dim=8,
+            dtype=torch.bfloat16,
+            seed=12,
+        )
+        with _env("PRIMUS_ROPE_TRITON", "1"):
+            assert is_triton_path_enabled()
+            out_on = apply_interleaved_partial_rope(x, cos, sin, rotary_dim=8)
+        with _env("PRIMUS_ROPE_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            out_off = apply_interleaved_partial_rope(x, cos, sin, rotary_dim=8)
+        atol, rtol = _dtype_tolerance(torch.bfloat16)
+        torch.testing.assert_close(out_on, out_off, atol=atol, rtol=rtol)
+
+    def test_apply_rope_interleaved_partial_dispatcher(self):
+        """The dispatcher in the kernel module mirrors the dual_rope
+        wiring -- pin both paths return the same answer.
+        """
+        x = _build_x(
+            leading_shape=(2, 16),
+            H=2,
+            head_dim=16,
+            dtype=torch.bfloat16,
+            seed=21,
+        )
+        cos, sin = _build_cos_sin(
+            leading_shape=(2, 16),
+            rotary_dim=8,
+            dtype=torch.bfloat16,
+            seed=22,
+        )
+        with _env("PRIMUS_ROPE_TRITON", "1"):
+            out_triton = apply_rope_interleaved_partial(x, cos, sin, rotary_dim=8)
+        with _env("PRIMUS_ROPE_TRITON", "0"):
+            out_eager = apply_rope_interleaved_partial(x, cos, sin, rotary_dim=8)
+        atol, rtol = _dtype_tolerance(torch.bfloat16)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py
new file mode 100644
index 000000000..d6f93a2f8
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py
@@ -0,0 +1,267 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-6 P39 G42 — V4 router post-logits Triton parity.
+
+Asserts that :class:`V4RouterPostFn` (Triton kernel from
+``primus.backends.megatron.core.transformer.moe._triton.v4_router_post``)
+produces ``(probs, routing_map)`` bit-equal to the eager body of
+:func:`primus.backends.megatron.core.transformer.moe.v4_topk_router._compute_route`
+across the 3 score functions × {with, without bias} × {hash router,
+learned router}.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip(
+        "v4_router_post Triton kernel requires CUDA / HIP",
+        allow_module_level=True,
+    )
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (  # noqa: E402
+    V4RouterPostFn,
+    is_triton_path_enabled,
+)
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (  # noqa: E402
+    DeepseekV4HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (  # noqa: E402
+    DeepseekV4LearnedRouter,
+)
+
+
+@contextmanager
+def _env(key: str, value: str | None):
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _eager_v4_score(logits: torch.Tensor, *, score_function: str):
+    if score_function == "softmax":
+        return torch.softmax(logits, dim=-1)
+    if score_function == "sigmoid":
+        return torch.sigmoid(logits)
+    if score_function == "sqrtsoftplus":
+        return torch.sqrt(torch.nn.functional.softplus(logits))
+    raise ValueError(score_function)
+
+
+def _eager_post(
+    logits: torch.Tensor,
+    indices: torch.Tensor,
+    *,
+    score_function: str,
+    topk_scaling_factor: float,
+    out_dtype: torch.dtype,
+):
+    scores = _eager_v4_score(logits.float(), score_function=score_function)
+    weights = scores.gather(1, indices)
+    if score_function != "softmax":
+        denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12)
+        weights = weights / denom
+    if topk_scaling_factor != 1.0:
+        weights = weights * float(topk_scaling_factor)
+    N, E = logits.shape
+    probs = torch.zeros(N, E, dtype=out_dtype, device=logits.device)
+    probs.scatter_(1, indices, weights.to(out_dtype))
+    rmap = torch.zeros(N, E, dtype=torch.bool, device=logits.device)
+    rmap.scatter_(1, indices, True)
+    return probs, rmap
+
+
+def _build_inputs(*, N: int, E: int, K: int, seed: int):
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    logits = torch.randn((N, E), dtype=torch.float32, device="cuda", generator=gen)
+    indices = torch.stack(
+        [torch.randperm(E, generator=gen, device="cuda")[:K] for _ in range(N)],
+        dim=0,
+    ).to(torch.int64)
+    return logits, indices
+
+
+# ---------------------------------------------------------------------------
+# G42: FWD parity vs eager (3 score fns × {fp32, bf16} × small / release)
+# ---------------------------------------------------------------------------
+
+
+class TestG42ForwardParity:
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    @pytest.mark.parametrize("scale", [1.0, 2.5])
+    @pytest.mark.parametrize("out_dtype", [torch.float32])
+    def test_fast_tier_fwd_eager_parity(self, score_function, scale, out_dtype):
+        N, E, K = 64, 32, 4
+        logits, indices = _build_inputs(N=N, E=E, K=K, seed=2000)
+        probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, scale, out_dtype)
+        probs_e, rmap_e = _eager_post(
+            logits, indices, score_function=score_function, topk_scaling_factor=scale, out_dtype=out_dtype
+        )
+        torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(rmap_t, rmap_e)
+        torch.testing.assert_close(probs_t.nonzero(), probs_e.nonzero(), check_dtype=False)
+
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    @pytest.mark.parametrize("shape", [(64, 256, 6), (64, 7, 3), (64, 100, 5)])
+    def test_non_pow2_e_k_fwd_parity(self, score_function, shape):
+        """Arbitrary (non-power-of-2) E / K — the production V4-Flash router
+        runs E=256, K=6, so K=6 is the load-bearing case.  Also cover a
+        non-power-of-2 E (7, 100)."""
+        N, E, K = shape
+        logits, indices = _build_inputs(N=N, E=E, K=K, seed=2500)
+        probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, 2.5, torch.float32)
+        probs_e, rmap_e = _eager_post(
+            logits, indices, score_function=score_function, topk_scaling_factor=2.5, out_dtype=torch.float32
+        )
+        torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(rmap_t, rmap_e)
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    def test_release_tier_fwd_eager_parity(self, score_function):
+        # V4-Flash production widths: E=256, K=6 (non-power-of-2 topk).
+        N, E, K = 4096, 256, 6
+        logits, indices = _build_inputs(N=N, E=E, K=K, seed=8000)
+        probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, 2.5, torch.float32)
+        probs_e, rmap_e = _eager_post(
+            logits, indices, score_function=score_function, topk_scaling_factor=2.5, out_dtype=torch.float32
+        )
+        torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(rmap_t, rmap_e)
+
+
+# ---------------------------------------------------------------------------
+# G42: BWD parity
+# ---------------------------------------------------------------------------
+
+
+class TestG42BackwardParity:
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    @pytest.mark.parametrize("scale", [1.0, 2.5])
+    @pytest.mark.parametrize("shape", [(64, 32, 4), (64, 256, 6), (64, 100, 5)])
+    def test_fast_tier_bwd_eager_parity(self, score_function, scale, shape):
+        N, E, K = shape
+        logits_e, indices = _build_inputs(N=N, E=E, K=K, seed=3100)
+        logits_e = logits_e.requires_grad_(True)
+        logits_t = logits_e.detach().clone().requires_grad_(True)
+
+        probs_t, _ = V4RouterPostFn.apply(logits_t, indices, score_function, scale, torch.float32)
+        probs_e, _ = _eager_post(
+            logits_e,
+            indices,
+            score_function=score_function,
+            topk_scaling_factor=scale,
+            out_dtype=torch.float32,
+        )
+
+        g = torch.randn_like(probs_t)
+        (probs_t * g).sum().backward()
+        (probs_e * g).sum().backward()
+
+        torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=1e-4, rtol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# G42: composed routers
+# ---------------------------------------------------------------------------
+
+
+class TestG42ComposedRouters:
+    """Both routers (`DeepseekV4LearnedRouter` and `DeepseekV4HashRouter`)
+    must produce bit-equal (probs, routing_map) between the
+    PRIMUS_V4_ROUTER_TRITON=1 and =0 paths.  This is the load-bearing
+    integration test for downstream MoEDispatch.
+    """
+
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    def test_learned_router_env_toggle_parity(self, score_function):
+        torch.manual_seed(123)
+        N, D, E, K = 64, 32, 32, 4
+        router = DeepseekV4LearnedRouter(
+            hidden_size=D, num_experts=E, topk=K, score_function=score_function
+        ).to("cuda")
+        hidden = torch.randn((1, N, D), dtype=torch.float32, device="cuda")
+
+        with _env("PRIMUS_V4_ROUTER_TRITON", "1"):
+            assert is_triton_path_enabled()
+            probs_t, rmap_t = router(hidden)
+        with _env("PRIMUS_V4_ROUTER_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            probs_e, rmap_e = router(hidden)
+
+        torch.testing.assert_close(probs_t, probs_e, atol=1e-4, rtol=1e-4)
+        torch.testing.assert_close(rmap_t, rmap_e)
+
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+    def test_hash_router_env_toggle_parity(self, score_function):
+        torch.manual_seed(124)
+        N, D, E, K = 64, 32, 32, 4
+        V = 1000
+        router = DeepseekV4HashRouter(
+            hidden_size=D,
+            num_experts=E,
+            topk=K,
+            vocab_size=V,
+            score_function=score_function,
+        ).to("cuda")
+        hidden = torch.randn((1, N, D), dtype=torch.float32, device="cuda")
+        token_ids = torch.randint(0, V, (1, N), dtype=torch.long, device="cuda")
+
+        with _env("PRIMUS_V4_ROUTER_TRITON", "1"):
+            probs_t, rmap_t = router(hidden, token_ids)
+        with _env("PRIMUS_V4_ROUTER_TRITON", "0"):
+            probs_e, rmap_e = router(hidden, token_ids)
+
+        torch.testing.assert_close(probs_t, probs_e, atol=1e-4, rtol=1e-4)
+        torch.testing.assert_close(rmap_t, rmap_e)
+
+
+# ---------------------------------------------------------------------------
+# G42: edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestG42EdgeCases:
+    def test_unknown_score_function_raises(self):
+        logits = torch.randn((4, 8), dtype=torch.float32, device="cuda")
+        indices = torch.zeros((4, 2), dtype=torch.int64, device="cuda")
+        with pytest.raises(ValueError, match="score_function"):
+            V4RouterPostFn.apply(logits, indices, "tanh", 1.0, torch.float32)
+
+    def test_cpu_tensor_asserts(self):
+        # v4_router_post_triton asserts CUDA tensors (no support predicate).
+        from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import (
+            v4_router_post_triton,
+        )
+
+        cpu_logits = torch.randn((4, 8), dtype=torch.float32)
+        cpu_idx = torch.zeros((4, 2), dtype=torch.int64)
+        with pytest.raises(AssertionError, match="CUDA"):
+            v4_router_post_triton(
+                cpu_logits,
+                cpu_idx,
+                score_function="softmax",
+                topk_scaling_factor=1.0,
+                out_dtype=torch.float32,
+            )
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py
new file mode 100644
index 000000000..934daf1ab
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py
@@ -0,0 +1,358 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-6 P36 G39 — `sinkhorn_normalize` Triton FWD/BWD parity.
+
+Asserts that :class:`SinkhornNormalizeFn` (Triton kernel from
+``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn``)
+matches the eager :func:`sinkhorn_normalize` body in
+``primus.backends.megatron.core.transformer.hyper_connection`` and the
+plan-5 P29 compiled path within dtype tolerance, FWD **and** BWD, at
+two tiers:
+
+* fast tier — `B=2, S=64, K=4`, exercising every code path in the
+  kernel (priming col-step + 19 row/col pairs + cached state buffer
+  round-trip) in milliseconds; parametrised over ``bf16`` (the
+  production compute dtype) and ``n_iters ∈ {5, 20}``;
+* release tier — `B=1, S=4096, K=4` bf16 (V4-Flash production shape),
+  behind ``pytest.mark.slow``.
+
+The **doubly-stochastic property check** is a model-quality contract
+independent of the eager path (row / col sums of the FWD output equal
+``1`` within ``eps * K``).
+
+Note on ``torch.autograd.gradcheck``: the eager body casts to fp32
+internally (``m = logits.float()``) and the Triton kernel matches that
+contract bit-for-bit.  ``gradcheck`` at fp64 input is therefore
+incompatible with both paths — the fp32 cast is lossy at the fp64
+finite-difference step size.  We instead pin the BWD via a direct
+parity check against ``torch.autograd`` of the eager body in fp64
+(where ``out.backward(grad)`` does exactly the right thing through the
+``.float()`` cast).
+
+GPU-only; CPU runs are ``pytest.skip``-ed at module collection time.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip(
+        "sinkhorn Triton kernel requires CUDA / HIP",
+        allow_module_level=True,
+    )
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.hyper_connection import (  # noqa: E402
+    _get_compiled_sinkhorn,
+    sinkhorn_normalize,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import (  # noqa: E402
+    SinkhornNormalizeFn,
+    eager_sinkhorn_normalize,
+    is_triton_kernel_supported,
+    is_triton_path_enabled,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def _env(key: str, value: str | None):
+    """Temporarily set / unset ``os.environ[key]``."""
+    prev = os.environ.get(key)
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        if prev is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = prev
+
+
+def _build_logits(
+    *, B: int, S: int, K: int, dtype: torch.dtype, seed: int, requires_grad: bool = False
+) -> torch.Tensor:
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    # Sinkhorn input is non-negative (softmax + eps in production).
+    x = torch.rand((B, S, K, K), dtype=dtype, device="cuda", generator=gen) + 1e-3
+    if requires_grad:
+        x.requires_grad_(True)
+    return x
+
+
+def _dtype_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.float32:
+        return 1e-5, 1e-5
+    if dtype == torch.float16:
+        return 1e-3, 1e-3
+    if dtype == torch.bfloat16:
+        return 1e-2, 1e-2
+    if dtype == torch.float64:
+        return 1e-7, 1e-7
+    raise ValueError(dtype)
+
+
+# ---------------------------------------------------------------------------
+# G39: FWD parity vs eager (and vs plan-5 P29 compiled)
+# ---------------------------------------------------------------------------
+
+
+class TestG39ForwardParity:
+    """FWD output matches eager and (separately) the compiled path."""
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16])
+    @pytest.mark.parametrize("n_iters", [5, 20])
+    def test_fast_tier_fwd_eager_parity(self, dtype, n_iters):
+        x = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=42)
+
+        out_triton = SinkhornNormalizeFn.apply(x.clone(), n_iters, 1e-6)
+        out_eager = eager_sinkhorn_normalize(x.clone(), n_iters=n_iters, eps=1e-6)
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16])
+    @pytest.mark.parametrize("n_iters", [5, 20])
+    def test_fast_tier_fwd_compiled_parity(self, dtype, n_iters):
+        """Triton FWD == plan-5 P29 compiled FWD within dtype tolerance.
+
+        Both paths share the same algorithm; this test pins them
+        together so a future P29 / P36 divergence raises a CI alarm.
+        """
+        x = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=43)
+
+        out_triton = SinkhornNormalizeFn.apply(x.clone(), n_iters, 1e-6)
+        compiled_fn = _get_compiled_sinkhorn(n_iters, 1e-6, dtype)
+        out_compiled = compiled_fn(x.clone())
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_compiled, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G39: BWD parity vs eager + gradcheck
+# ---------------------------------------------------------------------------
+
+
+class TestG39BackwardParity:
+    """BWD parity vs eager ``torch.autograd`` across dtypes."""
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16])
+    @pytest.mark.parametrize("n_iters", [5, 20])
+    def test_fast_tier_bwd_eager_parity(self, dtype, n_iters):
+        x_base = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=44)
+        x_t = x_base.detach().clone().requires_grad_(True)
+        x_e = x_base.detach().clone().requires_grad_(True)
+
+        out_triton = SinkhornNormalizeFn.apply(x_t, n_iters, 1e-6)
+        out_eager = eager_sinkhorn_normalize(x_e, n_iters=n_iters, eps=1e-6)
+
+        grad = torch.randn_like(out_triton)
+        out_triton.backward(grad)
+        out_eager.backward(grad.detach().clone())
+
+        atol, rtol = _dtype_tolerance(dtype)
+        assert x_t.grad is not None
+        assert x_e.grad is not None
+        torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol)
+
+    def test_fp64_input_bwd_eager_parity(self):
+        """fp64 input -> the kernel still does fp32-internal compute
+        (matches eager's ``m = logits.float()`` contract); we compare
+        Triton BWD vs ``torch.autograd``-of-eager rather than gradcheck
+        (gradcheck's fp64 finite-difference step is incompatible with
+        the lossy fp64->fp32 cast).
+        """
+        x_base = _build_logits(B=1, S=8, K=4, dtype=torch.float64, seed=100)
+        x_t = x_base.detach().clone().requires_grad_(True)
+        x_e = x_base.detach().clone().requires_grad_(True)
+
+        out_triton = SinkhornNormalizeFn.apply(x_t, 5, 1e-6)
+        out_eager = eager_sinkhorn_normalize(x_e, n_iters=5, eps=1e-6)
+
+        grad = torch.randn_like(out_triton)
+        out_triton.backward(grad)
+        out_eager.backward(grad.detach().clone())
+
+        # fp64 grads come back through a fp32-internal compute; tolerance
+        # is set to fp32 precision (the kernel can't be more accurate
+        # than its compute dtype).
+        torch.testing.assert_close(x_t.grad, x_e.grad, atol=1e-5, rtol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# G39: doubly-stochastic property (model-quality contract)
+# ---------------------------------------------------------------------------
+
+
+class TestG39DoublyStochastic:
+    """The FWD output's row and column sums equal ``1`` within
+    ``eps * K`` -- pinning the model-quality contract that
+    :func:`sinkhorn_normalize` is supposed to guarantee.
+
+    This is independent of the eager path; an algorithmic bug in the
+    Triton kernel that still happens to match the eager value would
+    fail this check.
+    """
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16])
+    @pytest.mark.parametrize("K", [4, 8])
+    def test_row_col_sums_close_to_1(self, dtype, K):
+        x = _build_logits(B=2, S=16, K=K, dtype=dtype, seed=55)
+        y = SinkhornNormalizeFn.apply(x, 20, 1e-6)
+
+        row_sums = y.sum(dim=-1)
+        col_sums = y.sum(dim=-2)
+        # eps*K = 4e-6 (fp32) or ~6e-5 (bf16 -> the cast back to bf16
+        # is the dominant error, so we widen by 16x for bf16).
+        tol = 4e-6 * K if dtype == torch.float32 else 1e-2
+        torch.testing.assert_close(row_sums, torch.ones_like(row_sums), atol=tol, rtol=tol)
+        torch.testing.assert_close(col_sums, torch.ones_like(col_sums), atol=tol, rtol=tol)
+
+
+# ---------------------------------------------------------------------------
+# G39: release-tier V4-Flash production shape
+# ---------------------------------------------------------------------------
+
+
+class TestG39ReleaseTier:
+    """V4-Flash production shape: `B=1, S=4096, K=4`, bf16, n_iters=20.
+
+    Marked ``slow``; pins both FWD and BWD against eager AND compiled
+    paths to make sure the in-production widths match.
+    """
+
+    @pytest.mark.slow
+    def test_v4_flash_fwd_bwd_parity(self):
+        dtype = torch.bfloat16
+        x_base = _build_logits(B=1, S=4096, K=4, dtype=dtype, seed=200)
+
+        x_t = x_base.detach().clone().requires_grad_(True)
+        x_e = x_base.detach().clone().requires_grad_(True)
+        x_c = x_base.detach().clone().requires_grad_(True)
+
+        out_triton = SinkhornNormalizeFn.apply(x_t, 20, 1e-6)
+        out_eager = eager_sinkhorn_normalize(x_e, n_iters=20, eps=1e-6)
+        compiled_fn = _get_compiled_sinkhorn(20, 1e-6, dtype)
+        out_compiled = compiled_fn(x_c)
+
+        atol, rtol = _dtype_tolerance(dtype)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+        torch.testing.assert_close(out_triton, out_compiled, atol=atol, rtol=rtol)
+
+        grad = torch.randn_like(out_triton)
+        out_triton.backward(grad.detach().clone())
+        out_eager.backward(grad.detach().clone())
+        out_compiled.backward(grad.detach().clone())
+        torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol)
+        torch.testing.assert_close(x_t.grad, x_c.grad, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# G39: edge cases + error paths
+# ---------------------------------------------------------------------------
+
+
+class TestG39EdgeCases:
+    """Defensive validation in :class:`SinkhornNormalizeFn`."""
+
+    def test_non_square_raises(self):
+        x = torch.rand(2, 4, 5, device="cuda", dtype=torch.float32) + 1e-3
+        with pytest.raises(ValueError, match="square"):
+            SinkhornNormalizeFn.apply(x, 20, 1e-6)
+
+    def test_unsupported_k_raises(self):
+        x = torch.rand(2, 32, 32, device="cuda", dtype=torch.float32) + 1e-3
+        with pytest.raises(ValueError, match="unsupported K"):
+            SinkhornNormalizeFn.apply(x, 20, 1e-6)
+
+    def test_n_iters_zero_raises(self):
+        x = torch.rand(2, 4, 4, device="cuda", dtype=torch.float32) + 1e-3
+        with pytest.raises(ValueError, match="n_iters must be"):
+            SinkhornNormalizeFn.apply(x, 0, 1e-6)
+
+    def test_kernel_supported_predicate(self):
+        good = torch.rand(2, 4, 4, device="cuda", dtype=torch.float32)
+        bad_k = torch.rand(2, 32, 32, device="cuda", dtype=torch.float32)
+        bad_dev = torch.rand(2, 4, 4, dtype=torch.float32)  # cpu
+        assert is_triton_kernel_supported(good)
+        assert not is_triton_kernel_supported(bad_k)
+        assert not is_triton_kernel_supported(bad_dev)
+
+
+# ---------------------------------------------------------------------------
+# G39: env-flag dispatch through hyper_connection.sinkhorn_normalize
+# ---------------------------------------------------------------------------
+
+
+class TestG39EnvFlagDispatch:
+    """The ``PRIMUS_SINKHORN_TRITON`` env knob flips the dispatcher
+    inside :func:`sinkhorn_normalize`.  Both paths agree within bf16
+    tolerance.
+    """
+
+    def test_env_on_uses_triton(self):
+        x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=11)
+
+        with _env("PRIMUS_SINKHORN_TRITON", "1"):
+            assert is_triton_path_enabled()
+            out_on = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6)
+        with _env("PRIMUS_SINKHORN_TRITON", "0"):
+            assert not is_triton_path_enabled()
+            out_eager = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6)
+
+        atol, rtol = _dtype_tolerance(torch.bfloat16)
+        torch.testing.assert_close(out_on, out_eager, atol=atol, rtol=rtol)
+
+    def test_use_triton_kwarg_overrides_env_off(self):
+        """``use_triton=True`` forces the Triton path even when the env
+        knob is off (used by the unit tests / by callers who want a
+        per-call override).
+        """
+        x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=12)
+        with _env("PRIMUS_SINKHORN_TRITON", "0"):
+            out_triton = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6, use_triton=True)
+            out_eager = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6)
+
+        atol, rtol = _dtype_tolerance(torch.bfloat16)
+        torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol)
+
+    def test_routing_precedence_triton_over_compiled(self):
+        """When both ``use_triton`` and ``use_compiled`` are set, the
+        Triton path wins (routing precedence:
+        ``use_triton > use_compiled > eager``).
+        """
+        x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=13)
+        with _env("PRIMUS_SINKHORN_TRITON", "0"):
+            out = sinkhorn_normalize(
+                x.clone(),
+                n_iters=20,
+                eps=1e-6,
+                use_triton=True,
+                use_compiled=True,
+            )
+            out_triton_only = sinkhorn_normalize(
+                x.clone(),
+                n_iters=20,
+                eps=1e-6,
+                use_triton=True,
+            )
+
+        torch.testing.assert_close(out, out_triton_only, atol=0, rtol=0)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py
new file mode 100644
index 000000000..ce70fcdfb
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py
@@ -0,0 +1,122 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Backend import / default / hardware-gating guarantees for V4 attention.
+
+These are the regression gates for the "gluon must not be a hard dependency"
+contract:
+
+* the shared config default is ``triton_v1`` (arch-portable), NOT the
+  gfx950-only ``gluon``;
+* importing ``v4_attention_kernels`` (and ``deepseek_v4_attention``) must NOT
+  eagerly import the gluon backend (``_gluon_dsa`` /
+  ``triton.experimental.gluon``), so ``eager`` / ``triton_v1`` / ``triton_v2``
+  work on any Triton build / GPU arch;
+* the gluon backend is loaded lazily via ``load_gluon_attention_backends`` and,
+  when selected, ``_require_gfx950`` rejects non-gfx950 devices.
+
+All tests here are hardware-independent (no CUDA required).
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+
+import pytest
+
+
+def test_config_defaults_are_triton_v1():
+    """Shared config default must be the arch-portable triton_v1 (not gluon)."""
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    fields = DeepSeekV4TransformerConfig.__dataclass_fields__
+    assert fields["use_v4_attention_backend"].default == "triton_v1"
+    assert fields["use_v4_csa_attention_backend"].default == "triton_v1"
+
+
+def test_kernels_package_exposes_lazy_gluon_loader():
+    """The package exposes the lazy loader and does NOT eagerly bind gluon entries."""
+    pkg = importlib.import_module("primus.backends.megatron.core.transformer.v4_attention_kernels")
+    # A prior gluon test may have called the lazy loader, which caches the
+    # backends as module attributes; drop them + reload so this asserts the
+    # import-time (no eager binding) contract independent of test order.
+    for attr in ("v4_attention_gluon", "v4_csa_attention_gluon"):
+        if hasattr(pkg, attr):
+            delattr(pkg, attr)
+    importlib.reload(pkg)
+    assert hasattr(pkg, "load_gluon_attention_backends")
+    # gluon entries must NOT be module-level attributes (would mean eager import).
+    assert not hasattr(pkg, "v4_attention_gluon")
+    assert not hasattr(pkg, "v4_csa_attention_gluon")
+
+
+def test_importing_kernels_package_does_not_pull_gluon():
+    """Reloading the kernels package must not import ``_gluon_dsa`` as a side effect."""
+    # Purge any previously-imported gluon modules so this asserts the *package*
+    # import path, independent of test ordering (e.g. the gfx950 gluon UT).
+    for name in [m for m in sys.modules if "_gluon_dsa" in m]:
+        del sys.modules[name]
+    pkg = importlib.import_module("primus.backends.megatron.core.transformer.v4_attention_kernels")
+    importlib.reload(pkg)
+    assert not any(
+        "_gluon_dsa" in m for m in sys.modules
+    ), "importing v4_attention_kernels must not eagerly import the gluon backend"
+
+
+def test_attention_module_uses_lazy_gluon_helpers():
+    """``deepseek_v4_attention`` imports the lazy loader + arch guard, not gluon entries."""
+    mod = importlib.import_module("primus.backends.megatron.core.transformer.deepseek_v4_attention")
+    assert hasattr(mod, "load_gluon_attention_backends")
+    assert hasattr(mod, "_require_gfx950")
+    # no eager module-level gluon entry bindings
+    assert not hasattr(mod, "v4_attention_gluon")
+    assert not hasattr(mod, "v4_csa_attention_gluon")
+
+
+def test_require_gfx950_rejects_non_gfx950(monkeypatch):
+    """``_require_gfx950`` raises a clear error on a non-gfx950 device."""
+    torch = pytest.importorskip("torch")
+    from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+        _require_gfx950,
+    )
+
+    class _FakeProps:
+        gcnArchName = "gfx942:sramecc+:xnack-"
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _idx: _FakeProps())
+    with pytest.raises(RuntimeError, match="gfx950"):
+        _require_gfx950()
+
+
+def test_require_gfx950_rejects_no_device(monkeypatch):
+    """``_require_gfx950`` raises when no accelerator is available."""
+    torch = pytest.importorskip("torch")
+    from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+        _require_gfx950,
+    )
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    with pytest.raises(RuntimeError, match="gfx950"):
+        _require_gfx950()
+
+
+def test_require_gfx950_accepts_gfx950(monkeypatch):
+    """``_require_gfx950`` passes on a gfx950 device."""
+    torch = pytest.importorskip("torch")
+    from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+        _require_gfx950,
+    )
+
+    class _FakeProps:
+        gcnArchName = "gfx950:sramecc+:xnack-"
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _idx: _FakeProps())
+    _require_gfx950()  # must not raise
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py
new file mode 100644
index 000000000..2f4505008
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py
@@ -0,0 +1,500 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-3 P22 — ``core_attention`` integration for DeepSeek-V4.
+
+The dense (``compress_ratio == 0``) layers route their softmax-and-attend
+through ``provider.core_attention()`` (PrimusTurboAttention when
+``use_turbo_attention=True``, TEDotProductAttention otherwise).  HCA
+(``compress_ratio == 128``) and CSA (``compress_ratio == 4``) layers
+**do not** get a ``core_attention`` slot — their joint softmax and
+per-query top-K gather can't be expressed as stock flash-attention.
+
+This file is the unit-test side of P22's test gates:
+
+* **G18 (spec surface)** — the V4 attention submodules dataclass exposes
+  a ``core_attention`` slot; the spec helper emits it for dense layers
+  only.
+* **G18a (alias contract)** — when V4's per-head sink is on AND the
+  built ``core_attention`` advertises ``use_sink_attention=True``, the
+  attention module aliases ``self.core_attention.sinks`` to
+  ``self.attn_sink`` so the released-checkpoint key path
+  ``layers.{i}.attn.attn_sink`` still loads.
+* **G18b (CPU forward equivalence)** — a ``MockTurbo`` core-attention
+  class that replicates the eager-Python scaled-dot-product math is
+  injected; the dense-path output must match the original eager-Python
+  forward within a tight numerical tolerance.  The full
+  Turbo-vs-eager-Python equivalence at full V4-Flash dims is exercised
+  by the P22 smoke gate on ``mi355-gpu-12``.
+* **G18c (no core_attention on HCA / CSA)** — the spec helper does not
+  emit ``core_attention`` for ``compress_ratio in {4, 128}``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import fields
+
+import pytest
+import torch
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+# ---------------------------------------------------------------------------
+# G18 — submodules surface
+# ---------------------------------------------------------------------------
+
+
+class TestSubmodulesSurface:
+    """``DeepseekV4AttentionSubmodules`` must expose a ``core_attention`` slot."""
+
+    def test_core_attention_field_present(self):
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: F401
+            DeepSeekV4TransformerConfig,
+        )
+        from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+            DeepseekV4AttentionSubmodules,
+        )
+
+        names = {f.name for f in fields(DeepseekV4AttentionSubmodules)}
+        assert "core_attention" in names, (
+            "Plan-3 P22 added ``core_attention`` to "
+            f"DeepseekV4AttentionSubmodules; current fields: {sorted(names)}."
+        )
+
+
+# ---------------------------------------------------------------------------
+# G18 — spec helper emits core_attention for dense only
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def _tp1_distributed():
+    """1-rank torch.distributed (gloo) with Megatron model-parallel state."""
+    import os
+
+    import torch.distributed as dist
+    from megatron.core import parallel_state
+
+    if dist.is_initialized():
+        yield
+        return
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29539")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+
+    dist.init_process_group(backend="gloo", world_size=1, rank=0)
+    try:
+        parallel_state.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+        )
+        yield
+    finally:
+        parallel_state.destroy_model_parallel()
+        dist.destroy_process_group()
+
+
+def _make_v4_cfg(*, attn_sink: bool, attn_sliding_window: int = 0):
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    return DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=64,
+        num_attention_heads=4,
+        ffn_hidden_size=128,
+        kv_channels=16,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=16,
+        attn_sink=attn_sink,
+        attn_sliding_window=attn_sliding_window,
+        qk_pos_emb_head_dim=8,
+        num_query_groups=1,
+        multi_latent_attention=False,
+        params_dtype=torch.float32,
+        init_method=lambda w: torch.nn.init.normal_(w, std=0.02),
+        output_layer_init_method=lambda w: torch.nn.init.normal_(w, std=0.02),
+        use_cpu_initialization=True,
+        perform_initialization=True,
+        tensor_model_parallel_size=1,
+        pipeline_model_parallel_size=1,
+    )
+
+
+def _make_v4_submods(cfg, compress_ratio: int):
+    from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+        DeepSeekV4SpecProvider,
+    )
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import (
+        _build_v4_attention_submodules,
+    )
+
+    provider = DeepSeekV4SpecProvider(config=cfg)
+    return _build_v4_attention_submodules(
+        config=cfg,
+        provider=provider,
+        compress_ratio=compress_ratio,
+    )
+
+
+class TestSpecEmission:
+    """The spec helper emits ``core_attention`` only for dense layers."""
+
+    def test_dense_emits_core_attention(self, _tp1_distributed):
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4)
+        submods = _make_v4_submods(cfg, compress_ratio=0)
+        assert submods.core_attention is not None, "compress_ratio == 0 must emit a ``core_attention`` spec."
+
+    def test_csa_no_core_attention(self, _tp1_distributed):
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4)
+        submods = _make_v4_submods(cfg, compress_ratio=4)
+        assert submods.core_attention is None, (
+            "compress_ratio == 4 (CSA) must NOT emit ``core_attention`` — "
+            "per-query top-K gather is not a flash-attn pattern."
+        )
+
+    def test_hca_no_core_attention(self, _tp1_distributed):
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4)
+        submods = _make_v4_submods(cfg, compress_ratio=128)
+        assert submods.core_attention is None, (
+            "compress_ratio == 128 (HCA) must NOT emit ``core_attention`` — "
+            "joint softmax across two key streams needs an LSE-returning kernel."
+        )
+
+
+# ---------------------------------------------------------------------------
+# G18 — alias contract + forward equivalence
+# ---------------------------------------------------------------------------
+
+
+class _MockTurboCoreAttention(torch.nn.Module):
+    """Stand-in for :class:`PrimusTurboAttention`.
+
+    Replicates the eager-Python scaled-dot-product math (with optional
+    learned per-head sinks) on CPU, so the V4 attention module's dense
+    forward can be exercised end-to-end without a CUDA / Turbo build.
+
+    Public surface that ``DeepseekV4Attention`` reads:
+
+    * ``self.use_sink_attention`` — whether learned sinks are honored.
+    * ``self.sinks`` — ``[num_attention_heads]`` parameter (aliased by
+      V4's ``self.attn_sink`` after construction).
+    * ``forward(q, k, v, mask, attn_mask_type=...)`` — accepts ``sbhd``
+      Q / K / V (K, V may be MQA-shape ``[S, B, 1, D]``).
+    """
+
+    def __init__(
+        self,
+        *,
+        config,
+        layer_number,
+        attn_mask_type,
+        attention_type,
+        softmax_scale,
+        k_channels=None,
+        v_channels=None,
+        cp_comm_type="p2p",
+        pg_collection=None,
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_number = layer_number
+        self.softmax_scale = float(softmax_scale)
+        self.use_sink_attention = bool(getattr(config, "attn_sink", False))
+        self._head_dim = int(k_channels or config.kv_channels)
+        self._num_heads = int(config.num_attention_heads)
+        if self.use_sink_attention:
+            self.sinks = torch.nn.Parameter(torch.zeros(self._num_heads))
+        else:
+            self.sinks = None
+
+    def forward(self, q, k, v, mask, attn_mask_type=None):
+        # Inputs in qkv_format="sbhd": [S, B, H_q/H_kv, D].
+        S_q, B, H_q, D = q.shape
+        S_k = k.shape[0]
+        H_kv = k.shape[2]
+
+        # MQA broadcast.
+        if H_kv != H_q:
+            k = k.expand(S_k, B, H_q, D)
+            v = v.expand(S_k, B, H_q, D)
+
+        # [S, B, H, D] -> [B, H, S, D].
+        q_bh = q.permute(1, 2, 0, 3).contiguous().float()
+        k_bh = k.permute(1, 2, 0, 3).contiguous().float()
+        v_bh = v.permute(1, 2, 0, 3).contiguous()
+
+        logits = torch.matmul(q_bh, k_bh.transpose(-2, -1)) * self.softmax_scale
+
+        # Causal mask.
+        if attn_mask_type == AttnMaskType.causal:
+            causal = torch.full((S_q, S_k), float("-inf"), device=q.device)
+            causal = torch.triu(causal, diagonal=1)
+            logits = logits + causal
+
+        # Sink column.
+        if self.use_sink_attention and self.sinks is not None:
+            sink_col = self.sinks.float().view(1, H_q, 1, 1).expand(B, H_q, S_q, 1)
+            logits_aug = torch.cat([logits, sink_col], dim=-1)
+            logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach()
+            probs = logits_aug.softmax(dim=-1)[..., :-1]
+        else:
+            logits = logits - logits.amax(dim=-1, keepdim=True).detach()
+            probs = logits.softmax(dim=-1)
+
+        out_bh = torch.matmul(probs.to(v_bh.dtype), v_bh)  # [B, H, S, D]
+
+        # [B, H, S, D] -> [S, B, H, D] -> [S, B, H*D].
+        out_sbh = out_bh.permute(2, 0, 1, 3).contiguous()
+        return out_sbh.view(S_q, B, H_q * D)
+
+
+def _build_v4_attention_with(cfg, *, core_attention_class):
+    """Build a 1L V4 dense attention with an explicit ``core_attention`` class."""
+    from primus.backends.megatron.core.transformer.deepseek_v4_attention import (
+        DeepseekV4Attention,
+    )
+    from primus.backends.megatron.core.transformer.dual_rope import DualRoPE
+
+    submods = _make_v4_submods(cfg, compress_ratio=0)
+    if core_attention_class is None:
+        submods.core_attention = None
+    else:
+        submods.core_attention = ModuleSpec(module=core_attention_class)
+
+    rope = DualRoPE(
+        rotary_dim=cfg.qk_pos_emb_head_dim,
+        rope_theta=10000.0,
+        compress_rope_theta=10000.0,
+    )
+    return DeepseekV4Attention(
+        config=cfg,
+        rope=rope,
+        compress_ratio=0,
+        submodules=submods,
+        layer_number=0,
+    )
+
+
+class TestG18AliasAndForward:
+    """``self.attn_sink`` aliasing + dense forward equivalence."""
+
+    def test_sink_alias_when_turbo_supports_sink(self, _tp1_distributed):
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=0)
+        attn = _build_v4_attention_with(cfg, core_attention_class=_MockTurboCoreAttention)
+
+        assert attn.core_attention is not None
+        assert attn._use_core_attention is True
+        assert attn.attn_sink is not None
+        assert attn.core_attention.sinks is attn.attn_sink, (
+            "Plan-3 P22: when the core-attention class advertises "
+            "``use_sink_attention=True`` and V4 has ``attn_sink=True``, "
+            "the V4 attention module must alias ``core_attention.sinks`` "
+            "to ``self.attn_sink`` so the released-checkpoint key path "
+            "``layers.{i}.attn.attn_sink`` still loads via state-dict."
+        )
+
+    def test_no_alias_when_core_attention_lacks_sink(self, _tp1_distributed):
+        # Default _build_v4_attention_submodules emits provider.core_attention()
+        # which is TEDotProductAttention (no use_sink_attention attr).
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4)
+        attn = _build_v4_attention_with(cfg, core_attention_class=None)
+        # Re-emit core_attention via provider so we exercise the TE path.
+        # _build_v4_attention_with(core_attention_class=None) drops the slot;
+        # the default emission already builds TEDotProductAttention which
+        # does NOT advertise use_sink_attention -> no alias, eager fallback.
+        # Reproduce the default emission and rebuild.
+        from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
+            DeepSeekV4SpecProvider,
+        )
+
+        provider = DeepSeekV4SpecProvider(config=cfg)
+        attn2 = _build_v4_attention_with(cfg, core_attention_class=provider.core_attention())
+        assert attn2.core_attention is not None, "TE core_attention must build"
+        assert attn2._use_core_attention is False, (
+            "When the built core_attention does not support learned sinks "
+            "(e.g. TEDotProductAttention) and V4 sink is on, the dense "
+            "path must fall back to eager-Python so the inline "
+            "softmax-with-sink math still produces the correct output."
+        )
+
+    @pytest.mark.skipif(
+        not torch.cuda.is_available(),
+        reason="V4 specs build TE-parallel linears whose forward requires CUDA.",
+    )
+    def test_dense_forward_matches_eager(self, _tp1_distributed):
+        torch.manual_seed(0)
+        cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=0)
+
+        attn_eager = _build_v4_attention_with(cfg, core_attention_class=None)
+        attn_turbo = _build_v4_attention_with(cfg, core_attention_class=_MockTurboCoreAttention)
+
+        # Copy weights so both modules have identical parameters.
+        with torch.no_grad():
+            for (n_e, p_e), (n_t, p_t) in zip(
+                attn_eager.named_parameters(),
+                attn_turbo.named_parameters(),
+            ):
+                if p_e.shape == p_t.shape:
+                    p_t.copy_(p_e)
+
+        # Sanity: both should have the same attn_sink storage.
+        assert attn_eager.attn_sink is not None
+        assert attn_turbo.core_attention.sinks is attn_turbo.attn_sink
+
+        # Move both modules onto CUDA — the V4 spec emits TE-parallel
+        # linears whose forward requires CUDA tensors.  V4's DualRoPE is
+        # held by reference (``self._rope = [rope]``) so it isn't moved
+        # by ``attn.to(device)``; move it explicitly.
+        device = torch.device("cuda")
+        attn_eager = attn_eager.to(device)
+        attn_turbo = attn_turbo.to(device)
+        attn_eager._rope[0].to(device)
+        attn_turbo._rope[0].to(device)
+
+        B, S, D = 2, 7, cfg.hidden_size
+        hidden = torch.randn(B, S, D, device=device)
+        position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, S)
+
+        attn_eager.eval()
+        attn_turbo.eval()
+        with torch.no_grad():
+            out_eager = attn_eager(hidden, position_ids)
+            out_turbo = attn_turbo(hidden, position_ids)
+
+        assert out_eager.shape == out_turbo.shape == (B, S, D)
+        max_abs = (out_eager - out_turbo).abs().max().item()
+        assert max_abs < 5e-3, (
+            f"Dense forward via core_attention diverged from eager-Python "
+            f"forward beyond tolerance: max-abs={max_abs}.  Plan-3 P22 "
+            "expects the two paths to compute identical math (the mock "
+            "core_attention replicates the eager softmax-with-sink kernel)."
+        )
+
+
+# ---------------------------------------------------------------------------
+# G18 — sink-attention args plumbing
+# ---------------------------------------------------------------------------
+
+
+class TestSinkAttentionArgsPlumbing:
+    """``deepseek_v4_builder._maybe_plumb_v4_sink_attention_args``."""
+
+    def test_plumbing_fires_when_turbo_attention_on_seq_le_window(self):
+        """Plumbing zeros out the window when it covers the full sequence
+        (mathematically equivalent to full causal — avoids the aiter
+        Triton SWA gap)."""
+        from types import SimpleNamespace
+
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+            _maybe_plumb_v4_sink_attention_args,
+        )
+
+        args = SimpleNamespace(
+            enable_primus_turbo=True,
+            use_turbo_attention=True,
+            attn_sink=True,
+            attn_sliding_window=128,
+            seq_length=128,
+        )
+        _maybe_plumb_v4_sink_attention_args(args)
+
+        assert args.use_sink_attention is True
+        # Window == seq_length -> drop window (full causal is identical).
+        assert args.sink_sliding_window == 0
+        assert args.sink_window_even_layers_only is False
+
+    def test_plumbing_warns_when_real_swa_requested(self):
+        """When the V4 config genuinely needs SWA (seq > window) the
+        plumbing zeros out the window with a warning — aiter Triton
+        flash-attn doesn't support SWA yet, so the V4 dense layer
+        attends to all causal tokens (deviation from V4-Flash math)."""
+        from types import SimpleNamespace
+
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+            _maybe_plumb_v4_sink_attention_args,
+        )
+
+        args = SimpleNamespace(
+            enable_primus_turbo=True,
+            use_turbo_attention=True,
+            attn_sink=True,
+            attn_sliding_window=128,
+            seq_length=4096,
+        )
+        _maybe_plumb_v4_sink_attention_args(args)
+
+        assert args.use_sink_attention is True
+        # SWA requested but kernel doesn't support it -> drop window.
+        assert args.sink_sliding_window == 0
+        assert args.sink_window_even_layers_only is False
+
+    def test_plumbing_skips_when_turbo_off(self):
+        from types import SimpleNamespace
+
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+            _maybe_plumb_v4_sink_attention_args,
+        )
+
+        args = SimpleNamespace(
+            enable_primus_turbo=False,
+            use_turbo_attention=False,
+            attn_sink=True,
+            attn_sliding_window=128,
+        )
+        _maybe_plumb_v4_sink_attention_args(args)
+
+        # Plumbing must not touch args when Turbo is off.
+        assert not hasattr(args, "use_sink_attention") or args.use_sink_attention in (
+            None,
+            False,
+        )
+
+    def test_plumbing_skips_when_attn_sink_off(self):
+        from types import SimpleNamespace
+
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+            _maybe_plumb_v4_sink_attention_args,
+        )
+
+        args = SimpleNamespace(
+            enable_primus_turbo=True,
+            use_turbo_attention=True,
+            attn_sink=False,
+            attn_sliding_window=128,
+        )
+        _maybe_plumb_v4_sink_attention_args(args)
+
+        assert not hasattr(args, "use_sink_attention") or args.use_sink_attention in (
+            None,
+            False,
+        )
+
+    def test_plumbing_respects_explicit_user_override(self):
+        from types import SimpleNamespace
+
+        from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import (
+            _maybe_plumb_v4_sink_attention_args,
+        )
+
+        args = SimpleNamespace(
+            enable_primus_turbo=True,
+            use_turbo_attention=True,
+            attn_sink=True,
+            attn_sliding_window=128,
+            seq_length=4096,
+            use_sink_attention=True,
+            sink_sliding_window=256,  # user override
+        )
+        _maybe_plumb_v4_sink_attention_args(args)
+
+        # Don't clobber user-set sink_sliding_window.
+        assert args.sink_sliding_window == 256
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py
new file mode 100644
index 000000000..64393d0e0
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py
@@ -0,0 +1,140 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for the DeepSeek-V4 FP8 (E4M3) Indexer QK path.
+
+CPU-friendly. Validates that:
+
+* :func:`fake_quantize_fp8_e4m3` is a bounded-error, dtype-preserving,
+  zero-safe fake-quantization.
+* The FP8 Indexer selects almost the same top-k compressed positions as the
+  BF16 reference (high top-k overlap) — the QK precision drop must not change
+  which positions the CSA selector picks for the vast majority of queries.
+* Output shapes / sentinel (-1) semantics are unchanged by the FP8 path.
+"""
+
+from __future__ import annotations
+
+import copy
+
+import pytest
+import torch
+
+from primus.backends.megatron.core.transformer.indexer import (
+    Indexer,
+    fake_quantize_fp8_e4m3,
+)
+
+_HAS_FP8 = hasattr(torch, "float8_e4m3fn")
+
+
+def _fp8_cast_works() -> bool:
+    if not _HAS_FP8:
+        return False
+    try:
+        torch.zeros(4).to(torch.float8_e4m3fn).to(torch.float32)
+        return True
+    except (RuntimeError, TypeError):
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _fp8_cast_works(), reason="torch.float8_e4m3fn cast unsupported on this build/device"
+)
+
+
+# ---------------------------------------------------------------------------
+# fake_quantize_fp8_e4m3
+# ---------------------------------------------------------------------------
+
+
+def test_fake_quant_preserves_dtype_and_shape():
+    x = torch.randn(3, 5, 7, dtype=torch.float32)
+    xq = fake_quantize_fp8_e4m3(x)
+    assert xq.dtype == x.dtype
+    assert xq.shape == x.shape
+
+
+def test_fake_quant_zero_input_is_safe():
+    x = torch.zeros(4, 4)
+    xq = fake_quantize_fp8_e4m3(x)
+    assert torch.equal(xq, x)
+
+
+def test_fake_quant_bounded_relative_error():
+    # E4M3 has 3 mantissa bits -> ~2^-3 = 12.5% worst-case step; with dynamic
+    # per-tensor scaling the typical relative error on sizable values is small.
+    torch.manual_seed(0)
+    x = torch.randn(2048, dtype=torch.float32) * 2.0
+    xq = fake_quantize_fp8_e4m3(x)
+    # Compare only on non-tiny values (tiny values have large relative error
+    # but negligible absolute impact on the QK dot product).
+    big = x.abs() > 0.1 * x.abs().max()
+    rel = ((xq[big] - x[big]).abs() / x[big].abs()).max().item()
+    assert rel < 0.2, f"FP8 fake-quant relative error too large: {rel}"
+
+
+# ---------------------------------------------------------------------------
+# Indexer FP8 QK path vs BF16 reference
+# ---------------------------------------------------------------------------
+
+
+def _make_indexer(use_fp8_qk: bool) -> Indexer:
+    return Indexer(
+        hidden_size=64,
+        index_head_dim=16,
+        index_n_heads=4,
+        index_topk=4,
+        compress_ratio=4,
+        use_fp8_qk=use_fp8_qk,
+    )
+
+
+def _topk_overlap(idx_a: torch.Tensor, idx_b: torch.Tensor) -> float:
+    """Mean per-query Jaccard overlap of selected (valid) pool positions."""
+    B, S, _K = idx_a.shape
+    overlaps = []
+    for b in range(B):
+        for s in range(S):
+            a = {int(i) for i in idx_a[b, s].tolist() if i >= 0}
+            c = {int(i) for i in idx_b[b, s].tolist() if i >= 0}
+            if not a and not c:
+                continue
+            union = a | c
+            overlaps.append(len(a & c) / max(1, len(union)))
+    return sum(overlaps) / max(1, len(overlaps))
+
+
+def test_fp8_indexer_topk_overlaps_bf16_reference():
+    torch.manual_seed(0)
+    ref = _make_indexer(use_fp8_qk=False)
+    fp8 = _make_indexer(use_fp8_qk=True)
+    # Share identical weights so the only difference is the FP8 QK rounding.
+    fp8.load_state_dict(copy.deepcopy(ref.state_dict()))
+
+    hidden = torch.randn(2, 32, 64, dtype=torch.float32)
+    with torch.no_grad():
+        idx_ref, _ = ref(hidden)
+        idx_fp8, _ = fp8(hidden)
+
+    assert idx_ref.shape == idx_fp8.shape == (2, 32, 4)
+    # Sentinel semantics preserved (same set of masked-out / early queries).
+    assert torch.equal(idx_ref < 0, idx_fp8 < 0)
+
+    overlap = _topk_overlap(idx_ref, idx_fp8)
+    assert overlap >= 0.6, f"FP8 indexer top-k overlap with BF16 too low: {overlap}"
+
+
+def test_fp8_flag_off_is_identical_to_reference():
+    torch.manual_seed(0)
+    a = _make_indexer(use_fp8_qk=False)
+    b = _make_indexer(use_fp8_qk=False)
+    b.load_state_dict(copy.deepcopy(a.state_dict()))
+    hidden = torch.randn(1, 16, 64, dtype=torch.float32)
+    with torch.no_grad():
+        idx_a, _ = a(hidden)
+        idx_b, _ = b(hidden)
+    assert torch.equal(idx_a, idx_b)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py
new file mode 100644
index 000000000..d76ac5548
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py
@@ -0,0 +1,300 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Gluon DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950).
+
+Unlike a raw sparse-MLA test (separate 64-rope, random absolute-token topk),
+this validates the *production* V4 invocation paths — the autograd adapters in
+``v4_csa_attention_gluon`` that DeepseekV4Attention actually dispatches to —
+against the eager V4 references for all three layer kinds:
+
+* ``compress_ratio == 0``   (dense / SWA)  -> :func:`v4_attention_gluon`
+* ``compress_ratio == 128`` (HCA)          -> :func:`v4_attention_gluon`
+* ``compress_ratio == 4``   (CSA)          -> :func:`v4_csa_attention_gluon`
+
+The V4 layout has ``head_dim = 512`` with RoPE applied *in-place* (K = V = 512,
+score over 512); the adapter feeds the 512 latent as the gluon "lora" with a
+zero rope pad and builds ``kv = [local ++ pool]`` / ``topk = [SWA window ++
+pool]`` (padded to a multiple of 64 so the gluon bwd dKV tiling is valid — HCA
+160 -> 192). We compare full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool,
+dsink) of the bf16 gluon adapter against the fp32 eager reference.
+
+GPU-only; skipped off gfx950 / when Gluon is unavailable. ``B = 2`` to exercise
+the per-batch token-flattening + topk-offset logic.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("Gluon DSA kernels require CUDA / HIP", allow_module_level=True)
+
+# Gluon is an experimental Triton submodule; skip cleanly if absent.
+pytest.importorskip("triton", reason="Triton not installed")
+try:
+    from triton.experimental import gluon  # noqa: F401
+except Exception:  # pragma: no cover - environment-dependent
+    pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True)
+
+_ARCH = torch.cuda.get_device_properties(0).gcnArchName
+if "gfx950" not in _ARCH:
+    pytest.skip(f"ported Gluon DSA kernels target gfx950; got {_ARCH}", allow_module_level=True)
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import (  # noqa: E402
+    eager_v4_attention,
+    eager_v4_csa_attention,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon import (  # noqa: E402
+    v4_attention_gluon,
+    v4_csa_attention_gluon,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_triton import (  # noqa: E402
+    v4_attention_v2,
+    v4_csa_attention_v2,
+)
+
+D = 512  # V4 head_dim (RoPE baked in-place)
+_SWA = 128
+
+# Fused single-latent (K == V) sparse-MLA backends sharing the V4-form adapter:
+# (dense/HCA attention wrapper, CSA-from-pool wrapper).
+_BACKENDS = {
+    "gluon": (v4_attention_gluon, v4_csa_attention_gluon),
+    "triton_v2": (v4_attention_v2, v4_csa_attention_v2),
+}
+
+
+# ---------------------------------------------------------------------------
+# Comparison helpers (bf16 kernel vs fp32 eager autograd)
+# ---------------------------------------------------------------------------
+
+
+def _stats(a, b, *, sig=1e-2):
+    a = a.float()
+    b = b.float()
+    d = (a - b).abs()
+    m = b.abs() > sig
+    rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0)
+    med_rel = rel.median().item() if rel.numel() else 0.0
+    av, bv = a.flatten(), b.flatten()
+    cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item()
+    return d.max().item(), med_rel, cos_err
+
+
+def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None):
+    max_abs, med_rel, cos_err = _stats(a, b, sig=sig)
+    print(f"    {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True)
+    assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}"
+    if med is not None:
+        assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}"
+    if cos is not None:
+        assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}"
+
+
+def _leaf(x, *, fp32):
+    """Detached leaf clone (fp32 for the eager ref, bf16 for the gluon adapter)."""
+    y = x.float() if fp32 else x.clone()
+    return y.detach().requires_grad_(True)
+
+
+# ---------------------------------------------------------------------------
+# (B, H, S, sink_init)
+# ---------------------------------------------------------------------------
+_CASES = [
+    (2, 16, 256, None),
+    (2, 16, 256, 1.0),
+    (1, 32, 512, -1.0),
+]
+
+
+@pytest.mark.parametrize("backend", list(_BACKENDS))
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_dense_matches_eager(B, H, S, sink_init, backend):
+    """cr=0 dense/SWA: sparse-MLA v4_attention_v1 fwd+bwd vs eager_v4_attention."""
+    attn_fn = _BACKENDS[backend][0]
+    g = torch.Generator(device="cuda").manual_seed(0)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = attn_fn(
+        qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale
+    )
+    og.backward(do)
+
+    qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = latf.unsqueeze(1).expand(B, H, S, D)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=_SWA,
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[{backend} V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("backend", list(_BACKENDS))
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_hca_matches_eager(B, H, S, sink_init, backend):
+    """cr=128 HCA: sparse-MLA v4_attention_v1 (local++pool) fwd+bwd vs eager_v4_attention."""
+    attn_fn = _BACKENDS[backend][0]
+    cr = 128
+    P = max(S // cr, 1)
+    g = torch.Generator(device="cuda").manual_seed(2)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    # HCA causal pool mask [S, P]: pool slot p visible to query s iff (p+1)*cr-1 <= s.
+    ti = torch.arange(S, device="cuda").view(S, 1)
+    ps = torch.arange(P, device="cuda").view(1, P)
+    pool_mask = torch.where(
+        ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda")
+    ).to(torch.bfloat16)
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    og = attn_fn(
+        qg,
+        kg,
+        kg,
+        sink=sg,
+        swa_window=_SWA,
+        additive_mask=pool_mask,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+        hca_local_seqlen=S,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32)
+    full_mask = torch.cat([local_mask, pool_mask.float()], dim=1)  # [S, S+P]
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=0,
+        additive_mask=full_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[{backend} V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("backend", list(_BACKENDS))
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_csa_matches_eager(B, H, S, sink_init, backend):
+    """cr=4 CSA: sparse-MLA v4_csa_attention_v1 fwd+bwd vs eager_v4_csa_attention."""
+    csa_fn = _BACKENDS[backend][1]
+    P = max(S // 4, 1)
+    K = min(128, P)
+    g = torch.Generator(device="cuda").manual_seed(3)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    # Per-query top-K pool indices in [0, P), with ~1/8 invalid (-1).
+    topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32)
+    drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125
+    topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs)
+    idx = topk_idxs.clamp(0, P - 1).long()
+    bidx = torch.arange(B, device="cuda").view(B, 1, 1)
+    sparse_mask_inf = torch.where(
+        topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda")
+    )  # [B, S, K]
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    klg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = csa_fn(
+        qg,
+        klg,
+        klg,
+        poolg,
+        topk_idxs=topk_idxs,
+        sink=sg,
+        swa_window=_SWA,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    klf = latf.unsqueeze(1).expand(B, H, S, D)
+    gathered = poolf[bidx, idx]  # [B, S, K, D] (differentiable gather into poolf)
+    of = eager_v4_csa_attention(
+        qf,
+        klf,
+        klf,
+        gathered,
+        sink=sf,
+        swa_window=_SWA,
+        sparse_mask=sparse_mask_inf.float(),
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[{backend} V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py
new file mode 100644
index 000000000..012e0d988
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py
@@ -0,0 +1,277 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""gluon_v2 DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950),
+against the fp32 eager references.
+
+gluon_v2 = Gluon sparse-MLA forward + plain-Triton chunked-gather backward. This
+validates the production V4 invocation paths (the autograd adapters in
+``v4_csa_attention_gluon_v2`` that DeepseekV4Attention would dispatch to) for all
+three layer kinds:
+
+* ``compress_ratio == 0``   (dense / SWA)  -> :func:`v4_attention_gluon_v2`
+* ``compress_ratio == 128`` (HCA)          -> :func:`v4_attention_gluon_v2`
+* ``compress_ratio == 4``   (CSA)          -> :func:`v4_csa_attention_gluon_v2`
+
+We compare full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool, dsink) of the
+bf16 gluon_v2 adapter against the fp32 eager reference (the same tolerances the
+gluon backend UT uses). GPU-only; skipped off gfx950 / when Gluon is unavailable.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("gluon_v2 kernels require CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+try:
+    from triton.experimental import gluon  # noqa: F401
+except Exception:  # pragma: no cover - environment-dependent
+    pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True)
+
+_ARCH = torch.cuda.get_device_properties(0).gcnArchName
+if "gfx950" not in _ARCH:
+    pytest.skip(f"gluon_v2 Gluon fwd targets gfx950; got {_ARCH}", allow_module_level=True)
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import (  # noqa: E402
+    eager_v4_attention,
+    eager_v4_csa_attention,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v2 import (  # noqa: E402
+    v4_attention_gluon_v2,
+    v4_csa_attention_gluon_v2,
+)
+
+# Fail early with the build hint if the installed triton can't compile the Gluon fwd.
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2.dsa_fwd_v4_gluon import (  # noqa: E402
+        _gluon_available,
+    )
+
+    if not _gluon_available():
+        pytest.skip("Gluon dialect unavailable for gluon_v2 fwd", allow_module_level=True)
+except Exception:  # noqa: BLE001
+    pass
+
+D = 512  # V4 head_dim (RoPE baked in-place)
+_SWA = 128
+
+
+def _stats(a, b, *, sig=1e-2):
+    a = a.float()
+    b = b.float()
+    d = (a - b).abs()
+    m = b.abs() > sig
+    rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0)
+    med_rel = rel.median().item() if rel.numel() else 0.0
+    av, bv = a.flatten(), b.flatten()
+    cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item()
+    return d.max().item(), med_rel, cos_err
+
+
+def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None):
+    max_abs, med_rel, cos_err = _stats(a, b, sig=sig)
+    print(f"    {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True)
+    assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}"
+    if med is not None:
+        assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}"
+    if cos is not None:
+        assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}"
+
+
+def _leaf(x, *, fp32):
+    y = x.float() if fp32 else x.clone()
+    return y.detach().requires_grad_(True)
+
+
+_CASES = [
+    (2, 16, 256, None),
+    (2, 16, 256, 1.0),
+    (1, 32, 512, -1.0),
+]
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_gluon_v2_dense_matches_eager(B, H, S, sink_init):
+    """cr=0 dense/SWA: gluon_v2 fwd+bwd vs eager_v4_attention."""
+    g = torch.Generator(device="cuda").manual_seed(0)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_attention_gluon_v2(
+        qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale
+    )
+    og.backward(do)
+
+    qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = latf.unsqueeze(1).expand(B, H, S, D)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=_SWA,
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v2 V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_gluon_v2_hca_matches_eager(B, H, S, sink_init):
+    """cr=128 HCA: gluon_v2 (local++pool) fwd+bwd vs eager_v4_attention."""
+    cr = 128
+    P = max(S // cr, 1)
+    g = torch.Generator(device="cuda").manual_seed(2)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    ti = torch.arange(S, device="cuda").view(S, 1)
+    ps = torch.arange(P, device="cuda").view(1, P)
+    pool_mask = torch.where(
+        ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda")
+    ).to(torch.bfloat16)
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    og = v4_attention_gluon_v2(
+        qg,
+        kg,
+        kg,
+        sink=sg,
+        swa_window=_SWA,
+        additive_mask=pool_mask,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+        hca_local_seqlen=S,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32)
+    full_mask = torch.cat([local_mask, pool_mask.float()], dim=1)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=0,
+        additive_mask=full_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v2 V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_gluon_v2_csa_matches_eager(B, H, S, sink_init):
+    """cr=4 CSA: gluon_v2 fwd+bwd vs eager_v4_csa_attention."""
+    P = max(S // 4, 1)
+    K = min(128, P)
+    g = torch.Generator(device="cuda").manual_seed(3)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32)
+    drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125
+    topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs)
+    idx = topk_idxs.clamp(0, P - 1).long()
+    bidx = torch.arange(B, device="cuda").view(B, 1, 1)
+    sparse_mask_inf = torch.where(
+        topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda")
+    )
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    klg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_csa_attention_gluon_v2(
+        qg,
+        klg,
+        klg,
+        poolg,
+        topk_idxs=topk_idxs,
+        sink=sg,
+        swa_window=_SWA,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    klf = latf.unsqueeze(1).expand(B, H, S, D)
+    gathered = poolf[bidx, idx]
+    of = eager_v4_csa_attention(
+        qf,
+        klf,
+        klf,
+        gathered,
+        sink=sf,
+        swa_window=_SWA,
+        sparse_mask=sparse_mask_inf.float(),
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v2 V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py
new file mode 100644
index 000000000..f9a31d52f
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py
@@ -0,0 +1,322 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""gluon_v3 DeepSeek-V4 attention fwd+bwd correctness against fp32 eager refs."""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("gluon_v3 kernels require CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+try:
+    from triton.experimental import gluon  # noqa: F401
+except Exception:  # pragma: no cover - environment-dependent
+    pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True)
+
+_ARCH = torch.cuda.get_device_properties(0).gcnArchName
+if "gfx950" not in _ARCH:
+    pytest.skip(f"gluon_v3 targets gfx950; got {_ARCH}", allow_module_level=True)
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import (  # noqa: E402
+    eager_v4_attention,
+    eager_v4_csa_attention,
+)
+
+try:
+    from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import (  # noqa: E402
+        sparse_mla_bwd_v4_gluon_v3,
+        sparse_mla_fwd_v4_gluon_v3,
+    )
+    from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v3 import (  # noqa: E402
+        v4_attention_gluon_v3,
+        v4_csa_attention_gluon_v3,
+    )
+except ImportError as exc:  # pragma: no cover - environment-dependent
+    pytest.skip(f"gluon_v3 unavailable: {exc}", allow_module_level=True)
+
+D = 512
+_ROPE = 64
+_SWA = 128
+
+
+def _stats(a, b, *, sig=1e-2):
+    a = a.float()
+    b = b.float()
+    d = (a - b).abs()
+    m = b.abs() > sig
+    rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0)
+    med_rel = rel.median().item() if rel.numel() else 0.0
+    av, bv = a.flatten(), b.flatten()
+    cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item()
+    return d.max().item(), med_rel, cos_err
+
+
+def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None):
+    max_abs, med_rel, cos_err = _stats(a, b, sig=sig)
+    print(f"    {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True)
+    assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}"
+    if med is not None:
+        assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}"
+    if cos is not None:
+        assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}"
+
+
+def _leaf(x, *, fp32):
+    y = x.float() if fp32 else x.clone()
+    return y.detach().requires_grad_(True)
+
+
+def _eager_sparse_mla(q, kv, topk, sink, *, scale):
+    q_lora = q[:, :, :D]
+    kv_lora = kv[:, 0, :D]
+    safe = topk.clamp(min=0).long()
+    gathered = kv_lora[safe]
+    valid = topk >= 0
+    scores = torch.einsum("thd,tkd->thk", q_lora, gathered) * scale
+    scores = torch.where(valid[:, None, :], scores, torch.full_like(scores, float("-inf")))
+    if sink is not None:
+        sink_scores = sink.view(1, -1, 1).expand(q.shape[0], q.shape[1], 1)
+        all_scores = torch.cat([scores, sink_scores], dim=-1)
+    else:
+        all_scores = scores
+    probs = torch.softmax(all_scores, dim=-1)
+    out = torch.einsum("thk,tkd->thd", probs[:, :, : topk.shape[1]], gathered)
+    lse = torch.logsumexp(all_scores, dim=-1)
+    return out, lse
+
+
+_CASES = [
+    (1, 16, 128, None),
+    (1, 32, 256, 1.0),
+]
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_gluon_v3_dense_matches_eager(B, H, S, sink_init):
+    g = torch.Generator(device="cuda").manual_seed(10)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_attention_gluon_v3(
+        qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale
+    )
+    og.backward(do)
+
+    qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = latf.unsqueeze(1).expand(B, H, S, D)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=_SWA,
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v3 V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True)
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+def test_v4_gluon_v3_hca_matches_eager():
+    B, H, S, cr = 1, 32, 256, 128
+    P = max(S // cr, 1)
+    g = torch.Generator(device="cuda").manual_seed(11)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), 1.0, dtype=torch.float32, device="cuda")
+
+    ti = torch.arange(S, device="cuda").view(S, 1)
+    ps = torch.arange(P, device="cuda").view(1, P)
+    pool_mask = torch.where(
+        ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda")
+    ).to(torch.bfloat16)
+
+    qg, latg, poolg, sg = (
+        _leaf(q, fp32=False),
+        _leaf(lat, fp32=False),
+        _leaf(pool, fp32=False),
+        _leaf(sink, fp32=False),
+    )
+    kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    og = v4_attention_gluon_v3(
+        qg,
+        kg,
+        kg,
+        sink=sg,
+        swa_window=_SWA,
+        additive_mask=pool_mask,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+        hca_local_seqlen=S,
+    )
+    og.backward(do)
+
+    qf, latf, poolf, sf = (
+        _leaf(q, fp32=True),
+        _leaf(lat, fp32=True),
+        _leaf(pool, fp32=True),
+        _leaf(sink, fp32=True),
+    )
+    kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    full_mask = torch.cat(
+        [sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32), pool_mask.float()], dim=1
+    )
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=0,
+        additive_mask=full_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v3 V4 HCA] B={B} H={H} S={S} P={P}", flush=True)
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+def test_v4_gluon_v3_csa_adapter_matches_eager():
+    B, H, S = 1, 32, 256
+    P, K = max(S // 4, 1), 64
+    g = torch.Generator(device="cuda").manual_seed(12)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), -1.0, dtype=torch.float32, device="cuda")
+    topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32)
+    topk_idxs = torch.where(
+        torch.rand(B, S, K, generator=g, device="cuda") < 0.125, torch.full_like(topk_idxs, -1), topk_idxs
+    )
+    idx = topk_idxs.clamp(0, P - 1).long()
+    bidx = torch.arange(B, device="cuda").view(B, 1, 1)
+    sparse_mask_inf = torch.where(
+        topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda")
+    )
+
+    qg, latg, poolg, sg = (
+        _leaf(q, fp32=False),
+        _leaf(lat, fp32=False),
+        _leaf(pool, fp32=False),
+        _leaf(sink, fp32=False),
+    )
+    klg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_csa_attention_gluon_v3(
+        qg,
+        klg,
+        klg,
+        poolg,
+        topk_idxs=topk_idxs,
+        sink=sg,
+        swa_window=_SWA,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+    )
+    og.backward(do)
+
+    qf, latf, poolf, sf = (
+        _leaf(q, fp32=True),
+        _leaf(lat, fp32=True),
+        _leaf(pool, fp32=True),
+        _leaf(sink, fp32=True),
+    )
+    klf = latf.unsqueeze(1).expand(B, H, S, D)
+    gathered = poolf[bidx, idx]
+    of = eager_v4_csa_attention(
+        qf,
+        klf,
+        klf,
+        gathered,
+        sink=sf,
+        swa_window=_SWA,
+        sparse_mask=sparse_mask_inf.float(),
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[gluon_v3 V4 CSA adapter] B={B} H={H} S={S} P={P} K={K}", flush=True)
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("H,pool_k,T", [(64, 512, 64), (128, 1024, 64)], ids=["h64_topk640", "h128_topk1152"])
+def test_v4_gluon_v3_round9_csa_formula_path_matches_eager(H, pool_k, T):
+    g = torch.Generator(device="cuda").manual_seed(13 + H)
+    scale = 1.0 / math.sqrt(D)
+    q512 = torch.randn(T, H, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    kv512 = torch.randn(T + pool_k, 1, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    q = torch.cat([q512, torch.zeros(T, H, _ROPE, device="cuda", dtype=torch.bfloat16)], dim=-1).contiguous()
+    kv = torch.cat(
+        [kv512, torch.zeros(T + pool_k, 1, _ROPE, device="cuda", dtype=torch.bfloat16)], dim=-1
+    ).contiguous()
+    do = torch.randn(T, H, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.randn(H, generator=g, device="cuda", dtype=torch.float32) * 0.1
+
+    ti = torch.arange(T, device="cuda").view(T, 1)
+    win = ti - _SWA + 1 + torch.arange(_SWA, device="cuda").view(1, _SWA)
+    win = torch.where(win >= 0, win, torch.full_like(win, -1))
+    pool_topk = T + torch.randint(0, pool_k, (T, pool_k), generator=g, device="cuda", dtype=torch.int64)
+    topk = torch.cat([win, pool_topk], dim=1).to(torch.int32).contiguous()
+
+    qg, kvg, sg = _leaf(q, fp32=False), _leaf(kv, fp32=False), _leaf(sink, fp32=False)
+    og, lseg = sparse_mla_fwd_v4_gluon_v3(qg, kvg, topk, attn_sink=sg, kv_lora_rank=D, scale=scale)
+    dqg, dkvg, dsg = sparse_mla_bwd_v4_gluon_v3(
+        qg, kvg, og, do, topk, lseg, attn_sink=sg, kv_lora_rank=D, scale=scale
+    )
+
+    qf, kvf, sf = _leaf(q, fp32=True), _leaf(kv, fp32=True), _leaf(sink, fp32=True)
+    of, lsef = _eager_sparse_mla(qf, kvf, topk, sf, scale=scale)
+    of.backward(do.float())
+
+    print(f"\n[gluon_v3 round9 CSA formula] T={T} H={H} TOPK={topk.shape[1]}", flush=True)
+    _check("O", og, of, abs_tol=5e-2, med=3e-2, cos=2e-3)
+    _check("LSE", lseg, lsef, abs_tol=1e-2, sig=1e-3, med=1e-3, cos=1e-6)
+    _check("dQ", dqg[:, :, :D], qf.grad[:, :, :D], abs_tol=1e-1, sig=1e-3, med=3e-2, cos=2e-3)
+    _check("dKV", dkvg[:, :, :D], kvf.grad[:, :, :D], abs_tol=3e-1, sig=1e-3, med=5e-2, cos=5e-3)
+    _check("dSink", dsg, sf.grad, abs_tol=1.0, sig=1e-3, med=8e-2, cos=2e-2)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_distill_loss.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_distill_loss.py
new file mode 100644
index 000000000..7d2aab33c
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_distill_loss.py
@@ -0,0 +1,643 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for the CSA indexer distillation loss.
+
+Covers the two halves separately:
+
+* the loss itself -- KL against a known target, the all-masked-row guard, and
+  the fact that it actually produces indexer gradients;
+* the wiring -- ``v4_indexer_distill_loss_coeff`` gates both the loss and
+  whether the indexer parameters are trainable at all.
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+mla_module = pytest.importorskip(
+    "megatron.core.transformer.multi_latent_attention",
+    reason="MLA base module not importable in this environment",
+)
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: E402
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import DualRoPE  # noqa: E402
+from primus.backends.megatron.core.transformer.indexer_distill_loss import (  # noqa: E402
+    INDEXER_DISTILL_LOSS_NAME,
+    V4IndexerLossAutoScaler,
+    compute_indexer_distill_loss,
+    log_indexer_distill_loss,
+)
+
+_DTYPE = torch.float32
+
+
+@pytest.fixture(autouse=True)
+def _reset_aux_loss_scale():
+    """Keep the class-level scale override from leaking across tests."""
+    V4IndexerLossAutoScaler.set_loss_scale(None)
+    yield
+    V4IndexerLossAutoScaler.set_loss_scale(None)
+
+
+# ---------------------------------------------------------------------------
+# The loss itself
+# ---------------------------------------------------------------------------
+
+
+def test_kl_is_zero_when_indexer_matches_attention():
+    """A perfectly-predicting indexer incurs no loss.
+
+    Build the indexer scores so their softmax equals the (single-head)
+    attention distribution over the selected entries; KL must vanish.
+    """
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 1, 1, 3, 4, 8, 6
+    query = torch.randn(B, H, S, Dh, dtype=_DTYPE)
+    pool = torch.randn(B, P, Dh, dtype=_DTYPE)
+    topk_idxs = torch.arange(K, dtype=torch.long).view(1, 1, K).expand(B, S, K).contiguous()
+    scale = 1.0 / math.sqrt(Dh)
+
+    gathered = pool[torch.arange(B).view(B, 1, 1), topk_idxs]  # [B,S,K,Dh]
+    attn_logits = torch.einsum("bhsd,bskd->bhsk", query, gathered) * scale
+    # One head, so the target distribution is exactly this row's softmax; give
+    # the indexer the same logits.
+    index_scores = attn_logits[:, 0]
+
+    loss = compute_indexer_distill_loss(
+        index_topk_scores=index_scores,
+        topk_idxs=topk_idxs,
+        query=query,
+        pool=pool,
+        softmax_scale=scale,
+        loss_coeff=1.0,
+    )
+    assert loss.item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_kl_is_positive_and_scales_with_coeff():
+    """A mismatched indexer is penalised, linearly in the coefficient."""
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 2, 4, 5, 3, 8, 7
+    query = torch.randn(B, H, S, Dh, dtype=_DTYPE)
+    pool = torch.randn(B, P, Dh, dtype=_DTYPE)
+    topk_idxs = torch.randint(0, P, (B, S, K), dtype=torch.long)
+    index_scores = torch.randn(B, S, K, dtype=_DTYPE)
+    scale = 1.0 / math.sqrt(Dh)
+
+    kwargs = dict(
+        index_topk_scores=index_scores,
+        topk_idxs=topk_idxs,
+        query=query,
+        pool=pool,
+        softmax_scale=scale,
+    )
+    loss_1 = compute_indexer_distill_loss(loss_coeff=1.0, **kwargs)
+    loss_2 = compute_indexer_distill_loss(loss_coeff=2.0, **kwargs)
+
+    assert loss_1.item() > 0.0
+    assert loss_2.item() == pytest.approx(2.0 * loss_1.item(), rel=1e-5)
+
+
+def test_fully_masked_rows_do_not_produce_nan():
+    """Early queries can have zero legal compressed entries.
+
+    Those rows softmax over all -inf; the loss must neutralise them rather
+    than emit NaN (which would poison the whole step).
+    """
+    B, H, S, K, Dh, P = 1, 2, 4, 3, 8, 5
+    query = torch.randn(B, H, S, Dh, dtype=_DTYPE)
+    pool = torch.randn(B, P, Dh, dtype=_DTYPE)
+
+    topk_idxs = torch.randint(0, P, (B, S, K), dtype=torch.long)
+    index_scores = torch.randn(B, S, K, dtype=_DTYPE)
+    # Row 0 has no legal entry at all; row 1 has a single one.
+    topk_idxs[0, 0, :] = -1
+    index_scores[0, 0, :] = float("-inf")
+    topk_idxs[0, 1, 1:] = -1
+    index_scores[0, 1, 1:] = float("-inf")
+
+    loss = compute_indexer_distill_loss(
+        index_topk_scores=index_scores,
+        topk_idxs=topk_idxs,
+        query=query,
+        pool=pool,
+        softmax_scale=1.0 / math.sqrt(Dh),
+        loss_coeff=1e-2,
+    )
+    assert torch.isfinite(loss), f"loss must stay finite, got {loss}"
+
+
+def test_loss_produces_indexer_gradients():
+    """The KL flows back into whatever produced the index scores."""
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 1, 2, 4, 3, 8, 6
+    query = torch.randn(B, H, S, Dh, dtype=_DTYPE)
+    pool = torch.randn(B, P, Dh, dtype=_DTYPE)
+    topk_idxs = torch.randint(0, P, (B, S, K), dtype=torch.long)
+
+    index_scores = torch.randn(B, S, K, dtype=_DTYPE, requires_grad=True)
+    loss = compute_indexer_distill_loss(
+        index_topk_scores=index_scores,
+        topk_idxs=topk_idxs,
+        query=query,
+        pool=pool,
+        softmax_scale=1.0 / math.sqrt(Dh),
+        loss_coeff=1e-2,
+    )
+    loss.backward()
+
+    assert index_scores.grad is not None
+    assert torch.isfinite(index_scores.grad).all()
+    assert index_scores.grad.abs().sum() > 0.0
+
+
+def test_loss_does_not_flow_into_the_target_side():
+    """Distillation is one-directional: the target must not be trainable.
+
+    ``query`` and ``pool`` describe the distribution the indexer is supposed to
+    imitate. If the KL could reach them, the main attention would be rewarded
+    for becoming easier to predict instead of the indexer for predicting
+    better.
+    """
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 1, 2, 4, 3, 8, 6
+    query = torch.randn(B, H, S, Dh, dtype=_DTYPE, requires_grad=True)
+    pool = torch.randn(B, P, Dh, dtype=_DTYPE, requires_grad=True)
+    topk_idxs = torch.randint(0, P, (B, S, K), dtype=torch.long)
+    index_scores = torch.randn(B, S, K, dtype=_DTYPE, requires_grad=True)
+
+    loss = compute_indexer_distill_loss(
+        index_topk_scores=index_scores,
+        topk_idxs=topk_idxs,
+        query=query,
+        pool=pool,
+        softmax_scale=1.0 / math.sqrt(Dh),
+        loss_coeff=1e-2,
+    )
+    loss.backward()
+
+    assert query.grad is None, "the KL target must be detached"
+    assert pool.grad is None, "the compressed pool must be detached"
+    # ...while the side that is supposed to learn still gets its gradient.
+    assert index_scores.grad is not None and index_scores.grad.abs().sum() > 0.0
+
+
+def test_head_sum_is_all_reduced_when_the_heads_are_sharded():
+    """The target distribution needs every head before it is renormalised.
+
+    V4 gathers its Q projection so the sum is already complete, but the loss has
+    to reduce when told the heads are sharded -- otherwise dropping the gather
+    would quietly turn the target into a partial sum.
+    """
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 1, 2, 4, 3, 8, 6
+    kwargs = dict(
+        index_topk_scores=torch.randn(B, S, K, dtype=_DTYPE),
+        topk_idxs=torch.randint(0, P, (B, S, K), dtype=torch.long),
+        query=torch.randn(B, H, S, Dh, dtype=_DTYPE),
+        pool=torch.randn(B, P, Dh, dtype=_DTYPE),
+        softmax_scale=1.0 / math.sqrt(Dh),
+        loss_coeff=1.0,
+    )
+
+    calls = []
+
+    class _FakeGroup:
+        pass
+
+    group = _FakeGroup()
+
+    import primus.backends.megatron.core.transformer.indexer_distill_loss as mod
+
+    real_world_size = torch.distributed.get_world_size
+    real_all_reduce = torch.distributed.all_reduce
+    torch.distributed.get_world_size = lambda g=None: 2 if g is group else 1
+    torch.distributed.all_reduce = lambda t, group=None: calls.append((t.shape, group))
+    try:
+        loss = mod.compute_indexer_distill_loss(head_reduce_group=group, **kwargs)
+    finally:
+        torch.distributed.get_world_size = real_world_size
+        torch.distributed.all_reduce = real_all_reduce
+
+    assert torch.isfinite(loss)
+    assert calls, "the head sum was not reduced"
+    assert calls[0][0] == (B, S, K), "must reduce the head-summed target, not the per-head scores"
+    assert calls[0][1] is group
+
+
+def test_head_sum_is_not_reduced_without_a_group():
+    """No group means no collective -- the common V4 case must stay comm-free."""
+    torch.manual_seed(0)
+    B, H, S, K, Dh, P = 1, 2, 4, 3, 8, 6
+    calls = []
+    real_all_reduce = torch.distributed.all_reduce
+    torch.distributed.all_reduce = lambda t, group=None: calls.append(t)
+    try:
+        compute_indexer_distill_loss(
+            index_topk_scores=torch.randn(B, S, K, dtype=_DTYPE),
+            topk_idxs=torch.randint(0, P, (B, S, K), dtype=torch.long),
+            query=torch.randn(B, H, S, Dh, dtype=_DTYPE),
+            pool=torch.randn(B, P, Dh, dtype=_DTYPE),
+            softmax_scale=1.0 / math.sqrt(Dh),
+            loss_coeff=1.0,
+        )
+    finally:
+        torch.distributed.all_reduce = real_all_reduce
+    assert not calls
+
+
+def test_head_group_is_none_when_q_is_gathered(monkeypatch):
+    """V4's column-parallel Q gathers its output, so no group is resolved."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    attn = _make_csa_attention(coeff=1e-2)
+    assert attn.num_attention_heads_per_partition == int(attn.config.num_attention_heads)
+    assert attn._indexer_loss_head_group() is None
+
+
+def test_auto_scaler_is_transparent_forward_and_seeds_aux_backward():
+    """The scaler passes the tensor through and differentiates the aux loss."""
+    V4IndexerLossAutoScaler.set_loss_scale(torch.tensor(1.0))
+    x = torch.randn(3, 4, requires_grad=True)
+    aux_src = torch.randn(2, requires_grad=True)
+    aux_loss = aux_src.sum()
+
+    out = V4IndexerLossAutoScaler.apply(x, aux_loss)
+    torch.testing.assert_close(out, x, rtol=0, atol=0)
+
+    out.sum().backward()
+    # x keeps its own gradient, and the aux subgraph got seeded with ones.
+    assert x.grad is not None and torch.allclose(x.grad, torch.ones_like(x))
+    assert aux_src.grad is not None and torch.allclose(aux_src.grad, torch.ones_like(aux_src))
+
+
+def test_auto_scaler_applies_the_explicit_scale():
+    """An explicit override is the gradient seeded into the aux loss."""
+    V4IndexerLossAutoScaler.set_loss_scale(torch.tensor(0.25))
+    x = torch.randn(2, 2, requires_grad=True)
+    aux_src = torch.randn(3, requires_grad=True)
+
+    V4IndexerLossAutoScaler.apply(x, aux_src.sum()).sum().backward()
+
+    assert aux_src.grad is not None
+    torch.testing.assert_close(aux_src.grad, torch.full_like(aux_src, 0.25))
+
+
+def test_auto_scaler_follows_the_moe_aux_loss_scale(monkeypatch):
+    """Default behaviour: inherit the per-microbatch MoE aux-loss scale.
+
+    Seeding a bare 1.0 would make the effective coefficient
+    ``num_microbatches`` times too large under gradient accumulation, so the
+    scaler reads the quantity Megatron's schedule already computes.
+    """
+    import primus.backends.megatron.core.transformer.indexer_distill_loss as mod
+
+    monkeypatch.setattr(mod, "_moe_aux_loss_scale", lambda: torch.tensor(0.125))
+
+    aux_src = torch.randn(3, requires_grad=True)
+    x = torch.randn(2, 2, requires_grad=True)
+    V4IndexerLossAutoScaler.apply(x, aux_src.sum()).sum().backward()
+
+    assert aux_src.grad is not None
+    torch.testing.assert_close(aux_src.grad, torch.full_like(aux_src, 0.125))
+
+
+def test_auto_scaler_falls_back_to_one_without_megatron(monkeypatch):
+    """No override and no MoE scale available -> neutral seed, never a crash."""
+    import primus.backends.megatron.core.transformer.indexer_distill_loss as mod
+
+    monkeypatch.setattr(mod, "_moe_aux_loss_scale", lambda: None)
+
+    aux_src = torch.randn(3, requires_grad=True)
+    x = torch.randn(2, 2, requires_grad=True)
+    V4IndexerLossAutoScaler.apply(x, aux_src.sum()).sum().backward()
+
+    assert aux_src.grad is not None
+    torch.testing.assert_close(aux_src.grad, torch.ones_like(aux_src))
+
+
+def test_explicit_scale_wins_over_the_moe_scale(monkeypatch):
+    """``set_loss_scale`` is an override, not a default."""
+    import primus.backends.megatron.core.transformer.indexer_distill_loss as mod
+
+    monkeypatch.setattr(mod, "_moe_aux_loss_scale", lambda: torch.tensor(0.125))
+    V4IndexerLossAutoScaler.set_loss_scale(torch.tensor(2.0))
+
+    aux_src = torch.randn(3, requires_grad=True)
+    x = torch.randn(2, 2, requires_grad=True)
+    V4IndexerLossAutoScaler.apply(x, aux_src.sum()).sum().backward()
+
+    torch.testing.assert_close(aux_src.grad, torch.full_like(aux_src, 2.0))
+
+    # ...and clearing it restores the MoE-following default.
+    V4IndexerLossAutoScaler.set_loss_scale(None)
+    aux_src2 = torch.randn(3, requires_grad=True)
+    V4IndexerLossAutoScaler.apply(x, aux_src2.sum()).sum().backward()
+    torch.testing.assert_close(aux_src2.grad, torch.full_like(aux_src2, 0.125))
+
+
+# ---------------------------------------------------------------------------
+# Wiring: the coefficient gates training of the indexer
+# ---------------------------------------------------------------------------
+
+
+def _make_csa_attention(coeff: float, layer_number: int = 1) -> DeepseekV4Attention:
+    config = DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=64,
+        num_attention_heads=4,
+        num_query_groups=1,
+        kv_channels=16,
+        qk_pos_emb_head_dim=8,
+        qk_head_dim=8,
+        v_head_dim=16,
+        kv_lora_rank=16,
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sliding_window=0,
+        attn_sink=True,
+        compress_ratios=None,
+        compress_rope_theta=40000.0,
+        use_v4_attention_backend="eager",
+        use_v4_csa_attention_backend="eager",
+        layernorm_epsilon=1e-6,
+        norm_epsilon=1e-6,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        v4_indexer_distill_loss_coeff=coeff,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    return DeepseekV4Attention(
+        config, rope=rope, compress_ratio=4, submodules=None, layer_number=layer_number
+    )
+
+
+def test_indexer_frozen_when_coeff_is_zero(monkeypatch):
+    """Default (0.0): no loss, and the indexer stays out of the grad buckets."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    attn = _make_csa_attention(coeff=0.0)
+
+    assert attn.indexer_distill_enabled is False
+    assert attn.indexer is not None
+    assert not any(p.requires_grad for p in attn.indexer.parameters())
+
+
+def test_indexer_trainable_when_coeff_positive(monkeypatch):
+    """A positive coefficient unfreezes the indexer."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    attn = _make_csa_attention(coeff=1e-2)
+
+    assert attn.indexer_distill_enabled is True
+    assert attn.indexer is not None
+    assert all(p.requires_grad for p in attn.indexer.parameters())
+
+
+def test_forward_backward_reaches_indexer_weights(monkeypatch):
+    """End to end: a CSA step with the loss on gives the indexer gradients.
+
+    Without the distillation loss the indexer is unreachable from the output
+    (only argTopK indices are consumed), so a non-zero gradient here is proof
+    the aux objective is wired into the main backward.
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=1e-2).to(_DTYPE)
+    attn.train()
+
+    B, S = 1, 8  # P = S // 4 = 2 pool entries
+    hidden = torch.randn(B, S, attn.config.hidden_size, dtype=_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    out = attn(hidden, position_ids)
+    out.sum().backward()
+
+    assert attn.last_indexer_distill_loss is not None
+    assert torch.isfinite(attn.last_indexer_distill_loss)
+
+    grads = [p.grad for p in attn.indexer.parameters() if p.grad is not None]
+    assert grads, "indexer received no gradient at all"
+    total = sum(g.abs().sum().item() for g in grads)
+    assert math.isfinite(total) and total > 0.0, f"indexer gradient is degenerate: {total}"
+
+
+def test_kl_leaves_the_backbone_gradients_untouched(monkeypatch):
+    """Turning the loss on must not change a single backbone gradient.
+
+    The indexer is fed a detached hidden state and the KL target is detached,
+    so the auxiliary objective is entirely confined to the indexer's own
+    parameters. Running the same weights and the same input with the loss off
+    and with an absurdly large coefficient therefore has to produce identical
+    gradients everywhere outside ``indexer.*`` -- including on the input, which
+    stands in for every layer below.
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    torch.manual_seed(0)
+    off = _make_csa_attention(coeff=0.0).to(_DTYPE)
+    on = _make_csa_attention(coeff=1e3).to(_DTYPE)
+    on.load_state_dict(copy.deepcopy(off.state_dict()))
+    off.train()
+    on.train()
+
+    B, S = 1, 8
+    torch.manual_seed(1)
+    base = torch.randn(B, S, off.config.hidden_size, dtype=_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    grads = {}
+    for tag, module in (("off", off), ("on", on)):
+        hidden = base.clone().requires_grad_(True)
+        module(hidden, position_ids).sum().backward()
+        grads[tag] = {
+            "__input__": hidden.grad.clone(),
+            **{
+                name: p.grad.clone()
+                for name, p in module.named_parameters()
+                if p.grad is not None and not name.startswith("indexer.")
+            },
+        }
+
+    assert grads["off"].keys() == grads["on"].keys()
+    for name in grads["off"]:
+        torch.testing.assert_close(
+            grads["on"][name],
+            grads["off"][name],
+            rtol=0,
+            atol=0,
+            msg=lambda s, n=name: f"indexer KL perturbed backbone gradient {n}: {s}",
+        )
+
+    # Guard against a vacuous pass: the loss really was active and large.
+    assert on.last_indexer_distill_loss is not None
+    assert on.last_indexer_distill_loss.abs().item() > 0.0
+    indexer_grad = sum(p.grad.abs().sum().item() for p in on.indexer.parameters() if p.grad is not None)
+    assert indexer_grad > 0.0, "the indexer itself received no gradient"
+
+
+def test_loss_is_reported_to_the_aux_loss_tracker(monkeypatch):
+    """Without this the loss is unobservable once enabled."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    recorded = []
+
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    monkeypatch.setattr(
+        moe_utils,
+        "save_to_aux_losses_tracker",
+        lambda name, loss, layer_number, num_layers, **kw: recorded.append(
+            (name, float(loss), layer_number, num_layers)
+        ),
+    )
+
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=1e-2).to(_DTYPE)
+    attn.train()
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, attn.config.hidden_size, dtype=_DTYPE)
+    attn(hidden, torch.arange(S).unsqueeze(0).expand(B, S)).sum().backward()
+
+    assert recorded, "nothing reached the aux-loss tracker"
+    name, value, layer_number, num_layers = recorded[0]
+    assert name == INDEXER_DISTILL_LOSS_NAME
+    assert value == pytest.approx(attn.last_indexer_distill_loss.item(), rel=1e-6)
+    assert layer_number == attn.layer_number
+    assert num_layers >= attn.config.num_layers
+
+
+def test_nothing_is_reported_when_the_loss_is_off(monkeypatch):
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    recorded = []
+
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    monkeypatch.setattr(
+        moe_utils,
+        "save_to_aux_losses_tracker",
+        lambda *a, **kw: recorded.append(a),
+    )
+
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=0.0).to(_DTYPE)
+    attn.train()
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, attn.config.hidden_size, dtype=_DTYPE)
+    attn(hidden, torch.arange(S).unsqueeze(0).expand(B, S)).sum().backward()
+
+    assert not recorded, "the tracker must stay untouched at coeff 0"
+
+
+def test_layers_without_an_indexer_still_report_a_zero(monkeypatch):
+    """Every V4 layer must report, or the cross-PP reduction diverges.
+
+    ``reduce_aux_losses_tracker_across_ranks`` all-reduces over whatever keys a
+    rank holds, so a key present only on the ranks owning a CSA layer would hang
+    the collective. Non-CSA layers therefore contribute an explicit zero.
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    recorded = []
+
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    monkeypatch.setattr(
+        moe_utils,
+        "save_to_aux_losses_tracker",
+        lambda name, loss, layer_number, num_layers, **kw: recorded.append((name, float(loss))),
+    )
+
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=1e-2).to(_DTYPE)
+    # An HCA layer owns no indexer, but shares the coefficient.
+    attn.compress_ratio = 128
+    attn.indexer = None
+    attn.last_indexer_distill_loss = None
+    attn.train()
+
+    attn._log_indexer_distill_loss(torch.device("cpu"))
+
+    assert recorded == [(INDEXER_DISTILL_LOSS_NAME, 0.0)]
+
+
+@pytest.mark.parametrize("layer_number", [None, 0])
+def test_unnumbered_layers_are_not_reported(monkeypatch, layer_number):
+    """``layer_number`` indexes the tracker, so the sentinel must be rejected.
+
+    ``0`` is the default on a standalone attention module and would otherwise
+    land on ``values[-1]``, corrupting the last layer's entry.
+    """
+    recorded = []
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    monkeypatch.setattr(moe_utils, "save_to_aux_losses_tracker", lambda *a, **kw: recorded.append(a))
+
+    log_indexer_distill_loss(
+        torch.tensor(1.0),
+        layer_number=layer_number,
+        num_layers=8,
+        device=torch.device("cpu"),
+    )
+    assert not recorded
+
+
+def test_eval_does_not_report_a_stale_loss(monkeypatch):
+    """A value from an earlier training step must not leak into eval logs."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    recorded = []
+
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    monkeypatch.setattr(moe_utils, "save_to_aux_losses_tracker", lambda *a, **kw: recorded.append(a))
+
+    attn = _make_csa_attention(coeff=1e-2).to(_DTYPE)
+    attn.last_indexer_distill_loss = torch.tensor(7.0)
+    attn.eval()
+
+    attn._log_indexer_distill_loss(torch.device("cpu"))
+
+    assert not recorded
+
+
+def test_no_indexer_loss_in_eval(monkeypatch):
+    """Eval must not build the aux graph."""
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=1e-2).to(_DTYPE)
+    attn.eval()
+    attn.last_indexer_distill_loss = None
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, attn.config.hidden_size, dtype=_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+    with torch.no_grad():
+        attn(hidden, position_ids)
+
+    assert attn.last_indexer_distill_loss is None
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_rope_hadamard.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_rope_hadamard.py
new file mode 100644
index 000000000..26d373766
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_rope_hadamard.py
@@ -0,0 +1,335 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for the indexer's partial RoPE + Hadamard rotation.
+
+The open-source reference prepares both scoring operands before the ReLU'd dot
+product: partial RoPE at the compressed-branch base (queries at their token
+positions, compressed keys at ``s * compress_ratio``) followed by a normalised
+Hadamard rotation. It does this to the indexer queries directly, and inside the
+indexer's own compressor, which is built with rotation enabled -- unlike the main
+compressed pool, which is not rotated.
+
+Covers the rotation primitive, the two operand paths, and the wiring that keeps
+a CSA layer from silently constructing an unrotated indexer.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+mla_module = pytest.importorskip(
+    "megatron.core.transformer.multi_latent_attention",
+    reason="MLA base module not importable in this environment",
+)
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: E402
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import (  # noqa: E402
+    DualRoPE,
+    apply_interleaved_partial_rope,
+)
+from primus.backends.megatron.core.transformer.hadamard_rotation import (  # noqa: E402
+    rotate_activation,
+)
+from primus.backends.megatron.core.transformer.indexer import Indexer  # noqa: E402
+
+_HIDDEN = 64
+_HEAD_DIM = 16
+_N_HEADS = 4
+_TOPK = 4
+_RATIO = 4
+_ROTARY_DIM = 8
+_THETA = 40000.0
+
+
+def _make_rope() -> DualRoPE:
+    return DualRoPE(
+        rotary_dim=_ROTARY_DIM,
+        rope_theta=10000.0,
+        compress_rope_theta=_THETA,
+        yarn_factor=1.0,
+        original_max_position_embeddings=2048,
+    )
+
+
+def _make_indexer(*, rotated: bool = True) -> Indexer:
+    rope = _make_rope().get_rope(compress_ratio=_RATIO) if rotated else None
+    return Indexer(
+        hidden_size=_HIDDEN,
+        index_head_dim=_HEAD_DIM,
+        index_n_heads=_N_HEADS,
+        index_topk=_TOPK,
+        compress_ratio=_RATIO,
+        rope=rope,
+        rotary_dim=_ROTARY_DIM if rotated else 0,
+    )
+
+
+# ---------------------------------------------------------------------------
+# The rotation primitive
+# ---------------------------------------------------------------------------
+
+
+def test_hadamard_rotation_preserves_inner_products():
+    """It is orthogonal, so scores are unchanged in exact arithmetic.
+
+    This is why applying it to both operands is safe: the value it adds is
+    numerical (spreading channel energy for the low-precision QK), not a change
+    to what the indexer computes.
+    """
+    torch.manual_seed(0)
+    a = torch.randn(5, 32, dtype=torch.float64)
+    b = torch.randn(7, 32, dtype=torch.float64)
+
+    before = a @ b.T
+    after = rotate_activation(a) @ rotate_activation(b).T
+
+    torch.testing.assert_close(after, before, rtol=1e-10, atol=1e-10)
+
+
+def test_hadamard_rotation_is_an_involution():
+    """``H / sqrt(n)`` is symmetric and orthogonal, so applying it twice is identity."""
+    torch.manual_seed(0)
+    x = torch.randn(3, 16, dtype=torch.float64)
+    torch.testing.assert_close(rotate_activation(rotate_activation(x)), x, rtol=1e-10, atol=1e-10)
+
+
+def test_hadamard_rotation_actually_mixes_channels():
+    """Guard the two assertions above from passing on an identity transform."""
+    x = torch.zeros(1, 16, dtype=torch.float64)
+    x[0, 0] = 1.0
+    rotated = rotate_activation(x)
+    assert (rotated.abs() > 0).all(), "a basis vector must spread over every channel"
+
+
+def test_hadamard_rotation_rejects_non_power_of_two():
+    with pytest.raises(ValueError):
+        rotate_activation(torch.randn(2, 12))
+
+
+# ---------------------------------------------------------------------------
+# Operand preparation
+# ---------------------------------------------------------------------------
+
+
+def test_keys_are_rotated_at_compressed_positions():
+    """Compressed key ``s`` must be rotated at original token ``s * ratio``.
+
+    Same convention as the main compressed pool; using the bare block index
+    would put the keys in a different coordinate system from the queries.
+    """
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    hidden = torch.randn(1, 16, _HIDDEN, dtype=torch.float32)
+
+    with torch.no_grad():
+        pooled = indexer.indexer_compressor(hidden)  # [B, P, Hd]
+        actual = indexer._rotate_keys(pooled)
+
+        P = pooled.shape[1]
+        positions = torch.arange(P) * _RATIO
+        cos, sin = indexer.rope(positions)
+        expected = rotate_activation(
+            apply_interleaved_partial_rope(
+                pooled.unsqueeze(2),
+                cos.unsqueeze(0).expand(pooled.shape[0], -1, -1),
+                sin.unsqueeze(0).expand(pooled.shape[0], -1, -1),
+                rotary_dim=_ROTARY_DIM,
+            ).squeeze(2)
+        )
+
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+    # Not a no-op: block indices would give a different answer.
+    with torch.no_grad():
+        cos_b, sin_b = indexer.rope(torch.arange(P))
+        block_indexed = rotate_activation(
+            apply_interleaved_partial_rope(
+                pooled.unsqueeze(2),
+                cos_b.unsqueeze(0).expand(pooled.shape[0], -1, -1),
+                sin_b.unsqueeze(0).expand(pooled.shape[0], -1, -1),
+                rotary_dim=_ROTARY_DIM,
+            ).squeeze(2)
+        )
+    assert not torch.allclose(actual, block_indexed, rtol=1e-4, atol=1e-5)
+
+
+def test_queries_are_rotated_at_token_positions():
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    B, S = 1, 16
+    q_i = torch.randn(B, S, _N_HEADS, _HEAD_DIM, dtype=torch.float32)
+    position_ids = torch.arange(S)
+
+    with torch.no_grad():
+        actual = indexer._rotate_queries(q_i, position_ids)
+        cos, sin = indexer.rope(position_ids)
+        expected = rotate_activation(
+            apply_interleaved_partial_rope(
+                q_i,
+                cos.expand(B, S, -1),
+                sin.expand(B, S, -1),
+                rotary_dim=_ROTARY_DIM,
+            )
+        )
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+
+
+def test_query_rotation_defaults_to_arange_positions():
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    q_i = torch.randn(1, 16, _N_HEADS, _HEAD_DIM, dtype=torch.float32)
+    with torch.no_grad():
+        implicit = indexer._rotate_queries(q_i, None)
+        explicit = indexer._rotate_queries(q_i, torch.arange(16))
+    torch.testing.assert_close(implicit, explicit, rtol=0, atol=0)
+
+
+def test_position_ids_change_the_selection_scores():
+    """Rotating the queries has to depend on where the token actually is."""
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    hidden = torch.randn(1, 16, _HIDDEN, dtype=torch.float32)
+
+    with torch.no_grad():
+        _, scores_a = indexer(hidden, torch.arange(16))
+        _, scores_b = indexer(hidden, torch.arange(16) + 64)
+
+    finite = torch.isfinite(scores_a) & torch.isfinite(scores_b)
+    assert finite.any()
+    assert not torch.allclose(scores_a[finite], scores_b[finite], rtol=1e-4, atol=1e-5)
+
+
+def test_rotation_changes_the_scores_versus_the_unrotated_indexer():
+    """The whole point: adding RoPE + Hadamard is not a no-op on the scores."""
+    torch.manual_seed(0)
+    rotated = _make_indexer(rotated=True)
+    plain = _make_indexer(rotated=False)
+    plain.load_state_dict(rotated.state_dict())
+
+    hidden = torch.randn(1, 16, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        _, scores_rotated = rotated(hidden)
+        _, scores_plain = plain(hidden)
+
+    finite = torch.isfinite(scores_rotated) & torch.isfinite(scores_plain)
+    assert finite.any()
+    assert not torch.allclose(scores_rotated[finite], scores_plain[finite], rtol=1e-4, atol=1e-5)
+
+
+def test_sentinels_and_shapes_are_unchanged_by_the_rotation():
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    hidden = torch.randn(1, 8, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        idxs, scores = indexer(hidden)
+
+    assert idxs.shape == scores.shape == (1, 8, _TOPK)
+    invalid = idxs < 0
+    assert invalid.any()
+    assert torch.isneginf(scores[invalid]).all()
+
+
+# ---------------------------------------------------------------------------
+# Wiring
+# ---------------------------------------------------------------------------
+
+
+def _make_csa_attention() -> DeepseekV4Attention:
+    config = DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=64,
+        num_attention_heads=4,
+        num_query_groups=1,
+        kv_channels=16,
+        qk_pos_emb_head_dim=8,
+        qk_head_dim=8,
+        v_head_dim=16,
+        kv_lora_rank=16,
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sliding_window=0,
+        attn_sink=True,
+        compress_ratios=None,
+        compress_rope_theta=_THETA,
+        use_v4_attention_backend="eager",
+        use_v4_csa_attention_backend="eager",
+        layernorm_epsilon=1e-6,
+        norm_epsilon=1e-6,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    return DeepseekV4Attention(config, rope=rope, compress_ratio=4, submodules=None)
+
+
+def test_csa_layer_always_builds_a_rotated_indexer():
+    """Guard against a silently unrotated indexer in production.
+
+    ``Indexer`` tolerates ``rope=None`` for isolated unit tests, so the layer
+    that builds it is where the contract has to be enforced.
+    """
+    attn = _make_csa_attention()
+    assert attn.indexer is not None
+    assert attn.indexer.rope is not None, "CSA layer built an indexer with no RoPE"
+    assert attn.indexer.rotary_dim == attn.rotary_dim > 0
+
+
+def test_indexer_shares_the_layers_compressed_rope_cache():
+    """It must be the same object, so base / YaRN can never diverge."""
+    attn = _make_csa_attention()
+    assert attn.indexer.rope is attn.rope.get_rope(compress_ratio=4)
+    assert attn.indexer.rope is not attn.rope.main_rope
+
+
+def test_csa_forward_threads_position_ids_into_the_indexer():
+    """Shifting the positions must reach the selector, not just the attention."""
+    torch.manual_seed(0)
+    attn = _make_csa_attention().to(torch.float32)
+    attn.eval()
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, attn.config.hidden_size, dtype=torch.float32)
+
+    captured = []
+    original = attn.indexer.forward
+
+    def _spy(h, position_ids=None):
+        captured.append(position_ids)
+        return original(h, position_ids)
+
+    attn.indexer.forward = _spy
+    try:
+        with torch.no_grad():
+            attn(hidden, torch.arange(S).unsqueeze(0).expand(B, S))
+    finally:
+        attn.indexer.forward = original
+
+    assert captured, "the indexer was never called"
+    assert captured[0] is not None, "position_ids did not reach the indexer"
+    assert captured[0].shape[-1] == S
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_score_scale.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_score_scale.py
new file mode 100644
index 000000000..2c4fc833f
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_score_scale.py
@@ -0,0 +1,178 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for the CSA indexer score temperature.
+
+The index score ``I_{t,s}`` carries a ``sqrt(index_n_heads * index_head_dim)``
+temperature that the open-source reference applies in two places:
+``index_n_heads ** -0.5`` when building the per-head weights, plus the indexer's
+own ``index_head_dim ** -0.5`` on the way into the distillation loss.
+
+The scale is invisible to ``topk`` -- which is why it went missing -- and only
+becomes load-bearing once the scores enter a softmax. Both properties are
+pinned here: the scale must be exactly the reference value, it must not perturb
+the selection, and it must actually change the softmax it feeds.
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from primus.backends.megatron.core.transformer.indexer import Indexer  # noqa: E402
+
+_HIDDEN = 64
+_HEAD_DIM = 16
+_N_HEADS = 4
+_TOPK = 4
+_RATIO = 4
+
+
+def _make_indexer() -> Indexer:
+    return Indexer(
+        hidden_size=_HIDDEN,
+        index_head_dim=_HEAD_DIM,
+        index_n_heads=_N_HEADS,
+        index_topk=_TOPK,
+        compress_ratio=_RATIO,
+    )
+
+
+def test_score_scale_matches_reference_formula():
+    """``score_scale == (index_n_heads * index_head_dim) ** -0.5``."""
+    indexer = _make_indexer()
+    expected = (_N_HEADS**-0.5) * (_HEAD_DIM**-0.5)
+    assert indexer.score_scale == pytest.approx(expected, rel=1e-12)
+    # Guard against the assertion going vacuous on a degenerate config.
+    assert expected < 0.5, "pick a config where the scale is clearly not 1"
+
+
+def test_scale_is_a_pure_constant_on_the_scores():
+    """Scaling only changes the temperature, never the relative scores.
+
+    Runs the same weights twice, once with the scale neutralised, and checks
+    every finite score differs by exactly the constant.
+    """
+    torch.manual_seed(0)
+    scaled = _make_indexer()
+    unscaled = _make_indexer()
+    unscaled.load_state_dict(copy.deepcopy(scaled.state_dict()))
+    unscaled.score_scale = 1.0
+
+    hidden = torch.randn(2, 32, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        idx_scaled, scores_scaled = scaled(hidden)
+        idx_unscaled, scores_unscaled = unscaled(hidden)
+
+    # top-k selection is invariant under a positive constant.
+    assert torch.equal(idx_scaled, idx_unscaled)
+
+    finite = torch.isfinite(scores_scaled)
+    assert finite.any(), "no finite scores to compare"
+    torch.testing.assert_close(
+        scores_scaled[finite],
+        scores_unscaled[finite] * scaled.score_scale,
+        rtol=1e-5,
+        atol=1e-7,
+    )
+
+
+def test_scale_changes_the_softmax_it_feeds():
+    """The distillation loss softmaxes these scores, so the scale must bite.
+
+    Without the scale the distribution is ~90x too sharp at the V4 widths; a
+    smaller config still has to show a clearly less peaked distribution.
+    """
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    hidden = torch.randn(1, 32, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        _, scores = indexer(hidden)
+
+    finite_rows = torch.isfinite(scores).all(dim=-1)
+    assert finite_rows.any(), "need at least one fully-valid row"
+    rows = scores[finite_rows].float()
+
+    peaked = torch.softmax(rows / indexer.score_scale, dim=-1).max(dim=-1).values
+    tempered = torch.softmax(rows, dim=-1).max(dim=-1).values
+
+    assert (tempered <= peaked + 1e-6).all()
+    assert tempered.mean() < peaked.mean(), "the scale must soften the distribution"
+
+
+def test_sentinel_semantics_survive_the_scale():
+    """Masked slots stay ``-inf`` / ``-1``: the scale is applied pre-mask."""
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    # S small enough that early queries have fewer than topk legal entries.
+    hidden = torch.randn(1, 8, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        idxs, scores = indexer(hidden)
+
+    invalid = idxs < 0
+    assert invalid.any(), "expected some early-query sentinels in this shape"
+    assert torch.isneginf(scores[invalid]).all()
+    assert torch.isfinite(scores[~invalid]).all()
+
+
+def test_scale_survives_the_fused_and_unfused_projection_layouts(monkeypatch):
+    """``w_dq_w`` (fused) and ``w_dq`` + ``w_w`` must land on the same scores."""
+    torch.manual_seed(0)
+    monkeypatch.setenv("PRIMUS_INDEXER_FUSE_PROJ", "1")
+    fused = _make_indexer()
+    monkeypatch.setenv("PRIMUS_INDEXER_FUSE_PROJ", "0")
+    split = _make_indexer()
+
+    # ``_load_from_state_dict`` bridges the two layouts.
+    split.load_state_dict(copy.deepcopy(fused.state_dict()))
+
+    hidden = torch.randn(1, 16, _HIDDEN, dtype=torch.float32)
+    with torch.no_grad():
+        idx_fused, scores_fused = fused(hidden)
+        idx_split, scores_split = split(hidden)
+
+    assert torch.equal(idx_fused, idx_split)
+    finite = torch.isfinite(scores_fused)
+    torch.testing.assert_close(scores_fused[finite], scores_split[finite], rtol=1e-5, atol=1e-6)
+
+
+def test_scale_is_not_double_counted_against_a_hand_computed_reference():
+    """Pin the whole scoring chain against an explicit reference expression."""
+    torch.manual_seed(0)
+    indexer = _make_indexer()
+    hidden = torch.randn(1, 16, _HIDDEN, dtype=torch.float32)
+
+    with torch.no_grad():
+        _, scores = indexer(hidden)
+
+        k_icomp = indexer.indexer_compressor(hidden)  # [B, P, Hd]
+        if indexer._fuse_qw_proj:
+            dqw = indexer.w_dq_w(hidden)
+            q_q = dqw[..., : indexer.dq_rank]
+            w_i = dqw[..., indexer.dq_rank :]
+        else:
+            q_q = indexer.w_dq(hidden)
+            w_i = indexer.w_w(hidden)
+        B, S, _ = hidden.shape
+        q_i = indexer.w_iuq(q_q).view(B, S, _N_HEADS, _HEAD_DIM)
+
+        dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp)
+        expected = (torch.relu(dot) * w_i.unsqueeze(-1)).sum(dim=2) * ((_N_HEADS * _HEAD_DIM) ** -0.5)
+        P = k_icomp.shape[1]
+        mask = indexer._causal_mask(S, P, expected.device, expected.dtype)
+        expected = expected + mask.unsqueeze(0)
+
+        topk_eff = min(_TOPK, P)
+        expected_topk = expected.topk(topk_eff, dim=-1).values
+
+    finite = torch.isfinite(expected_topk)
+    torch.testing.assert_close(scores[..., :topk_eff][finite], expected_topk[finite], rtol=1e-5, atol=1e-6)
+    # Sanity: the reference expression is not accidentally scale-free.
+    assert math.isclose(indexer.score_scale, (_N_HEADS * _HEAD_DIM) ** -0.5, rel_tol=1e-12)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_keep_in_fp32.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_keep_in_fp32.py
new file mode 100644
index 000000000..0272897b1
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_keep_in_fp32.py
@@ -0,0 +1,332 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for the keep-in-FP32 contract on ``ape`` and ``attn_sink``.
+
+Both are FP32 in the released DeepSeek-V4 checkpoint, so this mechanism can hold
+them at FP32 through the framework's blanket ``module.bfloat16()``. Covers the
+mechanism in isolation, the two V4 parameters that use it, and the property that
+actually matters when it is on: the FP32 values survive *bit-exactly* rather than
+being round-tripped through BF16.
+
+The mechanism is **off by default** (a second parameter dtype breaks the
+distributed optimizer's single-dtype assumption -- see the module docstring), so
+everything here opts in explicitly via the fixture below. The default-off
+behaviour has its own tests at the bottom.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+import torch.nn as nn  # noqa: E402
+
+mla_module = pytest.importorskip(
+    "megatron.core.transformer.multi_latent_attention",
+    reason="MLA base module not importable in this environment",
+)
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: E402
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.compressor import (  # noqa: E402
+    Compressor,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import DualRoPE  # noqa: E402
+from primus.backends.megatron.core.transformer.keep_in_fp32 import (  # noqa: E402
+    ENABLE_ENV_VAR,
+    KeepInFp32Mixin,
+    is_enabled,
+    is_marked_keep_in_fp32,
+    mark_keep_in_fp32,
+    unmark_keep_in_fp32,
+)
+
+
+@pytest.fixture(autouse=True)
+def _enable_keep_in_fp32(monkeypatch):
+    """Opt in for this module: the mechanism ships disabled."""
+    monkeypatch.setenv(ENABLE_ENV_VAR, "1")
+
+
+# ---------------------------------------------------------------------------
+# The mechanism
+# ---------------------------------------------------------------------------
+
+
+class _Pinned(KeepInFp32Mixin, nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.pinned = nn.Parameter(torch.randn(4, dtype=torch.float32))
+        self.floating = nn.Parameter(torch.randn(4, dtype=torch.float32))
+        mark_keep_in_fp32(self.pinned)
+
+
+def test_bfloat16_skips_the_pinned_parameter():
+    module = _Pinned()
+    module.bfloat16()
+
+    assert module.pinned.dtype == torch.float32
+    assert module.floating.dtype == torch.bfloat16
+
+
+def test_pinned_values_are_not_round_tripped_through_bf16():
+    """The whole point is precision, so restoring must not go via BF16.
+
+    A cast-down-then-up would leave the values quantised to BF16 steps; the
+    saved FP32 copy has to come back bit-exact.
+    """
+    module = _Pinned()
+    # Values with mantissa bits BF16 cannot hold, so a round trip is visible.
+    with torch.no_grad():
+        module.pinned.copy_(torch.tensor([1.0000001, 3.1415927, -2.7182818, 1e-8]))
+    before = module.pinned.detach().clone()
+
+    module.bfloat16()
+
+    assert not torch.equal(before.bfloat16().float(), before), "pick values BF16 actually loses"
+    torch.testing.assert_close(module.pinned, before, rtol=0, atol=0)
+
+
+def test_mark_survives_repeated_conversions():
+    """``_apply`` may replace the Parameter object, dropping the attribute."""
+    module = _Pinned()
+    module.bfloat16()
+    assert is_marked_keep_in_fp32(module.pinned)
+
+    # A second conversion must still be protected.
+    module.half()
+    assert module.pinned.dtype == torch.float32
+    assert module.floating.dtype == torch.float16
+
+
+def test_unmark_lets_the_parameter_follow_the_model():
+    module = _Pinned()
+    unmark_keep_in_fp32(module.pinned)
+    module.bfloat16()
+    assert module.pinned.dtype == torch.bfloat16
+
+
+def test_float_conversion_of_an_unpinned_module_is_unaffected():
+    """Sanity: the mixin is inert when nothing is marked."""
+
+    class _Plain(KeepInFp32Mixin, nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.w = nn.Parameter(torch.randn(3))
+
+    module = _Plain()
+    module.bfloat16()
+    assert module.w.dtype == torch.bfloat16
+
+
+# ---------------------------------------------------------------------------
+# Compressor.ape
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("ratio", [4, 128])
+def test_compressor_ape_stays_fp32(ratio: int):
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=ratio)
+    assert comp.ape.dtype == torch.float32
+
+    comp.bfloat16()
+    assert comp.ape.dtype == torch.float32, "ape must not follow the model dtype"
+    # ...and the projections did convert, so the test is not vacuous.
+    other = comp.wkv_gate.weight if hasattr(comp, "wkv_gate") else comp.wkv.weight
+    assert other.dtype == torch.bfloat16
+
+
+def test_compressor_forward_accepts_bf16_activations_with_fp32_ape():
+    """Mixed dtypes must not break the pooling path.
+
+    The eager body promotes via ``score + ape``; the Triton kernel loads ape
+    with an explicit ``.to(tl.float32)``. Only the eager path is reachable on
+    CPU, which is the one that relies on PyTorch's promotion rules.
+    """
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=4)
+    comp.bfloat16()
+    hidden = torch.randn(1, 8, 32, dtype=torch.bfloat16)
+
+    pooled = comp(hidden)
+
+    assert pooled.shape == (1, 2, 16)
+    assert torch.isfinite(pooled.float()).all()
+
+
+def test_compressor_state_dict_carries_fp32_ape():
+    """Checkpoint parity: the saved tensor dtype has to match the reference."""
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=4)
+    comp.bfloat16()
+    assert comp.state_dict()["ape"].dtype == torch.float32
+
+
+# ---------------------------------------------------------------------------
+# DeepseekV4Attention.attn_sink
+# ---------------------------------------------------------------------------
+
+
+def _make_attention(compress_ratio: int) -> DeepseekV4Attention:
+    config = DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=64,
+        num_attention_heads=4,
+        num_query_groups=1,
+        kv_channels=16,
+        qk_pos_emb_head_dim=8,
+        qk_head_dim=8,
+        v_head_dim=16,
+        kv_lora_rank=16,
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sliding_window=0,
+        attn_sink=True,
+        compress_ratios=None,
+        compress_rope_theta=40000.0,
+        use_v4_attention_backend="eager",
+        use_v4_csa_attention_backend="eager",
+        layernorm_epsilon=1e-6,
+        norm_epsilon=1e-6,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+    )
+    config.index_topk = 2
+    config.index_head_dim = 16
+    config.index_n_heads = 2
+
+    rope = DualRoPE(
+        rotary_dim=config.qk_pos_emb_head_dim,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    return DeepseekV4Attention(config, rope=rope, compress_ratio=compress_ratio, submodules=None)
+
+
+@pytest.mark.parametrize("compress_ratio", [0, 4, 128])
+def test_attn_sink_stays_fp32(compress_ratio: int):
+    attn = _make_attention(compress_ratio)
+    assert attn.attn_sink is not None
+    assert attn.attn_sink.dtype == torch.float32
+
+    attn.bfloat16()
+    assert attn.attn_sink.dtype == torch.float32, "attn_sink must not follow the model dtype"
+    # ...and an ordinary weight did convert, so the assertion is not vacuous.
+    assert attn.linear_q_up_proj.weight.dtype == torch.bfloat16
+
+
+def test_attn_sink_fp32_matches_the_gluon_kernel_contract():
+    """The gluon / flydsl_v1 kernels assert contiguous FP32 sink of shape [H].
+
+    Those kernels are satisfied by the promotion their callers already do, so
+    this only pins that *when* the mechanism is on, the parameter arrives in the
+    shape and layout they expect rather than needing a second cast.
+    """
+    attn = _make_attention(4)
+    attn.bfloat16()
+
+    sink = attn.attn_sink
+    assert sink.dtype == torch.float32
+    assert sink.is_contiguous()
+    assert sink.shape == (attn.num_heads,)
+
+
+def test_nested_compressor_ape_survives_the_parent_conversion():
+    """``_apply`` recurses, so the CSA compressor / indexer must be covered too."""
+    attn = _make_attention(4)
+    attn.bfloat16()
+
+    assert attn.compressor is not None
+    assert attn.compressor.ape.dtype == torch.float32
+    assert attn.indexer is not None
+    assert attn.indexer.indexer_compressor.ape.dtype == torch.float32
+
+
+def test_attention_state_dict_dtypes_match_the_reference_checkpoint():
+    """``ape`` / ``attn_sink`` save as FP32; the weight matrices do not."""
+    attn = _make_attention(4)
+    attn.bfloat16()
+    state = attn.state_dict()
+
+    pinned = {k for k in state if k == "attn_sink" or k.endswith(".ape")}
+    assert pinned, "expected attn_sink and at least one compressor ape"
+    for key in pinned:
+        assert state[key].dtype == torch.float32, f"{key} should be saved as FP32"
+
+    weights = {k for k in state if k.endswith(".weight")}
+    assert weights, "expected some weight matrices to compare against"
+    for key in weights:
+        assert state[key].dtype == torch.bfloat16, f"{key} should follow the model dtype"
+
+
+# ---------------------------------------------------------------------------
+# Default-off behaviour
+# ---------------------------------------------------------------------------
+
+
+def test_disabled_by_default(monkeypatch):
+    """Shipping default: a second parameter dtype breaks the distributed
+    optimizer (``single dtype supported, for now``), and on 4N/PP4 it aborted
+    every GPU with a memory access fault. So the model must come up uniform.
+    """
+    monkeypatch.delenv(ENABLE_ENV_VAR, raising=False)
+    assert not is_enabled()
+
+
+@pytest.mark.parametrize("value", ["0", "", "false", "True", "yes"])
+def test_only_exactly_one_enables(monkeypatch, value: str):
+    """Opt-in is ``"1"`` and nothing else -- a stray value must not half-enable
+    the mechanism, since the failure it causes is a GPU abort rather than an
+    exception.
+    """
+    monkeypatch.setenv(ENABLE_ENV_VAR, value)
+    assert not is_enabled()
+
+
+def test_parameters_follow_the_model_dtype_when_disabled(monkeypatch):
+    """The whole model is one dtype when off -- this is what makes the
+    distributed optimizer's single-dtype assumption hold.
+    """
+    monkeypatch.setenv(ENABLE_ENV_VAR, "0")
+
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=4)
+    comp.bfloat16()
+    assert comp.ape.dtype == torch.bfloat16
+
+    attn = _make_attention(4)
+    attn.bfloat16()
+    assert attn.attn_sink.dtype == torch.bfloat16
+    assert attn.compressor.ape.dtype == torch.bfloat16
+    assert attn.indexer.indexer_compressor.ape.dtype == torch.bfloat16
+
+    dtypes = {p.dtype for p in attn.parameters()}
+    assert dtypes == {torch.bfloat16}, f"expected a single parameter dtype, got {dtypes}"
+
+
+def test_disabled_pooling_forward_still_runs(monkeypatch):
+    """Guard the path that actually matters: ``score.float()`` before the
+    softmax means a BF16 ``ape`` does not change the pooling result's finiteness.
+    """
+    monkeypatch.setenv(ENABLE_ENV_VAR, "0")
+
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=4)
+    comp.bfloat16()
+    pooled = comp(torch.randn(1, 8, 32, dtype=torch.bfloat16))
+
+    assert pooled.shape == (1, 2, 16)
+    assert torch.isfinite(pooled.float()).all()
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py
new file mode 100644
index 000000000..dd7859335
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py
@@ -0,0 +1,504 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for V4 MoE forward (G5, plan-2 §04).
+
+Pins :class:`DeepseekV4MoE` against the HF reference at
+``DeepSeek-V4-Flash/inference/model.py:MoE.forward``. The test runs on
+GPU (CUDA/HIP), fp32, with ``pg_collection=None`` so the V4 MoE uses its
+local-experts path (per-expert dispatch loop, no Megatron dispatcher).
+
+Pass criteria (G5):
+* 1L MoE forward agrees with the inline HF reference to <= 1e-3 max-abs
+  for both routing modes (learned + hash).
+* Output dtype / shape match input.
+* Gradient flows from MoE output back into the gate ``weight`` (router)
+  AND the expert ``w1`` / ``w2`` / ``w3`` Linears.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.clamped_swiglu import ClampedSwiGLUMLP
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+    DeepseekV4HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_moe import (
+    DeepseekV4MoE,
+    DeepseekV4MoESubmodules,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    DeepseekV4LearnedRouter,
+)
+
+
+@pytest.fixture(autouse=True)
+def _v4_moe_on_cuda(monkeypatch):
+    """The V4 MoE routers are GPU-only: both the Triton and eager paths require
+    CUDA/HIP tensors, so default tensor creation to the CUDA device and skip on
+    a CPU-only host.  Force the eager router path (``PRIMUS_V4_ROUTER_TRITON=0``)
+    so this stays an exact MoE-vs-HF reference check.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("DeepSeek-V4 MoE requires a CUDA/HIP device")
+    monkeypatch.setenv("PRIMUS_V4_ROUTER_TRITON", "0")
+    torch.set_default_device("cuda")
+    try:
+        yield
+    finally:
+        torch.set_default_device("cpu")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_moe_config(
+    *,
+    hidden_size: int,
+    moe_intermediate_size: int,
+    shared_expert_intermediate_size: Optional[int],
+    num_experts: int,
+    topk: int,
+    swiglu_limit: float,
+    score_function: str,
+    num_hash_layers: int,
+    vocab_size: int,
+    enable_expert_bias: bool = False,
+    topk_scaling_factor: float = 1.0,
+) -> DeepSeekV4TransformerConfig:
+    """Minimal V4 config for the MoE CPU smoke."""
+    return DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=hidden_size,
+        num_attention_heads=4,
+        num_query_groups=1,
+        kv_channels=32,
+        ffn_hidden_size=hidden_size * 4,  # only consumed by attention path
+        moe_ffn_hidden_size=moe_intermediate_size,
+        moe_intermediate_size=moe_intermediate_size,
+        moe_shared_expert_intermediate_size=shared_expert_intermediate_size,
+        num_moe_experts=num_experts,
+        moe_router_topk=topk,
+        moe_router_score_function=score_function,
+        moe_router_enable_expert_bias=enable_expert_bias,
+        moe_router_topk_scaling_factor=topk_scaling_factor,
+        swiglu_limit=swiglu_limit,
+        num_hash_layers=num_hash_layers,
+        hash_routing_seed=11,
+        vocab_size=vocab_size,
+        padded_vocab_size=vocab_size,
+        layernorm_epsilon=1.0e-6,
+        norm_epsilon=1.0e-6,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        # MLATransformerConfig requirements (not exercised by MoE-only test)
+        qk_pos_emb_head_dim=8,
+        qk_head_dim=24,
+        v_head_dim=32,
+        kv_lora_rank=32,
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+    )
+
+
+def _make_moe(
+    config: DeepSeekV4TransformerConfig,
+    *,
+    layer_idx: int,
+) -> DeepseekV4MoE:
+    """Build a CPU-friendly V4 MoE (no pg_collection -> local-experts path)."""
+    return DeepseekV4MoE(
+        config=config,
+        layer_idx=layer_idx,
+        pg_collection=None,
+        submodules=DeepseekV4MoESubmodules(),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Inline HF reference
+# ---------------------------------------------------------------------------
+
+
+def _hf_gate_forward(
+    *,
+    hidden_flat: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    score_function: str,
+    topk: int,
+    route_scale: float,
+    tid2eid: Optional[torch.Tensor],
+    token_ids: Optional[torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Mirrors ``Gate.forward`` exactly. Returns dense ``(weights, indices)``."""
+    scores = F.linear(hidden_flat.float(), weight.float())
+    if score_function == "softmax":
+        scores = scores.softmax(dim=-1)
+    elif score_function == "sigmoid":
+        scores = scores.sigmoid()
+    else:
+        scores = F.softplus(scores).sqrt()
+    original_scores = scores
+    if bias is not None:
+        scores = scores + bias.float()
+    if tid2eid is not None:
+        assert token_ids is not None
+        indices = tid2eid[token_ids.reshape(-1).long()].long()
+    else:
+        indices = scores.topk(topk, dim=-1).indices
+    weights = original_scores.gather(1, indices)
+    if score_function != "softmax":
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+    weights = weights * route_scale
+    return weights, indices
+
+
+def _hf_expert_forward(
+    expert: ClampedSwiGLUMLP,
+    x: torch.Tensor,
+    weights: Optional[torch.Tensor],
+    *,
+    swiglu_limit: float,
+) -> torch.Tensor:
+    """Mirrors ``Expert.forward`` (pre-mul clamp + post-w1/w3 weight scaling)."""
+    gate = expert.w1(x).float()
+    up = expert.w3(x).float()
+    if swiglu_limit > 0.0:
+        up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
+        gate = gate.clamp(max=swiglu_limit)
+    h = F.silu(gate) * up
+    if weights is not None:
+        h = weights * h
+    return expert.w2(h.to(x.dtype))
+
+
+def _hf_moe_forward(
+    *,
+    moe: DeepseekV4MoE,
+    hidden: torch.Tensor,
+    token_ids: Optional[torch.Tensor],
+    score_function: str,
+    swiglu_limit: float,
+    route_scale: float,
+) -> torch.Tensor:
+    """Inline reference for the MoE forward, sharing weights with ``moe``.
+
+    Builds the (weights, indices) dense form via ``_hf_gate_forward`` and
+    runs the per-expert dispatch loop the way the released reference
+    does, then adds the shared-expert contribution. All tensors share
+    storage with the live :class:`DeepseekV4MoE` so a single set of
+    weights drives both forward paths.
+    """
+    shape = hidden.shape
+    flat_hidden = hidden.reshape(-1, moe.hidden_size)
+
+    # Gate
+    if moe.use_hash_router:
+        router = moe.router
+        assert isinstance(router, DeepseekV4HashRouter)
+        weights, indices = _hf_gate_forward(
+            hidden_flat=flat_hidden,
+            weight=router.weight,
+            bias=None,
+            score_function=score_function,
+            topk=moe.moe_router_topk,
+            route_scale=route_scale,
+            tid2eid=router.tid2eid,
+            token_ids=token_ids,
+        )
+    else:
+        router = moe.learned_router
+        assert isinstance(router, DeepseekV4LearnedRouter)
+        weights, indices = _hf_gate_forward(
+            hidden_flat=flat_hidden,
+            weight=router.weight,
+            bias=router.expert_bias,
+            score_function=score_function,
+            topk=moe.moe_router_topk,
+            route_scale=route_scale,
+            tid2eid=None,
+            token_ids=None,
+        )
+
+    y = torch.zeros_like(flat_hidden, dtype=torch.float32)
+    n_routed = moe.num_routed_experts
+    counts = torch.bincount(indices.flatten(), minlength=n_routed).tolist()
+    assert moe.local_experts is not None
+    for local_i, global_i in enumerate(moe.local_expert_indices):
+        if counts[global_i] == 0:
+            continue
+        expert = moe.local_experts[local_i]
+        idx, top = torch.where(indices == global_i)
+        y[idx] += _hf_expert_forward(
+            expert,
+            flat_hidden[idx],
+            weights[idx, top, None],
+            swiglu_limit=swiglu_limit,
+        ).float()
+
+    if moe.shared_expert is not None:
+        assert isinstance(moe.shared_expert, ClampedSwiGLUMLP)
+        y = (
+            y
+            + _hf_expert_forward(
+                moe.shared_expert,
+                flat_hidden,
+                None,
+                swiglu_limit=swiglu_limit,
+            ).float()
+        )
+
+    return y.type_as(flat_hidden).view(*shape)
+
+
+# ---------------------------------------------------------------------------
+# Construction sanity
+# ---------------------------------------------------------------------------
+
+
+def test_v4_moe_subclasses_megatron_module() -> None:
+    """``DeepseekV4MoE`` is a ``MegatronModule`` (config plumbing parity)."""
+    from megatron.core.transformer.module import MegatronModule
+
+    assert issubclass(DeepseekV4MoE, MegatronModule)
+
+
+def test_v4_moe_cpu_path_builds_local_experts_and_shared_expert() -> None:
+    """When ``pg_collection`` is None the MoE skips the dispatcher and
+    builds local :class:`ClampedSwiGLUMLP` experts plus a single
+    :class:`ClampedSwiGLUMLP` shared expert (matches HF reference)."""
+    config = _make_moe_config(
+        hidden_size=12,
+        moe_intermediate_size=24,
+        shared_expert_intermediate_size=24,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function="sqrtsoftplus",
+        num_hash_layers=0,
+        vocab_size=16,
+    )
+    moe = _make_moe(config, layer_idx=0)
+
+    assert moe.token_dispatcher is None
+    assert moe.grouped_experts is None
+    assert moe.local_experts is not None
+    assert len(moe.local_experts) == 4
+    assert all(isinstance(e, ClampedSwiGLUMLP) for e in moe.local_experts)
+    assert isinstance(moe.shared_expert, ClampedSwiGLUMLP)
+    assert moe.local_expert_indices == [0, 1, 2, 3]
+
+
+def test_v4_moe_set_layer_number_updates_attribute() -> None:
+    """``set_layer_number`` mirrors :class:`BaseMoELayer`."""
+    config = _make_moe_config(
+        hidden_size=12,
+        moe_intermediate_size=24,
+        shared_expert_intermediate_size=None,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function="sqrtsoftplus",
+        num_hash_layers=0,
+        vocab_size=16,
+    )
+    moe = _make_moe(config, layer_idx=0)
+    moe.set_layer_number(7)
+    assert moe.layer_number == 7
+
+
+# ---------------------------------------------------------------------------
+# G5: numerical alignment vs HF reference
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"])
+@pytest.mark.parametrize("with_shared_expert", [True, False])
+def test_v4_moe_learned_layer_matches_hf_reference(score_function: str, with_shared_expert: bool) -> None:
+    """G5: learned router + clamped-SwiGLU MoE forward agrees with
+    the inline HF reference on a 1L toy (CPU fp32)."""
+    torch.manual_seed(2025)
+    H = 16
+    I = 24
+    config = _make_moe_config(
+        hidden_size=H,
+        moe_intermediate_size=I,
+        shared_expert_intermediate_size=I if with_shared_expert else None,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function=score_function,
+        num_hash_layers=0,  # learned router for layer_idx >= 0
+        vocab_size=32,
+    )
+    moe = _make_moe(config, layer_idx=0)
+
+    hidden = torch.randn(2, 5, H, dtype=torch.float32) * 0.5
+    out = moe(hidden, token_ids=None)
+    ref = _hf_moe_forward(
+        moe=moe,
+        hidden=hidden,
+        token_ids=None,
+        score_function=score_function,
+        swiglu_limit=7.0,
+        route_scale=1.0,
+    )
+
+    assert out.shape == hidden.shape
+    assert torch.isfinite(out).all()
+    max_abs = (out - ref).abs().max().item()
+    assert max_abs <= 1.0e-3, f"max-abs vs HF reference = {max_abs}"
+
+
+@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"])
+def test_v4_moe_hash_layer_matches_hf_reference(score_function: str) -> None:
+    """G5: hash router + clamped-SwiGLU MoE forward agrees with
+    the inline HF reference on a 1L toy (CPU fp32)."""
+    torch.manual_seed(7)
+    H = 16
+    I = 24
+    V = 32
+    config = _make_moe_config(
+        hidden_size=H,
+        moe_intermediate_size=I,
+        shared_expert_intermediate_size=I,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function=score_function,
+        num_hash_layers=1,  # hash router for layer_idx == 0
+        vocab_size=V,
+    )
+    moe = _make_moe(config, layer_idx=0)
+
+    hidden = torch.randn(2, 5, H, dtype=torch.float32) * 0.5
+    token_ids = torch.randint(0, V, (2, 5), dtype=torch.long)
+
+    out = moe(hidden, token_ids=token_ids)
+    ref = _hf_moe_forward(
+        moe=moe,
+        hidden=hidden,
+        token_ids=token_ids,
+        score_function=score_function,
+        swiglu_limit=7.0,
+        route_scale=1.0,
+    )
+
+    assert out.shape == hidden.shape
+    assert torch.isfinite(out).all()
+    max_abs = (out - ref).abs().max().item()
+    assert max_abs <= 1.0e-3, f"max-abs vs HF reference = {max_abs}"
+
+
+def test_v4_moe_route_scale_propagates_to_output() -> None:
+    """``moe_router_topk_scaling_factor`` (HF ``route_scale``) is honored."""
+    torch.manual_seed(3)
+    H = 12
+    config = _make_moe_config(
+        hidden_size=H,
+        moe_intermediate_size=24,
+        shared_expert_intermediate_size=None,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function="sqrtsoftplus",
+        num_hash_layers=0,
+        vocab_size=16,
+        topk_scaling_factor=2.5,
+    )
+    moe = _make_moe(config, layer_idx=0)
+    hidden = torch.randn(1, 3, H, dtype=torch.float32) * 0.5
+
+    out = moe(hidden, token_ids=None)
+    ref = _hf_moe_forward(
+        moe=moe,
+        hidden=hidden,
+        token_ids=None,
+        score_function="sqrtsoftplus",
+        swiglu_limit=7.0,
+        route_scale=2.5,
+    )
+    assert (out - ref).abs().max().item() <= 1.0e-3
+
+
+def test_v4_moe_gradient_flows_to_router_and_experts() -> None:
+    """G5 follow-up: backward pass populates grads on ``router.weight``
+    and on at least one expert's ``w1`` / ``w2`` / ``w3`` weight."""
+    torch.manual_seed(13)
+    H = 12
+    config = _make_moe_config(
+        hidden_size=H,
+        moe_intermediate_size=24,
+        shared_expert_intermediate_size=24,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function="sqrtsoftplus",
+        num_hash_layers=0,
+        vocab_size=16,
+    )
+    moe = _make_moe(config, layer_idx=0)
+    hidden = torch.randn(2, 4, H, dtype=torch.float32, requires_grad=False) * 0.5
+
+    out = moe(hidden, token_ids=None)
+    out.sum().backward()
+
+    assert moe.learned_router is not None
+    assert moe.learned_router.weight.grad is not None
+    assert torch.isfinite(moe.learned_router.weight.grad).all()
+
+    # At least one routed expert must have received gradient. Some experts
+    # may be unselected in a small toy run; assert >= 1 has a non-zero grad.
+    any_expert_grad = False
+    assert moe.local_experts is not None
+    for expert in moe.local_experts:
+        if (
+            expert.w1.weight.grad is not None
+            and expert.w1.weight.grad.abs().sum().item() > 0.0
+            and expert.w2.weight.grad is not None
+            and expert.w3.weight.grad is not None
+        ):
+            any_expert_grad = True
+            break
+    assert any_expert_grad, "expected at least one routed expert to receive gradient"
+
+    # Shared expert is always-on so it always sees gradient.
+    assert moe.shared_expert is not None
+    assert moe.shared_expert.w1.weight.grad is not None
+    assert moe.shared_expert.w1.weight.grad.abs().sum().item() > 0.0
+
+
+def test_v4_moe_hash_layer_requires_token_ids() -> None:
+    """Hash-routed layers raise a clear error if ``token_ids`` is missing."""
+    config = _make_moe_config(
+        hidden_size=12,
+        moe_intermediate_size=24,
+        shared_expert_intermediate_size=None,
+        num_experts=4,
+        topk=2,
+        swiglu_limit=7.0,
+        score_function="sqrtsoftplus",
+        num_hash_layers=1,
+        vocab_size=16,
+    )
+    moe = _make_moe(config, layer_idx=0)
+    hidden = torch.randn(1, 3, 12, dtype=torch.float32)
+    with pytest.raises(ValueError, match="token_ids is required"):
+        moe(hidden, token_ids=None)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py
new file mode 100644
index 000000000..fff550143
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py
@@ -0,0 +1,361 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for V4 MTP integration with upstream
+:class:`MultiTokenPredictionBlock` (P16) plus plan-2 P17 cleanup
+guarantees.
+
+What this file covers (CPU-friendly):
+
+* :func:`get_v4_mtp_block_spec` returns a well-formed ``ModuleSpec``:
+  the outer module is :class:`MultiTokenPredictionBlock`, its
+  ``submodules.layer_specs`` is a list of length ``mtp_num_layers``,
+  and each entry is a :class:`MultiTokenPredictionLayer` ``ModuleSpec``
+  whose ``mtp_model_layer`` is the V4 hybrid-layer spec we passed in.
+* Per-MTP-layer submodules pull V4 RMSNorm / column-parallel linear
+  from the V4 spec provider — the helper does not silently fall back
+  to TE / vanilla impls.
+* The V4 hybrid layer's ``forward`` returns the upstream-compatible
+  ``(hidden_states, None)`` tuple so it can plug into
+  :meth:`MultiTokenPredictionLayer._proj_and_transformer_layer` which
+  unpacks ``hidden_states, _ = self.mtp_model_layer(...)``.
+* The V4 attention spec advertises ``attn_mask_type`` so upstream MTP
+  validation passes (V4 manages its own SWA / sink mask internally,
+  but the field is required by the upstream pre-build assertion).
+* The legacy primus-owned ``DeepseekV4MTPBlock`` is **gone** in plan-2
+  P17 — its module is no longer importable, the
+  ``v4_use_custom_mtp_block`` config flag is removed, and
+  ``DeepseekV4Model.__init__`` no longer references either of them.
+* :class:`DeepseekV4Model.forward` wires :func:`process_mtp_loss` and
+  :class:`MultiTokenPredictionBlock` (AST audit; full distributed run
+  is gated G7 in P19).
+
+The full ``mtp_num_layers=0`` vs ``mtp_num_layers=1`` main-LM loss
+invariance gate (G7) requires distributed init and is tracked into
+P19 distributed re-validation.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib
+import inspect
+from dataclasses import is_dataclass
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlock,
+    MultiTokenPredictionBlockSubmodules,
+    MultiTokenPredictionLayer,
+    MultiTokenPredictionLayerSubmodules,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules
+
+from primus.backends.megatron.core.models.deepseek_v4 import (
+    DeepseekV4HybridLayer,
+    DeepseekV4HybridLayerSubmodules,
+    DeepseekV4TransformerBlock,
+    DeepseekV4TransformerBlockSubmodules,
+    get_v4_mtp_block_spec,
+)
+from primus.backends.megatron.core.models.deepseek_v4.build_context import (
+    resolve_v4_provider,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_layer import (
+    DeepseekV4MTPLayer,
+)
+
+_REPO_ROOT = Path(__file__).resolve().parents[5]
+_MODEL_PATH = _REPO_ROOT / ("primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py")
+_SPECS_PATH = _REPO_ROOT / ("primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py")
+
+
+# ---------------------------------------------------------------------------
+# Minimal V4 config / layer-spec fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_v4_config(*, mtp_num_layers: int = 1):
+    """Tiny V4 config with the fields the MTP spec helper reads."""
+    cfg = MagicMock()
+    cfg.hidden_size = 64
+    cfg.num_layers = 2
+    cfg.mtp_num_layers = mtp_num_layers
+    cfg.hc_mult = 1
+    cfg.qk_pos_emb_head_dim = 16
+    cfg.attn_sliding_window = 128
+    cfg.norm_epsilon = 1e-6
+    cfg.layernorm_epsilon = 1e-6
+    cfg.use_mup = False
+    cfg.tensor_model_parallel_size = 1
+    cfg.sequence_parallel = False
+    cfg.bf16 = False
+    cfg.fp16 = False
+    return cfg
+
+
+def _placeholder_layer_spec() -> ModuleSpec:
+    """Stand-in for a V4 hybrid-layer spec; the helper just threads it through."""
+    return ModuleSpec(
+        module=DeepseekV4HybridLayer,
+        params={"layer_idx": 0, "compress_ratio": 0},
+        submodules=DeepseekV4HybridLayerSubmodules(),
+    )
+
+
+# ---------------------------------------------------------------------------
+# get_v4_mtp_block_spec — structural assertions
+# ---------------------------------------------------------------------------
+
+
+def test_helper_returns_multi_token_prediction_block_spec() -> None:
+    cfg = _make_v4_config(mtp_num_layers=1)
+    spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec())
+    assert isinstance(spec, ModuleSpec)
+    assert spec.module is MultiTokenPredictionBlock
+    assert isinstance(spec.submodules, MultiTokenPredictionBlockSubmodules)
+
+
+@pytest.mark.parametrize("mtp_num_layers", [1, 2, 3])
+def test_helper_emits_one_layer_spec_per_depth(mtp_num_layers: int) -> None:
+    cfg = _make_v4_config(mtp_num_layers=mtp_num_layers)
+    spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec())
+    layer_specs = spec.submodules.layer_specs
+    assert isinstance(layer_specs, list)
+    assert len(layer_specs) == mtp_num_layers
+    for layer_spec in layer_specs:
+        assert isinstance(layer_spec, ModuleSpec)
+        # V4 uses its own mHC-aware MTP layer (per-depth HyperHead), not the
+        # vanilla upstream MultiTokenPredictionLayer.
+        assert layer_spec.module is DeepseekV4MTPLayer
+        assert issubclass(DeepseekV4MTPLayer, MultiTokenPredictionLayer)
+
+
+def test_helper_threads_v4_inner_layer_unchanged() -> None:
+    cfg = _make_v4_config(mtp_num_layers=2)
+    inner = _placeholder_layer_spec()
+    spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=inner)
+    for layer_spec in spec.submodules.layer_specs:
+        sub: MultiTokenPredictionLayerSubmodules = layer_spec.submodules
+        assert sub.mtp_model_layer is inner, (
+            "MTP helper must thread the V4 hybrid layer spec through unchanged "
+            "so MTP depths share HC / hash-routing / clamped-SwiGLU with the main decoder."
+        )
+
+
+def test_helper_pulls_norm_and_linear_from_v4_provider() -> None:
+    cfg = _make_v4_config(mtp_num_layers=1)
+    # Resolve the provider through the same singleton helper the MTP spec
+    # builder uses (``resolve_v4_provider`` caches one provider on the config),
+    # so the identity comparison is against the exact provider that wired the
+    # spec — not a second, independently-constructed provider instance.
+    provider = resolve_v4_provider(cfg)
+    spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec())
+    sub = spec.submodules.layer_specs[0].submodules
+    expected_norm = provider.v4_norm_module()
+    expected_col = provider.column_parallel_linear()
+    assert sub.enorm is expected_norm
+    assert sub.hnorm is expected_norm
+    assert sub.layer_norm is expected_norm
+    assert sub.eh_proj is expected_col
+
+
+def test_helper_extracts_last_hybrid_layer_from_block_spec() -> None:
+    """When handed the decoder *block* spec (as DeepseekV4Model does), the
+    MTP helper must extract a single hybrid-layer spec (the last one) as the
+    MTP inner layer — not thread the whole block spec through (which would
+    trip MultiTokenPredictionLayer's TransformerLayerSubmodules validation)."""
+    cfg = _make_v4_config(mtp_num_layers=1)
+    first = _placeholder_layer_spec()
+    last = _placeholder_layer_spec()
+    block_spec = ModuleSpec(
+        module=DeepseekV4TransformerBlock,
+        submodules=DeepseekV4TransformerBlockSubmodules(layer_specs=[first, last]),
+    )
+    spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=block_spec)
+    sub = spec.submodules.layer_specs[0].submodules
+    assert sub.mtp_model_layer is last, (
+        "MTP inner layer must be the last hybrid-layer spec extracted from the "
+        "decoder block spec (mirrors upstream GPT spec.layer_specs[-1])."
+    )
+
+
+def test_helper_rejects_zero_mtp_num_layers() -> None:
+    cfg = _make_v4_config(mtp_num_layers=0)
+    with pytest.raises(ValueError, match="mtp_num_layers >= 1"):
+        get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec())
+
+
+# ---------------------------------------------------------------------------
+# DeepseekV4HybridLayer — upstream-compatible tuple return
+# ---------------------------------------------------------------------------
+
+
+def test_layer_submodules_extends_transformer_layer_submodules() -> None:
+    """Required so MultiTokenPredictionLayer.__init__'s submodules
+    isinstance check picks up the GPT path (not Mamba)."""
+    assert is_dataclass(DeepseekV4HybridLayerSubmodules)
+    assert issubclass(DeepseekV4HybridLayerSubmodules, TransformerLayerSubmodules)
+
+
+def test_layer_forward_signature_returns_tuple() -> None:
+    """``forward`` is annotated as a tuple-returning callable so
+    ``MultiTokenPredictionLayer._proj_and_transformer_layer`` can unpack
+    ``hidden_states, _ = self.mtp_model_layer(...)`` without error."""
+    src = inspect.getsource(DeepseekV4HybridLayer.forward)
+    assert "return x, None" in src, (
+        "DeepseekV4HybridLayer.forward must return (hidden_states, None) for "
+        "upstream MultiTokenPredictionLayer compatibility."
+    )
+
+
+# ---------------------------------------------------------------------------
+# V4 attention spec — declares an attn_mask_type for upstream MTP validation
+# ---------------------------------------------------------------------------
+
+
+def test_attention_spec_declares_supported_attn_mask_type() -> None:
+    src = _SPECS_PATH.read_text()
+    assert "AttnMaskType.causal" in src, (
+        "V4 attention spec must declare attn_mask_type=AttnMaskType.causal "
+        "(in spec.params) so MultiTokenPredictionLayer's pre-build "
+        "validation accepts the V4 inner layer."
+    )
+    # Sanity: import works and value resolves.
+    assert AttnMaskType.causal.name == "causal"
+
+
+# ---------------------------------------------------------------------------
+# Legacy V4 MTP block — retired in plan-2 P17 (G14 dead-code audit)
+# ---------------------------------------------------------------------------
+
+
+def test_legacy_mtp_block_module_is_gone() -> None:
+    """plan-2 P17 deletes :mod:`...deepseek_v4_mtp` outright.
+
+    Importing the module must raise :class:`ImportError`; the symbol
+    must not be re-exported from the package surface either.
+    """
+    with pytest.raises(ImportError):
+        importlib.import_module("primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp")
+
+    # Package surface — the legacy class must not leak back via __init__.py.
+    pkg = importlib.import_module("primus.backends.megatron.core.models.deepseek_v4")
+    assert not hasattr(pkg, "DeepseekV4MTPBlock"), (
+        "DeepseekV4MTPBlock must be retired in plan-2 P17; the package " "must not re-export it."
+    )
+    assert "DeepseekV4MTPBlock" not in getattr(pkg, "__all__", []), (
+        "DeepseekV4MTPBlock must be removed from the package __all__ " "(plan-2 P17 dead-code audit)."
+    )
+
+
+# ---------------------------------------------------------------------------
+# DeepseekV4Model.forward — process_mtp_loss + MTP block wired (AST audit)
+# ---------------------------------------------------------------------------
+
+
+def test_model_forward_calls_process_mtp_loss() -> None:
+    src = _MODEL_PATH.read_text()
+    tree = ast.parse(src)
+    has_call = False
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            if node.func.id == "process_mtp_loss":
+                has_call = True
+                break
+    assert has_call, (
+        "DeepseekV4Model.forward must call process_mtp_loss to wire MTP "
+        "into the LM-loss pipeline (plan-2 P16)."
+    )
+
+
+def test_model_imports_upstream_mtp_machinery() -> None:
+    src = _MODEL_PATH.read_text()
+    for needle in (
+        "from megatron.core.transformer.multi_token_prediction import",
+        "MultiTokenPredictionBlock",
+        "process_mtp_loss",
+        "mtp_on_this_rank",
+        "get_v4_mtp_block_spec",
+    ):
+        assert needle in src, f"DeepseekV4Model must import {needle!r} (P16)."
+
+
+def test_model_init_routes_through_v4_mtp_spec_helper() -> None:
+    src = _MODEL_PATH.read_text()
+    assert "get_v4_mtp_block_spec(" in src, (
+        "DeepseekV4Model.__init__ must use get_v4_mtp_block_spec to build "
+        "self.mtp; the spec is then handed to MultiTokenPredictionBlock."
+    )
+
+
+def test_model_no_longer_references_legacy_mtp_block() -> None:
+    """Plan-2 P17 retires the legacy MTP block AND the
+    ``v4_use_custom_mtp_block`` config flag; ``DeepseekV4Model``'s
+    source must not reference either of them anywhere outside a
+    historical comment that explicitly says the field is gone.
+    """
+    src = _MODEL_PATH.read_text()
+    tree = ast.parse(src)
+
+    bad_attr = []
+    bad_name = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Attribute) and node.attr == "v4_use_custom_mtp_block":
+            bad_attr.append(getattr(node, "lineno", -1))
+        if isinstance(node, ast.Name) and node.id == "DeepseekV4MTPBlock":
+            bad_name.append(getattr(node, "lineno", -1))
+        if isinstance(node, (ast.ImportFrom, ast.Import)):
+            module = getattr(node, "module", None) or ""
+            if module.endswith("deepseek_v4_mtp"):
+                bad_attr.append(getattr(node, "lineno", -1))
+
+    assert not bad_attr, (
+        "DeepseekV4Model must not reference v4_use_custom_mtp_block / "
+        "deepseek_v4_mtp after plan-2 P17 (lines: %s)." % bad_attr
+    )
+    assert not bad_name, (
+        "DeepseekV4Model must not call/reference DeepseekV4MTPBlock after "
+        "plan-2 P17 (lines: %s)." % bad_name
+    )
+
+
+def test_v4_config_no_longer_carries_legacy_mtp_fields() -> None:
+    """Plan-2 P17 dead-code audit: ``v4_use_custom_mtp_block`` and
+    ``mtp_compress_ratios`` are removed from the V4 config dataclass."""
+    from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (
+        DeepSeekV4TransformerConfig,
+    )
+
+    fields = {f.name for f in DeepSeekV4TransformerConfig.__dataclass_fields__.values()}
+    assert "v4_use_custom_mtp_block" not in fields, (
+        "v4_use_custom_mtp_block must be removed from DeepSeekV4TransformerConfig "
+        "(plan-2 P17 retired the legacy MTP block path)."
+    )
+    assert "mtp_compress_ratios" not in fields, (
+        "mtp_compress_ratios was only consumed by the legacy MTP block; "
+        "remove it alongside the block (plan-2 P17)."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Tiny CPU-only smoke: model __init__ in a no-MTP / no-distributed config
+# does not crash and leaves self.mtp = None
+# ---------------------------------------------------------------------------
+
+
+def test_model_init_no_mtp_path_does_not_build_mtp() -> None:
+    """When ``mtp_num_layers == 0`` the model must not construct an MTP
+    block (regardless of distributed state)."""
+    # We don't build a real model here (it requires Megatron's full init
+    # machinery); we just confirm the __init__ source has the
+    # mtp_num_layers > 0 guard so the no-MTP path stays inert.
+    src = _MODEL_PATH.read_text()
+    assert "mtp_num_layers > 0" in src
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_parity_suite.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_parity_suite.py
new file mode 100644
index 000000000..f2a98cc61
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_parity_suite.py
@@ -0,0 +1,390 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""DeepSeek-V4 CSA parity suite, split into two kinds of assertion.
+
+The two failure modes these guard against are genuinely different, and mixing
+them is how a wiring bug hides behind green maths:
+
+**Group 1 -- numerical parity.** Rebuild the CSA forward from the module's own
+parameters using plain matmuls, independent of the production helpers, and
+compare element-wise. This is where a wrong RoPE position, a missing score
+scale, a dropped Hadamard rotation or an absent output de-rotation shows up.
+Two steps are reused rather than rewritten: the joint local-SWA + sparse softmax
+(``eager_v4_csa_attention``) and the compressor pooling. Both are shared
+contracts with their own dedicated tests -- every CSA backend is pinned against
+the former, and the pooling has its own parity coverage -- so duplicating them
+here would test a copy instead of the contract. Everything the recent fixes
+touched (RoPE positions, the indexer's operand preparation, the output
+de-rotation, the softmax temperature) is rebuilt from scratch.
+
+**Group 2 -- gradient topology.** Assert which parameters each objective may and
+may not reach. Scale and detach bugs are invisible to Group 1 -- the forward is
+bit-identical either way -- and equally invisible to a loss-value assertion. The
+only thing that catches them is asking where the gradient went.
+
+Everything here is CPU-friendly: the attention is built with the ``eager``
+backends and no process group.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+mla_module = pytest.importorskip(
+    "megatron.core.transformer.multi_latent_attention",
+    reason="MLA base module not importable in this environment",
+)
+
+import torch.nn.functional as F  # noqa: E402
+
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import (  # noqa: E402
+    DeepSeekV4TransformerConfig,
+)
+from primus.backends.megatron.core.transformer.deepseek_v4_attention import (  # noqa: E402
+    DeepseekV4Attention,
+)
+from primus.backends.megatron.core.transformer.dual_rope import (  # noqa: E402
+    DualRoPE,
+    apply_interleaved_partial_rope,
+)
+from primus.backends.megatron.core.transformer.hadamard_rotation import (  # noqa: E402
+    rotate_activation,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import (  # noqa: E402
+    eager_v4_csa_attention,
+)
+
+_DTYPE = torch.float32
+
+_HIDDEN = 64
+_HEADS = 4
+_HEAD_DIM = 16
+_ROTARY_DIM = 8
+_RATIO = 4
+_INDEX_TOPK = 2
+_INDEX_HEAD_DIM = 16
+_INDEX_N_HEADS = 2
+
+
+def _make_csa_attention(coeff: float = 0.0) -> DeepseekV4Attention:
+    config = DeepSeekV4TransformerConfig(
+        num_layers=1,
+        hidden_size=_HIDDEN,
+        num_attention_heads=_HEADS,
+        num_query_groups=1,
+        kv_channels=_HEAD_DIM,
+        qk_pos_emb_head_dim=_ROTARY_DIM,
+        qk_head_dim=_ROTARY_DIM,
+        v_head_dim=_HEAD_DIM,
+        kv_lora_rank=_HEAD_DIM,
+        rope_type="rope",
+        rotary_base=10000.0,
+        rotary_scaling_factor=1.0,
+        rotary_percent=1.0,
+        original_max_position_embeddings=2048,
+        q_lora_rank=32,
+        o_groups=2,
+        o_lora_rank=8,
+        attn_sliding_window=0,
+        attn_sink=True,
+        compress_ratios=None,
+        compress_rope_theta=40000.0,
+        use_v4_attention_backend="eager",
+        use_v4_csa_attention_backend="eager",
+        layernorm_epsilon=1e-6,
+        norm_epsilon=1e-6,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        v4_indexer_distill_loss_coeff=coeff,
+    )
+    config.index_topk = _INDEX_TOPK
+    config.index_head_dim = _INDEX_HEAD_DIM
+    config.index_n_heads = _INDEX_N_HEADS
+
+    rope = DualRoPE(
+        rotary_dim=_ROTARY_DIM,
+        rope_theta=config.rotary_base,
+        compress_rope_theta=config.compress_rope_theta,
+        yarn_factor=1.0,
+        original_max_position_embeddings=config.original_max_position_embeddings,
+    )
+    attn = DeepseekV4Attention(config, rope=rope, compress_ratio=_RATIO, submodules=None, layer_number=1)
+    return attn.to(_DTYPE)
+
+
+def _rms_per_head(x: torch.Tensor, eps: float) -> torch.Tensor:
+    x32 = x.float()
+    return (x32 * torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps)).to(x.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Group 1: numerical parity
+# ---------------------------------------------------------------------------
+
+
+def _reference_csa_forward(attn: DeepseekV4Attention, hidden, position_ids):
+    """Independent CSA reference assembled from the module's parameters."""
+    B, S, _ = hidden.shape
+    H, Dh, rd, eps = attn.num_heads, attn.head_dim, attn.rotary_dim, attn.norm_eps
+    compress_rope = attn.rope.get_rope(compress_ratio=_RATIO)
+
+    # --- Q / KV branches -------------------------------------------------
+    q_c = hidden @ attn.linear_q_down_proj.weight.t()
+    q_c = _rms_per_head(q_c, eps) * attn.q_layernorm.weight
+    q = _rms_per_head((q_c @ attn.linear_q_up_proj.weight.t()).view(B, S, H, Dh), eps)
+
+    kv = hidden @ attn.linear_kv.weight.t()
+    kv = (_rms_per_head(kv, eps) * attn.kv_layernorm.weight).view(B, S, 1, Dh)
+
+    q = attn.rope.apply_rope(q, position_ids=position_ids, compress_ratio=_RATIO)
+    kv = attn.rope.apply_rope(kv, position_ids=position_ids, compress_ratio=_RATIO)
+
+    # --- compressed pool, rotated at s * ratio ---------------------------
+    pool = attn.compressor(hidden)  # [B, P, Dh]
+    P = pool.shape[1]
+    comp_pos = torch.arange(P, device=hidden.device) * _RATIO
+    cos, sin = compress_rope(comp_pos)
+    pool_rot = apply_interleaved_partial_rope(
+        pool.unsqueeze(2),
+        cos.unsqueeze(0).expand(B, -1, -1),
+        sin.unsqueeze(0).expand(B, -1, -1),
+        rotary_dim=rd,
+    ).squeeze(2)
+
+    # --- indexer: RoPE + Hadamard on both operands, then the scaled score -
+    idx = attn.indexer
+    k_icomp = idx.indexer_compressor(hidden)  # [B, P, Hd_i]
+    k_icomp = rotate_activation(
+        apply_interleaved_partial_rope(
+            k_icomp.unsqueeze(2),
+            cos.unsqueeze(0).expand(B, -1, -1),
+            sin.unsqueeze(0).expand(B, -1, -1),
+            rotary_dim=rd,
+        ).squeeze(2)
+    )
+
+    if idx._fuse_qw_proj:
+        dqw = idx.w_dq_w(hidden)
+        q_q, w_i = dqw[..., : idx.dq_rank], dqw[..., idx.dq_rank :]
+    else:
+        q_q, w_i = idx.w_dq(hidden), idx.w_w(hidden)
+    q_i = idx.w_iuq(q_q).view(B, S, _INDEX_N_HEADS, _INDEX_HEAD_DIM)
+    icos, isin = compress_rope(position_ids)
+    if tuple(icos.shape[:-1]) != (B, S):
+        icos, isin = icos.expand(B, S, -1), isin.expand(B, S, -1)
+    q_i = rotate_activation(apply_interleaved_partial_rope(q_i, icos, isin, rotary_dim=rd))
+    w_i = w_i * ((_INDEX_N_HEADS**-0.5) * (_INDEX_HEAD_DIM**-0.5))
+
+    scores = (F.relu(torch.einsum("bshd,bpd->bshp", q_i, k_icomp)) * w_i.unsqueeze(-1)).sum(dim=2)
+    t_idx = torch.arange(S, device=hidden.device).unsqueeze(1)
+    s_end = (torch.arange(P, device=hidden.device).unsqueeze(0) + 1) * _RATIO - 1
+    scores = scores + torch.where(s_end <= t_idx, 0.0, float("-inf")).to(scores.dtype)
+
+    topk_eff = min(_INDEX_TOPK, P)
+    topk_scores, topk_idxs = scores.topk(topk_eff, dim=-1)
+    topk_idxs = torch.where(torch.isneginf(topk_scores), -1, topk_idxs)
+
+    # --- joint softmax over local SWA + the selected compressed entries --
+    gathered = pool_rot[torch.arange(B).view(B, 1, 1), topk_idxs.clamp_min(0)]
+    sparse_mask = torch.where(topk_idxs >= 0, 0.0, float("-inf")).to(hidden.dtype)
+    out_bh = eager_v4_csa_attention(
+        q.transpose(1, 2),
+        kv.expand(B, S, H, Dh).transpose(1, 2),
+        kv.expand(B, S, H, Dh).transpose(1, 2),
+        gathered,
+        sink=attn.attn_sink,
+        swa_window=int(attn.attn_sliding_window),
+        sparse_mask=sparse_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=1.0 / math.sqrt(Dh),
+    )
+    out = out_bh.transpose(1, 2).contiguous()
+
+    # --- de-rotate at -t, then the grouped O projection ------------------
+    dcos, dsin = compress_rope(position_ids)
+    if tuple(dcos.shape[:-1]) != (B, S):
+        dcos, dsin = dcos.expand(B, S, -1), dsin.expand(B, S, -1)
+    out = apply_interleaved_partial_rope(out, dcos, -dsin, rotary_dim=rd)
+
+    G, r = attn.o_groups, attn.o_lora_rank
+    o = torch.einsum(
+        "bsgd,grd->bsgr",
+        out.reshape(B, S, G, (H * Dh) // G),
+        attn.linear_o_a.weight.view(G, r, (H * Dh) // G),
+    )
+    return o.flatten(2) @ attn.linear_o_b.weight.t()
+
+
+def test_csa_forward_matches_the_independent_reference():
+    """End-to-end CSA parity against a reference built from scratch."""
+    torch.manual_seed(0)
+    attn = _make_csa_attention()
+    attn.eval()
+
+    B, S = 2, 16
+    hidden = torch.randn(B, S, _HIDDEN, dtype=_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    with torch.no_grad():
+        ours = attn(hidden, position_ids)
+        ref = _reference_csa_forward(attn, hidden, position_ids)
+
+    torch.testing.assert_close(ours, ref, rtol=1e-4, atol=1e-5)
+
+
+def test_output_actually_depends_on_the_indexer_selection():
+    """Guard the parity test from passing with the indexer bypassed.
+
+    Also marks the boundary of what parity can cover: the indexer's score scale
+    is a positive constant, so it cannot move ``topk`` and therefore cannot show
+    up in the output at all -- which is precisely why it went missing. Pinning
+    that value needs a direct assertion (``test_v4_indexer_score_scale.py``),
+    not this test.
+    """
+    torch.manual_seed(0)
+    attn = _make_csa_attention()
+    attn.eval()
+
+    B, S = 1, 16
+    hidden = torch.randn(B, S, _HIDDEN, dtype=_DTYPE)
+    position_ids = torch.arange(S)
+
+    with torch.no_grad():
+        base = _reference_csa_forward(attn, hidden, position_ids)
+        # Flip one indexer head's weight so the selection genuinely differs.
+        attn.indexer.w_dq_w.weight[attn.indexer.dq_rank].mul_(-3.0)
+        perturbed = _reference_csa_forward(attn, hidden, position_ids)
+
+    assert not torch.allclose(base, perturbed, rtol=1e-3, atol=1e-4)
+
+
+@pytest.mark.parametrize("shift", [0, 7])
+def test_parity_holds_at_shifted_positions(shift: int):
+    """Positions enter in three places; a shift must move all of them together.
+
+    Q / KV rotation, the compressed pool's ``s * ratio``, and the output
+    de-rotation all read the position axis. A parity test only at
+    ``arange(S)`` would miss an off-by-origin in any of them.
+    """
+    torch.manual_seed(1)
+    attn = _make_csa_attention()
+    attn.eval()
+
+    B, S = 1, 16
+    hidden = torch.randn(B, S, _HIDDEN, dtype=_DTYPE)
+    position_ids = torch.arange(S) + shift
+
+    with torch.no_grad():
+        ours = attn(hidden, position_ids)
+        ref = _reference_csa_forward(attn, hidden, position_ids)
+
+    torch.testing.assert_close(ours, ref, rtol=1e-4, atol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# Group 2: gradient topology
+# ---------------------------------------------------------------------------
+
+
+def _backbone_names(attn: DeepseekV4Attention) -> set:
+    return {n for n, _ in attn.named_parameters() if not n.startswith("indexer.")}
+
+
+def test_main_loss_reaches_the_backbone_but_not_the_frozen_indexer(monkeypatch):
+    """With the distillation loss off, the indexer must be unreachable.
+
+    ``topk`` is not differentiable and the scores are discarded, so a gradient
+    arriving at the indexer here would mean something else is consuming them.
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    torch.manual_seed(0)
+    attn = _make_csa_attention(coeff=0.0)
+    attn.train()
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, _HIDDEN, dtype=_DTYPE)
+    attn(hidden, torch.arange(S).unsqueeze(0).expand(B, S)).sum().backward()
+
+    got_grad = {n for n, p in attn.named_parameters() if p.grad is not None}
+    assert got_grad & _backbone_names(attn), "the main loss did not reach the backbone"
+    assert not any(n.startswith("indexer.") for n in got_grad)
+
+
+def test_distillation_loss_reaches_only_the_indexer(monkeypatch):
+    """The KL is one-directional, so its gradient must stop at the indexer.
+
+    Run the same weights and input with the loss off and with a large
+    coefficient: every backbone gradient has to be bit-identical, while the
+    indexer goes from no gradient to a non-zero one.
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    torch.manual_seed(0)
+    off = _make_csa_attention(coeff=0.0)
+    on = _make_csa_attention(coeff=1e3)
+    on.load_state_dict(off.state_dict())
+    off.train()
+    on.train()
+
+    B, S = 1, 8
+    torch.manual_seed(2)
+    base = torch.randn(B, S, _HIDDEN, dtype=_DTYPE)
+    position_ids = torch.arange(S).unsqueeze(0).expand(B, S)
+
+    grads = {}
+    inputs = {}
+    for tag, module in (("off", off), ("on", on)):
+        hidden = base.clone().requires_grad_(True)
+        module(hidden, position_ids).sum().backward()
+        inputs[tag] = hidden.grad.clone()
+        grads[tag] = {n: p.grad.clone() for n, p in module.named_parameters() if p.grad is not None}
+
+    for name in _backbone_names(off):
+        if name not in grads["off"]:
+            continue
+        torch.testing.assert_close(
+            grads["on"][name],
+            grads["off"][name],
+            rtol=0,
+            atol=0,
+            msg=lambda s, n=name: f"KL perturbed backbone gradient {n}: {s}",
+        )
+    torch.testing.assert_close(inputs["on"], inputs["off"], rtol=0, atol=0)
+
+    indexer_on = {n for n in grads["on"] if n.startswith("indexer.")}
+    assert indexer_on, "the indexer received no gradient at coeff > 0"
+    assert not any(n.startswith("indexer.") for n in grads["off"])
+
+
+def test_ape_and_sink_receive_fp32_gradients(monkeypatch):
+    """With the pinning opted in, it must survive into the gradients too.
+
+    ``PRIMUS_V4_KEEP_FP32`` ships off, so this asserts the opt-in path rather
+    than the default one -- a mixed-dtype model is exactly what the distributed
+    optimizer cannot take (see ``keep_in_fp32``).
+    """
+    monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False)
+    monkeypatch.setenv("PRIMUS_V4_KEEP_FP32", "1")
+    torch.manual_seed(0)
+    attn = _make_csa_attention()
+    attn.bfloat16()
+    attn.train()
+
+    B, S = 1, 8
+    hidden = torch.randn(B, S, _HIDDEN, dtype=torch.bfloat16)
+    attn(hidden, torch.arange(S).unsqueeze(0).expand(B, S)).sum().backward()
+
+    assert attn.attn_sink.grad is not None
+    assert attn.attn_sink.grad.dtype == torch.float32
+    assert attn.compressor.ape.grad is not None
+    assert attn.compressor.ape.grad.dtype == torch.float32
+    # ...and the ordinary weights still carry BF16 gradients.
+    assert attn.linear_kv.weight.grad.dtype == torch.bfloat16
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py
new file mode 100644
index 000000000..7e2e9b39e
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py
@@ -0,0 +1,381 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for V4 routers (G4, plan-2 §04).
+
+Pins both routers to the HF reference at
+``DeepSeek-V4-Flash/inference/model.py:Gate.forward``. Two routers
+share the same scoring path; they only differ in *selection* (top-K
+argmax vs ``tid2eid`` lookup).
+
+Pass criteria (G4):
+* Identical weights -> identical sparse ``(probs, routing_map)`` to the
+  HF reference (max-abs <= 1e-6 fp32).
+* Routing weight gradient flows to ``weight`` (the learned gate).
+* For the hash router, ``tid2eid`` is a parameter with
+  ``requires_grad=False`` (frozen) so checkpoint round-trips preserve
+  it without polluting the optimizer state.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from primus.backends.megatron.core.transformer.moe.v4_hash_router import (
+    DeepseekV4HashRouter,
+    HashRouter,
+)
+from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+    DeepseekV4LearnedRouter,
+    V4TopKRouter,
+    v4_score_fn,
+)
+
+
+@pytest.fixture(autouse=True)
+def _v4_router_on_cuda(monkeypatch):
+    """The V4 routers are GPU-only: both the Triton and eager paths require
+    CUDA/HIP tensors, so default tensor creation to the CUDA device and skip on
+    a CPU-only host.  Force the eager path (``PRIMUS_V4_ROUTER_TRITON=0``) so
+    these stay exact router-vs-HF reference checks; the fused Triton path has
+    its own parity test (``test_router_post_triton``).
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("DeepSeek-V4 routers require a CUDA/HIP device")
+    monkeypatch.setenv("PRIMUS_V4_ROUTER_TRITON", "0")
+    torch.set_default_device("cuda")
+    try:
+        yield
+    finally:
+        torch.set_default_device("cpu")
+
+
+# ---------------------------------------------------------------------------
+# HF reference inline
+# ---------------------------------------------------------------------------
+
+
+def _hf_gate_forward(
+    *,
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    score_function: str,
+    topk: int,
+    route_scale: float,
+    tid2eid: Optional[torch.Tensor],
+    token_ids: Optional[torch.Tensor],
+):
+    """Mirrors HF reference ``Gate.forward`` exactly.
+
+    Returns dense ``(weights, indices)`` with shapes ``[N, K]``.
+    """
+    flat_hidden = hidden.reshape(-1, hidden.shape[-1])
+    scores = F.linear(flat_hidden.float(), weight.float())
+    if score_function == "softmax":
+        scores = scores.softmax(dim=-1)
+    elif score_function == "sigmoid":
+        scores = scores.sigmoid()
+    else:
+        scores = F.softplus(scores).sqrt()
+    original_scores = scores
+    if bias is not None:
+        scores = scores + bias.float()
+    if tid2eid is not None:
+        assert token_ids is not None
+        flat_ids = token_ids.reshape(-1).long()
+        indices = tid2eid[flat_ids].long()
+    else:
+        indices = scores.topk(topk, dim=-1).indices
+    weights = original_scores.gather(1, indices)
+    if score_function != "softmax":
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+    weights = weights * route_scale
+    return weights, indices
+
+
+def _sparse_from_dense(
+    *,
+    weights: torch.Tensor,
+    indices: torch.Tensor,
+    num_experts: int,
+):
+    """Pack dense ``(weights, indices)`` into the sparse ``(probs,
+    routing_map)`` format our routers return."""
+    N = weights.shape[0]
+    probs = torch.zeros(N, num_experts, dtype=weights.dtype, device=weights.device)
+    probs.scatter_(1, indices, weights)
+    routing_map = torch.zeros(N, num_experts, dtype=torch.bool, device=weights.device)
+    routing_map.scatter_(1, indices, True)
+    return probs, routing_map
+
+
+# ---------------------------------------------------------------------------
+# Score function
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
+def test_v4_score_fn_matches_inline_reference(score_function: str) -> None:
+    torch.manual_seed(123)
+    logits = torch.randn(8, 16, dtype=torch.float32)
+
+    out = v4_score_fn(logits, score_function=score_function)
+    if score_function == "softmax":
+        ref = F.softmax(logits, dim=-1)
+    elif score_function == "sigmoid":
+        ref = torch.sigmoid(logits)
+    else:
+        ref = F.softplus(logits).sqrt()
+    assert (out - ref).abs().max().item() <= 1.0e-6
+
+
+# ---------------------------------------------------------------------------
+# Learned router
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"])
+@pytest.mark.parametrize("enable_expert_bias", [False, True])
+def test_learned_router_matches_hf_reference(score_function: str, enable_expert_bias: bool) -> None:
+    """G4 learned-router gate: matches HF reference exactly (fp32)."""
+    torch.manual_seed(2025)
+    H, E, K = 12, 8, 2
+    route_scale = 1.5
+
+    router = DeepseekV4LearnedRouter(
+        hidden_size=H,
+        num_experts=E,
+        topk=K,
+        score_function=score_function,
+        enable_expert_bias=enable_expert_bias,
+        topk_scaling_factor=route_scale,
+    )
+    if enable_expert_bias:
+        with torch.no_grad():
+            router.expert_bias.copy_(torch.randn(E, dtype=torch.float32))
+
+    hidden = torch.randn(2, 5, H, dtype=torch.float32)
+
+    probs, routing_map = router(hidden)
+
+    weights_hf, indices_hf = _hf_gate_forward(
+        hidden=hidden,
+        weight=router.weight,
+        bias=router.expert_bias,
+        score_function=score_function,
+        topk=K,
+        route_scale=route_scale,
+        tid2eid=None,
+        token_ids=None,
+    )
+    probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E)
+
+    assert probs.shape == probs_hf.shape
+    assert routing_map.shape == routing_map_hf.shape
+    assert torch.equal(routing_map, routing_map_hf), "routing_map differs"
+    max_abs = (probs - probs_hf).abs().max().item()
+    assert max_abs <= 1.0e-6, f"probs max-abs vs HF reference = {max_abs}"
+
+
+def test_learned_router_back_compat_alias() -> None:
+    """``V4TopKRouter`` is an alias for the renamed router."""
+    assert V4TopKRouter is DeepseekV4LearnedRouter
+
+
+def test_learned_router_grad_flows_to_gate_weight() -> None:
+    """Probs gradient propagates back into ``weight`` (the gate)."""
+    torch.manual_seed(11)
+    H, E, K = 6, 4, 2
+    router = DeepseekV4LearnedRouter(hidden_size=H, num_experts=E, topk=K, score_function="sqrtsoftplus")
+    hidden = torch.randn(3, 4, H, dtype=torch.float32, requires_grad=False)
+    probs, _ = router(hidden)
+    loss = probs.sum()
+    loss.backward()
+    assert router.weight.grad is not None
+    assert torch.isfinite(router.weight.grad).all()
+    assert router.weight.grad.abs().sum().item() > 0.0
+
+
+def test_learned_router_expert_bias_does_not_contribute_to_probs_grad() -> None:
+    """``expert_bias`` is selection-only: probs use un-biased scores.
+
+    With ``enable_expert_bias=True`` the bias enters the *selection*
+    path (top-K) but not the gathered weights. The gradient on the bias
+    can therefore be exactly zero when the chosen indices stay the same
+    if you only differentiate the routing weights at the selected
+    experts. We assert ``expert_bias.grad is None`` (no graph) **after**
+    a forward + backward, since the bias never enters the autograd
+    chain in our implementation.
+    """
+    torch.manual_seed(42)
+    H, E, K = 6, 4, 2
+    router = DeepseekV4LearnedRouter(
+        hidden_size=H,
+        num_experts=E,
+        topk=K,
+        score_function="sqrtsoftplus",
+        enable_expert_bias=True,
+    )
+    hidden = torch.randn(3, 4, H, dtype=torch.float32)
+    probs, _ = router(hidden)
+    probs.sum().backward()
+    # Bias is detached from the probs graph — grad stays None (or zero).
+    assert (router.expert_bias.grad is None) or (router.expert_bias.grad.abs().max().item() == 0.0)
+
+
+def test_learned_router_softmax_skips_renormalization() -> None:
+    """Softmax probs already sum to 1, so the post-topK renorm is skipped.
+
+    We assert by setting ``route_scale != 1`` and confirming the gathered
+    weights match HF (i.e. they are *not* re-normalized to 1 before the
+    scale).
+    """
+    torch.manual_seed(0)
+    H, E, K = 6, 4, 2
+    router = DeepseekV4LearnedRouter(
+        hidden_size=H,
+        num_experts=E,
+        topk=K,
+        score_function="softmax",
+        topk_scaling_factor=2.0,
+    )
+    hidden = torch.randn(2, 3, H, dtype=torch.float32)
+    probs, routing_map = router(hidden)
+
+    weights_hf, indices_hf = _hf_gate_forward(
+        hidden=hidden,
+        weight=router.weight,
+        bias=None,
+        score_function="softmax",
+        topk=K,
+        route_scale=2.0,
+        tid2eid=None,
+        token_ids=None,
+    )
+    probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E)
+    assert torch.equal(routing_map, routing_map_hf)
+    assert (probs - probs_hf).abs().max().item() <= 1.0e-6
+
+
+# ---------------------------------------------------------------------------
+# Hash router
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"])
+def test_hash_router_matches_hf_reference(score_function: str) -> None:
+    """G4 hash-router: matches HF reference exactly (fp32)."""
+    torch.manual_seed(7)
+    H, E, K, V = 12, 8, 2, 32
+
+    router = DeepseekV4HashRouter(
+        hidden_size=H,
+        num_experts=E,
+        topk=K,
+        vocab_size=V,
+        seed=17,
+        score_function=score_function,
+        topk_scaling_factor=1.0,
+    )
+    hidden = torch.randn(2, 5, H, dtype=torch.float32)
+    token_ids = torch.randint(0, V, (2, 5), dtype=torch.long)
+
+    probs, routing_map = router(hidden, token_ids)
+
+    weights_hf, indices_hf = _hf_gate_forward(
+        hidden=hidden,
+        weight=router.weight,
+        bias=None,
+        score_function=score_function,
+        topk=K,
+        route_scale=1.0,
+        tid2eid=router.tid2eid,
+        token_ids=token_ids,
+    )
+    probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E)
+    assert torch.equal(routing_map, routing_map_hf), "routing_map differs"
+    max_abs = (probs - probs_hf).abs().max().item()
+    assert max_abs <= 1.0e-6, f"probs max-abs vs HF reference = {max_abs}"
+
+
+def test_hash_router_back_compat_alias() -> None:
+    """``HashRouter`` is an alias for the renamed router."""
+    assert HashRouter is DeepseekV4HashRouter
+
+
+def test_hash_router_tid2eid_is_frozen_parameter() -> None:
+    """``tid2eid`` is a Parameter with requires_grad=False.
+
+    Matches the HF reference layout (released checkpoint stores it as
+    a parameter) so a state-dict round-trip preserves it without
+    pulling it into the optimizer state.
+    """
+    router = DeepseekV4HashRouter(
+        hidden_size=8,
+        num_experts=4,
+        topk=2,
+        vocab_size=16,
+        seed=0,
+    )
+    assert "tid2eid" in dict(router.named_parameters())
+    assert router.tid2eid.requires_grad is False
+    assert router.tid2eid.dtype == torch.int32
+    assert tuple(router.tid2eid.shape) == (16, 2)
+
+
+def test_hash_router_state_dict_keys() -> None:
+    """State-dict exposes ``weight`` and ``tid2eid`` (matches HF gate keys)."""
+    router = DeepseekV4HashRouter(hidden_size=8, num_experts=4, topk=2, vocab_size=16, seed=0)
+    keys = set(router.state_dict().keys())
+    assert "weight" in keys
+    assert "tid2eid" in keys
+
+
+def test_hash_router_grad_flows_to_gate_weight() -> None:
+    """Even with static expert ids, the routing weights' gradient flows
+    back into the learned gate ``weight``."""
+    torch.manual_seed(13)
+    H, E, K, V = 6, 4, 2, 16
+    router = DeepseekV4HashRouter(
+        hidden_size=H,
+        num_experts=E,
+        topk=K,
+        vocab_size=V,
+        seed=0,
+        score_function="sqrtsoftplus",
+    )
+    hidden = torch.randn(3, 4, H, dtype=torch.float32)
+    token_ids = torch.randint(0, V, (3, 4), dtype=torch.long)
+    probs, _ = router(hidden, token_ids)
+    probs.sum().backward()
+    assert router.weight.grad is not None
+    assert torch.isfinite(router.weight.grad).all()
+    assert router.weight.grad.abs().sum().item() > 0.0
+    # tid2eid should never accumulate gradient.
+    assert router.tid2eid.grad is None
+
+
+def test_hash_router_deterministic_table_across_seeds() -> None:
+    """Same seed -> identical table; different seeds -> different table."""
+    a = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=42)
+    b = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=42)
+    c = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=43)
+    assert torch.equal(a.tid2eid, b.tid2eid)
+    assert not torch.equal(a.tid2eid, c.tid2eid)
+
+
+def test_hash_router_rejects_shape_mismatch() -> None:
+    """``hidden`` and ``token_ids`` must flatten to the same length."""
+    router = DeepseekV4HashRouter(hidden_size=4, num_experts=4, topk=2, vocab_size=8, seed=0)
+    hidden = torch.randn(2, 3, 4, dtype=torch.float32)  # N=6
+    token_ids = torch.zeros(1, 5, dtype=torch.long)  # N=5
+    with pytest.raises(ValueError, match="flatten to the same length"):
+        router(hidden, token_ids)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_selector_precision.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_selector_precision.py
new file mode 100644
index 000000000..70b85af62
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_selector_precision.py
@@ -0,0 +1,135 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""The CSA selector path stays high precision unless explicitly opted in.
+
+The open-source reference wraps the DeepSeek-V4 compressor and the indexer's
+weight projection in an fp8-disabled context so they keep running in BF16 even
+when the enclosing layer is under an FP8 autocast: the indexer decides *which*
+512 compressed entries each query attends to, so quantization error there
+changes the selection rather than nudging a value.
+
+These tests pin the two ways that contract can be lost:
+
+* the indexer projections must not be quantized as a side effect of enabling
+  FP8 on the *attention* projections -- those are separate decisions;
+* the compressor must have no quantized branch at all.
+
+Torch-only and CPU-friendly: the FP8 branches are gated on env knobs read at
+call time, so the gating can be exercised without a GPU.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from primus.backends.megatron.core.transformer import (  # noqa: E402
+    indexer as indexer_mod,
+)
+from primus.backends.megatron.core.transformer.compressor import (  # noqa: E402
+    Compressor,
+)
+
+# ---------------------------------------------------------------------------
+# Indexer projections
+# ---------------------------------------------------------------------------
+
+
+def test_indexer_proj_fp8_is_off_by_default(monkeypatch):
+    monkeypatch.delenv("PRIMUS_V4_FP8_INDEXER_PROJ", raising=False)
+    monkeypatch.delenv("PRIMUS_V4_FP8_ATTN_PROJ", raising=False)
+    assert indexer_mod._indexer_fp8_proj_enabled() is False
+
+
+def test_attention_proj_fp8_does_not_quantize_the_selector(monkeypatch):
+    """Enabling FP8 on the attention projections must leave the indexer alone.
+
+    This is the regression: the knob used to be shared, so turning on FP8 for
+    the attention projections silently quantized the selector too.
+    """
+    monkeypatch.delenv("PRIMUS_V4_FP8_INDEXER_PROJ", raising=False)
+    monkeypatch.setenv("PRIMUS_V4_FP8_ATTN_PROJ", "1")
+    assert indexer_mod._indexer_fp8_proj_enabled() is False
+
+
+def test_indexer_proj_fp8_reads_only_its_own_knob(monkeypatch):
+    """The opt-in still exists for QAT experiments, under its own name.
+
+    Asserts on which environment variables the helper actually consults rather
+    than on its source text: the docstring legitimately names the attention flag
+    while explaining that it is *not* what gates this.
+    """
+    consulted = []
+    real_get = os.environ.get
+
+    def _spy(key, *args, **kwargs):
+        consulted.append(key)
+        return real_get(key, *args, **kwargs)
+
+    monkeypatch.setattr(os.environ, "get", _spy)
+    monkeypatch.setenv("PRIMUS_V4_FP8_INDEXER_PROJ", "1")
+    indexer_mod._indexer_fp8_proj_enabled()
+
+    assert "PRIMUS_V4_FP8_INDEXER_PROJ" in consulted
+    assert "PRIMUS_V4_FP8_ATTN_PROJ" not in consulted
+
+
+def test_indexer_forward_is_unaffected_by_the_attention_fp8_knob(monkeypatch):
+    """End to end: same scores with and without PRIMUS_V4_FP8_ATTN_PROJ."""
+    torch.manual_seed(0)
+    idx = indexer_mod.Indexer(
+        hidden_size=32,
+        index_head_dim=16,
+        index_n_heads=4,
+        index_topk=4,
+        compress_ratio=4,
+    )
+    hidden = torch.randn(1, 16, 32, dtype=torch.float32)
+
+    monkeypatch.delenv("PRIMUS_V4_FP8_ATTN_PROJ", raising=False)
+    with torch.no_grad():
+        idxs_a, scores_a = idx(hidden)
+
+    monkeypatch.setenv("PRIMUS_V4_FP8_ATTN_PROJ", "1")
+    with torch.no_grad():
+        idxs_b, scores_b = idx(hidden)
+
+    assert torch.equal(idxs_a, idxs_b)
+    finite = torch.isfinite(scores_a)
+    torch.testing.assert_close(scores_a[finite], scores_b[finite], rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# Compressor
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("ratio", [4, 128])
+def test_compressor_projections_are_plain_linears(ratio: int):
+    """No quantized branch: plain ``nn.Linear`` ignores any FP8 autocast."""
+    comp = Compressor(hidden_size=32, head_dim=16, ratio=ratio)
+    linears = [m for m in comp.modules() if isinstance(m, torch.nn.Linear)]
+    assert linears, "compressor must own its KV / gate projections"
+    for lin in linears:
+        assert type(lin) is torch.nn.Linear, f"unexpected projection type {type(lin).__name__}"
+
+
+def test_compressor_has_no_fp8_code_path():
+    """Guard against a quantized branch being added without revisiting this."""
+    from primus.backends.megatron.core.transformer import compressor as compressor_mod
+
+    source = inspect.getsource(compressor_mod.Compressor.forward)
+    for token in ("fp8", "FP8", "fp4", "FP4", "quant"):
+        assert token not in source, (
+            f"Compressor.forward references {token!r}: the reference keeps the "
+            "compressor in high precision, so a low-precision branch needs an "
+            "explicit decision (and this test updated)."
+        )
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_seq_balance_loss.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_seq_balance_loss.py
new file mode 100644
index 000000000..2ea06999b
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_seq_balance_loss.py
@@ -0,0 +1,268 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for the MoE sequence-wise load-balancing loss.
+
+The V4 yaml declares ``moe_router_load_balancing_type: seq_aux_loss`` with
+``moe_aux_loss_coeff: 0.001``, matching the paper's "slight sequence-wise
+balance loss". These tests pin the formula (whose normalisation makes a
+perfectly balanced sequence score exactly ``alpha``), the fact that it is
+per-sequence rather than per-batch, and the sharding contract.
+
+The loss function itself is plain torch and runs on CPU; the router that drives
+it is GPU-only (``_compute_route`` asserts CUDA), so the wiring assertions are
+made against the router's constructor and the config plumbing rather than a
+forward pass.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from primus.backends.megatron.core.transformer.moe.v4_seq_balance_loss import (  # noqa: E402
+    SEQ_BALANCE_LOSS_NAME,
+    log_seq_balance_loss,
+    normalise_affinities,
+    sequence_balance_loss,
+)
+
+
+def _balanced(batch: int, seq: int, experts: int, topk: int):
+    """Perfectly balanced routing with uniform affinities."""
+    scores = torch.full((batch, seq, experts), 1.0 / experts, dtype=torch.float32)
+    routing_map = torch.zeros(batch, seq, experts, dtype=torch.bool)
+    for b in range(batch):
+        for t in range(seq):
+            picks = [(t * topk + k) % experts for k in range(topk)]
+            routing_map[b, t, picks] = True
+    return scores, routing_map
+
+
+# ---------------------------------------------------------------------------
+# The formula
+# ---------------------------------------------------------------------------
+
+
+def test_balanced_sequence_scores_exactly_the_coefficient():
+    """``f`` sums to E and ``P`` to 1, so balanced routing gives ``alpha``.
+
+    This is what makes the logged number readable: anything above the
+    coefficient is measurable imbalance.
+    """
+    scores, routing_map = _balanced(batch=2, seq=8, experts=4, topk=2)
+    loss = sequence_balance_loss(scores=scores, routing_map=routing_map, topk=2, coeff=1e-3)
+    assert loss.item() == pytest.approx(1e-3, rel=1e-5)
+
+
+def test_collapsed_sequence_is_penalised():
+    """Routing every token to the same experts must cost more than balanced."""
+    batch, seq, experts, topk = 1, 8, 4, 2
+    balanced_scores, balanced_map = _balanced(batch, seq, experts, topk)
+
+    collapsed_map = torch.zeros(batch, seq, experts, dtype=torch.bool)
+    collapsed_map[..., :topk] = True
+    collapsed_scores = torch.zeros(batch, seq, experts, dtype=torch.float32)
+    collapsed_scores[..., :topk] = 0.5
+
+    balanced = sequence_balance_loss(scores=balanced_scores, routing_map=balanced_map, topk=topk, coeff=1.0)
+    collapsed = sequence_balance_loss(
+        scores=collapsed_scores, routing_map=collapsed_map, topk=topk, coeff=1.0
+    )
+    assert collapsed.item() > balanced.item()
+    # Fully collapsed onto topk of E experts costs E / topk.
+    assert collapsed.item() == pytest.approx(experts / topk, rel=1e-5)
+
+
+def test_loss_is_linear_in_the_coefficient():
+    scores, routing_map = _balanced(1, 8, 4, 2)
+    scores = scores + torch.rand_like(scores) * 0.1
+    scores = normalise_affinities(scores)
+    a = sequence_balance_loss(scores=scores, routing_map=routing_map, topk=2, coeff=1.0)
+    b = sequence_balance_loss(scores=scores, routing_map=routing_map, topk=2, coeff=3.0)
+    assert b.item() == pytest.approx(3.0 * a.item(), rel=1e-6)
+
+
+def test_loss_is_per_sequence_not_per_batch():
+    """A batch of two collapsed-but-opposite sequences must still be penalised.
+
+    Pooled over the batch the two sequences look balanced; per-sequence they are
+    each fully collapsed. Catching this is the entire reason the loss exists.
+    """
+    seq, experts, topk = 8, 4, 2
+    routing_map = torch.zeros(2, seq, experts, dtype=torch.bool)
+    routing_map[0, :, :topk] = True
+    routing_map[1, :, topk:] = True
+
+    scores = torch.zeros(2, seq, experts, dtype=torch.float32)
+    scores[0, :, :topk] = 0.5
+    scores[1, :, topk:] = 0.5
+
+    loss = sequence_balance_loss(scores=scores, routing_map=routing_map, topk=topk, coeff=1.0)
+    # Each sequence is collapsed onto topk of E, so the mean is E / topk.
+    assert loss.item() == pytest.approx(experts / topk, rel=1e-5)
+
+    # Whereas the batch-pooled view of the same routing looks perfectly balanced.
+    pooled_map = routing_map.reshape(1, 2 * seq, experts)
+    pooled_scores = scores.reshape(1, 2 * seq, experts)
+    pooled = sequence_balance_loss(scores=pooled_scores, routing_map=pooled_map, topk=topk, coeff=1.0)
+    assert pooled.item() < loss.item()
+
+
+def test_loss_is_differentiable_through_the_affinities():
+    scores, routing_map = _balanced(1, 8, 4, 2)
+    scores = scores.clone().requires_grad_(True)
+    sequence_balance_loss(scores=scores, routing_map=routing_map, topk=2, coeff=1.0).backward()
+
+    assert scores.grad is not None
+    assert torch.isfinite(scores.grad).all()
+    assert scores.grad.abs().sum() > 0.0
+
+
+def test_rejects_flat_inputs():
+    """The loss is per-sequence, so a collapsed [N, E] view is a bug."""
+    with pytest.raises(ValueError):
+        sequence_balance_loss(
+            scores=torch.rand(8, 4),
+            routing_map=torch.zeros(8, 4, dtype=torch.bool),
+            topk=2,
+            coeff=1.0,
+        )
+
+
+def test_normalise_affinities_produces_a_distribution():
+    raw = torch.rand(2, 5, 7) + 0.1
+    normed = normalise_affinities(raw)
+    torch.testing.assert_close(normed.sum(dim=-1), torch.ones(2, 5), rtol=1e-5, atol=1e-6)
+
+
+def test_normalise_affinities_is_safe_on_all_zero_rows():
+    normed = normalise_affinities(torch.zeros(1, 2, 4))
+    assert torch.isfinite(normed).all()
+
+
+# ---------------------------------------------------------------------------
+# Sharding
+# ---------------------------------------------------------------------------
+
+
+def test_expert_counts_are_reduced_across_the_sequence_shards():
+    """With the sequence split, ``f`` must be global while ``P`` stays local."""
+    batch, seq, experts, topk = 1, 8, 4, 2
+    scores, routing_map = _balanced(batch, seq, experts, topk)
+
+    class _Group:
+        pass
+
+    group = _Group()
+    calls = []
+    real_world_size = torch.distributed.get_world_size
+    real_all_reduce = torch.distributed.all_reduce
+    torch.distributed.get_world_size = lambda g=None: 2 if g is group else 1
+    torch.distributed.all_reduce = lambda t, group=None: calls.append((tuple(t.shape), group))
+    try:
+        sharded = sequence_balance_loss(
+            scores=scores, routing_map=routing_map, topk=topk, coeff=1.0, reduce_group=group
+        )
+    finally:
+        torch.distributed.get_world_size = real_world_size
+        torch.distributed.all_reduce = real_all_reduce
+
+    assert calls == [((batch, experts), group)], "must reduce the [B, E] expert counts"
+    # The fake all-reduce is a no-op, so with T doubled and counts unchanged the
+    # local share is a quarter of the unsharded value (T appears squared).
+    local = sequence_balance_loss(scores=scores, routing_map=routing_map, topk=topk, coeff=1.0)
+    assert sharded.item() == pytest.approx(local.item() / 4.0, rel=1e-5)
+
+
+def test_no_group_means_no_collective():
+    batch, seq, experts, topk = 1, 8, 4, 2
+    scores, routing_map = _balanced(batch, seq, experts, topk)
+    calls = []
+    real_all_reduce = torch.distributed.all_reduce
+    torch.distributed.all_reduce = lambda t, group=None: calls.append(t)
+    try:
+        sequence_balance_loss(scores=scores, routing_map=routing_map, topk=topk, coeff=1.0)
+    finally:
+        torch.distributed.all_reduce = real_all_reduce
+    assert not calls
+
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+
+
+def test_loss_is_reported_under_the_expected_name(monkeypatch):
+    pytest.importorskip("megatron.core.transformer.moe.moe_utils")
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    recorded = []
+    monkeypatch.setattr(
+        moe_utils,
+        "save_to_aux_losses_tracker",
+        lambda name, loss, layer_number, num_layers, **kw: recorded.append(
+            (name, float(loss), layer_number, num_layers, kw.get("reduce_group"))
+        ),
+    )
+
+    log_seq_balance_loss(torch.tensor(0.5), layer_number=3, num_layers=44, reduce_group=None)
+
+    assert recorded == [(SEQ_BALANCE_LOSS_NAME, 0.5, 3, 44, None)]
+
+
+@pytest.mark.parametrize("layer_number", [None, 0])
+def test_unnumbered_layers_are_not_reported(monkeypatch, layer_number):
+    pytest.importorskip("megatron.core.transformer.moe.moe_utils")
+    import megatron.core.transformer.moe.moe_utils as moe_utils
+
+    recorded = []
+    monkeypatch.setattr(moe_utils, "save_to_aux_losses_tracker", lambda *a, **kw: recorded.append(a))
+    log_seq_balance_loss(torch.tensor(1.0), layer_number=layer_number, num_layers=44)
+    assert not recorded
+
+
+# ---------------------------------------------------------------------------
+# Router / config wiring
+# ---------------------------------------------------------------------------
+
+
+def test_router_defaults_to_the_loss_being_off():
+    from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+        DeepseekV4LearnedRouter,
+    )
+
+    router = DeepseekV4LearnedRouter(hidden_size=16, num_experts=8, topk=2)
+    assert router.seq_balance_loss_coeff == 0.0
+    assert router.seq_balance_loss_enabled is False
+
+
+def test_router_enables_the_loss_from_its_coefficient():
+    from primus.backends.megatron.core.transformer.moe.v4_topk_router import (
+        DeepseekV4LearnedRouter,
+    )
+
+    router = DeepseekV4LearnedRouter(hidden_size=16, num_experts=8, topk=2, seq_balance_loss_coeff=1e-3)
+    assert router.seq_balance_loss_enabled is True
+
+
+def test_v4_yaml_declares_a_balancing_type_the_router_implements():
+    """The yaml said ``seq_aux_loss`` long before anything read it.
+
+    Pin the two together so the declaration cannot drift back into being
+    decorative, and so a switch to a balancing type this router does not
+    implement is caught here.
+    """
+    from pathlib import Path
+
+    from primus.core.config.yaml_loader import parse_yaml
+
+    yaml_dir = Path(__file__).resolve().parents[5] / "primus" / "configs" / "models" / "megatron"
+    parsed = parse_yaml(str(yaml_dir / "deepseek_v4_base.yaml"))
+
+    assert parsed["moe_router_load_balancing_type"] == "seq_aux_loss"
+    assert float(parsed["moe_aux_loss_coeff"]) > 0.0
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_shared_expert_overlap.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_shared_expert_overlap.py
new file mode 100644
index 000000000..d18a30e27
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_shared_expert_overlap.py
@@ -0,0 +1,151 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Tests for shared-expert / all-to-all overlap wiring.
+
+``moe_shared_expert_overlap: true`` in the model yaml asks for the shared
+expert's GEMMs to be interleaved with the dispatch and combine all-to-alls.
+The token dispatcher owns the hooks, and once it owns the shared expert it also
+**adds the shared output into the combine result** -- so the risky part of this
+wiring is not the overlap itself but the double-count it would cause if the MoE
+layer kept adding the shared expert as well.
+
+These tests drive ``_enable_shared_expert_overlap`` and the forward's add-or-not
+decision against fakes, so they need no GPU or process group.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from primus.backends.megatron.core.transformer.moe.v4_moe import (  # noqa: E402
+    DeepseekV4MoE,
+)
+
+
+class _FakeConfig:
+    def __init__(self, overlap: bool):
+        self.moe_shared_expert_overlap = overlap
+
+
+class _OverlapCapableSharedExpert:
+    """Stands in for ``SharedExpertMLP``: has the overlap protocol."""
+
+    def pre_forward_comm(self, x):  # pragma: no cover - never called here
+        raise AssertionError
+
+
+class _PlainSharedExpert:
+    """Stands in for the CPU ``ClampedSwiGLUMLP``: no overlap protocol."""
+
+
+class _Dispatcher:
+    def __init__(self):
+        self.received = None
+
+    def set_shared_experts(self, shared_experts):
+        self.received = shared_experts
+
+
+class _FlexDispatcher:
+    def set_shared_experts(self, shared_experts):
+        raise NotImplementedError("Shared expert overlap is not supported")
+
+
+def _layer(*, overlap: bool, shared_expert, dispatcher) -> DeepseekV4MoE:
+    """A bare object with only what ``_enable_shared_expert_overlap`` reads."""
+    layer = DeepseekV4MoE.__new__(DeepseekV4MoE)
+    layer.config = _FakeConfig(overlap)
+    layer.shared_expert = shared_expert
+    layer.token_dispatcher = dispatcher
+    layer.layer_idx = 3
+    return layer
+
+
+def test_overlap_hands_the_shared_expert_to_the_dispatcher():
+    dispatcher = _Dispatcher()
+    shared = _OverlapCapableSharedExpert()
+    layer = _layer(overlap=True, shared_expert=shared, dispatcher=dispatcher)
+
+    assert layer._enable_shared_expert_overlap() is True
+    assert dispatcher.received is shared
+
+
+def test_overlap_off_leaves_the_dispatcher_alone():
+    dispatcher = _Dispatcher()
+    layer = _layer(overlap=False, shared_expert=_OverlapCapableSharedExpert(), dispatcher=dispatcher)
+
+    assert layer._enable_shared_expert_overlap() is False
+    assert dispatcher.received is None
+
+
+def test_unsupported_dispatcher_falls_back_to_serial():
+    """The flex dispatcher rejects overlap; that must not be fatal."""
+    layer = _layer(
+        overlap=True,
+        shared_expert=_OverlapCapableSharedExpert(),
+        dispatcher=_FlexDispatcher(),
+    )
+    assert layer._enable_shared_expert_overlap() is False
+
+
+def test_shared_expert_without_the_protocol_is_not_handed_over():
+    """The CPU shared expert has no ``pre_forward_comm``."""
+    dispatcher = _Dispatcher()
+    layer = _layer(overlap=True, shared_expert=_PlainSharedExpert(), dispatcher=dispatcher)
+
+    assert layer._enable_shared_expert_overlap() is False
+    assert dispatcher.received is None
+
+
+def test_no_shared_expert_means_no_overlap():
+    layer = _layer(overlap=True, shared_expert=None, dispatcher=_Dispatcher())
+    assert layer._enable_shared_expert_overlap() is False
+
+
+# ---------------------------------------------------------------------------
+# The add-or-not decision in forward
+# ---------------------------------------------------------------------------
+
+
+def _forward_layer(*, overlap: bool):
+    """A layer stub exercising only the production branch of ``forward``."""
+    layer = DeepseekV4MoE.__new__(DeepseekV4MoE)
+    layer.local_experts = None
+    layer.hidden_size = 4
+    layer.shared_expert_overlap = overlap
+    layer.shared_expert = lambda x: torch.ones_like(x)
+    layer._dispatcher_forward = lambda hidden, probs, routing_map: torch.zeros_like(hidden)
+    layer._route = lambda hidden, token_ids: (None, None)
+    return layer
+
+
+def test_forward_does_not_add_the_shared_expert_under_overlap():
+    """The dispatcher already folded it in; adding again would double-count."""
+    layer = _forward_layer(overlap=True)
+    hidden = torch.zeros(2, 3, 4)
+    out = DeepseekV4MoE.forward(layer, hidden)
+    torch.testing.assert_close(out, torch.zeros_like(hidden))
+
+
+def test_forward_adds_the_shared_expert_without_overlap():
+    layer = _forward_layer(overlap=False)
+    hidden = torch.zeros(2, 3, 4)
+    out = DeepseekV4MoE.forward(layer, hidden)
+    torch.testing.assert_close(out, torch.ones_like(hidden))
+
+
+def test_yaml_requests_overlap():
+    """Pin the declaration the wiring above now honours."""
+    from pathlib import Path
+
+    from primus.core.config.yaml_loader import parse_yaml
+
+    yaml_dir = Path(__file__).resolve().parents[5] / "primus" / "configs" / "models" / "megatron"
+    parsed = parse_yaml(str(yaml_dir / "deepseek_v4_base.yaml"))
+    assert parsed["moe_shared_expert_overlap"] is True
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py
new file mode 100644
index 000000000..6d3c6e65c
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py
@@ -0,0 +1,397 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-3 P23 — Turbo DeepEP dispatcher in V4 specs.
+
+Today the V4 spec build captures
+:class:`megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher`
+at module-import time, while the Primus turbo patch
+(``primus.backends.megatron.patches.turbo.moe_dispatcher_patches``) only
+rebinds that module attribute at ``before_train``.  The patch fires
+*after* V4 spec build, so V4 silently runs the upstream
+:class:`MoEFlexTokenDispatcher` even when the user opted into
+``use_turbo_deepep=True``.
+
+P23 fixes this by resolving the dispatcher class **at V4 spec-build time**
+through :func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs._pick_v4_dispatcher_cls`,
+which checks ``args.enable_primus_turbo``, ``args.use_turbo_deepep``,
+``args.tensor_model_parallel_size`` and the import status of the
+``primus_turbo`` package.  Two side-effects ensure the rest of the
+stack agrees:
+
+* :func:`_maybe_plumb_v4_turbo_deepep_args` (in
+  ``deepseek_v4_builders.py``) mutates ``args.moe_enable_deepep`` and
+  ``args.moe_token_dispatcher_type`` BEFORE
+  ``core_transformer_config_from_args``, so the V4 ``config`` carries
+  the right ``moe_token_dispatcher_type``.
+* :meth:`DeepseekV4MoE._resolve_dispatcher_type_from_spec` recognises
+  ``PrimusTurboDeepEPTokenDispatcher`` (by class name, no
+  ``primus_turbo`` import required) and returns ``"flex"``, so the
+  per-layer log line ``"dispatcher active via …"`` reports the
+  correct type.
+
+Test gates exercised here:
+
+* **G20a — gating predicate**
+  :func:`is_v4_turbo_deepep_active` matches the four conditions used
+  by the upstream patch.
+* **G20b — args plumbing**
+  :func:`_maybe_plumb_v4_turbo_deepep_args` mutates only when all
+  gates pass; respects an explicit ``"allgather"`` opt-in.
+* **G20c — class resolution**
+  :func:`_pick_v4_dispatcher_cls` returns the right
+  ``(cls, type_name)`` tuple for each
+  ``(config.moe_token_dispatcher_type, args)`` combination.
+* **G20d — V4 MoE resolver recognises Turbo class**
+  :meth:`DeepseekV4MoE._resolve_dispatcher_type_from_spec` returns
+  ``"flex"`` for a class named
+  ``PrimusTurboDeepEPTokenDispatcher`` (mocked when ``primus_turbo``
+  is not installed).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from megatron.core.transformer.moe.token_dispatcher import (
+    MoEAllGatherTokenDispatcher,
+    MoEAlltoAllTokenDispatcher,
+    MoEFlexTokenDispatcher,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+from primus.backends.megatron.core.models.deepseek_v4 import (
+    deepseek_v4_builders,
+    deepseek_v4_layer_specs,
+)
+from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import (
+    _pick_v4_dispatcher_cls,
+    is_v4_turbo_deepep_active,
+)
+from primus.backends.megatron.core.transformer.moe.v4_moe import DeepseekV4MoE
+
+# ---------------------------------------------------------------------------
+# Test fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_args(
+    *,
+    enable_primus_turbo: bool = False,
+    use_turbo_deepep: bool = False,
+    tensor_model_parallel_size: int = 1,
+    moe_token_dispatcher_type: str = "alltoall",
+    moe_enable_deepep: bool = False,
+):
+    """Minimal ``args`` namespace mirroring Megatron's runtime args."""
+    return SimpleNamespace(
+        enable_primus_turbo=enable_primus_turbo,
+        use_turbo_deepep=use_turbo_deepep,
+        tensor_model_parallel_size=tensor_model_parallel_size,
+        moe_token_dispatcher_type=moe_token_dispatcher_type,
+        moe_enable_deepep=moe_enable_deepep,
+    )
+
+
+def _make_cfg(*, moe_token_dispatcher_type: str = "alltoall"):
+    """Minimal ``config`` namespace with the dispatcher-type field only."""
+    return SimpleNamespace(moe_token_dispatcher_type=moe_token_dispatcher_type)
+
+
+@pytest.fixture
+def primus_turbo_available(monkeypatch):
+    """Pretend the ``primus_turbo`` package is importable."""
+    real_find_spec = deepseek_v4_layer_specs.importlib.util.find_spec
+
+    def _fake(name):
+        if name == "primus_turbo":
+            return MagicMock()
+        return real_find_spec(name)
+
+    monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake)
+
+
+@pytest.fixture
+def primus_turbo_unavailable(monkeypatch):
+    """Pretend the ``primus_turbo`` package is NOT importable."""
+    real_find_spec = deepseek_v4_layer_specs.importlib.util.find_spec
+
+    def _fake(name):
+        if name == "primus_turbo":
+            return None
+        return real_find_spec(name)
+
+    monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake)
+
+
+# ``type(..., (), {})`` produces a class whose ``__name__`` is the
+# string we pass — using a plain ``class`` block would set ``__name__``
+# to the local symbol (e.g. ``_FakeTurboDispatcher``) regardless of
+# any class-level ``__name__`` assignment.  Only the class name
+# matters for ``_resolve_dispatcher_type_from_spec``; construction is
+# never exercised in these unit tests.
+_FakeTurboDispatcher = type("PrimusTurboDeepEPTokenDispatcher", (), {})
+
+
+@pytest.fixture
+def fake_turbo_class(monkeypatch):
+    """Inject a fake ``PrimusTurboDeepEPTokenDispatcher`` import target."""
+    monkeypatch.setattr(
+        deepseek_v4_layer_specs,
+        "_import_primus_turbo_deepep_dispatcher_cls",
+        lambda: _FakeTurboDispatcher,
+    )
+    return _FakeTurboDispatcher
+
+
+# ---------------------------------------------------------------------------
+# G20a — gating predicate (is_v4_turbo_deepep_active)
+# ---------------------------------------------------------------------------
+
+
+class TestGatingPredicate:
+    """Mirror of the conditions enforced by ``moe_dispatcher_patches``."""
+
+    def test_all_gates_open(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=1,
+        )
+        assert is_v4_turbo_deepep_active(args) is True
+
+    def test_primus_turbo_missing(self, primus_turbo_unavailable):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=1,
+        )
+        assert is_v4_turbo_deepep_active(args) is False
+
+    def test_enable_primus_turbo_off(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=False,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=1,
+        )
+        assert is_v4_turbo_deepep_active(args) is False
+
+    def test_use_turbo_deepep_off(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=False,
+            tensor_model_parallel_size=1,
+        )
+        assert is_v4_turbo_deepep_active(args) is False
+
+    def test_tp_gt_1_blocks(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=2,
+        )
+        assert is_v4_turbo_deepep_active(args) is False
+
+
+# ---------------------------------------------------------------------------
+# G20b — args plumbing (_maybe_plumb_v4_turbo_deepep_args)
+# ---------------------------------------------------------------------------
+
+
+class TestArgsPlumbing:
+    """The plumbing helper sets ``moe_enable_deepep`` + dispatcher type."""
+
+    def test_plumb_when_active(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            moe_token_dispatcher_type="alltoall",
+            moe_enable_deepep=False,
+        )
+        deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args)
+        assert args.moe_enable_deepep is True
+        assert args.moe_token_dispatcher_type == "flex"
+
+    def test_no_plumb_when_inactive(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=False,
+            use_turbo_deepep=True,
+            moe_token_dispatcher_type="alltoall",
+            moe_enable_deepep=False,
+        )
+        deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args)
+        assert args.moe_enable_deepep is False
+        assert args.moe_token_dispatcher_type == "alltoall"
+
+    def test_no_plumb_when_package_missing(self, primus_turbo_unavailable):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            moe_token_dispatcher_type="alltoall",
+        )
+        deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args)
+        assert args.moe_token_dispatcher_type == "alltoall"
+
+    def test_explicit_allgather_preserved(self, primus_turbo_available):
+        """User opted into ``allgather`` — never silently overridden."""
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            moe_token_dispatcher_type="allgather",
+        )
+        deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args)
+        # ``moe_enable_deepep`` may flip (Turbo wants it for any deepep
+        # path), but the user's dispatcher choice is preserved.
+        assert args.moe_token_dispatcher_type == "allgather"
+
+    def test_idempotent_when_already_flex(self, primus_turbo_available):
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            moe_token_dispatcher_type="flex",
+            moe_enable_deepep=True,
+        )
+        deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args)
+        assert args.moe_token_dispatcher_type == "flex"
+        assert args.moe_enable_deepep is True
+
+
+# ---------------------------------------------------------------------------
+# G20c — class resolution (_pick_v4_dispatcher_cls)
+# ---------------------------------------------------------------------------
+
+
+class TestPickDispatcherCls:
+    """Exhaustive class-resolution table."""
+
+    def test_alltoall_default(self):
+        cfg = _make_cfg(moe_token_dispatcher_type="alltoall")
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args())
+        assert cls is MoEAlltoAllTokenDispatcher
+        assert type_name == "alltoall"
+
+    def test_allgather_explicit(self):
+        cfg = _make_cfg(moe_token_dispatcher_type="allgather")
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args())
+        assert cls is MoEAllGatherTokenDispatcher
+        assert type_name == "allgather"
+
+    def test_flex_without_turbo(self):
+        cfg = _make_cfg(moe_token_dispatcher_type="flex")
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args())
+        assert cls is MoEFlexTokenDispatcher
+        assert type_name == "flex"
+
+    def test_flex_with_turbo_active(self, primus_turbo_available, fake_turbo_class):
+        cfg = _make_cfg(moe_token_dispatcher_type="flex")
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=1,
+        )
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args)
+        assert cls is fake_turbo_class
+        assert type_name == "flex"
+
+    def test_flex_with_turbo_active_but_class_missing(self, primus_turbo_available, monkeypatch):
+        """Gracefully fall back to the upstream class with a warning."""
+        monkeypatch.setattr(
+            deepseek_v4_layer_specs,
+            "_import_primus_turbo_deepep_dispatcher_cls",
+            lambda: None,
+        )
+        cfg = _make_cfg(moe_token_dispatcher_type="flex")
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=1,
+        )
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args)
+        assert cls is MoEFlexTokenDispatcher
+        assert type_name == "flex"
+
+    def test_flex_with_turbo_inactive_tp_gt_1(self, primus_turbo_available, fake_turbo_class):
+        """TP > 1 keeps the upstream Flex dispatcher."""
+        cfg = _make_cfg(moe_token_dispatcher_type="flex")
+        args = _make_args(
+            enable_primus_turbo=True,
+            use_turbo_deepep=True,
+            tensor_model_parallel_size=2,
+        )
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args)
+        assert cls is MoEFlexTokenDispatcher
+        assert type_name == "flex"
+
+    def test_unknown_type_falls_back_to_alltoall(self, caplog):
+        cfg = _make_cfg(moe_token_dispatcher_type="random_string")
+        with caplog.at_level("WARNING", logger=deepseek_v4_layer_specs.__name__):
+            cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args())
+        assert cls is MoEAlltoAllTokenDispatcher
+        assert type_name == "alltoall"
+        assert any("unsupported moe_token_dispatcher_type" in rec.message for rec in caplog.records)
+
+    def test_args_none_falls_back_when_megatron_not_initialised(
+        self, primus_turbo_available, fake_turbo_class, monkeypatch
+    ):
+        """``args=None`` + no Megatron args available → non-turbo branch."""
+
+        def _raises():
+            raise RuntimeError("Megatron not initialised in unit test")
+
+        # Patch the lazy import target inside the helper. ``megatron.training``
+        # may not have ``get_args`` bound yet depending on suite import order
+        # (it is populated lazily), so ``raising=False`` forces the "raises"
+        # behaviour either way; monkeypatch restores/deletes it on teardown.
+        import megatron.training as _megatron_training  # noqa: WPS433
+
+        monkeypatch.setattr(_megatron_training, "get_args", _raises, raising=False)
+        cfg = _make_cfg(moe_token_dispatcher_type="flex")
+        cls, type_name = _pick_v4_dispatcher_cls(cfg, args=None)
+        assert cls is MoEFlexTokenDispatcher
+        assert type_name == "flex"
+
+
+# ---------------------------------------------------------------------------
+# G20d — V4 MoE resolver recognises Turbo class
+# ---------------------------------------------------------------------------
+
+
+class TestV4MoEResolver:
+    """``DeepseekV4MoE._resolve_dispatcher_type_from_spec`` must label
+    :class:`PrimusTurboDeepEPTokenDispatcher` as ``"flex"``."""
+
+    def test_turbo_class_resolves_to_flex(self):
+        spec = ModuleSpec(module=_FakeTurboDispatcher)
+        assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "flex"
+
+    def test_turbo_class_bare_resolves_to_flex(self):
+        # The resolver also accepts a bare class (some call sites pass
+        # the type directly, not a ModuleSpec).
+        assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(_FakeTurboDispatcher) == "flex"
+
+    def test_alltoall_unchanged(self):
+        spec = ModuleSpec(module=MoEAlltoAllTokenDispatcher)
+        assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "alltoall"
+
+    def test_flex_unchanged(self):
+        spec = ModuleSpec(module=MoEFlexTokenDispatcher)
+        assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "flex"
+
+    def test_allgather_unchanged(self):
+        spec = ModuleSpec(module=MoEAllGatherTokenDispatcher)
+        assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "allgather"
+
+    def test_unknown_class_falls_back_to_alltoall(self, caplog):
+        class _Unknown:
+            __name__ = "_Unknown"
+
+        spec = ModuleSpec(module=_Unknown)
+        with caplog.at_level("WARNING"):
+            result = DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec)
+        assert result == "alltoall"
+        assert any("unsupported dispatcher module" in rec.message for rec in caplog.records)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py
new file mode 100644
index 000000000..3f799a93a
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py
@@ -0,0 +1,298 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""``turbo`` DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950),
+against the fp32 eager references.
+
+The ``turbo`` backend (:mod:`..._turbo_flydsl`) is the Primus-Turbo native-FlyDSL
+sparse-MLA v2 kernel-pair reached through the **turbo API**
+(``primus_turbo.flydsl.attention.kernels.sparse_mla_v2``), bound to the V4
+autograd adapters (:mod:`v4_csa_attention_turbo_flydsl`). It is the backend
+selected by ``use_v4_attention_backend`` / ``use_v4_csa_attention_backend =
+"turbo"``. This validates the production V4 invocation paths for all three layer
+kinds:
+
+* ``compress_ratio == 0``   (dense / SWA)  -> :func:`v4_attention_turbo`
+* ``compress_ratio == 128`` (HCA)          -> :func:`v4_attention_turbo`
+* ``compress_ratio == 4``   (CSA)          -> :func:`v4_csa_attention_turbo`
+
+Full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool, dsink) of the bf16 turbo
+adapter vs the fp32 eager reference (same tolerances the gluon_v2 backend UT uses;
+turbo shares the identical fused single-latent sparse-MLA-with-sink math). GPU-only;
+skipped off gfx950 / when primus_turbo's flydsl attention or ``flydsl`` is absent.
+
+FlyDSL fixes the head-block at 64, so ``num_heads`` must be a multiple of 64 here
+(the kernel asserts ``num_heads % 32 == 0``; H=64/128 are the production sizes).
+``S`` is a multiple of 128 so the cr=128 closed-form pool (``pool_cr = S/P = 128``)
+is exact.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("turbo sparse-MLA kernels require CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("flydsl", reason="flydsl pip package not installed")
+# The turbo API: primus_turbo must carry the flydsl sparse-MLA attention submodule.
+# The upstream module layout changed across releases; accept either (new flat
+# ``sparse_mla_fwd`` module, or the old ``kernels.sparse_mla_v2`` package).
+try:
+    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
+except ImportError:
+    pytest.importorskip(
+        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",
+        reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)",
+    )
+
+_ARCH = torch.cuda.get_device_properties(0).gcnArchName
+if "gfx950" not in _ARCH:
+    pytest.skip(f"turbo native-FlyDSL sparse-MLA targets gfx950; got {_ARCH}", allow_module_level=True)
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import (  # noqa: E402
+    eager_v4_attention,
+    eager_v4_csa_attention,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_turbo_flydsl import (  # noqa: E402
+    v4_attention_turbo,
+    v4_csa_attention_turbo,
+)
+
+D = 512  # V4 head_dim (RoPE baked in-place)
+_SWA = 128
+
+
+def _stats(a, b, *, sig=1e-2):
+    a = a.float()
+    b = b.float()
+    d = (a - b).abs()
+    m = b.abs() > sig
+    rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0)
+    med_rel = rel.median().item() if rel.numel() else 0.0
+    av, bv = a.flatten(), b.flatten()
+    cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item()
+    return d.max().item(), med_rel, cos_err
+
+
+def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None):
+    # Guard against a broken kernel returning NaN/Inf (e.g. the cr=4 fast_path
+    # overflow / bwd race the extraction found) — surface it as a clear failure.
+    assert torch.isfinite(
+        a.float()
+    ).all(), f"{name} has NaN/Inf ({int((~torch.isfinite(a.float())).sum())} bad)"
+    max_abs, med_rel, cos_err = _stats(a, b, sig=sig)
+    print(f"    {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True)
+    assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}"
+    if med is not None:
+        assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}"
+    if cos is not None:
+        assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}"
+
+
+def _leaf(x, *, fp32):
+    y = x.float() if fp32 else x.clone()
+    return y.detach().requires_grad_(True)
+
+
+# H must be a multiple of 64 (FlyDSL head-block); S a multiple of 128 (cr=128 pool_cr).
+# B is fixed to 1: the flydsl dense/HCA "banded" path uses a closed-form SWA window
+# [i-127..i] over the flat token axis, which crosses batch boundaries for B>1 (it assumes
+# a single contiguous sequence). Production / bench_v4_attention.py run B(mbs)=1; multi-batch
+# dense/HCA is a known flydsl-banded limitation (CSA/cr=4 is fine for B>1 — it uses the
+# per-batch-offset topk). Vary H (64/128), S (256/512), and the sink instead.
+_CASES = [
+    (1, 64, 256, None),
+    (1, 64, 512, 1.0),
+    (1, 128, 512, -1.0),
+]
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_turbo_dense_matches_eager(B, H, S, sink_init):
+    """cr=0 dense/SWA: turbo fwd+bwd vs eager_v4_attention."""
+    g = torch.Generator(device="cuda").manual_seed(0)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_attention_turbo(
+        qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale
+    )
+    og.backward(do)
+
+    qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = latf.unsqueeze(1).expand(B, H, S, D)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=_SWA,
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[turbo V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude
+    # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards.
+    _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_turbo_hca_matches_eager(B, H, S, sink_init):
+    """cr=128 HCA: turbo (local++pool) fwd+bwd vs eager_v4_attention."""
+    cr = 128
+    P = max(S // cr, 1)
+    g = torch.Generator(device="cuda").manual_seed(2)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    ti = torch.arange(S, device="cuda").view(S, 1)
+    ps = torch.arange(P, device="cuda").view(1, P)
+    pool_mask = torch.where(
+        ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda")
+    ).to(torch.bfloat16)
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    og = v4_attention_turbo(
+        qg,
+        kg,
+        kg,
+        sink=sg,
+        swa_window=_SWA,
+        additive_mask=pool_mask,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+        hca_local_seqlen=S,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2)
+    local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32)
+    full_mask = torch.cat([local_mask, pool_mask.float()], dim=1)
+    of = eager_v4_attention(
+        qf,
+        kf,
+        kf,
+        sink=sf,
+        swa_window=0,
+        additive_mask=full_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[turbo V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude
+    # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards.
+    _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
+
+
+@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v))
+def test_v4_turbo_csa_matches_eager(B, H, S, sink_init):
+    """cr=4 CSA: turbo fwd+bwd vs eager_v4_csa_attention."""
+    P = max(S // 4, 1)
+    K = min(128, P)
+    g = torch.Generator(device="cuda").manual_seed(3)
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16)
+    sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None
+
+    topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32)
+    drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125
+    topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs)
+    idx = topk_idxs.clamp(0, P - 1).long()
+    bidx = torch.arange(B, device="cuda").view(B, 1, 1)
+    sparse_mask_inf = torch.where(
+        topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda")
+    )
+
+    qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False)
+    sg = _leaf(sink, fp32=False) if sink is not None else None
+    klg = latg.unsqueeze(1).expand(B, H, S, D)
+    og = v4_csa_attention_turbo(
+        qg,
+        klg,
+        klg,
+        poolg,
+        topk_idxs=topk_idxs,
+        sink=sg,
+        swa_window=_SWA,
+        attn_dropout=0.0,
+        training=True,
+        scale=scale,
+    )
+    og.backward(do)
+
+    qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True)
+    sf = _leaf(sink, fp32=True) if sink is not None else None
+    klf = latf.unsqueeze(1).expand(B, H, S, D)
+    gathered = poolf[bidx, idx]
+    of = eager_v4_csa_attention(
+        qf,
+        klf,
+        klf,
+        gathered,
+        sink=sf,
+        swa_window=_SWA,
+        sparse_mask=sparse_mask_inf.float(),
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    of.backward(do.float())
+
+    print(f"\n[turbo V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True)
+    assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16
+    _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3)
+    _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3)
+    # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude
+    # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards.
+    _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    _check("dpool", poolg.grad, poolf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3)
+    if sink is not None:
+        _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2)
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py
new file mode 100644
index 000000000..302ed8f38
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py
@@ -0,0 +1,439 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-4 P25 G24 — `v4_attention_v1` Triton BWD equivalence to autograd-on-eager.
+
+Asserts that gradients (``dq``, ``dk``, ``dv``, ``dsink``) returned by
+:func:`v4_attention_v1`'s autograd Function match the gradients computed
+by autograd-on-:func:`eager_v4_attention` within the plan-4 tolerance
+budget across the same shape envelope as G23 (V4-Flash + V4-Pro,
+``compress_ratio ∈ {0, 128}``, fp32 + bf16, sink_on / sink_off, MQA /
+MHA layouts).
+
+The sink gradient is asserted **per-head**: ``dsink`` is shape ``[H]``
+and each head's gradient must match independently.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("v4_attention_v1 Triton kernel requires CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (  # noqa: E402
+    eager_v4_attention,
+    v4_attention_v1,
+)
+
+# ---------------------------------------------------------------------------
+# Shape envelope (same as G23)
+# ---------------------------------------------------------------------------
+
+
+# See ``test_v4_p25_v4_attention_fwd._BASE_SHAPES`` for the fast vs
+# release tier rationale (G28 plan-4 release-tier shape gate).
+_BASE_SHAPES = [
+    ("v4_flash_small", 1, 8, 64, 64, 32),
+    ("v4_pro_small", 1, 4, 64, 64, 32),
+    pytest.param(
+        "v4_flash_release",
+        1,
+        64,
+        1024,
+        512,
+        128,
+        marks=pytest.mark.slow,
+    ),
+    pytest.param(
+        "v4_pro_release",
+        1,
+        128,
+        512,
+        512,
+        128,
+        marks=pytest.mark.slow,
+    ),
+]
+_DTYPES = [torch.float32, torch.bfloat16]
+_SINK_MODES = [True, False]
+_KV_LAYOUTS = ["mqa", "mha"]
+
+
+def _is_release_tier(variant: str) -> bool:
+    """Release-tier shapes are tagged by name (``*_release``)."""
+    return variant.endswith("_release")
+
+
+def _bwd_tol(dtype: torch.dtype, *, release: bool = False) -> dict:
+    """Plan-4 BWD tolerance budget.
+
+    Release-tier ``head_dim=512`` causes:
+    * ~sqrt(8) ≈ 2.8x more matmul-accumulation noise than the
+      ``head_dim=64`` fast tier;
+    * non-deterministic ``tl.atomic_add`` contributions to ``dk / dv``
+      (sliding-window cells receive ~SWA contributions, summed in a
+      non-deterministic thread order) which adds another ~``sqrt(SWA)``
+      bf16 jitter on top of matmul noise.
+
+    Empirically the worst-case outlier at V4-Flash / V4-Pro release
+    sits around ``0.15-0.18`` for bf16 ``dk``; we set ``atol=2e-1`` to
+    absorb the long tail while keeping fast-tier shapes tight.
+    """
+    if dtype == torch.float32:
+        return {"atol": 1e-4, "rtol": 1e-4}
+    if dtype == torch.bfloat16:
+        return {"atol": 2e-1, "rtol": 2e-1} if release else {"atol": 5e-2, "rtol": 5e-2}
+    raise ValueError(f"unsupported dtype {dtype!r}")
+
+
+def _sink_tol(dtype: torch.dtype, *, release: bool = False) -> dict:
+    """Per-head sink gradient tolerance.
+
+    ``dsink[h] = sum_{b, t} dprobs_at_sink_column[b, h, t]`` — the
+    reduction is over ``B * Sq`` softmax-derivative terms, so cumulative
+    bf16 rounding scales linearly with ``Sq``. Release-tier ``Sq=1024``
+    is 16x the fast-tier ``Sq=64``; loosen the bf16 budget accordingly.
+    """
+    if dtype == torch.float32:
+        return {"atol": 1e-4, "rtol": 1e-4}
+    if dtype == torch.bfloat16:
+        return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 5e-3, "rtol": 5e-3}
+    raise ValueError(f"unsupported dtype {dtype!r}")
+
+
+def _build_inputs(
+    *,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    sink_on: bool,
+    dtype: torch.dtype,
+    kv_layout: str,
+    use_hca: bool = False,
+    seed: int = 4321,
+):
+    """Build (q, k, v, sink, additive_mask, swa_window) with leaves on requires_grad.
+
+    The dense path uses ``swa_window > 0, additive_mask=None``. The HCA
+    path concatenates a tiny pool to the keys / values and supplies a
+    pre-built ``[Sq, Sq+P]`` joint additive mask.
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    device = "cuda"
+    K_H = 1 if kv_layout == "mqa" else H
+
+    q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True)
+    if use_hca:
+        P = 4
+        k_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+        v_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+        pool_k = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype)
+        pool_v = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype)
+        k = torch.cat([k_local, pool_k], dim=2).requires_grad_(True)
+        v = torch.cat([v_local, pool_v], dim=2).requires_grad_(True)
+        # Joint additive mask: local SWA-causal + pool causal-on-stride
+        local_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype)
+        ratio = 4
+        t = torch.arange(S, device=device).unsqueeze(1)
+        s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1
+        pool_mask = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype)
+        full_mask = torch.cat([local_mask, pool_mask], dim=-1)
+        kernel_swa_window = swa_window
+        eager_swa_window = 0
+        eager_mask: Optional[torch.Tensor] = full_mask
+        kernel_mask: Optional[torch.Tensor] = pool_mask
+        hca_local_seqlen = S
+    else:
+        k = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True)
+        v = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True)
+        # Eager and kernel both use the in-kernel SWA mask (eager via
+        # pre-built additive_mask, kernel via swa_window).
+        eager_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype)
+        kernel_mask = None
+        eager_swa_window = 0
+        kernel_swa_window = swa_window
+        hca_local_seqlen = 0
+
+    sink = (
+        torch.randn(H, generator=g, device=device, dtype=torch.float32, requires_grad=True) * 0.1
+        if sink_on
+        else None
+    )
+    return dict(
+        q=q,
+        k=k,
+        v=v,
+        sink=sink,
+        eager_mask=eager_mask,
+        eager_swa_window=eager_swa_window,
+        kernel_mask=kernel_mask,
+        kernel_swa_window=kernel_swa_window,
+        hca_local_seqlen=hca_local_seqlen,
+    )
+
+
+def _grads_from(model_out: torch.Tensor, *leaves: torch.Tensor) -> list[Optional[torch.Tensor]]:
+    """Run ``model_out.sum().backward()`` and pull ``leaf.grad`` into a list.
+
+    Cleans up ``leaf.grad`` on each leaf so the caller can re-use the
+    leaves for a second autograd pass.
+    """
+    real_leaves = [lf for lf in leaves if lf is not None]
+    grads_out = torch.ones_like(model_out)
+    grads = torch.autograd.grad(
+        outputs=model_out,
+        inputs=real_leaves,
+        grad_outputs=grads_out,
+        retain_graph=False,
+        create_graph=False,
+        allow_unused=False,
+    )
+    # Re-pad with None for missing leaves
+    out: list[Optional[torch.Tensor]] = []
+    j = 0
+    for lf in leaves:
+        if lf is None:
+            out.append(None)
+        else:
+            out.append(grads[j])
+            j += 1
+    return out
+
+
+# ---------------------------------------------------------------------------
+# G24 — dense (compress_ratio == 0)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS)
+def test_g24_dense_bwd_matches_eager(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+    kv_layout: str,
+):
+    """Dense path: kernel BWD matches autograd-on-eager."""
+    # Build two independent leaf sets so the two BWD passes do not
+    # interfere with each other's ``grad`` slots.
+    ref_inp = _build_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+        use_hca=False,
+        seed=4321,
+    )
+    cand_inp = _build_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+        use_hca=False,
+        seed=4321,
+    )
+
+    scale = 1.0 / math.sqrt(D)
+
+    out_ref = eager_v4_attention(
+        ref_inp["q"],
+        ref_inp["k"],
+        ref_inp["v"],
+        sink=ref_inp["sink"],
+        swa_window=ref_inp["eager_swa_window"],
+        additive_mask=ref_inp["eager_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    dq_ref, dk_ref, dv_ref, dsink_ref = _grads_from(
+        out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"]
+    )
+
+    out_cand = v4_attention_v1(
+        cand_inp["q"],
+        cand_inp["k"],
+        cand_inp["v"],
+        sink=cand_inp["sink"],
+        swa_window=cand_inp["kernel_swa_window"],
+        additive_mask=cand_inp["kernel_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+        hca_local_seqlen=cand_inp["hca_local_seqlen"],
+    )
+    dq_cand, dk_cand, dv_cand, dsink_cand = _grads_from(
+        out_cand, cand_inp["q"], cand_inp["k"], cand_inp["v"], cand_inp["sink"]
+    )
+
+    release = _is_release_tier(variant)
+    tol = _bwd_tol(dtype, release=release)
+    torch.testing.assert_close(dq_cand, dq_ref, **tol)
+    torch.testing.assert_close(dk_cand, dk_ref, **tol)
+    torch.testing.assert_close(dv_cand, dv_ref, **tol)
+    if sink_on:
+        # Per-head sink gradient: shape [H]
+        assert dsink_ref.shape == dsink_cand.shape == (H,)
+        torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release))
+    else:
+        assert dsink_ref is None and dsink_cand is None
+
+
+# ---------------------------------------------------------------------------
+# G24 — HCA (compress_ratio == 128)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS)
+def test_g24_hca_style_bwd_matches_eager(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+    kv_layout: str,
+):
+    """HCA path: caller-supplied additive mask, kernel BWD matches autograd-on-eager."""
+    ref_inp = _build_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+        use_hca=True,
+        seed=4321,
+    )
+    cand_inp = _build_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+        use_hca=True,
+        seed=4321,
+    )
+
+    scale = 1.0 / math.sqrt(D)
+
+    out_ref = eager_v4_attention(
+        ref_inp["q"],
+        ref_inp["k"],
+        ref_inp["v"],
+        sink=ref_inp["sink"],
+        swa_window=0,
+        additive_mask=ref_inp["eager_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    dq_ref, dk_ref, dv_ref, dsink_ref = _grads_from(
+        out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"]
+    )
+
+    out_cand = v4_attention_v1(
+        cand_inp["q"],
+        cand_inp["k"],
+        cand_inp["v"],
+        sink=cand_inp["sink"],
+        swa_window=cand_inp["kernel_swa_window"],
+        additive_mask=cand_inp["kernel_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+        hca_local_seqlen=cand_inp["hca_local_seqlen"],
+    )
+    dq_cand, dk_cand, dv_cand, dsink_cand = _grads_from(
+        out_cand, cand_inp["q"], cand_inp["k"], cand_inp["v"], cand_inp["sink"]
+    )
+
+    release = _is_release_tier(variant)
+    tol = _bwd_tol(dtype, release=release)
+    torch.testing.assert_close(dq_cand, dq_ref, **tol)
+    torch.testing.assert_close(dk_cand, dk_ref, **tol)
+    torch.testing.assert_close(dv_cand, dv_ref, **tol)
+    if sink_on:
+        assert dsink_ref.shape == dsink_cand.shape == (H,)
+        torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release))
+    else:
+        assert dsink_ref is None and dsink_cand is None
+
+
+# ---------------------------------------------------------------------------
+# Sanity: dq, dk, dv finite-ness (no NaN / Inf)
+# ---------------------------------------------------------------------------
+
+
+def test_g24_no_nan_in_grads_dense_fp32():
+    """Sanity: kernel BWD does not introduce NaN / Inf for a typical fp32 dense case."""
+    inp = _build_inputs(
+        B=1,
+        H=4,
+        S=64,
+        D=64,
+        swa_window=32,
+        sink_on=True,
+        dtype=torch.float32,
+        kv_layout="mha",
+        use_hca=False,
+        seed=4321,
+    )
+    out_cand = v4_attention_v1(
+        inp["q"],
+        inp["k"],
+        inp["v"],
+        sink=inp["sink"],
+        swa_window=inp["kernel_swa_window"],
+        additive_mask=inp["kernel_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=1.0 / math.sqrt(64),
+        hca_local_seqlen=inp["hca_local_seqlen"],
+    )
+    grads = _grads_from(out_cand, inp["q"], inp["k"], inp["v"], inp["sink"])
+    for name, g in zip(("q", "k", "v", "sink"), grads):
+        if g is None:
+            continue
+        assert torch.isfinite(g).all(), f"d{name} contains NaN/Inf"
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py
new file mode 100644
index 000000000..d28d0cc59
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py
@@ -0,0 +1,409 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-4 P25 G23 — `v4_attention_v1` Triton FWD equivalence to eager.
+
+Asserts that :func:`v4_attention_v1` (Triton kernel from
+``primus...transformer.v4_attention_kernels.v4_attention``) produces
+forward output equal to :func:`eager_v4_attention` within the plan-4
+tolerance budget across:
+
+* V4-Flash and V4-Pro shape envelopes (head_dim=512, H ∈ {64, 128});
+* ``compress_ratio ∈ {0, 128}`` — dense + SWA + sink (no bias) and HCA
+  (joint-softmax additive bias, no in-kernel SWA);
+* fp32 and bf16 inputs;
+* ``sink ∈ {None, learned [H]}``;
+* MQA (``K_H == 1``) and MHA (``K_H == HQ``) layouts.
+
+The test is GPU-only (Triton requires CUDA / HIP); CPU runs are
+``pytest.skip``-ed at module collection time.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("v4_attention_v1 Triton kernel requires CUDA / HIP", allow_module_level=True)
+
+# Triton import (must be importable in the env).
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.sliding_window_kv import (  # noqa: E402
+    sliding_window_causal_mask,
+)
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (  # noqa: E402
+    eager_v4_attention,
+    v4_attention_v1,
+)
+
+# ---------------------------------------------------------------------------
+# Shape envelope (V4-Flash / V4-Pro)
+# ---------------------------------------------------------------------------
+
+
+# Two tiers of shapes are exposed to the parametrise decorator below:
+#
+# * Fast tier — toy ``head_dim=64`` / ``H ∈ {4, 8}`` / ``S=64`` shapes
+#   that exercise every code path in the kernel (in-kernel SWA mask,
+#   caller-supplied additive_mask, MQA/MHA, sink) in milliseconds. Run
+#   on every ``pytest`` invocation.
+#
+# * Release tier (``pytest.mark.slow``) — production V4 dimensions
+#   (``head_dim=512``, real ``H``, real ``swa_window``) calibrated so
+#   the eager fp32 reference fits MI355X HBM. The release tier is the
+#   plan-4 G28 gate that empirically confirms kernel correctness at the
+#   exact ``head_dim`` that plan-4 exists to solve. ``S`` is calibrated
+#   per variant: V4-Flash @ S=1024 / V4-Pro @ S=512 keep peak fp32
+#   memory below ~1 GiB per test even for the MHA layout.
+#
+# The full ``S=4096`` smoke is owned by G30 (P27 smoke gate); these
+# unit tests intentionally stay below that to keep the per-test cost
+# in the ``-m slow`` budget.
+_BASE_SHAPES = [
+    # (variant, B, H, S, head_dim, swa_window)
+    ("v4_flash_small", 1, 8, 64, 64, 32),
+    ("v4_pro_small", 1, 4, 64, 64, 32),
+    pytest.param(
+        "v4_flash_release",
+        1,
+        64,
+        1024,
+        512,
+        128,
+        marks=pytest.mark.slow,
+    ),
+    pytest.param(
+        "v4_pro_release",
+        1,
+        128,
+        512,
+        512,
+        128,
+        marks=pytest.mark.slow,
+    ),
+]
+_DTYPES = [torch.float32, torch.bfloat16]
+_SINK_MODES = [True, False]
+_KV_LAYOUTS = ["mqa", "mha"]  # K_H == 1 vs K_H == HQ
+
+
+def _is_release_tier(variant: str) -> bool:
+    """Release-tier shapes are tagged by name (``*_release``)."""
+    return variant.endswith("_release")
+
+
+def _fwd_tol(dtype: torch.dtype, *, release: bool = False) -> dict:
+    if dtype == torch.float32:
+        return {"atol": 1e-4, "rtol": 1e-4}
+    if dtype == torch.bfloat16:
+        # Release-tier ``head_dim=512`` accumulates ~8x more terms than
+        # the ``head_dim=64`` fast tier; bf16 long-tail rounding scales
+        # as ~sqrt(N). Empirically the worst-case outlier sits between
+        # 2.0e-2 and 5.0e-2 at release-tier dims; we loosen by ~2.5x.
+        return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 2e-2, "rtol": 2e-2}
+    raise ValueError(f"unsupported dtype {dtype!r}")
+
+
+def _make_dense_inputs(
+    *,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    sink_on: bool,
+    dtype: torch.dtype,
+    kv_layout: str,
+    seed: int = 1234,
+):
+    """Build inputs for a dense (compress_ratio=0) test case.
+
+    The kernel is invoked with ``swa_window > 0, additive_mask=None``;
+    the eager reference is invoked with the equivalent ``additive_mask``
+    pre-built (so both go through identical mask math).
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    device = "cuda"
+    K_H = 1 if kv_layout == "mqa" else H
+
+    q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype)
+    k = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+    v = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+    sink = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 if sink_on else None
+    return dict(B=B, H=H, S=S, D=D, swa_window=swa_window, q=q, k=k, v=v, sink=sink)
+
+
+def _make_hca_inputs(
+    *,
+    B: int,
+    H: int,
+    S: int,
+    P: int,
+    D: int,
+    swa_window: int,
+    sink_on: bool,
+    dtype: torch.dtype,
+    kv_layout: str,
+    seed: int = 1234,
+):
+    """Build inputs for an HCA (compress_ratio=128) test case.
+
+    HCA pre-concatenates a length-``P`` compressed-pool key/value to the
+    length-``S`` local key/value, and the caller pre-builds the joint
+    additive mask so the kernel does NOT apply SWA / causal in-kernel.
+    """
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    device = "cuda"
+    K_H = 1 if kv_layout == "mqa" else H
+
+    q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype)
+    k_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+    v_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype)
+    pool_k = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype)
+    pool_v = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype)
+    k_full = torch.cat([k_local, pool_k], dim=2)
+    v_full = torch.cat([v_local, pool_v], dim=2)
+
+    # Local SWA-causal mask: [S, S]
+    local_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype)
+    # Pool causal-on-stride mask: pool[s] is visible at query t iff
+    # (s+1)*ratio - 1 <= t. Use a fixed ratio of 4 for the test.
+    ratio = 4
+    t = torch.arange(S, device=device).unsqueeze(1)
+    s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1
+    pool_mask = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype)
+    full_mask = torch.cat([local_mask, pool_mask], dim=-1)  # [S, S+P]
+
+    sink = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 if sink_on else None
+    return dict(
+        B=B,
+        H=H,
+        S=S,
+        P=P,
+        D=D,
+        q=q,
+        k=k_full,
+        v=v_full,
+        sink=sink,
+        full_mask=full_mask,
+        pool_mask=pool_mask,
+    )
+
+
+# ---------------------------------------------------------------------------
+# G23 — dense (compress_ratio == 0)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS)
+def test_g23_dense_fwd_matches_eager(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+    kv_layout: str,
+):
+    """Dense path: kernel uses in-kernel SWA-causal mask + optional sink."""
+    toy = _make_dense_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+    )
+
+    scale = 1.0 / math.sqrt(D)
+
+    # Reference: eager_v4_attention with pre-built SWA additive_mask
+    eager_mask = sliding_window_causal_mask(S, swa_window, device=toy["q"].device, dtype=dtype)
+    out_ref = eager_v4_attention(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=0,
+        additive_mask=eager_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+
+    # Candidate: v4_attention_v1 with swa_window > 0, additive_mask=None
+    out_cand = v4_attention_v1(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=swa_window,
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+
+    assert out_ref.shape == out_cand.shape == toy["q"].shape
+    assert out_ref.dtype == out_cand.dtype == dtype
+    torch.testing.assert_close(
+        out_cand,
+        out_ref,
+        **_fwd_tol(dtype, release=_is_release_tier(variant)),
+    )
+
+
+# ---------------------------------------------------------------------------
+# G23 — HCA (compress_ratio == 128) — caller-supplied additive mask
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS)
+def test_g23_hca_style_fwd_matches_eager(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+    kv_layout: str,
+):
+    """HCA path: caller pre-concatenates pool keys + supplies joint additive mask."""
+    P = 4  # tiny pool — exercises the additive_mask branch w/o full HCA setup
+    toy = _make_hca_inputs(
+        B=B,
+        H=H,
+        S=S,
+        P=P,
+        D=D,
+        swa_window=swa_window,
+        sink_on=sink_on,
+        dtype=dtype,
+        kv_layout=kv_layout,
+    )
+
+    scale = 1.0 / math.sqrt(D)
+
+    out_ref = eager_v4_attention(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=0,
+        additive_mask=toy["full_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+
+    out_cand = v4_attention_v1(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=swa_window,
+        additive_mask=toy["pool_mask"],
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+        hca_local_seqlen=S,
+    )
+
+    assert out_ref.shape == out_cand.shape == toy["q"].shape
+    assert out_ref.dtype == out_cand.dtype == dtype
+    torch.testing.assert_close(
+        out_cand,
+        out_ref,
+        **_fwd_tol(dtype, release=_is_release_tier(variant)),
+    )
+
+
+# ---------------------------------------------------------------------------
+# G25 — determinism with attn_dropout=0.0
+# ---------------------------------------------------------------------------
+
+
+def test_g25_determinism_fp32_mha():
+    """Repeated FWD calls with the same inputs produce bit-identical output (fp32 / MHA)."""
+    toy = _make_dense_inputs(
+        B=1,
+        H=4,
+        S=64,
+        D=64,
+        swa_window=32,
+        sink_on=True,
+        dtype=torch.float32,
+        kv_layout="mha",
+    )
+    scale = 1.0 / math.sqrt(toy["D"])
+
+    out_a = v4_attention_v1(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=toy["swa_window"],
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    out_b = v4_attention_v1(
+        toy["q"],
+        toy["k"],
+        toy["v"],
+        sink=toy["sink"],
+        swa_window=toy["swa_window"],
+        additive_mask=None,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    assert torch.equal(out_a, out_b), "v4_attention_v1 FWD is non-deterministic at fp32 / MHA"
+
+
+def test_g25_dropout_with_training_is_rejected():
+    """``attn_dropout > 0`` with ``training=True`` raises (kernel does not implement dropout)."""
+    toy = _make_dense_inputs(
+        B=1,
+        H=4,
+        S=64,
+        D=64,
+        swa_window=32,
+        sink_on=False,
+        dtype=torch.float32,
+        kv_layout="mha",
+    )
+    scale = 1.0 / math.sqrt(toy["D"])
+    with pytest.raises(NotImplementedError, match="dropout"):
+        v4_attention_v1(
+            toy["q"],
+            toy["k"],
+            toy["v"],
+            sink=None,
+            swa_window=toy["swa_window"],
+            additive_mask=None,
+            attn_dropout=0.1,
+            training=True,
+            scale=scale,
+        )
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py
new file mode 100644
index 000000000..0f2d8c371
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py
@@ -0,0 +1,280 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Plan-5 P31 — CSA in-kernel top-K gather equivalence.
+
+The P26 CSA kernel consumed a materialised ``gathered`` tensor. P31 moves
+that gather into the Triton kernel and scatters gradients directly back
+to the compressed pool. These tests compare the new pool/topk API against
+the eager reference plus PyTorch gather autograd.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+if not torch.cuda.is_available():
+    pytest.skip("v4_csa_attention_v0 Triton kernel requires CUDA / HIP", allow_module_level=True)
+
+pytest.importorskip("triton", reason="Triton not installed")
+
+from primus.backends.megatron.core.transformer.v4_attention_kernels import (  # noqa: E402
+    eager_v4_csa_attention,
+    v4_csa_attention_v1,
+)
+
+_SHAPES = [
+    ("fast", 1, 4, 32, 64, 32, 16, 16),
+    pytest.param("release_head_dim", 1, 16, 128, 512, 128, 64, 64, marks=pytest.mark.slow),
+]
+_DTYPES = [torch.float32, torch.bfloat16]
+_SINK_MODES = [True, False]
+
+
+def _tol(dtype: torch.dtype, *, release: bool = False) -> dict:
+    if dtype == torch.float32:
+        return {"atol": 1e-4, "rtol": 1e-4}
+    if dtype == torch.bfloat16:
+        return {"atol": 2e-1, "rtol": 2e-1} if release else {"atol": 5e-2, "rtol": 5e-2}
+    raise ValueError(f"unsupported dtype {dtype!r}")
+
+
+def _sink_tol(dtype: torch.dtype, *, release: bool = False) -> dict:
+    if dtype == torch.float32:
+        return {"atol": 1e-4, "rtol": 1e-4}
+    if dtype == torch.bfloat16:
+        return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 5e-3, "rtol": 5e-3}
+    raise ValueError(f"unsupported dtype {dtype!r}")
+
+
+def _make_inputs(
+    *,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    P: int,
+    K_topk: int,
+    sink_on: bool,
+    dtype: torch.dtype,
+    requires_grad: bool,
+    seed: int = 20260509,
+):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    device = "cuda"
+
+    q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad)
+    k_local = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad)
+    v_local = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad)
+    pool = torch.randn(B, P, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad)
+
+    topk_idxs = torch.randint(0, P, (B, S, K_topk), generator=g, device=device, dtype=torch.int64)
+    if K_topk >= 4:
+        # Exercise both invalid slots and duplicate slots, which are the
+        # two cases the in-kernel scatter-add must handle correctly.
+        topk_idxs[..., 0] = -1
+        topk_idxs[..., 1] = topk_idxs[..., 2]
+
+    sink = None
+    if sink_on:
+        sink_raw = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1
+        sink = sink_raw.detach().clone().requires_grad_(requires_grad)
+
+    return dict(q=q, k_local=k_local, v_local=v_local, pool=pool, topk_idxs=topk_idxs, sink=sink)
+
+
+def _gather_from_pool(pool: torch.Tensor, topk_idxs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    B, P, D = pool.shape
+    _, S, K = topk_idxs.shape
+    valid = topk_idxs >= 0
+    safe_idx = topk_idxs.clamp(min=0)
+    gathered = torch.gather(
+        pool.unsqueeze(1).expand(B, S, P, D),
+        dim=2,
+        index=safe_idx.unsqueeze(-1).expand(B, S, K, D),
+    )
+    gathered = gathered * valid.unsqueeze(-1).to(pool.dtype)
+    sparse_mask = torch.where(valid, 0.0, float("-inf")).to(pool.dtype)
+    return gathered, sparse_mask
+
+
+def _grads_from(model_out: torch.Tensor, *leaves: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]:
+    real_leaves = [leaf for leaf in leaves if leaf is not None]
+    grads = torch.autograd.grad(
+        outputs=model_out,
+        inputs=real_leaves,
+        grad_outputs=torch.ones_like(model_out),
+        retain_graph=False,
+        create_graph=False,
+        allow_unused=False,
+    )
+    out: list[Optional[torch.Tensor]] = []
+    idx = 0
+    for leaf in leaves:
+        if leaf is None:
+            out.append(None)
+        else:
+            out.append(grads[idx])
+            idx += 1
+    return out
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,P,K_topk,swa_window", _SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+def test_p31_csa_pool_fwd_matches_gathered_reference(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    P: int,
+    K_topk: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+):
+    inp = _make_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        P=P,
+        K_topk=K_topk,
+        sink_on=sink_on,
+        dtype=dtype,
+        requires_grad=False,
+    )
+    gathered, sparse_mask = _gather_from_pool(inp["pool"], inp["topk_idxs"])
+    scale = 1.0 / math.sqrt(D)
+
+    out_ref = eager_v4_csa_attention(
+        inp["q"],
+        inp["k_local"],
+        inp["v_local"],
+        gathered,
+        sink=inp["sink"],
+        swa_window=swa_window,
+        sparse_mask=sparse_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    out_cand = v4_csa_attention_v1(
+        inp["q"],
+        inp["k_local"],
+        inp["v_local"],
+        inp["pool"],
+        topk_idxs=inp["topk_idxs"],
+        sink=inp["sink"],
+        swa_window=swa_window,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+
+    torch.testing.assert_close(out_cand, out_ref, **_tol(dtype, release=variant != "fast"))
+
+
+@pytest.mark.parametrize("variant,B,H,S,D,P,K_topk,swa_window", _SHAPES)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1])
+@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"])
+def test_p31_csa_pool_bwd_matches_gathered_reference(
+    variant: str,
+    B: int,
+    H: int,
+    S: int,
+    D: int,
+    P: int,
+    K_topk: int,
+    swa_window: int,
+    dtype: torch.dtype,
+    sink_on: bool,
+):
+    ref_inp = _make_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        P=P,
+        K_topk=K_topk,
+        sink_on=sink_on,
+        dtype=dtype,
+        requires_grad=True,
+        seed=4242,
+    )
+    cand_inp = _make_inputs(
+        B=B,
+        H=H,
+        S=S,
+        D=D,
+        P=P,
+        K_topk=K_topk,
+        sink_on=sink_on,
+        dtype=dtype,
+        requires_grad=True,
+        seed=4242,
+    )
+    scale = 1.0 / math.sqrt(D)
+
+    gathered, sparse_mask = _gather_from_pool(ref_inp["pool"], ref_inp["topk_idxs"])
+    out_ref = eager_v4_csa_attention(
+        ref_inp["q"],
+        ref_inp["k_local"],
+        ref_inp["v_local"],
+        gathered,
+        sink=ref_inp["sink"],
+        swa_window=swa_window,
+        sparse_mask=sparse_mask,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    dq_ref, dkl_ref, dvl_ref, dpool_ref, dsink_ref = _grads_from(
+        out_ref,
+        ref_inp["q"],
+        ref_inp["k_local"],
+        ref_inp["v_local"],
+        ref_inp["pool"],
+        ref_inp["sink"],
+    )
+
+    out_cand = v4_csa_attention_v1(
+        cand_inp["q"],
+        cand_inp["k_local"],
+        cand_inp["v_local"],
+        cand_inp["pool"],
+        topk_idxs=cand_inp["topk_idxs"],
+        sink=cand_inp["sink"],
+        swa_window=swa_window,
+        attn_dropout=0.0,
+        training=False,
+        scale=scale,
+    )
+    dq_cand, dkl_cand, dvl_cand, dpool_cand, dsink_cand = _grads_from(
+        out_cand,
+        cand_inp["q"],
+        cand_inp["k_local"],
+        cand_inp["v_local"],
+        cand_inp["pool"],
+        cand_inp["sink"],
+    )
+
+    release = variant != "fast"
+    tol = _tol(dtype, release=release)
+    torch.testing.assert_close(dq_cand, dq_ref, **tol)
+    torch.testing.assert_close(dkl_cand, dkl_ref, **tol)
+    torch.testing.assert_close(dvl_cand, dvl_ref, **tol)
+    torch.testing.assert_close(dpool_cand, dpool_ref, **tol)
+    if sink_on:
+        torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release))
+    else:
+        assert dsink_ref is None and dsink_cand is None
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py
new file mode 100644
index 000000000..5254a2b95
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py
@@ -0,0 +1,260 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""V4-Flash / V4-Pro attention shape fixtures (plan-4 P24).
+
+This module exposes the **single source-of-truth** shape tables every
+plan-4 test parametrises over. The canonical config knobs come from:
+
+* ``primus/configs/models/megatron/deepseek_v4_flash.yaml``
+* ``primus/configs/models/megatron/deepseek_v4_pro.yaml``
+* ``primus/configs/models/megatron/deepseek_v4_base.yaml``
+
+Each variant exposes three sequence-length tiers:
+
+* ``small`` (S=128) — fast CI tier; runs in milliseconds on CPU
+* ``medium`` (S=512) — moderate tier; runs in seconds on a single GPU
+* ``large`` (S=4096) — release tier; gated behind ``pytest.mark.slow``
+
+Each shape is parametrised by the ``compress_ratio`` of the layer
+under test (``0`` dense + SWA + sink, ``128`` HCA, ``4`` CSA). Pool
+size ``P = S // compress_ratio`` for compressed branches; the indexer
+top-K (``K``) follows the per-variant ``index_topk`` knob (clamped to
+``P`` when ``P < K``, mirroring how the indexer behaves at small
+sequences).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Optional, Tuple
+
+SeqTier = str  # "small" | "medium" | "large"
+
+_SEQ_TIERS: Tuple[SeqTier, ...] = ("small", "medium", "large")
+_SEQ_LENGTHS = {
+    "small": 128,
+    "medium": 512,
+    "large": 4096,
+}
+
+
+@dataclass(frozen=True)
+class V4AttnShape:
+    """Single-shape parametrisation for a V4 attention layer test.
+
+    Attributes follow the V4 attention class:
+
+    * ``B`` — batch size
+    * ``S`` — sequence length (post-RoPE; raw token count)
+    * ``H`` — number of attention heads (per-rank; tests run TP=1)
+    * ``head_dim`` — per-head ``kv_channels``
+    * ``q_pe_dim`` — partial-RoPE rotary head dim (channels with RoPE
+      applied; the remaining ``head_dim - q_pe_dim`` channels are
+      "NOPE")
+    * ``attn_sliding_window`` — SWA window for dense / CSA local branch
+    * ``sink`` — ``True`` enables the per-head ``[H]`` learned softmax
+      sink (matches V4 yaml)
+    * ``compress_ratio`` — ``0`` dense, ``4`` CSA, ``128`` HCA
+    * ``P`` — compressed pool size (``S // compress_ratio`` when
+      ``compress_ratio > 0``; ``0`` for dense)
+    * ``K`` — indexer top-K (``min(index_topk, P)`` for CSA;
+      meaningless / ``0`` for dense / HCA but kept on the dataclass
+      for API uniformity)
+    * ``variant`` — display name ("v4_flash" / "v4_pro") for test ids
+    """
+
+    variant: str
+    B: int
+    S: int
+    H: int
+    head_dim: int
+    q_pe_dim: int
+    attn_sliding_window: int
+    sink: bool
+    compress_ratio: int
+    P: int
+    K: int
+
+    @property
+    def Sk(self) -> int:
+        """Effective key axis length per branch.
+
+        * dense (cr=0): ``S`` local keys
+        * HCA (cr=128): ``S + P`` (local + compressed pool)
+        * CSA (cr=4): ``S`` local keys (the sparse top-K branch is
+          handled separately via ``gathered``)
+        """
+        if self.compress_ratio == 128:
+            return self.S + self.P
+        return self.S
+
+    def shape_id(self) -> str:
+        """A short human-readable test-id for pytest parametrisation."""
+        return (
+            f"{self.variant}-cr{self.compress_ratio}-S{self.S}"
+            f"-H{self.H}-D{self.head_dim}"
+            f"-sink{int(self.sink)}-w{self.attn_sliding_window}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# V4-Flash and V4-Pro variant tables
+# ---------------------------------------------------------------------------
+
+# Source: deepseek_v4_flash.yaml + deepseek_v4_base.yaml
+_V4_FLASH_DEFAULTS = dict(
+    H=64,
+    head_dim=512,
+    q_pe_dim=64,
+    attn_sliding_window=128,
+    sink=True,
+    index_topk=512,
+)
+
+# Source: deepseek_v4_pro.yaml + deepseek_v4_base.yaml
+_V4_PRO_DEFAULTS = dict(
+    H=128,
+    head_dim=512,
+    q_pe_dim=64,
+    attn_sliding_window=128,
+    sink=True,
+    index_topk=1024,
+)
+
+
+def _make_shape(
+    variant: str,
+    defaults: dict,
+    *,
+    compress_ratio: int,
+    seq: int,
+    batch: int,
+    sink: Optional[bool] = None,
+) -> V4AttnShape:
+    H = int(defaults["H"])
+    head_dim = int(defaults["head_dim"])
+    q_pe_dim = int(defaults["q_pe_dim"])
+    attn_sliding_window = int(defaults["attn_sliding_window"])
+    sink_default = bool(defaults["sink"])
+    index_topk = int(defaults["index_topk"])
+
+    if compress_ratio == 0:
+        P = 0
+        K = 0
+    elif compress_ratio == 4:
+        P = max(seq // compress_ratio, 1)
+        K = min(index_topk, P)
+    elif compress_ratio == 128:
+        P = max(seq // compress_ratio, 1)
+        K = 0  # not used by HCA
+    else:
+        raise ValueError(f"compress_ratio must be in {{0, 4, 128}}; got {compress_ratio}.")
+
+    return V4AttnShape(
+        variant=variant,
+        B=int(batch),
+        S=int(seq),
+        H=H,
+        head_dim=head_dim,
+        q_pe_dim=q_pe_dim,
+        attn_sliding_window=attn_sliding_window,
+        sink=sink_default if sink is None else bool(sink),
+        compress_ratio=int(compress_ratio),
+        P=int(P),
+        K=int(K),
+    )
+
+
+def v4_flash_shape(
+    *,
+    compress_ratio: int,
+    seq_tier: SeqTier = "small",
+    batch: int = 1,
+    sink: Optional[bool] = None,
+) -> V4AttnShape:
+    """Build a single V4-Flash :class:`V4AttnShape`."""
+    if seq_tier not in _SEQ_LENGTHS:
+        raise ValueError(f"seq_tier must be one of {_SEQ_TIERS}; got {seq_tier!r}.")
+    return _make_shape(
+        "v4_flash",
+        _V4_FLASH_DEFAULTS,
+        compress_ratio=compress_ratio,
+        seq=_SEQ_LENGTHS[seq_tier],
+        batch=batch,
+        sink=sink,
+    )
+
+
+def v4_pro_shape(
+    *,
+    compress_ratio: int,
+    seq_tier: SeqTier = "small",
+    batch: int = 1,
+    sink: Optional[bool] = None,
+) -> V4AttnShape:
+    """Build a single V4-Pro :class:`V4AttnShape`."""
+    if seq_tier not in _SEQ_LENGTHS:
+        raise ValueError(f"seq_tier must be one of {_SEQ_TIERS}; got {seq_tier!r}.")
+    return _make_shape(
+        "v4_pro",
+        _V4_PRO_DEFAULTS,
+        compress_ratio=compress_ratio,
+        seq=_SEQ_LENGTHS[seq_tier],
+        batch=batch,
+        sink=sink,
+    )
+
+
+def v4_attention_shape_grid(
+    *,
+    variants: Iterable[str] = ("v4_flash", "v4_pro"),
+    compress_ratios: Iterable[int] = (0, 4, 128),
+    seq_tiers: Iterable[SeqTier] = ("small",),
+    batch: int = 1,
+    sinks: Iterable[bool] = (True, False),
+) -> list[V4AttnShape]:
+    """Cartesian product of (variant × compress_ratio × seq_tier × sink).
+
+    Returns a flat list of :class:`V4AttnShape` suitable for direct
+    pytest parametrisation. Plan-4 tests typically pin to a small
+    subset to keep CI green:
+
+    * G22 (P24 refactor safety net): all three compress ratios at
+      ``small`` seq tier with both sink modes.
+    * G23 / G24 (P25 fwd / bwd): cr ∈ {0, 128}, seq ∈ {small, medium}.
+    * G26 / G27 (P26 CSA fwd / bwd): cr == 4, seq ∈ {small, medium}.
+    """
+    out: list[V4AttnShape] = []
+    builders = {
+        "v4_flash": v4_flash_shape,
+        "v4_pro": v4_pro_shape,
+    }
+    for variant in variants:
+        if variant not in builders:
+            raise ValueError(f"variant must be one of {tuple(builders)}; got {variant!r}.")
+        builder = builders[variant]
+        for cr in compress_ratios:
+            for tier in seq_tiers:
+                for sink in sinks:
+                    out.append(
+                        builder(
+                            compress_ratio=cr,
+                            seq_tier=tier,
+                            batch=batch,
+                            sink=sink,
+                        )
+                    )
+    return out
+
+
+__all__ = [
+    "SeqTier",
+    "V4AttnShape",
+    "v4_attention_shape_grid",
+    "v4_flash_shape",
+    "v4_pro_shape",
+]
diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py
new file mode 100644
index 000000000..f7e33a17a
--- /dev/null
+++ b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py
@@ -0,0 +1,221 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Forward + backward equivalence harness (plan-4 P24).
+
+Every plan-4 test that compares an eager-Python reference against a
+candidate (P25 / P26 Triton kernels, or the new P24 reference op
+itself) runs through :func:`compare_fwd_bwd`. The harness:
+
+* clones every input that ``requires_grad`` so reference and candidate
+  run on independent autograd graphs (one ``.grad`` accumulator each);
+* invokes the reference and candidate;
+* asserts forward output matches within the supplied tolerance;
+* calls ``.sum().backward()`` on each (with ``retain_graph=False``);
+* asserts every cloned-and-grad-bearing input's gradient matches
+  within the supplied tolerance;
+* on mismatch, prints a structured diff (max abs / max rel error per
+  leaf, with the eight worst entries by magnitude).
+
+Tolerance budget defaults follow the plan-4 P24 design:
+
+* ``fp32`` — ``atol=1e-5, rtol=1e-5`` for forward, same for backward.
+* ``bf16`` — ``atol=2e-2, rtol=2e-2`` for forward, ``atol=5e-2,
+  rtol=5e-2`` for backward (the looser bwd budget absorbs MQA
+  atomic-add reordering and other accumulator-order delta).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Iterable, Mapping, Optional
+
+import torch
+
+# ---------------------------------------------------------------------------
+# Tolerance budget
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Tol:
+    """Numerical tolerance pair for ``torch.allclose``-style comparisons."""
+
+    atol: float
+    rtol: float
+
+    def as_kwargs(self) -> dict:
+        return {"atol": self.atol, "rtol": self.rtol}
+
+
+FP32_TOL = Tol(atol=1e-5, rtol=1e-5)
+BF16_FWD_TOL = Tol(atol=2e-2, rtol=2e-2)
+BF16_BWD_TOL = Tol(atol=5e-2, rtol=5e-2)
+
+
+def default_tols(dtype: torch.dtype) -> tuple[Tol, Tol]:
+    """Return ``(fwd_tol, bwd_tol)`` for the supplied input dtype."""
+    if dtype == torch.float32:
+        return FP32_TOL, FP32_TOL
+    if dtype == torch.bfloat16:
+        return BF16_FWD_TOL, BF16_BWD_TOL
+    if dtype == torch.float16:
+        # Treat fp16 like bf16 for tolerance purposes (slightly tighter
+        # mantissa, slightly tighter exponent — same ballpark).
+        return BF16_FWD_TOL, BF16_BWD_TOL
+    raise ValueError(f"No default tolerance preset for dtype {dtype!r}.")
+
+
+# ---------------------------------------------------------------------------
+# Input cloning + gradient comparison helpers
+# ---------------------------------------------------------------------------
+
+
+def _clone_for_grad(inputs: Mapping[str, Any]) -> dict:
+    """Return a shallow-copied mapping with leaf tensors deep-cloned.
+
+    Tensors that ``requires_grad`` are replaced with detached clones
+    that re-enable ``requires_grad`` so the autograd graph for
+    reference / candidate is independent. Non-tensor values pass
+    through unchanged. Tensors that do NOT require grad pass through
+    by reference (so the caller can share read-only tensors like masks
+    across reference / candidate cheaply).
+    """
+    out: dict = {}
+    for k, v in inputs.items():
+        if isinstance(v, torch.Tensor) and v.requires_grad:
+            out[k] = v.detach().clone().requires_grad_(True)
+        else:
+            out[k] = v
+    return out
+
+
+def _diag(name: str, ref: torch.Tensor, cand: torch.Tensor, tol: Tol) -> str:
+    """Format a structured diff for an off-tolerance ``(ref, cand)`` pair."""
+    if ref.shape != cand.shape:
+        return f"{name}: shape mismatch — ref={tuple(ref.shape)} " f"vs cand={tuple(cand.shape)}"
+    diff = (ref.float() - cand.float()).abs()
+    # rel uses ref magnitude as denominator (matches torch.allclose
+    # semantics: atol + rtol * |ref|).
+    rel = diff / (ref.float().abs() + 1e-12)
+    max_abs = diff.max().item()
+    max_rel = rel.max().item()
+
+    flat_diff = diff.flatten()
+    k = min(8, flat_diff.numel())
+    topk = torch.topk(flat_diff, k=k).indices.tolist()
+    worst = []
+    flat_ref = ref.float().flatten()
+    flat_cand = cand.float().flatten()
+    for idx in topk:
+        worst.append(
+            f"  [{idx}]: ref={flat_ref[idx].item():+.6g} "
+            f"cand={flat_cand[idx].item():+.6g} "
+            f"abs={(flat_ref[idx] - flat_cand[idx]).abs().item():+.6g}"
+        )
+    worst_block = "\n".join(worst)
+    return (
+        f"{name}: max_abs={max_abs:.6g} (atol={tol.atol:.6g}), "
+        f"max_rel={max_rel:.6g} (rtol={tol.rtol:.6g}); 8 worst entries:\n" + worst_block
+    )
+
+
+def _assert_close(
+    name: str,
+    ref: torch.Tensor,
+    cand: torch.Tensor,
+    tol: Tol,
+) -> None:
+    if ref is None and cand is None:
+        return
+    if ref is None or cand is None:
+        raise AssertionError(
+            f"{name}: one side is None (ref is None: {ref is None}, " f"cand is None: {cand is None})"
+        )
+    # Match dtype before comparing — gradient dtypes can differ
+    # slightly across paths (one side accumulates fp32, the other
+    # bf16), so coerce to fp32 for the comparison itself.
+    ref32 = ref.detach().float()
+    cand32 = cand.detach().float()
+    if ref32.shape != cand32.shape:
+        raise AssertionError(_diag(name, ref32, cand32, tol))
+    if torch.allclose(ref32, cand32, **tol.as_kwargs()):
+        return
+    raise AssertionError(_diag(name, ref32, cand32, tol))
+
+
+# ---------------------------------------------------------------------------
+# Public entry point
+# ---------------------------------------------------------------------------
+
+
+def compare_fwd_bwd(
+    *,
+    reference: Callable[..., torch.Tensor],
+    candidate: Callable[..., torch.Tensor],
+    inputs: Mapping[str, Any],
+    fwd_tol: Tol,
+    bwd_tol: Tol,
+    grad_keys: Optional[Iterable[str]] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run reference + candidate on independent input clones; assert match.
+
+    Args:
+        reference: callable taking ``**inputs`` and returning a single
+            tensor (the "eager truth").
+        candidate: callable with the same signature (the path under
+            test).
+        inputs: kwargs forwarded to both. Tensor values that
+            ``requires_grad`` are independently cloned so the two
+            backward passes do not interfere.
+        fwd_tol: forward output tolerance.
+        bwd_tol: backward gradient tolerance.
+        grad_keys: optional override for which input keys to assert
+            gradients on. Defaults to every tensor in ``inputs`` that
+            ``requires_grad``.
+
+    Returns:
+        ``(out_ref, out_cand)`` for follow-up inspection by the caller
+        (e.g., asserting non-NaN, asserting shape, etc.).
+
+    Raises:
+        AssertionError: on first forward or backward mismatch, with a
+        structured diff in the message.
+    """
+    ref_inputs = _clone_for_grad(inputs)
+    cand_inputs = _clone_for_grad(inputs)
+
+    out_ref = reference(**ref_inputs)
+    out_cand = candidate(**cand_inputs)
+
+    _assert_close("forward output", out_ref, out_cand, fwd_tol)
+
+    out_ref.sum().backward()
+    out_cand.sum().backward()
+
+    if grad_keys is None:
+        grad_keys = [k for k, v in inputs.items() if isinstance(v, torch.Tensor) and v.requires_grad]
+    for k in grad_keys:
+        ref_t = ref_inputs[k]
+        cand_t = cand_inputs[k]
+        if not isinstance(ref_t, torch.Tensor) or not isinstance(cand_t, torch.Tensor):
+            raise AssertionError(
+                f"grad_keys[{k!r}] points at a non-tensor input "
+                f"(ref={type(ref_t).__name__}, cand={type(cand_t).__name__})"
+            )
+        _assert_close(f"gradient[{k}]", ref_t.grad, cand_t.grad, bwd_tol)
+
+    return out_ref, out_cand
+
+
+__all__ = [
+    "Tol",
+    "FP32_TOL",
+    "BF16_FWD_TOL",
+    "BF16_BWD_TOL",
+    "compare_fwd_bwd",
+    "default_tols",
+]
diff --git a/tests/unit_tests/test_backend_loader.py b/tests/unit_tests/test_backend_loader.py
index efee65af5..d143aabff 100644
--- a/tests/unit_tests/test_backend_loader.py
+++ b/tests/unit_tests/test_backend_loader.py
@@ -8,12 +8,11 @@
 import shutil
 import sys
 import tempfile
-import types
 from pathlib import Path
 
 import pytest
 
-from primus.pretrain import load_backend_trainer, setup_backend_path
+from primus.pretrain import setup_backend_path
 
 
 @pytest.fixture
@@ -64,23 +63,3 @@ def test_setup_backend_path_failure():
     """Test that FileNotFoundError is raised when no valid backend path exists."""
     with pytest.raises(FileNotFoundError):
         setup_backend_path(framework="nonexistent_backend")
-
-
-def test_load_backend_trainer_supported(monkeypatch):
-    """Test that load_backend_trainer returns correct class for supported frameworks."""
-    # Mock: Define a dummy class to simulate MegatronPretrainTrainer
-    dummy_class = type("DummyTrainer", (), {})
-    dummy_module = types.ModuleType("primus.modules.trainer.megatron.pre_trainer")
-    dummy_module.MegatronPretrainTrainer = dummy_class
-
-    # Inject into sys.modules to bypass real import
-    sys.modules["primus.modules.trainer.megatron.pre_trainer"] = dummy_module
-
-    trainer_cls = load_backend_trainer("megatron")
-    assert trainer_cls is dummy_class
-
-
-def test_load_backend_trainer_unsupported():
-    """Test that unsupported framework raises ValueError."""
-    with pytest.raises(ValueError, match="Unsupported framework"):
-        load_backend_trainer("invalid_framework")
diff --git a/tests/unit_tests/tools/test_preflight_bisect_slurm.py b/tests/unit_tests/tools/test_preflight_bisect_slurm.py
new file mode 100644
index 000000000..a4267e250
--- /dev/null
+++ b/tests/unit_tests/tools/test_preflight_bisect_slurm.py
@@ -0,0 +1,223 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Integration tests for tools/preflight_bisect/bisect.py. These require a live
+# Slurm allocation and skip automatically when the required environment is
+# missing.
+#
+# Required env vars
+# -----------------
+# Both tests require a Slurm nodelist. If BISECT_NODELIST is unset, the tests
+# use SLURM_NODELIST from the current allocation.
+#
+# Running from inside a Slurm allocation
+# --------------------------
+#   cd ~/Primus
+#   export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+#   python3 -m pytest tests/unit_tests/tools/test_preflight_bisect_slurm.py -v
+# More complex example
+# --------------------------
+#   export NCCL_SOCKET_IFNAME=tw-eth0
+#   export GLOO_SOCKET_IFNAME=tw-eth0
+#   export NCCL_IB_HCA="rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7:1"
+#   export BISECT_PREFLIGHT_ENV="USING_AINIC=1 NCCL_IB_GID_INDEX=3 NCCL_CROSS_NIC=0 NCCL_PXN_DISABLE=0"
+#
+#   python3 -m pytest \
+#       tests/unit_tests/tools/test_preflight_bisect_slurm.py \
+#       -v \
+#       --basetemp=/tmp/preflight-bisect-pytest
+#
+# Inspect pytest artifacts:
+#   ls -R /tmp/preflight-bisect-pytest
+#   cat /tmp/preflight-bisect-pytest/test_bisect_all_nodes_pass*/trial-*.log
+#   cat /tmp/preflight-bisect-pytest/test_bisect_all_nodes_pass*/summary.txt
+#
+# Defaults
+# --------
+# - BISECT_NODELIST defaults to SLURM_NODELIST.
+# - BISECT_PARTITION defaults to SLURM_JOB_PARTITION; otherwise --partition is omitted.
+# - BISECT_BAD_NODE defaults to the last host in the resolved nodelist.
+#
+# Optional tuning:
+#   BISECT_TRIAL_TIMEOUT, BISECT_SLURM_TIME, BISECT_PREFLIGHT_ENV
+# Keep comma-containing values such as NCCL_IB_HCA as normal exported env vars;
+# BISECT_PREFLIGHT_ENV is passed through srun --export=ALL,...
+#
+###############################################################################
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+BISECT_PY = REPO_ROOT / "tools" / "preflight_bisect" / "bisect.py"
+FAKE_RUNNER = REPO_ROOT / "tools" / "preflight_bisect" / "fake_runner.sh"
+REAL_RUNNER = REPO_ROOT / "runner" / "primus-cli"
+
+
+def _require_env(*names: str) -> dict[str, str]:
+    """Return a dict of the requested env var values, skipping the test if any are missing."""
+    missing = [n for n in names if not os.environ.get(n)]
+    if missing:
+        pytest.skip(f"Required env var(s) not set: {', '.join(missing)}")
+    return {n: os.environ[n] for n in names}
+
+
+def _resolve_nodelist() -> str:
+    nodelist = os.environ.get("BISECT_NODELIST") or os.environ.get("SLURM_NODELIST")
+    if not nodelist:
+        pytest.skip("Set BISECT_NODELIST or run from inside a Slurm allocation with SLURM_NODELIST set")
+    return nodelist
+
+
+def _resolve_partition() -> str:
+    return os.environ.get("BISECT_PARTITION") or os.environ.get("SLURM_JOB_PARTITION") or ""
+
+
+def _resolve_hosts(nodelist: str) -> list[str]:
+    try:
+        result = subprocess.run(
+            ["scontrol", "show", "hostnames", nodelist],
+            capture_output=True,
+            check=True,
+            text=True,
+        )
+    except FileNotFoundError:
+        pytest.skip("scontrol not found; run this test on a Slurm login/head node")
+    except subprocess.CalledProcessError as exc:
+        pytest.skip(f"scontrol failed to resolve nodelist {nodelist!r}: {exc.stderr.strip()}")
+
+    hosts = [line.strip() for line in result.stdout.splitlines() if line.strip()]
+    if not hosts:
+        pytest.skip(f"scontrol produced no hostnames for nodelist {nodelist!r}")
+    return hosts
+
+
+def _resolve_bad_node(nodelist: str) -> str:
+    return os.environ.get("BISECT_BAD_NODE") or _resolve_hosts(nodelist)[-1]
+
+
+def _resolve_bisect_env() -> dict[str, str]:
+    env = {"BISECT_NODELIST": _resolve_nodelist()}
+    partition = _resolve_partition()
+    if partition:
+        env["BISECT_PARTITION"] = partition
+    return env
+
+
+def _run_bisect(extra_args: list[str], env: dict[str, str], tmp_path: Path, timeout: int) -> str:
+    """Invoke bisect.py as a subprocess and return the contents of summary.txt."""
+    preflight_env_args: list[str] = []
+    for kv in os.environ.get("BISECT_PREFLIGHT_ENV", "").split():
+        preflight_env_args += ["--preflight-env", kv]
+
+    cmd = [
+        sys.executable,
+        str(BISECT_PY),
+        "--nodelist",
+        env["BISECT_NODELIST"],
+        "--output-dir",
+        str(tmp_path),
+    ]
+    if env.get("BISECT_PARTITION"):
+        cmd.extend(["--partition", env["BISECT_PARTITION"]])
+    cmd.extend([*preflight_env_args, *extra_args])
+
+    result = subprocess.run(
+        cmd,
+        env={**os.environ, **env},
+        capture_output=True,
+        text=True,
+        timeout=timeout,
+    )
+    summary_path = tmp_path / "summary.txt"
+    summary_text = summary_path.read_text(encoding="utf-8") if summary_path.exists() else ""
+
+    if result.returncode != 0:
+        pytest.fail(
+            f"bisect.py exited {result.returncode} and wrote no summary.\n"
+            f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+        )
+
+    return summary_text
+
+
+def test_bisect_all_nodes_pass(tmp_path):
+    """Run bisect.py with the real preflight runner against a known-healthy nodeset.
+
+    Requires VENV_ACTIVATE and either BISECT_NODELIST or SLURM_NODELIST.
+    All nodes are expected to pass the preflight perf-test, so bisect should
+    report SUSPECT_NODES: (none).
+    """
+    env = {**_resolve_bisect_env(), **_require_env("VENV_ACTIVATE")}
+    trial_timeout = int(os.environ.get("BISECT_TRIAL_TIMEOUT", "600"))
+    slurm_time = os.environ.get("BISECT_SLURM_TIME", "00:15:00")
+
+    # subprocess timeout: give a comfortable margin above the per-trial timeout
+    # to account for bisect recursion and Slurm scheduling overhead.
+    # Each level of bisection can run up to 2 concurrent trials, and a nodeset
+    # of N nodes has at most log2(N)+1 levels, so multiply generously.
+    subprocess_timeout = trial_timeout * 8
+
+    summary_text = _run_bisect(
+        extra_args=[
+            "--trial-timeout-sec",
+            str(trial_timeout),
+            "--slurm-time",
+            slurm_time,
+            "--runner",
+            str(REAL_RUNNER),
+        ],
+        env=env,
+        tmp_path=tmp_path,
+        timeout=subprocess_timeout,
+    )
+
+    assert (
+        "SUSPECT_NODES: (none)" in summary_text
+    ), f"Expected no suspect nodes for a healthy nodeset, but got:\n{summary_text}"
+
+
+def test_bisect_identifies_bad_node(tmp_path):
+    """Run bisect.py with fake_runner.sh seeding one bad node.
+
+    Requires either BISECT_NODELIST or SLURM_NODELIST.
+    BISECT_BAD_NODE can override the default bad node, which is the last host in
+    the resolved nodelist.
+    bisect.py is expected to identify exactly that node as the sole suspect.
+    """
+    env = _resolve_bisect_env()
+    bad_node = _resolve_bad_node(env["BISECT_NODELIST"])
+    env["BISECT_BAD_NODE"] = bad_node
+    trial_timeout = int(os.environ.get("BISECT_TRIAL_TIMEOUT", "30"))
+    slurm_time = os.environ.get("BISECT_SLURM_TIME", "00:02:00")
+
+    subprocess_timeout = trial_timeout * 8
+
+    summary_text = _run_bisect(
+        extra_args=[
+            "--trial-timeout-sec",
+            str(trial_timeout),
+            "--slurm-time",
+            slurm_time,
+            "--runner",
+            str(FAKE_RUNNER),
+            "--preflight-env",
+            f"BAD_NODE={bad_node}",
+            "--max-concurrent-trials",
+            "2",
+        ],
+        env=env,
+        tmp_path=tmp_path,
+        timeout=subprocess_timeout,
+    )
+
+    assert (
+        f"SUSPECT_NODES: {bad_node}" in summary_text
+    ), f"Expected '{bad_node}' to be identified as the sole suspect, but got:\n{summary_text}"
diff --git a/tests/utils.py b/tests/utils.py
index 9255b75a6..83490432b 100755
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -10,17 +10,58 @@
 import sys
 import time
 import unittest
-import warnings
 from typing import Optional
 
 from primus.core.utils import logger
 
 TRAINING_COMPLETED_MARKER = "Training completed."
 
-# run_patcher's failure line ("[Patch] \u2717 Patch '<id>' failed..."); distinct
-# from a patch's own graceful "[SKIP]". run_patches doesn't fail training on it
-# (stop_on_error=False), so we surface it ourselves.
-PATCH_FAILURE_MARKER = "\u2717 Patch '"
+
+def skip_if_no_cuda(reason: str = "requires GPU (primus_turbo initializes CUDA at import)") -> None:
+    """Skip the calling test module at collection time when CUDA is unavailable.
+
+    Several Flux/diffusion test modules import ``primus_turbo`` (directly or
+    transitively), which initializes CUDA at import and raises on CPU-only
+    hosts. Call this at module scope *before* those imports so collection
+    succeeds without a GPU. It is a no-op when CUDA is present, so GPU CI still
+    runs every test.
+    """
+    import pytest
+    import torch
+
+    if not torch.cuda.is_available():
+        pytest.skip(reason, allow_module_level=True)
+
+
+def install_aiter_deepbind_hook() -> None:
+    """Install Primus' production RTLD_DEEPBIND import hook for aiter's mha kernels.
+
+    This invokes the exact mechanism the ``megatron.turbo.aiter_deepbind``
+    before_train patch uses in real training: it wraps ``importlib.import_module``
+    so aiter's pinned mha extensions bind their own ``aiter::mha_bwd`` instead of
+    transformer_engine's stale vendored ``libmha`` (ROCm/aiter#1332). On
+    gfx942/gfx950 that stale-symbol interposition otherwise makes Turbo's hd128
+    backward launch with an invalid grid config and crash the process.
+
+    Unit tests call ``flash_attn_func`` directly and never run the before_train
+    phase, so without this hook they hit the same crash the production patch
+    prevents. Call this at diffusion ``conftest`` import time so the hook is in
+    place before any test first imports the aiter mha modules (which happens
+    lazily on the first attention op). No-op when CUDA is unavailable or the hook
+    cannot be installed (e.g. CPU-only host).
+    """
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return
+        from primus.backends.megatron.patches.turbo.aiter_deepbind_patches import (
+            _install_deepbind_import_hook,
+        )
+    except Exception:
+        return
+
+    _install_deepbind_import_hook()
 
 
 class PrimusUT(unittest.TestCase):
@@ -51,25 +92,6 @@ def tearDown(self):
         pass
 
 
-def _warn_on_patch_failures(tag: str, stdout_output: str) -> None:
-    """Warn (don't fail) when patches failed to apply during training.
-
-    A failed patch doesn't crash training and a 3-step smoke run can't prove it
-    was harmless, so we emit a GitHub Actions warning to keep it visible without
-    blocking CI. No-op for backends that don't emit the marker.
-    """
-    if PATCH_FAILURE_MARKER not in stdout_output:
-        return
-    failed = [ln.strip() for ln in stdout_output.splitlines() if PATCH_FAILURE_MARKER in ln]
-    msg = (
-        f"[{tag}] {len(failed)} patch(es) failed to apply during training "
-        f"(training still completed): " + " | ".join(failed[:10])
-    )
-    # GitHub Actions annotation: shows on the run/PR without failing the job.
-    print(f"::warning title=Primus patch failed to apply::{msg}")
-    warnings.warn(msg)
-
-
 def run_training_script(
     tag: str,
     cmd: list[str],
@@ -129,8 +151,6 @@ def run_training_script(
                 f"Log file: {train_log_path}"
             )
 
-        _warn_on_patch_failures(tag, stdout_output)
-
         return stdout_output, ""
 
     except subprocess.CalledProcessError as e:
diff --git a/third_party/Automodel b/third_party/Automodel
new file mode 160000
index 000000000..c852b16ff
--- /dev/null
+++ b/third_party/Automodel
@@ -0,0 +1 @@
+Subproject commit c852b16ff97221bd97a408e224f011a817cbc8e9
diff --git a/third_party/mamba b/third_party/mamba
new file mode 160000
index 000000000..4d67c534d
--- /dev/null
+++ b/third_party/mamba
@@ -0,0 +1 @@
+Subproject commit 4d67c534d10b7299194ede55b36718cc7a1dd472
diff --git a/third_party/maxtext b/third_party/maxtext
index 851c0935a..a7c6c7e53 160000
--- a/third_party/maxtext
+++ b/third_party/maxtext
@@ -1 +1 @@
-Subproject commit 851c0935a4e9c9e4cd0771603e3bdf75a6578ad1
+Subproject commit a7c6c7e5377cce39be51cbe2a48fe0f86fbab8dc
diff --git a/third_party/torchtitan b/third_party/torchtitan
index 5fb7cc2e3..73a0e6979 160000
--- a/third_party/torchtitan
+++ b/third_party/torchtitan
@@ -1 +1 @@
-Subproject commit 5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021
+Subproject commit 73a0e6979dd10b6b1904098eb3c8f62c18ab87ce
diff --git a/tools/README.md b/tools/README.md
index e69de29bb..495effa93 100644
--- a/tools/README.md
+++ b/tools/README.md
@@ -0,0 +1,24 @@
+# Primus Tools
+
+Auxiliary tools for analysis, benchmarking, visualization, and diagnostics.
+
+## Available Tools
+
+| Tool | Directory | Description |
+|------|-----------|-------------|
+| **IRLens** | `tools/IRLens/` | Parses XLA HLO text dumps; prints execution skeleton with control flow, communication vs compute ops |
+| **model_stats** | `tools/model_stats/` | Generates charts from the model config registry under `primus/configs/models` |
+| **Pipeline Visualization** | `tools/visualization/pp_vis/` | Visualizes pipeline parallelism schedules from dumped data or PP simulator JSON via a local web UI |
+| **Auto Benchmark** | `tools/auto_benchmark/` | Interactive benchmark menu for Megatron/TorchTitan on MI300X/MI355X with metrics collection |
+| **Daily Report** | `tools/daily/` | Benchmark summary CSV generation used by CI workflows |
+| **Docker Helpers** | `tools/docker/` | Container startup and proxy scripts |
+| **Profile Trace** | `tools/profile_trace/` | Trace file merging utility |
+
+## Per-Tool Documentation
+
+Each tool has its own README with usage instructions:
+
+- [IRLens README](./IRLens/README.md)
+- [model_stats README](./model_stats/README.md)
+- [Pipeline Visualization README](./visualization/pp_vis/README.md)
+- [Auto Benchmark README](./auto_benchmark/Primus_Auto_Benchmark_README.md)
diff --git a/tools/auto_benchmark/metrics.py b/tools/auto_benchmark/metrics.py
new file mode 100644
index 000000000..a38a98a71
--- /dev/null
+++ b/tools/auto_benchmark/metrics.py
@@ -0,0 +1,718 @@
+#!/usr/bin/env python3
+
+import argparse
+import csv
+import os
+import re
+import shutil
+from collections import defaultdict
+from datetime import datetime
+from statistics import mean
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+PRIMUS_ROOT = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
+RESULTS_DIR = os.path.join(SCRIPT_DIR, "results")
+
+WARMUP_SKIP = 5
+
+ERROR_STATUS = "Error found - check log"
+OK_STATUS = "OK"
+
+ANSI_ESCAPE_REGEX = re.compile(r"\x1b\[[0-9;]*m")
+LOG_EXIT_CODE_REGEX = re.compile(r"primus launcher exited with code (\d+)", re.IGNORECASE)
+
+LOG_ERROR_PATTERNS = (
+    re.compile(r"Traceback \(most recent call last\)", re.IGNORECASE),
+    re.compile(r"\[ERROR\]", re.IGNORECASE),
+    re.compile(r"\bRuntimeError:", re.IGNORECASE),
+    re.compile(r"\bOutOfMemoryError:", re.IGNORECASE),
+    re.compile(r"\bCUDA out of memory\b", re.IGNORECASE),
+    re.compile(r"\bHIP out of memory\b", re.IGNORECASE),
+    re.compile(r"\bNCCL error\b", re.IGNORECASE),
+    re.compile(r"\bfatal error\b", re.IGNORECASE),
+)
+
+LOG_ERROR_EXCLUSIONS = (
+    re.compile(r"error_injection", re.IGNORECASE),
+    re.compile(r"\[SKIP\].*Import failed", re.IGNORECASE),
+    re.compile(r"avoid ImportError", re.IGNORECASE),
+    re.compile(r"TORCH_NCCL_ASYNC_ERROR_HANDLING", re.IGNORECASE),
+    re.compile(r"destroy_process_group\(\)", re.IGNORECASE),
+)
+
+RUN_LABEL_REGEX = re.compile(r"_run(\d+)\.log$", re.IGNORECASE)
+LEGACY_TIMESTAMP_REGEX = re.compile(r"_(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.log$")
+LEGACY_SUFFIX_REGEX = re.compile(r"_MI\d+X?_(.+)\.log$", re.IGNORECASE)
+
+ENV_DEFAULT_REGEX = re.compile(r"\$\{[A-Z0-9_]+:(\d+)\}")
+PLAIN_INT_REGEX = re.compile(r"^\d+$")
+
+MEGATRON_BS_KEYS = ("micro_batch_size",)
+MEGATRON_SEQ_KEYS = ("seq_length",)
+MEGATRON_GBS_KEYS = ("global_batch_size",)
+
+TORCHTITAN_BS_KEYS = ("local_batch_size",)
+TORCHTITAN_SEQ_KEYS = ("seq_len", "seq_length")
+TORCHTITAN_GBS_KEYS = ("global_batch_size",)
+
+MEGATRON_NOTE_TEXT = """
+NOTE:
+- Results are saved to results/metrics_megatron.csv (latest) and a timestamped snapshot.
+- "Run" is run1, run2, ... per model and device, ordered oldest to newest.
+- "BS", "Seq", and "GBS" come from the benchmark yaml when available.
+- Megatron logs often print each iteration twice (Primus log forwarding); metrics
+  deduplicate by iteration number before averaging.
+- Timing/throughput fields use the current value before "/" (e.g. 5896.1/5913.1 -> 5896.1).
+- Warm-up: the first five iterations are excluded before averaging.
+- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed.
+- Numeric fields may contain commas in logs; commas are removed before averaging.
+"""
+
+TORCHTITAN_NOTE_TEXT = """
+NOTE:
+- Results are saved to results/metrics_torchtitan.csv (latest) and a timestamped snapshot.
+- "Run" is run1, run2, ... per model and device, ordered oldest to newest.
+- "BS", "Seq", and "GBS" come from the benchmark yaml when available.
+- "Steps" is the number of training steps used after dropping the first five warm-up steps.
+- TPS and TFLOPS values may contain commas in logs; commas are removed before averaging.
+- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed.
+"""
+
+MEGATRON_NUM = r"[\d,]+(?:\.\d+)?"
+MEGATRON_METRIC_VALUE = rf"({MEGATRON_NUM})(?:\s*/\s*{MEGATRON_NUM})?"
+
+MEGATRON_ITERATION_REGEX = re.compile(
+    rf"iteration\s+(\d+)/\s*\d+.*?" rf"elapsed time per iteration \(ms\):\s*{MEGATRON_METRIC_VALUE}.*?"
+    # Compute (TFLOP/s/GPU) accepts both the current "compute per GPU (...)" label
+    # and the legacy "throughput per GPU (...)" label. Only the instantaneous value
+    # is captured; the "(avg Y)" / "/Y" suffix is absorbed by the trailing ".*?".
+    rf"(?:compute|throughput) per GPU \(TFLOP/s/GPU\):\s*{MEGATRON_METRIC_VALUE}.*?"
+    # Tokens accepts both the current "tokens/s/GPU inst/harmonic mean:" label and
+    # the legacy "tokens per GPU (tokens/s/GPU):" label.
+    rf"(?:(?:tokens per GPU \(tokens/s/GPU\)|tokens/s/GPU inst/harmonic mean):\s*{MEGATRON_METRIC_VALUE}.*?)?"
+    rf"global batch size:\s*(\d+)",
+    re.IGNORECASE,
+)
+
+MEGATRON_FILENAME_REGEX = re.compile(r"(?P<model>.+?)_megatron_(?P<device>MI\d+X?)", re.IGNORECASE)
+
+TORCHTITAN_STEP_REGEX = re.compile(
+    r"step:\s*(\d+).*?"
+    r"memory:\s*([\d.]+)GiB.*?"
+    r"tps:\s*([\d,]+(?:\.\d+)?).*?"
+    r"tflops:\s*([\d,]+(?:\.\d+)?).*?"
+    r"mfu:\s*([\d.]+)%"
+)
+
+TORCHTITAN_BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)")
+TORCHTITAN_SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)")
+
+TORCHTITAN_FILENAME_REGEX = re.compile(r"(?P<model>.+?)_torchtitan_(?P<device>MI\d+X?)", re.IGNORECASE)
+
+PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE)
+
+
+def _parse_scalar(value):
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, (int, float)):
+        return int(value) if float(value).is_integer() else value
+    if isinstance(value, str):
+        stripped = value.strip()
+        if PLAIN_INT_REGEX.match(stripped):
+            return int(stripped)
+        m = ENV_DEFAULT_REGEX.search(stripped)
+        if m:
+            return int(m.group(1))
+    return None
+
+
+def _deep_find(node, keys):
+    if isinstance(node, dict):
+        for key in keys:
+            if key in node:
+                parsed = _parse_scalar(node[key])
+                if parsed is not None:
+                    return parsed
+        for child in node.values():
+            found = _deep_find(child, keys)
+            if found is not None:
+                return found
+    elif isinstance(node, list):
+        for item in node:
+            found = _deep_find(item, keys)
+            if found is not None:
+                return found
+    return None
+
+
+def _regex_find(text, keys):
+    for key in keys:
+        m = re.search(
+            rf"^\s*{re.escape(key)}:\s*(.+?)(?:\s+#.*)?$",
+            text,
+            re.MULTILINE,
+        )
+        if not m:
+            continue
+        parsed = _parse_scalar(m.group(1).strip())
+        if parsed is not None:
+            return parsed
+    return None
+
+
+def load_yaml_dict(path):
+    try:
+        import yaml
+    except ImportError:
+        return None
+
+    try:
+        with open(path, "r", errors="ignore") as f:
+            data = yaml.safe_load(f)
+    except Exception:
+        return None
+
+    return data if isinstance(data, dict) else None
+
+
+def resolve_config_yaml(log_fname, backend, device, model, log_dir):
+    base = log_fname[:-4] if log_fname.endswith(".log") else log_fname
+
+    for suffix in ("_override", "_edited", ""):
+        candidate = os.path.join(log_dir, f"{base}{suffix}.yaml")
+        if os.path.isfile(candidate):
+            return candidate
+
+    config_dir = os.path.join(PRIMUS_ROOT, "examples", backend, "configs", device)
+    return os.path.join(config_dir, f"{model}.yaml")
+
+
+def load_training_params(backend, device, model, log_fname, log_dir):
+    path = resolve_config_yaml(log_fname, backend, device, model, log_dir)
+    if backend == "megatron":
+        bs_keys, seq_keys, gbs_keys = MEGATRON_BS_KEYS, MEGATRON_SEQ_KEYS, MEGATRON_GBS_KEYS
+    else:
+        bs_keys, seq_keys, gbs_keys = TORCHTITAN_BS_KEYS, TORCHTITAN_SEQ_KEYS, TORCHTITAN_GBS_KEYS
+
+    bs = seq = gbs = None
+
+    data = load_yaml_dict(path)
+    if data is not None:
+        bs = _deep_find(data, bs_keys)
+        seq = _deep_find(data, seq_keys)
+        gbs = _deep_find(data, gbs_keys)
+
+    if os.path.isfile(path) and (bs is None or seq is None or gbs is None):
+        with open(path, "r", errors="ignore") as f:
+            text = f.read()
+        if bs is None:
+            bs = _regex_find(text, bs_keys)
+        if seq is None:
+            seq = _regex_find(text, seq_keys)
+        if gbs is None:
+            gbs = _regex_find(text, gbs_keys)
+
+    return (
+        bs if bs is not None else "-",
+        seq if seq is not None else "-",
+        gbs if gbs is not None else "-",
+    )
+
+
+def log_chronological_key(fname, path):
+    m = RUN_LABEL_REGEX.search(fname)
+    if m:
+        return (0, int(m.group(1)), fname)
+
+    m = LEGACY_TIMESTAMP_REGEX.search(fname)
+    if m:
+        return (1, m.group(1), fname)
+
+    m = LEGACY_SUFFIX_REGEX.search(fname)
+    if m and m.group(1):
+        return (1, m.group(1), fname)
+
+    return (2, os.path.getmtime(path), fname)
+
+
+def assign_run_labels(entries):
+    grouped = defaultdict(list)
+    for entry in entries:
+        grouped[(entry["model"], entry["device"])].append(entry)
+
+    for group in grouped.values():
+        group.sort(key=lambda entry: log_chronological_key(entry["fname"], entry["path"]))
+        for index, entry in enumerate(group, start=1):
+            entry["run"] = f"run{index}"
+
+    return entries
+
+
+def run_sort_key(run_label):
+    m = re.fullmatch(r"run(\d+)", run_label)
+    if m:
+        return int(m.group(1))
+    return 0
+
+
+def terminal_width(default=120):
+    try:
+        return shutil.get_terminal_size().columns
+    except OSError:
+        return default
+
+
+def print_table(headers, rows, max_col_width=36):
+    if not rows:
+        return
+
+    widths = []
+    for i, header in enumerate(headers):
+        col_width = max(len(str(header)), *(len(str(row[i])) for row in rows))
+        widths.append(min(col_width, max_col_width))
+
+    def fmt(row):
+        cells = []
+        for i, value in enumerate(row):
+            text = str(value)
+            if len(text) > widths[i]:
+                text = text[: widths[i] - 3] + "..."
+            cells.append(text.ljust(widths[i]))
+        return "| " + " | ".join(cells) + " |"
+
+    sep = "+-" + "-+-".join("-" * w for w in widths) + "-+"
+    print(sep)
+    print(fmt(headers))
+    print(sep)
+    for row in rows:
+        print(fmt(row))
+    print(sep)
+
+
+def save_csv(headers, rows, backend):
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    latest_path = os.path.join(RESULTS_DIR, f"metrics_{backend}.csv")
+    snapshot_path = os.path.join(RESULTS_DIR, f"metrics_{backend}_{timestamp}.csv")
+
+    for path in (latest_path, snapshot_path):
+        with open(path, "w", newline="", encoding="utf-8") as handle:
+            writer = csv.writer(handle)
+            writer.writerow(headers)
+            writer.writerows(rows)
+
+    return latest_path, snapshot_path
+
+
+def strip_ansi(text):
+    return ANSI_ESCAPE_REGEX.sub("", text)
+
+
+def log_has_error(path):
+    exit_code = None
+    saw_error_line = False
+
+    with open(path, "r", errors="ignore") as f:
+        for line in f:
+            plain = strip_ansi(line)
+
+            m = LOG_EXIT_CODE_REGEX.search(plain)
+            if m:
+                exit_code = int(m.group(1))
+
+            if any(pattern.search(plain) for pattern in LOG_ERROR_EXCLUSIONS):
+                continue
+
+            if any(pattern.search(plain) for pattern in LOG_ERROR_PATTERNS):
+                saw_error_line = True
+
+    if exit_code not in (None, 0):
+        return True
+    return saw_error_line
+
+
+def render_results(backend, headers, rows, note_text):
+    if not rows:
+        print(f"No {backend} logs found.")
+        return None
+
+    latest_path, snapshot_path = save_csv(headers, rows, backend)
+    print_table(headers, rows)
+
+    error_rows = [row for row in rows if len(row) > 2 and row[2] == ERROR_STATUS]
+    if error_rows:
+        print(
+            f"\n{len(error_rows)} run(s) reported errors or incomplete metrics."
+            " Open the corresponding log file for details."
+        )
+
+    table_width = sum(len(str(header)) for header in headers) + 3 * len(headers)
+    if table_width > terminal_width():
+        print("\n(Table may be wider than the terminal. Open the CSV for the full view.)")
+
+    print(note_text)
+    print("\nMetrics saved to:")
+    print(f"  Latest:   {latest_path}")
+    print(f"  Snapshot: {snapshot_path}")
+
+    return latest_path
+
+
+def is_megatron(filename):
+    name = filename.lower()
+    return "megatron" in name or "megatorn" in name
+
+
+def megatron_parse_filename(filename):
+    m = MEGATRON_FILENAME_REGEX.search(filename)
+    if not m:
+        return None
+
+    p = PRECISION_REGEX.search(filename)
+    precision = p.group(1).upper() if p else "-"
+
+    return {
+        "model": m.group("model"),
+        "device": m.group("device"),
+        "precision": precision,
+    }
+
+
+def megatron_to_float(num_str):
+    return float(num_str.replace(",", ""))
+
+
+def megatron_parse_log_file(path):
+    records_by_iter = {}
+
+    with open(path, "r", errors="ignore") as f:
+        for line in f:
+            if "iteration" not in line:
+                continue
+
+            m = MEGATRON_ITERATION_REGEX.search(line)
+            if not m:
+                continue
+
+            iter_num = int(m.group(1))
+            if iter_num in records_by_iter:
+                continue
+
+            records_by_iter[iter_num] = {
+                "iter": iter_num,
+                "elapsed_ms": megatron_to_float(m.group(2)),
+                "tflops_gpu": megatron_to_float(m.group(3)),
+                "tokens_gpu": megatron_to_float(m.group(4)) if m.group(4) else None,
+                "gbs": int(m.group(5)),
+            }
+
+    return [records_by_iter[i] for i in sorted(records_by_iter)]
+
+
+def megatron_compute_averages(records):
+    if len(records) <= WARMUP_SKIP:
+        return None
+
+    records = records[WARMUP_SKIP:]
+    token_values = [r["tokens_gpu"] for r in records if r["tokens_gpu"] is not None]
+
+    return {
+        "count": len(records),
+        "elapsed_ms": mean(r["elapsed_ms"] for r in records),
+        "tflops_gpu": mean(r["tflops_gpu"] for r in records),
+        "tokens_gpu": mean(token_values) if token_values else None,
+        "gbs": records[0]["gbs"],
+    }
+
+
+def megatron_collect_rows():
+    log_dir = os.path.join(RESULTS_DIR, "logs_megatron")
+    entries = []
+
+    if not os.path.isdir(log_dir):
+        return []
+
+    for fname in sorted(os.listdir(log_dir)):
+        if not fname.endswith(".log") or not is_megatron(fname):
+            continue
+
+        meta = megatron_parse_filename(fname)
+        if not meta:
+            continue
+
+        path = os.path.join(log_dir, fname)
+        has_error = log_has_error(path)
+        records = megatron_parse_log_file(path)
+        stats = megatron_compute_averages(records)
+
+        bs, seq, gbs = load_training_params("megatron", meta["device"], meta["model"], fname, log_dir)
+
+        if has_error or not stats:
+            if gbs == "-" and records:
+                gbs = records[0]["gbs"]
+            values = [
+                ERROR_STATUS,
+                "megatron",
+                meta["device"],
+                bs,
+                seq,
+                gbs,
+                meta["precision"],
+                "-",
+                "-",
+                "-",
+                "-",
+            ]
+        else:
+            if gbs == "-":
+                gbs = stats["gbs"]
+
+            tokens_gpu = f"{stats['tokens_gpu']:.2f}" if stats["tokens_gpu"] is not None else "-"
+            values = [
+                OK_STATUS,
+                "megatron",
+                meta["device"],
+                bs,
+                seq,
+                gbs,
+                meta["precision"],
+                stats["count"],
+                f"{stats['elapsed_ms']:.2f}",
+                f"{stats['tflops_gpu']:.2f}",
+                tokens_gpu,
+            ]
+
+        entries.append(
+            {
+                "model": meta["model"],
+                "device": meta["device"],
+                "fname": fname,
+                "path": path,
+                "values": values,
+            }
+        )
+
+    assign_run_labels(entries)
+
+    rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries]
+    rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7]))
+    return rows
+
+
+def megatron_main():
+    headers = [
+        "Model",
+        "Run",
+        "Status",
+        "Backend",
+        "Device",
+        "BS",
+        "Seq",
+        "GBS",
+        "Precision",
+        "Iterations",
+        "Iter Time (ms)",
+        "TFLOPS/GPU",
+        "Tokens/GPU",
+    ]
+    render_results("megatron", headers, megatron_collect_rows(), MEGATRON_NOTE_TEXT)
+
+
+def is_torchtitan(filename):
+    return "torchtitan" in filename.lower()
+
+
+def torchtitan_parse_filename(filename):
+    m = TORCHTITAN_FILENAME_REGEX.search(filename)
+    if not m:
+        return None
+
+    p = PRECISION_REGEX.search(filename)
+    precision = p.group(1).upper() if p else "-"
+
+    return {
+        "model": m.group("model"),
+        "device": m.group("device"),
+        "precision": precision,
+    }
+
+
+def torchtitan_parse_log_file(path):
+    steps_by_num = {}
+
+    with open(path, "r", errors="ignore") as f:
+        for line in f:
+            m = TORCHTITAN_STEP_REGEX.search(line)
+            if not m:
+                continue
+
+            step_num = int(m.group(1))
+            if step_num in steps_by_num:
+                continue
+
+            steps_by_num[step_num] = {
+                "step": step_num,
+                "memory": float(m.group(2)),
+                "tps": float(m.group(3).replace(",", "")),
+                "tflops": float(m.group(4).replace(",", "")),
+                "mfu": float(m.group(5)),
+            }
+
+    return [steps_by_num[i] for i in sorted(steps_by_num)]
+
+
+def torchtitan_parse_log_training_fallback(path):
+    bs = None
+    seq = None
+
+    with open(path, "r", errors="ignore") as f:
+        for line in f:
+            if bs is None:
+                m = TORCHTITAN_BS_REGEX.search(line)
+                if m:
+                    bs = int(m.group(1))
+            if seq is None:
+                m = TORCHTITAN_SEQ_REGEX.search(line)
+                if m:
+                    seq = int(m.group(1))
+            if bs is not None and seq is not None:
+                break
+
+    return bs, seq
+
+
+def torchtitan_compute_averages(steps):
+    if len(steps) <= WARMUP_SKIP:
+        return None
+
+    steps = steps[WARMUP_SKIP:]
+
+    return {
+        "count": len(steps),
+        "memory": mean(s["memory"] for s in steps),
+        "tps": mean(s["tps"] for s in steps),
+        "tflops": mean(s["tflops"] for s in steps),
+        "mfu": mean(s["mfu"] for s in steps),
+    }
+
+
+def torchtitan_collect_rows():
+    log_dir = os.path.join(RESULTS_DIR, "logs_torchtitan")
+    entries = []
+
+    if not os.path.isdir(log_dir):
+        return []
+
+    for fname in sorted(os.listdir(log_dir)):
+        if not fname.endswith(".log") or not is_torchtitan(fname):
+            continue
+
+        meta = torchtitan_parse_filename(fname)
+        if not meta:
+            continue
+
+        path = os.path.join(log_dir, fname)
+        has_error = log_has_error(path)
+        steps = torchtitan_parse_log_file(path)
+        stats = torchtitan_compute_averages(steps)
+
+        bs, seq, gbs = load_training_params("torchtitan", meta["device"], meta["model"], fname, log_dir)
+        if bs == "-":
+            log_bs, _ = torchtitan_parse_log_training_fallback(path)
+            if log_bs is not None:
+                bs = log_bs
+        if seq == "-":
+            _, log_seq = torchtitan_parse_log_training_fallback(path)
+            if log_seq is not None:
+                seq = log_seq
+
+        if has_error or not stats:
+            values = [
+                ERROR_STATUS,
+                "torchtitan",
+                meta["device"],
+                bs,
+                seq,
+                gbs,
+                meta["precision"],
+                "-",
+                "-",
+                "-",
+                "-",
+                "-",
+            ]
+        else:
+            values = [
+                OK_STATUS,
+                "torchtitan",
+                meta["device"],
+                bs,
+                seq,
+                gbs,
+                meta["precision"],
+                stats["count"],
+                f"{stats['memory']:.2f}",
+                f"{stats['tps']:.2f}",
+                f"{stats['tflops']:.2f}",
+                f"{stats['mfu']:.2f}",
+            ]
+
+        entries.append(
+            {
+                "model": meta["model"],
+                "device": meta["device"],
+                "fname": fname,
+                "path": path,
+                "values": values,
+            }
+        )
+
+    assign_run_labels(entries)
+
+    rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries]
+    rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7]))
+    return rows
+
+
+def torchtitan_main():
+    headers = [
+        "Model",
+        "Run",
+        "Status",
+        "Backend",
+        "Device",
+        "BS",
+        "Seq",
+        "GBS",
+        "Precision",
+        "Steps",
+        "Mem(GiB)",
+        "TPS",
+        "TFLOPS",
+        "MFU(%)",
+    ]
+    render_results("torchtitan", headers, torchtitan_collect_rows(), TORCHTITAN_NOTE_TEXT)
+
+
+BACKENDS = {
+    "megatron": megatron_main,
+    "torchtitan": torchtitan_main,
+}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate benchmark metrics tables and CSVs.")
+    parser.add_argument(
+        "backend",
+        choices=sorted(BACKENDS),
+        help="Training backend to process logs for",
+    )
+    args = parser.parse_args()
+    BACKENDS[args.backend]()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/auto_benchmark/metrics_megatron.py b/tools/auto_benchmark/metrics_megatron.py
deleted file mode 100644
index ff70c9c79..000000000
--- a/tools/auto_benchmark/metrics_megatron.py
+++ /dev/null
@@ -1,212 +0,0 @@
-#!/usr/bin/env python3
-
-import os
-import re
-from statistics import mean
-
-LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_megatron"
-
-# ============================================================
-# NOTE ON ITERATIONS AND AVERAGING
-# ============================================================
-
-NOTE_TEXT = """
-NOTE:
-- "Iterations" represents the number of training iterations USED to compute
-  the averages shown above (after warm-up removal).
-
-- All iteration records are extracted from the log file, sorted by iteration
-  number, and the FIRST TWO iterations are discarded to remove warm-up effects.
-
-- Iteration numbers do NOT need to be sequential. Valid examples:
-    * iteration 1, iteration 2, iteration 3
-    * iteration 1, iteration 5, iteration 10
-    * iteration 1, iteration 10, iteration 20
-
-- Numeric fields (elapsed time, TFLOPS/GPU, tokens/GPU) may appear as:
-    1234
-    1,234
-    1,234.56
-  Commas are removed before averaging.
-"""
-
-# ============================================================
-# Regex patterns (comma-safe)
-# ============================================================
-
-NUM = r"[\d,]+(?:\.\d+)?"
-
-ITERATION_REGEX = re.compile(
-    rf"iteration\s+(\d+)/\s*\d+.*?"
-    rf"elapsed time per iteration \(ms\):\s*({NUM}).*?"
-    rf"throughput per GPU \(TFLOP/s/GPU\):\s*({NUM}).*?"
-    rf"tokens per GPU \(tokens/s/GPU\):\s*({NUM}).*?"
-    rf"global batch size:\s*(\d+)",
-    re.IGNORECASE,
-)
-
-# Filename metadata
-FILENAME_REGEX = re.compile(r"(?P<model>.+?)_megatron_(?P<device>MI\d+X?)", re.IGNORECASE)
-
-PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE)
-
-# ============================================================
-# Helpers
-# ============================================================
-
-
-def is_megatron(filename: str) -> bool:
-    name = filename.lower()
-    return "megatron" in name or "megatorn" in name
-
-
-def parse_filename(filename: str):
-    m = FILENAME_REGEX.search(filename)
-    if not m:
-        return None
-
-    p = PRECISION_REGEX.search(filename)
-    precision = p.group(1).upper() if p else "-"
-
-    return {
-        "model": m.group("model"),
-        "device": m.group("device"),
-        "precision": precision,
-    }
-
-
-def to_float(num_str: str) -> float:
-    """Convert numbers like '1,234.56' safely to float."""
-    return float(num_str.replace(",", ""))
-
-
-# ============================================================
-# Log parsing (FULL FILE SCAN)
-# ============================================================
-
-
-def parse_log_file(path):
-    records = []
-
-    with open(path, "r", errors="ignore") as f:
-        for line in f:
-            m = ITERATION_REGEX.search(line)
-            if m:
-                records.append(
-                    {
-                        "iter": int(m.group(1)),
-                        "elapsed_ms": to_float(m.group(2)),
-                        "tflops_gpu": to_float(m.group(3)),
-                        "tokens_gpu": to_float(m.group(4)),
-                        "gbs": int(m.group(5)),
-                    }
-                )
-
-    return records
-
-
-def compute_averages(records):
-    if len(records) <= 2:
-        return None
-
-    # Sort by iteration number (non-sequential safe)
-    records = sorted(records, key=lambda x: x["iter"])
-
-    # Drop first two warm-up iterations
-    records = records[2:]
-
-    return {
-        "count": len(records),
-        "elapsed_ms": mean(r["elapsed_ms"] for r in records),
-        "tflops_gpu": mean(r["tflops_gpu"] for r in records),
-        "tokens_gpu": mean(r["tokens_gpu"] for r in records),
-        "gbs": records[0]["gbs"],
-    }
-
-
-# ============================================================
-# Table printer
-# ============================================================
-
-
-def print_table(headers, rows):
-    widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
-
-    def fmt(row):
-        return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |"
-
-    sep = "+-" + "-+-".join("-" * w for w in widths) + "-+"
-
-    print(sep)
-    print(fmt(headers))
-    print(sep)
-    for r in rows:
-        print(fmt(r))
-    print(sep)
-
-
-def print_note():
-    print(NOTE_TEXT)
-
-
-# ============================================================
-# Main
-# ============================================================
-
-
-def main():
-    rows = []
-
-    for fname in sorted(os.listdir(LOG_DIR)):
-        if not fname.endswith(".log"):
-            continue
-
-        if not is_megatron(fname):
-            continue
-
-        meta = parse_filename(fname)
-        if not meta:
-            continue
-
-        path = os.path.join(LOG_DIR, fname)
-
-        records = parse_log_file(path)
-        stats = compute_averages(records)
-        if not stats:
-            continue
-
-        rows.append(
-            [
-                meta["model"],
-                "megatron",
-                meta["device"],
-                meta["precision"],
-                stats["count"],
-                f"{stats['elapsed_ms']:.2f}",
-                f"{stats['tflops_gpu']:.2f}",
-                f"{stats['tokens_gpu']:.2f}",
-                stats["gbs"],
-            ]
-        )
-
-    headers = [
-        "Model",
-        "Backend",
-        "Device",
-        "Precision",
-        "Iterations",
-        "Iter Time (ms)",
-        "TFLOPS/GPU",
-        "Tokens/GPU",
-        "GBS",
-    ]
-
-    if rows:
-        print_table(headers, rows)
-        print_note()
-    else:
-        print("No Megatron logs found.")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tools/auto_benchmark/metrics_torchtitan.py b/tools/auto_benchmark/metrics_torchtitan.py
deleted file mode 100644
index 88770f72f..000000000
--- a/tools/auto_benchmark/metrics_torchtitan.py
+++ /dev/null
@@ -1,224 +0,0 @@
-#!/usr/bin/env python3
-
-import os
-import re
-from statistics import mean
-
-LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_torchtitan"
-
-# ============================================================
-# Regex patterns
-# ============================================================
-
-STEP_REGEX = re.compile(
-    r"step:\s*(\d+).*?"
-    r"memory:\s*([\d.]+)GiB.*?"
-    r"tps:\s*([\d,]+(?:\.\d+)?).*?"
-    r"tflops:\s*([\d,]+(?:\.\d+)?).*?"
-    r"mfu:\s*([\d.]+)%"
-)
-
-# Dot-aligned config values
-BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)")
-
-SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)")
-
-# Filename metadata
-FILENAME_REGEX = re.compile(r"(?P<model>.+?)_torchtitan_(?P<device>MI\d+X?)", re.IGNORECASE)
-
-PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE)
-
-# ============================================================
-# Helpers
-# ============================================================
-
-
-def is_torchtitan(filename):
-    return "torchtitan" in filename.lower()
-
-
-def parse_filename(filename):
-    m = FILENAME_REGEX.search(filename)
-    if not m:
-        return None
-
-    p = PRECISION_REGEX.search(filename)
-    precision = p.group(1).upper() if p else "-"
-
-    return {
-        "model": m.group("model"),
-        "device": m.group("device"),
-        "precision": precision,
-    }
-
-
-# ============================================================
-# Log parsing (FULL FILE SCAN)
-# ============================================================
-
-
-def parse_log_file(path):
-    """
-    Fully scans the log file to extract:
-      - BS from training.local_batch_size
-      - SEQ from training.seq_len
-      - Per-step performance metrics
-    """
-    bs = None
-    seq = None
-    steps = []
-
-    with open(path, "r", errors="ignore") as f:
-        for line in f:
-            # Batch size
-            if bs is None:
-                m = BS_REGEX.search(line)
-                if m:
-                    bs = int(m.group(1))
-
-            # Sequence length
-            if seq is None:
-                m = SEQ_REGEX.search(line)
-                if m:
-                    seq = int(m.group(1))
-
-            # Step metrics
-            m = STEP_REGEX.search(line)
-            if m:
-                steps.append(
-                    {
-                        "step": int(m.group(1)),
-                        "memory": float(m.group(2)),
-                        "tps": float(m.group(3).replace(",", "")),
-                        "tflops": float(m.group(4).replace(",", "")),
-                        "mfu": float(m.group(5)),
-                    }
-                )
-
-    return bs if bs is not None else "-", seq if seq is not None else "-", steps
-
-
-def compute_averages(steps):
-    if len(steps) <= 2:
-        return None
-
-    steps = sorted(steps, key=lambda x: x["step"])
-    steps = steps[2:]  # drop first two warm-up steps
-
-    return {
-        "count": len(steps),
-        "memory": mean(s["memory"] for s in steps),
-        "tps": mean(s["tps"] for s in steps),
-        "tflops": mean(s["tflops"] for s in steps),
-        "mfu": mean(s["mfu"] for s in steps),
-    }
-
-
-# ============================================================
-# Table printer
-# ============================================================
-
-
-def print_table(headers, rows):
-    widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
-
-    def fmt(row):
-        return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |"
-
-    sep = "+-" + "-+-".join("-" * w for w in widths) + "-+"
-
-    print(sep)
-    print(fmt(headers))
-    print(sep)
-    for r in rows:
-        print(fmt(r))
-    print(sep)
-
-
-# ============================================================
-# Note printer
-# ============================================================
-
-
-def print_note():
-    print(
-        """
-NOTE:
-- "Steps" represents the number of training steps USED to compute averages.
-- Steps are sorted by step number and the FIRST TWO steps are dropped
-  to remove warm-up effects.
-- Step numbers do NOT need to be sequential. Valid examples:
-    * step 1, step 2, step 3
-    * step 1, step 5, step 10
-    * step 1, step 10, step 20
-- TPS and TFLOPS values may contain commas (e.g. 1,202.17);
-  commas are removed before averaging.
-"""
-    )
-
-
-# ============================================================
-# Main
-# ============================================================
-
-
-def main():
-    rows = []
-
-    for fname in sorted(os.listdir(LOG_DIR)):
-        if not fname.endswith(".log"):
-            continue
-
-        if not is_torchtitan(fname):
-            continue
-
-        meta = parse_filename(fname)
-        if not meta:
-            continue
-
-        path = os.path.join(LOG_DIR, fname)
-
-        bs, seq, steps = parse_log_file(path)
-        stats = compute_averages(steps)
-        if not stats:
-            continue
-
-        rows.append(
-            [
-                meta["model"],
-                "torchtitan",
-                meta["device"],
-                bs,
-                seq,
-                meta["precision"],
-                stats["count"],
-                f"{stats['memory']:.2f}",
-                f"{stats['tps']:.2f}",
-                f"{stats['tflops']:.2f}",
-                f"{stats['mfu']:.2f}",
-            ]
-        )
-
-    headers = [
-        "Model",
-        "Backend",
-        "Device",
-        "BS",
-        "Seq",
-        "Precision",
-        "Steps",
-        "Mem(GiB)",
-        "TPS",
-        "TFLOPS",
-        "MFU(%)",
-    ]
-
-    if rows:
-        print_table(headers, rows)
-        print_note()
-    else:
-        print("No TorchTitan logs found.")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tools/auto_benchmark/run_primus_autobenchmark.sh b/tools/auto_benchmark/run_primus_autobenchmark.sh
index 53201b858..4af17ecfe 100644
--- a/tools/auto_benchmark/run_primus_autobenchmark.sh
+++ b/tools/auto_benchmark/run_primus_autobenchmark.sh
@@ -1,6 +1,11 @@
 #!/usr/bin/env bash
 # Removed: set -e (allow script to continue on errors in benchmarks)
 
+if [[ -z "${BASH_VERSION:-}" ]]; then
+    echo "This script must be run with bash (not sh)." >&2
+    exit 1
+fi
+
 # Set up trap to debug unexpected exits
 trap 'echo "[DEBUG] Script exiting at line $LINENO with exit code $?"' EXIT
 
@@ -29,6 +34,7 @@ PRIMUS_ROOT="/workspace/Primus"
 MEGATRON_BASE_DIR="${PRIMUS_ROOT}/examples/megatron/configs"
 TORCHTITAN_BASE_DIR="${PRIMUS_ROOT}/examples/torchtitan/configs"
 RUN_SCRIPT="${PRIMUS_ROOT}/examples/run_pretrain.sh"
+VALID_DEVICES=(MI300X MI325X MI355X)
 
 # Check if run_pretrain.sh exists, otherwise try run_pretrain_1.sh
 if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; then
@@ -36,6 +42,213 @@ if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; t
     echo "[DEBUG] Using run_pretrain_1.sh instead"
 fi
 
+# ------------------------------------------
+# Helpers
+# ------------------------------------------
+install_vim_editor() {
+    echo -e "${YELLOW}⚠ No editor found. Installing vim...${RESET}"
+
+    if [[ $EUID -eq 0 ]]; then
+        apt-get update && apt-get install -y vim
+    elif command -v sudo &>/dev/null; then
+        sudo apt-get update && sudo apt-get install -y vim
+    else
+        echo -e "${RED}✗ Cannot install vim: not root and sudo is unavailable.${RESET}"
+        echo -e "   ${DOT} Install vim manually, then re-run config editing:"
+        echo -e "   ${CYAN}apt-get update && apt-get install -y vim${RESET}"
+        return 1
+    fi
+}
+
+open_config_editor() {
+    local config_file="$1"
+    local candidate editor_bin editor_args editor_label
+
+    for candidate in \
+        "${EDITOR:-}" \
+        "${VISUAL:-}" \
+        "vim" \
+        "vi" \
+        "nano" \
+        "emacs -nw" \
+        "code --wait" \
+        "cursor --wait"; do
+        if [[ -z "$candidate" ]]; then
+            continue
+        fi
+
+        editor_bin="${candidate%% *}"
+        if ! command -v "$editor_bin" &>/dev/null; then
+            continue
+        fi
+
+        editor_args="${candidate#"$editor_bin"}"
+        editor_label="$candidate"
+        echo -e "   ${DOT} Using editor: ${CYAN}$editor_label${RESET}"
+        # shellcheck disable=SC2086
+        "$editor_bin" $editor_args "$config_file"
+        return 0
+    done
+
+    if install_vim_editor && command -v vim &>/dev/null; then
+        echo -e "   ${DOT} Using editor: ${CYAN}vim${RESET}"
+        vim "$config_file"
+        return 0
+    fi
+
+    echo -e "${RED}✗ Failed to open an editor for:${RESET} ${CYAN}$config_file${RESET}"
+    return 1
+}
+
+next_run_number() {
+    local model_name="$1"
+    local prefix="${model_name}_${BACKEND}_${DEVICE}"
+    local max_run=0
+    local f bn run_n legacy_count=0
+
+    shopt -s nullglob
+    for f in "$LOG_DIR"/"${prefix}"_run*.log; do
+        bn=$(basename "$f" .log)
+        if [[ "$bn" =~ _run([0-9]+)$ ]]; then
+            run_n="${BASH_REMATCH[1]}"
+            if (( run_n > max_run )); then
+                max_run=$run_n
+            fi
+        fi
+    done
+
+    for f in "$LOG_DIR"/"${prefix}"_*.log; do
+        bn=$(basename "$f" .log)
+        if [[ "$bn" =~ _run[0-9]+$ ]]; then
+            continue
+        fi
+        legacy_count=$((legacy_count + 1))
+    done
+    shopt -u nullglob
+
+    echo $((max_run + legacy_count + 1))
+}
+
+prepare_benchmark_artifacts() {
+    local cfg_file="$1"
+    local model_name run_num artifact_prefix
+
+    PREP_CFG_FILE="$cfg_file"
+    PREP_MODEL_NAME=$(basename "$cfg_file" .yaml)
+    model_name="$PREP_MODEL_NAME"
+    run_num=$(next_run_number "$model_name")
+    PREP_RUN_LABEL="run${run_num}"
+    artifact_prefix="${model_name}_${BACKEND}_${DEVICE}_${PREP_RUN_LABEL}"
+
+    PREP_LOG_FILE="$LOG_DIR/${artifact_prefix}.log"
+
+    if [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then
+        PREP_WORKING_CONFIG="${EDITED_CONFIGS[$cfg_file]}"
+    else
+        PREP_WORKING_CONFIG="$cfg_file"
+    fi
+
+    if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then
+        PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_override.yaml"
+        cp "${EDITED_CONFIGS[$cfg_file]:-$cfg_file}" "$PREP_WORKING_CONFIG"
+        for KEY in "${!PARAM_OVERRIDES[@]}"; do
+            sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${PARAM_OVERRIDES[$KEY]}|g" "$PREP_WORKING_CONFIG"
+        done
+    elif [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then
+        PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_edited.yaml"
+        cp "${EDITED_CONFIGS[$cfg_file]}" "$PREP_WORKING_CONFIG"
+    fi
+}
+
+execute_benchmark_run() {
+    local cfg_file="$1"
+    local working_config="$2"
+    local log_file="$3"
+    local model_name="$4"
+    local current="$5"
+    local total="$6"
+
+    local original_config_backup=""
+    local run_exit_code=0
+
+    echo -e "${STAR} ${BOLD}Starting Benchmark ${current}/${total}...${RESET}"
+    echo -e "   ${DOT} Model: ${CYAN}$model_name${RESET}"
+    echo -e "   ${DOT} Backend: ${CYAN}$BACKEND${RESET}"
+    echo -e "   ${DOT} Device: ${CYAN}$DEVICE${RESET}"
+    echo -e "   ${DOT} Config: ${YELLOW}$working_config${RESET}"
+    echo -e "   ${DOT} Log: ${YELLOW}$log_file${RESET}\n"
+
+    if [[ "$working_config" != "$cfg_file" ]]; then
+        original_config_backup="${cfg_file}.backup_$$"
+        cp "$cfg_file" "$original_config_backup"
+        cp "$working_config" "$cfg_file"
+        echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$cfg_file${RESET}"
+    fi
+
+    EXP="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$cfg_file")"
+    export EXP
+    echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n"
+
+    echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}"
+    cd "$PRIMUS_ROOT" || return 1
+
+    set +e
+    bash "$RUN_SCRIPT" 2>&1 | tee "$log_file" || true
+    run_exit_code=$?
+    set +e
+
+    cd "$SCRIPT_DIR" || return 1
+
+    if [[ -n "$original_config_backup" && -f "$original_config_backup" ]]; then
+        mv "$original_config_backup" "$cfg_file"
+        echo -e " ${CHECK} Restored original config file"
+    fi
+
+    echo
+    echo -e "${GREEN}==========================================${RESET}"
+    if [[ $run_exit_code -eq 0 ]]; then
+        echo -e " ${BOLD}${GREEN}✓ Benchmark ${current}/${total} Completed Successfully!${RESET}"
+    else
+        echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${current}/${total} Completed with Exit Code: $run_exit_code${RESET}"
+    fi
+    echo -e " Log saved at:"
+    echo -e "   ${CYAN}$log_file${RESET}"
+    if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then
+        echo -e " Override config saved at:"
+        echo -e "   ${CYAN}$working_config${RESET}"
+    fi
+    echo -e "${GREEN}==========================================${RESET}"
+    echo
+
+    return "$run_exit_code"
+}
+
+generate_metrics_table() {
+    local metrics_script="metrics.py"
+
+    echo
+    echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n"
+
+    if [[ -f "$SCRIPT_DIR/$metrics_script" && ( "$BACKEND" == "megatron" || "$BACKEND" == "torchtitan" ) ]]; then
+        echo -e " ${CHECK} Running: ${CYAN}python $metrics_script $BACKEND${RESET}\n"
+        metrics_output=$(cd "$SCRIPT_DIR" && python "$metrics_script" "$BACKEND")
+        metrics_status=$?
+        printf '%s\n' "$metrics_output"
+        echo
+        if [[ $metrics_status -eq 0 ]]; then
+            csv_path=$(printf '%s\n' "$metrics_output" | awk '/^  Latest:/{print $2}')
+            echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}"
+            if [[ -n "$csv_path" ]]; then
+                echo -e " ${DOT} CSV: ${CYAN}$csv_path${RESET}"
+            fi
+        else
+            echo -e " ${RED}✗ Metrics generation failed${RESET}"
+        fi
+    else
+        echo -e " ${RED}✗ Metrics script not found: ${metrics_script:-unknown}${RESET}"
+    fi
+}
+
 # ------------------------------------------
 # Banner
 # ------------------------------------------
@@ -85,14 +298,49 @@ sleep 0.2
 # ------------------------------------------
 echo -e "${STAR} ${BOLD}Detecting Device...${RESET}"
 
-DEVICE=$(/opt/rocm/bin/rocminfo | grep "AMD Instinct" | head -n1 | awk '{print $5}')
-echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}"
+ROCMINFO=""
+for candidate in \
+    "$(command -v rocminfo 2>/dev/null)" \
+    "/opt/rocm/bin/rocminfo" \
+    "${ROCM_PATH:+$ROCM_PATH/bin/rocminfo}" \
+    /opt/rocm-*/bin/rocminfo; do
+    if [[ -n "$candidate" && -x "$candidate" ]]; then
+        ROCMINFO="$candidate"
+        break
+    fi
+done
 
-if [[ -z "$DEVICE" || "$DEVICE" != "MI300X" && "$DEVICE" != "MI355X" ]]; then
-  ARCH=$(/opt/rocm/bin/rocminfo | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]')
+is_valid_device() {
+    local candidate="$1"
+    for dev in "${VALID_DEVICES[@]}"; do
+        if [[ "$candidate" == "$dev" ]]; then
+            return 0
+        fi
+    done
+    return 1
+}
+
+if [[ -z "$ROCMINFO" ]]; then
+    echo -e " ${YELLOW}⚠ rocminfo not found (checked PATH, /opt/rocm/bin, ROCM_PATH)${RESET}"
+    DEVICE=""
+else
+    echo -e " ${DOT} Using rocminfo: ${CYAN}$ROCMINFO${RESET}"
+    DEVICE=$("$ROCMINFO" 2>/dev/null | grep -oE 'MI3[0-9]{2}X' | head -n1)
+    if [[ -z "$DEVICE" ]]; then
+        DEVICE=$("$ROCMINFO" 2>/dev/null | grep "AMD Instinct" | head -n1 | awk '{print $5}')
+    fi
+    echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}"
+fi
+
+if ! is_valid_device "$DEVICE"; then
+  if [[ -n "$ROCMINFO" ]]; then
+    ARCH=$("$ROCMINFO" 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]')
+  else
+    ARCH=""
+  fi
   case "$ARCH" in
-    "gfx942") DEVICE="MI300X" ;;
     "gfx950") DEVICE="MI355X" ;;
+    # gfx942 is shared by MI300X and MI325X; require manual selection if marketing name is missing
     *) DEVICE="" ;;
   esac
 fi
@@ -101,7 +349,8 @@ if [[ -z "$DEVICE" ]]; then
     echo -e "${RED}✗ Could not detect device automatically${RESET}"
     echo -e "${STAR} ${BOLD}Please select Device manually:${RESET}"
     echo -e "  ${DOT} 1) MI300X"
-    echo -e "  ${DOT} 2) MI355X"
+    echo -e "  ${DOT} 2) MI325X"
+    echo -e "  ${DOT} 3) MI355X"
 
     echo -en " ${ARROW} Enter number or name: "
     read -r DEV_IN
@@ -110,7 +359,10 @@ if [[ -z "$DEVICE" ]]; then
         1|MI300X|mi300x|Mi300x)
             DEVICE="MI300X"
             ;;
-        2|MI355X|mi355x|Mi355x)
+        2|MI325X|mi325x|Mi325x)
+            DEVICE="MI325X"
+            ;;
+        3|MI355X|mi355x|Mi355x)
             DEVICE="MI355X"
             ;;
         *)
@@ -175,10 +427,10 @@ SELECTED_CONFIGS=()
 if [[ "$CFG_NUM" == "all" ]]; then
     # Select all configs
     SELECTED_CONFIGS=("${CONFIG_LIST[@]}")
-elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then
+elif [[ "$CFG_NUM" =~ ^[0-9]+-[0-9]+$ ]]; then
     # Handle range input like 4-8
-    START="${BASH_REMATCH[1]}"
-    END="${BASH_REMATCH[2]}"
+    START="${CFG_NUM%%-*}"
+    END="${CFG_NUM##*-}"
 
     if [[ $START -lt 1 || $END -gt ${#CONFIG_LIST[@]} || $START -gt $END ]]; then
         echo -e "${RED}✗ Invalid range: $START-$END${RESET}"
@@ -190,7 +442,9 @@ elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then
     done
 else
     # Handle comma-separated input
+    _saved_ifs=$IFS
     IFS=',' read -ra CFG_NUMS <<< "$CFG_NUM"
+    IFS=$_saved_ifs
 
     for num in "${CFG_NUMS[@]}"; do
         # Trim whitespace
@@ -205,7 +459,8 @@ else
     done
 fi
 
-echo -e " ${CHECK} Selected ${GREEN}${#SELECTED_CONFIGS[@]}${RESET} config(s):"
+SELECTED_CONFIG_COUNT=${#SELECTED_CONFIGS[@]}
+echo -e " ${CHECK} Selected ${GREEN}${SELECTED_CONFIG_COUNT}${RESET} configs:"
 for cfg in "${SELECTED_CONFIGS[@]}"; do
     echo -e "    ${DOT} $(basename "$cfg")"
 done
@@ -255,7 +510,9 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then
         if [[ "$EDIT_SELECTION" == "all" ]]; then
             MODELS_TO_EDIT=("${!SELECTED_CONFIGS[@]}")
         else
+            _saved_ifs=$IFS
             IFS=',' read -ra EDIT_NUMS <<< "$EDIT_SELECTION"
+            IFS=$_saved_ifs
             MODELS_TO_EDIT=()
             for num in "${EDIT_NUMS[@]}"; do
                 num=$(echo "$num" | xargs)
@@ -277,18 +534,7 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then
             TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml"
             cp "$cfg" "$TEMP_EDIT_CONFIG"
 
-            # Try to find an editor
-            if command -v nano &> /dev/null; then
-                nano "$TEMP_EDIT_CONFIG"
-            elif command -v vim &> /dev/null; then
-                vim "$TEMP_EDIT_CONFIG"
-            elif command -v vi &> /dev/null; then
-                vi "$TEMP_EDIT_CONFIG"
-            elif command -v code &> /dev/null; then
-                code --wait "$TEMP_EDIT_CONFIG"
-            else
-                ${EDITOR:-vi} "$TEMP_EDIT_CONFIG"
-            fi
+            open_config_editor "$TEMP_EDIT_CONFIG"
 
             # Store the edited config
             EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG"
@@ -311,18 +557,7 @@ elif [[ ${#SELECTED_CONFIGS[@]} -eq 1 ]]; then
         TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml"
         cp "$cfg" "$TEMP_EDIT_CONFIG"
 
-        # Try to find an editor
-        if command -v nano &> /dev/null; then
-            nano "$TEMP_EDIT_CONFIG"
-        elif command -v vim &> /dev/null; then
-            vim "$TEMP_EDIT_CONFIG"
-        elif command -v vi &> /dev/null; then
-            vi "$TEMP_EDIT_CONFIG"
-        elif command -v code &> /dev/null; then
-            code --wait "$TEMP_EDIT_CONFIG"
-        else
-            ${EDITOR:-vi} "$TEMP_EDIT_CONFIG"
-        fi
+        open_config_editor "$TEMP_EDIT_CONFIG"
 
         # Store the edited config
         EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG"
@@ -360,7 +595,8 @@ if [[ "$OVERRIDE_PARAMS" == "y" || "$OVERRIDE_PARAMS" == "Y" ]]; then
     done
 
     if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then
-        echo -e "\n ${CHECK} ${GREEN}${#PARAM_OVERRIDES[@]}${RESET} parameter(s) will be overridden\n"
+        PARAM_OVERRIDE_COUNT=${#PARAM_OVERRIDES[@]}
+        echo -e "\n ${CHECK} ${GREEN}${PARAM_OVERRIDE_COUNT}${RESET} parameters will be overridden\n"
     fi
 fi
 
@@ -398,7 +634,8 @@ if [[ "$ADD_ENV_VARS" == "y" || "$ADD_ENV_VARS" == "Y" ]]; then
     done
 
     if [[ ${#DEVICE_ENV_VARS[@]} -gt 0 ]]; then
-        echo -e "\n ${CHECK} ${GREEN}${#DEVICE_ENV_VARS[@]}${RESET} environment variable(s) will be set\n"
+        DEVICE_ENV_VAR_COUNT=${#DEVICE_ENV_VARS[@]}
+        echo -e "\n ${CHECK} ${GREEN}${DEVICE_ENV_VAR_COUNT}${RESET} environment variables will be set\n"
     fi
 fi
 
@@ -453,105 +690,30 @@ for CFG_FILE in "${SELECTED_CONFIGS[@]}"; do
     echo -e "${MAGENTA}${BOLD}║  CONFIG FILE: $(basename "$CFG_FILE")${RESET}"
     echo -e "${MAGENTA}${BOLD}╚════════════════════════════════════════════════════════════╝${RESET}\n"
 
-    # Extract full filename without extension to preserve all details
-    CONFIG_FILENAME=$(basename "$CFG_FILE" .yaml)
-    MODEL_NAME="$CONFIG_FILENAME"
-    TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
+    prepare_benchmark_artifacts "$CFG_FILE"
 
-    # Use the full config filename in the log name to preserve all details
-    LOG_FILE="$LOG_DIR/${CONFIG_FILENAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}.log"
-
-    # Use edited config if available, otherwise use original
-    if [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then
-        WORKING_CONFIG="${EDITED_CONFIGS[$CFG_FILE]}"
-        echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$MODEL_NAME${RESET}"
-    else
-        WORKING_CONFIG="$CFG_FILE"
+    if [[ -n "${EDITED_CONFIGS[$CFG_FILE]:-}" && ${#PARAM_OVERRIDES[@]} -eq 0 ]]; then
+        echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$PREP_MODEL_NAME${RESET}"
     fi
-
-    # Apply parameter overrides if any
+    echo -e "   ${DOT} Run label: ${CYAN}$PREP_RUN_LABEL${RESET}"
     if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then
-        OVERRIDE_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_override.yaml"
-        cp "$WORKING_CONFIG" "$OVERRIDE_CONFIG"
-
         echo -e "${STAR} ${BOLD}Applying parameter overrides...${RESET}"
         for KEY in "${!PARAM_OVERRIDES[@]}"; do
-            VALUE="${PARAM_OVERRIDES[$KEY]}"
-            # Use sed to replace the parameter value in YAML (handles both 'key: value' and 'key:value')
-            sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${VALUE}|g" "$OVERRIDE_CONFIG"
-            echo -e "   ${DOT} ${CYAN}$KEY${RESET}: $VALUE"
+            echo -e "   ${DOT} ${CYAN}$KEY${RESET}: ${PARAM_OVERRIDES[$KEY]}"
         done
-        WORKING_CONFIG="$OVERRIDE_CONFIG"
-        echo -e " ${CHECK} Override config saved: ${YELLOW}$OVERRIDE_CONFIG${RESET}\n"
-    elif [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then
-        # Save edited config to logs directory
-        SAVED_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_edited.yaml"
-        cp "$WORKING_CONFIG" "$SAVED_CONFIG"
-        WORKING_CONFIG="$SAVED_CONFIG"
+        echo -e " ${CHECK} Override config saved: ${YELLOW}$PREP_WORKING_CONFIG${RESET}\n"
     fi
 
-    echo -e "${STAR} ${BOLD}Starting Benchmark ${CURRENT}/${TOTAL_CONFIGS}...${RESET}"
-    echo -e "   ${DOT} Model: ${CYAN}$MODEL_NAME${RESET}"
-    echo -e "   ${DOT} Backend: ${CYAN}$BACKEND${RESET}"
-    echo -e "   ${DOT} Device: ${CYAN}$DEVICE${RESET}"
-    echo -e "   ${DOT} Config: ${YELLOW}$WORKING_CONFIG${RESET}"
-    echo -e "   ${DOT} Log: ${YELLOW}$LOG_FILE${RESET}\n"
-
-    # Set EXP to the working config (edited or overridden version if available)
-    # For edited/overridden configs, we need to copy them to the expected location
-    if [[ "$WORKING_CONFIG" != "$CFG_FILE" ]]; then
-        # Config was edited or overridden, copy it to the original location temporarily
-        ORIGINAL_CONFIG_BACKUP="${CFG_FILE}.backup_$$"
-        cp "$CFG_FILE" "$ORIGINAL_CONFIG_BACKUP"
-        cp "$WORKING_CONFIG" "$CFG_FILE"
-        echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$CFG_FILE${RESET}"
-    fi
-
-    # Set EXP to the device-specific config path (now contains edited content if applicable)
-    EXP_CONFIG_PATH="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$CFG_FILE")"
-    export EXP="$EXP_CONFIG_PATH"
-    echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n"
-
-    # Change to Primus root directory before running the script
-    echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}"
-    cd "$PRIMUS_ROOT"
-
-    # Run the script and capture exit code, but don't stop on failure
-    # Use 'set +e' locally to ensure we continue even on errors
-    set +e
-    bash $RUN_SCRIPT 2>&1 | tee "$LOG_FILE" || true
-    RUN_EXIT_CODE=$?
-    set -e
-
-    # Return to script directory
-    cd "$SCRIPT_DIR"
-
-    # Restore original config if it was temporarily replaced
-    if [[ -n "$ORIGINAL_CONFIG_BACKUP" && -f "$ORIGINAL_CONFIG_BACKUP" ]]; then
-        mv "$ORIGINAL_CONFIG_BACKUP" "$CFG_FILE"
-        echo -e " ${CHECK} Restored original config file"
-        unset ORIGINAL_CONFIG_BACKUP
-    fi
-
-    echo
-    echo -e "${GREEN}==========================================${RESET}"
-    if [[ $RUN_EXIT_CODE -eq 0 ]]; then
-        echo -e " ${BOLD}${GREEN}✓ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed Successfully!${RESET}"
-    else
-        echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed with Exit Code: $RUN_EXIT_CODE${RESET}"
-    fi
-    echo -e " Log saved at:"
-    echo -e "   ${CYAN}$LOG_FILE${RESET}"
-    if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then
-        echo -e " Override config saved at:"
-        echo -e "   ${CYAN}$OVERRIDE_CONFIG${RESET}"
-    fi
-    echo -e "${GREEN}==========================================${RESET}"
-    echo
+    execute_benchmark_run \
+        "$PREP_CFG_FILE" \
+        "$PREP_WORKING_CONFIG" \
+        "$PREP_LOG_FILE" \
+        "$PREP_MODEL_NAME" \
+        "$CURRENT" \
+        "$TOTAL_CONFIGS" || true
 
     CURRENT=$((CURRENT + 1))
 
-    # Add a short delay between runs
     if [[ $CURRENT -le $TOTAL_CONFIGS ]]; then
         echo -e "${YELLOW}Preparing next benchmark...${RESET}\n"
         echo -e "${INFO} ${BOLD}Next: Config ${CURRENT}/${TOTAL_CONFIGS}${RESET}\n"
@@ -561,26 +723,10 @@ done
 
 echo
 echo -e "${MAGENTA}${BOLD}=========================================${RESET}"
-echo -e "${MAGENTA}${BOLD}  All ${TOTAL_CONFIGS} Benchmark(s) Completed!${RESET}"
+echo -e "${MAGENTA}${BOLD}  All ${TOTAL_CONFIGS} benchmarks completed!${RESET}"
 echo -e "${MAGENTA}${BOLD}=========================================${RESET}"
 
 # ------------------------------------------
 # 7. GENERATE METRICS TABLE
 # ------------------------------------------
-echo
-echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n"
-
-if [[ "$BACKEND" == "megatron" ]]; then
-    METRICS_SCRIPT="metrics_megatron.py"
-elif [[ "$BACKEND" == "torchtitan" ]]; then
-    METRICS_SCRIPT="metrics_torchtitan.py"
-fi
-
-if [[ -f "$METRICS_SCRIPT" ]]; then
-    echo -e " ${CHECK} Running: ${CYAN}python $METRICS_SCRIPT${RESET}\n"
-    python "$METRICS_SCRIPT"
-    echo
-    echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}"
-else
-    echo -e " ${RED}✗ Metrics script not found: $METRICS_SCRIPT${RESET}"
-fi
+generate_metrics_table
diff --git a/tools/backend_gap_report/build_site_bundle.py b/tools/backend_gap_report/build_site_bundle.py
index 134124fce..998b574eb 100644
--- a/tools/backend_gap_report/build_site_bundle.py
+++ b/tools/backend_gap_report/build_site_bundle.py
@@ -14,6 +14,11 @@
 REPO_ROOT = Path(__file__).resolve().parents[2]
 BACKEND_GAP_DOCS_ROOT = REPO_ROOT / "docs" / "backend-gap"
 SITE_SOURCE_ROOT = REPO_ROOT / "tools" / "backend_gap_report" / "site"
+# Declarative registry of extra, independent static sites mounted as subpaths of
+# the single Pages bundle. A repo has exactly ONE Pages site, so this one builder
+# aggregates every section; each is decoupled (add a manifest entry, no code
+# change). See pages-sections.json.
+PAGES_SECTIONS_MANIFEST = Path(__file__).resolve().with_name("pages-sections.json")
 SOURCE_DASHBOARD_DATA_DIR = BACKEND_GAP_DOCS_ROOT / "dashboard-data"
 METADATA_REPORTS_DIR = SOURCE_DASHBOARD_DATA_DIR / "reports"
 PDF_TEMPLATE = REPO_ROOT / "tools" / "backend_gap_report" / "templates" / "pdf-report.css"
@@ -254,6 +259,89 @@ def validate_bundle(bundle_dir: Path) -> None:
     if not isinstance(combined_payload["reports"], list):
         fail("combined reports index field 'reports' must be a list")
 
+    sections_index = bundle_dir / "sections.json"
+    if not sections_index.exists():
+        fail(f"bundle missing sections manifest: {sections_index}")
+    sections_payload = load_json(sections_index)
+    if not isinstance(sections_payload.get("sections"), list):
+        fail("sections.json field 'sections' must be a list")
+    if not isinstance(sections_payload.get("links", []), list):
+        fail("sections.json field 'links' must be a list")
+    for section in sections_payload["sections"]:
+        section_dir = bundle_dir / section["name"]
+        if not (section_dir / "index.html").exists():
+            fail(f"mounted section '{section['name']}' missing index.html in bundle")
+
+
+def mount_extra_sections(output_dir: Path) -> None:
+    """Aggregate manifest-declared sections/links into the bundle.
+
+    ``sections`` are static dirs copied to ``/<name>/``; ``links`` are nav-only
+    entries. Both are written to ``sections.json`` so the shell renders nav
+    without hardcoding any entry. A missing section source is skipped with a
+    warning. See pages-sections.json.
+    """
+
+    mounted: list[dict] = []
+    links: list[dict] = []
+    if PAGES_SECTIONS_MANIFEST.exists():
+        manifest = load_json(PAGES_SECTIONS_MANIFEST)
+        sections = manifest.get("sections", [])
+        if not isinstance(sections, list):
+            fail(f"{PAGES_SECTIONS_MANIFEST}: 'sections' must be a list")
+
+        raw_links = manifest.get("links", [])
+        if not isinstance(raw_links, list):
+            fail(f"{PAGES_SECTIONS_MANIFEST}: 'links' must be a list")
+        for link in raw_links:
+            if not isinstance(link, dict):
+                fail(f"{PAGES_SECTIONS_MANIFEST}: each link must be an object")
+            title = link.get("title")
+            path = link.get("path")
+            if not isinstance(title, str) or not title.strip():
+                fail(f"{PAGES_SECTIONS_MANIFEST}: link missing valid 'title'")
+            if not isinstance(path, str) or not path.strip():
+                fail(f"{PAGES_SECTIONS_MANIFEST}: link '{title}' missing valid 'path'")
+            entry = {"title": title, "path": path}
+            if link.get("probe"):
+                entry["probe"] = True
+            if isinstance(link.get("hint"), str) and link["hint"].strip():
+                entry["hint"] = link["hint"]
+            links.append(entry)
+        seen: set[str] = set()
+        for entry in sections:
+            if not isinstance(entry, dict):
+                fail(f"{PAGES_SECTIONS_MANIFEST}: each section must be an object")
+            name = entry.get("name")
+            source = entry.get("source")
+            if not isinstance(name, str) or not name.strip():
+                fail(f"{PAGES_SECTIONS_MANIFEST}: section missing valid 'name'")
+            if "/" in name or name in {".", ".."}:
+                fail(f"{PAGES_SECTIONS_MANIFEST}: section 'name' must be a single path segment: {name!r}")
+            if not isinstance(source, str) or not source.strip():
+                fail(f"{PAGES_SECTIONS_MANIFEST}: section '{name}' missing valid 'source'")
+            if name in seen:
+                fail(f"{PAGES_SECTIONS_MANIFEST}: duplicate section name: {name}")
+            seen.add(name)
+
+            source_dir = ensure_within(REPO_ROOT, REPO_ROOT / source, f"section '{name}' source")
+            title = entry.get("title") or name
+            if not source_dir.exists():
+                print(
+                    f"[sections] skip '{name}': source not present in checkout ({source})",
+                    flush=True,
+                )
+                continue
+            if not (source_dir / "index.html").exists():
+                fail(f"section '{name}' source has no index.html: {source}")
+            print(f"[sections] mount '{name}' -> /{name}/ (from {source})", flush=True)
+            copy_tree(source_dir, output_dir / name)
+            mounted.append({"name": name, "title": title, "path": f"./{name}/"})
+
+    (output_dir / "sections.json").write_text(
+        json.dumps({"sections": mounted, "links": links}, indent=2) + "\n", encoding="utf-8"
+    )
+
 
 def build_site(output_dir: Path) -> None:
     run_build_index()
@@ -271,6 +359,7 @@ def build_site(output_dir: Path) -> None:
     copy_tree(SOURCE_DASHBOARD_DATA_DIR, output_dir / "dashboard-data")
     build_combined_reports_index(output_dir)
     build_pdf_artifacts(output_dir)
+    mount_extra_sections(output_dir)
     print("[backend-gap] Validate standalone dashboard bundle", flush=True)
     validate_bundle(output_dir)
 
diff --git a/tools/backend_gap_report/pages-sections.json b/tools/backend_gap_report/pages-sections.json
new file mode 100644
index 000000000..85822bb7d
--- /dev/null
+++ b/tools/backend_gap_report/pages-sections.json
@@ -0,0 +1,18 @@
+{
+  "$schema-note": "Declarative registry for the single Pages bundle, aggregated by build_site_bundle.py. 'sections' = static site dirs (repo-relative 'source' with an index.html) mounted at /<name>/ — also add the source path to the deploy workflow push triggers. 'links' = homepage nav links only (nothing copied), e.g. the pip index; probe=true shows the link only when its path is reachable (hidden in local previews).",
+  "sections": [
+    {
+      "name": "deepseek-v4-projection",
+      "title": "DeepSeek-V4 Projection",
+      "source": "examples/deepseek-v4/projection/site"
+    }
+  ],
+  "links": [
+    {
+      "title": "Install (pip)",
+      "path": "./simple/",
+      "probe": true,
+      "hint": "pip install primus --extra-index-url <this-site>/simple/"
+    }
+  ]
+}
diff --git a/tools/backend_gap_report/site/assets/dashboard.css b/tools/backend_gap_report/site/assets/dashboard.css
index ccae7bc2d..60c9c0de1 100644
--- a/tools/backend_gap_report/site/assets/dashboard.css
+++ b/tools/backend_gap_report/site/assets/dashboard.css
@@ -163,6 +163,58 @@ code {
   font-weight: 600;
 }
 
+/* ---------- Hero nav (data-driven links to sibling sites) ---------- */
+.hero__nav {
+  margin-top: 1.4rem;
+  display: flex;
+  align-items: center;
+  gap: 0.75rem;
+  flex-wrap: wrap;
+}
+
+.hero__nav[hidden] {
+  display: none;
+}
+
+.hero__nav-label {
+  font-size: 0.72rem;
+  letter-spacing: 0.14em;
+  text-transform: uppercase;
+  color: rgba(255, 255, 255, 0.6);
+  font-weight: 600;
+}
+
+.hero__nav-links {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 0.5rem;
+}
+
+.hero__nav-link {
+  display: inline-flex;
+  align-items: center;
+  gap: 0.4rem;
+  padding: 0.4rem 0.85rem;
+  border-radius: 999px;
+  border: 1px solid rgba(255, 255, 255, 0.28);
+  background: rgba(255, 255, 255, 0.08);
+  color: #ffffff;
+  font-size: 0.85rem;
+  font-weight: 600;
+}
+
+.hero__nav-link:hover {
+  text-decoration: none;
+  background: rgba(255, 255, 255, 0.16);
+  border-color: rgba(255, 255, 255, 0.5);
+}
+
+.hero__nav-link::after {
+  content: "\2197";
+  font-size: 0.85em;
+  opacity: 0.7;
+}
+
 /* ---------- Section tabs ---------- */
 .section-tabs {
   display: flex;
diff --git a/tools/backend_gap_report/site/assets/dashboard.js b/tools/backend_gap_report/site/assets/dashboard.js
index 60578ad04..9f624187e 100644
--- a/tools/backend_gap_report/site/assets/dashboard.js
+++ b/tools/backend_gap_report/site/assets/dashboard.js
@@ -1,6 +1,10 @@
 "use strict";
 
 const BACKEND_GAP_DATA_URL = "./dashboard-data/index.json";
+// Data-driven navigation to independent sibling sites mounted as subpaths.
+// Generated by build_site_bundle.py from pages-sections.json — the shell never
+// hardcodes any section, so new content appears here automatically.
+const SECTIONS_DATA_URL = "./sections.json";
 // Unified weekly + monthly index. The `weekly` namespace below renders all
 // periodic engineering reports regardless of cadence; cadence is a data field.
 const REPORTS_DATA_URL = "./reports-data/index.json";
@@ -30,6 +34,8 @@ const el = (id) => document.getElementById(id);
 
 const elements = {
   heroChips: el("hero-chips"),
+  heroNav: el("hero-nav"),
+  heroNavLinks: el("hero-nav-links"),
 
   weekly: {
     statRow: el("weekly-stat-row"),
@@ -728,10 +734,64 @@ async function fetchJsonOptional(url) {
   }
 }
 
+function addNavLink(href, text, title) {
+  if (!elements.heroNavLinks) return;
+  const link = document.createElement("a");
+  link.className = "hero__nav-link";
+  link.href = href;
+  link.textContent = text;
+  if (title) link.title = title;
+  elements.heroNavLinks.append(link);
+}
+
+function syncNavVisibility() {
+  if (!elements.heroNav || !elements.heroNavLinks) return;
+  elements.heroNav.hidden = elements.heroNavLinks.children.length === 0;
+}
+
+async function isReachable(url) {
+  try {
+    const res = await fetch(url, { method: "HEAD", cache: "no-store" });
+    return res.ok;
+  } catch (_) {
+    return false;
+  }
+}
+
+// Renders the homepage nav from the generated sections.json: `sections` are
+// mounted static subpaths (always present); `links` are declared nav links
+// (e.g. the pip index produced by the deploy workflow). A link with `probe` is
+// shown only if reachable, so it stays hidden in local previews. Both come from
+// pages-sections.json, so the shell hardcodes no specific section or link.
+async function renderSectionNav(payload) {
+  if (!elements.heroNav || !elements.heroNavLinks) return;
+  const sections = (payload && Array.isArray(payload.sections) && payload.sections) || [];
+  sections.forEach((section) => {
+    if (!section || !section.path) return;
+    addNavLink(section.path, section.title || section.name || section.path);
+  });
+  syncNavVisibility();
+
+  const links = (payload && Array.isArray(payload.links) && payload.links) || [];
+  await Promise.all(
+    links.map(async (link) => {
+      if (!link || !link.path || !link.title) return;
+      if (link.probe && !(await isReachable(link.path))) return;
+      addNavLink(link.path, link.title, link.hint);
+      syncNavVisibility();
+    })
+  );
+}
+
 async function loadDashboard() {
   attachWeeklyFilters();
   attachBackendGapFilters();
 
+  const sectionsResult = await fetchJsonOptional(SECTIONS_DATA_URL);
+  if (sectionsResult.ok) {
+    renderSectionNav(sectionsResult.payload);
+  }
+
   const reportsResult = await fetchJsonOptional(REPORTS_DATA_URL);
   const gapResult = await fetchJsonOptional(BACKEND_GAP_DATA_URL);
 
diff --git a/tools/backend_gap_report/site/index.html b/tools/backend_gap_report/site/index.html
index d3a4090ee..6d1400e12 100644
--- a/tools/backend_gap_report/site/index.html
+++ b/tools/backend_gap_report/site/index.html
@@ -18,6 +18,10 @@ <h1>Primus Dashboard</h1>
             sync recommendations that need attention.
           </p>
           <div class="hero__chips" id="hero-chips"></div>
+          <nav class="hero__nav" id="hero-nav" aria-label="Other Primus sites" hidden>
+            <span class="hero__nav-label">Explore</span>
+            <div class="hero__nav-links" id="hero-nav-links"></div>
+          </nav>
         </div>
       </div>
     </header>
diff --git a/tools/checkpoint_conversion/convert_flux_hf_to_primus.py b/tools/checkpoint_conversion/convert_flux_hf_to_primus.py
new file mode 100755
index 000000000..b73b95b9a
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_flux_hf_to_primus.py
@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Convert HuggingFace Flux checkpoint to Primus format.
+
+This tool converts HuggingFace Diffusers Flux transformer checkpoints to
+Primus/Megatron-Core compatible format. It handles:
+- QKV weight fusion for grouped-query attention (GQA)
+- Key mapping from HF to Primus naming conventions
+- Single block proj_out splitting
+- Multi-file safetensors loading
+
+Usage:
+    # Convert FLUX.1-dev checkpoint
+    python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\
+        --input black-forest-labs/FLUX.1-dev \\
+        --output checkpoints/primus_flux_12b.safetensors \\
+        --variant flux_12b
+
+    # Convert with custom architecture
+    python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\
+        --input path/to/checkpoint \\
+        --output checkpoints/primus_custom.safetensors \\
+        --variant custom \\
+        --num-joint-layers 10 \\
+        --num-single-layers 20
+
+Example:
+    $ python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\
+        --input black-forest-labs/FLUX.1-dev \\
+        --output checkpoints/primus_flux_12b.safetensors \\
+        --variant flux_12b
+
+    Converting flux_12b checkpoint
+      Input: black-forest-labs/FLUX.1-dev
+      Output: checkpoints/primus_flux_12b.safetensors
+      Architecture: 19 joint + 38 single layers
+    Loading HuggingFace checkpoint from: black-forest-labs/FLUX.1-dev
+    ...
+    Conversion complete!
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+# Add primus to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from primus.backends.megatron.core.models.diffusion.flux import (
+    FluxConfig,
+    convert_hf_checkpoint,
+)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert HuggingFace Flux checkpoint to Primus format",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Convert FLUX.1-dev (12B parameters)
+  %(prog)s --input black-forest-labs/FLUX.1-dev \\
+           --output checkpoints/primus_flux_12b.safetensors \\
+           --variant flux_12b
+
+  # Convert local checkpoint directory
+  %(prog)s --input /path/to/flux/checkpoint \\
+           --output primus_flux.safetensors \\
+           --variant flux_12b
+
+  # Convert with custom architecture
+  %(prog)s --input /path/to/checkpoint \\
+           --output primus_custom.safetensors \\
+           --variant custom \\
+           --num-joint-layers 10 \\
+           --num-single-layers 20
+        """,
+    )
+
+    parser.add_argument(
+        "--input",
+        required=True,
+        help="Path to HF checkpoint (file, directory, or HF model ID like 'black-forest-labs/FLUX.1-dev')",
+    )
+    parser.add_argument("--output", required=True, help="Output path for Primus checkpoint (.safetensors)")
+    parser.add_argument(
+        "--variant",
+        choices=["flux_535m", "flux_12b", "custom"],
+        default="flux_12b",
+        help="Flux variant (determines architecture). Default: flux_12b",
+    )
+    parser.add_argument(
+        "--num-joint-layers", type=int, help="Number of joint layers (only with --variant custom)"
+    )
+    parser.add_argument(
+        "--num-single-layers", type=int, help="Number of single layers (only with --variant custom)"
+    )
+
+    args = parser.parse_args()
+
+    # Validate custom variant arguments
+    if args.variant == "custom":
+        if not args.num_joint_layers or not args.num_single_layers:
+            parser.error("--num-joint-layers and --num-single-layers are required with --variant custom")
+    else:
+        if args.num_joint_layers or args.num_single_layers:
+            parser.error("--num-joint-layers and --num-single-layers can only be used with --variant custom")
+
+    # Create config based on variant
+    if args.variant == "flux_535m":
+        config = FluxConfig.flux_535m()
+    elif args.variant == "flux_12b":
+        config = FluxConfig.flux_12b()
+    else:  # custom
+        config = FluxConfig(num_joint_layers=args.num_joint_layers, num_single_layers=args.num_single_layers)
+
+    # Print conversion info
+    print("=" * 80)
+    print(f"Converting {args.variant} checkpoint")
+    print(f"  Input:  {args.input}")
+    print(f"  Output: {args.output}")
+    print(f"  Architecture: {config.num_joint_layers} joint + {config.num_single_layers} single layers")
+    print("=" * 80)
+    print()
+
+    # Convert checkpoint
+    try:
+        convert_hf_checkpoint(
+            checkpoint_path=args.input,
+            flux_config=config,
+            save_to=args.output,
+        )
+
+        print()
+        print("=" * 80)
+        print("✓ Conversion complete!")
+        print(f"Primus checkpoint saved to: {args.output}")
+        print("=" * 80)
+
+        return 0
+
+    except Exception as e:
+        print()
+        print("=" * 80)
+        print(f"✗ Conversion failed: {e}")
+        print("=" * 80)
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/ci/check_config_schema.py b/tools/ci/check_config_schema.py
new file mode 100644
index 000000000..1b56dac2b
--- /dev/null
+++ b/tools/ci/check_config_schema.py
@@ -0,0 +1,322 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Fail CI when a Primus YAML sets a backend key upstream no longer defines.
+
+Every backend adapter accepts unknown keys silently, so a field that upstream
+renamed or moved keeps parsing while doing nothing. Prints the drift as Markdown
+for `$GITHUB_STEP_SUMMARY`:
+
+    python tools/ci/check_config_schema.py --warn-only >> "$GITHUB_STEP_SUMMARY"
+
+A key upstream looks to have merely renamed gets a table and a row of its own:
+that is a fix someone applies one key at a time. The rest share a table where a
+row is a set of configs, because dozens of dead keys are usually one preset's
+worth of rot, and repeating the same three long paths for each of them buries
+the part a reader has to act on.
+
+TorchTitan, MaxText and Megatron each have their schema read from their own
+`third_party/` submodule. A backend whose submodule is not checked out has
+nothing to compare against and is skipped, which exits non-zero even under
+`--warn-only`: a checkout that fetches only some of them drops that part of the
+check, and a check that did not run must not look like one that passed.
+
+A config that cannot be loaded (broken `extends:` path, unparseable YAML) fails
+the check just like drift does -- an unchecked config must never be counted as a
+passing one. `--warn-only` suppresses both.
+
+Misplaced model-scoped keys get their own table. Such a key is not unknown --
+some Primus config class does declare it -- but that class is built for one
+model only, so setting the key on any other model is the same silent no-op.
+"""
+
+import argparse
+import sys
+import time
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Sequence
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from primus.core.config.schema_check import (  # noqa: E402
+    BACKENDS,
+    DEFAULT_ALLOWLIST,
+    KeyFinding,
+    build_allowlist,
+    check_manual_allowlist_consumers,
+    dedupe,
+    dedupe_scoped,
+    scan_backend,
+)
+
+SHOWN_PATTERNS = 40
+SHOWN_FILES = 3
+
+
+def shorten_paths(paths: Iterable[str]) -> dict[str, str]:
+    """Map every path to the fewest trailing segments no other path here also ends with.
+
+    Almost every config sits under one of a handful of roots, so the shared
+    `primus/configs/models/megatron/` is pure repetition -- but only where
+    dropping it still names one file, and `MI300X/` and `MI355X/` hold configs
+    with identical names.
+    """
+    parts = {path: path.split("/") for path in set(paths)}
+    tails = Counter("/".join(s[-d:]) for s in parts.values() for d in range(1, len(s) + 1))
+    short = {}
+    for path, segments in parts.items():
+        # The whole path always names itself, so it is the fallback rather than a candidate.
+        candidates = ("/".join(segments[-d:]) for d in range(1, len(segments)))
+        short[path] = next((tail for tail in candidates if tails[tail] == 1), path)
+    return short
+
+
+@dataclass(frozen=True)
+class DriftRow:
+    """One row of the drift table: the keys that live in exactly the same configs."""
+
+    backend: str
+    keys: tuple[str, ...]
+    count: int
+    files: tuple[str, ...]
+    suggestion: str | None
+
+    def sample(self, short: dict[str, str], limit: int = SHOWN_FILES) -> str:
+        shown = ", ".join(f"`{short.get(f, f)}`" for f in self.files[:limit])
+        extra = len(self.files) - limit
+        return f"{shown} (+{extra} more)" if extra > 0 else shown
+
+
+def group_rows(rows: Sequence[KeyFinding]) -> tuple[list[DriftRow], list[DriftRow]]:
+    """Split the findings into ``(renamed, merged)``, the two tables.
+
+    A key with a suggested rename keeps its own row: that one is a direct fix
+    someone applies key by key, and burying it in a list of twenty-five is the
+    opposite of the point. The rest say the same thing many times over, so the
+    ones that drifted out of the very same configs share a row -- which loses
+    nothing, the row's configs being every member's.
+    """
+    renamed, shared = [], {}
+    for row in rows:
+        if row.suggestion:
+            renamed.append(DriftRow(row.backend, (row.key,), row.count, row.files, row.suggestion))
+        else:
+            shared.setdefault((row.backend, row.count, row.files), []).append(row.key)
+    merged = [
+        DriftRow(backend, tuple(sorted(keys)), count, files, None)
+        for (backend, count, files), keys in shared.items()
+    ]
+
+    def order(row: DriftRow) -> tuple[str, int, str]:
+        return row.backend, -row.count, row.keys[0]
+
+    return sorted(renamed, key=order), sorted(merged, key=order)
+
+
+def subsection(title: str, lead: str) -> list[str]:
+    """A `###` heading and its lead-in, spaced the way Markdown wants them."""
+    return ["", f"### {title}", "", lead, ""]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--backend", action="append", choices=BACKENDS, help="Limit to one backend.")
+    parser.add_argument(
+        "--path",
+        action="append",
+        type=Path,
+        help="Scan these files/dirs instead of the default roots (requires a single --backend).",
+    )
+    parser.add_argument("--allowlist", type=Path, help="Override the allowlist YAML path.")
+    parser.add_argument(
+        "--warn-only",
+        action="store_true",
+        help="Report findings but exit 0. A skipped backend still fails.",
+    )
+    parser.add_argument(
+        "--show-allowlist", action="store_true", help="Also print the derived allowlist patterns."
+    )
+    args = parser.parse_args()
+    if args.path and len(args.backend or BACKENDS) != 1:
+        parser.error("--path requires exactly one --backend")
+    return args
+
+
+def main():
+    args = parse_args()
+    backends = args.backend or list(BACKENDS)
+    started = time.perf_counter()
+
+    lines = ["## Backend config schema drift", ""]
+    rows, skipped, errors, checked = [], [], [], 0
+    stale, scoped, model_scopes = [], [], []
+
+    allowlist_path = args.allowlist or (ROOT / DEFAULT_ALLOWLIST)
+    for backend in backends:
+        for problem in check_manual_allowlist_consumers(ROOT, allowlist_path, backend):
+            if problem not in stale:  # the `common` scope is read once per backend
+                stale.append(problem)
+
+        result = scan_backend(ROOT, backend, args.path, args.allowlist)
+        if not result.available:
+            skipped.append(backend)
+            continue
+        checked += result.checked
+        errors.extend(result.errors)
+        rows.extend(dedupe(result.findings, result.schema))
+        scoped.extend(dedupe_scoped(result.scoped))
+        model_scopes.extend((backend, scope) for scope in result.scopes)
+
+    # Loud, and its own section: a backend nobody compared against anything is
+    # the one outcome that must not read like a pass. It is also the one that
+    # `--warn-only` does not forgive, so say so where the reader is looking.
+    if skipped:
+        lines.append(f"### FAILED -- {len(skipped)} backend(s) not checked")
+        lines.append("")
+        lines.append(
+            f"**{', '.join(skipped)}**: upstream schema not found under `third_party/`, so no "
+            "config of theirs was compared against anything. Run `git submodule update --init` "
+            "to restore the check. Drift is warn-only; a check that never ran is not."
+        )
+        lines.append("")
+    # Reported even when no schema is available: an allowlist entry whose
+    # consumer is gone widens the legal key set, which is the blindness this
+    # check exists to remove, and verifying it needs no submodule.
+    if stale:
+        lines.append(f"**{len(stale)} allowlist entr(y/ies) no longer backed by their consumer:**")
+        lines.append("")
+        for problem in stale:
+            lines.append(f"- {problem}")
+        lines.append("")
+
+    if len(skipped) == len(backends):
+        lines.append("Nothing to check." if not stale else "No schema available; allowlist checked.")
+        print("\n".join(lines))
+        return 1
+
+    if rows:
+        renamed, merged = group_rows(rows)
+        short = shorten_paths(f for row in rows for f in row.files)
+        shared = len(rows) - len(renamed)
+        lines.append(
+            f"**{len(rows)} unknown key(s)** across {sum(r.count for r in rows)} occurrence(s). "
+            "Each one parses and then does nothing: fix the config, or add the key to "
+            "`tools/ci/config_schema_allowlist.yaml` with a reason and a consumer. Example files "
+            "are shortened to the shortest trailing path that names one config here."
+        )
+        # Two tables, because the two halves are two different jobs. A rename is
+        # a key-by-key edit with a known destination; the rest is a decision per
+        # set of configs, and a `Likely moved to` column of nothing but dashes.
+        if renamed:
+            lines += subsection(
+                "Keys that look renamed rather than dropped",
+                f"{len(renamed)} of them: upstream still defines a name this close, so each is one "
+                "direct edit and gets a row to itself.",
+            )
+            lines.append("| Backend | Unknown key | Configs | Likely moved to | Example files |")
+            lines.append("| --- | --- | ---: | --- | --- |")
+            for row in renamed:
+                lines.append(
+                    f"| {row.backend} | `{row.keys[0]}` | {row.count} | "
+                    f"`{row.suggestion}` | {row.sample(short)} |"
+                )
+        if merged:
+            lines += subsection(
+                "Keys with no obvious replacement",
+                f"{'The other' if renamed else 'All'} {shared} drifted out of only {len(merged)} "
+                "distinct set(s) of configs, so a row here is one such set and every key that "
+                "drifted out of exactly it.",
+            )
+            lines.append("| Backend | Configs | Example files | Unknown key(s) |")
+            lines.append("| --- | ---: | --- | --- |")
+            for row in merged:
+                keys = ", ".join(f"`{k}`" for k in row.keys)
+                lines.append(f"| {row.backend} | {row.count} | {row.sample(short)} | {keys} |")
+    elif not errors:
+        lines.append(
+            f"No drift: every key in {checked} config(s) exists upstream " "or is a known Primus extension."
+        )
+    elif checked:
+        lines.append(f"No drift in the {checked} config(s) that loaded.")
+
+    if scoped:
+        if lines[-1]:
+            lines.append("")
+        lines.append("### Keys set on a model that cannot read them")
+        lines.append("")
+        lines.append("| Backend | Key | Only reaches | Set instead on | Configs | Example |")
+        lines.append("| --- | --- | --- | --- | ---: | --- |")
+        for row in scoped:
+            models = ", ".join(f"`{m}`" for m in row.models)
+            lines.append(
+                f"| {row.backend} | `{row.key}` | `{row.scope.describe()}` "
+                f"({row.scope.config_class}) | {models} | {row.count} | {row.sample(2)} |"
+            )
+        lines.append("")
+        lines.append(
+            f"**{len(scoped)} model-scoped key(s)** set on {sum(r.count for r in scoped)} config(s) "
+            "that build a different model. The key parses, no config class carries it, and the "
+            "value is silently replaced by the upstream default: drop it, or use the field the "
+            "model in question actually reads."
+        )
+
+    # An unreadable config was never checked, so it cannot count towards a pass.
+    if errors:
+        if lines[-1]:
+            lines.append("")
+        lines.append(f"**{len(errors)} config(s) could not be loaded and were NOT checked:**")
+        lines.append("")
+        for error in errors:
+            lines.append(f"- `{error.file}`: {error.message}")
+
+    if args.show_allowlist:
+        lines.append("")
+        lines.append("<details><summary>Derived allowlist</summary>")
+        lines.append("")
+        for backend in backends:
+            # Megatron derives one pattern per name read off `args`, thousands of
+            # them, so show a sample rather than flooding the summary.
+            patterns = build_allowlist(ROOT, backend, args.allowlist).patterns()
+            shown = ", ".join(f"`{p}`" for p in patterns[:SHOWN_PATTERNS])
+            extra = len(patterns) - SHOWN_PATTERNS
+            lines.append(
+                f"- **{backend}** ({len(patterns)}): {shown}"
+                + (f", ... (+{extra} more)" if extra > 0 else "")
+            )
+        lines.append("")
+        lines.append("</details>")
+
+        # A scope narrows the legal key set rather than widening it, so it earns
+        # the same scrutiny as an allowlist entry: show what it covers.
+        if model_scopes:
+            lines.append("")
+            lines.append("<details><summary>Derived model scopes</summary>")
+            lines.append("")
+            for backend, scope in model_scopes:
+                keys = ", ".join(f"`{k}`" for k in sorted(scope.keys))
+                lines.append(
+                    f"- **{backend}** `{scope.describe()}` -> `{scope.config_class}` "
+                    f"(from `{scope.source}`): {keys}"
+                )
+            lines.append("")
+            lines.append("</details>")
+
+    lines.append("")
+    scope = f"{checked} config(s)" if not errors else f"{checked} of {checked + len(errors)} config(s)"
+    lines.append(f"_Checked {scope} in {time.perf_counter() - started:.2f}s._")
+    print("\n".join(lines))
+    # A missing submodule is an infrastructure failure, not a config one, so it
+    # gets none of the grace period `--warn-only` buys the findings.
+    if skipped:
+        return 1
+    return 1 if (rows or scoped or errors or stale) and not args.warn_only else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/tools/ci/config_schema_allowlist.yaml b/tools/ci/config_schema_allowlist.yaml
new file mode 100644
index 000000000..7f1c5389f
--- /dev/null
+++ b/tools/ci/config_schema_allowlist.yaml
@@ -0,0 +1,79 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Hand-written escape hatch for `tools/ci/check_config_schema.py`.
+#
+# Keep this file small. Primus extension fields are derived automatically from
+# `get_param(ctx, "<path>", ...)` calls in the backend patch packages, from the
+# names the Megatron backend package reads off `args`, and from the Primus
+# dataclasses that extend a backend's config class, so a new patch consuming a
+# new field needs no edit here. Only keys that no automatic source can see
+# belong below.
+#
+# Every entry requires:
+#   key      - exact dotted path, or an fnmatch pattern (e.g. `model.*`)
+#   reason   - why the key is legal even though upstream does not define it
+#   consumer - the file that actually reads it (a reviewer must be able to
+#              confirm the key is not dead). Give a path with no line number:
+#              the check verifies the file exists and still mentions the key,
+#              which survives refactoring, whereas a line number does not.
+
+common:
+  - key: work_group
+    reason: Primus experiment envelope, never forwarded to the backend.
+    consumer: primus/core/launcher/parser.py
+  - key: user_name
+    reason: Primus experiment envelope, never forwarded to the backend.
+    consumer: primus/core/launcher/parser.py
+  - key: exp_name
+    reason: Primus experiment envelope, never forwarded to the backend.
+    consumer: primus/core/launcher/parser.py
+  - key: workspace
+    reason: Primus experiment envelope, never forwarded to the backend.
+    consumer: primus/core/launcher/parser.py
+  - key: platform
+    reason: Selects the platform preset and its overrides.
+    consumer: primus/core/launcher/parser.py
+  - key: extends
+    reason: Preset composition directive, consumed and stripped while loading.
+    consumer: primus/core/config/yaml_loader.py
+  - key: trainable
+    reason: Primus module flag declared by primus/configs/modules/module_base.yaml.
+    consumer: primus/core/base_module.py
+  - key: sink_level
+    reason: Primus logging level, declared by primus/configs/modules/module_base.yaml.
+    consumer: primus/core/base_module.py
+  - key: file_sink_level
+    reason: Primus logging level, declared by primus/configs/modules/module_base.yaml.
+    consumer: primus/core/base_module.py
+  - key: stderr_sink_level
+    reason: Primus logging level, declared by primus/configs/modules/module_base.yaml.
+    consumer: primus/core/base_module.py
+
+torchtitan:
+  - key: model.*
+    reason: >-
+      Everything under `model` that is not a JobConfig.Model field is a model-arg
+      override, validated at runtime against the flavor dataclass (an unknown
+      name raises AttributeError there, so it cannot fail silently).
+    consumer: primus/backends/torchtitan/patches/model_override_patches.py
+
+maxtext:
+  - key: base_config
+    reason: pyconfig directive naming the MaxText base config to inherit from; not a field of base.yml itself.
+    consumer: third_party/maxtext/src/maxtext/configs/pyconfig.py
+
+megatron:
+  - key: configs
+    reason: >-
+      Top-level section of a TE-precision quantization recipe, which Megatron
+      loads from a file named by `te_precision_config_file` rather than from a
+      module preset. Such a recipe is not a Primus config and has its own schema.
+    consumer: third_party/Megatron-LM/megatron/core/quantization/quant_config.py
+  - key: matchers
+    reason: >-
+      Top-level section of a TE-precision quantization recipe (see `configs`).
+    consumer: third_party/Megatron-LM/megatron/core/quantization/quant_config.py
diff --git a/tools/ci/coverage_summary.py b/tools/ci/coverage_summary.py
index a7477ace8..400136a74 100644
--- a/tools/ci/coverage_summary.py
+++ b/tools/ci/coverage_summary.py
@@ -9,51 +9,54 @@
 Modes (by number of report arguments):
   1 report   -> single "Coverage" column (e.g. JAX MaxText E2E).
   2+ reports -> "Unit" vs "Unit+E2E". The 1st is unit; the rest are E2E reports,
-                merged per module by taking the max covered lines. Each E2E
-                report should be a `coverage combine` of unit + that job's E2E
-                data (line-level). Taking the max across jobs avoids double
-                counting and lets torch (megatron/torchtitan) and jax (maxtext)
-                E2E - which cover near-disjoint modules - share one table.
-
-Layout: top-level groups (core, backends, modules, agents, cli, ...) are bold
-rows at the same level, sorted by coverage. core/ and backends/ also get
-indented sub-rows per area, sorted by coverage. __init__.py is dropped;
-tools/, platforms/ and the top-level pretrain.py entrypoint are excluded (ops
-tooling / env abstraction / thin CLI glue, exercised by E2E and shell tests
-rather than unit tests). runner/ is bash, covered by the tests/runner/ shell
-tests.
-
-Single-report mode reflects what a partial run (e.g. MaxText E2E) actually
-executed: modules with zero covered lines are hidden and the total is computed
-over the executed modules only, so the headline number is meaningful instead of
-diluted by code that run can never touch. The two-report comparison keeps every
-module (unit gives the full denominator).
+                merged per module by max covered lines (avoids double-counting
+                torch's near-disjoint megatron/torchtitan E2E).
+
+Layout: top-level groups are bold rows sorted by coverage; core/backends also
+get indented per-area detail rows. The headline percentage gets a tier emoji
+(_TIER_THRESHOLDS) and a curated "Notes" column flags widely-shared infra or a
+known low-coverage cause (NOTES) -- both are reading aids, not quality gates.
+
+What counts toward a total vs. what gets its own row are deliberately
+decoupled decisions -- see classify().
+
+Single-report mode hides untouched modules and totals only the executed ones,
+so a partial run (e.g. MaxText E2E) isn't diluted by code it can't reach; the
+two-report comparison keeps the full denominator.
 """
 
 import json
 import sys
 from collections import defaultdict
 
-OMIT_MODULES = {"tools", "platforms", "pretrain.py"}
+OMIT_MODULES = {"tools", "platforms"}
 # Top-level groups whose sub-packages are shown as indented detail rows; every
-# other group (modules, agents, cli, ...) is a single bold row.
+# other group (agents, cli, ...) is a single bold row.
 DETAILED_GROUPS = ("core", "backends")
 
 
 def classify(path: str):
-    """Return (group, detail) for a covered file, or None to skip it.
-
-    group is the top-level row key; detail is the sub-row key for
-    DETAILED_GROUPS (e.g. core/projection), else None.
+    """Return (group, detail) for a file, or None to drop it (not counted).
+
+    group is the top-level row key; detail is the sub-row key within
+    DETAILED_GROUPS (e.g. "core/projection"). detail=None means "counted
+    toward group's total, no row of its own" -- true for every
+    non-DETAILED_GROUPS group, and for a bare file with no sub-package of its
+    own (pretrain.py -> folds into "primus (top-level)"; core/base_module.py ->
+    folds into "core"). Folding is by path *depth*, not filename, so a future
+    file added the same way folds the same way for free.
+
+    Dropping entirely (OMIT_MODULES above) is the only *policy* exclusion --
+    orthogonal to the structural folding here.
     """
     seg = (path[path.find("primus/") :] if "primus/" in path else path).split("/")
     if seg[-1] == "__init__.py":
         return None
-    if len(seg) < 2:
-        return "(top-level)", None
-    if len(seg) == 2:  # primus/<file>.py, e.g. pretrain.py
-        return seg[1], None
+    if len(seg) <= 2:  # outside primus/ entirely, or primus/<file>.py directly (no sub-package)
+        return "primus (top-level)", None
     if seg[1] in DETAILED_GROUPS:
+        if len(seg) == 3:  # primus/<group>/<file>.py directly: no sub-package of its own
+            return seg[1], None
         return seg[1], seg[1] + "/" + seg[2]
     return seg[1], None
 
@@ -62,6 +65,52 @@ def _pct(covered: int, total: int) -> float:
     return (100.0 * covered / total) if total else 0.0
 
 
+# Best-effort visual highlight for the headline column: GitHub strips CSS color
+# from Action run summaries, so emoji is the portable substitute. Thresholds
+# are a rough reading aid, not a quality gate.
+_TIER_THRESHOLDS = ((50.0, "\U0001F7E2"), (25.0, "\U0001F7E1"))  # >=50% green, >=25% yellow, else red
+_TIER_RED = "\U0001F534"
+
+
+def _tier(pct: float) -> str:
+    for threshold, emoji in _TIER_THRESHOLDS:
+        if pct >= threshold:
+            return emoji
+    return _TIER_RED
+
+
+# Coverage priority: lower sorts first (ties keep coverage order), pinning the
+# most important groups to the top. core (Primus's core) ranks highest.
+PRIORITY = {"core": 0}
+_DEFAULT_PRIORITY = 1
+
+
+def _priority(group: str) -> int:
+    return PRIORITY.get(group, _DEFAULT_PRIORITY)
+
+
+# Curated, best-effort context -- not exhaustive. Flags infra whose coverage
+# matters more than its size suggests, and low-coverage areas with a known,
+# persistent cause, so it isn't re-litigated every read. Keyed like the table
+# (group, or "<group>/<detail>"); update alongside any fix or new finding.
+# <br> forces cell wrapping so one long note can't stretch the whole column.
+NOTES = {
+    "primus (top-level)": "loose primus/ modules, no sub-package;<br>auto-folded by path depth",
+    "core": "\U0001F511 Primus core;<br>imported by every run",
+    "backends/megatron": (
+        "100+ patches gated by fp8 / MoE /<br>zero-bubble-pp / fsdp2 flags;<br>"
+        "CI E2E runs only 1-2 configs"
+    ),
+    "backends/transformer_engine": "fp8 GEMM / attn-overlap kernels;<br>only hit when an E2E enables fp8",
+    "backends/diffusion": "no E2E trainer suite yet<br>(unit-tested only)",
+}
+
+_LEGEND = (
+    "\U0001F7E2 >=50% / \U0001F7E1 >=25% / \U0001F534 <25% (next to module) "
+    "&nbsp;\u00b7&nbsp; \U0001F511 Primus core\n"
+)
+
+
 def _aggregate(report: dict):
     """Return {group: {detail|group: [covered, statements]}} for kept modules."""
     agg = defaultdict(lambda: defaultdict(lambda: [0, 0]))
@@ -89,27 +138,34 @@ def e2e_cov(group, key):
         # modules, so max avoids double counting the shared core code).
         return max((sa.get(group, {}).get(key, [0])[0] for sa in sas), default=0) if two else 0
 
-    def row(label, cov, stmts, e2e, bold=False):
+    def row(label, key, cov, stmts, e2e, bold=False):
+        # Tier dot rides next to the module name (leftmost), so it reads as a
+        # per-module health mark and leaves the right-aligned % columns clean.
+        # It goes *after* any &emsp; indent so detail-row dots stay indented.
+        headline = _pct(e2e if two else cov, stmts)
+        dot = _tier(headline) + " " if stmts else ""
         if two:
             vals = [format(stmts, ","), "%.1f%%" % _pct(cov, stmts), "%.1f%%" % _pct(e2e, stmts)]
         else:
             vals = [format(cov, ","), format(stmts, ","), "%.1f%%" % _pct(cov, stmts)]
         w = "**" if bold else ""
-        return "| " + " | ".join("%s%s%s" % (w, x, w) for x in [label] + vals) + " |"
+        indent = "&emsp;" if label.startswith("&emsp;") else ""
+        rest = label[len(indent) :]
+        label_cell = "%s%s%s%s%s" % (indent, dot, w, rest, w)
+        cells = [label_cell] + ["%s%s%s" % (w, x, w) for x in vals]
+        cells.append(NOTES.get(key, ""))
+        return "| " + " | ".join(cells) + " |"
 
     def group_totals(group):
-        # Single-report mode counts only executed entries (cov > 0) so a partial
-        # run isn't diluted by sub-modules it never touched; two-report keeps all.
+        # Two-report mode keeps every entry (full denominator); single-report
+        # mode counts only executed ones so a partial run isn't diluted.
         entries = [(k, v) for k, v in pa[group].items() if two or v[0] > 0]
         cov = sum(v[0] for _, v in entries)
         stmts = sum(v[1] for _, v in entries)
         e2e = sum(e2e_cov(group, k) for k, _ in entries) if two else 0
         return cov, stmts, e2e
 
-    # Single-report mode hides modules with zero coverage and totals over the
-    # executed modules only, so a partial run (e.g. MaxText E2E) isn't diluted by
-    # code it can never touch. The two-report comparison keeps the full denominator.
-    def group_executed(group):
+    def group_executed(group):  # single-report mode: hide groups a partial run never touched
         return group_totals(group)[0] > 0 if not two else True
 
     groups = [g for g in pa if group_totals(g)[1] > 0 and group_executed(g)]
@@ -136,25 +192,27 @@ def group_executed(group):
         out.append(
             "_Including all modules (nothing excluded): Unit %.1f%% -> Unit+E2E %.1f%%._\n" % (p_all, s_all)
         )
-        out += ["| Module | Stmts | Unit | Unit+E2E |", "|---|--:|--:|--:|"]
+        out.append(_LEGEND)
+        out += ["| Module | Stmts | Unit | Unit+E2E | Notes |", "|---|--:|--:|--:|---|"]
     else:
         out.append(
             "**Total line coverage: %.1f%%** (%s / %s statements; excludes %s)\n"
             % (_pct(tc, tn), format(tc, ","), format(tn, ","), excl)
         )
-        out += ["| Module | Covered | Stmts | Coverage |", "|---|--:|--:|--:|"]
+        out.append(_LEGEND)
+        out += ["| Module | Covered | Stmts | Coverage | Notes |", "|---|--:|--:|--:|---|"]
 
-    # Top-level groups, sorted by coverage (desc).
-    for group in sorted(groups, key=lambda g: -_pct(group_totals(g)[0], group_totals(g)[1])):
+    for group in sorted(groups, key=lambda g: (_priority(g), -_pct(group_totals(g)[0], group_totals(g)[1]))):
         cov, stmts, e2e = group_totals(group)
-        out.append(row("`%s`" % group, cov, stmts, e2e, bold=True))
+        out.append(row("`%s`" % group, group, cov, stmts, e2e, bold=True))
         if group in DETAILED_GROUPS:
-            # In single-report mode, hide sub-rows that were never executed.
-            details = ((k, v) for k, v in pa[group].items() if v[1] > 0 and (two or v[0] > 0))
+            # k == group is a folded loose file (see classify()), already
+            # counted in group_totals() above -- no row of its own here.
+            details = ((k, v) for k, v in pa[group].items() if k != group and v[1] > 0 and (two or v[0] > 0))
             for k, v in sorted(details, key=lambda kv: -_pct(kv[1][0], kv[1][1])):
-                out.append(row("&emsp;`%s`" % k, v[0], v[1], e2e_cov(group, k)))
+                out.append(row("&emsp;`%s`" % k, k, v[0], v[1], e2e_cov(group, k)))
 
-    out.append(row("TOTAL", tc, tn, te, bold=True))
+    out.append(row("TOTAL", None, tc, tn, te, bold=True))
     return "\n".join(out)
 
 
diff --git a/tools/ci/junit_summary.py b/tools/ci/junit_summary.py
index b84936674..121fdda67 100644
--- a/tools/ci/junit_summary.py
+++ b/tools/ci/junit_summary.py
@@ -25,6 +25,17 @@
 COLUMNS = ("tests", "passed", "failed", "errors", "skipped")
 
 
+def _fmt_time(seconds):
+    seconds = int(seconds)
+    hours, rem = divmod(seconds, 3600)
+    minutes, secs = divmod(rem, 60)
+    if hours:
+        return f"{hours}h{minutes:02d}m{secs:02d}s"
+    if minutes:
+        return f"{minutes}m{secs:02d}s"
+    return f"{secs}s"
+
+
 def parse_file(path):
     """Return (label, stats|None, failures) for one JUnit XML file."""
     label = os.path.splitext(os.path.basename(path))[0]
@@ -69,15 +80,30 @@ def render(reports, title=None):
             lines.append("| `%s` | _no report_ |  |  |  |  |  |" % label)
             continue
         lines.append(
-            "| `%s` | %d | %d | %d | %d | %d | %.1fs |"
-            % (label, st["tests"], st["passed"], st["failed"], st["errors"], st["skipped"], st["time"])
+            "| `%s` | %d | %d | %d | %d | %d | %s |"
+            % (
+                label,
+                st["tests"],
+                st["passed"],
+                st["failed"],
+                st["errors"],
+                st["skipped"],
+                _fmt_time(st["time"]),
+            )
         )
         for k in total:
             total[k] += st[k]
         failures += [(label, *f) for f in fails]
     lines.append(
-        "| **TOTAL** | **%d** | **%d** | **%d** | **%d** | **%d** | **%.1fs** |"
-        % (total["tests"], total["passed"], total["failed"], total["errors"], total["skipped"], total["time"])
+        "| **TOTAL** | **%d** | **%d** | **%d** | **%d** | **%d** | **%s** |"
+        % (
+            total["tests"],
+            total["passed"],
+            total["failed"],
+            total["errors"],
+            total["skipped"],
+            _fmt_time(total["time"]),
+        )
     )
     lines.append("")
 
diff --git a/tools/ci/runtime_summary.py b/tools/ci/runtime_summary.py
index dd55ff08c..4d8635f97 100644
--- a/tools/ci/runtime_summary.py
+++ b/tools/ci/runtime_summary.py
@@ -4,14 +4,32 @@
 # See LICENSE for license information.
 ###############################################################################
 
-"""Render CI stage wall-clock times (a stage<TAB>seconds TSV) as a Markdown table.
+"""Render the *current* job's own step wall-clock times (fetched from the
+Actions API) as a Markdown table for the CI run summary.
 
-Complements junit_summary (test time) by surfacing the heavy build/install
-stages. A missing/empty TSV renders nothing and exits 0 so the step never fails
-the job; stage order (i.e. execution order) is preserved.
+Auto-discovers every step from the job's live metadata, so adding, removing,
+or renaming a step in ci.yaml needs no matching edit here and no hand-rolled
+per-step timing block -- unlike manually timing "just this stage" into a
+file, which silently misses whatever nobody remembered to wrap.
+
+Requires the calling step to export GITHUB_TOKEN with `actions: read` (see
+ci.yaml); a missing token, network failure, or running outside Actions all
+degrade to "print nothing, exit 0" so this can never fail the job:
+
+    GITHUB_TOKEN=... python tools/ci/runtime_summary.py --title torch >> "$GITHUB_STEP_SUMMARY"
+
+Steps shorter than --min-seconds (default 5s, almost always a banner `echo`
+step) are hidden so the table stays focused.
 """
 
 import argparse
+import json
+import os
+import urllib.error
+import urllib.request
+from datetime import datetime, timezone
+
+API_TIMEOUT_S = 30
 
 
 def fmt(seconds):
@@ -25,44 +43,84 @@ def fmt(seconds):
     return f"{secs}s"
 
 
-def parse(path):
+def _parse_ts(s):
+    if not s:
+        return None
+    return datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
+
+
+def fetch_steps(repo, run_id, job_name, runner_name, token):
+    """Return the step list for the run's job named `job_name`, preferring the
+    one running on `runner_name` (in case the name is ever repeated, e.g. a
+    future matrix), or [] if no such job is found."""
+    url = f"https://api.github.com/repos/{repo}/actions/runs/{run_id}/jobs?per_page=100"
+    req = urllib.request.Request(
+        url,
+        headers={
+            "Authorization": f"Bearer {token}",
+            "Accept": "application/vnd.github+json",
+            "X-GitHub-Api-Version": "2022-11-28",
+        },
+    )
+    with urllib.request.urlopen(req, timeout=API_TIMEOUT_S) as resp:  # noqa: S310 (fixed https:// API host)
+        jobs = json.load(resp).get("jobs", [])
+    matches = [j for j in jobs if j.get("name") == job_name]
+    exact = [j for j in matches if j.get("runner_name") == runner_name]
+    job = (exact or matches or [None])[0]
+    return job.get("steps", []) if job else []
+
+
+def step_durations(steps, min_seconds=5.0):
+    """[(name, seconds), ...] in run order, for steps that actually completed
+    (excludes the still-running caller itself, any not-yet-run step, and
+    anything skipped) and cleared the min_seconds noise floor."""
     rows = []
-    try:
-        with open(path) as handle:
-            for line in handle:
-                parts = line.rstrip("\n").split("\t")
-                if len(parts) != 2 or not parts[0].strip():
-                    continue
-                try:
-                    rows.append((parts[0].strip(), float(parts[1].strip())))
-                except ValueError:
-                    continue
-    except OSError:
-        return []
+    for step in steps:
+        if step.get("conclusion") == "skipped":
+            continue
+        start, end = _parse_ts(step.get("started_at")), _parse_ts(step.get("completed_at"))
+        if not start or not end:
+            continue  # still running (incl. this very step) or never reached
+        secs = (end - start).total_seconds()
+        if secs >= min_seconds:
+            rows.append((step.get("name", "?"), secs))
     return rows
 
 
-def render(rows, title=None):
+def render(steps, title=None, min_seconds=5.0):
+    rows = step_durations(steps, min_seconds)
+    if not rows:
+        return ""
     suffix = f" - {title}" if title else ""
     lines = [f"## CI runtime{suffix}\n", "| Stage | Time |", "|---|--:|"]
-    total = 0.0
-    for stage, secs in rows:
-        lines.append(f"| {stage} | {fmt(secs)} |")
-        total += secs
-    lines.append(f"| **TOTAL (timed stages)** | **{fmt(total)}** |")
+    lines += [f"| {name} | {fmt(secs)} |" for name, secs in rows]
+    lines.append(f"| **TOTAL (shown stages)** | **{fmt(sum(s for _, s in rows))}** |")
     return "\n".join(lines)
 
 
 def main():
-    ap = argparse.ArgumentParser(description="Render a stage<TAB>seconds TSV as a Markdown runtime table.")
-    ap.add_argument("tsv", help="Path to the runtime TSV (stage<TAB>seconds per line).")
+    ap = argparse.ArgumentParser(description="Render this job's own step timings as a Markdown table.")
     ap.add_argument("--title", default=None, help="Optional section title (e.g. torch).")
+    ap.add_argument(
+        "--min-seconds", type=float, default=5.0, help="Hide steps shorter than this (default 5s)."
+    )
     args = ap.parse_args()
 
-    rows = parse(args.tsv)
-    if not rows:
-        return 0  # nothing timed; don't emit an empty table or fail the step
-    print(render(rows, args.title))
+    token = os.environ.get("GITHUB_TOKEN")
+    repo, run_id, job_name = (
+        os.environ.get("GITHUB_REPOSITORY"),
+        os.environ.get("GITHUB_RUN_ID"),
+        os.environ.get("GITHUB_JOB"),
+    )
+    if not (token and repo and run_id and job_name):
+        return 0  # not running in Actions (or token not wired) -- nothing to render
+    try:
+        steps = fetch_steps(repo, run_id, job_name, os.environ.get("RUNNER_NAME"), token)
+    except (urllib.error.URLError, TimeoutError, ValueError, OSError):
+        return 0  # a reporting step must never fail the job
+    out = render(steps, args.title, args.min_seconds)
+    if out:
+        print(out)
     return 0
 
 
diff --git a/tools/ci/select_tests.py b/tools/ci/select_tests.py
index cfff17855..cdff9d014 100644
--- a/tools/ci/select_tests.py
+++ b/tools/ci/select_tests.py
@@ -4,19 +4,20 @@
 # See LICENSE for license information.
 ###############################################################################
 
-"""Map a PR's changed files (stdin, one per line) to the tests to run.
-
-Default prints the minimal unit-test paths; --e2e prints the E2E trainer suites
-(or "all"). A single classify() decides each path's blast radius and both
-selections build on it. Conventions over hard-coded tables:
-  - unit dirs mirror the source tree (primus/<x> -> tests/unit_tests/<x>),
-    resolved by walking up to the nearest existing dir;
+"""Map a PR's changed files (stdin, one per line) to the E2E trainer suites to
+run (or "all"). A single classify() decides each path's blast radius.
+Conventions over hard-coded tables:
   - E2E suites are auto-discovered from tests/trainer/test_<name>_trainer.py;
   - a backend is named by its dir (primus/backends/<X> or examples/<X>).
-Fail-safe is the only invariant: anything global, unlocatable, or a backend
-without a trainer expands to everything -- over-select, never under-select.
+Fail-safe is the only invariant: anything global, a non-backend source change,
+or a backend without a trainer expands to everything -- over-select, never
+under-select.
+
+    git diff --name-only "$BASE" HEAD | python tools/ci/select_tests.py
 
-    git diff --name-only "$BASE" HEAD | python tools/ci/select_tests.py [--e2e]
+Unit tests are deliberately NOT selected/narrowed here: the whole suite only
+takes ~5 minutes (vs. each E2E suite's tens of minutes), so narrowing risked
+under-selection for little wall-clock gain. ci.yaml always runs it in full.
 """
 
 import re
@@ -24,28 +25,24 @@
 from pathlib import Path
 
 ROOT = Path(__file__).resolve().parents[2]
-FULL = "tests/unit_tests/"
 
-# The only hard-coded list: changes whose blast radius is the whole repo. Being
-# absent here only ever falls back to other fail-safe paths, never to "skip".
+# The only hard-coded list: changes whose blast radius is the whole repo.
+# Scoped to paths *outside* primus/, tests/unit_tests/, and examples/ --
+# anything under those trees that isn't a recognized backend already falls
+# back to "component" below, which select_e2e() treats the same as "global"
+# anyway, so listing e.g. primus/core/launcher/ here would be redundant.
+# Being absent here only ever falls back to other fail-safe paths, never "skip".
 GLOBAL_TRIGGERS = (
     ".github/",
     "tools/",
     "runner/",  # the launcher drives all training
     "pyproject.toml",
-    "primus/__init__.py",
-    "primus/core/launcher/",
-    "primus/core/utils/",
-    "primus/core/config/",
     "tests/utils.py",
     "tests/conftest.py",
     "tests/unit_tests/conftest.py",
     "tests/run_unit_tests.py",
 )
 
-# megatron's GPU-operator tests aren't path-isomorphic to the backend source.
-_BACKEND_EXTRA_UNIT = {"megatron": ("tests/unit_tests/megatron/",)}
-
 _TRAINER_RE = re.compile(r"test_(.+)_trainer")
 _BACKEND_RE = re.compile(r"(?:primus/backends|examples)/([^/]+)/")
 
@@ -64,20 +61,8 @@ def _is_global(path):
     return any(path == t or path.startswith(t) for t in GLOBAL_TRIGGERS)
 
 
-def _nearest_unit_dir(rel):
-    # rel is source-relative (under primus/ or tests/unit_tests/); walk up to the
-    # nearest existing tests/unit_tests/<...> dir, or None if none exists.
-    parts = rel.split("/")[:-1]
-    while parts:
-        cand = "tests/unit_tests/" + "/".join(parts) + "/"
-        if (ROOT / cand).is_dir():
-            return cand
-        parts.pop()
-    return None
-
-
 def classify(path):
-    """('global', None) | ('backend', name) | ('component', unit_dir|None) | ('ignore', None)."""
+    """('global', None) | ('backend', name) | ('component', None) | ('ignore', None)."""
     if _is_global(path):
         return ("global", None)
     backend = _BACKEND_RE.match(path)  # primus/backends/<X>/ or examples/<X>/
@@ -86,43 +71,14 @@ def classify(path):
     if path.startswith("tests/trainer/"):
         m = _TRAINER_RE.search(path)
         return ("backend", m.group(1)) if m else ("ignore", None)
-    for root in ("primus/", "tests/unit_tests/"):
-        if path.startswith(root):
-            if not path.endswith(".py"):
-                return ("global", None)  # non-.py here (configs, fixtures) -> fail-safe
-            return ("component", _nearest_unit_dir(path[len(root) :]))
+    # A bare examples/<file> (no backend subdir, so it didn't match _BACKEND_RE
+    # above) is shared launcher plumbing, not per-backend -- e.g. maxtext's E2E
+    # shells out to examples/run_pretrain.sh directly.
+    if path.startswith("primus/") or path.startswith("tests/unit_tests/") or path.startswith("examples/"):
+        return ("component", None)  # any other source/unit-test/launcher change
     return ("ignore", None)  # docs, README, ... outside the source/test trees
 
 
-def select_targets(files):
-    files = [f.strip() for f in files if f.strip()]
-    if not files:
-        return [FULL]
-    targets = []
-
-    def add(d):
-        if d and d not in targets:
-            targets.append(d)
-
-    for path in files:
-        kind, val = classify(path)
-        if kind == "global":
-            return [FULL]
-        if kind == "backend":
-            base = f"tests/unit_tests/backends/{val}/"
-            if not (ROOT / base).is_dir():
-                return [FULL]  # backend without a unit dir (e.g. transformer_engine) -> safe
-            add(base)
-            for extra in _BACKEND_EXTRA_UNIT.get(val, ()):
-                add(extra)
-        elif kind == "component":
-            if val is None:
-                return [FULL]  # couldn't localize a unit dir -> safe
-            add(val)
-        # ignore -> skip
-    return targets or [FULL]
-
-
 def select_e2e(files, suites=None):
     suites = discover_e2e_suites() if suites is None else set(suites)
     files = [f.strip() for f in files if f.strip()]
@@ -146,12 +102,9 @@ def select_e2e(files, suites=None):
 
 def main():
     files = sys.stdin.read().splitlines()
-    if "--e2e" in sys.argv[1:]:
-        suites = discover_e2e_suites()
-        e2e = select_e2e(files, suites)
-        print("all" if suites and set(e2e) == suites else " ".join(e2e))
-    else:
-        print(" ".join(select_targets(files)))
+    suites = discover_e2e_suites()
+    e2e = select_e2e(files, suites)
+    print("all" if suites and set(e2e) == suites else " ".join(e2e))
     return 0
 
 
diff --git a/tools/daily/daily_report.py b/tools/daily/daily_report.py
index f0b6b4fd7..3ccce498a 100644
--- a/tools/daily/daily_report.py
+++ b/tools/daily/daily_report.py
@@ -17,8 +17,15 @@
 MEGATRON_ITERATION_PATTERN = re.compile(
     r"iteration\s+\d+\s*/\s*\d+.*?"
     r"elapsed time per iteration\s*\(ms\):\s*(?P<elapsed_ms>[\d.]+)(?:/(?P<elapsed_ms_avg>[\d.]+))?.*?"
-    r"throughput per GPU\s*\(TFLOP/s/GPU\):\s*(?P<tflops>[\d.]+)/(?P<tflops_avg>[\d.]+).*?"
-    r"tokens per GPU\s*\(tokens/s/GPU\):\s*(?P<tokens>[\d.]+)/(?P<tokens_avg>[\d.]+).*?"
+    # Compute (TFLOP/s/GPU) accepts both the current "compute per GPU (...): X (avg Y)"
+    # label and the legacy "throughput per GPU (...): X/Y" label. The average is
+    # optional so raw (un-averaged) Megatron lines still parse.
+    r"(?:compute|throughput) per GPU\s*\(TFLOP/s/GPU\):\s*(?P<tflops>[\d.]+)"
+    r"(?:\s*(?:/|\(avg\s*)\s*(?P<tflops_avg>[\d.]+)\)?)?.*?"
+    # Tokens accepts both the current "tokens/s/GPU inst/harmonic mean: X/Y" label
+    # and the legacy "tokens per GPU (tokens/s/GPU): X/Y" label.
+    r"(?:tokens per GPU\s*\(tokens/s/GPU\)|tokens/s/GPU\s+inst/harmonic mean):\s*"
+    r"(?P<tokens>[\d.]+)(?:\s*/\s*(?P<tokens_avg>[\d.]+))?.*?"
     r"hip mem usage/free/total/usage_ratio:\s*"
     r"(?P<hip_used>[\d.]+)\s*G(?:i)?B/(?P<hip_free>[\d.]+)\s*G(?:i)?B/(?P<hip_total>[\d.]+)\s*G(?:i)?B/(?P<hip_ratio>[\d.]+)%",
     re.DOTALL,
@@ -138,10 +145,12 @@ def parse_last_metrics_from_log(file_path: str) -> Dict[str, float]:
     if not last_m:
         raise ValueError(f"No valid iteration metrics found in {file_path}")
     metric = last_m.groupdict()
-    step_time_s_avg = float(metric["elapsed_ms_avg"]) / 1000.0
+    # Averages are optional in the pattern (raw Megatron lines omit them); fall
+    # back to the instantaneous value when an average is not present.
+    step_time_s_avg = float(metric["elapsed_ms_avg"] or metric["elapsed_ms"]) / 1000.0
     mem_usage = float(metric["hip_used"])
-    tflops_avg = float(metric["tflops_avg"])
-    tokens_avg = float(metric["tokens_avg"])
+    tflops_avg = float(metric["tflops_avg"] or metric["tflops"])
+    tokens_avg = float(metric["tokens_avg"] or metric["tokens"])
 
     return {
         "TFLOP/s/GPU": round(tflops_avg, 2),
diff --git a/tools/docker/start_container.sh b/tools/docker/start_container.sh
index d29289122..2fe081557 100755
--- a/tools/docker/start_container.sh
+++ b/tools/docker/start_container.sh
@@ -6,7 +6,7 @@
 ###############################################################################
 
 PRIMUS_PATH=$(realpath "$(dirname "$0")/../..")
-DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.3"}
+DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/rocm/primus:v26.4"}
 DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"}
 SANITIZED_USER=$(echo "${USER:-unknown}" | tr -cd '[:alnum:]_-')
 if [ -z "$SANITIZED_USER" ]; then
diff --git a/tools/generate_empty_encodings.py b/tools/generate_empty_encodings.py
new file mode 100644
index 000000000..10e55781d
--- /dev/null
+++ b/tools/generate_empty_encodings.py
@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Generate empty T5 and CLIP encodings for CFG dropout.
+
+Runs T5-XXL and CLIP-L on the empty string ("") and saves the resulting
+embeddings as .npy files. These are loaded by FluxPretrainTrainer at init
+to replace text embeddings during classifier-free guidance dropout.
+
+Using real model outputs (instead of torch.randn) is critical for training
+convergence — see NVIDIA MLPerf Training v5.1 reference.
+
+Usage:
+    python tools/generate_empty_encodings.py \
+        --output_dir /path/to/empty_encodings \
+        --t5_model google/t5-v1_1-xxl \
+        --clip_model openai/clip-vit-large-patch14 \
+        --t5_max_length 256
+
+Output files:
+    t5_empty.npy   - shape (1, seq_len, 4096)
+    clip_empty.npy - shape (1, 768)
+"""
+
+import argparse
+import os
+
+import numpy as np
+import torch
+
+
+def generate_t5_empty(model_name: str, max_length: int, device: str) -> np.ndarray:
+    """Generate T5-XXL encoding for empty string."""
+    from transformers import T5EncoderModel, T5Tokenizer
+
+    print(f"Loading T5 tokenizer: {model_name}")
+    tokenizer = T5Tokenizer.from_pretrained(model_name)
+
+    print(f"Loading T5 model: {model_name}")
+    model = T5EncoderModel.from_pretrained(model_name, torch_dtype=torch.float32)
+    model = model.to(device).eval()
+
+    with torch.no_grad():
+        inputs = tokenizer(
+            "",
+            max_length=max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        ).to(device)
+
+        outputs = model(**inputs)
+        # outputs.last_hidden_state: (1, seq_len, 4096)
+        embeddings = outputs.last_hidden_state.cpu().numpy()
+
+    print(f"T5 empty encoding shape: {embeddings.shape}")
+    del model
+    torch.cuda.empty_cache()
+    return embeddings
+
+
+def generate_clip_empty(model_name: str, device: str) -> np.ndarray:
+    """Generate CLIP-L pooled encoding for empty string."""
+    from transformers import CLIPTextModel, CLIPTokenizer
+
+    print(f"Loading CLIP tokenizer: {model_name}")
+    tokenizer = CLIPTokenizer.from_pretrained(model_name)
+
+    print(f"Loading CLIP model: {model_name}")
+    model = CLIPTextModel.from_pretrained(model_name, torch_dtype=torch.float32)
+    model = model.to(device).eval()
+
+    with torch.no_grad():
+        inputs = tokenizer(
+            "",
+            max_length=tokenizer.model_max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        ).to(device)
+
+        outputs = model(**inputs)
+        # outputs.pooler_output: (1, 768)
+        pooled = outputs.pooler_output.cpu().numpy()
+
+    print(f"CLIP empty encoding shape: {pooled.shape}")
+    del model
+    torch.cuda.empty_cache()
+    return pooled
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate empty T5/CLIP encodings for CFG dropout")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Directory to save t5_empty.npy and clip_empty.npy",
+    )
+    parser.add_argument(
+        "--t5_model",
+        type=str,
+        default="google/t5-v1_1-xxl",
+        help="T5 model name or path (default: google/t5-v1_1-xxl)",
+    )
+    parser.add_argument(
+        "--clip_model",
+        type=str,
+        default="openai/clip-vit-large-patch14",
+        help="CLIP model name or path (default: openai/clip-vit-large-patch14)",
+    )
+    parser.add_argument(
+        "--t5_max_length",
+        type=int,
+        default=256,
+        help="Max sequence length for T5 encoding (default: 256 for schnell)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device to run models on (default: cuda if available)",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    t5_path = os.path.join(args.output_dir, "t5_empty.npy")
+    clip_path = os.path.join(args.output_dir, "clip_empty.npy")
+
+    t5_embeddings = generate_t5_empty(args.t5_model, args.t5_max_length, args.device)
+    np.save(t5_path, t5_embeddings)
+    print(f"Saved T5 empty encodings to: {t5_path}")
+
+    clip_embeddings = generate_clip_empty(args.clip_model, args.device)
+    np.save(clip_path, clip_embeddings)
+    print(f"Saved CLIP empty encodings to: {clip_path}")
+
+    print("\nDone! Add to your YAML config:")
+    print(f"  empty_encodings_path: {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/chat_zebra_llama.py b/tools/hybrid/chat_zebra_llama.py
new file mode 100644
index 000000000..1a45ed635
--- /dev/null
+++ b/tools/hybrid/chat_zebra_llama.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""
+Simple interactive chat for Zebra-Llama (HF-converted checkpoint).
+
+Loads:
+  - model code from: tools/hybrid/modeling_zebra_llama.py
+  - weights from:    output/zebra_llama_1B_hf_iter_0150000
+
+Notes:
+  - KV cache is disabled in the model implementation, so generation is slower.
+  - Chat formatting here is intentionally simple; adjust the prompt template as needed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import List, Tuple
+
+import torch
+from transformers import AutoTokenizer
+
+
+def build_prompt(history: List[Tuple[str, str]]) -> str:
+    """
+    Very simple chat prompt:
+      User: ...
+      Assistant: ...
+    """
+    parts: List[str] = []
+    for role, text in history:
+        if role == "user":
+            parts.append(f"User: {text}")
+        else:
+            parts.append(f"Assistant: {text}")
+    parts.append("Assistant:")
+    return "\n".join(parts)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Chat with Zebra-Llama (converted HF checkpoint)")
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default="output/zebra_llama_1B_hf_iter_0200000",
+        help="Path to converted HF checkpoint directory",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=str,
+        default="meta-llama/Llama-3.2-1B",
+        help="Tokenizer name/path",
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--device", type=str, default=None, help="cpu/cuda (default: auto)")
+    args = parser.parse_args()
+
+    primus_root = Path(__file__).resolve().parent.parent.parent
+    sys.path.insert(0, str(primus_root))
+
+    from tools.hybrid.modeling_zebra_llama import (  # noqa: E402
+        ZebraLlamaConfig,
+        ZebraLlamaForCausalLM,
+    )
+
+    ckpt_dir = (
+        (primus_root / args.checkpoint).resolve()
+        if not Path(args.checkpoint).is_absolute()
+        else Path(args.checkpoint)
+    )
+    if not ckpt_dir.exists():
+        raise FileNotFoundError(f"Checkpoint dir not found: {ckpt_dir}")
+
+    config_path = ckpt_dir / "config.json"
+    weights_path = ckpt_dir / "pytorch_model.bin"
+    if not config_path.exists():
+        raise FileNotFoundError(f"Missing config.json: {config_path}")
+    if not weights_path.exists():
+        raise FileNotFoundError(f"Missing pytorch_model.bin: {weights_path}")
+
+    torch.manual_seed(args.seed)
+
+    # Device / dtype
+    if args.device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    else:
+        device = args.device
+
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+    print(f"[Info] Loading tokenizer: {args.tokenizer}")
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"[Info] Loading config: {config_path}")
+    cfg_dict = json.loads(config_path.read_text())
+    config = ZebraLlamaConfig(**cfg_dict)
+
+    print(f"[Info] Building model on {device} ({dtype})")
+    model = ZebraLlamaForCausalLM(config).to(device=device, dtype=dtype)
+    model.eval()
+
+    print(f"[Info] Loading weights: {weights_path}")
+    state = torch.load(weights_path, map_location="cpu", weights_only=True)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing:
+        print(f"[Warn] Missing keys ({len(missing)}). Showing first 20:")
+        for k in missing[:20]:
+            print(f"  - {k}")
+    if unexpected:
+        print(f"[Warn] Unexpected keys ({len(unexpected)}). Showing first 20:")
+        for k in unexpected[:20]:
+            print(f"  - {k}")
+
+    history: List[Tuple[str, str]] = []
+
+    print("\n=== Zebra-Llama chat ===")
+    print("Type your message and press enter.")
+    print("Commands: /reset, /exit\n")
+
+    # while True:
+    #     try:
+    #         user = input("User> ").strip()
+    #     except (EOFError, KeyboardInterrupt):
+    #         print("\n[Exit]")
+    #         return
+
+    #     if not user:
+    #         continue
+    #     if user == "/exit":
+    #         return
+    #     if user == "/reset":
+    #         history.clear()
+    #         print("[Info] History cleared.\n")
+    # continue
+
+    # history.append(("user", user))
+    # prompt = build_prompt(history)
+
+    prompt = "Once upon a time, "
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(device)
+    attention_mask = inputs.get("attention_mask", None)
+    if attention_mask is not None:
+        attention_mask = attention_mask.to(device)
+
+    with torch.no_grad():
+        out_ids = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=args.temperature > 0,
+            temperature=1.0,
+            top_p=args.top_p,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            use_cache=False,
+        )
+
+    # Decode only the newly generated portion
+    gen_ids = out_ids[0, input_ids.shape[1] :]
+    assistant = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
+    # history.append(("assistant", assistant))
+    print(f"OUTPUT> {prompt}{assistant}\n")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/compare_hybrid_eval.py b/tools/hybrid/compare_hybrid_eval.py
new file mode 100644
index 000000000..361b74a81
--- /dev/null
+++ b/tools/hybrid/compare_hybrid_eval.py
@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""Side-by-side compare lm-eval results from two output dirs.
+
+Produces the table the team uses for parity reports:
+
+    Task             Metric    Primus  ±stderr   FLA    ±stderr   Δ abs   σ (|Δ|/√(σ1²+σ2²))
+    arc_easy         acc       ...     ±...      ...    ±...     -0.008   0.58
+    arc_easy         acc_norm  ...     ±...      ...    ±...      0.006   0.44
+    ...
+    Mean of N metrics                                             ...     ...
+    Mean |Δ|                                                      ...     ...
+    Max  |Δ|                                                       ... (which task)
+
+`mmlu` is reported as a single line averaged over its 57 sub-tasks
+(matches FLA's `mmlu (avg of 57)` row).
+"""
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+# (display_name, lm-eval task id, [metrics])  — order = display order
+DEFAULT_REPORT = [
+    ("arc_easy", "arc_easy", ["acc", "acc_norm"]),
+    ("arc_challenge", "arc_challenge", ["acc", "acc_norm"]),
+    ("hellaswag", "hellaswag", ["acc", "acc_norm"]),
+    ("openbookqa", "openbookqa", ["acc", "acc_norm"]),
+    ("piqa", "piqa", ["acc", "acc_norm"]),
+    ("winogrande", "winogrande", ["acc"]),
+    ("mmlu (avg of 57)", "mmlu", ["acc"]),
+    ("race", "race", ["acc"]),
+]
+
+
+def _find_results_json(out_dir: Path):
+    candidates = sorted(out_dir.rglob("results*.json"), key=lambda p: p.stat().st_mtime)
+    return candidates[-1] if candidates else None
+
+
+def _metric_lookup(task_results, metric):
+    """Return (value, stderr) for `metric` in lm-eval's results dict for one task."""
+    if task_results is None:
+        return None, None
+    val = task_results.get(f"{metric},none", task_results.get(metric))
+    err = task_results.get(f"{metric}_stderr,none", task_results.get(f"{metric}_stderr"))
+    if val is None:
+        return None, None
+    if err is None or (isinstance(err, str) and err.upper() == "N/A"):
+        err = 0.0
+    return float(val), float(err)
+
+
+def _mmlu_aggregate(results):
+    """lm-eval reports mmlu as 57 leaf tasks (mmlu_abstract_algebra, ...) plus
+    parent rollup keys (mmlu, mmlu_humanities, mmlu_stem, ...).
+    Prefer the parent 'mmlu' if present; otherwise average the 57 leaves
+    (acc) and propagate stderr via √(Σσ²)/N."""
+    if "mmlu" in results:
+        return results["mmlu"]
+    leaves = []
+    for k, v in results.items():
+        if not k.startswith("mmlu_"):
+            continue
+        # parent groups have no `acc,none`, skip
+        if not isinstance(v, dict):
+            continue
+        if any(kk.startswith("acc") for kk in v):
+            leaves.append(v)
+    if not leaves:
+        return None
+    accs = [_metric_lookup(v, "acc")[0] for v in leaves]
+    errs = [_metric_lookup(v, "acc")[1] for v in leaves]
+    accs = [a for a in accs if a is not None]
+    errs = [e for e in errs if e is not None]
+    if not accs:
+        return None
+    return {
+        "acc,none": sum(accs) / len(accs),
+        "acc_stderr,none": math.sqrt(sum(e * e for e in errs)) / len(accs),
+    }
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--primus-dir", required=True)
+    ap.add_argument("--fla-dir", required=True)
+    ap.add_argument("--name-primus", default="Primus")
+    ap.add_argument("--name-fla", default="FLA")
+    args = ap.parse_args()
+
+    p_json = _find_results_json(Path(args.primus_dir))
+    f_json = _find_results_json(Path(args.fla_dir))
+    if not p_json or not f_json:
+        print(f"ERROR: missing results json. primus={p_json}, fla={f_json}")
+        sys.exit(1)
+
+    with open(p_json) as f:
+        primus_results = json.load(f).get("results", {})
+    with open(f_json) as f:
+        fla_results = json.load(f).get("results", {})
+
+    # Materialize mmlu rollup if it's split across 57 leaf tasks
+    if "mmlu" not in primus_results:
+        agg = _mmlu_aggregate(primus_results)
+        if agg is not None:
+            primus_results = {**primus_results, "mmlu": agg}
+    if "mmlu" not in fla_results:
+        agg = _mmlu_aggregate(fla_results)
+        if agg is not None:
+            fla_results = {**fla_results, "mmlu": agg}
+
+    print(f"{args.name_primus:8} results: {p_json}")
+    print(f"{args.name_fla:8} results: {f_json}")
+    print()
+
+    HDR = f"{'Task':22} {'Metric':9} {args.name_primus:>8} {'±stderr':>9}  {args.name_fla:>8} {'±stderr':>9}  {'Δ abs':>8}  {'σ-units':>8}"
+    print(HDR)
+    print("-" * len(HDR))
+
+    abs_deltas = []
+    z_scores = []
+    pvals = []
+    fvals = []
+    rows_for_max = []
+
+    for display_name, task_id, metrics in DEFAULT_REPORT:
+        ptask = primus_results.get(task_id)
+        ftask = fla_results.get(task_id)
+        for m_i, metric in enumerate(metrics):
+            pv, pe = _metric_lookup(ptask, metric)
+            fv, fe = _metric_lookup(ftask, metric)
+            if pv is None or fv is None:
+                print(f"{(display_name if m_i==0 else ''):22} {metric:9}    (missing)")
+                continue
+            d = pv - fv
+            denom = math.sqrt(pe * pe + fe * fe) if (pe or fe) else 0.0
+            z = (abs(d) / denom) if denom > 0 else 0.0
+            label_left = display_name if m_i == 0 else ""
+            print(
+                f"{label_left:22} {metric:9} {pv:8.4f} ±{pe:7.4f}  {fv:8.4f} ±{fe:7.4f}  "
+                f"{d:+8.4f}  {z:8.2f}"
+            )
+            abs_deltas.append(abs(d))
+            z_scores.append(z)
+            pvals.append(pv)
+            fvals.append(fv)
+            rows_for_max.append((abs(d), f"{display_name} {metric}", d))
+
+    print("-" * len(HDR))
+    n = len(abs_deltas)
+    if n:
+        mean_p = sum(pvals) / n
+        mean_f = sum(fvals) / n
+        mean_d = mean_p - mean_f
+        mean_abs_d = sum(abs_deltas) / n
+        mean_z = sum(z_scores) / n
+        max_abs_d, max_label, max_d_signed = max(rows_for_max, key=lambda r: r[0])
+        max_z = max(z_scores)
+        print(
+            f"{'Mean of '+str(n)+' metrics':22} {'':9} {mean_p:8.4f} {'':>8}  {mean_f:8.4f} {'':>8}  "
+            f"{mean_d:+8.4f}  {mean_z:8.2f}"
+        )
+        print(f"{'Mean |Δ|':22} {'':9} {'':>8} {'':>8}  {'':>8} {'':>8}  " f"{mean_abs_d:8.4f}  {'':>8}")
+        print(
+            f"{'Max  |Δ|':22} {'':9} {'':>8} {'':>8}  {'':>8} {'':>8}  "
+            f"{max_d_signed:+8.4f}  {max_z:8.2f}    ({max_label})"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/consolidate_distcp_to_torch.py b/tools/hybrid/consolidate_distcp_to_torch.py
new file mode 100644
index 000000000..92bdfa405
--- /dev/null
+++ b/tools/hybrid/consolidate_distcp_to_torch.py
@@ -0,0 +1,256 @@
+#!/usr/bin/env python3
+"""
+Consolidate a Primus FSDP-dtensor (distcp) checkpoint into the legacy
+single-rank ``mp_rank_00/model_optim_rng.pt`` layout that the existing
+``tools/convert_*_to_fla_hf.py`` converters consume.
+
+Background
+----------
+When Primus trains with ``ckpt_format: fsdp_dtensor`` (Megatron-FSDP path,
+used for the 1B Pure-GDN 100B run), Megatron writes one ``__N_0.distcp``
+shard per data-parallel rank plus a ``.metadata`` file via
+``torch.distributed.checkpoint``. Two transforms happen relative to the
+legacy ``torch`` ckpt format:
+
+1. All model parameters are nested under ``model.module.`` (the FSDP
+   wrapper class layer).
+2. Fused SwiGLU ``linear_fc1.weight`` (shape ``[2*intermediate, hidden]``)
+   is **split** by FSDP into two halves at save time:
+       * ``linear_fc1.weight_w`` = first half (gate proj, shape
+         ``[intermediate, hidden]``)
+       * ``linear_fc1.weight_v`` = second half (up proj, same shape)
+   See ``third_party/Megatron-LM/megatron/core/transformer/fsdp_dtensor_checkpoint.py``
+   ``split_swiglu_linear_fc1`` for the source-of-truth split.
+
+This script reverses both transforms and writes a single
+``mp_rank_00/model_optim_rng.pt`` file in the format
+``convert_gdn_to_fla_hf.py`` / ``convert_kda_to_fla_hf.py`` /
+``convert_gdn_hybrid_to_fla_hf.py`` already understand:
+
+    {"model": {<flat parameter dict, no FSDP prefix, fused fc1>},
+     "iteration": <int>,
+     "checkpoint_version": <float>}
+
+Usage
+-----
+    python3 tools/hybrid/consolidate_distcp_to_torch.py \
+        --distcp-dir output/amd/root/zebra_llama_1B_gdn_pure_100B-pretrain/checkpoints/iter_0095368 \
+        --output-dir output/amd/root/zebra_llama_1B_gdn_pure_100B-pretrain/checkpoints_consolidated/iter_0095368
+
+The output dir gets a ``mp_rank_00/model_optim_rng.pt`` file ready for the
+HF converters. Memory: ~14 GB peak (full 1B model in fp32/bf16 on CPU).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+from torch.distributed.checkpoint import FileSystemReader
+from torch.distributed.checkpoint.state_dict_loader import _load_state_dict_from_keys
+
+FSDP_PREFIX = "model.module."
+SWIGLU_W_SUFFIX = ".weight_w"
+SWIGLU_V_SUFFIX = ".weight_v"
+
+
+def _is_model_weight_key(key: str) -> bool:
+    """Filter out optimizer / scheduler / RNG / iteration / args keys."""
+    return key.startswith(FSDP_PREFIX)
+
+
+def _strip_fsdp_prefix(key: str) -> str:
+    assert key.startswith(FSDP_PREFIX), key
+    return key[len(FSDP_PREFIX) :]
+
+
+def _enumerate_model_keys(distcp_dir: Path) -> tuple[list[str], int]:
+    """Return (model_weight_keys, iteration) from the .metadata file."""
+    reader = FileSystemReader(str(distcp_dir))
+    md = reader.read_metadata()
+    keys = sorted(k for k in md.state_dict_metadata.keys() if _is_model_weight_key(k))
+    iteration = 0
+    iter_re = re.search(r"iter_0*(\d+)$", distcp_dir.name)
+    if iter_re:
+        iteration = int(iter_re.group(1))
+    return keys, iteration
+
+
+def _flatten(prefix: str, obj, out: dict[str, torch.Tensor]) -> None:
+    """Recursively flatten a nested dict (with dotted keys) into a single map."""
+    if isinstance(obj, dict):
+        for k, v in obj.items():
+            new_prefix = f"{prefix}.{k}" if prefix else str(k)
+            _flatten(new_prefix, v, out)
+    else:
+        out[prefix] = obj
+
+
+def _load_tensors(distcp_dir: Path, keys: list[str]) -> dict[str, torch.Tensor]:
+    """Pull the listed keys out of the distcp shards into a single CPU dict.
+
+    Uses torch's single-process DCP loader (``_load_state_dict_from_keys``);
+    no torch.distributed init is required.  The loader returns a NESTED dict
+    (because dotted keys are interpreted as paths), so we flatten it back.
+    """
+    print(f"[load] reading {len(keys)} tensors from {distcp_dir} ...")
+    nested = _load_state_dict_from_keys(
+        keys=set(keys),
+        checkpoint_id=str(distcp_dir),
+        storage_reader=FileSystemReader(str(distcp_dir)),
+    )
+    flat: dict[str, torch.Tensor] = {}
+    _flatten("", nested, flat)
+    # Drop anything we didn't ask for (the loader may pull adjacent BytesMetadata)
+    flat = {k: v for k, v in flat.items() if k in set(keys)}
+    print(f"[load] got {len(flat)} tensors back (flattened from nested dict)")
+    if len(flat) != len(keys):
+        missing = set(keys) - set(flat)
+        if missing:
+            raise RuntimeError(
+                f"Loader returned {len(flat)}/{len(keys)} requested keys. "
+                f"Missing examples: {list(missing)[:5]}"
+            )
+    return flat
+
+
+def _refuse_swiglu(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Re-concat ``...linear_fc1.weight_w`` + ``...linear_fc1.weight_v``
+    pairs back into the original fused ``...linear_fc1.weight``.
+
+    FSDP split puts ``weight_w`` first (gate) and ``weight_v`` second (up).
+    Megatron's SwiGLU expects the same order: ``cat([gate, up], dim=0)``.
+    """
+    out = OrderedDict()
+    pending_w: dict[str, torch.Tensor] = {}
+    pending_v: dict[str, torch.Tensor] = {}
+
+    for key, tensor in state.items():
+        if key.endswith(SWIGLU_W_SUFFIX):
+            base = key[: -len(SWIGLU_W_SUFFIX)] + ".weight"
+            pending_w[base] = tensor
+        elif key.endswith(SWIGLU_V_SUFFIX):
+            base = key[: -len(SWIGLU_V_SUFFIX)] + ".weight"
+            pending_v[base] = tensor
+        else:
+            out[key] = tensor
+
+    fused = 0
+    for base, w in pending_w.items():
+        v = pending_v.pop(base, None)
+        if v is None:
+            raise KeyError(
+                f"Found {base}_w but no matching {base}_v in checkpoint. " "SwiGLU split is incomplete."
+            )
+        if w.shape != v.shape:
+            raise ValueError(
+                f"Shape mismatch for SwiGLU pair {base}: "
+                f"weight_w={tuple(w.shape)} vs weight_v={tuple(v.shape)}"
+            )
+        out[base] = torch.cat([w, v], dim=0)
+        fused += 1
+
+    if pending_v:
+        raise KeyError(f"Dangling weight_v entries with no weight_w: {list(pending_v)[:5]}")
+
+    if fused:
+        print(f"[fuse] re-fused {fused} SwiGLU linear_fc1 pairs " f"(weight_w + weight_v -> weight)")
+    return out
+
+
+def consolidate(distcp_dir: Path, output_dir: Path) -> Path:
+    """Materialize a legacy mp_rank_00/model_optim_rng.pt from distcp shards."""
+    keys, iteration = _enumerate_model_keys(distcp_dir)
+    if not keys:
+        raise RuntimeError(
+            f"No model weight keys (model.module.*) found in {distcp_dir}. "
+            "Is this really an FSDP-dtensor checkpoint?"
+        )
+
+    raw = _load_tensors(distcp_dir, keys)
+
+    # Strip the FSDP wrapper prefix.
+    stripped = OrderedDict((_strip_fsdp_prefix(k), v) for k, v in raw.items())
+    print(f"[strip] removed '{FSDP_PREFIX}' prefix from {len(stripped)} keys")
+
+    # Re-fuse SwiGLU fc1 if present.
+    fused = _refuse_swiglu(stripped)
+
+    # Pack into the converter-expected envelope.
+    payload = {
+        "model": fused,
+        "iteration": iteration,
+        "checkpoint_version": 3.0,  # arbitrary, matches Megatron's writer
+        "args": None,  # converters don't read this
+    }
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    rank_dir = output_dir / "mp_rank_00"
+    rank_dir.mkdir(exist_ok=True)
+    out_path = rank_dir / "model_optim_rng.pt"
+    torch.save(payload, out_path)
+
+    # Also drop a manifest for debugging.
+    manifest = {
+        "source_distcp_dir": str(distcp_dir),
+        "iteration": iteration,
+        "num_tensors": len(fused),
+        "tensor_keys": sorted(fused.keys()),
+    }
+    with open(output_dir / "consolidation_manifest.json", "w") as f:
+        json.dump(manifest, f, indent=2)
+
+    size_mb = out_path.stat().st_size / 1e6
+    print(f"\n[save] {out_path}  ({size_mb:.1f} MB, {len(fused)} tensors, iter={iteration})")
+    print(f"[save] {output_dir / 'consolidation_manifest.json'}")
+    return out_path
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument(
+        "--distcp-dir",
+        required=True,
+        type=Path,
+        help="Primus iter dir holding the .distcp shards " "(e.g. .../checkpoints/iter_0095368)",
+    )
+    ap.add_argument(
+        "--output-dir",
+        required=True,
+        type=Path,
+        help="Where to write the consolidated mp_rank_00 layout. "
+        "Typically '<distcp-dir-parent>_consolidated/<iter_name>'.",
+    )
+    args = ap.parse_args()
+
+    if not args.distcp_dir.is_dir():
+        ap.error(f"--distcp-dir does not exist or is not a directory: {args.distcp_dir}")
+    if not (args.distcp_dir / ".metadata").is_file():
+        ap.error(f"No .metadata file in {args.distcp_dir} — " "is this an FSDP-dtensor checkpoint?")
+
+    print("=" * 78)
+    print(" Primus FSDP-dtensor -> legacy mp_rank_00 consolidator")
+    print("=" * 78)
+    print(f"  source   = {args.distcp_dir}")
+    print(f"  dest     = {args.output_dir}")
+    print()
+
+    consolidate(args.distcp_dir, args.output_dir)
+    print()
+    print("Done. Feed this directory to any of the FLA HF converters, e.g.:")
+    print()
+    print(f"    python3 tools/hybrid/convert_gdn_to_fla_hf.py \\")
+    print(f"        --checkpoint-path {args.output_dir} \\")
+    print(f"        --output-dir output/gdn_pure_1B_fla_hf \\")
+    print(
+        f"        --config /home/<user>/flash-linear-attention/legacy/training/configs/gated_deltanet_1B_pure_100B.json"
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/hybrid/convert_fla_gdn_init_to_megatron.py b/tools/hybrid/convert_fla_gdn_init_to_megatron.py
new file mode 100644
index 000000000..336266ae0
--- /dev/null
+++ b/tools/hybrid/convert_fla_gdn_init_to_megatron.py
@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""
+Initialize a Megatron (Primus) checkpoint from FLA's GatedDeltaNet model weights.
+
+Creates a fake Megatron checkpoint directory that Primus can load via --load,
+ensuring both frameworks start from identical weights for loss-curve comparison.
+
+Usage:
+    python tools/hybrid/convert_fla_gdn_init_to_megatron.py \
+        --fla-config /path/to/gated_deltanet_300M_pure.json \
+        --output-dir output/fla_init_ckpt_300M \
+        --seed 42 \
+        --no-te   # use no-TE key names (WrappedTorchNorm, ColumnParallelLinear)
+
+This script was reconstructed verbatim from the agent transcript dated 2026-05-13
+(see docs/04-technical-guides/hybrid-models/gdn-fla-parity.md §"Files in the repo for this work").  It is one of the
+"forensics scripts" kept untracked in tools/hybrid/ per the parity-doc footnote.
+"""
+
+import argparse
+import json
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+_primus_root = Path(__file__).resolve().parents[2]
+_fla_root = _primus_root.parent / "flash-linear-attention"
+if str(_fla_root) not in sys.path:
+    sys.path.insert(0, str(_fla_root))
+
+
+def init_fla_model(config_path, seed=42):
+    """Initialize an FLA GatedDeltaNet model and return its state_dict."""
+    torch.manual_seed(seed)
+
+    from fla.models.gated_deltanet import GatedDeltaNetConfig, GatedDeltaNetForCausalLM
+
+    with open(config_path) as f:
+        cfg_dict = json.load(f)
+
+    config = GatedDeltaNetConfig(**cfg_dict)
+    # Disable fusions for the init pass — we just want the raw weights, not the
+    # fused-kernel-specific buffer initialisation.  Primus will turn fusions on
+    # at training time via PRIMUS_FLA_* env vars regardless of these.
+    config.fuse_cross_entropy = False
+    config.fuse_norm = False
+    config.fuse_swiglu = False
+
+    model = GatedDeltaNetForCausalLM(config)
+    model.eval()
+
+    print(f"FLA model initialized: {sum(p.numel() for p in model.parameters()):,} params")
+    return model.state_dict(), config
+
+
+def convert_fla_to_megatron(fla_state, fla_cfg, use_te=True):
+    """Convert FLA HF state dict to Megatron naming convention.
+
+    Each FLA layer becomes TWO Megatron decoder layers: the mixer sub-layer
+    (GDN) at index 2k and the MLP sub-layer at index 2k+1.  This mirrors
+    HybridStack's split layout (see primus/backends/megatron/core/models/hybrid).
+    """
+    hidden_size = fla_cfg.hidden_size
+    num_heads = fla_cfg.num_heads
+    num_v_heads = getattr(fla_cfg, "num_v_heads", num_heads)
+    head_dim = fla_cfg.head_dim
+    expand_v = getattr(fla_cfg, "expand_v", 1.0)
+    intermediate_size = fla_cfg.intermediate_size
+    num_hidden_layers = fla_cfg.num_hidden_layers
+
+    head_k_dim = head_dim
+    head_v_dim = int(head_dim * expand_v)
+    key_dim = num_heads * head_k_dim
+    value_dim = num_v_heads * head_v_dim
+
+    print(f"\nConverting FLA -> Megatron (use_te={use_te}):")
+    print(f"  hidden_size={hidden_size}, num_heads={num_heads}, num_v_heads={num_v_heads}")
+    print(f"  head_dim={head_dim}, expand_v={expand_v}")
+    print(f"  key_dim={key_dim}, value_dim={value_dim}")
+    print(f"  intermediate_size={intermediate_size}, num_hidden_layers={num_hidden_layers}")
+
+    mg_state = OrderedDict()
+
+    mg_state["embedding.word_embeddings.weight"] = fla_state["model.embeddings.weight"].clone()
+
+    for fla_idx in range(num_hidden_layers):
+        gdn_idx = fla_idx * 2
+        mlp_idx = fla_idx * 2 + 1
+        fp = f"model.layers.{fla_idx}"
+
+        # ── GDN sub-layer ──
+        if use_te:
+            mg_state[f"decoder.layers.{gdn_idx}.mixer.in_proj.layer_norm_weight"] = fla_state[
+                f"{fp}.attn_norm.weight"
+            ].clone()
+        else:
+            mg_state[f"decoder.layers.{gdn_idx}.norm.weight"] = fla_state[f"{fp}.attn_norm.weight"].clone()
+
+        # Fuse separate projections into in_proj: [q, k, v, gate, beta, alpha]
+        q_w = fla_state[f"{fp}.attn.q_proj.weight"]
+        k_w = fla_state[f"{fp}.attn.k_proj.weight"]
+        v_w = fla_state[f"{fp}.attn.v_proj.weight"]
+        g_w = fla_state[f"{fp}.attn.g_proj.weight"]
+        b_w = fla_state[f"{fp}.attn.b_proj.weight"]
+        a_w = fla_state[f"{fp}.attn.a_proj.weight"]
+        in_proj_w = torch.cat([q_w, k_w, v_w, g_w, b_w, a_w], dim=0)
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.in_proj.weight"] = in_proj_w
+
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.A_log"] = fla_state[f"{fp}.attn.A_log"].clone()
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.dt_bias"] = fla_state[f"{fp}.attn.dt_bias"].clone()
+
+        # Fuse conv1d: [q_conv, k_conv, v_conv]
+        q_conv = fla_state[f"{fp}.attn.q_conv1d.weight"]
+        k_conv = fla_state[f"{fp}.attn.k_conv1d.weight"]
+        v_conv = fla_state[f"{fp}.attn.v_conv1d.weight"]
+        conv_w = torch.cat([q_conv, k_conv, v_conv], dim=0)
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.conv1d.weight"] = conv_w
+
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.out_norm.weight"] = fla_state[
+            f"{fp}.attn.o_norm.weight"
+        ].clone()
+        mg_state[f"decoder.layers.{gdn_idx}.mixer.out_proj.weight"] = fla_state[
+            f"{fp}.attn.o_proj.weight"
+        ].clone()
+
+        # ── MLP sub-layer ──
+        if use_te:
+            mg_state[f"decoder.layers.{mlp_idx}.mlp.linear_fc1.layer_norm_weight"] = fla_state[
+                f"{fp}.mlp_norm.weight"
+            ].clone()
+        else:
+            mg_state[f"decoder.layers.{mlp_idx}.pre_mlp_layernorm.weight"] = fla_state[
+                f"{fp}.mlp_norm.weight"
+            ].clone()
+
+        gate_w = fla_state[f"{fp}.mlp.gate_proj.weight"]
+        up_w = fla_state[f"{fp}.mlp.up_proj.weight"]
+        fc1_w = torch.cat([gate_w, up_w], dim=0)
+        mg_state[f"decoder.layers.{mlp_idx}.mlp.linear_fc1.weight"] = fc1_w
+
+        mg_state[f"decoder.layers.{mlp_idx}.mlp.linear_fc2.weight"] = fla_state[
+            f"{fp}.mlp.down_proj.weight"
+        ].clone()
+
+    mg_state["decoder.final_norm.weight"] = fla_state["model.norm.weight"].clone()
+
+    print(f"  Converted {len(mg_state)} parameter tensors")
+    return mg_state
+
+
+def save_megatron_checkpoint(mg_state, output_dir, iteration=0):
+    """Save as a Megatron checkpoint directory structure."""
+    ckpt_dir = Path(output_dir) / f"iter_{iteration:07d}" / "mp_rank_00"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = {
+        "iteration": iteration,
+        "model": mg_state,
+        "checkpoint_version": 3.0,
+    }
+
+    ckpt_path = ckpt_dir / "model_optim_rng.pt"
+    torch.save(checkpoint, ckpt_path)
+
+    iter_file = Path(output_dir) / "latest_checkpointed_iteration.txt"
+    iter_file.write_text(str(iteration))
+
+    print(f"\nSaved Megatron checkpoint to: {ckpt_dir}")
+    print(f"  {ckpt_path}: {ckpt_path.stat().st_size / 1e6:.1f} MB")
+    return str(Path(output_dir))
+
+
+def verify_conversion(fla_state, mg_state, fla_cfg, use_te):
+    """Verify total parameter count matches."""
+    fla_params = sum(v.numel() for v in fla_state.values())
+    mg_params = sum(v.numel() for v in mg_state.values())
+
+    # FLA has both model.embeddings.weight and lm_head.weight (tied).
+    # Megatron only has embedding.word_embeddings.weight (output_layer shares).
+    fla_unique = fla_params - fla_state["model.embeddings.weight"].numel()
+
+    print(f"\nVerification:")
+    print(f"  FLA params (unique): {fla_unique:,}")
+    print(f"  Megatron params:     {mg_params:,}")
+
+    if fla_unique != mg_params:
+        print(f"  WARNING: param count mismatch! Diff = {abs(fla_unique - mg_params):,}")
+    else:
+        print(f"  OK -- param counts match")
+
+    emb_match = torch.equal(
+        fla_state["model.embeddings.weight"], mg_state["embedding.word_embeddings.weight"]
+    )
+    norm_key = (
+        "decoder.layers.0.norm.weight" if not use_te else "decoder.layers.0.mixer.in_proj.layer_norm_weight"
+    )
+    norm_match = torch.equal(fla_state["model.layers.0.attn_norm.weight"], mg_state[norm_key])
+    print(f"  Embedding match: {emb_match}")
+    print(f"  Layer 0 norm match: {norm_match}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Initialize Primus checkpoint from FLA weights")
+    parser.add_argument("--fla-config", type=str, required=True, help="Path to FLA config JSON")
+    parser.add_argument(
+        "--output-dir", type=str, required=True, help="Output directory for Megatron checkpoint"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for FLA model initialization")
+    parser.add_argument(
+        "--no-te", action="store_true", help="Use no-TE key names (WrappedTorchNorm, ColumnParallelLinear)"
+    )
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print("FLA -> Primus (Megatron) Weight Initialization")
+    print("=" * 70)
+
+    fla_state, fla_cfg = init_fla_model(args.fla_config, args.seed)
+
+    use_te = not args.no_te
+    mg_state = convert_fla_to_megatron(fla_state, fla_cfg, use_te=use_te)
+
+    verify_conversion(fla_state, mg_state, fla_cfg, use_te)
+
+    ckpt_path = save_megatron_checkpoint(mg_state, args.output_dir)
+
+    print(f"\n{'='*70}")
+    print("Done! To use in Primus, add to your training config:")
+    print(f"  load: {ckpt_path}")
+    print(f"  finetune: true    # load weights only, fresh optimizer")
+    print(f"  no_load_optim: true")
+    print(f"  no_load_rng: true")
+    print(f"{'='*70}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_fla_kda_init_to_megatron.py b/tools/hybrid/convert_fla_kda_init_to_megatron.py
new file mode 100644
index 000000000..966923b20
--- /dev/null
+++ b/tools/hybrid/convert_fla_kda_init_to_megatron.py
@@ -0,0 +1,354 @@
+#!/usr/bin/env python3
+"""
+Generate an FLA-equivalent random-init KDA-300M checkpoint and convert it
+into the Megatron sharded format that Primus loads.
+
+This is the KDA counterpart of the GDN init-checkpoint dance documented in
+docs/04-technical-guides/hybrid-models/gdn-fla-parity.md.  It is the *only* code change that closed the residual
+loss-curve gap for GDN, and the same is expected to apply to KDA:
+
+  1. Re-seed PyTorch with FLA's training seed (default 42).
+  2. Instantiate FLA's `KDAForCausalLM` from the same JSON config FLA used
+     for the 4768-iter training run (`kda_300M_pure.json`).  FLA's
+     `_init_weights` is called automatically by `PreTrainedModel.__init__`,
+     reproducing the exact initial weights FLA saw at iter 0.
+  3. Re-map the resulting state_dict into Primus's Megatron layout:
+       - 12 FLA `KDABlock` blocks  →  24 alternating Megatron sublayers
+         (even = KDA mixer, odd = MLP).
+       - `attn.q_proj / k_proj / v_proj` concatenated row-wise into the
+         fused `mixer.in_proj.weight` Primus uses.
+       - `attn.q_conv1d / k_conv1d / v_conv1d` concatenated row-wise into
+         the fused `mixer.conv1d.weight` Primus uses.
+       - `mlp.gate_proj / up_proj` concatenated row-wise into the fused
+         SwiGLU `mlp.linear_fc1.weight` Primus uses.
+       - `attn.A_log` reshaped `[H] → [1, 1, H, 1]` to match Primus's
+         per-head storage layout.
+       - All other tensors copy 1:1.
+  4. Pack the renamed state_dict into a Megatron checkpoint pickle
+     (`iter_0000000/mp_rank_00/model_optim_rng.pt`) and write
+     `latest_checkpointed_iteration.txt`.
+
+After running this tool, the YAML flips to:
+
+    spec: ['...', 'kda_hybrid_stack_spec_no_te']
+    use_fla_kda_in_kernel_gate: true
+    use_fla_fused_norm_gated: true
+    finetune: true
+    auto_continue_train: false
+    no_load_optim: true
+    no_load_rng: true
+    load: <OUT_DIR from this script>
+
+…and Primus's iter-1 loss should be bit-identical to FLA's
+(`11.965` ≈ `95.74 / 8`), proving the kernel + init are both matched.
+
+Usage
+-----
+    PYTHONPATH=/home/<user>/flash-linear-attention \\
+      python3 tools/hybrid/convert_fla_kda_init_to_megatron.py \\
+        --fla-config /home/<user>/flash-linear-attention/legacy/training/configs/kda_300M_pure.json \\
+        --output-dir output/fla_init_kda_300M \\
+        --seed 42
+
+Verification (post-run)
+-----------------------
+The script prints a per-tensor `OK / MISSING / SHAPE-MISMATCH` table.  All
+tensors should say `OK`.  If anything is `MISSING`, the YAML's `spec:` must
+be `kda_hybrid_stack_spec_no_te` (the TE spec uses different keys for the
+pre-norm because it fuses it into `in_proj`).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+
+def _expected_megatron_keys(num_layers: int) -> set[str]:
+    """The full set of Megatron keys we expect to produce for a 12-layer KDA-300M
+    pure model under the *no-TE* spec.  Used as a self-check at the end of the
+    conversion so we crash early if Primus ever changes its layout.
+    """
+    keys = {
+        "embedding.word_embeddings.weight",
+        "decoder.final_norm.weight",
+    }
+    for i in range(num_layers):
+        kda = 2 * i
+        mlp = 2 * i + 1
+        keys.update(
+            {
+                f"decoder.layers.{kda}.norm.weight",
+                f"decoder.layers.{kda}.mixer.in_proj.weight",
+                f"decoder.layers.{kda}.mixer.conv1d.weight",
+                f"decoder.layers.{kda}.mixer.f_b_proj.weight",
+                f"decoder.layers.{kda}.mixer.A_log",
+                f"decoder.layers.{kda}.mixer.dt_bias",
+                f"decoder.layers.{kda}.mixer.g_b_proj.weight",
+                f"decoder.layers.{kda}.mixer.g_b_proj.bias",
+                f"decoder.layers.{kda}.mixer.out_norm.weight",
+                f"decoder.layers.{kda}.mixer.out_proj.weight",
+                f"decoder.layers.{mlp}.pre_mlp_layernorm.weight",
+                f"decoder.layers.{mlp}.mlp.linear_fc1.weight",
+                f"decoder.layers.{mlp}.mlp.linear_fc2.weight",
+            }
+        )
+    return keys
+
+
+def build_fla_init(fla_config_path: Path, seed: int) -> tuple[OrderedDict, dict]:
+    """Instantiate FLA's KDAForCausalLM and return (state_dict, config_dict).
+
+    The transformers `set_seed` re-seeds Python, NumPy, and PyTorch.
+    FLA's `_init_weights` runs inside `__init__` via HuggingFace's
+    `PreTrainedModel.post_init()`, so the resulting state is identical to
+    what FLA's training loop starts with at step 0.
+    """
+    # FLA's top-level `import fla` pulls in every Triton kernel registration
+    # (and therefore needs a working Triton).  We only need the model class
+    # for state-dict construction (no forward pass), so bypass `fla/__init__.py`
+    # entirely and import the model module directly.  This lets the converter
+    # run on any Python that has `transformers` + a basic `torch` install —
+    # in particular on the Primus host outside the ROCm container.
+    try:
+        from transformers import set_seed
+    except ImportError as exc:
+        raise RuntimeError("transformers not installed: `pip install transformers`.") from exc
+
+    fla_root = os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention"))
+    if fla_root not in sys.path:
+        sys.path.insert(0, fla_root)
+    try:
+        from fla.models.kda.configuration_kda import KDAConfig
+        from fla.models.kda.modeling_kda import KDAForCausalLM
+    except ImportError as exc:
+        raise RuntimeError(
+            "Could not import FLA's KDA model module.  Set FLA_ROOT to a "
+            "checkout of https://github.com/fla-org/flash-linear-attention "
+            "(default: ~/flash-linear-attention)."
+        ) from exc
+
+    set_seed(seed)
+
+    with open(fla_config_path) as f:
+        cfg_dict = json.load(f)
+    config = KDAConfig(**cfg_dict)
+    # Force bf16 to match FLA training (configs sometimes default to fp32 here).
+    config.torch_dtype = "bfloat16"
+
+    print(f"[init] instantiating FLA KDAForCausalLM (seed={seed})…")
+    model = KDAForCausalLM(config).to(dtype=torch.bfloat16)
+    print(f"[init] params = {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
+
+    sd = OrderedDict()
+    for k, v in model.state_dict().items():
+        sd[k] = v.detach().contiguous().cpu()
+    return sd, cfg_dict
+
+
+def convert_fla_to_megatron(fla_sd: OrderedDict, cfg: dict) -> OrderedDict:
+    """Map FLA HF state_dict → Primus Megatron-sharded state_dict (no-TE spec)."""
+    hidden_size = cfg["hidden_size"]
+    num_heads = cfg["num_heads"]
+    num_v_heads = cfg.get("num_v_heads") or num_heads
+    head_dim = cfg["head_dim"]
+    expand_v = cfg.get("expand_v", 1.0)
+    intermediate = cfg["intermediate_size"]
+    num_layers = cfg["num_hidden_layers"]
+
+    head_v_dim = int(head_dim * expand_v)
+    qk_dim = num_heads * head_dim  # 256
+    v_dim = num_v_heads * head_v_dim  # 512
+
+    print(
+        f"[map ] hidden={hidden_size} num_heads={num_heads} num_v_heads={num_v_heads}\n"
+        f"       head_dim={head_dim} expand_v={expand_v} head_v_dim={head_v_dim}\n"
+        f"       qk_dim={qk_dim} v_dim={v_dim} intermediate={intermediate} "
+        f"layers={num_layers}"
+    )
+
+    out = OrderedDict()
+    out["embedding.word_embeddings.weight"] = fla_sd["model.embeddings.weight"]
+    out["decoder.final_norm.weight"] = fla_sd["model.norm.weight"]
+
+    for i in range(num_layers):
+        kda = 2 * i
+        mlp = 2 * i + 1
+        src = f"model.layers.{i}"
+
+        # ── KDA sublayer ────────────────────────────────────────────────
+        out[f"decoder.layers.{kda}.norm.weight"] = fla_sd[f"{src}.attn_norm.weight"]
+
+        # in_proj (fully fused, GDN-style): concat
+        #     [ q_proj | k_proj | v_proj | f_proj.0 | g_proj.0 | b_proj ]
+        # along dim 0.  Resulting shape:
+        #     [qk_dim*2 + v_dim + head_v_dim + head_v_dim + num_v_heads, hidden]
+        # This is the same matrix that ColumnParallelLinear would write — by
+        # concatenating FLA's six independent `hidden_states → X` projections
+        # in this exact order we make Primus's forward bit-identical to FLA's
+        # while paying only ONE matmul launch per layer instead of six.
+        q_w = fla_sd[f"{src}.attn.q_proj.weight"]
+        k_w = fla_sd[f"{src}.attn.k_proj.weight"]
+        v_w = fla_sd[f"{src}.attn.v_proj.weight"]
+        f_a_w = fla_sd[f"{src}.attn.f_proj.0.weight"]
+        g_a_w = fla_sd[f"{src}.attn.g_proj.0.weight"]
+        b_w = fla_sd[f"{src}.attn.b_proj.weight"]
+        assert q_w.shape == (qk_dim, hidden_size), q_w.shape
+        assert k_w.shape == (qk_dim, hidden_size), k_w.shape
+        assert v_w.shape == (v_dim, hidden_size), v_w.shape
+        assert f_a_w.shape == (head_v_dim, hidden_size), f_a_w.shape
+        assert g_a_w.shape == (head_v_dim, hidden_size), g_a_w.shape
+        assert b_w.shape == (num_v_heads, hidden_size), b_w.shape
+        out[f"decoder.layers.{kda}.mixer.in_proj.weight"] = torch.cat(
+            [q_w, k_w, v_w, f_a_w, g_a_w, b_w], dim=0
+        )
+
+        # conv1d: concat [q_conv1d, k_conv1d, v_conv1d] along dim 0.
+        # FLA stores each as [channels, 1, kernel]; concatenation gives
+        # [qk_dim*2 + v_dim, 1, kernel] = [1024, 1, 4].
+        qc = fla_sd[f"{src}.attn.q_conv1d.weight"]
+        kc = fla_sd[f"{src}.attn.k_conv1d.weight"]
+        vc = fla_sd[f"{src}.attn.v_conv1d.weight"]
+        out[f"decoder.layers.{kda}.mixer.conv1d.weight"] = torch.cat([qc, kc, vc], dim=0)
+
+        # f_proj.1 (low-rank gate expander):  Linear(head_v_dim, gate_dim)
+        out[f"decoder.layers.{kda}.mixer.f_b_proj.weight"] = fla_sd[f"{src}.attn.f_proj.1.weight"]
+
+        # A_log: FLA stores [num_v_heads]; Primus stores [1, 1, num_heads_local_tp, 1].
+        out[f"decoder.layers.{kda}.mixer.A_log"] = (
+            fla_sd[f"{src}.attn.A_log"].view(1, 1, num_v_heads, 1).clone()
+        )
+
+        # dt_bias: same shape (gate_dim = num_v_heads * head_dim).
+        out[f"decoder.layers.{kda}.mixer.dt_bias"] = fla_sd[f"{src}.attn.dt_bias"]
+
+        # g_proj.1 (low-rank output-gate expander):
+        #   Linear(head_v_dim, value_dim, bias=True)
+        out[f"decoder.layers.{kda}.mixer.g_b_proj.weight"] = fla_sd[f"{src}.attn.g_proj.1.weight"]
+        out[f"decoder.layers.{kda}.mixer.g_b_proj.bias"] = fla_sd[f"{src}.attn.g_proj.1.bias"]
+
+        # output norm:  FusedRMSNormGated.weight shape == [head_v_dim]
+        out[f"decoder.layers.{kda}.mixer.out_norm.weight"] = fla_sd[f"{src}.attn.o_norm.weight"]
+
+        # output projection:  Linear(value_dim, hidden_size)
+        out[f"decoder.layers.{kda}.mixer.out_proj.weight"] = fla_sd[f"{src}.attn.o_proj.weight"]
+
+        # ── MLP sublayer ────────────────────────────────────────────────
+        out[f"decoder.layers.{mlp}.pre_mlp_layernorm.weight"] = fla_sd[f"{src}.mlp_norm.weight"]
+
+        # SwiGLU fc1: concat [gate_proj, up_proj] along dim 0
+        gp = fla_sd[f"{src}.mlp.gate_proj.weight"]
+        up = fla_sd[f"{src}.mlp.up_proj.weight"]
+        assert gp.shape == (intermediate, hidden_size), gp.shape
+        assert up.shape == (intermediate, hidden_size), up.shape
+        out[f"decoder.layers.{mlp}.mlp.linear_fc1.weight"] = torch.cat([gp, up], dim=0)
+
+        # SwiGLU fc2: Linear(intermediate, hidden_size)
+        out[f"decoder.layers.{mlp}.mlp.linear_fc2.weight"] = fla_sd[f"{src}.mlp.down_proj.weight"]
+
+    return out
+
+
+def cross_check(mg_sd: OrderedDict, cfg: dict) -> None:
+    expected = _expected_megatron_keys(cfg["num_hidden_layers"])
+    got = set(mg_sd.keys())
+    missing = expected - got
+    extra = got - expected
+    if missing:
+        print(f"\n[FAIL] {len(missing)} expected Megatron keys are MISSING from the converted state_dict:")
+        for k in sorted(missing)[:25]:
+            print(f"    - {k}")
+        raise SystemExit(1)
+    if extra:
+        print(f"\n[warn] {len(extra)} unexpected extra keys in the converted state_dict (probably harmless):")
+        for k in sorted(extra)[:10]:
+            print(f"    + {k}")
+    print(f"[chk ] all {len(expected)} expected Megatron keys present ✓")
+
+
+def write_megatron_checkpoint(mg_sd: OrderedDict, output_dir: Path) -> None:
+    """Write `output_dir/iter_0000000/mp_rank_00/model_optim_rng.pt`.
+
+    The Megatron loader (`megatron.training.checkpointing.load_checkpoint`)
+    expects:
+        OUTPUT_DIR/
+            latest_checkpointed_iteration.txt  ← contains "0"
+            iter_0000000/
+                mp_rank_00/
+                    model_optim_rng.pt          ← pickle with keys
+                        {'iteration', 'model', 'checkpoint_version'}
+    With `finetune: true; no_load_optim: true; no_load_rng: true` Primus
+    only reads the `model` field of the pickle.
+    """
+    iter_dir = output_dir / "iter_0000000" / "mp_rank_00"
+    iter_dir.mkdir(parents=True, exist_ok=True)
+    ckpt_path = iter_dir / "model_optim_rng.pt"
+
+    ckpt = {
+        "iteration": 0,
+        "model": mg_sd,
+        "checkpoint_version": 3.0,
+    }
+    torch.save(ckpt, ckpt_path)
+    print(f"[save] wrote {ckpt_path}  ({ckpt_path.stat().st_size / 1e6:.1f} MB)")
+
+    # latest_checkpointed_iteration.txt — Megatron uses this to discover the
+    # most recent checkpoint when `auto_continue_train: false` is unset.
+    (output_dir / "latest_checkpointed_iteration.txt").write_text("0\n")
+    print(f"[save] wrote {output_dir / 'latest_checkpointed_iteration.txt'}")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    fla_root = os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention"))
+    p.add_argument(
+        "--fla-config",
+        type=Path,
+        default=Path(fla_root) / "legacy" / "training" / "configs" / "kda_300M_pure.json",
+    )
+    p.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("output/fla_init_kda_300M"),
+    )
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    print("=" * 78)
+    print("  FLA  KDA-300M  ─→  Primus Megatron sharded checkpoint")
+    print("=" * 78)
+    print(f"  fla_config = {args.fla_config}")
+    print(f"  output_dir = {args.output_dir}")
+    print(f"  seed       = {args.seed}")
+    print()
+
+    fla_sd, cfg = build_fla_init(args.fla_config, args.seed)
+    mg_sd = convert_fla_to_megatron(fla_sd, cfg)
+    cross_check(mg_sd, cfg)
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    write_megatron_checkpoint(mg_sd, args.output_dir)
+
+    print()
+    print("✓ done. Now update the YAML to:")
+    print(
+        f"    spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'kda_hybrid_stack_spec_no_te']"
+    )
+    print(f"    use_fla_kda_in_kernel_gate: true")
+    print(f"    use_fla_fused_norm_gated: true")
+    print(f"    finetune: true")
+    print(f"    auto_continue_train: false")
+    print(f"    no_load_optim: true")
+    print(f"    no_load_rng: true")
+    print(f"    load: {args.output_dir}")
+    print()
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/hybrid/convert_fla_to_megatron.py b/tools/hybrid/convert_fla_to_megatron.py
new file mode 100644
index 000000000..620f9ab81
--- /dev/null
+++ b/tools/hybrid/convert_fla_to_megatron.py
@@ -0,0 +1,121 @@
+"""
+Convert FLA's preprocessed Arrow dataset to Megatron binary format (.bin + .idx).
+
+FAST: reads Arrow shard files directly with PyArrow, avoids HuggingFace datasets overhead.
+Each FLA 2048-token sequence becomes one "document" in Megatron.
+"""
+
+import glob
+import os
+import struct
+import time
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FLA_ROOT = os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention"))
+
+FLA_DATA = os.environ.get(
+    "FLA_DATA",
+    f"{_FLA_ROOT}/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train",
+)
+OUT_PREFIX = os.environ.get(
+    "OUT_PREFIX",
+    str(_REPO_ROOT / "data" / "fla_aligned" / "fla_fineweb_edu_10BT_text_sentence"),
+)
+
+_INDEX_HEADER = b"MMIDIDX\x00\x00"
+DTYPE_CODE_INT32 = 4
+SEQ_LEN = 2048
+
+
+def main():
+    import pyarrow.ipc as ipc
+
+    t0 = time.time()
+    out_dir = Path(OUT_PREFIX).parent
+    out_dir.mkdir(parents=True, exist_ok=True)
+    bin_path = f"{OUT_PREFIX}.bin"
+    idx_path = f"{OUT_PREFIX}.idx"
+
+    # Find all Arrow data shard files (sorted by shard number)
+    shard_files = sorted(glob.glob(os.path.join(FLA_DATA, "data-*.arrow")))
+    print(f"Found {len(shard_files)} Arrow shard files")
+
+    total_tokens = 0
+    num_samples = 0
+
+    print(f"Writing {bin_path}...")
+    with open(bin_path, "wb") as f_bin:
+        for i, shard_path in enumerate(shard_files):
+            t1 = time.time()
+            # Read Arrow IPC stream (HF datasets memory-mapped format)
+            mmap = pa.memory_map(shard_path, "r")
+            table = ipc.open_stream(mmap).read_all()
+            col = table.column("input_ids")
+
+            # Each element is a list of int32. Flatten per-chunk.
+            for chunk in col.chunks:
+                # chunk is a ListArray; .values gives the flat child array
+                flat = chunk.values.to_numpy(zero_copy_only=False)
+                if flat.dtype != np.int32:
+                    flat = flat.astype(np.int32)
+                f_bin.write(flat.tobytes())
+                total_tokens += len(flat)
+                num_samples += len(chunk)
+
+            elapsed = time.time() - t0
+            shard_elapsed = time.time() - t1
+            pct = (i + 1) / len(shard_files) * 100
+            print(
+                f"  Shard {i+1:>3}/{len(shard_files)} ({pct:5.1f}%) | "
+                f"{num_samples:>10,} samples | {total_tokens:>13,} tokens | "
+                f"shard: {shard_elapsed:.1f}s | total: {elapsed:.0f}s"
+            )
+
+    expected_tokens = num_samples * SEQ_LEN
+    print(f"\n  Total: {num_samples:,} samples, {total_tokens:,} tokens")
+    if total_tokens != expected_tokens:
+        print(f"  WARNING: expected {expected_tokens:,} tokens (samples * {SEQ_LEN})")
+
+    # ── Write .idx ──
+    print(f"Writing {idx_path}...")
+    seq_lengths = np.full(num_samples, SEQ_LEN, dtype=np.int32)
+    seq_pointers = np.arange(num_samples, dtype=np.int64) * np.int64(SEQ_LEN * 4)
+    doc_indices = np.arange(1, num_samples + 1, dtype=np.int64)
+
+    with open(idx_path, "wb") as f:
+        f.write(_INDEX_HEADER)
+        f.write(struct.pack("<Q", 1))
+        f.write(struct.pack("<B", DTYPE_CODE_INT32))
+        f.write(struct.pack("<Q", num_samples))
+        f.write(struct.pack("<Q", num_samples))
+        f.write(seq_lengths.tobytes(order="C"))
+        f.write(seq_pointers.tobytes(order="C"))
+        f.write(doc_indices.tobytes(order="C"))
+
+    # ── Verify ──
+    print("\nVerifying...")
+    data = np.memmap(bin_path, dtype=np.int32, mode="r")
+    print(f"  Bin file tokens: {len(data):,}")
+    print(f"  First 10: {data[:10].tolist()}")
+    print(f"  Last 10:  {data[-10:].tolist()}")
+
+    # Cross-check with first shard (HF datasets writes IPC stream, not IPC file)
+    mmap = pa.memory_map(shard_files[0], "r")
+    tbl = ipc.open_stream(mmap).read_all()
+    first_sample = tbl.column("input_ids")[0].as_py()
+    assert data[:10].tolist() == first_sample[:10], "MISMATCH at start!"
+    print("  OK — tokens match!")
+
+    bin_sz = Path(bin_path).stat().st_size
+    idx_sz = Path(idx_path).stat().st_size
+    print(f"\n  {bin_path}: {bin_sz / 1e9:.2f} GB")
+    print(f"  {idx_path}: {idx_sz / 1e6:.2f} MB")
+    print(f"\nDone! Total: {time.time()-t0:.0f}s")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_gdn_hybrid_to_fla_hf.py b/tools/hybrid/convert_gdn_hybrid_to_fla_hf.py
new file mode 100644
index 000000000..d9b870798
--- /dev/null
+++ b/tools/hybrid/convert_gdn_hybrid_to_fla_hf.py
@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+"""Convert Primus (Megatron) 75% Hybrid GDN+MLA checkpoint to FLA HuggingFace format.
+
+The hybrid model has 12 FLA mixer blocks: MLA at indices [0, 4, 8] and GDN
+at the other 9 indices (matches FLA `gated_deltanet_300M_hybrid.json` exactly).
+
+Primus Megatron stores them as 24 sublayers alternating mixer/MLP:
+
+    sublayer   pattern  FLA mixer  what
+    --------   -------  ---------  ----
+    0          *        0          MLA
+    1          -                   MLP
+    2          M        1          GDN
+    3          -                   MLP
+    4          M        2          GDN
+    5          -                   MLP
+    6          M        3          GDN
+    7          -                   MLP
+    8          *        4          MLA
+    9          -                   MLP
+    ...                            (repeats with `*-M-M-M-` pattern)
+
+This converter walks the 12 FLA mixer indices, finds the corresponding Megatron
+sublayer (mixer + MLP), and emits FLA-format weights.
+
+The MLA mapping uses our spec's `q_layernorm=WrappedTorchNorm` / `kv_layernorm=
+WrappedTorchNorm` (added in the parity fix). Without those, the FLA-side `q_proj.1`
+and `kv_proj.1` RMSNorm weights would be missing and the FLA model would NaN.
+
+Usage (inside the container):
+
+    python tools/hybrid/convert_gdn_hybrid_to_fla_hf.py \
+        --checkpoint-path output/amd/root/zebra_llama_300M_gdn_hybrid-pretrain/checkpoints/iter_0004768 \
+        --output-dir output/gdn_hybrid_300M_fla_hf \
+        --config /home/<user>/flash-linear-attention/legacy/training/configs/gated_deltanet_300M_hybrid.json
+"""
+import argparse
+import json
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+_megatron_path = str(Path(__file__).resolve().parents[2] / "third_party" / "Megatron-LM")
+if _megatron_path not in sys.path:
+    sys.path.insert(0, _megatron_path)
+
+
+def load_megatron_checkpoint(checkpoint_path):
+    model_path = Path(checkpoint_path) / "mp_rank_00" / "model_optim_rng.pt"
+    if not model_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+    print(f"Loading checkpoint from: {model_path}")
+    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+    print(f"Loaded. Iteration: {checkpoint.get('iteration', '?')}")
+    return checkpoint
+
+
+def _get_first(state, *candidates):
+    for k in candidates:
+        if k in state:
+            return state[k], k
+    raise KeyError(
+        f"None of these expected keys were found:\n  "
+        + "\n  ".join(candidates)
+        + "\nFirst 30 keys in checkpoint:\n  "
+        + "\n  ".join(sorted(state.keys())[:30])
+    )
+
+
+def _build_layer_map(hybrid_pattern):
+    """Walk the hybrid_override_pattern and return:
+       fla_to_megatron[fla_layer_idx] = (mixer_megatron_sublayer, mlp_megatron_sublayer, kind)
+       kind: 'mla' | 'gdn'
+
+    `*` = MLA, `M` = GDN (Mamba slot), `-` = MLP."""
+    fla_to_megatron = []
+    cur_mixer = None
+    cur_kind = None
+    for i, c in enumerate(hybrid_pattern):
+        if c in ("*", "M"):
+            cur_mixer = i
+            cur_kind = "mla" if c == "*" else "gdn"
+        elif c == "-":
+            if cur_mixer is None:
+                # MLP without preceding mixer — should never happen for our pattern
+                continue
+            fla_to_megatron.append((cur_mixer, i, cur_kind))
+            cur_mixer = None
+            cur_kind = None
+    return fla_to_megatron
+
+
+def _megatron_to_fla_rope_channels(w_rope, qk_rope_head_dim):
+    """Permute the rope channels of a tensor whose LAST dim is `qk_rope_head_dim`
+    from Megatron's storage order to FLA's RoPE-ready order.
+
+    Megatron stores rope channels as [c0, c1, c2, ..., c_{d-1}] and on-the-fly
+    rearranges them to [c0, c2, c4, ..., c_{d-2}, c1, c3, ..., c_{d-1}] right
+    before applying RoPE (multi_latent_attention branch in rope_utils.py).
+    FLA's mla.py applies no such permutation — it assumes the weights are
+    already in the rearranged layout.
+
+    So when going Megatron→FLA we apply the permutation ONCE, baked into the
+    saved weights. The mapping is index `i` in Megatron → index
+    `i//2`               if i even   (lands in first half)
+    `qk_rope_head_dim//2 + i//2`  if i odd    (lands in second half)
+    which is equivalent to `cat([even-indexed, odd-indexed], dim=-1)`.
+    """
+    d = qk_rope_head_dim
+    assert w_rope.shape[-1] == d, (
+        f"_megatron_to_fla_rope_channels: last dim {w_rope.shape[-1]} != " f"qk_rope_head_dim {d}"
+    )
+    even = w_rope[..., 0::2]
+    odd = w_rope[..., 1::2]
+    return torch.cat([even, odd], dim=-1).contiguous()
+
+
+def convert_mla_block(state, sub_mixer, prefix, fla_cfg):
+    """MLA mixer Primus→FLA mapping. Writes:
+        {prefix}.attn.q_proj.0/1/2.weight
+        {prefix}.attn.kv_proj.0/1/2.weight
+        {prefix}.attn.k_rope.weight
+        {prefix}.attn.o_proj.weight
+
+    Critical: applies the [::2, 1::2] → [first-half, second-half] permutation
+    to the RoPE-carrying channels of `linear_q_up_proj` and `linear_kv_down_proj
+    [kv_lora_rank:]` to match FLA's RoPE convention (see rope_utils.py
+    multi_latent_attention branch which Megatron applies on-the-fly but FLA
+    does not).
+    """
+    out = OrderedDict()
+    attn_cfg = fla_cfg["attn"]
+    q_lora_rank = attn_cfg["q_lora_rank"]
+    kv_lora_rank = attn_cfg["kv_lora_rank"]
+    qk_rope_head_dim = attn_cfg["qk_rope_head_dim"]
+    num_heads = attn_cfg["num_heads"]
+    qk_nope_head_dim = attn_cfg["qk_nope_head_dim"]
+    v_head_dim = attn_cfg["v_head_dim"]
+    qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+    hidden_size = fla_cfg["hidden_size"]
+
+    # ── q path: down → norm → up ──
+    q_down_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_q_down_proj.weight"]
+    q_up_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_q_up_proj.weight"]
+    q_norm_w, _ = _get_first(
+        state,
+        f"decoder.layers.{sub_mixer}.self_attention.q_layernorm.weight",
+    )
+    assert q_down_w.shape == (q_lora_rank, hidden_size)
+    assert q_up_w.shape == (num_heads * qk_head_dim, q_lora_rank)
+    assert q_norm_w.shape == (q_lora_rank,)
+
+    # Permute the rope half of q_up_proj's output (per head). Reshape into
+    # (num_heads, qk_head_dim, q_lora_rank), split, permute rope, recombine.
+    q_up_3d = q_up_w.view(num_heads, qk_head_dim, q_lora_rank)
+    q_up_nope = q_up_3d[:, :qk_nope_head_dim, :]
+    q_up_rope_meg = q_up_3d[:, qk_nope_head_dim:, :]  # (h, rope, r)
+    q_up_rope_meg = q_up_rope_meg.transpose(-1, -2)  # (h, r, rope)
+    q_up_rope_fla = _megatron_to_fla_rope_channels(q_up_rope_meg, qk_rope_head_dim)
+    q_up_rope_fla = q_up_rope_fla.transpose(-1, -2)  # (h, rope, r)
+    q_up_fla = torch.cat([q_up_nope, q_up_rope_fla], dim=1)  # (h, qk_head_dim, r)
+    q_up_fla = q_up_fla.reshape(num_heads * qk_head_dim, q_lora_rank).contiguous()
+
+    out[f"{prefix}.attn.q_proj.0.weight"] = q_down_w
+    out[f"{prefix}.attn.q_proj.1.weight"] = q_norm_w
+    out[f"{prefix}.attn.q_proj.2.weight"] = q_up_fla
+
+    # ── kv path: fused down → split [kv_compressed | k_rope] → norm → up ──
+    kv_down_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_kv_down_proj.weight"]
+    expected_down_rows = kv_lora_rank + qk_rope_head_dim
+    assert kv_down_w.shape == (expected_down_rows, hidden_size)
+    out[f"{prefix}.attn.kv_proj.0.weight"] = kv_down_w[:kv_lora_rank].contiguous()
+
+    # k_rope's output channels need the same permutation. Shape (rope, hidden):
+    # rope dim is the output (last after transpose), so transpose → permute → back.
+    k_rope_meg = kv_down_w[kv_lora_rank:].contiguous()  # (rope, hidden)
+    k_rope_meg = k_rope_meg.transpose(0, 1)  # (hidden, rope)
+    k_rope_fla = _megatron_to_fla_rope_channels(k_rope_meg, qk_rope_head_dim)
+    k_rope_fla = k_rope_fla.transpose(0, 1).contiguous()  # (rope, hidden)
+    out[f"{prefix}.attn.k_rope.weight"] = k_rope_fla
+
+    kv_norm_w, _ = _get_first(
+        state,
+        f"decoder.layers.{sub_mixer}.self_attention.kv_layernorm.weight",
+    )
+    assert kv_norm_w.shape == (kv_lora_rank,)
+    out[f"{prefix}.attn.kv_proj.1.weight"] = kv_norm_w
+
+    kv_up_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_kv_up_proj.weight"]
+    expected_up_rows = num_heads * (qk_nope_head_dim + v_head_dim)
+    assert kv_up_w.shape == (expected_up_rows, kv_lora_rank)
+    out[f"{prefix}.attn.kv_proj.2.weight"] = kv_up_w
+
+    out[f"{prefix}.attn.o_proj.weight"] = state[
+        f"decoder.layers.{sub_mixer}.self_attention.linear_proj.weight"
+    ]
+    return out
+
+
+def convert_gdn_block(state, sub_mixer, prefix, fla_cfg):
+    """GDN mixer Primus→FLA mapping (same as convert_gdn_to_fla_hf.py)."""
+    out = OrderedDict()
+    hidden_size = fla_cfg["hidden_size"]
+    num_heads = fla_cfg["num_heads"]
+    num_v_heads = fla_cfg.get("num_v_heads", num_heads)
+    head_dim = fla_cfg["head_dim"]
+    expand_v = fla_cfg.get("expand_v", 1.0)
+    key_dim = num_heads * head_dim
+    value_dim = num_v_heads * int(head_dim * expand_v)
+
+    in_proj_w = state[f"decoder.layers.{sub_mixer}.mixer.in_proj.weight"]
+    expected_in = key_dim * 2 + value_dim * 2 + num_v_heads * 2
+    assert (
+        in_proj_w.shape[0] == expected_in
+    ), f"in_proj shape {tuple(in_proj_w.shape)} expected ({expected_in},{hidden_size})"
+    out[f"{prefix}.attn.q_proj.weight"] = in_proj_w[:key_dim]
+    out[f"{prefix}.attn.k_proj.weight"] = in_proj_w[key_dim : key_dim * 2]
+    out[f"{prefix}.attn.v_proj.weight"] = in_proj_w[key_dim * 2 : key_dim * 2 + value_dim]
+    out[f"{prefix}.attn.g_proj.weight"] = in_proj_w[key_dim * 2 + value_dim : key_dim * 2 + value_dim * 2]
+    out[f"{prefix}.attn.b_proj.weight"] = in_proj_w[
+        key_dim * 2 + value_dim * 2 : key_dim * 2 + value_dim * 2 + num_v_heads
+    ]
+    out[f"{prefix}.attn.a_proj.weight"] = in_proj_w[key_dim * 2 + value_dim * 2 + num_v_heads :]
+
+    out[f"{prefix}.attn.A_log"] = state[f"decoder.layers.{sub_mixer}.mixer.A_log"]
+    out[f"{prefix}.attn.dt_bias"] = state[f"decoder.layers.{sub_mixer}.mixer.dt_bias"]
+
+    conv_w = state[f"decoder.layers.{sub_mixer}.mixer.conv1d.weight"]
+    out[f"{prefix}.attn.q_conv1d.weight"] = conv_w[:key_dim]
+    out[f"{prefix}.attn.k_conv1d.weight"] = conv_w[key_dim : key_dim * 2]
+    out[f"{prefix}.attn.v_conv1d.weight"] = conv_w[key_dim * 2 :]
+
+    out[f"{prefix}.attn.o_norm.weight"] = state[f"decoder.layers.{sub_mixer}.mixer.out_norm.weight"]
+    out[f"{prefix}.attn.o_proj.weight"] = state[f"decoder.layers.{sub_mixer}.mixer.out_proj.weight"]
+    return out
+
+
+def convert(checkpoint, fla_config_path, hybrid_pattern):
+    state = checkpoint["model"]
+    with open(fla_config_path) as f:
+        fla_cfg = json.load(f)
+
+    num_hidden_layers = fla_cfg["num_hidden_layers"]
+    intermediate_size = fla_cfg.get("intermediate_size", fla_cfg["hidden_size"] * 4)
+
+    layer_map = _build_layer_map(hybrid_pattern)
+    assert len(layer_map) == num_hidden_layers, (
+        f"hybrid_pattern produced {len(layer_map)} FLA mixer blocks "
+        f"but FLA config wants {num_hidden_layers}"
+    )
+
+    print(f"\nFLA layer index → Megatron sublayer mapping:")
+    for fla_idx, (sub_mixer, sub_mlp, kind) in enumerate(layer_map):
+        print(f"  FLA layer {fla_idx} ({kind.upper()}): mixer=sub{sub_mixer}, mlp=sub{sub_mlp}")
+
+    hf_state = OrderedDict()
+    hf_state["model.embeddings.weight"] = state["embedding.word_embeddings.weight"]
+
+    for fla_idx, (sub_mixer, sub_mlp, kind) in enumerate(layer_map):
+        prefix = f"model.layers.{fla_idx}"
+
+        attn_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{sub_mixer}.mixer.in_proj.layer_norm_weight",
+            f"decoder.layers.{sub_mixer}.norm.weight",
+            f"decoder.layers.{sub_mixer}.input_layernorm.weight",
+        )
+        hf_state[f"{prefix}.attn_norm.weight"] = attn_norm_w
+
+        if kind == "mla":
+            hf_state.update(convert_mla_block(state, sub_mixer, prefix, fla_cfg))
+        else:
+            hf_state.update(convert_gdn_block(state, sub_mixer, prefix, fla_cfg))
+
+        # MLP sublayer (shared between MLA and GDN paths)
+        mlp_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{sub_mlp}.mlp.linear_fc1.layer_norm_weight",
+            f"decoder.layers.{sub_mlp}.pre_mlp_layernorm.weight",
+            f"decoder.layers.{sub_mlp}.input_layernorm.weight",
+        )
+        hf_state[f"{prefix}.mlp_norm.weight"] = mlp_norm_w
+
+        fc1_w = state[f"decoder.layers.{sub_mlp}.mlp.linear_fc1.weight"]
+        assert (
+            fc1_w.shape[0] == intermediate_size * 2
+        ), f"MLP fc1 shape {tuple(fc1_w.shape)} expected ({intermediate_size*2},...)"
+        hf_state[f"{prefix}.mlp.gate_proj.weight"] = fc1_w[:intermediate_size]
+        hf_state[f"{prefix}.mlp.up_proj.weight"] = fc1_w[intermediate_size:]
+        hf_state[f"{prefix}.mlp.down_proj.weight"] = state[f"decoder.layers.{sub_mlp}.mlp.linear_fc2.weight"]
+
+    final_norm_w, _ = _get_first(
+        state,
+        "decoder.final_norm.weight",
+        "decoder.final_layernorm.weight",
+        "decoder.norm.weight",
+    )
+    hf_state["model.norm.weight"] = final_norm_w
+
+    if "output_layer.weight" in state:
+        hf_state["lm_head.weight"] = state["output_layer.weight"]
+    else:
+        hf_state["lm_head.weight"] = state["embedding.word_embeddings.weight"]
+
+    return hf_state
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint-path", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument(
+        "--config",
+        default=os.path.join(
+            os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention")),
+            "legacy",
+            "training",
+            "configs",
+            "gated_deltanet_300M_hybrid.json",
+        ),
+    )
+    parser.add_argument(
+        "--hybrid-pattern",
+        default="*-M-M-M-*-M-M-M-*-M-M-M-",
+        help="Same as YAML hybrid_override_pattern (24 chars for 300M hybrid)",
+    )
+    args = parser.parse_args()
+
+    print("=" * 72)
+    print("Primus 75% Hybrid GDN+MLA → FLA HuggingFace Conversion")
+    print("=" * 72)
+
+    checkpoint = load_megatron_checkpoint(args.checkpoint_path)
+    hf_state = convert(checkpoint, args.config, args.hybrid_pattern)
+    print(f"\nConverted {len(hf_state)} parameters")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        from safetensors.torch import save_file
+
+        if (
+            "lm_head.weight" in hf_state
+            and "model.embeddings.weight" in hf_state
+            and hf_state["lm_head.weight"].data_ptr() == hf_state["model.embeddings.weight"].data_ptr()
+        ):
+            hf_state["lm_head.weight"] = hf_state["lm_head.weight"].clone()
+        save_file(hf_state, str(output_dir / "model.safetensors"))
+        print(f"  Saved {output_dir / 'model.safetensors'}")
+    except ImportError:
+        torch.save(hf_state, output_dir / "pytorch_model.bin")
+        print(f"  Saved {output_dir / 'pytorch_model.bin'}")
+
+    with open(args.config) as f:
+        config = json.load(f)
+    config["architectures"] = ["GatedDeltaNetForCausalLM"]
+    config["fuse_cross_entropy"] = False
+    config["fuse_norm"] = False
+    config["fuse_swiglu"] = False
+    config["fuse_linear_cross_entropy"] = False
+    config["torch_dtype"] = "bfloat16"
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(config, f, indent=2)
+    print(f"  Saved config.json")
+
+    tokenizer_config = {
+        "tokenizer_class": "PreTrainedTokenizerFast",
+        "model_max_length": 2048,
+    }
+    with open(output_dir / "tokenizer_config.json", "w") as f:
+        json.dump(tokenizer_config, f, indent=2)
+
+    print(f"\nDone. To evaluate:")
+    print(f"  lm_eval --model hf \\")
+    print(f"    --model_args pretrained={output_dir},trust_remote_code=True,dtype=bfloat16 \\")
+    print(f"    --tasks hellaswag,winogrande,piqa,arc_easy,arc_challenge,lambada_openai \\")
+    print(f"    --batch_size 16")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_gdn_to_fla_hf.py b/tools/hybrid/convert_gdn_to_fla_hf.py
new file mode 100644
index 000000000..02cf2d3bb
--- /dev/null
+++ b/tools/hybrid/convert_gdn_to_fla_hf.py
@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+"""
+Convert Primus (Megatron) Pure GDN checkpoint to FLA HuggingFace format.
+
+Primus uses fused projections (in_proj, conv1d, mlp.linear_fc1) and alternating
+GDN/MLP sublayers. FLA uses separate projections and combined layers.
+
+Usage:
+    python tools/hybrid/convert_gdn_to_fla_hf.py \
+        --checkpoint-path output/amd/root/zebra_llama_1B_gdn_pure-pretrain/checkpoints/iter_0076294 \
+        --output-dir output/gdn_pure_1B_fla_hf \
+        --config /path/to/gated_deltanet_1B_pure.json
+"""
+
+import argparse
+import json
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+# Ensure Megatron is importable (needed for torch.load to unpickle checkpoint)
+_megatron_path = str(Path(__file__).resolve().parents[2] / "third_party" / "Megatron-LM")
+if _megatron_path not in sys.path:
+    sys.path.insert(0, _megatron_path)
+
+
+def load_megatron_checkpoint(checkpoint_path):
+    model_path = Path(checkpoint_path) / "mp_rank_00" / "model_optim_rng.pt"
+    if not model_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+    print(f"Loading checkpoint from: {model_path}")
+    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+    print(f"Loaded. Iteration: {checkpoint.get('iteration', '?')}")
+    return checkpoint
+
+
+def _get_first(state, *candidates):
+    """Return state[k] for the first k present in candidates, else raise with a helpful message."""
+    for k in candidates:
+        if k in state:
+            return state[k], k
+    raise KeyError(
+        f"None of these expected keys were found in checkpoint:\n  "
+        + "\n  ".join(candidates)
+        + "\nFirst 30 actually-present keys:\n  "
+        + "\n  ".join(sorted(state.keys())[:30])
+    )
+
+
+def convert(checkpoint, fla_config_path):
+    """Convert Megatron state dict to FLA GatedDeltaNet HuggingFace format.
+
+    Handles both layouts:
+      * TE spec     (gdn_hybrid_stack_spec):
+          norm folded into linear → `mixer.in_proj.layer_norm_weight`,
+                                    `mlp.linear_fc1.layer_norm_weight`
+      * no-TE spec  (gdn_hybrid_stack_spec_no_te):
+          separate WrappedTorchNorm → `norm.weight`, `pre_mlp_layernorm.weight`
+    """
+    state = checkpoint["model"]
+
+    with open(fla_config_path) as f:
+        fla_cfg = json.load(f)
+
+    hidden_size = fla_cfg["hidden_size"]
+    num_heads = fla_cfg["num_heads"]
+    num_v_heads = fla_cfg.get("num_v_heads", num_heads)
+    head_dim = fla_cfg["head_dim"]
+    expand_v = fla_cfg.get("expand_v", 1.0)
+    intermediate_size = fla_cfg.get("intermediate_size", hidden_size * 4)
+    num_hidden_layers = fla_cfg["num_hidden_layers"]
+
+    head_k_dim = head_dim
+    head_v_dim = int(head_dim * expand_v)
+    key_dim = num_heads * head_k_dim
+    value_dim = num_v_heads * head_v_dim
+
+    print(f"\nModel config:")
+    print(f"  hidden_size={hidden_size}, num_heads={num_heads}, num_v_heads={num_v_heads}")
+    print(f"  head_dim={head_dim}, expand_v={expand_v}")
+    print(f"  key_dim={key_dim}, value_dim={value_dim}")
+    print(f"  intermediate_size={intermediate_size}, num_hidden_layers={num_hidden_layers}")
+    print(f"  in_proj_dim = {key_dim*2 + value_dim*2 + num_v_heads*2}")
+
+    hf_state = OrderedDict()
+
+    # Embeddings (FLA uses 'model.embeddings.weight', not 'model.embed_tokens.weight')
+    hf_state["model.embeddings.weight"] = state["embedding.word_embeddings.weight"]
+
+    for fla_layer_idx in range(num_hidden_layers):
+        gdn_idx = fla_layer_idx * 2
+        mlp_idx = fla_layer_idx * 2 + 1
+        prefix = f"model.layers.{fla_layer_idx}"
+
+        # ── GDN sublayer ──
+        attn_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{gdn_idx}.mixer.in_proj.layer_norm_weight",
+            f"decoder.layers.{gdn_idx}.norm.weight",
+            f"decoder.layers.{gdn_idx}.input_layernorm.weight",
+        )
+        hf_state[f"{prefix}.attn_norm.weight"] = attn_norm_w
+
+        # Fused in_proj split: [q(key_dim), k(key_dim), v(value_dim), gate(value_dim), beta(num_v_heads), alpha(num_v_heads)]
+        in_proj_w = state[f"decoder.layers.{gdn_idx}.mixer.in_proj.weight"]
+        assert (
+            in_proj_w.shape[0] == key_dim * 2 + value_dim * 2 + num_v_heads * 2
+        ), f"in_proj shape mismatch: {in_proj_w.shape}"
+
+        q_w = in_proj_w[:key_dim]
+        k_w = in_proj_w[key_dim : key_dim * 2]
+        v_w = in_proj_w[key_dim * 2 : key_dim * 2 + value_dim]
+        g_w = in_proj_w[key_dim * 2 + value_dim : key_dim * 2 + value_dim * 2]
+        b_w = in_proj_w[key_dim * 2 + value_dim * 2 : key_dim * 2 + value_dim * 2 + num_v_heads]
+        a_w = in_proj_w[key_dim * 2 + value_dim * 2 + num_v_heads :]
+
+        hf_state[f"{prefix}.attn.q_proj.weight"] = q_w
+        hf_state[f"{prefix}.attn.k_proj.weight"] = k_w
+        hf_state[f"{prefix}.attn.v_proj.weight"] = v_w
+        hf_state[f"{prefix}.attn.g_proj.weight"] = g_w
+        hf_state[f"{prefix}.attn.b_proj.weight"] = b_w
+        hf_state[f"{prefix}.attn.a_proj.weight"] = a_w
+
+        # A_log and dt_bias
+        hf_state[f"{prefix}.attn.A_log"] = state[f"decoder.layers.{gdn_idx}.mixer.A_log"]
+        hf_state[f"{prefix}.attn.dt_bias"] = state[f"decoder.layers.{gdn_idx}.mixer.dt_bias"]
+
+        # Fused conv1d split: [q_conv(key_dim, 1, 4), k_conv(key_dim, 1, 4), v_conv(value_dim, 1, 4)]
+        conv_key = f"decoder.layers.{gdn_idx}.mixer.conv1d.weight"
+        if conv_key in state:
+            conv_w = state[conv_key]  # (key_dim*2 + value_dim, 1, kernel_size)
+            q_conv = conv_w[:key_dim]
+            k_conv = conv_w[key_dim : key_dim * 2]
+            v_conv = conv_w[key_dim * 2 :]
+            hf_state[f"{prefix}.attn.q_conv1d.weight"] = q_conv
+            hf_state[f"{prefix}.attn.k_conv1d.weight"] = k_conv
+            hf_state[f"{prefix}.attn.v_conv1d.weight"] = v_conv
+
+        # Output norm (per-head RMSNorm)
+        hf_state[f"{prefix}.attn.o_norm.weight"] = state[f"decoder.layers.{gdn_idx}.mixer.out_norm.weight"]
+
+        # Output projection
+        hf_state[f"{prefix}.attn.o_proj.weight"] = state[f"decoder.layers.{gdn_idx}.mixer.out_proj.weight"]
+
+        # ── MLP sublayer ──
+        mlp_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{mlp_idx}.mlp.linear_fc1.layer_norm_weight",
+            f"decoder.layers.{mlp_idx}.pre_mlp_layernorm.weight",
+            f"decoder.layers.{mlp_idx}.input_layernorm.weight",
+        )
+        hf_state[f"{prefix}.mlp_norm.weight"] = mlp_norm_w
+
+        # Fused SwiGLU fc1 split: [gate_proj(intermediate), up_proj(intermediate)]
+        fc1_w = state[f"decoder.layers.{mlp_idx}.mlp.linear_fc1.weight"]
+        assert (
+            fc1_w.shape[0] == intermediate_size * 2
+        ), f"fc1 shape mismatch: {fc1_w.shape}, expected ({intermediate_size*2}, {hidden_size})"
+        gate_proj = fc1_w[:intermediate_size]
+        up_proj = fc1_w[intermediate_size:]
+        hf_state[f"{prefix}.mlp.gate_proj.weight"] = gate_proj
+        hf_state[f"{prefix}.mlp.up_proj.weight"] = up_proj
+
+        # Down projection
+        hf_state[f"{prefix}.mlp.down_proj.weight"] = state[f"decoder.layers.{mlp_idx}.mlp.linear_fc2.weight"]
+
+    final_norm_w, _ = _get_first(
+        state,
+        "decoder.final_norm.weight",
+        "decoder.final_layernorm.weight",
+        "decoder.norm.weight",
+    )
+    hf_state["model.norm.weight"] = final_norm_w
+
+    # LM head (tied with embeddings)
+    if "output_layer.weight" in state:
+        hf_state["lm_head.weight"] = state["output_layer.weight"]
+    else:
+        hf_state["lm_head.weight"] = state["embedding.word_embeddings.weight"]
+
+    return hf_state
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert Primus GDN to FLA HuggingFace format")
+    parser.add_argument("--checkpoint-path", type=str, required=True)
+    parser.add_argument("--output-dir", type=str, required=True)
+    parser.add_argument(
+        "--config", type=str, default=None, help="Path to FLA config JSON (gated_deltanet_1B_pure.json)"
+    )
+    args = parser.parse_args()
+
+    # Default config path — auto-detect model size from checkpoint
+    if args.config is None:
+        fla_root = os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention"))
+        fla_configs_dir = Path(fla_root) / "legacy" / "training" / "configs"
+        alt_dir = (
+            Path(__file__).parent.parent.parent
+            / "third_party"
+            / "flash-linear-attention"
+            / "legacy"
+            / "training"
+            / "configs"
+        )
+        configs_dir = fla_configs_dir if fla_configs_dir.exists() else alt_dir
+
+        # Detect from checkpoint path name
+        ckpt_str = str(args.checkpoint_path).lower()
+        if "300m" in ckpt_str:
+            args.config = str(configs_dir / "gated_deltanet_300M_pure.json")
+        else:
+            args.config = str(configs_dir / "gated_deltanet_1B_pure.json")
+
+    print("=" * 70)
+    print("Primus GDN → FLA HuggingFace Conversion")
+    print("=" * 70)
+
+    checkpoint = load_megatron_checkpoint(args.checkpoint_path)
+    hf_state = convert(checkpoint, args.config)
+
+    print(f"\nConverted {len(hf_state)} parameters")
+
+    # Save
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Saving to: {output_dir}")
+    torch.save(hf_state, output_dir / "model.safetensors.bin")
+
+    # Also save as safetensors if available
+    try:
+        from safetensors.torch import save_file
+
+        # Clone tied weights to avoid shared memory error
+        if "lm_head.weight" in hf_state and "model.embeddings.weight" in hf_state:
+            if hf_state["lm_head.weight"].data_ptr() == hf_state["model.embeddings.weight"].data_ptr():
+                hf_state["lm_head.weight"] = hf_state["lm_head.weight"].clone()
+        save_file(hf_state, str(output_dir / "model.safetensors"))
+        (output_dir / "model.safetensors.bin").unlink()
+        print("  Saved as safetensors format")
+    except ImportError:
+        os.rename(output_dir / "model.safetensors.bin", output_dir / "pytorch_model.bin")
+        print("  Saved as pytorch_model.bin (safetensors not available)")
+
+    # Save config.json (FLA format)
+    with open(args.config) as f:
+        config = json.load(f)
+    config["architectures"] = ["GatedDeltaNetForCausalLM"]
+    config["fuse_cross_entropy"] = False
+    config["fuse_norm"] = False
+    config["fuse_swiglu"] = False
+
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(config, f, indent=2)
+    print(f"  Saved config.json")
+
+    # Tokenizer config (use Llama-3.2-1B tokenizer)
+    tokenizer_config = {
+        "tokenizer_class": "PreTrainedTokenizerFast",
+        "model_max_length": 2048,
+    }
+    with open(output_dir / "tokenizer_config.json", "w") as f:
+        json.dump(tokenizer_config, f, indent=2)
+
+    print(f"\n{'='*70}")
+    print("Conversion complete!")
+    print(f"{'='*70}")
+    print(f"\nTo load:")
+    print(f"  from transformers import AutoModelForCausalLM")
+    print(f"  model = AutoModelForCausalLM.from_pretrained('{output_dir}', trust_remote_code=True)")
+    print(f"\nTo evaluate with lm-eval:")
+    print(f"  lm_eval --model hf \\")
+    print(f"    --model_args pretrained={output_dir},trust_remote_code=True \\")
+    print(f"    --tasks hellaswag,winogrande,piqa,arc_easy,arc_challenge \\")
+    print(f"    --batch_size 16")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_kda_to_fla_hf.py b/tools/hybrid/convert_kda_to_fla_hf.py
new file mode 100644
index 000000000..4ad236ac9
--- /dev/null
+++ b/tools/hybrid/convert_kda_to_fla_hf.py
@@ -0,0 +1,341 @@
+#!/usr/bin/env python3
+"""
+Convert Primus (Megatron) Pure KDA checkpoint → FLA HuggingFace format.
+
+This is the inverse of `tools/hybrid/convert_fla_kda_init_to_megatron.py`. It takes
+the trained `iter_NNNNNNN/mp_rank_00/model_optim_rng.pt` file Primus writes
+and emits a directory loadable by `transformers.AutoModelForCausalLM` via
+FLA's `KDAForCausalLM` (`trust_remote_code=True`).
+
+Primus's KDA layer fuses six FLA "hidden_states → X" projections into a
+single column-parallel `in_proj`:
+
+    Primus `decoder.layers.{kda}.mixer.in_proj.weight`
+        ←→  cat([q_proj, k_proj, v_proj, f_proj.0, g_proj.0, b_proj], dim=0)
+
+The converter splits this back apart and remaps the rest 1:1.
+
+Layer ordering: Primus's HybridStack interleaves
+    even  index → KDA mixer
+    odd   index → MLP
+so `fla_layer_idx` maps to `gdn_idx = 2*i` and `mlp_idx = 2*i+1`.
+
+Usage
+-----
+    python3 tools/hybrid/convert_kda_to_fla_hf.py \\
+        --checkpoint-path output/amd/root/zebra_llama_300M_kda_pure-pretrain/checkpoints/iter_0004768 \\
+        --output-dir output/kda_pure_300M_fla_hf \\
+        --config /home/<user>/flash-linear-attention/legacy/training/configs/kda_300M_pure.json
+
+Then evaluate with lm-eval:
+    lm_eval --model hf \\
+        --model_args pretrained=output/kda_pure_300M_fla_hf,trust_remote_code=True,dtype=bfloat16 \\
+        --tasks hellaswag,winogrande,piqa,arc_easy,arc_challenge \\
+        --batch_size 16
+"""
+
+import argparse
+import json
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+# Megatron must be on sys.path so torch.load can unpickle the ShardedTensor
+# wrappers Primus writes into the checkpoint.
+_megatron_path = str(Path(__file__).resolve().parents[2] / "third_party" / "Megatron-LM")
+if _megatron_path not in sys.path:
+    sys.path.insert(0, _megatron_path)
+
+
+def load_megatron_checkpoint(checkpoint_path: Path) -> dict:
+    model_path = Path(checkpoint_path) / "mp_rank_00" / "model_optim_rng.pt"
+    if not model_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {model_path}\n"
+            f"Expected layout: <checkpoint-path>/mp_rank_00/model_optim_rng.pt"
+        )
+    print(f"[load] {model_path}")
+    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+    print(f"[load] iteration={ckpt.get('iteration', '?')}")
+    return ckpt
+
+
+def _get_first(state, *candidates):
+    for k in candidates:
+        if k in state:
+            return state[k], k
+    raise KeyError(
+        f"None of these expected keys were found in checkpoint:\n  "
+        + "\n  ".join(candidates)
+        + "\nFirst 30 actually-present keys:\n  "
+        + "\n  ".join(sorted(state.keys())[:30])
+    )
+
+
+def convert(checkpoint: dict, fla_config_path: Path) -> OrderedDict:
+    """Map Primus KDA Megatron state_dict → FLA HF KDA state_dict."""
+    state = checkpoint["model"]
+
+    with open(fla_config_path) as f:
+        cfg = json.load(f)
+
+    hidden_size = cfg["hidden_size"]
+    num_heads = cfg["num_heads"]
+    num_v_heads = cfg.get("num_v_heads") or num_heads
+    head_dim = cfg["head_dim"]
+    expand_v = cfg.get("expand_v", 1.0)
+    intermediate_size = cfg["intermediate_size"]
+    num_hidden_layers = cfg["num_hidden_layers"]
+
+    head_v_dim = int(head_dim * expand_v)
+    qk_dim = num_heads * head_dim  # 256 for 300M
+    v_dim = num_v_heads * head_v_dim  # 512 for 300M
+    fused_in_proj_dim = (
+        qk_dim * 2  # q + k
+        + v_dim  # v
+        + head_v_dim  # f_a (low-rank gate bottleneck)
+        + head_v_dim  # g_a (low-rank output-gate bottleneck)
+        + num_v_heads  # beta
+    )
+    print(
+        f"[cfg ] hidden={hidden_size} num_heads={num_heads} num_v_heads={num_v_heads}\n"
+        f"       head_dim={head_dim} expand_v={expand_v} head_v_dim={head_v_dim}\n"
+        f"       qk_dim={qk_dim} v_dim={v_dim} intermediate={intermediate_size}\n"
+        f"       layers={num_hidden_layers} fused_in_proj_dim={fused_in_proj_dim}"
+    )
+
+    hf = OrderedDict()
+    hf["model.embeddings.weight"] = state["embedding.word_embeddings.weight"]
+
+    for fla_idx in range(num_hidden_layers):
+        kda_i = 2 * fla_idx
+        mlp_i = 2 * fla_idx + 1
+        dst = f"model.layers.{fla_idx}"
+
+        # ── attn pre-norm ─────────────────────────────────────────────
+        attn_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{kda_i}.mixer.in_proj.layer_norm_weight",  # TE-fused spec
+            f"decoder.layers.{kda_i}.norm.weight",  # no-TE spec
+            f"decoder.layers.{kda_i}.input_layernorm.weight",
+        )
+        hf[f"{dst}.attn_norm.weight"] = attn_norm_w
+
+        # ── fused in_proj split: [q | k | v | f_a | g_a | beta] ───────
+        in_proj_w = state[f"decoder.layers.{kda_i}.mixer.in_proj.weight"]
+        assert in_proj_w.shape == (fused_in_proj_dim, hidden_size), (
+            f"in_proj shape {tuple(in_proj_w.shape)} != "
+            f"expected ({fused_in_proj_dim}, {hidden_size}) for layer {kda_i}. "
+            "Did you train with the post-fusion KDA code?"
+        )
+        o = 0
+        q_w = in_proj_w[o : o + qk_dim]
+        o += qk_dim
+        k_w = in_proj_w[o : o + qk_dim]
+        o += qk_dim
+        v_w = in_proj_w[o : o + v_dim]
+        o += v_dim
+        f_a_w = in_proj_w[o : o + head_v_dim]
+        o += head_v_dim
+        g_a_w = in_proj_w[o : o + head_v_dim]
+        o += head_v_dim
+        b_w = in_proj_w[o : o + num_v_heads]
+        o += num_v_heads
+        assert o == fused_in_proj_dim, (o, fused_in_proj_dim)
+
+        hf[f"{dst}.attn.q_proj.weight"] = q_w
+        hf[f"{dst}.attn.k_proj.weight"] = k_w
+        hf[f"{dst}.attn.v_proj.weight"] = v_w
+        hf[f"{dst}.attn.b_proj.weight"] = b_w
+        hf[f"{dst}.attn.f_proj.0.weight"] = f_a_w
+        hf[f"{dst}.attn.g_proj.0.weight"] = g_a_w
+
+        # ── low-rank bottleneck expansions (still separate in Primus) ─
+        hf[f"{dst}.attn.f_proj.1.weight"] = state[f"decoder.layers.{kda_i}.mixer.f_b_proj.weight"]
+        hf[f"{dst}.attn.g_proj.1.weight"] = state[f"decoder.layers.{kda_i}.mixer.g_b_proj.weight"]
+        # g_proj.1 has bias=True (FLA reference: fla/layers/kda.py:189)
+        hf[f"{dst}.attn.g_proj.1.bias"] = state[f"decoder.layers.{kda_i}.mixer.g_b_proj.bias"]
+
+        # ── fused conv1d split: [q_conv | k_conv | v_conv] ─────────────
+        conv_w = state[f"decoder.layers.{kda_i}.mixer.conv1d.weight"]
+        # FLA stores each conv as [channels, 1, kernel]
+        q_conv = conv_w[:qk_dim]
+        k_conv = conv_w[qk_dim : qk_dim * 2]
+        v_conv = conv_w[qk_dim * 2 :]
+        hf[f"{dst}.attn.q_conv1d.weight"] = q_conv
+        hf[f"{dst}.attn.k_conv1d.weight"] = k_conv
+        hf[f"{dst}.attn.v_conv1d.weight"] = v_conv
+
+        # ── A_log / dt_bias ───────────────────────────────────────────
+        # Primus stores A_log as [1, 1, num_v_heads, 1]; FLA wants flat [num_v_heads].
+        A_log = state[f"decoder.layers.{kda_i}.mixer.A_log"].reshape(num_v_heads)
+        hf[f"{dst}.attn.A_log"] = A_log
+        hf[f"{dst}.attn.dt_bias"] = state[f"decoder.layers.{kda_i}.mixer.dt_bias"]
+
+        # ── output norm + projection ──────────────────────────────────
+        hf[f"{dst}.attn.o_norm.weight"] = state[f"decoder.layers.{kda_i}.mixer.out_norm.weight"]
+        hf[f"{dst}.attn.o_proj.weight"] = state[f"decoder.layers.{kda_i}.mixer.out_proj.weight"]
+
+        # ── MLP sublayer ──────────────────────────────────────────────
+        mlp_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{mlp_i}.mlp.linear_fc1.layer_norm_weight",  # TE spec
+            f"decoder.layers.{mlp_i}.pre_mlp_layernorm.weight",  # no-TE spec
+            f"decoder.layers.{mlp_i}.input_layernorm.weight",
+        )
+        hf[f"{dst}.mlp_norm.weight"] = mlp_norm_w
+
+        # SwiGLU fc1 = cat([gate_proj, up_proj])
+        fc1_w = state[f"decoder.layers.{mlp_i}.mlp.linear_fc1.weight"]
+        assert fc1_w.shape == (intermediate_size * 2, hidden_size), (
+            f"fc1 shape {tuple(fc1_w.shape)} != "
+            f"expected ({intermediate_size * 2}, {hidden_size}) for layer {mlp_i}"
+        )
+        hf[f"{dst}.mlp.gate_proj.weight"] = fc1_w[:intermediate_size]
+        hf[f"{dst}.mlp.up_proj.weight"] = fc1_w[intermediate_size:]
+        hf[f"{dst}.mlp.down_proj.weight"] = state[f"decoder.layers.{mlp_i}.mlp.linear_fc2.weight"]
+
+    # ── final norm + LM head ─────────────────────────────────────────
+    final_norm_w, _ = _get_first(
+        state,
+        "decoder.final_norm.weight",
+        "decoder.final_layernorm.weight",
+        "decoder.norm.weight",
+    )
+    hf["model.norm.weight"] = final_norm_w
+
+    if "output_layer.weight" in state:
+        hf["lm_head.weight"] = state["output_layer.weight"]
+    else:
+        # Llama-style tied embeddings
+        hf["lm_head.weight"] = state["embedding.word_embeddings.weight"]
+
+    return hf
+
+
+def save_hf_dir(hf_state: OrderedDict, output_dir: Path, fla_config_path: Path) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Save weights — prefer safetensors, fall back to pytorch_model.bin.
+    try:
+        from safetensors.torch import save_file
+
+        # Clone tied weights so they don't share storage (safetensors rejects that).
+        if "lm_head.weight" in hf_state and "model.embeddings.weight" in hf_state:
+            if hf_state["lm_head.weight"].data_ptr() == hf_state["model.embeddings.weight"].data_ptr():
+                hf_state["lm_head.weight"] = hf_state["lm_head.weight"].clone()
+        save_file(hf_state, str(output_dir / "model.safetensors"))
+        print(f"[save] model.safetensors ({(output_dir / 'model.safetensors').stat().st_size / 1e6:.1f} MB)")
+    except ImportError:
+        torch.save(hf_state, output_dir / "pytorch_model.bin")
+        print(f"[save] pytorch_model.bin (safetensors not installed)")
+
+    # config.json — copy FLA's pretrain config, force HF/eval-friendly toggles.
+    with open(fla_config_path) as f:
+        cfg = json.load(f)
+    cfg["architectures"] = ["KDAForCausalLM"]
+    cfg["model_type"] = cfg.get("model_type", "kda")
+    cfg["torch_dtype"] = "bfloat16"
+    # Disable fused-loss / fused-norm / fused-swiglu auto-detection at eval
+    # time — FLA's HF model has these as optional fast paths that aren't
+    # needed for evaluation and avoid extra triton compiles.
+    cfg["fuse_cross_entropy"] = False
+    cfg["fuse_norm"] = False
+    cfg["fuse_swiglu"] = False
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(cfg, f, indent=2)
+    print(f"[save] config.json (architectures={cfg['architectures']})")
+
+    # generation_config.json (so lm-eval can use HF .generate cleanly)
+    gen_cfg = {
+        "_from_model_config": True,
+        "transformers_version": None,
+        "pad_token_id": cfg.get("pad_token_id"),
+        "eos_token_id": cfg.get("eos_token_id"),
+        "bos_token_id": cfg.get("bos_token_id"),
+    }
+    gen_cfg = {k: v for k, v in gen_cfg.items() if v is not None}
+    with open(output_dir / "generation_config.json", "w") as f:
+        json.dump(gen_cfg, f, indent=2)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--checkpoint-path",
+        required=True,
+        type=Path,
+        help="Primus iter dir, e.g. .../checkpoints/iter_0004768",
+    )
+    p.add_argument("--output-dir", required=True, type=Path)
+    p.add_argument(
+        "--config",
+        type=Path,
+        default=None,
+        help="Path to FLA KDA config JSON. Defaults to kda_300M_pure.json "
+        "when '300m' appears in --checkpoint-path, else kda_1B_pure.json.",
+    )
+    p.add_argument(
+        "--tokenizer-src", type=Path, default=None, help="Optional tokenizer dir to copy into --output-dir."
+    )
+    args = p.parse_args()
+
+    if args.config is None:
+        fla_root = os.environ.get("FLA_ROOT", os.path.expanduser("~/flash-linear-attention"))
+        configs_dir = Path(fla_root) / "legacy" / "training" / "configs"
+        if "300m" in str(args.checkpoint_path).lower():
+            args.config = configs_dir / "kda_300M_pure.json"
+        else:
+            args.config = configs_dir / "kda_1B_pure.json"
+        print(f"[auto] --config defaulted to {args.config}")
+
+    print("=" * 78)
+    print(" Primus KDA → FLA HuggingFace converter")
+    print("=" * 78)
+    print(f"  checkpoint = {args.checkpoint_path}")
+    print(f"  output     = {args.output_dir}")
+    print(f"  fla config = {args.config}")
+    print()
+
+    ckpt = load_megatron_checkpoint(args.checkpoint_path)
+    hf_state = convert(ckpt, args.config)
+    print(f"[map ] converted {len(hf_state)} tensors")
+
+    save_hf_dir(hf_state, args.output_dir, args.config)
+
+    # Optionally copy tokenizer files
+    if args.tokenizer_src is not None:
+        import shutil
+
+        copied = 0
+        for name in (
+            "tokenizer.json",
+            "tokenizer.model",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "vocab.json",
+            "merges.txt",
+        ):
+            src = args.tokenizer_src / name
+            if src.exists():
+                shutil.copy(src, args.output_dir / name)
+                copied += 1
+        print(f"[tok ] copied {copied} tokenizer files from {args.tokenizer_src}")
+
+    print()
+    print("Done. To load:")
+    print(f"  from transformers import AutoModelForCausalLM")
+    print(f"  model = AutoModelForCausalLM.from_pretrained('{args.output_dir}', trust_remote_code=True)")
+    print()
+    print("To evaluate with lm-eval:")
+    print(f"  lm_eval --model hf \\")
+    print(f"    --model_args pretrained={args.output_dir},trust_remote_code=True,dtype=bfloat16 \\")
+    print(f"    --tasks hellaswag,winogrande,piqa,arc_easy,arc_challenge \\")
+    print(f"    --batch_size 16")
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/hybrid/convert_mamba_hybrid_to_fla_hf.py b/tools/hybrid/convert_mamba_hybrid_to_fla_hf.py
new file mode 100644
index 000000000..a10f987d2
--- /dev/null
+++ b/tools/hybrid/convert_mamba_hybrid_to_fla_hf.py
@@ -0,0 +1,566 @@
+#!/usr/bin/env python3
+"""Convert Primus (Megatron) 75% Hybrid Mamba2+MLA checkpoint to FLA HuggingFace format.
+
+The hybrid model has 12 FLA mixer blocks: MLA at indices [0, 4, 8] and Mamba2
+at the other 9 indices.
+
+Primus Megatron stores them as 24 sublayers alternating mixer/MLP, with
+hybrid_override_pattern = "*-M-M-M-*-M-M-M-*-M-M-M-":
+
+    sublayer   pattern  FLA mixer  what
+    --------   -------  ---------  ----
+    0          *        0          MLA
+    1          -                   MLP
+    2          M        1          Mamba2
+    3          -                   MLP
+    4          M        2          Mamba2
+    5          -                   MLP
+    6          M        3          Mamba2
+    7          -                   MLP
+    8          *        4          MLA
+    9          -                   MLP
+    ...                            (repeats with `*-M-M-M-` pattern)
+
+This converter walks the 12 FLA mixer indices, finds the corresponding Megatron
+sublayer (mixer + MLP), and emits weights for FLA's `Mamba2ForCausalLM`
+(which uses the `backbone.` prefix, NOT `model.` like `GatedDeltaNetForCausalLM`).
+
+The MLA mapping reuses the same channel-permutation fix from the GDN converter
+(see `_megatron_to_fla_rope_channels`).  The Mamba2 mixer mapping is a direct
+copy because both Megatron's MambaMixer and FLA's Mamba2Mixer wrap upstream
+`mamba_ssm` with the same `[z | x | B | C | dt]` `in_proj` layout.
+
+NOTE on dim mismatch with FLA's reference 300M:
+
+    Primus YAML     FLA mamba2_300M_hybrid.json
+    -----------     ---------------------------
+    hidden_size     1024              1216
+    intermediate    4096              4864
+    state_size      64                128
+    n_groups        8                 1
+    n_heads(Mamba)  32                38
+    total params    ~273M             ~350M
+
+We deliberately matched Primus to the GDN-hybrid 300M dims so the architecture
+is a drop-in mixer swap (Mamba2 ↔ GDN under the same MLA + MLP backbone).
+The emitted HF config.json therefore reflects PRIMUS's dims, not FLA's.
+Both models are independently valid HF Mamba2ForCausalLM checkpoints; the
+downstream lm-eval can score each on its own merits.
+
+Usage (inside the container):
+
+    python tools/hybrid/convert_mamba_hybrid_to_fla_hf.py \
+        --checkpoint-path output/amd/root/zebra_llama_300M_mamba_hybrid-pretrain/checkpoints/iter_0004768 \
+        --output-dir output/mamba_hybrid_300M_fla_hf \
+        --tokenizer /home/<user>/checkpoints/gdn_pure_300M_10B
+"""
+import argparse
+import json
+import os
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+_megatron_path = str(Path(__file__).resolve().parents[2] / "third_party" / "Megatron-LM")
+if _megatron_path not in sys.path:
+    sys.path.insert(0, _megatron_path)
+
+
+def load_megatron_checkpoint(checkpoint_path):
+    model_path = Path(checkpoint_path) / "mp_rank_00" / "model_optim_rng.pt"
+    if not model_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+    print(f"Loading checkpoint from: {model_path}")
+    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+    print(f"Loaded. Iteration: {checkpoint.get('iteration', '?')}")
+    return checkpoint
+
+
+def _get_first(state, *candidates):
+    for k in candidates:
+        if k in state:
+            return state[k], k
+    raise KeyError(
+        f"None of these expected keys were found:\n  "
+        + "\n  ".join(candidates)
+        + "\nFirst 30 keys in checkpoint:\n  "
+        + "\n  ".join(sorted(state.keys())[:30])
+    )
+
+
+def _build_layer_map(hybrid_pattern):
+    """Same as the GDN converter — `*` = MLA, `M` = Mamba2, `-` = MLP."""
+    fla_to_megatron = []
+    cur_mixer = None
+    cur_kind = None
+    for i, c in enumerate(hybrid_pattern):
+        if c in ("*", "M"):
+            cur_mixer = i
+            cur_kind = "mla" if c == "*" else "mamba2"
+        elif c == "-":
+            if cur_mixer is None:
+                continue
+            fla_to_megatron.append((cur_mixer, i, cur_kind))
+            cur_mixer = None
+            cur_kind = None
+    return fla_to_megatron
+
+
+def _megatron_to_fla_rope_channels(w_rope, qk_rope_head_dim):
+    """Permute the rope channels of a tensor whose LAST dim is `qk_rope_head_dim`
+    from Megatron's storage order to FLA's RoPE-ready order.
+
+    Megatron stores [c0, c1, c2, ..., c_{d-1}] and on-the-fly interleaves to
+    [c0, c2, ..., c1, c3, ...]. FLA assumes already-permuted weights.
+    """
+    d = qk_rope_head_dim
+    assert w_rope.shape[-1] == d, (
+        f"_megatron_to_fla_rope_channels: last dim {w_rope.shape[-1]} != " f"qk_rope_head_dim {d}"
+    )
+    even = w_rope[..., 0::2]
+    odd = w_rope[..., 1::2]
+    return torch.cat([even, odd], dim=-1).contiguous()
+
+
+def convert_mla_block(state, sub_mixer, prefix, fla_cfg):
+    """Identical to the GDN-hybrid MLA conversion (since the MLA spec is shared).
+    Writes:
+        {prefix}.mixer.q_proj.0/1/2.weight
+        {prefix}.mixer.kv_proj.0/1/2.weight
+        {prefix}.mixer.k_rope.weight
+        {prefix}.mixer.o_proj.weight
+    (Note: FLA's Mamba2 model uses `mixer.attn` ... actually it uses `mixer`
+    for both attention and SSM blocks — the type is dispatched by config.attn.layers.)
+    """
+    out = OrderedDict()
+    attn_cfg = fla_cfg["attn"]
+    q_lora_rank = attn_cfg["q_lora_rank"]
+    kv_lora_rank = attn_cfg["kv_lora_rank"]
+    qk_rope_head_dim = attn_cfg["qk_rope_head_dim"]
+    num_heads = attn_cfg["num_heads"]
+    qk_nope_head_dim = attn_cfg["qk_nope_head_dim"]
+    v_head_dim = attn_cfg["v_head_dim"]
+    qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+    hidden_size = fla_cfg["hidden_size"]
+
+    # ── q path ──
+    q_down_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_q_down_proj.weight"]
+    q_up_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_q_up_proj.weight"]
+    q_norm_w, _ = _get_first(
+        state,
+        f"decoder.layers.{sub_mixer}.self_attention.q_layernorm.weight",
+    )
+    assert q_down_w.shape == (q_lora_rank, hidden_size)
+    assert q_up_w.shape == (num_heads * qk_head_dim, q_lora_rank)
+    assert q_norm_w.shape == (q_lora_rank,)
+
+    # Permute rope half of q_up_proj per head.
+    q_up_3d = q_up_w.view(num_heads, qk_head_dim, q_lora_rank)
+    q_up_nope = q_up_3d[:, :qk_nope_head_dim, :]
+    q_up_rope_meg = q_up_3d[:, qk_nope_head_dim:, :].transpose(-1, -2)
+    q_up_rope_fla = _megatron_to_fla_rope_channels(q_up_rope_meg, qk_rope_head_dim).transpose(-1, -2)
+    q_up_fla = torch.cat([q_up_nope, q_up_rope_fla], dim=1)
+    q_up_fla = q_up_fla.reshape(num_heads * qk_head_dim, q_lora_rank).contiguous()
+
+    out[f"{prefix}.q_proj.0.weight"] = q_down_w
+    out[f"{prefix}.q_proj.1.weight"] = q_norm_w
+    out[f"{prefix}.q_proj.2.weight"] = q_up_fla
+
+    # ── kv path ──
+    kv_down_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_kv_down_proj.weight"]
+    expected_down_rows = kv_lora_rank + qk_rope_head_dim
+    assert kv_down_w.shape == (expected_down_rows, hidden_size)
+    out[f"{prefix}.kv_proj.0.weight"] = kv_down_w[:kv_lora_rank].contiguous()
+
+    k_rope_meg = kv_down_w[kv_lora_rank:].contiguous().transpose(0, 1)
+    k_rope_fla = _megatron_to_fla_rope_channels(k_rope_meg, qk_rope_head_dim).transpose(0, 1).contiguous()
+    out[f"{prefix}.k_rope.weight"] = k_rope_fla
+
+    kv_norm_w, _ = _get_first(
+        state,
+        f"decoder.layers.{sub_mixer}.self_attention.kv_layernorm.weight",
+    )
+    assert kv_norm_w.shape == (kv_lora_rank,)
+    out[f"{prefix}.kv_proj.1.weight"] = kv_norm_w
+
+    kv_up_w = state[f"decoder.layers.{sub_mixer}.self_attention.linear_kv_up_proj.weight"]
+    expected_up_rows = num_heads * (qk_nope_head_dim + v_head_dim)
+    assert kv_up_w.shape == (expected_up_rows, kv_lora_rank)
+    out[f"{prefix}.kv_proj.2.weight"] = kv_up_w
+
+    out[f"{prefix}.o_proj.weight"] = state[f"decoder.layers.{sub_mixer}.self_attention.linear_proj.weight"]
+    return out
+
+
+def convert_mamba2_block(state, sub_mixer, prefix, fla_cfg):
+    """Mamba2 mixer Primus→FLA mapping.
+
+    Both Megatron and FLA wrap upstream mamba_ssm's MambaMixer with the SAME
+    `[z | x | B | C | dt]` in_proj layout, so this is a direct copy of:
+
+        in_proj.weight     (hidden_intermediate*2 + 2*n_groups*state + n_heads, hidden)
+        conv1d.weight      (hidden_intermediate + 2*n_groups*state, 1, kernel)
+        conv1d.bias        (hidden_intermediate + 2*n_groups*state,)
+        A_log              (n_heads,)
+        D                  (n_heads,)
+        dt_bias            (n_heads,)
+        norm.weight        (hidden_intermediate,)         ← post-mixer gated-RMSNorm
+        out_proj.weight    (hidden, hidden_intermediate)
+
+    Validated against an emitted Primus ckpt with hidden=1024, expand=2,
+    n_heads=32, n_groups=8, state=64 → in_proj (5152, 1024), conv1d (3072,1,4).
+    """
+    out = OrderedDict()
+    hidden_size = fla_cfg["hidden_size"]
+    expand = fla_cfg["expand"]
+    fla_cfg["head_dim"]
+    n_groups = fla_cfg["n_groups"]
+    state_size = fla_cfg["state_size"]
+    num_heads = fla_cfg["num_heads"]
+
+    intermediate = expand * hidden_size
+    conv_dim = intermediate + 2 * n_groups * state_size
+    in_dim = 2 * intermediate + 2 * n_groups * state_size + num_heads
+
+    in_proj_w = state[f"decoder.layers.{sub_mixer}.mixer.in_proj.weight"]
+    assert in_proj_w.shape == (in_dim, hidden_size), (
+        f"Mamba2 in_proj shape {tuple(in_proj_w.shape)} != "
+        f"expected ({in_dim}, {hidden_size}); "
+        f"check expand={expand}, n_groups={n_groups}, state={state_size}, n_heads={num_heads}"
+    )
+    out[f"{prefix}.in_proj.weight"] = in_proj_w
+
+    conv1d_w = state[f"decoder.layers.{sub_mixer}.mixer.conv1d.weight"]
+    conv1d_b = state[f"decoder.layers.{sub_mixer}.mixer.conv1d.bias"]
+    assert (
+        conv1d_w.shape[0] == conv_dim
+    ), f"Mamba2 conv1d weight rows {conv1d_w.shape[0]} != expected {conv_dim}"
+    out[f"{prefix}.conv1d.weight"] = conv1d_w
+    out[f"{prefix}.conv1d.bias"] = conv1d_b
+
+    out[f"{prefix}.A_log"] = state[f"decoder.layers.{sub_mixer}.mixer.A_log"]
+    out[f"{prefix}.D"] = state[f"decoder.layers.{sub_mixer}.mixer.D"]
+    out[f"{prefix}.dt_bias"] = state[f"decoder.layers.{sub_mixer}.mixer.dt_bias"]
+
+    norm_w = state[f"decoder.layers.{sub_mixer}.mixer.norm.weight"]
+    assert norm_w.shape == (intermediate,)
+    out[f"{prefix}.norm.weight"] = norm_w
+
+    out[f"{prefix}.out_proj.weight"] = state[f"decoder.layers.{sub_mixer}.mixer.out_proj.weight"]
+    return out
+
+
+def build_fla_config(primus_args, base_attn_cfg):
+    """Build an HF Mamba2 config that matches the PRIMUS-trained dims.
+
+    `primus_args` is the `args` namespace stashed in the Megatron checkpoint,
+    or our YAML-derived overrides; we extract the actual dims so the emitted
+    HF model loads without shape mismatches.
+    """
+    hidden_size = primus_args["hidden_size"]
+    intermediate_size = primus_args["ffn_hidden_size"]
+    expand = primus_args.get("mamba_expand", 2)
+    head_dim = primus_args.get("mamba_head_dim", 64)
+    n_groups = primus_args.get("mamba_num_groups", 8)
+    state_size = primus_args.get("mamba_state_dim", 64)
+    num_heads = (expand * hidden_size) // head_dim
+
+    cfg = {
+        "model_type": "mamba2",
+        "architectures": ["Mamba2ForCausalLM"],
+        "vocab_size": primus_args.get("padded_vocab_size", 128256),
+        "hidden_size": hidden_size,
+        "intermediate_size": intermediate_size,
+        "state_size": state_size,
+        "num_hidden_layers": 12,
+        "head_dim": head_dim,
+        "expand": expand,
+        "n_groups": n_groups,
+        "num_heads": num_heads,
+        "conv_kernel": 4,
+        "use_bias": False,
+        "use_conv_bias": True,
+        "hidden_act": "silu",
+        "hidden_act_mlp": "swish",
+        "initializer_range": 0.02,
+        # NOTE: set to False (FLA default is True) — at bf16 inference the
+        # mixer's `in_proj` is bf16 but with residual_in_fp32=True the
+        # incoming residual is fp32, causing an `F.linear` dtype mismatch.
+        # Training was bf16 + residual_in_fp32=True (FLA's training default)
+        # which only matters during gradient accumulation; for eval-time
+        # forward the bf16 residual is mathematically equivalent up to
+        # rounding noise.
+        "residual_in_fp32": False,
+        "rmsnorm": True,
+        "norm_eps": primus_args.get("norm_epsilon", 1e-6),
+        "chunk_size": 256,
+        "D_has_hdim": False,
+        "norm_before_gate": False,
+        "rescale_prenorm_residual": True,
+        "max_position_embeddings": primus_args.get("max_position_embeddings", 131072),
+        "hidden_ratio": None,
+        # Disable fused kernels in HF (they require packages we may not have)
+        "fuse_norm": False,
+        "fuse_swiglu": False,
+        "fuse_cross_entropy": False,
+        "fuse_linear_cross_entropy": False,
+        "tie_word_embeddings": True,
+        "use_cache": True,
+        # MLA attention sub-block — verbatim from base_attn_cfg
+        "attn": base_attn_cfg,
+        # Mamba2 init knobs (HF defaults; not used at eval time)
+        "A_init_range": [1, 16],
+        "dt_init_floor": 0.0001,
+        "dt_limit": [0.0, float("inf")],
+        "dt_max": 0.1,
+        "dt_min": 0.001,
+        "conv_init": None,
+        "bos_token_id": 128000,
+        "eos_token_id": 128001,
+        "pad_token_id": 0,
+        "torch_dtype": "bfloat16",
+    }
+    return cfg
+
+
+def convert(checkpoint, hybrid_pattern, base_attn_cfg, primus_args_override=None):
+    state = checkpoint["model"]
+
+    # Pull Primus's actual dims out of the checkpoint so config matches weights.
+    ckpt_args = checkpoint.get("args", None)
+    primus_args = {}
+    if ckpt_args is not None:
+        for k in (
+            "hidden_size",
+            "ffn_hidden_size",
+            "padded_vocab_size",
+            "mamba_state_dim",
+            "mamba_head_dim",
+            "mamba_num_groups",
+            "mamba_expand",
+            "max_position_embeddings",
+            "norm_epsilon",
+        ):
+            v = getattr(ckpt_args, k, None)
+            if v is not None:
+                primus_args[k] = v
+    if primus_args_override:
+        primus_args.update(primus_args_override)
+
+    fla_cfg = build_fla_config(primus_args, base_attn_cfg)
+    num_hidden_layers = fla_cfg["num_hidden_layers"]
+
+    layer_map = _build_layer_map(hybrid_pattern)
+    assert len(layer_map) == num_hidden_layers, (
+        f"hybrid_pattern produced {len(layer_map)} FLA mixer blocks "
+        f"but FLA config wants {num_hidden_layers}"
+    )
+
+    print(f"\nFLA layer index → Megatron sublayer mapping:")
+    for fla_idx, (sub_mixer, sub_mlp, kind) in enumerate(layer_map):
+        print(f"  FLA layer {fla_idx} ({kind.upper():>6}): mixer=sub{sub_mixer}, mlp=sub{sub_mlp}")
+
+    hf_state = OrderedDict()
+    hf_state["backbone.embeddings.weight"] = state["embedding.word_embeddings.weight"]
+
+    hidden_size = fla_cfg["hidden_size"]
+    n_ones_mixer_norm = 0
+    for fla_idx, (sub_mixer, sub_mlp, kind) in enumerate(layer_map):
+        prefix = f"backbone.layers.{fla_idx}"
+
+        # Pre-mixer norm (FLA: `mixer_norm.weight`).
+        #
+        # ⚠ For Mamba2 mixer sublayers, Primus's spec used `MambaLayerSubmodules(norm=IdentityOp)`
+        # (the default) — there was NO learnable pre-mixer norm at training time.
+        # FLA's `Mamba2Block` ALWAYS applies a `mixer_norm` RMSNorm before the mixer.
+        # We emit `ones(hidden_size)` so the FLA HF model loads, but this introduces
+        # a "spurious" RMSNorm at inference time that wasn't present during training
+        # (division by RMS, scaled by 1.0).  Empirically this is a small perturbation
+        # — the eval scores should still be representative — but it is NOT bit-exact.
+        # For MLA mixer sublayers we have `input_layernorm.weight` and copy it.
+        try:
+            mixer_norm_w, src_key = _get_first(
+                state,
+                f"decoder.layers.{sub_mixer}.input_layernorm.weight",
+                f"decoder.layers.{sub_mixer}.mixer.in_proj.layer_norm_weight",
+                f"decoder.layers.{sub_mixer}.norm.weight",
+            )
+        except KeyError:
+            assert kind == "mamba2", (
+                f"FLA layer {fla_idx} ({kind}): missing pre-mixer norm — only "
+                f"Mamba2 layers are allowed to lack one (MLA always has input_layernorm)."
+            )
+            mixer_norm_w = torch.ones(hidden_size, dtype=torch.bfloat16)
+            n_ones_mixer_norm += 1
+        hf_state[f"{prefix}.mixer_norm.weight"] = mixer_norm_w
+
+        if kind == "mla":
+            hf_state.update(convert_mla_block(state, sub_mixer, f"{prefix}.mixer", fla_cfg))
+        else:
+            hf_state.update(convert_mamba2_block(state, sub_mixer, f"{prefix}.mixer", fla_cfg))
+
+        # Pre-MLP norm (FLA: `mlp_norm.weight`)
+        mlp_norm_w, _ = _get_first(
+            state,
+            f"decoder.layers.{sub_mlp}.pre_mlp_layernorm.weight",
+            f"decoder.layers.{sub_mlp}.mlp.linear_fc1.layer_norm_weight",
+            f"decoder.layers.{sub_mlp}.input_layernorm.weight",
+        )
+        hf_state[f"{prefix}.mlp_norm.weight"] = mlp_norm_w
+
+        fc1_w = state[f"decoder.layers.{sub_mlp}.mlp.linear_fc1.weight"]
+        intermediate_size = fla_cfg["intermediate_size"]
+        assert (
+            fc1_w.shape[0] == intermediate_size * 2
+        ), f"MLP fc1 shape {tuple(fc1_w.shape)} expected ({intermediate_size*2},...)"
+        hf_state[f"{prefix}.mlp.gate_proj.weight"] = fc1_w[:intermediate_size]
+        hf_state[f"{prefix}.mlp.up_proj.weight"] = fc1_w[intermediate_size:]
+        hf_state[f"{prefix}.mlp.down_proj.weight"] = state[f"decoder.layers.{sub_mlp}.mlp.linear_fc2.weight"]
+
+    final_norm_w, _ = _get_first(
+        state,
+        "decoder.final_norm.weight",
+        "decoder.final_layernorm.weight",
+        "decoder.norm.weight",
+    )
+    hf_state["backbone.norm_f.weight"] = final_norm_w
+
+    if "output_layer.weight" in state:
+        hf_state["lm_head.weight"] = state["output_layer.weight"]
+    else:
+        hf_state["lm_head.weight"] = state["embedding.word_embeddings.weight"]
+
+    if n_ones_mixer_norm:
+        print(
+            f"\n⚠ Emitted ones(hidden_size) for {n_ones_mixer_norm}/9 Mamba2 mixer_norm.weight "
+            f"entries (Primus spec had no pre-mixer norm). FLA inference will apply an "
+            f"unweighted RMSNorm here that the Primus model never saw during training."
+        )
+
+    return hf_state, fla_cfg
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint-path", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument(
+        "--hybrid-pattern",
+        default="*-M-M-M-*-M-M-M-*-M-M-M-",
+        help="Same as YAML hybrid_override_pattern (24 chars for 300M hybrid).",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        default=os.path.expanduser("~/checkpoints/gdn_pure_300M_10B"),
+        help="Source dir to copy tokenizer.json / tokenizer_config.json from.",
+    )
+    args = parser.parse_args()
+
+    print("=" * 72)
+    print("Primus 75% Hybrid Mamba2+MLA → FLA HuggingFace Conversion")
+    print("=" * 72)
+
+    # MLA sub-config — verbatim layout (matches both our spec and FLA's mla.py).
+    base_attn_cfg = {
+        "type": "mla",
+        "layers": [0, 4, 8],
+        "num_heads": 16,
+        "q_lora_rank": 672,
+        "kv_lora_rank": 64,
+        "qk_nope_head_dim": 32,
+        "qk_rope_head_dim": 32,
+        "qk_head_dim": 64,
+        "v_head_dim": 64,
+        "rope_theta": 500000.0,
+    }
+
+    checkpoint = load_megatron_checkpoint(args.checkpoint_path)
+    hf_state, fla_cfg = convert(checkpoint, args.hybrid_pattern, base_attn_cfg)
+    print(f"\nConverted {len(hf_state)} parameters")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        from safetensors.torch import save_file
+
+        # Untie lm_head <-> embeddings (safetensors disallows shared tensors).
+        if (
+            "lm_head.weight" in hf_state
+            and "backbone.embeddings.weight" in hf_state
+            and hf_state["lm_head.weight"].data_ptr() == hf_state["backbone.embeddings.weight"].data_ptr()
+        ):
+            hf_state["lm_head.weight"] = hf_state["lm_head.weight"].clone()
+        save_file(hf_state, str(output_dir / "model.safetensors"))
+        print(f"  Saved {output_dir / 'model.safetensors'}")
+    except ImportError:
+        torch.save(hf_state, output_dir / "pytorch_model.bin")
+        print(f"  Saved {output_dir / 'pytorch_model.bin'}")
+
+    # Tell HF AutoModel to load via our custom Mamba2FullMlpForCausalLM class
+    # (saved as modeling_mamba2_full_mlp.py next to the checkpoint).  Primus has
+    # MLPs on EVERY layer (both MLA and Mamba2 sublayers) — the stock FLA
+    # Mamba2Block only builds MLPs for MLA layers and would silently drop the
+    # other 9 MLPs from the safetensors at load time.
+    fla_cfg["architectures"] = ["Mamba2FullMlpForCausalLM"]
+    fla_cfg["auto_map"] = {
+        "AutoConfig": "modeling_mamba2_full_mlp.Mamba2Config",
+        "AutoModelForCausalLM": "modeling_mamba2_full_mlp.Mamba2FullMlpForCausalLM",
+    }
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(fla_cfg, f, indent=2)
+    print(f"  Saved config.json")
+    print(f"  hidden_size={fla_cfg['hidden_size']}, intermediate_size={fla_cfg['intermediate_size']}")
+    print(
+        f"  n_groups={fla_cfg['n_groups']}, state_size={fla_cfg['state_size']}, num_heads={fla_cfg['num_heads']}"
+    )
+
+    # Drop a copy of our custom modeling file next to the checkpoint so
+    # `from_pretrained(..., trust_remote_code=True)` can pick it up.
+    import shutil
+
+    src_modeling = Path(__file__).parent / "_primus_mamba2_modeling.py"
+    dst_modeling = output_dir / "modeling_mamba2_full_mlp.py"
+    shutil.copy2(src_modeling, dst_modeling)
+    print(f"  Copied modeling_mamba2_full_mlp.py (custom class with per-layer MLPs)")
+    # Re-export Mamba2Config from FLA in the same module so AutoConfig finds it.
+    with open(dst_modeling, "a") as f:
+        f.write("\n\n# Re-export the FLA config under the same module for auto_map\n")
+        f.write("from fla.models.mamba2.configuration_mamba2 import Mamba2Config\n")
+
+    # Copy tokenizer if a source was provided.
+    src_tok = Path(args.tokenizer)
+    if src_tok.exists():
+        import shutil
+
+        for name in (
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+        ):
+            f_src = src_tok / name
+            if f_src.exists():
+                shutil.copy2(f_src, output_dir / name)
+                print(f"  Copied {name} from {src_tok}")
+    else:
+        # Minimal fallback so lm-eval can still load.
+        with open(output_dir / "tokenizer_config.json", "w") as f:
+            json.dump({"tokenizer_class": "PreTrainedTokenizerFast", "model_max_length": 2048}, f, indent=2)
+
+    print()
+    print("Done. To sanity-check then evaluate:")
+    print(
+        f"  python -c \"from fla.models import Mamba2ForCausalLM; m = Mamba2ForCausalLM.from_pretrained('{output_dir}'); print(sum(p.numel() for p in m.parameters())/1e6, 'M params')\""
+    )
+    print()
+    print(
+        f"  lm_eval --model hf --model_args pretrained={output_dir},trust_remote_code=True,dtype=bfloat16 \\"
+    )
+    print(f"    --tasks hellaswag,winogrande,piqa,arc_easy,arc_challenge,lambada_openai \\")
+    print(f"    --batch_size 16 --output_path {output_dir}/lm_eval")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_zebra_llama_to_hf.py b/tools/hybrid/convert_zebra_llama_to_hf.py
new file mode 100644
index 000000000..bac81d30a
--- /dev/null
+++ b/tools/hybrid/convert_zebra_llama_to_hf.py
@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+"""
+Convert Megatron Zebra-Llama checkpoint to HuggingFace format.
+
+This script converts a trained Zebra-Llama model (Hybrid Mamba+MLA) from
+Megatron-LM format to a HuggingFace-compatible format for evaluation.
+"""
+
+import argparse
+import json
+import sys
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+
+# Add Megatron-LM to Python path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "third_party" / "Megatron-LM"))
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+
+def load_megatron_checkpoint(checkpoint_path):
+    """Load Megatron checkpoint from disk."""
+    print(f"Loading checkpoint from: {checkpoint_path}")
+
+    # Load the model weights
+    model_path = Path(checkpoint_path) / "mp_rank_00" / "model_optim_rng.pt"
+
+    if not model_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+
+    print("Loading checkpoint (this may take a moment)...")
+    try:
+        # Try with weights_only=True first (safer)
+        checkpoint = torch.load(model_path, map_location="cpu", weights_only=True)
+        print("✓ Loaded with weights_only=True")
+    except Exception as e:
+        print(f"weights_only=True failed ({e}), trying full load...")
+        checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+        print("✓ Loaded with weights_only=False")
+
+    print(f"Checkpoint keys: {checkpoint.keys()}")
+
+    return checkpoint
+
+
+def extract_model_state(checkpoint):
+    """Extract model state dict from Megatron checkpoint."""
+    if "model" in checkpoint:
+        model_state = checkpoint["model"]
+    elif "state_dict" in checkpoint:
+        model_state = checkpoint["state_dict"]
+    else:
+        # Try to find the model state in the checkpoint
+        for key in checkpoint.keys():
+            if "model" in key.lower():
+                model_state = checkpoint[key]
+                break
+        else:
+            model_state = checkpoint
+
+    print(f"Model state contains {len(model_state)} parameters")
+
+    # Print some example keys
+    print("\nExample parameter names:")
+    for i, (key, v) in enumerate(list(model_state.items())):
+        if v is not None and hasattr(v, "shape"):
+            print(f"  {key}: {v.shape} ({v.dtype})")
+        else:
+            print(f"  {key}: {type(v)} (non-tensor)")
+
+    return model_state
+
+
+def convert_to_hf_format(model_state, config):
+    """Convert Megatron model state to HuggingFace format."""
+    hf_state = OrderedDict()
+
+    # This is a template - you'll need to customize based on your model architecture
+    # The key mapping depends on how your Zebra-Llama model is structured
+
+    print("\nConverting to HuggingFace format...")
+
+    for key, value in model_state.items():
+        # Remove 'module.' prefix if present
+        if key.startswith("module."):
+            key = key[7:]
+
+        # Convert layer names
+        # Example mappings (customize for your architecture):
+        # decoder.layers.0.mixer.in_proj.weight -> model.layers.0.mamba.in_proj.weight
+        new_key = key
+        if key.startswith("decoder."):
+            if key.startswith("decoder.final_norm."):
+                new_key = key.replace("decoder.final_norm.", "model.norm.")
+            else:
+                new_key = key.replace("decoder.", "model.")
+        if key.startswith("embedding.word_embeddings."):
+            new_key = key.replace("embedding.word_embeddings.", "model.embed_tokens.")
+
+        if "linear_kv_up_proj.layer_norm_weight" in new_key:
+            new_key = new_key.replace("linear_kv_up_proj.layer_norm_weight", "kv_layernorm.weight")
+        if "linear_kv_up_proj.layer_norm_bias" in new_key:
+            new_key = new_key.replace("linear_kv_up_proj.layer_norm_bias", "kv_layernorm.bias")
+        if "linear_q_up_proj.layer_norm_weight" in new_key:
+            new_key = new_key.replace("linear_q_up_proj.layer_norm_weight", "q_layernorm.weight")
+        if "linear_q_up_proj.layer_norm_bias" in new_key:
+            new_key = new_key.replace("linear_q_up_proj.layer_norm_bias", "q_layernorm.bias")
+        if "linear_fc1.layer_norm_weight" in new_key:
+            new_key = new_key.replace("linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight")
+        if "linear_fc1.layer_norm_bias" in new_key:
+            new_key = new_key.replace("linear_fc1.layer_norm_bias", "pre_mlp_layernorm.bias")
+        if "mixer.in_proj.layer_norm_weight" in new_key:
+            new_key = new_key.replace("mixer.in_proj.layer_norm_weight", "norm.weight")
+        if "mixer.in_proj.layer_norm_bias" in new_key:
+            new_key = new_key.replace("mixer.in_proj.layer_norm_bias", "norm.bias")
+        if "mlp.pre_mlp_layernorm" in new_key:
+            new_key = new_key.replace("mlp.pre_mlp_layernorm", "pre_mlp_layernorm")
+
+        # if 'layer_norm_weight' in new_key:
+        #     new_key = new_key.replace('layer_norm_weight', 'weight')
+        # if 'layer_norm_bias' in new_key:
+        #     new_key = new_key.replace('layer_norm_bias', 'bias')
+
+        if "_extra_state" not in new_key:
+            hf_state[new_key] = value
+
+    # Ensure lm_head.weight exists (Megatron uses output_layer.weight)
+    # Prefer explicit output layer if present; otherwise fall back to tying with embeddings.
+    if "lm_head.weight" not in hf_state:
+        if "output_layer.weight" in model_state:
+            hf_state["lm_head.weight"] = model_state["output_layer.weight"]
+        elif "model.output_layer.weight" in hf_state:
+            hf_state["lm_head.weight"] = hf_state["model.output_layer.weight"]
+        elif "model.embed_tokens.weight" in hf_state:
+            # Fallback: tie lm_head to embeddings (common when weights are tied)
+            hf_state["lm_head.weight"] = hf_state["model.embed_tokens.weight"]
+        elif "embedding.word_embeddings.weight" in model_state:
+            hf_state["lm_head.weight"] = model_state["embedding.word_embeddings.weight"]
+
+    return hf_state
+
+
+def save_hf_checkpoint(hf_state, config, output_dir):
+    """Save checkpoint in HuggingFace format."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Save model weights
+    model_path = output_dir / "pytorch_model.bin"
+    print(f"\nSaving model weights to: {model_path}")
+    torch.save(hf_state, model_path)
+
+    # Save config
+    config_path = output_dir / "config.json"
+    print(f"Saving config to: {config_path}")
+    with open(config_path, "w") as f:
+        json.dump(config, f, indent=2)
+
+    ratio = config.get("hybrid_attention_ratio", 0.25)
+    if ratio <= 0.0:
+        arch_desc = "Pure KDA (Kimi Delta Attention)"
+    elif ratio >= 1.0:
+        arch_desc = "Pure MLA (Multi-Latent Attention)"
+    else:
+        arch_desc = f"Hybrid KDA + MLA (attention_ratio={ratio})"
+
+    readme_path = output_dir / "README.md"
+    with open(readme_path, "w") as f:
+        f.write(
+            f"""# Zebra-Llama {config.get('hidden_size', '?')}
+
+Converted from Megatron-LM checkpoint.
+
+## Model Details
+- Hidden Size: {config.get('hidden_size', 'N/A')}
+- Num Sublayers: {config.get('num_hidden_layers', 'N/A')}
+- Vocab Size: {config.get('vocab_size', 'N/A')}
+- Architecture: {arch_desc}
+- Tied Embeddings: {config.get('tie_word_embeddings', False)}
+"""
+        )
+
+    print(f"\n✓ Conversion complete! Saved to: {output_dir}")
+
+
+def _getattr(obj, name, default=None):
+    """getattr that also treats ``None`` as missing (falls back to *default*)."""
+    v = getattr(obj, name, None)
+    return v if v is not None else default
+
+
+def create_config_from_checkpoint(checkpoint, args):
+    """Create HuggingFace config by reading everything from ``checkpoint['args']``.
+
+    The CLI ``args`` are used only as final fallbacks when a field is missing
+    from the Megatron namespace (which almost never happens).
+    """
+    megatron_args = checkpoint["args"]
+    untie = getattr(megatron_args, "untie_embeddings_and_output_weights", True)
+
+    config = {
+        "architectures": ["ZebraLlamaForCausalLM"],
+        "model_type": "zebra_llama",
+        # Core dimensions
+        "hidden_size": _getattr(megatron_args, "hidden_size", args.hidden_size),
+        "num_hidden_layers": _getattr(megatron_args, "num_layers", args.num_layers),
+        "intermediate_size": _getattr(megatron_args, "ffn_hidden_size", 8192),
+        "vocab_size": _getattr(megatron_args, "padded_vocab_size", args.vocab_size),
+        "num_attention_heads": _getattr(megatron_args, "num_attention_heads", args.num_attention_heads),
+        # Norm / dtype
+        "layernorm_epsilon": _getattr(megatron_args, "norm_epsilon", 1e-5),
+        "torch_dtype": "bfloat16" if getattr(megatron_args, "bf16", False) else "float32",
+        # Embeddings
+        "tie_word_embeddings": not untie,
+        # Hybrid layout
+        "hybrid_attention_ratio": _getattr(megatron_args, "hybrid_attention_ratio", 0.25),
+        # KDA (Kimi Delta Attention) -- all derived from Megatron linear_* args
+        "kda_num_heads": _getattr(megatron_args, "linear_num_value_heads", 16),
+        "kda_head_dim": _getattr(megatron_args, "linear_value_head_dim", 64),
+        "kda_key_head_dim": _getattr(megatron_args, "linear_key_head_dim", 32),
+        "kda_num_key_heads": _getattr(megatron_args, "linear_num_key_heads", 16),
+        "kda_conv_kernel": _getattr(megatron_args, "linear_conv_kernel_dim", 4),
+        # MLA (Multi-Latent Attention)
+        "q_lora_rank": getattr(megatron_args, "q_lora_rank", None),
+        "kv_lora_rank": _getattr(megatron_args, "kv_lora_rank", 128),
+        "qk_head_dim": _getattr(megatron_args, "qk_head_dim", 32),
+        "qk_pos_emb_head_dim": _getattr(megatron_args, "qk_pos_emb_head_dim", 32),
+        "v_head_dim": _getattr(megatron_args, "v_head_dim", 64),
+        # Mamba (legacy compat)
+        "mamba_state_dim": _getattr(megatron_args, "mamba_state_dim", 64),
+        "mamba_head_dim": _getattr(megatron_args, "mamba_head_dim", 64),
+        "mamba_num_groups": _getattr(megatron_args, "mamba_num_groups", 8),
+        # RoPE / YaRN
+        "original_max_position_embeddings": _getattr(megatron_args, "original_max_position_embeddings", 4096),
+        "rotary_scaling_factor": _getattr(megatron_args, "rotary_scaling_factor", 1.0),
+        "rotary_base": _getattr(megatron_args, "rotary_base", 500000),
+        "mscale": _getattr(megatron_args, "mscale", 1.0),
+        "mscale_all_dim": _getattr(megatron_args, "mscale_all_dim", 1.0),
+        "beta_fast": _getattr(megatron_args, "beta_fast", 32.0),
+        "beta_slow": _getattr(megatron_args, "beta_slow", 1.0),
+    }
+
+    return config
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert Megatron Zebra-Llama checkpoint to HuggingFace format"
+    )
+
+    parser.add_argument(
+        "--checkpoint-path",
+        type=str,
+        required=True,
+        help="Path to Megatron checkpoint directory (e.g., output/zebra_llama_1B-pretrain/iter_0001000)",
+    )
+    parser.add_argument(
+        "--output-dir", type=str, required=True, help="Output directory for HuggingFace checkpoint"
+    )
+    parser.add_argument("--hidden-size", type=int, default=2048, help="Hidden size of the model")
+    parser.add_argument("--num-layers", type=int, default=24, help="Number of transformer layers")
+    parser.add_argument("--num-attention-heads", type=int, default=16, help="Number of attention heads")
+    parser.add_argument("--vocab-size", type=int, default=128256, help="Vocabulary size")
+
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print("Megatron to HuggingFace Checkpoint Conversion")
+    print("=" * 70)
+
+    # Step 1: Load Megatron checkpoint
+    checkpoint = load_megatron_checkpoint(args.checkpoint_path)
+
+    # Step 2: Extract model state
+    model_state = extract_model_state(checkpoint)
+
+    # Step 3: Create config
+    config = create_config_from_checkpoint(checkpoint, args)
+    print(f"\nModel config: {json.dumps(config, indent=2)}")
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    from modeling_zebra_llama import ZebraLlamaConfig, ZebraLlamaForCausalLM
+
+    zebra_config = ZebraLlamaConfig(**config)
+    model = ZebraLlamaForCausalLM(zebra_config)
+
+    # Step 4: Convert to HuggingFace format
+    hf_state = convert_to_hf_format(model_state, zebra_config)
+    sd = model.state_dict()
+
+    # Compare keys between converted state and model's expected keys
+    hf_keys = set(hf_state.keys())
+    model_keys = set(sd.keys())
+    missing_in_hf = sorted(model_keys - hf_keys)
+    extra_in_hf = sorted(hf_keys - model_keys)
+
+    print("\n" + "=" * 70)
+    print("State dict key comparison")
+    print("=" * 70)
+    print(f"HF state keys:     {len(hf_keys)}")
+    print(f"Model state keys:  {len(model_keys)}")
+    print(f"Missing in hf_state (expected by model): {len(missing_in_hf)}")
+    print(f"Extra in hf_state (not in model):        {len(extra_in_hf)}")
+
+    if missing_in_hf:
+        print("\nFirst missing keys:")
+        for k in missing_in_hf[:50]:
+            shape = tuple(sd[k].shape) if hasattr(sd[k], "shape") else None
+            dtype = str(sd[k].dtype) if hasattr(sd[k], "dtype") else None
+            print(f"  - {k}  shape={shape} dtype={dtype}")
+        if len(missing_in_hf) > 50:
+            print(f"  ... and {len(missing_in_hf) - 50} more")
+
+    if extra_in_hf:
+        print("\nFirst extra keys:")
+        for k in extra_in_hf[:50]:
+            shape = tuple(hf_state[k].shape) if hasattr(hf_state[k], "shape") else None
+            dtype = str(hf_state[k].dtype) if hasattr(hf_state[k], "dtype") else None
+            print(f"  - {k}  shape={shape} dtype={dtype}")
+        if len(extra_in_hf) > 50:
+            print(f"  ... and {len(extra_in_hf) - 50} more")
+
+    # Shape check for intersection
+    common = sorted(hf_keys & model_keys)
+    shape_mismatches = []
+    for k in common:
+        v_hf = hf_state[k]
+        v_md = sd[k]
+        if hasattr(v_hf, "shape") and hasattr(v_md, "shape") and tuple(v_hf.shape) != tuple(v_md.shape):
+            shape_mismatches.append((k, tuple(v_hf.shape), tuple(v_md.shape)))
+
+    print(f"\nShape mismatches on common keys: {len(shape_mismatches)}")
+    for k, sh_hf, sh_md in shape_mismatches[:50]:
+        print(f"  - {k}: hf_state{sh_hf} vs model{sh_md}")
+    if len(shape_mismatches) > 50:
+        print(f"  ... and {len(shape_mismatches) - 50} more")
+
+    model.to(torch.bfloat16)
+    # Use strict=False so we can see all missing/extra keys without crashing,
+    # but we expect lm_head.weight to be present now.
+    missing, unexpected = model.load_state_dict(hf_state, strict=False)
+
+    if missing:
+        print("\nMissing keys when loading into model (strict=False):")
+        for k in missing[:50]:
+            print(f"  - {k}")
+        if len(missing) > 50:
+            print(f"  ... and {len(missing) - 50} more")
+    if unexpected:
+        print("\nUnexpected keys when loading into model (strict=False):")
+        for k in unexpected[:50]:
+            print(f"  - {k}")
+        if len(unexpected) > 50:
+            print(f"  ... and {len(unexpected) - 50} more")
+
+    # Step 5: Save HuggingFace checkpoint
+    save_hf_checkpoint(hf_state, config, args.output_dir)
+
+    print("\n" + "=" * 70)
+    print("Next steps:")
+    print("=" * 70)
+    print("1. Review the converted checkpoint")
+    print("2. Create a custom modeling file for Zebra-Llama if needed")
+    print("3. Run evaluation with lm-eval-harness")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/convert_zebra_llama_to_hf.sh b/tools/hybrid/convert_zebra_llama_to_hf.sh
new file mode 100644
index 000000000..a89e67561
--- /dev/null
+++ b/tools/hybrid/convert_zebra_llama_to_hf.sh
@@ -0,0 +1,4 @@
+#!/bin/bash
+python tools/hybrid/convert_zebra_llama_to_hf.py \
+    --checkpoint-path output/zebra_llama_1B-pretrain/iter_0020000 \
+    --output-dir output/zebra_llama_1B_hf_iter_0020000 \
diff --git a/tools/hybrid/eval_gdn_lm_eval.py b/tools/hybrid/eval_gdn_lm_eval.py
new file mode 100644
index 000000000..ae66d320b
--- /dev/null
+++ b/tools/hybrid/eval_gdn_lm_eval.py
@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""
+lm-eval wrapper for GDN models converted to FLA HuggingFace format.
+
+Importing `fla` registers GatedDeltaNetConfig / GatedDeltaNetForCausalLM
+with transformers' AutoConfig / AutoModel, which lm-eval's --model hf
+path relies on. Without this import, AutoConfig.from_pretrained fails
+with 'model type gated_deltanet not recognized'.
+
+Usage (same CLI as lm_eval, just swap the command):
+
+    python tools/hybrid/eval_gdn_lm_eval.py \
+        --model hf \
+        --model_args pretrained=output/gdn_pure_1B_fla_hf,dtype=bfloat16,trust_remote_code=True \
+        --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+        --batch_size auto \
+        --output_path eval_results/gdn_pure_1B
+"""
+import fla  # noqa: F401  — registers GatedDeltaNet with AutoConfig/AutoModel
+
+# Monkey-patch FLA model classes to accept **kwargs (e.g. 'dtype')
+# that newer transformers (>=4.55) passes internally during from_pretrained
+from fla.models.gated_deltanet import GatedDeltaNetForCausalLM, GatedDeltaNetModel
+
+_orig_causal_init = GatedDeltaNetForCausalLM.__init__
+_orig_model_init = GatedDeltaNetModel.__init__
+
+
+def _patched_causal_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_causal_init(self, config, *args, **kwargs)
+
+
+def _patched_model_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_model_init(self, config, *args, **kwargs)
+
+
+GatedDeltaNetForCausalLM.__init__ = _patched_causal_init
+GatedDeltaNetModel.__init__ = _patched_model_init
+
+from lm_eval.__main__ import cli_evaluate
+
+if __name__ == "__main__":
+    cli_evaluate()
diff --git a/tools/hybrid/eval_kda_lm_eval.py b/tools/hybrid/eval_kda_lm_eval.py
new file mode 100644
index 000000000..0b9d6fce1
--- /dev/null
+++ b/tools/hybrid/eval_kda_lm_eval.py
@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""
+lm-eval wrapper for KDA models (Kimi Delta Attention).
+
+Importing `fla` registers KDAConfig / KDAForCausalLM with transformers'
+AutoConfig / AutoModel, which lm-eval's --model hf path relies on.
+Without this import, AutoConfig.from_pretrained fails with
+'model type kda not recognized'.
+
+Usage (same CLI as lm_eval, just swap the command):
+
+    python tools/hybrid/eval_kda_lm_eval.py \
+        --model hf \
+        --model_args pretrained=output/kda_pure_300M_fla_hf,dtype=bfloat16,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
+        --tasks arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race \
+        --batch_size auto \
+        --output_path output/kda_pure_300M_eval_results
+"""
+import fla  # noqa: F401  — registers KDA with AutoConfig/AutoModel
+
+# Monkey-patch FLA model classes to accept **kwargs (e.g. 'dtype')
+# that newer transformers (>=4.55) passes internally during from_pretrained
+from fla.models.kda import KDAForCausalLM, KDAModel
+
+_orig_causal_init = KDAForCausalLM.__init__
+_orig_model_init = KDAModel.__init__
+
+
+def _patched_causal_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_causal_init(self, config, *args, **kwargs)
+
+
+def _patched_model_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_model_init(self, config, *args, **kwargs)
+
+
+KDAForCausalLM.__init__ = _patched_causal_init
+KDAModel.__init__ = _patched_model_init
+
+from lm_eval.__main__ import cli_evaluate
+
+if __name__ == "__main__":
+    cli_evaluate()
diff --git a/tools/hybrid/eval_mamba2_lm_eval.py b/tools/hybrid/eval_mamba2_lm_eval.py
new file mode 100644
index 000000000..da1ce63bb
--- /dev/null
+++ b/tools/hybrid/eval_mamba2_lm_eval.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+lm-eval wrapper for Mamba2 hybrid models in FLA HuggingFace format.
+
+Adds the same three workarounds we needed for GDN, plus one extra:
+
+  1. Import `fla` so Mamba2Config / Mamba2ForCausalLM register with
+     transformers' AutoConfig / AutoModel.
+  2. Monkey-patch Mamba2ForCausalLM / Mamba2Model __init__ to swallow
+     the `dtype=...` kwarg that newer transformers passes internally.
+  3. Force `residual_in_fp32=False` on the loaded config — FLA's Mamba2
+     mixer in_proj is bf16; with residual_in_fp32=True the residual comes
+     back from mixer_norm in fp32 and the next mixer's in_proj F.linear
+     fails with "expected mat1 and mat2 to have the same dtype".
+  4. Default `trust_remote_code=True` so the Primus-converted checkpoint
+     auto-loads via our custom Mamba2FullMlpForCausalLM class (saved as
+     `modeling_mamba2_full_mlp.py` next to the checkpoint).
+
+Usage (same CLI as lm_eval, just swap the command):
+
+    python tools/hybrid/eval_mamba2_lm_eval.py \
+        --model hf \
+        --model_args pretrained=output/mamba_hybrid_300M_fla_hf,dtype=bfloat16,trust_remote_code=True \
+        --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
+        --batch_size auto \
+        --output_path eval_results/mamba_hybrid_300M_primus
+"""
+
+import fla  # noqa: F401  — registers Mamba2 with AutoConfig/AutoModel
+
+# 1. dtype-kwarg patch (same trick as eval_gdn_lm_eval.py)
+from fla.models.mamba2 import Mamba2ForCausalLM, Mamba2Model
+
+_orig_causal_init = Mamba2ForCausalLM.__init__
+_orig_model_init = Mamba2Model.__init__
+
+
+def _patched_causal_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_causal_init(self, config, *args, **kwargs)
+
+
+def _patched_model_init(self, config, *args, **kwargs):
+    kwargs.pop("dtype", None)
+    return _orig_model_init(self, config, *args, **kwargs)
+
+
+Mamba2ForCausalLM.__init__ = _patched_causal_init
+Mamba2Model.__init__ = _patched_model_init
+
+
+# 2. residual_in_fp32=False patch — most reliable place is in the model's
+#    __init__: override config.residual_in_fp32 before super().__init__ uses it.
+def _make_patch(orig):
+    def _patched(self, config, *args, **kwargs):
+        kwargs.pop("dtype", None)
+        if getattr(config, "residual_in_fp32", False):
+            config.residual_in_fp32 = False
+            print(f"[eval_mamba2_lm_eval] forced residual_in_fp32=False on {type(self).__name__}")
+        return orig(self, config, *args, **kwargs)
+
+    return _patched
+
+
+Mamba2ForCausalLM.__init__ = _make_patch(_orig_causal_init)
+Mamba2Model.__init__ = _make_patch(_orig_model_init)
+
+# Also patch Mamba2Block.__init__ so the per-block `self.residual_in_fp32`
+# attribute is set from the (now-patched) config, not from the original
+# True value.  Block is constructed inside Mamba2Model.__init__.
+from fla.models.mamba2.modeling_mamba2 import Mamba2Block as _Mamba2Block
+
+_orig_block_init = _Mamba2Block.__init__
+
+
+def _patched_block_init(self, config, layer_idx):
+    if getattr(config, "residual_in_fp32", False):
+        config.residual_in_fp32 = False
+    return _orig_block_init(self, config, layer_idx)
+
+
+_Mamba2Block.__init__ = _patched_block_init
+
+
+# 3. Hand off to lm-eval CLI
+from lm_eval.__main__ import cli_evaluate
+
+if __name__ == "__main__":
+    cli_evaluate()
diff --git a/tools/hybrid/eval_zebra_llama_lm_eval.sh b/tools/hybrid/eval_zebra_llama_lm_eval.sh
new file mode 100644
index 000000000..7f0127d8f
--- /dev/null
+++ b/tools/hybrid/eval_zebra_llama_lm_eval.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+
+args=("$@")
+for arg in "${args[@]}"; do
+    eval "$arg"
+done
+
+echo "model_path:   ${model_path:=output/zebra_llama_1B_kda_pure_hf}"
+echo "tokenizer:    ${tokenizer:=meta-llama/Llama-3.2-1B}"
+echo "batch_size:   ${batch_size:=auto}"
+echo "num_fewshot:  ${num_fewshot:=0}"
+echo "output_path:  ${output_path:=eval_results}"
+echo "device:       ${device:=cuda}"
+
+TASKS="arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande"
+
+pip install lm-eval --quiet 2>/dev/null
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "${SCRIPT_DIR}/../.." || exit 1
+
+python -c "
+import sys; sys.path.insert(0, 'tools/hybrid')
+from modeling_zebra_llama import ZebraLlamaForCausalLM, ZebraLlamaConfig
+from transformers import AutoConfig, AutoModelForCausalLM
+AutoConfig.register('zebra_llama', ZebraLlamaConfig)
+AutoModelForCausalLM.register(ZebraLlamaConfig, ZebraLlamaForCausalLM)
+import lm_eval
+results = lm_eval.simple_evaluate(
+    model='hf',
+    model_args='pretrained=${model_path},tokenizer=${tokenizer},trust_remote_code=True',
+    tasks='${TASKS}'.split(','),
+    batch_size='${batch_size}',
+    num_fewshot=${num_fewshot} if '${num_fewshot}' != 'None' else None,
+    device='${device}',
+    log_samples=True,
+)
+from lm_eval.utils import make_table
+print(make_table(results))
+import json, os
+os.makedirs('${output_path}', exist_ok=True)
+with open('${output_path}/results.json', 'w') as f:
+    json.dump(results['results'], f, indent=2, default=str)
+print(f\"Results saved to ${output_path}/results.json\")
+"
diff --git a/tools/hybrid/fla_order_dataset.py b/tools/hybrid/fla_order_dataset.py
new file mode 100644
index 000000000..c4064d905
--- /dev/null
+++ b/tools/hybrid/fla_order_dataset.py
@@ -0,0 +1,121 @@
+"""
+FLA-order GPT dataset for Megatron.
+
+Loads FLA's preprocessed HuggingFace dataset and serves samples in the exact
+same order that FLA's training pipeline (HuggingFace Trainer + DistributedSampler)
+would present them.  This lets Primus/Megatron reproduce FLA's training trajectory
+bit-for-bit (modulo bf16 numerics), proving that any remaining loss gap comes
+solely from data ordering.
+
+Usage:
+    Set PRIMUS_FLA_DATA=1 and PRIMUS_FLA_CACHE_DIR=<path_to_fla_preprocessed_data>
+    before launching training.
+
+    Example:
+        PRIMUS_FLA_DATA=1 \\
+        PRIMUS_FLA_CACHE_DIR=/home/<user>/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train \\
+        EXP=examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml \\
+        GPUS_PER_NODE=8 bash examples/run_pretrain.sh
+"""
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+class FLAOrderGPTDataset(Dataset):
+    """
+    Wraps a preprocessed HuggingFace dataset (2048-token sequences) and serves
+    samples in the exact order that FLA's DistributedSampler would.
+
+    Megatron's MegatronPretrainingSampler assigns *contiguous blocks* of
+    micro_batch_size to each rank:
+        rank r at step t gets global indices [t*G + r*M, ..., t*G + r*M + M-1]
+        where G = micro_batch_size * data_parallel_size, M = micro_batch_size.
+
+    FLA's DistributedSampler assigns *strided* indices from a shuffled permutation:
+        rank r at step t gets shuffled_dataset[perm[r + (t*M + j)*D]] for j in 0..M-1
+        where D = data_parallel_size, perm = randperm(N, seed=42).
+
+    This dataset pre-builds an index_map so that:
+        __getitem__(megatron_gidx) returns the same data FLA would serve
+        for the same (rank, step, position-in-batch) tuple.
+    """
+
+    def __init__(
+        self,
+        cache_dir: str,
+        seq_length: int,
+        micro_batch_size: int,
+        data_parallel_size: int,
+        seed: int = 42,
+        pad_token_id: int = 0,
+        eod_token: int = 128000,
+        eod_mask_loss: bool = False,
+    ):
+        from datasets import load_from_disk
+
+        raw_dataset = load_from_disk(cache_dir)
+        self.dataset = raw_dataset.shuffle(seed=seed)
+        self.N = len(self.dataset)
+        self.seq_length = seq_length
+        self.pad_token_id = pad_token_id
+        self.eod_token = eod_token
+        self.eod_mask_loss = eod_mask_loss
+
+        mbs = micro_batch_size
+        dp = data_parallel_size
+        gb = mbs * dp
+
+        # Replicate PyTorch DistributedSampler's permutation
+        g = torch.Generator()
+        g.manual_seed(seed)  # seed + epoch(0)
+        total_size = ((self.N + dp - 1) // dp) * dp
+        perm = torch.randperm(self.N, generator=g).tolist()
+        padding = total_size - self.N
+        perm = perm + perm[:padding]
+
+        # Build megatron_gidx → fla_dataset_idx mapping
+        self.index_map = np.zeros(total_size, dtype=np.int64)
+        for gidx in range(total_size):
+            t = gidx // gb
+            r = (gidx % gb) // mbs
+            j = gidx % mbs
+            fla_perm_idx = r + (t * mbs + j) * dp
+            if fla_perm_idx < len(perm):
+                self.index_map[gidx] = perm[fla_perm_idx] % self.N
+            else:
+                self.index_map[gidx] = 0
+
+        self._cached_loss_mask = None
+        self._cached_position_ids = None
+
+    def __len__(self):
+        return len(self.index_map)
+
+    def __getitem__(self, idx):
+        if idx is None:
+            idx = 0
+
+        fla_idx = int(self.index_map[idx % len(self.index_map)])
+        input_ids = self.dataset[fla_idx]["input_ids"]
+
+        tokens = torch.tensor(input_ids[: self.seq_length], dtype=torch.long)
+        labels = torch.roll(tokens, shifts=-1, dims=0)
+        labels[-1] = self.pad_token_id
+
+        if self._cached_loss_mask is None:
+            loss_mask = torch.ones(self.seq_length, dtype=torch.float32)
+            loss_mask[-1] = 0.0
+            if self.eod_mask_loss:
+                eod_positions = tokens == self.eod_token
+                loss_mask[eod_positions] = 0.0
+            self._cached_loss_mask = loss_mask
+            self._cached_position_ids = torch.arange(self.seq_length, dtype=torch.long)
+
+        return {
+            "tokens": tokens,
+            "labels": labels,
+            "loss_mask": self._cached_loss_mask.clone(),
+            "position_ids": self._cached_position_ids,
+        }
diff --git a/tools/hybrid/launch_hybrid_faflag.sh b/tools/hybrid/launch_hybrid_faflag.sh
new file mode 100644
index 000000000..f97eb8312
--- /dev/null
+++ b/tools/hybrid/launch_hybrid_faflag.sh
@@ -0,0 +1,74 @@
+#!/bin/bash
+###############################################################################
+# Foolproof launcher for the 75% Hybrid GDN run with the FULL FLA-parity stack
+# enabled.  Run this *inside the container*:
+#
+#   cd /workspace/Primus && bash tools/hybrid/launch_hybrid_faflag.sh
+#
+# FLA-parity flags exported (matched to docs/04-technical-guides/hybrid-models/gdn-fla-parity.md §A/B/D and the
+# pure-KDA config that hit 1.46 s/iter on MI300X):
+#
+#   PRIMUS_FLA_MLA_ATTN=1  → MLA `core_attention` calls `flash_attn_func`
+#                             directly, skipping TE's CK fallback (~30 ms/MLA
+#                             block on MI300X with flash-attn 2.8.3 > 2.8.1).
+#   PRIMUS_FUSED_CE=1      → FLA `FusedLinearCrossEntropyLoss` (chunked, no
+#                             full logits tensor — huge memory + speed win).
+#   PRIMUS_FLA_SWIGLU=1    → FLA's Triton SwiGLU instead of Megatron's naive
+#                             `silu * x` (saves ~20 ms/iter).
+#   PRIMUS_FLA_NORM=1      → (a) WrappedTorchNorm → FLA's Triton RMSNorm
+#                             (used by GDN, MLP pre-norm, AND MLA's LoRA
+#                             q_layernorm/kv_layernorm — fixes the +0.12 loss
+#                             gap vs FLA), (b) GDN out_norm → FLA's
+#                             FusedRMSNormGated (RMSNorm + sigmoid + mul in
+#                             one Triton kernel — saves ~50 ms/iter), and
+#                             (c) HybridStack fuses each GDN block's mixer-out
+#                             with the next MLP's pre-MLP layernorm (~30 ms).
+#   PRIMUS_FLA_CONV=1      → FLA's Triton causal_conv1d for GDN (matches FLA's
+#                             default; small speed win and bit-exact gradient).
+#   PRIMUS_FLA_DATA=1      → drive Megatron from FLA's `DistributedSampler`
+#                             token order so the loss curves are directly
+#                             comparable (no batch-order noise).
+#
+# Combined, on the 75% GDN+MLA hybrid these match FLA's 1.47 s/iter and
+# loss-curve from iter 100 onward (iter-1 was already bit-perfect).
+###############################################################################
+
+set -euo pipefail
+
+export PRIMUS_FLA_MLA_ATTN=1
+export PRIMUS_FUSED_CE=1
+export PRIMUS_FLA_SWIGLU=1
+export PRIMUS_FLA_NORM=1
+export PRIMUS_FLA_CONV=1
+export PRIMUS_FLA_DATA=1
+# PRIMUS_FLA_DATA is a no-op without PRIMUS_FLA_CACHE_DIR — the trainer's
+# FLAOrderGPTDataset is guarded by `fla_data_flag == "1" and fla_cache`
+# (primus/modules/trainer/megatron/trainer.py).  Point this at the FLA
+# fineweb-edu cache so Primus consumes tokens in the exact same order as
+# FLA's HF DistributedSampler — eliminates the dataloader-ordering bias
+# that drives the +0.02 late-training loss gap and the +2.4 warm-up bump.
+export PRIMUS_FLA_CACHE_DIR=${PRIMUS_FLA_CACHE_DIR:-$HOME/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train}
+
+EXP=${EXP:-examples/megatron/configs/MI300X/zebra_llama_300M_gdn_hybrid-pretrain.yaml}
+LOG=${LOG:-primus_gdn_hybrid_300M_faflag.log}
+
+echo "==========[launch_hybrid_faflag.sh]=========="
+echo "PRIMUS_FLA_MLA_ATTN  = ${PRIMUS_FLA_MLA_ATTN}"
+echo "PRIMUS_FUSED_CE      = ${PRIMUS_FUSED_CE}"
+echo "PRIMUS_FLA_SWIGLU    = ${PRIMUS_FLA_SWIGLU}"
+echo "PRIMUS_FLA_NORM      = ${PRIMUS_FLA_NORM}"
+echo "PRIMUS_FLA_CONV      = ${PRIMUS_FLA_CONV}"
+echo "PRIMUS_FLA_DATA      = ${PRIMUS_FLA_DATA}"
+echo "PRIMUS_FLA_CACHE_DIR = ${PRIMUS_FLA_CACHE_DIR}"
+echo "EXP                  = ${EXP}"
+echo "LOG                  = ${LOG}"
+echo "============================================="
+
+# Sanity check: warn loudly if the FLA cache dir is missing
+if [ ! -d "${PRIMUS_FLA_CACHE_DIR}" ]; then
+    echo "WARNING: PRIMUS_FLA_CACHE_DIR=${PRIMUS_FLA_CACHE_DIR} does not exist."
+    echo "         FLAOrderGPTDataset will silently fall back to Megatron's GPTDataset"
+    echo "         shuffler, and you'll see the +0.02 late-training loss gap reappear."
+fi
+
+EXP="${EXP}" bash examples/run_pretrain.sh 2>&1 | tee "${LOG}"
diff --git a/tools/hybrid/launch_mamba_hybrid_300M.sh b/tools/hybrid/launch_mamba_hybrid_300M.sh
new file mode 100644
index 000000000..7fc9059f2
--- /dev/null
+++ b/tools/hybrid/launch_mamba_hybrid_300M.sh
@@ -0,0 +1,65 @@
+#!/bin/bash
+###############################################################################
+# Launch the 75% Hybrid Mamba2 + MLA 300M pretrain (matched to the FLA
+# `train_mamba2_hybrid_300M.log` schedule: 4768 iter × 1024 batch × 2048 seq
+# ≈ 10B tokens on 8 GPUs).
+#
+# Full FLA-parity stack — same as tools/hybrid/launch_hybrid_faflag.sh (the GDN hybrid).
+# An early run without these flags reproduced the iter-1 bit-perfect parity
+# with FLA, then drifted +2.58 nats by iter 100 (the warm-up spike we already
+# debugged for GDN).  Enabling the fusion / data flags closes that gap and
+# also unlocks ~25 % more throughput.
+#
+#   PRIMUS_FLA_MLA_ATTN=1  → MLA core_attention → flash_attn_func directly
+#                             (~30 ms/MLA block on MI300X with flash-attn 2.8.3)
+#   PRIMUS_FUSED_CE=1      → FLA FusedLinearCrossEntropyLoss (chunked, big mem
+#                             + speed win regardless of mixer)
+#   PRIMUS_FLA_SWIGLU=1    → FLA Triton SwiGLU in MLP (~20 ms/iter)
+#   PRIMUS_FLA_NORM=1      → WrappedTorchNorm → FLA Triton RMSNorm (saves
+#                             ~30 ms/iter AND drops peak memory ~5 GB/rank,
+#                             unblocking the 98.8 % rocm pressure we hit on
+#                             the first run)
+#   PRIMUS_FLA_DATA=1      → drive Megatron from FLA's DistributedSampler
+#                             token order so the loss curves are directly
+#                             comparable iter-by-iter (no batch-order noise)
+#   PRIMUS_FLA_CACHE_DIR   → FLA's fineweb-edu cache that PRIMUS_FLA_DATA
+#                             reads from
+#
+# (PRIMUS_FLA_CONV is intentionally NOT exported — that flag targets GDN's
+# fused conv1d; Mamba2 has its own causal_conv1d path in upstream Megatron.)
+#
+# Run inside the rocm/primus:v26.2 container:
+#
+#   cd /home/<user>/Primus && bash tools/hybrid/launch_mamba_hybrid_300M.sh
+###############################################################################
+
+set -euo pipefail
+
+export PRIMUS_FLA_MLA_ATTN=1
+export PRIMUS_FUSED_CE=1
+export PRIMUS_FLA_SWIGLU=1
+export PRIMUS_FLA_NORM=1
+export PRIMUS_FLA_DATA=1
+export PRIMUS_FLA_CACHE_DIR=${PRIMUS_FLA_CACHE_DIR:-$HOME/flash-linear-attention/legacy/training/data/HuggingFaceFW/fineweb-edu/sample-10BT/train}
+
+EXP=${EXP:-examples/megatron/configs/MI300X/zebra_llama_300M_mamba_hybrid-pretrain.yaml}
+LOG=${LOG:-primus_mamba_hybrid_300M.log}
+
+echo "==========[launch_mamba_hybrid_300M.sh]=========="
+echo "PRIMUS_FLA_MLA_ATTN  = ${PRIMUS_FLA_MLA_ATTN}"
+echo "PRIMUS_FUSED_CE      = ${PRIMUS_FUSED_CE}"
+echo "PRIMUS_FLA_SWIGLU    = ${PRIMUS_FLA_SWIGLU}"
+echo "PRIMUS_FLA_NORM      = ${PRIMUS_FLA_NORM}"
+echo "PRIMUS_FLA_DATA      = ${PRIMUS_FLA_DATA}"
+echo "PRIMUS_FLA_CACHE_DIR = ${PRIMUS_FLA_CACHE_DIR}"
+echo "EXP                  = ${EXP}"
+echo "LOG                  = ${LOG}"
+echo "=================================================="
+
+if [ ! -d "${PRIMUS_FLA_CACHE_DIR}" ]; then
+    echo "WARNING: PRIMUS_FLA_CACHE_DIR=${PRIMUS_FLA_CACHE_DIR} does not exist."
+    echo "         FLAOrderGPTDataset will silently fall back to Megatron's GPTDataset"
+    echo "         shuffler, and the loss curve will diverge from FLA's after iter 1."
+fi
+
+EXP="${EXP}" bash examples/run_pretrain.sh 2>&1 | tee "${LOG}"
diff --git a/tools/hybrid/lm_harness_eval.py b/tools/hybrid/lm_harness_eval.py
new file mode 100644
index 000000000..e476be005
--- /dev/null
+++ b/tools/hybrid/lm_harness_eval.py
@@ -0,0 +1,74 @@
+import json
+import os
+import sys
+from pathlib import Path
+
+import lm_eval
+from lm_eval.utils import make_table
+from transformers import AutoConfig, AutoModelForCausalLM
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from modeling_zebra_llama import ZebraLlamaConfig, ZebraLlamaForCausalLM
+
+AutoConfig.register("zebra_llama", ZebraLlamaConfig)
+AutoModelForCausalLM.register(ZebraLlamaConfig, ZebraLlamaForCausalLM)
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--tokenizer", type=str, default=None)
+    parser.add_argument(
+        "--tasks", type=str, default="arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande"
+    )
+    parser.add_argument("--batch_size", type=str, default="auto")
+    parser.add_argument("--num_fewshot", type=int, default=0)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--dtype", type=str, default="bfloat16")
+    parser.add_argument("--output_path", type=str, default="eval_results")
+    parser.add_argument("--limit", type=float, default=None)
+    args = parser.parse_args()
+
+    tokenizer_path = args.tokenizer or args.model_path
+    dtype_str = args.dtype if args.dtype != "auto" else "auto"
+    model_args = (
+        f"pretrained={args.model_path},tokenizer={tokenizer_path},trust_remote_code=True,dtype={dtype_str}"
+    )
+
+    print("=" * 72)
+    print("Zebra-Llama lm-eval")
+    print("=" * 72)
+    print(f"  Model path:  {args.model_path}")
+    print(f"  Tokenizer:   {tokenizer_path}")
+    print(f"  Tasks:       {args.tasks}")
+    print(f"  Batch size:  {args.batch_size}")
+    print(f"  Num fewshot: {args.num_fewshot}")
+    print(f"  Device:      {args.device}")
+    print(f"  Dtype:       {args.dtype}")
+    print(f"  Output:      {args.output_path}")
+    print("=" * 72)
+
+    results = lm_eval.simple_evaluate(
+        model="hf",
+        model_args=model_args,
+        tasks=args.tasks.split(","),
+        batch_size=args.batch_size,
+        num_fewshot=args.num_fewshot if args.num_fewshot > 0 else None,
+        device=args.device,
+        log_samples=True,
+        limit=args.limit,
+    )
+
+    print(make_table(results))
+
+    os.makedirs(args.output_path, exist_ok=True)
+    out_file = os.path.join(args.output_path, "results.json")
+    with open(out_file, "w") as f:
+        json.dump(results["results"], f, indent=2, default=str)
+    print(f"\nResults saved to {out_file}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/hybrid/megatron_forward_zebra_llama.py b/tools/hybrid/megatron_forward_zebra_llama.py
new file mode 100644
index 000000000..57e864a07
--- /dev/null
+++ b/tools/hybrid/megatron_forward_zebra_llama.py
@@ -0,0 +1,1162 @@
+"""
+Run a single forward pass with a Megatron (mcore) Zebra-Llama checkpoint.
+
+This is intended for *numerical parity* checks against the HF implementation in
+`tools/hybrid/modeling_zebra_llama.py`:
+  - same tokenizer ids
+  - same logits for a fixed prompt (within dtype tolerance)
+
+Usage (1 GPU):
+  cd /vfs/silo/mingyyan/home_backup/Primus
+  export PYTHONPATH="$(pwd):$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"
+
+  torchrun --nproc_per_node=1 tools/hybrid/megatron_forward_zebra_llama.py \
+    --load output/zebra_llama_1B-pretrain/iter_0150000 \
+    --prompt "The capital of France is" \
+    --topk 10
+
+Notes:
+  - For TP/PP > 1, you must launch with the matching world size and pass the
+    parallelism args; this script focuses on the common TP=1, PP=1 debug case.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from argparse import ArgumentParser
+from functools import partial
+from pathlib import Path
+
+import torch
+
+
+def _setup_sys_path() -> None:
+    """Make sure Primus + Megatron are importable when run from anywhere."""
+    primus_root = Path(__file__).resolve().parent.parent.parent
+    megatron_root = primus_root / "third_party" / "Megatron-LM"
+    tools_root = primus_root / "tools" / "hybrid"
+
+    for p in (str(primus_root), str(megatron_root), str(tools_root)):
+        if p not in sys.path:
+            sys.path.insert(0, p)
+
+
+def _ensure_primus_logger(rank: int, world_size: int) -> None:
+    """
+    Primus helpers (e.g., tokenizer builder) log via `primus.core.utils.logger`.
+    In full Primus training runs, the logger is initialized by the launcher/BaseModule.
+    This standalone script must initialize it explicitly.
+    """
+    from primus.core.utils import logger as primus_logger
+    from primus.core.utils.module_utils import set_logging_rank
+
+    # Make log_rank_0 / log_rank_last behave correctly before torch.distributed init.
+    set_logging_rank(rank, world_size)
+
+    # Avoid double-init (setup_logger is call_once, but also keep it cheap).
+    if getattr(primus_logger, "_logger", None) is not None:
+        return
+
+    exp_root = os.environ.get("PRIMUS_EXP_ROOT", "/tmp/primus_megatron_forward")
+    work_group = os.environ.get("PRIMUS_TEAM", "local")
+    user_name = os.environ.get("PRIMUS_USER", os.environ.get("USER", "user"))
+    exp_name = os.environ.get("PRIMUS_EXP_NAME", "megatron_forward")
+
+    cfg = primus_logger.LoggerConfig(
+        exp_root_path=exp_root,
+        work_group=work_group,
+        user_name=user_name,
+        exp_name=exp_name,
+        module_name="megatron_forward",
+        file_sink_level="INFO",
+        stderr_sink_level="INFO",
+        node_ip=os.environ.get("MASTER_ADDR", "localhost"),
+        rank=rank,
+        world_size=world_size,
+    )
+    primus_logger.setup_logger(cfg, is_head=False)
+
+
+def _ensure_rocm_validate_args_compat(args) -> None:
+    """
+    `primus.modules.trainer.megatron.utils.validate_args_on_rocm()` assumes a Primus
+    YAML-backed args object that contains some Primus-specific flags.
+    When running this standalone script, those attributes may not exist.
+    Set safe defaults so validation can run without crashing.
+    """
+
+    def _setdefault(name: str, value) -> None:
+        if not hasattr(args, name):
+            setattr(args, name, value)
+
+    # Determinism / MoE flags
+    _setdefault("deterministic_mode", False)
+    _setdefault("moe_grouped_gemm", False)
+
+    # FP8 / turbo linear flags
+    _setdefault("fp8", False)
+    _setdefault("use_turbo_parallel_linear", False)
+    _setdefault("fp8_recipe", "tensorwise")
+
+    # Pipeline debug
+    _setdefault("dump_pp_data", False)
+
+    # PrimusTurbo / MoE extras (keep disabled)
+    _setdefault("turbo_sync_free_moe_stage", 0)
+    _setdefault("enable_primus_turbo", False)
+    _setdefault("moe_use_legacy_grouped_gemm", False)
+    _setdefault("use_turbo_deepep", False)
+    _setdefault("moe_shared_expert_overlap", False)
+    _setdefault("moe_router_dtype", "fp32")
+    _setdefault("expert_model_parallel_size", 1)
+    _setdefault("turbo_deepep_num_cu", 0)
+
+
+def _apply_zebra_defaults_if_needed(args) -> None:
+    """
+    Megatron's argparse provides many non-None defaults (e.g. hybrid_attention_ratio=0.0),
+    which means `validate_args(args, args_defaults)` will not override them.
+    In Primus training, these values come from YAML and are set correctly.
+
+    For this standalone script, if the user is using the Zebra hybrid spec and has not
+    set key Zebra knobs explicitly, apply the Zebra defaults so model construction matches
+    the training config.
+    """
+
+    zebra_spec = [
+        "primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs",
+        "hybrid_stack_spec",
+    ]
+    if not hasattr(args, "spec") or args.spec != zebra_spec:
+        return
+
+    # ------------------------------------------------------------------
+    # Zebra-Llama 1B config parity knobs (from primus/configs/models/megatron/zebra_llama_1B.yaml)
+    # Ensure these are set before model construction.
+    # ------------------------------------------------------------------
+    # Model size
+    args.num_layers = 32
+    args.hidden_size = 2048
+    args.ffn_hidden_size = 8192
+    args.num_attention_heads = 32
+
+    # Hybrid / Mamba
+    args.is_hybrid_model = True
+    args.hybrid_attention_ratio = 0.25
+    args.mamba_state_dim = 64
+    args.mamba_head_dim = 64
+    args.mamba_num_groups = 8
+
+    # MLA
+    args.group_query_attention = False
+    args.swiglu = True
+    args.num_query_groups = None
+    args.multi_latent_attention = True
+    args.q_lora_rank = 1344
+    args.kv_lora_rank = 128
+    args.qk_head_dim = 32
+    args.qk_pos_emb_head_dim = 32
+    args.v_head_dim = 64
+
+    # RoPE / position settings
+    args.normalization = "RMSNorm"
+    args.rotary_base = 500000
+    args.rotary_scaling_factor = 1.0
+    args.mscale = 1.0
+    args.mscale_all_dim = 1.0
+    args.position_embedding_type = "none"
+    args.add_position_embedding = True
+    args.use_rotary_position_embeddings = False
+    args.original_max_position_embeddings = 2048
+
+    args.add_bias_linear = False
+    args.mamba_hidden_act = "silu"
+
+    # Extra Mamba knobs (not in zebra_llama_1B.yaml but required by spec)
+    if getattr(args, "mamba_expand", None) is None:
+        args.mamba_expand = 1
+    if getattr(args, "mamba_d_conv", None) is None:
+        args.mamba_d_conv = 4
+
+
+def _normalize_load_path(args) -> None:
+    """
+    Megatron expects `--load` to point to the *root* checkpoint directory that contains
+    `latest_checkpointed_iteration.txt`. Users often pass an `iter_XXXXXXX/` subdir.
+    If so, rewrite:
+      --load <root>/iter_XXXXXXX  ->  --load <root>  and set --ckpt-step XXXX
+    """
+    load_dir = getattr(args, "load", None)
+    if not load_dir:
+        return
+    p = Path(str(load_dir)).expanduser()
+    name = p.name
+    if name.startswith("iter_"):
+        try:
+            step = int(name[len("iter_") :])
+        except ValueError:
+            return
+        # Only override if user didn't already request a specific step.
+        if getattr(args, "ckpt_step", None) in (None, 0):
+            setattr(args, "ckpt_step", step)
+        setattr(args, "load", str(p.parent))
+
+
+def _add_script_args(parser: ArgumentParser) -> ArgumentParser:
+    group = parser.add_argument_group(title="zebra forward debug")
+    group.add_argument("--prompt", type=str, default="The capital of France is")
+    group.add_argument("--topk", type=int, default=10)
+    group.add_argument("--max-prompt-tokens", type=int, default=256)
+    group.add_argument(
+        "--save-logits", type=str, default=None, help="Optional .pt path to save logits tensor."
+    )
+    group.add_argument(
+        "--hf-dir",
+        type=str,
+        default=None,
+        help="If set, also run the HuggingFace Zebra-Llama forward pass from this converted HF checkpoint dir "
+        "(must contain config.json + pytorch_model.bin) and compare logits.",
+    )
+    group.add_argument(
+        "--hf-dtype",
+        type=str,
+        default="bfloat16",
+        choices=["float32", "float16", "bfloat16"],
+        help="dtype for HF model weights/compute during comparison.",
+    )
+    group.add_argument(
+        "--compare-atol",
+        type=float,
+        default=1e-2,
+        help="Abs tolerance used only for reporting (not an assert).",
+    )
+    group.add_argument(
+        "--compare-layerwise",
+        action="store_true",
+        default=False,
+        help="If set with --hf-dir, compare intermediate hidden states layer-by-layer.",
+    )
+    group.add_argument(
+        "--compare-layer-isolated",
+        action="store_true",
+        default=False,
+        help="If set with --hf-dir, feed each layer the SAME hidden-state tensor (taken from Megatron's layer input) "
+        "and compare that layer's output. This isolates per-layer numeric drift from upstream accumulation.",
+    )
+    group.add_argument(
+        "--layerwise-token",
+        type=str,
+        default="last",
+        choices=["last", "mean"],
+        help="Which token representation to compare per layer: last token vector or mean over sequence.",
+    )
+    group.add_argument(
+        "--runtime-gather-output",
+        action="store_true",
+        default=True,
+        help="Gather full-vocab logits at runtime (useful for TP>1). Default: true.",
+    )
+    group.add_argument(
+        "--position-ids-mode",
+        type=str,
+        default="normal",
+        choices=["normal", "zeros"],
+        help="Test RoPE/positional-embedding impact by controlling position_ids. "
+        "'zeros' forces all tokens to use position 0 (often makes rotary a no-op).",
+    )
+    group.add_argument(
+        "--torch-profiler",
+        type=str,
+        default="off",
+        choices=["off", "megatron", "hf", "both"],
+        help="Enable torch.profiler tracing for selected forward pass(es). Exports Chrome trace JSON.",
+    )
+    group.add_argument(
+        "--torch-profiler-dir",
+        type=str,
+        default=None,
+        help="Directory to write torch.profiler traces to. Default: ./profiler_traces",
+    )
+    group.add_argument(
+        "--torch-profiler-all-ranks",
+        action="store_true",
+        default=False,
+        help="If set, write one trace per rank. Default: only rank0 writes traces.",
+    )
+    group.add_argument(
+        "--torch-profiler-record-shapes",
+        action="store_true",
+        default=False,
+        help="Record operator input shapes (larger traces).",
+    )
+    group.add_argument(
+        "--torch-profiler-memory",
+        action="store_true",
+        default=False,
+        help="Record memory usage (larger traces).",
+    )
+    group.add_argument(
+        "--torch-profiler-with-stack",
+        action="store_true",
+        default=False,
+        help="Record Python stacks (much larger/slower).",
+    )
+    return parser
+
+
+def _profile_forward(
+    *,
+    enabled: bool,
+    name: str,
+    out_path: Path,
+    record_shapes: bool,
+    profile_memory: bool,
+    with_stack: bool,
+    fn,
+):
+    """Run `fn()` under torch.profiler and export a Chrome trace JSON."""
+    if not enabled:
+        return fn()
+
+    from torch.profiler import ProfilerActivity, profile, record_function
+
+    activities = [ProfilerActivity.CPU]
+    # On ROCm, torch.cuda APIs still work and profiler uses CUDA activity to mean GPU.
+    if torch.cuda.is_available():
+        activities.append(ProfilerActivity.CUDA)
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+    with profile(
+        activities=activities,
+        record_shapes=record_shapes,
+        profile_memory=profile_memory,
+        with_stack=with_stack,
+    ) as prof:
+        with record_function(name):
+            out = fn()
+
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    prof.export_chrome_trace(str(out_path))
+    return out
+
+
+@torch.inference_mode()
+def main() -> None:
+    _setup_sys_path()
+
+    # Megatron imports (after sys.path is set)
+    import megatron
+    from mamba_builders import mamba_builder
+    from megatron.training import get_args, get_model, get_tokenizer, print_rank_0
+    from megatron.training.arguments import parse_args, validate_args
+    from megatron.training.checkpointing import checkpoint_exists, load_checkpoint
+    from megatron.training.global_vars import set_global_variables
+    from megatron.training.initialize import (
+        _init_autoresume,
+        _initialize_distributed,
+        _set_random_seed,
+        setup_logging,
+    )
+    from megatron.training.utils import get_ltor_masks_and_position_ids
+
+    # Builders for MCore Mamba/hybrid models
+    from model_provider import model_provider
+
+    from primus.backends.megatron.training.global_vars import (
+        set_primus_global_variables,
+    )
+
+    # Primus adds a thin wrapper around Megatron tokenizer building (used in training).
+    from primus.backends.megatron.training.tokenizer.tokenizer import build_tokenizer
+    from primus.modules.trainer.megatron.utils import (
+        set_wandb_writer_patch,
+        validate_args_on_rocm,
+    )
+
+    # ---------------------------------------------------------------------
+    # Primus-style initialization (matches MegatronTrainer.initialize_megatron)
+    # ---------------------------------------------------------------------
+    args_defaults = {
+        # inference-ish defaults
+        "no_load_rng": True,
+        "no_load_optim": True,
+        "micro_batch_size": 1,
+        "global_batch_size": 1,
+        "exit_on_missing_checkpoint": True,
+        # Primus training default: skip compile_dependencies (avoids CUDA-only nvcc fused kernels)
+        "disable_compile_dependencies": True,
+        # zebra_llama_1B defaults (can be overridden from CLI)
+        "tensor_model_parallel_size": 1,
+        "pipeline_model_parallel_size": 1,
+        "num_layers": 32,
+        "hidden_size": 2048,
+        "ffn_hidden_size": 8192,
+        "num_attention_heads": 32,
+        "seq_length": 2048,
+        "max_position_embeddings": 2048,
+        "position_embedding_type": "none",
+        "normalization": "RMSNorm",
+        "tokenizer_type": "HuggingFaceTokenizer",
+        "tokenizer_model": "meta-llama/Llama-3.2-1B",
+        # Hybrid + MLA
+        "is_hybrid_model": True,
+        "hybrid_attention_ratio": 0.25,
+        "multi_latent_attention": True,
+        "q_lora_rank": 1344,
+        "kv_lora_rank": 128,
+        "qk_head_dim": 32,
+        "qk_pos_emb_head_dim": 32,
+        "v_head_dim": 64,
+        # Mamba
+        "mamba_state_dim": 64,
+        "mamba_head_dim": 64,
+        "mamba_num_groups": 8,
+        "mamba_expand": 1,
+        "mamba_d_conv": 4,
+        "mamba_hidden_act": "silu",
+        # Spec that defines the hybrid block layout
+        "spec": [
+            "primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs",
+            "hybrid_stack_spec",
+        ],
+    }
+
+    args = parse_args(extra_args_provider=_add_script_args, ignore_unknown_args=False)
+    _normalize_load_path(args)
+    validate_args(args, args_defaults)
+    _apply_zebra_defaults_if_needed(args)
+
+    # Primus utilities used below require logger to be initialized.
+    _ensure_primus_logger(rank=int(args.rank), world_size=int(args.world_size))
+
+    # Monkey-patch wandb writer hook (Primus does this before set_global_variables).
+    megatron.training.global_vars._set_wandb_writer = set_wandb_writer_patch
+
+    # Global vars (args/timers/etc), but build tokenizer ourselves (Primus wrapper).
+    set_global_variables(args, build_tokenizer=False)
+    set_primus_global_variables(args)
+    args = get_args()
+
+    # Build and register tokenizer the same way Primus does.
+    import megatron.training.global_vars as global_vars
+
+    global_vars._ensure_var_is_not_initialized(global_vars._GLOBAL_TOKENIZER, "tokenizer")
+    global_vars._GLOBAL_TOKENIZER = build_tokenizer(args)
+
+    setup_logging()
+
+    # Distributed init + seeding (Primus finish_mpu_init path, without compile_dependencies).
+    _initialize_distributed(None, None, None)
+    _set_random_seed(
+        args.seed,
+        args.data_parallel_random_init,
+        args.te_rng_tracker,
+        args.inference_rng_tracker,
+        use_cudagraphable_rng=bool(getattr(args, "enable_cuda_graph", False))
+        or bool(getattr(args, "external_cuda_graph", False)),
+    )
+    _init_autoresume()
+
+    # Mirror Primus trainer extra ROCm validations.
+    _ensure_rocm_validate_args_compat(args)
+    validate_args_on_rocm(args)
+
+    # Build model and load checkpoint
+    model_list = get_model(partial(model_provider, mamba_builder), wrap_with_ddp=False)
+
+    # Prove whether we actually load weights (common confusion is passing --load=.../iter_XXXXXXX).
+    def _fingerprint_first_param(m) -> tuple[float, float]:
+        p0 = next(m.parameters())
+        x = p0.detach().float()
+        return float(x.mean().item()), float(x.std().item())
+
+    if torch.distributed.get_rank() == 0:
+        print_rank_0(f"Checkpoint --load: {getattr(args, 'load', None)}")
+        if getattr(args, "load", None) is not None:
+            print_rank_0(f"Checkpoint tracker exists? {checkpoint_exists(getattr(args, 'load'))}")
+            tracker = Path(str(getattr(args, "load"))) / "latest_checkpointed_iteration.txt"
+            if tracker.exists():
+                print_rank_0(f"latest_checkpointed_iteration.txt: {tracker.read_text().strip()!r}")
+            if getattr(args, "ckpt_step", None):
+                print_rank_0(f"--ckpt-step: {getattr(args, 'ckpt_step')}")
+
+    fp_before = _fingerprint_first_param(model_list[0])
+    it, _ = load_checkpoint(model_list, None, None, strict=False)
+    fp_after = _fingerprint_first_param(model_list[0])
+    if torch.distributed.get_rank() == 0:
+        print_rank_0(f"load_checkpoint() iteration={it}")
+        print_rank_0(f"param_fingerprint mean/std before: {fp_before[0]:.6g}/{fp_before[1]:.6g}")
+        print_rank_0(f"param_fingerprint mean/std after:  {fp_after[0]:.6g}/{fp_after[1]:.6g}")
+    model = model_list[0]
+    model.eval()
+
+    # Build tokenizer
+    if args.legacy_tokenizer:
+        tokenizer = get_tokenizer()
+    else:
+        tokenizer = build_tokenizer(args)
+
+    prompt: str = getattr(args, "prompt")
+    topk: int = int(getattr(args, "topk"))
+    max_prompt_tokens: int = int(getattr(args, "max_prompt_tokens"))
+
+    # Tokenize (truncate for safety)
+    token_ids = tokenizer.tokenize(prompt)
+    if len(token_ids) > max_prompt_tokens:
+        token_ids = token_ids[:max_prompt_tokens]
+
+    input_ids = torch.tensor([token_ids], device=torch.cuda.current_device(), dtype=torch.long)
+
+    # Create standard causal attention mask + position_ids (no resets)
+    eos = getattr(tokenizer, "eos_id", None)
+    pad = getattr(tokenizer, "pad_id", None)
+    if eos is None:
+        eos = 0
+    if pad is None:
+        pad = eos
+
+    attention_mask, _loss_mask, position_ids = get_ltor_masks_and_position_ids(
+        data=input_ids,
+        eod_token=eos,
+        pad_token=pad,
+        reset_position_ids=False,
+        reset_attention_mask=False,
+        eod_mask_loss=False,
+        pad_mask_loss=False,
+    )
+
+    # Optional: control position_ids to isolate rotary/positional effects.
+    if str(getattr(args, "position_ids_mode", "normal")) == "zeros":
+        position_ids = torch.zeros_like(position_ids)
+
+    # Forward (logits)
+    prof_mode = str(getattr(args, "torch_profiler", "off"))
+    prof_dir = (
+        Path(str(getattr(args, "torch_profiler_dir", None) or "profiler_traces")).expanduser().resolve()
+    )
+    rank = int(torch.distributed.get_rank()) if torch.distributed.is_initialized() else 0
+    do_profile_rank = bool(getattr(args, "torch_profiler_all_ranks", False)) or (rank == 0)
+
+    logits = _profile_forward(
+        enabled=do_profile_rank and prof_mode in ("megatron", "both"),
+        name="megatron_forward_logits",
+        out_path=prof_dir / f"megatron_rank{rank}.json",
+        record_shapes=bool(getattr(args, "torch_profiler_record_shapes", False)),
+        profile_memory=bool(getattr(args, "torch_profiler_memory", False)),
+        with_stack=bool(getattr(args, "torch_profiler_with_stack", False)),
+        fn=lambda: model(
+            input_ids,
+            position_ids,
+            attention_mask,
+            labels=None,
+            runtime_gather_output=bool(getattr(args, "runtime_gather_output", True)),
+        ),
+    )
+
+    # logits: [b, s, vocab]
+    if torch.distributed.get_rank() == 0:
+        print_rank_0(f"Prompt tokens: {input_ids.shape[1]}")
+        print_rank_0(f"Logits shape:  {tuple(logits.shape)}")
+
+        # Megatron may use a padded vocab size; restrict to the *actual* tokenizer vocab
+        # so top-k and comparisons are meaningful.
+        tok_vocab_size = getattr(tokenizer, "vocab_size", None)
+        if tok_vocab_size is None:
+            tok_vocab_size = logits.shape[-1]
+        vocab_limit = min(int(tok_vocab_size), int(logits.shape[-1]))
+        print_rank_0(f"Tokenizer vocab_size: {int(tok_vocab_size)} (using logits[:{vocab_limit}])")
+
+        last_logits = logits[0, -1, :vocab_limit]
+        vals, idxs = torch.topk(last_logits, k=min(topk, last_logits.numel()))
+        print_rank_0("Top-k next-token candidates:")
+        for v, i in zip(vals.tolist(), idxs.tolist()):
+            # tokenizer.detokenize() may drop special tokens; still useful as a hint.
+            try:
+                piece = tokenizer.detokenize([i])
+            except Exception:
+                piece = ""
+            print_rank_0(f"  id={i:6d}  logit={v: .6f}  text={piece!r}")
+
+        # Optional HF comparison (requires a converted HF checkpoint dir)
+        hf_dir = getattr(args, "hf_dir", None)
+        if hf_dir:
+            from modeling_zebra_llama import ZebraLlamaConfig, ZebraLlamaForCausalLM
+            from transformers import AutoTokenizer
+
+            hf_dir_path = Path(str(hf_dir)).expanduser().resolve()
+            cfg_path = hf_dir_path / "config.json"
+            weights_path = hf_dir_path / "pytorch_model.bin"
+            if not cfg_path.exists() or not weights_path.exists():
+                raise FileNotFoundError(
+                    f"--hf-dir must contain config.json and pytorch_model.bin. Got: {hf_dir_path}"
+                )
+
+            cfg_dict = __import__("json").loads(cfg_path.read_text())
+            hf_config = ZebraLlamaConfig(**cfg_dict)
+
+            hf_model = ZebraLlamaForCausalLM(hf_config).eval()
+            # `use_return_dict` is a read-only property in Transformers configs.
+            # We pass `return_dict=True` in the forward call below; also try setting
+            # `return_dict` for any code paths that consult config defaults.
+            try:
+                hf_model.config.return_dict = True
+            except Exception:
+                pass
+
+            state = torch.load(weights_path, map_location="cpu", weights_only=True)
+            missing, unexpected = hf_model.load_state_dict(state, strict=False)
+            if missing:
+                print_rank_0(f"[HF warn] Missing keys: {len(missing)} (showing first 10)")
+                for k in missing[:10]:
+                    print_rank_0(f"  - {k}")
+            if unexpected:
+                print_rank_0(f"[HF warn] Unexpected keys: {len(unexpected)} (showing first 10)")
+                for k in unexpected[:10]:
+                    print_rank_0(f"  - {k}")
+
+            hf_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True)
+            hf_ids = hf_tokenizer(prompt, add_special_tokens=False).input_ids
+            if len(hf_ids) > max_prompt_tokens:
+                hf_ids = hf_ids[:max_prompt_tokens]
+
+            if hf_ids != token_ids:
+                print_rank_0("[HF compare] Tokenization mismatch between Megatron and HF tokenizers!")
+                print_rank_0(f"  Megatron ids[:32]: {token_ids[:32]}")
+                print_rank_0(f"  HF ids[:32]:       {hf_ids[:32]}")
+
+            # Build HF inputs from the *Megatron* ids to ensure identical tokens.
+            hf_input_ids = torch.tensor([token_ids], device=torch.cuda.current_device(), dtype=torch.long)
+            hf_attention_mask = torch.ones_like(hf_input_ids, dtype=torch.long)
+            hf_position_ids = torch.arange(
+                hf_input_ids.shape[1], device=hf_input_ids.device, dtype=torch.long
+            ).unsqueeze(0)
+            if str(getattr(args, "position_ids_mode", "normal")) == "zeros":
+                hf_position_ids = torch.zeros_like(hf_position_ids)
+
+            dtype_name = str(getattr(args, "hf_dtype", "bfloat16"))
+            dtype_map = {
+                "float32": torch.float32,
+                "float16": torch.float16,
+                "bfloat16": torch.bfloat16,
+            }
+            hf_model = hf_model.to(device=torch.device("cuda"), dtype=dtype_map[dtype_name])
+
+            def _extract_hidden(x):
+                """
+                Extract a tensor-like hidden state from various module outputs.
+                Megatron modules sometimes return tuples where the first element is not the hidden state.
+                """
+                # Common: WrappedTensor-like object.
+                if hasattr(x, "tensor") and isinstance(getattr(x, "tensor"), torch.Tensor):
+                    return x.tensor
+                # Common: tuple outputs (pick the first tensor-ish entry).
+                if isinstance(x, tuple):
+                    for y in x:
+                        if hasattr(y, "tensor") and isinstance(getattr(y, "tensor"), torch.Tensor):
+                            return y.tensor
+                        if isinstance(y, torch.Tensor):
+                            return y
+                    return x
+                return x
+
+            def _to_bsh(x: torch.Tensor, seq_len: int, batch_size: int = 1) -> torch.Tensor:
+                # Normalize hidden states to [B,S,H] for comparison.
+                if x.dim() != 3:
+                    raise ValueError(f"Expected 3D hidden states, got shape={tuple(x.shape)}")
+                # Heuristic: Megatron often uses [S,B,H], HF uses [B,S,H]
+                if x.shape[0] == seq_len and x.shape[1] == batch_size:
+                    return x.transpose(0, 1).contiguous()
+                return x
+
+            def _repr_from_bsh(x_bsh: torch.Tensor, mode: str) -> torch.Tensor:
+                if mode == "last":
+                    return x_bsh[:, -1, :]
+                # mode == "mean"
+                return x_bsh.mean(dim=1)
+
+            class _StopForward(Exception):
+                """Internal control-flow to stop after a specific layer."""
+
+            # Optional: layerwise comparison of intermediate activations.
+            if bool(getattr(args, "compare_layerwise", False)):
+                print_rank_0("Layerwise comparison enabled: collecting intermediate activations…")
+                seq_len = int(hf_input_ids.shape[1])
+
+                mg_layer_vecs: list[torch.Tensor] = []
+                mg_layer_names: list[str] = []
+
+                def _mk_mg_hook(name: str):
+                    def _hook(_m, _inp, out):
+                        h = _extract_hidden(out)
+                        if not isinstance(h, torch.Tensor):
+                            return
+                        h_bsh = _to_bsh(h, seq_len=seq_len, batch_size=1).float()
+                        mg_layer_vecs.append(
+                            _repr_from_bsh(h_bsh, getattr(args, "layerwise_token", "last")).cpu()
+                        )
+                        mg_layer_names.append(name)
+
+                    return _hook
+
+                hf_layer_vecs: list[torch.Tensor] = []
+                hf_layer_names: list[str] = []
+
+                def _mk_hf_hook(name: str):
+                    def _hook(_m, _inp, out):
+                        h = _extract_hidden(out)
+                        if not isinstance(h, torch.Tensor):
+                            return
+                        h_bsh = _to_bsh(h, seq_len=seq_len, batch_size=1).float()
+                        hf_layer_vecs.append(
+                            _repr_from_bsh(h_bsh, getattr(args, "layerwise_token", "last")).cpu()
+                        )
+                        hf_layer_names.append(name)
+
+                    return _hook
+
+                mg_hooks = []
+                try:
+                    mg_layers = getattr(model, "decoder", None)
+                    mg_layers = getattr(mg_layers, "layers", None)
+                    if mg_layers is None:
+                        raise RuntimeError("Megatron model.decoder.layers not found for hooks.")
+                    for i, layer in enumerate(mg_layers):
+                        mg_hooks.append(
+                            layer.register_forward_hook(_mk_mg_hook(f"mg[{i}]:{layer.__class__.__name__}"))
+                        )
+                except Exception as e:
+                    print_rank_0(f"[Layerwise] Failed to attach Megatron hooks: {e}")
+
+                hf_hooks = []
+                try:
+                    hf_layers = getattr(hf_model, "model", None)
+                    hf_layers = getattr(hf_layers, "layers", None)
+                    if hf_layers is None:
+                        raise RuntimeError("HF model.model.layers not found for hooks.")
+                    for i, layer in enumerate(hf_layers):
+                        hf_hooks.append(
+                            layer.register_forward_hook(_mk_hf_hook(f"hf[{i}]:{layer.__class__.__name__}"))
+                        )
+                except Exception as e:
+                    print_rank_0(f"[Layerwise] Failed to attach HF hooks: {e}")
+
+                # Run forwards again to collect activations
+                _ = model(
+                    input_ids,
+                    position_ids,
+                    attention_mask,
+                    labels=None,
+                    runtime_gather_output=bool(getattr(args, "runtime_gather_output", True)),
+                )
+                _ = hf_model(
+                    input_ids=hf_input_ids,
+                    attention_mask=hf_attention_mask,
+                    position_ids=hf_position_ids,
+                    use_cache=False,
+                    return_dict=True,
+                )
+
+                for h in mg_hooks:
+                    try:
+                        h.remove()
+                    except Exception:
+                        pass
+                for h in hf_hooks:
+                    try:
+                        h.remove()
+                    except Exception:
+                        pass
+
+                n = min(len(mg_layer_vecs), len(hf_layer_vecs))
+                print_rank_0(
+                    f"Layerwise vectors collected: mg={len(mg_layer_vecs)} hf={len(hf_layer_vecs)} compare_n={n}"
+                )
+                if n > 0:
+                    print_rank_0("Per-layer diff (vector max/mean abs, cosine):")
+                    for i in range(n):
+                        a = mg_layer_vecs[i].squeeze(0)
+                        b = hf_layer_vecs[i].squeeze(0)
+                        d = (a - b).abs()
+                        # cosine similarity
+                        cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
+                        print_rank_0(
+                            f"  {i:03d} {mg_layer_names[i]}  vs  {hf_layer_names[i]} | "
+                            f"max={d.max().item():.6g} mean={d.mean().item():.6g} cos={cos:.6g}"
+                        )
+
+            # Optional: pre/post of first norm inside each block.
+            if bool(getattr(args, "compare_prenorm", False)):
+                print_rank_0("Pre/Post first-norm comparison enabled: collecting norm IO…")
+                seq_len = int(hf_input_ids.shape[1])
+                mode = str(getattr(args, "layerwise_token", "last"))
+
+                mg_pre, mg_post, mg_names = [], [], []
+                hf_pre, hf_post, hf_names = [], [], []
+
+                def _mk_io_hook(store_pre, store_post, store_name, name: str):
+                    def _hook(_m, inp, out):
+                        if not inp:
+                            return
+                        x_in = _extract_hidden(inp[0])
+                        x_out = _extract_hidden(out)
+                        if not (isinstance(x_in, torch.Tensor) and isinstance(x_out, torch.Tensor)):
+                            return
+                        x_in_bsh = _to_bsh(x_in, seq_len=seq_len, batch_size=1).float()
+                        x_out_bsh = _to_bsh(x_out, seq_len=seq_len, batch_size=1).float()
+                        store_pre.append(_repr_from_bsh(x_in_bsh, mode).cpu())
+                        store_post.append(_repr_from_bsh(x_out_bsh, mode).cpu())
+                        store_name.append(name)
+
+                    return _hook
+
+                mg_io_hooks = []
+                try:
+                    mg_layers = getattr(model, "decoder", None)
+                    mg_layers = getattr(mg_layers, "layers", None)
+                    for i, layer in enumerate(mg_layers):
+                        nname, nmod = _find_first_norm_like(layer)
+                        if nmod is None:
+                            continue
+                        mg_io_hooks.append(
+                            nmod.register_forward_hook(
+                                _mk_io_hook(mg_pre, mg_post, mg_names, f"mg[{i}].{nname}")
+                            )
+                        )
+                except Exception as e:
+                    print_rank_0(f"[PreNorm] Failed to attach Megatron norm hooks: {e}")
+
+                hf_io_hooks = []
+                try:
+                    hf_layers = getattr(hf_model, "model", None)
+                    hf_layers = getattr(hf_layers, "layers", None)
+                    for i, layer in enumerate(hf_layers):
+                        nname, nmod = _find_first_norm_like(layer)
+                        if nmod is None:
+                            continue
+                        hf_io_hooks.append(
+                            nmod.register_forward_hook(
+                                _mk_io_hook(hf_pre, hf_post, hf_names, f"hf[{i}].{nname}")
+                            )
+                        )
+                except Exception as e:
+                    print_rank_0(f"[PreNorm] Failed to attach HF norm hooks: {e}")
+
+                # Re-run forwards to collect norm IO.
+                _ = model(
+                    input_ids,
+                    position_ids,
+                    attention_mask,
+                    labels=None,
+                    runtime_gather_output=bool(getattr(args, "runtime_gather_output", True)),
+                )
+                _ = hf_model(
+                    input_ids=hf_input_ids,
+                    attention_mask=hf_attention_mask,
+                    position_ids=hf_position_ids,
+                    use_cache=False,
+                    return_dict=True,
+                )
+
+                for h in mg_io_hooks:
+                    try:
+                        h.remove()
+                    except Exception:
+                        pass
+                for h in hf_io_hooks:
+                    try:
+                        h.remove()
+                    except Exception:
+                        pass
+
+                n = min(len(mg_pre), len(hf_pre), len(mg_post), len(hf_post))
+                print_rank_0(
+                    f"Pre/post norm vectors collected: mg={len(mg_pre)} hf={len(hf_pre)} compare_n={n}"
+                )
+                if n > 0:
+                    print_rank_0("Per-norm diff (PRE then POST): max/mean abs, cosine")
+                    for i in range(n):
+                        a0 = mg_pre[i].squeeze(0)
+                        b0 = hf_pre[i].squeeze(0)
+                        # PRE should generally be comparable (both are the incoming hidden state vector).
+                        d0 = (a0 - b0).abs()
+                        cos0 = torch.nn.functional.cosine_similarity(a0, b0, dim=0).item()
+                        a1 = mg_post[i].squeeze(0)
+                        b1 = hf_post[i].squeeze(0)
+                        # POST may be non-comparable if Megatron uses a fused LN+Linear module
+                        # (output dim != hidden_size). In that case, report PRE only and skip POST.
+                        if a1.numel() != b1.numel():
+                            print_rank_0(
+                                f"  {i:03d} {mg_names[i]} vs {hf_names[i]} | "
+                                f"PRE max={d0.max().item():.6g} mean={d0.mean().item():.6g} cos={cos0:.6g} || "
+                                f"POST skipped (dim mismatch: mg={a1.numel()} hf={b1.numel()})"
+                            )
+                        else:
+                            d1 = (a1 - b1).abs()
+                            cos1 = torch.nn.functional.cosine_similarity(a1, b1, dim=0).item()
+                            print_rank_0(
+                                f"  {i:03d} {mg_names[i]} vs {hf_names[i]} | "
+                                f"PRE max={d0.max().item():.6g} mean={d0.mean().item():.6g} cos={cos0:.6g} || "
+                                f"POST max={d1.max().item():.6g} mean={d1.mean().item():.6g} cos={cos1:.6g}"
+                            )
+
+            # Optional: isolated per-layer compare (force identical input hidden per layer).
+            if bool(getattr(args, "compare_layer_isolated", False)):
+                print_rank_0(
+                    "Isolated layer comparison enabled: capturing Megatron per-layer inputs/outputs…"
+                )
+                seq_len = int(hf_input_ids.shape[1])
+                mode = str(getattr(args, "layerwise_token", "last"))
+
+                mg_in_bsh_cpu: dict[int, torch.Tensor] = {}
+                mg_out_bsh_cpu: dict[int, torch.Tensor] = {}
+                mg_names: dict[int, str] = {}
+
+                # Capture Megatron layer input+output in a single forward pass.
+                mg_io_hooks = []
+                try:
+                    mg_layers = getattr(model, "decoder", None)
+                    mg_layers = getattr(mg_layers, "layers", None)
+                    if mg_layers is None:
+                        raise RuntimeError("Megatron model.decoder.layers not found for hooks.")
+
+                    def _mk_mg_pre_hook(i: int, name: str):
+                        def _pre(_m, inp, kwargs=None):
+                            # Megatron may call layers with positional args or kwargs.
+                            x_src = None
+                            if inp:
+                                x_src = inp[0]
+                            elif kwargs:
+                                # Prefer common hidden-state kwarg names.
+                                for k in ("hidden_states", "hidden", "x", "input", "input_tensor"):
+                                    if k in kwargs:
+                                        x_src = kwargs[k]
+                                        break
+                                if x_src is None:
+                                    # Fallback: first tensor-ish value.
+                                    for v in kwargs.values():
+                                        if isinstance(v, torch.Tensor) or (
+                                            hasattr(v, "tensor")
+                                            and isinstance(getattr(v, "tensor"), torch.Tensor)
+                                        ):
+                                            x_src = v
+                                            break
+                            if x_src is None:
+                                return
+
+                            x = _extract_hidden(x_src)
+                            if not isinstance(x, torch.Tensor):
+                                return
+                            x_bsh = _to_bsh(x, seq_len=seq_len, batch_size=1)
+                            # store float32 on CPU for stable injection + low GPU memory
+                            mg_in_bsh_cpu[i] = x_bsh.detach().float().cpu()
+                            mg_names[i] = name
+
+                        return _pre
+
+                    def _mk_mg_post_hook(_i: int):
+                        def _post(_m, _inp, out):
+                            x = _extract_hidden(out)
+                            if not isinstance(x, torch.Tensor):
+                                return
+                            x_bsh = _to_bsh(x, seq_len=seq_len, batch_size=1)
+                            mg_out_bsh_cpu[_i] = x_bsh.detach().float().cpu()
+
+                        return _post
+
+                    for i, layer in enumerate(mg_layers):
+                        name = f"mg[{i}]:{layer.__class__.__name__}"
+                        # Some Megatron layers are invoked with kwargs; capture those too.
+                        try:
+                            mg_io_hooks.append(
+                                layer.register_forward_pre_hook(_mk_mg_pre_hook(i, name), with_kwargs=True)
+                            )
+                        except TypeError:
+                            mg_io_hooks.append(layer.register_forward_pre_hook(_mk_mg_pre_hook(i, name)))
+                        mg_io_hooks.append(layer.register_forward_hook(_mk_mg_post_hook(i)))
+                except Exception as e:
+                    print_rank_0(f"[Isolated] Failed to attach Megatron IO hooks: {e}")
+
+                # Run Megatron once to populate mg_in_bsh_cpu / mg_out_bsh_cpu.
+                _ = model(
+                    input_ids,
+                    position_ids,
+                    attention_mask,
+                    labels=None,
+                    runtime_gather_output=bool(getattr(args, "runtime_gather_output", True)),
+                )
+
+                for h in mg_io_hooks:
+                    try:
+                        h.remove()
+                    except Exception:
+                        pass
+
+                common = sorted(set(mg_in_bsh_cpu.keys()) & set(mg_out_bsh_cpu.keys()))
+                print_rank_0(
+                    f"[Isolated] Captured Megatron IO: pre={len(mg_in_bsh_cpu)} post={len(mg_out_bsh_cpu)} common={len(common)}"
+                )
+                if len(common) == 0:
+                    print_rank_0("[Isolated] No Megatron layers captured; skipping isolated compare.")
+                else:
+                    # Run HF once per layer, overriding that layer's input with Megatron's captured input.
+                    hf_layers = getattr(hf_model, "model", None)
+                    hf_layers = getattr(hf_layers, "layers", None)
+                    if hf_layers is None:
+                        raise RuntimeError("HF model.model.layers not found for hooks.")
+
+                    # Compare only up to common layer count.
+                    n = min(len(common), len(hf_layers))
+                    common = common[:n]
+                    print_rank_0(
+                        f"[Isolated] Comparing {n} layers (mg_common={len(common)}, hf={len(hf_layers)})."
+                    )
+                    print_rank_0("Per-layer isolated diff (output vector max/mean abs, cosine):")
+
+                    hf_param_dtype = next(hf_model.parameters()).dtype
+                    hf_device = next(hf_model.parameters()).device
+
+                    for i in common:
+                        inj_bsh = mg_in_bsh_cpu[i].to(device=hf_device, dtype=hf_param_dtype)
+                        got_out: dict[str, torch.Tensor] = {}
+
+                        def _hf_pre(_m, inp, kwargs=None):
+                            # HF layers are usually called positionally, but support kwargs for safety.
+                            if inp:
+                                new_args = (inj_bsh,) + tuple(inp[1:])
+                                # When registered with `with_kwargs=True`, must return (new_args, new_kwargs).
+                                if kwargs is not None:
+                                    return new_args, kwargs
+                                return new_args
+                            if kwargs is not None:
+                                kwargs = dict(kwargs)
+                                # Try common hidden-state names.
+                                for k in ("hidden_states", "hidden", "x", "input", "input_tensor"):
+                                    if k in kwargs:
+                                        kwargs[k] = inj_bsh
+                                        return (), kwargs
+                                # If we couldn't find a name to overwrite, still satisfy hook contract.
+                                return (), kwargs
+                            return inp
+
+                        def _hf_post(_m, _inp, out):
+                            x = _extract_hidden(out)
+                            if isinstance(x, torch.Tensor):
+                                got_out["out"] = x
+                            raise _StopForward()
+
+                        try:
+                            h_pre = hf_layers[i].register_forward_pre_hook(_hf_pre, with_kwargs=True)
+                        except TypeError:
+                            h_pre = hf_layers[i].register_forward_pre_hook(_hf_pre)
+                        h_post = hf_layers[i].register_forward_hook(_hf_post)
+                        try:
+                            try:
+                                _ = hf_model(
+                                    input_ids=hf_input_ids,
+                                    attention_mask=hf_attention_mask,
+                                    position_ids=hf_position_ids,
+                                    use_cache=False,
+                                    return_dict=True,
+                                )
+                            except _StopForward:
+                                pass
+
+                            if "out" not in got_out:
+                                print_rank_0(f"  {i:03d} {mg_names[i]}  vs  hf[{i}] | no output captured")
+                                continue
+
+                            hf_out = _extract_hidden(got_out["out"])
+                            if not isinstance(hf_out, torch.Tensor):
+                                print_rank_0(f"  {i:03d} {mg_names[i]}  vs  hf[{i}] | non-tensor output")
+                                continue
+
+                            hf_out_bsh = _to_bsh(hf_out, seq_len=seq_len, batch_size=1).float().cpu()
+                            mg_out_bsh = mg_out_bsh_cpu[i]
+
+                            a = _repr_from_bsh(mg_out_bsh, mode).squeeze(0)
+                            b = _repr_from_bsh(hf_out_bsh, mode).squeeze(0)
+                            d = (a - b).abs()
+                            cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
+                            print_rank_0(
+                                f"  {i:03d} {mg_names.get(i, f'mg[{i}]')}  vs  hf[{i}]:{hf_layers[i].__class__.__name__} | "
+                                f"max={d.max().item():.6g} mean={d.mean().item():.6g} cos={cos:.6g}"
+                            )
+                        finally:
+                            try:
+                                h_pre.remove()
+                            except Exception:
+                                pass
+                            try:
+                                h_post.remove()
+                            except Exception:
+                                pass
+
+            hf_out = _profile_forward(
+                enabled=do_profile_rank and prof_mode in ("hf", "both"),
+                name="hf_forward_logits",
+                out_path=prof_dir / f"hf_rank{rank}.json",
+                record_shapes=bool(getattr(args, "torch_profiler_record_shapes", False)),
+                profile_memory=bool(getattr(args, "torch_profiler_memory", False)),
+                with_stack=bool(getattr(args, "torch_profiler_with_stack", False)),
+                fn=lambda: hf_model(
+                    input_ids=hf_input_ids,
+                    attention_mask=hf_attention_mask,
+                    position_ids=hf_position_ids,
+                    use_cache=False,
+                    return_dict=True,
+                ),
+            )
+            hf_logits = hf_out.logits  # [b, s, vocab]
+
+            # Compare on shared vocab range (Megatron vocab may be padded)
+            v = min(hf_logits.shape[-1], last_logits.shape[-1])
+            mg_last = last_logits[:v].float()
+            hf_last = hf_logits[0, -1, :v].float()
+            diff = (mg_last - hf_last).abs()
+
+            atol = float(getattr(args, "compare_atol", 1e-2))
+            print_rank_0("HF comparison (last-token logits):")
+            print_rank_0(f"  shared_vocab={v}")
+            print_rank_0(f"  max_abs_diff={diff.max().item():.6g}")
+            print_rank_0(f"  mean_abs_diff={diff.mean().item():.6g}")
+            print_rank_0(f"  pct(|diff|<=atol {atol:g})={(diff <= atol).float().mean().item()*100:.2f}%")
+
+            mg_top1 = int(torch.argmax(mg_last).item())
+            hf_top1 = int(torch.argmax(hf_last).item())
+            print_rank_0(f"  top1_match={mg_top1 == hf_top1}  (mg={mg_top1}, hf={hf_top1})")
+
+            mg_vals, mg_idxs = torch.topk(mg_last, k=min(topk, mg_last.numel()))
+            hf_vals, hf_idxs = torch.topk(hf_last, k=min(topk, hf_last.numel()))
+            print_rank_0("  top-k (Megatron vs HF):")
+            for rank, (mvi, mii, hvi, hii) in enumerate(
+                zip(mg_vals.tolist(), mg_idxs.tolist(), hf_vals.tolist(), hf_idxs.tolist()),
+                start=1,
+            ):
+                try:
+                    mtxt = hf_tokenizer.decode([mii])
+                except Exception:
+                    mtxt = ""
+                try:
+                    htxt = hf_tokenizer.decode([hii])
+                except Exception:
+                    htxt = ""
+                print_rank_0(
+                    f"   #{rank:02d}  mg id={mii:6d} logit={mvi: .6f} text={mtxt!r} | "
+                    f"hf id={hii:6d} logit={hvi: .6f} text={htxt!r}"
+                )
+
+        if getattr(args, "save_logits", None):
+            out_path = Path(str(getattr(args, "save_logits"))).expanduser()
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+            torch.save({"prompt": prompt, "input_ids": input_ids.cpu(), "logits": logits.cpu()}, out_path)
+            print_rank_0(f"Saved logits to: {str(out_path)}")
+
+
+if __name__ == "__main__":
+    # Megatron relies on torchrun/torch.distributed init; initialize_megatron handles it.
+    main()
diff --git a/tools/hybrid/modeling_zebra_llama.py b/tools/hybrid/modeling_zebra_llama.py
new file mode 100644
index 000000000..dca2810b9
--- /dev/null
+++ b/tools/hybrid/modeling_zebra_llama.py
@@ -0,0 +1,1222 @@
+"""
+Zebra-Llama: Hybrid Mamba + Multi-Latent Attention (MLA) Model (HuggingFace).
+
+This is a pragmatic HF implementation intended for:
+- loading converted Megatron checkpoints
+- running `generate()` / lm-eval
+
+Notes / simplifications:
+- KV cache is intentionally NOT supported (generation works, but is slower).
+- Mamba mixer is a Megatron-shaped *simplified* implementation (not the full SSM scan).
+- MLA follows Megatron’s tensorization and YaRN RoPE knobs, but uses PyTorch ops.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.utils import logging
+
+# KDA does not require mamba_ssm; pure PyTorch chunked attention is used.
+
+logger = logging.get_logger(__name__)
+
+try:
+    # Optional: FlashAttention v2 (CUDA-only in most environments).
+    from flash_attn import flash_attn_func  # type: ignore
+
+    _HAVE_FLASH_ATTN = True
+except Exception:
+    flash_attn_func = None
+    _HAVE_FLASH_ATTN = False
+
+try:
+    # Optional: Transformer Engine (used by Megatron for many fused ops).
+    import transformer_engine.pytorch as te  # type: ignore
+
+    _HAVE_TE = True
+except Exception:
+    te = None
+    _HAVE_TE = False
+
+
+# =============================================================================
+# Config
+# =============================================================================
+
+
+class ZebraLlamaConfig(PretrainedConfig):
+    model_type = "zebra_llama"
+
+    def __init__(
+        self,
+        vocab_size: int = 128256,
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 32,
+        intermediate_size: int = 8192,
+        num_attention_heads: int = 32,
+        # MLA
+        multi_latent_attention: bool = True,
+        q_lora_rank: int = 1344,
+        kv_lora_rank: int = 128,
+        qk_head_dim: int = 32,
+        qk_pos_emb_head_dim: int = 32,
+        v_head_dim: int = 64,
+        # Hybrid
+        is_hybrid_model: bool = True,
+        hybrid_attention_ratio: float = 0.25,
+        # KDA (Kimi Delta Attention)
+        kda_num_heads: int = 16,
+        kda_head_dim: int = 64,
+        kda_key_head_dim: int = 32,
+        kda_num_key_heads: int = 16,
+        kda_conv_kernel: int = 4,
+        # Mamba (legacy, kept for config compat)
+        mamba_expand: int = 1,
+        mamba_state_dim: int = 64,
+        mamba_head_dim: int = 64,
+        mamba_num_groups: int = 8,
+        mamba_d_conv: int = 4,
+        # Norm
+        normalization: str = "RMSNorm",  # "LayerNorm" or "RMSNorm"
+        layernorm_epsilon: float = 1e-5,
+        # Dropout / residual
+        hidden_dropout: float = 0.0,
+        attention_dropout: float = 0.0,
+        residual_in_fp32: bool = False,
+        bias_dropout_fusion: bool = True,
+        # RoPE / YaRN
+        rope_type: str = "yarn",  # "rope" or "yarn"
+        rope_theta: float = 500000.0,  # megatron: rotary_base
+        rotary_scaling_factor: float = 1.0,
+        original_max_position_embeddings: int = 4096,
+        beta_fast: float = 32.0,
+        beta_slow: float = 1.0,
+        mscale: float = 1.0,
+        mscale_all_dim: float = 1.0,
+        # HF misc
+        max_position_embeddings: int = 131072,
+        use_cache: bool = False,  # intentionally disabled
+        tie_word_embeddings: bool = False,
+        pad_token_id: Optional[int] = None,
+        bos_token_id: int = 128000,
+        eos_token_id: int = 128001,
+        torch_dtype: str | torch.dtype | None = "bfloat16",
+        # Optional: use Transformer Engine ops in HF model
+        use_transformer_engine: bool = True,
+        mamba_hidden_act: str = "silu",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+
+        self.multi_latent_attention = multi_latent_attention
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_head_dim = qk_head_dim
+        self.qk_pos_emb_head_dim = qk_pos_emb_head_dim
+        self.v_head_dim = v_head_dim
+
+        self.is_hybrid_model = is_hybrid_model
+        self.hybrid_attention_ratio = hybrid_attention_ratio
+
+        self.kda_num_heads = kda_num_heads
+        self.kda_head_dim = kda_head_dim
+        self.kda_key_head_dim = kda_key_head_dim
+        self.kda_num_key_heads = kda_num_key_heads
+        self.kda_conv_kernel = kda_conv_kernel
+
+        self.mamba_expand = mamba_expand
+        self.mamba_state_dim = mamba_state_dim
+        self.mamba_head_dim = mamba_head_dim
+        self.mamba_num_groups = mamba_num_groups
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_hidden_act = mamba_hidden_act
+
+        self.normalization = normalization
+        self.layernorm_epsilon = layernorm_epsilon
+
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.residual_in_fp32 = residual_in_fp32
+        self.bias_dropout_fusion = bias_dropout_fusion
+
+        self.rope_type = rope_type
+        self.rope_theta = rope_theta
+        self.rotary_scaling_factor = rotary_scaling_factor
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        self.mscale = mscale
+        self.mscale_all_dim = mscale_all_dim
+
+        self.max_position_embeddings = max_position_embeddings
+        self.use_cache = False  # force off
+        self.torch_dtype = torch_dtype
+        self.use_transformer_engine = use_transformer_engine
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    # --- Derived / computed properties -------------------------------------------
+
+    @property
+    def kda_v_dim(self) -> int:
+        return self.kda_num_heads * self.kda_head_dim
+
+    @property
+    def kda_qk_dim(self) -> int:
+        return self.kda_num_key_heads * self.kda_key_head_dim
+
+    @property
+    def kda_gate_dim(self) -> int:
+        """Forget-gate dimension: always num_heads * key_head_dim."""
+        return self.kda_num_heads * self.kda_key_head_dim
+
+    @property
+    def kda_proj_dim(self) -> int:
+        return self.kda_qk_dim * 2 + self.kda_v_dim
+
+    @property
+    def is_pure_kda(self) -> bool:
+        return self.hybrid_attention_ratio <= 0.0
+
+    @property
+    def is_pure_mla(self) -> bool:
+        return self.hybrid_attention_ratio >= 1.0
+
+
+# =============================================================================
+# Utils
+# =============================================================================
+
+
+def _params_dtype_from_config(config: ZebraLlamaConfig) -> torch.dtype:
+    td = getattr(config, "torch_dtype", None)
+    if td is torch.bfloat16 or td == "bfloat16":
+        return torch.bfloat16
+    if td is torch.float16 or td in ("float16", "fp16"):
+        return torch.float16
+    if td is torch.float32 or td == "float32":
+        return torch.float32
+    return torch.bfloat16
+
+
+def _te_enabled(config: ZebraLlamaConfig) -> bool:
+    return bool(_HAVE_TE and getattr(config, "use_transformer_engine", False))
+
+
+def _filter_kwargs_for_init(cls, kwargs: dict) -> dict:
+    """Filter kwargs to only those accepted by cls.__init__."""
+    import inspect
+
+    try:
+        sig = inspect.signature(cls.__init__)
+    except Exception:
+        return kwargs
+    allowed = set(sig.parameters.keys())
+    # remove self
+    allowed.discard("self")
+    return {k: v for k, v in kwargs.items() if k in allowed}
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.variance_epsilon)
+        return (self.weight.to(torch.float32) * x).to(orig_dtype)
+
+
+class RMSNormWithBias(nn.Module):
+    """RMSNorm with optional bias, matching Transformer Engine's RMSNorm."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.variance_epsilon)
+        return (self.weight.to(torch.float32) * x + self.bias.to(torch.float32)).to(orig_dtype)
+
+
+def _get_norm_eps(norm: nn.Module, default: float = 1e-5) -> float:
+    # Different norm impls use different attribute names.
+    for k in ("variance_epsilon", "eps", "epsilon"):
+        v = getattr(norm, k, None)
+        if v is not None:
+            try:
+                return float(v)
+            except Exception:
+                pass
+    return float(default)
+
+
+def build_norm(config: ZebraLlamaConfig, hidden_size: int) -> nn.Module:
+    # Prefer Transformer Engine norms when explicitly enabled.
+    if _te_enabled(config):
+        eps = float(getattr(config, "layernorm_epsilon", 1e-5))
+        if getattr(config, "normalization", "LayerNorm") == "RMSNorm":
+            te_rms = getattr(te, "RMSNorm", None) if te is not None else None
+            if te_rms is not None:
+                return te_rms(
+                    **_filter_kwargs_for_init(
+                        te_rms, {"hidden_size": hidden_size, "eps": eps, "epsilon": eps}
+                    )
+                )
+            # TE may not ship RMSNorm; fall back below.
+        te_ln = getattr(te, "LayerNorm", None) if te is not None else None
+        if te_ln is not None:
+            return te_ln(
+                **_filter_kwargs_for_init(te_ln, {"hidden_size": hidden_size, "eps": eps, "epsilon": eps})
+            )
+
+    if getattr(config, "normalization", "LayerNorm") == "RMSNorm":
+        return RMSNorm(hidden_size, eps=getattr(config, "layernorm_epsilon", 1e-5))
+    return nn.LayerNorm(hidden_size, eps=getattr(config, "layernorm_epsilon", 1e-5), elementwise_affine=True)
+
+
+def build_linear(config: ZebraLlamaConfig, in_features: int, out_features: int, bias: bool = False):
+    """
+    Build a Linear module, optionally using Transformer Engine.
+    Keeps parameter names compatible with `nn.Linear` so checkpoint loading works.
+    """
+    if _te_enabled(config):
+        te_linear = getattr(te, "Linear", None) if te is not None else None
+        if te_linear is not None:
+            kwargs = {
+                "in_features": in_features,
+                "out_features": out_features,
+                "bias": bias,
+                # Try to align TE param dtype with config.
+                "params_dtype": _params_dtype_from_config(config),
+            }
+            return te_linear(**_filter_kwargs_for_init(te_linear, kwargs))
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def allocate_hybrid_layers(num_layers: int, attention_ratio: float) -> List[str]:
+    """Return a list of 'kda' / 'attention' / 'mlp' tags for each Megatron sublayer.
+
+    ``num_layers`` is the **total** number of Megatron sublayers (KDA/attn + MLP).
+    ``attention_ratio`` controls how many of the non-MLP sublayers are full
+    MLA-attention vs KDA:
+      - 0.0 => pure KDA
+      - 1.0 => pure MLA (no KDA layers)
+      - 0.25 => ~25 % MLA, ~75 % KDA
+    """
+    num_pairs = num_layers // 2
+    if attention_ratio <= 0.0:
+        return ["kda", "mlp"] * num_pairs
+    if attention_ratio >= 1.0:
+        return ["attention", "mlp"] * num_pairs
+
+    num_attn = max(1, int(round(num_pairs * attention_ratio)))
+    num_attn = min(num_attn, num_pairs)
+    spacing = max(1, int(round(num_pairs / num_attn)))
+    types: List[str] = []
+    attn_used = 0
+    for i in range(num_pairs):
+        if (i % spacing == 0) and (attn_used < num_attn):
+            types.append("attention")
+            types.append("mlp")
+            attn_used += 1
+        else:
+            types.append("kda")
+            types.append("mlp")
+    return types
+
+
+# =============================================================================
+# YaRN RoPE (ported from Megatron)
+# =============================================================================
+
+
+def _yarn_find_correction_dim(
+    num_rotations: float, dim: int, rotary_base: float = 10000.0, max_position_embeddings: int = 2048
+) -> float:
+    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+        2 * math.log(rotary_base)
+    )
+
+
+def _yarn_find_correction_range(
+    low_rot: float,
+    high_rot: float,
+    dim: int,
+    rotary_base: float = 10000.0,
+    max_position_embeddings: int = 2048,
+    round_to_int: bool = True,
+) -> Tuple[int, int]:
+    low = _yarn_find_correction_dim(low_rot, dim, rotary_base, max_position_embeddings)
+    high = _yarn_find_correction_dim(high_rot, dim, rotary_base, max_position_embeddings)
+    if round_to_int:
+        low = math.floor(low)
+        high = math.ceil(high)
+    return max(low, 0), min(high, dim - 1)
+
+
+def _yarn_linear_ramp_mask(min_v: float, max_v: float, dim: int) -> torch.Tensor:
+    if min_v == max_v:
+        max_v += 0.001
+    linear_func = (torch.arange(dim, dtype=torch.float32) - min_v) / (max_v - min_v)
+    return torch.clamp(linear_func, 0.0, 1.0)
+
+
+def _yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float:
+    if scale <= 1.0:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def _yarn_get_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float:
+    return float(_yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale(scaling_factor, mscale_all_dim))
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._cache = {}
+
+    def forward(self, x: torch.Tensor, seq_len: int, offset: int = 0):
+        key = (seq_len, x.device.type, x.device.index, x.dtype, offset)
+        if key in self._cache:
+            return self._cache[key]
+        inv_freq = self.inv_freq.to(device=x.device)
+        t = torch.arange(seq_len, device=x.device, dtype=inv_freq.dtype) + offset
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos().to(dtype=x.dtype)
+        sin = emb.sin().to(dtype=x.dtype)
+        self._cache[key] = (cos, sin)
+        return cos, sin
+
+
+class YarnRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        rotary_base: float,
+        scaling_factor: float,
+        original_max_position_embeddings: int,
+        beta_fast: float,
+        beta_slow: float,
+        mscale: float,
+        mscale_all_dim: float,
+        correction_range_round_to_int: bool = True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.rotary_base = float(rotary_base)
+        self.scaling_factor = float(scaling_factor)
+        self.original_max_position_embeddings = int(original_max_position_embeddings)
+        self.beta_fast = float(beta_fast)
+        self.beta_slow = float(beta_slow)
+        self.mscale = float(mscale)
+        self.mscale_all_dim = float(mscale_all_dim)
+        self.correction_range_round_to_int = bool(correction_range_round_to_int)
+
+        inv_freq_extra = 1.0 / (
+            self.rotary_base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim)
+        )
+        inv_freq_inter = 1.0 / (
+            self.scaling_factor
+            * self.rotary_base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim)
+        )
+        self.register_buffer("inv_freq_extra", inv_freq_extra, persistent=False)
+        self.register_buffer("inv_freq_inter", inv_freq_inter, persistent=False)
+        self._cache = {}
+
+    def _inv_freq(self, device: torch.device) -> torch.Tensor:
+        low, high = _yarn_find_correction_range(
+            self.beta_fast,
+            self.beta_slow,
+            self.dim,
+            self.rotary_base,
+            self.original_max_position_embeddings,
+            self.correction_range_round_to_int,
+        )
+        inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, self.dim // 2).to(
+            device=device, dtype=torch.float32
+        )
+        inv_freq = (
+            self.inv_freq_inter.to(device=device) * (1.0 - inv_freq_mask)
+            + self.inv_freq_extra.to(device=device) * inv_freq_mask
+        )
+        return inv_freq
+
+    def forward(self, x: torch.Tensor, seq_len: int, offset: int = 0):
+        key = (seq_len, x.device.type, x.device.index, x.dtype, offset)
+        if key in self._cache:
+            return self._cache[key]
+
+        inv_freq = self._inv_freq(device=x.device)
+        t = torch.arange(seq_len, device=x.device, dtype=inv_freq.dtype) + offset
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        mscale = _yarn_get_concentration_factor(self.scaling_factor, self.mscale, self.mscale_all_dim)
+        cos = (emb.cos() * mscale).to(dtype=x.dtype)
+        sin = (emb.sin() * mscale).to(dtype=x.dtype)
+        self._cache[key] = (cos, sin)
+        return cos, sin
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _mla_rope_reorder(t: torch.Tensor) -> torch.Tensor:
+    """
+    Match Megatron's MLA RoPE path (`multi_latent_attention=True`) where features are
+    reordered as [even dims..., odd dims...] before applying rotate-half.
+    """
+    return torch.cat((t[..., 0::2], t[..., 1::2]), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    position_ids: torch.Tensor,
+    *,
+    multi_latent_attention: bool = False,
+):
+    # q,k: [B, H, L, D]; cos/sin: [S, D]
+    cos = cos[position_ids].unsqueeze(1)  # [B,1,L,D]
+    sin = sin[position_ids].unsqueeze(1)
+
+    if multi_latent_attention:
+        q = _mla_rope_reorder(q)
+        k = _mla_rope_reorder(k)
+
+    return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)
+
+
+# =============================================================================
+# MLA Attention
+# =============================================================================
+
+
+class IdentityNorm(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+class MLAAttention(nn.Module):
+    def __init__(self, config: ZebraLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+
+        self.q_lora_rank = config.q_lora_rank
+        self.kv_lora_rank = config.kv_lora_rank
+        self.qk_head_dim = config.qk_head_dim
+        self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim
+        self.v_head_dim = config.v_head_dim
+        self.q_head_dim = self.qk_head_dim + self.qk_pos_emb_head_dim
+        self.use_q_lora = self.q_lora_rank is not None and self.q_lora_rank > 0
+
+        # Q projection (LoRA or direct)
+        if self.use_q_lora:
+            self.linear_q_down_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False)
+            self.linear_q_up_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False)
+            self.q_layernorm = RMSNormWithBias(self.q_lora_rank, eps=config.layernorm_epsilon)
+        else:
+            self.linear_q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False)
+
+        # KV LoRA (+ positional part)
+        self.linear_kv_down_proj = nn.Linear(
+            self.hidden_size, self.kv_lora_rank + self.qk_pos_emb_head_dim, bias=False
+        )
+        self.linear_kv_up_proj = nn.Linear(
+            self.kv_lora_rank, self.num_heads * (self.qk_head_dim + self.v_head_dim), bias=False
+        )
+
+        self.kv_layernorm = RMSNormWithBias(self.kv_lora_rank, eps=config.layernorm_epsilon)
+
+        self.linear_proj = nn.Linear(self.num_heads * self.v_head_dim, self.hidden_size, bias=True)
+
+        if getattr(config, "rope_type", "rope") == "yarn":
+            self.rotary_emb = YarnRotaryEmbedding(
+                dim=self.qk_pos_emb_head_dim,
+                rotary_base=config.rope_theta,
+                scaling_factor=config.rotary_scaling_factor,
+                original_max_position_embeddings=config.original_max_position_embeddings,
+                beta_fast=config.beta_fast,
+                beta_slow=config.beta_slow,
+                mscale=config.mscale,
+                mscale_all_dim=config.mscale_all_dim,
+            )
+        else:
+            self.rotary_emb = RotaryEmbedding(dim=self.qk_pos_emb_head_dim, base=config.rope_theta)
+
+        mscale = _yarn_get_mscale(config.rotary_scaling_factor, config.mscale_all_dim)
+        self.softmax_scale = mscale * mscale / math.sqrt(self.q_head_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # [B, L, D]
+        attention_mask: Optional[torch.Tensor] = None,  # additive mask [B,1,1,S] with 0 or -inf
+        position_ids: Optional[torch.Tensor] = None,  # [B, L]
+        past_key_value=None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ):
+        # no cache support
+        _ = past_key_value
+
+        bsz, q_len, _ = hidden_states.shape
+        if position_ids is None:
+            position_ids = torch.arange(q_len, device=hidden_states.device, dtype=torch.long)[None, :]
+
+        # Q projection (LoRA or direct)
+        if self.use_q_lora:
+            q_comp = self.q_layernorm(self.linear_q_down_proj(hidden_states))
+            q = self.linear_q_up_proj(q_comp).view(bsz, q_len, self.num_heads, self.q_head_dim)
+        else:
+            q = self.linear_q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.q_head_dim)
+
+        # KV down projections
+        kv_combined = self.linear_kv_down_proj(hidden_states)
+        kv_comp = self.kv_layernorm(kv_combined[..., : self.kv_lora_rank])
+        k_pos_emb = kv_combined[..., self.kv_lora_rank :]  # [B,L,pos]
+
+        # KV up projection
+        kv = self.linear_kv_up_proj(kv_comp).view(
+            bsz, q_len, self.num_heads, self.qk_head_dim + self.v_head_dim
+        )
+
+        # Split
+        q_no_pe = q[..., : self.qk_head_dim]
+        q_pos = q[..., self.qk_head_dim :]
+        k_no_pe = kv[..., : self.qk_head_dim]
+        v = kv[..., self.qk_head_dim :]
+
+        # RoPE on positional components only
+        kv_seq_len = q_len
+        cos, sin = self.rotary_emb(q_pos, seq_len=kv_seq_len)
+
+        k_pos = k_pos_emb.unsqueeze(2).expand(bsz, q_len, self.num_heads, self.qk_pos_emb_head_dim)
+        q_pos = q_pos.transpose(1, 2)  # [B,H,L,pos]
+        k_pos = k_pos.transpose(1, 2)
+        q_pos, k_pos = apply_rotary_pos_emb(q_pos, k_pos, cos, sin, position_ids, multi_latent_attention=True)
+
+        # Concatenate and transpose to [B,H,L,D]
+        q = torch.cat([q_no_pe.transpose(1, 2), q_pos], dim=-1)
+        k = torch.cat([k_no_pe.transpose(1, 2), k_pos], dim=-1)
+        v = v.transpose(1, 2)
+
+        # Attention
+        if output_attentions:
+            attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.q_head_dim)
+            if attention_mask is not None:
+                attn_weights = attn_weights + attention_mask
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(dtype=q.dtype)
+            attn_out = torch.matmul(attn_weights, v)
+        else:
+            dropout_p = float(getattr(self.config, "attention_dropout", 0.0)) if self.training else 0.0
+
+            can_use_flash = (
+                _HAVE_FLASH_ATTN
+                and q.is_cuda
+                and q.dtype in (torch.float16, torch.bfloat16)
+                and self.q_head_dim == self.v_head_dim
+            )
+            if can_use_flash:
+                # flash_attn expects [B,L,H,D]
+                attn_out = (
+                    flash_attn_func(
+                        q.transpose(1, 2).contiguous(),
+                        k.transpose(1, 2).contiguous(),
+                        v.transpose(1, 2).contiguous(),
+                        dropout_p=dropout_p,
+                        softmax_scale=self.softmax_scale,
+                        causal=True,
+                    )
+                    .transpose(1, 2)
+                    .contiguous()
+                )
+            else:
+                sdpa_mask = None
+                if attention_mask is not None:
+                    sdpa_mask = attention_mask < 0  # boolean mask
+                try:
+                    attn_out = F.scaled_dot_product_attention(
+                        q, k, v, attn_mask=sdpa_mask, dropout_p=dropout_p, is_causal=True
+                    )
+                except TypeError:
+                    attn_out = F.scaled_dot_product_attention(
+                        q, k, v, attn_mask=sdpa_mask, dropout_p=dropout_p, is_causal=False
+                    )
+            attn_weights = None
+
+        # Project back to hidden size
+        attn_out = attn_out.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.v_head_dim)
+        attn_out = self.linear_proj(attn_out)
+
+        return attn_out, attn_weights, None
+
+
+class MLAAttentionLayer(nn.Module):
+    """
+    Legacy wrapper kept for compatibility.
+
+    Our primary implementation uses `ZebraLlamaDecoderLayer` directly.
+    This wrapper is not used by the current model, but keeping it as a valid class
+    avoids syntax/import issues if referenced elsewhere.
+    """
+
+    def __init__(self, config: ZebraLlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.input_layernorm = RMSNormWithBias(config.hidden_size, eps=config.layernorm_epsilon)
+        self.self_attention = MLAAttention(config, layer_idx=layer_idx or 0)
+        self.hidden_dropout = float(getattr(config, "hidden_dropout", 0.0))
+        self.dropout = nn.Dropout(self.hidden_dropout)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value=None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ):
+        # No KV cache support for now.
+        _ = past_key_value
+
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights, _ = self.self_attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=None,
+            output_attentions=output_attentions,
+            use_cache=False,
+            **kwargs,
+        )
+        hidden_states = self.dropout(hidden_states) + residual
+        return hidden_states, self_attn_weights
+
+
+# =============================================================================
+# KDA (Kimi Delta Attention)
+# =============================================================================
+
+
+def _torch_kda_gate(
+    g: torch.Tensor, A_log: torch.Tensor, dt_bias: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Pure-PyTorch KDA gate: -exp(A_log) * softplus(g + dt_bias)."""
+    H = g.shape[-2]
+    g = g.float()
+    if dt_bias is not None:
+        g = g + dt_bias.view(H, -1)
+    return -A_log.view(H, 1).float().exp() * F.softplus(g)
+
+
+def _torch_chunk_kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: Optional[float] = None,
+    chunk_size: int = 64,
+    use_qk_l2norm_in_kernel: bool = False,
+) -> torch.Tensor:
+    """Pure-PyTorch chunked KDA forward for inference."""
+    initial_dtype = q.dtype
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    BT = chunk_size
+
+    if scale is None:
+        scale = K**-0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6)
+        k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6)
+
+    pad_size = (BT - T % BT) % BT
+    if pad_size > 0:
+        q = F.pad(q, (0, 0, 0, 0, 0, pad_size))
+        k = F.pad(k, (0, 0, 0, 0, 0, pad_size))
+        v = F.pad(v, (0, 0, 0, 0, 0, pad_size))
+        g = F.pad(g, (0, 0, 0, 0, 0, pad_size))
+        beta = F.pad(beta, (0, 0, 0, pad_size))
+
+    total_T = T + pad_size
+    NT = total_T // BT
+
+    q, k, v, g, beta = [x.transpose(1, 2).contiguous().float() for x in (q, k, v, g, beta)]
+    q = q * scale
+
+    q = q.reshape(B, H, NT, BT, K)
+    k = k.reshape(B, H, NT, BT, K)
+    v = v.reshape(B, H, NT, BT, V)
+    g = g.reshape(B, H, NT, BT, K)
+    beta = beta.reshape(B, H, NT, BT)
+
+    g = g.cumsum(dim=-2)
+    k_eg = k * g.exp()
+
+    A = torch.zeros(B, H, NT, BT, BT, dtype=torch.float, device=q.device)
+    for j in range(BT):
+        k_j = k[..., j, :]
+        g_j = g[..., j : j + 1, :]
+        decay = (g - g_j).clamp(max=0).exp()
+        A[..., j] = (k * decay * k_j.unsqueeze(-2)).sum(-1)
+
+    A = A * beta.unsqueeze(-1)
+    mask_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    A = -A.masked_fill(mask_upper, 0)
+
+    for i in range(1, BT):
+        A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
+
+    A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta.unsqueeze(-2)
+    w = A @ k_eg
+    u = A @ v
+
+    mask_causal = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
+    S = q.new_zeros(B, H, K, V)
+    o = torch.zeros_like(v)
+
+    for i in range(NT):
+        q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i]
+        A_qk = torch.zeros(B, H, BT, BT, dtype=torch.float, device=q.device)
+        for j in range(BT):
+            k_j = k_i[..., j, :]
+            g_j = g_i[..., j : j + 1, :]
+            decay = (g_i - g_j).clamp(max=0).exp()
+            A_qk[..., j] = (q_i * decay * k_j.unsqueeze(-2)).sum(-1)
+        A_qk = A_qk.masked_fill(mask_causal, 0)
+        v_i = u_i - w_i @ S
+        o[:, :, i] = (q_i * g_i.exp()) @ S + A_qk @ v_i
+        g_last = g_i[:, :, -1]
+        S = S * g_last.unsqueeze(-1).exp()
+        k_dec = (g_last.unsqueeze(-2) - g_i).exp() * k_i
+        S = S + k_dec.transpose(-1, -2) @ v_i
+
+    o = o.reshape(B, H, -1, V)[:, :, :T]
+    o = o.transpose(1, 2).contiguous().to(initial_dtype)
+    return o
+
+
+class KDAMixer(nn.Module):
+    """Kimi Delta Attention mixer for HF inference.
+
+    All dimensions are derived from :class:`ZebraLlamaConfig` so that the same
+    class works for any KDA head configuration (symmetric / asymmetric key-value
+    dimensions, grouped-query heads, etc.).
+    """
+
+    def __init__(self, config: ZebraLlamaConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.kda_num_heads
+        self.head_dim = config.kda_head_dim  # value head dim
+        self.head_k_dim = config.kda_key_head_dim  # key/query head dim
+        self.num_k_heads = config.kda_num_key_heads  # may differ from num_heads (GQA-style)
+        self.conv_kernel = config.kda_conv_kernel
+
+        self.v_dim = config.kda_v_dim  # num_heads * head_dim
+        self.qk_dim = config.kda_qk_dim  # num_k_heads * head_k_dim
+        self.gate_dim = config.kda_gate_dim  # num_heads * head_k_dim
+        proj_dim = config.kda_proj_dim  # qk_dim*2 + v_dim
+
+        self.in_proj = nn.Linear(self.hidden_size, proj_dim, bias=False)
+        self.conv1d = nn.Conv1d(
+            proj_dim,
+            proj_dim,
+            kernel_size=self.conv_kernel,
+            groups=proj_dim,
+            padding=self.conv_kernel - 1,
+            bias=False,
+        )
+
+        self.gate_norm = RMSNormWithBias(self.hidden_size, eps=config.layernorm_epsilon)
+
+        # Forget gate (low-rank): hidden -> head_dim -> gate_dim (num_heads * head_k_dim)
+        self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.f_b_proj = nn.Linear(self.head_dim, self.gate_dim, bias=False)
+
+        self.A_log = nn.Parameter(torch.zeros(1, 1, self.num_heads, 1))
+        self.dt_bias = nn.Parameter(torch.zeros(self.gate_dim))
+
+        # Beta gate: per-head scalar
+        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
+
+        # Output gate (low-rank): hidden -> head_dim -> v_dim (num_heads * head_dim)
+        self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.g_b_proj = nn.Linear(self.head_dim, self.v_dim, bias=False)
+
+        self.out_norm = RMSNormWithBias(self.head_dim, eps=config.layernorm_epsilon)
+        self.out_proj = nn.Linear(self.v_dim, self.hidden_size, bias=False)
+
+    def forward(self, normed_hidden: torch.Tensor, raw_hidden: torch.Tensor) -> torch.Tensor:
+        bsz, seq_len, _ = normed_hidden.shape
+
+        qkv = self.in_proj(normed_hidden)
+        qkv = qkv.transpose(1, 2).contiguous()
+        qkv = F.silu(self.conv1d(qkv)[..., :seq_len])
+        qkv = qkv.transpose(1, 2)
+
+        q, k, v = torch.split(qkv, [self.qk_dim, self.qk_dim, self.v_dim], dim=-1)
+        q = q.reshape(bsz, seq_len, self.num_k_heads, self.head_k_dim)
+        k = k.reshape(bsz, seq_len, self.num_k_heads, self.head_k_dim)
+        v = v.reshape(bsz, seq_len, self.num_heads, self.head_dim)
+
+        if self.num_heads // self.num_k_heads > 1:
+            q = q.repeat_interleave(self.num_heads // self.num_k_heads, dim=2)
+            k = k.repeat_interleave(self.num_heads // self.num_k_heads, dim=2)
+
+        h = self.gate_norm(raw_hidden)
+
+        g = self.f_b_proj(self.f_a_proj(h))
+        g = g.reshape(bsz, seq_len, self.num_heads, self.head_k_dim)
+        g = _torch_kda_gate(g, self.A_log.view(-1), dt_bias=self.dt_bias)
+
+        beta = self.b_proj(h).float().sigmoid()
+
+        core_out = _torch_chunk_kda_fwd(
+            q.contiguous(),
+            k.contiguous(),
+            v.contiguous(),
+            g,
+            beta,
+            use_qk_l2norm_in_kernel=True,
+        )
+
+        gate = self.g_b_proj(self.g_a_proj(h))
+        gate = gate.reshape(bsz, seq_len, -1, self.head_dim)
+
+        x_dtype = core_out.dtype
+        y = self.out_norm(core_out.reshape(-1, self.head_dim))
+        y = y * gate.reshape(-1, self.head_dim).float().sigmoid()
+        y = y.to(x_dtype).reshape(bsz, seq_len, -1)
+
+        return self.out_proj(y)
+
+
+class KDALayer(nn.Module):
+    def __init__(self, config: ZebraLlamaConfig):
+        super().__init__()
+        self.config = config
+        self.norm = RMSNormWithBias(config.hidden_size, eps=config.layernorm_epsilon)
+        self.mixer = KDAMixer(config)
+        self.dropout = nn.Dropout(float(getattr(config, "hidden_dropout", 0.0)))
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs):
+        residual = hidden_states
+        normed = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+        out = self.mixer(normed, hidden_states)
+        return self.dropout(out) + residual, None
+
+
+# =============================================================================
+# MLP
+# =============================================================================
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config: ZebraLlamaConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.linear_fc1 = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=True)
+        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_fc1 = self.linear_fc1(x)
+        x_glu = x_fc1[:, :, : self.intermediate_size]
+        x_lin = x_fc1[:, :, self.intermediate_size :]
+        return self.linear_fc2(F.silu(x_glu) * x_lin)
+
+
+class MLPLayer(nn.Module):
+    def __init__(self, config: ZebraLlamaConfig):
+        super().__init__()
+        self.config = config
+        self.mlp = SwiGLUMLP(config)
+        self.dropout = nn.Dropout(float(getattr(config, "hidden_dropout", 0.0)))
+        self.pre_mlp_layernorm = RMSNormWithBias(config.hidden_size, eps=config.layernorm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        inference_context=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.pre_mlp_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.dropout(hidden_states) + residual
+        return hidden_states, None
+
+
+# =============================================================================
+# Base model
+# =============================================================================
+
+
+class ZebraLlamaModel(PreTrainedModel):
+    config_class = ZebraLlamaConfig
+
+    def __init__(self, config: ZebraLlamaConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.padding_idx = config.pad_token_id
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+
+        layer_types = allocate_hybrid_layers(config.num_hidden_layers, config.hybrid_attention_ratio)
+        self.layers = nn.ModuleList()
+        for i, t in enumerate(layer_types):
+            if t == "kda":
+                self.layers.append(KDALayer(config))
+            elif t == "attention":
+                self.layers.append(MLAAttentionLayer(config, i))
+            elif t == "mlp":
+                self.layers.append(MLPLayer(config))
+        self.norm = RMSNormWithBias(config.hidden_size, eps=config.layernorm_epsilon)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        # No cache support
+        _ = past_key_values
+
+        output_attentions = (
+            output_attentions if output_attentions is not None else bool(self.config.output_attentions)
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else bool(self.config.output_hidden_states)
+        )
+        return_dict = return_dict if return_dict is not None else bool(self.config.use_return_dict)
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("Specify only one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len, _ = inputs_embeds.shape
+        hidden_states = inputs_embeds
+
+        # attention_mask: [B,L] -> additive [B,1,1,L]
+        attn_mask = None
+        if attention_mask is not None:
+            if attention_mask.dim() == 2:
+                attn_mask = attention_mask[:, None, None, :].to(dtype=hidden_states.dtype)
+                attn_mask = (1.0 - attn_mask) * torch.finfo(hidden_states.dtype).min
+            else:
+                attn_mask = attention_mask
+
+        if position_ids is None:
+            device = hidden_states.device
+            position_ids = torch.arange(seq_len, device=device, dtype=torch.long)[None, :]
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states, attn = layer(
+                hidden_states,
+                attention_mask=attn_mask,
+                position_ids=position_ids,
+                output_attentions=output_attentions,
+            )
+            if output_attentions:
+                all_attns += (attn,)
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, None, all_hidden_states, all_attns] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=all_hidden_states,
+            attentions=all_attns,
+        )
+
+
+# =============================================================================
+# Causal LM
+# =============================================================================
+
+
+class ZebraLlamaForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = ZebraLlamaConfig
+
+    def __init__(self, config: ZebraLlamaConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.model = ZebraLlamaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self._tied_weights_keys = ["lm_head.weight"]
+        self.post_init()
+        if getattr(config, "tie_word_embeddings", False):
+            self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=None,
+            inputs_embeds=inputs_embeds,
+            use_cache=False,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states).float()
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1).to(shift_logits.device)
+            )
+
+        if not return_dict:
+            out = (logits,) + outputs[1:]
+            return (loss,) + out if loss is not None else out
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, inputs_embeds=None, **kwargs):
+        # No KV cache: do NOT slice to last token. Always recompute from scratch.
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 0)
+
+        if inputs_embeds is not None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "position_ids": position_ids,
+                "use_cache": False,
+                "past_key_values": None,
+            }
+        )
+        return model_inputs
+
+
+# =============================================================================
+# Optional local registration (helps when importing this module directly)
+# =============================================================================
+
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM  # noqa: E402
+
+AutoConfig.register("zebra_llama", ZebraLlamaConfig)
+AutoModel.register(ZebraLlamaConfig, ZebraLlamaModel)
+AutoModelForCausalLM.register(ZebraLlamaConfig, ZebraLlamaForCausalLM)
diff --git a/tools/hybrid/patch_fla_triton_autotune_hang.sh b/tools/hybrid/patch_fla_triton_autotune_hang.sh
new file mode 100755
index 000000000..8c01b801f
--- /dev/null
+++ b/tools/hybrid/patch_fla_triton_autotune_hang.sh
@@ -0,0 +1,512 @@
+#!/bin/bash
+# Patch FLA Triton kernels to fix autotuning hangs on MI300X/ROCm.
+#
+# Two kernels hang during Triton autotuning with num_stages >= 3:
+#   1. chunk_bwd_kernel_dh      in fla/ops/common/chunk_h.py
+#   2. chunk_kda_bwd_kernel_intra in fla/ops/kda/chunk_intra.py
+#
+# Root cause: Triton's multi-stage software pipelining (num_stages >= 3)
+# generates broken HIP code for kernels with dynamic loops, debug barriers,
+# and complex control flow on MI300X.
+#
+# Fix: restrict num_stages to [2] for affected kernels.
+
+set -e
+
+echo "=== FLA Triton MI300X/ROCm Patch ==="
+echo ""
+
+########################################################################
+# Step 0: Purge Triton cache + set up per-rank cache isolation
+########################################################################
+
+# echo "[Step 0] Purging Triton caches ..."
+# rm -rf ~/.triton/cache 2>/dev/null && echo "  Removed ~/.triton/cache" || true
+# rm -rf /tmp/triton* 2>/dev/null && echo "  Removed /tmp/triton*" || true
+# echo "  Done."
+# echo ""
+
+# echo "[Step 0b] Installing per-rank Triton cache isolation ..."
+
+# PERRANK_SCRIPT="/opt/venv/lib/python3.10/site-packages/fla/_triton_per_rank_cache.py"
+# cat > "$PERRANK_SCRIPT" << 'PERRANK_EOF'
+# """
+# Per-rank Triton cache isolation for multi-GPU training.
+
+# Import this BEFORE any Triton kernel usage (e.g. at the top of your
+# training script or in an __init__.py) to give each GPU rank its own
+# Triton cache directory, preventing race conditions during autotune.
+
+# Works with torchrun, deepspeed, and bare CUDA_VISIBLE_DEVICES.
+# """
+# import os
+
+# def _setup_per_rank_triton_cache():
+#     rank = os.environ.get("LOCAL_RANK") or os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK") or "0"
+#     base = os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton/cache"))
+#     per_rank = os.path.join(base, f"rank_{rank}")
+#     os.makedirs(per_rank, exist_ok=True)
+#     os.environ["TRITON_CACHE_DIR"] = per_rank
+
+# _setup_per_rank_triton_cache()
+# PERRANK_EOF
+
+# FLA_INIT="/opt/venv/lib/python3.10/site-packages/fla/__init__.py"
+# if [ -f "$FLA_INIT" ]; then
+#     cp "$FLA_INIT" "${FLA_INIT}.bak" 2>/dev/null || true
+#     IMPORT_LINE="import fla._triton_per_rank_cache  # per-rank cache isolation"
+#     if ! grep -qF "$IMPORT_LINE" "$FLA_INIT"; then
+#         sed -i "1i\\$IMPORT_LINE" "$FLA_INIT"
+#         echo "  Injected per-rank cache import into $FLA_INIT"
+#     else
+#         echo "  Per-rank cache import already present in $FLA_INIT"
+#     fi
+# else
+#     echo "  WARNING: $FLA_INIT not found, per-rank cache not auto-injected."
+#     echo "  Add 'import fla._triton_per_rank_cache' to the top of your training script."
+# fi
+# echo "  Each rank gets: ~/.triton/cache/rank_\$LOCAL_RANK/"
+# echo ""
+
+########################################################################
+# Locate fla package (supports editable installs and site-packages)
+########################################################################
+
+FLA_ROOT=$(python3 -c "import fla, os; print(os.path.dirname(fla.__file__))" 2>/dev/null)
+if [ -z "$FLA_ROOT" ]; then
+    echo "ERROR: Could not locate the fla package. Is flash-linear-attention installed?"
+    exit 1
+fi
+echo "FLA package found at: $FLA_ROOT"
+echo ""
+
+########################################################################
+# Patch 1: chunk_kda_bwd_kernel_intra (fla/ops/kda/chunk_intra.py)
+#   The autotuning tries num_stages in [2, 3, 4]. num_stages >= 3 hangs.
+########################################################################
+
+INTRA_TARGET="$FLA_ROOT/ops/kda/chunk_intra.py"
+
+if [ -f "$INTRA_TARGET" ]; then
+    cp "$INTRA_TARGET" "${INTRA_TARGET}.bak"
+    echo "[Patch 1] Patching $INTRA_TARGET"
+
+    sed -i 's/for num_stages in \[2, 3, 4\]/for num_stages in [2]/g' "$INTRA_TARGET"
+
+    echo "  - Restricted num_stages to [2] for all autotune configs"
+    echo "  - Backed up original to ${INTRA_TARGET}.bak"
+    echo ""
+else
+    echo "[Patch 1] SKIP: $INTRA_TARGET not found"
+    echo ""
+fi
+
+########################################################################
+# Patch 2: chunk_bwd_kernel_dh (fla/ops/common/chunk_h.py)
+#   The backward kernel autotuning hangs on ROCm.
+#   Fix: replace autotuning with fixed config (num_warps=4, num_stages=2).
+########################################################################
+
+# TARGET="/opt/venv/lib/python3.10/site-packages/fla/ops/common/chunk_h.py"
+
+# if [ ! -f "$TARGET" ]; then
+#     echo "[Patch 2] SKIP: $TARGET not found"
+# else
+
+# cp "$TARGET" "${TARGET}.bak"
+# echo "[Patch 2] Patching $TARGET"
+
+# cat > "$TARGET" << 'PATCH_EOF'
+# # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# # Patched for MI300X/ROCm: disabled autotuning on backward kernel to prevent hang.
+
+# import torch
+# import triton
+# import triton.language as tl
+
+# from fla.ops.utils import prepare_chunk_offsets
+# from fla.ops.utils.op import exp
+# from fla.utils import autotune_cache_kwargs, check_shared_mem
+
+# BKV_LIST = [32, 64] if check_shared_mem() else [16, 32]
+
+
+# @triton.heuristics({
+#     'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
+#     'STORE_FINAL_STATE': lambda args: args['ht'] is not None,
+#     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+# })
+# @triton.autotune(
+#     configs=[
+#         triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages)
+#         for BK in BKV_LIST
+#         for BV in BKV_LIST
+#         for num_warps in [4, 8]
+#         for num_stages in [2, 3]
+#     ],
+#     key=['BT', 'USE_G', 'USE_GK', 'USE_GV'],
+#     **autotune_cache_kwargs,
+# )
+# @triton.jit(do_not_specialize=['T'])
+# def chunk_fwd_kernel_h(
+#     k,
+#     v,
+#     h,
+#     g,
+#     g_gamma,
+#     gk,
+#     gv,
+#     h0,
+#     ht,
+#     cu_seqlens,
+#     split_offsets,
+#     T,
+#     H: tl.constexpr,
+#     K: tl.constexpr,
+#     V: tl.constexpr,
+#     BT: tl.constexpr,
+#     BS: tl.constexpr,
+#     BK: tl.constexpr,
+#     BV: tl.constexpr,
+#     USE_G: tl.constexpr,
+#     USE_G_GAMMA: tl.constexpr,
+#     USE_GK: tl.constexpr,
+#     USE_GV: tl.constexpr,
+#     USE_INITIAL_STATE: tl.constexpr,
+#     STORE_FINAL_STATE: tl.constexpr,
+#     IS_VARLEN: tl.constexpr,
+# ):
+#     i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+#     i_n, i_h = i_nh // H, i_nh % H
+#     if IS_VARLEN:
+#         bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+#         T = eos - bos
+#         NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS)
+#         boh = tl.load(split_offsets + i_n).to(tl.int32)
+#     else:
+#         bos, eos = i_n * T, i_n * T + T
+#         NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS)
+#         boh = i_n * NS
+#     NTS = BS // BT
+
+#     if USE_G_GAMMA:
+#         b_gamma = tl.load(g_gamma + i_h)
+#         b_g = b_gamma * (tl.arange(0, BT) + 1)
+
+#     # [BK, BV]
+#     b_h = tl.zeros([BK, BV], dtype=tl.float32)
+#     if USE_INITIAL_STATE:
+#         p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+#         b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32)
+
+#     for i_t in range(NT):
+#         i_s = i_t // NTS
+#         p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+#         p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+
+#         o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V
+#         p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+
+#         if i_t % NTS == 0:
+#             tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
+#         # [BK, BT]
+#         b_k = tl.load(p_k, boundary_check=(0, 1))
+#         # [BT, BV]
+#         b_v = tl.load(p_v, boundary_check=(0, 1))
+#         last_idx = min((i_t + 1) * BT, T) - 1
+
+#         # scalar decay
+#         if USE_G:
+#             b_g_last = tl.load(g + bos * H + last_idx * H + i_h)
+#             p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h
+#             b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.)
+#             b_h *= exp(b_g_last)
+#             b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype)
+
+#         if USE_G_GAMMA:
+#             b_g_last = b_gamma * min(BT, T - i_t * BT)
+#             b_h *= exp(b_g_last)
+#             b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype)
+
+#         # vector decay, h = Diag(gk) @ h
+#         if USE_GK:
+#             p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+#             p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK)
+
+#             b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.)
+#             b_h *= exp(b_gk_last)[:, None]
+
+#             b_gk = tl.load(p_gk, boundary_check=(0, 1))
+#             b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype)
+
+#         # vector decay, h = h @ Diag(gv)
+#         if USE_GV:
+#             p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+#             p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV)
+
+#             b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.)
+#             b_h *= exp(b_gv_last)[None, :]
+
+#             b_gv = tl.load(p_gv, boundary_check=(0, 1))
+#             b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype)
+
+#         b_h += tl.dot(b_k, b_v)
+
+#     if STORE_FINAL_STATE:
+#         p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+#         tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+
+
+# # PATCHED: Fixed config for backward kernel to avoid autotuning hang on MI300X/ROCm.
+# # Original used @triton.autotune which hangs during config search on AMD GPUs.
+# @triton.heuristics({
+#     'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None,
+#     'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None,
+#     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+# })
+# @triton.jit(do_not_specialize=['T'])
+# def chunk_bwd_kernel_dh(
+#     q,
+#     g,
+#     g_gamma,
+#     gk,
+#     gv,
+#     do,
+#     dh,
+#     dht,
+#     dh0,
+#     cu_seqlens,
+#     split_offsets,
+#     scale,
+#     T,
+#     HQ: tl.constexpr,
+#     H: tl.constexpr,
+#     K: tl.constexpr,
+#     V: tl.constexpr,
+#     BT: tl.constexpr,
+#     BS: tl.constexpr,
+#     BK: tl.constexpr,
+#     BV: tl.constexpr,
+#     NG: tl.constexpr,
+#     USE_G: tl.constexpr,
+#     USE_G_GAMMA: tl.constexpr,
+#     USE_GK: tl.constexpr,
+#     USE_GV: tl.constexpr,
+#     STORE_INITIAL_STATE_GRADIENT: tl.constexpr,
+#     USE_FINAL_STATE_GRADIENT: tl.constexpr,
+#     IS_VARLEN: tl.constexpr,
+# ):
+#     i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+#     i_n, i_hq = i_nh // HQ, i_nh % HQ
+#     i_h = i_hq // NG
+#     if IS_VARLEN:
+#         bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+#         T = eos - bos
+#         NT = tl.cdiv(T, BT)
+#         NS = tl.cdiv(T, BS)
+#         boh = tl.load(split_offsets + i_n).to(tl.int32)
+#     else:
+#         bos, eos = i_n * T, i_n * T + T
+#         NT = tl.cdiv(T, BT)
+#         NS = tl.cdiv(T, BS)
+#         boh = i_n * NS
+
+#     if USE_G_GAMMA:
+#         b_gamma = tl.load(g_gamma + i_h)
+#         b_g = b_gamma * (tl.arange(0, BT) + 1)
+
+#     # [BK, BV]
+#     b_dh = tl.zeros([BK, BV], dtype=tl.float32)
+#     if USE_FINAL_STATE_GRADIENT:
+#         p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+#         b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32)
+
+#     for i_t in range(NT - 1, -1, -1):
+#         i_s = i_t // (BS // BT)
+#         o_dh = ((boh + i_s) * H + i_h).to(tl.int64) * K*V
+#         p_dh = tl.make_block_ptr(dh + o_dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+
+#         if i_t % (BS // BT) == 0:
+#             tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1))
+#         last_idx = min(i_t * BT + BT, T) - 1
+#         # [BK, BT]
+#         p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+#         p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+#         b_q = tl.load(p_q, boundary_check=(0, 1))
+#         b_q = (b_q * scale).to(b_q.dtype)
+#         # [BT, BV]
+#         b_do = tl.load(p_do, boundary_check=(0, 1))
+
+#         if USE_G:
+#             p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h
+#             b_g_last = tl.load(g + (bos + last_idx) * H + i_h)
+#             b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.)
+#             b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype)
+#             b_dh *= exp(b_g_last)
+
+#         if USE_G_GAMMA:
+#             b_g_last = b_gamma * min(BT, T - i_t * BT)
+#             b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype)
+#             b_dh *= exp(b_g_last)
+
+#         if USE_GK:
+#             p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+#             p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK)
+
+#             b_gk = tl.load(p_gk, boundary_check=(0, 1))
+#             b_q = (b_q * exp(b_gk)).to(b_q.dtype)
+#             b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.)
+#             b_dh *= exp(b_gk_last)[:, None]
+
+#         if USE_GV:
+#             p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+#             p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV)
+
+#             b_gv = tl.load(p_gv, boundary_check=(0, 1))
+#             b_do = (b_do * exp(b_gv))
+
+#             b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.)
+#             b_dh *= exp(b_gv_last)[None, :]
+
+#         b_dh += tl.dot(b_q, b_do.to(b_q.dtype))
+
+#     if STORE_INITIAL_STATE_GRADIENT:
+#         p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+#         tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1))
+
+
+# def chunk_fwd_h(
+#     k: torch.Tensor,
+#     v: torch.Tensor,
+#     g: torch.Tensor | None = None,
+#     g_gamma: torch.Tensor | None = None,
+#     gk: torch.Tensor | None = None,
+#     gv: torch.Tensor | None = None,
+#     h0: torch.Tensor | None = None,
+#     output_final_state: bool = False,
+#     cu_seqlens: torch.Tensor | None = None,
+#     chunk_size: int = 64,
+#     split_size: int | None = None,
+#     states_in_fp32: bool = False,
+# ) -> tuple[torch.Tensor, torch.Tensor]:
+#     B, T, H, K, V = *k.shape, v.shape[-1]
+#     BT = chunk_size
+#     BS = BT if split_size is None else split_size
+#     assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}"
+#     if cu_seqlens is None:
+#         N, NS, split_offsets = B, triton.cdiv(T, BS), None
+#     else:
+#         split_offsets = prepare_chunk_offsets(cu_seqlens, BS)
+#         N, NS = len(cu_seqlens) - 1, split_offsets[-1].item()
+
+#     h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float)
+#     ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None
+#     def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H)
+#     chunk_fwd_kernel_h[grid](
+#         k=k,
+#         v=v,
+#         h=h,
+#         g=g,
+#         g_gamma=g_gamma,
+#         gk=gk,
+#         gv=gv,
+#         h0=h0,
+#         ht=ht,
+#         cu_seqlens=cu_seqlens,
+#         split_offsets=split_offsets,
+#         T=T,
+#         H=H,
+#         K=K,
+#         V=V,
+#         BT=BT,
+#         BS=BS,
+#         USE_G=g is not None,
+#         USE_G_GAMMA=g_gamma is not None,
+#         USE_GK=gk is not None,
+#         USE_GV=gv is not None,
+#     )
+#     return h, ht
+
+
+# def chunk_bwd_dh(
+#     q: torch.Tensor,
+#     k: torch.Tensor,
+#     v: torch.Tensor,
+#     do: torch.Tensor,
+#     h0: torch.Tensor,
+#     dht: torch.Tensor,
+#     scale: float,
+#     g: torch.Tensor | None = None,
+#     g_gamma: torch.Tensor | None = None,
+#     gk: torch.Tensor | None = None,
+#     gv: torch.Tensor | None = None,
+#     cu_seqlens: torch.Tensor | None = None,
+#     chunk_size: int = 64,
+#     split_size: int | None = None,
+#     states_in_fp32: bool = False,
+# ) -> tuple[torch.Tensor, torch.Tensor]:
+#     B, T, H, K, V = *k.shape, v.shape[-1]
+#     HQ = q.shape[2]
+#     BT = chunk_size
+#     BS = BT if split_size is None else split_size
+#     assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}"
+#     if cu_seqlens is None:
+#         N, NS, split_offsets = B, triton.cdiv(T, BS), None
+#     else:
+#         split_offsets = prepare_chunk_offsets(cu_seqlens, BS)
+#         N, NS = len(cu_seqlens) - 1, split_offsets[-1].item()
+#     NG = HQ // H
+
+#     dh = k.new_empty(B, NS, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float)
+#     dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None
+
+#     # PATCHED: Use fixed BK/BV instead of autotuning to avoid hang on MI300X/ROCm.
+#     BK_fixed = 64 if K >= 64 else 32
+#     BV_fixed = 64 if V >= 64 else 32
+#     def grid(meta): return (triton.cdiv(K, BK_fixed), triton.cdiv(V, BV_fixed), N * H)
+#     chunk_bwd_kernel_dh[grid](
+#         q=q,
+#         g=g,
+#         g_gamma=g_gamma,
+#         gk=gk,
+#         gv=gv,
+#         do=do,
+#         dh=dh,
+#         dht=dht,
+#         dh0=dh0,
+#         cu_seqlens=cu_seqlens,
+#         split_offsets=split_offsets,
+#         scale=scale,
+#         T=T,
+#         HQ=HQ,
+#         H=H,
+#         K=K,
+#         V=V,
+#         BT=BT,
+#         BS=BS,
+#         BK=BK_fixed,
+#         BV=BV_fixed,
+#         NG=NG,
+#         USE_G=g is not None,
+#         USE_G_GAMMA=g_gamma is not None,
+#         USE_GK=gk is not None,
+#         USE_GV=gv is not None,
+#         num_warps=4,
+#         num_stages=2,
+#     )
+#     return dh, dh0
+# PATCH_EOF
+
+# echo "  - Forward kernel: unchanged (autotuning kept)"
+# echo "  - Backward kernel: removed @triton.autotune, using fixed BK/BV with num_warps=4, num_stages=2"
+# echo "  - Backed up original to ${TARGET}.bak"
+
+# fi
+
+echo ""
+echo "=== Patch complete ==="
+echo ""
+echo "Per-rank cache: each GPU rank now uses ~/.triton/cache/rank_\$LOCAL_RANK/"
+echo ""
+echo "To restore originals:"
+echo "  cp ${INTRA_TARGET}.bak $INTRA_TARGET"
diff --git a/tools/hybrid/run_hybrid_eval.sh b/tools/hybrid/run_hybrid_eval.sh
new file mode 100644
index 000000000..e45768cc3
--- /dev/null
+++ b/tools/hybrid/run_hybrid_eval.sh
@@ -0,0 +1,104 @@
+#!/bin/bash
+###############################################################################
+# End-to-end evaluation of the 75% Hybrid GDN+MLA 300M model.
+#
+# Run this *inside the rocm/primus container* with the repo mounted at
+# /home/<user>/Primus:
+#
+#   cd /home/<user>/Primus && bash tools/hybrid/run_hybrid_eval.sh
+#
+# What it does:
+#   1. Converts the Primus Megatron checkpoint → FLA HuggingFace format
+#      (uses tools/hybrid/convert_gdn_hybrid_to_fla_hf.py — handles the 3 MLA + 9 GDN
+#      sublayer mix and FLA's nn.Sequential(Linear→RMSNorm→Linear) LoRA packing)
+#   2. Runs lm-eval on the Primus-converted HF model
+#   3. Runs lm-eval on FLA's HF checkpoint (apples-to-apples comparison)
+#   4. Diffs the two scoreboards
+###############################################################################
+set -euo pipefail
+
+PRIMUS_CKPT=${PRIMUS_CKPT:-output/amd/root/zebra_llama_300M_gdn_hybrid-pretrain/checkpoints/iter_0004768}
+PRIMUS_HF_DIR=${PRIMUS_HF_DIR:-output/gdn_hybrid_300M_fla_hf}
+FLA_HF_DIR=${FLA_HF_DIR:-$HOME/checkpoints/gdn_hybrid_300M_10B/checkpoint-4768}
+FLA_CONFIG=${FLA_CONFIG:-$HOME/flash-linear-attention/legacy/training/configs/gated_deltanet_300M_hybrid.json}
+RESULTS_DIR=${RESULTS_DIR:-output/gdn_hybrid_300M_eval_results}
+TASKS=${TASKS:-arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race}
+BATCH_SIZE=${BATCH_SIZE:-auto}
+TOKENIZER=${TOKENIZER:-$HOME/checkpoints/gdn_pure_300M_10B}
+
+echo "==========[run_hybrid_eval.sh]=========="
+echo "PRIMUS_CKPT     = ${PRIMUS_CKPT}"
+echo "PRIMUS_HF_DIR   = ${PRIMUS_HF_DIR}"
+echo "FLA_HF_DIR      = ${FLA_HF_DIR}"
+echo "FLA_CONFIG      = ${FLA_CONFIG}"
+echo "RESULTS_DIR     = ${RESULTS_DIR}"
+echo "TASKS           = ${TASKS}"
+echo "BATCH_SIZE      = ${BATCH_SIZE}"
+echo "TOKENIZER       = ${TOKENIZER}"
+echo "========================================="
+
+mkdir -p "${RESULTS_DIR}"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 1: Convert Primus checkpoint to FLA HF format
+# ─────────────────────────────────────────────────────────────────────────────
+if [ ! -f "${PRIMUS_HF_DIR}/model.safetensors" ]; then
+    echo
+    echo "==========[Step 1] Converting Primus → FLA HF =========="
+    python tools/hybrid/convert_gdn_hybrid_to_fla_hf.py \
+        --checkpoint-path "${PRIMUS_CKPT}" \
+        --output-dir "${PRIMUS_HF_DIR}" \
+        --config "${FLA_CONFIG}"
+else
+    echo "Primus HF checkpoint already exists at ${PRIMUS_HF_DIR}, skipping conversion."
+fi
+
+# Copy FLA tokenizer into both eval dirs so lm-eval can load them with --tokenizer
+# (the FLA hybrid config doesn't ship a tokenizer; it shares Llama-3.2 tokenizer)
+for tdir in "${PRIMUS_HF_DIR}" "${FLA_HF_DIR}"; do
+    if [ ! -f "${tdir}/tokenizer.json" ] && [ -f "${TOKENIZER}/tokenizer.json" ]; then
+        cp "${TOKENIZER}/tokenizer.json" \
+           "${TOKENIZER}/tokenizer_config.json" \
+           "${TOKENIZER}/special_tokens_map.json" \
+           "${tdir}/" 2>/dev/null || true
+    fi
+done
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 2: Evaluate Primus-converted HF
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 2] lm-eval on Primus HF =========="
+python tools/hybrid/eval_gdn_lm_eval.py \
+    --model hf \
+    --model_args "pretrained=${PRIMUS_HF_DIR},dtype=bfloat16,trust_remote_code=True" \
+    --tasks "${TASKS}" \
+    --batch_size "${BATCH_SIZE}" \
+    --output_path "${RESULTS_DIR}/primus" \
+    2>&1 | tee "${RESULTS_DIR}/primus_eval.log"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 3: Evaluate FLA HF reference
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 3] lm-eval on FLA HF =========="
+python tools/hybrid/eval_gdn_lm_eval.py \
+    --model hf \
+    --model_args "pretrained=${FLA_HF_DIR},dtype=bfloat16,trust_remote_code=True" \
+    --tasks "${TASKS}" \
+    --batch_size "${BATCH_SIZE}" \
+    --output_path "${RESULTS_DIR}/fla" \
+    2>&1 | tee "${RESULTS_DIR}/fla_eval.log"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 4: Compact side-by-side scoreboard
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 4] Side-by-side scoreboard =========="
+python tools/hybrid/compare_hybrid_eval.py \
+    --primus-dir "${RESULTS_DIR}/primus" \
+    --fla-dir "${RESULTS_DIR}/fla" \
+    2>&1 | tee "${RESULTS_DIR}/scoreboard.txt"
+
+echo
+echo "DONE. Results saved to ${RESULTS_DIR}/"
diff --git a/tools/hybrid/run_mamba_hybrid_eval.sh b/tools/hybrid/run_mamba_hybrid_eval.sh
new file mode 100644
index 000000000..0b8228a2f
--- /dev/null
+++ b/tools/hybrid/run_mamba_hybrid_eval.sh
@@ -0,0 +1,127 @@
+#!/bin/bash
+###############################################################################
+# End-to-end evaluation of the 75% Hybrid Mamba2+MLA 300M model.
+#
+# Run this *inside the rocm/primus container* with the repo mounted at
+# /home/<user>/Primus:
+#
+#   cd /home/<user>/Primus && bash tools/hybrid/run_mamba_hybrid_eval.sh
+#
+# What it does:
+#   1. Converts the Primus Megatron checkpoint → FLA HuggingFace
+#      `Mamba2ForCausalLM` format using tools/hybrid/convert_mamba_hybrid_to_fla_hf.py
+#      (handles the 3 MLA + 9 Mamba2 sublayer mix; MLA path reuses the same
+#      channel-permutation fix as the GDN-hybrid converter)
+#   2. Runs lm-eval on the Primus-converted HF model
+#   3. Runs lm-eval on FLA's reference HF checkpoint
+#   4. Diffs the two scoreboards
+#
+# ARCHITECTURE NOTE — the two models are NOT bit-compatible:
+#
+#                   Primus YAML          FLA mamba2_300M_hybrid.json
+#   hidden_size     1024                 1216
+#   intermediate    4096                 4864
+#   state_size      64                   128
+#   n_groups        8                    1
+#   total params   ~273M                ~350M
+#
+#   We sized Primus to match the GDN-hybrid 300M so the Mamba2 mixer is a
+#   drop-in swap into the same backbone.  The eval below scores each model
+#   on its own arch; treat the comparison as a "what 300M-ish hybrids land
+#   at this loss" benchmark, not a checkpoint-conversion parity check.
+###############################################################################
+set -euo pipefail
+
+PRIMUS_CKPT=${PRIMUS_CKPT:-output/amd/root/zebra_llama_300M_mamba_hybrid-pretrain/checkpoints/iter_0004768}
+PRIMUS_HF_DIR=${PRIMUS_HF_DIR:-output/mamba_hybrid_300M_fla_hf}
+FLA_HF_DIR=${FLA_HF_DIR:-$HOME/checkpoints/mamba2_hybrid_300M_10B/checkpoint-4768}
+RESULTS_DIR=${RESULTS_DIR:-output/mamba_hybrid_300M_eval_results}
+TASKS=${TASKS:-arc_easy,arc_challenge,hellaswag,openbookqa,piqa,winogrande,mmlu,race}
+BATCH_SIZE=${BATCH_SIZE:-auto}
+TOKENIZER=${TOKENIZER:-$HOME/checkpoints/gdn_pure_300M_10B}
+
+echo "==========[run_mamba_hybrid_eval.sh]=========="
+echo "PRIMUS_CKPT     = ${PRIMUS_CKPT}"
+echo "PRIMUS_HF_DIR   = ${PRIMUS_HF_DIR}"
+echo "FLA_HF_DIR      = ${FLA_HF_DIR}"
+echo "RESULTS_DIR     = ${RESULTS_DIR}"
+echo "TASKS           = ${TASKS}"
+echo "BATCH_SIZE      = ${BATCH_SIZE}"
+echo "TOKENIZER       = ${TOKENIZER}"
+echo "================================================"
+
+mkdir -p "${RESULTS_DIR}"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 1: Convert Primus checkpoint to FLA HF format
+# ─────────────────────────────────────────────────────────────────────────────
+if [ ! -f "${PRIMUS_HF_DIR}/model.safetensors" ]; then
+    echo
+    echo "==========[Step 1] Converting Primus → FLA HF (Mamba2ForCausalLM) =========="
+    python tools/hybrid/convert_mamba_hybrid_to_fla_hf.py \
+        --checkpoint-path "${PRIMUS_CKPT}" \
+        --output-dir "${PRIMUS_HF_DIR}" \
+        --tokenizer "${TOKENIZER}"
+else
+    echo "Primus HF checkpoint already exists at ${PRIMUS_HF_DIR}, skipping conversion."
+fi
+
+# Ensure tokenizer files are present in both eval dirs
+for tdir in "${PRIMUS_HF_DIR}" "${FLA_HF_DIR}"; do
+    if [ ! -f "${tdir}/tokenizer.json" ] && [ -f "${TOKENIZER}/tokenizer.json" ]; then
+        cp "${TOKENIZER}/tokenizer.json" \
+           "${TOKENIZER}/tokenizer_config.json" \
+           "${TOKENIZER}/special_tokens_map.json" \
+           "${tdir}/" 2>/dev/null || true
+        echo "Copied tokenizer from ${TOKENIZER} → ${tdir}"
+    fi
+done
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 1b: Sanity-load both HF models, confirm forward + param count
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 1b] Sanity-load both HF models =========="
+# Uses AutoModelForCausalLM with trust_remote_code so the Primus-converted
+# checkpoint picks up the custom Mamba2FullMlpForCausalLM class (which keeps
+# MLPs on every layer; stock FLA Mamba2 only has MLPs on MLA layers).
+PRIMUS_HF_DIR="${PRIMUS_HF_DIR}" FLA_HF_DIR="${FLA_HF_DIR}" python tools/_sanity_load_mamba.py
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 2: Evaluate Primus-converted HF
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 2] lm-eval on Primus HF (mamba2 hybrid) =========="
+python tools/hybrid/eval_mamba2_lm_eval.py \
+    --model hf \
+    --model_args "pretrained=${PRIMUS_HF_DIR},dtype=bfloat16,trust_remote_code=True" \
+    --tasks "${TASKS}" \
+    --batch_size "${BATCH_SIZE}" \
+    --output_path "${RESULTS_DIR}/primus" \
+    2>&1 | tee "${RESULTS_DIR}/primus_eval.log"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 3: Evaluate FLA HF reference
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 3] lm-eval on FLA HF (mamba2 hybrid) =========="
+python tools/hybrid/eval_mamba2_lm_eval.py \
+    --model hf \
+    --model_args "pretrained=${FLA_HF_DIR},dtype=bfloat16,trust_remote_code=True" \
+    --tasks "${TASKS}" \
+    --batch_size "${BATCH_SIZE}" \
+    --output_path "${RESULTS_DIR}/fla" \
+    2>&1 | tee "${RESULTS_DIR}/fla_eval.log"
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Step 4: Compact side-by-side scoreboard
+# ─────────────────────────────────────────────────────────────────────────────
+echo
+echo "==========[Step 4] Side-by-side scoreboard =========="
+python tools/hybrid/compare_hybrid_eval.py \
+    --primus-dir "${RESULTS_DIR}/primus" \
+    --fla-dir "${RESULTS_DIR}/fla" \
+    2>&1 | tee "${RESULTS_DIR}/scoreboard.txt"
+
+echo
+echo "DONE. Results saved to ${RESULTS_DIR}/"
diff --git a/tools/hybrid/run_zebra_eval.sh b/tools/hybrid/run_zebra_eval.sh
new file mode 100644
index 000000000..5de2f738e
--- /dev/null
+++ b/tools/hybrid/run_zebra_eval.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+python3 tools/hybrid/lm_harness_eval.py --model zebra_llama \
+    --model_args pretrained=output/zebra_llama_1B_hf_iter_0200000,dtype=bfloat16  \
+    --tasks arc_easy,arc_challenge,hellaswag,winogrande,piqa,race,openbookqa \
+    --batch_size 32
diff --git a/tools/hybrid/verify_gdn_conversion.py b/tools/hybrid/verify_gdn_conversion.py
new file mode 100644
index 000000000..7458c552b
--- /dev/null
+++ b/tools/hybrid/verify_gdn_conversion.py
@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""
+Sanity check for Primus GDN → FLA HuggingFace conversion.
+
+Loads the converted model and verifies:
+1. All weights loaded without errors
+2. Forward pass produces reasonable loss (~2-4 for a trained model)
+3. Top predictions are sensible English tokens
+
+Usage:
+    python tools/hybrid/verify_gdn_conversion.py --model-path output/gdn_pure_1B_fla_hf
+"""
+
+import argparse
+import sys
+
+import torch
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Verify GDN checkpoint conversion")
+    parser.add_argument(
+        "--model-path", type=str, required=True, help="Path to converted FLA HuggingFace model directory"
+    )
+    parser.add_argument("--tokenizer", type=str, default="meta-llama/Llama-3.2-1B", help="Tokenizer to use")
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("GDN Conversion Verification")
+    print("=" * 60)
+
+    from fla.models.gated_deltanet import GatedDeltaNetConfig, GatedDeltaNetForCausalLM
+    from transformers import AutoTokenizer
+
+    print(f"\nLoading config from: {args.model_path}")
+    config = GatedDeltaNetConfig.from_pretrained(args.model_path)
+    print(
+        f"  hidden_size={config.hidden_size}, num_heads={config.num_heads}, "
+        f"num_v_heads={config.num_v_heads}, num_layers={config.num_hidden_layers}"
+    )
+
+    print(f"\nLoading model...")
+    model = (
+        GatedDeltaNetForCausalLM.from_pretrained(
+            args.model_path,
+            config=config,
+            torch_dtype=torch.bfloat16,
+        )
+        .cuda()
+        .eval()
+    )
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"  Parameters: {num_params / 1e9:.3f}B")
+
+    print(f"\nLoading tokenizer: {args.tokenizer}")
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+
+    prompts = [
+        "The capital of France is",
+        "Machine learning is a field of",
+        "The largest planet in our solar system is",
+    ]
+
+    all_passed = True
+    for prompt in prompts:
+        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            out = model(**inputs, labels=inputs["input_ids"])
+
+        loss = out.loss.item()
+        logits = out.logits[0, -1]
+        top5 = torch.topk(logits, 5)
+
+        status = "PASS" if loss < 6.0 else "FAIL"
+        if loss > 6.0:
+            all_passed = False
+
+        # Greedy continuation — most intuitive coherence signal
+        with torch.no_grad():
+            gen = model.generate(
+                **inputs,
+                max_new_tokens=40,
+                do_sample=False,
+                temperature=1.0,
+                pad_token_id=tokenizer.eos_token_id or 0,
+            )
+        cont = tokenizer.decode(gen[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
+
+        print(f'\n  Prompt: "{prompt}"')
+        print(f"  Loss: {loss:.4f} [{status}]")
+        print(f"  Top-5: {[tokenizer.decode(t) for t in top5.indices]}")
+        print(f'  Greedy continuation: "{cont}"')
+
+    print(f"\n{'=' * 60}")
+    if all_passed:
+        print("RESULT: ALL CHECKS PASSED")
+        print("  - Model loads cleanly")
+        print("  - Forward pass produces reasonable loss")
+        print("  - Predictions are sensible")
+        print(f"{'=' * 60}")
+        return 0
+    else:
+        print("RESULT: SOME CHECKS FAILED")
+        print("  - Loss too high (>6.0) suggests weight mapping issue")
+        print(f"{'=' * 60}")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/installation-jax/README.md b/tools/installation-jax/README.md
new file mode 100644
index 000000000..82f6be57f
--- /dev/null
+++ b/tools/installation-jax/README.md
@@ -0,0 +1,165 @@
+# Primus JAX / MaxText environment in a venv (no docker, no sudo) — v26.5
+
+Reproduces the Primus **v26.5 JAX training Dockerfile** in a Python virtual
+environment. Same package pins as the Dockerfile, adapted for a bare-metal host
+with no root and no containers.
+
+This is the JAX/MaxText counterpart to `tools/installation/` (the PyTorch /
+Megatron / TorchTitan stack). Use this one if you want to run the **JAX MaxText**
+training backend of Primus.
+
+## Why it differs from the Dockerfile
+
+| Constraint | Dockerfile | Here |
+|---|---|---|
+| ROCm | release **tarball** → `/opt/rocm` | same tarball → **`$ROCM_DIR`** (`$PRIMUS_JAX_BASE/rocm`, user-writable, **no sudo**) |
+| GPU arch | gfx942 + gfx950 | **auto-detected** from the host (`rocminfo`/KFD sysfs) for source builds/TE |
+| Build dir | container FS | venv on **`$PRIMUS_JAX_BASE`** (persistent); transient sources on local `/tmp` |
+| System deps | `apt install ...` | **skipped** (no sudo); documented in the guide's Section 2 |
+| MaxText setup | `setup.sh` (runs `apt`, prompts for a venv) | Python steps only (apt is a one-time root action; venv already made) |
+
+The key reason no sudo is required: the ROCm tarball is extracted into a
+user-writable dir (`$ROCM_DIR`) instead of `/opt/rocm`, so we don't depend on
+system ROCm or apt for the ROCm toolchain itself.
+
+> **Python 3.12+ required (3.12 preferred).** MaxText requires Python ≥ 3.12 (the
+> PyTorch recipe works on 3.10; this one does not). 3.12 is preferred because the
+> prebuilt ROCm wheels are `cp312`. You do **not** need `sudo` or a PPA:
+> - The scripts use `python3.12` (preferred) or `python3.13` if on PATH.
+> - Otherwise, if [`uv`](https://docs.astral.sh/uv/) is installed, `env.sh`
+>   auto-detects a uv-managed `>= 3.12` interpreter and `setup.sh` runs
+>   `uv python install 3.12` for you when none exists yet.
+> - No `uv`? Install it once (no root) and re-run:
+>   `python3 -m pip install --user uv` (or `curl -LsSf https://astral.sh/uv/install.sh | sh`).
+> - To force a specific interpreter: `export PRIMUS_PYTHON=/path/to/python3.12`.
+>
+> On Ubuntu 22.04, `apt install python3.12` fails (jammy has no such package) —
+> use the `uv` path above instead.
+
+> **Host OS: Ubuntu 24.04 / glibc ≥ 2.38 recommended, but 22.04 works.** The
+> prebuilt `transformer_engine_rocm_jax` wheel is built against the Dockerfile's
+> `ubuntu:24.04` base and needs `glibc ≥ 2.38` + `GLIBCXX_3.4.32` (GCC 13/14).
+> On Ubuntu 22.04 (glibc 2.35) the prebuilt TE wheel can't load
+> (`version 'GLIBC_2.38' not found`). **`setup.sh` handles this automatically:**
+> the `te` stage detects the host glibc and builds TransformerEngine from source
+> when it is < 2.38, so `bash setup.sh` works on both 22.04 and 24.04. (Only the
+> prebuilt TE wheel needs glibc ≥ 2.38; the ROCm JAX/PJRT wheels load on glibc
+> 2.35.) Check with `ldd --version`.
+
+## Run it
+
+```bash
+cd tools/installation-jax
+# PRIMUS_JAX_BASE is REQUIRED (no default). Point it at a directory you can write
+# to with tens of GB free — the venv, ROCm tarball, and checkouts all live here:
+export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env
+bash setup.sh                 # all default stages
+```
+
+If any stage fails the script stops immediately and prints which stage failed;
+fix the cause and re-run just that stage. Stages are idempotent:
+
+```bash
+bash setup.sh --list          # show stages
+bash setup.sh te              # reinstall just TransformerEngine
+bash setup.sh venv rocm jax   # venv + ROCm + JAX only
+```
+
+## Use the environment afterward
+
+```bash
+# Set the SAME PRIMUS_JAX_BASE you built with (required — env.sh errors without it)
+export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env
+source tools/installation-jax/env.sh   # activates venv + sets ROCm / NVTE / XLA env vars
+python -c "import jax; print(jax.devices())"
+
+# Primus is checked out at $WORKSPACE_DIR/Primus; MaxText at $MAXTEXT_DIR.
+cd "$WORKSPACE_DIR/Primus"
+./primus-cli direct -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+```
+
+`env.sh` exports `MAXTEXT_PATH=$MAXTEXT_DIR`, so Primus runs the same MaxText
+checkout we installed the dependencies for.
+
+## Stages (default order, v26.5)
+
+`venv` → `rocm` → `maxtext` → `tf_source` → `jax` → `te` → `primus`
+→ `jaxreqs` → `rccl` → `manifest`
+
+- **venv** — create the venv (Python ≥ 3.12) and bootstrap `cmake`/`ninja`/`uv`.
+- **rocm** — download + extract the TheRock ROCm release tarball into `$ROCM_DIR`
+  (+ `amdsmi`).
+- **maxtext** — clone ROCm/MaxText (`release/v26.5`) and install its deps (the
+  Python part of MaxText's `setup.sh`) + the editable MaxText package.
+- **tf_source** — build **tensorflow-cpu 2.21 from source** (bazel, ~30–60 min);
+  fixes the ROCm-vs-TF LLVM symbol clash (SIGSEGV) and drops bundled NCCL.
+- **jax** — `jax`/`jaxlib` 0.10.0 + the ROCm `jax_rocm7_pjrt` / `jax_rocm7_plugin`
+  (installed after MaxText to override its stock jax).
+- **te** — prebuilt `transformer_engine_rocm_jax` wheel (+ `flax`, `pydantic`, ...);
+  auto-falls-back to a from-source build (`te_source`) on glibc < 2.38 hosts.
+- **primus** — clone Primus, init the `third_party/maxtext` submodule, drop the
+  stale `dataclasses` backports.
+- **jaxreqs** — install Primus' `requirements-jax.txt` (loguru, wandb, ...).
+- **rccl** — build **RCCL from source** and install it into `$ROCM_PATH/lib`.
+- **manifest** — dump `pip list` / `env` for reproducibility.
+
+Optional / alternative stages:
+
+- **te_source** — force the from-source TransformerEngine build regardless of
+  glibc. Normally unnecessary: the default `te` stage **auto-detects the host
+  glibc** and builds from source when it is < 2.38 (e.g. Ubuntu 22.04), where the
+  prebuilt wheel fails to load with `version 'GLIBC_2.38' not found`. Heavy build
+  (~30–60 min, compiles CK fused-attention kernels). Only the prebuilt TE wheel
+  needs glibc ≥ 2.38 — the ROCm JAX/PJRT wheels load fine on glibc 2.35.
+- **tf_cpu_fix** — lighter alternative to `tf_source`: `pip install tensorflow-cpu`
+  instead of the bazel build (avoids the bundled-NCCL clash; may still hit the
+  LLVM-symbol SIGSEGV on some ROCm 7.14 configs).
+
+```bash
+# Ubuntu 22.04 (glibc < 2.38): `bash setup.sh` already builds TE from source
+# automatically. To also skip the heavy tf_source bazel build, swap in tf_cpu_fix:
+bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest
+```
+
+> **MaxText v26.5 and Primus.** MaxText v26.5 uses a **2-value**
+> `initialize()`/`run()` API. Primus `main` handles it — `MaxTextPretrainTrainer`
+> forwards `initialize()`'s tuple verbatim to `run()` (fix #912) — so this
+> recipe's pinned `release/v26.5` trains out of the box. Override
+> `MAXTEXT_BRANCH` only if you deliberately need a different MaxText release.
+
+## What is SKIPPED (needs sudo / apt — not reproducible here)
+
+- **System packages** (build toolchain, `numactl`, `gcsfuse`, RDMA/verbs libs):
+  a one-time root action, documented in Section 2 of the guide.
+- **AINIC** (`add-apt-repository`, `libionic-dev`): apt-only. Skipped.
+- **UCX + OpenMPI**: **optional for MaxText** — JAX uses its own distributed
+  coordinator + RCCL, not `mpirun`. Carried over from the reference image for
+  other/MPI-launched JAX workloads; skipped here. See Section 4 of the guide.
+- **gcsfuse**: only needed to mount GCS buckets for data; not required for
+  synthetic-data or local-data runs.
+
+## Notes / gotchas vs. the Dockerfile
+
+- **Stage order matters (v26.5):** `maxtext` → `tf_source` → `jax` → `te`. MaxText's
+  `setup.sh` pulls in a stock `jax`/`tensorflow`; TF is then rebuilt from source and
+  the ROCm JAX/plugin is installed after (overriding MaxText's), and TE must come
+  after JAX or `jaxlib` gets clobbered. The stage order enforces this.
+- **TensorFlow + RCCL are built from source** in the default flow (matching the
+  v26.5 image). These are the two long builds; use `tf_cpu_fix` if you want to skip
+  the TF bazel build.
+- **No FlashAttention/aiter/torch stack.** The JAX MaxText path does not build the
+  PyTorch kernel libraries; attention fusion comes from `transformer_engine_rocm_jax`
+  + XLA.
+- **Two MaxText checkouts in the Docker image** (`/workspace/maxtext` and
+  `Primus/third_party/maxtext`) are collapsed here: we install deps from one
+  checkout and point `MAXTEXT_PATH` at it.
+- **Runtime-validated.** The default flow — with `te_source` substituted for the
+  prebuilt `te` — has been run end-to-end for single-node MaxText pretraining on
+  **gfx942 / Ubuntu 22.04 (glibc 2.35) / Python 3.12**, and a **2-node** run has
+  also been exercised (over the JAX distributed coordinator + RCCL, no
+  UCX/OpenMPI). On Ubuntu 24.04 (glibc ≥ 2.38) the prebuilt `te` wheel works
+  directly.
+
+Treat the reference `Dockerfile` as the authoritative, tested version
+combination; if you bump one pin you may need to bump the others.
diff --git a/tools/installation-jax/env.sh b/tools/installation-jax/env.sh
new file mode 100644
index 000000000..dc84d002d
--- /dev/null
+++ b/tools/installation-jax/env.sh
@@ -0,0 +1,202 @@
+#!/usr/bin/env bash
+# env.sh — Primus JAX/MaxText venv environment (v26.5).
+# Source this both during the build (setup.sh does it) and every time you
+# want to USE the environment:   source env.sh
+#
+# This mirrors the v26.5 JAX training Dockerfile, adapted for a no-docker /
+# no-sudo bare-metal install. In v26.5 ROCm is delivered as a RELEASE TARBALL
+# (TheRock dist) extracted to a user-writable dir ($ROCM_DIR) — NOT the pip
+# `rocm-sdk-devel` wheels used by the v26.4 recipe. No system ROCm is required.
+#
+# NOTE: PRIMUS_JAX_BASE is REQUIRED (no default) — export it to a big-disk dir
+#       before sourcing this file. MaxText requires Python >= 3.12; see
+#       PRIMUS_PYTHON below.
+
+# ---- Install location (persistent) + transient build sources ----
+# PRIMUS_JAX_BASE is REQUIRED (there is intentionally no default): it is where
+# the venv, the extracted ROCm tarball, and the kept checkouts live. Point it at
+# a directory you can write to with tens of GB free, then (re-)source this file:
+#     export PRIMUS_JAX_BASE=/big/disk/primus-jax-env
+if [ -z "${PRIMUS_JAX_BASE:-}" ]; then
+    echo "[env] ERROR: PRIMUS_JAX_BASE is not set (no default). Set it to a directory" >&2
+    echo "[env]        you can write to with tens of GB free, then re-run:" >&2
+    echo "[env]            export PRIMUS_JAX_BASE=/big/disk/primus-jax-env" >&2
+    # env.sh is always sourced (by setup.sh and to use the env); stop sourcing
+    # without killing an interactive shell. Fall back to exit if run directly.
+    # shellcheck disable=SC2317  # 'exit 1' is reachable when this file is executed, not sourced
+    return 1 2>/dev/null || exit 1
+fi
+export PRIMUS_JAX_BASE
+export VENV_DIR="${VENV_DIR:-$PRIMUS_JAX_BASE/venv}"
+export WORKSPACE_DIR="${WORKSPACE_DIR:-$PRIMUS_JAX_BASE/workspace}"  # kept checkouts (maxtext, Primus)
+# Transient build sources: put on fast LOCAL /tmp (NFS is slow for compile I/O;
+# these dirs are deleted after each build anyway).
+export SRC_DIR="${SRC_DIR:-/tmp/primus-jax-build}"
+# ROCm release tarball (TheRock dist) is extracted here. This replaces the
+# Dockerfile's /opt/rocm and keeps the whole install no-sudo / user-writable.
+export ROCM_DIR="${ROCM_DIR:-$PRIMUS_JAX_BASE/rocm}"
+
+# MaxText clone (deps + editable install live here). Primus resolves the
+# MaxText backend from MAXTEXT_PATH (else its own third_party/maxtext), so we
+# point it at the checkout we actually installed the deps for.
+export MAXTEXT_DIR="${MAXTEXT_DIR:-$WORKSPACE_DIR/maxtext}"
+export MAXTEXT_PATH="${MAXTEXT_PATH:-$MAXTEXT_DIR}"
+
+# Python interpreter used to CREATE the venv. MaxText needs >= 3.12.
+# 3.12 is PREFERRED: it is the version the reference image (ubuntu:24.04) uses,
+# and the prebuilt ROCm wheels (transformer_engine_rocm_jax, jax_rocm7_*) are
+# built for cp312 — a 3.13 venv would need those built from source instead.
+# Resolution order:
+#   1. An explicit python3.12 / python3.13 on PATH (3.12 first).
+#   2. A uv-managed interpreter (>= 3.12), if `uv` is installed. This is the
+#      no-sudo path on distros whose apt has no python3.12 (e.g. Ubuntu 22.04):
+#      `uv python install 3.12` provides one without root. setup.sh will run
+#      that install automatically if none is found yet.
+#   3. Fall back to python3 (setup.sh errors clearly if that is < 3.12).
+if [ -z "${PRIMUS_PYTHON:-}" ]; then
+    for _py in python3.12 python3.13; do
+        if command -v "$_py" >/dev/null 2>&1; then export PRIMUS_PYTHON="$_py"; break; fi
+    done
+    if [ -z "${PRIMUS_PYTHON:-}" ] && command -v uv >/dev/null 2>&1; then
+        _uv_py="$(uv python find '>=3.12' 2>/dev/null || true)"
+        [ -n "$_uv_py" ] && [ -x "$_uv_py" ] && export PRIMUS_PYTHON="$_uv_py"
+    fi
+    export PRIMUS_PYTHON="${PRIMUS_PYTHON:-python3}"
+fi
+
+# Build parallelism.
+export MAX_JOBS="${MAX_JOBS:-128}"
+
+# ---- Target GPU architecture (auto-detected) ----
+# PYTORCH_ROCM_ARCH controls the gfx targets used for the source builds (TE,
+# RCCL) and TE's runtime NVTE_ROCM_ARCH. The name is kept for parity with the
+# Dockerfile / TE which read it. If you set it yourself it is respected as-is:
+#     export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+# Otherwise it is auto-detected from the GPUs on this host. Detection needs no
+# sudo and no system ROCm: it uses `rocminfo` when available (e.g. after the
+# ROCm tarball is extracted) and otherwise falls back to the kernel KFD sysfs
+# topology, so it works even on a fresh machine before any install.
+
+# Decode a KFD gfx_target_version integer (e.g. 90402) into a gfx name (gfx942).
+_primus_gfxver_to_arch() {
+    local v="$1"
+    [ -n "$v" ] && [ "$v" != "0" ] || return 0
+    printf 'gfx%d%x%x\n' "$(( v / 10000 ))" "$(( (v / 100) % 100 ))" "$(( v % 100 ))"
+}
+
+# Echo a ";"-separated, de-duplicated list of detected gfx architectures.
+_primus_detect_gpu_arch() {
+    local arches=""
+    if command -v rocminfo >/dev/null 2>&1; then
+        arches="$(rocminfo 2>/dev/null \
+            | awk '/^[[:space:]]*Name:[[:space:]]*gfx/{print $2}' \
+            | sort -u | paste -sd';' -)"
+    fi
+    if [ -z "$arches" ] && [ -d /sys/class/kfd/kfd/topology/nodes ]; then
+        local f v a list=""
+        for f in /sys/class/kfd/kfd/topology/nodes/*/properties; do
+            [ -f "$f" ] || continue
+            v="$(awk '/^gfx_target_version/{print $2}' "$f" 2>/dev/null)"
+            a="$(_primus_gfxver_to_arch "$v")"
+            [ -n "$a" ] && list="${list}${a}"$'\n'
+        done
+        arches="$(printf '%s' "$list" | sort -u | sed '/^$/d' | paste -sd';' -)"
+    fi
+    printf '%s' "$arches"
+}
+
+if [ -z "${PYTORCH_ROCM_ARCH:-}" ]; then
+    _DETECTED_ARCH="$(_primus_detect_gpu_arch)"
+    if [ -n "$_DETECTED_ARCH" ]; then
+        export PYTORCH_ROCM_ARCH="$_DETECTED_ARCH"
+        echo "[env] detected GPU arch: PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH" >&2
+    else
+        export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+        echo "[env] WARNING: no GPU detected; defaulting PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH (override by exporting it)" >&2
+    fi
+fi
+
+# Comma-separated variant for the tools/vars that expect commas.
+_ARCH_CSV="${PYTORCH_ROCM_ARCH//;/,}"
+export ROCM_AMDGPU_TARGETS="${ROCM_AMDGPU_TARGETS:-$_ARCH_CSV}"
+export GPU_ARCHS="${GPU_ARCHS:-$PYTORCH_ROCM_ARCH}"
+export HCC_AMDGPU_TARGET="${HCC_AMDGPU_TARGET:-$_ARCH_CSV}"
+export HIP_ARCHITECTURES="${HIP_ARCHITECTURES:-$_ARCH_CSV}"
+
+# Keep caches OFF the tiny home quota. Use local /tmp.
+export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-jax-cache/pip}"
+export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/primus-jax-cache/xdg}"
+export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/primus-jax-cache/uv}"
+mkdir -p "$PIP_CACHE_DIR" "$XDG_CACHE_HOME" "$UV_CACHE_DIR" 2>/dev/null || true
+
+# Activate the venv if it exists
+if [ -f "$VENV_DIR/bin/activate" ]; then
+    # shellcheck disable=SC1091
+    source "$VENV_DIR/bin/activate"
+fi
+
+# NOTE: the two HSA_* vars below are NOT in the v26.5 Dockerfile — they are a
+# carryover workaround for HSA_STATUS_ERROR_OUT_OF_RESOURCES. Harmless, but
+# remove them if you want an exact match to the image.
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+export HSA_NO_SCRATCH_RECLAIM=1
+
+# Flags to fix the ROCm profiler hang issue (from the JAX Dockerfile base stage)
+export ROCPROFILER_QUEUE_INTERPOSITION=0
+export DEBUG_HIP_DYNAMIC_QUEUES=0
+
+# ---- ROCm path: from the extracted TheRock release tarball (v26.5) ----
+# setup.sh's `rocm` stage downloads therock-dist-*.tar.gz and extracts it into
+# $ROCM_DIR. Once present, wire the HIP/ROCm env at it (mirrors the Dockerfile's
+# /opt/rocm layout, whose LD path is lib + lib/rocm_sysdeps/lib).
+if [ -d "$ROCM_DIR/lib" ] || [ -x "$ROCM_DIR/bin/hipcc" ]; then
+    export ROCM_PATH="$ROCM_DIR"
+    export ROCM_HOME="$ROCM_DIR"   # Primus uses ROCM_HOME; set both
+    export HIP_PLATFORM=amd
+    export HIP_PATH="$ROCM_PATH"
+    export HIP_CLANG_PATH="$ROCM_PATH/llvm/bin"
+    export HIP_INCLUDE_PATH="$ROCM_PATH/include"
+    export HIP_LIB_PATH="$ROCM_PATH/lib"
+    export HIP_DEVICE_LIB_PATH="$ROCM_PATH/lib/llvm/amdgcn/bitcode"
+    # PATH: the Dockerfile puts /opt/rocm/lib on PATH too; mirror that.
+    export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
+    export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:${LD_LIBRARY_PATH:-}"
+    export LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64"
+    export CPATH="$HIP_INCLUDE_PATH"
+    export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+fi
+
+# ---- TransformerEngine (ROCm) runtime settings for JAX ----
+# The transformer_engine_rocm_jax wheel is prebuilt, so these are runtime knobs.
+export NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH"
+export NVTE_USE_ROCM=1
+export NVTE_USE_HIPBLASLT=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
+export NVTE_FUSED_ATTN=1
+export NVTE_CK_USES_BWD_V3=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+export NVTE_CK_HOW_V3_BF16_CVT=2
+
+# ---- AMD GPU runtime knobs (from the JAX Dockerfile) ----
+export GPU_MAX_HW_QUEUES=2
+export HIP_FORCE_DEV_KERNARG=1
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_DEBUG=VERSION
+
+# Bare-metal-only: force RCCL to use its built-in ROCm IB/RoCE transport instead
+# of any host NCCL net plugin. The rocm/jax-training image ships no
+# librccl-net.so and sets no NCCL_NET_PLUGIN, but on a host the system
+# /usr/local/lib/librccl-net.so is on the default loader path and is
+# ABI-incompatible with the from-source RCCL (undefined symbol
+# ncclNetPlugin_v11/_v10 -> falls back to v9 -> segfault at clique init).
+# NCCL_NET_PLUGIN=none makes RCCL ignore that host plugin, matching the image.
+export NCCL_NET_PLUGIN=none
+
+# NOT in the v26.5 Dockerfile: avoids a NaN-loss issue when training on gfx950
+# (no-op on gfx942). Keep it if you train on MI350/MI355; drop for exact parity.
+export RCCL_WARP_SPEED_AUTO=0
+
+# ---- XLA / JAX runtime settings ----
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.9
+export XLA_FLAGS="${XLA_FLAGS:---xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_enable_command_buffer=''}"
diff --git a/tools/installation-jax/setup.sh b/tools/installation-jax/setup.sh
new file mode 100644
index 000000000..a388450a4
--- /dev/null
+++ b/tools/installation-jax/setup.sh
@@ -0,0 +1,412 @@
+#!/usr/bin/env bash
+# setup.sh — Reproduce the Primus JAX/MaxText training environment (v26.5) in a
+# Python venv (no sudo, no docker). Mirrors the v26.5 JAX training Dockerfile,
+# adapted for bare metal:
+#   * ROCm from a RELEASE TARBALL (TheRock dist) extracted to $ROCM_DIR, under
+#     $PRIMUS_JAX_BASE (no sudo, no /opt/rocm) — replaces the v26.4 pip wheels
+#   * builds/checkouts on a big disk (home quota is usually tiny)
+#   * GPU arch auto-detected (gfx942 and/or gfx950); apt/sudo steps skipped
+#   * MaxText's own setup.sh apt/interactive steps skipped (system packages are
+#     a one-time root action, documented in the guide's Section 2)
+#   * TensorFlow (2.21 CPU) and RCCL are built from source in the default flow
+#     (matching the v26.5 image)
+#
+# Usage:
+#   bash setup.sh                # run all default stages in order
+#   bash setup.sh <stage>...     # run only specific stage(s), e.g.
+#   bash setup.sh venv rocm jax
+#
+# Stages are re-runnable. List them with:  bash setup.sh --list
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# shellcheck disable=SC1091
+source "$SCRIPT_DIR/env.sh"
+
+log()  { echo -e "\n\033[1;36m[setup] $*\033[0m"; }
+die()  { echo -e "\033[1;31m[setup][ERROR] $*\033[0m" >&2; exit 1; }
+# shellcheck disable=SC1091
+reload_env() { source "$SCRIPT_DIR/env.sh"; }
+
+# ---- pinned versions / commits (from the v26.5 JAX Dockerfile) ----
+# ROCm: TheRock RELEASE TARBALL (not pip wheels). Extracted into $ROCM_DIR.
+# See: https://repo.amd.com/rocm/tarball-multi-arch/
+THEROCK_TARBALL="${THEROCK_TARBALL:-https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz}"
+
+# JAX + ROCm PJRT/plugin
+# See: https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-pjrt/
+#      https://repo.amd.com/rocm/whl-multi-arch/jax-rocm7-plugin/
+JAX_VERSION="0.10.0"
+JAX_PJRT_VERSION="0.10.0+rocm7.14.0"
+JAX_PLUGIN_VERSION="0.10.0+rocm7.14.0"
+JAX_INDEX="https://repo.amd.com/rocm/whl-multi-arch/"
+
+# TransformerEngine (prebuilt ROCm JAX wheel)
+# See: https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/transformer-engine-rocm-jax/
+TE_VERSION="2.15.0.dev0+rocm7.15.0a20260707.72d01a0"
+TE_INDEX="https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/"
+# From-source TransformerEngine (used by the `te_source` stage). Needed on hosts
+# whose glibc is older than the prebuilt wheel requires (< 2.38, e.g. Ubuntu
+# 22.04): building against the host toolchain links TE to the local glibc so it
+# actually loads. (Commit is the last known-good ROCm TE JAX source build; bump
+# to match the prebuilt wheel above when a matching source tag is published.)
+TE_REPO="https://github.com/ROCm/TransformerEngine.git"
+TE_SOURCE_COMMIT="${TE_SOURCE_COMMIT:-635d7c085c39a6d9bfe4881c7d3efab7a46d7129}"
+
+# TensorFlow (CPU) built from source — replaces the stock PyPI wheel, whose
+# bundled LLVM collides with ROCm's libLLVM in Grain workers (SIGSEGV on
+# `import tensorflow` after `import jax`). A CPU build also drops the bundled
+# NCCL, preserving the XLA->RCCL collective fix.
+TF_REPO="https://github.com/ROCm/tensorflow-upstream.git"
+TF_BRANCH="upstream-v2.21.0"
+BAZELISK_VERSION="v1.25.0"
+
+# RCCL built from source (rocm-systems), installed into the ROCm tree.
+RCCL_REPO="https://github.com/ROCm/rocm-systems.git"
+RCCL_COMMIT="9e5e4084a4b8e1e86551b0eb054725c62354a926"
+
+# MaxText (ROCm fork)
+# v26.5 switched MaxText's initialize()/run() to a 2-value API (config, recorder);
+# v26.4 and earlier returned 3 values (adds diagnostic_config). Primus main
+# supports BOTH: MaxTextPretrainTrainer forwards initialize()'s tuple verbatim to
+# run() (fix #912), so v26.5 trains out of the box with Primus main. Override
+# MAXTEXT_BRANCH only if you deliberately need a different MaxText release.
+MAXTEXT_REPO="https://github.com/ROCm/maxtext.git"
+MAXTEXT_BRANCH="${MAXTEXT_BRANCH:-release/v26.5}"
+# Which MaxText requirements set to install. The reference Docker image runs
+# MaxText's setup.sh with defaults (DEVICE=tpu), which — on ROCm — pulls the
+# framework-agnostic deps WITHOUT any CUDA packages. Override to `cuda12` only
+# if you specifically need that set.
+MAXTEXT_DEVICE="${MAXTEXT_DEVICE:-tpu}"
+
+# Primus
+PRIMUS_REPO="https://github.com/AMD-AGI/Primus.git"
+PRIMUS_BRANCH="main"
+
+PIP="python -m pip"
+UVPIP="python -m uv pip"
+
+# Fresh clone helper: clone into transient SRC_DIR, build, then optionally clean.
+fresh_clone() {  # fresh_clone <url> <dir> [extra git clone args...]
+    local url="$1"; local dir="$2"; shift 2
+    rm -rf "${SRC_DIR:?}/$dir"
+    git clone "$@" "$url" "$SRC_DIR/$dir"
+}
+
+# Return 0 if the host glibc is >= 2.38 (what the prebuilt TE/JAX wheels need).
+# On failure to parse (unknown libc) we conservatively return non-zero so the
+# caller falls back to the always-works from-source build.
+_glibc_ge_238() {
+    local v; v="$(ldd --version 2>/dev/null | head -1 | grep -oE '[0-9]+\.[0-9]+$')"
+    [ -n "$v" ] || return 1
+    awk -v v="$v" 'BEGIN{split(v,a,"."); exit !(a[1]>2 || (a[1]==2 && a[2]>=38))}'
+}
+
+# ============================ STAGES ============================
+
+stage_venv() {
+    # MaxText requires Python >= 3.12.
+    _venv_python_ok() {  # _venv_python_ok <python>  -> 0 if >= 3.12
+        local v; v="$("$1" -c 'import sys; print("%d.%d" % sys.version_info[:2])' 2>/dev/null)" || return 1
+        case "$v" in 3.1[2-9]|3.[2-9][0-9]|[4-9].*) return 0 ;; *) return 1 ;; esac
+    }
+
+    # If the resolved interpreter is too old, try to provision Python 3.12 with
+    # uv (no sudo) — the standard fix on Ubuntu 22.04 where apt has no 3.12.
+    if ! _venv_python_ok "$PRIMUS_PYTHON"; then
+        if command -v uv >/dev/null 2>&1; then
+            log "PRIMUS_PYTHON ($PRIMUS_PYTHON) is < 3.12; provisioning Python 3.12 via uv (no sudo)"
+            uv python install 3.12 || die "uv python install 3.12 failed"
+            PRIMUS_PYTHON="$(uv python find '>=3.12' 2>/dev/null || true)"
+            if [ -z "$PRIMUS_PYTHON" ] || [ ! -x "$PRIMUS_PYTHON" ]; then
+                die "could not locate a uv-managed Python >= 3.12 after install"
+            fi
+        else
+            die "MaxText requires Python >= 3.12 but no suitable interpreter was found. Options (no sudo): install uv (pip install --user uv) then re-run — setup.sh will fetch Python 3.12 automatically; or install python3.12 yourself and re-run with PRIMUS_PYTHON=python3.12"
+        fi
+    fi
+    log "Creating venv at $VENV_DIR (interpreter: $PRIMUS_PYTHON -> $("$PRIMUS_PYTHON" --version 2>&1))"
+    mkdir -p "$PRIMUS_JAX_BASE" "$SRC_DIR" "$WORKSPACE_DIR" || die "could not create the install dirs under PRIMUS_JAX_BASE=$PRIMUS_JAX_BASE (SRC_DIR=$SRC_DIR, WORKSPACE_DIR=$WORKSPACE_DIR). Set PRIMUS_JAX_BASE to a directory you can write to with tens of GB free, e.g.  export PRIMUS_JAX_BASE=/big/disk/primus-jax-env"
+    # Recreate the venv if it exists but was built with a too-old Python.
+    if [ -f "$VENV_DIR/bin/activate" ] && ! _venv_python_ok "$VENV_DIR/bin/python"; then
+        log "Existing venv at $VENV_DIR is < Python 3.12; recreating it"
+        rm -rf "$VENV_DIR"
+    fi
+    [ -f "$VENV_DIR/bin/activate" ] || "$PRIMUS_PYTHON" -m venv "$VENV_DIR"
+    reload_env
+    $PIP install --upgrade pip
+    # Build front-end tooling (matches the Dockerfile; add uv, used by MaxText).
+    $PIP uninstall -y wheel || true
+    $PIP install \
+        cmake==3.31.6 \
+        ninja==1.11.1.3 \
+        wheel==0.45.1 \
+        packaging==25.0 \
+        setuptools==69.5.1 \
+        uv
+    rm -rf /root/.cache 2>/dev/null || true
+}
+
+stage_rocm() {
+    reload_env
+    # v26.5: ROCm comes from a TheRock RELEASE TARBALL extracted into $ROCM_DIR
+    # (user-writable, no sudo, no /opt/rocm). This also sidesteps the pip
+    # rocm-sdk-* version-skew that broke hipBLASLt GEMMs in the v26.4 recipe.
+    log "Installing ROCm from tarball into $ROCM_DIR"
+    mkdir -p "$ROCM_DIR" "$SRC_DIR"
+    if [ -x "$ROCM_DIR/bin/hipcc" ] || [ -d "$ROCM_DIR/lib" ]; then
+        log "ROCm already present at $ROCM_DIR (skipping download; rm -rf it to force)"
+    else
+        local tb="$SRC_DIR/therock-dist.tar.gz"
+        log "Downloading $THEROCK_TARBALL"
+        wget -O "$tb" "$THEROCK_TARBALL" || die "ROCm tarball download failed"
+        log "Extracting to $ROCM_DIR"
+        tar -xf "$tb" -C "$ROCM_DIR" || die "ROCm tarball extract failed"
+        rm -f "$tb"
+    fi
+    reload_env
+    [ -n "${ROCM_PATH:-}" ] || die "ROCM_PATH not resolved after extraction (check $ROCM_DIR layout: expected $ROCM_DIR/bin, $ROCM_DIR/lib)"
+    log "ROCM_PATH=$ROCM_PATH"
+    # amdsmi (the Dockerfile installs it via pip after the ROCm tarball)
+    $PIP install amdsmi==7.0.2 || log "WARNING: amdsmi install failed (non-fatal)"
+    ( "$ROCM_PATH/bin/hipcc" --version || hipcc --version ) 2>/dev/null || log "hipcc not found yet; re-source env.sh"
+}
+
+stage_jax() {
+    reload_env
+    log "Installing JAX ${JAX_VERSION} + ROCm PJRT/plugin (v26.5)"
+    # Note: JAX and related libraries need to be installed BEFORE TE and AFTER
+    # MaxText (whose setup.sh pulls in a stock jax/tensorflow we override here).
+    $PIP install "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}" scipy==1.16
+    $PIP install \
+        --index-url "$JAX_INDEX" \
+        --pre "jax_rocm7_pjrt==${JAX_PJRT_VERSION}" \
+        --pre "jax_rocm7_plugin==${JAX_PLUGIN_VERSION}"
+    python -c "import jax; print('jax', jax.__version__); print('devices:', jax.devices())" || \
+        log "WARNING: jax.devices() failed (expected if no GPU is visible on this build host)"
+}
+
+stage_te() {
+    reload_env
+    # The prebuilt wheel targets the Dockerfile's ubuntu:24.04 base (glibc>=2.38).
+    # On older hosts (e.g. Ubuntu 22.04, glibc 2.35) it cannot load, so fall back
+    # transparently to the from-source build, which links against the host glibc.
+    # This keeps `bash setup.sh` (defaults) working on both 22.04 and 24.04.
+    if ! _glibc_ge_238; then
+        log "host glibc < 2.38 (prebuilt TE needs >= 2.38): building TransformerEngine from source instead (te_source)"
+        stage_te_source
+        return
+    fi
+    log "Installing TransformerEngine (JAX) ${TE_VERSION} (prebuilt ROCm wheel)"
+    $PIP install \
+        pybind11==3.0.4 \
+        importlib-metadata==8.7.1 \
+        pydantic==2.13.4 \
+        flax==0.12.2
+    $PIP install \
+        --index-url "$TE_INDEX" \
+        --pre \
+        --no-build-isolation \
+        "transformer_engine_rocm_jax==${TE_VERSION}"
+
+    # The prebuilt wheel targets the Dockerfile's ubuntu:24.04 base (glibc>=2.38,
+    # GCC 13/14 libstdc++). Verify it actually loads NOW — otherwise training
+    # dies later with a silent exit (the launcher swallows the OSError).
+    log "Verifying TransformerEngine loads"
+    if ! python -c "import transformer_engine.jax" 2>/tmp/te_import_err; then
+        cat /tmp/te_import_err >&2 || true
+        if grep -q "GLIBC_2" /tmp/te_import_err 2>/dev/null; then
+            die "TransformerEngine failed to import: your host glibc ($(ldd --version 2>/dev/null | head -1 | grep -oE '[0-9]+\.[0-9]+$')) is older than the prebuilt wheel requires (needs glibc>=2.38 / Ubuntu 24.04). Options: run on Ubuntu 24.04, or build TransformerEngine from source (see docs Section 3.7). glibc cannot be side-loaded via LD_LIBRARY_PATH."
+        fi
+        die "TransformerEngine failed to import (see error above). See docs Section 3.7."
+    fi
+    log "TransformerEngine import OK"
+}
+
+stage_te_source() {
+    reload_env
+    # Build TransformerEngine (JAX) from source, linking against the HOST glibc.
+    # Use this instead of `te` on hosts whose glibc is older than the prebuilt
+    # wheel needs (< 2.38, e.g. Ubuntu 22.04). NOTE: this is a heavy build (CK
+    # fused-attention kernels for your arch) — expect ~30-60 min.
+    log "Building TransformerEngine (JAX) from source @ $TE_SOURCE_COMMIT"
+    $PIP install pybind11==3.0.4 importlib-metadata==8.7.1 pydantic==2.13.4 flax==0.12.2
+    # Remove ALL prebuilt/stale TE variants so TE's install sanity-check sees a
+    # single, self-consistent from-source package.
+    $PIP uninstall -y transformer_engine transformer-engine transformer_engine_rocm7 \
+        transformer_engine_rocm_jax transformer_engine_jax 2>/dev/null || true
+    fresh_clone "$TE_REPO" TransformerEngine --recursive
+    ( cd "$SRC_DIR/TransformerEngine" \
+        && git checkout "$TE_SOURCE_COMMIT" \
+        && git submodule update --init --recursive \
+        && USE_ROCM=1 NVTE_FRAMEWORK=jax NVTE_USE_ROCM=1 \
+           HIP_PATH="$ROCM_PATH" NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH" \
+           CMAKE_BUILD_PARALLEL_LEVEL="$MAX_JOBS" \
+           PYTHONPATH="$SRC_DIR/TransformerEngine/3rdparty/hipify_torch:${PYTHONPATH:-}" \
+           python3 setup.py bdist_wheel \
+        && $PIP install --no-deps --force-reinstall dist/*.whl ) || die "TransformerEngine source build failed"
+    log "Verifying TransformerEngine (from source) loads"
+    python -c "import transformer_engine.jax" || die "TransformerEngine (source) import failed"
+    log "TransformerEngine (source) import OK"
+    rm -rf "$SRC_DIR/TransformerEngine"
+}
+
+stage_maxtext() {
+    reload_env
+    log "Installing MaxText @ $MAXTEXT_BRANCH (kept in $MAXTEXT_DIR)"
+    rm -rf "$MAXTEXT_DIR"
+    git clone "$MAXTEXT_REPO" "$MAXTEXT_DIR" || die "MaxText clone failed"
+    ( cd "$MAXTEXT_DIR" && git checkout "$MAXTEXT_BRANCH" ) || die "MaxText checkout failed"
+
+    # Replicate the Python portion of MaxText's src/dependencies/scripts/setup.sh
+    # for MODE=stable, WORKFLOW=pre-training (the reference image's default).
+    # The apt/gcsfuse and interactive-venv steps of that script are skipped:
+    # system packages are a one-time root action (see the guide's Section 2)
+    # and the venv already exists here.
+    local req="src/dependencies/requirements/generated_requirements/${MAXTEXT_DEVICE}-requirements.txt"
+    [ -f "$MAXTEXT_DIR/$req" ] || die "MaxText requirements not found: $MAXTEXT_DIR/$req"
+    ( cd "$MAXTEXT_DIR" \
+        && $PIP install -U setuptools wheel uv \
+        && $UVPIP install --resolution=lowest -r "$req" \
+        && python -m src.dependencies.scripts.install_pre_train_extra_deps \
+        && $UVPIP install --no-deps -e . ) || die "MaxText dependency install failed"
+}
+
+stage_tf_cpu_fix() {
+    reload_env
+    # Fix: TF 2.20 sets RTLD_GLOBAL which exposes its bundled CUDA-targeting NCCL
+    # symbols globally, causing XLA's ncclCommInitRankConfig to resolve to TF's
+    # NCCL instead of the system ROCm RCCL. tensorflow-cpu has no bundled NCCL.
+    local tfver
+    tfver="$($PIP show tensorflow 2>/dev/null | awk '/^Version:/{print $2}')"
+    if [ -z "$tfver" ]; then
+        log "tensorflow not installed; skipping tensorflow-cpu override (run 'maxtext' stage first)"
+        return 0
+    fi
+    log "Overriding tensorflow with tensorflow-cpu==$tfver (--no-deps) to avoid NCCL symbol clash"
+    $PIP install --no-deps "tensorflow-cpu==${tfver}"
+}
+
+stage_tf_source() {
+    reload_env
+    # v26.5: build tensorflow-cpu 2.21 from ROCm's fork (bazel). The stock PyPI
+    # TF wheel bundles an LLVM whose symbols collide with ROCm's libLLVM in Grain
+    # "spawn" workers -> SIGSEGV on `import tensorflow` after `import jax`. A CPU
+    # build has correct symbol visibility and no bundled NCCL. HEAVY: ~30-60 min.
+    # Needs a host clang/lld (clang-18) and unzip/zip — see the guide's Section 2.
+    # Idempotent: if the CPU build is already installed, skip the ~30-60 min
+    # bazel rebuild (pip uninstall tensorflow-cpu first to force a rebuild).
+    if $PIP show tensorflow-cpu >/dev/null 2>&1; then
+        log "tensorflow-cpu already installed ($($PIP show tensorflow-cpu 2>/dev/null | awk '/^Version:/{print $2}')); skipping source build"
+        return 0
+    fi
+    log "Building tensorflow-cpu ($TF_BRANCH) from source with bazel (~30-60 min)"
+    local bz="$PRIMUS_JAX_BASE/bin/bazel"
+    mkdir -p "$PRIMUS_JAX_BASE/bin"
+    if [ ! -x "$bz" ]; then
+        log "Fetching bazelisk $BAZELISK_VERSION -> $bz"
+        wget -O "$bz" "https://github.com/bazelbuild/bazelisk/releases/download/${BAZELISK_VERSION}/bazelisk-linux-amd64" \
+            || die "bazelisk download failed"
+        chmod +x "$bz"
+    fi
+    export PATH="$PRIMUS_JAX_BASE/bin:$PATH"
+    fresh_clone "$TF_REPO" tensorflow-upstream --depth 1 --branch "$TF_BRANCH"
+    # Build the wheel for the venv's ACTUAL Python (the reference image uses 3.12;
+    # deriving the version keeps the recipe working on hosts whose interpreter is
+    # 3.13). A cp312 wheel will not install into a cp313 venv and vice versa.
+    local pyver pytag
+    pyver="$(python -c 'import sys; print("%d.%d" % sys.version_info[:2])')"
+    pytag="cp${pyver//./}"   # 3.12 -> cp312, 3.13 -> cp313
+    # Keep bazel's (large) output tree off the tiny home quota.
+    ( cd "$SRC_DIR/tensorflow-upstream" \
+        && "$bz" --output_user_root="$SRC_DIR/bazel" build //tensorflow/tools/pip_package:wheel \
+            --repo_env=WHEEL_NAME=tensorflow_cpu \
+            --repo_env=HERMETIC_PYTHON_VERSION="$pyver" ) || die "TensorFlow bazel build failed"
+    $PIP uninstall -y tensorflow tensorflow-cpu tensorflow_cpu || true
+    $PIP install --no-deps "$SRC_DIR"/tensorflow-upstream/bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow_cpu-*-"${pytag}"-"${pytag}"-linux_x86_64.whl \
+        || die "TensorFlow wheel install failed"
+    python -c "import tensorflow as tf; print('tensorflow', tf.__version__)" || log "WARNING: tensorflow import failed"
+    # bazel marks its external/downloaded trees read-only (dirs lose the write
+    # bit), so a plain `rm -rf` fails with "Permission denied". Make them
+    # writable first, and never let cleanup abort the install (TF is already
+    # built + installed at this point).
+    chmod -R u+w "$SRC_DIR/tensorflow-upstream" "$SRC_DIR/bazel" 2>/dev/null || true
+    rm -rf "$SRC_DIR/tensorflow-upstream" "$SRC_DIR/bazel" 2>/dev/null || true
+}
+
+stage_rccl() {
+    reload_env
+    # v26.5: build RCCL from source (rocm-systems) and drop the libraries into the
+    # ROCm tree so JAX/XLA's collectives use it. Requires the ROCm toolchain
+    # (hipcc) from the extracted tarball -> run AFTER the `rocm` stage.
+    [ -n "${ROCM_PATH:-}" ] || die "ROCM_PATH not set; run the 'rocm' stage first"
+    log "Building RCCL from source @ $RCCL_COMMIT -> $ROCM_PATH/lib"
+    fresh_clone "$RCCL_REPO" rocm-systems
+    ( cd "$SRC_DIR/rocm-systems" \
+        && git checkout "$RCCL_COMMIT" \
+        && cd projects/rccl \
+        && ./install.sh -l --prefix build/ --amdgpu_targets="$PYTORCH_ROCM_ARCH" \
+        && cp -r build/release/librccl* "$ROCM_PATH/lib/" ) || die "RCCL build failed"
+    rm -rf "$SRC_DIR/rocm-systems"
+}
+
+stage_primus() {
+    reload_env
+    log "Installing Primus @ $PRIMUS_BRANCH (kept in $WORKSPACE_DIR/Primus)"
+    rm -rf "$WORKSPACE_DIR/Primus"
+    git clone --recurse-submodules "$PRIMUS_REPO" "$WORKSPACE_DIR/Primus" || die "Primus clone failed"
+    ( cd "$WORKSPACE_DIR/Primus" \
+        && git checkout "$PRIMUS_BRANCH" \
+        && git submodule update --init third_party/maxtext/ ) || die "Primus checkout failed"
+    # The JAX Dockerfile does NOT pip install Primus' (torch-oriented)
+    # requirements.txt; the JAX runtime deps live in requirements-jax.txt
+    # (installed by the `jaxreqs` stage). It also removes stale dataclasses
+    # backports that conflict on modern Python.
+    $PIP uninstall -y dataclasses dataclasses_json || true
+}
+
+stage_jaxreqs() {
+    reload_env
+    # Primus' JAX/MaxText runtime deps (normally installed by the Primus
+    # MaxText pre-train hook at launch time). Front-load them here.
+    local req="$WORKSPACE_DIR/Primus/requirements-jax.txt"
+    if [ -f "$req" ]; then
+        log "Installing Primus JAX requirements from $req"
+        $PIP install -r "$req"
+    else
+        log "Primus requirements-jax.txt not found (skipping); run 'primus' stage first"
+    fi
+}
+
+stage_manifest() {
+    reload_env
+    log "Writing manifest to $WORKSPACE_DIR/.manifest"
+    mkdir -p "$WORKSPACE_DIR/.manifest"
+    env > "$WORKSPACE_DIR/.manifest/env.txt"
+    $PIP list > "$WORKSPACE_DIR/.manifest/requirements.txt"
+    cp "$SCRIPT_DIR/env.sh" "$WORKSPACE_DIR/.manifest/env.sh"
+}
+
+# v26.5 order mirrors the Dockerfile: ROCm tarball -> MaxText -> TF-from-source
+# -> JAX (overrides what MaxText pulled) -> TE -> Primus -> RCCL-from-source.
+DEFAULT_STAGES=(venv rocm maxtext tf_source jax te primus jaxreqs rccl manifest)
+
+run_stage() { local s="$1"; local fn="stage_$s"; declare -F "$fn" >/dev/null || die "unknown stage: $s"; "$fn"; }
+
+main() {
+    if [ "${1:-}" = "--list" ]; then
+        echo "default: ${DEFAULT_STAGES[*]}"
+        echo "note:     the 'te' stage auto-falls-back to a from-source build on glibc < 2.38 hosts (e.g. Ubuntu 22.04)"
+        echo "optional: te_source  (force the from-source TransformerEngine build regardless of glibc)"
+        echo "optional: tf_cpu_fix (lighter alternative to tf_source: pip tensorflow-cpu instead of the ~30-60 min bazel build)"
+        exit 0
+    fi
+    local stages=("$@"); [ ${#stages[@]} -eq 0 ] && stages=("${DEFAULT_STAGES[@]}")
+    log "Base dir: $PRIMUS_JAX_BASE | arch: $PYTORCH_ROCM_ARCH | stages: ${stages[*]}"
+    df -h "$PRIMUS_JAX_BASE" 2>/dev/null | tail -1 || true
+    for s in "${stages[@]}"; do run_stage "$s"; done
+    log "DONE. Activate later with:  source $SCRIPT_DIR/env.sh"
+}
+
+main "$@"
diff --git a/tools/installation/README.md b/tools/installation/README.md
index f87faef4c..b1cf8c1de 100644
--- a/tools/installation/README.md
+++ b/tools/installation/README.md
@@ -1,47 +1,140 @@
 # Primus environment in a venv (no docker, no sudo)
 
-Reproduces the Primus training Dockerfile in a Python virtual environment on
-this server. Same package pins/commits as the Dockerfile, adapted for the
-constraints of this machine.
+Reproduces the Primus **v26.5** training image in a Python virtual environment on
+a bare-metal host. Derived from
+[`.github/workflows/docker-release/Dockerfile.primus-v26.5`](../../.github/workflows/docker-release/Dockerfile.primus-v26.5),
+using the same package pins and commits, adapted for the constraints of a machine
+where we have no root.
+
+A handful of places deliberately diverge from the Dockerfile, because copying it
+exactly produces an environment that does not work outside the container. Each one
+is listed under [Fixes applied](#fixes-applied-gotchas-vs-the-dockerfile) with the
+failure it avoids — worth reading before "correcting" any of them back.
+
+## Python 3.12 is required
+
+This is a hard requirement from upstream packaging, not a preference. The pinned
+torch nightly `2.12.0+rocm7.15.0a20260720` ships a **cp312 Linux wheel only** —
+that nightly date published no cp310 Linux build — and the v26.5
+TransformerEngine wheels are cp312-only as well.
+
+Ubuntu 22.04 hosts only have `python3.10`, so `setup.sh` provisions a standalone
+CPython 3.12 with [`uv`](https://docs.astral.sh/uv/). No sudo, no apt. Interpreters
+land in `$PRIMUS_BASE/python`, keeping the whole environment in one place.
+
+You do not need `uv` beforehand: if it is missing, `setup.sh` downloads it into
+`$PRIMUS_BASE/bin` from `astral.sh`. Pre-install it yourself if that download is
+blocked, or if you would rather not pipe a script into a shell:
+
+```bash
+python3 -m pip install --user uv     # ~/.local/bin, one of the paths setup.sh searches
+# or: curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+`setup.sh` looks for `uv` on `PATH`, then in `$PRIMUS_BASE/bin`, then `~/.local/bin`.
+If you cannot install it at all, supply your own interpreter instead and it is never
+used: `PRIMUS_PYTHON=/path/to/python3.12 bash setup.sh`.
+
+> **Migrating from the v26.4-based scripts:** a Python 3.10 venv cannot be
+> upgraded in place. `setup.sh` detects the mismatch and stops with instructions;
+> remove the old venv and rebuild:
+> `rm -rf "$PRIMUS_BASE/venv" && bash setup.sh`
+
+## TransformerEngine: wheel on glibc ≥ 2.38, otherwise built from source
+
+v26.5 installs TransformerEngine from the ROCm staging index instead of building
+it. Those wheels are produced on Ubuntu 24.04, and `libtransformer_engine.so`
+requires **glibc ≥ 2.38** plus `GLIBCXX_3.4.32`. Ubuntu 22.04 has glibc 2.35, and
+unlike libstdc++, glibc cannot be side-loaded through `LD_LIBRARY_PATH` — so on a
+22.04 host the v26.5 wheels install fine but fail at import with:
+
+```
+OSError: /lib/x86_64-linux-gnu/libc.so.6: version `GLIBC_2.38' not found
+```
+
+`stage_te` therefore picks its install path from the host's glibc:
+
+| Host glibc | Path | What happens |
+|---|---|---|
+| ≥ 2.38 (e.g. Ubuntu 24.04) | wheel | exactly what v26.5 does |
+| < 2.38 (e.g. Ubuntu 22.04) | source | clones `ROCm/TransformerEngine` and builds commit `a07e607f14a…`, the same commit the `TE_VERSION` label refers to |
+
+Either way you get the same TE version; the source build just links against the
+host toolchain. It takes considerably longer than the wheel install. Override the
+choice with `PRIMUS_TE_MODE=wheel|source` (default `auto`). After installing,
+`stage_te` imports `transformer_engine` and fails loudly if it cannot load, rather
+than letting the problem surface later during training.
 
 ## Why it differs from the Dockerfile
 
 | Constraint | Dockerfile | Here |
 |---|---|---|
-| Python | 3.12 (Ubuntu 24.04) | **3.10** (Ubuntu 22.04) — cp310 wheels exist |
+| Python | 3.12 (Ubuntu 24.04) | 3.12, auto-provisioned by `uv` (host has only 3.10) |
 | ROCm | pip `rocm-sdk-devel` | same (self-contained, **no sudo needed**) |
 | GPU arch | gfx942 + gfx950 | **auto-detected** from the host (`rocminfo`/KFD sysfs); builds only what's present → faster builds |
-| Build dir | container FS | venv on **`/it-share-4/envs/primus-env`** (persistent, 1.1 TB free); build sources on local `/tmp` |
+| Build dir | container FS | venv under **`$PRIMUS_BASE`** (persistent, you choose it); build sources on local `/tmp` |
 | System deps | `apt install ...` | **skipped** (no sudo); pip provides what's needed |
+| torchaudio / torchvision / apex | floating (`torchvision==0.27`) | pinned to the exact `a20260720` nightly builds — see below |
 
 The key reason no sudo is required: the `rocm-sdk-devel` pip wheel ships a full
 ROCm toolchain inside the venv, so we don't depend on system ROCm or apt.
 
+### Companion wheels are pinned, the Dockerfile's floating versions are not reproducible
+
+The Dockerfile pins `torch` exactly but lets `torchaudio` and `apex` float and
+resolves torchvision as `==0.27`. That set no longer resolves at all: newer
+`rocm10.x` nightlies now publish matching version numbers, so a floating
+torchvision selects a build for a different ROCm line, conflicts with the pinned
+torch, and pip dies after a long backtracking storm. `setup.sh` pins each
+companion wheel to the version the `a20260720` nightly published, i.e. exactly
+what the Dockerfile picked up on the day it was built.
+
 ## Run it
 
+`PRIMUS_BASE` is required and has no default — the right location is
+site-specific, and quietly falling back to some other host's path only hides the
+problem. Point it at a writable directory on a disk with tens of GB free; it
+holds the venv, the provisioned interpreter, and the Primus/aiter checkouts.
+
 ```bash
 cd tools/installation
-# Default base is /it-share-4/envs/primus-env (persistent shared disk).
-# This default is site-specific; override it for your host:
-#   export PRIMUS_BASE=/some/big/disk/primus-env
-bash setup.sh                 # all default stages
+export PRIMUS_BASE="$HOME/envs/primus-env"   # required, pick your own location
+bash setup.sh                                # all default stages
 ```
 
-This takes a long time (source builds: TransformerEngine, aiter, Primus-Turbo,
-etc.). If any stage fails the script stops immediately and prints which stage
-failed; fix the cause and re-run just that stage. Stages are idempotent:
+The build compiles aiter, Primus-Turbo, flash-attention and TransformerEngine
+from source, so budget hours rather than minutes — on glibc < 2.38 that means all
+of TE, not just its torch glue layer. Running it detached avoids losing it to a
+dropped connection:
 
 ```bash
-bash setup.sh --list          # show stages
-bash setup.sh te              # rebuild just TransformerEngine
+nohup bash setup.sh > ~/primus-setup.log 2>&1 &
+tail -f ~/primus-setup.log
+```
+
+If any stage fails the script stops immediately and prints which stage failed.
+Stages are idempotent, so fix the cause and re-run just that one — exporting the
+same `PRIMUS_BASE` again, since that is how it finds the venv:
+
+```bash
+bash setup.sh --list          # show stages (works without PRIMUS_BASE)
+bash setup.sh te              # reinstall just TransformerEngine
 bash setup.sh venv torch      # venv + torch only
 ```
 
+When cherry-picking stages, note that `te` depends on `torch` and `flash_attn`
+having run first: the TE staging index serves only TE packages, so every other
+dependency must already be installed.
+
 ## Use the environment afterward
 
+Export the SAME `PRIMUS_BASE` you built with, then source `env.sh` — it activates
+the venv and sets every ROCm/NVTE variable. Without `PRIMUS_BASE` it stops with an
+error rather than guessing.
+
 ```bash
-# Use the SAME PRIMUS_BASE you built with (export it first if you overrode the default)
-source tools/installation/env.sh   # activates venv + sets all ROCm/NVTE env vars
+export PRIMUS_BASE="$HOME/envs/primus-env"
+source tools/installation/env.sh
 python -c "import torch; print(torch.cuda.is_available())"
 # Primus is checked out at $WORKSPACE_DIR/Primus
 ```
@@ -50,52 +143,163 @@ python -c "import torch; print(torch.cuda.is_available())"
 
 `venv` → `torch` → `flash_attn` → `te` → `torchtune` → `torchao` → `pydeps`
 → `grouped_gemm` → `causal_conv1d` → `mamba` → `primus` → `aiter` → `turbo`
-→ `boto` → `manifest`
+→ `boto` → `cleanup` → `manifest`
 
 Optional: `torchrec` (DLRM/recommendation stack).
 
+The order matters for `te`: see the note on the staging index below.
+
+## What changed from the v26.4-based scripts
+
+- **TransformerEngine now comes from the ROCm staging index where it can.**
+  `stage_te` installs `transformer_engine_rocm_torch==2.15.0.dev0+rocm7.15.0a20260716.a07e607`
+  from `rocm.frameworks-nightlies.amd.com/whl-staging/device-all/`, which pulls a
+  prebuilt core library and compiles only the thin torch glue layer. On hosts
+  whose glibc is too old for those wheels it builds the same commit from source
+  instead — see the section above.
+- **`stage_te` must run after `stage_flash_attn`.** That staging index serves
+  *only* the transformer-engine packages, so while `--index-url` points at it pip
+  cannot reach PyPI for anything else. Every TE dependency has to be installed
+  beforehand. The Dockerfile gets `einops` transitively from flash-attention and
+  `onnx` from `onnxscript`; `stage_te` installs both explicitly rather than
+  relying on another package's transitive deps.
+- **`ck_jit_compile.sh` is patched** the same way v26.5 patches it inside the
+  image: concurrent CK JIT compiles race to publish the same `.so`, and without
+  the patch the loser aborts the run. The script locates the file through
+  `sysconfig` (rather than a hardcoded `/opt/venv/lib/python3.12` path) and the
+  patch is idempotent.
+- **New `cleanup` stage** uninstalls `tilelang`, mirroring v26.5. `mamba-ssm`
+  pulls it in as a hard dependency and v26.5 removes it again once everything is
+  built.
+- **`apache-tvm-ffi` now follows the Dockerfile pin (0.1.11).** The old scripts
+  held it at 0.1.6 to dodge a `tilelang` import failure; with `tilelang`
+  uninstalled at the end, that workaround is obsolete.
+- **`GPU_ARCHS` is `native` at runtime.** v26.5 sets `ENV GPU_ARCHS=native` after
+  the last build stage so aiter's JIT compiles for the GPU actually in the box.
+  `env.sh` therefore defaults it to `native`, and `setup.sh` overrides it to the
+  full arch list for the individual stages that cross-compile.
+- **Build-only NVTE vars dropped from the runtime environment.** `NVTE_USE_ROCM`,
+  `NVTE_FRAMEWORK`, `NVTE_ROCM_ARCH` and `NVTE_USE_HIPBLASLT` configure the
+  from-source TE build and are gone from v26.5, so `env.sh` no longer exports
+  them; `stage_te_source` still sets them inline for the duration of its build.
+  The runtime performance knobs (`NVTE_CK_*`,
+  `NVTE_USE_CAST_TRANSPOSE_TRITON`, `NVTE_FLASH_ATTN=0`, `NVTE_FUSED_ATTN=1`)
+  stay.
+- **Updated pins:** torch `2.12.0+rocm7.15.0a20260720`, Primus
+  `b511d1b66b0068715308ea9bfe8ba147ea1a3860` (`release/v26.5`), Primus-Turbo
+  `edc8d2ccb0be4888e80ee7c6e765fd3956026a32`. flash-attention, torchtune,
+  torchao, grouped_gemm, causal-conv1d, mamba and aiter are unchanged in v26.5.
+
 ## What is SKIPPED (needs sudo / apt — not reproducible here)
 
 - **AINIC** (`add-apt-repository`, `libionic-dev`): apt-only. Skipped.
-- **UCX + OpenMPI + rocSHMEM**: autotools source builds need `libtool`
-  (missing, apt-only) and RDMA dev headers. Only needed for multi-node
-  rocSHMEM/GDA; single-node training works without them. Skipped.
+- **UCX + OpenMPI**: autotools source builds needing `libtool` and RDMA dev
+  headers, both apt-only. Single-node training works without them. Skipped.
+  (v26.5 removed rocSHMEM from the image entirely, so nothing else needs them.)
+- **DeepEP internode in Primus-Turbo**: needs rocSHMEM, so `stage_turbo` disables
+  it by pointing `ROCSHMEM_HOME` at a non-existent path, and the build prints a
+  warning saying internode DeepEP is off. This is deliberate — left to
+  auto-detect, Primus-Turbo mistakes the pip ROCm SDK directory for a rocSHMEM
+  install (it has rocshmem headers and device bitcode, but no host
+  `librocshmem.a`) and the link then fails. Intranode DeepEP is unaffected.
+  Export a real `ROCSHMEM_HOME` to opt back in.
+- **MLPerf `primus_mllog`**: v26.5 installs it from `training_results_v6.0`. Not
+  part of the core training path, so it is not installed here.
 - **DLRM / FBGEMM / Flux**: not part of the default Primus training path.
-  `torchrec` is provided as an optional stage; FBGEMM needs more apt deps.
+  `torchrec` is provided as an optional stage; FBGEMM additionally needs apt
+  `libtbb-dev`.
 - Misc apt runtime packages (`numactl`, `pciutils`, `libz3-dev`, `ffmpeg`,
   `gfortran`): not installed. Install via sudo later if a specific workload
   needs them.
 
 ## Fixes applied (gotchas vs. the Dockerfile)
 
-These were needed because of Python 3.10 / no-sudo and are baked into the scripts:
+These are needed because of the no-sudo/bare-metal setting and are baked into the
+scripts:
 
+- **pip is constrained so it cannot swap the ROCm GPU stack for a CUDA one.**
+  After installing torch, `setup.sh` writes `$PRIMUS_BASE/pip-constraints.txt`
+  pinning the installed `torch`/`triton`, and every later `pip install` passes
+  `-c`. Many packages depend on a bare `torch`; a resolver that decides to
+  "upgrade" it silently replaces the whole GPU stack with `nvidia-*` wheels.
+  With the constraint, such an attempt fails loudly instead.
+- **Primus-Turbo is installed with `--no-deps`, keeping ROCm's triton.** Its
+  `setup.py` hard-pins upstream `triton==3.7.0`, while torch is built against
+  ROCm's `triton==3.7.1+git…rocm…`. The unconstrained Dockerfile lets the upstream
+  wheel win, but doing that here breaks the environment: ROCm's HIP runtime
+  already loads its own `libLLVM.so.23` from the pip SDK, and the upstream triton
+  wheel bundles a second, statically linked LLVM, so importing it segfaults inside
+  LLVM's static initialisers — taking down `torch._dynamo`, `aiter`, `torchao` and
+  `mamba_ssm` with it. `stage_turbo` therefore installs Primus-Turbo without deps
+  and supplies its real runtime requirements (`scipy`, `flydsl`) explicitly. This
+  is a deliberate, tested deviation from v26.5.
 - **mamba built with pip, not `python setup.py install`.** The legacy
   `easy_install` path ignores pip-installed packages and re-fetches the *latest*
   of every unpinned dep as `.egg`s — it clobbered `transformers` (→5.x), removed
-  `accelerate`/`trl`, and pulled NVIDIA CUDA packages. `stage_mamba` now uses
+  `accelerate`/`trl`, and pulled NVIDIA CUDA packages. `stage_mamba` uses
   `pip install --no-build-isolation .` which respects the pins.
-- **`apache-tvm-ffi` pinned to 0.1.6.** `mamba_ssm` → `tilelang 0.1.8` declares
-  it unpinned; the latest (0.1.12) has a registry bug
-  (`attribute '__dict__' of 'type' objects is not writable`). 0.1.6 works.
-- **`libz3.so` from pip.** `mamba_ssm` needs it at runtime (Dockerfile uses apt
-  `libz3-dev`, unavailable without sudo). `stage_pydeps` installs the `z3-solver`
-  pip package (which ships `libz3.so`), and `env.sh` adds its directory to
-  `LD_LIBRARY_PATH`.
-- **Megatron `helpers_cpp` compiled.** Megatron's dataset indexing imports a
-  pybind11 C++ module (`megatron.core.datasets.helpers_cpp`) that must be built,
-  or training dies with `ModuleNotFoundError: ...helpers_cpp` (surfacing as
-  `MockGPTDataset failed to build`). `stage_primus` runs `make` in
-  `megatron/core/datasets` for both the workspace and `~/.cache/Primus` copies.
+- **`NVTE_CK_IS_V3_ATOMIC_FP32` defaults to `1` on gfx942, not the Dockerfile's `0`.**
+  Paired with `NVTE_CK_USES_BWD_V3=1`, turning fp32 atomics off makes the CK v3
+  backward attention kernel emit Inf gradients on MI300X/MI325X, killing training at
+  the first step. Primus's own tuning guide already prescribes fp32 atomics for these
+  GPUs (the MI300X/MI325X block in
+  [docs/02-user-guide/training-recipes.md](../../docs/02-user-guide/training-recipes.md)),
+  so `env.sh` applies it from the detected architecture and leaves `0` for gfx950,
+  which the Dockerfile value targets. `NVTE_CK_USES_BWD_V3` itself stays at the
+  Dockerfile's `1` — it is worth roughly 15% throughput, and the atomic mode is what
+  makes it safe.
+- **Megatron's `helpers_cpp` is built with an explicit `LIBEXT`, for every
+  checkout.** Its Makefile derives the output filename from `python3-config
+  --extension-suffix`, but a venv does not ship `python3-config`, so it silently
+  falls back to the system interpreter and writes a `cpython-310` name that a 3.12
+  venv will never import. `stage_primus` passes `LIBEXT` from the venv's own
+  `EXT_SUFFIX`, and does so for the workspace clone, `~/.cache/Primus`, *and* the
+  checkout containing these scripts — training is usually launched from the latter.
+  It also deletes any pre-existing extension of the target name first: a checkout
+  that has been bind-mounted into the training container may hold a root-owned
+  Ubuntu 24.04 build that needs `GLIBCXX_3.4.32`, cannot load here, is not
+  writable, and otherwise takes precedence. Either problem surfaces as
+  `MockGPTDataset failed to build as a mock data generator`.
+- **`flydsl` pinned to 0.2.4 so aiter keeps its CK/HIP kernels.** aiter pins
+  `flydsl==0.1.7` and Primus-Turbo requires `flydsl>=0.2.0`, so one of the two is
+  always unsatisfied and Turbo, installing last, wins. aiter only needs
+  `flydsl.expr.vector` at runtime, and that survived until 0.3.0 removed it —
+  after which aiter prints `ROCm/HIP JIT runtime not available … CK and HIP ops
+  are disabled. Triton ops remain available.` and quietly runs Triton-only. 0.2.4
+  is the newest release that satisfies Turbo and keeps aiter whole, and is what
+  the Dockerfile itself resolved to, since 0.3.0 was published after v26.5 was
+  built. `stage_turbo` asserts the symbol is importable afterwards. This is about
+  capability parity with the image, not speed: on llama3.1_8B BF16 an A/B of 0.2.4
+  against 0.3.0 measured 535.9 vs 534.5 TFLOP/s/GPU, i.e. no difference. Other
+  workloads that lean on aiter's CK/HIP kernels are the ones that would notice.
+- **`nvidia-cutlass-dsl` pinned to 4.5.3 so `import mamba_ssm` works.**
+  `mamba_ssm` pins `quack-kernels==0.3.1`, which requires `nvidia-cutlass-dsl>=4.4.1`
+  with no upper bound. 4.6.x restructured the DSL and removed `cute.core.ThrCopy`
+  that quack 0.3.1 imports, so the unpinned resolution ends in
+  `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrCopy'`. 4.5.3 is
+  the newest release quack 0.3.1 still works with, and pip cannot choose it
+  unaided because 4.5.3 was published *after* 4.6.0 yet sorts lower. v26.5 does
+  not pin this, so the released image is affected by the same break.
+- **`patchelf` from pip**, since the apt one is unavailable.
+- **`libz3.so` from pip** (`z3-solver`, added to `LD_LIBRARY_PATH` by `env.sh`).
+  With `tilelang` now uninstalled this is only a safety net, kept so that
+  re-installing `tilelang` by hand does not leave a broken environment.
 - **Megatron** is not pip-installed; it's bundled at
   `$WORKSPACE_DIR/Primus/third_party/Megatron-LM` and added to the path by Primus
   at runtime (or set `PYTHONPATH` yourself for standalone `import megatron`).
 
 ## Caveats
 
-- **Persistence**: the venv lives on `/it-share-4/envs` (persistent). Transient
-  build sources go to local `/tmp` for speed and are deleted after each build.
-- **Disk**: full build needs tens of GB; the shared disk has ~1.1 TB free.
+- **Persistence**: the venv, the provisioned interpreter and the kept checkouts
+  live under `$PRIMUS_BASE` (persistent). Transient build sources go to local
+  `/tmp` for speed and are deleted after each build.
+- **Disk**: a full build needs tens of GB.
+- **Nightly indexes are pruned.** The pinned torch/TE nightlies will eventually
+  disappear from the ROCm index. When that happens, pick a new nightly date that
+  publishes a *complete* cp312 set and update the pin block at the top of
+  `setup.sh` — all of `torch`, `amd-torch-device-*`, `rocm-sdk-*`, `torchaudio`,
+  `torchvision`, `amd-torchvision-device-*` and `apex` should come from the same
+  date.
 - **GPU arch is auto-detected** (`env.sh` reads `rocminfo`, else the kernel KFD
   sysfs `gfx_target_version`), and `stage_torch` installs the matching device
   wheels for whatever it finds (gfx942 and/or gfx950). To force a target — e.g.
diff --git a/tools/installation/env.sh b/tools/installation/env.sh
index 3c7188121..aa29cc4e6 100644
--- a/tools/installation/env.sh
+++ b/tools/installation/env.sh
@@ -1,12 +1,38 @@
 #!/usr/bin/env bash
-# env.sh — Primus venv environment (Python 3.10, ROCm via pip rocm-sdk-devel)
+# env.sh — Primus venv environment (Python 3.12, ROCm via pip rocm-sdk-devel)
+# Derived from .github/workflows/docker-release/Dockerfile.primus-v26.5
+#
 # Source this both during the build (setup.sh does it) and every time you
 # want to USE the environment:   source env.sh
 #
-# NOTE: choose a BASE on the big disk. Home quota here is tiny (~8 GB free),
-# while /tmp lives on a 28 TB device. Override by exporting PRIMUS_BASE first.
+# PRIMUS_BASE is REQUIRED and intentionally has no default: the right location is
+# site-specific, and silently falling back to some other host's path just moves
+# the failure somewhere less obvious. Point it at a writable directory on a disk
+# with tens of GB free; it holds the venv, the provisioned interpreter and the
+# Primus/aiter checkouts.
+
+if [ -z "${PRIMUS_BASE:-}" ]; then
+    echo "[env] ERROR: PRIMUS_BASE is not set." >&2
+    echo "[env]   Export it first, pointing at a writable dir on a big disk:" >&2
+    echo "[env]       export PRIMUS_BASE=\"\$HOME/envs/primus-env\"" >&2
+    echo "[env]   Use the same value for building and for every later 'source env.sh'." >&2
+    # This file is meant to be sourced, where `return` is the only way to stop
+    # without killing the caller's shell. The exit covers running it directly.
+    # shellcheck disable=SC2317
+    return 1 2>/dev/null || exit 1
+fi
 
-export PRIMUS_BASE="${PRIMUS_BASE:-/it-share-4/envs/primus-env}"   # persistent shared disk
+case "$PRIMUS_BASE" in
+    /*) ;;
+    *)
+        echo "[env] ERROR: PRIMUS_BASE must be an absolute path, got '$PRIMUS_BASE'." >&2
+        echo "[env]   Build stages cd into subdirectories, so a relative base breaks them." >&2
+        # shellcheck disable=SC2317
+        return 1 2>/dev/null || exit 1
+        ;;
+esac
+
+export PRIMUS_BASE
 export VENV_DIR="${VENV_DIR:-$PRIMUS_BASE/venv}"
 export WORKSPACE_DIR="${WORKSPACE_DIR:-$PRIMUS_BASE/workspace}"  # kept checkouts (Primus, etc.)
 # Transient build sources: put on fast LOCAL /tmp (NFS is slow for compile I/O;
@@ -16,6 +42,18 @@ export SRC_DIR="${SRC_DIR:-/tmp/primus-build}"
 # Build parallelism.
 export MAX_JOBS="${MAX_JOBS:-128}"
 
+# ---- Python version ----
+# 3.12 is REQUIRED, not a preference. The pinned torch nightly
+# (2.12.0+rocm7.15.0a20260720) published a cp312 Linux wheel and nothing else,
+# so this holds even where TransformerEngine is built from source. In wheel mode
+# TE reinforces it: its prebuilt core library `transformer_engine_rocm7` is also
+# cp312-only, with no sdist and no other cp3xx build. See README.md.
+export PRIMUS_PYTHON_VERSION="${PRIMUS_PYTHON_VERSION:-3.12}"
+# Interpreters provisioned by `uv python install` land under PRIMUS_BASE rather
+# than uv's default in ~/.local/share, so the whole environment stays in one
+# place and off whatever quota the home directory has.
+export UV_PYTHON_INSTALL_DIR="${UV_PYTHON_INSTALL_DIR:-$PRIMUS_BASE/python}"
+
 # ---- Target GPU architecture (auto-detected) ----
 # PYTORCH_ROCM_ARCH controls which gfx targets we build and install device
 # wheels for. If you set it yourself it is respected as-is, e.g.:
@@ -67,15 +105,22 @@ fi
 # Comma-separated variant for the tools/vars that expect commas.
 _ARCH_CSV="${PYTORCH_ROCM_ARCH//;/,}"
 export ROCM_AMDGPU_TARGETS="${ROCM_AMDGPU_TARGETS:-$_ARCH_CSV}"
-export GPU_ARCHS="${GPU_ARCHS:-$PYTORCH_ROCM_ARCH}"
 export HCC_AMDGPU_TARGET="${HCC_AMDGPU_TARGET:-$_ARCH_CSV}"
 export HIP_ARCHITECTURES="${HIP_ARCHITECTURES:-$_ARCH_CSV}"
 
-# Keep caches OFF the tiny home quota (~8 GB). Use local /tmp.
+# GPU_ARCHS is deliberately `native` here, matching the `ENV GPU_ARCHS=native`
+# that v26.5 sets after the last build stage: aiter's JIT has to compile for the
+# GPU it actually runs on. setup.sh overrides it to the full arch list for the
+# individual build stages that cross-compile.
+export GPU_ARCHS="${GPU_ARCHS:-native}"
+
+# Keep build caches off the home directory (they reach tens of GB). Local /tmp is
+# both roomier and much faster than NFS for this.
 export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-cache/pip}"
 export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/primus-cache/xdg}"
 export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-/tmp/primus-cache/triton}"
-mkdir -p "$PIP_CACHE_DIR" "$XDG_CACHE_HOME" "$TRITON_CACHE_DIR" 2>/dev/null || true
+export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/primus-cache/uv}"
+mkdir -p "$PIP_CACHE_DIR" "$XDG_CACHE_HOME" "$TRITON_CACHE_DIR" "$UV_CACHE_DIR" 2>/dev/null || true
 
 # Activate the venv if it exists
 if [ -f "$VENV_DIR/bin/activate" ]; then
@@ -83,12 +128,21 @@ if [ -f "$VENV_DIR/bin/activate" ]; then
     source "$VENV_DIR/bin/activate"
 fi
 
+# Constrain pip so nothing can silently swap the ROCm torch/triton builds for
+# upstream ones. setup.sh writes this file after installing torch; several
+# packages here depend on a bare `torch`, and a resolver that decides to
+# "upgrade" it replaces the whole GPU stack (see README, mamba egg incident).
+# The triton pin matters just as much: ROCm's HIP runtime loads its own
+# libLLVM.so, so an upstream triton (which bundles a second, statically linked
+# LLVM) segfaults on import and takes torch._dynamo/aiter/torchao with it.
+export PRIMUS_PIP_CONSTRAINTS="${PRIMUS_PIP_CONSTRAINTS:-$PRIMUS_BASE/pip-constraints.txt}"
+
 # Workaround for HSA_STATUS_ERROR_OUT_OF_RESOURCES
 export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
 export HSA_NO_SCRATCH_RECLAIM=1
 
 # ---- ROCm path: comes from the pip-installed _rocm_sdk_devel package ----
-# Computed dynamically so it works for any Python minor version (3.10 here).
+# Computed dynamically so it works for any Python minor version.
 if command -v python >/dev/null 2>&1; then
     _ROCM_SDK="$(python - <<'PY' 2>/dev/null
 try:
@@ -116,17 +170,30 @@ if [ -n "${_ROCM_SDK:-}" ]; then
     export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
 fi
 
-# ---- TransformerEngine ROCm settings ----
-export NVTE_USE_HIPBLASLT=1
-export NVTE_FRAMEWORK=pytorch
-export NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH"
-export NVTE_USE_CAST_TRANSPOSE_TRITON=1
-export NVTE_CK_IS_V3_ATOMIC_FP32=0
-export NVTE_CK_USES_BWD_V3=1
-export NVTE_CK_USES_FWD_V3=1
-export CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2
-export NVTE_CK_HOW_V3_BF16_CVT=2
-export NVTE_USE_ROCM=1
+# ---- TransformerEngine tuning ----
+# Only the runtime performance knobs v26.5 keeps. The NVTE_USE_ROCM /
+# NVTE_FRAMEWORK / NVTE_ROCM_ARCH / NVTE_USE_HIPBLASLT vars that v26.4 exported
+# are build-time switches, which is why v26.5 dropped them and why they are not
+# exported here; stage_te_source sets them inline for the duration of its build.
+export NVTE_USE_CAST_TRANSPOSE_TRITON="${NVTE_USE_CAST_TRANSPOSE_TRITON:-1}"
+export NVTE_CK_USES_FWD_V3="${NVTE_CK_USES_FWD_V3:-1}"
+export NVTE_CK_USES_BWD_V3="${NVTE_CK_USES_BWD_V3:-1}"
+export CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT="${CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT:-2}"
+export NVTE_CK_HOW_V3_BF16_CVT="${NVTE_CK_HOW_V3_BF16_CVT:-2}"
+
+# The Dockerfile sets NVTE_CK_IS_V3_ATOMIC_FP32=0, which is correct for gfx950.
+# On gfx942 (MI300X/MI325X) the CK v3 *backward* attention kernel needs fp32
+# atomics: with 0 it produces Inf gradients and llama3.1_8B BF16 pretrain dies at
+# iteration 1 with "found Inf in local grad norm for bucket #0 in backward pass".
+# Primus's own tuning guide prescribes 1 for these GPUs
+# (docs/02-user-guide/training-recipes.md), so default it per detected arch.
+# Measured on MI325X: BWD_V3=1 + fp32 atomics = 614 TFLOP/s/GPU, versus 535 with
+# BWD_V3 disabled, so this is worth ~15% as well as being the correct fix.
+case ";${PYTORCH_ROCM_ARCH};" in
+    *";gfx942;"*) _PRIMUS_CK_ATOMIC_FP32=1 ;;
+    *)            _PRIMUS_CK_ATOMIC_FP32=0 ;;
+esac
+export NVTE_CK_IS_V3_ATOMIC_FP32="${NVTE_CK_IS_V3_ATOMIC_FP32:-$_PRIMUS_CK_ATOMIC_FP32}"
 # Required post-v26.2 to resolve Primus attention backend issues
 export NVTE_FLASH_ATTN=0
 export NVTE_FUSED_ATTN=1
@@ -135,8 +202,10 @@ export NVTE_FUSED_ATTN=1
 export CAUSAL_CONV1D_FORCE_BUILD=TRUE
 export MAMBA_FORCE_BUILD=TRUE
 
-# mamba_ssm needs libz3.so at runtime (Dockerfile gets it from apt libz3-dev;
-# no sudo here, so use the copy shipped by the pip `z3-solver` package).
+# libz3.so safety net. v26.5 uninstalls tilelang (the mamba_ssm dep that needs
+# z3) at the end of the build, so this is normally unused; it stays because the
+# Dockerfile still ships apt libz3-dev and re-installing tilelang by hand should
+# not leave a broken environment behind. pip `z3-solver` supplies the library.
 if command -v python >/dev/null 2>&1; then
     _Z3_LIB="$(python - <<'PY' 2>/dev/null
 import glob, os
@@ -144,5 +213,12 @@ for p in glob.glob(os.path.join(os.path.dirname(__import__("site").getsitepackag
     print(os.path.dirname(p)); break
 PY
 )"
-    [ -n "${_Z3_LIB:-}" ] && export LD_LIBRARY_PATH="$_Z3_LIB:${LD_LIBRARY_PATH:-}"
+    if [ -n "${_Z3_LIB:-}" ]; then
+        export LD_LIBRARY_PATH="$_Z3_LIB:${LD_LIBRARY_PATH:-}"
+    fi
 fi
+
+# Keep the final status clean. setup.sh runs under `set -e` and re-sources this
+# file at the start of every stage, so a falsy test on the last line would abort
+# the build.
+:
diff --git a/tools/installation/setup.sh b/tools/installation/setup.sh
index 3abb6ceeb..a02a7278d 100644
--- a/tools/installation/setup.sh
+++ b/tools/installation/setup.sh
@@ -1,12 +1,16 @@
 #!/usr/bin/env bash
 # setup.sh — Reproduce the Primus training environment in a Python venv
-# (no sudo, no docker). Mirrors the Dockerfile pins, adapted for:
-#   * Python 3.10 (cp310 wheels)
+# (no sudo, no docker). Mirrors the pins of
+#   .github/workflows/docker-release/Dockerfile.primus-v26.5
+# adapted for:
+#   * Python 3.12, auto-provisioned with `uv` (the pinned torch nightly is
+#     cp312-only on Linux; see README.md)
 #   * ROCm provided by the pip `rocm-sdk-devel` wheel (no system ROCm needed)
-#   * Builds on the big /tmp disk (home quota is tiny)
+#   * Transient build sources on the local /tmp disk (fast, and freed afterwards)
 #   * GPU arch auto-detected (gfx942 and/or gfx950); apt/sudo steps skipped
 #
 # Usage:
+#   export PRIMUS_BASE=/big/disk/primus-env   # REQUIRED, no default
 #   bash setup.sh                # run all default stages in order
 #   bash setup.sh <stage>...     # run only specific stage(s), e.g.
 #   bash setup.sh venv torch flash_attn
@@ -16,6 +20,42 @@
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# The Primus checkout these scripts live in. Training is usually launched from
+# here rather than from the workspace clone, so it needs the same treatment.
+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+
+DEFAULT_STAGES=(venv torch flash_attn te torchtune torchao pydeps \
+                grouped_gemm causal_conv1d mamba primus aiter turbo boto \
+                cleanup manifest)
+OPTIONAL_STAGES=(torchrec)
+
+usage() {
+    cat <<EOF
+setup.sh — build the Primus v26.5 training environment in a venv.
+
+PRIMUS_BASE must be exported first; it has no default because the right location
+is site-specific. Point it at a writable dir on a disk with tens of GB free:
+
+    export PRIMUS_BASE="\$HOME/envs/primus-env"
+    bash setup.sh
+
+Run selected stages only (they are idempotent, so this is how you resume after a
+failure):
+
+    bash setup.sh te
+    bash setup.sh venv torch
+
+default stages: ${DEFAULT_STAGES[*]}
+optional stages: ${OPTIONAL_STAGES[*]}
+EOF
+}
+
+# Answer informational flags before sourcing env.sh, so listing the stages does
+# not require PRIMUS_BASE.
+case "${1:-}" in
+    --list|-h|--help) usage; exit 0 ;;
+esac
+
 # shellcheck disable=SC1091
 source "$SCRIPT_DIR/env.sh"
 
@@ -24,14 +64,33 @@ die()  { echo -e "\033[1;31m[setup][ERROR] $*\033[0m" >&2; exit 1; }
 # shellcheck disable=SC1091
 reload_env() { source "$SCRIPT_DIR/env.sh"; }
 
-# ---- pinned versions / commits (from the Dockerfile) ----
-PYTORCH_VERSION="2.12.0+rocm7.14.0a20260608"
+# ---- pinned versions / commits (from Dockerfile.primus-v26.5) ----
 TORCH_INDEX="https://rocm.nightlies.amd.com/whl-multi-arch"
+# The Dockerfile pins only `torch` and lets torchaudio/apex float and
+# torchvision resolve as `==0.27`. That no longer resolves: newer rocm10.x
+# nightlies now publish matching version numbers, so a floating torchvision
+# drags in a build for a different ROCm line and the resolve dies on a
+# backtracking storm. These are the versions the a20260720 nightly published,
+# i.e. the exact set the Dockerfile picked up when it was built.
+PYTORCH_VERSION="2.12.0+rocm7.15.0a20260720"
+ROCM_SDK_VERSION="7.15.0a20260720"
+TORCHAUDIO_VERSION="2.11.0+rocm7.15.0a20260720"
+TORCHVISION_VERSION="0.27.0+rocm7.15.0a20260720"
+APEX_VERSION="1.12.0+rocm7.15.0a20260720"
 
 FA_REPO="https://github.com/ROCm/flash-attention.git"
 FA_BRANCH="6387433156558135a998d5568a9d74c1778666d8"
+# v26.5 stopped building TransformerEngine from source and installs it from the
+# ROCm staging index instead. Those wheels are built on Ubuntu 24.04 though, and
+# libtransformer_engine.so needs glibc >= 2.38, so they cannot load on a 22.04
+# host. stage_te falls back to building the same TE commit from source; see
+# PRIMUS_TE_MODE and the README.
+TE_INDEX="https://rocm.frameworks-nightlies.amd.com/whl-staging/device-all/"
+TE_VERSION="2.15.0.dev0+rocm7.15.0a20260716.a07e607"
+TE_WHEEL_MIN_GLIBC="2.38"
 TE_REPO="https://github.com/ROCm/TransformerEngine.git"
-TE_BRANCH="e6ede467a49cfda1859b145e045109e2f330bccc"
+# The commit the TE_VERSION local label refers to (…a20260716.a07e607).
+TE_COMMIT="a07e607f14a5330807ffdafeeb6224f2d7dffacc"
 TORCHTUNE_REPO="https://github.com/pytorch/torchtune.git"
 TORCHTUNE_BRANCH="b4c98ac2a37f0397d64c22579aed415ce7264db6"
 TORCHAO_REPO="https://github.com/pytorch/ao.git"
@@ -42,17 +101,40 @@ CAUSAL_CONV1D_REPO="https://github.com/Dao-AILab/causal-conv1d"
 CAUSAL_CONV1D_BRANCH="e940ead2fd962c56854455017541384909ca669f"
 MAMBA_REPO="https://github.com/AndreasKaratzas/mamba.git"
 MAMBA_BRANCH="enable-primus-hybrid-models"
-# tilelang 0.1.8 (mamba_ssm dep) breaks with apache-tvm-ffi 0.1.12; pin lower.
-TVM_FFI_VERSION="0.1.6"
+TVM_FFI_VERSION="0.1.11"
+# mamba_ssm pins quack-kernels==0.3.1, which asks for `nvidia-cutlass-dsl>=4.4.1`
+# with no upper bound. 4.6.x restructured the DSL and dropped cute.core.ThrCopy,
+# which quack 0.3.1 imports, so `import mamba_ssm` dies with AttributeError.
+# 4.5.3 is the newest release quack 0.3.1 still works with. Note pip cannot pick
+# it on its own: 4.5.3 was published after 4.6.0 but sorts lower.
+CUTLASS_DSL_VERSION="4.5.3"
 PRIMUS_REPO="https://github.com/AMD-AGI/Primus.git"
-PRIMUS_BRANCH="9a1547cd5885c4de6ad0935a5d59c08303dd0674"
+# Latest commit on `release/v26.5` branch. Committed on 2026-07-22.
+PRIMUS_BRANCH="b511d1b66b0068715308ea9bfe8ba147ea1a3860"
 AITER_REPO="https://github.com/ROCm/aiter.git"
 AITER_COMMIT="0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e"
 TURBO_REPO="https://github.com/AMD-AGI/Primus-Turbo.git"
-TURBO_COMMIT="3974fc246be594d989156dd83e91da618274b0c8"
+# Latest commit on `main` branch. Committed on 2026-07-20.
+TURBO_COMMIT="edc8d2ccb0be4888e80ee7c6e765fd3956026a32"
+# aiter pins `flydsl==0.1.7` and Primus-Turbo wants `flydsl>=0.2.0`, so one of
+# them is always unsatisfied; Turbo installs last and wins. aiter only needs
+# `flydsl.expr.vector` at runtime, which survived until 0.3.0 removed it -- and
+# with it aiter's whole CK/HIP JIT path ("CK and HIP ops are disabled. Triton ops
+# remain available."). 0.2.4 is the newest release that satisfies Turbo and still
+# keeps aiter whole; it is also what the Dockerfile resolved to, since 0.3.0 was
+# published after v26.5 was built.
+FLYDSL_VERSION="0.2.4"
 
 PIP="python -m pip"
 
+# pip install that honours the torch constraint file once stage_torch has
+# written it, so no later resolve can swap the ROCm stack for a CUDA one.
+pipi() {
+    local args=()
+    [ -f "${PRIMUS_PIP_CONSTRAINTS:-}" ] && args+=( -c "$PRIMUS_PIP_CONSTRAINTS" )
+    $PIP install "${args[@]}" "$@"
+}
+
 # Fresh clone helper: clone into transient SRC_DIR, build, then optionally clean.
 fresh_clone() {  # fresh_clone <url> <dir> [extra git clone args...]
     local url="$1"; local dir="$2"; shift 2
@@ -60,12 +142,106 @@ fresh_clone() {  # fresh_clone <url> <dir> [extra git clone args...]
     git clone "$@" "$url" "$SRC_DIR/$dir"
 }
 
+# ======================== PYTHON PROVISIONING ========================
+# v26.5 needs cp312 wheels but Ubuntu 22.04 ships only python3.10 and we have no
+# sudo, so fetch a standalone CPython with uv (bootstrapping uv if needed).
+
+UV_BIN=""
+
+py_is_target() {  # py_is_target <interpreter>
+    local p="${1:-}"
+    [ -n "$p" ] && [ -x "$p" ] || return 1
+    "$p" -c "import sys; raise SystemExit(0 if '%d.%d' % sys.version_info[:2] == '$PRIMUS_PYTHON_VERSION' else 1)" 2>/dev/null
+}
+
+find_uv() {
+    local c
+    for c in "$(command -v uv 2>/dev/null || true)" "$PRIMUS_BASE/bin/uv" "$HOME/.local/bin/uv"; do
+        [ -n "$c" ] && [ -x "$c" ] && { UV_BIN="$c"; return 0; }
+    done
+    return 1
+}
+
+bootstrap_uv() {
+    log "uv not found; installing it into $PRIMUS_BASE/bin (no sudo required)"
+    mkdir -p "$PRIMUS_BASE/bin"
+    curl -LsSf https://astral.sh/uv/install.sh \
+        | env UV_INSTALL_DIR="$PRIMUS_BASE/bin" INSTALLER_NO_MODIFY_PATH=1 sh \
+        || die "could not install uv. Provide Python $PRIMUS_PYTHON_VERSION yourself and re-run with PRIMUS_PYTHON=/path/to/python$PRIMUS_PYTHON_VERSION"
+    UV_BIN="$PRIMUS_BASE/bin/uv"
+}
+
+resolve_python() {
+    if [ -n "${PRIMUS_PYTHON:-}" ]; then
+        py_is_target "$PRIMUS_PYTHON" \
+            || die "PRIMUS_PYTHON=$PRIMUS_PYTHON is not Python $PRIMUS_PYTHON_VERSION"
+        return 0
+    fi
+    local cand
+    cand="$(command -v "python$PRIMUS_PYTHON_VERSION" 2>/dev/null || true)"
+    py_is_target "$cand" && { PRIMUS_PYTHON="$cand"; return 0; }
+    for cand in "$UV_PYTHON_INSTALL_DIR"/cpython-"$PRIMUS_PYTHON_VERSION"*/bin/"python$PRIMUS_PYTHON_VERSION"; do
+        py_is_target "$cand" && { PRIMUS_PYTHON="$cand"; return 0; }
+    done
+    return 1
+}
+
+provision_python() {
+    resolve_python && return 0
+    find_uv || bootstrap_uv
+    log "Provisioning CPython $PRIMUS_PYTHON_VERSION via uv into $UV_PYTHON_INSTALL_DIR"
+    "$UV_BIN" python install "$PRIMUS_PYTHON_VERSION" || die "uv python install $PRIMUS_PYTHON_VERSION failed"
+    resolve_python && return 0
+    PRIMUS_PYTHON="$("$UV_BIN" python find "$PRIMUS_PYTHON_VERSION" 2>/dev/null || true)"
+    py_is_target "${PRIMUS_PYTHON:-}" || die "could not provision Python $PRIMUS_PYTHON_VERSION"
+}
+
+# Pin the GPU stack so a later `pip install` cannot quietly replace the ROCm
+# torch/triton builds with upstream ones (several deps ask for a bare `torch`,
+# and Primus-Turbo asks for an exact upstream `triton`).
+write_pip_constraints() {
+    python - "$PRIMUS_PIP_CONSTRAINTS" <<'PY'
+import sys
+from importlib.metadata import PackageNotFoundError, version
+
+pins = []
+for pkg in ("torch", "triton", "pytorch-triton-rocm"):
+    try:
+        pins.append(f"{pkg}=={version(pkg)}")
+    except PackageNotFoundError:
+        pass
+with open(sys.argv[1], "w") as fh:
+    fh.write("# Generated by tools/installation/setup.sh.\n")
+    fh.write("# Keeps pip from replacing the ROCm GPU stack with upstream builds.\n")
+    fh.write("\n".join(pins) + "\n")
+print("constraints: " + ", ".join(pins))
+PY
+}
+
+# Regenerate the constraint file if it is missing or predates the torch install,
+# so resuming at a single late stage still gets the guard.
+ensure_pip_constraints() {
+    grep -q '^torch==' "${PRIMUS_PIP_CONSTRAINTS:-/nonexistent}" 2>/dev/null && return 0
+    write_pip_constraints
+}
+
 # ============================ STAGES ============================
 
 stage_venv() {
-    log "Creating venv at $VENV_DIR (Python: $(python3 --version))"
     mkdir -p "$PRIMUS_BASE" "$SRC_DIR" "$WORKSPACE_DIR"
-    [ -f "$VENV_DIR/bin/activate" ] || python3 -m venv "$VENV_DIR"
+    provision_python
+    log "Creating venv at $VENV_DIR (interpreter: $PRIMUS_PYTHON, $("$PRIMUS_PYTHON" --version))"
+    if [ -f "$VENV_DIR/bin/python" ]; then
+        local have
+        have="$("$VENV_DIR/bin/python" -c 'import sys; print("%d.%d" % sys.version_info[:2])' 2>/dev/null || echo unknown)"
+        [ "$have" = "$PRIMUS_PYTHON_VERSION" ] || die \
+"existing venv at $VENV_DIR is Python $have, but v26.5 requires $PRIMUS_PYTHON_VERSION.
+  The pinned torch nightly ships a cp312 Linux wheel only, so an older venv
+  cannot be upgraded in place. Remove it and re-run:
+      rm -rf '$VENV_DIR' && bash setup.sh"
+    else
+        "$PRIMUS_PYTHON" -m venv "$VENV_DIR" || die "venv creation failed"
+    fi
     reload_env
     $PIP install --upgrade pip
     # Build front-end tooling. patchelf via pip (system one is missing/no sudo).
@@ -78,13 +254,14 @@ stage_venv() {
         packaging==25.0 \
         setuptools==75.1.0 \
         patchelf
-    rm -rf /root/.cache 2>/dev/null || true
 }
 
 stage_torch() {
     reload_env
     log "Installing PyTorch ${PYTORCH_VERSION} + ROCm SDK (arch: $PYTORCH_ROCM_ARCH) from nightly index"
-    # Early pip deps (Dockerfile block before torch)
+    # Early pip deps (Dockerfile block before torch). apex declares cxxfilt and
+    # pytest requirements that the nightly index does not serve, so these have
+    # to land before the torch install or its resolve fails.
     $PIP install \
         cxxfilt==0.3.0 \
         tqdm==4.67.3 \
@@ -104,8 +281,8 @@ stage_torch() {
     for _arch in "${_arches[@]}"; do
         _arch="${_arch// /}"; [ -z "$_arch" ] && continue
         arch_args+=( "amd-torch-device-${_arch}==${PYTORCH_VERSION}" \
-                     "rocm-sdk-device-${_arch}" \
-                     "amd-torchvision-device-${_arch}" )
+                     "rocm-sdk-device-${_arch}==${ROCM_SDK_VERSION}" \
+                     "amd-torchvision-device-${_arch}==${TORCHVISION_VERSION}" )
     done
     [ ${#arch_args[@]} -gt 0 ] || die "no GPU arch resolved; export PYTORCH_ROCM_ARCH (e.g. gfx942;gfx950)"
     log "Installing device wheels for: $PYTORCH_ROCM_ARCH"
@@ -114,10 +291,10 @@ stage_torch() {
         --index-url "$TORCH_INDEX" \
         --pre \
         "torch==${PYTORCH_VERSION}" \
-        rocm-sdk-devel \
-        torchaudio \
-        torchvision \
-        apex \
+        "rocm-sdk-devel==${ROCM_SDK_VERSION}" \
+        "torchaudio==${TORCHAUDIO_VERSION}" \
+        "torchvision==${TORCHVISION_VERSION}" \
+        "apex==${APEX_VERSION}" \
         "${arch_args[@]}"
 
     log "Running rocm-sdk init"
@@ -125,6 +302,7 @@ stage_torch() {
     reload_env
     [ -n "${ROCM_PATH:-}" ] || die "ROCM_PATH not resolved after rocm-sdk init"
     log "ROCM_PATH=$ROCM_PATH"
+    write_pip_constraints
     python -c "import torch; print('torch', torch.__version__, 'cuda avail', torch.cuda.is_available())"
 }
 
@@ -132,25 +310,133 @@ stage_flash_attn() {
     reload_env
     log "Building flash-attention @ $FA_BRANCH"
     fresh_clone "$FA_REPO" flash-attention --recursive
+    # GPU_ARCHS defaults to `native` for runtime; cross-compile the full set here.
     ( cd "$SRC_DIR/flash-attention" \
         && git checkout "$FA_BRANCH" \
         && git submodule update --init --recursive \
-        && python setup.py install ) || die "flash-attention build failed"
+        && GPU_ARCHS="$PYTORCH_ROCM_ARCH" python setup.py install ) || die "flash-attention build failed"
     rm -rf "$SRC_DIR/flash-attention"
 }
 
-stage_te() {
-    reload_env
-    log "Building TransformerEngine @ $TE_BRANCH"
+# Concurrent CK JIT compiles race to publish the same .so; without this the
+# loser aborts the run. Same patch v26.5 applies inside the image.
+# The single quotes below are deliberate: $_TMP_SO and $OUTPUT are literal text
+# in the target script, not variables to expand here.
+# shellcheck disable=SC2016
+patch_te_ck_jit() {
+    local f
+    f="$(python - <<'PY'
+import os, sysconfig
+for base in {sysconfig.get_paths()[k] for k in ("purelib", "platlib")}:
+    p = os.path.join(base, "transformer_engine", "lib", "ck_jit", "ck_jit_compile.sh")
+    if os.path.exists(p):
+        print(p)
+        break
+PY
+)"
+    if [ -z "$f" ]; then
+        # Both install paths ship this today; tolerate a layout change upstream
+        # rather than failing the build over a race-condition mitigation.
+        log "ck_jit_compile.sh not present, nothing to patch"
+        return 0
+    fi
+    if grep -qF 'mv -n "$_TMP_SO" "$OUTPUT" 2>/dev/null' "$f"; then
+        log "ck_jit_compile.sh already patched"
+        return 0
+    fi
+    sed -i 's|    mv -n "$_TMP_SO" "$OUTPUT"$|    mv -n "$_TMP_SO" "$OUTPUT" 2>/dev/null \|\| true|' "$f"
+    grep -qF 'mv -n "$_TMP_SO" "$OUTPUT" 2>/dev/null' "$f" \
+        || die "ck_jit_compile.sh patch did not apply; upstream changed $f"
+    log "Patched $f"
+}
+
+host_glibc() { ldd --version 2>/dev/null | head -1 | grep -oE '[0-9]+\.[0-9]+$'; }
+
+# version_ge <a> <b> -> true when a >= b
+version_ge() { [ "$(printf '%s\n%s\n' "$2" "$1" | sort -V | head -1)" = "$2" ]; }
+
+# The v26.5 wheels only load where glibc is new enough; otherwise build the same
+# commit locally. Force either path with PRIMUS_TE_MODE=wheel|source.
+resolve_te_mode() {
+    local mode="${PRIMUS_TE_MODE:-auto}"
+    if [ "$mode" != auto ]; then echo "$mode"; return 0; fi
+    local glibc; glibc="$(host_glibc)"
+    if [ -n "$glibc" ] && version_ge "$glibc" "$TE_WHEEL_MIN_GLIBC"; then
+        echo wheel
+    else
+        echo source
+    fi
+}
+
+# TE's own dependencies, needed by both install paths. TE_INDEX serves only the
+# transformer-engine packages, so pip cannot fall back to PyPI for these while
+# --index-url points there; they must already be installed. The Dockerfile gets
+# einops transitively from flash-attention and onnx from onnxscript -- install
+# both explicitly rather than relying on that.
+install_te_deps() {
+    pipi \
+        pybind11==3.0.4 \
+        importlib-metadata==8.7.1 \
+        onnxscript==0.7.0 \
+        pydantic==2.13.4 \
+        nvdlfw_inspect==0.2.2 \
+        einops \
+        onnx
+}
+
+stage_te_wheel() {
+    log "Installing TransformerEngine $TE_VERSION from the ROCm staging index"
+    install_te_deps
+    # transformer_engine_rocm_torch ships as an sdist: the torch glue layer is
+    # compiled here against the installed torch (hence --no-build-isolation),
+    # and it pulls the prebuilt transformer_engine_rocm7 core wheel.
+    MAX_JOBS="$MAX_JOBS" GPU_ARCHS="$PYTORCH_ROCM_ARCH" pipi \
+        --index-url "$TE_INDEX" \
+        --pre \
+        --no-build-isolation \
+        "transformer_engine_rocm_torch==${TE_VERSION}" \
+        || die "TransformerEngine install failed"
+}
+
+stage_te_source() {
+    log "Building TransformerEngine from source @ $TE_COMMIT (glibc $(host_glibc) < $TE_WHEEL_MIN_GLIBC, so the v26.5 wheels cannot load)"
+    # Drop any previously wheel-installed TE, otherwise the unusable prebuilt
+    # core library stays behind and keeps winning the import.
+    $PIP uninstall -y transformer_engine_rocm_torch transformer_engine_rocm7 \
+        transformer_engine transformer_engine_torch >/dev/null 2>&1 || true
+    install_te_deps
+    pipi psutil
     fresh_clone "$TE_REPO" TransformerEngine --recursive
+    # NVTE_FRAMEWORK/NVTE_USE_ROCM/NVTE_ROCM_ARCH/NVTE_USE_HIPBLASLT only matter
+    # for a from-source build, which is why env.sh no longer exports them.
     ( cd "$SRC_DIR/TransformerEngine" \
-        && git checkout "$TE_BRANCH" \
+        && git checkout "$TE_COMMIT" \
         && git submodule update --init --recursive \
-        && $PIP install psutil \
-        && MAX_JOBS="$MAX_JOBS" $PIP install --no-build-isolation . ) || die "TransformerEngine build failed"
+        && MAX_JOBS="$MAX_JOBS" \
+           NVTE_FRAMEWORK=pytorch \
+           NVTE_USE_ROCM=1 \
+           NVTE_USE_HIPBLASLT=1 \
+           NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH" \
+           GPU_ARCHS="$PYTORCH_ROCM_ARCH" \
+           pipi --no-build-isolation . ) || die "TransformerEngine source build failed"
     rm -rf "$SRC_DIR/TransformerEngine"
 }
 
+stage_te() {
+    reload_env
+    local mode; mode="$(resolve_te_mode)"
+    case "$mode" in
+        wheel)  stage_te_wheel ;;
+        source) stage_te_source ;;
+        *) die "PRIMUS_TE_MODE must be one of auto, wheel, source (got '$mode')" ;;
+    esac
+    patch_te_ck_jit
+    python -c "import transformer_engine" 2>/dev/null \
+        || die "TransformerEngine installed but cannot be imported. Re-run with
+  PRIMUS_TE_MODE=source to build it against this host's toolchain:
+      bash setup.sh te"
+}
+
 stage_torchtune() {
     reload_env
     log "Installing torchtune @ $TORCHTUNE_BRANCH (with use_grouped_mm patch)"
@@ -158,7 +444,7 @@ stage_torchtune() {
     ( cd "$SRC_DIR/torchtune" \
         && git checkout "$TORCHTUNE_BRANCH" \
         && sed -i 's/use_grouped_mm = True/use_grouped_mm = False/g' torchtune/modules/moe/utils.py \
-        && $PIP install . ) || die "torchtune install failed"
+        && pipi . ) || die "torchtune install failed"
     rm -rf "$SRC_DIR/torchtune"
 }
 
@@ -170,14 +456,14 @@ stage_torchao() {
         && git checkout "$TORCHAO_BRANCH" \
         && sed -i 's/pad_inner_dim: bool = False/pad_inner_dim: bool = True/g' torchao/float8/config.py \
         && sed -i 's/if defined(HIPBLASLT_VEC_EXT)/if false/g' torchao/csrc/rocm/swizzle/swizzle.cpp \
-        && $PIP install --no-build-isolation . ) || die "torchao build failed"
+        && pipi --no-build-isolation . ) || die "torchao build failed"
     rm -rf "$SRC_DIR/ao"
 }
 
 stage_pydeps() {
     reload_env
     log "Installing main pip dependency set"
-    $PIP install \
+    pipi \
         datasets==3.6.0 \
         av==16.0.1 \
         transformers==4.55.0 \
@@ -218,7 +504,7 @@ stage_grouped_gemm() {
     ( cd "$SRC_DIR/grouped_gemm" \
         && git checkout "$GROUPED_GEMM_BRANCH" \
         && git submodule update --init --recursive \
-        && $PIP install --no-build-isolation . ) || die "grouped_gemm build failed"
+        && pipi --no-build-isolation . ) || die "grouped_gemm build failed"
     rm -rf "$SRC_DIR/grouped_gemm"
 }
 
@@ -228,7 +514,7 @@ stage_causal_conv1d() {
     fresh_clone "$CAUSAL_CONV1D_REPO" causal-conv1d
     ( cd "$SRC_DIR/causal-conv1d" \
         && git checkout "$CAUSAL_CONV1D_BRANCH" \
-        && $PIP install --no-build-isolation . ) || die "causal-conv1d build failed"
+        && pipi --no-build-isolation . ) || die "causal-conv1d build failed"
     rm -rf "$SRC_DIR/causal-conv1d"
 }
 
@@ -241,12 +527,47 @@ stage_mamba() {
     # of every unpinned dep as .egg files, clobbering our pins (it pulled
     # transformers 5.x, removed accelerate/trl, and dragged in NVIDIA CUDA
     # packages). pip respects already-installed versions.
-    # tilelang (a hard mamba_ssm dep) needs an apache-tvm-ffi that predates the
-    # 0.1.12 registry breakage; pin it here.
+    # Install the CUTLASS DSL pin before mamba, not after: quack's `>=4.4.1` is
+    # then already satisfied so pip keeps 4.5.3 instead of pulling 4.6.x and
+    # needing a downgrade, which would strand the 4.6-only libs-core/libs-cu12
+    # packages as orphans that make `pip check` complain.
     ( cd "$SRC_DIR/mamba" \
-        && $PIP install "apache-tvm-ffi==${TVM_FFI_VERSION}" \
-        && $PIP install --no-build-isolation . ) || die "mamba build failed"
+        && pipi "apache-tvm-ffi==${TVM_FFI_VERSION}" \
+        && pipi "nvidia-cutlass-dsl==${CUTLASS_DSL_VERSION}" \
+        && pipi --no-build-isolation . ) || die "mamba build failed"
     rm -rf "$SRC_DIR/mamba"
+    python -c "import mamba_ssm" 2>/dev/null \
+        || die "mamba_ssm cannot be imported. If this is the cute/CUTLASS path,
+  check whether quack-kernels still works with nvidia-cutlass-dsl==${CUTLASS_DSL_VERSION}."
+}
+
+# Megatron's dataset indexing imports a pybind11 extension that has to be
+# compiled, or training dies with `MockGPTDataset failed to build as a mock data
+# generator`. Two traps:
+#  * The Makefile takes the output filename from `python3-config
+#    --extension-suffix`, but a venv does not ship `python3-config`, so it falls
+#    through to the system interpreter and emits a `cpython-310` name that this
+#    3.12 venv will never import. Pass LIBEXT explicitly.
+#  * A checkout that has been bind-mounted into the training container may already
+#    hold a root-owned .so built on Ubuntu 24.04. That one needs GLIBCXX_3.4.32
+#    and cannot load here, yet it takes precedence over anything we build. It is
+#    also not writable by us, so delete it before rebuilding.
+compile_megatron_helpers() {  # compile_megatron_helpers <primus_checkout>
+    local root="$1"
+    local dir="$root/third_party/Megatron-LM/megatron/core/datasets"
+    [ -f "$dir/Makefile" ] || return 0
+    local ext
+    ext="$(python -c 'import sysconfig; print(sysconfig.get_config_var("EXT_SUFFIX"))')"
+    [ -n "$ext" ] || die "could not determine this interpreter's extension suffix"
+    log "Compiling Megatron helpers_cpp in $dir"
+    rm -f "$dir/helpers_cpp$ext"
+    make -C "$dir" LIBEXT="$ext" || die "helpers_cpp build failed in $dir"
+    python - "$dir/helpers_cpp$ext" <<'PY' || die "helpers_cpp built but will not load in $dir"
+import importlib.util, sys
+
+spec = importlib.util.spec_from_file_location("helpers_cpp", sys.argv[1])
+spec.loader.exec_module(importlib.util.module_from_spec(spec))
+PY
 }
 
 stage_primus() {
@@ -257,16 +578,13 @@ stage_primus() {
     ( cd "$WORKSPACE_DIR/Primus" \
         && git checkout "$PRIMUS_BRANCH" \
         && git submodule update --init --recursive \
-        && $PIP install -r requirements.txt ) || die "Primus install failed"
-    # Megatron's dataset indexing needs the pybind11 helpers_cpp extension
-    # compiled, else training fails with:
-    #   ModuleNotFoundError: No module named 'megatron.core.datasets.helpers_cpp'
-    local mlm
-    for mlm in "$WORKSPACE_DIR/Primus/third_party/Megatron-LM" "$HOME/.cache/Primus/third_party/Megatron-LM"; do
-        if [ -f "$mlm/megatron/core/datasets/Makefile" ]; then
-            log "Compiling Megatron helpers_cpp in $mlm"
-            make -C "$mlm/megatron/core/datasets" || die "helpers_cpp build failed in $mlm"
-        fi
+        && pipi -r requirements.txt ) || die "Primus install failed"
+    # Build the Megatron dataset extension for every checkout that might be used.
+    local seen=" " root
+    for root in "$WORKSPACE_DIR/Primus" "$HOME/.cache/Primus" "$REPO_ROOT"; do
+        case "$seen" in *" $root "*) continue ;; esac
+        seen="$seen$root "
+        compile_megatron_helpers "$root"
     done
 }
 
@@ -279,25 +597,68 @@ stage_aiter() {
     ( cd "$WORKSPACE_DIR/aiter" \
         && git checkout "$AITER_COMMIT" \
         && git submodule update --init --recursive \
-        && PREBUILD_KERNELS=3 $PIP install --no-cache-dir --use-pep517 . ) || die "aiter build failed"
+        && PREBUILD_KERNELS=3 GPU_ARCHS="$PYTORCH_ROCM_ARCH" \
+           pipi --no-cache-dir --use-pep517 . ) || die "aiter build failed"
 }
 
 stage_turbo() {
     reload_env
+    ensure_pip_constraints
     log "Building Primus-Turbo @ $TURBO_COMMIT"
+    # Installed with --no-deps on purpose. Primus-Turbo's setup.py hard-pins
+    # upstream `triton==3.7.0`, but torch is built against ROCm's
+    # `triton==3.7.1+git...rocm...`. Letting the upstream wheel win (as the
+    # unconstrained Dockerfile does) breaks this environment: ROCm's HIP runtime
+    # already loads its own libLLVM.so.23 from the pip SDK, and upstream triton
+    # bundles a second, statically linked LLVM, so importing it segfaults inside
+    # LLVM's static initialisers. That takes down torch._dynamo, aiter, torchao
+    # and mamba_ssm. Keeping ROCm's triton fixes all of them, so install without
+    # deps and supply the real runtime requirements explicitly.
+    #
+    # Disable the DeepEP internode path. Primus-Turbo probes for rocSHMEM and,
+    # finding none, falls back to the pip ROCm SDK dir, which does exist and even
+    # ships rocshmem headers and device bitcode but no host `librocshmem.a`. With
+    # a system OpenMPI on PATH that false positive turns the internode link on,
+    # which then fails on the missing librocshmem.a. v26.5 dropped rocSHMEM from
+    # the image entirely, so make the probe fail cleanly instead: it treats a
+    # non-existent ROCSHMEM_HOME as "not found" and warns that internode DeepEP is
+    # disabled. Export ROCSHMEM_HOME yourself if you do have a real rocSHMEM.
+    local rocshmem_home="${ROCSHMEM_HOME:-$PRIMUS_BASE/.rocshmem-not-installed}"
+    if [ -z "${ROCSHMEM_HOME:-}" ] && [ -e "$rocshmem_home" ]; then
+        die "the placeholder path $rocshmem_home exists, so rocSHMEM detection
+  would not be disabled as intended. Delete it, or export ROCSHMEM_HOME to point
+  at a real rocSHMEM install."
+    fi
     fresh_clone "$TURBO_REPO" Primus-Turbo --recursive
     ( cd "$SRC_DIR/Primus-Turbo" \
         && git checkout "$TURBO_COMMIT" \
         && git submodule update --init --recursive \
-        && $PIP install -r requirements.txt \
-        && $PIP install --no-build-isolation . -v ) || die "Primus-Turbo build failed"
+        && pipi -r requirements.txt \
+        && pipi scipy "flydsl==${FLYDSL_VERSION}" \
+        && GPU_ARCHS="$PYTORCH_ROCM_ARCH" \
+           ROCSHMEM_HOME="$rocshmem_home" \
+           pipi --no-build-isolation --no-deps . -v ) || die "Primus-Turbo build failed"
     rm -rf "$SRC_DIR/Primus-Turbo"
+    # aiter silently downgrades itself to Triton-only ops if this symbol is gone,
+    # so fail here instead of losing the CK/HIP kernels at training time.
+    python -c "from flydsl.expr import vector" 2>/dev/null \
+        || die "flydsl ${FLYDSL_VERSION} has no flydsl.expr.vector, so aiter's CK/HIP
+  ops would be disabled. Pick a flydsl that still exports it and satisfies
+  Primus-Turbo's flydsl>=0.2.0."
 }
 
 stage_boto() {
     reload_env
     log "Installing boto3/botocore"
-    $PIP install boto3==1.35.42 botocore==1.35.99
+    pipi boto3==1.35.42 botocore==1.35.99
+}
+
+stage_cleanup() {
+    reload_env
+    # mamba-ssm depends on tilelang, which v26.5 removes again once everything
+    # is built to keep it out of the dependency graph at runtime.
+    log "Uninstalling tilelang (pulled in by mamba-ssm)"
+    $PIP uninstall -y tilelang || true
 }
 
 stage_manifest() {
@@ -306,31 +667,31 @@ stage_manifest() {
     mkdir -p "$WORKSPACE_DIR/.manifest"
     env > "$WORKSPACE_DIR/.manifest/env.txt"
     $PIP list > "$WORKSPACE_DIR/.manifest/requirements.txt"
+    python --version > "$WORKSPACE_DIR/.manifest/python_version"
+    echo "Dockerfile.primus-v26.5" > "$WORKSPACE_DIR/.manifest/derived_from"
     cp "$SCRIPT_DIR/env.sh" "$WORKSPACE_DIR/.manifest/env.sh"
+    [ -f "$PRIMUS_PIP_CONSTRAINTS" ] && cp "$PRIMUS_PIP_CONSTRAINTS" "$WORKSPACE_DIR/.manifest/"
+    log "Environment ready. torch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo '??')"
 }
 
 # ---- Optional stages (DLRM / recommendation stack); not in default run ----
 stage_torchrec() {
     reload_env
     log "Installing torchrec stack (optional)"
-    $PIP install --no-deps torchrec
-    $PIP install tensordict iopath torchmetrics==1.0.3 \
+    pipi --no-deps torchrec
+    pipi tensordict iopath torchmetrics==1.0.3 \
         "git+https://github.com/mlperf/logging.git" \
         --extra-index-url "$TORCH_INDEX"
 }
 
-DEFAULT_STAGES=(venv torch flash_attn te torchtune torchao pydeps \
-                grouped_gemm causal_conv1d mamba primus aiter turbo boto manifest)
-
 run_stage() { local s="$1"; local fn="stage_$s"; declare -F "$fn" >/dev/null || die "unknown stage: $s"; "$fn"; }
 
 main() {
-    if [ "${1:-}" = "--list" ]; then
-        echo "default: ${DEFAULT_STAGES[*]}"; echo "optional: torchrec"; exit 0
-    fi
     local stages=("$@"); [ ${#stages[@]} -eq 0 ] && stages=("${DEFAULT_STAGES[@]}")
-    log "Base dir: $PRIMUS_BASE | arch: $PYTORCH_ROCM_ARCH | stages: ${stages[*]}"
-    df -h "$PRIMUS_BASE" 2>/dev/null | tail -1
+    log "Base dir: $PRIMUS_BASE | python: $PRIMUS_PYTHON_VERSION | arch: $PYTORCH_ROCM_ARCH | stages: ${stages[*]}"
+    # Tolerate a base dir that does not exist yet: under `pipefail` a failing df
+    # would otherwise abort the whole run before the first stage.
+    df -h "$PRIMUS_BASE" 2>/dev/null | tail -1 || true
     for s in "${stages[@]}"; do run_stage "$s"; done
     log "DONE. Activate later with:  source $SCRIPT_DIR/env.sh"
 }
diff --git a/tools/preflight_bisect/bisect.py b/tools/preflight_bisect/bisect.py
new file mode 100644
index 000000000..3c0c9f2c9
--- /dev/null
+++ b/tools/preflight_bisect/bisect.py
@@ -0,0 +1,352 @@
+#!/usr/bin/env python3
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+# Minimal Slurm nodelist bisection for Primus preflight --perf-test.
+# Run from repo root (or any cwd; script resolves repo root for runner/ path).
+#
+# Usage (typical):
+#   export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate
+#   cd /path/to/Primus
+#   python tools/preflight_bisect/bisect.py --nodelist "node[01-32]" -p gpus ...
+#
+# Caveats (see also --help):
+#   - Scale-only hangs: subsets may all PASS while full N fails; suspects may be empty.
+#   - Multiple bad nodes: union of singleton suspects from failing subtrees.
+#   - Tune --trial-timeout-sec to ~2-3x healthy full-N runtime; too short causes false HANG.
+#   - --scancel-user-on-hang kills ALL your Slurm jobs; do not use if you have other work.
+###############################################################################
+from __future__ import annotations
+
+import argparse
+import os
+import signal
+import subprocess
+import sys
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+def _repo_root() -> Path:
+    # tools/preflight_bisect/bisect.py -> repo root is parent.parent.parent
+    return Path(__file__).resolve().parent.parent.parent
+
+
+def expand_nodelist(nodelist: str) -> list[str]:
+    out = subprocess.check_output(
+        ["scontrol", "show", "hostnames", nodelist],
+        text=True,
+        stderr=subprocess.PIPE,
+    )
+    hosts = [ln.strip() for ln in out.splitlines() if ln.strip()]
+    if not hosts:
+        raise SystemExit(f"scontrol produced no hostnames for nodelist={nodelist!r}")
+    return hosts
+
+
+def _format_node_range(nodes: list[str]) -> str:
+    if not nodes:
+        return ""
+    if len(nodes) == 1:
+        return nodes[0]
+    return f"{nodes[0]}..{nodes[-1]}"
+
+
+def _positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed < 1:
+        raise argparse.ArgumentTypeError("must be >= 1")
+    return parsed
+
+
+@dataclass
+class BisectState:
+    max_concurrent_trials: int
+    idx: int = 0
+    trials: list[dict[str, Any]] = field(default_factory=list)
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
+    _trial_slots: threading.BoundedSemaphore = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self._trial_slots = threading.BoundedSemaphore(self.max_concurrent_trials)
+
+    def next_trial_idx(self) -> int:
+        with self._lock:
+            idx = self.idx
+            self.idx += 1
+            return idx
+
+    def record_trial(self, idx: int, nodes: list[str], status: str) -> None:
+        record = {"idx": idx, "n": len(nodes), "status": status, "nodes": list(nodes)}
+        with self._lock:
+            self.trials.append(record)
+
+    def ordered_trials(self) -> list[dict[str, Any]]:
+        with self._lock:
+            return sorted(self.trials, key=lambda trial: trial["idx"])
+
+    def acquire_trial_slot(self) -> None:
+        self._trial_slots.acquire()
+
+    def release_trial_slot(self) -> None:
+        self._trial_slots.release()
+
+
+def run_trial(
+    nodes: list[str],
+    trial_idx: int,
+    state: BisectState,
+    args: argparse.Namespace,
+    runner: Path,
+    out_dir: Path,
+) -> str:
+    """Run one preflight perf trial. Returns 'pass', 'fail', or 'hang'."""
+    subset = ",".join(nodes)
+    log_path = out_dir / f"trial-{trial_idx:03d}.log"
+    cmd: list[str] = [
+        "srun",
+        f"-N{len(nodes)}",
+        f"--nodelist={subset}",
+        "-n",
+        str(len(nodes)),
+        "--ntasks-per-node=1",
+        "-c",
+        str(args.cpus_per_task),
+        f"--gres=gpu:{args.gpus_per_node}",
+        f"-t{args.slurm_time}",
+    ]
+    if args.partition:
+        cmd.extend(["-p", args.partition])
+    # Per-trial env overrides are propagated via srun --export so every rank on
+    # every node sees them (the consolidated primus-cli launcher does accept
+    # --env KEY=VALUE, but only rank 0 would see those; --export covers all
+    # ranks). ALL keeps the caller's environment (notably VENV_ACTIVATE)
+    # intact; the trailing K=V pairs override / add on top. SLURM tokenizes
+    # --export on commas, so values must not contain ',' or whitespace
+    # (NCCL flags never do).
+    export_val = "ALL"
+    if args.preflight_env:
+        export_val += "," + ",".join(args.preflight_env)
+    cmd.append(f"--export={export_val}")
+    cmd.append(str(runner))
+    cmd.extend(
+        [
+            "direct",
+            "--",
+            "preflight",
+            "--perf-test",
+            "--report-file-name",
+            f"trial-{trial_idx:03d}",
+        ]
+    )
+
+    state.acquire_trial_slot()
+    try:
+        header = (
+            f"CMD: {' '.join(cmd)}\n"
+            f"NODES ({len(nodes)}): {subset}\n"
+            f"START: {datetime.now(timezone.utc).isoformat()}\n\n"
+        )
+        print(
+            f"[trial {trial_idx:03d}] N={len(nodes)} {_format_node_range(nodes)} -> {log_path.name}",
+            flush=True,
+        )
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        with log_path.open("wb", buffering=0) as logf:
+            logf.write(header.encode())
+
+            proc = subprocess.Popen(
+                cmd,
+                cwd=str(_repo_root()),
+                stdout=logf,
+                stderr=subprocess.STDOUT,
+                start_new_session=True,
+            )
+            try:
+                rc = proc.wait(timeout=args.trial_timeout_sec)
+                return "pass" if rc == 0 else "fail"
+            except subprocess.TimeoutExpired:
+                try:
+                    os.killpg(proc.pid, signal.SIGKILL)
+                except ProcessLookupError:
+                    pass
+                try:
+                    proc.wait(timeout=60)
+                except subprocess.TimeoutExpired:
+                    pass  # process is stuck in D-state; SIGKILL is pending, move on
+                if args.scancel_user_on_hang:
+                    subprocess.run(
+                        ["scancel", "--signal=KILL", "--user", os.environ.get("USER", "")],
+                        stderr=subprocess.DEVNULL,
+                        stdout=subprocess.DEVNULL,
+                    )
+                return "hang"
+    finally:
+        state.release_trial_slot()
+
+
+def bisect(
+    nodes: list[str],
+    state: BisectState,
+    args: argparse.Namespace,
+    runner: Path,
+    out_dir: Path,
+) -> list[str]:
+    idx = state.next_trial_idx()
+    status = run_trial(nodes, idx, state, args, runner, out_dir)
+    state.record_trial(idx, nodes, status)
+
+    if status == "pass":
+        return []
+    if len(nodes) == 1:
+        return list(nodes)
+
+    mid = len(nodes) // 2
+    if args.max_concurrent_trials == 1:
+        left = bisect(nodes[:mid], state, args, runner, out_dir)
+        right = bisect(nodes[mid:], state, args, runner, out_dir)
+        return left + right
+
+    with ThreadPoolExecutor(max_workers=1, thread_name_prefix="preflight-bisect") as executor:
+        right_future = executor.submit(bisect, nodes[mid:], state, args, runner, out_dir)
+        left = bisect(nodes[:mid], state, args, runner, out_dir)
+        right = right_future.result()
+    return left + right
+
+
+def write_summary(out_dir: Path, nodes: list[str], suspects: list[str], trials: list[dict[str, Any]]) -> None:
+    path = out_dir / "summary.txt"
+    lines = [
+        f"{datetime.now(timezone.utc).isoformat()} bisect nodes={len(nodes)}",
+    ]
+    for t in sorted(trials, key=lambda trial: trial["idx"]):
+        nlist = t["nodes"]
+        r = _format_node_range(nlist)
+        lines.append(f"[{t['idx']:03d}] N={t['n']:2d} {t['status'].upper():4s}  nodes={r}")
+    if suspects:
+        lines.append("SUSPECT_NODES: " + " ".join(sorted(set(suspects))))
+    else:
+        lines.append("SUSPECT_NODES: (none)")
+    summary = "\n".join(lines) + "\n"
+    path.write_text(summary, encoding="utf-8")
+    print(summary, end="")
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(
+        description="Recursively bisect a Slurm nodelist using Primus preflight --perf-test.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Environment:
+  Export VENV_ACTIVATE to your venv activate script before running (required by
+  runner/primus-cli direct -> primus-cli-direct.sh). See docs/preflight-direct.md.
+
+Caveats:
+  - If the hang only reproduces at full scale, all subsets may PASS -> SUSPECT_NODES empty.
+  - Multiple faulty nodes yield a union of suspects from failing singleton trials.
+  - By default, failing sibling subsets launch in parallel (up to 2 concurrent trials).
+  - --trial-timeout-sec too low marks healthy runs as HANG.
+  - --scancel-user-on-hang cancels ALL jobs for $USER; only use it with --max-concurrent-trials=1.
+""",
+    )
+    p.add_argument("--nodelist", required=True, help='Slurm nodelist expression, e.g. "node[01-32]"')
+    p.add_argument("-p", "--partition", default="", help="Slurm partition (-p), optional")
+    p.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("bisect-out"),
+        help="Directory for trial-*.log and summary.txt (default: ./bisect-out)",
+    )
+    p.add_argument(
+        "--trial-timeout-sec",
+        type=int,
+        default=900,
+        help="Wall-clock timeout per trial in seconds (default: 900)",
+    )
+    p.add_argument(
+        "--slurm-time",
+        default="00:45:00",
+        help="srun -t limit per trial (default: 00:45:00)",
+    )
+    p.add_argument(
+        "--max-concurrent-trials",
+        type=_positive_int,
+        default=2,
+        help="Maximum concurrent subset trials (default: 2). Set to 1 to force sequential execution.",
+    )
+    p.add_argument("--cpus-per-task", type=int, default=128, help="srun -c (default: 128)")
+    p.add_argument(
+        "--gpus-per-node",
+        type=int,
+        default=8,
+        help="GPUs per node; emitted as srun --gres=gpu:N (default: 8)",
+    )
+    p.add_argument(
+        "--preflight-env",
+        action="append",
+        default=[],
+        metavar="KEY=VALUE",
+        help=(
+            "Repeatable; propagated into each trial via "
+            "'srun --export=ALL,KEY=VALUE,...'. Values must not contain ',' or whitespace."
+        ),
+    )
+    p.add_argument(
+        "--runner",
+        type=Path,
+        default=None,
+        help=(
+            "Override path to the launcher invoked per trial "
+            "(default: repo runner/primus-cli). The bisector always appends "
+            "'direct -- preflight --perf-test --report-file-name ...' after "
+            "the runner path; custom runners that ignore positional args "
+            "(e.g. fake_runner.sh) work transparently."
+        ),
+    )
+    p.add_argument(
+        "--scancel-user-on-hang",
+        action="store_true",
+        help="On timeout, also run: scancel --signal=KILL --user $USER (DANGEROUS)",
+    )
+    args = p.parse_args()
+    if args.scancel_user_on_hang and args.max_concurrent_trials > 1:
+        p.error("--scancel-user-on-hang is only supported with --max-concurrent-trials=1")
+
+    if not os.environ.get("VENV_ACTIVATE"):
+        print(
+            "ERROR: VENV_ACTIVATE is not set. Export the path to your venv's "
+            "bin/activate before running, e.g.\n"
+            "  export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate",
+            file=sys.stderr,
+        )
+        return 2
+
+    runner = args.runner or (_repo_root() / "runner" / "primus-cli")
+    if not runner.is_file():
+        print(f"ERROR: runner not found: {runner}", file=sys.stderr)
+        return 2
+
+    try:
+        hosts = expand_nodelist(args.nodelist)
+    except subprocess.CalledProcessError as e:
+        print(f"ERROR: scontrol failed: {e}", file=sys.stderr)
+        return 2
+    except FileNotFoundError:
+        print("ERROR: scontrol not found; run this script on a Slurm login/head node.", file=sys.stderr)
+        return 2
+
+    out_dir = args.output_dir.resolve()
+    state = BisectState(max_concurrent_trials=args.max_concurrent_trials)
+    suspects = bisect(hosts, state, args, runner, out_dir)
+    write_summary(out_dir, hosts, suspects, state.ordered_trials())
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/tools/preflight_bisect/fake_runner.sh b/tools/preflight_bisect/fake_runner.sh
new file mode 100755
index 000000000..4c7fb9127
--- /dev/null
+++ b/tools/preflight_bisect/fake_runner.sh
@@ -0,0 +1,54 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+# fake_runner.sh — drop-in runner replacement for bisect.py testing.
+#
+# Simulates a bad node without running a real preflight, so that
+# bisect.py's identification logic can be validated on a live cluster.
+#
+# Environment variables (pass via bisect.py --preflight-env):
+#   BAD_NODE      hostname of the node that should appear faulty (required)
+#   HANG_SECONDS  if >0, the bad node sleeps this long before exiting 1,
+#                 which lets you exercise the --trial-timeout-sec hang path
+#                 (default: 0 → instant failure)
+#   GOOD_SLEEP    seconds a healthy node sleeps to simulate a real run
+#                 (default: 2)
+#
+# Usage with bisect.py:
+#   python tools/preflight_bisect/bisect.py \
+#       --nodelist "node001,node002" \
+#       --partition mi355x \
+#       --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \
+#       --trial-timeout-sec 600 \
+#       --slurm-time 00:15:00 \
+#       --runner tools/preflight_bisect/fake_runner.sh \
+#       --preflight-env BAD_NODE=node002
+###############################################################################
+set -euo pipefail
+
+BAD_NODE="${BAD_NODE:-}"
+HANG_SECONDS="${HANG_SECONDS:-0}"
+GOOD_SLEEP="${GOOD_SLEEP:-2}"
+
+THIS_HOST="$(hostname -s)"
+
+if [[ -z "$BAD_NODE" ]]; then
+    echo "[fake_runner] ERROR: BAD_NODE is not set. Pass --preflight-env BAD_NODE=<hostname> to bisect.py." >&2
+    exit 2
+fi
+
+if [[ "$THIS_HOST" == "$BAD_NODE" ]]; then
+    echo "[fake_runner] $THIS_HOST == BAD_NODE ($BAD_NODE): simulating failure"
+    if (( HANG_SECONDS > 0 )); then
+        echo "[fake_runner] sleeping ${HANG_SECONDS}s to simulate a hang"
+        sleep "$HANG_SECONDS"
+    fi
+    exit 1
+fi
+
+echo "[fake_runner] $THIS_HOST != BAD_NODE ($BAD_NODE): simulating healthy run (sleep ${GOOD_SLEEP}s)"
+sleep "$GOOD_SLEEP"
+exit 0